[TRTLLM-16622][feat] DeepSeek-V4.1: Perf optimizations for DS-V4.1-Flash-NVFP4 by ziyixiong-nv · Pull Request #21438 · NVIDIA/TensorRT-LLM

ziyixiong-nv · 2026-11-02T09:50:58Z

This PR ports the DeepSeek-V4.1 performance optimizations onto main.

It brings:

DeepSeek-V4.1 on the PyTorch backend: compressed KV, mHC, Engram and the DSpark draft;
CSA2 sparse attention: decode and prefill kernels, the indexer and exact Top-K;
bounded encoder and decoder replay, with SWA scratch reuse;
MegaMoE: the one-round push router, the fused shared expert, eager launches without PDL, and warmup expert placement (moe_config.load_balancer.warmup_placement, standard mode only);
GB300 (SM103) and SM107 tuning for disaggregated serving, with the executor, KV-cache and test changes these need.

Shared code keeps main's behavior for other models. Where this PR changes a shared path for DeepSeek-V4/V4.1, the model opts in:

RuntimeHooks set by DeepseekV4ForCausalLM: stage_decode_inputs, prioritize_decode_graphs, and loop_on_execution_stream with speculative decoding;
use_fp8_small_m_gemm for the FP8 small-M GEMM race;
DeepSeek-V4 routing for the NVLink one-sided pull combine, one-round MegaMoE dispatch, CuteDSL FC1 zeroing and the spread tuning profile;
caller flags for the wide Top-K configurations (CSA2 only).

@coderabbitai summary

Description

After applying the optimizations, the AgentX perf is as below.

Behavior changes to know about:

DeepSeek-V4.1 runs Encoder replay by default when Decoder bounded replay runs with CED precompute; TRTLLM_V41_ENCODER_REPLAY=0 keeps the required Encoder state.
CSA2 candidate-consumer layers need an SM100-family GPU, DeepGEMM's sparse MQA-logits kernels, 32 index heads of 128 dims and candidate_block_size a multiple of 8; model construction fails otherwise.
gen_prompt_tail_tokens needs KV cache manager V2 and transceiver_runtime: PYTHON (set explicitly; auto is rejected).
The DeepGEMM mqa_logits patch gives bitwise-identical logits for finite inputs; NaN and -inf scores now contribute 0 instead of NaN (DSA prefill: DeepSeek-V3.2, GLM-5).

Answers to review questions:

The CSA2 cutover constants (_DENSE_O_B_MIN_ROWS, PROJECTION_CUTE_DSL_ROWS, _N256_MIN_ROWS, the fitted constants in csa2SparseDecode.cu) and HOLD_MAX_ROWS were tuned with DeepSeek-V4.1 on GB300 (SM103); they were not measured on B200. They are reached only by DeepSeek-V4.1.
GVR self-sampling decode on SM107, against the exact decode Top-K, DeepSeek-V4.1 TP8 synthetic-KV generation throughput: C1 -1.27%, C256 -5.95%, C8192 -6.69%. SM107 decode therefore keeps the exact Top-K; prefill uses self-sampling.
gen_prompt_tail_tokens >= 256: generation workers read prompt ids only near the end (the Engram lookback and stop sequences that straddle the prompt end); 256 covers both with margin.

Test Coverage

New L0 entries:

l0_b200 (SM100): DS-V4.1 bounded replay, snapshot and decoder-global model tests; DSpark draft step, acceptance, argmax chain and sampling buffers; fused DS-V4 route+quant (TRTLLM, MegaMoE); FP8 small-M MLP; vocab-slice embedding; LL all-gather; deferred gen-init receive.
l0_b300 (SM103): all of test_mhc.py and test_engram.py, replay model tests, DSpark CuTe DSL attention, draft step, argmax chain, WKV projection, DS-V4 route+quant.
l0_gb200_multi_gpus: MegaMoE e2e, warmup placement and combine tactics; DSpark Markov chain and forked draft prep; peer buffers; mHC all-reduce boundary; MNNVL RMSNorm; attention-DP decoder-pass graph and Engram sharding; autotuner cache disagreement; CSA2 two-rank TP.
l0_gb300: DSpark Markov chain and forked draft prep. l0_gb300_multi_gpus (post-merge): MegaMoE e2e and placement, attention-DP decoder/encoder-pass graphs, breakable prefill graphs against eager prefill.
l0_gr100 (SM107, post-merge): DSpark CuTe DSL attention and RMSNorm/RoPE.
Accuracy: TestDeepSeekV41::test_gsm8k_dspark[nvfp4-tp4] / [nvfp4-dp4] (l0_gb300) and test_breakable_prefill_cuda_graph.
l0_dgx_h100 (existing entry): test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap now also runs the disaggregated scheduling fields (gen-init admission cap, client I/O during the GPU wait, prompt-tail trimming, prefill-aware conversation placement) and checks the answer to a long prompt that the generation workers receive trimmed.

Run on GB300 (SM103) at dfcfb6e, rebased on main ee1e8e4, with its C++ changes built: the CSA2 kernel, module, decode (with and without PDL), indexer and cache tests; the DS-V4.1 model tests; attention-DP, Engram sharding, MegaMoE e2e and warmup placement; the DSpark, speculative and sharded-sampling tests; GSM8K with DSpark (TP4, DP4); the extended attention-DP overlap disaggregated test; CSA2 tensor parallel; breakable prefill; and Qwen3.8-Flash-Next MMMU. Not run here: SM100 (B200) and the 8-GPU cases (helix CP, world-8 attention DP).

