[TRTLLM-14703][fix] 1st cherrypick [combine MR99+MR100+MR102 on latest feat (with MR100 dispatch wiring)], 2nd cherrypick [deferred-finalize scales fix, fused-epilogue batch gate, hardening]

[moraxu]

@coderabbitai summary

Description

Cherry-picks two internal (squash-merged) MRs onto feat/kimi_k3. Each commit's original MR description follows verbatim.

NOTE: The third commit bb62b67 is just pre-commit application.

Commit 574265eecd — [None][perf] Kimi K3: combine MR99+MR100+MR102 on latest feat (with MR100 dispatch wiring)

Goal

Combine the three pending Kimi K3 perf MRs on top of the latest feat/golden_prairie (493ca9d724, which already includes MR !97 / tali/k3-night-opt-0726), make the combination actually functional, and validate performance + accuracy with the in-repo sweep scripts (examples/kimi_k3/perf_sweep/). This MR supersedes merging !99 / !100 / !102 individually.

Contents (cherry-picks + one fix)

1. !99 (rosong) — fused down-proj+MXFP8-quant kernel, fused MoE finalize + MNNVL AR + RMSNorm, MoE TP4+EP4 measurement notes. Cherry-picked as 260d001385.
2. !102 (sephw) — FP8 KDA fused prefill QKV (packed qkvg_proj reuse in _forward_prefill). Cherry-picked as 301e885a9c.
3. !100 (sephw) — KDA decode kernel selection heuristic (SM103 compact-heads selector). Cherry-picked as 042194de21, plus a new wiring commit bb3f29d984: as pushed, !100 only added shouldUseCompactHeads() and tests — invokeKdaDecode still used the legacy H==2 rule and the compact launcher was pinned to dispatch_kda_decode_layout<true, 2> (its own dispatch tests fail on SM103, and naive wiring would be silently wrong for H≠2 because kFixedHeads is baked into the static-layout path). The wiring commit adds the SM103-exclusive selector call and instantiates the compact path for every supported head count.

Analysis notes (activation conditions, overlap with the !97 optimizations, expected per-config gains): agent-runs/0726-1455-mr-analysis/mr-analysis-report.md in the project workspace. Key expectations on this baseline: !99's fused finalize+AR+RMSNorm is the main additive piece (TEP-only, decode ≤16 tokens); !99's down-proj fusion is inert under default KIMI_K3_FP8_WEIGHT_READ=1; !102 is prefill/TTFT-only and inert at TEP16 (recipe sets KIMI_K3_FP8_WEIGHT_READ_KDA=0); DEP16 is expected ~unchanged.

Validation plan (in flight)

