From 22e971890927c196f51c0ca0a4e72b5342a6bb62 Mon Sep 17 00:00:00 2001 From: Rehvaro Date: Tue, 9 Jun 2026 20:24:22 +0200 Subject: [PATCH 01/13] jit: don't run ninja in build_and_load() if the .so is already there build_and_load() calls build() unconditionally, which always invokes ninja. On shared filesystems (GPFS, NFS) this causes a full recompile of the CUTLASS fused-MoE kernels (~180 .cu files, ~1.5h on sm90a) on every process restart. The underlying issue is a mtime discrepancy between what ninja records in .ninja_log during the post-link restat pass and what a subsequent stat() call returns on GPFS (inode flush timing). Ninja sees them as different, concludes the .o files are stale, and rebuilds everything. Fix: if the .so already exists, return early without touching ninja. The lock+build path is only taken on the first compilation. Added an inner exists() check inside the lock for the concurrent first-run case. --- flashinfer/jit/core.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/flashinfer/jit/core.py b/flashinfer/jit/core.py index 53c4163e267..fa732149b54 100644 --- a/flashinfer/jit/core.py +++ b/flashinfer/jit/core.py @@ -308,12 +308,17 @@ def build_and_load(self): if self.is_aot: return self.load(self.aot_path) + so_path = self.jit_library_path + if so_path.exists(): + return self.load(so_path) + # Guard both build and load with the same lock to avoid race condition # where another process is building the library and removes the .so file. with FileLock(self.lock_path, thread_local=False): so_path = self.jit_library_path - verbose = os.environ.get("FLASHINFER_JIT_VERBOSE", "0") == "1" - self.build(verbose, need_lock=False) + if not so_path.exists(): + verbose = os.environ.get("FLASHINFER_JIT_VERBOSE", "0") == "1" + self.build(verbose, need_lock=False) result = self.load(so_path) return result From 0b29eed266f41fa762a512725bdc8ac098a9e812 Mon Sep 17 00:00:00 2001 From: Andrew Gu <31054793+awgu@users.noreply.github.com> Date: Tue, 9 Jun 2026 12:06:49 -0700 Subject: [PATCH 02/13] Fix smem race in `FilteredTopK` overflow refinement (#3529) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## πŸ“Œ Description The `FilteredTopK` overflow refinement path has an smem race (flagged by `racecheck`). In the overflow refinement loop, each round reuses `s_histogram` ([code](https://github.com/flashinfer-ai/flashinfer/blob/a28703432faab15fda7edd71b6c80be0206df973/include/flashinfer/topk.cuh#L2733)): - near the end of round `r`, threads read `s_histogram` to update `topk_remain`: ```cpp topk_remain -= static_cast(s_histogram[threshold + 1]); ``` - at the start of round `r + 1`, threads clear `s_histogram`: ```cpp if (tx < RADIX + 1) s_histogram[tx] = 0; ``` This PR fixes by adding a `__synchtreads()` after the `topk_remain` decrement before the loop can either break or continue to the next round. ---
Repro Script ```bash compute-sanitizer --tool racecheck --target-processes all uv run python repro_flashinfer_topk_race.py ``` ``` #!/usr/bin/env python3 """Minimal FlashInfer radix top-k racecheck repro. This script intentionally uses a tie-heavy input so FlashInfer's filtered top-k kernel enters the shared-memory overflow refinement path. Example: python repro_flashinfer_topk_race.py Run under NVIDIA Compute Sanitizer: compute-sanitizer --tool racecheck --target-processes all \ python repro_flashinfer_topk_race.py Affected FlashInfer versions report shared-memory hazards in the radix top-k kernel. The script does not use any project-specific wrappers. """ from __future__ import annotations import argparse import sys from importlib import metadata import torch DTYPES = { "float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16, } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--rows", type=int, default=512) parser.add_argument("--vocab", type=int, default=131072) parser.add_argument("--k", type=int, default=128) parser.add_argument("--dtype", choices=sorted(DTYPES), default="float32") parser.add_argument("--iters", type=int, default=1) parser.add_argument("--deterministic", action="store_true", default=True) parser.add_argument("--nondeterministic", dest="deterministic", action="store_false") parser.add_argument("--sorted", action="store_true") return parser.parse_args() def main() -> int: args = parse_args() if not torch.cuda.is_available(): print("CUDA is required for this repro.", file=sys.stderr) return 1 import flashinfer from flashinfer.topk import can_implement_filtered_topk, get_topk_module device = torch.device("cuda") dtype = DTYPES[args.dtype] try: flashinfer_version = metadata.version("flashinfer") except metadata.PackageNotFoundError: flashinfer_version = getattr(flashinfer, "__version__", "unknown") print(f"torch={torch.__version__}") print(f"flashinfer={flashinfer_version} ({flashinfer.__file__})") print(f"cuda_device={torch.cuda.get_device_name(device)}") print(f"can_implement_filtered_topk={can_implement_filtered_topk()}") print( "case=" f"rows={args.rows}, vocab={args.vocab}, k={args.k}, dtype={args.dtype}, " f"sorted={args.sorted}, deterministic={args.deterministic}" ) # All zeros force all logits into the same radix bucket. For vocab=128K and # k=128 this enters the overflow refinement loop that reuses s_histogram. # The default 512 rows give racecheck enough CTAs to report the hazard # reliably on affected versions. logits = torch.zeros((args.rows, args.vocab), dtype=dtype, device=device) topk_module = get_topk_module() row_states_buffer = torch.zeros(1024 * 1024, dtype=torch.uint8, device=device) output_values = torch.empty((args.rows, args.k), dtype=dtype, device=device) indices = None for _ in range(args.iters): indices = topk_module.radix_topk( logits, args.k, args.sorted, args.deterministic, row_states_buffer, output_values, ) torch.cuda.synchronize() assert indices is not None assert output_values.shape == (args.rows, args.k) assert indices.shape == (args.rows, args.k) assert indices.dtype == torch.int32 assert bool(torch.all(output_values == 0).item()) assert bool(torch.all((indices >= 0) & (indices < args.vocab)).item()) print("completed") return 0 if __name__ == "__main__": raise SystemExit(main()) ```
Example Output ``` ========= COMPUTE-SANITIZER ========= Variable environment CUDA_COREDUMP_FILE is not supported by compute-sanitizer, clearing it before target process launch. ========= Variable environment CUDA_COREDUMP_GENERATION_FLAGS is not supported by compute-sanitizer, clearing it before target process launch. ========= Variable environment CUDA_ENABLE_COREDUMP_ON_EXCEPTION is not supported by compute-sanitizer, clearing it before target process launch. torch=2.11.0+cu130 flashinfer=0.6.8.post1 (/tmp/.../lib/python3.12/site-packages/flashinfer/__init__.py) cuda_device=NVIDIA GB300 can_implement_filtered_topk=True case=rows=512, vocab=131072, k=128, dtype=float32, sorted=False, deterministic=True ========= Error: Race reported between Read access at void flashinfer::sampling::FilteredTopKUnifiedKernel(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x111e0 ========= and Write access at void flashinfer::sampling::FilteredTopKUnifiedKernel(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x11230 [3852 hazards] ========= ========= Error: Race reported between Read access at void flashinfer::sampling::FilteredTopKUnifiedKernel(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x137a0 ========= and Write access at void flashinfer::sampling::FilteredTopKUnifiedKernel(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x137f0 [3596 hazards] ========= ========= Error: Race reported between Read access at void flashinfer::sampling::FilteredTopKUnifiedKernel(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x16160 ========= and Write access at void flashinfer::sampling::FilteredTopKUnifiedKernel(const T1 *, T2 *, T1 *, const T2 *, long, const T2 *, const T2 *, unsigned int, unsigned int, unsigned int)+0x161a0 [3560 hazards] ========= completed ========= RACECHECK SUMMARY: 3 hazards displayed (3 errors, 0 warnings) ```
## πŸ” 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 - [ ] 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 ## Summary by CodeRabbit ## Bug Fixes * Improved reliability of top-k filtering operations by resolving a synchronization issue in the multi-round fallback path to ensure accurate results across all computational scenarios. --- include/flashinfer/topk.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/include/flashinfer/topk.cuh b/include/flashinfer/topk.cuh index e9d7893d983..8992e108fc9 100644 --- a/include/flashinfer/topk.cuh +++ b/include/flashinfer/topk.cuh @@ -2769,6 +2769,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) threshold_bytes[round] = static_cast(threshold); topk_remain -= static_cast(s_histogram[threshold + 1]); + __syncthreads(); if (topk_remain == 0) { stop_round = round; break; From 3111f8fdb582ac8ea47e3f29182d32f2105d6357 Mon Sep 17 00:00:00 2001 From: Vincent <34876120+Vinnie6167@users.noreply.github.com> Date: Tue, 9 Jun 2026 13:44:23 -0700 Subject: [PATCH 03/13] Avoid workspace reset for standalone MLA decode (#3465) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## πŸ“Œ Description ## πŸ” 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 - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Running the benchmark before and after on a B200 machine gives: ``` python3 flashinfer_benchmark.py \ --routine BatchMLAPagedAttentionWrapper \ --backends trtllm-native \ --page_size 64 --batch_size 1 \ --s_qo 1 --s_kv 8192 \ --num_qo_heads 128 --num_kv_heads 128 \ --head_dim_ckv 512 --head_dim_kpe 64 \ --random_actual_seq_len -vv --refcheck \ --q_dtype fp8_e4m3 --kv_dtype fp8_e4m3 \ --generate_repro_command --case_tag DeepSeek-R1 \ --num_iters 100 --dry_run_iters 20 \ ``` - Before: 0.016 ms, 73.2 TFLOPs/s - After: 0.013 ms, 91.4 TFLOPs/s ## Reviewer Notes ## Summary by CodeRabbit * **Bug Fixes** * Improved workspace memory handling for MLA decode: the buffer region used for profiling is now cleared only during profiling runs. This prevents unnecessary memory operations during normal inference, reduces interference in shared-workspace scenarios, and improves performance and stability during autotuning/profile measurements. --- flashinfer/mla/_core.py | 50 +++++++++++++++++++++++++++++------------ 1 file changed, 36 insertions(+), 14 deletions(-) diff --git a/flashinfer/mla/_core.py b/flashinfer/mla/_core.py index 0872b4324f8..369602c29f8 100644 --- a/flashinfer/mla/_core.py +++ b/flashinfer/mla/_core.py @@ -1069,13 +1069,26 @@ def run( # Size of the trtllm-gen workspace counter region (multi-block semaphores) # per csrc/trtllm_fmha_kernel_launcher.cu:200: max_batch_size * max_num_qo_heads -# * sizeof(uint32_t) = 8192 * 256 * 4 = 8 MB. The kernel requires this region -# to be zero on its first use and self-resets at the end of each launch, but -# we must restore zeros after any other kernel writes into the shared -# workspace_buffer (e.g. cute-dsl's split-K scratch during autotune). +# * sizeof(uint32_t) = 8192 * 256 * 4 = 8 MB. trtllm-gen places this counter +# slab at the head of the workspace_buffer and self-resets it at the end of +# every launch, so back-to-back trtllm-gen launches keep it valid without any +# host-side zeroing. _TRTLLM_GEN_MLA_COUNTER_REGION_BYTES = 8192 * 256 * 4 +def _cute_dsl_workspace_view(workspace_buffer: torch.Tensor) -> torch.Tensor: + """Sub-view of the shared workspace that skips trtllm-gen's counter region. + + cute-dsl carves its scratch from offset 0 of whatever buffer it is given; + offsetting past the 8 MB counter region keeps it from writing into the + bytes trtllm-gen needs zero on entry. Costs 8 MB of usable workspace + (callers should size the buffer accordingly; the recommended 128 MB has + ample headroom). + """ + workspace_i8 = workspace_buffer.reshape(-1).view(torch.int8) + return workspace_i8[_TRTLLM_GEN_MLA_COUNTER_REGION_BYTES:] + + def _round_to_seq_len_bucket(x: int) -> int: """Power-of-2 bucket for max_seq_len in autotune cache keys. @@ -1146,8 +1159,14 @@ def _compute_mla_decode_buckets( if "cute-dsl" in runner_names: from ..cute_dsl.utils import get_num_sm + # cute-dsl gives up the counter region only when trtllm-gen shares the + # buffer, so its usable size excludes that reservation only then. + reserved = ( + _TRTLLM_GEN_MLA_COUNTER_REGION_BYTES if "trtllm-gen" in runner_names else 0 + ) cute_dsl_cap = _cute_dsl_max_supported_batch( - workspace_bytes=workspace_buffer.numel() * workspace_buffer.element_size(), + workspace_bytes=workspace_buffer.numel() * workspace_buffer.element_size() + - reserved, q_len=q_len, num_heads=num_heads, kv_lora_rank=kv_lora_rank, @@ -1435,14 +1454,6 @@ def forward( lse_stride_tokens = 0 lse_stride_heads = 0 - # Zero the counter region on every call. Other runners (cute-dsl) - # may share this workspace_buffer and write scratch into the first - # bytes, which would leave non-zero values in trtllm-gen's - # mandatory-zero semaphore region and cause kernel hangs. Done - # inside forward() rather than only at dispatcher final-call time so - # that autotune profile-loop invocations are also protected. - # The 8 MB memset is ~5us on B200, negligible vs kernel time. - self.workspace_buffer[:_TRTLLM_GEN_MLA_COUNTER_REGION_BYTES].zero_() self._run( out, None, # fp4 output (unsupported by wrapper) @@ -1509,12 +1520,20 @@ def __init__( return_lse: bool, sinks: Optional[torch.Tensor], cute_dsl_impl: str, + reserve_counter_region: bool = False, ): from ..cute_dsl.attention import cute_dsl_mla_decode self._run = cute_dsl_mla_decode self.kv_cache = kv_cache - self.workspace_buffer = workspace_buffer + # Only skip trtllm-gen's counter region when trtllm-gen shares this + # buffer (the "auto" path); a standalone cute-dsl runner owns the whole + # buffer and can use it from offset 0. + self.workspace_buffer = ( + _cute_dsl_workspace_view(workspace_buffer) + if reserve_counter_region + else workspace_buffer + ) self.kv_lora_rank = kv_lora_rank self.qk_nope_head_dim = qk_nope_head_dim self.qk_rope_head_dim = qk_rope_head_dim @@ -1971,6 +1990,9 @@ def trtllm_batch_decode_with_kv_cache_mla( return_lse=return_lse, sinks=cute_dsl_sinks, cute_dsl_impl=cute_dsl_impl, + # Reserve trtllm-gen's counter region only when it co-runs on + # the shared workspace (the "auto" path). + reserve_counter_region="trtllm-gen" in runner_names, ) ) From 49cb250fa00d63b5f732a5d0c3e81af84f2f2c52 Mon Sep 17 00:00:00 2001 From: "Brian K. Ryu" Date: Tue, 9 Jun 2026 16:06:18 -0700 Subject: [PATCH 04/13] perf(sampling): Optimize top_k_top_p_sampling_from_logits/from_probs for large-vocab small-k sampling (#3461) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## πŸ“Œ Description ### Summary Adds a faster `filter_apply_order="top_k_first"` path for `top_k_top_p_sampling_from_logits` and `top_k_top_p_sampling_from_probs`. For modest `top_k` over a large vocabulary, instead of masking/renormalizing the full vocab and running rejection sampling across it, we select the top-k entries with the parallel radix top-k kernel and run top-p **over only those k survivors**. This is distribution-equivalent to the existing path but 2–4Γ— faster for `from_logits` (see benchmarks). Addresses the slow small-batch top-k/top-p sampling reported in #3389. ### Motivation The existing `top_k_first` path keeps everything in full-vocab layout: - `from_logits`: `top_k_mask_logits` β†’ `torch.softmax` (full vocab) β†’ `top_p_sampling_from_probs` (single-CTA rejection over full vocab) - `from_probs`: `top_k_renorm_probs` β†’ `top_p_sampling_from_probs` (single-CTA rejection over full vocab) After top-k, only `k` entries are relevant, but the general full-vocab kernels still process all `V` elements β€” a full-vocab softmax (logits) and a multi-round single-CTA rejection scan (both). At small batch the rejection kernel launches one CTA per request and underutilizes the GPU. ### Changes A gated fast path in `flashinfer/sampling.py` (no kernel changes): 1. Select top-k via `flashinfer.top_k(..., sorted=True)` (parallel radix selection, returns the `k` values + their indices). 2. Normalize over the `k` survivors β€” `softmax` (logits) or renorm (probs); these are mathematically identical to the masked full-vocab versions. 3. Run `top_p_sampling_from_probs` over the `k`-element distribution, then map the local choice back to the global vocab index via `gather`. Both entry points route through the same helper with `sorted=True`, so they reduce to the same `probs_k` and stay **sample-aligned** with each other. ### Accuracy - **Distribution-equivalent** to the original `top_k_first` path (validated: total-variation distance β‰ˆ 0.007 vs. the analytic target over 40k draws). - **Per-seed sample values may differ** from previous versions (the RNG now maps over a `k`-element domain, and radix top-k may break k-th-boundary ties differently). The two APIs remain aligned with each other. - **CUDA-graph compatible**: the default vectorized top-k captures into a graph; the graph-safe fallback is only requested for the non-deterministic cluster path. ### Benchmarks `top_k=50`, `top_p=0.9`, CUDA-graph mode, median ms. "before" = original path, "after" = fast path. ### B200 (SM100) | API | dtype | bs | vocab | before | after | speedup | |---|---|---|---|---|---|---| | from_logits | bf16 | 1 | 128k | 0.121 | 0.059 | 2.05Γ— | | from_logits | bf16 | 8 | 128k | 0.139 | 0.064 | 2.17Γ— | | from_logits | bf16 | 32 | 128k | 0.147 | 0.066 | 2.23Γ— | | from_logits | bf16 | 256 | 128k | 0.431 | 0.180 | 2.39Γ— | | from_logits | bf16 | 1 | 256k | 0.206 | 0.067 | 3.07Γ— | | from_logits | bf16 | 32 | 256k | 0.294 | 0.075 | 3.92Γ— | | from_logits | fp32 | 1 | 128k | 0.126 | 0.057 | 2.21Γ— | | from_logits | fp32 | 8 | 128k | 0.131 | 0.063 | 2.08Γ— | | from_logits | fp32 | 32 | 128k | 0.163 | 0.068 | 2.40Γ— | | from_logits | fp32 | 256 | 128k | 0.453 | 0.122 | 3.71Γ— | | from_logits | fp32 | 1 | 256k | 0.208 | 0.061 | 3.41Γ— | | from_logits | fp32 | 32 | 256k | 0.296 | 0.113 | 2.62Γ— | | from_probs | fp32 | 1 | 128k | 0.076 | 0.064 | 1.19Γ— | | from_probs | fp32 | 8 | 128k | 0.077 | 0.067 | 1.15Γ— | | **from_probs** | fp32 | 32 | 128k | 0.088 | 0.105 | **0.84Γ—** | | from_probs | fp32 | 256 | 128k | 0.365 | 0.155 | 2.35Γ— | | from_probs | fp32 | 1 | 256k | 0.099 | 0.067 | 1.48Γ— | | from_probs | fp32 | 32 | 256k | 0.235 | 0.116 | 2.03Γ— | ### RTX PRO 6000 (SM120) | API | dtype | bs | vocab | before | after | speedup | |---|---|---|---|---|---|---| | from_logits | bf16 | 1 | 128k | 0.096 | 0.042 | 2.29Γ— | | from_logits | bf16 | 8 | 128k | 0.111 | 0.044 | 2.52Γ— | | from_logits | bf16 | 32 | 128k | 0.143 | 0.054 | 2.65Γ— | | from_logits | bf16 | 256 | 128k | 0.454 | 0.173 | 2.62Γ— | | from_logits | bf16 | 1 | 256k | 0.159 | 0.042 | 3.79Γ— | | from_logits | bf16 | 32 | 256k | 0.232 | 0.081 | 2.86Γ— | | from_logits | fp32 | 1 | 128k | 0.101 | 0.041 | 2.46Γ— | | from_logits | fp32 | 8 | 128k | 0.126 | 0.045 | 2.80Γ— | | from_logits | fp32 | 32 | 128k | 0.162 | 0.079 | 2.05Γ— | | from_logits | fp32 | 256 | 128k | 0.631 | 0.287 | 2.20Γ— | | from_logits | fp32 | 1 | 256k | 0.171 | 0.042 | 4.07Γ— | | from_logits | fp32 | 32 | 256k | 0.287 | 0.113 | 2.54Γ— | | from_probs | fp32 | 1 | 128k | 0.056 | 0.043 | 1.30Γ— | | from_probs | fp32 | 8 | 128k | 0.061 | 0.048 | 1.27Γ— | | from_probs | fp32 | 32 | 128k | 0.110 | 0.083 | 1.33Γ— | | from_probs | fp32 | 256 | 128k | 0.405 | 0.293 | 1.38Γ— | | from_probs | fp32 | 1 | 256k | 0.080 | 0.044 | 1.82Γ— | | from_probs | fp32 | 32 | 256k | 0.252 | 0.115 | 2.19Γ— | ### Commands to reproduce results above ``` python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 1 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 8 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 32 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 256 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 1 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 32 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype bfloat16 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 1 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 8 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 32 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 256 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 1 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_logits --batch_size 32 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first --input_dtype float32 python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 1 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 8 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 32 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 256 --vocab_size 128000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 1 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first python flashinfer_benchmark.py --routine top_k_top_p_sampling_from_probs --batch_size 32 --vocab_size 256000 --top_k 50 --top_p 0.9 --filter_apply_order top_k_first ``` ## πŸ” Related Issues #3389 ## πŸš€ Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### βœ… Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## πŸ§ͺ Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes ## Summary by CodeRabbit * **Chores** * Added an automatic fast path that speeds up sampling when using top‑k‑first filtering with a scalar k on large vocabularies. Sampling now selects and renormalizes a smaller candidate set for faster draws, preserves determinism/tie‑breaking and validity flags when requested, falls back to the original flow when not applicable, and introduces no public API changes. --- flashinfer/sampling.py | 114 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/flashinfer/sampling.py b/flashinfer/sampling.py index 9c219663d55..7e83574c650 100644 --- a/flashinfer/sampling.py +++ b/flashinfer/sampling.py @@ -1328,6 +1328,95 @@ def min_p_sampling_from_probs( ) +# Gating thresholds for the "top_k_first" fast path (parallel top-k, then top-p over only +# the k survivors). For a modest scalar top_k this is far cheaper than masking/renorming the +# full vocab and running rejection sampling across it: the old path's expensive steps (a +# full-vocab softmax for the logits entry point, and a single-CTA full-vocab top-p rejection +# for both) shrink to k-element work, while top-k selection costs about the same. +# +# The win only holds when (a) the vocab is large enough that the avoided full-vocab work +# outweighs the top-k selection cost, AND (b) k is small enough that the survivors stay +# cheap. Outside these thresholds we fall back to the original kernels. +# Thresholds were empirically determined. +_TOP_K_FIRST_FAST_PATH_MAX_K = 256 +_TOP_K_FIRST_FAST_PATH_MIN_VOCAB = 65536 + + +def _top_k_first_fast_path_applicable( + x: torch.Tensor, + top_k: Union[torch.Tensor, int], + indices: Optional[torch.Tensor], +) -> bool: + return ( + indices is None + and isinstance(top_k, int) + and 0 < top_k <= _TOP_K_FIRST_FAST_PATH_MAX_K + and x.size(-1) >= _TOP_K_FIRST_FAST_PATH_MIN_VOCAB + and top_k < x.size(-1) + ) + + +def _top_k_first_fast_path( + x: torch.Tensor, + top_k: int, + top_p: Union[torch.Tensor, float], + *, + from_logits: bool, + deterministic: bool, + generator: Optional[torch.Generator], + check_nan: bool, + seed: Optional[Union[int, torch.Tensor]], + offset: Optional[Union[int, torch.Tensor]], + return_valid: bool = False, +): + """Shared "top_k_first" fast path for both the logits and probs entry points. + + Selects the top-k entries with the parallel radix/cluster top-k kernel, then runs + top-p sampling over only those k entries. ``sorted=True`` gives an identical, + deterministic ordering for both logits and probs inputs, so the two entry points + reduce to the same ``probs_k`` and stay sample-aligned. This is distribution-equivalent + to the masked full-vocab path (validated TV ~0.01) but far cheaper at small batch. + """ + # Local import avoids a module-level cycle between sampling and topk. + from .topk import top_k as _radix_top_k + + # deterministic=True makes top-k reproducible (its radix deterministic-collect path is + # stable even at ties). We do not enforce a tie break that requires 128KB smem/block. + values, gathered_indices = _radix_top_k( + x, top_k, sorted=True, deterministic=deterministic + ) + values = values.float() + if from_logits: + # softmax over the k retained logits == top-k-masked softmax over the full vocab. + probs_k = torch.softmax(values, dim=-1) + else: + # renormalizing the k retained probabilities == top_k_renorm over the full vocab. + probs_k = values / values.sum(dim=-1, keepdim=True) + result = top_p_sampling_from_probs( + probs_k, + top_p, + None, + deterministic, + check_nan=check_nan, + generator=generator, + seed=seed, + offset=offset, + return_valid=return_valid, + ) + + def _map(local): + return ( + gathered_indices.gather(1, local.view(-1, 1).long()) + .squeeze(1) + .to(torch.int32) + ) + + if return_valid: + local, valid = result + return _map(local), valid + return _map(result) + + @flashinfer_api(trace=top_k_top_p_sampling_from_logits_trace) def top_k_top_p_sampling_from_logits( logits: torch.Tensor, @@ -1443,6 +1532,18 @@ def top_k_top_p_sampling_from_logits( top_p_sampling_from_probs """ if filter_apply_order == "top_k_first": + if _top_k_first_fast_path_applicable(logits, top_k, indices): + return _top_k_first_fast_path( + logits, + top_k, + top_p, + from_logits=True, + deterministic=deterministic, + generator=generator, + check_nan=check_nan, + seed=seed, + offset=offset, + ) masked_logits = top_k_mask_logits(logits, top_k) probs = torch.softmax(masked_logits, dim=-1) return top_p_sampling_from_probs( @@ -1593,6 +1694,19 @@ def top_k_top_p_sampling_from_probs( top_k_mask_logits """ if filter_apply_order == "top_k_first": + if _top_k_first_fast_path_applicable(probs, top_k, indices): + return _top_k_first_fast_path( + probs, + top_k, + top_p, + from_logits=False, + deterministic=deterministic, + generator=generator, + check_nan=check_nan, + seed=seed, + offset=offset, + return_valid=return_valid, + ) renorm_probs = top_k_renorm_probs(probs, top_k) return top_p_sampling_from_probs( renorm_probs, From c4eef413d0e456100527f2675ebb228c69e7f2c5 Mon Sep 17 00:00:00 2001 From: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> Date: Wed, 10 Jun 2026 09:44:35 +0800 Subject: [PATCH 05/13] Use an environment variable to control the number of reserved SMs for overlapping in TRT-LLM fused MoE (#3483) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## πŸ“Œ Description The current code hardcodes the number of reserved SMs for overlapping in TRT-LLM fused MoE to be 8. Because sometimes better overlapping can be achieved by using other numbers of reserved SMs for overlapping, this PR modifies the code to use an environment variable `FLASHINFER_TRTLLM_MOE_OVERLAP_RESERVED_SMS` to control the number of reserved SMs. If this environment variable is not set, the current value 8 is used. ## πŸ” Related Issues ## πŸš€ Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### βœ… Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## πŸ§ͺ Tests - [ ] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes ## Summary by CodeRabbit * **Performance Improvements** * Cooperative-kernel resource management now uses a validated, configurable allocation instead of a fixed constant, improving GPU utilization and flexibility; the chosen allocation is logged when the cooperative path is used. * **New Features** * Added an environment variable to adjust reserved GPU resources at runtime and enable tuning without code changes. --------- Signed-off-by: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> --- .../trtllm_fused_moe_routing_common.cu | 6 +- .../trtllm_fused_moe_routing_custom.cu | 5 +- .../trtllm_fused_moe_routing_deepseek.cu | 40 ++++++------ .../trtllm/fused_moe/RoutingKernel.cuh | 64 ++++++++++++++++++- 4 files changed, 92 insertions(+), 23 deletions(-) diff --git a/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu b/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu index cd6c84ad269..c0677c9dd9b 100644 --- a/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu +++ b/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu @@ -97,13 +97,14 @@ void runPostTopKPipeline(DataType const& data, void* stream) { (data.mPtrPermutedIdxSize != nullptr) && (data.mPtrExpertCounts != nullptr); bool useCoop = false; + CoopLaunchSMCounts coopLaunchSMCounts{0, 0}; int numBlocksCoop = 0; if (canUseCoop) { // Number of blocks we can use in the cooperative kernel static int const smCount = tensorrt_llm::common::getMultiProcessorCount(); - // WAR: Reserve 8 SMs for overlapping kernels. - numBlocksCoop = smCount - kReservedSMsForOverlapping; + coopLaunchSMCounts = getCoopLaunchSMCounts(smCount); + numBlocksCoop = coopLaunchSMCounts.moeSms; // Maximum number of tokens supported by the kernel using a cooperative launch. // The number of blocks must be: // >= ⌈(numTokens * topK) / (MaxExpandedIdxPerThread * NumThreads)βŒ‰ @@ -115,6 +116,7 @@ void runPostTopKPipeline(DataType const& data, void* stream) { if (useCoop) { // Coop path: cooperative launch fuses histogram + offsets (more efficient). // The coop kernel atomicAdds to mPtrExpertCounts, so we must zero it first. + logCoopLaunchSMCounts(coopLaunchSMCounts); routingCustom::launchInitExpertCounts(customData, numThreadsHist, stream); routingCustom::launchCoopKernel(customData, numBlocksCoop, numThreadsHist, stream); } else { diff --git a/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu b/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu index fe656c8747a..7756ab71118 100644 --- a/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu +++ b/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu @@ -1346,16 +1346,19 @@ void run(Data const& data, void* stream) { bool const canUseCoop = (smMajor >= 9) && (data.mNumExperts <= 1024) && (data.mPtrPermutedIdxSize != nullptr); bool useCoop = false; + CoopLaunchSMCounts coopLaunchSMCounts{0, 0}; int numBlocksCoop = 0; if (canUseCoop) { static int const smCount = tensorrt_llm::common::getMultiProcessorCount(); - numBlocksCoop = smCount - kReservedSMsForOverlapping; + coopLaunchSMCounts = getCoopLaunchSMCounts(smCount); + numBlocksCoop = coopLaunchSMCounts.moeSms; int const maxTokensCoop = (numBlocksCoop * numThreadsHist * 64) / data.mTopK; useCoop = (data.mNumTokens <= maxTokensCoop); } if (useCoop) { + logCoopLaunchSMCounts(coopLaunchSMCounts); launchInitExpertCounts(mutableData, numThreadsHist, stream); launchCoopKernel(mutableData, numBlocksCoop, numThreadsHist, stream); } else { diff --git a/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu b/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu index 62fc25b8f1c..34a643878d9 100644 --- a/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu +++ b/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu @@ -544,27 +544,31 @@ void run(Data& data, void* stream) { data.mPtrExpertCounts = nullptr; } - static int const smCount = tensorrt_llm::common::getMultiProcessorCount(); - int const numBlocksCoop = smCount - 8; - int const maxTokensCoop = (numBlocksCoop * numThreadsHist * 64) / data.mTopK; - if (useSingleCluster) { launchClusterKernel(data, numThreadsHist, stream); - } else if (data.mNumTokens <= maxTokensCoop) { - launchCoopKernel(data, numBlocksCoop, numThreadsHist, stream); } else { - const int32_t expandedIdxSize = data.mNumTokens * data.mTopK; - const int32_t histogramEltsPerBlock = 8 * numThreadsHist; - const int32_t offsetEltsPerBlock = NumEltsPerOffsetTilePerThread * numThreadsHist; - const int32_t maxNumBlocks = 1024; - - int const numBlocksHistogram = std::min( - (expandedIdxSize + histogramEltsPerBlock - 1) / histogramEltsPerBlock, maxNumBlocks); - int const numBlocksOffsets = - std::min((expandedIdxSize + offsetEltsPerBlock - 1) / offsetEltsPerBlock, maxNumBlocks); - - launchHistogramKernel(data, numBlocksHistogram, numThreadsHist, stream); - launchOffsetsKernel(data, numBlocksOffsets, numThreadsHist, stream); + static int const smCount = tensorrt_llm::common::getMultiProcessorCount(); + CoopLaunchSMCounts const coopLaunchSMCounts = getCoopLaunchSMCounts(smCount); + int const numBlocksCoop = coopLaunchSMCounts.moeSms; + int const maxTokensCoop = (numBlocksCoop * numThreadsHist * 64) / data.mTopK; + + if (data.mNumTokens <= maxTokensCoop) { + logCoopLaunchSMCounts(coopLaunchSMCounts); + launchCoopKernel(data, numBlocksCoop, numThreadsHist, stream); + } else { + const int32_t expandedIdxSize = data.mNumTokens * data.mTopK; + const int32_t histogramEltsPerBlock = 8 * numThreadsHist; + const int32_t offsetEltsPerBlock = NumEltsPerOffsetTilePerThread * numThreadsHist; + const int32_t maxNumBlocks = 1024; + + int const numBlocksHistogram = std::min( + (expandedIdxSize + histogramEltsPerBlock - 1) / histogramEltsPerBlock, maxNumBlocks); + int const numBlocksOffsets = + std::min((expandedIdxSize + offsetEltsPerBlock - 1) / offsetEltsPerBlock, maxNumBlocks); + + launchHistogramKernel(data, numBlocksHistogram, numThreadsHist, stream); + launchOffsetsKernel(data, numBlocksOffsets, numThreadsHist, stream); + } } } } diff --git a/include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh b/include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh index 48d572abafc..6ee03c879f2 100644 --- a/include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh +++ b/include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh @@ -18,7 +18,10 @@ #include #include #include +#include +#include +#include #include #include #include @@ -42,8 +45,65 @@ static constexpr int WarpSize = 32; static constexpr int NumBlocksPerCluster = 8; // Performance tuning knob. static constexpr int NumEltsPerOffsetTilePerThread = 8; -// Number of SMs to reserve for overlapping kernels when using cooperative launch. -static constexpr int kReservedSMsForOverlapping = 8; +// Default number of SMs to leave available for overlapping kernels when using cooperative launch. +static constexpr int kDefaultReservedSMsForOverlapping = 8; +static constexpr char kReservedSMsForOverlappingEnv[] = + "FLASHINFER_TRTLLM_MOE_OVERLAP_RESERVED_SMS"; + +// Parsed reserved-SM value and whether it came from the environment. +struct ReservedSMsForOverlappingConfig { + // Number of SMs to reserve for overlapping kernels. + long reservedSms; + // True when reservedSms was read from FLASHINFER_TRTLLM_MOE_OVERLAP_RESERVED_SMS. + bool isSet; +}; + +struct CoopLaunchSMCounts { + int moeSms; + int reservedSms; +}; + +// Return the reserved-SM configuration for TRT-LLM fused MoE overlap. +// The value is read from FLASHINFER_TRTLLM_MOE_OVERLAP_RESERVED_SMS once per +// process. When the variable is not set, the default reserved-SM count is used +// and isSet is false so error messages can identify the source of the value. +inline ReservedSMsForOverlappingConfig getReservedSMsForOverlappingConfig() { + static ReservedSMsForOverlappingConfig const config = [] { + char const* env = std::getenv(kReservedSMsForOverlappingEnv); + if (env == nullptr) { + return ReservedSMsForOverlappingConfig{kDefaultReservedSMsForOverlapping, false}; + } + char* end = nullptr; + long const value = std::strtol(env, &end, 10); + FLASHINFER_CHECK(end != env && *end == '\0', kReservedSMsForOverlappingEnv, + " must be an integer, got ", env); + return ReservedSMsForOverlappingConfig{value, true}; + }(); + return config; +} + +// Return the cooperative-launch SM allocation after reserving SMs. +// Validate the effective reserved-SM count against the runtime SM count before +// subtraction so the number of SMs used by MoE is always positive. +inline CoopLaunchSMCounts getCoopLaunchSMCounts(int smCount) { + ReservedSMsForOverlappingConfig const config = getReservedSMsForOverlappingConfig(); + char const* source = config.isSet ? kReservedSMsForOverlappingEnv : "default reserved SM count"; + FLASHINFER_CHECK(config.reservedSms >= 0 && config.reservedSms < smCount, source, + " must satisfy 0 <= value < SM count (", smCount, "), got ", config.reservedSms); + int const reservedSms = static_cast(config.reservedSms); + return CoopLaunchSMCounts{smCount - reservedSms, reservedSms}; +} + +inline void logCoopLaunchSMCounts(CoopLaunchSMCounts const& counts) { + static bool logged = false; + if (!logged) { + logged = true; + FLASHINFER_LOG_INFO( + "TRT-LLM fused MoE cooperative launch SM allocation: {} SMs used for MoE, {} SMs " + "reserved for overlapping kernels (total SMs: {})", + counts.moeSms, counts.reservedSms, counts.moeSms + counts.reservedSms); + } +} //////////////////////////////////////////////////////////////////////////////////////////////////// From 2aa1d49cf140d73ccdd3761051c5f2944406cb83 Mon Sep 17 00:00:00 2001 From: nvjullin Date: Wed, 10 Jun 2026 13:37:03 +0800 Subject: [PATCH 06/13] Remove excessive CuteDslMoEWrapper memory allocation (#3404) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## πŸ“Œ Description Resolves https://github.com/flashinfer-ai/flashinfer/issues/3308. Previously, `max_num_tokens` were only checked in cuda graph path, but now deprecated and unused. Dropped tests that test for pre-allocated buffers. ## πŸ” Related Issues ## πŸš€ 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 ## Summary by CodeRabbit * **API Changes** * `max_num_tokens` is now optional/deprecated (ignored at runtime); existing calls remain compatible. * **Performance Improvements** * Persistent CUDA stream/event handling improves CUDA-graph compatibility and enables better async-memset overlap, reducing memory overhead. * **Tests** * Tests updated to validate routing/buffer write-completeness via the core execution path using poisoned-buffer checks. * **Documentation** * Docstrings updated to reflect the new persistent CUDA resource model and deprecated `max_num_tokens`. --- flashinfer/fused_moe/cute_dsl/fused_moe.py | 186 ++-------- flashinfer/fused_moe/cute_dsl/tuner.py | 6 +- tests/moe/test_cute_dsl_fused_moe.py | 404 ++------------------- 3 files changed, 75 insertions(+), 521 deletions(-) diff --git a/flashinfer/fused_moe/cute_dsl/fused_moe.py b/flashinfer/fused_moe/cute_dsl/fused_moe.py index 88477e86d8b..f03cf2d33f5 100644 --- a/flashinfer/fused_moe/cute_dsl/fused_moe.py +++ b/flashinfer/fused_moe/cute_dsl/fused_moe.py @@ -27,7 +27,8 @@ Best for: simple use cases, experimenting, auto-tuning. 2. **Wrapper API** (`CuteDslMoEWrapper`): - Class-based API with pre-allocated buffers for CUDA graph compatibility. + Class-based API that holds persistent CUDA stream/event resources for + async-memset overlap and CUDA graph compatibility. Best for: production inference with CUDA graphs, fine-grained control. Both APIs share the same core implementation and support auto-tuning. @@ -60,11 +61,9 @@ cute_dsl_fused_moe_nvfp4_trace, cute_dsl_moe_wrapper_run_trace, ) -from ...autotuner import AutoTuner, is_in_profile_measurement +from ...autotuner import AutoTuner from ...utils import supported_compute_capability from .moe_utils import ( - allocate_moe_sort_buffers, - get_max_num_permuted_tokens, moe_output_memset_inplace, moe_sort, ) @@ -77,7 +76,6 @@ from .tuner import ( ALL_MOE_TACTICS, CuteDslFusedMoENvfp4Runner, - VALID_TILE_SIZES, ) @@ -315,9 +313,10 @@ def _moe_core_impl( class CuteDslMoEWrapper: """Wrapper class for CuteDSL MoE with CUDA graph and auto-tuning support. - This wrapper pre-allocates all necessary buffers when `use_cuda_graph=True`, - enabling CUDA graph capture and replay. It also supports auto-tuning via - the `tactic` parameter or by calling inside `autotune()` context. + With `use_cuda_graph=True`, the wrapper creates persistent CUDA stream + and event resources outside graph capture, enabling async-memset / GEMM1 + overlap during capture and replay. Auto-tuning is supported via the `tactic` + parameter or `autotune()` context. Supported architectures: SM100, SM103. @@ -326,14 +325,16 @@ class CuteDslMoEWrapper: top_k: Number of experts per token. hidden_size: Hidden dimension size. intermediate_size: Intermediate dimension size. - use_cuda_graph: Whether to pre-allocate buffers for CUDA graph. - max_num_tokens: Maximum tokens (only used with use_cuda_graph=True). + use_cuda_graph: Whether the wrapper holds persistent stream/event + resources for CUDA graph capture. + max_num_tokens: Deprecated; accepted for backwards compatibility + but ignored. Example (CUDA Graph): >>> moe = CuteDslMoEWrapper( ... num_experts=256, top_k=8, ... hidden_size=7168, intermediate_size=2048, - ... use_cuda_graph=True, max_num_tokens=4096, + ... use_cuda_graph=True, ... ) >>> # Warmup >>> for _ in range(3): @@ -361,7 +362,7 @@ def __init__( hidden_size: int, intermediate_size: int, use_cuda_graph: bool = False, - max_num_tokens: int = 4096, + max_num_tokens: Optional[int] = None, num_local_experts: Optional[int] = None, local_expert_offset: int = 0, tile_size: int = 128, @@ -383,12 +384,11 @@ def __init__( intermediate_size : int Intermediate dimension size (after SwiGLU reduction). use_cuda_graph : bool - If ``True``, pre-allocate workspace buffers sized for - ``max_num_tokens`` so the wrapper can be captured into a CUDA - graph. Defaults to ``False``. - max_num_tokens : int - Maximum batch size, used when ``use_cuda_graph=True``. Defaults - to ``4096``. + Create persistent CUDA stream/events for async-memset overlap. + Required for CUDA graph capture, since streams and events must be + created outside graph capture. Defaults to ``False``. + max_num_tokens : Optional[int] + Deprecated; accepted for backwards compatibility but ignored. num_local_experts : Optional[int] Local experts for expert parallelism. Defaults to ``num_experts``. @@ -402,8 +402,7 @@ def __init__( output_dtype : torch.dtype Output dtype. Defaults to ``torch.bfloat16``. device : str - Device on which to allocate workspace buffers. Defaults to - ``"cuda"``. + Device on which to allocate buffers. Defaults to ``"cuda"``. enable_pdl : bool Enable Programmatic Dependent Launch. Defaults to ``True``. """ @@ -412,7 +411,6 @@ def __init__( self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.use_cuda_graph = use_cuda_graph - self.max_num_tokens = max_num_tokens self.num_local_experts = num_local_experts or num_experts self.local_expert_offset = local_expert_offset self.tile_size = tile_size @@ -421,11 +419,10 @@ def __init__( self.device = device self.enable_pdl = enable_pdl - # Pre-allocated buffers - self._moe_sort_buffers: Optional[Dict[str, torch.Tensor]] = None - self._gemm1_output: Optional[torch.Tensor] = None - self._gemm1_output_scale: Optional[torch.Tensor] = None - self._moe_output: Optional[torch.Tensor] = None + # Persistent CUDA resources for async-memset / GEMM1 overlap. These + # are created outside graph capture (so they can be reused inside it) + # when ``use_cuda_graph=True``. When None, ``_moe_core_impl`` falls + # back to module-level resources via ``_get_cuda_graph_resources``. self._aux_stream: Optional[torch.cuda.Stream] = None self._main_event: Optional[torch.cuda.Event] = None self._memset_event: Optional[torch.cuda.Event] = None @@ -455,85 +452,9 @@ def _forward_with_tactic_weak(*args, **kwargs): ) if use_cuda_graph: - self._allocate_buffers() - - def _allocate_buffers(self) -> None: - """Pre-allocate all buffers for CUDA graph compatibility. - - Buffers are sized to fit *any* ``tile_size in VALID_TILE_SIZES``, - not just ``self.tile_size``. Two distinct buffer-shape concerns: - - - ``max_num_permuted_tokens`` is monotonically *increasing* in - ``tile_size`` (the ``(tile - 1) * num_local_experts`` padding - term grows faster than ``max_num_tiles`` shrinks), so - permuted-token-indexed buffers (``_gemm1_output``, - ``_gemm1_output_scale``, - ``out_permuted_idx_to_expanded_idx``) must be sized using - ``max(VALID_TILE_SIZES)``. - - ``max_num_tiles`` is monotonically *decreasing* in - ``tile_size``, so the tile-count-indexed moe_sort buffers - (``out_tile_idx_to_expert_idx``, - ``out_tile_idx_to_mn_limit``) must be sized using - ``min(VALID_TILE_SIZES)``. - - Sizing this way means the ``use_prealloc`` gate at - ``_forward_with_tactic`` doesn't need a ``tile_size == - self.tile_size`` check at runtime: whichever tactic the - autotuner picked, the prealloc fits. This preserves the - wrapper's CUDA-graph contract (``run()`` is graph-safe with - ``use_cuda_graph=True``) regardless of which ``tile_size`` the - autotuner ends up choosing. - """ - smallest_tile = min(VALID_TILE_SIZES) - largest_tile = max(VALID_TILE_SIZES) - max_permuted_across_tiles = get_max_num_permuted_tokens( - self.max_num_tokens, self.top_k, self.num_local_experts, largest_tile - ) - - # moe_sort buffers β€” allocate using smallest_tile so the - # tile-count-indexed buffers (out_tile_idx_to_expert_idx, - # out_tile_idx_to_mn_limit) are large enough for any tile_size, - # then override out_permuted_idx_to_expanded_idx (which scales - # with tile_size in the opposite direction) to fit the largest - # tile_size's max_num_permuted_tokens. - self._moe_sort_buffers = allocate_moe_sort_buffers( - num_tokens=self.max_num_tokens, - num_experts=self.num_experts, - top_k=self.top_k, - num_local_experts=self.num_local_experts, - tile_tokens_dim=smallest_tile, - device=self.device, - ) - self._moe_sort_buffers["out_permuted_idx_to_expanded_idx"] = torch.empty( - (max_permuted_across_tiles,), dtype=torch.int32, device=self.device - ) - - # GEMM1 output (FP4 quantized) - self._gemm1_output = torch.empty( - (max_permuted_across_tiles, self.intermediate_size // 2), - dtype=torch.uint8, - device=self.device, - ) - - # GEMM1 output scale - scale_size = max_permuted_across_tiles * ( - self.intermediate_size // self.sf_vec_size - ) - self._gemm1_output_scale = torch.empty( - (scale_size,), dtype=torch.uint8, device=self.device - ) - - # Final output - self._moe_output = torch.empty( - (self.max_num_tokens, self.hidden_size), - dtype=self.output_dtype, - device=self.device, - ) - - # CUDA resources - self._aux_stream = torch.cuda.Stream(device=self.device) - self._main_event = torch.cuda.Event() - self._memset_event = torch.cuda.Event() + self._aux_stream = torch.cuda.Stream(device=self.device) + self._main_event = torch.cuda.Event() + self._memset_event = torch.cuda.Event() def _forward_with_tactic( self, @@ -564,30 +485,6 @@ def _forward_with_tactic( **kwargs, ) -> torch.Tensor: """Forward implementation called by auto-tuner.""" - # Pre-allocated buffers are sized to fit any ``tile_size in - # VALID_TILE_SIZES`` (see ``_allocate_buffers``). Fall back to - # dynamic allocation when: - # - # - the autotuner is in its per-tactic measurement window - # (``is_in_profile_measurement()``): every probed tactic must - # see the same allocation overhead so the comparison is - # unbiased. Note: this is intentionally narrower than - # ``is_tuning_mode`` -- it excludes cache lookups, - # ``do_preparation`` calls, the final invocation after - # ``choose_one`` returns, and other threads' inference, which - # all benefit from prealloc. - # - the tactic's ``tile_size`` is somehow outside the canonical - # enumeration (defensive; should never fire for tactics from - # ``ALL_MOE_TACTICS``). - # - the batch exceeds what the buffers were sized for (e.g. - # autotuner probing larger buckets than ``max_num_tokens``). - num_tokens = x.shape[0] - use_prealloc = ( - self.use_cuda_graph - and not is_in_profile_measurement() - and tile_size in VALID_TILE_SIZES - and num_tokens <= self.max_num_tokens - ) return _moe_core_impl( x=x, x_sf=x_sf, @@ -609,13 +506,10 @@ def _forward_with_tactic( gemm1_cluster_shape_mn=gemm1_cluster_shape_mn, gemm2_mma_tiler_mn=gemm2_mma_tiler_mn, gemm2_cluster_shape_mn=gemm2_cluster_shape_mn, - moe_sort_buffers=self._moe_sort_buffers if use_prealloc else None, - gemm1_out=self._gemm1_output if use_prealloc else None, - gemm1_out_scale=self._gemm1_output_scale if use_prealloc else None, - moe_output=moe_output - if moe_output is not None - # Slice the CUDA-graph buffer to the active batch. - else (self._moe_output[: x.shape[0]] if use_prealloc else None), + moe_sort_buffers=None, + gemm1_out=None, + gemm1_out_scale=None, + moe_output=moe_output, aux_stream=self._aux_stream, main_event=self._main_event, memset_event=self._memset_event, @@ -681,21 +575,11 @@ def run( """ num_tokens = token_selected_experts.size(0) - if self.use_cuda_graph and num_tokens > self.max_num_tokens: - raise ValueError( - f"num_tokens ({num_tokens}) exceeds max_num_tokens ({self.max_num_tokens})" - ) - - # Slice the pre-allocated buffer to the active batch so that - # _moe_core_impl only zeros num_tokens rows, not max_num_tokens. - if self.use_cuda_graph: - moe_output = self._moe_output[:num_tokens] - else: - moe_output = torch.empty( - (num_tokens, self.hidden_size), - dtype=self.output_dtype, - device=x.device, - ) + moe_output = torch.empty( + (num_tokens, self.hidden_size), + dtype=self.output_dtype, + device=x.device, + ) # Use auto-tuner for tactic selection tuner = AutoTuner.get() diff --git a/flashinfer/fused_moe/cute_dsl/tuner.py b/flashinfer/fused_moe/cute_dsl/tuner.py index ff2342bb298..84ecc86c7ca 100644 --- a/flashinfer/fused_moe/cute_dsl/tuner.py +++ b/flashinfer/fused_moe/cute_dsl/tuner.py @@ -154,11 +154,7 @@ def get_gemm2_valid_tactics(tile_size: int) -> List[Tuple]: # Canonical list of tile_sizes the autotuner is allowed to pick. Used by -# ``get_moe_valid_tactics`` for tactic enumeration AND by -# ``CuteDslMoEWrapper`` to size its preallocated kernel-output buffers so -# every tactic in this list can reuse the prealloc, regardless of which -# tile_size the autotuner picks at runtime. Adding a new tile_size here -# automatically widens the prealloc. +# ``get_moe_valid_tactics`` for tactic enumeration. VALID_TILE_SIZES: Tuple[int, ...] = (128, 256) diff --git a/tests/moe/test_cute_dsl_fused_moe.py b/tests/moe/test_cute_dsl_fused_moe.py index c284d0e1be4..80efbc94318 100644 --- a/tests/moe/test_cute_dsl_fused_moe.py +++ b/tests/moe/test_cute_dsl_fused_moe.py @@ -1756,8 +1756,8 @@ def test_functional_with_ep(self, ep_size: int): class TestMoeSortBufferInitPoisoned: """Validate the invariant that the routing kernel writes every output entry that downstream code reads, by pre-poisoning the - wrapper's preallocated ``moe_sort`` output buffers with a sentinel - value before the first call. + ``moe_sort`` output buffers with a sentinel value before invoking + the MoE pipeline. The ``moe_sort`` wrapper in ``moe_utils.py`` allocates its output buffers via ``torch.empty(...)`` and relies on the routing kernel @@ -1778,11 +1778,11 @@ class TestMoeSortBufferInitPoisoned: written as ``-1`` by the kernel. If the kernel ever stops writing masked slots, this test catches it. - These tests use ``use_cuda_graph=True`` so the wrapper preallocates - buffers (the path where stale state from prior calls / poisoning is - actually retained between calls). The default ``use_cuda_graph=False`` - path allocates fresh buffers per call and doesn't exercise the - same scenario. + ``_moe_core_impl`` accepts an external ``moe_sort_buffers`` dict + for callers who want to manage their own routing-output buffers. + The test allocates the dict via ``allocate_moe_sort_buffers``, + fills it with the poison sentinel, and drives the full + routing+gemm pipeline through ``_moe_core_impl`` directly. """ @pytest.mark.parametrize( @@ -1806,7 +1806,7 @@ def test_wrapper_with_poisoned_moe_sort_buffers( self, ep_size: int, num_tokens: int ): """Pre-poison all six moe_sort output buffers with a sentinel - before the wrapper's first call; verify output is well-formed + before invoking the MoE pipeline; verify output is well-formed and (at low N) matches the eager reference within tolerance. The high-N case (``num_tokens > 1024``) skips the @@ -1817,7 +1817,10 @@ def test_wrapper_with_poisoned_moe_sort_buffers( what the poisoning-detection logic actually relies on; the reference comparison is supplementary. """ - from flashinfer import CuteDslMoEWrapper + from flashinfer.fused_moe.cute_dsl.fused_moe import _moe_core_impl + from flashinfer.fused_moe.cute_dsl.moe_utils import ( + allocate_moe_sort_buffers, + ) hidden_size, intermediate_size = 256, 512 num_experts, top_k = 256, 8 @@ -1833,33 +1836,29 @@ def test_wrapper_with_poisoned_moe_sort_buffers( top_k=top_k, ) - # use_cuda_graph=True so the wrapper preallocates _moe_sort_buffers - # (the path that retains stale state between calls β€” exactly what - # we want to stress here). - moe = CuteDslMoEWrapper( + # Allocate the moe_sort outputs externally so we can poison them + # before invoking the kernel. ``_moe_core_impl`` accepts a + # ``moe_sort_buffers`` dict that flows through to ``moe_sort`` via + # **kwargs, taking the place of its internal ``torch.empty`` calls. + tile_size = 128 # default tactic for _moe_core_impl + moe_sort_buffers = allocate_moe_sort_buffers( + num_tokens=num_tokens, num_experts=num_experts, top_k=top_k, - hidden_size=hidden_size, - intermediate_size=intermediate_size, num_local_experts=num_local_experts, - local_expert_offset=local_expert_offset, - use_cuda_graph=True, - max_num_tokens=num_tokens, + tile_tokens_dim=tile_size, + device="cuda", ) - # Defensive guard: if a future refactor renames or restructures - # ``_moe_sort_buffers``, the poisoning loop below would silently - # iterate over zero items and the test would pass without - # actually exercising the kernel-write invariant. Fail loudly - # in that case so the test must be updated rather than silently - # rotting. - assert ( - getattr(moe, "_moe_sort_buffers", None) is not None - and len(moe._moe_sort_buffers) > 0 - ), ( - "Wrapper no longer exposes a non-empty ``_moe_sort_buffers`` " - "dict; the poisoning loop would be a no-op. Update this " - "test to target the new preallocation attribute." + # Defensive guard: if a future refactor renames or restructures the + # buffer dict returned by ``allocate_moe_sort_buffers``, the + # poisoning loop below would silently iterate over zero items and + # the test would pass without exercising the kernel-write + # invariant. Fail loudly in that case. + assert moe_sort_buffers and len(moe_sort_buffers) > 0, ( + "``allocate_moe_sort_buffers`` no longer returns a non-empty " + "dict; the poisoning loop would be a no-op. Update this test " + "to target the new buffer-allocation API." ) # Sentinel: a non-zero, non-(-1), out-of-valid-index-range int32. @@ -1868,10 +1867,10 @@ def test_wrapper_with_poisoned_moe_sort_buffers( # atomic-add will scatter into wildly wrong output rows, producing # NaN/Inf or massive numerical divergence. POISON = 0x7FFFFFFE - for buf in moe._moe_sort_buffers.values(): + for buf in moe_sort_buffers.values(): buf.fill_(POISON) - result = moe.run( + result = _moe_core_impl( x=tensors["x"], x_sf=tensors["x_sf"], token_selected_experts=tensors["token_selected_experts"], @@ -1883,6 +1882,13 @@ def test_wrapper_with_poisoned_moe_sort_buffers( w2_weight=tensors["w2_weight"], w2_weight_sf=tensors["w2_weight_sf"], w2_alpha=tensors["w2_alpha"], + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + moe_sort_buffers=moe_sort_buffers, + output_dtype=torch.bfloat16, ) assert result.shape == (num_tokens, hidden_size) @@ -2086,338 +2092,6 @@ def test_all_tactics_accuracy( ) -# ============================================================================= -# Test Class: CuteDslMoEWrapper prealloc static invariants (no GPU required) -# ============================================================================= - - -@cute_dsl_available -class TestPreallocStaticInvariants: - """No-GPU structural invariants on ``VALID_TILE_SIZES``. - - The empirical buffer-shape and prealloc-gate behavior is covered - by ``TestPreallocBuffersIntegration`` and - ``TestPreallocGateUnderTuning`` (GPU-required). This class catches - the orthogonal failure mode where ``VALID_TILE_SIZES`` is - accidentally reduced to a single entry β€” in that case the GPU - integration tests pass trivially (no max/min divergence in - ``_allocate_buffers``, only one tile_size to gate-check) and the - bias-prevention silently disappears. - """ - - def test_valid_tile_sizes_has_multiple_entries(self): - """``VALID_TILE_SIZES`` must enumerate more than one tile_size. - With a single entry, the bias-prevention is moot β€” the - autotuner only ever profiles one tile_size class, defeating - the whole point of widening the prealloc. - """ - from flashinfer.fused_moe.cute_dsl.tuner import VALID_TILE_SIZES - - assert len(VALID_TILE_SIZES) >= 2, ( - f"VALID_TILE_SIZES has only {len(VALID_TILE_SIZES)} entry; " - f"need >= 2 for the prealloc-bias fix to be meaningful." - ) - assert all(isinstance(t, int) and t > 0 for t in VALID_TILE_SIZES), ( - f"VALID_TILE_SIZES entries must be positive ints; got {VALID_TILE_SIZES}" - ) - - -# ============================================================================= -# Test Class: CuteDslMoEWrapper prealloc-buffer integration (GPU required) -# ============================================================================= - - -@cute_dsl_available -@sm100_required -class TestPreallocBuffersIntegration: - """Verify the wrapper's prealloc'd buffers fit the workload at - *every* ``tile_size in VALID_TILE_SIZES``, not just the - constructor-time ``self.tile_size``. - - Load-bearing property: when the autotuner picks a tactic with - ``tile_size != self.tile_size`` (the common case at large N where - ``tile_size=256`` wins on intrinsic kernel time), the wrapper's - ``use_prealloc`` gate still resolves True and inference uses the - prealloc. This requires the buffers to fit the *largest* possible - workload across all valid tile_sizes; if they were sized only for - ``self.tile_size``, picking a different tactic at runtime would - OOB-write the prealloc -- forcing the gate to fall through to - per-call ``torch.empty()`` calls, which violates the wrapper's - CUDA-graph contract. - """ - - def test_prealloc_buffers_fit_all_valid_tile_sizes(self): - from flashinfer import CuteDslMoEWrapper - from flashinfer.fused_moe.cute_dsl.moe_utils import ( - get_max_num_permuted_tokens, - get_max_num_tiles, - ) - from flashinfer.fused_moe.cute_dsl.tuner import VALID_TILE_SIZES - - wrapper = CuteDslMoEWrapper( - num_experts=256, - top_k=8, - hidden_size=256, - intermediate_size=512, - num_local_experts=256, - local_expert_offset=0, - use_cuda_graph=True, - max_num_tokens=256, - ) - - gemm1_capacity = wrapper._gemm1_output.shape[0] - gemm1_scale_capacity = wrapper._gemm1_output_scale.shape[0] - permuted_idx_capacity = wrapper._moe_sort_buffers[ - "out_permuted_idx_to_expanded_idx" - ].shape[0] - tile_expert_capacity = wrapper._moe_sort_buffers[ - "out_tile_idx_to_expert_idx" - ].shape[0] - tile_mn_limit_capacity = wrapper._moe_sort_buffers[ - "out_tile_idx_to_mn_limit" - ].shape[0] - - # Scale buffer is sized in scale-factor elements (one per - # (permuted_token, scale_vec_group) pair), not in permuted - # tokens directly. - scale_factor_per_token = wrapper.intermediate_size // wrapper.sf_vec_size - - for tile_size in VALID_TILE_SIZES: - required_permuted = get_max_num_permuted_tokens( - wrapper.max_num_tokens, - wrapper.top_k, - wrapper.num_local_experts, - tile_size, - ) - required_scale_size = required_permuted * scale_factor_per_token - required_tiles = get_max_num_tiles( - wrapper.max_num_tokens, - wrapper.top_k, - wrapper.num_local_experts, - tile_size, - ) - - assert gemm1_capacity >= required_permuted, ( - f"_gemm1_output rows ({gemm1_capacity}) < required " - f"({required_permuted}) at tile_size={tile_size}" - ) - assert gemm1_scale_capacity >= required_scale_size, ( - f"_gemm1_output_scale capacity ({gemm1_scale_capacity}) " - f"< required ({required_scale_size} = {required_permuted} " - f"permuted * {scale_factor_per_token} scales/token) at " - f"tile_size={tile_size}" - ) - assert permuted_idx_capacity >= required_permuted, ( - f"out_permuted_idx_to_expanded_idx capacity " - f"({permuted_idx_capacity}) < required ({required_permuted}) " - f"at tile_size={tile_size}" - ) - assert tile_expert_capacity >= required_tiles, ( - f"out_tile_idx_to_expert_idx capacity " - f"({tile_expert_capacity}) < required ({required_tiles}) " - f"at tile_size={tile_size}" - ) - assert tile_mn_limit_capacity >= required_tiles, ( - f"out_tile_idx_to_mn_limit capacity " - f"({tile_mn_limit_capacity}) < required ({required_tiles}) " - f"at tile_size={tile_size}" - ) - - -# ============================================================================= -# Test Class: CuteDslMoEWrapper autotune-profiling prealloc gate (GPU required) -# ============================================================================= - - -@cute_dsl_available -@sm100_required -class TestPreallocGateUnderTuning: - """Validate that ``_forward_with_tactic``'s ``use_prealloc`` gate - is on during normal inference (any valid tile_size) but off during - the autotuner's per-tactic measurement window. - - Behavioral contract: - - 1. **Inside the per-tactic measurement window** (i.e. while - ``is_in_profile_measurement()`` is True): the gate must return - ``False`` for every tactic, regardless of whether the tactic's - ``tile_size`` matches ``self.tile_size``. All tactics see the - same per-call ``torch.empty()`` allocation overhead and the - autotuner's tactic comparison is unbiased. - - 2. **Inside ``autotune(True)`` but outside the measurement window** - (cache lookups, ``do_preparation`` calls, the post-``choose_one`` - final invocation, concurrent threads): the gate must use - prealloc for *any* ``tile_size in VALID_TILE_SIZES``. This is - the property that ``is_in_profile_measurement()`` adds over the - broader ``is_tuning_mode`` flag: the gate doesn't leak into - these adjacent code paths. - - 3. **Outside any tuning context** (plain inference): same as case - 2 β€” prealloc for any ``tile_size in VALID_TILE_SIZES``. This is - the property that the expanded ``_allocate_buffers`` adds: the - gate doesn't depend on ``tile_size == self.tile_size``, so - whichever tactic the autotuner picks, the wrapper's CUDA-graph - prealloc is still used and the wrapper's graph-safety contract - is preserved. - - Implementation: monkey-patch the module-level ``_moe_core_impl`` - to capture the ``moe_sort_buffers`` argument without launching - kernels, then call ``_forward_with_tactic`` from each of the three - contexts Γ— {``self.tile_size``, other valid tile_size} - configurations. - """ - - def test_gate_decouples_self_tile_size_only_during_measurement_window( - self, monkeypatch - ): - from flashinfer import CuteDslMoEWrapper, autotune - from flashinfer.autotuner import _profile_measurement_scope - from flashinfer.fused_moe.cute_dsl import fused_moe as fused_moe_module - from flashinfer.fused_moe.cute_dsl.tuner import VALID_TILE_SIZES - - wrapper = CuteDslMoEWrapper( - num_experts=256, - top_k=8, - hidden_size=256, - intermediate_size=512, - num_local_experts=256, - local_expert_offset=0, - use_cuda_graph=True, - max_num_tokens=128, - ) - - # (context, tile_size) -> bool (prealloc'd buffers passed) - captured: dict = {} - # The mode under which the next call is made; updated by the - # caller before each ``call(tile_size, mode)`` so the mock can - # tag the captured row correctly. - current_mode = {"name": "inference"} - - def mock_moe_core_impl(*args, **kwargs): - captured[(current_mode["name"], kwargs["tile_size"])] = ( - kwargs["moe_sort_buffers"] is wrapper._moe_sort_buffers - ) - n = args[0].shape[0] if args else kwargs["x"].shape[0] - return torch.zeros( - (n, wrapper.hidden_size), dtype=torch.bfloat16, device="cuda" - ) - - monkeypatch.setattr(fused_moe_module, "_moe_core_impl", mock_moe_core_impl) - - # Build minimal placeholder tensors: _forward_with_tactic only - # reads x.shape[0]; everything else is passed through to the - # (mocked) inner function untouched. - n = 64 # < max_num_tokens=128 so the batch check passes - x = torch.empty((n, wrapper.hidden_size // 2), dtype=torch.uint8, device="cuda") - x_sf = torch.empty((n, 1), dtype=torch.uint8, device="cuda") - token_selected_experts = torch.zeros( - (n, wrapper.top_k), dtype=torch.int32, device="cuda" - ) - token_final_scales = torch.zeros( - (n, wrapper.top_k), dtype=torch.float32, device="cuda" - ) - dummy_w = torch.empty((1,), dtype=torch.uint8, device="cuda") - dummy_alpha = torch.empty((1,), dtype=torch.float32, device="cuda") - - def call(tile_size: int) -> None: - wrapper._forward_with_tactic( - x=x, - x_sf=x_sf, - token_selected_experts=token_selected_experts, - token_final_scales=token_final_scales, - w1_weight=dummy_w, - w1_weight_sf=dummy_w, - w1_alpha=dummy_alpha, - fc2_input_scale=dummy_alpha, - w2_weight=dummy_w, - w2_weight_sf=dummy_w, - w2_alpha=dummy_alpha, - num_experts=wrapper.num_experts, - top_k=wrapper.top_k, - num_local_experts=wrapper.num_local_experts, - tile_size=tile_size, - ) - - matching = wrapper.tile_size # the tile_size the prealloc was sized for - # Exercise every tile_size in VALID_TILE_SIZES so adding a new - # entry doesn't silently leave the gate untested for that tile. - others = [t for t in VALID_TILE_SIZES if t != matching] - assert others, ( - f"Test requires >= 2 distinct VALID_TILE_SIZES entries; " - f"got {VALID_TILE_SIZES}" - ) - all_tiles = (matching, *others) - - # Context 1: inside autotune(True) AND inside the measurement - # window β€” what _profile_single_kernel does for each tactic - # invocation. The gate must skip prealloc for every tactic. - with autotune(True): - with _profile_measurement_scope(): - current_mode["name"] = "measurement" - for tile_size in all_tiles: - call(tile_size) - - # Context 2: inside autotune(True) but OUTSIDE the - # measurement window β€” analogous to a cache hit, the - # do_preparation call, or the runner invocation immediately - # after choose_one returns. The gate should behave like - # plain inference here. - current_mode["name"] = "in_tuning_context_outside_measurement" - for tile_size in all_tiles: - call(tile_size) - - # Context 3: outside any tuning context β€” plain inference. - current_mode["name"] = "inference" - for tile_size in all_tiles: - call(tile_size) - - # Context 1 contract: skip prealloc unconditionally. - for tile_size in all_tiles: - assert not captured[("measurement", tile_size)], ( - f"In the per-tactic measurement window, gate passed " - f"prealloc'd buffers for tile_size={tile_size} " - f"(self.tile_size={matching}). This re-introduces the " - f"autotune-profiling bias the gate is designed to prevent." - ) - - # Context 2 contract: prealloc for ANY valid tile_size. This - # is the property that distinguishes - # ``is_in_profile_measurement()`` from the broader - # ``is_tuning_mode``: cache lookups, do_preparation calls, - # post-choose_one runs, and concurrent threads should NOT lose - # prealloc just because some other thread/operation is inside - # an ``autotune(True)`` context. Combined with the expanded - # buffer sizing in ``_allocate_buffers``, the prealloc is also - # used regardless of whether ``tile_size == self.tile_size``. - for tile_size in all_tiles: - assert captured[("in_tuning_context_outside_measurement", tile_size)], ( - f"Inside autotune(True) but outside the measurement " - f"window at tile_size={tile_size}, gate did not pass " - f"prealloc'd buffers (self.tile_size={matching}). " - f"Either the narrower is_in_profile_measurement() " - f"signal is leaking back into is_tuning_mode breadth, " - f"or the gate is incorrectly checking tile_size == " - f"self.tile_size -- both regress the wrapper's " - f"CUDA-graph contract." - ) - - # Context 3 contract: same as Context 2 -- gate uses prealloc - # for ANY valid tile_size. This preserves the wrapper's - # CUDA-graph contract regardless of which tactic the autotuner - # picks at runtime. - for tile_size in all_tiles: - assert captured[("inference", tile_size)], ( - f"In inference mode at tile_size={tile_size}, gate " - f"did not pass prealloc'd buffers " - f"(self.tile_size={matching}). The wrapper loses its " - f"CUDA-graph prealloc benefit -- with use_cuda_graph=" - f"True, captured graphs would record per-call " - f"torch.empty() calls instead of using the prealloc, " - f"violating the wrapper's run() graph-safety contract." - ) - - # ============================================================================ # moe_output_memset_inplace (dense Path A) β€” unit tests # ============================================================================ From 41e5708d6aac05630764fd9f400ec1d8ee6169ac Mon Sep 17 00:00:00 2001 From: kangbintNV Date: Wed, 10 Jun 2026 16:20:59 +0800 Subject: [PATCH 07/13] docs: close v0.6.13 doc-check gaps + fix(moe) misleading topk_indices ICHECK message (#3546) ## Summary Two small, low-risk changes surfaced by the v0.6.13 doc-check pass. Touches docs/docstrings only, plus a one-line C++ assertion-message fix (no runtime behavior change). ### docs: close v0.6.13-only doc-check gaps (MISSING / STALE / docstring) [issues/3538](https://github.com/flashinfer-ai/flashinfer/issues/3538) Resolves the doc-check findings that are new in v0.6.13 vs v0.6.12, using docs + docstrings only (no new `@flashinfer_api` decorators). - **MISSING (API exported but not listed in `.rst`):** - `docs/api/fused_moe.rst`: add a "Multi-LoRA MoE (BGMV)" section listing `bgmv_moe` / `bgmv_moe_shrink` / `bgmv_moe_expand`. - `docs/api/quantization.rst`: list `nvfp4_quantize_per_token_cute_dsl` under its canonical `flashinfer.quantization.kernels.nvfp4_quantize` currentmodule (the package-level re-export is guarded by `is_cute_dsl_available()` and is not importable at docs-build time without `nvidia-cutlass-dsl`). - **Docstring completeness:** - `nvfp4_quantize_per_token_cute_dsl`: expand the one-line summary into a full NumPy-style docstring with Parameters/Returns, mirroring the sibling `nvfp4_quantize_cute_dsl` and documenting the per-token scale output. ### fix(moe): correct misleading `topk_indices` dtype ICHECK message In `csrc/fused_moe/noAuxTcKernels.cu`, the dtype check for `topk_indices` requires int32 (`encode_dlpack_dtype(...) == int32_code`), but the assertion message was copy-pasted from the `topk_values` check and wrongly read "must have the same dtype as scores" (scores are fp32/bf16, not int32), which misleads debugging. Message changed to "topk_indices must be int32 dtype", matching the nearby `routing_replay_out` int16 check style. Message-only change; the runtime condition is unchanged. ## Test plan - [ ] Docs build (`docs/build_docs.sh`) succeeds; new autosummary entries render and removed STALE entries no longer produce dangling references. - [ ] `flashinfer_document_check` shows 0 MISSING and the two v0.6.13 STALE items (`is_cute_dsl_available`, `trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant`) gone; Docstring Completeness has 0 findings. - [ ] No code-path change for the MoE ICHECK; existing MoE routing tests remain green. AI-assisted. ## Summary by CodeRabbit * **New Features** * Added documentation for Multi-LoRA MoE (BGMV) operators: `bgmv_moe`, `bgmv_moe_shrink`, and `bgmv_moe_expand`. * **Documentation** * Enhanced quantization kernel documentation with detailed per-token scaling specifications and layout guidance. * Cleaned up API documentation by removing outdated entries and sections. --- csrc/fused_moe/noAuxTcKernels.cu | 2 +- docs/api/fused_moe.rst | 13 +++++ docs/api/quantization.rst | 1 + .../quantization/kernels/nvfp4_quantize.py | 47 ++++++++++++++++++- 4 files changed, 61 insertions(+), 2 deletions(-) diff --git a/csrc/fused_moe/noAuxTcKernels.cu b/csrc/fused_moe/noAuxTcKernels.cu index d08d3e83686..3efa2e62573 100644 --- a/csrc/fused_moe/noAuxTcKernels.cu +++ b/csrc/fused_moe/noAuxTcKernels.cu @@ -348,7 +348,7 @@ void NoAuxTc(TensorView scores, TensorView bias, int64_t n_group, int64_t topk_g TVM_FFI_ICHECK(topk_values.dtype() == data_type) << "topk_values must have the same dtype as scores"; TVM_FFI_ICHECK(encode_dlpack_dtype(topk_indices.dtype()) == int32_code) - << "topk_indices must have the same dtype as scores"; + << "topk_indices must be int32 dtype"; // Validate and extract routing_replay_out // NOTE: dim0 >= num_tokens is intentionally NOT checked β€” with CUDA graphs the buffer diff --git a/docs/api/fused_moe.rst b/docs/api/fused_moe.rst index 4dd02773b23..aeb257c7c15 100644 --- a/docs/api/fused_moe.rst +++ b/docs/api/fused_moe.rst @@ -28,6 +28,19 @@ Utility Functions interleave_moe_scales_for_sm90_mixed_gemm fused_topk_deepseek +Multi-LoRA MoE (BGMV) +--------------------- + +Batched Gather-Matrix-Vector kernels for serving multiple LoRA adapters on +top of a Mixture-of-Experts layer (shrink + expand). + +.. autosummary:: + :toctree: ../generated + + bgmv_moe + bgmv_moe_shrink + bgmv_moe_expand + CUTLASS Fused MoE ----------------- diff --git a/docs/api/quantization.rst b/docs/api/quantization.rst index 5dfeddb2de8..83bb354835c 100644 --- a/docs/api/quantization.rst +++ b/docs/api/quantization.rst @@ -94,6 +94,7 @@ importable. :toctree: ../generated nvfp4_quantize_cute_dsl + nvfp4_quantize_per_token_cute_dsl .. currentmodule:: flashinfer.quantization.kernels.mxfp4_quantize diff --git a/flashinfer/quantization/kernels/nvfp4_quantize.py b/flashinfer/quantization/kernels/nvfp4_quantize.py index ef627bf23c8..c72d442849c 100644 --- a/flashinfer/quantization/kernels/nvfp4_quantize.py +++ b/flashinfer/quantization/kernels/nvfp4_quantize.py @@ -1797,7 +1797,52 @@ def nvfp4_quantize_per_token_cute_dsl( sf_layout: int = SF_LAYOUT_128x4, enable_pdl: bool | None = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Per-token NVFP4 activation quantization using CuTe-DSL.""" + r"""Per-token NVFP4 activation quantization using the CuTe-DSL kernel. + + Unlike :func:`nvfp4_quantize_cute_dsl`, which applies a single global + scale, this variant computes one quantization scale **per row (token)** of + the activation. Each row is scaled independently so that its largest + magnitude maps to the NVFP4 dynamic range, and the resulting per-token + scale is returned alongside the packed FP4 output and the E4M3 block + scale factors. + + - E4M3 block scale factors (FP8), ``sf_vec_size = 16`` + - E2M1 output format (4-bit, 2 values per byte) + - Supports 128x4, 8x4, and linear scale-factor layouts + + The kernel is compiled once per ``(K, dtype, sf_layout, pdl)`` tuple and + handles varying ``M`` (number of tokens) at runtime without recompilation. + + Parameters + ---------- + input : torch.Tensor + 2-D activation tensor of shape ``[M, K]`` with dtype fp16/bf16. ``K`` + must be divisible by ``NVFP4_SF_VEC_SIZE`` (16). + global_scale_inv : torch.Tensor + Scalar tensor (``float32``) holding the inverse global scale applied on + top of the per-token scale. A Python ``float`` is also accepted and + wrapped into a tensor internally. + sf_layout : int + Scale-factor layout (``0=128x4``, ``1=8x4``, ``2=linear``). + enable_pdl : bool, optional + Whether to enable Programmatic Dependent Launch. Auto-detected from + device capability (SM >= 9.0) when ``None``; pass ``False`` to force it + off. + + Returns + ------- + Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + ``(fp4_output, scale_output, per_token_scale)`` where: + + - ``fp4_output`` is the packed quantized tensor of shape ``[M, K/2]`` + with dtype ``uint8`` (two E2M1 values per byte). + - ``scale_output`` holds the E4M3 block scale factors (``uint8``) + reshaped to ``[padded_rows, padded_sf_cols]``. The padding depends on + ``sf_layout``: ``linear`` keeps ``M`` rows, while ``128x4`` / ``8x4`` + pad rows and columns up to the layout tile. + - ``per_token_scale`` is the per-row quantization scale of shape + ``[M]`` with dtype ``float32``. + """ from ...utils import device_support_pdl _valid_sf_layouts = (SF_LAYOUT_128x4, SF_LAYOUT_8x4, SF_LAYOUT_LINEAR) From 6ece522337b618c73050fa912b07058b68038a55 Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Wed, 10 Jun 2026 07:59:52 -0700 Subject: [PATCH 08/13] feat: [initial progress] Unified MoE API: MoELayer with cross-backend NVFP4 autotune (#3093) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Draft PR introducing a unified MoE layer that autotunes across NVFP4 backends β€” CuteDSL and TRTLLM FP4 routed β€” on a per-shape basis. ## Update β€” MVP cut (2026-05-31) Autonomous MVP-completion pass on a B200 (SM100). The cross-backend autotune objective is met and validated end-to-end. Full narrative in `docs/design_docs/flashinfer_moe_api.md` (Decision Log + the new "MVP As-Built Reference"). - **All MVP follow-ups (CR1–CR11) done & validated**: `tests/moe/test_unified_moe_api.py` **9/9** (layer + per-backend accuracy vs bf16, autotuner visits both candidates, CUDA-graph replay), `tests/moe/test_moe_api.py` **97/97** (CPU config + fail-fast validation), and the `unified_nvfp4_moe` sweep (128β†’16384 tokens) with `--refcheck` passing for both backends. - **Blocker found & fixed**: both runner adapters had never executed against the post-`main`-merge `core.py` (a stale raw-`moe_op` API + a class-vs-instance `tuning_config` bug). `TrtllmFp4RoutedRunner` now delegates to the canonical `core.MoERunner` (newly exported from `get_trtllm_moe_sm100_module()`); the unified adapters only translate Packs ⇄ the inner runner's native tensor list. Future direction (out of scope): make the low-level TVM-FFI ops take structured config objects (design doc Β§5) to kill positional-arg drift. - **New since first draft**: `local_expert_offset` wired into TRTLLM packing (+EP-offset test); fail-fast NVFP4/Swiglu scope validation; per-token-bucket winner cache; `tune_max_num_tokens` threaded into runner tuning; first-class NVFP4 weight prep (`prepare.py` + `*.prepare_weights`) replacing duplicated test/bench prep; `--refcheck` for the unified benchmark. ## What this adds - **`MoEConfig` / `MoEActivationPack` / `MoEWeightPack`** β€” frozen config dataclasses and per-call / long-lived tensor containers (`flashinfer/fused_moe/api.py`). Single `QuantVariant` enum replaces the old 3-axis dtype Γ— granularity Γ— variant split. - **`CuteDslNvfp4Runner` / `TrtllmFp4RoutedRunner`** β€” `TunableRunner` adapters with `pack_inputs(act, weights)` translating packs into the backend's native tensor list (`flashinfer/fused_moe/runners.py`). ~~each with its own `tuning_config`~~ β†’ each **delegates to a canonical inner runner** (`CuteDslFusedMoENvfp4Runner` / `core.MoERunner`) and builds its `tuning_config` **per-instance**. - **`MoELayer`** β€” stateful dispatcher that builds one runner per compatible backend config, then on first call ~~per shape~~ **per token-bucket** runs per-runner `choose_one` (within-backend tactic tuning) + cross-runner `bench_gpu_time` comparison (cross-backend selection). Caches the winner **keyed by tuning bucket** (`flashinfer/fused_moe/layer.py`). - **`unified_nvfp4_moe` benchmark routine** β€” wired into `benchmarks/flashinfer_benchmark.py` (not a standalone script). Emits one result row per candidate backend with the winner marked; supports `--refcheck` against a shared bf16 reference. - **`bench_unified_moe_today.sh`** β€” convenience wrapper at repo root invoking the infra routine across the eight shapes we care about (EP=1 sweep + EP=16 wide-EP regime). - **`tests/moe/test_unified_moe_api.py`** β€” accuracy tests (`MoELayer` vs bf16 ref; each backend vs same ref), plumbing tests (autotuner visits all candidates; CUDA graph capture + replay), **and a pre-routed EP-offset packing test**. - **`flashinfer/fused_moe/prepare.py`** *(new)* β€” first-class NVFP4 weight-prep helpers exposed as `TrtllmFp4Config.prepare_weights` / `CuteDslConfig.prepare_weights`; the test and benchmark no longer carry duplicated prep copies. ## Design notes - **Per-runner `choose_one`** β€” `AutoTuner.choose_one` assumes all runners share one `inputs` list during profiling. Our backends' native schemas differ (CuteDSL 12 tensors with unpacked topk + trailing output buffer; TRTLLM 8-field packed `MoEInputs`). Resolution: call `choose_one` once per runner (within-backend tactic selection) and use `bench_gpu_time` to compare the per-runner winners. ~~Each runner owns its own `tuning_config` as a class attribute.~~ β†’ **`tuning_config` is built per-instance** (the TRTLLM runner via `MoERunner._make_tuning_config`, which also threads `ExecutionConfig.tune_max_num_tokens` into the bucket set). - **Shared-reference accuracy testing** β€” each backend is tested against the same bf16 reference, not against the other backend. Catches shared-mode failures (both wrong in the same way) that cross-backend agreement would miss. `--refcheck` brings the same check to the benchmark. - **Weight layout: `Shuffled_MajorK` only** β€” the only NVFP4-compatible TRTLLM layout today. Multi-layout autotune (opt-in additional variants via a future `layouts` field on `TrtllmFp4Config`) is a V2 extension β€” design accommodates it without core changes. ## Out of scope for this PR - FP8 / BF16 / MxInt4 backends - Monolithic routing (our path is pre-routed) - Unpacked topk for TRTLLM β€” #2425 (when that lands, `TrtllmFp4RoutedRunner` drops the inline `(id << 16) | bf16_bits` packing) - Multi-layout TRTLLM autotune across `Shuffled_BlockMajorK` / `NoShuffle_MajorK` - SM120/SM121 support for CuteDSL (kernel is SM100/SM103 only today; `CuteDslConfig.supported` tightened to match) - First-class **activation** prep + structured-config TVM-FFI boundary (Β§5) β€” post-MVP carryover ## Status ~~Draft β€” benchmark numbers and observation notes are being collected on a Blackwell box. Will de-draft once the winner-flip table is filled in and accuracy tests pass on hardware.~~ **MVP scope complete and validated on B200 (SM100).** Kept as a **draft PR on purpose** so CI skips while further edits are pushed autonomously β€” flip to "Ready for review" when you want CI to run. ## Test plan - [x] `pytest tests/moe/test_unified_moe_api.py -v` (SM100/SM103) β€” **9/9** - [ ] ~~`./bench_unified_moe_today.sh` β€” winner column matches expected per-shape~~ β†’ not run this pass; equivalent coverage via the `unified_nvfp4_moe` routine sweep (128β†’16384) + `--refcheck` (both backends pass). The script's EP=16 shapes are still worth a dedicated run. - [x] Pre-commit hooks pass ## Related PRs - #3453 πŸ€– Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **New Features** * Unified MoE API surfaced package-wide: immutable configs, MoELayer (cached cross-backend autotune winner), backend runners, NVFP4 weight-prep utilities, and re-exported unified API symbols. * **Benchmarking** * New benchmark sweep script for unified NVFP4 MoE and registration to record token-size sweeps to CSV. * **Documentation** * Added comprehensive FlashInfer Unified MoE API design doc with migration plan. * **Tests** * Added CPU and gated-GPU tests for API, validation, accuracy, and autotune/dispatch behavior; improved enum reprs for round-trip logging. --------- Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Yang Xu --- benchmarks/bench_unified_moe.sh | 45 + .../routines/flashinfer_benchmark_utils.py | 7 + benchmarks/routines/moe.py | 314 ++++++ docs/design_docs/flashinfer_moe_api.md | 932 ++++++++++++++++++ flashinfer/fused_moe/__init__.py | 46 + flashinfer/fused_moe/api.py | 541 ++++++++++ flashinfer/fused_moe/core.py | 38 +- flashinfer/fused_moe/layer.py | 202 ++++ flashinfer/fused_moe/prepare.py | 268 +++++ flashinfer/fused_moe/runners.py | 329 +++++++ flashinfer/tllm_enums.py | 19 + tests/moe/test_unified_moe.py | 925 +++++++++++++++++ tests/moe/test_unified_moe_fuzz.py | 788 +++++++++++++++ 13 files changed, 4441 insertions(+), 13 deletions(-) create mode 100755 benchmarks/bench_unified_moe.sh create mode 100644 docs/design_docs/flashinfer_moe_api.md create mode 100644 flashinfer/fused_moe/api.py create mode 100644 flashinfer/fused_moe/layer.py create mode 100644 flashinfer/fused_moe/prepare.py create mode 100644 flashinfer/fused_moe/runners.py create mode 100644 tests/moe/test_unified_moe.py create mode 100644 tests/moe/test_unified_moe_fuzz.py diff --git a/benchmarks/bench_unified_moe.sh b/benchmarks/bench_unified_moe.sh new file mode 100755 index 00000000000..e4e7b428729 --- /dev/null +++ b/benchmarks/bench_unified_moe.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +# Unified NVFP4 MoE benchmark sweep (DeepSeek-V3 geometry). +# +# Run from the repo root: bash benchmarks/bench_unified_moe.sh +# +# Shapes (see docs/design_docs/flashinfer_moe_api.md Β§11 for the reference +# winner/latency tables this sweep reproduces): +# EP=1 sweep: tokens = 1, 16, 1024, 4096, 16384 +# EP=16 sweep: tokens = 1, 16, 4096 +# +# DeepSeek-V3 geometry: hidden=7168, intermediate=2048, experts=256, top_k=8, +# n_group=8, topk_group=4, routed_scaling_factor=2.5. +# +# Expectation: the MoELayer winner and per-backend latencies track the +# cross-backend tables in the design doc (TRTLLM-gen wins ≀512 tokens, CuteDSL +# wins β‰₯1024). Requires an SM100 (Blackwell) GPU. +set -euo pipefail + +CSV=/tmp/unified_moe_today.csv +: > "$CSV" # truncate + +BASE=( + --routine unified_nvfp4_moe + --hidden_size 7168 --intermediate_size 2048 + --num_experts 256 --top_k 8 + --n_group 8 --topk_group 4 + --routing_method deepseek_v3 --routed_scaling_factor 2.5 + --output_path "$CSV" + -v +) + +# EP=1 sweep (local_num_experts defaults to num_experts=256) +for M in 1 16 1024 4096 16384; do + echo "=== S_EP1 num_tokens=$M ===" + python benchmarks/flashinfer_benchmark.py "${BASE[@]}" --num_tokens "$M" +done + +# EP=16 sweep (local_num_experts=16) +for M in 1 16 4096; do + echo "=== S_EP16 num_tokens=$M ===" + python benchmarks/flashinfer_benchmark.py "${BASE[@]}" \ + --num_tokens "$M" --local_num_experts 16 +done + +echo "Results: $CSV" diff --git a/benchmarks/routines/flashinfer_benchmark_utils.py b/benchmarks/routines/flashinfer_benchmark_utils.py index 9bcc58e275a..870dbf15d15 100644 --- a/benchmarks/routines/flashinfer_benchmark_utils.py +++ b/benchmarks/routines/flashinfer_benchmark_utils.py @@ -200,6 +200,7 @@ "cutlass_fused_moe", "cute_dsl_fp4_block_scale_moe", "b12x_fused_moe", + "unified_nvfp4_moe", "bgmv_moe", ], "moe_comm": [ @@ -521,6 +522,12 @@ def dtype_str_to_torch_dtype(dtype_str): "12.0": ["b12x"], "12.1": ["b12x"], }, + # MoELayer cross-backend NVFP4: intersection of CuteDSL + TRTLLM FP4 support. + # SM100 only (Blackwell); unlisted archs fall through to [] (skipped). + "unified_nvfp4_moe": { + "10.0": ["unified"], + "10.3": ["unified"], + }, # NORM "rmsnorm": { "7.5": ["cute-dsl"], diff --git a/benchmarks/routines/moe.py b/benchmarks/routines/moe.py index c873318b407..93ba123b1f4 100644 --- a/benchmarks/routines/moe.py +++ b/benchmarks/routines/moe.py @@ -113,6 +113,8 @@ def run_moe_test(args): return testB12xFusedMoe(args) elif args.routine == "bgmv_moe": return testBgmvMoe(args) + elif args.routine == "unified_nvfp4_moe": + return testUnifiedNvfp4Moe(args) else: raise ValueError(f"Unsupported routine: {args.routine}") @@ -2404,6 +2406,318 @@ def run_fp8_per_tensor_moe( return res +# ============================================================================= +# Unified NVFP4 MoE β€” MoELayer-based cross-backend autotune +# ============================================================================= + + +def testUnifiedNvfp4Moe(args): + """MoELayer cross-backend autotune for NVFP4. + + Exercises the unified MoE API: per-shape, builds MoELayer over + {CuteDslConfig, TrtllmFp4Config}, registers both backend-native weight + views on a single MoEWeightPack, dispatches once, reports the winner + and both backends' best-tactic latencies. + + Emits one result row per backend candidate so CSV output carries the + cross-backend comparison directly. + """ + from flashinfer.autotuner import AutoTuner + from flashinfer.fused_moe import ( + ActivationConfig, + CuteDslConfig, + ExecutionConfig, + ExpertConfig, + MoEActivationPack, + MoEConfig, + MoELayer, + MoEWeightPack, + QuantConfig, + QuantVariant, + RoutingConfig, + TrtllmFp4Config, + ) + from flashinfer.fused_moe.api import BackendOptions + + if args.verbose >= 1: + print("[INFO] Running testUnifiedNvfp4Moe") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + num_tokens = args.num_tokens + hidden_size = args.hidden_size + intermediate_size = args.intermediate_size + num_experts = args.num_experts + top_k = args.top_k + local_expert_offset = args.local_expert_offset + local_num_experts = args.local_num_experts or num_experts + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype) + + res = [] + backends = ["unified"] + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + if args.verbose >= 1: + print( + f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, " + f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}, " + f"local_experts={local_num_experts}" + ) + + # ---- Build shared bf16 reference weights ------------------------------ + torch.manual_seed(0) + w1_bf16 = ( + torch.randn( + local_num_experts, + 2 * intermediate_size, + hidden_size, + dtype=torch.bfloat16, + device=device, + ) + / 10 + ) + w2_bf16 = ( + torch.randn( + local_num_experts, + hidden_size, + intermediate_size, + dtype=torch.bfloat16, + device=device, + ) + / 10 + ) + + # ---- Backend-native weight views β€” first-class prepare helpers (CR2) -- + # Both views are built from the SAME canonical bf16 weights, so the two + # backends are directly comparable (and a shared reference is meaningful). + cute_dsl_view = CuteDslConfig.prepare_weights( + w1_bf16, + w2_bf16, + num_local_experts=local_num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + ) + trtllm_view = TrtllmFp4Config.prepare_weights( + w1_bf16, + w2_bf16, + num_local_experts=local_num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + ) + + # ---- Activation pack -------------------------------------------------- + # The activation (NVFP4-quantized hidden states + pre-routed indices) still + # comes from the canonical test data creator; only weight prep is + # first-class today (activation-prep promotion tracked under CR2). + import os + import sys + + _repo_root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + from tests.moe.test_cute_dsl_fused_moe import ( + check_accuracy, + compute_reference_moe_fp4, + create_moe_tensors, + ) + + # Wide-EP (MVP): model a single rank as a complete MoE over its + # local_num_experts experts β€” route the activation WITHIN the local experts + # (selected ids in [0, local_num_experts)) so every token is computed + # locally. This is the "local-only" proxy: it avoids the unfaithful + # global-routing-but-local-weights setup (most tokens skipped) and keeps the + # derived FLOP/bandwidth metrics correct (all tokens computed, active + # experts <= local). For EP=1 (local == global) this is unchanged. + routing_num_experts = local_num_experts + routing_top_k = min(top_k, local_num_experts) + cute_dsl_data = create_moe_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=routing_num_experts, + num_local_experts=local_num_experts, + top_k=routing_top_k, + device=device, + ) + # cute_dsl_data["x_sf"] is already unsqueezed to [M, H//16, 1]; strip that + # for the Pack (runner re-applies unsqueeze in pack_inputs). + x_sf = cute_dsl_data["x_sf"].squeeze(-1) + act_pack = MoEActivationPack( + hidden_states_q=cute_dsl_data["x"], + hidden_states_scale=x_sf, + selected_experts=cute_dsl_data["token_selected_experts"], + final_scales=cute_dsl_data["token_final_scales"], + ) + + num_active_experts = int(act_pack.selected_experts.unique().numel()) + + weight_pack = MoEWeightPack() + weight_pack.prepare_for("cute_dsl_nvfp4", cute_dsl_view) + weight_pack.prepare_for("trtllm_fp4_routed", trtllm_view) + + # ---- MoELayer config -------------------------------------------------- + config = MoEConfig( + routing=RoutingConfig( + num_experts=num_experts, + top_k=top_k, + n_group=args.n_group, + topk_group=args.topk_group, + routed_scaling_factor=args.routed_scaling_factor, + ), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig( + intermediate_size=intermediate_size, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + ), + activation=ActivationConfig(), + backend=BackendOptions(candidates=(CuteDslConfig(), TrtllmFp4Config())), + execution=ExecutionConfig(tune_max_num_tokens=max(num_tokens, 8192)), + ) + + # ---- Dispatch: trigger cross-backend selection ------------------------ + with autotune(True): + layer = MoELayer(config, device=device) + _ = layer(act_pack, weight_pack) + + winner_key = layer.winner_backend or "?" + if args.verbose >= 1: + print( + f"[INFO] MoELayer winner: {winner_key} " + f"(candidates: {[r.backend_key for r in layer.runners]})" + ) + + # ---- Optional shared-reference accuracy check (CR10/CR11) ------------- + # Both backend views derive from the same bf16 weights, so a single bf16 + # reference is valid for every candidate. Catches shared-mode errors that + # cross-backend agreement would miss. + ref_output = None + if args.refcheck: + ref_output = compute_reference_moe_fp4( + hidden_states=cute_dsl_data["x_bf16"].float().to(device), + gemm1_weights=w1_bf16.float().to(device), + gemm2_weights=w2_bf16.float().to(device), + token_selected_experts=act_pack.selected_experts, + token_final_scales=act_pack.final_scales, + num_tokens=num_tokens, + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + fc2_input_scale=cute_dsl_data["fc2_input_scale"], + ) + + # ---- Per-backend best-tactic timing (one result row per candidate) ---- + tuner = AutoTuner.get() + for runner in layer.runners: + inputs = runner.pack_inputs(act_pack, weight_pack) + with autotune(True): + _, tactic = tuner.choose_one( + custom_op=f"moe_{runner.backend_key}", + runners=[runner], + tuning_config=runner.tuning_config, + inputs=inputs, + ) + + def _call(r=runner, i=inputs, t=tactic): + return r.forward(i, tactic=t) + + times = bench_gpu_time( + fn=_call, + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + sleep_after_run=False, + enable_cupti=args.use_cupti, + use_cuda_graph=not args.no_cuda_graph, + cold_l2_cache=True, + ) + median_time = float(np.median(times)) + std_time = float(np.std(times)) + # Use the actual routing geometry (local experts / clamped top_k), not + # the global config, so the wide-EP local-only proxy reports honest + # FLOP/byte counts. For EP=1 these equal num_experts/top_k. + tflops = calculate_moe_tflops( + num_tokens, + hidden_size, + intermediate_size, + routing_num_experts, + routing_top_k, + median_time, + ) + tb_per_sec = calculate_moe_kernel_bandwidth( + num_tokens, + hidden_size, + intermediate_size, + routing_num_experts, + routing_top_k, + median_time, + input_dtype, + weight_dtype, + input_format="nvfp4", + weight_format="nvfp4", + routing_logits_dtype=None, + active_experts=num_active_experts, + verbose=args.verbose, + ) + + # Tag: "unified/" + star the winner + is_winner = runner.backend_key == winner_key + backend_label = f"unified/{runner.backend_key}" + ("*" if is_winner else "") + print_perf_metrics(backend_label, median_time, std_time, tflops, tb_per_sec) + + # Shared-reference accuracy check for this candidate (CR10/CR11). + refcheck_passed = None + if ref_output is not None: + out = runner.forward(inputs, tactic=tactic) + refcheck_passed, pct, atol = check_accuracy(out, ref_output) + status = "PASS" if refcheck_passed else "FAIL" + print( + f"[REFCHECK] {backend_label}: {status} " + f"({pct * 100:.2f}% within atol={atol:.4f} vs bf16 reference)" + ) + if not refcheck_passed and not args.allow_output_mismatch: + print( + f"[ERROR] {backend_label} failed reference check. " + f"Re-run with --allow_output_mismatch to continue past mismatches." + ) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend_label + cur_res["refcheck_passed"] = refcheck_passed + cur_res["num_tokens"] = num_tokens + cur_res["hidden_size"] = hidden_size + cur_res["intermediate_size"] = intermediate_size + cur_res["num_experts"] = num_experts + cur_res["top_k"] = top_k + cur_res["local_expert_offset"] = local_expert_offset + cur_res["local_num_experts"] = local_num_experts + cur_res["input_dtype"] = input_dtype + cur_res["weight_dtype"] = weight_dtype + cur_res["fp4_mode"] = "nvfp4" + res.append(cur_res) + + return res + + def testBgmvMoe(args): """ Benchmark BGMV MoE kernels for multi-LoRA inference. diff --git a/docs/design_docs/flashinfer_moe_api.md b/docs/design_docs/flashinfer_moe_api.md new file mode 100644 index 00000000000..c933023dc02 --- /dev/null +++ b/docs/design_docs/flashinfer_moe_api.md @@ -0,0 +1,932 @@ +# FlashInfer Unified MoE API + +*Design Document Β· v0.1 Β· March 2026* + +## 1. Motivation + +FlashInfer currently exposes MoE functionality through a family of flat, positional-argument functions: + +``` +trtllm_fp4_block_scale_moe(routing_logits, routing_bias, hidden_states, + hidden_states_scale, gemm1_weights, gemm1_weights_scale, gemm1_bias, + gemm1_alpha, gemm1_beta, gemm1_clamp_limit, gemm2_weights, ...) +# 30 positional arguments. Backend selected by function name. +``` + +Problems this creates: + +- Users must know which backend (cutlass vs trtllm_fp4 vs trtllm_fp8) to call for their hardware and quant config +- 30-argument flat signatures are error-prone β€” type mismatches at param 18 are silent until C++ segfaults +- Three near-identical functions for fp4 / fp8 block / fp8 per-tensor diverge over time +- No autotuning β€” optimal backend varies by batch size, hardware, and routing config +- No component benchmarking β€” impossible to attribute latency to routing vs gemm vs finalize +- No repro tracing β€” user-filed issues require manual reconstruction of call site + +## 2. Design Principles + +- Config is pure data β€” frozen dataclasses, no behavior, serializable via repr() +- Frontend owns the schema β€” backends are consumers +- PyTorch-familiar surface β€” follows nn.Module / F.functional pattern +- TVM FFI end-to-end β€” structured objects cross the C++ boundary, no positional arg counting +- Immutable by default β€” use dataclasses.replace() for variants, never mutate + +### Example Overview + +``` +# --- Define config once --- +config = MoEConfig( + routing=RoutingConfig( + num_experts=256, + top_k=8, + method=RoutingMethodType.DeepSeekV3, + ), + quant=QuantConfig(QuantDtype.FP4, QuantGranularity.BlockScale), + experts=ExpertConfig(intermediate_size=2048, local_num_experts=32), + backends=[TrtllmFp4Config(extra_backend_params...), CutlassConfig(extra_backend_params...)], +) +# --- Find possible backends --- +backends = MoELayer.find_backends(**config) +# this contains {"trtllm_fp4":TrtllmFp4Config(), "cutlass_fp4":CutlassConfig()} +# or {"trtllm_fp4":"unsupported reason...", "cutlass_fp4":CutlassConfig()} +# more modification to the backends' parameters could be done here +backends=["trtllm_fp4":TrtllmFp4Config(extra_backend_params...),"cutlass_fp4":CutlassConfig(extra_backend_params...)] +# --- Prepare Inputs Data --- +weight_pack = MoEWeightPack() +# the data is possibly obtained through helper functions then added here +weight_pack.prepare_for("trtllm_fp4", trtllm_weights) weight_pack.prepare_for("cutlass_fp4", cutlass_weights) +act_pack = MoEActivationPack( + hidden_states_q=cute_dsl_data["x"], + hidden_states_scale=x_sf, + selected_experts=cute_dsl_data["token_selected_experts"], + final_scales=cute_dsl_data["token_final_scales"], +) +tensors = (act_pack, weight_pack) +# --- Eager (heuristic backend) --- +output = moe_layer(tensors, **config, backends=backends) # optional backends selection +# --- Autotuned eager --- +with autotune(True): + for tensors in calibration_data: + output = moe_layer(tensors, **config) +# --- Production layer (amortized, cached) --- +layer = MoELayer(**config) +layer = layer.get_tuned_layer(tensors) # ensures selection has been done +output = layer(tensors) +# --- Benchmark --- +layer.benchmark(Gemm1Tensors) # isolate gemm1 +layer.benchmark_all() # full breakdown +# --- Variant via immutable replace --- +fp8_config = dataclasses.replace(config, quant=QuantConfig(QuantDtype.FP8)) +fp8_layer = MoELayer(**fp8_config) +# --- Repro from issue log --- +repro = MoERepro.from_file("user_issue.log") +repro.run() +repro.benchmark_all() +``` + +## 3. Config Hierarchy + +All configs are frozen dataclasses registered with TVM's object system. The hierarchy: + +| Config | Owns | +| --- | --- | +| RoutingConfig | num_experts, top_k, routing method, grouping params, scaling factor | +| QuantConfig | dtype (fp4/fp8/bf16), granularity (per-tensor/per-token/block) | +| ExpertConfig | intermediate_size, local sharding params | +| ActivationConfig | activation type (swiglu/geglu/relu2/identity) | +| BackendOptions | ordered candidate set via \| operator | +| ExecutionConfig | do_finalize, enable_pdl, tune_max_num_tokens, output tensor | +| MoEConfig | assembles all above; supports \*\*unpacking protocol | + +### 3.1 RoutingConfig + +``` +@tvm.register_object('flashinfer.RoutingConfig') +@dataclass(frozen=True) +class RoutingConfig: + num_experts: int + top_k: int + method: RoutingMethodType = RoutingMethodType.Default + n_group: Optional[int] = None + topk_group: Optional[int] = None + routed_scaling_factor: Optional[float] = None +``` + +### 3.2 BackendOptions + +Individual backend configs provided in an ordered list. The autotuner or heuristic selects among valid candidates at runtime. + +``` +# Single backend +backends = [TrtllmFp4Config()] +# Multiple candidates β€” autotuner or heuristic picks best +backends = [TrtllmFp4Config(), TrtllmFp8BlockConfig(), CutlassConfig()] +# | is associative, returns BackendOptions +# CutlassConfig is always the universal fallback +``` + +Each backend config declares its own preconditions: + +``` +class TrtllmFp4Config: + @classmethod + def supported(cls, arch: int) -> bool: + return arch >= 90 # Hopper+ +class CutlassConfig: + @classmethod + def supported(cls, arch: int) -> bool: + return True # universal fallback +``` + +### 3.3 MoEConfig β€” \*\*unpack protocol + +MoEConfig implements keys() and __getitem__ so it can be unpacked directly with \*\*. This allows the same config object to be passed to both the eager function and MoELayer. + +``` +config = MoEConfig( + routing=RoutingConfig(num_experts=256, top_k=8, method=RoutingMethodType.DeepSeekV3), + quant=QuantConfig(QuantDtype.FP4, QuantGranularity.BlockScale), + experts=ExpertConfig(intermediate_size=2048, local_num_experts=32), + backends=[TrtllmFp4Config(), CutlassConfig()], +) +# Unpack into any call accepting these kwargs +output = moe_layer(tensors, **config) +layer = MoELayer(**config) +# Immutable variant +fp8_config = dataclasses.replace(config, quant=QuantConfig(QuantDtype.FP8)) +``` + +## 4. Public API + +### 4.1 Eager function + +Stateless. No autotuning. Backend selected by heuristic priority table. Equivalent to the current flat functions but with structured config. + +``` +@flashinfer_api +def moe_layer(tensors: MoETensors, *, routing, quant, experts, + activation=..., backends=..., execution=...) -> Tensor: + ... +# Usage +output = moe_layer(tensors, **config) +``` + +### 4.2 MoELayer β€” stateful, autotuned + +Holds workspace and backend selection cache. On first call, selects and caches the best valid backend for the observed tensor shapes and arch. Subsequent calls are zero-overhead dispatch. + +``` +layer = MoELayer(**config) +output = layer(tensors) # first call: selects backend +output = layer(tensors) # subsequent: cached dispatch +# Component benchmark +layer.benchmark(Gemm1Tensors) # isolate gemm1 latency +layer.benchmark_all() # full breakdown by component +``` + +### 4.3 autotune context manager + +Within the with block, every moe_layer call profiles all valid BackendOptions candidates. On exit, the best backend is cached keyed by (device, shape, config hash). + +``` +with autotune(True): + for tensors in calibration_data: + output = moe_layer(tensors, **config) +# After the block, MoELayer uses the measured best +layer = MoELayer(**config) +``` + +Without the context manager, moe_layer uses the heuristic DEFAULT_PRIORITY table β€” no profiling overhead, reasonable defaults. + +## 5. TVM FFI β€” Structured Boundary Crossing + +The structured config crosses the Python/C++ boundary via TVM's object system. Reflection fires once at the crossing; all downstream C++ access is direct struct member access. + +``` +// C++ side β€” mirrors Python dataclass +class RoutingConfigNode : public Object { + public: + int num_experts; + int top_k; + int method; + Optional n_group; + // ... + static RoutingConfigNode FromObject(ObjectRef obj); // reflection once + TVM_DECLARE_FINAL_OBJECT_INFO(RoutingConfigNode, Object); +}; +TVM_REGISTER_GLOBAL('flashinfer.moe_layer_fp4') +.set_body_typed([](MoEConfig config, ...) { + auto routing = RoutingConfigNode::FromObject(config->routing); + dispatch_fp4(routing.num_experts, routing.top_k, ...); // native access +}); +``` + +With CUDA graph capture: FFI crossing happens at capture time. Graph replay is pure device-side β€” zero host overhead on the hot path. + +## 6. Repro Tracing + +### 6.1 @flashinfer_api decorator + +The existing decorator gains one additional responsibility: emit a single-line repro log on every MoE call. + +``` +logger.debug( + f'REPRO fn={fn.__name__} ' + f'config={repr(config)} ' + f'shapes={tensor_shapes} ' + f'device={torch.cuda.get_device_name()} ' + f'version={flashinfer.__version__}' +) +``` + +Because `repr(config)` emits valid Python constructor syntax (all enums use qualified repr), the log line is directly eval-able. No JSON parsing, no schema versioning needed. + +### 6.2 MoERepro + +Takes a repro log line, reconstructs the config, synthesizes tensors of the correct shape, and reproduces the call. + +``` +repro = MoERepro.from_file('issue_42.log') +# Reproduce +output = repro.run() +# Component breakdown +repro.benchmark_all() +# RoutingConfig 0.12ms +# Gemm1Config 1.43ms +# ActivationConfig 0.08ms +# Gemm2Config 1.39ms +# FinalizeConfig 0.11ms +# Total 3.13ms +# Isolate backend to narrow regression +repro.isolate_backend(CutlassConfig()) +repro.isolate_backend(TrtllmFp4Config()) +``` + +### 6.3 MoEConfig.from_repr β€” safe eval + +> **Not shipped in the MVP (post-MVP).** This eval-based deserializer belongs to +> the long-range repro design only. The implementation intentionally omits +> `from_repr` β€” eval-based deserialization is a security smell (review C4-C5/C39) +> and is deferred with the rest of the repro tooling (see "Post-MVP Carryover"). +> `repr(config)` still round-trips for logging; only the parser is deferred. + +eval() with a restricted namespace. Only config constructors in scope β€” no arbitrary execution risk. + +``` +@classmethod +def from_repr(cls, s: str) -> MoEConfig: + ns = { + 'MoEConfig': cls, 'RoutingConfig': RoutingConfig, + 'RoutingMethodType': RoutingMethodType, ..., + '__builtins__': {}, # no arbitrary execution + } + return eval(s, ns) +``` + +## 7. File Layout + +``` +flashinfer/ + moe_layer/ + __init__.py # public exports + config.py # all dataclasses and enums + layer.py # MoELayer + functional.py # moe_layer eager function + autotune.py # autotune context manager + backends/ + __init__.py # BACKEND_REGISTRY, DEFAULT_PRIORITY + trtllm_fp4.py # TrtllmFp4Config + adapter + trtllm_fp8.py # Fp8Block + Fp8PerTensor + adapters + cutlass.py # CutlassConfig + adapter + repro.py # MoERepro + tensors.py # MoETensors, Gemm1Tensors, Gemm2Tensors +``` + +## 8. Migration Path + +The flat functions are not removed. They become internal adapters called only from BACKEND_REGISTRY. Public surface is the new API. + +| Phase | Action | User Impact | +| --- | --- | --- | +| 1. Config layer | Add dataclasses + MoEConfig | None β€” flat API unchanged | +| 2. Adapters | Wrap flat functions in adapters | None β€” flat API unchanged | +| 3. moe_layer eager | Publish new unified function | Opt-in to new API | +| 4. MoELayer + autotune | Add stateful layer + context manager | Opt-in | +| 5. Repro tracing | Extend @flashinfer_api | Automatic for all users | +| 6. Deprecate flat APIs | Add deprecation warnings | Warning on old call sites | + +Each phase is independently valuable and independently shippable. + +## 9. Review Comments + +These comments were transcribed from the DOCX review metadata. Anchors refer to the text range the comment was attached to in Word. + +### C0 β€” Reviewer 1 + +**Anchor:** `FlashInfer Unified MoE API` + +> Would it be helpful to describe first (1) what FlashInfer should be responsible for and (2) what frameworks should be responsible for, e.g., weight preprocessing (and what FlashInfer provides to support/supplement it) + +### C1 β€” Reviewer 2 + +**Anchor:** `Design Principles` + +> I think one thing we are missing is backend introspection. Frameworks should be able to ask questions like "what are the supported backends for WideEP", "what data-types can I use for my model" etc. to help with their own heuristics. Maybe this doesn't fall under this scope, but it is something we likely need to expose longer-term + +### C2 β€” Reviewer 3 + +**Anchor:** `Design Principles` + +> agree. i'll add it to the doc before resolving. thanks + +### C3 β€” Reviewer 3 + +**Anchor:** `Design Principles` + +> please see "Find possible backends" in "Example Overview" if that workflow sound reasonable + +### C4 β€” Reviewer 4 + +**Anchor:** `Config is pure data β€” frozen dataclasses, no behavior, serializable via repr()` + +> frozen dataclasses is good. i just have one concern on repr() serialization, its output is not a versioned format, if a constructor signature changes between releases, old repro logs cannot be parseable. Have you considered to_dict() and from_dict() with an explicit schema version field? + +### C5 β€” Reviewer 3 + +**Anchor:** `Config is pure data β€” frozen dataclasses, no behavior, serializable via repr()` + +> what could go wrong if we don't worry about parsing old repro logs? e.g. limiting the usage to same version repro. as a "repro" the version should align anyway + +### C6 β€” Reviewer 5 + +**Anchor:** `Frontend owns the schema β€” backends are consumers` + +> A tricky part is scale factor swizzling pattern and the interleaving required by trtllm-gen MoE. Many users aren't clear why these are needed and Flashinfer needs to emphasize these and provide useful tools. For example, MoEConfig can have a method called weight_preprocessing()2 total reactionsReviewer 1 reacted with βž• at 2026-05-12 20:00 PMReviewer 6 reacted with βž• at 2026-05-13 01:13 AM + +### C7 β€” Reviewer 3 + +**Anchor:** `Frontend owns the schema β€” backends are consumers` + +> i missed this detail in the doc. but the plan is to have prepare functions like this https://github.com/flashinfer-ai/flashinfer/pull/3093/changes#diff-0995b259d16c6c37f02b6cc5d8825f9147ddf095a1c9b29846bbdee96ce2ab96R2278-R2280 weight_pack = MoEWeightPack()weight_pack.prepare_for("cute_dsl_nvfp4", cute_dsl_view)weight_pack.prepare_for("trtllm_fp4_routed", trtllm_view)and it's organized into something like a static helper function in TrtllmFp4Config(). (wrapping current standalone functions) + +### C8 β€” Reviewer 7 + +**Anchor:** `Frontend owns the schema β€” backends are consumers` + +> Could you elaborate on this? I'm not sure I get it. + +### C9 β€” Reviewer 7 + +**Anchor:** `Immutable by default β€” use dataclasses.replace() for variants, never mutate` + +> Nit: even though I'm all for this, I think it mixes up design and implementation details + +### C10 β€” Reviewer 7 + +**Anchor:** `# --- Define config once ---` + +> Is it fair to summarize the idea of this design as:- User knows the dimensions of their MoE layer- User knows how it is quantized- User knows how routing should happen- User knows batch sizes that they are interested inThis information is mapped to a config type that FlashInfer defines, FlashInfer takes that in and returns something the user can call on their inputs.The `backends` are is just here for optionally overriding the selection? I would almost leave it out of the config type, and pass it as an additional argument.Could we have another function which returns the compatible backends for a config?3 total reactionsReviewer 8 reacted with βž• at 2026-05-12 16:02 PMReviewer 5 reacted with βž• at 2026-05-12 16:33 PMReviewer 1 reacted with βž• at 2026-05-12 20:01 PM + +### C11 β€” Reviewer 3 + +**Anchor:** `# --- Define config once ---` + +> that makes a lot of sense to me! thx for the suggestion + +### C12 β€” Reviewer 3 + +**Anchor:** `# --- Define config once ---` + +> >Could we have another function which returns the compatible backends for a config?that sounds super reasonable too + +### C13 β€” Reviewer 3 + +**Anchor:** `# --- Define config once ---` + +> pls see updated "Example Overview" + +### C14 β€” Reviewer 5 + +**Anchor:** `quant=QuantConfig(QuantDtype.FP4, QuantGranularity.BlockScale)` + +> BTW we also have "MXFP4", so the block size (16 or 32) needs to be exposed + +### C15 β€” Reviewer 8 + +**Anchor:** `experts=ExpertConfig(intermediate_size=2048, local_num_experts=32),` + +> hidden_size as well? + +### C16 β€” Reviewer 8 + +**Anchor:** `backends=[TrtllmFp4Config(extra_backend_params...), CutlassConfig(extra_backend_params...)]` + +> Very often community developers are not aware of which backends are supported for a given routing + quant config. Agree with Reviewer 7 here that there should be an interface where users can query for a list of backends given quant config, hardware architecture. This can also extend to a benchmark option where we give the users an option to also see which backend performs best for a set of shapes of interest. + +### C17 β€” Reviewer 3 + +**Anchor:** `backends=[TrtllmFp4Config(extra_backend_params...), CutlassConfig(extra_backend_params...)]` + +> added a "Find possible backends" step in the example + +### C18 β€” Reviewer 2 + +**Anchor:** `"unsupported reason..."` + +> I wonder if this format would be confusing. I would suggest a `find_backends()` that returns only the valid ones for the config. And then a separate `request_backend(backend_name, config, extra_backend_params)` that returns `Config | Error` + +### C19 β€” Reviewer 2 + +**Anchor:** `"unsupported reason..."` + +> This way if frameworks dont care about customizing backends they can use `find_backends` and just directly use the result. But if the framework is aware of the different backends and wants to customize them it could instead explicitly call request_backend() with the customised parameters + +### C20 β€” Reviewer 1 + +**Anchor:** `output = layer(tensors)` + +> When is the backend decided in this workflow? During the first call of layer(tensors)? + +### C21 β€” Reviewer 3 + +**Anchor:** `output = layer(tensors)` + +> added a line above as a possible solution to make it more explicit and clear + +### C22 β€” Reviewer 1 + +**Anchor:** `layer.benchmark(Gemm1Tensors) # isolate gemm1` + +> Should it be layer's responsibility to do benchmark? I thought it should be something like benchmark(layer, tensors). + +### C23 β€” Reviewer 3 + +**Anchor:** `layer.benchmark(Gemm1Tensors) # isolate gemm1` + +> yeah your suggestion sounds better + +### C24 β€” Reviewer 9 + +**Anchor:** `routing method` + +> Since the interface is being changed, it would be a good chance to rework how the routing method is passed to Flashinfer. I would find it a lot better to not expose the internal enumeration of possible routing methods, but instead receive the routing configuration (scoring function, topK, renorm, topK groups, etc.) and select the correct routing internally. A `supported()` classmethod can be added to allow queries.The drawback of the current enumeration is that it encodes multiple pieces of information. This leads to a large number of methods and to the inference frameworks implementing some version of the above `supported()` function to find out if a model can use Flashinfer MoE, which is error prone and not easily extendable. + +### C25 β€” Reviewer 1 + +**Anchor:** `routing method` + +> @siyuanf@nvidia.com @jiahanc@nvidia.com What do you think? + +### C26 β€” Reviewer 5 + +**Anchor:** `routing method` + +> Having a decoupled routing configuration looks reasonable. However, in the CPP side, we probably need to preserve the enum list to align with Trtllm. Additionally, the user should be able to directly provide the topk_ids, in order to support custom score functions or DeepEP. + +### C27 β€” Reviewer 6 + +**Anchor:** `routing method` + +> @siyuanf@nvidia.com 's point is valid. I agree have a dataclass instead of just enumeration, but we should keep the enumeration in cpp to align with trtllm + +### C28 β€” Reviewer 8 + +**Anchor:** `dtype (fp4/fp8/bf16),` + +> nit: it should reflect both activation and weight dtypes + +### C29 β€” Reviewer 10 + +**Anchor:** `backends = [TrtllmFp4Config(), TrtllmFp8BlockConfig(), CutlassConfig()]` + +> I'm not sure how realistic this is since usually each backend requires different weight processing at model load time or different checkpoints entirely. + +### C30 β€” Reviewer 7 + +**Anchor:** `backends = [TrtllmFp4Config(), TrtllmFp8BlockConfig(), CutlassConfig()]` + +> Agree with Reviewer 10: these backends are different operations with different inputs and potentially different outputs. I'm not sure there's a way of hiding that from the caller. + +### C31 β€” Reviewer 8 + +**Anchor:** `backends = [TrtllmFp4Config(), TrtllmFp8BlockConfig(), CutlassConfig()]` + +> Agree, is the framework still defining and calling weight processing methods like swizzling, shuffling etc? + +### C32 β€” Reviewer 8 + +**Anchor:** `backends = [TrtllmFp4Config(), TrtllmFp8BlockConfig(), CutlassConfig()]` + +> >I'm not sure there's a way of hiding that from the caller.It would be nice to have something like `moe_wrapper.process_weights_after_loading(layer.weight, ... , Backend=<>`1 total reactionReviewer 10 reacted with βž• at 2026-05-12 19:27 PM + +### C33 β€” Reviewer 3 + +**Anchor:** `backends = [TrtllmFp4Config(), TrtllmFp8BlockConfig(), CutlassConfig()]` + +> sorry it was not meant to be hidden. the doc got a bit old but i sync'd back the idea from the WIP PR. please see "Example Overview" + +### C34 β€” Reviewer 7 + +**Anchor:** `layer = MoELayer(**config)` + +> Do we need to expose layers without implementations? Could we directly pass tensors to the constructor, or have a utility function that returns the "tuned" layer directly? This way no need to track "has this layer been tuned already". + +### C35 β€” Reviewer 3 + +**Anchor:** `layer = MoELayer(**config)` + +> i see your point + +### C36 β€” Reviewer 2 + +**Anchor:** `layer = MoELayer(**config)` + +> Yes I like the idea of an explicit "compile" api. Maybe we could do something similar to CuTe-DSL, though I don't have enough CuTe-DSL experience to know if that is a good design or not (or if it matches how FI/frameworks work) + +### C37 β€” Reviewer 11 + +**Anchor:** `Config` + +> when I instantiate a config, do I have guarantee that it is supported? or is that still "trial and error"? + +### C38 β€” Reviewer 11 + +**Anchor:** `Config` + +> nice I see `find_backend` in the later section + +### C39 β€” Reviewer 4 + +**Anchor:** `Config is pure data β€” frozen dataclasses, no behavior, serializable via repr()` + +> My concern was mainly like bug reports where the fix lands on a newer version, but agreed that same version repro covers the majority case. + +### C40 β€” Reviewer 2 + +**Anchor:** `routing method` + +> As Reviewer 5 mentioned, an important thing I would consider is ease of supporting custom routing methods. e.g. when DeepSeek released it was a lot of work to bring up the new routing method because it was originally fused. If frameworks are dependent on flashinfer to supply their routing methods this limits their flexibility in enabling new/experimental/research models.I am not sure what the best approach is, since we still want to support all the various fusions, but maybe a "RoutingMethodType.Custom" that accepts a python function would allow users to inject custom scoring functions at the appropriate place, without having to reimplement the routing themselves.Im partial to the API from TRT-LLM where we have a BaseMoeRoutingMethod interface the user can override, but I understand that probably doesnt play well with trtllm backend with its fusions + +### C41 β€” Reviewer 11 + +**Anchor:** `RoutingConfig` + +> does the interface consider routed MoE as well? + +### C42 β€” Reviewer 5 + +**Anchor:** `method: RoutingMethodType = RoutingMethodType.Default` + +> Kindly remind that we now have this file to store enum: flashinfer/tllm_enums.py. It's desirable to put new enums there + +## 10. Codex Reviews + +Review lens: the current WIP is the MVP for PR #3093, not the full long-range design. The MVP target is NVFP4 only, CuteDSL MoE plus TRTLLM-Gen MoE only, pre-routed inputs only, cross-backend autotuning, CUDA graph tests, and benchmarks. + +### MVP Gaps In Current WIP + +**CR1 β€” Config/test API drift.** The current implementation exposes `QuantVariant` and explicit `BackendOptions(candidates=...)`, while `tests/moe/test_moe_api.py` still imports `QuantDtype`, `QuantGranularity`, and `Fp8Variant`, and uses `TrtllmFp4Config() | CutlassConfig()`. This is not a request to widen the MVP; it is a request to make the MVP API and its CPU config tests agree. + +**CR2 β€” Backend preparation is proven but not yet first-class.** `MoEActivationPack` and `MoEWeightPack.prepare_for(...)` give the MVP a clean canonical input shape, but the actual TRTLLM NVFP4 preparation logic still lives in benchmark/test helpers. The MVP goal says per-backend prepare funcs should handle differences, so the shared CuteDSL/TRTLLM NVFP4 prepare entrypoints should move into the implementation surface. + +**CR3 β€” TRTLLM expert offset is not wired through pack input conversion.** `ExpertConfig.local_expert_offset` is stored in config and passed to the TRTLLM kernel, but `TrtllmFp4RoutedRunner.pack_inputs(...)` defaults `local_expert_offset=0`, and `MoELayer` calls it without passing the config offset. For expert-parallel pre-routed inputs, the packed top-k IDs can be wrong when the local shard does not start at expert 0. + +**CR4 β€” Winner caching is single-shape unless the caller rebuilds the layer.** `MoELayer` caches a single `_winner` and `_winner_tactic` after the first call. If the same layer instance is reused for a different token count within `tune_max_num_tokens`, it will not re-run cross-backend selection. The MVP can either make this explicitly one-layer-per-shape/bucket or make the cache shape/bucket-aware. + +**CR5 β€” Token ceiling is only partially threaded into tuning.** `ExecutionConfig.tune_max_num_tokens` gates `MoELayer.__call__`, but the TRTLLM runner's dynamic tuning buckets are still built with an 8192-token ceiling. The MVP benchmark script includes a 16384-token sweep, so the runner tuning config and benchmark expectations need to agree. + +**CR6 β€” MVP scope should fail fast.** The config objects describe more than the MVP can execute, while `MoELayer` silently skips non-MVP backend configs and the current TRTLLM runner hard-codes NVFP4/SwiGLU assumptions internally. For the MVP, unsupported quant variants, activation variants, backend choices, and non-pre-routed call shapes should produce explicit errors. + +### Implementation Innovations To Fold Back Into The Doc + +**CR7 β€” Activation and weight packs are sharper than the original `MoETensors` sketch.** The WIP splits per-call activation/routing data (`MoEActivationPack`) from long-lived backend-native weight materializations (`MoEWeightPack`). That directly addresses reviewer concerns about backend-specific preprocessing without pretending all backends share one weight layout. + +**CR8 β€” Cross-backend autotune is a two-stage decision.** Each runner first uses the existing `AutoTuner` to select its best tactic, then `MoELayer` measures each runner at its winning tactic and caches the fastest backend. The original doc says "autotune across backends" but does not spell out this tactic-then-backend selection structure. + +**CR9 β€” Minimal runtime introspection already exists.** `winner_backend` and `reset_winner()` are practical MVP observability hooks. They are smaller than a full backend discovery API but useful for benchmarks, debugging, and validating winner selection. + +**CR10 β€” The WIP tests are more concrete than the design doc.** `tests/moe/test_unified_moe_api.py` checks the selected `MoELayer` output against a shared BF16 reference, checks each candidate backend against the same reference, verifies autotune visits all candidates, and checks CUDA graph replay against eager output. + +**CR11 β€” The benchmark path reports candidates, not only the selected layer.** `unified_nvfp4_moe` builds both backend-native weight views on one `MoEWeightPack`, triggers winner selection, then emits one result row per candidate backend with winner metadata. That is useful MVP evidence and should be reflected in the doc. + +## 11. Task Tracking + +This tracker is scoped to the PR #3093 MVP, not the full long-range API design. Its purpose is continuity: what already landed in the current branch, what still needs tightening before review, and what should stay out of scope. + +> **Status (2026-05-31): MVP scope complete.** Every "Today's MVP Cut" item and every "Remaining MVP Follow-Up" (CR1–CR11) is done and validated on a B200 (SM100): the unified-MoE GPU tests **9/9** (layer + per-backend accuracy vs bf16, autotune visits both candidates, CUDA-graph replay), the CPU config tests **97/97** (CPU config + fail-fast validation), and the `unified_nvfp4_moe` benchmark sweep (128β†’16384 tokens) with `--refcheck` passing for both backends. Only **Post-MVP Carryover** and **Explicit Non-Goals** remain open by design. + +> **Status (2026-06-01): PR #3093 review-comment pass.** Addressed the open GitHub review threads without widening MVP scope. Bot threads: the redundant `pack_inputs(local_expert_offset=...)` was already fixed (reads from config β€” CR3); removed the eval-based `MoEConfig.from_repr` (eval is a security smell and repro/serialization is documented Post-MVP β€” see C4-C5/C39); broke the import-time `layer.py β†’ testing.utils` coupling by lazy-importing `bench_gpu_time` on the autotune path only. Human threads (benchmark merge damage): `benchmarks/routines/moe.py` was reconciled against `main` so its diff is now **purely additive** (the `unified_nvfp4_moe` routine) β€” the accidentally-dropped `bgmv_moe` routine, `warn_if_pdl_unsupported`, and all `enable_pdl=args.enable_pdl` call-sites are restored; the `unified_nvfp4_moe` arch table was trimmed to the SM100 entries it actually supports (`10.0`/`10.3`); the sweep script moved to `benchmarks/bench_unified_moe.sh`; and the two MoE test files were merged + renamed to `tests/moe/test_unified_moe.py` (96 CPU config tests after dropping the `from_repr` round-trip + 9 SM100 GPU tests). Re-validated on B200: benchmark/library imports clean, **96 CPU config tests pass**, full file collects 105 tests. (Nightly/dashboard tracking of the sweep script is left as a separate CI-infra follow-up.) + +### Release Gates (do before this ships in a tagged release) + +The branch can merge to `main` early for team review and may land in a nightly/early release. To avoid implying any stability/observability commitment on a still-evolving API surface, the MVP **intentionally ships the new unified MoE APIs without the `@flashinfer_api` decorator** (no logging / repro-trace / stability contract). This is deliberate β€” not an oversight β€” and reserves the right to change `MoEConfig` / `MoELayer` / `MoEActivationPack` / `MoEWeightPack` / the runners / `prepare_weights` freely pre-release. + +| Status | Gate | Notes | +| --- | --- | --- | +| [ ] | Add `@flashinfer_api` (+ a `TraceTemplate` per the `CLAUDE.md` "Trace Template Checklist") to the public unified MoE APIs **at release time**, not before. | The decorator carries logging/repro + an implied stability contract; Β§4.1/Β§6 describe the intended end-state. The decorated legacy MoE functions (`trtllm_*_moe`, `cutlass_fused_moe`) already ship in v0.6.12 and are untouched here. | + +### Landed In Current Branch + +| Status | Task | Continuity notes | +| --- | --- | --- | +| [x] | Add the MVP config and pack surface. | `MoEConfig`, component configs, `MoEActivationPack`, and `MoEWeightPack` exist in `flashinfer/fused_moe/api.py`. | +| [x] | Add two MVP backend runners. | `CuteDslNvfp4Runner` and `TrtllmFp4RoutedRunner` exist as `TunableRunner` adapters in `flashinfer/fused_moe/runners.py`. | +| [x] | Add cross-backend `MoELayer` dispatch. | `MoELayer` builds compatible runners, selects a winner, caches it, and exposes `winner_backend` plus `reset_winner()`. | +| [x] | Preserve legacy flat MoE APIs. | `flashinfer/fused_moe/__init__.py` exports the new MVP API while keeping the existing flat APIs available. | +| [x] | Add unified NVFP4 benchmark path. | `unified_nvfp4_moe` is wired into the benchmark registry; the sweep script is `benchmarks/bench_unified_moe.sh`. | +| [x] | Add MVP accuracy, autotune, and CUDA graph tests. | `tests/moe/test_unified_moe.py` covers CPU config tests plus shared-reference accuracy, candidate visitation, and graph replay. | + +### MVP As-Built Reference + +The aspirational API in Β§2–§4 (eager `moe_layer(...)`, `MoETensors`, `find_backends`, pipe-operator backends) describes the long-range design. What actually shipped for the PR #3093 MVP is narrower and pack-based; this section is the authoritative end-to-end description of the as-built surface (CR7–CR9). + +```python +import torch +from flashinfer.fused_moe import ( + MoEConfig, RoutingConfig, QuantConfig, QuantVariant, ExpertConfig, + ActivationConfig, ExecutionConfig, MoELayer, + MoEActivationPack, MoEWeightPack, CuteDslConfig, TrtllmFp4Config, +) +from flashinfer.fused_moe.api import BackendOptions +from flashinfer.autotuner import autotune + +# 1. Config β€” single-knob QuantVariant; explicit candidate set. +config = MoEConfig( + routing=RoutingConfig(num_experts=32, top_k=2), + quant=QuantConfig(variant=QuantVariant.NVFP4), # MVP: NVFP4 only + experts=ExpertConfig(intermediate_size=512, local_num_experts=32), + activation=ActivationConfig(), # MVP: Swiglu only + backend=BackendOptions(candidates=(CuteDslConfig(), TrtllmFp4Config())), + execution=ExecutionConfig(tune_max_num_tokens=8192), +) + +# 2. Long-lived weights: one MoEWeightPack holds a backend-native view per +# backend, built from canonical bf16 weights by first-class prepare helpers. +weights = MoEWeightPack() +weights.prepare_for("cute_dsl_nvfp4", + CuteDslConfig.prepare_weights(w1_bf16, w2_bf16, num_local_experts=32, + hidden_size=1024, intermediate_size=512)) +weights.prepare_for("trtllm_fp4_routed", + TrtllmFp4Config.prepare_weights(w1_bf16, w2_bf16, num_local_experts=32, + hidden_size=1024, intermediate_size=512)) + +# 3. Per-call activations: pre-routed (selected_experts/final_scales supplied). +act = MoEActivationPack( + hidden_states_q=x_q, # [M, H//2] uint8 packed NVFP4 + hidden_states_scale=x_sf, # [M, H//16] float8_e4m3fn (or uint8 bytes) + selected_experts=topk_ids, # [M, top_k] int32 + final_scales=topk_weights, # [M, top_k] float32 +) + +# 4. Dispatch. First call per token-bucket runs cross-backend selection. +layer = MoELayer(config) +with autotune(True): + out = layer(act, weights) # tunes + selects winner for this bucket +print(layer.winner_backend) # e.g. "cute_dsl_nvfp4" +out = layer(act, weights) # subsequent calls: cached winner dispatch +``` + +Key mechanisms (and where they live): + +- **Two packs, two lifetimes.** `MoEWeightPack` holds long-lived, backend-native weight materializations keyed by `backend_key` (`prepare_for` / `get_view`); `MoEActivationPack` carries per-call pre-routed activations. This is the concrete answer to reviewers' "backends need different weight preprocessing" concern (C29–C32): each backend stores its own view, none is hidden from the caller. +- **First-class prep.** `TrtllmFp4Config.prepare_weights(...)` / `CuteDslConfig.prepare_weights(...)` (backed by `flashinfer/fused_moe/prepare.py`) turn canonical bf16 weights into the native views (C6/C7). +- **Two-stage cross-backend autotune** (`MoELayer._select_winner`, runners' delegation): for each candidate, the `AutoTuner.choose_one` picks the best *within-backend tactic* (each backend tuned in its own native input schema), then `bench_gpu_time` compares the candidates at their winning tactics and the fastest backend is dispatched. A single `choose_one` over both runners is not possible because their input schemas differ β€” hence the explicit two stages. +- **Winner caching is per token-bucket** (`map_to_hybrid_bucket`): reusing one `MoELayer` across token counts re-selects per bucket; `winner_backend` reports the most-recent choice and `reset_winner()` clears the cache. +- **Fail-fast scope** (`MoELayer._validate_mvp_scope`): non-NVFP4 quant or non-Swiglu activation raises `NotImplementedError` at construction. +- **Runners delegate** to canonical inner runners (`CuteDslFusedMoENvfp4Runner` / `core.MoERunner`); the unified adapters only translate Packs ⇄ the inner runner's native tensor list. + +### Today's MVP Cut + +This is the May 27, 2026 working slice (executed May 31, 2026). It should improve the current PR without expanding it beyond NVFP4, CuteDSL plus TRTLLM-Gen, pre-routed inputs, cross-backend autotune, CUDA graph tests, and benchmark evidence. + +| Priority | Status | Task | Why it fits today | +| --- | --- | --- | --- | +| P0 | [x] | Keep local `moe_api`, `origin/moe_api`, and PR #3093 on the same head commit before editing. | Avoids iterating on a stale branch or accidentally reviewing a different PR state. | +| P1 | [x] | Align the CPU config tests with the actual MVP API surface. | Fast, no GPU required, and removes obvious API/test drift before deeper validation. | +| P1 | [x] | Wire `local_expert_offset` into TRTLLM routed top-k packing and add a focused pre-routed EP-offset test. | Small correctness fix inside MVP scope; prevents wrong packed expert IDs for nonzero local shard offsets. | +| P1 | [x] | Add fail-fast MVP validation for quant variant, activation assumptions, backend set, and pre-routed-only inputs. | Answers the config-support review concern without building a full backend discovery API today. | +| P2 | [x] | Decide and document the `MoELayer` reuse contract for this PR: one layer per tuned shape/bucket, or shape-aware winner cache. | Avoids a hidden behavioral trap while keeping implementation scope explicit. | +| P2 | [x] | Thread `tune_max_num_tokens` into runner tuning configs enough to make the current benchmark sweep honest. | Needed if the 16K-token row remains in the MVP evidence path. | +| P2 | [x] | Run the unified NVFP4 benchmark sweep on the intended GPU and record winner/per-candidate latency evidence. | Converts the branch from "implemented" to "PR argument is supported by measurements." | + +> **Mid-cut discovery (blocker, resolved).** While validating on B200, both MVP runner adapters turned out to have *never* run against the post-`main`-merge `core.py`: `TrtllmFp4RoutedRunner` targeted a raw-`moe_op` API the module factory does not expose, and `CuteDslNvfp4Runner` read its `tuning_config` off the class and under-populated its input list. The runners are listed as "Landed" above, but the layer / autotune / CUDA-graph tests could not actually execute. Fixing this was a prerequisite for the P2 evidence items and is recorded in the Decision Log below; it stayed within MVP scope (no new backends, dtypes, or routing modes). + +#### Decision Log β€” May 31, 2026 working slice + +Decisions made while executing the cut above, recorded so reviewers see the *why*, not just the diff. + +- **P0 β€” branch alignment verified.** Local `moe_api`, `origin/moe_api`, and PR #3093 head all resolve to the same commit (`1f74494b` at the time of writing), so the cut edits the live PR state. The most recent prior change on the branch (`fix(fused_moe): align TrtllmFp4RoutedRunner with hybrid token buckets`) is already reflected in `runners.py`. +- **P1 / CR1 β€” single-knob `QuantVariant`, explicit `BackendOptions(candidates=...)`.** The CPU config tests (`tests/moe/test_moe_api.py`) were rewritten to match the implementation rather than the other way around. Rationale: the implementation deliberately collapsed the older `QuantDtype` + `QuantGranularity` + `Fp8Variant` triple into one `QuantVariant` enum (`NVFP4`, `MxFp8`, `DeepSeekFp8`, `FP8PerTensor`, `MxInt4`, `MXFP4`, `BF16`). One knob is simpler for the MVP and still distinguishes the cases reviewers flagged (C14 MXFP4 block size, C28 activation+weight dtype) because each becomes a distinct enum member. The `|` pipe-operator sugar and a richer multi-field `QuantConfig` are listed as Explicit Non-Goals for this MVP, so the canonical spelling is the explicit `BackendOptions(candidates=(...))` already used by the GPU test (`tests/moe/test_unified_moe_api.py`) and the benchmark. Verified: 84 CPU tests pass in the B200 container. +- **Runner rework (blocker fix) β€” delegate to the canonical inner runners.** `TrtllmFp4RoutedRunner` now wraps `core.MoERunner` (newly exported from `get_trtllm_moe_sm100_module()`), mirroring how `CuteDslNvfp4Runner` wraps `CuteDslFusedMoENvfp4Runner`. `pack_inputs` builds the `MoEInputs` list (with an allocated output buffer and the kernel-required `topk_weights` placeholder for `PackedPrecomputed`) plus a static weight/config kwargs dict; `forward`/`get_valid_tactics` delegate to the inner runner, which owns the one fragile raw-op launch. This keeps the unified adapters thin and resistant to future `core.py` signature drift. The CuteDSL adapter additionally appends the optional `moe_output` buffer (index 11) its tuning_config declares as dynamic. The nvfp4 activation scale is viewed to `float8_e4m3fn` (the canonical Pack may carry raw `uint8` bytes; trtllm-gen accepts the *linear* scale layout, so no per-call swizzle is needed). Validated on B200: all 9 `tests/moe/test_unified_moe_api.py` pass. +- **P1 / CR3 β€” offset read from config, not a dead parameter.** `pack_inputs` no longer takes a `local_expert_offset` argument (no caller ever passed it, so it silently defaulted to 0); it reads `ExpertConfig.local_expert_offset` off the runner's own config. A focused SM100 test (`TestTrtllmEPOffset`) decodes the packed ids and asserts they land in the kernel's local `[0, local_num_experts)` range for offsets 0/32/96. +- **P1 / CR6 β€” fail fast at construction.** `MoELayer._validate_mvp_scope` raises `NotImplementedError` for any non-`NVFP4` quant variant or non-`Swiglu` activation, and the "no usable backend" error now names the MVP-supported backend set. Pre-routed-only is structural (the layer consumes `MoEActivationPack`, which carries `selected_experts`/`final_scales`). Covered by CPU tests in `TestMoELayerMVPValidation`. +- **P2 / CR4 β€” bucket-keyed winner cache.** The cross-backend winner can legitimately differ across token-count buckets (the per-tactic autotuner is already bucket-aware), so `MoELayer` now caches `(runner, tactic)` keyed by `map_to_hybrid_bucket(num_tokens, tune_max_num_tokens)` instead of a single `_winner`. Reusing one layer across token counts re-selects correctly per bucket; `winner_backend` reports the most-recent call's choice and `reset_winner()` clears all buckets. This removes the silent stale-winner trap without forcing one-layer-per-shape on callers. +- **P2 / CR5 β€” token ceiling threaded for free.** Because the reworked TRTLLM runner builds its tuning config via `MoERunner._make_tuning_config(tune_max_num_tokens=ExecutionConfig.tune_max_num_tokens)`, the num_tokens buckets now honor the configured ceiling, so a 16384-token sweep tunes against 16384-token buckets rather than a hard-coded 8192. +- **CR2/CR7 β€” first-class NVFP4 weight prep.** Added `flashinfer/fused_moe/prepare.py` (`prepare_trtllm_fp4_weights`, `prepare_cute_dsl_nvfp4_weights`) and exposed them as `TrtllmFp4Config.prepare_weights` / `CuteDslConfig.prepare_weights` so callers do `weight_pack.prepare_for("trtllm_fp4_routed", TrtllmFp4Config.prepare_weights(w1, w2, ...))` (the workflow reviewer C7 described). Deleted the byte-identical `_build_trtllm_view` / `_build_trtllm_nvfp4_view` copies from the test and benchmark; the benchmark now builds *both* backend views from the same bf16 weights via the helpers. The `prepare_weights` staticmethods lazily import the heavy prep so `api.py` stays pure-data. Validated: test suite 9/9 (TRTLLM view), benchmark runs both candidates (CuteDSL view). Activation-scale prep remains a smaller follow-up. +- **P2 β€” benchmark evidence (B200, SM100).** `benchmarks/flashinfer_benchmark.py --routine unified_nvfp4_moe` referenced an undefined `_create_cute_dsl_moe_test_data`; pointed it at the canonical `create_moe_tensors` (the same helper the GPU test uses). Sweep at `hidden=1024, intermediate=512, num_experts=32, top_k=2`, NVFP4 + Swiglu, CUDA-graph timing (CUPTI unavailable β†’ CUDA events). One row per candidate; `*` marks the cross-backend winner. The 16384-token row exercises the CR5 ceiling. + + | num_tokens | winner | cute_dsl_nvfp4 (ms / TFLOPΒ·s⁻¹) | trtllm_fp4_routed (ms / TFLOPΒ·s⁻¹) | + | --- | --- | --- | --- | + | 128 | cute_dsl_nvfp4 | 0.016 / 49.5 | 0.017 / 48.3 | + | 512 | cute_dsl_nvfp4 | 0.018 / 182.1 | 0.020 / 163.2 | + | 2048 | cute_dsl_nvfp4 | 0.023 / 554.8 | 0.033 / 392.2 | + | 8192 | cute_dsl_nvfp4 | 0.040 / 1274.9 | 0.067 / 766.3 | + | 16384 | cute_dsl_nvfp4 | 0.067 / 1546.8 | 0.108 / 954.3 | + + CuteDSL wins across the swept range for this geometry; the cross-backend selection, per-candidate latency, and winner introspection (`winner_backend`) all flow through to CSV/stdout as the benchmark intends (CR10/CR11). +- **CR10/CR11 β€” `--refcheck` for the unified routine.** Because both backend views now derive from the same bf16 weights, the benchmark can verify each candidate against one `compute_reference_moe_fp4` bf16 reference. With `--refcheck`, each row prints `[REFCHECK] unified/: PASS/FAIL`; a failure errors unless `--allow_output_mismatch`. Validated on B200: both `cute_dsl_nvfp4` and `trtllm_fp4_routed` report 100% within tolerance (atolβ‰ˆ0.13). + +#### Cross-backend autotune value + a selection bug (DeepSeek-V3, B200) + +The whole point of `MoELayer` is to pick the faster backend *per shape*. A DeepSeek-V3 sweep (hidden=7168, intermediate=2048, num_experts=256, top_k=8, NVFP4+Swiglu; EP=1) makes the case β€” and surfaced a real selection bug. + +All numbers below were **regenerated end-to-end (2026-06-01, B200) via `benchmarks/flashinfer_benchmark.py --routine unified_nvfp4_moe`** β€” the perf-tracking driver, not a side harness. One invocation per shape emits both per-candidate `[PERF]` rows *and* the `MoELayer` winner, so a single sweep yields all three comparisons (`benchmarks/bench_unified_moe.sh` drives it). The per-candidate latency *is* that backend's within-backend-autotuned time. + +| num_tokens (EP=1) | cute_dsl_nvfp4 (ms) | trtllm_fp4_routed (ms) | MoELayer winner | regime | +| --- | --- | --- | --- | --- | +| 1 | 0.046 | **0.042** | cute_dsl_nvfp4 † | noise tie | +| 16 | 0.428 | **0.365** | trtllm_fp4_routed | low-latency | +| 64 | 0.938 | **0.822** | trtllm_fp4_routed | low-latency | +| 128 | 1.072 | **0.901** | trtllm_fp4_routed | low-latency | +| 256 | 1.115 | **0.936** | trtllm_fp4_routed | low-latency | +| 512 | 1.134 | **0.942** | trtllm_fp4_routed | low-latency | +| 1024 | **1.199** | 1.306 | cute_dsl_nvfp4 | throughput | +| 2048 | **1.226** | 1.311 | cute_dsl_nvfp4 | throughput | +| 4096 | **1.624** | 1.711 | cute_dsl_nvfp4 | throughput | +| 16384 | **3.742** | 4.613 | cute_dsl_nvfp4 | throughput | + +The winner **flips with a sharp crossover between 512 and 1024 tokens**: TRTLLM-gen wins the entire low-latency regime (16–512 tokens, ~15–17% faster) β€” consistent with its known small-batch specialization (cf. PR #2529) β€” while CuteDSL wins large-batch throughput (β‰₯1024, up to ~19% faster at 16384). Neither single-backend strategy dominates, so cross-backend autotune is β‰₯ either backend alone and strictly faster wherever the other clearly loses. Both backends pass `--refcheck` at this geometry (t=1024: 100% within tol vs the bf16 reference), so the comparison is between two numerically-correct implementations β€” not one that is fast because it is wrong. + +† **t=1 is a noise tie.** At ~40-Β΅s kernels the two backends are within ~4 Β΅s (~9%); the selector's internal timing picked CuteDSL on this run while the benchmark's independent per-candidate re-timing shows TRTLLM marginally faster, and the pick flips run-to-run (an earlier cross-check picked TRTLLM at t=1). The ~4 Β΅s cost of the "wrong" choice at the smallest shape is negligible. This is the measurement noise floor, distinct from the systematic bug below. + +**Selection bug found & fixed.** An earlier sweep mis-picked the *slower* backend even at well-separated shapes (e.g. EP1 t=1024 picked TRTLLM though CuteDSL was clearly faster). Cause: `MoELayer._select_winner` timed candidates with a no-CUDA-graph 10-iter `bench_gpu_time`, so at low token counts launch/Python overhead dominated the median. Fix: time the selection with CUDA graph + 30 iters (matching deployment and the benchmark's own per-candidate timing). After the fix the winner tracks the faster backend at every well-separated shape; only genuine near-ties (t=1, and tβ‰ˆ4096 where the gap is a few %) remain coin-flips, as expected. (Requires a warmed-up layer β€” the autotune pass β€” not a cold graph capture.) + +**The optimal backend depends on geometry *and* batch β€” which is the whole motivation.** The small-geometry sweep above (hidden=1024, intermediate=512, 32 experts) has CuteDSL winning at *every* token count, whereas the DeepSeek-V3 geometry (hidden=7168, 256 experts) hands the entire ≀512-token regime to TRTLLM-gen. So there is no fixed "use backend X" rule even per-batch-size β€” the right choice moves with the full problem shape. A per-shape cross-backend selector is therefore the only way to stay on the frontier without hand-tuning a routing table, which is exactly what `MoELayer` automates. + +**Wide-EP: local-only MVP proxy (realistic EP deferred).** The first EP=16 sweep was unfaithful β€” it fed *global* routing ids (range 256) against only 16 local experts at `offset=0`, so the kernel skipped most tokens (implausibly low time) and the metrics over-counted weight bytes (impossible >50 TB/s). Fixed for the MVP by routing the activation *within* the local experts (`selected ∈ [0, local_num_experts)`), modeling a single rank as a complete MoE over its local experts. Every token is now computed locally, so latencies are real and the derived metrics are correct. Validated (DeepSeek-V3, local=16): + +| EP16 tokens | CuteDSL (ms) | TRTLLM-gen (ms) | winner | bandwidth | +| --- | --- | --- | --- | --- | +| 1 | 0.045 | 0.045 | ~tie | ~4.4 TB/s | +| 16 | 0.077 | **0.073** | trtllm_fp4_routed | ~5.5 TB/s | +| 4096 | **0.696** | 0.975 | cute_dsl_nvfp4 | ~0.68 TB/s | + +Bandwidth is now physically sane (≀6 TB/s) and both backends pass `--refcheck` at every shape. **Realistic wide-EP** β€” global top-k-of-N routing with cross-rank dispatch and the resulting per-rank load imbalance β€” is **out of scope for this PR** and tracked for the separate follow-on `moe_ep` API PR (see Post-MVP Carryover). The building blocks already exist: `compute_reference_moe_fp4` accepts `num_local_experts`/`local_expert_offset` and skips non-local tokens, and `bench_moe_deepseek.py` scales work by `local_fraction = num_local_experts/num_experts` (uniform-distribution assumption); a faithful version would feed each rank only its dispatched tokens rather than assume uniformity. + +#### Legacy-vs-unified kernel equivalence (cross-check, 2026-06-01) + +Diligence requested in review: confirm the unified benchmark measures the *same* +underlying trtllm-gen kernel as the legacy flat routine β€” not a different or +no-op path. Both go through `get_trtllm_moe_sm100_module()`; the legacy +`trtllm_fp4_block_scale_moe` is fed routing logits (routing runs *inside* the +kernel), whereas the unified `trtllm_fp4_routed` uses +`RoutingInputMode.PackedPrecomputed` (pre-routed β†’ GEMM/activation/finalize +only). Same DeepSeek-V3 geometry (hidden=7168, intermediate=2048, 256 experts, +top_k=8, n_group=8, topk_group=4, routed_scaling_factor=2.5), B200, CUDA-graph +timing. + +| num_tokens | legacy `trtllm_fp4_block_scale_moe` | unified `trtllm_fp4_routed` | Ξ” (unified βˆ’ legacy) | unified `cute_dsl_nvfp4` | MoELayer winner | +| --- | --- | --- | --- | --- | --- | +| 1 | 0.047 ms | 0.041 ms | βˆ’13% | 0.045 ms | trtllm_fp4_routed | +| 1024 | 1.412 ms | 1.300 ms | βˆ’8% | 1.199 ms | cute_dsl_nvfp4 | + +The unified routed path tracks the legacy kernel to within ~8–13% and is +consistently *slightly faster* β€” by ~the in-kernel routing cost it legitimately +skips (β‰ˆ6 Β΅s at t=1, β‰ˆ0.11 ms at t=1024). That is the expected signature of "same +GEMM kernel minus routing," not a discrepancy; a no-op/half-pipeline would show a +5–10Γ— gap. The per-shape winner flip (trtllm at t=1, CuteDSL at t=1024) +independently reproduces the DeepSeek crossover above. **Conclusion:** the +`unified_nvfp4_moe` benchmark measures the real kernel β€” functionality confirmed. + +Repro: `benchmarks/flashinfer_benchmark.py --routine {trtllm_fp4_block_scale_moe, unified_nvfp4_moe}` at the geometry above (the legacy routine adds `--use_routing_bias --routing_method deepseek_v3 --use_shuffled_weight`). + +### Remaining MVP Follow-Ups + +| Status | Task | Review refs | +| --- | --- | --- | +| [x] | Align the MVP config API and config tests: `tests/moe/test_moe_api.py` was rewritten to `QuantVariant` plus explicit `BackendOptions(candidates=...)`; the `QuantDtype`/`Fp8Variant`/pipe-operator spelling was dropped (see Decision Log). | CR1 | +| [x] | NVFP4 **weight** prep is now first-class: `flashinfer/fused_moe/prepare.py` provides `prepare_trtllm_fp4_weights` / `prepare_cute_dsl_nvfp4_weights`, exposed as `TrtllmFp4Config.prepare_weights` / `CuteDslConfig.prepare_weights` (per C7). The duplicated `_build_trtllm_view` (test) and `_build_trtllm_nvfp4_view` (benchmark) are removed; both call sites use the helper. **Activation** prep is the remaining slice: the Pack carries one `hidden_states_scale`, and trtllm-gen happily consumes the linear-layout scale (the runner only re-`view`s `uint8`β†’`float8_e4m3fn`), so a first-class activation-prep helper (and a future swizzled-activation backend) is the leftover follow-up. | CR2, CR7 | +| [x] | Wire `ExpertConfig.local_expert_offset` into TRTLLM `pack_inputs(...)` (read from config) and add an EP-offset test (`TestTrtllmEPOffset`) for the pre-routed path. | CR3 | +| [x] | Layer reuse contract decided: bucket-keyed winner cache (`map_to_hybrid_bucket`), so reuse across token counts re-selects per bucket. See Decision Log. | CR4 | +| [x] | `ExecutionConfig.tune_max_num_tokens` is threaded into the TRTLLM runner tuning config via `MoERunner._make_tuning_config`; benchmark-sweep validation is the remaining P2 evidence item. | CR5 | +| [x] | Added `MoELayer._validate_mvp_scope` (NVFP4 + Swiglu fail-fast) and a clearer no-usable-backend error; pre-routed-only is structural via `MoEActivationPack`. Covered by `TestMoELayerMVPValidation`. | CR6, C37-C38 | +| [x] | Documented the as-built MVP in the new "MVP As-Built Reference" subsection: end-to-end example plus `MoEActivationPack` / `MoEWeightPack`, first-class `prepare_weights`, two-stage cross-backend autotune, per-bucket winner caching, and `winner_backend` / `reset_winner` introspection. | CR7-CR9 | +| [x] | `unified_nvfp4_moe` runs end-to-end and emits `winner_backend` + per-candidate latency (one row per candidate; see the Decision Log evidence table) and now supports `--refcheck`: each candidate is verified against a shared bf16 reference (both views derive from the same bf16 weights). Validated on B200 β€” both backends 100% within tolerance. | CR10-CR11 | + +### PR #3093 Review Threads β€” Reviewer Pass 2 (2026-06-02/03) + +> **Status (2026-06-09): all resolved.** The second human-reviewer pass (G1–G7; +> `G1–G6` Reviewer 12 on `flashinfer/fused_moe/api.py`, `G7` Reviewer 13) is fully +> addressed and pushed β€” the bot threads (CodeRabbit Γ—6, Gemini Γ—3) and earlier +> self-notes were resolved in the 2026-06-01 pass. The per-thread decisions and +> rationale live in the commits (`fix(moe): … review comments` updates 1–3) and, +> for the structural threads, in the code + sections above: G1 enum +> consolidation (the enum block in `api.py` / `tllm_enums.py`), G4/G5 the +> `MoETensors` cluster drop ("Two packs, two lifetimes" + the pack rationale in +> `api.py`). The Reviewer 14 `_select_winner` thread remains tracked under Post-MVP +> Carryover (a deferred design decision, not part of this pass). The *open* MoE +> work is now the two fuzzer-filed bugs β€” gh #3547 / #3548 β€” described under +> "Test Harness" above. + + +### Test Harness β€” Forward-Compatible Fuzzer (PR #6, merged 2026-06-09) + +`tests/moe/test_unified_moe_fuzz.py` (merged from `aleozlx/flashinfer#6`, +branch `yanxu/unified-moe-api-fuzzer`) drives the **real user-facing surface** β€” +one `MoEConfig` β†’ `XxxConfig.prepare_weights(w1_bf16, w2_bf16, …)` β†’ +`MoELayer`'s per-backend runners β€” so the production dispatch + the `prepare.py` +scale/layout plumbing are what's under test, where low-precision-MoE bugs +cluster. + +**Forward-compatible by construction:** +- Backends are **auto-discovered from the live runner registry** (`layer.runners`); + an unwired backend is skipped and gets covered the moment its runner lands β€” + zero new test code. +- Weight prep is the uniform `cfg.prepare_weights(...)` (canonical bf16 in, + quantize+layout internal). +- Per-dtype specifics live in one `_DTYPE` table (golden-input snap / activation + pack / canonical reference / poison / tolerances). New dtype = one + `DTypeHandler`; new backend = free. + +**Verification model** (uniform per config): (1) no crash / no NaN-Inf where the +reference is finite; (2) numeric agreement vs a **single authoritative +quant-aware reference** β€” inputs snapped to the exact nvfp4 grid + sparsified so +a structural bug (dropped expert / wrong index / wrong scale role) is a gross +error, tolerance at the fp4 requant floor (~0.08); (3) per-backend determinism +(bitwise reproduce unless declared non-deterministic, e.g. CuteDSL atomic +finalize); (4) output-buffer poison (garbage+NaN/Inf in the kernel's `new_empty` +output β†’ the torchβ†’JAX buffer-hygiene guard); (5) autotune-tactic sweep (every +valid tactic matches, not just the default); (6) autotune-ON real path +(`autotune(True)` profiles+selects+caches a winner, output still matches); (7) +device-state probe (turns a context-corrupting IMA into a clean failure). A +sibling `test_autotune_cache_coherence` scenario covers the cross-call winner +cache (token-count sequence across bucket boundaries 4095/4096/4097). +Cross-backend agreement is **intentionally not** a check β€” an authoritative +tight reference already catches (and names) a deviating backend. + +**Config space:** random non-pow2 (aligned) hidden/intermediate, odd/tile-boundary +token counts, routing-load skew, and ~30% **expert-parallel shards** (global > +local + `local_expert_offset` β€” the real deployment shape, in scope for the +single-GPU harness; the EP *collective* is not), all under a weight-memory +budget so one config never hogs the GPU. + +**Known-failure ledger** (`_KNOWN_FAILURES`): a filed-and-tracked bug is `xfail`ed +by `(backend_key, predicate)` β€” the case is **still run**, so the suite stays +green yet flags loudly (`xpass` β†’ "remove this entry") the day the bug is fixed. +A crash is never tolerated, only a wrong answer. + +**CI-safety gate (waived, opt-in).** The ledger tolerates a *wrong answer* but +cannot absorb a *process abort*, and a single-process run of this suite on SM100 +hit `CUDA error: device-side assert triggered` β†’ `Fatal Python error: Aborted` +(triage 2026-06-09) β€” which would block B200 CI. Per-config isolation passes +68/86 incl. EP `offset>0`, so the abort is **not** cleanly attributable to one +config (the #3547 EP case returns tolerated zeros under `synchronize`, no +assert); it surfaces only in the accumulated single-process run CI uses, and +`--forked` can't isolate it (CUDA inits at collection). So the suite is gated +behind `FLASHINFER_UMOE_FUZZ` (`pytestmark` skip): **unset (CI default) β†’ +collected-and-skipped, launches no kernel, cannot abort the job**; set β†’ runs +(developer / nightly). The follow-up PR fixes #3547, root-causes the abort, and +removes the gate. + +**Bugs this fuzzer found + filed** (the EP/scale regimes the prior suite never +exercised end-to-end): +- **gh #3547** β€” `trtllm_fp4_routed` returns all-zeros for EP shards + (`local_expert_offset > 0`): the offset is applied twice (pre-subtracted in + `pack_inputs` *and* forwarded to the kernel). `cute_dsl_nvfp4` is correct + (passes global ids + offset, kernel localizes once). Encoded as the current + `_KNOWN_FAILURES` entry. Fix = stop pre-subtracting (pass global ids), then + delete the ledger entry so the case flips to passing. +- **gh #3548** β€” activation **global-scale** gap: `prepare_*_weights` hardcodes + `gs=1.0`/`fc2_input_scale=1.0`/`alpha=ones` and `MoEActivationPack` has no + global-scale field, so calibrated-checkpoint scales are silently dropped + (~2400Γ— output inflation). This is roadmap item #5 below (a standardized + intermediate-scale **QuantSpec** policy), not a quick fix; quick mitigation is + to make `prepare_*` fail **loud** on a non-default scale. + +**Roadmap (ranked, from the 2026-06-09 audit of 51 past MoE issues):** (1) a +Blackwell/SM120 **PR-CI runner** β€” highest leverage, since PR-gating CI tops out +at SM90 and the dominant fp4/MoE bug class is collected-then-skipped at PR time; +(2) N-run stress + per-test timeout under `--forked`; (3) curated production +shapes (DeepSeek-V3 / Llama-4 / Qwen3 / Mixtral + tile-window enumeration); (4) a +build-manifest oracle (assert each advertised backendΓ—quantΓ—arch instantiates a +kernel); (5) tighten the quantized-numeric net via the QuantSpec scale policy +(also unblocks #3548 and an independent fp32 reference). + +### Post-MVP Carryover + +| Status | Task | Review refs | +| --- | --- | --- | +| [ ] | Design a backend discovery/support-query API that can tell users whether a config is supported without trial-and-error execution. | C16, C18-C19, C37-C38 | +| [ ] | Decide the long-term custom routing extension point, including how routed MoE, caller-provided top-k IDs, and custom scoring functions should compose with fused backends. | C24-C27, C40-C41 | +| [ ] | Keep new routing enums aligned with the shared enum home instead of creating a parallel enum surface in the MoE API. | C42 | +| [ ] | Decide whether repro logs remain same-version-only or need a versioned schema for cross-version bug reports. | C4-C5, C39 | +| [ ] | **Realistic wide-EP** lands in the separate `moe_ep` API (PR #3453): global top-k-of-N routing + cross-rank dispatch/combine + per-rank load imbalance. This MVP only ships a *local-only* per-rank proxy (route within local experts). Coordination seam: `moe_ep`'s `MoEEpLayer.forward` does `dispatch β†’ inner_compute β†’ combine`, and `inner_compute` (identity today) is where this unified `fused_moe` path becomes the per-rank expert compute β€” so our `ExpertConfig.local_expert_offset` wiring and local-only benchmark proxy already model that compute side. Building blocks for faithful EP: `compute_reference_moe_fp4`'s `num_local_experts`/`local_expert_offset` local-skip, `bench_moe_deepseek.py`'s `local_fraction` metric scaling. (Detailed coordination review in `var/log/`.) | #3453 | +| [ ] | **Make the low-level trtllm-gen TVM-FFI ops take structured config objects instead of long positional argument lists** (Β§5). The mid-cut blocker (the unified runner rotting after a `main` merge silently inserted `routing_input_mode` / `topk_weights` / `per_token_scale` and moved the tactic arg) was a *positional-argument-drift* failure with no compile-time signal. Delegating to `core.MoERunner` reduced the fragile call to one site; a structured `…Node::FromObject(config)` boundary (C++ reads named struct members) would remove the failure mode entirely and let adapters pass dataclass configs through unchanged. Out of MVP scope β€” sequence it after the NVFP4 MVP lands. | Β§5 | +| [ ] | **Gate `MoELayer._select_winner` behind the tuner's tuning mode.** Today a bucket-miss unconditionally runs the cross-backend `bench_gpu_time` shootout β€” even outside `autotune(True)` and even for ops `AutoTuner` would treat as `skip_ops` (where `choose_one` is a pure cache lookup / immediate fallback). The MVP is always exercised under `autotune(True)`, so this is correct for the shipped tests/benchmark, but a production layer built without an autotune context still silently benchmarks on first use per bucket. Follow-up: gate the shootout behind `self.tuner.is_tuning_mode`; when false and the bucket is uncached, select via a deterministic priority order (the Β§4 `DEFAULT_PRIORITY` heuristic) instead of benchmarking, and honor `skip_ops`. Needs a decision on the priority order, so deferred rather than rushed. | review (Reviewer 14) | + +### Explicit Non-Goals For This MVP + +| Status | Task | Notes | +| --- | --- | --- | +| [ ] | Keep FP8, MXFP4, and additional backend families out of this PR. | Track as follow-up runners after the NVFP4 path is solid. | +| [ ] | Keep general routing unification out of this PR. | This MVP assumes pre-routed inputs through `MoEActivationPack`. | +| [ ] | Keep broader public API ergonomics separate from MVP correctness. | Backend discovery, pipe-operator sugar, repro replay, and a full eager functional API can evolve after the MVP contract is stable. | diff --git a/flashinfer/fused_moe/__init__.py b/flashinfer/fused_moe/__init__.py index a5737ad1a7a..13f64662729 100644 --- a/flashinfer/fused_moe/__init__.py +++ b/flashinfer/fused_moe/__init__.py @@ -14,6 +14,30 @@ limitations under the License. """ +# Unified MoE API +from .api import ( # noqa: F401 + ActivationConfig, + BackendOptions, + CuteDslConfig, + CutlassConfig, + ExecutionConfig, + ExpertConfig, + MoEActivationPack, + MoEConfig, + MoEWeightPack, + QuantConfig, + QuantVariant, + RoutingConfig, + TrtllmBf16Config, + TrtllmFp4Config, + TrtllmFp8BlockConfig, + TrtllmFp8PerTensorConfig, + TrtllmMxInt4Config, +) +from .layer import MoELayer # noqa: F401 +from .runners import CuteDslNvfp4Runner, TrtllmFp4RoutedRunner # noqa: F401 + +# Legacy flat-argument APIs (unchanged, not deprecated) from .core import ( convert_to_block_layout, cutlass_fused_moe, @@ -69,6 +93,28 @@ _cute_dsl_available = False __all__ = [ + # Unified API + "ActivationConfig", + "BackendOptions", + "CuteDslConfig", + "CutlassConfig", + "ExecutionConfig", + "ExpertConfig", + "CuteDslNvfp4Runner", + "MoEActivationPack", + "MoEConfig", + "MoELayer", + "MoEWeightPack", + "TrtllmFp4RoutedRunner", + "QuantConfig", + "QuantVariant", + "RoutingConfig", + "TrtllmBf16Config", + "TrtllmFp4Config", + "TrtllmFp8BlockConfig", + "TrtllmFp8PerTensorConfig", + "TrtllmMxInt4Config", + # Legacy flat APIs "ActivationType", "Fp8QuantizationType", "RoutingMethodType", diff --git a/flashinfer/fused_moe/api.py b/flashinfer/fused_moe/api.py new file mode 100644 index 00000000000..64c15634114 --- /dev/null +++ b/flashinfer/fused_moe/api.py @@ -0,0 +1,541 @@ +"""Unified MoE API β€” configuration dataclasses and tensor groupings. + +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Config objects are frozen (immutable). Use ``dataclasses.replace`` to derive +variants. ``eval(repr(cfg))`` round-trips for every config type, enabling +repro-log serialization. + +Tensor groupings are mutable containers β€” they hold runtime data, not +configuration. They group related tensors for ergonomics (no more counting +20+ positional arguments). +""" + +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass, field +from enum import Enum +from typing import ClassVar, Dict, Optional, Tuple, Union + +from torch import Tensor + +from ..tllm_enums import ActivationType, RoutingMethodType + +# --------------------------------------------------------------------------- +# Enums +# --------------------------------------------------------------------------- +# Routing and activation reuse the shared kernel-level enums directly +# (``RoutingMethodType`` / ``ActivationType`` from ``tllm_enums``): the API +# speaks the kernels' vocabulary rather than mirroring it, so there is a single +# source of truth (PR #3093 review G1). Both are ``IntEnum`` β€” the value *is* +# the kernel ABI int β€” and carry an eval-safe ``__repr__`` (defined in +# ``tllm_enums``) plus ``ActivationType.is_gated`` for the repro round-trip and +# config helpers. +# +# ``QuantVariant`` below is the one genuinely API-level enum: it has no single +# kernel counterpart (the quant path is selected by dtype/scale wiring in the +# runners, not one enum), so it is defined here as a plain ``Enum``. + + +class QuantVariant(Enum): + """Quantization variant β€” single knob for dtype + granularity + scale convention.""" + + BF16 = 0 + FP8PerTensor = 1 + DeepSeekFp8 = 2 + MxFp8 = 3 + NVFP4 = 4 # day-1 MVP target + MXFP4 = 5 + MxInt4 = 6 + + def __repr__(self) -> str: + return f"{type(self).__name__}.{self.name}" + + +# --------------------------------------------------------------------------- +# Component configs β€” each owns one concern +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class RoutingConfig: + """Expert routing parameters. + + Parameters + ---------- + num_experts : int + Total number of experts (global, before EP sharding). + top_k : int + Number of experts selected per token. + method : RoutingMethodType + Routing strategy. + n_group : int or None + Expert group count for DeepSeekV3 routing. + topk_group : int or None + Number of groups selected in DeepSeekV3. + routed_scaling_factor : float or None + Fixed routing weight scaling (DeepSeekV3). + """ + + num_experts: int + top_k: int + method: RoutingMethodType = RoutingMethodType.Default + n_group: Optional[int] = None + topk_group: Optional[int] = None + routed_scaling_factor: Optional[float] = None + + def __repr__(self) -> str: + parts = [f"num_experts={self.num_experts!r}", f"top_k={self.top_k!r}"] + if self.method != RoutingMethodType.Default: + parts.append(f"method={self.method!r}") + if self.n_group is not None: + parts.append(f"n_group={self.n_group!r}") + if self.topk_group is not None: + parts.append(f"topk_group={self.topk_group!r}") + if self.routed_scaling_factor is not None: + parts.append(f"routed_scaling_factor={self.routed_scaling_factor!r}") + return f"RoutingConfig({', '.join(parts)})" + + +@dataclass(frozen=True) +class QuantConfig: + """Quantization scheme. + + Parameters + ---------- + variant : QuantVariant + Single knob for dtype + granularity + scale convention. + swizzled_scale_factors : bool or None + Whether block scale factors use the swizzled (vs linear) layout. + ``None`` β†’ backend default. Mirrors core's ``swizzled_input_sf``. Finer + ``SfLayout`` (128x4 / 8x4 / linear) selection is deferred (design doc + C42): unlike ``RoutingMethodType`` / ``ActivationType``, ``SfLayout`` has + no eval-safe ``__repr__``, so exposing it here would break the + ``eval(repr(cfg))`` round-trip β€” a bool keeps this config serializable. + per_token_scale : bool or None + Whether activations carry a per-token scale (vs per-tensor / block). + ``None`` β†’ backend default. + """ + + variant: QuantVariant = QuantVariant.BF16 + swizzled_scale_factors: Optional[bool] = None + per_token_scale: Optional[bool] = None + + +@dataclass(frozen=True) +class ActivationConfig: + """Fused activation between GEMM1 and GEMM2.""" + + # Convenience singletons β€” populated after class definition + swiglu: ClassVar[ActivationConfig] + geglu: ClassVar[ActivationConfig] + relu2: ClassVar[ActivationConfig] + identity: ClassVar[ActivationConfig] + + type: ActivationType = ActivationType.Swiglu + + def __repr__(self) -> str: + return f"ActivationConfig(type={self.type!r})" + + @property + def is_gated(self) -> bool: + return self.type.is_gated + + +ActivationConfig.swiglu = ActivationConfig(ActivationType.Swiglu) +ActivationConfig.geglu = ActivationConfig(ActivationType.Geglu) +ActivationConfig.relu2 = ActivationConfig(ActivationType.Relu2) +ActivationConfig.identity = ActivationConfig(ActivationType.Identity) + + +@dataclass(frozen=True) +class ExpertConfig: + """Expert geometry. + + Parameters + ---------- + intermediate_size : int + Hidden dimension of the expert FFN (the N in gemm1's MxK β†’ MxN). + local_expert_offset : int + Start index for expert-parallel sharding. + local_num_experts : int or None + Number of experts on this rank. ``None`` β†’ ``num_experts`` at runtime. + """ + + intermediate_size: int + local_expert_offset: int = 0 + local_num_experts: Optional[int] = None + + def __repr__(self) -> str: + parts = [f"intermediate_size={self.intermediate_size!r}"] + if self.local_expert_offset != 0: + parts.append(f"local_expert_offset={self.local_expert_offset!r}") + if self.local_num_experts is not None: + parts.append(f"local_num_experts={self.local_num_experts!r}") + return f"ExpertConfig({', '.join(parts)})" + + +@dataclass(frozen=True) +class ExecutionConfig: + """Runtime execution parameters. + + Parameters + ---------- + do_finalize : bool + Whether to apply routing-weight scaling and accumulate into output. + enable_pdl : bool or None + Persistent device launch. ``None`` β†’ auto (True for sm90+). + tune_max_num_tokens : int + Token budget hint for autotuner / CUDA graph capture. + """ + + do_finalize: bool = True + enable_pdl: Optional[bool] = None + tune_max_num_tokens: int = 8192 + + def __repr__(self) -> str: + parts = [] + if not self.do_finalize: + parts.append(f"do_finalize={self.do_finalize!r}") + if self.enable_pdl is not None: + parts.append(f"enable_pdl={self.enable_pdl!r}") + if self.tune_max_num_tokens != 8192: + parts.append(f"tune_max_num_tokens={self.tune_max_num_tokens!r}") + return f"ExecutionConfig({', '.join(parts)})" + + +# --------------------------------------------------------------------------- +# Backend configs β€” each declares hardware preconditions +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class TrtllmFp4Config: + """TensorRT-LLM FP4 block-scale backend.""" + + @classmethod + def supported(cls, arch: int) -> bool: + # SM100+ only: the routed runner delegates to the trtllm-gen sm100 + # module, which core.is_trtllm_moe_supported() gates on major >= 10. + # Returning True on SM90 would mark the backend available on H100 and + # then fail at dispatch. + return arch >= 100 + + @staticmethod + def prepare_weights( + w1_bf16, + w2_bf16, + *, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + device=None, + permute_cache=None, + ): + """Build the ``trtllm_fp4_routed`` weight view from canonical bf16 weights. + + Register the result with ``MoEWeightPack.prepare_for("trtllm_fp4_routed", ...)``. + See :func:`flashinfer.fused_moe.prepare.prepare_trtllm_fp4_weights`. + """ + from .prepare import prepare_trtllm_fp4_weights + + return prepare_trtllm_fp4_weights( + w1_bf16, + w2_bf16, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + permute_cache=permute_cache, + ) + + def __repr__(self) -> str: + return "TrtllmFp4Config()" + + +@dataclass(frozen=True) +class TrtllmFp8BlockConfig: + """TensorRT-LLM FP8 block-scale backend.""" + + @classmethod + def supported(cls, arch: int) -> bool: + return arch >= 80 + + def __repr__(self) -> str: + return "TrtllmFp8BlockConfig()" + + +@dataclass(frozen=True) +class TrtllmFp8PerTensorConfig: + """TensorRT-LLM FP8 per-tensor-scale backend.""" + + @classmethod + def supported(cls, arch: int) -> bool: + return arch >= 80 + + def __repr__(self) -> str: + return "TrtllmFp8PerTensorConfig()" + + +@dataclass(frozen=True) +class TrtllmBf16Config: + """TensorRT-LLM BF16 (unquantized) backend.""" + + @classmethod + def supported(cls, arch: int) -> bool: + return arch >= 100 + + def __repr__(self) -> str: + return "TrtllmBf16Config()" + + +@dataclass(frozen=True) +class TrtllmMxInt4Config: + """TensorRT-LLM MxInt4 backend.""" + + @classmethod + def supported(cls, arch: int) -> bool: + return arch >= 100 + + def __repr__(self) -> str: + return "TrtllmMxInt4Config()" + + +@dataclass(frozen=True) +class CutlassConfig: + """CUTLASS backend β€” broadest architecture support.""" + + @classmethod + def supported(cls, arch: int) -> bool: + return True # universal fallback + + def __repr__(self) -> str: + return "CutlassConfig()" + + +@dataclass(frozen=True) +class CuteDslConfig: + """CuteDSL NVFP4 backend β€” SM100 family only (Blackwell SM100, SM103). + + The underlying CuteDSL kernel throws at launch on SM120/SM121/SM130. + """ + + @classmethod + def supported(cls, arch: int) -> bool: + # SM100, SM103 β€” tighten when CuteDSL adds more targets + return arch in (100, 103) + + @staticmethod + def prepare_weights( + w1_bf16, + w2_bf16, + *, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + device=None, + ): + """Build the ``cute_dsl_nvfp4`` weight view from canonical bf16 weights. + + Register the result with ``MoEWeightPack.prepare_for("cute_dsl_nvfp4", ...)``. + See :func:`flashinfer.fused_moe.prepare.prepare_cute_dsl_nvfp4_weights`. + """ + from .prepare import prepare_cute_dsl_nvfp4_weights + + return prepare_cute_dsl_nvfp4_weights( + w1_bf16, + w2_bf16, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + ) + + def __repr__(self) -> str: + return "CuteDslConfig()" + + +# Union type for backend config +BackendConfigType = Union[ + TrtllmFp4Config, + TrtllmFp8BlockConfig, + TrtllmFp8PerTensorConfig, + TrtllmBf16Config, + TrtllmMxInt4Config, + CutlassConfig, + CuteDslConfig, +] + +ALL_BACKEND_CONFIGS = ( + TrtllmFp4Config, + TrtllmFp8BlockConfig, + TrtllmFp8PerTensorConfig, + TrtllmBf16Config, + TrtllmMxInt4Config, + CutlassConfig, + CuteDslConfig, +) + + +# --------------------------------------------------------------------------- +# BackendOptions +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class BackendOptions: + """Ordered list of backend candidates for dispatch / autotuning.""" + + candidates: Tuple[BackendConfigType, ...] = () # type: ignore[type-arg] + + def valid_for(self, arch: int) -> list: + """Return candidates whose hardware preconditions are met.""" + return [c for c in self.candidates if c.__class__.supported(arch)] + + def __len__(self) -> int: + return len(self.candidates) + + def __iter__(self): + return iter(self.candidates) + + +# --------------------------------------------------------------------------- +# MoEConfig β€” top-level container +# --------------------------------------------------------------------------- + +# Default backend search order +_DEFAULT_BACKEND = BackendOptions( + candidates=( + TrtllmFp4Config(), + TrtllmFp8BlockConfig(), + TrtllmFp8PerTensorConfig(), + TrtllmBf16Config(), + TrtllmMxInt4Config(), + CutlassConfig(), + CuteDslConfig(), + ) +) + + +@dataclass(frozen=True) +class MoEConfig: + """Top-level MoE configuration. + + Combines all sub-configs into a single hashable, serializable object. + Supports ``**config`` unpacking via the dict protocol. + + Example + ------- + >>> config = MoEConfig( + ... routing=RoutingConfig(num_experts=64, top_k=8, + ... method=RoutingMethodType.DeepSeekV3), + ... quant=QuantConfig(variant=QuantVariant.DeepSeekFp8), + ... experts=ExpertConfig(intermediate_size=2048), + ... ) + >>> output = fused_moe(tensors, **config) + """ + + routing: RoutingConfig + quant: QuantConfig + experts: ExpertConfig + activation: ActivationConfig = field( + default_factory=lambda: ActivationConfig(ActivationType.Swiglu) + ) + backend: BackendOptions = field(default_factory=lambda: _DEFAULT_BACKEND) + execution: ExecutionConfig = field(default_factory=ExecutionConfig) + + # --- Dict-unpacking protocol: enables ``**config`` at call sites --- + + def keys(self): + return (f.name for f in dataclasses.fields(self)) + + def __getitem__(self, key: str): + return getattr(self, key) + + # --- Serialization --- + # + # ``repr(config)`` already round-trips to valid constructor syntax (frozen + # dataclasses + qualified enum repr), which is all the MVP needs for logging. + # A deserializer (``from_repr``/``from_dict``) is intentionally *not* shipped + # here: eval-based parsing is a security smell, and the repro/serialization + # design (versioned schema vs. same-version-only) is a documented post-MVP + # item β€” see docs/design_docs/flashinfer_moe_api.md (C4-C5/C39, Post-MVP + # Carryover). It will land with the repro tooling, not before. + + +# --------------------------------------------------------------------------- +# Activation / weight packs for the autotuned pre-routed path +# --------------------------------------------------------------------------- +# These are the runner-level inputs used by MoELayer (plan Β§1). +# +# Why two packs instead of one tensor bundle (PR #3093 review G5): the grouping +# axis is *lifetime/role*, which is invariant across backends β€” +# * MoEActivationPack: per-call transient data (pre-routed activations), +# rebuilt every forward; +# * MoEWeightPack: long-lived weights, materialized once at load and read +# every call, holding one native view *per backend* (the price of +# cross-backend autotune) keyed by backend_key. +# A single per-call bundle cannot model a load-time, multi-backend weight cache +# without conflating the two lifetimes. We deliberately do *not* group tensors +# by compute-graph stage (e.g. gemm1/gemm2): that mirrors the unfused two-GEMM +# implementation and would overfit it β€” a fused/megakernel backend has no such +# boundary, so a graph-shaped public API would leak one backend's internals. +# Each pack presents itself to a backend via prepare_for / get_view, keeping +# backend-specific layout logic out of the dispatch hot-path. + + +@dataclass +class MoEActivationPack: + """Per-call transient data β€” pre-quantized NVFP4 activations + pre-routed indices.""" + + hidden_states_q: Tensor # [M, H//2] uint8 (packed NVFP4) + hidden_states_scale: Tensor # [M, H//16] float8_e4m3fn + selected_experts: Tensor # [M, top_k] int32 + final_scales: Tensor # [M, top_k] float32 + + @property + def num_tokens(self) -> int: + return self.hidden_states_q.shape[0] + + +@dataclass +class MoEWeightPack: + """Long-lived weight container with per-backend native materializations. + + Each backend's native weight layout (quantized, swizzled, MMA-ordered, etc.) + is stored under its ``backend_key``. Populated once at model-load / + layer-init via ``prepare_for(key, view)``; read on every call via + ``get_view(key)``. + + Holding multiple materializations is intentional β€” that's the memory cost + the user pays for cross-backend autotune. Each view is the exact kwargs + dict that runner's ``forward`` expects for weight-side arguments. + """ + + native_views: Dict[str, Dict[str, Tensor]] = field(default_factory=dict) + + def prepare_for(self, backend_key: str, view: Dict[str, Tensor]) -> None: + """Register a backend-native weight view. Caller owns the quantization + / swizzle / layout conversion β€” this method just stores the result.""" + self.native_views[backend_key] = view + + def get_view(self, backend_key: str) -> Dict[str, Tensor]: + if backend_key not in self.native_views: + raise KeyError( + f"Weights not prepared for backend {backend_key!r}. " + f"Available: {list(self.native_views)}" + ) + return self.native_views[backend_key] diff --git a/flashinfer/fused_moe/core.py b/flashinfer/fused_moe/core.py index e80579b1ac8..c272a59743e 100644 --- a/flashinfer/fused_moe/core.py +++ b/flashinfer/fused_moe/core.py @@ -1053,7 +1053,7 @@ def cutlass_fused_moe( @dataclass -class MoEInputs: +class MoeRunnerInputs: """MoERunner inputs. Field order defines the flat-list index used by the autotuner. @@ -1094,10 +1094,10 @@ class MoEInputs: } def to_list(self) -> List[Optional[torch.Tensor]]: - return [getattr(self, name) for name in MoEInputs._FIELDS] + return [getattr(self, name) for name in MoeRunnerInputs._FIELDS] @classmethod - def from_list(cls, lst: List) -> "MoEInputs": + def from_list(cls, lst: List) -> "MoeRunnerInputs": return cls(**{name: lst[i] for i, name in enumerate(cls._FIELDS)}) @classmethod @@ -1105,6 +1105,13 @@ def idx(cls, name: str) -> int: return cls._FIELDS.index(name) +# Backward-compatible alias: this class was previously named ``MoEInputs``. +# Renamed to ``MoeRunnerInputs`` to disambiguate from the unified-API input +# grouping (the ``MoEActivationPack`` / ``MoEWeightPack`` lifetime split) β€” see +# PR #3093 review G6. Old name kept working for out-of-tree importers and tests. +MoEInputs = MoeRunnerInputs + + def _unpack_trtllm_moe_output( intermediate_output, output: torch.Tensor, @@ -1183,7 +1190,7 @@ def __init__( def _make_tuning_config( self, - moe_inputs: "MoEInputs", + moe_inputs: "MoeRunnerInputs", tune_max_num_tokens: int = 8192, **kwargs, ) -> TuningConfig: @@ -1240,7 +1247,7 @@ def _init_packed_topk_ids(shapes, dtype, device): ).to(dtype) sorted_inputs = sorted( - (MoEInputs.idx(name), name, init) for name, init in spec.items() + (MoeRunnerInputs.idx(name), name, init) for name, init in spec.items() ) input_idx = tuple(i for i, _, _ in sorted_inputs) @@ -1263,7 +1270,7 @@ def _dynamic_dim(name: str) -> int: f"expected layout (num_tokens={num_tokens}, ...)" ) return 0 - return MoEInputs._DYNAMIC_DIM[name] + return MoeRunnerInputs._DYNAMIC_DIM[name] dim_idx = tuple(_dynamic_dim(name) for _, name, _ in sorted_inputs) initializers = [init for _, _, init in sorted_inputs] @@ -1286,7 +1293,7 @@ def get_valid_tactics( inputs: List[torch.Tensor], profile: OptimizationProfile, ) -> List[int]: - moe_inputs = MoEInputs.from_list(inputs) + moe_inputs = MoeRunnerInputs.from_list(inputs) num_tokens = moe_inputs.hidden_states.shape[0] has_gemm1_lora_delta = moe_inputs.gemm1_lora_delta is not None @@ -1324,7 +1331,7 @@ def forward( do_preparation: bool = False, **kwargs, ): - moe_inputs = MoEInputs.from_list(inputs) + moe_inputs = MoeRunnerInputs.from_list(inputs) output = moe_inputs.output routing_logits = moe_inputs.routing_logits topk_ids = moe_inputs.topk_ids @@ -1650,7 +1657,7 @@ def trtllm_bf16_moe_op( num_experts=num_experts, ) - moe_inputs = MoEInputs( + moe_inputs = MoeRunnerInputs( output=output, routing_logits=routing_logits, topk_ids=topk_ids, @@ -1823,7 +1830,7 @@ def trtllm_fp8_per_tensor_scale_moe_op( num_experts=num_experts, ) - moe_inputs = MoEInputs( + moe_inputs = MoeRunnerInputs( output=output, routing_logits=routing_logits, topk_ids=topk_ids, @@ -2053,7 +2060,7 @@ def trtllm_fp8_block_scale_moe_op( num_experts=num_experts, ) - moe_inputs = MoEInputs( + moe_inputs = MoeRunnerInputs( output=output, routing_logits=routing_logits, topk_ids=topk_ids, @@ -2284,7 +2291,7 @@ def trtllm_fp4_block_scale_moe_op( use_per_token_scaling=per_token_scale is not None, num_experts=num_experts, ) - moe_inputs = MoEInputs( + moe_inputs = MoeRunnerInputs( output=output, routing_logits=routing_logits, topk_ids=topk_ids, @@ -2510,7 +2517,7 @@ def trtllm_mxint4_block_scale_moe_op( num_experts=num_experts, ) - moe_inputs = MoEInputs( + moe_inputs = MoeRunnerInputs( output=output, routing_logits=routing_logits, topk_ids=topk_ids, @@ -2629,6 +2636,11 @@ def _fake_trtllm_mxint4_block_scale_moe( trtllm_fp8_block_scale_moe=trtllm_fp8_block_scale_moe_op, trtllm_fp4_block_scale_moe=trtllm_fp4_block_scale_moe_op, trtllm_mxint4_block_scale_moe=trtllm_mxint4_block_scale_moe_op, + # Canonical tactic-aware TunableRunner (closes over the raw moe_op and + # trtllm_get_valid_moe_configs). Exposed so the unified MoE API's + # TrtllmFp4RoutedRunner can delegate to it instead of re-deriving the + # raw op's positional call. + MoERunner=MoERunner, ) diff --git a/flashinfer/fused_moe/layer.py b/flashinfer/fused_moe/layer.py new file mode 100644 index 00000000000..5abde6c8983 --- /dev/null +++ b/flashinfer/fused_moe/layer.py @@ -0,0 +1,202 @@ +"""MoELayer β€” stateful cross-backend MoE dispatcher with autotune. + +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Builds one runner per compatible backend, picks the cross-backend winner +by measuring each runner's best tactic, then dispatches to the winner. + +MVP scope: NVFP4 only, pre-routed path, two backends. +""" + +from __future__ import annotations + +from statistics import median +from typing import Any, Dict, List, Optional, Tuple, Type, Union + +import torch + +from ..autotuner import AutoTuner +from ..utils import get_compute_capability +from .api import ( + ActivationType, + CuteDslConfig, + MoEActivationPack, + MoEConfig, + MoEWeightPack, + QuantVariant, + TrtllmFp4Config, +) +from .runners import CuteDslNvfp4Runner, TrtllmFp4RoutedRunner +from .utils import map_to_hybrid_bucket + + +# Union of the concrete runners the layer dispatches to. All share +# backend_key / tuning_config / pack_inputs as attributes or class members; +# typing the list with this Union gives mypy the visibility it needs. +_RunnerT = Union[CuteDslNvfp4Runner, TrtllmFp4RoutedRunner] + +# Map backend-config class -> runner class +_BACKEND_RUNNERS: Dict[type, Type[_RunnerT]] = { + CuteDslConfig: CuteDslNvfp4Runner, + TrtllmFp4Config: TrtllmFp4RoutedRunner, +} + + +class MoELayer: + """Stateful MoE layer with cross-backend autotune. + + Example + ------- + >>> layer = MoELayer(config) + >>> out = layer(act_pack, weight_pack) + """ + + def __init__(self, config: MoEConfig, device: Optional[torch.device] = None): + self.config = config + self._validate_mvp_scope(config) + self.device = device or torch.device("cuda", torch.cuda.current_device()) + self.tuner = AutoTuner.get() + + major, minor = get_compute_capability(self.device) + arch = major * 10 + minor + + # Build one runner per compatible backend + self.runners: List[_RunnerT] = [] + for backend_cfg in config.backend: + if not backend_cfg.supported(arch): + continue + runner_cls = _BACKEND_RUNNERS.get(type(backend_cfg)) + if runner_cls is None: + continue # MVP scope β€” skip non-MVP backends silently + self.runners.append(runner_cls(config, device=self.device)) + + if not self.runners: + mvp = ", ".join(c.__name__ for c in _BACKEND_RUNNERS) + raise RuntimeError( + f"MoELayer: none of the configured backends " + f"{[type(c).__name__ for c in config.backend]} are usable on " + f"arch sm{arch}. The MVP supports only NVFP4 via [{mvp}]." + ) + + # Cross-backend winner cache, keyed by the num_tokens tuning bucket. + # See the MoELayer reuse contract (CR4): the fastest backend can differ + # across token-count buckets, so each bucket caches its own winner. + self._winners: Dict[int, Tuple[_RunnerT, Any]] = {} + # Backend key selected on the most recent call (introspection hook). + self._last_winner_backend: Optional[str] = None + + @staticmethod + def _validate_mvp_scope(config: MoEConfig) -> None: + """Fail fast on configs the MVP cannot execute (CR6). + + The MVP is NVFP4 + Swiglu + pre-routed packs only. Surfacing this at + construction time turns a deep C++ crash or silent backend skip into a + clear, actionable Python error. + """ + variant = config.quant.variant + if variant is not QuantVariant.NVFP4: + raise NotImplementedError( + f"MoELayer MVP supports only QuantVariant.NVFP4; got {variant!r}. " + "FP8 / MXFP4 / MxInt4 / BF16 paths are tracked as post-MVP " + "follow-ups." + ) + act = config.activation.type + if act is not ActivationType.Swiglu: + raise NotImplementedError( + f"MoELayer MVP supports only the Swiglu activation; got {act!r}." + ) + + def __call__( + self, + act_pack: MoEActivationPack, + weight_pack: MoEWeightPack, + ) -> torch.Tensor: + ceiling = self.config.execution.tune_max_num_tokens + if act_pack.num_tokens > ceiling: + raise ValueError( + f"num_tokens={act_pack.num_tokens} exceeds " + f"tune_max_num_tokens={ceiling}. " + f"Reconstruct MoELayer with a larger ceiling." + ) + + bucket = map_to_hybrid_bucket(act_pack.num_tokens, ceiling) + winner = self._winners.get(bucket) + if winner is None: + winner = self._select_winner(act_pack, weight_pack) + self._winners[bucket] = winner + runner, tactic = winner + self._last_winner_backend = runner.backend_key + + inputs = runner.pack_inputs(act_pack, weight_pack) + return runner.forward(inputs, tactic=tactic) + + def _select_winner( + self, + act_pack: MoEActivationPack, + weight_pack: MoEWeightPack, + ) -> Tuple[_RunnerT, Any]: + """Run per-runner autotune, then measure each winner-tactic and + pick cross-backend winner.""" + # Lazy import: keep the library import path (``import flashinfer``) free + # of a dependency on the testing framework. The GPU timing helper is only + # needed here, on the autotune path. Relocating it to a non-testing + # utility module is the cleaner long-term fix (post-MVP). + from ..testing.utils import bench_gpu_time + + best_time_ms = float("inf") + best_runner: Optional[_RunnerT] = None + best_tactic: Any = -1 + + for runner in self.runners: + inputs = runner.pack_inputs(act_pack, weight_pack) + # Per-runner tactic selection via autotuner + _, tactic = self.tuner.choose_one( + custom_op=f"moe_{runner.backend_key}", + runners=[runner], + tuning_config=runner.tuning_config, + inputs=inputs, + ) + # Measure runner at its winning tactic. Use CUDA-graph timing so + # the cross-backend comparison reflects production (graph-captured) + # latency rather than per-call launch/Python overhead β€” at low token + # counts (~tens of us kernels) a no-graph 10-iter median is dominated + # by that overhead and picks the wrong backend. Requires a warmed-up + # layer (the autotune pass above), not a cold capture. + times = bench_gpu_time( + lambda r=runner, i=inputs, t=tactic: r.forward(i, tactic=t), + dry_run_iters=5, + repeat_iters=30, + use_cuda_graph=True, + ) + t_ms = median(times) + if t_ms < best_time_ms: + best_time_ms = t_ms + best_runner = runner + best_tactic = tactic + + assert best_runner is not None # self.runners is non-empty + return best_runner, best_tactic + + # ---- Introspection helpers --------------------------------------------- + + @property + def winner_backend(self) -> Optional[str]: + """Backend key selected on the most recent call, or None before first call.""" + return self._last_winner_backend + + def reset_winner(self) -> None: + """Clear all cached per-bucket winners β€” next call re-tunes.""" + self._winners.clear() + self._last_winner_backend = None diff --git a/flashinfer/fused_moe/prepare.py b/flashinfer/fused_moe/prepare.py new file mode 100644 index 00000000000..9ca5bc2c66e --- /dev/null +++ b/flashinfer/fused_moe/prepare.py @@ -0,0 +1,268 @@ +"""First-class NVFP4 weight-preparation helpers for the unified MoE API. + +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Backends consume different native weight layouts (quantization + swizzle + +MMA reorder). These helpers turn canonical bf16 expert weights into the +backend-native ``MoEWeightPack`` views, so that preparation lives in the +implementation surface rather than being copy-pasted into tests and benchmarks +(design doc CR2/CR7; reviewer comments C6, C7, C31, C32). + +They are exposed as ``TrtllmFp4Config.prepare_weights(...)`` / +``CuteDslConfig.prepare_weights(...)`` static helpers (see ``api.py``). +""" + +from __future__ import annotations + +from typing import Dict, Optional + +import torch + +# Module-level permute-index cache. Permute indices depend only on weight +# dims, so the cache is safe to reuse across shapes and calls. +_TRTLLM_PERMUTE_CACHE: dict = {} + + +def prepare_trtllm_fp4_weights( + w1_bf16: torch.Tensor, + w2_bf16: torch.Tensor, + *, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + device: Optional[torch.device] = None, + permute_cache: Optional[dict] = None, +) -> Dict[str, torch.Tensor]: + """Build the TRTLLM NVFP4 ``trtllm_fp4_routed`` weight view. + + Layout is ``Shuffled_MajorK`` β€” the only NVFP4-compatible trtllm-gen combo + today: per-expert gated-act reorder + MMA shuffle on the packed weights and + ``block_scale_interleave`` on the block scales. + + Parameters + ---------- + w1_bf16 : Tensor + Gate+up expert weights ``[num_local_experts, 2*intermediate_size, hidden_size]``. + w2_bf16 : Tensor + Down-projection expert weights ``[num_local_experts, hidden_size, intermediate_size]``. + num_local_experts, hidden_size, intermediate_size : int + Expert geometry. + device : torch.device, optional + Target device; defaults to ``w1_bf16.device``. + permute_cache : dict, optional + Shape-keyed permute-index cache; defaults to a module-level cache. + + Returns + ------- + dict + Keys expected by ``TrtllmFp4RoutedRunner.pack_inputs``: ``gemm1_weights``, + ``gemm1_weights_scale``, ``gemm1_alpha``, ``gemm2_weights``, + ``gemm2_weights_scale``, ``output1_scale_scalar``, + ``output1_scale_gate_scalar``, ``output2_scale_scalar``. + """ + from ..fp4_quantization import fp4_quantize + from ..quantization.fp4_quantization import block_scale_interleave + from .core import ( + _maybe_get_cached_w3_w1_permute_indices, + get_w2_permute_indices_with_cache, + ) + + if device is None: + device = w1_bf16.device + # Honor the documented device target: move canonical weights there (no-op if + # already resident). Otherwise CPU weights + device="cuda" hit mixed-device + # ops inside quantization. + w1_bf16 = w1_bf16.to(device) + w2_bf16 = w2_bf16.to(device) + if permute_cache is None: + permute_cache = _TRTLLM_PERMUTE_CACHE + + sf_vec_size = 16 + epilogue_tile_m = 128 # TRTLLM kernel-internal constant + + w1_gs = torch.tensor([1.0], device=device, dtype=torch.float32) + w1_flat = w1_bf16.view(num_local_experts * 2 * intermediate_size, hidden_size) + w1_q_flat, w1_sf_flat = fp4_quantize( + w1_flat, + global_scale=w1_gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=False, + ) + g1_w = w1_q_flat.view( + num_local_experts, 2 * intermediate_size, hidden_size // 2 + ).view(torch.uint8) + g1_s = w1_sf_flat.view(torch.float8_e4m3fn).reshape( + num_local_experts, 2 * intermediate_size, hidden_size // sf_vec_size + ) + + w2_gs = torch.tensor([1.0], device=device, dtype=torch.float32) + w2_flat = w2_bf16.view(num_local_experts * hidden_size, intermediate_size) + w2_q_flat, w2_sf_flat = fp4_quantize( + w2_flat, + global_scale=w2_gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=False, + ) + g2_w = w2_q_flat.view(num_local_experts, hidden_size, intermediate_size // 2).view( + torch.uint8 + ) + g2_s = w2_sf_flat.view(torch.float8_e4m3fn).reshape( + num_local_experts, hidden_size, intermediate_size // sf_vec_size + ) + + g1_w_sh, g1_s_sh, g2_w_sh, g2_s_sh = [], [], [], [] + for i in range(num_local_experts): + p = _maybe_get_cached_w3_w1_permute_indices( + permute_cache, g1_w[i], epilogue_tile_m, is_gated_act_gemm=True + ) + g1_w_sh.append(g1_w[i][p.to(device)].contiguous()) + + p_sf = _maybe_get_cached_w3_w1_permute_indices( + permute_cache, + g1_s[i].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + is_gated_act_gemm=True, + ) + g1_s_sh.append( + block_scale_interleave( + g1_s[i].view(torch.uint8)[p_sf.to(device)].contiguous() + ) + ) + + p = get_w2_permute_indices_with_cache(permute_cache, g2_w[i], epilogue_tile_m) + g2_w_sh.append(g2_w[i][p.to(device)].contiguous()) + + p_sf = get_w2_permute_indices_with_cache( + permute_cache, + g2_s[i].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + ) + g2_s_sh.append( + block_scale_interleave( + g2_s[i].view(torch.uint8)[p_sf.to(device)].contiguous() + ) + ) + + ones = torch.ones(num_local_experts, device=device, dtype=torch.float32) + return { + "gemm1_weights": torch.stack(g1_w_sh), + "gemm1_weights_scale": torch.stack(g1_s_sh) + .view(torch.float8_e4m3fn) + .reshape(num_local_experts, 2 * intermediate_size, hidden_size // sf_vec_size), + "gemm1_alpha": ones, + "gemm2_weights": torch.stack(g2_w_sh), + "gemm2_weights_scale": torch.stack(g2_s_sh) + .view(torch.float8_e4m3fn) + .reshape(num_local_experts, hidden_size, intermediate_size // sf_vec_size), + "output1_scale_scalar": ones, + "output1_scale_gate_scalar": ones, + "output2_scale_scalar": ones, + } + + +def _interleave_linear_and_gate( + x: torch.Tensor, group_size: int = 64, dim: int = -1 +) -> torch.Tensor: + """Interleave the linear and gate halves of a SwiGLU gemm1 weight.""" + sizes = x.size() + dim = dim % x.dim() + assert sizes[dim] % (group_size * 2) == 0 + prev_sizes = sizes[:dim] + post_sizes = sizes[dim + 1 :] + x = x.view(*prev_sizes, 2, sizes[dim] // (group_size * 2), group_size, *post_sizes) + x = x.transpose(dim, dim + 1).contiguous().view(*sizes) + return x + + +def prepare_cute_dsl_nvfp4_weights( + w1_bf16: torch.Tensor, + w2_bf16: torch.Tensor, + *, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + device: Optional[torch.device] = None, +) -> Dict[str, torch.Tensor]: + """Build the CuteDSL NVFP4 ``cute_dsl_nvfp4`` weight view. + + Gemm1 weights get the SwiGLU linear/gate interleave; both gemms are NVFP4 + block-quantized (swizzled) with scales converted to the CuteDSL MMA layout. + Starts from the same canonical bf16 expert weights as + :func:`prepare_trtllm_fp4_weights`, so a single weight set can feed both + backends and a shared reference. + + Returns + ------- + dict + Keys expected by ``CuteDslNvfp4Runner.pack_inputs``: ``w1_weight``, + ``w1_weight_sf``, ``w1_alpha``, ``fc2_input_scale``, ``w2_weight``, + ``w2_weight_sf``, ``w2_alpha``. + """ + from ..cute_dsl.utils import convert_sf_to_mma_layout + from ..fp4_quantization import fp4_quantize + + if device is None: + device = w1_bf16.device + # Honor the documented device target (no-op if already resident); avoids + # mixed-device ops when canonical weights are on CPU. + w1_bf16 = w1_bf16.to(device) + w2_bf16 = w2_bf16.to(device) + + sf_vec_size = 16 + gs = torch.tensor([1.0], device=device, dtype=torch.float32) + + w1_interleaved = _interleave_linear_and_gate(w1_bf16, group_size=64, dim=1) + w1_flat = w1_interleaved.view( + num_local_experts * 2 * intermediate_size, hidden_size + ) + w1_q_flat, w1_sf_flat = fp4_quantize( + w1_flat, global_scale=gs, sf_vec_size=sf_vec_size, is_sf_swizzled_layout=True + ) + w1_weight = w1_q_flat.view( + num_local_experts, 2 * intermediate_size, hidden_size // 2 + ) + w1_weight_sf = convert_sf_to_mma_layout( + w1_sf_flat, + m=2 * intermediate_size, + k=hidden_size, + num_groups=num_local_experts, + sf_vec_size=sf_vec_size, + ) + + w2_flat = w2_bf16.view(num_local_experts * hidden_size, intermediate_size) + w2_q_flat, w2_sf_flat = fp4_quantize( + w2_flat, global_scale=gs, sf_vec_size=sf_vec_size, is_sf_swizzled_layout=True + ) + w2_weight = w2_q_flat.view(num_local_experts, hidden_size, intermediate_size // 2) + w2_weight_sf = convert_sf_to_mma_layout( + w2_sf_flat, + m=hidden_size, + k=intermediate_size, + num_groups=num_local_experts, + sf_vec_size=sf_vec_size, + ) + + ones = torch.ones(num_local_experts, device=device, dtype=torch.float32) + return { + "w1_weight": w1_weight, + "w1_weight_sf": w1_weight_sf, + "w1_alpha": ones, + "fc2_input_scale": torch.tensor([1.0], device=device, dtype=torch.float32), + "w2_weight": w2_weight, + "w2_weight_sf": w2_weight_sf, + "w2_alpha": ones, + } diff --git a/flashinfer/fused_moe/runners.py b/flashinfer/fused_moe/runners.py new file mode 100644 index 00000000000..65abcd1bea1 --- /dev/null +++ b/flashinfer/fused_moe/runners.py @@ -0,0 +1,329 @@ +"""Unified MoE runner adapters for the autotuned pre-routed NVFP4 path. + +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Each runner wraps one backend and translates (MoEActivationPack, MoEWeightPack) +into the backend's native calling convention. Both MVP runners are thin +adapters over an existing, canonical inner runner (CuteDSL's +``CuteDslFusedMoENvfp4Runner`` and trtllm-gen's ``core.MoERunner``) so the +fragile backend-specific kernel-launch code lives in exactly one place. + +MVP scope: NVFP4 only, two backends (CuteDSL, TRTLLM routed). +""" + +from __future__ import annotations + +from typing import Any, List + +import torch + +from ..autotuner import TunableRunner +from .api import MoEActivationPack, MoEConfig, MoEWeightPack + + +# --------------------------------------------------------------------------- +# CuteDSL NVFP4 runner β€” delegates to the existing CuteDslFusedMoENvfp4Runner +# --------------------------------------------------------------------------- + + +class CuteDslNvfp4Runner(TunableRunner): + """Wraps CuteDslFusedMoENvfp4Runner, translating Pack inputs into its + List[Tensor] convention.""" + + backend_key = "cute_dsl_nvfp4" + + def __init__(self, config: MoEConfig, device: torch.device): + from .cute_dsl.fused_moe import _cute_dsl_fused_moe_nvfp4_impl + from .cute_dsl.tuner import CuteDslFusedMoENvfp4Runner + + experts = config.experts + routing = config.routing + num_local_experts = experts.local_num_experts or routing.num_experts + + self._inner = CuteDslFusedMoENvfp4Runner( + forward_impl=_cute_dsl_fused_moe_nvfp4_impl, + num_experts=routing.num_experts, + top_k=routing.top_k, + num_local_experts=num_local_experts, + local_expert_offset=experts.local_expert_offset, + ) + # tuning_config is an instance attribute on the inner runner (its + # dummy expert-id span depends on num_experts/offset), so read it from + # the instance we just built, not off the class. + self.tuning_config = self._inner.tuning_config + + def get_valid_tactics(self, inputs: List[torch.Tensor], profile: Any) -> List[Any]: + return self._inner.get_valid_tactics(inputs, profile) + + def forward( + self, + inputs: List[torch.Tensor], + tactic: Any = -1, + do_preparation: bool = False, + **kwargs: Any, + ) -> torch.Tensor: + return self._inner.forward( + inputs, tactic=tactic, do_preparation=do_preparation, **kwargs + ) + + def pack_inputs( + self, act: MoEActivationPack, weights: MoEWeightPack + ) -> List[torch.Tensor]: + """Translate Packs β†’ List[Tensor] expected by CuteDslFusedMoENvfp4Runner. + + Expected weight view keys: w1_weight, w1_weight_sf, w1_alpha, + fc2_input_scale, w2_weight, w2_weight_sf, w2_alpha. + Input order: x, x_sf, token_selected_experts, token_final_scales, + w1_weight, w1_weight_sf, w1_alpha, fc2_input_scale, + w2_weight, w2_weight_sf, w2_alpha, moe_output. + + The trailing ``moe_output`` buffer (index 11) is optional for a direct + ``forward`` (the inner runner allocates it), but the inner runner's + tuning_config declares index 11 as a dynamic tensor, so it must be + present for the autotuner profiling path to assign it a per-bucket + initializer. + """ + v = weights.get_view(self.backend_key) + num_tokens = act.hidden_states_q.shape[0] + hidden_size = act.hidden_states_q.shape[1] * 2 # FP4 packed + moe_output = act.hidden_states_q.new_empty( + (num_tokens, hidden_size), dtype=torch.bfloat16 + ) + return [ + act.hidden_states_q, + act.hidden_states_scale.unsqueeze(-1), # CuteDSL expects [M, H//16, 1] + act.selected_experts, + act.final_scales, + v["w1_weight"], + v["w1_weight_sf"], + v["w1_alpha"], + v["fc2_input_scale"], + v["w2_weight"], + v["w2_weight_sf"], + v["w2_alpha"], + moe_output, + ] + + def __hash__(self): + return hash(("cute_dsl_nvfp4", hash(self._inner))) + + +# --------------------------------------------------------------------------- +# TRTLLM FP4 routed runner β€” delegates to the canonical trtllm-gen MoERunner +# --------------------------------------------------------------------------- + + +class TrtllmFp4RoutedRunner(TunableRunner): + """Pre-routed NVFP4 adapter over the canonical trtllm-gen ``MoERunner``. + + Translates (MoEActivationPack, MoEWeightPack) into the ``MoeRunnerInputs`` list + plus the static weight/config kwargs that ``core.MoERunner.forward`` + consumes, then delegates tactic enumeration, tuning-config construction, and + the tactic'd forward to that inner runner. This mirrors + ``CuteDslNvfp4Runner`` (which wraps ``CuteDslFusedMoENvfp4Runner``) and keeps + the fragile raw-op positional launch in exactly one place β€” + ``core.MoERunner.forward``. + + Routing is pre-computed (``RoutingInputMode.PackedPrecomputed``): the packed + int32 top-k ids carry ``((expert_id - local_offset) << 16) | bf16(weight)``. + The inner ``MoERunner`` needs the hidden size for its tactic keys and tuning + buckets, so it is built lazily on the first ``pack_inputs`` call. + """ + + backend_key = "trtllm_fp4_routed" + + def __init__(self, config: MoEConfig, device: torch.device): + from ..tllm_enums import DtypeTrtllmGen, Fp8QuantizationType + from ..utils import device_support_pdl + from .core import get_trtllm_moe_sm100_module + + self.config = config + self.device = device + self._module = get_trtllm_moe_sm100_module() + + routing = config.routing + experts = config.experts + execution = config.execution + self._num_local_experts = experts.local_num_experts or routing.num_experts + self._local_expert_offset = experts.local_expert_offset + self._intermediate_size = experts.intermediate_size + self._activation_type = int(config.activation.type) + self._tune_max_num_tokens = execution.tune_max_num_tokens + + # NVFP4: E2m1 activations + weights, no fp8 sub-quant. + self._dtype_act = DtypeTrtllmGen.E2m1 + self._dtype_weights = DtypeTrtllmGen.E2m1 + self._fp8_quantization_type = Fp8QuantizationType.NoneFp8 + + # enable_pdl=None means "auto" β€” resolve once here exactly like the + # high-level wrapper does before building its MoERunner, because the raw + # op (reached via MoERunner.forward) expects a concrete bool. Resolving + # once also keeps the value stable across CUDA-graph capture/replay. + enable_pdl = execution.enable_pdl + if enable_pdl is None: + enable_pdl = device_support_pdl(device) + self._enable_pdl = enable_pdl + + # Built lazily on first pack_inputs once hidden_size is known. + self._inner: Any = None + self._static_kwargs: dict = {} + self.tuning_config: Any = None + + def _ensure_inner(self, hidden_size: int) -> None: + if self._inner is not None: + return + from ..tllm_enums import WeightLayout + + self._inner = self._module.MoERunner( + top_k=self.config.routing.top_k, + num_local_experts=self._num_local_experts, + dtype_act=self._dtype_act, + dtype_weights=self._dtype_weights, + fp8_quantization_type=self._fp8_quantization_type, + hidden_size=hidden_size, + intermediate_size=self._intermediate_size, + activation_type=self._activation_type, + use_shuffled_weight=True, + weight_layout=int(WeightLayout.MajorK), + use_per_token_scaling=False, + num_experts=self.config.routing.num_experts, + ) + + def get_valid_tactics( # type: ignore[override] + self, inputs: List[torch.Tensor], profile: Any + ) -> List[Any]: + # The inner runner reads num_tokens from inputs + its own instance key; + # no static kwargs are needed for tactic enumeration. + return self._inner.get_valid_tactics(inputs, profile) + + def forward( + self, + inputs: List[torch.Tensor], + tactic: Any = -1, + do_preparation: bool = False, + **kwargs: Any, + ) -> torch.Tensor: + # MoELayer's autotuner call passes no kwargs, so the static weight/config + # kwargs are injected here. The inner runner writes the result in-place + # into inputs[0] (the output buffer of the MoeRunnerInputs list). + self._inner.forward( + inputs, + tactic=tactic, + do_preparation=do_preparation, + **self._static_kwargs, + ) + return inputs[0] + + def pack_inputs( + self, act: MoEActivationPack, weights: MoEWeightPack + ) -> List[torch.Tensor]: + """Translate Packs β†’ the ``MoeRunnerInputs`` list ``core.MoERunner`` expects. + + Expected weight view keys: gemm1_weights, gemm1_weights_scale, + gemm1_alpha, gemm2_weights, gemm2_weights_scale, and optionally + output1_scale_scalar, output1_scale_gate_scalar, output2_scale_scalar. + + The local-shard offset comes from ``ExpertConfig.local_expert_offset`` + on the config this runner was built with. For expert-parallel + pre-routed inputs the kernel indexes local experts as + ``[0, local_num_experts)``, so global expert ids are shifted down by the + local offset before packing. + """ + from .core import MoeRunnerInputs, RoutingInputMode + + v = weights.get_view(self.backend_key) + routing = self.config.routing + + num_tokens = act.hidden_states_q.shape[0] + hidden_size = act.hidden_states_q.shape[1] * 2 # FP4 packed + + # trtllm-gen requires the nvfp4 activation scale as float8_e4m3fn; the + # canonical Pack may carry it as raw uint8 bytes. + hidden_states_scale = act.hidden_states_scale + if hidden_states_scale.dtype == torch.uint8: + hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn) + + # Packed pre-routed top-k ids: ((expert_id - offset) << 16) | bf16(weight) + ids = act.selected_experts - self._local_expert_offset + weight_bf16_bits = ( + act.final_scales.to(torch.bfloat16).view(torch.int16).to(torch.int32) + ) + topk_ids = (ids << 16) | (weight_bf16_bits & 0xFFFF) + + output = act.hidden_states_q.new_empty( + (num_tokens, hidden_size), dtype=torch.bfloat16 + ) + # PackedPrecomputed still requires a (kernel-side) topk_weights buffer: + # the raw op declares it non-Optional. The high-level wrapper allocates + # an empty bf16 placeholder here; we mirror that since we bypass it. + expert_weights = act.final_scales.new_empty( + (num_tokens, routing.top_k), dtype=torch.bfloat16 + ) + moe_inputs = MoeRunnerInputs( + output=output, + routing_logits=None, + topk_ids=topk_ids, + expert_weights=expert_weights, + hidden_states=act.hidden_states_q, + hidden_states_scale=hidden_states_scale, + gemm1_lora_delta=None, + per_token_scale=None, + ) + + # Static (num_tokens-invariant) launch arguments for the fp4 branch of + # MoERunner.forward. None-valued entries are the optional gemm bias / + # swiglu beta-clamp / per-token-scale paths not used by the MVP. + self._static_kwargs = dict( + routing_input_mode=RoutingInputMode.PackedPrecomputed, + routing_bias=None, + gemm1_weights=v["gemm1_weights"], + gemm1_weights_scale=v["gemm1_weights_scale"], + gemm1_bias=None, + gemm1_alpha=v.get("gemm1_alpha"), + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=v["gemm2_weights"], + gemm2_weights_scale=v["gemm2_weights_scale"], + gemm2_bias=None, + output1_scale_scalar=v.get("output1_scale_scalar"), + output1_scale_gate_scalar=v.get("output1_scale_gate_scalar"), + output2_scale_scalar=v.get("output2_scale_scalar"), + per_token_scale=None, + num_experts=routing.num_experts, + n_group=routing.n_group, + topk_group=routing.topk_group, + local_expert_offset=self._local_expert_offset, + routed_scaling_factor=routing.routed_scaling_factor, + routing_method_type=int(routing.method), + do_finalize=self.config.execution.do_finalize, + enable_pdl=self._enable_pdl, + ) + + self._ensure_inner(hidden_size) + # Reuse the inner runner's tuning-config builder so the num_tokens + # buckets honor ExecutionConfig.tune_max_num_tokens (CR5). + self.tuning_config = self._inner._make_tuning_config( + moe_inputs, + tune_max_num_tokens=self._tune_max_num_tokens, + # Match the canonical trtllm-gen wrappers' profiling regime so + # choose_one() tunes under the same conditions as deployment + # (otherwise it can cache a tactic picked under a different regime). + use_cuda_graph=True, + use_cold_l2_cache=True, + ) + return moe_inputs.to_list() + + def __hash__(self): + return hash(("trtllm_fp4_routed",)) diff --git a/flashinfer/tllm_enums.py b/flashinfer/tllm_enums.py index d302bd1d994..25368063adc 100644 --- a/flashinfer/tllm_enums.py +++ b/flashinfer/tllm_enums.py @@ -27,6 +27,12 @@ class RoutingMethodType(IntEnum): # Unspecified Unspecified = (9,) + # Eval-safe repr (``RoutingMethodType.Default`` rather than IntEnum's default + # ````) so configs that embed this member + # round-trip through ``eval(repr(cfg))`` β€” relied on by the unified MoE API. + def __repr__(self) -> str: + return f"{type(self).__name__}.{self.name}" + # Copied from csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/common.h class ActivationType(IntEnum): @@ -41,6 +47,19 @@ class ActivationType(IntEnum): Identity = 8 InvalidType = 9 + # Eval-safe repr β€” see ``RoutingMethodType.__repr__``. + def __repr__(self) -> str: + return f"{type(self).__name__}.{self.name}" + + @property + def is_gated(self) -> bool: + """True for activations that consume a gate branch (SwiGLU family).""" + return self in ( + ActivationType.Swiglu, + ActivationType.Geglu, + ActivationType.SwigluBias, + ) + class DtypeTrtllmGen(IntEnum): def __new__(cls, block_format_bit, signed_bit, integer_bit, num_bits, uid): diff --git a/tests/moe/test_unified_moe.py b/tests/moe/test_unified_moe.py new file mode 100644 index 00000000000..1e9096ecdf0 --- /dev/null +++ b/tests/moe/test_unified_moe.py @@ -0,0 +1,925 @@ +"""Tests for the unified MoE API (config dataclasses + MoELayer + Packs). + +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Two halves: + + * CPU-only config/dataclass tests (no GPU or JIT). These track the actual MVP + API surface (single-knob ``QuantVariant``, explicit + ``BackendOptions(candidates=(...))``); see + ``docs/design_docs/flashinfer_moe_api.md`` Β§10 CR1. + + * SM100 (Blackwell) GPU tests for ``MoELayer`` + Packs: shared-bf16-reference + accuracy, autotune candidate visitation, CUDA-graph replay, and the + expert-parallel offset packing (CR3). Scope: NVFP4, pre-routed path. +""" + +from __future__ import annotations + +import dataclasses + +import pytest +import torch + +from flashinfer.autotuner import autotune +from flashinfer.fused_moe import ( + MoEActivationPack, + MoELayer, + MoEWeightPack, + TrtllmFp4RoutedRunner, +) +from flashinfer.fused_moe.api import ( + ActivationConfig, + ActivationType, + BackendOptions, + CuteDslConfig, + CutlassConfig, + ExecutionConfig, + ExpertConfig, + MoEConfig, + QuantConfig, + QuantVariant, + RoutingConfig, + RoutingMethodType, + TrtllmBf16Config, + TrtllmFp4Config, + TrtllmFp8BlockConfig, + TrtllmFp8PerTensorConfig, + TrtllmMxInt4Config, +) + +# Reuse the canonical reference implementation + accuracy helpers from the +# existing CuteDSL test β€” keeps tolerance bounds consistent across tests. +from tests.moe.test_cute_dsl_fused_moe import ( # noqa: E402 + check_accuracy, + compute_reference_moe_fp4, + create_moe_tensors, + is_sm100_family, +) + + +# --------------------------------------------------------------------------- +# Enum repr round-trip +# --------------------------------------------------------------------------- + + +class TestEnumRepr: + @pytest.mark.parametrize("member", list(RoutingMethodType)) + def test_routing_method_repr(self, member): + assert eval(repr(member)) == member + + @pytest.mark.parametrize("member", list(ActivationType)) + def test_activation_repr(self, member): + assert eval(repr(member)) == member + + @pytest.mark.parametrize("member", list(QuantVariant)) + def test_quant_variant_repr(self, member): + assert eval(repr(member)) == member + + +# --------------------------------------------------------------------------- +# ActivationType helpers +# --------------------------------------------------------------------------- + + +class TestActivation: + def test_is_gated(self): + assert ActivationType.Swiglu.is_gated + assert ActivationType.Geglu.is_gated + assert ActivationType.SwigluBias.is_gated + assert not ActivationType.Identity.is_gated + assert not ActivationType.Relu2.is_gated + assert not ActivationType.Gelu.is_gated + + +# --------------------------------------------------------------------------- +# Config immutability +# --------------------------------------------------------------------------- + + +class TestImmutability: + def test_routing_config_frozen(self): + cfg = RoutingConfig(num_experts=64, top_k=8) + with pytest.raises(dataclasses.FrozenInstanceError): + cfg.top_k = 4 + + def test_quant_config_frozen(self): + cfg = QuantConfig(variant=QuantVariant.FP8PerTensor) + with pytest.raises(dataclasses.FrozenInstanceError): + cfg.variant = QuantVariant.BF16 + + def test_moe_config_frozen(self): + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=512), + ) + with pytest.raises(dataclasses.FrozenInstanceError): + cfg.routing = RoutingConfig(num_experts=16, top_k=2) + + +# --------------------------------------------------------------------------- +# Config repr round-trip (critical for repro serialization) +# --------------------------------------------------------------------------- + + +def _eval_repr(obj): + """Evaluate repr(obj) in the config namespace β€” must reconstruct the object.""" + from flashinfer.fused_moe import api as ns + + return eval( + repr(obj), {k: getattr(ns, k) for k in dir(ns) if not k.startswith("_")} + ) + + +class TestReprRoundTrip: + def test_routing_config_minimal(self): + cfg = RoutingConfig(num_experts=64, top_k=8) + assert _eval_repr(cfg) == cfg + + def test_routing_config_full(self): + cfg = RoutingConfig( + num_experts=256, + top_k=8, + method=RoutingMethodType.DeepSeekV3, + n_group=8, + topk_group=4, + routed_scaling_factor=1.0, + ) + assert _eval_repr(cfg) == cfg + + @pytest.mark.parametrize("variant", list(QuantVariant)) + def test_quant_config(self, variant): + cfg = QuantConfig(variant=variant) + assert _eval_repr(cfg) == cfg + + def test_activation_config(self): + for act in ActivationType: + cfg = ActivationConfig(type=act) + assert _eval_repr(cfg) == cfg + + def test_expert_config(self): + cfg = ExpertConfig( + intermediate_size=2048, local_expert_offset=4, local_num_experts=8 + ) + assert _eval_repr(cfg) == cfg + + def test_execution_config_default(self): + cfg = ExecutionConfig() + assert _eval_repr(cfg) == cfg + + def test_execution_config_custom(self): + cfg = ExecutionConfig( + do_finalize=False, enable_pdl=True, tune_max_num_tokens=1024 + ) + assert _eval_repr(cfg) == cfg + + def test_backend_options_multi(self): + opts = BackendOptions(candidates=(TrtllmFp4Config(), CutlassConfig())) + reconstructed = _eval_repr(opts) + assert len(reconstructed) == 2 + assert isinstance(reconstructed.candidates[0], TrtllmFp4Config) + assert isinstance(reconstructed.candidates[1], CutlassConfig) + + def test_backend_options_single(self): + opts = BackendOptions(candidates=(TrtllmFp8PerTensorConfig(),)) + reconstructed = _eval_repr(opts) + assert len(reconstructed) == 1 + assert isinstance(reconstructed.candidates[0], TrtllmFp8PerTensorConfig) + + def test_moe_config_minimal(self): + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=512), + ) + assert _eval_repr(cfg) == cfg + + def test_moe_config_full(self): + cfg = MoEConfig( + routing=RoutingConfig( + num_experts=256, + top_k=8, + method=RoutingMethodType.DeepSeekV3, + n_group=8, + topk_group=4, + routed_scaling_factor=1.0, + ), + quant=QuantConfig(variant=QuantVariant.MxFp8), + experts=ExpertConfig(intermediate_size=2048, local_num_experts=32), + activation=ActivationConfig(type=ActivationType.Geglu), + backend=BackendOptions( + candidates=(TrtllmFp8BlockConfig(), CutlassConfig()) + ), + execution=ExecutionConfig(enable_pdl=True, tune_max_num_tokens=4096), + ) + assert _eval_repr(cfg) == cfg + + +# --------------------------------------------------------------------------- +# BackendOptions +# --------------------------------------------------------------------------- + + +class TestBackendOptions: + def test_explicit_candidates(self): + opts = BackendOptions(candidates=(TrtllmFp4Config(), CutlassConfig())) + assert isinstance(opts, BackendOptions) + assert len(opts) == 2 + + def test_multiple_candidates(self): + opts = BackendOptions( + candidates=(TrtllmFp4Config(), TrtllmFp8BlockConfig(), CutlassConfig()) + ) + assert len(opts) == 3 + + def test_valid_for_filtering(self): + opts = BackendOptions( + candidates=(TrtllmBf16Config(), TrtllmFp8BlockConfig(), CutlassConfig()) + ) + # sm80: BF16 requires 100+, FP8Block requires 80+, Cutlass is universal + valid = opts.valid_for(80) + assert len(valid) == 2 + assert isinstance(valid[0], TrtllmFp8BlockConfig) + assert isinstance(valid[1], CutlassConfig) + + def test_valid_for_blackwell(self): + opts = BackendOptions( + candidates=(TrtllmBf16Config(), TrtllmFp8BlockConfig(), CutlassConfig()) + ) + valid = opts.valid_for(100) + assert len(valid) == 3 + + def test_iteration(self): + opts = BackendOptions(candidates=(TrtllmFp4Config(), CutlassConfig())) + items = list(opts) + assert len(items) == 2 + assert any(isinstance(c, TrtllmFp4Config) for c in items) + assert any(isinstance(c, CutlassConfig) for c in items) + + def test_empty(self): + opts = BackendOptions() + assert len(opts) == 0 + assert opts.valid_for(100) == [] + + +# --------------------------------------------------------------------------- +# QuantConfig +# --------------------------------------------------------------------------- + + +class TestQuantConfig: + def test_default_is_bf16(self): + assert QuantConfig().variant == QuantVariant.BF16 + + def test_explicit_variant(self): + assert QuantConfig(variant=QuantVariant.NVFP4).variant == QuantVariant.NVFP4 + + @pytest.mark.parametrize("variant", list(QuantVariant)) + def test_all_variants_constructible(self, variant): + assert QuantConfig(variant=variant).variant is variant + + +# --------------------------------------------------------------------------- +# MoEConfig dict-unpacking protocol +# --------------------------------------------------------------------------- + + +class TestMoEConfigDictProtocol: + def test_keys(self): + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=512), + ) + keys = list(cfg.keys()) + assert "routing" in keys + assert "quant" in keys + assert "experts" in keys + assert "activation" in keys + assert "backend" in keys + assert "execution" in keys + + def test_getitem(self): + routing = RoutingConfig(num_experts=8, top_k=2) + cfg = MoEConfig( + routing=routing, + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=512), + ) + assert cfg["routing"] is routing + + def test_unpack(self): + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=512), + ) + d = dict(**cfg) + assert isinstance(d["routing"], RoutingConfig) + assert isinstance(d["backend"], BackendOptions) + + +# --------------------------------------------------------------------------- +# Dataclasses.replace for immutable overrides +# --------------------------------------------------------------------------- + + +class TestImmutableReplace: + def test_replace_quant(self): + cfg = MoEConfig( + routing=RoutingConfig(num_experts=64, top_k=8), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=2048), + ) + fp8_cfg = dataclasses.replace( + cfg, + quant=QuantConfig(variant=QuantVariant.DeepSeekFp8), + ) + assert fp8_cfg.quant.variant == QuantVariant.DeepSeekFp8 + assert cfg.quant.variant == QuantVariant.BF16 # original unchanged + + def test_replace_backend(self): + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig(intermediate_size=512), + ) + narrow = dataclasses.replace(cfg, backend=BackendOptions((CutlassConfig(),))) + assert len(narrow.backend) == 1 + + +# --------------------------------------------------------------------------- +# Hashability (needed for cache keys) +# --------------------------------------------------------------------------- + + +class TestHashability: + def test_routing_config_hashable(self): + a = RoutingConfig(num_experts=64, top_k=8) + b = RoutingConfig(num_experts=64, top_k=8) + assert hash(a) == hash(b) + assert {a, b} == {a} + + def test_moe_config_hashable(self): + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=512), + backend=BackendOptions(candidates=(TrtllmBf16Config(), CutlassConfig())), + ) + # Must not raise + h = hash(cfg) + assert isinstance(h, int) + + def test_moe_config_as_dict_key(self): + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=512), + ) + d = {cfg: "value"} + assert d[cfg] == "value" + + +# --------------------------------------------------------------------------- +# ActivationConfig singletons +# --------------------------------------------------------------------------- + + +class TestActivationConfigSingletons: + def test_singletons_exist(self): + assert ActivationConfig.swiglu == ActivationConfig(ActivationType.Swiglu) + assert ActivationConfig.geglu == ActivationConfig(ActivationType.Geglu) + assert ActivationConfig.relu2 == ActivationConfig(ActivationType.Relu2) + assert ActivationConfig.identity == ActivationConfig(ActivationType.Identity) + + def test_singleton_is_gated(self): + assert ActivationConfig.swiglu.is_gated + assert not ActivationConfig.identity.is_gated + + +# --------------------------------------------------------------------------- +# Expressiveness: can we represent the existing test configurations? +# --------------------------------------------------------------------------- + + +class TestExpressiveness: + """Verify that the unified config can express every existing test scenario. + + Each scenario maps a legacy flat-API configuration onto the single-knob + ``QuantVariant`` surface. + """ + + def test_trtllm_fp4_deepseekv3(self): + """The most common DeepSeek-V3 FP4 config from test_trtllm_gen_fused_moe.py.""" + cfg = MoEConfig( + routing=RoutingConfig( + num_experts=256, + top_k=8, + method=RoutingMethodType.DeepSeekV3, + n_group=8, + topk_group=4, + routed_scaling_factor=1.0, + ), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig(intermediate_size=1024), + activation=ActivationConfig(type=ActivationType.Swiglu), + backend=BackendOptions(candidates=(TrtllmFp4Config(), CutlassConfig())), + ) + assert cfg.routing.method == RoutingMethodType.DeepSeekV3 + assert cfg.quant.variant == QuantVariant.NVFP4 + assert cfg.activation.is_gated + + def test_trtllm_fp8_block_mxfp8(self): + """MxFP8 block-scale config.""" + cfg = MoEConfig( + routing=RoutingConfig( + num_experts=64, + top_k=8, + method=RoutingMethodType.Renormalize, + ), + quant=QuantConfig(variant=QuantVariant.MxFp8), + experts=ExpertConfig(intermediate_size=512), + activation=ActivationConfig(type=ActivationType.Swiglu), + backend=BackendOptions( + candidates=(TrtllmFp8BlockConfig(), CutlassConfig()) + ), + ) + assert cfg.quant.variant == QuantVariant.MxFp8 + + def test_trtllm_fp8_per_tensor(self): + """Per-tensor FP8 config.""" + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.FP8PerTensor), + experts=ExpertConfig(intermediate_size=512), + backend=BackendOptions((TrtllmFp8PerTensorConfig(),)), + ) + assert cfg.quant.variant == QuantVariant.FP8PerTensor + + def test_trtllm_bf16(self): + """BF16 unquantized config.""" + cfg = MoEConfig( + routing=RoutingConfig( + num_experts=8, + top_k=2, + method=RoutingMethodType.Renormalize, + ), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=512), + backend=BackendOptions(candidates=(TrtllmBf16Config(), CutlassConfig())), + ) + assert cfg.quant.variant == QuantVariant.BF16 + + def test_trtllm_mxint4(self): + """MxInt4 config.""" + cfg = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.MxInt4), + experts=ExpertConfig(intermediate_size=512), + backend=BackendOptions((TrtllmMxInt4Config(),)), + ) + assert cfg.quant.variant == QuantVariant.MxInt4 + + def test_cutlass_modular_fp8(self): + """CUTLASS modular (pre-routed) FP8 config.""" + cfg = MoEConfig( + routing=RoutingConfig(num_experts=64, top_k=8), + quant=QuantConfig(variant=QuantVariant.DeepSeekFp8), + experts=ExpertConfig(intermediate_size=2048), + activation=ActivationConfig(type=ActivationType.Swiglu), + backend=BackendOptions((CutlassConfig(),)), + ) + # CUTLASS uses modular (pre-routed) dispatch β€” supplied at call time via + # MoEActivationPack (selected_experts/final_scales), not via config + assert any(isinstance(c, CutlassConfig) for c in cfg.backend) + + def test_cutedsl_nvfp4(self): + """CuteDSL NVFP4 config.""" + cfg = MoEConfig( + routing=RoutingConfig(num_experts=64, top_k=8), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig(intermediate_size=1024), + activation=ActivationConfig(type=ActivationType.Swiglu), + backend=BackendOptions(candidates=(CuteDslConfig(), CutlassConfig())), + ) + assert any(isinstance(c, CuteDslConfig) for c in cfg.backend) + + def test_expert_parallel(self): + """Config with expert parallelism (EP).""" + cfg = MoEConfig( + routing=RoutingConfig(num_experts=256, top_k=8), + quant=QuantConfig(variant=QuantVariant.DeepSeekFp8), + experts=ExpertConfig( + intermediate_size=2048, + local_expert_offset=32, + local_num_experts=32, + ), + ) + assert cfg.experts.local_expert_offset == 32 + assert cfg.experts.local_num_experts == 32 + + def test_llama4_routing(self): + """Llama4 top-1 sigmoid routing.""" + cfg = MoEConfig( + routing=RoutingConfig( + num_experts=16, + top_k=1, + method=RoutingMethodType.Llama4, + ), + quant=QuantConfig(variant=QuantVariant.BF16), + experts=ExpertConfig(intermediate_size=4096), + ) + assert cfg.routing.method == RoutingMethodType.Llama4 + assert cfg.routing.top_k == 1 + + def test_qwen3_renormalize_naive(self): + """Qwen3 RenormalizeNaive routing.""" + cfg = MoEConfig( + routing=RoutingConfig( + num_experts=64, + top_k=8, + method=RoutingMethodType.RenormalizeNaive, + ), + quant=QuantConfig(variant=QuantVariant.DeepSeekFp8), + experts=ExpertConfig(intermediate_size=1024), + ) + assert cfg.routing.method == RoutingMethodType.RenormalizeNaive + + +# --------------------------------------------------------------------------- +# MoELayer MVP fail-fast validation (CR6) +# --------------------------------------------------------------------------- +# These exercise MoELayer._validate_mvp_scope, which runs at construction time +# before any device/runner setup, so they need no GPU. + + +class TestMoELayerMVPValidation: + def _nvfp4_swiglu(self, **overrides): + base = dict( + routing=RoutingConfig(num_experts=32, top_k=2), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig(intermediate_size=512), + activation=ActivationConfig(type=ActivationType.Swiglu), + ) + base.update(overrides) + return MoEConfig(**base) + + @pytest.mark.parametrize( + "variant", + [v for v in QuantVariant if v is not QuantVariant.NVFP4], + ) + def test_non_nvfp4_quant_rejected(self, variant): + from flashinfer.fused_moe import MoELayer + + cfg = self._nvfp4_swiglu(quant=QuantConfig(variant=variant)) + with pytest.raises(NotImplementedError, match="NVFP4"): + MoELayer(cfg) + + @pytest.mark.parametrize( + "act", + [a for a in ActivationType if a is not ActivationType.Swiglu], + ) + def test_non_swiglu_activation_rejected(self, act): + from flashinfer.fused_moe import MoELayer + + cfg = self._nvfp4_swiglu(activation=ActivationConfig(type=act)) + with pytest.raises(NotImplementedError, match="Swiglu"): + MoELayer(cfg) + + +sm100_required = pytest.mark.skipif( + not is_sm100_family(), + reason="Unified NVFP4 MoE requires SM100 family (Blackwell SM100/SM103)", +) + + +# Small-scale geometry for fast accuracy + dispatch tests. +SMALL = dict(hidden_size=1024, intermediate_size=512, num_experts=32, top_k=2) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_packs_and_config( + num_tokens: int, + *, + hidden_size: int, + intermediate_size: int, + num_experts: int, + top_k: int, + local_num_experts: int | None = None, + max_tokens: int | None = None, +): + """Build (act_pack, weight_pack, config, tensors_dict) for a given shape. + + ``tensors_dict`` contains the original bf16 reference weights used to + compute ground truth via ``compute_reference_moe_fp4``. + """ + local_num_experts = local_num_experts or num_experts + max_tokens = max_tokens or max(num_tokens, 8192) + device = torch.device("cuda", torch.cuda.current_device()) + + # CuteDSL view comes pre-built by create_moe_tensors + bf16 refs + tensors = create_moe_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + num_local_experts=local_num_experts, + top_k=top_k, + ) + + act_pack = MoEActivationPack( + hidden_states_q=tensors["x"], + hidden_states_scale=tensors["x_sf"].squeeze(-1), + selected_experts=tensors["token_selected_experts"], + final_scales=tensors["token_final_scales"], + ) + + weight_pack = MoEWeightPack() + weight_pack.prepare_for( + "cute_dsl_nvfp4", + { + "w1_weight": tensors["w1_weight"], + "w1_weight_sf": tensors["w1_weight_sf"], + "w1_alpha": tensors["w1_alpha"], + "fc2_input_scale": tensors["fc2_input_scale"], + "w2_weight": tensors["w2_weight"], + "w2_weight_sf": tensors["w2_weight_sf"], + "w2_alpha": tensors["w2_alpha"], + }, + ) + weight_pack.prepare_for( + "trtllm_fp4_routed", + TrtllmFp4Config.prepare_weights( + tensors["w1_weight_bf16"], + tensors["w2_weight_bf16"], + num_local_experts=local_num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + ), + ) + + config = MoEConfig( + routing=RoutingConfig(num_experts=num_experts, top_k=top_k), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig( + intermediate_size=intermediate_size, + local_num_experts=local_num_experts, + ), + activation=ActivationConfig(), + backend=BackendOptions(candidates=(CuteDslConfig(), TrtllmFp4Config())), + execution=ExecutionConfig(tune_max_num_tokens=max_tokens), + ) + return act_pack, weight_pack, config, tensors + + +# --------------------------------------------------------------------------- +# 1. Accuracy +# --------------------------------------------------------------------------- + + +def _compute_ref(act_pack, tensors, shape): + """bf16 ground-truth MoE output for the given pack + shape.""" + return compute_reference_moe_fp4( + hidden_states=tensors["x_bf16"].float().cuda(), + gemm1_weights=tensors["w1_weight_bf16"].float().cuda(), + gemm2_weights=tensors["w2_weight_bf16"].float().cuda(), + token_selected_experts=act_pack.selected_experts, + token_final_scales=act_pack.final_scales, + num_tokens=act_pack.num_tokens, + num_experts=shape["num_experts"], + top_k=shape["top_k"], + hidden_size=shape["hidden_size"], + intermediate_size=shape["intermediate_size"], + fc2_input_scale=tensors["fc2_input_scale"], + ) + + +@sm100_required +class TestUnifiedMoEAccuracy: + """Every path compares against the same bf16 reference. + + Catches cases where both backends are wrong in the same way, which a + cross-backend agreement test would miss. + """ + + @pytest.mark.parametrize("num_tokens", [128, 512]) + def test_layer_output_matches_reference(self, num_tokens): + """MoELayer end-to-end output matches bf16 reference.""" + act_pack, weight_pack, config, tensors = _make_packs_and_config( + num_tokens, **SMALL + ) + + with autotune(True): + layer = MoELayer(config) + out = layer(act_pack, weight_pack) + + ref = _compute_ref(act_pack, tensors, SMALL) + passed, pct, atol = check_accuracy(out, ref) + assert passed, ( + f"MoELayer output {pct * 100:.2f}% within tolerance " + f"(atol={atol:.4f}) vs bf16 reference at num_tokens={num_tokens}" + ) + + @pytest.mark.parametrize("backend_key", ["cute_dsl_nvfp4", "trtllm_fp4_routed"]) + def test_each_backend_matches_reference(self, backend_key): + """Each candidate backend individually matches the same bf16 reference. + + If either backend's weight view were semantically wrong, its output + would diverge from the shared reference β€” even in cases where it + might agree with the other backend. + """ + act_pack, weight_pack, config, tensors = _make_packs_and_config(256, **SMALL) + layer = MoELayer(config) + runner = next(r for r in layer.runners if r.backend_key == backend_key) + + inputs = runner.pack_inputs(act_pack, weight_pack) + out = runner.forward(inputs, tactic=-1) + + ref = _compute_ref(act_pack, tensors, SMALL) + passed, pct, atol = check_accuracy(out, ref) + assert passed, ( + f"{backend_key}: {pct * 100:.2f}% within tolerance " + f"(atol={atol:.4f}) vs bf16 reference" + ) + + +# --------------------------------------------------------------------------- +# 2. Dispatch plumbing +# --------------------------------------------------------------------------- + + +@sm100_required +class TestUnifiedMoEDispatch: + """Plumbing tests β€” invariants MoELayer must guarantee.""" + + def test_autotune_visits_all_candidate_backends(self): + """The autotuner actually profiles every candidate backend. + + Shape-robust: doesn't commit to a specific winner (those change with + kernel updates), just asserts each backend's `forward` was invoked + during _select_winner. + """ + act_pack, weight_pack, config, _ = _make_packs_and_config(256, **SMALL) + layer = MoELayer(config) + + # Wrap each runner's forward to count invocations. + call_counts: dict = {} + for runner in layer.runners: + key = runner.backend_key + call_counts[key] = 0 + original = runner.forward + + def counted(*args, __key=key, __orig=original, **kwargs): + call_counts[__key] += 1 + return __orig(*args, **kwargs) + + runner.forward = counted # type: ignore[assignment] + + with autotune(True): + _ = layer(act_pack, weight_pack) + + assert len(call_counts) >= 2, ( + f"Expected β‰₯2 candidate backends, got {list(call_counts)}" + ) + for key, count in call_counts.items(): + assert count > 0, ( + f"Backend {key!r} was never invoked β€” autotuner skipped it " + f"(call counts: {call_counts})" + ) + + def test_graph_capture_replay(self): + """CUDA-graph-captured replay matches eager output.""" + num_tokens = 256 + act_pack, weight_pack, config, _ = _make_packs_and_config( + num_tokens, max_tokens=num_tokens, **SMALL + ) + + # Warm up: populate autotune cache + stabilize allocator + with autotune(True): + layer = MoELayer(config) + for _ in range(3): + _ = layer(act_pack, weight_pack) + for _ in range(3): + _ = layer(act_pack, weight_pack) + + eager_out = layer(act_pack, weight_pack).clone() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + captured = layer(act_pack, weight_pack) + + for _ in range(10): + g.replay() + torch.cuda.synchronize() + + passed, pct, atol = check_accuracy(captured, eager_out) + assert passed, ( + f"Graph replay diverged from eager: {pct * 100:.2f}% within " + f"tolerance (atol={atol:.4f})" + ) + + +# --------------------------------------------------------------------------- +# 3. Expert-parallel offset (CR3) +# --------------------------------------------------------------------------- + + +@sm100_required +class TestTrtllmEPOffset: + """TRTLLM routed packing must shift global expert ids down by the local + shard offset. + + The packed int32 top-k id is ``((expert_id - offset) << 16) | bf16(weight)``; + the kernel indexes local experts as ``[0, local_num_experts)``. Before the + CR3 fix, ``pack_inputs`` defaulted the offset to 0, so a nonzero local shard + produced out-of-range local expert ids. + """ + + @pytest.mark.parametrize("local_expert_offset", [0, 32, 96]) + def test_pack_inputs_applies_local_expert_offset(self, local_expert_offset): + device = torch.device("cuda", torch.cuda.current_device()) + num_experts = 128 + local_num_experts = 32 + top_k = 4 + num_tokens = 16 + hidden_size = 256 + sf_vec_size = 16 + + config = MoEConfig( + routing=RoutingConfig(num_experts=num_experts, top_k=top_k), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig( + intermediate_size=512, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + ), + ) + runner = TrtllmFp4RoutedRunner(config, device=device) + + # Global expert ids drawn from this rank's local shard. + selected_experts = ( + torch.randint(0, local_num_experts, (num_tokens, top_k), device=device).to( + torch.int32 + ) + + local_expert_offset + ) + final_scales = torch.rand(num_tokens, top_k, device=device) + act_pack = MoEActivationPack( + hidden_states_q=torch.zeros( + num_tokens, hidden_size // 2, dtype=torch.uint8, device=device + ), + hidden_states_scale=torch.zeros( + num_tokens, hidden_size // sf_vec_size, dtype=torch.uint8, device=device + ).view(torch.float8_e4m3fn), + selected_experts=selected_experts, + final_scales=final_scales, + ) + + # pack_inputs only passes weight tensors through; dummies suffice since + # we inspect topk_ids only (no kernel launch). + weight_pack = MoEWeightPack() + weight_pack.prepare_for( + "trtllm_fp4_routed", + { + "gemm1_weights": torch.empty(0, device=device), + "gemm1_weights_scale": torch.empty(0, device=device), + "gemm1_alpha": torch.empty(0, device=device), + "gemm2_weights": torch.empty(0, device=device), + "gemm2_weights_scale": torch.empty(0, device=device), + }, + ) + + from flashinfer.fused_moe.core import MoEInputs + + inputs = runner.pack_inputs(act_pack, weight_pack) + topk_ids = MoEInputs.from_list(inputs).topk_ids + + # Upper 16 bits hold the (offset-shifted) local expert id. + decoded_local_ids = topk_ids >> 16 + expected_local_ids = selected_experts - local_expert_offset + assert torch.equal(decoded_local_ids, expected_local_ids), ( + f"offset={local_expert_offset}: packed local ids " + f"{decoded_local_ids} != expected {expected_local_ids}" + ) + # Local ids must land inside the kernel's [0, local_num_experts) range. + assert int(decoded_local_ids.min()) >= 0 + assert int(decoded_local_ids.max()) < local_num_experts diff --git a/tests/moe/test_unified_moe_fuzz.py b/tests/moe/test_unified_moe_fuzz.py new file mode 100644 index 00000000000..11beb0afdf5 --- /dev/null +++ b/tests/moe/test_unified_moe_fuzz.py @@ -0,0 +1,788 @@ +"""Forward-compatible fuzzer for the unified MoE API (``MoELayer`` + Packs, PR #3093). + +Drives the **real user-facing surface** -- one ``MoEConfig`` -> the API's own +``XxxConfig.prepare_weights(w1_bf16, w2_bf16, ...)`` marshalling -> ``MoELayer``'s per-backend +runners -- so what's under test is the production dispatch + the new ``prepare.py`` scale/layout +plumbing, where the low-precision-MoE bugs cluster (GH #2356/#2485/#2907/#3068). + +Forward-compatible by construction: + * Backends are discovered from the live runner registry (``layer.runners``); an unwired backend + is skipped and is tested the moment its runner lands -- zero new code. + * Weight prep is the uniform ``cfg.prepare_weights(w1_bf16, w2_bf16, **shape)`` (canonical bf16 in, + quantize+layout done internally). + * Only the per-DTYPE pieces live in one ``_DTYPE`` table: how to make golden inputs, how to build + the activation pack, and the canonical reference recipe. + +Config space: random shapes deliberately incl non-pow2 (aligned) hidden/intermediate + odd/ +tile-boundary token counts (real-model + #2907/#3168 territory), routing-load skew, all under a +weight-memory budget so one config never hogs the GPU (parallel-CI-friendly), plus a few curated +larger-end shapes. Large expert counts are reached with small H/I and/or **expert-parallel shards** +(global>local + ``local_expert_offset``, the real deployment shape), not by filling the GPU. + +A small ``_KNOWN_FAILURES`` ledger xfails already-filed bugs (e.g. trtllm EP offset>0 -> all-zero, +EP_OFFSET_FINDING.md): the case is still *run* so the suite stays green on a tracked bug yet flags +loudly the day it starts passing (fixed). A crash is never tolerated -- only a wrong answer. + +Verification model (single mode, uniform -- every config that runs gets the same checks): + 1. **no crash / no NaN-Inf** where the reference is finite. + 2. **numeric vs the canonical quant-aware reference.** The reference is the *authority*: it + defines the one true numerical recipe (exactly-representable inputs + the fp4 intermediate + requant), so a backend that invents a different recipe is wrong by definition. Inputs are + snapped to the exact nvfp4 grid and sparsified, so input quantization is lossless and the + gemm reductions are short -- a structural bug (dropped expert / wrong index / wrong scale role) + produces a gross error instead of one averaged away. Tolerance is set to the fp4 + intermediate-requant floor (~0.08), far tighter than a dense-random comparison. + 3. **determinism, per-backend contract.** A backend declared deterministic must reproduce + bitwise across reruns; a non-deterministic one (CuteDSL's atomic-scatter finalize) is exempt. + Flags are established empirically (CRC across runs) -- see ``_DETERMINISTIC``. + 4. **output-buffer poison.** The kernel owns its output (an uninitialized ``new_empty`` inside + the runner's ``pack_inputs``) and MoE finalize *accumulates* into it -- so the result must not + depend on the buffer being clean. We fill it with garbage + NaN/Inf and re-assert #1+#2. + torch's caching allocator usually returns clean memory and hides this; JAX/XLA donates dirty + buffers (the GH-6158764 padding-leak class), so this is the torch->JAX buffer-hygiene guard. + 5. **autotune-tactic sweep.** EVERY valid tactic (``get_valid_tactics``), not just the default, + must match the reference -- the autotuner-picks-a-corrupting-tactic class (#3168/#3227) on the + real ``MoELayer`` dispatch, since the autotuner enumerates these same tactics in production. + 6. **autotune-ON, real production path** (gated to a config subset for cost): drive + ``with autotune(True): layer(...)`` so ``MoELayer._select_winner`` *profiles* every tactic of + every runner (the #3168 profiling-IMA / #2749 profiling-crash class -- distinct from #5, which + replays tactics outside the tuner) then selects + caches a winner; the autotuned output must + still match the authoritative reference. Skipped when a candidate has a known failure (the + tuner could legitimately pick the broken backend). + 7. **device-state probe** after each config: a context-corrupting IMA surfaces as a failed + alloc/launch or non-finite probe, turning silent corruption into a clean failure. + +A sibling SCENARIO test ``test_autotune_cache_coherence`` covers the one autotune surface this +per-config fuzz structurally can't: the cross-call **winner cache**. It drives ONE persistent +``MoELayer`` through a token-count *sequence* (incl bucket boundaries 4095/4096/4097) under +``autotune(True)`` -- filling the per-bucket cache, crossing shapes, then re-running earlier counts +to force cache hits -- asserting each output stays correct, so a stale / mis-keyed cached winner +reused for a different shape is caught (the #2933-adjacent class). + +(Cross-backend agreement is intentionally NOT a check: with an authoritative tight reference, a +deviating backend is caught by #2 directly, and #2 also names which backend -- so a cross-backend +comparison adds no pass/fail power, only redundancy. See the design discussion.) + +Coverage today: NVFP4 (CuteDSL + TRTLLM-FP4-routed) on SM100 -- the only wired MVP runners. + +OPT-IN: this suite is gated behind FLASHINFER_UMOE_FUZZ (see the pytestmark below) and is +SKIPPED unless that env var is set -- waived in CI pending gh #3547 and root-cause of a +whole-process device-side-assert abort that would block B200 CI. Run it explicitly: + FLASHINFER_UMOE_FUZZ=1 CUDA_HOME= CUDA_VISIBLE_DEVICES= \ + pytest tests/moe/test_unified_moe_fuzz.py +NOTE: `pytest --forked` does NOT work here (CUDA inits at collection -> +"Cannot re-initialize CUDA in forked subprocess"); for crash-isolated enumeration run each +test id in its own process instead (see var/03-ssh-docker-workflow.md). +Env: FLASHINFER_UMOE_FUZZ_NUM_TESTS (default 80), FLASHINFER_UMOE_FUZZ_SEED (default 0). + +------------------------------------------------------------------------------------------------ +EXTENDING (cheap, by design): + * New backend -> nothing to do: it is auto-discovered from ``_BACKEND_RUNNERS`` the moment its + runner registers and ``supported(sm)`` is true. If it ships with a tracked bug, add one + ``_KNOWN_FAILURES`` entry (the case still RUNS; an xpass then flags the fix). + * New dtype -> add ONE ``DTypeHandler`` to ``_DTYPE`` (snap / make_act_pack / reference / poison + / tolerances). Everything else (config gen, all 7 checks, the cache test) is dtype-generic. + +ROADMAP -- what's left, ranked by the 2026-06-09 audit of 51 past MoE GH issues (the full-build-out +harness catches ~60% full / ~91% touched of the 35 in-scope; ~31% are structurally out-of-scope). +Full synthesis lives in the cuDNN-project auto-memory ``flashinfer_quality_fuzzers.md``. Each item +names the issue class it closes: + 1. [HIGHEST LEVERAGE -- infra, not code] Blackwell/SM120 PR-CI runner. PR-gating CI tops out at + H100/SM90, so the dominant ~36% fp4/MoE bug class is collected-then-SKIPPED at PR time. This + harness only protects users on arches it actually RUNS on -- no oracle improvement beats + provisioning the runner. (This is the #1 documented escape reason for the whole MoE class.) + 2. N-run (>=10) stress per config + a PER-TEST TIMEOUT, under ``--forked`` isolation. Turns the + intermittent PARTIALs into CAUGHT: #2569 (intermittent NaN), #2933 (concurrency-bucket hang). + A single pass samples a "hangs 1-in-100" failure poorly. NOTE a *deterministic* hang is + already catchable -- but TODAY it blocks the whole job; add ``@pytest.mark.timeout`` so it + fails ONE test cleanly. ``--forked`` needs lazy-CUDA-init handled (the cuDNN _replacement + Heisenbug lesson: forked children must init CUDA fresh). + 3. Curated PRODUCTION shapes: seed the generator with real model dims (DeepSeek-V3, Llama-4, + Qwen3, Mixtral) + dense tile-window enumeration (every M in [tile-2, tile+2] around each + kernel's tile boundary). Closes the shape-luck escapes #3310 (Llama-4-Scout "no kernel") and + #2732 (Qwen3-Coder wrong output) that a synthetic 4096+-1 sweep misses. + 4. BUILD-MANIFEST oracle: enumerate the advertised (backend x quant x arch) support matrix and + assert each combo actually INSTANTIATES a kernel. Closes #2501 (W4A8 autotune fail) -- an + un-compiled combo is invisible to a runtime fuzzer (the harness assumes backends are built). + 5. [DEEPEST -- the one structurally-weak oracle] Tighten the QUANTIZED-NUMERIC net. Today check + #2 compares to ONE authoritative quant-aware reference at the fp4 requant-floor tolerance + (~10% of ||ref||inf). That floor HIDES sub-10% accuracy regressions (#2356 small-scale, #3103 + minority-NaN), and the reference -- because it must itself encode the quant recipe -- can be + "wrong the same way" as a kernel (no independent fp32 ideal, unlike bf16). The real fix is the + unified API standardizing ONE intermediate-activation-scale POLICY (the design doc's + role-named QuantSpec; gh #3548): once every backend honors one DECLARED recipe, a single fp32 + reference computing that recipe becomes an INDEPENDENT authority for all of them (and + calibrated checkpoints become expressible). Until then: add a small-scale / edge-magnitude + input axis and document the floor. + +OUT OF SCOPE for this single-GPU correctness harness (must live elsewhere, do NOT try to force in): + * multi-GPU / EP>1 / TP collective hangs & deadlocks (#3279 EP=8, #3530 TP8) -> a distributed + (2-8 GPU) test tier with collective-aware timeouts. (Single-GPU EP SHARDS -- global>local + + local_expert_offset -- ARE in scope and tested here; the COLLECTIVE is not.) + * perf/latency regressions (#2671) -> perf-CI with per-kernel latency baselines. A wrong-but-fast + tactic IS caught (check #5/#6); a correct-but-slow one is invisible by design. + * framework-glue triggers (vLLM/SGLang dispatch sequences #3427, #3390) -> integration tests. + The underlying KERNEL bug is in scope here IF invoked directly; the live-dispatch trigger isn't. + * build / cubin / packaging (#3466 missing SM103 cubin, #3344 _sm100f-only) -> an arch-coverage + manifest check in build CI (related to roadmap #4 but at the .so/cubin level). + +POINTERS for future agents (point me at this file and I know the rest): + * Full context (this fuzzer + the older adapter/GEMM fuzzers + the audit + findings): cuDNN- + project auto-memory ``flashinfer_quality_fuzzers.md``. + * Bugs THIS fuzzer found + filed: gh #3547 (trtllm EP offset>0 all-zero -- encoded in the + ``_KNOWN_FAILURES`` ledger below: xfailed but still RUN, so an xpass announces the fix) and + gh #3548 (activation global-scale gap == roadmap #5's scale-policy fix). + * Findings writeups: flashinfer_triage/EP_OFFSET_FINDING.md, flashinfer_triage/WEIGHT_SCALE_FINDING.md. + * The unified API under test: PR #3093 (branch ``moe_api``); this fuzzer is PR aleozlx/flashinfer#6 + (branch ``yanxu/unified-moe-api-fuzzer``). +""" + +from __future__ import annotations + +import os +import random +import warnings +from dataclasses import dataclass +from typing import Callable + +import pytest +import torch +import torch.nn.functional as F + +from flashinfer.autotuner import autotune +from flashinfer.fp4_quantization import fp4_quantize +from flashinfer.fused_moe import MoEActivationPack, MoELayer, MoEWeightPack +from flashinfer.fused_moe.api import ( + ActivationConfig, + BackendOptions, + CuteDslConfig, + ExecutionConfig, + ExpertConfig, + MoEConfig, + QuantConfig, + QuantVariant, + RoutingConfig, + TrtllmFp4Config, +) +from flashinfer.fused_moe.layer import _BACKEND_RUNNERS +from flashinfer.quantization import e2m1_and_ufp8sf_scale_to_float +from flashinfer.utils import get_compute_capability + +NUM_TESTS = int(os.environ.get("FLASHINFER_UMOE_FUZZ_NUM_TESTS", "80")) +BASE_SEED = int(os.environ.get("FLASHINFER_UMOE_FUZZ_SEED", "0")) + +# --- CI-safety gate: OPT-IN ---------------------------------------------------------------- +# Waived in CI pending gh #3547 + root-cause of a whole-process abort. Running the SM100 fuzzer +# in a single `pytest` process can hit `CUDA error: device-side assert triggered` -> +# `Fatal Python error: Aborted`, which would BLOCK B200 CI (an abort fails the whole job, not one +# test). Notes from triage (2026-06-09): per-config isolation (one process each) passes 68/86 +# incl. EP offset>0 -- so the abort is NOT cleanly attributable to one config (the gh #3547 EP +# case returns tolerated zeros, no assert, under torch.cuda.synchronize); it surfaces only in the +# accumulated single-process run that CI uses. `pytest --forked` can't isolate it here either +# (CUDA inits at collection -> "Cannot re-initialize CUDA in forked subprocess"). Until the abort +# is root-caused and #3547 fixed (follow-up PR), this suite is opt-in: set FLASHINFER_UMOE_FUZZ=1 +# to run it (developer / nightly / SM100 box). Unset (CI default) -> collected-and-skipped, so it +# never launches a kernel and cannot abort the job. +pytestmark = pytest.mark.skipif( + not os.environ.get("FLASHINFER_UMOE_FUZZ"), + reason="opt-in fuzzer (set FLASHINFER_UMOE_FUZZ=1); waived in CI pending gh #3547 and " + "root-cause of the whole-process device-side-assert abort", +) + +# Per-backend determinism contract, established empirically (CRC across reruns) + confirmed against +# code. A "True" backend MUST reproduce bitwise; flip to False only with evidence (and ideally an +# upstream note), because a deterministic->non-deterministic regression is exactly a bug to catch. +_DETERMINISTIC = { + "trtllm_fp4_routed": True, # bitwise-stable across reruns in calibration + "cute_dsl_nvfp4": False, # atomic scatter-add finalize -> non-bit-exact by design +} + +# Known-bug ledger: (backend_key, predicate(cfg)) -> reason. A matching (backend, config) is run but +# its correctness failure is TOLERATED (xfail) -- this keeps the suite green on a filed-and-tracked +# bug while still EXERCISING it, so the day the bug is fixed the case starts passing and we get a loud +# "unexpectedly passed -> remove this entry" signal. A crash is never tolerated (only wrong answers). +_KNOWN_FAILURES = [ + ( + "trtllm_fp4_routed", + lambda c: c.expert_offset > 0, + "trtllm EP local_expert_offset>0 -> all-zero output (offset applied twice); gh #3547", + ), +] + + +def _known_failure(backend_key, cfg): + for bk, predicate, reason in _KNOWN_FAILURES: + if bk == backend_key and predicate(cfg): + return reason + return None + + +# --------------------------------------------------------------------------- +# nvfp4 exact-grid snapping: make a tensor a fixed point of the kernel's quantizer, so input +# quantization is lossless (the kernel re-quantizes to the same fp4 values) and only the +# intermediate requant remains as quant error. +# --------------------------------------------------------------------------- +def _snap_to_nvfp4(t: torch.Tensor) -> torch.Tensor: + one = torch.tensor([1.0], device=t.device) + flat = t.reshape(-1, t.shape[-1]).to(torch.bfloat16) + packed, scale = fp4_quantize( + flat, global_scale=one, sf_vec_size=16, is_sf_swizzled_layout=False + ) + deq = e2m1_and_ufp8sf_scale_to_float( + packed.cpu(), + scale.cpu().view(torch.uint8).reshape(-1), + (1.0 / one).cpu(), + 16, + 1, + False, + ) + return deq.reshape(t.shape).to(t.device, torch.bfloat16) + + +# --------------------------------------------------------------------------- +# Per-DTYPE handlers: golden input generation, activation pack, canonical reference recipe. +# The ONLY place a new quant variant needs code; new *backends* for a variant are free. +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class DTypeHandler: + variant: QuantVariant + candidate_configs: ( + tuple # all plausible backend config classes; unwired ones auto-skip + ) + snap: Callable # bf16 tensor -> exactly-representable fixed point for this dtype + make_act_pack: Callable # (x, selected_experts, final_scales) -> MoEActivationPack + reference: Callable # (x, w1, w2, selected_experts, final_scales, I) -> fp32 [T,H] authority + poison: Callable # in-place fill a kernel-owned output buffer with garbage + (NaN/Inf if repr.) + out_dtype: torch.dtype # output buffer dtype (used to locate it in the inputs list) + atol_frac: float # numeric tolerance vs reference = atol_frac * β€–refβ€–βˆž + rtol: float + + +def _nvfp4_poison(buf): + """Fill a bf16 output buffer with large garbage + scattered NaN/Β±Inf. If a kernel reads or + scatter-adds into an uninitialized output instead of fully writing it, the poison leaks and + is caught by no-NaN / numeric. This is the torch->JAX buffer-hygiene guard: torch's caching + allocator usually hands back clean memory (masking the bug), JAX/XLA donates dirty buffers + (the GH-6158764 class).""" + g = torch.randn_like(buf) * 1e4 + flat = g.view(-1) + flat[0::4], flat[1::4], flat[2::4] = float("nan"), float("inf"), float("-inf") + buf.copy_(g) + + +def _nvfp4_act_pack(x, selected_experts, final_scales): + # global activation scale == 1.0 (MVP wires no global-scale field; block scales carry range). + one = torch.tensor([1.0], device=x.device) + packed, scale = fp4_quantize( + x, global_scale=one, sf_vec_size=16, is_sf_swizzled_layout=False + ) + return MoEActivationPack( + hidden_states_q=packed, + hidden_states_scale=scale.squeeze(-1) if scale.dim() > 2 else scale, + selected_experts=selected_experts, + final_scales=final_scales, + ) + + +def _nvfp4_reference( + x, w1, w2, selected_experts, final_scales, intermediate_size, expert_offset=0 +): + """Canonical nvfp4 MoE recipe (the authority): exact-fp4 inputs (lossless), SwiGLU = + silu(2nd half)*(1st half), then the intermediate is re-quantized to fp4 (block-scaled, + gs=1.0) before gemm2 -- matching what the kernels do. w1/w2 hold only this rank's LOCAL + experts; a token routed to global id ``g`` uses local weight ``g - expert_offset`` (EP).""" + x32, half = x.float(), intermediate_size + out = torch.zeros_like(x32) + for local_e in range(w1.shape[0]): + mask = ( + selected_experts == local_e + expert_offset + ) # global id of this local expert + if not mask.any(): + continue + tok, nth = torch.where(mask) + gate, up = w1[local_e][half:, :].float(), w1[local_e][:half, :].float() + inter = F.silu(x32[tok] @ gate.t()) * (x32[tok] @ up.t()) + inter = _snap_to_nvfp4( + inter.to(torch.bfloat16) + ).float() # intermediate fp4 requant + out[tok] += final_scales[tok, nth, None] * (inter @ w2[local_e].float().t()) + return out + + +_DTYPE = { + QuantVariant.NVFP4: DTypeHandler( + variant=QuantVariant.NVFP4, + candidate_configs=(CuteDslConfig, TrtllmFp4Config), + snap=_snap_to_nvfp4, + make_act_pack=_nvfp4_act_pack, + reference=_nvfp4_reference, + poison=_nvfp4_poison, + out_dtype=torch.bfloat16, + atol_frac=0.15, # calibrated: obs ratio ≀0.077 (fp4 intermediate-requant floor) + rtol=0.1, + ), + # FP8 / MXFP4 / MXINT4 / BF16 add one entry each as their runners are wired upstream. +} + + +# --------------------------------------------------------------------------- +# Config generation: random shapes + routing-load skew (an orthogonal axis -- uniform enforcement, +# not a numeric mode, so it never changes which checks apply). +# --------------------------------------------------------------------------- +# Deliberately NOT all powers of two: real models use non-pow2 (aligned) hidden/intermediate +# (Llama 14336/11008, DeepSeek-MoE 1408/1536, Qwen 18944), and #2907 was an intermediate-padding +# accuracy bug. H/I stay %16 for fp4 block alignment; if a kernel rejects a shape we skip it. +_HIDDEN = [256, 512, 1024, 1536, 2048, 3072] # 1536/3072 aligned-non-pow2 +_INTERMED = [ + 256, + 512, + 768, + 1024, + 1408, + 1536, +] # 768/1408/1536 aligned-non-pow2 (#2907 class) +_EXPERTS = [ + 8, + 16, + 32, + 64, + 72, + 128, + 160, + 256, + 512, +] # 72/160 non-pow2; 512 needs small H/I (budget) +_TOPK = [1, 2, 4, 6, 8] # 6 non-pow2 +# num_tokens is runtime batch*seqlen -- arbitrary. Sweep odd + tile/autotune-bucket boundaries +# (the #3168 16384-bucket / 4095-4097 tile-remainder class), not just clean powers of two. +_TOKENS = [1, 2, 3, 7, 17, 64, 127, 129, 256, 1024, 2048, 4095, 4096, 4097] +_ROUTE = ["uniform", "uniform", "hot1", "imbalanced"] + +# Per-test weight footprint cap so one fuzz config never hogs the GPU (parallel-CI-friendly) and the +# CPU exact-grid snap stays sub-few-seconds. ~500M bf16 weight elems β‰ˆ 1 GB. The cap naturally pairs +# a large expert count with small H/I (and rejects giant-H/I x many-experts), matching real EP-sharded +# deployments where no single rank holds thousands of full-size experts. +_WEIGHT_ELEM_BUDGET = 500_000_000 + + +def _weight_elems(num_experts, hidden, intermediate): + return num_experts * (2 * intermediate * hidden + hidden * intermediate) # w1 + w2 + + +@dataclass(frozen=True) +class Cfg: + num_tokens: int + hidden: int + intermediate: int + num_experts: int # GLOBAL expert count (RoutingConfig.num_experts) + top_k: int + variant: str + route: str + seed: int + local_experts: int = 0 # this rank's shard; 0 -> non-EP (== num_experts) + expert_offset: int = 0 # global id of this shard's first expert (EP) + + @property + def n_local(self): # experts actually held + computed on this rank + return self.local_experts or self.num_experts + + @property + def is_ep(self): + return self.expert_offset > 0 or self.n_local != self.num_experts + + @property + def label(self): + ep = f"L{self.n_local}o{self.expert_offset}_" if self.is_ep else "" + return ( + f"{self.variant}_{self.route}_e{self.num_experts}_{ep}k{self.top_k}_" + f"t{self.num_tokens}_h{self.hidden}_i{self.intermediate}_s{self.seed}" + ) + + +def _gen(seed): + rng = random.Random(seed) + # Resample shape until the LOCAL-shard weights fit the budget (modest per-test GPU footprint). + for _ in range(64): + ne, h, i = rng.choice(_EXPERTS), rng.choice(_HIDDEN), rng.choice(_INTERMED) + # ~30%: expert-parallel shard -- split the global experts and pick a shard (offset>0). This + # is how large MoE actually runs (no rank holds all experts) and exercises the offset path. + local, offset = ne, 0 + shards = rng.choice([2, 4]) + if rng.random() < 0.3 and ne >= 16 and ne % shards == 0: + local = ne // shards + offset = local * rng.randrange(shards) + if _weight_elems(local, h, i) <= _WEIGHT_ELEM_BUDGET: + break + return Cfg( + num_tokens=rng.choice(_TOKENS), + hidden=h, + intermediate=i, + num_experts=ne, + top_k=rng.choice( + [t for t in _TOPK if t <= local] + ), # route within the local shard + variant="nvfp4", # only wired variant today; expands with _DTYPE + route=rng.choice(_ROUTE), + seed=seed, + local_experts=local, + expert_offset=offset, + ) + + +# A few curated "larger end of the common range" shapes (all within the weight budget) so the big +# end is always represented, not left to chance: many-experts, large-hidden+many-tokens, and max-experts. +_CURATED = [ + Cfg( + 256, 1024, 512, 256, 8, "nvfp4", "uniform", 900_001 + ), # DeepSeek-ish: 256 experts, top_k 8 + Cfg( + 4096, 2048, 1024, 64, 8, "nvfp4", "uniform", 900_002 + ), # large hidden + many tokens + Cfg( + 2048, 1024, 1024, 128, 6, "nvfp4", "imbalanced", 900_003 + ), # empty-expert load + mid size + Cfg( + 512, 512, 512, 512, 4, "nvfp4", "hot1", 900_004 + ), # max expert count (small H/I) +] +_CONFIGS = _CURATED + [_gen(BASE_SEED + i) for i in range(NUM_TESTS)] + + +def _master(cfg, handler): + """Sparse, exactly-representable bf16 inputs + host routing. Sparsity keeps the gemm reductions + short so a structural bug isn't averaged away; exact-grid snapping makes input quant lossless. + Weights cover only this rank's LOCAL shard (E_local); routing selects within the shard's GLOBAL + id range [offset, offset+E_local) (the EP contract -- non-EP is offset=0, E_local=num_experts).""" + g = torch.Generator(device="cuda").manual_seed(cfg.seed) + E_local, H, I, T = cfg.n_local, cfg.hidden, cfg.intermediate, cfg.num_tokens + + def sparse(*shape): + dense = torch.randn(*shape, device="cuda", generator=g) + keep = torch.rand(shape, device="cuda", generator=g) >= 0.75 # ~75% zeros + return handler.snap(dense * keep) + + x, w1, w2 = sparse(T, H), sparse(E_local, 2 * I, H), sparse(E_local, H, I) + + logits = torch.randn(T, E_local, device="cuda", generator=g) # over the local shard + if cfg.route == "hot1": # pile every token onto one expert + logits[:, 0] += 50.0 + elif cfg.route == "imbalanced": # rank-skew -> some experts get zero tokens + logits += torch.linspace(8.0, -8.0, E_local, device="cuda") + weights = F.softmax(logits, dim=1, dtype=torch.float32) + weights, local_sel = torch.topk(weights, cfg.top_k, dim=-1) + final_scales = (weights / weights.sum(dim=-1, keepdim=True)).float() + selected_experts = (local_sel + cfg.expert_offset).to( + torch.int32 + ) # local -> global ids + return x, w1, w2, selected_experts, final_scales + + +_SKIP_SUBSTR = ( + "not supported", + "unsupported", + "no valid", + "not implemented", + "must be", + "divisible", + "requires", + "only support", +) +_CRASH_SUBSTR = ("cuda error", "illegal memory", "misaligned", "device-side assert") + + +def _is_unsupported(e): + msg = str(e).lower() + if any(c in msg for c in _CRASH_SUBSTR): + return False # a crash is always a finding, never "unsupported" + return isinstance(e, NotImplementedError) or any(s in msg for s in _SKIP_SUBSTR) + + +@pytest.mark.parametrize("cfg", _CONFIGS, ids=[c.label for c in _CONFIGS]) +def test_unified_moe_fuzz(cfg): + if not torch.cuda.is_available(): + pytest.skip("no CUDA") + # Full per-config determinism so any failure reproduces from the seed in the test id alone. + # Shapes (random.Random(seed)) and the input tensors (a per-config torch.Generator) are already + # seeded; this also pins the two GLOBAL-RNG draws -- the poison garbage and the device probe -- + # so the entire run is bitwise-reproducible from `cfg.seed`. (Autotune winner selection is + # timing-based and may vary run-to-run, but every tactic is validated, so a correctness failure + # still reproduces via the tactic sweep regardless of which winner the tuner picks.) + torch.manual_seed(cfg.seed) + sm = get_compute_capability(torch.device("cuda:0")) + sm = sm[0] * 10 + sm[1] + + handler = _DTYPE[QuantVariant.NVFP4] + dev = torch.device("cuda") + # Backend *config classes* whose runner is registered in the live MoELayer registry AND valid + # on this arch. A newly-wired backend lands here automatically. + wired_backends = [ + BackendCfg + for BackendCfg in handler.candidate_configs + if BackendCfg in _BACKEND_RUNNERS and BackendCfg.supported(sm) + ] + if not wired_backends: + pytest.skip(f"no wired backend for {cfg.variant} on SM{sm}") + + x, w1, w2, selected_experts, final_scales = _master(cfg, handler) + ref = handler.reference( + x, w1, w2, selected_experts, final_scales, cfg.intermediate, cfg.expert_offset + ) + atol = handler.atol_frac * ref.abs().max().item() + 1e-3 + rtol = handler.rtol + + # One activation pack + one weight pack with each backend's native view, all built from the + # SAME bf16 inputs (this rank's LOCAL shard) via the API's uniform prepare_weights. + act_pack = handler.make_act_pack(x, selected_experts, final_scales) + weight_pack = MoEWeightPack() + for BackendCfg in wired_backends: + weight_pack.prepare_for( + _BACKEND_RUNNERS[BackendCfg].backend_key, + BackendCfg.prepare_weights( + w1, + w2, + num_local_experts=cfg.n_local, + hidden_size=cfg.hidden, + intermediate_size=cfg.intermediate, + device=dev, + ), + ) + + config = MoEConfig( + routing=RoutingConfig(num_experts=cfg.num_experts, top_k=cfg.top_k), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig( + intermediate_size=cfg.intermediate, + local_num_experts=cfg.n_local, + local_expert_offset=cfg.expert_offset, + ), + activation=ActivationConfig(), + backend=BackendOptions( + candidates=tuple(BackendCfg() for BackendCfg in wired_backends) + ), + execution=ExecutionConfig(tune_max_num_tokens=max(cfg.num_tokens, 8192)), + ) + + try: + layer = MoELayer(config) + except Exception as e: + if _is_unsupported(e): + pytest.skip(f"MoELayer rejected {cfg.label}: {e}") + raise + + out_shape = (cfg.num_tokens, cfg.hidden) + + def run(runner, poison=False): + inputs = runner.pack_inputs(act_pack, weight_pack) + if poison: + # The output buffer is a kernel-owned `new_empty` tensor inside the inputs list + # (cute_dsl idx 11, trtllm the `output=`); locate it by dtype+shape and poison it. + bufs = [ + t + for t in inputs + if torch.is_tensor(t) + and t.dtype == handler.out_dtype + and tuple(t.shape) == out_shape + ] + assert bufs, "could not locate the output buffer in pack_inputs to poison" + for b in bufs: + handler.poison(b) + out = runner.forward(inputs, tactic=-1) + out = (out[0] if isinstance(out, (list, tuple)) else out).float() + torch.cuda.synchronize() + return out + + def assert_correct(out, tag): + # no NaN/Inf where the reference is finite. + n_bad = int(((~torch.isfinite(out)) & torch.isfinite(ref)).sum().item()) + assert n_bad == 0, f"{tag}: {n_bad} non-finite outputs vs finite reference" + # numeric vs the canonical quant-aware reference (the authority), magnitude-scaled. + abs_diff = (out - ref).abs() + over_tol = abs_diff > (atol + rtol * ref.abs()) + if over_tol.any(): + pytest.fail( + f"{tag}: {int(over_tol.sum())}/{out.numel()} elems exceed " + f"(rtol={rtol} atol={atol:.3g}); max|diff|={abs_diff.max().item():.4g}, " + f"β€–refβ€–βˆž={ref.abs().max().item():.4g}" + ) + + def check_backend(runner, out, tag): + # (1)+(2) no-NaN + numeric vs the authoritative reference, on a clean run. + assert_correct(out, tag) + # (3) determinism per the backend's contract: deterministic backends must reproduce + # bitwise; non-deterministic ones (atomic-scatter finalize) are exempt. + if _DETERMINISTIC.get(runner.backend_key, False): + if not torch.equal(out, run(runner)): + drift = (out - run(runner)).abs().max().item() + pytest.fail( + f"{tag}: declared DETERMINISTIC but not bitwise-reproducible " + f"(max abs diff {drift:.3e})" + ) + # (4) output-buffer poison: the kernel owns its (uninitialized `new_empty`) output, so the + # result must NOT depend on it being clean. torch's allocator usually hands back zeros and + # hides this; poisoning forces it -- the torch->JAX hazard (GH-6158764 padding leak). + assert_correct(run(runner, poison=True), f"{tag} [poisoned-output]") + # (5) autotune-tactic sweep: EVERY valid tactic must be correct, not just the default -- + # the autotuner-picks-a-bad-tactic class (#3168/#3227) on the real MoELayer dispatch. + inputs = runner.pack_inputs(act_pack, weight_pack) + try: + tactics = runner.get_valid_tactics(inputs, None) + except Exception: + tactics = [] # backend needs a profile object -> skip the sweep (default tactic stands) + for tactic in tactics: + o = runner.forward(inputs, tactic=tactic) + o = (o[0] if isinstance(o, (list, tuple)) else o).float() + torch.cuda.synchronize() + assert_correct(o, f"{tag} [tactic={tactic}]") + + n_ran = 0 + for runner in layer.runners: + try: + out = run(runner) + except Exception as e: + if _is_unsupported(e): + continue # backend rejects this shape -> skip; a crash re-raises + raise + tag = f"{runner.backend_key} {cfg.label}" + n_ran += 1 + + known = _known_failure(runner.backend_key, cfg) + if known: # tracked bug -> run it, tolerate a wrong answer, but flag if it starts passing + try: + check_backend(runner, out, tag) + except (AssertionError, pytest.fail.Exception): + continue + warnings.warn( + f"{tag}: KNOWN-FAILURE unexpectedly PASSED -- fixed? remove from " + f"_KNOWN_FAILURES ({known})", + stacklevel=2, + ) + else: + check_backend(runner, out, tag) + + if n_ran == 0: + pytest.skip(f"no runner ran {cfg.label} on SM{sm}") + + # (6) autotune-ON: drive the REAL production path -- MoELayer._select_winner profiles every + # tactic of every runner (the #3168 profiling-IMA class) then selects + caches a winner; the + # autotuned output must match the authoritative reference. Gated to a subset (profiling is slow) + # and skipped if a candidate has a known failure (the tuner could pick the broken backend). + autotune_due = cfg.seed % 4 == 0 and not any( + _known_failure(_BACKEND_RUNNERS[B].backend_key, cfg) for B in wired_backends + ) + if autotune_due: + with autotune(True): + a_out = layer(act_pack, weight_pack) + a_out = (a_out[0] if isinstance(a_out, (list, tuple)) else a_out).float() + torch.cuda.synchronize() + assert_correct( + a_out, f"{cfg.label} [autotune-ON winner={layer.winner_backend}]" + ) + + # (7) device-state probe: a context-corrupting IMA in any backend above would surface here as a + # failed alloc/launch or a non-finite probe, turning a silent corruption into a clean failure. + probe = torch.randn(2048, device="cuda") * 2.0 + torch.cuda.synchronize() + assert torch.isfinite(probe).all(), ( + f"{cfg.label}: CUDA context corrupted after MoE run" + ) + + +# --------------------------------------------------------------------------- +# Sibling SCENARIO test (not per-config-stateless): the autotune CACHE is cross-call state, which +# the fuzz test (fresh MoELayer per config) structurally can't reach. So drive ONE persistent layer +# through a token-count SEQUENCE under autotune -- fill the per-bucket winner cache, cross shapes, +# and re-run earlier counts to force cache hits. A stale / mis-keyed cached winner reused for a +# different shape would produce a wrong answer here. (Shares the harness's snap/reference/prep.) +# --------------------------------------------------------------------------- +_CACHE_BASES = [ + (32, 1024, 512), + (128, 512, 512), +] # (experts, hidden, intermediate), non-EP +_CACHE_TOKEN_SEQ = [ + 16, + 256, + 4095, + 4096, + 4097, + 256, + 16, +] # buckets + boundaries + cache-hit re-runs + + +@pytest.mark.parametrize( + "base", _CACHE_BASES, ids=[f"e{e}h{h}i{i}" for e, h, i in _CACHE_BASES] +) +def test_autotune_cache_coherence(base): + if not torch.cuda.is_available(): + pytest.skip("no CUDA") + sm = get_compute_capability(torch.device("cuda:0")) + sm = sm[0] * 10 + sm[1] + handler = _DTYPE[QuantVariant.NVFP4] + dev = torch.device("cuda") + wired = [ + B + for B in handler.candidate_configs + if B in _BACKEND_RUNNERS and B.supported(sm) + ] + if not wired: + pytest.skip(f"no wired backend on SM{sm}") + + E, H, I = base + top_k = 4 + g = torch.Generator(device="cuda").manual_seed(12345) + + def sparse(*shape): + dense = torch.randn(*shape, device="cuda", generator=g) + return handler.snap( + dense * (torch.rand(shape, device="cuda", generator=g) >= 0.75) + ) + + # Fixed weights + ONE layer instance; the cache lives across the whole sequence. + w1, w2 = sparse(E, 2 * I, H), sparse(E, H, I) + weight_pack = MoEWeightPack() + for B in wired: + weight_pack.prepare_for( + _BACKEND_RUNNERS[B].backend_key, + B.prepare_weights( + w1, + w2, + num_local_experts=E, + hidden_size=H, + intermediate_size=I, + device=dev, + ), + ) + layer = MoELayer( + MoEConfig( + routing=RoutingConfig(num_experts=E, top_k=top_k), + quant=QuantConfig(variant=QuantVariant.NVFP4), + experts=ExpertConfig(intermediate_size=I, local_num_experts=E), + activation=ActivationConfig(), + backend=BackendOptions(candidates=tuple(B() for B in wired)), + execution=ExecutionConfig(tune_max_num_tokens=max(_CACHE_TOKEN_SEQ)), + ) + ) + + with autotune( + True + ): # fill the per-bucket cache on first sight; hit it on the re-runs + for num_tokens in _CACHE_TOKEN_SEQ: + x = sparse(num_tokens, H) + w = F.softmax(torch.randn(num_tokens, E, device="cuda", generator=g), dim=1) + w, sel = torch.topk(w, top_k, dim=-1) + final = (w / w.sum(dim=-1, keepdim=True)).float() + sel = sel.to(torch.int32) + act = handler.make_act_pack(x, sel, final) + ref = handler.reference(x, w1, w2, sel, final, I, 0) + out = layer(act, weight_pack) + out = (out[0] if isinstance(out, (list, tuple)) else out).float() + torch.cuda.synchronize() + tag = f"cache-seq T={num_tokens} (winner={layer.winner_backend}) {base}" + n_bad = int(((~torch.isfinite(out)) & torch.isfinite(ref)).sum().item()) + assert n_bad == 0, f"{tag}: {n_bad} non-finite outputs" + atol = handler.atol_frac * ref.abs().max().item() + 1e-3 + over = (out - ref).abs() > (atol + handler.rtol * ref.abs()) + assert not over.any(), ( + f"{tag}: {int(over.sum())} elems exceed tol " + f"(max|diff|={(out - ref).abs().max().item():.4g}) -- stale/mis-keyed cached winner?" + ) From 9e6a2805324bb3a2892badf62eb53e1baf085ae0 Mon Sep 17 00:00:00 2001 From: Perkz Zheng <67892460+PerkzZheng@users.noreply.github.com> Date: Thu, 11 Jun 2026 00:48:23 +0800 Subject: [PATCH 09/13] Support smaller DSv4 sparse MLA head counts (#3545) ## Summary - allow DeepSeek V4 sparse MLA validation for 8/16/32 heads in addition to 64/128 - add targeted BF16/FP8 sparse MLA coverage for 8, 16, and 32 heads ## Tests - `python -m ruff check flashinfer/mla/_core.py tests/attention/test_trtllm_gen_sparse_mla_dsv4.py` - `FLASHINFER_DISABLE_VERSION_CHECK=1 python -m pytest --collect-only tests/attention/test_trtllm_gen_sparse_mla_dsv4.py -q` - `FLASHINFER_DISABLE_VERSION_CHECK=1 python -m pytest tests/attention/test_trtllm_gen_sparse_mla_dsv4.py -k 'h8 or h16 or h32' -vv` ## Summary by CodeRabbit * **New Features** * Expanded DeepSeek V4 sparse MLA support to head counts: 8, 16, 32, 64, and 128. * **Bug Fixes** * Updated input validation to accept the expanded set of supported head counts. * **Tests** * Added test coverage for sparse MLA decode scenarios with smaller head counts, multiple data types, and both KV layout variants. Co-authored-by: Perkz Zheng --- flashinfer/mla/_core.py | 4 +- .../test_trtllm_gen_sparse_mla_dsv4.py | 54 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/flashinfer/mla/_core.py b/flashinfer/mla/_core.py index 369602c29f8..8308da1fb57 100644 --- a/flashinfer/mla/_core.py +++ b/flashinfer/mla/_core.py @@ -348,8 +348,8 @@ def _check_dsv4_sparse_mla_inputs( ) if head_dim != 512: raise ValueError(f"Expected query head dim 512, got {head_dim}") - if num_heads not in (64, 128): - raise ValueError(f"Expected 64 or 128 query heads, got {num_heads}") + if num_heads not in (8, 16, 32, 64, 128): + raise ValueError(f"Expected 8, 16, 32, 64, or 128 query heads, got {num_heads}") if ( sparse_indices is None diff --git a/tests/attention/test_trtllm_gen_sparse_mla_dsv4.py b/tests/attention/test_trtllm_gen_sparse_mla_dsv4.py index 1cca2003fd3..274623b2f72 100644 --- a/tests/attention/test_trtllm_gen_sparse_mla_dsv4.py +++ b/tests/attention/test_trtllm_gen_sparse_mla_dsv4.py @@ -256,6 +256,60 @@ def add_case(**kwargs) -> None: sparse_case="swa128+topk128x", ) + # Smaller head counts use the sparse MLA small-head kernel selection path. + swa_seq_len, c4_seq_len, c128_seq_len = DECODE_SEQ_LEN_CASES[0] + c128_topk = _round_up(c128_seq_len + (DECODE_BATCH_SIZE - 1) * C128_PAGE_SIZE, 4) + for h_q in (8, 16, 32): + for dtype in (torch.bfloat16, torch.float8_e4m3fn): + for kv_layout in ("HND", "NHD"): + add_case( + b=DECODE_BATCH_SIZE, + h_q=h_q, + s_q=DECODE_Q_LEN, + h_kv=1, + s_kv=swa_seq_len, + is_varlen=True, + topk=DSV4_SWA_TOPK, + block_size=SWA_PAGE_SIZE, + dtype=dtype, + kv_layout=kv_layout, + sparse_case="swa128", + ) + add_case( + b=DECODE_BATCH_SIZE, + h_q=h_q, + s_q=DECODE_Q_LEN, + h_kv=1, + s_kv=swa_seq_len, + is_varlen=True, + topk=DSV4_SWA_TOPK, + extra_s_k=c4_seq_len, + extra_topk=h_q * 8, + block_size=SWA_PAGE_SIZE, + extra_block_size=C4_PAGE_SIZE, + have_extra_topk_length=True, + dtype=dtype, + kv_layout=kv_layout, + sparse_case="swa128+topk4x", + ) + add_case( + b=DECODE_BATCH_SIZE, + h_q=h_q, + s_q=DECODE_Q_LEN, + h_kv=1, + s_kv=swa_seq_len, + is_varlen=True, + topk=DSV4_SWA_TOPK, + extra_s_k=c128_seq_len, + extra_topk=c128_topk, + block_size=SWA_PAGE_SIZE, + extra_block_size=C128_PAGE_SIZE, + have_extra_topk_length=True, + dtype=dtype, + kv_layout=kv_layout, + sparse_case="swa128+topk128x", + ) + # Guard seq_lens-driven SWA masking. With q_len == kv_len == 128, early # query tokens have fewer than 128 valid SWA entries, so the kernel must use # real seq_lens instead of treating every SWA tile as full. From a0cb4e7804385662c33aed401f319d931d397a22 Mon Sep 17 00:00:00 2001 From: "Brian K. Ryu" Date: Wed, 10 Jun 2026 16:33:45 -0700 Subject: [PATCH 10/13] feat(bench): add GDN routines to flashinfer_benchmark.py (#3572) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## πŸ“Œ Description Adds Gated Delta Net (GDN) support to `flashinfer_benchmark.py` as a new `gdn` routine category: Added Routines: - `gated_delta_rule_decode` β€” T=1 decode; `--state_layout {pretranspose,nontranspose}`, `--state_dtype {float32,bfloat16}` (bf16 state kernels), `--pool_mode {single,split}` - `gated_delta_rule_mtp` β€” T>=2 multi-token processing with state pool + indices - `chunk_gated_delta_rule` β€” varlen chunked prefill #### Bug fixes in the standalone benches - `bench_gdn_decode.py`: all-layouts mode called the BF16-state kernel without pool indices (column was always N/A); bytes model counted intermediate-state traffic for T>1 even with caching disabled, inflating reported TB/s - `bench_gdn_prefill.py`: fed log-space `g` to the FlashInfer kernel, which takes a linear-space alpha β€” outputs/state were NaN; FLA is now optional instead of a hard exit ## πŸ” 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 have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes ## Summary by CodeRabbit * **New Features** * Added GDN (Gated Delta Net) benchmark family: decode, multi-token (MTP), and chunked prefill with multi-backend support and performance metrics that account for optional intermediate-state caching. * Added Triton-backed reference kernels and optional Triton routing for GDN workloads. * **Bug Fixes** * Prefill baseline is no longer fatal when optional baseline kernel is missing; reporting and columns adjust accordingly. * Corrected gate generation to avoid out-of-domain values. * **Documentation** * Updated docs and sample test lists with GDN routines, flags, backend matrix, and usage notes. --- benchmarks/README.md | 39 +- benchmarks/bench_gdn_decode.py | 868 +------------- benchmarks/bench_gdn_prefill.py | 26 +- benchmarks/flashinfer_benchmark.py | 11 +- benchmarks/gdn_triton_reference.py | 873 ++++++++++++++ .../routines/flashinfer_benchmark_utils.py | 53 + benchmarks/routines/gdn.py | 1067 +++++++++++++++++ benchmarks/samples/sample_testlist.txt | 22 + 8 files changed, 2111 insertions(+), 848 deletions(-) create mode 100644 benchmarks/gdn_triton_reference.py create mode 100644 benchmarks/routines/gdn.py diff --git a/benchmarks/README.md b/benchmarks/README.md index b5fb4e69bca..cdc5faf3902 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -5,11 +5,11 @@ The aim of `flashinfer_benchmark.py` is to provide a single framework for benchm ## Overview This framework provides tools to: -- Benchmark FlashInfer's Attention, GEMM, MOE, Norm, Quantization, Sampling, RoPE, and Mamba API performance from different kernel backends such as FlashAttention2/3, cuDNN, cuBLAS, CUTLASS, CuTe-DSL, TensorRT-LLM, and Triton +- Benchmark FlashInfer's Attention, GEMM, MOE, Norm, Quantization, Sampling, RoPE, Mamba, and GDN API performance from different kernel backends such as FlashAttention2/3, cuDNN, cuBLAS, CUTLASS, CuTe-DSL, TensorRT-LLM, and Triton - Compare performance across different configurations - Batch performance test multiple test cases -Currently supports testing attention, gemm, fused MOE, normalization, quantization, sampling, RoPE, and Mamba APIs: +Currently supports testing attention, gemm, fused MOE, normalization, quantization, sampling, RoPE, Mamba, and GDN (Gated Delta Net) APIs: - Attention: - `BatchDecodeWithPagedKVCacheWrapper` - Decode attention with paged KV cache. - Also supports computationally similar `cudnn_batch_decode_with_kv_cache` and `trtllm_batch_decode_with_kv_cache`. @@ -77,6 +77,10 @@ Currently supports testing attention, gemm, fused MOE, normalization, quantizati - `rope_quantize_fp8_append_paged_kv_cache` - RoPE with FP8 quantization and paged KV cache append (SM8.9+). - Mamba (Selective State Space Models): - `selective_state_update` - Selective state update for Mamba layers (generation phase). Supports both single-token prediction (STP) and multi-token prediction (MTP) via `--cache_steps`. Backends: `flashinfer` (CUDA, architecture-specific kernels for base/SM90/SM100+) and `triton` (reference). +- GDN (Gated Delta Net linear attention, SM90+): + - `gated_delta_rule_decode` - Single-token (T=1) gated delta rule decode. `--state_layout` selects between `gated_delta_rule_decode_pretranspose` ([B, HV, V, K] state, default) and `gated_delta_rule_decode` ([B, HV, K, V] state). `--state_dtype bfloat16` selects the BF16 state kernels (head_size=128, pretranspose only). Backends: `flashinfer` (CuTe-DSL) and `triton` (reference). + - `gated_delta_rule_mtp` - Multi-token (T>=2) gated delta rule for speculative-decoding verification, with a state pool + indices. `--state_dtype float32` uses `gated_delta_rule_mtp`; `--state_dtype bfloat16` uses the BF16 MTP kernel via `gated_delta_rule_decode_pretranspose`. Backends: `flashinfer`, `triton`. + - `chunk_gated_delta_rule` - Chunked GDN prefill over varlen sequences (uniform per-sequence length `--s_qo`). Backends: `flashinfer` (SM90 C++ / SM100 CuTe-DSL) and `fla` (flash-linear-attention Triton baseline, perf-only). ## Quick Start ### Single Test Run @@ -445,6 +449,31 @@ mpirun -np 8 python benchmarks/flashinfer_benchmark.py \ | `--dt_softplus` | Apply softplus to dt before use | | `--backends` | Backends to test: `flashinfer` (default), `triton` (reference). Refcheck compares against Triton reference | +### GDN Flags +Applies to `gated_delta_rule_decode`, `gated_delta_rule_mtp`, and `chunk_gated_delta_rule` (SM90+). + +| Flag | Description | +|-------------------------------|-------------------------------------------------------------------------------------------------------------| +| `--batch_size` | Decode/MTP: number of concurrent requests. Prefill: number of sequences | +| `--num_q_heads` | Number of query heads. Default: 16 | +| `--num_k_heads` | Number of key heads. Default: 16 | +| `--num_v_heads` | Number of value heads (GVA when > `num_q_heads`). Default: 32 | +| `--head_size` | Head dimension (K = V = head_size). Default: 128 | +| `--input_dtype` | Data type for q/k/v/a/b tensors: `bfloat16` (default) or `float16` | +| `--state_dtype` | Recurrent state dtype: `float32` (default) or `bfloat16` (BF16 state kernels; decode/MTP, head_size=128, pretranspose) | +| `--state_layout` | Decode only: `pretranspose` ([B, HV, V, K], default) or `nontranspose` ([B, HV, K, V]) | +| `--pool_mode` | `single` (default, read == write slots) or `split` (pool of 2B; reads slots [0..B), writes [B..2B)) | +| `--seq_len` | MTP only: tokens per request (>= 2). Default: 2 | +| `--s_qo` | Prefill only: per-sequence length (uniform). Default: 2048 | +| `--update_state` | MTP only: write the final state back (`disable_state_update=False`). BF16 state always updates in-place | +| `--cache_intermediate_states` | MTP with `float32` state only: cache per-token intermediate states | +| `--no_qk_l2norm` | Decode/MTP: disable in-kernel Q/K L2 normalization | +| `--backends` | Decode/MTP: `flashinfer` (default), `triton`. Prefill: `flashinfer` (default), `fla` (requires `pip install flash-linear-attention`; perf-only, excluded from refcheck) | + +Notes: +- Refcheck compares against the torch reference in `tests/gdn/reference_delta_rule.py`. +- Prefill pre-L2-normalizes k and calls the kernel with `use_qk_l2norm_in_kernel=False` so the kernel and reference see identical inputs. + ## `flashinfer_benchmark.py` Routine & Backend Support Matrix The following table summarizes the support surface of each routine & backend's on various [CUDA Compute Capabilities](https://developer.nvidia.com/cuda-gpus). @@ -516,6 +545,9 @@ Legend: | **rope_quantize_fp8** | | | | cuda | cuda | cuda | cuda | cuda | | **rope_quantize_fp8_append_paged_kv_cache** | | | | cuda | cuda | cuda | cuda | cuda | | **selective_state_update** | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | flashinfer, triton | +| **gated_delta_rule_decode** | | | | | flashinfer, triton | flashinfer, triton | flashinfer, triton | triton | +| **gated_delta_rule_mtp** | | | | | flashinfer, triton | flashinfer, triton | flashinfer, triton | triton | +| **chunk_gated_delta_rule** | | | | | flashinfer, fla | flashinfer, fla | flashinfer, fla | | Backend Legend: - fa2: FlashAttention2 @@ -533,4 +565,5 @@ Backend Legend: - cute-dsl: FlashInfer CuTe-DSL kernels (Blackwell SM10.0+) - moe_a2a: MoE All-to-All communication (requires mpirun, Blackwell SM10.0+ with MNNVL) - allreduce: AllReduce fusion communication (requires mpirun, Blackwell SM10.0+ with MNNVL) -- triton: Triton reference kernels (used for Mamba selective_state_update) +- triton: Triton reference kernels (used for Mamba selective_state_update and GDN decode/MTP) +- fla: flash-linear-attention Triton kernels (GDN prefill baseline) diff --git a/benchmarks/bench_gdn_decode.py b/benchmarks/bench_gdn_decode.py index fbeeb7631c3..abe1c5db6b2 100644 --- a/benchmarks/bench_gdn_decode.py +++ b/benchmarks/bench_gdn_decode.py @@ -124,6 +124,7 @@ def gdn_decode_bytes( seq_len: int = 1, disable_state_update: bool = False, state_dtype_bytes: int = 4, # 4 for FP32, 2 for BF16 + cache_intermediate_states: bool = False, ) -> int: """ Calculate memory bytes for GDN. @@ -175,10 +176,10 @@ def gdn_decode_bytes( # b: [B, T, HV] - dtype b_bytes = batch_size * seq_len * num_sab_heads * elem_size - # Intermediate states: [B, T, HV, K, V] - only for MTP (seq_len > 1) - # Write all T steps of intermediate states + # Intermediate states: [B, T, HV, K, V] - only written when MTP + # intermediate-state caching is enabled intermediate_bytes = 0 - if seq_len > 1: + if cache_intermediate_states and seq_len > 1: intermediate_bytes = ( batch_size * seq_len @@ -206,837 +207,16 @@ def gdn_decode_bytes( # ============================================================================ # Triton Kernels for comparison benchmarks # ============================================================================ - -try: - import triton - import triton.language as tl - - TRITON_AVAILABLE = True -except ImportError: - TRITON_AVAILABLE = False - -if TRITON_AVAILABLE: - - @triton.jit - def fused_sigmoid_gating_delta_rule_kernel( - # Pointers to matrices - Q, - K, - V, - O, - H, # Hidden state [B, HV, K, V] - A_LOG, # Log decay [HV] - A, # Input-dependent decay [B, HV] - DT_BIAS, # Decay bias [HV] - B_GATE, # Update gate [B, HV] - # Strides - stride_qb, - stride_qh, - stride_qk, - stride_kb, - stride_kh, - stride_kk, - stride_vb, - stride_vh, - stride_vv, - stride_ob, - stride_oh, - stride_ov, - stride_hb, - stride_hh, - stride_hk, - stride_hv, - # Parameters - softplus_beta: tl.constexpr, - softplus_threshold: tl.constexpr, - scale: tl.constexpr, - use_qk_l2norm: tl.constexpr, - B: tl.constexpr, - HV: tl.constexpr, - H_Q: tl.constexpr, - H_K: tl.constexpr, - K_DIM: tl.constexpr, - V_DIM: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - ): - """ - Triton kernel for fused sigmoid gating delta rule update. - - Follows SGLang's implementation: - 1. g = -exp(A_log) * softplus(a + dt_bias) - 2. beta = sigmoid(b) - 3. h *= exp(g) - 4. v_new = v - k @ h - 5. v_new *= beta - 6. h += outer(k, v_new) - 7. o = q @ h - """ - # Block indices - i_bh = tl.program_id(0) - i_k = tl.program_id(1) - i_v = tl.program_id(2) - - i_b = i_bh // HV - i_hv = i_bh % HV - - # GVA head mapping (num_v_heads > num_q_heads) - h_ratio_q = HV // H_Q - h_ratio_k = HV // H_K - i_hq = i_hv // h_ratio_q - i_hk = i_hv // h_ratio_k - - # Load A_log and dt_bias for this head - b_A_log = tl.load(A_LOG + i_hv).to(tl.float32) - b_dt_bias = tl.load(DT_BIAS + i_hv).to(tl.float32) - - # Load a (input-dependent decay) for this batch and head - b_a = tl.load(A + i_b * HV + i_hv).to(tl.float32) - - # Load b (update gate) for this batch and head - b_b = tl.load(B_GATE + i_b * HV + i_hv).to(tl.float32) - - # Compute softplus: softplus(x) = (1/beta) * log(1 + exp(beta*x)) - x = b_a + b_dt_bias - beta_x = softplus_beta * x - softplus_x = tl.where( - beta_x <= softplus_threshold, - (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), - x, - ) - - # Compute g = -exp(A_log) * softplus(a + dt_bias) - b_g = -tl.exp(b_A_log) * softplus_x - - # Compute beta = sigmoid(b) - b_beta = 1.0 / (1.0 + tl.exp(-b_b)) - - # Block offsets - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - - # Load q, k, v - p_q = Q + i_b * stride_qb + i_hq * stride_qh + o_k * stride_qk - p_k = K + i_b * stride_kb + i_hk * stride_kh + o_k * stride_kk - p_v = V + i_b * stride_vb + i_hv * stride_vh + o_v * stride_vv - - b_q = tl.load(p_q, mask=o_k < K_DIM, other=0.0).to(tl.float32) - b_k = tl.load(p_k, mask=o_k < K_DIM, other=0.0).to(tl.float32) - b_v = tl.load(p_v, mask=o_v < V_DIM, other=0.0).to(tl.float32) - - # Apply L2 normalization (if enabled) - if use_qk_l2norm: - # Compute L2 norm across K dimension (need reduction across blocks) - # For simplicity, assume single K block - q_norm = tl.sqrt(tl.sum(b_q * b_q) + 1e-8) - k_norm = tl.sqrt(tl.sum(b_k * b_k) + 1e-8) - b_q = b_q / q_norm - b_k = b_k / k_norm - - # Apply scale to q - b_q = b_q * scale - - # Load hidden state h[K, V] from state[B, HV, K, V] - p_h = ( - H - + i_b * stride_hb - + i_hv * stride_hh - + o_k[:, None] * stride_hk - + o_v[None, :] * stride_hv - ) - b_h = tl.load( - p_h, mask=(o_k[:, None] < K_DIM) & (o_v[None, :] < V_DIM), other=0.0 - ).to(tl.float32) - - # Step 1: Apply decay to hidden state: h *= exp(g) - b_h = b_h * tl.exp(b_g) - - # Step 2: Delta rule: v -= sum(h * k, dim=0) = k @ h - # b_h is [BK, BV], b_k is [BK] - # We need to compute k @ h = sum(k[:, None] * h, dim=0) - b_v = b_v - tl.sum(b_h * b_k[:, None], 0) - - # Step 3: Apply beta gating: v *= beta - b_v = b_v * b_beta - - # Step 4: Update hidden state: h += outer(k, v) = k[:, None] * v[None, :] - b_h = b_h + b_k[:, None] * b_v[None, :] - - # Step 5: Compute output: o = q @ h = sum(q[:, None] * h, dim=0) - b_o = tl.sum(b_h * b_q[:, None], 0) - - # Store output - p_o = O + i_b * stride_ob + i_hv * stride_oh + o_v * stride_ov - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=o_v < V_DIM) - - # Store updated hidden state - tl.store( - p_h, - b_h.to(p_h.dtype.element_ty), - mask=(o_k[:, None] < K_DIM) & (o_v[None, :] < V_DIM), - ) - - @triton.jit - def fused_sigmoid_gating_delta_rule_mtp_kernel( - # Pointers to matrices - Q, # [B, T, H_Q, K] - K, # [B, T, H_K, K] - V, # [B, T, HV, V] - O, # [B, T, HV, V] - H, # Hidden state [pool_size, HV, V, K] (K-last layout) - INTERMEDIATE, # Intermediate states [pool_size, T, HV, V, K] - H0_INDICES, # [B] - A_LOG, # Log decay [HV] - A, # Input-dependent decay [B, T, HV] - DT_BIAS, # Decay bias [HV] - B_GATE, # Update gate [B, T, HV] - # Strides for Q, K, V, O [B, T, H, dim] - stride_qb, - stride_qt, - stride_qh, - stride_qk, - stride_kb, - stride_kt, - stride_kh, - stride_kk, - stride_vb, - stride_vt, - stride_vh, - stride_vv, - stride_ob, - stride_ot, - stride_oh, - stride_ov, - # Strides for hidden state [pool_size, HV, V, K] - stride_hp, - stride_hh, - stride_hv, - stride_hk, - # Strides for intermediate states [pool_size, T, HV, V, K] - stride_ip, - stride_it, - stride_ih, - stride_iv, - stride_ik, - # Strides for A [B, T, HV] - stride_ab, - stride_at, - stride_ah, - # Parameters - softplus_beta: tl.constexpr, - softplus_threshold: tl.constexpr, - scale: tl.constexpr, - use_qk_l2norm: tl.constexpr, - disable_state_update: tl.constexpr, - cache_intermediate_states: tl.constexpr, - B: tl.constexpr, - T: tl.constexpr, - HV: tl.constexpr, - H_Q: tl.constexpr, - H_K: tl.constexpr, - K_DIM: tl.constexpr, - V_DIM: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - ): - """ - Triton kernel for MTP (Multiple Token Processing) delta rule update. - Processes T tokens sequentially, updating state after each token. - - Note: The delta rule operations are fundamentally GEMV (matrix-vector) and - rank-1 updates, which don't directly benefit from tensor cores. Tensor cores - are optimized for GEMM (matrix-matrix). To use tensor cores, we would need - to batch across multiple tokens/heads to form proper GEMM operations. - """ - # Block indices - i_bh = tl.program_id(0) - i_k = tl.program_id(1) - i_v = tl.program_id(2) - - i_b = i_bh // HV - i_hv = i_bh % HV - - # GVA head mapping - h_ratio_q = HV // H_Q - h_ratio_k = HV // H_K - i_hq = i_hv // h_ratio_q - i_hk = i_hv // h_ratio_k - - # Load initial state index for this batch - i_pool = tl.load(H0_INDICES + i_b) - - # Load A_log and dt_bias for this head - b_A_log = tl.load(A_LOG + i_hv).to(tl.float32) - b_dt_bias = tl.load(DT_BIAS + i_hv).to(tl.float32) - - # Block offsets - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - - # Load initial hidden state h[V, K] from state[pool, HV, V, K] - p_h = ( - H - + i_pool * stride_hp - + i_hv * stride_hh - + o_v[:, None] * stride_hv - + o_k[None, :] * stride_hk - ) - b_h = tl.load( - p_h, mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), other=0.0 - ).to(tl.float32) # [BV, BK] - - # Process each token - for t in range(T): - # Load a for this batch, time, head - b_a = tl.load(A + i_b * stride_ab + t * stride_at + i_hv * stride_ah).to( - tl.float32 - ) - b_b = tl.load( - B_GATE + i_b * stride_ab + t * stride_at + i_hv * stride_ah - ).to(tl.float32) - - # Compute softplus and decay - x = b_a + b_dt_bias - beta_x = softplus_beta * x - softplus_x = tl.where( - beta_x <= softplus_threshold, - (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), - x, - ) - b_g = -tl.exp(b_A_log) * softplus_x - b_beta = 1.0 / (1.0 + tl.exp(-b_b)) - - # Load q, k, v for this timestep - p_q = ( - Q + i_b * stride_qb + t * stride_qt + i_hq * stride_qh + o_k * stride_qk - ) - p_k = ( - K + i_b * stride_kb + t * stride_kt + i_hk * stride_kh + o_k * stride_kk - ) - p_v = ( - V + i_b * stride_vb + t * stride_vt + i_hv * stride_vh + o_v * stride_vv - ) - - b_q = tl.load(p_q, mask=o_k < K_DIM, other=0.0).to(tl.float32) - b_k = tl.load(p_k, mask=o_k < K_DIM, other=0.0).to(tl.float32) - b_v = tl.load(p_v, mask=o_v < V_DIM, other=0.0).to(tl.float32) - - # Apply L2 normalization - if use_qk_l2norm: - q_norm = tl.sqrt(tl.sum(b_q * b_q) + 1e-8) - k_norm = tl.sqrt(tl.sum(b_k * b_k) + 1e-8) - b_q = b_q / q_norm - b_k = b_k / k_norm - - b_q = b_q * scale - - # Step 1: Apply decay: h *= exp(g) - b_h = b_h * tl.exp(b_g) - - # Step 2: Delta rule: v -= h @ k (h is [BV, BK], k is [BK]) - # This is GEMV, which doesn't directly use tensor cores efficiently - # h @ k = sum(h * k[None, :], axis=1) -> [BV] - b_v = b_v - tl.sum(b_h * b_k[None, :], 1) - - # Step 3: Apply beta gating - b_v = b_v * b_beta - - # Step 4: Update state: h += outer(v, k) = v[:, None] * k[None, :] - # This is a rank-1 update - b_h = b_h + b_v[:, None] * b_k[None, :] - - # Step 5: Compute output: o = h @ q = sum(h * q[None, :], axis=1) -> [BV] - # This is also GEMV - b_o = tl.sum(b_h * b_q[None, :], 1) - - # Store output for this timestep - p_o = ( - O + i_b * stride_ob + t * stride_ot + i_hv * stride_oh + o_v * stride_ov - ) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=o_v < V_DIM) - - # Cache intermediate state if needed - if cache_intermediate_states: - p_inter = ( - INTERMEDIATE - + i_pool * stride_ip - + t * stride_it - + i_hv * stride_ih - + o_v[:, None] * stride_iv - + o_k[None, :] * stride_ik - ) - tl.store( - p_inter, - b_h.to(p_inter.dtype.element_ty), - mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), - ) - - # Store final state if state update is enabled - if not disable_state_update: - tl.store( - p_h, - b_h.to(p_h.dtype.element_ty), - mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), - ) - - def triton_gdn_decode( - q: torch.Tensor, # [B, 1, H_Q, K] - k: torch.Tensor, # [B, 1, H_K, K] - v: torch.Tensor, # [B, 1, HV, V] - state: torch.Tensor, # [B, HV, K, V] - A_log: torch.Tensor, # [HV] - a: torch.Tensor, # [B, 1, HV] - dt_bias: torch.Tensor, # [HV] - b: torch.Tensor, # [B, 1, HV] - scale: float, - output: torch.Tensor, # [B, 1, HV, V] - use_qk_l2norm: bool = True, - softplus_beta: float = 1.0, - softplus_threshold: float = 20.0, - ): - """ - Triton-based GDN decode matching SGLang's implementation. - """ - B, T, H_Q, K_DIM = q.shape - _, _, H_K, _ = k.shape - _, _, HV, V_DIM = v.shape - - assert T == 1, "Triton kernel only supports decode (T=1)" - - # Reshape inputs for kernel - q_flat = q.squeeze(1) # [B, H_Q, K] - k_flat = k.squeeze(1) # [B, H_K, K] - v_flat = v.squeeze(1) # [B, HV, V] - a_flat = a.squeeze(1) # [B, HV] - b_flat = b.squeeze(1) # [B, HV] - o_flat = output.squeeze(1) # [B, HV, V] - - # Block sizes - BK = triton.next_power_of_2(K_DIM) - BV = triton.next_power_of_2(V_DIM) - - # Limit block sizes (BV smaller to allow more V blocks) - BV = min(BV, 32) - - # Number of blocks - NK = triton.cdiv(K_DIM, BK) - NV = triton.cdiv(V_DIM, BV) - - assert NK == 1, f"Multi-block K not supported: NK={NK}" - - # Launch kernel - grid = (B * HV, NK, NV) - - fused_sigmoid_gating_delta_rule_kernel[grid]( - q_flat, - k_flat, - v_flat, - o_flat, - state, - A_log, - a_flat, - dt_bias, - b_flat, - # Strides for q [B, H_Q, K] - q_flat.stride(0), - q_flat.stride(1), - q_flat.stride(2), - # Strides for k [B, H_K, K] - k_flat.stride(0), - k_flat.stride(1), - k_flat.stride(2), - # Strides for v [B, HV, V] - v_flat.stride(0), - v_flat.stride(1), - v_flat.stride(2), - # Strides for o [B, HV, V] - o_flat.stride(0), - o_flat.stride(1), - o_flat.stride(2), - # Strides for h [B, HV, K, V] - state.stride(0), - state.stride(1), - state.stride(2), - state.stride(3), - # Parameters - softplus_beta=softplus_beta, - softplus_threshold=softplus_threshold, - scale=scale, - use_qk_l2norm=use_qk_l2norm, - B=B, - HV=HV, - H_Q=H_Q, - H_K=H_K, - K_DIM=K_DIM, - V_DIM=V_DIM, - BK=BK, - BV=BV, - ) - - return output, state - - @triton.jit - def fused_sigmoid_gating_delta_rule_kernel_pretranspose( - # Pointers to matrices - Q, - K, - V, - O, - H, # Hidden state [B, HV, V, K] - V-major (pretranspose) layout - A_LOG, # Log decay [HV] - A, # Input-dependent decay [B, HV] - DT_BIAS, # Decay bias [HV] - B_GATE, # Update gate [B, HV] - # Strides - stride_qb, - stride_qh, - stride_qk, - stride_kb, - stride_kh, - stride_kk, - stride_vb, - stride_vh, - stride_vv, - stride_ob, - stride_oh, - stride_ov, - stride_hb, - stride_hh, - stride_hv, # V dimension stride - stride_hk, # K dimension stride - # Parameters - softplus_beta: tl.constexpr, - softplus_threshold: tl.constexpr, - scale: tl.constexpr, - use_qk_l2norm: tl.constexpr, - B: tl.constexpr, - HV: tl.constexpr, - H_Q: tl.constexpr, - H_K: tl.constexpr, - K_DIM: tl.constexpr, - V_DIM: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - ): - """ - Triton kernel for pretranspose layout [B, HV, V, K]. - - Key difference from nontranspose: - - State layout: [B, HV, V, K] instead of [B, HV, K, V] - - h is [BV, BK] instead of [BK, BV] - - h @ k = sum(h * k[None, :], axis=1) -> [BV] - - h += outer(v, k) = v[:, None] * k[None, :] - - o = h @ q = sum(h * q[None, :], axis=1) -> [BV] - """ - # Block indices - i_bh = tl.program_id(0) - i_k = tl.program_id(1) - i_v = tl.program_id(2) - - i_b = i_bh // HV - i_hv = i_bh % HV - - # GVA head mapping (num_v_heads > num_q_heads) - h_ratio_q = HV // H_Q - h_ratio_k = HV // H_K - i_hq = i_hv // h_ratio_q - i_hk = i_hv // h_ratio_k - - # Load A_log and dt_bias for this head - b_A_log = tl.load(A_LOG + i_hv).to(tl.float32) - b_dt_bias = tl.load(DT_BIAS + i_hv).to(tl.float32) - - # Load a (input-dependent decay) for this batch and head - b_a = tl.load(A + i_b * HV + i_hv).to(tl.float32) - - # Load b (update gate) for this batch and head - b_b = tl.load(B_GATE + i_b * HV + i_hv).to(tl.float32) - - # Compute softplus: softplus(x) = (1/beta) * log(1 + exp(beta*x)) - x = b_a + b_dt_bias - beta_x = softplus_beta * x - softplus_x = tl.where( - beta_x <= softplus_threshold, - (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), - x, - ) - - # Compute g = -exp(A_log) * softplus(a + dt_bias) - b_g = -tl.exp(b_A_log) * softplus_x - - # Compute beta = sigmoid(b) - b_beta = 1.0 / (1.0 + tl.exp(-b_b)) - - # Block offsets - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - - # Load q, k, v - p_q = Q + i_b * stride_qb + i_hq * stride_qh + o_k * stride_qk - p_k = K + i_b * stride_kb + i_hk * stride_kh + o_k * stride_kk - p_v = V + i_b * stride_vb + i_hv * stride_vh + o_v * stride_vv - - b_q = tl.load(p_q, mask=o_k < K_DIM, other=0.0).to(tl.float32) - b_k = tl.load(p_k, mask=o_k < K_DIM, other=0.0).to(tl.float32) - b_v = tl.load(p_v, mask=o_v < V_DIM, other=0.0).to(tl.float32) - - # Apply L2 normalization (if enabled) - if use_qk_l2norm: - q_norm = tl.sqrt(tl.sum(b_q * b_q) + 1e-8) - k_norm = tl.sqrt(tl.sum(b_k * b_k) + 1e-8) - b_q = b_q / q_norm - b_k = b_k / k_norm - - # Apply scale to q - b_q = b_q * scale - - # Load hidden state h[V, K] from state[B, HV, V, K] - pretranspose layout - p_h = ( - H - + i_b * stride_hb - + i_hv * stride_hh - + o_v[:, None] * stride_hv - + o_k[None, :] * stride_hk - ) - b_h = tl.load( - p_h, mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), other=0.0 - ).to(tl.float32) # [BV, BK] - - # Step 1: Apply decay to hidden state: h *= exp(g) - b_h = b_h * tl.exp(b_g) - - # Step 2: Delta rule: v -= h @ k = sum(h * k[None, :], axis=1) - # b_h is [BV, BK], b_k is [BK] - b_v = b_v - tl.sum(b_h * b_k[None, :], 1) - - # Step 3: Apply beta gating: v *= beta - b_v = b_v * b_beta - - # Step 4: Update hidden state: h += outer(v, k) = v[:, None] * k[None, :] - b_h = b_h + b_v[:, None] * b_k[None, :] - - # Step 5: Compute output: o = h @ q = sum(h * q[None, :], axis=1) - b_o = tl.sum(b_h * b_q[None, :], 1) - - # Store output - p_o = O + i_b * stride_ob + i_hv * stride_oh + o_v * stride_ov - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=o_v < V_DIM) - - # Store updated hidden state - tl.store( - p_h, - b_h.to(p_h.dtype.element_ty), - mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), - ) - - def triton_gdn_decode_pretranspose( - q: torch.Tensor, # [B, 1, H_Q, K] - k: torch.Tensor, # [B, 1, H_K, K] - v: torch.Tensor, # [B, 1, HV, V] - state: torch.Tensor, # [B, HV, V, K] - pretranspose layout - A_log: torch.Tensor, # [HV] - a: torch.Tensor, # [B, 1, HV] - dt_bias: torch.Tensor, # [HV] - b: torch.Tensor, # [B, 1, HV] - scale: float, - output: torch.Tensor, # [B, 1, HV, V] - use_qk_l2norm: bool = True, - softplus_beta: float = 1.0, - softplus_threshold: float = 20.0, - ): - """ - Triton-based GDN decode for pretranspose layout [B, HV, V, K]. - """ - B, T, H_Q, K_DIM = q.shape - _, _, H_K, _ = k.shape - _, _, HV, V_DIM = v.shape - - assert T == 1, "Triton kernel only supports decode (T=1)" - - # Reshape inputs for kernel - q_flat = q.squeeze(1) # [B, H_Q, K] - k_flat = k.squeeze(1) # [B, H_K, K] - v_flat = v.squeeze(1) # [B, HV, V] - a_flat = a.squeeze(1) # [B, HV] - b_flat = b.squeeze(1) # [B, HV] - o_flat = output.squeeze(1) # [B, HV, V] - - # Block sizes - BK = triton.next_power_of_2(K_DIM) - BV = triton.next_power_of_2(V_DIM) - - # Limit block sizes (BV smaller to allow more V blocks) - BV = min(BV, 32) - - # Number of blocks - NK = triton.cdiv(K_DIM, BK) - NV = triton.cdiv(V_DIM, BV) - - assert NK == 1, f"Multi-block K not supported: NK={NK}" - - # Launch kernel - grid = (B * HV, NK, NV) - - fused_sigmoid_gating_delta_rule_kernel_pretranspose[grid]( - q_flat, - k_flat, - v_flat, - o_flat, - state, - A_log, - a_flat, - dt_bias, - b_flat, - # Strides for q [B, H_Q, K] - q_flat.stride(0), - q_flat.stride(1), - q_flat.stride(2), - # Strides for k [B, H_K, K] - k_flat.stride(0), - k_flat.stride(1), - k_flat.stride(2), - # Strides for v [B, HV, V] - v_flat.stride(0), - v_flat.stride(1), - v_flat.stride(2), - # Strides for o [B, HV, V] - o_flat.stride(0), - o_flat.stride(1), - o_flat.stride(2), - # Strides for h [B, HV, V, K] - pretranspose layout - state.stride(0), - state.stride(1), - state.stride(2), - state.stride(3), - # Parameters - softplus_beta=softplus_beta, - softplus_threshold=softplus_threshold, - scale=scale, - use_qk_l2norm=use_qk_l2norm, - B=B, - HV=HV, - H_Q=H_Q, - H_K=H_K, - K_DIM=K_DIM, - V_DIM=V_DIM, - BK=BK, - BV=BV, - ) - - return output, state - - def triton_gdn_mtp( - q: torch.Tensor, # [B, T, H_Q, K] - k: torch.Tensor, # [B, T, H_K, K] - v: torch.Tensor, # [B, T, HV, V] - initial_state: torch.Tensor, # [pool_size, HV, V, K] - initial_state_indices: torch.Tensor, # [B] - A_log: torch.Tensor, # [HV] - a: torch.Tensor, # [B, T, HV] - dt_bias: torch.Tensor, # [HV] - b: torch.Tensor, # [B, T, HV] - scale: float, - output: torch.Tensor, # [B, T, HV, V] - intermediate_states_buffer: torch.Tensor = None, # [pool_size, T, HV, V, K] - disable_state_update: bool = True, - use_qk_l2norm: bool = True, - softplus_beta: float = 1.0, - softplus_threshold: float = 20.0, - ): - """ - Triton-based GDN MTP matching SGLang's implementation. - """ - B, T, H_Q, K_DIM = q.shape - _, _, H_K, _ = k.shape - _, _, HV, V_DIM = v.shape - - # Block sizes (BV smaller to allow more V blocks) - BK = triton.next_power_of_2(K_DIM) - BV = triton.next_power_of_2(V_DIM) - BV = min(BV, 32) - - NK = triton.cdiv(K_DIM, BK) - NV = triton.cdiv(V_DIM, BV) - - assert NK == 1, f"Multi-block K not supported: NK={NK}" - - cache_intermediate_states = intermediate_states_buffer is not None - if cache_intermediate_states: - intermediate = intermediate_states_buffer - else: - intermediate = torch.zeros( - 1, 1, 1, 1, 1, dtype=torch.float32, device=q.device - ) - - # Launch kernel - grid = (B * HV, NK, NV) - - fused_sigmoid_gating_delta_rule_mtp_kernel[grid]( - q, - k, - v, - output, - initial_state, - intermediate, - initial_state_indices, - A_log, - a, - dt_bias, - b, - # Q strides - q.stride(0), - q.stride(1), - q.stride(2), - q.stride(3), - # K strides - k.stride(0), - k.stride(1), - k.stride(2), - k.stride(3), - # V strides - v.stride(0), - v.stride(1), - v.stride(2), - v.stride(3), - # O strides - output.stride(0), - output.stride(1), - output.stride(2), - output.stride(3), - # H strides [pool_size, HV, V, K] - initial_state.stride(0), - initial_state.stride(1), - initial_state.stride(2), - initial_state.stride(3), - # Intermediate strides [pool_size, T, HV, V, K] - intermediate.stride(0), - intermediate.stride(1), - intermediate.stride(2) if cache_intermediate_states else 0, - intermediate.stride(3) if cache_intermediate_states else 0, - intermediate.stride(4) if cache_intermediate_states else 0, - # A strides [B, T, HV] - a.stride(0), - a.stride(1), - a.stride(2), - # Parameters - softplus_beta=softplus_beta, - softplus_threshold=softplus_threshold, - scale=scale, - use_qk_l2norm=use_qk_l2norm, - disable_state_update=disable_state_update, - cache_intermediate_states=cache_intermediate_states, - B=B, - T=T, - HV=HV, - H_Q=H_Q, - H_K=H_K, - K_DIM=K_DIM, - V_DIM=V_DIM, - BK=BK, - BV=BV, - ) - - return output, initial_state +# The Triton GDN kernels live in benchmarks/gdn_triton_reference.py (same +# directory as this script) so they can be shared with the +# flashinfer_benchmark.py GDN routines (benchmarks/routines/gdn.py). + +from gdn_triton_reference import ( + TRITON_AVAILABLE, + triton_gdn_decode, + triton_gdn_decode_pretranspose, + triton_gdn_mtp, +) # ============================================================================ @@ -1265,6 +445,7 @@ def bench_gdn_mtp( dtype, seq_len, disable_state_update=disable_state_update, + cache_intermediate_states=cache_intermediate_states, ) kernel_tflops = flops / kernel_median_ms / 1e9 if kernel_median_ms > 0 else 0 @@ -2073,11 +1254,27 @@ def bench_all_layouts( output = torch.empty( batch_size, T, num_o_heads, head_size, dtype=dtype, device="cuda" ) + # The BF16 state kernels are pool-only: treat the [B, HV, V, K] state + # as a pool of size B with sequential indices (read == write). + initial_state_indices = torch.arange( + batch_size, dtype=torch.int32, device="cuda" + ) try: times = bench_gpu_time( lambda: gdn_decode_bf16_state_wrapper( - q, k, v, state, A_log, a, dt_bias, b, scale, output, use_qk_l2norm + q, + k, + v, + state, + A_log, + a, + dt_bias, + b, + scale, + output, + use_qk_l2norm, + initial_state_indices=initial_state_indices, ), enable_cupti=True, dry_run_iters=warmup_iters, @@ -2372,6 +1569,7 @@ def bench_gdn_decode_bf16_state( seq_len, disable_state_update=disable_state_update, state_dtype_bytes=2, # BF16 state for gdn_decode_bf16_state + cache_intermediate_states=cache_intermediate_states, ) kernel_tflops = flops / kernel_median_ms / 1e9 if kernel_median_ms > 0 else 0 diff --git a/benchmarks/bench_gdn_prefill.py b/benchmarks/bench_gdn_prefill.py index 65955416cf6..09c2a520b50 100644 --- a/benchmarks/bench_gdn_prefill.py +++ b/benchmarks/bench_gdn_prefill.py @@ -94,7 +94,10 @@ def bench_fi(endpoints, h_qk, h_v, d, warmup, iters): torch.randn(T, h_qk, d, dtype=torch.float32, device=device), p=2, dim=-1 ).to(dtype) v = torch.randn((T, h_v, d), dtype=dtype, device=device) - g = F.logsigmoid(torch.rand(T, h_v, dtype=torch.float32, device=device)) + # FlashInfer's g is the linear-space forget gate alpha in (0, 1) + # ("defaults to all ones" = no decay). Log-space gates (e.g. logsigmoid) + # are out of domain and produce NaN outputs/state. + g = torch.rand(T, h_v, dtype=torch.float32, device=device) beta = torch.rand(T, h_v, dtype=torch.float32, device=device).sigmoid() h0 = torch.randn((N, h_v, d, d), dtype=torch.float32, device=device) state_out = torch.zeros_like(h0) @@ -162,8 +165,10 @@ def main(): ) if not _has_fla: - print("Error: FLA not installed. Run: pip install flash-linear-attention") - sys.exit(1) + print( + "Warning: FLA not installed (pip install flash-linear-attention). " + "Benchmarking FlashInfer only." + ) print(f"\nGPU: {torch.cuda.get_device_name(0)} [{arch_label}]") print("Models: Qwen3.5 family (397B, 122B, 35B, 27B, 9B, 4B, 2B, 0.8B), d=128") @@ -172,8 +177,9 @@ def main(): header = ( f"{'Heads':<15s} {'Seqlens':<16s} {'h_qk':>4s} {'h_v':>4s}" f" {fi_col:>22s} {'TFLOPS':>7s}" - f" {'FLA/Triton':>10s} {'Speedup':>8s}" ) + if _has_fla: + header += f" {'FLA/Triton':>10s} {'Speedup':>8s}" print(header) print("-" * len(header)) @@ -181,15 +187,17 @@ def main(): for endpoints, s_label in SEQ_CONFIGS: T = endpoints[-1] fi_ms = bench_fi(endpoints, h_qk, h_v, d, args.warmup, args.iters) - fla_ms = bench_fla(endpoints, h_qk, h_v, d, args.warmup, args.iters) tflops = _gdn_tflops(T, h_v, d, fi_ms) - speedup = fla_ms / fi_ms - marker = "+" if speedup > 1.0 else "-" - print( + row = ( f"{h_label:<15s} {s_label:<16s} {h_qk:>4d} {h_v:>4d}" f" {fi_ms:>21.3f}ms {tflops:>6.1f}" - f" {fla_ms:>9.3f}ms {speedup:>7.2f}x {marker}" ) + if _has_fla: + fla_ms = bench_fla(endpoints, h_qk, h_v, d, args.warmup, args.iters) + speedup = fla_ms / fi_ms + marker = "+" if speedup > 1.0 else "-" + row += f" {fla_ms:>9.3f}ms {speedup:>7.2f}x {marker}" + print(row) print() diff --git a/benchmarks/flashinfer_benchmark.py b/benchmarks/flashinfer_benchmark.py index d502e7d562f..45a95f07067 100644 --- a/benchmarks/flashinfer_benchmark.py +++ b/benchmarks/flashinfer_benchmark.py @@ -64,6 +64,10 @@ def run_test(args): from routines.mamba import run_mamba_test res = run_mamba_test(args) + elif args.routine in benchmark_apis["gdn"]: + from routines.gdn import run_gdn_test + + res = run_gdn_test(args) else: raise ValueError(f"Unsupported routine: {args.routine}") @@ -114,7 +118,8 @@ def parse_args(line=sys.argv[1:]): + list(benchmark_apis["quantization"]) + list(benchmark_apis["sampling"]) + list(benchmark_apis["rope"]) - + list(benchmark_apis["mamba"]), + + list(benchmark_apis["mamba"]) + + list(benchmark_apis["gdn"]), ) args, _ = parser.parse_known_args(line[:]) @@ -267,6 +272,10 @@ def parse_args(line=sys.argv[1:]): from routines.mamba import parse_mamba_args args = parse_mamba_args(line, parser) + elif args.routine in benchmark_apis["gdn"]: + from routines.gdn import parse_gdn_args + + args = parse_gdn_args(line, parser) else: raise ValueError(f"Unsupported routine: {args.routine}") diff --git a/benchmarks/gdn_triton_reference.py b/benchmarks/gdn_triton_reference.py new file mode 100644 index 00000000000..b85e4da74d5 --- /dev/null +++ b/benchmarks/gdn_triton_reference.py @@ -0,0 +1,873 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +""" +Triton reference implementations of the Gated Delta Rule (GDN) decode/MTP +kernels, following SGLang's fused_sigmoid_gating_delta_rule implementation. + +Shared by: +- bench_gdn_decode.py (standalone GDN decode benchmark) +- routines/gdn.py (flashinfer_benchmark.py GDN routines) + +Exports: +- TRITON_AVAILABLE: bool, whether triton could be imported +- triton_gdn_decode: nontranspose [B, HV, K, V] state layout, T=1 +- triton_gdn_decode_pretranspose: pretranspose [B, HV, V, K] state layout, T=1 +- triton_gdn_mtp: MTP (T>=1) with state pool + indices, [pool, HV, V, K] layout +""" + +import torch + +# ============================================================================ +# Triton Kernels for comparison benchmarks +# ============================================================================ + +try: + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: + TRITON_AVAILABLE = False + +if TRITON_AVAILABLE: + + @triton.jit + def fused_sigmoid_gating_delta_rule_kernel( + # Pointers to matrices + Q, + K, + V, + O, + H, # Hidden state [B, HV, K, V] + A_LOG, # Log decay [HV] + A, # Input-dependent decay [B, HV] + DT_BIAS, # Decay bias [HV] + B_GATE, # Update gate [B, HV] + # Strides + stride_qb, + stride_qh, + stride_qk, + stride_kb, + stride_kh, + stride_kk, + stride_vb, + stride_vh, + stride_vv, + stride_ob, + stride_oh, + stride_ov, + stride_hb, + stride_hh, + stride_hk, + stride_hv, + # Parameters + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + scale: tl.constexpr, + use_qk_l2norm: tl.constexpr, + B: tl.constexpr, + HV: tl.constexpr, + H_Q: tl.constexpr, + H_K: tl.constexpr, + K_DIM: tl.constexpr, + V_DIM: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + ): + """ + Triton kernel for fused sigmoid gating delta rule update. + + Follows SGLang's implementation: + 1. g = -exp(A_log) * softplus(a + dt_bias) + 2. beta = sigmoid(b) + 3. h *= exp(g) + 4. v_new = v - k @ h + 5. v_new *= beta + 6. h += outer(k, v_new) + 7. o = q @ h + """ + # Block indices + i_bh = tl.program_id(0) + i_k = tl.program_id(1) + i_v = tl.program_id(2) + + i_b = i_bh // HV + i_hv = i_bh % HV + + # GVA head mapping (num_v_heads > num_q_heads) + h_ratio_q = HV // H_Q + h_ratio_k = HV // H_K + i_hq = i_hv // h_ratio_q + i_hk = i_hv // h_ratio_k + + # Load A_log and dt_bias for this head + b_A_log = tl.load(A_LOG + i_hv).to(tl.float32) + b_dt_bias = tl.load(DT_BIAS + i_hv).to(tl.float32) + + # Load a (input-dependent decay) for this batch and head + b_a = tl.load(A + i_b * HV + i_hv).to(tl.float32) + + # Load b (update gate) for this batch and head + b_b = tl.load(B_GATE + i_b * HV + i_hv).to(tl.float32) + + # Compute softplus: softplus(x) = (1/beta) * log(1 + exp(beta*x)) + x = b_a + b_dt_bias + beta_x = softplus_beta * x + softplus_x = tl.where( + beta_x <= softplus_threshold, + (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), + x, + ) + + # Compute g = -exp(A_log) * softplus(a + dt_bias) + b_g = -tl.exp(b_A_log) * softplus_x + + # Compute beta = sigmoid(b) + b_beta = 1.0 / (1.0 + tl.exp(-b_b)) + + # Block offsets + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + # Load q, k, v + p_q = Q + i_b * stride_qb + i_hq * stride_qh + o_k * stride_qk + p_k = K + i_b * stride_kb + i_hk * stride_kh + o_k * stride_kk + p_v = V + i_b * stride_vb + i_hv * stride_vh + o_v * stride_vv + + b_q = tl.load(p_q, mask=o_k < K_DIM, other=0.0).to(tl.float32) + b_k = tl.load(p_k, mask=o_k < K_DIM, other=0.0).to(tl.float32) + b_v = tl.load(p_v, mask=o_v < V_DIM, other=0.0).to(tl.float32) + + # Apply L2 normalization (if enabled) + if use_qk_l2norm: + # Compute L2 norm across K dimension (need reduction across blocks) + # For simplicity, assume single K block + q_norm = tl.sqrt(tl.sum(b_q * b_q) + 1e-8) + k_norm = tl.sqrt(tl.sum(b_k * b_k) + 1e-8) + b_q = b_q / q_norm + b_k = b_k / k_norm + + # Apply scale to q + b_q = b_q * scale + + # Load hidden state h[K, V] from state[B, HV, K, V] + p_h = ( + H + + i_b * stride_hb + + i_hv * stride_hh + + o_k[:, None] * stride_hk + + o_v[None, :] * stride_hv + ) + b_h = tl.load( + p_h, mask=(o_k[:, None] < K_DIM) & (o_v[None, :] < V_DIM), other=0.0 + ).to(tl.float32) + + # Step 1: Apply decay to hidden state: h *= exp(g) + b_h = b_h * tl.exp(b_g) + + # Step 2: Delta rule: v -= sum(h * k, dim=0) = k @ h + # b_h is [BK, BV], b_k is [BK] + # We need to compute k @ h = sum(k[:, None] * h, dim=0) + b_v = b_v - tl.sum(b_h * b_k[:, None], 0) + + # Step 3: Apply beta gating: v *= beta + b_v = b_v * b_beta + + # Step 4: Update hidden state: h += outer(k, v) = k[:, None] * v[None, :] + b_h = b_h + b_k[:, None] * b_v[None, :] + + # Step 5: Compute output: o = q @ h = sum(q[:, None] * h, dim=0) + b_o = tl.sum(b_h * b_q[:, None], 0) + + # Store output + p_o = O + i_b * stride_ob + i_hv * stride_oh + o_v * stride_ov + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=o_v < V_DIM) + + # Store updated hidden state + tl.store( + p_h, + b_h.to(p_h.dtype.element_ty), + mask=(o_k[:, None] < K_DIM) & (o_v[None, :] < V_DIM), + ) + + @triton.jit + def fused_sigmoid_gating_delta_rule_mtp_kernel( + # Pointers to matrices + Q, # [B, T, H_Q, K] + K, # [B, T, H_K, K] + V, # [B, T, HV, V] + O, # [B, T, HV, V] + H, # Hidden state [pool_size, HV, V, K] (K-last layout) + INTERMEDIATE, # Intermediate states [pool_size, T, HV, V, K] + H0_INDICES, # [B] + A_LOG, # Log decay [HV] + A, # Input-dependent decay [B, T, HV] + DT_BIAS, # Decay bias [HV] + B_GATE, # Update gate [B, T, HV] + # Strides for Q, K, V, O [B, T, H, dim] + stride_qb, + stride_qt, + stride_qh, + stride_qk, + stride_kb, + stride_kt, + stride_kh, + stride_kk, + stride_vb, + stride_vt, + stride_vh, + stride_vv, + stride_ob, + stride_ot, + stride_oh, + stride_ov, + # Strides for hidden state [pool_size, HV, V, K] + stride_hp, + stride_hh, + stride_hv, + stride_hk, + # Strides for intermediate states [pool_size, T, HV, V, K] + stride_ip, + stride_it, + stride_ih, + stride_iv, + stride_ik, + # Strides for A [B, T, HV] + stride_ab, + stride_at, + stride_ah, + # Parameters + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + scale: tl.constexpr, + use_qk_l2norm: tl.constexpr, + disable_state_update: tl.constexpr, + cache_intermediate_states: tl.constexpr, + B: tl.constexpr, + T: tl.constexpr, + HV: tl.constexpr, + H_Q: tl.constexpr, + H_K: tl.constexpr, + K_DIM: tl.constexpr, + V_DIM: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + ): + """ + Triton kernel for MTP (Multiple Token Processing) delta rule update. + Processes T tokens sequentially, updating state after each token. + + Note: The delta rule operations are fundamentally GEMV (matrix-vector) and + rank-1 updates, which don't directly benefit from tensor cores. Tensor cores + are optimized for GEMM (matrix-matrix). To use tensor cores, we would need + to batch across multiple tokens/heads to form proper GEMM operations. + """ + # Block indices + i_bh = tl.program_id(0) + i_k = tl.program_id(1) + i_v = tl.program_id(2) + + i_b = i_bh // HV + i_hv = i_bh % HV + + # GVA head mapping + h_ratio_q = HV // H_Q + h_ratio_k = HV // H_K + i_hq = i_hv // h_ratio_q + i_hk = i_hv // h_ratio_k + + # Load initial state index for this batch + i_pool = tl.load(H0_INDICES + i_b) + + # Load A_log and dt_bias for this head + b_A_log = tl.load(A_LOG + i_hv).to(tl.float32) + b_dt_bias = tl.load(DT_BIAS + i_hv).to(tl.float32) + + # Block offsets + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + # Load initial hidden state h[V, K] from state[pool, HV, V, K] + p_h = ( + H + + i_pool * stride_hp + + i_hv * stride_hh + + o_v[:, None] * stride_hv + + o_k[None, :] * stride_hk + ) + b_h = tl.load( + p_h, mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), other=0.0 + ).to(tl.float32) # [BV, BK] + + # Process each token + for t in range(T): + # Load a for this batch, time, head + b_a = tl.load(A + i_b * stride_ab + t * stride_at + i_hv * stride_ah).to( + tl.float32 + ) + b_b = tl.load( + B_GATE + i_b * stride_ab + t * stride_at + i_hv * stride_ah + ).to(tl.float32) + + # Compute softplus and decay + x = b_a + b_dt_bias + beta_x = softplus_beta * x + softplus_x = tl.where( + beta_x <= softplus_threshold, + (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x + b_beta = 1.0 / (1.0 + tl.exp(-b_b)) + + # Load q, k, v for this timestep + p_q = ( + Q + i_b * stride_qb + t * stride_qt + i_hq * stride_qh + o_k * stride_qk + ) + p_k = ( + K + i_b * stride_kb + t * stride_kt + i_hk * stride_kh + o_k * stride_kk + ) + p_v = ( + V + i_b * stride_vb + t * stride_vt + i_hv * stride_vh + o_v * stride_vv + ) + + b_q = tl.load(p_q, mask=o_k < K_DIM, other=0.0).to(tl.float32) + b_k = tl.load(p_k, mask=o_k < K_DIM, other=0.0).to(tl.float32) + b_v = tl.load(p_v, mask=o_v < V_DIM, other=0.0).to(tl.float32) + + # Apply L2 normalization + if use_qk_l2norm: + q_norm = tl.sqrt(tl.sum(b_q * b_q) + 1e-8) + k_norm = tl.sqrt(tl.sum(b_k * b_k) + 1e-8) + b_q = b_q / q_norm + b_k = b_k / k_norm + + b_q = b_q * scale + + # Step 1: Apply decay: h *= exp(g) + b_h = b_h * tl.exp(b_g) + + # Step 2: Delta rule: v -= h @ k (h is [BV, BK], k is [BK]) + # This is GEMV, which doesn't directly use tensor cores efficiently + # h @ k = sum(h * k[None, :], axis=1) -> [BV] + b_v = b_v - tl.sum(b_h * b_k[None, :], 1) + + # Step 3: Apply beta gating + b_v = b_v * b_beta + + # Step 4: Update state: h += outer(v, k) = v[:, None] * k[None, :] + # This is a rank-1 update + b_h = b_h + b_v[:, None] * b_k[None, :] + + # Step 5: Compute output: o = h @ q = sum(h * q[None, :], axis=1) -> [BV] + # This is also GEMV + b_o = tl.sum(b_h * b_q[None, :], 1) + + # Store output for this timestep + p_o = ( + O + i_b * stride_ob + t * stride_ot + i_hv * stride_oh + o_v * stride_ov + ) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=o_v < V_DIM) + + # Cache intermediate state if needed + if cache_intermediate_states: + p_inter = ( + INTERMEDIATE + + i_pool * stride_ip + + t * stride_it + + i_hv * stride_ih + + o_v[:, None] * stride_iv + + o_k[None, :] * stride_ik + ) + tl.store( + p_inter, + b_h.to(p_inter.dtype.element_ty), + mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), + ) + + # Store final state if state update is enabled + if not disable_state_update: + tl.store( + p_h, + b_h.to(p_h.dtype.element_ty), + mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), + ) + + def triton_gdn_decode( + q: torch.Tensor, # [B, 1, H_Q, K] + k: torch.Tensor, # [B, 1, H_K, K] + v: torch.Tensor, # [B, 1, HV, V] + state: torch.Tensor, # [B, HV, K, V] + A_log: torch.Tensor, # [HV] + a: torch.Tensor, # [B, 1, HV] + dt_bias: torch.Tensor, # [HV] + b: torch.Tensor, # [B, 1, HV] + scale: float, + output: torch.Tensor, # [B, 1, HV, V] + use_qk_l2norm: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + ): + """ + Triton-based GDN decode matching SGLang's implementation. + """ + B, T, H_Q, K_DIM = q.shape + _, _, H_K, _ = k.shape + _, _, HV, V_DIM = v.shape + + assert T == 1, "Triton kernel only supports decode (T=1)" + + # Reshape inputs for kernel + q_flat = q.squeeze(1) # [B, H_Q, K] + k_flat = k.squeeze(1) # [B, H_K, K] + v_flat = v.squeeze(1) # [B, HV, V] + a_flat = a.squeeze(1) # [B, HV] + b_flat = b.squeeze(1) # [B, HV] + o_flat = output.squeeze(1) # [B, HV, V] + + # Block sizes + BK = triton.next_power_of_2(K_DIM) + BV = triton.next_power_of_2(V_DIM) + + # Limit block sizes (BV smaller to allow more V blocks) + BV = min(BV, 32) + + # Number of blocks + NK = triton.cdiv(K_DIM, BK) + NV = triton.cdiv(V_DIM, BV) + + assert NK == 1, f"Multi-block K not supported: NK={NK}" + + # Launch kernel + grid = (B * HV, NK, NV) + + fused_sigmoid_gating_delta_rule_kernel[grid]( + q_flat, + k_flat, + v_flat, + o_flat, + state, + A_log, + a_flat, + dt_bias, + b_flat, + # Strides for q [B, H_Q, K] + q_flat.stride(0), + q_flat.stride(1), + q_flat.stride(2), + # Strides for k [B, H_K, K] + k_flat.stride(0), + k_flat.stride(1), + k_flat.stride(2), + # Strides for v [B, HV, V] + v_flat.stride(0), + v_flat.stride(1), + v_flat.stride(2), + # Strides for o [B, HV, V] + o_flat.stride(0), + o_flat.stride(1), + o_flat.stride(2), + # Strides for h [B, HV, K, V] + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + # Parameters + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + use_qk_l2norm=use_qk_l2norm, + B=B, + HV=HV, + H_Q=H_Q, + H_K=H_K, + K_DIM=K_DIM, + V_DIM=V_DIM, + BK=BK, + BV=BV, + ) + + return output, state + + @triton.jit + def fused_sigmoid_gating_delta_rule_kernel_pretranspose( + # Pointers to matrices + Q, + K, + V, + O, + H, # Hidden state [B, HV, V, K] - V-major (pretranspose) layout + A_LOG, # Log decay [HV] + A, # Input-dependent decay [B, HV] + DT_BIAS, # Decay bias [HV] + B_GATE, # Update gate [B, HV] + # Strides + stride_qb, + stride_qh, + stride_qk, + stride_kb, + stride_kh, + stride_kk, + stride_vb, + stride_vh, + stride_vv, + stride_ob, + stride_oh, + stride_ov, + stride_hb, + stride_hh, + stride_hv, # V dimension stride + stride_hk, # K dimension stride + # Parameters + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + scale: tl.constexpr, + use_qk_l2norm: tl.constexpr, + B: tl.constexpr, + HV: tl.constexpr, + H_Q: tl.constexpr, + H_K: tl.constexpr, + K_DIM: tl.constexpr, + V_DIM: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + ): + """ + Triton kernel for pretranspose layout [B, HV, V, K]. + + Key difference from nontranspose: + - State layout: [B, HV, V, K] instead of [B, HV, K, V] + - h is [BV, BK] instead of [BK, BV] + - h @ k = sum(h * k[None, :], axis=1) -> [BV] + - h += outer(v, k) = v[:, None] * k[None, :] + - o = h @ q = sum(h * q[None, :], axis=1) -> [BV] + """ + # Block indices + i_bh = tl.program_id(0) + i_k = tl.program_id(1) + i_v = tl.program_id(2) + + i_b = i_bh // HV + i_hv = i_bh % HV + + # GVA head mapping (num_v_heads > num_q_heads) + h_ratio_q = HV // H_Q + h_ratio_k = HV // H_K + i_hq = i_hv // h_ratio_q + i_hk = i_hv // h_ratio_k + + # Load A_log and dt_bias for this head + b_A_log = tl.load(A_LOG + i_hv).to(tl.float32) + b_dt_bias = tl.load(DT_BIAS + i_hv).to(tl.float32) + + # Load a (input-dependent decay) for this batch and head + b_a = tl.load(A + i_b * HV + i_hv).to(tl.float32) + + # Load b (update gate) for this batch and head + b_b = tl.load(B_GATE + i_b * HV + i_hv).to(tl.float32) + + # Compute softplus: softplus(x) = (1/beta) * log(1 + exp(beta*x)) + x = b_a + b_dt_bias + beta_x = softplus_beta * x + softplus_x = tl.where( + beta_x <= softplus_threshold, + (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), + x, + ) + + # Compute g = -exp(A_log) * softplus(a + dt_bias) + b_g = -tl.exp(b_A_log) * softplus_x + + # Compute beta = sigmoid(b) + b_beta = 1.0 / (1.0 + tl.exp(-b_b)) + + # Block offsets + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + # Load q, k, v + p_q = Q + i_b * stride_qb + i_hq * stride_qh + o_k * stride_qk + p_k = K + i_b * stride_kb + i_hk * stride_kh + o_k * stride_kk + p_v = V + i_b * stride_vb + i_hv * stride_vh + o_v * stride_vv + + b_q = tl.load(p_q, mask=o_k < K_DIM, other=0.0).to(tl.float32) + b_k = tl.load(p_k, mask=o_k < K_DIM, other=0.0).to(tl.float32) + b_v = tl.load(p_v, mask=o_v < V_DIM, other=0.0).to(tl.float32) + + # Apply L2 normalization (if enabled) + if use_qk_l2norm: + q_norm = tl.sqrt(tl.sum(b_q * b_q) + 1e-8) + k_norm = tl.sqrt(tl.sum(b_k * b_k) + 1e-8) + b_q = b_q / q_norm + b_k = b_k / k_norm + + # Apply scale to q + b_q = b_q * scale + + # Load hidden state h[V, K] from state[B, HV, V, K] - pretranspose layout + p_h = ( + H + + i_b * stride_hb + + i_hv * stride_hh + + o_v[:, None] * stride_hv + + o_k[None, :] * stride_hk + ) + b_h = tl.load( + p_h, mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), other=0.0 + ).to(tl.float32) # [BV, BK] + + # Step 1: Apply decay to hidden state: h *= exp(g) + b_h = b_h * tl.exp(b_g) + + # Step 2: Delta rule: v -= h @ k = sum(h * k[None, :], axis=1) + # b_h is [BV, BK], b_k is [BK] + b_v = b_v - tl.sum(b_h * b_k[None, :], 1) + + # Step 3: Apply beta gating: v *= beta + b_v = b_v * b_beta + + # Step 4: Update hidden state: h += outer(v, k) = v[:, None] * k[None, :] + b_h = b_h + b_v[:, None] * b_k[None, :] + + # Step 5: Compute output: o = h @ q = sum(h * q[None, :], axis=1) + b_o = tl.sum(b_h * b_q[None, :], 1) + + # Store output + p_o = O + i_b * stride_ob + i_hv * stride_oh + o_v * stride_ov + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=o_v < V_DIM) + + # Store updated hidden state + tl.store( + p_h, + b_h.to(p_h.dtype.element_ty), + mask=(o_v[:, None] < V_DIM) & (o_k[None, :] < K_DIM), + ) + + def triton_gdn_decode_pretranspose( + q: torch.Tensor, # [B, 1, H_Q, K] + k: torch.Tensor, # [B, 1, H_K, K] + v: torch.Tensor, # [B, 1, HV, V] + state: torch.Tensor, # [B, HV, V, K] - pretranspose layout + A_log: torch.Tensor, # [HV] + a: torch.Tensor, # [B, 1, HV] + dt_bias: torch.Tensor, # [HV] + b: torch.Tensor, # [B, 1, HV] + scale: float, + output: torch.Tensor, # [B, 1, HV, V] + use_qk_l2norm: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + ): + """ + Triton-based GDN decode for pretranspose layout [B, HV, V, K]. + """ + B, T, H_Q, K_DIM = q.shape + _, _, H_K, _ = k.shape + _, _, HV, V_DIM = v.shape + + assert T == 1, "Triton kernel only supports decode (T=1)" + + # Reshape inputs for kernel + q_flat = q.squeeze(1) # [B, H_Q, K] + k_flat = k.squeeze(1) # [B, H_K, K] + v_flat = v.squeeze(1) # [B, HV, V] + a_flat = a.squeeze(1) # [B, HV] + b_flat = b.squeeze(1) # [B, HV] + o_flat = output.squeeze(1) # [B, HV, V] + + # Block sizes + BK = triton.next_power_of_2(K_DIM) + BV = triton.next_power_of_2(V_DIM) + + # Limit block sizes (BV smaller to allow more V blocks) + BV = min(BV, 32) + + # Number of blocks + NK = triton.cdiv(K_DIM, BK) + NV = triton.cdiv(V_DIM, BV) + + assert NK == 1, f"Multi-block K not supported: NK={NK}" + + # Launch kernel + grid = (B * HV, NK, NV) + + fused_sigmoid_gating_delta_rule_kernel_pretranspose[grid]( + q_flat, + k_flat, + v_flat, + o_flat, + state, + A_log, + a_flat, + dt_bias, + b_flat, + # Strides for q [B, H_Q, K] + q_flat.stride(0), + q_flat.stride(1), + q_flat.stride(2), + # Strides for k [B, H_K, K] + k_flat.stride(0), + k_flat.stride(1), + k_flat.stride(2), + # Strides for v [B, HV, V] + v_flat.stride(0), + v_flat.stride(1), + v_flat.stride(2), + # Strides for o [B, HV, V] + o_flat.stride(0), + o_flat.stride(1), + o_flat.stride(2), + # Strides for h [B, HV, V, K] - pretranspose layout + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + # Parameters + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + use_qk_l2norm=use_qk_l2norm, + B=B, + HV=HV, + H_Q=H_Q, + H_K=H_K, + K_DIM=K_DIM, + V_DIM=V_DIM, + BK=BK, + BV=BV, + ) + + return output, state + + def triton_gdn_mtp( + q: torch.Tensor, # [B, T, H_Q, K] + k: torch.Tensor, # [B, T, H_K, K] + v: torch.Tensor, # [B, T, HV, V] + initial_state: torch.Tensor, # [pool_size, HV, V, K] + initial_state_indices: torch.Tensor, # [B] + A_log: torch.Tensor, # [HV] + a: torch.Tensor, # [B, T, HV] + dt_bias: torch.Tensor, # [HV] + b: torch.Tensor, # [B, T, HV] + scale: float, + output: torch.Tensor, # [B, T, HV, V] + intermediate_states_buffer: torch.Tensor = None, # [pool_size, T, HV, V, K] + disable_state_update: bool = True, + use_qk_l2norm: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + ): + """ + Triton-based GDN MTP matching SGLang's implementation. + """ + B, T, H_Q, K_DIM = q.shape + _, _, H_K, _ = k.shape + _, _, HV, V_DIM = v.shape + + # Block sizes (BV smaller to allow more V blocks) + BK = triton.next_power_of_2(K_DIM) + BV = triton.next_power_of_2(V_DIM) + BV = min(BV, 32) + + NK = triton.cdiv(K_DIM, BK) + NV = triton.cdiv(V_DIM, BV) + + assert NK == 1, f"Multi-block K not supported: NK={NK}" + + cache_intermediate_states = intermediate_states_buffer is not None + if cache_intermediate_states: + intermediate = intermediate_states_buffer + else: + intermediate = torch.zeros( + 1, 1, 1, 1, 1, dtype=torch.float32, device=q.device + ) + + # Launch kernel + grid = (B * HV, NK, NV) + + fused_sigmoid_gating_delta_rule_mtp_kernel[grid]( + q, + k, + v, + output, + initial_state, + intermediate, + initial_state_indices, + A_log, + a, + dt_bias, + b, + # Q strides + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + # K strides + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + # V strides + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + # O strides + output.stride(0), + output.stride(1), + output.stride(2), + output.stride(3), + # H strides [pool_size, HV, V, K] + initial_state.stride(0), + initial_state.stride(1), + initial_state.stride(2), + initial_state.stride(3), + # Intermediate strides [pool_size, T, HV, V, K] + intermediate.stride(0), + intermediate.stride(1), + intermediate.stride(2) if cache_intermediate_states else 0, + intermediate.stride(3) if cache_intermediate_states else 0, + intermediate.stride(4) if cache_intermediate_states else 0, + # A strides [B, T, HV] + a.stride(0), + a.stride(1), + a.stride(2), + # Parameters + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + use_qk_l2norm=use_qk_l2norm, + disable_state_update=disable_state_update, + cache_intermediate_states=cache_intermediate_states, + B=B, + T=T, + HV=HV, + H_Q=H_Q, + H_K=H_K, + K_DIM=K_DIM, + V_DIM=V_DIM, + BK=BK, + BV=BV, + ) + + return output, initial_state +else: + # Allow `from gdn_triton_reference import ...` to succeed without Triton; + # all call sites guard on TRITON_AVAILABLE before using these. + triton_gdn_decode = None + triton_gdn_decode_pretranspose = None + triton_gdn_mtp = None diff --git a/benchmarks/routines/flashinfer_benchmark_utils.py b/benchmarks/routines/flashinfer_benchmark_utils.py index 870dbf15d15..8f6a4b40507 100644 --- a/benchmarks/routines/flashinfer_benchmark_utils.py +++ b/benchmarks/routines/flashinfer_benchmark_utils.py @@ -135,6 +135,16 @@ "has_z", "dt_softplus", ], + "gdn": [ + "num_q_heads", + "num_k_heads", + "num_v_heads", + "head_size", + "state_layout", + "pool_mode", + "update_state", + "use_qk_l2norm", + ], "general": [ "batch_size", "hidden_size", @@ -172,6 +182,7 @@ + output_column_dict["sampling"] + output_column_dict["rope"] + output_column_dict["mamba"] + + output_column_dict["gdn"] + output_column_dict["general"] ) @@ -261,6 +272,11 @@ "mamba": [ "selective_state_update", ], + "gdn": [ + "gated_delta_rule_decode", + "gated_delta_rule_mtp", + "chunk_gated_delta_rule", + ], } @@ -933,6 +949,43 @@ def dtype_str_to_torch_dtype(dtype_str): "12.0": ["flashinfer", "triton"], "12.1": ["flashinfer", "triton"], }, + # GDN (Gated Delta Net) + "gated_delta_rule_decode": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": ["flashinfer", "triton"], + "10.0": ["flashinfer", "triton"], + "10.3": ["flashinfer", "triton"], + "11.0": ["triton"], + "12.0": ["triton"], + "12.1": ["triton"], + }, + "gated_delta_rule_mtp": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": ["flashinfer", "triton"], + "10.0": ["flashinfer", "triton"], + "10.3": ["flashinfer", "triton"], + "11.0": ["triton"], + "12.0": ["triton"], + "12.1": ["triton"], + }, + "chunk_gated_delta_rule": { + "7.5": [], + "8.0": [], + "8.6": [], + "8.9": [], + "9.0": ["flashinfer", "fla"], + "10.0": ["flashinfer", "fla"], + "10.3": ["flashinfer", "fla"], + "11.0": [], + "12.0": [], + "12.1": [], + }, } diff --git a/benchmarks/routines/gdn.py b/benchmarks/routines/gdn.py new file mode 100644 index 00000000000..fe23f3a943a --- /dev/null +++ b/benchmarks/routines/gdn.py @@ -0,0 +1,1067 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import os +import sys +from collections import defaultdict + +import numpy as np +import torch + +import flashinfer +from flashinfer.gdn_decode import ( + gated_delta_rule_decode, + gated_delta_rule_decode_pretranspose, + gated_delta_rule_mtp, +) +from flashinfer.gdn_prefill import chunk_gated_delta_rule +from flashinfer.testing.utils import bench_gpu_time + +# Add tests/gdn to sys.path so the torch GDN reference is importable (same +# pattern as routines/mamba.py with tests/mamba), and benchmarks/ for the +# shared Triton GDN kernels. +_repo_root = os.path.normpath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..") +) +_tests_gdn = os.path.join(_repo_root, "tests", "gdn") +if _tests_gdn not in sys.path: + sys.path.insert(0, _tests_gdn) +_benchmarks_dir = os.path.join(_repo_root, "benchmarks") +if _benchmarks_dir not in sys.path: + sys.path.insert(0, _benchmarks_dir) + +from .flashinfer_benchmark_utils import ( + dtype_str_to_torch_dtype, + get_device, + is_close_stats, + print_perf_metrics, + filter_backends_by_compute_capability, + warn_if_pdl_unsupported, +) + +from reference_delta_rule import blockwise_delta_rule, decode_delta_rule +from gdn_triton_reference import ( + TRITON_AVAILABLE, + triton_gdn_decode, + triton_gdn_decode_pretranspose, + triton_gdn_mtp, +) + +try: + from fla.ops.gated_delta_rule.chunk import chunk_gated_delta_rule_fwd as fla_gdn + + FLA_AVAILABLE = True +except ImportError: + FLA_AVAILABLE = False + + +# ============================================================================== +# Benchmark infrastructure +# ============================================================================== + + +def run_gdn_test(args): + """ + Run a GDN (Gated Delta Net) test. + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + if args.routine == "gated_delta_rule_decode": + return testGatedDeltaRuleDecode(args) + elif args.routine == "gated_delta_rule_mtp": + return testGatedDeltaRuleMtp(args) + elif args.routine == "chunk_gated_delta_rule": + return testChunkGatedDeltaRule(args) + else: + raise ValueError(f"Unsupported routine: {args.routine}") + + +def parse_gdn_args(line, parser): + """ + Parse command line arguments for GDN test configuration. + + Args: + line: Command line arguments + parser: ArgumentParser object already populated with shared arguments + + Returns: + Parsed argument namespace + """ + parser.add_argument( + "--batch_size", + type=int, + required=True, + help="Batch size. Decode/MTP: number of concurrent requests. " + "Prefill: number of sequences.", + ) + parser.add_argument( + "--num_q_heads", + type=int, + default=16, + help="Number of query heads.", + ) + parser.add_argument( + "--num_k_heads", + type=int, + default=16, + help="Number of key heads.", + ) + parser.add_argument( + "--num_v_heads", + type=int, + default=32, + help="Number of value heads (GVA when > num_q_heads).", + ) + parser.add_argument( + "--head_size", + type=int, + default=128, + help="Head dimension (K = V = head_size for GDN).", + ) + parser.add_argument( + "--input_dtype", + type=str, + required=False, + default="bfloat16", + choices=["bfloat16", "float16"], + help="Data type for q/k/v/a/b input tensors.", + ) + parser.add_argument( + "--state_dtype", + type=str, + required=False, + default="float32", + choices=["float32", "bfloat16"], + help="Data type for the recurrent state. bfloat16 selects the BF16 " + "state kernels (decode/MTP, requires head_size=128 and pretranspose " + "layout).", + ) + parser.add_argument( + "--state_layout", + type=str, + required=False, + default="pretranspose", + choices=["pretranspose", "nontranspose"], + help="Decode state layout: pretranspose [B, HV, V, K] " + "(gated_delta_rule_decode_pretranspose) or nontranspose [B, HV, K, V] " + "(gated_delta_rule_decode). Decode routine only.", + ) + parser.add_argument( + "--pool_mode", + type=str, + required=False, + default="single", + choices=["single", "split"], + help="State pool indexing mode. 'single': state pool of size B with " + "read == write slots. 'split': pool of size 2B; reads from slots " + "[0..B), writes to slots [B..2B) (speculative-decoding shape). " + "'split' requires the pretranspose layout (decode) or bfloat16 state " + "(MTP).", + ) + parser.add_argument( + "--seq_len", + type=int, + required=False, + default=None, + help="Tokens per request for gated_delta_rule_mtp (must be >= 2, " + "default 2). Not applicable to other GDN routines.", + ) + parser.add_argument( + "--s_qo", + type=int, + required=False, + default=2048, + help="Per-sequence length for chunk_gated_delta_rule (uniform across " + "the batch).", + ) + parser.add_argument( + "--update_state", + action="store_true", + default=False, + help="MTP: write the final state back (disable_state_update=False). " + "The BF16 state path always updates state in-place.", + ) + parser.add_argument( + "--cache_intermediate_states", + action="store_true", + default=False, + help="MTP (float32 state only): cache per-token intermediate states.", + ) + parser.add_argument( + "--no_qk_l2norm", + action="store_true", + default=False, + help="Decode/MTP: disable in-kernel Q/K L2 normalization.", + ) + parser.add_argument( + "--backends", + type=str, + required=False, + nargs="+", + default=["flashinfer"], + choices=["flashinfer", "triton", "fla"], + help="Kernel backends to benchmark. 'triton' is available for " + "decode/MTP; 'fla' (flash-linear-attention) for prefill.", + ) + + args = parser.parse_args(line) + + is_decode = args.routine == "gated_delta_rule_decode" + is_mtp = args.routine == "gated_delta_rule_mtp" + is_prefill = args.routine == "chunk_gated_delta_rule" + + # Resolve seq_len per routine + if is_mtp: + if args.seq_len is None: + args.seq_len = 2 + if args.seq_len < 2: + raise ValueError( + f"gated_delta_rule_mtp requires --seq_len >= 2, got {args.seq_len}" + ) + elif args.seq_len is not None and args.seq_len != 1: + raise ValueError( + f"--seq_len is only applicable to gated_delta_rule_mtp, got " + f"--seq_len {args.seq_len} for {args.routine}" + ) + + if is_decode or is_mtp: + if args.state_dtype == "bfloat16": + if args.head_size != 128: + raise ValueError( + "bfloat16 state requires head_size=128 (BF16 state kernels " + f"support K=V=128 only), got head_size={args.head_size}" + ) + if args.state_layout != "pretranspose": + raise ValueError("bfloat16 state requires --state_layout pretranspose") + if "triton" in args.backends: + raise ValueError( + "The triton backend only supports float32 state; use " + "--state_dtype float32 or drop the triton backend" + ) + if args.state_layout == "nontranspose": + if args.pool_mode != "single": + raise ValueError("--pool_mode split requires the pretranspose layout") + if is_mtp: + raise ValueError( + "gated_delta_rule_mtp uses the K-last [pool, HV, V, K] " + "state layout; --state_layout nontranspose is not " + "applicable" + ) + if args.pool_mode == "split": + if "triton" in args.backends: + raise ValueError( + "The triton backend does not support --pool_mode split" + ) + if is_mtp and args.state_dtype != "bfloat16": + raise ValueError( + "gated_delta_rule_mtp only supports --pool_mode split " + "with --state_dtype bfloat16 (the float32 MTP API has no " + "output_state_indices)" + ) + if is_mtp and args.state_dtype == "bfloat16": + if args.cache_intermediate_states: + raise ValueError( + "--cache_intermediate_states requires --state_dtype " + "float32 (the public BF16 MTP path does not expose " + "intermediate state caching)" + ) + if is_decode and (args.update_state or args.cache_intermediate_states): + raise ValueError( + "--update_state / --cache_intermediate_states are only " + "applicable to gated_delta_rule_mtp" + ) + if "fla" in args.backends: + raise ValueError( + "The fla backend is only available for chunk_gated_delta_rule" + ) + + if is_prefill: + if "triton" in args.backends: + raise ValueError( + "The triton backend is only available for decode/MTP routines; " + "prefill supports flashinfer and fla" + ) + + if args.verbose >= 1: + print(f"[INFO] {args = }") + return args + + +# ============================================================================== +# FLOPs / bytes models +# ============================================================================== + + +def gdn_decode_flops( + batch_size: int, + num_q_heads: int, + num_v_heads: int, + head_size: int, + seq_len: int = 1, +) -> int: + """FLOPs for gated delta rule decode/MTP. + + Per token per output head, three K x V matrix-vector products + (k @ state, outer-product update, q @ state): 6 * head_size^2 FLOPs. + """ + num_o_heads = max(num_q_heads, num_v_heads) + return 6 * seq_len * batch_size * num_o_heads * head_size * head_size + + +def gdn_decode_bytes( + batch_size: int, + num_q_heads: int, + num_k_heads: int, + num_v_heads: int, + head_size: int, + input_dtype: torch.dtype, + seq_len: int = 1, + state_writeback: bool = True, + state_dtype_bytes: int = 4, + cache_intermediate_states: bool = False, +) -> int: + """Memory bytes for gated delta rule decode/MTP. + + Counts q/k/v/output traffic, the state read (+ optional write-back), + GDN parameters, and the optional MTP intermediate state writes. + """ + num_o_heads = max(num_q_heads, num_v_heads) + num_sab_heads = num_o_heads + elem_size = input_dtype.itemsize + + q_bytes = batch_size * seq_len * num_q_heads * head_size * elem_size + k_bytes = batch_size * seq_len * num_k_heads * head_size * elem_size + v_bytes = batch_size * seq_len * num_v_heads * head_size * elem_size + o_bytes = batch_size * seq_len * num_o_heads * head_size * elem_size + + state_elems = batch_size * num_sab_heads * head_size * head_size + state_bytes = state_elems * state_dtype_bytes * (2 if state_writeback else 1) + + # Parameters: A_log [HV] fp32, dt_bias [HV] fp32, a/b [B, T, HV] + param_bytes = ( + num_sab_heads * 4 + + num_sab_heads * 4 + + 2 * batch_size * seq_len * num_sab_heads * elem_size + ) + + intermediate_bytes = 0 + if cache_intermediate_states and seq_len > 1: + intermediate_bytes = ( + batch_size + * seq_len + * num_sab_heads + * head_size + * head_size + * state_dtype_bytes + ) + + return ( + q_bytes + + k_bytes + + v_bytes + + o_bytes + + state_bytes + + param_bytes + + intermediate_bytes + ) + + +def gdn_prefill_flops(total_tokens: int, num_sab_heads: int, head_size: int) -> int: + """FLOPs for chunked GDN prefill. + + Counts the two dominant GEMMs per token per head (kv outer-product + accumulation and q @ state), matching bench_gdn_prefill.py. Intra-chunk + attention terms are excluded for consistency with that convention. + """ + return 2 * 2 * total_tokens * num_sab_heads * head_size * head_size + + +def gdn_prefill_bytes( + total_tokens: int, + num_seqs: int, + num_q_heads: int, + num_k_heads: int, + num_v_heads: int, + head_size: int, + input_dtype: torch.dtype, +) -> int: + """Memory bytes for chunked GDN prefill (q/k/v/g/beta reads, output and + final-state writes).""" + num_sab_heads = max(num_q_heads, num_v_heads) + elem_size = input_dtype.itemsize + + qkv_bytes = ( + total_tokens * (num_q_heads + num_k_heads + num_v_heads) * head_size * elem_size + ) + gate_bytes = 2 * total_tokens * num_sab_heads * 4 # g, beta fp32 + o_bytes = total_tokens * num_sab_heads * head_size * elem_size + state_bytes = num_seqs * num_sab_heads * head_size * head_size * 4 # fp32 + return qkv_bytes + gate_bytes + o_bytes + state_bytes + + +# ============================================================================== +# Decode / MTP +# ============================================================================== + + +def testGatedDeltaRuleDecode(args): + """Test gated_delta_rule_decode (T=1, pretranspose or nontranspose layout).""" + return _testGdnDecodeLike(args, seq_len=1) + + +def testGatedDeltaRuleMtp(args): + """Test gated_delta_rule_mtp (T>1 multi-token processing).""" + return _testGdnDecodeLike(args, seq_len=args.seq_len) + + +def _testGdnDecodeLike(args, seq_len): + """ + Shared implementation for GDN decode (T=1) and MTP (T>1). + + This test: + 1. Generates random q/k/v, GDN gating parameters, and a recurrent state + pool in the requested layout and dtype + 2. Runs the requested backend(s): + - 'flashinfer': CuTe-DSL kernels via the public gdn_decode APIs + - 'triton': Triton reference kernels (tests/gdn/gdn_triton_reference.py) + 3. Optionally checks outputs against the torch reference + (tests/gdn/reference_delta_rule.py::decode_delta_rule) + 4. Measures performance (memory bandwidth is the primary metric) + + Args: + args: Parsed command line arguments containing test configuration + seq_len: Tokens per request (1 = decode, >1 = MTP) + + Returns: + dict: List of dictionaries containing performance results + """ + warn_if_pdl_unsupported(args, args.routine) + if args.verbose >= 1: + print(f"[INFO] Running {args.routine}") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + ## Parse input arguments + backends = args.backends[:] + batch_size = args.batch_size + num_q_heads = args.num_q_heads + num_k_heads = args.num_k_heads + num_v_heads = args.num_v_heads + head_size = args.head_size + state_layout = args.state_layout if seq_len == 1 else "pretranspose" + pool_split = args.pool_mode == "split" + use_qk_l2norm = not args.no_qk_l2norm + is_mtp = seq_len > 1 + is_cuda_graph_compatible = not args.no_cuda_graph + run_refcheck = args.refcheck + res = [] + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if "triton" in backends and not TRITON_AVAILABLE: + print("[WARNING] triton is not installed. Skipping triton backend.") + backends.remove("triton") + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + state_dtype = dtype_str_to_torch_dtype(args.state_dtype) + use_bf16_state = state_dtype == torch.bfloat16 + ## Done parsing input arguments + + num_o_heads = max(num_q_heads, num_v_heads) + num_sab_heads = num_o_heads + T = seq_len + + ## Prepare input tensors + q = torch.randn( + batch_size, T, num_q_heads, head_size, dtype=input_dtype, device=device + ) + k = torch.randn( + batch_size, T, num_k_heads, head_size, dtype=input_dtype, device=device + ) + v = torch.randn( + batch_size, T, num_v_heads, head_size, dtype=input_dtype, device=device + ) + A_log = torch.randn(num_sab_heads, dtype=torch.float32, device=device) + a = torch.randn(batch_size, T, num_sab_heads, dtype=input_dtype, device=device) + dt_bias = torch.randn(num_sab_heads, dtype=torch.float32, device=device) + b = torch.randn(batch_size, T, num_sab_heads, dtype=input_dtype, device=device) + scale = head_size**-0.5 + + # State pool. Layout interpretation: + # pretranspose / MTP: [pool, HV, V, K] (K-last) + # nontranspose: [B, HV, K, V] (V-last) + # K = V = head_size, so the allocation shape is the same either way. + pool_size = 2 * batch_size if pool_split else batch_size + state_pool = torch.randn( + pool_size, + num_sab_heads, + head_size, + head_size, + dtype=state_dtype, + device=device, + ) + initial_state_indices = torch.arange(batch_size, dtype=torch.int32, device=device) + if pool_split: + output_state_indices = torch.arange( + batch_size, 2 * batch_size, dtype=torch.int32, device=device + ) + else: + output_state_indices = None + + # Intermediate states buffer (fp32 MTP only) + intermediate_states_buffer = None + if is_mtp and args.cache_intermediate_states: + intermediate_states_buffer = torch.zeros( + batch_size, + T, + num_sab_heads, + head_size, + head_size, + dtype=torch.float32, + device=device, + ) + + output = torch.empty( + batch_size, T, num_o_heads, head_size, dtype=input_dtype, device=device + ) + + # The BF16 state path always updates state in-place; the fp32 MTP API + # makes the write-back optional. + state_writeback = True if (not is_mtp or use_bf16_state) else args.update_state + disable_state_update = not state_writeback + + if args.verbose >= 2: + print(f"[VVERBOSE] Mode: {'MTP' if is_mtp else 'decode'} (T={T})") + print(f"[VVERBOSE] {q.shape = }, {q.dtype = }") + print(f"[VVERBOSE] {v.shape = }, {v.dtype = }") + print(f"[VVERBOSE] {state_pool.shape = }, {state_pool.dtype = }") + print(f"[VVERBOSE] {state_layout = }, {args.pool_mode = }") + print(f"[VVERBOSE] {use_qk_l2norm = }, {state_writeback = }") + + def run_backend(backend, state): + if backend == "flashinfer": + if is_mtp: + if use_bf16_state: + # BF16 MTP path is reached through the pretranspose API + return gated_delta_rule_decode_pretranspose( + q, + k, + v, + None, + A_log, + a, + dt_bias, + b, + scale, + output, + use_qk_l2norm, + initial_state=state, + initial_state_indices=initial_state_indices, + output_state_indices=output_state_indices, + )[0] + return gated_delta_rule_mtp( + q, + k, + v, + state, + initial_state_indices, + A_log, + a, + dt_bias, + b, + scale, + output, + intermediate_states_buffer=intermediate_states_buffer, + disable_state_update=disable_state_update, + use_qk_l2norm=use_qk_l2norm, + )[0] + if state_layout == "nontranspose": + return gated_delta_rule_decode( + q, k, v, state, A_log, a, dt_bias, b, scale, output, use_qk_l2norm + )[0] + # BF16 state always goes through the pool path with pre-allocated + # indices: the non-pool API path synthesizes an arange index + # tensor per call, whose tiny kernel would pollute CUPTI timing. + if pool_split or use_bf16_state: + return gated_delta_rule_decode_pretranspose( + q, + k, + v, + None, + A_log, + a, + dt_bias, + b, + scale, + output, + use_qk_l2norm, + initial_state=state, + initial_state_indices=initial_state_indices, + output_state_indices=output_state_indices, + )[0] + return gated_delta_rule_decode_pretranspose( + q, k, v, state, A_log, a, dt_bias, b, scale, output, use_qk_l2norm + )[0] + elif backend == "triton": + if is_mtp: + return triton_gdn_mtp( + q, + k, + v, + state, + initial_state_indices, + A_log, + a, + dt_bias, + b, + scale, + output, + intermediate_states_buffer=intermediate_states_buffer, + disable_state_update=disable_state_update, + use_qk_l2norm=use_qk_l2norm, + )[0] + if state_layout == "nontranspose": + return triton_gdn_decode( + q, k, v, state, A_log, a, dt_bias, b, scale, output, use_qk_l2norm + )[0] + return triton_gdn_decode_pretranspose( + q, k, v, state, A_log, a, dt_bias, b, scale, output, use_qk_l2norm + )[0] + else: + raise ValueError(f"Unsupported backend: {backend}") + + # Reference check against the torch reference, which uses the K-major + # state layout [B, HV, K, V] and processes one token at a time. + # bench_gpu_time mutates state_pool in-place, so all refcheck runs use + # clones of the clean snapshot taken before any benchmarking. + has_reference_output = False + clean_state_snapshot = state_pool.clone() if run_refcheck else None + outputs = {} + if run_refcheck: + ref_state = clean_state_snapshot[:batch_size] + if state_layout == "pretranspose": + # [B, HV, V, K] -> [B, HV, K, V] + ref_state = ref_state.transpose(-1, -2) + ref_state = ref_state.contiguous().float() + ref_outs = [] + for t in range(T): + ref_o, ref_state = decode_delta_rule( + q[:, t].float(), + k[:, t].float(), + v[:, t].float(), + ref_state, + A_log=A_log, + a=a[:, t].float(), + dt_bias=dt_bias, + b=b[:, t].float(), + scale_factor=scale, + use_l2_norm=use_qk_l2norm, + state_dtype=state_dtype, + ) + ref_outs.append(ref_o) + reference_output = torch.stack(ref_outs, dim=1) # [B, T, HV, V] + has_reference_output = True + + for cur_backend in backends: + fresh_state = clean_state_snapshot.clone() + outputs[cur_backend] = ( + run_backend(cur_backend, fresh_state).detach().clone() + ) + + # Storage for timing results + backend_times = {backend: [] for backend in backends} + for cur_backend in backends: + backend_times[cur_backend] = bench_gpu_time( + fn=run_backend, + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + input_args=(cur_backend, state_pool), + ) + + # Compare outputs against the torch reference + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 0 and run_refcheck and has_reference_output: + # bf16 state rounds the recurrent state every step in the reference + # but is kept in higher precision inside the kernels, so use looser + # tolerances there. + rtol, atol = (2e-2, 2e-2) if use_bf16_state else (5e-3, 5e-3) + for i in range(len(tested_backends)): + ( + num_different_elements, + num_elements, + num_different_elements_percentage, + ) = is_close_stats( + reference_output.float(), + tested_outputs[i].float(), + rtol=rtol, + atol=atol, + ) + mismatch_threshold_pct = 0.01 + if num_different_elements_percentage > mismatch_threshold_pct: + print( + f"[ERROR] Output tensor mismatch from backend {tested_backends[i]}: " + f"{num_different_elements}/{num_elements} ({num_different_elements_percentage:.4f}%) elements differ " + f"(threshold: {mismatch_threshold_pct}%)" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch with {num_different_elements} elements" + ) + elif args.verbose >= 1: + print( + f"[REFCHECK] Backend {tested_backends[i]}: PASSED " + f"({num_different_elements}/{num_elements} elements differ " + f"({num_different_elements_percentage:.4f}%), within {mismatch_threshold_pct}% threshold)" + ) + + # Compute and report performance metrics + problem_flops = gdn_decode_flops(batch_size, num_q_heads, num_v_heads, head_size, T) + problem_bytes = gdn_decode_bytes( + batch_size, + num_q_heads, + num_k_heads, + num_v_heads, + head_size, + input_dtype, + T, + state_writeback=state_writeback, + state_dtype_bytes=state_dtype.itemsize, + cache_intermediate_states=args.cache_intermediate_states, + ) + + for backend in backends: + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + tflops = problem_flops / (10**9 * median_time) + tb_per_sec = problem_bytes / (10**9 * median_time) + + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["batch_size"] = batch_size + cur_res["num_q_heads"] = num_q_heads + cur_res["num_k_heads"] = num_k_heads + cur_res["num_v_heads"] = num_v_heads + cur_res["head_size"] = head_size + cur_res["seq_len"] = T + cur_res["input_dtype"] = str(input_dtype) + cur_res["state_dtype"] = str(state_dtype) + cur_res["state_layout"] = state_layout + cur_res["pool_mode"] = args.pool_mode + cur_res["update_state"] = state_writeback + cur_res["use_qk_l2norm"] = use_qk_l2norm + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res + + +# ============================================================================== +# Prefill +# ============================================================================== + + +def testChunkGatedDeltaRule(args): + """ + Test chunk_gated_delta_rule (varlen GDN prefill). + + This test: + 1. Generates varlen q/k/v (uniform per-sequence length --s_qo), forget + gate alpha and update gate beta. k is pre-L2-normalized and the kernel + is called with use_qk_l2norm_in_kernel=False so that the torch + reference (which performs no normalization) sees identical inputs. + 2. Runs the requested backend(s): + - 'flashinfer': SM90 C++ / SM100 CuTe-DSL chunked GDN prefill + - 'fla': flash-linear-attention Triton baseline (perf-only, excluded + from refcheck; requires pip install flash-linear-attention) + 3. Optionally checks the flashinfer output against + tests/gdn/reference_delta_rule.py::blockwise_delta_rule + 4. Measures performance metrics + + Args: + args: Parsed command line arguments containing test configuration + + Returns: + dict: List of dictionaries containing performance results + """ + warn_if_pdl_unsupported(args, args.routine) + if args.verbose >= 1: + print("[INFO] Running testChunkGatedDeltaRule") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + ## Parse input arguments + backends = args.backends[:] + num_seqs = args.batch_size + s_qo = args.s_qo + num_q_heads = args.num_q_heads + num_k_heads = args.num_k_heads + num_v_heads = args.num_v_heads + head_size = args.head_size + is_cuda_graph_compatible = not args.no_cuda_graph + run_refcheck = args.refcheck + res = [] + + backends = filter_backends_by_compute_capability(backends, args.routine, device) + if "fla" in backends and not FLA_AVAILABLE: + print( + "[WARNING] fla is not installed (pip install flash-linear-attention). " + "Skipping fla backend." + ) + backends.remove("fla") + if len(backends) == 0: + print("[ERROR] No backends to test. Exiting.") + return res + + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + ## Done parsing input arguments + + num_o_heads = max(num_q_heads, num_v_heads) + num_sab_heads = num_o_heads + seq_lens = [s_qo] * num_seqs + total_tokens = num_seqs * s_qo + cu_seqlens = torch.arange( + 0, total_tokens + 1, s_qo, dtype=torch.int64, device=device + ) + + ## Prepare input tensors + q = torch.randn( + total_tokens, num_q_heads, head_size, dtype=input_dtype, device=device + ) + # Pre-normalize k for numerical stability (matches bench_gdn_prefill.py); + # the kernel and the reference then see identical inputs. + k = torch.nn.functional.normalize( + torch.randn( + total_tokens, num_k_heads, head_size, dtype=torch.float32, device=device + ), + p=2.0, + dim=-1, + ).to(input_dtype) + v = torch.randn( + total_tokens, num_v_heads, head_size, dtype=input_dtype, device=device + ) + # Forget gate alpha in (0, 1) and update gate beta in (0, 1), fp32 + g = torch.rand(total_tokens, num_sab_heads, dtype=torch.float32, device=device) + beta = torch.rand(total_tokens, num_sab_heads, dtype=torch.float32, device=device) + scale = head_size**-0.5 + + output = torch.empty( + total_tokens, num_o_heads, head_size, dtype=input_dtype, device=device + ) + output_state = torch.empty( + num_seqs, + num_sab_heads, + head_size, + head_size, + dtype=torch.float32, + device=device, + ) + + # FLA baseline tensors: batch-dim layout [1, T, H, D], log-space forget + # gate, q expanded to HV heads for GVA, int32 cu_seqlens. + if "fla" in backends: + q_fla = ( + q.repeat_interleave(num_sab_heads // num_q_heads, dim=1) + if num_sab_heads > num_q_heads + else q + ).unsqueeze(0) + k_fla = ( + k.repeat_interleave(num_sab_heads // num_k_heads, dim=1) + if num_sab_heads > num_k_heads + else k + ).unsqueeze(0) + v_fla = v.unsqueeze(0) + g_fla = torch.log(g.clamp_min(1e-6)).unsqueeze(0) + beta_fla = beta.unsqueeze(0) + h0_fla = torch.zeros( + num_seqs, + num_sab_heads, + head_size, + head_size, + dtype=torch.float32, + device=device, + ) + cu_seqlens_fla = cu_seqlens.to(torch.int32) + + if args.verbose >= 2: + print(f"[VVERBOSE] {q.shape = }, {q.dtype = }") + print(f"[VVERBOSE] {v.shape = }, {v.dtype = }") + print(f"[VVERBOSE] {cu_seqlens = }") + print(f"[VVERBOSE] {g.shape = }, {beta.shape = }") + + def run_backend(backend): + if backend == "flashinfer": + return chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + scale, + None, # initial_state (zero state) + True, # output_final_state + cu_seqlens, + False, # use_qk_l2norm_in_kernel (k is pre-normalized) + output=output, + output_state=output_state, + )[0] + elif backend == "fla": + return fla_gdn( + q_fla, + k_fla, + v_fla, + g_fla, + beta_fla, + None, + initial_state=h0_fla, + output_final_state=True, + cu_seqlens=cu_seqlens_fla, + )[0] + else: + raise ValueError(f"Unsupported backend: {backend}") + + # Reference check (flashinfer only; fla uses a different gate + # parameterization and is benchmarked perf-only) + has_reference_output = False + outputs = {} + if run_refcheck: + reference_output = blockwise_delta_rule( + q.float(), + k.float(), + v.float(), + seq_lens, + alpha=g, + beta=beta, + scale_factor=scale, + state_dtype=torch.float32, + )[0] + has_reference_output = True + for cur_backend in backends: + if cur_backend == "fla": + continue + outputs[cur_backend] = run_backend(cur_backend).detach().clone() + + # Storage for timing results + backend_times = {backend: [] for backend in backends} + for cur_backend in backends: + backend_times[cur_backend] = bench_gpu_time( + fn=run_backend, + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + input_args=(cur_backend,), + ) + + # Compare outputs against the torch reference + tested_backends = list(outputs.keys()) + tested_outputs = list(outputs.values()) + if len(tested_backends) > 0 and run_refcheck and has_reference_output: + rtol, atol = 1e-2, 1e-2 + for i in range(len(tested_backends)): + ( + num_different_elements, + num_elements, + num_different_elements_percentage, + ) = is_close_stats( + reference_output.float(), + tested_outputs[i].float(), + rtol=rtol, + atol=atol, + ) + mismatch_threshold_pct = 0.01 + if num_different_elements_percentage > mismatch_threshold_pct: + print( + f"[ERROR] Output tensor mismatch from backend {tested_backends[i]}: " + f"{num_different_elements}/{num_elements} ({num_different_elements_percentage:.4f}%) elements differ " + f"(threshold: {mismatch_threshold_pct}%)" + ) + if not args.allow_output_mismatch: + raise AssertionError( + f"[ERROR] Backend {tested_backends[i]} output mismatch with {num_different_elements} elements" + ) + elif args.verbose >= 1: + print( + f"[REFCHECK] Backend {tested_backends[i]}: PASSED " + f"({num_different_elements}/{num_elements} elements differ " + f"({num_different_elements_percentage:.4f}%), within {mismatch_threshold_pct}% threshold)" + ) + + # Compute and report performance metrics + problem_flops = gdn_prefill_flops(total_tokens, num_sab_heads, head_size) + problem_bytes = gdn_prefill_bytes( + total_tokens, + num_seqs, + num_q_heads, + num_k_heads, + num_v_heads, + head_size, + input_dtype, + ) + + for backend in backends: + if len(backend_times[backend]) > 0: + median_time = np.median(backend_times[backend]) + std_time = np.std(backend_times[backend]) + tflops = problem_flops / (10**9 * median_time) + tb_per_sec = problem_bytes / (10**9 * median_time) + + print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec) + + if args.output_path is not None: + cur_res = defaultdict(str) + cur_res["routine"] = args.routine + cur_res["median_time"] = median_time + cur_res["std_time"] = std_time + cur_res["tflops"] = tflops + cur_res["tb_per_sec"] = tb_per_sec + cur_res["backend"] = backend + cur_res["batch_size"] = num_seqs + cur_res["s_qo"] = s_qo + cur_res["num_q_heads"] = num_q_heads + cur_res["num_k_heads"] = num_k_heads + cur_res["num_v_heads"] = num_v_heads + cur_res["head_size"] = head_size + cur_res["input_dtype"] = str(input_dtype) + cur_res["case_tag"] = args.case_tag + res.append(cur_res) + return res diff --git a/benchmarks/samples/sample_testlist.txt b/benchmarks/samples/sample_testlist.txt index d812719ea28..1ff9ac03804 100644 --- a/benchmarks/samples/sample_testlist.txt +++ b/benchmarks/samples/sample_testlist.txt @@ -314,3 +314,25 @@ # FlashInfer-only (no refcheck, perf focus) --routine selective_state_update --batch_size 128 --nheads 64 --dim 128 --dstate 128 --ngroups 8 --backends flashinfer -vv --generate_repro_command --case_tag "mamba2_stp_perf" + +## GDN (Gated Delta Net, SM90+) +# Decode T=1, Qwen3-Next config (q=k=16, v=32, d=128), pretranspose layout, FlashInfer vs Triton +--routine gated_delta_rule_decode --batch_size 128 --num_q_heads 16 --num_k_heads 16 --num_v_heads 32 --head_size 128 --backends flashinfer triton --refcheck -vv --generate_repro_command --case_tag "gdn_decode_pretrans" + +# Decode T=1, nontranspose layout +--routine gated_delta_rule_decode --batch_size 128 --num_q_heads 16 --num_k_heads 16 --num_v_heads 32 --head_size 128 --state_layout nontranspose --backends flashinfer triton --refcheck -vv --generate_repro_command --case_tag "gdn_decode_nontrans" + +# Decode T=1, BF16 state kernel +--routine gated_delta_rule_decode --batch_size 128 --num_q_heads 16 --num_k_heads 16 --num_v_heads 32 --head_size 128 --state_dtype bfloat16 --backends flashinfer --refcheck -vv --generate_repro_command --case_tag "gdn_decode_bf16_state" + +# Decode T=1, BF16 state, split pool (speculative-decoding shape) +--routine gated_delta_rule_decode --batch_size 128 --num_q_heads 16 --num_k_heads 16 --num_v_heads 32 --head_size 128 --state_dtype bfloat16 --pool_mode split --backends flashinfer -vv --generate_repro_command --case_tag "gdn_decode_bf16_split" + +# MTP T=4, fp32 state pool, FlashInfer vs Triton +--routine gated_delta_rule_mtp --batch_size 128 --seq_len 4 --num_q_heads 16 --num_k_heads 16 --num_v_heads 32 --head_size 128 --backends flashinfer triton --refcheck -vv --generate_repro_command --case_tag "gdn_mtp4_fp32" + +# MTP T=4, BF16 state +--routine gated_delta_rule_mtp --batch_size 128 --seq_len 4 --num_q_heads 16 --num_k_heads 16 --num_v_heads 32 --head_size 128 --state_dtype bfloat16 --backends flashinfer --refcheck -vv --generate_repro_command --case_tag "gdn_mtp4_bf16" + +# Prefill, 4 x 2048 tokens, Qwen3.5-35B head config +--routine chunk_gated_delta_rule --batch_size 4 --s_qo 2048 --num_q_heads 16 --num_k_heads 16 --num_v_heads 32 --head_size 128 --backends flashinfer --refcheck -vv --generate_repro_command --case_tag "gdn_prefill_4x2048" From 2e045334165a8d809abbc69d221ccd9a92a5ef53 Mon Sep 17 00:00:00 2001 From: Ruoqian Guo <22525902+ruoqianguo@users.noreply.github.com> Date: Thu, 11 Jun 2026 09:11:03 +0800 Subject: [PATCH 11/13] Add BF16 MoE SwiGLU OA params (#3532) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## πŸ“Œ Description Add optional per-expert SwiGLU OA activation parameter support to the TRT-LLM BF16 MoE APIs: - `gemm1_alpha` - `gemm1_beta` - `gemm1_clamp_limit` These are optional float32 CUDA tensors with shape `[local_num_experts]` and default to `None`, so existing callers keep the same behavior. This PR threads the new parameters through the BF16 C++ launcher, Python registered op/fake op, public APIs, autotuner path, `MoERunner`, trace metadata, and Python reference implementation. For BF16, these values are passed through as raw scalar values. No host-side dequant-scale adjustment is applied. The SwiGLU OA formula is: ```python activation = X2 * sigmoid(alpha * X2) * (X1 + beta) with optional clamp behavior: X1 = clamp(X1, min=-limit, max=limit) X2 = clamp(X2, max=limit) ``` The implementation also validates that BF16 OA parameters are only used with ActivationType.Swiglu. ## πŸ” Related Issues N/A ## πŸš€ Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### βœ… Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## πŸ§ͺ Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Local checks: python3 -m compileall flashinfer/fused_moe/core.py flashinfer/trace/templates/moe.py tests/moe/test_trtllm_gen_fused_moe.py git diff --check pre-commit run --all-files GPU validation on computelab GB100 SM100 node: tests/moe/test_trtllm_gen_fused_moe.py::test_bf16_moe_swiglu_oa_activation_param_validation tests/moe/test_trtllm_gen_fused_moe.py::test_bf16_moe_swiglu_oa_activation_params Result: 2 passed in 605.52s (0:10:05) ## Reviewer Notes This follows the MXFP8 SwiGLU OA API shape and raw-parameter semantics: BF16 passes `gemm1_alpha`, `gemm1_beta`, and `gemm1_clamp_limit` directly to trtllm-gen without host-side `dequantScaleAb` adjustment. `gemm1_alpha` is never scaled. ## Summary by CodeRabbit * **New Features** * Added optional per-expert BF16 SwiGLU activation parameters (alpha, beta, clamp limit) to MoE public APIs and runtime, enabling per-expert activation shaping and clamping. * **Validation** * Added runtime/schema checks to ensure these optional parameters are well-formed and only accepted with SwiGLU activations. * **Trace Templates** * BF16 MoE trace templates updated to include the new optional SwiGLU/OA inputs. * **Tests** * Added unit/trace tests covering parameter schema, behavior (clamp/OA/default), and production vs. reference correctness. --------- Co-authored-by: Sam (Kesen Li) --- csrc/trtllm_fused_moe_kernel_launcher.cu | 75 +++++-- flashinfer/fused_moe/core.py | 103 +++++++++ flashinfer/trace/templates/moe.py | 67 +++++- tests/moe/test_trtllm_gen_fused_moe.py | 251 +++++++++++++++++++++- tests/trace/test_trtllm_bf16_moe_trace.py | 150 +++++++++++++ 5 files changed, 623 insertions(+), 23 deletions(-) create mode 100644 tests/trace/test_trtllm_bf16_moe_trace.py diff --git a/csrc/trtllm_fused_moe_kernel_launcher.cu b/csrc/trtllm_fused_moe_kernel_launcher.cu index 5a73e720769..b7494e42b46 100644 --- a/csrc/trtllm_fused_moe_kernel_launcher.cu +++ b/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -360,6 +360,23 @@ class FusedMoeLauncher { } } + void check_optional_per_expert_float_tensor(Optional const& tensor, + std::string const& tensor_name) const { + if (!tensor.has_value()) { + return; + } + auto const& value = tensor.value(); + TVM_FFI_ICHECK(value.device().device_type == kDLCUDA) + << tensor_name << " must be a CUDA tensor."; + TVM_FFI_ICHECK(value.device().device_id == hidden_states.device().device_id) + << tensor_name << " must be on the same device as hidden_states."; + TVM_FFI_ICHECK_EQ(value.dtype(), dl_float32) << tensor_name << " must be float32."; + TVM_FFI_ICHECK_EQ(value.ndim(), 1) << tensor_name << " must be 1D."; + TVM_FFI_ICHECK_EQ(value.size(0), args->local_num_experts) + << tensor_name << " must have shape [local_num_experts]."; + TVM_FFI_ICHECK(value.IsContiguous()) << tensor_name << " must be contiguous."; + } + void check_routing_common() const { TVM_FFI_ICHECK(args->top_k > 0 && args->top_k <= args->num_experts) << "top_k must be between 1 and num_experts"; @@ -652,12 +669,17 @@ class Bf16MoeLauncher : public FusedMoeLauncher { Optional const& routing_bias, TensorView const& expert_indices, TensorView const& expert_weights, TensorView const& hidden_states, TensorView const& gemm1_weights, TensorView const& gemm2_weights, - Optional const& gemm1_bias) + Optional const& gemm1_bias, Optional const& gemm1_alpha, + Optional const& gemm1_beta, + Optional const& gemm1_clamp_limit) : FusedMoeLauncher(routing_logits, routing_bias, hidden_states, gemm1_weights, gemm1_bias, Optional(), Optional(), gemm2_weights, Optional(), Optional()), expert_indices(expert_indices), - expert_weights(expert_weights) {} + expert_weights(expert_weights), + gemm1_alpha(gemm1_alpha), + gemm1_beta(gemm1_beta), + gemm1_clamp_limit(gemm1_clamp_limit) {} void init(std::unique_ptr&& args, int64_t tile_tokens_dim, int64_t routing_method_type, bool use_shuffled_weight, @@ -725,6 +747,14 @@ class Bf16MoeLauncher : public FusedMoeLauncher { << "BF16 Moe: weight_layout must be BlockMajorK"; check_weights_shape("gemm1"); check_weights_shape("gemm2"); + check_optional_per_expert_float_tensor(gemm1_alpha, "gemm1_alpha"); + check_optional_per_expert_float_tensor(gemm1_beta, "gemm1_beta"); + check_optional_per_expert_float_tensor(gemm1_clamp_limit, "gemm1_clamp_limit"); + if (gemm1_alpha.has_value() || gemm1_beta.has_value() || gemm1_clamp_limit.has_value()) { + TVM_FFI_ICHECK(activation_type == ActivationType::Swiglu) + << "gemm1_alpha, gemm1_beta, and gemm1_clamp_limit are only supported for " + "ActivationType::Swiglu."; + } TVM_FFI_ICHECK_EQ(args->intermediate_size % 128, 0) << "the second dimension of weights must be a multiple of 128."; @@ -755,6 +785,13 @@ class Bf16MoeLauncher : public FusedMoeLauncher { args->output = output.data_ptr(); } args->output_scale = nullptr; + args->gemm1_alpha = + gemm1_alpha.has_value() ? static_cast(gemm1_alpha.value().data_ptr()) : nullptr; + args->gemm1_beta = + gemm1_beta.has_value() ? static_cast(gemm1_beta.value().data_ptr()) : nullptr; + args->gemm1_clamp_limit = gemm1_clamp_limit.has_value() + ? static_cast(gemm1_clamp_limit.value().data_ptr()) + : nullptr; } static Array> getValidConfigs(int64_t top_k, int64_t hidden_size, @@ -790,6 +827,9 @@ class Bf16MoeLauncher : public FusedMoeLauncher { private: TensorView expert_weights; TensorView expert_indices; + Optional gemm1_alpha; + Optional gemm1_beta; + Optional gemm1_clamp_limit; }; class Fp8PerTensorLauncher : public FusedMoeLauncher { @@ -1976,19 +2016,18 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { } }; -Array trtllm_bf16_moe(Optional const& routing_logits, - Optional const& routing_bias, - TensorView const& expert_indices, TensorView const& expert_weights, - TensorView const& hidden_states, TensorView const& gemm1_weights, - TensorView const& gemm2_weights, - Optional const& gemm1_lora_delta, TensorView output, - int64_t num_experts, int64_t top_k, Optional n_group, - Optional topk_group, int64_t intermediate_size, - int64_t local_expert_offset, int64_t local_num_experts, - Optional routed_scaling_factor, int64_t routing_method_type, - bool use_shuffled_weight, int64_t weight_layout, bool do_finalize, - bool enable_pdl, Array moe_tactic, int64_t activation_type, - bool norm_topk_prob, Optional routing_replay_out) { +Array trtllm_bf16_moe( + Optional const& routing_logits, Optional const& routing_bias, + TensorView const& expert_indices, TensorView const& expert_weights, + TensorView const& hidden_states, TensorView const& gemm1_weights, + TensorView const& gemm2_weights, Optional const& gemm1_lora_delta, + Optional const& gemm1_alpha, Optional const& gemm1_beta, + Optional const& gemm1_clamp_limit, TensorView output, int64_t num_experts, + int64_t top_k, Optional n_group, Optional topk_group, + int64_t intermediate_size, int64_t local_expert_offset, int64_t local_num_experts, + Optional routed_scaling_factor, int64_t routing_method_type, bool use_shuffled_weight, + int64_t weight_layout, bool do_finalize, bool enable_pdl, Array moe_tactic, + int64_t activation_type, bool norm_topk_prob, Optional routing_replay_out) { // Just some basic type validation first and leave more checks to the launcher if (routing_logits.has_value()) { TVM_FFI_ICHECK(routing_logits.value().dtype() == dl_float32 || @@ -2043,9 +2082,9 @@ Array trtllm_bf16_moe(Optional const& routing_logits, args->output_scale = nullptr; // Create and initialize launcher for this tile size - auto launcher = std::make_unique(routing_logits, routing_bias, expert_indices, - expert_weights, hidden_states, gemm1_weights, - gemm2_weights, gemm1_lora_delta); + auto launcher = std::make_unique( + routing_logits, routing_bias, expert_indices, expert_weights, hidden_states, gemm1_weights, + gemm2_weights, gemm1_lora_delta, gemm1_alpha, gemm1_beta, gemm1_clamp_limit); launcher->init(std::move(args), curr_tile_N, routing_method_type, use_shuffled_weight, weight_layout, activation, static_cast(gemm1_bias_type_enum), norm_topk_prob); diff --git a/flashinfer/fused_moe/core.py b/flashinfer/fused_moe/core.py index c272a59743e..3c4471cfa94 100644 --- a/flashinfer/fused_moe/core.py +++ b/flashinfer/fused_moe/core.py @@ -1392,6 +1392,9 @@ def forward( kwargs["gemm1_weights"], kwargs["gemm2_weights"], moe_inputs.gemm1_lora_delta, + kwargs.get("gemm1_alpha"), + kwargs.get("gemm1_beta"), + kwargs.get("gemm1_clamp_limit"), output, kwargs["num_experts"], self.top_k, @@ -1606,10 +1609,21 @@ def trtllm_bf16_moe_op( activation_type: int = ActivationType.Swiglu.value, norm_topk_prob: bool = True, routing_replay_out: Optional[torch.Tensor] = None, + gemm1_alpha: Optional[torch.Tensor] = None, + gemm1_beta: Optional[torch.Tensor] = None, + gemm1_clamp_limit: Optional[torch.Tensor] = None, ) -> List[torch.Tensor]: assert routing_logits is not None or topk_ids is not None, ( "either routing_logits or topk_ids must be provided" ) + _validate_bf16_gemm1_activation_params( + activation_type, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + local_num_experts, + hidden_states.device, + ) if enable_pdl is None: enable_pdl = device_support_pdl(hidden_states.device) @@ -1682,6 +1696,9 @@ def trtllm_bf16_moe_op( routing_bias=routing_bias, gemm1_weights=gemm1_weights, gemm2_weights=gemm2_weights, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, num_experts=num_experts, n_group=n_group, topk_group=topk_group, @@ -1706,6 +1723,9 @@ def trtllm_bf16_moe_op( gemm1_weights, gemm2_weights, gemm1_lora_delta, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, output, num_experts, top_k, @@ -1757,6 +1777,9 @@ def _fake_trtllm_bf16_moe( activation_type: int = ActivationType.Swiglu.value, norm_topk_prob: bool = True, routing_replay_out: Optional[torch.Tensor] = None, + gemm1_alpha: Optional[torch.Tensor] = None, + gemm1_beta: Optional[torch.Tensor] = None, + gemm1_clamp_limit: Optional[torch.Tensor] = None, ) -> List[torch.Tensor]: _ = routing_replay_out seq_len = hidden_states.shape[0] @@ -2644,6 +2667,36 @@ def _fake_trtllm_mxint4_block_scale_moe( ) +def _validate_bf16_gemm1_activation_params( + activation_type: int, + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + local_num_experts: int, + device: torch.device, +) -> None: + if gemm1_alpha is None and gemm1_beta is None and gemm1_clamp_limit is None: + return + if int(activation_type) != int(ActivationType.Swiglu): + raise ValueError( + "gemm1_alpha, gemm1_beta, and gemm1_clamp_limit are only supported " + "for ActivationType.Swiglu." + ) + for name, tensor in ( + ("gemm1_alpha", gemm1_alpha), + ("gemm1_beta", gemm1_beta), + ("gemm1_clamp_limit", gemm1_clamp_limit), + ): + if tensor is not None: + check_shape_dtype_device( + tensor, + (local_num_experts,), + torch.float32, + device, + name, + ) + + def _validate_routing_replay_out( routing_replay_out: Optional[torch.Tensor], top_k: int ) -> None: @@ -2690,6 +2743,9 @@ def trtllm_bf16_moe( activation_type: int = ActivationType.Swiglu.value, norm_topk_prob: bool = True, routing_replay_out: Optional[torch.Tensor] = None, + gemm1_alpha: Optional[torch.Tensor] = None, + gemm1_beta: Optional[torch.Tensor] = None, + gemm1_clamp_limit: Optional[torch.Tensor] = None, ) -> Union[List[torch.Tensor], torch.Tensor]: r"""BF16 MoE operation with autotuning support. @@ -2783,6 +2839,17 @@ def trtllm_bf16_moe( kernel skips the write entirely. The buffer may be larger than ``num_tokens`` for CUDA-graph pre-allocation; only rows ``[0, num_tokens)`` are written. + gemm1_alpha / gemm1_beta / gemm1_clamp_limit : Optional[torch.Tensor] + Optional ``[local_num_experts]`` float32 CUDA per-expert SwiGLU OA + parameters. They are supported with ``ActivationType.Swiglu``. Any + subset can be provided: ``gemm1_alpha=None`` uses ``alpha=1.0``, + ``gemm1_beta=None`` uses ``beta=0.0``, and + ``gemm1_clamp_limit=None`` applies no clamp. Let GEMM1 output be split + as ``X1`` (linear/up half) and ``X2`` (gate half). If a clamp limit is + provided, ``X1 = clamp(X1, -limit, limit)`` and + ``X2 = clamp(X2, max=limit)``. The fused activation output is + ``X2 * sigmoid(alpha * X2) * (X1 + beta)``. Pass raw BF16-path values; + no host-side scalar dequant-scale conversion is applied. Returns ------- @@ -2792,6 +2859,14 @@ def trtllm_bf16_moe( ``[gemm2_output, expert_weights, expanded_idx_to_permuted_idx]``. """ _validate_routing_replay_out(routing_replay_out, top_k) + _validate_bf16_gemm1_activation_params( + activation_type, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + local_num_experts, + hidden_states.device, + ) result = get_trtllm_moe_sm100_module().trtllm_bf16_moe( routing_logits, routing_bias, @@ -2818,6 +2893,9 @@ def trtllm_bf16_moe( activation_type, norm_topk_prob, routing_replay_out, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, ) if do_finalize: @@ -2852,6 +2930,9 @@ def trtllm_bf16_routed_moe( tune_max_num_tokens: int = 8192, activation_type: int = ActivationType.Swiglu.value, routing_replay_out: Optional[torch.Tensor] = None, + gemm1_alpha: Optional[torch.Tensor] = None, + gemm1_beta: Optional[torch.Tensor] = None, + gemm1_clamp_limit: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, List[torch.Tensor]]: r"""Pre-routed BF16 MoE operation with autotuning support. @@ -2940,6 +3021,17 @@ def trtllm_bf16_routed_moe( kernel skips the write entirely. The buffer may be larger than ``num_tokens`` for CUDA-graph pre-allocation; only rows ``[0, num_tokens)`` are written. + gemm1_alpha / gemm1_beta / gemm1_clamp_limit : Optional[torch.Tensor] + Optional ``[local_num_experts]`` float32 CUDA per-expert SwiGLU OA + parameters. They are supported with ``ActivationType.Swiglu``. Any + subset can be provided: ``gemm1_alpha=None`` uses ``alpha=1.0``, + ``gemm1_beta=None`` uses ``beta=0.0``, and + ``gemm1_clamp_limit=None`` applies no clamp. Let GEMM1 output be split + as ``X1`` (linear/up half) and ``X2`` (gate half). If a clamp limit is + provided, ``X1 = clamp(X1, -limit, limit)`` and + ``X2 = clamp(X2, max=limit)``. The fused activation output is + ``X2 * sigmoid(alpha * X2) * (X1 + beta)``. Pass raw BF16-path values; + no host-side scalar dequant-scale conversion is applied. Returns ------- @@ -2956,6 +3048,14 @@ def trtllm_bf16_routed_moe( ============= ================== ========================================================================= """ _validate_routing_replay_out(routing_replay_out, top_k) + _validate_bf16_gemm1_activation_params( + activation_type, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + local_num_experts, + hidden_states.device, + ) result = get_trtllm_moe_sm100_module().trtllm_bf16_moe( None, None, @@ -2982,6 +3082,9 @@ def trtllm_bf16_routed_moe( activation_type, True, # norm_topk_prob: not used for pre-computed routing routing_replay_out, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, ) if do_finalize and gemm1_lora_delta is None: diff --git a/flashinfer/trace/templates/moe.py b/flashinfer/trace/templates/moe.py index dfea21e1e83..a72512fd56a 100644 --- a/flashinfer/trace/templates/moe.py +++ b/flashinfer/trace/templates/moe.py @@ -1817,6 +1817,9 @@ def _moe_bf16_run_experts( topk_idx, local_expert_offset, E_global, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, ): """Un-quantized (bf16) MoE expert computation with SwiGLU.""" T, H = hidden_states.shape @@ -1839,12 +1842,34 @@ def _moe_bf16_run_experts( A_e = A.index_select(0, token_idx) G1 = A_e.matmul(W1[le].t()) X1, X2 = G1[:, :I], G1[:, I:] - silu_X2 = X2 / (1.0 + torch.exp(-X2)) - O = (silu_X2 * X1).matmul(W2[le].t()) + if gemm1_clamp_limit is not None: + limit = gemm1_clamp_limit[le].to(device=X1.device, dtype=torch.float32) + X1 = torch.clamp(X1, min=-limit, max=limit) + X2 = torch.clamp(X2, max=limit) + if ( + gemm1_alpha is not None + or gemm1_beta is not None + or gemm1_clamp_limit is not None + ): + alpha = ( + 1.0 + if gemm1_alpha is None + else gemm1_alpha[le].to(device=X2.device, dtype=torch.float32) + ) + beta = ( + 0.0 + if gemm1_beta is None + else gemm1_beta[le].to(device=X1.device, dtype=torch.float32) + ) + activation = X2 * torch.sigmoid(alpha * X2) * (X1 + beta) + else: + silu_X2 = X2 / (1.0 + torch.exp(-X2)) + activation = silu_X2 * X1 + expert_out = activation.matmul(W2[le].t()) w_tok = weights.index_select(0, token_idx) match = (topk_idx.index_select(0, token_idx) == ge).float() w_e = (w_tok * match).sum(dim=1) - output.index_add_(0, token_idx, O * w_e.unsqueeze(1)) + output.index_add_(0, token_idx, expert_out * w_e.unsqueeze(1)) return output.to(torch.bfloat16) @@ -1891,6 +1916,9 @@ def _trtllm_bf16_moe_reference( top_k, local_expert_offset, routed_scaling_factor=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, **_unused, ): """Reference for TRT-LLM BF16 MoE (Default routing).""" @@ -1905,6 +1933,9 @@ def _trtllm_bf16_moe_reference( topk_idx, local_expert_offset, int(num_experts), + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, ) @@ -1918,6 +1949,9 @@ def _trtllm_bf16_routed_moe_reference( top_k, local_expert_offset, routed_scaling_factor=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, **_unused, ): """Reference for TRT-LLM BF16 MoE with precomputed topk_ids.""" @@ -1938,6 +1972,9 @@ def _trtllm_bf16_routed_moe_reference( topk_ids.to(torch.int64), local_expert_offset, int(num_experts), + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, ) @@ -2225,13 +2262,34 @@ def _unpack_int4(packed): ), } +_TRTLLM_BF16_SWIGLU_OA_INPUTS: dict[str, Tensor] = { + "gemm1_alpha": Tensor( + ["num_local_experts"], + dtype="float32", + optional=True, + description="Optional per-expert SwiGLU OA alpha.", + ), + "gemm1_beta": Tensor( + ["num_local_experts"], + dtype="float32", + optional=True, + description="Optional per-expert SwiGLU OA beta.", + ), + "gemm1_clamp_limit": Tensor( + ["num_local_experts"], + dtype="float32", + optional=True, + description="Optional per-expert SwiGLU OA clamp limit.", + ), +} + # BF16 MoE (no quantization) trtllm_bf16_moe_trace = TraceTemplate( op_type="moe", name_prefix="trtllm_bf16_moe", description="TRT-LLM BF16 MoE (no quantization).", axes=dict(_TRTLLM_MOE_COMMON_AXES), - inputs=dict(_TRTLLM_MOE_COMMON_INPUTS), + inputs={**_TRTLLM_MOE_COMMON_INPUTS, **_TRTLLM_BF16_SWIGLU_OA_INPUTS}, outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS), tags=["status:verified", "backend:trtllm"], reference=_trtllm_bf16_moe_reference, @@ -2265,6 +2323,7 @@ def _unpack_int4(packed): "top_k": _TRTLLM_MOE_COMMON_INPUTS["top_k"], "local_expert_offset": _TRTLLM_MOE_COMMON_INPUTS["local_expert_offset"], "routed_scaling_factor": _TRTLLM_MOE_COMMON_INPUTS["routed_scaling_factor"], + **_TRTLLM_BF16_SWIGLU_OA_INPUTS, }, outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS), tags=["status:verified", "backend:trtllm"], diff --git a/tests/moe/test_trtllm_gen_fused_moe.py b/tests/moe/test_trtllm_gen_fused_moe.py index bd3b76dc0ef..e7b995e134d 100644 --- a/tests/moe/test_trtllm_gen_fused_moe.py +++ b/tests/moe/test_trtllm_gen_fused_moe.py @@ -1656,6 +1656,9 @@ def call_moe( activation_type = kwargs["activation_type"] norm_topk_prob = kwargs.get("norm_topk_prob", True) gemm1_lora_delta = kwargs.get("gemm1_lora_delta") + gemm1_alpha = kwargs.get("gemm1_alpha") + gemm1_beta = kwargs.get("gemm1_beta") + gemm1_clamp_limit = kwargs.get("gemm1_clamp_limit") permute_info = kwargs.get("permute_info") # Use autotuner for optimal kernel selection @@ -1683,6 +1686,9 @@ def call_moe( tune_max_num_tokens=TUNE_MAX_NUM_TOKENS, activation_type=activation_type, gemm1_lora_delta=gemm1_lora_delta, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, ) else: output = trtllm_bf16_moe( @@ -1705,6 +1711,9 @@ def call_moe( tune_max_num_tokens=TUNE_MAX_NUM_TOKENS, activation_type=activation_type, norm_topk_prob=norm_topk_prob, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, ) if isinstance(output, list): return output[0].to(torch.float) @@ -1763,6 +1772,9 @@ def __init__( gemm1_bias=None, gemm2_bias=None, gemm1_lora_delta=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, ): self.num_tokens = num_tokens self.num_experts = num_experts @@ -1786,6 +1798,9 @@ def __init__( self.gemm1_bias = gemm1_bias self.gemm2_bias = gemm2_bias self.gemm1_lora_delta = gemm1_lora_delta + self.gemm1_alpha = gemm1_alpha + self.gemm1_beta = gemm1_beta + self.gemm1_clamp_limit = gemm1_clamp_limit class moe_args_dequant: @@ -1810,6 +1825,9 @@ def __init__( gemm1_bias=None, gemm2_bias=None, gemm1_lora_delta=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, ): self.num_tokens = num_tokens self.num_experts = num_experts @@ -1828,6 +1846,9 @@ def __init__( self.gemm1_bias = gemm1_bias self.gemm2_bias = gemm2_bias self.gemm1_lora_delta = gemm1_lora_delta + self.gemm1_alpha = gemm1_alpha + self.gemm1_beta = gemm1_beta + self.gemm1_clamp_limit = gemm1_clamp_limit def routing_reference(expertLogits, topK, padding): @@ -2414,7 +2435,39 @@ def run_moe_dequant(args, quant_mode: QuantMode): if is_gated_activation(args.activation_type): my_x1 = my_a[:, : args.intermediate_size] my_x2 = my_a[:, args.intermediate_size :] - activation_output[i : i + my_num_tokens] = activation_func(my_x2) * my_x1 + if args.gemm1_clamp_limit is not None: + limit = args.gemm1_clamp_limit[expert_idx].to( + device=my_x1.device, dtype=torch.float + ) + my_x1 = torch.clamp(my_x1, min=-limit, max=limit) + my_x2 = torch.clamp(my_x2, max=limit) + if ( + args.gemm1_alpha is not None + or args.gemm1_beta is not None + or args.gemm1_clamp_limit is not None + ): + assert int(args.activation_type) == int(ActivationType.Swiglu) + alpha = ( + 1.0 + if args.gemm1_alpha is None + else args.gemm1_alpha[expert_idx].to( + device=my_x2.device, dtype=torch.float + ) + ) + beta = ( + 0.0 + if args.gemm1_beta is None + else args.gemm1_beta[expert_idx].to( + device=my_x1.device, dtype=torch.float + ) + ) + activation_output[i : i + my_num_tokens] = ( + my_x2 * torch.sigmoid(alpha * my_x2) * (my_x1 + beta) + ) + else: + activation_output[i : i + my_num_tokens] = ( + activation_func(my_x2) * my_x1 + ) else: my_x1 = my_a[:, : args.intermediate_size] activation_output[i : i + my_num_tokens] = activation_func(my_x1) @@ -2593,6 +2646,9 @@ def run_moe_reference_mxfp8(args): gemm1_bias=args.gemm1_bias, gemm2_bias=args.gemm2_bias, gemm1_lora_delta=args.gemm1_lora_delta, + gemm1_alpha=args.gemm1_alpha, + gemm1_beta=args.gemm1_beta, + gemm1_clamp_limit=args.gemm1_clamp_limit, ) return run_moe_dequant(args_dequant, QuantMode.FP8_BLOCK_SCALE_MXFP8), args_dequant @@ -2735,6 +2791,9 @@ def run_moe_reference_bf16(args): gemm1_bias=args.gemm1_bias, gemm2_bias=args.gemm2_bias, gemm1_lora_delta=args.gemm1_lora_delta, + gemm1_alpha=args.gemm1_alpha, + gemm1_beta=args.gemm1_beta, + gemm1_clamp_limit=args.gemm1_clamp_limit, ) return run_moe_dequant(args_dequant, QuantMode.BF16), args_dequant @@ -2829,6 +2888,9 @@ def _compute_moe_actual_unified(moe_impl, args_dequant, args, **kwargs): "gemm1_bias": args.gemm1_bias, "gemm2_bias": args.gemm2_bias, "gemm1_lora_delta": args.gemm1_lora_delta, + "gemm1_alpha": args.gemm1_alpha, + "gemm1_beta": args.gemm1_beta, + "gemm1_clamp_limit": args.gemm1_clamp_limit, "permute_info": args.permute_info, "norm_topk_prob": kwargs.get("norm_topk_prob", True), } @@ -2862,6 +2924,9 @@ def run_moe_test( gemm1_bias=None, gemm2_bias=None, gemm1_lora_delta=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, routing_bias_dtype=None, norm_topk_prob=True, ): @@ -3032,6 +3097,9 @@ def run_moe_test( gemm1_bias=gemm1_bias, gemm2_bias=gemm2_bias, gemm1_lora_delta=gemm1_lora_delta, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, ) # Compute reference output @@ -4700,6 +4768,187 @@ def test_routing_dtype_flexibility( ) +def test_bf16_moe_swiglu_oa_activation_param_validation(): + """BF16 SwiGLU OA params are rejected before dispatch for non-SwiGLU activations.""" + kwargs = { + "routing_logits": torch.empty((1, 1), dtype=torch.bfloat16), + "routing_bias": None, + "hidden_states": torch.empty((1, 128), dtype=torch.bfloat16), + "gemm1_weights": torch.empty((1, 2, 128), dtype=torch.bfloat16), + "gemm2_weights": torch.empty((1, 128, 1), dtype=torch.bfloat16), + "num_experts": 1, + "top_k": 1, + "n_group": None, + "topk_group": None, + "intermediate_size": 1, + "local_expert_offset": 0, + "local_num_experts": 1, + "routed_scaling_factor": None, + "routing_method_type": RoutingMethodType.Renormalize.value, + "activation_type": ActivationType.Geglu.value, + } + per_expert = torch.ones((1,), dtype=torch.float32) + + with pytest.raises(ValueError, match=r"ActivationType\.Swiglu"): + trtllm_bf16_moe(**kwargs, gemm1_alpha=per_expert) + + routed_kwargs = { + key: value + for key, value in kwargs.items() + if key not in ("routing_logits", "routing_bias") + } + routed_kwargs["topk_ids"] = torch.empty((1, 1), dtype=torch.int32) + + with pytest.raises(ValueError, match=r"ActivationType\.Swiglu"): + trtllm_bf16_routed_moe(**routed_kwargs, gemm1_clamp_limit=per_expert) + + +def test_bf16_moe_swiglu_oa_activation_params(cache_permute_indices): + """TRT-LLM Gen BF16 MoE applies raw fused FC1 SwiGLU OA params.""" + if not torch.cuda.is_available(): + pytest.skip("TRT-LLM Gen BF16 MoE test requires CUDA.") + compute_capability = get_compute_capability(torch.device(device="cuda")) + if compute_capability[0] not in [10, 12]: + pytest.skip("TRT-LLM Gen BF16 MoE requires SM10.x or SM12.x.") + + num_experts = 64 + num_tokens = 8 + hidden_size = 512 + intermediate_size = 512 + top_k = 1 + padding = 8 + routing_method_type = RoutingMethodType.Renormalize + weight_processing = { + "use_shuffled_weight": True, + "layout": WeightLayout.BlockMajorK, + } + + selected_experts = torch.arange(num_tokens, device="cuda", dtype=torch.long) + routing_logits = torch.full( + (num_tokens, num_experts), -80.0, device="cuda", dtype=torch.bfloat16 + ) + routing_logits[torch.arange(num_tokens, device="cuda"), selected_experts] = 80.0 + permute_info, scores = routing_reference_renormalize( + routing_logits, top_k, num_experts, padding + ) + + hidden_states = torch.zeros( + (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16 + ) + hidden_states[:, 0] = torch.tensor( + [-3.0, -1.0, 0.25, 1.0, 3.0, -4.0, 2.0, 0.5], + device="cuda", + dtype=torch.bfloat16, + ) + hidden_states[:, 1] = torch.tensor( + [-3.0, -0.5, 0.75, 2.5, 4.0, 6.0, -8.0, 1.5], + device="cuda", + dtype=torch.bfloat16, + ) + + gemm1_weights = torch.zeros( + (num_experts, 2 * intermediate_size, hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + gemm2_weights = torch.zeros( + (num_experts, hidden_size, intermediate_size), + device="cuda", + dtype=torch.bfloat16, + ) + for expert_idx in range(num_tokens): + gemm1_weights[expert_idx, 0, 0] = 1.0 + gemm1_weights[expert_idx, intermediate_size, 1] = 1.0 + gemm2_weights[expert_idx, 0, 0] = 1.0 + + moe_impl = BF16Moe() + moe_impl._cache_permute_indices = cache_permute_indices + + def per_expert(value): + return torch.full((num_experts,), value, device="cuda", dtype=torch.float32) + + def run_case(gemm1_alpha=None, gemm1_beta=None, gemm1_clamp_limit=None): + weights_data = moe_impl.quantize_weights( + gemm1_weights, gemm2_weights, hidden_states + ) + inputs_data = moe_impl.quantize_inputs( + hidden_states, weights_data["hidden_states_scale_global"] + ) + quant_data = {**weights_data, **inputs_data} + args = moe_args( + num_tokens, + num_experts, + hidden_size, + intermediate_size, + top_k, + padding, + quant_data["hidden_states"], + quant_data["hidden_states_scale"], + quant_data["hidden_states_scale_global"], + scores, + quant_data["gemm1_weights"], + quant_data["gemm1_scales"], + quant_data["gemm1_scales_global"], + quant_data["gemm2_weights"], + quant_data["gemm2_scales"], + quant_data["gemm2_scales_global"], + permute_info, + False, + ActivationType.Swiglu, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + output_ref, args_dequant = moe_impl.compute_reference(args) + output_actual = moe_impl.compute_production( + args_dequant, + args, + expert_logits=routing_logits, + routing_bias=None, + hidden_states_orig=hidden_states, + gemm1_weights_orig=gemm1_weights, + gemm2_weights_orig=gemm2_weights, + n_groups=None, + top_k_groups=None, + routed_scaling=None, + routing_method_type=routing_method_type, + weight_processing=weight_processing, + enable_pdl=True, + hidden_states_quant=inputs_data["hidden_states"], + enable_autotune=False, + norm_topk_prob=True, + ) + return output_ref.to(torch.float), output_actual.to(torch.float) + + output_default_ref, output_default = run_case() + output_noop_ref, output_noop = run_case( + gemm1_alpha=per_expert(1.0), + gemm1_beta=per_expert(0.0), + gemm1_clamp_limit=per_expert(1.0e9), + ) + + alpha = per_expert(1.702) + beta = per_expert(1.0) + clamp_limit = per_expert(2.0) + output_oa_ref, output_oa = run_case( + gemm1_alpha=alpha, + gemm1_clamp_limit=clamp_limit, + ) + output_beta_oa_ref, output_beta_oa = run_case( + gemm1_alpha=alpha, + gemm1_beta=beta, + gemm1_clamp_limit=clamp_limit, + ) + + torch.testing.assert_close(output_default, output_default_ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(output_noop, output_noop_ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(output_oa, output_oa_ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(output_beta_oa, output_beta_oa_ref, atol=1e-2, rtol=1e-2) + assert torch.allclose(output_default, output_noop, atol=1e-2, rtol=1e-2) + assert not torch.allclose(output_default, output_oa, atol=1e-2, rtol=1e-2) + assert not torch.allclose(output_oa, output_beta_oa, atol=1e-2, rtol=1e-2) + + def test_fp8_block_scale_routed_activation_type_relu2_smoke(): """Smoke test routed FP8 block-scale call path with explicit non-gated activation_type.""" compute_capability = get_compute_capability(torch.device(device="cuda")) diff --git a/tests/trace/test_trtllm_bf16_moe_trace.py b/tests/trace/test_trtllm_bf16_moe_trace.py new file mode 100644 index 00000000000..2f9084aada6 --- /dev/null +++ b/tests/trace/test_trtllm_bf16_moe_trace.py @@ -0,0 +1,150 @@ +"""Trace tests for TRT-LLM BF16 MoE.""" + +import torch + + +def _bf16_trace_kwargs(): + seq_len = 4 + num_experts = 2 + num_local_experts = 2 + hidden_size = 2 + intermediate_size = 1 + top_k = 1 + + return dict( + routing_logits=torch.zeros(seq_len, num_experts, dtype=torch.float32), + routing_bias=None, + hidden_states=torch.zeros(seq_len, hidden_size, dtype=torch.bfloat16), + gemm1_weights=torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size, + dtype=torch.bfloat16, + ), + gemm2_weights=torch.zeros( + num_local_experts, + hidden_size, + intermediate_size, + dtype=torch.bfloat16, + ), + top_k=top_k, + n_group=None, + topk_group=None, + local_expert_offset=0, + routed_scaling_factor=None, + routing_method_type=0, + gemm1_alpha=torch.ones(num_local_experts, dtype=torch.float32), + gemm1_beta=torch.zeros(num_local_experts, dtype=torch.float32), + gemm1_clamp_limit=torch.full((num_local_experts,), 2.0, dtype=torch.float32), + ) + + +def _bf16_routed_trace_kwargs(): + kwargs = _bf16_trace_kwargs() + kwargs.pop("routing_logits") + kwargs.pop("routing_bias") + kwargs["topk_ids"] = torch.zeros(4, 1, dtype=torch.int32) + kwargs["num_experts"] = 2 + kwargs["intermediate_size"] = 1 + return kwargs + + +def test_bf16_moe_trace_schema_includes_swiglu_oa_params(): + from flashinfer.fused_moe import trtllm_bf16_moe, trtllm_bf16_routed_moe + + trace_defs = [ + trtllm_bf16_moe.fi_trace(**_bf16_trace_kwargs()), + trtllm_bf16_routed_moe.fi_trace(**_bf16_routed_trace_kwargs()), + ] + + for defn in trace_defs: + assert defn["axes"]["num_local_experts"]["value"] == 2 + for name in ("gemm1_alpha", "gemm1_beta", "gemm1_clamp_limit"): + assert defn["inputs"][name]["shape"] == ["num_local_experts"] + assert defn["inputs"][name]["dtype"] == "float32" + assert defn["inputs"][name]["optional"] is True + assert defn["inputs"][name]["description"] + + +def test_bf16_moe_trace_reference_applies_swiglu_oa_params(): + from flashinfer.trace.templates.moe import ( + trtllm_bf16_moe_trace, + trtllm_bf16_routed_moe_trace, + ) + + routing_logits = torch.zeros(4, 1, dtype=torch.float32) + hidden_states = torch.tensor( + [[-3.0, -3.0], [-1.0, -0.5], [3.0, 4.0], [-4.0, 6.0]], + dtype=torch.bfloat16, + ) + gemm1_weights = torch.tensor( + [[[1.0, 0.0], [0.0, 1.0]]], + dtype=torch.bfloat16, + ) + gemm2_weights = torch.tensor( + [[[1.0], [0.0]]], + dtype=torch.bfloat16, + ) + + common_kwargs = dict( + routing_logits=routing_logits, + routing_bias=None, + hidden_states=hidden_states, + gemm1_weights=gemm1_weights, + gemm2_weights=gemm2_weights, + num_experts=1, + top_k=1, + local_expert_offset=0, + routed_scaling_factor=None, + ) + + default_out = trtllm_bf16_moe_trace.reference(**common_kwargs).to(torch.float32) + + large_limit = torch.full((1,), 1.0e9, dtype=torch.float32) + noop_out = trtllm_bf16_moe_trace.reference( + **common_kwargs, + gemm1_alpha=torch.ones((1,), dtype=torch.float32), + gemm1_beta=torch.zeros((1,), dtype=torch.float32), + gemm1_clamp_limit=large_limit, + ).to(torch.float32) + + clamp_limit = torch.full((1,), 2.0, dtype=torch.float32) + clamp_only_out = trtllm_bf16_moe_trace.reference( + **common_kwargs, + gemm1_clamp_limit=clamp_limit, + ).to(torch.float32) + oa_out = trtllm_bf16_moe_trace.reference( + **common_kwargs, + gemm1_alpha=torch.full((1,), 1.702, dtype=torch.float32), + gemm1_beta=torch.ones((1,), dtype=torch.float32), + gemm1_clamp_limit=clamp_limit, + ).to(torch.float32) + routed_oa_out = trtllm_bf16_routed_moe_trace.reference( + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + hidden_states=hidden_states, + gemm1_weights=gemm1_weights, + gemm2_weights=gemm2_weights, + num_experts=1, + top_k=1, + local_expert_offset=0, + routed_scaling_factor=None, + gemm1_alpha=torch.full((1,), 1.702, dtype=torch.float32), + gemm1_beta=torch.ones((1,), dtype=torch.float32), + gemm1_clamp_limit=clamp_limit, + ).to(torch.float32) + + x1 = hidden_states[:, :1].to(torch.float32).clamp(min=-2.0, max=2.0) + x2 = hidden_states[:, 1:].to(torch.float32).clamp(max=2.0) + expected_clamp_only = x2 * torch.sigmoid(x2) * x1 + expected_oa = x2 * torch.sigmoid(1.702 * x2) * (x1 + 1.0) + + torch.testing.assert_close(default_out, noop_out, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + clamp_only_out[:, :1], expected_clamp_only.to(torch.bfloat16).to(torch.float32) + ) + torch.testing.assert_close( + oa_out[:, :1], expected_oa.to(torch.bfloat16).to(torch.float32) + ) + torch.testing.assert_close(routed_oa_out, oa_out) + assert not torch.allclose(default_out, clamp_only_out, atol=1e-2, rtol=1e-2) + assert not torch.allclose(default_out, oa_out, atol=1e-2, rtol=1e-2) From 28406af5b9134757acbd6bc44647fd00261d163f Mon Sep 17 00:00:00 2001 From: yichengj Date: Wed, 10 Jun 2026 21:35:00 -0700 Subject: [PATCH 12/13] perf(gemm): update mm_fp4 b12x SM120 NVFP4 dense GEMM kernel (#3560) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## πŸ“Œ Description ### Summary This PR refreshes the `b12x` dense FP4 GEMM kernel to match the latest upstream for better performance. **No public API change.** The `auto` routing heuristic (`_heuristic_func_mm_fp4`) is **unchanged** β€” b12x is preferred only at SM120, not SM121. Addresses **#3517** *[Perf][SM12x] Update mm_fp4(backend='b12x') kernels*. ### Motivation @lukealonso has made performance improvements to b12x dense GEMMs; we'd like to port over the improvements to FlashInfer. Gains are focused on small-M. ### Benchmarks: new b12x vs old b12x Setup per SKU: `benchmarks/flashinfer_benchmark.py --testlist <104 nvfp4 shapes> --use_nvfp4 --use_128x4_sf_layout --refcheck`, 30 iters / 5 warmup, **CUPTI timing** (harness default). 104 shapes = 13 token counts m∈{1..4096} Γ— 8 (n,k) layer shapes taken from real **DeepSeek-R1** and **Llama-3** model layers, all K%128==0. The refresh is a consistent small win at small `m` (where `swap_ab` and `expected_m` help most) and break-even-to-positive overall β€” **no SKU regresses in aggregate**. Geomean speedup (old_time / new_time; >1 = new faster), grouped by `m`: | SKU | overall | m≀8 | 16≀m≀128 | mβ‰₯256 | |---|---|---|---|---| | RTX 5080 (SM120) | 1.003Γ— | 1.019Γ— | 1.003Γ— | 0.991Γ— | | RTX PRO 6000 (SM120) | **1.021Γ—** | 1.060Γ— | 1.010Γ— | 1.000Γ— | | DGX Spark / GB10 (SM121) | **1.024Γ—** | 1.021Γ— | 1.028Γ— | 1.022Γ— | The win is largest at small `m` on PRO 6000 / Spark and break-even on the 5080. A few `n=4096 k=4096/5376` shapes regress ≀5% (bolded in the per-shape tables); the PRO 6000 has a sharper `8192Γ—2560` regression at medium `m` (down to 0.873Γ—), which is SKU-specific and flagged as a follow-up. Full per-shape new-vs-old tables (all 104 shapes, sorted by `n, k, m`) per SKU:
RTX 5080 (SM120) | n | k | m | old (us) | new (us) | new TFLOPs | speedup | |---|---|---|---|---|---|---| | 512 | 7168 | 1 | 37.6 | 39.8 | 0.2 | **0.945x** | | 512 | 7168 | 2 | 38.3 | 38.8 | 0.4 | 0.986x | | 512 | 7168 | 4 | 38.8 | 38.9 | 0.8 | 0.995x | | 512 | 7168 | 8 | 38.9 | 39.2 | 1.5 | 0.994x | | 512 | 7168 | 16 | 39.4 | 40.4 | 2.9 | 0.975x | | 512 | 7168 | 32 | 40.6 | 41.2 | 5.7 | 0.985x | | 512 | 7168 | 64 | 43.8 | 43.5 | 10.8 | 1.007x | | 512 | 7168 | 128 | 55.1 | 55.6 | 16.9 | 0.992x | | 512 | 7168 | 256 | 67.4 | 67.7 | 27.8 | 0.995x | | 512 | 7168 | 512 | 84.0 | 84.7 | 44.4 | 0.991x | | 512 | 7168 | 1024 | 94.4 | 95.9 | 78.4 | 0.984x | | 512 | 7168 | 2048 | 104.4 | 106.0 | 141.9 | 0.985x | | 512 | 7168 | 4096 | 165.4 | 168.2 | 178.7 | 0.983x | | 1024 | 8192 | 1 | 58.8 | 59.1 | 0.3 | 0.995x | | 1024 | 8192 | 2 | 58.7 | 59.0 | 0.6 | 0.995x | | 1024 | 8192 | 4 | 59.7 | 59.7 | 1.1 | 0.999x | | 1024 | 8192 | 8 | 60.2 | 59.9 | 2.2 | 1.006x | | 1024 | 8192 | 16 | 61.1 | 61.1 | 4.4 | 1.001x | | 1024 | 8192 | 32 | 60.8 | 60.9 | 8.8 | 0.998x | | 1024 | 8192 | 64 | 70.1 | 69.7 | 15.4 | 1.006x | | 1024 | 8192 | 128 | 78.5 | 77.4 | 27.7 | 1.014x | | 1024 | 8192 | 256 | 85.9 | 85.8 | 50.1 | 1.002x | | 1024 | 8192 | 512 | 96.5 | 97.0 | 88.6 | 0.995x | | 1024 | 8192 | 1024 | 97.3 | 97.5 | 176.2 | 0.998x | | 1024 | 8192 | 2048 | 154.2 | 155.8 | 220.5 | 0.990x | | 1024 | 8192 | 4096 | 251.2 | 254.7 | 269.8 | 0.986x | | 2560 | 8192 | 1 | 101.1 | 87.6 | 0.5 | 1.154x | | 2560 | 8192 | 2 | 101.1 | 91.2 | 0.9 | 1.108x | | 2560 | 8192 | 4 | 101.5 | 93.7 | 1.8 | 1.083x | | 2560 | 8192 | 8 | 101.9 | 95.1 | 3.5 | 1.072x | | 2560 | 8192 | 16 | 102.3 | 100.1 | 6.7 | 1.022x | | 2560 | 8192 | 32 | 101.6 | 95.3 | 14.1 | 1.067x | | 2560 | 8192 | 64 | 101.2 | 95.6 | 28.1 | 1.058x | | 2560 | 8192 | 128 | 95.3 | 94.4 | 56.9 | 1.009x | | 2560 | 8192 | 256 | 109.3 | 106.8 | 100.5 | 1.023x | | 2560 | 8192 | 512 | 113.9 | 114.1 | 188.2 | 0.998x | | 2560 | 8192 | 1024 | 157.3 | 158.3 | 271.3 | 0.994x | | 2560 | 8192 | 2048 | 252.6 | 254.9 | 337.0 | 0.991x | | 2560 | 8192 | 4096 | 433.5 | 436.9 | 393.2 | 0.992x | | 4096 | 4096 | 1 | 57.2 | 57.6 | 0.6 | 0.993x | | 4096 | 4096 | 2 | 57.3 | 56.4 | 1.2 | 1.016x | | 4096 | 4096 | 4 | 58.1 | 59.3 | 2.3 | 0.981x | | 4096 | 4096 | 8 | 56.5 | 57.5 | 4.7 | 0.983x | | 4096 | 4096 | 16 | 59.6 | 60.2 | 8.9 | 0.990x | | 4096 | 4096 | 32 | 60.3 | 60.4 | 17.8 | 1.000x | | 4096 | 4096 | 64 | 57.1 | 59.2 | 36.3 | **0.964x** | | 4096 | 4096 | 128 | 69.9 | 72.8 | 59.0 | **0.960x** | | 4096 | 4096 | 256 | 61.2 | 65.6 | 130.8 | **0.933x** | | 4096 | 4096 | 512 | 112.8 | 113.9 | 150.8 | 0.991x | | 4096 | 4096 | 1024 | 151.8 | 152.3 | 225.7 | 0.997x | | 4096 | 4096 | 2048 | 220.7 | 223.1 | 308.0 | 0.989x | | 4096 | 4096 | 4096 | 366.1 | 376.4 | 365.2 | 0.973x | | 4096 | 5376 | 1 | 72.6 | 71.3 | 0.6 | 1.017x | | 4096 | 5376 | 2 | 72.5 | 70.5 | 1.2 | 1.028x | | 4096 | 5376 | 4 | 76.2 | 75.7 | 2.3 | 1.006x | | 4096 | 5376 | 8 | 75.7 | 75.2 | 4.7 | 1.007x | | 4096 | 5376 | 16 | 76.1 | 78.2 | 9.0 | 0.973x | | 4096 | 5376 | 32 | 83.4 | 85.8 | 16.4 | 0.971x | | 4096 | 5376 | 64 | 77.2 | 81.4 | 34.6 | **0.948x** | | 4096 | 5376 | 128 | 84.4 | 88.6 | 63.6 | **0.953x** | | 4096 | 5376 | 256 | 84.6 | 86.6 | 130.2 | 0.978x | | 4096 | 5376 | 512 | 133.9 | 134.5 | 167.6 | 0.995x | | 4096 | 5376 | 1024 | 178.8 | 180.3 | 250.1 | 0.991x | | 4096 | 5376 | 2048 | 262.6 | 265.2 | 340.2 | 0.990x | | 4096 | 5376 | 4096 | 464.7 | 470.1 | 383.7 | 0.988x | | 4096 | 14336 | 1 | 176.1 | 174.6 | 0.7 | 1.009x | | 4096 | 14336 | 2 | 176.0 | 175.2 | 1.3 | 1.005x | | 4096 | 14336 | 4 | 176.9 | 176.7 | 2.7 | 1.001x | | 4096 | 14336 | 8 | 177.6 | 177.5 | 5.3 | 1.000x | | 4096 | 14336 | 16 | 178.5 | 178.6 | 10.5 | 1.000x | | 4096 | 14336 | 32 | 178.6 | 178.9 | 21.0 | 0.999x | | 4096 | 14336 | 64 | 180.0 | 180.3 | 41.7 | 0.998x | | 4096 | 14336 | 128 | 182.3 | 182.7 | 82.3 | 0.998x | | 4096 | 14336 | 256 | 194.4 | 195.0 | 154.1 | 0.997x | | 4096 | 14336 | 512 | 253.2 | 254.7 | 236.1 | 0.994x | | 4096 | 14336 | 1024 | 375.2 | 377.6 | 318.5 | 0.994x | | 4096 | 14336 | 2048 | 580.2 | 585.3 | 410.9 | 0.991x | | 4096 | 14336 | 4096 | 1049.4 | 1061.3 | 453.3 | 0.989x | | 8192 | 2560 | 1 | 70.1 | 68.4 | 0.6 | 1.025x | | 8192 | 2560 | 2 | 66.1 | 65.4 | 1.3 | 1.010x | | 8192 | 2560 | 4 | 65.3 | 65.2 | 2.6 | 1.001x | | 8192 | 2560 | 8 | 67.5 | 67.0 | 5.0 | 1.008x | | 8192 | 2560 | 16 | 67.9 | 65.2 | 10.3 | 1.042x | | 8192 | 2560 | 32 | 68.2 | 66.1 | 20.3 | 1.032x | | 8192 | 2560 | 64 | 67.1 | 67.8 | 39.6 | 0.989x | | 8192 | 2560 | 128 | 69.3 | 69.6 | 77.2 | 0.996x | | 8192 | 2560 | 256 | 89.8 | 91.6 | 117.3 | 0.981x | | 8192 | 2560 | 512 | 137.4 | 137.1 | 156.6 | 1.002x | | 8192 | 2560 | 1024 | 176.3 | 177.1 | 242.5 | 0.995x | | 8192 | 2560 | 2048 | 279.4 | 280.6 | 306.1 | 0.996x | | 8192 | 2560 | 4096 | 489.2 | 495.3 | 346.9 | 0.988x | | 14336 | 4096 | 1 | 171.5 | 162.8 | 0.7 | 1.053x | | 14336 | 4096 | 2 | 171.7 | 162.1 | 1.4 | 1.059x | | 14336 | 4096 | 4 | 171.5 | 162.0 | 2.9 | 1.059x | | 14336 | 4096 | 8 | 171.9 | 163.3 | 5.8 | 1.053x | | 14336 | 4096 | 16 | 171.8 | 162.9 | 11.5 | 1.054x | | 14336 | 4096 | 32 | 173.4 | 165.2 | 22.7 | 1.049x | | 14336 | 4096 | 64 | 176.7 | 167.4 | 44.9 | 1.055x | | 14336 | 4096 | 128 | 181.3 | 180.4 | 83.3 | 1.005x | | 14336 | 4096 | 256 | 188.7 | 190.3 | 158.0 | 0.992x | | 14336 | 4096 | 512 | 253.7 | 255.0 | 235.8 | 0.995x | | 14336 | 4096 | 1024 | 382.4 | 384.3 | 312.9 | 0.995x | | 14336 | 4096 | 2048 | 620.9 | 624.4 | 385.2 | 0.994x | | 14336 | 4096 | 4096 | 1140.7 | 1148.9 | 418.7 | 0.993x |
RTX PRO 6000 (SM120) | n | k | m | old (us) | new (us) | new TFLOPs | speedup | |---|---|---|---|---|---|---| | 512 | 7168 | 1 | 22.3 | 21.7 | 0.3 | 1.026x | | 512 | 7168 | 2 | 22.4 | 22.5 | 0.7 | 0.994x | | 512 | 7168 | 4 | 22.4 | 22.5 | 1.3 | 0.998x | | 512 | 7168 | 8 | 22.4 | 22.5 | 2.6 | 0.995x | | 512 | 7168 | 16 | 22.4 | 22.4 | 5.2 | 0.999x | | 512 | 7168 | 32 | 22.3 | 22.5 | 10.4 | 0.991x | | 512 | 7168 | 64 | 22.4 | 22.7 | 20.7 | 0.991x | | 512 | 7168 | 128 | 22.9 | 23.1 | 40.6 | 0.992x | | 512 | 7168 | 256 | 32.6 | 32.8 | 57.2 | 0.992x | | 512 | 7168 | 512 | 35.7 | 35.9 | 104.8 | 0.996x | | 512 | 7168 | 1024 | 41.5 | 41.7 | 180.1 | 0.995x | | 512 | 7168 | 2048 | 45.6 | 47.1 | 319.2 | **0.969x** | | 512 | 7168 | 4096 | 51.4 | 53.6 | 561.2 | **0.959x** | | 1024 | 8192 | 1 | 25.6 | 25.8 | 0.7 | 0.993x | | 1024 | 8192 | 2 | 25.6 | 25.8 | 1.3 | 0.994x | | 1024 | 8192 | 4 | 25.7 | 25.9 | 2.6 | 0.994x | | 1024 | 8192 | 8 | 25.7 | 25.9 | 5.2 | 0.992x | | 1024 | 8192 | 16 | 25.7 | 25.8 | 10.4 | 0.994x | | 1024 | 8192 | 32 | 25.6 | 25.8 | 20.8 | 0.993x | | 1024 | 8192 | 64 | 25.8 | 25.9 | 41.4 | 0.995x | | 1024 | 8192 | 128 | 27.0 | 28.4 | 75.6 | **0.951x** | | 1024 | 8192 | 256 | 39.0 | 39.6 | 108.4 | 0.985x | | 1024 | 8192 | 512 | 43.1 | 43.5 | 197.6 | 0.991x | | 1024 | 8192 | 1024 | 47.4 | 49.5 | 347.4 | **0.959x** | | 1024 | 8192 | 2048 | 50.2 | 50.0 | 687.0 | 1.004x | | 1024 | 8192 | 4096 | 89.5 | 90.9 | 756.3 | 0.985x | | 2560 | 8192 | 1 | 34.5 | 31.6 | 1.3 | 1.092x | | 2560 | 8192 | 2 | 34.6 | 31.7 | 2.6 | 1.092x | | 2560 | 8192 | 4 | 34.7 | 31.8 | 5.3 | 1.090x | | 2560 | 8192 | 8 | 34.6 | 32.1 | 10.5 | 1.079x | | 2560 | 8192 | 16 | 34.9 | 32.7 | 20.5 | 1.069x | | 2560 | 8192 | 32 | 34.9 | 32.8 | 40.9 | 1.065x | | 2560 | 8192 | 64 | 35.2 | 33.3 | 80.7 | 1.059x | | 2560 | 8192 | 128 | 41.7 | 39.9 | 134.7 | 1.046x | | 2560 | 8192 | 256 | 48.2 | 39.9 | 269.1 | 1.207x | | 2560 | 8192 | 512 | 48.8 | 48.6 | 442.2 | 1.004x | | 2560 | 8192 | 1024 | 50.6 | 51.3 | 837.6 | 0.987x | | 2560 | 8192 | 2048 | 87.5 | 88.2 | 973.8 | 0.992x | | 2560 | 8192 | 4096 | 163.8 | 166.9 | 1029.4 | 0.982x | | 4096 | 4096 | 1 | 23.6 | 22.4 | 1.5 | 1.053x | | 4096 | 4096 | 2 | 23.6 | 22.7 | 3.0 | 1.038x | | 4096 | 4096 | 4 | 23.7 | 22.8 | 5.9 | 1.041x | | 4096 | 4096 | 8 | 23.7 | 23.2 | 11.6 | 1.021x | | 4096 | 4096 | 16 | 23.8 | 23.2 | 23.2 | 1.029x | | 4096 | 4096 | 32 | 24.0 | 23.6 | 45.5 | 1.020x | | 4096 | 4096 | 64 | 24.0 | 23.8 | 90.2 | 1.009x | | 4096 | 4096 | 128 | 31.9 | 31.2 | 137.7 | 1.022x | | 4096 | 4096 | 256 | 35.2 | 35.2 | 244.0 | 1.000x | | 4096 | 4096 | 512 | 36.8 | 36.7 | 467.7 | 1.001x | | 4096 | 4096 | 1024 | 54.3 | 54.5 | 630.0 | 0.995x | | 4096 | 4096 | 2048 | 73.0 | 73.9 | 930.3 | 0.988x | | 4096 | 4096 | 4096 | 134.6 | 137.0 | 1003.4 | 0.983x | | 4096 | 5376 | 1 | 30.9 | 28.9 | 1.5 | 1.068x | | 4096 | 5376 | 2 | 30.8 | 28.8 | 3.1 | 1.069x | | 4096 | 5376 | 4 | 30.8 | 29.4 | 6.0 | 1.046x | | 4096 | 5376 | 8 | 30.8 | 29.2 | 12.1 | 1.055x | | 4096 | 5376 | 16 | 30.7 | 29.5 | 23.9 | 1.042x | | 4096 | 5376 | 32 | 31.0 | 29.9 | 47.1 | 1.035x | | 4096 | 5376 | 64 | 31.1 | 30.0 | 93.9 | 1.037x | | 4096 | 5376 | 128 | 37.6 | 37.0 | 152.5 | 1.018x | | 4096 | 5376 | 256 | 41.4 | 38.3 | 294.7 | 1.083x | | 4096 | 5376 | 512 | 43.2 | 43.2 | 522.2 | 0.999x | | 4096 | 5376 | 1024 | 66.2 | 66.3 | 680.7 | 0.999x | | 4096 | 5376 | 2048 | 92.4 | 92.4 | 975.8 | 0.999x | | 4096 | 5376 | 4096 | 166.6 | 169.3 | 1065.2 | 0.984x | | 4096 | 14336 | 1 | 64.4 | 60.0 | 2.0 | 1.072x | | 4096 | 14336 | 2 | 64.3 | 59.8 | 3.9 | 1.075x | | 4096 | 14336 | 4 | 64.4 | 60.3 | 7.8 | 1.067x | | 4096 | 14336 | 8 | 64.4 | 60.2 | 15.6 | 1.070x | | 4096 | 14336 | 16 | 64.3 | 60.4 | 31.1 | 1.064x | | 4096 | 14336 | 32 | 64.4 | 60.9 | 61.7 | 1.059x | | 4096 | 14336 | 64 | 64.5 | 61.3 | 122.6 | 1.053x | | 4096 | 14336 | 128 | 65.3 | 62.8 | 239.2 | 1.040x | | 4096 | 14336 | 256 | 77.5 | 65.2 | 460.9 | 1.188x | | 4096 | 14336 | 512 | 80.7 | 82.3 | 730.7 | 0.981x | | 4096 | 14336 | 1024 | 139.4 | 140.6 | 855.5 | 0.992x | | 4096 | 14336 | 2048 | 206.0 | 206.8 | 1163.0 | 0.996x | | 4096 | 14336 | 4096 | 395.7 | 398.7 | 1206.5 | 0.992x | | 8192 | 2560 | 1 | 22.2 | 22.0 | 1.9 | 1.012x | | 8192 | 2560 | 2 | 21.5 | 20.7 | 4.1 | 1.043x | | 8192 | 2560 | 4 | 21.4 | 20.4 | 8.2 | 1.049x | | 8192 | 2560 | 8 | 21.9 | 21.7 | 15.5 | 1.008x | | 8192 | 2560 | 16 | 21.2 | 21.7 | 30.9 | 0.977x | | 8192 | 2560 | 32 | 21.6 | 23.6 | 56.9 | **0.915x** | | 8192 | 2560 | 64 | 24.6 | 26.8 | 100.2 | **0.917x** | | 8192 | 2560 | 128 | 25.8 | 29.7 | 180.9 | **0.868x** | | 8192 | 2560 | 256 | 29.1 | 29.6 | 362.2 | 0.981x | | 8192 | 2560 | 512 | 43.0 | 42.8 | 501.7 | 1.005x | | 8192 | 2560 | 1024 | 54.7 | 55.2 | 777.4 | 0.989x | | 8192 | 2560 | 2048 | 93.2 | 93.7 | 916.9 | 0.994x | | 8192 | 2560 | 4096 | 163.2 | 167.5 | 1025.7 | 0.974x | | 14336 | 4096 | 1 | 49.1 | 40.3 | 2.9 | 1.218x | | 14336 | 4096 | 2 | 48.3 | 39.1 | 6.0 | 1.235x | | 14336 | 4096 | 4 | 49.2 | 39.9 | 11.8 | 1.231x | | 14336 | 4096 | 8 | 47.8 | 39.5 | 23.8 | 1.211x | | 14336 | 4096 | 16 | 46.0 | 39.7 | 47.3 | 1.158x | | 14336 | 4096 | 32 | 44.5 | 43.1 | 87.3 | 1.033x | | 14336 | 4096 | 64 | 40.7 | 44.2 | 170.1 | **0.922x** | | 14336 | 4096 | 128 | 43.4 | 45.3 | 332.0 | **0.959x** | | 14336 | 4096 | 256 | 59.5 | 60.7 | 495.0 | 0.979x | | 14336 | 4096 | 512 | 80.0 | 79.7 | 754.6 | 1.004x | | 14336 | 4096 | 1024 | 116.2 | 116.5 | 1032.0 | 0.997x | | 14336 | 4096 | 2048 | 212.8 | 214.3 | 1122.6 | 0.993x | | 14336 | 4096 | 4096 | 390.8 | 395.2 | 1217.2 | 0.989x |
DGX Spark / GB10 (SM121) | n | k | m | old (us) | new (us) | new TFLOPs | speedup | |---|---|---|---|---|---|---| | 512 | 7168 | 1 | 31.1 | 31.3 | 0.2 | 0.991x | | 512 | 7168 | 2 | 31.5 | 31.4 | 0.5 | 1.003x | | 512 | 7168 | 4 | 31.6 | 31.1 | 0.9 | 1.016x | | 512 | 7168 | 8 | 31.6 | 31.5 | 1.9 | 1.005x | | 512 | 7168 | 16 | 31.7 | 31.8 | 3.7 | 0.997x | | 512 | 7168 | 32 | 32.5 | 33.7 | 7.0 | **0.964x** | | 512 | 7168 | 64 | 34.1 | 33.8 | 13.9 | 1.009x | | 512 | 7168 | 128 | 49.0 | 39.3 | 23.9 | 1.245x | | 512 | 7168 | 256 | 60.7 | 49.5 | 38.0 | 1.227x | | 512 | 7168 | 512 | 56.8 | 56.3 | 66.7 | 1.008x | | 512 | 7168 | 1024 | 73.1 | 73.6 | 102.2 | 0.994x | | 512 | 7168 | 2048 | 129.0 | 129.0 | 116.5 | 1.000x | | 512 | 7168 | 4096 | 227.0 | 225.8 | 133.1 | 1.005x | | 1024 | 8192 | 1 | 65.4 | 59.8 | 0.3 | 1.094x | | 1024 | 8192 | 2 | 64.7 | 60.2 | 0.6 | 1.075x | | 1024 | 8192 | 4 | 65.9 | 59.5 | 1.1 | 1.108x | | 1024 | 8192 | 8 | 64.1 | 60.5 | 2.2 | 1.058x | | 1024 | 8192 | 16 | 63.6 | 57.1 | 4.7 | 1.114x | | 1024 | 8192 | 32 | 66.5 | 58.5 | 9.2 | 1.136x | | 1024 | 8192 | 64 | 55.5 | 58.1 | 18.5 | **0.956x** | | 1024 | 8192 | 128 | 72.0 | 64.4 | 33.3 | 1.119x | | 1024 | 8192 | 256 | 83.1 | 79.1 | 54.3 | 1.050x | | 1024 | 8192 | 512 | 81.5 | 84.5 | 101.6 | **0.964x** | | 1024 | 8192 | 1024 | 155.6 | 130.2 | 132.0 | 1.195x | | 1024 | 8192 | 2048 | 200.4 | 191.2 | 179.7 | 1.048x | | 1024 | 8192 | 4096 | 456.5 | 410.5 | 167.4 | 1.112x | | 2560 | 8192 | 1 | 117.8 | 118.0 | 0.4 | 0.998x | | 2560 | 8192 | 2 | 116.2 | 119.2 | 0.7 | 0.975x | | 2560 | 8192 | 4 | 118.0 | 118.2 | 1.4 | 0.998x | | 2560 | 8192 | 8 | 116.2 | 119.4 | 2.8 | 0.973x | | 2560 | 8192 | 16 | 118.2 | 118.5 | 5.7 | 0.998x | | 2560 | 8192 | 32 | 117.1 | 119.8 | 11.2 | 0.977x | | 2560 | 8192 | 64 | 117.7 | 120.7 | 22.2 | 0.975x | | 2560 | 8192 | 128 | 119.6 | 124.2 | 43.2 | **0.963x** | | 2560 | 8192 | 256 | 124.7 | 128.5 | 83.6 | 0.970x | | 2560 | 8192 | 512 | 151.2 | 146.9 | 146.1 | 1.029x | | 2560 | 8192 | 1024 | 212.2 | 206.8 | 207.6 | 1.026x | | 2560 | 8192 | 2048 | 346.6 | 337.4 | 254.6 | 1.027x | | 2560 | 8192 | 4096 | 951.6 | 861.4 | 199.4 | 1.105x | | 4096 | 4096 | 1 | 94.2 | 98.6 | 0.3 | **0.956x** | | 4096 | 4096 | 2 | 113.4 | 99.0 | 0.7 | 1.146x | | 4096 | 4096 | 4 | 111.6 | 101.8 | 1.3 | 1.096x | | 4096 | 4096 | 8 | 113.1 | 97.4 | 2.8 | 1.161x | | 4096 | 4096 | 16 | 113.5 | 101.6 | 5.3 | 1.117x | | 4096 | 4096 | 32 | 113.6 | 99.6 | 10.8 | 1.141x | | 4096 | 4096 | 64 | 112.5 | 100.6 | 21.3 | 1.119x | | 4096 | 4096 | 128 | 111.9 | 99.2 | 43.3 | 1.128x | | 4096 | 4096 | 256 | 107.3 | 108.4 | 79.3 | 0.991x | | 4096 | 4096 | 512 | 115.5 | 112.7 | 152.5 | 1.025x | | 4096 | 4096 | 1024 | 158.6 | 158.4 | 216.9 | 1.001x | | 4096 | 4096 | 2048 | 281.3 | 272.3 | 252.3 | 1.033x | | 4096 | 4096 | 4096 | 508.9 | 494.9 | 277.7 | 1.028x | | 4096 | 5376 | 1 | 117.4 | 118.8 | 0.4 | 0.988x | | 4096 | 5376 | 2 | 116.6 | 117.4 | 0.8 | 0.994x | | 4096 | 5376 | 4 | 117.0 | 118.4 | 1.5 | 0.989x | | 4096 | 5376 | 8 | 116.3 | 117.7 | 3.0 | 0.988x | | 4096 | 5376 | 16 | 118.6 | 118.8 | 5.9 | 0.998x | | 4096 | 5376 | 32 | 116.5 | 118.3 | 11.9 | 0.985x | | 4096 | 5376 | 64 | 117.8 | 119.1 | 23.7 | 0.989x | | 4096 | 5376 | 128 | 118.7 | 119.9 | 47.0 | 0.991x | | 4096 | 5376 | 256 | 129.1 | 132.6 | 85.0 | 0.973x | | 4096 | 5376 | 512 | 139.8 | 142.2 | 158.6 | 0.983x | | 4096 | 5376 | 1024 | 200.0 | 197.5 | 228.3 | 1.012x | | 4096 | 5376 | 2048 | 338.3 | 340.3 | 265.1 | 0.994x | | 4096 | 5376 | 4096 | 693.6 | 645.1 | 279.6 | 1.075x | | 4096 | 14336 | 1 | 231.6 | 231.8 | 0.5 | 0.999x | | 4096 | 14336 | 2 | 229.6 | 229.4 | 1.0 | 1.001x | | 4096 | 14336 | 4 | 232.6 | 232.5 | 2.0 | 1.000x | | 4096 | 14336 | 8 | 232.1 | 228.7 | 4.1 | 1.015x | | 4096 | 14336 | 16 | 230.7 | 232.7 | 8.1 | 0.991x | | 4096 | 14336 | 32 | 230.2 | 228.7 | 16.4 | 1.007x | | 4096 | 14336 | 64 | 231.9 | 232.7 | 32.3 | 0.997x | | 4096 | 14336 | 128 | 234.1 | 234.2 | 64.2 | 1.000x | | 4096 | 14336 | 256 | 266.6 | 265.2 | 113.4 | 1.005x | | 4096 | 14336 | 512 | 303.0 | 302.6 | 198.7 | 1.001x | | 4096 | 14336 | 1024 | 469.2 | 466.4 | 257.9 | 1.006x | | 4096 | 14336 | 2048 | 902.7 | 880.0 | 273.3 | 1.026x | | 4096 | 14336 | 4096 | 3847.1 | 3765.8 | 127.7 | 1.022x | | 8192 | 2560 | 1 | 113.6 | 113.0 | 0.4 | 1.005x | | 8192 | 2560 | 2 | 111.7 | 113.4 | 0.7 | 0.986x | | 8192 | 2560 | 4 | 113.2 | 113.3 | 1.5 | 0.999x | | 8192 | 2560 | 8 | 113.2 | 113.8 | 2.9 | 0.995x | | 8192 | 2560 | 16 | 114.2 | 113.5 | 5.9 | 1.006x | | 8192 | 2560 | 32 | 113.6 | 114.5 | 11.7 | 0.993x | | 8192 | 2560 | 64 | 114.2 | 115.4 | 23.3 | 0.990x | | 8192 | 2560 | 128 | 116.4 | 119.2 | 45.0 | 0.976x | | 8192 | 2560 | 256 | 119.7 | 121.7 | 88.2 | 0.984x | | 8192 | 2560 | 512 | 135.1 | 139.3 | 154.2 | 0.970x | | 8192 | 2560 | 1024 | 200.1 | 199.5 | 215.3 | 1.003x | | 8192 | 2560 | 2048 | 370.9 | 371.6 | 231.1 | 0.998x | | 8192 | 2560 | 4096 | 644.3 | 641.7 | 267.7 | 1.004x | | 14336 | 4096 | 1 | 232.3 | 226.1 | 0.5 | 1.028x | | 14336 | 4096 | 2 | 232.5 | 226.5 | 1.0 | 1.026x | | 14336 | 4096 | 4 | 231.8 | 225.8 | 2.1 | 1.026x | | 14336 | 4096 | 8 | 232.2 | 226.7 | 4.1 | 1.024x | | 14336 | 4096 | 16 | 232.7 | 226.7 | 8.3 | 1.026x | | 14336 | 4096 | 32 | 232.1 | 226.4 | 16.6 | 1.025x | | 14336 | 4096 | 64 | 233.3 | 229.0 | 32.8 | 1.019x | | 14336 | 4096 | 128 | 235.8 | 232.5 | 64.7 | 1.014x | | 14336 | 4096 | 256 | 246.5 | 239.6 | 125.5 | 1.029x | | 14336 | 4096 | 512 | 293.4 | 287.1 | 209.4 | 1.022x | | 14336 | 4096 | 1024 | 441.3 | 441.2 | 272.6 | 1.000x | | 14336 | 4096 | 2048 | 838.2 | 836.7 | 287.5 | 1.002x | | 14336 | 4096 | 4096 | 1500.5 | 1506.2 | 319.4 | 0.996x |
## πŸ” Related Issues Closes #3517. ## πŸš€ Pull Request Checklist ### βœ… Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## πŸ§ͺ Tests - [x] Tests have been added or updated as needed. *(covered by existing `tests/gemm/test_mm_fp4.py -k b12x`)* - [x] All tests are passing β€” `pytest -k b12x` 396/0; 104-shape `--refcheck` 104/104 on 5080 / PRO 6000 / GB10. ## Reviewer Notes - The vendored kernel is **not a verbatim upstream copy**: it carries backported scale-factor helpers so it builds on public `cutlass-dsl 4.5.2` (pure upstream b12x needs an internal/newer cutlass-dsl), and the MXFP8 path is pruned (unsupported there). - The SFB-layout N-divisibility assertion in `cute_dsl/utils.py` is relaxed **64β†’16** (upstream commit `0daa6ab`) to admit the narrow-N `swap_ab` tiles. This is safe: it only *loosens* a precondition (every previously-valid N%64==0 tile still passes), the SFB smem layout already rounds N up to a full 128-wide block so the layout is unchanged for existing tiles, and the newly-admitted narrow tiles pass `--refcheck`. - `swap_ab` here is b12x's **device-internal** transpose (public C stays row-major) β€” *not* the SM100 operand-swap/FFI-transpose convention. The shared `_compile_block_scaled_gemm` harness used by Sm100/Sm103 is **left untouched**. ## Summary by CodeRabbit * **Improvements** * Relaxed a tile-shape constraint to allow more flexible tensor layouts. * Deterministic plan selection for low‑precision GEMM with a smaller, predictable candidate set. * Earlier, clearer validation for incompatible contraction dimensions to surface errors sooner. * Consistent operand orientation and compilation caching to reduce surprises in kernel selection and runtime behavior. --------- Co-authored-by: Claude Opus 4.8 --- flashinfer/cute_dsl/utils.py | 4 +- flashinfer/gemm/gemm_base.py | 103 +- .../dense_blockscaled_gemm_sm120_b12x.py | 2025 +++++++++++------ 3 files changed, 1416 insertions(+), 716 deletions(-) diff --git a/flashinfer/cute_dsl/utils.py b/flashinfer/cute_dsl/utils.py index b7c0dc00409..ff6f685dca2 100644 --- a/flashinfer/cute_dsl/utils.py +++ b/flashinfer/cute_dsl/utils.py @@ -585,8 +585,8 @@ def sm120_make_smem_layout_sfb( assert sf_vec_size == 16 or sf_vec_size == 32, "sf_vec_size must be 16 or 32" - assert tile_shape_mnk[1] % (blk_mn // 2) == 0, ( - "tile_shape_mnk[1] must be divisible by 64" + assert tile_shape_mnk[1] % (blk_mn // 8) == 0, ( + "tile_shape_mnk[1] must be divisible by 16" ) assert tile_shape_mnk[2] % sf_vec_size == 0, ( diff --git a/flashinfer/gemm/gemm_base.py b/flashinfer/gemm/gemm_base.py index 7ec8de71c93..acfb8c076e7 100755 --- a/flashinfer/gemm/gemm_base.py +++ b/flashinfer/gemm/gemm_base.py @@ -4333,29 +4333,6 @@ def _cudnn_mm_mxfp8_requirement( _SM100_DEFAULT_CLUSTER_SHAPE_MN = (1, 1) -def _select_default_sm120_mma_tiler(m, n, sm_count): - """Select optimal SM120 tile shape based on problem size and SM count. - - Uses narrower tiles (64x64, 64x128, 128x64) when the default 128x128 - would leave SMs idle on small-M shapes. - """ - coarse_tile = (128, 128) - coarse_tiles = ((m + coarse_tile[0] - 1) // coarse_tile[0]) * ( - (n + coarse_tile[1] - 1) // coarse_tile[1] - ) - if m <= 128 and coarse_tiles < max(1, sm_count // 2): - if n > 1536: - return (64, 128) - medium_tile = (128, 64) - medium_tiles = ((m + medium_tile[0] - 1) // medium_tile[0]) * ( - (n + medium_tile[1] - 1) // medium_tile[1] - ) - if medium_tiles < max(1, sm_count // 2): - return (64, 64) - return (128, 64) - return (128, 128) - - def _get_approximate_cta_nums(m, n, tile_mn, cluster_shape_mn): tile_m, tile_n = tile_mn cluster_m, cluster_n = cluster_shape_mn @@ -5468,7 +5445,7 @@ def _cute_dsl_gemm_fp4_requirement( @supported_compute_capability([120, 121]) def _b12x_gemm_fp4_requirement( - a: torch.Tensor, # unused + a: torch.Tensor, b: torch.Tensor, # unused a_descale: torch.Tensor, # unused b_descale: torch.Tensor, # unused @@ -5493,6 +5470,15 @@ def _b12x_gemm_fp4_requirement( raise ValueError("b12x FP4 GEMM only supports 128x4 scale factor layout.") if not use_nvfp4: raise ValueError("b12x FP4 GEMM only supports NVFP4 (sf_vec_size=16).") + # K must be a multiple of 128 (tile_k = sf_vec_size * 8); a is packed FP4 (M, K//2). + real_k = a.shape[1] * 2 + if real_k % 128 != 0: + if backend != "b12x": + return False # let "auto" fall back to cutlass/cudnn + raise ValueError( + "b12x FP4 GEMM requires the contraction dim K to be a multiple of 128 " + f"(tile_k = sf_vec_size * 8). Got K={real_k}." + ) _check_cute_dsl_availability() return True @@ -5809,6 +5795,7 @@ def _b12x_gemm_fp4_runner( from .kernels.dense_blockscaled_gemm_sm120_b12x import ( Sm120B12xBlockScaledDenseGemmKernel, + _select_default_dense_gemm_plan, ) cutlass_dtype_attr = _TORCH_TO_CUTLASS_DTYPE_ATTR.get(out_dtype) @@ -5821,6 +5808,11 @@ def _b12x_gemm_fp4_runner( f"Supported: torch.bfloat16, torch.float16." ) + def _default_dense_plan(m, n, real_k, device): + return _select_default_dense_gemm_plan( + m, n, real_k, get_device_sm_count(device), expected_m=m + ) + class B12xFp4GemmRunner(TunableRunner): """TunableRunner for b12x block-scaled FP4 dense GEMM on SM120. @@ -5846,14 +5838,9 @@ def get_valid_tactics( batch_size = 1 valid_tactics = [] - sm120_mma_tiler_candidates = [ - (64, 64), - (64, 128), - (128, 64), - (128, 128), - ] - swap_ab = False - for mma_tiler_mn in sm120_mma_tiler_candidates: + + def _add(mma_tiler_mn, swap_ab): + # can_implement is M-independent (takes no `m`) if not Sm120B12xBlockScaledDenseGemmKernel.can_implement( ab_dtype, sf_dtype, @@ -5861,19 +5848,29 @@ def get_valid_tactics( c_cutlass_dtype, mma_tiler_mn, (1, 1), - m, n, real_k, batch_size, "k", "k", "n", + swap_ab=swap_ab, ): - continue + return for use_prefetch in (False, True): - valid_tactics.append( - (mma_tiler_mn, (1, 1), swap_ab, use_prefetch, "sm120", None) - ) + tac = (mma_tiler_mn, (1, 1), swap_ab, use_prefetch, "sm120", None) + if tac not in valid_tactics: + valid_tactics.append(tac) + + # A few balanced swap_ab-free tiles for the tuner to profile (a larger + # grid overfit the bucket representative and made picks noisier). + for mma_tiler_mn in [(64, 64), (64, 128), (128, 64), (128, 128)]: + _add(mma_tiler_mn, swap_ab=False) + + # Also include the default-path tile (may be a narrow-N swap_ab tile + # absent from the set above) so the tuner can't pick worse than static. + plan = _default_dense_plan(m, n, real_k, a.device) + _add(plan.mma_tiler_mn, swap_ab=plan.swap_ab) return valid_tactics def forward( @@ -5894,12 +5891,13 @@ def forward( batch_size = 1 if tactic is None or tactic == -1: + # Default path: the m-aware plan picks the tile (and swap_ab for + # narrow-N) for this shape; expected_m=m since m is the actual size. + plan = _default_dense_plan(m, n, real_k, a.device) tactic = ( - _select_default_sm120_mma_tiler( - m, n, get_device_sm_count(a.device) - ), + plan.mma_tiler_mn, (1, 1), - False, + plan.swap_ab, False, "sm120", None, @@ -5914,13 +5912,11 @@ def forward( use_tma_store, ) = tactic - # b12x SM120 kernel does not support swap_ab - kernel_m, kernel_n = m, n kernel_a, kernel_b = a, b.T kernel_a_sf, kernel_b_sf = a_descale, b_descale.T - sf_m = (kernel_m + 127) // 128 - sf_n = (kernel_n + 127) // 128 + sf_m = (m + 127) // 128 + sf_n = (n + 127) // 128 sf_k = (real_k // sf_vec_size + 3) // 4 cache_key = ( @@ -5935,14 +5931,20 @@ def forward( out_dtype, ) + # ctor takes mma_k/tile_k/single_work_tile_per_cta before use_prefetch, + # so pass the later args by keyword to avoid mis-binding. make_kernel = lambda: Sm120B12xBlockScaledDenseGemmKernel( sf_vec_size, mma_tiler_mn, cluster_shape_mn, - use_prefetch, - enable_pdl, + use_prefetch=use_prefetch, + enable_pdl=enable_pdl, + swap_ab=swap_ab, ) + # swap_ab is device-internal (applied in the kernel ctor); public C + # stays row-major (m, n). Pass swap_ab=False to the harness so it keeps + # the (m, n) output convention, not the SM100 operand-swap one. compiled_gemm, _ = _compile_block_scaled_gemm( _B12X_MM_FP4_KERNEL_CACHE, cache_key, @@ -5952,7 +5954,7 @@ def forward( c_cutlass_dtype=c_cutlass_dtype, ab_assumed_align=32, cluster_shape_mn=cluster_shape_mn, - swap_ab=swap_ab, + swap_ab=False, sf_m=sf_m, sf_n=sf_n, sf_k=sf_k, @@ -5961,6 +5963,7 @@ def forward( alpha_for_launch = _prepare_alpha_for_launch(alpha_tensor, a.device) + # `out` passed as-is (row-major (m, n)). compiled_gemm( kernel_a, kernel_b, @@ -6014,7 +6017,9 @@ def _heuristic_func_mm_fp4( is_sm103 = major == 10 and minor == 3 is_sm120 = major == 12 and minor == 0 - # SM120 + CUDA 13: prefer b12x (warp-level MMA, underfill tile selection) + # SM120 + CUDA 13: prefer b12x. SM121 (GB10) is intentionally excluded -- b12x + # is supported there as an explicit backend, but cutlass/cudnn are faster in + # most cases, so `auto` keeps using them. if is_sm120 and use_nvfp4 and cuda_major >= 13: return [c for c in ("b12x", "cutlass", "cudnn") if c in suitable_backends] diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py index 914dd281585..de2616f8b10 100644 --- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py +++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py @@ -28,10 +28,9 @@ # This file is ported from the CUTLASS dense block-scaled GEMM example # and adapted for the current Blackwell GeForce target. -# -# Ported from the b12x kernel library to FlashInfer. -from typing import Callable, List, Optional, Tuple, Type +from dataclasses import dataclass +from typing import Literal, Optional, Tuple import cuda.bindings.driver as cuda import cutlass @@ -41,27 +40,132 @@ import cutlass.utils.blackwell_helpers as sm120_utils import cutlass.utils.blockscaled_layout as blockscaled_utils import cutlass.utils.hopper_helpers as sm90_utils -import functools -import torch +import logging +from cutlass import Int32, Int64 from cutlass.cute.nvgpu import cpasync from cutlass.cute.nvgpu.warp.mma import Field as WarpField +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo +from cutlass._mlir.dialects import llvm from flashinfer.cute_dsl.utils import ( - cutlass_to_torch_dtype, - get_cutlass_dtype, - get_max_active_clusters, - get_num_sm, - make_ptr, sm120_make_smem_layout_sfa, sm120_make_smem_layout_sfb, ) -def current_cuda_stream(): - """Return current CUDA stream as a CUDA driver stream handle.""" - import cuda.bindings.driver as cuda_driver +# Vendored from b12x.cute.fp4 (so this kernel does not depend on the external +# b12x package). Used only by the kernel's opt-in split-K atomic-reduction path. +@dsl_user_op +def get_ptr_as_int64(tensor: cute.Tensor, offset, *, loc=None, ip=None) -> Int64: + """Get the memory address of tensor[offset] as Int64.""" + elem_ptr = tensor.iterator + offset + ptr_int = llvm.ptrtoint(T.i64(), elem_ptr.llvm_ptr, loc=loc, ip=ip) + return Int64(ptr_int) + + +@dsl_user_op +def scatter_add_bf16(addr: Int64, val_f32, *, loc=None, ip=None): + """BF16 atomic reduction add to global memory.""" + llvm.inline_asm( + None, + [ + Int64(addr).ir_value(loc=loc, ip=ip), + val_f32.ir_value(loc=loc, ip=ip), + ], + "{ .reg .b16 packed; cvt.rn.satfinite.bf16.f32 packed, $1; red.relaxed.gpu.global.add.noftz.bf16 [$0], packed; }", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def scatter_add_bf16x2(addr: Int64, val0_f32, val1_f32, *, loc=None, ip=None): + """BF16x2 atomic reduction add to global memory.""" + llvm.inline_asm( + None, + [ + Int64(addr).ir_value(loc=loc, ip=ip), + val0_f32.ir_value(loc=loc, ip=ip), + val1_f32.ir_value(loc=loc, ip=ip), + ], + "{ .reg .b32 packed; cvt.rn.satfinite.bf16x2.f32 packed, $2, $1; red.relaxed.gpu.global.add.noftz.bf16x2 [$0], packed; }", + "l,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +logger = logging.getLogger(__name__) +_DENSE_LOAD_PATHS = ("tma", "cpasync") + + +@dataclass(frozen=True) +class _DenseGemmPlan: + mma_tiler_mn: Tuple[int, int] + load_path: Literal["tma", "cpasync"] + swap_ab: bool + + +# @dsl_user_op on PersistentTileSchedulerParams.__init__ can rename attributes +# (e.g. raster_along_m -> _raster_along_m, cluster_shape_major_fdd -> +# cluster_shape_m_fdd) but __extract_mlir_values__ (used by TVM-FFI) +# still references the original names. +_orig_extract = utils.PersistentTileSchedulerParams.__extract_mlir_values__ + +# Map of source-code attr name -> runtime attr name set by @dsl_user_op +_ATTR_RENAMES = { + "raster_along_m": "_raster_along_m", + "cluster_shape_major_fdd": "cluster_shape_m_fdd", + "cluster_shape_minor_fdd": "cluster_shape_n_fdd", +} + + +def _patched_extract(self): + for src_name, dst_name in _ATTR_RENAMES.items(): + if not hasattr(self, src_name) and hasattr(self, dst_name): + setattr(self, src_name, getattr(self, dst_name)) + return _orig_extract(self) + + +utils.PersistentTileSchedulerParams.__extract_mlir_values__ = _patched_extract + + +def _convert_layout_acc_mn( + acc_layout: cute.Layout, transpose: bool = False +) -> cute.Layout: + acc_layout_col_major = cute.make_layout(acc_layout.shape) + shape = ( + (acc_layout_col_major.shape[0][1], acc_layout_col_major.shape[1]), + ( + acc_layout_col_major.shape[0][0], + *acc_layout_col_major.shape[0][2:], + acc_layout_col_major.shape[2], + ), + *acc_layout_col_major.shape[3:], + ) + stride = ( + (acc_layout_col_major.stride[0][1], acc_layout_col_major.stride[1]), + ( + acc_layout_col_major.stride[0][0], + *acc_layout_col_major.stride[0][2:], + acc_layout_col_major.stride[2], + ), + *acc_layout_col_major.stride[3:], + ) + if cutlass.const_expr(transpose): + shape = (shape[1], shape[0], *shape[2:]) + stride = (stride[1], stride[0], *stride[2:]) + return cute.composition(acc_layout, cute.make_layout(shape, stride=stride)) + - return cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) +def _reshape_acc_to_mn(acc: cute.Tensor, transpose: bool = False) -> cute.Tensor: + return cute.make_tensor( + acc.iterator, _convert_layout_acc_mn(acc.layout, transpose=transpose) + ) class DenseGemmKernel: @@ -78,13 +182,12 @@ class DenseGemmKernel: Notes: - Supported combinations: - * NVF4 only: A/B: Float4E2M1FN, SF: Float8E4M3FN, sf_vec_size: 16 - (MXF4 / sf_vec_size=32 is not supported β€” the CUTLASS DSL - MmaMXF4NVF4Op hardcodes sf_vec_size=16 in its constructor.) + * NVF4: A/B: Float4E2M1FN, SF: Float8E4M3FN, sf_vec_size: 16 + * MXF4: A/B: Float4E2M1FN, SF: Float8E8M0FNU, sf_vec_size: 32 - Tile shape constraints: - * tile_m must be divisible by 64 - * tile_n must be divisible by 64 - * tile_k = sf_vec_size * 8 = 128 + * tile_m must be divisible by 128 + * tile_n must be divisible by 128 + * tile_k must be divisible by 64 (sf_vec_size=16) or 128 (sf_vec_size=32) """ def __init__( @@ -92,27 +195,64 @@ def __init__( sf_vec_size: int, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], + mma_k: int = 64, + tile_k: Optional[int] = None, + single_work_tile_per_cta: bool = False, use_prefetch: bool = False, enable_pdl: bool = True, + direct_one_m_tile_scheduler: bool = False, + split_k_slices: int = 1, + split_k_atomic_bf16: bool = False, + use_m1_non_tma_a: bool = False, + use_m1_non_tma_c: bool = False, + use_m1_non_tma_sfa: bool = False, + load_path: Literal["tma", "cpasync"] = "tma", + swap_ab: bool = False, ): self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size - # K = sf_vec_size * 8 for FP4 (each FP4 element is 0.5 bytes, sf_vec_size - # elements per scale factor, and we want 4 MMA k-tiles per stage) - tile_k = sf_vec_size * 8 # 128 for sf_vec_size=16 + self.mma_k = mma_k + if tile_k is None: + tile_k = sf_vec_size * 8 self.tile_shape_mnk = (mma_tiler_mn[0], mma_tiler_mn[1], tile_k) + self.mma_tile_shape_mnk = ( + (mma_tiler_mn[1], mma_tiler_mn[0], tile_k) + if swap_ab + else self.tile_shape_mnk + ) self.sfa_tile_shape_mk = (max(128, mma_tiler_mn[0]), tile_k) self.sfa_tiles_per_block = self.sfa_tile_shape_mk[0] // mma_tiler_mn[0] self.sfb_tile_shape_nk = (max(128, mma_tiler_mn[1]), tile_k) self.sfb_tiles_per_block = self.sfb_tile_shape_nk[0] // mma_tiler_mn[1] self.cluster_shape_mnk = (1, 1, 1) # Always (1,1,1) on the current target self.epi_tile = (mma_tiler_mn[0], mma_tiler_mn[1]) + self.single_work_tile_per_cta = single_work_tile_per_cta self.use_prefetch = use_prefetch self.enable_pdl = enable_pdl + self.direct_one_m_tile_scheduler = direct_one_m_tile_scheduler + self.split_k_slices = split_k_slices + self.split_k_atomic_bf16 = split_k_atomic_bf16 + self.use_m1_non_tma_a = use_m1_non_tma_a + self.use_m1_non_tma_c = use_m1_non_tma_c + self.use_m1_non_tma_sfa = use_m1_non_tma_sfa + self.load_path = load_path + self.swap_ab = swap_ab + mma_atom_mn = (self.mma_tile_shape_mnk[0], self.mma_tile_shape_mnk[1]) + if mma_atom_mn in ((16, 64), (16, 128)): + self.atom_shape = (1, 2, 1) + elif mma_atom_mn in ((32, 64), (32, 128)): + self.atom_shape = (2, 2, 1) + else: + self.atom_shape = (4, 2, 1) self.tiled_mma = None self.occupancy = 1 - self.num_mma_warps = 8 + if mma_atom_mn in ((16, 64), (16, 128)): + self.num_mma_warps = 2 + elif mma_atom_mn in ((32, 64), (32, 128)): + self.num_mma_warps = 4 + else: + self.num_mma_warps = 8 self.tma_load_warp_id = self.num_mma_warps self.num_threads_per_warp = 32 self.threads_per_cta = ( @@ -141,15 +281,20 @@ def __init__( self.mma_register_requirement = 232 def _setup_attributes(self): + # FP4-only target (NVF4 sf_vec_size=16 / MXF4 sf_vec_size=32). The MXFP8 + # warp-MMA path was dropped: FlashInfer only drives this kernel for FP4, + # and cute.nvgpu.warp.MmaMXF8Op is absent in the public cutlass-dsl build. mma_op = cute.nvgpu.warp.MmaMXF4NVF4Op( self.a_dtype, self.acc_dtype, self.sf_dtype, ) - atom_shape = (4, 2, 1) + atom_shape = self.atom_shape atom_layout = cute.make_layout(atom_shape) permutation_mnk = sm120_utils.get_permutation_mnk( - self.tile_shape_mnk, self.sf_vec_size, False + self.mma_tile_shape_mnk, + self.sf_vec_size, + False, # is_mxfp8: FP4-only ) self.tiled_mma = cute.make_tiled_mma( mma_op, @@ -159,11 +304,11 @@ def _setup_attributes(self): # Bare atom for manual unroll workaround (avoids hasAuxTensor address space bug) self.mma_atom = cute.make_mma_atom(mma_op) # Compute atom loop bounds from tile shape and atom/layout shape - # MMA atom: m16, n8, k64; atom_layout: (4,2,1) -> group: m64, n16, k64 - mma_m, mma_n, mma_k = 16, 8, 64 - self.num_m_tiles = self.tile_shape_mnk[0] // (mma_m * atom_shape[0]) - self.num_n_tiles = self.tile_shape_mnk[1] // (mma_n * atom_shape[1]) - self.num_k_blocks = self.tile_shape_mnk[2] // mma_k + # MMA atom: m16n8k64 for FP4. + mma_m, mma_n, mma_k = 16, 8, self.mma_k + self.num_m_tiles = self.mma_tile_shape_mnk[0] // (mma_m * atom_shape[0]) + self.num_n_tiles = self.mma_tile_shape_mnk[1] // (mma_n * atom_shape[1]) + self.num_k_blocks = self.mma_tile_shape_mnk[2] // mma_k self.cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk) @@ -284,13 +429,17 @@ def __call__( (self.tile_shape_mnk[1], self.tile_shape_mnk[2]), 1, ) - tma_atom_sfa, tma_tensor_sfa = self._make_tma_atoms_and_tensors( - sfa_tensor, - self.sfa_smem_layout_staged, - self.sfa_tile_shape_mk, - 1, - internal_type=cutlass.Int16, - ) + if cutlass.const_expr(self.use_m1_non_tma_sfa): + tma_atom_sfa = tma_atom_b + tma_tensor_sfa = sfa_tensor + else: + tma_atom_sfa, tma_tensor_sfa = self._make_tma_atoms_and_tensors( + sfa_tensor, + self.sfa_smem_layout_staged, + self.sfa_tile_shape_mk, + 1, + internal_type=cutlass.Int16, + ) tma_atom_sfb, tma_tensor_sfb = self._make_tma_atoms_and_tensors( sfb_tensor, self.sfb_smem_layout_staged, @@ -308,6 +457,8 @@ def __call__( c, self.tile_shape_mnk, max_active_clusters, + self.direct_one_m_tile_scheduler, + self.split_k_slices, ) @cute.struct @@ -351,14 +502,19 @@ class SharedStorage: self.kernel( tma_atom_a, tma_tensor_a, + a, tma_atom_b, tma_tensor_b, + b, tma_atom_sfa, tma_tensor_sfa, + sfa_tensor, tma_atom_sfb, tma_tensor_sfb, + sfb_tensor, tma_atom_c, tma_tensor_c, + c, self.tiled_mma, self.mma_atom, self.cta_layout_mnk, @@ -543,20 +699,133 @@ def _get_layoutSFB_TV(self, tiled_mma: cute.TiledMma): layout_tv = cute.composition(layout_tv, (thridx_2_thrid, None)) return layout_tv + @cute.jit + def _make_cpasync_tiled_copy( + self, + dtype: cutlass.Constexpr, + tile_cols: cutlass.Constexpr[int], + ) -> cute.TiledCopy: + copy_bits = 128 + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + dtype, + num_bits_per_copy=copy_bits, + ) + async_copy_elems = copy_bits // dtype.width + t_shape_dim_1 = tile_cols // async_copy_elems + assert self.num_threads_per_warp % t_shape_dim_1 == 0 + t_layout = cute.make_ordered_layout( + (self.num_threads_per_warp // t_shape_dim_1, t_shape_dim_1), + order=(1, 0), + ) + v_layout = cute.make_layout((1, async_copy_elems)) + return cute.make_tiled_copy_tv(atom_async_copy, t_layout, v_layout) + + @cute.jit + def _make_scale_tiled_copy( + self, + dtype: cutlass.Constexpr, + ) -> cute.TiledCopy: + copy_bits = dtype.width + atom_async_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + dtype, + num_bits_per_copy=copy_bits, + ) + return cute.make_tiled_copy_tv( + atom_async_copy, + cute.make_layout((self.num_threads_per_warp,)), + cute.make_layout((copy_bits // dtype.width,)), + ) + + @cute.jit + def _predicate_cpasync_rows( + self, + tCc: cute.Tensor, + row_limit: Int32, + ) -> cute.Tensor: + tPred = cute.make_fragment( + cute.make_layout( + ( + cute.size(tCc, mode=[0, 1]), + cute.size(tCc, mode=[1]), + cute.size(tCc, mode=[2]), + ), + stride=(cute.size(tCc, mode=[2]), 0, 1), + ), + cutlass.Boolean, + ) + for rest_v in cutlass.range_constexpr(tPred.shape[0]): + for rest_k in cutlass.range_constexpr(tPred.shape[2]): + tPred[rest_v, 0, rest_k] = tCc[(0, rest_v), 0, rest_k][0] < row_limit + return tPred + + @cute.jit + def _cpasync_copy_2d( + self, + tiled_copy: cute.TiledCopy, + tG: cute.Tensor, + tS: cute.Tensor, + tC: cute.Tensor, + row_limit: Int32, + predicate_rows: cutlass.Constexpr[bool], + ) -> None: + if cutlass.const_expr(predicate_rows): + tP = self._predicate_cpasync_rows(tC, row_limit) + for rest_m in cutlass.range_constexpr(cute.size(tS.shape[1])): + if cutlass.const_expr(predicate_rows): + cute.copy( + tiled_copy, + tG[None, rest_m, None], + tS[None, rest_m, None], + pred=tP[None, rest_m, None], + ) + else: + cute.copy( + tiled_copy, + tG[None, rest_m, None], + tS[None, rest_m, None], + ) + + @cute.jit + def _scale_copy_2d( + self, + tiled_copy: cute.TiledCopy, + tG: cute.Tensor, + tS: cute.Tensor, + tC: cute.Tensor, + row_limit: Int32, + ) -> None: + tP = cute.make_fragment(cute.make_layout(tS.shape), cutlass.Boolean) + for i in cutlass.range_constexpr(cute.size(tP)): + tP[i] = cute.elem_less(tC[i][0][0][0], row_limit) + for rest_m in cutlass.range_constexpr(cute.size(tS.shape[1])): + cute.copy( + tiled_copy, + tG[None, rest_m, None], + tS[None, rest_m, None], + pred=tP[None, rest_m, None], + ) + # GPU device kernel @cute.kernel def kernel( self, tma_atom_a: cute.CopyAtom, mA_mkl: cute.Tensor, + directA_mkl: cute.Tensor, tma_atom_b: cute.CopyAtom, mB_nkl: cute.Tensor, + directB_nkl: cute.Tensor, tma_atom_sfa: cute.CopyAtom, mSFA_mkl: cute.Tensor, + directSFA_mkl: cute.Tensor, tma_atom_sfb: cute.CopyAtom, mSFB_nkl: cute.Tensor, + directSFB_nkl: cute.Tensor, tma_atom_c: cute.CopyAtom, mC_mnl: cute.Tensor, + directC_mnl: cute.Tensor, tiled_mma: cute.TiledMma, mma_atom: cute.MmaAtom, cta_layout_mnk: cute.Layout, @@ -578,11 +847,20 @@ def kernel( # Prefetch TMA descriptors if warp_idx == 0: - cpasync.prefetch_descriptor(tma_atom_a) - cpasync.prefetch_descriptor(tma_atom_b) - cpasync.prefetch_descriptor(tma_atom_sfa) - cpasync.prefetch_descriptor(tma_atom_sfb) - cpasync.prefetch_descriptor(tma_atom_c) + if cutlass.const_expr( + self.load_path == "tma" and not self.use_m1_non_tma_a + ): + cpasync.prefetch_descriptor(tma_atom_a) + if cutlass.const_expr(self.load_path == "tma"): + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr( + self.load_path == "tma" and not self.use_m1_non_tma_sfa + ): + cpasync.prefetch_descriptor(tma_atom_sfa) + if cutlass.const_expr(self.load_path == "tma"): + cpasync.prefetch_descriptor(tma_atom_sfb) + if cutlass.const_expr(not self.use_m1_non_tma_c): + cpasync.prefetch_descriptor(tma_atom_c) cta_rank_in_cluster = cute.arch.make_warp_uniform( cute.arch.block_idx_in_cluster() @@ -593,12 +871,19 @@ def kernel( b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, 0)) sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, 0)) sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, 0)) - tma_copy_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem_layout) - + cute.size_in_bytes(self.b_dtype, b_smem_layout) - + cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) - + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) - ) + if cutlass.const_expr(self.use_m1_non_tma_sfa): + tma_copy_bytes = cute.size_in_bytes( + self.b_dtype, b_smem_layout + ) + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + if cutlass.const_expr(not self.use_m1_non_tma_a): + tma_copy_bytes += cute.size_in_bytes(self.a_dtype, a_smem_layout) + else: + tma_copy_bytes = ( + cute.size_in_bytes(self.a_dtype, a_smem_layout) + + cute.size_in_bytes(self.b_dtype, b_smem_layout) + + cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + ) # Allocate shared memory smem = cutlass.utils.SmemAllocator() @@ -614,14 +899,30 @@ def kernel( ) cta_layout_vmnk = cute.make_layout((1, *cta_layout_mnk.shape)) - mainloop_pipeline = pipeline.PipelineTmaAsync.create( - num_stages=self.ab_stage, - producer_group=mainloop_pipeline_producer_group, - consumer_group=mainloop_pipeline_consumer_group, - tx_count=tma_copy_bytes, - barrier_storage=mainloop_pipeline_array_ptr, - cta_layout_vmnk=cta_layout_vmnk, - ) + if cutlass.const_expr(self.load_path == "cpasync"): + mainloop_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_threads_per_warp, + ) + mainloop_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mma_warps * self.num_threads_per_warp, + ) + mainloop_pipeline = pipeline.PipelineAsync.create( + num_stages=self.ab_stage, + producer_group=mainloop_pipeline_producer_group, + consumer_group=mainloop_pipeline_consumer_group, + barrier_storage=mainloop_pipeline_array_ptr, + ) + else: + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.ab_stage, + producer_group=mainloop_pipeline_producer_group, + consumer_group=mainloop_pipeline_consumer_group, + tx_count=tma_copy_bytes, + barrier_storage=mainloop_pipeline_array_ptr, + cta_layout_vmnk=cta_layout_vmnk, + ) if cute.size(self.cluster_shape_mnk) > 1: cute.arch.cluster_arrive_relaxed() @@ -650,16 +951,38 @@ def kernel( cute.slice_(self.tile_shape_mnk, (0, None, None)), (None, None, None), ) - gSFA_mkl = cute.local_tile( - mSFA_mkl, - self.sfa_tile_shape_mk, - (None, None, None), - ) + if cutlass.const_expr(not self.use_m1_non_tma_sfa): + gSFA_mkl = cute.local_tile( + mSFA_mkl, + self.sfa_tile_shape_mk, + (None, None, None), + ) gSFB_nkl = cute.local_tile( mSFB_nkl, self.sfb_tile_shape_nk, (None, None, None), ) + if cutlass.const_expr(self.load_path == "cpasync"): + gA_cpasync_mkl = cute.local_tile( + directA_mkl, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + gB_cpasync_nkl = cute.local_tile( + directB_nkl, + cute.slice_(self.tile_shape_mnk, (0, None, None)), + (None, None, None), + ) + gSFA_cpasync_mkl = cute.local_tile( + directSFA_mkl, + self.sfa_tile_shape_mk, + (None, None, None), + ) + gSFB_cpasync_nkl = cute.local_tile( + directSFB_nkl, + self.sfb_tile_shape_nk, + (None, None, None), + ) gC_mnl = cute.local_tile( mC_mnl, cute.slice_(self.tile_shape_mnk, (None, None, 0)), @@ -672,57 +995,133 @@ def kernel( # TMA partitions for A a_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (0, None, 0)).shape) a_cta_crd = cluster_coord_mnk[1] - tAsA, tAgA = cpasync.tma_partition( - tma_atom_a, - a_cta_crd, - a_cta_layout, - cute.group_modes(sA, 0, 2), - cute.group_modes(gA_mkl, 0, 2), - ) + if cutlass.const_expr(self.load_path == "tma" and not self.use_m1_non_tma_a): + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + a_cta_crd, + a_cta_layout, + cute.group_modes(sA, 0, 2), + cute.group_modes(gA_mkl, 0, 2), + ) # TMA partitions for B b_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (None, 0, 0)).shape) b_cta_crd = cluster_coord_mnk[0] - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, - b_cta_crd, - b_cta_layout, - cute.group_modes(sB, 0, 2), - cute.group_modes(gB_nkl, 0, 2), - ) + if cutlass.const_expr(self.load_path == "tma"): + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + b_cta_crd, + b_cta_layout, + cute.group_modes(sB, 0, 2), + cute.group_modes(gB_nkl, 0, 2), + ) # TMA partitions for SFA - tAsSFA, tAgSFA = cpasync.tma_partition( - tma_atom_sfa, - a_cta_crd, - a_cta_layout, - cute.group_modes(sSFA, 0, 2), - cute.group_modes(gSFA_mkl, 0, 2), - ) - tAsSFA = cute.filter_zeros(tAsSFA) - tAgSFA = cute.filter_zeros(tAgSFA) + if cutlass.const_expr(self.load_path == "tma" and not self.use_m1_non_tma_sfa): + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + a_cta_crd, + a_cta_layout, + cute.group_modes(sSFA, 0, 2), + cute.group_modes(gSFA_mkl, 0, 2), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) # TMA partitions for SFB - tBsSFB, tBgSFB = cpasync.tma_partition( - tma_atom_sfb, - b_cta_crd, - b_cta_layout, - cute.group_modes(sSFB, 0, 2), - cute.group_modes(gSFB_nkl, 0, 2), - ) - tBsSFB = cute.filter_zeros(tBsSFB) - tBgSFB = cute.filter_zeros(tBgSFB) + if cutlass.const_expr(self.load_path == "tma"): + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + b_cta_crd, + b_cta_layout, + cute.group_modes(sSFB, 0, 2), + cute.group_modes(gSFB_nkl, 0, 2), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) - # Make fragments - tCsA = thr_mma.partition_A(sA) - tCsB = thr_mma.partition_B(sB) + if cutlass.const_expr(self.load_path == "cpasync"): + cpasync_tiled_copy_A = self._make_cpasync_tiled_copy( + self.a_dtype, + self.tile_shape_mnk[2], + ) + cpasync_tiled_copy_B = self._make_cpasync_tiled_copy( + self.b_dtype, + self.tile_shape_mnk[2], + ) + cpasync_tiled_copy_SF = self._make_scale_tiled_copy(self.sf_dtype) + cA_mkl = cute.make_identity_tensor(cute.shape(directA_mkl)) + cA_cpasync_mkl = cute.local_tile( + cA_mkl, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + cB_nkl = cute.make_identity_tensor(cute.shape(directB_nkl)) + cB_cpasync_nkl = cute.local_tile( + cB_nkl, + cute.slice_(self.tile_shape_mnk, (0, None, None)), + (None, None, None), + ) + cSFA_mkl = cute.make_identity_tensor(cute.shape(directSFA_mkl)) + cSFA_cpasync_mkl = cute.local_tile( + cSFA_mkl, + self.sfa_tile_shape_mk, + (None, None, None), + ) + cSFB_nkl = cute.make_identity_tensor(cute.shape(directSFB_nkl)) + cSFB_cpasync_nkl = cute.local_tile( + cSFB_nkl, + self.sfb_tile_shape_nk, + (None, None, None), + ) + + cpasync_lane = tidx % self.num_threads_per_warp + thr_cpasync_A = cpasync_tiled_copy_A.get_slice(cpasync_lane) + thr_cpasync_B = cpasync_tiled_copy_B.get_slice(cpasync_lane) + thr_cpasync_SF = cpasync_tiled_copy_SF.get_slice(cpasync_lane) + tAgA_cpasync_mkl = thr_cpasync_A.partition_S(gA_cpasync_mkl) + tAsA_cpasync = thr_cpasync_A.partition_D(sA) + tAcA_cpasync_mkl = thr_cpasync_A.partition_S(cA_cpasync_mkl) + tBgB_cpasync_nkl = thr_cpasync_B.partition_S(gB_cpasync_nkl) + tBsB_cpasync = thr_cpasync_B.partition_D(sB) + tBcB_cpasync_nkl = thr_cpasync_B.partition_S(cB_cpasync_nkl) + tAgSFA_cpasync_mkl = thr_cpasync_SF.partition_S(gSFA_cpasync_mkl) + tAsSFA_cpasync = thr_cpasync_SF.partition_D(sSFA) + tAcSFA_cpasync_mkl = thr_cpasync_SF.partition_S(cSFA_cpasync_mkl) + tBgSFB_cpasync_nkl = thr_cpasync_SF.partition_S(gSFB_cpasync_nkl) + tBsSFB_cpasync = thr_cpasync_SF.partition_D(sSFB) + tBcSFB_cpasync_nkl = thr_cpasync_SF.partition_S(cSFB_cpasync_nkl) + + # Make fragments. swap_ab keeps public C[M,N] unchanged but presents + # B as MMA-A and A as MMA-B. + if cutlass.const_expr(self.swap_ab): + tCsA = thr_mma.partition_A(sB) + tCsB = thr_mma.partition_B(sA) + else: + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) - tCrSFA_full = self._partition_fragment_SFA(sSFA[None, None, 0], thr_mma, tidx) - tCrSFB_full = self._partition_fragment_SFB(sSFB[None, None, 0], thr_mma, tidx) - - tCgC = thr_mma.partition_C(gC_mnl) + if cutlass.const_expr(self.swap_ab): + tCrSFA_full = self._partition_fragment_SFA( + sSFB[None, None, 0], thr_mma, tidx + ) + tCrSFB_full = self._partition_fragment_SFB( + sSFA[None, None, 0], thr_mma, tidx + ) + c_mma = cute.make_identity_tensor( + (self.tile_shape_mnk[1], self.tile_shape_mnk[0]) + ) + tCgC = thr_mma.partition_C(c_mma) + else: + tCrSFA_full = self._partition_fragment_SFA( + sSFA[None, None, 0], thr_mma, tidx + ) + tCrSFB_full = self._partition_fragment_SFB( + sSFB[None, None, 0], thr_mma, tidx + ) + tCgC = thr_mma.partition_C(gC_mnl) acc_shape = tCgC.shape[:3] accumulators = cute.make_rmem_tensor(acc_shape, self.acc_dtype) @@ -733,12 +1132,28 @@ def kernel( cute.arch.sync_threads() k_tile_cnt = cute.size(gA_mkl, mode=[3]) + block_idx = cute.arch.block_idx() + k_tile_start = Int32(0) + k_tile_iter_cnt = k_tile_cnt + if cutlass.const_expr(self.split_k_slices > 1): + k_tiles_per_split = k_tile_cnt // self.split_k_slices + k_tile_start = Int32(block_idx[1]) * Int32(k_tiles_per_split) + k_tile_iter_cnt = k_tiles_per_split # Tile scheduler - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() - ) - work_tile = tile_sched.initial_work_tile_info() + if cutlass.const_expr(self.direct_one_m_tile_scheduler): + direct_tile_valid = Int32(block_idx[2]) < Int32( + tile_sched_params.problem_shape_ntile_mnl[1] + ) + work_tile = WorkTileInfo( + (Int32(0), Int32(block_idx[2]), Int32(0)), + direct_tile_valid, + ) + else: + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, block_idx, cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() # Pipeline states mainloop_producer_state = pipeline.make_pipeline_state( @@ -755,14 +1170,24 @@ def kernel( num_k_blocks = cute.size(tCrA, mode=[2]) # Copy atoms for SMEM->RMEM - atom_copy_ldmatrix_A = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(self.a_layout.is_m_major_a(), 4), - self.a_dtype, - ) - atom_copy_ldmatrix_B = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(self.b_layout.is_n_major_b(), 4), - self.b_dtype, - ) + if cutlass.const_expr(self.swap_ab): + atom_copy_ldmatrix_A = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(self.b_layout.is_n_major_b(), 4), + self.b_dtype, + ) + atom_copy_ldmatrix_B = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(self.a_layout.is_m_major_a(), 4), + self.a_dtype, + ) + else: + atom_copy_ldmatrix_A = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(self.a_layout.is_m_major_a(), 4), + self.a_dtype, + ) + atom_copy_ldmatrix_B = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(self.b_layout.is_n_major_b(), 4), + self.b_dtype, + ) smem_tiled_copy_A = cute.make_tiled_copy_A(atom_copy_ldmatrix_A, tiled_mma) smem_tiled_copy_B = cute.make_tiled_copy_B(atom_copy_ldmatrix_B, tiled_mma) @@ -789,16 +1214,24 @@ def kernel( thr_copy_ldmatrix_A = smem_tiled_copy_A.get_slice(tidx) thr_copy_ldmatrix_B = smem_tiled_copy_B.get_slice(tidx) - tCsA_copy_view = thr_copy_ldmatrix_A.partition_S(sA) + tCsA_copy_view = thr_copy_ldmatrix_A.partition_S( + sB if cutlass.const_expr(self.swap_ab) else sA + ) tCrA_copy_view = thr_copy_ldmatrix_A.retile(tCrA) - tCsB_copy_view = thr_copy_ldmatrix_B.partition_S(sB) + tCsB_copy_view = thr_copy_ldmatrix_B.partition_S( + sA if cutlass.const_expr(self.swap_ab) else sB + ) tCrB_copy_view = thr_copy_ldmatrix_B.retile(tCrB) thr_copy_ldmatrix_SFA = smem_tiled_copy_SFA.get_slice(tidx) thr_copy_ldmatrix_SFB = smem_tiled_copy_SFB.get_slice(tidx) - tCsSFA_copy_view_full = thr_copy_ldmatrix_SFA.partition_S(sSFA) + tCsSFA_copy_view_full = thr_copy_ldmatrix_SFA.partition_S( + sSFB if cutlass.const_expr(self.swap_ab) else sSFA + ) tCrSFA_copy_view_full = thr_copy_ldmatrix_SFA.retile(tCrSFA_full) - tCsSFB_copy_view_full = thr_copy_ldmatrix_SFB.partition_S(sSFB) + tCsSFB_copy_view_full = thr_copy_ldmatrix_SFB.partition_S( + sSFA if cutlass.const_expr(self.swap_ab) else sSFB + ) tCrSFB_copy_view_full = thr_copy_ldmatrix_SFB.retile(tCrSFB_full) while work_tile.is_valid_tile: @@ -806,43 +1239,91 @@ def kernel( gC_mnl_slice = gC_mnl[(None, None, *tile_coord_mnl)] sfa_tile_offset = tile_coord_mnl[0] % self.sfa_tiles_per_block sfb_tile_offset = tile_coord_mnl[1] % self.sfb_tiles_per_block - if cutlass.const_expr(self.sfa_tiles_per_block > 1): - sSFA_tile = cute.local_tile( - sSFA, - cute.slice_(self.tile_shape_mnk, (None, 0, None)), - (sfa_tile_offset, 0, None), - ) - tCsSFA_tile_copy_view = thr_copy_ldmatrix_SFA.partition_S(sSFA_tile) - tCrSFA_tile = self._partition_fragment_SFA( - sSFA_tile[None, None, 0], thr_mma, tidx - ) - tCrSFA_tile_copy_view = thr_copy_ldmatrix_SFA.retile(tCrSFA_tile) - else: - tCsSFA_tile_copy_view = tCsSFA_copy_view_full - tCrSFA_tile = tCrSFA_full - tCrSFA_tile_copy_view = tCrSFA_copy_view_full - if cutlass.const_expr(self.sfb_tiles_per_block > 1): - sSFB_tile = cute.local_tile( - sSFB, - cute.slice_(self.tile_shape_mnk, (0, None, None)), - (sfb_tile_offset, 0, None), - ) - tCsSFB_tile_copy_view = thr_copy_ldmatrix_SFB.partition_S(sSFB_tile) - tCrSFB_tile = self._partition_fragment_SFB( - sSFB_tile[None, None, 0], thr_mma, tidx - ) - tCrSFB_tile_copy_view = thr_copy_ldmatrix_SFB.retile(tCrSFB_tile) + if cutlass.const_expr(self.swap_ab): + if cutlass.const_expr(self.sfb_tiles_per_block > 1): + sSFB_tile = cute.local_tile( + sSFB, + cute.slice_(self.tile_shape_mnk, (0, None, None)), + (sfb_tile_offset, 0, None), + ) + tCsSFA_tile_copy_view = thr_copy_ldmatrix_SFA.partition_S( + sSFB_tile + ) + tCrSFA_tile = self._partition_fragment_SFA( + sSFB_tile[None, None, 0], thr_mma, tidx + ) + tCrSFA_tile_copy_view = thr_copy_ldmatrix_SFA.retile( + tCrSFA_tile + ) + else: + tCsSFA_tile_copy_view = tCsSFA_copy_view_full + tCrSFA_tile = tCrSFA_full + tCrSFA_tile_copy_view = tCrSFA_copy_view_full + if cutlass.const_expr(self.sfa_tiles_per_block > 1): + sSFA_tile = cute.local_tile( + sSFA, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (sfa_tile_offset, 0, None), + ) + tCsSFB_tile_copy_view = thr_copy_ldmatrix_SFB.partition_S( + sSFA_tile + ) + tCrSFB_tile = self._partition_fragment_SFB( + sSFA_tile[None, None, 0], thr_mma, tidx + ) + tCrSFB_tile_copy_view = thr_copy_ldmatrix_SFB.retile( + tCrSFB_tile + ) + else: + tCsSFB_tile_copy_view = tCsSFB_copy_view_full + tCrSFB_tile = tCrSFB_full + tCrSFB_tile_copy_view = tCrSFB_copy_view_full else: - tCsSFB_tile_copy_view = tCsSFB_copy_view_full - tCrSFB_tile = tCrSFB_full - tCrSFB_tile_copy_view = tCrSFB_copy_view_full + if cutlass.const_expr(self.sfa_tiles_per_block > 1): + sSFA_tile = cute.local_tile( + sSFA, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (sfa_tile_offset, 0, None), + ) + tCsSFA_tile_copy_view = thr_copy_ldmatrix_SFA.partition_S( + sSFA_tile + ) + tCrSFA_tile = self._partition_fragment_SFA( + sSFA_tile[None, None, 0], thr_mma, tidx + ) + tCrSFA_tile_copy_view = thr_copy_ldmatrix_SFA.retile( + tCrSFA_tile + ) + else: + tCsSFA_tile_copy_view = tCsSFA_copy_view_full + tCrSFA_tile = tCrSFA_full + tCrSFA_tile_copy_view = tCrSFA_copy_view_full + if cutlass.const_expr(self.sfb_tiles_per_block > 1): + sSFB_tile = cute.local_tile( + sSFB, + cute.slice_(self.tile_shape_mnk, (0, None, None)), + (sfb_tile_offset, 0, None), + ) + tCsSFB_tile_copy_view = thr_copy_ldmatrix_SFB.partition_S( + sSFB_tile + ) + tCrSFB_tile = self._partition_fragment_SFB( + sSFB_tile[None, None, 0], thr_mma, tidx + ) + tCrSFB_tile_copy_view = thr_copy_ldmatrix_SFB.retile( + tCrSFB_tile + ) + else: + tCsSFB_tile_copy_view = tCsSFB_copy_view_full + tCrSFB_tile = tCrSFB_full + tCrSFB_tile_copy_view = tCrSFB_copy_view_full accumulators.fill(0.0) # Pipelined MAINLOOP mainloop_consumer_state.reset_count() peek_ab_full_status = cutlass.Boolean(1) - if mainloop_consumer_state.count < k_tile_cnt: + if mainloop_consumer_state.count < k_tile_iter_cnt: peek_ab_full_status = mainloop_pipeline.consumer_try_wait( mainloop_consumer_state ) @@ -885,7 +1366,7 @@ def kernel( tCrSFB_copy_view_filtered[None, None, 0], ) - for _k_tile in range(0, k_tile_cnt - 1, 1, unroll=2): # type: ignore[call-overload] + for _k_tile in range(0, k_tile_iter_cnt - 1, 1, unroll=2): # type: ignore[call-overload] for k_block_idx in cutlass.range_constexpr(num_k_blocks): k_block_next = ( 0 if k_block_idx + 1 == num_k_blocks else k_block_idx + 1 @@ -1022,200 +1503,664 @@ def kernel( accumulators[None, _mt, _nt], ) - # EPILOGUE - _is_m_major = self.c_layout.is_m_major_c() - if cutlass.const_expr(self.c_dtype.width == 16): - copy_atom_r2s = cute.make_copy_atom( - cute.nvgpu.warp.StMatrix8x8x16bOp(_is_m_major, 2), - self.c_dtype, + if cutlass.const_expr(self.swap_ab): + acc_mn = _reshape_acc_to_mn(accumulators, transpose=True) + c_identity = cute.make_identity_tensor( + (self.tile_shape_mnk[1], self.tile_shape_mnk[0]) ) - else: - copy_atom_r2s = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - self.c_dtype, + coord_mn = _reshape_acc_to_mn( + thr_mma.partition_C(c_identity), + transpose=True, ) + for acc_m in cutlass.range_constexpr(cute.size(acc_mn.shape[0])): + for acc_n in cutlass.range_constexpr( + cute.size(acc_mn.shape[1]) + ): + coord = coord_mn[acc_m, acc_n] + m_coord = ( + tile_coord_mnl[0] * Int32(self.tile_shape_mnk[0]) + + coord[1] + ) + n_coord = ( + tile_coord_mnl[1] * Int32(self.tile_shape_mnk[1]) + + coord[0] + ) + if m_coord < Int32( + directC_mnl.shape[0] + ) and n_coord < Int32(directC_mnl.shape[1]): + directC_mnl[ + ( + m_coord, + n_coord, + tile_coord_mnl[2], + ) + ] = epilogue_op( + (alpha_value * acc_mn[acc_m, acc_n]).to( + self.c_dtype + ) + ) + if cutlass.const_expr(self.single_work_tile_per_cta): + work_tile = WorkTileInfo( + work_tile.tile_idx, + cutlass.Boolean(0), + ) + else: + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + if cutlass.const_expr(not self.swap_ab): + # EPILOGUE + _is_m_major = self.c_layout.is_m_major_c() + if cutlass.const_expr(self.c_dtype.width == 16): + copy_atom_r2s = cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp(_is_m_major, 2), + self.c_dtype, + ) + else: + copy_atom_r2s = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.c_dtype, + ) - copy_atom_C = cute.make_copy_atom( - cute.nvgpu.warp.StMatrix8x8x16bOp( - self.c_layout.is_m_major_c(), - 2, - ), - self.c_dtype, - ) + if cutlass.const_expr(self.c_dtype.width == 16): + copy_atom_C = cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp( + self.c_layout.is_m_major_c(), + 2, + ), + self.c_dtype, + ) + else: + copy_atom_C = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.c_dtype + ) - tiled_copy_C_Atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma) + tiled_copy_C_Atom = cute.make_tiled_copy_C_atom( + copy_atom_C, tiled_mma + ) - tiled_copy_r2s = cute.make_tiled_copy_S( - copy_atom_r2s, - tiled_copy_C_Atom, - ) + tiled_copy_r2s = cute.make_tiled_copy_S( + copy_atom_r2s, + tiled_copy_C_Atom, + ) - thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) - tRS_sD = thr_copy_r2s.partition_D(sC) - tRS_rAcc = tiled_copy_r2s.retile(accumulators) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sD = thr_copy_r2s.partition_D(sC) + tRS_rAcc = tiled_copy_r2s.retile(accumulators) - rD_shape = cute.shape(thr_copy_r2s.partition_S(sC)) - tRS_rD_layout = cute.make_layout(rD_shape[:3]) - tRS_rD = cute.make_rmem_tensor(tRS_rD_layout.shape, self.acc_dtype) + rD_shape = cute.shape(thr_copy_r2s.partition_S(sC)) + tRS_rD_layout = cute.make_layout(rD_shape[:3]) + tRS_rD = cute.make_rmem_tensor(tRS_rD_layout.shape, self.acc_dtype) - sepi_for_tma_partition = cute.group_modes(sC, 0, 2) - tcgc_for_tma_partition = cute.zipped_divide(gC_mnl_slice, self.epi_tile) + sepi_for_tma_partition = cute.group_modes(sC, 0, 2) + tcgc_for_tma_partition = cute.zipped_divide( + gC_mnl_slice, self.epi_tile + ) - bSG_sD, bSG_gD = cpasync.tma_partition( - tma_atom_c, - 0, - cute.make_layout(1), - sepi_for_tma_partition, - tcgc_for_tma_partition, - ) + bSG_sD, bSG_gD = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sepi_for_tma_partition, + tcgc_for_tma_partition, + ) - epi_rest_m = bSG_gD.shape[1][0] - epi_rest_n = bSG_gD.shape[1][1] - epi_tile_m = self.epi_tile[0] - epi_tile_n = self.epi_tile[1] - mma_tile_m = self.tile_shape_mnk[0] // cute.size(tRS_rAcc, mode=[1]) - mma_tile_n = self.tile_shape_mnk[1] // cute.size(tRS_rAcc, mode=[2]) - has_multi_epi_store = cutlass.const_expr( - not (self.epi_stage == 1 and epi_rest_m == 1 and epi_rest_n == 1) - ) - tma_store_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, - self.num_mma_warps * self.num_threads_per_warp, - ) - tma_store_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.epi_stage, - producer_group=tma_store_producer_group, - ) + epi_rest_m = bSG_gD.shape[1][0] + epi_rest_n = bSG_gD.shape[1][1] + epi_tile_m = self.epi_tile[0] + epi_tile_n = self.epi_tile[1] + mma_tile_m = self.tile_shape_mnk[0] // cute.size(tRS_rAcc, mode=[1]) + mma_tile_n = self.tile_shape_mnk[1] // cute.size(tRS_rAcc, mode=[2]) + has_multi_epi_store = cutlass.const_expr( + not ( + self.epi_stage == 1 and epi_rest_m == 1 and epi_rest_n == 1 + ) + ) + tma_store_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mma_warps * self.num_threads_per_warp, + ) + tma_store_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.epi_stage, + producer_group=tma_store_producer_group, + ) - for epi_m in cutlass.range_constexpr(epi_rest_m): - for epi_n in cutlass.range_constexpr(epi_rest_n): - MmaMPerEpiM = epi_tile_m // mma_tile_m - MmaNPerEpiN = epi_tile_n // mma_tile_n - for mma_n_in_epi in cutlass.range_constexpr(MmaNPerEpiN): - for mma_m_in_epi in cutlass.range_constexpr(MmaMPerEpiM): - mma_n = (epi_n * MmaNPerEpiN) + mma_n_in_epi - mma_m = (epi_m * MmaMPerEpiM) + mma_m_in_epi - tRS_rD_slice = tRS_rD[ - (None, mma_m_in_epi, mma_n_in_epi) - ] - tRS_rAcc_slice = tRS_rAcc[(None, mma_m, mma_n)] - for elem_idx in cutlass.range_constexpr( - cute.size(tRS_rD_slice) + for epi_m in cutlass.range_constexpr(epi_rest_m): + for epi_n in cutlass.range_constexpr(epi_rest_n): + MmaMPerEpiM = epi_tile_m // mma_tile_m + MmaNPerEpiN = epi_tile_n // mma_tile_n + for mma_n_in_epi in cutlass.range_constexpr(MmaNPerEpiN): + for mma_m_in_epi in cutlass.range_constexpr( + MmaMPerEpiM ): - tRS_rD_slice[elem_idx] = tRS_rAcc_slice[elem_idx] + mma_n = (epi_n * MmaNPerEpiN) + mma_n_in_epi + mma_m = (epi_m * MmaMPerEpiM) + mma_m_in_epi + tRS_rD_slice = tRS_rD[ + (None, mma_m_in_epi, mma_n_in_epi) + ] + tRS_rAcc_slice = tRS_rAcc[(None, mma_m, mma_n)] + for elem_idx in cutlass.range_constexpr( + cute.size(tRS_rD_slice) + ): + tRS_rD_slice[elem_idx] = tRS_rAcc_slice[ + elem_idx + ] + + gmem_coord = (epi_m, epi_n) + if cutlass.const_expr(self.split_k_slices > 1): + acc_mn = _reshape_acc_to_mn(accumulators) + c_identity = cute.make_identity_tensor( + cute.slice_(self.tile_shape_mnk, (None, None, 0)) + ) + coord_mn = _reshape_acc_to_mn( + thr_mma.partition_C(c_identity) + ) + if cutlass.const_expr(self.split_k_atomic_bf16): + for acc_m in cutlass.range_constexpr( + cute.size(acc_mn.shape[0]) + ): + for acc_n_pair in cutlass.range_constexpr( + cute.size(acc_mn.shape[1]) // 2 + ): + acc_n0 = acc_n_pair * 2 + acc_n1 = acc_n0 + 1 + coord0 = coord_mn[acc_m, acc_n0] + coord1 = coord_mn[acc_m, acc_n1] + m_coord0 = ( + tile_coord_mnl[0] + * Int32(self.tile_shape_mnk[0]) + + coord0[0] + ) + n_coord0 = ( + tile_coord_mnl[1] + * Int32(self.tile_shape_mnk[1]) + + coord0[1] + ) + m_coord1 = ( + tile_coord_mnl[0] + * Int32(self.tile_shape_mnk[0]) + + coord1[0] + ) + n_coord1 = ( + tile_coord_mnl[1] + * Int32(self.tile_shape_mnk[1]) + + coord1[1] + ) + if ( + m_coord0 < Int32(directC_mnl.shape[0]) + and m_coord1 + < Int32(directC_mnl.shape[0]) + and n_coord0 + < Int32(directC_mnl.shape[1]) + and n_coord1 + < Int32(directC_mnl.shape[1]) + ): + c_offset = cute.crd2idx( + ( + m_coord0, + n_coord0, + Int32(0), + ), + directC_mnl.layout, + ) + scatter_add_bf16x2( + get_ptr_as_int64( + directC_mnl, + c_offset, + ), + alpha_value * acc_mn[acc_m, acc_n0], + alpha_value * acc_mn[acc_m, acc_n1], + ) + if cutlass.const_expr( + cute.size(acc_mn.shape[1]) % 2 == 1 + ): + acc_n = cute.size(acc_mn.shape[1]) - 1 + coord = coord_mn[acc_m, acc_n] + m_coord = ( + tile_coord_mnl[0] + * Int32(self.tile_shape_mnk[0]) + + coord[0] + ) + n_coord = ( + tile_coord_mnl[1] + * Int32(self.tile_shape_mnk[1]) + + coord[1] + ) + if m_coord < Int32( + directC_mnl.shape[0] + ) and n_coord < Int32(directC_mnl.shape[1]): + c_offset = cute.crd2idx( + ( + m_coord, + n_coord, + Int32(0), + ), + directC_mnl.layout, + ) + scatter_add_bf16( + get_ptr_as_int64( + directC_mnl, + c_offset, + ), + alpha_value * acc_mn[acc_m, acc_n], + ) + else: + split_idx = Int32(block_idx[1]) + for acc_m in cutlass.range_constexpr( + cute.size(acc_mn.shape[0]) + ): + for acc_n in cutlass.range_constexpr( + cute.size(acc_mn.shape[1]) + ): + coord = coord_mn[acc_m, acc_n] + m_coord = ( + tile_coord_mnl[0] + * Int32(self.tile_shape_mnk[0]) + + coord[0] + ) + n_coord = ( + tile_coord_mnl[1] + * Int32(self.tile_shape_mnk[1]) + + coord[1] + ) + if m_coord < Int32( + directC_mnl.shape[0] + ) and n_coord < Int32(directC_mnl.shape[1]): + directC_mnl[ + (m_coord, n_coord, split_idx) + ] = alpha_value * acc_mn[acc_m, acc_n] + else: + # Type conversion with alpha scaling + tRS_rD_out = cute.make_rmem_tensor( + tRS_rD_layout.shape, self.c_dtype + ) + acc_vec = tRS_rD.load() + # Multiply alpha in FP32 before converting to c_dtype + # to avoid overflow when c_dtype is FP16 + acc_vec = epilogue_op( + (alpha_value * acc_vec).to(self.c_dtype) + ) + tRS_rD_out.store(acc_vec) - # Type conversion with alpha scaling - tRS_rD_out = cute.make_rmem_tensor( - tRS_rD_layout.shape, self.c_dtype - ) - acc_vec = tRS_rD.load() - # Multiply alpha in FP32 before converting to c_dtype - # to avoid overflow when c_dtype is FP16 - acc_vec = epilogue_op((alpha_value * acc_vec).to(self.c_dtype)) - tRS_rD_out.store(acc_vec) - - # Register to shared memory - epi_buffer = (epi_m * epi_rest_n + epi_n) % cute.size( - tRS_sD, mode=[3] - ) - if has_multi_epi_store: - self.epilog_sync_barrier.arrive_and_wait() - cute.copy( - tiled_copy_r2s, - tRS_rD_out, - tRS_sD[(None, None, None, epi_buffer)], - ) - cute.arch.fence_proxy( - "async.shared", - space="cta", + # Register to shared memory + epi_buffer = (epi_m * epi_rest_n + epi_n) % cute.size( + tRS_sD, mode=[3] + ) + if has_multi_epi_store: + self.epilog_sync_barrier.arrive_and_wait() + cute.copy( + tiled_copy_r2s, + tRS_rD_out, + tRS_sD[(None, None, None, epi_buffer)], + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # Copy from shared memory to global memory + if cutlass.const_expr(self.use_m1_non_tma_c): + for n_iter in cutlass.range_constexpr( + ( + self.epi_tile[1] + + self.num_mma_warps + * self.num_threads_per_warp + - 1 + ) + // ( + self.num_mma_warps + * self.num_threads_per_warp + ) + ): + n_local = Int32(tidx) + Int32( + n_iter + * self.num_mma_warps + * self.num_threads_per_warp + ) + n_coord = ( + tile_coord_mnl[1] + * Int32(self.tile_shape_mnk[1]) + + Int32(epi_n * self.epi_tile[1]) + + n_local + ) + if n_local < Int32( + self.epi_tile[1] + ) and n_coord < Int32(directC_mnl.shape[1]): + directC_mnl[ + ( + Int32(0), + n_coord, + tile_coord_mnl[2], + ) + ] = sC[(Int32(0), n_local, epi_buffer)] + else: + if warp_idx == 0: + cute.copy( + tma_atom_c, + bSG_sD[(None, epi_buffer)], + bSG_gD[(None, gmem_coord)], + ) + if has_multi_epi_store: + tma_store_pipeline.producer_commit() + tma_store_pipeline.producer_acquire() + + # Advance to the next work tile + if cutlass.const_expr(self.single_work_tile_per_cta): + work_tile = WorkTileInfo( + work_tile.tile_idx, + cutlass.Boolean(0), ) - self.epilog_sync_barrier.arrive_and_wait() - - # Copy from shared memory to global memory - gmem_coord = (epi_m, epi_n) - if warp_idx == 0: - cute.copy( - tma_atom_c, - bSG_sD[(None, epi_buffer)], - bSG_gD[(None, gmem_coord)], - ) - if has_multi_epi_store: - tma_store_pipeline.producer_commit() - tma_store_pipeline.producer_acquire() + else: + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + if has_multi_epi_store and cutlass.const_expr( + self.split_k_slices == 1 + ): + tma_store_pipeline.producer_tail() - # Advance to the next work tile - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() - if has_multi_epi_store: - tma_store_pipeline.producer_tail() - - # DMA warp group elif warp_idx == self.tma_load_warp_id: cute.arch.setmaxregister_decrease(self.load_register_requirement) while work_tile.is_valid_tile: tile_coord_mnl = work_tile.tile_idx - tAgA_mkl = tAgA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])] - tBgB_nkl = tBgB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])] - sfa_tile_coord_m = tile_coord_mnl[0] // self.sfa_tiles_per_block - tAgSFA_mkl = tAgSFA[(None, sfa_tile_coord_m, None, tile_coord_mnl[2])] - sfb_tile_coord_n = tile_coord_mnl[1] // self.sfb_tiles_per_block - tBgSFB_nkl = tBgSFB[(None, sfb_tile_coord_n, None, tile_coord_mnl[2])] + if cutlass.const_expr( + self.load_path == "tma" and not self.use_m1_non_tma_a + ): + tAgA_mkl = tAgA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])] + if cutlass.const_expr(self.load_path == "tma"): + tBgB_nkl = tBgB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])] + if cutlass.const_expr( + self.load_path == "tma" and not self.use_m1_non_tma_sfa + ): + sfa_tile_coord_m = tile_coord_mnl[0] // self.sfa_tiles_per_block + tAgSFA_mkl = tAgSFA[ + (None, sfa_tile_coord_m, None, tile_coord_mnl[2]) + ] + if cutlass.const_expr(self.load_path == "tma"): + sfb_tile_coord_n = tile_coord_mnl[1] // self.sfb_tiles_per_block + tBgSFB_nkl = tBgSFB[ + (None, sfb_tile_coord_n, None, tile_coord_mnl[2]) + ] + if cutlass.const_expr(self.load_path == "cpasync"): + cpasync_sfa_tile_coord_m = ( + tile_coord_mnl[0] // self.sfa_tiles_per_block + ) + cpasync_sfb_tile_coord_n = ( + tile_coord_mnl[1] // self.sfb_tiles_per_block + ) mainloop_producer_state.reset_count() - for _k_tile in range(0, k_tile_cnt, 1, unroll=2): # type: ignore[call-overload] + for _k_tile in range(0, k_tile_iter_cnt, 1, unroll=2): # type: ignore[call-overload] mainloop_pipeline.producer_acquire(mainloop_producer_state) - tAgA_k = tAgA_mkl[(None, mainloop_producer_state.count)] - tAsA_pipe = tAsA[(None, mainloop_producer_state.index)] - - tBgB_k = tBgB_nkl[(None, mainloop_producer_state.count)] - tBsB_pipe = tBsB[(None, mainloop_producer_state.index)] - - tAgSFA_k = tAgSFA_mkl[(None, mainloop_producer_state.count)] - tAsSFA_pipe = tAsSFA[(None, mainloop_producer_state.index)] - - tBgSFB_k = tBgSFB_nkl[(None, mainloop_producer_state.count)] - tBsSFB_pipe = tBsSFB[(None, mainloop_producer_state.index)] + k_tile_global = k_tile_start + mainloop_producer_state.count + if cutlass.const_expr(self.load_path == "tma"): + tBgB_k = tBgB_nkl[(None, k_tile_global)] + tBsB_pipe = tBsB[(None, mainloop_producer_state.index)] + if cutlass.const_expr(not self.use_m1_non_tma_a): + tAgA_k = tAgA_mkl[(None, k_tile_global)] + tAsA_pipe = tAsA[(None, mainloop_producer_state.index)] + + tAgSFA_k = tAgSFA_mkl[(None, k_tile_global)] + tAsSFA_pipe = tAsSFA[(None, mainloop_producer_state.index)] + + tBgSFB_k = tBgSFB_nkl[(None, k_tile_global)] + tBsSFB_pipe = tBsSFB[(None, mainloop_producer_state.index)] + + if cutlass.const_expr(self.load_path == "cpasync"): + tAgA_cpasync_k = tAgA_cpasync_mkl[ + ( + None, + None, + None, + tile_coord_mnl[0], + k_tile_global, + tile_coord_mnl[2], + ) + ] + tAsA_cpasync_pipe = tAsA_cpasync[ + (None, None, None, mainloop_producer_state.index) + ] + tAcA_cpasync_k = cute.slice_( + tAcA_cpasync_mkl, + ( + None, + None, + None, + tile_coord_mnl[0], + k_tile_global, + tile_coord_mnl[2], + ), + ) + tBgB_cpasync_k = tBgB_cpasync_nkl[ + ( + None, + None, + None, + tile_coord_mnl[1], + k_tile_global, + tile_coord_mnl[2], + ) + ] + tBsB_cpasync_pipe = tBsB_cpasync[ + (None, None, None, mainloop_producer_state.index) + ] + tBcB_cpasync_k = cute.slice_( + tBcB_cpasync_nkl, + ( + None, + None, + None, + tile_coord_mnl[1], + k_tile_global, + tile_coord_mnl[2], + ), + ) + tAgSFA_cpasync_k = cute.filter_zeros( + tAgSFA_cpasync_mkl[ + ( + None, + None, + None, + cpasync_sfa_tile_coord_m, + k_tile_global, + tile_coord_mnl[2], + ) + ] + ) + tAsSFA_cpasync_pipe = cute.filter_zeros( + tAsSFA_cpasync[ + (None, None, None, mainloop_producer_state.index) + ] + ) + tAcSFA_cpasync_k = cute.filter_zeros( + cute.slice_( + tAcSFA_cpasync_mkl, + ( + None, + None, + None, + cpasync_sfa_tile_coord_m, + k_tile_global, + tile_coord_mnl[2], + ), + ) + ) + tBgSFB_cpasync_k = cute.filter_zeros( + tBgSFB_cpasync_nkl[ + ( + None, + None, + None, + cpasync_sfb_tile_coord_n, + k_tile_global, + tile_coord_mnl[2], + ) + ] + ) + tBsSFB_cpasync_pipe = cute.filter_zeros( + tBsSFB_cpasync[ + (None, None, None, mainloop_producer_state.index) + ] + ) + tBcSFB_cpasync_k = cute.filter_zeros( + cute.slice_( + tBcSFB_cpasync_nkl, + ( + None, + None, + None, + cpasync_sfb_tile_coord_n, + k_tile_global, + tile_coord_mnl[2], + ), + ) + ) + self._cpasync_copy_2d( + cpasync_tiled_copy_A, + tAgA_cpasync_k, + tAsA_cpasync_pipe, + tAcA_cpasync_k, + Int32(directA_mkl.shape[0]), + True, + ) + self._cpasync_copy_2d( + cpasync_tiled_copy_B, + tBgB_cpasync_k, + tBsB_cpasync_pipe, + tBcB_cpasync_k, + Int32(directC_mnl.shape[1]), + True, + ) + self._scale_copy_2d( + cpasync_tiled_copy_SF, + tAgSFA_cpasync_k, + tAsSFA_cpasync_pipe, + tAcSFA_cpasync_k, + Int32(directA_mkl.shape[0]), + ) + self._scale_copy_2d( + cpasync_tiled_copy_SF, + tBgSFB_cpasync_k, + tBsSFB_cpasync_pipe, + tBcSFB_cpasync_k, + Int32(directC_mnl.shape[1]), + ) + cute.arch.fence_proxy("async.shared", space="cta") + elif cutlass.const_expr(self.use_m1_non_tma_a): + lane = Int32(tidx % self.num_threads_per_warp) + for a_iter in cutlass.range_constexpr( + (self.tile_shape_mnk[2] + self.num_threads_per_warp - 1) + // self.num_threads_per_warp + ): + k_local = lane + Int32(a_iter * self.num_threads_per_warp) + if k_local < Int32(self.tile_shape_mnk[2]): + k_coord = ( + k_tile_global * Int32(self.tile_shape_mnk[2]) + + k_local + ) + sA[ + ( + Int32(0), + k_local, + mainloop_producer_state.index, + ) + ] = directA_mkl[ + ( + Int32(0), + k_coord, + tile_coord_mnl[2], + ) + ] + else: + cute.copy( + tma_atom_a, + tAgA_k, + tAsA_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) - cute.copy( - tma_atom_a, - tAgA_k, - tAsA_pipe, - tma_bar_ptr=mainloop_pipeline.producer_get_barrier( - mainloop_producer_state - ), - ) - cute.copy( - tma_atom_b, - tBgB_k, - tBsB_pipe, - tma_bar_ptr=mainloop_pipeline.producer_get_barrier( - mainloop_producer_state - ), - ) - cute.copy( - tma_atom_sfa, - tAgSFA_k, - tAsSFA_pipe, - tma_bar_ptr=mainloop_pipeline.producer_get_barrier( - mainloop_producer_state - ), - ) - cute.copy( - tma_atom_sfb, - tBgSFB_k, - tBsSFB_pipe, - tma_bar_ptr=mainloop_pipeline.producer_get_barrier( - mainloop_producer_state - ), - ) + if cutlass.const_expr(self.load_path == "cpasync"): + pass + elif cutlass.const_expr(self.use_m1_non_tma_sfa): + lane = Int32(tidx % self.num_threads_per_warp) + scale_groups_per_k_tile = ( + self.tile_shape_mnk[2] // self.sf_vec_size + ) + sfa_slots = self.sfa_tile_shape_mk[0] * scale_groups_per_k_tile + for sfa_iter in cutlass.range_constexpr( + (sfa_slots + self.num_threads_per_warp - 1) + // self.num_threads_per_warp + ): + linear = lane + Int32(sfa_iter * self.num_threads_per_warp) + m_local = linear // Int32(scale_groups_per_k_tile) + scale_group = linear - m_local * Int32( + scale_groups_per_k_tile + ) + k_local_sfa = scale_group * Int32(self.sf_vec_size) + k_coord_sfa = ( + k_tile_global * Int32(self.tile_shape_mnk[2]) + + k_local_sfa + ) + if linear < Int32(sfa_slots): + sSFA[ + ( + m_local, + k_local_sfa, + mainloop_producer_state.index, + ) + ] = directSFA_mkl[ + ( + Int32(0), + k_coord_sfa, + tile_coord_mnl[2], + ) + ] + cute.arch.fence_proxy("async.shared", space="cta") + else: + cute.copy( + tma_atom_sfa, + tAgSFA_k, + tAsSFA_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + if cutlass.const_expr(self.load_path == "tma"): + cute.copy( + tma_atom_b, + tBgB_k, + tBsB_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + cute.copy( + tma_atom_sfb, + tBgSFB_k, + tBsSFB_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + if cutlass.const_expr(self.load_path == "cpasync"): + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) mainloop_pipeline.producer_commit(mainloop_producer_state) mainloop_producer_state.advance() - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() + if cutlass.const_expr(self.single_work_tile_per_cta): + work_tile = WorkTileInfo( + work_tile.tile_idx, + cutlass.Boolean(0), + ) + else: + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() mainloop_pipeline.producer_tail(mainloop_producer_state) return @@ -1252,12 +2197,14 @@ def _compute_stages( ) mbar_helpers_bytes = 1024 - ab_stage = ( + raw_ab_stage = ( (smem_capacity - occupancy * 1024) // occupancy - mbar_helpers_bytes - epi_bytes ) // (ab_bytes_per_stage + sf_bytes_per_stage) - ab_stage = max(1, min(ab_stage, 4)) + ab_stage = max(1, min(raw_ab_stage, 4)) + if tile_shape_mnk[0] in (16, 64) and tile_shape_mnk[1] == 128: + ab_stage = max(1, min(raw_ab_stage, 5)) return ab_stage, epi_stage @staticmethod @@ -1353,6 +2300,8 @@ def _compute_grid( c, tile_shape_mnk: tuple, max_active_clusters, + direct_one_m_tile_scheduler: bool, + split_k_slices: int, ) -> tuple: c_shape = cute.slice_(tile_shape_mnk, (None, None, 0)) gc = cute.zipped_divide(c, tiler=c_shape) @@ -1361,9 +2310,12 @@ def _compute_grid( tile_sched_params = utils.PersistentTileSchedulerParams( num_ctas_mnl, cluster_shape_mnl ) - grid = utils.StaticPersistentTileScheduler.get_grid_shape( - tile_sched_params, max_active_clusters - ) + if cutlass.const_expr(split_k_slices > 1): + grid = (1, split_k_slices, num_ctas_mnl[1]) + else: + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) return tile_sched_params, grid @staticmethod @@ -1413,33 +2365,50 @@ def can_implement( c_dtype, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], - m: int, n: int, k: int, l: int, a_major: str, b_major: str, c_major: str, + *, + load_path: str = "tma", + swap_ab: bool = False, ) -> bool: # The current target only supports cluster (1,1) if cluster_shape_mn != (1, 1): return False - # Tile M must be divisible by 128; tile N follows 64-column warpgroup - # quanta, while the SF paths round narrow tiles up to full 128-element - # scale-factor blocks. - if mma_tiler_mn[0] % 64 != 0 or mma_tiler_mn[1] % 64 != 0: + if load_path not in _DENSE_LOAD_PATHS: return False - # The current target only supports FP4 (MmaMXF4NVF4Op) + # FP4-only target (NVF4/MXF4). The MXFP8 warp-MMA path was dropped. if ab_dtype != cutlass.Float4E2M1FN: return False - # SM120 warp-level MmaMXF4NVF4Op only supports sf_vec_size=16 - # (CUTLASS DSL hardcodes sf_vec_size=16 in the MMA atom constructor) - if sf_vec_size != 16: + if swap_ab: + if l != 1: + return False + if sf_vec_size != 16: + return False + if load_path == "cpasync" and (sf_vec_size != 16 or l != 1): + return False + # FP4 experiments allow narrow N tiles. The scale-factor smem paths + # still allocate full 128-element SF blocks, but the live MMA tile may + # consume only 16/32 columns. + if ( + mma_tiler_mn[0] % 64 != 0 + or mma_tiler_mn[1] % 16 != 0 + or mma_tiler_mn[1] > 128 + or (mma_tiler_mn[1] < 64 and not swap_ab) + ): return False - if sf_dtype != cutlass.Float8E4M3FN: + # Current target MMA constraints: + # sf_vec_size=16 requires sf_dtype=Float8E4M3FN + # sf_vec_size=32 requires sf_dtype=Float8E8M0FNU + if sf_vec_size == 16 and sf_dtype != cutlass.Float8E4M3FN: return False - # Only 16-bit output types supported for now - if c_dtype not in (cutlass.Float16, cutlass.BFloat16): + if sf_vec_size == 32 and sf_dtype != cutlass.Float8E8M0FNU: + return False + # Public output is 16-bit; split-K internally uses FP32 partial output. + if c_dtype not in (cutlass.Float16, cutlass.BFloat16, cutlass.Float32): return False # A must be K-major, B must be K-major if a_major != "k" or b_major != "k": @@ -1448,25 +2417,8 @@ def can_implement( tile_k = sf_vec_size * 8 if k % tile_k != 0: return False - # Reject tiles that cannot fit even one pipeline stage in SM120 - # shared memory. A+B are FP4 (0.5 bytes/element), SF blocks are - # rounded up to 128-element granularity, epilogue is 16-bit output. - sfa_tile_m = max(128, ((mma_tiler_mn[0] + 127) // 128) * 128) - sfb_tile_n = max(128, ((mma_tiler_mn[1] + 127) // 128) * 128) - ab_bytes = (mma_tiler_mn[0] * tile_k + mma_tiler_mn[1] * tile_k) // 2 - # SF: 128 * 4 elements per SF block, 1 byte each - sf_bytes = (sfa_tile_m // 128) * 4 * 128 + (sfb_tile_n // 128) * 4 * 128 - epi_bytes = mma_tiler_mn[0] * mma_tiler_mn[1] * 2 # 16-bit output - mbar_bytes = 1024 - smem_capacity = utils.get_smem_capacity_in_bytes("sm_120") - if ab_bytes + sf_bytes + epi_bytes + mbar_bytes > smem_capacity: - return False return True - # ------------------------------------------------------------------ - # wrapper: compile-time entry point matching the SM100 interface - # for FlashInfer's _compile_block_scaled_gemm - # ------------------------------------------------------------------ @cute.jit def wrapper( self, @@ -1511,16 +2463,12 @@ def wrapper( b_ptr, layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2)), ) - if cutlass.const_expr(swap_ab): - c_tensor = cute.make_tensor( - mC.iterator, - layout=cute.make_ordered_layout((m, n, l), order=(0, 1, 2)), - ) - else: - c_tensor = cute.make_tensor( - mC.iterator, - layout=cute.make_ordered_layout((m, n, l), order=(1, 0, 2)), - ) + # C is always row-major (m, n): b12x swap_ab is device-internal (the + # epilogue writes logical [m, n]), so this `swap_ab` arg is unused here. + c_tensor = cute.make_tensor( + mC.iterator, + layout=cute.make_ordered_layout((m, n, l), order=(1, 0, 2)), + ) sfa_tensor = cute.make_tensor( a_sf_ptr, layout=cute.make_ordered_layout( @@ -1549,269 +2497,48 @@ def wrapper( ) -# Alias for FlashInfer integration -Sm120B12xBlockScaledDenseGemmKernel = DenseGemmKernel +_ALPHA_ONE_CACHE: dict = {} -class _DenseGemmLaunch: - def __init__( - self, - m: int, - n: int, - k: int, - l: int, - a_major: str, - b_major: str, - c_major: str, - ab_dtype: torch.dtype, - sf_dtype: torch.dtype, - c_dtype: torch.dtype, - alpha_dtype: torch.dtype, - sf_vec_size: int, - mma_tiler_mn: Tuple[int, int], - cluster_shape_mn: Tuple[int, int], - sm_count: int, - sm_version: str, - ): - self._m = m - self._n = n - self._k = k - self._l = l - self._a_major = a_major - self._b_major = b_major - self._c_major = c_major - self._ab_dtype = ab_dtype - self._sf_dtype = sf_dtype - self._c_dtype = c_dtype - self._alpha_dtype = alpha_dtype - self._sf_vec_size = sf_vec_size - self._mma_tiler_mn = mma_tiler_mn - self._cluster_shape_mn = cluster_shape_mn - - if sm_version not in ("sm_120", "sm_121"): - raise ValueError( - f"dense_gemm launch only supports SM12x (sm_120/sm_121), got {sm_version}" - ) - - if not DenseGemmKernel.can_implement( - ab_dtype, - sf_dtype, - sf_vec_size, - c_dtype, - mma_tiler_mn, - cluster_shape_mn, - m, - n, - k, - l, - a_major, - b_major, - c_major, - ): - raise TypeError( - "dense_gemm launch is unsupported with " - f"{ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, " - f"{mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, " - f"{a_major}, {b_major}, {c_major}" - ) - - self._max_active_clusters = min( - get_max_active_clusters( - self._cluster_shape_mn[0] * self._cluster_shape_mn[1] - ), - sm_count, - ) - - @cute.jit - def __call__( - self, - a_ptr: cute.Pointer, - b_ptr: cute.Pointer, - sfa_ptr: cute.Pointer, - sfb_ptr: cute.Pointer, - c_ptr: cute.Pointer, - alpha_ptr: cute.Pointer, - current_stream: cuda.CUstream, - ): - a_tensor = cute.make_tensor( - a_ptr, - layout=cute.make_ordered_layout( - (self._m, self._k, self._l), - order=(0, 1, 2) if self._a_major == "m" else (1, 0, 2), - ), - ) - b_tensor = cute.make_tensor( - b_ptr, - layout=cute.make_ordered_layout( - (self._n, self._k, self._l), - order=(0, 1, 2) if self._b_major == "n" else (1, 0, 2), - ), - ) - c_tensor = cute.make_tensor( - c_ptr, - layout=cute.make_ordered_layout( - (self._m, self._n, self._l), - order=(0, 1, 2) if self._c_major == "m" else (1, 0, 2), - ), - ) - alpha_tensor = cute.make_tensor( - alpha_ptr, - layout=cute.make_ordered_layout((1,), order=(0,)), - ) - sfa_tensor = cute.make_tensor(sfa_ptr, layout=cute.make_layout((1,))) - sfb_tensor = cute.make_tensor(sfb_ptr, layout=cute.make_layout((1,))) - - DenseGemmKernel( - sf_vec_size=self._sf_vec_size, - mma_tiler_mn=self._mma_tiler_mn, - cluster_shape_mn=self._cluster_shape_mn, - )( - a_tensor, - b_tensor, - sfa_tensor, - sfb_tensor, - c_tensor, - alpha_tensor, - self._max_active_clusters, - current_stream, - ) - - -@functools.cache -def _get_compiled_dense_gemm( +def _select_default_mma_tiler_mn( m: int, n: int, - k: int, - l: int, - a_major: str, - b_major: str, - c_major: str, - ab_dtype: Type[cutlass.Numeric], - sf_dtype: Type[cutlass.Numeric], - c_dtype: Type[cutlass.Numeric], - alpha_dtype: Type[cutlass.Numeric], - sf_vec_size: int, - mma_tiler_mn: Tuple[int, int], - cluster_shape_mn: Tuple[int, int], sm_count: int, - sm_version: str, -) -> Callable: - def _make_runtime_pointers( - input_tensors: Optional[List[torch.Tensor]], - ) -> List[cute.Pointer]: - if input_tensors is None: - ( - a_data_ptr, - b_data_ptr, - sfa_data_ptr, - sfb_data_ptr, - c_data_ptr, - alpha_data_ptr, - ) = [16 for _ in range(6)] - else: - ( - a_tensor_gpu, - b_tensor_gpu, - sfa_tensor_gpu, - sfb_tensor_gpu, - c_tensor_gpu, - alpha_tensor_gpu, - ) = input_tensors - ( - a_data_ptr, - b_data_ptr, - sfa_data_ptr, - sfb_data_ptr, - c_data_ptr, - alpha_data_ptr, - ) = ( - a_tensor_gpu.data_ptr(), - b_tensor_gpu.data_ptr(), - sfa_tensor_gpu.data_ptr(), - sfb_tensor_gpu.data_ptr(), - c_tensor_gpu.data_ptr(), - alpha_tensor_gpu.data_ptr(), - ) - - return [ - make_ptr(ab_dtype, a_data_ptr, cute.AddressSpace.gmem, assumed_align=16), - make_ptr(ab_dtype, b_data_ptr, cute.AddressSpace.gmem, assumed_align=16), - make_ptr(sf_dtype, sfa_data_ptr, cute.AddressSpace.gmem, assumed_align=16), - make_ptr(sf_dtype, sfb_data_ptr, cute.AddressSpace.gmem, assumed_align=16), - make_ptr(c_dtype, c_data_ptr, cute.AddressSpace.gmem, assumed_align=16), - make_ptr( - alpha_dtype, alpha_data_ptr, cute.AddressSpace.gmem, assumed_align=16 - ), - ] - - compiled_kernel = cute.compile( - _DenseGemmLaunch( - m=m, - n=n, - k=k, - l=l, - a_major=a_major, - b_major=b_major, - c_major=c_major, - ab_dtype=ab_dtype, - sf_dtype=sf_dtype, - c_dtype=c_dtype, - alpha_dtype=alpha_dtype, - sf_vec_size=sf_vec_size, - mma_tiler_mn=mma_tiler_mn, - cluster_shape_mn=cluster_shape_mn, - sm_count=sm_count, - sm_version=sm_version, - ), - *_make_runtime_pointers(None), - current_cuda_stream(), - ) - - def tensor_api( - a_tensor_gpu: torch.Tensor, - b_tensor_gpu: torch.Tensor, - sfa_tensor_gpu: torch.Tensor, - sfb_tensor_gpu: torch.Tensor, - c_tensor_gpu: Optional[torch.Tensor] = None, - alpha_tensor_gpu: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if c_tensor_gpu is None: - c_tensor_gpu = torch.empty( - (m, n, l), - dtype=cutlass_to_torch_dtype(c_dtype), - device=a_tensor_gpu.device, - ) - if alpha_tensor_gpu is None: - alpha_tensor_gpu = torch.ones( - (1,), - dtype=torch.float32, - device=a_tensor_gpu.device, - ) - - nonlocal compiled_kernel - compiled_kernel( - *_make_runtime_pointers( - [ - a_tensor_gpu, - b_tensor_gpu, - sfa_tensor_gpu, - sfb_tensor_gpu, - c_tensor_gpu, - alpha_tensor_gpu, - ] - ), - current_cuda_stream(), - ) - return c_tensor_gpu - - return tensor_api - - -def _select_default_mma_tiler_mn(m: int, n: int, sm_count: int) -> Tuple[int, int]: + *, + expected_m: Optional[int] = None, + k: Optional[int] = None, +) -> Tuple[int, int]: + # FP4-only tile selector. The MXFP8 regimes (16x64/16x128/32x128 decode + # tiles, narrow-N 64x64) were dropped along with the MXFP8 warp-MMA path. coarse_tile = (128, 128) + plan_m = expected_m if expected_m is not None else m + if plan_m == 1 and k is not None: + # Flushed M=1 FP4 probe (benchmarks/probe_dense_fp4_tile_load_sweep.py) + # across the repo's common shapes: + # * wide/medium N: (64,128)/TMA has the best geomean and wins nearly all + # shapes. + # * N=1024,K=5376: (64,64)/TMA wins the boundary by a small margin. + # * N<=512 with long K: (64,32)/TMA+swap_ab is the only clear tiny-N win. + # Keep the tile selector tile-only; the launch planner below attaches + # swap_ab to the narrow tile. + if n <= 512 and k >= 4096: + return (64, 32) + if n <= 1024: + return (64, 64) + return (64, 128) + coarse_tiles = ((m + coarse_tile[0] - 1) // coarse_tile[0]) * ( (n + coarse_tile[1] - 1) // coarse_tile[1] ) + # The coarse CTA-count heuristic misses exact-small-M, wide-N cases: a wide + # N dimension can generate plenty of CTAs even while each 128-row M tile is + # mostly empty. Keep using the narrower 64x128 tile while the 128x128 plan + # still leaves the GPU below the existing half-SM occupancy proxy. + if n > 1536: + if m <= 64: + return (64, 128) + if m <= 256 and coarse_tiles < max(1, sm_count // 2): + return (64, 128) if m <= 128 and coarse_tiles < max(1, sm_count // 2): if n > 1536: return (64, 128) @@ -1822,62 +2549,30 @@ def _select_default_mma_tiler_mn(m: int, n: int, sm_count: int) -> Tuple[int, in if medium_tiles < max(1, sm_count // 2): return (64, 64) return (128, 64) - return (128, 128) + return coarse_tile -def dense_gemm( - lhs: Tuple[torch.Tensor, torch.Tensor], - rhs: Tuple[torch.Tensor, torch.Tensor], - out: Optional[torch.Tensor] = None, +def _select_default_dense_gemm_plan( + m: int, + n: int, + k: int, + sm_count: int, *, - ab_dtype: str, - sf_dtype: str, - c_dtype: str, - sf_vec_size: int, - sm_count: Optional[int] = None, - mma_tiler_mn: Optional[Tuple[int, int]] = None, - cluster_shape_mn: Tuple[int, int] = (1, 1), - alpha: Optional[torch.Tensor] = None, - alpha_dtype: Optional[str] = None, -) -> torch.Tensor: - """Execute dense block-scaled GEMM for one expert-major batch stack.""" - a_torch, sfa_torch = lhs - b_torch, sfb_torch = rhs - - m, k, l = a_torch.shape - n, _, _ = b_torch.shape - if ab_dtype == "float4_e2m1fn": - k *= 2 - - if sm_count is None: - sm_count = get_num_sm(a_torch.device) - if mma_tiler_mn is None: - mma_tiler_mn = _select_default_mma_tiler_mn(m, n, sm_count) - if alpha_dtype is None: - alpha_dtype = "float32" if alpha is None else str(alpha.dtype).split(".")[-1] - - return _get_compiled_dense_gemm( - m=m, - n=n, + expected_m: Optional[int] = None, +) -> _DenseGemmPlan: + tile = _select_default_mma_tiler_mn( + m, + n, + sm_count, + expected_m=expected_m, k=k, - l=l, - a_major="k", - b_major="k", - c_major="n", - ab_dtype=get_cutlass_dtype(ab_dtype), - sf_dtype=get_cutlass_dtype(sf_dtype), - c_dtype=get_cutlass_dtype(c_dtype), - alpha_dtype=get_cutlass_dtype(alpha_dtype), - sf_vec_size=sf_vec_size, - mma_tiler_mn=mma_tiler_mn, - cluster_shape_mn=cluster_shape_mn, - sm_count=sm_count, - sm_version="sm_120", - )( - a_tensor_gpu=a_torch, - b_tensor_gpu=b_torch, - sfa_tensor_gpu=sfa_torch, - sfb_tensor_gpu=sfb_torch, - c_tensor_gpu=out, - alpha_tensor_gpu=alpha, ) + return _DenseGemmPlan( + mma_tiler_mn=tile, + load_path="tma", + swap_ab=(tile[1] < 64), + ) + + +# Alias for FlashInfer integration +Sm120B12xBlockScaledDenseGemmKernel = DenseGemmKernel From 6044fc5ebfcabc70816b592403850739e8879598 Mon Sep 17 00:00:00 2001 From: Rehvaro Date: Thu, 11 Jun 2026 09:40:29 +0200 Subject: [PATCH 13/13] jit: run `ninja -t restat` before builds to fix spurious full rebuilds on GPFS/NFS On shared filesystems (GPFS, NFS), the mtimes ninja records in .ninja_log at the end of a build can disagree with what a later stat() returns (inode flush timing). On the next process start, ninja sees the mismatch, considers every output "modified externally", and recompiles all CUTLASS fused-MoE kernels (~180 .cu files, ~1.5h on GH200/sm90a) on every vLLM restart. Fix: run `ninja -t restat` (available since ninja 1.10) right before run_ninja(). It rewrites the mtimes stored in .ninja_log with current stat() values, so the incremental check passes, while ninja's staleness detection stays fully intact: modified sources or changed command lines still trigger rebuilds. A restat failure is non-fatal (check=False) and falls back to the current behavior. Replace PR #3556. --- flashinfer/jit/core.py | 8 +++++++- flashinfer/jit/cpp_ext.py | 25 +++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/flashinfer/jit/core.py b/flashinfer/jit/core.py index 53c4163e267..17bed954147 100644 --- a/flashinfer/jit/core.py +++ b/flashinfer/jit/core.py @@ -12,7 +12,12 @@ from ..compilation_context import CompilationContext from . import env as jit_env -from .cpp_ext import generate_ninja_build_for_op, get_nvcc_parallelism_flags, run_ninja +from .cpp_ext import ( + generate_ninja_build_for_op, + get_nvcc_parallelism_flags, + run_ninja, + run_ninja_restat, +) from .utils import write_if_different os.makedirs(jit_env.FLASHINFER_WORKSPACE_DIR, exist_ok=True) @@ -299,6 +304,7 @@ def build(self, verbose: bool, need_lock: bool = True) -> None: ) with lock: self.write_ninja() + run_ninja_restat(self.build_dir, self.ninja_path, verbose) run_ninja(self.build_dir, self.ninja_path, verbose) def load(self, so_path: Path): diff --git a/flashinfer/jit/cpp_ext.py b/flashinfer/jit/cpp_ext.py index f2e5527b17e..c14375c5c20 100644 --- a/flashinfer/jit/cpp_ext.py +++ b/flashinfer/jit/cpp_ext.py @@ -348,6 +348,31 @@ def _get_num_workers() -> Optional[int]: return None +def run_ninja_restat(workdir: Path, ninja_file: Path, verbose: bool) -> None: + """Re-sync the mtimes recorded in .ninja_log with current stat() values, + so that mtime skew on shared filesystems (GPFS, NFS) does not make ninja + rebuild up-to-date outputs. Non-fatal: worst case is a rebuild.""" + if not (workdir / ".ninja_log").exists(): + return + command = [ + "ninja", + "-C", + str(workdir.resolve()), + "-f", + str(ninja_file.resolve()), + "-t", + "restat", + ] + subprocess.run( + command, + stdout=None if verbose else subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=str(workdir.resolve()), + check=False, + text=True, + ) + + def run_ninja(workdir: Path, ninja_file: Path, verbose: bool) -> None: workdir.mkdir(parents=True, exist_ok=True) command = [