fix(minimax-m3): unblock fp8 KV cache and EAGLE3 spec decode on the vLLM plugin

.github/workflows/atom-vllm-test.yaml

@@ -144,11 +144,12 @@ jobs:
             lm_eval_num_fewshot: 20
             accuracy_test_threshold: 0.91
             runner: spur-runner-mi355x-8gpu
-          - display_name: "MiniMax-M3-MXFP4 TP4"
-            model_name: "MiniMax-M3-MXFP4"
+          - display_name: "MiniMax-M3-MXFP4 EAGLE3 TP4"
+            model_name: "MiniMax-M3-MXFP4-EAGLE3"
             model_path: "amd/MiniMax-M3-MXFP4"
+            draft_model_path: "Inferact/MiniMax-M3-EAGLE3-GQA"
             client_command: "lm_eval --model local-chat-completions --apply_chat_template --fewshot_as_multiturn --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=32,max_gen_toks=2048 --tasks gsm8k --num_fewshot 5 --batch_size 65 --output_path ${OUTPUT_PATH}"
-            extra_args: "--tensor-parallel-size 4 --kv-cache-dtype auto --no-trust-remote-code --gpu-memory-utilization 0.85 --block-size 128 --max-model-len 32768 --max-num-seqs 128 --max-num-batched-tokens 32768 --no-enable_prefix_caching --no-async-scheduling --language-model-only --hf-overrides '{\"use_index_cache\": true, \"index_topk_freq\": 4}' --compilation-config '{\"cudagraph_mode\": \"FULL_AND_PIECEWISE\"}'"
+            extra_args: "--tensor-parallel-size 4 --kv-cache-dtype fp8 --no-trust-remote-code --gpu-memory-utilization 0.85 --block-size 128 --max-model-len 32768 --max-num-seqs 128 --max-num-batched-tokens 32768 --no-enable_prefix_caching --no-async-scheduling --language-model-only --hf-overrides '{\"use_index_cache\": true, \"index_topk_freq\": 4}' --compilation-config '{\"cudagraph_mode\": \"FULL_AND_PIECEWISE\"}' --speculative-config '{\"method\":\"eagle3\",\"model\":\"/models/Inferact/MiniMax-M3-EAGLE3-GQA\",\"num_speculative_tokens\":3}'"
             env_vars: "AITER_QUICK_REDUCE_QUANTIZATION=INT4"
             lm_eval_num_fewshot: 5
             accuracy_test_threshold: 0.93

atom/model_ops/minimax_m3/index_topk.py

@@ -877,9 +877,14 @@ def minimax_m3_index_topk_decode(
     assert (
         num_idx_heads == num_kv_heads
     ), "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
-    assert (
-        total_q % max_query_len == 0
-    ), f"total_q {total_q} not divisible by max_query_len {max_query_len}"
+    assert total_q % max_query_len == 0, (
+        f"total_q {total_q} not divisible by max_query_len {max_query_len}: the "
+        "decode segment must be uniform, i.e. every request contributes exactly "
+        "max_query_len rows, because the kernels below recover a request as "
+        "row // max_query_len. A ragged segment (spec decode where some request "
+        "carries no draft tokens) has to be routed to the prefill path by the "
+        "attention metadata builder."
+    )
     batch = seq_lens.shape[0]
     max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
     _require_packable(max_block)

atom/plugin/vllm/attention/metadata.py

@@ -449,6 +449,37 @@ class MinimaxM3SparseMetadata:
     decode: MinimaxM3SparseDecodeMetadata | None = None
 
 
+def _uniform_decode_query_len(
+    query_start_loc_cpu: torch.Tensor | None, num_decodes: int
+) -> int | None:
+    """Query length shared by every decode request, or None if they differ.
+
+    ATOM's decode kernels recover a request from a flat query row by dividing:
+    the M3 index-topk kernels use ``row // max_query_len`` (and the causal
+    cutoff ``seq_len - max_query_len + tok + 1``), and aiter's gluon paged
+    decode reshapes ``q`` to ``[q.shape[0] // max_query_len, max_query_len,
+    ...]``. Both only hold when every decode request contributes exactly
+    ``max_query_len`` rows.
+
+    Speculative decode breaks that on both sides. On the target model a request
+    that joins without draft tokens contributes one row next to requests
+    verifying ``num_spec + 1``; on the EAGLE/MTP draft a request contributes
+    however many tokens the last step accepted. vLLM keeps all of them in the
+    decode segment because each query length is still within the reorder
+    threshold, so the builders have to check the lengths themselves.
+    """
+    if num_decodes <= 0 or query_start_loc_cpu is None:
+        return None
+    starts = query_start_loc_cpu[: num_decodes + 1]
+    if starts.numel() != num_decodes + 1:
+        return None
+    query_lens = starts[1:] - starts[:-1]
+    first = int(query_lens[0])
+    if not bool(torch.all(query_lens == first)):
+        return None
+    return first
+
+
 class MinimaxM3SparseAttentionMetadataBuilder(AttentionMetadataBuilder):
     # Uniform decode batches are safe to capture, including spec-decode verify
     # (query_len == num_spec + 1): the decode index-topk and sparse-attn kernels
@@ -519,9 +550,34 @@ def build(
 
         # Plain decode has max_query_len == 1, while MTP/spec decode verifies
         # num_spec+1 tokens per request. Both should use the decode path, but only
-        # when the split says there are no prefill/extend requests in the batch.
-        if num_decodes > 0 and num_extends == 0 and num_prefills == 0:
-            return self._build_uniform_decode_metadata(common_attn_metadata)
+        # when the split says there are no prefill/extend requests in the batch
+        # AND every decode request carries the same number of query tokens -- the
+        # decode kernels index a request as `row // max_query_len`, so a ragged
+        # decode segment would read the wrong request and the wrong causal
+        # cutoff. Spec decode produces such a segment whenever a request without
+        # draft tokens sits next to requests verifying num_spec+1 tokens.
+        decode_query_len = _uniform_decode_query_len(
+            common_attn_metadata.query_start_loc_cpu, num_decodes
+        )
+        if (
+            num_decodes > 0
+            and num_extends == 0
+            and num_prefills == 0
+            and decode_query_len is not None
+        ):
+            return self._build_uniform_decode_metadata(
+                common_attn_metadata, decode_query_len
+            )
+
+        if num_decodes > 0 and decode_query_len is None:
+            # Ragged decode segment: hand those requests to the prefill kernel,
+            # which derives causality from cu_seqlens_q/context_lens and so
+            # accepts variable query lengths. The prefill slice below starts at
+            # num_decodes/num_decode_tokens, so zeroing both widens it to the
+            # whole batch.
+            num_prefills += num_decodes
+            num_decodes = 0
+            num_decode_tokens = 0
 
         num_tokens = common_attn_metadata.num_actual_tokens
         num_prefills_total = num_extends + num_prefills
@@ -564,10 +620,14 @@ def build(
 
         decode_metadata: MinimaxM3SparseDecodeMetadata | None = None
         if num_decodes > 0:
+            # decode_query_len is the measured length, not the reorder
+            # threshold: a mixed batch whose decode segment is plain decode has
+            # query_len == 1 even when the threshold is num_spec + 1, and
+            # feeding the threshold to the kernels would mis-map every row.
             decode_metadata = MinimaxM3SparseDecodeMetadata(
                 seq_lens=seq_lens[:num_decodes],
                 block_table=block_table[:num_decodes],
-                max_query_len=self.reorder_batch_threshold,
+                max_query_len=decode_query_len,
             )
 
         return MinimaxM3SparseMetadata(
@@ -585,12 +645,21 @@ def build(
             decode=decode_metadata,
         )
 
-    def _build_uniform_decode_metadata(self, common_attn_metadata):
+    def _build_uniform_decode_metadata(
+        self, common_attn_metadata, decode_query_len: int | None = None
+    ):
         assert common_attn_metadata is not None
 
         num_reqs = common_attn_metadata.num_reqs
         num_tokens = common_attn_metadata.num_actual_tokens
-        max_query_len = common_attn_metadata.max_query_len
+        # Callers that already measured the per-request query length pass it in;
+        # cudagraph capture builds a uniform batch by construction, so falling
+        # back to the batch maximum is exact there.
+        max_query_len = (
+            decode_query_len
+            if decode_query_len is not None
+            else common_attn_metadata.max_query_len
+        )
         seq_lens = common_attn_metadata.seq_lens
         block_table = common_attn_metadata.block_table_tensor
 
@@ -737,6 +806,23 @@ def build(
             num_prefill_tokens,
         ) = split_ret
 
+        # aiter's gluon decode kernel derives the batch as
+        # `q.shape[0] // max_query_len`, so a decode segment whose requests
+        # disagree on query length cannot go through it. EAGLE/MTP propose steps
+        # produce exactly that: a request's query length there is however many
+        # draft tokens the previous step accepted. Send such a segment to the
+        # extend path instead, which is varlen (driven by cu_seqlens_q). Decode
+        # requests sort before extends, so widening the extend segment covers
+        # them without reordering.
+        decode_query_len = _uniform_decode_query_len(
+            common_attn_metadata.query_start_loc_cpu, num_decodes
+        )
+        if num_decodes > 0 and decode_query_len is None:
+            num_extends += num_decodes
+            num_extend_tokens += num_decode_tokens
+            num_decodes = 0
+            num_decode_tokens = 0
+
         prefill_only = num_decodes == 0 and num_extends == 0 and num_prefills > 0
         decode_only = num_decodes > 0 and num_extends == 0 and num_prefills == 0
         mixed = not (prefill_only or decode_only)
@@ -767,7 +853,7 @@ def build(
                     - prefill_query_start_loc[prefill_start]
                 )
             if num_decodes > 0:
-                decode_max_query_len = query_lens_cpu[:num_decodes].max().item()
+                decode_max_query_len = decode_query_len
                 decode_max_seq_len = seq_lens[:num_decodes].max().item()
                 decode_query_start_loc = decode_query_start_loc[: num_decodes + 1]
 
@@ -965,30 +1051,38 @@ def build_for_drafting(
         """
         Build attention metadata for draft model without CPU-GPU sync.
 
-        During EAGLE/MTP drafting all requests are uniform decodes, so we can
-        skip split_decodes_prefills_and_extends() and avoid all .cpu() /
-        .item() calls that would otherwise break CUDA graph capture.
+        Drafting is usually a uniform decode, and then we can skip
+        split_decodes_prefills_and_extends() and avoid all .cpu() / .item()
+        calls that would otherwise break CUDA graph capture.
+
+        It is not always uniform, though: a request's query length here is
+        however many draft tokens the previous step accepted, so one request
+        can bring fewer rows than its neighbours. The gluon decode kernel
+        reshapes q to [q.shape[0] // max_query_len, max_query_len, ...] and
+        would fail on such a batch, so fall back to the full build(), which
+        routes a ragged segment to the varlen extend path.
         """
         query_start_loc = common_attn_metadata.query_start_loc_cpu
         query_lens = query_start_loc[1:] - query_start_loc[:-1]
         is_prefill = query_lens > self.reorder_batch_threshold
 
-        if torch.any(is_prefill):
+        num_reqs = common_attn_metadata.num_reqs
+        decode_query_len = _uniform_decode_query_len(query_start_loc, num_reqs)
+        if torch.any(is_prefill) or decode_query_len is None:
             return self.build(
                 common_prefix_len=0, common_attn_metadata=common_attn_metadata
             )
 
-        num_reqs = common_attn_metadata.num_reqs
         num_tokens = common_attn_metadata.num_actual_tokens
         decode_metadata = AiterMhaPhaseMetadata(
-            max_query_len=common_attn_metadata.max_query_len,
+            max_query_len=decode_query_len,
             max_seq_len=common_attn_metadata.max_seq_len,
             query_start_loc=common_attn_metadata.query_start_loc,
         )
         return AiterMhaMetadataForVllm(
             num_actual_tokens=num_tokens,
             num_actual_kv_tokens=0,
-            max_query_len=common_attn_metadata.max_query_len,
+            max_query_len=decode_query_len,
             query_start_loc=common_attn_metadata.query_start_loc,
             max_seq_len=common_attn_metadata.max_seq_len,
             seq_lens=common_attn_metadata.seq_lens,

atom/plugin/vllm/attention/minimax_m3_attnetion.py

@@ -65,6 +65,40 @@ def minimax_m3_sparse_attention(
     return layer._forward_with_output(qkv, positions, output)
 
 
+def _index_cache_torch_dtype(kv_cache_dtype: str, model_config) -> torch.dtype:
+    """Torch dtype for the MiniMax-M3 index cache.
+
+    Everything except fp8 is vLLM's own mapping. fp8 is the exception, and it
+    belongs here rather than in ``kv_cache_dtype_str_to_dtype``: that table is
+    vLLM's, it is shared by every layer of every model, and its ``fp8 ->
+    torch.uint8`` entry is deliberate -- vLLM's kernels take an fp8 KV cache as
+    a byte buffer and reinterpret it, and so do the aiter paged kernels behind
+    ATOM's main sparse and dense caches (see
+    ``_page16_shuffle_cache_for_sparse_kernel``, which does the ``.view()``
+    itself). Relabelling fp8 globally would change those caches too.
+
+    The index cache is read by a kernel that dispatches on the tensor instead:
+    ``_index_block_score_kernel`` branches on ``k.dtype.is_fp8()``, so a
+    uint8-labelled cache goes down the bf16 branch and dots bf16 against uint8
+    (a Triton compile error on the first request). The bytes there are already
+    fp8 -- aiter's ``fused_qknorm_idxrqknorm`` writes them under
+    ``kv_cache_dtype="fp8"`` -- so only the label was wrong. ``dtypes.d_dtypes``
+    is the same handle the native server resolves this cache through
+    (``atom.model_ops.attentions.aiter_attention._resolve_index_cache_dtype``),
+    which keeps the two paths on one arch-correct fp8 dtype instead of a
+    hard-coded ``torch.float8_e4m3fn``.
+
+    The element stays one byte either way, so vLLM's page-size accounting is
+    unchanged, and vLLM already allows a real fp8 dtype in this position -- its
+    own ``fp8_inc`` entry maps to ``torch.float8_e4m3fn``.
+    """
+    from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
+
+    if str(kv_cache_dtype).startswith("fp8"):
+        return dtypes.d_dtypes["fp8"]
+    return kv_cache_dtype_str_to_dtype(kv_cache_dtype, model_config)
+
+
 class MiniMaxM3SparseIndexerCache(nn.Module, AttentionLayerBase):
     """Key-only index cache owned by MiniMax-M3 sparse attention."""
 
@@ -76,7 +110,6 @@ def __init__(
         kv_cache_dtype: str,
     ) -> None:
         from vllm.v1.attention.backend import AttentionType
-        from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
 
         super().__init__()
         atom_config = get_current_atom_config()
@@ -86,7 +119,7 @@ def __init__(
         self.attn_type = AttentionType.DECODER
         self.attn_backend = SparseMHAIndexerBackend
         self.kv_cache_dtype = kv_cache_dtype
-        self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
+        self.kv_cache_torch_dtype = _index_cache_torch_dtype(
             kv_cache_dtype, vllm_config.model_config
         )
         self.num_kv_heads = 1

tests/plugin/test_minimax_m3_decode_uniformity.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""The M3 decode path may only be handed a batch whose query lengths agree.
+
+Its index-topk kernels recover a request from a flat query row as
+``row // max_query_len`` and the row's causal cutoff as ``seq_len -
+max_query_len + tok + 1``. Both are only true when every decode request
+contributes exactly ``max_query_len`` rows.
+
+Speculative decode breaks that: a request that joins the batch without draft
+tokens contributes one row while its neighbours contribute ``num_spec + 1``,
+and vLLM keeps all of them in the decode segment because each query length is
+still within the reorder threshold. The batch then has ``total_q`` rows that no
+longer divide by ``max_query_len`` -- which is how this surfaced in practice,
+as ``total_q 121 not divisible by max_query_len 4`` after a few hundred gsm8k
+requests, while short runs happened to stay uniform and passed.
+
+So the property under test is not "does the assert fire" but "does the builder
+recognise a ragged decode segment before the kernels see it".
+"""
+
+import pytest
+import torch
+
+pytest.importorskip("vllm")
+
+from atom.plugin.vllm.attention.metadata import _uniform_decode_query_len
+
+
+def _starts(query_lens):
+    """cu_seqlens_q for the given per-request query lengths."""
+    return torch.tensor(
+        [0] + list(torch.tensor(query_lens).cumsum(0)), dtype=torch.int32
+    )
+
+
+@pytest.mark.parametrize("query_len", [1, 2, 4, 8])
+def test_uniform_segment_reports_its_query_len(query_len):
+    starts = _starts([query_len] * 6)
+    assert _uniform_decode_query_len(starts, 6) == query_len
+
+
+def test_single_request_is_uniform():
+    assert _uniform_decode_query_len(_starts([4]), 1) == 4
+
+
+@pytest.mark.parametrize(
+    "query_lens",
+    [
+        [4, 4, 1, 4],  # a request joined without draft tokens
+        [1, 4, 4, 4],  # ... at the head of the segment
+        [4, 4, 4, 1],  # ... at the tail
+        [4, 2, 4, 4],  # partially accepted drafts
+    ],
+)
+def test_ragged_segment_is_rejected(query_lens):
+    assert _uniform_decode_query_len(_starts(query_lens), len(query_lens)) is None
+
+
+def test_ragged_segment_whose_total_still_divides_is_rejected():
+    """total_q % max_query_len == 0 is not enough to make a segment uniform.
+
+    [4, 4, 6, 2] sums to 16 == 4 * 4, so the kernel's own assert would pass and
+    it would then read rows 8..11 as request 2 -- silently wrong instead of
+    loud. The builder has to reject on the query lengths themselves.
+    """
+    query_lens = [4, 4, 6, 2]
+    starts = _starts(query_lens)
+    assert int(starts[-1]) % 4 == 0
+    assert _uniform_decode_query_len(starts, len(query_lens)) is None
+
+
+def test_only_the_decode_prefix_is_inspected():
+    """Trailing prefill requests must not make a uniform decode segment look ragged."""
+    starts = _starts([4, 4, 4, 512, 300])
+    assert _uniform_decode_query_len(starts, 3) == 4
+
+
+def test_no_decode_requests():
+    assert _uniform_decode_query_len(_starts([512]), 0) is None
+    assert _uniform_decode_query_len(None, 4) is None
+
+
+def test_truncated_query_start_loc_is_rejected():
+    """Fewer offsets than num_decodes+1 means the caller cannot be trusted."""
+    assert _uniform_decode_query_len(_starts([4, 4]), 5) is None