perf(cutedsl): optimize Hopper FLSCE forward (#1370)

## Summary - replaces the SM90 M-outer and M-fast FLSCE forwards with one fixed cluster-M2 fragment kernel - computes a logical `M128 x N320 x K64` tile as two N160 accumulators and four N80 epilogue slices - adds a measured runtime split-N lookup for the profiled 16K/32K long-sequence shapes - supports differentiable BF16 entropy through online weighted statistics and a small split finalizer - preserves `m_tiles_per_cluster` for API compatibility while removing its slower schedule variants - keeps the existing one-kernel N256 backward unchanged - extracts shared SM90 forward helpers and removes 2,320 lines of obsolete forward code ## Kernel design - WG0: 24-register TMA producer - WG1/WG2: 240-register WGMMA plus online-softmax epilogue - cluster-M2 pairs process adjacent token tiles and multicast shared weight panels - safe one-stage AB release lag prevents TMA overwriting live WGMMA operands - two padded FP16 SMEM buffers stage four N80 slices; logits never round-trip through HBM - cluster pairs emit compact `(max, sum, target[, weighted])` statistics; a small kernel finalizes NLL/LSE/entropy - unprofiled shapes derive split count from active cluster capacity, including uneven M-pair assignments The H100 lookup for `H=4096`, `V=131072` selects: | M | Entropy | `split_n` | |---:|:---:|---:| | 16384 | No | 4 | | 16384 | Yes | 2 | | 32768 | No | 2 | | 32768 | Yes | 2 | ## H100 performance BF16, `H=4096`, `V=131072`, cuTile 1.5.0. Effective TFLOPS count the common projection work: forward `2*M*H*V`, backward `6*M*H*V`, and full `8*M*H*V`. ### Forward Medians from 60 interleaved samples per provider. | M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback | |---:|:---:|---:|---:|---:| | 2048 | No | **3.12 ms / 706 TFLOPS** | 3.19 ms / 690 TFLOPS | 11.36 ms / 194 TFLOPS | | 2048 | Yes | **3.13 ms / 703 TFLOPS** | 3.25 ms / 676 TFLOPS | 11.37 ms / 193 TFLOPS | | 4096 | No | **6.05 ms / 727 TFLOPS** | 6.28 ms / 701 TFLOPS | 22.62 ms / 194 TFLOPS | | 4096 | Yes | **6.11 ms / 720 TFLOPS** | 6.39 ms / 688 TFLOPS | 22.68 ms / 194 TFLOPS | | 8192 | No | **12.25 ms / 718 TFLOPS** | 12.41 ms / 709 TFLOPS | 45.35 ms / 194 TFLOPS | | 8192 | Yes | **12.30 ms / 715 TFLOPS** | 12.70 ms / 693 TFLOPS | 45.59 ms / 193 TFLOPS | | 16384 | No | **23.84 ms / 738 TFLOPS** | 25.11 ms / 700 TFLOPS | 90.81 ms / 194 TFLOPS | | 16384 | Yes | **24.09 ms / 730 TFLOPS** | 26.10 ms / 674 TFLOPS | 90.97 ms / 193 TFLOPS | | 32768 | No | **49.69 ms / 708 TFLOPS** | 54.31 ms / 648 TFLOPS | 180.01 ms / 195 TFLOPS | | 32768 | Yes | **49.85 ms / 706 TFLOPS** | 55.69 ms / 632 TFLOPS | 179.85 ms / 196 TFLOPS | The tuned 16K/32K CuTe forward is 5.4-11.7% faster than cuTile. ### Backward Medians from 30 interleaved samples per provider. | M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback | |---:|:---:|---:|---:|---:| | 8192 | No | **37.93 ms / 696 TFLOPS** | 42.55 ms / 620 TFLOPS | 85.17 ms / 310 TFLOPS | | 8192 | Yes | **37.89 ms / 696 TFLOPS** | 41.43 ms / 637 TFLOPS | 121.00 ms / 218 TFLOPS | | 16384 | No | **79.42 ms / 665 TFLOPS** | 83.78 ms / 630 TFLOPS | 166.85 ms / 316 TFLOPS | | 16384 | Yes | **78.20 ms / 675 TFLOPS** | 84.12 ms / 627 TFLOPS | 239.39 ms / 220 TFLOPS | | 32768 | No | 163.34 ms / 646 TFLOPS | **163.21 ms / 647 TFLOPS** | 331.08 ms / 319 TFLOPS | | 32768 | Yes | **160.64 ms / 657 TFLOPS** | 161.73 ms / 653 TFLOPS | 477.51 ms / 221 TFLOPS | ### Full forward + backward | M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback | |---:|:---:|---:|---:|---:| | 8192 | No | **50.39 ms / 698 TFLOPS** | 57.15 ms / 616 TFLOPS | 128.76 ms / 273 TFLOPS | | 8192 | Yes | **50.22 ms / 701 TFLOPS** | 56.86 ms / 619 TFLOPS | 165.63 ms / 212 TFLOPS | | 16384 | No | **104.77 ms / 672 TFLOPS** | 112.03 ms / 628 TFLOPS | 254.91 ms / 276 TFLOPS | | 16384 | Yes | **103.65 ms / 679 TFLOPS** | 111.21 ms / 633 TFLOPS | 328.37 ms / 214 TFLOPS | | 32768 | No | **211.30 ms / 666 TFLOPS** | 221.34 ms / 636 TFLOPS | 508.33 ms / 277 TFLOPS | | 32768 | Yes | **210.60 ms / 668 TFLOPS** | 221.39 ms / 636 TFLOPS | 655.89 ms / 215 TFLOPS | Against the removed generic CuTe forward, the fragment path measured: | M | Generic | Fragment | Speedup | |---:|---:|---:|---:| | 2048 | 4.56 ms | 3.28 ms | **1.39x** | | 4096 | 7.58 ms | 6.41 ms | **1.18x** | ## Validation - H100: `47 passed, 19 skipped` ```bash pytest -q \ test/cutedsl/test_fused_scaled_cross_entropy_forward_fragment_sm90.py \ test/cutedsl/test_fused_scaled_cross_entropy_sm90.py \ test/transformers/test_fused_linear_scaled_cross_entropy.py ``` - independent Torch oracle coverage for NLL, entropy, temperature, ignore index, ragged M/H/V, K64/K4096, final-column targets, hardware/uneven split assignments, entropy-only backward, and runtime lookup selection - `make checkstyle` passes ## Compatibility The public input/output contract and original N256/wave-8 backward are unchanged. Entropy and unprofiled shapes use the same fragment forward rather than retaining slower fallback schedules. --------- Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]>
2026-10-12 02:17:47 +00:00 · 2026-08-11 18:14:17 +00:00
parent ee1ea1157e
commit 62b691b721
8 changed files with 1508 additions and 2380 deletions
@@ -235,15 +235,48 @@ loss = nll.sum() / (target != -100).sum().clamp_min(1)
 The implementations share this public contract but use different schedules:

 - **cuTile** (`LIGER_KERNEL_IMPL=cutile`) supports matching floating-point `input` and `weight` tensors on CUDA and a finite positive scalar `temperature`. The portable temporary-logits budget is 256 MiB. Large FP16/BF16 Blackwell workloads (`M >= 4096`, `V >= 131072`) automatically use 512 MiB to improve GEMM utilization. Set `LIGER_CUTILE_SCALED_CE_WORKSPACE_MB` to another positive MiB value for workload-specific tuning; for example, 1024 can help large combined NLL-plus-entropy workloads but is not universally faster. Backward reuses one workspace and writes `dX` and accumulates `dW` directly into their final tensors. `m_tiles_per_cluster` is accepted for API compatibility but does not change the cuTile schedule.
- **CuTe SM90** uses `LigerFusedScaledCrossEntropySM90Function` for BF16 inputs on Hopper. Forward never writes logits to HBM; `m_tiles_per_cluster=1` selects the fastest M-outer schedule and larger values select the M-fast schedule. Backward runs `dZ`, `dX`, and `dW` in one persistent cluster kernel with a reusable 1024-token `dZ` workspace.
+- **CuTe SM90** uses `LigerFusedScaledCrossEntropySM90Function` for BF16 inputs on Hopper. Its sole forward uses the fixed cluster-M2 N160 fragment kernel, with a measured split-N lookup for profiled long-sequence shapes, and never writes logits to HBM; `m_tiles_per_cluster` remains accepted for API compatibility but does not change that schedule. Backward runs `dZ`, `dX`, and `dW` in one persistent cluster kernel with a reusable 1024-token `dZ` workspace.
 - **Fallback** uses a 512-token chunked PyTorch implementation adapted from Verl's fused PPO formulas when the default frontend cannot use the SM90 kernel.

-H100 measurements for per-token NLL at `M=4096`, `H=4096`, `V=131072`, BF16:
+H100 BF16 forward medians from 60 interleaved samples per provider at
+`H=4096`, `V=131072`. Effective TFLOPS count the common projection work,
+`2*M*H*V`:

-| Implementation | Forward | Backward | Full | Peak full memory |
-|---|---:|---:|---:|---:|
-| cuTile | 6.45 ms | 20.70 ms | 26.94 ms | 2.43 GiB |
-| CuTe SM90 | 7.26 ms | 19.81 ms | 27.74 ms | 2.43 GiB |
+| M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback |
+|---:|:---:|---:|---:|---:|
+| 2048 | No | **3.12 ms / 706 TFLOPS** | 3.19 ms / 690 TFLOPS | 11.36 ms / 194 TFLOPS |
+| 2048 | Yes | **3.13 ms / 703 TFLOPS** | 3.25 ms / 676 TFLOPS | 11.37 ms / 193 TFLOPS |
+| 4096 | No | **6.05 ms / 727 TFLOPS** | 6.28 ms / 701 TFLOPS | 22.62 ms / 194 TFLOPS |
+| 4096 | Yes | **6.11 ms / 720 TFLOPS** | 6.39 ms / 688 TFLOPS | 22.68 ms / 194 TFLOPS |
+| 8192 | No | **12.25 ms / 718 TFLOPS** | 12.41 ms / 709 TFLOPS | 45.35 ms / 194 TFLOPS |
+| 8192 | Yes | **12.30 ms / 715 TFLOPS** | 12.70 ms / 693 TFLOPS | 45.59 ms / 193 TFLOPS |
+| 16384 | No | **23.84 ms / 738 TFLOPS** | 25.11 ms / 700 TFLOPS | 90.81 ms / 194 TFLOPS |
+| 16384 | Yes | **24.09 ms / 730 TFLOPS** | 26.10 ms / 674 TFLOPS | 90.97 ms / 193 TFLOPS |
+| 32768 | No | **49.69 ms / 708 TFLOPS** | 54.31 ms / 648 TFLOPS | 180.01 ms / 195 TFLOPS |
+| 32768 | Yes | **49.85 ms / 706 TFLOPS** | 55.69 ms / 632 TFLOPS | 179.85 ms / 196 TFLOPS |
+
+Backward medians use 30 interleaved samples per provider. Effective TFLOPS
+count `6*M*H*V`:
+
+| M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback |
+|---:|:---:|---:|---:|---:|
+| 8192 | No | **37.93 ms / 696 TFLOPS** | 42.55 ms / 620 TFLOPS | 85.17 ms / 310 TFLOPS |
+| 8192 | Yes | **37.89 ms / 696 TFLOPS** | 41.43 ms / 637 TFLOPS | 121.00 ms / 218 TFLOPS |
+| 16384 | No | **79.42 ms / 665 TFLOPS** | 83.78 ms / 630 TFLOPS | 166.85 ms / 316 TFLOPS |
+| 16384 | Yes | **78.20 ms / 675 TFLOPS** | 84.12 ms / 627 TFLOPS | 239.39 ms / 220 TFLOPS |
+| 32768 | No | 163.34 ms / 646 TFLOPS | **163.21 ms / 647 TFLOPS** | 331.08 ms / 319 TFLOPS |
+| 32768 | Yes | **160.64 ms / 657 TFLOPS** | 161.73 ms / 653 TFLOPS | 477.51 ms / 221 TFLOPS |
+
+Full forward-and-backward effective TFLOPS count `8*M*H*V`:
+
+| M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback |
+|---:|:---:|---:|---:|---:|
+| 8192 | No | **50.39 ms / 698 TFLOPS** | 57.15 ms / 616 TFLOPS | 128.76 ms / 273 TFLOPS |
+| 8192 | Yes | **50.22 ms / 701 TFLOPS** | 56.86 ms / 619 TFLOPS | 165.63 ms / 212 TFLOPS |
+| 16384 | No | **104.77 ms / 672 TFLOPS** | 112.03 ms / 628 TFLOPS | 254.91 ms / 276 TFLOPS |
+| 16384 | Yes | **103.65 ms / 679 TFLOPS** | 111.21 ms / 633 TFLOPS | 328.37 ms / 214 TFLOPS |
+| 32768 | No | **211.30 ms / 666 TFLOPS** | 221.34 ms / 636 TFLOPS | 508.33 ms / 277 TFLOPS |
+| 32768 | Yes | **210.60 ms / 668 TFLOPS** | 221.39 ms / 636 TFLOPS | 655.89 ms / 215 TFLOPS |

 B200 measurements for the same shape, using automatic cuTile workspace selection:

@@ -0,0 +1,79 @@
+"""Shared host and DSL helpers for the SM90 scaled cross entropy forward."""
+
+import inspect
+
+import cuda.bindings.driver as cuda
+import cutlass
+import torch
+
+from cutlass import Float32
+from cutlass._mlir.dialects import nvvm
+from cutlass.cutlass_dsl import T
+from cutlass.cutlass_dsl import dsl_user_op
+
+LOG2_E = 1.4426950408889634
+LN2 = 0.6931471805599453
+NEG_INF_F32 = -1.0e38
+MASK_F32 = -3.0e38
+HOPPER_MAX_SMEM_BYTES = 227 * 1024
+
+_stream_cache = {}
+_max_active_clusters_cache = {}
+
+try:
+    _FMAX_NEEDS_RESULT_TYPE = next(iter(inspect.signature(nvvm.fmax).parameters)) == "res"
+except Exception:
+    _FMAX_NEEDS_RESULT_TYPE = False
+
+
+@dsl_user_op
+def _fmax(a, b, *, loc=None, ip=None) -> Float32:
+    av = Float32(a).ir_value(loc=loc, ip=ip)
+    bv = Float32(b).ir_value(loc=loc, ip=ip)
+    if _FMAX_NEEDS_RESULT_TYPE:
+        return Float32(nvvm.fmax(T.f32(), av, bv, loc=loc, ip=ip))
+    return Float32(nvvm.fmax(av, bv, loc=loc, ip=ip))
+
+
+def _cute_stream():
+    raw = torch.cuda.current_stream().cuda_stream
+    stream = _stream_cache.get(raw)
+    if stream is None:
+        stream = cuda.CUstream(raw)
+        _stream_cache[raw] = stream
+    return stream
+
+
+def _max_active_clusters(cluster_size):
+    key = (torch.cuda.current_device(), cluster_size)
+    value = _max_active_clusters_cache.get(key)
+    if value is None:
+        torch.cuda.current_stream()
+        value = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size)
+        _max_active_clusters_cache[key] = value
+    return value
+
+
+def _validate(x, weight, target):
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability(x.device) != (9, 0):
+        raise RuntimeError("SM90 fused scaled cross entropy forward requires a Hopper (sm90) GPU")
+    if x.device != weight.device or x.device != target.device:
+        raise ValueError("input, weight, and target must be on the same CUDA device")
+    if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16:
+        raise TypeError("SM90 fused scaled cross entropy forward supports BF16 input and weight only")
+    if x.ndim != 2 or weight.ndim != 2 or target.ndim != 1:
+        raise ValueError("expected input[BT,H], weight[V,H], and target[BT]")
+    if x.shape[0] != target.shape[0] or x.shape[1] != weight.shape[1]:
+        raise ValueError("input, weight, and target shapes are incompatible")
+
+
+def _pad_hidden(x, weight, tile_k):
+    hidden_size = x.shape[1]
+    padded = (hidden_size + tile_k - 1) // tile_k * tile_k
+    if padded == hidden_size:
+        return x.contiguous(), weight.contiguous(), hidden_size
+    return (
+        torch.nn.functional.pad(x, (0, padded - hidden_size)),
+        torch.nn.functional.pad(weight, (0, padded - hidden_size)),
+        hidden_size,
+    )
@@ -32,24 +32,22 @@ 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``
+=========================  ==================================================
+forward                    ``_fused_scaled_cross_entropy_forward_fragment_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.
+The fixed forward uses cluster-M2 weight multicast, two N160 accumulators and
+four N80 online-softmax folds per logical N320 vocabulary tile.
+``m_tiles_per_cluster`` remains accepted for API compatibility, but every
+positive value selects this same forward.  Profiled long-sequence shapes use
+a split-N lookup table; all other shapes retain the hardware-derived split.

 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
