| # flashrt-adaptive-norms |
|
|
| Tensor-facing FlashRT adaptive normalization kernels for Hugging Face `kernels`. |
|
|
| This package targets DiT, VLA, video, and world-model blocks that combine |
| RMSNorm with per-row style scale/shift/gate parameters: |
|
|
| ```text |
| AdaRMSNorm: x -> RMSNorm(x, weight) -> norm * (1 + style_scale) + style_shift |
| Fused gate path: residual += x * gate -> AdaRMSNorm(residual) -> static FP8 |
| ``` |
|
|
| ## Exported APIs |
|
|
| - `ada_rms_norm_style_bf16(x, weight, style, eps=1e-6, out=None, gate_out=None)` |
| - `gate_residual_ada_norm_fp8_static_bf16(residual, x, gate, weight, style, scale, eps=1e-6, out=None, gate_out=None)` |
|
|
| ## Tensor Conventions |
|
|
| - `x`, `residual`, `gate`: contiguous BF16 matrices, shape `(rows, dim)` |
| - `weight`: contiguous BF16 vector, shape `(dim,)` |
| - `style`: contiguous BF16 matrix, shape `(rows, 3 * dim)` or `(1, 3 * dim)` |
| for a style row broadcast across all rows |
| - first `dim`: style scale |
| - second `dim`: style shift |
| - third `dim`: gate output |
| - `scale`: CUDA FP32 scalar tensor for static FP8 quantization |
| - FP8 output dtype: `torch.float8_e4m3fn` |
| - `dim` must be even. |
|
|
| ## Minimal Usage |
|
|
| ```python |
| from kernels import get_kernel |
| import torch |
| |
| ops = get_kernel( |
| "flashrt/flashrt-adaptive-norms", |
| version=1, |
| trust_remote_code=True, |
| ) |
| |
| rows, dim = 2520, 3072 |
| x = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16) |
| residual = torch.randn_like(x) |
| gate = torch.randn_like(x) |
| weight = torch.ones((dim,), device="cuda", dtype=torch.bfloat16) |
| style = torch.randn((rows, 3 * dim), device="cuda", dtype=torch.bfloat16) |
| scale = torch.tensor([0.04], device="cuda", dtype=torch.float32) |
| |
| out_bf16, gate_out = ops.ada_rms_norm_style_bf16(x, weight, style) |
| residual, out_fp8, gate_out = ops.gate_residual_ada_norm_fp8_static_bf16( |
| residual, |
| x, |
| gate, |
| weight, |
| style, |
| scale, |
| ) |
| ``` |
|
|
| For prompt/style values shared by every token, pass a single row directly; |
| the kernel broadcasts it without materializing a `rows x 3*dim` tensor: |
|
|
| ```python |
| shared_style = torch.randn((1, 3 * dim), device="cuda", dtype=torch.bfloat16) |
| out_bf16, gate_out = ops.ada_rms_norm_style_bf16( |
| x, weight, shared_style |
| ) |
| ``` |
|
|
| ## Validation |
|
|
| ```bash |
| python flashrt-adaptive-norms/tests/test_adaptive_norms.py --backend source --mode full |
| python flashrt-adaptive-norms/benchmarks/benchmark.py --backend source --shapes all |
| ``` |
|
|
| RTX 5090 and Thor SM110 source-extension validation passes the full source |
| grid for both per-row and single-row broadcast style. |
| Residual and gate outputs are bit-level exact. FP8 output uses a boundary-aware |
| gate because PyTorch and CUDA FP8 casters can choose adjacent FP8 values for |
| rare tie cases; p99_abs is zero across the source grid. |
| |