Since then, on GB300 with the C++ changes built:

at 15e5184 (main 1399e99): the PyExecutor, LLM API and API-stability unit tests; main's new engine, sampler, speculative and QSA tests; the MXFP8 and MiniMax-M3 dispatch tests; GSM8K with DSpark (TP4);
at 9c7e147 (main 7a06809): the mHC, DSv4 compressor, DSpark engine-runtime, decoder-runner, DS-V4.1 model, disaggregated extractor, mixed-precision KV and serve-schema tests; GSM8K with DSpark (TP4).

The head 66f5889 is 9c7e147 rebased onto main d782229.

AgentX on GB300 at 9c7e147, this PR against the bring-up branch it ports (first 30 minutes of each run):

hi (c7000): +1.3 to +1.8% throughput, +1.6 to +2.6% tokens/s per user;
mid (c7600): +0.5 to +3.7% throughput, +1.0 to +5.2% tokens/s per user;
ll (TP8, c84): +0.0 to +0.2% throughput, +0.3 to +0.4% tokens/s per user.

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.

ziyixiong-nv · 2026-11-02T09:51:56Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-02T09:58:56Z

PR_Github #81973 [ run ] triggered by Bot. Commit: 12faa4c Link to invocation

tensorrt-cicd · 2026-11-02T17:34:47Z

PR_Github #81973 [ run ] completed with state SUCCESS. Commit: 12faa4c
/LLM/main/L0_MergeRequest_PR pipeline #67615 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

sunnyqgg

Requesting changes. This review has 41 blockers, 32 possibly-blocking items and 18 required items: 87 are inline, and the 4 below have no anchor in the diff. The most urgent are the TRTLLMINF-435 TEMPORARY CI override, stages and test lists, which must not reach main, the version regression to 1.4.0a0, and the shared-default changes that alter behavior for models other than DS-V4.1 without data. Some areas got only a light pass: the inner pipelines of the large CuTe DSL kernels (CSA2 decode/prefill, MegaMoE, mHC); the ~28k lines of new tests, checked for registration, whether they can fail, hangs and env leaks but not test logic; and deepgemm.patch, not compared line by line with the pinned upstream source. Nothing was run on a GPU, and CI had not reported for PR_Github #81973 when this was written.

[Possibly blocking] cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/indexConversionKernels.cu:160 (outside the diff). The PR changes the SWA scratch rule to "a block in the scratch range that already has a regular page keeps it". It updates four page-table builders:
- PageIndexConverter (kvCacheManager.cpp:74);
- the Python table (kv_cache_manager_v2.py:1485);
- both DeepSeek-V4 block-table kernels (deepseekV4BlockTable.cu:205/284-285);
- the CSA2 converter (csa2/kernel.py:444).
It does not update indexConversionScratchKernel, which Batch::_convertIndices uses (batch.cpp:882-896); that kernel still maps every block in [scratchBeg, scratchEnd) to scratch. The two can disagree whenever a speculative reservation has kept a regular page in the scratch range. A caller of the public tensorrt_llm.runtime.kv_cache_manager_v2.Batch index tables with SWA scratch reuse then gets a page table that differs from the manager's own converter: KV goes to scratch while the persistent page goes stale. This is latent today, because the PyTorch executor doesn't build a Batch with index tables. Separately, the reference in tests/microbenchmarks/dsv4_block_table_perf.py:77-78 still uses the old rule and fails its first comparison with the default --bad-ratio 0.05. Please add raw[block] == kBadPageIndex && to the scratch predicate, update the microbenchmark reference, and consider one shared helper for the five copies of the rule.
[Required] The PR title, description and scope don't follow the repo rules:
- Scope. The title says "Perf optimizations for DS-V4.1-Flash-NVFP4", but the PR ports the whole bring-up branch: 389 files, +76k/-7.7k. It also includes two [fix] commits for KV cache manager v2 that the description doesn't mention.
- Ticket. It is [None]. CONTRIBUTING.md:91 says: "For NVIDIA developers, please include the JIRA number or NVBUG ID in the PR title whenever possible."
- Testing. Test Coverage is empty, and the description says the full CI run "has not been re-run on this commit".
- Undescribed shared changes. Several changes affect all models:
  - logits processors are rejected for every one-model speculative mode;
  - attached frontends start with CUDA_VISIBLE_DEVICES="";
  - decode graphs with ≤256 tokens are captured on a priority stream;
  - every BF16 embedding on SM100 runs a new CuTe kernel;
  - every TRTLLM-backend model gets a HostUploader;
  - the autotuner adds a per-profile cross-rank sync and cold-L2 profiling;
  - MNNVL all-reduce returns None as a fallback;
  - the DeepGEMM sm100_mqa_logits patch.
- Reviewability. Owners of the shared code (executor, sampler, serve, autotuner, embedding, distributed) can't evaluate their parts. Once squashed, any regression bisects to a single 76k-line commit.
- Move plus rewrite. dspark_rmsnorm_rope.py → dspark/rmsnorm_rope.py moves the file and rewrites it in the same change (+184/-100).
AGENTS.md: "One concern per PR — avoid scope creep. If a PR touches unrelated areas, split it."