+shared-memory arena: dZ uses its static persistent 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.
@@ -74,15 +72,14 @@ 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._fused_scaled_cross_entropy_forward_fragment_sm90 import ScaledCEForwardFragmentConfig
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_fragment_sm90 import scaled_ce_forward_fragment
 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.
+# The forward and every backward phase tile the hidden dimension by 64, so one
+# padding of H covers the full operation.
 HIDDEN_TILE = 64
 # dZ and dX both tile the token dimension by 128, so a backward wave is counted
 # in M128 tiles.
@@ -92,6 +89,19 @@ BACKWARD_M_TILE = 128
 BACKWARD_M_TILES_PER_WAVE = 8
 VALIDATE_TARGETS_ENV = "LIGER_SCALED_CE_SM90_VALIDATE_TARGETS"

+# Measured H100 winners for the representative large-vocabulary workload.
+# Unlisted shapes retain the fragment kernel's hardware-derived split.
+_FORWARD_CONFIG_BY_SHAPE = {
+    (16384, 4096, 131072, False): ScaledCEForwardFragmentConfig(split_n=4),
+    (16384, 4096, 131072, True): ScaledCEForwardFragmentConfig(split_n=2),
+    (32768, 4096, 131072, False): ScaledCEForwardFragmentConfig(split_n=2),
+    (32768, 4096, 131072, True): ScaledCEForwardFragmentConfig(split_n=2),
+}
+
+
+def _select_forward_config(tokens, hidden_size, vocab_size, return_entropy):
+    return _FORWARD_CONFIG_BY_SHAPE.get((tokens, hidden_size, vocab_size, return_entropy))
+

 def _validate_temperature(temperature):
    if isinstance(temperature, bool) or not isinstance(temperature, (int, float)):
