[None][feat] Kimi K3 optimizations on Rubin

[reasonsolo]

Dev Engineer Review

The inspected diff contains import-formatting changes in three Python files. It does not show runtime behavior, configuration, API, or test-list changes. The broader Kimi K3 and Rubin feature claims in the supplied objectives are not supported by this diff.

QA Engineer Review

No test changes.

Per-File QA Perspective

• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py: Import statements were reformatted. No observable behavior change appears in the inspected diff.
• tensorrt_llm/_torch/pyexecutor/_util.py: Import statements were reformatted. No observable behavior change appears in the inspected diff.
• tensorrt_llm/llmapi/llm_args.py: Import statements were reformatted, including moving the same Pydantic imports into a grouped import. No observable behavior change appears in the inspected diff.

Description

Based on main ff7ea92dfb.

No longer stacked. #19959 has merged (ff7ea92dfb); this branch now holds only its own 23 commits on top of main.

Kimi K3 work for Rubin (SM107).

• MXFP8 FP8 weight-read linear. On SM107 with Rubin CuTe DSL, _Fp8BlockScaleWeightReadLinear packs K32 R128c4 UE8M0 scales and runs cute_dsl_mxfp8_gemm_rubin with new K3 fine-grained low-M tuning buckets (fine_grained_m). SM100/SM103 keep DeepGEMM. Adds forward_prequantized.
• CUTEDSL_FC12 SiTU. The fused FC1+FC2 Rubin kernel specializes its gated epilogue for SiTU at trace time. The runner, op and leaf class forward the SiTU soft-caps and key the kernel and autotuner on the epilogue. The CuteDSL runners prime inner tactics only for the locality-domain path (prime_inner_tactics), so the epilogue-keyed FC12 runner no longer raises TypeError in autotuner warmup. Kimi K3 accepts moe_config.backend: CUTEDSL_FC12.
• CuteDSL NVFP4 SiTU on SM107. Drops the SM107 rejection added in [None][feat] Kimi K3: unlock the CUTEDSL MoE backend for NVFP4 SiTU #19003. The Rubin act-fusion kernel already implements SiTU.
• CuTe DSL MLA FP8 KV scales on the device. The FP8 decode kernel and op take the BMM1/BMM2 scales as device tensors. This removes the per-layer host read.
• Helix per-token bounds coverage for the CuTe DSL MLA decode. The kernel change is already on main via [None][feat] Helix speculative verify groups: fp8 + fp4 MLA and DSpark #19273.
• DWDP for Kimi K3 on CUTEDSL / CUTEDSL_FC12, plus compact_packing in submit_dwdp.py. Needs [None][feat] Rubin merge-back: SM107 MoE, CFT release op, MegaMoE multi-node bootstrap, core runtime fixes #19183 for the dwdp/setup.py part.
• Tensor-core (tcgen05) KDA MTP verify kernel. It has an optional bf16 tensor-core operand path (TRTLLM_KDA_MTP_BF16_MATH=1, off by default). The replay caches become kda_k_cache (bf16) + kda_g_cache (fp32) + bf16 kda_v_cache / kda_beta_cache / conv caches. bf16 and fp32 state pools are both accepted. SM107 is enabled for the optimized KDA kernels. The mixer now uses the kernel's fused attention-residual epilogue: packed token layout, gated RMSNorm in the epilogue (fuse_output_norm), and on SM107, for verify-only batches outside the breakable CUDA graph, an MXFP8 core with R128c4 scales that goes straight to o_proj.forward_prequantized (Rubin CuTe MXFP8 GEMM). Without the output-gate section or finalized o_norm weights, the previous launch with the Python output gate is kept as the fallback.
• KDA suspend on the tensor-core kernel. [None][fix] Kimi K3 / KDA / MNNVL / DFlash: fixes for defects on main #19988 commits a suspended request's pending KDA drafts through a replay-only variant of the verify kernel. That mode is ported to the tcgen05 kernel (REPLAY_ONLY: the chunk is the accepted drafts, only the state is stored), and kda_mtp_commit_pending_drafts, commit_kda_pending_drafts and the V2 manager's _commit_kda_pending_drafts read the split caches. The verify after a resume starts a new WY-form chunk, so it matches an unsuspended run to the kernel's accuracy rather than bit for bit; the conv windows and replay caches stay bit-exact.
• CuTe DSL MLA decode split-KV sizing for skewed decode KV. About 8 CTA waves over the packed 128-row query tiles (ceil(H * S / 128)). Non-persistent only. The CUDA-graph workspace is sized once for every batch bucket, in 128-row tiles.
• K3 MoE routing prologue:
  • a tensor-core 896-expert router GEMM;
  • a redux-based small-batch noAuxTc top-k;
  • post-topK routing tiers and launch choices;
  • async NVFP4 input quantization on the MoE aux stream for CuteDSL (TLLM_MOE_ASYNC_INPUT_QUANT).
