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..f60ac8a4d44 --- /dev/null +++ b/benchmarks/bench_fused_add_rmsnorm_fp8_block_quant.py @@ -0,0 +1,164 @@ +"""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 (M*H bf16) + weight (H bf16); + # write residual+normed (M*H bf16) + fp8 out (M*H) + fp32 block scale. + bytes_f = ( + (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( + 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..734a87cead2 100644 --- a/csrc/norm.cu +++ b/csrc/norm.cu @@ -186,6 +186,74 @@ 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); + // 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); + 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)) + // 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); + TVM_FFI_ICHECK_EQ(block_scale.size(1), m_pad); + // 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()); + + // 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), + 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; + }); +} + void gemma_rmsnorm(TensorView output, TensorView input, TensorView weight, double eps, bool enable_pdl) { CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); 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 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..2f5cd42b3ad 100644 --- a/flashinfer/norm/__init__.py +++ b/flashinfer/norm/__init__.py @@ -389,6 +389,76 @@ 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"), +) +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 (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 + ``(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..7217d72c855 100644 --- a/include/flashinfer/norm.cuh +++ b/include/flashinfer/norm.cuh @@ -675,6 +675,143 @@ 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 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 + 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_normed + 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. 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 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); + 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, stride_normed, 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, diff --git a/tests/utils/test_norm.py b/tests/utils/test_norm.py index d2cd35be427..3061e2ac70c 100644 --- a/tests/utils/test_norm.py +++ b/tests/utils/test_norm.py @@ -276,6 +276,144 @@ 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 = 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=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 + ) + 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( + "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])