• Editable GB300 build (OCI-AGA)
• MR100 dispatch unit tests + MR102 packed-prefill tests on GB300
• 17-point perf sweep via submit_perf_sweep.sh (tep16 / tep8 / dep16-lo / dep16-hi, simultaneous submission) vs the !97 baseline sweep
• GSM8K accuracy sweep via submit_acc_sweep.sh (band 96.5±0.5; main numerics change = !99's bf16 finalize accumulation)

Results + recipe recommendations will be posted as a comment when the runs complete.

---

---

Commit 2e4c343922 — [TRTLLM-14703][fix] Kimi K3: unblock 574265eecd - deferred-finalize scales fix, fused-epilogue batch gate, hardening

Goal

Unblock 574265eecd (stacked MR: targets its source branch tali/k3-feat-3mrs-0726). Fixes the two blockers from the 574265eecd close-out — the tep8 GSM8K=0.00 hard blocker and the tep16 c4/c8 decode regression — plus merge-safety hardening identified during review.

Root cause (tep8 GSM8K=0.00)

With separated routing (TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING=1, tep8-only recipe), the fp4/mxfp4 block-scale MoE runners consume the caller's topk_weights through expert_weights_ptr and the routing cubin treats the weights buffer as input-only, so the locally allocated expert_weights tensor is never written — yet the do_finalize=False path returned it. The Kimi fused finalize+MNNVL AR+RMSNorm epilogue then applied uninitialized memory as per-(token,k) routing scales: silent garbage, no crash. tep16 (in-kernel routing) and dep16 (fusion gated off) were unaffected, matching their ~97 scores.

Contents (oldest first)

1. f1828a58ed [fix] Return host topk_weights from the block-scale MoE runners when do_finalize=False (both mxFp4BlockScaleMoe.cpp and fp4BlockScaleMoe.cpp; the nvfp4 file has the same latent defect reachable upstream via EPLB) + backend-level regression test.
2. 43981ead8b [perf] Gate the fused finalize+AR+RMSNorm to decode batch <= 2 (measured-win region; c4 attribution confirmed by the partial no-fusion A/B) and expose TLLM_K3_FUSED_MOE_FINALIZE_AR_RMS_MAX_TOKENS for rebuild-free sweeps.
3. 74ecff544b [chore] Sweep containers: PYTHONNOUSERSITE=1 (user-site torch 2.13 ABI poisoning) + guarded HF_HOME for the acc sweep.
4. 63f58901ca [fix] MNNVL fused-finalize input validation + TP{2,4,8,16}/bf16-fp16 capability gating with graceful fallback + rejection tests.
5. bb67ec734d [fix] dsv3 fused A-GEMM dense-layout/16-byte-alignment enforcement (silent wrong-memory hazard) + tests.
6. ddc3d2368c [test] Off-SM103 KDA selector test expects legacy compact dispatch (test was wrong at 574265eecd head; fails on H100/B200 runners).
7. 6eb7751ae6 [fix] mx MoE runner fake signature: missing trailing tune_max_num_tokens/use_dp (TypeError under FakeTensorMode).
8. 6750f6d5d2 [chore] Env-knob safe parsing (clamp [0,16], fail-fast on non-integer) + W4A8_NVFP4_FP8 regression-test coverage.
9. 7d1fa67c24 [test] Supported layer shape for the fp8-activation deferred-scale test param (the fp8-act cubins ship no configs for the small default shape); clean skip when a shape has no kernel config.
10. ed059bf872 [chore] Point the GSM8K acc sweep at the local dataset snapshot via --dataset_path (offline compute nodes; the hub-id load path needs a materialized arrow cache the snapshot does not provide).
11. 4f1cfd7e24 [chore] Give the perf-sweep benchmark client the same cache redirects as the server (flashinfer creates HOME-based dirs and HOME is unwritable in the client container).
12. 9f9cb9493f [chore] Trim the fused-finalize bound's code comment to the knob's contract; the motivating sweep numbers and kernel-mechanism analysis moved to the internal GitLab MR discussion.
13. 57d1a4b91d + 0679fd508b precommit and its revert: the hook run reformatted only pre-existing code (none of this MR's own hunks — 97% of it the one-time whole-file reformat of modeling_kimi_linear.py), so it was reverted to keep the MR diff reviewable. Net diff of the pair is empty; the whole-file formatter pass stays deferred to upstreaming, as listed above.

Deferred by review (separate MR): tactic-independent deferred GEMM2 capacity + fake row-math package (torch.compile-only; needs its own hardening), whole-file formatter churn.

Validation matrix (complete)

Metric of record for perf rows: tps/user = 1000 / median TPOT (ms). Rows 3-11 were initially scored against the !97 baseline sweep TRTLLM-sweep-Jul-26-commit-82f9780e; rows 13-14 re-measure the baseline same-moment because that comparison proved unreliable (see verdict). Concurrency 3 is deliberate: it straddles the batch<=2 fusion gate.

#	Stage	Mode / shape	Extra env	Concurrencies	Pass criteria	Result
0	Wheel rebuild + install + deps (SM103)	1 node / 4 GPUs	—	—	wheel builds clean, 103-real	PASS (incremental rebuild, both changed C++ files recompiled)
1	GB300 unit tests (deferred-scale regression x3 quants, MR100 KDA, MR102 prefill, MNNVL rejection, dsv3 A-GEMM)	1 node / 4 GPUs	—	—	all pass	PASS (regression 2/2 + 1 shape-skip, KDA 19/1skip, prefill 3, dsv3 14, MNNVL capability 7)
2	GSM8K dataset preflight in container	1 node	HF_HOME -> shared cache	—	1319 test / 7473 train records visible	PASS (offline load via --dataset_path snapshot)
3	tep8 GSM8K accuracy	2 nodes / 8 GPUs	default (bound 2)	—	96.5 +/- 0.5 (was 0.00); full 1319 responses	PASS: 96.44 (strict & flexible, 1319/1319)
4	tep16 perf, fused bound 2 (default)	4 nodes / 16 GPUs	—	1 2 3 4 8 16	c1/c2 retain ~+9/+5%; c4/c8 back to >= parity	c1 +8.0% c2 +3.2% c4 +1.6% PASS vs stale baseline; c8/c16 deficit resolved by row 13 as baseline drift
5	tep16 perf, fused bound 16 (attribution A/B)	4 nodes / 16 GPUs	..._MAX_TOKENS=16	1 2 3 4 8 16	== row 4 at c1/c2; c4/c8 regression reappears	c1 identical (+0.0%); fused effect small everywhere (c4 +1.5%, c8 +3.2%, c16 +0.4%, c3 -5.5%); c2 point anomalous (-9.3% on identical config); old c4 fusion penalty NOT reproduced -> c8/c16 deficit is NOT the fusion
6	tep8 perf, fused bound 2 (default)	2 nodes / 8 GPUs	—	1 2 3 4 8 16	vs row 7: differs only at c1/c2	PASS: c1 -0.1% c2 +0.7% (fusion neutral at TP8); vs base +1.6/+6.0/+10.8/+1.6/-2.1; old c1 -15.5% NOT reproduced
7	tep8 perf, fusion disabled (TP8 gating decision)	2 nodes / 8 GPUs	..._MAX_TOKENS=0	1 2 3 4 8 16	negative control; converges with row 6 at c3+	PASS: c3+ within +/-1.6% of row 6 (config-identical there)
8	dep16-lo perf	4 nodes / 16 GPUs	—	16 32 64 128 256	parity vs !97 baseline	recorded: 38.76/32.64/29.48/23.26/15.84 tps/user (baseline absolutes not in workspace; original campaign: c16 +4.1%, c>=32 parity)
9	dep16-hi perf	4 nodes / 16 GPUs	—	512 1024	parity	recorded: c512 10.98, c1024 6.63 tps/user (baseline absolutes not in workspace; original campaign: parity)
10	tep16 GSM8K	4 nodes / 16 GPUs	—	—	~97.0	PASS: 96.44 (band 96.5+/-0.5)
11	dep16 GSM8K	4 nodes / 16 GPUs	—	—	~97.0	PASS: 97.04
12	(conditional) nsys KDA compact-vs-many-heads at c4/c8	4 nodes / 16 GPUs	—	4 8	only if row 4 c4/c8 do not recover	NOT NEEDED — row 13 resolved c8/c16 as baseline drift
13	same-moment baseline A/B: fresh 493ca9d724 wheel vs MR, simultaneous launch	2x 4 nodes / 16 GPUs	—	1 2 3 4 8 16	isolate MR effect from baseline staleness	DONE: MR vs same-moment base = c2 -1.5% c3 -1.5% c4 -1.9% c8 +3.7% c16 +0.2%; base-today itself -7.6% at c8 vs its stale numbers -> c8/c16 "regression" was drift. c1 point anomalous (-14.8% vs 3 same-day MR repeats at ~76.9) -> row 14
14	c1/c2 targeted rerun of the same-moment pair	2x 4 nodes / 16 GPUs	—	1 2	settle the anomalous c1 point	DONE: c1 base 75.82 / mr 76.68 -> MR +1.1%; MR c1 repeats 76.91/76.92/76.68 confirm the 64.54 point was bad. c2 -3.3% (parity within noise)

Rows 0–2 are prerequisites; rows 3–5 are the first wave (decides the fix); rows 6–7 decide TP8 fusion gating; rows 8–11 are full-qualification regression coverage; rows 12–14 resolve baseline comparability.

Verdict

• Accuracy: fully fixed. tep8 GSM8K 0.00 -> 96.44, tep16 96.44, dep16 97.04 (all full 1319-response evals, band 96.5 +/- 0.5).
• Performance: neutral-to-positive, no regressions. Against a freshly built merge-base (493ca9d724) wheel launched at the same moment, the branch measures c1 +1.1%, c2/c3/c4 parity within noise (-1.5..-3%), c8 +3.7%, c16 +0.2% at tep16; tep8 shows no regression and the fusion is neutral there.
• Measurement-drift finding (important for future sweeps): the original campaign's headline deltas — tep16 c1 +9.4%, c4/c8 -5.5/-9.6%, tep8 c1 -15.5% — did not reproduce and are attributable to cross-campaign drift: the baseline code itself measures -7.6% at c8 today vs its own Jul-26 numbers, and individual low-concurrency points can go bad by 9-15% between allocations. Same-moment paired launches with repeats are the only reliable comparison on this cluster.
• The fused finalize+AR+RMSNorm epilogue stays enabled with the batch<=2 default (small c1 win, accuracy-neutral, rebuild-free tunable for future sweeps).

Test Coverage

See the "Validation plan" / "Validation matrix" sections above: GB300 unit tests (deferred-scale regression, KDA dispatch, packed prefill, MNNVL rejection, dsv3 A-GEMM), GSM8K accuracy at tep8/tep16/dep16, and tep/dep perf sweeps.

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.

[moraxu]

/bot run

[tensorrt-cicd]

PR_Github #62879 [ run ] triggered by Bot. Commit: 7ebd907 Link to invocation

[tensorrt-cicd]

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