From 55dbfa8edc39ca28632bc24959a6e4ab0e59d24d Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 12 Aug 2026 09:45:39 -0700 Subject: [PATCH 1/7] feat(norm): fused Add+RMSNorm+1x128 fp8 block-quant producer Add FusedAddRMSNormFP8BlockQuant, the block-scaled sibling of FusedAddRMSNormQuant: instead of a scalar per-tensor scale it emits a dynamic per-1x128-block fp32 scale in the column-major TMA-aligned (hidden_size/128, round_up(batch, 4)) layout that the fp8 block-scaled GEMMs (flashinfer/deep_gemm.py check_sf_layout) consume. This fuses the 2-kernel producer (fused_add_rmsnorm -> a separate per-token-group 1x128 fp8 quant) into a single memory pass, emitting: fp8 activation, the fp32 block scale, the pre-norm residual (in place), and the bf16 normed output. One CTA per token row, single-wave (num_threads = d/VEC_SIZE, VEC8 for d<=8192 / VEC16 to 16384); 128/VEC_SIZE consecutive warp-local lanes own one 1x128 block so its amax is a register-only warp-subgroup shfl. Mirrors FusedAddRMSNormQuant's PDL + cudaLaunchKernelEx. fp8 e4m3 (block-scale is e4m3); the fp8 is quantized from the bf16-rounded normed so it matches normed_out exactly. Memory-optimal (~90% of H200 HBM peak): one launch instead of two, ~2x less memory traffic than the 2-kernel path. Adds the csrc DLPack wrapper + TVM-FFI export, the public flashinfer.fused_add_rmsnorm_fp8_block_quant op, and a benchmark with a torch-reference correctness check. Co-Authored-By: Claude Opus 4.8 --- ...bench_fused_add_rmsnorm_fp8_block_quant.py | 163 ++++++++++++++++++ csrc/flashinfer_norm_binding.cu | 5 + csrc/norm.cu | 51 ++++++ flashinfer/__init__.py | 3 + flashinfer/norm/__init__.py | 68 ++++++++ include/flashinfer/norm.cuh | 133 ++++++++++++++ 6 files changed, 423 insertions(+) create mode 100644 benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py diff --git a/benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py b/benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py new file mode 100644 index 00000000000..228e5ce0cda --- /dev/null +++ b/benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py @@ -0,0 +1,163 @@ +"""Benchmark: fused Add+RMSNorm+1x128-fp8-block-quant (one kernel) vs the 2-kernel producer +(flashinfer.fused_add_rmsnorm then a per-token-group 1x128 fp8 quant), which is what runs before +fp8 block-scaled GEMMs in deepseek_v2-family / MoE models. + +Shows the kernel-level win: the fused kernel does add+norm+quant in a single memory pass, saving +the intermediate bf16 round-trip. +""" + +import numpy as np +import torch + +import flashinfer + + +def _round_up(x, m): + return (x + m - 1) // m * m + + +def _run_fused( + out, block_scale, normed_out, input, residual, weight, eps, enable_pdl=False +): + """Call the fused op via the public API if present, else the JIT module export directly.""" + if hasattr(flashinfer, "fused_add_rmsnorm_fp8_block_quant"): + flashinfer.fused_add_rmsnorm_fp8_block_quant( + out, + block_scale, + normed_out, + input, + residual, + weight, + eps, + enable_pdl=enable_pdl, + ) + else: + flashinfer.norm.get_norm_module().fused_add_rmsnorm_fp8_block_quant( + out, block_scale, normed_out, input, residual, weight, eps, enable_pdl + ) + + +def _ref_group_quant_fp8(normed, block=128): + """Reference 1x128 per-token-group fp8 quant (torch) -> (fp8, scale[M, H/128]).""" + M, H = normed.shape + x = normed.float().view(M, H // block, block) + amax = x.abs().amax(dim=-1).clamp_min(1e-4) + scale = amax / 448.0 + q = (x / scale.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn) + return q.view(M, H), scale + + +def bench_one(M, H, dtype=torch.bfloat16, iters=100, warmup=20): + dev = "cuda" + input = torch.randn(M, H, dtype=dtype, device=dev) * 0.1 + residual = torch.randn(M, H, dtype=dtype, device=dev) * 0.1 + weight = torch.randn(H, dtype=dtype, device=dev) + eps = 1e-6 + + # fused outputs + out = torch.empty(M, H, dtype=torch.float8_e4m3fn, device=dev) + block_scale = torch.empty( + H // 128, _round_up(M, 4), dtype=torch.float32, device=dev + ) + normed_out = torch.empty(M, H, dtype=dtype, device=dev) + + # try to use a production per-token-group quant for the baseline's 2nd kernel; else torch ref. + try: + from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8 + + def baseline_quant(n): + return per_token_group_quant_fp8(n, 128) + except Exception: + + def baseline_quant(n): + return _ref_group_quant_fp8(n) + + # Timing loops run the kernels in place (no setup clones); residual grows ~linearly (stays bounded), + # and memory traffic per call is identical regardless of values -> fair kernel-time comparison. + + def fused(): + _run_fused(out, block_scale, normed_out, input, residual, weight, eps) + + def baseline(): + flashinfer.fused_add_rmsnorm( + input, residual, weight, eps + ) # residual <- add+norm (bf16), in place + baseline_quant(residual) # 2nd kernel: 1x128 fp8 quant of the bf16 normed + + def time_it(fn): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + s, e = torch.cuda.Event(True), torch.cuda.Event(True) + ts = [] + for _ in range(iters): + s.record() + fn() + e.record() + torch.cuda.synchronize() + ts.append(s.elapsed_time(e)) + return float(np.median(ts)) # ms + + t_f = time_it(fused) + t_b = time_it(baseline) + # fused byte model: read input+residual+weight (bf16), write residual+normed(bf16)+fp8+fp32 scale + bytes_f = ( + (3 * M * H * 2) + + (2 * M * H * 2) + + (M * H * 1) + + (H // 128 * _round_up(M, 4) * 4) + ) + gbps = bytes_f / (t_f * 1e-3) / 1e9 + print( + f"M={M:>5} H={H:>5} | fused {t_f * 1e3:7.1f}us 2-kernel {t_b * 1e3:7.1f}us " + f"speedup {t_b / t_f:4.2f}x | fused {gbps:6.0f} GB/s" + ) + return t_b / t_f + + +def check_correctness(M=99, H=4096, dtype=torch.bfloat16): + dev = "cuda" + input = torch.randn(M, H, dtype=dtype, device=dev) * 0.1 + residual = torch.randn(M, H, dtype=dtype, device=dev) * 0.1 + weight = torch.randn(H, dtype=dtype, device=dev) + eps = 1e-6 + # reference + x = input.float() + residual.float() + var = x.pow(2).mean(-1, keepdim=True) + normed_ref = (x * torch.rsqrt(var + eps) * weight.float()).to( + dtype + ) # round to bf16 like kernel + q_ref, scale_ref = _ref_group_quant_fp8(normed_ref) # scale_ref [M, H/128] + # fused + out = torch.empty(M, H, dtype=torch.float8_e4m3fn, device=dev) + block_scale = torch.empty( + H // 128, _round_up(M, 4), dtype=torch.float32, device=dev + ) + normed_out = torch.empty(M, H, dtype=dtype, device=dev) + r = residual.clone() + _run_fused(out, block_scale, normed_out, input, r, weight, eps) + # checks + torch.testing.assert_close(r, x.to(dtype), rtol=1e-2, atol=1e-2) # residual + torch.testing.assert_close( + normed_out.float(), normed_ref.float(), rtol=1e-2, atol=1e-2 + ) # normed + scale_got = block_scale.transpose(0, 1)[:M] # -> [M, H/128] + torch.testing.assert_close( + scale_got, scale_ref, rtol=1e-3, atol=1e-6 + ) # block scale + # dequant fp8 vs reference normed + deq = out.float().view(M, H // 128, 128) * scale_got.unsqueeze(-1) + torch.testing.assert_close(deq.view(M, H), normed_ref.float(), rtol=0.1, atol=0.3) + print("correctness: OK (residual, normed_out, block_scale, dequant all match)") + + +if __name__ == "__main__": + print(f"flashinfer {flashinfer.__version__} | {torch.cuda.get_device_name()}") + check_correctness() + speedups = [] + for M in (256, 1024, 4096): + for H in (4096, 6144, 8192): + speedups.append(bench_one(M, H)) + print( + f"\ngeomean speedup (fused vs 2-kernel): {np.exp(np.mean(np.log(speedups))):.2f}x" + ) diff --git a/csrc/flashinfer_norm_binding.cu b/csrc/flashinfer_norm_binding.cu index cb7f0c3c694..3163ec5c890 100644 --- a/csrc/flashinfer_norm_binding.cu +++ b/csrc/flashinfer_norm_binding.cu @@ -26,6 +26,10 @@ void fused_add_rmsnorm(TensorView input, TensorView residual, TensorView weight, void fused_add_rmsnorm_quant(TensorView output, TensorView input, TensorView residual, TensorView weight, TensorView scale, double eps, bool enable_pdl); +void fused_add_rmsnorm_fp8_block_quant(TensorView output, TensorView block_scale, + TensorView normed_out, TensorView input, TensorView residual, + TensorView weight, double eps, bool enable_pdl); + void gemma_rmsnorm(TensorView out, TensorView input, TensorView weight, double eps, bool enable_pdl); @@ -58,6 +62,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm, rmsnorm); TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm_quant, rmsnorm_quant); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm, fused_add_rmsnorm); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm_quant, fused_add_rmsnorm_quant); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm_fp8_block_quant, fused_add_rmsnorm_fp8_block_quant); TVM_FFI_DLL_EXPORT_TYPED_FUNC(gemma_rmsnorm, gemma_rmsnorm); TVM_FFI_DLL_EXPORT_TYPED_FUNC(gemma_fused_add_rmsnorm, gemma_fused_add_rmsnorm); TVM_FFI_DLL_EXPORT_TYPED_FUNC(layernorm, layernorm); diff --git a/csrc/norm.cu b/csrc/norm.cu index 8df877b2551..e0a16816845 100644 --- a/csrc/norm.cu +++ b/csrc/norm.cu @@ -186,6 +186,57 @@ void fused_add_rmsnorm_quant(TensorView output, TensorView input, TensorView res }); } +// Add-residual + RMSNorm + 1x128 fp8 block-quant. Emits fp8 `output` [M,d], fp32 `block_scale` +// (d/128, round_up(M,4)) column-major TMA layout, bf16 `normed_out` [M,d], and updates `residual` +// in place with (input+residual). See norm::FusedAddRMSNormFP8BlockQuant. +void fused_add_rmsnorm_fp8_block_quant(TensorView output, TensorView block_scale, + TensorView normed_out, TensorView input, TensorView residual, + TensorView weight, double eps, bool enable_pdl) { + CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(residual); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(weight); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(output); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(normed_out); + CHECK_DEVICE(input, residual); + CHECK_DEVICE(input, weight); + CHECK_DEVICE(input, output); + CHECK_DEVICE(input, block_scale); + CHECK_DEVICE(input, normed_out); + CHECK_DIM(2, input); // (batch_size, hidden_size) + CHECK_DIM(2, residual); // (batch_size, hidden_size) + CHECK_DIM(1, weight); // (hidden_size) + CHECK_DIM(2, output); // fp8 (batch_size, hidden_size) + CHECK_DIM(2, normed_out); + CHECK_DIM(2, block_scale); // (hidden_size/128, round_up(batch_size, 4)) + unsigned int batch_size = input.size(0); + unsigned int hidden_size = input.size(1); + unsigned int m_pad = (batch_size + 3) & ~3u; + TVM_FFI_ICHECK_EQ(residual.size(0), batch_size); + TVM_FFI_ICHECK_EQ(residual.size(1), hidden_size); + TVM_FFI_ICHECK_EQ(weight.size(0), hidden_size); + TVM_FFI_ICHECK_EQ(normed_out.size(0), batch_size); + TVM_FFI_ICHECK_EQ(normed_out.size(1), hidden_size); + TVM_FFI_ICHECK_EQ(block_scale.size(0), hidden_size / 128); + TVM_FFI_ICHECK_EQ(block_scale.size(1), m_pad); + TVM_FFI_ICHECK_EQ(hidden_size % 128, 0); + ffi::CUDADeviceGuard device_guard(input.device().device_id); + const cudaStream_t stream = get_stream(input.device()); + + // 1x128 block-scale activations are e4m3 (the layout DeepGEMM/flashinfer fp8 block GEMMs + // consume). + DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(input.dtype(), c_type, [&] { + cudaError_t status = norm::FusedAddRMSNormFP8BlockQuant( + static_cast(input.data_ptr()), static_cast(residual.data_ptr()), + static_cast(weight.data_ptr()), static_cast<__nv_fp8_e4m3*>(output.data_ptr()), + static_cast(block_scale.data_ptr()), static_cast(normed_out.data_ptr()), + batch_size, hidden_size, input.stride(0), residual.stride(0), output.stride(0), + block_scale.stride(0), eps, enable_pdl, stream); + TVM_FFI_ICHECK(status == cudaSuccess) + << "FusedAddRMSNormFP8BlockQuant failed with error code " << cudaGetErrorString(status); + return true; + }); +} + void gemma_rmsnorm(TensorView output, TensorView input, TensorView weight, double eps, bool enable_pdl) { CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 1948950f99b..c66975b43d8 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -135,6 +135,9 @@ from . import msa_ops as msa_ops from .norm import fused_add_rmsnorm as fused_add_rmsnorm from .norm import fused_add_rmsnorm_quant as fused_add_rmsnorm_quant +from .norm import ( + fused_add_rmsnorm_fp8_block_quant as fused_add_rmsnorm_fp8_block_quant, +) from .norm import layernorm as layernorm from .norm import layernorm_quant as layernorm_quant from .norm import gemma_fused_add_rmsnorm as gemma_fused_add_rmsnorm diff --git a/flashinfer/norm/__init__.py b/flashinfer/norm/__init__.py index 6efd0dbdea5..0188c710b12 100644 --- a/flashinfer/norm/__init__.py +++ b/flashinfer/norm/__init__.py @@ -389,6 +389,74 @@ def _fused_add_rmsnorm_quant_fake( pass +@register_custom_op( + "flashinfer::fused_add_rmsnorm_fp8_block_quant", + mutates_args=("out", "block_scale", "normed_out", "residual"), +) +def fused_add_rmsnorm_fp8_block_quant( + out: torch.Tensor, + block_scale: torch.Tensor, + normed_out: torch.Tensor, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, +) -> None: + r"""Fused add-residual + RMSNorm + 1x128 fp8 block quantization, in one pass. + + Replaces the two-kernel producer (:func:`fused_add_rmsnorm` then a per-token-group 1x128 fp8 + quant) used before fp8 block-scaled GEMMs. Unlike :func:`fused_add_rmsnorm_quant` (scalar + per-tensor scale) this emits a **dynamic per-1x128-block** fp32 scale in the column-major + TMA-aligned layout that :mod:`flashinfer.deep_gemm` consumes. + + Step 1: ``residual = input + residual`` (pre-norm, written in place). + Step 2: ``normed = RMSNorm(residual) * weight`` (bf16, written to ``normed_out``). + Step 3: ``out, block_scale = quant_1x128_fp8(normed)`` (fp8 from the bf16-rounded normed, so + ``out`` matches re-quantizing ``normed_out``). + + Parameters + ---------- + out : torch.Tensor + fp8 output, shape ``(batch_size, hidden_size)``, dtype float8_e4m3fn / float8_e5m2. + block_scale : torch.Tensor + fp32 block scales, shape ``(hidden_size // 128, round_up(batch_size, 4))``, contiguous. + This is the column-major (MN-major, TMA-aligned) buffer deep_gemm expects: the logical + ``(batch_size, hidden_size // 128)`` scale is ``block_scale.transpose(0, 1)[:batch_size]``. + normed_out : torch.Tensor + bf16/fp16 normed output, shape ``(batch_size, hidden_size)`` (for consumers that need the + pre-quant activation, e.g. an MoE router). + input, residual : torch.Tensor + Shape ``(batch_size, hidden_size)``. ``residual`` is updated in place with ``input+residual``. + weight : torch.Tensor + RMSNorm weight, shape ``(hidden_size,)``. ``hidden_size`` must be a multiple of 128 + (and of 256 for hidden_size<=8192, 512 for 8192 None: + pass + + @flashinfer_api(trace=gemma_rmsnorm_trace) def gemma_rmsnorm( input: torch.Tensor, diff --git a/include/flashinfer/norm.cuh b/include/flashinfer/norm.cuh index 298f4fe40cf..7c8a17673b1 100644 --- a/include/flashinfer/norm.cuh +++ b/include/flashinfer/norm.cuh @@ -675,6 +675,139 @@ cudaError_t FusedAddRMSNormQuant(T* input, T* residual, T* weight, O* output, ui return cudaSuccess; } +// Fused Add-residual + RMSNorm + 1x128 fp8 block-quant producer. The block-scaled sibling of +// FusedAddRMSNormQuant: instead of a scalar per-tensor scale it emits one fp32 scale per 1x128 +// block, in the column-major TMA-aligned (d/128, round_up(M,4)) layout that flashinfer's fp8 +// block-scaled GEMMs consume (flashinfer/deep_gemm.py check_sf_layout: stride(-2)==1, +// stride(-1)==round_up(M,4)). One CTA per token row; num_threads = d/VEC_SIZE cover the row in a +// single wave (host picks VEC_SIZE so num_threads<=1024). kLanesPerBlock = 128/VEC_SIZE +// consecutive, warp-local threads own one 1x128 block, so its amax is a register-only warp-subgroup +// reduction. Requires d % (32*VEC_SIZE) == 0 and d % 128 == 0. Emits: fp8 output[M,d]; fp32 +// block_scale; bf16 residual (hidden+residual, pre-norm, in place); bf16 normed_out. The fp8 is +// quantized from the T-rounded normed so it exactly matches normed_out (a bf16 consumer's re-quant +// is identical). +template +__global__ void FusedAddRMSNormFP8BlockQuantKernel( + T* __restrict__ input, T* __restrict__ residual, T* __restrict__ weight, O* __restrict__ output, + float* __restrict__ block_scale, T* __restrict__ normed_out, const uint32_t d, + const uint32_t stride_input, const uint32_t stride_residual, const uint32_t stride_output, + const uint32_t scale_stride_m, float eps) { + constexpr uint32_t kQuantBlock = 128; + constexpr uint32_t kLanesPerBlock = kQuantBlock / VEC_SIZE; // threads covering one 1x128 block + const uint32_t bx = blockIdx.x; // token row m + const uint32_t tx = threadIdx.x, ty = threadIdx.y; + constexpr uint32_t warp_size = 32; + const uint32_t num_warps = blockDim.y; + const uint32_t thread_id = tx + ty * warp_size; + const uint32_t col = thread_id * VEC_SIZE; +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + // Pass 1: x = input + residual; write pre-norm residual in place; accumulate sum of squares. + vec_t input_vec, residual_vec, weight_vec; + input_vec.load(input + bx * stride_input + col); + residual_vec.load(residual + bx * stride_residual + col); + weight_vec.load(weight + col); + float x[VEC_SIZE]; + float sum_sq = 0.f; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + x[j] = float(input_vec[j]) + float(residual_vec[j]); + sum_sq += x[j] * x[j]; + residual_vec[j] = (T)x[j]; + } + residual_vec.store(residual + bx * stride_residual + col); + + // Cross-CTA sum reduction (warp shfl -> smem -> first-warp shfl), then rms reciprocal. + extern __shared__ float smem[]; +#pragma unroll + for (uint32_t offset = warp_size / 2; offset > 0; offset /= 2) + sum_sq += math::shfl_xor_sync(sum_sq, offset); + smem[ty] = sum_sq; + __syncthreads(); + if (ty == 0) { + float s = (tx < num_warps) ? smem[tx] : 0.f; +#pragma unroll + for (uint32_t offset = warp_size / 2; offset > 0; offset /= 2) + s += math::shfl_xor_sync(s, offset); + smem[0] = s; + } + __syncthreads(); + const float rms_rcp = math::rsqrt(smem[0] / float(d) + eps); + constexpr float max_val = 448.0f; // fp8 e4m3 amax (1x128 block-scale activations are e4m3) + + // Pass 2: normed = round_T(x * rms * weight); write bf16 normed_out; per-128-block amax. + float normed[VEC_SIZE]; + vec_t normed_vec; + float lamax = 0.f; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + T nb = (T)(x[j] * rms_rcp * float(weight_vec[j])); // round to T before amax/quant + normed[j] = float(nb); + normed_vec[j] = nb; + lamax = fmaxf(lamax, fabsf(normed[j])); + } + normed_vec.store(normed_out + bx * stride_output + col); + + // amax over the 1x128 block: kLanesPerBlock consecutive lanes reduce within their warp subgroup. +#pragma unroll + for (uint32_t offset = kLanesPerBlock / 2; offset > 0; offset /= 2) + lamax = fmaxf(lamax, math::shfl_xor_sync(lamax, offset)); + const float amax = fmaxf(lamax, 1e-4f); + if ((thread_id % kLanesPerBlock) == 0) // one scale per 1x128 block, column-major (TMA) layout + block_scale[(col / kQuantBlock) * scale_stride_m + bx] = amax / max_val; + + // Quantize the T-rounded normed to fp8. + const float inv = max_val / amax; + vec_t out_vec; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) out_vec[j] = normed[j] * inv; + out_vec.cast_store(output + bx * stride_output + col); + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +// Host launcher. VEC_SIZE is chosen by d so num_threads = d/VEC_SIZE <= 1024 (single wave), with +// 128 % VEC_SIZE == 0 for the block reduction. Caller must ensure d % (32*VEC_SIZE) == 0 and +// d % 128 == 0 (VEC_SIZE=8 for d<=8192 -> d%256==0; VEC_SIZE=16 for d<=16384 -> d%512==0). +// scale_stride_m = round_up(batch_size, 4) (the block_scale column stride, elements). +template +cudaError_t FusedAddRMSNormFP8BlockQuant(T* input, T* residual, T* weight, O* output, + float* block_scale, T* normed_out, uint32_t batch_size, + uint32_t d, uint32_t stride_input, + uint32_t stride_residual, uint32_t stride_output, + uint32_t scale_stride_m, float eps = 1e-5, + bool enable_pdl = false, cudaStream_t stream = 0) { + const uint32_t vec_size = (d <= 8192) ? 8 : 16; // num_threads = d/vec_size <= 1024 + const uint32_t num_threads = d / vec_size; + const uint32_t num_warps = ceil_div(num_threads, 32); + dim3 nblks(batch_size); + dim3 nthrs(32, num_warps); + const uint32_t smem_size = ceil_div(num_warps, 4) * 4 * sizeof(float); + + cudaLaunchConfig_t config; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + auto kernel = FusedAddRMSNormFP8BlockQuantKernel; + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, input, residual, weight, output, + block_scale, normed_out, d, stride_input, + stride_residual, stride_output, scale_stride_m, eps)); + }); + return cudaSuccess; +} + template cudaError_t GemmaRMSNorm(T* input, T* weight, T* output, uint32_t batch_size, uint32_t d, uint32_t stride_input, uint32_t stride_output, float eps = 1e-5, From 1501b44f4af3d6b7eb7674783f7594db093ca989 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 12 Aug 2026 10:16:50 -0700 Subject: [PATCH 2/7] =?UTF-8?q?fix(bench):=20correct=20byte=20model=20?= =?UTF-8?q?=E2=80=94=20weight=20is=20H,=20not=20M=C3=97H?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fused byte model counted `weight` as a full M×H read; it has only H elements. Real traffic is ~9·M·H, not ~11·M·H, so the reported GB/s was ~22% high (e.g. M=4096,H=8192: 4397 -> 3599 GB/s, ~75% of H200 HBM peak). Latencies and speedups are unchanged. Co-Authored-By: Claude Opus 4.8 --- benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py b/benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py index 228e5ce0cda..f60ac8a4d44 100644 --- a/benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py +++ b/benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py @@ -100,12 +100,13 @@ def time_it(fn): t_f = time_it(fused) t_b = time_it(baseline) - # fused byte model: read input+residual+weight (bf16), write residual+normed(bf16)+fp8+fp32 scale + # fused byte model: read input+residual (M*H bf16) + weight (H bf16); + # write residual+normed (M*H bf16) + fp8 out (M*H) + fp32 block scale. bytes_f = ( - (3 * M * H * 2) - + (2 * M * H * 2) - + (M * H * 1) - + (H // 128 * _round_up(M, 4) * 4) + (4 * M * H * 2) # input, residual reads + residual, normed writes + + (H * 2) # weight (H elements, not M*H) + + (M * H * 1) # fp8 out + + (H // 128 * _round_up(M, 4) * 4) # block scale ) gbps = bytes_f / (t_f * 1e-3) / 1e9 print( From 58635645fe2d14b0be63b3bcd50c43c752371fca Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 12 Aug 2026 10:23:35 -0700 Subject: [PATCH 3/7] fix(norm): validate hidden_size + use normed_out's own row stride Two robustness fixes in fused_add_rmsnorm_fp8_block_quant from review: - Enforce the kernel's single-wave contract. The kernel covers each row with num_threads = d/vec_size and has no per-thread bounds guard, but only hidden_size % 128 == 0 was checked. A valid-looking size like 1152 (% 256 != 0) launched a warp-padded thread count whose extra threads read/wrote past every buffer. The wrapper (and the host launcher) now require hidden_size <= 16384 and hidden_size % (32*vec_size) == 0, else raise / return cudaErrorInvalidValue. - Thread a separate stride_normed. normed_out was written using the fp8 output's row stride; a normed_out that is last-dim contiguous but row-strided differently would write to wrong addresses. It now uses its own stride(0). Co-Authored-By: Claude Opus 4.8 --- csrc/norm.cu | 10 ++++++++-- include/flashinfer/norm.cuh | 24 ++++++++++++++---------- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/csrc/norm.cu b/csrc/norm.cu index e0a16816845..944dcbfc92b 100644 --- a/csrc/norm.cu +++ b/csrc/norm.cu @@ -218,7 +218,13 @@ void fused_add_rmsnorm_fp8_block_quant(TensorView output, TensorView block_scale TVM_FFI_ICHECK_EQ(normed_out.size(1), hidden_size); TVM_FFI_ICHECK_EQ(block_scale.size(0), hidden_size / 128); TVM_FFI_ICHECK_EQ(block_scale.size(1), m_pad); - TVM_FFI_ICHECK_EQ(hidden_size % 128, 0); + // Kernel covers each row single-wave with no per-thread bounds guard: hidden_size must be a whole + // number of warps of vectors -> <= 16384 and a multiple of 32*vec_size (256 for <=8192, 512 + // above). + unsigned int vec_size = hidden_size <= 8192 ? 8 : 16; + TVM_FFI_ICHECK(hidden_size <= 16384 && hidden_size % (32 * vec_size) == 0) + << "fused_add_rmsnorm_fp8_block_quant: hidden_size " << hidden_size + << " unsupported (must be <= 16384 and a multiple of " << (32 * vec_size) << ")"; ffi::CUDADeviceGuard device_guard(input.device().device_id); const cudaStream_t stream = get_stream(input.device()); @@ -230,7 +236,7 @@ void fused_add_rmsnorm_fp8_block_quant(TensorView output, TensorView block_scale static_cast(weight.data_ptr()), static_cast<__nv_fp8_e4m3*>(output.data_ptr()), static_cast(block_scale.data_ptr()), static_cast(normed_out.data_ptr()), batch_size, hidden_size, input.stride(0), residual.stride(0), output.stride(0), - block_scale.stride(0), eps, enable_pdl, stream); + normed_out.stride(0), block_scale.stride(0), eps, enable_pdl, stream); TVM_FFI_ICHECK(status == cudaSuccess) << "FusedAddRMSNormFP8BlockQuant failed with error code " << cudaGetErrorString(status); return true; diff --git a/include/flashinfer/norm.cuh b/include/flashinfer/norm.cuh index 7c8a17673b1..7217d72c855 100644 --- a/include/flashinfer/norm.cuh +++ b/include/flashinfer/norm.cuh @@ -691,7 +691,7 @@ __global__ void FusedAddRMSNormFP8BlockQuantKernel( T* __restrict__ input, T* __restrict__ residual, T* __restrict__ weight, O* __restrict__ output, float* __restrict__ block_scale, T* __restrict__ normed_out, const uint32_t d, const uint32_t stride_input, const uint32_t stride_residual, const uint32_t stride_output, - const uint32_t scale_stride_m, float eps) { + const uint32_t stride_normed, const uint32_t scale_stride_m, float eps) { constexpr uint32_t kQuantBlock = 128; constexpr uint32_t kLanesPerBlock = kQuantBlock / VEC_SIZE; // threads covering one 1x128 block const uint32_t bx = blockIdx.x; // token row m @@ -748,7 +748,7 @@ __global__ void FusedAddRMSNormFP8BlockQuantKernel( normed_vec[j] = nb; lamax = fmaxf(lamax, fabsf(normed[j])); } - normed_vec.store(normed_out + bx * stride_output + col); + normed_vec.store(normed_out + bx * stride_normed + col); // amax over the 1x128 block: kLanesPerBlock consecutive lanes reduce within their warp subgroup. #pragma unroll @@ -771,17 +771,21 @@ __global__ void FusedAddRMSNormFP8BlockQuantKernel( } // Host launcher. VEC_SIZE is chosen by d so num_threads = d/VEC_SIZE <= 1024 (single wave), with -// 128 % VEC_SIZE == 0 for the block reduction. Caller must ensure d % (32*VEC_SIZE) == 0 and -// d % 128 == 0 (VEC_SIZE=8 for d<=8192 -> d%256==0; VEC_SIZE=16 for d<=16384 -> d%512==0). -// scale_stride_m = round_up(batch_size, 4) (the block_scale column stride, elements). +// 128 % VEC_SIZE == 0 for the block reduction. The kernel has no per-thread bounds guard, so d must +// be a whole number of warps of vectors: d <= 16384 and d % (32*VEC_SIZE) == 0 (VEC_SIZE=8 for +// d<=8192 -> d%256==0; VEC_SIZE=16 for d<=16384 -> d%512==0). Returns cudaErrorInvalidValue +// otherwise. scale_stride_m = round_up(batch_size, 4) (the block_scale column stride, elements). template cudaError_t FusedAddRMSNormFP8BlockQuant(T* input, T* residual, T* weight, O* output, float* block_scale, T* normed_out, uint32_t batch_size, uint32_t d, uint32_t stride_input, uint32_t stride_residual, uint32_t stride_output, - uint32_t scale_stride_m, float eps = 1e-5, - bool enable_pdl = false, cudaStream_t stream = 0) { + uint32_t stride_normed, uint32_t scale_stride_m, + float eps = 1e-5, bool enable_pdl = false, + cudaStream_t stream = 0) { const uint32_t vec_size = (d <= 8192) ? 8 : 16; // num_threads = d/vec_size <= 1024 + if (d == 0 || d > 16384 || d % (32 * vec_size) != 0) + return cudaErrorInvalidValue; // kernel covers the row single-wave with no bounds guard const uint32_t num_threads = d / vec_size; const uint32_t num_warps = ceil_div(num_threads, 32); dim3 nblks(batch_size); @@ -801,9 +805,9 @@ cudaError_t FusedAddRMSNormFP8BlockQuant(T* input, T* residual, T* weight, O* ou DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { auto kernel = FusedAddRMSNormFP8BlockQuantKernel; - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, input, residual, weight, output, - block_scale, normed_out, d, stride_input, - stride_residual, stride_output, scale_stride_m, eps)); + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &config, kernel, input, residual, weight, output, block_scale, normed_out, d, stride_input, + stride_residual, stride_output, stride_normed, scale_stride_m, eps)); }); return cudaSuccess; } From f0e9fbef560a68b7b17a3a89545bcc44fa5800cf Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 12 Aug 2026 12:23:36 -0700 Subject: [PATCH 4/7] test(norm): unit tests for fused_add_rmsnorm_fp8_block_quant Mirror test_fused_add_rmsnorm_quant: a torch reference (add -> RMSNorm -> per-1x128-block e4m3 quant) checked against the kernel's residual, normed_out, block_scale, and dequantized fp8 across batch_size x hidden_size x {fp16,bf16} x enable_pdl, plus a reject test that unsupported hidden sizes (1152, 32768) raise. 130 cases pass on H200. Co-Authored-By: Claude Opus 4.8 --- tests/utils/test_norm.py | 89 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/tests/utils/test_norm.py b/tests/utils/test_norm.py index d2cd35be427..32fafe61aaf 100644 --- a/tests/utils/test_norm.py +++ b/tests/utils/test_norm.py @@ -276,6 +276,95 @@ def test_fused_add_rmsnorm_quant( torch.testing.assert_close(residual_fused, residual_native, rtol=1e-3, atol=1e-3) +def fused_add_rms_norm_fp8_block_quant(x, residual, weight, eps, block=128): + """Reference: add residual, RMSNorm, then per-1x128-block fp8-e4m3 quant. + + Returns (fp8 out, scale[M, H/block], pre-norm residual, bf16/fp16 normed). + """ + orig_dtype = x.dtype + xf = x.to(torch.float32) + residual.to(torch.float32) + residual = xf.to(orig_dtype) # pre-norm residual + variance = xf.pow(2).mean(dim=-1, keepdim=True) + normed = (xf * torch.rsqrt(variance + eps) * weight.float()).to( + orig_dtype + ) # round to T + M, H = normed.shape + nb = normed.float().view(M, H // block, block) + amax = nb.abs().amax(dim=-1).clamp_min(1e-4) # (M, H/block) + scale = amax / 448.0 # e4m3 max + q = (nb / scale.unsqueeze(-1)).to(torch.float8_e4m3fn).view(M, H) + return q, scale, residual, normed + + +@pytest.mark.parametrize("batch_size", [1, 19, 99, 989]) +@pytest.mark.parametrize("hidden_size", [256, 512, 1024, 3072, 4096, 6144, 8192, 16384]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("enable_pdl", [True, False]) +def test_fused_add_rmsnorm_fp8_block_quant(batch_size, hidden_size, dtype, enable_pdl): + eps = 1e-6 + x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda") * 0.1 + if enable_pdl and not device_support_pdl(x.device): + pytest.skip("PDL is only available for Hopper and later GPUs") + residual = torch.randn_like(x) * 0.1 + weight = torch.randn(hidden_size, dtype=dtype, device="cuda") + + q_ref, scale_ref, residual_ref, normed_ref = fused_add_rms_norm_fp8_block_quant( + x.clone(), residual.clone(), weight, eps + ) + + m_pad = (batch_size + 3) & ~3 + out = torch.empty(batch_size, hidden_size, dtype=torch.float8_e4m3fn, device="cuda") + block_scale = torch.empty( + hidden_size // 128, m_pad, dtype=torch.float32, device="cuda" + ) + normed_out = torch.empty_like(x) + residual_fused = residual.clone() + flashinfer.norm.fused_add_rmsnorm_fp8_block_quant( + out, + block_scale, + normed_out, + x, + residual_fused, + weight, + eps, + enable_pdl=enable_pdl, + ) + + # pre-norm residual and bf16/fp16 normed: tight + torch.testing.assert_close(residual_fused, residual_ref, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(normed_out, normed_ref, rtol=1e-2, atol=1e-2) + # per-1x128-block scale, read back from the (H/128, round_up(M,4)) column-major buffer + scale_got = block_scale.transpose(0, 1)[:batch_size] + torch.testing.assert_close(scale_got, scale_ref, rtol=1e-3, atol=1e-6) + # dequantized fp8 vs the reference normed (loose fp8 tolerance) + deq = out.float().view(batch_size, hidden_size // 128, 128) * scale_got.unsqueeze( + -1 + ) + torch.testing.assert_close( + deq.view(batch_size, hidden_size), normed_ref.float(), rtol=0.1, atol=0.3 + ) + + +@pytest.mark.parametrize("hidden_size", [1152, 32768]) +def test_fused_add_rmsnorm_fp8_block_quant_rejects_unsupported_hidden(hidden_size): + # single-wave contract: hidden_size must be <= 16384 and a multiple of 32*vec_size + # (256 for <=8192). 1152 is %128==0 but %256!=0; 32768 exceeds 16384. Both must raise. + batch_size = 64 + x = torch.randn(batch_size, hidden_size, dtype=torch.bfloat16, device="cuda") + residual = torch.randn_like(x) + weight = torch.randn(hidden_size, dtype=torch.bfloat16, device="cuda") + m_pad = (batch_size + 3) & ~3 + out = torch.empty(batch_size, hidden_size, dtype=torch.float8_e4m3fn, device="cuda") + block_scale = torch.empty( + hidden_size // 128, m_pad, dtype=torch.float32, device="cuda" + ) + normed_out = torch.empty_like(x) + with pytest.raises(RuntimeError, match="hidden_size"): + flashinfer.norm.fused_add_rmsnorm_fp8_block_quant( + out, block_scale, normed_out, x, residual, weight, 1e-6, enable_pdl=False + ) + + @pytest.mark.parametrize("batch_size", [1, 19, 99, 989]) @pytest.mark.parametrize("hidden_size", [111, 500, 1024, 3072, 3584, 4096, 8192, 16384]) @pytest.mark.parametrize("dtype", [torch.float16]) From ccc3a558cdb674626be656a6b41ad07d6bf888ba Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 12 Aug 2026 13:19:59 -0700 Subject: [PATCH 5/7] feat(norm): mark fused_add_rmsnorm_fp8_block_quant as a public flashinfer_api Add @flashinfer_api to the public op so it participates in API logging/tracing like the other norm entry points. Also widen the unit test's per-block scale tolerance to rtol=2e-2/atol=1e-5: scale = amax/448 and amax can differ by ~2 low-dtype ULP between kernel and reference (fast rsqrt + a different reduction order flip the max element's bf16 rounding), which the prior 1e-3 rtol did not allow. 130 cases pass on H200. Co-Authored-By: Claude Opus 4.8 --- flashinfer/norm/__init__.py | 1 + tests/utils/test_norm.py | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/flashinfer/norm/__init__.py b/flashinfer/norm/__init__.py index 0188c710b12..7286f2b57b8 100644 --- a/flashinfer/norm/__init__.py +++ b/flashinfer/norm/__init__.py @@ -389,6 +389,7 @@ def _fused_add_rmsnorm_quant_fake( pass +@flashinfer_api @register_custom_op( "flashinfer::fused_add_rmsnorm_fp8_block_quant", mutates_args=("out", "block_scale", "normed_out", "residual"), diff --git a/tests/utils/test_norm.py b/tests/utils/test_norm.py index 32fafe61aaf..023be852c44 100644 --- a/tests/utils/test_norm.py +++ b/tests/utils/test_norm.py @@ -333,9 +333,12 @@ def test_fused_add_rmsnorm_fp8_block_quant(batch_size, hidden_size, dtype, enabl # pre-norm residual and bf16/fp16 normed: tight torch.testing.assert_close(residual_fused, residual_ref, rtol=1e-3, atol=1e-3) torch.testing.assert_close(normed_out, normed_ref, rtol=1e-2, atol=1e-2) - # per-1x128-block scale, read back from the (H/128, round_up(M,4)) column-major buffer + # per-1x128-block scale, read back from the (H/128, round_up(M,4)) column-major buffer. + # scale = block amax / 448; amax can differ by up to ~2 low-dtype ULP between kernel and + # reference (fast rsqrt + a different reduction order flip the max element's rounding), so + # use a dtype-appropriate rtol rather than an exact match. scale_got = block_scale.transpose(0, 1)[:batch_size] - torch.testing.assert_close(scale_got, scale_ref, rtol=1e-3, atol=1e-6) + torch.testing.assert_close(scale_got, scale_ref, rtol=2e-2, atol=1e-5) # dequantized fp8 vs the reference normed (loose fp8 tolerance) deq = out.float().view(batch_size, hidden_size // 128, 128) * scale_got.unsqueeze( -1 From 3268a087ab13ee6fde767a0609f04804c3e54e47 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 26 Aug 2026 14:15:32 -0700 Subject: [PATCH 6/7] fix(norm): validate out/block_scale/normed_out in fp8 block-quant binding The binding reinterprets every output buffer with a fixed type -- `output` as __nv_fp8_e4m3, `block_scale` as float, and `normed_out`/`residual`/`weight` as the dispatched `c_type` -- and indexes block_scale as [(d/128)*stride(0) + m], i.e. unit-stride along m. None of that was checked host-side, so a wrong dtype or an m-major block_scale silently corrupted memory instead of raising. Add the missing checks: output shape + e4m3 dtype, block_scale fp32 dtype + last-dim contiguity, and same-dtype-as-input for residual/weight/normed_out. Also correct the `out` docstring, which claimed e5m2 was accepted. Covered by test_fused_add_rmsnorm_fp8_block_quant_rejects_bad_outputs. Addresses review feedback on PR #4480. AI-assisted (Claude Code). Co-Authored-By: Claude Opus 5 (1M context) --- csrc/norm.cu | 11 +++++++++ flashinfer/norm/__init__.py | 3 ++- tests/utils/test_norm.py | 46 +++++++++++++++++++++++++++++++++++++ 3 files changed, 59 insertions(+), 1 deletion(-) diff --git a/csrc/norm.cu b/csrc/norm.cu index 944dcbfc92b..734a87cead2 100644 --- a/csrc/norm.cu +++ b/csrc/norm.cu @@ -197,6 +197,8 @@ void fused_add_rmsnorm_fp8_block_quant(TensorView output, TensorView block_scale CHECK_LAST_DIM_CONTIGUOUS_INPUT(weight); CHECK_LAST_DIM_CONTIGUOUS_INPUT(output); CHECK_LAST_DIM_CONTIGUOUS_INPUT(normed_out); + // The kernel indexes block_scale as [(d/128) * stride(0) + m], i.e. unit-stride along m. + CHECK_LAST_DIM_CONTIGUOUS_INPUT(block_scale); CHECK_DEVICE(input, residual); CHECK_DEVICE(input, weight); CHECK_DEVICE(input, output); @@ -208,12 +210,21 @@ void fused_add_rmsnorm_fp8_block_quant(TensorView output, TensorView block_scale CHECK_DIM(2, output); // fp8 (batch_size, hidden_size) CHECK_DIM(2, normed_out); CHECK_DIM(2, block_scale); // (hidden_size/128, round_up(batch_size, 4)) + // residual/weight/normed_out share the kernel's compute type with input; output is e4m3 and + // block_scale is fp32 (the dtypes the fp8 block-scaled GEMMs consume). + CHECK_SAME_DTYPE(residual, input); + CHECK_SAME_DTYPE(weight, input); + CHECK_SAME_DTYPE(normed_out, input); + CHECK_INPUT_TYPE(output, dl_float8_e4m3fn); + CHECK_INPUT_TYPE(block_scale, dl_float32); unsigned int batch_size = input.size(0); unsigned int hidden_size = input.size(1); unsigned int m_pad = (batch_size + 3) & ~3u; TVM_FFI_ICHECK_EQ(residual.size(0), batch_size); TVM_FFI_ICHECK_EQ(residual.size(1), hidden_size); TVM_FFI_ICHECK_EQ(weight.size(0), hidden_size); + TVM_FFI_ICHECK_EQ(output.size(0), batch_size); + TVM_FFI_ICHECK_EQ(output.size(1), hidden_size); TVM_FFI_ICHECK_EQ(normed_out.size(0), batch_size); TVM_FFI_ICHECK_EQ(normed_out.size(1), hidden_size); TVM_FFI_ICHECK_EQ(block_scale.size(0), hidden_size / 128); diff --git a/flashinfer/norm/__init__.py b/flashinfer/norm/__init__.py index 7286f2b57b8..2f5cd42b3ad 100644 --- a/flashinfer/norm/__init__.py +++ b/flashinfer/norm/__init__.py @@ -419,7 +419,8 @@ def fused_add_rmsnorm_fp8_block_quant( Parameters ---------- out : torch.Tensor - fp8 output, shape ``(batch_size, hidden_size)``, dtype float8_e4m3fn / float8_e5m2. + fp8 output, shape ``(batch_size, hidden_size)``, dtype float8_e4m3fn (the activation + dtype the fp8 block-scaled GEMMs consume). block_scale : torch.Tensor fp32 block scales, shape ``(hidden_size // 128, round_up(batch_size, 4))``, contiguous. This is the column-major (MN-major, TMA-aligned) buffer deep_gemm expects: the logical diff --git a/tests/utils/test_norm.py b/tests/utils/test_norm.py index 023be852c44..3061e2ac70c 100644 --- a/tests/utils/test_norm.py +++ b/tests/utils/test_norm.py @@ -368,6 +368,52 @@ def test_fused_add_rmsnorm_fp8_block_quant_rejects_unsupported_hidden(hidden_siz ) +@pytest.mark.parametrize( + "bad", + [ + "out_dtype", + "out_shape", + "block_scale_dtype", + "block_scale_stride", + "normed_out_dtype", + ], +) +def test_fused_add_rmsnorm_fp8_block_quant_rejects_bad_outputs(bad): + # The kernel reinterprets the output buffers with fixed types (e4m3 out, fp32 block_scale, + # input-dtype normed_out) and indexes block_scale unit-stride along m, so each of these would + # silently corrupt memory without a host-side check. + batch_size, hidden_size = 64, 1024 + dtype = torch.bfloat16 + x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda") + residual = torch.randn_like(x) + weight = torch.randn(hidden_size, dtype=dtype, device="cuda") + m_pad = (batch_size + 3) & ~3 + out = torch.empty(batch_size, hidden_size, dtype=torch.float8_e4m3fn, device="cuda") + block_scale = torch.empty( + hidden_size // 128, m_pad, dtype=torch.float32, device="cuda" + ) + normed_out = torch.empty_like(x) + + if bad == "out_dtype": + out = out.to(torch.float8_e5m2) + elif bad == "out_shape": + out = out[:, : hidden_size // 2] + elif bad == "block_scale_dtype": + block_scale = block_scale.to(torch.float16) + elif bad == "block_scale_stride": + # transposed (m-major) view: last dim is no longer unit-stride + block_scale = torch.empty( + m_pad, hidden_size // 128, dtype=torch.float32, device="cuda" + ).t() + elif bad == "normed_out_dtype": + normed_out = normed_out.to(torch.float32) + + with pytest.raises(RuntimeError): + flashinfer.norm.fused_add_rmsnorm_fp8_block_quant( + out, block_scale, normed_out, x, residual, weight, 1e-6, enable_pdl=False + ) + + @pytest.mark.parametrize("batch_size", [1, 19, 99, 989]) @pytest.mark.parametrize("hidden_size", [111, 500, 1024, 3072, 3584, 4096, 8192, 16384]) @pytest.mark.parametrize("dtype", [torch.float16]) From 2b2e6cbdc3ad09515dcd4279c0420c470039968d Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 26 Aug 2026 14:15:32 -0700 Subject: [PATCH 7/7] docs(norm): list fused_add_rmsnorm_fp8_block_quant in the API reference The new public API was missing from docs/api/norm.rst, so it never made it into the generated docs. Addresses review feedback on PR #4480. AI-assisted (Claude Code). Co-Authored-By: Claude Opus 5 (1M context) --- docs/api/norm.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/api/norm.rst b/docs/api/norm.rst index 7f722583b28..ccb5c9cf5c1 100644 --- a/docs/api/norm.rst +++ b/docs/api/norm.rst @@ -14,6 +14,7 @@ Kernels for normalization layers. rmsnorm_quant fused_add_rmsnorm fused_add_rmsnorm_quant + fused_add_rmsnorm_fp8_block_quant gemma_rmsnorm gemma_fused_add_rmsnorm layernorm