Frontend APIFrontend OSS APIs

RMSNorm + RHT + Amax (SM100)

View as Markdown

This is an experimental API and subject to change.

Overview

RMSNorm + RHT + amax: A fused CUTE DSL kernel for NVIDIA Blackwell GPUs (SM100+) that applies RMS normalization, a block-diagonal Hadamard transform with fixed block size 16, and a per-CTA amax reduction.

This frontend integration exposes the kernel as a standard FE-OSS Python API with:

  • a class API (RmsNormRhtAmaxSm100)
  • a wrapper API (rmsnorm_rht_amax_wrapper_sm100)
  • grouped-gemm-style regression coverage for compile/execute, wrapper use, and cache reuse

Shapes

  • Inputs

    • X: activation tensor, shape (M, N)
    • W: RMSNorm scale tensor, shape (N,)

  • Outputs

    • O: fused RMSNorm + RHT output tensor, shape (M, N)
    • Amax: per-CTA max-abs tensor, shape (M / rows_per_cta,)

rows_per_cta is the number of rows reduced into each amax element.

Equations

For each row m:

RMS(Xm)=1N∑n=0N−1X[m,n]2+ϵ\mathrm{RMS}(X_m) = \sqrt{\frac{1}{N}\sum_{n=0}^{N-1} X[m, n]^2 + \epsilon} Y[m,n]=X[m,n]RMS(Xm)⋅W[n]Y[m, n] = \frac{X[m, n]}{\mathrm{RMS}(X_m)} \cdot W[n]

Then apply the fixed Hadamard transform blockwise over 16-wide chunks:

O[m,b]=Y[m,b]⋅H16/16O[m, b] = Y[m, b] \cdot H_{16} / \sqrt{16}

where H_16 is the 16 x 16 Hadamard matrix and b indexes each 16-element block in the hidden dimension.

For each CTA covering rows_per_cta rows:

Amax[c]=max⁡∣O∣\mathrm{Amax}[c] = \max |O|

over every element produced by that CTA.

API Usage

High-level wrapper

from cudnn import rmsnorm_rht_amax_wrapper_sm100
result = rmsnorm_rht_amax_wrapper_sm100(
    x_tensor=x,
    w_tensor=w,
    eps=1e-5,
    num_threads=None,   # optional override
    rows_per_cta=None,  # optional override
    current_stream=None,
)
o_tensor, amax_tensor = result

When no overrides are supplied, the wrapper uses the upstream-tuned thread table when available and an upstream-style rows_per_cta heuristic.

Class API

from cudnn import RmsNormRhtAmaxSm100
op = RmsNormRhtAmaxSm100(
    sample_x=x,
    sample_w=w,
    sample_o=o,
    sample_amax=amax,
    eps=1e-5,
    num_threads=128,
    rows_per_cta=2,
)
assert op.check_support()
op.compile()
op.execute(
    x_tensor=x,
    w_tensor=w,
    o_tensor=o,
    amax_tensor=amax,
    current_stream=None,
)

Parameters

Input and output tensors

  • x_tensor / sample_x
    • Shape: (M, N)
    • Layout: row-major contiguous
    • Dtype: torch.bfloat16
  • w_tensor / sample_w
    • Shape: (N,)
    • Layout: contiguous
    • Dtype: torch.bfloat16
  • o_tensor / sample_o
    • Shape: (M, N)
    • Layout: row-major contiguous
    • Dtype: torch.bfloat16
  • amax_tensor / sample_amax
    • Shape: (M / rows_per_cta,)
    • Dtype: torch.float32

Common parameters

  • eps: float
    • RMSNorm epsilon. Default: 1e-5
  • num_threads: Optional[int]
    • Threads per CTA. If omitted, the API uses the upstream-tuned table when possible, otherwise a valid fallback search.
  • rows_per_cta: Optional[int]
    • Rows processed by each CTA. If omitted, the wrapper uses the upstream-style heuristic over {2, 4, 8}.
  • CUDA stream (current_stream)

Wrapper return values

Returns a TupleDict with keys:

  • o_tensor
  • amax_tensor

Tuple unpacking order is (o_tensor, amax_tensor).

Support surface and constraints

  • Requires SM100+.
  • N must be divisible by 16.
  • N must be divisible by the resolved num_threads.
  • EPT = N / num_threads must be at least 8 and divisible by 8.
  • M must be divisible by rows_per_cta.
  • Inputs and output are currently bf16 only.
  • The frontend integration matches the upstream RMSNorm kernel semantics; it does not expose full LayerNorm mean/bias behavior.

Verification

Focused correctness and cache coverage live in:

  • test/python/fe_api/norm/test_rmsnorm_rht_amax.py