Repository navigation
finegrained-kernels: add rowwise FP8 across GEMMs and MoE - #1240
Open
IlyasMoutawwakil wants to merge 26 commits into
Open
IlyasMoutawwakil wants to merge 26 commits into
IlyasMoutawwakil wants to merge 26 commits into
Conversation
…token activations) E4M3 weights with one fp32 scale per output row and dynamic per-token activation scales (the Llama 4 FP8 recipe) run on the tensor-wide kernels: the weight scale is constant along K, so it folds after the K-loop like a per-tensor one, read per row (a per-tensor scale is stride 0). - formats: is_rowwise / per_token_activations / row_scale_stride; weight_block_size routes (E, N, 1) to the tensor-wide ops - kernels: apply_row_scale in the 2D / batched / grouped tensor-wide kernels (gate|up rows carry their scales interleaved) - 2D tensor-wide kernel: host-TMA A/B arms, SWAP_AB decode rows, PDL, derived per-token rows padded to the largest M tile, native_mma_m_pruner, 200 trials - matmul_2d: cuBLAS RowWise route for 65 <= M < 4096; the scaled_mm routes are plain functions inside the opaque ops - MoE: per-token formats hand the down a bf16 intermediate - quant: fp8_act_quant_tensor_wide gains amax_upper_bound and row_multiple - backward: row-wise 2D dgrad on the natural-tile kernel (scale groups wider than the K tile); grouped tensor-wide / row-wise dgrad registered - fixes: the grouped prefill expansion carries per-token scales with their rows; WS + descriptor compile failures fenced on the grouped tensor-wide kernel - tests: fp8_rowwise registry row and scenarios, scaled_mm parity, MoE cells, amax bound, row-wise dgrad, forced-config sweep cells Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Preserve compact prequantized inputs with safe TMA padding, mask tail activation scales, and normalize leading activation dimensions. Separate gated batched tuning keys and keep the measured decode tile pruning specific to Blackwell. Add public-loader eager and fullgraph regressions and document the batched training boundary.
|
This PR changes source code for the following kernel(s): If you have the kernel-bot build permissions, then trigger |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Support per-output-row FP8 weight scales and per-token activation scales across dense, batched, grouped, and MoE GEMMs, reusing existing loading, scaling, scheduling, and epilogue helpers. Fused MoE requires matching gate/up and down recipes. Dgrad stays in its separate stream.
Dense rowwise dispatch uses public
torch.nn.functional.scaled_mmin its measured M65–8192 range, preserving Triton decode and accumulation precision. Eager, fullgraph compilation, and CUDA graph capture were checked to reach the public route with identical outputs.bench_moe.pyadds Llama 4 Scout routed experts and its actual GQA QKV shape (N7168/K5120), comparing Torch, FBGEMM/MSLK, and vLLM. All 48 timing cells passed. The Torch reference uses the public API; MoE uses per-expert dense GEMMs with fixed routing bounds prepared outside timing. Router/shared expert are excluded and sigmoid input weights are prepared equally for every arm. FlashInfer has no compatible rowwise arm in the inspected backends; vLLM MoE uses its heuristic config.B200 decode graph / eager prefill: routed MoE 36.0 / 1188.8 µs; dense QKV 7.8 / 254.9 µs. Routed MoE is 1.92× / 1.88× faster than Torch and 1.60× / 2.02× faster than vLLM in those modes. A separate matched-input, normally tuned A/B measured about 6% lower full-MoE latency from grouping neighboring M tiles for weight reuse. These are synthetic operator measurements, not model throughput. Torch2.13+cu130, Triton3.7.1, vLLM0.29, MSLK1.3.
The refreshed plot includes the configured activation cap: Torch, MSLK and finegrained receive 1200; vLLM exposes no equivalent argument, but the benchmark values stay below the cap. All 48 timing cells completed. QKV graphed prefill ties Torch and trails vLLM; compiled MoE decode also trails native alternatives. Maximum cross-implementation relative maximum-element difference is 3.04% for MoE; each compiled arm passes its eager-versus-compiled check.
Tests extend existing Problem/MoEProblem and configuration-sweep tables, including compiled prequantized inputs and default output dtype. Ruff and diff checks passed. Companion #1237 improves generic tuning; no rowwise-specific tuning override is added.
Related issue
Requested by a maintainer for huggingface/transformers#49658.
Checklist