Promote latest kernel artifacts to main

CARD.md

@@ -0,0 +1,29 @@
+# FlashRT Adaptive Norms
+
+This package provides FlashRT adaptive normalization kernels for Hugging Face
+Kernel Hub.
+
+It is intended for DiT, VLA, video, and world-model blocks that use RMSNorm
+plus per-row or single-row-broadcast style scale/shift/gate parameters.
+
+## Kernels
+
+- `ada_rms_norm_style_bf16`: RMSNorm + style scale/shift, with BF16 output and
+  BF16 gate extraction.
+- `gate_residual_ada_norm_fp8_static_bf16`: in-place gated residual update,
+  AdaRMSNorm, static FP8 E4M3 output, and BF16 gate extraction.
+
+## Hardware
+
+- CUDA 12.8+
+- BF16-capable NVIDIA GPUs
+- FP8-capable NVIDIA GPUs for the FP8 output API
+
+The per-row and broadcast contracts are source-validated on RTX 5090 and
+Jetson AGX Thor SM110.
+
+## Upstream
+
+The serving source of truth remains FlashRT:
+
+https://github.com/LiangSu8899/FlashRT

README.md

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

SYNC.md

@@ -0,0 +1,20 @@
+# Source Sync
+
+Synced from `official/FlashRT`:
+
+- `csrc/kernels/norm.cu`: `ada_rms_norm_style` and `ada_rms_norm_style_fp8`
+  math.
+- `csrc/kernels/fusion.cu`: `gate_residual_ada_norm_fp8` fused residual,
+  AdaRMSNorm, and FP8 output path.
+
+Local adaptation:
+
+- Public package/API names use generic adaptive-norm terminology.
+- Only BF16 input and static FP8 output paths are included in this first public
+  package version.
+- Tensor-facing `torch.ops` bindings validate dtype, shape, contiguity, device,
+  and even hidden dimensions.
+- CUDA math uses explicit round-to-nearest add/multiply operations in key
+  BF16-reference paths for stable public correctness behavior.
+- FP8 validation gates on p99-zero and sparse boundary differences because
+  PyTorch and CUDA FP8 casters can select adjacent FP8 grid points on rare ties.

VALIDATION.md

@@ -0,0 +1,56 @@
+# Validation: flashrt-adaptive-norms
+
+Required before publishing this package:
+
+1. Source-extension correctness:
+
+   ```bash
+   python flashrt-adaptive-norms/tests/test_adaptive_norms.py \
+     --backend source \
+     --mode full
+   ```
+
+2. Source-extension benchmark:
+
+   ```bash
+   python flashrt-adaptive-norms/benchmarks/benchmark.py \
+     --backend source \
+     --shapes all
+   ```
+
+3. Kernel-builder artifact build:
+
+   ```bash
+   kernel-builder build-and-copy flashrt-adaptive-norms
+   ```
+
+4. Built-artifact correctness:
+
+   ```bash
+   PYTHONPATH=<artifact-path> \
+   python flashrt-adaptive-norms/tests/test_adaptive_norms.py \
+     --backend installed \
+     --mode full
+   ```
+
+5. Built-artifact benchmark:
+
+   ```bash
+   python flashrt-adaptive-norms/benchmarks/benchmark.py \
+     --backend installed \
+     --artifact <artifact-path> \
+     --shapes all
+   ```
+
+6. Multi-hardware matrix:
+
+   Add hardware claims only after the same correctness and benchmark commands
+   pass on that machine.
+
+## Style Broadcast Gate
+
+Both APIs accept `style` as either `(rows, 3*dim)` or `(1, 3*dim)`. The full
+source and installed-artifact runs must execute both layouts at rows
+`64/2520/4096`. The broadcast path additionally requires bit-identical gate
+output, exact residual update, FP8 p99 absolute error zero, CUDA Graph replay,
+and installed-wrapper `torch.compile(fullgraph=True)` parity.

benchmarks/README.md

@@ -0,0 +1,12 @@
+# benchmarks
+
+Package-level microbenchmarks for adaptive norm APIs.
+
+Run:
+
+```bash
+python flashrt-adaptive-norms/benchmarks/benchmark.py --backend source --shapes all
+```
+
+The benchmark compares the fused FlashRT Tensor API against a PyTorch eager
+reference chain that performs the same adaptive norm and FP8 quantization math.

benchmarks/RESULTS.md

@@ -0,0 +1,29 @@
+# Source Benchmark Results
+
+Environment: NVIDIA GeForce RTX 5090 local source-extension build.
+Baseline: PyTorch eager tensor reference with matching BF16/FP8 math contract.
+
+| Shape | Rows,Dim | Kernel | FlashRT us | Eager us | vs eager | p99 abs | Cosine | Status |
+|---|---:|---|---:|---:|---:|---:|---:|---|
+| small | 64,1024 | ada_rms_norm_style_bf16 | 4.172 | 61.407 | 14.72x | 0.000000 | 0.99999988 | PASS |
+| small | 64,1024 | gate_residual_ada_norm_fp8_static_bf16 | 4.150 | 81.997 | 19.76x | 0.000000 | 1.00000000 | PASS |
+| vla_2k | 2520,3072 | ada_rms_norm_style_bf16 | 15.404 | 265.453 | 17.23x | 0.000000 | 0.99999994 | PASS |
+| vla_2k | 2520,3072 | gate_residual_ada_norm_fp8_static_bf16 | 58.822 | 431.002 | 7.33x | 0.000000 | 1.00000000 | PASS |
+| vla_4k | 4096,3072 | ada_rms_norm_style_bf16 | 79.392 | 594.854 | 7.49x | 0.000000 | 1.00000000 | PASS |
+| vla_4k | 4096,3072 | gate_residual_ada_norm_fp8_static_bf16 | 119.475 | 901.242 | 7.54x | 0.000000 | 1.00000000 | PASS |
+
+## Thor Style-Broadcast Update
+
+Source benchmark on NVIDIA Thor, PyTorch 2.13.0+cu130, August 8, 2026. The
+single-row style path removes the expanded style tensor while preserving the
+same public operation.
+
+| Shape | Kernel | Per-row us | Broadcast us | Broadcast vs eager | P99 abs |
+|---|---|---:|---:|---:|---:|
+| `2520x3072` | AdaRMS style | 407.806 | 202.440 | 15.11x | 0 |
+| `2520x3072` | gate+residual+AdaRMS to FP8 | 592.928 | 398.299 | 12.58x | 0 |
+| `4096x3072` | AdaRMS style | 604.960 | 306.514 | 14.79x | 0 |
+| `4096x3072` | gate+residual+AdaRMS to FP8 | 864.846 | 553.539 | 13.69x | 0 |
+
+The full source gate passed both style layouts. Residual and gate tensors were
+exact; FP8 p99 absolute error was zero across all tested rows.

