From 448c489e1cfaa86434b65fc69c80a38f0e36f151 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Thu, 9 Jul 2026 03:56:40 +0000 Subject: [PATCH] [TRTLLM-14093][feat] Enable one-model Eagle3 speculative decoding for MiniMax-M3 Squashed port of feat/minimax-m3-eagle3 (PR #16021) onto feat/branch_m3. Target-side enablement for the Inferact/MiniMax-M3-EAGLE3 draft head on the reference (non-MSA) sparse backend: - Wire MiniMaxM3ForCausalLM as SpecDecOneEngineForCausalLM and capture Eagle3 aux hidden states at layer exit (fully TP-reduced; no cross-layer allreduce+norm fusion). - Rebase MiniMaxM3AttentionMetadata onto TrtllmAttentionMetadata (the shared per-step metadata is consumed by TRTLLM draft layers, the Eagle3 one-model worker, and engine isinstance gates; precedent: DSA and the M3 MSA metadata). - Route multi-token generation rows (spec verify: 1 + draft_len tokens) through the extend path; the decode branch stays reserved for batches where every row appends exactly one token. - Vectorized sync-free builder shared with a new on_update_kv_lens hook that re-derives seq_lens/prefix_lens/q_positions/out_cache_loc on device from the corrected kv_lens_cuda under overlap scheduler + spec (DSA pattern); in-bounds clamps cover the optimistic page-boundary overhang (max_seqlen_k SDPA width + slot gathers). - V1-family draft manager support in KVCacheManagerV2.add_dummy_requests (exception-safe); attention-DP dummy requests register in the draft manager; MiniMaxM3KVCacheManagerV2 opts out of shared draft layers under attention DP (its AttentionOp tensors are synthetic). - Creation-time guards: tree modes, disabled separate draft KV (disagg WAR), CUDA graphs, and MSA+spec (the in-builder MSA rejection is hoisted above routing so mixed batches cannot bypass it). - Accuracy test test_nvfp4_eagle3 (MMLU + GSM8K + acceptance probe, attention_dp parametrized) + reference rows. Validated on 4xB200 (NVFP4 tp4/ep4, draft_len=3, overlap scheduler on, eager) at this commit: - test_nvfp4_eagle3[attention_dp=False]: MMLU 85.50 / GSM8K 89.73 (refs 83/88), acceptance rate 0.709 / mean acceptance length 3.126 - test_nvfp4_eagle3[attention_dp=True]: MMLU 85.14 / GSM8K 91.32, acceptance rate 0.767 / mean acceptance length 3.302 - batch-1 greedy: 6.05 -> 17.05 tok/s (2.82x) - spec-off boot + generation clean (TRTLLM-Gen warmup now runs for M3 and is harmless) Signed-off-by: Zheyu Fu --- .../sparse/minimax_m3/cache_manager.py | 6 + .../sparse/minimax_m3/metadata.py | 242 ++++++++++++------ .../_torch/models/modeling_minimaxm3.py | 20 +- tensorrt_llm/_torch/pyexecutor/_util.py | 11 +- .../_torch/pyexecutor/kv_cache_manager_v2.py | 66 +++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 5 + .../_torch/pyexecutor/py_executor_creator.py | 44 +++- .../defs/accuracy/references/gsm8k.yaml | 3 + .../defs/accuracy/references/mmlu.yaml | 3 + .../defs/accuracy/test_llm_api_pytorch.py | 74 ++++++ .../test_lists/qa/llm_function_core.txt | 2 + 11 files changed, 374 insertions(+), 102 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py index 3785c9f16e97..a7808e880464 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py @@ -144,6 +144,12 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): * ``sparse_index_dim`` — width of the index-K/V vectors. """ + # The AttentionOp-facing tensors this manager builds are synthetic + # placeholders over INDEX_KEY-coalesced pools, so one-model speculative + # draft layers must live in a separate manager even under attention DP + # (read by ``_should_create_separate_draft_kv_cache``). + supports_shared_draft_layers = False + def __init__( self, *args, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py index 62286c761404..3b4ea2544f4e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py @@ -447,7 +447,13 @@ def prepare(self) -> None: if batch_size == 0: self.max_seqlen_k = 1 else: - self.max_seqlen_k = int(self.seq_lens_cpu[:batch_size].max().item()) + max_k = int(self.seq_lens_cpu[:batch_size].max().item()) + # Optimistic overlap+spec lengths can overhang the page table at + # a page boundary (see derive_q_positions_and_cache_slots); the + # dense fallback consumes max_seqlen_k as the exact SDPA + # mask/gather width, so an unclamped overhang crashes SDPA + # (mask 385 vs K 384 on MMLU). + self.max_seqlen_k = min(max_k, int(self.req_to_token.shape[1])) def ensure_metadata_on_device( @@ -621,6 +627,51 @@ def maybe_build_static_buffers_placeholder( return placeholder +def derive_q_positions_and_cache_slots( + req_to_token: torch.Tensor, + prefix_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + q_batch_row: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-Q-token K-side positions and KV slot ids, on device, sync-free. + + ``q_positions[t] = prefix_lens[row(t)] + (t - cu_seqlens_q[row(t)])``; + ``out_cache_loc[t] = req_to_token[row(t), q_positions[t]]``. Shared by + the metadata builder and the ``on_update_kv_lens`` re-derivation so the + two cannot drift. + """ + total_q = int(q_batch_row.shape[0]) + qbr = q_batch_row.to(torch.long) + tok = torch.arange(total_q, dtype=torch.int32, device=q_batch_row.device) + q_positions = prefix_lens[qbr] + (tok - cu_seqlens_q[qbr]) + # Overlap+spec: the first derivation runs with optimistic + # (full-acceptance) prefix_lens, which overhang the last allocated page + # when a request's true length sits at a page boundary. The overhanging + # slots are placeholders (on_update_kv_lens re-derives them from the + # corrected kv_lens before any forward consumes them), but the gather + # must stay in bounds: unclamped, the last row trips the CUDA gather + # assert and inner rows silently read the next request's slots. Clamp + # the table index only (non-inplace: ``.to`` aliases int64 inputs) — + # q_positions keep the optimistic values, whose mask width prepare() + # bounds separately. + idx = q_positions.to(torch.long).clamp(min=0, max=req_to_token.shape[1] - 1) + flat = qbr * req_to_token.shape[1] + idx + return q_positions, req_to_token.reshape(-1).index_select(0, flat) + + +def derive_decode_cache_slots(req_to_token: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor: + """Decode-row KV slot ids (new token at position ``seq_lens[b] - 1``). + + Same in-bounds-placeholder clamp contract as + :func:`derive_q_positions_and_cache_slots`; the ``min=0`` floor + additionally covers zero-length dummy rows indexing ``-1``. + """ + rows = torch.arange(seq_lens.shape[0], device=seq_lens.device, dtype=torch.long) + idx = (seq_lens.to(torch.long) - 1).clamp_(min=0, max=req_to_token.shape[1] - 1) + flat = rows * req_to_token.shape[1] + idx + return req_to_token.reshape(-1).index_select(0, flat) + + def build_runtime_metadata_from_kv_manager( *, kv_cache_manager, @@ -757,21 +808,13 @@ def build_runtime_metadata_from_kv_manager( req_to_token = req_to_token_fresh slot_ids = torch.arange(batch, device=device, dtype=torch.int32) - # Compute out_cache_loc: per-new-token slot ids, in flattened order - # matching the q-token order the model layer projects. The Python - # loops below run on CPU lists derived from the CPU-resident - # ``seq_lens_cpu`` / ``prefix_lens`` / ``extend_seq_lens_cpu``, so - # no GPU sync is needed at this point. The resulting - # ``out_cache_loc`` tensor is constructed directly on ``device``. - # The ``int(...item())`` reads against ``req_to_token`` are a CPU - # sync but only ever run from ``prepare()`` (outside any CUDA-graph - # capture window) — they are not in the forward path. + # out_cache_loc must be flattened in the q-token order the model layer + # projects, or K/V lands in the wrong requests' slots. if is_prefill: if extend_seq_lens_cpu is None: raise ValueError("prefill metadata requires extend_seq_lens_cpu") if prefix_lens is None: raise ValueError("prefill metadata requires prefix_lens") - prefix_lens_cpu = prefix_lens.to("cpu").tolist() if static_buffers is not None: prefix_buf = static_buffers["prefix_lens"] prefix_src = prefix_lens.to(device=device, dtype=torch.int32, non_blocking=True) @@ -781,17 +824,18 @@ def build_runtime_metadata_from_kv_manager( prefix_lens_dev = ( prefix_lens.to(device) if prefix_lens.device != device else prefix_lens ) - out_cache_loc_list: List[int] = [] cu_q: List[int] = [0] - req_to_token_cpu = req_to_token_fresh.to("cpu") - for b in range(batch): - pref = int(prefix_lens_cpu[b]) - ext = int(extend_seq_lens_cpu[b]) - for offset in range(ext): - slot = int(req_to_token_cpu[b, pref + offset].item()) - out_cache_loc_list.append(slot) - cu_q.append(cu_q[-1] + ext) + for ext in extend_seq_lens_cpu: + cu_q.append(cu_q[-1] + int(ext)) total_q = cu_q[-1] + cu_seqlens_q_src = torch.tensor(cu_q, dtype=torch.int32, device=device) + q_batch_row_src = torch.repeat_interleave( + torch.arange(batch, device=device, dtype=torch.int32), + torch.tensor(extend_seq_lens_cpu, dtype=torch.int64, device=device), + ) + q_positions_src, out_cache_loc_src = derive_q_positions_and_cache_slots( + req_to_token, prefix_lens_dev, cu_seqlens_q_src, q_batch_row_src + ) if static_buffers is not None: if total_q > static_buffers["max_num_tokens"]: raise ValueError( @@ -799,33 +843,22 @@ def build_runtime_metadata_from_kv_manager( f"is smaller than current total_q={total_q}" ) out_cache_loc_buf = static_buffers["out_cache_loc"] - out_cache_loc_src = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) out_cache_loc_buf[:total_q].copy_(out_cache_loc_src, non_blocking=True) out_cache_loc = out_cache_loc_buf[:total_q] cu_seqlens_q_buf = static_buffers["cu_seqlens_q"] - cu_seqlens_q_src = torch.tensor(cu_q, dtype=torch.int32, device=device) cu_seqlens_q_buf[: batch + 1].copy_(cu_seqlens_q_src, non_blocking=True) cu_seqlens_q = cu_seqlens_q_buf[: batch + 1] - # Populate persistent q_batch_row / q_positions in-place so - # the inner metadata's prepare() can leave them alone. q_batch_row_buf = static_buffers["q_batch_row"] q_positions_buf = static_buffers["q_positions"] - for b in range(batch): - start, end = cu_q[b], cu_q[b + 1] - if end > start: - q_batch_row_buf[start:end] = b - pref = int(prefix_lens_cpu[b]) - offsets = ( - torch.arange(start, end, device=device, dtype=torch.int32) - start + pref - ) - q_positions_buf[start:end].copy_(offsets, non_blocking=True) + q_batch_row_buf[:total_q].copy_(q_batch_row_src, non_blocking=True) + q_positions_buf[:total_q].copy_(q_positions_src, non_blocking=True) q_batch_row = q_batch_row_buf[:total_q] q_positions = q_positions_buf[:total_q] else: - out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) - cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device) - q_batch_row = None - q_positions = None + out_cache_loc = out_cache_loc_src + cu_seqlens_q = cu_seqlens_q_src + q_batch_row = q_batch_row_src + q_positions = q_positions_src meta = MiniMaxM3SparseAttentionMetadata( is_prefill=True, req_to_token=req_to_token, @@ -840,12 +873,7 @@ def build_runtime_metadata_from_kv_manager( ) else: # Decode: the new token sits at position seq_lens[b] - 1. - seq_lens_cpu_list = seq_lens_cpu.to("cpu").tolist() - out_cache_loc_list = [] - req_to_token_cpu = req_to_token_fresh.to("cpu") - for b in range(batch): - pos = int(seq_lens_cpu_list[b]) - 1 - out_cache_loc_list.append(int(req_to_token_cpu[b, pos].item())) + out_cache_loc_src = derive_decode_cache_slots(req_to_token, seq_lens_dev) if static_buffers is not None: if batch > static_buffers["max_num_tokens"]: raise ValueError( @@ -853,11 +881,10 @@ def build_runtime_metadata_from_kv_manager( f"is smaller than current batch={batch}" ) out_cache_loc_buf = static_buffers["out_cache_loc"] - out_cache_loc_src = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) out_cache_loc_buf[:batch].copy_(out_cache_loc_src, non_blocking=True) out_cache_loc = out_cache_loc_buf[:batch] else: - out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) + out_cache_loc = out_cache_loc_src meta = MiniMaxM3SparseAttentionMetadata( is_prefill=False, req_to_token=req_to_token, @@ -1001,10 +1028,31 @@ def build_m3_sparse_metadata_and_plans( use_msa = bool(getattr(kv_cache_manager, "use_msa", False)) - is_extend = num_contexts > 0 + new_tokens_per_seq = [int(seq_lens_cpu[b].item()) for b in range(batch_size)] + # The MSA path assumes a single query token per generation row (the + # decode wrapper packs qo_len=1, and the eager prefill plan's CPU + # values go stale under overlap kv-len correction). Speculative + # decoding emits 1 + draft_len query tokens per gen row; those rows + # route to the extend branch below, so reject BEFORE routing — a + # decode-branch check would silently miss mixed batches. + if use_msa and any(n > 1 for n in new_tokens_per_seq[num_contexts:]): + raise NotImplementedError( + "MiniMax-M3 MSA sparse attention does not support speculative " + "decoding (multiple query tokens per generation step). Disable " + "speculative decoding or use the non-MSA MiniMax-M3 backend." + ) + + # Any batch containing a context request takes the extend path. + # Pure-generation batches also take it when any row carries more than + # one new token this step: one-model speculative decoding (Eagle3) + # verifies gen rows with 1 + draft_len tokens, which the decode branch + # (one slot per row) cannot represent. Decode stays reserved for + # batches where every row appends exactly one token, keeping the + # non-speculative CUDA-graph decode geometry unchanged. + is_extend = num_contexts > 0 or any(n > 1 for n in new_tokens_per_seq) if is_extend: prefix_lens_list = [int(num_cached_per_seq[b]) for b in range(batch_size)] - extend_seq_lens_cpu = [kv_lens_cpu_list[b] - prefix_lens_list[b] for b in range(batch_size)] + extend_seq_lens_cpu = new_tokens_per_seq prefix_lens = torch.tensor(prefix_lens_list, dtype=torch.int32, device=cache_device) m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( kv_cache_manager=kv_cache_manager, @@ -1018,19 +1066,6 @@ def build_m3_sparse_metadata_and_plans( static_buffers=static_buffers, ) else: - # The MSA decode path assumes a single query token per request - # (decode_wrapper/dispatch.py packs qo_len=1). Speculative - # decoding emits multiple draft query tokens per generation step, - # which this path cannot represent, so reject it with a clear - # error rather than silently mis-staging out_cache_loc. The Triton - # reference path shares this builder but keeps its prior behavior, - # so gate the rejection on the MSA backend only. - if use_msa and batch_size > 0 and int(seq_lens_cpu[:batch_size].max().item()) > 1: - raise NotImplementedError( - "MiniMax-M3 MSA sparse attention does not support speculative " - "decoding (multiple query tokens per decode step). Disable " - "speculative decoding or use the non-MSA MiniMax-M3 backend." - ) m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( kv_cache_manager=kv_cache_manager, request_ids=request_ids, @@ -1064,18 +1099,38 @@ def build_m3_sparse_metadata_and_plans( @functools.lru_cache(maxsize=1) def get_minimax_m3_attention_metadata_cls(): - """Return :class:`MiniMaxM3AttentionMetadata` (lazy import).""" - from ...interface import AttentionMetadata - - class MiniMaxM3AttentionMetadata(AttentionMetadata): - """:class:`AttentionMetadata` that pre-builds MiniMax-M3 metadata. - - `prepare()` builds the per-forward `MiniMaxM3SparseAttentionMetadata` - and per-new-token `out_cache_loc` once per scheduler step, outside - the CUDA-graph capture window, and publishes them as - `self.m3_sparse_metadata` / `self.m3_out_cache_loc` so the forward - reads them with no capture-time CPU->GPU copies. Test paths may set - those attributes directly instead of going through `prepare()`. + """Return :class:`MiniMaxM3AttentionMetadata` (lazy import). + + Subclassing :class:`TrtllmAttentionMetadata` (precedent: + ``DSAtrtllmAttentionMetadata``, and ``MiniMaxM3MSATrtllmAttentionMetadata`` + on the MSA path) rather than the plain :class:`AttentionMetadata` is + required for one-model speculative decoding (Eagle3): the draft layers + live in the same engine and run :class:`TrtllmAttention` against this + shared per-step metadata, so it must carry the TRTLLM surface + (``kv_lens_cuda``, ``host_request_types``, ``kv_cache_block_offsets``, + ``draft_kv_cache_block_offsets``, ``update_spec_dec_param``, ...). The + engine also gates spec-dec plumbing (draft KV cache swap, + overlap-scheduler ``kv_lens_cuda`` fixups) on + ``isinstance(..., TrtllmAttentionMetadata)``. The M3 sparse layers keep + consuming only the published ``m3_sparse_metadata`` attachment; the + TRTLLM-side block-offset tensors built against the M3 cache manager are + consistent but unused by the target layers (see + ``MiniMaxM3KVCacheManagerV2._build_pool_mapping_tensors``). + """ + from ...trtllm import TrtllmAttentionMetadata + + class MiniMaxM3AttentionMetadata(TrtllmAttentionMetadata): + """:class:`TrtllmAttentionMetadata` that pre-builds MiniMax-M3 metadata. + + `prepare()` first runs the full TRTLLM preparation (kv lens, request + types, block offsets — consumed by one-model speculative draft + layers), then builds the per-forward + `MiniMaxM3SparseAttentionMetadata` and per-new-token `out_cache_loc` + once per scheduler step, outside the CUDA-graph capture window, and + publishes them as `self.m3_sparse_metadata` / `self.m3_out_cache_loc` + so the forward reads them with no capture-time CPU->GPU copies. Test + paths may set those attributes directly instead of going through + `prepare()`. """ m3_sparse_metadata: Optional["MiniMaxM3SparseAttentionMetadata"] = None @@ -1100,6 +1155,47 @@ def prepare(self) -> None: self.m3_out_cache_loc = None build_m3_sparse_metadata_and_plans(self, geometry=get_global_msa_geometry()) + def on_update_kv_lens(self) -> None: + """Re-derive the M3 attachment from the corrected ``kv_lens_cuda``. + + Under the overlap scheduler + speculative decoding, prepare() + runs with optimistic cached counts (full draft acceptance) and + the engine corrects ``kv_lens_cuda`` on device before invoking + this hook (same pattern as ``DSAtrtllmAttentionMetadata``). + On-device, sync-free, and idempotent; ``seq_lens_cpu`` / + ``max_seqlen_k`` keep the optimistic values — they only bound + arange widths that the kernels mask by ``seq_lens``. + """ + super().on_update_kv_lens() + meta = self.m3_sparse_metadata + if meta is None: + return + batch = int(meta.slot_ids.shape[0]) + kv_lens = self.kv_lens_cuda[:batch] + meta.seq_lens[:batch].copy_(kv_lens) + if meta.is_prefill: + # Only the K-side prefix moves with rejections; the Q-side + # structure (cu_seqlens_q, q_batch_row) is fixed per step. + total_q = int(meta.q_positions.shape[0]) + cu = meta.cu_seqlens_q + meta.prefix_lens[:batch].copy_(kv_lens - (cu[1 : batch + 1] - cu[:batch])) + q_positions, out_cache_loc = derive_q_positions_and_cache_slots( + meta.req_to_token, + meta.prefix_lens[:batch], + cu, + meta.q_batch_row[:total_q], + ) + meta.q_positions[:total_q].copy_(q_positions) + self.m3_out_cache_loc[:total_q].copy_(out_cache_loc) + else: + # Reached today only as an identity (the hook also fires + # pre-correction on ordinary decode steps); re-deriving + # keeps corrected 0-draft steps correct once dynamic + # draft lengths make them reachable. + self.m3_out_cache_loc[:batch].copy_( + derive_decode_cache_slots(meta.req_to_token, kv_lens) + ) + return MiniMaxM3AttentionMetadata @@ -1110,6 +1206,8 @@ def prepare(self) -> None: "build_m3_sparse_metadata_and_plans", "build_runtime_metadata_from_kv_manager", "build_stable_kv_indices", + "derive_decode_cache_slots", + "derive_q_positions_and_cache_slots", "ensure_metadata_on_device", "get_global_msa_geometry", "get_minimax_m3_attention_metadata_cls", diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index c4b84ac558a4..91726b671ced 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -45,6 +45,7 @@ from ..modules.linear import Linear, TensorParallelMode, copy_weight, load_weight_shard from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..speculative import SpecMetadata from ..utils import ( ActivationType, AuxStreamType, @@ -52,7 +53,8 @@ get_model_extra_attrs, is_torch_compiling, ) -from .modeling_utils import DecoderModel, DecoderModelForCausalLM, ModelConfig, register_auto_model +from .modeling_speculative import SpecDecOneEngineForCausalLM +from .modeling_utils import DecoderModel, ModelConfig, register_auto_model # Dense layers use SDPA with non-contiguous Q/K/V and a bool attn_mask. # Limit backends to memory-efficient and math; cuDNN SDPA fails for this layout, @@ -1426,6 +1428,7 @@ def forward( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: if residual is None: @@ -1446,6 +1449,10 @@ def forward( hidden_states = self.block_sparse_moe(hidden_states, attn_metadata) else: hidden_states = self.mlp(hidden_states) + # hidden_states is fully TP-reduced at layer exit (no cross-layer + # allreduce+norm fusion). + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual) return hidden_states, residual @@ -1496,6 +1503,7 @@ def forward( input_ids: Optional[torch.IntTensor] = None, position_ids: Optional[torch.IntTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: if (input_ids is None) ^ (inputs_embeds is not None): @@ -1512,6 +1520,7 @@ def forward( hidden_states=hidden_states, attn_metadata=attn_metadata, residual=residual, + spec_metadata=spec_metadata, ) hidden_states, _ = self.norm(hidden_states, residual) @@ -1533,19 +1542,14 @@ def forward( @register_auto_model("MiniMaxM3SparseForCausalLM") -class MiniMaxM3ForCausalLM(DecoderModelForCausalLM[MiniMaxM3Model, PretrainedConfig]): +class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]): """Text-only M3 model.""" def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): raw_pretrained = model_config.pretrained_config if is_minimax_m3_vl_config(raw_pretrained): model_config = get_text_model_config(model_config) - super().__init__( - MiniMaxM3Model(model_config), - config=model_config, - hidden_size=model_config.pretrained_config.hidden_size, - vocab_size=model_config.pretrained_config.vocab_size, - ) + super().__init__(MiniMaxM3Model(model_config), model_config) def load_weights(self, weights, *args, **kwargs): # Merge the M3-specific gate-bias rename into any caller- diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 02ffb035705a..a2f17a1a4d4e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -952,10 +952,13 @@ def _should_create_separate_draft_kv_cache(self) -> bool: in the target model and don't produce a separate ModelConfig. We fall back to the target model's config via _get_effective_draft_config(). """ - if self._mapping.enable_attention_dp: - logger.info( - "Attention DP is enabled, separate draft KV cache is not supported." - ) + if self._mapping.enable_attention_dp and getattr( + self._kv_cache_manager_cls, 'supports_shared_draft_layers', + True): + # Back-compat: attention DP keeps the shared-manager layout + # existing deployments were validated with. + logger.info("Attention DP: draft layers share the target KV " + "cache manager.") return False return should_use_separate_draft_kv_cache(self._speculative_config) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 0911c5041652..aa4c94804ba6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2716,22 +2716,39 @@ def release_resources( return None draft_kv_cache = None if draft_kv_cache_manager is not None: - draft_kv_cache = draft_kv_cache_manager._create_kv_cache( - req.py_request_id, req.lora_task_id, input_tokens, is_dummy=req.is_dummy - ) - # Dummy path: see comment above, no salt. - if draft_kv_cache is None: - release_resources(req) - return None - success = draft_kv_cache.resume(draft_kv_cache_manager._stream.cuda_stream) - if not success: - release_resources(req, free_draft_resources=True) - return None - draft_kv_cache.stop_committing() - success = draft_kv_cache.resize(dummy_capacity) - if not success: - release_resources(req, free_draft_resources=True) - return None + if isinstance(draft_kv_cache_manager, KVCacheManagerV2): + draft_kv_cache = draft_kv_cache_manager._create_kv_cache( + req.py_request_id, req.lora_task_id, input_tokens, is_dummy=req.is_dummy + ) + # Dummy path: see comment above, no salt. + if draft_kv_cache is None: + release_resources(req) + return None + success = draft_kv_cache.resume(draft_kv_cache_manager._stream.cuda_stream) + if not success: + release_resources(req, free_draft_resources=True) + return None + draft_kv_cache.stop_committing() + success = draft_kv_cache.resize(dummy_capacity) + if not success: + release_resources(req, free_draft_resources=True) + return None + else: + # V1-family draft manager (no per-request cache handles); + # mirrors KVCacheManager.add_dummy_requests. The C++ side + # raises on allocation failure rather than returning a + # status, so release before propagating. + draft_seq_added = False + try: + draft_kv_cache_manager.impl.add_sequence_batch( + [(req.py_request_id, token_num, beam_width)], [req] + ) + draft_seq_added = True + for _ in range(self.num_extra_kv_tokens): + draft_kv_cache_manager.impl.add_token(req.py_request_id) + except Exception: + release_resources(req, free_draft_resources=draft_seq_added) + raise if is_gen: req.state = LlmRequestState.GENERATION_IN_PROGRESS @@ -2742,13 +2759,28 @@ def release_resources( new_capacity = kv_cache.capacity + _kv_draft + 1 success = kv_cache.resize(new_capacity, history_length=history_hint) if not success: - release_resources(req, free_draft_resources=draft_kv_cache is not None) + # V1-family draft allocations have no draft_kv_cache + # handle, so key on the manager, not the handle. + release_resources( + req, + free_draft_resources=draft_kv_cache_manager is not None, + ) return None if draft_kv_cache is not None: success = draft_kv_cache.resize(new_capacity) if not success: release_resources(req, free_draft_resources=True) return None + elif draft_kv_cache_manager is not None: + # Gen dummies must expose a 1 + draft_len kv span to + # the draft layers; a V1 manager only grows a + # sequence via add_token. + try: + for _ in range(_kv_draft): + draft_kv_cache_manager.impl.add_token(req.py_request_id) + except Exception: + release_resources(req, free_draft_resources=True) + raise if use_mrope: _populate_dummy_mrope_config(req, token_num, is_gen) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index fcd91503d86f..2bc3eebead10 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4787,12 +4787,17 @@ def _pad_attention_dp_dummy_request(self): and self.max_num_tokens is not None): token_nums = [self.max_num_tokens] dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID] + # A separate draft KV cache manager must also see the dummy, or + # its prepare_resources hits an unknown request id. + draft_kv_cache_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) llm_request = self.kv_cache_manager.add_dummy_requests( request_ids=dummy_request_ids, token_nums=token_nums, is_gen=self._adp_dummy_is_gen, prepare_resource=True, max_num_draft_tokens=self.max_total_draft_tokens, + draft_kv_cache_manager=draft_kv_cache_manager, )[0] llm_request.is_attention_dp_dummy = True spec_resource_manager = self.resource_manager.get_resource_manager( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 797f2fd48666..4c2fc95b4370 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -31,7 +31,8 @@ from ..attention_backend.trtllm import TrtllmAttention from ..distributed import Distributed from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter, - get_spec_resource_manager) + get_spec_resource_manager, + should_use_separate_draft_kv_cache) from ..virtual_memory import scope as virtual_memory_scope from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, create_py_executor_instance, instantiate_sampler, is_mla, @@ -632,6 +633,47 @@ def drafting_loop_wrapper(model): max_num_tokens = model_engine.max_num_tokens sparse_attention_config = model_engine.sparse_attention_config + # The MSA kernel path packs one query token per generation row and its + # prefill plan stages CPU values that go stale under overlap kv-len + # correction, so it cannot verify draft tokens (two-model spec also + # emits multi-token target rows) — reject at creation instead of at the + # first verify step inside prepare(). + if (sparse_attention_config is not None + and sparse_attention_config.algorithm == "minimax_m3" + and getattr(sparse_attention_config, "sparse_use_msa", False) + and spec_config is not None): + raise NotImplementedError( + "Speculative decoding is not supported with the MiniMax-M3 MSA " + "kernel path (sparse_use_msa=True): MSA packs one query token " + "per generation row. Set sparse_use_msa=False to use the " + "reference sparse backend with speculative decoding.") + + if (sparse_attention_config is not None + and sparse_attention_config.algorithm == "minimax_m3" + and spec_config is not None + and spec_config.spec_dec_mode.is_eagle3_one_model()): + if not spec_config.is_linear_tree: + raise NotImplementedError( + "Tree-based speculative decoding (eagle_choices / " + "use_dynamic_tree) is not supported with MiniMax-M3 sparse " + "attention: the M3 sparse kernels implement linear-chain " + "verification only. Remove eagle_choices / use_dynamic_tree " + "from the speculative config.") + if llm_args.cuda_graph_config is not None: + raise NotImplementedError( + "CUDA graphs are not supported with MiniMax-M3 sparse " + "attention and speculative decoding: multi-token verify " + "routes through the M3 extend path, which is not " + "capture-safe yet. Set cuda_graph_config to null.") + if not should_use_separate_draft_kv_cache(spec_config): + raise NotImplementedError( + "One-model speculative decoding with MiniMax-M3 sparse " + "attention requires a separate draft KV cache manager, but " + "it is disabled for this configuration (e.g. disaggregated " + "serving disables it as a WAR for nvbug 5807902). Use " + "two-model speculative decoding (eagle3_one_model=False) " + "instead.") + # Set default value for cache_transceiver_config.max_tokens_in_buffer if cache_transceiver_config and cache_transceiver_config.max_tokens_in_buffer is None: cache_transceiver_config.max_tokens_in_buffer = net_max_seq_len diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index f6c8404baa51..10e52aacb833 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -457,6 +457,9 @@ MiniMaxAI/MiniMax-M3-MXFP8: nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION accuracy: 88 + - quant_algo: MIXED_PRECISION + spec_dec_algo: Eagle3 + accuracy: 88 nvidia/NVIDIA-Nemotron-Nano-9B-v2: - accuracy: 85.027 - quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index d79fa2b0e41b..a965c962f8b2 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -284,6 +284,9 @@ MiniMaxAI/MiniMax-M3-MXFP8: nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION accuracy: 83 + - quant_algo: MIXED_PRECISION + spec_dec_algo: Eagle3 + accuracy: 83 moonshotai/Kimi-K2-Instruct: - quant_algo: FP8_BLOCK_SCALES accuracy: 87.65 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ea78c4b79df7..cd4c7222adcd 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7686,6 +7686,80 @@ def test_nvfp4(self, use_msa): task = GSM8K(model_name) task.evaluate(llm) + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(140000) + @parametrize_with_ids("overlap_scheduler", [True]) + @parametrize_with_ids("attention_dp", [False, True]) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) + def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, + overlap_scheduler): + model_name = "nvidia/MiniMax-M3-NVFP4" + model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4" + max_draft_len = 3 + spec_config = Eagle3DecodingConfig( + max_draft_len=max_draft_len, + speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3", + ) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, + enable_block_reuse=False) + with LLM(model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + sparse_attention_config=MiniMaxM3SparseAttentionConfig(), + moe_config=MoeConfig(backend="CUTLASS"), + max_seq_len=4096, + speculative_config=spec_config, + cuda_graph_config=None, + disable_overlap_scheduler=not overlap_scheduler, + enable_attention_dp=attention_dp, + trust_remote_code=True) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + task = MMLU(model_name) + task.evaluate(llm) + task = GSM8K(model_name) + task.evaluate(llm) + + # Acceptance probe (pattern: TestNemotronV3Ultra + # test_nvfp4_4gpu_mtp_ar): stream a few greedy prompts and + # derive per-step acceptance from the token increments. + raw_prompts = [ + "Solve step by step: what is 12 times 17?", + "Write a Python function that reverses a linked list.", + "The capital of France is", + ] + prompts = [ + llm.tokenizer.apply_chat_template( + [{ + "role": "user", + "content": p + }], + tokenize=False, + add_generation_prompt=True, + ) for p in raw_prompts + ] + tok_ids = [llm.tokenizer.encode(p) for p in prompts] + sampling_params = SamplingParams(max_tokens=128, temperature=0) + total_drafted = 0 + total_accepted = 0 + total_steps = 0 + for i in range(len(tok_ids)): + num_tokens = 0 + for output in llm.generate_async(tok_ids[i], + sampling_params, + streaming=True): + new_tokens = output.outputs[0].token_ids + total_drafted += max_draft_len + total_accepted += len(new_tokens) - num_tokens - 1 + total_steps += 1 + num_tokens = len(new_tokens) + accept_rate = total_accepted / total_drafted + accept_length = 1 + total_accepted / total_steps + print(f"MiniMax-M3 Eagle3 acceptance: rate={accept_rate:.3f}, " + f"mean acceptance length={accept_length:.3f}") + assert accept_rate > 0.25, \ + f"Eagle3 acceptance rate too low: {accept_rate:.3f}" + @skip_pre_blackwell class TestGLM5FP8(LlmapiAccuracyTestHarness): diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 4789118aefe2..4bc149a95b08 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -682,6 +682,8 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_si accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm]