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..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 @@ -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,17 @@ 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.""" + """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): - return self._m3_fused_draft_loop( + next_draft_tokens = self._m3_fused_draft_loop( inputs, attn_metadata, spec_metadata, @@ -4302,6 +4311,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 +4324,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..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 @@ -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,17 @@ 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.""" + """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): - return self._m3_fused_draft_loop( + next_draft_tokens = self._m3_fused_draft_loop( inputs, attn_metadata, spec_metadata, @@ -3971,6 +3980,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 +3993,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/waives.txt b/tests/integration/test_lists/waives.txt index 29181236dfc4..ca3ee102abfd 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -1,5 +1,3 @@ -accuracy/modeling_v2/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_accuracy.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_eagle3_gsm8k_and_acceptance SKIP (https://nvbugs/7028061) -accuracy/modeling_v2/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_accuracy.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_eagle3_gsm8k_and_acceptance SKIP (https://nvbugs/7028061) accuracy/test_disaggregated_serving.py::TestDeepSeekR1::test_kv_cache_v2_nixl_python[cache_mgr_v2] SKIP (https://nvbugs/7066121) 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) 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 c8625ac32358..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 @@ -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,8 @@ 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")) +assert MODEL_PATHS, f"No MiniMax-M3 NVFP4 modeling.py files found under {MODEL_ROOT}" def _dispatch(class_node): @@ -29,7 +33,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 +45,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 +78,22 @@ 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): + """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} # Use every required parent argument, including the real accepted tokens and @@ -89,14 +101,19 @@ 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 + 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()) -@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): + """Verify external shared-KV drafting returns tokens without a next-input buffer.""" + class ExternalWorker: __slots__ = () _uses_external_shared_target_kv = True @@ -105,32 +122,111 @@ 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") + for n in forward.body + if isinstance(n, ast.If) + and ast.unparse(n.test) == "self._uses_external_shared_target_kv" ) - 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 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["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"): + 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