Skip to content

[None][feat] MiniMax-M3 NVFP4 low latency decode megakernel for Rubin (sm_107) TP2 / TP4 / TP8 in modeling_v2 - #20623

Merged
peihu-nv merged 9 commits into
NVIDIA:mainfrom
peihu-nv:peihengh/m3-vr-sm107-decode
Nov 4, 2026
Merged

peihu-nv merged 9 commits into
NVIDIA:mainfrom
peihu-nv:peihengh/m3-vr-sm107-decode

Conversation

@peihu-nv

@peihu-nv peihu-nv commented Oct 19, 2026 •

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Description

Depends on #20306 (flat staircase layout of the sm_103 targets). This PR adds only the last commit; please review that commit.

This PR brings the MiniMax-M3 NVFP4 decode targets to Rubin (sm_107): minimax_m3_nvfp4__sm_107__tp2, __tp4 and __tp8. They follow the flat layout of the sm_103 targets: one modeling.py per target, a weights.py weight table, and an ops/ catalog entry with certified cells and a test for every kernel the model calls. The router picks a target only for the matching checkpoint, SM and width, and modeling_v2: require makes it mandatory.

Performance

On Rubin (VR200), MiniMax-M3 NVFP4 AgentX with Eagle3, against the stock decode path (modeling_v2 off):

  • End to end: about 2.8× interactivity vs the best stock interactivity point (TP8 disaggregated, 1,112 tok/s/user, vs stock TP4 aggregated, 401).
  • Decode step: about 1.5× (TP2), 2.0× (TP4) and 2.4–2.7× (TP8) faster at batch 1.