benchmarks/benchmark.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import argparse
 import ctypes
 import ctypes.util
+import importlib
 import json
 import math
 import os
@@ -135,16 +136,22 @@ def rms_norm(x, weight, eps):
 def ref_ada(x, weight, style, eps):
     dim = x.shape[1]
     normed = rms_norm(x, weight, eps)
-    y = normed * (1.0 + style[:, :dim].float()) + style[:, dim : 2 * dim].float()
-    return y.to(torch.bfloat16), style[:, 2 * dim :].contiguous().to(torch.bfloat16)
+    style_scale = style[:, :dim].float().expand(x.shape[0], -1)
+    style_shift = style[:, dim : 2 * dim].float().expand(x.shape[0], -1)
+    style_gate = style[:, 2 * dim :].expand(x.shape[0], -1)
+    y = normed * (1.0 + style_scale) + style_shift
+    return y.to(torch.bfloat16), style_gate.contiguous().to(torch.bfloat16)
 
 
 def ref_fused(residual, x, gate, weight, style, scale, eps):
     updated = (residual.float() + x.float() * gate.float()).to(torch.bfloat16)
     dim = updated.shape[1]
     normed = rms_norm(updated, weight, eps)
-    y = (normed * (1.0 + style[:, :dim].float()) + style[:, dim : 2 * dim].float()) / scale.float().reshape(())
-    return updated, y.to(torch.float8_e4m3fn), style[:, 2 * dim :].contiguous().to(torch.bfloat16)
+    style_scale = style[:, :dim].float().expand(updated.shape[0], -1)
+    style_shift = style[:, dim : 2 * dim].float().expand(updated.shape[0], -1)
+    style_gate = style[:, 2 * dim :].expand(updated.shape[0], -1)
+    y = (normed * (1.0 + style_scale) + style_shift) / scale.float().reshape(())
+    return updated, y.to(torch.float8_e4m3fn), style_gate.contiguous().to(torch.bfloat16)
 
 
 def time_us(fn, warmup, iters):
@@ -174,9 +181,12 @@ def metrics(got, expected):
     )
 
 
-def run_one(ops, name, shape, args):
+def run_one(ops, name, shape, args, broadcast: bool):
     rows, dim = shape
     x, residual, gate, weight, style, scale = make_case(rows, dim)
+    if broadcast:
+        style = style[:1].contiguous()
+        name = f"{name}_broadcast"
     out = torch.empty_like(x)
     gate_out = torch.empty_like(x)
     fp8_out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
@@ -235,6 +245,10 @@ def main():
     parser.add_argument("--warmup", type=int, default=5)
     parser.add_argument("--iters", type=int, default=20)
     parser.add_argument("--eps", type=float, default=1e-6)
+    parser.add_argument(
+        "--style-mode", choices=["per-row", "broadcast", "both"],
+        default="both"
+    )
     parser.add_argument("--output", default=None)
     parser.add_argument("--markdown", default=None)
     args = parser.parse_args()
@@ -244,7 +258,10 @@ def main():
     ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
     results = []
     for name in SHAPE_GROUPS[args.shapes]:
-        results.extend(run_one(ops, name, SHAPES[name], args))
+        if args.style_mode in ("per-row", "both"):
+            results.extend(run_one(ops, name, SHAPES[name], args, False))
+        if args.style_mode in ("broadcast", "both"):
+            results.extend(run_one(ops, name, SHAPES[name], args, True))
     for r in results:
         print(
             f"{r.status} {r.shape}/{r.kernel}: flashrt={r.flashrt_us:.3f}us "

build.toml

@@ -0,0 +1,28 @@
+[general]
+name = "flashrt-adaptive-norms"
+license = "Apache-2.0"
+version = 1
+backends = ["cuda"]
+
+[general.cuda]
+minver = "12.8"
+
+[general.hub]
+repo-id = "flashrt/flashrt-adaptive-norms"
+
+[torch]
+include = ["csrc"]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+
+[kernel.flashrt_adaptive_norms]
+backend = "cuda"
+depends = ["torch"]
+include = ["csrc"]
+cuda-minver = "12.8"
+src = [
+  "csrc/adaptive_norms.cu",
+  "csrc/adaptive_norms.cuh",
+]

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

@@ -21,8 +21,8 @@ def _check_adaptive_shapes(
         raise RuntimeError("x.shape[1] must be even")
     if weight.shape != (dim,):
         raise RuntimeError("weight must have shape (dim,)")
-    if style.shape != (rows, 3 * dim):
-        raise RuntimeError("style must have shape (rows, 3 * dim)")
+    if style.shape not in ((rows, 3 * dim), (1, 3 * dim)):
+        raise RuntimeError("style must have shape (rows, 3 * dim) or (1, 3 * dim)")
     if out.shape != x.shape or gate_out.shape != x.shape:
         raise RuntimeError("out and gate_out must have the same shape as x")
 
@@ -68,7 +68,10 @@ def ada_rms_norm_style_bf16(
     out: torch.Tensor | None = None,
     gate_out: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Apply RMSNorm, style scale/shift, and return the style gate."""
+    """Apply RMSNorm and style modulation.
+
+    ``style`` may be per-row ``(rows, 3 * dim)`` or broadcast ``(1, 3 * dim)``.
+    """
 
     if out is None:
         out = torch.empty_like(x)
@@ -89,7 +92,10 @@ def gate_residual_ada_norm_fp8_static_bf16(
     out: torch.Tensor | None = None,
     gate_out: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Update residual in place, apply AdaRMSNorm, and emit static-scale FP8."""
+    """Update residual in place, apply AdaRMSNorm, and emit static-scale FP8.
+
+    ``style`` may be per-row ``(rows, 3 * dim)`` or broadcast ``(1, 3 * dim)``.
+    """
 
     if out is None:
         out = torch.empty_like(residual, dtype=torch.float8_e4m3fn)

build/torch213-cxx11-cu130-aarch64-linux/{flashrt_adaptive_norms_test.abi3.so → _flashrt_adaptive_norms_cuda_503100c.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:de09a9bab79d28c6a9dd109c47c8e5792aa773c5d069455bd8fc808afb4c30cb
-size 240872
+oid sha256:68dbeca9872b314f4228db24b9b1b26464742c2c8ad39406ca30083041e6ead4
+size 240944

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

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

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

@@ -1,32 +1,28 @@
 {
   "name": "flashrt-adaptive-norms",
-  "id": "flashrt_adaptive_norms_test",
+  "id": "_flashrt_adaptive_norms_cuda_503100c",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
-      "11.0"
+      "11.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "kv6yHmnnIUpOpDGE4dHU9J32michTE3ygV62h/+D+RI=",
-      "flashrt_adaptive_norms_test.abi3.so": "3gmpuredKMap3RCcR8jleSqnc8XQaUVb2PyAivtMMMs=",
-      "_ops.py": "JOw3+PHqZUGoT770CFlNJojgiXwMOTvFMrvsxZXy7Xc=",
+      "__init__.py": "gldeiAjCxeb9fN4IzR7vobgvtJVqN6naBsEaE6KpOyQ=",
+      "_flashrt_adaptive_norms_cuda_503100c.abi3.so": "aNvsqYcrMU9CKNskubGyZGR0LCyK05QGyjAIMEHm6tQ=",
+      "_ops.py": "3DRost+49dPPYAX9CRnUzvmWa/YXC6pp7hI3m5DvNqY=",
       "flashrt_adaptive_norms/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "456d297",
+      "sha": "503100c",
       "dirty": false
-    },
-    "validation": {
-      "torch": "2.13.0+cu130",
-      "cuda": "13.0"
     }
   }
 }

csrc/README.md

@@ -0,0 +1,3 @@
+# csrc
+
+CUDA source for adaptive norm Tensor API wrappers.

csrc/adaptive_norms.cu

@@ -0,0 +1,218 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Tensor-facing adaptive norm kernels. The math follows FlashRT
+// csrc/kernels/norm.cu::ada_rms_norm_style and
+// csrc/kernels/fusion.cu::gate_residual_ada_norm_fp8.
+
+#include "adaptive_norms.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace adaptive_norms {
+namespace {
+
+__device__ __forceinline__ float bf16_to_f32(__nv_bfloat16 x) {
+  return __bfloat162float(x);
+}
+
+__device__ __forceinline__ __nv_bfloat16 f32_to_bf16(float x) {
+  return __float2bfloat16(x);
+}
+
+__device__ __forceinline__ float block_reduce_sum(float value, float* shared) {
+  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 + 31) >> 5)) ? shared[lane] : 0.0f;
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_xor_sync(0xffffffffu, value, offset);
+    }
+    if (lane == 0) {
+      shared[0] = value;
+    }
+  }
+  __syncthreads();
+  return shared[0];
+}
+
+__global__ void ada_rms_norm_style_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ weight,
+    const __nv_bfloat16* __restrict__ style,
+    __nv_bfloat16* __restrict__ out,
+    __nv_bfloat16* __restrict__ gate_out,
+    int style_rows,
+    int dim,
+    float eps) {
+  const int row = blockIdx.x;
+  const int dim2 = dim >> 1;
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + static_cast<long long>(row) * dim);
+  const __nv_bfloat162* w2 = reinterpret_cast<const __nv_bfloat162*>(weight);
+  const int style_row_index = style_rows == 1 ? 0 : row;
+  const __nv_bfloat16* style_row =
+      style + static_cast<long long>(style_row_index) * 3 * dim;
+  const __nv_bfloat162* scale2 = reinterpret_cast<const __nv_bfloat162*>(style_row);
+  const __nv_bfloat162* shift2 = reinterpret_cast<const __nv_bfloat162*>(style_row + dim);
+  const __nv_bfloat162* gate2 = reinterpret_cast<const __nv_bfloat162*>(style_row + 2 * dim);
+  __nv_bfloat162* out2 =
+      reinterpret_cast<__nv_bfloat162*>(out + static_cast<long long>(row) * dim);
+  __nv_bfloat162* gate_out2 =
+      reinterpret_cast<__nv_bfloat162*>(gate_out + static_cast<long long>(row) * dim);
+
+  extern __shared__ float shared[];
+  float local_sum = 0.0f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    const __nv_bfloat162 xv = x2[i];
+    const float x0 = bf16_to_f32(xv.x);
+    const float x1 = bf16_to_f32(xv.y);
+    local_sum = __fadd_rn(local_sum, __fadd_rn(__fmul_rn(x0, x0), __fmul_rn(x1, x1)));
+  }
+  const float rms = rsqrtf(block_reduce_sum(local_sum, shared) / static_cast<float>(dim) + eps);
+
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    const __nv_bfloat162 xv = x2[i];
+    const __nv_bfloat162 wv = w2[i];
+    const __nv_bfloat162 sv = scale2[i];
+    const __nv_bfloat162 hv = shift2[i];
+    const __nv_bfloat162 gv = gate2[i];
+    const float n0 = __fmul_rn(__fmul_rn(bf16_to_f32(xv.x), rms), bf16_to_f32(wv.x));
+    const float n1 = __fmul_rn(__fmul_rn(bf16_to_f32(xv.y), rms), bf16_to_f32(wv.y));
+    const float o0 = __fadd_rn(__fmul_rn(n0, __fadd_rn(1.0f, bf16_to_f32(sv.x))), bf16_to_f32(hv.x));
+    const float o1 = __fadd_rn(__fmul_rn(n1, __fadd_rn(1.0f, bf16_to_f32(sv.y))), bf16_to_f32(hv.y));
+    out2[i] = __halves2bfloat162(f32_to_bf16(o0), f32_to_bf16(o1));
+    gate_out2[i] = gv;
+  }
+}
+
+__global__ void gate_residual_ada_norm_fp8_static_bf16_kernel(
+    __nv_bfloat16* __restrict__ residual,
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ gate,
+    const __nv_bfloat16* __restrict__ weight,
+    const __nv_bfloat16* __restrict__ style,
+    __nv_fp8_e4m3* __restrict__ out,
+    __nv_bfloat16* __restrict__ gate_out,
+    int style_rows,
+    int dim,
+    float eps,
+    const float* __restrict__ scale) {
+  const int row = blockIdx.x;
+  const int dim2 = dim >> 1;
+  __nv_bfloat162* residual2 =
+      reinterpret_cast<__nv_bfloat162*>(residual + static_cast<long long>(row) * dim);
+  const __nv_bfloat162* x2 =
+      reinterpret_cast<const __nv_bfloat162*>(x + static_cast<long long>(row) * dim);
+  const __nv_bfloat162* input_gate2 =
+      reinterpret_cast<const __nv_bfloat162*>(gate + static_cast<long long>(row) * dim);
+  const __nv_bfloat162* w2 = reinterpret_cast<const __nv_bfloat162*>(weight);
+  const int style_row_index = style_rows == 1 ? 0 : row;
+  const __nv_bfloat16* style_row =
+      style + static_cast<long long>(style_row_index) * 3 * dim;
+  const __nv_bfloat162* style_scale2 = reinterpret_cast<const __nv_bfloat162*>(style_row);
+  const __nv_bfloat162* shift2 = reinterpret_cast<const __nv_bfloat162*>(style_row + dim);
+  const __nv_bfloat162* style_gate2 = reinterpret_cast<const __nv_bfloat162*>(style_row + 2 * dim);
+  __nv_fp8_e4m3* out_row = out + static_cast<long long>(row) * dim;
+  __nv_bfloat162* gate_out2 =
+      reinterpret_cast<__nv_bfloat162*>(gate_out + static_cast<long long>(row) * dim);
+
+  extern __shared__ float shared[];
+  float local_sum = 0.0f;
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    const __nv_bfloat162 rv = residual2[i];
+    const __nv_bfloat162 xv = x2[i];
+    const __nv_bfloat162 gv = input_gate2[i];
+    const float r0 = __fadd_rn(bf16_to_f32(rv.x), __fmul_rn(bf16_to_f32(xv.x), bf16_to_f32(gv.x)));
+    const float r1 = __fadd_rn(bf16_to_f32(rv.y), __fmul_rn(bf16_to_f32(xv.y), bf16_to_f32(gv.y)));
+    const __nv_bfloat16 r0_b = f32_to_bf16(r0);
+    const __nv_bfloat16 r1_b = f32_to_bf16(r1);
+    residual2[i] = __halves2bfloat162(r0_b, r1_b);
+    const float rr0 = bf16_to_f32(r0_b);
+    const float rr1 = bf16_to_f32(r1_b);
+    local_sum = __fadd_rn(local_sum, __fadd_rn(__fmul_rn(rr0, rr0), __fmul_rn(rr1, rr1)));
+  }
+  const float rms = rsqrtf(block_reduce_sum(local_sum, shared) / static_cast<float>(dim) + eps);
+  const float inv_scale = 1.0f / (*scale);
+
+  for (int i = threadIdx.x; i < dim2; i += blockDim.x) {
+    const __nv_bfloat162 rv = residual2[i];
+    const __nv_bfloat162 wv = w2[i];
+    const __nv_bfloat162 sv = style_scale2[i];
+    const __nv_bfloat162 hv = shift2[i];
+    const __nv_bfloat162 gv = style_gate2[i];
+    const float n0 = __fmul_rn(__fmul_rn(bf16_to_f32(rv.x), rms), bf16_to_f32(wv.x));
+    const float n1 = __fmul_rn(__fmul_rn(bf16_to_f32(rv.y), rms), bf16_to_f32(wv.y));
+    const float v0 = __fmul_rn(__fadd_rn(__fmul_rn(n0, __fadd_rn(1.0f, bf16_to_f32(sv.x))), bf16_to_f32(hv.x)), inv_scale);
+    const float v1 = __fmul_rn(__fadd_rn(__fmul_rn(n1, __fadd_rn(1.0f, bf16_to_f32(sv.y))), bf16_to_f32(hv.y)), inv_scale);
+    out_row[2 * i] = __nv_fp8_e4m3(fminf(fmaxf(v0, -448.0f), 448.0f));
+    out_row[2 * i + 1] = __nv_fp8_e4m3(fminf(fmaxf(v1, -448.0f), 448.0f));
+    gate_out2[i] = gv;
+  }
+}
+
+}  // namespace
+
+void ada_rms_norm_style_bf16(
+    const void* x,
+    const void* weight,
+    const void* style,
+    void* out,
+    void* gate_out,
+    int rows,
+    int style_rows,
+    int dim,
+    float eps,
+    cudaStream_t stream) {
+  if (rows <= 0 || dim <= 0) return;
+  ada_rms_norm_style_bf16_kernel<<<rows, 256, 256 * sizeof(float), stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(x),
+      reinterpret_cast<const __nv_bfloat16*>(weight),
+      reinterpret_cast<const __nv_bfloat16*>(style),
+      reinterpret_cast<__nv_bfloat16*>(out),
+      reinterpret_cast<__nv_bfloat16*>(gate_out),
+      style_rows,
+      dim,
+      eps);
+}
+
+void gate_residual_ada_norm_fp8_static_bf16(
+    void* residual,
+    const void* x,
+    const void* gate,
+    const void* weight,
+    const void* style,
+    void* out,
+    void* gate_out,
+    int rows,
+    int style_rows,
+    int dim,
+    float eps,
+    const float* scale,
+    cudaStream_t stream) {
+  if (rows <= 0 || dim <= 0) return;
+  gate_residual_ada_norm_fp8_static_bf16_kernel<<<rows, 256, 256 * sizeof(float), stream>>>(
+      reinterpret_cast<__nv_bfloat16*>(residual),
+      reinterpret_cast<const __nv_bfloat16*>(x),
+      reinterpret_cast<const __nv_bfloat16*>(gate),
+      reinterpret_cast<const __nv_bfloat16*>(weight),
+      reinterpret_cast<const __nv_bfloat16*>(style),
+      reinterpret_cast<__nv_fp8_e4m3*>(out),
+      reinterpret_cast<__nv_bfloat16*>(gate_out),
+      style_rows,
+      dim,
+      eps,
+      scale);
+}
+
+}  // namespace adaptive_norms
+}  // namespace flash_rt

csrc/adaptive_norms.cuh

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace adaptive_norms {
+
+void ada_rms_norm_style_bf16(
+    const void* x,
+    const void* weight,
+    const void* style,
+    void* out,
+    void* gate_out,
+    int rows,
+    int style_rows,
+    int dim,
+    float eps,
+    cudaStream_t stream);
+
+void gate_residual_ada_norm_fp8_static_bf16(
+    void* residual,
+    const void* x,
+    const void* gate,
+    const void* weight,
+    const void* style,
+    void* out,
+    void* gate_out,
+    int rows,
+    int style_rows,
+    int dim,
+    float eps,
+    const float* scale,
+    cudaStream_t stream);
+
+}  // namespace adaptive_norms
+}  // namespace flash_rt

examples/README.md

@@ -0,0 +1,9 @@
+# examples
+
+Run the minimal adaptive-norm example after installing `kernels` or after
+adding a built artifact to `PYTHONPATH`:
+
+```bash
+python flashrt-adaptive-norms/examples/adaptive_norms.py
+python flashrt-adaptive-norms/examples/adaptive_norms.py --compile
+```

examples/adaptive_norms.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+"""Minimal HF Hub-style adaptive norms example."""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+
+try:
+    from kernels import get_kernel
+except ModuleNotFoundError:
+    get_kernel = None
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo-id", default="flashrt/flashrt-adaptive-norms")
+    parser.add_argument("--version", type=int, default=1)
+    parser.add_argument("--rows", type=int, default=2520)
+    parser.add_argument("--dim", type=int, default=3072)
+    parser.add_argument("--compile", action="store_true")
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+
+    if get_kernel is not None:
+        ops = get_kernel(args.repo_id, version=args.version, trust_remote_code=True)
+    else:
+        import flashrt_adaptive_norms as ops
+
+    x = torch.randn((args.rows, args.dim), device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    gate = torch.randn_like(x)
+    weight = torch.ones((args.dim,), device="cuda", dtype=torch.bfloat16)
+    style = torch.randn((args.rows, 3 * args.dim), device="cuda", dtype=torch.bfloat16)
+    scale = torch.tensor([0.04], device="cuda", dtype=torch.float32)
+
+    fn = ops.gate_residual_ada_norm_fp8_static_bf16
+    if args.compile:
+        fn = torch.compile(fn, fullgraph=True, mode="reduce-overhead")
+
+    residual, out_fp8, gate_out = fn(residual, x, gate, weight, style, scale)
+    torch.cuda.synchronize()
+    mode = "torch.compile(fullgraph=True)" if args.compile else "eager"
+    print(
+        f"{mode} residual={tuple(residual.shape)} out={tuple(out_fp8.shape)} "
+        f"gate={tuple(gate_out.shape)} dtype={out_fp8.dtype}"
+    )
+
+
+if __name__ == "__main__":
+    main()

flake.lock

@@ -0,0 +1,118 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1767039857,
+        "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      },
+      "locked": {
+        "lastModified": 1780582651,
+        "narHash": "sha256-WLVlcWJONDeYTw/m4/mbouy0yBqM1LRKGRrKwN6JTJ0=",
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "432702bfbfbb17d3a1bd2c2743d004e21e769ab7",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "432702bfbfbb17d3a1bd2c2743d004e21e769ab7",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1776927958,
+        "narHash": "sha256-XOzEtft7E0P6TgQViLUOQeGHlEYiQ0+FY24BPEksj6s=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "fec2c46cca5bf9767486a290abae51200b656d69",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1776914043,
+        "narHash": "sha256-qug5r56yW1qOsjSI99l3Jm15JNT9CvS2otkXNRNtrPI=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "2d35c4358d7de3a0e606a6e8b27925d981c01cc3",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

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

scripts/README.md

@@ -0,0 +1,3 @@
+# scripts
+
+Package-local sync or audit scripts can live here.

tests/README.md

@@ -0,0 +1,4 @@
+# tests
+
+Correctness tests for adaptive norm APIs, including BF16 AdaRMSNorm/style and
+fused residual/AdaRMSNorm/static-FP8 output.

tests/test_adaptive_norms.py

@@ -0,0 +1,266 @@
+#!/usr/bin/env python3
+"""Correctness tests for flashrt-adaptive-norms."""
+
+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 / "flashrt-adaptive-norms"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def ada_rms_norm_style_bf16(self, x, weight, style, eps=1e-6, out=None, gate_out=None):
+        if out is None:
+            out = torch.empty_like(x)
+        if gate_out is None:
+            gate_out = torch.empty_like(x)
+        self._ops.ada_rms_norm_style_bf16(x, weight, style, float(eps), out, gate_out)
+        return out, gate_out
+
+    def gate_residual_ada_norm_fp8_static_bf16(
+        self, residual, x, gate, weight, style, scale, eps=1e-6, out=None, gate_out=None
+    ):
+        if out is None:
+            out = torch.empty_like(residual, dtype=torch.float8_e4m3fn)
+        if gate_out is None:
+            gate_out = torch.empty_like(residual)
+        self._ops.gate_residual_ada_norm_fp8_static_bf16(
+            residual, x, gate, weight, style, scale, float(eps), out, gate_out
+        )
+        return residual, out, gate_out
+
+
+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)
+    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}")
+    _preload_cublaslt()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
+    namespace = "flashrt_adaptive_norms_test"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "adaptive_norms.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("flashrt_adaptive_norms")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def make_case(rows: int, dim: int):
+    x = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    gate = torch.randn_like(x)
+    weight = (1.0 + 0.1 * torch.randn((dim,), device="cuda", dtype=torch.bfloat16)).contiguous()
+    style = torch.randn((rows, 3 * dim), device="cuda", dtype=torch.bfloat16) * 0.05
+    scale = torch.tensor([0.04], device="cuda", dtype=torch.float32)
+    return x, residual, gate, weight, style.contiguous(), scale
+
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float):
+    rms = torch.rsqrt(torch.mean(x.float() * x.float(), dim=-1, keepdim=True) + eps)
+    return x.float() * rms * weight.float()
+
+
+def ref_ada(x, weight, style, eps):
+    dim = x.shape[1]
+    scale = style[:, :dim].float()
+    shift = style[:, dim : 2 * dim].float()
+    gate = style[:, 2 * dim :].expand(x.shape[0], -1).contiguous()
+    normed = rms_norm(x, weight, eps)
+    return (normed * (1.0 + scale) + shift).to(torch.bfloat16), gate.to(torch.bfloat16)
+
+
+def ref_gate_residual_fp8(residual, x, gate, weight, style, scale, eps):
+    updated = (residual.float() + x.float() * gate.float()).to(torch.bfloat16)
+    dim = updated.shape[1]
+    style_scale = style[:, :dim].float()
+    shift = style[:, dim : 2 * dim].float()
+    gate_out = (
+        style[:, 2 * dim :]
+        .expand(updated.shape[0], -1)
+        .contiguous()
+        .to(torch.bfloat16)
+    )
+    normed = rms_norm(updated, weight, eps)
+    fp8 = ((normed * (1.0 + style_scale) + shift) / scale.float().reshape(())).to(torch.float8_e4m3fn)
+    return updated, fp8, gate_out
+
+
+def percentile(x: torch.Tensor, q: float) -> torch.Tensor:
+    flat = x.flatten()
+    k = max(1, min(flat.numel(), math.ceil(q * flat.numel())))
+    return flat.kthvalue(k).values
+
+
+def assert_close_distribution(name: str, got: torch.Tensor, expected: torch.Tensor, p99_limit: float, max_limit: float) -> None:
+    diff = (got.float() - expected.float()).abs().flatten()
+    max_abs = float(diff.max().item())
+    mean_abs = float(diff.mean().item())
+    p99_abs = float(percentile(diff, 0.99).item())
+    cosine = float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item())
+    if p99_abs > p99_limit or max_abs > max_limit:
+        raise AssertionError(
+            f"{name} failed: max_abs={max_abs} mean_abs={mean_abs} p99_abs={p99_abs} cosine={cosine}"
+        )
+    print(
+        f"PASS {name}: max_abs={max_abs:.6f} mean_abs={mean_abs:.6f} "
+        f"p99_abs={p99_abs:.6f} cosine={cosine:.8f}"
+    )
+
+
+def assert_fp8_boundary_distribution(name: str, got: torch.Tensor, expected: torch.Tensor) -> None:
+    diff = (got.float() - expected.float()).abs().flatten()
+    nonzero = int((diff != 0).sum().item())
+    max_nonzero = max(8, diff.numel() // 100000)
+    max_abs = float(diff.max().item())
+    mean_abs = float(diff.mean().item())
+    p99_abs = float(percentile(diff, 0.99).item())
+    cosine = float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item())
+    if p99_abs > 0.0 or nonzero > max_nonzero:
+        raise AssertionError(
+            f"{name} failed: max_abs={max_abs} mean_abs={mean_abs} p99_abs={p99_abs} "
+            f"nonzero={nonzero}/{diff.numel()} cosine={cosine}"
+        )
+    print(
+        f"PASS {name}: max_abs={max_abs:.6f} mean_abs={mean_abs:.6f} "
+        f"p99_abs={p99_abs:.6f} nonzero={nonzero}/{diff.numel()} cosine={cosine:.8f}"
+    )
+
+
+def run_shape(ops, label: str, rows: int, dim: int, eps: float, broadcast: bool = False) -> None:
+    x, residual, gate, weight, style, scale = make_case(rows, dim)
+    if broadcast:
+        style = style[:1].contiguous()
+    got, got_gate = ops.ada_rms_norm_style_bf16(x, weight, style, eps)
+    exp, exp_gate = ref_ada(x, weight, style, eps)
+    assert_close_distribution(f"{label}/ada_out", got, exp, 0.0, 0.015625)
+    assert_close_distribution(f"{label}/ada_gate", got_gate, exp_gate, 0.0, 0.0)
+
+    residual_input = residual.clone()
+    got_residual, got_fp8, got_gate2 = ops.gate_residual_ada_norm_fp8_static_bf16(
+        residual, x, gate, weight, style, scale, eps
+    )
+    exp_residual, exp_fp8, exp_gate2 = ref_gate_residual_fp8(
+        residual_input, x, gate, weight, style, scale, eps
+    )
+    assert_close_distribution(f"{label}/fused_residual", got_residual, exp_residual, 0.0, 0.0)
+    assert_fp8_boundary_distribution(f"{label}/fused_fp8", got_fp8.float(), exp_fp8.float())
+    assert_close_distribution(f"{label}/fused_gate", got_gate2, exp_gate2, 0.0, 0.0)
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+    torch.manual_seed(47)
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    shapes = {
+        "small": (64, 1024),
+        "vla_2k": (2520, 3072),
+        "vla_4k": (4096, 3072),
+    }
+    if args.mode == "smoke":
+        shapes = {"small": shapes["small"]}
+    for label, (rows, dim) in shapes.items():
+        run_shape(ops, label, rows, dim, args.eps)
+        run_shape(ops, f"{label}_broadcast", rows, dim, args.eps, broadcast=True)
+
+    x, residual, gate, weight, style, scale = make_case(64, 1024)
+    style = style[:1].contiguous()
+    eager_out, eager_gate = ops.ada_rms_norm_style_bf16(
+        x, weight, style, args.eps
+    )
+    graph_out = torch.empty_like(x)
+    graph_gate = torch.empty_like(x)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        ops.ada_rms_norm_style_bf16(
+            x, weight, style, args.eps, graph_out, graph_gate
+        )
+    graph.replay()
+    if not torch.equal(graph_out, eager_out) or not torch.equal(
+        graph_gate, eager_gate
+    ):
+        raise AssertionError("broadcast AdaRMS CUDA Graph mismatch")
+
+    if args.backend == "installed":
+        def invoke(value, norm_weight, style_value):
+            return ops.ada_rms_norm_style_bf16(
+                value, norm_weight, style_value, args.eps
+            )
+
+        compiled_out, compiled_gate = torch.compile(
+            invoke, fullgraph=True
+        )(x, weight, style)
+        if not torch.equal(compiled_out, eager_out) or not torch.equal(
+            compiled_gate, eager_gate
+        ):
+            raise AssertionError("broadcast AdaRMS torch.compile mismatch")
+
+
+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-6)
+    args = parser.parse_args()
+    run(args)
+
+
+if __name__ == "__main__":
+    main()

