[chenfeiz0326]

@coderabbitai summary

Draft, experimental, not for merge. Off by default: TLLM_JIT_AHEAD=1 turns it on. No LLM-args change. Report with per-model data (NVIDIA internal).

Design

• Kernel selection is untouched. Helpers only compile; Triton's autotuner, DeepGEMM's heuristic and TRT-LLM's autotuner still choose. The serving loop is unchanged: no per-iteration step.
• Isolated from the runtime code. The model loader, the engine's warmup and its serving forwards only call in at fixed points; jit_prefetch_engine.py (EngineJitAhead) runs each startup phase, and each kernel family's module decides, from the model and the engine limits, which kernels to build or tune (start_model_builds, start_weight_load_builds, start_warmup_builds, register_providers).
• Warmup is not extended. It skips the passes that only compiled, whose kernels the helpers now build or that built nothing new: the DeepGEMM swap-AB sweep and the DSA metadata sweep (DeepGEMM helpers; on SM100 the metadata kernel has no batch-size template, so the sweep rebuilt one kernel per next_n), the multi-sequence one of [None][perf] Close Mamba hybrid warmup gap in autotuner warmup #16177's two Mamba passes (Triton helpers; with tune ahead, the other one, whose autotune moves to the end of warmup), the gen-FMHA NVRTC grid on disaggregated context workers (they run no generation request), and repeated MXFP8 graph tuning (each graph shape is tuned once and reuses its tactics). A pass whose helpers are not running still runs; TLLM_WARMUP_JIT_COVERAGE=1 runs them all.
• Dead or idle helpers cost nothing. One that exits releases its builds, so a waiting launch builds in place, and a warmup pass skipped for its kernels runs again. With every helper of a kind killed as it starts, a server is left with what main's warmup leaves: on Nemotron Nano the same 11 kernels compiled while serving as with the switch off (none with live helpers), on DeepSeek-V3.2 none. Idle helpers stop once serving runs.
• With the switch off, two changes remain, both against MiniMax-M3 recompiling in pass 2: its MSA KV-pool views are marked dynamic (no decoder recompile), and its MSA decode kernels do not specialize on the block table's stride (no decode-kernel recompiles).
• Switches: TLLM_JIT_AHEAD=1; TLLM_JIT_TUNE_AHEAD=0|1 overrides tune ahead (=1 alone: in-process, no helpers); TLLM_JIT_STATS=1 logs every compile or autotune left on the executor, and tests/scripts/jit_prefetch_coverage.py turns that log into a coverage report. TLLM_JIT_PREFETCH is a deprecated alias. Knobs: TLLM_JIT_PREFETCH_{,DG_,CD_,FMHA_}WORKERS, TLLM_JIT_PREFETCH_WAIT_TIMEOUT_S, TLLM_JIT_TUNE_AHEAD_THREADS, TLLM_JIT_CUTE_DSL_DIR.

Performance Comparison

GB300, every server from empty JIT caches. Before = switch off, after = TLLM_JIT_AHEAD=1.

