[TRTLLM-14093][feat] Eagle3 for M3 non-MSA path + overlap scheduler compatibility

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,

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -45,14 +45,16 @@
 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,
     EventType,
     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-

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)
 

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)

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(

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

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

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

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):

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]