Promote latest kernel artifacts to main

CARD.md

@@ -0,0 +1,65 @@
+# flashrt/adaptive-layernorm-producers
+
+FlashRT native CUDA adaptive LayerNorm producer kernels for DiT, Wan-style
+video diffusion, and VLA/runtime demo pipelines.
+
+This package fuses normalization/modulation and low-precision activation
+production before FP8 or NVFP4 GEMM consumers.
+
+## Functions
+
+- `ada_layer_norm_quant_fp8_bf16`
+- `ada_layer_norm_quant_fp8_ptok_bf16`
+- `ada_layer_norm_quant_fp8_ptok_table_bf16`
+- `ada_layer_norm_quant_fp8_modfp8_bf16`
+- `awq_ada_layer_norm_quant_fp8_bf16`
+- `ada_layer_norm_quant_nvfp4_swizzled_bf16`
+- `ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16`
+- `layer_norm_no_affine_quant_fp8_static_bf16`
+- `layer_norm_no_affine_quant_nvfp4_swizzled_bf16`
+- `adaln_modulation6_bf16`
+- `swizzled_sf_size`
+
+## Quick Start
+
+```python
+from kernels import get_kernel
+import torch
+
+ops = get_kernel("flashrt/adaptive-layernorm-producers", version=1, trust_remote_code=True)
+
+x = torch.randn((2520, 3072), device="cuda", dtype=torch.bfloat16)
+scale = torch.zeros((3072,), device="cuda", dtype=torch.bfloat16)
+shift = torch.zeros((3072,), 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 and table-fused modulation are available for video DiT
+and VLA blocks whose shift/scale values vary by token:
+
+```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
+)
+
+chunks = 6
+temb = torch.zeros((x.shape[0], chunks, x.shape[1]), device="cuda",
+                   dtype=torch.bfloat16)
+table = torch.zeros((chunks, x.shape[1]), 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
+)
+```
+
+The `torch211-cxx11-cu130-aarch64-linux` variant is a native fat binary with
+SM87 (Jetson AGX Orin) and SM110a (Jetson AGX Thor) code objects. The SM110a
+path is runtime-validated on Thor; SM87 runtime validation remains pending on
+Orin hardware. Runtime variant selection is handled by `get_kernel`.
+
+Use `README.md` for tensor contracts and `VALIDATION.md` for correctness and
+benchmark status.

README.md

@@ -1,9 +1,151 @@
-# flashrt/adaptive-layernorm-producers
+# adaptive-layernorm-producers
 
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
+FlashRT native CUDA producer kernels for DiT/Wan-style adaptive LayerNorm
+blocks.
 
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/adaptive-layernorm-producers
+This package fuses the producer side of the block:
 
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.
+```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.

SYNC.md

@@ -0,0 +1,19 @@
+# Source Sync
+
+- Upstream: `flashrt-project/FlashRT`
+- GROOT N1.7 sync commit:
+  `24df793f4fa2d50780aea03b644208c6e0cb4162`
+- Copied sources: `csrc/dit_norm_fp4_sfa.cu/.cuh`
+
+The package adds `sm110_fp4_dispatch.cu/.cuh` and Tensor-facing bindings around
+the native producer. The SM110 implementation is an independent CUDA 13,
+CUTLASS 4.4 target. Existing SM120 FP8/FP4 paths are unchanged.
+
+The SM110 FP4 scale contract is E4M3 round-to-nearest. It must be checked
+against the staged native BF16 norm plus native FP4 quantizer, including exact
+packed values and swizzled scale bytes. Do not reuse the older SM120 ceil-scale
+reference for this target.
+
+`tests/native_reference/dit_bf16.cu` is an exact, test-only copy from the same
+FlashRT revision. It keeps source and installed-artifact gates self-contained;
+it is not listed in `build.toml` and is never linked into release artifacts.

benchmarks/README.md

@@ -0,0 +1,17 @@
+# Benchmarks
+
+Run source benchmarks:
+
+```bash
+python adaptive-layernorm-producers/benchmarks/benchmark.py --backend source --iters 100
+```
+
+Run installed-artifact benchmarks after Hub packaging:
+
+```bash
+python adaptive-layernorm-producers/benchmarks/benchmark.py --backend installed --iters 100
+```
+
+`RESULTS.md` records the current local source-build numbers. Public benchmark
+claims should be refreshed after installed-artifact testing on the target
+hardware.

benchmarks/native_binding.cpp

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: Apache-2.0
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include "adaln_modulation6.cuh"
+
+void adaln_modulation6(
+    uintptr_t params, uintptr_t modulation, uintptr_t out0, uintptr_t out1,
+    uintptr_t out2, uintptr_t out3, uintptr_t out4, uintptr_t out5,
+    int batch, int sequence, int dim) {
+  flash_rt::adaln_producers::adaln_modulation6_bf16(
+      reinterpret_cast<const float*>(params),
+      reinterpret_cast<const float*>(modulation),
+      reinterpret_cast<void*>(out0), reinterpret_cast<void*>(out1),
+      reinterpret_cast<void*>(out2), reinterpret_cast<void*>(out3),
+      reinterpret_cast<void*>(out4), reinterpret_cast<void*>(out5),
+      batch, sequence, dim, at::cuda::getCurrentCUDAStream().stream());
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("adaln_modulation6", &adaln_modulation6);
+}

build.toml

@@ -0,0 +1,50 @@
+[general]
+name = "adaptive-layernorm-producers"
+license = "Apache-2.0"
+version = 1
+backends = ["cuda"]
+
+[general.cuda]
+minver = "12.8"
+
+[general.hub]
+repo-id = "flashrt/adaptive-layernorm-producers"
+
+[torch]
+include = ["csrc"]
+src = [
+  "torch-ext/torch_binding.cpp",
+]
+
+[kernel.adaptive_layernorm_producers]
+backend = "cuda"
+depends = ["torch"]
+include = ["csrc"]
+cuda-minver = "12.8"
+src = [
+  "csrc/ada_layer_norm_fp8.cu",
+  "csrc/ada_layer_norm_fp8_ptok.cu",
+  "csrc/ada_layer_norm_fp8.cuh",
+  "csrc/dit_layer_norm_fp8.cu",
+  "csrc/dit_layer_norm_fp8.cuh",
+  "csrc/adaln_modulation6.cu",
+  "csrc/adaln_modulation6.cuh",
+]
+
+[kernel.adaptive_layernorm_producers_sm110_fp4]
+backend = "cuda"
+depends = ["torch", "cutlass_4_4"]
+include = ["csrc"]
+cuda-minver = "13"
+cuda-capabilities = ["11.0a"]
+cuda-flags = [
+  "-O3",
+  "--use_fast_math",
+  "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
+]
+src = [
+  "csrc/dit_norm_fp4_sfa.cu",
+  "csrc/dit_norm_fp4_sfa.cuh",
+  "csrc/sm110_fp4_dispatch.cu",
+  "csrc/sm110_fp4_dispatch.cuh",
+]

build/torch211-cxx11-cu130-aarch64-linux/__init__.py

