feat: NVFP4 standard format quantized inference with optimized bitwise decode kernels by AUXStar · Pull Request #11 · BlinkDL/Albatross

AUXStar · 2026-07-09T05:45:42Z

NVFP4 量化推理 + GEMV kernel 优化

变更概述

在 RWKV-7 v3a 推理引擎上实现 NVFP4 (E2M1 + E4M3 block scale + FP32 tensor scale) 4-bit 量化推理，并逐步优化 decode 性能。

优化历程

Phase 1-4: 推理路径融合

Phase 1: CUDA Graph 自动捕获 (B≤2)，kernel launch 5μs→1.5μs
Phase 2: r/k/v 三个 GEMV 融合为 1 个 kernel (launch 3→1)，p50 +22.8%
Phase 3: WKV+LNX 融合 (消除 HBM roundtrip)，p50 +5%
Phase 4: kk_a_gate 融入 WKV kernel (launch 3→1)
Phase 5: Grid sync megakernel 分析，因 block 数 > SM 数且数据依赖严格而不可行

Phase 6: NVFP4 GEMV kernel 优化

v2: shared memory LUT (byte→__half2) + __hfma2 块内 FP16 累加，指令数 112→18/block
OutTile 2→4: x 向量复用 4 次，block 数减半
废弃方案: dual-block ILP (寄存器溢出), OutTile=8 (寄存器溢出)

性能测试结果

硬件: RTX 5070 Ti Laptop (12GB GDDR7, 448 GB/s, sm_120)

模型	FP16	NVFP4	VRAM	B=1 tok/s	B=2 tok/s
2.9B	5.5GB	2.2GB	3.4GiB	174.8	130.0
7.2B	14GB	4.8GB	6.1GiB	98.4	72.0
13.3B	25GB	8.3GB	10.0GiB	63.4	39.2

13.3B kernel 优化前后对比 (B=1):

版本	ms/tok	tok/s	vs v1
v1 (原始)	19.35	51.7	baseline
v2 (LUT+hfma2)	17.50	57.1	+10.4%
v2+OutTile4	15.76	63.4	+22.6%

带宽利用率: 96.1% (HBM 8.35GB/tok, 理论 18.6ms)

正确性

所有模型 3 步 decode 无 NaN/Inf
NVFP4 量化 CosSim ≥ 0.9999 (与 FP16 对比)

新增文件

quantize_stream.py: mmap 流式量化工具 (13.3B 量化 171 秒，峰值 12.9GB RAM)

约束

head.weight 保持 FP16 (防止 RL 训练崩坏)
ffn.value.weight 走 NVFP4 cmix_sparse kernel
低秩权重 (w1/w2/a1/a2/g1/g2/v1/v2) 保持 FP16
推理路径根据权重 dtype 自动选择 NVFP4/FP16 kernel

详细报告

完整性能报告见 docs/features/nvfp4/performance_report.md

Add INT8 weight quantization support for RWKV-7 v3a inference engine: - New CUDA kernels: int8 GEMV (M=1/2), 4-wide K-loop GEMM (M>=3), vectorized dequant (8-wide int32 load + half2 write) - Offline quantization tool (per-channel symmetric INT8) - Dispatch optimization: group-aware INT8 threshold (att/head <=16, ffn_key <=12), RowTile/OutTile tuning (M<=4 -> (4,4), M>4 -> (8,2)) - 4-wide K loop reduces loop count by 50% for K%4==0 - Head dequant 13x faster, total dequant 30% faster B=128 throughput: 553 -> 979 tok/s (+77%) B=64 throughput: 852 -> 1013 tok/s (+19%) 1x4 prefill: 36% faster than 2-wide baseline

- Migrate from bitsandbytes NF4 to standard NVFP4 (E2M1 + E4M3 block scale + FP32 tensor scale) - Add e2m1_decode_f() and e2m1_decode_h() bitwise IEEE 754 decode (FP32/FP16) - Replace all constant-memory LUT lookups with bitwise decode to eliminate warp-level address divergence serialization bottleneck - Add optimized blk16 GEMV kernels (uint2 vectorized, 1-warp, K%16 aligned) for M=1 and M=2 decode paths - Update quantize_nf4.py to generate standard NVFP4 format - Update dispatch in rwkv7_fast_v3a.py to use optimized blk16 kernels Performance (2.9B model, RTX 5070 Ti Laptop): - T=1 decode: 71.69 -> 207 tok/s (2.9x) - T=16: 574 -> 983 tok/s (1.7x) - Kernel bandwidth: 80 -> 350 GB/s (4.4x) - VRAM: 3.40 GiB (unchanged) Performance (7.2B model): - T=1 decode: 106 tok/s - T=128: 797 tok/s - VRAM: 6.02 GiB (from 14GB original) - CosSim: ~0.9952 across all 192 quantized weights

This reverts commit 4adf2ea.

This reverts commit 4bbe65d.

This reverts commit c6a859e.

…speedup)

…ith baseline)

… count > SM limit, data dependencies block fusion)

- Replace 16x e2m1_decode_f (bitwise) + fmaf with shared mem byte→__half2 LUT - Use __hfma2 for FP16 block accumulation (8 instr instead of 16 fmaf) - Apply block scale once per block instead of per element - FP32 cross-block accumulation preserves precision - Both row1 and rkv fused kernels optimized 13.3B decode B=1: 19.35ms → 17.50ms (-9.6%, +10.6% tok/s) 51.7 → 57.1 tok/s

