From e59db3f81937a96d322e9aa2d34cecfd27766218 Mon Sep 17 00:00:00 2001 From: perzhang Date: Wed, 2 Sep 2026 23:42:09 -0500 Subject: [PATCH 1/4] fix(minimax-m3): keep the vLLM-plugin index cache in a real fp8 dtype Serving MiniMax-M3 through the vLLM plugin backend with --kv-cache-dtype fp8 kills EngineCore on the first request: triton.compiler.errors.CompilationError: at 73:17 qk = tl.dot(q, k, out_dtype=tl.float32) * sm_scale_log2e Both operands must be same dtype. Got bf16 and uint8 MiniMaxM3SparseIndexerCache takes its KV-cache spec dtype from vLLM's kv_cache_dtype_str_to_dtype(), and vLLM maps every fp8 kv-cache-dtype to torch.uint8 -- a byte buffer its own kernels reinterpret. ATOM's _index_block_score_kernel dispatches on k.dtype.is_fp8() instead, so a uint8-labelled index cache sends it down the bf16 branch, where it dots bf16 against uint8 and fails to compile. The bytes in that buffer are already fp8: aiter's fused_qknorm_idxrqknorm writes them under kv_cache_dtype="fp8", and the native ATOM server allocates the same cache as dtypes.d_dtypes["fp8"]. Only the torch dtype label was wrong, so label it aiter.dtypes.fp8 (torch.float8_e4m3fn). The element stays one byte, so vLLM's page-size accounting is unchanged, and vLLM already allows a true fp8 dtype in this position -- its own "fp8_inc" entry maps to torch.float8_e4m3fn. Scoped to the indexer cache deliberately: the main sparse and dense KV caches keep their uint8 label because the aiter paged kernels read them as raw bytes. Verified on 4x MI355X (gfx950), TP=4, amd/MiniMax-M3-MXFP4, vLLM 0.27.2.dev0+g6e448d0ea, --kv-cache-dtype fp8. Before: the first chat request kills EngineCore. After: the index cache is allocated as (num_blocks, 128, 128) torch.float8_e4m3fn, chat completions answer normally, and a 90K-token 90%-prefix-hit aiperf run completes clean (32/32 requests, TTFT p50 2.28 s, TPOT p50 14.25 ms at concurrency 8). Co-Authored-By: Claude Opus 5 (1M context) --- .../vllm/attention/minimax_m3_attnetion.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/atom/plugin/vllm/attention/minimax_m3_attnetion.py b/atom/plugin/vllm/attention/minimax_m3_attnetion.py index e2755556f2..dc4c22392d 100644 --- a/atom/plugin/vllm/attention/minimax_m3_attnetion.py +++ b/atom/plugin/vllm/attention/minimax_m3_attnetion.py @@ -86,9 +86,20 @@ 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( - kv_cache_dtype, vllm_config.model_config - ) + if str(kv_cache_dtype).startswith("fp8"): + # vLLM maps every fp8 kv-cache-dtype to torch.uint8, a byte buffer + # its own kernels reinterpret. The index-topk kernel dispatches on + # ``k.dtype.is_fp8()`` instead, so a uint8-labelled cache sends it + # down the bf16 branch and it dots bf16 against uint8. aiter already + # writes real fp8 bytes here and the native server labels this cache + # ``dtypes.d_dtypes["fp8"]``: same 1-byte element, correct label. + from aiter import dtypes as aiter_dtypes + + self.kv_cache_torch_dtype = aiter_dtypes.fp8 + else: + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + kv_cache_dtype, vllm_config.model_config + ) self.num_kv_heads = 1 self.head_size = head_dim self.head_size_v = head_dim From d8cbdcdc9a667b9b1533ef424737dc80d7ee4938 Mon Sep 17 00:00:00 2001 From: perzhang Date: Thu, 3 Sep 2026 05:59:33 -0500 Subject: [PATCH 2/4] fix(minimax-m3): give the index-cache dtype exception one named home Review feedback: why not handle this in kv_cache_dtype_str_to_dtype. That table is vLLM's, and it is read by every attention layer of every model. Its fp8 -> torch.uint8 entry is deliberate: vLLM's own 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 -- which is why _page16_shuffle_cache_for_sparse_kernel does the .view() itself. Relabelling fp8 there would relabel those caches too, on every model ATOM serves through the plugin, to fix one side cache. The exception is layout-specific, so it belongs where the layout is known. vLLM resolves the same way for the same reason: _resolve_dsv4_kv_cache_dtype returns uint8 for the packed fp8_ds_mla layout and torch.float8_e4m3fn for the plain-row one from the same --kv-cache-dtype fp8, and vLLM's own MiniMaxM3IndexerCache skips kv_cache_dtype_str_to_dtype entirely and labels its index cache torch.float8_e4m3fn. So keep the exception here, but collapse the inline branch into one named function with the reasoning attached, and take the fp8 handle from dtypes.d_dtypes -- the same handle the native server resolves this cache through (_resolve_index_cache_dtype). aiter picks that dtype per gfx target, so the two paths stay on one arch-correct fp8 label instead of a hard-coded torch.float8_e4m3fn. No behavior change: dtypes.d_dtypes["fp8"] is dtypes.fp8, and non-fp8 strings still go to vLLM's mapping. Re-checked in the same container the fix was verified in (gfx950, aiter fp8 = torch.float8_e4m3fn): the helper returns float8_e4m3fn for "fp8"/"fp8_e4m3", float16 for "auto", bfloat16 for "bfloat16". Co-Authored-By: Claude Opus 5 (1M context) --- .../vllm/attention/minimax_m3_attnetion.py | 52 +++++++++++++------ 1 file changed, 37 insertions(+), 15 deletions(-) diff --git a/atom/plugin/vllm/attention/minimax_m3_attnetion.py b/atom/plugin/vllm/attention/minimax_m3_attnetion.py index dc4c22392d..5f70548c26 100644 --- a/atom/plugin/vllm/attention/minimax_m3_attnetion.py +++ b/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,20 +119,9 @@ def __init__( self.attn_type = AttentionType.DECODER self.attn_backend = SparseMHAIndexerBackend self.kv_cache_dtype = kv_cache_dtype - if str(kv_cache_dtype).startswith("fp8"): - # vLLM maps every fp8 kv-cache-dtype to torch.uint8, a byte buffer - # its own kernels reinterpret. The index-topk kernel dispatches on - # ``k.dtype.is_fp8()`` instead, so a uint8-labelled cache sends it - # down the bf16 branch and it dots bf16 against uint8. aiter already - # writes real fp8 bytes here and the native server labels this cache - # ``dtypes.d_dtypes["fp8"]``: same 1-byte element, correct label. - from aiter import dtypes as aiter_dtypes - - self.kv_cache_torch_dtype = aiter_dtypes.fp8 - else: - self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( - kv_cache_dtype, vllm_config.model_config - ) + self.kv_cache_torch_dtype = _index_cache_torch_dtype( + kv_cache_dtype, vllm_config.model_config + ) self.num_kv_heads = 1 self.head_size = head_dim self.head_size_v = head_dim From 1a3d27c2e0003c77f47026a473fb002e2a9d9d22 Mon Sep 17 00:00:00 2001 From: perzhang Date: Thu, 3 Sep 2026 08:04:06 -0500 Subject: [PATCH 3/4] [MiniMax-M3][Spec decode] Check that a decode segment is uniform before using it MiniMax-M3 with an EAGLE3 draft dies part way into a long run: atom/model_ops/minimax_m3/index_topk.py minimax_m3_index_topk_decode AssertionError: total_q 121 not divisible by max_query_len 4 Three ATOM decode paths divide a flat query row by max_query_len to recover the request it belongs to: - the M3 index-topk kernels use `row // max_query_len`, and the row's causal cutoff `seq_len - max_query_len + tok + 1` - aiter's gluon paged decode reshapes q to `[q.shape[0] // max_query_len, max_query_len, ...]` All three only hold when every decode request contributes exactly max_query_len rows, and none of them checked. Speculative decode breaks the assumption on both models. On the target, 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 (121 = 30*4 + 1). On the draft, a request contributes however many tokens the previous step accepted, so the segment is ragged by construction -- which surfaced as aiter/ops/triton/gluon/pa_decode_gluon.py RuntimeError: shape '[27, 4, 1, 16, 128]' is invalid for input of size 225280 (110 rows where 27 requests * 4 = 108 were assumed). Add `_uniform_decode_query_len()` and consult it in all three builders: - MinimaxM3SparseAttentionMetadataBuilder.build(): only take the uniform decode fast path when the segment really is uniform; otherwise hand those requests to the prefill kernel, which derives causality from cu_seqlens_q/context_lens and accepts variable query lengths. - AiterMhaMetadataBuilderForVllm.build(): route a ragged decode segment to the extend path, which is varlen. Decode requests sort before extends, so widening the extend segment covers them without reordering. - AiterMhaMetadataBuilderForVllm.build_for_drafting(): its docstring claimed "during EAGLE/MTP drafting all requests are uniform decodes" and it only tested for prefills, so it never reached build(). Fall back to build() when the batch is ragged. Both mixed-batch branches also stopped reporting a max_query_len they had not measured -- the M3 one passed reorder_batch_threshold (always num_spec+1) and the MHA one a max over the segment. A plain-decode segment under a spec-decode threshold would take 4 there, and with a row count that happens to divide by 4 the kernels would mis-map every row silently instead of asserting. The kernel-side assert stays as a backstop, now spelling out the invariant. Verified on 4x MI355X TP4 with MiniMax-M3-MXFP8 + Inferact/MiniMax-M3-EAGLE3-GQA (num_speculative_tokens=3), gsm8k 5-shot chat over the full 1319 questions: before: crashes (total_q 121 not divisible by max_query_len 4) after: 0.9530 flexible-extract, mean acceptance length 3.27, no errors no-spec regression, same build: 0.9484 (0.9477 before this change) Co-Authored-By: Claude Opus 5 (1M context) --- atom/model_ops/minimax_m3/index_topk.py | 11 +- atom/plugin/vllm/attention/metadata.py | 122 ++++++++++++++++-- .../test_minimax_m3_decode_uniformity.py | 87 +++++++++++++ 3 files changed, 203 insertions(+), 17 deletions(-) create mode 100644 tests/plugin/test_minimax_m3_decode_uniformity.py diff --git a/atom/model_ops/minimax_m3/index_topk.py b/atom/model_ops/minimax_m3/index_topk.py index e346fc4282..1686b03984 100644 --- a/atom/model_ops/minimax_m3/index_topk.py +++ b/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) diff --git a/atom/plugin/vllm/attention/metadata.py b/atom/plugin/vllm/attention/metadata.py index 6099bc0815..a9726bbab3 100644 --- a/atom/plugin/vllm/attention/metadata.py +++ b/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, diff --git a/tests/plugin/test_minimax_m3_decode_uniformity.py b/tests/plugin/test_minimax_m3_decode_uniformity.py new file mode 100644 index 0000000000..c5e6429d3e --- /dev/null +++ b/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 From 8f76794c4f14344c23cab063aaf323ce0e111492 Mon Sep 17 00:00:00 2001 From: perzhang Date: Thu, 3 Sep 2026 08:19:48 -0500 Subject: [PATCH 4/4] [ATOM Plugin CI] Cover M3 under fp8 KV + EAGLE3 in the vLLM accuracy matrix The M3 cell ran with --kv-cache-dtype auto and no speculative config, so it exercised neither fix in this PR: the index cache never reached the fp8 label path, and the decode segment was uniform by construction. Replace it with the combination that did break -- fp8 KV cache plus an EAGLE3 draft (Inferact/MiniMax-M3-EAGLE3-GQA, num_speculative_tokens=3). Everything else about the cell is unchanged. Measured on 4x MI355X TP4 with MiniMax-M3-MXFP8 and the same draft, gsm8k 5-shot chat over the full 1319 questions: 0.9530 flexible-extract with a mean acceptance length of 3.27, against 0.9484 for the same build without spec decode. The threshold stays 0.93; the native accuracy catalog lists 0.9469 for MXFP4 + EAGLE3 at the same threshold. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/atom-vllm-test.yaml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/atom-vllm-test.yaml b/.github/workflows/atom-vllm-test.yaml index 46e284f61a..e0e124b38f 100644 --- a/.github/workflows/atom-vllm-test.yaml +++ b/.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