# adaptive-layernorm-producers

FlashRT native CUDA producer kernels for DiT/Wan-style adaptive LayerNorm
blocks.

This package fuses the producer side of the block:

```text
LayerNorm / adaptive LayerNorm modulation -> low-precision activation
```

into a single Kernel Hub operator. It is intended for transformer and diffuser
runtimes where the following GEMM already consumes FP8 or NVFP4 activations.
Without this fusion, the runtime pays for separate normalization, modulation,
BF16 materialization, and quantization launches before the GEMM.

The source kernels are derived from the production FlashRT runtime:

- `official/FlashRT/csrc/quantize/ada_layer_norm_fp8.cu`
- `official/FlashRT/csrc/kernels/dit_bf16.cu::layer_norm_no_affine_fp8_static_bf16`

The full FlashRT model runtime and serving pipeline live upstream at
[LiangSu8899/FlashRT](https://github.com/LiangSu8899/FlashRT).

## Available Functions

- `ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, eps=1e-5, out=None)`
- `ada_layer_norm_quant_fp8_ptok_bf16(x, scale, shift, act_scale, eps=1e-5, out=None)`
- `ada_layer_norm_quant_fp8_ptok_table_bf16(x, temb, table, act_scale, shift_idx, scale_idx, eps=1e-5, out=None)`
- `ada_layer_norm_quant_fp8_modfp8_bf16(x, scale_fp8, shift_fp8, scale_deq, shift_deq, act_scale, eps=1e-5, out=None)`
- `awq_ada_layer_norm_quant_fp8_bf16(x, scale, shift, inv_s, act_scale, eps=1e-5, out=None)`
- `ada_layer_norm_quant_nvfp4_swizzled_bf16(x, scale, shift, eps=1e-5, packed=None, sf_swizzled=None)`
- `ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(x, scale_fp8, shift_fp8, scale_deq, shift_deq, eps=1e-5, packed=None, sf_swizzled=None)`
- `layer_norm_no_affine_quant_fp8_static_bf16(x, act_scale, eps=1e-5, out=None)`
- `layer_norm_no_affine_quant_nvfp4_swizzled_bf16(x, eps=1e-5, packed=None, sf_swizzled=None)`
- `adaln_modulation6_bf16(adaln_params, layer_modulation, out=None)`
- `swizzled_sf_size(rows, dim)`

## Tensor Contract

- `x`: contiguous CUDA `torch.bfloat16`, shape `(rows, dim)`.
- Broadcast `scale`, `shift`, and `inv_s`: contiguous CUDA
  `torch.bfloat16`, shape `(dim,)`.
- Per-token `scale` and `shift`: contiguous CUDA `torch.bfloat16`, shape
  `(rows, dim)`.
- Table form `temb`: contiguous CUDA `torch.bfloat16`, shape
  `(rows, chunks, dim)`; `table`: contiguous CUDA `torch.float32`, shape
  `(chunks, dim)`.
- `scale_fp8`, `shift_fp8`: contiguous CUDA `torch.float8_e4m3fn`, shape `(dim,)`.
- `scale_deq`, `shift_deq`, `act_scale`: CUDA `torch.float32` scalar tensors.
- FP8 outputs use `torch.float8_e4m3fn`, shape `(rows, dim)`.
- NVFP4 packed output uses `torch.uint8`, shape `(rows, dim // 2)`.
- NVFP4 scale output uses the FlashRT/CUTLASS 128x4 swizzled uint8 scale
  layout. Allocate with `swizzled_sf_size(rows, dim)`.
- `dim` must be even for FP8 outputs and divisible by 16 for NVFP4 outputs.
- The package targets CUDA 12.8+ and Blackwell-class deployment paths.
- `adaln_modulation6_bf16` takes FP32 `(batch, sequence, 6, dim)` parameters
  plus FP32 `(6, dim)` layer modulation and returns six BF16
  `(batch, sequence, dim)` tensors.

## Minimal Usage

```python
from kernels import get_kernel
import torch

ops = get_kernel("flashrt/adaptive-layernorm-producers", version=1, trust_remote_code=True)

rows, dim = 2520, 3072
x = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16)
scale = torch.zeros((dim,), device="cuda", dtype=torch.bfloat16)
shift = torch.zeros((dim,), device="cuda", dtype=torch.bfloat16)
act_scale = torch.tensor([0.025], device="cuda", dtype=torch.float32)

x_fp8 = ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale)
```

Per-token modulation without materializing a broadcast buffer:

```python
ptok_scale = torch.zeros_like(x)
ptok_shift = torch.zeros_like(x)
x_fp8 = ops.ada_layer_norm_quant_fp8_ptok_bf16(
    x, ptok_scale, ptok_shift, act_scale
)
```

Fused host modulation table add and chunk selection:

```python
chunks = 6
temb = torch.zeros((rows, chunks, dim), device="cuda", dtype=torch.bfloat16)
table = torch.zeros((chunks, dim), device="cuda", dtype=torch.float32)
x_fp8 = ops.ada_layer_norm_quant_fp8_ptok_table_bf16(
    x, temb, table, act_scale, shift_idx=0, scale_idx=1
)
```

Static-buffer usage for CUDA Graph capture:

```python
out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, out=out)
```

NVFP4 producer:

```python
packed, sf = ops.ada_layer_norm_quant_nvfp4_swizzled_bf16(x, scale, shift)

# GROOT N1.7 pre-FFN producer on Thor.
packed_ln, sf_ln = ops.layer_norm_no_affine_quant_nvfp4_swizzled_bf16(x)
```

Six-way DiT modulation:

```python
params = torch.randn((1, rows, 6, dim), device="cuda", dtype=torch.float32)
modulation = torch.randn((6, dim), device="cuda", dtype=torch.float32)
outputs = ops.adaln_modulation6_bf16(params, modulation)
```

## Validation

```bash
python adaptive-layernorm-producers/tests/test_adaptive_layernorm_producers.py --backend source --mode full
python adaptive-layernorm-producers/benchmarks/benchmark.py --backend source --iters 100
```

The validation suite checks exact FP8 output for small producer shapes, exact
NVFP4 packed/scalefactor output against the native staged BF16-norm plus FP4
quantizer contract on SM110, and strict FP8
boundary accounting for long video shapes. The long-shape policy requires
`p99_abs == 0` and only a tiny count of adjacent FP8-code boundary differences
against the eager reference.

See `VALIDATION.md` and `benchmarks/RESULTS.md` for the current local RTX 5090
source-build results.

The SM110 FP4 producer uses the native GROOT N1.7 contract from FlashRT commit
`24df793f4fa2d50780aea03b644208c6e0cb4162`. Its scale conversion and FP4 bin
boundaries intentionally differ from the older SM120 producer; release tests
compare each architecture against its own production contract.

The aarch64 release variant `torch211-cxx11-cu130-aarch64-linux` contains
native SM87 and SM110a code objects. Its SM110a path was validated on NVIDIA
Thor from the exact Hub-downloaded artifact over decode, DiT, vision, and
2K/4K video rows. The two per-token entry points were also compared directly
with their raw registered ops using A-B-B-A timing; all six wrapper/raw rows
were bit-exact and the measured ratio was `0.997-1.000`. SM87 runtime
validation remains pending on Jetson AGX Orin hardware.