Skip to content

vllm.models.minimax_m3.amd.sparse_attention_mxfp4

MiniMax M3 block-sparse attention over an MXFP4 K/V cache on ROCm (gfx950).

Each head's K and V are two slots of one 8704 B page (mxfp4_pages). The top-k block ids the indexer wrote are turned into one page table row per query token, and pa_decode (MXFP4 K/V) attends each row as a length-1 sequence with the query as e4m3.

Classes:

Functions:

MiniMaxM3SparseRocmMxfp4Backend

Bases: MiniMaxM3SparseBackend

MiniMax M3 sparse attend over MXFP4 K/V pages (ROCm, gfx950).

Methods:

  • customize_spec –

    Store each head's K and V as two slots of 128 tokens x 68 B.

Source code in vllm/models/minimax_m3/amd/sparse_attention_mxfp4.py 
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
class MiniMaxM3SparseRocmMxfp4Backend(MiniMaxM3SparseBackend):
    """MiniMax M3 sparse attend over MXFP4 K/V pages (ROCm, gfx950)."""

    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["mxfp4"]

    @staticmethod
    def get_builder_cls() -> type[MiniMaxM3SparseRocmNvfp4MetadataBuilder]:
        # Adds each prefill token's request and absolute position, which the
        # sparse block-table builder reads.
        return MiniMaxM3SparseRocmNvfp4MetadataBuilder

    @classmethod
    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
        """Store each head's K and V as two slots of 128 tokens x 68 B."""
        if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_mxfp4:
            return spec
        return replace(
            spec,
            num_head_slots=2 * spec.num_kv_heads,
            state_content_bytes=TOKEN_BYTES,
        )

    @classmethod
    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
        return (KVCacheLayout.LBHNC,)

customize_spec(spec) classmethod

Store each head's K and V as two slots of 128 tokens x 68 B.

Source code in vllm/models/minimax_m3/amd/sparse_attention_mxfp4.py 
57
58
59
60
61
62
63
64
65
66
@classmethod
def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
    """Store each head's K and V as two slots of 128 tokens x 68 B."""
    if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_mxfp4:
        return spec
    return replace(
        spec,
        num_head_slots=2 * spec.num_kv_heads,
        state_content_bytes=TOKEN_BYTES,
    )

MiniMaxM3SparseRocmMxfp4Impl

Bases: MiniMaxM3SparseImpl

pa_decode over the pages each query token's top-k selects.

Source code in vllm/models/minimax_m3/amd/sparse_attention_mxfp4.py 
129
130
131
132
133
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
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
class MiniMaxM3SparseRocmMxfp4Impl(MiniMaxM3SparseImpl):
    """``pa_decode`` over the pages each query token's top-k selects."""

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        reason = rocm_mxfp4_attention_unsupported_reason(
            self.num_heads, self.num_kv_heads, self.head_size, self.block_size
        )
        if reason is not None:
            raise ValueError(f"MiniMax M3 MXFP4 KV cache on ROCm {reason}")

    def forward(
        self,
        layer: AttentionLayer,
        query: torch.Tensor,
        kv_cache: torch.Tensor,
        output: torch.Tensor,
    ) -> torch.Tensor:
        attn_metadata = get_forward_context().attn_metadata
        if not isinstance(attn_metadata, dict):
            return output  # profiling run; caches unbound
        md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
        assert isinstance(md, MiniMaxM3SparseMetadata)

        nd = md.num_decode_tokens
        num_tokens = md.num_actual_tokens
        topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
        assert topk is not None
        assert query.dtype == torch.float8_e4m3fn
        q = query[:num_tokens].view(-1, self.num_heads, self.head_size)
        out = output[:num_tokens].view(-1, self.num_heads, self.head_size)

        if md.num_decodes > 0:
            d = md.decode
            assert d is not None
            sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_decode(
                topk[:, :nd, :],
                d.block_table,
                d.seq_lens,
                d.decode_query_len,
                block_page_stride=1,
                pages_per_block=1,
            )
            mxfp4_sparse_attend(
                q[:nd], kv_cache, sparse_bt, sparse_ctx, out[:nd], self.scale
            )
        if md.num_prefills > 0:
            p = md.prefill
            assert p is not None
            assert p.query_req_id is not None and p.query_abs_pos is not None
            sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_prefill(
                topk[:, nd:num_tokens, :],
                p.block_table,
                p.query_req_id,
                p.query_abs_pos,
                block_page_stride=1,
                pages_per_block=1,
            )
            mxfp4_sparse_attend(
                q[nd:], kv_cache, sparse_bt, sparse_ctx, out[nd:], self.scale
            )
        return output

mxfp4_sparse_attend(q, kv_cache, sparse_bt, sparse_ctx, out, scale)

Attend each row to its selected pages as a length-1 sequence.

Source code in vllm/models/minimax_m3/amd/sparse_attention_mxfp4.py 
 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 mxfp4_sparse_attend(
    q: torch.Tensor,  # [rows, H, 128] e4m3
    kv_cache: torch.Tensor,  # [pages, 2, 128, 68] uint8, one KV head
    sparse_bt: torch.Tensor,  # [rows, topk] int32 page ids
    sparse_ctx: torch.Tensor,  # [rows] int32 tokens to attend
    out: torch.Tensor,  # [rows, H, 128] bf16
    scale: float,
) -> None:
    """Attend each row to its selected pages as a length-1 sequence."""
    from aiter.ops.flydsl.pa_decode import pa_decode

    pages = kv_cache.view(kv_cache.shape[0], 2, PAGE_BYTES)
    k_data, k_scale = mxfp4_page_views(pages[:, :1])
    v_data, v_scale = mxfp4_page_views(pages[:, 1:])
    pa_decode(
        out,
        q,
        k_data,
        v_data,
        sparse_ctx,
        sparse_bt,
        scale,
        1,
        key_scale=k_scale,
        value_scale=v_scale,
        work_plan=None,
        max_context_length=sparse_bt.shape[1] * PAGE_TOKENS,
    )

rocm_mxfp4_attention_unsupported_reason(num_heads, num_kv_heads, head_dim, block_size)

Why the ROCm MXFP4 attend cannot serve this layer, None when it can.

Source code in vllm/models/minimax_m3/amd/sparse_attention_mxfp4.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def rocm_mxfp4_attention_unsupported_reason(
    num_heads: int, num_kv_heads: int, head_dim: int, block_size: int
) -> str | None:
    """Why the ROCm MXFP4 attend cannot serve this layer, None when it can."""
    if not current_platform.is_rocm():
        return "needs ROCm"
    from vllm.platforms.rocm import on_gfx950

    if not on_gfx950():
        return "needs gfx950"
    if (
        head_dim != HEAD_DIM
        or block_size != PAGE_TOKENS
        or num_kv_heads != 1
        or num_heads % num_kv_heads
        or num_heads // num_kv_heads > MAX_GROUP_SIZE
    ):
        return (
            f"needs head_dim {HEAD_DIM}, block_size {PAGE_TOKENS}, one KV head "
            f"and at most {MAX_GROUP_SIZE} query heads per rank (got "
            f"num_heads={num_heads}, num_kv_heads={num_kv_heads}, "
            f"head_dim={head_dim}, block_size={block_size})"
        )
    return rocm_mxfp4_aiter_unavailable_reason()