@@ -181,26 +191,20 @@ def fused_scaled_cross_entropy_forward(
    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,
+    nll, entropy, lse = scaled_ce_forward_fragment(
+        x_padded,
+        weight_padded,
+        target,
+        temperature,
+        ignore_index,
+        return_entropy,
+        config=_select_forward_config(
+            x_padded.shape[0],
+            x_padded.shape[1],
+            weight_padded.shape[0],
            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


@@ -260,8 +264,8 @@ class LigerFusedScaledCrossEntropySM90Function(torch.autograd.Function):
    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.
+    per-token upstream gradient ``[M]``.  ``m_tiles_per_cluster`` is retained
+    for API compatibility; backward always uses fixed 1024-token waves.
    """

    @staticmethod
@@ -276,6 +280,7 @@ class LigerFusedScaledCrossEntropySM90Function(torch.autograd.Function):
        m_tiles_per_cluster=1,
        return_entropy=False,
    ):
+        ctx.set_materialize_grads(False)
        _validate_temperature(temperature)
        nll, entropy, lse, x_padded, weight_padded, hidden_size = fused_scaled_cross_entropy_forward(
            _input,
@@ -299,6 +304,8 @@ class LigerFusedScaledCrossEntropySM90Function(torch.autograd.Function):
    @amp_custom_bwd
    def backward(ctx, grad_output, grad_entropy=None):
        x_padded, weight_padded, target, lse, entropy = ctx.saved_tensors
+        if grad_output is None:
+            grad_output = torch.zeros(target.shape[0], device=target.device, dtype=torch.float32)
        grad_input, grad_weight = fused_scaled_cross_entropy_backward(
            grad_output,
            x_padded,
@@ -0,0 +1,265 @@
+"""Correctness checks for the sole SM90 fragment forward."""
+
+import pytest
+import torch
+
+import liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_fragment_sm90 as fragment_module
+
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_fragment_sm90 import CLUSTER_M
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_fragment_sm90 import EPILOGUE_SLICE_N
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_fragment_sm90 import LOGICAL_N
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_fragment_sm90 import ScaledCEForwardFragmentConfig
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_fragment_sm90 import _resolve_split
+from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_forward_fragment_sm90 import scaled_ce_forward_fragment
+from test.utils import assert_verbose_allclose
+from test.utils import set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0),
+    reason="the fragment forward requires a Hopper (SM90) GPU",
+)
+
+
+def _torch_oracle(x, weight, target, temperature, ignore_index):
+    logits = (x.float() @ weight.float().t()) / temperature
+    lse = torch.logsumexp(logits, dim=-1)
+    probs = torch.softmax(logits, dim=-1)
+    weighted = (probs * logits).sum(dim=-1)
+    safe_target = target.masked_fill(target == ignore_index, 0)
+    target_logits = logits.gather(1, safe_target[:, None]).squeeze(1)
+    valid = target != ignore_index
+    zeros = torch.zeros_like(lse)
+    nll = torch.where(valid, lse - target_logits, zeros)
+    entropy = torch.where(valid, lse - weighted, zeros)
+    return nll, lse, entropy
+
+
+def _assert_matches_oracle(
+    x,
+    weight,
+    target,
+    temperature,
+    ignore_index,
+    return_entropy,
+    config,
+):
+    actual_nll, actual_entropy, actual_lse = scaled_ce_forward_fragment(
+        x,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+        return_entropy,
+        config,
+    )
+    expected_nll, expected_lse, expected_entropy = _torch_oracle(
+        x,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+    )
+
+    assert actual_nll.dtype == torch.float32
+    assert actual_lse.dtype == torch.float32
+    assert_verbose_allclose(actual_nll, expected_nll, atol=5e-2, rtol=2e-2)
+    assert_verbose_allclose(actual_lse, expected_lse, atol=5e-2, rtol=2e-2)
+    if return_entropy:
+        assert actual_entropy.dtype == torch.bfloat16
+        assert_verbose_allclose(actual_entropy.float(), expected_entropy, atol=5e-2, rtol=5e-2)
+    else:
+        assert actual_entropy is None
+    return actual_nll, actual_entropy, actual_lse
+
+
+@pytest.mark.parametrize(
+    ("tokens", "expected"),
+    [(2048, (9, 8, 2)), (4096, (5, 4, 2))],
+)
+def test_fragment_explicit_cluster_target_uses_uneven_split(tokens, expected):
+    config = ScaledCEForwardFragmentConfig(target_cluster_pairs=66)
+    split_n, base_split_n, extra_m_pairs, num_m_tiles, _ = _resolve_split(config, tokens, 131072)
+    num_m_pairs = (num_m_tiles + CLUSTER_M - 1) // CLUSTER_M
+
+    assert (split_n, base_split_n, extra_m_pairs) == expected
+    assert num_m_pairs * base_split_n + extra_m_pairs == 66
+
+
+def test_fragment_auto_split_uses_hardware_cluster_capacity(monkeypatch):
+    calls = []
+
+    def fake_max_active_clusters(cluster_m):
+        calls.append(cluster_m)
+        return 66
+
+    monkeypatch.setattr(fragment_module, "_max_active_clusters", fake_max_active_clusters)
+    split_n, base_split_n, extra_m_pairs, _, _ = _resolve_split(
+        ScaledCEForwardFragmentConfig(),
+        2048,
+        131072,
+    )
+
+    assert calls == [CLUSTER_M]
+    assert (split_n, base_split_n, extra_m_pairs) == (9, 8, 2)
+
+
+def test_fragment_uneven_split_entropy_matches_torch_oracle():
+    set_seed(40)
+    tokens = 512
+    hidden_size = 64
+    vocab_size = 1600
+    temperature = 0.7
+    ignore_index = -17
+    config = ScaledCEForwardFragmentConfig(target_cluster_pairs=7)
+    split_n, base_split_n, extra_m_pairs, _, num_logical_n_tiles = _resolve_split(
+        config,
+        tokens,
+        vocab_size,
+    )
+
+    assert (split_n, base_split_n, extra_m_pairs) == (4, 3, 1)
+    assert num_logical_n_tiles == 5
+
+    x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) * 0.15
+    weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
+    target = torch.randint(vocab_size, (tokens,), device="cuda", dtype=torch.long)
+    target[0] = vocab_size - 1
+    target[1] = ignore_index
+
+    actual_nll, actual_entropy, _ = _assert_matches_oracle(
+        x,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+        True,
+        config,
+    )
+    assert actual_nll[1].item() == 0.0
+    assert actual_entropy[1].item() == 0.0
+
+
+@pytest.mark.parametrize(
+    ("tokens", "hidden_size", "vocab_size", "temperature", "ignore_index", "return_entropy"),
+    [
+        (129, 64, 257, 1.1, -11, True),
+        (256, 64, 1600, 1.0, -7, False),
+        (256, 64, 1600, 0.8, -7, True),
+        (257, 70, 1937, 0.7, -23, False),
+        (257, 70, 1937, 1.3, -23, True),
+    ],
+    ids=["small-v-entropy", "k64-nll", "k64-entropy", "ragged-mhv-nll", "ragged-mhv-entropy"],
+)
+def test_fragment_matches_independent_torch_oracle(
+    tokens,
+    hidden_size,
+    vocab_size,
+    temperature,
+    ignore_index,
+    return_entropy,
+):
+    set_seed(41 + hidden_size + int(return_entropy))
+    x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) * 0.15
+    weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
+    target = torch.randint(vocab_size, (tokens,), device="cuda", dtype=torch.long)
+    target[0] = vocab_size - 1
+    target[1] = ignore_index
+    target[-1] = vocab_size - 1
+
+    actual_nll, actual_entropy, _ = _assert_matches_oracle(
+        x,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+        return_entropy,
+        ScaledCEForwardFragmentConfig(split_n=min(2, (vocab_size + LOGICAL_N - 1) // LOGICAL_N)),
+    )
+    assert actual_nll[1].item() == 0.0
+    if return_entropy:
+        assert actual_entropy[1].item() == 0.0
+
+
+def test_fragment_two_n160_accumulators_four_n80_slices_k4096():
+    set_seed(160)
+    tokens = 133
+    hidden_size = 4095
+    vocab_size = 1277
+    temperature = 0.75
+    ignore_index = -29
+    config = ScaledCEForwardFragmentConfig(split_n=2)
+    x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) * 0.15
+    weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
+    target = torch.randint(vocab_size, (tokens,), device="cuda", dtype=torch.long)
+    target[:8] = torch.tensor(
+        [79, 80, 159, 160, 239, 240, 319, vocab_size - 1],
+        device="cuda",
+    )
+    target[8] = ignore_index
+    target[-1] = vocab_size - 1
+
+    actual_nll, actual_entropy, _ = _assert_matches_oracle(
+        x,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+        True,
+        config,
+    )
+
+    assert LOGICAL_N == 320
+    assert EPILOGUE_SLICE_N == 80
+    assert (vocab_size + LOGICAL_N - 1) // LOGICAL_N == 4
+    assert (vocab_size - 1) % LOGICAL_N >= 3 * EPILOGUE_SLICE_N
+    assert actual_nll[8].item() == 0.0
+    assert actual_entropy[8].item() == 0.0
+
+
+def test_fragment_repeated_calls_preserve_nll_and_refresh_partial_statistics():
+    set_seed(161)
+    tokens = 129
+    hidden_size = 64
+    vocab_size = 641
+    temperature = 0.9
+    ignore_index = -37
+    config = ScaledCEForwardFragmentConfig(split_n=2)
+    weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
+    target = torch.randint(vocab_size, (tokens,), device="cuda", dtype=torch.long)
+    target[0] = vocab_size - 1
+    target[1] = ignore_index
+    x0 = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) * 0.1
+    x1 = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) * 0.2
+
+    nll_only = _assert_matches_oracle(
+        x0,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+        False,
+        config,
+    )
+    first = _assert_matches_oracle(
+        x0,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+        True,
+        config,
+    )
+    second = _assert_matches_oracle(
+        x1,
+        weight,
+        target,
+        temperature,
+        ignore_index,
+        True,
+        config,
+    )
+
+    assert_verbose_allclose(first[0], nll_only[0], atol=1e-5, rtol=1e-5)
+    assert_verbose_allclose(first[2], nll_only[2], atol=1e-5, rtol=1e-5)
+    assert not torch.equal(first[0], second[0])
+    assert not torch.equal(first[1], second[1])
@@ -199,6 +199,34 @@ def test_sm90_differentiable_entropy_matches_torch(m_tiles_per_cluster):
    assert_verbose_allclose(wa.grad.float(), wr.grad.float(), atol=3e-2, rtol=1.5e-1)


+def test_sm90_entropy_only_backward_matches_torch():
+    set_seed(144)
+    x, weight, target = _make_inputs(193, 70, 517, ignore_stride=7)
+    grad_entropy = torch.rand(193, device="cuda", dtype=torch.float32) - 0.5
+
+    function = _frontend_function()
+    actual_x = x.detach().clone().requires_grad_(True)
+    actual_weight = weight.detach().clone().requires_grad_(True)
+    _, actual_entropy = function.apply(
+        actual_x,
+        actual_weight,
+        target,
+        0.8,
+        -100,
+        1,
+        True,
+    )
+    actual_entropy.backward(grad_entropy.to(actual_entropy.dtype))
+
+    expected_x = x.detach().clone().requires_grad_(True)
+    expected_weight = weight.detach().clone().requires_grad_(True)
+    _, expected_entropy = _reference_nll_entropy(expected_x, expected_weight, target, -100, 0.8)
+    expected_entropy.backward(grad_entropy)
+
+    assert_verbose_allclose(actual_x.grad.float(), expected_x.grad.float(), atol=3e-2, rtol=1.5e-1)
+    assert_verbose_allclose(actual_weight.grad.float(), expected_weight.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)
@@ -386,30 +414,44 @@ def test_sm90_backward_uses_one_fused_kernel(monkeypatch):


@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):
+@pytest.mark.parametrize("return_entropy", [False, True], ids=["nll", "entropy"])
+def test_sm90_m_tiles_values_route_fragment_forward(monkeypatch, m_tiles_per_cluster, return_entropy):
    module = _sm90_module()
-    calls = {}
+    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_fragment(x, weight, target, temperature, ignore_index, return_entropy, config):
+        calls.append((temperature, ignore_index, return_entropy, config))
+        zeros = torch.zeros(x.shape[0], device=x.device)
+        entropy = zeros.to(torch.bfloat16) if return_entropy else None
+        return zeros, entropy, zeros

-    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)
+    monkeypatch.setattr(module, "scaled_ce_forward_fragment", fake_fragment)

    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)
+    module.fused_scaled_cross_entropy_forward(
+        x,
+        weight,
+        target,
+        1.0,
+        -100,
+        m_tiles_per_cluster,
+        return_entropy,
+    )

-    if m_tiles_per_cluster == 1:
-        assert calls == {"m_outer": (1.0, -100)}
-    else:
-        assert calls == {"m_fast": (1.0, m_tiles_per_cluster)}
+    assert calls == [(1.0, -100, return_entropy, None)]
+
+
+def test_sm90_long_sequence_forward_tuning_table():
+    module = _sm90_module()
+
+    assert module._select_forward_config(8192, 4096, 131072, False) is None
+    assert module._select_forward_config(16384, 4096, 131072, False).split_n == 4
+    assert module._select_forward_config(16384, 4096, 131072, True).split_n == 2
+    assert module._select_forward_config(32768, 4096, 131072, False).split_n == 2
+    assert module._select_forward_config(32768, 4096, 131072, True).split_n == 2
+    assert module._select_forward_config(32768, 4096, 65536, False) is None


 def test_sm90_rejects_scalar_and_mismatched_grad_output():