Please add a JIRA ticket and split the PR into:
- infra;
- each shared default change, with non-DS-V4.1 perf/accuracy data;
- the generic disaggregated-serving features;
- the KV cache manager v2 fixes;
- the DS-V4.1 model work.
Then fill in Description and Test Coverage, and re-run CI on the final commit.
[Required] The user docs contradict the code, and new features are undocumented. The docs/source files are outside the diff.
- Encoder replay. docs/source/features/sparse-attention.md:353-368 and docs/source/features/feature-combination-matrix.md:7-15 still say Encoder replay is opt-in ("Encoder replay is opt-in via TRTLLM_V41_ENCODER_REPLAY=1").
- Hardware. Those two pages and the CSA2 README hardware table (csa2/README.md:21-23/43/110-112/352) still list SM90 and SM120/121. Candidate-consumer layers now need SM100.
- SM103 decode path. The CSA2 README leaves out the SM103 gen_fmha decode path (backend.py:354-355/453, gen_fmha.py:28), and two of its statements are not true for that path.
- New knobs. None of these has a docs/source page:
  - the env var TRTLLM_V41_BOUNDED_REPLAY_TOKENS;
  - prompt_token_cache_config, gen_prompt_tail_tokens and enable_client_io_during_gpu_wait;
  - max_gen_init_admissions_per_pass and host_cache_bind_to_gpu_numa_node;
  - megamoe_fuse_shared_expert and warmup_placement.
- Supported models. supported-models.md has no V4.1 hardware note.
As a result, operators will:
- assume approximate Encoder recovery is off;
- deploy checkpoints with candidate layers on Hopper or SM120, where they fail at warmup;
- look at csa2_sparse_decode while debugging GB300 large-batch decode, when an undocumented third kernel is running.
Please:
- update sparse-attention.md, the feature matrix and the CSA2 README (encoder replay default, the SM100 requirement for candidate layers, the gen_fmha path);
- document the replay env var or move it into the config;
- document the new knobs;
- add the V4.1 hardware footnote to supported-models.md.
[Required] tests/unittest/llmapi/test_trtllm_serve_config_schemas.py is unchanged, although this PR adds serving YAML fields and constraints:
- gen_prompt_tail_tokens (ge=256);
- the MoeWarmupPlacementConfig bounds;
- PromptTokenCacheConfig;
- the least_prefill value;
- CSA2's _validate_gvr_emission.
AGENTS.md: "Changes to serving YAML fields, aliases, defaults, types, or constraints ... must regenerate and inspect the editor JSON Schemas ... and add regression coverage for the affected accepted/rejected values."

Without those tests:
- the schema and the runtime can drift apart silently;
- the rejection list in validate_gen_prompt_tail_tokens has no unit test;
- test_load_balancer_schema_matches_runtime doesn't cover warmup_placement.
Please add accepted and rejected schema cases for each new field and validator, and extend test_load_balancer_schema_matches_runtime to cover warmup_placement.

ziyixiong-nv · 2026-11-03T07:45:01Z

Thanks for the review. The fixes are in fbe5b2a, and each inline thread has a reply. On the items in the review body:

Scope, JIRA, description, split

We'll keep this as one PR: splitting the bring-up into five dependent PRs would serialize their reviews and CI and push the merge out much further. To keep the shared code reviewable, the CI/infra changes are reverted, and the shared-default changes you listed now keep main's behavior for other models and apply only to DeepSeek-V4/V4.1 (see the individual threads). The title now carries TRTLLM-16622, the description now lists the behavior changes, the answers to your questions and the Test Coverage, and CI will be re-run on the final commit.

indexConversionKernels.cu (SWA scratch predicate)

Fixed: indexConversionScratchKernel now takes a scratch slot only for a block whose raw entry has no regular page, the same rule as PageIndexConverter and the other three builders, and the microbenchmark reference applies it too. The existing test that compares a Batch's converted table with PageIndexConverter now also runs with a speculative reservation whose regular page lies inside the next chunk's scratch range. The five copies live in C++, two CUDA kernels in different libraries, torch and CuTe DSL, so a shared helper would cover only two of them; each keeps the same one-line predicate.

User docs vs code

Thanks, the docs now follow the code. sparse-attention.md and the feature matrix say that Encoder replay is on by default whenever Decoder bounded replay runs with CED precompute, with TRTLLM_V41_ENCODER_REPLAY=0 to keep required Encoder state and the fallbacks for other layouts and small max_num_tokens; they document TRTLLM_V41_BOUNDED_REPLAY_TOKENS (default: 140, or the required window plus one) and limit SM90 and SM120/121 to checkpoints without candidate-consumer layers, and supported-models.md has a matching DeepSeek-V4.1 footnote. The CSA2 README's routing table and text now cover the SM103 gen_fmha decode path for large ratio-2 generation tiles and no longer say that every 64-head generation tile skips staging for the fused decode. The new knobs are documented next to their features: the prompt token cache with the prefix-tokenization cache, client I/O with the overlap scheduler, prompt-tail transport and the admission bound with disaggregated serving, host-tier NUMA binding with KV cache offloading, warmup placement with EPLB, the fused shared expert in the MoE performance guide, and TRTLLM_ENGRAM_EVICT_PAGE_CACHE in the model footnote. The replay switches remain environment variables in this PR, like the existing TRTLLM_V41_* switches.