• Routed RMSNorm fused into the routed all-reduce. RMS_NORM gets MNNVL support without a residual. Other strategies fall back.
• MegaMoE startup deadlocks:
  • an FC12 scheduler prologue fix;
  • with fused-comm MEGAMOE_CUTEDSL under TP, the shared-expert all-reduce is ordered after the routed MoE.
• Fused router + latent-down GEMM (trtllm::dsv3_router_latent_gemm_op), SM107, TP MEGAMOE_CUTEDSL.
• Latent down-projection overlapped with routing for CUTEDSL / CUTEDSL_FC12, via an input_ready_event that ConfigurableMoE and the schedulers join at the first read of x. The FC12 op zeroes its counters with one fill.
• KDA standalone QKVG codes released at fusion time. The standalone Q/K/V/G FP8 codes are freed as each projection is fused, avoiding about 22.6 GiB of peak loading memory on a DEP8 context rank.
• No blocking HtoD copies in KDA replay slot updates. The disaggregated-generation admission updates of the KDA replay slots no longer issue blocking host-to-device copies.
• DSpark draft KV page size. The separate draft KV pool of DFlash and standalone DSpark gets its own 32-token page size, so a 128-token target page size no longer breaks the trtllm-gen MLA drafter.
• DSpark MLA drafter defaults to the CuTe DSL block decode (AUTO). About 36% faster in the drafter region.
• Draft KV GPU budget cap. Opt-in TLLM_DRAFT_KV_GPU_BUDGET_GIB caps the draft KV GPU budget on disaggregated generation workers and gives the excess to the target pool. It never goes below the draft's allocator floor and does nothing when unset.
• Helix CP GEN autotuner warmup. Helix CP generation workers run the autotuner warmup on one generation-only batch instead of skipping it. PARALLEL tactics are split across the TP×CP group, so each Helix rank profiles 1/(TP×CP) of the tactics and the group keeps the fastest per shape.
• CuTe DSL MLA decode split-KV planned per request. A one-launch Triton kernel plans each request's split count from the step's KV lengths. The decode ops take an optional block_split_kvs; a request planned onto one split stores its output directly and the reducer skips its rows.
• FP8 H=96 MLA verify groups go to CuTe DSL on SM100/SM103. A non-Helix multi-token FP8-KV MLA generation batch with 64 < H < 128 (DSpark under attention DP keeps all 96 heads per rank) was not in the CuTe DSL FP8 perf table, so it fell through to FlashInfer TRTLLM-Gen, which rejects that head range and failed executor warmup. The Helix H=96 verify-group carve-out now also covers these batches on SM100/SM103. SM107 is unchanged.
• Tests: test_kimi3.py::TestKimiK3::test_gsm8k_nvfp4_dep8 with a gsm8k reference row. It runs with
  KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1 (passed as env_overrides): attention DP replicates the shared-expert and latent
  MoE MLPs, and in BF16 they leave DEP8 no KV-cache room. The reference was measured with this cast on.

Not included:

• BMM/GEMM cubin updates. Main's own export covers K3's SiTU lookups.
• Making FP8 KV the default K3 MLA path. Its hook was removed by [TRTLLM-15033][fix] Revert #17800 FlashInfer CuTeDSL MLA dispatch #18653.
• K3 FP4 (NVFP4 KV) MLA. Main rejects FP4 MLA on hybrid linear-attention models and in disaggregated serving. The K3 FP4 MLA context, Helix, split-KV and CTX-cache work therefore belongs in a separate PR.
• MNNVL RMS_NORM_QUANT_FP8_BLOCKSCALE. Nothing uses it.
• The KDA prefill scratch-pinning fix and the runtime fla dependency. Main has both since [None][fix] Kimi K3 / KDA / MNNVL / DFlash: fixes for defects on main #19988 (_scratch_ct, fla-core in requirements.txt).
• The attention-residual ws_stats double buffer and slot-release __syncwarp. Main has them, plus
  cross-proxy fences, since [None][fix] Kimi K3 attn_res kernels: order the ws_stats reads and fence the slot releases #19829.