- Increase OutTile from 2 to 4: each block produces 4 outputs instead of 2 - x vector loaded once per block, reused 4x instead of 2x - Block count halved: 2048→1024 for N=4096 (row1), 6144→3072 (rkv) - OutTile=8 tested but causes register spilling (24ms vs 16ms) - OutTile=4 is sweet spot: 4 FP32 acc + 4×8 weight temps fit in registers 13.3B decode B=1: 17.50ms → 16.25ms (-7.1%, +7.7% tok/s) Combined v1→v2+OutTile4: 19.35ms → 16.25ms (-16.0%, +19.0% tok/s) 51.7 → 61.5 tok/s

- 2.9B: 174.8 tok/s (B=1), 130.0 tok/s (B=2) - 7.2B: 98.4 tok/s (B=1), 72.0 tok/s (B=2) - 13.3B: 63.4 tok/s (B=1), 39.2 tok/s (B=2) - All models pass correctness (no NaN/Inf) - 13.3B B=1: 51.7→63.4 tok/s (+22.6% vs v1)

Root cause: NVFP4 cmix kernel expects [F, C/2] layout but weights stay [C, F/2]. Fix: dequant -> transpose -> requant to NVFP4 [F, C/2]. Preserves NVFP4 format. Also fix dense path (T>19): dequant with transpose=False since weight is already [F, C/2]. Tested: 2.9B loss=1.79 VRAM=2.91GB, 7.2B loss=1.63 VRAM=7.09GB, 13.3B loss=1.54 VRAM=11.03GB

AUXStar · 2026-07-13T15:19:22Z

Fix: NVFP4 cmix kernel weight layout mismatch (transpose-requant)

Root Cause

NVFP4 cmix kernel (cmix_sparse_down_relu_one_nf4 etc.) accesses weight as w_nf4[actual_f * C2 + col2], expecting [F, C/2] layout. But NVFP4 (uint8) weights are excluded from the transpose step in weight preprocessing, keeping [C, F/2] layout. This mismatch causes complete model failure (loss ~11 vs expected ~1.3).

Fix

During weight preprocessing, for each NVFP4 ffn.value.weight:

Dequant [C, F/2] → FP16 [F, C] (transpose=True)
Re-quantize → NVFP4 [F, C/2] with new scales
Also fix dense path (T>19): use transpose=False since weight is already [F, C/2]

Preserves NVFP4 format, zero VRAM increase.

Test Results (all 3 model sizes, RTX 5070 Ti Laptop 12GB)

Model	Loss (avg)	Peak VRAM	12GB Fit	Generation
2.9B	1.79	2.91 GB	✅	✅ coherent
7.2B	1.63	7.09 GB	✅	✅ coherent
13.3B	1.54	11.03 GB	✅	✅ coherent

Generation sample (7.2B): "The meaning of life is to find your gift. The purpose of life is to give it away." – Pablo Picasso

Commit: 07428dd

AUXStar added 18 commits July 5, 2026 19:37

fix:

fe8c22b

docs:

4bbe65d

docs:

4adf2ea

Revert "docs:"

88dfde9

This reverts commit 4adf2ea.

Revert "docs:"

c6a859e

This reverts commit 4bbe65d.

Reapply "docs:"

4f98e4e

This reverts commit c6a859e.

Phase2:

6dd1383

Phase3: fused WKV+LNX kernel for T=1 decode (CosSim=0.99997, ~5% p50 …

43109cd

…speedup)

Phase4: fused kk_a_gate+WKV+LNX kernel (CosSim=0.99998, perf parity w…

fb8da29

…ith baseline)

Phase5: feasibility analysis - grid_sync megakernel not viable (block…

6f7281e

… count > SM limit, data dependencies block fusion)

backup before nvfp4 kernel optimization

6d0b254

docs: update performance report with full regression test results

4c06c27

- 2.9B: 174.8 tok/s (B=1), 130.0 tok/s (B=2) - 7.2B: 98.4 tok/s (B=1), 72.0 tok/s (B=2) - 13.3B: 63.4 tok/s (B=1), 39.2 tok/s (B=2) - All models pass correctness (no NaN/Inf) - 13.3B B=1: 51.7→63.4 tok/s (+22.6% vs v1)

Merge remote-tracking branch 'origin/main' into feat/nvfp4-quantization

d954f25

Navigation Menu

feat: NVFP4 standard format quantized inference with optimized bitwise decode kernels - #11

feat: NVFP4 standard format quantized inference with optimized bitwise decode kernels#11
AUXStar wants to merge 18 commits into
BlinkDL:mainfrom
AUXStar:feat/nvfp4-quantization

AUXStar commented Jul 9, 2026 •

edited

Loading

Uh oh!

AUXStar commented Jul 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

AUXStar commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

NVFP4 量化推理 + GEMV kernel 优化

变更概述

优化历程

Phase 1-4: 推理路径融合

Phase 6: NVFP4 GEMV kernel 优化

性能测试结果

正确性

新增文件

约束

详细报告

Uh oh!

AUXStar commented Jul 13, 2026

Fix: NVFP4 cmix kernel weight layout mismatch (transpose-requant)

Root Cause

Fix

Test Results (all 3 model sizes, RTX 5070 Ti Laptop 12GB)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

AUXStar commented Jul 9, 2026 •

edited

Loading