test_trtllm_serve_config_schemas.py coverage

Added schema regression cases in test_trtllm_serve_config_schemas.py. test_config_validation covers gen_prompt_tail_tokens (256 accepted, 255 rejected, also under generation_servers), PromptTokenCacheConfig, least_prefill and the CSA2 field set, and test_load_balancer_schema_matches_runtime now covers warmup_placement (accepted configs, the required start_step, the interval_steps/decay/max_moved_fraction bounds, an unknown planner and per-layer dynamic mode). Rules that span fields, which the editor schema does not express (validate_gen_prompt_tail_tokens, CSA2 _validate_gvr_emission, warmup placement with online EPLB), have a new test that validates each YAML against the schema and then builds TorchLlmArgs from it. The checked-in schema snapshots are regenerated.

ziyixiong-nv · 2026-11-03T07:50:56Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-10T08:38:52Z

PR_Github #83717 [ run ] completed with state ABORTED. Commit: 1962e7e

Link to invocation

…ptured sampling variants One-engine speculative decoding captures one set of decode CUDA graphs per sampling variant: greedy, advanced and, where the model samples vocabulary-sharded logits, temperature_only. A workload that sends one variant still pays the warmup, capture time and graph memory of all of them at every startup. cuda_graph_config.sampling_variants lists the variants to capture; unset keeps the current behavior. A decode batch of an unlisted variant runs eagerly, and the first such batch of each variant logs a warning with the captured set (warmup batches are not reported). The variant check is part of the attention-DP graph agreement, so one rank's unlisted variant sends every rank eager in that step, as replay assumes (it freezes all_rank_num_tokens); the graph key and the check read the variant from one helper. Without temperature_only in the list, temperature-only batches use the advanced graphs. Argument validation rejects an empty or repeated list, a list without one-engine speculative decoding (AUTO is accepted: it becomes suffix-automaton decoding), a list with DSpark confidence scheduling (it plans each step's graph before the step's variant is known), and temperature_only without MTP-Eagle one-model decoding with tensor parallelism, no attention DP and no relaxed acceptance. A variant the model cannot capture is rejected at the first capture. The golden manifest and the serve config schemas gain the new field. Measured on GB300, DeepSeek-R1-0528 NVFP4, DEP4 with MTP (3 draft tokens), 21 decode batch sizes, two sets of runs on two GB300 nodes, stock runs before and after: sampling_variants=["greedy"] captures 21 graphs instead of 42; the final executor's CUDA graph phase took 18.5 s against 36.2 / 35.1 s and 16.4 s against 33.2 / 32.3 s (rank 0; -48% and -50%), and the KV-cache estimation executor's capture drops the same way, about 30 s per startup in all. With the check, DEP4 with [greedy], forced all-gather MoE communication and a mixed batch (five requests, one sampled) completed 5/5. Without it the same probe also completed here, with ranks mixing replay and eager in one step; whether that hangs or miscounts depends on the communication backend and the shapes. Tested on main 0bdadaa plus NVIDIA#21438, which also changes cuda_graph_runner.py, graphs.py, runner.py and metadata.py. Test: tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_sampling_variants.py runs Llama-3.1-8B-Instruct with an Eagle3 one-model draft: sampling_variants=["advanced"] captures exactly the advanced graphs of the default engine, and its greedy requests run eagerly under the once-per-variant warning and decode the same tokens as an engine without CUDA graphs; it fails without this change. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…timizations onto main This commit is the DeepSeek-V4.1 bring-up branch's change against main 5d86481, the main commit that branch was last synced to, applied to c67c0f4 with a three-way merge. It brings: - DeepSeek-V4.1 on the PyTorch backend: compressed KV, mHC, Engram and the DSpark draft; - CSA2 sparse attention: decode and prefill kernels, the indexer and exact Top-K; - bounded encoder and decoder replay, with SWA scratch reuse; - MegaMoE: the one-round push router, the fused shared expert, eager launches without PDL, and warmup expert placement (moe_config.load_balancer.warmup_placement, standard mode only); - DenseGEMM MoE fixes for padded tiles; - GB300 (SM103) and SM107 tuning for disaggregated serving, with the executor, KV-cache and test changes these need. Main gained 48 commits after 5d86481. Four files conflicted: - tensorrt_llm/inputs/registry.py: both sides extended the typing import; the resolution keeps both. - jenkins/Build.groovy: main dropped buildWheelInContainer's cpver argument; the Debug build keeps the bring-up's NVRTC_DYNAMIC_LINKING=ON. - mega_moe_cute_dsl.py: main made MegaMoE's MPI bootstrap of torch.distributed bind its store on rank 0, leave MASTER_ADDR/MASTER_PORT alone and refuse a loopback store across hosts. The bring-up had moved that bootstrap into pg_utils.init_torch_distributed_under_mpi, which warmup placement also uses, so main's version now lives there. - kv_cache_manager_v2.py: main syncs the init status across aggregated helix CP ranks; the bring-up's disk-tier error stays on the unsynced path. Test: the same change on 5d86481 passed the bring-up branch's full pre-merge CI, including the GB300 DeepSeek-V4.1 stage (GSM8K with DSpark, breakable prefill CUDA graphs). It has not been re-run on this commit. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…planner ConversationManager skipped plan_committed_block_drop() whenever the KV cache had a PRIVATE dense sliding-window group, on the assumption that the plan was then always empty. The native planner no longer works that way: it skips PRIVATE life cycles, still collects the pages of the other windowed groups and SSM, and its radix match only needs REQUIRED groups covered. DeepSeek-V4.1 with bounded decoder replay is such a cache: its decoder SWA layers are PRIVATE, while its encoder SWA layers and compressor states are windowed OPTIONAL or REQUIRED groups. With per_conversation block reuse the shortcut kept a finished turn's tail-window pages of those groups resident until LRU eviction instead of dropping them once the next turn completed. Remove the shortcut, so ConversationManager matches main again and the native planner decides what a plan holds. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

