flashinfer.norm.fused_add_rmsnorm_fp8_block_quant

flashinfer.norm.fused_add_rmsnorm_fp8_block_quant(out: Tensor, block_scale: Tensor, normed_out: Tensor, input: Tensor, residual: Tensor, weight: Tensor, eps: float = 1e-06, enable_pdl: bool | None = None) → None

Fused add-residual + RMSNorm + 1x128 fp8 block quantization, in one pass.

Replaces the two-kernel producer (fused_add_rmsnorm() then a per-token-group 1x128 fp8 quant) used before fp8 block-scaled GEMMs. Unlike 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 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 (torch.Tensor) – Shape (batch_size, hidden_size). residual is updated in place with input+residual.

  • 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<hidden_size<=16384).

  • eps (float) – Epsilon for numerical stability.

  • enable_pdl (bool) – Whether to enable programmatic dependent launch.