Merge branch 'main' into fix/rope-dead-constexpr-batch-size-dynamic-shapes

README.md

@@ -150,7 +150,7 @@ pip install -e ".[dev]" --extra-index-url https://triton-ascend.osinfra.cn/pypi/
 - `transformers >= 4.x`: Required if you plan to use the transformers models patching APIs. The specific model you are working will dictate the minimum version of transformers.
 - `cuda-tile`: Required when enabling the optional cuTile backend on CUDA. Use this when your environment already provides CUDA Toolkit 13.1 or newer, or an existing tileiras compiler installation.
 - `cuda-tile[tileiras]`: Required when enabling the optional cuTile backend with the tileiras compiler installed directly into your Python environment.
-- `nvidia-cutlass-dsl`: Required when enabling the optional CuTe DSL backend on CUDA (the CUDA-only Python DSL shipped with NVIDIA CUTLASS, `import cutlass.cute`). Targets Hopper (SM90) and Blackwell (SM100/SM110).
+- `nvidia-cutlass-dsl >= 4.5.2`: Required when enabling the optional CuTe DSL backend on CUDA (the CUDA-only Python DSL shipped with NVIDIA CUTLASS, `import cutlass.cute`). Targets Hopper (SM90) and Blackwell (SM100/SM110).
 
 > **Note:**
 > Our kernels inherit the full spectrum of hardware compatibility offered by [Triton](https://github.com/triton-lang/triton).
@@ -217,14 +217,42 @@ pip install "liger-kernel[cutedsl]"
 LIGER_KERNEL_IMPL=cutedsl python your_script.py
 ```
 
-It currently provides genuine `cutlass.cute` implementations of **RMSNorm** and **cross entropy**. Any op without a CuTe DSL kernel transparently falls back to the default Triton kernel, so selecting the backend is always safe. Selecting it on a non-CUDA device, or without `nvidia-cutlass-dsl` installed, raises an error.
-
-
-## Getting Started
-
-There are a couple of ways to apply Liger kernels, depending on the level of customization required.
-
-### 1. Use AutoLigerKernelForCausalLM
+It currently provides genuine `cutlass.cute` implementations of **RMSNorm**, **cross entropy**, and
+Hopper **fused scaled cross entropy** (`LigerFusedScaledCrossEntropySM90Function`). The SM90 scaled
+kernel is an *additional* operator, not a replacement for the Triton `LigerFusedLinearCrossEntropyFunction`:
+it fuses `x @ weight.T` with the softmax, takes BF16 `input`/`weight` plus `ignore_index`, and returns
+the per-token negative log-likelihood `[M]` plus optional vocabulary entropy — reductions are composed in PyTorch. Backward
+therefore receives the per-token upstream gradient `[M]`, which is used verbatim as the row scale of
+`dZ`; a scalar `grad_output` is rejected rather than silently reduced. `temperature` defaults to
+`1.0` and applies the same `logits / temperature` semantics as Verl.
+Set `return_entropy=True` to additionally return differentiable per-token
+vocabulary entropy in BF16; its backward contribution follows Verl's PPO
+formula. Both NLL and entropy are zeroed for `ignore_index` rows.
+`m_tiles_per_cluster` (>= 1,
+default `1`) selects the forward schedule: `1` uses the fastest M-outer self-TMA forward, values `> 1`
+use the M-fast forward with that many live M tiles per cluster. Backward executes `dZ`, `dX = dZ @ W`,
+and `dW = dZ.T @ X` inside one persistent cluster-2 CUDA kernel. Its device-side wave loop aliases
+phase-specific shared memory, uses W multicast for dX, transposed TMA for dW, and HBM atomics between
+waves. Backward uses a fixed reusable BF16 `dZ` workspace covering 1024 tokens; wave zero
+plain-stores BF16 `dW`, and later waves use Hopper BF16 TMA reduce-add. Ops without a CuTe DSL kernel
+transparently fall back to the default Triton kernel. Selecting the backend on a non-CUDA device, or
+without `nvidia-cutlass-dsl` installed, raises an error.
+
+```python
+from liger_kernel.ops.cutedsl.ops import LigerFusedScaledCrossEntropySM90Function
+
+nll, entropy = LigerFusedScaledCrossEntropySM90Function.apply(
+    x, weight, target, 1.0, -100, 1, True
+)  # [M], [M]
+loss = nll.sum() / (target != -100).sum().clamp_min(1)
+```
+
+
+## Getting Started
+
+There are a couple of ways to apply Liger kernels, depending on the level of customization required.
+
+### 1. Use AutoLigerKernelForCausalLM
 
 Using the `AutoLigerKernelForCausalLM` is the simplest approach, as you don't have to import a model-specific patching API. If the model type is supported, the modeling code will be automatically patched using the default settings.
 
@@ -459,5 +487,3 @@ url={https://openreview.net/forum?id=36SjAIT42G}
         ↑ Back to Top ↑
     </a>
 </p>
-
-

benchmark/scripts/benchmark_fused_scaled_cross_entropy_sm90.py

@@ -0,0 +1,287 @@
+"""Benchmark the Hopper (SM90) CuTe DSL fused *scaled* cross entropy.
+
+The measured operation returns the **per-token NLL** ``[M]`` (never a mean or a
+sum), so the backward is driven by a *fixed* per-token gradient vector ``[M]``
+-- the same vector for every provider, every repetition and every operation
+mode -- instead of the shared harness' fresh ``randn_like`` draw.  That keeps
+the three providers numerically comparable and keeps the gradient allocation
+out of the timed region.
+
+Sweep parameters (all optional; defaults reproduce the previous behaviour)::
+
+    --tokens M [M ...]              total tokens (the x axis)
+    --hidden H                      hidden size
+    --vocab V                       vocabulary size
+    --providers NAME [NAME ...]     explicit provider list
+
+Providers: ``torch``, ``liger`` (Triton FLCE, ``reduction="none"``),
+and ``cutedsl-sm90`` (fixed 1024-token wave-batched backward).
+
+Example::
+
+    python benchmark_fused_scaled_cross_entropy_sm90.py \\
+        --tokens 4096 --hidden 4096 --vocab 131072 \\
+        --providers torch liger cutedsl-sm90
+"""
+
+import argparse
+import sys
+
+import torch
+
+from benchmark_model_configs import MODEL_REGISTRY
+from benchmark_model_configs import build_model_config_sweep
+from benchmark_model_configs import build_token_length_sweep
+from benchmark_model_configs import get_benchmark_model_config
+from utils import QUANTILES
+from utils import SingleBenchmarkRunInput
+from utils import SingleBenchmarkRunOutput
+from utils import _test_memory
+from utils import parse_benchmark_script_args
+from utils import run_benchmarks
+
+from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import LigerFusedScaledCrossEntropySM90Function
+from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+from liger_kernel.utils import infer_device
+
+device = infer_device()
+
+REPRESENTATIVE_CONFIG = {
+    "hidden_size": 4096,
+    "vocab_size": 131072,
+    "dtype": torch.bfloat16,
+}
+
+CUTEDSL_PREFIX = "cutedsl-sm90"
+# Seed of the fixed per-token upstream gradient, so every provider and every
+# repetition sees byte-identical ``d(NLL)/d(loss)`` values.
+GRAD_SEED = 1234
+
+
+class TorchLMHeadCE(torch.nn.Module):
+    """Per-token (``reduction="none"``) torch baseline."""
+
+    def __init__(self, hidden_size: int, vocab_size: int, dtype: torch.dtype):
+        super().__init__()
+        self.linear = torch.nn.Linear(hidden_size, vocab_size, bias=False, dtype=dtype)
+
+    def forward(self, x, target):
+        logits = self.linear(x).float()
+        return torch.nn.functional.cross_entropy(logits, target, reduction="none")
+
+
+class TritonLMHeadCE(torch.nn.Module):
+    """Triton FLCE with ``reduction="none"`` so the semantics match."""
+
+    def __init__(self, hidden_size: int, vocab_size: int, dtype: torch.dtype):
+        super().__init__()
+        self.linear = torch.nn.Linear(hidden_size, vocab_size, bias=False, dtype=dtype)
+        self.cross_entropy = LigerFusedLinearCrossEntropyLoss(reduction="none", accum_dtype=torch.float32)
+
+    def forward(self, x, target):
+        return self.cross_entropy(self.linear.weight, x, target)
+
+
+class CuteDSLHopperLMHeadScaledCE(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        vocab_size: int,
+        dtype: torch.dtype,
+        temperature: float = 1.0,
+        m_tiles_per_cluster: int = 1,
+    ):
+        super().__init__()
+        if dtype != torch.bfloat16:
+            raise TypeError(f"SM90 fused scaled cross entropy requires bfloat16, got {dtype}")
+        self.weight = torch.nn.Parameter(torch.empty(vocab_size, hidden_size, dtype=dtype))
+        self.temperature = temperature
+        self.m_tiles_per_cluster = m_tiles_per_cluster
+        torch.nn.init.normal_(self.weight, std=hidden_size**-0.5)
+
+    def forward(self, x, target):
+        return LigerFusedScaledCrossEntropySM90Function.apply(
+            x,
+            self.weight,
+            target,
+            self.temperature,
+            -100,
+            self.m_tiles_per_cluster,
+        )
+
+
+def fixed_grad_output(tokens: int) -> torch.Tensor:
+    """The fixed ``[M]`` upstream gradient of the per-token NLL."""
+    generator = torch.Generator(device=device).manual_seed(GRAD_SEED)
+    return torch.rand(tokens, generator=generator, device=device, dtype=torch.float32) + 0.25
+
+
+def setup_fused_scaled_cross_entropy_sm90(input: SingleBenchmarkRunInput):
+    """Return ``(x, forward_fn, grad_output)`` for one benchmark point."""
+    cfg = input.extra_benchmark_config
+    if isinstance(input.x, str):
+        model_cfg = MODEL_REGISTRY[input.x]
+        total_tokens = cfg["seq_len"] * cfg["bsz"]
+        hidden_size = model_cfg.hidden_size
+        vocab_size = model_cfg.vocab_size
+        dtype = model_cfg.dtype
+    else:
+        total_tokens = input.x
+        hidden_size = cfg["hidden_size"]
+        vocab_size = cfg["vocab_size"]
+        dtype = cfg["dtype"]
+
+    x = torch.randn(total_tokens, hidden_size, requires_grad=True, dtype=dtype, device=device)
+    target = torch.randint(vocab_size, (total_tokens,), dtype=torch.long, device=device)
+
+    provider = input.kernel_provider
+    if provider == CUTEDSL_PREFIX:
+        layer = CuteDSLHopperLMHeadScaledCE(hidden_size, vocab_size, dtype)
+    elif provider == "liger":
+        layer = TritonLMHeadCE(hidden_size, vocab_size, dtype)
+    elif provider == "torch":
+        layer = TorchLMHeadCE(hidden_size, vocab_size, dtype)
+    else:
+        raise ValueError(f"Unknown provider {provider!r}")
+
+    layer = layer.to(device)
+    return x, lambda: layer(x, target), fixed_grad_output(total_tokens)
+
+
+def probe_forward_fn(x, fwd_fn, grad_output):
+    """Adapter for the shared memory-probing helpers (setup returns a 3-tuple)."""
+    return fwd_fn()
+
+
+def bench_speed_fused_scaled_cross_entropy_sm90(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    import triton
+
+    x, fwd_fn, grad_output = setup_fused_scaled_cross_entropy_sm90(input)
+    mode = input.kernel_operation_mode
+
+    if mode == "forward":
+        bench_fn = fwd_fn
+    elif mode == "backward":
+        y = fwd_fn()
+
+        def bench_fn():
+            y.backward(grad_output, retain_graph=True)
+    elif mode == "full":
+
+        def bench_fn():
+            fwd_fn().backward(grad_output)
+    else:
+        raise ValueError(f"Unsupported mode: {mode}")
+
+    ms_50, ms_20, ms_80 = triton.testing.do_bench(bench_fn, grad_to_none=[x], rep=10, quantiles=QUANTILES)
+    return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80)
+
+
+def bench_memory_fused_scaled_cross_entropy_sm90(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    x, fwd_fn, grad_output = setup_fused_scaled_cross_entropy_sm90(input)
+    mode = input.kernel_operation_mode
+
+    if mode == "forward":
+        bench_fn = fwd_fn
+    elif mode == "backward":
+        y = fwd_fn()
+
+        def bench_fn():
+            x.grad = None
+            y.backward(grad_output, retain_graph=True)
+    elif mode == "full":
+
+        def bench_fn():
+            x.grad = None
+            fwd_fn().backward(grad_output)
+    else:
+        raise ValueError(f"Unsupported mode: {mode}")
+
+    mem_50, mem_20, mem_80 = _test_memory(bench_fn, quantiles=QUANTILES)
+    return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80)
+
+
+def parse_sweep_args():
+    """Script-specific flags, stripped from ``argv`` before the shared parser."""
+    parser = argparse.ArgumentParser(add_help=False, description="fused scaled cross entropy sweep options")
+    parser.add_argument("--tokens", type=int, nargs="+", default=None, help="Total tokens (M) to sweep.")
+    parser.add_argument("--hidden", type=int, default=None, help="Hidden size (H).")
+    parser.add_argument("--vocab", type=int, default=None, help="Vocabulary size (V).")
+    parser.add_argument("--providers", type=str, nargs="+", default=None, help="Explicit provider list.")
+    if any(flag in ("-h", "--help") for flag in sys.argv[1:]):
+        parser.print_help()
+        print()
+    sweep_args, remaining = parser.parse_known_args()
+    sys.argv = [sys.argv[0], *remaining]
+    return sweep_args
+
+
+def resolve_providers(sweep_args):
+    if sweep_args.providers:
+        for provider in sweep_args.providers:
+            if provider not in ("torch", "liger", CUTEDSL_PREFIX):
+                raise ValueError(f"Unknown provider {provider!r}")
+        return list(sweep_args.providers)
+    return ["torch", "liger", CUTEDSL_PREFIX]
+
+
+if __name__ == "__main__":
+    sweep_args = parse_sweep_args()
+    args = parse_benchmark_script_args()
+
+    if args.sweep_mode == "model_config":
+        common_configs = build_model_config_sweep(
+            kernel_name="fused_scaled_cross_entropy_sm90",
+            setup_fn=setup_fused_scaled_cross_entropy_sm90,
+            model_keys=["hidden_size", "vocab_size", "dtype"],
+            forward_fn=probe_forward_fn,
+            probe_provider="torch",
+            extra_configs={},
+            probe_dim="BT",
+            bt=args.bt,
+            overwrite=args.overwrite,
+        )
+    else:
+        model = get_benchmark_model_config(args.model)
+        common_configs = build_token_length_sweep(
+            kernel_name="fused_scaled_cross_entropy_sm90",
+            probe_x=1024,
+            model=model,
+            setup_fn=setup_fused_scaled_cross_entropy_sm90,
+            model_keys=["hidden_size", "vocab_size", "dtype"],
+            forward_fn=probe_forward_fn,
+            extra_configs={},
+            scale_dim="BT",
+            x_label="total tokens",
+            probe_provider="torch",
+            overwrite=args.overwrite,
+        )
+        shape = dict(REPRESENTATIVE_CONFIG)
+        if sweep_args.hidden is not None:
+            shape["hidden_size"] = sweep_args.hidden
+        if sweep_args.vocab is not None:
+            shape["vocab_size"] = sweep_args.vocab
+        if sweep_args.tokens:
+            # An explicit M sweep replaces the model-derived token lengths.
+            common_configs["x_values"] = sorted(set(sweep_args.tokens))
+            common_configs["extra_benchmark_configs"] = [shape]
+        else:
+            common_configs["x_values"] = sorted(set([*common_configs["x_values"], 4096]))
+            common_configs["extra_benchmark_configs"].append(shape)
+
+    common_configs["kernel_providers"] = resolve_providers(sweep_args)
+
+    run_benchmarks(
+        bench_test_fn=bench_speed_fused_scaled_cross_entropy_sm90,
+        kernel_operation_modes=["forward", "backward", "full"],
+        metric_name="speed",
+        metric_unit="ms",
+        **common_configs,
+    )
+    run_benchmarks(
+        bench_test_fn=bench_memory_fused_scaled_cross_entropy_sm90,
+        kernel_operation_modes=["forward", "backward", "full"],
+        metric_name="memory",
+        metric_unit="MB",
+        **common_configs,
+    )

setup.py

@@ -37,7 +37,7 @@ def get_optional_dependencies():
         "cuda-tile[tileiras]",
     ]
     cutedsl_deps = [
-        "nvidia-cutlass-dsl",
+        "nvidia-cutlass-dsl>=4.5.2",
     ]
     dev_deps = [
         "transformers>=4.52.0",

src/liger_kernel/ops/cutedsl/ops/__init__.py

@@ -14,14 +14,24 @@ except ImportError as exc:
 from liger_kernel.ops.cutedsl.ops.cross_entropy import LigerCrossEntropyFunction
 from liger_kernel.ops.cutedsl.ops.cross_entropy import cross_entropy_backward
 from liger_kernel.ops.cutedsl.ops.cross_entropy import cross_entropy_forward
+from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import LigerFusedScaledCrossEntropySM90Function
+from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import fused_scaled_cross_entropy_backward
+from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import fused_scaled_cross_entropy_forward
 from liger_kernel.ops.cutedsl.ops.rms_norm import LigerRMSNormFunction
 from liger_kernel.ops.cutedsl.ops.rms_norm import rms_norm_backward
 from liger_kernel.ops.cutedsl.ops.rms_norm import rms_norm_forward
 
+# ``LigerFusedScaledCrossEntropySM90Function`` is an *additional* CuTe DSL
+# operator (per-token NLL only, Hopper BF16); it deliberately does not replace
+# or alias the Triton ``LigerFusedLinearCrossEntropyFunction``, which keeps its
+# reduction and legacy-option surface.
 __all__ = [
     "LigerCrossEntropyFunction",
     "cross_entropy_backward",
     "cross_entropy_forward",
+    "LigerFusedScaledCrossEntropySM90Function",
+    "fused_scaled_cross_entropy_backward",
+    "fused_scaled_cross_entropy_forward",
     "LigerRMSNormFunction",
     "rms_norm_backward",
     "rms_norm_forward",

src/liger_kernel/ops/cutedsl/ops/fused_scaled_cross_entropy_sm90.py

@@ -0,0 +1,322 @@
+"""
+Hopper (SM90) CuTe DSL **fused scaled cross entropy**.
+
+This operator fuses the classifier projection ``Z = X @ W.T`` with the cross
+entropy softmax so the ``[M, V]`` logits are never written to HBM in the
+forward pass.  It is deliberately *narrower* than the Triton
+``LigerFusedLinearCrossEntropyFunction``:
+
+* the forward returns the per-token negative log-likelihood ``nll[M]`` and,
+  when requested, differentiable per-token vocabulary entropy -- never a
+  ``mean``/``sum`` reduction;
+* the backward consumes the **per-token upstream gradient** ``grad_output[M]``,
+  which is exactly the row scale of ``dZ``.  A scalar ``grad_output`` is
+  rejected instead of being silently interpreted as a reduction;
+* there is no ``bias``, ``ce_weight``, label smoothing, ``softcap``,
+  ``lse_square_scale``, z-loss or token-scaling support.
+
+Anything that needs a reduction composes it in PyTorch, e.g.::
+
+    nll = LigerFusedScaledCrossEntropySM90Function.apply(x, w, target)
+    loss = nll.sum() / (target != -100).sum().clamp_min(1)
+
+which keeps the reduction (and therefore the per-token upstream gradient) in
+autograd where it is cheap and unambiguous.
+
+Requirements
+------------
+Hopper (compute capability 9.0), BF16 ``input``/``weight``, ``input[M, H]``,
+``weight[V, H]``, ``target[M]`` int64.  ``H`` and ``V`` may be ragged; they are
+padded internally and the returned gradients are sliced back to the caller's
+shapes and dtypes.
+
+Kernels
+-------
+=========================  =========================================
+phase                      module
+=========================  =========================================
+forward (M-outer)          ``_fused_scaled_cross_entropy_forward_sm90``
+forward (M-fast)           ``_fused_scaled_cross_entropy_forward_mfast_sm90``
+backward dZ+dX+dW          ``_fused_scaled_cross_entropy_backward_fused_sm90``
+=========================  =========================================
+
+``m_tiles_per_cluster`` (>= 1, default 1) is the forward scheduling knob.  The
+default dispatches to the M-outer self-TMA forward, which is the fastest
+measured forward at ``M=4096, H=4096, V=131072``.  Values > 1 dispatch to the
+M-fast (N-outer / M-inner) forward with that many live ``M128`` online-softmax
+slots per cluster; it is correct everywhere but measured slower at the
+representative shape, and exists for tuning other shapes.
+
+Backward runs in one persistent cluster-2 CUDA kernel.  Its device-side wave
+loop executes dZ, dX, and dW with phase-local schedulers over a phase-serial
+shared-memory arena: dZ uses the static M-fast scheduler, dX interprets each
+cluster pair as W-multicast peers, and dW treats both CTAs as independent
+output-tile workers with transposed (``order=(1, 0, 2)``) TMA.  A GPU-wide HBM
+atomic barrier publishes dZ and protects the reusable workspace between waves.
+
+``temperature`` (default ``1.0``) follows Verl semantics: softmax and NLL use
+``logits / temperature`` and backward applies the corresponding
+``1 / temperature`` chain-rule factor.
+
+Backward always processes eight ``M128`` tiles (1024 tokens) per device-side
+wave, reusing one BF16 dZ buffer.  At the representative shape this holds a
+~256 MiB workspace instead of the ~1 GiB full dZ.  Wave zero plain-stores BF16
+dW; later waves use Hopper BF16 TMA reduce-add.  FP32 WGMMA accumulators are
+rounded in the epilogue, with no FP32 HBM accumulator or final cast.
+``LIGER_SCALED_CE_SM90_VALIDATE_TARGETS=0`` disables the (host-synchronising)
+out-of-range target check.
+"""
+
+import math
+import os
+
+import torch
+
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_backward_fused_sm90 import FusedBackwardConfig
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_backward_fused_sm90 import fused_backward_one_kernel
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_mfast_sm90 import ScaledCEForwardMFastConfig
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_mfast_sm90 import scaled_ce_forward_mfast
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_sm90 import scaled_ce_forward
+from liger_kernel.ops.cutedsl.ops.utils import ensure_cuda_context
+from liger_kernel.ops.utils import amp_custom_bwd
+from liger_kernel.ops.utils import amp_custom_fwd
+
+# All five kernels tile the hidden dimension by 64, so one padding of H covers
+# the forward, dZ, dX and dW phases.
+HIDDEN_TILE = 64
+# dZ and dX both tile the token dimension by 128, so a backward wave is counted
+# in M128 tiles.
+BACKWARD_M_TILE = 128
+# 8 M128 tiles == 1024 tokens per wave: a ~256 MiB BF16 dZ buffer at
+# V = 131072 instead of the ~1 GiB a 4096-token batch would need.
+BACKWARD_M_TILES_PER_WAVE = 8
+VALIDATE_TARGETS_ENV = "LIGER_SCALED_CE_SM90_VALIDATE_TARGETS"
+
+
+def _validate_temperature(temperature):
+    if isinstance(temperature, bool) or not isinstance(temperature, (int, float)):
+        raise TypeError("temperature must be a real number")
+    if not math.isfinite(temperature) or temperature <= 0:
+        raise ValueError("temperature must be finite and > 0")
+
+
+def _validate_targets_enabled() -> bool:
+    return os.environ.get(VALIDATE_TARGETS_ENV, "1") not in ("0", "false", "False")
+
+
+def _validate_inputs(_input, weight, target, ignore_index, m_tiles_per_cluster):
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability(_input.device) != (9, 0):
+        raise RuntimeError("CuTe DSL fused scaled cross entropy requires a Hopper compute capability 9.0 GPU")
+    if _input.device != weight.device or _input.device != target.device:
+        raise ValueError("input, weight, and target must be on the same CUDA device")
+    if _input.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16:
+        raise TypeError("CuTe DSL fused scaled cross entropy supports BF16 input and weight only")
+    if target.dtype != torch.int64:
+        raise TypeError("target must be an int64 tensor")
+    if _input.ndim != 2 or weight.ndim != 2 or target.ndim != 1:
+        raise ValueError("expected input[M,H], weight[V,H], and target[M]")
+    if _input.shape[0] != target.shape[0] or _input.shape[1] != weight.shape[1]:
+        raise ValueError(
+            f"input {tuple(_input.shape)}, weight {tuple(weight.shape)} and target "
+            f"{tuple(target.shape)} shapes are incompatible"
+        )
+    if not isinstance(m_tiles_per_cluster, int) or isinstance(m_tiles_per_cluster, bool):
+        raise TypeError("m_tiles_per_cluster must be an int")
+    if m_tiles_per_cluster < 1:
+        raise ValueError("m_tiles_per_cluster must be >= 1")
+    if not _validate_targets_enabled():
+        return
+    vocab_size = weight.shape[0]
+    out_of_range = ((target < 0) | (target >= vocab_size)) & (target != ignore_index)
+    if bool(out_of_range.any()):
+        raise ValueError(f"target contains values outside [0, {vocab_size}) that are not ignore_index={ignore_index}")
+
+
+def _validate_grad_output(grad_output, tokens):
+    if not isinstance(grad_output, torch.Tensor):
+        raise TypeError("grad_output must be a tensor of per-token gradients")
+    if grad_output.ndim != 1 or grad_output.shape[0] != tokens:
+        raise ValueError(
+            "CuTe DSL fused scaled cross entropy backward expects the per-token upstream gradient of shape "
+            f"[{tokens}], got {tuple(grad_output.shape)}. Apply any reduction outside the kernel so that "
+            "autograd supplies a per-token gradient."
+        )
+
+
+def _pad_hidden(x, weight):
+    """Pad ``H`` up to the shared ``HIDDEN_TILE`` so every phase tiles exactly."""
+    hidden_size = x.shape[1]
+    padded = (hidden_size + HIDDEN_TILE - 1) // HIDDEN_TILE * HIDDEN_TILE
+    if padded == hidden_size:
+        return x.contiguous(), weight.contiguous(), hidden_size
+    pad = (0, padded - hidden_size)
+    return (
+        torch.nn.functional.pad(x.contiguous(), pad),
+        torch.nn.functional.pad(weight.contiguous(), pad),
+        hidden_size,
+    )
+
+
+def fused_scaled_cross_entropy_forward(
+    _input,
+    weight,
+    target,
+    temperature=1.0,
+    ignore_index=-100,
+    m_tiles_per_cluster=1,
+    return_entropy=False,
+):
+    """Per-token fused scaled cross entropy forward.
+
+    Returns ``(nll, entropy, lse, x_padded, weight_padded, hidden_size)``.
+    ``nll`` is the ``[M]`` FP32 negative log-likelihood, ``entropy`` is BF16
+    when requested (otherwise ``None``), and ``lse`` is the FP32 log-sum-exp
+    reused by backward.
+    """
+    _validate_inputs(_input, weight, target, ignore_index, m_tiles_per_cluster)
+    _validate_temperature(temperature)
+    if not isinstance(return_entropy, bool):
+        raise TypeError("return_entropy must be a bool")
+    ensure_cuda_context()
+    x_padded, weight_padded, hidden_size = _pad_hidden(_input, weight)
+    target = target.contiguous()
+
+    if m_tiles_per_cluster == 1:
+        nll, entropy, lse = scaled_ce_forward(
+            x_padded,
+            weight_padded,
+            target,
+            temperature,
+            ignore_index,
+            return_entropy,
+        )
+    else:
+        config = ScaledCEForwardMFastConfig(slots_per_wave=m_tiles_per_cluster)
+        nll, entropy, lse = scaled_ce_forward_mfast(
+            x_padded,
+            weight_padded,
+            target,
+            temperature,
+            ignore_index,
+            return_entropy,
+            config=config,
+        )
+    return nll, entropy, lse, x_padded, weight_padded, hidden_size
+
+
+def fused_scaled_cross_entropy_backward(
+    grad_output,
+    x_padded,
+    weight_padded,
+    target,
+    lse,
+    temperature,
+    ignore_index,
+    hidden_size,
+    entropy=None,
+    grad_entropy=None,
+    input_dtype=torch.bfloat16,
+):
+    """Backward of :func:`fused_scaled_cross_entropy_forward`.
+
+    ``grad_output`` is the ``[M]`` upstream gradient of the per-token NLL and is
+    divided by ``temperature`` and used as the row scale of ``dZ``.  Returns
+    ``(grad_input[M, H], grad_weight[V, H])`` in the caller's dtypes.
+
+    Backward streams tokens through a reusable BF16 dZ buffer in fixed
+    1024-token waves.
+    """
+    tokens = x_padded.shape[0]
+    _validate_grad_output(grad_output, tokens)
+    _validate_temperature(temperature)
+    # Backward usually runs on an autograd worker thread, which has no driver
+    # context bound; the CuTe DSL occupancy helpers need one.
+    ensure_cuda_context()
+
+    scale = (grad_output.detach().to(torch.float32) / temperature).contiguous()
+    entropy_scale = None
+    if grad_entropy is not None:
+        entropy_scale = (grad_entropy.detach().to(torch.float32) / temperature).contiguous()
+    total_m_tiles = (tokens + BACKWARD_M_TILE - 1) // BACKWARD_M_TILE
+    m_tiles_per_wave = min(BACKWARD_M_TILES_PER_WAVE, total_m_tiles)
+    return fused_backward_one_kernel(
+        scale,
+        x_padded,
+        weight_padded,
+        target,
+        lse,
+        entropy,
+        entropy_scale,
+        ignore_index,
+        hidden_size,
+        input_dtype,
+        temperature,
+        config=FusedBackwardConfig(m_tiles_per_wave=m_tiles_per_wave),
+    )
+
+
+class LigerFusedScaledCrossEntropySM90Function(torch.autograd.Function):
+    """``apply(_input, weight, target, temperature=1.0, ignore_index=-100,
+    m_tiles_per_cluster=1, return_entropy=False)``.
+
+    Forward returns the per-token NLL ``[M]`` (FP32).  Backward expects the
+    per-token upstream gradient ``[M]``.  ``m_tiles_per_cluster`` schedules the
+    forward; backward always uses fixed 1024-token waves.
+    """
+
+    @staticmethod
+    @amp_custom_fwd
+    def forward(
+        ctx,
+        _input,
+        weight,
+        target,
+        temperature=1.0,
+        ignore_index=-100,
+        m_tiles_per_cluster=1,
+        return_entropy=False,
+    ):
+        _validate_temperature(temperature)
+        nll, entropy, lse, x_padded, weight_padded, hidden_size = fused_scaled_cross_entropy_forward(
+            _input,
+            weight,
+            target,
+            temperature,
+            ignore_index,
+            m_tiles_per_cluster,
+            return_entropy,
+        )
+        saved_entropy = entropy if entropy is not None else torch.empty(0, device=_input.device, dtype=torch.bfloat16)
+        ctx.save_for_backward(x_padded, weight_padded, target, lse, saved_entropy)
+        ctx.ignore_index = ignore_index
+        ctx.temperature = temperature
+        ctx.hidden_size = hidden_size
+        ctx.input_dtype = _input.dtype
+        ctx.return_entropy = return_entropy
+        return (nll, entropy) if return_entropy else nll
+
+    @staticmethod
+    @amp_custom_bwd
+    def backward(ctx, grad_output, grad_entropy=None):
+        x_padded, weight_padded, target, lse, entropy = ctx.saved_tensors
+        grad_input, grad_weight = fused_scaled_cross_entropy_backward(
+            grad_output,
+            x_padded,
+            weight_padded,
+            target,
+            lse,
+            ctx.temperature,
+            ctx.ignore_index,
+            ctx.hidden_size,
+            entropy=entropy if ctx.return_entropy else None,
+            grad_entropy=grad_entropy,
+            input_dtype=ctx.input_dtype,
+        )
+        return grad_input, grad_weight, None, None, None, None, None
+
+
+__all__ = [
+    "LigerFusedScaledCrossEntropySM90Function",
+    "fused_scaled_cross_entropy_backward",
+    "fused_scaled_cross_entropy_forward",
+]

src/liger_kernel/ops/cutedsl/ops/utils.py

@@ -2,9 +2,37 @@
 Shared helpers for the CuTe DSL backend ops.
 """
 
+import cuda.bindings.driver as driver
+import torch
+
 from cutlass.cute.runtime import from_dlpack
 
 
+def ensure_cuda_context():
+    """Make the CUDA primary context current on the calling thread.
+
+    ``torch.autograd`` runs backward nodes on worker threads.  PyTorch binds
+    those threads through the CUDA *runtime* API, which leaves the *driver*
+    API's current context unset, so CuTe DSL helpers that call the driver
+    directly (notably ``cutlass.utils.HardwareInfo``) fail with
+    ``CUDA_ERROR_INVALID_CONTEXT``.  Retaining and setting the device's primary
+    context -- the very context PyTorch itself uses -- fixes that and is a no-op
+    once a context is already current.
+    """
+    err, ctx = driver.cuCtxGetCurrent()
+    if err == driver.CUresult.CUDA_SUCCESS and ctx and int(ctx) != 0:
+        return
+    err, device = driver.cuDeviceGet(torch.cuda.current_device())
+    if err != driver.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"cuDeviceGet failed: {err}")
+    err, primary = driver.cuDevicePrimaryCtxRetain(device)
+    if err != driver.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"cuDevicePrimaryCtxRetain failed: {err}")
+    (err,) = driver.cuCtxSetCurrent(primary)
+    if err != driver.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"cuCtxSetCurrent failed: {err}")
+
+
 def _next_power_of_2(n: int) -> int:
     """Return the smallest power of 2 greater than or equal to n."""
     n -= 1

test/cutedsl/test_fused_scaled_cross_entropy_sm90.py

@@ -0,0 +1,482 @@
+"""Correctness tests for the Hopper fused scaled cross entropy CuTe DSL kernel."""
+
+import os
+import subprocess
+import sys
+import textwrap
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from test.utils import assert_verbose_allclose
+from test.utils import set_seed
+
+# Every test in this module launches SM90 CuTe DSL kernels, so gate the whole
+# file at collection time: on a CPU-only, non-NVIDIA or non-Hopper host nothing
+# here is collectable.
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0),
+    reason="CuTe DSL fused scaled cross entropy requires a Hopper (SM90) GPU",
+)
+
+
+def _sm90_module():
+    try:
+        import liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 as module
+    except ImportError as exc:
+        pytest.skip(f"CuTe DSL backend is not installed: {exc}")
+    return module
+
+
+def _reference_nll(x, weight, target, ignore_index, temperature=1.0):
+    logits = (x.float() @ weight.float().t()) / temperature
+    valid = target != ignore_index
+    safe_target = target.masked_fill(~valid, 0)
+    lse = torch.logsumexp(logits, dim=-1)
+    target_logits = logits.gather(1, safe_target[:, None]).squeeze(1)
+    return torch.where(valid, lse - target_logits, torch.zeros_like(lse))
+
+
+def _reference_nll_entropy(x, weight, target, ignore_index, temperature=1.0):
+    logits = (x.float() @ weight.float().t()) / temperature
+    log_probs = torch.log_softmax(logits, dim=-1)
+    probs = log_probs.exp()
+    entropy = -(probs * log_probs).sum(dim=-1)
+    valid = target != ignore_index
+    safe_target = target.masked_fill(~valid, 0)
+    nll = -log_probs.gather(1, safe_target[:, None]).squeeze(1)
+    zeros = torch.zeros_like(nll)
+    return torch.where(valid, nll, zeros), torch.where(valid, entropy, zeros)
+
+
+def _verl_reference(x, weight, target, grad_nll, grad_entropy, temperature):
+    """Dependency-free equivalent of Verl's experimental fused PPO function."""
+    logits = ((x @ weight.t()) / temperature).float()
+    probs = logits.softmax(dim=-1)
+    log_probs = logits.log_softmax(dim=-1)
+    token_log_probs = log_probs.gather(-1, target[:, None]).squeeze(-1)
+    entropy = (torch.logsumexp(logits, dim=-1) - torch.sum(probs * logits, dim=-1)).to(x.dtype)
+
+    one_hot = torch.zeros_like(logits).scatter_(-1, target[:, None], 1)
+    dlogits = (-grad_nll).float()[:, None] * (one_hot - probs)
+    dlogits += probs * (log_probs + entropy.float()[:, None]) * (-grad_entropy.float()[:, None])
+    dlogits = dlogits.to(x.dtype) / temperature
+    return -token_log_probs, entropy, dlogits @ weight, dlogits.t() @ x
+
+
+def _make_inputs(bt, hidden_size, vocab_size, ignore_stride=None):
+    x = torch.randn(bt, hidden_size, device="cuda", dtype=torch.bfloat16) * 0.25
+    weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
+    target = torch.randint(vocab_size, (bt,), device="cuda", dtype=torch.long)
+    if ignore_stride is not None:
+        target[::ignore_stride] = -100
+    return x, weight, target
+
+
+def _run_liger(
+    x,
+    weight,
+    target,
+    grad_output,
+    ignore_index=-100,
+    temperature=1.0,
+    m_tiles_per_cluster=1,
+):
+    module = _sm90_module()
+    x = x.detach().clone().requires_grad_(True)
+    weight = weight.detach().clone().requires_grad_(True)
+    nll = module.LigerFusedScaledCrossEntropySM90Function.apply(
+        x,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+        m_tiles_per_cluster,
+        False,
+    )
+    if grad_output is not None:
+        nll.backward(grad_output)
+        return nll.detach(), x.grad.detach(), weight.grad.detach()
+    return nll.detach(), None, None
+
+
+def _run_reference(x, weight, target, grad_output, ignore_index=-100, temperature=1.0):
+    x = x.detach().clone().requires_grad_(True)
+    weight = weight.detach().clone().requires_grad_(True)
+    nll = _reference_nll(x, weight, target, ignore_index, temperature)
+    if grad_output is not None:
+        nll.backward(grad_output)
+        return nll.detach(), x.grad.detach(), weight.grad.detach()
+    return nll.detach(), None, None
+
+
+@pytest.mark.parametrize("shape", [(256, 128, 512), (193, 70, 517)], ids=["aligned", "ragged"])
+@pytest.mark.parametrize("m_tiles_per_cluster", [1, 2], ids=["mtiles1", "mtiles2"])
+def test_sm90_per_token_nll_and_grads_match_torch(shape, m_tiles_per_cluster):
+    set_seed(42)
+    bt, hidden_size, vocab_size = shape
+    x, weight, target = _make_inputs(bt, hidden_size, vocab_size, ignore_stride=7)
+    grad_output = torch.rand(bt, device="cuda", dtype=torch.float32) + 0.25
+
+    actual = _run_liger(
+        x,
+        weight,
+        target,
+        grad_output,
+        m_tiles_per_cluster=m_tiles_per_cluster,
+    )
+    expected = _run_reference(x, weight, target, grad_output)
+
+    assert actual[0].shape == (bt,)
+    assert actual[0].dtype == torch.float32
+    assert actual[1].shape == x.shape and actual[1].dtype == x.dtype
+    assert actual[2].shape == weight.shape and actual[2].dtype == weight.dtype
+    assert_verbose_allclose(actual[0].float(), expected[0].float(), atol=5e-2, rtol=5e-2)
+    assert_verbose_allclose(actual[1].float(), expected[1].float(), atol=2e-2, rtol=1e-1)
+    assert_verbose_allclose(actual[2].float(), expected[2].float(), atol=2e-2, rtol=1e-1)
+
+
+@pytest.mark.parametrize("temperature", [0.7, 1.5])
+def test_sm90_temperature_matches_torch(temperature):
+    set_seed(43)
+    x, weight, target = _make_inputs(193, 70, 517, ignore_stride=7)
+    grad_output = torch.rand(193, device="cuda", dtype=torch.float32) + 0.25
+
+    actual = _run_liger(x, weight, target, grad_output, temperature=temperature)
+    expected = _run_reference(x, weight, target, grad_output, temperature=temperature)
+
+    assert_verbose_allclose(actual[0].float(), expected[0].float(), atol=5e-2, rtol=5e-2)
+    assert_verbose_allclose(actual[1].float(), expected[1].float(), atol=2e-2, rtol=1e-1)
+    assert_verbose_allclose(actual[2].float(), expected[2].float(), atol=2e-2, rtol=1e-1)
+
+
+@pytest.mark.parametrize("m_tiles_per_cluster", [1, 2])
+def test_sm90_differentiable_entropy_matches_torch(m_tiles_per_cluster):
+    set_seed(44)
+    x, weight, target = _make_inputs(193, 70, 517, ignore_stride=7)
+    temperature = 0.8
+    grad_nll = torch.rand(193, device="cuda", dtype=torch.float32) + 0.25
+    grad_entropy = torch.rand(193, device="cuda", dtype=torch.float32) - 0.5
+
+    module = _sm90_module()
+    xa = x.detach().clone().requires_grad_(True)
+    wa = weight.detach().clone().requires_grad_(True)
+    nll, entropy = module.LigerFusedScaledCrossEntropySM90Function.apply(
+        xa,
+        wa,
+        target,
+        temperature,
+        -100,
+        m_tiles_per_cluster,
+        True,
+    )
+    torch.autograd.backward((nll, entropy), (grad_nll, grad_entropy.to(entropy.dtype)))
+
+    xr = x.detach().clone().requires_grad_(True)
+    wr = weight.detach().clone().requires_grad_(True)
+    ref_nll, ref_entropy = _reference_nll_entropy(xr, wr, target, -100, temperature)
+    torch.autograd.backward((ref_nll, ref_entropy), (grad_nll, grad_entropy))
+
+    assert entropy.dtype == torch.bfloat16
+    assert entropy[target == -100].count_nonzero().item() == 0
+    assert_verbose_allclose(nll.float(), ref_nll.float(), atol=5e-2, rtol=5e-2)
+    assert_verbose_allclose(entropy.float(), ref_entropy.float(), atol=5e-2, rtol=5e-2)
+    assert_verbose_allclose(xa.grad.float(), xr.grad.float(), atol=3e-2, rtol=1.5e-1)
+    assert_verbose_allclose(wa.grad.float(), wr.grad.float(), atol=3e-2, rtol=1.5e-1)
+
+
+def test_sm90_matches_verl_fused_ppo_formulas():
+    set_seed(45)
+    x, weight, target = _make_inputs(193, 70, 517)
+    temperature = 0.8
+    grad_nll = torch.rand(193, device="cuda", dtype=torch.float32) + 0.25
+    grad_entropy = torch.rand(193, device="cuda", dtype=torch.float32) - 0.5
+
+    module = _sm90_module()
+    xa = x.detach().clone().requires_grad_(True)
+    wa = weight.detach().clone().requires_grad_(True)
+    nll, entropy = module.LigerFusedScaledCrossEntropySM90Function.apply(
+        xa,
+        wa,
+        target,
+        temperature,
+        -100,
+        1,
+        True,
+    )
+    torch.autograd.backward((nll, entropy), (grad_nll, grad_entropy.to(entropy.dtype)))
+
+    ref_nll, ref_entropy, ref_dx, ref_dw = _verl_reference(
+        x,
+        weight,
+        target,
+        grad_nll,
+        grad_entropy,
+        temperature,
+    )
+    assert_verbose_allclose(nll.float(), ref_nll.float(), atol=5e-2, rtol=5e-2)
+    assert_verbose_allclose(entropy.float(), ref_entropy.float(), atol=5e-2, rtol=5e-2)
+    assert_verbose_allclose(xa.grad.float(), ref_dx.float(), atol=3e-2, rtol=1.5e-1)
+    assert_verbose_allclose(wa.grad.float(), ref_dw.float(), atol=3e-2, rtol=1.5e-1)
+
+
+def test_sm90_ignored_entropy_rows_stay_zero_with_large_logits():
+    set_seed(46)
+    x = (torch.randn(128, 64, device="cuda", dtype=torch.bfloat16) * 20).requires_grad_(True)
+    weight = (torch.randn(256, 64, device="cuda", dtype=torch.bfloat16) * 20).requires_grad_(True)
+    target = torch.full((128,), -100, device="cuda", dtype=torch.long)
+
+    module = _sm90_module()
+    nll, entropy = module.LigerFusedScaledCrossEntropySM90Function.apply(
+        x,
+        weight,
+        target,
+        0.1,
+        -100,
+        1,
+        True,
+    )
+    torch.autograd.backward((nll, entropy), (torch.ones_like(nll), torch.ones_like(entropy)))
+
+    assert nll.count_nonzero().item() == 0
+    assert entropy.count_nonzero().item() == 0
+    assert torch.isfinite(x.grad).all() and x.grad.count_nonzero().item() == 0
+    assert torch.isfinite(weight.grad).all() and weight.grad.count_nonzero().item() == 0
+
+
+def test_sm90_ignore_index_rows_are_zero():
+    set_seed(0)
+    bt, hidden_size, vocab_size = 192, 96, 257
+    x, weight, target = _make_inputs(bt, hidden_size, vocab_size, ignore_stride=5)
+    grad_output = torch.ones(bt, device="cuda", dtype=torch.float32)
+
+    nll, dx, _ = _run_liger(x, weight, target, grad_output)
+    ignored = target == -100
+
+    assert nll[ignored].count_nonzero().item() == 0
+    assert nll[~ignored].min().item() > 0.0
+    assert dx[ignored].count_nonzero().item() == 0
+
+
+def test_sm90_zero_grad_rows_produce_zero_dx():
+    set_seed(1)
+    bt, hidden_size, vocab_size = 192, 96, 257
+    x, weight, target = _make_inputs(bt, hidden_size, vocab_size, ignore_stride=11)
+    grad_output = torch.zeros(bt, device="cuda", dtype=torch.float32)
+    grad_output[::5] = torch.rand((bt + 4) // 5, device="cuda", dtype=torch.float32) + 0.25
+
+    actual = _run_liger(x, weight, target, grad_output)
+    expected = _run_reference(x, weight, target, grad_output)
+
+    assert actual[1][grad_output == 0].count_nonzero().item() == 0
+    assert_verbose_allclose(actual[1].float(), expected[1].float(), atol=2e-2, rtol=1e-1)
+    assert_verbose_allclose(actual[2].float(), expected[2].float(), atol=2e-2, rtol=1e-1)
+
+
+def test_sm90_all_ignored_is_zero():
+    set_seed(2)
+    bt, hidden_size, vocab_size = 128, 64, 256
+    x, weight, _ = _make_inputs(bt, hidden_size, vocab_size)
+    target = torch.full((bt,), -100, device="cuda", dtype=torch.long)
+    grad_output = torch.ones(bt, device="cuda", dtype=torch.float32)
+
+    nll, dx, dw = _run_liger(x, weight, target, grad_output)
+
+    assert nll.count_nonzero().item() == 0
+    assert dx.count_nonzero().item() == 0
+    assert dw.count_nonzero().item() == 0
+
+
+def test_sm90_many_wave_dw_matches_torch_within_bf16_accumulation_error():
+    """Multi-wave BF16 dW accumulation must track the FP32 torch reference."""
+    set_seed(17)
+    bt, hidden_size, vocab_size = 2048, 256, 512
+    x, weight, target = _make_inputs(bt, hidden_size, vocab_size, ignore_stride=11)
+    grad_output = torch.rand(bt, device="cuda", dtype=torch.float32) + 0.25
+
+    waved = _run_liger(x, weight, target, grad_output)  # two 1024-token waves
+    expected = _run_reference(x, weight, target, grad_output)[2].float()
+
+    # Absolute tolerance is relaxed because every
+    # wave's contribution is rounded to BF16 (~2^-8 relative) before the add.
+    assert_verbose_allclose(waved[2].float(), expected, atol=5e-2, rtol=2e-1)
+
+    scale = expected.abs().mean()
+    wave_err = (waved[2].float() - expected).abs().mean() / scale
+    assert wave_err < 2e-2, wave_err
+
+
+def test_sm90_backward_uses_one_fused_kernel(monkeypatch):
+    """Backward dispatches one fused kernel with the fixed wave cap."""
+    module = _sm90_module()
+    set_seed(12)
+    bt, hidden_size, vocab_size = 512, 128, 256
+    x, weight, target = _make_inputs(bt, hidden_size, vocab_size)
+    x = x.requires_grad_(True)
+    weight = weight.requires_grad_(True)
+    grad_output = torch.rand(bt, device="cuda", dtype=torch.float32) + 0.25
+
+    calls = []
+    original = module.fused_backward_one_kernel
+
+    def wrapper(*args, **kwargs):
+        calls.append(kwargs["config"].m_tiles_per_wave)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(module, "fused_backward_one_kernel", wrapper)
+
+    module.LigerFusedScaledCrossEntropySM90Function.apply(x, weight, target).backward(grad_output)
+    assert calls == [4]
+
+
+@pytest.mark.parametrize("m_tiles_per_cluster", [1, 3], ids=["mtiles1", "mtiles3"])
+def test_sm90_m_tiles_dispatch_selects_expected_forward(monkeypatch, m_tiles_per_cluster):
+    module = _sm90_module()
+    calls = {}
+
+    def fake_m_outer(x, weight, target, temperature, ignore_index, return_entropy):
+        calls["m_outer"] = (temperature, ignore_index)
+        return torch.zeros(x.shape[0], device=x.device), None, torch.zeros(x.shape[0], device=x.device)
+
+    def fake_m_fast(x, weight, target, temperature, ignore_index, return_entropy, config=None):
+        calls["m_fast"] = (temperature, config.slots_per_wave)
+        return torch.zeros(x.shape[0], device=x.device), None, torch.zeros(x.shape[0], device=x.device)
+
+    monkeypatch.setattr(module, "scaled_ce_forward", fake_m_outer)
+    monkeypatch.setattr(module, "scaled_ce_forward_mfast", fake_m_fast)
+
+    x = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16)
+    target = torch.zeros(4, device="cuda", dtype=torch.long)
+    module.fused_scaled_cross_entropy_forward(x, weight, target, 1.0, -100, m_tiles_per_cluster)
+
+    if m_tiles_per_cluster == 1:
+        assert calls == {"m_outer": (1.0, -100)}
+    else:
+        assert calls == {"m_fast": (1.0, m_tiles_per_cluster)}
+
+
+def test_sm90_rejects_scalar_and_mismatched_grad_output():
+    set_seed(3)
+    module = _sm90_module()
+    bt, hidden_size, vocab_size = 64, 64, 128
+    x, weight, target = _make_inputs(bt, hidden_size, vocab_size)
+
+    _, _, lse, x_padded, weight_padded, hidden = module.fused_scaled_cross_entropy_forward(
+        x, weight, target, 1.0, -100, 1
+    )
+    for bad_grad in (
+        torch.tensor(1.0, device="cuda"),
+        torch.ones(bt - 1, device="cuda"),
+        torch.ones(bt, 1, device="cuda"),
+    ):
+        with pytest.raises(ValueError):
+            module.fused_scaled_cross_entropy_backward(
+                bad_grad, x_padded, weight_padded, target, lse, 1.0, -100, hidden
+            )
+    with pytest.raises(TypeError):
+        module.fused_scaled_cross_entropy_backward(1.0, x_padded, weight_padded, target, lse, 1.0, -100, hidden)
+
+
+def test_sm90_sum_reduction_composed_in_autograd():
+    """``sum().backward()`` is the supported reduction path: autograd expands
+    the scalar into the per-token vector of ones before the kernel sees it."""
+    set_seed(4)
+    bt, hidden_size, vocab_size = 128, 64, 256
+    x, weight, target = _make_inputs(bt, hidden_size, vocab_size, ignore_stride=6)
+
+    module = _sm90_module()
+    xa = x.detach().clone().requires_grad_(True)
+    wa = weight.detach().clone().requires_grad_(True)
+    module.LigerFusedScaledCrossEntropySM90Function.apply(xa, wa, target).sum().backward()
+
+    xr = x.detach().clone().requires_grad_(True)
+    wr = weight.detach().clone().requires_grad_(True)
+    _reference_nll(xr, wr, target, -100).sum().backward()
+
+    assert_verbose_allclose(xa.grad.float(), xr.grad.float(), atol=2e-2, rtol=1e-1)
+    assert_verbose_allclose(wa.grad.float(), wr.grad.float(), atol=2e-2, rtol=1e-1)
+
+
+def test_sm90_rejects_bad_dtypes_shapes_and_options():
+    module = _sm90_module()
+    fn = module.LigerFusedScaledCrossEntropySM90Function
+    x = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16)
+    target = torch.zeros(4, device="cuda", dtype=torch.long)
+
+    with pytest.raises(TypeError):
+        fn.apply(x.to(torch.float16), weight.to(torch.float16), target)
+    with pytest.raises(TypeError):
+        fn.apply(x, weight, target.to(torch.int32))
+    with pytest.raises(ValueError):
+        fn.apply(x, weight, target[:2])
+    with pytest.raises(ValueError):
+        fn.apply(x, torch.randn(16, 8, device="cuda", dtype=torch.bfloat16), target)
+    with pytest.raises(ValueError):
+        fn.apply(x, weight, target, 0.0)
+    with pytest.raises(TypeError):
+        fn.apply(x, weight, target, "cold")
+    with pytest.raises(ValueError):
+        fn.apply(x, weight, target, float("inf"))
+    with pytest.raises(ValueError):
+        fn.apply(x, weight, target, 1.0, -100, 0)
+    with pytest.raises(TypeError):
+        fn.apply(x, weight, target, 1.0, -100, 1.5)
+    with pytest.raises(TypeError):
+        fn.apply(x, weight, target, 1.0, -100, 1, 2.0)
+
+
+def test_sm90_rejects_out_of_range_targets():
+    module = _sm90_module()
+    fn = module.LigerFusedScaledCrossEntropySM90Function
+    x = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16)
+    target = torch.zeros(4, device="cuda", dtype=torch.long)
+    target[1] = 16
+
+    with pytest.raises(ValueError):
+        fn.apply(x, weight, target)
+
+    target[1] = -100
+    fn.apply(x, weight, target)
+
+
+def test_sm90_exports_are_additive_and_do_not_alias_triton_flce():
+    _sm90_module()
+    repo_root = Path(__file__).resolve().parents[2]
+    pythonpath = os.pathsep.join([str(repo_root / "src"), str(repo_root), os.environ.get("PYTHONPATH", "")])
+    env = {
+        **os.environ,
+        "CUTE_DSL_ARCH": "sm_90a",
+        "LIGER_KERNEL_IMPL": "cutedsl",
+        "PYTHONPATH": pythonpath,
+    }
+    script = textwrap.dedent(
+        """
+        import liger_kernel.ops.cutedsl.ops as cutedsl_ops
+        from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction
+
+        expected = "liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90"
+        actual = cutedsl_ops.LigerFusedScaledCrossEntropySM90Function.__module__
+        if actual != expected:
+            raise AssertionError(f"Expected {expected}, got {actual}")
+        for name in (
+            "LigerFusedScaledCrossEntropySM90Function",
+            "fused_scaled_cross_entropy_forward",
+            "fused_scaled_cross_entropy_backward",
+        ):
+            if name not in cutedsl_ops.__all__:
+                raise AssertionError(f"{name} missing from cutedsl ops __all__")
+        if hasattr(cutedsl_ops, "LigerFusedLinearCrossEntropyFunction"):
+            raise AssertionError("cutedsl must not export a FusedLinearCrossEntropy override")
+        if LigerFusedLinearCrossEntropyFunction.__module__ != "liger_kernel.ops.fused_linear_cross_entropy":
+            raise AssertionError(
+                "Triton FusedLinearCrossEntropy must stay in place, got "
+                f"{LigerFusedLinearCrossEntropyFunction.__module__}"
+            )
+        """
+    )
+    subprocess.run([sys.executable, "-c", script], check=True, env=env, cwd=repo_root)