Skip to content

vllm.models.deepseek_v41.nvidia.ops.o_proj

DSV4.1 output projection with a small-batch SM100/SM103 fusion.

Functions:

_can_fuse_wo_a(layer)

Whether the attention layer matches the fused WO-A kernel's layout.

Source code in vllm/models/deepseek_v41/nvidia/ops/o_proj.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def _can_fuse_wo_a(layer: nn.Module) -> bool:
    """Whether the attention layer matches the fused WO-A kernel's layout."""
    # The fused kernel never materializes the BF16 input an adapter on wo_a
    # needs, so such a layer takes ``deep_gemm_fp8_o_proj`` instead.
    if getattr(layer.wo_a, "requires_unquantized_input", False):
        return False
    method = getattr(layer.wo_b, "scheme", layer.wo_b.quant_method)
    return (
        current_platform.is_device_capability_family(100)
        # WO-A is FP8 with per-32 scales, and WO-B takes MXFP8 input directly.
        and layer._einsum_recipe == (1, 1, 32)
        and get_input_quant_key(layer.wo_b) == kMxfp8Dynamic
        and isinstance(
            getattr(method, "kernel", None),
            (FlashInferCutedslMxfp8LinearKernel, TorchMxFp8LinearKernel),
        )
        # A group's heads form one portable (<= 8 CTA) cluster.
        and layer.n_local_heads // layer.n_local_groups <= 8
        and (layer.nope_head_dim, layer.rope_head_dim) == (448, 64)
        and layer.o_lora_rank % 128 == 0
    )

_fused_wo_a_shape(layer)

The fused WO-A kernel's warmup arguments for layer.

Source code in vllm/models/deepseek_v41/nvidia/ops/o_proj.py 
55
56
57
58
59
60
61
62
def _fused_wo_a_shape(layer: nn.Module) -> dict[str, int]:
    """The fused WO-A kernel's warmup arguments for ``layer``."""
    return dict(
        max_tokens=_fused_wo_a_max_tokens(layer),
        n_groups=layer.n_local_groups,
        heads_per_group=layer.n_local_heads // layer.n_local_groups,
        o_lora_rank=layer.o_lora_rank,
    )

dsv41_o_proj(layer, attn_out, positions, wo_b_proj=None, early_trigger=False)

deep_gemm_fp8_o_proj with WO-A fused for small SM100/SM103 batches; wo_b_proj replaces layer._wo_b_proj. With early_trigger, the fused WO-A launches wo_b_proj (PDL) right after its TMEM allocation, for a wo_b_proj that loads its weights before its grid wait.

Source code in vllm/models/deepseek_v41/nvidia/ops/o_proj.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def dsv41_o_proj(
    layer: nn.Module,
    attn_out: torch.Tensor,
    positions: torch.Tensor,
    wo_b_proj: Callable[..., torch.Tensor] | None = None,
    early_trigger: bool = False,
) -> torch.Tensor:
    """``deep_gemm_fp8_o_proj`` with WO-A fused for small SM100/SM103 batches;
    ``wo_b_proj`` replaces ``layer._wo_b_proj``. With ``early_trigger``, the
    fused WO-A launches ``wo_b_proj`` (PDL) right after its TMEM allocation,
    for a ``wo_b_proj`` that loads its weights before its grid wait."""
    wo_b_proj = wo_b_proj or layer._wo_b_proj
    o = attn_out[:, : layer.n_local_heads, :]
    cos_sin_cache = layer.rotary_emb.cos_sin_cache
    if 1 <= o.shape[0] <= _fused_wo_a_max_tokens(layer) and _can_fuse_wo_a(layer):
        from .fused_wo_a import _FUSED_WO_A_KERNEL

        q, scales = _FUSED_WO_A_KERNEL(
            x=o,
            positions=positions,
            rope=cos_sin_cache,
            weight=layer.wo_a.weight,
            weight_scale=layer.wo_a.weight_scale,
            early_trigger=early_trigger,
        )
        return wo_b_proj(
            QuantizedActivation(q, scales, o.dtype, q.shape, kMxfp8Dynamic)
        )
    z = layer.wo_a.grouped_forward(
        o,
        positions,
        cos_sin_cache,
        n_groups=layer.n_local_groups,
        heads_per_group=layer.n_local_heads // layer.n_local_groups,
        nope_dim=layer.nope_head_dim,
        rope_dim=layer.rope_head_dim,
        o_lora_rank=layer.o_lora_rank,
        einsum_recipe=layer._einsum_recipe,
        tma_aligned_scales=layer._tma_aligned_scales,
    )
    return wo_b_proj(z.flatten(1))

precompile_dsv41_o_proj_early_trigger(layer, max_tokens)

Compile the fused WO-A variant dsv41_o_proj dispatches with early_trigger, for its token counts up to max_tokens.

Source code in vllm/models/deepseek_v41/nvidia/ops/o_proj.py 
65
66
67
68
69
70
71
72
73
def precompile_dsv41_o_proj_early_trigger(layer: nn.Module, max_tokens: int) -> None:
    """Compile the fused WO-A variant ``dsv41_o_proj`` dispatches with
    ``early_trigger``, for its token counts up to ``max_tokens``."""
    if _can_fuse_wo_a(layer):
        from .fused_wo_a import _FUSED_WO_A_KERNEL

        shape = _fused_wo_a_shape(layer)
        shape["max_tokens"] = min(shape["max_tokens"], max_tokens)
        _FUSED_WO_A_KERNEL.warmup(**shape, early_trigger=True)

register_dsv41_o_proj_warmup(layer)

Warm the fused WO-A token counts dsv41_o_proj can dispatch.

Source code in vllm/models/deepseek_v41/nvidia/ops/o_proj.py 
76
77
78
79
80
81
82
83
def register_dsv41_o_proj_warmup(layer: nn.Module) -> None:
    """Warm the fused WO-A token counts ``dsv41_o_proj`` can dispatch."""
    # Draft models load outside the warmup registry; their layers share the
    # target's attention shapes, so the target's registration covers them too.
    if _can_fuse_wo_a(layer):
        from .fused_wo_a import _FUSED_WO_A_KERNEL

        _FUSED_WO_A_KERNEL.register_warmup(**_fused_wo_a_shape(layer))