[zhaoyuanh-nvidia]

Description

KDAReplayState.commit_pending_drafts (modules/kimi_kda/cache_manager.py) folds a slot's pending accepted drafts and its conv window into the persistent KDA states. Only suspend_request calls it (in KimiK3HybridCacheManager and Glm5NextCacheManager), and no warmup forward suspends a request. So the first suspend in serving JIT-compiled the kernels it launches, inside one scheduler step, while every request in flight waited:

• the replay-only kda_mtp_decode CuTe DSL variant, or _replay_kda_token_records_kernel when the verify kernels keep draft records;
• _store_kda_replay_conv_window_kernel.

No compile key of these kernels follows the batch: a commit always launches on one slot over the same per-layer pool views. In the Kimi K3 MXFP4 AgentX runs on GB300, _store_kda_replay_conv_window_kernel compiled exactly one variant per node, with the same Triton cache key in every run, and never during startup. At the same moment every rank logged one replay-only kda_mtp_decode compile. In a TP8 run with the host KV tier (FP8 KV, 32 sessions), the step that made the first suspend took 6.4 s instead of about 0.15 s. In three earlier runs of that point, the streaming gaps around these compiles were 5.4–29.9 s.

This PR:

• adds KDAReplayState.prewarm_commit_kernels, which bind calls once the replay views are bound, before any request holds a slot. It seeds slot 0's replay conv window from the live pool (copy_kda_replay_conv_window, as a prefill does), then commits zero drafts of slot 0 on every layer. The commit goes through _commit_drafts, the loop commit_pending_drafts now uses, so the views and argument types and therefore the compile keys are the same. Committing zero drafts after the seed leaves slot 0's SSM and conv states bit for bit as they were;
• has a prewarm failure only log a warning, like the decode-kernel prewarm hooks. The compiles then move back to the first suspend.

do_not_specialize would not help: each kernel compiles once per process, on first use.

main does not have this suspend-time commit yet, so no main PR is needed. The fix goes with the feature branch.

Test Coverage

• Regression test: test_v2_kda_replay_suspend_compiles_no_commit_kernel (replay caches and draft records) builds a V2 manager at Kimi K3's per-rank KDA shapes, forgets the process's earlier compiles of the commit kernels, suspends a request with two drafts pending, and asserts the suspend compiled nothing. Without the source change both cases fail: the suspend compiles 2 kernels. With it, they pass, and bind compiled 2.
• States unchanged: test_v2_kda_replay_commit_prewarm_leaves_the_slot_states (both modes) fills slot 0's SSM and conv states and its replay buffers with random values, runs the prewarm, and checks the states are bit-identical and the slot's pending-draft count is untouched.
• Affected tests on GB300: every unit-test file that builds a KDA replay manager or commits drafts: test_mamba_cache_manager.py, test_mamba_model_managers.py, test_mixed_precision_kv_cache.py, test_kimi_kda_fused_verify_parity.py, test_kda_mtp_decode_cute_parity.py, the K3 KDA modeling_v2 tests, test_modeling_v2_claims.py, test_glm_kpool.py, test_glm5_next_contracts.py and test_kda_mamba_transfer.py. With the fix, 1,287 passed and 6 skipped; at the base, 1,283 passed with the same 6 skips. The difference is the 4 new cases. Pre-commit's ruff, ruff-format and codespell checks pass.
• Serving: Kimi K3 MXFP4 with DSpark (K=5), forced acceptance length 3.44, InferenceX AgentX trace replay with a 30-minute measured window. TP8 with the host KV tier, FP8 KV, 32 sessions, one 8-GPU GB300 replica, on the 338b563 build. The run had this change plus [TRTLLM-16624][fix] Do not specialize the fused strict-accept kernel on batch-dependent arguments #22197; the baseline run had only [TRTLLM-16624][fix] Do not specialize the fused strict-accept kernel on batch-dependent arguments #22197.
  • _store_kda_replay_conv_window_kernel compiles after startup: 2 → 0. Both nodes now compile it during startup, under the same Triton cache key the earlier runs compiled mid-serving. Every rank compiles the replay-only kda_mtp_decode variant during startup instead of at the first suspend.
  • Serving time stalled on compiles: 0.59% → 0.25%. The rest is two _rms_norm_gated_fwd_multirow_kernel compiles, fixed separately. The baseline's two steps of 6.5 s and 7.6 s at its first suspend are gone.
  • 4,752 tok/s/GPU: 4,736 in the baseline run, and 4,743–4,858 in three earlier runs of 338b563. p90 TTFT 13.8 s (15.7 s), acceptance length 3.44, 0 errors.
  • Startup to healthy: 856 s vs 780 s, on different nodes. The prewarm itself takes about 5 s.
  • This run's polled iteration stats show no paused request, so they do not show a suspend on their own. A rerun of the same point with a debug-only log line in commit_pending_drafts (not part of this PR) logged one suspend-time commit on every rank, with 15–16 requests in flight. It compiled nothing; the steps around it took 0.12–0.64 s, as usual at that load, where the baseline's first suspend took 6.4 s. The rerun had no late Triton compiles at all, 4,749 tok/s/GPU, acceptance length 3.44 and 0 errors, and reached healthy in 721 s.

PR Checklist

• After merge, verify a Kimi K3 serving run on feat/helix-agg that suspends requests compiles no _store_kda_replay_conv_window_kernel or replay-only kda_mtp_decode variant after startup.