torch-ext/README.md

@@ -0,0 +1,3 @@
+# torch-ext
+
+PyTorch Tensor bindings and Python wrappers for `flashrt-adaptive-norms`.

torch-ext/flashrt_adaptive_norms/__init__.py

@@ -0,0 +1,121 @@
+"""FlashRT adaptive norm kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_adaptive_shapes(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    out: torch.Tensor,
+    gate_out: torch.Tensor,
+) -> None:
+    if x.dim() != 2:
+        raise RuntimeError("x must have shape (rows, dim)")
+    rows, dim = x.shape
+    if dim % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if weight.shape != (dim,):
+        raise RuntimeError("weight must have shape (dim,)")
+    if style.shape not in ((rows, 3 * dim), (1, 3 * dim)):
+        raise RuntimeError("style must have shape (rows, 3 * dim) or (1, 3 * dim)")
+    if out.shape != x.shape or gate_out.shape != x.shape:
+        raise RuntimeError("out and gate_out must have the same shape as x")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ada_rms_norm_style_bf16"))
+def _ada_rms_norm_style_bf16_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+    gate_out: torch.Tensor,
+) -> None:
+    _check_adaptive_shapes(x, weight, style, out, gate_out)
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gate_residual_ada_norm_fp8_static_bf16"))
+def _gate_residual_ada_norm_fp8_static_bf16_fake(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+    gate_out: torch.Tensor,
+) -> None:
+    _check_adaptive_shapes(residual, weight, style, out, gate_out)
+    if x.shape != residual.shape or gate.shape != residual.shape:
+        raise RuntimeError("x and gate must have the same shape as residual")
+    if scale.numel() != 1:
+        raise RuntimeError("scale must be a scalar tensor")
+    return None
+
+
+def ada_rms_norm_style_bf16(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    eps: float = 1e-6,
+    out: torch.Tensor | None = None,
+    gate_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply RMSNorm and style modulation.
+
+    ``style`` may be per-row ``(rows, 3 * dim)`` or broadcast ``(1, 3 * dim)``.
+    """
+
+    if out is None:
+        out = torch.empty_like(x)
+    if gate_out is None:
+        gate_out = torch.empty_like(x)
+    ops.ada_rms_norm_style_bf16(x, weight, style, float(eps), out, gate_out)
+    return out, gate_out
+
+
+def gate_residual_ada_norm_fp8_static_bf16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float = 1e-6,
+    out: torch.Tensor | None = None,
+    gate_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Update residual in place, apply AdaRMSNorm, and emit static-scale FP8.
+
+    ``style`` may be per-row ``(rows, 3 * dim)`` or broadcast ``(1, 3 * dim)``.
+    """
+
+    if out is None:
+        out = torch.empty_like(residual, dtype=torch.float8_e4m3fn)
+    if gate_out is None:
+        gate_out = torch.empty_like(residual)
+    ops.gate_residual_ada_norm_fp8_static_bf16(
+        residual,
+        x,
+        gate,
+        weight,
+        style,
+        scale,
+        float(eps),
+        out,
+        gate_out,
+    )
+    return residual, out, gate_out
+
+
+__all__ = [
+    "ada_rms_norm_style_bf16",
+    "gate_residual_ada_norm_fp8_static_bf16",
+]

torch-ext/torch_binding.cpp

@@ -0,0 +1,175 @@
+// 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 "adaptive_norms.cuh"
+#include "registration.h"
+#include "torch_binding.h"
+
+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_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_matrix(torch::Tensor const& tensor, const char* name) {
+  check_bf16(tensor, name);
+  TORCH_CHECK(tensor.dim() == 2, name, " must have shape (rows, dim)");
+  TORCH_CHECK(tensor.size(0) > 0 && tensor.size(1) > 0,
+              name, " rows and dim must be positive");
+  TORCH_CHECK((tensor.size(1) % 2) == 0, name, ".shape[1] must be even");
+}
+
+void check_common(
+    torch::Tensor const& x,
+    torch::Tensor const& weight,
+    torch::Tensor const& style,
+    torch::Tensor const& out,
+    torch::Tensor const& gate_out) {
+  check_matrix(x, "x");
+  check_bf16(weight, "weight");
+  check_bf16(style, "style");
+  check_bf16(gate_out, "gate_out");
+  const int64_t rows = x.size(0);
+  const int64_t dim = x.size(1);
+  TORCH_CHECK(weight.dim() == 1 && weight.size(0) == dim,
+              "weight must have shape (dim,)");
+  TORCH_CHECK(style.dim() == 2 && style.size(1) == 3 * dim &&
+                  (style.size(0) == 1 || style.size(0) == rows),
+              "style must have shape (rows, 3 * dim) or (1, 3 * dim)");
+  TORCH_CHECK(gate_out.sizes() == x.sizes(),
+              "gate_out must have the same shape as x");
+  check_same_device(x, weight, "x", "weight");
+  check_same_device(x, style, "x", "style");
+  check_same_device(x, out, "x", "out");
+  check_same_device(x, gate_out, "x", "gate_out");
+}
+
+}  // namespace
+
+void ada_rms_norm_style_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& weight,
+    torch::Tensor const& style,
+    double eps,
+    torch::Tensor& out,
+    torch::Tensor& gate_out) {
+  check_bf16(out, "out");
+  check_common(x, weight, style, out, gate_out);
+  TORCH_CHECK(out.sizes() == x.sizes(), "out must have the same shape as x");
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  flash_rt::adaptive_norms::ada_rms_norm_style_bf16(
+      x.data_ptr(),
+      weight.data_ptr(),
+      style.data_ptr(),
+      out.data_ptr(),
+      gate_out.data_ptr(),
+      static_cast<int>(x.size(0)),
+      static_cast<int>(style.size(0)),
+      static_cast<int>(x.size(1)),
+      static_cast<float>(eps),
+      stream);
+#else
+  TORCH_CHECK(false, "flashrt-adaptive-norms was not built with CUDA support");
+#endif
+}
+
+void gate_residual_ada_norm_fp8_static_bf16(
+    torch::Tensor& residual,
+    torch::Tensor const& x,
+    torch::Tensor const& gate,
+    torch::Tensor const& weight,
+    torch::Tensor const& style,
+    torch::Tensor const& scale,
+    double eps,
+    torch::Tensor& out,
+    torch::Tensor& gate_out) {
+  check_matrix(residual, "residual");
+  check_matrix(x, "x");
+  check_bf16(gate, "gate");
+  check_fp32(scale, "scale");
+  check_fp8(out, "out");
+  check_common(residual, weight, style, out, gate_out);
+  TORCH_CHECK(x.sizes() == residual.sizes(), "x must have the same shape as residual");
+  TORCH_CHECK(gate.sizes() == residual.sizes(), "gate must have the same shape as residual");
+  TORCH_CHECK(out.sizes() == residual.sizes(), "out must have the same shape as residual");
+  TORCH_CHECK(scale.numel() == 1, "scale must be a scalar tensor");
+  check_same_device(residual, x, "residual", "x");
+  check_same_device(residual, gate, "residual", "gate");
+  check_same_device(residual, scale, "residual", "scale");
+
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(residual.device());
+  auto stream = at::cuda::getCurrentCUDAStream(residual.get_device()).stream();
+  flash_rt::adaptive_norms::gate_residual_ada_norm_fp8_static_bf16(
+      residual.data_ptr(),
+      x.data_ptr(),
+      gate.data_ptr(),
+      weight.data_ptr(),
+      style.data_ptr(),
+      out.data_ptr(),
+      gate_out.data_ptr(),
+      static_cast<int>(residual.size(0)),
+      static_cast<int>(style.size(0)),
+      static_cast<int>(residual.size(1)),
+      static_cast<float>(eps),
+      reinterpret_cast<const float*>(scale.data_ptr()),
+      stream);
+#else
+  TORCH_CHECK(false, "flashrt-adaptive-norms was not built with CUDA support");
+#endif
+}
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("ada_rms_norm_style_bf16("
+          "Tensor x, Tensor weight, Tensor style, float eps, Tensor! out, Tensor! gate_out) -> ()");
+  ops.def("gate_residual_ada_norm_fp8_static_bf16("
+          "Tensor! residual, Tensor x, Tensor gate, Tensor weight, Tensor style, Tensor scale, "
+          "float eps, Tensor! out, Tensor! gate_out) -> ()");
+#if defined(CUDA_KERNEL)
+  ops.impl("ada_rms_norm_style_bf16",
+           torch::kCUDA,
+           &ada_rms_norm_style_bf16);
+  ops.impl("gate_residual_ada_norm_fp8_static_bf16",
+           torch::kCUDA,
+           &gate_residual_ada_norm_fp8_static_bf16);
+#endif
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <torch/all.h>
+
+void ada_rms_norm_style_bf16(
+    torch::Tensor const& x,
+    torch::Tensor const& weight,
+    torch::Tensor const& style,
+    double eps,
+    torch::Tensor& out,
+    torch::Tensor& gate_out);
+
+void gate_residual_ada_norm_fp8_static_bf16(
+    torch::Tensor& residual,
+    torch::Tensor const& x,
+    torch::Tensor const& gate,
+    torch::Tensor const& weight,
+    torch::Tensor const& style,
+    torch::Tensor const& scale,
+    double eps,
+    torch::Tensor& out,
+    torch::Tensor& gate_out);