Test Coverage

• tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py (incl. test_checkpoint_pair_carries_cute_scale, Rubin)
• tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py (incl. test_cutedsl_fc12_survives_autotune_warmup, SM107)
• tests/unittest/_torch/moe/test_moe_backend.py (incl. test_fc1_prologue_wait_preserves_exhaustion, SM107)
• tests/unittest/_torch/moe/test_kimi_k3_mlp.py:
  • test_kimi_k3_shared_expert_parallel_construction
  • test_kimi_k3_shared_allreduce_ordering
  • test_kimi_k3_latent_down_projection_overlaps_routing
• tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py: priming gate, FC12 warmup preparation, async input quantization
• tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py: 896 experts, top-k equality, fused router+latent op
• tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_rms_norm_fusion_without_residual
• tests/unittest/_torch/attention/test_cute_dsl_mla.py: FP8 scales, split-KV sizing, MLA generation routing for FP8 wide-head verify groups
• tests/unittest/_torch/attention/test_cute_dsl_mla_split_plan.py: per-request split plan
• tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py: draft KV GPU budget cap, separate draft page size
• tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py: DSpark drafter defaults
• tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py
• tests/unittest/_torch/modules/kimi_kda/: test_kda_mtp_decode_cute_parity.py, test_kimi_kda_fused_verify_parity.py (epilogue arguments, unfused fallback, MXFP8 gating, SM107 MXFP8-vs-BF16), test_kimi_kda_bf16_state_pool.py; the replay-only commit in test_kda_mtp_decode_cute_parity.py
• tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py
• tests/unittest/_torch/speculative/hw_agnostic/test_sa.py
• tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py: Helix CP generation-only autotuner warmup
• tests/unittest/_torch/misc/test_autotuner.py: PARALLEL tactics split across the Helix TP×CP group
• tests/integration/defs/accuracy/test_kimi3.py::TestKimiK3::test_gsm8k_nvfp4_dep8

Validation results

K3 disaggregated DSpark, GSM8K (DEP16 context + DEP16 generation, FP8 KV cache, DSpark with 7 drafts)

GPU	Commit	Flexible	Strict	Baseline	Result
SM107	42ddc1b	0.9629	0.9613	0.9629 (pre-rebase baseline)	PASS
GB300 / SM103	19693d7	0.9575	0.9575	0.9621	PASS
GB300 / SM103	ea33adc	0.9613	0.9613	0.9621	PASS

• Pass criterion: the repo's GSM8K hypothesis test against the baseline.
• Before 19693d7, GB300 failed at generation warmup: the H=96 FP8 multi-token MLA verify was routed to trtllm-gen, which rejects 64 < H < 128. Since 19693d7 it runs cute_dsl_mla_decode_fp8.
• SM107 GSM8K on ea33adc is pending.

K3 disaggregated perf gate (SM107, Helix CP8 generation-only, concurrency 16, MTP5) at d5ca963: all 6 gates pass against the baseline.

Metric	Result	Baseline	Gate
error_rate	0.000	0.000	≤ 0.001
output_throughput	98.30	77.44	≥ 73.57
user_throughput	6.148	4.843	≥ 4.600
mean_tpot_ms	5.691	6.555	≤ 7.211
avg_decoded_tokens_per_iter	3.593	3.595	≥ 3.415
device_step_time_ms	23.92	27.10	≤ 29.81

• Generation warmup took 92 s, with no autotuner phase.
• This gate predates the Helix CP GEN autotuner warmup (added after d5ca963). A perf rerun with the warmup is pending, as is a GPU run of test_helix_cp_runs_only_the_generation_autotuner_warmup and of the multi-rank Helix TP×CP tuning.
• The rerun on ea33adc (KDA cluster-barrier fix) is pending.

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.

[lucifer1004]

A data point for the preserve_section_strides change in KimiKDALinearAttention._project_verify_inputs (q/k/v/output-gate sections passed to kda_mtp_decode as strided views of the fused QKVG output):

We run the same change on Kimi K3 + DSpark (7 drafts, 8-token verify) on SM107, 8 GPUs over 2 nodes, with the current kda_mtp_decode kernel (the views pass _require_stride_layout, which only needs a unit last-dim stride):

