| # 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: |
|
|
| ```text |
| 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 |
|
|
| ```python |
| 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 |
|
|
| ```bash |
| 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. |
|
|