… in HostMemBacking HostMemBacking already places the host tier on a NUMA node, with libnuma, MPOL_PREFERRED or MPOL_BIND from allowRemoteNumaFallback, and a warning fallback where mbind is denied. HostCacheTierConfig.numa_node added a second binding: HostMem re-bound the whole reservation with a raw mbind syscall, always MPOL_BIND, after the backing had bound it. The later call won, so the effective policy depended on call order, ignored allowRemoteNumaFallback, and failed with its own error handling. HostMemBackingOptions now carries the node, and the backing's one binding uses it instead of the GPU's node when it is set. The policy stays allowRemoteNumaFallback's: kv_cache_config.host_cache_bind_to_gpu_numa_node passes its node with allow_remote_numa_fallback=False, which keeps its strict binding. HostMem and the host pools drop their numaNode parameters, and the strict-binding error now names both settings that relax it. Test: test_host_tier_bound_to_gpu_numa_node (tests/unittest/_torch/pyexecutor/ kv_cache/test_kvcm2_integration.py) checks that the host tier lands on the GPU's node while the building thread prefers another node. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…nt ones Drop this branch's CI wiring: the jenkins/ stage and build changes and the separate DS-V4.1 test-db lists go back to main. jenkins/license_cpp.json keeps the three csa2SparseDecode entries, since those sources carry DeepSeek's MIT header and the C++ license check fails without them. Run a DS-V4.1 subset in the standard L0 lists instead, where main's directory entries (attention/, pyexecutor/, kernels/) do not already collect it: - l0_b200 (SM100): bounded replay, snapshot and decoder-global model tests; DSpark draft step, acceptance, argmax chain and sampling buffers; fused DS-V4 route+quant for TRTLLM and MegaMoE; small-M FP8 MLP, vocab-slice embedding, LL all-gather, DeepGEMM JIT cache keys; deferred gen-init receive. - l0_b300 (SM103): all of test_mhc.py and test_engram.py, the replay model tests, DSpark CuTe DSL attention, draft step, argmax chain, WKV projection, and DS-V4 route+quant. - l0_gb200_multi_gpus: MegaMoE e2e, warmup placement and combine tactics; DSpark Markov chain and forked draft prep; peer buffers; mHC all-reduce boundary; MNNVL RMSNorm; attention-DP decoder-pass graph; CSA2 two-rank TP; disaggregated gen-init cap. - l0_gb300: DSpark Markov chain and forked draft prep. - l0_gb300_multi_gpus (post-merge): MegaMoE e2e and placement, attention-DP decoder/encoder-pass graphs, breakable prefill graphs against eager prefill. Trim the new tests: drop helper, decision-table, mock-wiring and duplicate tests, and keep the strongest parametrizations. The branch's test diff goes from +19.5K to +17.4K lines and its new test functions from 259 to 201. Fix tests that no list had run on this branch: - test_dspark_slot_reset_no_sync (main's) sets the worker's host slot mirror and no longer checks the scratch row prepare() stopped resetting; this branch's copy of that test is dropped. - test_verify_logits_dtype's metadata stub gains sharded_target_logits. - The DS-V4.1 MPI worker initializer installs the test tree's import finder, so workers can import test_modeling_deepseekv41 (its defs import) as pytest does. Test: on GB300, every new L0 entry passes, including MegaMoE e2e and placement, the attention-DP graph passes, CSA2 two-rank TP, breakable prefill against eager (17 min) and the disaggregated gen-init cap test. So does every test file this branch adds or changes (collect-only: 8,963 items, no errors) and main's DS-V4.1 L0 entries. SM100-only behavior was not run. The test-db lists pass scripts/check_test_list.py --validate. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…ations Fix the defects the review found for every model, and keep main's behavior on shared paths for other models while DeepSeek-V4/V4.1 opt in to their tuned paths. Opt-ins, so other models get main's behavior: - RuntimeHooks stage_decode_inputs (one-launch decode input staging) and prioritize_decode_graphs (high-priority decode capture), plus loop_on_execution_stream with speculative decoding, set by DeepseekV4ForCausalLM. - use_fp8_small_m_gemm for the FP8 small-M race; DeepSeek-V4 routing for the pull combine, one-round MegaMoE dispatch, FC1 zeroing and the spread tuning profile; caller flags for the wide Top-K configurations (CSA2 only); the private MQA logits pool only for compressed logits. Fixes: - Ranks agree: autotuner profile hits in one all-gather per call; peer buffers fall back to NCCL on every rank together; KV host-tier NUMA read errors reach the failure allreduce; MegaMoE shared-expert fusion is all-or-none across the target, MTP and DSpark draft layers. - Warmup placement: only backends that can re-place experts; round buffers reserved before the vote; no counting during KV estimation; MTP draft layers keep their placement; counters halve before int32 overflow. - DSpark: PyTorch fallbacks restored; seeded rows stay on the stock sampler; FP32 LM-head outputs; draft head copies built at load; logits processors rejected only for the embedded DSpark worker. - 64-bit offsets in CSA2 MXFP8 quantization; CuTe masked embedding only below 2^31 elements; 64-bit pull-combine epoch with a bounded handshake; FORCE_DETERMINISTIC mHC boundaries; SWA scratch predicate in the KVCM2 scratch converter. - Config: CSA2 config stops inheriting DSA-only fields; validators for prompt-tail transport, warmup placement and GVR emission; prototype status markers; serve schemas regenerated. Version back to 1.4.0rc3. - deepgemm.patch split into two described patches; the inactive SM107 config.h and JIT cache-key hunks are dropped. - Frontends hide GPUs only for input processors that never need one. - Docs: CSA2 hardware requirements, Encoder replay default, the new knobs next to their features. Test: on GB300 with these C++ changes built, the changed test files pass: pyexecutor 422, llmapi 520, serve schemas 171, KVCM2 340, DSA/CSA2 indexers, mHC, DSpark 76, and multi-GPU MegaMoE e2e/placement, shared MoE, peer buffers, DSpark draft prep, attention-DP engram, CSA2 two-rank TP 14/14. Breakable prefill (18 min), disaggregated gen-init cap, multi-frontend, GSM8K DSpark TP4/DP4 and an 80/80 DS-V4.1 multi-turn smoke pass. SM100 and 8-GPU helix were not run. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Pre-merge CI on this branch failed in single-GPU and CPU stages, and the rebase onto main's DSpark confidence scheduling (NVIDIA#20078) broke its new tests: - kv_cache_manager_v2: _remove_kv_cache discards the cache's pending stats before the lender's free hook again, as on main. A raise there must leave the lender's records (test_staging_lender). - sync_group_all_greedy_sample takes the CUDA-graph runner and batch as optional keywords, so main's confidence callers keep their signature (test_dspark_sampling_admission). - DecoderMetadata._host_uploader and DeepseekV41Model's Engram context-overlap attributes get class defaults, so objects built without __init__ (main's test_dspark_model_engine_runtime and test_deepseek_v41_multimodal_backbone) take the plain path. - Attached frontends get --gpus_per_node only when they hide the GPUs. The OpenEngine gRPC launcher's engine needs no GPU count otherwise (test_frontend_lifecycle). - test_fused_decode_select_matches_chain compares with the CuTe DSL decode chain only on SM100-family GPUs, where CSA2 uses it; it rejects the wide rows' launch on SM120. - Test stubs gain the attributes this branch reads: the FC12 warmup test's routing_method, nvfp4_experts in the embedded-confidence loading test, and enable_attention_dp in the sampling-admission tests' mappings. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…ing side streams A Full or Reindex layer that forks its side streams is the first to resolve the forward's SWA slots when an earlier layer has not: - when the forward starts at that layer; - right after the CED decoder boundary re-arms them, in bounded replay and in the remote-tail decoder entry of a disaggregated generation worker. The descriptor stream's resolution then refreshed the slots on that stream and marked them resolved. The model stream's cache write read them without joining the stream, so a decode step could write its SWA row into the previous step's slot. The layer now resolves the slots on the model stream before it records the fork event, so its side streams and the cache write see the same refresh. Test: test_native_module_auxiliary_stream_and_graph also runs two decode steps at new positions, overlapped with the descriptor stream held back and then serially, and requires bit-identical outputs. It fails without the fix. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

- FP4BlockScaleMoERunner.run_moe keeps main's arguments again: main's experimental modeling_v2 MoE ops call it positionally. The caller-built routing tables take a new run_moe_routing_tables method. - The CuTe DSL FC2 finalize runner passes the L2 evict-first hint only to kernel classes derived from its own: modeling_v2's runner subclasses swap in their own copy of the kernel, which does not take it. - test_fused_decode: the FP32 reference dequantizes only each chunk's selected rows, and the large decode test takes each launch's reference as it launches instead of keeping full cache copies (out of memory on B300). - test_megamoe_e2e: only SM103 has the shape-tuned default tactics that pick a different scheduler at the two buckets; on SM100 the default and the plain path are compared without requiring different schedules. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…rted GEN receive - DisaggTransferCoordinator.forget_request() also drops the request's started GEN receive. Otherwise only completed_gen_receives() pruned it, and only a forward calls that: a worker with repeated transfer failures and no forward kept the retired requests and their prompts. - The disaggregated scheduling fields move into the existing DeepSeek-V3-Lite attention-DP overlap test instead of a separate four-GPU pre-merge case: the gen-init admission cap, client I/O during the GPU wait, prompt-tail trimming and prefill-aware conversation placement. A long prompt that reaches the generation workers trimmed to its tail must still be answered, with its full prompt length reported. Test: the extended test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

- test_connector_adp: the executor stubs also set previous_batch, track_unique_tokens and _adp_dummy_is_gen, which PyExecutor.__init__ sets and this branch's loop, request fetch and ADP dummy padding read. - test_breakable_prefill_cuda_graph states its four GPUs with skip_less_device, as main's multi-GPU accuracy tests now do (NVIDIA#21541); the helpers it called are gone from the module. - Waive test_deepseek_v4_mixed_disagg_names_every_layer on B300 too: it fails with main's code, which waives it on DGX_B200 (nvbugs/7043616). Rebased onto main ee1e8e4. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Runtime fixes: - Staged responses keep the result queue looked up when they are gathered, so a terminal response reaches the supplied actor after its request ends. - SharedPromptStore eviction subtracts what it deleted from the current total when a write raced the scan, so no entry is counted twice. - PeerBuffers: the MNNVL vote uses the communicator's allgather (also ProcessGroupComm); the IPC path agrees on every rank's allocation in the handle exchange before any rank decodes or opens a handle. - DSpark fused acceptance and next-token preparation serve drafts of at most 31 tokens; wider drafts take the base path. Hardening: - The fused NVFP4 MoE inputs decline AWQ/ARC pre-quant scales; MegaMoE's shared expert takes the unscaled input. - The small-M GEMM race is decided on every rank of the tuning group. - CSA2: one setmaxnreg per warpgroup in the context kernel; decode schedules are validated and the combine wait traps after seconds; both kernels check their shared memory against the opt-in limit; the FP8 Q is skipped when the direct decode serves a step; packed norm rows are bounded to one warp. - Host-side checks at the entry points of the Engram, mHC, GEMM, sparse MLA, alltoallv_tokens and inverse-RoPE kernels and ops; spec store scatter and UE8M0 scale checks; deferred Markov heads ignore the hidden. - Deferred gen receives stay off with a KV connector; the Request pickle layout is validated once at startup; forward_MoE takes its all-reduce flags as arguments; the warmup placement keeps its validated config; MegaMoE recognizes the multimodal V4.1 architecture. - The DS-V4.1 perf-sanity cases keep main's KV dtype (auto). Test: on GB300 (SM103) with this change's C++ built: the CSA2 context, decode (with and without PDL), module, kernel, direct-decode, indexer and cache tests, CSA2 tensor parallel, inverse RoPE, the DS-V4.1 model tests, MegaMoE (incl. NVFP4 AWQ), attention-DP and Engram sharding; the speculative, DSpark, MoE route, Engram, mHC, low-M and small-M GEMM, autotuner, pyexecutor and connector tests; mHC all-reduce, draft prep, sharded sampling, peer buffers (also with one rank failing, with and without MPI), warmup placement and MegaMoE e2e; GSM8K with DSpark (TP4, DP4), the attention-DP overlap disaggregated test, breakable prefill and Qwen3.8-Flash-Next MMMU. Serve schemas and the LLM args manifest are unchanged. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Stub objects in main's tests lack fields this branch reads: - test_cuda_graph_runner, test_step_inputs: the runner config stubs set stage_decode_inputs, which CUDAGraphRunner.replay reads before use_mrope. - test_queued_context_cancellation: the executor fixture sets track_unique_tokens, which the request fetch passes to merge_requests. - test_llm_args TestMultimodalConfig: the args stub sets prompt_token_cache_config, which TorchLLM passes to the input processor. Rebased onto main 1399e99. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

- MiniMax-M3's MXFP8 dispatch state and the SM107 TP4 pinned-tactic path read tune_mxfp8_backends, the name this branch gives tune_decode_graph_backends; the old name raised AttributeError on every SM107 MXFP8 GEMM. - CSA2SparseAttentionConfig.indexer_max_chunk_size is a PositiveInt, so 0 and negative values are rejected instead of meaning 32768 or one row. The serve schemas are regenerated. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…in the ADP dummy rollback test Main's new ADP dummy rollback test calls PyExecutor._revert_gen_alloc, which this branch renamed to _revert_unexecuted_alloc (it also gives back the growth of unexecuted context chunks). The test's batch has no context requests, so the call is equivalent. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…main's PARALLEL autotuner test The test runs two ranks one after another over a fake tp_cp_allgather that kept only each rank's latest payload. This branch's rank agreement adds host all-gathers before the PARALLEL merge, so the second rank received the first rank's merge payload in place of its agreement payload. Pair the calls by index, as real collectives do. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

- prompts_to_owners (rank 0 sends each admitted prompt only to the rank that owns it under multi-rank MPI attention DP) is now the RuntimeHooks opt-in prompts_to_owners, which DeepSeek-V4 and V4.1 set. Other models keep broadcasting every prompt and skip the prompt-splice startup check. - The CuTe DSL masked-embedding lookup on SM100/103/107 is now the Embedding opt-in use_masked_embedding_kernel, which the DeepSeek-V4 and V4.1 embedding sets; MTP and DSpark reuse that table. Other models, including MiniMax-M3's direct pre_comm_embedding_ops callers, keep the previous lookup. - The queued-request counters read py_prompt_trimmed_prefix directly: ExecutorRequest always defines it as an int. The masked-embedding test opts in to keep covering the kernel, and the iter-stats test's queued request stub gains py_prompt_trimmed_prefix. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…g rank-local The rank agreement this PR added to AutoTuner.choose_one() (a host all-gather of the cache hit under collective strategies) also ran for a choose_one() reached from inside a runner the tuner is profiling. Under BROADCAST only rank 0 profiles, so its nested all-gather waits for ranks that are waiting for rank 0's broadcast. Qwen3.5 with the TRTLLM MoE backend hits this in warmup: the MoE decode front-end (BROADCAST) profiles a runner that calls the gate's low-M GEMM (PARALLEL), so TestQwen3_5_397B_A17B::test_nvfp4_mtp3_gdn_replay_tep4 hung until the CI job was cancelled. A choose_one() reached from _profile_runners_impl now skips the agreement and decides on this rank's cache, as before this PR; top-level calls keep it. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…executor state they reach These tests build PyExecutor without __init__. With this PR's executor they also need max_draft_len (read once per step by the spec-decode stats), the overlap loop's intake window, and a mock of _revert_unexecuted_alloc, which replaces _revert_gen_alloc. Their stub previous batches also need guided_failures, which _process_previous_batch reads on main as well. Drop a fixture line that repeats main's own previous_batch reset. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…n the merged draft attention Main's paged DSpark attention tests drive the merged code with CPU tensors and with draft models built without __init__: - dspark_attention_forward_batched runs its K/V half inline for CPU inputs, and the context back-fill takes the CuTe DSL kernel only for CUDA tensors. - The back-fill and context projections bound at load default to None on the class; the back-fill tries the Rubin stacked FP8 context projection first and stacks the K/V norms itself when they were not packed. - _prepare_tp_attention checks the draft before reading the mapping. - Main's tests patch _get_dspark_arch_str (is_sm_100f is no longer used here) and give forward_batched's stub model the block size and TP flags it reads; main's fused FC12 stub backend gets the routing method this PR's tuning profile reads. - l0_gr100 keeps main's entries for the two DSpark kernel tests instead of listing them twice. - Main's GEN-preemption test stubs get _cancelled_disagg_model_request_ids (its _schedule calls it on main too), this PR's ctx_decoder_batcher, and the DSpark worker's host slot mirror. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

…p loop's inputs The overlap-loop tests in test_guided_decoder_errors.py run the real loop on an executor built without __init__. This PR's loop passes defer_gen_receives to _prepare_and_schedule_batch and opens the request-intake window around the sampler wait, so the stub takes that argument and the shell gets the window. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

ziyixiong-nv · 2026-11-10T11:17:08Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-10T11:25:59Z

PR_Github #83785 [ run ] triggered by Bot. Commit: e3e0bb1 Link to invocation

tensorrt-cicd · 2026-11-10T11:28:59Z

PR_Github #83739 [ run ] completed with state ABORTED. Commit: a0a4992

Link to invocation

…tests to the runner and request APIs - test_context_logits_memory's graph stub takes the runtime draft length and sampling mode this PR's decoder runner passes. - test_overlap_admission_lifetime queues bound executor Requests: since the typed request transport, ExecutorRequest is an envelope without Request's keyword constructor. Its executor stub also gets _cancel_queued_context_requests, which _fetch_new_requests calls on main too, and this PR's prompt delivery and unique-token tracking. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

tensorrt-cicd · 2026-11-10T17:07:20Z

PR_Github #83785 [ run ] completed with state FAILURE. Commit: e3e0bb1
/LLM/main/L0_MergeRequest_PR pipeline #69280 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

github-actions Bot assigned ziyixiong-nv Nov 2, 2026

zhenhuaw-me mentioned this pull request Nov 2, 2026

[Do NOT Merge][Preview branches] DeepSeek-V4.1-Flash performance optimization #21017

Closed

1 task done

longlee0622 added the ci: full pre-merge approved label Nov 3, 2026

ziyixiong-nv force-pushed the ds-v4.1-perf-opt branch from 1812cb1 to ff6ef97 Compare November 3, 2026 00:49

sunnyqgg self-requested a review November 3, 2026 01:52

sunnyqgg approved these changes Nov 3, 2026

View reviewed changes

ziyixiong-nv changed the title ~~[None][feat] DeepSeek-V4.1: Perf optimizations for DS-V4.1-Flash-NVFP4~~ [TRTLLM-16622][feat] DeepSeek-V4.1: Perf optimizations for DS-V4.1-Flash-NVFP4 Nov 3, 2026

ziyixiong-nv marked this pull request as ready for review November 3, 2026 07:50

ziyixiong-nv requested review from a team and peihu-nv as code owners November 3, 2026 07:50

ziyixiong-nv requested review from a team as code owners November 3, 2026 07:50

ziyixiong-nv added 21 commits November 10, 2026 11:12

ziyixiong-nv force-pushed the ds-v4.1-perf-opt branch from a0a4992 to e3e0bb1 Compare November 10, 2026 11:16

Hang-get mentioned this pull request Nov 10, 2026

[Performance]: SpecSampler.update_requests converts every preallocated slot to Python lists each step #22016

Open

1 task done

Conversation

ziyixiong-nv commented Nov 2, 2026 • edited by longlee0622 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

ziyixiong-nv commented Nov 2, 2026

Uh oh!

tensorrt-cicd commented Nov 2, 2026

Uh oh!

tensorrt-cicd commented Nov 2, 2026

Uh oh!

sunnyqgg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ziyixiong-nv commented Nov 3, 2026

Uh oh!

ziyixiong-nv commented Nov 3, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

ziyixiong-nv commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

tensorrt-cicd commented Nov 10, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

ziyixiong-nv commented Nov 2, 2026 •

edited by longlee0622

Loading