• nsys: it removes 4 direct_copy kernels plus one small copy per KDA layer before kda_mtp_decode.
• Single-request decode step: -0.66 ms at 2K context, -0.69 ms at 32K, -0.31 ms at 128K.
• Outputs byte-identical to the build with the copies (8/8 probe cells, 120/120 needle-retrieval answers, 41/41 qualification prompts), same acceptance length.

Since it is a self-contained, output-identical win that does not depend on the rest of the Rubin stack, would it make sense to split it out of this draft into a small PR so it can land earlier? I am happy to open one if you prefer.

[lucifer1004]

On the K3 MoE routing prologue: we worked on the other routing path, the scores path where routingCustom::run gets the router logits and does top-k in the kernel. Sharing it here so we don't duplicate work.

For 5..8 tokens with more than 256 experts, that path uses routingIndicesCoopBlockMultiKernel, which runs on one CTA and handles the tokens one after another. With K3 parameters (896 experts, top-16, fp32 logits and bias, EP8 with 112 local experts, tile 8) on SM107 it takes 13.7 us at 8 tokens with PDL, and it runs once per MoE layer in every verify step. We replaced it in that range with an 8-CTA thread-block-cluster kernel: one CTA per token runs the same per-warp top-k / merge / normalization code. CTA 0 gathers the selected expert ids over DSMEM and builds the histogram, offsets and permutation from per-expert token bitmasks. A threshold pre-filter in the merge (the topK-th largest warp maximum is a lower bound for the true topK-th value) brings it to 3.5 us. All output buffers are bit-identical to the existing kernel over 3,200 random cases (5..8 tokens, tiles 3/8/16/32, EP8 local ranges and EP1, PDL on and off, heavy ties). End to end it saves about 0.85 ms per decode step at 2K-128K context, with byte-identical outputs.

Your PR moves K3 to a separate small-batch top-k plus the post-topK permutation (dyn-block at <= 12 tokens). Two questions, so we can decide where this should go:

1. After this PR, does any supported K3 configuration still route through the scores path (for example the TRTLLM-gen MoE backend fed with logits)? If not, the kernel would mostly help other > 256-expert models at 5..8 tokens, and I would send it as a standalone PR for the generic path.
2. At M = 8, how do the redux top-k and the post-topK dyn-block add up on SM107 together? If one fused scores-path kernel at about 3.5 us is competitive with that, it could also be an option for K3.

cc @reasonsolo

[lucifer1004]

A data point for widening is_kda_optimized_supported() to SM107 (_kda_kernels.py): on main it only accepts (100, 103), so on Rubin the KDA verify/decode falls back to FLA.

Same setup as my comment above (Kimi K3 + DSpark, 8-token verify, SM107, 8 GPUs over 2 nodes), only this gate changed, fresh server for each arm:

context	decode step, FLA	decode step, fused KDA	TPS at equal acceptance
2K	24.69 ms	17.96 ms	+37%
32K	25.97 ms	19.17 ms	+35%
128K	29.78 ms	22.95 ms	+30%
393K	39.85 ms	32.91 ms	+21%

So about -6.7 to -6.9 ms per step at every length; acceptance is set by the drafter and the target, not by the verify kernel. Prefill goes through FLA in both arms. TTFT went up by 0.26 s (+2.1%) at 128K and 0.23 s (+0.5%) at 393K; part of that is the per-chunk replay-window write the fused verify adds to prefill (0.8-1.7 ms per chunk), the rest we could not separate from server-to-server variation.

Outputs are not token-identical to the FLA path, as expected from a different kernel: the first generated token matched on 40/40 prompts, and the mean shared-prefix fraction of the greedy continuations against the FLA path was 0.658. The differences are deterministic (two cold starts of the fused path were token-identical on 40/40), and when I scored the divergence points with logprobs from a non-speculative reference run, none of the 18 scorable ones was a high-confidence divergence (largest |logit margin| 1.25 nats). Our accuracy checks did not change.

So +1 for enabling 107 here. If this PR stays in draft for a while, would you be OK with me sending just the one-line gate change as a separate small PR?

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #78046 [ run ] triggered by Bot. Commit: f09d991 Link to invocation

[tensorrt-cicd]

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

[reasonsolo]

/bot run

[tensorrt-cicd]