• MiniMax-M3 ctx TP2 on fd2095d, both arms at once on one node: the off arm built 2 kernels while serving (two ~99 s stalls) and served 54.0K input tok/s over the 420 s window; with TLLM_JIT_AHEAD=1, none and 110.7K tok/s.
• DeepSeek-V4.1-Flash TP4 on a1a4692 vs 8427b4b, same node, both on: the CuTe DSL MoE grouped GEMM tactics (12 per rank) and all but one of the 13 top-k and CSA2 variants now build in helpers, so pass 1 compiles 14 s less on each executor: pass 1 144.5 → 130.1 s, ready 380.8 → 370.6 s. On 474dc80 FlashInfer's CuTe DSL MXFP8 GEMMs (4 per rank, 2.4 s) also build in helpers; pass 1 moves within its run-to-run noise.
• DeepSeek-V3.2 TP4 on 27abdf9 vs df55628 (TRTLLM MoE on FlashInfer, MTP3), same node: FlashInfer's TRTLLM-Gen fused MoE launcher (15 min of nvcc per rank) and RoPE now build in helpers from model construction; ready 1332 → 1177 s and, in reverse order, 1263 → 1217 s. Serving is unchanged (699.4 vs 696.6 tok/s, TPOT 3.48 vs 3.49 ms) and no longer compiles CuTe DSL's occupancy probe (3–4 per server).
• Nothing new compiles while serving: after ready, no server measured here compiled or autotuned a kernel on its executor. That includes Kimi K3's 16 disaggregated context ranks, which skip the gen-FMHA grid: each built only its warmup forward's own FMHA kernel (6–7 s, before ready) and none in 46 min of serving at concurrency 360.
• Kernels are unchanged: per-step GPU time agrees within ±0.4 % on matched steps. Right after a cold start the PR serves faster, as nothing compiles: Super's first 160 requests (be3a227, ISL/OSL 1024, concurrency 32), 6,952 / 6,948 → 7,137 / 7,214 tok/s, TPOT 10.19 / 10.19 → 9.95 / 9.85 ms; the off arms compiled 15 kernels across ranks during them.
• Cost: CPU only. Helpers run at nice 19 (FlashInfer modules a layer needs at process priority). Weight loading runs slower while they build, 13–26 s on MiniMax-M3 and 22–36 s on Nemotron 3 Super; the ready times include it.
• Still built on the executor in a cold start, and why not in a helper: trtllm-gen FMHA's NVRTC kernels (19–45 s of attention_jit), which its prebuilt library compiles and loads in one call into an in-process module cache, with no compile-only or load-from-file entry; Kimi K3's low-M split-K GEMM (~36 s), which FlashInfer builds from tensors it allocates on the GPU (a helper prototype matched bitwise but stalled in the 4-node server); MiniMax-M3's MSA sparse-attention kernels (~3.5 s per rank), compiled inside the vendored fmha_sm100 package's plan/run; each sparse layer's first top-k variant (0.6 s), compiled while its siblings are queued; and about 3 s per rank of small CuTe DSL glue kernels on DeepSeek-V4.1 (CSA2's quantize-scatter, scale-finalize and compaction kernels, the QK norm), each of which needs its own builder and plan, as CSA2's candidate publication has.

🤖 Generated with Claude Code

[litaotju]

I am not sure if only enumerating the batch size is good enough to trigger all dsa kernels needed.

[litaotju]

same here.. how do we make sure this logic can trigger all needed kernels?

[xxi-nv]

FYI, @rosong11 is working on the removal of FI from TRTLLMGen.

[xxi-nv]

On fused_moe_deepgemm.py#L740-L748:

Question on the design here: why does the masked grouped GEMM only report to done_observer, without a pre-launch wait like the other JIT paths?

Every other path waits for a helper that is already building the variant it needs:

• FP8 swap-AB GEMM: note_launch -> launch_observer bumps the queued request to urgent and waits for its cubin (jit_prefetch.py#L1267-L1283, torch_custom_ops.py#L2206-L2222)
• CuTe DSL: waits if a helper has started the variant (jit_prefetch.py#L1653-L1667)
• Triton: same policy in cache_hook (jit_prefetch.py#L2480-L2502)

Here the only hook is done_observer, and it fires after the launch, so it can only count a compile. The comment at jit_prefetch_deep_gemm.py#L437-L441 states this, but not why. As far as I can tell, if a warmup forward reaches a (num_groups, N, K) that a DeepGEMM helper is still compiling, the executor compiles the same kernel again. The result stays correct, since DeepJIT writes its cache atomically, but the second compile is wasted host time on the executor thread.

Is skipping the wait intentional? For example:

• the window is too small to matter, because DeepGemmMoEProvider queues these requests as soon as prefetch starts; or
• _queued_deep_gemm_spec only resolves FP8 Linear specs (spec_for), and mapping -num_groups to moe_spec_for was judged not worth it.

If it is intentional, could you add a short comment stating that a duplicate in-place compile is accepted here, and why? If not, a note_launch(-num_groups, n, k) before the launch, plus a _queued_deep_gemm_spec that resolves MoE specs, would bring this path in line with the others.