From db8b81d3b7a23b36bd03659a3254e37d797d238b Mon Sep 17 00:00:00 2001
From: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Date: Fri, 6 Nov 2026 09:50:05 -0800
Subject: [PATCH 1/3] [https://nvbugs/7028061][fix] Align MiniMax M3 Eagle3
 draft-loop API

Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
---
 .../minimax_m3_nvfp4__sm_103__tp4/modeling.py |  19 ++-
 .../minimax_m3_nvfp4__sm_103__tp8/modeling.py |  19 ++-
 .../integration/test_lists/test-db/l0_cpu.yml |   1 +
 tests/integration/test_lists/waives.txt       |   2 -
 .../minimax_m3/test_eagle3_draft_loop.py      | 128 ++++++++++++++++++
 5 files changed, 155 insertions(+), 14 deletions(-)
 create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py
index c0bf5cd12b17..e52d8c9fa4ad 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py
@@ -3709,6 +3709,7 @@ def _forward_impl(
             )
 
         if self._uses_external_shared_target_kv:
+            next_new_tokens = None
             next_draft_tokens = self._forward_external_shared_target_kv_draft_loop(
                 position_ids=position_ids,
                 hidden_states=hidden_states,
@@ -3721,6 +3722,8 @@ def _forward_impl(
                 batch_size=batch_size,
             )
         else:
+            # Retain the sequence lengths before draft preparation replaces them.
+            pre_draft_seq_lens = attn_metadata._seq_lens_cuda
             # Save the old attn_metadata and spec_metadata
             self._prepare_attn_metadata_for_spec_dec(attn_metadata)
 
@@ -3743,7 +3746,7 @@ def _forward_impl(
             # Get the draft KV cache manager if using separate layouts
             draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager)
 
-            next_draft_tokens = self._forward_draft_loop(
+            next_draft_tokens, next_new_tokens = self._forward_draft_loop(
                 inputs,
                 attn_metadata,
                 spec_metadata,
@@ -3755,6 +3758,8 @@ def _forward_impl(
                 num_accepted_tokens,
                 original_all_rank_num_tokens,
                 resource_manager,
+                accepted_tokens=accepted_tokens,
+                pre_draft_seq_lens=pre_draft_seq_lens,
             )
             # restore attn_metadata to support cuda graph
             self._restore_attn_metadata_from_spec_dec(attn_metadata)
@@ -3763,10 +3768,7 @@ def _forward_impl(
                 attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens
 
         # prepare next new tokens to support overlap scheduler
-        if self._m3_next_new_tokens is not None:
-            # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..).
-            next_new_tokens = self._m3_next_new_tokens
-        else:
+        if next_new_tokens is None:
             next_new_tokens = self._prepare_next_new_tokens(
                 accepted_tokens,
                 next_draft_tokens,
@@ -4290,10 +4292,12 @@ def _forward_draft_loop(
         num_accepted_tokens,
         original_all_rank_num_tokens,
         resource_manager,
+        accepted_tokens,
+        pre_draft_seq_lens,
     ):
         """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop."""
         if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model):
-            return self._m3_fused_draft_loop(
+            next_draft_tokens = self._m3_fused_draft_loop(
                 inputs,
                 attn_metadata,
                 spec_metadata,
@@ -4302,6 +4306,7 @@ def _forward_draft_loop(
                 batch_size,
                 num_accepted_tokens,
             )
+            return next_draft_tokens, self._m3_next_new_tokens
         return super()._forward_draft_loop(
             inputs,
             attn_metadata,
@@ -4314,6 +4319,8 @@ def _forward_draft_loop(
             num_accepted_tokens,
             original_all_rank_num_tokens,
             resource_manager,
+            accepted_tokens=accepted_tokens,
+            pre_draft_seq_lens=pre_draft_seq_lens,
         )
 
     def _m3_fused_draft_loop(
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py
index 1e8c520ca9d6..d7cd18563211 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py
@@ -3429,6 +3429,7 @@ def _forward_impl(
             )
 
         if self._uses_external_shared_target_kv:
+            next_new_tokens = None
             next_draft_tokens = self._forward_external_shared_target_kv_draft_loop(
                 position_ids=position_ids,
                 hidden_states=hidden_states,
@@ -3441,6 +3442,8 @@ def _forward_impl(
                 batch_size=batch_size,
             )
         else:
+            # Retain the sequence lengths before draft preparation replaces them.
+            pre_draft_seq_lens = attn_metadata._seq_lens_cuda
             # Save the old attn_metadata and spec_metadata
             self._prepare_attn_metadata_for_spec_dec(attn_metadata)
 
@@ -3463,7 +3466,7 @@ def _forward_impl(
             # Get the draft KV cache manager if using separate layouts
             draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager)
 
-            next_draft_tokens = self._forward_draft_loop(
+            next_draft_tokens, next_new_tokens = self._forward_draft_loop(
                 inputs,
                 attn_metadata,
                 spec_metadata,
@@ -3475,6 +3478,8 @@ def _forward_impl(
                 num_accepted_tokens,
                 original_all_rank_num_tokens,
                 resource_manager,
+                accepted_tokens=accepted_tokens,
+                pre_draft_seq_lens=pre_draft_seq_lens,
             )
             # restore attn_metadata to support cuda graph
             self._restore_attn_metadata_from_spec_dec(attn_metadata)
@@ -3483,10 +3488,7 @@ def _forward_impl(
                 attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens
 
         # prepare next new tokens to support overlap scheduler
-        if self._m3_next_new_tokens is not None:
-            # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..).
-            next_new_tokens = self._m3_next_new_tokens
-        else:
+        if next_new_tokens is None:
             next_new_tokens = self._prepare_next_new_tokens(
                 accepted_tokens,
                 next_draft_tokens,
@@ -3959,10 +3961,12 @@ def _forward_draft_loop(
         num_accepted_tokens,
         original_all_rank_num_tokens,
         resource_manager,
+        accepted_tokens,
+        pre_draft_seq_lens,
     ):
         """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop."""
         if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model):
-            return self._m3_fused_draft_loop(
+            next_draft_tokens = self._m3_fused_draft_loop(
                 inputs,
                 attn_metadata,
                 spec_metadata,
@@ -3971,6 +3975,7 @@ def _forward_draft_loop(
                 batch_size,
                 num_accepted_tokens,
             )
+            return next_draft_tokens, self._m3_next_new_tokens
         return super()._forward_draft_loop(
             inputs,
             attn_metadata,
@@ -3983,6 +3988,8 @@ def _forward_draft_loop(
             num_accepted_tokens,
             original_all_rank_num_tokens,
             resource_manager,
+            accepted_tokens=accepted_tokens,
+            pre_draft_seq_lens=pre_draft_seq_lens,
         )
 
     def _m3_fused_draft_loop(
diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml
index 2e816594cd1f..df9b38804e35 100644
--- a/tests/integration/test_lists/test-db/l0_cpu.yml
+++ b/tests/integration/test_lists/test-db/l0_cpu.yml
@@ -19,6 +19,7 @@ l0_cpu:
   # MiniMax-M3 Rubin target unit tests.
   - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py
   # MiniMax-M3 target unit tests.
+  - unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py
   - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py
   - unittest/_torch/modules/test_linear_output_slice.py::test_fp8_weight_slice_cache
   - unittest/_torch/modules/test_linear_output_slice.py::test_prepare_output_slice_rejects_unsupported_projection_before_forward
diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index cc4b8c6f8047..66ddca19dc47 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -1,5 +1,3 @@
-accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_eagle3_gsm8k_and_acceptance SKIP (https://nvbugs/7028061)
-accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_eagle3_gsm8k_and_acceptance SKIP (https://nvbugs/7028061)
 accuracy/test_disaggregated_serving.py::TestGPTOSS::test_kv_cache_v2_nixl_python[cache_mgr_v2] SKIP (https://nvbugs/7023112)
 accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6402058)
diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py
new file mode 100644
index 000000000000..518618ab2ce9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py
@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Exercise the MiniMax worker's Python dispatch without importing GPU kernels."""
+
+import ast
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock, sentinel
+
+import pytest
+
+pytestmark = pytest.mark.cpu_only
+
+_TORCH_ROOT = Path(__file__).resolve().parents[6] / "tensorrt_llm" / "_torch"
+
+
+def _load_worker(tp: int) -> type:
+    # Execute the actual methods, including the base dispatcher and its signature.
+    # Keeping the class definitions preserves zero-argument super() in the override.
+    namespace = {}
+    sources = (
+        (_TORCH_ROOT / "speculative/eagle3.py", "Eagle3OneModelWorker", {"_forward_draft_loop"}),
+        (
+            _TORCH_ROOT
+            / f"_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp{tp}/modeling.py",
+            "M3Eagle3Worker",
+            {"_forward_impl", "_forward_draft_loop"},
+        ),
+    )
+    for path, name, methods in sources:
+        cls = next(
+            node
+            for node in ast.parse(path.read_text()).body
+            if isinstance(node, ast.ClassDef) and node.name == name
+        )
+        cls.body = [
+            node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in methods
+        ]
+        assert len(cls.body) == len(methods)
+        if name == "Eagle3OneModelWorker":
+            cls.bases = []
+        module = ast.Module(body=[cls], type_ignores=[])
+        exec(compile(module, str(path), "exec"), namespace)
+    return namespace["M3Eagle3Worker"]
+
+
+@pytest.mark.parametrize("tp", [4, 8])
+@pytest.mark.parametrize("path", ["fallback", "fallback_buffer", "fused", "external"])
+def test_eagle3_draft_loop_contract(tp: int, path: str) -> None:
+    worker = _load_worker(tp)()
+    worker._is_mamba_hybrid_cache = False
+    worker._auxiliary_state_handlers = []
+    worker.sa_enhancer = None
+    worker._uses_external_shared_target_kv = path == "external"
+    worker._execute_guided_decoder_if_present = Mock()
+    worker.sample_and_accept_draft_tokens = Mock(
+        return_value=(sentinel.accepted, sentinel.accepted_counts)
+    )
+    worker.draft_only_eligible = Mock(return_value=False)
+    metadata = SimpleNamespace(
+        num_seqs=2,
+        num_contexts=1,
+        _seq_lens_cuda=sentinel.pre_draft_seq_lens,
+        all_rank_num_tokens=sentinel.rank_tokens,
+    )
+    spec = SimpleNamespace(
+        runtime_draft_len=3,
+        spec_resource_manager=SimpleNamespace(),
+        batch_indices_cuda=sentinel.batch_indices,
+    )
+
+    def prepare_metadata(md: SimpleNamespace) -> None:
+        # The caller must retain the sequence lengths from before preparation.
+        md._seq_lens_cuda = sentinel.draft_seq_lens
+
+    worker._prepare_attn_metadata_for_spec_dec = Mock(side_effect=prepare_metadata)
+    worker._restore_attn_metadata_from_spec_dec = Mock()
+    worker.prepare_1st_drafter_inputs = Mock(return_value=sentinel.inputs)
+    worker.get_draft_kv_cache_manager = Mock(return_value=sentinel.cache)
+    worker._m3_fused_loop_ok = Mock(return_value=path == "fused")
+    prebuilt = sentinel.next_new if path == "fallback_buffer" else None
+    worker._forward_linear_draft_loop = Mock(return_value=(sentinel.drafts, prebuilt))
+
+    def fused_loop(*args: object) -> object:
+        worker._m3_next_new_tokens = sentinel.next_new
+        return sentinel.drafts
+
+    worker._m3_fused_draft_loop = Mock(side_effect=fused_loop)
+    worker._forward_external_shared_target_kv_draft_loop = Mock(return_value=sentinel.drafts)
+    worker._prepare_next_new_tokens = Mock(return_value=sentinel.next_new)
+
+    outputs = worker._forward_impl(
+        sentinel.input_ids,
+        Mock(),
+        sentinel.hidden_states,
+        sentinel.logits,
+        metadata,
+        spec,
+        sentinel.draft_model,
+        sentinel.resources,
+    )
+
+    assert outputs["next_draft_tokens"] is sentinel.drafts
+    assert outputs["next_new_tokens"] is sentinel.next_new
+    assert outputs["new_tokens"] is sentinel.accepted
+    assert metadata.use_spec_decoding
+    if path.startswith("fallback"):
+        worker._forward_linear_draft_loop.assert_called_once_with(
+            sentinel.inputs,
+            metadata,
+            spec,
+            sentinel.draft_model,
+            sentinel.cache,
+            1,
+            2,
+            sentinel.accepted_counts,
+            sentinel.rank_tokens,
+            sentinel.accepted,
+            sentinel.pre_draft_seq_lens,
+        )
+    else:
+        worker._forward_linear_draft_loop.assert_not_called()
+    if path in ("fallback_buffer", "fused"):
+        worker._prepare_next_new_tokens.assert_not_called()
+    else:
+        worker._prepare_next_new_tokens.assert_called_once_with(
+            sentinel.accepted, sentinel.drafts, sentinel.batch_indices, 2, sentinel.accepted_counts
+        )

From 5a9e48804e3512e108fbba7ec217704cf2003953 Mon Sep 17 00:00:00 2001
From: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Date: Fri, 6 Nov 2026 11:12:26 -0800
Subject: [PATCH 2/3] [https://nvbugs/7028061][test] Consolidate MiniMax Eagle3
 API coverage

Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
---
 .../minimax_m3_nvfp4__sm_103__tp4/modeling.py |   7 +-
 .../minimax_m3_nvfp4__sm_103__tp8/modeling.py |   7 +-
 .../integration/test_lists/test-db/l0_cpu.yml |   1 -
 .../test_modeling_v2_eagle3_worker_api.py     | 163 ++++++++++++++----
 .../minimax_m3/test_eagle3_draft_loop.py      | 128 --------------
 5 files changed, 138 insertions(+), 168 deletions(-)
 delete mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py

diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py
index e52d8c9fa4ad..213e676ef425 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py
@@ -4295,7 +4295,12 @@ def _forward_draft_loop(
         accepted_tokens,
         pre_draft_seq_lens,
     ):
-        """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop."""
+        """Run the fused draft loop on a covered step, otherwise the in-tree loop.
+
+        Returns:
+            (next_draft_tokens, next_new_tokens): the fused loop's draft and next-input
+            buffers, or the in-tree loop's pair. A None next-input buffer is built by the caller.
+        """
         if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model):
             next_draft_tokens = self._m3_fused_draft_loop(
                 inputs,
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py
index d7cd18563211..d27d4fadc7fa 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py
@@ -3964,7 +3964,12 @@ def _forward_draft_loop(
         accepted_tokens,
         pre_draft_seq_lens,
     ):
-        """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop."""
+        """Run the fused draft loop on a covered step, otherwise the in-tree loop.
+
+        Returns:
+            (next_draft_tokens, next_new_tokens): the fused loop's draft and next-input
+            buffers, or the in-tree loop's pair. A None next-input buffer is built by the caller.
+        """
         if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model):
             next_draft_tokens = self._m3_fused_draft_loop(
                 inputs,
diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml
index df9b38804e35..2e816594cd1f 100644
--- a/tests/integration/test_lists/test-db/l0_cpu.yml
+++ b/tests/integration/test_lists/test-db/l0_cpu.yml
@@ -19,7 +19,6 @@ l0_cpu:
   # MiniMax-M3 Rubin target unit tests.
   - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py
   # MiniMax-M3 target unit tests.
-  - unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py
   - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py
   - unittest/_torch/modules/test_linear_output_slice.py::test_fp8_weight_slice_cache
   - unittest/_torch/modules/test_linear_output_slice.py::test_prepare_output_slice_rejects_unsupported_projection_before_forward
diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py
index c8625ac32358..6bf3a51b0a76 100644
--- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py
+++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py
@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""Exercise the actual SM107 dispatch methods against the current parent API.
+"""Exercise every MiniMax-M3 target against the current Eagle3 worker API.
 
-Only the Python dispatch method is loaded, so this contract runs without CUDA.
+Only the dispatch and forward token-buffer selection are loaded, without CUDA.
 The strict parent signature comes from its real source rather than a separately
 maintained stub; GPU serving checks cover the model computations themselves.
 """
@@ -11,6 +11,8 @@
 import copy
 import inspect
 from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock, sentinel
 
 import pytest
 
@@ -19,6 +21,7 @@
 ROOT = next(p for p in Path(__file__).resolve().parents if (p / "tensorrt_llm").is_dir())
 MODEL_ROOT = ROOT / "tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3"
 PARENT_PATH = ROOT / "tensorrt_llm/_torch/speculative/eagle3.py"
+MODEL_PATHS = sorted(MODEL_ROOT.glob("minimax_m3_nvfp4__*/modeling.py"))
 
 
 def _dispatch(class_node):
@@ -29,7 +32,8 @@ def _dispatch(class_node):
     )
 
 
-def _load_dispatch(tp, fused):
+def _load_dispatch(path: Path, fused: bool):
+    """Load a target dispatcher with the real parent signature and stubbed GPU work."""
     parent_tree = ast.parse(PARENT_PATH.read_text())
     parent_class = next(
         n
@@ -40,7 +44,6 @@ def _load_dispatch(tp, fused):
     parent.body = ast.parse(
         "self.parent_arguments = locals().copy()\nreturn self.parent_result"
     ).body
-    path = MODEL_ROOT / f"minimax_m3_nvfp4__sm_107__tp{tp}/modeling.py"
     target_tree = ast.parse(path.read_text())
     target_class = next(
         n for n in target_tree.body if isinstance(n, ast.ClassDef) and n.name == "M3Eagle3Worker"
@@ -74,14 +77,21 @@ def _load_dispatch(tp, fused):
     worker._m3_wide_draft_logits_ok = lambda *args: False
     worker.parent_result = (object(), object())
     worker.fused_result = object()
-    worker._m3_fused_draft_loop = lambda *args: worker.fused_result
+    worker._m3_next_new_tokens = None
+    worker.fused_next_new = object()
+
+    def fused_loop(*args: object) -> object:
+        worker._m3_next_new_tokens = worker.fused_next_new
+        return worker.fused_result
+
+    worker._m3_fused_draft_loop = fused_loop
     return worker, namespace["Parent"]._forward_draft_loop
 
 
-@pytest.mark.parametrize("tp", [2, 4, 8])
+@pytest.mark.parametrize("path", MODEL_PATHS, ids=lambda path: path.parent.name)
 @pytest.mark.parametrize("fused", [False, True], ids=["stock", "fused"])
-def test_dispatch_matches_current_parent_signature_and_return_contract(tp, fused):
-    worker, parent = _load_dispatch(tp, fused)
+def test_dispatch_matches_current_parent_signature_and_return_contract(path: Path, fused: bool):
+    worker, parent = _load_dispatch(path, fused)
     parameters = list(inspect.signature(parent).parameters)[1:]
     arguments = {name: object() for name in parameters}
     # Use every required parent argument, including the real accepted tokens and
@@ -89,14 +99,16 @@ def test_dispatch_matches_current_parent_signature_and_return_contract(tp, fused
     result = worker._forward_draft_loop(**arguments)
     assert isinstance(result, tuple) and len(result) == 2
     if fused:
-        assert result[0] is worker.fused_result and result[1] is None
+        assert result[0] is worker.fused_result
+        # Targets may return the fused buffer or leave it on the worker for the caller.
+        assert result[1] is None or result[1] is worker.fused_next_new
     else:
         assert result is worker.parent_result
         assert all(worker.parent_arguments[name] is value for name, value in arguments.items())
 
 
-@pytest.mark.parametrize("tp", [2, 4, 8])
-def test_external_shared_kv_keeps_the_parent_tensor_return_contract(tp):
+@pytest.mark.parametrize("path", [PARENT_PATH, *MODEL_PATHS], ids=lambda path: path.parent.name)
+def test_external_shared_kv_keeps_the_parent_tensor_return_contract(path: Path):
     class ExternalWorker:
         __slots__ = ()
         _uses_external_shared_target_kv = True
@@ -105,32 +117,109 @@ class ExternalWorker:
         def _forward_external_shared_target_kv_draft_loop(self, **kwargs):
             return self.tokens
 
-    paths = [PARENT_PATH, MODEL_ROOT / f"minimax_m3_nvfp4__sm_107__tp{tp}/modeling.py"]
-    for path in paths:
-        tree = ast.parse(path.read_text())
-        cls = next(
+    tree = ast.parse(path.read_text())
+    cls = next(
+        n
+        for n in tree.body
+        if isinstance(n, ast.ClassDef) and n.name in ("Eagle3OneModelWorker", "M3Eagle3Worker")
+    )
+    forward = next(
+        n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "_forward_impl"
+    )
+    branch = copy.deepcopy(
+        next(
             n
-            for n in tree.body
-            if isinstance(n, ast.ClassDef) and n.name in ("Eagle3OneModelWorker", "M3Eagle3Worker")
-        )
-        forward = next(
-            n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "_forward_impl"
+            for n in forward.body
+            if isinstance(n, ast.If)
+            and ast.unparse(n.test) == "self._uses_external_shared_target_kv"
         )
-        branch = copy.deepcopy(
-            next(
-                n
-                for n in forward.body
-                if isinstance(n, ast.If)
-                and ast.unparse(n.test) == "self._uses_external_shared_target_kv"
-            )
+    )
+    # Isolate the actual external branch: it retains a Tensor return, while
+    # the normal draft-loop dispatcher now returns a two-item tuple.
+    branch.orelse = []
+    names = {n.id for n in ast.walk(branch) if isinstance(n, ast.Name)} - {"self"}
+    namespace = {name: object() for name in names}
+    namespace.update(self=ExternalWorker(), next_new_tokens=None)
+    code = ast.Module(body=[branch], type_ignores=[])
+    exec(compile(ast.fix_missing_locations(code), str(path), "exec"), namespace)
+    assert namespace["next_draft_tokens"] is ExternalWorker.tokens
+    assert namespace["next_new_tokens"] is None
+
+
+@pytest.mark.parametrize("path", MODEL_PATHS, ids=lambda path: path.parent.name)
+@pytest.mark.parametrize("mode", ["fallback", "fallback_buffer", "fused", "external"])
+def test_forward_preserves_draft_arguments_and_token_buffers(path: Path, mode: str) -> None:
+    """Exercise the caller, dispatch, and buffer selection without unrelated forward globals."""
+    worker, _ = _load_dispatch(path, fused=mode == "fused")
+    worker._uses_external_shared_target_kv = mode == "external"
+    worker.parent_result = (
+        sentinel.drafts,
+        sentinel.next_new if mode == "fallback_buffer" else None,
+    )
+    worker.fused_result = sentinel.drafts
+    worker.fused_next_new = sentinel.next_new
+    metadata = SimpleNamespace(
+        _seq_lens_cuda=sentinel.pre_draft_seq_lens,
+        all_rank_num_tokens=sentinel.rank_tokens,
+    )
+
+    def prepare_metadata(md: SimpleNamespace) -> None:
+        md._seq_lens_cuda = sentinel.draft_seq_lens
+
+    worker._prepare_attn_metadata_for_spec_dec = Mock(side_effect=prepare_metadata)
+    worker._restore_attn_metadata_from_spec_dec = Mock()
+    worker.prepare_1st_drafter_inputs = Mock(return_value=sentinel.inputs)
+    worker.get_draft_kv_cache_manager = Mock(return_value=sentinel.cache)
+    worker._forward_external_shared_target_kv_draft_loop = Mock(return_value=sentinel.drafts)
+    worker._prepare_next_new_tokens = Mock(return_value=sentinel.next_new)
+
+    tree = ast.parse(path.read_text())
+    cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "M3Eagle3Worker")
+    forward = next(
+        n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "_forward_impl"
+    )
+    start = next(
+        i
+        for i, node in enumerate(forward.body)
+        if isinstance(node, ast.If)
+        and ast.unparse(node.test) == "self._uses_external_shared_target_kv"
+    )
+    # Keep only drafting and next-input buffer selection. Acceptance, diagnostics,
+    # model state updates, and their module globals are outside this contract.
+    buffer_selection = next(
+        node
+        for node in forward.body[start + 1 :]
+        if isinstance(node, ast.If)
+        and ast.unparse(node.test)
+        in ("next_new_tokens is None", "self._m3_next_new_tokens is not None")
+    )
+    statements = [forward.body[start], buffer_selection]
+    namespace = dict(
+        self=worker,
+        input_ids=sentinel.input_ids,
+        position_ids=Mock(),
+        hidden_states=sentinel.hidden_states,
+        attn_metadata=metadata,
+        spec_metadata=SimpleNamespace(batch_indices_cuda=sentinel.batch_indices),
+        draft_model=sentinel.draft_model,
+        accepted_tokens=sentinel.accepted,
+        num_accepted_tokens=sentinel.accepted_counts,
+        num_contexts=1,
+        num_gens=1,
+        batch_size=2,
+        resource_manager=sentinel.resources,
+        next_new_tokens=None,
+    )
+    exec(compile(ast.Module(body=statements, type_ignores=[]), str(path), "exec"), namespace)
+
+    assert namespace["next_draft_tokens"] is sentinel.drafts
+    assert namespace["next_new_tokens"] is sentinel.next_new
+    if mode.startswith("fallback"):
+        assert worker.parent_arguments["accepted_tokens"] is sentinel.accepted
+        assert worker.parent_arguments["pre_draft_seq_lens"] is sentinel.pre_draft_seq_lens
+    if mode in ("fallback_buffer", "fused"):
+        worker._prepare_next_new_tokens.assert_not_called()
+    else:
+        worker._prepare_next_new_tokens.assert_called_once_with(
+            sentinel.accepted, sentinel.drafts, sentinel.batch_indices, 2, sentinel.accepted_counts
         )
-        # Isolate the actual external branch: it retains a Tensor return, while
-        # the normal draft-loop dispatcher now returns a two-item tuple.
-        branch.orelse = []
-        names = {n.id for n in ast.walk(branch) if isinstance(n, ast.Name)} - {"self"}
-        namespace = {name: object() for name in names}
-        namespace.update(self=ExternalWorker(), next_new_tokens=None)
-        code = ast.Module(body=[branch], type_ignores=[])
-        exec(compile(ast.fix_missing_locations(code), str(path), "exec"), namespace)
-        assert namespace["next_draft_tokens"] is ExternalWorker.tokens
-        assert namespace["next_new_tokens"] is None
diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py
deleted file mode 100644
index 518618ab2ce9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_eagle3_draft_loop.py
+++ /dev/null
@@ -1,128 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Exercise the MiniMax worker's Python dispatch without importing GPU kernels."""
-
-import ast
-from pathlib import Path
-from types import SimpleNamespace
-from unittest.mock import Mock, sentinel
-
-import pytest
-
-pytestmark = pytest.mark.cpu_only
-
-_TORCH_ROOT = Path(__file__).resolve().parents[6] / "tensorrt_llm" / "_torch"
-
-
-def _load_worker(tp: int) -> type:
-    # Execute the actual methods, including the base dispatcher and its signature.
-    # Keeping the class definitions preserves zero-argument super() in the override.
-    namespace = {}
-    sources = (
-        (_TORCH_ROOT / "speculative/eagle3.py", "Eagle3OneModelWorker", {"_forward_draft_loop"}),
-        (
-            _TORCH_ROOT
-            / f"_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp{tp}/modeling.py",
-            "M3Eagle3Worker",
-            {"_forward_impl", "_forward_draft_loop"},
-        ),
-    )
-    for path, name, methods in sources:
-        cls = next(
-            node
-            for node in ast.parse(path.read_text()).body
-            if isinstance(node, ast.ClassDef) and node.name == name
-        )
-        cls.body = [
-            node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in methods
-        ]
-        assert len(cls.body) == len(methods)
-        if name == "Eagle3OneModelWorker":
-            cls.bases = []
-        module = ast.Module(body=[cls], type_ignores=[])
-        exec(compile(module, str(path), "exec"), namespace)
-    return namespace["M3Eagle3Worker"]
-
-
-@pytest.mark.parametrize("tp", [4, 8])
-@pytest.mark.parametrize("path", ["fallback", "fallback_buffer", "fused", "external"])
-def test_eagle3_draft_loop_contract(tp: int, path: str) -> None:
-    worker = _load_worker(tp)()
-    worker._is_mamba_hybrid_cache = False
-    worker._auxiliary_state_handlers = []
-    worker.sa_enhancer = None
-    worker._uses_external_shared_target_kv = path == "external"
-    worker._execute_guided_decoder_if_present = Mock()
-    worker.sample_and_accept_draft_tokens = Mock(
-        return_value=(sentinel.accepted, sentinel.accepted_counts)
-    )
-    worker.draft_only_eligible = Mock(return_value=False)
-    metadata = SimpleNamespace(
-        num_seqs=2,
-        num_contexts=1,
-        _seq_lens_cuda=sentinel.pre_draft_seq_lens,
-        all_rank_num_tokens=sentinel.rank_tokens,
-    )
-    spec = SimpleNamespace(
-        runtime_draft_len=3,
-        spec_resource_manager=SimpleNamespace(),
-        batch_indices_cuda=sentinel.batch_indices,
-    )
-
-    def prepare_metadata(md: SimpleNamespace) -> None:
-        # The caller must retain the sequence lengths from before preparation.
-        md._seq_lens_cuda = sentinel.draft_seq_lens
-
-    worker._prepare_attn_metadata_for_spec_dec = Mock(side_effect=prepare_metadata)
-    worker._restore_attn_metadata_from_spec_dec = Mock()
-    worker.prepare_1st_drafter_inputs = Mock(return_value=sentinel.inputs)
-    worker.get_draft_kv_cache_manager = Mock(return_value=sentinel.cache)
-    worker._m3_fused_loop_ok = Mock(return_value=path == "fused")
-    prebuilt = sentinel.next_new if path == "fallback_buffer" else None
-    worker._forward_linear_draft_loop = Mock(return_value=(sentinel.drafts, prebuilt))
-
-    def fused_loop(*args: object) -> object:
-        worker._m3_next_new_tokens = sentinel.next_new
-        return sentinel.drafts
-
-    worker._m3_fused_draft_loop = Mock(side_effect=fused_loop)
-    worker._forward_external_shared_target_kv_draft_loop = Mock(return_value=sentinel.drafts)
-    worker._prepare_next_new_tokens = Mock(return_value=sentinel.next_new)
-
-    outputs = worker._forward_impl(
-        sentinel.input_ids,
-        Mock(),
-        sentinel.hidden_states,
-        sentinel.logits,
-        metadata,
-        spec,
-        sentinel.draft_model,
-        sentinel.resources,
-    )
-
-    assert outputs["next_draft_tokens"] is sentinel.drafts
-    assert outputs["next_new_tokens"] is sentinel.next_new
-    assert outputs["new_tokens"] is sentinel.accepted
-    assert metadata.use_spec_decoding
-    if path.startswith("fallback"):
-        worker._forward_linear_draft_loop.assert_called_once_with(
-            sentinel.inputs,
-            metadata,
-            spec,
-            sentinel.draft_model,
-            sentinel.cache,
-            1,
-            2,
-            sentinel.accepted_counts,
-            sentinel.rank_tokens,
-            sentinel.accepted,
-            sentinel.pre_draft_seq_lens,
-        )
-    else:
-        worker._forward_linear_draft_loop.assert_not_called()
-    if path in ("fallback_buffer", "fused"):
-        worker._prepare_next_new_tokens.assert_not_called()
-    else:
-        worker._prepare_next_new_tokens.assert_called_once_with(
-            sentinel.accepted, sentinel.drafts, sentinel.batch_indices, 2, sentinel.accepted_counts
-        )

From 51155046e8077b9f7c790fcd0cb153a02d10a6cd Mon Sep 17 00:00:00 2001
From: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
Date: Mon, 9 Nov 2026 10:16:39 -0800
Subject: [PATCH 3/3] Strengthen MiniMax worker API regression assertions

Signed-off-by: Maxim Gluhovskoi <mgluhovskoi@nvidia.com>
---
 .../minimax_m3/test_modeling_v2_eagle3_worker_api.py       | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_modeling_v2_eagle3_worker_api.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_modeling_v2_eagle3_worker_api.py
index 6bf3a51b0a76..71473270866a 100644
--- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_modeling_v2_eagle3_worker_api.py
+++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_modeling_v2_eagle3_worker_api.py
@@ -22,6 +22,7 @@
 MODEL_ROOT = ROOT / "tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3"
 PARENT_PATH = ROOT / "tensorrt_llm/_torch/speculative/eagle3.py"
 MODEL_PATHS = sorted(MODEL_ROOT.glob("minimax_m3_nvfp4__*/modeling.py"))
+assert MODEL_PATHS, f"No MiniMax-M3 NVFP4 modeling.py files found under {MODEL_ROOT}"
 
 
 def _dispatch(class_node):
@@ -91,6 +92,7 @@ def fused_loop(*args: object) -> object:
 @pytest.mark.parametrize("path", MODEL_PATHS, ids=lambda path: path.parent.name)
 @pytest.mark.parametrize("fused", [False, True], ids=["stock", "fused"])
 def test_dispatch_matches_current_parent_signature_and_return_contract(path: Path, fused: bool):
+    """Verify argument forwarding and token-buffer handoff for stock and fused dispatch."""
     worker, parent = _load_dispatch(path, fused)
     parameters = list(inspect.signature(parent).parameters)[1:]
     arguments = {name: object() for name in parameters}
@@ -102,6 +104,7 @@ def test_dispatch_matches_current_parent_signature_and_return_contract(path: Pat
         assert result[0] is worker.fused_result
         # Targets may return the fused buffer or leave it on the worker for the caller.
         assert result[1] is None or result[1] is worker.fused_next_new
+        assert worker._m3_next_new_tokens is worker.fused_next_new
     else:
         assert result is worker.parent_result
         assert all(worker.parent_arguments[name] is value for name, value in arguments.items())
@@ -109,6 +112,8 @@ def test_dispatch_matches_current_parent_signature_and_return_contract(path: Pat
 
 @pytest.mark.parametrize("path", [PARENT_PATH, *MODEL_PATHS], ids=lambda path: path.parent.name)
 def test_external_shared_kv_keeps_the_parent_tensor_return_contract(path: Path):
+    """Verify external shared-KV drafting returns tokens without a next-input buffer."""
+
     class ExternalWorker:
         __slots__ = ()
         _uses_external_shared_target_kv = True
@@ -215,6 +220,8 @@ def prepare_metadata(md: SimpleNamespace) -> None:
     assert namespace["next_draft_tokens"] is sentinel.drafts
     assert namespace["next_new_tokens"] is sentinel.next_new
     if mode.startswith("fallback"):
+        assert worker.parent_arguments["inputs"] is sentinel.inputs
+        assert worker.parent_arguments["draft_kv_cache_manager"] is sentinel.cache
         assert worker.parent_arguments["accepted_tokens"] is sentinel.accepted
         assert worker.parent_arguments["pre_draft_seq_lens"] is sentinel.pre_draft_seq_lens
     if mode in ("fallback_buffer", "fused"):