sm_107 enablement

  • Catalog and routing: Arch.SM_107, and routing entries for (10, 7) at 2, 4 and 8 ranks, with the same drafter and runtime declines as sm_103.
  • Compiled kernels: minimaxM3Sm107Tp{2,4,8} with their trtllm::minimax_m3_sm107* thop ops, plus two shared GEN kernels: the TF32 tensor-core router GEMV and the spec-tail draft-input step.
  • CuTe DSL kernels: the decode MoE engines (1–8 rows, plus a 9–13-row engine at TP4), the MX GEMV, and the fused embedding + first RMSNorm.
  • Decode runtime: the step-input pack, skipped unchanged uploads, the spec step fold and the MSA decode metadata reuse. Each piece is inert until a target's RuntimeHooks turns it on, so other models are unchanged.
  • PrimTS (vendored FlashInfer task-scheduled attention): the dense layers and the drafter run its split-KV decode. The vendor patch scopes PrimTS's global WorkTileInfo rebind to its own compiles, so other CuTe DSL kernels (e.g. FA4) are unaffected.
  • KV fallback: the sm_107 NVFP4 → FP8 KV fallback skips MiniMax-M3 MSA, which keeps NVFP4 only on its sparse layers.
  • FMHA fallback: on SM107 the fallback FMHA serves paged prefill at the TRT-LLM Gen extra page sizes (the 128-token KV pages the MiniMax-M3 caches opt into) instead of refusing it; SM100 / SM103 keep the refusal.
  • Runtime items (target opt-in): the pinned FlashInfer CuTe DSL MXFP8 tactics above 32 rows on SM107, the post-warmup late-compile detector, and first-content emission (a request's first content after two iterations) are armed only when the model's target opts in through its RuntimeHooks; the MiniMax-M3 sm_107 targets do, and other models are unchanged. The AllReduce tuner leaves NCCL_SYMMETRIC out of its candidates only with TRTLLM_AR_TUNE_NCCL_SYMMETRIC=0 (opt-in; the default keeps the old candidates).

The targets

All three share one structure; they differ in rank geometry and row ranges.

TP2 TP4 TP8
Decode step bound 192 rows (fast paths to 128) 48 rows 8 rows
MoE engine 1–8 rows (Triton finalize from 49) 1–8 + 9–13 rows 1–8 rows
Sparse decode attention m3_sparse_decode_lc m3_sparse_decode_lc m3_sparse_decode_lc8
  • Boundaries: multicast all-reduce fused with the next RMSNorm and with the MoE finalize.
  • GEMMs: an MX GEMV for low-row MXFP8 projections and drafter GEMMs; at TP2 the LM head is split across the two dies.
  • Speculation (Eagle3, three drafts): a fused draft step and the target's own sampled accept. At TP4 the sampled verify runs two rounds (single-round accept off), matching the deployed configuration.
  • No tuning knobs: the targets read no tuning knobs from the environment; the only variable is the opt-in runtime-value dump diagnostic (TRTLLM_MODELING_V2_VALUE_DUMP), as in the sm_103 targets. Every fast path is on by default and has one module constant that turns it off. Capacity bounds carry a # capacity: annotation.

Aggregated (wide) decode at TP2

  • Rows: the TP2 target serves aggregated decode steps of up to 48 requests (1..192 rows, each request verifying 1 + 3 drafts). Fast-path cells cover up to 128 rows; 129..192-row steps take the stock-exact paths.
  • Wide cells: steps above 32 rows run the generic layer with 48-row multicast boundaries and MC finalize, the tensor-core router (9..48 rows), the Triton MoE finalize (49+ rows) and the 33..48-row GEMVs.
  • Mixed steps: steps with context rows run under breakable prefill CUDA graphs, with the drafter's attention in an eager break.

Numerics. No lossy steps. Most kernels are bitwise equal to the stock op they replace; the rest reduce in a fixed order, are deterministic in eager and graph mode, and stay within the stock kernel's error band. One known divergence: the TP4 drafter all-reduce + RMSNorm always sums like ONESHOT, so under a forced NCCL strategy its rounding differs from stock NCCL. Against the earlier wide TP2 reference, target tokens are identical: at 12+ concurrent requests under Eagle3 some draft tokens can resolve near-ties differently (target logits bitwise), and without speculation (not a production configuration) decode steps of 9..48 rows are bitwise on the first such step and differ after it, with identical output tokens. Steps of 1..8 rows are bitwise for prompts up to about 12K tokens, and at TP8 also for a single 83K-token request (prefill and decode).

Not in this PR: batch 4/8 decode at TP4 / TP8 (follow-up). The prefill (context) optimizations are in #20629.

Known behaviours: The pinned SM107 MXFP8 GEMM tactics are armed process-wide by an executor whose model opts in, so another model in the same process at the same GEMM shape also uses them; its outputs are bitwise unchanged (every pinned tactic equals the default), and the LLM API runs one executor per worker process.

Test Coverage

  • Bitwise A/B vs the pre-restructure sm_107 targets (TP2/TP4/TP8, no speculation and Eagle3, every rank, prompts up to about 12K tokens; at TP8 also one 83K-token request): identical tokens and per-step output hashes. The only exception is the unread tail of new_tokens in Eagle3 context steps. Re-run on the final commit at TP2 and TP4 (no speculation and Eagle3): bitwise at every step of 1..8 rows; at TP2 the first difference is a 45-row context step, which this PR moves to the wide cells.
  • Wide TP2 vs the aggregated (wide) reference target at 9..128 rows, Eagle3 and no speculation: target tokens identical in every phase and logits bitwise on every recorded step under Eagle3; the two deltas above are the only differences.
  • FMHA fallback at 128-token pages on SM107: a GPU oracle (fresh contexts beside generation rows, cached prefixes intact, K/V stored at their slots) and a CPU gate, all cases in band.
  • Unit tests on Rubin, one session per file: TP2 1449, TP4 1659, TP8 1415 tests, all passing. Each target's drafter context cell is a strict expected failure because the SM107 non-paged context FMHA does not apply RoPE (shared runtime issue, tracked separately).
  • Decode runtime tests (executor/, speculative/, attention): 255 pass, 1 skipped.
  • Delivery and capacity checkers: TP2 99/99, TP4 107/107, TP8 97/97.
  • Accuracy at TP2 vs the pre-restructure targets: GSM8K (1319) McNemar p = 0.30 greedy, p = 1.00 sampled; the Eagle3 natural acceptance length over 264 long-context requests is identical. At TP4 and TP8, GSM8K (1319) vs the pre-restructure targets: TP4 p = 1.00 greedy, p = 1.00 sampled; TP8 p = 0.38 greedy, p = 0.29 sampled.
  • Decode perf vs the pre-restructure targets, same node, both run orders: TP2, TP4 and TP8 within noise (TP8 compared at equal decode work: the synthetic runs' generated text, and so the number of distinct routed experts per step, varies from server start to server start).

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

@peihu-nv peihu-nv changed the title [None][feat] MiniMax-M3 NVFP4 decode targets for Rubin (sm_107) TP2 / TP4 / TP8 in modeling_v2 [None][feat] MiniMax-M3 NVFP4 low latency decode megakernel for Rubin (sm_107) TP2 / TP4 / TP8 in modeling_v2 Oct 19, 2026
@peihu-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@peihu-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79103 [ run ] triggered by Bot. Commit: 5d24e81 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79103 [ run ] completed with state FAILURE. Commit: 5d24e81
/LLM/main/L0_MergeRequest_PR pipeline #65162 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@peihu-nv
peihu-nv force-pushed the peihengh/m3-vr-sm107-decode branch from 5d24e81 to 4d16da2 Compare October 20, 2026 12:17
@peihu-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@peihu-nv
peihu-nv force-pushed the peihengh/m3-vr-sm107-decode branch from 4d16da2 to 7333ac8 Compare October 20, 2026 15:21
@peihu-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79291 [ run ] triggered by Bot. Commit: 7333ac8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79291 [ run ] completed with state SUCCESS. Commit: 7333ac8
/LLM/main/L0_MergeRequest_PR pipeline #65314 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@peihu-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79385 [ run ] triggered by Bot. Commit: 8b49063 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79385 [ run ] completed with state SUCCESS. Commit: 8b49063
/LLM/main/L0_MergeRequest_PR pipeline #65380 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@peihu-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79616 [ run ] triggered by Bot. Commit: 8b49063 Link to invocation

@peihu-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

peihu-nv added a commit to peihu-nv/TensorRT-LLM that referenced this pull request Oct 21, 2026
Brings in NVIDIA#20623 at ffa76ab (the NVIDIA#20306 head 300c1c1 merged in, sm107 target kernels under ops). No change to the context path.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79695 [ run ] triggered by Bot. Commit: ffa76ab Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #79616 [ run ] completed with state ABORTED. Commit: 8b49063

Link to invocation

@peihu-nv
peihu-nv marked this pull request as ready for review October 21, 2026 23:49
@peihu-nv
peihu-nv requested review from a team as code owners October 21, 2026 23:49
@peihu-nv

peihu-nv commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #82111 [ run ] triggered by Bot. Commit: 23c0805 Link to invocation

@peihu-nv
peihu-nv enabled auto-merge (squash) November 3, 2026 00:29
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #82111 [ run ] completed with state SUCCESS. Commit: 23c0805
/LLM/main/L0_MergeRequest_PR pipeline #67744 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Either:

  • Wait for the PR to be fully approved — the label is added automatically once approval is complete. Having unresolved open comments is fine, or
  • If needed, ask a member of NVIDIA/trt-llm-ci-approvers to add the label manually.
    Then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@SimengLiu-nv SimengLiu-nv left a comment

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.

LGTM.

@peihu-nv

peihu-nv commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #82335 [ run ] triggered by Bot. Commit: 23c0805 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #82335 [ run ] completed with state FAILURE. Commit: 23c0805
/LLM/main/L0_MergeRequest_PR pipeline #67955 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Retain upstream single KV-metadata invalidation after overlap corrections,
including the model-specific fused offset path. Align the existing CPU
regression with the upstream behavior for both legacy hook settings.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
@peihu-nv

peihu-nv commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #82359 [ run ] triggered by Bot. Commit: 45b6918 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #82359 [ run ] completed with state SUCCESS. Commit: 45b6918
/LLM/main/L0_MergeRequest_PR pipeline #67979 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@peihu-nv

peihu-nv commented Nov 4, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #82417 [ run ] triggered by Bot. Commit: 45b6918 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #82417 [ run ] completed with state SUCCESS. Commit: 45b6918
/LLM/main/L0_MergeRequest_PR pipeline #68034 completed with status: 'SUCCESS'

CI Report

Link to invocation

@peihu-nv
peihu-nv merged commit 5074545 into NVIDIA:main Nov 4, 2026
9 checks passed
WeiHaocheng added a commit to WeiHaocheng/TensorRT-LLM that referenced this pull request Nov 5, 2026
The sm_107 targets (NVIDIA#20623) issued their step's copies through
step_inputs / h2d_cache with RuntimeHooks switches. They now declare a
CopyPlan like the sm_103 targets.

* CopyPlan.parallel_scatter: an order-free segment is scattered with one
  program per copy instead of the serial one-CTA scatter.
* CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every
  decode graph is captured with a prologue that issues a pinned window, and
  the final flush before a replay only fills it (early_copy: uploaded on a
  side copy stream; prologue_op: the model's one-launch prologue;
  step_epoch: a device counter bumped once per replay, read through
  device_copy.step_epoch()). Barrier segments stay eager: only the step's
  last segment rides the window, and only when it fits and is order-free.
  No prologue is captured when the graph runner has pre-replay hooks; a
  replay outside the copier's context gets an empty window.
* CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True):
  a changed merged 2-D upload stages only each row's changed span (MSA's
  block table on a page crossing).
* M3SpecCopySite: the spec fold's arming uploads, owned by
  DeviceCopyBuffers and merged by the sm_107 plans together with the
  speculative RNG sites.
* The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode
  shortcuts and the sampler's store copies are those components' own.
* MSA's in-forward paths are its fixed behaviour: the sub-page table is
  always expanded inside the forward, the block table always keeps its
  owner-side record, and the out_cache_loc tail is always filled in the
  forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and
  TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the
  eager tail fill are removed; nothing reads a merged destination before
  the replay, so no M3 plan needs a flush_after barrier.
* The step prologue ops read device_copy's graph scatter as reference.
* Ported tests: test_step_inputs' parallel-scatter and in-graph tests
  (test_device_copy_in_graph.py), test_h2d_changed_span
  (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests
  (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph
  runner's prologue test, the MSA skip-routing scan, the sm_107 contract
  scans and step-prologue tests; test_copy_plans.py checks every M3
  target's plan.

Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
WeiHaocheng added a commit to WeiHaocheng/TensorRT-LLM that referenced this pull request Nov 9, 2026
The sm_107 targets (NVIDIA#20623) issued their step's copies through
step_inputs / h2d_cache with RuntimeHooks switches. They now declare a
CopyPlan like the sm_103 targets.

* CopyPlan.parallel_scatter: an order-free segment is scattered with one
  program per copy instead of the serial one-CTA scatter.
* CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every
  decode graph is captured with a prologue that issues a pinned window, and
  the final flush before a replay only fills it (early_copy: uploaded on a
  side copy stream; prologue_op: the model's one-launch prologue;
  step_epoch: a device counter bumped once per replay, read through
  device_copy.step_epoch()). Barrier segments stay eager: only the step's
  last segment rides the window, and only when it fits and is order-free.
  No prologue is captured when the graph runner has pre-replay hooks; a
  replay outside the copier's context gets an empty window.
* CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True):
  a changed merged 2-D upload stages only each row's changed span (MSA's
  block table on a page crossing).
* M3SpecCopySite: the spec fold's arming uploads, owned by
  DeviceCopyBuffers and merged by the sm_107 plans together with the
  speculative RNG sites.
* The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode
  shortcuts and the sampler's store copies are those components' own.
* MSA's in-forward paths are its fixed behaviour: the sub-page table is
  always expanded inside the forward, the block table always keeps its
  owner-side record, and the out_cache_loc tail is always filled in the
  forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and
  TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the
  eager tail fill are removed; nothing reads a merged destination before
  the replay, so no M3 plan needs a flush_after barrier.
* The step prologue ops read device_copy's graph scatter as reference.
* Ported tests: test_step_inputs' parallel-scatter and in-graph tests
  (test_device_copy_in_graph.py), test_h2d_changed_span
  (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests
  (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph
  runner's prologue test, the MSA skip-routing scan, the sm_107 contract
  scans and step-prologue tests; test_copy_plans.py checks every M3
  target's plan.

Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
WeiHaocheng added a commit to WeiHaocheng/TensorRT-LLM that referenced this pull request Nov 11, 2026
The sm_107 targets (NVIDIA#20623) issued their step's copies through
step_inputs / h2d_cache with RuntimeHooks switches. They now declare a
CopyPlan like the sm_103 targets.

* CopyPlan.parallel_scatter: an order-free segment is scattered with one
  program per copy instead of the serial one-CTA scatter.
* CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every
  decode graph is captured with a prologue that issues a pinned window, and
  the final flush before a replay only fills it (early_copy: uploaded on a
  side copy stream; prologue_op: the model's one-launch prologue;
  step_epoch: a device counter bumped once per replay, read through
  device_copy.step_epoch()). Barrier segments stay eager: only the step's
  last segment rides the window, and only when it fits and is order-free.
  No prologue is captured when the graph runner has pre-replay hooks; a
  replay outside the copier's context gets an empty window.
* CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True):
  a changed merged 2-D upload stages only each row's changed span (MSA's
  block table on a page crossing).
* M3SpecCopySite: the spec fold's arming uploads, owned by
  DeviceCopyBuffers and merged by the sm_107 plans together with the
  speculative RNG sites.
* The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode
  shortcuts and the sampler's store copies are those components' own.
* MSA's in-forward paths are its fixed behaviour: the sub-page table is
  always expanded inside the forward, the block table always keeps its
  owner-side record, and the out_cache_loc tail is always filled in the
  forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and
  TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the
  eager tail fill are removed; nothing reads a merged destination before
  the replay, so no M3 plan needs a flush_after barrier.
* The step prologue ops read device_copy's graph scatter as reference.
* Ported tests: test_step_inputs' parallel-scatter and in-graph tests
  (test_device_copy_in_graph.py), test_h2d_changed_span
  (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests
  (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph
  runner's prologue test, the MSA skip-routing scan, the sm_107 contract
  scans and step-prologue tests; test_copy_plans.py checks every M3
  target's plan.

Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
WeiHaocheng added a commit to WeiHaocheng/TensorRT-LLM that referenced this pull request Nov 11, 2026
The sm_107 targets (NVIDIA#20623) issued their step's copies through
step_inputs / h2d_cache with RuntimeHooks switches. They now declare a
CopyPlan like the sm_103 targets.

* CopyPlan.parallel_scatter: an order-free segment is scattered with one
  program per copy instead of the serial one-CTA scatter.
* CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every
  decode graph is captured with a prologue that issues a pinned window, and
  the final flush before a replay only fills it (early_copy: uploaded on a
  side copy stream; prologue_op: the model's one-launch prologue;
  step_epoch: a device counter bumped once per replay, read through
  device_copy.step_epoch()). Barrier segments stay eager: only the step's
  last segment rides the window, and only when it fits and is order-free.
  No prologue is captured when the graph runner has pre-replay hooks; a
  replay outside the copier's context gets an empty window.
* CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True):
  a changed merged 2-D upload stages only each row's changed span (MSA's
  block table on a page crossing).
* M3SpecCopySite: the spec fold's arming uploads, owned by
  DeviceCopyBuffers and merged by the sm_107 plans together with the
  speculative RNG sites.
* The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode
  shortcuts and the sampler's store copies are those components' own.
* MSA's in-forward paths are its fixed behaviour: the sub-page table is
  always expanded inside the forward, the block table always keeps its
  owner-side record, and the out_cache_loc tail is always filled in the
  forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and
  TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the
  eager tail fill are removed; nothing reads a merged destination before
  the replay, so no M3 plan needs a flush_after barrier.
* The step prologue ops read device_copy's graph scatter as reference.
* Ported tests: test_step_inputs' parallel-scatter and in-graph tests
  (test_device_copy_in_graph.py), test_h2d_changed_span
  (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests
  (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph
  runner's prologue test, the MSA skip-routing scan, the sm_107 contract
  scans and step-prologue tests; test_copy_plans.py checks every M3
  target's plan.

Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
hanjingtian pushed a commit to WeiHaocheng/TensorRT-LLM that referenced this pull request Nov 11, 2026
The sm_107 targets (NVIDIA#20623) issued their step's copies through
step_inputs / h2d_cache with RuntimeHooks switches. They now declare a
CopyPlan like the sm_103 targets.

* CopyPlan.parallel_scatter: an order-free segment is scattered with one
  program per copy instead of the serial one-CTA scatter.
* CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every
  decode graph is captured with a prologue that issues a pinned window, and
  the final flush before a replay only fills it (early_copy: uploaded on a
  side copy stream; prologue_op: the model's one-launch prologue;
  step_epoch: a device counter bumped once per replay, read through
  device_copy.step_epoch()). Barrier segments stay eager: only the step's
  last segment rides the window, and only when it fits and is order-free.
  No prologue is captured when the graph runner has pre-replay hooks; a
  replay outside the copier's context gets an empty window.
* CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True):
  a changed merged 2-D upload stages only each row's changed span (MSA's
  block table on a page crossing).
* M3SpecCopySite: the spec fold's arming uploads, owned by
  DeviceCopyBuffers and merged by the sm_107 plans together with the
  speculative RNG sites.
* The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode
  shortcuts and the sampler's store copies are those components' own.
* MSA's in-forward paths are its fixed behaviour: the sub-page table is
  always expanded inside the forward, the block table always keeps its
  owner-side record, and the out_cache_loc tail is always filled in the
  forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and
  TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the
  eager tail fill are removed; nothing reads a merged destination before
  the replay, so no M3 plan needs a flush_after barrier.
* The step prologue ops read device_copy's graph scatter as reference.
* Ported tests: test_step_inputs' parallel-scatter and in-graph tests
  (test_device_copy_in_graph.py), test_h2d_changed_span
  (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests
  (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph
  runner's prologue test, the MSA skip-routing scan, the sm_107 contract
  scans and step-prologue tests; test_copy_plans.py checks every M3
  target's plan.

Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
hanjingtian pushed a commit to tensorrt-cicd/TensorRT-LLM that referenced this pull request Nov 11, 2026
…ngine's environment overrides

LLM writes env_overrides into the calling process's os.environ and does
not restore them, so the modeling_v2 accuracy tests open each engine
inside scoped_env_overrides, and test_modeling_v2_env_scope.py checks
that every LLM(...) call under accuracy/modeling_v2 is built that way.

TestModelingV2Glm52Nvfp4Sm103Tp8 passed env_overrides to an engine built
outside such a scope, so TRTLLM_MODELING_V2_VALIDATE and
TRTLLM_DSA_INDEXER_BF16 stayed set in the pytest process after the test,
and the check fails on
glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_accuracy.py. The file-level
waiver of this bug hid that failure; it was added when the check failed
on test_minimax_m3.py, which NVIDIA#20623 has since scoped.

Open the engine inside scoped_env_overrides and pass it the yielded
mapping, as _env.py documents, and remove the waiver.

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
WeiHaocheng added a commit to WeiHaocheng/TensorRT-LLM that referenced this pull request Nov 12, 2026
The sm_107 targets (NVIDIA#20623) issued their step's copies through
step_inputs / h2d_cache with RuntimeHooks switches. They now declare a
CopyPlan like the sm_103 targets.

* CopyPlan.parallel_scatter: an order-free segment is scattered with one
  program per copy instead of the serial one-CTA scatter.
* CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every
  decode graph is captured with a prologue that issues a pinned window, and
  the final flush before a replay only fills it (early_copy: uploaded on a
  side copy stream; prologue_op: the model's one-launch prologue;
  step_epoch: a device counter bumped once per replay, read through
  device_copy.step_epoch()). Barrier segments stay eager: only the step's
  last segment rides the window, and only when it fits and is order-free.
  No prologue is captured when the graph runner has pre-replay hooks; a
  replay outside the copier's context gets an empty window.
* CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True):
  a changed merged 2-D upload stages only each row's changed span (MSA's
  block table on a page crossing).
* M3SpecCopySite: the spec fold's arming uploads, owned by
  DeviceCopyBuffers and merged by the sm_107 plans together with the
  speculative RNG sites.
* The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode
  shortcuts and the sampler's store copies are those components' own.
* MSA's in-forward paths are its fixed behaviour: the sub-page table is
  always expanded inside the forward, the block table always keeps its
  owner-side record, and the out_cache_loc tail is always filled in the
  forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and
  TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the
  eager tail fill are removed; nothing reads a merged destination before
  the replay, so no M3 plan needs a flush_after barrier.
* The step prologue ops read device_copy's graph scatter as reference.
* Ported tests: test_step_inputs' parallel-scatter and in-graph tests
  (test_device_copy_in_graph.py), test_h2d_changed_span
  (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests
  (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph
  runner's prologue test, the MSA skip-routing scan, the sm_107 contract
  scans and step-prologue tests; test_copy_plans.py checks every M3
  target's plan.

Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

10 participants