vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn

DeepSeek V4.1 attention on FlashMLA's mega-attention kernel (SM100).

One kernel does Q RoPE, sparse attention, the inverse RoPE of the output and its FP8 cast, and writes straight into the buffer wo_a consumes. So the layer declares accepts_unnormed_unroped_query -- the kernel RoPEs Q itself, leaving Q only the zero-pad to the kernel's head count -- and its _alloc_attn_out / _o_proj pair speaks QuantizedActivation instead of bf16, since the output needs no rotation or quantization here.

wq_b rows and wo_a columns are permuted once at load so the surrounding GEMMs speak the kernel's chunk-interleaved layouts directly. A step's prefill and decode segments write disjoint token ranges of one output buffer pair, so a single wo_a einsum covers the whole step.

Small decode steps with 16 or 32 live heads run a swap-AB kernel with the same inputs and output instead (see _swapab_decode_supported).

Classes:

• DeepseekV4MegaAttnAttention –

  FlashMLA mega-attention layer for DeepSeek V4.1 (SM100).

Functions:

• alloc_mega_attn_output –

  Allocate the output buffer pair the mega-attention kernels write into.

• is_flashmla_mega_attn_supported –

  Whether FlashMLA's fused mega-attention kernels are usable here.

DeepseekV4MegaAttnAttention

Bases: DeepseekV4FlashMLAAttention

FlashMLA mega-attention layer for DeepSeek V4.1 (SM100).

Methods:

• is_available_for –

  Whether this layer can serve the configured model on this device.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention):
    """FlashMLA mega-attention layer for DeepSeek V4.1 (SM100)."""

    backend_cls = FlashMLAMegaAttnBackend

    @classmethod
    def is_available_for(cls, vllm_config: VllmConfig) -> bool:
        """Whether this layer can serve the configured model on this device.

        The kernel is SM100-only and wants ``WV_GROUP_SIZE`` heads per ``wo_a``
        group, which stops holding once TP divides the head count far enough
        (TP16 on a 64-head model leaves 4). ``__init__`` raises on both, so the
        default-backend selector asks here first rather than turning an
        unsupported topology into a startup crash.
        """
        if not is_flashmla_mega_attn_supported()[0]:
            return False
        config = vllm_config.model_config.hf_config
        config = getattr(config, "text_config", config)
        n_heads = getattr(config, "num_attention_heads", 0) or 0
        n_groups = getattr(config, "o_groups", 0) or 0
        tp_size = vllm_config.parallel_config.tensor_parallel_size
        if n_heads % tp_size or n_groups % tp_size:
            return False
        n_local_heads, n_local_groups = n_heads // tp_size, n_groups // tp_size
        if not n_local_groups or n_local_heads % WV_GROUP_SIZE:
            return False
        return n_local_heads // n_local_groups == WV_GROUP_SIZE

    def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
        super().__init__(vllm_config, *args, **kwargs)
        if self.n_local_heads % WV_GROUP_SIZE:
            raise ValueError(
                f"{self.prefix}: mega attention needs the local head count "
                f"({self.n_local_heads}) to be a multiple of {WV_GROUP_SIZE}."
            )
        if self.n_local_heads // self.n_local_groups != WV_GROUP_SIZE:
            raise ValueError(
                f"{self.prefix}: mega attention needs {WV_GROUP_SIZE} heads per "
                "wo_a group."
            )
        self.n_wv_group = self.padded_heads // WV_GROUP_SIZE
        q_perm = q_fused_permutation(self.n_local_heads, self.head_dim)
        permute_on_load(self.wq_b.weight, q_perm, dim=0)
        permute_on_load(self.wq_b.weight_scale, q_perm, dim=0)
        o_perm = o_fused_permutation(WV_GROUP_SIZE, self.head_dim)
        permute_on_load(self.wo_a.weight, o_perm, dim=1)
        o_chunk_perm = o_fused_chunk_permutation(WV_GROUP_SIZE, self.head_dim)
        permute_on_load(self.wo_a.weight_scale, o_chunk_perm, dim=1)
        self._use_swapab = _swapab_decode_supported(self)
        if self._use_swapab and vllm_config.kernel_config.enable_jit_warmup:
            self._register_swapab_warmup()

    # ---- interface contract ------------------------------------------------

    @property
    def accepts_unnormed_unroped_query(self) -> bool:
        return True

    @property
    def packed_kv_cache_dtype(self) -> CacheDType:
        # This kernel is the only one that reads an NVFP4 compressed cache, so
        # it is what an unspecific --kv-cache-dtype resolves to here.
        return "nvfp4_ds_mla"

    def _register_swapab_warmup(self) -> None:
        """Register the compile key this layer's swap-AB decode steps launch.

        Every field is fixed at construction. The compressed cache is read in
        kernel blocks of this backend's single size, which hold
        ``kernel_block // compress_ratio`` compressed slots each, whatever the
        KV cache manager's block size is. Layers sharing a key register it
        once.
        """
        from vllm.models.deepseek_v41.nvidia.ops.swapab_decode_attn import (
            _SWAPAB_DECODE_KERNEL,
        )

        (kernel_block,) = self.backend_cls.get_supported_kernel_block_sizes()
        assert isinstance(kernel_block, int)
        assert self.topk_indices_buffer is not None
        _SWAPAB_DECODE_KERNEL.register_warmup(
            heads=self.n_local_heads,
            swa_width=self.window_size,
            extra_width=self.topk_indices_buffer.shape[-1],
            extra_record=self.compressed_bytes_per_token,
            swa_page=self.swa_cache_layer.block_size,
            extra_page=kernel_block // self.compress_ratio,
        )

    def _alloc_attn_out(
        self, num_tokens: int, hidden_states: torch.Tensor
    ) -> QuantizedActivation:
        return alloc_mega_attn_output(num_tokens, self.n_wv_group, hidden_states.device)

    def _o_proj(
        self, attn_out: QuantizedActivation, positions: torch.Tensor
    ) -> torch.Tensor:
        """Grouped ``wo_a`` then ``wo_b`` over the kernel's quantized output.

        The inverse RoPE and the FP8 cast happened inside the attention
        kernel, and ``wo_a`` is permuted for its output layout, so there is
        nothing to rotate or quantize here and ``positions`` is unused. The
        scale already is DeepGEMM's packed-ue8m0 MN-major layout, so the
        einsum consumes the kernel's output with no repacking.
        """
        groups = self.n_local_groups
        num_tokens = attn_out.data.shape[0]
        quantize = self._can_quantize_o_proj
        z = torch.empty(
            (num_tokens, groups, self.o_lora_rank),
            dtype=torch.float8_e4m3fn if quantize else torch.bfloat16,
            device=attn_out.data.device,
        )
        if quantize:
            # wo_a writes wo_b's MXFP8 input with F8_128x4 UE8M0 scales. Padded
            # scale rows stay unwritten and only reach wo_b rows it never stores.
            z_scale = torch.empty(
                (cdiv(num_tokens, 128), groups * self.o_lora_rank // 128, 128),
                dtype=torch.int32,
                device=z.device,
            )
        fp8_einsum(
            "bhr,hdr->bhd",
            (attn_out.data[:, :groups], attn_out.scale[:, :groups]),
            (self.wo_a.weight, self.wo_a.weight_scale),
            (z, z_scale) if quantize else z,
            recipe=self._einsum_recipe,
            use_128x4_sfd_layout=quantize,
        )
        z = z.flatten(1)
        if not quantize:
            return self._wo_b_proj(z)
        scale = z_scale.view(-1).view(torch.uint8)
        return self._wo_b_proj(
            QuantizedActivation(z, scale, torch.bfloat16, z.shape, kMxfp8Dynamic)
        )

    @cached_property
    def _can_quantize_o_proj(self) -> bool:
        return (
            self._einsum_recipe == (1, 1, 32)
            and self.o_lora_rank % 128 == 0
            # The fused GEMM + reduce-scatter takes a BF16 z only.
            and self.gemm_rs is None
            and can_fuse_query_quant([self.wo_b])
        )

    # ---- forward -----------------------------------------------------------

    def forward_mqa(
        self,
        q: torch.Tensor,
        kv: torch.Tensor,
        positions: torch.Tensor,
        output: QuantizedActivation,
    ) -> None:
        attn_metadata = get_forward_context().attn_metadata
        if attn_metadata is None:
            # Warmup dummy run: reserve the prefill workspace, produce zeros.
            self._reserve_prefill_workspace(q)
            output.data.zero_()
            output.scale.zero_()
            return

        assert isinstance(attn_metadata, dict)
        # _fused_qnorm_rope_kv_insert zero-padded q to the kernel's head count.
        assert q.shape[1] == self.padded_heads
        flashmla_metadata = cast(
            DeepseekV4FlashMLAMetadata | None,
            attn_metadata.get(self.compressed_cache_prefix)
            if self.compressed_cache_prefix is not None
            else None,
        )
        swa_metadata = cast(
            "DeepseekSparseSWAMetadata", attn_metadata[self.swa_cache_layer.prefix]
        )
        num_decode_tokens = swa_metadata.num_decode_tokens

        if swa_metadata.num_prefills > 0:
            self._forward_prefill_mega(
                q[num_decode_tokens:],
                positions[num_decode_tokens:],
                flashmla_metadata,
                swa_metadata,
                output,
                num_decode_tokens,
            )
        if swa_metadata.num_decodes > 0:
            self._forward_decode_mega(
                q[:num_decode_tokens],
                positions[:num_decode_tokens],
                flashmla_metadata,
                swa_metadata,
                _token_slice(output, 0, num_decode_tokens),
            )

    def _reserve_prefill_workspace(self, q: torch.Tensor) -> None:
        swa_only = self.compress_ratio == 0
        n = 0 if swa_only else -(-self.max_model_len // self.compress_ratio)
        m = n + self.window_size + self.max_num_batched_tokens
        if swa_only:
            top_k = 0
        else:
            assert self.topk_indices_buffer is not None
            top_k = self.topk_indices_buffer.shape[-1]
        # One request covers the KV chunk (retained at the tail at runtime) and
        # the index scratch (taken from the head) side by side.
        current_workspace_manager().get_simultaneous(
            ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16),
            *self._prefill_index_workspace_specs(top_k),
        )

    def _prefill_index_workspace_specs(
        self, top_k: int
    ) -> tuple[tuple[tuple[int, ...], torch.dtype], ...]:
        combined_topk = round_up(top_k + self.window_size, 128)
        return (
            ((self.max_num_batched_tokens, combined_topk), torch.int32),
            ((self.max_num_batched_tokens,), torch.int32),
        )

    def _decode_compressed_kv_and_topk(
        self,
        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
        swa_metadata: "DeepseekSparseSWAMetadata",
    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
        """The compressed cache and its ``[n, topk]`` slot indices / lengths."""
        if self.compress_ratio == 0:
            return None, None, None
        assert flashmla_metadata is not None
        assert swa_metadata.is_valid_token is not None
        assert self.topk_indices_buffer is not None
        num_decode_tokens = swa_metadata.num_decode_tokens
        indices, lens = compute_global_topk_indices_and_lens(
            self.topk_indices_buffer[:num_decode_tokens],
            swa_metadata.token_to_req_indices,
            flashmla_metadata.block_table[: swa_metadata.num_decodes],
            flashmla_metadata.block_size // self.compress_ratio,
            swa_metadata.is_valid_token[:num_decode_tokens],
        )
        return self._compressed_kv_cache().unsqueeze(-2), indices, lens

    def _forward_decode_mega(
        self,
        q: torch.Tensor,
        positions: torch.Tensor,
        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
        swa_metadata: "DeepseekSparseSWAMetadata",
        out: QuantizedActivation,
    ) -> None:
        """Mega attention over the paged quantized caches, in place into ``out``.

        The caches are ``[num_blocks, page, 1, bytes]``, the record taken from
        ``bytes`` (528 V4.1 fp8, 288 V4.1 NVFP4; NVFP4 only as the compressed
        cache beside an fp8 SWA one), and the indices ``[s_q, topk]`` int32
        slot ids (``block * page + offset``, ``-1`` invalid). Where
        ``_use_swapab`` holds, steps of at most ``_SWAPAB_MAX_TOKENS`` tokens
        take the swap-AB kernel instead, which reads the same inputs and
        writes the same output layout.
        """
        extra_cache, extra_idx, extra_len = self._decode_compressed_kv_and_topk(
            flashmla_metadata, swa_metadata
        )
        assert swa_metadata.decode_swa_indices is not None
        swa_idx = swa_metadata.decode_swa_indices.view(q.shape[0], -1)
        if (
            self._use_swapab
            and q.shape[0] <= _SWAPAB_MAX_TOKENS
            and extra_cache is not None
            # The cluster and the warmed-up compile key are sized for the
            # window; DSpark's non-causal draft steps widen the SWA list past it.
            and swa_idx.shape[-1] == self.window_size
        ):
            from vllm.models.deepseek_v41.nvidia.ops.swapab_decode_attn import (
                _SWAPAB_DECODE_KERNEL,
            )

            _SWAPAB_DECODE_KERNEL(
                q=q,
                swa_cache=self.swa_cache_layer.kv_cache,
                swa_idx=swa_idx,
                swa_len=swa_metadata.decode_swa_lens,
                extra_cache=extra_cache.squeeze(-2),
                extra_idx=extra_idx,
                extra_len=extra_len,
                positions=positions,
                cos_sin_cache=self.rotary_emb.cos_sin_cache,
                sink=self.attn_sink,
                heads=self.n_local_heads,
                softmax_scale=self.scale,
                out_data=out.data,
                out_scale=out.scale,
            )
            return
        torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode(
            q,
            self.swa_cache_layer.kv_cache.unsqueeze(-2),
            swa_idx,
            self.scale,
            _HEAD_DIM_V,
            self.attn_sink,
            swa_metadata.decode_swa_lens,
            extra_cache,
            extra_idx,
            extra_len,
            False,  # enable_q_norm
            0.0,  # rms_norm_eps
            positions,
            *_ROPE_ARGS,
            self.rotary_emb.cos_sin_cache,
            self.n_wv_group,
            *_SF_ARGS,
            out.data,
            out.scale,
        )

    def _forward_prefill_mega(
        self,
        q: torch.Tensor,
        positions: torch.Tensor,
        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
        swa_metadata: "DeepseekSparseSWAMetadata",
        out: QuantizedActivation,
        token_base: int,
    ) -> None:
        swa_only = self.compress_ratio == 0
        num_decodes = swa_metadata.num_decodes
        num_decode_tokens = swa_metadata.num_decode_tokens
        seq_lens = swa_metadata.prefill_seq_lens
        gather_lens = swa_metadata.prefill_gather_lens
        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
        query_start_loc = swa_metadata.query_start_loc
        assert seq_lens is not None and gather_lens is not None
        assert query_start_loc_cpu is not None and query_start_loc is not None
        prefill_token_base = query_start_loc_cpu[num_decodes]
        assert self.topk_indices_buffer is not None
        topk_indices = self.topk_indices_buffer[num_decode_tokens:][
            : swa_metadata.num_prefill_tokens
        ]
        top_k = 0 if swa_only else topk_indices.shape[-1]
        chunk_plan = swa_metadata.get_prefill_chunk_plan(
            compress_ratio=self.compress_ratio,
            prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
            has_compressed=not swa_only,
        )
        assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
        workspace_manager = current_workspace_manager()
        index_specs = self._prefill_index_workspace_specs(top_k)
        for chunk in chunk_plan:
            chunk_start, chunk_end, chunk_n, chunk_m = chunk
            chunk_size = chunk_end - chunk_start
            kv_spec = ((chunk_size, chunk_m, q.shape[-1]), torch.bfloat16)
            if swa_only:
                kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous(
                    kv_spec, *index_specs
                )
            else:
                assert flashmla_metadata is not None
                idx_ws, lens_ws = workspace_manager.get_simultaneous(*index_specs)
                # The chunk's compressed columns stay in the workspace's tail
                # span, which the scratch of the layers in between leaves
                # alone. A kv source has just inserted this step's rows, so it
                # always gathers; the consumers after it in the same forward
                # reuse its columns while the span still holds this source and
                # chunk, and rewrite only their SWA columns.
                kv_ws, intact = workspace_manager.get_retained(
                    (self.compressed_cache_prefix, chunk), *kv_spec
                )
                if self.is_kv_source or not intact:
                    dequantize_and_gather_k_cache(
                        kv_ws[:, :chunk_n],
                        self._compressed_kv_cache(),
                        seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
                        gather_lens=None,
                        block_table=flashmla_metadata.block_table[num_decodes:][
                            chunk_start:chunk_end
                        ],
                        block_size=flashmla_metadata.block_size // self.compress_ratio,
                        offset=0,
                    )
            dequantize_and_gather_k_cache(
                kv_ws[:chunk_size],
                self.swa_cache_layer.kv_cache,
                seq_lens=seq_lens[chunk_start:chunk_end],
                gather_lens=gather_lens[chunk_start:chunk_end],
                block_table=swa_metadata.block_table[num_decodes:][
                    chunk_start:chunk_end
                ],
                block_size=swa_metadata.block_size,
                offset=chunk_n,
            )
            first, last = num_decodes + chunk_start, num_decodes + chunk_end
            qs = int(query_start_loc_cpu[first] - prefill_token_base)
            qe = int(query_start_loc_cpu[last] - prefill_token_base)
            combined_indices, combined_lens = combine_topk_swa_indices(
                topk_indices[qs:qe],
                query_start_loc[first : last + 1],
                seq_lens[chunk_start:chunk_end],
                gather_lens[chunk_start:chunk_end],
                self.window_size,
                self.compress_ratio,
                top_k,
                chunk_m,
                chunk_n,
                out=(idx_ws[: qe - qs], lens_ws[: qe - qs]),
            )
            chunk_out = _token_slice(out, token_base + qs, token_base + qe)
            # Mega attention over the gathered non-paged bf16 KV (RoPE already
            # applied), in place into this chunk's token range. `indices`
            # entries outside [0, s_kv) skip.
            torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_fwd(
                q[qs:qe],
                kv_ws.view(-1, 1, q.shape[-1]),
                combined_indices.unsqueeze(1),
                self.scale,
                _HEAD_DIM_V,
                self.attn_sink,
                combined_lens,
                False,  # enable_q_norm
                0.0,  # rms_norm_eps
                positions[qs:qe],
                *_ROPE_ARGS,
                self.rotary_emb.cos_sin_cache,
                self.n_wv_group,
                *_SF_ARGS,
                chunk_out.data,
                chunk_out.scale,
            )

_decode_compressed_kv_and_topk(flashmla_metadata, swa_metadata)

The compressed cache and its [n, topk] slot indices / lengths.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

def _decode_compressed_kv_and_topk(
    self,
    flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
    swa_metadata: "DeepseekSparseSWAMetadata",
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
    """The compressed cache and its ``[n, topk]`` slot indices / lengths."""
    if self.compress_ratio == 0:
        return None, None, None
    assert flashmla_metadata is not None
    assert swa_metadata.is_valid_token is not None
    assert self.topk_indices_buffer is not None
    num_decode_tokens = swa_metadata.num_decode_tokens
    indices, lens = compute_global_topk_indices_and_lens(
        self.topk_indices_buffer[:num_decode_tokens],
        swa_metadata.token_to_req_indices,
        flashmla_metadata.block_table[: swa_metadata.num_decodes],
        flashmla_metadata.block_size // self.compress_ratio,
        swa_metadata.is_valid_token[:num_decode_tokens],
    )
    return self._compressed_kv_cache().unsqueeze(-2), indices, lens

_forward_decode_mega(q, positions, flashmla_metadata, swa_metadata, out)

Mega attention over the paged quantized caches, in place into out.

The caches are [num_blocks, page, 1, bytes], the record taken from bytes (528 V4.1 fp8, 288 V4.1 NVFP4; NVFP4 only as the compressed cache beside an fp8 SWA one), and the indices [s_q, topk] int32 slot ids (block * page + offset, -1 invalid). Where _use_swapab holds, steps of at most _SWAPAB_MAX_TOKENS tokens take the swap-AB kernel instead, which reads the same inputs and writes the same output layout.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

def _forward_decode_mega(
    self,
    q: torch.Tensor,
    positions: torch.Tensor,
    flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
    swa_metadata: "DeepseekSparseSWAMetadata",
    out: QuantizedActivation,
) -> None:
    """Mega attention over the paged quantized caches, in place into ``out``.

    The caches are ``[num_blocks, page, 1, bytes]``, the record taken from
    ``bytes`` (528 V4.1 fp8, 288 V4.1 NVFP4; NVFP4 only as the compressed
    cache beside an fp8 SWA one), and the indices ``[s_q, topk]`` int32
    slot ids (``block * page + offset``, ``-1`` invalid). Where
    ``_use_swapab`` holds, steps of at most ``_SWAPAB_MAX_TOKENS`` tokens
    take the swap-AB kernel instead, which reads the same inputs and
    writes the same output layout.
    """
    extra_cache, extra_idx, extra_len = self._decode_compressed_kv_and_topk(
        flashmla_metadata, swa_metadata
    )
    assert swa_metadata.decode_swa_indices is not None
    swa_idx = swa_metadata.decode_swa_indices.view(q.shape[0], -1)
    if (
        self._use_swapab
        and q.shape[0] <= _SWAPAB_MAX_TOKENS
        and extra_cache is not None
        # The cluster and the warmed-up compile key are sized for the
        # window; DSpark's non-causal draft steps widen the SWA list past it.
        and swa_idx.shape[-1] == self.window_size
    ):
        from vllm.models.deepseek_v41.nvidia.ops.swapab_decode_attn import (
            _SWAPAB_DECODE_KERNEL,
        )

        _SWAPAB_DECODE_KERNEL(
            q=q,
            swa_cache=self.swa_cache_layer.kv_cache,
            swa_idx=swa_idx,
            swa_len=swa_metadata.decode_swa_lens,
            extra_cache=extra_cache.squeeze(-2),
            extra_idx=extra_idx,
            extra_len=extra_len,
            positions=positions,
            cos_sin_cache=self.rotary_emb.cos_sin_cache,
            sink=self.attn_sink,
            heads=self.n_local_heads,
            softmax_scale=self.scale,
            out_data=out.data,
            out_scale=out.scale,
        )
        return
    torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode(
        q,
        self.swa_cache_layer.kv_cache.unsqueeze(-2),
        swa_idx,
        self.scale,
        _HEAD_DIM_V,
        self.attn_sink,
        swa_metadata.decode_swa_lens,
        extra_cache,
        extra_idx,
        extra_len,
        False,  # enable_q_norm
        0.0,  # rms_norm_eps
        positions,
        *_ROPE_ARGS,
        self.rotary_emb.cos_sin_cache,
        self.n_wv_group,
        *_SF_ARGS,
        out.data,
        out.scale,
    )

_o_proj(attn_out, positions)

Grouped wo_a then wo_b over the kernel's quantized output.

The inverse RoPE and the FP8 cast happened inside the attention kernel, and wo_a is permuted for its output layout, so there is nothing to rotate or quantize here and positions is unused. The scale already is DeepGEMM's packed-ue8m0 MN-major layout, so the einsum consumes the kernel's output with no repacking.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

def _o_proj(
    self, attn_out: QuantizedActivation, positions: torch.Tensor
) -> torch.Tensor:
    """Grouped ``wo_a`` then ``wo_b`` over the kernel's quantized output.

    The inverse RoPE and the FP8 cast happened inside the attention
    kernel, and ``wo_a`` is permuted for its output layout, so there is
    nothing to rotate or quantize here and ``positions`` is unused. The
    scale already is DeepGEMM's packed-ue8m0 MN-major layout, so the
    einsum consumes the kernel's output with no repacking.
    """
    groups = self.n_local_groups
    num_tokens = attn_out.data.shape[0]
    quantize = self._can_quantize_o_proj
    z = torch.empty(
        (num_tokens, groups, self.o_lora_rank),
        dtype=torch.float8_e4m3fn if quantize else torch.bfloat16,
        device=attn_out.data.device,
    )
    if quantize:
        # wo_a writes wo_b's MXFP8 input with F8_128x4 UE8M0 scales. Padded
        # scale rows stay unwritten and only reach wo_b rows it never stores.
        z_scale = torch.empty(
            (cdiv(num_tokens, 128), groups * self.o_lora_rank // 128, 128),
            dtype=torch.int32,
            device=z.device,
        )
    fp8_einsum(
        "bhr,hdr->bhd",
        (attn_out.data[:, :groups], attn_out.scale[:, :groups]),
        (self.wo_a.weight, self.wo_a.weight_scale),
        (z, z_scale) if quantize else z,
        recipe=self._einsum_recipe,
        use_128x4_sfd_layout=quantize,
    )
    z = z.flatten(1)
    if not quantize:
        return self._wo_b_proj(z)
    scale = z_scale.view(-1).view(torch.uint8)
    return self._wo_b_proj(
        QuantizedActivation(z, scale, torch.bfloat16, z.shape, kMxfp8Dynamic)
    )

_register_swapab_warmup()

Register the compile key this layer's swap-AB decode steps launch.

Every field is fixed at construction. The compressed cache is read in kernel blocks of this backend's single size, which hold kernel_block // compress_ratio compressed slots each, whatever the KV cache manager's block size is. Layers sharing a key register it once.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

def _register_swapab_warmup(self) -> None:
    """Register the compile key this layer's swap-AB decode steps launch.

    Every field is fixed at construction. The compressed cache is read in
    kernel blocks of this backend's single size, which hold
    ``kernel_block // compress_ratio`` compressed slots each, whatever the
    KV cache manager's block size is. Layers sharing a key register it
    once.
    """
    from vllm.models.deepseek_v41.nvidia.ops.swapab_decode_attn import (
        _SWAPAB_DECODE_KERNEL,
    )

    (kernel_block,) = self.backend_cls.get_supported_kernel_block_sizes()
    assert isinstance(kernel_block, int)
    assert self.topk_indices_buffer is not None
    _SWAPAB_DECODE_KERNEL.register_warmup(
        heads=self.n_local_heads,
        swa_width=self.window_size,
        extra_width=self.topk_indices_buffer.shape[-1],
        extra_record=self.compressed_bytes_per_token,
        swa_page=self.swa_cache_layer.block_size,
        extra_page=kernel_block // self.compress_ratio,
    )

is_available_for(vllm_config) classmethod

Whether this layer can serve the configured model on this device.

The kernel is SM100-only and wants WV_GROUP_SIZE heads per wo_a group, which stops holding once TP divides the head count far enough (TP16 on a 64-head model leaves 4). __init__ raises on both, so the default-backend selector asks here first rather than turning an unsupported topology into a startup crash.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

@classmethod
def is_available_for(cls, vllm_config: VllmConfig) -> bool:
    """Whether this layer can serve the configured model on this device.

    The kernel is SM100-only and wants ``WV_GROUP_SIZE`` heads per ``wo_a``
    group, which stops holding once TP divides the head count far enough
    (TP16 on a 64-head model leaves 4). ``__init__`` raises on both, so the
    default-backend selector asks here first rather than turning an
    unsupported topology into a startup crash.
    """
    if not is_flashmla_mega_attn_supported()[0]:
        return False
    config = vllm_config.model_config.hf_config
    config = getattr(config, "text_config", config)
    n_heads = getattr(config, "num_attention_heads", 0) or 0
    n_groups = getattr(config, "o_groups", 0) or 0
    tp_size = vllm_config.parallel_config.tensor_parallel_size
    if n_heads % tp_size or n_groups % tp_size:
        return False
    n_local_heads, n_local_groups = n_heads // tp_size, n_groups // tp_size
    if not n_local_groups or n_local_heads % WV_GROUP_SIZE:
        return False
    return n_local_heads // n_local_groups == WV_GROUP_SIZE

_swapab_decode_supported(layer)

Whether the layer's small decode steps take the swap-AB kernel.

The mega kernel pads the query heads to 64, so at 32 or 16 live heads half or three quarters of its MMA work is padding. Batch-invariant runs keep the mega kernel at every step size, as the two kernels round the FP8 output differently.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

def _swapab_decode_supported(layer: "DeepseekV4MegaAttnAttention") -> bool:
    """Whether the layer's small decode steps take the swap-AB kernel.

    The mega kernel pads the query heads to 64, so at 32 or 16 live heads half
    or three quarters of its MMA work is padding. Batch-invariant runs keep the
    mega kernel at every step size, as the two kernels round the FP8 output
    differently.
    """
    if (
        layer.compress_ratio == 0
        or layer.n_local_heads not in (16, 32)
        or envs.VLLM_BATCH_INVARIANT
        or not current_platform.is_device_capability_family(100)
        # The kernel's module imports the CuTe DSL when it loads.
        or not has_cutedsl()
    ):
        return False
    from vllm.models.deepseek_v41.nvidia.ops.swapab_decode_attn import (
        swapab_decode_supported,
    )

    assert layer.topk_indices_buffer is not None
    return swapab_decode_supported(
        layer.n_local_heads, layer.window_size, layer.topk_indices_buffer.shape[-1]
    )

_token_slice(out, start, end)

The [start, end) token slice of a mega-attention output buffer.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

def _token_slice(out: QuantizedActivation, start: int, end: int) -> QuantizedActivation:
    """The ``[start, end)`` token slice of a mega-attention output buffer."""
    data = out.data[start:end]
    return replace(out, data=data, scale=out.scale[start:end], orig_shape=data.shape)

alloc_mega_attn_output(num_tokens, n_wv_group, device)

Allocate the output buffer pair the mega-attention kernels write into.

One pair per forward step: the prefill and decode segments fill disjoint token ranges, so a single fp8_einsum over all num_tokens can consume the result instead of one call per segment.

The scale buffer is MN-major -- stride 1 along tokens, head-dim stride ceil4(num_tokens) -- which is both what the kernel requires of a caller-provided buffer and what DeepGEMM expects for this N.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

def alloc_mega_attn_output(
    num_tokens: int,
    n_wv_group: int,
    device: torch.device,
) -> QuantizedActivation:
    """Allocate the output buffer pair the mega-attention kernels write into.

    One pair per forward step: the prefill and decode segments fill disjoint
    token ranges, so a single ``fp8_einsum`` over all ``num_tokens`` can
    consume the result instead of one call per segment.

    The scale buffer is MN-major -- stride 1 along tokens, head-dim stride
    ``ceil4(num_tokens)`` -- which is both what the kernel requires of a
    caller-provided buffer and what DeepGEMM expects for this N.
    """
    d = WV_GROUP_SIZE * _HEAD_DIM_V
    data = torch.empty(
        (num_tokens, n_wv_group, d), dtype=torch.float8_e4m3fn, device=device
    )
    aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize)
    scale = torch.empty(
        (n_wv_group, d // (_QUANT_GROUP * 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=kMxfp8Dynamic,
    )

is_flashmla_mega_attn_supported()

Whether FlashMLA's fused mega-attention kernels are usable here.

Source code in vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

def is_flashmla_mega_attn_supported() -> tuple[bool, str | None]:
    """Whether FlashMLA's fused mega-attention kernels are usable here."""
    is_available, maybe_reason = is_flashmla_sparse_supported()
    if not is_available:
        return False, maybe_reason
    if not current_platform.is_device_capability_family(100):
        return False, "FlashMLA mega attention requires sm_10x GPUs."
    if not hasattr(torch.ops._flashmla_C, "fused_norm_rope_attn_rope_cast_decode"):
        return False, "vllm._flashmla_C was built without the mega attention ops."
    if not hasattr(torch.ops._flashmla_C, "permute_q_b_proj"):
        return False, "vllm._flashmla_C predates the mega attention weight permutes."
    return True, None