File size: 2,712 Bytes
02c100f 3b968ce 02c100f 3b968ce 02c100f 3b968ce 02c100f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | # 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.
|