@@ -198,6 +198,22 @@ def _layer_norm_no_affine_quant_fp8_static_bf16_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_no_affine_quant_nvfp4_swizzled_bf16"))
+def _layer_norm_no_affine_quant_nvfp4_swizzled_bf16_fake(
+    x: torch.Tensor,
+    eps: float,
+    packed: torch.Tensor,
+    sf_swizzled: torch.Tensor,
+) -> None:
+    _check_x(x)
+    rows, dim = x.shape
+    if packed.shape != (rows, dim // 2):
+        raise RuntimeError("packed must have shape (rows, dim // 2)")
+    if sf_swizzled.numel() < swizzled_sf_size(rows, dim):
+        raise RuntimeError("sf_swizzled is too small")
+    return None
+
+
 def ada_layer_norm_quant_fp8_bf16(
     x: torch.Tensor,
     scale: torch.Tensor,
@@ -353,6 +369,21 @@ def layer_norm_no_affine_quant_fp8_static_bf16(
     return out
 
 
+def layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+    x: torch.Tensor,
+    eps: float = 1e-5,
+    packed: torch.Tensor | None = None,
+    sf_swizzled: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SM110 no-affine LayerNorm producing NVFP4 and CUTLASS SFA."""
+    if packed is None or sf_swizzled is None:
+        packed, sf_swizzled = _alloc_nvfp4(x)
+    ops.layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+        x, float(eps), packed, sf_swizzled
+    )
+    return packed, sf_swizzled
+
+
 @torch.library.register_fake(add_op_namespace_prefix("adaln_modulation6_bf16"))
 def _adaln_modulation6_bf16_fake(
     adaln_params: torch.Tensor,
@@ -418,6 +449,7 @@ __all__ = [
     "ada_layer_norm_quant_nvfp4_swizzled_bf16",
     "ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16",
     "layer_norm_no_affine_quant_fp8_static_bf16",
+    "layer_norm_no_affine_quant_nvfp4_swizzled_bf16",
     "adaln_modulation6_bf16",
     "swizzled_sf_size",
 ]

build/torch211-cxx11-cu130-aarch64-linux/{_adaptive_layernorm_producers_cuda_7781728.abi3.so → _adaptive_layernorm_producers_cuda_8a66d8b.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2b5542ae448332ebd4d1f2ca2f0030f88200061f8d2e64eb8581c314e7d7322c
-size 592256
+oid sha256:b7a2db2e007dbc80da262c206891058c7db8c1c7f249b92ad6185c1ce11b814f
+size 535776

build/torch211-cxx11-cu130-aarch64-linux/_ops.py

@@ -1,6 +1,6 @@
 import torch
-from . import _adaptive_layernorm_producers_cuda_7781728
-ops = torch.ops._adaptive_layernorm_producers_cuda_7781728
+from . import _adaptive_layernorm_producers_cuda_8a66d8b
+ops = torch.ops._adaptive_layernorm_producers_cuda_8a66d8b
 
 def add_op_namespace_prefix(op_name: str):
-    return f"_adaptive_layernorm_producers_cuda_7781728::{op_name}"
+    return f"_adaptive_layernorm_producers_cuda_8a66d8b::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/metadata.json

@@ -1,33 +1,28 @@
 {
   "name": "adaptive-layernorm-producers",
-  "id": "_adaptive_layernorm_producers_cuda_7781728",
+  "id": "_adaptive_layernorm_producers_cuda_8a66d8b",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
-      "8.7",
       "11.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "Y0RRm5MDP9m0wKQRpuGRFg0YewrbqdZnVIBItbh7DAc=",
-      "_adaptive_layernorm_producers_cuda_7781728.abi3.so": "K1VCrkSDMuvU0fLKLwAw+IIABh+NLmTrhYHDFOfXMiw=",
-      "_ops.py": "kG1EL+EuLDtXPuUtyj/ygYLPUZg6eszWjMCBtbLC0pw=",
+      "__init__.py": "/6Ld6WSyL7PjILcYt4P7x3fnXi2C+ovbSP/7nNFetb8=",
+      "_adaptive_layernorm_producers_cuda_8a66d8b.abi3.so": "t6LbLgB9vIDaJiwgaJEFjH24wcfySbkq1hhcHOEbgU8=",
+      "_ops.py": "8CjL8CxdQouAOnaox+O0AKwLERGSDc8mao+V3ArMuTk=",
       "adaptive_layernorm_producers/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "9b532b42ad71d0ef6bc79d1372e96910087faad4",
+      "sha": "8a66d8b",
       "dirty": false
-    },
-    "validation": {
-      "torch": "2.11.0+cu130",
-      "cuda": "13.0"
     }
   }
 }

csrc/ada_layer_norm_fp8.cu

@@ -0,0 +1,666 @@
+// SPDX-License-Identifier: Apache-2.0
+// G7.17 — Fused AdaLayerNorm + per-tensor FP8 quantize. See header.
+// Math mirrors csrc/kernels/dit_bf16.cu::ada_layer_norm_bf16 with the
+// final store path replaced by fp8 e4m3 quantize+pack.
+
+#include "ada_layer_norm_fp8.cuh"
+
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace quantize {
+
+namespace {
+
+constexpr float kFp8Max = 448.0f;
+
+__device__ __forceinline__ uint8_t f32_to_fp4_e2m1(float x) {
+  if (x == 0.0f) return 0;
+  uint8_t sign = x < 0 ? 0x8 : 0;
+  float ax = fabsf(x);
+  if (ax < 0.25f) return sign;
+  if (ax < 0.75f) return sign | 0x1;  // 0.5
+  if (ax < 1.5f) return sign | 0x2;   // 1
+  if (ax < 3.0f) return sign | 0x3;   // 2
+  if (ax < 5.0f) return sign | 0x6;   // 4
+  return sign | 0x7;                  // 6
+}
+
+__device__ __forceinline__ uint8_t f32_to_ue4m3_ceil_local(float x) {
+  if (x <= 0.0f) return 0;
+  if (x < 0.0009765625f) return 1;
+  int exp;
+  float mant = frexpf(x, &exp);
+  exp -= 1;
+  float frac = mant * 2.0f - 1.0f;
+  int mantissa = (int)ceilf(frac * 8.0f);
+  if (mantissa >= 8) {
+    mantissa = 0;
+    exp += 1;
+  }
+  int biased_exp = exp + 7;
+  if (biased_exp <= 0) return 1;
+  if (biased_exp >= 15) return 0x7F;
+  return (uint8_t)((biased_exp << 3) | mantissa);
+}
+
+__device__ __forceinline__ float ue4m3_to_f32_local(uint8_t v) {
+  if (v == 0) return 0.0f;
+  int exp = ((v >> 3) & 0xF) - 7;
+  int mant = v & 0x7;
+  float m = 1.0f + mant / 8.0f;
+  return ldexpf(m, exp);
+}
+
+__global__ void ada_layer_norm_fp8_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ scale,
+    const __nv_bfloat16* __restrict__ shift,
+    __nv_fp8_e4m3*       __restrict__ out,
+    const float*          __restrict__ act_scale_ptr,
+    int dim, float eps)
+{
+  const int row = blockIdx.x;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + (long long)row * dim);
+  const __nv_bfloat162* sc2 =
+      reinterpret_cast<const __nv_bfloat162*>(scale);
+  const __nv_bfloat162* sh2 =
+      reinterpret_cast<const __nv_bfloat162*>(shift);
+  __nv_fp8_e4m3* out_row = out + (long long)row * dim;
+  const int dim2 = dim >> 1;
+
+  extern __shared__ float shared[];
+
+  // ── Pass 1: mean ──
+  float local_sum = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    local_sum += __bfloat162float(v.x) + __bfloat162float(v.y);
+  }
+  float val = local_sum;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  const int lane = threadIdx.x & 31;
+  const int wid  = threadIdx.x >> 5;
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float mean = shared[0] / static_cast<float>(dim);
+
+  // ── Pass 2: variance ──
+  float local_var = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    const float d0 = __bfloat162float(v.x) - mean;
+    const float d1 = __bfloat162float(v.y) - mean;
+    local_var += d0 * d0 + d1 * d1;
+  }
+  val = local_var;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float inv_std = rsqrtf(shared[0] / static_cast<float>(dim) + eps);
+
+  // ── Pass 3: ada_modulate + fp8 quantize ──
+  // Register-level round-through bf16 BEFORE fp8 conversion to match
+  // the reference (ada_layer_norm_bf16 → bf16 store → quantize_fp8)
+  // bit-exactly while skipping the actual bf16 global-memory round-trip.
+  // This is the precision/memory tradeoff sweet spot: zero memory cost
+  // for the bf16 intermediate, exact cos match with the 2-launch chain.
+  const float inv_a = 1.0f / *act_scale_ptr;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 xv = x2[i], sv = sc2[i], hv = sh2[i];
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float v0_f32 = n0 * (1.0f + __bfloat162float(sv.x))
+                          + __bfloat162float(hv.x);
+    const float v1_f32 = n1 * (1.0f + __bfloat162float(sv.y))
+                          + __bfloat162float(hv.y);
+    // Round through bf16 in registers (no memory write).
+    const float v0 = __bfloat162float(__float2bfloat16(v0_f32));
+    const float v1 = __bfloat162float(__float2bfloat16(v1_f32));
+    float q0 = fminf(fmaxf(v0 * inv_a, -kFp8Max), kFp8Max);
+    float q1 = fminf(fmaxf(v1 * inv_a, -kFp8Max), kFp8Max);
+    out_row[2 * i]     = __nv_fp8_e4m3(q0);
+    out_row[2 * i + 1] = __nv_fp8_e4m3(q1);
+  }
+}
+
+__global__ void awq_ada_layer_norm_fp8_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ scale,
+    const __nv_bfloat16* __restrict__ shift,
+    const __nv_bfloat16* __restrict__ inv_s,
+    __nv_fp8_e4m3*       __restrict__ out,
+    const float*          __restrict__ act_scale_ptr,
+    int dim, float eps)
+{
+  const int row = blockIdx.x;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + (long long)row * dim);
+  const __nv_bfloat162* sc2 =
+      reinterpret_cast<const __nv_bfloat162*>(scale);
+  const __nv_bfloat162* sh2 =
+      reinterpret_cast<const __nv_bfloat162*>(shift);
+  const __nv_bfloat162* inv2 =
+      reinterpret_cast<const __nv_bfloat162*>(inv_s);
+  __nv_fp8_e4m3* out_row = out + (long long)row * dim;
+  const int dim2 = dim >> 1;
+
+  extern __shared__ float shared[];
+
+  float local_sum = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    local_sum += __bfloat162float(v.x) + __bfloat162float(v.y);
+  }
+  float val = local_sum;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  const int lane = threadIdx.x & 31;
+  const int wid  = threadIdx.x >> 5;
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float mean = shared[0] / static_cast<float>(dim);
+
+  float local_var = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    const float d0 = __bfloat162float(v.x) - mean;
+    const float d1 = __bfloat162float(v.y) - mean;
+    local_var += d0 * d0 + d1 * d1;
+  }
+  val = local_var;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float inv_std = rsqrtf(shared[0] / static_cast<float>(dim) + eps);
+
+  const float inv_a = 1.0f / *act_scale_ptr;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 xv = x2[i], sv = sc2[i], hv = sh2[i], iv = inv2[i];
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float v0 = __bfloat162float(__float2bfloat16(
+        n0 * (1.0f + __bfloat162float(sv.x)) + __bfloat162float(hv.x)));
+    const float v1 = __bfloat162float(__float2bfloat16(
+        n1 * (1.0f + __bfloat162float(sv.y)) + __bfloat162float(hv.y)));
+    float q0 = fminf(fmaxf(v0 * __bfloat162float(iv.x) * inv_a,
+                           -kFp8Max), kFp8Max);
+    float q1 = fminf(fmaxf(v1 * __bfloat162float(iv.y) * inv_a,
+                           -kFp8Max), kFp8Max);
+    out_row[2 * i]     = __nv_fp8_e4m3(q0);
+    out_row[2 * i + 1] = __nv_fp8_e4m3(q1);
+  }
+}
+
+__global__ void ada_layer_norm_fp8_modfp8_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_fp8_e4m3* __restrict__ scale,
+    const __nv_fp8_e4m3* __restrict__ shift,
+    const float* __restrict__ scale_deq,
+    const float* __restrict__ shift_deq,
+    __nv_fp8_e4m3* __restrict__ out,
+    const float* __restrict__ act_scale_ptr,
+    int dim, float eps)
+{
+  const int row = blockIdx.x;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + (long long)row * dim);
+  __nv_fp8_e4m3* out_row = out + (long long)row * dim;
+  const int dim2 = dim >> 1;
+
+  extern __shared__ float shared[];
+  float local_sum = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    local_sum += __bfloat162float(v.x) + __bfloat162float(v.y);
+  }
+  float val = local_sum;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  const int lane = threadIdx.x & 31;
+  const int wid  = threadIdx.x >> 5;
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float mean = shared[0] / static_cast<float>(dim);
+
+  float local_var = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    const float d0 = __bfloat162float(v.x) - mean;
+    const float d1 = __bfloat162float(v.y) - mean;
+    local_var += d0 * d0 + d1 * d1;
+  }
+  val = local_var;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float inv_std = rsqrtf(shared[0] / static_cast<float>(dim) + eps);
+
+  const float s_deq = *scale_deq;
+  const float h_deq = *shift_deq;
+  const float inv_a = 1.0f / *act_scale_ptr;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 xv = x2[i];
+    const int j = i << 1;
+    const float s0 = static_cast<float>(scale[j]) * s_deq;
+    const float s1 = static_cast<float>(scale[j + 1]) * s_deq;
+    const float h0 = static_cast<float>(shift[j]) * h_deq;
+    const float h1 = static_cast<float>(shift[j + 1]) * h_deq;
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float v0 = __bfloat162float(__float2bfloat16(n0 * (1.0f + s0) + h0));
+    const float v1 = __bfloat162float(__float2bfloat16(n1 * (1.0f + s1) + h1));
+    out_row[j] = __nv_fp8_e4m3(fminf(fmaxf(v0 * inv_a, -kFp8Max), kFp8Max));
+    out_row[j + 1] = __nv_fp8_e4m3(fminf(fmaxf(v1 * inv_a, -kFp8Max), kFp8Max));
+  }
+}
+
+__global__ void ada_layer_norm_nvfp4_swizzled_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ scale,
+    const __nv_bfloat16* __restrict__ shift,
+    uint8_t* __restrict__ packed,
+    uint8_t* __restrict__ sf_swizzled,
+    int dim, int num_blocks, int n_col_blocks, float eps)
+{
+  const int row = blockIdx.x;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + (long long)row * dim);
+  const __nv_bfloat162* sc2 =
+      reinterpret_cast<const __nv_bfloat162*>(scale);
+  const __nv_bfloat162* sh2 =
+      reinterpret_cast<const __nv_bfloat162*>(shift);
+  uint8_t* row_packed = packed + (long long)row * dim / 2;
+  const int dim2 = dim >> 1;
+
+  extern __shared__ float shared[];
+  float* red = shared;
+  float* blk_scale = shared + 256;
+
+  float local_sum = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    local_sum += __bfloat162float(v.x) + __bfloat162float(v.y);
+  }
+  float val = local_sum;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  const int lane = threadIdx.x & 31;
+  const int wid = threadIdx.x >> 5;
+  if (!lane) red[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? red[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) red[0] = val;
+  __syncthreads();
+  const float mean = red[0] / static_cast<float>(dim);
+
+  float local_var = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    const float d0 = __bfloat162float(v.x) - mean;
+    const float d1 = __bfloat162float(v.y) - mean;
+    local_var += d0 * d0 + d1 * d1;
+  }
+  val = local_var;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  if (!lane) red[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? red[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) red[0] = val;
+  __syncthreads();
+  const float inv_std = rsqrtf(red[0] / static_cast<float>(dim) + eps);
+
+  for (int b = threadIdx.x; b < num_blocks; b += blockDim.x)
+    blk_scale[b] = 0.0f;
+  __syncthreads();
+
+  for (int i2 = threadIdx.x; i2 < dim2; i2 += blockDim.x) {
+    __nv_bfloat162 xv = x2[i2], sv = sc2[i2], hv = sh2[i2];
+    const int i = i2 << 1;
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float v0 = __bfloat162float(__float2bfloat16(
+        n0 * (1.0f + __bfloat162float(sv.x)) + __bfloat162float(hv.x)));
+    const float v1 = __bfloat162float(__float2bfloat16(
+        n1 * (1.0f + __bfloat162float(sv.y)) + __bfloat162float(hv.y)));
+    atomicMax((int*)&blk_scale[i >> 4], __float_as_int(fabsf(v0)));
+    atomicMax((int*)&blk_scale[(i + 1) >> 4], __float_as_int(fabsf(v1)));
+  }
+  __syncthreads();
+
+  const int rb = row / 128;
+  const int ri = row % 128;
+  for (int b = threadIdx.x; b < num_blocks; b += blockDim.x) {
+    float amax = __int_as_float(*(int*)&blk_scale[b]);
+    uint8_t ue = f32_to_ue4m3_ceil_local(amax / 6.0f);
+    int cb = b / 4;
+    int ci = b % 4;
+    int out_idx = (rb * n_col_blocks + cb) * 512
+        + (ri % 32) * 16 + (ri / 32) * 4 + ci;
+    sf_swizzled[out_idx] = ue;
+    blk_scale[b] = ue4m3_to_f32_local(ue);
+  }
+  __syncthreads();
+
+  for (int p = threadIdx.x; p < dim / 2; p += blockDim.x) {
+    const int i = p << 1;
+    const int i2 = p;
+    __nv_bfloat162 xv = x2[i2], sv = sc2[i2], hv = sh2[i2];
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float m0 = __bfloat162float(__float2bfloat16(
+        n0 * (1.0f + __bfloat162float(sv.x)) + __bfloat162float(hv.x)));
+    const float m1 = __bfloat162float(__float2bfloat16(
+        n1 * (1.0f + __bfloat162float(sv.y)) + __bfloat162float(hv.y)));
+    const int b0 = i >> 4;
+    const int b1 = (i + 1) >> 4;
+    const float inv0 = blk_scale[b0] > 0.0f ? 1.0f / blk_scale[b0] : 0.0f;
+    const float inv1 = blk_scale[b1] > 0.0f ? 1.0f / blk_scale[b1] : 0.0f;
+    uint8_t lo = f32_to_fp4_e2m1(m0 * inv0);
+    uint8_t hi = f32_to_fp4_e2m1(m1 * inv1);
+    row_packed[p] = (hi << 4) | (lo & 0x0F);
+  }
+}
+
+__global__ void ada_layer_norm_nvfp4_swizzled_modfp8_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_fp8_e4m3* __restrict__ scale,
+    const __nv_fp8_e4m3* __restrict__ shift,
+    const float* __restrict__ scale_deq,
+    const float* __restrict__ shift_deq,
+    uint8_t* __restrict__ packed,
+    uint8_t* __restrict__ sf_swizzled,
+    int dim, int num_blocks, int n_col_blocks, float eps)
+{
+  const int row = blockIdx.x;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + (long long)row * dim);
+  uint8_t* row_packed = packed + (long long)row * dim / 2;
+  const int dim2 = dim >> 1;
+
+  extern __shared__ float shared[];
+  float* red = shared;
+  float* blk_scale = shared + 256;
+
+  float local_sum = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    local_sum += __bfloat162float(v.x) + __bfloat162float(v.y);
+  }
+  float val = local_sum;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  const int lane = threadIdx.x & 31;
+  const int wid = threadIdx.x >> 5;
+  if (!lane) red[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? red[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) red[0] = val;
+  __syncthreads();
+  const float mean = red[0] / static_cast<float>(dim);
+
+  float local_var = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    const float d0 = __bfloat162float(v.x) - mean;
+    const float d1 = __bfloat162float(v.y) - mean;
+    local_var += d0 * d0 + d1 * d1;
+  }
+  val = local_var;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  if (!lane) red[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? red[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) red[0] = val;
+  __syncthreads();
+  const float inv_std = rsqrtf(red[0] / static_cast<float>(dim) + eps);
+
+  for (int b = threadIdx.x; b < num_blocks; b += blockDim.x)
+    blk_scale[b] = 0.0f;
+  __syncthreads();
+
+  const float s_deq = *scale_deq;
+  const float h_deq = *shift_deq;
+  for (int i2 = threadIdx.x; i2 < dim2; i2 += blockDim.x) {
+    __nv_bfloat162 xv = x2[i2];
+    const int i = i2 << 1;
+    const float s0 = static_cast<float>(scale[i]) * s_deq;
+    const float s1 = static_cast<float>(scale[i + 1]) * s_deq;
+    const float h0 = static_cast<float>(shift[i]) * h_deq;
+    const float h1 = static_cast<float>(shift[i + 1]) * h_deq;
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float v0 = __bfloat162float(__float2bfloat16(n0 * (1.0f + s0) + h0));
+    const float v1 = __bfloat162float(__float2bfloat16(n1 * (1.0f + s1) + h1));
+    atomicMax((int*)&blk_scale[i >> 4], __float_as_int(fabsf(v0)));
+    atomicMax((int*)&blk_scale[(i + 1) >> 4], __float_as_int(fabsf(v1)));
+  }
+  __syncthreads();
+
+  const int rb = row / 128;
+  const int ri = row % 128;
+  for (int b = threadIdx.x; b < num_blocks; b += blockDim.x) {
+    float amax = __int_as_float(*(int*)&blk_scale[b]);
+    uint8_t ue = f32_to_ue4m3_ceil_local(amax / 6.0f);
+    int cb = b / 4;
+    int ci = b % 4;
+    int out_idx = (rb * n_col_blocks + cb) * 512
+        + (ri % 32) * 16 + (ri / 32) * 4 + ci;
+    sf_swizzled[out_idx] = ue;
+    blk_scale[b] = ue4m3_to_f32_local(ue);
+  }
+  __syncthreads();
+
+  for (int p = threadIdx.x; p < dim / 2; p += blockDim.x) {
+    const int i = p << 1;
+    __nv_bfloat162 xv = x2[p];
+    const float s0 = static_cast<float>(scale[i]) * s_deq;
+    const float s1 = static_cast<float>(scale[i + 1]) * s_deq;
+    const float h0 = static_cast<float>(shift[i]) * h_deq;
+    const float h1 = static_cast<float>(shift[i + 1]) * h_deq;
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float m0 = __bfloat162float(__float2bfloat16(n0 * (1.0f + s0) + h0));
+    const float m1 = __bfloat162float(__float2bfloat16(n1 * (1.0f + s1) + h1));
+    const int b0 = i >> 4;
+    const int b1 = (i + 1) >> 4;
+    const float inv0 = blk_scale[b0] > 0.0f ? 1.0f / blk_scale[b0] : 0.0f;
+    const float inv1 = blk_scale[b1] > 0.0f ? 1.0f / blk_scale[b1] : 0.0f;
+    uint8_t lo = f32_to_fp4_e2m1(m0 * inv0);
+    uint8_t hi = f32_to_fp4_e2m1(m1 * inv1);
+    row_packed[p] = (hi << 4) | (lo & 0x0F);
+  }
+}
+
+}  // namespace
+
+void ada_layer_norm_fp8(
+    const void*  x_bf16,
+    const void*  scale_bf16,
+    const void*  shift_bf16,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream)
+{
+  if (seq_len <= 0 || dim <= 0) return;
+  // dim must be even for bf162 vectorization (always true for our shapes).
+  ada_layer_norm_fp8_kernel<<<seq_len, 256, 256 * sizeof(float), stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(x_bf16),
+      reinterpret_cast<const __nv_bfloat16*>(scale_bf16),
+      reinterpret_cast<const __nv_bfloat16*>(shift_bf16),
+      reinterpret_cast<__nv_fp8_e4m3*>(out_fp8),
+      act_scale,
+      dim, eps);
+}
+
+void ada_layer_norm_fp8_modfp8(
+    const void*  x_bf16,
+    const void*  scale_fp8,
+    const void*  shift_fp8,
+    const float* scale_deq,
+    const float* shift_deq,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream)
+{
+  if (seq_len <= 0 || dim <= 0) return;
+  ada_layer_norm_fp8_modfp8_kernel<<<seq_len, 256, 256 * sizeof(float), stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(x_bf16),
+      reinterpret_cast<const __nv_fp8_e4m3*>(scale_fp8),
+      reinterpret_cast<const __nv_fp8_e4m3*>(shift_fp8),
+      scale_deq, shift_deq,
+      reinterpret_cast<__nv_fp8_e4m3*>(out_fp8),
+      act_scale,
+      dim, eps);
+}
+
+void awq_ada_layer_norm_fp8(
+    const void*  x_bf16,
+    const void*  scale_bf16,
+    const void*  shift_bf16,
+    const void*  inv_s_bf16,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream)
+{
+  if (seq_len <= 0 || dim <= 0) return;
+  awq_ada_layer_norm_fp8_kernel
+      <<<seq_len, 256, 256 * sizeof(float), stream>>>(
+          reinterpret_cast<const __nv_bfloat16*>(x_bf16),
+          reinterpret_cast<const __nv_bfloat16*>(scale_bf16),
+          reinterpret_cast<const __nv_bfloat16*>(shift_bf16),
+          reinterpret_cast<const __nv_bfloat16*>(inv_s_bf16),
+          reinterpret_cast<__nv_fp8_e4m3*>(out_fp8),
+          act_scale, dim, eps);
+}
+
+void ada_layer_norm_nvfp4_swizzled(
+    const void* x_bf16,
+    const void* scale_bf16,
+    const void* shift_bf16,
+    void* packed_u8,
+    void* sf_swizzled_u8,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream)
+{
+  if (seq_len <= 0 || dim <= 0) return;
+  const int num_blocks = (dim + 15) / 16;
+  const int n_col_blocks = (num_blocks + 3) / 4;
+  const int smem = (256 + num_blocks) * sizeof(float);
+  ada_layer_norm_nvfp4_swizzled_kernel<<<seq_len, 256, smem, stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(x_bf16),
+      reinterpret_cast<const __nv_bfloat16*>(scale_bf16),
+      reinterpret_cast<const __nv_bfloat16*>(shift_bf16),
+      reinterpret_cast<uint8_t*>(packed_u8),
+      reinterpret_cast<uint8_t*>(sf_swizzled_u8),
+      dim, num_blocks, n_col_blocks, eps);
+}
+
+void ada_layer_norm_nvfp4_swizzled_modfp8(
+    const void* x_bf16,
+    const void* scale_fp8,
+    const void* shift_fp8,
+    const float* scale_deq,
+    const float* shift_deq,
+    void* packed_u8,
+    void* sf_swizzled_u8,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream)
+{
+  if (seq_len <= 0 || dim <= 0) return;
+  const int num_blocks = (dim + 15) / 16;
+  const int n_col_blocks = (num_blocks + 3) / 4;
+  const int smem = (256 + num_blocks) * sizeof(float);
+  ada_layer_norm_nvfp4_swizzled_modfp8_kernel<<<seq_len, 256, smem, stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(x_bf16),
+      reinterpret_cast<const __nv_fp8_e4m3*>(scale_fp8),
+      reinterpret_cast<const __nv_fp8_e4m3*>(shift_fp8),
+      scale_deq, shift_deq,
+      reinterpret_cast<uint8_t*>(packed_u8),
+      reinterpret_cast<uint8_t*>(sf_swizzled_u8),
+      dim, num_blocks, n_col_blocks, eps);
+}
+
+}  // namespace quantize
+}  // namespace flash_rt

csrc/ada_layer_norm_fp8.cuh

@@ -0,0 +1,102 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// G7.17 — Fused AdaLayerNorm + per-tensor FP8 e4m3 quantize.
+//
+// Replaces the motus 2-launch chain
+//   ada_layer_norm_bf16(x, scale, shift) -> bf16 modulated
+//   quantize_fp8_static(modulated)       -> fp8
+// with a single kernel. Eliminates the bf16 (B*L, D) intermediate
+// (~15 MB at T=2520 D=3072 for Wan video), which is the dominant
+// cost — memory traffic round-trip dwarfs launch overhead.
+//
+// Math (per row, given mean μ and variance σ² of x):
+//   norm[i] = (x[i] - μ) / sqrt(σ² + eps)
+//   mod[i]  = norm[i] * (1 + scale[i]) + shift[i]
+//   out[i]  = clamp(mod[i] / *act_scale, ±448).to(fp8_e4m3)
+//
+// Layout: x, scale, shift are (B*L, D) bf16 row-major (or scale/shift
+// broadcast over rows when their leading dim == 1). out is (B*L, D)
+// fp8_e4m3 row-major. act_scale is a device fp32 scalar.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace quantize {
+
+void ada_layer_norm_fp8(
+    const void*  x_bf16,
+    const void*  scale_bf16,
+    const void*  shift_bf16,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream);
+
+void ada_layer_norm_fp8_modfp8(
+    const void*  x_bf16,
+    const void*  scale_fp8,
+    const void*  shift_fp8,
+    const float* scale_deq,
+    const float* shift_deq,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream);
+
+void awq_ada_layer_norm_fp8(
+    const void*  x_bf16,
+    const void*  scale_bf16,
+    const void*  shift_bf16,
+    const void*  inv_s_bf16,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream);
+
+void ada_layer_norm_nvfp4_swizzled(
+    const void* x_bf16,
+    const void* scale_bf16,
+    const void* shift_bf16,
+    void* packed_u8,
+    void* sf_swizzled_u8,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream);
+
+void ada_layer_norm_nvfp4_swizzled_modfp8(
+    const void* x_bf16,
+    const void* scale_fp8,
+    const void* shift_fp8,
+    const float* scale_deq,
+    const float* shift_deq,
+    void* packed_u8,
+    void* sf_swizzled_u8,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream);
+
+// Per-token modulation form: scale/shift are [seq_len, dim], one row per
+// activation row (video-DiT per-token timestep tables).
+void ada_layer_norm_ptok_fp8(
+    const void*  x_bf16,
+    const void*  scale_bf16,
+    const void*  shift_bf16,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream);
+
+// Table form: the block's [n_chunks, dim] modulation table is added to
+// the per-token embedding [seq_len, n_chunks, dim] inside the kernel —
+// the host-side per-block chunk materialization disappears.
+void ada_layer_norm_ptok_table_fp8(
+    const void*  x_bf16,
+    const void*  temb_bf16,
+    const float* table_f32,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, int n_chunks, int shift_idx, int scale_idx,
+    float eps, cudaStream_t stream);
+
+}  // namespace quantize
+}  // namespace flash_rt

csrc/ada_layer_norm_fp8_ptok.cu

@@ -0,0 +1,254 @@
+// SPDX-License-Identifier: Apache-2.0
+// Per-token AdaLayerNorm + per-tensor FP8 quantize.
+//
+// Same math as ada_layer_norm_fp8.cu::ada_layer_norm_fp8_kernel with one
+// difference: the modulation vectors are per token ([M, D], one row per
+// activation row) instead of broadcast per channel ([D]). This is the
+// producer form modern video DiTs need — their timestep embedding carries
+// a scale/shift *per token* (e.g. a [B, S, 6, D] table chunked per site),
+// so a [D] broadcast producer cannot serve the seam at all.
+
+#include "ada_layer_norm_fp8.cuh"
+
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace quantize {
+
+namespace {
+
+constexpr float kFp8MaxPtok = 448.0f;
+
+__global__ void ada_layer_norm_ptok_fp8_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ scale,
+    const __nv_bfloat16* __restrict__ shift,
+    __nv_fp8_e4m3*       __restrict__ out,
+    const float*          __restrict__ act_scale_ptr,
+    int dim, float eps)
+{
+  const int row = blockIdx.x;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + (long long)row * dim);
+  const __nv_bfloat162* sc2 =
+      reinterpret_cast<const __nv_bfloat162*>(scale + (long long)row * dim);
+  const __nv_bfloat162* sh2 =
+      reinterpret_cast<const __nv_bfloat162*>(shift + (long long)row * dim);
+  __nv_fp8_e4m3* out_row = out + (long long)row * dim;
+  const int dim2 = dim >> 1;
+
+  extern __shared__ float shared[];
+
+  // ── Pass 1: mean ──
+  float local_sum = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    local_sum += __bfloat162float(v.x) + __bfloat162float(v.y);
+  }
+  float val = local_sum;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  const int lane = threadIdx.x & 31;
+  const int wid  = threadIdx.x >> 5;
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float mean = shared[0] / static_cast<float>(dim);
+
+  // ── Pass 2: variance ──
+  float local_var = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    const float d0 = __bfloat162float(v.x) - mean;
+    const float d1 = __bfloat162float(v.y) - mean;
+    local_var += d0 * d0 + d1 * d1;
+  }
+  val = local_var;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float inv_std = rsqrtf(shared[0] / static_cast<float>(dim) + eps);
+
+  // ── Pass 3: per-token ada_modulate + fp8 quantize ──
+  // Register-level round-through bf16 before fp8, matching the family's
+  // reference chain semantics (see ada_layer_norm_fp8.cu).
+  const float inv_a = 1.0f / *act_scale_ptr;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 xv = x2[i], sv = sc2[i], hv = sh2[i];
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float v0_f32 = n0 * (1.0f + __bfloat162float(sv.x))
+                          + __bfloat162float(hv.x);
+    const float v1_f32 = n1 * (1.0f + __bfloat162float(sv.y))
+                          + __bfloat162float(hv.y);
+    const float v0 = __bfloat162float(__float2bfloat16(v0_f32));
+    const float v1 = __bfloat162float(__float2bfloat16(v1_f32));
+    float q0 = fminf(fmaxf(v0 * inv_a, -kFp8MaxPtok), kFp8MaxPtok);
+    float q1 = fminf(fmaxf(v1 * inv_a, -kFp8MaxPtok), kFp8MaxPtok);
+    out_row[2 * i]     = __nv_fp8_e4m3(q0);
+    out_row[2 * i + 1] = __nv_fp8_e4m3(q1);
+  }
+}
+
+}  // namespace
+
+void ada_layer_norm_ptok_fp8(
+    const void*  x_bf16,
+    const void*  scale_bf16,
+    const void*  shift_bf16,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, float eps,
+    cudaStream_t stream)
+{
+  if (seq_len <= 0 || dim <= 0) return;
+  ada_layer_norm_ptok_fp8_kernel<<<seq_len, 256, 256 * sizeof(float),
+                                   stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(x_bf16),
+      reinterpret_cast<const __nv_bfloat16*>(scale_bf16),
+      reinterpret_cast<const __nv_bfloat16*>(shift_bf16),
+      reinterpret_cast<__nv_fp8_e4m3*>(out_fp8),
+      act_scale,
+      dim, eps);
+}
+
+}  // namespace quantize
+}  // namespace flash_rt
+
+namespace flash_rt {
+namespace quantize {
+
+namespace {
+
+__global__ void ada_layer_norm_ptok_table_fp8_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ temb,   // [M, n_chunks, D]
+    const float*         __restrict__ table,  // [n_chunks, D]
+    __nv_fp8_e4m3*       __restrict__ out,
+    const float*          __restrict__ act_scale_ptr,
+    int dim, int n_chunks, int shift_idx, int scale_idx, float eps)
+{
+  const int row = blockIdx.x;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + (long long)row * dim);
+  const __nv_bfloat162* sc2 = reinterpret_cast<const __nv_bfloat162*>(
+      temb + ((long long)row * n_chunks + scale_idx) * dim);
+  const __nv_bfloat162* sh2 = reinterpret_cast<const __nv_bfloat162*>(
+      temb + ((long long)row * n_chunks + shift_idx) * dim);
+  const float2* tsc2 = reinterpret_cast<const float2*>(
+      table + (long long)scale_idx * dim);
+  const float2* tsh2 = reinterpret_cast<const float2*>(
+      table + (long long)shift_idx * dim);
+  __nv_fp8_e4m3* out_row = out + (long long)row * dim;
+  const int dim2 = dim >> 1;
+
+  extern __shared__ float shared[];
+
+  float local_sum = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    local_sum += __bfloat162float(v.x) + __bfloat162float(v.y);
+  }
+  float val = local_sum;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  const int lane = threadIdx.x & 31;
+  const int wid  = threadIdx.x >> 5;
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float mean = shared[0] / static_cast<float>(dim);
+
+  float local_var = 0.f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 v = x2[i];
+    const float d0 = __bfloat162float(v.x) - mean;
+    const float d1 = __bfloat162float(v.y) - mean;
+    local_var += d0 * d0 + d1 * d1;
+  }
+  val = local_var;
+  for (int o = 16; o > 0; o >>= 1)
+    val += __shfl_xor_sync(0xffffffffu, val, o);
+  if (!lane) shared[wid] = val;
+  __syncthreads();
+  if (!wid) {
+    val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.f;
+    for (int o = 16; o > 0; o >>= 1)
+      val += __shfl_xor_sync(0xffffffffu, val, o);
+  }
+  __syncthreads();
+  if (!threadIdx.x) shared[0] = val;
+  __syncthreads();
+  const float inv_std = rsqrtf(shared[0] / static_cast<float>(dim) + eps);
+
+  const float inv_a = 1.0f / *act_scale_ptr;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    __nv_bfloat162 xv = x2[i], sv = sc2[i], hv = sh2[i];
+    const float2 ts = tsc2[i], th = tsh2[i];
+    const float s0 = ts.x + __bfloat162float(sv.x);
+    const float s1 = ts.y + __bfloat162float(sv.y);
+    const float h0 = th.x + __bfloat162float(hv.x);
+    const float h1 = th.y + __bfloat162float(hv.y);
+    const float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    const float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    const float v0 = __bfloat162float(
+        __float2bfloat16(n0 * (1.0f + s0) + h0));
+    const float v1 = __bfloat162float(
+        __float2bfloat16(n1 * (1.0f + s1) + h1));
+    float q0 = fminf(fmaxf(v0 * inv_a, -kFp8MaxPtok), kFp8MaxPtok);
+    float q1 = fminf(fmaxf(v1 * inv_a, -kFp8MaxPtok), kFp8MaxPtok);
+    out_row[2 * i]     = __nv_fp8_e4m3(q0);
+    out_row[2 * i + 1] = __nv_fp8_e4m3(q1);
+  }
+}
+
+}  // namespace
+
+void ada_layer_norm_ptok_table_fp8(
+    const void*  x_bf16,
+    const void*  temb_bf16,
+    const float* table_f32,
+    void*        out_fp8,
+    const float* act_scale,
+    int seq_len, int dim, int n_chunks, int shift_idx, int scale_idx,
+    float eps, cudaStream_t stream)
+{
+  if (seq_len <= 0 || dim <= 0) return;
+  ada_layer_norm_ptok_table_fp8_kernel<<<seq_len, 256,
+                                         256 * sizeof(float), stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(x_bf16),
+      reinterpret_cast<const __nv_bfloat16*>(temb_bf16),
+      table_f32,
+      reinterpret_cast<__nv_fp8_e4m3*>(out_fp8),
+      act_scale,
+      dim, n_chunks, shift_idx, scale_idx, eps);
+}
+
+}  // namespace quantize
+}  // namespace flash_rt

csrc/adaln_modulation6.cu

@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include "adaln_modulation6.cuh"
+
+#include <cuda_bf16.h>
+
+namespace flash_rt {
+namespace adaln_producers {
+namespace {
+
+__global__ void adaln_modulation6_bf16_kernel(
+    const float* __restrict__ adaln_params,
+    const float* __restrict__ layer_modulation,
+    __nv_bfloat16* __restrict__ out0,
+    __nv_bfloat16* __restrict__ out1,
+    __nv_bfloat16* __restrict__ out2,
+    __nv_bfloat16* __restrict__ out3,
+    __nv_bfloat16* __restrict__ out4,
+    __nv_bfloat16* __restrict__ out5,
+    int dim,
+    long long elements) {
+  const long long index =
+      static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= elements) {
+    return;
+  }
+  const int d = static_cast<int>(index % dim);
+  const long long base = (index / dim) * 6LL * dim + d;
+  out0[index] = __float2bfloat16(
+      adaln_params[base] + layer_modulation[d]);
+  out1[index] = __float2bfloat16(
+      adaln_params[base + dim] + layer_modulation[dim + d]);
+  out2[index] = __float2bfloat16(
+      adaln_params[base + 2LL * dim] + layer_modulation[2LL * dim + d]);
+  out3[index] = __float2bfloat16(
+      adaln_params[base + 3LL * dim] + layer_modulation[3LL * dim + d]);
+  out4[index] = __float2bfloat16(
+      adaln_params[base + 4LL * dim] + layer_modulation[4LL * dim + d]);
+  out5[index] = __float2bfloat16(
+      adaln_params[base + 5LL * dim] + layer_modulation[5LL * dim + d]);
+}
+
+}  // namespace
+
+void adaln_modulation6_bf16(
+    const float* adaln_params,
+    const float* layer_modulation,
+    void* out0,
+    void* out1,
+    void* out2,
+    void* out3,
+    void* out4,
+    void* out5,
+    int batch,
+    int sequence,
+    int dim,
+    cudaStream_t stream) {
+  const long long elements =
+      static_cast<long long>(batch) * sequence * dim;
+  constexpr int threads = 256;
+  const unsigned blocks =
+      static_cast<unsigned>((elements + threads - 1) / threads);
+  adaln_modulation6_bf16_kernel<<<blocks, threads, 0, stream>>>(
+      adaln_params,
+      layer_modulation,
+      static_cast<__nv_bfloat16*>(out0),
+      static_cast<__nv_bfloat16*>(out1),
+      static_cast<__nv_bfloat16*>(out2),
+      static_cast<__nv_bfloat16*>(out3),
+      static_cast<__nv_bfloat16*>(out4),
+      static_cast<__nv_bfloat16*>(out5),
+      dim,
+      elements);
+}
+
+}  // namespace adaln_producers
+}  // namespace flash_rt

csrc/adaln_modulation6.cuh

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace adaln_producers {
+
+void adaln_modulation6_bf16(
+    const float* adaln_params,
+    const float* layer_modulation,
+    void* out0,
+    void* out1,
+    void* out2,
+    void* out3,
+    void* out4,
+    void* out5,
+    int batch,
+    int sequence,
+    int dim,
+    cudaStream_t stream);
+
+}  // namespace adaln_producers
+}  // namespace flash_rt

csrc/dit_layer_norm_fp8.cu

@@ -0,0 +1,124 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// BF16 no-affine LayerNorm -> static FP8 producer.
+// Ported from FlashRT csrc/kernels/dit_bf16.cu for Hub Tensor APIs.
+
+#include "dit_layer_norm_fp8.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace flash_rt {
+namespace adaln_producers {
+namespace {
+
+__global__ void layer_norm_no_affine_fp8_static_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    __nv_fp8_e4m3* __restrict__ out,
+    const float* __restrict__ d_scale,
+    int dim,
+    float eps) {
+  const int row = blockIdx.x;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + static_cast<long long>(row) * dim);
+  __nv_fp8_e4m3* out_row = out + static_cast<long long>(row) * dim;
+  const int dim2 = dim >> 1;
+
+  extern __shared__ float shared[];
+  float local_sum = 0.0f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    const __nv_bfloat162 v = x2[i];
+    local_sum += __bfloat162float(v.x) + __bfloat162float(v.y);
+  }
+  float value = local_sum;
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    value += __shfl_xor_sync(0xffffffffu, value, offset);
+  }
+  const int lane = threadIdx.x & 31;
+  const int wid = threadIdx.x >> 5;
+  if (lane == 0) {
+    shared[wid] = value;
+  }
+  __syncthreads();
+  if (wid == 0) {
+    value = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_xor_sync(0xffffffffu, value, offset);
+    }
+  }
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    shared[0] = value;
+  }
+  __syncthreads();
+  const float mean = shared[0] / static_cast<float>(dim);
+
+  float local_var = 0.0f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    const __nv_bfloat162 v = x2[i];
+    const float d0 = __bfloat162float(v.x) - mean;
+    const float d1 = __bfloat162float(v.y) - mean;
+    local_var += d0 * d0 + d1 * d1;
+  }
+  value = local_var;
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    value += __shfl_xor_sync(0xffffffffu, value, offset);
+  }
+  if (lane == 0) {
+    shared[wid] = value;
+  }
+  __syncthreads();
+  if (wid == 0) {
+    value = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_xor_sync(0xffffffffu, value, offset);
+    }
+  }
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    shared[0] = value;
+  }
+  __syncthreads();
+  const float inv_std = rsqrtf(shared[0] / static_cast<float>(dim) + eps);
+  const float inv_scale = 1.0f / fmaxf(*d_scale, 1e-12f);
+
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    const __nv_bfloat162 xv = x2[i];
+    float v0 = (__bfloat162float(xv.x) - mean) * inv_std;
+    float v1 = (__bfloat162float(xv.y) - mean) * inv_std;
+    v0 = __bfloat162float(__float2bfloat16(v0)) * inv_scale;
+    v1 = __bfloat162float(__float2bfloat16(v1)) * inv_scale;
+    __nv_fp8_e4m3 pair[2];
+    pair[0] = __nv_fp8_e4m3(fminf(fmaxf(v0, -448.0f), 448.0f));
+    pair[1] = __nv_fp8_e4m3(fminf(fmaxf(v1, -448.0f), 448.0f));
+    *reinterpret_cast<uint16_t*>(out_row + 2 * i) =
+        *reinterpret_cast<uint16_t*>(pair);
+  }
+}
+
+}  // namespace
+
+void layer_norm_no_affine_fp8_static_bf16(
+    const void* x_bf16,
+    void* out_fp8,
+    const float* scale,
+    int rows,
+    int dim,
+    float eps,
+    cudaStream_t stream) {
+  if (rows <= 0 || dim <= 0) {
+    return;
+  }
+  layer_norm_no_affine_fp8_static_bf16_kernel
+      <<<rows, 256, 256 * sizeof(float), stream>>>(
+          reinterpret_cast<const __nv_bfloat16*>(x_bf16),
+          reinterpret_cast<__nv_fp8_e4m3*>(out_fp8),
+          scale,
+          dim,
+          eps);
+}
+
+}  // namespace adaln_producers
+}  // namespace flash_rt

csrc/dit_layer_norm_fp8.cuh

@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace adaln_producers {
+
+void layer_norm_no_affine_fp8_static_bf16(
+    const void* x_bf16,
+    void* out_fp8,
+    const float* scale,
+    int rows,
+    int dim,
+    float eps,
+    cudaStream_t stream);
+
+}  // namespace adaln_producers
+}  // namespace flash_rt

csrc/dit_norm_fp4_sfa.cu

@@ -0,0 +1,208 @@
+// ============================================================================
+//  Fused DiT norms (bf16) + NVFP4 quantize + SFA write. See header.
+//
+//  Grid: (S,) one CTA per row, 128 threads. Row mean/var via the standard
+//  two-pass block reduction; 16-element quant blocks one per thread with
+//  16-byte loads. Scale selection and e2m1 rounding are identical to
+//  quantize_fp4_dynamic_sfa_bf16_vec.
+// ============================================================================
+#include "dit_norm_fp4_sfa.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(__CUDA_ARCH__)
+#  include "cutlass/cutlass.h"
+#  include "cutlass/detail/sm100_blockscaled_layout.hpp"
+#  include "cute/tensor.hpp"
+#  define FV_HAVE_CUTLASS 1
+#else
+#  define FV_HAVE_CUTLASS 0
+#endif
+
+namespace flash_rt {
+namespace fused_fp4 {
+
+#if FV_HAVE_CUTLASS
+
+namespace {
+
+using CfgDN = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+__device__ __forceinline__ uint8_t fp32_to_e2m1_dn(float x) {
+    uint8_t sign = (x < 0.f) ? 0x8u : 0x0u;
+    float ax = fabsf(x);
+    uint8_t mant;
+    if      (ax <= 0.25f) mant = 0u;
+    else if (ax <= 0.75f) mant = 1u;
+    else if (ax <= 1.25f) mant = 2u;
+    else if (ax <= 1.75f) mant = 3u;
+    else if (ax <= 2.5f)  mant = 4u;
+    else if (ax <= 3.5f)  mant = 5u;
+    else if (ax <= 5.0f)  mant = 6u;
+    else                  mant = 7u;
+    return sign | mant;
+}
+
+__device__ __forceinline__ float block_sum_dn(float v, float* sh) {
+    const int lane = threadIdx.x & 31;
+    const int warp = threadIdx.x >> 5;
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1)
+        v += __shfl_xor_sync(0xffffffffu, v, o);
+    if (lane == 0) sh[warp] = v;
+    __syncthreads();
+    if (warp == 0) {
+        v = (lane < ((blockDim.x + 31) >> 5)) ? sh[lane] : 0.f;
+        #pragma unroll
+        for (int o = 16; o > 0; o >>= 1)
+            v += __shfl_xor_sync(0xffffffffu, v, o);
+        if (lane == 0) sh[0] = v;
+    }
+    __syncthreads();
+    const float r = sh[0];
+    __syncthreads();
+    return r;
+}
+
+// HAS_MOD=true: AdaLN (modulate with (1+scale), shift). false: plain LN.
+template <bool HAS_MOD, class LayoutSF>
+__global__ void dit_norm_fp4_sfa_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ scale,   // [D] or nullptr
+    const __nv_bfloat16* __restrict__ shift,   // [D] or nullptr
+    uint2* __restrict__ packed,                // [S, D/16] 8-byte blocks
+    uint8_t* __restrict__ dst_sfa,
+    LayoutSF layout,
+    int D, float eps) {
+    const int r = blockIdx.x;
+    const __nv_bfloat162* row2 =
+        reinterpret_cast<const __nv_bfloat162*>(x + static_cast<long>(r) * D);
+    const int D2 = D >> 1;
+    __shared__ float sh[32];
+
+    float s = 0.f;
+    for (int i = threadIdx.x; i < D2; i += blockDim.x) {
+        const __nv_bfloat162 v = row2[i];
+        s += __bfloat162float(v.x) + __bfloat162float(v.y);
+    }
+    const float mean = block_sum_dn(s, sh) / D;
+
+    float var = 0.f;
+    for (int i = threadIdx.x; i < D2; i += blockDim.x) {
+        const __nv_bfloat162 v = row2[i];
+        const float d0 = __bfloat162float(v.x) - mean;
+        const float d1 = __bfloat162float(v.y) - mean;
+        var += d0 * d0 + d1 * d1;
+    }
+    const float rstd = rsqrtf(block_sum_dn(var, sh) / D + eps);
+
+    const int n_blocks = D >> 4;
+    const int4* x4 = reinterpret_cast<const int4*>(x + static_cast<long>(r) * D);
+    const int4* s4 = reinterpret_cast<const int4*>(scale);
+    const int4* h4 = reinterpret_cast<const int4*>(shift);
+    for (int blk = threadIdx.x; blk < n_blocks; blk += blockDim.x) {
+        const int4 xr[2] = {x4[2 * blk], x4[2 * blk + 1]};
+        const __nv_bfloat16* xh = reinterpret_cast<const __nv_bfloat16*>(xr);
+        int4 sr[2], hr[2];
+        const __nv_bfloat16* sc = nullptr;
+        const __nv_bfloat16* sf = nullptr;
+        if (HAS_MOD) {
+            sr[0] = s4[2 * blk]; sr[1] = s4[2 * blk + 1];
+            hr[0] = h4[2 * blk]; hr[1] = h4[2 * blk + 1];
+            sc = reinterpret_cast<const __nv_bfloat16*>(sr);
+            sf = reinterpret_cast<const __nv_bfloat16*>(hr);
+        }
+        float vals[16];
+        float amax = 0.f;
+        #pragma unroll
+        for (int i = 0; i < 16; ++i) {
+            float normed = (__bfloat162float(xh[i]) - mean) * rstd;
+            if (HAS_MOD)
+                normed = normed * (1.0f + __bfloat162float(sc[i])) +
+                         __bfloat162float(sf[i]);
+            // bf16 rounding keeps this bit-identical to the two-step
+            // (bf16 norm kernel, then quantize) reference chain.
+            vals[i] = __bfloat162float(__float2bfloat16(normed));
+            const float a = fabsf(vals[i]);
+            if (a > amax) amax = a;
+        }
+        float desired = amax / 6.f;
+        if (desired < 1e-12f) desired = 1e-12f;
+        __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f));
+        const float bs_dq = static_cast<float>(bs_q);
+        dst_sfa[layout(r, blk * 16, 0)] = *reinterpret_cast<uint8_t*>(&bs_q);
+        const float inv_bs = 1.f / bs_dq;
+        uint2 out;
+        uint8_t* ob = reinterpret_cast<uint8_t*>(&out);
+        #pragma unroll
+        for (int p = 0; p < 8; ++p) {
+            const uint8_t lo = fp32_to_e2m1_dn(vals[2 * p] * inv_bs);
+            const uint8_t hi = fp32_to_e2m1_dn(vals[2 * p + 1] * inv_bs);
+            ob[p] = static_cast<uint8_t>(lo | (hi << 4));
+        }
+        packed[static_cast<long>(r) * n_blocks + blk] = out;
+    }
+}
+
+}  // namespace
+
+#endif  // FV_HAVE_CUTLASS
+
+static int check_dn_args(const void* x, void* packed, int dim) {
+    if (dim % 16 != 0 || dim % 2 != 0) return -1;
+    if ((reinterpret_cast<uintptr_t>(x) & 15) ||
+        (reinterpret_cast<uintptr_t>(packed) & 7)) return -1;
+    return 0;
+}
+
+int ada_layer_norm_fp4_sfa_bf16(
+    const void* x, const void* scale, const void* shift,
+    void* packed, void* sfa,
+    int seq_len, int dim, float eps, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+    if (check_dn_args(x, packed, dim) != 0) return -1;
+    if ((reinterpret_cast<uintptr_t>(scale) & 15) ||
+        (reinterpret_cast<uintptr_t>(shift) & 15)) return -1;
+    auto shape = cute::make_shape(seq_len, 1, dim, 1);
+    auto layout = CfgDN::tile_atom_to_shape_SFA(shape);
+    dit_norm_fp4_sfa_kernel<true><<<seq_len, 128, 0, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(x),
+        reinterpret_cast<const __nv_bfloat16*>(scale),
+        reinterpret_cast<const __nv_bfloat16*>(shift),
+        reinterpret_cast<uint2*>(packed),
+        reinterpret_cast<uint8_t*>(sfa),
+        layout, dim, eps);
+    const cudaError_t e = cudaGetLastError();
+    return (e == cudaSuccess) ? 0 : -static_cast<int>(e);
+#else
+    (void)x; (void)scale; (void)shift; (void)packed; (void)sfa;
+    (void)seq_len; (void)dim; (void)eps; (void)stream;
+    return -2;
+#endif
+}
+
+int layer_norm_no_affine_fp4_sfa_bf16(
+    const void* x, void* packed, void* sfa,
+    int seq_len, int dim, float eps, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+    if (check_dn_args(x, packed, dim) != 0) return -1;
+    auto shape = cute::make_shape(seq_len, 1, dim, 1);
+    auto layout = CfgDN::tile_atom_to_shape_SFA(shape);
+    dit_norm_fp4_sfa_kernel<false><<<seq_len, 128, 0, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(x),
+        nullptr, nullptr,
+        reinterpret_cast<uint2*>(packed),
+        reinterpret_cast<uint8_t*>(sfa),
+        layout, dim, eps);
+    const cudaError_t e = cudaGetLastError();
+    return (e == cudaSuccess) ? 0 : -static_cast<int>(e);
+#else
+    (void)x; (void)packed; (void)sfa;
+    (void)seq_len; (void)dim; (void)eps; (void)stream;
+    return -2;
+#endif
+}
+
+}  // namespace fused_fp4
+}  // namespace flash_rt

csrc/dit_norm_fp4_sfa.cuh

@@ -0,0 +1,33 @@
+// ============================================================================
+//  FlashRT — fused DiT norms (bf16 in) + NVFP4 quantize + SFA write.
+//
+//  GR00T N1.7 DiT front-ends for NVFP4 GEMMs: the AdaLN-modulated norm1
+//  (QKV / cross-Q input) and the pre-FFN no-affine LayerNorm emit packed
+//  e2m1 + tile-interleaved UE4M3 scales directly, replacing a bf16 norm
+//  kernel followed by a separate quantize kernel. The normalized value is
+//  rounded through bf16 before quantization so the output is bit-identical
+//  to the two-step reference chain.
+//
+//  Additive: new symbols only.
+// ============================================================================
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace fused_fp4 {
+
+// packed/sfa = quantize(LN_no_affine(x[row]) * (1 + scale) + shift)
+// x: bf16 [S, D]; scale/shift: bf16 [D] (per-layer AdaLN modulators).
+int ada_layer_norm_fp4_sfa_bf16(
+    const void* x, const void* scale, const void* shift,
+    void* packed, void* sfa,
+    int seq_len, int dim, float eps, cudaStream_t stream);
+
+// packed/sfa = quantize(LN_no_affine(x[row]))
+int layer_norm_no_affine_fp4_sfa_bf16(
+    const void* x, void* packed, void* sfa,
+    int seq_len, int dim, float eps, cudaStream_t stream);
+
+}  // namespace fused_fp4
+}  // namespace flash_rt

csrc/sm110_fp4_dispatch.cu

@@ -0,0 +1,20 @@
+#include "sm110_fp4_dispatch.cuh"
+
+#include "dit_norm_fp4_sfa.cuh"
+
+namespace flash_rt::adaln_producers::hub {
+namespace {
+
+struct Registration {
+  Registration() {
+    ada_layer_norm_fp4_dispatch =
+        &flash_rt::fused_fp4::ada_layer_norm_fp4_sfa_bf16;
+    layer_norm_fp4_dispatch =
+        &flash_rt::fused_fp4::layer_norm_no_affine_fp4_sfa_bf16;
+  }
+};
+
+Registration registration;
+
+}  // namespace
+}  // namespace flash_rt::adaln_producers::hub

csrc/sm110_fp4_dispatch.cuh

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cuda_runtime_api.h>
+
+namespace flash_rt::adaln_producers::hub {
+
+using AdaLayerNormFp4Dispatch = int (*)(
+    const void* x, const void* scale, const void* shift,
+    void* packed, void* sfa, int rows, int dim, float eps,
+    cudaStream_t stream);
+
+using LayerNormFp4Dispatch = int (*)(
+    const void* x, void* packed, void* sfa, int rows, int dim,
+    float eps, cudaStream_t stream);
+
+extern AdaLayerNormFp4Dispatch ada_layer_norm_fp4_dispatch;
+extern LayerNormFp4Dispatch layer_norm_fp4_dispatch;
+
+}  // namespace flash_rt::adaln_producers::hub

examples/README.md

@@ -0,0 +1,10 @@
+# Examples
+
+`basic_usage.py` shows the public Kernel Hub loading path:
+
+```bash
+python adaptive-layernorm-producers/examples/basic_usage.py
+```
+
+It allocates static output buffers for the FP8 path and uses the package helper
+to allocate the NVFP4 packed activation and swizzled scale-factor buffers.

examples/basic_usage.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import torch
+from kernels import get_kernel
+
+
+def main() -> None:
+    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)
+
+    out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, out=out)
+
+    packed, sf = ops.ada_layer_norm_quant_nvfp4_swizzled_bf16(x, scale, shift)
+    print(out.shape, out.dtype)
+    print(packed.shape, packed.dtype, sf.numel())
+
+
+if __name__ == "__main__":
+    main()

flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for FlashRT adaptive LayerNorm producer kernels";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernels/870e825d881664e39f9287a27a74ef63ff3c545e";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

tests/README.md

@@ -0,0 +1,16 @@
+# Tests
+
+Run correctness tests from the repository root:
+
+```bash
+python adaptive-layernorm-producers/tests/test_adaptive_layernorm_producers.py --backend source --mode full
+```
+
+After Kernel Hub build/publish, run installed-artifact validation:
+
+```bash
+python adaptive-layernorm-producers/tests/test_adaptive_layernorm_producers.py --backend installed --mode full
+```
+
+The source mode compiles the package locally with `torch.utils.cpp_extension`.
+The installed mode imports the package exactly as exposed by Kernel Hub.

tests/native_reference/dit_bf16.cu

@@ -0,0 +1,315 @@
+// DiT bf16 helpers — Phase 5a-2 (R0: pure additions, no edits to existing
+// fp16 kernels). Mirrors layer_norm_no_affine_fp16 / ada_layer_norm_fp16 /
+// add_bias_fp16 but with __nv_bfloat162 vectorized math, plus cast helpers
+// for fp16 ↔ bf16 boundaries.
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cstdint>
+
+// ────────────────────────────────────────────────────────────────────
+// LayerNorm (no affine) — bf16
+// ────────────────────────────────────────────────────────────────────
+__global__ void layer_norm_no_affine_bf16_kernel(const __nv_bfloat16* __restrict__ x,
+                                                   __nv_bfloat16* __restrict__ out,
+                                                   int dim, float eps) {
+    int row = blockIdx.x;
+    const __nv_bfloat162* x2 = reinterpret_cast<const __nv_bfloat162*>(x + row * dim);
+    __nv_bfloat162* out2 = reinterpret_cast<__nv_bfloat162*>(out + row * dim);
+    int dim2 = dim >> 1;
+
+    extern __shared__ float shared[];
+    float local_sum = 0.0f;
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 val = x2[i];
+        local_sum += __bfloat162float(val.x) + __bfloat162float(val.y);
+    }
+    float val = local_sum;
+    for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
+    if (!lane) shared[wid] = val;
+    __syncthreads();
+    if (!wid) { val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0;
+                for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o); }
+    __syncthreads();
+    if (!threadIdx.x) shared[0] = val;
+    __syncthreads();
+    float mean = shared[0] / dim;
+
+    float local_var = 0.0f;
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 v = x2[i];
+        float d0 = __bfloat162float(v.x) - mean, d1 = __bfloat162float(v.y) - mean;
+        local_var += d0 * d0 + d1 * d1;
+    }
+    val = local_var;
+    for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    if (!lane) shared[wid] = val;
+    __syncthreads();
+    if (!wid) { val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0;
+                for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o); }
+    __syncthreads();
+    if (!threadIdx.x) shared[0] = val;
+    __syncthreads();
+    float inv_std = rsqrtf(shared[0] / dim + eps);
+
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 xv = x2[i];
+        float v0 = (__bfloat162float(xv.x) - mean) * inv_std;
+        float v1 = (__bfloat162float(xv.y) - mean) * inv_std;
+        out2[i] = __floats2bfloat162_rn(v0, v1);
+    }
+}
+
+void layer_norm_no_affine_bf16(const __nv_bfloat16* x, __nv_bfloat16* out,
+                                int seq_len, int dim, float eps,
+                                cudaStream_t stream) {
+    layer_norm_no_affine_bf16_kernel<<<seq_len, 256, 256 * sizeof(float), stream>>>(
+        x, out, dim, eps);
+}
+
+__global__ void layer_norm_no_affine_fp8_static_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    __nv_fp8_e4m3* __restrict__ out,
+    const float* __restrict__ d_scale,
+    int dim, float eps) {
+    int row = blockIdx.x;
+    const __nv_bfloat162* x2 =
+        reinterpret_cast<const __nv_bfloat162*>(x + (long long)row * dim);
+    __nv_fp8_e4m3* out_row = out + (long long)row * dim;
+    int dim2 = dim >> 1;
+
+    extern __shared__ float shared[];
+    float local_sum = 0.0f;
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 val = x2[i];
+        local_sum += __bfloat162float(val.x) + __bfloat162float(val.y);
+    }
+    float val = local_sum;
+    for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
+    if (!lane) shared[wid] = val;
+    __syncthreads();
+    if (!wid) {
+        val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
+        for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    }
+    __syncthreads();
+    if (!threadIdx.x) shared[0] = val;
+    __syncthreads();
+    float mean = shared[0] / dim;
+
+    float local_var = 0.0f;
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 v = x2[i];
+        float d0 = __bfloat162float(v.x) - mean;
+        float d1 = __bfloat162float(v.y) - mean;
+        local_var += d0 * d0 + d1 * d1;
+    }
+    val = local_var;
+    for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    if (!lane) shared[wid] = val;
+    __syncthreads();
+    if (!wid) {
+        val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
+        for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    }
+    __syncthreads();
+    if (!threadIdx.x) shared[0] = val;
+    __syncthreads();
+    float inv_std = rsqrtf(shared[0] / dim + eps);
+    float inv_a = 1.0f / fmaxf(*d_scale, 1e-12f);
+
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 xv = x2[i];
+        float v0 = (__bfloat162float(xv.x) - mean) * inv_std;
+        float v1 = (__bfloat162float(xv.y) - mean) * inv_std;
+        v0 = __bfloat162float(__float2bfloat16(v0)) * inv_a;
+        v1 = __bfloat162float(__float2bfloat16(v1)) * inv_a;
+        __nv_fp8_e4m3 pair[2];
+        pair[0] = __nv_fp8_e4m3(fminf(fmaxf(v0, -448.0f), 448.0f));
+        pair[1] = __nv_fp8_e4m3(fminf(fmaxf(v1, -448.0f), 448.0f));
+        *reinterpret_cast<uint16_t*>(out_row + 2 * i) =
+            *reinterpret_cast<uint16_t*>(pair);
+    }
+}
+
+void layer_norm_no_affine_fp8_static_bf16(
+    const __nv_bfloat16* x, __nv_fp8_e4m3* out, const float* d_scale,
+    int seq_len, int dim, float eps, cudaStream_t stream) {
+    layer_norm_no_affine_fp8_static_bf16_kernel
+        <<<seq_len, 256, 256 * sizeof(float), stream>>>(
+            x, out, d_scale, dim, eps);
+}
+
+// ────────────────────────────────────────────────────────────────────
+// AdaLayerNorm — bf16
+// out = LN(x, no_affine) * (1 + scale) + shift
+// ────────────────────────────────────────────────────────────────────
+__global__ void ada_layer_norm_bf16_kernel(const __nv_bfloat16* __restrict__ x,
+                                            const __nv_bfloat16* __restrict__ scale,
+                                            const __nv_bfloat16* __restrict__ shift,
+                                            __nv_bfloat16* __restrict__ out,
+                                            int dim, float eps) {
+    int row = blockIdx.x;
+    const __nv_bfloat162* x2 = reinterpret_cast<const __nv_bfloat162*>(x + row * dim);
+    const __nv_bfloat162* sc2 = reinterpret_cast<const __nv_bfloat162*>(scale);
+    const __nv_bfloat162* sh2 = reinterpret_cast<const __nv_bfloat162*>(shift);
+    __nv_bfloat162* out2 = reinterpret_cast<__nv_bfloat162*>(out + row * dim);
+    int dim2 = dim >> 1;
+
+    extern __shared__ float shared[];
+    float local_sum = 0.0f;
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 val = x2[i];
+        local_sum += __bfloat162float(val.x) + __bfloat162float(val.y);
+    }
+    float val = local_sum;
+    for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
+    if (!lane) shared[wid] = val;
+    __syncthreads();
+    if (!wid) { val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0;
+                for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o); }
+    __syncthreads(); if (!threadIdx.x) shared[0] = val; __syncthreads();
+    float mean = shared[0] / dim;
+
+    float local_var = 0.0f;
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 v = x2[i];
+        float d0 = __bfloat162float(v.x) - mean, d1 = __bfloat162float(v.y) - mean;
+        local_var += d0 * d0 + d1 * d1;
+    }
+    val = local_var;
+    for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    if (!lane) shared[wid] = val;
+    __syncthreads();
+    if (!wid) { val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0;
+                for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o); }
+    __syncthreads(); if (!threadIdx.x) shared[0] = val; __syncthreads();
+    float inv_std = rsqrtf(shared[0] / dim + eps);
+
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 xv = x2[i], sv = sc2[i], hv = sh2[i];
+        float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+        float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+        float v0 = n0 * (1.0f + __bfloat162float(sv.x)) + __bfloat162float(hv.x);
+        float v1 = n1 * (1.0f + __bfloat162float(sv.y)) + __bfloat162float(hv.y);
+        out2[i] = __floats2bfloat162_rn(v0, v1);
+    }
+}
+
+void ada_layer_norm_bf16(const __nv_bfloat16* x,
+                          const __nv_bfloat16* scale, const __nv_bfloat16* shift,
+                          __nv_bfloat16* out, int seq_len, int dim, float eps,
+                          cudaStream_t stream) {
+    ada_layer_norm_bf16_kernel<<<seq_len, 256, 256 * sizeof(float), stream>>>(
+        x, scale, shift, out, dim, eps);
+}
+
+// Per-row AdaLayerNorm — bf16
+// out[row] = LN(x[row], no_affine) * (1 + scale[row]) + shift[row]
+__global__ void ada_layer_norm_bf16_per_token_kernel(
+                                            const __nv_bfloat16* __restrict__ x,
+                                            const __nv_bfloat16* __restrict__ scale,
+                                            const __nv_bfloat16* __restrict__ shift,
+                                            __nv_bfloat16* __restrict__ out,
+                                            int dim, float eps) {
+    int row = blockIdx.x;
+    const __nv_bfloat162* x2 = reinterpret_cast<const __nv_bfloat162*>(x + row * dim);
+    const __nv_bfloat162* sc2 = reinterpret_cast<const __nv_bfloat162*>(scale + row * dim);
+    const __nv_bfloat162* sh2 = reinterpret_cast<const __nv_bfloat162*>(shift + row * dim);
+    __nv_bfloat162* out2 = reinterpret_cast<__nv_bfloat162*>(out + row * dim);
+    int dim2 = dim >> 1;
+
+    extern __shared__ float shared[];
+    float local_sum = 0.0f;
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 val = x2[i];
+        local_sum += __bfloat162float(val.x) + __bfloat162float(val.y);
+    }
+    float val = local_sum;
+    for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
+    if (!lane) shared[wid] = val;
+    __syncthreads();
+    if (!wid) { val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0;
+                for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o); }
+    __syncthreads(); if (!threadIdx.x) shared[0] = val; __syncthreads();
+    float mean = shared[0] / dim;
+
+    float local_var = 0.0f;
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 v = x2[i];
+        float d0 = __bfloat162float(v.x) - mean, d1 = __bfloat162float(v.y) - mean;
+        local_var += d0 * d0 + d1 * d1;
+    }
+    val = local_var;
+    for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o);
+    if (!lane) shared[wid] = val;
+    __syncthreads();
+    if (!wid) { val = (lane < (blockDim.x >> 5)) ? shared[lane] : 0;
+                for (int o = 16; o > 0; o >>= 1) val += __shfl_xor_sync(0xffffffff, val, o); }
+    __syncthreads(); if (!threadIdx.x) shared[0] = val; __syncthreads();
+    float inv_std = rsqrtf(shared[0] / dim + eps);
+
+    for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+        __nv_bfloat162 xv = x2[i], sv = sc2[i], hv = sh2[i];
+        float n0 = (__bfloat162float(xv.x) - mean) * inv_std;
+        float n1 = (__bfloat162float(xv.y) - mean) * inv_std;
+        float v0 = n0 * (1.0f + __bfloat162float(sv.x)) + __bfloat162float(hv.x);
+        float v1 = n1 * (1.0f + __bfloat162float(sv.y)) + __bfloat162float(hv.y);
+        out2[i] = __floats2bfloat162_rn(v0, v1);
+    }
+}
+
+void ada_layer_norm_bf16_per_token(
+                          const __nv_bfloat16* x,
+                          const __nv_bfloat16* scale, const __nv_bfloat16* shift,
+                          __nv_bfloat16* out, int seq_len, int dim, float eps,
+                          cudaStream_t stream) {
+    ada_layer_norm_bf16_per_token_kernel
+        <<<seq_len, 256, 256 * sizeof(float), stream>>>(
+            x, scale, shift, out, dim, eps);
+}
+
+// ────────────────────────────────────────────────────────────────────
+// Bias add: x[i] += b[i % D] — bf16
+// ────────────────────────────────────────────────────────────────────
+__global__ void add_bias_bf16_kernel(__nv_bfloat16* x, const __nv_bfloat16* b,
+                                      int S, int D) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < S * D) {
+        x[i] = __float2bfloat16(__bfloat162float(x[i]) + __bfloat162float(b[i % D]));
+    }
+}
+
+void add_bias_bf16(__nv_bfloat16* x, const __nv_bfloat16* b,
+                    int S, int D, cudaStream_t stream) {
+    add_bias_bf16_kernel<<<(S * D + 255) / 256, 256, 0, stream>>>(x, b, S, D);
+}
+
+// ────────────────────────────────────────────────────────────────────
+// Cast fp16 ↔ bf16 (DiT input/output boundary)
+// ────────────────────────────────────────────────────────────────────
+__global__ void cast_fp16_to_bf16_kernel(const __half* in, __nv_bfloat16* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = __float2bfloat16(__half2float(in[i]));
+}
+
+void cast_fp16_to_bf16(const __half* in, __nv_bfloat16* out, int n,
+                        cudaStream_t stream) {
+    cast_fp16_to_bf16_kernel<<<(n + 255) / 256, 256, 0, stream>>>(in, out, n);
+}
+
+__global__ void cast_bf16_to_fp16_kernel(const __nv_bfloat16* in, __half* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = __float2half(__bfloat162float(in[i]));
+}
+
+void cast_bf16_to_fp16(const __nv_bfloat16* in, __half* out, int n,
+                        cudaStream_t stream) {
+    cast_bf16_to_fp16_kernel<<<(n + 255) / 256, 256, 0, stream>>>(in, out, n);
+}

tests/sm110_reference_binding.cpp

@@ -0,0 +1,64 @@
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_bf16.h>
+
+#include "quantize_fp4_sfa_bf16.cuh"
+
+void ada_layer_norm_bf16(
+    const __nv_bfloat16* x, const __nv_bfloat16* scale,
+    const __nv_bfloat16* shift, __nv_bfloat16* out,
+    int seq_len, int dim, float eps, cudaStream_t stream);
+void layer_norm_no_affine_bf16(
+    const __nv_bfloat16* x, __nv_bfloat16* out, int seq_len, int dim, float eps,
+    cudaStream_t stream);
+
+namespace {
+
+void reference_adaln_fp4(
+    const torch::Tensor& x, const torch::Tensor& scale,
+    const torch::Tensor& shift, double eps, torch::Tensor& packed,
+    torch::Tensor& sf) {
+  c10::cuda::CUDAGuard guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  auto norm = torch::empty_like(x);
+  ada_layer_norm_bf16(
+      static_cast<const __nv_bfloat16*>(x.data_ptr()),
+      static_cast<const __nv_bfloat16*>(scale.data_ptr()),
+      static_cast<const __nv_bfloat16*>(shift.data_ptr()),
+      static_cast<__nv_bfloat16*>(norm.data_ptr()),
+      static_cast<int>(x.size(0)), static_cast<int>(x.size(1)),
+      static_cast<float>(eps), stream);
+  const int rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_bf16_vec(
+      norm.data_ptr(), packed.data_ptr(), sf.data_ptr(),
+      static_cast<int>(x.size(0)), static_cast<int>(x.size(1)), false, stream);
+  TORCH_CHECK(rc == 0, "SM110 staged AdaLN-FP4 reference failed: rc=", rc);
+}
+
+void reference_ln_fp4(
+    const torch::Tensor& x, double eps, torch::Tensor& packed,
+    torch::Tensor& sf) {
+  c10::cuda::CUDAGuard guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  auto norm = torch::empty_like(x);
+  layer_norm_no_affine_bf16(
+      static_cast<const __nv_bfloat16*>(x.data_ptr()),
+      static_cast<__nv_bfloat16*>(norm.data_ptr()),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)), static_cast<float>(eps), stream);
+  const int rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_bf16_vec(
+      norm.data_ptr(), packed.data_ptr(), sf.data_ptr(),
+      static_cast<int>(x.size(0)), static_cast<int>(x.size(1)), false, stream);
+  TORCH_CHECK(rc == 0, "SM110 staged LayerNorm-FP4 reference failed: rc=", rc);
+}
+
+}  // namespace
+
+TORCH_LIBRARY_FRAGMENT(adaptive_layernorm_producers_test, ops) {
+  ops.def("_reference_adaln_fp4(Tensor x, Tensor scale, Tensor shift, float eps, Tensor! packed, Tensor! sf) -> ()");
+  ops.def("_reference_ln_fp4(Tensor x, float eps, Tensor! packed, Tensor! sf) -> ()");
+  ops.impl("_reference_adaln_fp4", torch::kCUDA, &reference_adaln_fp4);
+  ops.impl("_reference_ln_fp4", torch::kCUDA, &reference_ln_fp4);
+}

tests/test_adaptive_layernorm_producers.py

@@ -0,0 +1,597 @@
+#!/usr/bin/env python3
+"""Correctness tests for adaptive-layernorm-producers."""
+
+from __future__ import annotations
+
+import argparse
+import ctypes
+import ctypes.util
+import importlib
+import math
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "adaptive-layernorm-producers"
+DEFAULT_CUTLASS_INCLUDE = ROOT.parent / "official" / "FlashRT" / "third_party" / "cutlass" / "include"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
+FP8_MAX = 448.0
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def ada_layer_norm_quant_fp8_bf16(self, x, scale, shift, act_scale, eps=1e-5, out=None):
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn) if out is None else out
+        self._ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, float(eps), out)
+        return out
+
+    def ada_layer_norm_quant_fp8_ptok_bf16(self, x, scale, shift, act_scale, eps=1e-5, out=None):
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn) if out is None else out
+        self._ops.ada_layer_norm_quant_fp8_ptok_bf16(x, scale, shift, act_scale, float(eps), out)
+        return out
+
+    def ada_layer_norm_quant_fp8_ptok_table_bf16(
+        self, x, temb, table, act_scale, shift_idx, scale_idx, eps=1e-5, out=None
+    ):
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn) if out is None else out
+        self._ops.ada_layer_norm_quant_fp8_ptok_table_bf16(
+            x, temb, table, act_scale, int(shift_idx), int(scale_idx), float(eps), out
+        )
+        return out
+
+    def ada_layer_norm_quant_fp8_modfp8_bf16(
+        self, x, scale_fp8, shift_fp8, scale_deq, shift_deq, act_scale, eps=1e-5, out=None
+    ):
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn) if out is None else out
+        self._ops.ada_layer_norm_quant_fp8_modfp8_bf16(
+            x, scale_fp8, shift_fp8, scale_deq, shift_deq, act_scale, float(eps), out
+        )
+        return out
+
+    def awq_ada_layer_norm_quant_fp8_bf16(self, x, scale, shift, inv_s, act_scale, eps=1e-5, out=None):
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn) if out is None else out
+        self._ops.awq_ada_layer_norm_quant_fp8_bf16(x, scale, shift, inv_s, act_scale, float(eps), out)
+        return out
+
+    def ada_layer_norm_quant_nvfp4_swizzled_bf16(self, x, scale, shift, eps=1e-5, packed=None, sf_swizzled=None):
+        packed, sf_swizzled = _nvfp4_out(x, packed, sf_swizzled)
+        self._ops.ada_layer_norm_quant_nvfp4_swizzled_bf16(x, scale, shift, float(eps), packed, sf_swizzled)
+        return packed, sf_swizzled
+
+    def ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(
+        self, x, scale_fp8, shift_fp8, scale_deq, shift_deq, eps=1e-5, packed=None, sf_swizzled=None
+    ):
+        packed, sf_swizzled = _nvfp4_out(x, packed, sf_swizzled)
+        self._ops.ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(
+            x, scale_fp8, shift_fp8, scale_deq, shift_deq, float(eps), packed, sf_swizzled
+        )
+        return packed, sf_swizzled
+
+    def layer_norm_no_affine_quant_fp8_static_bf16(self, x, act_scale, eps=1e-5, out=None):
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn) if out is None else out
+        self._ops.layer_norm_no_affine_quant_fp8_static_bf16(x, act_scale, float(eps), out)
+        return out
+
+    def layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+        self, x, eps=1e-5, packed=None, sf_swizzled=None
+    ):
+        packed, sf_swizzled = _nvfp4_out(x, packed, sf_swizzled)
+        self._ops.layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+            x, float(eps), packed, sf_swizzled
+        )
+        return packed, sf_swizzled
+
+    def adaln_modulation6_bf16(self, adaln_params, layer_modulation, out=None):
+        expected = (
+            adaln_params.shape[0],
+            adaln_params.shape[1],
+            adaln_params.shape[3],
+        )
+        if out is None:
+            out = tuple(
+                torch.empty(expected, device=adaln_params.device, dtype=torch.bfloat16)
+                for _ in range(6)
+            )
+        self._ops.adaln_modulation6_bf16(
+            adaln_params, layer_modulation, *out
+        )
+        return out
+
+    def reference_adaln_fp4(self, x, scale, shift, eps, packed, sf):
+        self._ops._reference_adaln_fp4(x, scale, shift, float(eps), packed, sf)
+        return packed, sf
+
+    def reference_ln_fp4(self, x, eps, packed, sf):
+        self._ops._reference_ln_fp4(x, float(eps), packed, sf)
+        return packed, sf
+
+
+def _preload_cublaslt() -> None:
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
+            return
+    library = ctypes.util.find_library("cublasLt")
+    if library:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+
+
+def _current_arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    cutlass_include = Path(
+        os.environ.get("FLASHRT_CUTLASS_INCLUDE", str(DEFAULT_CUTLASS_INCLUDE))
+    )
+    if not cutlass_include.is_dir():
+        raise RuntimeError(f"missing CUTLASS include path: {cutlass_include}")
+    _preload_cublaslt()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
+    namespace = "adaptive_layernorm_producers_test"
+    sm110_sources = []
+    if torch.cuda.get_device_capability(0) == (11, 0):
+        sm110_sources = [
+            str(PACKAGE / "csrc" / "dit_norm_fp4_sfa.cu"),
+            str(PACKAGE / "csrc" / "sm110_fp4_dispatch.cu"),
+            str(PACKAGE / "tests" / "native_reference" / "dit_bf16.cu"),
+            str(ROOT / "fp4-gemm" / "csrc" / "quantize" / "quantize_fp4_sfa_bf16.cu"),
+            str(PACKAGE / "tests" / "sm110_reference_binding.cpp"),
+        ]
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "ada_layer_norm_fp8.cu"),
+            str(PACKAGE / "csrc" / "ada_layer_norm_fp8_ptok.cu"),
+            str(PACKAGE / "csrc" / "dit_layer_norm_fp8.cu"),
+            str(PACKAGE / "csrc" / "adaln_modulation6.cu"),
+            *sm110_sources,
+        ],
+        extra_include_paths=[
+            str(PACKAGE / "csrc"), str(cutlass_include),
+            str(ROOT / "fp4-gemm" / "csrc" / "quantize"),
+            str(REGISTRATION_INCLUDE),
+        ],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=[
+            "-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL",
+            "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
+        ],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_sm110_reference_ops() -> SourceOps:
+    """Load the package-local staged native oracle for installed tests."""
+    from torch.utils.cpp_extension import load
+
+    cutlass_include = Path(
+        os.environ.get("FLASHRT_CUTLASS_INCLUDE", str(DEFAULT_CUTLASS_INCLUDE))
+    )
+    if not cutlass_include.is_dir():
+        raise RuntimeError(f"missing CUTLASS include path: {cutlass_include}")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
+    load(
+        name="adaptive_layernorm_producers_reference",
+        sources=[
+            str(PACKAGE / "tests" / "native_reference" / "dit_bf16.cu"),
+            str(ROOT / "fp4-gemm" / "csrc" / "quantize" / "quantize_fp4_sfa_bf16.cu"),
+            str(PACKAGE / "tests" / "sm110_reference_binding.cpp"),
+        ],
+        extra_include_paths=[
+            str(ROOT / "fp4-gemm" / "csrc" / "quantize"),
+            str(REGISTRATION_INCLUDE),
+            str(cutlass_include),
+        ],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=[
+            "-O3", "--expt-relaxed-constexpr", "--expt-extended-lambda",
+            "-DCUDA_KERNEL", "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
+        ],
+        is_python_module=False,
+        verbose=False,
+    )
+    return SourceOps("adaptive_layernorm_producers_test")
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("adaptive_layernorm_producers")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def swizzled_sf_size(rows: int, dim: int) -> int:
+    n_blocks = dim // 16
+    return ((rows + 127) // 128) * ((n_blocks + 3) // 4) * 128 * 64
+
+
+def _nvfp4_out(x: torch.Tensor, packed=None, sf_swizzled=None):
+    if packed is None:
+        packed = torch.empty((x.shape[0], x.shape[1] // 2), device=x.device, dtype=torch.uint8)
+    if sf_swizzled is None:
+        sf_swizzled = torch.zeros((swizzled_sf_size(x.shape[0], x.shape[1]),), device=x.device, dtype=torch.uint8)
+    return packed, sf_swizzled
+
+
+def make_case(rows: int, dim: int):
+    x = (torch.randn((rows, dim), device="cuda", dtype=torch.float32) * 0.75).to(torch.bfloat16).contiguous()
+    scale = (torch.randn((dim,), device="cuda", dtype=torch.float32) * 0.08).to(torch.bfloat16).contiguous()
+    shift = (torch.randn((dim,), device="cuda", dtype=torch.float32) * 0.08).to(torch.bfloat16).contiguous()
+    inv_s = (1.0 + torch.rand((dim,), device="cuda", dtype=torch.float32) * 0.2).to(torch.bfloat16).contiguous()
+    act_scale = torch.tensor([0.025], device="cuda", dtype=torch.float32)
+    scale_deq = torch.tensor([0.0075], device="cuda", dtype=torch.float32)
+    shift_deq = torch.tensor([0.0075], device="cuda", dtype=torch.float32)
+    scale_fp8 = torch.clamp(scale.float() / scale_deq, -FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn).contiguous()
+    shift_fp8 = torch.clamp(shift.float() / shift_deq, -FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn).contiguous()
+    return x, scale, shift, inv_s, act_scale, scale_fp8, shift_fp8, scale_deq, shift_deq
+
+
+def ref_layer_norm_no_affine(x: torch.Tensor, eps: float) -> torch.Tensor:
+    return ref_layer_norm_no_affine_f32(x, eps).to(torch.bfloat16)
+
+
+def ref_layer_norm_no_affine_f32(x: torch.Tensor, eps: float) -> torch.Tensor:
+    xf = x.float()
+    mean = xf.mean(dim=-1, keepdim=True)
+    var = ((xf - mean) * (xf - mean)).mean(dim=-1, keepdim=True)
+    return (xf - mean) * torch.rsqrt(var + eps)
+
+
+def ref_adaln(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float) -> torch.Tensor:
+    norm = ref_layer_norm_no_affine_f32(x, eps)
+    return (norm * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
+
+
+def ref_adaln_float_mod(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float) -> torch.Tensor:
+    norm = ref_layer_norm_no_affine_f32(x, eps)
+    return (norm * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
+
+
+def quant_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return torch.clamp(x.float() / scale.float().reshape(()), -FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+
+
+def f32_to_fp4_e2m1(x: float) -> int:
+    if x == 0.0:
+        return 0
+    sign = 0x8 if x < 0 else 0
+    ax = abs(x)
+    if ax < 0.25:
+        return sign
+    if ax < 0.75:
+        return sign | 0x1
+    if ax < 1.5:
+        return sign | 0x2
+    if ax < 3.0:
+        return sign | 0x3
+    if ax < 5.0:
+        return sign | 0x6
+    return sign | 0x7
+
+
+def f32_to_fp4_e2m1_sm110(x: float) -> int:
+    """Match #163's native SM110 FP4 producer bucket boundaries."""
+    sign = 0x8 if x < 0 else 0
+    ax = abs(x)
+    if ax <= 0.25:
+        mantissa = 0
+    elif ax <= 0.75:
+        mantissa = 1
+    elif ax <= 1.25:
+        mantissa = 2
+    elif ax <= 1.75:
+        mantissa = 3
+    elif ax <= 2.5:
+        mantissa = 4
+    elif ax <= 3.5:
+        mantissa = 5
+    elif ax <= 5.0:
+        mantissa = 6
+    else:
+        mantissa = 7
+    return sign | mantissa
+
+
+def f32_to_ue4m3_ceil(x: float) -> int:
+    if x <= 0.0:
+        return 0
+    if x < 0.0009765625:
+        return 1
+    mant, exp = math.frexp(x)
+    exp -= 1
+    frac = mant * 2.0 - 1.0
+    mantissa = math.ceil(frac * 8.0)
+    if mantissa >= 8:
+        mantissa = 0
+        exp += 1
+    biased_exp = exp + 7
+    if biased_exp <= 0:
+        return 1
+    if biased_exp >= 15:
+        return 0x7F
+    return (biased_exp << 3) | mantissa
+
+
+def ue4m3_to_f32(v: int) -> float:
+    if v == 0:
+        return 0.0
+    exp = ((v >> 3) & 0xF) - 7
+    mant = v & 0x7
+    return math.ldexp(1.0 + mant / 8.0, exp)
+
+
+def ref_nvfp4(
+    modulated: torch.Tensor, *, sm110_contract: bool = False
+) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, dim = modulated.shape
+    assert dim % 16 == 0
+    packed = torch.zeros((rows, dim // 2), dtype=torch.uint8)
+    sf = torch.zeros((swizzled_sf_size(rows, dim),), dtype=torch.uint8)
+    cpu = modulated.float().cpu()
+    n_blocks = dim // 16
+    n_col_blocks = (n_blocks + 3) // 4
+    for row in range(rows):
+        block_scales: list[float] = []
+        for block in range(n_blocks):
+            amax = float(cpu[row, block * 16 : (block + 1) * 16].abs().max().item())
+            desired = max(amax / 6.0, 1e-12)
+            if sm110_contract:
+                ue = int(
+                    torch.tensor(desired, dtype=torch.float32)
+                    .to(torch.float8_e4m3fn)
+                    .view(torch.uint8)
+                    .item()
+                )
+            else:
+                ue = f32_to_ue4m3_ceil(desired)
+            rb = row // 128
+            ri = row % 128
+            cb = block // 4
+            ci = block % 4
+            out_idx = (rb * n_col_blocks + cb) * 512 + (ri % 32) * 16 + (ri // 32) * 4 + ci
+            sf[out_idx] = ue
+            block_scales.append(ue4m3_to_f32(ue))
+        for p in range(dim // 2):
+            i = p * 2
+            s0 = block_scales[i // 16]
+            s1 = block_scales[(i + 1) // 16]
+            convert = f32_to_fp4_e2m1_sm110 if sm110_contract else f32_to_fp4_e2m1
+            lo = convert(float(cpu[row, i].item()) / s0 if s0 > 0 else 0.0)
+            hi = convert(float(cpu[row, i + 1].item()) / s1 if s1 > 0 else 0.0)
+            packed[row, p] = (hi << 4) | (lo & 0x0F)
+    return packed, sf
+
+
+def assert_exact(name: str, got: torch.Tensor, expected: torch.Tensor) -> None:
+    got_cpu = got.detach().cpu()
+    exp_cpu = expected.detach().cpu()
+    if not torch.equal(got_cpu, exp_cpu):
+        diff = (got_cpu.to(torch.int16) - exp_cpu.to(torch.int16)).abs()
+        raise AssertionError(f"{name} mismatch: nonzero={int((diff != 0).sum())} max={int(diff.max())}")
+    print(f"PASS {name}: exact")
+
+
+def assert_fp8_contract(name: str, got: torch.Tensor, expected: torch.Tensor) -> None:
+    diff = (got.float() - expected.float()).abs()
+    nonzero = int((diff != 0).sum().item())
+    max_abs = float(diff.max().item())
+    sorted_diff = diff.flatten().sort().values
+    p99_abs = float(sorted_diff[min(sorted_diff.numel() - 1, math.ceil(0.99 * sorted_diff.numel()) - 1)].item())
+    cosine = float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item())
+    max_nonzero = max(8, diff.numel() // 100000)
+    if p99_abs != 0.0 or nonzero > max_nonzero:
+        raise AssertionError(
+            f"{name} mismatch: nonzero={nonzero}/{diff.numel()} max_abs={max_abs} "
+            f"p99_abs={p99_abs} cosine={cosine}"
+        )
+    status = "exact" if nonzero == 0 else "fp8-boundary"
+    print(
+        f"PASS {name}: {status} nonzero={nonzero}/{diff.numel()} "
+        f"max_abs={max_abs:.6f} p99_abs={p99_abs:.6f} cosine={cosine:.8f}"
+    )
+
+
+def run_shape(ops, label: str, rows: int, dim: int, eps: float, reference_ops=None) -> None:
+    x, scale, shift, inv_s, act_scale, scale_fp8, shift_fp8, scale_deq, shift_deq = make_case(rows, dim)
+
+    mod = ref_adaln(x, scale, shift, eps)
+    got_fp8 = ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, eps)
+    assert_fp8_contract(f"{label}/ada_fp8", got_fp8, quant_fp8(mod, act_scale))
+
+    # per-token form: one modulation row per activation row; the reference
+    # is the same adaln math, broadcasting replaced by row alignment
+    scale_pt = (0.1 * torch.randn(rows, dim, device=x.device)).to(torch.bfloat16)
+    shift_pt = (0.1 * torch.randn(rows, dim, device=x.device)).to(torch.bfloat16)
+    mod_pt = ref_adaln(x, scale_pt, shift_pt, eps)
+    got_ptok = ops.ada_layer_norm_quant_fp8_ptok_bf16(x, scale_pt, shift_pt, act_scale, eps)
+    assert_fp8_contract(f"{label}/ada_fp8_ptok", got_ptok, quant_fp8(mod_pt, act_scale))
+
+    # table form: the block's [n_chunks, dim] fp32 table is added to a
+    # [rows, n_chunks, dim] bf16 per-token embedding inside the kernel;
+    # the reference materializes the chunks the way a host block would
+    n_chunks, s_idx, c_idx = 6, 0, 1
+    temb_pt = (0.1 * torch.randn(rows, n_chunks, dim, device=x.device)).to(torch.bfloat16)
+    table_pt = (torch.randn(n_chunks, dim, device=x.device) / dim**0.5).float()
+    shift_tab = (table_pt[s_idx] + temb_pt[:, s_idx].float())
+    scale_tab = (table_pt[c_idx] + temb_pt[:, c_idx].float())
+    mod_tab = ref_adaln_float_mod(x, scale_tab, shift_tab, eps)
+    got_tab = ops.ada_layer_norm_quant_fp8_ptok_table_bf16(
+        x, temb_pt, table_pt, act_scale, s_idx, c_idx, eps
+    )
+    assert_fp8_contract(f"{label}/ada_fp8_ptok_table", got_tab, quant_fp8(mod_tab, act_scale))
+
+    scale_mod = scale_fp8.float() * scale_deq
+    shift_mod = shift_fp8.float() * shift_deq
+    mod_fp8 = ref_adaln_float_mod(x, scale_mod, shift_mod, eps)
+    got_modfp8 = ops.ada_layer_norm_quant_fp8_modfp8_bf16(
+        x, scale_fp8, shift_fp8, scale_deq, shift_deq, act_scale, eps
+    )
+    assert_fp8_contract(f"{label}/ada_modfp8_fp8", got_modfp8, quant_fp8(mod_fp8, act_scale))
+
+    awq_ref = mod.float() * inv_s.float()
+    got_awq = ops.awq_ada_layer_norm_quant_fp8_bf16(x, scale, shift, inv_s, act_scale, eps)
+    assert_fp8_contract(f"{label}/awq_ada_fp8", got_awq, quant_fp8(awq_ref, act_scale))
+
+    got_noaffine = ops.layer_norm_no_affine_quant_fp8_static_bf16(x, act_scale, eps)
+    assert_fp8_contract(f"{label}/no_affine_fp8", got_noaffine, quant_fp8(ref_layer_norm_no_affine(x, eps), act_scale))
+
+    if dim % 16 == 0 and rows <= 260:
+        packed = torch.empty((rows, dim // 2), device=x.device, dtype=torch.uint8)
+        sf = torch.zeros((swizzled_sf_size(rows, dim),), device=x.device, dtype=torch.uint8)
+        got_packed, got_sf = ops.ada_layer_norm_quant_nvfp4_swizzled_bf16(
+            x, scale, shift, eps, packed=packed, sf_swizzled=sf
+        )
+        sm110_contract = torch.cuda.get_device_capability(0) == (11, 0)
+        oracle = reference_ops or ops
+        if sm110_contract and hasattr(oracle, "reference_adaln_fp4"):
+            exp_packed = torch.empty_like(packed)
+            exp_sf = torch.zeros_like(sf)
+            oracle.reference_adaln_fp4(
+                x, scale, shift, eps, exp_packed, exp_sf
+            )
+        else:
+            exp_packed, exp_sf = ref_nvfp4(
+                mod, sm110_contract=sm110_contract
+            )
+        assert_exact(f"{label}/nvfp4_packed", got_packed, exp_packed)
+        assert_exact(f"{label}/nvfp4_sf", got_sf, exp_sf)
+
+        if torch.cuda.get_device_capability(0) == (11, 0):
+            packed.zero_()
+            sf.zero_()
+            got_ln_packed, got_ln_sf = (
+                ops.layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+                    x, eps, packed=packed, sf_swizzled=sf
+                )
+            )
+            if hasattr(oracle, "reference_ln_fp4"):
+                exp_ln_packed = torch.empty_like(packed)
+                exp_ln_sf = torch.zeros_like(sf)
+                oracle.reference_ln_fp4(
+                    x, eps, exp_ln_packed, exp_ln_sf
+                )
+            else:
+                exp_ln_packed, exp_ln_sf = ref_nvfp4(
+                    ref_layer_norm_no_affine(x, eps), sm110_contract=True
+                )
+            assert_exact(
+                f"{label}/no_affine_nvfp4_packed", got_ln_packed,
+                exp_ln_packed,
+            )
+            assert_exact(
+                f"{label}/no_affine_nvfp4_sf", got_ln_sf, exp_ln_sf
+            )
+
+        packed.zero_()
+        sf.zero_()
+        got_packed2, got_sf2 = ops.ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(
+            x, scale_fp8, shift_fp8, scale_deq, shift_deq, eps, packed=packed, sf_swizzled=sf
+        )
+        exp_packed2, exp_sf2 = ref_nvfp4(mod_fp8)
+        assert_exact(f"{label}/nvfp4_modfp8_packed", got_packed2, exp_packed2)
+        assert_exact(f"{label}/nvfp4_modfp8_sf", got_sf2, exp_sf2)
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+    torch.manual_seed(2026)
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    reference_ops = None
+    if args.backend == "installed" and torch.cuda.get_device_capability(0) == (11, 0):
+        reference_ops = load_sm110_reference_ops()
+    shapes = {
+        "decode_action": (16, 2048),
+        "wan_video_short": (64, 3072),
+        "wan_video_ctx": (256, 3072),
+        "wan_video_2k": (2520, 3072),
+        "wan_video_4k": (4096, 3072),
+    }
+    if args.mode == "smoke":
+        shapes = {"wan_video_short": shapes["wan_video_short"]}
+    for label, (rows, dim) in shapes.items():
+        run_shape(ops, label, rows, dim, args.eps, reference_ops)
+
+    modulation_shapes = {
+        "boundary": (1, 1, 48),
+        "groot_n17_dit": (1, 41, 1536),
+        "groot_legacy_dit": (1, 51, 1536),
+        "motus": (1, 360, 3072),
+        "video_long": (2, 2520, 3072),
+    }
+    if args.mode == "smoke":
+        modulation_shapes = {"boundary": modulation_shapes["boundary"]}
+    for label, (batch, sequence, dim) in modulation_shapes.items():
+        params = torch.randn(
+            (batch, sequence, 6, dim), device="cuda", dtype=torch.float32
+        )
+        modulation = torch.randn(
+            (6, dim), device="cuda", dtype=torch.float32
+        )
+        expected = tuple(
+            (params[:, :, index, :] + modulation[index]).to(torch.bfloat16)
+            for index in range(6)
+        )
+        out = ops.adaln_modulation6_bf16(params, modulation)
+        for index, (got, ref) in enumerate(zip(out, expected, strict=True)):
+            if not torch.equal(got, ref):
+                diff = (got.float() - ref.float()).abs()
+                raise AssertionError(
+                    f"{label}/adaln_modulation6/out{index}: "
+                    f"max_abs={diff.max().item()}"
+                )
+        compiled = torch.compile(
+            lambda p, m: ops.adaln_modulation6_bf16(p, m),
+            fullgraph=True,
+        )
+        compiled_out = compiled(params, modulation)
+        for index, (got, ref) in enumerate(
+            zip(compiled_out, expected, strict=True)
+        ):
+            if not torch.equal(got, ref):
+                raise AssertionError(
+                    f"{label}/adaln_modulation6_compile/out{index} mismatch"
+                )
+        print(f"PASS {label}/adaln_modulation6: exact six-output parity")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="full")
+    parser.add_argument("--eps", type=float, default=1e-5)
+    args = parser.parse_args()
+    run(args)
+
+
+if __name__ == "__main__":
+    main()

torch-ext/adaptive_layernorm_producers/__init__.py

@@ -0,0 +1,455 @@
+"""FlashRT adaptive LayerNorm producer kernels for DiT/Wan-style blocks."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_x(x: torch.Tensor) -> None:
+    if x.dim() != 2:
+        raise RuntimeError("x must have shape (rows, dim)")
+    if x.shape[1] % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+
+
+def _check_mod(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> None:
+    _check_x(x)
+    dim = x.shape[1]
+    if scale.shape != (dim,) or shift.shape != (dim,):
+        raise RuntimeError("scale and shift must have shape (dim,)")
+
+
+def _check_scalar(scale: torch.Tensor, name: str) -> None:
+    if scale.numel() != 1:
+        raise RuntimeError(f"{name} must be a scalar tensor")
+
+
+def swizzled_sf_size(rows: int, dim: int) -> int:
+    """Return bytes required by the FlashRT/CUTLASS NVFP4 swizzled scale layout."""
+
+    if dim % 16 != 0:
+        raise RuntimeError("dim must be divisible by 16 for NVFP4 swizzled output")
+    n_blocks = dim // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 128 * 64
+
+
+def _alloc_fp8(x: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(x, dtype=torch.float8_e4m3fn)
+
+
+def _alloc_nvfp4(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, dim = x.shape
+    packed = torch.empty((rows, dim // 2), device=x.device, dtype=torch.uint8)
+    sf = torch.zeros((swizzled_sf_size(rows, dim),), device=x.device, dtype=torch.uint8)
+    return packed, sf
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_fp8_bf16"))
+def _ada_layer_norm_quant_fp8_bf16_fake(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_mod(x, scale, shift)
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_fp8_ptok_bf16"))
+def _ada_layer_norm_quant_fp8_ptok_bf16_fake(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_x(x)
+    if scale.shape != x.shape or shift.shape != x.shape:
+        raise RuntimeError(
+            "per-token scale and shift must have shape (seq_len, dim)")
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_fp8_ptok_table_bf16"))
+def _ada_layer_norm_quant_fp8_ptok_table_bf16_fake(
+    x: torch.Tensor,
+    temb: torch.Tensor,
+    table: torch.Tensor,
+    act_scale: torch.Tensor,
+    shift_idx: int,
+    scale_idx: int,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_x(x)
+    if temb.dim() != 3 or temb.shape[0] != x.shape[0] or temb.shape[2] != x.shape[1]:
+        raise RuntimeError("temb must have shape (seq_len, n_chunks, dim)")
+    if table.dim() != 2 or table.shape[0] != temb.shape[1] or table.shape[1] != x.shape[1]:
+        raise RuntimeError("table must have shape (n_chunks, dim)")
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_fp8_modfp8_bf16"))
+def _ada_layer_norm_quant_fp8_modfp8_bf16_fake(
+    x: torch.Tensor,
+    scale_fp8: torch.Tensor,
+    shift_fp8: torch.Tensor,
+    scale_deq: torch.Tensor,
+    shift_deq: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_mod(x, scale_fp8, shift_fp8)
+    _check_scalar(scale_deq, "scale_deq")
+    _check_scalar(shift_deq, "shift_deq")
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("awq_ada_layer_norm_quant_fp8_bf16"))
+def _awq_ada_layer_norm_quant_fp8_bf16_fake(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    inv_s: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_mod(x, scale, shift)
+    if inv_s.shape != (x.shape[1],):
+        raise RuntimeError("inv_s must have shape (dim,)")
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_nvfp4_swizzled_bf16"))
+def _ada_layer_norm_quant_nvfp4_swizzled_bf16_fake(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+    packed: torch.Tensor,
+    sf_swizzled: torch.Tensor,
+) -> None:
+    _check_mod(x, scale, shift)
+    rows, dim = x.shape
+    if packed.shape != (rows, dim // 2):
+        raise RuntimeError("packed must have shape (rows, dim // 2)")
+    if sf_swizzled.numel() < swizzled_sf_size(rows, dim):
+        raise RuntimeError("sf_swizzled is too small")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16"))
+def _ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16_fake(
+    x: torch.Tensor,
+    scale_fp8: torch.Tensor,
+    shift_fp8: torch.Tensor,
+    scale_deq: torch.Tensor,
+    shift_deq: torch.Tensor,
+    eps: float,
+    packed: torch.Tensor,
+    sf_swizzled: torch.Tensor,
+) -> None:
+    _check_mod(x, scale_fp8, shift_fp8)
+    _check_scalar(scale_deq, "scale_deq")
+    _check_scalar(shift_deq, "shift_deq")
+    rows, dim = x.shape
+    if packed.shape != (rows, dim // 2):
+        raise RuntimeError("packed must have shape (rows, dim // 2)")
+    if sf_swizzled.numel() < swizzled_sf_size(rows, dim):
+        raise RuntimeError("sf_swizzled is too small")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_no_affine_quant_fp8_static_bf16"))
+def _layer_norm_no_affine_quant_fp8_static_bf16_fake(
+    x: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_x(x)
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_no_affine_quant_nvfp4_swizzled_bf16"))
+def _layer_norm_no_affine_quant_nvfp4_swizzled_bf16_fake(
+    x: torch.Tensor,
+    eps: float,
+    packed: torch.Tensor,
+    sf_swizzled: torch.Tensor,
+) -> None:
+    _check_x(x)
+    rows, dim = x.shape
+    if packed.shape != (rows, dim // 2):
+        raise RuntimeError("packed must have shape (rows, dim // 2)")
+    if sf_swizzled.numel() < swizzled_sf_size(rows, dim):
+        raise RuntimeError("sf_swizzled is too small")
+    return None
+
+
+def ada_layer_norm_quant_fp8_bf16(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused no-affine LayerNorm + AdaLN scale/shift + static FP8 quantize."""
+
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, float(eps), out)
+    return out
+
+
+def ada_layer_norm_quant_fp8_ptok_bf16(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Per-token AdaLN producer: no-affine LayerNorm + [M, D] scale/shift
+    modulation + static FP8 quantize, one pass.
+
+    The per-token form serves hosts whose timestep embedding carries one
+    modulation vector per activation row (video DiTs with [B, S, 6, D]
+    tables); the broadcast form above cannot express that seam.
+    """
+
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.ada_layer_norm_quant_fp8_ptok_bf16(
+        x, scale, shift, act_scale, float(eps), out)
+    return out
+
+
+def ada_layer_norm_quant_fp8_ptok_table_bf16(
+    x: torch.Tensor,
+    temb: torch.Tensor,
+    table: torch.Tensor,
+    act_scale: torch.Tensor,
+    shift_idx: int,
+    scale_idx: int,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Per-token AdaLN producer, table form.
+
+    ``temb`` is the host's per-token modulation embedding
+    ``[seq_len, n_chunks, dim]`` (cast to BF16 once per call) and
+    ``table`` the block's own ``[n_chunks, dim]`` FP32 modulation
+    parameter. The table add, chunk selection, no-affine LayerNorm,
+    modulation and static FP8 quantize all happen in one pass — the
+    per-block six-chunk materialization the host would otherwise pay
+    never exists.
+    """
+
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.ada_layer_norm_quant_fp8_ptok_table_bf16(
+        x, temb, table, act_scale, int(shift_idx), int(scale_idx),
+        float(eps), out)
+    return out
+
+
+def ada_layer_norm_quant_fp8_modfp8_bf16(
+    x: torch.Tensor,
+    scale_fp8: torch.Tensor,
+    shift_fp8: torch.Tensor,
+    scale_deq: torch.Tensor,
+    shift_deq: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused AdaLN producer when modulation vectors are stored in FP8."""
+
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.ada_layer_norm_quant_fp8_modfp8_bf16(
+        x, scale_fp8, shift_fp8, scale_deq, shift_deq, act_scale, float(eps), out
+    )
+    return out
+
+
+def awq_ada_layer_norm_quant_fp8_bf16(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    inv_s: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused AdaLN producer with AWQ/SmoothQuant per-channel activation scale."""
+
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.awq_ada_layer_norm_quant_fp8_bf16(x, scale, shift, inv_s, act_scale, float(eps), out)
+    return out
+
+
+def ada_layer_norm_quant_nvfp4_swizzled_bf16(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float = 1e-5,
+    packed: torch.Tensor | None = None,
+    sf_swizzled: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused AdaLN producer to NVFP4 packed activations and swizzled scale factors."""
+
+    if packed is None or sf_swizzled is None:
+        packed, sf_swizzled = _alloc_nvfp4(x)
+    ops.ada_layer_norm_quant_nvfp4_swizzled_bf16(x, scale, shift, float(eps), packed, sf_swizzled)
+    return packed, sf_swizzled
+
+
+def ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(
+    x: torch.Tensor,
+    scale_fp8: torch.Tensor,
+    shift_fp8: torch.Tensor,
+    scale_deq: torch.Tensor,
+    shift_deq: torch.Tensor,
+    eps: float = 1e-5,
+    packed: torch.Tensor | None = None,
+    sf_swizzled: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused FP8-modulated AdaLN producer to NVFP4 packed activations."""
+
+    if packed is None or sf_swizzled is None:
+        packed, sf_swizzled = _alloc_nvfp4(x)
+    ops.ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(
+        x, scale_fp8, shift_fp8, scale_deq, shift_deq, float(eps), packed, sf_swizzled
+    )
+    return packed, sf_swizzled
+
+
+def layer_norm_no_affine_quant_fp8_static_bf16(
+    x: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused no-affine LayerNorm + static FP8 quantize."""
+
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.layer_norm_no_affine_quant_fp8_static_bf16(x, act_scale, float(eps), out)
+    return out
+
+
+def layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+    x: torch.Tensor,
+    eps: float = 1e-5,
+    packed: torch.Tensor | None = None,
+    sf_swizzled: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SM110 no-affine LayerNorm producing NVFP4 and CUTLASS SFA."""
+    if packed is None or sf_swizzled is None:
+        packed, sf_swizzled = _alloc_nvfp4(x)
+    ops.layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+        x, float(eps), packed, sf_swizzled
+    )
+    return packed, sf_swizzled
+
+
+@torch.library.register_fake(add_op_namespace_prefix("adaln_modulation6_bf16"))
+def _adaln_modulation6_bf16_fake(
+    adaln_params: torch.Tensor,
+    layer_modulation: torch.Tensor,
+    out0: torch.Tensor,
+    out1: torch.Tensor,
+    out2: torch.Tensor,
+    out3: torch.Tensor,
+    out4: torch.Tensor,
+    out5: torch.Tensor,
+) -> None:
+    if adaln_params.dim() != 4 or adaln_params.shape[2] != 6:
+        raise RuntimeError(
+            "adaln_params must have shape (batch, sequence, 6, dim)"
+        )
+    expected = (
+        adaln_params.shape[0],
+        adaln_params.shape[1],
+        adaln_params.shape[3],
+    )
+    if layer_modulation.shape != (6, expected[2]):
+        raise RuntimeError("layer_modulation must have shape (6, dim)")
+    for output in (out0, out1, out2, out3, out4, out5):
+        if output.shape != expected:
+            raise RuntimeError(
+                "each output must have shape (batch, sequence, dim)"
+            )
+    return None
+
+
+def adaln_modulation6_bf16(
+    adaln_params: torch.Tensor,
+    layer_modulation: torch.Tensor,
+    *,
+    out: tuple[torch.Tensor, ...] | None = None,
+) -> tuple[torch.Tensor, ...]:
+    """Add six FP32 AdaLN modulation vectors and return six BF16 tensors."""
+
+    batch, sequence, _, dim = adaln_params.shape
+    if out is None:
+        out = tuple(
+            torch.empty(
+                (batch, sequence, dim),
+                device=adaln_params.device,
+                dtype=torch.bfloat16,
+            )
+            for _ in range(6)
+        )
+    if len(out) != 6:
+        raise RuntimeError("out must contain six tensors")
+    ops.adaln_modulation6_bf16(
+        adaln_params, layer_modulation, *out
+    )
+    return out
+
+
+__all__ = [
+    "ada_layer_norm_quant_fp8_bf16",
+    "ada_layer_norm_quant_fp8_ptok_bf16",
+    "ada_layer_norm_quant_fp8_ptok_table_bf16",
+    "ada_layer_norm_quant_fp8_modfp8_bf16",
+    "awq_ada_layer_norm_quant_fp8_bf16",
+    "ada_layer_norm_quant_nvfp4_swizzled_bf16",
+    "ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16",
+    "layer_norm_no_affine_quant_fp8_static_bf16",
+    "layer_norm_no_affine_quant_nvfp4_swizzled_bf16",
+    "adaln_modulation6_bf16",
+    "swizzled_sf_size",
+]

torch-ext/torch_binding.cpp

@@ -0,0 +1,586 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include <torch/all.h>
+#include <torch/library.h>
+
+#if defined(CUDA_KERNEL)
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#endif
+
+#include "ada_layer_norm_fp8.cuh"
+#include "adaln_modulation6.cuh"
+#include "dit_layer_norm_fp8.cuh"
+#include "registration.h"
+#include "sm110_fp4_dispatch.cuh"
+
+flash_rt::adaln_producers::hub::AdaLayerNormFp4Dispatch
+    flash_rt::adaln_producers::hub::ada_layer_norm_fp4_dispatch = nullptr;
+flash_rt::adaln_producers::hub::LayerNormFp4Dispatch
+    flash_rt::adaln_producers::hub::layer_norm_fp4_dispatch = nullptr;
+
+namespace {
+
+void check_cuda_contiguous(torch::Tensor const& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+
+void check_bf16(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kBFloat16,
+              name, " must have dtype torch.bfloat16");
+}
+
+void check_fp32(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat32,
+              name, " must have dtype torch.float32");
+}
+
+void check_fp8(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat8_e4m3fn,
+              name, " must have dtype torch.float8_e4m3fn");
+}
+
+void check_u8(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kUInt8,
+              name, " must have dtype torch.uint8");
+}
+
+void check_same_device(torch::Tensor const& a,
+                       torch::Tensor const& b,
+                       const char* a_name,
+                       const char* b_name) {
+  TORCH_CHECK(a.get_device() == b.get_device(),
+              a_name, " and ", b_name, " must be on the same CUDA device");
+}
+
+void check_x(torch::Tensor const& x) {
+  check_bf16(x, "x");
+  TORCH_CHECK(x.dim() == 2, "x must have shape (rows, dim)");
+  TORCH_CHECK(x.size(0) > 0 && x.size(1) > 0,
+              "x rows and dim must be positive");
+  TORCH_CHECK((x.size(1) % 2) == 0, "x.shape[1] must be even");
+}
+
+void check_bf16_mod(torch::Tensor const& x,
+                    torch::Tensor const& scale,
+                    torch::Tensor const& shift) {
+  check_bf16(scale, "scale");
+  check_bf16(shift, "shift");
+  TORCH_CHECK(scale.dim() == 1 && scale.size(0) == x.size(1),
+              "scale must have shape (dim,)");
+  TORCH_CHECK(shift.dim() == 1 && shift.size(0) == x.size(1),
+              "shift must have shape (dim,)");
+  check_same_device(x, scale, "x", "scale");
+  check_same_device(x, shift, "x", "shift");
+}
+
+void check_fp8_mod(torch::Tensor const& x,
+                   torch::Tensor const& scale,
+                   torch::Tensor const& shift,
+                   torch::Tensor const& scale_deq,
+                   torch::Tensor const& shift_deq) {
+  check_fp8(scale, "scale_fp8");
+  check_fp8(shift, "shift_fp8");
+  check_fp32(scale_deq, "scale_deq");
+  check_fp32(shift_deq, "shift_deq");
+  TORCH_CHECK(scale.dim() == 1 && scale.size(0) == x.size(1),
+              "scale_fp8 must have shape (dim,)");
+  TORCH_CHECK(shift.dim() == 1 && shift.size(0) == x.size(1),
+              "shift_fp8 must have shape (dim,)");
+  TORCH_CHECK(scale_deq.numel() == 1, "scale_deq must be a scalar tensor");
+  TORCH_CHECK(shift_deq.numel() == 1, "shift_deq must be a scalar tensor");
+  check_same_device(x, scale, "x", "scale_fp8");
+  check_same_device(x, shift, "x", "shift_fp8");
+  check_same_device(x, scale_deq, "x", "scale_deq");
+  check_same_device(x, shift_deq, "x", "shift_deq");
+}
+
+void check_act_scale(torch::Tensor const& x, torch::Tensor const& act_scale) {
+  check_fp32(act_scale, "act_scale");
+  TORCH_CHECK(act_scale.numel() == 1, "act_scale must be a scalar tensor");
+  check_same_device(x, act_scale, "x", "act_scale");
+}
+
+int64_t swizzled_sf_bytes(int64_t rows, int64_t dim) {
+  TORCH_CHECK((dim % 16) == 0, "dim must be divisible by 16 for NVFP4 swizzled output");
+  const int64_t blocks = dim / 16;
+  const int64_t row_super = (rows + 127) / 128;
+  const int64_t col_super = (blocks + 3) / 4;
+  return row_super * col_super * 128 * 64;
+}
+
+void check_fp8_out(torch::Tensor const& x, torch::Tensor const& out) {
+  check_fp8(out, "out");
+  TORCH_CHECK(out.sizes() == x.sizes(), "out must have the same shape as x");
+  check_same_device(x, out, "x", "out");
+}
+
+void check_nvfp4_out(torch::Tensor const& x,
+                     torch::Tensor const& packed,
+                     torch::Tensor const& sf_swizzled) {
+  check_u8(packed, "packed");
+  check_u8(sf_swizzled, "sf_swizzled");
+  TORCH_CHECK((x.size(1) % 16) == 0, "x.shape[1] must be divisible by 16");
+  TORCH_CHECK(packed.dim() == 2 && packed.size(0) == x.size(0) &&
+                  packed.size(1) == x.size(1) / 2,
+              "packed must have shape (rows, dim // 2)");
+  TORCH_CHECK(sf_swizzled.numel() >= swizzled_sf_bytes(x.size(0), x.size(1)),
+              "sf_swizzled is too small for the swizzled NVFP4 scale layout");
+  check_same_device(x, packed, "x", "packed");
+  check_same_device(x, sf_swizzled, "x", "sf_swizzled");
+}
+
+}  // namespace
+
+void ada_layer_norm_quant_fp8_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& scale,
+    torch::Tensor const& shift,
+    torch::Tensor const& act_scale,
+    double eps,
+    torch::Tensor& out) {
+  check_x(x);
+  check_bf16_mod(x, scale, shift);
+  check_act_scale(x, act_scale);
+  check_fp8_out(x, out);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  flash_rt::quantize::ada_layer_norm_fp8(
+      x.data_ptr(),
+      scale.data_ptr(),
+      shift.data_ptr(),
+      out.data_ptr(),
+      static_cast<const float*>(act_scale.data_ptr()),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void check_bf16_mod_ptok(torch::Tensor const& x,
+                         torch::Tensor const& scale,
+                         torch::Tensor const& shift) {
+  check_bf16(scale, "scale");
+  check_bf16(shift, "shift");
+  TORCH_CHECK(scale.sizes() == x.sizes(),
+              "per-token scale must have shape (seq_len, dim)");
+  TORCH_CHECK(shift.sizes() == x.sizes(),
+              "per-token shift must have shape (seq_len, dim)");
+  check_same_device(x, scale, "x", "scale");
+  check_same_device(x, shift, "x", "shift");
+}
+
+void ada_layer_norm_quant_fp8_ptok_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& scale,
+    torch::Tensor const& shift,
+    torch::Tensor const& act_scale,
+    double eps,
+    torch::Tensor& out) {
+  check_x(x);
+  check_bf16_mod_ptok(x, scale, shift);
+  check_act_scale(x, act_scale);
+  check_fp8_out(x, out);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  flash_rt::quantize::ada_layer_norm_ptok_fp8(
+      x.data_ptr(),
+      scale.data_ptr(),
+      shift.data_ptr(),
+      out.data_ptr(),
+      static_cast<const float*>(act_scale.data_ptr()),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void ada_layer_norm_quant_fp8_ptok_table_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& temb,
+    torch::Tensor const& table,
+    torch::Tensor const& act_scale,
+    int64_t shift_idx,
+    int64_t scale_idx,
+    double eps,
+    torch::Tensor& out) {
+  check_x(x);
+  check_bf16(temb, "temb");
+  check_cuda_contiguous(temb, "temb");
+  check_fp32(table, "table");
+  check_cuda_contiguous(table, "table");
+  TORCH_CHECK(temb.dim() == 3 && temb.size(0) == x.size(0) &&
+                  temb.size(2) == x.size(1),
+              "temb must have shape (seq_len, n_chunks, dim)");
+  const int64_t n_chunks = temb.size(1);
+  TORCH_CHECK(table.dim() == 2 && table.size(0) == n_chunks &&
+                  table.size(1) == x.size(1),
+              "table must have shape (n_chunks, dim)");
+  TORCH_CHECK(shift_idx >= 0 && shift_idx < n_chunks &&
+                  scale_idx >= 0 && scale_idx < n_chunks,
+              "chunk indices must lie in [0, n_chunks)");
+  check_act_scale(x, act_scale);
+  check_fp8_out(x, out);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  flash_rt::quantize::ada_layer_norm_ptok_table_fp8(
+      x.data_ptr(),
+      temb.data_ptr(),
+      static_cast<const float*>(table.data_ptr()),
+      out.data_ptr(),
+      static_cast<const float*>(act_scale.data_ptr()),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<int>(n_chunks),
+      static_cast<int>(shift_idx),
+      static_cast<int>(scale_idx),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void ada_layer_norm_quant_fp8_modfp8_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& scale_fp8,
+    torch::Tensor const& shift_fp8,
+    torch::Tensor const& scale_deq,
+    torch::Tensor const& shift_deq,
+    torch::Tensor const& act_scale,
+    double eps,
+    torch::Tensor& out) {
+  check_x(x);
+  check_fp8_mod(x, scale_fp8, shift_fp8, scale_deq, shift_deq);
+  check_act_scale(x, act_scale);
+  check_fp8_out(x, out);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  flash_rt::quantize::ada_layer_norm_fp8_modfp8(
+      x.data_ptr(),
+      scale_fp8.data_ptr(),
+      shift_fp8.data_ptr(),
+      static_cast<const float*>(scale_deq.data_ptr()),
+      static_cast<const float*>(shift_deq.data_ptr()),
+      out.data_ptr(),
+      static_cast<const float*>(act_scale.data_ptr()),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void awq_ada_layer_norm_quant_fp8_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& scale,
+    torch::Tensor const& shift,
+    torch::Tensor const& inv_s,
+    torch::Tensor const& act_scale,
+    double eps,
+    torch::Tensor& out) {
+  check_x(x);
+  check_bf16_mod(x, scale, shift);
+  check_bf16(inv_s, "inv_s");
+  TORCH_CHECK(inv_s.dim() == 1 && inv_s.size(0) == x.size(1),
+              "inv_s must have shape (dim,)");
+  check_same_device(x, inv_s, "x", "inv_s");
+  check_act_scale(x, act_scale);
+  check_fp8_out(x, out);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  flash_rt::quantize::awq_ada_layer_norm_fp8(
+      x.data_ptr(),
+      scale.data_ptr(),
+      shift.data_ptr(),
+      inv_s.data_ptr(),
+      out.data_ptr(),
+      static_cast<const float*>(act_scale.data_ptr()),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void ada_layer_norm_quant_nvfp4_swizzled_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& scale,
+    torch::Tensor const& shift,
+    double eps,
+    torch::Tensor& packed,
+    torch::Tensor& sf_swizzled) {
+  check_x(x);
+  check_bf16_mod(x, scale, shift);
+  check_nvfp4_out(x, packed, sf_swizzled);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  auto const* props = at::cuda::getDeviceProperties(x.get_device());
+  if (props->major == 11 && props->minor == 0) {
+    TORCH_CHECK(
+        flash_rt::adaln_producers::hub::ada_layer_norm_fp4_dispatch != nullptr,
+        "SM110 AdaLayerNorm-to-FP4 source is not present in this build");
+    const int rc =
+        flash_rt::adaln_producers::hub::ada_layer_norm_fp4_dispatch(
+            x.data_ptr(), scale.data_ptr(), shift.data_ptr(),
+            packed.data_ptr(), sf_swizzled.data_ptr(),
+            static_cast<int>(x.size(0)), static_cast<int>(x.size(1)),
+            static_cast<float>(eps), stream);
+    TORCH_CHECK(rc == 0,
+                "ada_layer_norm_quant_nvfp4_swizzled_bf16 failed with rc=",
+                rc);
+    return;
+  }
+  flash_rt::quantize::ada_layer_norm_nvfp4_swizzled(
+      x.data_ptr(),
+      scale.data_ptr(),
+      shift.data_ptr(),
+      packed.data_ptr(),
+      sf_swizzled.data_ptr(),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+    torch::Tensor const& x,
+    double eps,
+    torch::Tensor& packed,
+    torch::Tensor& sf_swizzled) {
+  check_x(x);
+  check_nvfp4_out(x, packed, sf_swizzled);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto const* props = at::cuda::getDeviceProperties(x.get_device());
+  TORCH_CHECK(props->major == 11 && props->minor == 0,
+              "layer_norm_no_affine_quant_nvfp4_swizzled_bf16 currently "
+              "requires SM110; got SM", props->major, props->minor);
+  TORCH_CHECK(
+      flash_rt::adaln_producers::hub::layer_norm_fp4_dispatch != nullptr,
+      "SM110 LayerNorm-to-FP4 source is not present in this build");
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  const int rc = flash_rt::adaln_producers::hub::layer_norm_fp4_dispatch(
+      x.data_ptr(), packed.data_ptr(), sf_swizzled.data_ptr(),
+      static_cast<int>(x.size(0)), static_cast<int>(x.size(1)),
+      static_cast<float>(eps), stream);
+  TORCH_CHECK(
+      rc == 0,
+      "layer_norm_no_affine_quant_nvfp4_swizzled_bf16 failed with rc=", rc);
+#else
+  TORCH_CHECK(false,
+              "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& scale_fp8,
+    torch::Tensor const& shift_fp8,
+    torch::Tensor const& scale_deq,
+    torch::Tensor const& shift_deq,
+    double eps,
+    torch::Tensor& packed,
+    torch::Tensor& sf_swizzled) {
+  check_x(x);
+  check_fp8_mod(x, scale_fp8, shift_fp8, scale_deq, shift_deq);
+  check_nvfp4_out(x, packed, sf_swizzled);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  flash_rt::quantize::ada_layer_norm_nvfp4_swizzled_modfp8(
+      x.data_ptr(),
+      scale_fp8.data_ptr(),
+      shift_fp8.data_ptr(),
+      static_cast<const float*>(scale_deq.data_ptr()),
+      static_cast<const float*>(shift_deq.data_ptr()),
+      packed.data_ptr(),
+      sf_swizzled.data_ptr(),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void layer_norm_no_affine_quant_fp8_static_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& act_scale,
+    double eps,
+    torch::Tensor& out) {
+  check_x(x);
+  check_act_scale(x, act_scale);
+  check_fp8_out(x, out);
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  flash_rt::adaln_producers::layer_norm_no_affine_fp8_static_bf16(
+      x.data_ptr(),
+      out.data_ptr(),
+      static_cast<const float*>(act_scale.data_ptr()),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+void adaln_modulation6_bf16(
+    torch::Tensor const& adaln_params,
+    torch::Tensor const& layer_modulation,
+    torch::Tensor& out0,
+    torch::Tensor& out1,
+    torch::Tensor& out2,
+    torch::Tensor& out3,
+    torch::Tensor& out4,
+    torch::Tensor& out5) {
+  check_fp32(adaln_params, "adaln_params");
+  check_fp32(layer_modulation, "layer_modulation");
+  TORCH_CHECK(
+      adaln_params.dim() == 4 && adaln_params.size(2) == 6,
+      "adaln_params must have shape (batch, sequence, 6, dim)");
+  const auto batch = adaln_params.size(0);
+  const auto sequence = adaln_params.size(1);
+  const auto dim = adaln_params.size(3);
+  TORCH_CHECK(batch > 0 && sequence > 0 && dim > 0,
+              "batch, sequence and dim must be positive");
+  TORCH_CHECK(layer_modulation.dim() == 2 &&
+                  layer_modulation.size(0) == 6 &&
+                  layer_modulation.size(1) == dim,
+              "layer_modulation must have shape (6, dim)");
+  check_same_device(
+      adaln_params, layer_modulation, "adaln_params", "layer_modulation");
+  torch::Tensor* outputs[] = {&out0, &out1, &out2, &out3, &out4, &out5};
+  for (int index = 0; index < 6; ++index) {
+    const std::string name = "out" + std::to_string(index);
+    check_bf16(*outputs[index], name.c_str());
+    TORCH_CHECK(
+        outputs[index]->dim() == 3 &&
+            outputs[index]->size(0) == batch &&
+            outputs[index]->size(1) == sequence &&
+            outputs[index]->size(2) == dim,
+        name, " must have shape (batch, sequence, dim)");
+    check_same_device(
+        adaln_params, *outputs[index], "adaln_params", name.c_str());
+  }
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(adaln_params.device());
+  auto stream =
+      at::cuda::getCurrentCUDAStream(adaln_params.get_device()).stream();
+  flash_rt::adaln_producers::adaln_modulation6_bf16(
+      adaln_params.data_ptr<float>(),
+      layer_modulation.data_ptr<float>(),
+      out0.data_ptr(),
+      out1.data_ptr(),
+      out2.data_ptr(),
+      out3.data_ptr(),
+      out4.data_ptr(),
+      out5.data_ptr(),
+      static_cast<int>(batch),
+      static_cast<int>(sequence),
+      static_cast<int>(dim),
+      stream);
+#else
+  TORCH_CHECK(false, "adaptive-layernorm-producers was not built with CUDA support");
+#endif
+}
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("ada_layer_norm_quant_fp8_bf16("
+          "Tensor x, Tensor scale, Tensor shift, Tensor act_scale, float eps, Tensor! out) -> ()");
+  ops.def("ada_layer_norm_quant_fp8_ptok_bf16("
+          "Tensor x, Tensor scale, Tensor shift, Tensor act_scale, float eps, Tensor! out) -> ()");
+  ops.def("ada_layer_norm_quant_fp8_ptok_table_bf16("
+          "Tensor x, Tensor temb, Tensor table, Tensor act_scale, "
+          "int shift_idx, int scale_idx, float eps, Tensor! out) -> ()");
+  ops.def("ada_layer_norm_quant_fp8_modfp8_bf16("
+          "Tensor x, Tensor scale_fp8, Tensor shift_fp8, Tensor scale_deq, Tensor shift_deq, "
+          "Tensor act_scale, float eps, Tensor! out) -> ()");
+  ops.def("awq_ada_layer_norm_quant_fp8_bf16("
+          "Tensor x, Tensor scale, Tensor shift, Tensor inv_s, Tensor act_scale, float eps, Tensor! out) -> ()");
+  ops.def("ada_layer_norm_quant_nvfp4_swizzled_bf16("
+          "Tensor x, Tensor scale, Tensor shift, float eps, Tensor! packed, Tensor! sf_swizzled) -> ()");
+  ops.def("ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16("
+          "Tensor x, Tensor scale_fp8, Tensor shift_fp8, Tensor scale_deq, Tensor shift_deq, "
+          "float eps, Tensor! packed, Tensor! sf_swizzled) -> ()");
+  ops.def("layer_norm_no_affine_quant_nvfp4_swizzled_bf16("
+          "Tensor x, float eps, Tensor! packed, Tensor! sf_swizzled) -> ()");
+  ops.def("layer_norm_no_affine_quant_fp8_static_bf16("
+          "Tensor x, Tensor act_scale, float eps, Tensor! out) -> ()");
+  ops.def("adaln_modulation6_bf16("
+          "Tensor adaln_params, Tensor layer_modulation, "
+          "Tensor! out0, Tensor! out1, Tensor! out2, "
+          "Tensor! out3, Tensor! out4, Tensor! out5) -> ()");
+#if defined(CUDA_KERNEL)
+  ops.impl("ada_layer_norm_quant_fp8_bf16",
+           torch::kCUDA,
+           &ada_layer_norm_quant_fp8_bf16);
+  ops.impl("ada_layer_norm_quant_fp8_ptok_bf16",
+           torch::kCUDA,
+           &ada_layer_norm_quant_fp8_ptok_bf16);
+  ops.impl("ada_layer_norm_quant_fp8_ptok_table_bf16",
+           torch::kCUDA,
+           &ada_layer_norm_quant_fp8_ptok_table_bf16);
+  ops.impl("ada_layer_norm_quant_fp8_modfp8_bf16",
+           torch::kCUDA,
+           &ada_layer_norm_quant_fp8_modfp8_bf16);
+  ops.impl("awq_ada_layer_norm_quant_fp8_bf16",
+           torch::kCUDA,
+           &awq_ada_layer_norm_quant_fp8_bf16);
+  ops.impl("ada_layer_norm_quant_nvfp4_swizzled_bf16",
+           torch::kCUDA,
+           &ada_layer_norm_quant_nvfp4_swizzled_bf16);
+  ops.impl("ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16",
+           torch::kCUDA,
+           &ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16);
+  ops.impl("layer_norm_no_affine_quant_nvfp4_swizzled_bf16",
+           torch::kCUDA,
+           &layer_norm_no_affine_quant_nvfp4_swizzled_bf16);
+  ops.impl("layer_norm_no_affine_quant_fp8_static_bf16",
+           torch::kCUDA,
+           &layer_norm_no_affine_quant_fp8_static_bf16);
+  ops.impl("adaln_modulation6_bf16",
+           torch::kCUDA,
+           &adaln_modulation6_bf16);
+#endif
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)