PR_Github #78474 [ run ] triggered by Bot. Commit: 9df8056 Link to invocation

[reasonsolo]

/bot run

[tensorrt-cicd]

PR_Github #78480 [ run ] triggered by Bot. Commit: 029c8a2 Link to invocation

[tensorrt-cicd]

PR_Github #78474 [ run ] completed with state ABORTED. Commit: 9df8056

Link to invocation

[tensorrt-cicd]

PR_Github #78480 [ run ] completed with state SUCCESS. Commit: 029c8a2
/LLM/main/L0_MergeRequest_PR pipeline #64686 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

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #78510 [ run ] triggered by Bot. Commit: 029c8a2 Link to invocation

[Barry-Delaney]

Re-reviewed 029c8a2. Both previous blockers are resolved. The split-V pair now completes both CTAs cache reads before a cluster barrier, and only CTA 0 writes shared conv/replay caches; verified all ten cache reads precede the first cluster wait and all ten cache stores follow it under the single-writer guard. Draft counts above MAX_NUM_SPEC now disable replay-cache allocation through the shared executor/benchmark gate, retaining intermediate buffers and selecting sequential verification. Executed ten actual fallback-gate cases and both actual verify-dispatch branches. Reviewed the FP8 wide-head MLA carve-out, router gate tests, generation-only Helix warmup and TP/CP tuning partition. Eighteen actual partition groups and two cache-merge cases passed; Python parsing and git diff --check passed. No remaining blocking code finding. GPU race/numerical checks and native distributed tuning were not rerun locally.

[tensorrt-cicd]

PR_Github #78510 [ run ] completed with state SUCCESS. Commit: 029c8a2
/LLM/main/L0_MergeRequest_PR pipeline #64709 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

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #78589 [ run ] triggered by Bot. Commit: 029c8a2 Link to invocation

[tensorrt-cicd]

PR_Github #78589 [ run ] completed with state SUCCESS. Commit: 029c8a2
/LLM/main/L0_MergeRequest_PR pipeline #64745 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

[reasonsolo]

/bot run

[tensorrt-cicd]

PR_Github #78618 [ run ] triggered by Bot. Commit: 029c8a2 Link to invocation

[tensorrt-cicd]

PR_Github #78618 [ run ] completed with state SUCCESS. Commit: 029c8a2
/LLM/main/L0_MergeRequest_PR pipeline #64770 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

[reasonsolo]

/bot run

[tensorrt-cicd]

PR_Github #78626 [ run ] triggered by Bot. Commit: 029c8a2 Link to invocation

[tensorrt-cicd]

PR_Github #78626 [ run ] completed with state SUCCESS. Commit: 029c8a2
/LLM/main/L0_MergeRequest_PR pipeline #64778 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chenfeiz0326]

@reasonsolo @pengbowang-nv: since this merged, one case of tests/unittest/_torch/attention/test_cute_dsl_mla.py fails on main in every attention shard that runs it:

test_mla_generation_routes_wide_head_fp8_verify_away_from_trtllm_gen[103-96-1-True-False-cute_dsl_mla]
AssertionError: assert 'trtllm_gen' == 'cute_dsl_mla'

It hit DGX_B200-PyTorch-4 and DGX_H100-PyTorch-5 on L0 #64787 (#20250, which doesn't touch attention). The test is host-only, so it fails the same way on any GPU.

This is a semantic merge conflict between two PRs:

• This PR's tested head (029c8a2) was based before [https://nvbugs/6691772][fix] Update trtllm-gen fmha autotuner to enable 96 head MLA #18599. In that tree the FP8 perf table had (96, 1): 1, so the single-token H=96 case reached cute_dsl_mla at the test's batch_size = 4.
• [https://nvbugs/6691772][fix] Update trtllm-gen fmha autotuner to enable 96 head MLA #18599 (merged 10-16) changed that entry to (96, 1): 8 and removed the H=96 mixed-batch exemption.
• On main the 4-request single-token case is now below the floor and falls through to trtllm_gen. The test's comment on that row ("Single-token H=96 is unchanged (measured perf-table entry)") no longer holds.

The other 9 parameter rows still pass. Which side is right is your call. Either the row should expect trtllm_gen at batch 4, or use batch_size >= 8, or the table entry should go back to 1. I haven't changed anything. This is just a heads-up, since every PR's CI will fail on this case until it's fixed or waived.