Skip to content

vllm.models.deepseek_v41.nvidia.specialized_decode.attention.mxfp8

Low-batch MXFP8 attention projections.

At up to LOW_BATCH_MAX_ROWS rows the executor runs fused_wqa_wkv, wq_b and wo_b on FlashInfer mm_mxfp8 with its cutedsl_low_latency backend whenever their input is already MXFP8; the stock layers keep cute-dsl. Both backends read the same weights, weight scales and F8_128x4 activation scales.

The backend's tactics differ by under 1 us when tuned in isolation, but split-K tactics are slower inside the back-to-back PDL chain of a decode step, so the 5-8-row bucket of the GB300 TP4 shapes is tuned over one pinned tactic.

Functions:

  • low_latency_mm –

    layer's GEMM of x on cutedsl_low_latency for an eligible

  • mxfp8_eligible –

    Whether layer is a cute-dsl MXFP8 linear whose weights the

  • mxfp8_mm –

    layer's GEMM on cutedsl_low_latency.

  • pinned_tactics –

    Offer FlashInfer's tuner only the pinned tactic for pinned problems. The

  • run_zeros –

    Call linear on an input laid out like the fused quantize producers'.

  • shape_representatives –

    One eligible linear per distinct (N, K), in a rank-independent order.

_mxfp8_kernel(layer)

layer's kernel when it is exactly the cute-dsl MXFP8 one.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py 
77
78
79
80
81
def _mxfp8_kernel(layer: nn.Module) -> FlashInferCutedslMxfp8LinearKernel | None:
    """``layer``'s kernel when it is exactly the ``cute-dsl`` MXFP8 one."""
    method = getattr(layer, "scheme", getattr(layer, "quant_method", None))
    kernel = getattr(method, "kernel", None)
    return kernel if type(kernel) is FlashInferCutedslMxfp8LinearKernel else None

low_latency_mm(layer, x)

layer's GEMM of x on cutedsl_low_latency for an eligible layer, or None when x is not an MXFP8 input of the backend's layout.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
def low_latency_mm(
    layer: nn.Module, x: torch.Tensor | QuantizedActivation
) -> torch.Tensor | None:
    """``layer``'s GEMM of ``x`` on ``cutedsl_low_latency`` for an eligible
    layer, or None when ``x`` is not an MXFP8 input of the backend's layout."""
    kernel = _mxfp8_kernel(layer)
    assert kernel is not None
    qa = as_quantized_activation(x, kernel.input_quant_key())
    if qa is None or qa.scale.ndim != 1 or not qa.data.is_contiguous():
        return None
    return mxfp8_mm(layer, qa)

mxfp8_eligible(layer)

Whether layer is a cute-dsl MXFP8 linear whose weights the low-latency backend also takes.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py 
84
85
86
87
88
89
90
91
92
93
def mxfp8_eligible(layer: nn.Module) -> bool:
    """Whether ``layer`` is a ``cute-dsl`` MXFP8 linear whose weights the
    low-latency backend also takes."""
    kernel = _mxfp8_kernel(layer)
    return (
        kernel is not None
        and kernel.config.weight_shape[1] % 128 == 0
        and getattr(layer, "bias", None) is None
        and _backend_supported()
    )

mxfp8_mm(layer, qa)

layer's GEMM on cutedsl_low_latency.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py 
51
52
53
54
55
56
57
58
59
60
61
def mxfp8_mm(layer: nn.Module, qa: QuantizedActivation) -> torch.Tensor:
    """``layer``'s GEMM on ``cutedsl_low_latency``."""
    out = vllm_flashinfer.mm_mxfp8(
        qa.data,
        layer.weight,
        qa.scale,
        layer.weight_scale,
        out_dtype=qa.orig_dtype,
        backend=BACKEND,
    )
    return out.view(*qa.orig_shape[:-1], out.shape[-1])

pinned_tactics()

Offer FlashInfer's tuner only the pinned tactic for pinned problems. The tuner's tactic list is a FlashInfer internal: off the checked version, every tactic is tuned.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.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
159
160
161
162
163
164
165
166
167
168
169
170
171
@contextmanager
def pinned_tactics() -> Iterator[None]:
    """Offer FlashInfer's tuner only the pinned tactic for pinned problems. The
    tuner's tactic list is a FlashInfer internal: off the checked version,
    every tactic is tuned."""
    if not (
        current_platform.is_device_capability(103)
        and uses_checked_flashinfer("pinning the MXFP8 GEMM tactics")
    ):
        yield
        return
    module = importlib.import_module(
        "flashinfer.gemm.kernels.cute_dsl.low_latency_blockscaled_gemm"
    )
    original: Callable[..., list[Tactic]] | None = getattr(
        module, "autotune_tactics", None
    )
    if original is None:
        yield
        return
    enumerate_tactics = original

    def autotune_tactics(
        problem_sizes_mnkl: tuple[int, int, int, int], *args: object, **kwargs: object
    ) -> list[Tactic]:
        tactics = enumerate_tactics(problem_sizes_mnkl, *args, **kwargs)
        # The runner computes out.T: its problem is (N, M, K, L).
        n, rows, k, _ = problem_sizes_mnkl
        pinned = PINNED_TACTICS.get((n, k))
        if _PINNED_MIN_ROWS <= rows <= LOW_BATCH_MAX_ROWS and pinned in tactics:
            return [pinned]
        return tactics

    module.autotune_tactics = autotune_tactics  # type: ignore[attr-defined]
    try:
        yield
    finally:
        module.autotune_tactics = original  # type: ignore[attr-defined]

run_zeros(linear, rows)

Call linear on an input laid out like the fused quantize producers'.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
def run_zeros(linear: nn.Module, rows: int) -> torch.Tensor:
    """Call ``linear`` on an input laid out like the fused quantize producers'."""
    k = linear.weight.shape[0]  # column-major [K, N]
    device = linear.weight.device
    shape = torch.Size((rows, k))
    qa = QuantizedActivation(
        torch.zeros(shape, dtype=torch.float8_e4m3fn, device=device),
        torch.zeros(mxfp8_scale_bytes(rows, k), dtype=torch.uint8, device=device),
        torch.bfloat16,
        shape,
        kMxfp8Dynamic,
    )
    return mxfp8_mm(linear, qa)

shape_representatives(linears)

One eligible linear per distinct (N, K), in a rank-independent order.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py 
109
110
111
112
113
114
115
116
def shape_representatives(linears: Iterable[nn.Module]) -> list[nn.Module]:
    """One eligible linear per distinct (N, K), in a rank-independent order."""
    by_shape: dict[tuple[int, int], nn.Module] = {}
    for linear in linears:
        kernel = _mxfp8_kernel(linear)
        if kernel is not None and mxfp8_eligible(linear):
            by_shape.setdefault(kernel.config.weight_shape, linear)
    return [by_shape[shape] for shape in sorted(by_shape)]