Skip to content

vllm.models.deepseek_v41.nvidia.specialized_decode.attention.aux_projections

Low-batch GEMMs of the attention side projections.

Next to fused_wqa_wkv, the kv-source layers project the attention input through the compressor's fused_wkv_wgate (FP32 kv_score) and the index-source layers through the indexer's weights_proj. At decode sizes cuBLAS runs both as a split-K GEMM plus a reduce kernel, longer than fused_wqa_wkv itself. For low-batch FULL-graph steps they run the CuTe DSL ll_bf16 kernels instead: FP32 x @ W.T that differs from cuBLAS only in the K summation order. The indexer's wk (compressor latent to index key) takes the same kernels.

Functions:

aux_projection_gemm_accepts(weights)

Whether aux_projection_gemm takes every one of weights.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/aux_projections.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
def aux_projection_gemm_accepts(weights: list[torch.Tensor]) -> bool:
    """Whether ``aux_projection_gemm`` takes every one of ``weights``."""
    return (
        bool(weights)
        and ll_bf16.is_available()
        and all(
            w.dtype == torch.bfloat16
            and w.dim() == 2
            and w.is_contiguous()
            and w.shape[1] % 8 == 0
            for w in weights
        )
    )

warmup_aux_projection_gemms(weights)

Compile the kernels of every row count before graph capture.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/aux_projections.py 
83
84
85
86
def warmup_aux_projection_gemms(weights: Iterable[torch.Tensor]) -> None:
    """Compile the kernels of every row count before graph capture."""
    shapes = sorted({(w.shape[1], w.shape[0]) for w in weights})
    aux_projection_gemm.warmup(shapes=shapes, m_values=range(1, LOW_BATCH_MAX_ROWS + 1))