Skip to content

vllm.models.deepseek_v4.nvidia.ops.o_proj

Functions:

compute_fp8_einsum_recipe(block_size=128)

fp8_einsum recipe + scale layout for the current GPU arch.

SM90 keeps block-row FP32 scales. SM100 uses packed per-row E8M0 scales.

Returns (einsum_recipe, tma_aligned_scales) for deep_gemm_fp8_o_proj.

Source code in vllm/models/deepseek_v4/nvidia/ops/o_proj.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def compute_fp8_einsum_recipe(
    block_size: int = 128,
) -> tuple[tuple[int, int, int], bool]:
    """fp8_einsum recipe + scale layout for the current GPU arch.

    SM90 keeps block-row FP32 scales. SM100 uses packed per-row E8M0 scales.

    Returns ``(einsum_recipe, tma_aligned_scales)`` for ``deep_gemm_fp8_o_proj``.
    """
    cap = current_platform.get_device_capability()
    assert cap is not None, "DeepseekV4 attention requires a CUDA device"
    einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, block_size)
    tma_aligned_scales = cap.major >= 10
    return einsum_recipe, tma_aligned_scales

deep_gemm_fp8_o_proj(o, positions, cos_sin_cache, wo_a, wo_b, *, n_groups, heads_per_group, nope_dim, rope_dim, o_lora_rank, einsum_recipe, tma_aligned_scales)

O projection: inverse RoPE + grouped wo_a + wo_b.

Shared by the FlashMLA and FlashInfer CUDA backends. The attention layer selects the recipe at initialization. wo_b is any callable over the flattened z: the projection module itself, or a wrapper that also reduce-scatters its output (DeepSeek-V4.1 GEMM-RS).

A QuantizedActivation o was already rotated and cast by the attention kernel, one wo_a group per slot with the live groups first, so only the padding groups are sliced off.

Source code in vllm/models/deepseek_v4/nvidia/ops/o_proj.py 
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def deep_gemm_fp8_o_proj(
    o: torch.Tensor | QuantizedActivation,
    positions: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    wo_a: nn.Module,
    wo_b: Callable[[torch.Tensor], torch.Tensor],
    *,
    n_groups: int,
    heads_per_group: int,
    nope_dim: int,
    rope_dim: int,
    o_lora_rank: int,
    einsum_recipe: tuple[int, int, int],
    tma_aligned_scales: bool,
) -> torch.Tensor:
    """O projection: inverse RoPE + grouped wo_a + wo_b.

    Shared by the FlashMLA and FlashInfer CUDA backends. The attention
    layer selects the recipe at initialization. ``wo_b`` is any callable over
    the flattened ``z``: the projection module itself, or a wrapper that also
    reduce-scatters its output (DeepSeek-V4.1 GEMM-RS).

    A QuantizedActivation ``o`` was already rotated and cast by the attention
    kernel, one ``wo_a`` group per slot with the live groups first, so only
    the padding groups are sliced off.
    """
    use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn
    if isinstance(o, QuantizedActivation):
        assert use_fp8
        o_proj_input, o_scale = o.data[:, :n_groups], o.scale[:, :n_groups]
    else:
        o_proj_input, o_scale = fused_inv_rope_fp8_quant(
            o,
            positions,
            cos_sin_cache,
            n_groups=n_groups,
            heads_per_group=heads_per_group,
            nope_dim=nope_dim,
            rope_dim=rope_dim,
            quant_group_size=einsum_recipe[2],
            tma_aligned_scales=tma_aligned_scales,
            quantize=use_fp8,
        )
    z = torch.empty(
        (o_proj_input.shape[0], n_groups, o_lora_rank),
        device=o_proj_input.device,
        dtype=torch.bfloat16,
    )
    if use_fp8:
        weight_scale = (
            wo_a.weight_scale
            if hasattr(wo_a, "weight_scale")
            else wo_a.weight_scale_inv
        )
        fp8_einsum(
            "bhr,hdr->bhd",
            (o_proj_input, o_scale),
            (wo_a.weight, weight_scale),
            z,
            recipe=einsum_recipe,
        )
    else:
        grouped_weight = wo_a.weight.view(n_groups, o_lora_rank, -1)
        torch.bmm(
            o_proj_input.transpose(0, 1),
            grouped_weight.transpose(1, 2),
            out=z.transpose(0, 1),
        )
    return wo_b(z.flatten(1))

rope_quant_attn_out(num_tokens, device)

Allocate the output pair the RopeQuant kernel writes.

Values are group-major, [num_tokens, 16, 8 * 512] over physical [group, token, K]. Scales are MN-major packed UE8M0, one int32 per head over its four 128-wide blocks. Both are what fp8_einsum consumes.

Source code in vllm/models/deepseek_v4/nvidia/ops/o_proj.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def rope_quant_attn_out(num_tokens: int, device: torch.device) -> QuantizedActivation:
    """Allocate the output pair the RopeQuant kernel writes.

    Values are group-major, ``[num_tokens, 16, 8 * 512]`` over physical
    ``[group, token, K]``. Scales are MN-major packed UE8M0, one int32 per
    head over its four 128-wide blocks. Both are what ``fp8_einsum`` consumes.
    """
    n_groups = _ROPE_QUANT_NUM_HEADS // _ROPE_QUANT_HEADS_PER_GROUP
    group_width = _ROPE_QUANT_HEADS_PER_GROUP * _ROPE_QUANT_HEAD_DIM
    data = torch.empty(
        (n_groups, num_tokens, group_width), dtype=torch.float8_e4m3fn, device=device
    ).transpose(0, 1)
    aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize)
    scale = torch.empty(
        (n_groups, group_width // (_ROPE_QUANT_BLOCK * 4), aligned),
        dtype=torch.int32,
        device=device,
    ).permute(2, 0, 1)[:num_tokens]
    return QuantizedActivation(
        data=data,
        scale=scale,
        orig_dtype=torch.bfloat16,
        orig_shape=data.shape,
        quant_key=kFp8Dynamic128Sym,
    )

rope_quant_unsupported_reason(layer)

Why layer cannot take FlashInfer's fused inverse RoPE + FP8 output.

DSv4 always has 8 heads per wo_a group at any TP size, and the rest of the kernel's fixed shape (head/RoPE dims, per-128 einsum scales, FP32 RoPE cache) is what the SM100 layer uses anyway. What varies is the local head count and the configs the layer also serves unquantized.

Source code in vllm/models/deepseek_v4/nvidia/ops/o_proj.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def rope_quant_unsupported_reason(layer: torch.nn.Module) -> str | None:
    """Why ``layer`` cannot take FlashInfer's fused inverse RoPE + FP8 output.

    DSv4 always has 8 heads per ``wo_a`` group at any TP size, and the rest of
    the kernel's fixed shape (head/RoPE dims, per-128 einsum scales, FP32 RoPE
    cache) is what the SM100 layer uses anyway. What varies is the local head
    count and the configs the layer also serves unquantized.
    """
    if layer.n_local_heads != _ROPE_QUANT_NUM_HEADS:
        # Padding fewer heads up to the kernel's fixed 128 costs more
        # attention than the fused epilogue saves.
        return f"RopeQuant needs {_ROPE_QUANT_NUM_HEADS} local query heads"
    if layer.kv_cache_torch_dtype != torch.float8_e4m3fn:
        return "RopeQuant needs the per-tensor FP8 KV cache"
    if layer.wo_a.weight.dtype != torch.float8_e4m3fn:
        return "RopeQuant needs an FP8 wo_a"
    return None