diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index b8e722541..5f8ba64ad 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -6,8 +6,8 @@ tags: # finegrained-kernels -Triton GEMM + MoE kernels for fine-grained quantization — block/tensor FP8, MXFP8, MXFP4, and -two-level NVFP4 weights, with fused gate|up GLU epilogues, fused intermediate requantization, and +Triton GEMM + MoE kernels for fine-grained quantization — block/tensor/row-wise FP8, MXFP8, MXFP4, +two-level NVFP4, and INT2/INT4/INT8 weights, with fused gate|up GLU epilogues, fused intermediate requantization, and expert routing — developed as part of the HuggingFace Transformers FP8 + MoE optimization effort. All tuning and validation is on NVIDIA Blackwell (SM100 / B200), whose tcgen05 scaled-MMA the MX @@ -62,7 +62,7 @@ matmul_grouped(A, B, As, Bs, *, expert_start, gather_idx=None, scatter_idx=None, (see the format matrix below). - `activation_format` is the activations' quantization, in the same vocabulary as transformers' `FineGrainedConfig`: `None` (default) = the weights' own format, `"bf16"` = weight-only (W4A16 / - W8A16), or an explicit `"fp8"` / `"mxfp8"` / `"mxfp4"` / `"nvfp4"` (`"mxfp8"` on MXFP4 weights = + W8A16), or an explicit `"fp8"` / `"int8"` / `"mxfp8"` / `"mxfp4"` / `"nvfp4"` (`"mxfp8"` on MXFP4 weights = the W4A8 chain). The weight format is never named — it is read off `B`/`Bs`. - `gate=True, act_fn="silu", swiglu_alpha=..., swiglu_limit=...` fuses the gate|up GLU: `B` holds the stacked `(2N, K)` gate|up rows and the op returns the activated `(M, N)` intermediate; @@ -75,6 +75,8 @@ matmul_grouped(A, B, As, Bs, *, expert_start, gather_idx=None, scatter_idx=None, gate|up stack whose halves a checkpoint calibrated separately carries one global per half; the loader merges those into one per expert (the up half's folds into the down projection's global, SwiGLU being linear in it) rather than the ops taking a second layout. +- Weight-only FP8 (`activation_format="bf16"`) also accepts coarse block scales and one scale + per tensor/expert, with unquantized activations. Coarse FP8 entrypoints are inference-only. - Routing: `matmul_batched` reads one expert id per row (`expert_ids`, EP sentinels `>= num_experts` skipped); `matmul_grouped` takes the launch maps from `compute_grouped_scheduling(expert_ids, num_experts, top_k)` — no pre-sorted input required. @@ -86,9 +88,83 @@ matmul_grouped(A, B, As, Bs, *, expert_start, gather_idx=None, scatter_idx=None, | full precision | bf16 / fp16 | `None` | — | | block FP8 (128×128) | `float8_e4m3fn` | fp32 or UE8M0 `(⌈N/128⌉, K/128)` (the block is derived from the K dim, so K must divide) | — | | tensor FP8 | `float8_e4m3fn` | one scalar | — | +| row-wise FP8 | `float8_e4m3fn` | fp32 `(N, 1)`, one per output row | — | | MXFP8 | `float8_e4m3fn` | UE8M0 `(N, K/32)` (raw `uint8` accepted) | — | | MXFP4 | packed E2M1 (`int8`, 2/byte) | UE8M0 `(N, K/32)` | — | | NVFP4 | packed E2M1 (`int8`, 2/byte) | E4M3 `(N, K/16)` | fp32 global(s) | +| INT2 / INT4 / INT8 | interleaved `uint32` words (`32 / bits` values/word) | bf16 / fp16 / fp32 `(N, K/group)` | — | + +Integer weights use raw bf16/fp16 activations and optional `b_zero` values shaped like `Bs` +(default: `2 ** (bits - 1)`). Dequantization rounds `(q - zero) * scale` in the scale dtype, +then casts each weight to the activation dtype before multiplication, including per-channel +scales (`group=K`, one scale per output row). Group widths and K tiles must divide one another; +the K dimension must contain whole groups. Use `interleave_int_words` for naturally packed +words, `repack_gptq_int4` for GPTQ v1/v2 (including act-order), or `repack_awq_int4` for AWQ +GEMM checkpoints. Natural `int32` checkpoint words are rejected to prevent confusing the two layouts. +For ModelOpt's signed exports, `repack_modelopt_int4` reads `uint8` `(N/2, K)` with two +output channels per byte, and `repack_modelopt_int8` reads signed `int8` `(N, K)` with one +scale per output channel. Both also accept expert stacks and return `(weight, scale, None)` +for these same integer GEMMs. Pass an exported AWQ `pre_quant_scale` as `a_pre_quant_scale` +to an integer GEMM, or fuse it upstream. It multiplies the raw input columns before any +act-order permutation, rounded in the activation dtype. Dense scales are contiguous `(K,)`; +routed scales are shared `(K,)` / `(1, K)` or per expert `(E, K)`, on the activation's device +and in its dtype. MoE forwards accept `gate_up_proj_pre_quant_scale` and +`down_proj_pre_quant_scale`, including a separate down input vector per expert. A fused +gate/up stack uses a common input transform for both halves, as after ModelOpt's resmoothing. +For ModelOpt `W4A8_AWQ`, use `activation_format="fp8"`, `As=input_scale`, and +`b_global_scale=weight_scale_2` with the repacked symmetric INT4 weight and its original +group scales. The kernel reconstructs the weight, rounds it to E4M3 against `weight_scale_2`, +and restores both calibrated factors before FP32 bias/GLU. `activation_format="int8"` with +`As=input_scale` selects ModelOpt INT8 SmoothQuant on symmetric per-channel INT8 weights: +the dot accumulates in INT32, then applies per-channel weight scales and input scale in FP32. +Both recipes take raw activations and keep the intermediate unquantized between projections; +the down's own calibration and optional input transform apply before its quantization. + +ModelOpt `W4A8_MXFP4_FP8` and `W4A8_NVFP4_FP8` also use `activation_format="fp8"` and +`As=input_scale`. Both use group-32 FP4 weights; NVFP4's E4M3 block scales and separate +`weight_scale_2` differ from group-16 NVFP4 W4A4. Ordinary FP8 activations keep one calibrated +scale per tensor/expert; they are not re-quantized per block as MXFP8. The MX loops consume +unit activation group scales; affine NVFP4 prefill instead reuses the weight-upcast GEMM, +storing FP8-rounded inputs exactly in BF16. Both restore calibration in the FP32 epilogue. Input +`pre_quant_scale` also supports MX/FP4 activation-quantized paths, including NVFP4 AWQ. + +ModelOpt `NVFP4_SVD` (SVDQuant) is supported by dense `matmul_2d`: pass the BF16 +`svdquant_lora_a` `(rank, K)`, `svdquant_lora_b` `(N, rank)` and input `pre_quant_scale`: + +```python +out = matmul_2d( + x, weight.view(torch.int8), Bs=weight_scale, + a_global_scale=input_scale, b_global_scale=weight_scale_2, + a_pre_quant_scale=pre_quant_scale, + svdquant_lora_a=lora_a, svdquant_lora_b=lora_b, + simulate_unfused=True, +) +``` + +Smoothing, NVFP4 activation quantization and the BF16 low-rank down projection normally share one pass. +Large rank-32 prefill on SM10x quantizes separately from the down projection; both smooth the input +without materializing an extra activation buffer. The up projection is fused into the GEMM epilogue +before GLU/requant. All dot products accumulate in FP32. `simulate_unfused=True` matches the +BF16 materialization points in ModelOpt's two linear branches and their addition; the default +retains FP32 for the final sum. +Factors use the checkpoint's orientation without pre-folding smoothing or global scales into them. +Weight scales may be affine or pre-swizzled at load time. Gate/up weights and their up-factor rows +must use the usual interleaved order; independent gate/up down factors can be concatenated with +matching zero-padded up factors. Fused gate/up also requires shared smoothing and global scales. +This supports dense BF16 inference with ranks divisible by 32 and K divisible by 64; +routed SVDQuant MoE and Nunchaku's INT4/packed checkpoint layouts are not implemented. +`python bench/bench_moe.py SVDQuant` compares the complete linear on Gemma-4-31B up/down shapes +against FlashInfer, including smoothing, activation quantization and the correction projections. + +`repack_ggml_int(payload, "Q4_0" | "Q8_0")` reads canonical little-endian uint8 GGML blocks +`(..., N, K/32, 18 | 34)` and returns the same integer operand tuple with FP32 group scales. +Q8_0 is also a ModelOpt GGML recipe; Q4_0 follows upstream GGML. A GGUF loader supplies the +payload and its logical shape; this helper does not parse containers. IQ codebooks, Q4_1 +affine minima, and Q5 formats are unsupported. +GPTQ act-order repacking also returns `b_perm` when K reordering is required. +Fused integer MoE blocks require matching widths, group recipes, and zero-point presence across +gate/up and down. Both projections may use per-channel scales despite their different K dimensions; +scale dtypes and calibrated values may differ between projections. Routed ops carry the expert dim in front (`(E, ...)` weights and scales). MX scales may instead be passed **pre-swizzled** — the 5D output of `swizzle_mx_scales` — which is the deployment contract: @@ -150,7 +226,9 @@ the rows. `rms_norm_rows` exposes the standalone kernel. `nvfp4_quantize_two_level(weight)` returns `(packed_e2m1, e4m3_block_scales, fp32_global)` — the canonical two-level NVFP4 weight quant (`global = amax / (6·448)`). The `*_act_quant` helpers are the offline activation quants the ops use internally, exposed for pre-quantized (`As`) pipelines; -`fp8_act_quant_block_dynamic` / `fp8_act_quant_tensor_wide` are their block-FP8 siblings. +`fp8_act_quant_block_dynamic` / `fp8_act_quant_tensor_wide` are their block-FP8 siblings; +`fp8_act_quant_tensor_wide(x, K, amax_upper_bound=...)` is the per-token quant of the tensor-wide and +row-wise formats, the optional bound capping each row's amax. ## Autotuning diff --git a/finegrained-kernels/bench/README.md b/finegrained-kernels/bench/README.md index 0418cda20..5c6c9d2ac 100644 --- a/finegrained-kernels/bench/README.md +++ b/finegrained-kernels/bench/README.md @@ -6,14 +6,14 @@ reference implementations, on real model shapes. ## What it compares -Three figure rows, each a **decode | prefill** subplot pair: +One figure contains every model and quantization recipe, including Llama 4 rowwise FP8 and +FP8/integer weight-only cases. Its three rows are **decode | prefill** subplot pairs: - **quantized** — the MoE at every impl's best: `moe_fused_*` against finegrained-fp8, DeepGEMM, vLLM, TRT-LLM, `transformers@main` and `triton_kernels`. Each impl contributes the shape it has — fused where there is one, the two-GEMM shape for `transformers@main`, which has only that. There is no fused/unfused split in the figure: it is an internal distinction of ours that most of these - baselines do not have. The unfused arms (`moe_unfused_*`) stay in `ARMS`, but no problem's - baseline set enlists them, so they do not run by default. + baselines do not have. The unfused arm is included where a problem requests it. - **linear quantized** — one qkv-shaped `matmul_2d` linear per model, in its deployment format - **unquantized (BF16)** — fused vs transformers `grouped_mm`/`batched_mm`, SonicMoE, vLLM, DeepGEMM BF16 and megablocks @@ -27,8 +27,17 @@ torch/cuBLAS path), **SonicMoE**, the OpenAI **triton_kernels** MXFP4 path (GPT- import-guarded — a missing dependency skips that baseline instead of failing the run. Every cell runs in three modes: `eager`, `cudagraph` (decode's deployment mode), and -`compile` (`torch.compile(max-autotune, fullgraph)`). A cell that raises is a red ✕ marker; -the others still run. +`compile` (`torch.compile(max-autotune, fullgraph)`). Failures stay in the CSV; the other modes +still run. Decode shows CUDA graph latency and prefill shows eager latency, on logarithmic axes +so slow baselines do not flatten the faster kernels. Missing deployment-mode timings use red ✕ +markers; decode falls back to dotted eager bars when available. In both +columns, a hatched outline shows compiled latency when it is faster. Every format uses these +same plotting rules; the ordinary Torch weight-only arm also supplies its compiled measurement. + +Each case runs in a fresh process, including single-GPU runs, so an optional baseline's CUDA +fault cannot invalidate later cases. Pinned legacy finegrained-fp8 MX prefill readers use 32-bit +expert offsets: larger-than-2-GiB stacks are unavailable, including through Transformers' wrapper. +The fused reader also faults at the 2-GiB boundary and is excluded there. On CUDA, every mode cycles through identical weight banks whose **active** weights exceed **twice the device's L2 cache**. For MoE, active bytes count only the experts selected by the @@ -43,7 +52,7 @@ CSV rows record `cache_policy`; logs show bank counts, active bytes and checkpoi storage. Insufficient memory fails explicitly; it never falls back to warm weights. Replotting or merging old CSVs without matching cache metadata requires rerunning the benchmark. XPU retains its existing backend timing behavior and is labeled `backend_default`, not CUDA cold. -Previously generated figures used warm CUDA graph timing; regenerate them for cold-weight comparisons. +The committed figure uses this cold-weight policy. ## The pre-swizzled fast path @@ -53,7 +62,7 @@ so its numbers reflect the tcgen05 fast path. Only MX weights on 128-aligned dim to measure the affine path instead. Correctness is cross-checked in-run: each baseline's output is compared to the finegrained-kernels -anchor (`parity-vs-finegrained-kernels`, relative L2 distance). What the number means, by regime: +anchor (`parity-vs-finegrained-kernels`, normalized maximum absolute difference). What the number means, by regime: - **fp8 ~3e-3 / MX ~1e-2 / fp4 ~5e-2** — arms that share our activation quant: only reduction order and epilogue rounding differ. @@ -63,8 +72,7 @@ anchor (`parity-vs-finegrained-kernels`, relative L2 distance). What the number the two quantized intermediates go through a second GEMM. These are measured values, not a bound, and this oracle cannot tell that regime apart from a bug of the same magnitude — the model smokes (generation) and the torch-reference tests carry that. -- **of order 1 (measured 1.3–1.7)** — the outputs are uncorrelated: the relative distance between two - independent same-magnitude vectors is √2. A wrong gate|up row order (pairing the wrong halves in +- **of order 1 (measured 1.3–1.7)** — a likely layout or routing mismatch. A wrong gate|up row order (pairing the wrong halves in the GLU) or a routing error lands here, which is how the stacked-vs-interleaved feed was caught. The bench's gate|up weights are STACKED (`[gate; up]`); the finegrained-kernels arms are fed the @@ -74,7 +82,7 @@ row-interleaved view (`_interleave_gate_up`) because the kernels read gate|up in ```bash python bench/bench_moe.py # full grid, single GPU -> bench_moe.csv + bench_moe.png -python bench/bench_moe.py --gpus 8 # shard problems across 8 GPUs (one process per GPU), then merge + plot +python bench/bench_moe.py --gpus 8 # run cases across 8 GPUs, then merge + plot python bench/bench_moe.py --gpus 5 --devices 3,4,5,6,7 # pin the shards to GPUs 3-7 python bench/bench_moe.py --smoke # fast everything-compiles pass (3-trial tunes, 256-tok prefill) python bench/bench_moe.py --no-preswizzle # affine MX scales instead of the fast path diff --git a/finegrained-kernels/bench/bench_moe.png b/finegrained-kernels/bench/bench_moe.png index 73e612443..a8d0f1b92 100644 Binary files a/finegrained-kernels/bench/bench_moe.png and b/finegrained-kernels/bench/bench_moe.png differ diff --git a/finegrained-kernels/bench/bench_moe.py b/finegrained-kernels/bench/bench_moe.py index 125387501..174dd3bf5 100644 --- a/finegrained-kernels/bench/bench_moe.py +++ b/finegrained-kernels/bench/bench_moe.py @@ -16,7 +16,7 @@ The **finegrained-kernels** arm is the local kernel; **finegrained-fp8** is the upstream hub build (``kernels-community/finegrained-fp8`` @ ``v4``, which has the fused MoE + MX paths). By default the finegrained-kernels arm feeds PRE-SWIZZLED (SWIZZLE_32_4_4) MX weight scales, so -its numbers reflect the tcgen05 fast path (set ``PRESWIZZLE=0`` for the affine path). Writes +its numbers reflect the tcgen05 fast path (pass ``--no-preswizzle`` for the affine path). Writes ``bench/bench_moe.csv`` (all 3 modes) + ``bench/bench_moe.png`` beside this file. Rows (each row = decode | prefill subplot pair in the figure): @@ -25,7 +25,7 @@ transformers@main contributes its two-GEMM experts dispatch unquantized (BF16) finegrained-kernels fused vs transformers grouped_mm/batched_mm vs SonicMoE vs DeepGEMM grouped BF16 - linear quantized matmul_2d, one qkv-proj-shaped linear (N=3H, K=H) per model in + linear quantized matmul_2d, one qkv-proj-shaped linear per model in its deployment format — FP8 128x128 (finegrained-kernels/finegrained-fp8/DeepGEMM), MXFP4 W4A4 (finegrained-kernels W4A4, finegrained-fp8 W4A8, DeepGEMM FP4), NVFP4 (finegrained-kernels only), MXFP8 (finegrained-kernels/finegrained-fp8) @@ -41,12 +41,29 @@ block scales (the B200 deployment format; DeepGEMM SM100 rejects fp32 scales by design) MiniMaxAI/MiniMax-M3 MXFP8 W8A8 finegrained-fp8 (no TRT-LLM: clamped SwiGLU) + meta-llama/Llama-4-Scout rowwise FP8 W8A8 torch, FBGEMM, vLLM (TRT-LLM rowwise unavailable on B200) + unknown model FP8 W8A16 E128 H6144 I6144 top4, tensor/block scales -The TRT-LLM arms are FlashInfer's ``trtllm_*_routed_moe`` kernels — the backends vLLM's +The FP8/FP4 TRT-LLM arms are FlashInfer's ``trtllm_*_routed_moe`` kernels — the backends vLLM's serving path dispatches to on Blackwell for these quant modes (its own triton fused-MoE -covers only bf16 and block-FP8). Weight prep and activation quant come from THEIR stack, +covers bf16 and FP8). Weight prep and activation quant come from THEIR stack, so each bar is that integration end-to-end. +Integer weight-only examples (synthetic tensors, actual checkpoint recipes/shapes): + Qwen3-30B-A3B GPTQ INT4 symmetric and AWQ INT4 asymmetric MoE, group 128, + BF16 runtime; Qwen2.5-7B GPTQ INT4/INT8 and AWQ INT4 QKV, group 128, FP16 runtime. + vLLM Marlin covers all five; ExLlama (vLLM's adapted backend) covers FP16 GPTQ + INT4/INT8 dense rows, including its reconstruction+cuBLAS prefill path. FlashInfer TRT-LLM MXINT4 covers symmetric BF16 MoE + with the SAME quantized values, repeating g128 scales into g32; it cannot represent + AWQ zero points. Packing is offline; routing conversion stays in the timed call. + Normal runs tune FlashInfer tactics before timing; --smoke uses its default tactic. + These are layer microbenchmarks, not downloaded-checkpoint accuracy or serving throughput. + +ModelOpt NVFP4_SVD dense examples use Gemma-4-31B MLP up/down shapes and rank-32 +BF16 corrections, compared against FlashInfer's complete SVDQuant linear. Smoothing, +activation quantization and both correction projections are timed; checkpoint-only factor +transforms are prepared offline. The tensors are synthetic, not loaded checkpoint weights. + Every (row, problem, regime, impl) cell runs in THREE modes: eager do_bench on the plain call cudagraph do_bench_cudagraph (decode's deployment mode) @@ -70,7 +87,9 @@ python bench/bench_moe.py --gpus 5 --devices 3,4,5,6,7 (pin those shards to GPUs 3-7) python bench/bench_moe.py --smoke (fast compile check) python bench/bench_moe.py --no-preswizzle (affine MX scales instead of the fast path) + python bench/bench_moe.py GPTQ AWQ --smoke (integer weight-only examples) python bench/bench_moe.py gpt-oss (substring filter on row/problem names) + python bench/bench_moe.py Llama-4 (Llama4 rowwise FP8 routed experts + QKV linear) python bench/bench_moe.py --replot (rebuild the figure from bench_moe.csv) python bench/bench_moe.py --mock (no GPU: random latencies, validates the figure layout) python bench/bench_moe.py --help @@ -81,7 +100,10 @@ import random import subprocess import sys +from concurrent.futures import ThreadPoolExecutor +from functools import partial from types import SimpleNamespace +from textwrap import fill import argparse @@ -93,8 +115,8 @@ _parser.add_argument("filters", nargs="*", metavar="FILTER", help="substring filter(s) on row/problem names (e.g. gpt-oss, 'linear quantized')") _parser.add_argument("--gpus", type=int, default=1, - help="shard the problems across this many GPUs, one process per GPU; the " - "coordinator merges the shard CSVs and plots (default 1: run inline)") + help="run cases across this many GPUs, with a fresh process per case; " + "the coordinator merges the CSVs and plots (default 1)") _parser.add_argument("--devices", default="", help="comma-separated physical GPUs for the shards (default 0..gpus-1); shard g " "still owns tasks where i %% gpus == g, only the device it runs on changes") @@ -124,6 +146,10 @@ os.environ.setdefault("FINEGRAINED_AUTOTUNE_TRIALS", "3") import torch # noqa: E402 +try: + from torch.nn.functional import ScalingType, scaled_mm # noqa: E402 +except ImportError: + ScalingType = scaled_mm = None _HERE = os.path.dirname(os.path.abspath(__file__)) _ROOT = os.path.dirname(_HERE) @@ -131,6 +157,7 @@ sys.path.insert(0, os.path.join(_ROOT, "tests")) import finegrained_kernels as fgm # noqa: E402 local branch from finegrained_kernels.bench import do_bench, rotation_count # noqa: E402 +from finegrained_kernels.quant import fp8_act_quant_tensor_wide # noqa: E402 from kernels import get_kernel # noqa: E402 # kernels-community IS a trusted publisher, so the check passes whenever it can reach the Hub @@ -158,8 +185,8 @@ def _can_preswizzle(cfg): # the activation format fed to the forwards, in the kernels' public vocabulary: None follows the -# weight family (the op resolves it off the weight dtypes), "bf16" is weight-only (raw bf16 -# activations), an explicit format pins a chain (DeepSeek-V4 W4A8 pins "mxfp8"). +# weight family (the op resolves it off the weight dtypes), "bf16" is the weight-only +# mode (BF16 or FP16 inputs stay in their incoming dtype), an explicit format pins a chain (DeepSeek-V4 W4A8 pins "mxfp8"). def _activation_format(cfg): return cfg["activation_format"] @@ -214,7 +241,8 @@ def _preswizzle_moe_scale(scale): return fgm.swizzle_mx_scales(scale) -from utils import WEIGHTS, dq_nvfp4_two_level, make_weights # noqa: E402 tests/utils.py registry +from utils import INT_BITS, WEIGHTS, dq_grouped, dq_int, dq_nvfp4_two_level, make_weights # noqa: E402 tests/utils.py registry +from finegrained_kernels.quant import nvfp4_act_quant # noqa: E402 from transformers.integrations.deepgemm import ( # noqa: E402 deepgemm_bf16_experts_forward, deepgemm_fp8_fp4_experts_forward, @@ -229,22 +257,71 @@ def _preswizzle_moe_scale(scale): fp8_grouped_mm_experts_forward, ) from transformers.integrations.sonicmoe import sonicmoe_experts_forward # noqa: E402 +import transformers.integrations.mxfp4 as _tfmx # noqa: E402 # vLLM's triton fused-MoE kernel — its bf16 serving path (DeepGEMM/CUTLASS/FlashInfer are its -# quantized backends). Soft import off the local checkout; the arm is enlisted only when it -# resolves. Tuned tile configs for the roster's shapes ship in bench/vllm_configs — vLLM has +# quantized backends). The installed vLLM package is optional; the arm is enlisted only +# when it resolves. Tuned tile configs for the roster's shapes ship in bench/vllm_configs — vLLM has # no B200 bf16 JSONs for them and would silently fall back to heuristic defaults (measured # 12-23% slow at prefill). os.environ.setdefault("VLLM_TUNED_CONFIG_FOLDER", os.path.join(_HERE, "vllm_configs")) -sys.path.insert(0, os.path.expanduser("~/vllm")) try: from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts as vllm_fused_experts # noqa: E402 from vllm.model_executor.layers.fused_moe.activation import MoEActivation as _VllmAct # noqa: E402 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config as _vllm_fp8_qc # noqa: E402 from vllm.model_executor.layers.quantization.utils.fp8_utils import per_token_group_quant_fp8 as _vllm_group_quant # noqa: E402 from vllm._custom_ops import scaled_fp8_quant as _vllm_tensor_quant # noqa: E402 + from vllm._custom_ops import cutlass_scaled_mm as _vllm_scaled_mm # noqa: E402 except ImportError: vllm_fused_experts = None +try: + from vllm.model_executor.layers.fused_moe.experts.marlin_moe import fused_marlin_moe # noqa: E402 + from vllm.scalar_type import scalar_types # noqa: E402 +except ImportError: + fused_marlin_moe = None + +try: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( # noqa: E402 + apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin, prepare_fp8_moe_layer_for_marlin, + ) +except ImportError: + apply_fp8_marlin_linear = None + + +try: + import fbgemm_gpu.experimental.gen_ai # noqa: E402,F401 registers the native operators + _fbgemm_rowwise_grouped = torch.ops.fbgemm.f8f8bf16_rowwise_grouped_stacked + _fbgemm_rowwise = torch.ops.fbgemm.f8f8bf16_rowwise + _fbgemm_rowwise_quant = torch.ops.fbgemm.quantize_fp8_per_row +except (ImportError, AttributeError, OSError): + try: + import mslk.gemm # noqa: E402,F401 registers the grouped GEMM and its compile metadata + from mslk.quantize.triton.fp8_quantize import quantize_fp8_row as _fbgemm_rowwise_quant # noqa: E402 + _fbgemm_rowwise_grouped = torch.ops.mslk.f8f8bf16_rowwise_grouped_stacked + _fbgemm_rowwise = torch.ops.mslk.f8f8bf16_rowwise + except (ImportError, AttributeError, OSError): + _fbgemm_rowwise_grouped = None + _fbgemm_rowwise = None + +# vLLM's production GPTQ/AWQ Marlin path; repacking is load-time work. +try: + from vllm._custom_ops import gptq_marlin_repack + from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + apply_gptq_marlin_linear, + marlin_make_workspace_new, + marlin_permute_scales, + marlin_zero_points, + ) +except ImportError: + apply_gptq_marlin_linear = None + +# vLLM's ExLlama-derived GPTQ backend (its own ExllamaLinearKernel path). +try: + from vllm._custom_ops import gptq_gemm, gptq_shuffle +except ImportError: + gptq_gemm = None + + # kernels-community/nvfp4-gemm — the dense NVFP4 GEMM behind transformers' `NVFP4Linear` # (upstream #47883). Linear-only: it never sees an MoE, because experts are fused 3-D # parameters rather than nn.Linear. Benched here two ways — as a 2D GEMM against our @@ -252,16 +329,13 @@ def _preswizzle_moe_scale(scale): # an MoE would cost if it were served by that integration. try: from transformers.integrations.hub_kernels import lazy_load_kernel as _hub_kernel # noqa: E402 - _nvfp4_gemm = _hub_kernel("nvfp4") + _nvfp4_gemm = None # loaded only by an NVFP4 dense task except Exception: _nvfp4_gemm = None # megablocks dropless-MoE (dMoE) — the hub build, NOT a local pip install: megablocks pins an # older torch, so installing it from source downgrades the environment out from under every other # arm. Loaded like the rest of the roster. -try: - _megablocks = get_kernel("kernels-community/megablocks", version=1, trust_remote_code=TRUSTED) -except Exception: - _megablocks = None +_megablocks = None # loaded only by a BF16 task # FlashInfer TRT-LLM fused MoE (NVIDIA's Blackwell serving kernels) — installed --no-deps # (its PyPI pin would downgrade torch; the JIT has no torch ABI coupling). JIT needs @@ -324,6 +398,22 @@ def _check_cuda_home(): trtllm_fp8_block_scale_routed_moe = None +try: + from flashinfer import autotune as flashinfer_autotune + from flashinfer.fused_moe import trtllm_mxint4_block_scale_routed_moe + from flashinfer.fused_moe.core import convert_to_block_layout + from flashinfer.fused_moe.core import _maybe_get_cached_w3_w1_permute_indices + from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache + from flashinfer.fp4_quantization import block_scale_interleave +except (ImportError, OSError): + trtllm_mxint4_block_scale_routed_moe = None + +try: + from flashinfer import autotune as flashinfer_svdquant_autotune + from flashinfer.gemm import svdquant_linear as flashinfer_svdquant_linear +except (ImportError, OSError): + flashinfer_svdquant_linear = None + UPSTREAM_FP8_LABEL = "v4" # the legend suffix: the hub tag the pinned snapshot was published under # The `v4` tag was rebuilt on 2026-09-13 with a new op-namespace id (same kernel source). That # re-keys the upstream tuner's caches, and its cold search on the MX grouped rows (DeepSeek-V4 @@ -331,15 +421,10 @@ def _check_cuda_home(): # upstream tuner has no trap fences, so the fault poisons the shard's CUDA context and blanks every # later cell. Pin the snapshot the committed figure measured (identical source, warm crowns). UPSTREAM_FP8_REV = "29083040812e244b390757d6198e2889fe551d13" -upstream_fp8 = (None if (MOCK or REPLOT) - else get_kernel("kernels-community/finegrained-fp8", revision=UPSTREAM_FP8_REV, trust_remote_code=TRUSTED)) - -# OpenAI triton_kernels (matmul_ogs) — the MXFP4 experts path transformers uses for -# GPT-OSS. Loaded like finegrained-fp8; its module-level handle drives the mxfp4 swizzle helpers. -if not (MOCK or REPLOT): - import transformers.integrations.mxfp4 as _tfmx - triton_kernels_hub = get_kernel("kernels-community/gpt-oss-triton-kernels", version=1, trust_remote_code=TRUSTED) - _tfmx.triton_kernels_hub = triton_kernels_hub +# Load Hub dependencies only when a selected task needs them. Integer-only runs do not +# download unrelated FP8/GPT-OSS kernels or depend on their availability. +upstream_fp8 = None +triton_kernels_hub = None DEVICE = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda" if DEVICE not in ("cuda", "xpu"): @@ -354,7 +439,7 @@ def _check_cuda_home(): # fixed left-to-right model order for every figure row (matched by base-model prefix, # so GLM-5.2-NVFP4 and GLM-5.2 both land in the GLM-5.2 slot). Roughly most-baseline- # support first, finegrained-kernels-only (GPT-OSS, GLM-NVFP4) last. -CANONICAL_MODEL_ORDER = ["DeepSeek-V4", "DeepSeek-V3", "MiniMax-M3", "GPT-OSS-120B", "GLM-5.2", "Mistral-4"] +CANONICAL_MODEL_ORDER = ["DeepSeek-V4", "DeepSeek-V3", "MiniMax-M3", "GPT-OSS-120B", "GLM-5.2", "Mistral-4", "Llama-4-Scout", "Qwen3-30B-A3B", "Qwen2.5-7B", "Gemma-4-31B", "nemotron-diff", "unknown model"] # Low-VRAM parts run the same recipes and baselines at /8 experts, /2 dims, INSTEAD OF the # full roster rather than alongside it. Selected off the device, not a knob, so a small part @@ -433,8 +518,46 @@ def _check_cuda_home(): baselines=("finegrained-fp8",), fp8_block=None, block_size=None, act="silu", swiglu_alpha=1.702, swiglu_limit=7.0, ), + "meta-llama/Llama-4-Scout-17B-16E-Instruct FP8 rowwise (E16 H5120 I8192 top1)": dict( + # hugging-quants/Llama-4-Scout-17B-16E-Instruct-fbgemm: one scale per weight row, + # dynamic per-token activations. Compare routed experts only: sigmoid input routing + # is prepared once for all arms; the router projection and shared expert are excluded. + E=16, H=5120, I=8192, top_k=1, weights="fp8_rowwise", activation_format="fp8", + baselines=("torch_rowwise", "fbgemm", "vllm"), fp8_block=None, block_size=(1, 5120), + act="silu", swiglu_alpha=None, swiglu_limit=None, router_weight_on_input=True, + weight_std=0.02, activation_scale_ub=1200.0, + ), + "unknown model FP8 W8A16 tensor (E128 H6144 I6144 top4)": dict( + E=128, H=6144, I=6144, top_k=4, weights="fp8_tensor", activation_format="bf16", + baselines=("finegrained-kernels-unfused", "torch_weight_only", + "marlin"), fp8_block=None, block_size=None, + act="silu", swiglu_alpha=None, swiglu_limit=None, weight_std=0.02, + ), + "unknown model FP8 W8A16 block128 (E128 H6144 I6144 top4)": dict( + E=128, H=6144, I=6144, top_k=4, weights="fp8_128x128", activation_format="bf16", + baselines=("finegrained-kernels-unfused", "torch_weight_only", + "marlin"), fp8_block=[128, 128], block_size=(128, 128), + act="silu", swiglu_alpha=None, swiglu_limit=None, weight_std=0.02, + ), + # Synthetic weights in published checkpoint recipes, not downloaded model weights. Both + # Qwen3 rows run BF16 activations/scales; GPTQ's config defaults FP16, so this is the + # equivalent --dtype bfloat16 serving setting. Router renormalization matches Qwen3. + # https://huggingface.co/Qwen/Qwen3-30B-A3B-GPTQ-Int4/blob/main/config.json + # https://huggingface.co/QuantTrio/Qwen3-30B-A3B-Thinking-2507-AWQ/blob/main/config.json + "Qwen/Qwen3-30B-A3B-GPTQ-Int4 W4A16 g128 sym BF16 (E128 H2048 I768 top8)": dict( + E=128, H=2048, I=768, top_k=8, weights="int4", activation_format="bf16", + group=128, asymmetric=False, dtype=torch.bfloat16, renormalize=True, + baselines=("vllm_marlin", "flashinfer_mxint4"), fp8_block=None, block_size=None, + act="silu", swiglu_alpha=None, swiglu_limit=None, + ), + "QuantTrio/Qwen3-30B-A3B-Thinking-2507-AWQ W4A16 g128 asym BF16 (E128 H2048 I768 top8)": dict( + E=128, H=2048, I=768, top_k=8, weights="int4", activation_format="bf16", + group=128, asymmetric=True, dtype=torch.bfloat16, renormalize=True, + baselines=("vllm_marlin",), fp8_block=None, block_size=None, + act="silu", swiglu_alpha=None, swiglu_limit=None, + ), } -# the same base-model roster, run as if dequantized to BF16 (one shape per model) +# BF16 reference roster: one shape per base model. FULL_BF16_PROBLEMS = { "deepseek-ai/DeepSeek-V4 BF16 (E256 H4096 I2048 top6)": dict( E=256, H=4096, I=2048, top_k=6, weights="bf16", activation_format=None, @@ -535,6 +658,9 @@ def _check_cuda_home(): ), } +# Keep these modest integer expert banks at their actual geometry on smaller GPUs too. +SMALL_MOE_PROBLEMS.update({p: c for p, c in FULL_MOE_PROBLEMS.items() if c["weights"] in INT_BITS}) + MOE_PROBLEMS = SMALL_MOE_PROBLEMS if SMALL_GEOMETRY else FULL_MOE_PROBLEMS BF16_PROBLEMS = SMALL_BF16_PROBLEMS if SMALL_GEOMETRY else FULL_BF16_PROBLEMS # the roster is picked off the device, so say which one ran: the two are not comparable, and a @@ -551,9 +677,7 @@ def _check_cuda_home(): N=12288, K=4096, weights="fp8_128x128_ue8m0", activation_format=None, block=[128, 128], baselines=("finegrained-fp8", "deepgemm"), ), - # (No standalone dense 2D W4A16 row: no prominent model ships dense mxfp4 — gpt-oss's only mxfp4 - # is its experts (benchmarked as the MoE W4A16 row), and dense 4-bit-weight LLMs deploy GPTQ/AWQ - # int4, not mxfp4. The 2D weight-only kernel stays covered by test_ops + the unfused MoE path.) + # GPT-OSS deploys MXFP4 in experts; dense GPTQ/AWQ examples follow below. "nvidia/GLM-5.2-NVFP4 attn W4A4 qkv-shaped (N=18432 K=6144)": dict( N=18432, K=6144, weights="nvfp4", activation_format=None, block=None, baselines=(), # no baseline supports NVFP4 @@ -583,6 +707,53 @@ def _check_cuda_home(): N=18432, K=6144, weights="mxfp8", activation_format=None, block=None, baselines=("finegrained-fp8",), # DeepGEMM FP8 is 128-block, not group-32 MX ), + "meta-llama/Llama-4-Scout-17B-16E-Instruct attn FP8 rowwise qkv-shaped (N=7168 K=5120)": dict( + # Scout uses GQA: (40 query heads + 2 * 8 KV heads) * head_dim 128. + N=7168, K=5120, weights="fp8_rowwise", activation_format="fp8", + block=None, baselines=("fbgemm", "vllm_mm"), weight_std=0.02, activation_scale_ub=1200.0, + ), + "unknown model FP8 W8A16 tensor linear (N=8192 K=6144)": dict( + N=8192, K=6144, weights="fp8_tensor", activation_format="bf16", + block=None, baselines=("torch_weight_only", "marlin"), weight_std=0.02, + ), + "unknown model FP8 W8A16 block128 linear (N=8192 K=6144)": dict( + N=8192, K=6144, weights="fp8_128x128", activation_format="bf16", + block=[128, 128], baselines=("torch_weight_only", "marlin"), weight_std=0.02, + ), + # Actual fused QKV width, including GQA: (28 Q + 2 * 4 KV) * 128 = 4608, + # rather than the 3H proxy used by the older rows. Each checkpoint quantizes Q/K/V. + # https://huggingface.co/Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4/blob/main/config.json + # https://huggingface.co/Qwen/Qwen2.5-7B-Instruct-GPTQ-Int8/blob/main/config.json + # https://huggingface.co/Qwen/Qwen2.5-7B-Instruct-AWQ/blob/main/config.json + "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 qkv W4A16 g128 sym FP16 (N=4608 K=3584)": dict( + N=4608, K=3584, weights="int4", activation_format="bf16", group=128, + asymmetric=False, dtype=torch.float16, block=None, baselines=("vllm_marlin", "vllm_exllama"), + ), + "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int8 qkv W8A16 g128 sym FP16 (N=4608 K=3584)": dict( + N=4608, K=3584, weights="int8", activation_format="bf16", group=128, + asymmetric=False, dtype=torch.float16, block=None, baselines=("vllm_marlin", "vllm_exllama"), + ), + "Qwen/Qwen2.5-7B-Instruct-AWQ qkv W4A16 g128 asym FP16 (N=4608 K=3584)": dict( + N=4608, K=3584, weights="int4", activation_format="bf16", group=128, + asymmetric=True, dtype=torch.float16, block=None, baselines=("vllm_marlin",), + ), + # Native ModelOpt payload verified from this checkpoint's header: Q=4096, K=V=1024, + # signed INT4 across output-row pairs, FP32 group128 scales and scalar input_scale/weight_scale_2. + # https://huggingface.co/yury-kapoyko/nemotron-diff-vlm-8b-w4a8_awq/blob/main/hf_quant_config.json + "yury-kapoyko/nemotron-diff-vlm-8b-w4a8_awq ModelOpt QKV W4A8 g128 sym BF16 (N=6144 K=4096)": dict( + N=6144, K=4096, weights="int4", activation_format="fp8", group=128, + asymmetric=False, dtype=torch.bfloat16, block=None, baselines=(), static=True, + ), + # ModelOpt NVFP4_SVD checkpoint: actual MLP dimensions and BF16 rank-32 correction. + # https://huggingface.co/AEON-7/Gemma-4-31B-it-DECKARD-HERETIC-Uncensored-NVFP4-SVDQuant + "Gemma-4-31B ModelOpt SVDQuant up r32 (N=21504 K=5376)": dict( + N=21504, K=5376, weights="nvfp4", activation_format="nvfp4", svdquant_rank=32, + block=None, baselines=("flashinfer_svdquant",), + ), + "Gemma-4-31B ModelOpt SVDQuant down r32 (N=5376 K=21504)": dict( + N=5376, K=21504, weights="nvfp4", activation_format="nvfp4", svdquant_rank=32, + block=None, baselines=("flashinfer_svdquant",), + ), } MODES = ["eager", "cudagraph", "compile"] CACHE_POLICY = "rotating_gt_2l2" if DEVICE == "cuda" else "backend_default" @@ -596,36 +767,60 @@ def _check_cuda_home(): "transformers@main": "#7f7f7f", # transformers main on a quantized checkpoint (status quo) "sonicmoe": "#ff7f0e", "vllm": "#17becf", # vLLM triton fused-MoE (its bf16 serving kernel) + "vllm_marlin": "#bcbd22", + "vllm_exllama": "#8c564b", "trtllm": "#76b900", # FlashInfer TRT-LLM fused MoE (NVIDIA Blackwell serving kernels) + "flashinfer_mxint4": "#76b900", + "flashinfer_svdquant": "#76b900", "triton_kernels": "#8c564b", # OpenAI mxfp4 (GPT-OSS) reference "nvfp4_gemm": "#d62728", # transformers NVFP4Linear kernel (dense 2D row only) "megablocks": "#7f7f7f", # databricks dropless-MoE (bf16 reference) "torch": "#e377c2", # torch/cuBLAS F.scaled_grouped_mm (quantized prefill reference) + "finegrained-kernels-unfused": "#6baed6", + "torch_weight_only": "#e377c2", + "marlin": "#17becf", + "torch_rowwise": "#e377c2", + "fbgemm": "#bcbd22", "torch_mm": "#e377c2", # torch/cuBLAS F.scaled_mm (the attn-row sibling) + "vllm_mm": "#17becf", } -def _impl_label(impl, regime): +def _impl_label(impl): """Legend name. The reference forwards ARE plain torch ops — name them by the op: transformers batched_mm -> torch.bmm (decode) / grouped_mm -> torch._grouped_mm (bf16); ``torch`` -> torch.scaled_grouped_mm (the quantized cuBLAS path). Every other impl (finegrained-kernels, finegrained-fp8, deepgemm, ...) is its own label.""" + if impl == "finegrained-kernels-unfused": + return "finegrained-kernels unfused" + if impl == "torch_weight_only": + return "torch runtime dequant + BF16 GEMM" + if impl == "marlin": + return "vLLM Marlin FP8 W8A16" if impl == "transformers": - return "torch.bmm" if regime == "decode" else "torch._grouped_mm" + return "torch.bmm / torch._grouped_mm" if impl == "torch": return "torch.scaled_grouped_mm" + if impl == "torch_rowwise": + return "torch.scaled_mm (per expert)" + if impl == "fbgemm": + return "FBGEMM / MSLK" if impl == "torch_mm": return "torch.scaled_mm" - # transformers' NVFP4Linear kernel (kernels-community/nvfp4-gemm). One name in both rows: - # this label function only sees the regime, not the row, and "(per-expert)" — true of the MoE - # rows, where it is looped because that integration has no MoE kernel — is wrong on the dense - # attn row. The per-expert framing lives in the row docs instead of a misapplied suffix. if impl == "nvfp4_gemm": return "nvfp4-gemm" + if impl == "vllm_exllama": + return "ExLlama (vLLM)" + if impl == "vllm_marlin": + return "vLLM Marlin" + if impl == "flashinfer_mxint4": + return "FlashInfer TRT-LLM MXINT4" if impl == "vllm": return "vLLM fused_moe" + if impl == "vllm_mm": + return "vLLM CUTLASS" if impl == "trtllm": - return "TRT-LLM (FlashInfer)" + return "FlashInfer" if impl == "finegrained-fp8": return f"finegrained-fp8@{UPSTREAM_FP8_LABEL}" return impl @@ -641,11 +836,77 @@ def _mark_static(*tensors): torch._dynamo.mark_static_address(t) +def _integer_weight(cfg, n, k, experts): + weight, scale, _ = WEIGHTS[cfg["weights"]]["make"](n, k, experts, group=cfg["group"]) + scale = scale.to(cfg["dtype"]) + bits = INT_BITS[cfg["weights"]] + zero = (torch.randint(1, (1 << bits) - 1, scale.shape, device=DEVICE, dtype=torch.uint8) + if cfg["asymmetric"] else None) + return weight, scale, zero + + +def _integer_values(weight, bits): + """Internal interleaved words -> unsigned logical values, for offline backend packing.""" + vpw = 32 // bits + shifts = bits * torch.arange(vpw, device=weight.device, dtype=torch.int32) + q = (weight.view(torch.int32)[..., None] >> shifts) & ((1 << bits) - 1) + return torch.stack((q[..., :vpw // 2], q[..., vpw // 2:]), -1).flatten(-3) + + +def _gptq_words(weight, bits): + q = _integer_values(weight, bits).t().contiguous() # K, N + k, n = q.shape + vpw = 32 // bits + shifts = bits * torch.arange(vpw, device=q.device, dtype=torch.int32) + return (q.reshape(k // vpw, vpw, n) << shifts[None, :, None]).sum(1, dtype=torch.int32) + + +def _marlin_weight(cfg, weight, scale, zero): + """Same values/scales/zero points in Marlin's layout; no CPU reference packer.""" + bits = INT_BITS[cfg["weights"]] + gptq = _gptq_words(weight, bits) + k, n = gptq.shape[0] * (32 // bits), gptq.shape[1] + empty = torch.empty(0, device=weight.device, dtype=torch.int32) + packed = gptq_marlin_repack(gptq, empty, k, n, bits) + scales = marlin_permute_scales(scale.t().contiguous(), k, n, cfg["group"]) + zeros = (empty if zero is None else + marlin_zero_points(zero.t().int().contiguous(), k // cfg["group"], n, bits)) + return packed, scales, zeros + + +def _exllama_weight(cfg, weight, scale): + """The serving backend's offline shuffle and GPTQ-v1 (zero minus one) metadata.""" + bits = INT_BITS[cfg["weights"]] + assert not cfg["asymmetric"] and scale.dtype == torch.float16 + packed = _gptq_words(weight, bits) + groups, n = scale.shape[-1], scale.shape[-2] + vpw = 32 // bits + zero = (1 << (bits - 1)) - 1 + word = sum(zero << (bits * j) for j in range(vpw)) + zeros = torch.full((groups, n // vpw), word, dtype=torch.int32, device=weight.device) + empty = torch.empty(0, device=weight.device, dtype=torch.int32) + gptq_shuffle(packed, empty, bits) + return packed, zeros, scale.t().contiguous(), empty + + +def _marlin_type(cfg): + assert cfg["weights"] in ("int4", "int8") and (not cfg["asymmetric"] or cfg["weights"] == "int4") + return (scalar_types.uint4 if cfg["asymmetric"] else + scalar_types.uint4b8 if cfg["weights"] == "int4" else scalar_types.uint8b128) + + def build(cfg): """Quantized weights for one MoE problem, shared by every impl and row. The GLM-FP8 problem uses UE8M0 block scales (finegrained-fp8/finegrained-kernels decode them natively; DeepGEMM SM100 requires them).""" + torch.manual_seed(0) E, H, inter = cfg["E"], cfg["H"], cfg["I"] + if cfg["weights"] in INT_BITS: + torch.manual_seed(0) + gu, gus, guz = _integer_weight(cfg, 2 * inter, H, E) + dn, dns, dnz = _integer_weight(cfg, H, inter, E) + _mark_static(gu, gus, guz, dn, dns, dnz) + return gu, gus, dn, dns, None, None, guz, dnz if cfg["weights"] == "fp8_128x128_ue8m0": def make(n, k, e): return (*make_weights(n, k, DEVICE, [128, 128], @@ -656,15 +917,26 @@ def make(n, k, e): # single-level FP8/MX formats). Thread it through — dropping it silently ran nvfp4 at global=1. gu, gus, gu_g = make(2 * inter, H, E) dn, dns, dn_g = make(H, inter, E) + if "weight_std" in cfg: + gus, dns = gus * cfg["weight_std"], dns * cfg["weight_std"] _mark_static(gu, gus, dn, dns) - return gu, gus, dn, dns, gu_g, dn_g + return gu, gus, dn, dns, gu_g, dn_g, None, None def routing(cfg, tokens): torch.manual_seed(0) - hidden = torch.randn(tokens, cfg["H"], device=DEVICE, dtype=torch.bfloat16) + hidden = torch.randn(tokens, cfg["H"], device=DEVICE, dtype=cfg.get("dtype", torch.bfloat16)) logits = torch.randn(tokens, cfg["E"], device=DEVICE) - w, idx = torch.topk(torch.softmax(logits, -1), cfg["top_k"], dim=-1) + if cfg.get("router_weight_on_input"): + # Llama4's top-1 sigmoid weight multiplies the expert input, before SwiGLU. + assert cfg["top_k"] == 1 + scores, idx = torch.topk(logits, cfg["top_k"], dim=-1) + hidden = (hidden * scores.sigmoid()).to(hidden.dtype) + w = torch.ones_like(scores) + else: + w, idx = torch.topk(torch.softmax(logits, -1), cfg["top_k"], dim=-1) + if cfg.get("renormalize"): + w = w / w.sum(-1, keepdim=True) return hidden, idx.to(torch.int32), w, logits @@ -682,6 +954,9 @@ class _Experts: (grouped_mm/batched_mm, SonicMoE, DeepGEMM): our (E, out, in) layout is their ``is_transposed=False``; gate|up rows are stacked (concatenated), not interleaved.""" + config = SimpleNamespace(hidden_act="silu") + has_post_expert_norm = False + def __init__(self, cfg, gu, dn, gus=None, dns=None): self.num_experts = cfg["E"] self.has_gate, self.has_bias = True, False @@ -697,7 +972,6 @@ def __init__(self, cfg, gu, dn, gus=None, dns=None): self.activation_scheme = "dynamic" self._deepgemm_disabled = False self.act_fn = torch.nn.functional.silu - self.config = SimpleNamespace(hidden_act="silu") self.swiglu_alpha = cfg["swiglu_alpha"] self.swiglu_limit = cfg["swiglu_limit"] self._cfg = cfg @@ -747,7 +1021,7 @@ def _static_act_kwargs(cfg, hidden): return {"gate_up_proj_activation_scale": per_expert, "down_proj_activation_scale": per_expert} -def moe_fused_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, *_): +def moe_fused_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits, guz=None, dnz=None): """``activation_format`` sets the activation precision; None follows the weight format (mxfp4/nvfp4 -> the all-fp4 W4A4 chain, bf16 weights -> unquantized). dsv4 deploys W4A8, so it pins activation_format="mxfp8". Under ``PRESWIZZLE`` the MX weight scales are @@ -756,31 +1030,38 @@ def moe_fused_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, *_ nvfp4_kw = _nvfp4_kwargs(cfg, hidden, gu, gus, gu_g) # raw scales: before any preswizzle static_kw = _static_act_kwargs(cfg, hidden) gu, gus = _interleave_gate_up(gu, gus) # our kernels read gate|up interleaved + guz = _interleave_rows(guz) if _can_preswizzle(cfg): gus = _preswizzle_moe_scale(gus) # fused gate GEMM reads the interleaved layout dns = _preswizzle_moe_scale(dns) - _mark_static(gu, gus, dns) # derived closure tensors: unmarked, cudagraph trees re-copy them every compiled call + _mark_static(gu, gus, guz, dns) # derived closure tensors: unmarked, cudagraph trees re-copy them every compiled call kw = dict(act_fn=cfg["act"], swiglu_alpha=cfg["swiglu_alpha"], swiglu_limit=cfg["swiglu_limit"], activation_format=_activation_format(cfg), + activation_scale_ub=cfg.get("activation_scale_ub"), + simulate_unfused=True, gate_up_proj_weight_global_scale=gu_g, down_proj_weight_global_scale=dn_g, + gate_up_proj_weight_zero_point=guz, down_proj_weight_zero_point=dnz, **nvfp4_kw, **static_kw) return lambda: fn(hidden, idx, w, gu, dn, gus, dns, **kw) -def moe_unfused_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, *_): +def moe_unfused_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits, guz=None, dnz=None): fn = fgm.moe.moe_unfused_grouped if grouped else fgm.moe.moe_unfused_batched nvfp4_kw = _nvfp4_kwargs(cfg, hidden, gu, gus, gu_g) # raw scales: before any preswizzle static_kw = _static_act_kwargs(cfg, hidden) gu, gus = _interleave_gate_up(gu, gus) # our kernels read gate|up interleaved + guz = _interleave_rows(guz) if _can_preswizzle(cfg): # ONE checkpoint layout: gate_up scales are always the gate-interleaved artifact; the # unfused plain 2N GEMM reads it via the in-kernel INTERLEAVED_SCALES block remap. gus = _preswizzle_moe_scale(gus) dns = _preswizzle_moe_scale(dns) - _mark_static(gu, gus, dns) # derived closure tensors: unmarked, cudagraph trees re-copy them every compiled call + _mark_static(gu, gus, guz, dns) # derived closure tensors: unmarked, cudagraph trees re-copy them every compiled call kw = dict(act_fn=cfg["act"], swiglu_alpha=cfg["swiglu_alpha"], swiglu_limit=cfg["swiglu_limit"], activation_format=_activation_format(cfg), + activation_scale_ub=cfg.get("activation_scale_ub"), gate_up_proj_weight_global_scale=gu_g, down_proj_weight_global_scale=dn_g, + gate_up_proj_weight_zero_point=guz, down_proj_weight_zero_point=dnz, **nvfp4_kw, **static_kw) return lambda: fn(hidden, idx, w, gu, dn, gus, dns, **kw) @@ -826,8 +1107,17 @@ def _fp8_scales(t, block): return t +def _check_upstream_fp8_offsets(cfg, grouped, gu, dn, *, fused=False): + # Pinned MX grouped readers use int32 expert offsets; the fused path also faults at the boundary. + if grouped and cfg["weights"] in _MX_WEIGHTS: + size = max(t.numel() * t.element_size() for t in (gu, dn)) + if size > 2**31 or (fused and size == 2**31): + raise RuntimeError("pinned finegrained-fp8 MX prefill has unsafe 32-bit expert offsets for this weight stack") + + def fp8_fused_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): """Hub finegrained-fp8 fused forward (block_size positional; None = MX, [128,128] = block FP8).""" + _check_upstream_fp8_offsets(cfg, grouped, gu, dn, fused=True) fn = upstream_fp8.moe_fused_grouped if grouped else upstream_fp8.moe_fused_batched gus = _fp8_scales(gus, cfg["fp8_block"]) dns = _fp8_scales(dns, cfg["fp8_block"]) @@ -840,6 +1130,7 @@ def fp8_unfused_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): """finegrained-fp8 unfused: two plain finegrained-fp8 GEMMs + host GLU + weighted reduce. finegrained-fp8's grouped op wants A expert-sorted with per-expert cumulative row-ends (``offsets``) and counts; its batched op wants pre-expanded A + flat expert ids.""" + _check_upstream_fp8_offsets(cfg, grouped, gu, dn) E, top_k, T = cfg["E"], cfg["top_k"], hidden.shape[0] gus = _fp8_scales(gus, cfg["fp8_block"]) dns = _fp8_scales(dns, cfg["fp8_block"]) @@ -885,12 +1176,57 @@ def transformers_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): return lambda: fwd(mod, hidden, idx, w) +def dequantize_fp8_weight(weight, scale, block): + """Native dequantization of aligned blocks without materializing repeated scales.""" + if block is None: + return (weight.float() * scale.float()).to(torch.bfloat16) + n, k = weight.shape[-2:] + shape = (*weight.shape[:-2], n // block[0], block[0], k // block[1], block[1]) + values = weight.float().reshape(shape) * scale.float()[..., :, None, :, None] + return values.reshape(weight.shape).to(torch.bfloat16) + + +def torch_weight_only_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): + """Runtime dequantization: selected experts for decode, grouped BF16 GEMMs for prefill.""" + block = cfg["fp8_block"] + if grouped: + def run(): + experts = _Experts(cfg, dequantize_fp8_weight(gu, gus, block), dequantize_fp8_weight(dn, dns, block)) + return grouped_mm_experts_forward(experts, hidden, idx, w) + else: + def run(): + expert_ids = idx.flatten().long() + gu_bf16 = dequantize_fp8_weight(gu[expert_ids], gus[expert_ids], block) + dn_bf16 = dequantize_fp8_weight(dn[expert_ids], dns[expert_ids], block) + selected = hidden.repeat_interleave(cfg["top_k"], dim=0) + gate, up = torch.bmm(selected[:, None, :], gu_bf16.transpose(1, 2)).squeeze(1).chunk(2, dim=-1) + inter = _glu(gate, up, cfg).to(hidden.dtype) + down = torch.bmm(inter[:, None, :], dn_bf16.transpose(1, 2)).squeeze(1) + return (down.view(hidden.shape[0], cfg["top_k"], -1) * w[..., None]).sum(1).to(hidden.dtype) + return run + + +def marlin_weight_only_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): + layer = torch.nn.Module() + layer.num_experts, layer.hidden_size = cfg["E"], cfg["H"] + layer.intermediate_size_per_partition = cfg["I"] + layer.weight_block_size, layer.orig_dtype = cfg["fp8_block"], hidden.dtype + layer.w13_weight = torch.nn.Parameter(gu, requires_grad=False) + gu, dn, gus, dns = prepare_fp8_moe_layer_for_marlin(layer, gu, dn, gus, dns) + _mark_static(gu, dn, gus, dns, layer.workspace) + return lambda: fused_marlin_moe( + hidden, gu, dn, None, None, gus, dns, w, idx, + scalar_types.float8_e4m3fn.id, workspace=layer.workspace, + ) + + def transformers_main_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): """``transformers`` main on a quantized checkpoint — its ``finegrained_fp8`` integration (fp8_grouped_mm / fp8_batched_mm experts dispatch), which loads the finegrained-fp8 hub kernel behind a Python wrapper. This is the status quo an fgm integration would replace, so it is the reference for what integrating fgm buys end-to-end; the bare ``finegrained-fp8`` arm is the same kernel WITHOUT the wrapper, which separates integration overhead from kernel speed.""" + _check_upstream_fp8_offsets(cfg, grouped, gu, dn) mod = _Experts(cfg, gu, dn, gus, dns) fwd = fp8_grouped_mm_experts_forward if grouped else fp8_batched_mm_experts_forward return lambda: fwd(mod, hidden, idx, w) @@ -1034,7 +1370,7 @@ def run(): return run -def trtllm_per_tensor_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits): +def trtllm_per_tensor_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits, *_): """TRT-LLM (FlashInfer) ``trtllm_fp8_per_tensor_scale_moe`` — the entry point their serving path uses for the W8A8 scheme Mistral-4 ships: one FP8 scale per expert on each weight, one calibrated activation scale. Their gate|up order is [up; gate], so the halves swap offline. @@ -1139,6 +1475,57 @@ def run(): return run +def rowwise_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_, fbgemm=False): + """Native rowwise GEMMs, with sort/GLU/scatter included in the forward. + + Torch uses the shared Triton activation quantizer; FBGEMM uses its own per-row quantizer. + Both consume the same FP8 weights and per-output-channel scales without dequantizing them. + Torch 2.13's grouped rowwise kernel targets Hopper, so use its dense op per active expert. + Slice bounds are prepared once for this benchmark's fixed routing; no host sync is timed. + """ + gus, dns = gus.squeeze(-1).contiguous(), dns.squeeze(-1).contiguous() + upper_bound = cfg["activation_scale_ub"] + scale_ub = torch.tensor([upper_bound], device=hidden.device) + _mark_static(gus, dns, scale_ub) + slices = [] + if not fbgemm: + start = 0 + for expert, count in enumerate(torch.bincount(idx.flatten().long(), minlength=cfg["E"]).tolist()): + if count: + slices.append((expert, start, start + count)) + start += count + + def mm(a, weight, scale, counts): + if fbgemm: + aq, a_scale = _fbgemm_rowwise_quant(a, scale_ub=scale_ub) + return _fbgemm_rowwise_grouped(aq, weight, a_scale, scale, counts) + aq, a_scale = fp8_act_quant_tensor_wide(a, a.shape[-1], amax_upper_bound=upper_bound) + outputs = [ + scaled_mm( + aq[start:end], weight[expert].t(), + a_scale[start:end].reshape(-1, 1), ScalingType.RowWise, + scale[expert].unsqueeze(0), ScalingType.RowWise, + output_dtype=hidden.dtype, use_fast_accum=True, + ) + for expert, start, end in slices + ] + return outputs[0] if len(outputs) == 1 else torch.cat(outputs) + + def run(): + experts, order = torch.sort(idx.reshape(-1)) + counts = (torch.histc(experts.float(), bins=cfg["E"], min=0, max=cfg["E"] - 1).to(torch.int64) + if fbgemm else None) + tokens = order // cfg["top_k"] + gate_up = mm(hidden[tokens], gu, gus, counts) + gate, up = gate_up.chunk(2, dim=-1) + inter = _glu(gate, up, cfg).to(hidden.dtype) + down = mm(inter, dn, dns, counts) + weighted = down * w.reshape(-1)[order, None].to(down.dtype) + return torch.zeros_like(hidden).index_add_(0, tokens, weighted) + + return run + + def vllm_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): """vLLM triton fused-MoE — the kernel its serving path runs for bf16, for block-fp8 W8A8 (the DeepSeek scheme; dynamic per-token-group act quant inside) and for @@ -1152,7 +1539,10 @@ def vllm_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): tki = idx.to(torch.int32) tkw = w.to(torch.float32) qc = None - if cfg["fp8_block"] is not None: + if cfg["weights"] == "fp8_rowwise": + qc = _vllm_fp8_qc(w1_scale=gus, w2_scale=dns, + per_act_token_quant=True, per_out_ch_quant=True) + elif cfg["fp8_block"] is not None: qc = _vllm_fp8_qc(w1_scale=gus.float(), w2_scale=dns.float(), block_shape=list(cfg["fp8_block"])) elif cfg.get("static"): @@ -1173,6 +1563,67 @@ def vllm_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): activation=_VllmAct.SILU, quant_config=qc) +def marlin_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits, guz, dnz): + def pack(weight, scales, zeros): + triples = [_marlin_weight(cfg, weight[e], scales[e], None if zeros is None else zeros[e]) + for e in range(cfg["E"])] + return tuple(torch.stack(t) for t in zip(*triples)) + g, gs, gz = pack(gu, gus, guz) + d, ds, dz = pack(dn, dns, dnz) + workspace = marlin_make_workspace_new(hidden.device, 4) + kind = _marlin_type(cfg).id + _mark_static(g, gs, gz, d, ds, dz, workspace) + return lambda: fused_marlin_moe( + hidden, g, d, None, None, gs, ds, w, idx, kind, + w1_zeros=gz, w2_zeros=dz, workspace=workspace) + + +def flashinfer_int4_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits, guz, dnz): + """TRT-LLM signed INT4: g128 is exactly four repeated g32 scales. No requantization. + + Only the symmetric BF16 row qualifies; AWQ's per-group zero points cannot be + represented. Layout expansion/reordering happens once, outside the forward. + """ + assert cfg["weights"] == "int4" and guz is None and dnz is None + assert hidden.dtype == gus.dtype == torch.bfloat16 and cfg["group"] % 32 == 0 + cache = {} + def pack(weight, scales, gated): + weights_out, scales_out = [], [] + for e in range(cfg["E"]): + q = (_integer_values(weight[e], 4) - 8) & 15 + q = (q[:, ::2] | (q[:, 1::2] << 4)).to(torch.uint8) + scale = scales[e].repeat_interleave(cfg["group"] // 32, -1) + if gated: + # TRT-LLM takes [up; gate], while the shared checkpoint is [gate; up]. + half = q.shape[0] // 2 + q = torch.cat((q[half:], q[:half])) + scale = torch.cat((scale[half:], scale[:half])) + p = _maybe_get_cached_w3_w1_permute_indices(cache, q, 128) + sp = _maybe_get_cached_w3_w1_permute_indices(cache, scale, 128, num_elts_per_sf=32) + else: + p = get_w2_permute_indices_with_cache(cache, q, 128) + sp = get_w2_permute_indices_with_cache(cache, scale, 128, num_elts_per_sf=16) + weights_out.append(convert_to_block_layout(q[p.to(q.device)].contiguous(), 128)) + scales_out.append(block_scale_interleave(scale[sp.to(scale.device)].contiguous())) + return torch.stack(weights_out), torch.stack(scales_out).view(torch.bfloat16) + g, gs = pack(gu, gus, True) + d, ds = pack(dn, dns, False) + _mark_static(g, gs, d, ds) + def run(): + # Keep routing-weight conversion in the timed call: this backend consumes BF16 + # weights in its packed routing ids, while the shared router produces FP32 weights. + packed = (idx << 16) | (w.to(torch.bfloat16).view(torch.int16).int() & 0xffff) + return trtllm_mxint4_block_scale_routed_moe( + packed, hidden, g, gs, None, None, None, d, ds, + cfg["E"], cfg["top_k"], None, None, cfg["I"], 0, cfg["E"], None, + tune_max_num_tokens=hidden.shape[0])[0] + if not SMOKE: + # Search with their tuner before timing/capture; smoke only checks the default tactic. + with flashinfer_autotune(tune_mode=True): + run() + return run + + def deepgemm_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): """DeepGEMM M-grouped experts (transformers integration): FP8 128-block (UE8M0 scales on SM100) or FP4 (int8-packed weights, group-32 UE8M0) with FP8 acts.""" @@ -1185,7 +1636,7 @@ def deepgemm_bf16_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): return lambda: deepgemm_bf16_experts_forward(mod, hidden, idx, w) -def triton_kernels_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits): +def triton_kernels_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits, *_): """OpenAI triton_kernels MXFP4 experts (transformers' GPT-OSS path): fused ``matmul_ogs`` gate_up+SwiGLU → down over a swizzled MXFP4 weight layout. The quantize+swizzle is load-time weight prep (done here, before the timed closure — @@ -1293,6 +1744,9 @@ def run(): ARMS = { "finegrained-kernels": moe_fused_arm, "finegrained-kernels_unfused": moe_unfused_arm, + "finegrained-kernels-unfused": moe_unfused_arm, + "torch_weight_only": torch_weight_only_arm, + "marlin": marlin_weight_only_arm, "finegrained-fp8": fp8_fused_arm, "finegrained-fp8_unfused": fp8_unfused_arm, "transformers": transformers_arm, @@ -1300,11 +1754,15 @@ def run(): "sonicmoe": sonicmoe_arm, "trtllm_per_tensor": trtllm_per_tensor_arm, "vllm": vllm_moe_arm, + "vllm_marlin": marlin_moe_arm, + "flashinfer_mxint4": flashinfer_int4_arm, "trtllm": trtllm_moe_arm, "deepgemm": deepgemm_arm, "deepgemm_bf16": deepgemm_bf16_arm, "triton_kernels": triton_kernels_arm, "torch": torch_arm, + "torch_rowwise": rowwise_moe_arm, + "fbgemm": partial(rowwise_moe_arm, fbgemm=True), "megablocks": megablocks_arm, } @@ -1361,7 +1819,7 @@ def _weight_banks(weights, active_experts=None): return banks -def bench_modes(runs, tag): +def bench_modes(runs, tag, reference=None): """{mode: latency_us | None}, plus the eager output for parity. A mode that raises is a red CRASH cell — the other modes still run (fresh error printed inline).""" def run(): @@ -1374,19 +1832,27 @@ def run(): try: out = run() torch.accelerator.synchronize() # warm + tune before ANY timing/capture + if reference is not None: + for output in out: + error = rel_diff(output[:reference.shape[0]], reference) + if not torch.isfinite(output).all() or error > 0.05: + print(f" [{tag} dequant reference mismatch (rel {error:.2e}); skipping timings]", flush=True) + return {m: None for m in MODES}, None res["eager"] = do_bench(run, return_mode="min") * 1e3 / count print(f" {tag:14s} eager {res['eager']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} eager crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["eager"] = None - _context_poisoned(tag, "eager") + if _context_poisoned(tag, "eager"): + return dict.fromkeys(MODES), None try: res["cudagraph"] = do_bench(run, cudagraph=True, return_mode="min") * 1e3 / count print(f" {tag:14s} cudagraph {res['cudagraph']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} cudagraph crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["cudagraph"] = None - _context_poisoned(tag, "cudagraph") + if _context_poisoned(tag, "cudagraph"): + return dict.fromkeys(MODES), None try: # This wrapper is shared across every arm and regime; isolate its Dynamo guards. torch._dynamo.reset_code(run.__code__) @@ -1408,12 +1874,18 @@ def run(): drift = rel_diff(compiled, eager) if drift > 0.25: raise RuntimeError(f"compiled output diverges from eager (rel {drift:.2e})") + if reference is not None: + for output in cout: + error = rel_diff(output[:reference.shape[0]], reference) + if not torch.isfinite(output).all() or error > 0.05: + raise RuntimeError(f"compiled dequant reference mismatch (rel {error:.2e})") res["compile"] = do_bench(crun, return_mode="min") * 1e3 / count print(f" {tag:14s} compile {res['compile']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} compile crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["compile"] = None - _context_poisoned(tag, "compile") + if _context_poisoned(tag, "compile"): + return dict.fromkeys(MODES), None return res, None if out is None else out[0] @@ -1424,8 +1896,7 @@ def rel_diff(a, b): def _mock_rows(row, pname, arms, rows_out): """Figure-validation stand-in: plausible random latencies/parities through the - exact plotting path — decode vs prefill scales, a crashed finegrained-fp8 prefill-compile - cell (red X), and one wild parity (hatched bar) per row.""" + exact plotting path, including decode/prefill scales and missing compiled measurements.""" rng = random.Random(hash((row, pname)) & 0xFFFF) for regime, scale in (("decode", 100.0), ("prefill", 2000.0)): for i, name in enumerate(arms): @@ -1436,6 +1907,24 @@ def _mock_rows(row, pname, arms, rows_out): rows_out.append((row, pname, regime, _impl(name), res, parity)) +def _integer_moe_reference(cfg, hidden, idx, weights, gu, gus, dn, dns, guz, dnz): + # A short prefix exercises identical routes without dequantizing a permanent expert bank. + hidden, idx, weights = hidden[:32], idx[:32], weights[:32] + out = torch.zeros_like(hidden, dtype=torch.float32) + for expert in idx.unique().tolist(): + token, slot = torch.where(idx == expert) + gate_up = dq_int(gu[expert], gus[expert], None if guz is None else guz[expert], + dtype=hidden.dtype, group=cfg["group"]) + down = dq_int(dn[expert], dns[expert], None if dnz is None else dnz[expert], + dtype=hidden.dtype, group=cfg["group"]) + projected = (hidden[token].float() @ gate_up.t()).to(hidden.dtype) + gate, up = projected.chunk(2, -1) + intermediate = _glu(gate.float(), up.float(), cfg).to(hidden.dtype) + value = (intermediate.float() @ down.t()).to(hidden.dtype) + out.index_add_(0, token, value.float() * weights[token, slot, None]) + return out.to(hidden.dtype) + + def bench_problem_row(row, pname, cfg, arms, weights, rows_out): """One (row, problem): both regimes, finegrained-kernels-first (parity anchor), streaming prints.""" print(f"== [{row}] {pname}") @@ -1454,10 +1943,15 @@ def bench_problem_row(row, pname, cfg, arms, weights, rows_out): inputs = [(hidden, idx, w, logits)] + [ _clone_tensors((hidden, idx, w, logits)) for _ in banks[1:] ] + first_args = (cfg, grouped, hidden, idx, w, *banks[0][:6], logits, *banks[0][6:]) + reference = torch_weight_only_arm(*first_args)() if "torch_weight_only" in cfg["baselines"] else None + if cfg["weights"] in INT_BITS: + gu, gus, dn, dns, _, _, guz, dnz = banks[0] + reference = _integer_moe_reference(cfg, hidden, idx, w, gu, gus, dn, dns, guz, dnz) anchor_res, anchor_out = None, None for name in arms: try: - runs = [ARMS[name](cfg, grouped, h, i, rw, *bank, rl) + runs = [ARMS[name](cfg, grouped, h, i, rw, *bank[:6], rl, *bank[6:]) for bank, (h, i, rw, rl) in zip(banks, inputs)] except Exception as e: print(f" [{name} setup failed: {type(e).__name__}: {str(e)[:90]}]", @@ -1465,7 +1959,7 @@ def bench_problem_row(row, pname, cfg, arms, weights, rows_out): rows_out.append((row, pname, regime, _impl(name), {m: None for m in MODES}, None)) continue - res, out = bench_modes(runs, name) + res, out = bench_modes(runs, name, reference) parity = None if anchor_res is None: anchor_res, anchor_out = res, out @@ -1486,9 +1980,88 @@ def _impl(arm_name): return arm_name.replace("_unfused", "").replace("_bf16", "") +def _svdquant_arms(x, weights, act_global): + W, Ws, Wg, smooth, lora_a, lora_b = weights + Ws_local = fgm.swizzle_mx_scales(Ws) if PRESWIZZLE else Ws + _mark_static(*weights, Ws_local, act_global) + runs = {"finegrained-kernels": lambda: fgm.matmul_2d( + x, W, Bs=Ws_local, a_global_scale=act_global, b_global_scale=Wg, + a_pre_quant_scale=smooth, svdquant_lora_a=lora_a, svdquant_lora_b=lora_b, + simulate_unfused=True)} + if (flashinfer_svdquant_linear is not None and torch.version.cuda is not None + and torch.cuda.get_device_capability()[0] == 10): + # FlashInfer's public chain pre-folds these transforms once at load time. + alpha = act_global * Wg + fi_scales = fgm.swizzle_mx_scales(Ws).view(torch.uint8).flatten() + fi_down = (smooth[:, None] * lora_a.T).contiguous() + fi_up = (lora_b.float() / alpha).bfloat16() + fi_global = act_global.reciprocal() + _mark_static(alpha, fi_scales, fi_down, fi_up, fi_global) + runs["flashinfer_svdquant"] = lambda: flashinfer_svdquant_linear( + x, W.view(torch.uint8), fi_scales, alpha, smooth, fi_down, fi_up, fi_global) + if not SMOKE: + with flashinfer_svdquant_autotune(): + runs["flashinfer_svdquant"]() + return runs + + +def bench_svdquant_row(row, pname, cfg, rows_out): + """Whole dense linear: smoothing/activation quant, BF16 down, NVFP4 GEMM plus up.""" + print(f"== [{row}] {pname}") + if MOCK: + _mock_rows(row, pname, ("finegrained-kernels",) + cfg["baselines"], rows_out) + return + N, K, rank = cfg["N"], cfg["K"], cfg["svdquant_rank"] + torch.manual_seed(0) + W, Ws, Wg = WEIGHTS["nvfp4"]["make"](N, K, 1) + W, Ws = W[0], Ws[0] + smooth = torch.linspace(0.5, 1.5, K, device=DEVICE, dtype=torch.bfloat16) + lora_a = torch.randn(rank, K, device=DEVICE, dtype=torch.bfloat16) / K**0.5 + lora_b = torch.randn(N, rank, device=DEVICE, dtype=torch.bfloat16) * 20 + weights = (W, Ws, Wg, smooth, lora_a, lora_b) + for regime, tokens in (("decode", DECODE_TOKENS), ("prefill", PREFILL_TOKENS)): + torch._dynamo.reset() + x = torch.randn(tokens, K, device=DEVICE, dtype=torch.bfloat16) + # Stand-in for checkpoint calibration, outside timing. All arms and banks share it. + act_global = ((x * smooth).abs().amax().float() / (6 * 448)).reshape(1) + banks = _weight_banks(weights) + inputs = [x] + [x.clone() for _ in banks[1:]] + bank_arms = [_svdquant_arms(hidden, bank, act_global) for bank, hidden in zip(banks, inputs)] + # Sample rows to bound reference memory without omitting timed work. + xs = x[:32] * smooth + aq, asc = nvfp4_act_quant(xs, global_scale=act_global) + residual = (dq_grouped(aq, asc, 16) * act_global) @ dq_nvfp4_two_level(W, Ws, Wg)[0].T + down = (xs.float() @ lora_a.float().T).bfloat16() + correction = (down.float() @ lora_b.float().T).bfloat16() + reference = residual.bfloat16() + correction + for name in bank_arms[0]: + runs = [arms[name] for arms in bank_arms] + res, out = bench_modes(runs, name, reference) + parity = rel_diff(out[:reference.shape[0]], reference) if out is not None else None + rows_out.append((row, pname, regime, name, res, parity)) + del banks, inputs, bank_arms + + def _attn_arms(cfg, x, weights): N, K, block = cfg["N"], cfg["K"], cfg["block"] - W, Ws, W_g = weights + W, Ws, W_g, Wz = weights + marlin = None + if "vllm_marlin" in cfg["baselines"] and apply_gptq_marlin_linear is None: + print(" [vLLM Marlin unavailable: install a CUDA vLLM build]", flush=True) + if "vllm_marlin" in cfg["baselines"] and apply_gptq_marlin_linear is not None: + marlin = _marlin_weight(cfg, W, Ws, Wz) + marlin_workspace = marlin_make_workspace_new(W.device) + marlin_empty = torch.empty(0, device=W.device, dtype=torch.int32) + marlin_kind = _marlin_type(cfg) + _mark_static(*marlin, marlin_workspace, marlin_empty) + exllama = None + if "vllm_exllama" in cfg["baselines"]: + if gptq_gemm is None: + print(" [ExLlama unavailable: install a CUDA vLLM build]", flush=True) + else: + exllama = _exllama_weight(cfg, W, Ws) + exllama_bits = INT_BITS[cfg["weights"]] + _mark_static(*exllama) Ws_fp8 = _fp8_scales(Ws, block) dg_block = tuple(block) if block else None # deployment layout for the LOCAL arm only (every baseline reads the raw row-major scale): @@ -1525,7 +2098,7 @@ def _attn_arms(cfg, x, weights): if cfg["weights"] in _MX_WEIGHTS and _quantized_acts(cfg) and not (MOCK or REPLOT): # both are version-gated: scaled_mm's scaling enums land in recent torch only, and # torchao is an optional dep — importing either at module scope would gate the whole bench - from torch.nn.functional import ScalingType, SwizzleType + from torch.nn.functional import SwizzleType from torchao.prototype.mx_formats.utils import to_blocked nvfp4_row = cfg["weights"] == "nvfp4" @@ -1558,12 +2131,31 @@ def torch_mm(a, a_scale=None): return torch.nn.functional.scaled_mm( aq, mat_b, sa, rb, sb, rb, swizzle_a=swz, swizzle_b=swz, output_dtype=torch.bfloat16) + elif cfg["weights"] == "fp8_rowwise": + mat_b = W.t() + _mark_static(mat_b) + scale_ub = torch.tensor([cfg["activation_scale_ub"]], device=W.device) + _mark_static(scale_ub) + + def torch_mm(a, a_scale=None): + aq, a_s = fp8_act_quant_tensor_wide(a, K, amax_upper_bound=cfg["activation_scale_ub"]) + return scaled_mm(aq, mat_b, a_s, ScalingType.RowWise, Ws.t(), ScalingType.RowWise, + output_dtype=a.dtype, use_fast_accum=True) + + def fbgemm_mm(a): + aq, a_s = _fbgemm_rowwise_quant(a, scale_ub=scale_ub) + return _fbgemm_rowwise(aq, W, a_s, Ws.flatten(), use_fast_accum=True) + + def vllm_mm(a): + aq, a_s = _vllm_tensor_quant(a, scale_ub=scale_ub, use_per_token_if_dynamic=True) + return _vllm_scaled_mm(aq, mat_b, a_s, Ws.t(), out_dtype=a.dtype) + if scaled_mm is None: + torch_mm = None elif cfg["weights"] == "fp8_tensor" and cfg.get("static") and not (MOCK or REPLOT): # per-tensor static: one fp32 scale per operand IS scaled_mm's TensorWise, so this row # has a like-for-like cuBLAS bar after all. The calibrated activation scale is passed in # rather than derived, so neither arm pays an amax pass the other doesn't. from finegrained_kernels.quant import quantize_rows_static - from torch.nn.functional import ScalingType mat_b = W.t() _mark_static(mat_b) @@ -1575,10 +2167,18 @@ def torch_mm(a, a_scale): Ws.reshape(()), ScalingType.TensorWise, output_dtype=torch.bfloat16) + if "marlin" in cfg["baselines"] and fused_marlin_moe is not None: + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(W, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(Ws, requires_grad=False) + layer.output_size_per_partition, layer.input_size_per_partition = N, K + layer.weight_block_size, layer.orig_dtype = block, torch.bfloat16 + prepare_fp8_layer_for_marlin(layer, size_k_first=False) + _mark_static(layer.weight, layer.weight_scale, layer.workspace) # a STATIC row's calibrated activation scale replaces the inline quant: one value, since a # dense module is one quantized module (the MoE's per-expert form has no dense analogue) attn_act_scale = ( - (x.float().abs().amax() / 448.0).clamp(min=1e-12).reshape(1) + (x.float().abs().amax() / (127.0 if _activation_format(cfg) == "int8" else 448.0)).clamp(min=1e-12).reshape(1) if cfg.get("static") else None ) if attn_act_scale is not None: @@ -1587,11 +2187,31 @@ def torch_mm(a, a_scale): # activation precision (None follows the weight format). attn_arms = { "finegrained-kernels": lambda: fgm.matmul_2d( - x, W, attn_act_scale, Ws_fgm, activation_format=_activation_format(cfg), output_dtype=torch.bfloat16, - b_global_scale=W_g), + x, W, attn_act_scale, Ws_fgm, activation_format=_activation_format(cfg), output_dtype=x.dtype, + activation_scale_ub=cfg.get("activation_scale_ub"), + b_global_scale=W_g, b_zero=Wz), } + if exllama is not None: + attn_arms["vllm_exllama"] = lambda: gptq_gemm(x, *exllama, True, False, exllama_bits) + if marlin is not None: + attn_arms["vllm_marlin"] = lambda: apply_gptq_marlin_linear( + x, marlin[0], marlin[1], marlin[2], marlin_empty, marlin_empty, + marlin_workspace, marlin_kind, N, K, True) + if "torch_weight_only" in cfg["baselines"]: + attn_arms["torch_weight_only"] = lambda: torch.nn.functional.linear( + x, dequantize_fp8_weight(W, Ws, block) + ) + if "marlin" in cfg["baselines"] and fused_marlin_moe is not None: + attn_arms["marlin"] = lambda: apply_fp8_marlin_linear( + x, layer.weight, layer.weight_scale, layer.workspace, N, K, None + ) if torch_mm is not None: attn_arms["torch_mm"] = lambda: torch_mm(x, attn_act_scale) + if cfg["weights"] == "fp8_rowwise": + if "fbgemm" in cfg["baselines"] and _fbgemm_rowwise is not None: + attn_arms["fbgemm"] = lambda: fbgemm_mm(x) + if "vllm_mm" in cfg["baselines"] and vllm_fused_experts is not None: + attn_arms["vllm_mm"] = lambda: vllm_mm(x) if "finegrained-fp8" in cfg["baselines"]: attn_arms["finegrained-fp8"] = lambda: upstream_fp8.matmul_2d(x, W, Ws_fp8, block, torch.bfloat16) @@ -1622,14 +2242,27 @@ def torch_mm(a, a_scale): def bench_attn_row(row, pname, cfg, rows_out): """One attn linear per model, in its deployment format (same weights across impls; the finegrained-kernels arm's ``activation_format`` follows the model — GPT-OSS runs W4A4).""" + if cfg.get("svdquant_rank"): + return bench_svdquant_row(row, pname, cfg, rows_out) print(f"== [{row}] {pname}") if MOCK: _mock_rows(row, pname, ("finegrained-kernels",) + cfg["baselines"], rows_out) return torch._dynamo.reset() N, K = cfg["N"], cfg["K"] - W_g = None - if cfg["weights"] == "fp8_128x128_ue8m0": + W_g = Wz = None + if cfg["weights"] in INT_BITS: + torch.manual_seed(0) + W, Ws, Wz = _integer_weight(cfg, N, K, 1) + W, Ws = W[0], Ws[0] + Wz = None if Wz is None else Wz[0] + if _activation_format(cfg) in ("fp8", "int8"): + Ws = Ws.float() + if _activation_format(cfg) == "fp8": + weight = dq_int(W, Ws, Wz, group=cfg["group"]) + W_g = weight.abs().amax().div(448).clamp(min=1e-12).reshape(1) + del weight + elif cfg["weights"] == "fp8_128x128_ue8m0": W, Ws = make_weights(N, K, DEVICE, [128, 128], scale_dtype=torch.float8_e8m0fnu) else: @@ -1639,20 +2272,36 @@ def bench_attn_row(row, pname, cfg, rows_out): W, Ws, W_g = WEIGHTS[cfg["weights"]]["make"](N, K, 1) W, Ws = W[0], Ws[0] W_g = W_g if W_g is None else W_g[0] + if "weight_std" in cfg: + Ws = Ws * cfg["weight_std"] _mark_static(W, Ws) - weights = (W, Ws, W_g) + weights = (W, Ws, W_g, Wz) for regime, tokens in (("decode", DECODE_TOKENS), ("prefill", PREFILL_TOKENS)): print(f" -- {regime}") torch.manual_seed(0) - x = torch.randn(tokens, K, device=DEVICE, dtype=torch.bfloat16) + x = torch.randn(tokens, K, device=DEVICE, dtype=cfg.get("dtype", torch.bfloat16)) banks = _weight_banks(weights) inputs = [x] + [x.clone() for _ in banks[1:]] bank_arms = [_attn_arms(cfg, hidden, bank) for bank, hidden in zip(banks, inputs)] + reference = bank_arms[0]["torch_weight_only"]() if "torch_weight_only" in bank_arms[0] else None + if cfg["weights"] in INT_BITS: + fmt = _activation_format(cfg) + dequant = dq_int(W, Ws, Wz, dtype=None if fmt in ("fp8", "int8") else x.dtype, group=cfg["group"]) + reference_input = x[:32].float() + if cfg.get("static"): + act_scale = (x.float().abs().amax() / (127.0 if fmt == "int8" else 448.0)).clamp(min=1e-12).reshape(1) + if fmt == "fp8": + dequant = (dequant / W_g).clamp(-448, 448).to(torch.float8_e4m3fn).float() * W_g + reference_input = (reference_input / act_scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * act_scale + elif fmt == "int8": + reference_input = (reference_input / act_scale).round().clamp(-128, 127) * act_scale + reference = (reference_input @ dequant.t()).to(x.dtype) + del dequant anchor_res, anchor_out = None, None for name in bank_arms[0]: runs = [arms[name] for arms in bank_arms] - res, out = bench_modes(runs, name) + res, out = bench_modes(runs, name, reference) parity = None if anchor_res is None: anchor_res, anchor_out = res, out @@ -1671,7 +2320,7 @@ def bench_attn_row(row, pname, cfg, rows_out): print(f"device: {device_name} torch {torch.__version__} cache={CACHE_POLICY}" f"{' [SMOKE]' if SMOKE else ''}") print("finegrained-kernels = local build; baselines: finegrained-fp8 (upstream), DeepGEMM, " - "vLLM fused MoE, FlashInfer's TRT-LLM routed MoE, transformers grouped_mm/batched_mm, " + "vLLM fused MoE/Marlin/ExLlama, FlashInfer's TRT-LLM routed MoE/MXINT4, transformers grouped_mm/batched_mm, " "SonicMoE, OpenAI triton_kernels, torch.scaled_grouped_mm, " "nvfp4-gemm (transformers' NVFP4Linear kernel), megablocks (bf16 dMoE)" f"{f' | {GPUS} GPUs' if GPUS > 1 else ''}\n") @@ -1683,7 +2332,10 @@ def wanted(*names): # ── multi-GPU: GPUS>1 shards the per-problem tasks across GPUs, one process per GPU (each # owning one device via CUDA_VISIBLE_DEVICES). The coordinator spawns GPUS workers, each # writes a shard CSV, then the coordinator merges + plots. A single GPU (GPUS=1) runs inline. ── -_CSV = os.path.join(_HERE, "bench_moe.csv") +suffix = ARGS.suffix if ARGS.suffix is not None else ( + "_mock" if MOCK else ("_partial" if FILTERS else "") +) +_CSV = os.path.join(_HERE, f"bench_moe{suffix}.csv") _CSV_HEADER = ("category,problem,regime,impl,mode,latency_us,parity_vs_finegrained_kernels," "cache_policy\n") @@ -1738,6 +2390,25 @@ def _allowed(cfg): def _run_task(kind, pname, cfg, rows_out): """Run one problem's row(s) into rows_out (build its weights once, bench its arms).""" + global upstream_fp8, triton_kernels_hub, _megablocks, _nvfp4_gemm + if not (MOCK or REPLOT): + if "finegrained-fp8" in cfg["baselines"] and upstream_fp8 is None: + upstream_fp8 = get_kernel("kernels-community/finegrained-fp8", revision=UPSTREAM_FP8_REV, + trust_remote_code=TRUSTED) + if "triton_kernels" in cfg.get("fused_extra", ()) + cfg["baselines"] and triton_kernels_hub is None: + triton_kernels_hub = get_kernel("kernels-community/gpt-oss-triton-kernels", version=1, + trust_remote_code=TRUSTED) + _tfmx.triton_kernels_hub = triton_kernels_hub + if kind == "bf16" and _megablocks is None: + try: + _megablocks = get_kernel("kernels-community/megablocks", version=1, trust_remote_code=TRUSTED) + except Exception: + pass + if kind == "attn" and cfg["weights"] == "nvfp4" and not cfg.get("svdquant_rank") and _nvfp4_gemm is None: + try: + _nvfp4_gemm = _hub_kernel("nvfp4") + except Exception: + pass if kind in ("moe", "bf16"): try: weights = None if MOCK else build(cfg) @@ -1787,15 +2458,19 @@ def _run_task(kind, pname, cfg, rows_out): trtllm_gated_gap = (cfg["weights"] not in ("mxfp4", "nvfp4") and not (cfg["swiglu_alpha"] is None and cfg["swiglu_limit"] is None)) unavailable = (("vllm",) if vllm_fused_experts is None else ()) + ( - ("trtllm",) if trtllm_fp8_block_scale_routed_moe is None or trtllm_gated_gap else ()) - quant_arms = tuple(a for a in quant_arms if a not in unavailable) - # Arms whose kernel is known to TRAP asynchronously go LAST: a trap poisons the CUDA - # context, which is unrecoverable in-process, so every arm benched after it blanks - # through no fault of its own (that is how finegrained-fp8's v4 int32 expert-offset - # overflow — fixed in a5cf637f, awaiting the hub rebuild — took out the deepgemm/trtllm/ - # transformers@main prefill cells on DeepSeek-V4 W4A8). Ordering does not change any - # measurement (each arm is benched independently in the same warmed process); it only - # changes who a trap can hurt. Drop an arm from this set once its fixed build ships. + ("trtllm",) if trtllm_fp8_block_scale_routed_moe is None or trtllm_gated_gap else ()) + ( + ("fbgemm",) if _fbgemm_rowwise_grouped is None else ()) + ( + ("torch_rowwise",) if scaled_mm is None else ()) + unavailable += (("vllm_marlin",) if fused_marlin_moe is None or apply_gptq_marlin_linear is None else ()) + unavailable += (("flashinfer_mxint4",) if ( + trtllm_mxint4_block_scale_routed_moe is None or DEVICE != "cuda" + or ACCELERATOR.get_device_capability()[0] != 10) else ()) + for name in cfg["baselines"]: + if name in unavailable: + print(f" [{name} unavailable on this installation]", flush=True) + quant_arms = tuple(a for a in quant_arms if a not in unavailable + and (a != "marlin" or fused_marlin_moe is not None and apply_fp8_marlin_linear is not None)) + # Keep the legacy Hub arm last so an unexpected CUDA fault preserves other timings. _TRAP_RISK = ("finegrained-fp8",) quant_arms = (tuple(a for a in quant_arms if a not in _TRAP_RISK) + tuple(a for a in quant_arms if a in _TRAP_RISK)) @@ -1827,10 +2502,10 @@ def _run_task(kind, pname, cfg, rows_out): rows = [] if REPLOT: rows = _load_rows_csv(_CSV) -elif GPUS > 1 and _SHARD is None and not MOCK: - # COORDINATOR: fan the tasks across GPUS subprocesses (one device each), merge, then plot. +elif _SHARD is None and not MOCK: + # A fresh process per case keeps a faulting optional baseline from poisoning later cases. - shard_paths = [os.path.join(_HERE, f"bench_moe.shard{g}.csv") for g in range(GPUS)] + shard_paths = [os.path.join(_HERE, f"bench_moe{suffix}.shard{i}.csv") for i in range(len(TASKS))] for sp in shard_paths: # a leftover shard from a prior run must never merge as fresh data if os.path.exists(sp): os.unlink(sp) @@ -1840,15 +2515,23 @@ def _run_task(kind, pname, cfg, rows_out): if devices and len(devices) != GPUS: raise SystemExit(f"--devices lists {len(devices)} device(s) but --gpus is {GPUS}") worker_flags = (["--smoke"] if SMOKE else []) + ([] if PRESWIZZLE else ["--no-preswizzle"]) - procs = [subprocess.Popen( - [sys.executable, os.path.abspath(__file__), "--gpus", "1", *worker_flags, *FILTERS], - env={**os.environ, DEVICE_MASK_ENV: devices[g] if devices else str(g), - "BENCH_SHARD": f"{g}/{GPUS}"}) for g in range(GPUS)] - nfail = sum(p.wait() != 0 for p in procs) - missing = [g for g, sp in enumerate(shard_paths) if not os.path.exists(sp)] + if ARGS.suffix is not None: + worker_flags += ["--suffix", ARGS.suffix] + + def _run_device(g): + device = devices[g] if devices else (os.environ.get(DEVICE_MASK_ENV, "0") if GPUS == 1 else str(g)) + return sum(subprocess.run( + [sys.executable, os.path.abspath(__file__), "--gpus", "1", *worker_flags, *FILTERS], + env={**os.environ, DEVICE_MASK_ENV: device, "BENCH_SHARD": f"{i}/{len(TASKS)}"}, + ).returncode != 0 for i in range(g, len(TASKS), GPUS)) + + with ThreadPoolExecutor(max_workers=GPUS) as pool: + nfail = sum(pool.map(_run_device, range(GPUS))) + missing = [g for g, sp in enumerate(shard_paths) + if not os.path.exists(sp) or os.path.getsize(sp) <= len(_CSV_HEADER)] if nfail or missing: raise SystemExit( - f"{nfail}/{GPUS} shard worker(s) exited non-zero; shards missing: {missing} — " + f"{nfail}/{len(TASKS)} task worker(s) exited non-zero; shards missing/empty: {missing} — " "results are incomplete, not merging (rerun; a partial figure misreads as a full one)" ) with open(_CSV, "w") as out: # merge shard CSVs (skip repeated headers) @@ -1859,17 +2542,20 @@ def _run_task(kind, pname, cfg, rows_out): out.writelines(f) rows = _load_rows_csv(_CSV) else: - # single process (GPUS=1) OR one shard worker (BENCH_SHARD="g/n") + # Mock run or a worker assigned one case by the coordinator. shard, nshards = (int(x) for x in _SHARD.split("/")) if _SHARD else (0, 1) - for i, (kind, pname, cfg) in enumerate(TASKS): - if i % nshards == shard: - _run_task(kind, pname, cfg, rows) - if _SHARD is not None: # WORKER: write only this shard, then exit (coordinator merges+plots) - _write_rows_csv(os.path.join(_HERE, f"bench_moe.shard{shard}.csv"), rows) + try: + for i, (kind, pname, cfg) in enumerate(TASKS): + if i % nshards == shard: + _run_task(kind, pname, cfg, rows) + finally: + if _SHARD is not None: # Preserve completed measurements even if this case fails. + _write_rows_csv(os.path.join(_HERE, f"bench_moe{suffix}.shard{shard}.csv"), rows) + if _SHARD is not None: # The coordinator merges and plots. sys.exit(0) -# ── figure: ONE png — 4 rows x (decode | prefill), linear axes, red CRASH markers ── +# ── Shared figure: categories × (decode | prefill) ── import matplotlib # noqa: E402 @@ -1880,6 +2566,8 @@ def _run_task(kind, pname, cfg, rows_out): # has only the two-GEMM shape, so that is what it contributes. No fused/unfused split in the # figure — it is an internal distinction of ours that most baselines do not have. ROW_ORDER = ["quantized", "linear quantized", "unquantized"] +ROW_LABELS = {"quantized": "MoE quantized", "linear quantized": "Linear quantized", + "unquantized": "MoE unquantized (BF16)"} present_rows = [r for r in ROW_ORDER if any(rr == r for rr, *_ in rows)] # bars must be PHYSICALLY identical across every subplot of every figure: one slot @@ -1896,65 +2584,57 @@ def _run_task(kind, pname, cfg, rows_out): default=1, ) - 1 -# every (row, problem, regime, impl) x mode -> CSV, so the full 3-mode numbers survive even -# though the figure shows one deployment mode per regime. REPLOT reads this CSV as its source -# (no write), and the multi-GPU coordinator already merged the shard CSVs into it — so only the -# single-process run writes here. +# Measured runs already merged their per-case CSVs; replot only reads them. +# Mock timings still need their CSV written here. # --suffix lets filtered rows run concurrently on separate GPUs without racing for the same # partial CSV/PNG (splice the pieces into bench_moe.csv afterwards) -suffix = ARGS.suffix if ARGS.suffix is not None else ( - "_mock" if MOCK else ("_partial" if FILTERS else "") -) -_via_coordinator = GPUS > 1 and _SHARD is None and not MOCK # merged the CSV already -if not REPLOT and _SHARD is None and not _via_coordinator: +if MOCK and not REPLOT: _write_rows_csv(os.path.join(_HERE, f"bench_moe{suffix}.csv"), rows) -# ONE figure, 8 panels: rows = the 4 categories, cols = (decode | prefill). Each -# bar is colored by impl (finegrained-kernels leftmost, fixed slots). DECODE superposes its two -# graph-captured modes — cudagraph as the solid fill, compile as a black hatched -# outline over the same slot — so the gap between them is visible per impl. PREFILL -# is eager only (single solid) — prefill is not compiled in deployment, so a compile -# bar there would show a mode nobody ships. Red X = crashed (no latency). This chart is latency -# only; parity vs finegrained-kernels lives in the bench log + the CSV beside this png. +# Each implementation uses the same modes in every category and format. The compiled +# measurement is an outline over its eager/pure-graph counterpart when compilation wins. import matplotlib.patches as mpatches # noqa: E402 +from matplotlib.lines import Line2D # noqa: E402 -# (solid_mode, overlay_mode) per regime -REGIME_MODES = {"decode": ("cudagraph", "compile"), "prefill": ("eager", None)} +REGIME_MODES = {"decode": ("cudagraph", "compile"), "prefill": ("eager", "compile")} fig, axes = plt.subplots(max(len(present_rows), 1), 2, - figsize=(18, 4.6 * max(len(present_rows), 1)), + figsize=(max(18, 2.2 * (GLOBAL_SPAN + 1)), 4.6 * max(len(present_rows), 1)), squeeze=False) +any_overlay = any_eager_fallback = any_missing = False for ri, row in enumerate(present_rows): for ci, regime in enumerate(("decode", "prefill")): ax = axes[ri][ci] solid_mode, overlay_mode = REGIME_MODES[regime] cells = [(p, i, r, par) for (rr, p, reg, i, r, par) in rows if rr == row and reg == regime] + floor = min((v for _, _, res, _ in cells for v in res.values() if v is not None and v > 0), default=1) / 2 + ax.set_yscale("log") + ax.grid(axis="y", alpha=0.2, zorder=0) problems = list(dict.fromkeys(p for p, *_ in cells)) # ONE fixed model order across every row (was per-panel support-sort, which # reordered models between rows and read as confusing). Roughly most-supported # first, finegrained-kernels-only (GPT-OSS, GLM-NVFP4) last; keyed on the base model so the # quantized/attn (GLM-5.2-NVFP4) and unquantized (GLM-5.2) rows line up. def _model_rank(p): - name = p.split(" ")[0].split("/")[-1] + name = p.split("/")[-1] for i, m in enumerate(CANONICAL_MODEL_ORDER): if name.startswith(m): return i return len(CANONICAL_MODEL_ORDER) problems.sort(key=_model_rank) - # legend order stays canonical; the BARS are speed-ranked per group below - row_impls = list(dict.fromkeys(i for _, i, *_ in cells)) - labeled = set() - overlay_drawn = False # any compile-beats-cudagraph overlay in this panel? + overlay_drawn = False eager_only_drawn = False # any bar shown as eager because cudagraph is unsupported? ticks, ticklabels = [], [] for gi, pname in enumerate(problems): ticks.append(gi) - # compact 2-line tick: model short-name + format (full ids in the log) + # Wrap model and recipe names; full ids remain in the CSV. short = pname.split(" (")[0] model, _, rest = short.partition(" ") + if short.startswith("unknown model "): + model, rest = "unknown model", short.removeprefix("unknown model ") model = model.split("/")[-1] rest = rest.removeprefix("attn ").removesuffix(" qkv-shaped") - ticklabels.append(f"{model}\n{rest}" if rest else model) + ticklabels.append(fill(model, 18) + (f"\n{fill(rest, 18)}" if rest else "")) # bars speed-ranked within the group (fastest leftmost, crashed last), # centered on the group's own arm count — colors identify the impls group = [(i, r) for p, i, r, _par in cells if p == pname] @@ -1972,27 +2652,22 @@ def _drawn(res): eager_fb = (res.get("eager") if sval is None and regime == "decode" else None) if sval is not None: ax.bar(gi + off, sval, GLOBAL_WIDTH, color=IMPL_COLORS[impl], - label=impl if impl not in labeled else None, zorder=2) - labeled.add(impl) + zorder=2) elif eager_fb is not None: # Cannot be graph-captured (megablocks' routing does a device->host copy; # nvfp4_gemm's per-expert loop has data-dependent shapes) but DOES run. Show # the eager cost, dot-hatched so it is never read as a cudagraph number — # a red X here would say "broken" about something merely 5x slower. ax.bar(gi + off, eager_fb, GLOBAL_WIDTH, color=IMPL_COLORS[impl], - alpha=0.55, edgecolor="black", hatch="..", linewidth=0.7, zorder=2, - label=impl if impl not in labeled else None) - labeled.add(impl) + alpha=0.55, edgecolor="black", hatch="..", linewidth=0.7, zorder=2) eager_only_drawn = True else: - # nothing ran in any mode -> red X at the baseline - ax.plot(gi + off, 0, "x", color="red", markersize=9, + # Deployment mode unavailable: keep its marker above the log-axis floor. + ax.plot(gi + off, floor * 1.25, "x", color="red", markersize=9, clip_on=False, zorder=4) - # compile overlay (decode): hatched outline — shown ONLY when compile - # beats cudagraph (a real win); dropped otherwise (compile's usual - # decode regression is the fixed per-call wrapper overhead, which - # amortizes at model scale — see notes, not worth cluttering the bar) - oval = res.get(overlay_mode) if overlay_mode else None + any_missing = True + # Overlay compile wins on the same implementation's eager/graph bar. + oval = res.get(overlay_mode) if oval is not None and sval is not None and oval < sval: ax.bar(gi + off, oval, GLOBAL_WIDTH, facecolor="none", edgecolor="black", hatch="////", linewidth=0.8, zorder=3) @@ -2007,27 +2682,31 @@ def _drawn(res): ax.set_xlim(-0.6 - margin, extent + 0.6 + margin) ax.set_xticks(ticks) ax.set_xticklabels(ticklabels, fontsize=7) - ax.set_ylabel("latency (us)") + ax.set_ylim(bottom=floor) + ax.set_ylabel("latency (us, log scale)") tok = DECODE_TOKENS if regime == "decode" else PREFILL_TOKENS - mode_note = ("solid=cudagraph, hatch=compile-if-faster" if overlay_drawn - else (solid_mode if not overlay_mode else "cudagraph")) - ax.set_title(f"{row} — {regime} (T={tok}; {mode_note})") - # legend: impl colors + the mode/marker key (compile entry only if any - # panel bar actually had a compile win to show) - handles = [mpatches.Patch(color=IMPL_COLORS[i], label=_impl_label(i, regime)) - for i in row_impls] - if overlay_drawn: - handles.append(mpatches.Patch(facecolor="none", edgecolor="black", - hatch="////", label="compile (faster)")) - if eager_only_drawn: - handles.append(mpatches.Patch(facecolor="none", edgecolor="black", - hatch="..", label="eager (no cudagraph)")) - ax.legend(handles=handles, loc="upper left", fontsize=8) + mode_note = (f"solid={solid_mode}, hatch=compile-if-faster" if overlay_drawn + else solid_mode) + ax.set_title(f"{ROW_LABELS[row]} — {regime} (T={tok}; {mode_note})") + any_overlay |= overlay_drawn + any_eager_fallback |= eager_only_drawn + +# One legend serves all panels; none of the model groups are hidden behind it. +impls = list(dict.fromkeys(impl for _, _, _, impl, *_ in rows)) +handles = [mpatches.Patch(color=IMPL_COLORS[i], label=_impl_label(i)) for i in impls] +if any_overlay: + handles.append(mpatches.Patch(facecolor="none", edgecolor="black", hatch="////", label="compile (faster)")) +if any_eager_fallback: + handles.append(mpatches.Patch(facecolor="none", edgecolor="black", hatch="..", label="eager (no cudagraph)")) +if any_missing: + handles.append(Line2D([], [], color="red", marker="x", linestyle="none", label="unavailable")) +fig.legend(handles=handles, loc="lower center", ncol=4, fontsize=8, frameon=False) +legend_height = 0.25 * ((len(handles) + 3) // 4) / fig.get_figheight() cache_note = "cold weights (>2xL2)" if DEVICE == "cuda" else "backend cache defaults" -fig.suptitle(f"MoE bench — finegrained-kernels vs finegrained-fp8 + references " - f"({device_name}, real model shapes; decode=cudagraph+compile, " - f"prefill=eager; {cache_note})", y=0.9995) -fig.tight_layout(rect=(0, 0, 1, 0.99)) +fig.suptitle(f"GEMM + MoE quantization — finegrained-kernels vs references " + f"({device_name}, real model shapes; decode=cudagraph, prefill=eager; " + f"hatch=compile-if-faster; {cache_note})", y=0.9995) +fig.tight_layout(rect=(0, legend_height + 0.02, 1, 0.98)) out_png = os.path.join(_HERE, f"bench_moe{suffix}.png") fig.savefig(out_png, dpi=120) print(f"wrote {out_png}" + ("" if REPLOT else f" + bench_moe{suffix}.csv")) diff --git a/finegrained-kernels/tests/test_act_quant.py b/finegrained-kernels/tests/test_act_quant.py index 841887158..d01023f5a 100644 --- a/finegrained-kernels/tests/test_act_quant.py +++ b/finegrained-kernels/tests/test_act_quant.py @@ -21,7 +21,7 @@ import triton import triton.language as tl -from finegrained_kernels.quant import fp8_act_quant_block_dynamic, fp8_act_quant_tensor_wide, mx_act_quant_inline, mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant, nvfp4_quantize_two_level # type: ignore +from finegrained_kernels.quant import _quant_block_k_pruner, fp8_act_quant_block_dynamic, fp8_act_quant_tensor_wide, mx_act_quant_inline, mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant, nvfp4_quantize_two_level # type: ignore _FP8_DTYPE = torch.float8_e4m3fn @@ -415,8 +415,6 @@ def test_the_grouped_quant_grid_admits_only_its_tile_height(): ``BLOCK_T`` would cover fewer rows than the schedule holds and leave the tail unquantized. The dense grid has no such tie and keeps the whole range. Both scale stores reach the grouped grid, so the pin cannot hang off the swizzle.""" - from finegrained_kernels.quant import _quant_block_k_pruner - configs = [ triton.Config({"BLOCK_K": bk, "BLOCK_T": bt}, num_warps=4) for bk in (32, 64) diff --git a/finegrained-kernels/tests/test_backward.py b/finegrained-kernels/tests/test_backward.py index c9c0d8fd3..29ffa3ab8 100644 --- a/finegrained-kernels/tests/test_backward.py +++ b/finegrained-kernels/tests/test_backward.py @@ -35,6 +35,7 @@ SUPPORTS_FP8, TEST_DEVICE, WEIGHTS, + INT_BITS, dequantize_weight, make_weights, quant_dequant_a, @@ -45,6 +46,7 @@ dgrad_matmul_2d, dgrad_matmul_batched, dgrad_matmul_grouped, + dgrad_matmul_grouped_kernel, glu_backward, ) @@ -259,7 +261,8 @@ def test_inference_path_is_untouched(): assert not out.requires_grad and torch.isfinite(out.float()).all() -QUANT_WEIGHTS = [r for r in WEIGHTS if r not in ("bf16", "fp16")] +# Row-wise FP8 and integer formats leave gradients to the separate dgrad stream. +QUANT_WEIGHTS = [r for r in WEIGHTS if r not in ("bf16", "fp16", "fp8_rowwise") and r not in INT_BITS] @pytest.mark.parametrize("weights", QUANT_WEIGHTS) @@ -321,7 +324,6 @@ def test_dgrad_grouped_accumulates_over_top_k(): row is written exactly once, so a kernel that STORES and one that ACCUMULATES are indistinguishable — which is how a store survived the suite while silently dropping all but one expert's contribution per token on any real MoE.""" - from finegrained_kernels.backward import dgrad_matmul_grouped torch.manual_seed(0) T, tk, E, Ne, Ke = 8, 2, 4, 256, 256 @@ -365,8 +367,6 @@ def test_dgrad_grouped_every_admitted_config_accumulates(): benches configs by speed, so an arm that is wrong ONLY under these launch semantics gets crowned whenever it wins a near-tied tune — the descriptor-anchor bug this guards against read the wrong dY rows on exactly this launch while passing every identity-gather test.""" - from finegrained_kernels import backward as backward_module - from finegrained_kernels.backward import dgrad_matmul_grouped torch.manual_seed(0) T, tk, E, Ne, Ke = 8, 2, 4, 256, 256 @@ -394,7 +394,7 @@ def launch(): num_input_rows=T, ) - tuner = backward_module.dgrad_matmul_grouped_kernel + tuner = dgrad_matmul_grouped_kernel admitted = {} original_prune, original_configs = tuner.early_config_prune, tuner.configs diff --git a/finegrained-kernels/tests/test_moe.py b/finegrained-kernels/tests/test_moe.py index 15ca93ada..6f115dd5a 100644 --- a/finegrained-kernels/tests/test_moe.py +++ b/finegrained-kernels/tests/test_moe.py @@ -29,13 +29,20 @@ from utils import ( # type: ignore DTYPE_TAG, DTYPE_TO_TOL, + INT_BITS, + INT_GROUP, SUPPORTS_SWIZZLED_SCALES, TEST_DEVICE, WEIGHTS, + make_gptq_act_order, + maybe_compile, + quant_dequant_a, + dq_grouped, + dq_int, ) -from finegrained_kernels import moe, scheduling, swizzle_mx_scales # type: ignore -from finegrained_kernels.epilogue import fused_glu # type: ignore +from finegrained_kernels import fold_act_order_into_gate_up, moe, scheduling, swizzle_mx_scales # type: ignore +from finegrained_kernels.epilogue import apply_glu, fused_glu # type: ignore @dataclass(frozen=True) @@ -55,6 +62,8 @@ class MoEProblem: negative_sentinels: bool = False # off-rank routes as -1 rather than num_experts dtype: torch.dtype = torch.bfloat16 activation_format: Optional[str] = None + activation_scale_ub: Optional[float] = None + compile: bool = False swiglu_alpha: Optional[float] = None swiglu_limit: Optional[float] = None act_fn: Union[str, Callable] = "silu" # a callable runs on the host between the GEMMs @@ -64,6 +73,9 @@ class MoEProblem: static: bool = False # calibrated (static) activation scale, one per expert post_expert_norm: Optional[str] = None # a model's per-expert output norm, by name bias: bool = False # a per-expert output bias on both projections (GPT-OSS ships them) + zero_points: bool = False # asymmetric integer weights: a zero point per weight scale group + int_group: int | None = INT_GROUP # None spans each projection's K (per-channel) + pre_quant_scale: bool = False # shared gate/up AWQ input scale, per-expert down input scale @property def id(self): @@ -81,12 +93,16 @@ def id(self): return ( f"{self.weights}_T{self.num_tokens}_E{self.num_experts}_H{self.hidden_dim}" f"_I{self.intermediate_dim}_top{self.num_top_k}_{DTYPE_TAG[self.dtype]}" + f"{'_capped' if self.activation_scale_ub is not None else ''}{'_compile' if self.compile else ''}" f"{act}{fmt}{'_swizzled' if self.swizzled else ''}" f"{'_inputglobals' if self.input_globals else ''}" f"{'_expertglobals' if self.expert_globals else ''}" f"{'_static' if self.static else ''}" f"{'_' + self.post_expert_norm if self.post_expert_norm else ''}" f"{'_bias' if self.bias else ''}" + f"{'_zp' if self.zero_points else ''}" + f"{'_channel' if self.weights in INT_BITS and self.int_group is None else ''}" + f"{'_prequantscale' if self.pre_quant_scale else ''}" f"{'_sentinel' if self.sentinel_fraction > 0 else ''}" f"{'_negative' if self.negative_sentinels else ''}" ) @@ -119,6 +135,16 @@ def id(self): # out quantized instead (``quantize_routed_rows_per_expert``): 4096 routed rows, which the # 64-token cells above never reach, so only this one runs the kernel's pre-quantized arm MoEProblem(weights="fp8_tensor", num_tokens=1024, num_top_k=4, static=True), + # dynamic per-token activations (tensor-wide, row-wise): the intermediate's whole-row scale is + # beyond any gate_up epilogue tile, so the chain hands the down a bf16 intermediate to quantize + MoEProblem(weights="fp8_tensor", num_top_k=4), + # row-wise weights (Llama 4 FP8 experts, top-1): decode, small batch, and a prefill batch + # deep enough per expert for the grouped kernel's big tiles + MoEProblem(weights="fp8_rowwise", num_tokens=1, num_top_k=1), + MoEProblem(weights="fp8_rowwise", num_top_k=1), + MoEProblem(weights="fp8_rowwise", num_top_k=1, activation_scale_ub=0.5, compile=True), + MoEProblem(weights="fp8_rowwise", num_tokens=1024, num_top_k=1), + MoEProblem(weights="fp8_rowwise", num_top_k=4, static=True), # block-FP8 with UE8M0 (power-of-two) scales — the whole-model UE8M0 contract: acts, # weights, and the fused intermediate requant all power-of-two (DeepSeek-V4 attn / B200). MoEProblem(weights="fp8_128x128_ue8m0", num_tokens=1), @@ -129,6 +155,9 @@ def id(self): # against the down's, the down consumes it — asymmetric threading breaks parity hard ── MoEProblem(weights="nvfp4", input_globals=True), MoEProblem(weights="nvfp4", num_tokens=1, input_globals=True), + MoEProblem(weights="nvfp4", input_globals=True, expert_globals=True, pre_quant_scale=True), + MoEProblem(weights="nvfp4", activation_format="fp8", static=True, num_top_k=2, pre_quant_scale=True), + MoEProblem(weights="mxfp4", activation_format="fp8", static=True, num_top_k=2, pre_quant_scale=True), # ── pre-swizzled weight scales (swizzle once at load — the deployment contract): the fused # grouped chain then runs scatter-free gate_up -> swizzled Cs -> down's 5D-As fast path, and # batched decode reads the descriptor scale load. Values unchanged, so parity holds as-is. ── @@ -169,10 +198,12 @@ def id(self): MoEProblem(weights="mxfp4", hidden_dim=320, intermediate_dim=320), # ── explicit format forwarding: W4A8 chain on mxfp4 weights (default is W4A4) ── MoEProblem(weights="mxfp4", activation_format="mxfp8"), - # ── weight-only weight-only: bf16 acts × mxfp4 weights, dedicated dequant-then-bf16-dot - # kernels (the gpt-oss / matmul_ogs format); intermediate stays bf16 (no requant) ── + # ── weight-only: bf16 activations and intermediate (no requantization) ── MoEProblem(weights="mxfp4", activation_format="bf16", num_tokens=1), MoEProblem(weights="mxfp4", activation_format="bf16"), + MoEProblem(weights="fp8_128x128", activation_format="bf16", num_tokens=1, bias=True), + MoEProblem(weights="fp8_128x128", activation_format="bf16", num_tokens=5, bias=True), + MoEProblem(weights="fp8_tensor", activation_format="bf16", num_tokens=5, num_top_k=2), # ── clamped/scaled SwiGLU (GPT-OSS / MiniMax-M3); glu is format-independent ── MoEProblem(weights="mxfp8", swiglu_alpha=1.702, swiglu_limit=7.0), # alpha / limit are independent glu branches — cover each alone @@ -188,12 +219,31 @@ def id(self): MoEProblem(weights="fp8_128x128", bias=True), MoEProblem(weights="fp8_128x128", num_tokens=1, bias=True), MoEProblem(weights="fp8_tensor", num_top_k=4, static=True, bias=True), + MoEProblem(weights="fp8_rowwise", num_top_k=4, bias=True), MoEProblem(weights="mxfp8", bias=True), MoEProblem(weights="nvfp4", input_globals=True, bias=True), MoEProblem(weights="mxfp4", activation_format="bf16", swiglu_alpha=1.702, swiglu_limit=7.0, num_tokens=1, bias=True), + # ── int4 weight-only (W4A16), symmetric and with zero points: decode, batch, the GPT-OSS-style gate_up + # (clamped SwiGLU + bias), and EP sentinels ── + MoEProblem(weights="int4", num_tokens=1), + MoEProblem(weights="int4"), + MoEProblem(weights="int4", zero_points=True, num_tokens=1), + MoEProblem(weights="int4", zero_points=True, num_tokens=64, num_top_k=2), + MoEProblem(weights="int4", zero_points=True, swiglu_alpha=1.702, swiglu_limit=7.0, bias=True), + MoEProblem(weights="int4", num_tokens=8, sentinel_fraction=0.875), + MoEProblem(weights="int4", num_tokens=8, sentinel_fraction=0.875, negative_sentinels=True), + MoEProblem(weights="int8", zero_points=True, num_tokens=4), + MoEProblem(weights="int2", zero_points=True, num_tokens=64, num_top_k=2), + MoEProblem(weights="int4", int_group=None, zero_points=True), + MoEProblem(weights="int4", pre_quant_scale=True, num_tokens=4, num_top_k=2), + MoEProblem(weights="int4", activation_format="fp8", static=True, pre_quant_scale=True, + num_tokens=32, num_top_k=2), + MoEProblem(weights="int8", int_group=None, activation_format="int8", static=True, pre_quant_scale=True, + num_tokens=32, num_top_k=2), # ── expert parallelism: non-local experts sentinel-masked ── MoEProblem(weights="mxfp8", num_tokens=8, sentinel_fraction=0.875), MoEProblem(weights="fp8_128x128", num_tokens=8, sentinel_fraction=0.875), + MoEProblem(weights="fp8_rowwise", num_tokens=8, sentinel_fraction=0.875), MoEProblem(weights="mxfp8", num_tokens=8, sentinel_fraction=0.875, negative_sentinels=True), # swizzled MX scales quantize the activations once per token into every routed row's tile, which a token # routed only to other ranks' experts must not write @@ -222,11 +272,20 @@ def _make_moe_weights(problem: MoEProblem): read block pairs, the unfused plain GEMM remaps its block index in-kernel.""" make = WEIGHTS[problem.weights]["make"] gate_up, gate_up_s, gate_up_g = make( - 2 * problem.intermediate_dim, problem.hidden_dim, problem.num_experts + 2 * problem.intermediate_dim, problem.hidden_dim, problem.num_experts, + **({"group": problem.int_group or problem.hidden_dim} if problem.weights in INT_BITS else {}), ) down, down_s, down_g = make( - problem.hidden_dim, problem.intermediate_dim, problem.num_experts + problem.hidden_dim, problem.intermediate_dim, problem.num_experts, + **({"group": problem.int_group or problem.intermediate_dim} if problem.weights in INT_BITS else {}), ) + if problem.weights == "nvfp4" and problem.activation_format == "fp8": + gate_up_s, down_s = gate_up_s[..., ::2].contiguous(), down_s[..., ::2].contiguous() + if problem.weights in INT_BITS and problem.activation_format in ("fp8", "int8"): + gate_up_s, down_s = gate_up_s.float(), down_s.float() + if problem.activation_format == "fp8": + gate_up_g = dq_int(gate_up, gate_up_s, group=problem.int_group or problem.hidden_dim).abs().amax(dim=(-2, -1)).div(448).clamp(min=1e-12) + down_g = dq_int(down, down_s, group=problem.int_group or problem.intermediate_dim).abs().amax(dim=(-2, -1)).div(448).clamp(min=1e-12) if problem.swizzled: if not SUPPORTS_SWIZZLED_SCALES: pytest.skip("the SWIZZLE_32_4_4 scale layout is the tcgen05 fast path (CUDA-only)") @@ -288,7 +347,14 @@ def _static_scales(problem: MoEProblem, hidden, gate_up, gate_up_s, top_k_index) routed to it. ``None`` on a dynamic problem, which quantizes at runtime instead.""" if not problem.static: return None, None - weight = WEIGHTS[problem.weights]["dequant"](gate_up, gate_up_s).float() + weight = ( + dq_int(gate_up, gate_up_s, group=problem.int_group or problem.hidden_dim) + if problem.weights in INT_BITS else ( + dq_grouped(gate_up, gate_up_s, 32) + if problem.weights == "nvfp4" and problem.activation_format == "fp8" + else WEIGHTS[problem.weights]["dequant"](gate_up, gate_up_s).float() + ) + ) inter = fused_glu( torch.einsum("th,enh->etn", hidden.float(), weight), problem.act_fn, @@ -298,7 +364,8 @@ def _static_scales(problem: MoEProblem, hidden, gate_up, gate_up_s, top_k_index) routed = torch.stack([(top_k_index == e).any(dim=1) for e in range(problem.num_experts)]) def calibrate(seen): # (E, T, C) over every token -> (E,) over the ones each expert sees - return (seen.abs() * routed[..., None]).amax(dim=(1, 2)).div(448.0).clamp(min=1e-12).float() + max_quant = 127.0 if problem.activation_format == "int8" else 448.0 + return (seen.abs() * routed[..., None]).amax(dim=(1, 2)).div(max_quant).clamp(min=1e-12).float() scales = [calibrate(hidden.float().expand(problem.num_experts, -1, -1)), calibrate(inter)] return tuple(scale.contiguous() for scale in scales) @@ -314,9 +381,21 @@ def _biases(problem: MoEProblem): return gate_up_bias, down_bias +def _zero_points(problem: MoEProblem): + """The integer weights' zero points, one per scale group: ``(E, 2I, H // group)`` on gate_up and + ``(E, H, I // group)`` on down; ``None`` on a problem without them.""" + if not problem.zero_points: + return None, None + E, H, inter = problem.num_experts, problem.hidden_dim, problem.intermediate_dim + high = 1 << INT_BITS[problem.weights] + gate_up_zero = torch.randint(0, high, (E, 2 * inter, H // (problem.int_group or H)), device=TEST_DEVICE, dtype=torch.uint8) + down_zero = torch.randint(0, high, (E, H, inter // (problem.int_group or inter)), device=TEST_DEVICE, dtype=torch.uint8) + return gate_up_zero, down_zero + + def _common_kwargs(problem: MoEProblem, hidden, top_k_index, gate_up, gate_up_s, gate_up_g, down_g): - """The kwargs every forward takes: the per-expert biases, the weights' second-level globals, the calibrated - activation globals and scales, the GLU knobs, and a model's per-expert output norm. The fused chain + """The kwargs every forward takes: the per-expert biases, the weights' second-level globals and zero points, + the calibrated activation globals and scales, the GLU knobs, and a model's per-expert output norm. The fused chain folds a named norm into its reduce while the unfused reference normalizes the routed rows in a pass of their own, so handing both the same weight is what makes that a parity check.""" gate_up_in_g, down_in_g = _input_globals(problem, hidden) @@ -327,7 +406,13 @@ def _common_kwargs(problem: MoEProblem, hidden, top_k_index, gate_up, gate_up_s, else None ) gate_up_bias, down_bias = _biases(problem) - return dict( + gate_up_zero, down_zero = _zero_points(problem) + gate_up_pre_scale = down_pre_scale = None + if problem.pre_quant_scale: + gate_up_pre_scale = torch.linspace(0.5, 1.5, problem.hidden_dim, device=hidden.device, dtype=hidden.dtype) + down_pre_scale = torch.linspace(0.5, 1.5, problem.num_experts * problem.intermediate_dim, + device=hidden.device, dtype=hidden.dtype).view(problem.num_experts, problem.intermediate_dim) + kwargs = dict( gate_up_proj_bias=gate_up_bias, down_proj_bias=down_bias, gate_up_proj_weight_global_scale=gate_up_g, @@ -343,6 +428,13 @@ def _common_kwargs(problem: MoEProblem, hidden, top_k_index, gate_up, gate_up_s, post_expert_norm=problem.post_expert_norm, post_expert_norm_weight=norm_weight, ) + if problem.weights in INT_BITS: + kwargs.update(gate_up_proj_weight_zero_point=gate_up_zero, down_proj_weight_zero_point=down_zero) + if problem.pre_quant_scale: + kwargs.update(gate_up_proj_pre_quant_scale=gate_up_pre_scale, down_proj_pre_quant_scale=down_pre_scale) + if problem.activation_scale_ub is not None: + kwargs["activation_scale_ub"] = problem.activation_scale_ub + return kwargs def _assert_fused_correctness(out, ref, problem: MoEProblem): @@ -363,7 +455,7 @@ def _run_pair(problem: MoEProblem, fused_fn, unfused_fn): ref = unfused_fn( hidden, top_k_index, top_k_weights, gate_up, down, gate_up_s, down_s, **common ) - out = fused_fn( + out = maybe_compile(fused_fn, problem.compile)( hidden, top_k_index, top_k_weights, @@ -375,17 +467,30 @@ def _run_pair(problem: MoEProblem, fused_fn, unfused_fn): **common, ) _assert_fused_correctness(out, ref, problem) + if problem.activation_scale_ub is not None: + # Check both clipping sites independently: fused and unfused share their quantizer. + row = WEIGHTS[problem.weights] + ids = top_k_index.flatten().long() + aq = quant_dequant_a(hidden, problem.hidden_dim, activation_scale_ub=problem.activation_scale_ub) + aq = aq.repeat_interleave(problem.num_top_k, dim=0) + gu = torch.einsum("sk,snk->sn", aq, row["dequant"](gate_up, gate_up_s)[ids]).to(problem.dtype) + inter = apply_glu(gu[:, 0::2], gu[:, 1::2]) + iq = quant_dequant_a(inter, problem.intermediate_dim, activation_scale_ub=problem.activation_scale_ub) + rows = torch.einsum("sk,snk->sn", iq, row["dequant"](down, down_s)[ids]).to(problem.dtype) + torch_ref = (rows.view(problem.num_tokens, problem.num_top_k, -1) * top_k_weights[..., None]).sum(1) + _assert_fused_correctness(out, torch_ref, problem) @pytest.mark.kernels_ci @pytest.mark.skipif(TEST_DEVICE is None, reason="Accelerator not available") +@pytest.mark.parametrize("num_tokens", [300, 1025]) @pytest.mark.parametrize("sentinel_fraction", [0.5, 1.0]) @pytest.mark.parametrize("negative", [False, True], ids=["past_experts", "negative"]) -def test_grouped_scheduling_defines_every_route(sentinel_fraction, negative): +def test_grouped_scheduling_defines_every_route(sentinel_fraction, negative, num_tokens): """Both maps are a permutation of the ``S`` routes: the local rows expert-sorted up to ``expert_start[E]``, the EP-sentinel routes (another rank's experts, any id outside ``[0, E)``) after them, so a consumer over all ``S`` rows reads real slots and tokens.""" - E, T, K = 8, 300, 4 + E, T, K = 8, num_tokens, 4 g = torch.Generator(device=TEST_DEVICE).manual_seed(0) expert_ids = torch.randint(0, E, (T, K), device=TEST_DEVICE, generator=g) off_rank = torch.rand(T, K, device=TEST_DEVICE, generator=g) < sentinel_fraction @@ -507,8 +612,8 @@ def preblock(ws): # the baseline's own layout: torchao's rearrange, done once o assert rel < 0.05, f"torch baseline diverges from unfused reference: rel={rel:.3e}" -def _run_compiled_across_shapes(fused_fn): - """TWO different mxfp4 problems through ONE ``torch.compile(fullgraph=True)`` function +def _run_compiled_across_shapes(fused_fn, weights): + """TWO different problems through ONE ``torch.compile(fullgraph=True)`` function with no compiler reset in between: the recompile marks the weight shapes automatic-dynamic, and the family predicates must still return real bools — a lazy SymBool reaching ``is_x(gate) != is_x(down)`` builds a nested symbolic Eq that crashes @@ -519,10 +624,9 @@ def _run_compiled_across_shapes(fused_fn): torch.compiler.reset() compiled = torch.compile(fused_fn, fullgraph=True) for problem in ( - MoEProblem(weights="mxfp4", num_tokens=1), - MoEProblem( - weights="mxfp4", num_tokens=1, hidden_dim=320, intermediate_dim=320 - ), + MoEProblem(weights=weights, num_tokens=1), + MoEProblem(weights=weights, num_tokens=1, hidden_dim=640, intermediate_dim=384, int_group=None) + if weights in INT_BITS or weights.startswith("fp8") else MoEProblem(weights=weights, num_tokens=1, hidden_dim=320, intermediate_dim=320), ): torch.manual_seed(0) gate_up, gate_up_s, gate_up_g, down, down_s, down_g = _make_moe_weights(problem) @@ -540,17 +644,80 @@ def _run_compiled_across_shapes(fused_fn): @pytest.mark.kernels_ci @pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") -def test_fused_batched_compiles_across_shapes(): +@pytest.mark.parametrize("weights", ["mxfp4", "int4", "fp8_128x128", "fp8_rowwise"]) +def test_fused_batched_compiles_across_shapes(weights): """``moe_fused_batched`` through the shared two-shape compile check (see ``_run_compiled_across_shapes`` for the dynamo failure class it guards).""" - _run_compiled_across_shapes(moe.moe_fused_batched) + _run_compiled_across_shapes(moe.moe_fused_batched, weights) @pytest.mark.kernels_ci @pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") -def test_fused_grouped_compiles_across_shapes(): +@pytest.mark.parametrize("weights", ["mxfp4", "int4", "fp8_128x128", "fp8_rowwise"]) +def test_fused_grouped_compiles_across_shapes(weights): """``moe_fused_grouped`` through the same two-shape compile check — the grouped chain additionally puts ``compute_grouped_scheduling`` (an opaque custom op) inside the graph, which the batched sibling never exercises.""" - _run_compiled_across_shapes(moe.moe_fused_grouped) + _run_compiled_across_shapes(moe.moe_fused_grouped, weights) + + +@pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") +@pytest.mark.parametrize("fold", [False, True]) +@pytest.mark.parametrize("num_tokens", [1, 32]) +@pytest.mark.parametrize("forward", [moe.moe_fused_grouped, moe.moe_fused_batched], ids=["grouped", "batched"]) +def test_fused_int4_act_order(forward, num_tokens, fold): + """The fused MoE chain on act-order GPTQ experts: gate|up gathers hidden's columns by its per-expert K order, and + the down's order is either gathered at runtime too or folded into the gate|up output rows at load + (``fold_act_order_into_gate_up``), in which case the down runs in natural order. Reference: the torch chain on + the weights dequantized in natural K order.""" + + torch.manual_seed(0) + E, H, inter, top_k = 8, 512, 256, 2 + gu, gu_s, gu_z, gu_perm, gu_ref = make_gptq_act_order(2 * inter, H, E) + dn, dn_s, dn_z, dn_perm, dn_ref = make_gptq_act_order(H, inter, E) + hidden = torch.randn(num_tokens, H, device=TEST_DEVICE, dtype=torch.bfloat16) + top_k_index = torch.stack([torch.randperm(E, device=TEST_DEVICE)[:top_k] for _ in range(num_tokens)]).int() + top_k_weights = torch.rand(num_tokens, top_k, device=TEST_DEVICE, dtype=torch.bfloat16) + down_perm = dn_perm + if fold: + gu, gu_s, gu_z, _ = fold_act_order_into_gate_up(gu, gu_s, gu_z, None, dn_perm) + down_perm = None + out = forward( + hidden, top_k_index, top_k_weights, gu, dn, gu_s, dn_s, + gate_up_proj_weight_zero_point=gu_z, down_proj_weight_zero_point=dn_z, + gate_up_proj_weight_perm=gu_perm, down_proj_weight_perm=down_perm, + ) + ref = torch.zeros(num_tokens, H, device=TEST_DEVICE) + for t in range(num_tokens): + for k in range(top_k): + e = top_k_index[t, k].long() + z = hidden[t].float() @ gu_ref[e].T + h = apply_glu(z[0::2], z[1::2], "silu", None, None).to(torch.bfloat16).float() + ref[t] += top_k_weights[t, k].float() * (h @ dn_ref[e].T) + err = ((out.float() - ref).norm() / ref.norm()).item() + assert err < 1e-2, f"{forward} T={num_tokens} fold={fold}: act-order MoE vs reference {err:.2e}" + + +@pytest.mark.kernels_ci +@pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator required") +@pytest.mark.parametrize("forward", [moe.moe_fused_batched, moe.moe_fused_grouped]) +@pytest.mark.parametrize("mismatch", ["width", "group", "zero", "full_precision"]) +def test_fused_int_recipe_mismatch(forward, mismatch): + """Integer fused blocks reject unsupported mixed recipes before launching either GEMM.""" + problem = MoEProblem(weights="int4") + gu, gs, _, dn, ds, _ = _make_moe_weights(problem) + hidden, ids, weights = _make_moe_inputs(problem) + down_zero = None + if mismatch == "width": + dn, ds, _ = WEIGHTS["int2"]["make"](problem.hidden_dim, problem.intermediate_dim, problem.num_experts) + elif mismatch == "group": + ds = ds.repeat_interleave(2, dim=-1) + elif mismatch == "zero": + down_zero = torch.zeros_like(ds, dtype=torch.uint8) + else: + dn = torch.empty(problem.num_experts, problem.hidden_dim, problem.intermediate_dim, + device=hidden.device, dtype=hidden.dtype) + ds = None + with pytest.raises(ValueError, match="same integer quantization recipe"): + forward(hidden, ids, weights, gu, dn, gs, ds, down_proj_weight_zero_point=down_zero) diff --git a/finegrained-kernels/tests/test_ops.py b/finegrained-kernels/tests/test_ops.py index b75d3582e..83ec88a99 100644 --- a/finegrained-kernels/tests/test_ops.py +++ b/finegrained-kernels/tests/test_ops.py @@ -37,21 +37,42 @@ from utils import ( # type: ignore DTYPE_TO_TOL, + INT_BITS, + INT_GROUP, REQUANT_FN, REQUANT_GROUP, SUPPORTS_SWIZZLED_SCALES, TEST_DEVICE, WEIGHTS, dq_grouped, + dq_int, dq_scale, make_static_activation_scale, + make_gptq_act_order, maybe_compile, quant_dequant_a, ) import finegrained_kernels # type: ignore -from finegrained_kernels import swizzle_mx_scales, unswizzle_mx_scales # type: ignore +from finegrained_kernels import swizzle_mx_scales, unswizzle_mx_scales, repack_modelopt_int4, repack_modelopt_int8, repack_ggml_int, dequantize_int_weight, repack_awq_int4, repack_gptq_int4 # type: ignore +from finegrained_kernels.batched import ( # type: ignore + int_weight_only_matmul_batched_kernel, + mx_dynamic_matmul_batched_kernel, + w8a8_tensor_dynamic_fp8_matmul_batched_kernel, +) from finegrained_kernels.compat import NVFP4_SCALE_GROUP_K # type: ignore +from finegrained_kernels.grouped import ( # type: ignore + block_fp8_weight_only_matmul_grouped_kernel, + full_precision_matmul_grouped_kernel, + int_weight_only_matmul_grouped_kernel, + mx_dynamic_matmul_grouped_kernel, + w8a8_tensor_dynamic_fp8_matmul_grouped_kernel, +) +from finegrained_kernels.matmul import ( # type: ignore + int_weight_only_matmul_2d_kernel, + mx_dynamic_matmul_kernel, + w8a8_tensor_dynamic_fp8_matmul_kernel, +) from finegrained_kernels.quant import nvfp4_act_quant # type: ignore from finegrained_kernels.epilogue import apply_glu # type: ignore @@ -75,6 +96,7 @@ class Problem: swiglu_alpha: float | None = None swiglu_limit: float | None = None activation_format: str | None = None + activation_scale_ub: float | None = None quantize_output: bool = False per_expert_globals: bool = False # activation (and requant output) globals calibrated per expert prequant: bool = False # pass As explicitly (must be bit-identical to raw A) @@ -85,11 +107,19 @@ class Problem: empty_expert: bool = False compile: bool = False bias: bool = False # a per-expert output bias, added on the accumulator before the GLU + zero_points: bool = False # asymmetric integer weights: a zero point per weight scale group + int_group: int | None = INT_GROUP # None spans K (per-channel) + pre_quant_scale: bool = False # AWQ input transform, shared for dense or per expert when routed + svdquant_rank: int = 0 # dense NVFP4 residual plus BF16 low-rank correction dtype: torch.dtype = torch.bfloat16 @property def id(self): tag = self.weights + if self.activation_scale_ub is not None: + tag += "_capped" + if self.weights in INT_BITS and self.int_group != INT_GROUP: + tag += f"_group{self.int_group or 'channel'}" if self.gate: tag += f"_gate_{self.act_fn}" if self.swiglu_alpha is not None: @@ -118,6 +148,12 @@ def id(self): tag += "_compile" if self.bias: tag += "_bias" + if self.zero_points: + tag += "_zp" + if self.pre_quant_scale: + tag += "_prequantscale" + if self.svdquant_rank: + tag += f"_svdquant_r{self.svdquant_rank}" if self.dtype != torch.bfloat16: tag += f"_{str(self.dtype).rsplit('.', 1)[-1]}" # float16 / float32 return f"{tag}_S{self.S}_E{self.E}_N{self.N}_K{self.K}" @@ -148,11 +184,17 @@ def scenarios() -> list[Problem]: out.append( Problem(weights=w, gate=True, activation_format=fmt, quantize_output=True, swizzled=True) ) - default_in = {"fp8_128x128": "fp8", "fp8_tensor": "fp8"}.get(w) + default_in = {"fp8_128x128": "fp8", "fp8_tensor": "fp8", "fp8_rowwise": "fp8"}.get(w) for fmt in row["activation_formats"]: if fmt != default_in: out.append(Problem(weights=w, activation_format=fmt)) out += [ + # Coarse and per-tensor FP8 with raw BF16/FP16 activations. + Problem(weights="fp8_128x128", activation_format="bf16", S=1, N=192, K=384, bias=True, dtype=torch.float16), + Problem(weights="fp8_128x128_ue8m0", activation_format="bf16", S=8, empty_expert=True, sentinel_fraction=0.25), + Problem(weights="fp8_tensor", activation_format="bf16", S=1, gate=True, bias=True, compile=True), + Problem(weights="fp8_128x128", activation_format="bf16", S=257, N=256, K=512, gate=True, bias=True, compile=True), + Problem(weights="fp8_128x128", activation_format="bf16", S=1025, E=1, N=192, K=384, compile=True, bias=True), Problem(weights="fp16", dtype=torch.float16), # gate WITHOUT requant (raw GLU intermediate) — one quantized + one full format Problem(weights="mxfp8", gate=True), @@ -168,6 +210,13 @@ def scenarios() -> list[Problem]: Problem(weights="fp8_128x128", prequant=True), Problem(weights="mxfp8", prequant=True), Problem(weights="nvfp4", prequant=True), + Problem(weights="mxfp4", activation_format="fp8", static=True, gate=True, pre_quant_scale=True, compile=True), + Problem(weights="nvfp4", activation_format="fp8", static=True, gate=True, pre_quant_scale=True, compile=True), + Problem(weights="nvfp4", S=8, gate=True, pre_quant_scale=True, compile=True), + Problem(weights="nvfp4", S=1, svdquant_rank=32, pre_quant_scale=True, bias=True), + Problem(weights="nvfp4", S=64, N=129, K=320, svdquant_rank=128, pre_quant_scale=True, compile=True), + Problem(weights="nvfp4", S=128, svdquant_rank=32, pre_quant_scale=True, gate=True, + activation_format="nvfp4", quantize_output=True, swizzled=True), Problem(weights="mxfp8", sentinel_fraction=0.25), # the prefill expansion (>= 4096 routed rows below 1024 per expert) copies only the local rows: # values + 2D scales, and raw bf16 rows @@ -234,6 +283,7 @@ def scenarios() -> list[Problem]: Problem(weights="fp8_128x128", compile=True), Problem(weights="mxfp4", compile=True), Problem(weights="bf16", compile=True), # the fp kernel's pre_hook under compile + Problem(weights="fp8_rowwise", S=256, activation_scale_ub=0.5, compile=True), Problem(weights="fp8_tensor", static=True, compile=True), # the static arm's operand marshalling # calibrated (static) activation quant, reached when As is a calibrated scale rather than # per-block. One value per quantized module, so the OP fixes the shape: a dense linear @@ -243,6 +293,20 @@ def scenarios() -> list[Problem]: # per-TENSOR weights: the shipped static form (Ministral-3 dense, Mistral-4 MoE — both # write qscheme_act="TENSOR" with weight_block_size=None) Problem(weights="fp8_tensor", static=True), + # row-wise weights (Llama 4 FP8): the stacked gate|up rows carry their own scales, the swapped decode + # tile folds them on its other axis, a caller's per-token As rides as-is, and the launch-scale shape + # reaches the grouped kernel's descriptor / warp-specialized arms + Problem(weights="fp8_rowwise", gate=True), + Problem(weights="fp8_rowwise", S=8), + Problem(weights="fp8_rowwise", gate=True, S=8), + Problem(weights="fp8_rowwise", prequant=True, compile=True), + Problem(weights="fp8_rowwise", S=4096, E=8, prequant=True, noncontiguous=True), + Problem(weights="fp8_rowwise", static=True, S=17), + Problem(weights="fp8_rowwise", S=257, bias=True, empty_expert=True, compile=True), + Problem(weights="fp8_rowwise", S=2048, E=16, N=512, K=1024), + # the prefill expansion lays the routed rows out, and their per-token scales with them + Problem(weights="fp8_rowwise", S=4096, E=8, sentinel_fraction=0.5), + Problem(weights="fp8_tensor", S=4096, E=8), # non-aligned N (64-grid, off the 128-grid — gpt-oss H=I=2880 shape). matmul_2d masks the # N-tail; routed MX runs the affine arm (per-row scales, any BN|N); routed FP8 rejects it # (its scales are 128-blocked along N — raises pointing to matmul_2d). @@ -279,6 +343,33 @@ def scenarios() -> list[Problem]: Problem(weights="mxfp8", gate=True, activation_format="mxfp8", quantize_output=True, bias=True), Problem(weights="mxfp4", activation_format="bf16", gate=True, swiglu_alpha=1.702, swiglu_limit=7.0, bias=True), Problem(weights="mxfp4", activation_format="bf16", gate=True, swiglu_alpha=1.702, swiglu_limit=7.0, bias=True, S=8), + # int4 weight-only (the format matrix above gives each row its plain cell): the gated GLU with + # a bias, decode (the batched swap reduce at S=1), a K off the BK grid (the batched tail tile + # slides back) and fp16 activations, each with and without zero points + Problem(weights="int4", gate=True, swiglu_alpha=1.702, swiglu_limit=7.0, bias=True), + Problem(weights="int4", zero_points=True, gate=True, bias=True), + Problem(weights="int4", S=1), + Problem(weights="int4", zero_points=True, S=8), + Problem(weights="int4", zero_points=True, K=640), + Problem(weights="int4", dtype=torch.float16), + Problem(weights="int4", zero_points=True, dtype=torch.float16), + # the 8- and 2-bit widths ride the same kernels (the width follows from the weight's word count) + Problem(weights="int8", zero_points=True, S=1), + Problem(weights="int8", zero_points=True, gate=True, bias=True), + Problem(weights="int8", dtype=torch.float16), + Problem(weights="int2", zero_points=True, S=1), + Problem(weights="int2", zero_points=True, gate=True, bias=True), + # Wide groups and per-channel scales retain rounded weight dequantization. + Problem(weights="int2", int_group=512, K=1024), + Problem(weights="int4", int_group=512, K=1024, zero_points=True, gate=True, bias=True), + Problem(weights="int8", int_group=None, K=1024, zero_points=True), + Problem(weights="int4", int_group=None, K=1024, compile=True), + Problem(weights="int4", zero_points=True, compile=True), + Problem(weights="int4", S=8, gate=True, bias=True, pre_quant_scale=True, compile=True), + Problem(weights="int4", S=32, gate=True, bias=True, activation_format="fp8", static=True, + pre_quant_scale=True, compile=True), + Problem(weights="int8", int_group=None, gate=True, bias=True, activation_format="int8", static=True, + pre_quant_scale=True, compile=True), ] return out @@ -318,6 +409,17 @@ def _bias(problem: Problem, op): return bias[0] if op == "matmul" else bias +def _zero_points(problem: Problem, E): + """The scenario's integer zero points, else None: one per weight scale group, ``(E, rows, K // group)`` + over the stack the weights were made at. Seeded, so the reference and the op read the identical values.""" + if not problem.zero_points: + return None + g = torch.Generator(device=TEST_DEVICE).manual_seed(2) + rows = 2 * problem.N if problem.gate else problem.N + shape = (E, rows, problem.K // (problem.int_group or problem.K)) + return torch.randint(0, 1 << INT_BITS[problem.weights], shape, device=TEST_DEVICE, dtype=torch.uint8, generator=g) + + def _make_noncontig(x): base = torch.empty((x.numel(), 2), dtype=x.dtype, device=x.device) base[:, 0] = x @@ -334,6 +436,8 @@ def _static_scale(problem: Problem, op, A): if not problem.static: return None scale = make_static_activation_scale(A) + if problem.activation_format == "int8": + scale = scale * (448.0 / 127.0) if op == "matmul": return scale fan = torch.linspace(0.5, 2.0, problem.E, device=A.device, dtype=torch.float32) @@ -354,7 +458,7 @@ def _act_global(problem: Problem, A, expert_ids=None): vector modelopt writes (one ``input_scale`` per expert), fanned around the per-tensor value so an ignored entry shows up. Weight-only (bf16 activations) and non-nvfp4 formats have no second level (None).""" - if problem.weights != "nvfp4" or problem.activation_format == "bf16": + if problem.weights != "nvfp4" or problem.activation_format in ("bf16", "fp8"): return None g = _nvfp4_global(A) if not problem.per_expert_globals: @@ -376,6 +480,8 @@ def _dequant_a(problem: Problem, op, A, expert_ids=None): calibrated static scale of each row's expert), plus the pre-quantized ``(Aq, As)`` form for the prequant round-trip check (``None`` where ``A`` stays raw).""" row = WEIGHTS[problem.weights] + if problem.activation_scale_ub is not None: + return quant_dequant_a(A, problem.K, activation_scale_ub=problem.activation_scale_ub), None static_scale = _static_scale(problem, op, A) act_global = _act_global(problem, A, expert_ids) if static_scale is not None: # static (calibrated) activation quant @@ -421,6 +527,20 @@ def _act_dequant(problem: Problem, op, A, As=None, As_global=None, expert_ids=No return WEIGHTS[problem.weights]["dq_act"](A, As) +def _pre_quant_scale(problem: Problem, op): + E = 1 if op == "matmul" else problem.E + scale = torch.linspace(0.5, 1.5, E * problem.K, device=TEST_DEVICE, dtype=problem.dtype) + return scale if op == "matmul" else scale.view(E, problem.K) + + +def _svdquant_factors(problem: Problem): + g = torch.Generator(device=TEST_DEVICE).manual_seed(87) + rows = problem.N * (2 if problem.gate else 1) + a = torch.randn(problem.svdquant_rank, problem.K, generator=g, device=TEST_DEVICE, dtype=problem.dtype) / problem.K**0.5 + b = torch.randn(rows, problem.svdquant_rank, generator=g, device=TEST_DEVICE, dtype=problem.dtype) * 20 + return a, b + + def _fp32_intermediate(problem: Problem, op, A, expert_ids, B, Bs, Bs_global, As=None, As_global=None): """The fp32 GLU output on the format-quantized operands — the pre-requant oracle, shared by the global calibration (``_out_global``) and the reference (``_reference``), reading the SAME ``A``/ @@ -428,8 +548,42 @@ def _fp32_intermediate(problem: Problem, op, A, expert_ids, B, Bs, Bs_global, As gather ``W[expert]`` and zero sentinel rows; GLU in fp32 (the production epilogue applies it to the fp32 accumulator directly).""" row = WEIGHTS[problem.weights] - A_dq = _act_dequant(problem, op, A, As, As_global, expert_ids) - W = row["dequant"](B, Bs, Bs_global) # (E, rows, K) fp32 + A_raw = A + if problem.pre_quant_scale: + scale = _pre_quant_scale(problem, op) + if op != "matmul": + scale = scale[expert_ids.long().clamp(max=problem.E - 1)] + A_raw = A * scale + if problem.static and problem.activation_format in ("fp8", "int8"): + act_scale = _rowwise(problem, _static_scale(problem, op, A), expert_ids) + normalized = A_raw.float() / act_scale + A_quant = normalized.round().clamp(-128, 127) if problem.activation_format == "int8" else normalized.clamp(-448, 448).to(torch.float8_e4m3fn).float() + A_dq = A_quant * act_scale + elif problem.pre_quant_scale and problem.weights == "nvfp4" and As_global is not None: + aq, scales = nvfp4_act_quant(A_raw, global_scale=As_global, expert_index=expert_ids) + A_dq = dq_grouped(aq.view(torch.int8), scales, NVFP4_SCALE_GROUP_K) * _rowwise(problem, As_global, expert_ids) + else: + A_dq = _act_dequant(problem, op, A_raw, As, As_global, expert_ids) + # (E, rows, K) fp32 + if problem.weights in INT_BITS: # dequantized to the activation dtype, as the GEMM does + W = dq_int(B, Bs, _zero_points(problem, B.shape[0]), + None if problem.activation_format in ("fp8", "int8") else problem.dtype, problem.int_group or problem.K) + if problem.activation_format == "fp8": + global_scale = Bs_global.view(-1, 1, 1) + W = (W / global_scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * global_scale + elif problem.gate and problem.weights in ("fp8_128x128", "fp8_128x128_ue8m0"): + # Gate/up weight rows alternate per value; their calibrated scale rows alternate per block. + gate_w = row["dequant"](B[:, 0::2], Bs[:, 0::2]) + up_w = row["dequant"](B[:, 1::2], Bs[:, 1::2]) + W = torch.stack((gate_w, up_w), dim=-2).flatten(-3, -2) + else: + W = ( + dq_grouped(B, Bs, 32) * Bs_global.view(-1, 1, 1) + if problem.weights == "nvfp4" and problem.activation_format == "fp8" + else row["dequant"](B, Bs, Bs_global) + ) + if problem.activation_format == "bf16": + W = W.to(A.dtype).float() bias = _bias(problem, op) if op == "matmul": ref = A_dq @ W[0].T # single linear, no routing @@ -441,6 +595,11 @@ def _fp32_intermediate(problem: Problem, op, A, expert_ids, B, Bs, Bs_global, As if bias is not None: ref = ref + bias.float()[local] ref[expert_ids.long() >= problem.E] = 0 + if problem.svdquant_rank: + lora_a, lora_b = _svdquant_factors(problem) + down = (A_raw.float() @ lora_a.float().T).to(problem.dtype) + correction = (down.float() @ lora_b.float().T).to(problem.dtype) + ref = ref.to(problem.dtype) + correction if problem.gate: gate_v, up_v = ref[..., 0::2], ref[..., 1::2] ref = apply_glu( @@ -494,7 +653,12 @@ def _op(problem: Problem, op, A, expert_ids, B, Bs, Bs_global, As=None, As_globa it before requant — the same scalar the reference uses, so the two spaces match.""" # a requant cell names its format explicitly (the reference snaps to that grid) assert not problem.quantize_output or problem.activation_format not in (None, "bf16") - kw = dict(activation_format=problem.activation_format, quantize_output=problem.quantize_output) + kw = dict(activation_format=problem.activation_format, activation_scale_ub=problem.activation_scale_ub, quantize_output=problem.quantize_output) + if problem.svdquant_rank: + lora_a, lora_b = _svdquant_factors(problem) + kw.update(svdquant_lora_a=lora_a, svdquant_lora_b=lora_b, simulate_unfused=True) + if problem.pre_quant_scale: + kw["a_pre_quant_scale"] = _pre_quant_scale(problem, op) if problem.gate: kw.update( gate=True, @@ -503,9 +667,12 @@ def _op(problem: Problem, op, A, expert_ids, B, Bs, Bs_global, As=None, As_globa swiglu_limit=problem.swiglu_limit, ) if not problem.quantize_output: - kw["output_dtype"] = problem.dtype + kw["output_dtype"] = None if problem.prequant else problem.dtype if problem.bias: kw["bias"] = _bias(problem, op) + if problem.zero_points: + zero_points = _zero_points(problem, B.shape[0]) + kw["b_zero"] = zero_points[0] if op == "matmul" else zero_points if out_global is not None: # provided NVFP4 output global (next proj's input_scale) kw["output_global_scale"] = out_global if problem.static: # fused static activation quant — As is the calibrated scale @@ -600,6 +767,8 @@ def _check(problem: Problem, dq_out, ref_cmp, expert_ids, op): else expert_ids.long() < problem.E ) if not problem.quantize_output: + # Use output-dtype tolerances: Torch and tuned kernels may sum FP32 + # products in different orders, and GLU can amplify near-zero differences. atol, rtol = DTYPE_TO_TOL[problem.dtype] torch.testing.assert_close( dq_out[keep], ref_cmp[keep].to(problem.dtype).float(), atol=atol, rtol=rtol @@ -626,24 +795,38 @@ def _check(problem: Problem, dq_out, ref_cmp, expert_ids, op): def _make_weights(problem: Problem, row, E): """The scenario's weights: ``(B, Bs, Bs_global)`` at the stack's row count (doubled under gate).""" - return row["make"](2 * problem.N if problem.gate else problem.N, problem.K, E) - - -def _skip_moe_only(problem: Problem, op: str) -> None: - """matmul_2d is the single-GEMM sibling: skip only the scenarios it can't represent — expert - routing (sentinel / noncontiguous / empty-expert / the MoE prequant-As check) and non-MX - input/output format knobs (its FP8 paths infer the quant from the scale shape and return the - intermediate dense). Everything else — including full-precision (BF16/FP16) weights and a - shared static activation scale — runs on all three ops.""" + kwargs = {"group": problem.int_group or problem.K} if problem.weights in INT_BITS else {} + B, Bs, Bs_global = row["make"](2 * problem.N if problem.gate else problem.N, problem.K, E, **kwargs) + if problem.weights == "nvfp4" and problem.activation_format == "fp8": + Bs = Bs[..., ::2].contiguous() # valid group-32 E4M3 payload; both FP4 halves use one scale + if problem.weights in INT_BITS and problem.activation_format == "fp8": + Bs = Bs.float() + Bs_global = dq_int(B, Bs, group=problem.int_group or problem.K).abs().amax(dim=(-2, -1)).div(448).clamp(min=1e-12) + elif problem.weights in INT_BITS and problem.activation_format == "int8": + Bs = Bs.float() + return B, Bs, Bs_global + + +def _skip_unsupported(problem: Problem, op: str) -> None: + """The scenarios an op can't represent. matmul_2d is the single-GEMM + sibling: it skips expert routing (sentinel / noncontiguous / empty-expert / the MoE prequant-As + check) and non-MX input/output format knobs (its FP8 paths infer the quant from the scale shape and + return the intermediate dense). Everything else — including full-precision (BF16/FP16) weights and + a shared static activation scale — runs on all three ops.""" + if problem.svdquant_rank and op != "matmul": + pytest.skip("SVDQuant checkpoint support is dense") if op != "matmul": return - if problem.sentinel_fraction or problem.noncontiguous or problem.empty_expert or problem.prequant: + if problem.sentinel_fraction or problem.noncontiguous or problem.empty_expert: pytest.skip("expert-routing scenario (MoE only)") + if problem.prequant and problem.weights != "fp8_rowwise": + pytest.skip("prequantized block-scale inputs need routing") if problem.per_expert_globals: pytest.skip("per-expert globals need routing (MoE only)") mx = problem.weights in ("mxfp8", "mxfp8_u8", "mxfp4", "nvfp4") - if not mx and (problem.activation_format is not None or problem.quantize_output): + int_quantized = problem.weights in INT_BITS and problem.activation_format in ("fp8", "int8") + if not mx and not int_quantized and (problem.activation_format not in (None, "bf16") or problem.quantize_output): pytest.skip("an explicit activation format / requant is MX-only for matmul_2d") @@ -654,12 +837,13 @@ def _skip_moe_only(problem: Problem, op: str) -> None: def test_op_scenarios(problem: Problem, op): """Reference (the op written in torch) vs op (the kernel): same inputs, each returning the op's own output format, compared once through the shared ``_dequant``.""" - _skip_moe_only(problem, op) + _skip_unsupported(problem, op) A, expert_ids = _routed(problem) row = WEIGHTS[problem.weights] E = 1 if op == "matmul" else problem.E # matmul is a single weight matrix B, Bs, Bs_global = _make_weights(problem, row, E) - if op != "matmul" and problem.N % 128 != 0 and problem.weights.startswith("fp8"): + if (op != "matmul" and problem.N % 128 != 0 and problem.weights.startswith("fp8") + and problem.activation_format != "bf16"): # fp8 weight scales are 128-blocked along N, so routed fp8 rejects non-128 N. MX (per-row # scales, BN | N) handles it on the affine arm — falls through to the normal ref-vs-op run. with pytest.raises(ValueError, match="matmul_2d"): @@ -668,8 +852,8 @@ def test_op_scenarios(problem: Problem, op): # ref is the op written in torch; both take the same inputs — including the PROVIDED nvfp4 output # global — and return the op's own format, so _dequant reads them identically and _check compares # once. _assert_op_layout is the one op-only check (the reference never swizzles its output). - # prequant scenarios hand the op its activations already quantized (As set — routed ops only, - # _skip_moe_only excludes matmul); otherwise raw A + As=None and the op quantizes it. Reference + # prequant scenarios hand the op its activations already quantized (As set); otherwise + # raw A + As=None and the op quantizes it. Reference # and op take the identical (A, As): pre-quantizing just hands the op the values it would # otherwise compute, and _fp32_intermediate reads whichever form it's given. The nvfp4 activation # global comes off the RAW A (before prequant replaces it) and rides separately as As_global. @@ -686,6 +870,8 @@ def _run_ref_vs_op(problem: Problem, op, A, expert_ids, B, Bs, Bs_global, shared As_global = _act_global(problem, A, expert_ids) if problem.prequant: A, As = _prequant_args(problem, op, A, expert_ids) + if problem.noncontiguous and As is not None: + As = _make_noncontig(As.reshape(-1)).view(As.shape) g_out = _out_global(problem, op, A, expert_ids, B, Bs, Bs_global, As, As_global) ref = _reference(problem, op, A, expert_ids, B, Bs, Bs_global, As=As, As_global=As_global, out_global=g_out) if shared is not None: @@ -712,42 +898,55 @@ def _run_ref_vs_op(problem: Problem, op, A, expert_ids, B, Bs, Bs_global, shared # kernel (the dense gate fusion and its gate_pointer_only/WS pruners), and one nvfp4 cell # (E4M3 scales: the nvfp4_native_ok fence + software decode arms). _SWEEP_CELLS = [ - (Problem(weights="mxfp8", gate=True, activation_format="mxfp8", quantize_output=True), "grouped", "mx_dynamic_matmul_grouped_kernel"), + (Problem(weights="mxfp8", gate=True, activation_format="mxfp8", quantize_output=True), "grouped", mx_dynamic_matmul_grouped_kernel), # K off the BK=128 grid admits the BK=64 rows, the family a K=256 cell never reaches - (Problem(weights="mxfp4", K=320), "grouped", "mx_dynamic_matmul_grouped_kernel"), - (Problem(weights="mxfp8", K=320), "grouped", "mx_dynamic_matmul_grouped_kernel"), - (Problem(weights="mxfp4", S=8), "batched", "mx_dynamic_matmul_batched_kernel"), - (Problem(weights="mxfp8", gate=True, activation_format="mxfp8", quantize_output=True, swizzled=True), "grouped", "mx_dynamic_matmul_grouped_kernel"), - (Problem(weights="mxfp8", gate=True, activation_format="mxfp8", quantize_output=True, swizzled=True), "matmul", "mx_dynamic_matmul_kernel"), - (Problem(weights="nvfp4", gate=True, activation_format="nvfp4", quantize_output=True, swizzled=True, S=8), "batched", "mx_dynamic_matmul_batched_kernel"), + (Problem(weights="mxfp4", K=320), "grouped", mx_dynamic_matmul_grouped_kernel), + (Problem(weights="mxfp8", K=320), "grouped", mx_dynamic_matmul_grouped_kernel), + (Problem(weights="mxfp4", S=8), "batched", mx_dynamic_matmul_batched_kernel), + (Problem(weights="mxfp8", gate=True, activation_format="mxfp8", quantize_output=True, swizzled=True), "grouped", mx_dynamic_matmul_grouped_kernel), + (Problem(weights="mxfp8", gate=True, activation_format="mxfp8", quantize_output=True, swizzled=True), "matmul", mx_dynamic_matmul_kernel), + (Problem(weights="nvfp4", gate=True, activation_format="nvfp4", quantize_output=True, swizzled=True, S=8), "batched", mx_dynamic_matmul_batched_kernel), + (Problem(weights="fp8_128x128", activation_format="bf16"), "grouped", block_fp8_weight_only_matmul_grouped_kernel), # the calibrated (static) arm on per-tensor weights: it hands the kernel a RAW A to quantize # per tile, the one activation form that cannot ride the TMA gather, so its admitted set is # the one a memory-mode fence gets wrong (silently — the tuner forgives what will not lower) - (Problem(weights="fp8_tensor", static=True), "grouped", "w8a8_tensor_dynamic_fp8_matmul_grouped_kernel"), + (Problem(weights="fp8_tensor", static=True), "grouped", w8a8_tensor_dynamic_fp8_matmul_grouped_kernel), + # row-wise weights: the per-row scale vector each arm folds on its own weight-row axis — the + # stacked gate|up span, the swapped decode tile, and the 2D kernel's tile + (Problem(weights="fp8_rowwise", gate=True, S=257, bias=True, empty_expert=True), "grouped", w8a8_tensor_dynamic_fp8_matmul_grouped_kernel), + (Problem(weights="fp8_rowwise", gate=True, S=8), "batched", w8a8_tensor_dynamic_fp8_matmul_batched_kernel), + (Problem(weights="fp8_rowwise", gate=True, S=256), "matmul", w8a8_tensor_dynamic_fp8_matmul_kernel), + + # the int4 weight-only kernels: the 2D kernel's tile x WS x memory-mode space and the batched + # dot / swapped-reduce arms, asymmetric so the zero-point reads are in every config + (Problem(weights="int4", zero_points=True, gate=True), "matmul", int_weight_only_matmul_2d_kernel), + (Problem(weights="int4", zero_points=True, S=8), "batched", int_weight_only_matmul_batched_kernel), + # the 2D kernel's single-token scalar reduce: its narrow tiles and the gate|up rows + (Problem(weights="int4", zero_points=True, gate=True, S=1), "matmul", int_weight_only_matmul_2d_kernel), + (Problem(weights="int4", zero_points=True, gate=True), "grouped", int_weight_only_matmul_grouped_kernel), + (Problem(weights="int8", zero_points=True, S=4), "matmul", int_weight_only_matmul_2d_kernel), + (Problem(weights="int2", zero_points=True, S=4), "matmul", int_weight_only_matmul_2d_kernel), + # Sweeps include BK < group, so a tuned winner cannot hide wide-group addressing. + (Problem(weights="int4", int_group=512, K=1024, S=4, zero_points=True, gate=True), "matmul", int_weight_only_matmul_2d_kernel), + (Problem(weights="int4", int_group=512, K=1024, zero_points=True, gate=True), "batched", int_weight_only_matmul_batched_kernel), + (Problem(weights="int4", int_group=None, K=1024, zero_points=True, gate=True), "grouped", int_weight_only_matmul_grouped_kernel), ] @pytest.mark.slow @pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") @pytest.mark.parametrize( - "problem, op, kernel_name", + "problem, op, tuner", _SWEEP_CELLS, ids=[f"{p.id}_{op}" for p, op, _ in _SWEEP_CELLS], ) -def test_every_admitted_config_is_correct(problem: Problem, op, kernel_name): +def test_every_admitted_config_is_correct(problem: Problem, op, tuner): """Force-run EVERY config the pruners admit for one launch and hard-check each against the torch oracle. The tuner benches configs by SPEED, so a wrong-but-fast config that slips the pruners gets crowned silently — the class no winner-only test can see. A forced config may fail to compile/launch (the tuner's forgiven-inf path — reported, not asserted); a config that RUNS must be correct.""" - import finegrained_kernels.batched as batched_mod - import finegrained_kernels.grouped as grouped_mod - import finegrained_kernels.matmul as matmul_mod - import finegrained_kernels.quant as quant_mod - - # the op's own kernel, or one of the shared passes it launches (the activation quant) - op_mod = {"batched": batched_mod, "grouped": grouped_mod, "matmul": matmul_mod}[op] - tuner = getattr(op_mod, kernel_name, None) or getattr(quant_mod, kernel_name) + A, expert_ids = _routed(problem) row = WEIGHTS[problem.weights] B, Bs, Bs_global = _make_weights(problem, row, problem.E) @@ -765,7 +964,7 @@ def spy(configs, named_args, **kwargs): try: tuner.cache.clear() _run_ref_vs_op(problem, op, A, expert_ids, B, Bs, Bs_global, shared=shared) - assert admitted.get("configs"), "the spy never saw a prune pass — wrong kernel_name?" + assert admitted.get("configs"), "the spy never saw a prune pass" forgiven = [] for cfg in admitted["configs"]: tuner.configs = [cfg] @@ -796,3 +995,278 @@ def spy(configs, named_args, **kwargs): print(f"\n[sweep] {len(forgiven)} admitted config(s) failed to compile/run (forgiven):") for cfg, err in forgiven: print(f" {cfg}: {err}") + + +@pytest.mark.skipif(not SUPPORTS_SWIZZLED_SCALES, reason="native scaled MMA required") +@pytest.mark.parametrize("strategy", ["split", "persistent"]) +def test_svdquant_scheduling(monkeypatch, strategy): + """Force both schedules: a short correctness tune can miss split-K or multiple tile iterations.""" + tuner = mx_dynamic_matmul_kernel + persistent = strategy == "persistent" + config = next(c for c in tuner.configs if c.kwargs.get("SPLIT_K") == (1 if persistent else 4) + and c.kwargs["PERSISTENT"] == persistent + and c.kwargs["BLOCK_SIZE_N"] == c.kwargs["BLOCK_SIZE_K"] == 128 + and c.kwargs["WARP_SPEC"] and c.num_warps == 4 and c.num_stages == 3) + baseline = next(c for c in tuner.configs if c.kwargs == dict(config.kwargs, PERSISTENT=False) + and c.num_warps == config.num_warps and c.num_stages == config.num_stages) + monkeypatch.setattr(tuner, "configs", [config]) + monkeypatch.setattr(tuner, "early_config_prune", None) + tuner.cache.clear() + if persistent: + monkeypatch.setattr("finegrained_kernels.matmul.persistent_program_count", lambda _: 2) + problem = Problem(weights="nvfp4", S=384 if persistent else 1, N=256, + K=256 if persistent else 1024, svdquant_rank=32, + pre_quant_scale=True, bias=True, swizzled=True, compile=True) + A, ids = _routed(problem) + B, Bs, global_scale = _make_weights(problem, WEIGHTS[problem.weights], 1) + try: + if persistent: + # Scheduling must preserve every bit, including BF16 midpoint choices from FP32 sum order. + act_global = _act_global(problem, A, ids) + actual = _op(problem, "matmul", A, ids, B, Bs, global_scale, As_global=act_global) + tuner.configs = [baseline] + tuner.cache.clear() + expected = _op(problem, "matmul", A, ids, B, Bs, global_scale, As_global=act_global) + assert torch.equal(actual, expected) + else: + _run_ref_vs_op(problem, "matmul", A, ids, B, Bs, global_scale) + finally: + tuner.cache.clear() + + +@pytest.mark.skipif(TEST_DEVICE != "cuda", reason="two concurrent CUDA streams required") +def test_int4_split_k_concurrent_streams(): + """Two split-K int4 matmuls launched concurrently on two streams, many times over, must equal their serial results + bit for bit: each launch reduces through its own arrival counters and partials (shared counters would let one + launch's arrivals finish another's tile early and sum unwritten partials). Both streams first sleep, so the + launches queue up behind the sleep and then run overlapped, with small grids that co-reside on the GPU.""" + tuner = int_weight_only_matmul_2d_kernel + split = [c for c in tuner.configs if c.kwargs.get("SPLIT_K") == 8 and c.kwargs["BLOCK_SIZE_N"] == 16 + and c.kwargs["BLOCK_SIZE_K"] == 128 and c.kwargs["A_MEMORY_MODE"] == "pointer" + and c.num_warps == 2 and c.num_stages == 3] + assert split, "no SPLIT_K=8 decode config in the grid" + g = torch.Generator(device=TEST_DEVICE).manual_seed(3) + N, K, M = 512, 4096, 4 + problems = [] + for _ in range(2): + q = torch.randint(0, 16, (N, K), device=TEST_DEVICE, dtype=torch.int32, generator=g) + B = (q.reshape(N, K // 8, 8) << (4 * torch.arange(8, device=TEST_DEVICE, dtype=torch.int32))).sum(-1, dtype=torch.int32).view(torch.uint32) + Bs = (torch.rand(N, K // 128, device=TEST_DEVICE, generator=g) * 0.02 + 0.005).to(torch.bfloat16) + A = torch.randn(M, K, device=TEST_DEVICE, dtype=torch.bfloat16, generator=g) + problems.append((A, B, Bs)) + orig_prune, orig_configs = tuner.early_config_prune, tuner.configs + try: + tuner.configs, tuner.early_config_prune = split[:1], None + tuner.cache.clear() + serial = [finegrained_kernels.matmul_2d(A, B, None, Bs) for A, B, Bs in problems] + torch.cuda.synchronize() + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + outs = [[], []] + for stream in streams: + with torch.cuda.stream(stream): + torch.cuda._sleep(200_000_000) + for _ in range(100): + for i, (A, B, Bs) in enumerate(problems): + with torch.cuda.stream(streams[i]): + outs[i].append(finegrained_kernels.matmul_2d(A, B, None, Bs)) + torch.cuda.synchronize() + finally: + tuner.early_config_prune, tuner.configs = orig_prune, orig_configs + tuner.cache.clear() + for i in range(2): + for out in outs[i]: + assert torch.equal(out, serial[i]), "a concurrent split-K launch differs from its serial result" + + +@pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") +@pytest.mark.parametrize("op", ["matmul", "batched", "grouped"]) +def test_int_natural_layout_refused(op): + """int32 words, a checkpoint's natural packing, are refused by every integer op: only uint32 words, the + interleaved layout ``interleave_int_words`` and the repacks emit, are read.""" + B, Bs, _ = WEIGHTS["int4"]["make"](256, 512, 1) + A = torch.randn(4, 512, device=TEST_DEVICE, dtype=torch.bfloat16) + expert_ids = torch.zeros(4, device=TEST_DEVICE, dtype=torch.int32) + expert_start = torch.tensor([0, 4], device=TEST_DEVICE, dtype=torch.int32) + run = { + "matmul": lambda W: finegrained_kernels.matmul_2d(A, W[0], None, Bs[0]), + "batched": lambda W: finegrained_kernels.matmul_batched(A, W, None, Bs, expert_ids=expert_ids), + "grouped": lambda W: finegrained_kernels.matmul_grouped(A, W, None, Bs, expert_start=expert_start), + }[op] + run(B) + with pytest.raises(AssertionError, match="interleave_int_words"): + run(B.view(torch.int32)) + + +@pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") +@pytest.mark.parametrize("checkpoint", ["gptq", "gptq_v2", "gptq_symmetric", "awq", "modelopt_int4", "modelopt_int8", "ggml_q4_0", "ggml_q8_0"]) +def test_int_repacks(checkpoint): + """GPTQ / AWQ int4 linears packed by their own conventions (GPTQ: values along K, zero points along N minus 1, + 0 under ``gptq_v2``; AWQ: both along N in its ``[0, 2, 4, 6, 1, 3, 5, 7]`` column order) repack to weights whose + dequant is ``(q - zero) * scale`` of the source values. ModelOpt stores signed values, INT4 across + output-channel pairs or INT8 unpacked, and both must preserve the original signed-weight reference.""" + + g = torch.Generator(device=TEST_DEVICE).manual_seed(4) + K, N, group = 512, 256, 128 + if checkpoint.startswith(("modelopt", "ggml")): + ggml = checkpoint.startswith("ggml") + bits = 4 if checkpoint in ("modelopt_int4", "ggml_q4_0") else 8 + group = 32 if ggml else (128 if bits == 4 else K) + # Cover dense 2D and stacked experts without a separate codec/metadata test matrix. + leading = () if bits == 4 else (2,) + signed = torch.randint(-(1 << (bits - 1)), 1 << (bits - 1), (*leading, N, K), + device=TEST_DEVICE, dtype=torch.int8, generator=g) + scales = (torch.rand(*leading, N, K // group, device=TEST_DEVICE, generator=g) * 0.02 + 0.005).to( + torch.float16 if ggml else torch.float32 + ) + if ggml: + d = scales.half() + d[..., 1::2] *= -1 # canonical Q4_0 scales can be negative + d[..., 0] = 0 + blocks = torch.empty(*leading, N, K // group, 2 + 4 * bits, device=TEST_DEVICE, dtype=torch.uint8) + blocks[..., :2] = d.contiguous().view(torch.uint8).view(*leading, N, K // group, 2) + values = signed.reshape(*leading, N, K // group, group) + if bits == 4: + values = values.int() + 8 + blocks[..., 2:] = (values[..., :16] | (values[..., 16:] << 4)).to(torch.uint8) + else: + blocks[..., 2:] = values.view(torch.uint8) + W, Ws, Wz = repack_ggml_int(blocks, "Q4_0" if bits == 4 else "Q8_0") + scales = d.float() # GGML reconstructs q*d in FP32 before casting the weight. + elif bits == 4: + qweight = ((signed[0::2].int() & 15) | ((signed[1::2].int() & 15) << 4)).to(torch.uint8) + W, Ws, Wz = repack_modelopt_int4(qweight, scales) + else: + W, Ws, Wz = repack_modelopt_int8(signed, scales.squeeze(-1)) + ref = signed.to(scales.dtype) * scales.repeat_interleave(group, dim=-1) + assert W.dtype == torch.uint32 and Wz is None + assert torch.equal(dequantize_int_weight(W, Ws, K), ref) + A = torch.randn(8, K, device=TEST_DEVICE, dtype=torch.float16, generator=g) + if bits == 4: + out = finegrained_kernels.matmul_2d(A, W, None, Ws) + expected = A.float() @ ref.half().float().T + else: + expert_ids = torch.arange(8, device=TEST_DEVICE, dtype=torch.int32) % 2 + run = torch.compile(finegrained_kernels.matmul_batched, fullgraph=True) + out = run(A, W, None, Ws, expert_ids=expert_ids) + expected = torch.bmm(ref[expert_ids.long()].half().float(), A.float().unsqueeze(-1)).squeeze(-1) + atol, rtol = DTYPE_TO_TOL[torch.float16] + torch.testing.assert_close(out.float(), expected.half().float(), atol=atol, rtol=rtol) + return + q = torch.randint(0, 16, (K, N), device=TEST_DEVICE, dtype=torch.int32, generator=g) + z = torch.randint(0, 16, (K // group, N), device=TEST_DEVICE, dtype=torch.int32, generator=g) + if checkpoint == "gptq_symmetric": + z.fill_(8) + scales = (torch.rand(K // group, N, device=TEST_DEVICE, generator=g) * 0.02 + 0.005).to(torch.float16) + nibble = 4 * torch.arange(8, device=TEST_DEVICE, dtype=torch.int32) + + def pack_n(x, order=range(8)): # (rows, N) -> (rows, N // 8), word nibble i = column order[i] + cols = x.reshape(x.shape[0], N // 8, 8)[..., list(order)] + return (cols << nibble).sum(-1, dtype=torch.int32) + + if checkpoint == "awq": + awq_order = (0, 2, 4, 6, 1, 3, 5, 7) + W, Ws, Wz = repack_awq_int4(pack_n(q, awq_order), pack_n(z, awq_order), scales) + else: + qweight = (q.reshape(K // 8, 8, N) << nibble[None, :, None]).sum(1, dtype=torch.int32) + offset = 0 if checkpoint == "gptq_v2" else 1 + W, Ws, Wz = repack_gptq_int4( + qweight, pack_n((z - offset) & 0xF), scales, zeros_offset=offset, + symmetric=checkpoint == "gptq_symmetric", + ) + assert W.dtype == torch.uint32 + ref = (q.T - z.T.repeat_interleave(group, dim=1)).to(torch.float16) * scales.T.repeat_interleave(group, dim=1) + assert torch.equal(dequantize_int_weight(W, Ws, K, Wz), ref) + + +@pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") +@pytest.mark.parametrize( + "weights,tuner,activation_format", + [ + ("bf16", full_precision_matmul_grouped_kernel, None), + ("mxfp8", mx_dynamic_matmul_grouped_kernel, None), + ("int4", int_weight_only_matmul_grouped_kernel, None), + ("fp8_tensor", block_fp8_weight_only_matmul_grouped_kernel, "bf16"), + ], +) +def test_grouped_packed_schedule_built_by_its_configs(weights, tuner, activation_format): + """A ``PACKED_SCHEDULE`` config reads the tile schedule its own pre_hook builds (``fill_packed_schedule``), the + other configs launch without it: forced to one packed config, a grouped op equals its inline-layout twin (same + tile, same K order) bit for bit.""" + torch.manual_seed(0) + E, S, N, K = 4, 512, 256, 512 + B, Bs, _ = WEIGHTS[weights]["make"](N, K, E) + A = torch.randn(S, K, device=TEST_DEVICE, dtype=torch.bfloat16) + expert_ids = torch.randint(0, E, (S,), device=TEST_DEVICE, dtype=torch.int32) + order = torch.argsort(expert_ids, stable=True).to(torch.int32) + expert_start = torch.zeros(E + 1, device=TEST_DEVICE, dtype=torch.int32) + expert_start[1:] = torch.histc(expert_ids.float(), bins=E, min=0, max=E - 1).to(torch.int32).cumsum(0) + admitted: dict = {} + orig_prune, orig_configs = tuner.early_config_prune, tuner.configs + + def spy(configs, named_args, **kwargs): + admitted["configs"] = orig_prune(configs, named_args, **kwargs) + return admitted["configs"] + + def run(): + tuner.cache.clear() + return finegrained_kernels.matmul_grouped(A, B, None, Bs, expert_start=expert_start, gather_idx=order, activation_format=activation_format) + + try: + tuner.early_config_prune = spy + run() + packed = next(c for c in admitted["configs"] if c.kwargs.get("PACKED_SCHEDULE")) + twin = next( + c for c in orig_configs + if not c.kwargs.get("PACKED_SCHEDULE") and c.num_warps == packed.num_warps + and c.num_stages == packed.num_stages + and {k: v for k, v in c.kwargs.items() if k != "PACKED_SCHEDULE"} + == {k: v for k, v in packed.kwargs.items() if k != "PACKED_SCHEDULE"} + ) + tuner.early_config_prune = None + outs = [] + for c in (packed, twin): + tuner.configs = [c] + outs.append(run()) + finally: + tuner.early_config_prune, tuner.configs = orig_prune, orig_configs + tuner.cache.clear() + assert torch.equal(outs[0], outs[1]), "the packed-schedule config differs from its inline twin" + + +@pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") +@pytest.mark.parametrize("M", [1, 16, 256]) +@pytest.mark.parametrize("op", ["matmul", "grouped", "batched"]) +def test_int4_act_order(op, M): + """Act-order GPTQ (``desc_act=True``, groups scattered across K): repacked to contiguous groups plus a K order the + ops gather A's columns by (``b_perm``), every op matches ``A @ W_ref.T`` with the weight dequantized in A's own + order; an identity ``g_idx`` repacks exactly as no ``g_idx`` (three tensors, same bits).""" + + torch.manual_seed(0) + E, N, K = (1 if op == "matmul" else 4), 256, 512 + W, Ws, Wz, perm, W_ref = make_gptq_act_order(N, K, E) + A = torch.randn(M, K, device=TEST_DEVICE, dtype=torch.bfloat16) + if op == "matmul": + out = finegrained_kernels.matmul_2d(A, W[0], None, Ws[0], b_zero=Wz[0], b_perm=perm[0]) + ref = A.float() @ W_ref[0].T + else: + expert_ids = torch.randint(0, E, (M,), device=TEST_DEVICE, dtype=torch.int32) + ref = torch.einsum("sk,snk->sn", A.float(), W_ref[expert_ids.long()]) + if op == "batched": + out = finegrained_kernels.matmul_batched(A, W, None, Ws, expert_ids=expert_ids, b_zero=Wz, b_perm=perm) + else: + order = torch.argsort(expert_ids, stable=True) + expert_start = torch.zeros(E + 1, device=TEST_DEVICE, dtype=torch.int32) + expert_start[1:] = torch.histc(expert_ids.float(), bins=E, min=0, max=E - 1).to(torch.int32).cumsum(0) + out = finegrained_kernels.matmul_grouped( + A, W, None, Ws, expert_start=expert_start, gather_idx=order.to(torch.int32), b_zero=Wz, b_perm=perm, + ) + ref = ref[order] + err = ((out.float() - ref).norm() / ref.norm()).item() + assert err < 4e-3, f"{op} M={M}: act-order vs dequantized reference {err:.2e}" + + q = torch.randint(0, 2**31 - 1, (K // 8, N), device=TEST_DEVICE, dtype=torch.int32) + z = torch.randint(0, 2**31 - 1, (K // 128, N // 8), device=TEST_DEVICE, dtype=torch.int32) + s = torch.rand(K // 128, N, device=TEST_DEVICE).to(torch.bfloat16) + identity = repack_gptq_int4(q, z, s, g_idx=torch.arange(K, device=TEST_DEVICE, dtype=torch.int32) // 128) + plain = repack_gptq_int4(q, z, s) + assert len(identity) == 3 and all(torch.equal(x, y) for x, y in zip(identity, plain)) diff --git a/finegrained-kernels/tests/utils.py b/finegrained-kernels/tests/utils.py index 5a862bd03..c495985a8 100644 --- a/finegrained-kernels/tests/utils.py +++ b/finegrained-kernels/tests/utils.py @@ -209,6 +209,7 @@ def quant_dequant_a( *, scale: torch.Tensor | None = None, pow2_scale: bool = False, + activation_scale_ub: float | None = None, ) -> torch.Tensor: """Mirror a kernel's inline activation quant→dequant round-trip in fp32. @@ -222,7 +223,11 @@ def quant_dequant_a( return A_fp8.float() * scale M, K = A.shape groups = A.float().reshape(M, K // block_k, block_k) - s = groups.abs().amax(dim=-1, keepdim=True) / 448.0 + amax = groups.abs().amax(dim=-1, keepdim=True) + if activation_scale_ub is not None: + amax = amax.clamp(max=activation_scale_ub) + groups = groups.clamp(min=-amax, max=amax) + s = amax / 448.0 if pow2_scale: s = (_ue8m0_exp(s) << 23).view(torch.float32) A_fp8 = (groups / s.clamp(min=1e-12)).to(FP8_DTYPE) @@ -235,7 +240,7 @@ def quant_dequant_a( # (E, N, K); act_quant[activation_format] -> the host quant fn the ops themselves call (None = the # family default applied to a raw A); dq_act dequantizes its output for the torch oracle. -from finegrained_kernels.quant import fp8_act_quant_block_dynamic, fp8_act_quant_tensor_wide, mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant, nvfp4_quantize_two_level # type: ignore # noqa: E402 +from finegrained_kernels.quant import fp8_act_quant_block_dynamic, interleave_int_words, fp8_act_quant_tensor_wide, mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant, nvfp4_quantize_two_level, repack_gptq_int4 # type: ignore # noqa: E402 _E2M1_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] _E2M1_LUT = _E2M1_LUT + [-v for v in _E2M1_LUT] @@ -320,6 +325,42 @@ def make(N, K, E): return make +INT_GROUP = 128 # GPTQ / AWQ / compressed-tensors' usual group +INT_BITS = {"int2": 2, "int4": 4, "int8": 8} # the integer weight rows below and their value widths + + +def _make_int(bits): + def make(N, K, E, group=INT_GROUP): + """Integer weights in the GEMMs' layout: ``(E, N, K // VPW)`` uint32 words of ``VPW = 32 // bits`` unsigned + values along K in the kernels' interleaved order (``interleave_int_words``) and bf16 group scales; + single-level. A zero point is a scenario knob (``dq_int``).""" + vpw = 32 // bits + q = torch.randint(0, 1 << bits, (E, N, K), device=TEST_DEVICE, dtype=torch.int32) + shifts = bits * torch.arange(vpw, device=TEST_DEVICE, dtype=torch.int32) + natural = (q.reshape(E, N, K // vpw, vpw) << shifts).sum(-1, dtype=torch.int32) + scales = (torch.rand(E, N, K // group, device=TEST_DEVICE) * 0.02 + 0.005).to(torch.bfloat16) + return interleave_int_words(natural, bits), scales, None + + return make + + +def dq_int(B, Bs, Bz=None, dtype=None, group=INT_GROUP): + """Dequant of interleaved integer weight words the compressed-tensors way, the reference the GEMMs match bit for + bit: ``(q - zero)`` in the scale dtype times the per-group scale (rounded once, in that dtype), then cast to + ``dtype`` (the activation's; ``None`` keeps the scale dtype). The width follows from the shapes; the zero point + is ``2 ** (bits - 1)`` when ``Bz`` is None. Returned as fp32.""" + K = Bs.shape[-1] * group + vpw = K // B.shape[-1] + bits = 32 // vpw + p = torch.arange(vpw, device=B.device, dtype=torch.int32) + slots = (B.view(torch.int32).unsqueeze(-1) >> (bits * p)) & ((1 << bits) - 1) # (..., words, slot) + # slot s < vpw / 2 holds K offset 2s, slot s + vpw / 2 offset 2s + 1 + q = torch.stack([slots[..., : vpw // 2], slots[..., vpw // 2 :]], dim=-1).reshape(*B.shape[:-1], K) + zero = (1 << (bits - 1)) if Bz is None else torch.repeat_interleave(Bz.int(), group, dim=-1) + W = (q - zero).to(Bs.dtype) * torch.repeat_interleave(Bs, group, dim=-1) + return (W if dtype is None else W.to(dtype)).float() + + # None follows the weight family, so MXFP8 and MXFP4 rows resolve it differently _MX8_ACT = {None: mxfp8_act_quant, "mxfp8": mxfp8_act_quant, "mxfp4": mxfp4_act_quant} _MX4_ACT = {None: mxfp4_act_quant, "mxfp8": mxfp8_act_quant, "mxfp4": mxfp4_act_quant} @@ -328,9 +369,10 @@ def make(N, K, E): "fp8_128x128": dict( make=lambda N, K, E: (*make_weights(N, K, TEST_DEVICE, [128, 128], num_experts=E), None), dequant=lambda B, Bs, g=None: dq_block_fp8(B, Bs, 128, 128), - activation_formats=("fp8",), + activation_formats=("fp8", "bf16"), requant_formats=("fp8",), act_quant={ + "bf16": None, None: lambda A: fp8_act_quant_block_dynamic(A, 128), "fp8": lambda A: fp8_act_quant_block_dynamic(A, 128), }, @@ -343,9 +385,10 @@ def make(N, K, E): N, K, TEST_DEVICE, [128, 128], scale_dtype=torch.float8_e8m0fnu, num_experts=E ), None), dequant=lambda B, Bs, g=None: dq_block_fp8(B, Bs, 128, 128), - activation_formats=("fp8",), + activation_formats=("fp8", "bf16"), requant_formats=("fp8",), act_quant={ + "bf16": None, None: lambda A: fp8_act_quant_block_dynamic(A, 128, use_ue8m0=True), "fp8": lambda A: fp8_act_quant_block_dynamic(A, 128, use_ue8m0=True), }, @@ -356,6 +399,19 @@ def make(N, K, E): N, K, TEST_DEVICE, None, scale_layout="per_tensor_3d", num_experts=E ), None), dequant=lambda B, Bs, g=None: B.float() * Bs.float().reshape(-1, 1, 1), + activation_formats=("fp8", "bf16"), + requant_formats=(), + act_quant={ + "bf16": None, + None: lambda A: fp8_act_quant_tensor_wide(A, A.shape[-1]), + "fp8": lambda A: fp8_act_quant_tensor_wide(A, A.shape[-1]), + }, + dq_act=lambda q, s: q.float() * s.float().reshape(-1, 1), + ), + # Row-wise FP8: one fp32 scale per output row, (E, N, 1) — Llama 4 FP8; per-token activations. + "fp8_rowwise": dict( + make=lambda N, K, E: (*make_weights(N, K, TEST_DEVICE, [1, K], num_experts=E), None), + dequant=lambda B, Bs, g=None: B.float() * Bs.float(), activation_formats=("fp8",), requant_formats=(), act_quant={ @@ -401,6 +457,30 @@ def make(N, K, E): act_quant={None: nvfp4_act_quant, "nvfp4": nvfp4_act_quant}, dq_act=lambda q, s: dq_grouped(q, s, NVFP4_SCALE_GROUP_K), ), + "int4": dict( + make=_make_int(4), + dequant=lambda B, Bs, g=None: dq_int(B, Bs), + activation_formats=(), + requant_formats=(), + act_quant={None: None}, + dq_act=None, + ), + "int8": dict( + make=_make_int(8), + dequant=lambda B, Bs, g=None: dq_int(B, Bs), + activation_formats=(), + requant_formats=(), + act_quant={None: None}, + dq_act=None, + ), + "int2": dict( + make=_make_int(2), + dequant=lambda B, Bs, g=None: dq_int(B, Bs), + activation_formats=(), + requant_formats=(), + act_quant={None: None}, + dq_act=None, + ), "bf16": dict( make=_make_full(torch.bfloat16), dequant=lambda B, Bs, g=None: B.float(), @@ -455,6 +535,8 @@ def dequantize_weight( """A quantized weight back to fp32, for any format the ops accept: E4M3 or packed-E2M1 values, per-tensor / block / per-row-group scales, with the optional NVFP4 second-level global folded in. Host-side torch — the correctness path, not the fast one.""" + if B.dtype == torch.uint32: # integer words (interleaved), the GEMMs' rounded dequant + return dq_int(B, Bs) if B.dtype == torch.int8: # packed E2M1: two codes per byte along K lut = torch.tensor(_E2M1 + tuple(-v for v in _E2M1), device=B.device, dtype=torch.float32) b = B.view(torch.uint8) @@ -476,3 +558,21 @@ def dequantize_weight( return out +def make_gptq_act_order(N, K, E, group=INT_GROUP, scale_dtype=torch.bfloat16): + """``E`` act-order GPTQ int4 linears (``desc_act=True``: each one's groups scattered across K by a random + ``g_idx``), packed in GPTQ's own layout and repacked: ``(W, Ws, Wz, perm, W_ref)``, ``W``/``Ws``/``Wz`` stacked + ``(E, N, ...)``, ``perm`` ``(E, K)`` int32 and ``W_ref`` ``(E, N, K)`` fp32 the dequantized weight in A's natural + K order, ``(q - zero) * scale`` rounded in the scale dtype.""" + nibble = 4 * torch.arange(8, device=TEST_DEVICE, dtype=torch.int32) + out = [] + for _ in range(E): + q = torch.randint(0, 16, (K, N), device=TEST_DEVICE, dtype=torch.int32) + z = torch.randint(1, 16, (K // group, N), device=TEST_DEVICE, dtype=torch.int32) + s = (torch.rand(K // group, N, device=TEST_DEVICE) * 0.02 + 0.005).to(scale_dtype) + g_idx = (torch.arange(K, device=TEST_DEVICE, dtype=torch.int32) // group)[torch.randperm(K, device=TEST_DEVICE)] + qweight = (q.reshape(K // 8, 8, N) << nibble[None, :, None]).sum(1, dtype=torch.int32) + qzeros = (((z - 1) & 0xF).reshape(K // group, N // 8, 8) << nibble).sum(-1, dtype=torch.int32) + W, Ws, Wz, perm = repack_gptq_int4(qweight, qzeros, s, g_idx=g_idx) + ref = ((q - z[g_idx.long()]).to(scale_dtype) * s[g_idx.long()]).T.float() + out.append((W, Ws, Wz, perm, ref)) + return tuple(torch.stack(t) for t in zip(*out)) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/__init__.py b/finegrained-kernels/torch-ext/finegrained_kernels/__init__.py index f5bdfceb5..1eb71ff21 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/__init__.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/__init__.py @@ -21,7 +21,19 @@ from .formats import get_supported_act_fns from .norm import get_supported_norms, rms_norm_rows from .swizzle import swizzle_mx_scales, unswizzle_mx_scales -from .quant import mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant +from .quant import ( + mxfp8_act_quant, + mxfp4_act_quant, + nvfp4_act_quant, + dequantize_int_weight, + interleave_int_words, + fold_act_order_into_gate_up, + repack_gptq_int4, + repack_awq_int4, + repack_modelopt_int4, + repack_modelopt_int8, + repack_ggml_int, +) __all__ = [ # the three GEMM dispatchers (the weight format is read off the tensors; `activation_format` @@ -38,11 +50,20 @@ # the per-expert output norm a model can hand the forwards by name, and its standalone op "get_supported_norms", "rms_norm_rows", - # load-time helpers: the swizzled scale layout the SM100 scaled-MMA reads, and the row-wise - # quantizers a loader uses to quantize weights into the group formats + # load-time helpers: the swizzled scale layout the SM100 scaled-MMA reads, the row-wise + # quantizers a loader uses to quantize weights into the group formats, and the GPTQ / AWQ int4 + # repacks into the int4 layout "swizzle_mx_scales", "unswizzle_mx_scales", "mxfp8_act_quant", "mxfp4_act_quant", "nvfp4_act_quant", + "dequantize_int_weight", + "interleave_int_words", + "fold_act_order_into_gate_up", + "repack_gptq_int4", + "repack_awq_int4", + "repack_modelopt_int4", + "repack_modelopt_int8", + "repack_ggml_int", ] diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/batched.py b/finegrained-kernels/torch-ext/finegrained_kernels/batched.py index 46cdee294..0538e5011 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/batched.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/batched.py @@ -14,6 +14,8 @@ +from typing import List + import torch import triton import triton.language as tl @@ -24,21 +26,27 @@ from .bayesian_autotuner import bayesian_autotune from .bench import weight_input -from .compat import add_op_namespace_prefix, FP8_DTYPE, MX_SCALE_GROUP_K, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, tl_dtype, pdl_launch_kwargs +from .compat import add_op_namespace_prefix, FP8_DTYPE, MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, narrow_scalar_autotuning_configs, tl_dtype, pdl_launch_kwargs from .descriptors import rebind_batched_mx_bs_descriptor -from .formats import check_activation_format, global_scale_stride, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, is_per_expert_global, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, ue8m0_as_uint8, validate_dense_operands, weight_block_size, weight_format +from .formats import block_fp8_weight_scale_grid, _validate_quantization_args, global_scale_stride, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, is_per_expert_global, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, row_scale_stride, ue8m0_as_uint8, validate_dense_operands, weight_block_size, validate_int_operands, INT_WEIGHT_DTYPES from .epilogue import fused_glu -from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, static_expert_act_operands, tensor_wide_act_operands +from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, prepare_fp8_mx_inputs, permute_routed_rows, static_expert_act_operands, tensor_wide_act_operands from .swizzle import swizzled_scale_descriptor from .scheduling import is_off_rank -from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot +from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, plain_dot, static_dot, MMA_N_ATOM, int_weight_only_compute +from .loading.scales import apply_row_scale from .loading.tiles import ( advance_ptrs, load_act_block_dynamic, load_act_mx, + gather_act_cols, + load_act_k_permuted, load_act_plain, load_act_static, load_weight_block_dynamic, + load_weight_block_fp8, + load_weight_int, + load_weight_int_group, load_weight_mx, load_weight_plain, load_weight_static, @@ -47,7 +55,7 @@ weight_tile_ptrs, ) from .epilogue import acc_init, bias_strides, gemm_epilogue -from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, dot_scaled_staging_pruner, block_fits_dim_pruner, block_within_dim_pruner, compose_pruners, gate_tile_cap_pruner, mx_config_pruner, require_moe_dims_aligned, scale_subblock_pruner, smem_pruner, swizzled_scale_config_pruner, weight_only_swap_scope_pruner +from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, dot_scaled_staging_pruner, block_fits_dim_pruner, block_within_dim_pruner, compose_pruners, gate_tile_cap_pruner, mx_config_pruner, require_moe_dims_aligned, scale_subblock_pruner, single_row_tile_pruner, smem_pruner, swizzled_scale_config_pruner, weight_only_swap_scope_pruner, int_scale_group_tile_pruner @triton.jit @@ -383,7 +391,9 @@ def w8a8_block_static_fp8_matmul_batched_kernel( @bayesian_autotune( # S (routed rows) keyed like the block-dynamic/mx batched siblings — decode re-tunes per batch. get_accelerator_autotuning_configs(tune_block_nk=True, swap_ab=True), - ["N", "K", "S"], + # GATE keys the gate|up arm separately: its 2*BN tile is a different config space, and an + # ungated crown replayed under GATE can overflow shared memory + ["N", "K", "S", "GATE"], n_trials=100, benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, @@ -394,6 +404,7 @@ def w8a8_block_static_fp8_matmul_batched_kernel( fp8_dot_warp_pruner(), block_within_dim_pruner("K"), block_within_dim_pruner("N", "BLOCK_SIZE_N"), + single_row_tile_pruner(), ) }, ) @@ -402,7 +413,7 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel( A, # (S, K) pre-quantized FP8 activations As, # per-token activation scales (S,), or a calibrated (static) one: shared, or per expert B, # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack - Bs, # (num_experts, 1, 1) per-tensor weight scales + Bs, # (num_experts, N, 1) row-wise weight scales, or (num_experts, 1, 1) per tensor (stride_bs_n 0) C, # (S, N) output; under GATE the bf16 GLU intermediate Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None ExpertIds, # (S,) — which expert each batch element routes to @@ -421,6 +432,7 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel( stride_b_k, stride_b_n, stride_bs_e, + stride_bs_n, stride_c_m, stride_c_n, stride_bias_e, @@ -441,16 +453,16 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel( INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, PDL: tl.constexpr = False, ): - """Tensor-scale batched FP8 expert matmul kernel. + """Tensor-scale / row-wise batched FP8 expert matmul kernel. Activations are already quantized; the kernel applies per-token activation - scales and per-expert tensor weight scales. + scales and per-expert weight scales constant along K (per tensor, or per output row). ``SWAP_AB`` (tuner axis, M=1 decode): weight output rows in the MMA M dim (``B`` as ``[BN, BK]``, single token padded to N=16); column 0 of the ``[BN, 16]`` accumulator is the result. Both - scales are per-token/per-tensor scalars, applied once after the loop, orientation-agnostic. - ``GATE`` fuses the gate|up projection (``B`` the interleaved (2N, K) stack, one per-tensor - scale over both halves) into one tile + SwiGLU, emitting the bf16 intermediate; ``GATE=False`` + scales are applied once after the loop, the weight's on the acc's weight-row axis. + ``GATE`` fuses the gate|up projection (``B`` the interleaved (2N, K) stack, its scales + interleaved alike) into one tile + SwiGLU, emitting the bf16 intermediate; ``GATE=False`` is the plain GEMM (bit-identical).""" if PDL: gdc_wait() @@ -478,7 +490,6 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel( offs_k = tl.arange(0, BLOCK_SIZE_K) a_ptrs = operand_tile_ptrs(A, tl.arange(0, BLOCK_SIZE_M) * 0, offs_k, stride_a_m, stride_a_k, "pointer", True) b_ptrs = oriented_tile_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, SWAP_AB) - b_s = tl.load(Bs) a_s = tl.load(As + in_row * stride_as_m + expert_id * stride_as_e) accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB) @@ -494,11 +505,11 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel( "pointer", "pointer", False, False, False, ) - # per-tensor scales fold on the raw accumulator; `gemm_epilogue` owns the finalize. + # the scales fold on the raw accumulator; `gemm_epilogue` owns the finalize. # Finalizing here too was invisible under no-swap (a pass-through) and collapsed the # swapped tile twice, so every SWAP_AB config failed to compile and the tuner quietly # dropped the whole swap arm at this shape. - accumulator = accumulator * a_s * b_s + accumulator = apply_row_scale(accumulator * a_s, Bs, offs_bn, stride_bs_n, SWAP_AB) if PDL: gdc_launch_dependents() gemm_epilogue( @@ -867,6 +878,157 @@ def mx_weight_only_matmul_batched_kernel( ) +@bayesian_autotune( + # the MX weight-only batched kernel's space minus dot_scaled: the upcast dot, plus the swapped decode reduce at + # the MMA tiles and at narrow ones (8 / 16 rows: a decode grid of N / 8 programs, the measured winner) + get_accelerator_autotuning_configs(tune_block_nk=True, compute_modes=("dot", "scalar"), swap_ab=True) + + narrow_scalar_autotuning_configs(block_ns=(8, 16), SWAP_AB=True), + # the width and the scale group decide which tiles fit and compile: part of the key + ["N", "K", "S", "GATE", "INT_BITS", "SCALE_GROUP_K"], + n_trials=100, + benchmark_input=weight_input, + path_anchor_axes=PATH_ANCHOR_AXES, + prune_configs_by={ + # the tail K tile slides back in bounds (see the K-loop), so BK need only fit K + "early_config_prune": compose_pruners( + block_fits_dim_pruner("K"), + block_within_dim_pruner("N", "BLOCK_SIZE_N"), + weight_only_swap_scope_pruner(), + int_scale_group_tile_pruner(divides_k=False), + ) + }, +) +@triton.jit +def int_weight_only_matmul_batched_kernel( + A, # (rows, K) prepared bf16/fp16, FP8 (W4A8), or INT8 (SmoothQuant) activations + AsGlobal, # fp32 calibrated input scale; None for weight-only + B, # (num_experts, N, K // VPW) interleaved uint32 words of VPW = 32 // INT_BITS values; 2N under GATE + Bs, # (num_experts, N, K // SCALE_GROUP_K) float group scales + Bz, # (num_experts, N, K // SCALE_GROUP_K) integer zero points; None for a symmetric weight + BsGlobal, # fp32 W4A8 second-level weight scale; None otherwise + BPerm, # (num_experts, K) int32 act-order K orders, tile column j reading A[:, BPerm[e, j]]; None: natural order + C, + Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None + ExpertIds, + GatherIdx, + ScatterIdx, + S, + N, + K, + stride_a_m, + stride_a_k, + stride_asglobal_e, + stride_b_e, + stride_b_k, + stride_b_n, + stride_bs_e, + stride_bs_n, + stride_bs_k, + stride_bz_e, + stride_bz_n, + stride_bz_k, + stride_bsglobal_e, + stride_perm_e, + stride_c_m, + stride_c_n, + stride_bias_e, + stride_bias_n, + stride_eid, + num_experts, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SCALE_GROUP_K: tl.constexpr, + COMPUTE_MODE: tl.constexpr = "dot", + SWAP_AB: tl.constexpr = False, + GATE: tl.constexpr = False, + ACT_FN: tl.constexpr = "silu", + SWIGLU_ALPHA: tl.constexpr = None, + SWIGLU_LIMIT: tl.constexpr = None, + SIMULATE_UNFUSED: tl.constexpr = False, + INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, + PDL: tl.constexpr = False, + INT_BITS: tl.constexpr = 4, +): + """Integer batched expert matmul: bf16/fp16 weight-only, calibrated INT4/FP8, or INT8 SmoothQuant. + One routed row + one N-tile per + program (expert from ``ExpertIds``). ``SWAP_AB`` puts FP8 weights on the MMA lhs or selects the scalar + decode reduce. ``GATE`` fuses gate|up.""" + if PDL: + gdc_wait() + batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row = expert_setup( + A, B, C, Bs, ExpertIds, GatherIdx, ScatterIdx, + stride_a_m, stride_b_e, stride_c_m, stride_bs_e, stride_eid, ADVANCE_BS=False, + ) + if is_off_rank(expert_id, num_experts): # EP sentinel: non-local expert, output left uninit + return + n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N + # a non-128 N's last tile wraps back into the tensor; the epilogue's N_COLS masks those columns + offs_bn = (pid_n * n_width + tl.arange(0, n_width)) % (2 * N if GATE else N) + # Repeat this routed token across the swapped MMA's N atom; FAKE_BATCH stores column zero. + ACT_ROWS: tl.constexpr = MMA_N_ATOM if SWAP_AB and COMPUTE_MODE == "dot" else BLOCK_SIZE_M + a_ptrs = operand_tile_ptrs( + A, tl.arange(0, ACT_ROWS) * 0, tl.arange(0, BLOCK_SIZE_K), stride_a_m, stride_a_k, "pointer", True, + ) + k_perm = None + a_rows = None + if BPerm is not None: + k_perm = BPerm + expert_id * stride_perm_e + a_rows = A + (tl.arange(0, ACT_ROWS) * 0)[:, None] * stride_a_m + # output-rows-major: the scalar reduce and the BM <= 16 mma.sync tiles (int_dequant_tile) + VPW: tl.constexpr = 32 // INT_BITS # values per uint32 word + b_ptrs = weight_tile_ptrs(B, offs_bn, tl.arange(0, BLOCK_SIZE_K // VPW), stride_b_n, stride_b_k, True) + accumulator = acc_init(COMPUTE_MODE, BLOCK_SIZE_M, n_width, SWAP_AB, tl.int32 if A.dtype.element_ty == tl.int8 else tl.float32) + k_last = tl.cdiv(K, BLOCK_SIZE_K) - 1 + a_next = None + if BPerm is not None: # BK fits K: the first tile never slides + a_next = gather_act_cols(a_rows, k_perm, 0, None, stride_a_k, BLOCK_SIZE_K) + # maskless K-loop indexed off the base tile; the last tile slides back to end at K and zeroes the + # activation lanes the previous trip already accumulated + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + k_start = tl.minimum(k * BLOCK_SIZE_K, K - BLOCK_SIZE_K) + overlap = k * BLOCK_SIZE_K - k_start + a = load_act_k_permuted(a_ptrs + k_start * stride_a_k, 0, 0, k_start, None, 0, a_next, "pointer", False) + if BPerm is not None: + k_next = tl.minimum(tl.minimum(k + 1, k_last) * BLOCK_SIZE_K, K - BLOCK_SIZE_K) + a_next = gather_act_cols(a_rows, k_perm, k_next, None, stride_a_k, BLOCK_SIZE_K) + if overlap > 0: + a = tl.where(tl.arange(0, BLOCK_SIZE_K)[None, :] >= overlap, a, 0) + b_k = b_ptrs + (k_start // VPW) * stride_b_k + w, w_s = load_weight_int( + b_k, b_k, Bs, Bs, 0, 0, 0, expert_id, pid_n, k, N, K, stride_bs_e, stride_bs_n, stride_bs_k, + GATE, True, "pointer", True, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, + k_col_off=-(overlap // SCALE_GROUP_K), BITS=INT_BITS, + ) + w_z = None + if Bz is not None: + w_z = load_weight_int_group( + Bz, Bz, expert_id, pid_n, k, N, K, stride_bz_e, stride_bz_n, stride_bz_k, + GATE, True, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, k_col_off=-(overlap // SCALE_GROUP_K), + ) + accumulator = int_weight_only_compute( + accumulator, a, w, w_s, w_z, BLOCK_SIZE_K, n_width, SCALE_GROUP_K, COMPUTE_MODE, True, INT_BITS, SWAP_AB, + weight_global=None if BsGlobal is None else tl.load(BsGlobal + expert_id * stride_bsglobal_e), + ) + + if PDL: + gdc_launch_dependents() + output_scale = BsGlobal + output_scale_stride = stride_bsglobal_e + if A.dtype.element_ty == tl.int8: + output_scale = Bs + offs_bn * stride_bs_n + output_scale_stride = stride_bs_e + gemm_epilogue( + C, None, accumulator, out_row, pid_n, 0, out_row, 1, stride_c_n, 1, 1, + BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K, + ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, + COMPUTE_MODE=COMPUTE_MODE, SWAP_AB=SWAP_AB, FAKE_BATCH=True, N_COLS=N, global_row=expert_id, + Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n, + GlobalScale=output_scale, stride_global_e=output_scale_stride, + GlobalScaleA=AsGlobal, stride_global_a=stride_asglobal_e, + ) + + @bayesian_autotune( get_accelerator_autotuning_configs(tune_block_nk=True, swap_ab=True), # S (routed rows) keyed like the fp8/mx batched siblings — decode re-tunes per batch; @@ -985,7 +1147,129 @@ def full_precision_matmul_batched_kernel( ) +@bayesian_autotune( + get_accelerator_autotuning_configs(tune_block_nk=True, swap_ab=True, swap_block_sizes_m=(1, 2)), + # S (routed rows) keyed like the fp8/mx batched siblings — decode re-tunes per batch; + # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide). + ["N", "K", "S", "GATE", "BLOCK_N", "BLOCK_K", "GATE_SCALES"], + n_trials=100, + benchmark_input=weight_input, + path_anchor_axes=PATH_ANCHOR_AXES, + # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is + # row-masked only — veto non-dividing tiles on both. + prune_configs_by={ + "early_config_prune": compose_pruners( + block_within_dim_pruner("K"), + block_within_dim_pruner("N", "BLOCK_SIZE_N"), + smem_pruner(), + ) + }, +) +@triton.jit +def block_fp8_weight_only_matmul_batched_kernel( + A, # (rows, K) BF16/FP16 activations + B, # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack + Bs, # coarse FP8 weight scales + C, # (S, N) output; under GATE the GLU intermediate + Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None + ExpertIds, # (S,) — which expert each batch element routes to + GatherIdx, # (S,) int — batch_id -> source row of A; read only when not None + ScatterIdx, # (S,) int — batch_id -> destination row of C; read only when not None + # Shape + S, + N, + K, + num_experts, + # Strides + stride_a_m, + stride_a_k, + stride_b_e, + stride_b_k, + stride_b_n, + stride_bs_e, + stride_bs_n, + stride_bs_k, + stride_c_m, + stride_c_n, + stride_bias_e, + stride_bias_n, + stride_eid, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + SWAP_AB: tl.constexpr = False, + GATE: tl.constexpr = False, + GATE_SCALES: tl.constexpr = False, + ACT_FN: tl.constexpr = "silu", + SWIGLU_ALPHA: tl.constexpr = None, + SWIGLU_LIMIT: tl.constexpr = None, + SIMULATE_UNFUSED: tl.constexpr = False, + INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, + PDL: tl.constexpr = False, +): + """Batched FP8 W8A16 expert GEMM with optional fused gate/up activation.""" + if PDL: + gdc_wait() + batch_id, pid_n, expert_id, A, B, C, unused_scale_base, in_row, out_row = expert_setup( + A, + B, + C, + Bs, # the scale leaf applies the expert offset + ExpertIds, + GatherIdx, + ScatterIdx, + stride_a_m, + stride_b_e, + stride_c_m, + stride_bs_e, + stride_eid, + ADVANCE_BS=False, + ) + # EP sentinel: row routed to a non-local expert; output is left uninit. + if is_off_rank(expert_id, num_experts): + return + + n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N + offs_bn = pid_n * n_width + tl.arange(0, n_width) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = operand_tile_ptrs(A, tl.arange(0, BLOCK_SIZE_M) * 0, offs_k, stride_a_m, stride_a_k, "pointer", True) + b_ptrs = weight_tile_ptrs( + B, offs_bn, offs_k, stride_b_n, stride_b_k, SWAP_AB + ) + + accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a, unused_activation = load_act_plain(a_ptrs, 0, 0, 0, None, 0, "pointer", False) + w = load_weight_block_fp8( + b_ptrs, b_ptrs, Bs, 0, 0, 0, + expert_id, offs_bn, k * BLOCK_SIZE_K + offs_k, + stride_bs_e, stride_bs_n, stride_bs_k, + GATE, False, "pointer", SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K, + BLOCK_N, BLOCK_K, GATE_SCALES, a.dtype, + ) + accumulator = plain_dot(accumulator, a, w, SWAP_AB) + a_ptrs += BLOCK_SIZE_K * stride_a_k + b_ptrs += BLOCK_SIZE_K * stride_b_k + + if PDL: + gdc_launch_dependents() + gemm_epilogue( + C, None, accumulator, out_row, pid_n, 0, out_row, 1, stride_c_n, 1, 1, + BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K, + ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, + COMPUTE_MODE="dot", FAKE_BATCH=True, + Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n, + global_row=expert_id, + ) + + GATE_UNSTACK_MAX_S = 16 # the unstacked-gate decode band (see the dispatch in the wrapper) +# act-order routed rows the integer kernel gathers in-kernel; above, permute_routed_rows first (B200, Qwen3-30B-A3B +# gate|up: S=8 8.8 vs 9.8us, S=32 21.3 vs 21.6us, S=128 70.6 vs 68.0us, S=512 236.6 vs 224.4us) +ACT_ORDER_GATHER_MAX_S = 8 def unstacked_gate(gemm, output_dtype, simulate_unfused, act_fn, swiglu_alpha, swiglu_limit, @@ -1300,8 +1584,9 @@ def w8a8_tensor_dynamic_fp8_matmul_batched( gather_idx: torch.Tensor | None = None, scatter_idx: torch.Tensor | None = None, bias: torch.Tensor | None = None, + activation_scale_ub: float | None = None, ) -> torch.Tensor: - """Tensor-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T. ``A`` raw + """Tensor-scale / row-wise batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T. ``A`` raw (``As`` None) -> quantized here (offline, per-token); else pre-quantized (``As`` given). ``gather_idx``/``scatter_idx`` map the source row of A / destination row of C per program (None = row s). @@ -1310,7 +1595,8 @@ def w8a8_tensor_dynamic_fp8_matmul_batched( B: (num_experts, N, K) FP8 expert weights; under ``gate`` the (num_experts, 2N, K) stack As: (rows,) per-token scales alongside a pre-quantized A, or — on a raw A — the calibrated (static) activation scale: one value, or one per expert - Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales + Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales, or (num_experts, N, 1) + row-wise (2N rows under ``gate``) """ validate_dense_operands(A, B) @@ -1321,11 +1607,12 @@ def w8a8_tensor_dynamic_fp8_matmul_batched( # under gate|up fusion B is the (E, 2N, K) stack; N is the per-projection output width N = rows // 2 if gate else rows - # Normalize Bs to (num_experts, 1, 1) Bs = normalize_per_expert_scale(Bs, num_experts) bs_u8 = ue8m0_as_uint8(Bs) - qA, As, as_stride_m, as_stride_e = tensor_wide_act_operands(A, As, num_experts) + qA, As, as_stride_m, as_stride_e = tensor_wide_act_operands( + A, As, num_experts, activation_scale_ub=activation_scale_ub + ) C = A.new_empty(S, N, dtype=output_dtype) def grid(META): @@ -1356,6 +1643,7 @@ def grid(META): B.stride(2), B.stride(1), bs_u8.stride(0), + row_scale_stride(bs_u8), C.stride(0), C.stride(1), bias_stride_e, @@ -1396,6 +1684,7 @@ def mx_dynamic_matmul_batched( b_global_scale: torch.Tensor | None = None, output_global_scale: torch.Tensor | None = None, bias: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, ) -> list[torch.Tensor]: """Batched MX matmul ``C[s] = A[s] @ B[expert_ids[s]].T``; activations quantized inline in the kernel (decode: one act row per program, inline is free). The @@ -1420,6 +1709,22 @@ def mx_dynamic_matmul_batched( activation_format = resolve_activation_format(activation_format, B, Bs) output_format = activation_format if quantize_output else None requant = quantize_output + if a_pre_quant_scale is not None and activation_format != "fp8": + assert As is None, "pre_quant_scale applies before activation quantization; pass raw inputs" + A = permute_routed_rows( + A, None, expert_ids.numel(), gather_idx, expert_ids=expert_ids, + pre_quant_scale=a_pre_quant_scale, num_experts=B.shape[0], + ) + gather_idx = None + if activation_format == "fp8": + assert not quantize_output and a_global_scale is None, "ordinary FP8 inputs use As=input_scale and no fused requant" + output_dtype = resolve_output_dtype(output_dtype, A, None) + group = mx_scale_family(Bs, A.shape[-1]) + A, As, a_global_scale = prepare_fp8_mx_inputs( + A, As, ue8m0_as_uint8(Bs).dtype, group, B.shape[0], expert_ids.numel(), + gather_idx, expert_ids=expert_ids, pre_quant_scale=a_pre_quant_scale, + ) + gather_idx = None if As is not None: assert (As.dtype == torch.float8_e4m3fn) == (Bs.dtype == torch.float8_e4m3fn), ( f"activation scales ({As.dtype}) must match the weight scale family ({Bs.dtype})" @@ -1667,6 +1972,91 @@ def grid(META): return [C] +@compile_time_only_triton_op( + add_op_namespace_prefix("block_fp8_weight_only_matmul_batched"), + mutates_args=(), + opaque=True, +) +def block_fp8_weight_only_matmul_batched( + A: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + expert_ids: torch.Tensor, + gate: bool = False, + act_fn: str = "silu", + swiglu_alpha: float | None = None, + swiglu_limit: float | None = None, + simulate_unfused: bool = False, + activation_format: str | None = None, + quantize_output: bool = False, + output_dtype: torch.dtype | None = None, + gather_idx: torch.Tensor | None = None, + scatter_idx: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + *, + interleaved_scales: bool = False, +) -> List[torch.Tensor]: + """FP8 W8A16 expert GEMM with in-tile weight rescaling.""" + validate_dense_operands(A, B) + assert A.dtype in (torch.bfloat16, torch.float16), "FP8 weight-only takes raw BF16/FP16 activations" + Bs, scale_block_n, scale_block_k = block_fp8_weight_scale_grid(B, Bs, gate or interleaved_scales) + Bs = ue8m0_as_uint8(Bs) + assert activation_format in (None, "bf16") and not quantize_output, ( + "FP8 weight-only requires unquantized activations and output" + ) + + output_dtype = resolve_output_dtype(output_dtype, A, None) + K = A.shape[1] + S = expert_ids.shape[0] + num_experts, _, N = expert_weight_shape(B, gate) + C = A.new_empty(S, N, dtype=output_dtype) + + def grid(META): + return (S, triton.cdiv(N, META["BLOCK_SIZE_N"])) + + with device_context(A.device): + bias_stride_e, bias_stride_n = bias_strides(bias) + compile_time_only_triton_wrap(block_fp8_weight_only_matmul_batched_kernel)[grid]( + A, + B, + Bs, + C, + bias, + expert_ids, + gather_idx, # None = A is expert-sorted; read only when not None (folds at trace time) + scatter_idx, # None = C is expert-sorted; read only when not None (folds at trace time) + S, + N, + K, + num_experts, + A.stride(0), + A.stride(1), + B.stride(0), + B.stride(2), + B.stride(1), + Bs.stride(0), + Bs.stride(1), + Bs.stride(2), + C.stride(0), + C.stride(1), + bias_stride_e, + bias_stride_n, + expert_ids.stride(0), + BLOCK_N=scale_block_n, + BLOCK_K=scale_block_k, + GATE=gate, + GATE_SCALES=(gate or interleaved_scales) and Bs.shape[1] > 1, + ACT_FN=act_fn, + SWIGLU_ALPHA=swiglu_alpha, + SWIGLU_LIMIT=swiglu_limit, + SIMULATE_UNFUSED=simulate_unfused, + INTERMEDIATE_DTYPE=tl_dtype(output_dtype), + **pdl_launch_kwargs(), + ) + + return [C] + + @compile_time_only_triton_op( add_op_namespace_prefix("mx_weight_only_matmul_batched"), mutates_args=(), @@ -1720,6 +2110,7 @@ def mx_weight_only_matmul_batched( ) N = rows // 2 if gate else rows scale_group = mx_scale_family(Bs, K) + assert Bs.dtype != torch.float8_e4m3fn or scale_group == NVFP4_SCALE_GROUP_K, "group-32 E4M3 FP4 weights require calibrated FP8 activations" b_global_scale = normalize_global_scale(b_global_scale, num_experts) C = A.new_empty(S, N, dtype=output_dtype) b_u8 = e2m1_as_uint8(B) @@ -1772,6 +2163,118 @@ def grid(META): return [C] +@compile_time_only_triton_op( + add_op_namespace_prefix("int_weight_only_matmul_batched"), + mutates_args=(), + opaque=True, +) +def int_weight_only_matmul_batched( + A: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + expert_ids: torch.Tensor, + gate: bool = False, + act_fn: str = "silu", + swiglu_alpha: float | None = None, + swiglu_limit: float | None = None, + simulate_unfused: bool = False, + output_dtype: torch.dtype | None = None, + gather_idx: torch.Tensor | None = None, + scatter_idx: torch.Tensor | None = None, + b_zero: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + b_perm: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, + a_scale: torch.Tensor | None = None, + b_scale2: torch.Tensor | None = None, +) -> list[torch.Tensor]: + """Integer weight-only batched (decode) matmul: ``C[s] = A[s] @ B[expert_ids[s]].T`` with raw bf16/fp16 + activations against 2/4/8-bit integer weights, ``B`` ``(E, N, K // VPW)`` interleaved uint32 words + (``int_weight_only_matmul_2d``), ``Bs`` ``(E, N, K // group)`` float + group scales, ``b_zero`` the matching integer zero points (``None``: symmetric) and ``b_perm`` ``(E, K)`` int32 + the act-order K orders (``None``: natural): up to ``ACT_ORDER_GATHER_MAX_S`` routed rows the kernel gathers each + row's columns itself, above it one ``permute_routed_rows`` launch lays the rows out first. Returns ``[C]``. The weight + is dequantized to the activation dtype (``((q - zero) * scale)`` rounded in the scale dtype, then cast), as the torch + "dequantize, then matmul" reference does, so results match it up to fp32 summation order.""" + bits = validate_int_operands(A, B, Bs, b_zero, b_perm, a_scale, b_scale2) + smoothquant = a_scale is not None and b_scale2 is None + a_scale = normalize_global_scale(a_scale, B.shape[0]) + b_scale2 = normalize_global_scale(b_scale2, B.shape[0]) + # gate runs unstacked (``unstacked_gate``), like the MX weight-only op + output_dtype = resolve_output_dtype(output_dtype, A, None) + if gate: + return unstacked_gate( + lambda intermediate_dtype: int_weight_only_matmul_batched( + A, B, Bs, expert_ids, output_dtype=intermediate_dtype, + gather_idx=gather_idx, scatter_idx=scatter_idx, b_zero=b_zero, bias=bias, b_perm=b_perm, a_pre_quant_scale=a_pre_quant_scale, a_scale=a_scale, b_scale2=b_scale2, + ), + output_dtype, simulate_unfused, act_fn, swiglu_alpha, swiglu_limit, + ) + K = A.shape[1] + S = expert_ids.shape[0] + num_experts, N, K_b = B.shape + if a_scale is not None or a_pre_quant_scale is not None or (b_perm is not None and S > ACT_ORDER_GATHER_MAX_S): + A = permute_routed_rows( + A, b_perm, S, gather_idx, expert_ids=expert_ids, + pre_quant_scale=a_pre_quant_scale, num_experts=num_experts, quant_scale=a_scale, + quant_dtype=torch.int8 if smoothquant else FP8_DTYPE, + ) + gather_idx = b_perm = None + C = A.new_empty(S, N, dtype=output_dtype) + + def grid(META): + return (S, triton.cdiv(N, META["BLOCK_SIZE_N"])) + + with device_context(A.device): + bias_stride_e, bias_stride_n = bias_strides(bias) + compile_time_only_triton_wrap(int_weight_only_matmul_batched_kernel)[grid]( + A, + a_scale, + B, + Bs, + b_zero, + b_scale2, + b_perm, + C, + bias, + expert_ids, + gather_idx, + scatter_idx, + S, + N, + K, + A.stride(0), + A.stride(1), + global_scale_stride(a_scale), + B.stride(0), + B.stride(2), + B.stride(1), + Bs.stride(0), + Bs.stride(1), + Bs.stride(2), + *(b_zero.stride() if b_zero is not None else (0, 0, 0)), + global_scale_stride(b_scale2), + b_perm.stride(0) if b_perm is not None else 0, + C.stride(0), + C.stride(1), + bias_stride_e, + bias_stride_n, + expert_ids.stride(0), + num_experts=num_experts, + SCALE_GROUP_K=K // Bs.shape[2], + INT_BITS=bits, + GATE=gate, + ACT_FN=act_fn, + SWIGLU_ALPHA=swiglu_alpha, + SWIGLU_LIMIT=swiglu_limit, + SIMULATE_UNFUSED=simulate_unfused, + INTERMEDIATE_DTYPE=tl_dtype(output_dtype), + **pdl_launch_kwargs(), + ) + + return [C] + + def matmul_batched( A: torch.Tensor, B: torch.Tensor, @@ -1781,6 +2284,7 @@ def matmul_batched( expert_ids: torch.Tensor, bias: torch.Tensor | None = None, # (E, N_out) per-expert output bias; (N_out,) for 2D activation_format: str | None = None, + activation_scale_ub: float | None = None, gate: bool = False, act_fn: str = "silu", swiglu_alpha: float | None = None, @@ -1793,6 +2297,9 @@ def matmul_batched( a_global_scale: torch.Tensor | None = None, b_global_scale: torch.Tensor | None = None, output_global_scale: torch.Tensor | None = None, + b_zero: torch.Tensor | None = None, + b_perm: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Batched matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). Routes one program per routed row (``expert_ids`` gives its expert). @@ -1824,23 +2331,30 @@ def matmul_batched( - ``Bs`` None → ``full_precision_matmul_batched`` (plain dot, no scales anywhere). - MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0 group-32 ``Bs`` → ``mx_dynamic_matmul_batched``. - - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) → + - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) or per output row (``(E, N, 1)``) → ``w8a8_tensor_dynamic_fp8_matmul_batched``. - block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_batched``. + + ``activation_scale_ub`` caps the dynamic per-token amax for row-wise FP8; + it requires raw activations and applies before conversion to E4M3. + - uint32 integer weights → ``int_weight_only_matmul_batched`` (weight-only), ``b_zero`` their zero points + (``None``: symmetric), ``b_perm`` ``(E, K)`` their act-order K orders (``None``: natural). """ - check_activation_format(activation_format, quantize_output) - assert (a_global_scale is None and b_global_scale is None) or ( - Bs is not None and weight_format(B, Bs) == "nvfp4" - ), "two-level globals (a_global_scale / b_global_scale) are NVFP4-only" - assert output_global_scale is None or ( - Bs is not None and quantize_output and resolve_activation_format(activation_format, B, Bs) == "nvfp4" - ), ( - "output_global_scale is the NVFP4 requant second level — it requires quantize_output=True on NVFP4 " - "(the epilogue would otherwise normalize by it with nothing downstream to compensate)" + _validate_quantization_args( + A, As, B, Bs, a_global_scale, b_global_scale, output_global_scale, + b_zero, b_perm, a_pre_quant_scale, activation_format, quantize_output, activation_scale_ub, ) + if B.dtype == FP8_DTYPE and Bs is not None and not is_mx(B, Bs) and activation_format == "bf16": + assert As is None and not quantize_output, "weight-only takes raw activations without requantization" + out = block_fp8_weight_only_matmul_batched( + A, B, Bs, expert_ids, gate, act_fn, swiglu_alpha, swiglu_limit, + simulate_unfused, activation_format, quantize_output, output_dtype, + gather_idx, scatter_idx, bias=bias, + ) + return out[0] # a calibrated (static) scale on a raw A (see `tensor_wide_act_operands`): block-scale weights # have a dedicated static kernel, per-tensor ones read the same As on the tensor-wide arm - static_act = As is not None and As.ndim <= 1 and As.numel() in (1, B.shape[0]) and A.dtype != FP8_DTYPE + static_act = B.dtype not in INT_WEIGHT_DTYPES and not is_mx(B, Bs) and As is not None and As.ndim <= 1 and As.numel() in (1, B.shape[0]) and A.dtype != FP8_DTYPE if static_act: assert Bs is not None and not is_mx(B, Bs), ( "a calibrated (static) activation scale is an FP8 form — MX activations carry a scale " @@ -1888,6 +2402,11 @@ def matmul_batched( scatter_idx, bias=bias, ) + elif B.dtype in INT_WEIGHT_DTYPES: # packed integer weights + out = int_weight_only_matmul_batched( + A, B, Bs, expert_ids, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, output_dtype, + gather_idx, scatter_idx, b_zero=b_zero, bias=bias, b_perm=b_perm, a_pre_quant_scale=a_pre_quant_scale, a_scale=As, b_scale2=b_global_scale, + ) elif is_mx(B, Bs) and activation_format == "bf16": # weight-only: raw bf16 acts, MX weight upcast in-MMA assert As is None and a_global_scale is None and not quantize_output, ( "weight-only (activation_format='bf16') takes a raw activation, no As/global/requant" @@ -1931,7 +2450,7 @@ def matmul_batched( a_global_scale, b_global_scale, output_global_scale, - bias=bias, + bias=bias, a_pre_quant_scale=a_pre_quant_scale, ) elif (block_size := weight_block_size(B, Bs)) is None: assert activation_format in (None, "fp8") and not quantize_output, ( @@ -1940,6 +2459,7 @@ def matmul_batched( out = w8a8_tensor_dynamic_fp8_matmul_batched( A, B, As, Bs, expert_ids, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, output_dtype, gather_idx, scatter_idx, bias=bias, + activation_scale_ub=activation_scale_ub, ) else: out = w8a8_block_dynamic_fp8_matmul_batched( diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/compat.py b/finegrained-kernels/torch-ext/finegrained_kernels/compat.py index 06cf5917e..776b99d3c 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/compat.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/compat.py @@ -99,6 +99,12 @@ def add_op_namespace_prefix(name: str) -> str: +def launches_skipped() -> bool: + """Whether this call runs inside an opaque op's fake impl, where only output allocation counts + (see ``compile_time_only_triton_wrap``) — for the non-Triton launches an op makes.""" + return _SKIP_LAUNCHES_MIRROR + + def compile_time_only_triton_wrap(kernel): """``wrap_triton`` while torch.compile is tracing (required to capture a raw Triton launch into the graph), the bare kernel in eager. Every kernel launch goes through @@ -224,18 +230,26 @@ def sm_count(device_index: int) -> int: @functools.lru_cache(maxsize=8) -def persistent_program_count(device_index: int) -> int: - """Grid for the persistent grouped GEMMs, whose tile loop strides by a matching NUM_SMS - constexpr — the program count is a free parameter, not a shape. - - One program per processor is CUDA sizing: an SM runs the whole CTA and hides latency - within it. An Xe-core does not — it holds several hardware threads per vector engine and - needs more than one work-group to fill them. Measured on Xe, BF16 MoE prefill: 2x is - 1.6x faster than 1x (bit-identical output), and 4x and beyond give the win back.""" +def persistent_program_count(device_index: int, block_m: int | None = None) -> int: + """Size the persistent grid, including extra programs for small SM10x tiles. + + Small tiles need several programs per SM to hide memory latency. Xe uses + two work-groups per core regardless of the tile size. + """ n = sm_count(device_index) + if block_m is not None and block_m < 64 and is_sm10x(): + n *= persistent_programs_per_sm(block_m) return n * 2 if get_active_device_type() == "xpu" else n +def persistent_programs_per_sm(block_m: int) -> int: + """Overlap small persistent tiles' load latency with four programs per SM. + + Tiles of 128 rows already occupy an SM's shared memory, so extra programs only add scheduling overhead. + """ + return 1 if block_m >= 128 else 4 + + def tl_dtype(dtype: torch.dtype) -> tl.dtype: """The ``tl`` dtype matching a torch dtype (``torch.bfloat16`` → ``tl.bfloat16``) — the @@ -290,6 +304,8 @@ def get_accelerator_autotuning_configs( tune_block_n: bool = False, tune_block_nk: bool = False, swap_ab: bool = False, + swap_block_sizes_m=(1,), + swap_ab_prefill: bool = False, warp_spec: bool = False, packed_schedule: bool = False, compute_modes=None, @@ -326,6 +342,10 @@ def get_accelerator_autotuning_configs( bursts are the one in-block lever that helps (+12% on the dsv4 swap GEMV; deeper num_stages and split-K both measured WORSE). Prefill tunes carry the same rows as trial-budget dilution (descriptor_box_pruner drops the illegal e4m3 boxes). + - ``swap_block_sizes_m``: token tiles for swapped plain GEMMs that retain a normal + logical accumulator. Padded-single-token helpers require the default ``(1,)``. + - ``swap_ab_prefill``: append the measured plain-dot prefill tiles with computed + weight rows on the MMA lhs; dense/grouped callers share the same candidate set. - ``warp_spec``: warp-specialize the K-loop, crossed over every mode — pair it with ``warp_spec_compile_guard_pruner``, which owns the can't-compile regions (including dot_scaled + WS, a PassManager failure). Compile support is @@ -368,7 +388,6 @@ def get_accelerator_autotuning_configs( # # Axes measured DEAD and deliberately absent (do not re-add without fresh evidence): # num_ctas slower everywhere + miscompiles with TMA at num_stages<=4 - # flatten (tl.range) proxy +17% never survived e2e (<=1%; gate_up regressed) # persistent 2D grid-stride 2D prototype's win was an inline-vs-helper artifact # batched WARP_SPEC decode loops are launch/latency-bound; WS never won a batched # cell, so the axis is not emitted at the six batched sites @@ -448,21 +467,20 @@ def get_accelerator_autotuning_configs( # The coupled decode (BLOCK_SIZE_M, SWAP_AB) pairs — the single source of the swap/BM # coupling. At M=1 the MMA 16-atom is filled either by replicating the token (non-swap # BM=16, ~40% over the degenerate BM=1 on plain tl.dot) or by putting the weight's - # output rows in M (swap, BM=1). Swap REQUIRES BM=1 (its rhs is the single token - # padded to the N=16 atom), so (16, swap) never exists; the tuner picks. Prefill - # kernels leave SWAP_AB unset — the prefill swap arms measured losers everywhere - # (see the bd 2D kernel's tuner note). + # Swapped decode token tiles default to BM=1 for padded-token helpers. Plain GEMMs + # can opt into larger token tiles or the measured computed-lhs prefill candidates. if swap_ab and tune_block_m: - # tuned-BM kernels (the 2D mx GEMM): the prefill BM sweep keeps no-swap; the swap - # decode rows are APPENDED at their structural BM=1 (dedup collapses the copies) + # Keep the normal BM sweep and append the requested swapped token tiles. blocks = [{**b, "SWAP_AB": False} for b in blocks] + [ - {**b, "BLOCK_SIZE_M": 1, "SWAP_AB": True} for b in blocks + {**b, "BLOCK_SIZE_M": bm, "SWAP_AB": True} + for b in blocks for bm in swap_block_sizes_m ] elif swap_ab: + orientations = [(1, False), (16, False)] + [(bm, True) for bm in swap_block_sizes_m] blocks = [ {**b, "BLOCK_SIZE_M": bm, "SWAP_AB": sw} for b in blocks - for bm, sw in ((1, False), (16, False), (1, True)) + for bm, sw in orientations ] # scalar's BM is STRUCTURALLY 1 (a single-row GEVM reduce — a bigger-BM scalar @@ -481,11 +499,102 @@ def get_accelerator_autotuning_configs( for w in num_warps for s in num_stages ] + if swap_ab_prefill: + for config in configs: + config.kwargs["SWAP_AB"] = False + if is_sm10x(): + for stages in (2, 3): + kwargs = dict(BLOCK_SIZE_M=128, BLOCK_SIZE_N=128, BLOCK_SIZE_K=128, + COMPUTE_MODE="dot", WARP_SPEC=True, SWAP_AB=True, + A_MEMORY_MODE="host_descriptor", B_MEMORY_MODE="host_descriptor") + if packed_schedule: + kwargs["PACKED_SCHEDULE"] = False + configs.append(triton.Config(kwargs, num_warps=8, num_stages=stages, pre_hook=pre_hook)) # dedup: the scalar BM=1 pin and the MX "dot" BK override collapse swept axes onto # one value, leaving identical copies that would each be benched. return list({str(c): c for c in configs}.values()) +def pipeline_autotuning_configs(configs, pipelines, *, num_warps=None): + """Expand selected persistent GEMM tiles over compiler-managed output-tile pipelines.""" + return [ + triton.Config( + dict(config.kwargs, PIPELINE=pipeline), + num_warps=num_warps or config.num_warps, + num_stages=config.num_stages, + num_ctas=config.num_ctas, + maxnreg=config.maxnreg, + pre_hook=config.pre_hook, + ) + for config in configs + for pipeline in pipelines + ] + + +def narrow_scalar_autotuning_configs(block_ms=(1,), block_ns=(8, 16, 32), **fixed): + """The tiles of a ``scalar`` arm that unpacks a weight tile once for ``BLOCK_SIZE_M`` token rows and reduces it on + the CUDA cores (a weight-rows-major GEVM): 8-32 weight rows per program, so a decode grid fills the SMs where the + MMA tiles leave them idle (int4 at N=4096: 128 programs at BN=32, 512 at BN=8), and K bursts up to 1024. A tile + whose broadcast product (BM x BN x BK fp32) would not fit the registers is left out; ``fixed`` adds kwargs.""" + blocks = [ + {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "COMPUTE_MODE": "scalar", **fixed} + for bm in block_ms + for bn in block_ns + for bk in (128, 256, 512, 1024) + if bm * bn * bk <= 8192 + ] + return [triton.Config(b, num_warps=w, num_stages=s) for b in blocks for w in (2, 4) for s in (2, 3, 4, 5, 6)] + + + +def split_k_autotuning_configs(splits, memory_modes=("descriptor", "pointer"), pre_hook=None, configs=None, **fixed): + """The decode ``dot`` tiles (BM=16, BN {16, 32, 64}, BK {128, 256, 512}) crossed with a ``SPLIT_K`` axis: K split + across ``splits`` programs per output tile, reduced in-kernel. Matched memory modes, warps {2, 4}, stages 3-5. + ``configs`` instead expands caller-selected tiles, e.g. native scaled MMA at decode.""" + if configs is not None: + return [triton.Config(dict(c.kwargs, SPLIT_K=split), num_warps=c.num_warps, num_stages=c.num_stages, + num_ctas=c.num_ctas, maxnreg=c.maxnreg, pre_hook=c.pre_hook) + for c in configs for split in splits] + modes = resolve_memory_modes(memory_modes) + blocks = [ + {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "A_MEMORY_MODE": m, "B_MEMORY_MODE": m, + "SPLIT_K": sk, **fixed} + for m in modes + for bn in (16, 32, 64) + for bk in (128, 256, 512) + for sk in splits + ] + return [ + triton.Config(b, num_warps=w, num_stages=s, pre_hook=pre_hook) + for b in blocks + for w in (2, 4) + for s in (3, 4, 5) + ] + + +def narrow_dot_autotuning_configs( + memory_modes=("descriptor", "pointer"), pre_hook=None, block_ms=(16,), block_ns=(16,), **fixed, +): + """Narrow decode ``dot`` tiles, by default BM=16/BN=16: the union span's narrowest BN=32 leaves a + decode grid of ``N / 32`` programs, fewer than the SMs at N=4096 (128 on a 148-SM B200); BN=16 doubles it. BK + {128, 256, 512} (the longer bursts win where the K-loop is long), crossed with both operands' memory modes + (matched), warps {2, 4} and stages 2-6; ``fixed`` adds kwargs.""" + modes = resolve_memory_modes(memory_modes) + blocks = [ + {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "A_MEMORY_MODE": m, "B_MEMORY_MODE": m, **fixed} + for m in modes + for bm in block_ms + for bn in block_ns + for bk in (128, 256, 512) + ] + return [ + triton.Config(b, num_warps=w, num_stages=s, pre_hook=pre_hook) + for b in blocks + for w in (2, 4) + for s in (2, 3, 4, 5, 6) + ] + + @functools.cache def is_sm10x() -> bool: @@ -494,7 +603,7 @@ def is_sm10x() -> bool: driverless build boxes (``torch.cuda.is_available()`` short-circuits before the capability query, which would otherwise raise); cached (pruners call this per tune, and a process is pinned to one device).""" - return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + return torch.version.hip is None and torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 @functools.cache diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/descriptors.py b/finegrained-kernels/torch-ext/finegrained_kernels/descriptors.py index c97b73c7f..f4d88817a 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/descriptors.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/descriptors.py @@ -34,11 +34,11 @@ def rebind_operand_box(nargs, mode_key, desc_key, rows, cols): def maybe_descriptor(t: torch.Tensor, box: list[int]) -> TensorDescriptor | None: """A host-TMA descriptor over ``t``, or None when the layout cannot back one — a TMA box - needs a unit-stride innermost dim, which a transposed view (as a backward pass hands in to - contract over the other axis) does not have. ``None`` is the same signal the pointer arms - already take for an unused descriptor, and ``descriptor_box_pruner`` drops the configs that - would read it, so the tuner is left with the pointer arms the kernel serves stride-generally.""" - if t is None or (t.ndim >= 2 and t.stride(-1) != 1): + needs a 16-byte-aligned base and unit-stride innermost dim, which a transposed view (as a + backward pass hands in to contract over the other axis) does not have. ``None`` is the same + signal the pointer arms already take for an unused descriptor; ``descriptor_box_pruner`` drops + the configs that would read it, leaving the stride-general pointer arms.""" + if t is None or (t.ndim >= 2 and t.stride(-1) != 1) or (t.storage_offset() * t.element_size()) % 16: return None return TensorDescriptor.from_tensor(t, box) @@ -50,22 +50,40 @@ def rebind_bd_descriptors(nargs): rebind_operand_box(nargs, "B_MEMORY_MODE", "BDescriptor", (2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"], nargs["block_k"]) +def rebind_tiled_descriptors(nargs): + """Per-config pre_hook for the FP8 kernels with a tuned K tile (tensor-wide / row-wise): set the + A/B host-TMA boxes to ``[BM, BK]`` and ``[(2|1) * BN, BK]`` over the (rows, K) E4M3 matrices.""" + rebind_operand_box(nargs, "A_MEMORY_MODE", "ADescriptor", nargs["BLOCK_SIZE_M"], nargs["BLOCK_SIZE_K"]) + rebind_operand_box(nargs, "B_MEMORY_MODE", "BDescriptor", (2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"], nargs["BLOCK_SIZE_K"]) + + def rebind_weight_only_descriptors(nargs): """Per-config pre_hook for the weight-only kernels — set the A/B host-TMA boxes to the tuned tile: ``[BM, BK]`` over the (M, K) bf16 activation and ``[BN, BK // WEIGHT_VALUES_PER_BYTE]`` over the packed (N, K_bytes) weight (bytes: uint8 = packed E2M1, two values/byte). No scale boxes — weight-only scales are always affine (never swizzled), read through the pointer arm.""" - wvpb = 2 if nargs["B"].dtype == torch.uint8 else 1 + # uint32 integer weights: K // words values per word (2/4/8-bit) + wvpb = {torch.uint8: 2, torch.uint32: nargs["K"] // nargs["B"].shape[-1]}.get(nargs["B"].dtype, 1) bk = nargs["BLOCK_SIZE_K"] rebind_operand_box(nargs, "A_MEMORY_MODE", "ADescriptor", nargs["BLOCK_SIZE_M"], bk) rebind_operand_box(nargs, "B_MEMORY_MODE", "BDescriptor", (2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"], bk // wvpb) +def int_weight_only_pre_hook(nargs): + """Per-config pre_hook for the integer 2D kernel: the weight-only descriptor boxes, and, for a split-K config, the + zeroing of the launch's own arrival counters (only split configs read them, so the others skip the memset).""" + rebind_weight_only_descriptors(nargs) + if nargs.get("SPLIT_K", 1) > 1: + nargs["TileCounters"].zero_() + + def rebind_mx_descriptors(nargs): """Per-config pre_hook for the MX kernel — set the A/B host-TMA boxes to the tuned tile in BYTES over the (rows, K_bytes) packed matrices: ``[BM, BK // ACT_VALUES_PER_BYTE]`` and ``[BN, BK // WEIGHT_VALUES_PER_BYTE]`` (values-per-byte read off the operand dtype; uint8 = packed E2M1), plus the SWIZZLE_32_4_4 scale boxes on the offline (pre-quantized A) path.""" + if nargs.get("SPLIT_K", 1) > 1: + nargs["TileCounters"].zero_() avpb = 2 if nargs["A"].dtype == torch.uint8 else 1 wvpb = 2 if nargs["B"].dtype == torch.uint8 else 1 bk = nargs["BLOCK_SIZE_K"] @@ -113,11 +131,11 @@ def rebind_grouped_weight_descriptor(nargs): nargs["BDescriptor"], int ): return - values_per_byte = 2 if nargs["B"].dtype == torch.uint8 else 1 + values_per_element = {torch.uint8: 2, torch.uint32: nargs["K"] // nargs["B"].shape[-1]}.get(nargs["B"].dtype, 1) nargs["BDescriptor"].block_shape = [ 1, (2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"], - nargs["BLOCK_SIZE_K"] // values_per_byte, + nargs["BLOCK_SIZE_K"] // values_per_element, ] diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/epilogue.py b/finegrained-kernels/torch-ext/finegrained_kernels/epilogue.py index f4efb8142..983be891d 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/epilogue.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/epilogue.py @@ -100,21 +100,23 @@ def acc_init( BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, SWAP_AB: tl.constexpr, + DTYPE: tl.constexpr = tl.float32, ): """Zero accumulator shaped for the layout: swapped scalar reduces into ``[1, N]``; any other swapped mode keeps the ``[N, M]`` MMA acc (weight rows in M, act tile in N — padded up to the ``MMA_N_ATOM`` when the caller is the single-token decode GEVM, col 0 taken after the - K-loop); no-swap uses ``[M, N]``. ``COMPUTE_MODE`` matters only under swap — kernels with no + K-loop); no-swap uses ``[M, N]``. INT8 SmoothQuant selects INT32; other callers keep FP32. + ``COMPUTE_MODE`` matters only under swap — kernels with no mode axis (fp8 ``tl.dot``) pass ``"dot"``. ``N`` is the weight-output tile (``BLOCK_SIZE_H`` for the fp8 down projection). Single return: only the taken branch compiles.""" if SWAP_AB and COMPUTE_MODE == "scalar": - acc = tl.zeros((1, BLOCK_SIZE_N), dtype=tl.float32) + acc = tl.zeros((1, BLOCK_SIZE_N), dtype=DTYPE) elif SWAP_AB and BLOCK_SIZE_M < MMA_N_ATOM: - acc = tl.zeros((BLOCK_SIZE_N, MMA_N_ATOM), dtype=tl.float32) + acc = tl.zeros((BLOCK_SIZE_N, MMA_N_ATOM), dtype=DTYPE) elif SWAP_AB: - acc = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32) + acc = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=DTYPE) else: - acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=DTYPE) return acc @@ -374,8 +376,10 @@ def fused_glu( dtype=out_dtype or gate_up.dtype, device=gate_up.device) n = out.numel() with device_context(gate_up.device): + # Explicit defaults keep torch.compile's Triton argument analysis consistent. compile_time_only_triton_wrap(_glu_kernel)[(triton.cdiv(n, 1024),)]( gate_up, out, None, n, inter_dim, act_fn, swiglu_alpha, swiglu_limit, BLOCK=1024, + QUANT_GROUP=128, UE8M0=False, **pdl_launch_kwargs(), ) return out @@ -427,6 +431,9 @@ def split_gate_up_glu( """Gate|up epilogue in one step: split the finalized stacked tile into its (gate, up) pair (``split_gate_up``) and apply the ``ACT_FN``/SwiGLU gated linear unit (``glu``), returning the combined intermediate. See those two for the orientation and activation details.""" + # Materialize before the layout split so it can move narrow values through shared memory. + if SIMULATE_UNFUSED: + flat = flat.to(INTERMEDIATE_DTYPE) gate, up = split_gate_up(flat, BLOCK_SIZE_M, BLOCK_SIZE_N, SWAP_AB) return glu( gate, @@ -532,6 +539,46 @@ def _epilogue_requant_mx( ) +@triton.jit +def add_low_rank( + acc, LoraInput, LoraWeight, out_row, row_mask, pid_n, + stride_li_m, stride_li_r, stride_lw_n, stride_lw_r, + BLOCK_SIZE_M: tl.constexpr, WIDTH: tl.constexpr, RANK: tl.constexpr, SPLIT_K: tl.constexpr, + GATE: tl.constexpr, N_COLS: tl.constexpr, SIMULATE_UNFUSED: tl.constexpr, + INTERMEDIATE_DTYPE: tl.constexpr, +): + """SVDQuant's BF16 correction, after dequantization/bias and before GLU/requant.""" + if LoraInput is not None: + cols = pid_n * WIDTH + tl.arange(0, WIDTH) + col_mask = cols < (2 * N_COLS if GATE else N_COLS) + rank = tl.arange(0, triton.next_power_of_2(RANK)) + w = tl.load(LoraWeight + cols[:, None] * stride_lw_n + rank[None, :] * stride_lw_r, + mask=col_mask[:, None] & (rank[None, :] < RANK), other=0) + if SPLIT_K > 1: + # Decode has one row: finish the split reduction here, without a launch. + splits = tl.arange(0, triton.next_power_of_2(SPLIT_K)) + parts = tl.load(LoraInput + splits[:, None] * RANK + rank[None, :], + mask=(splits[:, None] < SPLIT_K) & (rank[None, :] < RANK), other=0) + x = tl.sum(parts, 0).to(LoraWeight.dtype.element_ty) + correction = tl.sum(w.to(tl.float32) * x.to(tl.float32)[None, :], 1)[None, :] + elif BLOCK_SIZE_M < 16: + x = tl.load(LoraInput + out_row[:, None] * stride_li_m + rank[None, :] * stride_li_r, + mask=row_mask[:, None] & (rank[None, :] < RANK), other=0) + correction = tl.sum(x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :], 2) + else: + x = tl.load(LoraInput + out_row[:, None] * stride_li_m + rank[None, :] * stride_li_r, + mask=row_mask[:, None] & (rank[None, :] < RANK), other=0) + correction = tl.dot(x, tl.trans(w)) + if SIMULATE_UNFUSED: + # ModelOpt materializes both linear outputs before their BF16 addition. + acc = acc.to(INTERMEDIATE_DTYPE).to(tl.float32) + correction = correction.to(INTERMEDIATE_DTYPE).to(tl.float32) + acc = acc + correction + if SIMULATE_UNFUSED: + acc = acc.to(INTERMEDIATE_DTYPE).to(tl.float32) + return acc + + @triton.jit def gemm_epilogue( C, @@ -563,7 +610,7 @@ def gemm_epilogue( CSDescriptor=0, # SWIZZLE_32_4_4 requant-scale descriptor; read only under SWIZZLED_OUT (else dummy) CsGlobal=None, # fp32 NVFP4 output global (the NEXT proj's provided input_scale), per tensor or per expert; normalizes the requant, None folds out stride_cs_global=0, # expert stride of CsGlobal (0 broadcasts a scalar) - GlobalScale=None, # fp32 NVFP4 weight global g_b, per tensor or per expert; applied PRE-GLU, None folds out + GlobalScale=None, # weight global or per-channel pointer tile; applied before bias/GLU, None folds out stride_global_e=0, # expert stride of GlobalScale (global_scale_stride; 0 broadcasts a scalar) global_row=0, # expert id indexing the globals and the bias (0 for the 2D case) GlobalScaleA=None, # fp32 NVFP4 activation global g_a, per tensor or per expert; multiplied in-register with GlobalScale @@ -574,6 +621,14 @@ def gemm_epilogue( Bias=None, # (E, N_out) per-expert output bias, N_out = 2N under GATE; row is `global_row` stride_bias_e=0, stride_bias_n=0, + LoraInput=None, + LoraWeight=None, + stride_li_m=0, + stride_li_r=0, + stride_lw_n=0, + stride_lw_r=0, + LORA_RANK: tl.constexpr = 0, + LORA_SPLIT_K: tl.constexpr = 1, ): """Unified output epilogue for grouped (a real scatter tile) and batched (fake-batch decode: one token replicated across the BM lanes) GEMMs. Plain: cast + store the accumulator. ``GATE``: @@ -599,6 +654,12 @@ def gemm_epilogue( Bias, stride_bias_e, stride_bias_n, global_row, pid_n, 2 * BLOCK_SIZE_N, 2 * N_COLS if N_COLS > 0 else 0, ) + flat = add_low_rank( + flat, LoraInput, LoraWeight, out_row, row_mask, pid_n, + stride_li_m, stride_li_r, stride_lw_n, stride_lw_r, + BLOCK_SIZE_M, 2 * BLOCK_SIZE_N, LORA_RANK, LORA_SPLIT_K, True, N_COLS, + SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, + ) # SAVE_PREACT: the GLU backward needs Z, the 2*BN-wide pre-activation. Interleaved, that # tile IS an ungated output tile at doubled width (pid_n * 2BN == 2 * pid_n * BN), so the # ordinary store serves it — no separate path, no gated store arm. @@ -634,4 +695,10 @@ def gemm_epilogue( ), Bias, stride_bias_e, stride_bias_n, global_row, pid_n, BLOCK_SIZE_N, N_COLS, ) + acc = add_low_rank( + acc, LoraInput, LoraWeight, out_row, row_mask, pid_n, + stride_li_m, stride_li_r, stride_lw_n, stride_lw_r, + BLOCK_SIZE_M, BLOCK_SIZE_N, LORA_RANK, LORA_SPLIT_K, False, N_COLS, + SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, + ) _store_out(C, acc, out_row, pid_n, row_mask, stride_c_m, stride_c_n, BLOCK_SIZE_M, BLOCK_SIZE_N, FAKE_BATCH, N_COLS) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/formats.py b/finegrained-kernels/torch-ext/finegrained_kernels/formats.py index 5f4db7f83..b7b779f5e 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/formats.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/formats.py @@ -12,11 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math + import torch import triton from .compat import MX_SCALE_GROUP_K, NIBBLES_PER_BYTE, NVFP4_SCALE_GROUP_K +# Amortize weight conversion across this many activation rows per GEMM. +FP8_WEIGHT_ONLY_DEQUANT_MIN_M = 1024 + def ue8m0_as_uint8(scale: torch.Tensor | None) -> torch.Tensor | None: """View UE8M0 (``float8_e8m0fnu``) weight scales as ``uint8`` for the Triton binder, which doesn't recognize the dtype; kernels decode ``2^(exp-127)`` @@ -75,15 +80,15 @@ def is_mxfp4(weight: torch.Tensor, scale: torch.Tensor) -> bool: def is_nvfp4(weight: torch.Tensor, scale: torch.Tensor) -> bool: """NVFP4 weight/scale pair: packed E2M1 weights (``int8``, two codes/byte) with E4M3 - group-16 block scales — the scale DTYPE is the format carrier (UE8M0 = MX, E4M3 = NV), and - the group falls out of the shape. This predicate reads the block scale; the per-tensor + group-16 block scales, or ModelOpt's group-32 FP4/FP8 recipe. The dtype identifies the scale + encoding, and the group falls out of the shape. This predicate reads the block scale; the per-tensor second-level global is a separate ``b_global_scale`` argument (``nvfp4_quantize_two_level``). Layout-agnostic questions ("is this NVFP4, swizzled or not?") go through ``weight_format``, which carries the same dtype rule without the shape check.""" return ( weight.dtype == torch.int8 and scale.dtype == torch.float8_e4m3fn - and _shapes_match(weight, scale, NVFP4_SCALE_GROUP_K) + and (_shapes_match(weight, scale, NVFP4_SCALE_GROUP_K) or _shapes_match(weight, scale, MX_SCALE_GROUP_K)) ) def is_preswizzled_mx(weight: torch.Tensor, scale: torch.Tensor) -> bool: @@ -121,14 +126,34 @@ def is_tensor_wide(block_size, weight: torch.Tensor) -> bool: block_size[0] == weight.shape[-2] and block_size[1] == weight.shape[-1] ) +def is_rowwise(weight: torch.Tensor, scale: torch.Tensor) -> bool: + """Row-wise FP8 weight/scale pair: E4M3 weights with one fp32 scale per output row — + ``scale`` ``(N, 1)`` for a 2D ``(N, K)`` weight, ``(E, N, 1)`` for a 3D stack (e.g. Llama 4 + FP8). The block is ``(1, K)``: the scale is constant along K, so it folds once after the + K-loop, like a per-tensor one.""" + if weight.dtype != torch.float8_e4m3fn or scale.dtype != torch.float32: + return False + if scale.shape != weight.shape[:-1] + (1,): + return False + return True + +def per_token_activations(B: torch.Tensor, Bs: torch.Tensor | None) -> bool: + """Whether the weights take one activation scale per token: FP8 with a K-constant scale + (tensor-wide or row-wise). The scale spans the whole row, so a tile-local epilogue cannot + form it — such a GEMM never fuses the requant of its output.""" + if Bs is None or is_mx(B, Bs): + return False + return Bs.numel() == (B.shape[0] if B.ndim == 3 else 1) or is_rowwise(B, Bs) + def weight_block_size(B: torch.Tensor, Bs: torch.Tensor) -> list[int] | None: """The fp8 weight-quantization block ``[block_n, block_k]``, derived from the scale tensor's shape — the data already says how it was quantized, so no ``block_size`` - parameter exists to disagree with it. ``None`` = tensor-wide (one scalar per expert: - ``Bs`` ``(E,)`` or ``(E, 1, 1)`` spanning the full ``(N, K)``). Expects a 3D + parameter exists to disagree with it. ``None`` = a scale constant along K, folded after + the K-loop by the tensor-wide kernels: one scalar per expert (``Bs`` ``(E,)`` or + ``(E, 1, 1)``) or one per output row (``(E, N, 1)``, ``is_rowwise``). Expects a 3D ``(E, N, K)`` fp8 ``B``; MX weights never reach here (format keyed by scale dtype).""" num_experts, n_rows, K = B.shape - if Bs.numel() == num_experts: + if Bs.numel() == num_experts or is_rowwise(B, Bs): return None assert Bs.ndim == 3 and K % Bs.shape[2] == 0, ( f"Bs shape {tuple(Bs.shape)} does not tile B {tuple(B.shape)} along K" @@ -203,7 +228,7 @@ def tokens_per_expert_bucket(S: int, num_experts: int) -> int: def normalize_global_scale( g: torch.Tensor | None, num_experts: int ) -> torch.Tensor | None: - """An NVFP4 second-level global in one of the kernels' fp32 load layouts: ``(1,)`` per tensor + """A calibrated/global scale in the kernels' fp32 load layouts: ``(1,)`` per tensor or ``(num_experts,)`` per expert (``(E, 1, 1)`` accepted for the latter). The kernels index it through ``global_scale_stride``, so a shared value broadcasts by stride and nothing is materialized per call. ``None`` passes through. @@ -233,8 +258,12 @@ def global_scale_stride(g: torch.Tensor | None) -> int: return 0 if g is None or g.numel() == 1 else 1 def normalize_per_expert_scale(Bs: torch.Tensor, num_experts: int) -> torch.Tensor: - """One per-tensor scale per expert, normalized to ``(num_experts, 1, 1)`` from - either that or a bare ``(num_experts,)``.""" + """A K-constant weight scale in the tensor-wide kernels' ``(num_experts, rows, 1)`` layout: + one per-tensor scale per expert (``rows`` 1, from that or a bare ``(num_experts,)``), or a + row-wise ``(num_experts, N, 1)`` passed through. The kernels read row ``n`` at + ``n * stride(1)``, so ``row_scale_stride`` broadcasts the per-tensor one.""" + if Bs.ndim == 3 and Bs.shape[0] == num_experts and Bs.shape[1] > 1 and Bs.shape[2] == 1: + return Bs if Bs.ndim == 1: assert Bs.shape[0] == num_experts, ( f"Bs shape {tuple(Bs.shape)} != expected ({num_experts},)" @@ -245,19 +274,23 @@ def normalize_per_expert_scale(Bs: torch.Tensor, num_experts: int) -> torch.Tens ) return Bs +def row_scale_stride(Bs: torch.Tensor) -> int: + """The output-row stride of a K-constant weight scale (``normalize_per_expert_scale``, or a + 2D ``(N, 1)`` / scalar one): 0 broadcasts a per-tensor scale to every row.""" + return 0 if Bs.ndim < 2 or Bs.shape[-2] == 1 else Bs.stride(-2) + def mx_scale_family(Bs: torch.Tensor, K: int) -> int: """The group size of an MX/NV weight-scale tensor, in either layout — the wrapper hands ``Bs`` as-is and this reads the group off its shape. Row-major (2D/3D): ``K // Bs.shape[-1]``. SWIZZLE_32_4_4 (5D ``(G, row_blocks, cols // 4, 2, 256)``): ``K // (Bs.shape[2] * 4)``. The scale - dtype IS the format carrier (E4M3 = NVFP4 group-16, UE8M0 = MX group-32) and the pairing is - validated; callers that need the format read it off ``Bs.dtype`` (``== torch.float8_e4m3fn`` - is NVFP4).""" + dtype identifies the scale encoding (E4M3 for NVFP4 group-16 or FP4/FP8 group-32; UE8M0 for + MX group-32), while shape identifies its group width. Both are validated.""" nvfp4 = Bs.dtype == torch.float8_e4m3fn assert nvfp4 or Bs.dtype in UE8M0_SCALE_DTYPES, ( f"Bs must be UE8M0 (float8_e8m0fnu/uint8) or E4M3 (NVFP4), got {Bs.dtype}" ) scale_group = K // (Bs.shape[2] * 4) if Bs.ndim == 5 else K // Bs.shape[-1] - assert scale_group == (NVFP4_SCALE_GROUP_K if nvfp4 else MX_SCALE_GROUP_K), ( + assert scale_group in ((NVFP4_SCALE_GROUP_K, MX_SCALE_GROUP_K) if nvfp4 else (MX_SCALE_GROUP_K,)), ( f"scale group {scale_group} does not match the scale dtype {Bs.dtype}" ) assert K % scale_group == 0, f"K (={K}) must be a multiple of {scale_group}" @@ -269,7 +302,7 @@ def get_supported_act_fns() -> tuple[str, ...]: return ("silu", "gelu", "relu") -ACTIVATION_FORMATS = (None, "bf16", "fp8", "mxfp8", "mxfp4", "nvfp4") +ACTIVATION_FORMATS = (None, "bf16", "int8", "fp8", "mxfp8", "mxfp4", "nvfp4") def check_activation_format(activation_format: str | None, quantize_output: bool) -> None: @@ -295,13 +328,15 @@ def check_activation_format(activation_format: str | None, quantize_output: bool ================== =============================================================== block-dynamic FP8 None / "fp8" (E4M3 + per-block scales); quantize_output ok tensor-wide FP8 None / "fp8" (E4M3 + per-token scales); no quantize_output - MXFP8 / MXFP4 None / "mxfp8" (E4M3 + UE8M0) / "mxfp4" (packed E2M1 + UE8M0) / "bf16" - NVFP4 None / "nvfp4" (packed E2M1 + E4M3 group-16) / "bf16" + row-wise FP8 None / "fp8" (E4M3 + per-token scales); no quantize_output + MXFP8 / MXFP4 None / "mxfp8" / "mxfp4" / "bf16"; MXFP4 also calibrated "fp8" + NVFP4 None / "nvfp4" (group-16) / "bf16"; group-32 requires calibrated "fp8" + integer words None / "bf16" (weight-only), calibrated "fp8" (INT4) or "int8" (INT8) full-precision None (nothing to quantize) ================== =============================================================== - (Tensor-wide's whole-row activation scale can't be formed by a tile-local epilogue; the unfused - path quantizes on the host between GEMMs.) A format the weight family doesn't support fails + (A per-token activation scale spans the whole row, which a tile-local epilogue can't form; the + next GEMM quantizes the intermediate instead.) A format the weight family doesn't support fails loudly at the op.""" assert activation_format in ACTIVATION_FORMATS, ( f"unknown activation_format {activation_format!r}; expected one of {ACTIVATION_FORMATS}" @@ -311,10 +346,136 @@ def check_activation_format(activation_format: str | None, quantize_output: bool ) +def _validate_quantization_args( + A, As, B, Bs, a_global_scale, b_global_scale, output_global_scale, + b_zero, b_perm, a_pre_quant_scale, activation_format, quantize_output, activation_scale_ub, +): + """Shared recipe and metadata checks for the dense and routed dispatchers.""" + assert a_pre_quant_scale is None or B.dtype in INT_WEIGHT_DTYPES or ( + Bs is not None and is_mx(B, Bs) and activation_format != "bf16" + ), "a_pre_quant_scale requires integer weights or an MX/FP4 activation-quantized path" + check_activation_format(activation_format, quantize_output) + check_activation_scale_ub(A, As, B, Bs, activation_format, activation_scale_ub) + assert b_zero is None or B.dtype in INT_WEIGHT_DTYPES, "zero points (b_zero) are the integer weights'" + assert b_perm is None or B.dtype in INT_WEIGHT_DTYPES, "an act-order K order (b_perm) is the integer weights'" + assert (a_global_scale is None and b_global_scale is None) or ( + Bs is not None and weight_format(B, Bs) == "nvfp4" + ) or ( + B.dtype in INT_WEIGHT_DTYPES and activation_format == "fp8" and a_global_scale is None + ), "global scales require NVFP4, or integer W4A8 weight_scale_2" + assert output_global_scale is None or ( + Bs is not None and quantize_output and resolve_activation_format(activation_format, B, Bs) == "nvfp4" + ), ( + "output_global_scale is the NVFP4 requant second level — it requires quantize_output=True on NVFP4 " + "(the epilogue would otherwise normalize by it with nothing downstream to compensate)" + ) + if B.dtype in INT_WEIGHT_DTYPES: + assert not quantize_output and ( + (activation_format in (None, "bf16") and As is None and b_global_scale is None) + or (activation_format == "fp8" and As is not None and b_global_scale is not None) + or (activation_format == "int8" and As is not None and b_global_scale is None) + ), "integer GEMMs support weight-only, calibrated INT4/FP8 AWQ, or calibrated INT8 SmoothQuant, without requant" + + +def check_activation_scale_ub(A, As, B, Bs, activation_format, activation_scale_ub): + """An amax cap controls dynamic per-token FP8 quantization on row-wise weights.""" + if activation_scale_ub is None: + return + if not math.isfinite(activation_scale_ub) or activation_scale_ub <= 0: + raise ValueError("activation_scale_ub must be finite and positive") + if Bs is None or not is_rowwise(B, Bs) or activation_format not in (None, "fp8"): + raise ValueError("activation_scale_ub requires row-wise FP8 weights and FP8 activations") + if As is not None or A.dtype == torch.float8_e4m3fn: + raise ValueError("activation_scale_ub requires raw activations without a calibrated activation scale") + + +INT_WEIGHT_BITS = (2, 4, 8) +# uint32: the kernels' interleaved layout; int32 routes to the integer ops only to be refused there (int_weight_bits) +INT_WEIGHT_DTYPES = (torch.uint32, torch.int32) + + +def int_weight_bits(K: int, B: torch.Tensor) -> int: + """The value width of an integer weight ``B``: its ``K`` values per row ride in ``B.shape[-1]`` uint32 words, + ``32 // bits`` each. The width is a fact of the shapes, so the data cannot disagree with a parameter. The dtype + is the layout: uint32 words are the kernels' interleaved order (``interleave_int_words``, the GPTQ / AWQ + repacks), and int32 words, a checkpoint's natural packing, are refused rather than read in the wrong order.""" + assert B.dtype == torch.uint32, ( + f"integer weights are uint32 words in the kernels' interleaved layout, got {B.dtype}: convert natural-order " + "int32 packing with interleave_int_words (GPTQ / AWQ: repack_gptq_int4 / repack_awq_int4)" + ) + words = B.shape[-1] + bits = 32 * words // K + assert bits in INT_WEIGHT_BITS and bits * K == 32 * words, ( + f"{words} uint32 words per row cannot hold K={K} values of 2, 4 or 8 bits" + ) + return bits + + +def validate_int_operands(A, B, Bs, b_zero, b_perm, a_scale, b_scale2) -> int: + """Shared dense/routed integer layout and calibrated activation recipe checks.""" + assert A.dtype in (torch.bfloat16, torch.float16), f"integer GEMMs take raw bf16/fp16 activations, got {A.dtype}" + K = A.shape[-1] + bits = int_weight_bits(K, B) + assert Bs.ndim == B.ndim and Bs.shape[:-1] == B.shape[:-1] and Bs.shape[-1] > 0 and K % Bs.shape[-1] == 0, ( + "weight scales must match the weight rows and divide K into whole groups" + ) + assert b_zero is None or b_zero.shape == Bs.shape, "zero points are shaped like the scales" + assert b_perm is None or (b_perm.shape == (*B.shape[:-2], K) and b_perm.dtype == torch.int32), ( + "act-order K orders must be int32 with one K-vector per weight matrix" + ) + assert b_scale2 is None or a_scale is not None, "weight_scale_2 needs a calibrated input_scale" + if a_scale is not None: + smoothquant = b_scale2 is None + assert bits == (8 if smoothquant else 4) and b_zero is None and b_perm is None, ( + "ModelOpt activation-quantized integer recipes are symmetric and have no act-order" + ) + if smoothquant: + assert Bs.shape[-1] == 1 and K * 128 * 128 < 2**31, ( + "INT8 SmoothQuant needs per-channel scales and a K that cannot overflow INT32" + ) + return bits + + +def validate_svdquant_operands(A, B, As, Bs, activation_format, output_dtype, lora_a, lora_b, pre_quant_scale): + if lora_a is None and lora_b is None: + return + assert lora_a is not None and lora_b is not None, "SVDQuant requires both low-rank factors" + assert ( + B.ndim == 2 and B.dtype == torch.int8 and Bs is not None + and weight_format(B, Bs) == "nvfp4" + and mx_scale_family(Bs, A.shape[-1]) == NVFP4_SCALE_GROUP_K + and activation_format in (None, "nvfp4") + ), "SVDQuant requires NVFP4 W4A4 weights" + assert As is None and A.dtype == torch.bfloat16 and output_dtype in (None, torch.bfloat16), ( + "SVDQuant requires raw BF16 activations and BF16 output" + ) + assert A.shape[-1] % 64 == 0, "SVDQuant's fused input preparation requires K divisible by 64" + assert ( + lora_a.ndim == lora_b.ndim == 2 and lora_a.shape[1] == A.shape[-1] + and lora_b.shape == (B.shape[0], lora_a.shape[0]) + ), "SVDQuant factors must have shapes (rank, K) and (N, rank)" + if pre_quant_scale is not None: + assert pre_quant_scale.shape == (A.shape[-1],) and pre_quant_scale.is_contiguous(), ( + "pre_quant_scale must be contiguous (K,)" + ) + assert pre_quant_scale.dtype == A.dtype and pre_quant_scale.device == A.device, ( + "pre_quant_scale must match the BF16 activation device and dtype" + ) + assert not torch.is_grad_enabled() or not any(t.requires_grad for t in (A, lora_a, lora_b)), ( + "SVDQuant currently supports inference only" + ) + assert lora_a.shape[0] > 0 and lora_a.shape[0] % 32 == 0, "SVDQuant rank must be a positive multiple of 32" + assert lora_a.dtype == lora_b.dtype == A.dtype and lora_a.device == lora_b.device == A.device, ( + "SVDQuant factors must match the BF16 activation device and dtype" + ) + + def weight_format(B: torch.Tensor, Bs: torch.Tensor) -> str: """The format the weights are ALREADY in, read off their dtypes — what ``activation_format=None`` resolves to. E4M3 scales -> ``"nvfp4"``; else packed E2M1 -> ``"mxfp4"``, MX E4M3 -> ``"mxfp8"``, - block/tensor-scaled E4M3 -> ``"fp8"``.""" + block/tensor-scaled E4M3 -> ``"fp8"``; int32 / uint32 words of integer values -> ``"int"`` (2/4/8-bit, weight-only).""" + if B.dtype in INT_WEIGHT_DTYPES: + return "int" if Bs.dtype == torch.float8_e4m3fn: return "nvfp4" if not is_mx(B, Bs): @@ -328,12 +489,22 @@ def weight_format(B: torch.Tensor, Bs: torch.Tensor) -> str: def resolve_activation_format(activation_format: str | None, B: torch.Tensor, Bs: torch.Tensor) -> str: """The format an MX/NVFP4 GEMM quantizes its activations to. ``None`` resolves to the weights' - own format (``weight_format``): NVFP4 weights pin ``"nvfp4"`` (the MMA kind needs matching + own format (``weight_format``): group-16 NVFP4 weights pin ``"nvfp4"`` (the MMA kind needs matching scale formats on both operands), packed E2M1 -> ``"mxfp4"`` (the all-fp4 W4A4 chain), else ``"mxfp8"``. An explicit name overrides (e.g. ``"mxfp8"`` on MXFP4 weights = the W4A8 chain) - and is validated against the family. ``"bf16"`` never reaches here — a weight-only chain is + and is validated against the family. Packed group-32 weights also support calibrated ordinary + ``"fp8"`` activations. ``"bf16"`` never reaches here — a weight-only chain is routed to its own kernel by the dispatcher.""" + if activation_format == "fp8": + assert B.dtype == torch.int8, "ordinary FP8 activations on MX weights require packed FP4 weights" + assert mx_scale_family(Bs, B.shape[-1] * 2) == MX_SCALE_GROUP_K, ( + "ModelOpt FP4/FP8 recipes use group-32 weight scales" + ) + return "fp8" if Bs.dtype == torch.float8_e4m3fn: + assert mx_scale_family(Bs, B.shape[-1] * 2) == NVFP4_SCALE_GROUP_K, ( + "E4M3 group-32 FP4 weights require activation_format='fp8'" + ) assert activation_format in (None, "nvfp4"), ( f"NVFP4 activations are packed E2M1 + E4M3 scales, got {activation_format!r}" ) @@ -355,3 +526,27 @@ def resolve_output_dtype( if output_dtype is not None: return output_dtype return activation.dtype if act_scale is None else torch.bfloat16 + + +def block_fp8_weight_scale_grid(B: torch.Tensor, Bs: torch.Tensor, interleaved_scales: bool = False): + """Normalize scalar/block FP8 scales and infer their block size from the grid. + + A partial square block on N retains the size implied by K. Separately + calibrated gate/up scale rows alternate per block under interleaved_scales. + """ + assert B.dtype == torch.float8_e4m3fn and B.ndim in (2, 3) + experts = B.shape[0] if B.ndim == 3 else 1 + rows, K = B.shape[-2:] + if Bs.numel() == experts: + return Bs.reshape(*B.shape[:-2], 1, 1), rows, K + assert Bs.ndim == B.ndim and Bs.shape[:-2] == B.shape[:-2] + assert K % Bs.shape[-1] == 0, "the weight scale grid must tile K evenly" + block_k = K // Bs.shape[-1] + if interleaved_scales and Bs.shape[-2] > 1: + assert rows % 2 == 0 and Bs.shape[-2] % 2 == 0, "gate/up scale rows must occur in pairs" + scale_rows, weight_rows = Bs.shape[-2] // 2, rows // 2 + else: + scale_rows, weight_rows = Bs.shape[-2], rows + block_n = block_k if ((weight_rows + block_k - 1) // block_k) == scale_rows else weight_rows // scale_rows + assert block_n > 0 and ((weight_rows + block_n - 1) // block_n) == scale_rows + return Bs, block_n, block_k diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py b/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py index 55cf77d6c..3acf836d7 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py @@ -13,6 +13,8 @@ # limitations under the License. +from typing import List + import torch import triton import triton.language as tl @@ -22,19 +24,22 @@ from .bayesian_autotuner import bayesian_autotune from .bench import weight_input -from .compat import add_op_namespace_prefix, FP8_DTYPE, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, persistent_program_count, tl_dtype +from .compat import add_op_namespace_prefix, FP8_DTYPE, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, persistent_program_count, tl_dtype, MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, pipeline_autotuning_configs, persistent_programs_per_sm, is_sm10x from .descriptors import build_grouped_operand_descriptors, rebind_grouped_descriptors, rebind_grouped_mx_descriptors -from .formats import check_activation_format, global_scale_stride, is_per_expert_global, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, routed_rows, tokens_per_expert_bucket, ue8m0_as_uint8, validate_dense_operands, weight_block_size, weight_format -from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, mx_act_quant_grouped, mx_act_quant_routed_rows, quantize_routed_rows_per_expert, static_expert_act_operands, swizzle_grouped_mx_scales, tensor_wide_act_operands +from .formats import FP8_WEIGHT_ONLY_DEQUANT_MIN_M, block_fp8_weight_scale_grid, _validate_quantization_args, global_scale_stride, is_per_expert_global, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, routed_rows, row_scale_stride, tokens_per_expert_bucket, ue8m0_as_uint8, validate_dense_operands, weight_block_size, validate_int_operands, INT_WEIGHT_DTYPES +from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, mx_act_quant_grouped, mx_act_quant_routed_rows, quantize_routed_rows_per_expert, static_expert_act_operands, swizzle_grouped_mx_scales, tensor_wide_act_operands, prepare_fp8_mx_inputs, permute_routed_rows from .swizzle import swizzled_scale_descriptor -from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot -from .scheduling import build_tile_layout, expand_gather_below_parity, expand_regime, build_packed_schedule, load_packed_schedule, prefetch_packed_entry, resolve_grouped_tile, resolve_grouped_tile_packed +from .loading.scales import apply_row_scale, dequant_block_fp8_weight +from .mma import block_dynamic_dot, fp8_dot, plain_dot, mx_compute, mx_weight_only_compute, static_dot, int_weight_only_compute +from .scheduling import build_tile_layout, expand_gather_below_parity, expand_regime, load_packed_schedule, grouped_tile_order, prefetch_packed_entry, resolve_grouped_tile, resolve_grouped_tile_packed, fill_packed_schedule, packed_schedule_buffer from .loading.tiles import ( load_act_block_dynamic, load_act_mx, load_act_plain, load_act_static, load_weight_block_dynamic, + load_weight_block_fp8, + load_weight_int_group, load_weight_mx, load_weight_plain, load_weight_static, @@ -43,7 +48,32 @@ weight_tile_ptrs, ) from .epilogue import acc_init, bias_strides, gemm_epilogue -from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, raw_activation_pointer_pruner, global_scale_warp_spec_pruner, warp_spec_memory_mode_pruner, packed_schedule_scope_pruner, affine_scale_warp_spec_pruner, block_dynamic_grouped_matmul_pruner, block_fits_dim_pruner, block_within_dim_pruner, compose_pruners, descriptor_box_pruner, gate_stacked_tmem_trap_pruner, gated_pointer_weight_warp_spec_pruner, mx_config_pruner, require_moe_dims_aligned, smem_pruner, swizzled_out_bm_pruner, swizzled_scale_config_pruner, swizzled_scales_bm_pruner, warp_spec_compile_guard_pruner +from .pruners import PATH_ANCHOR_AXES, int_scale_group_tile_pruner, fp8_dot_warp_pruner, raw_activation_pointer_pruner, warp_spec_memory_mode_pruner, packed_schedule_scope_pruner, affine_scale_warp_spec_pruner, block_dynamic_grouped_matmul_pruner, block_fits_dim_pruner, block_within_dim_pruner, compose_pruners, descriptor_box_pruner, gate_stacked_tmem_trap_pruner, gated_pointer_weight_warp_spec_pruner, mx_config_pruner, require_moe_dims_aligned, smem_pruner, swizzled_out_bm_pruner, swizzled_scale_config_pruner, swizzled_scales_bm_pruner, warp_spec_compile_guard_pruner + + +def grouped_pre_hook(nargs): + """Per-config pre_hook of the grouped kernels: the operand descriptor boxes, and the packed schedule for the + configs that read it.""" + rebind_grouped_descriptors(nargs) + fill_packed_schedule(nargs) + + +def grouped_mx_pre_hook(nargs): + """``grouped_pre_hook`` plus the MX scale boxes.""" + rebind_grouped_mx_descriptors(nargs) + fill_packed_schedule(nargs) + + +def _grouped_dot_configs(): + return get_accelerator_autotuning_configs( + packed_schedule=True, + tune_block_nk=True, + warp_spec=True, + tune_block_m=True, + a_memory_modes=("descriptor", "pointer"), + b_memory_modes=("descriptor", "pointer"), + pre_hook=grouped_pre_hook, + ) @bayesian_autotune( @@ -57,7 +87,7 @@ tune_block_m=True, a_memory_modes=("descriptor", "pointer"), b_memory_modes=("descriptor", "pointer"), - pre_hook=rebind_grouped_descriptors, + pre_hook=grouped_pre_hook, ), # GATE keys the gate|up arm separately: its dot is 2*BN wide, a different tile optimum. # BLOCK_SIZE_N/BLOCK_SIZE_K here are the launch-pinned quant block (not tuned) — they gate @@ -91,7 +121,7 @@ def w8a8_block_dynamic_fp8_matmul_grouped_kernel( GatherIdx, # (S,) int32 — sorted position -> source row of A; read only when not None ScatterIdx, # (S,) int32 — sorted position -> destination row of C; read only when not None ExpertStart, # (NUM_EXPERTS_POW2 + 1,) int32 — cumulative row starts, S sentinel - Schedule, # packed tile schedule (build_packed_schedule); read iff PACKED_SCHEDULE + Schedule, # packed tile schedule (fill_packed_schedule); read iff PACKED_SCHEDULE # Shape S, N, @@ -265,7 +295,7 @@ def w8a8_block_dynamic_fp8_matmul_grouped_kernel( tune_block_m=True, a_memory_modes=("descriptor", "pointer"), b_memory_modes=("descriptor", "pointer"), - pre_hook=rebind_grouped_descriptors, + pre_hook=grouped_pre_hook, ), ["N", "K", "tokens_per_expert_bit_length", "GATE", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], n_trials=100, @@ -297,7 +327,7 @@ def w8a8_block_static_fp8_matmul_grouped_kernel( GatherIdx, # (S,) int32 — sorted position -> source row of A; read only when not None ScatterIdx, # (S,) int32 — sorted position -> destination row of C; read only when not None ExpertStart, # (NUM_EXPERTS_POW2 + 1,) int32 — cumulative row starts, S sentinel - Schedule, # packed tile schedule (build_packed_schedule); read iff PACKED_SCHEDULE + Schedule, # packed tile schedule (fill_packed_schedule); read iff PACKED_SCHEDULE # Shape S, N, @@ -448,15 +478,7 @@ def w8a8_block_static_fp8_matmul_grouped_kernel( @bayesian_autotune( - get_accelerator_autotuning_configs( - packed_schedule=True, - tune_block_nk=True, - warp_spec=True, - tune_block_m=True, - a_memory_modes=("descriptor", "pointer"), - b_memory_modes=("descriptor", "pointer"), - pre_hook=rebind_grouped_descriptors, - ), + _grouped_dot_configs(), # GATE keys the gate|up arm separately (its dot is 2*BN wide, a different tile optimum). ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, @@ -474,6 +496,7 @@ def w8a8_block_static_fp8_matmul_grouped_kernel( block_within_dim_pruner("K"), block_within_dim_pruner("N", "BLOCK_SIZE_N"), warp_spec_compile_guard_pruner(), + warp_spec_memory_mode_pruner(), descriptor_box_pruner(), smem_pruner(), raw_activation_pointer_pruner(), @@ -487,14 +510,14 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( As, # per-token activation scales (S,), or a calibrated (static) one: shared, or per expert B, # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack BDescriptor, # host TMA descriptor over B viewed (E, 2N|N, K), box (1, (2|1)*BN, BK); read iff B_MEMORY_MODE != "pointer" - Bs, # (num_experts, 1, 1) per-tensor weight scales (one scalar covers the gate|up stack) + Bs, # (num_experts, N, 1) row-wise weight scales, or (num_experts, 1, 1) per tensor (stride_bs_n 0) C, # (S, N) output; under GATE the bf16 GLU intermediate Cs, # None (tensor-wide has no fused requant); kept for the shared epilogue signature Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None GatherIdx, # (S,) int32 — sorted position -> source row of A; read only when not None ScatterIdx, # (S,) int32 — sorted position -> destination row of C; read only when not None ExpertStart, # (NUM_EXPERTS_POW2 + 1,) int32 — cumulative row starts, S sentinel - Schedule, # packed tile schedule (build_packed_schedule); read iff PACKED_SCHEDULE + Schedule, # packed tile schedule (fill_packed_schedule); read iff PACKED_SCHEDULE # Shape S, N, @@ -508,6 +531,7 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( stride_b_k, stride_b_n, stride_bs_e, + stride_bs_n, stride_c_m, stride_c_n, stride_cs_m, @@ -521,7 +545,6 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, NUM_EXPERTS_POW2: tl.constexpr, - NUM_SMS: tl.constexpr, WARP_SPEC: tl.constexpr = False, # tuner axis: resolve M-tiles from the packed schedule (3 scalar loads) instead of the # E-wide register-resident layout — register relief, bit-identical tiles; scope pruner @@ -539,26 +562,29 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( SIMULATE_UNFUSED: tl.constexpr = False, INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, ): - """Tensor-scale grouped FP8 expert matmul kernel — persistent grid-stride over tiles. + """Tensor-scale / row-wise grouped FP8 expert matmul kernel — persistent grid-stride over tiles. Grouped expert scheduling with pre-quantized activations plus per-token activation scales - and one per-expert tensor weight scale. ``GATE`` fuses the gate|up projection (``B`` the - ``(E, 2N, K)`` stack, one scale for both) into a ``[BK, 2*BN]`` dot + SwiGLU ``glu``, - emitting the bf16 intermediate; ``GATE=False`` is the plain GEMM (bit-identical).""" + and a per-expert weight scale constant along K (per tensor, or per output row), both folded + after the loop. ``GATE`` fuses the gate|up projection (``B`` the ``(E, 2N, K)`` stack, its + scales interleaved alike) into a ``[BK, 2*BN]`` dot + SwiGLU ``glu``, emitting the bf16 + intermediate; ``GATE=False`` is the plain GEMM (bit-identical).""" + NUM_PROGRAMS = tl.num_programs(axis=0) start_pid = tl.program_id(axis=0) + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) if PACKED_SCHEDULE: total_m_tiles = load_packed_schedule(Schedule, BLOCK_SIZE_M) prefetched = prefetch_packed_entry( - Schedule, start_pid // tl.cdiv(N, BLOCK_SIZE_N), total_m_tiles + Schedule, grouped_tile_order(start_pid, total_m_tiles, num_n_tiles) // num_n_tiles, total_m_tiles ) else: exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = build_tile_layout( ExpertStart, NUM_EXPERTS_POW2, BLOCK_SIZE_M ) - num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) offs_k = tl.arange(0, BLOCK_SIZE_K) - for tile_id in tl.range(start_pid, total_m_tiles * num_n_tiles, NUM_SMS): + for logical_tile_id in tl.range(start_pid, total_m_tiles * num_n_tiles, NUM_PROGRAMS): + tile_id = grouped_tile_order(logical_tile_id, total_m_tiles, num_n_tiles) pid_n, _, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off, m_start = ( resolve_grouped_tile_packed( tile_id, num_n_tiles, prefetched, @@ -579,11 +605,11 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( ) if PACKED_SCHEDULE: prefetched = prefetch_packed_entry( - Schedule, (tile_id + NUM_SMS) // num_n_tiles, total_m_tiles + Schedule, grouped_tile_order(logical_tile_id + NUM_PROGRAMS, total_m_tiles, num_n_tiles) // num_n_tiles, total_m_tiles ) a_ptrs = operand_tile_ptrs(A, in_row, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) - # GATE stacks gate|up into one [BK, 2*BN] tile (one per-tensor scale covers both); - # rows are interleaved, so the tile is one 2*BN span. GATE=False -> plain [BK, BN]. + # GATE stacks gate|up into one [BK, 2*BN] tile; rows are interleaved, so the tile is + # one 2*BN span. GATE=False -> plain [BK, BN]. b_ptrs = weight_tile_ptrs( B + expert_id64 * stride_b_e, offs_bn, @@ -597,7 +623,6 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( a_s = tl.load( As + in_row * stride_as_m + expert_id64.to(tl.int32) * stride_as_e, mask=row_mask, other=1.0 ) - b_s = tl.load(Bs + expert_id64 * stride_bs_e) acc = acc_init("dot", BLOCK_SIZE_M, (2 if GATE else 1) * BLOCK_SIZE_N, False) for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC): @@ -612,8 +637,8 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( acc = acc + fp8_dot(a, w, False, BLOCK_SIZE_K) a_ptrs += BLOCK_SIZE_K * stride_a_k b_ptrs += BLOCK_SIZE_K * stride_b_k - # per-token activation scale and per-expert tensor weight scale, applied once post-loop - acc = acc * a_s[:, None] * b_s + # per-token activation scale and the expert's weight scale, applied once post-loop + acc = apply_row_scale(acc * a_s[:, None], Bs + expert_id64 * stride_bs_e, offs_bn, stride_bs_n) gemm_epilogue( C, @@ -653,7 +678,7 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( compute_modes=("dot_scaled", "dot"), a_memory_modes=("descriptor", "pointer"), b_memory_modes=("descriptor", "pointer"), - pre_hook=rebind_grouped_mx_descriptors, + pre_hook=grouped_mx_pre_hook, warp_spec=True, # dot_scaled+WS+TMA compiles+wins on a clean context (was a num_warps=2 misread) ), # prefill: no scalar branch; TMA descriptor vs pointer loads on both operands # the MXFP4/MXFP8 (and packed-activation) splits key themselves — the tuner appends @@ -706,7 +731,7 @@ def mx_dynamic_matmul_grouped_kernel( GatherIdx, # (S,) int32 — sorted position -> source row of A; read only when not None ScatterIdx, # (S,) int32 — sorted position -> destination row of C; read only when not None ExpertStart, # (NUM_EXPERTS_POW2 + 1,) int32 — cumulative row starts, S sentinel - Schedule, # packed tile schedule (build_packed_schedule); read iff PACKED_SCHEDULE + Schedule, # packed tile schedule (fill_packed_schedule); read iff PACKED_SCHEDULE # Shape S, N, @@ -921,13 +946,14 @@ def mx_dynamic_matmul_grouped_kernel( # WS-compile region; WS+pointer-A (masked load) is scored out. get_accelerator_autotuning_configs( packed_schedule=True, + swap_ab_prefill=True, tune_block_nk=True, tune_block_m=True, warp_spec=True, compute_modes=("dot", "dot_scaled"), a_memory_modes=("descriptor", "pointer"), b_memory_modes=("descriptor", "pointer"), - pre_hook=rebind_grouped_descriptors, + pre_hook=grouped_pre_hook, ), ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, @@ -943,7 +969,6 @@ def mx_dynamic_matmul_grouped_kernel( # columns masked by N_COLS in the epilogue), so non-dividing BN is correct — the # veto belongs to the maskless-store kernels (see OPTIMIZATION_LOG: n-veto-lift) mx_config_pruner("K", "N", block_within_k=False), # dot_scaled shape gates - global_scale_warp_spec_pruner(), gate_stacked_tmem_trap_pruner(), descriptor_box_pruner(), smem_pruner(), @@ -956,6 +981,7 @@ def mx_dynamic_matmul_grouped_kernel( def mx_weight_only_matmul_grouped_kernel( A, # (num_tokens, K) raw BF16/FP16 activations — NOT quantized ADescriptor, # host TMA descriptor over A (rows, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer" + AsGlobal, B, # (num_experts, N, K[/2]) MXFP4/NVFP4 (packed E2M1) / MXFP8 (E4M3) weights; 2N under GATE BDescriptor, # host TMA descriptor over B viewed (E, 2N|N, K_bytes), box (1, (2|1)*BN, BK_bytes); read iff B_MEMORY_MODE != "pointer" Bs, # (num_experts, N, K // SCALE_GROUP_K) group scales — UE8M0 (MX) or E4M3 (NVFP4) @@ -965,30 +991,30 @@ def mx_weight_only_matmul_grouped_kernel( GatherIdx, ScatterIdx, ExpertStart, - Schedule, # packed tile schedule (build_packed_schedule); read iff PACKED_SCHEDULE + Schedule, # packed tile schedule (fill_packed_schedule); read iff PACKED_SCHEDULE S, N, K, stride_a_m, stride_a_k, + stride_as_global: tl.constexpr, stride_b_e, stride_b_k, stride_b_n, stride_bs_e, stride_bs_n, stride_bs_k, + stride_bs_global, # expert stride of BsGlobal (``global_scale_stride``) stride_c_m, stride_c_n, stride_bias_e, stride_bias_n, - stride_bs_global, # expert stride of BsGlobal (``global_scale_stride``) num_experts, tokens_per_expert_bit_length, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, NUM_EXPERTS_POW2: tl.constexpr, - NUM_SMS: tl.constexpr, SCALE_GROUP_K: tl.constexpr, WEIGHT_VALUES_PER_BYTE: tl.constexpr, A_MEMORY_MODE: tl.constexpr = "pointer", @@ -1005,12 +1031,14 @@ def mx_weight_only_matmul_grouped_kernel( SWIGLU_LIMIT: tl.constexpr = None, SIMULATE_UNFUSED: tl.constexpr = False, INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, + SWAP_AB: tl.constexpr = False, ): """weight-only grouped expert matmul: raw bf16 activations against MXFP4/MXFP8 weights upcast to bf16 in-loop (unpack + per-group group-scale), plain ``tl.dot``. The ``matmul_ogs`` format. Persistent grid-stride; pointer or host-TMA operands, affine scales. ``GATE`` fuses the (E, 2N, K) gate|up stack + SwiGLU.""" n_width: tl.constexpr = (2 if GATE else 1) * BLOCK_SIZE_N + NUM_PROGRAMS = tl.num_programs(axis=0) start_pid = tl.program_id(axis=0) if PACKED_SCHEDULE: total_m_tiles = load_packed_schedule(Schedule, BLOCK_SIZE_M) @@ -1024,7 +1052,7 @@ def mx_weight_only_matmul_grouped_kernel( num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) offs_k = tl.arange(0, BLOCK_SIZE_K) offs_kb = tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) - for tile_id in tl.range(start_pid, total_m_tiles * num_n_tiles, NUM_SMS): + for tile_id in tl.range(start_pid, total_m_tiles * num_n_tiles, NUM_PROGRAMS): pid_n, _, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off, m_start = ( resolve_grouped_tile_packed( tile_id, num_n_tiles, prefetched, @@ -1036,15 +1064,15 @@ def mx_weight_only_matmul_grouped_kernel( ) if PACKED_SCHEDULE: prefetched = prefetch_packed_entry( - Schedule, (tile_id + NUM_SMS) // num_n_tiles, total_m_tiles + Schedule, (tile_id + NUM_PROGRAMS) // num_n_tiles, total_m_tiles ) offs_bn = offs_bn % (2 * N if GATE else N) # non-128 N: pointer rows wrap; wrapped cols masked (N_COLS) in epilogue a_ptrs = operand_tile_ptrs(A, in_row, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) b_ptrs = weight_tile_ptrs( B + expert_id64 * stride_b_e, offs_bn, offs_kb, - stride_b_n, stride_b_k, False, + stride_b_n, stride_b_k, SWAP_AB, ) - acc = acc_init("dot", BLOCK_SIZE_M, n_width, False) + acc = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB) kb_off = 0 # The tail K-tile slides back to end at K instead of running past the row, so BK need # only FIT K (see the batched sibling): the overlapped lanes are re-read and cancelled @@ -1061,26 +1089,29 @@ def mx_weight_only_matmul_grouped_kernel( w, w_s = load_weight_mx( b_ptrs, BDescriptor, Bs, None, None, 0, row0, n_off, kb_off, 0, expert_id64, pid_n, k, N, K, stride_bs_e, stride_bs_n, stride_bs_k, - GATE, True, False, B_MEMORY_MODE, False, False, + GATE, True, False, B_MEMORY_MODE, SWAP_AB, False, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, WEIGHT_VALUES_PER_BYTE, k_col_off=-(overlap // SCALE_GROUP_K), ) acc = mx_weight_only_compute( - acc, a, w, w_s, COMPUTE_MODE, BLOCK_SIZE_K, n_width, SCALE_GROUP_K + acc, a, w, w_s, COMPUTE_MODE, BLOCK_SIZE_K, n_width, SCALE_GROUP_K, + SWAP_AB=SWAP_AB, ) step = tl.minimum((k + 1) * BLOCK_SIZE_K, K - BLOCK_SIZE_K) - k_start a_ptrs += step * stride_a_k b_ptrs += (step // WEIGHT_VALUES_PER_BYTE) * stride_b_k kb_off = tl.minimum((k + 1) * BLOCK_SIZE_K, K - BLOCK_SIZE_K) // WEIGHT_VALUES_PER_BYTE + if SWAP_AB: + acc = tl.trans(acc) gemm_epilogue( C, C, acc, out_row, pid_n, tile_id // num_n_tiles, row_mask, stride_c_m, stride_c_n, 1, 1, BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, 1, ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, N_COLS=N, - GlobalScale=BsGlobal, - stride_global_e=stride_bs_global, + GlobalScale=BsGlobal, stride_global_e=stride_bs_global, + GlobalScaleA=AsGlobal, stride_global_a=stride_as_global, global_row=expert_id64, Bias=Bias, stride_bias_e=stride_bias_e, @@ -1088,16 +1119,263 @@ def mx_weight_only_matmul_grouped_kernel( ) +def _int_grouped_configs(): + configs = _grouped_dot_configs() + for config in configs: + config.kwargs["PROGRAMS_PER_SM"] = persistent_programs_per_sm(config.kwargs["BLOCK_SIZE_M"]) + config.kwargs["SWAP_AB"] = False + # These narrower tiles fit twice per SM when the initial WS register reservation is capped. + # Without the cap the vector partition reserves the whole register file, even with no spills. + if is_sm10x(): + for bn in (64, 128): + for packed in (False, True): + for stages in (2, 3): + configs.append(triton.Config( + dict(BLOCK_SIZE_M=128, BLOCK_SIZE_N=bn, BLOCK_SIZE_K=128, WARP_SPEC=True, + A_MEMORY_MODE="host_descriptor", B_MEMORY_MODE="host_descriptor", + PACKED_SCHEDULE=packed, PROGRAMS_PER_SM=2, SWAP_AB=False), + num_warps=4, num_stages=stages, maxnreg=128, pre_hook=grouped_pre_hook, + )) + # A computed MMA lhs can stay in tensor memory instead of staging decoded weights through shared memory. + for bn in (64, 128): + for packed in (False, True): + for stages in (2, 3): + configs.append(triton.Config( + dict(BLOCK_SIZE_M=128, BLOCK_SIZE_N=bn, BLOCK_SIZE_K=256, WARP_SPEC=True, + A_MEMORY_MODE="host_descriptor", B_MEMORY_MODE="host_descriptor", + PACKED_SCHEDULE=packed, PROGRAMS_PER_SM=1, SWAP_AB=True), + num_warps=4, num_stages=stages, pre_hook=grouped_pre_hook, + )) + for config in configs: + config.kwargs["PIPELINE"] = "k_loop" + # Tile overlap and flattening are separate choices: GLU register pressure can favor preserving the K loop. + for config in list(configs): + if config.kwargs["SWAP_AB"] and config.num_stages == 2: + configs.extend(pipeline_autotuning_configs([config], ("flattened",))) + if config.kwargs["BLOCK_SIZE_N"] == 64: + pipelines = ("tile_loop", "flattened") if config.kwargs["PACKED_SCHEDULE"] else ("tile_loop",) + configs.extend(pipeline_autotuning_configs([config], pipelines, num_warps=8)) + elif config.kwargs["PACKED_SCHEDULE"]: + # More epilogue warps relieve FP32 register pressure when the down projection adds bias. + configs.extend(pipeline_autotuning_configs([config], ("flattened",), num_warps=8)) + return configs + + @bayesian_autotune( - get_accelerator_autotuning_configs( - tune_block_nk=True, - warp_spec=True, - tune_block_m=True, - packed_schedule=True, - a_memory_modes=("descriptor", "pointer"), - b_memory_modes=("descriptor", "pointer"), - pre_hook=rebind_grouped_descriptors, - ), + # the MX weight-only grouped kernel's space minus dot_scaled (the scaled MMA takes float formats only) + _int_grouped_configs(), + # the width and the scale group decide which tiles fit and compile: part of the key + ["N", "K", "tokens_per_expert_bit_length", "GATE", "INT_BITS", "SCALE_GROUP_K", "SIMULATE_UNFUSED"], + n_trials=100, + benchmark_input=weight_input, + path_anchor_axes=PATH_ANCHOR_AXES, + prune_configs_by={ + "early_config_prune": compose_pruners( + warp_spec_memory_mode_pruner(weight_only=True), + packed_schedule_scope_pruner(), + # the tail tile slides back in bounds (see the K-loop) — BK only has to fit; no N veto, the N tail + # wraps and the epilogue masks it + block_fits_dim_pruner("K"), + int_scale_group_tile_pruner(), + gate_stacked_tmem_trap_pruner(), + descriptor_box_pruner(), + smem_pruner(), + warp_spec_compile_guard_pruner(), + gated_pointer_weight_warp_spec_pruner(), + ) + }, +) +@triton.jit +def int_weight_only_matmul_grouped_kernel( + A, # (num_tokens, K) prepared bf16/fp16, FP8 (W4A8), or INT8 (SmoothQuant) activations + ADescriptor, # host TMA descriptor over A (rows, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer" + AsGlobal, # fp32 calibrated input scale; None for weight-only + B, # (num_experts, N, K // VPW) interleaved uint32 words of VPW = 32 // INT_BITS values; 2N under GATE + BDescriptor, # host TMA descriptor over B viewed (E, 2N|N, K // VPW), box (1, (2|1)*BN, BK // VPW); read iff B_MEMORY_MODE != "pointer" + Bs, # (num_experts, N, K // SCALE_GROUP_K) float group scales + Bz, # (num_experts, N, K // SCALE_GROUP_K) integer zero points; None for a symmetric weight + BsGlobal, # fp32 W4A8 second-level weight scale; None otherwise + C, + Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None + GatherIdx, + ScatterIdx, + ExpertStart, + Schedule, # packed tile schedule (fill_packed_schedule); read iff PACKED_SCHEDULE + S, + N, + K, + stride_a_m, + stride_a_k, + stride_asglobal_e, + stride_b_e, + stride_b_k, + stride_b_n, + stride_bs_e, + stride_bs_n, + stride_bs_k, + stride_bz_e, + stride_bz_n, + stride_bz_k, + stride_bsglobal_e, + stride_c_m, + stride_c_n, + stride_bias_e, + stride_bias_n, + num_experts, + tokens_per_expert_bit_length, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + NUM_EXPERTS_POW2: tl.constexpr, + SCALE_GROUP_K: tl.constexpr, + A_MEMORY_MODE: tl.constexpr = "pointer", + B_MEMORY_MODE: tl.constexpr = "pointer", + WARP_SPEC: tl.constexpr = False, + PACKED_SCHEDULE: tl.constexpr = False, + GATE: tl.constexpr = False, + ACT_FN: tl.constexpr = "silu", + SWIGLU_ALPHA: tl.constexpr = None, + SWIGLU_LIMIT: tl.constexpr = None, + SIMULATE_UNFUSED: tl.constexpr = False, + INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, + INT_BITS: tl.constexpr = 4, + PROGRAMS_PER_SM: tl.constexpr = 1, + SWAP_AB: tl.constexpr = False, + PIPELINE: tl.constexpr = "k_loop", +): + """Integer grouped expert matmul: bf16/fp16 weight-only, calibrated INT4/FP8, or INT8 SmoothQuant. + Persistent grid-stride; pointer or host-TMA operands, affine scales. + ``GATE`` fuses the (E, 2N, K) gate|up stack + SwiGLU.""" + n_width: tl.constexpr = (2 if GATE else 1) * BLOCK_SIZE_N + # output-rows-major for the mma.sync tiles (BM < 64), whose rhs stays in registers; K-major for the tcgen05 tiles, + # whose rhs is staged through shared memory (int_dequant_tile) + ROWS_MAJOR: tl.constexpr = SWAP_AB or BLOCK_SIZE_M < 64 + NUM_PROGRAMS = tl.num_programs(axis=0) # persistent_programs_per_sm per SM + start_pid = tl.program_id(axis=0) + if PACKED_SCHEDULE: + total_m_tiles = load_packed_schedule(Schedule, BLOCK_SIZE_M) + if PIPELINE == "k_loop": + prefetched = prefetch_packed_entry( + Schedule, start_pid // tl.cdiv(N, BLOCK_SIZE_N), total_m_tiles + ) + else: + exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = build_tile_layout( + ExpertStart, NUM_EXPERTS_POW2, BLOCK_SIZE_M + ) + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + VPW: tl.constexpr = 32 // INT_BITS # values per uint32 word + offs_kb = tl.arange(0, BLOCK_SIZE_K // VPW) + for tile_id in tl.range( + start_pid, total_m_tiles * num_n_tiles, NUM_PROGRAMS, + flatten=PIPELINE == "flattened", warp_specialize=WARP_SPEC and PIPELINE != "k_loop", + ): + if PACKED_SCHEDULE and PIPELINE != "k_loop": + # Let the outer pipeline schedule this tile's entry instead of carrying a second prefetch across it. + prefetched = prefetch_packed_entry(Schedule, tile_id // num_n_tiles, total_m_tiles) + pid_n, _, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off, m_start = ( + resolve_grouped_tile_packed( + tile_id, num_n_tiles, prefetched, + GatherIdx, ScatterIdx, BLOCK_SIZE_N, BLOCK_SIZE_M, GATE, + ) if PACKED_SCHEDULE else resolve_grouped_tile( + tile_id, num_n_tiles, exp_start, freqs, tile_start_excl, e_offs, + GatherIdx, ScatterIdx, BLOCK_SIZE_N, BLOCK_SIZE_M, GATE, + ) + ) + if PACKED_SCHEDULE and PIPELINE == "k_loop": + prefetched = prefetch_packed_entry( + Schedule, (tile_id + NUM_PROGRAMS) // num_n_tiles, total_m_tiles + ) + offs_bn = offs_bn % (2 * N if GATE else N) # non-128 N: pointer rows wrap; wrapped cols masked (N_COLS) in epilogue + a_ptrs = operand_tile_ptrs(A, in_row, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) + b_ptrs = weight_tile_ptrs( + B + expert_id64 * stride_b_e, offs_bn, offs_kb, + stride_b_n, stride_b_k, ROWS_MAJOR, + ) + acc = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB, tl.int32 if A.dtype.element_ty == tl.int8 else tl.float32) + # zero points run one K step ahead (the 2D kernel's K-loop) + k_last = K // BLOCK_SIZE_K - 1 + w_s_next = load_weight_int_group( + Bs, Bs, expert_id64, pid_n, 0, N, K, stride_bs_e, stride_bs_n, stride_bs_k, + GATE, False, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, GROUPED=True, + ) + w_z_next = None + if Bz is not None: + w_z_next = load_weight_int_group( + Bz, Bz, expert_id64, pid_n, 0, N, K, stride_bz_e, stride_bz_n, stride_bz_k, + GATE, False, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, GROUPED=True, + ) + # BK divides K (int_scale_group_tile_pruner): no tail tile, and a constant pointer step keeps the weight + # rows' 16-byte alignment provable (a runtime step drops the async copies to 8 bytes) + for k in tl.range(0, K // BLOCK_SIZE_K, warp_specialize=WARP_SPEC and PIPELINE == "k_loop"): + k_start = k * BLOCK_SIZE_K + kb_off = k_start // VPW + a, _as = load_act_plain( + a_ptrs, ADescriptor, m_start, k_start, row_mask, in_row, + A_MEMORY_MODE, GatherIdx is not None, + ) + w_s = w_s_next + w_s_next = load_weight_int_group( + Bs, Bs, expert_id64, pid_n, tl.minimum(k + 1, k_last), N, K, + stride_bs_e, stride_bs_n, stride_bs_k, GATE, False, + BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, GROUPED=True, + ) + w, _w_unused = load_weight_plain( + b_ptrs, BDescriptor, row0, n_off, kb_off, GATE, True, B_MEMORY_MODE, ROWS_MAJOR, + BLOCK_SIZE_N, BLOCK_SIZE_K, VPW, + ) + w_z = w_z_next + if Bz is not None: + w_z_next = load_weight_int_group( + Bz, Bz, expert_id64, pid_n, tl.minimum(k + 1, k_last), N, K, stride_bz_e, stride_bz_n, + stride_bz_k, GATE, False, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, GROUPED=True, + ) + acc = int_weight_only_compute( + acc, a, w, w_s, w_z, BLOCK_SIZE_K, n_width, SCALE_GROUP_K, "dot", ROWS_MAJOR, INT_BITS, SWAP_AB, + weight_global=None if BsGlobal is None else tl.load(BsGlobal + expert_id64 * stride_bsglobal_e), + ) + a_ptrs += BLOCK_SIZE_K * stride_a_k + b_ptrs += (BLOCK_SIZE_K // VPW) * stride_b_k + + if SWAP_AB: + acc = tl.trans(acc) + output_scale = BsGlobal + output_scale_stride = stride_bsglobal_e + if A.dtype.element_ty == tl.int8: + output_scale = Bs + offs_bn * stride_bs_n + output_scale_stride = stride_bs_e + gemm_epilogue( + C, C, acc, out_row, pid_n, tile_id // num_n_tiles, row_mask, + stride_c_m, stride_c_n, 1, 1, + BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, 1, + ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, + N_COLS=N, + global_row=expert_id64, + GlobalScale=output_scale, stride_global_e=output_scale_stride, + GlobalScaleA=AsGlobal, stride_global_a=stride_asglobal_e, + Bias=Bias, + stride_bias_e=stride_bias_e, + stride_bias_n=stride_bias_n, + ) + + +def _full_precision_grouped_configs(): + configs = _grouped_dot_configs() + for config in configs: + config.kwargs["PIPELINE"] = "k_loop" + if is_sm10x(): + seeds = [ + c for c in configs + if c.kwargs["WARP_SPEC"] and c.kwargs["A_MEMORY_MODE"] == c.kwargs["B_MEMORY_MODE"] == "host_descriptor" + and c.kwargs["BLOCK_SIZE_M"] == 128 and c.kwargs["BLOCK_SIZE_K"] == 64 and c.num_stages == 4 + and (c.kwargs["BLOCK_SIZE_N"], c.num_warps) in ((128, 8), (256, 4)) + ] + configs.extend(pipeline_autotuning_configs(seeds, ("tile_loop", "flattened"))) + return configs + + +@bayesian_autotune( + _full_precision_grouped_configs(), # GATE keys the gate|up arm separately (its dot is 2*BN wide, a different tile optimum). ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, @@ -1111,6 +1389,7 @@ def mx_weight_only_matmul_grouped_kernel( warp_spec_compile_guard_pruner(), packed_schedule_scope_pruner(), descriptor_box_pruner(), + gate_stacked_tmem_trap_pruner(), smem_pruner(), ) }, @@ -1126,7 +1405,7 @@ def full_precision_matmul_grouped_kernel( GatherIdx, # (S,) int32 — sorted position -> source row of A; read only when not None ScatterIdx, # (S,) int32 — sorted position -> destination row of C; read only when not None ExpertStart, # (NUM_EXPERTS_POW2 + 1,) int32 — cumulative row starts, S sentinel - Schedule, # packed tile schedule (build_packed_schedule); read iff PACKED_SCHEDULE + Schedule, # packed tile schedule (fill_packed_schedule); read iff PACKED_SCHEDULE # Shape S, N, @@ -1148,7 +1427,6 @@ def full_precision_matmul_grouped_kernel( BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, NUM_EXPERTS_POW2: tl.constexpr, - NUM_SMS: tl.constexpr, WARP_SPEC: tl.constexpr = False, # tuner axis: resolve M-tiles from the packed schedule (3 scalar loads) instead of the # E-wide register-resident layout — pure register relief, bit-identical tiles; wins only @@ -1166,6 +1444,7 @@ def full_precision_matmul_grouped_kernel( SWIGLU_LIMIT: tl.constexpr = None, SIMULATE_UNFUSED: tl.constexpr = False, INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, + PIPELINE: tl.constexpr = "k_loop", ): """Full-precision grouped expert matmul kernel — persistent grid-stride over tiles. @@ -1173,19 +1452,26 @@ def full_precision_matmul_grouped_kernel( ``tl.dot`` with fp32 accumulation, no scales anywhere. ``GATE`` fuses the gate|up projection (``B`` the ``(E, 2N, K)`` stack) into a ``[BK, 2*BN]`` dot + SwiGLU ``glu``, emitting the intermediate; ``GATE=False`` is the plain GEMM (bit-identical).""" + NUM_PROGRAMS = tl.num_programs(axis=0) start_pid = tl.program_id(axis=0) if PACKED_SCHEDULE: total_m_tiles = load_packed_schedule(Schedule, BLOCK_SIZE_M) - prefetched = prefetch_packed_entry( - Schedule, start_pid // tl.cdiv(N, BLOCK_SIZE_N), total_m_tiles - ) + if PIPELINE == "k_loop": + prefetched = prefetch_packed_entry( + Schedule, start_pid // tl.cdiv(N, BLOCK_SIZE_N), total_m_tiles + ) else: exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = build_tile_layout( ExpertStart, NUM_EXPERTS_POW2, BLOCK_SIZE_M ) num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) offs_k = tl.arange(0, BLOCK_SIZE_K) - for tile_id in tl.range(start_pid, total_m_tiles * num_n_tiles, NUM_SMS): + for tile_id in tl.range( + start_pid, total_m_tiles * num_n_tiles, NUM_PROGRAMS, + flatten=PIPELINE == "flattened", warp_specialize=WARP_SPEC and PIPELINE != "k_loop", + ): + if PACKED_SCHEDULE and PIPELINE != "k_loop": + prefetched = prefetch_packed_entry(Schedule, tile_id // num_n_tiles, total_m_tiles) if PACKED_SCHEDULE: pid_n, _, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off, m_start = ( resolve_grouped_tile_packed( @@ -1209,9 +1495,9 @@ def full_precision_matmul_grouped_kernel( GATE, ) ) - if PACKED_SCHEDULE: + if PACKED_SCHEDULE and PIPELINE == "k_loop": prefetched = prefetch_packed_entry( - Schedule, (tile_id + NUM_SMS) // num_n_tiles, total_m_tiles + Schedule, (tile_id + NUM_PROGRAMS) // num_n_tiles, total_m_tiles ) a_ptrs = operand_tile_ptrs(A, in_row, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) # GATE reads the interleaved gate|up rows as one [BK, 2*BN] tile. @@ -1249,7 +1535,7 @@ def full_precision_matmul_grouped_kernel( # The box's rank, not the backend: only this kernel resolves its own 2D per-expert box # (above). The other grouped kernels run the same mode on Xe against host-built 3D boxes. B_DESCRIPTOR_IS_2D: tl.constexpr = B_MEMORY_MODE == "device_descriptor" - for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC): + for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC and PIPELINE == "k_loop"): a, _as = load_act_plain( a_ptrs, a_desc, m_start, k * BLOCK_SIZE_K, row_mask, in_row, A_MEMORY_MODE, GatherIdx is not None, @@ -1294,6 +1580,215 @@ def full_precision_matmul_grouped_kernel( ) +def _block_fp8_weight_only_grouped_configs(): + configs = _grouped_dot_configs() + for config in configs: + config.kwargs.update(SWAP_AB=False, PIPELINE="k_loop") + # Stage decoded weights on the MMA lhs and overlap consecutive output tiles. + if is_sm10x(): + for packed in (False, True): + seed = triton.Config( + dict(BLOCK_SIZE_M=128, BLOCK_SIZE_N=64, BLOCK_SIZE_K=256, + WARP_SPEC=True, PACKED_SCHEDULE=packed, SWAP_AB=True, + A_MEMORY_MODE="host_descriptor", B_MEMORY_MODE="host_descriptor"), + num_warps=8, num_stages=2, pre_hook=grouped_pre_hook, + ) + configs.extend(pipeline_autotuning_configs([seed], ("tile_loop", "flattened"))) + return configs + + +@bayesian_autotune( + _block_fp8_weight_only_grouped_configs(), + # GATE keys the gate|up arm separately (its dot is 2*BN wide, a different tile optimum). + ["N", "K", "tokens_per_expert_bit_length", "GATE", "BLOCK_N", "BLOCK_K", "GATE_SCALES"], + n_trials=100, + benchmark_input=weight_input, + path_anchor_axes=PATH_ANCHOR_AXES, + # BLOCK_SIZE_K is a tuned axis and the K-loop is maskless — veto non-dividing BKs. + prune_configs_by={ + "early_config_prune": compose_pruners( + block_within_dim_pruner("K"), + block_within_dim_pruner("N", "BLOCK_SIZE_N"), + warp_spec_compile_guard_pruner(), + warp_spec_memory_mode_pruner(gather_only=True), + packed_schedule_scope_pruner(), + descriptor_box_pruner(), + gate_stacked_tmem_trap_pruner(), + smem_pruner(), + ) + }, +) +@triton.jit +def block_fp8_weight_only_matmul_grouped_kernel( + A, # (num_tokens, K) BF16/FP16 activations, any row order + ADescriptor, # host TMA descriptor over A (rows, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer" + B, # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack + BDescriptor, # host TMA descriptor over B viewed (E, 2N|N, K), box (1, (2|1)*BN, BK); read iff B_MEMORY_MODE != "pointer" + Bs, # coarse FP8 weight scales + C, # (S, N) output; under GATE the GLU intermediate + Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None + GatherIdx, # (S,) int32 — sorted position -> source row of A; read only when not None + ScatterIdx, # (S,) int32 — sorted position -> destination row of C; read only when not None + ExpertStart, # (NUM_EXPERTS_POW2 + 1,) int32 — cumulative row starts, S sentinel + Schedule, # packed tile schedule; read iff PACKED_SCHEDULE + # Shape + S, + N, + K, + num_experts, + tokens_per_expert_bit_length, # autotune key only + # Strides + stride_a_m, + stride_a_k, + stride_b_e, + stride_b_k, + stride_b_n, + stride_bs_e, + stride_bs_n, + stride_bs_k, + stride_c_m, + stride_c_n, + stride_bias_e, + stride_bias_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_EXPERTS_POW2: tl.constexpr, + NUM_SMS: tl.constexpr, + WARP_SPEC: tl.constexpr = False, + PACKED_SCHEDULE: tl.constexpr = False, + B_MEMORY_MODE: tl.constexpr = "pointer", + A_MEMORY_MODE: tl.constexpr = "pointer", + PIPELINE: tl.constexpr = "k_loop", + SWAP_AB: tl.constexpr = False, + GATE: tl.constexpr = False, + GATE_SCALES: tl.constexpr = False, + ACT_FN: tl.constexpr = "silu", + SWIGLU_ALPHA: tl.constexpr = None, + SWIGLU_LIMIT: tl.constexpr = None, + SIMULATE_UNFUSED: tl.constexpr = False, + INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, +): + """Persistent FP8 W8A16 GEMM with optional fused gate/up activation.""" + start_pid = tl.program_id(axis=0) + if PACKED_SCHEDULE: + total_m_tiles = load_packed_schedule(Schedule, BLOCK_SIZE_M) + if PIPELINE == "k_loop": + prefetched = prefetch_packed_entry( + Schedule, start_pid // tl.cdiv(N, BLOCK_SIZE_N), total_m_tiles + ) + else: + exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = build_tile_layout( + ExpertStart, NUM_EXPERTS_POW2, BLOCK_SIZE_M + ) + num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + for tile_id in tl.range( + start_pid, total_m_tiles * num_n_tiles, NUM_SMS, + flatten=PIPELINE == "flattened", warp_specialize=WARP_SPEC and PIPELINE != "k_loop", + ): + if PACKED_SCHEDULE and PIPELINE != "k_loop": + prefetched = prefetch_packed_entry(Schedule, tile_id // num_n_tiles, total_m_tiles) + if PACKED_SCHEDULE: + pid_n, _, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off, m_start = ( + resolve_grouped_tile_packed( + tile_id, num_n_tiles, prefetched, + GatherIdx, ScatterIdx, BLOCK_SIZE_N, BLOCK_SIZE_M, GATE, + ) + ) + else: + pid_n, _, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off, m_start = ( + resolve_grouped_tile( + tile_id, + num_n_tiles, + exp_start, + freqs, + tile_start_excl, + e_offs, + GatherIdx, + ScatterIdx, + BLOCK_SIZE_N, + BLOCK_SIZE_M, + GATE, + ) + ) + if PACKED_SCHEDULE and PIPELINE == "k_loop": + prefetched = prefetch_packed_entry( + Schedule, (tile_id + NUM_SMS) // num_n_tiles, total_m_tiles + ) + a_ptrs = operand_tile_ptrs(A, in_row, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) + b_ptrs = weight_tile_ptrs( + B + expert_id64 * stride_b_e, + offs_bn, + offs_k, + stride_b_n, + stride_b_k, + SWAP_AB, + ) + + acc = acc_init("dot", BLOCK_SIZE_M, (2 if GATE else 1) * BLOCK_SIZE_N, False) + # Device-built weight descriptors address a 2D slab for this expert. + a_desc = operand_tile_descriptor( + ADescriptor, A, S, K, stride_a_m, 1, # K is the contiguous dim of (S, K) + BLOCK_SIZE_M, BLOCK_SIZE_K, A_MEMORY_MODE, + ) + b_desc = operand_tile_descriptor( + BDescriptor, B + expert_id64 * stride_b_e, N, K, stride_b_n, 1, # and of the slab + BLOCK_SIZE_N, BLOCK_SIZE_K, B_MEMORY_MODE, GATE, + ) + B_DESCRIPTOR_IS_2D: tl.constexpr = B_MEMORY_MODE == "device_descriptor" + for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC and PIPELINE == "k_loop"): + a, _as = load_act_plain( + a_ptrs, a_desc, m_start, k * BLOCK_SIZE_K, row_mask, in_row, + A_MEMORY_MODE, GatherIdx is not None, + ) + w = load_weight_block_fp8( + b_ptrs, b_desc, Bs, row0, n_off, k * BLOCK_SIZE_K, + expert_id64, offs_bn, k * BLOCK_SIZE_K + offs_k, + stride_bs_e, stride_bs_n, stride_bs_k, + GATE, True, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K, + BLOCK_N, BLOCK_K, GATE_SCALES, a.dtype, + B_DESCRIPTOR_IS_2D=B_DESCRIPTOR_IS_2D, + ) + acc = plain_dot(acc, a, w, SWAP_AB) + if A_MEMORY_MODE == "pointer": + a_ptrs += BLOCK_SIZE_K * stride_a_k + if B_MEMORY_MODE == "pointer": + b_ptrs += BLOCK_SIZE_K * stride_b_k + + gemm_epilogue( + C, + None, # no requant arm here + acc, + out_row, + pid_n, + tile_id // num_n_tiles, + row_mask, + stride_c_m, + stride_c_n, + 1, # Cs strides: unread, Cs is None + 1, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + GATE, + None, + 1, + ACT_FN, + SWIGLU_ALPHA, + SWIGLU_LIMIT, + SIMULATE_UNFUSED, + INTERMEDIATE_DTYPE, + Bias=Bias, + stride_bias_e=stride_bias_e, + stride_bias_n=stride_bias_n, + global_row=expert_id64, + ) + + + @compile_time_only_triton_op( add_op_namespace_prefix("w8a8_block_dynamic_fp8_matmul_grouped"), mutates_args=(), @@ -1390,9 +1885,7 @@ def w8a8_block_dynamic_fp8_matmul_grouped( A, B.view(num_experts, 2 * N if gate else N, K) ) - # one launch per call (graph-safe, ~2µs) so PACKED_SCHEDULE and inline configs share - # identical launch args during tunes - schedule = build_packed_schedule(expert_start, S, num_experts) + schedule = packed_schedule_buffer(S, num_experts, expert_start.device) with device_context(A.device): bias_stride_e, bias_stride_n = bias_strides(bias) compile_time_only_triton_wrap(w8a8_block_dynamic_fp8_matmul_grouped_kernel)[ @@ -1524,9 +2017,7 @@ def w8a8_block_static_fp8_matmul_grouped( A_q, B.view(num_experts, 2 * N if gate else N, K) ) - # one launch per call (graph-safe, ~2µs) so PACKED_SCHEDULE and inline configs share - # identical launch args during tunes - schedule = build_packed_schedule(expert_start, S, num_experts) + schedule = packed_schedule_buffer(S, num_experts, expert_start.device) with device_context(A.device): bias_stride_e, bias_stride_n = bias_strides(bias) compile_time_only_triton_wrap(w8a8_block_static_fp8_matmul_grouped_kernel)[ @@ -1603,8 +2094,9 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped( gather_idx: torch.Tensor | None = None, scatter_idx: torch.Tensor | None = None, bias: torch.Tensor | None = None, + activation_scale_ub: float | None = None, ) -> list[torch.Tensor]: - """Tensor-scale grouped FP8 matmul over expert-sorted positions (per-tile + """Tensor-scale / row-wise grouped FP8 matmul over expert-sorted positions (per-tile gather/scatter, the sort is virtual — see ``compute_grouped_scheduling`` for the maps). Raw activations are quantized here (``fp8_act_quant_tensor_wide``, one offline pass); pre-quantized ``A`` + ``As`` is read as-is. @@ -1615,7 +2107,8 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped( B: (num_experts, N, K) FP8 expert weights; under ``gate`` the (num_experts, 2N, K) stack As: (S,) per-token scales alongside a pre-quantized A, or — on a raw A — the calibrated (static) activation scale: one value, or one per expert - Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales + Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales, or (num_experts, N, 1) + row-wise (2N rows under ``gate``) expert_start: (num_experts_pow2 + 1,) int32 — cumulative sorted-row starts, S sentinel gather_idx: optional (S,) — sorted position -> source row of A; None = A is expert-sorted scatter_idx: optional (S,) — sorted position -> destination row of C; None = C stays expert-sorted @@ -1635,11 +2128,12 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped( num_experts, _, N = expert_weight_shape(B, gate) S = routed_rows(A, gather_idx, scatter_idx, expert_start, num_experts) - # Normalize Bs to (num_experts, 1, 1) — one per-tensor scale (covers the gate|up stack) Bs = normalize_per_expert_scale(Bs, num_experts) output_dtype = resolve_output_dtype(output_dtype, A, As) - A, As, as_stride_m, as_stride_e = tensor_wide_act_operands(A, As, num_experts) + A, As, as_stride_m, as_stride_e = tensor_wide_act_operands( + A, As, num_experts, activation_scale_ub=activation_scale_ub + ) if as_stride_e and expand_regime(S, num_experts): # a per-expert calibrated scale: lay the routed rows out quantized rather than read # them raw through every N-tile — the same copy-amortizes-at-prefill law below @@ -1648,21 +2142,23 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped( if A.dtype == FP8_DTYPE: # post-quant: trade the in-kernel gather for one packed-row copy where that wins. Keyed # on A being quantized, which the dynamic quant, a shared calibrated scale and the - # per-expert layout above all reach; raw rows (decode) keep the gather. - A, _, gather_idx = expand_gather_below_parity(A, None, gather_idx, num_experts, expert_start) + # per-expert layout above all reach; raw rows (decode) keep the gather. Per-token scales + # (stride_m 1) follow their rows into the copy. + row_scales = As.view(-1, 1) if as_stride_m else None + A, row_scales, gather_idx = expand_gather_below_parity(A, row_scales, gather_idx, num_experts, expert_start) + if as_stride_m: + As = row_scales.view(-1) + as_stride_m = As.stride(0) C = A.new_empty(S, N, dtype=output_dtype) - num_sms = persistent_program_count(A.device.index) a_descriptor, b_descriptor = build_grouped_operand_descriptors( A, B.view(num_experts, 2 * N if gate else N, K) ) - # one launch per call (graph-safe, ~2µs) so PACKED_SCHEDULE and inline configs share - # identical launch args during tunes - schedule = build_packed_schedule(expert_start, S, num_experts) + schedule = packed_schedule_buffer(S, num_experts, expert_start.device) with device_context(A.device): bias_stride_e, bias_stride_n = bias_strides(bias) compile_time_only_triton_wrap(w8a8_tensor_dynamic_fp8_matmul_grouped_kernel)[ - (num_sms,) + lambda meta: (persistent_program_count(A.device.index, meta["BLOCK_SIZE_M"]),) ]( A, a_descriptor, @@ -1688,6 +2184,7 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped( B.stride(2), B.stride(1), Bs.stride(0), + row_scale_stride(Bs), C.stride(0), C.stride(1), 1, # Cs strides: unread, Cs is None @@ -1697,7 +2194,6 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped( num_experts=num_experts, tokens_per_expert_bit_length=tokens_per_expert_bucket(S, num_experts), NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts), - NUM_SMS=num_sms, GATE=gate, ACT_FN=act_fn, SWIGLU_ALPHA=swiglu_alpha, @@ -1732,6 +2228,7 @@ def mx_dynamic_matmul_grouped( b_global_scale: torch.Tensor | None = None, output_global_scale: torch.Tensor | None = None, bias: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, ) -> list[torch.Tensor]: """Grouped MX matmul over expert-sorted positions (per-tile gather/scatter, the sort is virtual — see ``compute_grouped_scheduling`` for the maps). Activations arrive @@ -1749,6 +2246,33 @@ def mx_dynamic_matmul_grouped( scatter_idx: optional (S,) — sorted position -> destination row of C; None = C stays expert-sorted """ assert A.ndim == 2 and B.ndim == 3 and Bs.ndim in (3, 5) # 5D = pre-swizzled + if a_pre_quant_scale is not None and activation_format != "fp8": + assert As is None, "pre_quant_scale applies before activation quantization; pass raw inputs" + num_rows = routed_rows(A, gather_idx, scatter_idx, expert_start, B.shape[0]) + A = permute_routed_rows( + A, None, num_rows, gather_idx, expert_start=expert_start, + pre_quant_scale=a_pre_quant_scale, num_experts=B.shape[0], + ) + gather_idx = None + if activation_format == "fp8": + assert not quantize_output and a_global_scale is None, "ordinary FP8 inputs use As=input_scale and no fused requant" + output_dtype = resolve_output_dtype(output_dtype, A, None) + group = mx_scale_family(Bs, A.shape[-1]) + num_rows = routed_rows(A, gather_idx, scatter_idx, expert_start, B.shape[0]) + weight_upcast = Bs.ndim == 3 and Bs.dtype == torch.float8_e4m3fn and group == MX_SCALE_GROUP_K + A, As, a_global_scale = prepare_fp8_mx_inputs( + A, As, ue8m0_as_uint8(Bs).dtype, group, B.shape[0], num_rows, + gather_idx, expert_start=expert_start, pre_quant_scale=a_pre_quant_scale, storage_dtype=torch.bfloat16 if weight_upcast else FP8_DTYPE, + ) + gather_idx = None + if weight_upcast: + return mx_weight_only_matmul_grouped( + A, B, Bs, expert_start, output_dtype=output_dtype, gate=gate, act_fn=act_fn, + swiglu_alpha=swiglu_alpha, swiglu_limit=swiglu_limit, simulate_unfused=simulate_unfused, + scatter_idx=scatter_idx, + b_global_scale=b_global_scale, a_global_scale=a_global_scale, + bias=bias, + ) assert B.dtype in (torch.int8, torch.float8_e4m3fn), ( f"B must be int8 (packed E2M1) or float8_e4m3fn (E4M3), got {B.dtype}" ) @@ -1825,7 +2349,7 @@ def mx_dynamic_matmul_grouped( # rows the GEMM then reads without a gather. quantize_per_routed_row = is_per_expert_global(a_global_scale) and gather_idx is not None if a_global_scale is not None: - assert activation_format == "nvfp4", "an activation global is NVFP4-only" + assert activation_format in ("nvfp4", "fp8"), "an activation global requires NVFP4 or calibrated tensor-FP8" assert not quantize_per_routed_row or As is None, ( "a per-expert a_global_scale under a gather needs a raw A" ) @@ -1947,9 +2471,7 @@ def mx_dynamic_matmul_grouped( # (the SA/SB swizzled-scale descriptors and their placeholder boxes — re-bound per tuned # config by rebind_grouped_mx_descriptors — are built above with the scales.) - # one launch per call (graph-safe, ~2µs) so PACKED_SCHEDULE and inline configs share - # identical launch args during tunes - schedule = build_packed_schedule(expert_start, S, num_experts) + schedule = packed_schedule_buffer(S, num_experts, expert_start.device) with device_context(A.device): bias_stride_e, bias_stride_n = bias_strides(bias) compile_time_only_triton_wrap(mx_dynamic_matmul_grouped_kernel)[(num_sms,)]( @@ -2060,17 +2582,16 @@ def full_precision_matmul_grouped( output_dtype = resolve_output_dtype(output_dtype, A, None) C = A.new_empty(S, N, dtype=output_dtype) - num_sms = persistent_program_count(A.device.index) a_descriptor, b_descriptor = build_grouped_operand_descriptors( A, B.view(num_experts, 2 * N if gate else N, K) ) - # one launch per call (graph-safe, ~2µs) so PACKED_SCHEDULE and inline configs share - # identical launch args during tunes - schedule = build_packed_schedule(expert_start, S, num_experts) + schedule = packed_schedule_buffer(S, num_experts, expert_start.device) with device_context(A.device): bias_stride_e, bias_stride_n = bias_strides(bias) - compile_time_only_triton_wrap(full_precision_matmul_grouped_kernel)[(num_sms,)]( + compile_time_only_triton_wrap(full_precision_matmul_grouped_kernel)[ + lambda meta: (persistent_program_count(A.device.index, meta["BLOCK_SIZE_M"]),) + ]( A, a_descriptor, B, @@ -2096,8 +2617,109 @@ def full_precision_matmul_grouped( num_experts=num_experts, tokens_per_expert_bit_length=tokens_per_expert_bucket(S, num_experts), NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts), + GATE=gate, + ACT_FN=act_fn, + SWIGLU_ALPHA=swiglu_alpha, + SWIGLU_LIMIT=swiglu_limit, + SIMULATE_UNFUSED=simulate_unfused, + INTERMEDIATE_DTYPE=tl_dtype(output_dtype), + ) + + return [C] + + +@compile_time_only_triton_op( + add_op_namespace_prefix("block_fp8_weight_only_matmul_grouped"), + mutates_args=(), + opaque=True, +) +def block_fp8_weight_only_matmul_grouped( + A: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + expert_start: torch.Tensor, + gate: bool = False, + act_fn: str = "silu", + swiglu_alpha: float | None = None, + swiglu_limit: float | None = None, + simulate_unfused: bool = False, + activation_format: str | None = None, + quantize_output: bool = False, + output_dtype: torch.dtype | None = None, + gather_idx: torch.Tensor | None = None, + scatter_idx: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + *, + interleaved_scales: bool = False, +) -> List[torch.Tensor]: + """FP8 W8A16 grouped GEMM; stage weights only when routed reuse amortizes conversion.""" + validate_dense_operands(A, B) + assert A.dtype in (torch.bfloat16, torch.float16), "FP8 weight-only takes raw BF16/FP16 activations" + assert activation_format in (None, "bf16") and not quantize_output, ( + "FP8 weight-only requires unquantized activations and output" + ) + + _, K = A.shape + + num_experts, _, N = expert_weight_shape(B, gate) + S = routed_rows(A, gather_idx, scatter_idx, expert_start, num_experts) + if S >= FP8_WEIGHT_ONLY_DEQUANT_MIN_M * num_experts: + weights = dequant_block_fp8_weight( + B, Bs, A.dtype, gate=gate or interleaved_scales, expert_start=expert_start, + ) + return full_precision_matmul_grouped( + A, weights, expert_start, gate, act_fn, swiglu_alpha, swiglu_limit, + simulate_unfused, activation_format, quantize_output, output_dtype, + gather_idx, scatter_idx, bias=bias, + ) + Bs, scale_block_n, scale_block_k = block_fp8_weight_scale_grid(B, Bs, gate or interleaved_scales) + Bs = ue8m0_as_uint8(Bs) + + output_dtype = resolve_output_dtype(output_dtype, A, None) + C = A.new_empty(S, N, dtype=output_dtype) + num_sms = persistent_program_count(A.device.index) + a_descriptor, b_descriptor = build_grouped_operand_descriptors( + A, B.view(num_experts, 2 * N if gate else N, K) + ) + schedule = packed_schedule_buffer(S, num_experts, A.device) + + with device_context(A.device): + bias_stride_e, bias_stride_n = bias_strides(bias) + compile_time_only_triton_wrap(block_fp8_weight_only_matmul_grouped_kernel)[(num_sms,)]( + A, + a_descriptor, + B, + b_descriptor, + Bs, + C, + bias, + gather_idx, # None = A is expert-sorted; read only when not None (folds at trace time) + scatter_idx, # None = C is expert-sorted; read only when not None (folds at trace time) + expert_start, + schedule, + S, + N, + K, + num_experts, + tokens_per_expert_bucket(S, num_experts), + A.stride(0), + A.stride(1), + B.stride(0), + B.stride(2), + B.stride(1), + Bs.stride(0), + Bs.stride(1), + Bs.stride(2), + C.stride(0), + C.stride(1), + bias_stride_e, + bias_stride_n, + NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts), NUM_SMS=num_sms, + BLOCK_N=scale_block_n, + BLOCK_K=scale_block_k, GATE=gate, + GATE_SCALES=(gate or interleaved_scales) and Bs.shape[1] > 1, ACT_FN=act_fn, SWIGLU_ALPHA=swiglu_alpha, SWIGLU_LIMIT=swiglu_limit, @@ -2128,6 +2750,7 @@ def mx_weight_only_matmul_grouped( scatter_idx: torch.Tensor | None = None, b_global_scale: torch.Tensor | None = None, bias: torch.Tensor | None = None, + a_global_scale: torch.Tensor | None = None, ) -> list[torch.Tensor]: """weight-only grouped matmul: raw bf16/fp16 activations against MXFP4/NVFP4/MXFP8 weights upcast to bf16 in-loop — the ``matmul_ogs`` format (fp4/fp8 weight bytes, bf16 acts, weight upcast in-MMA). @@ -2147,11 +2770,14 @@ def mx_weight_only_matmul_grouped( ) N = rows // 2 if gate else rows scale_group = mx_scale_family(Bs, K) + assert Bs.dtype != torch.float8_e4m3fn or scale_group == NVFP4_SCALE_GROUP_K or a_global_scale is not None, ( + "group-32 E4M3 FP4 weights require calibrated FP8 activations" + ) b_global_scale = normalize_global_scale(b_global_scale, num_experts) + a_global_scale = normalize_global_scale(a_global_scale, num_experts) S = routed_rows(A, gather_idx, scatter_idx, expert_start, num_experts) output_dtype = resolve_output_dtype(output_dtype, A, None) C = A.new_empty(S, N, dtype=output_dtype) - num_sms = persistent_program_count(A.device.index) b_u8 = e2m1_as_uint8(B) bs_u8 = ue8m0_as_uint8(Bs) # Operand host-TMA descriptors (A over (S, K), B over the (E, 2N|N, K_bytes) weight view); @@ -2161,14 +2787,15 @@ def mx_weight_only_matmul_grouped( A, b_u8.view(num_experts, 2 * N if gate else N, K_b) ) - # one launch per call (graph-safe, ~2µs) so PACKED_SCHEDULE and inline configs share - # identical launch args during tunes - schedule = build_packed_schedule(expert_start, S, num_experts) + schedule = packed_schedule_buffer(S, num_experts, expert_start.device) with device_context(A.device): bias_stride_e, bias_stride_n = bias_strides(bias) - compile_time_only_triton_wrap(mx_weight_only_matmul_grouped_kernel)[(num_sms,)]( + compile_time_only_triton_wrap(mx_weight_only_matmul_grouped_kernel)[ + lambda meta: (persistent_program_count(A.device.index, meta["BLOCK_SIZE_M"]),) + ]( A, a_descriptor, + a_global_scale, b_u8, b_descriptor, bs_u8, @@ -2184,21 +2811,21 @@ def mx_weight_only_matmul_grouped( K, A.stride(0), A.stride(1), + global_scale_stride(a_global_scale), b_u8.stride(0), b_u8.stride(2), b_u8.stride(1), bs_u8.stride(0), bs_u8.stride(1), bs_u8.stride(2), + global_scale_stride(b_global_scale), C.stride(0), C.stride(1), bias_stride_e, bias_stride_n, - global_scale_stride(b_global_scale), num_experts=num_experts, tokens_per_expert_bit_length=tokens_per_expert_bucket(S, num_experts), NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts), - NUM_SMS=num_sms, SCALE_GROUP_K=scale_group, WEIGHT_VALUES_PER_BYTE=WEIGHT_VALUES_PER_BYTE, GATE=gate, @@ -2212,6 +2839,112 @@ def mx_weight_only_matmul_grouped( return [C] +@compile_time_only_triton_op( + add_op_namespace_prefix("int_weight_only_matmul_grouped"), + mutates_args=(), + opaque=True, +) +def int_weight_only_matmul_grouped( + A: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + expert_start: torch.Tensor, + gate: bool = False, + act_fn: str = "silu", + swiglu_alpha: float | None = None, + swiglu_limit: float | None = None, + simulate_unfused: bool = False, + output_dtype: torch.dtype | None = None, + gather_idx: torch.Tensor | None = None, + scatter_idx: torch.Tensor | None = None, + b_zero: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + b_perm: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, + a_scale: torch.Tensor | None = None, + b_scale2: torch.Tensor | None = None, +) -> list[torch.Tensor]: + """Integer weight-only grouped matmul: raw bf16/fp16 activations against 2/4/8-bit integer weights, ``B`` + ``(E, N, K // VPW)`` interleaved uint32 words, ``Bs`` ``(E, N, K // group)`` float group scales, ``b_zero`` the + matching integer zero points (``None``: symmetric) and ``b_perm`` ``(E, K)`` int32 the act-order K orders + (``None``: natural), read by one ``permute_routed_rows`` launch that lays the routed rows out contiguous and in + their experts' orders ahead of the natural-order kernel (B200, Qwen3-30B-A3B gate|up, T=1-512: 1.1-1.6x faster + than an in-kernel column gather). Same gather/scatter expert scheduling as the other grouped ops. Returns + ``[C]``. The weight is dequantized to the activation dtype (``((q - zero) * scale)`` rounded in the scale dtype, then cast), as the torch + "dequantize, then matmul" reference does, so results match it up to fp32 summation order.""" + bits = validate_int_operands(A, B, Bs, b_zero, b_perm, a_scale, b_scale2) + smoothquant = a_scale is not None and b_scale2 is None + a_scale = normalize_global_scale(a_scale, B.shape[0]) + b_scale2 = normalize_global_scale(b_scale2, B.shape[0]) + _, K = A.shape + num_experts, rows, K_b = B.shape + N = rows // 2 if gate else rows + S = routed_rows(A, gather_idx, scatter_idx, expert_start, num_experts) + output_dtype = resolve_output_dtype(output_dtype, A, None) + if b_perm is not None or a_pre_quant_scale is not None or a_scale is not None: + A = permute_routed_rows( + A, b_perm, S, gather_idx, expert_start=expert_start, + pre_quant_scale=a_pre_quant_scale, num_experts=num_experts, quant_scale=a_scale, + quant_dtype=torch.int8 if smoothquant else FP8_DTYPE, + ) + gather_idx = None + C = A.new_empty(S, N, dtype=output_dtype) + num_sms = persistent_program_count(A.device.index) + a_descriptor, b_descriptor = build_grouped_operand_descriptors(A, B) + schedule = packed_schedule_buffer(S, num_experts, expert_start.device) + with device_context(A.device): + bias_stride_e, bias_stride_n = bias_strides(bias) + compile_time_only_triton_wrap(int_weight_only_matmul_grouped_kernel)[ + lambda meta: (num_sms * meta["PROGRAMS_PER_SM"],) + ]( + A, + a_descriptor, + a_scale, + B, + b_descriptor, + Bs, + b_zero, + b_scale2, + C, + bias, + gather_idx, + scatter_idx, + expert_start, + schedule, + S, + N, + K, + A.stride(0), + A.stride(1), + global_scale_stride(a_scale), + B.stride(0), + B.stride(2), + B.stride(1), + Bs.stride(0), + Bs.stride(1), + Bs.stride(2), + *(b_zero.stride() if b_zero is not None else (0, 0, 0)), + global_scale_stride(b_scale2), + C.stride(0), + C.stride(1), + bias_stride_e, + bias_stride_n, + num_experts=num_experts, + tokens_per_expert_bit_length=tokens_per_expert_bucket(S, num_experts), + NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts), + SCALE_GROUP_K=K // Bs.shape[2], + INT_BITS=bits, + GATE=gate, + ACT_FN=act_fn, + SWIGLU_ALPHA=swiglu_alpha, + SWIGLU_LIMIT=swiglu_limit, + SIMULATE_UNFUSED=simulate_unfused, + INTERMEDIATE_DTYPE=tl_dtype(output_dtype), + ) + + return [C] + + def matmul_grouped( A: torch.Tensor, B: torch.Tensor, @@ -2221,6 +2954,7 @@ def matmul_grouped( expert_start: torch.Tensor, bias: torch.Tensor | None = None, # (E, N_out) per-expert output bias; (N_out,) for 2D activation_format: str | None = None, + activation_scale_ub: float | None = None, gate: bool = False, act_fn: str = "silu", swiglu_alpha: float | None = None, @@ -2233,6 +2967,9 @@ def matmul_grouped( a_global_scale: torch.Tensor | None = None, b_global_scale: torch.Tensor | None = None, output_global_scale: torch.Tensor | None = None, + b_zero: torch.Tensor | None = None, + b_perm: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Grouped matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). ``expert_start`` is the ``(E+1,)`` tiling schedule from one ``compute_grouped_scheduling`` @@ -2269,28 +3006,36 @@ def matmul_grouped( - ``Bs`` None → ``full_precision_matmul_grouped`` (plain dot, no scales anywhere). - MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0 group-32 ``Bs`` → ``mx_dynamic_matmul_grouped``. - - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) → + - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) or per output row (``(E, N, 1)``) → ``w8a8_tensor_dynamic_fp8_matmul_grouped``. - block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_grouped``. + + ``activation_scale_ub`` caps the dynamic per-token amax for row-wise FP8; + it requires raw activations and applies before conversion to E4M3. + - uint32 integer weights → ``int_weight_only_matmul_grouped`` (weight-only), ``b_zero`` their zero points + (``None``: symmetric), ``b_perm`` ``(E, K)`` their act-order K orders (``None``: natural). """ - check_activation_format(activation_format, quantize_output) - if gather_idx is not None and activation_format == "bf16" and As is None: + _validate_quantization_args( + A, As, B, Bs, a_global_scale, b_global_scale, output_global_scale, + b_zero, b_perm, a_pre_quant_scale, activation_format, quantize_output, activation_scale_ub, + ) + if gather_idx is not None and activation_format == "bf16" and As is None and a_pre_quant_scale is None: # weight-only: raw bf16 rows are what the kernel consumes, so the trade is decided # here; the QUANTIZED formats decide it in their family wrappers AFTER the act quant - # (expanding packed rows + scales — see ``expand_gather_below_parity``). + # (expanding packed rows + scales — see ``expand_gather_below_parity``). An AWQ input + # transform combines the sorted copy and column scaling in its preparation pass below. A, _, gather_idx = expand_gather_below_parity(A, None, gather_idx, B.shape[0], expert_start) - assert (a_global_scale is None and b_global_scale is None) or ( - Bs is not None and weight_format(B, Bs) == "nvfp4" - ), "two-level globals (a_global_scale / b_global_scale) are NVFP4-only" - assert output_global_scale is None or ( - Bs is not None and quantize_output and resolve_activation_format(activation_format, B, Bs) == "nvfp4" - ), ( - "output_global_scale is the NVFP4 requant second level — it requires quantize_output=True on NVFP4 " - "(the epilogue would otherwise normalize by it with nothing downstream to compensate)" - ) + if B.dtype == FP8_DTYPE and Bs is not None and not is_mx(B, Bs) and activation_format == "bf16": + assert As is None and not quantize_output, "weight-only takes raw activations without requantization" + out = block_fp8_weight_only_matmul_grouped( + A, B, Bs, expert_start, gate, act_fn, swiglu_alpha, swiglu_limit, + simulate_unfused, activation_format, quantize_output, output_dtype, + gather_idx, scatter_idx, bias=bias, + ) + return out[0] # a calibrated (static) scale on a raw A (see `tensor_wide_act_operands`): block-scale weights # have a dedicated static kernel, per-tensor ones read the same As on the tensor-wide arm - static_act = As is not None and As.ndim <= 1 and As.numel() in (1, B.shape[0]) and A.dtype != FP8_DTYPE + static_act = B.dtype not in INT_WEIGHT_DTYPES and not is_mx(B, Bs) and As is not None and As.ndim <= 1 and As.numel() in (1, B.shape[0]) and A.dtype != FP8_DTYPE if static_act: assert Bs is not None and not is_mx(B, Bs), ( "a calibrated (static) activation scale is an FP8 form — MX activations carry a scale " @@ -2338,6 +3083,11 @@ def matmul_grouped( scatter_idx, bias=bias, ) + elif B.dtype in INT_WEIGHT_DTYPES: # packed integer weights + out = int_weight_only_matmul_grouped( + A, B, Bs, expert_start, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, output_dtype, + gather_idx, scatter_idx, b_zero=b_zero, bias=bias, b_perm=b_perm, a_pre_quant_scale=a_pre_quant_scale, a_scale=As, b_scale2=b_global_scale, + ) elif is_mx(B, Bs) and activation_format == "bf16": # weight-only: raw bf16 acts, MX weight upcast in-MMA assert As is None and a_global_scale is None and not quantize_output, ( "weight-only (activation_format='bf16') takes a raw activation, no As/global/requant" @@ -2381,7 +3131,7 @@ def matmul_grouped( a_global_scale, b_global_scale, output_global_scale, - bias=bias, + bias=bias, a_pre_quant_scale=a_pre_quant_scale, ) elif (block_size := weight_block_size(B, Bs)) is None: out = w8a8_tensor_dynamic_fp8_matmul_grouped( @@ -2401,6 +3151,7 @@ def matmul_grouped( gather_idx, scatter_idx, bias=bias, + activation_scale_ub=activation_scale_ub, ) else: out = w8a8_block_dynamic_fp8_matmul_grouped( diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/loading/scales.py b/finegrained-kernels/torch-ext/finegrained_kernels/loading/scales.py index 237f45f5d..927010241 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/loading/scales.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/loading/scales.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import torch import triton import triton.language as tl +from ..compat import compile_time_only_triton_wrap, device_context +from ..formats import block_fp8_weight_scale_grid, ue8m0_as_uint8 from ..quant import fp8_act_quant_inline, mx_act_quant_inline @@ -261,13 +264,27 @@ def load_mx_act_tile( return a, a_scale +@triton.jit +def apply_row_scale(acc, Scale, offs_rows, stride_rows, SWAP_AB: tl.constexpr = False): + """Fold a K-constant weight scale onto the fp32 accumulator after the K-loop: weight row ``n`` + reads ``Scale + n * stride_rows`` — a row-wise scale, or a per-tensor one broadcast by stride + 0. It lands on the acc's weight-row axis (axis 0 under ``SWAP_AB``).""" + s = tl.load(Scale + offs_rows * stride_rows) + if SWAP_AB: + out = acc * s[:, None] + else: + out = acc * s[None, :] + return out + + @triton.jit def apply_global_scale(acc, GlobalScale, stride_gs_e, GlobalScaleA, stride_gsa_e, expert_id): - """Fold the NVFP4 second-level globals onto the fp32 accumulator: the weight global ``g_b`` - and the activation global ``g_a``, each per tensor (stride 0, broadcast) or per expert. + """Apply output scales before bias/GLU: a weight global or a tile of per-channel weight + scale pointers, then the activation global. Expert strides select each expert's scales; + stride 0 broadcasts a shared scale. INT32 accumulators convert here, after the full K sum. ``None`` folds the arm out.""" if GlobalScale is not None: - acc = acc * tl.load(GlobalScale + expert_id * stride_gs_e).to(tl.float32) + acc = acc.to(tl.float32) * tl.load(GlobalScale + expert_id * stride_gs_e).to(tl.float32) if GlobalScaleA is not None: acc = acc * tl.load(GlobalScaleA + expert_id * stride_gsa_e).to(tl.float32) return acc @@ -323,3 +340,102 @@ def mx_2d_scale_ptrs( as_mask = offs_am_lin < M bs_mask = offs_bn_lin < (2 * N if GATE else N) return as_ptrs, bs_ptrs, as_mask, bs_mask + + +@triton.jit +def dequant_block_fp8_weight_tile( + w, Scales, expert, offs_n, offs_k, + stride_scale_e, stride_scale_n, stride_scale_k, + BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GATE_SCALES: tl.constexpr = False, + SWAP_AB: tl.constexpr = False, + DTYPE: tl.constexpr = tl.bfloat16, +): + """Rescale an FP8 tile to the raw activation dtype, retaining coarse block scales. + + Gate/up weight rows alternate per value; their scale rows alternate per block. + ``offs_k`` is a contiguous K tile starting at a multiple of its width. + """ + scale_n = (offs_n // 2 // BLOCK_N) * 2 + offs_n % 2 if GATE_SCALES else offs_n // BLOCK_N + tile_k: tl.constexpr = offs_k.shape[0] + first_k = tl.min(offs_k, 0) + tl.device_assert(first_k % tile_k == 0, "weight K tiles must start on their tile grid") + if BLOCK_K % tile_k == 0 or tile_k % BLOCK_K == 0: + # Broadcast compact K-group scales rather than staging a full FP32 tile. + groups: tl.constexpr = max(tile_k // BLOCK_K, 1) + ptrs = (Scales + expert * stride_scale_e + scale_n[:, None] * stride_scale_n + + (first_k // BLOCK_K + tl.arange(0, groups)[None, :]) * stride_scale_k) + scales = decode_group_scale(tl.load(ptrs)) + if SWAP_AB: + w = (w.to(tl.float32).reshape(offs_n.shape[0], groups, tile_k // groups) + * scales[:, :, None]).reshape(offs_n.shape[0], tile_k) + else: + w = (w.to(tl.float32).reshape(groups, tile_k // groups, offs_n.shape[0]) + * tl.trans(scales)[:, None, :]).reshape(tile_k, offs_n.shape[0]) + else: + ptrs = (Scales + expert * stride_scale_e + + scale_n[:, None] * stride_scale_n + + (offs_k[None, :] // BLOCK_K) * stride_scale_k) + scales = decode_group_scale(tl.load(ptrs)) + if not SWAP_AB: + scales = tl.trans(scales) + w = w.to(tl.float32) * scales + return w.to(DTYPE) + + +@triton.jit +def _dequant_block_fp8_weight_kernel( + B, Bs, Out, ExpertStart, + N: tl.constexpr, K: tl.constexpr, + stride_b_e, stride_b_n, stride_b_k, + stride_bs_e, stride_bs_n, stride_bs_k, + BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GATE_SCALES: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + expert = tl.program_id(1).to(tl.int64) + active = True + if ExpertStart is not None: + active = tl.load(ExpertStart + expert) < tl.load(ExpertStart + expert + 1) + if active: + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + n, k = offsets // K, offsets % K + # Full tiles avoid masked-load codegen; the final partial tile stays bounded. + mask = offsets < N * K if (N * K) % BLOCK_SIZE else tl.full((BLOCK_SIZE,), True, tl.int1) + scale_n = (n // 2 // BLOCK_N) * 2 + n % 2 if GATE_SCALES else n // BLOCK_N + weight = tl.load(B + expert * stride_b_e + n * stride_b_n + k * stride_b_k, mask, 0.0) + scale = tl.load( + Bs + expert * stride_bs_e + scale_n * stride_bs_n + (k // BLOCK_K) * stride_bs_k, + mask, 0.0, + ) + value = weight.to(tl.float32) * decode_group_scale(scale) + tl.store(Out + expert * N * K + offsets, value, mask) + + +def dequant_block_fp8_weight(B, Bs, dtype, *, gate=False, expert_start=None): + """Materialize FP8 weights for one prefill GEMM, rounding after the FP32 scale multiply. + + Grouped weights retain their expert indices; inactive expert planes are allocated + but untouched, and the subsequent grouped GEMM never reads them. + """ + assert dtype in (torch.bfloat16, torch.float16) + Bs, block_n, block_k = block_fp8_weight_scale_grid(B, Bs, gate) + Bs = ue8m0_as_uint8(Bs) + experts = B.shape[0] if B.ndim == 3 else 1 + if expert_start is not None: + assert (B.ndim == 3 and expert_start.is_contiguous() + and expert_start.numel() >= experts + 1 and expert_start.device == B.device) + rows, K = B.shape[-2:] + output = B.new_empty(B.shape, dtype=dtype) + with device_context(B.device): + compile_time_only_triton_wrap(_dequant_block_fp8_weight_kernel)[ + (triton.cdiv(rows * K, 2048), experts) + ]( + B, Bs, output, expert_start, rows, K, + B.stride(0) if B.ndim == 3 else 0, B.stride(-2), B.stride(-1), + Bs.stride(0) if Bs.ndim == 3 else 0, Bs.stride(-2), Bs.stride(-1), + BLOCK_N=block_n, BLOCK_K=block_k, + GATE_SCALES=gate and Bs.shape[-2] > 1, + BLOCK_SIZE=2048, + ) + return output diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/loading/tiles.py b/finegrained-kernels/torch-ext/finegrained_kernels/loading/tiles.py index 1c30b3ed6..288c59566 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/loading/tiles.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/loading/tiles.py @@ -15,7 +15,7 @@ import triton import triton.language as tl -from .scales import load_block_fp8_act_tile, load_mx_act_tile, load_swizzled_scale, load_swizzled_scale_tile, load_weight_scale_tile +from .scales import dequant_block_fp8_weight_tile, load_block_fp8_act_tile, load_mx_act_tile, load_swizzled_scale, load_swizzled_scale_tile, load_weight_scale_tile # A-sub-tile byte budget for the 1-byte-activation grouped-swizzle depth cap (see # swizzle_offsets): the co-scheduled rows' A tile per K-step (depth*BM*BK bytes) must stay @@ -34,14 +34,16 @@ def swizzle_offsets( BLOCK_SIZE_K: tl.constexpr, WEIGHT_VALUES_PER_BYTE: tl.constexpr = 1, GATE: tl.constexpr = False, + SINGLE_M_TILE: tl.constexpr = False, + tile_id=None, ): """2D-grid tile scheduling shared by the kernels below: grouped-swizzle the ``(pid_m, pid_n)`` program ids for L2 locality on B, then build the operand offset vectors. Returns ``(pid_m, pid_n, offs_am, offs_bn, offs_k)`` — the swizzled ids - (reused by the output store) and the ``%``-wrapped row/col offsets plus the K range. + (reused by the output store) and the bounded row/col offsets plus the K range. ``offs_bn`` is the WEIGHT-side row range: under ``GATE`` the doubled span - ``[2*pid_n*BN, 2*(pid_n+1)*BN)`` wrapped at ``2N``. The grid is sized by the OUTPUT tiles. + ``[2*pid_n*BN, 2*(pid_n+1)*BN)`` over ``2N`` weight rows. The grid is sized by the OUTPUT tiles. The swizzle keeps the B (weight) column-tile L2-hot while the co-scheduled rows reuse it, so the depth cap is set by the WEIGHT footprint, capped at ``min(num_pid_m, .)``. With @@ -54,22 +56,34 @@ def swizzle_offsets( activation-driven). Uses an EXPLICIT grouped swizzle (linearize the 2D program ids, then group), NOT ``tl.swizzle2d`` (which degrades with group depth, -3pp+ MFU at depth 32). Same grid launch, same result set — only the program-id -> tile mapping changes.""" - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - pid = tl.program_id(axis=1) * num_pid_m + tl.program_id(axis=0) - if WEIGHT_VALUES_PER_BYTE == 2: - max_group = num_pid_m + if SINGLE_M_TILE: + # Grouped ordering is the identity when every row fits in one tile. + pid_m = 0 + pid_n = tl.program_id(axis=1) else: - max_group = SWIZZLE_GROUP_A_BYTES // (BLOCK_SIZE_M * BLOCK_SIZE_K) - num_pid_in_group = max_group * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * max_group - group_size_m = min(num_pid_m - first_pid_m, max_group) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + pid = tl.program_id(axis=1) * num_pid_m + tl.program_id(axis=0) if tile_id is None else tile_id + if WEIGHT_VALUES_PER_BYTE == 2: + max_group = num_pid_m + else: + max_group = SWIZZLE_GROUP_A_BYTES // (BLOCK_SIZE_M * BLOCK_SIZE_K) + num_pid_in_group = max_group * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * max_group + group_size_m = min(num_pid_m - first_pid_m, max_group) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m WEIGHT_ROWS: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N - offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_bn = (pid_n * WEIGHT_ROWS + tl.arange(0, WEIGHT_ROWS)) % (2 * N if GATE else N) + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * WEIGHT_ROWS + tl.arange(0, WEIGHT_ROWS) + if SINGLE_M_TILE: + # Unused tail lanes can reuse the last valid index; output stores mask them. + offs_am = tl.minimum(offs_am, M - 1) + offs_bn = tl.minimum(offs_bn, (2 * N if GATE else N) - 1) + else: + offs_am %= M + offs_bn %= 2 * N if GATE else N offs_k = tl.arange(0, BLOCK_SIZE_K) return pid_m, pid_n, offs_am, offs_bn, offs_k @@ -155,8 +169,8 @@ def load_grouped_weight_tile( """One K-major (optionally gate|up-stacked) MX weight K-tile for the grouped / batched loop: the explicit-pointer tile flattened to the ``[KB, (2|1)*BN]`` rhs (or, under ``SWAP_AB``, the ``[(2|1)*BN, KB]`` rows-major lhs — the batched-decode orientation), or the ``[(2|1), BN, KB]`` - descriptor box over the ``(E, 2N|N, K_bytes)`` weight view, reshaped and transposed to the same - form (the fused-era TMA arm: natural orientation + per-iteration trans; grouped/2D never swap). + descriptor box over the ``(E, 2N|N, K_bytes)`` weight view, reshaped to ``[(2|1)*BN, KB]`` and, unless + ``SWAP_AB`` (rows-major, as the int4 tiles always are), transposed to the same form. ``B_DESCRIPTOR_IS_2D`` is the box's RANK, which ``B_MEMORY_MODE`` does not imply: it says the caller resolved a 2D box over ONE expert's slab (base already advanced by the expert stride) rather than a 3D box over the stack, so the expert axis is gone from both the shape and the @@ -166,15 +180,13 @@ def load_grouped_weight_tile( arm compiles; the caller advances ``w_ptrs`` and passes the box offsets either way.""" if B_MEMORY_MODE == "pointer": w = tl.load(w_ptrs) - elif B_DESCRIPTOR_IS_2D: - w = tl.trans(w_descriptor.load([n_off, kb_off])) else: - w = tl.trans( - tl.reshape( - w_descriptor.load([row0, n_off, kb_off]), - ((2 if GATE else 1) * BLOCK_SIZE_N, KB), - ) - ) + if B_DESCRIPTOR_IS_2D: + w = w_descriptor.load([n_off, kb_off]) + else: + w = tl.reshape(w_descriptor.load([row0, n_off, kb_off]), ((2 if GATE else 1) * BLOCK_SIZE_N, KB)) + if not SWAP_AB: + w = tl.trans(w) return w @@ -307,6 +319,33 @@ def load_act_plain( return a, a +@triton.jit +def gather_act_cols(a_rows, k_perm, k_off, row_mask, stride_a_k, BLOCK_SIZE_K: tl.constexpr): + """An activation K-tile with its columns gathered by an act-order K order: ``a_rows`` holds the tile's + ``[BM, 1]`` row pointers and column ``j`` reads ``k_perm[k_off + j]``. ``row_mask`` None loads maskless.""" + ptrs = a_rows + tl.load(k_perm + k_off + tl.arange(0, BLOCK_SIZE_K))[None, :] * stride_a_k + if row_mask is None: + a = tl.load(ptrs) + else: + a = tl.load(ptrs, mask=row_mask[:, None], other=0.0) + return a + + +@triton.jit +def load_act_k_permuted( + a_ptrs, a_descriptor, m_off, k_off, row_mask, gather_rows, a_next, + A_MEMORY_MODE: tl.constexpr, A_GATHER: tl.constexpr, +): + """``load_act_plain``'s activation K-tile, or, under an act-order K order (``a_next`` not None, pointer arm), the + gathered tile the previous K step issued (``gather_act_cols``): the caller issues the next step's gather + alongside, so the dependent gather's latency overlaps a step of compute.""" + if a_next is None: + a, _ = load_act_plain(a_ptrs, a_descriptor, m_off, k_off, row_mask, gather_rows, A_MEMORY_MODE, A_GATHER) + else: + a = a_next + return a + + @triton.jit def _weight_scale_mx( bs_ptrs, bs_mask, bs_descriptor, bs_ptr, blk_idx, expert_id, pid_n, k, N, K, @@ -461,6 +500,55 @@ def load_weight_mx( return w, w_s +@triton.jit +def load_weight_int( + b_ptrs, b_descriptor, bs_ptrs, bs_ptr, row0, n_off, k_off, expert_id, pid_n, k, N, K, + stride_bs_e, stride_bs_n, stride_bs_k, + GATE: tl.constexpr, PER_EXPERT: tl.constexpr, B_MEMORY_MODE: tl.constexpr, SWAP_AB: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, + k_col_off=0, GROUPED: tl.constexpr = False, BITS: tl.constexpr = 4, +): + """The integer weight path: a tile of uint32 words (``32 // BITS`` values each) + its affine group scale. + ``GROUPED`` reads the expert's tile out of the 3D stack (the grouped kernel's box or pointer arm).""" + w = _weight_value( + b_ptrs, b_descriptor, row0, n_off, k_off, GATE, GROUPED, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, + BLOCK_SIZE_K // (32 // BITS), + ) + w_s = load_weight_int_group( + bs_ptrs, bs_ptr, expert_id, pid_n, k, N, K, stride_bs_e, stride_bs_n, stride_bs_k, + GATE, PER_EXPERT, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, k_col_off=k_col_off, GROUPED=GROUPED, + ) + return w, w_s + + +@triton.jit +def int_group_ptrs(ptrs, k, stride_k, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr): + """Scale / zero-point pointers for K-tile ``k``. Groups no wider than the tile: ``ptrs`` already walks them (the + caller advances it ``BLOCK_SIZE_K // SCALE_GROUP_K`` columns a tile) and comes back as is. A group wider than the + tile holds it whole: ``ptrs`` stays at column 0 and tile ``k`` reads column ``k * BK // group``.""" + if SCALE_GROUP_K <= BLOCK_SIZE_K: + out = ptrs + else: + out = ptrs + ((k * BLOCK_SIZE_K) // SCALE_GROUP_K) * stride_k + return out + + +@triton.jit +def load_weight_int_group( + bz_ptrs, bz_ptr, expert_id, pid_n, k, N, K, stride_bz_e, stride_bz_n, stride_bz_k, + GATE: tl.constexpr, PER_EXPERT: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, + SCALE_GROUP_K: tl.constexpr, k_col_off=0, GROUPED: tl.constexpr = False, +): + """Scale or zero-point columns covering this tile, including tiles inside a wider group.""" + SCALE_COLS: tl.constexpr = max(1, BLOCK_SIZE_K // SCALE_GROUP_K) + if SCALE_GROUP_K > BLOCK_SIZE_K: + k = (k * BLOCK_SIZE_K) // SCALE_GROUP_K + return _weight_scale_mx( + bz_ptrs, None, None, bz_ptr, 0, expert_id, pid_n, k, N, K, stride_bz_e, stride_bz_n, stride_bz_k, + GROUPED, GATE, PER_EXPERT, False, BLOCK_SIZE_N, SCALE_COLS, SCALE_GROUP_K, k_col_off=k_col_off, + ) + + @triton.jit def load_weight_block_dynamic( b_ptrs, b_descriptor, bs_ptrs, bs_mask, up_s_ptr, row0, n_off, k_off, k, stride_bs_k, @@ -475,6 +563,30 @@ def load_weight_block_dynamic( return w, w_s +@triton.jit +def load_weight_block_fp8( + b_ptrs, b_descriptor, bs_ptr, row0, n_off, k_off, + expert_id, offs_n, offs_k, stride_bs_e, stride_bs_n, stride_bs_k, + GATE: tl.constexpr, GROUPED: tl.constexpr, B_MEMORY_MODE: tl.constexpr, SWAP_AB: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, + BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GATE_SCALES: tl.constexpr, DTYPE: tl.constexpr, + B_DESCRIPTOR_IS_2D: tl.constexpr = False, +): + """Coarse FP8 weight-only tile, rescaled to the raw activation dtype.""" + w = _weight_value( + b_ptrs, b_descriptor, row0, n_off, k_off, + GATE, GROUPED, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K, + B_DESCRIPTOR_IS_2D, + ) + w = dequant_block_fp8_weight_tile( + w, bs_ptr, expert_id, offs_n, offs_k, + stride_bs_e, stride_bs_n, stride_bs_k, + BLOCK_N, BLOCK_K, GATE_SCALES, SWAP_AB, DTYPE, + ) + return w + + @triton.jit def load_weight_static( b_ptrs, b_descriptor, bs_ptrs, bs_mask, up_s_ptr, row0, n_off, k_off, k, stride_bs_k, diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py b/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py index f44d3b80c..7c43ef52f 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py @@ -14,21 +14,23 @@ import os +from typing import List import torch import triton import triton.language as tl +from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait from triton.tools.tensor_descriptor import TensorDescriptor from .bayesian_autotuner import bayesian_autotune from .bench import weight_input -from .compat import add_op_namespace_prefix, FP8_DTYPE, is_sm10x, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, tl_dtype -from .descriptors import maybe_descriptor, rebind_bd_descriptors, rebind_mx_descriptors, rebind_weight_only_descriptors -from .formats import check_activation_format, global_scale_stride, normalize_global_scale, e2m1_as_uint8, is_mx, mx_scale_family, resolve_activation_format, resolve_output_dtype, ue8m0_as_uint8, validate_dense_2d_operands, weight_format +from .compat import add_op_namespace_prefix, FP8_DTYPE, is_sm10x, launches_skipped, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, pdl_launch_kwargs, tl_dtype, MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, narrow_dot_autotuning_configs, narrow_scalar_autotuning_configs, split_k_autotuning_configs, persistent_program_count +from .descriptors import maybe_descriptor, rebind_bd_descriptors, rebind_mx_descriptors, rebind_tiled_descriptors, rebind_weight_only_descriptors, int_weight_only_pre_hook +from .formats import FP8_WEIGHT_ONLY_DEQUANT_MIN_M, block_fp8_weight_scale_grid, _validate_quantization_args, global_scale_stride, normalize_global_scale, e2m1_as_uint8, is_mx, is_rowwise, mx_scale_family, resolve_activation_format, resolve_output_dtype, row_scale_stride, ue8m0_as_uint8, validate_dense_2d_operands, weight_format, validate_int_operands, validate_svdquant_operands, INT_WEIGHT_DTYPES from .swizzle import swizzle_mx_scales, swizzled_scale_descriptor -from .quant import MX_ACT_QUANT, fp8_act_quant_block_dynamic, maybe_act_quant, mxfp8_act_quant, nvfp4_act_quant, quantize_rows_static, tensor_wide_act_operands -from .loading.scales import apply_global_scale, mx_2d_scale_ptrs -from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot +from .quant import MX_ACT_QUANT, fp8_act_quant_block_dynamic, maybe_act_quant, mxfp8_act_quant, nvfp4_act_quant, quantize_rows_static, tensor_wide_act_operands, prepare_fp8_mx_inputs, dequantize_int_weight, permute_act_cols, svdquant_act_quant +from .loading.scales import apply_row_scale, dequant_block_fp8_weight, mx_2d_scale_ptrs +from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, plain_dot, static_dot, int_weight_only_compute from .loading.tiles import ( advance_ptrs, load_act_block_dynamic, @@ -36,6 +38,10 @@ load_act_plain, load_act_static, load_weight_block_dynamic, + load_weight_block_fp8, + load_weight_int, + load_weight_int_group, + int_group_ptrs, load_weight_mx, load_weight_plain, load_weight_static, @@ -44,8 +50,8 @@ operand_tile_ptrs, swizzle_offsets, ) -from .epilogue import acc_init, add_bias, bias_strides, gemm_epilogue, store_masked, store_masked_oriented -from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, gate_stacked_tmem_trap_pruner, block_dynamic_2d_warp_spec_pruner, block_dynamic_mma_width_pruner, block_within_dim_pruner, compose_pruners, descriptor_box_pruner, descriptor_needs_prequant_pruner, matched_memory_modes_pruner, mx_2d_swap_scope_pruner, mx_config_pruner, scalar_max_m_pruner, smem_pruner, swizzled_out_bm_pruner, swizzled_scale_config_pruner, warp_spec_compile_guard_pruner +from .epilogue import acc_finalize, acc_init, add_bias, bias_strides, gemm_epilogue, store_masked, store_masked_oriented +from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, gate_stacked_tmem_trap_pruner, block_dynamic_2d_warp_spec_pruner, block_dynamic_mma_width_pruner, block_within_dim_pruner, compose_pruners, descriptor_box_pruner, descriptor_needs_prequant_pruner, matched_memory_modes_pruner, mx_2d_swap_scope_pruner, native_mma_m_pruner, mx_config_pruner, scalar_max_m_pruner, smem_pruner, swap_decode_scope_pruner, swap_max_m_pruner, swizzled_out_bm_pruner, swizzled_scale_config_pruner, warp_spec_compile_guard_pruner, warp_spec_memory_mode_pruner, mx_auxiliary_pruner, block_m_cap_pruner, int_2d_warp_spec_pruner, int_split_k_pruner, scalar_block_m_pruner, int_scale_group_tile_pruner # The 2D-grid kernels' L2-locality swizzle depth is derived per-tile inside # ``swizzle_offsets`` (see SWIZZLE_GROUP_A_BYTES there) — no per-kernel constant here. @@ -60,6 +66,8 @@ # inline either way. MX_MATMUL_ACT_PREQUANT_MIN_M = 16 STATIC_MATMUL_ACT_PREQUANT_MIN_M = 16 +# the largest M tile of the tensor-wide / row-wise 2D kernel; its derived activations are padded to it +TENSOR_MATMUL_A_ROW_MULTIPLE = 128 # Block-dynamic has NO gate: offline wins at EVERY M incl. M=1 (24.7 vs 31.0us, isolated # arm A/B, graph-timed, H=6144 b128) — the inline arm pays a per-tile fp32 amax+div. The # always-offline paths (bd 2D/batched/grouped, fused gate_ups, MX grouped) quantize @@ -278,9 +286,17 @@ def w8a8_block_dynamic_fp8_matmul_kernel( @bayesian_autotune( - # tune_block_m mirrors the block-dynamic kernel (adaptive BM=128 cost it 17% at prefill). + # tune_block_m mirrors the block-dynamic kernel (adaptive BM=128 cost it 17% at prefill), and + # so do the two memory regimes: pointer loads for decode / short K, host-TMA on both operands + # to feed the big prefill tiles. The K tile is free here (no scale block pins it). get_accelerator_autotuning_configs( - tune_block_nk=True, warp_spec=True, tune_block_m=True + tune_block_nk=True, + warp_spec=True, + tune_block_m=True, + swap_ab=True, + a_memory_modes=("descriptor", "pointer"), + b_memory_modes=("descriptor", "pointer"), + pre_hook=rebind_tiled_descriptors, ), # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide, a distinct config space). ["N", "K", "m_bit_length", "GATE"], @@ -296,15 +312,23 @@ def w8a8_block_dynamic_fp8_matmul_kernel( fp8_dot_warp_pruner(), block_within_dim_pruner("K"), warp_spec_compile_guard_pruner(), + warp_spec_memory_mode_pruner(), + matched_memory_modes_pruner(), + swap_decode_scope_pruner(), + native_mma_m_pruner(), + descriptor_box_pruner(), + smem_pruner(), ) }, ) @triton.jit def w8a8_tensor_dynamic_fp8_matmul_kernel( A, # (M, K) pre-quantized FP8 activations + ADescriptor, # host TMA descriptor over A (M, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer" As, # (M,) per-token activation scales, or one calibrated (static) scale (stride_as_m 0) B, # (N, K) FP8 weights - Bs, # scalar/(1,) per-tensor weight scale + BDescriptor, # host TMA descriptor over B (N, K), box ((2|1)*BN, BK); read iff B_MEMORY_MODE != "pointer" + Bs, # (N, 1) row-wise weight scales, or one per-tensor scale (stride_bs_n 0) C, # (M, N) output Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None # Shape @@ -318,6 +342,7 @@ def w8a8_tensor_dynamic_fp8_matmul_kernel( stride_as_m, stride_b_k, stride_b_n, + stride_bs_n, stride_c_m, stride_c_n, stride_bias_e, @@ -327,6 +352,9 @@ def w8a8_tensor_dynamic_fp8_matmul_kernel( BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, WARP_SPEC: tl.constexpr = False, + A_MEMORY_MODE: tl.constexpr = "pointer", + B_MEMORY_MODE: tl.constexpr = "pointer", + SWAP_AB: tl.constexpr = False, # Dense gate|up fusion: B is the (2N, K) gate|up stack, C the [M, N] GLU output. Every arm # folds out at compile time when GATE=False (the plain dense GEMM, unchanged). GATE: tl.constexpr = False, @@ -335,41 +363,58 @@ def w8a8_tensor_dynamic_fp8_matmul_kernel( SWIGLU_LIMIT: tl.constexpr = None, SIMULATE_UNFUSED: tl.constexpr = False, INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, + PDL: tl.constexpr = False, ): - """Tensor-scale FP8 GEMM kernel. + """Tensor-scale / row-wise FP8 GEMM kernel. - Computes ``C = A @ B.T`` with one activation scale per row and one - weight scale for the full matrix. + Computes ``C = A @ B.T`` with one activation scale per row and a weight scale constant + along K — one for the full matrix, or one per output row — both folded after the loop. Uses a 2D grid with swizzle for L2 cache locality on B tiles. + + ``SWAP_AB`` (decode, BM=1): the weight's output rows fill the MMA M dim and the single token + is padded to the N=16 atom; the accumulator collapses back to ``[1, (2|1)*BN]`` before the + scales, so everything after the loop is the unswapped path. """ + if PDL: + gdc_wait() n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets( M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GATE=GATE ) - a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, "pointer", True) - # Under GATE the weight is the stacked (2N, K) gate|up tile; one per-tensor weight scale - # covers both projections, applied (with the per-row act scale) before the GLU split. - b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, "pointer") + a_descriptor = operand_tile_descriptor( + ADescriptor, A, M, K, stride_a_m, stride_a_k, BLOCK_SIZE_M, BLOCK_SIZE_K, A_MEMORY_MODE + ) + b_descriptor = operand_tile_descriptor( + BDescriptor, B, N, K, stride_b_n, stride_b_k, BLOCK_SIZE_N, BLOCK_SIZE_K, B_MEMORY_MODE, GATE, + ) + a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, not SWAP_AB) + # Under GATE the weight is the stacked (2N, K) gate|up tile; its scales (per tensor, or per + # stacked row) are applied with the per-row act scale before the GLU split. + b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, B_MEMORY_MODE, SWAP_AB) - a_s = tl.load(As + offs_am * stride_as_m) - b_s = tl.load(Bs) + a_s = tl.load(As + offs_am * stride_as_m, mask=offs_am < M, other=0.0) # Accumulate raw dot products, apply scales once after the loop. - accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False) + accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB) for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC): - a, _as = load_act_plain(a_ptrs, a_ptrs, 0, 0, None, 0, "pointer", False) + a, _as = load_act_plain( + a_ptrs, a_descriptor, pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K, None, 0, A_MEMORY_MODE, False + ) b, _bs = load_weight_plain( - b_ptrs, b_ptrs, 0, pid_n * n_width, k * BLOCK_SIZE_K, - GATE, False, "pointer", False, BLOCK_SIZE_N, BLOCK_SIZE_K, + b_ptrs, b_descriptor, 0, pid_n * n_width, k * BLOCK_SIZE_K, + GATE, False, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K, ) - accumulator = accumulator + fp8_dot(a, b, False, BLOCK_SIZE_K) + accumulator = accumulator + fp8_dot(a, b, SWAP_AB, BLOCK_SIZE_K) a_ptrs, _asp, b_ptrs, _bsp, _, _ = advance_ptrs( a_ptrs, a_ptrs, b_ptrs, b_ptrs, b_ptrs, b_ptrs, BLOCK_SIZE_K * stride_a_k, 0, BLOCK_SIZE_K * stride_b_k, 0, - "pointer", "pointer", False, False, False, + A_MEMORY_MODE, B_MEMORY_MODE, False, False, False, ) - accumulator = accumulator * a_s[:, None] * b_s + accumulator = acc_finalize(accumulator, "dot", n_width, SWAP_AB) + accumulator = apply_row_scale(accumulator * a_s[:, None], Bs, offs_bn, stride_bs_n) + if PDL: + gdc_launch_dependents() if GATE: out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) @@ -554,17 +599,12 @@ def w8a8_block_static_fp8_matmul_kernel( ) -@bayesian_autotune( - # tune_block_m: BLOCK_SIZE_M becomes a config axis, so the tuner sizes the M tile - # per workload — small at decode, large at prefill. - # scalar is in the mode set for M=1 decode (attn projection): it avoids the MMA M->16 pad - # that held MXFP8 attn decode 35% over block-dynamic at identical weight bytes; fp4-scalar - # is dropped by the pruner (ALU-bound unpack). swap_ab emits the BM=1 swap decode rows, - # scoped by mx_2d_swap_scope_pruner to E4M3-scale (NVFP4) decode — its ONLY native M=1 - # arm (dot_scaled has no sub-M=128 E4M3 form; scalar is unpack-bound: 60-142µs vs swap's - # 22-50µs on the GLM attn shape, 2026-08-05). UE8M0 formats keep swap OFF: an 18-cell - # forced-swap sweep (cudagraph, M3+dsv4 decode) showed swap losing there (M3 attn −38%). - get_accelerator_autotuning_configs( +SPLIT_K_VALUES = (2, 4, 8) +SPLIT_K_MAX_M = 16 + + +def _mx_matmul_configs(): + configs = get_accelerator_autotuning_configs( mx=True, tune_block_nk=True, compute_modes=("dot_scaled", "dot", "scalar"), @@ -579,7 +619,30 @@ def w8a8_block_static_fp8_matmul_kernel( # WS tuner axis: dot_scaled + WS + TMA compiles + wins at prefill (1918 -> ~1995 @M=4k); # warp_spec_compile_guard_pruner keeps it to num_warps%4==0, BM>=64. warp_spec=True, - ), + ) + for c in configs: + c.kwargs["SPLIT_K"] = 1 + c.kwargs["PERSISTENT"] = False + native = [c for c in configs if c.kwargs["COMPUTE_MODE"] == "dot_scaled" + and c.kwargs["BLOCK_SIZE_M"] == 128 and c.kwargs["BLOCK_SIZE_N"] >= 128 + and c.kwargs["A_MEMORY_MODE"] == c.kwargs["B_MEMORY_MODE"] == "host_descriptor" + and not c.kwargs["SWAP_AB"] and c.kwargs["WARP_SPEC"] and c.num_warps in (4, 8)] + persistent = [triton.Config(dict(c.kwargs, PERSISTENT=True), num_warps=c.num_warps, + num_stages=c.num_stages, pre_hook=c.pre_hook) for c in native] + return configs + split_k_autotuning_configs(SPLIT_K_VALUES, configs=native) + persistent + + +@bayesian_autotune( + # tune_block_m: BLOCK_SIZE_M becomes a config axis, so the tuner sizes the M tile + # per workload — small at decode, large at prefill. + # scalar is in the mode set for M=1 decode (attn projection): it avoids the MMA M->16 pad + # that held MXFP8 attn decode 35% over block-dynamic at identical weight bytes; fp4-scalar + # is dropped by the pruner (ALU-bound unpack). swap_ab emits the BM=1 swap decode rows, + # scoped by mx_2d_swap_scope_pruner to E4M3-scale (NVFP4) decode — its ONLY native M=1 + # arm (dot_scaled has no sub-M=128 E4M3 form; scalar is unpack-bound: 60-142µs vs swap's + # 22-50µs on the GLM attn shape, 2026-08-05). UE8M0 formats keep swap OFF: an 18-cell + # forced-swap sweep (cudagraph, M3+dsv4 decode) showed swap losing there (M3 attn −38%). + _mx_matmul_configs(), # the MXFP4/MXFP8 split keys itself — the tuner appends every tensor arg's dtype to # its cache key (memory and disk); # m_bit_length (log2 M bucket) keys the M tile — the winner keeps shifting with M well past the @@ -590,10 +653,10 @@ def w8a8_block_static_fp8_matmul_kernel( # split W4A8 from W4A4 itself. OUTPUT_FORMAT/SWIZZLED_OUT key the requant epilogue # explicitly (today the C/Cs dtypes split them incidentally; a descriptor Cs has no # dtype token, so don't lean on the append). - ["N", "K", "m_bit_length", "ACTIVATION_FORMAT", "SWIZZLED_SCALES", "GATE", "OUTPUT_FORMAT", "SWIZZLED_OUT"], + ["N", "K", "m_bit_length", "ACTIVATION_FORMAT", "SWIZZLED_SCALES", "GATE", "OUTPUT_FORMAT", "SWIZZLED_OUT", "LORA_RANK", "LORA_SPLIT_K"], n_trials=100, benchmark_input=weight_input, - path_anchor_axes=PATH_ANCHOR_AXES, + path_anchor_axes=(*PATH_ANCHOR_AXES, "SPLIT_K", "PERSISTENT"), # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), # BK-within-K veto (the loop loads are unmasked) + the sm_10x dot_scaled shape guards @@ -605,6 +668,7 @@ def w8a8_block_static_fp8_matmul_kernel( gate_stacked_tmem_trap_pruner(), swizzled_out_bm_pruner(), mx_config_pruner("K"), + mx_auxiliary_pruner(), mx_2d_swap_scope_pruner(), scalar_max_m_pruner("M"), smem_pruner(), @@ -620,10 +684,14 @@ def mx_dynamic_matmul_kernel( ADescriptor, # host TMA descriptor over A (M, K_bytes), box (BM, BK_bytes); read iff A_MEMORY_MODE != "pointer" As, # activation scales: SWIZZLE_32_4_4 (swizzled arm — bulk via ASDescriptor at BM=128, or this pointer for the BM<128 gather) or (M, K // 32) UE8M0 affine ASDescriptor, # host TMA descriptor over the SWIZZLE_32_4_4 A scales (BM=128 bulk load); read iff SWIZZLED_SCALES + LoraInput, # BF16 down projection or decode FP32 partials; None without SVDQuant B, # (N, K) E4M3 (MXFP8) or (N, K // 2) packed E2M1 (MXFP4) weights BDescriptor, # host TMA descriptor over B (N, K_bytes), box (BN, BK_bytes); read iff B_MEMORY_MODE != "pointer" Bs, # weight scales: SWIZZLE_32_4_4 (swizzled arm — bulk via BSDescriptor at BN=128, or this pointer for the BN<128 gather) or (N, K // SCALE_GROUP_K) affine BSDescriptor, # host TMA descriptor over the SWIZZLE_32_4_4 B scales (BN=128 bulk load); read iff SWIZZLED_SCALES + LoraWeight, # (N_out, rank) BF16 correction weight + Partials, + TileCounters, C, # (M, N) output (the GLU intermediate under GATE) Cs, # (M, N // group) row-major requant output scale — written iff OUTPUT_FORMAT and not SWIZZLED_OUT Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None @@ -643,10 +711,14 @@ def mx_dynamic_matmul_kernel( stride_as_k, # activation-scale group-axis stride, in GROUPS (see the affine advance below): # 1 for the row-major (M, K // SCALE_GROUP_K) grid, a real stride under a transposed view. # Dead on the swizzled arm, which re-derives its block offsets instead of walking pointers. + stride_li_m, + stride_li_r, stride_b_k, stride_b_n, stride_bs_k, stride_bs_n, + stride_lw_n, + stride_lw_r, stride_c_m, stride_c_n, stride_cs_m, # requant output-scale strides (dead unless OUTPUT_FORMAT) @@ -661,6 +733,10 @@ def mx_dynamic_matmul_kernel( BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, + LORA_RANK: tl.constexpr, + LORA_SPLIT_K: tl.constexpr, + SPLIT_K: tl.constexpr, + PERSISTENT: tl.constexpr, COMPUTE_MODE: tl.constexpr, A_MEMORY_MODE: tl.constexpr = "pointer", B_MEMORY_MODE: tl.constexpr = "pointer", @@ -692,86 +768,111 @@ def mx_dynamic_matmul_kernel( rescale (wins at decode; FP4 unpacks E2M1->E4M3 first — lossless). 2D grid with swizzle for L2 reuse. """ - # uint8 A = packed-E2M1 activations (W4A4 — the dtype IS the activation format) - ACT_VALUES_PER_BYTE: tl.constexpr = 2 if A.dtype.element_ty == tl.uint8 else 1 - WEIGHT_VALUES_PER_BYTE: tl.constexpr = 2 if B.dtype.element_ty == tl.uint8 else 1 - n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N - SCALE_COLS: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K - offs_ka = tl.arange(0, BLOCK_SIZE_K // ACT_VALUES_PER_BYTE) - offs_kb = tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) - # packed-fp4 weights (WEIGHT_VALUES_PER_BYTE==2) route swizzle_offsets to full L2 grouping - pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets( - M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, WEIGHT_VALUES_PER_BYTE, GATE=GATE - ) - # Scales read affine off row/col offsets + bounds masks (the %-wrapped operand offsets would - # make the scale load a non-affine gather); the SWIZZLED arm reads them via the SA/BS - # descriptors in the loop instead, so its pointer tiles are dead (base scalars + null masks). - as_ptrs, bs_ptrs, as_mask, bs_mask = mx_2d_scale_ptrs( - As, Bs, pid_m, pid_n, M, N, stride_as_m, stride_bs_n, stride_bs_k, - BLOCK_SIZE_M, BLOCK_SIZE_N, SCALE_COLS, SWIZZLED_SCALES, GATE, - ) - # Operand tiles: activation + the (GATE-stacked-aware) weight, both single leaf calls (the - # descriptor-vs-pointer and gate|up-stack branches fold inside the leaves). The swap arm - # loads both naturally — mx_compute reorients at the dot under SWAP_AB, flattening the token tile. - if SWAP_AB: - tl.static_assert( - BLOCK_SIZE_M == 1, "the mx swap arm flattens the single-token tile" - ) - a_ptrs = operand_tile_ptrs(A, offs_am, offs_ka, stride_a_m, stride_a_k, A_MEMORY_MODE, True) - b_ptrs = matmul_weight_ptrs( - B, offs_bn, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE, SWAP_AB - ) + num_m_tiles = tl.cdiv(M, BLOCK_SIZE_M) + if PERSISTENT: + tile_start, tile_end, tile_step = tl.program_id(0), num_m_tiles * tl.cdiv(N, BLOCK_SIZE_N), tl.num_programs(0) + else: + tile_start = tl.program_id(1) * num_m_tiles + tl.program_id(0) + tile_end, tile_step = tile_start + 1, 1 + for tile_id in range(tile_start, tile_end, tile_step): + # uint8 A = packed-E2M1 activations (W4A4 — the dtype IS the activation format) + ACT_VALUES_PER_BYTE: tl.constexpr = 2 if A.dtype.element_ty == tl.uint8 else 1 + WEIGHT_VALUES_PER_BYTE: tl.constexpr = 2 if B.dtype.element_ty == tl.uint8 else 1 + n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N + SCALE_COLS: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K + offs_ka = tl.arange(0, BLOCK_SIZE_K // ACT_VALUES_PER_BYTE) + offs_kb = tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) + # packed-fp4 weights (WEIGHT_VALUES_PER_BYTE==2) route swizzle_offsets to full L2 grouping + pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets( + M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, WEIGHT_VALUES_PER_BYTE, GATE=GATE, tile_id=tile_id + ) + # Scales read affine off row/col offsets + bounds masks (the %-wrapped operand offsets would + # make the scale load a non-affine gather); the SWIZZLED arm reads them via the SA/BS + # descriptors in the loop instead, so its pointer tiles are dead (base scalars + null masks). + as_ptrs, bs_ptrs, as_mask, bs_mask = mx_2d_scale_ptrs( + As, Bs, pid_m, pid_n, M, N, stride_as_m, stride_bs_n, stride_bs_k, + BLOCK_SIZE_M, BLOCK_SIZE_N, SCALE_COLS, SWIZZLED_SCALES, GATE, + ) + # Operand tiles: activation + the (GATE-stacked-aware) weight, both single leaf calls (the + # descriptor-vs-pointer and gate|up-stack branches fold inside the leaves). The swap arm + # loads both naturally — mx_compute reorients at the dot under SWAP_AB, flattening the token tile. + if SWAP_AB: + tl.static_assert( + BLOCK_SIZE_M == 1, "the mx swap arm flattens the single-token tile" + ) + a_ptrs = operand_tile_ptrs(A, offs_am, offs_ka, stride_a_m, stride_a_k, A_MEMORY_MODE, True) + b_ptrs = matmul_weight_ptrs( + B, offs_bn, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE, SWAP_AB + ) + + split = tl.program_id(2) if SPLIT_K > 1 else 0 + a_ptrs += split * (BLOCK_SIZE_K // ACT_VALUES_PER_BYTE) * stride_a_k + b_ptrs += split * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) * stride_b_k + if not SWIZZLED_SCALES: + as_ptrs += split * (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_as_k + bs_ptrs += split * (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_bs_k + accumulator = acc_init(COMPUTE_MODE if SWAP_AB else "dot", BLOCK_SIZE_M, n_width, SWAP_AB) + for k in tl.range(split, tl.cdiv(K, BLOCK_SIZE_K), SPLIT_K, warp_specialize=WARP_SPEC): + a, a_s = load_act_mx( + a_ptrs, as_ptrs, AsGlobal, 0, None, as_mask, ADescriptor, pid_m * BLOCK_SIZE_M, + k * (BLOCK_SIZE_K // ACT_VALUES_PER_BYTE), ASDescriptor, As, 0, 0, pid_m, k, M, K, + A_MEMORY_MODE, False, False, SWIZZLED_SCALES, + BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K, ACTIVATION_FORMAT, + ) + b, b_s = load_weight_mx( + b_ptrs, BDescriptor, bs_ptrs, bs_mask, BSDescriptor, Bs, 0, pid_n * n_width, + k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE), 0, 0, pid_n, k, N, K, + 0, stride_bs_n, stride_bs_k, + GATE, False, False, B_MEMORY_MODE, SWAP_AB, SWIZZLED_SCALES, + BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, WEIGHT_VALUES_PER_BYTE, + ) + accumulator = mx_compute( + accumulator, a, a_s, b, b_s, COMPUTE_MODE, + n_width, BLOCK_SIZE_K, SCALE_GROUP_K, SWAP_AB, + ) + a_ptrs, as_ptrs, b_ptrs, bs_ptrs, _, _ = advance_ptrs( + a_ptrs, as_ptrs, b_ptrs, bs_ptrs, b_ptrs, bs_ptrs, + SPLIT_K * (BLOCK_SIZE_K // ACT_VALUES_PER_BYTE) * stride_a_k, + SPLIT_K * (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_as_k, + SPLIT_K * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) * stride_b_k, + SPLIT_K * (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_bs_k, + A_MEMORY_MODE, B_MEMORY_MODE, + not SWIZZLED_SCALES, (not SWIZZLED_SCALES) and not GATE, False, + ) - accumulator = acc_init(COMPUTE_MODE if SWAP_AB else "dot", BLOCK_SIZE_M, n_width, SWAP_AB) - for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC): - a, a_s = load_act_mx( - a_ptrs, as_ptrs, AsGlobal, 0, None, as_mask, ADescriptor, pid_m * BLOCK_SIZE_M, - k * (BLOCK_SIZE_K // ACT_VALUES_PER_BYTE), ASDescriptor, As, 0, 0, pid_m, k, M, K, - A_MEMORY_MODE, False, False, SWIZZLED_SCALES, - BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K, ACTIVATION_FORMAT, - ) - b, b_s = load_weight_mx( - b_ptrs, BDescriptor, bs_ptrs, bs_mask, BSDescriptor, Bs, 0, pid_n * n_width, - k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE), 0, 0, pid_n, k, N, K, - 0, stride_bs_n, stride_bs_k, - GATE, False, False, B_MEMORY_MODE, SWAP_AB, SWIZZLED_SCALES, - BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, WEIGHT_VALUES_PER_BYTE, - ) - accumulator = mx_compute( - accumulator, a, a_s, b, b_s, COMPUTE_MODE, - n_width, BLOCK_SIZE_K, SCALE_GROUP_K, SWAP_AB, - ) - a_ptrs, as_ptrs, b_ptrs, bs_ptrs, _, _ = advance_ptrs( - a_ptrs, as_ptrs, b_ptrs, bs_ptrs, b_ptrs, bs_ptrs, - (BLOCK_SIZE_K // ACT_VALUES_PER_BYTE) * stride_a_k, - (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_as_k, - (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) * stride_b_k, - (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_bs_k, - A_MEMORY_MODE, B_MEMORY_MODE, - not SWIZZLED_SCALES, (not SWIZZLED_SCALES) and not GATE, False, - ) - - # NVFP4 two-level: block e4m3 scales rode through dot_scaled; recover the combined per-tensor - # one epilogue for both arms: GATE splits+SwiGLUs (+ optional requant), plain casts+stores. - # affine output rows (swizzle's offs_am is %-wrapped — the store scatters to real rows); - # N_COLS masks the 2D-dense column tail. - out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - gemm_epilogue( - C, Cs, accumulator, out_row, pid_n, pid_m, out_row < M, - stride_c_m, stride_c_n, stride_cs_m, stride_cs_n, - BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, OUTPUT_FORMAT, SCALE_GROUP_K, - ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, - COMPUTE_MODE=COMPUTE_MODE, SWAP_AB=SWAP_AB, N_COLS=N, SWIZZLED_OUT=SWIZZLED_OUT, - CSDescriptor=CSDescriptor, CsGlobal=CsGlobal, stride_cs_global=stride_cs_global, - GlobalScale=AsBsGlobal, - stride_global_e=stride_bs_global, - GlobalScaleA=AsGlobal, - stride_global_a=stride_as_global, - global_row=0, - Bias=Bias, - stride_bias_e=stride_bias_e, - stride_bias_n=stride_bias_n, - ) + if SPLIT_K > 1: + # Decode's padded MMA rows are identical; reduce only the real row. + first = tl.gather(accumulator, tl.full((1, n_width), 0, tl.int32), 0) + accumulator, last = split_k_reduce(first, Partials, TileCounters, + tl.program_id(1), tl.num_programs(1), 1, n_width, SPLIT_K) + else: + last = True + STORE_M: tl.constexpr = 1 if SPLIT_K > 1 else BLOCK_SIZE_M + if last: + # NVFP4 two-level: block e4m3 scales rode through dot_scaled; recover the combined per-tensor + # one epilogue for both arms: GATE splits+SwiGLUs (+ optional requant), plain casts+stores. + # affine output rows (swizzle's offs_am is %-wrapped — the store scatters to real rows); + # N_COLS masks the 2D-dense column tail. + out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, STORE_M) + gemm_epilogue( + C, Cs, accumulator, out_row, pid_n, pid_m, out_row < M, + stride_c_m, stride_c_n, stride_cs_m, stride_cs_n, + STORE_M, BLOCK_SIZE_N, GATE, OUTPUT_FORMAT, SCALE_GROUP_K, + ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, + COMPUTE_MODE=COMPUTE_MODE, SWAP_AB=False if SPLIT_K > 1 else SWAP_AB, N_COLS=N, SWIZZLED_OUT=SWIZZLED_OUT, + CSDescriptor=CSDescriptor, CsGlobal=CsGlobal, stride_cs_global=stride_cs_global, + GlobalScale=AsBsGlobal, + stride_global_e=stride_bs_global, + GlobalScaleA=AsGlobal, + stride_global_a=stride_as_global, + global_row=0, + Bias=Bias, + stride_bias_e=stride_bias_e, + stride_bias_n=stride_bias_n, + LoraInput=LoraInput, LoraWeight=LoraWeight, + stride_li_m=stride_li_m, stride_li_r=stride_li_r, + stride_lw_n=stride_lw_n, stride_lw_r=stride_lw_r, LORA_RANK=LORA_RANK, LORA_SPLIT_K=LORA_SPLIT_K, + ) @bayesian_autotune( @@ -783,6 +884,7 @@ def mx_dynamic_matmul_kernel( # groups (the dequant is in-tile), so no BK==group constraint. get_accelerator_autotuning_configs( tune_block_nk=True, + swap_ab_prefill=True, tune_block_m=True, warp_spec=True, compute_modes=("dot", "dot_scaled"), @@ -803,6 +905,7 @@ def mx_dynamic_matmul_kernel( "early_config_prune": compose_pruners( block_within_dim_pruner("K"), mx_config_pruner("K", "N"), # dot_scaled shape gates + gate_stacked_tmem_trap_pruner(), warp_spec_compile_guard_pruner(), descriptor_box_pruner("BLOCK_SIZE_K"), smem_pruner("BLOCK_SIZE_K"), @@ -813,6 +916,7 @@ def mx_dynamic_matmul_kernel( def mx_weight_only_matmul_2d_kernel( A, # (M, K) raw BF16/FP16 activations — NOT quantized ADescriptor, # host TMA descriptor over A (M, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer" + AsGlobal, B, # (N, K[/2]) MXFP4/NVFP4 (packed E2M1) / MXFP8 (E4M3) weights; under GATE the (2N, ...) gate|up stack BDescriptor, # host TMA descriptor over B (N, K_bytes), box (BN, BK_bytes); read iff B_MEMORY_MODE != "pointer" Bs, # (N, K // SCALE_GROUP_K) group scales — UE8M0 (MX) or E4M3 (NVFP4) @@ -825,15 +929,16 @@ def mx_weight_only_matmul_2d_kernel( m_bit_length, # autotune key only (log2 M bucket) stride_a_m, stride_a_k, + stride_as_global: tl.constexpr, stride_b_n, stride_b_k, stride_bs_n, stride_bs_k, + stride_bs_global, # expert stride of BsGlobal (``global_scale_stride``) stride_c_m, stride_c_n, stride_bias_e, stride_bias_n, - stride_bs_global, # expert stride of BsGlobal (``global_scale_stride``) BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, @@ -849,6 +954,7 @@ def mx_weight_only_matmul_2d_kernel( SWIGLU_LIMIT: tl.constexpr = None, SIMULATE_UNFUSED: tl.constexpr = False, INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, + SWAP_AB: tl.constexpr = False, ): """weight-only dense matmul: ``C = A @ B.T`` with a RAW bf16 activation and an MXFP4/MXFP8 weight upcast to bf16 in-loop (unpack E2M1->E4M3, per-group group-scale multiply), then a plain bf16 @@ -875,9 +981,9 @@ def mx_weight_only_matmul_2d_kernel( ) bs_mask = None a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) - b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE) + b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE, SWAP_AB) - accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False) + accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB) for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC): # raw bf16 [BM, BK]: pointer tile or the host-TMA box (no scale — the second slot is dead; # a distinct name, not `_`, which is already bound pre-loop by mx_2d_scale_ptrs) @@ -888,18 +994,251 @@ def mx_weight_only_matmul_2d_kernel( b_ptrs, BDescriptor, bs_ptrs, bs_mask, None, Bs, 0, pid_n * n_width, k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE), 0, 0, pid_n, k, N, K, 0, stride_bs_n, stride_bs_k, - GATE, False, False, B_MEMORY_MODE, False, False, + GATE, False, False, B_MEMORY_MODE, SWAP_AB, False, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, WEIGHT_VALUES_PER_BYTE, ) accumulator = mx_weight_only_compute( - accumulator, a, b, b_s, COMPUTE_MODE, BLOCK_SIZE_K, n_width, SCALE_GROUP_K + accumulator, a, b, b_s, COMPUTE_MODE, BLOCK_SIZE_K, n_width, SCALE_GROUP_K, SWAP_AB ) # descriptor arms read by absolute box offset (m/k above); the pointer arms advance here a_ptrs += BLOCK_SIZE_K * stride_a_k b_ptrs += (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) * stride_b_k bs_ptrs += SCALE_COLS * stride_bs_k - if GATE: + if SWAP_AB: + accumulator = tl.trans(accumulator) + out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + gemm_epilogue( + C, None, accumulator, out_row, pid_n, pid_m, out_row < M, + stride_c_m, stride_c_n, 0, 0, + BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K, + ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, + COMPUTE_MODE="dot", N_COLS=N, + GlobalScale=BsGlobal, stride_global_e=stride_bs_global, + GlobalScaleA=AsGlobal, stride_global_a=stride_as_global, + Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n, + ) + + +@triton.jit +def split_k_reduce(acc, Partials, TileCounters, tile, num_tiles, BM: tl.constexpr, BN: tl.constexpr, SPLIT_K: tl.constexpr): + """In-kernel split-K reduction for one output tile: store this split's fp32 partial, count the arrival + (``acq_rel`` at GPU scope, after a CTA barrier so the whole partial is written first) on the launch's own zeroed + counter, and in the last program to arrive sum every split's partial in split order (fixed order: the result does + not depend on arrival order). Returns ``(acc, is_last)``.""" + offs = tl.arange(0, BM)[:, None] * BN + tl.arange(0, BN)[None, :] + split = tl.program_id(2) + tl.store(Partials + (split * num_tiles + tile) * (BM * BN) + offs, acc) + tl.debug_barrier() + arrived = tl.atomic_add(TileCounters + tile, 1, sem="acq_rel", scope="gpu") + last = arrived == SPLIT_K - 1 + if last: + total = tl.zeros((BM, BN), tl.float32) + for s in tl.static_range(SPLIT_K): + total += tl.load(Partials + (s * num_tiles + tile) * (BM * BN) + offs, cache_modifier=".cg") + acc = total + return acc, last + + +_SM10X = tl.constexpr(is_sm10x()) + +# split-K over the decode band (M <= 16): a decode tile walks its whole K serially, latency-bound per step + + +@bayesian_autotune( + # the MX weight-only kernel's tuning space minus dot_scaled (the scaled MMA takes float formats only): tile + WS + + # host-TMA memory modes on both operands, scales affine; plus the decode-sized scalar reduce's narrow tiles + get_accelerator_autotuning_configs( + tune_block_nk=True, + tune_block_m=True, + warp_spec=True, + a_memory_modes=("descriptor", "pointer"), + b_memory_modes=("descriptor", "pointer"), + pre_hook=int_weight_only_pre_hook, + ) + # same branch-axis values as the union grid's BM=16 tiles, so coordinate descent steps BN 32 -> 16 within one basin + + narrow_dot_autotuning_configs(pre_hook=int_weight_only_pre_hook, WARP_SPEC=False) + + narrow_dot_autotuning_configs( + block_ms=(1, 2, 4, 8), block_ns=(32, 64), pre_hook=int_weight_only_pre_hook, WARP_SPEC=False, SWAP_AB=True, + ) + + split_k_autotuning_configs(SPLIT_K_VALUES, pre_hook=int_weight_only_pre_hook, WARP_SPEC=False) + + narrow_scalar_autotuning_configs(block_ms=(1, 2, 4)), + # the width and the scale group decide which tiles fit and compile: part of the key + ["N", "K", "m_bit_length", "GATE", "INT_BITS", "SCALE_GROUP_K"], + n_trials=100, + benchmark_input=weight_input, + path_anchor_axes=PATH_ANCHOR_AXES, + finite_check_args=("C", "Cs"), + prune_configs_by={ + "early_config_prune": compose_pruners( + block_within_dim_pruner("K"), + int_scale_group_tile_pruner(), + # the CUDA-core reduce (int_words_dot_bf16) beats the BM=16 mma.sync tile up to M=4 (INT4, B200, 4096x4096: + # M=2 7.7 vs 11.2us, M=4 9.9 vs 11.2us) + scalar_max_m_pruner("M", max_m=4), + scalar_block_m_pruner("M"), + int_split_k_pruner("M", SPLIT_K_MAX_M), + block_m_cap_pruner("M"), + warp_spec_compile_guard_pruner(), + int_2d_warp_spec_pruner(), + descriptor_box_pruner("BLOCK_SIZE_K"), + smem_pruner("BLOCK_SIZE_K"), + ) + }, +) +@triton.jit +def int_weight_only_matmul_2d_kernel( + A, # (M, K) prepared bf16/fp16, FP8 (W4A8), or INT8 (SmoothQuant) activations + ADescriptor, # host TMA descriptor over A (M, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer" + AsGlobal, # fp32 calibrated input scale; None for weight-only + B, # (N, K // VPW) interleaved uint32 words of VPW = 32 // INT_BITS values; under GATE the (2N, ...) gate|up stack + BDescriptor, # host TMA descriptor over B, box (BN, BK // VPW); read iff B_MEMORY_MODE != "pointer" + Bs, # (N, K // SCALE_GROUP_K) float group scales + Bz, # (N, K // SCALE_GROUP_K) integer zero points; None for a symmetric weight + BsGlobal, # fp32 W4A8 second-level weight scale; None otherwise + C, # (M, N) output (GLU intermediate under GATE) + Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None + Partials, # (SPLIT_K_MAX, tiles, BM, n_width) fp32 split-K partial sums; read iff SPLIT_K > 1 + TileCounters, # (tiles,) int32 split-K arrival counters, zeroed per launch; read iff SPLIT_K > 1 + M, + N, + K, + m_bit_length: tl.constexpr, # ceil(log2 M), also proves when a single row tile covers M + stride_a_m, + stride_a_k, + stride_b_n, + stride_b_k, + stride_bs_n, + stride_bs_k, + stride_bz_n, + stride_bz_k, + stride_c_m, + stride_c_n, + stride_bias_e, + stride_bias_n, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SCALE_GROUP_K: tl.constexpr, + A_MEMORY_MODE: tl.constexpr = "pointer", + B_MEMORY_MODE: tl.constexpr = "pointer", + COMPUTE_MODE: tl.constexpr = "dot", + SWAP_AB: tl.constexpr = False, + WARP_SPEC: tl.constexpr = False, + GATE: tl.constexpr = False, + ACT_FN: tl.constexpr = "silu", + SWIGLU_ALPHA: tl.constexpr = None, + SWIGLU_LIMIT: tl.constexpr = None, + SIMULATE_UNFUSED: tl.constexpr = False, + INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, + SPLIT_K: tl.constexpr = 1, + INT_BITS: tl.constexpr = 4, +): + """Integer dense matmul: bf16/fp16 weight-only, calibrated INT4/FP8, or INT8 SmoothQuant. + ``int_weight_only_compute`` runs the MMA or small-token scalar reduce. ``GATE`` loads the stacked (2N, K) + gate|up weight as one ``2*BN``-wide tile + applies the GLU. ``SPLIT_K`` (decode) splits K across that many + programs per output tile (``program_id(2)``); each stores its fp32 partial, and the last to arrive sums all of + them in split order (deterministic) and runs the epilogue, in the same launch.""" + n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N + n_rows = 2 * N if GATE else N + # Wider groups keep the pointer base; int_group_ptrs selects the tile's group. + SCALE_STEP: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K + VPW: tl.constexpr = 32 // INT_BITS # values per uint32 word + offs_kb = tl.arange(0, BLOCK_SIZE_K // VPW) + pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets( + M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, 8, GATE=GATE, + SINGLE_M_TILE=BLOCK_SIZE_M >= (1 << m_bit_length), + ) + # N-clamped, unmasked scale and zero-point reads (WS-lowerable); the epilogue masks the tail + offs_bn_lin = tl.minimum(pid_n * n_width + tl.arange(0, n_width), n_rows - 1) + offs_sk = tl.arange(0, max(1, SCALE_STEP)) + bs_ptrs = Bs + offs_bn_lin[:, None] * stride_bs_n + offs_sk[None, :] * stride_bs_k + bz_ptrs = None + if Bz is not None: + bz_ptrs = Bz + offs_bn_lin[:, None] * stride_bz_n + offs_sk[None, :] * stride_bz_k + a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) + # output-rows-major for the scalar reduce and the mma.sync tiles (BM < 64), whose rhs stays in registers; K-major + # for the tcgen05 tiles, whose rhs is staged through shared memory (int_dequant_tile) + ROWS_MAJOR: tl.constexpr = COMPUTE_MODE == "scalar" or BLOCK_SIZE_M < 64 + b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE, ROWS_MAJOR) + + accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False, tl.int32 if A.dtype.element_ty == tl.int8 else tl.float32) + if SWAP_AB and COMPUTE_MODE == "dot": + accumulator = tl.trans(accumulator) + # this program's share of the K tiles (int_split_k_pruner: the split divides them) + k_per_split = tl.cdiv(K, BLOCK_SIZE_K) + k_first = 0 + if SPLIT_K > 1: + k_per_split = k_per_split // SPLIT_K + k_first = tl.program_id(2) * k_per_split + a_ptrs += k_first * BLOCK_SIZE_K * stride_a_k + b_ptrs += k_first * (BLOCK_SIZE_K // VPW) * stride_b_k + bs_ptrs += k_first * SCALE_STEP * stride_bs_k + if Bz is not None: + bz_ptrs += k_first * SCALE_STEP * stride_bz_k + if Bz is None: + BYTE_ZERO: tl.constexpr = False + else: + BYTE_ZERO: tl.constexpr = Bz.dtype.element_ty == tl.uint8 or Bz.dtype.element_ty == tl.int8 + WHOLE_WORD: tl.constexpr = _SM10X and SWAP_AB and INT_BITS == 4 and ROWS_MAJOR and BYTE_ZERO + # zero points run one K step ahead: their narrow per-row read is not one the pipeliner stages, so each step issues + # the next step's and consumes the previous one's instead of stalling on its own + k_last = k_first + k_per_split - 1 + b_z_next = None + if Bz is not None and not WHOLE_WORD: + b_z_next = load_weight_int_group( + int_group_ptrs(bz_ptrs, k_first, stride_bz_k, BLOCK_SIZE_K, SCALE_GROUP_K), Bz, 0, pid_n, k_first, N, K, + 0, stride_bz_n, stride_bz_k, GATE, False, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, + ) + for k in tl.range(k_first, k_first + k_per_split, warp_specialize=WARP_SPEC): + if WHOLE_WORD: + # Expose the byte load inside this iteration to preserve whole-word MMA ownership; issue before TMA waits. + b_z = load_weight_int_group( + int_group_ptrs(bz_ptrs, k, stride_bz_k, BLOCK_SIZE_K, SCALE_GROUP_K), Bz, 0, pid_n, k, N, K, + 0, stride_bz_n, stride_bz_k, GATE, False, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, + ) + else: + b_z = b_z_next + a, a_scale_dead = load_act_plain( + a_ptrs, ADescriptor, pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K, None, 0, A_MEMORY_MODE, False + ) + b, b_s = load_weight_int( + b_ptrs, BDescriptor, int_group_ptrs(bs_ptrs, k, stride_bs_k, BLOCK_SIZE_K, SCALE_GROUP_K), Bs, 0, + pid_n * n_width, k * (BLOCK_SIZE_K // VPW), 0, pid_n, k, N, K, 0, stride_bs_n, stride_bs_k, GATE, False, + B_MEMORY_MODE, ROWS_MAJOR, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, BITS=INT_BITS, + ) + if Bz is not None and not WHOLE_WORD: + k_next = tl.minimum(k + 1, k_last) + b_z_next = load_weight_int_group( + int_group_ptrs( + bz_ptrs + (k_next - k) * SCALE_STEP * stride_bz_k, k_next, stride_bz_k, BLOCK_SIZE_K, SCALE_GROUP_K + ), + Bz, 0, pid_n, k_next, N, K, 0, stride_bz_n, stride_bz_k, GATE, False, BLOCK_SIZE_N, BLOCK_SIZE_K, + SCALE_GROUP_K, + ) + accumulator = int_weight_only_compute( + accumulator, a, b, b_s, b_z, BLOCK_SIZE_K, n_width, SCALE_GROUP_K, COMPUTE_MODE, ROWS_MAJOR, + INT_BITS, SWAP_AB, weight_global=None if BsGlobal is None else tl.load(BsGlobal), + ) + a_ptrs += BLOCK_SIZE_K * stride_a_k + b_ptrs += (BLOCK_SIZE_K // VPW) * stride_b_k + bs_ptrs += SCALE_STEP * stride_bs_k + if Bz is not None: + bz_ptrs += SCALE_STEP * stride_bz_k + + if SWAP_AB and COMPUTE_MODE == "dot": + accumulator = tl.trans(accumulator) + if SPLIT_K > 1: + accumulator, last = split_k_reduce( + accumulator, Partials, TileCounters, tl.program_id(0) * tl.num_programs(1) + tl.program_id(1), + tl.num_programs(0) * tl.num_programs(1), BLOCK_SIZE_M, n_width, SPLIT_K, + ) + else: + last = True + if last: + output_scale = BsGlobal + if A.dtype.element_ty == tl.int8: + output_scale = Bs + offs_bn_lin * stride_bs_n out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) gemm_epilogue( C, None, accumulator, out_row, pid_n, pid_m, out_row < M, @@ -907,29 +1246,102 @@ def mx_weight_only_matmul_2d_kernel( BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K, ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, COMPUTE_MODE="dot", N_COLS=N, - GlobalScale=BsGlobal, stride_global_e=stride_bs_global, + GlobalScale=output_scale, GlobalScaleA=AsGlobal, Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n, ) - else: - # explicit (not the epilogue fold): this arm exits through store_masked, which has no - # GlobalScale slot - accumulator = apply_global_scale( - accumulator, BsGlobal, stride_bs_global, None, 0, 0 + + +@bayesian_autotune( + get_accelerator_autotuning_configs( + tune_block_nk=True, tune_block_m=True, warp_spec=True, + swap_ab=True, swap_block_sizes_m=(1, 2, 4, 8), + a_memory_modes=("descriptor", "pointer"), + b_memory_modes=("descriptor", "pointer"), + pre_hook=rebind_weight_only_descriptors, + ), + ["N", "K", "m_bit_length", "GATE", "BLOCK_N", "BLOCK_K", "GATE_SCALES"], + n_trials=100, + benchmark_input=weight_input, + path_anchor_axes=PATH_ANCHOR_AXES, + coupled_axes=(("BLOCK_SIZE_K", "num_stages"),), + finite_check_args=("C",), + prune_configs_by={ + "early_config_prune": compose_pruners( + block_within_dim_pruner("K"), + swap_max_m_pruner("M", max_m=8), + gate_stacked_tmem_trap_pruner(), + warp_spec_compile_guard_pruner(), + descriptor_box_pruner(), + smem_pruner(), ) - accumulator = add_bias( - accumulator, Bias, stride_bias_e, stride_bias_n, 0, pid_n, - BLOCK_SIZE_N, N, False, + }, +) +@triton.jit +def block_fp8_weight_only_matmul_2d_kernel( + A, ADescriptor, B, BDescriptor, Bs, C, Bias, + M, N, K, m_bit_length, + stride_a_m, stride_a_k, stride_b_n, stride_b_k, + stride_bs_n, stride_bs_k, stride_c_m, stride_c_n, + stride_bias_e, stride_bias_n, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + A_MEMORY_MODE: tl.constexpr = "pointer", + B_MEMORY_MODE: tl.constexpr = "pointer", + WARP_SPEC: tl.constexpr = False, + SWAP_AB: tl.constexpr = False, + GATE: tl.constexpr = False, + GATE_SCALES: tl.constexpr = False, + ACT_FN: tl.constexpr = "silu", + SWIGLU_ALPHA: tl.constexpr = None, + SWIGLU_LIMIT: tl.constexpr = None, + SIMULATE_UNFUSED: tl.constexpr = False, + INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16, +): + """Dense FP8 W8A16 GEMM with in-tile weight rescaling.""" + n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N + pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets( + M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GATE=GATE + ) + a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) + b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, B_MEMORY_MODE, SWAP_AB) + accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False) + for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC): + a, _ = load_act_plain( + a_ptrs, ADescriptor, pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K, None, 0, + A_MEMORY_MODE, False, ) - store_masked( - C, accumulator, pid_m, pid_n, M, N, stride_c_m, stride_c_n, - BLOCK_SIZE_M, BLOCK_SIZE_N, + w = load_weight_block_fp8( + b_ptrs, BDescriptor, Bs, 0, pid_n * n_width, k * BLOCK_SIZE_K, + 0, offs_bn, k * BLOCK_SIZE_K + offs_k, + 0, stride_bs_n, stride_bs_k, + GATE, False, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K, + BLOCK_N, BLOCK_K, GATE_SCALES, a.dtype, ) + accumulator = plain_dot(accumulator, a, w, SWAP_AB) + a_ptrs += BLOCK_SIZE_K * stride_a_k + b_ptrs += BLOCK_SIZE_K * stride_b_k + out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + gemm_epilogue( + C, None, accumulator, out_row, pid_n, pid_m, out_row < M, + stride_c_m, stride_c_n, 0, 0, + BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K, + ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE, + COMPUTE_MODE="dot", N_COLS=N, + Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n, + ) @bayesian_autotune( # Same 2D-swizzle loop physics as the tensor kernel — tune the tile, WS a pure perf axis - # (compile-guarded), no memory-mode/scale axes (unquantized, no scales anywhere). - get_accelerator_autotuning_configs(tune_block_nk=True, warp_spec=True, tune_block_m=True), + # The same pointer/descriptor loaders as the quantized GEMMs, without scale loads. + get_accelerator_autotuning_configs( + tune_block_nk=True, warp_spec=True, tune_block_m=True, + a_memory_modes=("descriptor", "pointer"), b_memory_modes=("descriptor", "pointer"), + pre_hook=rebind_weight_only_descriptors, + ), # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide). ["N", "K", "m_bit_length", "GATE"], n_trials=100, @@ -940,6 +1352,10 @@ def mx_weight_only_matmul_2d_kernel( prune_configs_by={ "early_config_prune": compose_pruners( block_within_dim_pruner("K"), + matched_memory_modes_pruner(), + descriptor_box_pruner(), + gate_stacked_tmem_trap_pruner(), + smem_pruner(), warp_spec_compile_guard_pruner(), ) }, @@ -947,7 +1363,9 @@ def mx_weight_only_matmul_2d_kernel( @triton.jit def full_precision_matmul_2d_kernel( A, # (M, K) BF16/FP16 activations + ADescriptor, B, # (N, K) weights in A's dtype; under GATE the (2N, K) gate|up stack + BDescriptor, C, # (M, N) output; under GATE the GLU intermediate Bias, # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None # Shape @@ -969,6 +1387,8 @@ def full_precision_matmul_2d_kernel( BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, WARP_SPEC: tl.constexpr = False, + A_MEMORY_MODE: tl.constexpr = "pointer", + B_MEMORY_MODE: tl.constexpr = "pointer", # Dense gate|up fusion: B is the (2N, K) gate|up stack, C the [M, N] GLU output. Folds out at # compile time when GATE=False (the plain dense GEMM). GATE: tl.constexpr = False, @@ -986,15 +1406,17 @@ def full_precision_matmul_2d_kernel( pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets( M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GATE=GATE ) - a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, "pointer", True) - b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, "pointer") + a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True) + b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, B_MEMORY_MODE) + a_desc = operand_tile_descriptor(ADescriptor, A, M, K, stride_a_m, stride_a_k, BLOCK_SIZE_M, BLOCK_SIZE_K, A_MEMORY_MODE) + b_desc = operand_tile_descriptor(BDescriptor, B, N, K, stride_b_n, stride_b_k, BLOCK_SIZE_N, BLOCK_SIZE_K, B_MEMORY_MODE, GATE) accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False) for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC): - a = tl.load(a_ptrs) + a, _ = load_act_plain(a_ptrs, a_desc, pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K, None, offs_am, A_MEMORY_MODE, False) w, _ = load_weight_plain( - b_ptrs, b_ptrs, 0, pid_n * n_width, k * BLOCK_SIZE_K, - GATE, False, "pointer", False, BLOCK_SIZE_N, BLOCK_SIZE_K, + b_ptrs, b_desc, 0, pid_n * n_width, k * BLOCK_SIZE_K, + GATE, False, B_MEMORY_MODE, False, BLOCK_SIZE_N, BLOCK_SIZE_K, ) accumulator = accumulator + fp8_dot(a, w, False, BLOCK_SIZE_K) a_ptrs += BLOCK_SIZE_K * stride_a_k @@ -1267,27 +1689,30 @@ def w8a8_tensor_dynamic_fp8_matmul( swiglu_limit: float | None = None, simulate_unfused: bool = False, bias: torch.Tensor | None = None, + activation_scale_ub: float | None = None, ) -> list[torch.Tensor]: - """Tensor-scale FP8 matmul: ``C = A @ B.T``; activations quantized offline per row. + """Tensor-scale / row-wise FP8 matmul: ``C = A @ B.T``; activations quantized offline per row. A: (..., K) raw activations, bf16/fp16/fp32 (flattened to (M, K) internally) — per-row scales computed via ``fp8_act_quant_tensor_wide(A, K)``, or quantized against ``As`` when a calibrated (static) one is given. - B: (N, K) FP8 weights — under ``gate`` the ``(2N, K)`` gate|up stack (one per-tensor scale). - Bs: scalar, (1,), or (1, 1) — single tensor-scale weight scale. - As: the calibrated (static) activation scale, one value for the whole matmul; ``None`` - derives one per row from the data. + B: (N, K) FP8 weights — under ``gate`` the ``(2N, K)`` gate|up stack. + Bs: scalar, (1,), or (1, 1) — single tensor-scale weight scale; or (N, 1) row-wise, one + per (stacked) weight row. + As: per-token scales for pre-quantized A, or a scalar calibrated scale for raw A; + ``None`` derives one per row from the data. ``gate`` fuses the gate|up projection into one stacked GEMM + SwiGLU, returning the ``[..., N]`` GLU intermediate. Returns a one-element list (mirrors the MX/grouped op). """ validate_dense_2d_operands(A, B) + output_dtype = resolve_output_dtype(output_dtype, A, As) # cuBLAS per-tensor fast path, INSIDE the op for the same reasons as the MX one below: # every caller gets it, and a None falls through to the Triton launch unchanged. if not gate and not simulate_unfused and bias is None: _smm = _torch_scaled_mm_2d_tensor_wise( - A, B, As, Bs, None, resolve_output_dtype(output_dtype, A, None) + A, B, As, Bs, None, output_dtype, activation_scale_ub ) if _smm is not None: return [_smm] @@ -1297,13 +1722,33 @@ def w8a8_tensor_dynamic_fp8_matmul( N = rows // 2 if gate else rows M = A.numel() // A.shape[-1] - assert Bs.numel() == 1, f"Bs must be scalar or (1,), got {tuple(Bs.shape)}" + assert Bs.numel() == 1 or is_rowwise(B, Bs), ( + f"Bs must be scalar, (1,) or row-wise ({rows}, 1), got {tuple(Bs.shape)}" + ) - # one scale per row (derived here), or the calibrated one every row reads (stride 0) - qA, As, as_stride_m, _ = tensor_wide_act_operands(A, As) - Bs = Bs.reshape(1) + # One scale per row, derived here into rows padded to the largest M tile, or the calibrated + # one every row reads (stride 0). The A box then reads the padded extent in bounds: a box past + # the tensor's rows takes the TMA's zero-fill path, which costs the decode GEMV ~4us with cold + # weights (B200, N=16384 K=5120, M=1). A pre-quantized A whose storage lacks the padded rows + # (``fp8_act_quant_tensor_wide(row_multiple=...)`` provides them) is copied into a padded + # buffer, cheaper than that path. An unaligned view also needs a copy for TMA's 16-byte base + # alignment. Rows past M are never stored. + qA, As, as_stride_m, _ = tensor_wide_act_operands( + A.reshape(M, K), None if As is None else As.reshape(-1), + row_multiple=TENSOR_MATMUL_A_ROW_MULTIPLE, + activation_scale_ub=activation_scale_ub, + ) + a_rows = triton.cdiv(M, TENSOR_MATMUL_A_ROW_MULTIPLE) * TENSOR_MATMUL_A_ROW_MULTIPLE + if (qA.storage_offset() * qA.element_size()) % 16 or ( + qA.untyped_storage().nbytes() < (qA.storage_offset() + a_rows * K) * qA.element_size() + ): + qA = torch.empty(a_rows, K, device=qA.device, dtype=qA.dtype)[:M].copy_(qA) + Bs = Bs.reshape(-1, 1) C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype) + # placeholder boxes, re-bound per tuned config by rebind_tiled_descriptors + a_descriptor = TensorDescriptor(qA, [a_rows, K], [K, 1], [1, 16]) + b_descriptor = maybe_descriptor(B, [1, 16]) def grid(META): return ( @@ -1315,8 +1760,10 @@ def grid(META): bias_stride_e, bias_stride_n = bias_strides(bias) compile_time_only_triton_wrap(w8a8_tensor_dynamic_fp8_matmul_kernel)[grid]( qA, + a_descriptor, As, B, + b_descriptor, Bs, C, bias, @@ -1329,6 +1776,7 @@ def grid(META): as_stride_m, B.stride(1), B.stride(0), + row_scale_stride(Bs), C.stride(-2), C.stride(-1), bias_stride_e, @@ -1338,7 +1786,8 @@ def grid(META): SWIGLU_ALPHA=swiglu_alpha, SWIGLU_LIMIT=swiglu_limit, SIMULATE_UNFUSED=simulate_unfused, - INTERMEDIATE_DTYPE=tl_dtype(resolve_output_dtype(output_dtype, A, None)), + INTERMEDIATE_DTYPE=tl_dtype(output_dtype), + **pdl_launch_kwargs(), ) return [C] @@ -1364,6 +1813,9 @@ def mx_dynamic_matmul( b_global_scale: torch.Tensor | None = None, output_global_scale: torch.Tensor | None = None, bias: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, + svdquant_lora_a: torch.Tensor | None = None, + svdquant_lora_b: torch.Tensor | None = None, ) -> list[torch.Tensor]: """MX/NVFP4 matmul ``C = A @ B.T``; activations quantized offline above the ``maybe_act_quant`` M threshold, inline below it. ``activation_format`` sets the @@ -1377,6 +1829,23 @@ def mx_dynamic_matmul( leading dims are flattened to (M, K) and restored on the output """ assert B.ndim == 2 and Bs.ndim in (2, 5, 6) # 5D = pre-swizzled; 6D = the gate|up artifact + if a_pre_quant_scale is not None and activation_format != "fp8" and svdquant_lora_a is None: + assert As is None, "pre_quant_scale applies before activation quantization; pass raw inputs" + A = permute_act_cols(A, None, a_pre_quant_scale) + lora_input = None + if activation_format == "fp8": + assert not quantize_output and a_global_scale is None, "ordinary FP8 inputs use As=input_scale and no fused requant" + output_dtype = resolve_output_dtype(output_dtype, A, None) + group = mx_scale_family(Bs, A.shape[-1]) + weight_upcast = Bs.ndim == 2 and Bs.dtype == torch.float8_e4m3fn and group == MX_SCALE_GROUP_K and A.numel() // A.shape[-1] > 16 + A, As, a_global_scale = prepare_fp8_mx_inputs(A, As, ue8m0_as_uint8(Bs).dtype, group, pre_quant_scale=a_pre_quant_scale, storage_dtype=torch.bfloat16 if weight_upcast else FP8_DTYPE) + if weight_upcast: + return mx_weight_only_matmul_2d( + A, B, Bs, output_dtype=output_dtype, gate=gate, act_fn=act_fn, + swiglu_alpha=swiglu_alpha, swiglu_limit=swiglu_limit, simulate_unfused=simulate_unfused, + b_global_scale=b_global_scale, a_global_scale=a_global_scale, + bias=bias, + ) assert B.dtype in (torch.int8, torch.float8_e4m3fn), ( f"B must be int8 (packed E2M1) or float8_e4m3fn (E4M3), got {B.dtype}" ) @@ -1394,7 +1863,7 @@ def mx_dynamic_matmul( # hangs on this op and saves INPUTS, so the registered dgrad serves a scaled_mm forward # unchanged. Measured 1.4-4.8x over the Triton kernel at every M (B200, 2026-08-27); a None # falls through to the Triton launch exactly as before. - if not gate and not quantize_output and bias is None and not simulate_unfused: + if not gate and not quantize_output and bias is None and not simulate_unfused and svdquant_lora_a is None: _smm = _torch_scaled_mm_2d( A, B, As, Bs, activation_format, a_global_scale, b_global_scale, resolve_output_dtype(output_dtype, A, None), @@ -1443,8 +1912,10 @@ def mx_dynamic_matmul( # NVFP4 two-level: a_global_scale (the calibrated activation global) normalizes the quant on both # arms — the offline kernel divides before the block quant, the inline arm via AsGlobal. act_global_scale = normalize_global_scale(a_global_scale, 1) # g_a; folds in-kernel - if act_global_scale is not None: - assert activation_format == "nvfp4", "an activation global is NVFP4-only" + if activation_format == "fp8": + act_quant = None # As carries unit group scales; A is already quantized. + elif act_global_scale is not None: + assert activation_format == "nvfp4", "an activation global requires NVFP4" act_quant = lambda a: MX_ACT_QUANT[activation_format]( # noqa: E731 a, swizzled=swizzled_scales, global_scale=act_global_scale ) @@ -1457,7 +1928,10 @@ def mx_dynamic_matmul( out_global_scale = normalize_global_scale(output_global_scale, 1) # As given ⇒ A is already quantized (the routed-op parity: a pre-quantized activation + its # scales); else quantize raw A (offline above the M threshold, inline in the kernel below it). - if As is not None: + if svdquant_lora_a is not None: + a_vals, as_scales, lora_input = svdquant_act_quant( + A, svdquant_lora_a, a_pre_quant_scale, act_global_scale, swizzled_scales) + elif As is not None: assert (As.dtype == torch.float8_e4m3fn) == (Bs.dtype == torch.float8_e4m3fn), ( f"activation scales ({As.dtype}) must match the weight scale family ({Bs.dtype})" ) @@ -1517,11 +1991,21 @@ def mx_dynamic_matmul( bs_descriptor = swizzled_scale_descriptor(bs_u8) if swizzled_scales else b_descriptor def grid(META): + if META["PERSISTENT"]: + tiles = triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]) + return (min(tiles, persistent_program_count(A.device.index or 0)),) return ( triton.cdiv(M, META["BLOCK_SIZE_M"]), triton.cdiv(N, META["BLOCK_SIZE_N"]), + META["SPLIT_K"], ) + # Only the one real decode row is stored; native split tiles span 128 or 256 columns. + partials = counters = None + if lora_input is not None and M == 1 and not gate: + partials = torch.empty(SPLIT_K_VALUES[-1] * triton.cdiv(N, 256) * 256, device=A.device, dtype=torch.float32) + counters = torch.empty(triton.cdiv(N, 128), device=A.device, dtype=torch.int32) + with device_context(A.device): bias_stride_e, bias_stride_n = bias_strides(bias) compile_time_only_triton_wrap(mx_dynamic_matmul_kernel)[grid]( @@ -1529,10 +2013,14 @@ def grid(META): a_descriptor, as_u8, as_descriptor, + lora_input, b_u8, b_descriptor, bs_u8, bs_descriptor, + svdquant_lora_b, + partials, + counters, C, Cs, bias, @@ -1550,10 +2038,14 @@ def grid(META): A_q.stride(1), as_u8.stride(0), # As row stride (affine act-scale read; dead on the swizzled arm) as_u8.stride(-1), # As group-axis stride (1 when row-major; see stride_as_k) + lora_input.stride(-2) if lora_input is not None else 0, + lora_input.stride(-1) if lora_input is not None else 0, b_u8.stride(1), b_u8.stride(0), bs_u8.stride(1), bs_u8.stride(0), + svdquant_lora_b.stride(0) if svdquant_lora_b is not None else 0, + svdquant_lora_b.stride(1) if svdquant_lora_b is not None else 0, C.stride(-2), C.stride(-1), stride_cs_m, @@ -1564,6 +2056,8 @@ def grid(META): global_scale_stride(input_global_scale), global_scale_stride(out_global_scale), SCALE_GROUP_K=scale_group, + LORA_RANK=svdquant_lora_a.shape[0] if svdquant_lora_a is not None else 0, + LORA_SPLIT_K=lora_input.shape[0] if lora_input is not None and lora_input.ndim == 3 else 1, ACTIVATION_FORMAT=activation_format, SWIZZLED_SCALES=swizzled_scales, SWIZZLED_OUT=swizzled_out, @@ -1607,6 +2101,9 @@ def full_precision_matmul_2d( M = A.numel() // A.shape[-1] C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype) + a_descriptor = maybe_descriptor(A.view(M, K), [16, 32]) + b_descriptor = maybe_descriptor(B, [16, 32]) + def grid(META): return ( triton.cdiv(M, META["BLOCK_SIZE_M"]), @@ -1617,7 +2114,9 @@ def grid(META): bias_stride_e, bias_stride_n = bias_strides(bias) compile_time_only_triton_wrap(full_precision_matmul_2d_kernel)[grid]( A, + a_descriptor, B, + b_descriptor, C, bias, M, @@ -1660,6 +2159,7 @@ def mx_weight_only_matmul_2d( simulate_unfused: bool = False, b_global_scale: torch.Tensor | None = None, bias: torch.Tensor | None = None, + a_global_scale: torch.Tensor | None = None, ) -> list[torch.Tensor]: """weight-only dense matmul ``C = A @ B.T``: a RAW bf16/fp16 activation (no quant) against an MXFP4/NVFP4/MXFP8 weight upcast to bf16 in-loop (NVFP4's per-tensor fp32 ``b_global_scale`` @@ -1681,6 +2181,9 @@ def mx_weight_only_matmul_2d( f"K (={K}) must equal {WEIGHT_VALUES_PER_BYTE} * B.shape[1] (={K_b})" ) scale_group = mx_scale_family(Bs, K) + assert Bs.dtype != torch.float8_e4m3fn or scale_group == NVFP4_SCALE_GROUP_K or a_global_scale is not None, ( + "group-32 E4M3 FP4 weights require calibrated FP8 activations" + ) assert Bs.shape == (rows, K // scale_group), ( f"Bs shape {tuple(Bs.shape)} != ({rows}, {K // scale_group})" ) @@ -1688,6 +2191,7 @@ def mx_weight_only_matmul_2d( b_u8 = e2m1_as_uint8(B) bs_u8 = ue8m0_as_uint8(Bs) b_global_scale = normalize_global_scale(b_global_scale, 1) + a_global_scale = normalize_global_scale(a_global_scale, 1) C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype) # Host-TMA descriptors over the (M, K) bf16 A and packed (N, K_bytes) weight — placeholder box # rebound per tuned config by rebind_weight_only_descriptors, read only by the descriptor configs the @@ -1706,6 +2210,7 @@ def grid(META): compile_time_only_triton_wrap(mx_weight_only_matmul_2d_kernel)[grid]( A, a_descriptor, + a_global_scale, b_u8, b_descriptor, bs_u8, @@ -1718,15 +2223,16 @@ def grid(META): int(M).bit_length(), A.stride(-2), A.stride(-1), + global_scale_stride(a_global_scale), b_u8.stride(0), b_u8.stride(1), bs_u8.stride(0), bs_u8.stride(1), + global_scale_stride(b_global_scale), C.stride(-2), C.stride(-1), bias_stride_e, bias_stride_n, - global_scale_stride(b_global_scale), SCALE_GROUP_K=scale_group, WEIGHT_VALUES_PER_BYTE=WEIGHT_VALUES_PER_BYTE, GATE=gate, @@ -1740,6 +2246,177 @@ def grid(META): return [C] +@compile_time_only_triton_op( + add_op_namespace_prefix("block_fp8_weight_only_matmul_2d"), mutates_args=(), opaque=True, +) +def block_fp8_weight_only_matmul_2d( + A: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + output_dtype: torch.dtype | None = None, + gate: bool = False, + act_fn: str = "silu", + swiglu_alpha: float | None = None, + swiglu_limit: float | None = None, + simulate_unfused: bool = False, + bias: torch.Tensor | None = None, +) -> List[torch.Tensor]: + """Block/per-tensor FP8 weights with unquantized activations (ModelOpt FP8_PB_WO).""" + assert A.dtype in (torch.bfloat16, torch.float16), "weight-only FP8 expects BF16/FP16 activations" + assert A.is_contiguous(), "A must be contiguous" + assert B.ndim == 2 and B.dtype == torch.float8_e4m3fn + assert Bs.ndim == 2 or Bs.numel() == 1, "weight scales must be a 2D block grid or a scalar" + rows, K = B.shape + assert A.shape[-1] == K + M = A.numel() // K + N = rows // 2 if gate else rows + if M >= FP8_WEIGHT_ONLY_DEQUANT_MIN_M: + weights = dequant_block_fp8_weight(B, Bs, A.dtype, gate=gate) + if (not gate and output_dtype in (None, A.dtype) + and (bias is None or bias.dtype == A.dtype)): + native_bias = bias[0] if bias is not None and bias.ndim == 2 else bias + with torch.autocast(A.device.type, enabled=False): + return [torch.nn.functional.linear(A, weights, native_bias)] + return full_precision_matmul_2d( + A, weights, + output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, bias=bias, + ) + Bs, scale_block_n, scale_block_k = block_fp8_weight_scale_grid(B, Bs, gate) + Bs = ue8m0_as_uint8(Bs) + dtype = resolve_output_dtype(output_dtype, A, None) + C = A.new_empty(A.shape[:-1] + (N,), dtype=dtype) + a_descriptor = maybe_descriptor(A.view(M, K), [1, 32]) + b_descriptor = maybe_descriptor(B, [1, 32]) + with device_context(A.device): + bias_stride_e, bias_stride_n = bias_strides(bias) + compile_time_only_triton_wrap(block_fp8_weight_only_matmul_2d_kernel)[ + lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE_M"]), triton.cdiv(N, meta["BLOCK_SIZE_N"])) + ]( + A, a_descriptor, B, b_descriptor, Bs, C, bias, M, N, K, + # M=2 and M=3/4 need different one-tile swaps and cached configs. + max(int(M) - 1, 0).bit_length(), + A.stride(-2), A.stride(-1), B.stride(0), B.stride(1), + Bs.stride(0), Bs.stride(1), C.stride(-2), C.stride(-1), + bias_stride_e, bias_stride_n, + BLOCK_N=scale_block_n, BLOCK_K=scale_block_k, + GATE=gate, ACT_FN=act_fn, SWIGLU_ALPHA=swiglu_alpha, + SWIGLU_LIMIT=swiglu_limit, SIMULATE_UNFUSED=simulate_unfused, + INTERMEDIATE_DTYPE=tl_dtype(dtype), + GATE_SCALES=gate and Bs.shape[0] > 1, + ) + return [C] + + +@compile_time_only_triton_op( + add_op_namespace_prefix("int_weight_only_matmul_2d"), + mutates_args=(), + opaque=True, +) +def int_weight_only_matmul_2d( + A: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + output_dtype: torch.dtype | None = None, + gate: bool = False, + act_fn: str = "silu", + swiglu_alpha: float | None = None, + swiglu_limit: float | None = None, + simulate_unfused: bool = False, + b_zero: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + b_perm: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, + a_scale: torch.Tensor | None = None, + b_scale2: torch.Tensor | None = None, +) -> list[torch.Tensor]: + """Dense integer GEMM with interleaved uint32 weights and optional fused GLU. + + Weight width follows from the packed shape. Dequantization rounds `(q - zero) * scale` + in the scale dtype before casting to the activation dtype. `b_perm` gathers act-order + columns after the optional input smoothing; calibrated recipes quantize this prepared input. + Prefill can dequantize weights once per call; decode keeps conversion inline. + """ + bits = validate_int_operands(A, B, Bs, b_zero, b_perm, a_scale, b_scale2) + smoothquant = a_scale is not None and b_scale2 is None + a_scale = normalize_global_scale(a_scale, 1) + b_scale2 = normalize_global_scale(b_scale2, 1) + output_dtype = resolve_output_dtype(output_dtype, A, None) + rows, K_b = B.shape + K = A.shape[-1] + M = A.numel() // K + N = rows // 2 if gate else rows + if b_perm is not None or a_pre_quant_scale is not None or a_scale is not None: + A = permute_act_cols(A, b_perm, a_pre_quant_scale, a_scale, torch.int8 if smoothquant else FP8_DTYPE) + assert A.is_contiguous(), "A must be contiguous (the host-TMA descriptor arm reads its box)" + # Large prefill amortizes a per-call dequantization and reuses the unquantized Triton GEMM. + # Smaller projections keep conversion inline to avoid the extra dispatch cost in eager mode. + if ( + a_scale is None and A.ndim == 2 and M >= 4096 and min(N, K) >= 2048 + and not gate and bias is None and output_dtype in (None, A.dtype) + ): + return full_precision_matmul_2d(A, dequantize_int_weight(B, Bs, K, b_zero, dtype=A.dtype), output_dtype) + C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype) + # placeholder boxes, rebound per tuned config by rebind_weight_only_descriptors + a_descriptor = TensorDescriptor.from_tensor(A.view(M, K), [1, 32]) + b_descriptor = TensorDescriptor.from_tensor(B, [1, 32]) + + def grid(META): + return (triton.cdiv(M, META["BLOCK_SIZE_M"]), triton.cdiv(N, META["BLOCK_SIZE_N"]), META.get("SPLIT_K", 1)) + + # split-K scratch (decode band only, int_split_k_pruner), private to this launch: fp32 partials sized for the + # widest split and the narrowest tile, and arrival counters the split configs' pre_hook zeroes. Both come from + # the caching allocator on the launch's stream (a concurrent launch on another stream gets its own), and a + # captured graph replays the memset. + partials = counters = None + if M <= SPLIT_K_MAX_M: + tiles = triton.cdiv(M, 16) * triton.cdiv(N, 16) + partials = torch.empty(SPLIT_K_VALUES[-1] * tiles * 16 * 16 * (2 if gate else 1), device=A.device, dtype=torch.float32) + counters = torch.empty(tiles, device=A.device, dtype=torch.int32) # zeroed by int_weight_only_pre_hook + + with device_context(A.device): + bias_stride_e, bias_stride_n = bias_strides(bias) + compile_time_only_triton_wrap(int_weight_only_matmul_2d_kernel)[grid]( + A, + a_descriptor, + a_scale, + B, + b_descriptor, + Bs, + b_zero, + b_scale2, + C, + bias, + partials, + counters, + M, + N, + K, + # M=2 and M=3 need different matched row tiles and cached configs. + max(int(M) - 1, 0).bit_length(), + A.stride(-2), + A.stride(-1), + B.stride(0), + B.stride(1), + Bs.stride(0), + Bs.stride(1), + *(b_zero.stride() if b_zero is not None else (0, 0)), + C.stride(-2), + C.stride(-1), + bias_stride_e, + bias_stride_n, + SCALE_GROUP_K=K // Bs.shape[1], + INT_BITS=bits, + GATE=gate, + ACT_FN=act_fn, + SWIGLU_ALPHA=swiglu_alpha, + SWIGLU_LIMIT=swiglu_limit, + SIMULATE_UNFUSED=simulate_unfused, + INTERMEDIATE_DTYPE=tl_dtype(resolve_output_dtype(output_dtype, A, None)), + ) + + return [C] + + # ── torch scaled_mm fast path (2D dense, Blackwell) ────────────────────────── # # cuBLAS's block-scaled GEMM beats our Triton 2D kernel at EVERY measured M on the @@ -1759,17 +2436,19 @@ def grid(META): # from Triton, not a tuning gap. Relerr 1.2e-5 against the Triton arm. # # Inductor refuses to lower scaled_mm with SWIZZLE_32_4_4 scales ("does not yet support -# non-trivial swizzles" — repros/scaled_mm_swizzle_compile.py), so the call lives behind an -# OPAQUE custom op: dynamo captures the op as a leaf and the real scaled_mm runs at runtime, -# which serves compiled callers too. F.scaled_mm has no autograd formula while our ops carry -# registered dgrads, so differentiating callers fall through to the Triton ops. CUDA-graph +# non-trivial swizzles" — repros/scaled_mm_swizzle_compile.py). It never sees the call: the routes +# run inside the OPAQUE ops, which dynamo captures as leaves, so compiled callers get the cuBLAS +# path at runtime and eager pays no op dispatch. F.scaled_mm has no autograd formula while our ops +# carry registered dgrads, so differentiating callers fall through to the Triton ops. CUDA-graph # capture is fine (measured under do_bench_cudagraph). _SCALED_MM_MIN_M = int(os.environ.get("FINEGRAINED_SCALED_MM_MIN_M", "2")) +# Row-wise FP8 keeps Triton decode; the public path wins through the measured M=8192. +_SCALED_MM_ROWWISE_MIN_M = int(os.environ.get("FINEGRAINED_SCALED_MM_ROWWISE_MIN_M", "65")) +_SCALED_MM_ROWWISE_MAX_M = int(os.environ.get("FINEGRAINED_SCALED_MM_ROWWISE_MAX_M", "8193")) -@torch.library.custom_op(add_op_namespace_prefix("scaled_mm_2d"), mutates_args=()) -def _scaled_mm_2d_op( +def _scaled_mm_2d( Aq: torch.Tensor, As: torch.Tensor, B: torch.Tensor, @@ -1778,17 +2457,18 @@ def _scaled_mm_2d_op( g_b: torch.Tensor | None, fmt: str, ) -> torch.Tensor: - """``F.scaled_mm`` behind an OPAQUE custom op. Inductor cannot lower scaled_mm with - SWIZZLE_32_4_4 scales (repros/scaled_mm_swizzle_compile.py), but it never has to: dynamo - captures this op as a leaf and the real call runs at runtime — so compiled callers get the - cuBLAS fast path too, not just eager ones.""" + """``F.scaled_mm`` in the format's scaling; inside an opaque op's fake impl only the output is + allocated, as the Triton launches are skipped there.""" + if launches_skipped(): + return Aq.new_empty(Aq.shape[0], B.shape[0], dtype=torch.bfloat16) F = torch.nn.functional ST, SW = F.ScalingType, F.SwizzleType - if fmt == "fp8": + if fmt == "fp8": # one fp32 per operand, or one per row of each (per token / per output row) + scaling = ST.TensorWise if Bs.numel() == 1 else ST.RowWise return F.scaled_mm( Aq, B.t(), - As.reshape(()), ST.TensorWise, - Bs.reshape(()), ST.TensorWise, + As.reshape(()) if Bs.numel() == 1 else As.reshape(-1, 1), scaling, + Bs.reshape(()) if Bs.numel() == 1 else Bs.reshape(1, -1), scaling, output_dtype=torch.bfloat16, ) if fmt == "mxfp8": @@ -1813,23 +2493,20 @@ def _scaled_mm_2d_op( ) -@_scaled_mm_2d_op.register_fake -def _(Aq, As, B, Bs, g_a, g_b, fmt): - return Aq.new_empty(Aq.shape[0], B.shape[0], dtype=torch.bfloat16) - - def _smm_reject(reason): if os.environ.get("FINEGRAINED_SCALED_MM_DEBUG"): print(f"[scaled_mm route] rejected: {reason}", flush=True) return None -def _torch_scaled_mm_2d_tensor_wise(A, B, As, Bs, activation_format, out_dtype): - """The scaled_mm route for per-tensor static FP8, or ``None`` to take the Triton ops. +def _torch_scaled_mm_2d_tensor_wise(A, B, As, Bs, activation_format, out_dtype, activation_scale_ub=None): + """The scaled_mm route for the K-constant FP8 scales, or ``None`` to take the Triton ops. - One fp32 value per operand IS scaled_mm's ``TensorWise``, so this arm needs neither swizzled - scales nor a block-aligned K — only a STATIC activation scale, since a dynamic one would cost - the amax pass the Triton kernel fuses into its own loop. + One fp32 value per operand IS scaled_mm's ``TensorWise`` and one per row of each its + ``RowWise``, so this arm needs neither swizzled scales nor a block-aligned K: + - per-tensor weights with a STATIC activation scale (``TensorWise``); + - row-wise weights with per-token activations, derived here by the same quant the Triton + arm runs, or handed in pre-quantized (``RowWise``), in the measured M window. """ if os.environ.get("FINEGRAINED_DISABLE_SCALED_MM"): return _smm_reject("env-disabled") @@ -1841,18 +2518,33 @@ def _torch_scaled_mm_2d_tensor_wise(A, B, As, Bs, activation_format, out_dtype): return _smm_reject("activation_format") if out_dtype is not torch.bfloat16: return _smm_reject("out-dtype") - if As is None or As.numel() != 1: + K = A.shape[-1] + M = A.numel() // K + rowwise = is_rowwise(B, Bs) + if rowwise: + if As is not None and As.numel() != M: + return _smm_reject("static activation scale on row-wise weights") + if not _SCALED_MM_ROWWISE_MIN_M <= M < _SCALED_MM_ROWWISE_MAX_M: + return _smm_reject("M") + elif As is None or As.numel() != 1: return _smm_reject("not-static-per-tensor") - if As.dtype is not torch.float32 or Bs.dtype is not torch.float32: + if (As is not None and As.dtype is not torch.float32) or Bs.dtype is not torch.float32: # a UE8M0 container's scalar is an EXPONENT, which TensorWise would read as a multiplier return _smm_reject("non-fp32 scale") - K = A.shape[-1] if K % 16 or B.shape[0] % 16: return _smm_reject("shape") - M = A.numel() // K - As = As.reshape(-1).float() - Aq = A.reshape(M, K) if A.dtype == FP8_DTYPE else quantize_rows_static(A.reshape(M, K), As) - out = _scaled_mm_2d_op(Aq, As, B, Bs.reshape(-1).float(), None, None, "fp8") + if B.stride(-1) != 1 or (B.storage_offset() * B.element_size()) % 16: + return _smm_reject("weight-layout") + if rowwise: + Aq, As, _, _ = tensor_wide_act_operands( + A.reshape(M, K), As, activation_scale_ub=activation_scale_ub + ) + else: + As = As.reshape(-1).float() + Aq = A.reshape(M, K) if A.dtype == FP8_DTYPE else quantize_rows_static(A.reshape(M, K), As) + if (Aq.storage_offset() * Aq.element_size()) % 16: + return _smm_reject("activation-alignment") + out = _scaled_mm_2d(Aq, As.contiguous(), B, Bs.contiguous(), None, None, "fp8") return out.reshape(*A.shape[:-1], B.shape[0]) @@ -1898,7 +2590,7 @@ def _torch_scaled_mm_2d(A, B, As, Bs, activation_format, a_global_scale, b_globa Aq = A.reshape(M, K // 2) if As is not None else None if Aq is None: Aq, As = nvfp4_act_quant(A.reshape(M, K), swizzled=True, global_scale=a_global_scale)[:2] - out = _scaled_mm_2d_op(Aq, As, B, Bs, a_global_scale, b_global_scale, fmt) + out = _scaled_mm_2d(Aq, As, B, Bs, a_global_scale, b_global_scale, fmt) return out.reshape(*A.shape[:-1], B.shape[0]) @@ -1910,6 +2602,7 @@ def matmul_2d( *, bias: torch.Tensor | None = None, # (E, N_out) per-expert output bias; (N_out,) for 2D activation_format: str | None = None, + activation_scale_ub: float | None = None, gate: bool = False, act_fn: str = "silu", swiglu_alpha: float | None = None, @@ -1920,6 +2613,11 @@ def matmul_2d( a_global_scale: torch.Tensor | None = None, b_global_scale: torch.Tensor | None = None, output_global_scale: torch.Tensor | None = None, + b_zero: torch.Tensor | None = None, + b_perm: torch.Tensor | None = None, + a_pre_quant_scale: torch.Tensor | None = None, + svdquant_lora_a: torch.Tensor | None = None, + svdquant_lora_b: torch.Tensor | None = None, ) -> torch.Tensor | list[torch.Tensor]: """Dense (2D) quantized matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4) — the single-GEMM sibling of ``matmul_grouped``/``matmul_batched``, taking the same ``As``/``Bs`` scale spec, @@ -1934,7 +2632,8 @@ def matmul_2d( accumulator (``a_global_scale`` rides raw ``A`` or a pre-quantized ``As`` alike). ``output_global_scale`` (NVFP4 requant only) is the NEXT proj's provided ``input_scale``: the fused requant normalizes the GLU intermediate by it, returning ``(C, Cs)`` the down-proj - consumes as ``As=Cs, a_global_scale=output_global_scale``. + consumes as ``As=Cs, a_global_scale=output_global_scale``. ``b_zero`` is an integer weight's zero points + (``None``: symmetric) and ``b_perm`` its act-order K order (``None``: natural). ``activation_format`` is the activations' quantization: ``None`` (default) = the weights' own format, ``"bf16"`` = weight-only (no activation quant), or an explicit ``"fp8"`` / ``"mxfp8"`` / @@ -1944,21 +2643,26 @@ def matmul_2d( ``[..., N]`` intermediate; ``quantize_output=True`` (MX weights only) stores that intermediate in ``activation_format`` and returns ``(C, Cs)`` for the next op. - Routes by weight dtype + inferred block: ``Bs`` None → full-precision; MX/NVFP4 (int8/E4M3 weights - with UE8M0 group-32 or E4M3 group-16 ``Bs``) → ``mx_dynamic_matmul``; tensor-wide FP8 (scalar - ``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``; block FP8 → the static variant when ``As`` is a - per-tensor scalar, else dynamic. + ``svdquant_lora_a`` (rank, K) and ``svdquant_lora_b`` (N_out, rank) add ModelOpt's + BF16 low-rank correction to dense NVFP4 W4A4. Both branches read the smoothed input before + quantization. The correction is fused before GLU/requant; gate/up factor rows use the same + interleaved order as B. ``simulate_unfused=True`` also matches ModelOpt's materialized + BF16 linear outputs and addition; otherwise the final sum stays FP32 until the output cast. + This recipe currently supports inference only. + + Routes by weight dtype + inferred block: ``Bs`` None → full-precision; uint32 words → integer GEMM; MX/NVFP4 (int8/E4M3 weights + with UE8M0 group-32 or E4M3 group-16 ``Bs``) → ``mx_dynamic_matmul``; tensor-wide or row-wise FP8 + (scalar or ``(N, 1)`` ``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``; block FP8 → the static variant + when ``As`` is a per-tensor scalar, else dynamic. + + ``activation_scale_ub`` caps the dynamic per-token amax for row-wise FP8; + it requires raw activations and applies before conversion to E4M3. """ - check_activation_format(activation_format, quantize_output) - assert (a_global_scale is None and b_global_scale is None) or ( - Bs is not None and weight_format(B, Bs) == "nvfp4" - ), "two-level globals (a_global_scale / b_global_scale) are NVFP4-only" - assert output_global_scale is None or ( - Bs is not None and quantize_output and resolve_activation_format(activation_format, B, Bs) == "nvfp4" - ), ( - "output_global_scale is the NVFP4 requant second level — it requires quantize_output=True on NVFP4 " - "(the epilogue would otherwise normalize by it with nothing downstream to compensate)" + _validate_quantization_args( + A, As, B, Bs, a_global_scale, b_global_scale, output_global_scale, + b_zero, b_perm, a_pre_quant_scale, activation_format, quantize_output, activation_scale_ub, ) + validate_svdquant_operands(A, B, As, Bs, activation_format, output_dtype, svdquant_lora_a, svdquant_lora_b, a_pre_quant_scale) def _unwrap(ret: list[torch.Tensor]) -> torch.Tensor | tuple[torch.Tensor, ...]: # tuple, matching the routed siblings — callers branch on isinstance(out, tuple) @@ -1973,6 +2677,13 @@ def _unwrap(ret: list[torch.Tensor]) -> torch.Tensor | tuple[torch.Tensor, ...]: A, B, output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, bias=bias) ) + if B.dtype in INT_WEIGHT_DTYPES: # integer words: weight-only + return _unwrap( + int_weight_only_matmul_2d( + A, B, Bs, output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, + b_zero=b_zero, bias=bias, b_perm=b_perm, a_pre_quant_scale=a_pre_quant_scale, a_scale=As, b_scale2=b_global_scale) + ) + if is_mx(B, Bs): if activation_format == "bf16": # weight-only: raw bf16 activation, MX weight upcast in-MMA assert As is None and a_global_scale is None and not quantize_output, ( @@ -1990,7 +2701,15 @@ def _unwrap(ret: list[torch.Tensor]) -> torch.Tensor | tuple[torch.Tensor, ...]: mx_dynamic_matmul( A, B, As, Bs, output_dtype, activation_format, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, quantize_output, - a_global_scale, b_global_scale, output_global_scale, bias=bias) + a_global_scale, b_global_scale, output_global_scale, bias=bias, a_pre_quant_scale=a_pre_quant_scale, + svdquant_lora_a=svdquant_lora_a, svdquant_lora_b=svdquant_lora_b) + ) + if activation_format == "bf16": + assert As is None and not quantize_output, "weight-only takes raw activations without requantization" + return _unwrap( + block_fp8_weight_only_matmul_2d( + A, B, Bs, output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit, + simulate_unfused, bias=bias) ) # FP8 activations are always E4M3 with the weight-implied scale granularity, so "fp8" is a # no-op format name (accepted for symmetry with the MoE ops); no other name applies here. @@ -2002,16 +2721,17 @@ def _unwrap(ret: list[torch.Tensor]) -> torch.Tensor | tuple[torch.Tensor, ...]: # dim K tiles evenly, so block_k is exact; N may be non-aligned (a partial last block, so # N % Bs.shape[0] != 0 and plain division would under-count) — recover block_n from the even dim, # the square-block convention for these FP8 weights. - if Bs.numel() == 1: + if Bs.numel() == 1 or is_rowwise(B, Bs): block_size = None else: block_k = B.shape[1] // Bs.shape[1] block_n = B.shape[0] // Bs.shape[0] if B.shape[0] % Bs.shape[0] == 0 else block_k block_size = [block_n, block_k] - if block_size is None: # tensor-wide (per-tensor) scale + if block_size is None: # K-constant weight scale: per tensor or per row return _unwrap( w8a8_tensor_dynamic_fp8_matmul( - A, B, Bs, As, output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, bias=bias) + A, B, Bs, As, output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, + bias=bias, activation_scale_ub=activation_scale_ub) ) # Block-wise FP8: a per-tensor scalar As is the static (calibrated) activation scale; else dynamic. if As is not None: diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/mma.py b/finegrained-kernels/torch-ext/finegrained_kernels/mma.py index 2efc0a250..2e17b6b78 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/mma.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/mma.py @@ -15,8 +15,9 @@ import triton import triton.language as tl -from .quant import e2m1_cols_to_e4m3, e2m1_cols_to_f32, e2m1_to_bf16, e2m1_to_e4m3, e2m1_to_f32 +from .quant import e2m1_cols_to_bf16, e2m1_cols_to_e4m3, e2m1_cols_to_f32, e2m1_to_bf16, e2m1_to_e4m3, e2m1_to_f32 from .loading.scales import decode_group_scale +from .compat import get_active_device_type, is_sm10x @triton.jit @@ -98,10 +99,11 @@ def mx_dot_rescale( @triton.jit def mx_weight_upcast(w, w_scale, BLOCK_SIZE_K: tl.constexpr, N: tl.constexpr, SCALE_GROUP_K: tl.constexpr, - OUT_DTYPE: tl.constexpr = tl.bfloat16): + OUT_DTYPE: tl.constexpr = tl.bfloat16, ROWS_MAJOR: tl.constexpr = False): """Upcast one MXFP4/MXFP8 weight K-tile ``[BK, N]`` to ``OUT_DTYPE`` for the weight-only path: unpack E2M1 -> bf16 directly (fp8 passes through a cast) and apply each ``SCALE_GROUP_K``-row K-group's - UE8M0 group scale. ``w_scale`` is ``[N, BK // SCALE_GROUP_K]`` (per-N-row, per-group); transposed + UE8M0 group scale. ``ROWS_MAJOR`` decodes and returns ``[N, BK]`` directly for the + computed MMA lhs; the default returns ``[BK, N]``. ``w_scale`` is ``[N, BK // SCALE_GROUP_K]`` (per-N-row, per-group); transposed and broadcast across each group's rows to ``[BK, N]``. Unlike the ``dot`` arm (BK == group, scale folded onto the [M,N] product) this dequantizes IN the tile, so BK spans any number of groups — a full-BK bf16 ``tl.dot`` (the matmul_ogs format: fp4 weight, bf16 acts). @@ -111,15 +113,16 @@ def mx_weight_upcast(w, w_scale, BLOCK_SIZE_K: tl.constexpr, N: tl.constexpr, SC into the bits as an int add + zero-mask ``tl.where`` measured ~13% SLOWER (gate_up 4.68->5.29ms) — the mask + uint16<->int32 bitcast chain costs more than the multiply it replaces.""" ng: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K - wq = e2m1_to_bf16(w) if w.dtype == tl.uint8 else w.to(tl.bfloat16) # [BK, N] bf16 (direct, no E4M3) - # Reshape the (real, contiguous) weight to [ng, g, N] and broadcast the group scale [ng, 1, N] - # across each group's g rows in the multiply, then reshape back to [BK, N]. (Reshaping a - # broadcasted stride-0 tensor doesn't lower, so broadcast happens IN the op, not before it.) - ws = tl.trans(decode_group_scale(w_scale)).to(tl.bfloat16)[:, None, :] # [ng, 1, N] - w3 = wq.reshape(ng, SCALE_GROUP_K, N) * ws - # dequant + scale happen in bf16 (exact: E2M1/E4M3 codes and power-of-two UE8M0 scales all fit), - # then widen if the activation is wider — bf16->fp32 is lossless and keeps the multiply cheap. - return w3.reshape(BLOCK_SIZE_K, N).to(OUT_DTYPE) + if ROWS_MAJOR: + wq = e2m1_cols_to_bf16(w) if w.dtype == tl.uint8 else w.to(tl.bfloat16) + ws = decode_group_scale(w_scale).to(tl.bfloat16)[:, :, None] + out = (wq.reshape(N, ng, SCALE_GROUP_K) * ws).reshape(N, BLOCK_SIZE_K) + else: + wq = e2m1_to_bf16(w) if w.dtype == tl.uint8 else w.to(tl.bfloat16) + ws = tl.trans(decode_group_scale(w_scale)).to(tl.bfloat16)[:, None, :] + out = (wq.reshape(ng, SCALE_GROUP_K, N) * ws).reshape(BLOCK_SIZE_K, N) + # Decode and group scaling round in bf16 before the activation-dtype cast, in either orientation. + return out.to(OUT_DTYPE) @triton.jit @@ -134,23 +137,24 @@ def mx_weight_only_compute( SCALE_GROUP_K: tl.constexpr, SWAP_AB: tl.constexpr = False, ): - """Weight-only (W4A16/W8A16) MMA step: raw bf16/fp16 activations against an MXFP4/MXFP8 weight - tile, dispatched on ``COMPUTE_MODE``. ``"dot_scaled"`` hands the packed weight and its UE8M0 - group scales straight to the tcgen05 scaled MMA with an UNSCALED activation operand (the - ``lhs_scale=None`` / bf16-format form), so the tensor core does the fp4 decode and the group - rescale; ``"dot"`` upcasts the tile to the activation dtype in-loop and runs a plain - ``tl.dot`` (the matmul_ogs Hopper format). Both are correct everywhere and the tuner picks per - workload — the two differ by arm: measured on gpt-oss K=2880, ``dot_scaled`` wins the down - projection by 18.5% while the stacked gate|up tile prefers ``dot`` by 7.3``%``. ``SWAP_AB`` - is the decode form: weight output rows lead the tile and the FMA reduce replaces the MMA - (both dot arms keep the token in M). Single return — only the taken branch compiles.""" - if SWAP_AB: + """Weight-only MMA step with raw activations and MXFP4/MXFP8 weights. + + ``dot_scaled`` delegates weight upcasting and group scaling to Triton; an unscaled + BF16 activation does not imply native block-scaled MMA. ``dot`` explicitly upcasts + the weight tile before ``tl.dot``; the tuner chooses between these lowerings. + ``SWAP_AB`` puts weight rows on the MMA lhs; scalar decode instead reduces + each weight row against one token. Callers restore output orientation after the K loop. + """ + if SWAP_AB and COMPUTE_MODE == "scalar": acc = mx_weight_only_scalar_swapped( acc, tl.reshape(a, (BLOCK_SIZE_K,)), w, w_scale, N_WIDTH, BLOCK_SIZE_K, SCALE_GROUP_K, ) elif COMPUTE_MODE == "dot_scaled": acc = mx_dot_scaled(acc, a, None, w, w_scale) + elif SWAP_AB: + lhs = mx_weight_upcast(w, w_scale, BLOCK_SIZE_K, N_WIDTH, SCALE_GROUP_K, a.dtype, ROWS_MAJOR=True) + acc = tl.dot(lhs, tl.trans(a), acc) else: acc = acc + tl.dot( a, mx_weight_upcast(w, w_scale, BLOCK_SIZE_K, N_WIDTH, SCALE_GROUP_K, a.dtype) @@ -186,6 +190,467 @@ def mx_weight_only_scalar_swapped( ) +# PTX inline asm (round_f32_to_bf16_f32, the packed zero-point subtract) on the CUDA backend; elsewhere plain Triton +_PTX = tl.constexpr(get_active_device_type() == "cuda") +# Keep architecture-specific integer fast paths on the measured sm_10x target. +_SM10X = tl.constexpr(is_sm10x()) + + +@triton.jit +def int_word_pairs(w, BITS: tl.constexpr, K_MAJOR: tl.constexpr): + """The value pairs of uint32 words in the kernels' INTERLEAVED layout (``interleave_int_words``): of a word's + ``32 // BITS`` values, position ``s`` holds the value at K offset ``2s`` and position ``s + 16 // BITS`` the one + at ``2s + 1``, so one shift and one mask put K-adjacent values in the low and high 16-bit halves of a lane, + the register pair an MMA operand holds. ``w`` ``[R, C]`` words along C (``K_MAJOR``: along R); returns the raw + pairs ``[R, C, 16 // BITS]`` (``K_MAJOR``: ``[R, 16 // BITS, C]``), pair ``s`` = K offsets ``(2s, 2s + 1)``.""" + HALF: tl.constexpr = 16 // BITS + MASK: tl.constexpr = (1 << BITS) - 1 + shifts = tl.arange(0, HALF) * BITS + if K_MAJOR: + pairs = (w[:, None, :] >> shifts[None, :, None]) & (MASK | (MASK << 16)) + else: + pairs = (w[:, :, None] >> shifts[None, None, :]) & (MASK | (MASK << 16)) + return pairs + + +@triton.jit +def int8_words_to_int8(w, K_MAJOR: tl.constexpr): + """Offset-binary interleaved words to signed INT8, in natural K order.""" + pairs = int_word_pairs(w ^ 0x80808080, 8, K_MAJOR) + v = tl.join((pairs & 0xFF).to(tl.int8), (pairs >> 16).to(tl.int8)) + if K_MAJOR: + out = tl.trans(v, 0, 1, 3, 2).reshape(4 * w.shape[0], w.shape[1]) + else: + out = v.reshape(w.shape[0], 4 * w.shape[1]) + return out + + +@triton.jit +def int_words_to_offset(w, zero_words, BITS: tl.constexpr, K_MAJOR: tl.constexpr, WHOLE_WORD: tl.constexpr = False, FP16: tl.constexpr = False): + """Interleaved uint32 words (``int_word_pairs``) to their values as ``OFFSET + q``, exactly, with no integer-to-float + conversion: a value ORed under the offset's exponent bits IS that float, bf16 ``128 + q`` for 2- and 4-bit values + (q < 128), fp16 ``1024 + q`` for FP16 scales or 8-bit values (bf16 cannot hold 128 + 255). ``zero_words`` (``None``: no + zero point) is one word per weight word ``[R, C]`` carrying its group's zero point ``z`` in both 16-bit lanes + (``int_zero_words``): on CUDA the float ``-(OFFSET + z)``, added to each pair by one packed FMA, so the values + are ``q - z``; elsewhere the integer lift ``2 ** BITS - z`` added before the OR, so they are + ``OFFSET + 2 ** BITS + q - z``. Either way the zero point costs no per-value float op of its own. Output-rows + ``[R, C]`` -> ``[R, (32 // BITS) C]`` (``K_MAJOR``: ``[R, C]`` -> ``[(32 // BITS) R, C]``), natural K order.""" + MAGIC: tl.constexpr = 0x64006400 if FP16 or BITS == 8 else 0x43004300 + if _PTX and BITS == 8: + # Keep both byte pairs with their word before orienting the MMA tile. + lo, hi = tl.inline_asm_elementwise( + "prmt.b32 $0, $2, 0x64646464, 0x5250; prmt.b32 $1, $2, 0x64646464, 0x5351;", + "=r,=r,r", [w], dtype=(tl.uint32, tl.uint32), is_pure=True, pack=1, + ) + pairs = tl.join(lo, hi) + if K_MAJOR: + pairs = tl.trans(pairs, 0, 2, 1) + elif WHOLE_WORD: + # Keep all four pairs in the word's lane. A shared mask avoids materializing one per asm expansion. + mask = tl.full(w.shape, 0x000f000f, tl.uint32) + exponent: tl.constexpr = "0x64006400" if FP16 else "0x43004300" + p0, p1, p2, p3 = tl.inline_asm_elementwise( + "{ .reg .b32 t1, t2, t3, mask; mov.b32 mask, $5; " + "shr.u32 t1, $4, 4; shr.u32 t2, $4, 8; shr.u32 t3, $4, 12; " + "lop3.b32 $0, $4, mask, " + exponent + ", 0xEA; " + "lop3.b32 $1, t1, mask, " + exponent + ", 0xEA; " + "lop3.b32 $2, t2, mask, " + exponent + ", 0xEA; " + "lop3.b32 $3, t3, mask, " + exponent + ", 0xEA; }", + "=r,=r,=r,=r,r,r", [w, mask], dtype=(tl.uint32,) * 4, is_pure=True, pack=1, + ) + pairs = tl.join(tl.join(p0, p2), tl.join(p1, p3)).reshape(w.shape[0], w.shape[1], 4) + if K_MAJOR: + pairs = tl.trans(pairs, 0, 2, 1) + else: + pairs = int_word_pairs(w, BITS, K_MAJOR) + if zero_words is not None: + zero_words = tl.broadcast_to(zero_words[:, None, :] if K_MAJOR else zero_words[:, :, None], pairs.shape) + if zero_words is not None and not _PTX: + pairs += zero_words + if not WHOLE_WORD: + pairs = pairs | MAGIC + if zero_words is not None and _PTX: + # fma(pair, 1, -(OFFSET + z)): the exact difference; fma.rn.bf16x2 / f16x2 lower from sm_80 / sm_53 + pairs = tl.inline_asm_elementwise( + "{ .reg .b32 one; mov.b32 one, 0x3F803F80; fma.rn.bf16x2 $0, $1, one, $2; }" if BITS < 8 and not FP16 + else "{ .reg .b32 one; mov.b32 one, 0x3C003C00; fma.rn.f16x2 $0, $1, one, $2; }", + "=r,r,r", [pairs, zero_words], dtype=tl.uint32, is_pure=True, pack=1, + ) + ODT: tl.constexpr = tl.float16 if FP16 or BITS == 8 else tl.bfloat16 + lo = (pairs & 0xFFFF).to(tl.uint16).to(ODT, bitcast=True) + hi = (pairs >> 16).to(tl.uint16).to(ODT, bitcast=True) + v = tl.join(lo, hi) # [.., pair, half] (K_MAJOR: [R, pair, C, half]) + if K_MAJOR: + out = tl.trans(v, 0, 1, 3, 2).reshape((32 // BITS) * w.shape[0], w.shape[1]) + else: + out = v.reshape(w.shape[0], (32 // BITS) * w.shape[1]) + return out + + +@triton.jit +def int8_words_to_bf16(w, K_MAJOR: tl.constexpr): + """Interleaved INT8 words to exact bf16 ``q - 128``. + ``K_MAJOR`` maps ``[K // 4, N]`` to ``[K, N]``; otherwise ``[N, K // 4]`` to ``[N, K]``. + Return both K-adjacent pairs together to keep a whole word per lane and enable ldmatrix shared loads.""" + # 128+(q&127) minus 128 or 256 is exact, without fp16 -> fp32 -> bf16 conversion. + p0, p1 = tl.inline_asm_elementwise( + "{ .reg .b32 q, z, t; shr.u32 t, $2, 8; lop3.b32 q, $2, 0x007f007f, 0x43004300, 0xEA; " + "lop3.b32 z, $2, 0x00800080, 0x43804380, 0x6A; sub.rn.bf16x2 $0, q, z; " + "lop3.b32 q, t, 0x007f007f, 0x43004300, 0xEA; " + "lop3.b32 z, t, 0x00800080, 0x43804380, 0x6A; sub.rn.bf16x2 $1, q, z; }", + "=r,=r,r", [w], dtype=(tl.uint32, tl.uint32), is_pure=True, pack=1, + ) + pairs = tl.join(p0, p1) + if K_MAJOR: + pairs = tl.trans(pairs, 0, 2, 1) + lo = (pairs & 0xffff).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (pairs >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + v = tl.join(lo, hi) + if K_MAJOR: + out = tl.trans(v, 0, 1, 3, 2).reshape(4 * w.shape[0], w.shape[1]) + else: + out = v.reshape(w.shape[0], 4 * w.shape[1]) + return out + + +@triton.jit +def int_zero_words(w_zero, WORDS_PER_GROUP: tl.constexpr, BITS: tl.constexpr, K_MAJOR: tl.constexpr, FP16: tl.constexpr = False): + """``[R, G]`` per-group zero points to ``int_words_to_offset``'s ``[R, G * WORDS_PER_GROUP]`` zero words + (``K_MAJOR``: ``[G, R]`` to ``[G * WORDS_PER_GROUP, R]``): CUDA, the packed float ``-(OFFSET + z)`` pair (bf16 + ``0xC300 | z``, fp16 ``0xE400 | z``); elsewhere the lift ``2 ** BITS - z`` in both lanes.""" + z = w_zero.to(tl.uint32) + if _PTX: + lane = (z | (0xC300 if BITS < 8 and not FP16 else 0xE400)) * 0x00010001 + else: + lane = ((1 << BITS) - z) * 0x00010001 + if K_MAJOR: + G: tl.constexpr = w_zero.shape[0] + R: tl.constexpr = w_zero.shape[1] + out = tl.broadcast_to(lane[:, None, :], (G, WORDS_PER_GROUP, R)).reshape(G * WORDS_PER_GROUP, R) + else: + R: tl.constexpr = w_zero.shape[0] + G: tl.constexpr = w_zero.shape[1] + out = tl.broadcast_to(lane[:, :, None], (R, G, WORDS_PER_GROUP)).reshape(R, G * WORDS_PER_GROUP) + return out + + +@triton.jit +def int_dequant(q, zero, scale, OUT_DTYPE: tl.constexpr, divisor=None): + """``((q - zero) * scale)`` rounded in the scale's dtype, then cast to ``OUT_DTYPE``: the dequantize-to-a-float + reference (compressed-tensors' ``_dequantize``, what transformers and Marlin users compare against), bit for bit. + ``q`` holds ``q - z`` plus a constant offset and ``zero`` (a scalar) that offset, both exact in ``q``'s dtype; + their difference is exact in every float dtype, and a bf16/fp16 scale's product with it is exact in fp32, so its + one rounding is the reference's.""" + qz = (q - zero).to(scale.dtype) + value = qz * scale + if divisor is not None: + # ModelOpt's second sequential quantizer rounds the reconstructed INT4 weight to E4M3. + value = tl.clamp(value.to(tl.float32) / divisor, -448.0, 448.0) + return value.to(OUT_DTYPE) + + +@triton.jit +def int_dequant_tile( + w, w_scale, w_zero, ROWS_MAJOR: tl.constexpr, SCALE_GROUP_K: tl.constexpr, BITS: tl.constexpr, + OUT_DTYPE: tl.constexpr, WHOLE_WORD: tl.constexpr = False, divisor=None, +): + """One ``BITS``-wide integer weight K-tile of interleaved uint32 words (``int_word_pairs``), output-rows-major + ``[N, BK // VPW]`` (``ROWS_MAJOR``) or K-major ``[BK // VPW, N]`` (VPW = ``32 // BITS`` values per word), + dequantized (``int_dequant``) into the K-major ``[BK, N]`` MMA rhs. The tile spans ``BK // SCALE_GROUP_K`` whole + groups: ``w_scale`` and ``w_zero`` (``None``: symmetric, zero point ``2 ** (BITS - 1)``) are + ``[N, BK // SCALE_GROUP_K]``.""" + VPW: tl.constexpr = 32 // BITS + FP16: tl.constexpr = w_scale.dtype == tl.float16 + OFFSET: tl.constexpr = 128.0 if BITS < 8 and not FP16 else 1024.0 # int_words_to_offset + # what int_words_to_offset leaves on top of q - z + if w_zero is None: + ZERO: tl.constexpr = OFFSET + (1 << (BITS - 1)) + elif _PTX: + ZERO: tl.constexpr = 0.0 + else: + ZERO: tl.constexpr = OFFSET + (1 << BITS) + WPG: tl.constexpr = SCALE_GROUP_K // VPW # words per scale group + if ROWS_MAJOR: + N: tl.constexpr = w.shape[0] + BK: tl.constexpr = VPW * w.shape[1] + NG: tl.constexpr = BK // SCALE_GROUP_K + if _SM10X and BITS == 8 and w_scale.dtype == tl.bfloat16: + q = int8_words_to_bf16(w, False).reshape(N, NG, SCALE_GROUP_K) + center = 0 if w_zero is None else (w_zero.to(tl.bfloat16) - 128)[:, :, None] + v = int_dequant(q, center, w_scale[:, :, None], OUT_DTYPE, divisor) + else: + zero_words = None if w_zero is None else int_zero_words(w_zero, WPG, BITS, False, FP16) + q = int_words_to_offset(w, zero_words, BITS, False, WHOLE_WORD, FP16).reshape(N, NG, SCALE_GROUP_K) + v = int_dequant(q, ZERO, w_scale[:, :, None], OUT_DTYPE, divisor) + + out = tl.trans(v.reshape(N, BK)) + else: + BK: tl.constexpr = VPW * w.shape[0] + N: tl.constexpr = w.shape[1] + NG: tl.constexpr = BK // SCALE_GROUP_K + if _SM10X and BITS == 8 and w_scale.dtype == tl.bfloat16: + q = int8_words_to_bf16(w, True).reshape(NG, SCALE_GROUP_K, N) + center = 0 if w_zero is None else (tl.trans(w_zero).to(tl.bfloat16) - 128)[:, None, :] + v = int_dequant(q, center, tl.trans(w_scale)[:, None, :], OUT_DTYPE, divisor) + else: + zero_words = None if w_zero is None else int_zero_words(tl.trans(w_zero), WPG, BITS, True, FP16) + q = int_words_to_offset(w, zero_words, BITS, True, WHOLE_WORD, FP16).reshape(NG, SCALE_GROUP_K, N) + v = int_dequant(q, ZERO, tl.trans(w_scale)[:, None, :], OUT_DTYPE, divisor) + out = v.reshape(BK, N) + return out + + +@triton.jit +def round_f32_to_bf16_f32(x): + """fp32 ``x`` rounded to bf16 (round-to-nearest-even) and widened back, two lanes per packed convert + (``cvt.rn.bf16x2.f32``) instead of one ``F2F`` per value.""" + return tl.inline_asm_elementwise( + "{ .reg .b32 d; cvt.rn.bf16x2.f32 d, $3, $2; shl.b32 $0, d, 16; and.b32 $1, d, 0xFFFF0000; }", + "=r,=r,r,r", + [x.to(tl.uint32, bitcast=True)], + dtype=tl.uint32, + is_pure=True, + pack=2, + ).to(tl.float32, bitcast=True) + + +def _int_word_dot_asm(rows: int, bits: int, fp16: bool = False) -> str: + """Pair-unpack interleaved 2/4/8-bit words, dequantize in bf16/fp16, then FMA into each token row's fp32 sum. + Operands: row sums, word, packed zero/scale, then ``16 // bits`` activation pairs per row.""" + w, z, sc = rows, rows + 1, rows + 2 + dtype = "f16" if fp16 else "bf16" + magic = "0x64006400" if fp16 else "0x43004300" + lines = [ + "{", + ".reg .b32 t, x;", + ".reg .b16 wl, wh, al, ah;", + ".reg .b32 " + ", ".join(f"c{m}" for m in range(rows)) + ";", + f"mov.b32 t, ${w};", + ] + [f"mov.b32 c{m}, 0;" for m in range(rows)] + for p in range(16 // bits): + lines += [ + f"lop3.b32 x, t, {((1 << bits) - 1) * 0x10001:#010x}, {magic}, 0xEA;", + f"sub.rn.{dtype}x2 x, x, ${z}; mul.rn.{dtype}x2 x, x, ${sc}; mov.b32 {{wl, wh}}, x;", + ] + for m in range(rows): + lines += [ + f"mov.b32 {{al, ah}}, ${rows + 3 + (16 // bits) * m + p};", + f"fma.rn.f32.{dtype} c{m}, wl, al, c{m}; fma.rn.f32.{dtype} c{m}, wh, ah, c{m};", + ] + if p < 16 // bits - 1: + lines.append(f"shr.u32 t, t, {bits};") + lines += [f"mov.b32 ${m}, c{m};" for m in range(rows)] + ["}"] + return "\n".join(lines) + + +# BF16 supports widths (2, 4), FP16 (2, 4, 8), each with token-row counts (1, 2, 4). +_INT_WORD_DOT_ASMS = tl.constexpr(tuple( + _int_word_dot_asm(rows, bits, fp16) + for fp16 in (False, True) for bits in ((2, 4, 8) if fp16 else (2, 4)) for rows in (1, 2, 4) +)) +_INT_WORD_DOT_CONSTRAINTS = tl.constexpr(tuple( + ",".join(["=r"] * rows + ["r"] * (3 + (16 // bits) * rows)) + for bits in (2, 4, 8) for rows in (1, 2, 4) +)) + + +@triton.jit +def _bcast(x, shape: tl.constexpr): + return tl.broadcast_to(x[None, :], shape) + + +@triton.jit +def _row_split(p, BM: tl.constexpr): + """``[BM, KW]`` -> ``BM`` tensors ``[KW]`` (row order), BM in {1, 2, 4}.""" + if BM == 1: + out = (p.reshape(p.shape[1]),) + elif BM == 2: + r0, r1 = tl.split(tl.trans(p)) + out = (r0, r1) + else: + e, o = tl.split(tl.trans(p).reshape(p.shape[1], 2, 2)) # row = 2 * hi + lo: split lo + r0, r2 = tl.split(e) + r1, r3 = tl.split(o) + out = (r0, r1, r2, r3) + return out + + +@triton.jit +def int_words_dot(w, zero2, scale2, a, BM: tl.constexpr, BITS: tl.constexpr): + """Interleaved 2/4/8-bit words ``[R, KW]`` against ``BM`` bf16/fp16 token rows. Dequantize each pair once for all + rows with the scale's rounding, then accumulate the products in fp32 on sm_100. + ``zero2`` holds packed ``offset + zero`` (128 for bf16, 1024 for fp16); ``scale2`` holds each word's packed group scale. + Returns fp32 ``[BM, R]`` for ``BM`` in {1, 2, 4}.""" + R: tl.constexpr = w.shape[0] + KW: tl.constexpr = w.shape[1] + HALF: tl.constexpr = 16 // BITS + shape: tl.constexpr = (R, KW) + lo, hi = tl.split(a.to(tl.uint16, bitcast=True).to(tl.uint32).reshape(BM, KW, HALF, 2)) + packed = lo | (hi << 16) + if BITS == 8: + p0, p1 = tl.split(packed.reshape(BM, KW, 2)) + pairs = (p0, p1) + elif BITS == 4: + even, odd = tl.split(packed.reshape(BM, KW, 2, 2)) + p0, p2 = tl.split(even) + p1, p3 = tl.split(odd) + pairs = (p0, p1, p2, p3) + else: + even, odd = tl.split(packed.reshape(BM, KW, 4, 2)) + ee, eo = tl.split(even.reshape(BM, KW, 2, 2)) + oe, oo = tl.split(odd.reshape(BM, KW, 2, 2)) + p0, p4 = tl.split(ee) + p2, p6 = tl.split(eo) + p1, p5 = tl.split(oe) + p3, p7 = tl.split(oo) + pairs = (p0, p1, p2, p3, p4, p5, p6, p7) + row_pairs = () + for p in tl.static_range(HALF): + row_pairs += (_row_split(pairs[p], BM),) + operands = (w.to(tl.uint32, bitcast=True), zero2, scale2) + for m in tl.static_range(BM): + for p in tl.static_range(HALF): + operands += (_bcast(row_pairs[p][m], shape),) + IDX: tl.constexpr = (0 if BITS == 2 else (3 if BITS == 4 else 6)) + (0 if BM == 1 else (1 if BM == 2 else 2)) + ASM_IDX: tl.constexpr = IDX + (6 if a.dtype == tl.float16 else 0) + out = tl.inline_asm_elementwise( + _INT_WORD_DOT_ASMS[ASM_IDX], _INT_WORD_DOT_CONSTRAINTS[IDX], operands, + dtype=(tl.uint32,) * BM, is_pure=True, pack=1, + ) + sums = () + for m in tl.static_range(BM): + sums += (tl.sum(out[m].to(tl.float32, bitcast=True), axis=1),) + if BM == 1: + result = sums[0][None, :] + elif BM == 2: + result = tl.trans(tl.join(sums[0], sums[1])) + else: + result = tl.permute(tl.join(tl.join(sums[0], sums[1]), tl.join(sums[2], sums[3])), (2, 1, 0)).reshape(4, R) + return result + + +@triton.jit +def int8_words_dot(w, a): + """Signed INT8 dot from offset-binary interleaved words, accumulating exactly in INT32.""" + BM: tl.constexpr = a.shape[0] + KW: tl.constexpr = a.shape[1] // 4 + if _SM10X: + # Weight bytes hold K offsets (0, 2, 1, 3); pack the activation bytes in the same order. + even, odd = tl.split(a.to(tl.uint8, bitcast=True).to(tl.uint32).reshape(BM, KW, 2, 2)) + a0, a2 = tl.split(even) + a1, a3 = tl.split(odd) + packed_a = a0 | (a2 << 8) | (a1 << 16) | (a3 << 24) + signed_w = w.to(tl.uint32, bitcast=True) ^ 0x80808080 + prod = tl.inline_asm_elementwise( + "dp4a.s32.s32 $0, $1, $2, 0;", "=r,r,r", [signed_w[None, :, :], packed_a[:, None, :]], + dtype=tl.int32, is_pure=True, pack=1, + ) + else: + rhs = int8_words_to_int8(w, False) + prod = a.to(tl.int32)[:, None, :] * rhs.to(tl.int32)[None, :, :] + return tl.sum(prod, axis=2) + + +@triton.jit +def int_weight_only_scalar( + acc, + a, + w, + w_scale, + w_zero, + ROWS_W: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SCALE_GROUP_K: tl.constexpr, + BITS: tl.constexpr, +): + """The integer weight-only CUDA-core reduce: ``w`` is output-rows-major ``[ROWS_W, BK // VPW]`` interleaved int32 + words, dequantized to ``a``'s dtype with ``int_dequant``'s rounding once for all ``BM`` token rows of ``a`` + ``[BM, BK]``, and reduced in fp32. Matching bf16/fp16 on sm_100 (2/4 bits, plus 8 bits for fp16) runs ``int_words_dot``; otherwise the dequant is + computed in fp32 (the reduce's operand dtype). Returns ``acc + [BM, ROWS_W]``.""" + NG: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K + VPW: tl.constexpr = 32 // BITS + KW: tl.constexpr = BLOCK_SIZE_K // VPW + BM: tl.constexpr = a.shape[0] + if ( + _SM10X and (BITS == 2 or BITS == 4 or a.dtype == tl.float16) and (BM == 1 or BM == 2 or BM == 4) + and a.dtype == w_scale.dtype and (a.dtype == tl.bfloat16 or a.dtype == tl.float16) + ): + WPG: tl.constexpr = SCALE_GROUP_K // VPW # words per scale group + s16 = w_scale.to(tl.uint16, bitcast=True).to(tl.uint32) + scale2 = tl.broadcast_to((s16 | (s16 << 16))[:, :, None], (ROWS_W, NG, WPG)).reshape(ROWS_W, KW) + OFFSET: tl.constexpr = 1024.0 if a.dtype == tl.float16 else 128.0 + MAGIC: tl.constexpr = 0x6400 if a.dtype == tl.float16 else 0x4300 + if w_zero is None: + zero2 = tl.full((ROWS_W, KW), (MAGIC | (1 << (BITS - 1))) * 0x10001, tl.uint32) + else: + z16 = (w_zero.to(tl.float32) + OFFSET).to(a.dtype).to(tl.uint16, bitcast=True).to(tl.uint32) + zero2 = tl.broadcast_to((z16 | (z16 << 16))[:, :, None], (ROWS_W, NG, WPG)).reshape(ROWS_W, KW) + return acc + int_words_dot(w, zero2, scale2, a, BM, BITS) + # a value ORed under fp32 2^23's exponent bits IS 2^23 + q: no int-to-float conversion + pairs = int_word_pairs(w, BITS, False) + lo = ((pairs & 0xFFFF) | 0x4B000000).to(tl.float32, bitcast=True) + hi = ((pairs >> 16) | 0x4B000000).to(tl.float32, bitcast=True) + v = tl.join(lo, hi).reshape(ROWS_W, NG, SCALE_GROUP_K) + zero = 8388608.0 + (1 << (BITS - 1)) if w_zero is None else (w_zero.to(tl.float32) + 8388608.0)[:, :, None] + # q - zero is exact, its product with a bf16/fp16 scale exact in fp32: one rounding, in the scale dtype + wt = (v - zero) * w_scale.to(tl.float32)[:, :, None] + if _PTX and w_scale.dtype == tl.bfloat16 and a.dtype == tl.bfloat16: + wt = round_f32_to_bf16_f32(wt) + else: + wt = wt.to(w_scale.dtype).to(a.dtype).to(tl.float32) + prod = wt.reshape(ROWS_W, BLOCK_SIZE_K)[None] * a.to(tl.float32)[:, None, :] + return acc + tl.sum(prod, axis=2) + + +@triton.jit +def int_weight_only_compute( + acc, + a, + w, + w_scale, + w_zero, + BLOCK_SIZE_K: tl.constexpr, + N_WIDTH: tl.constexpr, + SCALE_GROUP_K: tl.constexpr, + COMPUTE_MODE: tl.constexpr = "dot", + ROWS_MAJOR: tl.constexpr = False, + BITS: tl.constexpr = 4, + SWAP_AB: tl.constexpr = False, + weight_global=None, +): + """Integer-weight GEMM step: dequantize a ``BITS``-wide tile to the activation dtype (``int_dequant``), + then accumulate with the MMA (``"dot"``) or the small-token CUDA-core reduce (``"scalar"``). + Floating operands accumulate in FP32; SmoothQuant keeps its signed integer dot in INT32 until the epilogue. + ``ROWS_MAJOR`` is the weight tile's orientation (``int_dequant_tile``); the scalar arm's is always + output-rows-major. ``SWAP_AB`` puts token rows on the MMA's N axis; keep that accumulator orientation + across the K-loop and transpose it once before the epilogue.""" + # A tile inside a wider scale group reads one scale and zero-point column. + TILE_GROUP: tl.constexpr = min(SCALE_GROUP_K, BLOCK_SIZE_K) + if COMPUTE_MODE == "scalar" and a.dtype == tl.int8: + acc += int8_words_dot(w, a) + elif COMPUTE_MODE == "scalar": + acc = int_weight_only_scalar(acc, a, w, w_scale, w_zero, N_WIDTH, BLOCK_SIZE_K, TILE_GROUP, BITS) + else: + # Keep a complete INT4 word with its owner through either orientation. + WHOLE_WORD: tl.constexpr = _SM10X and BITS == 4 and ( + (SWAP_AB and ROWS_MAJOR) or (not ROWS_MAJOR and w_zero is None) + ) + if a.dtype == tl.int8: + # SmoothQuant's per-channel scale applies after the integer dot, not to its INT8 operands. + rhs = int8_words_to_int8(w, not ROWS_MAJOR) + if ROWS_MAJOR: + rhs = tl.trans(rhs) + else: + rhs = int_dequant_tile(w, w_scale, w_zero, ROWS_MAJOR, TILE_GROUP, BITS, a.dtype, WHOLE_WORD, weight_global) + DOT_DTYPE: tl.constexpr = tl.int32 if a.dtype == tl.int8 else tl.float32 + if SWAP_AB: + # Put the token rows on the MMA's smaller N axis (8 lanes instead of 16). + acc = tl.dot(tl.trans(rhs), tl.trans(a), acc, out_dtype=DOT_DTYPE) + else: + acc = tl.dot(a, rhs, acc, out_dtype=DOT_DTYPE) + return acc + + @triton.jit def mx_scalar_reduce( acc, @@ -296,6 +761,20 @@ def swap_pad_rhs(a, BLOCK_SIZE_K: tl.constexpr): ) +@triton.jit +def plain_dot(acc, a, b, SWAP_AB: tl.constexpr = False): + """Plain GEMM step with a logical [M, N] accumulator in either operand orientation. + + Swapping uses row-major weights [N, K] as the MMA lhs and transposes the + [M, K] activations, avoiding the padded token columns of single-token helpers. + """ + if SWAP_AB: + out = tl.trans(tl.dot(b, tl.trans(a), tl.trans(acc))) + else: + out = tl.dot(a, b, acc) + return out + + @triton.jit def fp8_dot(a, b, SWAP_AB: tl.constexpr, BLOCK_SIZE_K: tl.constexpr): """Swap-aware plain ``tl.dot`` for the fp8 (block/tensor-dynamic) paths — no scaling; the caller diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/moe.py b/finegrained-kernels/torch-ext/finegrained_kernels/moe.py index 7b41d4a90..a982c9b25 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/moe.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/moe.py @@ -44,8 +44,8 @@ import triton.language as tl from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait -from .grouped import matmul_grouped -from .batched import GATE_UNSTACK_MAX_S, matmul_batched +from .grouped import block_fp8_weight_only_matmul_grouped, matmul_grouped +from .batched import GATE_UNSTACK_MAX_S, block_fp8_weight_only_matmul_batched, matmul_batched from .bayesian_autotuner import bayesian_autotune from .compat import ( MX_SCALE_GROUP_K, @@ -56,7 +56,7 @@ device_context, pdl_launch_kwargs, ) -from .formats import get_supported_act_fns, is_mx, is_mxfp4, is_per_expert_global, mx_scale_family, normalize_global_scale, ue8m0_as_uint8, weight_format +from .formats import get_supported_act_fns, is_mx, is_per_expert_global, is_rowwise, mx_scale_family, normalize_global_scale, per_token_activations, ue8m0_as_uint8, weight_block_size, weight_format, INT_WEIGHT_DTYPES, int_weight_bits from .norm import norm_column_factor, rms_inv_rows, rms_norm_rows from .quant import _launch_act_quant, MX_ACT_QUANT, mx_act_quant_routed_scales from .scheduling import compute_grouped_scheduling, is_off_rank, off_rank @@ -192,21 +192,57 @@ def weighted_reduce( return reduced -def _validate_moe(gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale): - """gate_up and down must share the format (both MX or both block-dynamic FP8 — the - intermediate handed between them carries one quant format). Returns whether the format - is MX (the fused dispatchers branch on it); the fp8 quantization block is derived from - the scale shapes (``weight_block_size``), never passed. Scales are the pure block scales - (per the decoupled API — per-tensor globals ride as separate ``*_global_scale`` args); - the format predicates read the block scale's dtype/grouping.""" +def _validate_moe( + gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale, gate, + gate_up_zero, down_zero, +): + """Both projections share a weight format, scale grouping and integer zero-point recipe.""" + gate_up_int = gate_up_proj.dtype in INT_WEIGHT_DTYPES + down_int = down_proj.dtype in INT_WEIGHT_DTYPES + if gate_up_int or down_int: + error = "gate_up_proj and down_proj must use the same integer quantization recipe" + if not (gate_up_int and down_int) or gate_up_proj_scale is None or down_proj_scale is None: + raise ValueError(error) + hidden = down_proj.shape[-2] + intermediate = gate_up_proj.shape[-2] // (2 if gate else 1) + gate_up_bits = int_weight_bits(hidden, gate_up_proj) + down_bits = int_weight_bits(intermediate, down_proj) + gate_up_groups, down_groups = gate_up_proj_scale.shape[-1], down_proj_scale.shape[-1] + # Per-channel recipes span each projection's own K; H and I need not match. + same_groups = ( + gate_up_groups == down_groups == 1 + or hidden // gate_up_groups == intermediate // down_groups + ) + if gate_up_bits != down_bits or not same_groups or (gate_up_zero is None) != (down_zero is None): + raise ValueError(error) + return gate_up_is_mx = is_mx(gate_up_proj, gate_up_proj_scale) - if gate_up_is_mx != is_mx(down_proj, down_proj_scale): + if gate_up_proj_scale is None or down_proj_scale is None: + matching = ( + gate_up_proj_scale is None + and down_proj_scale is None + and gate_up_proj.dtype == down_proj.dtype + ) + else: + matching = weight_format(gate_up_proj, gate_up_proj_scale) == weight_format(down_proj, down_proj_scale) + if matching and not gate_up_is_mx: + matching = ( + is_rowwise(gate_up_proj, gate_up_proj_scale) == is_rowwise(down_proj, down_proj_scale) + and weight_block_size(gate_up_proj, gate_up_proj_scale) == weight_block_size(down_proj, down_proj_scale) + and gate_up_proj_scale.element_size() == down_proj_scale.element_size() + ) + if not matching: raise ValueError( - "gate_up_proj and down_proj must use the same format (both MX or both block-dynamic FP8)." + "gate_up_proj and down_proj must use the same quantization recipe " + "(weight format and scale granularity)." ) - if is_mxfp4(gate_up_proj, gate_up_proj_scale) != is_mxfp4(down_proj, down_proj_scale): - raise ValueError("gate_up_proj and down_proj must use the same MX format (both MXFP4 or both MXFP8).") - return gate_up_is_mx + if gate_up_is_mx: + hidden = down_proj.shape[-2] + intermediate = gate_up_proj.shape[-2] // (2 if gate else 1) + if ( + mx_scale_family(gate_up_proj_scale, hidden) != mx_scale_family(down_proj_scale, intermediate) + ): + raise ValueError("gate_up_proj and down_proj must use the same MX format and scale group.") def _gather_idx(top_k_index: torch.Tensor) -> torch.Tensor: @@ -276,16 +312,21 @@ def _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate) -> di # ── Fused (gate_up epilogue owns SwiGLU + intermediate requant) ────────────── -def _block_format(gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale, activation_format): +def _block_format( + gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale, activation_format, + gate, gate_up_zero, down_zero, +): """The MoE block's activation format, resolved to a name: validates the weight pairing; an explicit ``activation_format`` (``"bf16"`` weight-only, or a format) is respected as-is, ``None`` follows the weights (fp8 / mxfp8 / mxfp4 / nvfp4 — mxfp4 weights default to mxfp4 activations, the all-fp4 W4A4 chain; unquantized BF16/FP16 weights carry no scales, so their - activations stay ``"bf16"``, the full-precision path).""" - _validate_moe(gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale) + activations stay ``"bf16"``, the full-precision path; nor do int4 weights, which run weight-only).""" + _validate_moe( + gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale, gate, gate_up_zero, down_zero, + ) if activation_format is not None: return activation_format - if gate_up_proj_scale is None: + if gate_up_proj_scale is None or gate_up_proj.dtype in INT_WEIGHT_DTYPES: return "bf16" return weight_format(gate_up_proj, gate_up_proj_scale) @@ -315,12 +356,13 @@ def _fused_post_expert_norm(down_out, post_expert_norm, weight, eps): return _post_expert_norm(down_out, post_expert_norm, weight, eps), {} -def _gate_up_activations(hidden, fmt, gate_up_scale, global_scale, activation_scale, scatter_idx, expert_start): +def _gate_up_activations(hidden, fmt, gate_up_scale, global_scale, activation_scale, scatter_idx, expert_start, pre_quant_scale=None): """``(A, As)`` for the grouped gate_up. Under a shared global on swizzled MX weights, ``hidden`` is quantized once with its scales stored straight into each routed row's expert tile, which the op reads as is while gathering the values; anywhere else the op quantizes ``hidden``.""" if ( - fmt not in MX_ACT_QUANT + pre_quant_scale is not None + or fmt not in MX_ACT_QUANT or gate_up_scale.ndim != 5 or activation_scale is not None or is_per_expert_global(global_scale) @@ -362,6 +404,13 @@ def moe_fused_grouped( simulate_unfused: bool = False, activation_format: str | None = None, gate: bool = True, + activation_scale_ub: float | None = None, + gate_up_proj_weight_zero_point: torch.Tensor | None = None, # integer weights' (E, 2I, H // group) + down_proj_weight_zero_point: torch.Tensor | None = None, # (E, H, I // group) + gate_up_proj_weight_perm: torch.Tensor | None = None, # act-order integer weights' (E, H) int32 K orders + down_proj_weight_perm: torch.Tensor | None = None, # (E, I); None once fold_act_order_into_gate_up folded it + gate_up_proj_pre_quant_scale: torch.Tensor | None = None, + down_proj_pre_quant_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Fused grouped MoE (prefill): gather gate_up + SiLU + requant epilogue → quantized expert-ordered intermediate → grouped down → routing-weighted top-k reduce. Returns @@ -373,14 +422,18 @@ def moe_fused_grouped( ``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's calibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own, requants the intermediate against the down's, and the down consumes that as its activation - global. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each step through - the activation dtype so the output matches the unfused reference to reduce order. ``act_fn`` is a + global. ``None`` is dynamic quant everywhere else. Tensor-wide / row-wise FP8 hand the + down a bf16 intermediate it quantizes per token (a whole-row scale no epilogue tile sees). + ``simulate_unfused`` (testing) rounds each step through the activation dtype so the output + matches the unfused reference to reduce order. ``act_fn`` is a ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D) - — the swizzled layout is for the dot_scaled arm; weight-only chains take affine.""" + — the swizzled layout is for the dot_scaled arm; weight-only chains take affine. + ``activation_scale_ub`` caps dynamic row-wise FP8 activations at both projections.""" fmt = _block_format( - gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format + gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format, + gate, gate_up_proj_weight_zero_point, down_proj_weight_zero_point, ) num_top_k = top_k_index.size(-1) glu = _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate) @@ -393,9 +446,12 @@ def moe_fused_grouped( # intermediate. scatter_idx=None: the down reads the intermediate in place. (C, Cs) under a # requant format; a bare Tensor otherwise. static_act = gate_up_proj_activation_scale is not None or down_proj_activation_scale is not None + # a per-token (tensor-wide / row-wise) intermediate scale spans the whole row, which the + # gate_up epilogue's tiles never see: the down quantizes the bf16 intermediate instead + per_token = per_token_activations(gate_up_proj, gate_up_proj_scale_inv) gate_up_in, gate_up_in_scale = _gate_up_activations( hidden_states, fmt, gate_up_proj_scale_inv, gate_up_proj_input_global_scale, - gate_up_proj_activation_scale, scatter_idx, expert_start, + gate_up_proj_activation_scale, scatter_idx, expert_start, gate_up_proj_pre_quant_scale, ) gate_up_out = matmul_grouped( gate_up_in, @@ -404,16 +460,20 @@ def moe_fused_grouped( Bs=gate_up_proj_scale_inv, a_global_scale=gate_up_proj_input_global_scale, b_global_scale=gate_up_proj_weight_global_scale, + b_zero=gate_up_proj_weight_zero_point, + b_perm=gate_up_proj_weight_perm, + a_pre_quant_scale=gate_up_proj_pre_quant_scale, # the intermediate requant normalizes against the DOWN's calibrated input global, # which the down then consumes as its activation global — the two-level handoff - output_global_scale=down_proj_input_global_scale, + output_global_scale=down_proj_input_global_scale if down_proj_pre_quant_scale is None else None, expert_start=expert_start, **glu, bias=gate_up_proj_bias, # fmt is the resolved format; "bf16" (weight-only) leaves the GLU intermediate bf16, # no requant. So does static, whose epilogue has no calibrated scale to requant against. activation_format=fmt, - quantize_output=bool(glu) and fmt != "bf16" and not static_act, + activation_scale_ub=activation_scale_ub, + quantize_output=bool(glu) and fmt != "bf16" and not static_act and not per_token and down_proj_pre_quant_scale is None, output_dtype=hidden_states.dtype, gather_idx=gather_idx, ) @@ -431,10 +491,14 @@ def moe_fused_grouped( Bs=down_proj_scale_inv, a_global_scale=down_proj_input_global_scale, b_global_scale=down_proj_weight_global_scale, + b_zero=down_proj_weight_zero_point, + b_perm=down_proj_weight_perm, + a_pre_quant_scale=down_proj_pre_quant_scale, expert_start=expert_start, bias=down_proj_bias, # weight-only: the intermediate is bf16 (As None) — the down goes weight-only too; a # quantized intermediate carries its format in its own dtypes. + activation_scale_ub=activation_scale_ub, activation_format=fmt if inter_scale is None else None, output_dtype=hidden_states.dtype, scatter_idx=scatter_idx, @@ -478,6 +542,13 @@ def moe_fused_batched( simulate_unfused: bool = False, activation_format: str | None = None, gate: bool = True, + activation_scale_ub: float | None = None, + gate_up_proj_weight_zero_point: torch.Tensor | None = None, # integer weights' (E, 2I, H // group) + down_proj_weight_zero_point: torch.Tensor | None = None, # (E, H, I // group) + gate_up_proj_weight_perm: torch.Tensor | None = None, # act-order integer weights' (E, H) int32 K orders + down_proj_weight_perm: torch.Tensor | None = None, # (E, I); None once fold_act_order_into_gate_up folded it + gate_up_proj_pre_quant_scale: torch.Tensor | None = None, + down_proj_pre_quant_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Fused batched MoE (decode): gate_up + SiLU + requant epilogue → per-row quantized intermediate → batched down → routing-weighted top-k reduce. Returns @@ -490,15 +561,18 @@ def moe_fused_batched( ``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's calibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own, requants the intermediate against the down's, and the down consumes that as its activation - global. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each - step through the activation dtype so the output matches the unfused reference to + global. ``None`` is dynamic quant everywhere else. Tensor-wide / row-wise FP8 hand the + down a bf16 intermediate it quantizes per token (a whole-row scale no epilogue tile sees). + ``simulate_unfused`` (testing) rounds each step through the activation dtype so the output matches the unfused reference to reduce order. ``act_fn`` is a ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D) - — the swizzled layout is for the dot_scaled arm; weight-only chains take affine.""" + — the swizzled layout is for the dot_scaled arm; weight-only chains take affine. + ``activation_scale_ub`` caps dynamic row-wise FP8 activations at both projections.""" fmt = _block_format( - gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format + gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format, + gate, gate_up_proj_weight_zero_point, down_proj_weight_zero_point, ) glu = _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate) NUM_EXPERTS = gate_up_proj.size(0) @@ -510,6 +584,8 @@ def moe_fused_batched( # row from the unexpanded hidden in-kernel (no copy). # (C, Cs) under a requant format; a bare Tensor on the full-precision path static_act = gate_up_proj_activation_scale is not None or down_proj_activation_scale is not None + # A per-token scale spans the whole intermediate row, so the down projection quantizes it. + per_token = per_token_activations(gate_up_proj, gate_up_proj_scale_inv) gate_up_out = matmul_batched( hidden_states, gate_up_proj, @@ -517,8 +593,11 @@ def moe_fused_batched( Bs=gate_up_proj_scale_inv, a_global_scale=gate_up_proj_input_global_scale, b_global_scale=gate_up_proj_weight_global_scale, + b_zero=gate_up_proj_weight_zero_point, + b_perm=gate_up_proj_weight_perm, + a_pre_quant_scale=gate_up_proj_pre_quant_scale, # the two-level handoff, as in the grouped sibling - output_global_scale=down_proj_input_global_scale, + output_global_scale=down_proj_input_global_scale if down_proj_pre_quant_scale is None else None, expert_ids=expert_ids, **glu, bias=gate_up_proj_bias, @@ -530,11 +609,14 @@ def moe_fused_batched( # (down inline-quants) stays the win there. Static keeps bf16 either way, its epilogue # having no calibrated scale to requant against. activation_format=fmt, + activation_scale_ub=activation_scale_ub, quantize_output=( bool(glu) and fmt != "bf16" and (fmt != "fp8" or expert_ids.numel() <= GATE_UNSTACK_MAX_S) and not static_act + and not per_token + and down_proj_pre_quant_scale is None ), output_dtype=hidden_states.dtype, gather_idx=gather_idx, @@ -553,11 +635,15 @@ def moe_fused_batched( Bs=down_proj_scale_inv, a_global_scale=down_proj_input_global_scale, b_global_scale=down_proj_weight_global_scale, + b_zero=down_proj_weight_zero_point, + b_perm=down_proj_weight_perm, + a_pre_quant_scale=down_proj_pre_quant_scale, expert_ids=expert_ids, bias=down_proj_bias, # weight-only / block-FP8: the intermediate is bf16 (As is None), so the down carries the # format and quantizes it, mirroring the unfused sibling; a quantized MX intermediate # carries its format in its own dtypes. + activation_scale_ub=activation_scale_ub, activation_format=fmt if inter_scale is None or fmt == "fp8" else None, output_dtype=hidden_states.dtype, ) @@ -601,22 +687,31 @@ def moe_unfused_grouped( swiglu_limit: float | None = None, activation_format: str | None = None, gate: bool = True, + activation_scale_ub: float | None = None, + gate_up_proj_weight_zero_point: torch.Tensor | None = None, # integer weights' (E, 2I, H // group) + down_proj_weight_zero_point: torch.Tensor | None = None, # (E, H, I // group) + gate_up_proj_weight_perm: torch.Tensor | None = None, # act-order integer weights' (E, H) int32 K orders + down_proj_weight_perm: torch.Tensor | None = None, # (E, I); None once fold_act_order_into_gate_up folded it + gate_up_proj_pre_quant_scale: torch.Tensor | None = None, + down_proj_pre_quant_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Unfused grouped MoE: gate_up (plain grouped GEMM, gather hidden) → host ``apply_glu`` → down (plain grouped GEMM, scatter to routed rows) → routing-weighted reduce. Same math as ``moe_fused_grouped`` but the SwiGLU + intermediate quant happen between two plain GEMMs rather than inside the gate_up epilogue; each GEMM quantizes its raw input in ``activation_format`` (``None`` follows the weight format, mirroring the fused forward — mxfp4 - weights run the all-fp4 W4A4 chain). All formats route through the shared ``matmul_grouped``. + weights run the all-fp4 W4A4 chain). The NVFP4 activation globals thread the same way as the fused sibling: each GEMM quantizes its raw input against its ``*_input_global_scale``, and a ``*_activation_scale`` quantizes it against that calibrated scale instead of a runtime one. ``act_fn`` is a ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D) - — the swizzled layout is for the dot_scaled arm; weight-only chains take affine.""" + — the swizzled layout is for the dot_scaled arm; weight-only chains take affine. + ``activation_scale_ub`` caps dynamic row-wise FP8 activations at both projections.""" fmt = _block_format( - gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format + gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format, + gate, gate_up_proj_weight_zero_point, down_proj_weight_zero_point, ) num_top_k = top_k_index.size(-1) @@ -628,21 +723,36 @@ def moe_unfused_grouped( # gate_up as a plain GEMM (no gate epilogue) over gathered hidden -> expert-ordered (S, 2I). gate_up_in, gate_up_in_scale = _gate_up_activations( hidden_states, fmt, gate_up_proj_scale_inv, gate_up_proj_input_global_scale, - gate_up_proj_activation_scale, scatter_idx, expert_start, - ) - gate_up_out = matmul_grouped( - gate_up_in, - gate_up_proj, - As=gate_up_in_scale, - Bs=gate_up_proj_scale_inv, - a_global_scale=gate_up_proj_input_global_scale, - b_global_scale=gate_up_proj_weight_global_scale, - expert_start=expert_start, - bias=gate_up_proj_bias, - activation_format=fmt, - output_dtype=hidden_states.dtype, - gather_idx=gather_idx, + gate_up_proj_activation_scale, scatter_idx, expert_start, gate_up_proj_pre_quant_scale, ) + if (gate and fmt == "bf16" and gate_up_proj.dtype == torch.float8_e4m3fn + and not is_mx(gate_up_proj, gate_up_proj_scale_inv) and gate_up_proj_pre_quant_scale is None): + assert (gate_up_proj_activation_scale is None and gate_up_proj_input_global_scale is None + and gate_up_proj_weight_global_scale is None), "FP8 weight-only takes raw activations and block scales" + # Keep calibrated gate/up scale blocks interleaved while leaving the GLU unfused. + gate_up_out = block_fp8_weight_only_matmul_grouped( + gate_up_in, gate_up_proj, gate_up_proj_scale_inv, expert_start, + output_dtype=hidden_states.dtype, gather_idx=gather_idx, + bias=gate_up_proj_bias, interleaved_scales=True, + )[0] + else: + gate_up_out = matmul_grouped( + gate_up_in, + gate_up_proj, + As=gate_up_in_scale, + Bs=gate_up_proj_scale_inv, + a_global_scale=gate_up_proj_input_global_scale, + b_global_scale=gate_up_proj_weight_global_scale, + b_zero=gate_up_proj_weight_zero_point, + b_perm=gate_up_proj_weight_perm, + a_pre_quant_scale=gate_up_proj_pre_quant_scale, + expert_start=expert_start, + bias=gate_up_proj_bias, + activation_format=fmt, + activation_scale_ub=activation_scale_ub, + output_dtype=hidden_states.dtype, + gather_idx=gather_idx, + ) inter = _host_glu(gate_up_out, act_fn, swiglu_alpha, swiglu_limit, gate) # down over the expert-ordered intermediate (quantized in the same format), scattering # to routed rows. @@ -653,9 +763,13 @@ def moe_unfused_grouped( Bs=down_proj_scale_inv, a_global_scale=down_proj_input_global_scale, b_global_scale=down_proj_weight_global_scale, + b_zero=down_proj_weight_zero_point, + b_perm=down_proj_weight_perm, + a_pre_quant_scale=down_proj_pre_quant_scale, expert_start=expert_start, bias=down_proj_bias, activation_format=fmt, + activation_scale_ub=activation_scale_ub, output_dtype=hidden_states.dtype, scatter_idx=scatter_idx, ) @@ -878,40 +992,64 @@ def moe_unfused_batched( swiglu_limit: float | None = None, activation_format: str | None = None, gate: bool = True, + activation_scale_ub: float | None = None, + gate_up_proj_weight_zero_point: torch.Tensor | None = None, # integer weights' (E, 2I, H // group) + down_proj_weight_zero_point: torch.Tensor | None = None, # (E, H, I // group) + gate_up_proj_weight_perm: torch.Tensor | None = None, # act-order integer weights' (E, H) int32 K orders + down_proj_weight_perm: torch.Tensor | None = None, # (E, I); None once fold_act_order_into_gate_up folded it + gate_up_proj_pre_quant_scale: torch.Tensor | None = None, + down_proj_pre_quant_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Unfused batched MoE: gate_up (plain batched GEMM, gather hidden) → host ``apply_glu`` → down (plain batched GEMM) → routing-weighted reduce. Same math as ``moe_fused_batched`` but the SwiGLU + intermediate quant happen between two plain GEMMs; each GEMM quantizes its raw - input in ``activation_format`` (``None`` follows the weight format, ``"bf16"`` is weight-only). All - formats route through the shared ``matmul_batched``. The NVFP4 activation globals thread the - same way as the fused sibling: each GEMM quantizes its raw input against its + input in ``activation_format`` (``None`` follows the weight format, ``"bf16"`` is weight-only). + The NVFP4 activation globals thread the same way as the fused sibling: each GEMM + quantizes its raw input against its ``*_input_global_scale``, and a ``*_activation_scale`` quantizes it against that calibrated scale instead of a runtime one. ``act_fn`` is a ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D) - — the swizzled layout is for the dot_scaled arm; weight-only chains take affine.""" + — the swizzled layout is for the dot_scaled arm; weight-only chains take affine. + ``activation_scale_ub`` caps dynamic row-wise FP8 activations at both projections.""" fmt = _block_format( - gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format + gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format, + gate, gate_up_proj_weight_zero_point, down_proj_weight_zero_point, ) NUM_EXPERTS = gate_up_proj.size(0) expert_ids = top_k_index.reshape(-1) gather_idx = _gather_idx(top_k_index) # gate_up as a plain GEMM (no gate epilogue) over gathered hidden -> (S, 2I). - gate_up_out = matmul_batched( - hidden_states, - gate_up_proj, - As=gate_up_proj_activation_scale, - Bs=gate_up_proj_scale_inv, - a_global_scale=gate_up_proj_input_global_scale, - b_global_scale=gate_up_proj_weight_global_scale, - expert_ids=expert_ids, - bias=gate_up_proj_bias, - activation_format=fmt, - output_dtype=hidden_states.dtype, - gather_idx=gather_idx, - ) + if (gate and fmt == "bf16" and gate_up_proj.dtype == torch.float8_e4m3fn + and not is_mx(gate_up_proj, gate_up_proj_scale_inv) and gate_up_proj_pre_quant_scale is None): + assert (gate_up_proj_activation_scale is None and gate_up_proj_input_global_scale is None + and gate_up_proj_weight_global_scale is None), "FP8 weight-only takes raw activations and block scales" + # Keep calibrated gate/up scale blocks interleaved while leaving the GLU unfused. + gate_up_out = block_fp8_weight_only_matmul_batched( + hidden_states, gate_up_proj, gate_up_proj_scale_inv, expert_ids, + output_dtype=hidden_states.dtype, gather_idx=gather_idx, + bias=gate_up_proj_bias, interleaved_scales=True, + )[0] + else: + gate_up_out = matmul_batched( + hidden_states, + gate_up_proj, + As=gate_up_proj_activation_scale, + Bs=gate_up_proj_scale_inv, + a_global_scale=gate_up_proj_input_global_scale, + b_global_scale=gate_up_proj_weight_global_scale, + b_zero=gate_up_proj_weight_zero_point, + b_perm=gate_up_proj_weight_perm, + a_pre_quant_scale=gate_up_proj_pre_quant_scale, + expert_ids=expert_ids, + bias=gate_up_proj_bias, + activation_format=fmt, + activation_scale_ub=activation_scale_ub, + output_dtype=hidden_states.dtype, + gather_idx=gather_idx, + ) inter = _host_glu(gate_up_out, act_fn, swiglu_alpha, swiglu_limit, gate) # down over the intermediate (quantized in the same format), routed-order output. down_out = matmul_batched( @@ -921,9 +1059,13 @@ def moe_unfused_batched( Bs=down_proj_scale_inv, a_global_scale=down_proj_input_global_scale, b_global_scale=down_proj_weight_global_scale, + b_zero=down_proj_weight_zero_point, + b_perm=down_proj_weight_perm, + a_pre_quant_scale=down_proj_pre_quant_scale, expert_ids=expert_ids, bias=down_proj_bias, activation_format=fmt, + activation_scale_ub=activation_scale_ub, output_dtype=hidden_states.dtype, ) return _torch_weighted_reduce( diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/pruners.py b/finegrained-kernels/torch-ext/finegrained_kernels/pruners.py index cd102585a..03957c96c 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/pruners.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/pruners.py @@ -31,9 +31,9 @@ # scalar "scalar" -> BM=1 the scalar GEVM broadcasts wrong at BM>1 (config builder) # descriptor_needs_prequant_pruner raw-A inline-quant arm reads a_ptrs, dead under a # descriptor mode — static 2D -# weight_only_swap_scope_pruner SWAP_AB iff COMPUTE_MODE == "scalar" — the swapped dot -# arms are unimplemented and a non-swapped scalar has no -# accumulator shape — weight-only batched +# weight_only_swap_scope_pruner swapped dot is implemented for batched integer FP8; +# the other batched weight-only kernels swap only for scalar +# decode, which requires the weight-rows-major layout # mma_trap_ok swapped dot_scaled BK<128 (sticky 716), single-trip # accumulator-init miscompile, width > 256 (sticky trap), # packed-fp4 GATE BN<64 silent zeros — mx 2D/batched/grouped; @@ -55,12 +55,10 @@ # weight-only grouped # affine_scale_warp_spec_pruner affine scale gather's other=0.0 can't take # ttg.partition under WS — mx-dynamic grouped -# global_scale_warp_spec_pruner the NVFP4 global load can't take ttg.partition -# under WS — weight-only grouped, BsGlobal-gated # nvfp4_native_ok (mx_config_pruner) E4M3-scale dot_scaled PassManager-fails off the # native staging: M operand < 128, num_warps outside {4, 8}, # or a weight-only kernel (no As — never native) -# descriptor_box_pruner TMA box > 256/dim is illegal — all descriptor-axis kernels +# descriptor_box_pruner Descriptor strides and gather architecture — descriptor-axis kernels # smem_pruner IR allocation bounds (B200-verified, keep-side error) — # bd/static/mx/weight-only 2D; tensor/mx/weight-only/fp grouped; # mx batched @@ -111,7 +109,7 @@ # (WARP_SPEC, PACKED_SCHEDULE, memory modes) share the tile basin shape and stay out — # per-basin anchors over the full branch cross starved a 100-trial budget (see # _basin_anchor_indices). -PATH_ANCHOR_AXES = ("COMPUTE_MODE", "SWAP_AB") +PATH_ANCHOR_AXES = ("COMPUTE_MODE", "SWAP_AB", "PIPELINE") def config_dim(c, all_args, name): @@ -186,6 +184,29 @@ def raise_no_dividing_block(configs, args): return config_filter(ok, when=when, on_empty=raise_no_dividing_block) +def int_scale_group_tile_pruner(divides_k: bool = True): + """``early_config_prune`` for the integer weight-only kernels' K tile: whole scale groups (each group's scale + broadcasts over its columns of the dequantized tile) or a whole number of tiles per group (a group wider than the + tile, per-channel included), and a tile that divides K wherever the K-loop has no tail (``divides_k``: the 2D and + grouped kernels) or a group spans tiles (a slid-back tail tile would straddle two groups). The batched loop slides + its last tile back in bounds otherwise. Newly admitted within-group tiles start at 128, preserving + the existing group-128 grid; narrower tiles add K-loop and metadata traffic.""" + + def ok(c, args): + # The packed scalar decode helpers support bf16/fp16 and INT8, not FP8 operands. + if args["A"].dtype == torch.float8_e4m3fn and c.kwargs.get("COMPUTE_MODE") == "scalar": + return False + # INT32 partials must stay integral through the final K sum; the split-K workspace is FP32. + if args["A"].dtype == torch.int8 and c.kwargs.get("SPLIT_K", 1) != 1: + return False + bk, group = c.kwargs["BLOCK_SIZE_K"], args["SCALE_GROUP_K"] + if group > bk: + return bk >= 128 and group % bk == 0 and args["K"] % bk == 0 + return bk % group == 0 and (not divides_k or args["K"] % bk == 0) + + return config_filter(ok) + + def block_fits_dim_pruner(dim_arg: str, block_key: str = "BLOCK_SIZE_K"): """``early_config_prune`` for a loop whose tail tile SLIDES back in bounds instead of running past the row (the weight-only batched K-loop): the tile need not divide the dim, @@ -232,6 +253,9 @@ def warp_spec_compile_guard_pruner(): def ok(c, args): if not c.kwargs.get("WARP_SPEC"): return True + # The tile-loop WS candidates require Triton 3.8's partitioning passes. + if c.kwargs.get("PIPELINE", "k_loop") != "k_loop" and tuple(int(n) for n in triton.__version__.split(".")[:2]) < (3, 8): + return False # WS needs num_warps a multiple of 4 (its async producer/consumer partitions) and enough M # work — the MMA's M operand, which SWAP_AB moves to BLOCK_SIZE_N (so BM=1 decode can still # WS when swapped). dot_scaled + WS + TMA DOES compile here (measured 1995 on a clean @@ -510,6 +534,9 @@ def nvfp4_native_ok(c, args): # token in N. if c.kwargs.get("COMPUTE_MODE") != "dot_scaled": return True + if args.get("SCALE_GROUP_K") == 32: + # Mixed tensor-FP8 / group-32 E4M3 FP4 is not the group-16 FP4xFP4 MMA. + return False if not acts_are_scaled(args) or not 4 <= c.num_warps <= 8: return False if c.kwargs.get("SWAP_AB"): @@ -558,6 +585,7 @@ def scales_are_e4m3(args): dot_arm_ok, when=lambda args: is_sm10x() and scaled_mma_available(args) + and not (scales_are_e4m3(args) and args.get("SCALE_GROUP_K") == 32) and acts_are_scaled(args), ), config_filter(mma_poison_ok, when=lambda args: is_sm10x()), @@ -785,7 +813,7 @@ def _warp_spec_needs_descriptor_tile() -> bool: return tuple(int(n) for n in triton.__version__.split(".")[:2]) >= (3, 8) -def warp_spec_memory_mode_pruner(weight_only: bool = False): +def warp_spec_memory_mode_pruner(weight_only: bool = False, gather_only: bool = False): """``early_config_prune`` dropping ``warp_specialize`` configs the WS passes cannot lower (``TritonGPUOptimizePartitionWarps`` / ``RelayoutTritonGPU`` -> ``PassManager::run failed``). @@ -805,6 +833,17 @@ def warp_spec_memory_mode_pruner(weight_only: bool = False): REGRESSED in 3.8 and pointer+WS was FIXED. So each version gets the rule its own sweep measured, rather than an intersection that would fence off a whole memory mode on both. + ``gather_only`` applies this restriction only to gathered descriptor A. The raw-activation + grouped kernel reproduces the TMA gather layout failure with both BF16 and rescaled FP8 + weights, while non-gathered descriptor boxes compile at the other warp counts. Pointer A + and non-gathered launches retain their existing search space. + + Weight-only dot prefill additionally supports the measured w8 128x128x128 tile at + stages 2/3 with both host descriptors, inline scheduling and an already expert-sorted A. Both original + and swapped operand orientations lower; descriptor gather is outside this exception. + Integer outer pipelines also support eight warps for the 128x64x256 gate tile and + the packed, flattened 128x128x256 down tile at stage 2. + ``weight_only`` adds the rows that kernel charted for itself and the sweep did not re-probe: ``BN == 32`` with WS, which failed 9/9 cells there, and (below 3.8) its trapping pointer/pointer warp counts. @@ -815,6 +854,10 @@ def warp_spec_memory_mode_pruner(weight_only: bool = False): width law and belongs wherever gather width is decided.""" def ok(c, args): + if gather_only and ( + args.get("GatherIdx") is None or c.kwargs.get("A_MEMORY_MODE", "pointer") == "pointer" + ): + return True if not c.kwargs.get("WARP_SPEC"): return True if weight_only and config_dim(c, args, "BLOCK_SIZE_N") == 32: @@ -822,8 +865,32 @@ def ok(c, args): a, b = c.kwargs.get("A_MEMORY_MODE"), c.kwargs.get("B_MEMORY_MODE") descriptor = a != "pointer" or b != "pointer" if _warp_spec_needs_descriptor_tile(): - # 3.8: one surviving descriptor cell, and pointer/pointer is wide open - return not descriptor or (c.num_warps == 4 and config_dim(c, args, "BLOCK_SIZE_M") == 128) + prefill_tile = ( + weight_only and c.kwargs.get("COMPUTE_MODE") == "dot" + and a == b == "host_descriptor" and args.get("GatherIdx") is None + and not c.kwargs.get("PACKED_SCHEDULE") + and config_dim(c, args, "BLOCK_SIZE_N") == 128 + and config_dim(c, args, "BLOCK_SIZE_K") == 128 + and c.num_stages in (2, 3) + ) + pipeline_tile = ( + weight_only and c.kwargs.get("PIPELINE") in ("tile_loop", "flattened") + and c.kwargs.get("SWAP_AB") + and a == b == "host_descriptor" and args.get("GatherIdx") is None + and config_dim(c, args, "BLOCK_SIZE_K") == 256 and c.num_stages == 2 + and ( + args.get("GATE", False) and config_dim(c, args, "BLOCK_SIZE_N") == 64 + or ( + args.get("INT_BITS") in (2, 4, 8) and not args.get("GATE", False) + and config_dim(c, args, "BLOCK_SIZE_N") == 128 + and c.kwargs.get("PACKED_SCHEDULE") and c.kwargs.get("PIPELINE") == "flattened" + ) + ) + ) + return not descriptor or ( + config_dim(c, args, "BLOCK_SIZE_M") == 128 + and (c.num_warps == 4 or ((prefill_tile or pipeline_tile) and c.num_warps == 8)) + ) if descriptor: return True traps = _WEIGHT_ONLY_POINTER_TRAPS if weight_only else {} @@ -832,16 +899,29 @@ def ok(c, args): return config_filter(ok, when=lambda args: is_sm10x()) +def int_2d_warp_spec_pruner(): + """``early_config_prune`` for the integer 2D kernel: ``warp_specialize`` with a host-TMA activation and a pointer + weight (A descriptor, B pointer) fails ``TritonGPUOptimizePartitionWarps`` (``PassManager::run failed``). + Charted on B200 / Triton 3.8 over the kernel's own WS neighbourhood (4- and 8-bit, warps {4, 8, 16}, BM {64, 128}, + BK {128, 256, 512}, BN {32, 64, 128}, stages {2, 4}, all four operand-mode pairs): every lowering failure is in + that pair, at warps 4 and 8 (68 of the 100 cells there that fit shared memory; warps 16 lowers), and none in the + other pairs. The pair is fenced whole: its cells that do lower never came within 1.6x of a crown in any recorded + tune (decode 5-10x off), so the exact boundary (tile bytes, width, warps) buys nothing. CUDA-only.""" + + def ok(c, args): + mixed = c.kwargs.get("A_MEMORY_MODE", "pointer") != "pointer" and c.kwargs.get("B_MEMORY_MODE", "pointer") == "pointer" + return not (c.kwargs.get("WARP_SPEC") and mixed) + + return config_filter(ok, when=lambda args: get_active_device_type() == "cuda") + + def weight_only_swap_scope_pruner(): - """``early_config_prune`` binding the weight-only kernel's ``SWAP_AB`` axis to its one - implemented arm: the CUDA-core reduce. Swap exists here to drop the MMA entirely — with - raw bf16 activations there is no native fp4 x bf16 MMA, so both dot arms pay an M->16 pad - and a full-tile upcast, and a swapped ``dot_scaled`` measured no better than the row-major - one (35.2 vs 38.0us, gpt-oss gate_up decode). So ``SWAP_AB`` iff ``COMPUTE_MODE`` is - ``"scalar"``: the swapped-dot rows would compile into an arm the compute helper does not - implement, and a non-swapped scalar row has no accumulator shape.""" + """Batched FP8 integer-weight dots support the swapped MMA; the other batched weight-only + kernels swap only for scalar decode. The scalar reduce requires the swapped row layout.""" def ok(c, args): + if args["A"].dtype == FP8_DTYPE and c.kwargs.get("COMPUTE_MODE") == "dot" and c.kwargs.get("SWAP_AB"): + return c.kwargs["BLOCK_SIZE_M"] <= MMA_N_ATOM_WIDTH return bool(c.kwargs.get("SWAP_AB")) == (c.kwargs.get("COMPUTE_MODE") == "scalar") return config_filter(ok) @@ -878,6 +958,41 @@ def ok(c, args): return config_filter(ok) +def native_mma_m_pruner(): + """Keep native tcgen05 M tiles for plain FP8 on sm_10x. Narrow unswapped + tiles lower to warp MMA and lose warm-cache performance; swapped decode tiles + use the weight-row dimension as MMA M and remain eligible.""" + + def ok(c, args): + return bool(c.kwargs.get("SWAP_AB")) or config_dim(c, args, "BLOCK_SIZE_M") >= 64 + + return config_filter(ok, when=lambda args: is_sm10x()) + + +def single_row_tile_pruner(): + """Avoid replicating a single token across a BM=16 warp-MMA tile on sm_10x. + BM=1 retains the same compute path with less replicated work.""" + + def ok(c, args): + return config_dim(c, args, "BLOCK_SIZE_M") == 1 + + return config_filter(ok, when=lambda args: is_sm10x()) + + +def swap_decode_scope_pruner(max_m: int = 16): + """``early_config_prune`` scoping a 2D plain-dot kernel's ``SWAP_AB`` rows (BM=1: the weight's + output rows in the MMA M dim, the token padded to the N=16 atom) to decode M, where they run + one program row per token, and to a pointer-mode activation (an A-side descriptor under + ``SWAP_AB`` traps on Triton 3.8, see ``mx_2d_swap_scope_pruner``). Unswapped rows pass.""" + + def ok(c, args): + if not c.kwargs.get("SWAP_AB"): + return True + return args["M"] <= max_m and c.kwargs.get("A_MEMORY_MODE", "pointer") == "pointer" + + return config_filter(ok) + + def mx_2d_swap_scope_pruner(max_m: int = 16): """``early_config_prune`` scoping the 2D mx kernel's ``SWAP_AB`` rows to the ONE regime they exist for: E4M3-scale (NVFP4) single-token decode. NVFP4 has no other native M=1 @@ -1097,6 +1212,79 @@ def ok(c, args): return config_filter(ok, when=lambda args: args[m_arg] > max_m) + +def swap_max_m_pruner(m_arg: str, max_m: int): + """Keep measured small-row swaps in one M tile, retaining BM1/2 candidates for M1.""" + + def ok(c, args): + if not c.kwargs.get("SWAP_AB"): + return True + m = args[m_arg] + bm = c.kwargs["BLOCK_SIZE_M"] + return m <= max_m and m <= bm <= triton.next_power_of_2(max(2, m)) + + return config_filter(ok) + + +def block_m_cap_pruner(m_arg: str, min_bm: int = 16): + """``early_config_prune`` dropping configs whose ``BLOCK_SIZE_M`` exceeds the launch's row count rounded up to a + power of two (floor ``min_bm``, the MMA atom): a dense GEMM with M below the tile has one M tile either way, so + the extra rows only pad the MMA and shrink the grid, and benching them dilutes a decode tune (int4 2D, B200: + the 4096x12288 M=16 tune crowned 28.2us against the 25.2us its BM=16 basin holds). Non-empty: ``min_bm`` stays.""" + + def ok(c, args): + bm = c.kwargs.get("BLOCK_SIZE_M") + return bm is None or bm <= max(min_bm, triton.next_power_of_2(args[m_arg])) + + return config_filter(ok) + + +def scalar_block_m_pruner(m_arg: str): + """``early_config_prune`` keeping only the ``scalar`` configs whose ``BLOCK_SIZE_M`` is the launch's row count + rounded up to a power of two: the reduce dequantizes each weight word once for its BM rows, so a smaller BM + re-dequantizes per program and a larger one reduces padding rows (int4 2D, B200, 4096x4096: M=2 BM 1/2/4 = + 9.0/7.7/10.2us, M=4 14.2/11.2/9.9us), and the TPE missed the matched tile at the default budget. Non-empty: + the dot configs stay.""" + + def ok(c, args): + return c.kwargs.get("COMPUTE_MODE") != "scalar" or c.kwargs["BLOCK_SIZE_M"] == triton.next_power_of_2(args[m_arg]) + + return config_filter(ok) + + +def mx_auxiliary_pruner(): + """SVDQuant native MMA: split the decode reduction or persist across prefill tiles.""" + def ok(c, args): + split = c.kwargs.get("SPLIT_K", 1) + persistent = c.kwargs.get("PERSISTENT", False) + if split == 1 and not persistent: + return True + if not is_sm10x() or args["GATE"] or args.get("LORA_RANK", 0) == 0: + return False + if persistent: + return args["M"] >= 128 + return args["M"] == 1 and args["K"] // c.kwargs["BLOCK_SIZE_K"] >= 2 * split + return config_filter(ok) + + +def int_split_k_pruner(m_arg: str, max_m: int): + """``early_config_prune`` for a split-K axis: splits only in the decode band (``m_arg`` <= ``max_m``, where the + launch allocates the partials) and only where the split divides the K tiles. Sub-16 dot tiles cover the + launch's rows in one tile, keeping the counter grid within its allocation and avoiding repeated weight reads.""" + + def ok(c, args): + bm = c.kwargs.get("BLOCK_SIZE_M", 16) + if bm < 16 and c.kwargs.get("COMPUTE_MODE", "dot") != "scalar" and bm != triton.next_power_of_2(args[m_arg]): + return False + split = c.kwargs.get("SPLIT_K", 1) + if split == 1: + return True + k_tiles = args["K"] // c.kwargs["BLOCK_SIZE_K"] + return args[m_arg] <= max_m and k_tiles % split == 0 and k_tiles >= split + + return config_filter(ok) + + def matched_memory_modes_pruner(): """Drop configs that mix pointer and descriptor operand loads — for a DENSE 2D GEMM only (contiguous A, whole M in one tile). Measured across M=1..8192 on the block-scale @@ -1134,57 +1322,35 @@ def ok(c, args): def descriptor_box_pruner(k_dim="BLOCK_SIZE_K"): - """Drop descriptor configs whose TMA box exceeds the tensormap's 256-per-dim limit: - the weight inner dim is the K tile in BYTES over values-per-byte (packed E2M1 - halves it; E4M3 is one byte per value), so BK=512 e4m3 boxes are illegal while - BK=512 packed ones fit. Pointer configs pass untouched. ``k_dim`` names the K-tile - axis — ``BLOCK_SIZE_K`` for the tune_block_nk kernels, ``block_k`` for the block-scale - kernels whose K tile is the caller's fixed scale block.""" + """Fence descriptor strides and gather support. Triton splits logical tiles into + legal physical TMA boxes; the logical tile may exceed 256 elements per dimension. + Pointer configs pass untouched.""" def ok(c, args): - # tiles may be tuned (config kwargs) or pinned launch args — config_dim reads both - bk = config_dim(c, args, k_dim) if c.kwargs.get("B_MEMORY_MODE", "pointer") != "pointer": - weight_vpb = ( - 2 if getattr(args.get("B"), "dtype", None) == torch.uint8 else 1 - ) - if bk // weight_vpb > 256 or config_dim(c, args, "BLOCK_SIZE_N") > 256: - return False # A TMA descriptor needs a unit-stride innermost dim; a transposed operand # (column-major view, as a backward pass hands in) has none, so only the # pointer arm can serve it. The kernel itself is stride-general. b = args.get("B") if b is not None and getattr(b, "ndim", 0) == 2 and b.stride(-1) != 1: return False + if "BDescriptor" in args and args["BDescriptor"] is None: + return False if c.kwargs.get("A_MEMORY_MODE", "pointer") != "pointer": + if "ADescriptor" in args and args["ADescriptor"] is None: + return False # descriptor A, both arms tuner-routed per the tokens-per-expert key: # contiguous boxes (no gather) win at dense routing (-7.1/-7.4%) and lose at # skew; tma gather4 (gathered rows, 1-row box, sm_100 only) is MONOTONIC in # tokens/expert — 2x loss at 32/expert, parity at 1024, -5.3/-4.7% at 4096 # (32k tokens, 2026-07-16). Only the arch is fenced. - act_vpb = 2 if getattr(args.get("A"), "dtype", None) == torch.uint8 else 1 if args.get("GatherIdx") is not None and not is_sm10x(): return False - if bk // act_vpb > 256: - return False return True return config_filter(ok) -def global_scale_warp_spec_pruner(): - """Drop WARP_SPEC configs when the launch carries an NVFP4 global (``BsGlobal`` not None): - the per-expert global load cannot take ``ttg.partition`` under warp specialization — the - same Triton lowering gap as the affine scale gather's ``other=0.0`` — so every such config - is a benign compile-failure inf. Fenced to save the compiles; single-level launches - (``BsGlobal`` None) keep their WS rows.""" - - def ok(c, args): - return not (c.kwargs.get("WARP_SPEC") and args.get("BsGlobal") is not None) - - return config_filter(ok) - - def gate_stacked_tmem_trap_pruner(): """Miscompile gate: drop GATE configs at exactly ``BLOCK_SIZE_N == 256, BLOCK_SIZE_M == 64`` — the stacked gate|up accumulator (``2*BN = 512`` columns) exceeds sm_100 tensor memory, and @@ -1283,9 +1449,9 @@ def ok(c, args): + (28 if c.kwargs.get("WARP_SPEC") else 0) + (32 if c.kwargs.get("B_MEMORY_MODE", "pointer") != "pointer" else 0) ) - need = fixed + stages * bk * ( - a.element_size() * bm + b.element_size() * tiles * bn - ) + # an integer weight is uint32 words of 32 // bits K values: bits / 8 bytes per K per row + b_bytes_per_k = 4 * b.shape[-1] / args["K"] if b.dtype == torch.uint32 else b.element_size() + need = fixed + stages * bk * (a.element_size() * bm + b_bytes_per_k * tiles * bn) elif c.kwargs.get("COMPUTE_MODE") == "dot_scaled": # KEEP-SIDE bound = HALF of one buffer set: operands stored PACKED (fp4 = 2 # values/byte) plus both group scale tiles (BK/group cols; group 32 UE8M0, 16 diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/quant.py b/finegrained-kernels/torch-ext/finegrained_kernels/quant.py index 363e6c80e..f4c133c29 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/quant.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/quant.py @@ -18,7 +18,7 @@ from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait from .bayesian_autotuner import bayesian_autotune -from .formats import global_scale_stride, is_per_expert_global +from .formats import global_scale_stride, int_weight_bits, is_per_expert_global, normalize_global_scale from .compat import FP8_DTYPE, MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, compile_time_only_triton_wrap, device_context, is_sm10x, pdl_launch_kwargs from .swizzle import swizzle_store_block from .scheduling import build_tile_layout, resolve_tile_inline @@ -26,6 +26,20 @@ # ── Triton-side helpers (inlined by ``@triton.jit`` callers) ────────────────── +@triton.jit +def _quantize_modelopt_input(vals, scale, OUT_DTYPE: tl.constexpr): + value = vals.to(tl.float32) / scale + if OUT_DTYPE == tl.int8: + value = tl.clamp(value, -128.0, 127.0) + # Round to nearest even, including negative half-way values, on every backend. + lower = tl.floor(value) + fraction = value - lower + value = lower + ((fraction > 0.5) | ((fraction == 0.5) & ((lower.to(tl.int32) & 1) != 0))).to(tl.float32) + else: + value = tl.clamp(value, -448.0, 448.0) + return value.to(tl.int8 if OUT_DTYPE == tl.int8 else tl.float8e4nv) + + @triton.jit def fp8_act_quant_inline( a_raw, TRANSPOSED: tl.constexpr = False, UE8M0: tl.constexpr = False @@ -429,6 +443,7 @@ def store_mx_act_scales( @triton.jit def _mx_act_quant_kernel( X, + PreQuantScale, Y, S, # (T, K // SCALE_GROUP_K) row-major scales; None on the swizzled path SOut, # flat SWIZZLE_32_4_4 scale buffer (1, n_tiles, cb, 2, 256); None on the plain path @@ -510,6 +525,9 @@ def _mx_act_quant_kernel( mask=row_mask[:, None], other=0.0, ).to(tl.float32) + if PreQuantScale is not None: + # SVDQuant rounds the smoothed input to BF16 before either branch. + x = (x * tl.load(PreQuantScale + offs)[None, :]).to(X.dtype.element_ty).to(tl.float32) if GlobalScale is not None: # NVFP4 two-level: normalize by the calibrated global before the block quant — block # scales are then amax/6 of x/g (the canonical two-step); the GEMM folds g back onto the @@ -586,6 +604,7 @@ def mx_act_quant_grouped( lambda META: (n_m_tiles, K // META["BLOCK_K"]) ]( x, + None, # PreQuantScale y, None if swizzled else scales, # S: row-major scales, affine arm only scales if swizzled else None, # SOut: SWIZZLE_32_4_4 buffer, swizzled arm only @@ -683,6 +702,7 @@ def mx_act_quant_routed_rows( lambda META: (triton.cdiv(S, META["BLOCK_T"]), K // META["BLOCK_K"]) ]( x, + None, # PreQuantScale y, None if store_swizzled else scales, scales if store_swizzled else None, @@ -742,6 +762,7 @@ def mx_act_quant_routed_scales( lambda META: (triton.cdiv(T, META["BLOCK_T"]), K // META["BLOCK_K"]) ]( x, + None, # PreQuantScale y, None, scales, @@ -905,6 +926,118 @@ def nvfp4_act_quant( ) +def _svdquant_pruner(configs, named_args, **kwargs): + args = {**named_args, **kwargs} + configs = [c for c in configs if args["K"] % c.kwargs["BLOCK_K"] == 0] + # One K tile per split cannot amortize padded MMA setup for a rank-32 decode matvec. + single_tile_decode = (args["M"] == 1 and args["RANK"] == 32 and args["Y"] is not None + and args["K"] <= args["SPLIT_K"] * max((c.kwargs["BLOCK_K"] for c in configs), default=0)) + return [c for c in configs if (c.kwargs["BLOCK_M"] > 1 or args["M"] <= 16) + and not (c.kwargs["BLOCK_M"] == 1 and c.kwargs["SWAP_AB"]) + and (not single_tile_decode or c.kwargs["BLOCK_M"] == 1)] + + +@bayesian_autotune( + [triton.Config({"BLOCK_M": bm, "BLOCK_K": bk, "SWAP_AB": swap}, num_warps=w, num_stages=s) + for bm in (1, 16, 32, 64) for bk in (64, 128, 256) for w in (4, 8) for s in (1, 3) + for swap in (False, True)], + ["K", "RANK", "m_bucket", "SPLIT_K", "SWIZZLED"], + n_trials=100, + prune_configs_by={"early_config_prune": _svdquant_pruner}, +) +@triton.jit +def _svdquant_act_quant_kernel( + X, PreQuantScale, GlobalScale, + LoraA, + Y, S, SOut, Down, Partials, + M, m_bucket, + stride_x_m, stride_x_k, stride_la_r, stride_la_k, + K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, RANK: tl.constexpr, SPLIT_K: tl.constexpr, SWIZZLED: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr, SWAP_AB: tl.constexpr, +): + """BF16 LoRA-down with optional fused NVFP4 quantization of the smoothed input.""" + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + split = tl.program_id(1) + rank = tl.arange(0, triton.next_power_of_2(RANK)) + if SWAP_AB: + acc = tl.full((triton.next_power_of_2(RANK), BLOCK_M), 0, tl.float32) + else: + acc = tl.full((BLOCK_M, triton.next_power_of_2(RANK)), 0, tl.float32) + if GlobalScale is not None and Y is not None: + global_scale = tl.load(GlobalScale) + for kb in range(split, K // BLOCK_K, SPLIT_K): + cols = kb * BLOCK_K + tl.arange(0, BLOCK_K) + x = tl.load(X + rows[:, None] * stride_x_m + cols[None, :] * stride_x_k, + mask=rows[:, None] < M, other=0) + if PreQuantScale is not None: + x = (x.to(tl.float32) * tl.load(PreQuantScale + cols)[None, :]).to(X.dtype.element_ty) + w = tl.load(LoraA + rank[:, None] * stride_la_r + cols[None, :] * stride_la_k, + mask=rank[:, None] < RANK, other=0) + if BLOCK_M == 1: + acc += tl.sum(x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :], 2) + elif SWAP_AB: + acc = tl.dot(w, tl.trans(x), acc) + else: + acc = tl.dot(x, tl.trans(w), acc) + if Y is not None: + normalized = x.to(tl.float32) + if GlobalScale is not None: + normalized /= global_scale + y, scales = mx_act_quant_inline(normalized, BLOCK_M, BLOCK_K, SCALE_GROUP_K, "nvfp4") + packed_cols = kb * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2) + tl.store(Y + rows[:, None] * (K // 2) + packed_cols[None, :], y, mask=rows[:, None] < M) + store_mx_act_scales(S, SOut, rows, scales, rows < M, kb, M, K, + SCALE_GROUP_K, BLOCK_K, SWIZZLED, "rows") + if SWAP_AB: + acc = tl.trans(acc) + mask = (rows[:, None] < M) & (rank[None, :] < RANK) + if SPLIT_K == 1: + tl.store(Down + rows[:, None] * RANK + rank[None, :], acc, mask=mask) + else: + tl.store(Partials + (split * M + rows[:, None]) * RANK + rank[None, :], acc, mask=mask) + + +@triton.jit +def _svdquant_down_reduce_kernel(Partials, Down, M, RANK: tl.constexpr, SPLIT_K: tl.constexpr, BLOCK: tl.constexpr): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + splits = tl.arange(0, triton.next_power_of_2(SPLIT_K)) + parts = tl.load(Partials + splits[:, None] * M * RANK + offsets[None, :], + mask=(splits[:, None] < SPLIT_K) & (offsets[None, :] < M * RANK), other=0) + tl.store(Down + offsets, tl.sum(parts, 0), mask=offsets < M * RANK) + + +def svdquant_act_quant(x, lora_a, pre_quant_scale, global_scale, swizzled): + """Quantized residual input and BF16 down projection from the same pre-quantized input.""" + K, rank = x.shape[-1], lora_a.shape[0] + x = x.reshape(-1, K) + M = x.shape[0] + separate_quant = is_sm10x() and rank == 32 and M >= 4096 and M * K >= 1 << 25 + if separate_quant: + # Independent quantization tiles avoid keeping the large input in both quant and MMA layouts. + values, scales = _launch_act_quant( + x, "nvfp4", NVFP4_SCALE_GROUP_K, torch.float8_e4m3fn, + swizzled, global_scale, pre_quant_scale=pre_quant_scale) + else: + values = torch.empty(M, K // 2, device=x.device, dtype=torch.int8) + scale_shape = (1, triton.cdiv(M, 128), triton.cdiv(K // NVFP4_SCALE_GROUP_K, 4), 2, 256) if swizzled else (M, K // NVFP4_SCALE_GROUP_K) + scales = torch.empty(scale_shape, device=x.device, dtype=torch.float8_e4m3fn) + split_k = min(32, max(1, triton.cdiv(256, triton.cdiv(M, 32)))) + down = torch.empty(M, rank, device=x.device, dtype=x.dtype) if M > 1 or split_k == 1 else None + partials = torch.empty(split_k, M, rank, device=x.device, dtype=torch.float32) if split_k > 1 else None + with device_context(x.device): + compile_time_only_triton_wrap(_svdquant_act_quant_kernel)[lambda meta: (triton.cdiv(M, meta["BLOCK_M"]), split_k)]( + x, pre_quant_scale, global_scale, lora_a, None if separate_quant else values, + None if separate_quant or swizzled else scales, scales if swizzled and not separate_quant else None, down, partials, + M, int(M).bit_length(), x.stride(0), x.stride(1), lora_a.stride(0), lora_a.stride(1), + K=K, SCALE_GROUP_K=NVFP4_SCALE_GROUP_K, RANK=rank, SPLIT_K=split_k, SWIZZLED=swizzled, + ) + if split_k > 1 and M > 1: + compile_time_only_triton_wrap(_svdquant_down_reduce_kernel)[(triton.cdiv(M * rank, 128),)]( + partials, down, M, RANK=rank, SPLIT_K=split_k, BLOCK=128, + ) + return values, scales, partials if M == 1 and split_k > 1 else down + + def nvfp4_quantize_two_level( weight: torch.Tensor, swizzled: bool = False ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -936,7 +1069,7 @@ def nvfp4_quantize_two_level( def _launch_act_quant( x, activation_format, scale_group, scale_dtype, swizzled=False, global_scale=None, - expert_index=None, gather_idx=None, + expert_index=None, gather_idx=None, pre_quant_scale=None, ): """One-pass activation quant for every format (``mxfp8`` = E4M3 values, else packed E2M1) and both scale layouts. ``swizzled=True`` writes the scale straight into the SWIZZLE_32_4_4 buffer @@ -970,6 +1103,7 @@ def _launch_act_quant( lambda META: (triton.cdiv(T, META["BLOCK_T"]), K // META["BLOCK_K"]) ]( x, + pre_quant_scale, # PreQuantScale values, None if swizzled else scales, # S: row-major scales, plain arm only scales if swizzled else None, # SOut: SWIZZLE_32_4_4 buffer, swizzled arm only @@ -1080,6 +1214,7 @@ def grid(META): @triton.jit def _fp8_act_quant_kernel( x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr, PADDED_BLOCK: tl.constexpr, + AMAX_UPPER_BOUND: tl.constexpr = None, PDL: tl.constexpr = False, ): if PDL: @@ -1092,7 +1227,12 @@ def _fp8_act_quant_kernel( mask = cols < BLOCK_SIZE offs = pid * BLOCK_SIZE + cols x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) - s = tl.max(tl.abs(x)) / 448.0 # float8_e4m3fn max + amax = tl.max(tl.abs(x)) + if AMAX_UPPER_BOUND is not None: + # the amax is capped, and the values past it saturate + amax = tl.minimum(amax, AMAX_UPPER_BOUND) + x = tl.clamp(x, -amax, amax) + s = amax / 448.0 # float8_e4m3fn max y = (x / tl.maximum(s, 1e-12)).to(y_ptr.dtype.element_ty) tl.store(y_ptr + offs, y, mask=mask) if PDL: @@ -1120,7 +1260,11 @@ def static_expert_act_operands(A, As, num_experts): # not alias an input), and it hands back Python strides. The quant kernels it calls carry their # own compile handling. def tensor_wide_act_operands( - A: torch.Tensor, As: torch.Tensor | None, num_experts: int | None = None + A: torch.Tensor, + As: torch.Tensor | None, + num_experts: int | None = None, + row_multiple: int = 1, + activation_scale_ub: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor, int, int]: """Activation operands of the tensor-scale FP8 arms: ``(A, As, stride_m, stride_e)``. @@ -1130,10 +1274,13 @@ def tensor_wide_act_operands( one per expert, which stays raw for the kernel to quantize each tile against its own expert — a gathered row serves several experts, so it has no single pre-quantized form until the routed rows are laid out (``quantize_routed_rows_per_expert``). ``num_experts`` ``None`` on - the 2D op, whose rows belong to no expert. + the 2D op, whose rows belong to no expert. ``row_multiple`` pads the derived rows' storage + (``fp8_act_quant_tensor_wide``). ``activation_scale_ub`` caps their dynamic amax. """ if As is None: - A, As = fp8_act_quant_tensor_wide(A, A.shape[-1]) + A, As = fp8_act_quant_tensor_wide( + A, A.shape[-1], amax_upper_bound=activation_scale_ub, row_multiple=row_multiple + ) As = As.reshape(-1) return A, As, As.stride(0), 0 if A.dtype == FP8_DTYPE: @@ -1217,11 +1364,19 @@ def quantize_routed_rows_per_expert(A, As, gather_idx, expert_start, num_experts def fp8_act_quant_tensor_wide( - x: torch.Tensor, block_size: int = 128 + x: torch.Tensor, + block_size: int = 128, + amax_upper_bound: float | None = None, + row_multiple: int = 1, ) -> tuple[torch.Tensor, torch.Tensor]: + """E4M3 + one ``amax / 448`` fp32 scale per ``block_size`` span (``block_size = K``: per + token). ``amax_upper_bound`` caps the amax (values past it saturate); ``None`` leaves it + uncapped. ``row_multiple`` backs the E4M3 rows with storage rounded up to that many rows + (left unwritten), so a TMA box over the padded extent reads in bounds.""" assert x.is_contiguous() assert x.shape[-1] % block_size == 0 - y = torch.empty_like(x, dtype=FP8_DTYPE) + padded_rows = triton.cdiv(x.numel() // x.shape[-1], row_multiple) * row_multiple + y = x.new_empty(padded_rows * x.shape[-1], dtype=FP8_DTYPE)[: x.numel()].view(x.shape) grid = (triton.cdiv(x.numel(), block_size),) s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32) @@ -1232,9 +1387,541 @@ def fp8_act_quant_tensor_wide( s, BLOCK_SIZE=block_size, PADDED_BLOCK=triton.next_power_of_2(block_size), + AMAX_UPPER_BOUND=amax_upper_bound, **pdl_launch_kwargs(), ) return y, s + + +# ── int4 checkpoint repacks (load time) ────────────────────────────────────── + + +@triton.jit +def _repack_int4_kernel( + QWeight, # GPTQ (K // 8, N) words packed along K; AWQ (K, N // 8) words packed along N + QZeros, # (K // group, N // 8) words packed along N; read iff not SYMMETRIC + Scales, # (K // group, N) + Perm, # (K,) int32 act-order K order: output K row j is source row Perm[j]; None: identity (GPTQ only) + W, # out: (N, K // 8) uint32, eight K values per word in the interleaved order (interleave_int_words) + Ws, # out: (N, K // group) + Wz, # out: (N, K // group) uint8; written iff not SYMMETRIC + N, + stride_qw_0, + stride_qw_1, + stride_qz_0, + stride_qz_1, + stride_s_0, + stride_s_1, + ZEROS_OFFSET: tl.constexpr, + AWQ: tl.constexpr, + SYMMETRIC: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUPS: tl.constexpr, # scale groups per K tile +): + """One ``[BLOCK_SIZE_K, BLOCK_SIZE_N]`` tile of a GPTQ or AWQ int4 linear into the GEMMs' layout, its groups' + scales and zero points with it: each value's nibble out of the source word, then eight per word along K in the + kernels' interleaved order.""" + pid_n = tl.program_id(0) + pid_k = tl.program_id(1) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = pid_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + n_mask = (offs_n < N)[None, :] + # AWQ interleaves a word's eight columns: column c sits in nibble c // 2 + 4 * (c % 2) + n_shift = 4 * ((offs_n % 8) // 2 + 4 * (offs_n % 2)) if AWQ else 4 * (offs_n % 8) + if AWQ: + src = tl.load(QWeight + offs_k[:, None] * stride_qw_0 + (offs_n // 8)[None, :] * stride_qw_1, mask=n_mask) + q = (src >> n_shift[None, :]) & 0xF + else: + src_k = offs_k if Perm is None else tl.load(Perm + offs_k) + src = tl.load(QWeight + (src_k // 8)[:, None] * stride_qw_0 + offs_n[None, :] * stride_qw_1, mask=n_mask) + q = (src >> (4 * (src_k % 8))[:, None]) & 0xF + # the kernels' interleaved word (interleave_int_words): K offset j of the word in nibble j // 2 + 4 * (j % 2) + j = tl.arange(0, 8) + words = tl.sum( + q.reshape(BLOCK_SIZE_K // 8, 8, BLOCK_SIZE_N) << (4 * (j // 2 + 4 * (j % 2)))[None, :, None], axis=1 + ) + offs_w = pid_k * (BLOCK_SIZE_K // 8) + tl.arange(0, BLOCK_SIZE_K // 8) + K_WORDS = tl.num_programs(1) * (BLOCK_SIZE_K // 8) + tl.store(W + offs_n[None, :] * K_WORDS + offs_w[:, None], words, mask=n_mask) + + offs_g = pid_k * GROUPS + tl.arange(0, GROUPS) + n_groups = tl.num_programs(1) * GROUPS + out_g = offs_n[None, :] * n_groups + offs_g[:, None] + s = tl.load(Scales + offs_g[:, None] * stride_s_0 + offs_n[None, :] * stride_s_1, mask=n_mask) + tl.store(Ws + out_g, s, mask=n_mask) + if not SYMMETRIC: + zw = tl.load(QZeros + offs_g[:, None] * stride_qz_0 + (offs_n // 8)[None, :] * stride_qz_1, mask=n_mask) + z = (((zw >> n_shift[None, :]) & 0xF) + ZEROS_OFFSET) & 0xF # a zero of 0 packs as 0xF under offset 1 + tl.store(Wz + out_g, z.to(tl.uint8), mask=n_mask) + + +def _repack_int4(qweight, qzeros, scales, K, awq, zeros_offset, symmetric, perm=None): + groups, N = scales.shape + group = K // groups + assert group * groups == K and group % 8 == 0, f"K={K} is not a whole number of groups of {group}" + block_k = group * max(1, 64 // group) # a K tile of whole groups, at least 64 values + assert K % block_k == 0, f"K={K} must be a multiple of {block_k}" + W = torch.empty(N, K // 8, dtype=torch.uint32, device=qweight.device) + Ws = torch.empty(N, groups, dtype=scales.dtype, device=scales.device) + Wz = None if symmetric else torch.empty(N, groups, dtype=torch.uint8, device=qweight.device) + grid = (triton.cdiv(N, 64), K // block_k) + with device_context(qweight.device): + _repack_int4_kernel[grid]( + qweight, None if symmetric else qzeros, scales, perm, W, Ws, Wz, N, + qweight.stride(0), qweight.stride(1), + *((0, 0) if symmetric else qzeros.stride()), + scales.stride(0), scales.stride(1), + ZEROS_OFFSET=zeros_offset, AWQ=awq, SYMMETRIC=symmetric, + BLOCK_SIZE_N=64, BLOCK_SIZE_K=block_k, GROUPS=block_k // group, + ) + return W, Ws, Wz + + +@triton.jit +def _permute_cols_kernel(X, Perm, Scale, QuantScale, UnitScale, Out, M, K, stride_x_m, stride_o_m, BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr, UNIT_GROUP: tl.constexpr): + """Gather original input columns and optionally apply their AWQ scales in the same pass.""" + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K) + mask = (rows < M)[:, None] & (cols < K)[None, :] + src = cols if Perm is None else tl.load(Perm + cols, mask=cols < K, other=0) + rows64 = rows[:, None].to(tl.int64) + vals = tl.load(X + rows64 * stride_x_m + src[None, :], mask=mask, other=0) + if Scale is not None: + vals = (vals.to(tl.float32) * tl.load(Scale + src, mask=cols < K, other=0)[None, :]).to(X.dtype.element_ty) + if QuantScale is not None: + vals = _quantize_modelopt_input(vals, tl.load(QuantScale), Out.dtype.element_ty) + tl.store(Out + rows64 * stride_o_m + cols[None, :], vals, mask=mask) + if UnitScale is not None: + if tl.program_id(0) == 0: + group = tl.program_id(1) * (BLOCK_K // UNIT_GROUP) + tl.arange(0, BLOCK_K // UNIT_GROUP) + tl.store(UnitScale + group, 127 if UnitScale.dtype.element_ty == tl.uint8 else 1.0, mask=group < K // UNIT_GROUP) + + +def permute_act_cols( + x: torch.Tensor, perm: torch.Tensor | None, pre_quant_scale: torch.Tensor | None = None, + quant_scale: torch.Tensor | None = None, + quant_dtype: torch.dtype = FP8_DTYPE, + unit_scale: torch.Tensor | None = None, + unit_group: int = MX_SCALE_GROUP_K, + storage_dtype: torch.dtype | None = None, +) -> torch.Tensor: + """``x[..., perm]`` for an act-order weight's K order (``b_perm``), one launch: the activation the natural-order + kernels (TMA arms included) then read. ``pre_quant_scale`` optionally multiplies the original columns + before permutation, rounded in the activation dtype; ``perm=None`` only applies the transform.""" + K = x.shape[-1] + if pre_quant_scale is not None: + assert pre_quant_scale.shape == (K,) and pre_quant_scale.is_contiguous(), "pre_quant_scale must be contiguous (K,)" + assert pre_quant_scale.device == x.device and pre_quant_scale.dtype == x.dtype, "pre_quant_scale must match the activation device and dtype" + x2 = x.reshape(-1, K) + out = torch.empty_like(x2, dtype=(storage_dtype or quant_dtype) if quant_scale is not None else x.dtype) + M = x2.shape[0] + # one row per program; a long column span per program once rows are many (B200, K=12288: M=1 1.4 vs 2.3us at + # 512 vs 4096 columns, M=2048 73.0 vs 46.7us) + block_k = min(512 if M <= 64 else 4096, triton.next_power_of_2(K)) + with device_context(x.device): + compile_time_only_triton_wrap(_permute_cols_kernel)[(M, triton.cdiv(K, block_k))]( + x2, perm, pre_quant_scale, quant_scale, unit_scale, out, M, K, x2.stride(0), out.stride(0), + BLOCK_M=1, BLOCK_K=block_k, UNIT_GROUP=unit_group, num_warps=4, + ) + return out.view(x.shape) + + +@triton.jit +def _permute_routed_rows_kernel( + X, Perm, Scale, QuantScale, UnitScale, GatherIdx, ExpertStart, ExpertIds, Out, K, stride_x_m, stride_perm_e, stride_scale_e, stride_quant_e, num_experts, + NUM_EXPERTS_POW2: tl.constexpr, BLOCK_K: tl.constexpr, + UNIT_GROUP: tl.constexpr, +): + """``Out[s, j] = X[src(s), Perm[e(s), j]]`` over one ``BLOCK_K`` column span of routed row ``s``: ``src`` from + ``GatherIdx`` (None: ``s``), the expert from ``ExpertIds`` (batched) or the sorted ranges of ``ExpertStart`` + (grouped). An EP-sentinel row (off-rank expert) is left unwritten, as the GEMM leaves its output.""" + s = tl.program_id(0) + if UnitScale is not None: + if s == 0: + group = tl.program_id(1) * (BLOCK_K // UNIT_GROUP) + tl.arange(0, BLOCK_K // UNIT_GROUP) + tl.store(UnitScale + group, 127 if UnitScale.dtype.element_ty == tl.uint8 else 1.0, mask=group < K // UNIT_GROUP) + if ExpertIds is not None: + e = tl.load(ExpertIds + s) + else: + ends = tl.load(ExpertStart + 1 + tl.arange(0, NUM_EXPERTS_POW2), mask=tl.arange(0, NUM_EXPERTS_POW2) < num_experts, other=2**31 - 1) + e = tl.sum((ends <= s).to(tl.int32)) + if (e >= 0) & (e < num_experts): + src = s if GatherIdx is None else tl.load(GatherIdx + s) + cols = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K) + mask = cols < K + idx = cols if Perm is None else tl.load(Perm + e.to(tl.int64) * stride_perm_e + cols, mask=mask, other=0) + vals = tl.load(X + src.to(tl.int64) * stride_x_m + idx, mask=mask) + if Scale is not None: + scale = tl.load(Scale + e.to(tl.int64) * stride_scale_e + idx, mask=mask, other=0) + vals = (vals.to(tl.float32) * scale).to(X.dtype.element_ty) + if QuantScale is not None: + denom = tl.load(QuantScale + e * stride_quant_e) + vals = _quantize_modelopt_input(vals, denom, Out.dtype.element_ty) + tl.store(Out + s.to(tl.int64) * K + cols, vals, mask=mask) + + +def permute_routed_rows( + x: torch.Tensor, perm: torch.Tensor | None, num_rows: int, gather_idx: torch.Tensor | None = None, + expert_start: torch.Tensor | None = None, expert_ids: torch.Tensor | None = None, + pre_quant_scale: torch.Tensor | None = None, num_experts: int | None = None, + quant_scale: torch.Tensor | None = None, + quant_dtype: torch.dtype = FP8_DTYPE, + unit_scale: torch.Tensor | None = None, + unit_group: int = MX_SCALE_GROUP_K, + storage_dtype: torch.dtype | None = None, +) -> torch.Tensor: + """The ``num_rows`` routed activation rows, each read from ``x`` (through ``gather_idx``) in its expert's + act-order K order, ``perm`` ``(E, K)``: ``out[s] = x[gather_idx[s], perm[expert(s)]]`` with the expert from + ``expert_ids`` (batched) or ``expert_start``'s sorted ranges (grouped). One launch; the routed GEMMs then read + contiguous natural-order rows. ``pre_quant_scale`` additionally multiplies each original column by + its shared or per-expert AWQ input scale, rounded in the activation dtype; ``perm=None`` retains K order.""" + K = x.shape[-1] + out = x.new_empty(num_rows, K, dtype=(storage_dtype or quant_dtype) if quant_scale is not None else x.dtype) + if num_experts is None: + num_experts = perm.shape[0] + stride_scale_e = 0 + if pre_quant_scale is not None: + assert pre_quant_scale.shape in ((K,), (1, K), (num_experts, K)) and pre_quant_scale.is_contiguous(), ( + "pre_quant_scale must be contiguous (K,), (1, K), or (E, K)" + ) + assert pre_quant_scale.device == x.device and pre_quant_scale.dtype == x.dtype, "pre_quant_scale must match the activation device and dtype" + stride_scale_e = K if pre_quant_scale.numel() > K else 0 + block_k = min(512, triton.next_power_of_2(K)) + num_warps = 4 + if num_rows > 64 and is_sm10x(): + block_k = min(4096 if perm is not None else 2048, triton.next_power_of_2(K)) + num_warps = 2 if block_k > 2048 else 1 + with device_context(x.device): + compile_time_only_triton_wrap(_permute_routed_rows_kernel)[(num_rows, triton.cdiv(K, block_k))]( + x, perm, pre_quant_scale, quant_scale, unit_scale, gather_idx, expert_start, expert_ids, out, K, x.stride(0), + perm.stride(0) if perm is not None else 0, stride_scale_e, global_scale_stride(quant_scale), num_experts, + NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts), BLOCK_K=block_k, UNIT_GROUP=unit_group, num_warps=num_warps, + ) + return out + + +def prepare_fp8_mx_inputs( + A, input_scale, scale_dtype, group, num_experts=1, num_rows=None, + gather_idx=None, expert_start=None, expert_ids=None, pre_quant_scale=None, storage_dtype=FP8_DTYPE, +): + """Calibrated tensor-FP8 inputs for FP4/MX loops: quantized rows, unit group scales, and calibration. + + Constant unit group scales preserve ordinary FP8 quantization (one calibrated scale per tensor/expert), + rather than re-quantizing each block as MXFP8. Routed quantization and the sorted copy share one pass. + BF16 stores FP8-rounded inputs exactly, as well as every finite E2M1 x E4M3 weight-scale product. + This lets the wide-K weight-upcast GEMM use the same real operands without unit group scales. + """ + assert input_scale is not None and A.dtype in (torch.bfloat16, torch.float16, torch.float32), ( + "FP4 with ordinary FP8 activations requires raw inputs and a calibrated input_scale" + ) + assert storage_dtype in (FP8_DTYPE, torch.bfloat16) + calibration = normalize_global_scale(input_scale, num_experts) + K = A.shape[-1] + unit = None if storage_dtype == torch.bfloat16 else torch.empty((1, K // group), device=A.device, dtype=scale_dtype) + if num_rows is None: + values = permute_act_cols(A.reshape(-1, K), None, pre_quant_scale, quant_scale=calibration, unit_scale=unit, unit_group=group, storage_dtype=storage_dtype) + else: + values = permute_routed_rows( + A, None, num_rows, gather_idx, expert_start, expert_ids, + num_experts=num_experts, quant_scale=calibration, unit_scale=unit, unit_group=group, pre_quant_scale=pre_quant_scale, storage_dtype=storage_dtype, + ) + # One row, written by the preparation pass, broadcasts across tokens without a full scale grid. + scales = None if storage_dtype == torch.bfloat16 else unit.expand(values.shape[0], K // group) + return values.reshape(A.shape) if num_rows is None else values, scales, calibration + + +def interleave_int_words(words: torch.Tensor, bits: int) -> torch.Tensor: + """Natural-order integer weight words (``(..., K // VPW)`` int32, ``VPW = 32 // bits`` unsigned values per word + along K, the first lowest: compressed-tensors' ``pack-quantized`` layout) into the integer GEMMs' interleaved + layout, uint32 words (the dtype marks the layout): within each word, the value at K offset ``j`` moves to slot + ``j // 2 + (j % 2) * VPW // 2``, so one shift and one mask yield K-adjacent value pairs in a lane's two 16-bit + halves (``int_word_pairs``). A loader-time conversion, done once per weight (the GPTQ / AWQ repacks emit this + layout directly).""" + assert words.dtype == torch.int32 and bits in (2, 4, 8), f"int32 words of 2/4/8-bit values, got {words.dtype}/{bits}" + vpw = 32 // bits + mask = (1 << bits) - 1 + j = torch.arange(vpw, device=words.device, dtype=torch.int32) + values = (words.unsqueeze(-1) >> (bits * j)) & mask # (..., words, vpw), natural slot j + slot = j // 2 + (j % 2) * (vpw // 2) + return (values << (bits * slot)).sum(-1, dtype=torch.int32).view(torch.uint32) + + +@triton.jit +def _dequantize_int_kernel( + W, # (rows, K // VPW) uint32 words, interleaved (interleave_int_words) + Ws, # (rows, K // SCALE_GROUP_K) float group scales + Wz, # (rows, K // SCALE_GROUP_K) integer zero points; None for the symmetric 2 ** (BITS - 1) + Out, # (rows, K) + rows, + K_WORDS, + BITS: tl.constexpr, + SCALE_GROUP_K: tl.constexpr, + BLOCK_SIZE_R: tl.constexpr, + BLOCK_SIZE_W: tl.constexpr, +): + """One ``[BLOCK_SIZE_R, BLOCK_SIZE_W]`` word tile back to its ``[BLOCK_SIZE_R, BLOCK_SIZE_W * VPW]`` values: + ``((q - zero) * scale)`` rounded in the scale dtype, then cast to ``Out``'s.""" + VPW: tl.constexpr = 32 // BITS + HALF: tl.constexpr = VPW // 2 + MASK: tl.constexpr = (1 << BITS) - 1 + offs_r = tl.program_id(0) * BLOCK_SIZE_R + tl.arange(0, BLOCK_SIZE_R) + offs_w = tl.program_id(1) * BLOCK_SIZE_W + tl.arange(0, BLOCK_SIZE_W) + r_mask = (offs_r < rows)[:, None] + w = tl.load(W + offs_r[:, None].to(tl.int64) * K_WORDS + offs_w[None, :], mask=r_mask & (offs_w < K_WORDS)[None, :]) + shifts = (tl.arange(0, HALF) * BITS)[None, None, :] + # slot s < HALF holds K offset 2s, slot s + HALF offset 2s + 1 + lo = (w[:, :, None] >> shifts) & MASK + hi = (w[:, :, None] >> (shifts + HALF * BITS)) & MASK + q = tl.join(lo, hi).reshape(BLOCK_SIZE_R, BLOCK_SIZE_W * VPW).to(tl.int32) + offs_k = tl.program_id(1) * BLOCK_SIZE_W * VPW + tl.arange(0, BLOCK_SIZE_W * VPW) + mask = r_mask & (offs_k < K_WORDS * VPW)[None, :] + groups = K_WORDS * VPW // SCALE_GROUP_K + g_ptrs = offs_r[:, None].to(tl.int64) * groups + (offs_k // SCALE_GROUP_K)[None, :] + scale = tl.load(Ws + g_ptrs, mask=mask) + zero = (1 << (BITS - 1)) if Wz is None else tl.load(Wz + g_ptrs, mask=mask).to(tl.int32) + v = (q - zero).to(scale.dtype) * scale + out_ptrs = Out + offs_r[:, None].to(tl.int64) * (K_WORDS * VPW) + offs_k[None, :] + tl.store(out_ptrs, v.to(Out.dtype.element_ty), mask=mask) + + +def dequantize_int_weight( + words: torch.Tensor, scales: torch.Tensor, K: int, zeros: torch.Tensor | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """An integer weight in the GEMMs' interleaved layout (``(..., N, K // VPW)`` uint32, ``interleave_int_words``) + back to ``(..., N, K)`` the way the GEMMs dequantize it: ``(q - zero)`` in the scale dtype times the per-group + scale, rounded once, then cast to ``dtype`` (``None``: the scale dtype). The width follows from ``K`` and the + word count; ``zeros`` ``None`` is the symmetric zero point ``2 ** (bits - 1)``. One fused launch; the integer + dgrad's operand.""" + bits = int_weight_bits(K, words) + k_words = words.shape[-1] + group = K // scales.shape[-1] + rows = words.numel() // k_words + assert group * scales.shape[-1] == K and group % (32 // bits) == 0, f"{scales.shape[-1]} groups cannot split K={K}" + assert zeros is None or zeros.shape == scales.shape, "zero points are shaped like the scales" + out = torch.empty(*words.shape[:-1], K, device=words.device, dtype=dtype or scales.dtype) + + grid = (triton.cdiv(rows, 32), triton.cdiv(k_words, 16)) + with device_context(words.device): + compile_time_only_triton_wrap(_dequantize_int_kernel)[grid]( + words.contiguous(), scales.contiguous(), None if zeros is None else zeros.contiguous(), out, + rows, k_words, BITS=bits, SCALE_GROUP_K=group, BLOCK_SIZE_R=32, BLOCK_SIZE_W=16, + ) + return out + + +def repack_gptq_int4( + qweight: torch.Tensor, + qzeros: torch.Tensor, + scales: torch.Tensor, + g_idx: torch.Tensor | None = None, + zeros_offset: int = 1, + symmetric: bool = False, +) -> tuple[torch.Tensor, ...]: + """A GPTQ int4 linear (``qweight`` ``(K // 8, N)``, ``qzeros`` ``(K // group, N // 8)``, ``scales`` + ``(K // group, N)``) as the int4 GEMMs' ``(weight, scale, zero)``: ``(N, K // 8)`` uint32, ``(N, K // group)``, + ``(N, K // group)`` uint8 or ``None`` for a ``symmetric`` checkpoint (``sym=True``: every zero point 8). + ``zeros_offset`` is what the checkpoint subtracted from each zero point before packing: 1 in the original GPTQ + format, 0 in ``"gptq_v2"``. An act-order ``g_idx`` (``desc_act=True``: groups scattered across K) sorts K by + group, so the weight holds ordinary contiguous groups, and returns that order too, ``(weight, scale, zero, + perm)``: ``perm`` ``(K,)`` int32 is the ops' ``b_perm`` (the GEMM contracts ``A[..., perm]``). An identity + ``g_idx`` (or none) keeps the three-tensor form.""" + K = qweight.shape[0] * 8 + group = K // scales.shape[0] + natural = torch.arange(K, device=qweight.device, dtype=torch.int32) // group + if g_idx is None or torch.equal(g_idx.to(torch.int32), natural): + return _repack_int4(qweight, qzeros, scales, K, False, zeros_offset, symmetric) + perm = torch.argsort(g_idx, stable=True).to(torch.int32) + if not torch.equal(g_idx.to(torch.int32)[perm.long()], natural): + raise ValueError(f"g_idx does not assign exactly {group} K rows to each of its {K // group} groups") + return (*_repack_int4(qweight, qzeros, scales, K, False, zeros_offset, symmetric, perm), perm) + + +def fold_act_order_into_gate_up( + gate_up: torch.Tensor, + gate_up_scale: torch.Tensor, + gate_up_zero: torch.Tensor | None, + gate_up_bias: torch.Tensor | None, + down_perm: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Fold an act-order down projection's K order into the gate|up projection that feeds it, at load: the down + contracts its input in ``down_perm`` order, so gate|up output pair ``j`` (rows ``2j`` gate, ``2j + 1`` up, + interleaved) takes rows ``2 * down_perm[j]`` and ``2 * down_perm[j] + 1``, with their scales, zero points and + bias. The GLU output then comes out already in the down's order and the down runs with ``b_perm=None``. + ``(E, 2I, ...)`` stacks with ``down_perm`` ``(E, I)``, or one linear: ``(2I, ...)`` with ``(I,)``.""" + rows = torch.stack([2 * down_perm, 2 * down_perm + 1], dim=-1).flatten(-2).long() # (..., 2I) + + def take_rows(t): + return None if t is None else torch.gather(t, rows.ndim - 1, rows[..., None].expand(*rows.shape, t.shape[-1])) + + bias = None if gate_up_bias is None else torch.gather(gate_up_bias, rows.ndim - 1, rows) + return take_rows(gate_up), take_rows(gate_up_scale), take_rows(gate_up_zero), bias + + +def repack_awq_int4( + qweight: torch.Tensor, qzeros: torch.Tensor, scales: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """An AWQ (GEMM) int4 linear (``qweight`` ``(K, N // 8)``, ``qzeros`` ``(K // group, N // 8)``, both packed along N + in AWQ's interleaved nibble order, ``scales`` ``(K // group, N)``) as the int4 GEMMs' ``(weight, scale, zero)``.""" + return _repack_int4(qweight, qzeros, scales, qweight.shape[0], True, 0, False) + + +@triton.jit +def _pack_interleaved_int_tile(q, BITS: tl.constexpr): + VPW: tl.constexpr = 32 // BITS + j = tl.arange(0, VPW) + slot = j // 2 + (j % 2) * (VPW // 2) + return tl.sum(q << (BITS * slot)[None, None, :], axis=2) + + +@triton.jit +def _repack_modelopt_int_kernel( + QWeight, W, N, K, stride_e, stride_n, stride_k, + BITS: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_W: tl.constexpr, +): + """Signed ModelOpt values to offset-binary, K-interleaved words without an unpacked weight bank.""" + VPW: tl.constexpr = 32 // BITS + n = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + w = tl.program_id(1) * BLOCK_W + tl.arange(0, BLOCK_W) + j = tl.arange(0, VPW) + k = w[:, None] * VPW + j[None, :] + if BITS == 4: + src_n = n // 2 + else: + src_n = n + raw = tl.load( + QWeight + tl.program_id(2).to(tl.int64) * stride_e + + src_n[:, None, None].to(tl.int64) * stride_n + k[None, :, :] * stride_k, + mask=(n < N)[:, None, None] & (k < K)[None, :, :], other=0, + ).to(tl.uint32) + if BITS == 4: + raw >>= (4 * (n % 2))[:, None, None] + q = (raw & ((1 << BITS) - 1)) ^ (1 << (BITS - 1)) + packed = _pack_interleaved_int_tile(q, BITS) + tl.store( + W + (tl.program_id(2).to(tl.int64) * N + n[:, None]) * (K // VPW) + w[None, :], + packed, mask=(n < N)[:, None] & (w < K // VPW)[None, :], + ) + + +def _repack_modelopt_int(qweight, scales, bits): + expected_dtype = torch.uint8 if bits == 4 else torch.int8 + assert qweight.dtype == expected_dtype and qweight.ndim in (2, 3), ( + f"ModelOpt INT{bits} weights must be 2D or 3D {expected_dtype}, got {qweight.dtype}/{qweight.ndim}D" + ) + assert scales.device == qweight.device and scales.dtype in (torch.bfloat16, torch.float16, torch.float32), ( + "ModelOpt weight scales must be bf16/fp16/fp32 on the weight device" + ) + N = qweight.shape[-2] * (2 if bits == 4 else 1) + K = qweight.shape[-1] + leading = qweight.shape[:-2] + if bits == 8 and scales.shape == (*leading, N): + scales = scales.unsqueeze(-1) + assert scales.ndim == qweight.ndim and scales.shape[:-1] == (*leading, N), ( + f"ModelOpt weight scales must have shape {(*leading, N, 'K/group')}, got {tuple(scales.shape)}" + ) + groups = scales.shape[-1] + vpw = 32 // bits + assert groups > 0 and K % groups == 0 and (K // groups) % vpw == 0, ( + f"K={K} must contain whole scale groups aligned to {vpw} values" + ) + assert bits == 4 or groups == 1, "ModelOpt INT8 weight-only uses one scale per output channel" + W = torch.empty(*leading, N, K // vpw, dtype=torch.uint32, device=qweight.device) + experts = qweight.shape[0] if qweight.ndim == 3 else 1 + with device_context(qweight.device): + _repack_modelopt_int_kernel[(triton.cdiv(N, 16), triton.cdiv(K // vpw, 16), experts)]( + qweight, W, N, K, qweight.stride(0) if leading else 0, *qweight.stride()[-2:], + BITS=bits, BLOCK_N=16, BLOCK_W=16, num_warps=4, + ) + return W, scales.contiguous(), None + + +def repack_modelopt_int4( + qweight: torch.Tensor, scales: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, None]: + """ModelOpt ``INT4_AWQ`` / blockwise weight-only export to ``(weight, scale, None)`` for the integer GEMMs. + + ``qweight`` is uint8 ``(..., N // 2, K)``: output row ``2n`` is the low signed two's-complement nibble, + row ``2n+1`` the high nibble. ``scales`` is ``(..., N, K // group)``. Returns K-interleaved uint32 + words ``(..., N, K // 8)`` with implicit zero point 8; scales keep their dtype and calibrated values. + Pass the checkpoint's ``pre_quant_scale``, when present, to the GEMM as ``a_pre_quant_scale`` (or fuse it + upstream), exactly as in the exported model. This conversion neither applies it nor requantizes weights. + ``W4A8_AWQ`` uses the same packing, plus ``activation_format='fp8'``, ``As=input_scale``, and + ``b_global_scale=weight_scale_2`` on the GEMM; a repack alone does not apply that recipe. + """ + return _repack_modelopt_int(qweight, scales, 4) + + +def repack_modelopt_int8( + qweight: torch.Tensor, scales: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, None]: + """ModelOpt ``INT8_WEIGHT_ONLY`` signed int8 ``(..., N, K)`` to the integer GEMMs' K-interleaved uint32 + words, implicit zero point 128. Per-channel ``scales`` is ``(..., N)`` or ``(..., N, 1)``; returned as + ``(..., N, 1)`` without changing its dtype or values. This is W8A16, not INT8 SmoothQuant W8A8. + """ + return _repack_modelopt_int(qweight, scales, 8) + + +@triton.jit +def _repack_ggml_int_kernel( + Blocks, W, Scales, ROWS, GROUPS, BITS: tl.constexpr, BLOCK_R: tl.constexpr, BLOCK_W: tl.constexpr, +): + VPW: tl.constexpr = 32 // BITS + BLOCK_BYTES: tl.constexpr = 2 + 32 * BITS // 8 + r = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R) + w = tl.program_id(1) * BLOCK_W + tl.arange(0, BLOCK_W) + k = w[:, None] * VPW + tl.arange(0, VPW)[None, :] + block = r[:, None, None].to(tl.int64) * GROUPS + k[None, :, :] // 32 + qoffset = k % 16 if BITS == 4 else k % 32 + raw = tl.load( + Blocks + block * BLOCK_BYTES + 2 + qoffset[None, :, :], + mask=(r < ROWS)[:, None, None] & (k < GROUPS * 32)[None, :, :], other=0, + ).to(tl.uint32) + if BITS == 4: + q = (raw >> (4 * ((k % 32) // 16))[None, :, :]) & 15 + else: + q = raw ^ 128 # signed byte to offset-binary + words = _pack_interleaved_int_tile(q, BITS) + tl.store(W + r[:, None].to(tl.int64) * (GROUPS * 32 // VPW) + w[None, :], words, + mask=(r < ROWS)[:, None] & (w < GROUPS * 32 // VPW)[None, :]) + group = w // (32 // VPW) + scale_ptr = Blocks + (r[:, None].to(tl.int64) * GROUPS + group[None, :]) * BLOCK_BYTES + mask = (r < ROWS)[:, None] & (group < GROUPS)[None, :] + lo = tl.load(scale_ptr, mask=mask, other=0).to(tl.uint16) + hi = tl.load(scale_ptr + 1, mask=mask, other=0).to(tl.uint16) + d = (lo | (hi << 8)).to(tl.float16, bitcast=True).to(tl.float32) + tl.store(Scales + r[:, None].to(tl.int64) * GROUPS + group[None, :], d, + mask=mask & (w % (32 // VPW) == 0)[None, :]) + + +def repack_ggml_int(packed: torch.Tensor, quant_type: str) -> tuple[torch.Tensor, torch.Tensor, None]: + """GGML ``Q4_0`` / ``Q8_0`` payload blocks to the existing integer GEMMs, without a GGUF parser. + + ``packed`` is uint8 ``(..., N, K // 32, block_bytes)`` (18 for Q4_0, 34 for Q8_0), in canonical + little-endian GGML block order. A GGUF loader reshapes a tensor's raw payload to this shape first. + Q4_0 stores the first 16 values in the low nibbles, the next 16 in the high ones; its zero is 8. + Q8_0 stores signed int8 values; its zero becomes 128. Both use one FP16 scale per 32 values. + Returns ``(interleaved_uint32, fp32_scales, None)``. Scales expand to FP32 because GGML's + reference reconstructs weights in FP32 before the GEMM casts them to its activation dtype. + No requantization, dense weight bank, or codebook-format approximation is performed. + """ + assert quant_type in ("Q4_0", "Q8_0"), "only affine symmetric GGML Q4_0 and Q8_0 blocks are supported" + bits = 4 if quant_type == "Q4_0" else 8 + block_bytes = 2 + 4 * bits + assert packed.dtype == torch.uint8 and packed.ndim in (3, 4) and packed.shape[-1] == block_bytes, ( + f"{quant_type} needs uint8 (..., N, K // 32, {block_bytes}) payload blocks" + ) + shape = packed.shape[:-2] + groups = packed.shape[-2] + rows = packed.numel() // (groups * block_bytes) + W = torch.empty(*shape, groups * bits, device=packed.device, dtype=torch.uint32) + scales = torch.empty(*shape, groups, device=packed.device, dtype=torch.float32) + with device_context(packed.device): + _repack_ggml_int_kernel[(triton.cdiv(rows, 16), triton.cdiv(groups * bits, 16))]( + packed.contiguous(), W, scales, rows, groups, BITS=bits, BLOCK_R=16, BLOCK_W=16, num_warps=4, + ) + return W, scales, None diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/scheduling.py b/finegrained-kernels/torch-ext/finegrained_kernels/scheduling.py index 49b1358df..9d06c5702 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/scheduling.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/scheduling.py @@ -63,10 +63,8 @@ def resolve_tile_inline( sorted_indices = tl.max_contiguous(m_start + offs, BLOCK_SIZE_M) return expert_id, sorted_indices, row_mask -# Flat-slot tile per program for the O(S) routing kernels (count + scatter). These are small -# latency-bound atomic kernels that want many programs: a sweep over {256..4096} x prefill shapes -# put 256 best (or within ~1%) for both, with 1024 up to ~1.5x slower. The grid derives from it -# so the two can't drift. Power of 2. +# Scatter tiles and the per-slot atomic fallback favor many small programs. Aggregated +# counting uses larger tiles to reduce global atomic traffic further. _ROUTING_BLOCK_SIZE = 256 @@ -106,10 +104,11 @@ def _scatter_kernel( NUM_TOP_K: tl.constexpr, NUM_EXPERTS: tl.constexpr, BLOCK_SIZE: tl.constexpr, + AGGREGATE: tl.constexpr = False, ): - """Counting-sort scatter: each flat slot atomically claims the next slot of its expert - (``expert_start[e] + counter[e]++``). O(S), replaces an O(S·logS) argsort. Within-expert - order is arbitrary (atomic race) — fine, the per-token reduce is order-invariant. Slots whose + """Counting-sort scatter: reserve per-expert slots, aggregating reservations within a + block for large routes. O(S), replaces an O(S·logS) argsort. Within-expert + order is arbitrary — the per-token reduce is order-invariant. Slots whose expert is non-local (an ``is_off_rank`` EP sentinel) fill the tail past ``expert_start[E]``, one atomic per block, so every entry of both maps is a real slot / token and a consumer that runs over all ``S`` rows reads in bounds; the tiles never reach them.""" @@ -118,39 +117,55 @@ def _scatter_kernel( expert_id = tl.load(ExpertIds + offs, mask=in_range, other=NUM_EXPERTS) sentinel = in_range & is_off_rank(expert_id, NUM_EXPERTS) valid = in_range & ~sentinel - local = tl.load(ExpertStart + expert_id, mask=valid, other=0) + tl.atomic_add( - Counters + expert_id, 1, mask=valid, sem="relaxed" - ) - dest = local - n_sentinel = tl.sum(sentinel.to(tl.int32), 0) - if n_sentinel > 0: # without EP no block takes this branch - rank = tl.cumsum(sentinel.to(tl.int32), 0) - sentinel.to(tl.int32) - tail = tl.load(ExpertStart + NUM_EXPERTS) + tl.atomic_add(Counters + NUM_EXPERTS, n_sentinel, sem="relaxed") - dest = tl.where(valid, local, tail + rank) - stored = valid | sentinel + if AGGREGATE: + # Contiguous local expert runs reserve one range each, including the off-rank tail. + expert_id = tl.where(valid, expert_id, NUM_EXPERTS) + packed = tl.sort((expert_id.to(tl.uint64) << 32) | offs.to(tl.uint64), descending=False) + expert_id, offs = (packed >> 32).to(tl.int32), packed.to(tl.uint32) + counts = tl.histogram(expert_id, 2 * NUM_EXPERTS, mask=offs < S) + bins = tl.arange(0, 2 * NUM_EXPERTS) + base = tl.atomic_add(Counters + bins, counts, mask=(bins <= NUM_EXPERTS) & (counts > 0), sem="relaxed") + base += tl.load(ExpertStart + bins, mask=bins <= NUM_EXPERTS, other=0) + base -= tl.cumsum(counts, 0) - counts + dest = tl.gather(base, expert_id, axis=0) + tl.arange(0, BLOCK_SIZE) + stored = offs < S + else: + local = tl.load(ExpertStart + expert_id, mask=valid, other=0) + tl.atomic_add( + Counters + expert_id, 1, mask=valid, sem="relaxed" + ) + dest = local + n_sentinel = tl.sum(sentinel.to(tl.int32), 0) + if n_sentinel > 0: # without EP no block takes this branch + rank = tl.cumsum(sentinel.to(tl.int32), 0) - sentinel.to(tl.int32) + tail = tl.load(ExpertStart + NUM_EXPERTS) + tl.atomic_add(Counters + NUM_EXPERTS, n_sentinel, sem="relaxed") + dest = tl.where(valid, local, tail + rank) + stored = valid | sentinel tl.store(Perm + dest, offs, mask=stored) tl.store(PermToken + dest, offs // NUM_TOP_K, mask=stored) @triton.jit def _count_kernel( - ExpertIds, ExpertFreq, S, NUM_EXPERTS: tl.constexpr, BLOCK_SIZE: tl.constexpr + ExpertIds, ExpertFreq, S, NUM_EXPERTS: tl.constexpr, BLOCK_SIZE: tl.constexpr, AGGREGATE: tl.constexpr = False, ): - """Per-expert token count via atomics — replaces ``torch.histc`` (no float cast), fixed - ``(NUM_EXPERTS,)`` output stays CUDA-graph friendly. ``ExpertFreq`` is pre-zeroed.""" + """Per-expert counts, aggregated locally before global atomics on large routes. + Fixed ``(NUM_EXPERTS,)`` output stays graph friendly; ``ExpertFreq`` is pre-zeroed.""" offs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offs < S expert_id = tl.load(ExpertIds + offs, mask=mask, other=NUM_EXPERTS) - tl.atomic_add( - ExpertFreq + expert_id, 1, mask=mask & ~is_off_rank(expert_id, NUM_EXPERTS), sem="relaxed" - ) + valid = mask & ~is_off_rank(expert_id, NUM_EXPERTS) + if AGGREGATE: + counts = tl.histogram(expert_id, NUM_EXPERTS, mask=valid) + tl.atomic_add(ExpertFreq + tl.arange(0, NUM_EXPERTS), counts, mask=counts > 0, sem="relaxed") + else: + tl.atomic_add(ExpertFreq + expert_id, 1, mask=valid, sem="relaxed") def compute_grouped_scheduling( expert_ids: torch.Tensor, num_experts: int, num_top_k: int ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """On-device routing: expert-sorted index (no copy of the activations) via two Triton - launches — exclusive offsets + an atomic counting-sort scatter (replaces host ``argsort``). + """On-device routing: count, exclusive offsets and counting-sort scatter build an + expert-sorted index without copying the activations (replaces host ``argsort``). Run it once per layer and pass the results to every grouped GEMM of that layer. Returns ``(expert_start, gather_idx, scatter_idx)``: @@ -193,15 +208,19 @@ def _compute_grouped_scheduling( counters = torch.empty(num_experts + 1, dtype=torch.int32, device=device) perm = torch.empty(num_routed_tokens, dtype=torch.int32, device=device) perm_token = torch.empty(num_routed_tokens, dtype=torch.int32, device=device) + # Bound local histogram storage; small routes and larger expert sets keep per-slot atomics. + aggregate = num_routed_tokens >= 4096 and num_experts <= 2 * _ROUTING_BLOCK_SIZE + count_block = 1024 if aggregate else _ROUTING_BLOCK_SIZE with device_context(device): compile_time_only_triton_wrap(_count_kernel)[ - (triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),) + (triton.cdiv(num_routed_tokens, count_block),) ]( expert_ids, expert_freq, num_routed_tokens, NUM_EXPERTS=num_experts, - BLOCK_SIZE=_ROUTING_BLOCK_SIZE, + BLOCK_SIZE=count_block, + AGGREGATE=aggregate, ) compile_time_only_triton_wrap(_exclusive_offsets_kernel)[(1,)]( expert_freq, @@ -221,6 +240,7 @@ def _compute_grouped_scheduling( NUM_TOP_K=num_top_k, NUM_EXPERTS=num_experts, BLOCK_SIZE=_ROUTING_BLOCK_SIZE, + AGGREGATE=aggregate, ) return perm_token, perm, expert_start @@ -237,20 +257,23 @@ def packed_schedule_capacity(num_routed_tokens: int, num_experts: int) -> int: return (num_routed_tokens + PACKED_SCHEDULE_BM - 1) // PACKED_SCHEDULE_BM + num_experts -def build_packed_schedule( - expert_start: torch.Tensor, num_routed_tokens: int, num_experts: int -) -> torch.Tensor: - """Allocate + fill the packed tile schedule (one launch, fixed shape — graph-safe).""" - schedule = torch.empty( - 1 + 3 * packed_schedule_capacity(num_routed_tokens, num_experts), - dtype=torch.int32, device=expert_start.device, - ) - pow2 = triton.next_power_of_2(num_experts) +def packed_schedule_buffer(num_routed_tokens: int, num_experts: int, device: torch.device) -> torch.Tensor: + """The packed tile schedule's buffer: fixed shape (graph-safe), shared by every config of a launch so the tuned + configs see identical args; ``fill_packed_schedule`` writes it for the configs that read it.""" + return torch.empty(1 + 3 * packed_schedule_capacity(num_routed_tokens, num_experts), dtype=torch.int32, device=device) + + +def fill_packed_schedule(nargs): + """Per-config pre_hook: build ``Schedule`` from ``ExpertStart`` (one launch) for a ``PACKED_SCHEDULE`` config. + The other configs never read it, so their launches (every decode crown) carry no schedule kernel.""" + if not nargs.get("PACKED_SCHEDULE"): + return + expert_start = nargs["ExpertStart"] + pow2 = triton.next_power_of_2(expert_start.shape[0] - 1) with device_context(expert_start.device): compile_time_only_triton_wrap(_packed_schedule_kernel)[(pow2,)]( - expert_start, schedule, NUM_EXPERTS=pow2 + expert_start, nargs["Schedule"], NUM_EXPERTS=pow2 ) - return schedule @triton.jit @@ -288,6 +311,18 @@ def load_packed_schedule(Schedule, BLOCK_SIZE_M: tl.constexpr): return tl.load(Schedule) +@triton.jit +def grouped_tile_order(tile_id, num_m_tiles, num_n_tiles): + """Reuse each weight tile across up to eight adjacent M-tiles before advancing N.""" + group_first_m = tile_id // (8 * num_n_tiles) * 8 + # Keep a trailing prefetch past the M-grid well-defined. + group_size_m = tl.maximum(1, tl.minimum(8, num_m_tiles - group_first_m)) + group_offset = tile_id % (8 * num_n_tiles) + pid_m = group_first_m + group_offset % group_size_m + pid_n = group_offset // group_size_m + return pid_m * num_n_tiles + pid_n + + @triton.jit def prefetch_packed_entry(Schedule, pid_m, total_m_tiles): """Issue a future M-tile's ``(expert_id, m_start, row_end)`` loads (masked past the @@ -403,7 +438,7 @@ def _gather_local_rows_kernel(Src, Dst, GatherIdx, ExpertStart, row_words, NUM_E cols = c + tl.arange(0, BLOCK_C) mask = row_mask[:, None] & (cols < row_words)[None, :] vals = tl.load(Src + src[:, None] * row_words + cols[None, :], mask=mask) - tl.store(Dst + rows.to(tl.int64)[:, None] * row_words + cols[None, :], vals, mask=mask) + tl.store(Dst + rows.to(tl.int64)[:, None] * row_words + cols[None, :], vals, mask=mask, cache_modifier=".cs") def _gather_local_rows(x, gather_idx, expert_start, num_experts):