flashrt-residual-norm-quant

Tensor-facing FlashRT residual, RMSNorm, and static FP8 quantization kernels for Hugging Face kernels.

This package is runtime glue for model hot paths that need to produce the next layer's FP8 activation without falling back to multiple PyTorch operations:

BF16 residual/x -> residual add -> RMSNorm -> static-scale FP8 E4M3 activation

Exported APIs

• rms_norm_bf16(x, weight, eps=1e-6, out=None)
• rms_norm_quant_fp8_static_bf16(x, weight, scale, eps=1e-6, out=None)
• residual_add_rms_norm_quant_fp8_static_bf16(residual, x, weight, scale, eps=1e-6, out=None)
• residual_add_rms_norm_bf16(residual, x, weight, eps=1e-6, out=None)

The residual API updates residual in place with residual += x, rounded to BF16, then emits the normalized FP8 activation. The no-quant twin preserves the same in-place BF16 residual contract and emits BF16, for runtimes whose next consumer is not low precision.

Tensor Conventions

• x: BF16 tensor, shape (rows, dim)
• residual: BF16 tensor, shape (rows, dim), in-place updated
• weight: BF16 tensor, shape (dim,)
• scale: CUDA float32 scalar tensor used for static FP8 quantization
• out: FP8 E4M3 tensor, shape (rows, dim)

The hidden dimension must be even because this version uses FlashRT's packed BF16 pair path.

Minimal Usage

from kernels import get_kernel
import torch

ops = get_kernel(
    "flashrt/flashrt-residual-norm-quant",
    version=1,
    trust_remote_code=True,
)

x = torch.randn((10, 1024), device="cuda", dtype=torch.bfloat16)
residual = torch.randn_like(x)
weight = torch.ones((1024,), device="cuda", dtype=torch.bfloat16)
scale = torch.tensor([0.04], device="cuda", dtype=torch.float32)

x_fp8 = ops.residual_add_rms_norm_quant_fp8_static_bf16(
    residual,
    x,
    weight,
    scale,
    eps=1e-6,
)

Validation

python flashrt-residual-norm-quant/tests/test_residual_norm_quant.py --backend source --mode full
python flashrt-residual-norm-quant/benchmarks/benchmark.py --backend source --shapes all

Current RTX 5090 source-extension rows pass the package's strict BF16/FP8 distribution gates across the PI0.5, VLA, video, and Cosmos3-Edge shape grid. Built artifacts must pass the same gate before publication.

The no-quant BF16 path is also gated at the Cosmos3-Edge production hidden size (rows=128, dim=2048) and preserves the in-place residual contract. Its static-output CUDA Graph replay is checked bitwise.