feat: add FP8 blockwise dequantize/GEMM operators and FP8(E4M3) KV cache support in paged attention - #1568
Open
shsaihdsaiudh wants to merge 11 commits into
Open
feat: add FP8 blockwise dequantize/GEMM operators and FP8(E4M3) KV cache support in paged attention#1568shsaihdsaiudh wants to merge 11 commits into
shsaihdsaiudh wants to merge 11 commits into
Conversation
Elementwise dequantization for block-quantized FP8 (E4M3) weights: out[i,j] = fp8_e4m3_decode(q[i,j]) * scales[i/BM, j/BN], with the block shape (BM, BN) inferred from tensor shapes (e.g. 128x128 for DeepSeek/Qwen3-FP8 checkpoints). q is INFINI_DTYPE_F8, scales F32, output F16/BF16/F32; no workspace. Backends: cpu (portable bit-manipulation decode) and nvidia (reuses infiniopFp8E4m3Decode). infinicore C++ op layer included for InfiniLM consumption. infiniop tests cover 5 block shapes x 3 output dtypes, passing on cpu and nvidia (sm_89, sm_120).
- xmake/nvidia.lua: exclude avg_pool3d nvidia .cu when cudnn=n (its implementation requires cuDNN; mixing PyTorch's bundled cuDNN with the system CUDA runtime breaks the FA-enabled build) - avg_pool3d/operator.cc: gate NVIDIA/QY branches on ENABLE_CUDNN_API - xmake.lua: add sm_120 to cuda_arch values (RTX 5060 Ti)
out[m,n] = sum_k a[m,k] * fp8_e4m3_decode(q[n,k]) * scales[n/BN, k/BK],
computed without materializing the dequantized weight (decode-oriented;
InfiniLM routes M<=8 here). CPU reference backend + full infiniop /
infinicore / ctypes surface.
nvidia kernel: warp-per-row GEMV (4 warps/block, grid N/4), each lane
holds 16B of FP8 with next-group register prefetch, per-128-chunk scale
application, FP32 accumulate with warp shuffle reduce; M_TILE in
{1,2,4,8,16} templates. test/infiniop/fp8_blockwise_gemm.py passes on
sm_120 for 7 shapes x F16/BF16/F32 (incl. block 64x128 / 128x256 and
K-tail cases).
Measured on RTX 5090 with Qwen3-8B-FP8 in InfiniLM: decode bs=1 goes
142.96 tok/s vs 92.34 for the BF16 baseline (1.55x) and 28.49 for the
naive dequantize+GEMM path (5.02x); op-level streaming bandwidth
1338-1731 GB/s on the model shapes.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ends Extend the ascend/bang/metax/moore descriptor create/calculate signatures of paged_caching, paged_attention and paged_attention_prefill with the k_scale/v_scale parameters added for FP8(E4M3) KV caches, forwarding the descriptors to the shared info validation. These backends do not implement the FP8 path, so create explicitly returns INFINI_STATUS_NOT_IMPLEMENTED for F8 caches instead of silently misinterpreting the data.
The FP8 KV decode kernel assigns one CTA per (sequence, q_head), so at bs=1 only 32 CTAs (Qwen3-8B) occupy ~19% of the SMs and each CTA scans the whole token range serially; the FP8/BF16 overhead grows with context length (1.05x @ 1k -> 1.27x @ 16k). Shard the token dimension across up to 8 CTAs per (seq, head), each emitting a partial (max, sum_exp, weighted sum) triple, and merge across CTAs in a second online-softmax reduction kernel (log2-domain merge, standard flash-decoding). - kernel_fp8.cuh: partial_acc/m/l outputs + num_splits, token shard windows, -inf guard for empty shards, split-kv combine kernel - paged_attention_fp8.cu: two-stage launch, auto heuristic, INFINIOP_FLASH_DECODE_SPLITKV/NUM_SPLITS/DEBUG_SPLITS env overrides - paged_attention_nvidia.cu: FP8 workspace 0 -> splits*seqs*heads* (value_size+2)*4B - test: (1,8,8,128,16,4096,True) alibi+split case Validated with an exact Python simulator only (dev_fp8/results/ fp8_splitkv_sim.py, ALL PASS); GPU build/bench pending (server down). See dev_fp8/w7_kv_fp8_splitkv_report.md in the InfiniLM repo.
The SIMT warp-per-row kernel re-reads each activation row per weight
element and goes instruction-throughput bound once M grows (W4: fused
loses to dequant+cuBLAS from M >= 16). Add an mma.m16n8k16.row.col
kernel for the decode range 9 <= M <= 32 with F16/BF16 activations:
- CTA = 4 warps x n8 (N_TILE 32), M_BLOCKS in {1, 2} covering M <= 32
- K streamed in 128-wide chunks (exactly one scale sub-chunk; the shape
constraints are already enforced by Fp8BlockwiseGemmInfo), global ->
register staging -> double-buffered smem, one __syncthreads per chunk
- per-chunk mma partials promoted with the block scale into an FP32
accumulator (c_fin = fmaf(s, c_part, c_fin)), same numeric structure
as the SIMT path
- FP8 e4m3 decoded in registers with a bit-placement trick giving the
true value times 2^-120 (BF16) / 2^-8 (F16); the power-of-two factor
folds into the promote scale, so mma inputs are bit-exact with no
per-element multiply
- dispatch: F16/BF16 and 9 <= M <= 32 -> mma; INFINIOP_FP8_GEMM_MMA=0
forces the SIMT kernels (A/B); F32 stays on SIMT M_TILE
Validated with an exact Python simulator (10 cases: M tails, N tails,
block_n/block_k variants, M_BLOCKS 1/2; ALL PASS) plus static review of
fragment layouts against the PTX ISA; GPU build/bench pending (server
down). See dev_fp8/w8_mma_report.md in the InfiniLM repo.
The W8 mma commit extracted the SIMT dispatch into launch_mtile but dropped its int M_TILE template parameter (all call sites already pass two arguments), breaking the GPU build with 'identifier M_TILE is undefined'. Found during the first server build of the mma path.
The device-link step of -rdc=true targets did not see the -gencode flags that were only added to cuflags; nvcc >= 13 device-links for sm_75 by default in that case, producing binaries without SASS for the actual GPU (every kernel launch fails with cudaErrorNoKernelImageForDevice). - apply_cuda_arch_flags() now registers cugencodes (compile + devlink) - fix the nvidia-smi auto-detect: os.iorunv returns (stdout, stderr); the version match ran against stderr and always fell back to native. Auto-detect now builds for the detected GPU's own arch (sm_90 -> 90a) - infinirt-nvidia gains the same on_load arch handling - drop the unused boost add_requires (nothing in src/ references it)
The decode kernels returned early for seq_len == 0 without writing anything: the direct path left the caller's output buffer untouched, and the split-kv path never published partials, so the combine kernel merged stale workspace from previous calls (reproduced on sm_120: the len-0 row came back bit-identical to the previous call's row 0). - direct kernels (v1, v2 warp/CTA/pipelined/GQA, MLA, FP8) now write zeros for len-0 rows - split kernels (v2 warp/CTA, MLA, FP8) publish the neutral element (m=-inf, l=0, acc=0) so combine emits a zero row - v2/MLA combine kernels guard exp2f(ms - m_total) against ms=-inf so an all-empty row cannot produce NaN
The flash path only fell back when scale tensors were present; an F8 cache without scales slipped into a path that cannot quantize on write and now fails loudly (Bad Parameter) instead of silently corrupting.
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.
Summary
fp8_blockwise_dequantize: FP8(E4M3) weight[M,N]+ F32 block scales[ceil(M/128), ceil(N/128)]-> F16/BF16/F32. Elementwise, no workspace. Backends: NVIDIA CUDA kernel + CPU reference. Files:include/infiniop/ops/fp8_blockwise_dequantize.h,src/infiniop/ops/fp8_blockwise_dequantize/,include/infinicore/ops/fp8_blockwise_dequantize.hpp,src/infinicore/ops/fp8_blockwise_dequantize/.fp8_blockwise_gemm: fused FP8 blockwise GEMM,out[m,n] = sum_k a[m,k] * e4m3_decode(q[n,k]) * scales[n/128, k/128], without materializing a BF16 weight matrix. Backends: NVIDIA + CPU reference. Three NVIDIA dispatch paths:M_TILE in {1,2,4,8}template, forM <= 8(decode/GEMV-shaped; FP32 accumulate, per-128-K-block scale application).mma.m16n8k16.row.col, F16/BF16) for9 <= M <= 32: 4 warps x n8 = N_TILE 32,M_BLOCKS in {1,2}, K pipelined in 128-chunks through double-buffered smem, per-chunk mma partials promoted with the block scale in FP32. E4M3 code bits are placed directly into BF16/F16 bit patterns (decode = true value x 2^-120 / 2^-8, folded into the scale), so mma inputs are bit-exact with no per-element multiply.INFINIOP_FP8_GEMM_MMA=0forces SIMT.M_TILE=16fallback otherwise (F32, orM > 32).K % 128 == 0,block_k % 128 == 0,block_n % 16 == 0) are enforced inFp8BlockwiseGemmInfo::create, so a bad shape fails at descriptor creation instead of silently misrouting.paged_caching,paged_attention,paged_attention_prefill):k_scale/v_scaledescriptors and data pointers (NULL = absent; existing callers unchanged).info.henforces: cache dtype F8 => scales required, F32, shape[num_blocks, num_kv_heads, block_size]; cache dtype != F8 => scales must be NULL. BF16/F16 paths are bit-identical to before.scale = amax(|x[0:head_dim]|) / 448(amax == 0 => scale = 1); K stored post-RoPE; decode/encode reuse the existing portableinfiniopFp8E4m3Encode/Decode.paged_caching: new quantization write kernel (grid over (token, kv_head), per-block amax reduction -> scale write -> encode).paged_attention(decode): new clean-room FP8 kernel family (cuda/kernel_fp8.cuh,nvidia/paged_attention_fp8.cu), CTA per (seq, q_head), 32 warps, dequant-on-load, log2-domain online softmax, hd 64/128. Includes split-kv flash-decoding: token dimension sharded across CTAs (grid z-dim, up to 8 splits), partials combined by a log2-domain combine kernel; defaultautoFA2-style waves heuristic (no split when the grid is already large). Env overrides:INFINIOP_FLASH_DECODE_SPLITKV=0/1/auto,INFINIOP_FLASH_NUM_SPLITS=1..8/auto,INFINIOP_FLASH_DEBUG_SPLITS=1.paged_attention_prefill: gather-dequant of referenced pages to BF16 scratch (identity block_tables, scratch via workspace), reusing the existing BF16 prefill kernel unchanged.NOT_IMPLEMENTEDon these backends.int M_TILEtemplate parameter onlaunch_mtile(build break when extracted to a function).cugencodes(not rawcuflags) so the device-link step inherits the arch — nvcc >= 13 otherwise device-links for sm_75 by default and produces binaries without SASS for the actual GPU (xmake/nvidia.lua).cuda/kernel_v2.cuh).paged_cachingdispatch (paged_caching_infiniops.cc).cudnn=n(exclude cuDNN-dependentavg_pool3dsources); addsm_120to--cuda_arch.test/infiniop/fp8_blockwise_dequantize.py,test/infiniop/fp8_blockwise_gemm.py; FP8 cases added totest/infiniop/paged_attention.py,paged_attention_prefill.py,paged_caching.py,paged_caching_prefill.py; ctypes signatures registered intest/infiniop/libinfiniop/op_register.py.Motivation
quant/dequantare INT8-only,scaled_mmis INT8 (per-row x per-col F32 scale), and the Marlin port's FP8 support (kFE4M3fn) is per-tensor/per-channel scale — and the Marlin kernel itself is broken on sm_120 for all bit widths (deadlock under CUDA 12.8, wrong output under CUDA 13.2; documented with minimal repro in the downstream repo'sdev_fp8/marlin_sm120_issue.md).weight_block_size: [128,128]), and is the target format of the downstream InfiniLM integration (feat: add FP8 blockwise weight quantization and FP8(E4M3) KV cache support InfiniLM#586).Tests and reproduction
Operator tests (all run on RTX 5090, sm_120,
--nvidia; torch CPU/numpy references):fp8_blockwise_dequantize.pyfp8_blockwise_gemm.pyINFINIOP_FP8_GEMM_MMA=1/0(mma path and SIMT A/B both green)paged_caching.py(FP8)paged_attention.py(1, 8, 8, 128, 16, 4096, alibi=True)triggers num_splits=8 on GPU, large-grid cases correctly stay at num_splits=1paged_attention_prefill.py(FP8)paged_caching_prefill.pyReproduce:
Additional pre-GPU verification artifacts (in downstream repo
dev_fp8/results/):fp8_mma_sim.py(bit-trick exact over all 256 non-NaN E4M3 codes, BF16/F16; 10 simulated GEMM cases vs float64 reference, ALL PASS) andfp8_splitkv_sim.py(split/combine vs direct softmax, rel err <= 1e-6 over 300 random cases incl. empty shards). Both were later confirmed on GPU by the operator tests above.Results
Operator-level performance (NVIDIA, RTX 5090):
fp8_blockwise_gemmSIMT kernel: 1317 GB/s at 4096x4096 — 88% of the 1505 GB/s pure-stream ceiling measured on the same card (HBM-resident rotation benchmark; earlier L2-contaminated microbenchmarks were discarded, see downstreamdev_fp8/w4_fused_gemm_report.md).End-to-end results via the downstream InfiniLM integration (InfiniTensor/InfiniLM#586), Qwen3-8B-FP8,
examples/bench.py, paged attention:E2E correctness (downstream): fused-vs-naive greedy generation token-identical (3 prompts x 64 tokens); FP8 KV vs BF16 KV greedy token-identical (4020-token prompt + 48 tokens, split-kv active).
Platform support status
fp8_blockwise_dequantizeandfp8_blockwise_gemmincluded; a standalone--cpubackend run was not performed (CPU code serves as the reference/fallback path).k_scale/v_scalethrough; F8 cache explicitly returnsNOT_IMPLEMENTED. Compile-level change only — no hardware access, untested. FP8 weight/KV features remain NVIDIA-only.Notes
main@9855c0c6 (2026-08-25, the tip at development time). Per course instruction the branch is intentionally not rebased across the recent operator-layer refactor onmain; happy to rebase if maintainers prefer.INFINIOP_FLASH_DECODE_SPLITKV=0avoids it.