Skip to content

[https://nvbugs/7028061][fix] Align MiniMax M3 Eagle3 draft-loop API #21815

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mgluhovskoi wants to merge 4 commits into
base: main
Choose a base branch
from
+171 −53
Open

[https://nvbugs/7028061][fix] Align MiniMax M3 Eagle3 draft-loop API #21815

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
db8b81d
[https://nvbugs/7028061][fix] Align MiniMax M3 Eagle3 draft-loop API
Mgluhovskoi Nov 6, 2026
5a9e488
[https://nvbugs/7028061][test] Consolidate MiniMax Eagle3 API coverage
Mgluhovskoi Nov 6, 2026
d69e4ff
Merge main and resolve MiniMax test waiver conflict
Mgluhovskoi Nov 9, 2026
5115504
Strengthen MiniMax worker API regression assertions
Mgluhovskoi Nov 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 19 additions & 7 deletions ...rch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)

Expand All @@ -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,
Expand All @@ -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)
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -4302,6 +4311,7 @@ def _forward_draft_loop(
batch_size,
num_accepted_tokens,
)
return next_draft_tokens, self._m3_next_new_tokens
Comment thread
Mgluhovskoi marked this conversation as resolved.
Resolved
return super()._forward_draft_loop(
inputs,
attn_metadata,
Expand All @@ -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(
Expand Down
26 changes: 19 additions & 7 deletions ...rch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)

Expand All @@ -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,
Expand All @@ -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)
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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(
Expand Down
2 changes: 0 additions & 2 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)

@sunnyqgg sunnyqgg Nov 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Required] The two waive lines above are removed before any run has shown these tests pass with the fix.

If either case still misses its GSM8K or acceptance floor, GB300 post-merge goes red on main as soon as this lands.

The stage-list run you started (PR_Github #83129 → L0_MergeRequest_PR #68680) was still in progress when I looked. Because this is a PR pipeline and the PR has no ci: full pre-merge approved label, its SBSA multi-GPU stages will be marked [Test-SBSA-Multi-GPU] Blocked instead of running these tests (jenkins/L0_MergeRequest.groovy:3056-3064).

Please:

  • ask a trt-llm-ci-approvers member for the label, then re-run /bot run --stage-list "GB300-4_GPUs-PyTorch-Post-Merge-*,GB300-8_GPUs-2_Nodes-PyTorch-SingleNvlinkDomain-Post-Merge-1";
  • link the passing run for both cases in the description before merge, or keep the failing case's waive line and un-waive it in a follow-up;
  • also run a full /bot run: #83128 was aborted, so the new l0_cpu entry has not run in CI yet.

accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664)
Expand Down
Loading
Loading