feat: wire flashinfer autotuner - #820
Conversation
|
Waiting for flashinfer-ai/flashinfer#4237 |
|
The autotuner increases the server startup time by 3~5 mins. @zhyncs |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2d16982c40
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self.prefill_graph.warmup_eager(self.forward_step) | ||
|
|
||
| self._autotune() |
There was a problem hiding this comment.
When chunked_prefill_size is greater than the 8192 default, this ordering runs warmup_eager() before _autotune() calls set_autotune_max_num_tokens(). All of the MoE wrappers changed in this commit now read get_autotune_max_num_tokens(), so the pre-autotune warmup executes a large prefill with a stale 8192 tune_max_num_tokens, producing the wrong FlashInfer bucket mapper/cache key or fallback for that first large shape; set the process max before any prefill warmup that can call FlashInfer kernels.
Useful? React with 👍 / 👎.
| with autotune(), maybe_inference_mode(): | ||
| ctx = self.prefill_graph.make_dummy_batch(num_tokens, self.forward_step) |
There was a problem hiding this comment.
tokenspeed_kernel.ops.tuning.autotune() is explicitly a no-op when flashinfer.autotuner cannot be imported, but this block still runs a full chunked_prefill_size model forward after the existing best-effort warmup. On vendor-neutral or non-FlashInfer deployments that means no cache can be populated, yet startup now pays an uncaught full-prefill allocation/forward that can OOM; return before the dummy forward when the tuning backend is unavailable.
Useful? React with 👍 / 👎.
Replace the freeze-once lazy autotune lifecycle with a single explicit pass: ModelExecutor runs one dummy extend forward at chunked_prefill_size inside an autotune() context, before any CUDA graph capture. flashinfer enumerates every smaller shape bucket from that one forward, so no decode-sized pass is needed; and a captured graph records the tactic chosen while it was recorded, so tuning must precede capture to have any effect on a replay. ops/tuning.py now exposes autotune(), set_autotune_max_num_tokens() and get_autotune_max_num_tokens() in place of autotune_frozen() / freeze_autotuning(), whose default-permissive state could tune during serving. The per-call gates in trtllm_mxfp4 and cutlass_unquant are gone, and all eight flashinfer MoE sites pass the process-wide constant rather than a per-batch next_power_of_2(x.shape[0]): flashinfer builds the serve-time bucket mapper from that value, so it has to stay constant for lookups to resolve to the bucket they were tuned at. CudaGraphWrapper and PrefillGraph no longer capture in __init__; the executor calls capture() explicitly once tuning has run. make_dummy_batch is public and now sets gather_ids, so a full model forward can run its logits tail. Also annotate @contextmanager returns as Generator instead of the deprecated Iterator form. Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… shape The fused NVFP4 FC1 kernel picked its mma tiler and cluster shape from a single threshold on M, sampling two of the configurations the vendored kernel supports. Wrap it in a TunableRunner so the autotuner selects per shape bucket instead: a tactic is an (mma_tiler_mn, cluster_shape_mn) pair, -1 keeps the previous heuristic as the mandatory untuned fallback, and get_valid_tactics filters candidates through can_implement. An nsys sweep at the two production shapes (DSv3/Kimi-K2.5 TP4 shared and dense MLP, cold L2) puts the win at up to 9.6% for the shared expert below M=128, where the heuristic's 1-CTA branch loses to a 2-CTA tile, and 2-4% at large M, where a 256-wide N tile wins on both shapes. The tuning config is a module-level constant keyed off the observed M, so it needs no process-wide token count and joins the existing startup window with no runtime wiring. tile_n=64 is excluded from the candidates: can_implement() accepts it but the kernel then writes roughly three quarters of the output. All nine tile_n=64 configurations disagree with the fallback while all twelve kept ones are bitwise identical to it, which is what the new tactics-agree-with-heuristic test asserts. cluster_m is derived as mma_m // 128 rather than swept -- can_implement requires 1 for a 1-CTA tile and an even value for 2-CTA, and cluster_m=4 won no cell in the sweep. Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The freeze switch (autotune_frozen/freeze_autotuning) and the MXFP4 SiTU lazy in-situ tuning pass existed to keep serving-time tuning survivable. With the startup autotune window, serving-time tuning is impossible by construction -- the window is closed outside engine startup -- so both are dead complexity, and the lazy pass carried the TP rank-divergence hazard the window was built to remove. - ops/tuning.py: drop autotune_frozen/freeze_autotuning and the tuning-cache-active flag. Pre-swept table loads are fire-and-forget cache seeding: covered buckets become cache hits the window skips automatically, everything else is tuned by the window. Logger calls converted to f-strings. - trtllm_mxfp4.py: drop the lazy full-ladder pass and the fixed SITU_TUNE_MAX_NUM_TOKENS; the op keys its buckets from get_autotune_max_num_tokens() like every other tunable op. - moe_tactic_sweep.py: derive the sweep tune_max from ops.tuning so swept tables key the same bucket ladder as the runtime (identical for chunked_prefill_size <= 8192; resweep above that). - model_executor/prefill_graph/kimi_k3: remove the freeze call and stale wording; delete test_autotune_freeze.py with the API it tested. Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 39d322c3d2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| ib = self.input_buffers | ||
| tic = time.time() | ||
| with autotune(), maybe_inference_mode(): | ||
| ctx = self.prefill_graph.make_dummy_batch(num_tokens, self.forward_step) |
There was a problem hiding this comment.
When a model/backend's dummy extend batch is not covered, make_dummy_batch() or the dummy full forward can raise the same NotImplementedError/AttributeError/KeyError/RuntimeError that the prefill warmup and capture paths intentionally catch to fall back to eager. This new unguarded autotune pass runs after that best-effort warmup, so those deployments now fail ModelExecutor initialization instead of starting with heuristic kernel choices; skip or catch this pass when the dummy-forward seam is unavailable.
Useful? React with 👍 / 👎.
…ntext length make_dummy_batch fabricated a single request of chunked_prefill_size tokens, which is illegal when that exceeds max_model_len: every per-request structure is sized for the context length, and DSv4-Flash's fp4 indexer paged gather indexes its block table (17 columns at max-model-len 4096, 4:1 compressed, 64-token pages) with the fabricated sequence span (32 columns at 8192 tokens) and no bounds check -- an OOB read whose garbage page ids fault or stay silent depending on what the allocator left past the table. Both the eager prefill warmup and the startup autotune window built this batch; the warmup's broad RuntimeError guard then swallowed the AcceleratorError, so the dead context resurfaced inside the autotune window and pointed the blame at the autotuner. Split the dummy batch into ceil(num_tokens / context_len) requests, each within the context length -- gather_ids, positions, and all per-request metadata now match how a real multi-request extend batch is built -- and re-raise AcceleratorError from warmup_eager so a CUDA fault fails the boot at its true site. Validated on DeepSeek-V4-Flash TP4/EP4 + MTP with chunked-prefill 8192 and max-model-len 4096 (previously an IMA at startup on every attempt): warmup and tuning complete, server serves short and 3k-token prefills. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bbb696614f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| """ | ||
| ib = self.input_buffers | ||
| max_req_tokens = max(1, int(self.config.context_len)) | ||
| bs = max(1, -(-num_tokens // max_req_tokens)) |
There was a problem hiding this comment.
When chunked_prefill_size exceeds context_len * input_buffers.max_bs (for example --max-model-len 512 --max-num-seqs 1 with the default 8192 chunk), this split fabricates ceil(chunk/context_len) dummy requests even though the buffers were allocated only for max_num_seqs per rank. The subsequent req_pool_indices_buf[:bs] / seq_lens_buf[:bs] copies then see a shorter destination slice and raise during startup warmup/autotune, although the scheduler could never produce such a real batch; clamp the dummy token window or bs to the allocated request capacity.
Useful? React with 👍 / 👎.
… in nvfp4_gemm_swiglu The CuTeDSL StaticPersistentTileScheduler ceil-divides the CTA grid by the cluster shape but computes each CTA's tile coordinate as cluster_coord * cluster_shape + cta_id_in_cluster with no extent check, and validity is only tracked per cluster. When cluster_n does not divide the N tile count, the trailing CTAs of the last cluster compute a phantom tile and store it over real output rows and scale factors. The vendored can_implement() checked n % tile_n but not this, so the autotuner could select a corrupting tactic -- on Qwen3.5-35B-A3B-NVFP4 (dense MLP n=512, two 256-wide N tiles) the cluster_n=4 tactics win the profile and serving generates garbage, failing the runtime-2gpu CI generation test. Upstream TensorRT-LLM main has the same gap. Add the tile-count divisibility check to can_implement(), and replace the m-keyed heuristic fallback with the always-safe (128,128) tile on a 1x1 cluster; the former heuristic tiles remain in the tactic sweep, so tuned deployments still get them per bucket where they win. A semantic sweep (through-FC2 comparison against the unfused reference) across four shapes, four M values, and all twelve tactics separates good from bad cells exactly on this predicate. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2e22df498b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| logger.info(f"Kernel tuning with a dummy prefill of {num_tokens} tokens") | ||
| ib = self.input_buffers | ||
| tic = time.time() | ||
| with autotune(), maybe_inference_mode(): |
There was a problem hiding this comment.
When world_size > 1, this startup window lets each rank profile FlashInfer tactics independently. FlashInfer 0.6.16 documents set_autotune_process_group as the mechanism that all-reduces per-tactic timings so every rank picks the same tactic, and warns that leaving it unset can let timing noise make ranks diverge and deadlock symmetric-memory/NCCL allocation (see the set_autotune_process_group docstring in https://raw.githubusercontent.com/flashinfer-ai/flashinfer/v0.6.16/flashinfer/autotuner/autotuner.py). Since this full-model dummy tuning pass now runs before graph capture on distributed boots, set the tuner process group (prefer the existing gloo world group) around this block and clear it afterward.
Useful? React with 👍 / 👎.
| # runtime floors chunked_prefill_size at the ops.tuning default (8192), so | ||
| # the default sweep matches every deployment with chunked prefill <= 8192. | ||
| # Resweep with a matching value for larger prefill configurations. | ||
| SITU_TUNE_MAX_NUM_TOKENS = get_autotune_max_num_tokens() |
There was a problem hiding this comment.
For deployments with chunked_prefill_size > 8192, the runtime now calls set_autotune_max_num_tokens() and all SiTU MoE calls key FlashInfer's cache with that larger bucket ladder, but this module-level constant is still fixed to the default and there is no CLI path to change it before _candidate_tactics, _run, and _make_tokens use it. That means the documented “resweep with a matching value” path always writes an 8192-keyed JSON table, so larger-prefill deployments won't hit the offline cache and will fall back to live startup tuning.
Useful? React with 👍 / 👎.
2 participants
Summary
Opens a single autotune window during engine startup — one dummy prefill at
chunked_prefill_sizetokens, before CUDA-graph capture — and makes it theonly place kernel tuning can happen. Serving is a pure cache-lookup /
heuristic-fallback path by construction.
Changes
ops/tuning.py:autotune()context manager +set/get_autotune_max_num_tokens();ModelExecutor._autotune()runs the dummy prefill under it. Flashinfer skipscache hits, so shapes seeded from a pre-swept table cost nothing.
nvfp4_gemm_swiglu_nvfp4_quant(CuteDSL) joins the autotune framework:12
(mma_tiler, cluster)tuple tactics with aTunableRunner, replacing thefixed heuristic.
autotune_frozen/freeze_autotuning(a switch to suppressserving-time tuning admits serving-time tuning is possible — it no longer is),
the MXFP4 SiTU lazy in-situ tuning pass (carried the TP rank-divergence
deadlock hazard the window removes), and the tuning-cache-active flag.
Pre-swept table loads become fire-and-forget cache seeding. The sweeper and
all MoE ops now key buckets from
get_autotune_max_num_tokens().Measured (B200)
(tuned vs. default tactic, nsys kernel-level, cold L2).
Notes
an OOB read in the trtllm-gen dynB kernels' route map (IMA on Kimi-K2.5),
root-caused with compute-sanitizer and fixed upstream in
fix(moe): pad trtllm-gen route map by one element to avoid OOB read flashinfer-ai/flashinfer#4237 (cherry-picked into v0.6.16).
tune_max ≠ runtimeno longer match bucket keys; the default(8192 floor) covers all
chunked_prefill_size ≤ 8192deployments.