[None][feat] MiniMax-M3 eagle+disagg pipe clean: share Eagle3 draft layers via a sub-page view

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -1797,9 +1797,15 @@ def _build_per_layer_num_kv_heads(
         draft_pretrained, 'num_key_value_heads',
         getattr(draft_pretrained, 'num_attention_heads', None))
 
-    if draft_num_kv_heads is None or draft_num_kv_heads == num_key_value_heads:
+    if draft_num_kv_heads is None:
         return num_key_value_heads
 
+    # Return the extended per-layer list even when the drafter's head count
+    # equals the target's: the list's draft tail is what tells shared-draft
+    # managers that the appended layers exist. An equal-head drafter (e.g.
+    # MiniMax-M3's GQA Eagle head, 4 KV heads like the target) otherwise
+    # disappears into the target's layer range and gets routed through the
+    # target's attention machinery.
     num_spec_layers = get_num_spec_layers(spec_config)
     logger.info(f"Per-layer KV heads for speculative decoding: "
                 f"target={num_key_value_heads} x {num_hidden_layers} layers, "
@@ -1980,6 +1986,19 @@ def _create_kv_cache_manager(
     manager_extra_kwargs = {}
     if issubclass(kv_cache_manager_cls, KVCacheManagerV2):
         manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats
+        # One-model spec with shared draft layers appends the drafter's
+        # layers to this manager; tell the manager how many. Anchor on the
+        # pretrained TARGET layer count — local num_hidden_layers may already
+        # include the draft tail. Consumed by managers with a draft sub-page
+        # view (MiniMax-M3); others ignore it. Masked/cross flows yield a
+        # non-positive delta and correctly report 0.
+        target_num_layers = getattr(config, "num_hidden_layers", None)
+        num_appended_draft_layers = (len(per_layer_num_kv_heads) -
+                                     target_num_layers
+                                     if isinstance(per_layer_num_kv_heads, list)
+                                     and target_num_layers is not None else 0)
+        manager_extra_kwargs["num_one_model_draft_layers"] = max(
+            0, num_appended_draft_layers)
     if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2):
         manager_extra_kwargs["is_disagg"] = is_disagg
 

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -514,23 +514,11 @@ def create_py_executor(
         # WAR for https://nvbugs/5807902
         # Disable separate draft KV cache in disaggregated mode
         # Enable separate pool for None DI + Non-KVBM and Aggregated + KVBM
-        #
-        # MiniMax-M3 is exempt: its cache manager forbids sharing draft
-        # layers (supports_shared_draft_layers=False), so this WAR's
-        # shared-manager fallback does not exist for it. Both worker roles
-        # must keep the separate manager, or their target pool layouts
-        # diverge and disaggregated KV transfer breaks. Checked via the
-        # sparse-attention algorithm because the manager class is not yet
-        # resolved here ("minimax_m3" maps 1:1 to MiniMaxM3KVCacheManagerV2).
+        # (The shared-manager fallback is exactly what MiniMax-M3 wants here
+        # — its drafter shares the target manager and rides prefix reuse and
+        # KV transfer natively — so M3's earlier #17341 exemption is retired.)
         if cache_transceiver_config is not None:
-            if is_minimax_m3(m3_sparse_config):
-                logger.warning(
-                    "Disaggregated MiniMax-M3 keeps the separate draft KV "
-                    "cache manager; draft-layer KV is not transferred, so "
-                    "generation-side acceptance is reduced until the drafter "
-                    "rebuilds its context window.")
-            else:
-                spec_config._allow_separate_draft_kv_cache = False
+            spec_config._allow_separate_draft_kv_cache = False
 
     # chunk_unit_size may be changed to 64 when using flash mla
     attn_runtime_features = AttentionRuntimeFeatures(

tensorrt_llm/_torch/speculative/interface.py

@@ -32,8 +32,7 @@
 from ..attention_backend.trtllm import (AttentionBackend, TrtllmAttention,
                                         TrtllmAttentionMetadata)
 from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE
-from ..pyexecutor.resource_manager import (BaseResourceManager,
-                                           ResourceManagerType)
+from ..pyexecutor.resource_manager import BaseResourceManager
 
 if TYPE_CHECKING:
     from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
@@ -2109,12 +2108,13 @@ def _prepare_context_input_ids(self, input_ids, num_ctx_tokens, gather_ids,
 
     def get_draft_kv_cache_manager(self, resource_manager):
         """
-        Get the draft KV cache manager if using separate KV cache layouts.
+        Get the draft-side KV manager (separate manager or the target
+        manager's draft sub-page view); see resolve_draft_kv_cache_manager.
         """
-        if self.use_separate_draft_kv_cache and resource_manager is not None:
-            return resource_manager.get_resource_manager(
-                ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
-        return None
+        if resource_manager is None:
+            return None
+        from .utils import resolve_draft_kv_cache_manager
+        return resolve_draft_kv_cache_manager(resource_manager)
 
     @contextmanager
     def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager):

tensorrt_llm/_torch/speculative/utils.py

@@ -523,19 +523,41 @@ def get_num_extra_kv_tokens(spec_config):
     return 0
 
 
+def resolve_draft_kv_cache_manager(resource_manager):
+    """Resolve the draft-side KV manager for one-model speculative decoding.
+
+    The registered separate manager is the ground truth when present;
+    otherwise ask the target manager for its draft sub-page view (managers
+    without one — e.g. same-geometry shared drafters — resolve to None and
+    the drafter attends the shared manager directly).
+    """
+    from ..pyexecutor.resource_manager import ResourceManagerType
+
+    draft_manager = resource_manager.get_resource_manager(
+        ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
+    if draft_manager is not None:
+        return draft_manager
+    target_manager = resource_manager.get_resource_manager(
+        ResourceManagerType.KV_CACHE_MANAGER)
+    # getattr fetches the method without executing it, so a failure inside
+    # view construction propagates from the call itself instead of being
+    # silently swallowed into "no view".
+    get_view = getattr(target_manager, "get_draft_subpage_view", None)
+    return get_view() if get_view is not None else None
+
+
 def get_draft_kv_cache_manager(spec_config, resource_manager):
     """
     Returns the draft KV cache manager only in one-model speculative decoding
-    mode where the target model manages a separate draft KV cache.
+    mode: the separate manager when the target manages one, or the target
+    manager's draft sub-page view when shared draft layers run at a smaller
+    kernel page size (e.g. MiniMax-M3). See resolve_draft_kv_cache_manager.
     """
-    from ..pyexecutor.resource_manager import ResourceManagerType
-
     if spec_config is None:
         return None
     if not spec_config.spec_dec_mode.use_one_engine():
         return None
-    return resource_manager.get_resource_manager(
-        ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
+    return resolve_draft_kv_cache_manager(resource_manager)
 
 
 def update_spec_config_from_model_config(spec_config, model_config):

tests/integration/defs/accuracy/test_disaggregated_serving.py

@@ -46,7 +46,9 @@ def result(self):
         return self
 
 
-DuckLLM = namedtuple('DuckLLM', ['args', 'tokenizer', 'generate_async'])
+DuckLLM = namedtuple('DuckLLM',
+                     ['args', 'tokenizer', 'generate_async', 'router_url'],
+                     defaults=(None, ))
 
 # Timeout for the entire test
 DEFAULT_TEST_TIMEOUT = 3600
@@ -511,7 +513,8 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100):
 
         tokenizer = load_hf_tokenizer(model_name)
         try:
-            yield DuckLLM(args, tokenizer, generate_async)
+            yield DuckLLM(args, tokenizer, generate_async,
+                          f"http://localhost:{serve_port}")
         finally:
             if enable_perf:
                 _show_kvcache_time(kv_cache_perf_dir)

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -7883,16 +7883,203 @@ def test_nvfp4(self, use_msa, eval_mode):
                 task = GSM8K(model_name)
                 task.evaluate(llm)
 
+    def _run_nvfp4_eagle3_disagg(self, model_name, model_path, max_draft_len,
+                                 inferencemax, attention_dp, overlap_scheduler,
+                                 use_msa, cuda_graph):
+        """Disaggregated arm of test_nvfp4_eagle3.
+
+        CI coverage for Eagle3 + the unified (shared) draft KV cache
+        crossing the disaggregated transceiver on one 4-GPU node: context
+        TEP2 -> generation TEP2 over NIXL, block reuse on the context
+        server. ``attention_dp`` selects the generation flavor: True is
+        the AgentX-submission shape (attention-DP generation), False the
+        TEP production-candidate shape. Every path-gating knob mirrors
+        the production disaggregated configs (fix15a production-candidate
+        sweep, GB300); values that are workload-tuned (1M-context sizes)
+        or GPU-generation-tuned (memory fractions) are CI-adjusted and
+        called out inline. Accuracy is asserted through the router; the
+        drafter's KV rides the shared logical blocks, so a corrupted or
+        dropped drafter cache collapses accuracy. Acceptance stats stay
+        with the aggregated arm, which exercises the same view code.
+        """
+        if not (overlap_scheduler and cuda_graph and use_msa):
+            pytest.skip("the disagg arm pins the production serving shape "
+                        "(overlap scheduler + CUDA graphs + MSA)")
+        from .test_disaggregated_serving import launch_disaggregated_llm
+        speculative_config = {
+            "decoding_type": "Eagle3",
+            "max_draft_len": max_draft_len,
+            "speculative_model": f"{llm_models_root()}/MiniMax-M3-EAGLE3",
+            "eagle3_one_model": True,
+        }
+        common_config = {
+            "tensor_parallel_size": 2,
+            "moe_expert_parallel_size": 2,
+            "speculative_config": speculative_config,
+            "sparse_attention_config": {
+                "algorithm": "minimax_m3",
+                "implementation": "msa",
+                "fuse_qkv_index_projection": True,
+                "indexer_kv_dtype": "fp8",
+            },
+            "cache_transceiver_config": {
+                "backend": "NIXL",
+                "transceiver_runtime": "PYTHON",
+                "kv_cache_bounce_size_mb": 0,
+                "kv_transfer_timeout_ms": 600000,
+            },
+            "moe_config": {
+                "backend": "TRTLLM"
+            },
+            "scheduler_config": {
+                "capacity_scheduler_policy": "MAX_UTILIZATION"
+            },
+            "enable_autotuner": True,
+            "return_perf_metrics": True,
+            "perf_metrics_max_requests": 1000,
+            # The InferenceMAX protocol needs room for thinking output;
+            # production serves 1M-token contexts here.
+            "max_seq_len": 16384 if inferencemax else 4096,
+            "trust_remote_code": True,
+        }
+        kv_cache_common = {
+            "dtype": "fp8",
+            "tokens_per_block": 128,
+            "use_kv_cache_manager_v2": True,
+            "event_buffer_max_size": 0,
+            # Production runs 0.7 (ctx) / 0.9 (gen) on GB300; one derated
+            # value keeps headroom on the B200 CI stage.
+            "free_gpu_memory_fraction": 0.7,
+        }
+        ctx_server_config = {
+            **common_config,
+            "disable_overlap_scheduler": not overlap_scheduler,
+            "enable_attention_dp": False,
+            "enable_chunked_prefill": True,
+            "attention_dp_config": {
+                "enable_kv_cache_aware_routing": False,
+                "kv_cache_routing_conversation_affinity": True,
+                "kv_cache_routing_max_sessions": 65536,
+            },
+            "kv_cache_config": {
+                **kv_cache_common, "enable_block_reuse": True
+            },
+            "max_batch_size": 4,
+            "max_num_tokens": 32768,
+            "cuda_graph_config": None,
+        }
+        gen_server_config = {
+            **common_config,
+            "disable_overlap_scheduler": not overlap_scheduler,
+            "enable_attention_dp": attention_dp,
+            "enable_lm_head_tp_in_adp": False,
+            "kv_cache_config": {
+                **kv_cache_common, "enable_block_reuse": False
+            },
+            "max_batch_size": 16,
+            # Decode-only token budget: (1 + draft_len) verify tokens per
+            # request x max_batch_size (the production sweep's value is
+            # its no-spec arm's).
+            "max_num_tokens": (1 + max_draft_len) * 16,
+            "num_postprocess_workers": 4,
+            "stream_interval": 100,
+            "enable_iter_perf_stats": True,
+            "cuda_graph_config": {
+                "enable_padding": True,
+                "batch_sizes": [1, 2, 4, 8, 16],
+            },
+        }
+        disaggregated_server_config = {
+            "hostname": "localhost",
+            "backend": "pytorch",
+            "context_servers": {
+                "num_instances": 1
+            },
+            "generation_servers": {
+                "num_instances": 1
+            },
+        }
+        with launch_disaggregated_llm(disaggregated_server_config,
+                                      ctx_server_config,
+                                      gen_server_config,
+                                      model_path,
+                                      server_waiting_timeout=1800) as llm:
+            # The launcher stamps quant_algo=NVFP4 from the model name; the
+            # checkpoint's hf_quant_config is MIXED_PRECISION (which is what
+            # the in-process arm asserts and the accuracy references key on).
+            llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION
+            task = (GSM8KInferenceMax(model_name)
+                    if inferencemax else GSM8K(model_name))
+            task.evaluate(llm)
+
+            # Chat-format acceptance probe — the same workload and
+            # thresholds as the aggregated arm's (200 GSM8K questions,
+            # chat template, greedy, 512 tokens; drafter card reference
+            # rate 0.839 / length 3.518). The disagg fixture has no
+            # get_stats, so the probe window comes from the generation
+            # worker's /metrics buffer (enable_iter_perf_stats), drained
+            # right before the probe; the dataset loads from the models
+            # root so the probe works offline like the eval does.
+            import requests
+            info = requests.get(f"{llm.router_url}/cluster_info",
+                                timeout=30).json()
+            gen_worker = info["current_workers"]["generation_servers"][0]
+            gen_url = f'{gen_worker["host"]}:{gen_worker["port"]}'
+            questions = [
+                r["question"]
+                for r in load_dataset(GSM8K.DATASET_DIR, "main", split="test")
+            ][:200]
+            chat_prompts = [
+                llm.tokenizer.apply_chat_template([{
+                    "role": "user",
+                    "content": q
+                }],
+                                                  tokenize=False,
+                                                  add_generation_prompt=True)
+                for q in questions
+            ]
+            requests.get(f"http://{gen_url}/metrics", timeout=30)  # drain
+            probe_params = SamplingParams(max_tokens=512, temperature=0)
+            for future in [
+                    llm.generate_async(prompt, probe_params)
+                    for prompt in chat_prompts
+            ]:
+                future.result()
+            records = requests.get(f"http://{gen_url}/metrics",
+                                   timeout=30).json()
+            drafted = accepted = steps = 0
+            for record in records:
+                stats = record.get("specDecodingStats") or {}
+                drafted += stats.get("numDraftTokens", 0)
+                accepted += stats.get("numAcceptedTokens", 0)
+                steps += stats.get("numRequestsWithDraftTokens", 0)
+            assert steps > 0, "no speculative iterations in /metrics"
+            chat_rate = accepted / drafted
+            chat_length = 1 + accepted / steps
+            print(f"MiniMax-M3 Eagle3 disagg chat-GSM8K acceptance: rate="
+                  f"{chat_rate:.3f}, mean acceptance length="
+                  f"{chat_length:.3f} ({steps} spec iterations)")
+            assert chat_rate > 0.78, \
+                f"Eagle3 chat-GSM8K acceptance rate too low: " \
+                f"{chat_rate:.3f} (threshold 0.78, reference 0.839 from " \
+                f"the drafter card)"
+            assert chat_length > 3.3, \
+                f"Eagle3 chat-GSM8K acceptance length too low: " \
+                f"{chat_length:.3f} (threshold 3.3, reference 3.518 from " \
+                f"the drafter card)"
+
     @pytest.mark.skip_less_device(4)
     @pytest.mark.skip_less_device_memory(140000)
+    @parametrize_with_ids("disagg", [False, True])
     @parametrize_with_ids("eval_mode", ["default", "inferencemax"])
     @parametrize_with_ids("cuda_graph", [True])
     @parametrize_with_ids("use_msa", [True])
     @parametrize_with_ids("overlap_scheduler", [False, 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, use_msa, cuda_graph, eval_mode):
+                          overlap_scheduler, use_msa, cuda_graph, eval_mode,
+                          disagg):
         if use_msa:
             from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import \
                 msa_package_available
@@ -7902,6 +8089,12 @@ def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp,
         model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4"
         inferencemax = eval_mode == "inferencemax"
         max_draft_len = 3
+        if disagg:
+            self._run_nvfp4_eagle3_disagg(model_name, model_path, max_draft_len,
+                                          inferencemax, attention_dp,
+                                          overlap_scheduler, use_msa,
+                                          cuda_graph)
+            return
         spec_config = Eagle3DecodingConfig(
             max_draft_len=max_draft_len,
             speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3",

tests/integration/test_lists/qa/llm_function_core.txt

@@ -665,8 +665,8 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] TIMEO
 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-eval_mode=inferencemax] TIMEOUT (180)
 accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True-eval_mode=inferencemax] TIMEOUT (60)
-accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-use_msa=True-cuda_graph=True-eval_mode=inferencemax] TIMEOUT (180)
-accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-use_msa=True-cuda_graph=True-eval_mode=inferencemax] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-use_msa=True-cuda_graph=True-eval_mode=inferencemax-disagg=False] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-use_msa=True-cuda_graph=True-eval_mode=inferencemax-disagg=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]