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 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/atom/plugin/vllm/attention/minimax_m3_attnetion.py b/atom/plugin/vllm/attention/minimax_m3_attnetion.py index e2755556f2..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,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 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