Record the v0.6.18 merge that #344's squash discarded - #345
Merged
Conversation
<!-- .github/pull_request_template.md --> ## 📌 Description Move JIT symlinks from `FLASHINFER_CUBIN_DIR` (read-only when installed system-wide) to `FLASHINFER_GEN_SRC_DIR` (always user-writable). Add `FLASHINFER_GEN_SRC_DIR` to `extra_include_paths` so C++ resolution is unchanged. Symlink targets still point into `FLASHINFER_CUBIN_DIR`. ## 🔍 Related Issues - fix flashinfer-ai#2834 ## 🚀 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 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Improved internal build infrastructure for JIT compilation by reorganizing generated source caching and header inclusion paths to enhance build consistency and modularity. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…flashinfer-ai#3759) # [comm] Add FP8 quantized two-shot AllReduce via symmetric memory ## Summary Adds `flashinfer.comm.quantized_all_reduce()`, a Triton kernel that halves AllReduce transfer volume by quantizing activations to FP8 (1 byte + scale per group) instead of BF16 (2 bytes) before P2P transfer over symmetric memory. ## Algorithm Two-shot with 2 barriers: 1. Each rank quantizes BF16 → FP8 + FP32 inverse scale, writes to symmetric memory 2. Barrier 3. Each rank reduces its 1/W stripe: reads FP8 from all W peers, dequants (multiply by inverse scale), sums in FP32, writes BF16 to output + re-quantizes for Phase 3 4. Barrier 5. Each rank dequants other W-1 stripes from FP8 → BF16 ## Results Hardware: 8x H200 NVLink, PyTorch 2.12, Triton 3.7. Requirement: SM90+ (Hopper or later), NVSwitch topology. ### Latency ```bash mpirun -np 8 python benchmarks/comm/bench_quantized_allreduce.py --latency-only --warmup 20 --iters 30 ``` Method: `bench_gpu_time(use_cuda_graph=True, num_iters_within_graph=10)`. Reports median across iterations. ```python # NCCL baseline (pre-allocated output) out.copy_(inp) dist.all_reduce(out) # symm_mem baseline (NVSwitch hardware multicast) buf.copy_(inp) torch.ops.symm_mem.multimem_all_reduce_(buf, "sum", group) # This PR quantized_all_reduce(inp, group) ``` | Size | NCCL | symm_mem | FP8 Quant | vs NCCL | vs symm_mem | |------|------|----------|-----------|---------|-------------| | 2KB | 16.6us | 8.5us | 11.6us | 1.44x | 0.74x | | 8KB | 17.4us | 8.5us | 11.6us | 1.50x | 0.73x | | 32KB | 19.7us | 8.7us | 12.8us | 1.54x | 0.68x | | 128KB | 20.3us | 9.6us | 15.3us | 1.33x | 0.63x | | 512KB | 20.8us | 11.3us | 15.9us | 1.31x | 0.71x | | 2MB | 40.0us | 18.2us | 23.1us | 1.73x | 0.79x | | 8MB | 80.6us | 46.1us | **41.4us** | **1.95x** | **1.11x** | | 32MB | 198.8us | 166.0us | **149.8us** | **1.33x** | **1.11x** | | 128MB | 626.5us | 629.9us | **523.8us** | **1.20x** | **1.20x** | | 512MB | 2444us | 2755us | **1833us** | **1.33x** | **1.50x** | | 2GB | 9518us | 10950us | **7220us** | **1.32x** | **1.52x** | Effective bandwidth = `message_size_bytes / median_time`. Peak: 297 GB/s at 2GB (vs NCCL 226, symm_mem 196). ### Accuracy ```bash mpirun -np 8 python benchmarks/comm/bench_quantized_allreduce.py --error-only ``` Reference: BF16 NCCL allreduce. 5 independent trials at 8MB (4.2M elements). Relative error computed only over elements where |ref| > 0.01. Full sweep (32KB–128MB) available via `--error-only`. | Distribution | mean_rel | max_abs | |---|---|---| | gelu(x), x ~ N(0,1) | 3.7% ± 0.00% | 0.63 ± 0.03 | | sparse (90% zero) | 4.0% ± 0.01% | 0.35 ± 0.02 | | silu(x) * y, x,y ~ N(0,1) | 6.9% ± 0.01% | 0.64 ± 0.05 | | uniform [-1,1] | 8.7% ± 0.01% | 0.35 ± 0.01 | | heavy_tail (1% spikes at 100x) | 9.0% ± 0.01% | 20.4 ± 1.5 | | normal N(0,1) | 9.5% ± 0.02% | 0.69 ± 0.03 | ## When to use Best above 8MB where it beats both NCCL and symm_mem. Below that, `symm_mem.multimem_all_reduce_` is faster (NVSwitch hardware multicast, no quantization overhead). At large sizes, NVLink is the bottleneck: SMs are idle waiting for data. This kernel uses those idle cycles to quantize/dequant, transferring half the bytes. At small sizes, the fixed overhead of 2 cross-GPU barriers dominates. ## Attribution Kernel design and PTX barrier derived from [Kraken's two-shot AllReduce](https://github.com/meta-pytorch/kraken/blob/main/kraken/comm/two_shot_all_reduce.py). This PR adds per-group FP8 quantization on top of that protocol. ## Follow-up ideas - **INT8 quantization**: Same 2× bandwidth reduction but works on A100 (SM80) which has NVLink but no FP8 hardware. Trivial change (`tl.int8`, QMAX=127). Also works on Hopper/Blackwell as a portable alternative. - **4-bit quantization**: 4× bandwidth reduction via packing 2 elements per byte. Requires accuracy mitigation (e.g., spike reservation) as naive 4-bit degrades model quality significantly. ## Test plan - [x] Correctness: 2 GPU and 8 GPU (sizes 32KB–64MB) - [x] Scale group sweep: 128, 256, 512, 1024 - [x] Edge cases: minimum numel, exact block boundaries - [x] CI: 14/20 passed on NVIDIA internal CI (`/bot run`)
…ge path (flashinfer-ai#3062) ## What's broken? When `flashinfer-cubin` is installed (all pre-built container images), setting `FLASHINFER_CUBIN_DIR` to redirect cubin storage to a writable directory is silently ignored. Non-root containers crash with `PermissionError` on startup. ## Who is affected? All non-root K8s/OpenShift deployments using pre-built images with `flashinfer-cubin` (e.g., vLLM containers). No workaround exists. ## Why does it happen? `_get_cubin_dir()` checks the `flashinfer-cubin` package **before** the env var. When the package is installed, the env var is never reached. Introduced in PR flashinfer-ai#1718 (intentional priority, but didn't consider non-root containers). Every other env var in `env.py` follows "env var overrides automatic discovery." This was the sole exception. ## How did we fix it? Swap priority: env var -> package -> default cache. Pure code-block reorder (~8 lines moved + docstring/comments updated). Zero new logic. **Before:** ```python def _get_cubin_dir(): if has_flashinfer_cubin(): # always wins return package_path env_dir = os.getenv("FLASHINFER_CUBIN_DIR") # unreachable ... ``` **After:** ```python def _get_cubin_dir(): env_dir = os.getenv("FLASHINFER_CUBIN_DIR") # user override first if env_dir: return pathlib.Path(env_dir) if has_flashinfer_cubin(): # package fallback return package_path ... ``` Default behavior (no env var set) is **unchanged**. > **Note:** If you previously had a stale `FLASHINFER_CUBIN_DIR` that was silently ignored, it will now take effect. Unset it if unneeded. ## How do we know it works? Added `tests/test_env.py` with 4 regression tests covering the full priority matrix: - env var + package -> env var wins (the bug) - package only -> package wins - env var only -> env var wins - neither -> default cache dir All 4 pass. Pre-commit hooks (ruff, mypy, format) all pass. Fixes flashinfer-ai#2976 Related: flashinfer-ai#2834 @aleozlx @yzh119 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved CUBIN directory resolution to honor `FLASHINFER_CUBIN_DIR` first, with fallback to the installed `flashinfer_cubin` location when unset. * When `FLASHINFER_CUBIN_DIR` is set alongside the installed package, a warning is emitted to clarify which source is used. * **Tests** * Added regression tests covering CUBIN directory selection priority across environment and package availability scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: jimmzhou <jimmzhou@nvidia.com>
## 📌 Description Add TRTLLM per-tensor FP8 execution to the unified `MoELayer` API. Also fixes the MXFP8 scale-layout issue reported in flashinfer-ai#4087. ### What changed * Added `TrtllmFp8PerTensorRunner` to the unified `MoELayer` API with: * `FromLogits` in-kernel routing * SM100 and SM103 support * Autotuning and CUDA graphs * Llama4 routing-scale-on-input behavior * Added per-tensor FP8 weight and activation preparation: * Per-expert E4M3 weight quantization * Calibrated per-tensor activation quantization * TRTLLM gated-row reorder and shuffled MajorK weights * Per-expert GEMM1 linear/gate and GEMM2 epilogue scales * Registered `QuantVariant.FP8PerTensor` with `MoELayer`. * Added independent-reference, routing-replay, CUDA-graph, architecture-gating, and unified-fuzzer coverage. * Fixed unified MXFP8 block-scale preparation by converting row-permuted GEMM1/GEMM2 scale tensors into TRTLLM’s required 128×4 interleaved layout. ## 🔍 Scope and follow-up Unified quantized MoE support is split into separate PRs because the formats use different kernels, scaling conventions, architectures, and execution contracts: * PR 1: block-scale DeepSeek FP8 and MXFP8 — merged in flashinfer-ai#4026 * PR 2: per-tensor FP8 — this PR * PR 3: CUTLASS W4A8 — follow-up This PR keeps per-tensor FP8 `FromLogits`-only, matching the existing TRTLLM kernel entry point. The legacy flat APIs remain unchanged. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks * I have installed `pre-commit` by running `pip install pre-commit` (or used my 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. ## 🧪 Tests Validated on SM100: * `18 passed` — unified block-FP8 and per-tensor-FP8 conformance suites * `7 passed` — focused unified API/config validation * `1 passed` — curated per-tensor FP8 fuzzer profile with production autotuning * `1 passed` — MXFP8 seed `900013` regression, including valid-tactic coverage * `1 passed` — exact reported MXFP8 seed-100 configuration, including valid-tactic coverage ### MXFP8 flashinfer-ai#4087 validation The fuzzer and legacy MXFP8 references agree to within `1.16e-10`. Unified and legacy weight payloads were identical, but the unified scale tensors were missing the 128×4 physical interleave. With the corrected scale preparation: ```text before: 110 / 524288 elements over tolerance after: 0 / 524288 elements over tolerance ``` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added TensorRT-LLM FP8 per-tensor quantization for Mixture-of-Experts, including integration into the cross-backend MoE layer. * Exposed a new FP8 per-tensor runner via the public API (FromLogits routing). * Added calibrated global scaling support for FP8 weights and activations. * **Bug Fixes** * Tightened FP8/MXFP8 preparation validation for MXFP8 and clarified the FP8 layout constraints. * Refined hardware support to SM100 family values (supported: 100/103; unsupported: 90/120). * **Tests** * Expanded FP8 per-tensor correctness and replay coverage, plus new negative/behavioral tests. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
<!-- .github/pull_request_template.md --> ## 📌 Description Update CODEOWNERS to add @feih-nv and @Aneureka <!-- 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** * Updated code ownership assignments to expand review coverage for core, MOE, and MOE_EP areas. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…idential Computing (flashinfer-ai#3993) ## Motivation Under NVIDIA Confidential Computing (CC), the symmetric-memory allocator's `cuMulticast` setup fails (the bounce-buffer path can't complete the fabric/multicast rendezvous), so the TRT-LLM AllReduce-fusion workspace can't be created the usual way. ## What this changes `create_allreduce_fusion_workspace` / `TRTLLMAllReduceFusionWorkspace` now allocate a **multicast-free IPC workspace** instead of symmetric device memory whenever CC is detected — auto-detected via `is_confidential_compute()` (overridable with `FLASHINFER_CONFIDENTIAL_COMPUTE`). - The trtllm one-shot Lamport and two-shot sync fusion kernels are **both multicast-free** (0 `multimem` in `trtllm_allreduce_fusion.cuh`), and the IPC workspace is sized identically to the symmetric one, so both strategies run on it — only the allocator differs. - The shorter `(ipc_handles, workspace_tensor, metadata)` return tuple is handled, and `mem_handles` is set to `[]` so the "handles attached?" guard in `allreduce_fusion` and `destroy()` iterate a no-op. - The `mnnvl` backend requires NVLink multicast (unavailable under CC), so `create_allreduce_fusion_workspace` raises there rather than failing deep inside workspace creation. ## Tests `tests/comm/test_allreduce_unified_api.py`: - **`test_allreduce_trtllm_multicast_free`** — forces CC via `FLASHINFER_CONFIDENTIAL_COMPUTE=1`, exercising both kernels (`use_oneshot` True/False), both AR patterns, and fp16/bf16. - **`test_mnnvl_raises_under_cc`** — asserts the mnnvl+CC guard. Verified 17/17 on 2×B200. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
<!-- .github/pull_request_template.md -->
## 📌 Description
FlashInfer provides fused norm + FP8 quantization for the RMSNorm family
(`rmsnorm_quant`, `fused_add_rmsnorm_quant`), but not for LayerNorm.
Models that use LayerNorm (OPT, Falcon, BERT family, ViT/DiT variants,
Whisper/CLIP style encoders) currently have to run layernorm and
quantization as two kernels, which costs about 7 bytes of traffic per
element instead of 3.
The kernel side is already there: `generalLayerNorm` in
`include/flashinfer/norm.cuh` implements per-tensor and per-token
quantization paths, but the `LayerNorm` host launcher passes nullptr for
all quant arguments, with a note:
// TODO(kaixih): add support for fp8 quantization if needed
This PR connects the existing per-tensor path rather than writing a new
kernel, so the diff stays small.
- Adds `flashinfer.layernorm_quant(out, input, gemma, beta, scale,
eps)`. Output dtype (float8_e4m3fn or float8_e5m2) is taken from the
preallocated `out` tensor, same as `rmsnorm_quant`.
- Adds fp8 `QuantTypeStaticVals` specializations (constants match
TRT-LLM `quantTypeUtils.cuh`) and the missing e5m2 `cuda_cast`
specializations. The FP8 dispatch macro instantiates both e4m3 and e5m2,
so the e5m2 casts are needed for compilation.
- Adds a `LayerNormQuant` host launcher, guarded by `ENABLE_FP8`. The
csrc launcher keeps a non-FP8 build working
(`test_norm_compilation_without_fp8` passes).
- Changes the per-tensor scale in `generalLayerNorm` to be applied as
division (`out = normed / scale`). The kernel originally multiplied, but
this path was unreachable from every existing call site (the pointer is
always nullptr), so no current behavior changes. Keeping the multiply
semantics would make `layernorm_quant` and `rmsnorm_quant` disagree on
what `scale` means, which seemed worse for fusion passes that target
both. Happy to adjust if there is a concern I am missing here.
- Trace template, docs entry, tests, and a `layernorm_quant` benchmark
routine.
## Performance
RTX 5090, bf16 input, e4m3 output, eps 1e-6. Fused kernel vs. unfused
baseline (`flashinfer.layernorm` followed by eager torch
div/clamp/cast), median over 100 iterations under CUDA graph:
| batch | hidden | unfused (us) | fused (us) | speedup |
|------:|-------:|-------------:|-----------:|--------:|
| 256 | 4096 | 13.2 | 6.9 | 1.91x |
| 1024 | 4096 | 36.9 | 21.4 | 1.73x |
| 4096 | 4096 | 178.9 | 71.8 | 2.49x |
| 256 | 8192 | 21.1 | 8.4 | 2.52x |
| 1024 | 8192 | 69.2 | 24.9 | 2.77x |
| 4096 | 8192 | 499.6 | 80.6 | 6.20x |
| 1024 | 16384 | 167.2 | 43.8 | 3.82x |
| 4096 | 16384 | 1066.5 | 148.7 | 7.17x |
The eager baseline materializes intermediate tensors for div/clamp/cast,
so large shapes gain more than the pure traffic ratio (7B vs 3B per
element, about 2.3x). At batch 4096 / hidden 16384 the fused kernel
reaches about 1.35 TB/s on this card.
Reproduce with:
python benchmarks/flashinfer_benchmark.py --routine layernorm_quant \
--batch_size 1024 --hidden_size 4096 --refcheck
## 🔍 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
- [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/utils/test_norm.py::test_layernorm_quant`: e4m3 and e5m2, batch
1 to 989, hidden 111 to 16384 (odd sizes cover the non-vectorized path),
scale 0.01/1.0/10.0, tolerance rtol=1/atol=1 following the existing FP8
norm tests. 168 cases pass.
- `test_layernorm_quant_invalid_inputs`: rejects fp16 input,
non-contiguous input, and non-scalar scale.
- `test_norm_compilation_without_fp8` passes unchanged (norm module
still builds without `ENABLE_FP8`).
- Full `tests/utils/test_norm.py` passes (2887 tests).
- `pytest tests/trace/` template consistency and end-to-end tests pass
for the new trace.
- CUDA graph capture/replay smoke test passes; scale is read on device
at kernel execution time, so updating the scale tensor between replays
works.
- [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
* **New Features**
* Added `flashinfer.layernorm_quant` (LayerNorm + FP8 quantization) with
CUDA support across supported compute capabilities.
* Exposed the API at both `flashinfer.norm` and the top-level
`flashinfer` namespace.
* Added benchmark and trace coverage for the quantized LayerNorm
operation.
* **Documentation**
* Updated normalization API docs to include `layernorm_quant`.
* **Tests**
* Added correctness/invalid-input tests, plus trace example coverage and
expected outputs for `layernorm_quant`.
* **Bug Fixes**
* Improved trace output dtype resolution when traced tensors are not
available.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
<!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> Add @Anerudhan to communication kernel code owners. ## 🔍 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 * **Chores** * Updated review ownership for communication-related code and tests to include an additional required reviewer. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…nce (flashinfer-ai#3952) # CuTe-DSL modular prefill: sliding-window support + band-mask performance ## Summary The modular CuTe-DSL prefill (SM100) mapped `causal=True + window_left >= 0` to a plain causal mask — wrong results, and full O(s²/2) work. This PR fixes the mask semantics with composable per-row bands, removes the banded regime's dead work, makes window sizes runtime arguments (one compiled kernel per mask kind), and routes windowed plans to this path. It also adds native LSE output and mixed K/V dtype support (bf16/fp16 Q/K + fp8 V) to the modular kernel, so windowed plans never leave it at run time. ## Changes - **Composable band masks (correctness fix)**: the `MaskType` enum becomes a band `[lo(q), hi(q))` — causal composes with `window_left`; `window_left` alone is a left-bounded non-causal window (FA2 semantics); new `window_right` plan() parameter for symmetric windows. Loader / MMA / softmax / correction all derive trip counts from the same band helpers. - **Dead work removal**: per-half band iteration (head/tail peel) — each 128-row half of the CTA runs exactly its own band's trips instead of the union with ~1 dead step per stage; warp-ballot rescale skip in correction (skips the O-tile TMEM roundtrip when the factor is exactly 1). - **Mask codegen**: hoisted row-invariant bounds, trace-time-constant K coordinates, `cutlass.select_` for masked writes; one shared path for all four mask kinds (avoids an MLIR lowering pathology that caused register spills on the release toolchain). - **Runtime window values**: `MaskSpec` is two presence bits; window sizes are runtime `Int32` kernel args, causal is a right bound with value 0, and the `k < seqlen_k` tail check is always on. Mask compile matrix: 4 kernels. - **Scheduling by mask kind**: banded masks default to non-persistent launch, unmasked stays persistent (`FLASHINFER_CUTE_PREFILL_PERSISTENT` overrides). - **`sm_scale` default fix**: cute-dsl wrapper plan() now defaults to `1/sqrt(head_dim_qk)` (was 1.0, silently unscaled). - **Vendored FMHA windowed compile fix**: the flashinfer-ai#3857 path crashed on any `window_left >= 0` (no window axis in the trace signature, compile cache, or artifact names); it now JIT-compiles with window-presence bits. - **Native LSE**: `run(return_lse=True)` is computed in the correction warp from the final (row_sum, row_max) — log2-domain, matching the vendored kernel's convention. Zero cost when off (the non-LSE kernel's SASS is byte-identical) and ≲1% when on. Supported for score-mod and sink variants; logits-transform variants (no softmax) raise. - **Mixed K/V dtypes**: bf16/fp16 Q/K with fp8 V. P converts to V's dtype for the PV MMA, V slots on the shared K/V smem ring re-arm their barrier with V's byte count, and the wrapper JIT-compiles per-V-dtype kernel variants at run time (V's dtype is a run-time property; plan()'s `kv_data_type` declares K and V together). - **Routing**: `variant is None and head_dim_qk == 128 and window_left < 0` routes to the flashinfer-ai#3857 vendored kernel; everything else — windows, variants, other head dims, any V dtype, `return_lse` — runs the modular path. The gate is decided entirely at plan(); there is no run-time fallback. ## Correctness New causal+window, left-only, and symmetric-window tests in `tests/attention/test_modular_fmha_prefill.py`, including windowed masks over mixed-length ragged batches (each sequence gets its own band geometry; short sequences exercise the item-skip path). Validated on B200 against an exact f32 reference over {8/8, 32/8, 64/8, 32/32} heads × 5 shapes × 4 masks in both scheduler modes (152/152), with bitwise-identical outputs across scheduler modes. Also new: uniform-fp8 (e4m3) tests calibrated against the vendored kernel's fp8 error on identical inputs (both ~0.066 max vs f32); mixed bf16-QK/fp8-V tests against a dequantized-V f32 reference (~0.05 max — S stays bf16-exact); LSE tests (log2-domain, f32 logsumexp references) covering windowed, ragged varlen, ALiBi, and sink variants, with the output tensor bit-identical between LSE and non-LSE runs. ## Performance vs trtllm-gen B200 (SM100a), torch 2.11+cu130, `nvidia-cutlass-dsl` 4.6.0. CUPTI kernel-only, median of 30 iters, cold L2; each cell benches both backends back-to-back in one process, median of 3 passes, single node and session. GQA 32/8 heads, head_dim 128, bf16. cute-dsl = `BatchPrefillWithRaggedKVCacheWrapper(backend="cute-dsl")`, i.e. what the public wrapper runs per mask kind: vendored route for unmasked/causal, modular kernel for windows. trtllm-gen = paged wrapper, HND, page 16, preallocated output. Ratio = cute-dsl / trtllm (< 1: cute-dsl faster). GPU clocks are not lockable on these nodes, so absolute times are only comparable within a table (each table is built from back-to-back same-process measurements; the ratios are clock-robust). Modular vs vendored, why windows route to the modular kernel (B=1 S=16K, median of 3 passes): | config | modular (µs) | vendored (µs) | modular/vendored | |--|--:|--:|:--:| | unmasked | 3540.4 | 2893.8 | 1.22 | | causal | 1808.3 | 1532.7 | 1.18 | | causal + w=127 | 265.0 | 359.8 | 0.74 | | causal + w=511 | 314.4 | 620.9 | 0.51 | | causal + w=1023 | 412.0 | 748.1 | 0.55 | Shape grid (increasing S; per shape: unmasked, causal, windows in increasing size). The kernel column shows which kernel the wrapper routes to: windowed rows run this PR's modular path; unmasked/causal rows run the flashinfer-ai#3857 vendored kernel and are included for context only. | B | S | mask | kernel | cute-dsl (µs) | trtllm (µs) | ratio | |--:|--:|--|--|--:|--:|:--:| | 64 | 512 | unmasked | vendored | 330.9 | 280.8 | 1.18 | | 64 | 512 | causal | vendored | 511.7 | 269.4 | 1.90 | | 64 | 512 | w=256 | modular | 570.9 | 300.6 | 1.90 | | 16 | 1024 | unmasked | vendored | 249.9 | 236.4 | 1.06 | | 16 | 1024 | causal | vendored | 318.9 | 185.7 | 1.72 | | 16 | 1024 | w=256 | modular | 302.7 | 169.2 | 1.79 | | 1 | 4096 | unmasked | vendored | 201.8 | 227.4 | 0.89 | | 1 | 4096 | causal | vendored | 159.3 | 150.0 | 1.06 | | 1 | 4096 | w=256 | modular | 76.4 | 52.5 | 1.46 | | 1 | 4096 | w=1024 | modular | 115.6 | 87.9 | 1.32 | | 1 | 4096 | w=2048 | modular | 151.1 | 121.0 | 1.25 | | 4 | 4096 | unmasked | vendored | 824.4 | 825.0 | 0.99 | | 4 | 4096 | causal | vendored | 557.9 | 483.7 | 1.15 | | 4 | 4096 | w=256 | modular | 288.3 | 172.4 | 1.67 | | 4 | 4096 | w=1024 | modular | 428.5 | 284.5 | 1.51 | | 4 | 4096 | w=2048 | modular | 519.2 | 399.7 | 1.30 | | 1 | 16384 | unmasked | vendored | 2845.0 | 3376.5 | 0.86 | | 1 | 16384 | causal | vendored | 1568.4 | 1758.7 | 0.86 | | 1 | 16384 | w=256 | modular | 292.8 | 178.1 | 1.62 | | 1 | 16384 | w=1024 | modular | 415.0 | 304.8 | 1.36 | | 1 | 16384 | w=2048 | modular | 576.9 | 480.3 | 1.20 | | 1 | 16384 | w=4096 | modular | 881.7 | 808.0 | 1.09 | | 1 | 16384 | w=8192 | modular | 1402.0 | 1314.6 | 1.06 | | 1 | 32768 | unmasked | vendored | 12151.0 | 13847.2 | 0.88 | | 1 | 32768 | causal | vendored | 6353.5 | 7176.0 | 0.89 | | 1 | 32768 | w=256 | modular | 580.1 | 334.8 | 1.77 | | 1 | 32768 | w=1024 | modular | 798.7 | 603.6 | 1.33 | | 1 | 32768 | w=2048 | modular | 1145.3 | 983.6 | 1.16 | | 1 | 32768 | w=4096 | modular | 1955.4 | 1740.8 | 1.11 | | 1 | 32768 | w=8192 | modular | 3443.3 | 3392.4 | 1.01 | Window-size sweep (B=1, S=16384): | window | cute-dsl (µs) | trtllm (µs) | ratio | |--:|--:|--:|:--:| | 127 | 236.2 | 145.6 | 1.62 | | 255 | 269.7 | 167.9 | 1.61 | | 511 | 312.9 | 216.8 | 1.44 | | 1023 | 411.3 | 304.6 | 1.35 | | 2047 | 578.0 | 480.0 | 1.20 | | 4095 | 883.8 | 804.9 | 1.10 | | 8191 | 1387.7 | 1320.7 | 1.04 | Mixed V dtype (bf16 Q/K + fp8 V, B=1 S=8192, same protocol): mixed inputs force the vendored path onto JIT compilation (its prebuilt artifact matrix has no mixed-dtype axis), where the modular kernel also wins causal: | config | modular (µs) | vendored-JIT (µs) | modular/vendored | |--|--:|--:|:--:| | unmasked | 739.4 | 706.3 | 1.05 | | causal | 481.8 | 548.1 | 0.88 | | causal + w=127 | 117.7 | 183.9 | 0.64 | | causal + w=1023 | 208.4 | 395.2 | 0.53 | <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added flexible attention band masking, including causal, symmetric, left-only, and right-bounded sliding windows. * Added optional log-sum-exp (LSE) output, including variable-length, ALiBi, and attention-sink scenarios. * Added support for mixed Q/K/V data types in prefill attention. * Added automatic scale selection when no scale is provided. * **Bug Fixes** * Improved kernel selection and fallback behavior for windowed attention and unsupported data-type combinations. * **Tests** * Expanded coverage for masking, FP8, mixed dtypes, LSE results, and ragged inputs. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…hinfer-ai#3960) ## Summary - Compile SM12x GDN CuteDSL kernels with a device-matched arch (`sm_120a` / `sm_121a`). Hardcoded `sm_120a` left DGX Spark (SM121) without a CuteDSL execution engine. - Align GDN prefill test gates with `is_sm12x_supported` so SM121 is exercised. Validated on NVIDIA GB10 (SM121, CUDA 13.0): representative GDN prefill/decode cases pass after the arch fix; b12x MoE suite was also green in the same sweep (`142 passed`). Relates to flashinfer-ai#3170 (item 5). ## Test plan - [x] `pytest tests/gdn/test_prefill_delta_rule.py -k 'test_prefill_kernel_basic and float16-64-seq_lens0-4-1-1-128'` on GB10 — 24 passed - [x] Small GDN decode smoke (`basic_pretranspose` / `basic_nontranspose` B=1,4) on GB10 — 4 passed - [x] `pytest tests/utils/test_jit_example.py::test_dump_logits --runxfail` still fails as expected on SM121 - [ ] CI / SM120 regression sanity if available Made with [Cursor](https://cursor.com) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added dynamic SM12x GPU architecture targeting for Delta Rule kernel compilation. * Extended SM12x support across prefill and context-parallel Delta Rule operations. * **Bug Fixes** * Improved SM12x PTX validation to prevent caching unsupported generated PTX. * Ensured compilation options are selected based on the actual tensor device for SM12x. * **Tests** * Updated architecture gating in Delta Rule prefill tests to recognize SM12x (and adjust context-parallel skip logic accordingly). <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
flashinfer-ai#4139) ## 📌 Description Follow-up to flashinfer-ai#4075 (merged). Two fixes to the `flashinfer.moe_ep` **NIXL-EP** LL transport that flashinfer-ai#4075 did not include — both are combine-phase deadlocks that only surface under real DP-EP serving load, not in the single-shot multirank roundtrip test. Found while integrating `flashinfer_ep_nixl` into vLLM (vllm-project/vllm#47948). Symptom in both cases: worker rank 0 stops sending its combine data, every peer times out on `NIXL-EP timeout for combine receive ... src_rank 0`, and the EP group hangs. ### 1. Don't wrap NIXL Buffer calls in a user-stream context (`3272f338`) Honoring `HandleAlgoKnobUserStream` by running dispatch/combine/complete inside `torch.cuda.stream(ExternalStream(...))` breaks the NIXL Buffer's own async RDMA completion signaling. The Buffer takes no stream argument and manages its own streams; it must run on the natural current stream. Drop the wrapper — the knob is a no-op for nixl_ep (unlike nccl_ep, which binds the stream explicitly). ### 2. Drive dispatch/combine synchronously — `async_finish=False` + recv hook (`b848e0ae`) The NIXL MVP's `async_finish=True` completion event does not guarantee the RDMA transfer landed, so under load the combine data isn't ready when peers read it. vLLM's native `nixl_ep` backend avoids this with `async_finish=False` + explicit recv hook — the proven path. Match it: `async_finish=False, return_recv_hook=True`, drain the hook synchronously on the non-staged path (defer to `complete()` only when staged/DBO). Fix 1 cleared short sequences; fix 2 was needed for large token counts. ## 🧪 Tests Validated end-to-end on **B200 DP8-EP** (Qwen3-30B-A3B, DP8 + EP) via the `flashinfer_ep_nixl` vLLM backend: - **Before:** 2048-prefill / 2048-decode / GSM8K hang (109–112 `combine receive src_rank 0` timeouts). **After:** all complete. - GSM8K 5-shot **0.848 / 0.891** (flex/strict); throughput (sum of 8 ranks) **9,617 / 23,370 / 5,911** tok/s for 128/128 · 2048/128 · 128/2048 — competitive with NCCL-EP LL. - Mock suite updated for the no-stream-redirect contract; pre-commit clean. The single-shot multirank roundtrip still passes but doesn't exercise this — the deadlocks are load-dependent (~256-prompt sustained load), which is why flashinfer-ai#4075 didn't catch them. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved NIXL dispatch and combine completion handling for more reliable asynchronous communication. * Ensured operations consistently use the caller’s active CUDA stream. * Added safer handling for varying completion responses. * **Tests** * Updated coverage to verify dispatch uses the current CUDA stream rather than a configured side stream. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
<!-- .github/pull_request_template.md --> ## 📌 Description Adds SM107 (compute capability 10.7) support to FlashInfer, so its existing kernels and backends compile and run on SM107 GPUs. Key changes: - SM107a enablement across existing backends — capability lists, JIT architecture flags, and dispatch paths extended to recognize and target sm_107a. - CuTe DSL device-support check with sm_100f family fallback — before selecting a CuTe DSL backend, probe whether the installed DSL can target the device; when it lacks a native SM107 target, fall back to the sm_100f family target so kernels still compile and run correctly. - Auto-detected CUTLASS SM107 support — for CUTLASS-based modules, the sm_107a→sm_100f mapping is now driven by probing the bundled CUTLASS headers, so it upgrades to native sm_107a automatically once CUTLASS gains SM107, with no code change. - Toolchain-aware architecture probes (NVRTC, Triton) — availability checks so unsupported compiler/DSL toolchains skip cleanly instead of failing deep inside compilation. - Arch-aware test handling — tests gate on, or adjust tolerances for, SM107 where hardware-dependent behavior differs. ## 🔍 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** * Expanded SM107 (compute capability 10.7) support across FP4 quantization, GEMM, MoE, and attention workflows. * Added optional `use_fp16_softmax` and `uses_spcompress` controls for paged and ragged attention variant selection. * Added runtime CuTe DSL architecture validation before CuTe-based execution. * **Bug Fixes** * Improved TensorRT-LLM attention kernel selection and shared-memory configuration, including new variant traits. * **Documentation** * Updated source-install instructions to include SM107 in the CUDA architecture list. * **Tests** * Enhanced architecture-aware test skipping and SM107 coverage; relaxed a few numerical tolerances. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Vinnie6167 <Vinnie6167@users.noreply.github.com> Co-authored-by: jimmzhou <jimmzhou@nvidia.com>
…lashinfer-ai#4088) ## 📌 Description FP8 per-tensor routed fused-MoE is supported at the kernel level. This PR exposes it and tests it. ## 🚀 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. - [X] All tests are passing (`unittest`, etc.). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added FP8 per-tensor MoE pre-routed execution that accepts packed expert top-k indices (routing logits can be omitted). - Introduced a new routed FP8 per-tensor MoE API and exposed it at the Python package top level. - Supports non-zero local expert offsets for correct shifted expert dispatch. - **Tests** - Added parity tests comparing logits-based vs pre-routed outputs, including renormalize variants. - Added correctness coverage for non-zero expert offset scenarios. - **Trace** - Added routed MoE trace template, example output, and trace-schema/reference tests. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: jdebache <jdebache@nvidia.com>
<!-- .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 - [ ] 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 * **Compatibility** * Updated GPU architecture support and kernel selection across attention, GEMM, MoE, quantization, and sampling workloads. * Removed dedicated support for select SM107/10.7a builds and refreshed compatible architecture targets. * NVFP4 KV-cache operations now validate scale factors without an additional SM107 restriction. * **Bug Fixes** * Improved TensorRT-LLM kernel compatibility and shared-memory configuration. * Improved optional backend detection and reduced failures in environments with limited runtime support. * **Documentation** * Updated installation examples to reflect the revised CUDA architecture list. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
## What Optimizes the Blackwell SM100 chunk-stream Gated Delta Net (GDN) prefill kernel. Coauthored by @guangyunh-nv @jhjpark ## Test environment - **GPU:** NVIDIA B300 SXM6 AC — Blackwell (SM100) - **Compile:** `--enable-tvm-ffi --opt-level 3` - **Dtype:** bf16 in/out, fp32 state + accumulation - **Benchmark:** `benchmarks/bench_gdn_prefill.py`, 2 warmup + 8 timed iters, CUDA graph - **Baseline ("pre-PR"):** the initial functional port of this kernel (`--opt-level 2`) ## Performance 120 configurations (8 head configs x 15 sequence configs). **Median 1.25x** speedup (range 1.07x-1.32x). No regressions (every config is faster than the baseline). <details> <summary>Full 120-configuration results (h_qk/h_v, seqlen, pre-PR ms, this-PR ms, speedup)</summary> | h_qk/h_v | seqlen | pre-PR (ms) | this PR (ms) | speedup | |---|---|---|---|---| | 2/8 | 1x65536 | 2.352 | 1.827 | 1.29x | | 2/8 | 1x32768 | 1.184 | 0.923 | 1.28x | | 2/8 | 1x16384 | 0.601 | 0.470 | 1.28x | | 2/8 | 1x8192 | 0.311 | 0.242 | 1.28x | | 2/8 | 1x4096 | 0.165 | 0.129 | 1.28x | | 2/8 | 1x2048 | 0.091 | 0.069 | 1.31x | | 2/8 | 6144+2048 | 0.237 | 0.187 | 1.27x | | 2/8 | 4096+4096 | 0.165 | 0.131 | 1.26x | | 2/8 | 2048+6144 | 0.237 | 0.187 | 1.27x | | 2/8 | 1024+7168 | 0.274 | 0.215 | 1.27x | | 2/8 | 2048x4 | 0.093 | 0.074 | 1.25x | | 2/8 | 1024x8 | 0.057 | 0.047 | 1.22x | | 2/8 | 8192x8 | 0.315 | 0.246 | 1.28x | | 2/8 | 8192x16 | 0.317 | 0.248 | 1.28x | | 2/8 | 8192x32 | 0.627 | 0.490 | 1.28x | | 4/16 | 1x65536 | 2.349 | 1.830 | 1.28x | | 4/16 | 1x32768 | 1.185 | 0.925 | 1.28x | | 4/16 | 1x16384 | 0.602 | 0.471 | 1.28x | | 4/16 | 1x8192 | 0.311 | 0.245 | 1.27x | | 4/16 | 1x4096 | 0.165 | 0.131 | 1.26x | | 4/16 | 1x2048 | 0.092 | 0.074 | 1.25x | | 4/16 | 6144+2048 | 0.238 | 0.189 | 1.26x | | 4/16 | 4096+4096 | 0.166 | 0.133 | 1.25x | | 4/16 | 2048+6144 | 0.239 | 0.188 | 1.27x | | 4/16 | 1024+7168 | 0.276 | 0.216 | 1.27x | | 4/16 | 2048x4 | 0.094 | 0.077 | 1.23x | | 4/16 | 1024x8 | 0.060 | 0.050 | 1.21x | | 4/16 | 8192x8 | 0.317 | 0.249 | 1.27x | | 4/16 | 8192x16 | 0.626 | 0.493 | 1.27x | | 4/16 | 8192x32 | 1.235 | 1.031 | 1.20x | | 8/32 | 1x65536 | 2.361 | 1.836 | 1.29x | | 8/32 | 1x32768 | 1.190 | 0.928 | 1.28x | | 8/32 | 1x16384 | 0.605 | 0.475 | 1.27x | | 8/32 | 1x8192 | 0.312 | 0.247 | 1.26x | | 8/32 | 1x4096 | 0.166 | 0.133 | 1.24x | | 8/32 | 1x2048 | 0.092 | 0.075 | 1.24x | | 8/32 | 6144+2048 | 0.240 | 0.191 | 1.26x | | 8/32 | 4096+4096 | 0.167 | 0.134 | 1.24x | | 8/32 | 2048+6144 | 0.241 | 0.190 | 1.27x | | 8/32 | 1024+7168 | 0.278 | 0.218 | 1.28x | | 8/32 | 2048x4 | 0.097 | 0.079 | 1.23x | | 8/32 | 1024x8 | 0.113 | 0.094 | 1.20x | | 8/32 | 8192x8 | 0.628 | 0.508 | 1.24x | | 8/32 | 8192x16 | 1.242 | 1.033 | 1.20x | | 8/32 | 8192x32 | 2.185 | 1.933 | 1.13x | | 16/64 | 1x65536 | 2.383 | 1.835 | 1.30x | | 16/64 | 1x32768 | 1.202 | 0.927 | 1.30x | | 16/64 | 1x16384 | 0.611 | 0.473 | 1.29x | | 16/64 | 1x8192 | 0.316 | 0.247 | 1.28x | | 16/64 | 1x4096 | 0.168 | 0.133 | 1.26x | | 16/64 | 1x2048 | 0.094 | 0.077 | 1.23x | | 16/64 | 6144+2048 | 0.243 | 0.192 | 1.26x | | 16/64 | 4096+4096 | 0.171 | 0.137 | 1.26x | | 16/64 | 2048+6144 | 0.244 | 0.192 | 1.27x | | 16/64 | 1024+7168 | 0.281 | 0.220 | 1.27x | | 16/64 | 2048x4 | 0.186 | 0.152 | 1.23x | | 16/64 | 1024x8 | 0.219 | 0.185 | 1.19x | | 16/64 | 8192x8 | 1.246 | 1.034 | 1.21x | | 16/64 | 8192x16 | 2.189 | 1.930 | 1.13x | | 16/64 | 8192x32 | 4.404 | 4.001 | 1.10x | | 16/32 | 1x65536 | 2.364 | 1.839 | 1.29x | | 16/32 | 1x32768 | 1.192 | 0.930 | 1.28x | | 16/32 | 1x16384 | 0.605 | 0.475 | 1.27x | | 16/32 | 1x8192 | 0.312 | 0.247 | 1.26x | | 16/32 | 1x4096 | 0.166 | 0.134 | 1.24x | | 16/32 | 1x2048 | 0.092 | 0.075 | 1.23x | | 16/32 | 6144+2048 | 0.240 | 0.191 | 1.25x | | 16/32 | 4096+4096 | 0.167 | 0.134 | 1.25x | | 16/32 | 2048+6144 | 0.241 | 0.191 | 1.27x | | 16/32 | 1024+7168 | 0.278 | 0.219 | 1.27x | | 16/32 | 2048x4 | 0.097 | 0.080 | 1.21x | | 16/32 | 1024x8 | 0.112 | 0.095 | 1.18x | | 16/32 | 8192x8 | 0.629 | 0.499 | 1.26x | | 16/32 | 8192x16 | 1.244 | 1.065 | 1.17x | | 16/32 | 8192x32 | 2.190 | 1.990 | 1.10x | | 16/48 | 1x65536 | 2.335 | 1.832 | 1.27x | | 16/48 | 1x32768 | 1.177 | 0.927 | 1.27x | | 16/48 | 1x16384 | 0.599 | 0.473 | 1.27x | | 16/48 | 1x8192 | 0.310 | 0.246 | 1.26x | | 16/48 | 1x4096 | 0.165 | 0.133 | 1.25x | | 16/48 | 1x2048 | 0.092 | 0.075 | 1.23x | | 16/48 | 6144+2048 | 0.239 | 0.191 | 1.25x | | 16/48 | 4096+4096 | 0.168 | 0.134 | 1.25x | | 16/48 | 2048+6144 | 0.240 | 0.190 | 1.26x | | 16/48 | 1024+7168 | 0.277 | 0.219 | 1.27x | | 16/48 | 2048x4 | 0.184 | 0.150 | 1.23x | | 16/48 | 1024x8 | 0.165 | 0.138 | 1.19x | | 16/48 | 8192x8 | 0.930 | 0.780 | 1.19x | | 16/48 | 8192x16 | 1.852 | 1.566 | 1.18x | | 16/48 | 8192x32 | 3.393 | 3.059 | 1.11x | | 16/16 | 1x65536 | 2.377 | 1.836 | 1.29x | | 16/16 | 1x32768 | 1.198 | 0.928 | 1.29x | | 16/16 | 1x16384 | 0.608 | 0.473 | 1.29x | | 16/16 | 1x8192 | 0.314 | 0.246 | 1.28x | | 16/16 | 1x4096 | 0.167 | 0.132 | 1.26x | | 16/16 | 1x2048 | 0.092 | 0.074 | 1.25x | | 16/16 | 6144+2048 | 0.240 | 0.189 | 1.27x | | 16/16 | 4096+4096 | 0.168 | 0.132 | 1.27x | | 16/16 | 2048+6144 | 0.242 | 0.189 | 1.28x | | 16/16 | 1024+7168 | 0.279 | 0.217 | 1.28x | | 16/16 | 2048x4 | 0.095 | 0.076 | 1.24x | | 16/16 | 1024x8 | 0.060 | 0.050 | 1.20x | | 16/16 | 8192x8 | 0.319 | 0.252 | 1.27x | | 16/16 | 8192x16 | 0.632 | 0.507 | 1.25x | | 16/16 | 8192x32 | 1.243 | 1.107 | 1.12x | | 32/32 | 1x65536 | 2.398 | 1.844 | 1.30x | | 32/32 | 1x32768 | 1.209 | 0.932 | 1.30x | | 32/32 | 1x16384 | 0.614 | 0.476 | 1.29x | | 32/32 | 1x8192 | 0.316 | 0.247 | 1.28x | | 32/32 | 1x4096 | 0.168 | 0.133 | 1.26x | | 32/32 | 1x2048 | 0.093 | 0.075 | 1.24x | | 32/32 | 6144+2048 | 0.242 | 0.191 | 1.27x | | 32/32 | 4096+4096 | 0.168 | 0.134 | 1.25x | | 32/32 | 2048+6144 | 0.242 | 0.191 | 1.27x | | 32/32 | 1024+7168 | 0.280 | 0.219 | 1.28x | | 32/32 | 2048x4 | 0.098 | 0.080 | 1.22x | | 32/32 | 1024x8 | 0.113 | 0.096 | 1.18x | | 32/32 | 8192x8 | 0.634 | 0.537 | 1.18x | | 32/32 | 8192x16 | 1.254 | 1.114 | 1.13x | | 32/32 | 8192x32 | 2.238 | 2.097 | 1.07x | ## 🚀 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 * **Bug Fixes** * Improved Blackwell prefill reliability across padded and partial chunks. * Corrected final-state handling for sequences with no tokens. * Improved state loading and storing for Blackwell layouts. * Enhanced numerical handling for inverse calculations and edge cases. * **Performance** * Reworked kernel processing and staging to improve execution efficiency. * Optimized compilation settings for faster runtime performance. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
## Summary - migrate the WY output-only GDN decode kernel from deprecated `cute.experimental` decorators to the unified CuTe DSL decorators - compile SM12x kernels for the device-specific `sm_120a` / `sm_121a` target - include compute capability in the in-process compile cache key The kernel uses no experimental-only APIs; the unified DSL supports its TMA, mbarrier, cluster-launch, and occupancy controls. The experimental path emitted a Core MLIR `sm_121` attribute that `nvidia-cutlass-dsl==4.6.1` rejects before compilation. Closes flashinfer-ai#3995. ## Test plan - [x] H100 (SM90): `pytest tests/gdn/test_decode_delta_rule.py::test_gdn_decode_bf16_wy_output_only_mtp_kernel -q` — 15 passed - [x] GB10 (SM121), `nvidia-cutlass-dsl==4.6.1`: same test — 15 passed - [x] pre-commit hooks: mypy, ruff check, ruff format Related: flashinfer-ai#3960 uses the same device-matched SM12x compile target for GDN prefill. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Updated GPU kernel compilation to use stable compilation interfaces. * Improved compilation behavior for newer GPU architectures. * Enhanced kernel caching to prevent reuse of incompatible compiled binaries. * **Bug Fixes** * Improved reliability when compiling and running the kernel across different GPU variants. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…-ai#4167) ## 📌 Description `layernorm_quant` can miss the FP8 e4m3 correctness bound when the quantization scale is small. Reported for A10/SM86 in the nightly pipeline. The packed `float2` to `__nv_fp8x2_e4m3` cast in `cudaTypeUtils.cuh` rounded through bf16 before quantizing: return __nv_fp8x2_e4m3(bf1622float2(float22bf162(val))); `generalLayerNorm` applies the per-tensor scale before this cast, so the mantissa bits dropped by the bf16 step are scaled up by `1/scale`. At `scale=0.01` that is a 100x amplification, enough to push values into a neighbouring e4m3 bucket. The test reference models the bf16 rounding that the kernel does before scaling, but not this second rounding after it, so the two diverge. This PR converts directly with `__nv_cvt_float2_to_fp8x2`, matching the style already used in `vec_dtypes.cuh`. The e5m2 packed cast already converts without a bf16 round trip, so it is unchanged, which matches the failures being e4m3 only. The cast is only reachable through `generalLayerNorm`, whose sole caller is `LayerNormQuant`, so `layernorm_quant` is the only affected API. This PR also seeds `test_layernorm_quant`. The bf16 round trip is architecture independent, so the error was always present on the packed path; the check bounds the fraction of drifting elements, and unseeded inputs left that fraction varying per run, which is why it surfaced as an intermittent SM86 failure. ## 🔍 Related Issues flashinfer-ai#4160 ## 🚀 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 Verified on an RTX A6000 (SM86, same architecture as the A10 in the nightly pipeline). Before and after, over 30 seeds per shape at `scale=0.01`, e4m3, batch 1: | hidden | seeds over 1% bound (before) | after | worst mismatch (before) | after | |-------:|-----------------------------:|------:|------------------------:|------:| | 500 | 4/30 | 0/30 | 1.20% | 0.00% | | 1024 | 3/30 | 0/30 | 1.27% | 0.00% | The two cases named in the report, `hidden=500` and `hidden=1024` at `scale=0.01`, pass after the fix. - Full `tests/utils/test_norm.py`: 1543 passed, 1344 skipped. Skips are PDL cases that need Hopper or newer. - `pytest tests/trace/ -k "norm"`: 123 passed, 21 skipped. - `pre-commit run` on the changed files: all hooks pass. Unrelated pre-existing failures: `pytest tests/trace/ -k "rope"` reports 24 failures on this machine both with and without this change, so they are not caused by it. - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved FP8 conversion accuracy and handling by converting values directly with saturation. * **Tests** * Stabilized quantized layer normalization test results by making input generation deterministic. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…~4x (flashinfer-ai#4119) <!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> `tests/attention/test_trtllm_gen_attention_decode_xqa.py` takes **40+ minutes** in CI where profiling shows the wall time is ~95% JIT compilation, not kernel execution. The file's parametrize matrix expands to cases that map to 408 unique XQA kernel compilations, all of which occur serially, one at a time. Each compilation takes around 5 seconds. Current PR adds a `tests/attention/conftest.py` that, right before the xqa test file runs, enumerates the exact kernels the collected cases will need and builds them **all in a single parallel `ninja` invocation** via `flashinfer.jit.build_jit_specs`, then stages the artifacts so the tests load them directly instead of recompiling. **Measured on GB10 / SM121, cold cache (full file):** | | Wall time | |---|---| | Before (serial first-touch compile) | ~40 min `4896 passed, 66912 skipped, 2 warnings in 2305.72s (0:38:25)` | | After (parallel prebuild + load-only) | **628 s (10:28)** — `4896 passed, 66912 skipped ... in (0:10:28)` | The parallel build compiles all 408 kernels in ~4.6 min (0.67 s/kernel amortized vs 5.9 s serial); the rest is unchanged test execution. No behavior change — same pass/skip counts. ## 🔍 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 * **Tests** * Improved attention test startup time by bulk-precompiling supported FlashInfer XQA decode kernels during pytest collection. * Reduced repeated compilation overhead by compiling a deduplicated set of module specs in parallel. * Added optional progress/status reporting when terminal output is available. * Improved robustness by staging built artifacts ahead of runtime using safe atomic moves/hardlinks, reusing existing files when possible, and falling back to normal JIT behavior if bulk setup encounters any issues. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…flashinfer-ai#4152) <!-- .github/pull_request_template.md --> ## 📌 Description - Add Tier<896, 16> to the routing and post-TopK policy tables, avoiding fallback to Tier<1024, 32>. Static-block dimensions now follow the selected policy tier. - Add `reduceTopKForLane`, where each lane retains only its assigned TopK rank instead of materializing K outputs per thread. This preserves comparison semantics while reducing register and local-memory pressure. - Parallelize block-per-token TopK For E=512..1024 and tier TopK 9..16, worker warps process 128-expert partitions and warp 0 merges their lane-owned candidates. Other shapes retain the generic fallback. - Use four expanded indices per thread when the high-expert/high-top-K workload fits the bounded cooperative capacity. Larger inputs retain the existing 64-entry fallback. - Derive launch dimensions from the selected tier, use wider clusters for bounded permutation-only launches, and split TopK from permutation at the measured boundary. Reuse histogram initialization already performed by the score kernel. ## Performance on affected shapes ### NoOp + Softmax | E | K | Tokens | Baseline (us) | Best opt (us) | Speedup | Gain | |---:|---:|---:|---:|---:|---:|---:| | 256 | 8 | 1,024 | 9.219 | 8.192 | 1.1253x | +12.53% | | 256 | 8 | 2,048 | 10.247 | 10.238 | 1.0008x | +0.08% | | 256 | 8 | 4,096 | 12.193 | 10.238 | 1.1909x | +19.09% | | 256 | 8 | 8,192 | 14.332 | 12.331 | 1.1623x | +16.23% | | 512 | 8 | 512 | 10.239 | 8.869 | 1.1545x | +15.45% | | 512 | 8 | 1,024 | 12.296 | 12.281 | 1.0012x | +0.12% | | 512 | 8 | 2,048 | 12.287 | 10.492 | 1.1711x | +17.11% | | 512 | 8 | 4,096 | 14.529 | 13.613 | 1.0673x | +6.73% | | 512 | 8 | 8,192 | 18.443 | 16.394 | 1.1250x | +12.50% | | 512 | 10 | 1,024 | 12.308 | 12.104 | 1.0169x | +1.69% | | 512 | 10 | 2,048 | 12.296 | 11.052 | 1.1125x | +11.25% | | 512 | 10 | 4,096 | 14.389 | 14.289 | 1.0070x | +0.70% | | 512 | 10 | 8,192 | 20.490 | 18.442 | 1.1111x | +11.11% | | 896 | 16 | 8 | 10.247 | 10.238 | 1.0009x | +0.09% | | 896 | 16 | 16 | 11.267 | 10.247 | 1.0995x | +9.95% | | 896 | 16 | 32 | 12.181 | 10.238 | 1.1897x | +18.97% | | 896 | 16 | 64 | 12.296 | 10.243 | 1.2004x | +20.04% | | 896 | 16 | 128 | 10.247 | 10.243 | 1.0004x | +0.04% | | 896 | 16 | 256 | 12.290 | 10.248 | 1.1993x | +19.93% | | 896 | 16 | 512 | 13.341 | 10.248 | 1.3018x | +30.18% | | 896 | 16 | 1,024 | 16.574 | 14.332 | 1.1564x | +15.64% | | 896 | 16 | 2,048 | 18.385 | 14.345 | 1.2816x | +28.16% | | 896 | 16 | 4,096 | 24.582 | 20.957 | 1.1729x | +17.29% | | 896 | 16 | 8,192 | 34.805 | 34.025 | 1.0229x | +2.29% | | 1024 | 32 | 8 | 12.285 | 10.719 | 1.1461x | +14.61% | | 1024 | 32 | 16 | 12.286 | 12.285 | 1.0001x | +0.01% | | 1024 | 32 | 32 | 12.297 | 12.296 | 1.0001x | +0.01% | | 1024 | 32 | 64 | 14.332 | 14.332 | 1.0000x | +0.00% | | 1024 | 32 | 128 | 12.286 | 12.285 | 1.0001x | +0.01% | | 1024 | 32 | 256 | 14.312 | 13.640 | 1.0493x | +4.93% | | 1024 | 32 | 512 | 14.334 | 12.297 | 1.1657x | +16.57% | | 1024 | 32 | 1,024 | 20.292 | 17.813 | 1.1392x | +13.92% | | 1024 | 32 | 2,048 | 24.570 | 24.433 | 1.0056x | +0.56% | | 1024 | 32 | 4,096 | 33.614 | 32.765 | 1.0259x | +2.59% | | 1024 | 32 | 8,192 | 53.229 | 52.324 | 1.0173x | +1.73% | | 2048 | 32 | 8 | 24.590 | 24.567 | 1.0009x | +0.09% | | 2048 | 32 | 16 | 26.545 | 24.586 | 1.0797x | +7.97% | | 2048 | 32 | 32 | 28.037 | 26.624 | 1.0531x | +5.31% | | 2048 | 32 | 64 | 30.549 | 28.688 | 1.0649x | +6.49% | | 2048 | 32 | 128 | 22.520 | 20.834 | 1.0810x | +8.10% | | 2048 | 32 | 256 | 20.476 | 20.473 | 1.0001x | +0.01% | | 2048 | 32 | 512 | 35.067 | 33.208 | 1.0560x | +5.60% | | 2048 | 32 | 1,024 | 49.486 | 47.084 | 1.0510x | +5.10% | | 2048 | 32 | 2,048 | 75.816 | 70.466 | 1.0759x | +7.59% | | 2048 | 32 | 4,096 | 126.997 | 116.489 | 1.0902x | +9.02% | | 2048 | 32 | 8,192 | 227.442 | 208.942 | 1.0885x | +8.85% | ### Softmax + SumNormalize | E | K | Tokens | Baseline (us) | Best opt (us) | Speedup | Gain | |---:|---:|---:|---:|---:|---:|---:| | 256 | 8 | 256 | 10.238 | 8.480 | 1.2073x | +20.73% | | 256 | 8 | 512 | 10.238 | 8.197 | 1.2490x | +24.90% | | 256 | 8 | 1,024 | 10.286 | 10.238 | 1.0046x | +0.46% | | 256 | 8 | 2,048 | 12.286 | 10.247 | 1.1989x | +19.89% | | 256 | 8 | 4,096 | 14.332 | 12.285 | 1.1666x | +16.66% | | 256 | 8 | 8,192 | 16.416 | 16.319 | 1.0059x | +0.59% | | 512 | 8 | 512 | 12.290 | 10.247 | 1.1994x | +19.94% | | 512 | 8 | 1,024 | 17.698 | 14.346 | 1.2337x | +23.37% | | 512 | 8 | 2,048 | 14.378 | 14.332 | 1.0032x | +0.32% | | 512 | 8 | 4,096 | 20.473 | 18.426 | 1.1111x | +11.11% | | 512 | 8 | 8,192 | 28.660 | 26.614 | 1.0769x | +7.69% | | 512 | 10 | 512 | 12.287 | 10.239 | 1.2000x | +20.00% | | 512 | 10 | 1,024 | 18.393 | 14.346 | 1.2821x | +28.21% | | 512 | 10 | 2,048 | 15.371 | 14.332 | 1.0725x | +7.25% | | 512 | 10 | 4,096 | 20.587 | 19.934 | 1.0328x | +3.28% | | 512 | 10 | 8,192 | 30.211 | 27.942 | 1.0812x | +8.12% | | 896 | 16 | 8 | 22.541 | 22.520 | 1.0010x | +0.10% | | 896 | 16 | 16 | 24.590 | 23.437 | 1.0492x | +4.92% | | 896 | 16 | 32 | 24.765 | 22.541 | 1.0987x | +9.87% | | 896 | 16 | 64 | 26.056 | 22.541 | 1.1559x | +15.59% | | 896 | 16 | 128 | 24.565 | 22.541 | 1.0898x | +8.98% | | 896 | 16 | 256 | 38.915 | 36.857 | 1.0558x | +5.58% | | 896 | 16 | 512 | 61.465 | 53.276 | 1.1537x | +15.37% | | 896 | 16 | 1,024 | 102.407 | 92.649 | 1.1053x | +10.53% | | 896 | 16 | 2,048 | 53.234 | 47.081 | 1.1307x | +13.07% | | 896 | 16 | 4,096 | 92.188 | 83.926 | 1.0984x | +9.84% | | 896 | 16 | 8,192 | 153.474 | 140.932 | 1.0890x | +8.90% | | 1024 | 32 | 8 | 26.590 | 24.567 | 1.0823x | +8.23% | | 1024 | 32 | 16 | 28.661 | 26.634 | 1.0761x | +7.61% | | 1024 | 32 | 32 | 28.662 | 26.639 | 1.0759x | +7.59% | | 1024 | 32 | 64 | 28.675 | 26.640 | 1.0764x | +7.64% | | 1024 | 32 | 128 | 26.614 | 24.592 | 1.0822x | +8.22% | | 1024 | 32 | 256 | 43.438 | 42.989 | 1.0105x | +1.05% | | 1024 | 32 | 512 | 71.633 | 67.604 | 1.0596x | +5.96% | | 1024 | 32 | 1,024 | 120.501 | 114.753 | 1.0501x | +5.01% | | 1024 | 32 | 2,048 | 219.059 | 211.049 | 1.0380x | +3.80% | | 1024 | 32 | 4,096 | 412.856 | 399.442 | 1.0336x | +3.36% | | 1024 | 32 | 8,192 | 807.110 | 781.711 | 1.0325x | +3.25% | | 2048 | 32 | 8 | 34.807 | 34.519 | 1.0083x | +0.83% | | 2048 | 32 | 16 | 37.639 | 36.849 | 1.0214x | +2.14% | | 2048 | 32 | 32 | 38.897 | 36.902 | 1.0541x | +5.41% | | 2048 | 32 | 64 | 40.958 | 39.007 | 1.0500x | +5.00% | | 2048 | 32 | 128 | 32.787 | 32.755 | 1.0010x | +0.10% | | 2048 | 32 | 256 | 49.131 | 47.129 | 1.0425x | +4.25% | | 2048 | 32 | 512 | 84.027 | 81.859 | 1.0265x | +2.65% | | 2048 | 32 | 1,024 | 135.455 | 131.144 | 1.0329x | +3.29% | | 2048 | 32 | 2,048 | 241.258 | 233.501 | 1.0332x | +3.32% | | 2048 | 32 | 4,096 | 448.555 | 434.204 | 1.0330x | +3.30% | | 2048 | 32 | 8,192 | 863.408 | 837.840 | 1.0305x | +3.05% | ### SigmoidBias + ScaledSumNormalize | E | K | Tokens | Baseline (us) | Best opt (us) | Speedup | Gain | |---:|---:|---:|---:|---:|---:|---:| | 256 | 8 | 1,024 | 10.238 | 8.220 | 1.2456x | +24.56% | | 256 | 8 | 2,048 | 12.282 | 10.238 | 1.1996x | +19.96% | | 256 | 8 | 4,096 | 14.332 | 11.788 | 1.2159x | +21.59% | | 256 | 8 | 8,192 | 16.427 | 14.347 | 1.1450x | +14.50% | | 512 | 8 | 1,024 | 14.333 | 12.297 | 1.1656x | +16.56% | | 512 | 8 | 2,048 | 16.379 | 14.346 | 1.1418x | +14.18% | | 512 | 8 | 4,096 | 20.491 | 20.473 | 1.0009x | +0.09% | | 512 | 8 | 8,192 | 28.688 | 28.662 | 1.0009x | +0.09% | | 896 | 16 | 8 | 14.333 | 10.245 | 1.3990x | +39.90% | | 896 | 16 | 16 | 14.369 | 10.247 | 1.4023x | +40.23% | | 896 | 16 | 32 | 14.530 | 10.238 | 1.4192x | +41.92% | | 896 | 16 | 64 | 15.812 | 10.238 | 1.5444x | +54.44% | | 896 | 16 | 128 | 14.333 | 10.247 | 1.3987x | +39.87% | | 896 | 16 | 256 | 16.317 | 12.285 | 1.3282x | +32.82% | | 896 | 16 | 512 | 20.473 | 12.296 | 1.6651x | +66.51% | | 896 | 16 | 1,024 | 24.633 | 14.345 | 1.7172x | +71.72% | | 896 | 16 | 2,048 | 28.688 | 18.439 | 1.5558x | +55.58% | | 896 | 16 | 4,096 | 45.069 | 28.104 | 1.6036x | +60.36% | | 896 | 16 | 8,192 | 70.405 | 42.111 | 1.6719x | +67.19% | | 1024 | 32 | 8 | 15.235 | 14.345 | 1.0620x | +6.20% | | 1024 | 32 | 16 | 16.379 | 16.107 | 1.0169x | +1.69% | | 1024 | 32 | 32 | 16.394 | 16.379 | 1.0009x | +0.09% | | 1024 | 32 | 64 | 17.410 | 16.394 | 1.0620x | +6.20% | | 1024 | 32 | 128 | 16.380 | 14.346 | 1.1418x | +14.18% | | 1024 | 32 | 256 | 16.402 | 16.394 | 1.0005x | +0.05% | | 1024 | 32 | 512 | 22.536 | 20.492 | 1.0997x | +9.97% | | 1024 | 32 | 1,024 | 28.810 | 28.376 | 1.0153x | +1.53% | | 1024 | 32 | 2,048 | 43.058 | 40.987 | 1.0505x | +5.05% | | 1024 | 32 | 4,096 | 73.783 | 71.667 | 1.0295x | +2.95% | | 1024 | 32 | 8,192 | 134.244 | 130.549 | 1.0283x | +2.83% | ## 🔍 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 * **New Features** * Enhanced Mixture-of-Experts routing with opt-in lane-owned TopK support for eligible postprocess policies. * Extended expert-tier coverage up to **896 experts**, including additional routing dispatch branches for specific configurations. * **Performance** * Improved TopK selection and permutation/coop launch behavior for bounded high-expert/high-TopK ranges, including SM90+ heuristics and tighter cooperative state limits. * Enhanced robustness for non-divisible expert tiers via padded per-thread expert coverage. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
…e decode (flashinfer-ai#4137) ## 📌 Description XQA is the FlashInfer decode kernel used on SM120/121 for models with attention sinks (flashinfer-ai#4070). Serving those models with speculative decoding runs the draft-verification step through XQA as well, and two gaps blocked that: - Every request in a batch had to verify the same number of draft tokens. The kernel indexes queries and masks through cumulative lengths internally, but the argument was never exposed. - Sliding-window masking computed one window start from the last draft token's position and applied it to the whole draft block. Each draft token sits at its own position, so this masked out KV that the earlier draft tokens should still attend to. Fixes: - Plumb `q_cu_seq_lens` through the wrapper and Python API so each request can verify a different number of draft tokens, with host-side input validation. The SM90 fp8 path rejects ragged Q rather than run an unvalidated path. - Compute the window per draft row: whole KV tiles are skipped conservatively, and the exact per-row edge is masked in the kernel. - Extend tests and the benchmark to cover both draft-block mask modes (causal and full, selected with `--spec_dec_mask`), long contexts (split-KV path), zero-length drafts, and GQA group ratio 16. Kernel changes and investigation by @bkryu. ## 📈 Performance The table shows XQA's speedup over each baseline, measured as kernel time with CUPTI at batch size 1 with head_dim 128, 32 query heads over 2 KV heads, and a bf16 KV cache. Each cell gives the speedup at context lengths 1k, 4k, and 32k. The column m is the number of query tokens per request: m=1 is plain decode, and m=4 or 8 is draft-block verification in speculative decode. Each baseline is measured under full attention and under sliding-window attention with a 1024-token window (SWA 1024). Baselines: - vLLM's Triton unified-attention kernel, sinks enabled on both sides. - FlashInfer fa2 wrappers, sinks disabled on both sides (fa2 has no sink support). | GPU | m | vs Triton, full attn | vs Triton, SWA 1024 | vs fa2, full attn | vs fa2, SWA 1024 | |---|---|---|---|---|---| | GB10 | 1 | 1.7 / 0.9 / 1.1 | 1.5 / 1.2 / 2.3 | 0.8 / 0.9 / 1.0 | 0.9 / 0.7 / 0.8 | | GB10 | 4 | 1.2 / 2.2 / 5.0 | 1.2 / 1.2 / 1.2 | 2.7 / 1.4 / 1.2 | 2.3 / 1.9 / 1.9 | | GB10 | 8 | 1.1 / 1.8 / 4.0 | 1.2 / 1.1 / 0.9 | 1.8 / 1.4 / 1.0 | 2.5 / 2.1 / 1.5 | | RTX PRO 6000 | 1 | 3.0 / 2.4 / 3.2 | 3.3 / 3.0 / 3.6 | 1.0 / 0.8 / 0.8 | 1.0 / 1.0 / 0.6 | | RTX PRO 6000 | 4 | 2.1 / 6.0 / 18.0 | 1.9 / 1.9 / 1.2 | 3.7 / 2.3 / 1.5 | 2.7 / 3.0 / 2.1 | | RTX PRO 6000 | 8 | 2.1 / 5.5 / 14.2 | 1.8 / 2.0 / 1.0 | 3.3 / 2.4 / 1.1 | 2.8 / 2.8 / 1.6 | ## 🔍 Related Issues flashinfer-ai#4070 (attention sinks on SM120/121). ## 🚀 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. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). `tests/attention/test_xqa_batch_decode.py` passes on SM120 (RTX 5080, RTX PRO 6000) and SM121 (GB10). Outputs cross-checked against an independent reference kernel up to 64k context. ## Reviewer Notes - Speculative-decode builds now assume the draft tokens form a linear chain rather than a tree (the `IS_SPEC_DEC_TREE` compile flag flips from 1 to 0). This is deliberate: the per-row window needs each draft token's sequence position, and only a linear chain defines one. The assumption only matters in sliding-window builds, and tree-shaped drafts never worked correctly with sliding windows, so no working caller changes behavior. - Not validated on SM90/SM100 hardware. The changed paths are spec-dec only, and the SM90 fp8 wrapper rejects ragged explicitly; CI runs the tests on those arches. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added ragged-Q speculative decoding support via optional `q_cu_seq_lens`. - Introduced configurable speculative draft attention masking with `--spec_dec_mask` (`causal`/`full`), including trace support. - Extended JIT/custom-op pathways to support ragged-Q specialization. - **Bug Fixes** - Improved speculative decoding masking for sliding-window cases and fixed an empty-query edge case for ragged requests. - **Tests** - Expanded coverage for ragged-Q, `causal`/`full` mask modes, sliding-window behavior, and KV cache variants. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: bryu <bryu@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…#4187) ## 📌 Description Fixes flashinfer-ai#4166. `tests/utils/test_gen_module_symlink_race_condition.py` fails on `release-v0.6.16` rc1 (`unit_test_b300`, cu129/cu130): ```text AssertionError: Expected /tmp/flashinfer_test_fused_moe_symlink_ng4ln5vs/flashinfer/trtllm/ batched_gemm/trtllmGen_bmm_export to be a symlink ``` ### Root cause flashinfer-ai#3468 ("move JIT symlinks to writable gen dir", `b9890c85`) moved the `trtllmGen_bmm_export` symlink from `FLASHINFER_CUBIN_DIR` to `FLASHINFER_GEN_SRC_DIR` in `flashinfer/jit/fused_moe.py`, `flashinfer/jit/moe_utils.py`, and `flashinfer/jit/gemm/core.py`. That PR touched three production files and no tests. The test redirects only `FLASHINFER_CUBIN_DIR`, then asserts the symlink under it, so it now checks a location the code no longer writes to. There is a second, quieter consequence. Since `FLASHINFER_GEN_SRC_DIR` is never redirected, the symlink is created in the *real* workspace directory and points into the test's temp directory, which `shutil.rmtree` deletes on teardown. The between-iteration `symlink_path.unlink()` also targeted the temp path, so it never removed the symlink actually in use — meaning the race stopped being re-triggered after the first iteration, and the test would have silently under-tested even with the assertion path corrected. ### Fix - Redirect `FLASHINFER_GEN_SRC_DIR` alongside `FLASHINFER_CUBIN_DIR` in the worker, and assert the symlink there. `tests/jit/test_jit_cpp_ext.py` already redirects `FLASHINFER_GEN_SRC_DIR` this way, and `flashinfer/aot.py` reassigns the same attribute, so this is the established pattern. - Keep the artifact cache and generated-source dir as **separate** temp subdirectories, so the assertion still fails if the symlink ever moves back under `FLASHINFER_CUBIN_DIR`. Pointing both at one directory would make the assertion pass either way. - Restore the race semantics and confine all symlinks to the temp tree. - Correct the `ensure_symlink()` docstring, which still cited the pre-flashinfer-ai#3468 `CUBIN_DIR` location. No production behavior changes; the only non-test edit is a docstring. ## 🔍 Related Issues - Fixes flashinfer-ai#4166 - Root cause: flashinfer-ai#3468 ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used my 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. All hooks pass on the changed files (`ruff check`, `ruff format`, `mypy`, whitespace/EOL). ## 🧪 Tests - [x] Tests have been updated as needed. The test requires SM100/SM12x and returns early elsewhere, so it cannot be executed end-to-end on the SM89/90 host available to me. **It still needs a B300/GB200 run to confirm the CI failure is cleared.** To verify the path logic without that hardware, I exercised the real `gen_trtllm_gen_fused_moe_sm100_module()` with only artifact download/checksum verification stubbed (the symlink is created before the nvcc-arch check raises), and recorded the `ensure_symlink()` link path. A stand-in directory represents the shared workspace dir so nothing touches the real cache. Redirecting `FLASHINFER_CUBIN_DIR` only, as the test did before: ```text [ensure_symlink] link=/tmp/verify_before_.../shared_workspace_generated/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export symlink present in cubin_dir (old assertion target) : False <-- assertion fails, matches flashinfer-ai#4166 symlink present in shared gen dir (real workspace) : True <-- leaks outside the temp tree symlink present in test gen dir (temp, isolated) : False ``` Redirecting both, as this PR does: ```text [ensure_symlink] link=/tmp/verify_after_.../generated/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export symlink present in cubin_dir (old assertion target) : False symlink present in shared gen dir (real workspace) : False <-- no leak symlink present in test gen dir (temp, isolated) : True <-- assertion passes ``` Collection and the non-SM100 early-return path were also confirmed green: `pytest -q tests/utils/test_gen_module_symlink_race_condition.py` → `1 passed`. ## Reviewer Notes `gen_fused_moe_worker_process` now takes a `(cubin_dir, gen_src_dir)` tuple, since `Pool.map` passes a single argument. `flashinfer/jit/moe_utils.py` and `flashinfer/jit/gemm/core.py` have the same symlink layout but no test asserts their paths, so nothing else needed updating. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Corrected an example path in the symlink helper documentation. * **Tests** * Improved race-condition coverage for generated module symlinks. * Reduced test flakiness by separating generated files from cached build artifacts. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
## 📌 Description This PR adds the head-scatter / sequence-gather all-to-all used by **Ulysses sequence parallelism** (video DiT / long-sequence attention workloads), consisting of a fused NVLink-P2P kernel and a topology-aware public API with automatic NCCL fallback. **Fused-transpose NVLink-P2P kernel** (`include/flashinfer/comm/ulysses_all_to_all.cuh`, `csrc/ulysses_all_to_all.cu`): the Ulysses layout permutation is folded directly into the cross-GPU writes over CUDA IPC, replacing the `permute → nccl all_to_all_single → permute` sequence with a single coalesced push kernel. Adapted from [ThunderKittens' NVLink all-to-all](https://github.com/HazyResearch/ThunderKittens/blob/main/kernels/parallel/all_to_all/all_to_all.cu); reuses the vLLM-derived `Signal`/`multi_gpu_barrier` machinery already in FlashInfer. Supports world sizes 2/4/6/8, fp16/bf16/fp32. **Public API** (`flashinfer/comm/ulysses.py`): `UlyssesCommunicator(group, max_elems=..., dtype=..., backend="auto"|"nvlink"|"nccl")` with `scatter_heads()` / `gather_heads()` deriving geometry from tensor shapes: ```python with UlyssesCommunicator(group, max_elems=B*S_local*H*D, dtype=torch.bfloat16) as comm: q_ = comm.scatter_heads(q) # [B,S_local,H,D] -> [B,S_global,H_local,D] o = comm.gather_heads(o_) # inverse print(comm.backend, comm.fallback_reason) # e.g. "nvlink", None ``` **Topology-aware backend selection** (`flashinfer/comm/ulysses_topology.py`): the fused kernel is only enabled after a group-wide handshake verifies single host, unique physical GPU identity (NVML UUID), **pair-wise** P2P *and* NVLink for every ordered GPU pair, and a supported world size — all **before any IPC allocation or JIT compilation**. Anything unknown or unverifiable (multi-node, PCIe-only pairs, probe errors, old torch without device UUIDs, single visible GPU per process) conservatively falls back to NCCL with the reason exposed. Constructor and teardown are collective-failure safe: rank-local errors are exchanged as group outcomes so all ranks jointly clean up and raise or fall back — no rank can strand its peers; `close()` is staged (peer mappings closed group-wide before any export is freed), drain-retried, and retryable. **Performance** — 8×H20 (full NVLink NV18), bf16, Wan2.1-14B geometry (B=1, S_global=32760, 40 heads × 128; divisible standalone workload for W=6). One sample = the attention pattern 3×scatter + 1×gather; p50 of 150 rank-max samples; reproducible via `benchmarks/bench_ulysses_a2a.py` with a fail-closed ≤3% regression gate (full methodology in `benchmarks/ulysses_a2a_m3_report.md`): | W | fused (public API) | NCCL reference | speedup | |---|---|---|---| | 2 | 2.093 ms | 2.631 ms | 1.26× | | 4 | 1.083 ms | 1.709 ms | 1.58× | | 6 | 1.024 ms | 1.472 ms | 1.44× | | 8 | 0.682 ms | 0.997 ms | 1.46× | The public API adds no measurable per-call overhead vs an inline NCCL implementation (−1.2%..+0.1%). End-to-end: Wan2.1-14B self-attention forward 27.43 → 27.10 ms with bit-identical outputs; measured earlier in SGLang serving (LingBot-World, 8×H200, ulysses-degree 8): −4.7% DiT forward median vs NCCL, bit-identical frames. **Also included**: `docs/api/comm.rst` section (layouts, backend policy / fallback-reason table, lifecycle & stream constraints, known limitations); `examples/pytorch/wan` integration (Wan2.1/2.2 attention, microbenchmark, full-video generation with pre-spawn config preflight) consuming the public API. **Known limitations** (documented): old PyTorch without `device_properties.uuid` or one-visible-GPU-per-process setups fall back to NCCL; coordinated shutdown after a mid-run distributed exception follows standard distributed failure-abort; a per-world-size block-size dispatch (1024 threads helps W=2 by ~30% but regresses W=8) is recorded as future work in the M3 report. ## 🔍 Related Issues Please merge this PR after flashinfer-ai#3819. ## 🚀 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 - ~180 tests across six suites (verified on 8×H20; the topology decision layer, benchmark-methodology and example-preflight suites are pure CPU): kernel correctness at W=2/4/6/8 × 3 dtypes × aligned/unaligned shapes against independent all-gather references for **both** directions; 2-rank collective-safety fault injection (invalid/inconsistent configs, probe/decision/IPC/JIT/init faults at every stage with resource-ledger balance checks; all ranks reach the same outcome within a deadline and exit naturally); retryable-close scenarios including the resource-less-rank retry deadlock; device binding contract (`cuda:rank`, bare `cuda`, current-device switches); fallback paths proven to never touch IPC/JIT via booby-trapped entry points. - Review focus suggestions: the collective-safe outcome protocol in `resolve_ulysses_backend` and the staged init/teardown transaction in `UlyssesCommunicator` (`flashinfer/comm/ulysses.py`) — these are the parts designed so that no single-rank failure can deadlock peers. - Sphinx `-W` builds the new Ulysses pages with zero warnings; performance artifacts (per-sample JSON/CSV with provenance) were audited during development and the committed report contains the reproduction commands. 🤖 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** * Added Ulysses context-parallel all-to-all for `[B, S, H, D]` tensors with head-scatter/gather semantics. * Added an NVLink-accelerated path with automatic backend selection, topology probing, and safe NCCL fallback. * Added optional Ulysses integration to the WAN transformer example. * Exposed a UlyssesCommunicator API plus advanced raw init/dispose and kernel entry points. * **Bug Fixes** * Improved CUDA IPC teardown using explicit IPC handle closing and stronger synchronization. * **Documentation** * Added detailed Ulysses API docs, constraints, and usage guidance. * **Tests** * Added correctness, lifecycle, validation, topology-decision, and fault-injection test coverage. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: forrestl <forrestl@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…, fused quant+stage launch, persistent knob cache, and prequantized weight packs (flashinfer-ai#4079) <!-- .github/pull_request_template.md --> ## 📌 Description ### Summary Follow-up to flashinfer-ai#3980. This PR makes the MegaMoE (`moe_ep`) path consumable by serving engines (vLLM-style integration): full CUDA-graph capture/replay support, a fused single-launch quant+stage hot path, engine-friendly weight lifecycle (prequantized packs, post-preprocess source release), pooled symmetric-buffer workspaces shared across layers, and a persistent knob cache so production sessions resolve tuned knobs with a pure lookup. No in-engine autotuning. ### Performance Highlights All numbers measured on a **single node, 4x GB200 (EP=4)**. The built-in knob heuristic currently supports **GB200 only** (offline-tuned knob caches expected to work on any supported device). **Microbenchmark** — DeepSeek-V3 geometry (256 experts, top-8, hidden 7168, intermediate 2048), `e2e_pipelined` p50 µs, heuristic knobs, best variant (`nvfp4 + combine_nvfp4`) vs `deep_gemm_mega` baseline: | tok/rank | deep_gemm_mega | nvfp4 + combine_nvfp4 | Speedup | |---|---:|---:|---:| | 512 | 340.0 | 321.1 | 1.06x | | 1024 | 468.0 | 363.5 | 1.29x | | 2048 | 817.2 | 529.4 | 1.54x | | 4096 | 1473.5 | 862.7 | 1.71x | | 8192 | 2993.7 | 1677.3 | 1.78x | The sweep also covers real-model MoE geometries (DeepSeek V3 / V4-Flash / V4-Pro, Kimi K2.6, Qwen3.5-397B, gpt-oss-120b — the last enabled by the %64 alignment relaxation); the pattern holds everywhere: deep_gemm-parity at small token counts, 1.6–1.9x for fp4 combine-wire at large tokens on 7168-hidden shapes. Full sweep tables (all variants), accuracy, and benchmark methodology: [`kernel_src/cutedsl_megamoe/TUNING.md`](https://github.com/mhoqueanik/flashinfer-moe_ep/blob/fi-moe_ep-framework-integration/flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md#real-model-geometry-sweep-2026-07-21). **End-to-end (vLLM)** — early e2e benchmarks with this backend integrated into vLLM 0.25.1 (DeepSeek-V4-Flash, 4x GB200 TP4/EP4, CUDA graphs capturing all recurring step shapes including prefill chunks, per-role offline knob caches) show **~18% prefill and ~7% decode throughput gain over native vLLM** (i.e.: vLLMs built-in deep_gemm mega path): | Workload | native vLLM | fi MegaMoE (nvfp4) | Speedup | |---|---:|---:|---:| | Prefill, 8k-token chunks (tok/s) | 45,701 | 53,962 | 1.18x | | Decode, 1024-seq concurrency (output tok/s) | 21,049 | 22,614 | 1.07x | | GSM8K (200q, greedy) | 0.965 | 0.975 | — | Numbers reproduce within 1% at this branch tip; details and caveats (the decode number requires graph capture to cover prefill chunk shapes) are in TUNING.md's e2e section. The vLLM integration itself will land as a separate PR. **Accuracy** — per-variant microbench rel-L2 (`acc_loss_pct` vs fp32 dense-MoE reference) is unchanged from flashinfer-ai#3980; see the accuracy table in TUNING.md. On real-model distributions the nvfp4 path is at parity-or-better (GSM8K above), consuming the single-quant nvfp4 checkpoint directly via the prequantized weight-pack path. ### Key Changes **CUDA Graph Support (cutedsl mega path)** - Warmup contract + capture guards for graph-safe forward passes - Tail-mask memo made graph-capture-aware; records the actual staged count (fixes a view-slicing regression on capture-touched buffers) - 2-rank lockstep CUDA-graph replay test for the nvfp4 mega layer **Launch-Path Optimizations** - Fused single-launch quant+stage for the cutedsl mega path; fused staging for `deep_gemm` mega with a 16B-alignment fallback (bit-identical to the multi-kernel path, toggle: `FLASHINFER_MEGA_FUSED_STAGE`) - Zero-copy output + memoized tail mask + caller-owned ikr buffers to cut per-call launch count - Cached stream-aware launch thunk in the nvfp4 backend hot path **Engine Integration: Weights & Workspaces** - `MoEWeightPack` split into `Unquantized | Prequantized` variants - Source weight pack released after preprocess (fixes OOM at model load) - Pooled symmetric-buffer workspaces shared across mega layers **Tuning Infrastructure** - Persistent knob cache (`FLASHINFER_MOE_EP_KNOB_CACHE`): offline-tuned winners resolved by pure lookup before the built-in heuristic; populated via the new offline tuner CLI `python -m flashinfer.moe_ep.tune`. Per-role cache files (prefill-tuned / decode-tuned) are the validated deployment pattern. - Skew-aware tuning (`--skew`), schedule-axis sweep (`--sweep schedule`), `--live-tokens` **Shape Coverage & Guardrails** - cutedsl mega alignment relaxed to %64 (dg keeps %128) — enables gpt-oss-120b-class geometries - Warn-once import check for the CuTe-DSL 4.6.1 performance floor **Determinism Investigation (docs)** - Reported `fi_dg` cross-run nondeterminism traced to vLLM batch-formation timing (schedule diff, 1-tok vs 8-tok batch) — engine-level and backend-independent; closed as a FlashInfer issue with probe tests retained ### Directories Affected - `flashinfer/moe_ep/backends/mega/` — CUDA graph support, fused staging, launch caching - `flashinfer/moe_ep/kernel_src/cutedsl_megamoe/` — shim knob cache, TUNING.md - `flashinfer/moe_ep/core/`, `flashinfer/moe_ep/modes/` — workspace pooling, weight-pack lifecycle - `flashinfer/moe_ep/tune.py` — offline tuner CLI - `tests/moe_ep/` — CUDA graph (single- and multi-rank), fused stage, knob cache, weight-pack, workspace-pool, determinism-probe coverage ### Next Steps (WIP) - Generalized tuning: extend the knob heuristic / offline-tuned profiles beyond GB200 - Performance support for CuTe-DSL 4.5.2 (currently 34–54% slower than 4.6.1, which is treated as the perf floor with a warn-once import check) ## 🔍 Related Issues Follow-up to flashinfer-ai#3980 (MegaMoE kernel restructure). ### ✅ 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.). New tests: `test_mega_cuda_graph.py`, `test_mega_cuda_graph_multirank.py`, `test_fused_quant_stage.py`, `test_knob_cache.py`, `test_weight_pack_union.py`, `test_workspace_pool.py`, `test_moe_ep_deep_gemm_skew_determinism.py`. Multi-rank tests run under `tests/moe_ep/run_tests.sh`. ## Reviewer Notes - The fused quant+stage path is bit-identical to the original multi-kernel path (`tests/moe_ep/test_fused_quant_stage.py`); `FLASHINFER_MEGA_FUSED_STAGE=0` reverts it for bisection. - The knob cache keeps autotuning strictly offline: `knobs=None` sessions do a pure cache lookup, falling back to the heuristic on miss — no measurement runs inside an engine process. <!-- 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 fused input staging for supported MoE quantization workflows, with an environment-variable toggle. * Added persistent autotuning cache support and an offline tuning command. * Added workspace sharing to reduce memory usage across compatible layers. * Added optional zero-copy output handling and CUDA Graph warmup support. * Added explicit pre-quantized and unquantized weight types. * **Bug Fixes** * Improved CUDA Graph capture safety, token tracking, alignment support, and workspace cleanup. * **Documentation** * Updated configuration and tuning guidance, including cache settings and performance information. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
## Summary - document FLASHINFER_CUTE_PREFILL_PERSISTENT in the CLAUDE.md environment-variable quick reference - clarify that 0 selects non-persistent scheduling and 1 selects persistent scheduling ## Testing - pre-commit run --files CLAUDE.md - git diff --check <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Documented the `FLASHINFER_CUTE_PREFILL_PERSISTENT` environment variable for controlling CuTe-DSL prefill scheduling. * Added explanations for the supported values: `0` for non-persistent scheduling and `1` for persistent scheduling. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…lashinfer-ai#4175) ## Summary - add the missing `trtllm_fp8_per_tensor_scale_routed_moe` entry to `docs/api/fused_moe.rst` - fix the doc coverage regression reported by the doc checker ## Testing - not run (docs autosummary list change only) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Added the routed FP8 per-tensor scale API to the TensorRT-LLM Fused MoE reference documentation. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
) ## Description Document the two `checkpointing_ssu` arguments still reported by the documentation checker in flashinfer-ai#4061: - `cu_seqlens`: packed variable-length boundaries, dtype/device requirements, and packed `x` shape - `max_seqlen`: required per-sequence upper bound, JIT specialization/ring-window role, and fixed-length exclusion The wording follows the wrapper validation and ring-buffer contract in the implementation. This is documentation-only and does not change runtime behavior. ## Validation - `git diff --check` - `pre-commit run --files flashinfer/mamba/checkpointing_ssu.py` - mypy passed - Ruff check/format passed - all remaining configured hooks passed Addresses the remaining `checkpointing_ssu` argument-consistency finding in flashinfer-ai#4061. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Clarified variable-length mode inputs for checkpointing. * Specified that packed sequence boundaries (`cu_seqlens`) use a 1D `int32` CUDA tensor of shape `(batch + 1,)`. * Updated required input shapes when sequence boundaries are provided, and clarified `max_seqlen` requirements (must be provided with boundaries, omitted otherwise). <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: hebo1221 <hebo1221@users.noreply.github.com>
<!-- .github/pull_request_template.md --> Fix ci failure issue introduced by incorrect testing tensor layout: tests/gemm/test_unified_gemm_fuzz.py Fix ci failure issue introduced by testing on unsupported cudnn version: tests/grouped_mm/test_grouped_mm_bf16.py tests/grouped_mm/test_grouped_mm_fp8.py ## 🚀 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 - **Tests** - Improved NVFP4 quantization validation by aligning test inputs with the exact reference dequantization grid. - Updated grouped matrix multiplication tests to consistently detect supported cuDNN versions and capabilities. - Improved test skipping behavior for environments that lack required mixture-of-experts functionality. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yanqin Zhai <yanqinz@nvidia.com>
…ointer packing (flashinfer-ai#3903) <!-- .github/pull_request_template.md --> ## 📌 Description Resolves item 14 of the SM121 support audit (flashinfer-ai#3170). The fused TRT-LLM allreduce module (`trtllm_comm`) refuses to build on SM12x because its JIT gate only lists SM90 and SM100, so those GPUs get no fused allreduce at all. The gate came from the original TRT-LLM port, not from a hardware limit. SM12x has everything these kernels use (verified on RTX 5080 and DGX Spark: thread-block clusters, distributed shared memory, programmatic dependent launch). Two real bugs sat behind the gate: - **Empty kernel on SM12x.** Separate from the JIT gate, the kernel header excludes SM12x with a compile-time check. Widening the JIT gate alone would build the lamport kernel (the low-latency path for small batches) with an empty body: small-batch decode calls would launch it, do nothing, and return garbage with no error. - **Wrong pointer layout in the host wrapper (all architectures).** The wrapper writes the lamport buffer pointers into a fixed-size array with one spacing between buffer sets, and the kernels read them back with another. The two layouts only agree at world size 16. At any other world size the kernel reads slots that were never written, then crashes or corrupts. No test ever caught this: the lamport kernel only runs at 16 tokens or fewer, and the existing tests use 64 and 128. Fixes: - The wrapper now writes the pointer array with the same spacing the kernels read. - The compile-time exclusion is removed, so the lamport kernel gets a real body on SM12x. - The JIT gate now includes major version 12, and AOT builds include the module for every arch it supports. - Launch cluster sizes are clamped to what the device reports through `cudaOccupancyMaxPotentialClusterSize`, instead of assuming a size per arch. SM90, SM100, SM120, and SM121 all report 8, so behavior is unchanged where the kernels already ran. - The test suites skip cleanly on unsupported machines instead of erroring. SM90 and SM100 compile to the same machine code as before; the compile-time check change does not affect them. ## 🔍 Related Issues flashinfer-ai#3170 (item 14, cross-cutting issue 5). ## 🧪 Tests On DGX Spark (GB10, SM121, CUDA 13, aarch64) and RTX 5080 (SM120, CUDA 13.3, x86_64): - JIT build through `gen_trtllm_comm_module` (`sm_121a` and `sm_120f`). - A functional test with two ranks sharing a single GPU over CUDA IPC. The lamport, oneshot, twoshot, and fusion paths all match fp32 references on both GPUs, including the small-batch lamport shapes the existing suites never reach. Before the pointer fix, those shapes crash or corrupt. - Both GPUs report max cluster size 8 for the real kernels and pass cluster-size-8 launches. - Cross-compilation for `sm_90a` and `sm_100a` is clean. - `pre-commit run` is clean. ## Reviewer Notes - The lamport dispatch thresholds (16 tokens or fewer, hidden size 256 or more) carry over from TRT-LLM's datacenter tuning. Whether lamport beats plain oneshot on SM12x needs a real multi-GPU measurement. - No API change. Test changes are only skip guards. 🤖 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** * Broadened TRT-LLM communication-kernel build support to additional GPU generations (SM90/100/100f/103/120/120f/121). * Added CUDA major version 12 handling for NVCC flag selection. * **Bug Fixes** * Added topology-limit validation for all-reduce IPC packing. * Corrected triple-buffer lamport IPC pointer packing stride. * Expanded architecture-guarded synchronization/launch-completion behavior. * Updated fusion and MoE cluster sizing to clamp using each device’s maximum supported cluster size. * **Tests** * Tests now run only when CUDA is available and compute capability is supported (9/10/12), with lamport buffers re-initialized per shape and expanded token coverage (including 8). <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
## 📌 Description Adds unified `UnpackedPrecomputed` routing support for `TrtllmFp4RoutedRunner`. Callers can provide separate contiguous routing tensors directly: - `topk_ids`: `[T, K]` `int32` global expert IDs - `topk_weights`: `[T, K]` BF16 and FP32 weights The runner forwards both tensors to the existing TRT-LLM FP4 Mode-3 launcher without constructing packed `(expert_id << 16) | bf16(weight)` IDs. Existing `PackedPrecomputed` and `FromLogits` behavior is unchanged. Coverage includes input validation, zero-copy forwarding, nonzero EP offsets, CUDA graph replay, and unified-fuzzer cases. ## 🚀 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. - [x] All tests are passing. SM100 results: ```bash pytest -v \ tests/moe/test_unified_moe.py::TestActivationPackValidation \ tests/moe/test_unified_moe.py::TestTrtllmFp4UnpackedContract \ tests/moe/test_unified_moe.py::TestTrtllmEPOffset 26 passed, 2 warnings ``` ```bash FLASHINFER_UMOE_FUZZ=1 \ FLASHINFER_UMOE_FUZZ_ONLY_SEED=900017 \ pytest -v tests/moe/test_unified_moe_fuzz.py::test_unified_moe_fuzz 2 passed, 4 warnings ``` ## Reviewer Notes This PR intentionally supports only the TRT-LLM FP4 unified runner. FP8 block, BF16, and MxInt4 launchers still use packed expert IDs and are out of scope. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added support for unpacked precomputed routing in TRTLLM FP4 MoE execution. - Uses separate `int32` expert-id tensors and BF16/FP32 weight tensors (no repacking), preserving caller-provided tensors. - Enforces routing contracts for dtype, shape, contiguity, and unsupported-field rejection. - Extended autotuning and CUDA-graph execution to cover the new routing mode. - **Tests** - Expanded unit, fuzz, expert-parallel, autotuner, and CUDA-graph coverage for unpacked routing. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…ai#4178) ## Why Kimi K3 MLA uses 96 global query heads with one KV head, TP-local head counts of 48/24/12/6, speculative query lengths up to 8, variable request-local query lengths, and context parallelism for long contexts. The existing Blackwell decode paths do not cover that combination efficiently or completely. ## Changes - Pack `(query_token, query_head)` rows continuously into CuTeDSL M128 tiles, so partial head groups can share a tile across tokens and only the final tile is padded. - Add compact variable-Q support to monolithic CuTeDSL with `[total_q, Hq, D]`, `cum_seq_lens_q`, and `max_q_len`. - Extend TRTLLM-GEN dense and sparse MLA to non-power-of-two `Hq <= 64` with padded Q8/Q16/Q32/Q64 tiles; `65 <= Hq <= 127` is routed to CuTeDSL. - Normalize split-KV to nonempty K partitions and parallelize underfilled D512 reduction across 1/2/4 output bands. - Classify dense versus causal work per M128 query tile so fully visible K tiles skip masking. - Add static cyclic decode context parallelism to monolithic BF16/FP16/FP8 CuTeDSL MLA. `cp_world` is compile-time, `cp_rank` is runtime, physical KV lengths remain rank-local, and `causal_seqlens_kv_global` supplies the global causal bound. The kernel returns rank-local output/LSE states for caller-side cross-rank merging. ## Correctness fixes - Close FP8 two-softmax named-barrier generations before a persistent 2-CTA cluster advances, fixing the upstream multi-work-item hang. - Bound partial-query TMA descriptors and predicate output/LSE tails. - Keep split count, workspace sizing, scheduler geometry, and autotuner keys consistent. - Handle DCP ranks or splits with no visible local keys as the neutral state: `O=0`, `LSE=-inf`. - Reject incompatible DCP combinations explicitly: DCP is monolithic fixed-Q only, requires `return_lse=True`, and cannot be combined with sinks or compact variable Q. ## Performance B200 CUDA-graph measurements for the packed-query and reducer changes showed: - up to 1.27x for flat query packing at B148/K8K; - 1.06x at B1 and 1.08x at B16 for H96/Sq8; - 9.7-13.0% end-to-end improvement from adaptive D2/D4 reduction; - no repeatable regression in a cold-L2 CUDA-graph sweep over `B={1,4,16,32,128}`, `K={1K,16K,128K}`, `Hq={12,24,48,96}`, `Sq={2,4,8}`, and BF16/FP8. DCP is statically specialized, so its disabled path does not add runtime DCP work. ## Validation Validated on B200 with CuTeDSL 4.5.0: - DCP matrix: 27 passed, covering BF16/FP16/FP8, `Hq={6,12,24,48,96,128}`, `Sq={1,4,8}`, `cp_world={1,2,4,8,16}`, heterogeneous/empty local KV, split-KV, PDL, and CUDA graphs. - Existing CuTeDSL MLA decode module: 678 passed. - Remaining MLA attention modules, including TRTLLM-GEN, DeepSeek MLA, paging, sparse MLA, XQA, and backend selection: 14,404 passed; 24,506 architecture-specific skips. - Full `pre-commit run --all-files`: passed, including clang-format, mypy, Ruff check, and Ruff format. ## Current limitations - CuTeDSL MLA requires SM100+, `Hq <= 128`, latent dimension 512, and RoPE dimension 64. - DCP requires monolithic fixed-Q and external cross-rank output/LSE reduction. - Runtime per-request variable split-KV is not qualified; automatic fixed split-KV is. - Partial final M128 rows remain compute lanes, although their memory accesses and stores are bounded or predicated. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added compact variable-length query support for MLA decode, including ragged query batches and inferred maximum query lengths. * Added static distributed context parallelism (DCP) support for monolithic MLA decoding. * Improved automatic backend selection for variable-query, DCP, and incompatible feature combinations. * Added support for non-power-of-two head configurations across MLA kernels. * **Bug Fixes** * Improved query masking, workspace sizing, output layouts, and handling of empty or inactive query rows. * Added clearer validation for unsupported or conflicting decoder options. * **Tests** * Expanded coverage for variable queries, DCP, CUDA graphs, FP8, backend routing, and boundary conditions. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Perkz Zheng <PerkzZheng@users.noreply.github.com> Co-authored-by: mingyangw <mingyangw@nvidia.com>
## Summary Follow-up to flashinfer-ai#4758. - Default B200/B300 (and moe_ep) jobs still uninstall and install `==4.6.2`. - other internal jobs on the laternate branch skip the pin and leave the image/prep package. - Same skip in `docker/install/build_flashinfer_ep_pytorch.sh` so CUDA 13 moe_ep setup does not re-pin those jobs. A global `>=4.6.2` would leave default CI on the baked 4.7.0; this keeps the exact pin everywhere except the Rubin env flags GitLab already passes into the container.
…al (flashinfer-ai#4667) 0.6.18 has no kda_prefill_cute / flashinfer-ai#4605 auto CuTe prefill path (public recurrent_kda already uses Cake), so only the shared is_cute_dsl_experimental_available probe is cherry-picked.
…infer-ai#4732) Alternate the math/IO barrier ids by tile parity so a leading math warp cannot double-arrive one phase (flashinfer-ai#3700).
…107)" This reverts commit c3b96d0.
…er-ai#4469) (cherry picked from commit 9a0e83b)
## 📌 Description CUDA 13.4 changes `nvcc --dryrun` output in a way that sccache v0.17.0 parses incorrectly. The missing compile steps later surface as `fatbinary` failures because the expected cubins were never produced. - Bypass sccache for cu134 NVCC invocations while keeping host C++ compilation cached. - Accept both CUDA version forms used by the release/nightly (`13.4`) and PR (`134`) build paths. - Trigger the Release dry-run matrix when the shared JIT-cache helper changes, so both cu134 architecture jobs exercise this workaround before merge. - Remove the guard once the pinned sccache release includes the upstream CUDA 13.3+ fix. ## 🔍 Related Issues - Upstream fix: [mozilla/sccache#2722](mozilla/sccache#2722) - Failing nightly: [run 32544126473](https://github.com/flashinfer-ai/flashinfer/actions/runs/32544126473) ## 🚀 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.). Focused validation passed: - `bash -n scripts/jit_cache_build_common.sh` - `shellcheck scripts/jit_cache_build_common.sh` - Mocked launcher checks for CUDA `13.4`, `134`, and `13.0` - Release workflow YAML parse - `pre-commit run --all-files` - `git diff --check` ## Reviewer Notes This intentionally disables only the affected NVCC launcher for cu134. The sccache server and host C++ launcher remain enabled so safe cache hits are preserved. The Release workflow should provide the end-to-end cu134 x86_64 and aarch64 validation. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved CUDA 13.4 build compatibility by avoiding an incompatible compiler-cache path. * Preserved compiler caching for other supported CUDA versions. * **Chores** * Updated release automation to recognize changes affecting shared build tooling. * Added clearer build logs showing which compiler-cache launchers are enabled. <!-- end of auto-generated comment: release notes by coderabbit.ai --> (cherry picked from commit cf9a04d)
(cherry picked from commit ac1c275)
…i#4750) ## 📌 Description Keep `test_supported_jit_cache_versions_match_cuda_config` active in source-tree presubmit runs, but skip it when `ci/cuda-versions.json` is intentionally absent from the isolated installed-package layout used by Nightly Release. This fixes the identical cu129/cu130 shard 5 failures in Nightly Release flashinfer-ai#360 without copying repository-only CI metadata into the installed-package test directory. ## 🔍 Related Issues - Nightly Release flashinfer-ai#360: https://github.com/flashinfer-ai/flashinfer/actions/runs/32920335499 ## 🚀 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.). Validation performed: - `python3 -m compileall -q tests/cli/test_cli_cmds.py` - `pre-commit run --files tests/cli/test_cli_cmds.py` - `git diff --check` Targeted pytest was not run locally because the existing repository virtualenv does not contain a usable pytest installation. PR CI should confirm both the source-tree pass and isolated-nightly skip behavior. ## Reviewer Notes The test still protects the duplicated CLI/config CUDA-version invariant during presubmit. The skip applies only when the repository-only `ci/cuda-versions.json` file is unavailable, as it is in the Nightly Release installed-wheel test directory. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Improved CUDA version consistency testing. * Tests now skip gracefully when the required configuration file is unavailable. * Added support for running tests with `pytest`. <!-- end of auto-generated comment: release notes by coderabbit.ai --> (cherry picked from commit e4b7fa4)
Drop SM75 from all cache wheels and SM121a from the cache arch lists so published jit-cache artifacts stay under GitHub's release-asset limit.
…uilds (flashinfer-ai#4760) Skip the single-request attention modules in AOT jit-cache so those kernels stay JIT-only and the cache wheel stays smaller.
attention-ts context and decode reject compute capability 10.7 outright:
NotImplementedError: attention-ts context requires an SM100a/B200 or
SM103a/B300 GPU; device cuda:0 has compute capability (10, 7)
The exclusion carries no rationale. Unlike the tinygemm2 capability tuple it
has no explaining comment, the module does not declare native_only=True (the
mechanism for kernels that need arch-specific tcgen05 instructions), and its
tcgen05 usage is limited to tcgen05_alloc/dealloc, which are family-portable.
Widen both tuples to include (10, 7) and gate on the CuTe DSL arch, so a DSL
older than 4.8 without CUTE_DSL_ARCH=sm_100f gets the actionable message from
require_cute_dsl_arch instead of KeyError: sm_107a during compilation.
Measured on SM107 hardware, tests/attention/test_attention_ts_context.py:
test_attention_ts_context_plan_rejects_critical_public_contracts
before: 3 failed, 2 passed
after : 5 passed
That test is a negative test -- it feeds deliberately invalid inputs and
asserts plan() rejects each with a specific message. plan() only validates and
builds scheduling metadata, so it never compiles a kernel; with the gate closed
the device check short-circuits the contract checks mid-sequence, which is why
two cases passed and three reported the arch error instead.
KNOWN LIMITATION: this fixes the gate, not attention-ts on Rubin. The decode
suite is unchanged at 4 failed / 45 passed / 76 skipped before and after,
because those four fail earlier at kernel construction in the exhaustive
deadlock/race checker -- a failure that also reproduces on B200 (cc 10.0) and
on SM107 with stock code, so it is not Rubin-specific. Landing this alone
trades a clear NotImplementedError for an opaque checker ValueError on the
decode path.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
(cherry picked from commit dd0646f)
…infer-ai#4761) Decline Rubin MoE tactics when the CuTe DSL lacks rubin_helpers, and pin the multi-arch Rubin batched-GEMM cubin package.
…ass (flashinfer-ai#4753) Move availability probes to a cutlass-free module so import flashinfer does not hard-require nvidia-cutlass-dsl. Paths that do not exist on 0.6.18 (kda_prefill_cute, SM120 KDA prefill runtime, attn_scores) are omitted.
(cherry picked from commit fd433de)
install_and_verify uninstalled the image 4.7.0 stack then pip install --no-deps nvidia-cutlass-dsl==4.6.2, which is only the metapackage. import cutlass then failed (GitLab job 413590603). Keep --no-deps for requirements.txt; reinstall 4.6.2 with deps.
Keep architecture resolution in the cutlass-free availability module and re-export it from utils so SM100 CUDA 13 GDN tests can import and collect.
…t K-smoothing (flashinfer-ai#4654) Cherry-pick of 857bc3a. The FMHA cubin pin was already 2d6a5a02 from flashinfer-ai#4648; keep that and take the sageQuant ragged/batch-size fix needed by that pack.
) (flashinfer-ai#4786) ## Summary Fixes flashinfer-ai#4773: `mm_fp8` / `test_mm_fp8_replay` SIGSEGV in `cuModuleGetFunction` during low-latency GEMM autotune on SM100 (B200/GB200) and SM103 (B300). After flashinfer-ai#4648 the trtllm-gen GEMM pack is a single multi-arch publish, so `flashinferMetaInfo.h` now lists sm100f **and** sm107a configs in the manifest the Blackwell module compiles against. Previously the Rubin cubins lived in a separate `TRTLLM_GEN_GEMM_RUBIN` pin, so the non-Rubin manifest was sm100-only. `trtllm_low_latency_gemm_runner.cu` was the one trtllm-gen runner without the arch filter that flashinfer-ai#4280 added to its siblings. With the consolidated pack, `getValidTactics()` returned 16 tactics on Blackwell (8 sm100f + 8 sm107a); `cuModuleLoadData` fails on the first sm107a cubin, the `CUresult` is ignored by the generated `GemmInterface`, and `cuModuleGetFunction` faults on the uninitialised `CUmodule`. Two commits, no cubin regeneration: 1. **`isArchCompatible` filter** when building `mPassingConfigIndices`, identical to `csrc/trtllm_gemm_runner.cu` (`Sm107a` only under `TLLM_RUBIN_FEATURES`, `Sm100f` allowed on sm100 and sm103). 2. **`checkPassingConfigIndex` in `run()`**, also mirroring `trtllm_gemm_runner.cu`. Tactic ids are manifest indices, and the autotuner's file-config key (`custom_op`, `runner_class_name`, `nearest_profile`, `extras`) does not include the device arch, so a config saved via `save_configs()` / `autotune(cache=...)` on other hardware — or an explicit FFI tactic — could still hand a foreign-arch index straight to the cubin loader. It now raises instead of faulting. Other ops touched by flashinfer-ai#4648 already have the equivalent guard, so no further coverage is needed: | Consumer | Arch filter | |---|---| | `trtllm_gemm_runner.cu` | `isArchCompatible` + `checkPassingConfigIndex` (flashinfer-ai#4280) | | `trtllm_batched_gemm_runner.cu` (trtllm-gen MoE backend) | `isArchCompatible` + `checkPassingConfigIndex` (flashinfer-ai#4280) | | trtllm-gen FMHA | `isSMCompatible()` in `fmhaKernels.cuh`, with explicit sm107 rules | | `trtllm_low_latency_gemm_runner.cu` | **missing — this PR** | ## Test plan Local B200 (SM100, CUDA 13.0, Python 3.10), on `release-v0.6.18` + these commits: - [x] Before the filter: 16 valid tactics (sm100f indices `0,2,3,4,5,7,10,11` + sm107a `93,95,96,97,101,102,104,109`); SIGSEGV on the first sm107a cubin load. - [x] After: 8 sm100f tactics only; `mm_fp8` passes under `autotune()` and on the heuristic `tactic=-1` path. - [x] Forced sm107a tactic (`93`) now raises `RuntimeError: Tactic 93 is not in this runner's compatible config set` instead of SIGSEGV. - [x] `pytest tests/gemm/test_mm_fp8.py tests/utils/test_logging_replay.py` → 44 passed, 2 skipped (includes `test_mm_fp8_replay`, the test that crashed in CI). - [ ] GitLab `unit_test_b300` / GB200 jobs covering `tests/gemm/test_mm_fp8.py` and `tests/utils/test_logging_replay.py`.
…oE JIT build time (flashinfer-ai#4789) <!-- .github/pull_request_template.md --> ## 📌 Description Cherry-picks the trtllm-gen MoE compile-time fix onto `release-v0.6.18`. The branch already carries the artifact consolidation from flashinfer-ai#4648; this adds the two follow-up changes it is missing. **1. Arch-filtered kernel manifest.** Consolidating the packs doubled the BMM `flashinferMetaInfo.h` from 3476 to ~6850 entries. That header declares one `BatchedGemmConfig` initializer per cubin, and `BatchedGemmConfig` embeds `BatchedGemmOptions`, which has non-trivial members (`std::string`, `std::vector<int>`) — so the array cannot be constant-folded into `.rodata` and the host compiler emits a dynamic initializer for every entry, at superlinear cost. Exactly one TU includes it (`csrc/trtllm_batched_gemm_runner.cu`), and it became 99.96% of the critical path of the fused-MoE JIT build. The runtime never dispatches across architecture families (`isArchCompatible()` accepts `Sm100a/Sm100f/Sm103a` only on sm100/sm103, `Sm107a` only on sm107), so the other family's entries are dead weight. `flashinfer/jit/trtllm_gen_metainfo.py` strips them by their `}, gemm::SmVersion::SmXXX},` terminator and rewrites the declared `...ListLen`, raising rather than silently truncating if the manifest shape ever changes. `gen_trtllm_gen_fused_moe_sm100_module` writes the filtered copy into its per-module `gen_root` and drops the raw artifact `include/` dir from the include path, so exactly one manifest is reachable. The filter keys on the **module variant** (`enable_rubin`), not the GPU present at build time — AOT builds both variants on one machine. No cubin is dropped: the variants partition the manifest exactly (3468 + 3378 = 6846, disjoint arch sets). **2. `TRTLLM_GEN_BMM` pin bump.** cubin_publishing `8ec29a98` ("[fix] Remove invalid config") is the direct child of the currently pinned `2d6a5a02` and deletes one line from `batched_gemm_config.json` (`[false, "none", "relu2", "bf16", "bf16", "128x4", true]`), removing 16 kernels. Both packs come from the same trtllm-gen commit (`fa419f4`), so this is purely the invalid-config removal. `main` already carries this pin via flashinfer-ai#4618; the release branch was still on the older one. Measured on B200 (CUDA 13.0), cold JIT build of `fused_moe_trtllm_sm100`: | | before | after | speedup | |---|---|---|---| | manifest entries | 6862 (29.2 MB) | 3476 (14.8 MB) | 1.97x smaller | | `trtllm_batched_gemm_runner.cu` | 1175.2 s | 205.1 s | **5.7x** | | whole module build | 1175.7 s | 226.4 s | **5.2x** | After the change that TU is no longer the critical path — the build is bounded by an unrelated routing kernel at 225.9 s. ## 🔍 Related Issues Cherry-pick of the compile-time follow-up in flashinfer-ai#4648. Fixes the trtllm-gen MoE test timeout seen in v0.6.18rc9. ## 🚀 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. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Validated on B200 (`umb-b200-261`, CUDA 13.0): - `tests/jit/test_trtllm_gen_metainfo.py` — 11 passed (new; pure text-transform, no GPU/network) - `tests/moe/test_trtllm_gen_moe_autotune_tactics.py` — 82 passed (25 min). This is the meaningful gate: its `*_all_tactics_are_correct` cases enumerate every valid config index and check numerics per tactic. ## Reviewer Notes - Filtering renumbers `configIndex` into the manifest. In-process this is consistent (`getValidConfigIndices` -> autotuner -> `checkPassingConfigIndex`), and `prioritizePredefinedConfigs` matches by kernel name, not index. A *persisted* tactic recorded against an unfiltered build (a `FLASHINFER_TACTICS_BLOCKLIST` JSON, or a saved autotune result) would be stale — the same hazard as any pin bump. No blocklist JSONs are checked in, and the autotune-tactics suite above covers the in-process path. - trtllm-gen FMHA uses a different compat rule (`isSMCompatible()` in `fmhaKernels.cuh` *does* accept `kSM_100f` on sm107) and its manifest is flat POD that constant-folds — measured at 0.98 s for all 35396 entries. It does not need this treatment. The GEMM manifest is 186 entries, also negligible. - The speedup table was measured against the previous pin (6862 entries) before the bump in commit 2; the new pack is 6846, a 0.2% difference. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
## 📌 Description Backport of two SM107 (Rubin) fixes from flashinfer-ai#4787 to `release-v0.6.18`. Two files, +82 lines. **This is not identical to flashinfer-ai#4787**, deliberately: * The `cute_dsl/tuner.py` guard from flashinfer-ai#4787 is **omitted** — this branch already carries it via flashinfer-ai#4761. * The `isArchCompatible`/`Sm100f` change from flashinfer-ai#4787 is **omitted** — see *Deliberately excluded* below. * `flashinfer-ai#4474` was considered and **excluded** — see *Not included* below. * The `fix(moe)` commit shares flashinfer-ai#4787's title but **is not the same code** — see *Why the shared commit differs* below. ### 1. `fix(moe)` — decline the CuTe DSL NVFP4 backend on SM107 without CuTe DSL 4.8 `CuteDslNvfp4Runner._check_support()` checked only the activation type and the W4A16 per-token scale. On a public CuTe DSL 4.7.0 stack the runner therefore passed the support check, survived `build()`, and entered `MoELayer.runners` — and the failure surfaced from inside `forward()` instead of from backend selection: ``` NotImplementedError: The SM107 (Rubin) CuTe DSL gather/activation-fusion grouped GEMM requires CuTe DSL >= 4.8, which provides cutlass.utils.rubin_helpers; the installed CuTe DSL does not have it. ``` Probing the DSL at support-check time lets `MoELayer` drop the backend at build time, so `auto` routes elsewhere and callers that enumerate backends see it absent rather than failing mid-call. `tests/moe/test_unified_moe.py::test_each_backend_matches_reference` already anticipates exactly this — it skips a backend that is not in `layer.runners` — but nothing made that true for the DSL-version case. The probe is **arch-conditional on purpose**: only the SM107 kernels need `rubin_helpers`, so an older DSL remains fully usable on SM100/SM103. Gating unconditionally would drop a working backend on Blackwell. This complements the tactic-level guard already on this branch, which covers the autotuning path. A direct `forward(tactic=-1)` bypasses tactic filtering entirely, so the two guards cover different entry points and neither subsumes the other. ### 2. `test(moe)` — skip SM107 parameterizations the CuTe DSL kernels do not implement `fused_moe/cute_dsl/rubin/` holds a narrower specialisation of the Blackwell kernels: the gather kernel hardcodes SwiGLU and exposes no `activation_type`, its wrapper has no `a_per_token_scale_ptr`, and the finalize kernel has no unfused path. The `NotImplementedError`s for `use_a_per_token_scale`, `use_fused_finalize=False` and `GegluTanh` are accurate statements about kernel code that does not exist — product gaps, not defects — but they report as failures on every SM107 run (~162 occurrences per job). No dispatch-level fix is possible, and that was measured rather than assumed: the affected tests call `cute_dsl_fused_moe_nvfp4()` directly and contain zero `MoELayer` references, so there is no backend selection to influence. Tuner tactic predicates were never executed (Rubin branch hit count 0), `_check_support()` declines regressed two passing tests without fixing any, and a dispatch catch-and-fall-back fired zero times. The skip is decided **from the parameterization, before the test body runs**, so it cannot absorb a genuine regression — anything failing for a different reason still fails. The three parameters are parametrized only in this file, so no other MoE test is affected. ### Why the shared commit differs from flashinfer-ai#4787 Same title, different body. This branch has the cutlass-free `flashinfer/cute_dsl/availability.py`; `main` does not have it yet (it arrives with flashinfer-ai#4753), and there `cute_dsl/utils.py` imports `cutlass` at module scope. Here (correct for this branch): ```python from ..cute_dsl.availability import is_rubin_cute_dsl_available if not is_rubin_cute_dsl_available(): ``` On flashinfer-ai#4787 (correct for `main` until flashinfer-ai#4753 lands): ```python try: from ..cute_dsl.utils import is_rubin_cute_dsl_available rubin_dsl_available = is_rubin_cute_dsl_available() except ImportError: rubin_dsl_available = False ``` Importing `cute_dsl.utils` on this branch would reintroduce the hard `cutlass` dependency flashinfer-ai#4753 removed, so a user with no CuTe DSL would get `ModuleNotFoundError` from a *support check* rather than a graceful decline. That is why this branch must use `availability` and `main` currently cannot. Once flashinfer-ai#4753 merges, flashinfer-ai#4787 collapses to the same two lines used here. **Consequence for review:** these two PRs are not a change and its backport — they are two branch-specific responses to real divergence. A review comment on one does not automatically apply to the other, and the `Sm100f` change is reviewable only on flashinfer-ai#4787. ## 🔍 Related Issues Backport of flashinfer-ai#4787. **Deliberately excluded — `isArchCompatible` accepting `Sm100f` on SM107.** flashinfer-ai#4787 carries a change widening the `Sm100f` case in `csrc/trtllm_batched_gemm_runner.cu` and `csrc/trtllm_gemm_runner.cu` to accept `smVersion == 107`. It is omitted here for two reasons: 1. **It would be inert.** flashinfer-ai#4789 landed on this branch after rc9 and filters the manifest per module variant, with `RUBIN_CUBIN_ARCHS = ("Sm107a",)` — so the Rubin module's manifest contains no `Sm100f` entries for a widened check to match. 2. **It contradicts flashinfer-ai#4789's stated premise.** That change documents *"sm100f cubins are NOT loadable on Rubin for BMM/GEMM — unlike trtllm-gen FMHA, whose `isSMCompatible()` does accept `kSM_100f` on `kSM_107`."* flashinfer-ai#4787 reads the same asymmetry the opposite way. I have measurements that appear to contradict that premise (an A/B on SM107 hardware where widening the check took a MoE tactics suite from 158 arch-rejection errors and 2 passing tests to 0 errors and 56 passing, correctness assertions included) — but also one segfault in two patched runs, which is exactly the hazard "not loadable" would predict. That disagreement should be resolved with the author of flashinfer-ai#4789 rather than by landing opposing changes on two branches, so it is not part of this PR. **Not included — the exhaustive-checker aliasing race.** `tests/attention/test_attention_ts_decode.py::test_attention_ts_decode_keeps_alias_schedule_is_race_free` fails all four parameterizations on this branch: ``` ValueError: Exhaustive checker found 1 aliasing race(s) after exploring {89505, 93178, 121713, 134567} states: Softmax0Task writes tmemSoftmaxLocal0 vs MmaTask prod tmemS0 ``` `TmemSoftmaxLocalResource.get_tmem_requirements()` declares a TMEM allocation the kernel never uses when `keeps_stats_via_smem` is set, so the checker correctly flags an overlap with a resource that is not really touched. flashinfer-ai#4474 stops declaring it, and cherry-picking flashinfer-ai#4474 onto this branch was verified to take the four tests from failing to passing on SM107 hardware. It is **excluded** because flashinfer-ai#4474 is a feature commit — *"add PrimTS Q64/KV256 and paged GQA block-sparse attention"*, 39 files, ~17.9k insertions — so it would have made this PR 99.6% unrelated payload to close a **test-only** failure that affects no runtime behaviour. A minimal extraction is not clean either: removing only the `tmem_softmax_stats.py` guard from `main` makes all four fail again with a *different* error (`TMEM usage (576 columns) exceeds hardware capacity (512)`), because flashinfer-ai#4474 also drops alias-group wiring this branch still relies on. That leaves it as a scoping decision for the prims_ts owner rather than something to smuggle in here. The failure remains open on this branch, which is its status today. ## 🚀 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. - [ ] All tests are passing (`unittest`, etc.). Validated on SM107 hardware: | Change | Evidence | |---|---| | `test(moe)` skips | Affected subset on SM107 hardware — before: **3 failed / 11 passed / 0 skipped**; after: **0 failed / 11 passed / 3 skipped**. All three outcome counts checked against baseline, so the skips are exactly the three known gaps and no passing test was lost | | `fix(moe)` decline | Lazy-import paths **executed** on SM107 hardware: `_assert_rubin_cute_dsl_available()` runs, the `cute_dsl.utils` probe resolves and returns `True`, and the guarded `.rubin` kernel import succeeds. The *declining* branch is still unexercised — it needs a public CuTe DSL < 4.8 stack | ## Reviewer Notes * One case is matched by function identity, not by parameter.** `test_geglu_tanh_accuracy` sets its activation in the test body rather than via a parameter, so it is matched with `request.node.function.__name__ == "test_geglu_tanh_accuracy"`. An earlier revision used a substring match on `"geglu_tanh"`, which also caught `test_geglu_tanh_activation_is_supported` — a pure-Python assertion about `normalize_cute_dsl_moe_activation_type` that touches no kernel and passes on SM107. That silently cost one passing test; the exact match restores it, confirmed by the pass count above. * **`fix(moe)`'s decline branch is still unexercised.** Its import paths were executed on SM107 hardware, but the container ships CuTe DSL 4.8, so the probe returns `True` and the decline never fires. Only a public CuTe DSL 4.7.0 stack exercises it. An earlier revision of this commit imported `is_rubin_cute_dsl_available` from `cute_dsl.availability`, which exists on `release-v0.6.18` but **not** on `main`; because the import is function-local, `py_compile` and `ruff` both passed and only a runtime call would have caught it. It now imports from `cute_dsl.utils`, which provides the symbol on both branches, verified by execution.
…er-ai#4792) ## 📌 Description Five independent SM107 (Rubin) fixes for `release-v0.6.18`, found by a sharded validation sweep of this branch at `89cafe9a4`. Together they remove **133 test failures**. Each commit stands alone and can be reviewed separately. Every fix was A/B-validated on SM107 hardware, comparing **failed, passed and skipped** counts — not just the failure count. ### 1. `fix(moe)` — `check_support()` must not require a bound device Follow-up to flashinfer-ai#4790, which introduced a regression this branch is carrying. The SM107 CuTe DSL probe read `self.device` **before** deciding whether the architecture was relevant, so `check_support()` raised on a runner with no device attached: ``` AttributeError: 'CuteDslNvfp4Runner' object has no attribute 'device' ``` `tests/moe/test_unified_moe.py::TestMoERunnerSupport` builds runners with `__new__` and attaches only a config — a reasonable way to exercise a pure configuration check. **Not SM107-specific:** the attribute access precedes any compute-capability test, so it raised on every architecture, Blackwell included. Treats a missing device as "nothing arch-specific to decide". `MoELayer` always sets a device in `__init__` before calling `check_support()`, so the dispatch path is unaffected. *Validated: `TestMoERunnerSupport` 64 passed (was 2 failed).* ### 2. `test` — skip SM107 CuTe-DSL cases when the installed DSL predates them Two unrelated failures with one cause: the public stack ships CuTe DSL 4.7.0, which has no `sm_107` in its `Arch` enum and no `cutlass.utils.rubin_helpers`. * **`tests/gemm/test_bmm_fp8.py`** reported a *problem-shape* error for an unavailable backend. `_can_implement_config_sm107` instantiates the kernel class to call `can_implement`; on DSL 4.7.0 that raises `NotImplementedError`, and a surrounding `except Exception: return False` turns it into "this config is invalid" — so every entry of `SM107_AUTOTUNE_CONFIGS` is rejected and the user sees `No valid cute-dsl SM107 bmm_fp8 config for problem (...)`. The geometry was fine. Same pipeline, same test unit, same 7 nodes: internal DSL 4.8 → 7 passed; public DSL 4.7.0 → 6 passed / 1 failed. * **`TestCuteDslMoeW4A16`** is the only GPU-executing DSL class in `tests/moe/test_cute_dsl_fused_moe.py` missing `pytestmark = _requires_dsl_arch` — all 16 classes audited. Its entry point calls `require_cute_dsl_arch(..., native_only=True)`, exactly what that marker tests. *Validated both directions: with 4.7.0 both skip; with 4.8 neither skips and the W4A16 test proceeds into kernel compilation, so it does not over-skip.* ### 3. `test(moe)` — skip the hardcoded unfused-finalize case The SM107 skip fixture keys on parameterization, but `test_deterministic_finalize_numerical_accuracy` passes `use_fused_finalize=False` in its *body*, so it escaped and still hit the `NotImplementedError` the fixture exists to absorb. Matched by function identity, like `test_geglu_tanh_accuracy` — still decided before the body runs, so it cannot absorb a genuine regression. `test_route_tile_boundary_accuracy` and `test_weight_scale_mapping` also hardcode that flag and are **deliberately not skipped**: neither fails in CI, and they use the W4A16 entry point rather than the blockscaled finalize path, so skipping them would drop real coverage. ### 4. `fix(trtllm-gen)` — accept SM107 for `Sm100f` in the low-latency GEMM filter ``` RuntimeError: Check failed: (it != mPassingConfigIndices.end()) is false: Tactic 0 is not in this runner's compatible config set ``` `isArchCompatible` in `csrc/trtllm_low_latency_gemm_runner.cu` mapped `CudaArch::Sm100f` to `smVersion == 100 || smVersion == 103`, omitting 107, so every family-conditional cubin was discarded on Rubin. The mechanism was verified on hardware rather than assumed. `TLLM_RUBIN_FEATURES` **is** defined for the SM107 module (nvcc flags dumped), and the passing set is non-empty unpatched — 8 native `Sm107a` tactics. But `select_kernel()`, twelve lines above the filter, names its heuristic kernels *literally* and every name ends in `_sm100f`, so the index it returns is always one the filter dropped, and `checkPassingConfigIndex()` converts that into the error. Patched, the passing set goes 8 → 16 by adding exactly the `Sm100f` entries. *Validated: `test_mm_fp8` **30 failed → 30 passed**; `test_logging_replay` **1 failed → 16 passed**. Both arms genuinely rebuilt (12 s each, no timing asymmetry). `cos_sim > 0.99` assertions hold on the `Sm100f` kernels.* **Deliberately excluded — the cuDNN mixed-form seqlens failure (64 occurrences).** `_cudnn_supports_direct_seqlens(mixed=True)` authorises a paged path that every currently published `cudnn-frontend` rejects at graph validation, so the request fails with `Padding mask requires seq_len_q/seq_len_kv`. Verified against a wheel installed fresh from PyPI (`nvidia-cudnn-frontend==1.27.0`, the newest published): the mixed form is rejected there too, so this is **not Rubin-specific and not a container artifact** — it affects any architecture taking that path on a stock install. It is excluded from this Rubin-scoped PR and should be fixed separately. A capability-probe fix exists and is validated (32 failed -> 0, passed unchanged) but is held for its own PR. **Deliberately excluded — the PrimTS decode exhaustive-checker race (4 occurrences).** `TmemSoftmaxLocalResource.get_tmem_requirements()` declares a TMEM allocation the kernel never performs when `keeps_stats_via_smem` is set, so the exhaustive checker correctly reports an aliasing race against `tmemS0` and kernel construction fails with `ValueError: Exhaustive checker found 1 aliasing race(s)`. Scope, stated accurately: this is **not** test-only. The production path (`_run_decode_gen_active`) enables the checker via `not (cfg.use_keeps_mma_ab and cfg.num_insts_kv == 1)`, and all four affected profiles measure `num_insts_kv == 2`, so the checker runs for them outside tests too. The failure is a hard construction-time error, not a silent wrong answer, and no runtime behaviour changes either way — a candidate fix was verified to leave TMEM offsets and column counts byte-identical across ten decode profiles. It is excluded here as a deliberate release-management decision: the change is a scheduling-model edit in the PrimTS engine, which is owned elsewhere, and the risk of touching it unfamiliar outweighs a loud, characterised construction failure. A validated minimal backport exists (2 files, +51/-21, 4 failed -> 4 passed) and should be routed to the PrimTS owner rather than landed here. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit`. - [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. - [x] All tests are passing (`unittest`, etc.). | Fix | Before → after on SM107 | |---|---| | 1 `check_support` | `TestMoERunnerSupport` 2 failed → **64 passed** | | 2 DSL-4.7 skips | 1 failed → skipped (bmm_fp8); 1 failed → skipped (W4A16) | | 4 unfused-finalize skip | 1 failed → skipped | | 5 low-latency GEMM | 30 failed → **30 passed**; 1 failed → **16 passed** | ## Reviewer Notes * **`Sm100f` is not universally family-valid, and the file says so.** A few lines below the filter changed in commit 5 there is already a carve-out — `sm103` must fall back to `Sm103a` for the f2fp patch. sm107 is empirically fine here (correct numerics, verified), but a similar per-feature gap on Rubin would need the same treatment. * **Two copies of the same defect remain on this branch**, in `csrc/trtllm_gemm_runner.cu` and `csrc/trtllm_batched_gemm_runner.cu`, both still `100 || 103`. They are not producing failures on the current cubin pin and are deliberately out of scope here. * **Commit 5 is release-only by construction.** `csrc/trtllm_low_latency_gemm_runner.cu` has no `isArchCompatible` on `main`; flashinfer-ai#4773/flashinfer-ai#4786 added it to this branch only.
PDL in the trtllm-gen fused-MoE pipeline intermittently fails on Rubin with
"unspecified launch failure", and occasionally hangs. A/B stress runs on SM107
hardware isolated the trigger: with PDL enabled the renormalize-routing tests
crash across routing modes (split-topK on and off), dtypes (BF16/MxFP4/MxInt4)
and autotune on/off -- 4 crashes in ~21 full-file runs -- while the same loop
with PDL fully disabled ran clean. The crash surfaces at the
routingIndicesClusterKernel launch in trtllm_fused_moe_routing_custom.cu.
Clamp enable_pdl to off on SM107 at the six trtllm-gen MoE entry points,
including when the caller passes enable_pdl=True explicitly: an illegal memory
access is not something a caller should be able to opt into while the
underlying launch-dependency chain is unaudited for Rubin timing.
Scope is deliberately narrow:
* Only compute capability (10, 7). device_support_pdl() returns True for every
major >= 9, so gating there instead would cost Hopper and Blackwell for a
Rubin-only defect.
* Only the trtllm-gen path. The three CUTLASS MoE sites are untouched -- that
path has soaked with PDL enabled on Rubin for 9+ nights without a crash.
Measured cost of disabling PDL on SM107, trtllm-gen NvFP4xNvFP4 routed MoE
(128 experts, top-k 8, hidden 2048, intermediate 768), autotuned:
tokens PDL on PDL off cost
8 0.027 ms 0.028 ms +3.7%
64 0.046 ms 0.048 ms +4.3%
512 0.055 ms 0.061 ms +10.9%
4096 0.156 ms 0.158 ms +1.3%
The cost is launch-overlap latency, so it is largest mid-range and negligible
at 4096 tokens where the kernels are long enough that overlap stops mattering.
Single run per arm with CUDA-event timing (CUPTI unavailable), so the 1-4%
points are within noise; the 512-token gap reproduces in both the autotuned and
untuned columns.
Revert once the PDL launch-dependency chain is audited for Rubin.
(cherry picked from commit 1c0f82f)
PR #344 was squash-merged, so de79634 has a single parent and neither v0.6.18 nor the branch head is an ancestor of amd-integration. The tree is byte-identical -- the code landed correctly -- but git no longer knows the fork has merged v0.6.18, and three things read that ancestry: scripts/upstream_canary.py --upstream-ref v0.6.18 before: 27 conflicted, 19 of them code; 14 of 25 forked headers drifting after : clean merge -- nothing to do; no drift git merge-base amd-integration v0.6.18 before: 2628beb (2025-11-20, the pre-sync base) after : 69ff11f (bump version to 0.6.18) scripts/amd_coverage.py's derived base follows the same merge-base, so the ratchet still measures against v0.5.3 and reports the whole tree as owned. Cutting v0.6.18+amd.1 does not fix that on its own -- measured. Every future upstream sync would re-resolve the 27 conflicts already resolved in this tree. records the parent and changes no file. IMPORTANT: this must NOT be squash-merged. A squash drops the second parent, which is the entire content of this change, and would land an empty commit that looks like it worked.
There was a problem hiding this comment.
Copilot wasn't able to review any files in this pull request.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
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
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.
Summary
Records the
v0.6.18merge that #344's squash discarded. Zero files change —git diff de7963489..1539da2b8is empty, and the tree hash is identical. The entire content of this PR is the commit's second parent.Why
#344 was squash-merged, so
de7963489has one parent and neitherv0.6.18nor the PR's branch head is an ancestor ofamd-integration. The code landed correctly; git just no longer knows the fork has merged upstream. Three things read that ancestry, and all three are wrong today:scripts/upstream_canary.py --upstream-ref v0.6.18clean merge -- nothing to do; no driftgit merge-base amd-integration v0.6.182628bebcf(2025-11-20, the pre-sync base)69ff11fc4(bump version to 0.6.18)upstream/mainv0.6.18scripts/amd_coverage.pyderives its base from the same merge-base, so the ratchet still measures againstv0.5.3and reports the whole tree as owned. Cuttingv0.6.18+amd.1does not fix that on its own — measured: with the tag present but the ancestry absent,_resolve_basestill returns2628bebcf.The recurring cost is the first row. Every future upstream sync re-resolves the 27 conflicts that are already resolved in this tree.
The residual 101 commits behind
mainare genuine, not an artifact:v0.6.18is cut from a release branch that diverged frommainat61a6c6518, somainhas moved on independently. Closing that is a separate sync, not this PR.Reviewer notes
This must be merged with "Create a merge commit" — not squash. A squash drops the second parent, which is the whole change, and would land an empty commit that looks like it worked. The
amd-integrationruleset hasrequired_linear_history, which forbids merge commits;demandal25is a bypass actor withbypass_mode: pull_request, so merging this as a merge commit is the intended route and no rule needs to be relaxed.Longer term this is worth a decision: a fork whose job is periodic upstream merges is structurally at odds with
required_linear_history, and the alternative to a bypass every time is teachingupstream_canary.pyandamd_coverage.pyto read a recorded base instead ofgit merge-base.