Conversation
build_and_load() calls build() unconditionally, which always invokes ninja. On shared filesystems (GPFS, NFS) this causes a full recompile of the CUTLASS fused-MoE kernels (~180 .cu files, ~1.5h on sm90a) on every process restart. The underlying issue is a mtime discrepancy between what ninja records in .ninja_log during the post-link restat pass and what a subsequent stat() call returns on GPFS (inode flush timing). Ninja sees them as different, concludes the .o files are stale, and rebuilds everything. Fix: if the .so already exists, return early without touching ninja. The lock+build path is only taken on the first compilation. Added an inner exists() check inside the lock for the concurrent first-run case.
Rehvaro
requested review from
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
June 9, 2026 18:36
Contributor
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (44)
📝 WalkthroughWalkthroughAdds GDN benchmarks and Triton reference kernels; implements a unified MoE API (configs, weight preparation, runners, MoELayer) with BF16 SwiGLU OA parameter support threaded through C++/Python/trace/tests; changes CuteDSL wrapper CUDA-graph behavior; fixes sampling/top-k paths; improves JIT build with ninja restat and lock-free existence check; and updates many tests/docs. ChangesMonolithic PR: unified MoE, GDN benchmarks, BF16 OA, infra fixes
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request attempts to optimize the JIT compilation and loading process by checking if the shared library already exists before acquiring a file lock. However, the reviewer identified a critical race condition where checking for the file's existence outside of the file lock could lead to loading a partially written library in multi-process or multi-GPU environments. It is recommended to remove the pre-lock check and only perform the existence check inside the lock to ensure safety.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+311
to
324
| so_path = self.jit_library_path | ||
| if so_path.exists(): | ||
| return self.load(so_path) | ||
|
|
||
| # Guard both build and load with the same lock to avoid race condition | ||
| # where another process is building the library and removes the .so file. | ||
| with FileLock(self.lock_path, thread_local=False): | ||
| so_path = self.jit_library_path | ||
| verbose = os.environ.get("FLASHINFER_JIT_VERBOSE", "0") == "1" | ||
| self.build(verbose, need_lock=False) | ||
| if not so_path.exists(): | ||
| verbose = os.environ.get("FLASHINFER_JIT_VERBOSE", "0") == "1" | ||
| self.build(verbose, need_lock=False) | ||
| result = self.load(so_path) | ||
|
|
||
| return result |
Contributor
There was a problem hiding this comment.
Checking so_path.exists() outside of the lock introduces a race condition in multi-process/multi-GPU environments (such as Tensor Parallelism setups).
When one process is compiling and linking the .so library, the file is created on disk before the write/link operation is fully complete. If another process starts up at this moment, it will see so_path.exists() == True outside the lock, bypass the lock, and attempt to load the partially written/incomplete .so file, leading to a crash (e.g., OSError or segmentation fault).
To prevent this race condition while still avoiding the GPFS recompilation issue, we should remove the check outside the lock and rely solely on the check inside the lock. Since the lock is already acquired, checking so_path.exists() inside the lock is completely safe and still skips the expensive self.build() call once the library is fully built.
# Guard both build and load with the same lock to avoid race condition
# where another process is building the library and removes the .so file.
with FileLock(self.lock_path, thread_local=False):
so_path = self.jit_library_path
if not so_path.exists():
verbose = os.environ.get("FLASHINFER_JIT_VERBOSE", "0") == "1"
self.build(verbose, need_lock=False)
result = self.load(so_path)
return result![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/jit/core.py (1)
311-313: 💤 Low valueEarly return optimization looks good; consider documenting the TOCTOU tradeoff.
The lock-free existence check provides a valuable optimization for the shared-filesystem use case described in the PR. However, there's a narrow TOCTOU window: the
.socould theoretically be deleted between theexists()check (line 312) and theload()call (line 313), causingload()to fail without fallback.Given the PR context—stable shared filesystems where the
.sopersists once built—this risk is negligible and the performance benefit justifies the tradeoff. The locked double-check below (line 319) ensures only one process builds on first run.Optional: Document the known tradeoff
Consider adding a brief inline comment acknowledging this design choice:
# Fast path: if .so exists, load immediately without lock. # TOCTOU risk (deletion between check and load) is acceptable for this use case. so_path = self.jit_library_path if so_path.exists(): return self.load(so_path)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/jit/core.py` around lines 311 - 313, Add a brief inline comment next to the fast-path existence check that documents the intentional TOCTOU tradeoff: identify the so_path = self.jit_library_path check and the immediate if so_path.exists(): return self.load(so_path) fast path, state that it deliberately avoids locking for performance on stable shared filesystems, and note the small risk that the .so could be removed between exists() and load() (acceptable given the locked double-check fallback via the builder). Keep the comment short and colocated with the exists()/load() lines.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@flashinfer/jit/core.py`:
- Around line 311-313: Add a brief inline comment next to the fast-path
existence check that documents the intentional TOCTOU tradeoff: identify the
so_path = self.jit_library_path check and the immediate if so_path.exists():
return self.load(so_path) fast path, state that it deliberately avoids locking
for performance on stable shared filesystems, and note the small risk that the
.so could be removed between exists() and load() (acceptable given the locked
double-check fallback via the builder). Keep the comment short and colocated
with the exists()/load() lines.
<!-- .github/pull_request_template.md --> ## 📌 Description The `FilteredTopK` overflow refinement path has an smem race (flagged by `racecheck`). In the overflow refinement loop, each round reuses `s_histogram` ([code](https://github.com/flashinfer-ai/flashinfer/blob/a28703432faab15fda7edd71b6c80be0206df973/include/flashinfer/topk.cuh#L2733)): - near the end of round `r`, threads read `s_histogram` to update `topk_remain`: ```cpp topk_remain -= static_cast<uint32_t>(s_histogram[threshold + 1]); ``` - at the start of round `r + 1`, threads clear `s_histogram`: ```cpp if (tx < RADIX + 1) s_histogram[tx] = 0; ``` This PR fixes by adding a `__synchtreads()` after the `topk_remain` decrement before the loop can either break or continue to the next round. --- <details> <summary> Repro Script </summary> ```bash compute-sanitizer --tool racecheck --target-processes all uv run python repro_flashinfer_topk_race.py ``` ``` #!/usr/bin/env python3 """Minimal FlashInfer radix top-k racecheck repro. This script intentionally uses a tie-heavy input so FlashInfer's filtered top-k kernel enters the shared-memory overflow refinement path. Example: python repro_flashinfer_topk_race.py Run under NVIDIA Compute Sanitizer: compute-sanitizer --tool racecheck --target-processes all \ python repro_flashinfer_topk_race.py Affected FlashInfer versions report shared-memory hazards in the radix top-k kernel. The script does not use any project-specific wrappers. """ from __future__ import annotations import argparse import sys from importlib import metadata import torch DTYPES = { "float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16, } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--rows", type=int, default=512) parser.add_argument("--vocab", type=int, default=131072) parser.add_argument("--k", type=int, default=128) parser.add_argument("--dtype", choices=sorted(DTYPES), default="float32") parser.add_argument("--iters", type=int, default=1) parser.add_argument("--deterministic", action="store_true", default=True) parser.add_argument("--nondeterministic", dest="deterministic", action="store_false") parser.add_argument("--sorted", action="store_true") return parser.parse_args() def main() -> int: args = parse_args() if not torch.cuda.is_available(): print("CUDA is required for this repro.", file=sys.stderr) return 1 import flashinfer from flashinfer.topk import can_implement_filtered_topk, get_topk_module device = torch.device("cuda") dtype = DTYPES[args.dtype] try: flashinfer_version = metadata.version("flashinfer") except metadata.PackageNotFoundError: flashinfer_version = getattr(flashinfer, "__version__", "unknown") print(f"torch={torch.__version__}") print(f"flashinfer={flashinfer_version} ({flashinfer.__file__})") print(f"cuda_device={torch.cuda.get_device_name(device)}") print(f"can_implement_filtered_topk={can_implement_filtered_topk()}") print( "case=" f"rows={args.rows}, vocab={args.vocab}, k={args.k}, dtype={args.dtype}, " f"sorted={args.sorted}, deterministic={args.deterministic}" ) # All zeros force all logits into the same radix bucket. For vocab=128K and # k=128 this enters the overflow refinement loop that reuses s_histogram. # The default 512 rows give racecheck enough CTAs to report the hazard # reliably on affected versions. logits = torch.zeros((args.rows, args.vocab), dtype=dtype, device=device) topk_module = get_topk_module() row_states_buffer = torch.zeros(1024 * 1024, dtype=torch.uint8, device=device) output_values = torch.empty((args.rows, args.k), dtype=dtype, device=device) indices = None for _ in range(args.iters): indices = topk_module.radix_topk( logits, args.k, args.sorted, args.deterministic, row_states_buffer, output_values, ) torch.cuda.synchronize() assert indices is not None assert output_values.shape == (args.rows, args.k) assert indices.shape == (args.rows, args.k) assert indices.dtype == torch.int32 assert bool(torch.all(output_values == 0).item()) assert bool(torch.all((indices >= 0) & (indices < args.vocab)).item()) print("completed") return 0 if __name__ == "__main__": raise SystemExit(main()) ``` </details> <details> <summary> Example Output </summary> ``` ========= COMPUTE-SANITIZER ========= Variable environment CUDA_COREDUMP_FILE is not supported by compute-sanitizer, clearing it before target process launch. ========= Variable environment CUDA_COREDUMP_GENERATION_FLAGS is not supported by compute-sanitizer, clearing it before target process launch. ========= Variable environment CUDA_ENABLE_COREDUMP_ON_EXCEPTION is not supported by compute-sanitizer, clearing it before target process launch. torch=2.11.0+cu130 flashinfer=0.6.8.post1 (/tmp/.../lib/python3.12/site-packages/flashinfer/__init__.py) cuda_device=NVIDIA GB300 can_implement_filtered_topk=True case=rows=512, vocab=131072, k=128, dtype=float32, sorted=False, deterministic=True ========= Error: Race reported between Read access at void flashinfer::sampling::FilteredTopKUnifiedKernel<float, int, (int)4, (bool)1, (flashinfer::sampling::FilteredTopKMode)0>(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x111e0 ========= and Write access at void flashinfer::sampling::FilteredTopKUnifiedKernel<float, int, (int)4, (bool)1, (flashinfer::sampling::FilteredTopKMode)0>(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x11230 [3852 hazards] ========= ========= Error: Race reported between Read access at void flashinfer::sampling::FilteredTopKUnifiedKernel<float, int, (int)4, (bool)1, (flashinfer::sampling::FilteredTopKMode)0>(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x137a0 ========= and Write access at void flashinfer::sampling::FilteredTopKUnifiedKernel<float, int, (int)4, (bool)1, (flashinfer::sampling::FilteredTopKMode)0>(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x137f0 [3596 hazards] ========= ========= Error: Race reported between Read access at void flashinfer::sampling::FilteredTopKUnifiedKernel<float, int, (int)4, (bool)1, (flashinfer::sampling::FilteredTopKMode)0>(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x16160 ========= and Write access at void flashinfer::sampling::FilteredTopKUnifiedKernel<float, int, (int)4, (bool)1, (flashinfer::sampling::FilteredTopKMode)0>(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x161a0 [3560 hazards] ========= completed ========= RACECHECK SUMMARY: 3 hazards displayed (3 errors, 0 warnings) ``` </details> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Bug Fixes * Improved reliability of top-k filtering operations by resolving a synchronization issue in the multi-round fallback path to ensure accurate results across all computational scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
Collaborator
|
I have some concerns about blindly skipping the build based solely on os.path.exists. During local development, when we modify .cu or .cuh source files, we expect the build system to recompile and reflect the latest changes. Simply skipping the build because an old .so exists will lead to a painful developer experience ("code changed but not taking effect"). |
<!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Running the benchmark before and after on a B200 machine gives: ``` python3 flashinfer_benchmark.py \ --routine BatchMLAPagedAttentionWrapper \ --backends trtllm-native \ --page_size 64 --batch_size 1 \ --s_qo 1 --s_kv 8192 \ --num_qo_heads 128 --num_kv_heads 128 \ --head_dim_ckv 512 --head_dim_kpe 64 \ --random_actual_seq_len -vv --refcheck \ --q_dtype fp8_e4m3 --kv_dtype fp8_e4m3 \ --generate_repro_command --case_tag DeepSeek-R1 \ --num_iters 100 --dry_run_iters 20 \ ``` - Before: 0.016 ms, 73.2 TFLOPs/s - After: 0.013 ms, 91.4 TFLOPs/s ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved workspace memory handling for MLA decode: the buffer region used for profiling is now cleared only during profiling runs. This prevents unnecessary memory operations during normal inference, reduces interference in shared-workspace scenarios, and improves performance and stability during autotuning/profile measurements. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
Author
Hmmm. That is a good point. |
…for large-vocab small-k sampling (#3461) <!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ### Summary Adds a faster `filter_apply_order="top_k_first"` path for `top_k_top_p_sampling_from_logits` and `top_k_top_p_sampling_from_probs`. For modest `top_k` over a large vocabulary, instead of masking/renormalizing the full vocab and running rejection sampling across it, we select the top-k entries with the parallel radix top-k kernel and run top-p **over only those k survivors**. This is distribution-equivalent to the existing path but 2–4× faster for `from_logits` (see benchmarks). Addresses the slow small-batch top-k/top-p sampling reported in #3389. ### Motivation The existing `top_k_first` path keeps everything in full-vocab layout: - `from_logits`: `top_k_mask_logits` → `torch.softmax` (full vocab) → `top_p_sampling_from_probs` (single-CTA rejection over full vocab) - `from_probs`: `top_k_renorm_probs` → `top_p_sampling_from_probs` (single-CTA rejection over full vocab) After top-k, only `k` entries are relevant, but the general full-vocab kernels still process all `V` elements — a full-vocab softmax (logits) and a multi-round single-CTA rejection scan (both). At small batch the rejection kernel launches one CTA per request and underutilizes the GPU. ### Changes A gated fast path in `flashinfer/sampling.py` (no kernel changes): 1. Select top-k via `flashinfer.top_k(..., sorted=True)` (parallel radix selection, returns the `k` values + their indices). 2. Normalize over the `k` survivors — `softmax` (logits) or renorm (probs); these are mathematically identical to the masked full-vocab versions. 3. Run `top_p_sampling_from_probs` over the `k`-element distribution, then map the local choice back to the global vocab index via `gather`. Both entry points route through the same helper with `sorted=True`, so they reduce to the same `probs_k` and stay **sample-aligned** with each other. ### Accuracy - **Distribution-equivalent** to the original `top_k_first` path (validated: total-variation distance ≈ 0.007 vs. the analytic target over 40k draws). - **Per-seed sample values may differ** from previous versions (the RNG now maps over a `k`-element domain, and radix top-k may break k-th-boundary ties differently). The two APIs remain aligned with each other. - **CUDA-graph compatible**: the default vectorized top-k captures into a graph; the graph-safe fallback is only requested for the non-deterministic cluster path. ### Benchmarks `top_k=50`, `top_p=0.9`, CUDA-graph mode, median ms. "before" = original path, "after" = fast path. ### B200 (SM100) | API | dtype | bs | vocab | before | after | speedup | |---|---|---|---|---|---|---| | from_logits | bf16 | 1 | 128k | 0.121 | 0.059 | 2.05× | | from_logits | bf16 | 8 | 128k | 0.139 | 0.064 | 2.17× | | from_logits | bf16 | 32 | 128k | 0.147 | 0.066 | 2.23× | | from_logits | bf16 | 256 | 128k | 0.431 | 0.180 | 2.39× | | from_logits | bf16 | 1 | 256k | 0.206 | 0.067 | 3.07× | | from_logits | bf16 | 32 | 256k | 0.294 | 0.075 | 3.92× | | from_logits | fp32 | 1 | 128k | 0.126 | 0.057 | 2.21× | | from_logits | fp32 | 8 | 128k | 0.131 | 0.063 | 2.08× | | from_logits | fp32 | 32 | 128k | 0.163 | 0.068 | 2.40× | | from_logits | fp32 | 256 | 128k | 0.453 | 0.122 | 3.71× | | from_logits | fp32 | 1 | 256k | 0.208 | 0.061 | 3.41× | | from_logits | fp32 | 32 | 256k | 0.296 | 0.113 | 2.62× | | from_probs | fp32 | 1 | 128k | 0.076 | 0.064 | 1.19× | | from_probs | fp32 | 8 | 128k | 0.077 | 0.067 | 1.15× | | **from_probs** | fp32 | 32 | 128k | 0.088 | 0.105 | **0.84×** | | from_probs | fp32 | 256 | 128k | 0.365 | 0.155 | 2.35× | | from_probs | fp32 | 1 | 256k | 0.099 | 0.067 | 1.48× | | from_probs | fp32 | 32 | 256k | 0.235 | 0.116 | 2.03× | ### RTX PRO 6000 (SM120) | API | dtype | bs | vocab | before | after | speedup | |---|---|---|---|---|---|---| | from_logits | bf16 | 1 | 128k | 0.096 | 0.042 | 2.29× | | from_logits | bf16 | 8 | 128k | 0.111 | 0.044 | 2.52× | | from_logits | bf16 | 32 | 128k | 0.143 | 0.054 | 2.65× | | from_logits | bf16 | 256 | 128k | 0.454 | 0.173 | 2.62× | | from_logits | bf16 | 1 | 256k | 0.159 | 0.042 | 3.79× | | from_logits | bf16 | 32 | 256k | 0.232 | 0.081 | 2.86× | | from_logits | fp32 | 1 | 128k | 0.101 | 0.041 | 2.46× | | from_logits | fp32 | 8 | 128k | 0.126 | 0.045 | 2.80× | | from_logits | fp32 | 32 | 128k | 0.162 | 0.079 | 2.05× | | from_logits | fp32 | 256 | 128k | 0.631 | 0.287 | 2.20× | | from_logits | fp32 | 1 | 256k | 0.171 | 0.042 | 4.07× | | from_logits | fp32 | 32 | 256k | 0.287 | 0.113 | 2.54× | | from_probs | fp32 | 1 | 128k | 0.056 | 0.043 | 1.30× | | from_probs | fp32 | 8 | 128k | 0.061 | 0.048 | 1.27× | | from_probs | fp32 | 32 | 128k | 0.110 | 0.083 | 1.33× | | from_probs | fp32 | 256 | 128k | 0.405 | 0.293 | 1.38× | | from_probs | fp32 | 1 | 256k | 0.080 | 0.044 | 1.82× | | from_probs | fp32 | 32 | 256k | 0.252 | 0.115 | 2.19× | ### Commands to reproduce results above ``` python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 1 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 8 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 32 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 256 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 1 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 32 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 1 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 8 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 32 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 256 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 1 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 32 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 1 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 8 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 32 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 256 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 1 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 32 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first ``` ## 🔍 Related Issues <!-- Link any related issues here --> #3389 ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Added an automatic fast path that speeds up sampling when using top‑k‑first filtering with a scalar k on large vocabularies. Sampling now selects and renormalizes a smaller candidate set for faster draws, preserves determinism/tie‑breaking and validity flags when requested, falls back to the original flow when not applicable, and introduces no public API changes. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
… overlapping in TRT-LLM fused MoE (#3483) <!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> The current code hardcodes the number of reserved SMs for overlapping in TRT-LLM fused MoE to be 8. Because sometimes better overlapping can be achieved by using other numbers of reserved SMs for overlapping, this PR modifies the code to use an environment variable `FLASHINFER_TRTLLM_MOE_OVERLAP_RESERVED_SMS` to control the number of reserved SMs. If this environment variable is not set, the current value 8 is used. ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Performance Improvements** * Cooperative-kernel resource management now uses a validated, configurable allocation instead of a fixed constant, improving GPU utilization and flexibility; the chosen allocation is logged when the cooperative path is used. * **New Features** * Added an environment variable to adjust reserved GPU resources at runtime and enable tuning without code changes. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com>
<!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> Resolves #3308. Previously, `max_num_tokens` were only checked in cuda graph path, but now deprecated and unused. Dropped tests that test for pre-allocated buffers. ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **API Changes** * `max_num_tokens` is now optional/deprecated (ignored at runtime); existing calls remain compatible. * **Performance Improvements** * Persistent CUDA stream/event handling improves CUDA-graph compatibility and enables better async-memset overlap, reducing memory overhead. * **Tests** * Tests updated to validate routing/buffer write-completeness via the core execution path using poisoned-buffer checks. * **Documentation** * Docstrings updated to reflect the new persistent CUDA resource model and deprecated `max_num_tokens`. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
… ICHECK message (#3546) ## Summary Two small, low-risk changes surfaced by the v0.6.13 doc-check pass. Touches docs/docstrings only, plus a one-line C++ assertion-message fix (no runtime behavior change). ### docs: close v0.6.13-only doc-check gaps (MISSING / STALE / docstring) [issues/3538](#3538) Resolves the doc-check findings that are new in v0.6.13 vs v0.6.12, using docs + docstrings only (no new `@flashinfer_api` decorators). - **MISSING (API exported but not listed in `.rst`):** - `docs/api/fused_moe.rst`: add a "Multi-LoRA MoE (BGMV)" section listing `bgmv_moe` / `bgmv_moe_shrink` / `bgmv_moe_expand`. - `docs/api/quantization.rst`: list `nvfp4_quantize_per_token_cute_dsl` under its canonical `flashinfer.quantization.kernels.nvfp4_quantize` currentmodule (the package-level re-export is guarded by `is_cute_dsl_available()` and is not importable at docs-build time without `nvidia-cutlass-dsl`). - **Docstring completeness:** - `nvfp4_quantize_per_token_cute_dsl`: expand the one-line summary into a full NumPy-style docstring with Parameters/Returns, mirroring the sibling `nvfp4_quantize_cute_dsl` and documenting the per-token scale output. ### fix(moe): correct misleading `topk_indices` dtype ICHECK message In `csrc/fused_moe/noAuxTcKernels.cu`, the dtype check for `topk_indices` requires int32 (`encode_dlpack_dtype(...) == int32_code`), but the assertion message was copy-pasted from the `topk_values` check and wrongly read "must have the same dtype as scores" (scores are fp32/bf16, not int32), which misleads debugging. Message changed to "topk_indices must be int32 dtype", matching the nearby `routing_replay_out` int16 check style. Message-only change; the runtime condition is unchanged. ## Test plan - [ ] Docs build (`docs/build_docs.sh`) succeeds; new autosummary entries render and removed STALE entries no longer produce dangling references. - [ ] `flashinfer_document_check` shows 0 MISSING and the two v0.6.13 STALE items (`is_cute_dsl_available`, `trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant`) gone; Docstring Completeness has 0 findings. - [ ] No code-path change for the MoE ICHECK; existing MoE routing tests remain green. AI-assisted. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added documentation for Multi-LoRA MoE (BGMV) operators: `bgmv_moe`, `bgmv_moe_shrink`, and `bgmv_moe_expand`. * **Documentation** * Enhanced quantization kernel documentation with detailed per-token scaling specifications and layout guidance. * Cleaned up API documentation by removing outdated entries and sections. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
… NVFP4 autotune (#3093) ## Summary Draft PR introducing a unified MoE layer that autotunes across NVFP4 backends — CuteDSL and TRTLLM FP4 routed — on a per-shape basis. ## Update — MVP cut (2026-05-31) Autonomous MVP-completion pass on a B200 (SM100). The cross-backend autotune objective is met and validated end-to-end. Full narrative in `docs/design_docs/flashinfer_moe_api.md` (Decision Log + the new "MVP As-Built Reference"). - **All MVP follow-ups (CR1–CR11) done & validated**: `tests/moe/test_unified_moe_api.py` **9/9** (layer + per-backend accuracy vs bf16, autotuner visits both candidates, CUDA-graph replay), `tests/moe/test_moe_api.py` **97/97** (CPU config + fail-fast validation), and the `unified_nvfp4_moe` sweep (128→16384 tokens) with `--refcheck` passing for both backends. - **Blocker found & fixed**: both runner adapters had never executed against the post-`main`-merge `core.py` (a stale raw-`moe_op` API + a class-vs-instance `tuning_config` bug). `TrtllmFp4RoutedRunner` now delegates to the canonical `core.MoERunner` (newly exported from `get_trtllm_moe_sm100_module()`); the unified adapters only translate Packs ⇄ the inner runner's native tensor list. Future direction (out of scope): make the low-level TVM-FFI ops take structured config objects (design doc §5) to kill positional-arg drift. - **New since first draft**: `local_expert_offset` wired into TRTLLM packing (+EP-offset test); fail-fast NVFP4/Swiglu scope validation; per-token-bucket winner cache; `tune_max_num_tokens` threaded into runner tuning; first-class NVFP4 weight prep (`prepare.py` + `*.prepare_weights`) replacing duplicated test/bench prep; `--refcheck` for the unified benchmark. ## What this adds - **`MoEConfig` / `MoEActivationPack` / `MoEWeightPack`** — frozen config dataclasses and per-call / long-lived tensor containers (`flashinfer/fused_moe/api.py`). Single `QuantVariant` enum replaces the old 3-axis dtype × granularity × variant split. - **`CuteDslNvfp4Runner` / `TrtllmFp4RoutedRunner`** — `TunableRunner` adapters with `pack_inputs(act, weights)` translating packs into the backend's native tensor list (`flashinfer/fused_moe/runners.py`). ~~each with its own `tuning_config`~~ → each **delegates to a canonical inner runner** (`CuteDslFusedMoENvfp4Runner` / `core.MoERunner`) and builds its `tuning_config` **per-instance**. - **`MoELayer`** — stateful dispatcher that builds one runner per compatible backend config, then on first call ~~per shape~~ **per token-bucket** runs per-runner `choose_one` (within-backend tactic tuning) + cross-runner `bench_gpu_time` comparison (cross-backend selection). Caches the winner **keyed by tuning bucket** (`flashinfer/fused_moe/layer.py`). - **`unified_nvfp4_moe` benchmark routine** — wired into `benchmarks/flashinfer_benchmark.py` (not a standalone script). Emits one result row per candidate backend with the winner marked; supports `--refcheck` against a shared bf16 reference. - **`bench_unified_moe_today.sh`** — convenience wrapper at repo root invoking the infra routine across the eight shapes we care about (EP=1 sweep + EP=16 wide-EP regime). - **`tests/moe/test_unified_moe_api.py`** — accuracy tests (`MoELayer` vs bf16 ref; each backend vs same ref), plumbing tests (autotuner visits all candidates; CUDA graph capture + replay), **and a pre-routed EP-offset packing test**. - **`flashinfer/fused_moe/prepare.py`** *(new)* — first-class NVFP4 weight-prep helpers exposed as `TrtllmFp4Config.prepare_weights` / `CuteDslConfig.prepare_weights`; the test and benchmark no longer carry duplicated prep copies. ## Design notes - **Per-runner `choose_one`** — `AutoTuner.choose_one` assumes all runners share one `inputs` list during profiling. Our backends' native schemas differ (CuteDSL 12 tensors with unpacked topk + trailing output buffer; TRTLLM 8-field packed `MoEInputs`). Resolution: call `choose_one` once per runner (within-backend tactic selection) and use `bench_gpu_time` to compare the per-runner winners. ~~Each runner owns its own `tuning_config` as a class attribute.~~ → **`tuning_config` is built per-instance** (the TRTLLM runner via `MoERunner._make_tuning_config`, which also threads `ExecutionConfig.tune_max_num_tokens` into the bucket set). - **Shared-reference accuracy testing** — each backend is tested against the same bf16 reference, not against the other backend. Catches shared-mode failures (both wrong in the same way) that cross-backend agreement would miss. `--refcheck` brings the same check to the benchmark. - **Weight layout: `Shuffled_MajorK` only** — the only NVFP4-compatible TRTLLM layout today. Multi-layout autotune (opt-in additional variants via a future `layouts` field on `TrtllmFp4Config`) is a V2 extension — design accommodates it without core changes. ## Out of scope for this PR - FP8 / BF16 / MxInt4 backends - Monolithic routing (our path is pre-routed) - Unpacked topk for TRTLLM — #2425 (when that lands, `TrtllmFp4RoutedRunner` drops the inline `(id << 16) | bf16_bits` packing) - Multi-layout TRTLLM autotune across `Shuffled_BlockMajorK` / `NoShuffle_MajorK` - SM120/SM121 support for CuteDSL (kernel is SM100/SM103 only today; `CuteDslConfig.supported` tightened to match) - First-class **activation** prep + structured-config TVM-FFI boundary (§5) — post-MVP carryover ## Status ~~Draft — benchmark numbers and observation notes are being collected on a Blackwell box. Will de-draft once the winner-flip table is filled in and accuracy tests pass on hardware.~~ **MVP scope complete and validated on B200 (SM100).** Kept as a **draft PR on purpose** so CI skips while further edits are pushed autonomously — flip to "Ready for review" when you want CI to run. ## Test plan - [x] `pytest tests/moe/test_unified_moe_api.py -v` (SM100/SM103) — **9/9** - [ ] ~~`./bench_unified_moe_today.sh` — winner column matches expected per-shape~~ → not run this pass; equivalent coverage via the `unified_nvfp4_moe` routine sweep (128→16384) + `--refcheck` (both backends pass). The script's EP=16 shapes are still worth a dedicated run. - [x] Pre-commit hooks pass ## Related PRs - #3453 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Unified MoE API surfaced package-wide: immutable configs, MoELayer (cached cross-backend autotune winner), backend runners, NVFP4 weight-prep utilities, and re-exported unified API symbols. * **Benchmarking** * New benchmark sweep script for unified NVFP4 MoE and registration to record token-size sweeps to CSV. * **Documentation** * Added comprehensive FlashInfer Unified MoE API design doc with migration plan. * **Tests** * Added CPU and gated-GPU tests for API, validation, accuracy, and autotune/dispatch behavior; improved enum reprs for round-trip logging. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Yang Xu <yanxu@nvidia.com>
## Summary - allow DeepSeek V4 sparse MLA validation for 8/16/32 heads in addition to 64/128 - add targeted BF16/FP8 sparse MLA coverage for 8, 16, and 32 heads ## Tests - `python -m ruff check flashinfer/mla/_core.py tests/attention/test_trtllm_gen_sparse_mla_dsv4.py` - `FLASHINFER_DISABLE_VERSION_CHECK=1 python -m pytest --collect-only tests/attention/test_trtllm_gen_sparse_mla_dsv4.py -q` - `FLASHINFER_DISABLE_VERSION_CHECK=1 python -m pytest tests/attention/test_trtllm_gen_sparse_mla_dsv4.py -k 'h8 or h16 or h32' -vv` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Expanded DeepSeek V4 sparse MLA support to head counts: 8, 16, 32, 64, and 128. * **Bug Fixes** * Updated input validation to accept the expanded set of supported head counts. * **Tests** * Added test coverage for sparse MLA decode scenarios with smaller head counts, multiple data types, and both KV layout variants. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Perkz Zheng <perkzz@users.noreply.github.com>
<!-- .github/pull_request_template.md -->
## 📌 Description
Adds Gated Delta Net (GDN) support to `flashinfer_benchmark.py` as a new
`gdn` routine category:
Added Routines:
- `gated_delta_rule_decode` — T=1 decode; `--state_layout
{pretranspose,nontranspose}`, `--state_dtype {float32,bfloat16}` (bf16
state kernels), `--pool_mode {single,split}`
- `gated_delta_rule_mtp` — T>=2 multi-token processing with state pool +
indices
- `chunk_gated_delta_rule` — varlen chunked prefill
#### Bug fixes in the standalone benches
- `bench_gdn_decode.py`: all-layouts mode called the BF16-state kernel
without pool indices (column was always N/A); bytes model counted
intermediate-state traffic for T>1 even with caching disabled, inflating
reported TB/s
- `bench_gdn_prefill.py`: fed log-space `g` to the FlashInfer kernel,
which takes a linear-space alpha — outputs/state were NaN; FLA is now
optional instead of a hard exit
<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->
## 🔍 Related Issues
<!-- Link any related issues here -->
## 🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.
### ✅ Pre-commit Checks
- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.
> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).
## 🧪 Tests
- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).
## Reviewer Notes
<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Added GDN (Gated Delta Net) benchmark family: decode, multi-token
(MTP), and chunked prefill with multi-backend support and performance
metrics that account for optional intermediate-state caching.
* Added Triton-backed reference kernels and optional Triton routing for
GDN workloads.
* **Bug Fixes**
* Prefill baseline is no longer fatal when optional baseline kernel is
missing; reporting and columns adjust accordingly.
* Corrected gate generation to avoid out-of-domain values.
* **Documentation**
* Updated docs and sample test lists with GDN routines, flags, backend
matrix, and usage notes.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
<!-- .github/pull_request_template.md --> ## 📌 Description Add optional per-expert SwiGLU OA activation parameter support to the TRT-LLM BF16 MoE APIs: - `gemm1_alpha` - `gemm1_beta` - `gemm1_clamp_limit` These are optional float32 CUDA tensors with shape `[local_num_experts]` and default to `None`, so existing callers keep the same behavior. This PR threads the new parameters through the BF16 C++ launcher, Python registered op/fake op, public APIs, autotuner path, `MoERunner`, trace metadata, and Python reference implementation. For BF16, these values are passed through as raw scalar values. No host-side dequant-scale adjustment is applied. The SwiGLU OA formula is: ```python activation = X2 * sigmoid(alpha * X2) * (X1 + beta) with optional clamp behavior: X1 = clamp(X1, min=-limit, max=limit) X2 = clamp(X2, max=limit) ``` The implementation also validates that BF16 OA parameters are only used with ActivationType.Swiglu. ## 🔍 Related Issues N/A ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Local checks: python3 -m compileall flashinfer/fused_moe/core.py flashinfer/trace/templates/moe.py tests/moe/test_trtllm_gen_fused_moe.py git diff --check pre-commit run --all-files GPU validation on computelab GB100 SM100 node: tests/moe/test_trtllm_gen_fused_moe.py::test_bf16_moe_swiglu_oa_activation_param_validation tests/moe/test_trtllm_gen_fused_moe.py::test_bf16_moe_swiglu_oa_activation_params Result: 2 passed in 605.52s (0:10:05) ## Reviewer Notes This follows the MXFP8 SwiGLU OA API shape and raw-parameter semantics: BF16 passes `gemm1_alpha`, `gemm1_beta`, and `gemm1_clamp_limit` directly to trtllm-gen without host-side `dequantScaleAb` adjustment. `gemm1_alpha` is never scaled. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added optional per-expert BF16 SwiGLU activation parameters (alpha, beta, clamp limit) to MoE public APIs and runtime, enabling per-expert activation shaping and clamping. * **Validation** * Added runtime/schema checks to ensure these optional parameters are well-formed and only accepted with SwiGLU activations. * **Trace Templates** * BF16 MoE trace templates updated to include the new optional SwiGLU/OA inputs. * **Tests** * Added unit/trace tests covering parameter schema, behavior (clamp/OA/default), and production vs. reference correctness. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Sam (Kesen Li) <lsam@nvidia.com>
<!-- .github/pull_request_template.md --> ## 📌 Description ### Summary This PR refreshes the `b12x` dense FP4 GEMM kernel to match the latest upstream for better performance. **No public API change.** The `auto` routing heuristic (`_heuristic_func_mm_fp4`) is **unchanged** — b12x is preferred only at SM120, not SM121. Addresses **#3517** *[Perf][SM12x] Update mm_fp4(backend='b12x') kernels*. ### Motivation @lukealonso has made performance improvements to b12x dense GEMMs; we'd like to port over the improvements to FlashInfer. Gains are focused on small-M. ### Benchmarks: new b12x vs old b12x Setup per SKU: `benchmarks/flashinfer_benchmark.py --testlist <104 nvfp4 shapes> --use_nvfp4 --use_128x4_sf_layout --refcheck`, 30 iters / 5 warmup, **CUPTI timing** (harness default). 104 shapes = 13 token counts m∈{1..4096} × 8 (n,k) layer shapes taken from real **DeepSeek-R1** and **Llama-3** model layers, all K%128==0. The refresh is a consistent small win at small `m` (where `swap_ab` and `expected_m` help most) and break-even-to-positive overall — **no SKU regresses in aggregate**. Geomean speedup (old_time / new_time; >1 = new faster), grouped by `m`: | SKU | overall | m≤8 | 16≤m≤128 | m≥256 | |---|---|---|---|---| | RTX 5080 (SM120) | 1.003× | 1.019× | 1.003× | 0.991× | | RTX PRO 6000 (SM120) | **1.021×** | 1.060× | 1.010× | 1.000× | | DGX Spark / GB10 (SM121) | **1.024×** | 1.021× | 1.028× | 1.022× | The win is largest at small `m` on PRO 6000 / Spark and break-even on the 5080. A few `n=4096 k=4096/5376` shapes regress ≤5% (bolded in the per-shape tables); the PRO 6000 has a sharper `8192×2560` regression at medium `m` (down to 0.873×), which is SKU-specific and flagged as a follow-up. Full per-shape new-vs-old tables (all 104 shapes, sorted by `n, k, m`) per SKU: <details> <summary><b>RTX 5080 (SM120)</b></summary> | n | k | m | old (us) | new (us) | new TFLOPs | speedup | |---|---|---|---|---|---|---| | 512 | 7168 | 1 | 37.6 | 39.8 | 0.2 | **0.945x** | | 512 | 7168 | 2 | 38.3 | 38.8 | 0.4 | 0.986x | | 512 | 7168 | 4 | 38.8 | 38.9 | 0.8 | 0.995x | | 512 | 7168 | 8 | 38.9 | 39.2 | 1.5 | 0.994x | | 512 | 7168 | 16 | 39.4 | 40.4 | 2.9 | 0.975x | | 512 | 7168 | 32 | 40.6 | 41.2 | 5.7 | 0.985x | | 512 | 7168 | 64 | 43.8 | 43.5 | 10.8 | 1.007x | | 512 | 7168 | 128 | 55.1 | 55.6 | 16.9 | 0.992x | | 512 | 7168 | 256 | 67.4 | 67.7 | 27.8 | 0.995x | | 512 | 7168 | 512 | 84.0 | 84.7 | 44.4 | 0.991x | | 512 | 7168 | 1024 | 94.4 | 95.9 | 78.4 | 0.984x | | 512 | 7168 | 2048 | 104.4 | 106.0 | 141.9 | 0.985x | | 512 | 7168 | 4096 | 165.4 | 168.2 | 178.7 | 0.983x | | 1024 | 8192 | 1 | 58.8 | 59.1 | 0.3 | 0.995x | | 1024 | 8192 | 2 | 58.7 | 59.0 | 0.6 | 0.995x | | 1024 | 8192 | 4 | 59.7 | 59.7 | 1.1 | 0.999x | | 1024 | 8192 | 8 | 60.2 | 59.9 | 2.2 | 1.006x | | 1024 | 8192 | 16 | 61.1 | 61.1 | 4.4 | 1.001x | | 1024 | 8192 | 32 | 60.8 | 60.9 | 8.8 | 0.998x | | 1024 | 8192 | 64 | 70.1 | 69.7 | 15.4 | 1.006x | | 1024 | 8192 | 128 | 78.5 | 77.4 | 27.7 | 1.014x | | 1024 | 8192 | 256 | 85.9 | 85.8 | 50.1 | 1.002x | | 1024 | 8192 | 512 | 96.5 | 97.0 | 88.6 | 0.995x | | 1024 | 8192 | 1024 | 97.3 | 97.5 | 176.2 | 0.998x | | 1024 | 8192 | 2048 | 154.2 | 155.8 | 220.5 | 0.990x | | 1024 | 8192 | 4096 | 251.2 | 254.7 | 269.8 | 0.986x | | 2560 | 8192 | 1 | 101.1 | 87.6 | 0.5 | 1.154x | | 2560 | 8192 | 2 | 101.1 | 91.2 | 0.9 | 1.108x | | 2560 | 8192 | 4 | 101.5 | 93.7 | 1.8 | 1.083x | | 2560 | 8192 | 8 | 101.9 | 95.1 | 3.5 | 1.072x | | 2560 | 8192 | 16 | 102.3 | 100.1 | 6.7 | 1.022x | | 2560 | 8192 | 32 | 101.6 | 95.3 | 14.1 | 1.067x | | 2560 | 8192 | 64 | 101.2 | 95.6 | 28.1 | 1.058x | | 2560 | 8192 | 128 | 95.3 | 94.4 | 56.9 | 1.009x | | 2560 | 8192 | 256 | 109.3 | 106.8 | 100.5 | 1.023x | | 2560 | 8192 | 512 | 113.9 | 114.1 | 188.2 | 0.998x | | 2560 | 8192 | 1024 | 157.3 | 158.3 | 271.3 | 0.994x | | 2560 | 8192 | 2048 | 252.6 | 254.9 | 337.0 | 0.991x | | 2560 | 8192 | 4096 | 433.5 | 436.9 | 393.2 | 0.992x | | 4096 | 4096 | 1 | 57.2 | 57.6 | 0.6 | 0.993x | | 4096 | 4096 | 2 | 57.3 | 56.4 | 1.2 | 1.016x | | 4096 | 4096 | 4 | 58.1 | 59.3 | 2.3 | 0.981x | | 4096 | 4096 | 8 | 56.5 | 57.5 | 4.7 | 0.983x | | 4096 | 4096 | 16 | 59.6 | 60.2 | 8.9 | 0.990x | | 4096 | 4096 | 32 | 60.3 | 60.4 | 17.8 | 1.000x | | 4096 | 4096 | 64 | 57.1 | 59.2 | 36.3 | **0.964x** | | 4096 | 4096 | 128 | 69.9 | 72.8 | 59.0 | **0.960x** | | 4096 | 4096 | 256 | 61.2 | 65.6 | 130.8 | **0.933x** | | 4096 | 4096 | 512 | 112.8 | 113.9 | 150.8 | 0.991x | | 4096 | 4096 | 1024 | 151.8 | 152.3 | 225.7 | 0.997x | | 4096 | 4096 | 2048 | 220.7 | 223.1 | 308.0 | 0.989x | | 4096 | 4096 | 4096 | 366.1 | 376.4 | 365.2 | 0.973x | | 4096 | 5376 | 1 | 72.6 | 71.3 | 0.6 | 1.017x | | 4096 | 5376 | 2 | 72.5 | 70.5 | 1.2 | 1.028x | | 4096 | 5376 | 4 | 76.2 | 75.7 | 2.3 | 1.006x | | 4096 | 5376 | 8 | 75.7 | 75.2 | 4.7 | 1.007x | | 4096 | 5376 | 16 | 76.1 | 78.2 | 9.0 | 0.973x | | 4096 | 5376 | 32 | 83.4 | 85.8 | 16.4 | 0.971x | | 4096 | 5376 | 64 | 77.2 | 81.4 | 34.6 | **0.948x** | | 4096 | 5376 | 128 | 84.4 | 88.6 | 63.6 | **0.953x** | | 4096 | 5376 | 256 | 84.6 | 86.6 | 130.2 | 0.978x | | 4096 | 5376 | 512 | 133.9 | 134.5 | 167.6 | 0.995x | | 4096 | 5376 | 1024 | 178.8 | 180.3 | 250.1 | 0.991x | | 4096 | 5376 | 2048 | 262.6 | 265.2 | 340.2 | 0.990x | | 4096 | 5376 | 4096 | 464.7 | 470.1 | 383.7 | 0.988x | | 4096 | 14336 | 1 | 176.1 | 174.6 | 0.7 | 1.009x | | 4096 | 14336 | 2 | 176.0 | 175.2 | 1.3 | 1.005x | | 4096 | 14336 | 4 | 176.9 | 176.7 | 2.7 | 1.001x | | 4096 | 14336 | 8 | 177.6 | 177.5 | 5.3 | 1.000x | | 4096 | 14336 | 16 | 178.5 | 178.6 | 10.5 | 1.000x | | 4096 | 14336 | 32 | 178.6 | 178.9 | 21.0 | 0.999x | | 4096 | 14336 | 64 | 180.0 | 180.3 | 41.7 | 0.998x | | 4096 | 14336 | 128 | 182.3 | 182.7 | 82.3 | 0.998x | | 4096 | 14336 | 256 | 194.4 | 195.0 | 154.1 | 0.997x | | 4096 | 14336 | 512 | 253.2 | 254.7 | 236.1 | 0.994x | | 4096 | 14336 | 1024 | 375.2 | 377.6 | 318.5 | 0.994x | | 4096 | 14336 | 2048 | 580.2 | 585.3 | 410.9 | 0.991x | | 4096 | 14336 | 4096 | 1049.4 | 1061.3 | 453.3 | 0.989x | | 8192 | 2560 | 1 | 70.1 | 68.4 | 0.6 | 1.025x | | 8192 | 2560 | 2 | 66.1 | 65.4 | 1.3 | 1.010x | | 8192 | 2560 | 4 | 65.3 | 65.2 | 2.6 | 1.001x | | 8192 | 2560 | 8 | 67.5 | 67.0 | 5.0 | 1.008x | | 8192 | 2560 | 16 | 67.9 | 65.2 | 10.3 | 1.042x | | 8192 | 2560 | 32 | 68.2 | 66.1 | 20.3 | 1.032x | | 8192 | 2560 | 64 | 67.1 | 67.8 | 39.6 | 0.989x | | 8192 | 2560 | 128 | 69.3 | 69.6 | 77.2 | 0.996x | | 8192 | 2560 | 256 | 89.8 | 91.6 | 117.3 | 0.981x | | 8192 | 2560 | 512 | 137.4 | 137.1 | 156.6 | 1.002x | | 8192 | 2560 | 1024 | 176.3 | 177.1 | 242.5 | 0.995x | | 8192 | 2560 | 2048 | 279.4 | 280.6 | 306.1 | 0.996x | | 8192 | 2560 | 4096 | 489.2 | 495.3 | 346.9 | 0.988x | | 14336 | 4096 | 1 | 171.5 | 162.8 | 0.7 | 1.053x | | 14336 | 4096 | 2 | 171.7 | 162.1 | 1.4 | 1.059x | | 14336 | 4096 | 4 | 171.5 | 162.0 | 2.9 | 1.059x | | 14336 | 4096 | 8 | 171.9 | 163.3 | 5.8 | 1.053x | | 14336 | 4096 | 16 | 171.8 | 162.9 | 11.5 | 1.054x | | 14336 | 4096 | 32 | 173.4 | 165.2 | 22.7 | 1.049x | | 14336 | 4096 | 64 | 176.7 | 167.4 | 44.9 | 1.055x | | 14336 | 4096 | 128 | 181.3 | 180.4 | 83.3 | 1.005x | | 14336 | 4096 | 256 | 188.7 | 190.3 | 158.0 | 0.992x | | 14336 | 4096 | 512 | 253.7 | 255.0 | 235.8 | 0.995x | | 14336 | 4096 | 1024 | 382.4 | 384.3 | 312.9 | 0.995x | | 14336 | 4096 | 2048 | 620.9 | 624.4 | 385.2 | 0.994x | | 14336 | 4096 | 4096 | 1140.7 | 1148.9 | 418.7 | 0.993x | </details> <details> <summary><b>RTX PRO 6000 (SM120)</b></summary> | n | k | m | old (us) | new (us) | new TFLOPs | speedup | |---|---|---|---|---|---|---| | 512 | 7168 | 1 | 22.3 | 21.7 | 0.3 | 1.026x | | 512 | 7168 | 2 | 22.4 | 22.5 | 0.7 | 0.994x | | 512 | 7168 | 4 | 22.4 | 22.5 | 1.3 | 0.998x | | 512 | 7168 | 8 | 22.4 | 22.5 | 2.6 | 0.995x | | 512 | 7168 | 16 | 22.4 | 22.4 | 5.2 | 0.999x | | 512 | 7168 | 32 | 22.3 | 22.5 | 10.4 | 0.991x | | 512 | 7168 | 64 | 22.4 | 22.7 | 20.7 | 0.991x | | 512 | 7168 | 128 | 22.9 | 23.1 | 40.6 | 0.992x | | 512 | 7168 | 256 | 32.6 | 32.8 | 57.2 | 0.992x | | 512 | 7168 | 512 | 35.7 | 35.9 | 104.8 | 0.996x | | 512 | 7168 | 1024 | 41.5 | 41.7 | 180.1 | 0.995x | | 512 | 7168 | 2048 | 45.6 | 47.1 | 319.2 | **0.969x** | | 512 | 7168 | 4096 | 51.4 | 53.6 | 561.2 | **0.959x** | | 1024 | 8192 | 1 | 25.6 | 25.8 | 0.7 | 0.993x | | 1024 | 8192 | 2 | 25.6 | 25.8 | 1.3 | 0.994x | | 1024 | 8192 | 4 | 25.7 | 25.9 | 2.6 | 0.994x | | 1024 | 8192 | 8 | 25.7 | 25.9 | 5.2 | 0.992x | | 1024 | 8192 | 16 | 25.7 | 25.8 | 10.4 | 0.994x | | 1024 | 8192 | 32 | 25.6 | 25.8 | 20.8 | 0.993x | | 1024 | 8192 | 64 | 25.8 | 25.9 | 41.4 | 0.995x | | 1024 | 8192 | 128 | 27.0 | 28.4 | 75.6 | **0.951x** | | 1024 | 8192 | 256 | 39.0 | 39.6 | 108.4 | 0.985x | | 1024 | 8192 | 512 | 43.1 | 43.5 | 197.6 | 0.991x | | 1024 | 8192 | 1024 | 47.4 | 49.5 | 347.4 | **0.959x** | | 1024 | 8192 | 2048 | 50.2 | 50.0 | 687.0 | 1.004x | | 1024 | 8192 | 4096 | 89.5 | 90.9 | 756.3 | 0.985x | | 2560 | 8192 | 1 | 34.5 | 31.6 | 1.3 | 1.092x | | 2560 | 8192 | 2 | 34.6 | 31.7 | 2.6 | 1.092x | | 2560 | 8192 | 4 | 34.7 | 31.8 | 5.3 | 1.090x | | 2560 | 8192 | 8 | 34.6 | 32.1 | 10.5 | 1.079x | | 2560 | 8192 | 16 | 34.9 | 32.7 | 20.5 | 1.069x | | 2560 | 8192 | 32 | 34.9 | 32.8 | 40.9 | 1.065x | | 2560 | 8192 | 64 | 35.2 | 33.3 | 80.7 | 1.059x | | 2560 | 8192 | 128 | 41.7 | 39.9 | 134.7 | 1.046x | | 2560 | 8192 | 256 | 48.2 | 39.9 | 269.1 | 1.207x | | 2560 | 8192 | 512 | 48.8 | 48.6 | 442.2 | 1.004x | | 2560 | 8192 | 1024 | 50.6 | 51.3 | 837.6 | 0.987x | | 2560 | 8192 | 2048 | 87.5 | 88.2 | 973.8 | 0.992x | | 2560 | 8192 | 4096 | 163.8 | 166.9 | 1029.4 | 0.982x | | 4096 | 4096 | 1 | 23.6 | 22.4 | 1.5 | 1.053x | | 4096 | 4096 | 2 | 23.6 | 22.7 | 3.0 | 1.038x | | 4096 | 4096 | 4 | 23.7 | 22.8 | 5.9 | 1.041x | | 4096 | 4096 | 8 | 23.7 | 23.2 | 11.6 | 1.021x | | 4096 | 4096 | 16 | 23.8 | 23.2 | 23.2 | 1.029x | | 4096 | 4096 | 32 | 24.0 | 23.6 | 45.5 | 1.020x | | 4096 | 4096 | 64 | 24.0 | 23.8 | 90.2 | 1.009x | | 4096 | 4096 | 128 | 31.9 | 31.2 | 137.7 | 1.022x | | 4096 | 4096 | 256 | 35.2 | 35.2 | 244.0 | 1.000x | | 4096 | 4096 | 512 | 36.8 | 36.7 | 467.7 | 1.001x | | 4096 | 4096 | 1024 | 54.3 | 54.5 | 630.0 | 0.995x | | 4096 | 4096 | 2048 | 73.0 | 73.9 | 930.3 | 0.988x | | 4096 | 4096 | 4096 | 134.6 | 137.0 | 1003.4 | 0.983x | | 4096 | 5376 | 1 | 30.9 | 28.9 | 1.5 | 1.068x | | 4096 | 5376 | 2 | 30.8 | 28.8 | 3.1 | 1.069x | | 4096 | 5376 | 4 | 30.8 | 29.4 | 6.0 | 1.046x | | 4096 | 5376 | 8 | 30.8 | 29.2 | 12.1 | 1.055x | | 4096 | 5376 | 16 | 30.7 | 29.5 | 23.9 | 1.042x | | 4096 | 5376 | 32 | 31.0 | 29.9 | 47.1 | 1.035x | | 4096 | 5376 | 64 | 31.1 | 30.0 | 93.9 | 1.037x | | 4096 | 5376 | 128 | 37.6 | 37.0 | 152.5 | 1.018x | | 4096 | 5376 | 256 | 41.4 | 38.3 | 294.7 | 1.083x | | 4096 | 5376 | 512 | 43.2 | 43.2 | 522.2 | 0.999x | | 4096 | 5376 | 1024 | 66.2 | 66.3 | 680.7 | 0.999x | | 4096 | 5376 | 2048 | 92.4 | 92.4 | 975.8 | 0.999x | | 4096 | 5376 | 4096 | 166.6 | 169.3 | 1065.2 | 0.984x | | 4096 | 14336 | 1 | 64.4 | 60.0 | 2.0 | 1.072x | | 4096 | 14336 | 2 | 64.3 | 59.8 | 3.9 | 1.075x | | 4096 | 14336 | 4 | 64.4 | 60.3 | 7.8 | 1.067x | | 4096 | 14336 | 8 | 64.4 | 60.2 | 15.6 | 1.070x | | 4096 | 14336 | 16 | 64.3 | 60.4 | 31.1 | 1.064x | | 4096 | 14336 | 32 | 64.4 | 60.9 | 61.7 | 1.059x | | 4096 | 14336 | 64 | 64.5 | 61.3 | 122.6 | 1.053x | | 4096 | 14336 | 128 | 65.3 | 62.8 | 239.2 | 1.040x | | 4096 | 14336 | 256 | 77.5 | 65.2 | 460.9 | 1.188x | | 4096 | 14336 | 512 | 80.7 | 82.3 | 730.7 | 0.981x | | 4096 | 14336 | 1024 | 139.4 | 140.6 | 855.5 | 0.992x | | 4096 | 14336 | 2048 | 206.0 | 206.8 | 1163.0 | 0.996x | | 4096 | 14336 | 4096 | 395.7 | 398.7 | 1206.5 | 0.992x | | 8192 | 2560 | 1 | 22.2 | 22.0 | 1.9 | 1.012x | | 8192 | 2560 | 2 | 21.5 | 20.7 | 4.1 | 1.043x | | 8192 | 2560 | 4 | 21.4 | 20.4 | 8.2 | 1.049x | | 8192 | 2560 | 8 | 21.9 | 21.7 | 15.5 | 1.008x | | 8192 | 2560 | 16 | 21.2 | 21.7 | 30.9 | 0.977x | | 8192 | 2560 | 32 | 21.6 | 23.6 | 56.9 | **0.915x** | | 8192 | 2560 | 64 | 24.6 | 26.8 | 100.2 | **0.917x** | | 8192 | 2560 | 128 | 25.8 | 29.7 | 180.9 | **0.868x** | | 8192 | 2560 | 256 | 29.1 | 29.6 | 362.2 | 0.981x | | 8192 | 2560 | 512 | 43.0 | 42.8 | 501.7 | 1.005x | | 8192 | 2560 | 1024 | 54.7 | 55.2 | 777.4 | 0.989x | | 8192 | 2560 | 2048 | 93.2 | 93.7 | 916.9 | 0.994x | | 8192 | 2560 | 4096 | 163.2 | 167.5 | 1025.7 | 0.974x | | 14336 | 4096 | 1 | 49.1 | 40.3 | 2.9 | 1.218x | | 14336 | 4096 | 2 | 48.3 | 39.1 | 6.0 | 1.235x | | 14336 | 4096 | 4 | 49.2 | 39.9 | 11.8 | 1.231x | | 14336 | 4096 | 8 | 47.8 | 39.5 | 23.8 | 1.211x | | 14336 | 4096 | 16 | 46.0 | 39.7 | 47.3 | 1.158x | | 14336 | 4096 | 32 | 44.5 | 43.1 | 87.3 | 1.033x | | 14336 | 4096 | 64 | 40.7 | 44.2 | 170.1 | **0.922x** | | 14336 | 4096 | 128 | 43.4 | 45.3 | 332.0 | **0.959x** | | 14336 | 4096 | 256 | 59.5 | 60.7 | 495.0 | 0.979x | | 14336 | 4096 | 512 | 80.0 | 79.7 | 754.6 | 1.004x | | 14336 | 4096 | 1024 | 116.2 | 116.5 | 1032.0 | 0.997x | | 14336 | 4096 | 2048 | 212.8 | 214.3 | 1122.6 | 0.993x | | 14336 | 4096 | 4096 | 390.8 | 395.2 | 1217.2 | 0.989x | </details> <details> <summary><b>DGX Spark / GB10 (SM121)</b></summary> | n | k | m | old (us) | new (us) | new TFLOPs | speedup | |---|---|---|---|---|---|---| | 512 | 7168 | 1 | 31.1 | 31.3 | 0.2 | 0.991x | | 512 | 7168 | 2 | 31.5 | 31.4 | 0.5 | 1.003x | | 512 | 7168 | 4 | 31.6 | 31.1 | 0.9 | 1.016x | | 512 | 7168 | 8 | 31.6 | 31.5 | 1.9 | 1.005x | | 512 | 7168 | 16 | 31.7 | 31.8 | 3.7 | 0.997x | | 512 | 7168 | 32 | 32.5 | 33.7 | 7.0 | **0.964x** | | 512 | 7168 | 64 | 34.1 | 33.8 | 13.9 | 1.009x | | 512 | 7168 | 128 | 49.0 | 39.3 | 23.9 | 1.245x | | 512 | 7168 | 256 | 60.7 | 49.5 | 38.0 | 1.227x | | 512 | 7168 | 512 | 56.8 | 56.3 | 66.7 | 1.008x | | 512 | 7168 | 1024 | 73.1 | 73.6 | 102.2 | 0.994x | | 512 | 7168 | 2048 | 129.0 | 129.0 | 116.5 | 1.000x | | 512 | 7168 | 4096 | 227.0 | 225.8 | 133.1 | 1.005x | | 1024 | 8192 | 1 | 65.4 | 59.8 | 0.3 | 1.094x | | 1024 | 8192 | 2 | 64.7 | 60.2 | 0.6 | 1.075x | | 1024 | 8192 | 4 | 65.9 | 59.5 | 1.1 | 1.108x | | 1024 | 8192 | 8 | 64.1 | 60.5 | 2.2 | 1.058x | | 1024 | 8192 | 16 | 63.6 | 57.1 | 4.7 | 1.114x | | 1024 | 8192 | 32 | 66.5 | 58.5 | 9.2 | 1.136x | | 1024 | 8192 | 64 | 55.5 | 58.1 | 18.5 | **0.956x** | | 1024 | 8192 | 128 | 72.0 | 64.4 | 33.3 | 1.119x | | 1024 | 8192 | 256 | 83.1 | 79.1 | 54.3 | 1.050x | | 1024 | 8192 | 512 | 81.5 | 84.5 | 101.6 | **0.964x** | | 1024 | 8192 | 1024 | 155.6 | 130.2 | 132.0 | 1.195x | | 1024 | 8192 | 2048 | 200.4 | 191.2 | 179.7 | 1.048x | | 1024 | 8192 | 4096 | 456.5 | 410.5 | 167.4 | 1.112x | | 2560 | 8192 | 1 | 117.8 | 118.0 | 0.4 | 0.998x | | 2560 | 8192 | 2 | 116.2 | 119.2 | 0.7 | 0.975x | | 2560 | 8192 | 4 | 118.0 | 118.2 | 1.4 | 0.998x | | 2560 | 8192 | 8 | 116.2 | 119.4 | 2.8 | 0.973x | | 2560 | 8192 | 16 | 118.2 | 118.5 | 5.7 | 0.998x | | 2560 | 8192 | 32 | 117.1 | 119.8 | 11.2 | 0.977x | | 2560 | 8192 | 64 | 117.7 | 120.7 | 22.2 | 0.975x | | 2560 | 8192 | 128 | 119.6 | 124.2 | 43.2 | **0.963x** | | 2560 | 8192 | 256 | 124.7 | 128.5 | 83.6 | 0.970x | | 2560 | 8192 | 512 | 151.2 | 146.9 | 146.1 | 1.029x | | 2560 | 8192 | 1024 | 212.2 | 206.8 | 207.6 | 1.026x | | 2560 | 8192 | 2048 | 346.6 | 337.4 | 254.6 | 1.027x | | 2560 | 8192 | 4096 | 951.6 | 861.4 | 199.4 | 1.105x | | 4096 | 4096 | 1 | 94.2 | 98.6 | 0.3 | **0.956x** | | 4096 | 4096 | 2 | 113.4 | 99.0 | 0.7 | 1.146x | | 4096 | 4096 | 4 | 111.6 | 101.8 | 1.3 | 1.096x | | 4096 | 4096 | 8 | 113.1 | 97.4 | 2.8 | 1.161x | | 4096 | 4096 | 16 | 113.5 | 101.6 | 5.3 | 1.117x | | 4096 | 4096 | 32 | 113.6 | 99.6 | 10.8 | 1.141x | | 4096 | 4096 | 64 | 112.5 | 100.6 | 21.3 | 1.119x | | 4096 | 4096 | 128 | 111.9 | 99.2 | 43.3 | 1.128x | | 4096 | 4096 | 256 | 107.3 | 108.4 | 79.3 | 0.991x | | 4096 | 4096 | 512 | 115.5 | 112.7 | 152.5 | 1.025x | | 4096 | 4096 | 1024 | 158.6 | 158.4 | 216.9 | 1.001x | | 4096 | 4096 | 2048 | 281.3 | 272.3 | 252.3 | 1.033x | | 4096 | 4096 | 4096 | 508.9 | 494.9 | 277.7 | 1.028x | | 4096 | 5376 | 1 | 117.4 | 118.8 | 0.4 | 0.988x | | 4096 | 5376 | 2 | 116.6 | 117.4 | 0.8 | 0.994x | | 4096 | 5376 | 4 | 117.0 | 118.4 | 1.5 | 0.989x | | 4096 | 5376 | 8 | 116.3 | 117.7 | 3.0 | 0.988x | | 4096 | 5376 | 16 | 118.6 | 118.8 | 5.9 | 0.998x | | 4096 | 5376 | 32 | 116.5 | 118.3 | 11.9 | 0.985x | | 4096 | 5376 | 64 | 117.8 | 119.1 | 23.7 | 0.989x | | 4096 | 5376 | 128 | 118.7 | 119.9 | 47.0 | 0.991x | | 4096 | 5376 | 256 | 129.1 | 132.6 | 85.0 | 0.973x | | 4096 | 5376 | 512 | 139.8 | 142.2 | 158.6 | 0.983x | | 4096 | 5376 | 1024 | 200.0 | 197.5 | 228.3 | 1.012x | | 4096 | 5376 | 2048 | 338.3 | 340.3 | 265.1 | 0.994x | | 4096 | 5376 | 4096 | 693.6 | 645.1 | 279.6 | 1.075x | | 4096 | 14336 | 1 | 231.6 | 231.8 | 0.5 | 0.999x | | 4096 | 14336 | 2 | 229.6 | 229.4 | 1.0 | 1.001x | | 4096 | 14336 | 4 | 232.6 | 232.5 | 2.0 | 1.000x | | 4096 | 14336 | 8 | 232.1 | 228.7 | 4.1 | 1.015x | | 4096 | 14336 | 16 | 230.7 | 232.7 | 8.1 | 0.991x | | 4096 | 14336 | 32 | 230.2 | 228.7 | 16.4 | 1.007x | | 4096 | 14336 | 64 | 231.9 | 232.7 | 32.3 | 0.997x | | 4096 | 14336 | 128 | 234.1 | 234.2 | 64.2 | 1.000x | | 4096 | 14336 | 256 | 266.6 | 265.2 | 113.4 | 1.005x | | 4096 | 14336 | 512 | 303.0 | 302.6 | 198.7 | 1.001x | | 4096 | 14336 | 1024 | 469.2 | 466.4 | 257.9 | 1.006x | | 4096 | 14336 | 2048 | 902.7 | 880.0 | 273.3 | 1.026x | | 4096 | 14336 | 4096 | 3847.1 | 3765.8 | 127.7 | 1.022x | | 8192 | 2560 | 1 | 113.6 | 113.0 | 0.4 | 1.005x | | 8192 | 2560 | 2 | 111.7 | 113.4 | 0.7 | 0.986x | | 8192 | 2560 | 4 | 113.2 | 113.3 | 1.5 | 0.999x | | 8192 | 2560 | 8 | 113.2 | 113.8 | 2.9 | 0.995x | | 8192 | 2560 | 16 | 114.2 | 113.5 | 5.9 | 1.006x | | 8192 | 2560 | 32 | 113.6 | 114.5 | 11.7 | 0.993x | | 8192 | 2560 | 64 | 114.2 | 115.4 | 23.3 | 0.990x | | 8192 | 2560 | 128 | 116.4 | 119.2 | 45.0 | 0.976x | | 8192 | 2560 | 256 | 119.7 | 121.7 | 88.2 | 0.984x | | 8192 | 2560 | 512 | 135.1 | 139.3 | 154.2 | 0.970x | | 8192 | 2560 | 1024 | 200.1 | 199.5 | 215.3 | 1.003x | | 8192 | 2560 | 2048 | 370.9 | 371.6 | 231.1 | 0.998x | | 8192 | 2560 | 4096 | 644.3 | 641.7 | 267.7 | 1.004x | | 14336 | 4096 | 1 | 232.3 | 226.1 | 0.5 | 1.028x | | 14336 | 4096 | 2 | 232.5 | 226.5 | 1.0 | 1.026x | | 14336 | 4096 | 4 | 231.8 | 225.8 | 2.1 | 1.026x | | 14336 | 4096 | 8 | 232.2 | 226.7 | 4.1 | 1.024x | | 14336 | 4096 | 16 | 232.7 | 226.7 | 8.3 | 1.026x | | 14336 | 4096 | 32 | 232.1 | 226.4 | 16.6 | 1.025x | | 14336 | 4096 | 64 | 233.3 | 229.0 | 32.8 | 1.019x | | 14336 | 4096 | 128 | 235.8 | 232.5 | 64.7 | 1.014x | | 14336 | 4096 | 256 | 246.5 | 239.6 | 125.5 | 1.029x | | 14336 | 4096 | 512 | 293.4 | 287.1 | 209.4 | 1.022x | | 14336 | 4096 | 1024 | 441.3 | 441.2 | 272.6 | 1.000x | | 14336 | 4096 | 2048 | 838.2 | 836.7 | 287.5 | 1.002x | | 14336 | 4096 | 4096 | 1500.5 | 1506.2 | 319.4 | 0.996x | </details> ## 🔍 Related Issues Closes #3517. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. *(covered by existing `tests/gemm/test_mm_fp4.py -k b12x`)* - [x] All tests are passing — `pytest -k b12x` 396/0; 104-shape `--refcheck` 104/104 on 5080 / PRO 6000 / GB10. ## Reviewer Notes - The vendored kernel is **not a verbatim upstream copy**: it carries backported scale-factor helpers so it builds on public `cutlass-dsl 4.5.2` (pure upstream b12x needs an internal/newer cutlass-dsl), and the MXFP8 path is pruned (unsupported there). - The SFB-layout N-divisibility assertion in `cute_dsl/utils.py` is relaxed **64→16** (upstream commit `0daa6ab`) to admit the narrow-N `swap_ab` tiles. This is safe: it only *loosens* a precondition (every previously-valid N%64==0 tile still passes), the SFB smem layout already rounds N up to a full 128-wide block so the layout is unchanged for existing tiles, and the newly-admitted narrow tiles pass `--refcheck`. - `swap_ab` here is b12x's **device-internal** transpose (public C stays row-major) — *not* the SM100 operand-swap/FFI-transpose convention. The shared `_compile_block_scaled_gemm` harness used by Sm100/Sm103 is **left untouched**. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Improvements** * Relaxed a tile-shape constraint to allow more flexible tensor layouts. * Deterministic plan selection for low‑precision GEMM with a smaller, predictable candidate set. * Earlier, clearer validation for incompatible contraction dimensions to surface errors sooner. * Consistent operand orientation and compilation caching to reduce surprises in kernel selection and runtime behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Author
|
I have checked and discover the 'restat' ninja option, it could correct this problem in a better way. |
…s on GPFS/NFS On shared filesystems (GPFS, NFS), the mtimes ninja records in .ninja_log at the end of a build can disagree with what a later stat() returns (inode flush timing). On the next process start, ninja sees the mismatch, considers every output "modified externally", and recompiles all CUTLASS fused-MoE kernels (~180 .cu files, ~1.5h on GH200/sm90a) on every vLLM restart. Fix: run `ninja -t restat` (available since ninja 1.10) right before run_ninja(). It rewrites the mtimes stored in .ninja_log with current stat() values, so the incremental check passes, while ninja's staleness detection stays fully intact: modified sources or changed command lines still trigger rebuilds. A restat failure is non-fatal (check=False) and falls back to the current behavior. Replace PR #3556.
Rehvaro
added a commit
to Rehvaro/flashinfer
that referenced
this pull request
Jun 11, 2026
…s on GPFS/NFS On shared filesystems (GPFS, NFS), the mtimes ninja records in .ninja_log at the end of a build can disagree with what a later stat() returns (inode flush timing). On the next process start, ninja sees the mismatch, considers every output "modified externally", and recompiles all CUTLASS fused-MoE kernels (~180 .cu files, ~1.5h on GH200/sm90a) on every vLLM restart. Fix: run `ninja -t restat` (available since ninja 1.10) right before run_ninja(). It rewrites the mtimes stored in .ninja_log with current stat() values, so the incremental check passes, while ninja's staleness detection stays fully intact: modified sources or changed command lines still trigger rebuilds. A restat failure is non-fatal (check=False) and falls back to the current behavior. Replace flashinfer-ai#3556.
5 tasks
Author
|
New PR #3579 created with what I believe is a better solution. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
13 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
On shared filesystems (GPFS, NFS), every vLLM restart triggers a full recompilation of the FlashInfer CUTLASS fused-MoE kernels (~180 .cu files). On GH200/sm90a this takes around 1.5 hours, making the cache completely useless.
In details :
build_and_load() always calls build() regardless of whether the .so is already present. build() always runs ninja. On a local filesystem
ninja's incremental check works fine — it compares the mtime stored in .ninja_log against stat() of the outputs, finds them equal, exits in under a second.
On GPFS the two values don't match. When ninja finishes a build it does a restat pass: stat() on each output, records the mtime in .ninja_log.
On GPFS this returns the time the inode was committed to storage, which can be tens of minutes after the compiler actually wrote the file. The next time ninja runs (second TP worker, or any subsequent restart) it calls stat() again and gets the original write-time mtime from the kernel page cache, a completely different value. Ninja sees the mismatch, decides the files were modified xternally, and recompiles.
If the .so already exists, skip the whole thing and load it directly. The expensive path is only needed on the first ever compilation. An inner exists() check is added inside the lock to handle concurrent workers both hitting the first-run case at the same time.
Tested on GLM-4.7-Flash (MoE BF16, TP=2) on GH200 with GPFS home. Before: full 1.5h recompile on every restart. After: loads in a few seconds.
🔍 Related Issues
/
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.Everything passed, except things I think aren't related to my PR :
🧪 Tests
unittest, etc.).No Tests
Reviewer Notes
I'm not a developer, only a 'jack-of-all-trades'.
If this PR does not correspond to the quality standard, I hope at least the underlying issue will allow someone more qualify than me to patch it correctly with a correct PR.
Summary by CodeRabbit
Bug Fixes
Performance
New Features
Behavioral
Tests & Docs