From 62b691b7219114dad2fa2f795ddfd6ec6727e5b5 Mon Sep 17 00:00:00 2001 From: Jari Kolehmainen Date: Tue, 11 Aug 2026 11:14:17 -0700 Subject: [PATCH] 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 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 45 +- ...led_cross_entropy_forward_fragment_sm90.py | 1022 ++++++++++++++ ...scaled_cross_entropy_forward_mfast_sm90.py | 1216 ----------------- ...fused_scaled_cross_entropy_forward_sm90.py | 1104 --------------- .../_fused_scaled_cross_entropy_utils_sm90.py | 79 ++ .../ops/fused_scaled_cross_entropy_sm90.py | 83 +- ...led_cross_entropy_forward_fragment_sm90.py | 265 ++++ .../test_fused_scaled_cross_entropy_sm90.py | 74 +- 8 files changed, 1508 insertions(+), 2380 deletions(-) create mode 100644 src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_fragment_sm90.py delete mode 100644 src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_mfast_sm90.py delete mode 100644 src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_sm90.py create mode 100644 src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_utils_sm90.py create mode 100644 test/cutedsl/test_fused_scaled_cross_entropy_forward_fragment_sm90.py diff --git a/README.md b/README.md index 6199ba9d..5e3fce33 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_fragment_sm90.py b/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_fragment_sm90.py new file mode 100644 index 00000000..0a960afe --- /dev/null +++ b/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_fragment_sm90.py @@ -0,0 +1,1022 @@ +"""SM90 two-fragment fused scaled cross entropy forward. + +Each cluster-M2 CTA computes an M128xN320 logical tile with a three-stage K64 +pipeline. WG0 is the 24-register TMA producer and multicasts two N160 weight +panels to the cluster peer. WG1 and WG2 each own 64 token rows, two FP32 N160 +accumulators, and the online-softmax epilogue. + +The consumer read and release states are deliberately separate. Stage zero is +issued without a release. Each later stage is committed, ``wait_group(1)`` +retires the previous WGMMA group, and only then is that previous AB stage +released. The tail uses ``wait_group(0)`` before the final release. This +one-stage lag prevents TMA from overwriting operands still consumed by WGMMA. + +Each N160 accumulator is staged as two N80 slices through two padded FP16 SMEM +buffers, producing four online-softmax folds per logical N320 tile. Vocabulary +tiles are split over cluster pairs; each CTA writes compact FP32 +``(max, sum, target[, weighted])`` statistics to HBM and a small finalizer +reduces the splits into per-token NLL/LSE and optional BF16 entropy. Entropy +uses the same online max rescaling as the softmax sum and is compiled out when +not requested. +""" + +from dataclasses import dataclass +from dataclasses import replace + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch + +from cutlass import BFloat16 +from cutlass import Float16 +from cutlass import Float32 +from cutlass import Int32 +from cutlass import pipeline +from cutlass.pipeline.helpers import pipeline_init_arrive +from cutlass.pipeline.helpers import pipeline_init_wait +from cutlass.utils import LayoutEnum +from cutlass.utils import hopper_helpers + +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import HOPPER_MAX_SMEM_BYTES +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import LN2 +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import LOG2_E +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import MASK_F32 +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import NEG_INF_F32 +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import _cute_stream +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import _fmax +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import _max_active_clusters +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import _pad_hidden +from liger_kernel.ops.cutedsl.ops._fused_scaled_cross_entropy_utils_sm90 import _validate +from liger_kernel.ops.cutedsl.ops.utils import to_cute_tensor + +THREADS_PER_CTA = 384 +PRODUCER_WARP_GROUP = 0 +NUM_MMA_WARP_GROUPS = 2 +WARPS_PER_WARP_GROUP = 4 +TILE_M = 128 +TILE_K = 64 +STAGES = 3 +CLUSTER_M = 2 +ACCUMULATOR_N = 160 +NUM_ACCUMULATORS = 2 +LOGICAL_N = ACCUMULATOR_N * NUM_ACCUMULATORS +EPILOGUE_SLICE_N = 80 +CHUNK_N = 40 +LOGIT_BUFFERS = 2 +LOGIT_PAD = 8 +LOGIT_STRIDE = EPILOGUE_SLICE_N + LOGIT_PAD +EPILOGUE_BARRIER_ID = 4 +PRODUCER_REGISTERS = 24 +MMA_REGISTERS = 240 +USABLE_REGISTER_BUDGET = 64512 +SHARED_MEMORY_BYTES_BASE = ( + (TILE_M + LOGICAL_N) * TILE_K * 2 * STAGES + + LOGIT_BUFFERS * TILE_M * LOGIT_STRIDE * 2 + + 3 * TILE_M * 4 + + 16 * STAGES + + 3 * 1024 +) + +_compile_cache = {} + + +@dataclass(frozen=True) +class ScaledCEForwardFragmentConfig: + """Vocabulary-split tuning for the fixed M128xN320 fragment kernel. + + ``return_entropy`` compile-time enables weighted softmax statistics. + """ + + fast_math: bool = True + return_entropy: bool = False + split_n: int = 0 + base_split_n: int = 0 + extra_m_pairs: int = 0 + target_cluster_pairs: int = 0 + max_split_n: int = 9 + + def smem_bytes(self): + return SHARED_MEMORY_BYTES_BASE + (TILE_M if self.return_entropy else 1) * 4 + + def register_total(self): + return 128 * PRODUCER_REGISTERS + NUM_MMA_WARP_GROUPS * 128 * MMA_REGISTERS + + +class _ScaledCEForwardFragmentSM90: + def __init__(self, config: ScaledCEForwardFragmentConfig): + self.fast_math = config.fast_math + self.return_entropy = config.return_entropy + self.split_n = config.split_n + self.base_split_n = config.base_split_n + self.extra_m_pairs = config.extra_m_pairs + self.a_tile_shape = (TILE_M, ACCUMULATOR_N, TILE_K) + self.b_tile_shape = (TILE_M, ACCUMULATOR_N, TILE_K) + self.buffer_align = 1024 + self.epilogue_barrier = pipeline.NamedBarrier( + barrier_id=EPILOGUE_BARRIER_ID, + num_threads=NUM_MMA_WARP_GROUPS * 128, + ) + + @staticmethod + def _make_tma(tensor, smem_layout_staged, smem_tile, multicast=1): + smem_layout = cute.slice_(smem_layout_staged, (None, None, 0)) + op = ( + cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + if multicast == 1 + else cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp() + ) + return cute.nvgpu.cpasync.make_tiled_tma_atom( + op, + tensor, + smem_layout, + smem_tile, + num_multicast=multicast, + ) + + @cute.jit + def __call__( + self, + x: cute.Tensor, + weight: cute.Tensor, + target: cute.Tensor, + partial_max: cute.Tensor, + partial_sum: cute.Tensor, + partial_target: cute.Tensor, + partial_weighted: cute.Tensor, + inverse_temperature: Float32, + stream: cuda.CUstream, + ): + x_layout = LayoutEnum.from_tensor(x) + weight_layout = LayoutEnum.from_tensor(weight) + tiled_mma = hopper_helpers.make_trivial_tiled_mma( + BFloat16, + BFloat16, + x_layout.sm90_mma_major_mode(), + weight_layout.sm90_mma_major_mode(), + Float32, + (NUM_MMA_WARP_GROUPS, 1, 1), + (64, ACCUMULATOR_N), + ) + a_smem_layout = hopper_helpers.make_smem_layout_a( + x_layout, + self.a_tile_shape, + BFloat16, + STAGES, + ) + b_smem_layout = hopper_helpers.make_smem_layout_b( + weight_layout, + self.b_tile_shape, + BFloat16, + STAGES, + ) + tma_x, tma_tensor_x = self._make_tma(x, a_smem_layout, (TILE_M, TILE_K)) + tma_w, tma_tensor_w = self._make_tma( + weight, + b_smem_layout, + (ACCUMULATOR_N, TILE_K), + multicast=CLUSTER_M, + ) + + @cute.struct + class SharedStorage: + pipeline: cute.struct.MemRange[cutlass.Int64, STAGES * 2] + sX: cute.struct.Align[ + cute.struct.MemRange[BFloat16, cute.cosize(a_smem_layout)], + self.buffer_align, + ] + sW0: cute.struct.Align[ + cute.struct.MemRange[BFloat16, cute.cosize(b_smem_layout)], + self.buffer_align, + ] + sW1: cute.struct.Align[ + cute.struct.MemRange[BFloat16, cute.cosize(b_smem_layout)], + self.buffer_align, + ] + sLogits: cute.struct.Align[ + cute.struct.MemRange[ + Float16, + LOGIT_BUFFERS * TILE_M * LOGIT_STRIDE, + ], + self.buffer_align, + ] + segment_max: cute.struct.Align[cute.struct.MemRange[Float32, TILE_M], 16] + segment_sum: cute.struct.Align[cute.struct.MemRange[Float32, TILE_M], 16] + segment_target: cute.struct.Align[cute.struct.MemRange[Float32, TILE_M], 16] + segment_weighted: cute.struct.Align[ + cute.struct.MemRange[Float32, TILE_M if self.return_entropy else 1], + 16, + ] + + self.shared_storage = SharedStorage + num_m_tiles = cute.ceil_div(x.shape[0], TILE_M) + num_m_pairs = cute.ceil_div(num_m_tiles, CLUSTER_M) + num_cluster_pairs = num_m_pairs * self.base_split_n + self.extra_m_pairs + grid = (CLUSTER_M, 1, num_cluster_pairs) + self.kernel( + tma_x, + tma_tensor_x, + tma_w, + tma_tensor_w, + target, + partial_max, + partial_sum, + partial_target, + partial_weighted, + inverse_temperature, + tiled_mma, + a_smem_layout, + b_smem_layout, + ).launch( + grid=grid, + block=(THREADS_PER_CTA, 1, 1), + cluster=(CLUSTER_M, 1, 1), + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.jit + def _stage_slice( + self, + accum: cute.Tensor, + s_logits: cute.Tensor, + buffer: Int32, + row_base: Int32, + col_base: Int32, + slice_base: Int32, + ): + for i in cutlass.range_constexpr(cute.size(accum)): + row = row_base + ((i % 4) // 2) * 8 + accum_col = col_base + (i // 4) * 8 + (i % 2) + if accum_col >= slice_base and accum_col < slice_base + EPILOGUE_SLICE_N: + s_logits[buffer, row, accum_col - slice_base] = Float16(accum[i]) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_acq_rel_cta() + self.epilogue_barrier.arrive_and_wait() + + @cute.jit + def _fold_slice( + self, + s_logits: cute.Tensor, + buffer: Int32, + row: Int32, + segment: Int32, + slice_global_base: Int32, + valid_cols: Int32, + target_col: Int32, + inverse_temperature: Float32, + run_max: Float32, + run_sum: Float32, + run_target: Float32, + run_weighted: Float32, + ): + segment_n = EPILOGUE_SLICE_N // 2 + segment_base = segment * segment_n + values = cute.make_rmem_tensor((CHUNK_N,), Float32) + for ci in cutlass.range_constexpr(segment_n // CHUNK_N): + c0 = segment_base + ci * CHUNK_N + chunk_max = Float32(MASK_F32) + if c0 + CHUNK_N <= valid_cols: + for j in cutlass.range_constexpr(CHUNK_N): + value = Float32(s_logits[buffer, row, c0 + j]) * inverse_temperature + values[j] = value + chunk_max = _fmax(chunk_max, value) + else: + for j in cutlass.range_constexpr(CHUNK_N): + value = Float32(s_logits[buffer, row, c0 + j]) * inverse_temperature + if c0 + j >= valid_cols: + value = Float32(MASK_F32) + values[j] = value + chunk_max = _fmax(chunk_max, value) + + new_max = _fmax(run_max, chunk_max) + chunk_sum = Float32(0.0) + if cutlass.const_expr(self.return_entropy): + chunk_weighted = Float32(0.0) + for j in cutlass.range_constexpr(CHUNK_N): + exp_value = cute.math.exp2( + (values[j] - new_max) * LOG2_E, + fastmath=self.fast_math, + ) + chunk_sum += exp_value + chunk_weighted += exp_value * values[j] + old_scale = cute.math.exp2( + (run_max - new_max) * LOG2_E, + fastmath=self.fast_math, + ) + run_sum = run_sum * old_scale + chunk_sum + run_weighted = run_weighted * old_scale + chunk_weighted + else: + for j in cutlass.range_constexpr(CHUNK_N): + chunk_sum += cute.math.exp2( + (values[j] - new_max) * LOG2_E, + fastmath=self.fast_math, + ) + run_sum = ( + run_sum + * cute.math.exp2( + (run_max - new_max) * LOG2_E, + fastmath=self.fast_math, + ) + + chunk_sum + ) + run_max = new_max + + target_begin = slice_global_base + segment_base + target_end = target_begin + segment_n + if target_col >= target_begin and target_col < target_end: + run_target += Float32(s_logits[buffer, row, target_col - slice_global_base]) * inverse_temperature + return run_max, run_sum, run_target, run_weighted + + @cute.kernel + def kernel( + self, + tma_x: cute.CopyAtom, + x: cute.Tensor, + tma_w: cute.CopyAtom, + weight: cute.Tensor, + target: cute.Tensor, + output_partial_max: cute.Tensor, + output_partial_sum: cute.Tensor, + output_partial_target: cute.Tensor, + output_partial_weighted: cute.Tensor, + inverse_temperature: Float32, + tiled_mma: cute.TiledMma, + a_smem_layout: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + ): + tid, _, _ = cute.arch.thread_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + warp_group = cute.arch.make_warp_uniform(tid // 128) + lane = tid % 32 + local_warp = (tid % 128) // 32 + local_tid = tid % 128 + cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + _, _, cluster_work = cute.arch.block_idx() + + num_m_tiles = cute.ceil_div(target.shape[0], TILE_M) + if cutlass.const_expr(self.extra_m_pairs == 0): + m_pair = cluster_work // self.base_split_n + split_id = cluster_work % self.base_split_n + split_count = self.base_split_n + else: + extra_work = self.extra_m_pairs * self.split_n + m_pair = Int32(0) + split_id = Int32(0) + split_count = Int32(self.base_split_n) + if cluster_work < extra_work: + m_pair = cluster_work // self.split_n + split_id = cluster_work % self.split_n + split_count = Int32(self.split_n) + else: + normal_work = cluster_work - extra_work + m_pair = self.extra_m_pairs + normal_work // self.base_split_n + split_id = normal_work % self.base_split_n + split_count = Int32(self.base_split_n) + raw_pid_m = m_pair * CLUSTER_M + cluster_rank + store_enabled = raw_pid_m < num_m_tiles + pid_m = cutlass.min(raw_pid_m, num_m_tiles - 1) + output_work = raw_pid_m * self.split_n + split_id + + if warp == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_x) + cute.nvgpu.cpasync.prefetch_descriptor(tma_w) + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + s_x = storage.sX.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner) + s_w0 = storage.sW0.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) + s_w1 = storage.sW1.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) + s_logits = storage.sLogits.get_tensor( + cute.make_layout( + (LOGIT_BUFFERS, TILE_M, EPILOGUE_SLICE_N), + stride=(TILE_M * LOGIT_STRIDE, LOGIT_STRIDE, 1), + ) + ) + segment_max = storage.segment_max.get_tensor(cute.make_layout(TILE_M)) + segment_sum = storage.segment_sum.get_tensor(cute.make_layout(TILE_M)) + segment_target = storage.segment_target.get_tensor(cute.make_layout(TILE_M)) + segment_weighted = storage.segment_weighted.get_tensor( + cute.make_layout(TILE_M if cutlass.const_expr(self.return_entropy) else 1) + ) + + g_x = cute.local_tile( + x, + cute.slice_(self.a_tile_shape, (None, 0, None)), + (None, None, None), + ) + g_w = cute.local_tile( + weight, + cute.slice_(self.b_tile_shape, (0, None, None)), + (None, None, None), + ) + one = cute.make_layout(1) + t_x_s, t_x_g = cute.nvgpu.cpasync.tma_partition( + tma_x, + 0, + one, + cute.group_modes(s_x, 0, 2), + cute.group_modes(g_x, 0, 2), + ) + + cta_layout = cute.make_layout((CLUSTER_M, 1, 1)) + cluster_coord = cta_layout.get_flat_coord(cluster_rank) + b_cta_layout = cute.make_layout(cute.slice_(cta_layout, (None, 0, 0)).shape) + b_mcast_mask = cute.make_layout_image_mask(cta_layout, cluster_coord, mode=0) + b_mcast_rank = cluster_coord[0] + t_w_s0, t_w_g0 = cute.nvgpu.cpasync.tma_partition( + tma_w, + b_mcast_rank, + b_cta_layout, + cute.group_modes(s_w0, 0, 2), + cute.group_modes(g_w, 0, 2), + ) + t_w_s1, t_w_g1 = cute.nvgpu.cpasync.tma_partition( + tma_w, + b_mcast_rank, + b_cta_layout, + cute.group_modes(s_w1, 0, 2), + cute.group_modes(g_w, 0, 2), + ) + x_stage_bytes = cute.size_in_bytes(BFloat16, cute.slice_(a_smem_layout, (None, None, 0))) + w_stage_bytes = cute.size_in_bytes(BFloat16, cute.slice_(b_smem_layout, (None, None, 0))) + mainloop = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.pipeline.data_ptr(), + num_stages=STAGES, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + NUM_MMA_WARP_GROUPS * WARPS_PER_WARP_GROUP * CLUSTER_M, + ), + tx_count=x_stage_bytes + NUM_ACCUMULATORS * w_stage_bytes, + cta_layout_vmnk=cute.make_layout((1, CLUSTER_M, 1, 1)), + tidx=tid, + defer_sync=True, + ) + pipeline_init_arrive(cluster_shape_mn=(CLUSTER_M, 1), is_relaxed=True) + pipeline_init_wait(cluster_shape_mn=(CLUSTER_M, 1)) + + num_k_tiles = cute.size(g_x, mode=[3]) + num_logical_n_tiles = cute.ceil_div(weight.shape[0], LOGICAL_N) + num_split_tiles = cute.ceil_div(num_logical_n_tiles - split_id, split_count) + + if warp_group == PRODUCER_WARP_GROUP: + cute.arch.setmaxregister_decrease(PRODUCER_REGISTERS) + load_index = Int32(0) + load_phase = Int32(1) + if warp == 0: + t_x_g_m = t_x_g[(None, pid_m, None, 0)] + for local_n in range(num_split_tiles): + logical_n = split_id + local_n * split_count + fragment_base = logical_n * NUM_ACCUMULATORS + for k_tile in range(num_k_tiles): + pstate = pipeline.PipelineState( + STAGES, + Int32(0), + load_index, + load_phase, + ) + mainloop.producer_acquire(pstate) + barrier = mainloop.producer_get_barrier(pstate) + cute.copy( + tma_x, + t_x_g_m[(None, k_tile)], + t_x_s[(None, load_index)], + tma_bar_ptr=barrier, + mcast_mask=0, + ) + cute.copy( + tma_w, + t_w_g0[(None, fragment_base + 0, None, 0)][(None, k_tile)], + t_w_s0[(None, load_index)], + tma_bar_ptr=barrier, + mcast_mask=b_mcast_mask, + ) + cute.copy( + tma_w, + t_w_g1[(None, fragment_base + 1, None, 0)][(None, k_tile)], + t_w_s1[(None, load_index)], + tma_bar_ptr=barrier, + mcast_mask=b_mcast_mask, + ) + mainloop.producer_commit(pstate) + pstate.advance() + load_index = pstate.index + load_phase = pstate.phase + tail = pipeline.PipelineState( + STAGES, + Int32(0), + load_index, + load_phase, + ) + mainloop.producer_tail(tail) + else: + cute.arch.setmaxregister_increase(MMA_REGISTERS) + mma_wg = warp_group - 1 + mma_wg_layout = cute.make_layout(NUM_MMA_WARP_GROUPS, stride=128) + thr_mma = tiled_mma.get_slice(mma_wg_layout(Int32(mma_wg))) + r_x = tiled_mma.make_fragment_A(thr_mma.partition_A(s_x)) + r_w0 = tiled_mma.make_fragment_B(thr_mma.partition_B(s_w0)) + r_w1 = tiled_mma.make_fragment_B(thr_mma.partition_B(s_w1)) + accum_shape = thr_mma.partition_shape_C((TILE_M, ACCUMULATOR_N)) + accum0 = cute.make_rmem_tensor(accum_shape, Float32) + accum1 = cute.make_rmem_tensor(accum_shape, Float32) + + row_base = mma_wg * 64 + local_warp * 16 + lane // 4 + col_base = (lane % 4) * 2 + epi_row = mma_wg * 64 + local_tid % 64 + epi_segment = local_tid // 64 + global_row = raw_pid_m * TILE_M + epi_row + target_col = Int32(-1) + if store_enabled and global_row < target.shape[0]: + target_col = Int32(target[global_row]) + + run_max = Float32(NEG_INF_F32) + run_sum = Float32(0.0) + run_target = Float32(0.0) + run_weighted = Float32(0.0) + read_index = Int32(0) + read_phase = Int32(0) + release_index = Int32(0) + release_phase = Int32(0) + tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, True) + + for local_n in range(num_split_tiles): + logical_n = split_id + local_n * split_count + accum0.fill(0.0) + accum1.fill(0.0) + + # K stage 0: issue the first multi-fragment WGMMA group. No AB + # stage is released because no earlier group has retired. + rstate = pipeline.PipelineState( + STAGES, + Int32(0), + read_index, + read_phase, + ) + mainloop.consumer_wait(rstate) + cute.nvgpu.warpgroup.fence() + for k_block in cutlass.range_constexpr(TILE_K // 16): + coord = (None, None, k_block, read_index) + cute.gemm(tiled_mma, accum0, r_x[coord], r_w0[coord], accum0) + cute.gemm(tiled_mma, accum1, r_x[coord], r_w1[coord], accum1) + cute.nvgpu.warpgroup.commit_group() + rstate.advance() + read_index = rstate.index + read_phase = rstate.phase + + # The release state intentionally trails the read state. Once + # wait_group(1) retires the previous WGMMA group, and only + # then, its AB stage is returned to the producer. + for _ in range(1, num_k_tiles): + rstate = pipeline.PipelineState( + STAGES, + Int32(0), + read_index, + read_phase, + ) + mainloop.consumer_wait(rstate) + cute.nvgpu.warpgroup.fence() + for k_block in cutlass.range_constexpr(TILE_K // 16): + coord = (None, None, k_block, read_index) + cute.gemm(tiled_mma, accum0, r_x[coord], r_w0[coord], accum0) + cute.gemm(tiled_mma, accum1, r_x[coord], r_w1[coord], accum1) + cute.nvgpu.warpgroup.commit_group() + cute.nvgpu.warpgroup.wait_group(1) + release_state = pipeline.PipelineState( + STAGES, + Int32(0), + release_index, + release_phase, + ) + mainloop.consumer_release(release_state) + release_state.advance() + release_index = release_state.index + release_phase = release_state.phase + rstate.advance() + read_index = rstate.index + read_phase = rstate.phase + + # Drain the final live WGMMA group before releasing its stage. + cute.nvgpu.warpgroup.wait_group(0) + release_state = pipeline.PipelineState( + STAGES, + Int32(0), + release_index, + release_phase, + ) + mainloop.consumer_release(release_state) + release_state.advance() + release_index = release_state.index + release_phase = release_state.phase + + logical_base = logical_n * LOGICAL_N + + # Each N160 accumulator is folded as two N80 slices. Both + # buffers are fully consumed before either is overwritten. + self._stage_slice(accum0, s_logits, Int32(0), row_base, col_base, Int32(0)) + self._stage_slice( + accum0, + s_logits, + Int32(1), + row_base, + col_base, + Int32(EPILOGUE_SLICE_N), + ) + run_max, run_sum, run_target, run_weighted = self._fold_slice( + s_logits, + Int32(0), + epi_row, + epi_segment, + logical_base, + Int32(weight.shape[0]) - logical_base, + target_col, + inverse_temperature, + run_max, + run_sum, + run_target, + run_weighted, + ) + self.epilogue_barrier.arrive_and_wait() + self._stage_slice(accum1, s_logits, Int32(0), row_base, col_base, Int32(0)) + run_max, run_sum, run_target, run_weighted = self._fold_slice( + s_logits, + Int32(1), + epi_row, + epi_segment, + logical_base + EPILOGUE_SLICE_N, + Int32(weight.shape[0]) - (logical_base + EPILOGUE_SLICE_N), + target_col, + inverse_temperature, + run_max, + run_sum, + run_target, + run_weighted, + ) + self.epilogue_barrier.arrive_and_wait() + self._stage_slice( + accum1, + s_logits, + Int32(1), + row_base, + col_base, + Int32(EPILOGUE_SLICE_N), + ) + run_max, run_sum, run_target, run_weighted = self._fold_slice( + s_logits, + Int32(0), + epi_row, + epi_segment, + logical_base + ACCUMULATOR_N, + Int32(weight.shape[0]) - (logical_base + ACCUMULATOR_N), + target_col, + inverse_temperature, + run_max, + run_sum, + run_target, + run_weighted, + ) + self.epilogue_barrier.arrive_and_wait() + run_max, run_sum, run_target, run_weighted = self._fold_slice( + s_logits, + Int32(1), + epi_row, + epi_segment, + logical_base + ACCUMULATOR_N + EPILOGUE_SLICE_N, + Int32(weight.shape[0]) - (logical_base + ACCUMULATOR_N + EPILOGUE_SLICE_N), + target_col, + inverse_temperature, + run_max, + run_sum, + run_target, + run_weighted, + ) + self.epilogue_barrier.arrive_and_wait() + + if epi_segment == 1: + segment_max[epi_row] = run_max + segment_sum[epi_row] = run_sum + segment_target[epi_row] = run_target + if cutlass.const_expr(self.return_entropy): + segment_weighted[epi_row] = run_weighted + self.epilogue_barrier.arrive_and_wait() + + if epi_segment == 0 and store_enabled: + other_max = segment_max[epi_row] + combined_max = _fmax(run_max, other_max) + if cutlass.const_expr(self.return_entropy): + run_scale = cute.math.exp2( + (run_max - combined_max) * LOG2_E, + fastmath=self.fast_math, + ) + other_scale = cute.math.exp2( + (other_max - combined_max) * LOG2_E, + fastmath=self.fast_math, + ) + combined_sum = run_sum * run_scale + segment_sum[epi_row] * other_scale + combined_weighted = run_weighted * run_scale + segment_weighted[epi_row] * other_scale + else: + combined_sum = run_sum * cute.math.exp2( + (run_max - combined_max) * LOG2_E, + fastmath=self.fast_math, + ) + segment_sum[epi_row] * cute.math.exp2( + (other_max - combined_max) * LOG2_E, + fastmath=self.fast_math, + ) + output_partial_max[output_work, epi_row] = combined_max + output_partial_sum[output_work, epi_row] = combined_sum + output_partial_target[output_work, epi_row] = run_target + segment_target[epi_row] + if cutlass.const_expr(self.return_entropy): + output_partial_weighted[output_work, epi_row] = combined_weighted + + +class _ScaledCEForwardFragmentFinalizeSM90: + """Combine split vocabulary statistics into final NLL/LSE/entropy.""" + + def __init__(self, split_n, base_split_n, extra_m_pairs, fast_math, return_entropy): + self.split_n = split_n + self.base_split_n = base_split_n + self.extra_m_pairs = extra_m_pairs + self.fast_math = fast_math + self.return_entropy = return_entropy + + @cute.jit + def __call__( + self, + partial_max: cute.Tensor, + partial_sum: cute.Tensor, + partial_target: cute.Tensor, + partial_weighted: cute.Tensor, + target: cute.Tensor, + lse: cute.Tensor, + token_loss: cute.Tensor, + entropy: cute.Tensor, + ignore_index: Int32, + stream: cuda.CUstream, + ): + num_m_tiles = cute.ceil_div(target.shape[0], TILE_M) + self.kernel( + partial_max, + partial_sum, + partial_target, + partial_weighted, + target, + lse, + token_loss, + entropy, + ignore_index, + ).launch( + grid=(num_m_tiles, 1, 1), + block=(TILE_M, 1, 1), + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + partial_max: cute.Tensor, + partial_sum: cute.Tensor, + partial_target: cute.Tensor, + partial_weighted: cute.Tensor, + target: cute.Tensor, + lse: cute.Tensor, + token_loss: cute.Tensor, + entropy: cute.Tensor, + ignore_index: Int32, + ): + tid, _, _ = cute.arch.thread_idx() + pid_m, _, _ = cute.arch.block_idx() + row = pid_m * TILE_M + tid + if cutlass.const_expr(self.extra_m_pairs == 0): + split_count = self.base_split_n + else: + m_pair = pid_m // CLUSTER_M + split_count = Int32(self.base_split_n) + if m_pair < self.extra_m_pairs: + split_count = Int32(self.split_n) + row_max = Float32(NEG_INF_F32) + row_target = Float32(0.0) + for split_id in cutlass.range_constexpr(self.split_n): + if split_id < split_count: + work_id = pid_m * self.split_n + split_id + row_max = _fmax(row_max, partial_max[work_id, tid]) + row_target += partial_target[work_id, tid] + + row_sum = Float32(0.0) + row_weighted = Float32(0.0) + for split_id in cutlass.range_constexpr(self.split_n): + if split_id < split_count: + work_id = pid_m * self.split_n + split_id + if cutlass.const_expr(self.return_entropy): + split_scale = cute.math.exp2( + (partial_max[work_id, tid] - row_max) * LOG2_E, + fastmath=self.fast_math, + ) + row_sum += partial_sum[work_id, tid] * split_scale + row_weighted += partial_weighted[work_id, tid] * split_scale + else: + row_sum += partial_sum[work_id, tid] * cute.math.exp2( + (partial_max[work_id, tid] - row_max) * LOG2_E, + fastmath=self.fast_math, + ) + + if row < target.shape[0]: + row_lse = row_max + cute.math.log2(row_sum, fastmath=self.fast_math) * LN2 + lse[row] = row_lse + if target[row] == ignore_index: + token_loss[row] = Float32(0.0) + if cutlass.const_expr(self.return_entropy): + entropy[row] = BFloat16(0.0) + else: + token_loss[row] = row_lse - row_target + if cutlass.const_expr(self.return_entropy): + entropy[row] = BFloat16(row_lse - row_weighted / row_sum) + + +def _resolve_split(config, tokens, vocab_size): + num_m_tiles = (tokens + TILE_M - 1) // TILE_M + num_m_pairs = (num_m_tiles + CLUSTER_M - 1) // CLUSTER_M + num_logical_n_tiles = (vocab_size + LOGICAL_N - 1) // LOGICAL_N + if config.base_split_n: + return ( + config.split_n, + config.base_split_n, + config.extra_m_pairs, + num_m_tiles, + num_logical_n_tiles, + ) + if config.split_n: + return config.split_n, config.split_n, 0, num_m_tiles, num_logical_n_tiles + + target_cluster_pairs = config.target_cluster_pairs or _max_active_clusters(CLUSTER_M) + max_cluster_pairs = num_m_pairs * min(config.max_split_n, num_logical_n_tiles) + num_cluster_pairs = max( + num_m_pairs, + min(target_cluster_pairs, max_cluster_pairs), + ) + base_split_n, extra_m_pairs = divmod(num_cluster_pairs, num_m_pairs) + split_n = base_split_n + int(extra_m_pairs != 0) + return split_n, base_split_n, extra_m_pairs, num_m_tiles, num_logical_n_tiles + + +def scaled_ce_forward_fragment( + _input, + weight, + target, + temperature=1.0, + ignore_index=-100, + return_entropy=False, + config: ScaledCEForwardFragmentConfig = None, +): + """Run the fixed two-N160/four-N80 fragment forward.""" + _validate(_input, weight, target) + if config is None: + config = ScaledCEForwardFragmentConfig() + if config.return_entropy != return_entropy: + config = replace(config, return_entropy=return_entropy) + if min(config.split_n, config.base_split_n, config.extra_m_pairs, config.target_cluster_pairs) < 0: + raise ValueError("split and cluster tuning values must be non-negative") + if config.max_split_n < 1: + raise ValueError("max_split_n must be positive") + if config.extra_m_pairs and not config.base_split_n: + raise ValueError("extra_m_pairs requires an explicit base_split_n") + if config.smem_bytes() > HOPPER_MAX_SMEM_BYTES: + raise ValueError( + f"fragment kernel needs approximately {config.smem_bytes()} B of shared memory, " + f"over the {HOPPER_MAX_SMEM_BYTES} B Hopper limit" + ) + if config.register_total() > USABLE_REGISTER_BUDGET: + raise ValueError( + f"register budget {config.register_total()} exceeds the usable {USABLE_REGISTER_BUDGET}-register budget" + ) + + target = target.contiguous() + x_padded, weight_padded, _ = _pad_hidden(_input, weight, TILE_K) + tokens = x_padded.shape[0] + split_n, base_split_n, extra_m_pairs, num_m_tiles, num_logical_n_tiles = _resolve_split( + config, + tokens, + weight_padded.shape[0], + ) + if base_split_n < 1 or split_n > num_logical_n_tiles: + raise ValueError("split_n must be between 1 and the number of logical vocabulary tiles") + num_m_pairs = (num_m_tiles + CLUSTER_M - 1) // CLUSTER_M + if extra_m_pairs >= num_m_pairs or split_n != base_split_n + int(extra_m_pairs != 0): + raise ValueError("uneven split tuning must use split_n=base_split_n+1 for a proper subset of M pairs") + config = replace( + config, + split_n=split_n, + base_split_n=base_split_n, + extra_m_pairs=extra_m_pairs, + ) + + partial_shape = (num_m_tiles * split_n, TILE_M) + partial_max = torch.empty(partial_shape, device=_input.device, dtype=torch.float32) + partial_sum = torch.empty_like(partial_max) + partial_target = torch.empty_like(partial_max) + partial_weighted = ( + torch.empty_like(partial_max) + if return_entropy + else torch.empty( + 1, + device=_input.device, + dtype=torch.float32, + ) + ) + lse = torch.empty(tokens, device=_input.device, dtype=torch.float32) + token_loss = torch.empty(tokens, device=_input.device, dtype=torch.float32) + entropy = torch.empty( + tokens if return_entropy else 1, + device=_input.device, + dtype=torch.bfloat16, + ) + + x_cute = to_cute_tensor(x_padded.unsqueeze(-1), leading_dim=1, assumed_align=16) + weight_cute = to_cute_tensor(weight_padded.unsqueeze(-1), leading_dim=1, assumed_align=16) + target_cute = to_cute_tensor(target, assumed_align=8) + partial_max_cute = to_cute_tensor(partial_max, leading_dim=1, assumed_align=4) + partial_sum_cute = to_cute_tensor(partial_sum, leading_dim=1, assumed_align=4) + partial_target_cute = to_cute_tensor(partial_target, leading_dim=1, assumed_align=4) + partial_weighted_cute = to_cute_tensor( + partial_weighted, + leading_dim=1 if return_entropy else None, + assumed_align=4, + ) + lse_cute = to_cute_tensor(lse, assumed_align=4) + token_loss_cute = to_cute_tensor(token_loss, assumed_align=4) + entropy_cute = to_cute_tensor(entropy, assumed_align=2) + stream = _cute_stream() + inverse_temperature = Float32(1.0 / temperature) + + key = ("scaled_ce_forward_fragment", config) + compiled = _compile_cache.get(key) + if compiled is None: + compiled = cute.compile( + _ScaledCEForwardFragmentSM90(config), + x_cute, + weight_cute, + target_cute, + partial_max_cute, + partial_sum_cute, + partial_target_cute, + partial_weighted_cute, + inverse_temperature, + stream, + ) + _compile_cache[key] = compiled + compiled( + x_cute, + weight_cute, + target_cute, + partial_max_cute, + partial_sum_cute, + partial_target_cute, + partial_weighted_cute, + inverse_temperature, + stream, + ) + + finalize_key = ( + "scaled_ce_forward_fragment_finalize", + split_n, + base_split_n, + extra_m_pairs, + config.fast_math, + config.return_entropy, + ) + finalize = _compile_cache.get(finalize_key) + if finalize is None: + finalize = cute.compile( + _ScaledCEForwardFragmentFinalizeSM90( + split_n, + base_split_n, + extra_m_pairs, + config.fast_math, + config.return_entropy, + ), + partial_max_cute, + partial_sum_cute, + partial_target_cute, + partial_weighted_cute, + target_cute, + lse_cute, + token_loss_cute, + entropy_cute, + Int32(ignore_index), + stream, + ) + _compile_cache[finalize_key] = finalize + finalize( + partial_max_cute, + partial_sum_cute, + partial_target_cute, + partial_weighted_cute, + target_cute, + lse_cute, + token_loss_cute, + entropy_cute, + Int32(ignore_index), + stream, + ) + return token_loss, entropy if return_entropy else None, lse diff --git a/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_mfast_sm90.py b/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_mfast_sm90.py deleted file mode 100644 index a80a7f25..00000000 --- a/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_mfast_sm90.py +++ /dev/null @@ -1,1216 +0,0 @@ -"""SM90 fused scaled cross entropy **M-fast forward** kernel: the default -"MMA-warp issues TMA + dedicated SMEM-logits epilogue" kernel with an -**M-fast (N-outer / M-inner) persistent work order**. - -This kernel backs ``m_tiles_per_cluster > 1`` in -:mod:`liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90`; the public -``m_tiles_per_cluster`` is this module's ``slots_per_wave``. It keeps several -``M128`` online-softmax states live so that one L2-resident ``W`` panel is -shared by several M tiles. It is correct at every ``slots_per_wave`` but was -measured **slower** than the M-outer default at the representative -``M=4096, H=4096, V=131072`` shape, hence the default of 1. Everything except -:func:`to_cute_tensor` is self-contained. - -Work order -========== - -The default forward (``_fused_scaled_cross_entropy_forward_sm90.py``) is -*already* statically persistent: a cluster owns ``pid_m = cluster_idx + m_slot * -num_persistent_clusters`` and walks its slots one at a time. Its work nesting -is:: - - for m_slot in owned_slots: # M outer - for n_round in range(N/(cluster_n*tile_n)): # N inner - for k in range(H/tile_k): - WGMMA(M128 x N256 x K64) - reduce + store row statistics for this m_slot - -Softmax is a *row* reduction, so it only has to happen after **all** N tiles of -a row have been seen -- the N tiles themselves may be visited in any order. -This module therefore reverses the nesting, exactly as the dZ scheduler -(``_fused_scaled_cross_entropy_dz_sm90.py``) does for the -backward:: - - for n_round in range(N/(cluster_n*tile_n)): # N outer - for m_slot in owned_slots: # M inner - for k in range(H/tile_k): - WGMMA(M128 x N256 x K64) - for m_slot in owned_slots: # after ALL N work - reduce + store row statistics - -so the epilogue warp group carries **one online-softmax state per owned M -slot** across the whole N sweep instead of one state at a time. - -Locality hypothesis under test (**disproven** -- see "Status: not the default") ------------------------------------------------------------------------------- - -At ``M=4096, H=4096, V=131072`` with ``cluster_n=16`` and 132 SMs, 7 clusters -are resident and each owns 4-5 ``M128`` slots. Per N round a CTA touches one -``256 x 4096`` BF16 W panel = 2 MiB (32 MiB per cluster, and *all* clusters -touch the same 32 MiB in the same round). The motivating hypothesis was that in -M-outer order that panel is re-fetched from DRAM once per M slot, so M-fast -order would cut W traffic ~4x by reusing the L2-resident panel across all owned -M slots. - -The measurement says otherwise: because all resident clusters are on the *same* -N round in *both* orders, W is already fetched once and shared through L2, so -there is nothing to save; meanwhile the live X footprint grows from ~7 MiB -(1 block per cluster) to ~32 MiB (4-5 blocks per cluster), which together with -the 32 MiB W panel set overflows the 60 MiB L2 and makes X re-stream every -round. Mainloop DRAM reads go **up** 4.3% and the kernel is 2.2-3.2% slower. -Full numbers and the mechanism are at the bottom of this docstring. - -Everything else is held fixed ------------------------------ - -Identical to the accepted forward, deliberately: WG0+WG1 cooperative -``atom_layout=(2,1,1)`` WGMMA over a ``(64, 256)`` atom (``M128 x N256 x K64``), -warp 0 of WG0 issuing every TMA under a live WGMMA group, ``wait_group(1)`` in -the steady state with a full drain only at the tile tail, WG2 as a dedicated -padded-SMEM (row stride ``tile_n + 8``) single-pass online-softmax epilogue, -3 stages, ``cluster_n = 16``, **no multicast**, 208/88 register split. -``cluster_n`` sweeps are explicitly *not* repeated here: ``cluster_n=4`` was -already measured at 10.44 ms / 421 TFLOPS against 7.02 ms / 627 for -``cluster_n=16`` under matched timing. - -Implementation notes specific to the new order -============================================== - -Waves ------ - -Per-slot softmax state lives in registers, so the number of simultaneously live -slots must be a compile-time constant. Owned slots are therefore processed in -*waves* of at most ``wave_slots = min(config.slots_per_wave, ceil(num_m_tiles / -num_persistent_clusters))`` slots (a host-computed constant, 4 at the target -shape, so the target shape is a single wave):: - - for wave: # dynamic, cluster-uniform - for n_round: # dynamic - for slot in range(cur): # dynamic; cur = slots in this wave - one M128 x N256 tile - for slot in range(cur): # reduce + store - # register state: run_max[wave_slots], run_sum[wave_slots], - # run_target[wave_slots], tgt_tile[..], tgt_col[..] - -Only 1 wave is needed until a cluster owns more than ``slots_per_wave`` M -blocks, so the target shape never pays for the outer loop. - -Keeping the epilogue code single-instance ------------------------------------------ - -A naive implementation makes the slot loop ``range_constexpr`` so the register -arrays can be indexed statically -- which replicates the *entire* 256-column -softmax body ``wave_slots`` times (~4x SASS growth in WG2, straight into the -I-cache). Instead the slot loop stays **dynamic** and the heavy body computes a -purely tile-local ``(tile_max, tile_sum, tile_target)`` in plain registers; only -the ~6-instruction fold into the per-slot state is wrapped in a -``range_constexpr`` + ``if slot == j`` select:: - - for j in range_constexpr(wave_slots): - if slot == j: - new = fmax(run_max[j], tile_max) - run_sum[j] = run_sum[j]*exp2((run_max[j]-new)*log2e) - + tile_sum *exp2((tile_max -new)*log2e) - run_max[j] = new - run_target[j] += tile_target - -The same trick reads ``tgt_tile/tgt_col`` before the body and ``run_*`` before -the cluster reduction. Net effect: register-resident per-slot state with the -epilogue body emitted exactly once, i.e. the same SASS footprint as the accepted -kernel. - -The two-sentinel masking trick is preserved: the tile-local max starts at -``NEG_INF_F32 = -1e38`` while out-of-vocabulary columns are ``MASK_F32 = --3e38``, so a fully-masked tile contributes ``exp2((-3e38 + 1e38) * log2 e) = 0`` -rather than 1, and folding a ``(max=-1e38, sum=0)`` tile into any state is a -no-op (``0 * exp2(...) == 0``, never ``0 * inf``). - -TMA stream ----------- - -The load stream is a **single linear stream over the whole persistent kernel**: -one prologue of ``stages`` loads at kernel entry, one refill per consumer -release, one ``producer_tail`` at exit. The stream cursor is -``(k, slot, n_round, wave)`` advanced k-fastest, mirroring the consumer nesting -exactly, so ``load_index == release_index`` holds at every refill with no -modular arithmetic -- and unlike the accepted kernel there is no per-M-tile -prologue/drain at all, because M tiles are no longer loop boundaries. -``wait_group(1)`` and the tail-only ``wait_group(0)`` are unchanged. - -Ragged tails ------------- - -``num_real_slots = ceil((num_m_tiles - cluster_idx) / num_persistent_clusters)`` -depends only on ``blockIdx.z``, which is identical for all 16 CTAs of a cluster, -so every loop bound and every predicate above is **cluster-uniform**; named -barriers (``MMA_BARRIER_ID`` over 256 threads, ``EPI_BARRIER_ID`` over 128), -``__syncthreads``, the cluster barriers and both pipelines therefore stay -balanced even when the last cluster owns fewer M slots than the others. All 16 -ranks of a cluster are always on the same M slot for a given tile. - -Status: not the default ------------------------ -The locality hypothesis is **disproven**. Matched, interleaved measurements -put this order ~3 % behind the M-outer forward at the representative shape -(independently reconfirmed at ~3.8 %), because keeping several M tiles live per -cluster shrinks the per-tile ``W`` reuse window without buying back enough -``X`` reuse. It is retained as the implementation of the public -``m_tiles_per_cluster > 1`` scheduling knob (``slots_per_wave``); the default -``m_tiles_per_cluster == 1`` dispatches to the M-outer forward instead. - -""" - -import inspect - -from dataclasses import dataclass - -import cuda.bindings.driver as cuda -import cutlass -import cutlass.cute as cute -import torch - -from cutlass import BFloat16 -from cutlass import Float16 -from cutlass import Float32 -from cutlass import Int32 -from cutlass import cutlass_dsl -from cutlass import pipeline -from cutlass._mlir.dialects import llvm -from cutlass._mlir.dialects import nvvm -from cutlass.cutlass_dsl import T -from cutlass.cutlass_dsl import dsl_user_op -from cutlass.pipeline.helpers import pipeline_init_arrive -from cutlass.pipeline.helpers import pipeline_init_wait -from cutlass.utils import LayoutEnum -from cutlass.utils import hopper_helpers - -from liger_kernel.ops.cutedsl.ops.utils import to_cute_tensor - -THREADS_PER_CTA = 384 -NUM_MMA_WARP_GROUPS = 2 -WARPS_PER_WARP_GROUP = 4 - -# Named barriers. 0 is reserved for __syncthreads(). -MMA_BARRIER_ID = 4 # both MMA warp groups, 256 threads -EPI_BARRIER_ID = 5 # epilogue warp group, 128 threads - -LOG2_E = 1.4426950408889634 -LN2 = 0.6931471805599453 -NEG_INF_F32 = -1.0e38 -MASK_F32 = -3.0e38 -HOPPER_MAX_SMEM_BYTES = 227 * 1024 - -_compile_cache = {} -_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)) - - -@dsl_user_op -def _map_dsmem_addr(local_ptr, peer_rank, *, loc=None, ip=None) -> Int32: - """Map a local SMEM pointer to the same object in a peer CTA's SMEM. - - ``llvm.nvvm.mapa`` ICEs this toolchain ("only supported when pointer size is - >= 64 bits"), so the PTX is emitted directly, exactly as - ``cutlass.cute.arch.map_dsmem_ptr`` does. - """ - addr = llvm.ptrtoint(T.i32(), local_ptr.llvm_ptr, loc=loc, ip=ip) - return Int32( - llvm.inline_asm( - T.i32(), - [addr, Int32(peer_rank).ir_value(loc=loc, ip=ip)], - "mapa.shared::cluster.u32 $0, $1, $2;", - "=r,r,r", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) - - -def _asm(res_ty, args, tmpl, cons, side=True, *, loc=None, ip=None): - return llvm.inline_asm( - res_ty, - args, - tmpl, - cons, - has_side_effects=side, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def _dsmem_load_f32(addr, *, loc=None, ip=None) -> Float32: - return Float32( - _asm( - T.f32(), - [Int32(addr).ir_value(loc=loc, ip=ip)], - "ld.shared::cluster.f32 $0, [$1];", - "=f,r", - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _dsmem_atomic_add_f32(addr, value, *, loc=None, ip=None) -> Float32: - return Float32( - _asm( - T.f32(), - [ - Int32(addr).ir_value(loc=loc, ip=ip), - Float32(value).ir_value(loc=loc, ip=ip), - ], - "atom.shared::cluster.add.f32 $0, [$1], $2;", - "=f,r,f", - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _dsmem_atomic_max_s32(addr, value, *, loc=None, ip=None) -> Int32: - return Int32( - _asm( - T.i32(), - [ - Int32(addr).ir_value(loc=loc, ip=ip), - Int32(value).ir_value(loc=loc, ip=ip), - ], - "atom.shared::cluster.max.s32 $0, [$1], $2;", - "=r,r,r", - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _dsmem_atomic_min_u32(addr, value, *, loc=None, ip=None) -> Int32: - return Int32( - _asm( - T.i32(), - [ - Int32(addr).ir_value(loc=loc, ip=ip), - Int32(value).ir_value(loc=loc, ip=ip), - ], - "atom.shared::cluster.min.u32 $0, [$1], $2;", - "=r,r,r", - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _dsmem_atomic_fmax(addr, value, *, loc=None, ip=None) -> Float32: - """FP32 max via the monotonic-bit-pattern trick (``red.max.f32`` has no PTX).""" - bits = llvm.bitcast(T.i32(), Float32(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) - old = cutlass_dsl.if_generate( - Int32(bits) < 0, - lambda: _dsmem_atomic_min_u32(addr, Int32(bits), loc=loc, ip=ip), - lambda: _dsmem_atomic_max_s32(addr, Int32(bits), loc=loc, ip=ip), - [], - [Int32], - loc=loc, - ip=ip, - ) - return Float32(llvm.bitcast(T.f32(), old.ir_value(loc=loc, ip=ip), 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_n): - key = (torch.cuda.current_device(), cluster_n) - value = _max_active_clusters_cache.get(key) - if value is None: - torch.cuda.current_stream() - value = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_n) - _max_active_clusters_cache[key] = value - return value - - -@dataclass(frozen=True) -class ScaledCEForwardMFastConfig: - """Tuning knobs for the M-fast (N-outer / M-inner) forward. - - Everything except ``slots_per_wave`` matches the accepted forward kernel's - defaults and should not be swept here -- this experiment is about the work - order only. - - ``slots_per_wave`` caps how many owned ``M128`` slots keep a live - online-softmax state (and therefore how many M slots share one L2-resident W - panel). The kernel is compiled with ``wave_slots = min(slots_per_wave, - ceil(num_m_tiles / num_persistent_clusters))``, so it never costs registers - the shape cannot use. ``slots_per_wave=1`` reproduces the accepted - M-outer order exactly and is the in-module A/B control. - - Register budget: ``2 * 128 * mma_registers + 128 * epi_registers <= 64512`` - (requesting the full 65536-register file deadlocks ``setmaxnreg.inc``). - """ - - tile_m: int = 128 - tile_n: int = 256 - tile_k: int = 64 - stages: int = 3 - cluster_n: int = 16 - fast_math: bool = True - logit_dtype: str = "f16" - logit_pad: int = 8 - chunk_n: int = 32 - mma_registers: int = 208 - epi_registers: int = 88 - slots_per_wave: int = 8 - return_entropy: bool = False - - @property - def logit_stride(self): - return self.tile_n + self.logit_pad - - def smem_bytes(self): - ab = (self.tile_m + self.tile_n) * self.tile_k * 2 * self.stages - logits = self.tile_m * self.logit_stride * 2 - reduction = (3 + int(self.return_entropy)) * self.tile_m * 4 - barriers = 16 * self.stages + 16 - return ab + logits + reduction + barriers + 3 * 1024 - - def register_total(self): - return 2 * 128 * self.mma_registers + 128 * self.epi_registers - - -class _ScaledCEForwardMFastSM90: - """M-fast persistent forward: N-outer / M-inner over each cluster's slots.""" - - def __init__(self, config: ScaledCEForwardMFastConfig, wave_slots: int): - self.tile_m = config.tile_m - self.tile_n = config.tile_n - self.tile_k = config.tile_k - self.stages = config.stages - self.cluster_n = config.cluster_n - self.fast_math = config.fast_math - self.logit_pad = config.logit_pad - self.logit_stride = config.tile_n + config.logit_pad - self.chunk_n = config.chunk_n - self.mma_registers = config.mma_registers - self.epi_registers = config.epi_registers - self.return_entropy = config.return_entropy - self.logit_dtype = Float16 if config.logit_dtype == "f16" else BFloat16 - self.wave_slots = int(wave_slots) - self.tile_shape = (self.tile_m, self.tile_n, self.tile_k) - self.buffer_align = 1024 - self.mma_barrier = pipeline.NamedBarrier( - barrier_id=MMA_BARRIER_ID, - num_threads=NUM_MMA_WARP_GROUPS * 128, - ) - self.epi_barrier = pipeline.NamedBarrier( - barrier_id=EPI_BARRIER_ID, - num_threads=128, - ) - - @staticmethod - def _make_tma(tensor, smem_layout_staged, smem_tile): - smem_layout = cute.slice_(smem_layout_staged, (None, None, 0)) - return cute.nvgpu.cpasync.make_tiled_tma_atom( - cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(), - tensor, - smem_layout, - smem_tile, - num_multicast=1, - ) - - @cute.jit - def __call__( - self, - x: cute.Tensor, - weight: cute.Tensor, - target: cute.Tensor, - lse: cute.Tensor, - token_loss: cute.Tensor, - entropy: cute.Tensor, - inverse_temperature: Float32, - ignore_index: Int32, - max_active_clusters: cutlass.Constexpr, - stream: cuda.CUstream, - ): - x_layout = LayoutEnum.from_tensor(x) - weight_layout = LayoutEnum.from_tensor(weight) - tiled_mma = hopper_helpers.make_trivial_tiled_mma( - BFloat16, - BFloat16, - x_layout.sm90_mma_major_mode(), - weight_layout.sm90_mma_major_mode(), - Float32, - (NUM_MMA_WARP_GROUPS, 1, 1), - (64, self.tile_n), - ) - a_smem_layout = hopper_helpers.make_smem_layout_a( - x_layout, - self.tile_shape, - BFloat16, - self.stages, - ) - b_smem_layout = hopper_helpers.make_smem_layout_b( - weight_layout, - self.tile_shape, - BFloat16, - self.stages, - ) - tma_x, tma_tensor_x = self._make_tma(x, a_smem_layout, (self.tile_m, self.tile_k)) - tma_w, tma_tensor_w = self._make_tma(weight, b_smem_layout, (self.tile_n, self.tile_k)) - - logit_dtype = self.logit_dtype - - @cute.struct - class SharedStorage: - pipeline: cute.struct.MemRange[cutlass.Int64, self.stages * 2] - logits_pipeline: cute.struct.MemRange[cutlass.Int64, 2] - sX: cute.struct.Align[ - cute.struct.MemRange[BFloat16, cute.cosize(a_smem_layout)], - self.buffer_align, - ] - sW: cute.struct.Align[ - cute.struct.MemRange[BFloat16, cute.cosize(b_smem_layout)], - self.buffer_align, - ] - sLogits: cute.struct.Align[ - cute.struct.MemRange[ - logit_dtype, - self.tile_m * self.logit_stride, - ], - self.buffer_align, - ] - cluster_max: cute.struct.Align[cute.struct.MemRange[Float32, self.tile_m], 16] - cluster_sum: cute.struct.Align[cute.struct.MemRange[Float32, self.tile_m], 16] - cluster_target: cute.struct.Align[cute.struct.MemRange[Float32, self.tile_m], 16] - cluster_weighted: cute.struct.Align[ - cute.struct.MemRange[Float32, self.tile_m if self.return_entropy else 1], - 16, - ] - - self.shared_storage = SharedStorage - cluster_shape = (1, self.cluster_n, 1) - num_m_tiles = cute.ceil_div(x.shape[0], self.tile_m) - num_persistent_clusters = cutlass.min(num_m_tiles, max_active_clusters) - grid = (1, self.cluster_n, num_persistent_clusters) - self.kernel( - tma_x, - tma_tensor_x, - tma_w, - tma_tensor_w, - target, - lse, - token_loss, - entropy, - inverse_temperature, - tiled_mma, - a_smem_layout, - b_smem_layout, - ignore_index, - ).launch( - grid=grid, - block=(THREADS_PER_CTA, 1, 1), - cluster=cluster_shape, - min_blocks_per_mp=1, - stream=stream, - ) - - @cute.kernel - def kernel( - self, - tma_x: cute.CopyAtom, - x: cute.Tensor, - tma_w: cute.CopyAtom, - weight: cute.Tensor, - target: cute.Tensor, - lse: cute.Tensor, - token_loss: cute.Tensor, - entropy: cute.Tensor, - inverse_temperature: Float32, - tiled_mma: cute.TiledMma, - a_smem_layout: cute.ComposedLayout, - b_smem_layout: cute.ComposedLayout, - ignore_index: Int32, - ): - tid, _, _ = cute.arch.thread_idx() - warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - warp_group = cute.arch.make_warp_uniform(tid // 128) - lane = tid % 32 - local_warp = (tid % 128) // 32 - local_tid = tid % 128 - cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) - _, _, cluster_idx = cute.arch.block_idx() - _, _, num_persistent_clusters = cute.arch.grid_dim() - num_m_tiles = cute.ceil_div(target.shape[0], self.tile_m) - - if warp == 0: - cute.nvgpu.cpasync.prefetch_descriptor(tma_x) - cute.nvgpu.cpasync.prefetch_descriptor(tma_w) - - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(self.shared_storage) - sX = storage.sX.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner) - sW = storage.sW.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) - sLogits = storage.sLogits.get_tensor( - cute.make_layout( - (self.tile_m, self.tile_n), - stride=(self.logit_stride, 1), - ) - ) - cluster_max = storage.cluster_max.get_tensor(cute.make_layout(self.tile_m)) - cluster_sum = storage.cluster_sum.get_tensor(cute.make_layout(self.tile_m)) - cluster_target = storage.cluster_target.get_tensor(cute.make_layout(self.tile_m)) - cluster_weighted = storage.cluster_weighted.get_tensor( - cute.make_layout(self.tile_m if cutlass.const_expr(self.return_entropy) else 1) - ) - - cta_layout = cute.make_layout((1, self.cluster_n, 1)) - cluster_coord = cta_layout.get_flat_coord(cluster_rank) - - gX = cute.local_tile( - x, - cute.slice_(self.tile_shape, (None, 0, None)), - (None, None, None), - ) - gW = cute.local_tile( - weight, - cute.slice_(self.tile_shape, (0, None, None)), - (None, None, None), - ) - tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( - tma_x, - 0, - cute.make_layout(1), - cute.group_modes(sX, 0, 2), - cute.group_modes(gX, 0, 2), - ) - w_cta_layout = cute.make_layout(cute.slice_(cta_layout, (None, 0, 0)).shape) - tWsW, tWgW = cute.nvgpu.cpasync.tma_partition( - tma_w, - cluster_coord[0], - w_cta_layout, - cute.group_modes(sW, 0, 2), - cute.group_modes(gW, 0, 2), - ) - - tma_bytes = cute.size_in_bytes(BFloat16, cute.slice_(a_smem_layout, (None, None, 0))) + cute.size_in_bytes( - BFloat16, cute.slice_(b_smem_layout, (None, None, 0)) - ) - - mainloop = pipeline.PipelineTmaAsync.create( - barrier_storage=storage.pipeline.data_ptr(), - num_stages=self.stages, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, - NUM_MMA_WARP_GROUPS * WARPS_PER_WARP_GROUP, - ), - tx_count=tma_bytes, - cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), - defer_sync=True, - ) - logits_pipeline = pipeline.PipelineAsync.create( - barrier_storage=storage.logits_pipeline.data_ptr(), - num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - defer_sync=True, - ) - pipeline_init_arrive(cluster_shape_mn=(1, self.cluster_n), is_relaxed=True) - pipeline_init_wait(cluster_shape_mn=(1, self.cluster_n)) - - num_k_tiles = cute.size(gX, mode=[3]) - num_k_blocks = self.tile_k // 16 - num_pair_tiles = cute.ceil_div(weight.shape[0], self.cluster_n * self.tile_n) - - # ---- M-slot ownership (cluster-uniform: depends only on blockIdx.z) --- - num_real_slots = cute.ceil_div(num_m_tiles - cluster_idx, num_persistent_clusters) - num_waves = cute.ceil_div(num_real_slots, self.wave_slots) - - root_max_addr = _map_dsmem_addr(storage.cluster_max.data_ptr(), 0) - root_sum_addr = _map_dsmem_addr(storage.cluster_sum.data_ptr(), 0) - root_target_addr = _map_dsmem_addr(storage.cluster_target.data_ptr(), 0) - root_weighted_addr = _map_dsmem_addr(storage.cluster_weighted.data_ptr(), 0) - - if warp_group < NUM_MMA_WARP_GROUPS: - # ================= MMA warp groups (WG0 + WG1) ================= - cute.arch.setmaxregister_increase(self.mma_registers) - - mma_wg_layout = cute.make_layout(NUM_MMA_WARP_GROUPS, stride=128) - thr_mma = tiled_mma.get_slice(mma_wg_layout(Int32(warp_group))) - tCrX = tiled_mma.make_fragment_A(thr_mma.partition_A(sX)) - tCrW = tiled_mma.make_fragment_B(thr_mma.partition_B(sW)) - accum = cute.make_rmem_tensor( - thr_mma.partition_shape_C((self.tile_m, self.tile_n)), - Float32, - ) - n_accum = cute.size(accum) - row_base_in_tile = warp_group * 64 + local_warp * 16 + lane // 4 - col_base_in_tile = (lane % 4) * 2 - - read_index = Int32(0) - read_phase = Int32(0) - rel_index = Int32(0) - rel_phase = Int32(0) - load_index = Int32(0) - load_phase = Int32(1) - logit_w_index = Int32(0) - logit_w_phase = Int32(1) - - # TMA cursor: (k, slot-in-wave, n_round, wave), advanced k-fastest. - load_k = Int32(0) - load_j = Int32(0) - load_round = Int32(0) - load_wave = Int32(0) - load_wave_slots = cutlass.min(Int32(self.wave_slots), num_real_slots) - - cute.arch.sync_threads() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - # ---- prologue: S complete TMA stages before any WGMMA ---------- - if warp == 0: - total_loads = Int32(num_k_tiles) * Int32(num_pair_tiles) * num_real_slots - prologue_stages = cutlass.min(Int32(self.stages), total_loads) - for _ in range(prologue_stages): - pstate = pipeline.PipelineState(self.stages, Int32(0), load_index, load_phase) - mainloop.producer_acquire(pstate) - bar = mainloop.producer_get_barrier(pstate) - load_pid_m = cluster_idx + (load_wave * self.wave_slots + load_j) * num_persistent_clusters - tXgX_m = tXgX[(None, load_pid_m, None, 0)] - cute.copy( - tma_x, - tXgX_m[(None, load_k)], - tXsX[(None, load_index)], - tma_bar_ptr=bar, - mcast_mask=0, - ) - tWgW_n = tWgW[(None, load_round * self.cluster_n + cluster_rank, None, 0)] - cute.copy( - tma_w, - tWgW_n[(None, load_k)], - tWsW[(None, load_index)], - tma_bar_ptr=bar, - ) - mainloop.producer_commit(pstate) - pstate.advance() - load_index = pstate.index - load_phase = pstate.phase - load_k += 1 - if load_k == num_k_tiles: - load_k = Int32(0) - load_j += 1 - if load_j == load_wave_slots: - load_j = Int32(0) - load_round += 1 - if load_round == num_pair_tiles: - load_round = Int32(0) - load_wave += 1 - load_wave_slots = cutlass.min( - Int32(self.wave_slots), - num_real_slots - load_wave * self.wave_slots, - ) - if load_wave_slots < 1: - load_wave_slots = Int32(1) - - for wave in range(num_waves): - cur_slots = cutlass.min( - Int32(self.wave_slots), - num_real_slots - wave * self.wave_slots, - ) - - for _n_round in range(num_pair_tiles): - for _slot in range(cur_slots): - # ------------- one M128 x N256 output tile ---------- - accum.fill(0.0) - tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, True) - - read_state = pipeline.PipelineState(self.stages, Int32(0), read_index, read_phase) - mainloop.consumer_wait(read_state) - cute.nvgpu.warpgroup.fence() - for k_block in cutlass.range_constexpr(num_k_blocks): - coord = (None, None, k_block, read_index) - cute.gemm(tiled_mma, accum, tCrX[coord], tCrW[coord], accum) - cute.nvgpu.warpgroup.commit_group() - read_state.advance() - read_index = read_state.index - read_phase = read_state.phase - - for _ in range(1, num_k_tiles): - read_state = pipeline.PipelineState(self.stages, Int32(0), read_index, read_phase) - mainloop.consumer_wait(read_state) - cute.nvgpu.warpgroup.fence() - for k_block in cutlass.range_constexpr(num_k_blocks): - coord = (None, None, k_block, read_index) - cute.gemm(tiled_mma, accum, tCrX[coord], tCrW[coord], accum) - cute.nvgpu.warpgroup.commit_group() - # Retire only the previous group; the group just - # issued stays IN FLIGHT across release + refill. - cute.nvgpu.warpgroup.wait_group(1) - rstate = pipeline.PipelineState(self.stages, Int32(0), rel_index, rel_phase) - mainloop.consumer_release(rstate) - rstate.advance() - rel_index = rstate.index - rel_phase = rstate.phase - if warp == 0: - if load_wave < num_waves: - pstate = pipeline.PipelineState(self.stages, Int32(0), load_index, load_phase) - mainloop.producer_acquire(pstate) - bar = mainloop.producer_get_barrier(pstate) - load_pid_m = ( - cluster_idx + (load_wave * self.wave_slots + load_j) * num_persistent_clusters - ) - tXgX_m = tXgX[(None, load_pid_m, None, 0)] - cute.copy( - tma_x, - tXgX_m[(None, load_k)], - tXsX[(None, load_index)], - tma_bar_ptr=bar, - mcast_mask=0, - ) - tWgW_n = tWgW[(None, load_round * self.cluster_n + cluster_rank, None, 0)] - cute.copy( - tma_w, - tWgW_n[(None, load_k)], - tWsW[(None, load_index)], - tma_bar_ptr=bar, - ) - mainloop.producer_commit(pstate) - pstate.advance() - load_index = pstate.index - load_phase = pstate.phase - load_k += 1 - if load_k == num_k_tiles: - load_k = Int32(0) - load_j += 1 - if load_j == load_wave_slots: - load_j = Int32(0) - load_round += 1 - if load_round == num_pair_tiles: - load_round = Int32(0) - load_wave += 1 - load_wave_slots = cutlass.min( - Int32(self.wave_slots), - num_real_slots - load_wave * self.wave_slots, - ) - if load_wave_slots < 1: - load_wave_slots = Int32(1) - read_state.advance() - read_index = read_state.index - read_phase = read_state.phase - - # ---- tile tail: the only WGMMA drain --------------- - cute.nvgpu.warpgroup.wait_group(0) - rstate = pipeline.PipelineState(self.stages, Int32(0), rel_index, rel_phase) - mainloop.consumer_release(rstate) - rstate.advance() - rel_index = rstate.index - rel_phase = rstate.phase - if warp == 0: - if load_wave < num_waves: - pstate = pipeline.PipelineState(self.stages, Int32(0), load_index, load_phase) - mainloop.producer_acquire(pstate) - bar = mainloop.producer_get_barrier(pstate) - load_pid_m = ( - cluster_idx + (load_wave * self.wave_slots + load_j) * num_persistent_clusters - ) - tXgX_m = tXgX[(None, load_pid_m, None, 0)] - cute.copy( - tma_x, - tXgX_m[(None, load_k)], - tXsX[(None, load_index)], - tma_bar_ptr=bar, - mcast_mask=0, - ) - tWgW_n = tWgW[(None, load_round * self.cluster_n + cluster_rank, None, 0)] - cute.copy( - tma_w, - tWgW_n[(None, load_k)], - tWsW[(None, load_index)], - tma_bar_ptr=bar, - ) - mainloop.producer_commit(pstate) - pstate.advance() - load_index = pstate.index - load_phase = pstate.phase - load_k += 1 - if load_k == num_k_tiles: - load_k = Int32(0) - load_j += 1 - if load_j == load_wave_slots: - load_j = Int32(0) - load_round += 1 - if load_round == num_pair_tiles: - load_round = Int32(0) - load_wave += 1 - load_wave_slots = cutlass.min( - Int32(self.wave_slots), - num_real_slots - load_wave * self.wave_slots, - ) - if load_wave_slots < 1: - load_wave_slots = Int32(1) - - wstate = pipeline.PipelineState(1, Int32(0), logit_w_index, logit_w_phase) - logits_pipeline.producer_acquire(wstate) - for i in cutlass.range_constexpr(n_accum): - row = row_base_in_tile + ((i % 4) // 2) * 8 - col = col_base_in_tile + (i // 4) * 8 + (i % 2) - sLogits[row, col] = self.logit_dtype(accum[i]) - cute.arch.fence_proxy("async.shared", space="cta") - cute.arch.fence_acq_rel_cta() - self.mma_barrier.arrive_and_wait() - if tid == 0: - logits_pipeline.producer_commit(wstate) - wstate.advance() - logit_w_index = wstate.index - logit_w_phase = wstate.phase - - # ---- per-slot cluster reduction: barrier participation only -- - for _slot in range(cur_slots): - cute.arch.sync_threads() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - if warp == 0: - ptail = pipeline.PipelineState(self.stages, Int32(0), load_index, load_phase) - mainloop.producer_tail(ptail) - else: - # ==================== epilogue warp group ====================== - cute.arch.setmaxregister_decrease(self.epi_registers) - - row = local_tid - chunk = cute.make_rmem_tensor((self.chunk_n,), Float32) - # Per-owned-M-slot online softmax state, register resident. - run_max = cute.make_rmem_tensor((self.wave_slots,), Float32) - run_sum = cute.make_rmem_tensor((self.wave_slots,), Float32) - run_target = cute.make_rmem_tensor((self.wave_slots,), Float32) - run_weighted = cute.make_rmem_tensor( - (self.wave_slots if self.return_entropy else 1,), - Float32, - ) - tgt_tile = cute.make_rmem_tensor((self.wave_slots,), Int32) - tgt_col = cute.make_rmem_tensor((self.wave_slots,), Int32) - logit_r_index = Int32(0) - logit_r_phase = Int32(0) - - cute.arch.sync_threads() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - for wave in range(num_waves): - cur_slots = cutlass.min( - Int32(self.wave_slots), - num_real_slots - wave * self.wave_slots, - ) - slot_base = wave * self.wave_slots - - for j in cutlass.range_constexpr(self.wave_slots): - run_max[j] = Float32(NEG_INF_F32) - run_sum[j] = Float32(0.0) - run_target[j] = Float32(0.0) - if cutlass.const_expr(self.return_entropy): - run_weighted[j] = Float32(0.0) - tgt_tile[j] = Int32(-1) - tgt_col[j] = Int32(-1) - if cur_slots > j: - grow = (cluster_idx + (slot_base + j) * num_persistent_clusters) * self.tile_m + row - if grow < target.shape[0]: - tcol = Int32(target[grow]) - if tcol >= 0: - tgt_tile[j] = tcol // self.tile_n - tgt_col[j] = tcol % self.tile_n - - for n_round in range(num_pair_tiles): - n_tile = n_round * self.cluster_n + cluster_rank - valid_cols = Int32(weight.shape[0]) - n_tile * self.tile_n - - for slot in range(cur_slots): - # Pull this slot's target coordinates into plain - # registers so the heavy body below stays a single - # SASS instance (see module docstring). - cur_tgt_tile = Int32(-1) - cur_tgt_col = Int32(-1) - for j in cutlass.range_constexpr(self.wave_slots): - if slot == j: - cur_tgt_tile = tgt_tile[j] - cur_tgt_col = tgt_col[j] - - rstate = pipeline.PipelineState(1, Int32(0), logit_r_index, logit_r_phase) - logits_pipeline.consumer_wait(rstate) - cute.arch.fence_acq_rel_cta() - - tile_max = Float32(NEG_INF_F32) - tile_sum = Float32(0.0) - tile_target = Float32(0.0) - tile_weighted = Float32(0.0) - for ci in cutlass.range_constexpr(self.tile_n // self.chunk_n): - c0 = ci * self.chunk_n - chunk_max = Float32(MASK_F32) - if c0 + self.chunk_n <= valid_cols: - for jj in cutlass.range_constexpr(self.chunk_n): - v = Float32(sLogits[row, c0 + jj]) * inverse_temperature - chunk[jj] = v - chunk_max = _fmax(chunk_max, v) - else: - for jj in cutlass.range_constexpr(self.chunk_n): - v = Float32(sLogits[row, c0 + jj]) * inverse_temperature - if c0 + jj >= valid_cols: - v = Float32(MASK_F32) - chunk[jj] = v - chunk_max = _fmax(chunk_max, v) - new_max = _fmax(tile_max, chunk_max) - acc = Float32(0.0) - weighted = Float32(0.0) - for jj in cutlass.range_constexpr(self.chunk_n): - exp_value = cute.math.exp2( - (chunk[jj] - new_max) * LOG2_E, - fastmath=self.fast_math, - ) - acc += exp_value - if cutlass.const_expr(self.return_entropy): - weighted += exp_value * chunk[jj] - old_scale = cute.math.exp2( - (tile_max - new_max) * LOG2_E, - fastmath=self.fast_math, - ) - tile_sum = tile_sum * old_scale + acc - if cutlass.const_expr(self.return_entropy): - tile_weighted = tile_weighted * old_scale + weighted - tile_max = new_max - - if cur_tgt_tile == n_tile: - tile_target = Float32(sLogits[row, cur_tgt_col]) * inverse_temperature - - self.epi_barrier.arrive_and_wait() - if local_tid == 0: - logits_pipeline.consumer_release(rstate) - rstate.advance() - logit_r_index = rstate.index - logit_r_phase = rstate.phase - - # Fold the tile-local result into this slot's state. - for j in cutlass.range_constexpr(self.wave_slots): - if slot == j: - folded = _fmax(run_max[j], tile_max) - run_sum[j] = run_sum[j] * cute.math.exp2( - (run_max[j] - folded) * LOG2_E, - fastmath=self.fast_math, - ) + tile_sum * cute.math.exp2( - (tile_max - folded) * LOG2_E, - fastmath=self.fast_math, - ) - if cutlass.const_expr(self.return_entropy): - run_weighted[j] = run_weighted[j] * cute.math.exp2( - (run_max[j] - folded) * LOG2_E, - fastmath=self.fast_math, - ) + tile_weighted * cute.math.exp2( - (tile_max - folded) * LOG2_E, - fastmath=self.fast_math, - ) - run_max[j] = folded - run_target[j] = run_target[j] + tile_target - - # ---- all N work for this wave is done: reduce every slot ---- - for slot in range(cur_slots): - cur_max = Float32(NEG_INF_F32) - cur_sum = Float32(0.0) - cur_target = Float32(0.0) - cur_weighted = Float32(0.0) - for j in cutlass.range_constexpr(self.wave_slots): - if slot == j: - cur_max = run_max[j] - cur_sum = run_sum[j] - cur_target = run_target[j] - if cutlass.const_expr(self.return_entropy): - cur_weighted = run_weighted[j] - - global_row = (cluster_idx + (slot_base + slot) * num_persistent_clusters) * self.tile_m + row - - if cluster_rank == 0: - cluster_max[row] = Float32(NEG_INF_F32) - cluster_sum[row] = Float32(0.0) - cluster_target[row] = Float32(0.0) - if cutlass.const_expr(self.return_entropy): - cluster_weighted[row] = Float32(0.0) - cute.arch.sync_threads() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - _dsmem_atomic_fmax(root_max_addr + row * 4, cur_max) - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - global_max = _dsmem_load_f32(root_max_addr + row * 4) - scaled_sum = cur_sum * cute.math.exp2( - (cur_max - global_max) * LOG2_E, - fastmath=self.fast_math, - ) - _dsmem_atomic_add_f32(root_sum_addr + row * 4, scaled_sum) - _dsmem_atomic_add_f32(root_target_addr + row * 4, cur_target) - if cutlass.const_expr(self.return_entropy): - scaled_weighted = cur_weighted * cute.math.exp2( - (cur_max - global_max) * LOG2_E, - fastmath=self.fast_math, - ) - _dsmem_atomic_add_f32(root_weighted_addr + row * 4, scaled_weighted) - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - if cluster_rank == 0: - if global_row < target.shape[0]: - row_lse = cluster_max[row] + cute.math.log2(cluster_sum[row], fastmath=self.fast_math) * LN2 - lse[global_row] = row_lse - if target[global_row] == ignore_index: - token_loss[global_row] = Float32(0.0) - if cutlass.const_expr(self.return_entropy): - entropy[global_row] = BFloat16(0.0) - else: - token_loss[global_row] = row_lse - cluster_target[row] - if cutlass.const_expr(self.return_entropy): - entropy[global_row] = BFloat16(row_lse - cluster_weighted[row] / cluster_sum[row]) - - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - -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 M-fast 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 M-fast 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, - ) - - -def scaled_ce_forward_mfast( - _input, - weight, - target, - temperature=1.0, - ignore_index=-100, - return_entropy=False, - config: ScaledCEForwardMFastConfig = None, -): - """M-fast (N-outer / M-inner) forward. Returns per-token ``(nll[M], lse[M])``.""" - _validate(_input, weight, target) - if config is None: - config = ScaledCEForwardMFastConfig() - if config.return_entropy != return_entropy: - config = ScaledCEForwardMFastConfig(**{**config.__dict__, "return_entropy": return_entropy}) - if config.tile_m != 128: - raise ValueError("tile_m must be 128 (one epilogue thread per row)") - if config.tile_n % 8 != 0 or config.tile_n % config.chunk_n != 0: - raise ValueError("tile_n must be a multiple of 8 and of chunk_n") - if config.slots_per_wave < 1: - raise ValueError("slots_per_wave must be >= 1") - if config.smem_bytes() > HOPPER_MAX_SMEM_BYTES: - raise ValueError( - f"config needs ~{config.smem_bytes()} B of shared memory, " - f"which exceeds the {HOPPER_MAX_SMEM_BYTES} B Hopper opt-in limit" - ) - if config.register_total() > 64512: - raise ValueError(f"register budget {config.register_total()} exceeds the usable 64512-register budget") - - target = target.contiguous() - x_padded, weight_padded, _ = _pad_hidden(_input, weight, config.tile_k) - bt = x_padded.shape[0] - - lse = torch.empty(bt, device=_input.device, dtype=torch.float32) - token_loss = torch.empty(bt, device=_input.device, dtype=torch.float32) - entropy = torch.empty(bt if return_entropy else 1, device=_input.device, dtype=torch.bfloat16) - x_cute = to_cute_tensor(x_padded.unsqueeze(-1), leading_dim=1, assumed_align=16) - weight_cute = to_cute_tensor(weight_padded.unsqueeze(-1), leading_dim=1, assumed_align=16) - target_cute = to_cute_tensor(target, assumed_align=8) - lse_cute = to_cute_tensor(lse, assumed_align=4) - token_loss_cute = to_cute_tensor(token_loss, assumed_align=4) - entropy_cute = to_cute_tensor(entropy, assumed_align=2) - stream = _cute_stream() - inverse_temperature = Float32(1.0 / temperature) - max_active_clusters = _max_active_clusters(config.cluster_n) - - # ``wave_slots`` must be a compile-time constant (it sizes the per-slot - # register arrays), so it is derived on the host from the launch geometry. - num_m_tiles = (bt + config.tile_m - 1) // config.tile_m - num_persistent_clusters = min(num_m_tiles, max_active_clusters) - max_m_slots = (num_m_tiles + num_persistent_clusters - 1) // num_persistent_clusters - wave_slots = max(1, min(config.slots_per_wave, max_m_slots)) - - key = ("scaled_ce_forward_mfast", config, max_active_clusters, wave_slots) - compiled = _compile_cache.get(key) - if compiled is None: - compiled = cute.compile( - _ScaledCEForwardMFastSM90(config, wave_slots), - x_cute, - weight_cute, - target_cute, - lse_cute, - token_loss_cute, - entropy_cute, - inverse_temperature, - Int32(ignore_index), - max_active_clusters, - stream, - ) - _compile_cache[key] = compiled - compiled( - x_cute, - weight_cute, - target_cute, - lse_cute, - token_loss_cute, - entropy_cute, - inverse_temperature, - Int32(ignore_index), - stream, - ) - - return token_loss, entropy if return_entropy else None, lse diff --git a/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_sm90.py b/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_sm90.py deleted file mode 100644 index d467a0c2..00000000 --- a/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_forward_sm90.py +++ /dev/null @@ -1,1104 +0,0 @@ -"""SM90 fused scaled cross entropy **forward** kernel: producer-less -"MMA-warp issues TMA" mainloop + dedicated SMEM-logits epilogue warp group. - -This is the default forward used by -:mod:`liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90` -(``m_tiles_per_cluster == 1``). It computes the per-token negative -log-likelihood ``nll[M]`` and the log-sum-exp ``lse[M]`` of ``X @ W.T`` without -ever writing the logits to HBM. The module is self-contained: apart from -:func:`to_cute_tensor` it imports nothing from the rest of the backend. - -Architecture -============ - -384 threads / 3 warp groups, but the split is **not** the production -producer/consumer split: - -=========== ========================================================= -warp group role -=========== ========================================================= -WG0 (0-127) WGMMA, rows ``[0, 64)`` of the M128 tile. **Warp 0 of WG0 - is additionally the *TMA main warp*: it issues every - global->SMEM ``cp.async.bulk.tensor`` for X and W.** -WG1 (128-255) WGMMA, rows ``[64, 128)`` of the M128 tile. -WG2 (256-383) Epilogue: online softmax max/exp/sum + target logit read - from a staged SMEM logits tile, one row per thread, plus - the CTA/cluster DSMEM reduction and the loss/lse stores. -=========== ========================================================= - -There is **no dedicated producer warp group**, so the whole register file is -split two ways instead of three and 2/3 of the threads do WGMMA instead of 1/2. -The MMA topology itself is unchanged from production: a cooperative -``atom_layout = (2, 1, 1)`` tiled MMA over a ``(64, tile_n)`` WGMMA atom, i.e. -``M128 x N256 x K64`` split along M between the two MMA warp groups, 3 stages, -``cluster_n = 16``, no multicast anywhere. - -Mainloop TMA software pipeline ------------------------------- - -The interesting part is that the TMA-issuing warp is a *member of a WGMMA warp -group*, so it must be back with its warp group before every collective -(``.aligned``) WGMMA instruction or the warp group hangs. The loop is arranged -so that this is automatic and so that a WGMMA group is always in flight while -the main warp refills a stage:: - - prologue (main warp only): issue S complete TMA stages [k = 0 .. S-1] - - per output tile: - k = 0: consumer_wait(stage 0) ; fence ; WGMMA ; commit - k = 1 .. K-1: - consumer_wait(stage k) ; fence ; WGMMA ; commit - wait_group(1) <- group k-1 retired, group k IN FLIGHT - consumer_release(stage k-1) - if warp == 0: producer_acquire(stage k-1) ; TMA for k-1+S - (issued with group k still executing) - -> all four warps meet again at the next wgmma.fence - tile tail: - wait_group(0) <- only here is the WGMMA pipe drained - consumer_release(last stage) ; refill it - -Warps 1-3 of WG0 simply wait on the *stage* mbarrier while warp 0 refills; they -never wait on warp 0 through a barrier that warp 0 can only reach after a WGMMA, -so the cycle that would deadlock a naive implementation does not exist. The -load stream is **linear across output tiles** (load ``n`` always goes to stage -``n % S``), so a tile boundary costs no pipeline drain and no TMA bubble: the -main warp is already prefetching the next output tile's W while the current -tile's last WGMMAs run. - -Every stage refill is paired 1:1 with the release that freed it, which keeps -``producer_index == release_index`` at every refill without any modular -arithmetic. The last ``S`` refills of an m-tile are skipped (their data would -be past the end of the stream); those ``S`` unconsumed empty-credits are exactly -the ``S`` credits the next m-tile's prologue consumes, so the pipeline stays -balanced across the persistent m-tile loop. - -SMEM logits staging and the register->(row, col) map ----------------------------------------------------- - -After ``wait_group(0)`` the MMA warp groups convert their FP32 accumulator -straight into a **single** 16-bit SMEM logits tile (no second FP32 tensor). The -register-to-tile map cannot be taken from ``thr_mma.partition_C(identity)``: -with the cooperative ``atom_layout = (2, 1, 1)`` slicing the production kernel -uses (``get_slice`` of a *constant* warp-group base), ``coords[i]`` is identical -for all 128 threads of a warp group -- it carries no lane term at all, so an -explicit map is required. The map used here was validated against a real -WGMMA on an H200 (0 / 32768 mismatches):: - - row = wg * 64 + local_warp * 16 + ((i % 4) // 2) * 8 + lane // 4 - col = (i // 4) * 8 + (lane % 4) * 2 + (i % 2) - -for ``i`` in ``[0, 128)``, ``wg`` in ``{0, 1}``. - -Bank conflicts are the whole ball game here. The production 2-idle-warp -epilogue ablation was correct but ran at 10.73 ms / 410 TFLOPS, and the reason -is the un-padded ``[128, 256]`` 16-bit tile: with one row per thread, every -thread of a warp lands on the *same* bank (row stride 512 B = 128 words ≡ 0 mod -32), i.e. a 32-way conflict on every access, twice (max pass + sum pass). This -module pads the row stride to ``tile_n + 8`` 16-bit elements = 528 B = 132 -words, which makes **both** directions conflict-free: - -* MMA-side stores: element ``i`` and ``i+1`` of the accumulator are the same row - and adjacent columns, so they pack into one 32-bit ``STS``. Its word index is - ``row * 132 + (i // 4) * 4 + lane % 4``; across a warp ``row`` varies as - ``lane // 4`` so the bank is ``4 * (lane // 4) + lane % 4 (mod 32)`` -- a - bijection onto the 32 banks. -* epilogue-side loads: 8 consecutive 16-bit logits are 16 B and 16 B-aligned - (``528 = 33 * 16``), so they issue as ``LDS.128``; the hardware splits that - into phases of 8 lanes and lanes ``t .. t+7`` cover banks - ``4t .. 4t+3 (mod 32)`` -- again exactly the 32 banks, once each. - -The epilogue also does a **single** pass over the tile instead of production's -max-pass-then-sum-pass: it walks the row in ``chunk_n``-column chunks and folds -each chunk into the running online-softmax state, so the SMEM traffic is halved -and the cross-tile accumulation comes for free. - -Because there is exactly one epilogue warp group and it owns a whole row per -thread, there is no per-warp-group partial to combine: thread ``t`` of WG2 holds -the final CTA statistics for row ``t`` and drives the DSMEM cluster reduction -and the ``lse`` / ``token_loss`` stores directly. No ``wg_max/sum/target`` -staging array exists in this kernel. - -Masking sentinels ------------------ - -Columns past ``V`` read whatever the out-of-range TMA left in SMEM (zeros), so -they must be masked. They are replaced by ``MASK_F32 = -3e38`` while the online -state initialises to ``NEG_INF_F32 = -1e38``. The two-sentinel trick removes -the second predicate from the exp pass: ``exp2((-3e38 - m) * log2 e)`` flushes to -zero for every reachable ``m >= -1e38``, including a fully-masked tile where -``m`` is still ``-1e38``, while ``exp2((-1e38 - (-1e38)) * log2 e) == 1`` would -have silently added one per masked column had a single sentinel been used. - -No multicast is used anywhere. There is no HBM logits store in the forward -pass; the only asynchronous stores in this kernel are the TMA *loads*, and the -register->SMEM logits staging is a plain generic ``STS``. - -Accepted configuration ----------------------- -``M=4096 H=4096 V=131072`` BF16: cooperative ``tile_n = 256`` with a 3-deep AB -pipeline and ``cluster_n = 2``. ``tile_n = 128``/``192``, ``stages = 2`` and -``cluster_n >= 8`` were all measurably worse; ``stages = 2`` is catastrophic -(no prefetch depth). NCU on the accepted config: 168 registers/thread, zero -local ld/st (no accumulator spill), tensor pipe ~97 % active. - -Two gotchas worth recording ---------------------------- - -1. **Register-pool deadlock.** ``2*128*mma_registers + 128*epi_registers`` must - be ``<= 64512``, *not* ``<= 65536``. At exactly 65536 (216/80) the hardware - granted ``setmaxnreg.inc`` to only one MMA warp per SM sub-partition and the - other four spun forever in ``USETMAXREG.TRY_ALLOC.CTAPOOL``; both production - (232/40) and the pingpong ablation (240/24) also total 64512. Diagnosed by - attaching ``cuda-gdb -p -ex "info cuda warps"`` to the hung process. -2. **``llvm.nvvm.mapa`` ICEs** this toolchain ("only supported when pointer size - is >= 64 bits"), independent of the SMEM offset. The DSMEM helpers below - therefore use inline PTX ``mapa.shared::cluster.u32`` with 32-bit - ``.shared::cluster`` loads/atomics, exactly as - ``cutlass/cute/arch/smem.py::map_dsmem_ptr`` does for the same reason. Also - note ``red.max.f32`` does not exist in PTX -- the cluster max reduction uses - the monotonic-bit-pattern trick (``atom.max.s32`` / ``atom.min.u32``). -""" - -import inspect - -from dataclasses import dataclass - -import cuda.bindings.driver as cuda -import cutlass -import cutlass.cute as cute -import torch - -from cutlass import BFloat16 -from cutlass import Float16 -from cutlass import Float32 -from cutlass import Int32 -from cutlass import cutlass_dsl -from cutlass import pipeline -from cutlass._mlir.dialects import llvm -from cutlass._mlir.dialects import nvvm -from cutlass.cutlass_dsl import T -from cutlass.cutlass_dsl import dsl_user_op -from cutlass.pipeline.helpers import pipeline_init_arrive -from cutlass.pipeline.helpers import pipeline_init_wait -from cutlass.utils import LayoutEnum -from cutlass.utils import hopper_helpers - -from liger_kernel.ops.cutedsl.ops.utils import to_cute_tensor - -THREADS_PER_CTA = 384 -NUM_MMA_WARP_GROUPS = 2 -WARPS_PER_WARP_GROUP = 4 -EPILOGUE_WARP_GROUP = 2 - -# Named barriers. 0 is reserved for __syncthreads(). -MMA_BARRIER_ID = 4 # both MMA warp groups, 256 threads -EPI_BARRIER_ID = 5 # epilogue warp group, 128 threads - -LOG2_E = 1.4426950408889634 -LN2 = 0.6931471805599453 -NEG_INF_F32 = -1.0e38 -# Distinct, *more* negative sentinel for out-of-vocabulary columns; see module -# docstring ("Masking sentinels"). -MASK_F32 = -3.0e38 -HOPPER_MAX_SMEM_BYTES = 227 * 1024 - -_compile_cache = {} -_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)) - - -@dsl_user_op -def _map_dsmem_addr(local_ptr, peer_rank, *, loc=None, ip=None) -> Int32: - """Map a local SMEM pointer to the same object in a peer CTA's SMEM. - - Returns a **32-bit** ``shared::cluster`` address. The production module - uses the ``nvvm.mapa`` intrinsic on a generic pointer, which libNVVM in this - toolchain rejects for this kernel ("this intrinsic is only supported when - pointer size is >= 64 bits") once the pointer folds into a constant - ``addrspacecast`` expression. ``cutlass.cute.arch.map_dsmem_ptr`` avoids the - intrinsic entirely by emitting the PTX directly; the same trick is used here - so every DSMEM access below is a plain 32-bit shared-window address. - """ - addr = llvm.ptrtoint(T.i32(), local_ptr.llvm_ptr, loc=loc, ip=ip) - return Int32( - llvm.inline_asm( - T.i32(), - [addr, Int32(peer_rank).ir_value(loc=loc, ip=ip)], - "mapa.shared::cluster.u32 $0, $1, $2;", - "=r,r,r", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) - - -def _asm(res_ty, args, tmpl, cons, side=True, *, loc=None, ip=None): - return llvm.inline_asm( - res_ty, - args, - tmpl, - cons, - has_side_effects=side, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def _dsmem_load_f32(addr, *, loc=None, ip=None) -> Float32: - return Float32( - _asm( - T.f32(), - [Int32(addr).ir_value(loc=loc, ip=ip)], - "ld.shared::cluster.f32 $0, [$1];", - "=f,r", - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _dsmem_atomic_add_f32(addr, value, *, loc=None, ip=None) -> Float32: - return Float32( - _asm( - T.f32(), - [ - Int32(addr).ir_value(loc=loc, ip=ip), - Float32(value).ir_value(loc=loc, ip=ip), - ], - "atom.shared::cluster.add.f32 $0, [$1], $2;", - "=f,r,f", - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _dsmem_atomic_max_s32(addr, value, *, loc=None, ip=None) -> Int32: - return Int32( - _asm( - T.i32(), - [ - Int32(addr).ir_value(loc=loc, ip=ip), - Int32(value).ir_value(loc=loc, ip=ip), - ], - "atom.shared::cluster.max.s32 $0, [$1], $2;", - "=r,r,r", - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _dsmem_atomic_min_u32(addr, value, *, loc=None, ip=None) -> Int32: - return Int32( - _asm( - T.i32(), - [ - Int32(addr).ir_value(loc=loc, ip=ip), - Int32(value).ir_value(loc=loc, ip=ip), - ], - "atom.shared::cluster.min.u32 $0, [$1], $2;", - "=r,r,r", - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _dsmem_atomic_fmax(addr, value, *, loc=None, ip=None) -> Float32: - """FP32 max via the monotonic-bit-pattern trick (identical to production). - - ``red/atom.max.f32`` does not exist in PTX, so positive floats are max'd as - signed ints and negative floats are min'd as unsigned ints. - """ - bits = llvm.bitcast(T.i32(), Float32(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) - old = cutlass_dsl.if_generate( - Int32(bits) < 0, - lambda: _dsmem_atomic_min_u32(addr, Int32(bits), loc=loc, ip=ip), - lambda: _dsmem_atomic_max_s32(addr, Int32(bits), loc=loc, ip=ip), - [], - [Int32], - loc=loc, - ip=ip, - ) - return Float32(llvm.bitcast(T.f32(), old.ir_value(loc=loc, ip=ip), 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_n): - key = (torch.cuda.current_device(), cluster_n) - value = _max_active_clusters_cache.get(key) - if value is None: - torch.cuda.current_stream() - value = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_n) - _max_active_clusters_cache[key] = value - return value - - -@dataclass(frozen=True) -class ScaledCEForwardConfig: - """Tuning knobs for the MMA-issues-TMA + SMEM-epilogue forward. - - ``tile_m`` is fixed at 128 (one row per epilogue thread) and ``tile_n`` at - 256 (the production cooperative WGMMA topology). ``logit_pad`` is the extra - 16-bit columns appended to every SMEM logits row; 8 makes the row stride - 528 B and both the MMA-side 32-bit stores and the epilogue-side 128-bit - loads bank-conflict-free (see the module docstring). ``chunk_n`` is how many - columns the epilogue folds into the online-softmax state at a time; it is a - pure register/rescale trade (``tile_n / chunk_n`` extra rescales per tile). - - Register budget: ``2 * 128 * mma_registers + 128 * epi_registers`` must not - exceed **64512** registers. The SM register file holds 65536 registers and - the CTA is resident alone, but requesting the full file deadlocks: the - hardware could only satisfy one ``setmaxnreg.inc`` per sub-partition and the - remaining four MMA warps spun forever in ``USETMAXREG.TRY_ALLOC.CTAPOOL`` - (observed with cuda-gdb at 216/80 = exactly 65536). Production (232/40) and - the pingpong ablation (240/24) both stop at 64512, i.e. they leave one - 1024-register (8 regs/thread x 128 threads) reserve. - """ - - tile_m: int = 128 - tile_n: int = 256 - tile_k: int = 64 - stages: int = 3 - cluster_n: int = 16 - fast_math: bool = True - logit_dtype: str = "f16" - logit_pad: int = 8 - chunk_n: int = 32 - mma_registers: int = 208 - epi_registers: int = 88 - return_entropy: bool = False - - @property - def logit_stride(self): - return self.tile_n + self.logit_pad - - def smem_bytes(self): - ab = (self.tile_m + self.tile_n) * self.tile_k * 2 * self.stages - logits = self.tile_m * self.logit_stride * 2 - reduction = (3 + int(self.return_entropy)) * self.tile_m * 4 - barriers = 16 * self.stages + 16 - return ab + logits + reduction + barriers + 3 * 1024 - - def register_total(self): - return 2 * 128 * self.mma_registers + 128 * self.epi_registers - - -class _ScaledCEForwardSM90: - """MMA-warp-issues-TMA forward with a dedicated SMEM-logits epilogue WG.""" - - def __init__(self, config: ScaledCEForwardConfig): - self.tile_m = config.tile_m - self.tile_n = config.tile_n - self.tile_k = config.tile_k - self.stages = config.stages - self.cluster_n = config.cluster_n - self.fast_math = config.fast_math - self.logit_pad = config.logit_pad - self.logit_stride = config.tile_n + config.logit_pad - self.chunk_n = config.chunk_n - self.mma_registers = config.mma_registers - self.epi_registers = config.epi_registers - self.return_entropy = config.return_entropy - self.logit_dtype = Float16 if config.logit_dtype == "f16" else BFloat16 - self.tile_shape = (self.tile_m, self.tile_n, self.tile_k) - self.buffer_align = 1024 - self.mma_barrier = pipeline.NamedBarrier( - barrier_id=MMA_BARRIER_ID, - num_threads=NUM_MMA_WARP_GROUPS * 128, - ) - self.epi_barrier = pipeline.NamedBarrier( - barrier_id=EPI_BARRIER_ID, - num_threads=128, - ) - - @staticmethod - def _make_tma(tensor, smem_layout_staged, smem_tile): - smem_layout = cute.slice_(smem_layout_staged, (None, None, 0)) - return cute.nvgpu.cpasync.make_tiled_tma_atom( - cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(), - tensor, - smem_layout, - smem_tile, - num_multicast=1, - ) - - @cute.jit - def __call__( - self, - x: cute.Tensor, - weight: cute.Tensor, - target: cute.Tensor, - lse: cute.Tensor, - token_loss: cute.Tensor, - entropy: cute.Tensor, - inverse_temperature: Float32, - ignore_index: Int32, - max_active_clusters: cutlass.Constexpr, - stream: cuda.CUstream, - ): - x_layout = LayoutEnum.from_tensor(x) - weight_layout = LayoutEnum.from_tensor(weight) - # Production's cooperative split-M topology, unchanged. - tiled_mma = hopper_helpers.make_trivial_tiled_mma( - BFloat16, - BFloat16, - x_layout.sm90_mma_major_mode(), - weight_layout.sm90_mma_major_mode(), - Float32, - (NUM_MMA_WARP_GROUPS, 1, 1), - (64, self.tile_n), - ) - a_smem_layout = hopper_helpers.make_smem_layout_a( - x_layout, - self.tile_shape, - BFloat16, - self.stages, - ) - b_smem_layout = hopper_helpers.make_smem_layout_b( - weight_layout, - self.tile_shape, - BFloat16, - self.stages, - ) - tma_x, tma_tensor_x = self._make_tma(x, a_smem_layout, (self.tile_m, self.tile_k)) - tma_w, tma_tensor_w = self._make_tma(weight, b_smem_layout, (self.tile_n, self.tile_k)) - - logit_dtype = self.logit_dtype - - @cute.struct - class SharedStorage: - pipeline: cute.struct.MemRange[cutlass.Int64, self.stages * 2] - logits_pipeline: cute.struct.MemRange[cutlass.Int64, 2] - sX: cute.struct.Align[ - cute.struct.MemRange[BFloat16, cute.cosize(a_smem_layout)], - self.buffer_align, - ] - sW: cute.struct.Align[ - cute.struct.MemRange[BFloat16, cute.cosize(b_smem_layout)], - self.buffer_align, - ] - sLogits: cute.struct.Align[ - cute.struct.MemRange[ - logit_dtype, - self.tile_m * self.logit_stride, - ], - self.buffer_align, - ] - cluster_max: cute.struct.Align[cute.struct.MemRange[Float32, self.tile_m], 16] - cluster_sum: cute.struct.Align[cute.struct.MemRange[Float32, self.tile_m], 16] - cluster_target: cute.struct.Align[cute.struct.MemRange[Float32, self.tile_m], 16] - cluster_weighted: cute.struct.Align[ - cute.struct.MemRange[Float32, self.tile_m if self.return_entropy else 1], - 16, - ] - - self.shared_storage = SharedStorage - cluster_shape = (1, self.cluster_n, 1) - num_m_tiles = cute.ceil_div(x.shape[0], self.tile_m) - num_persistent_clusters = cutlass.min(num_m_tiles, max_active_clusters) - grid = (1, self.cluster_n, num_persistent_clusters) - self.kernel( - tma_x, - tma_tensor_x, - tma_w, - tma_tensor_w, - target, - lse, - token_loss, - entropy, - inverse_temperature, - tiled_mma, - a_smem_layout, - b_smem_layout, - ignore_index, - ).launch( - grid=grid, - block=(THREADS_PER_CTA, 1, 1), - cluster=cluster_shape, - min_blocks_per_mp=1, - stream=stream, - ) - - @cute.kernel - def kernel( - self, - tma_x: cute.CopyAtom, - x: cute.Tensor, - tma_w: cute.CopyAtom, - weight: cute.Tensor, - target: cute.Tensor, - lse: cute.Tensor, - token_loss: cute.Tensor, - entropy: cute.Tensor, - inverse_temperature: Float32, - tiled_mma: cute.TiledMma, - a_smem_layout: cute.ComposedLayout, - b_smem_layout: cute.ComposedLayout, - ignore_index: Int32, - ): - tid, _, _ = cute.arch.thread_idx() - warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - warp_group = cute.arch.make_warp_uniform(tid // 128) - lane = tid % 32 - local_warp = (tid % 128) // 32 - local_tid = tid % 128 - cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) - _, _, pid_m = cute.arch.block_idx() - _, _, m_tile_stride = cute.arch.grid_dim() - num_m_tiles = cute.ceil_div(target.shape[0], self.tile_m) - - if warp == 0: - cute.nvgpu.cpasync.prefetch_descriptor(tma_x) - cute.nvgpu.cpasync.prefetch_descriptor(tma_w) - - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(self.shared_storage) - sX = storage.sX.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner) - sW = storage.sW.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner) - sLogits = storage.sLogits.get_tensor( - cute.make_layout( - (self.tile_m, self.tile_n), - stride=(self.logit_stride, 1), - ) - ) - cluster_max = storage.cluster_max.get_tensor(cute.make_layout(self.tile_m)) - cluster_sum = storage.cluster_sum.get_tensor(cute.make_layout(self.tile_m)) - cluster_target = storage.cluster_target.get_tensor(cute.make_layout(self.tile_m)) - cluster_weighted = storage.cluster_weighted.get_tensor( - cute.make_layout(self.tile_m if cutlass.const_expr(self.return_entropy) else 1) - ) - - cta_layout = cute.make_layout((1, self.cluster_n, 1)) - cluster_coord = cta_layout.get_flat_coord(cluster_rank) - - gX = cute.local_tile( - x, - cute.slice_(self.tile_shape, (None, 0, None)), - (None, None, None), - ) - gW = cute.local_tile( - weight, - cute.slice_(self.tile_shape, (0, None, None)), - (None, None, None), - ) - tXsX, tXgX = cute.nvgpu.cpasync.tma_partition( - tma_x, - 0, - cute.make_layout(1), - cute.group_modes(sX, 0, 2), - cute.group_modes(gX, 0, 2), - ) - w_cta_layout = cute.make_layout(cute.slice_(cta_layout, (None, 0, 0)).shape) - tWsW, tWgW = cute.nvgpu.cpasync.tma_partition( - tma_w, - cluster_coord[0], - w_cta_layout, - cute.group_modes(sW, 0, 2), - cute.group_modes(gW, 0, 2), - ) - - tma_bytes = cute.size_in_bytes(BFloat16, cute.slice_(a_smem_layout, (None, None, 0))) + cute.size_in_bytes( - BFloat16, cute.slice_(b_smem_layout, (None, None, 0)) - ) - - # Both MMA warp groups (8 warps) consume and release every AB stage. - mainloop = pipeline.PipelineTmaAsync.create( - barrier_storage=storage.pipeline.data_ptr(), - num_stages=self.stages, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, - # ``PipelineTmaAsync.consumer_release`` is guarded by - # ``is_signalling_thread`` (= lane 0 of each warp when the CTA - # layout has size 1), so the empty barrier receives exactly one - # arrival per consumer WARP, not per thread. - NUM_MMA_WARP_GROUPS * WARPS_PER_WARP_GROUP, - ), - tx_count=tma_bytes, - cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), - defer_sync=True, - ) - # One-stage hand-off of the staged SMEM logits tile: the MMA warp groups - # are the producer, the epilogue warp group the consumer. Single-thread - # commit / release on both sides, everybody waits. - logits_pipeline = pipeline.PipelineAsync.create( - barrier_storage=storage.logits_pipeline.data_ptr(), - num_stages=1, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), - defer_sync=True, - ) - pipeline_init_arrive(cluster_shape_mn=(1, self.cluster_n), is_relaxed=True) - pipeline_init_wait(cluster_shape_mn=(1, self.cluster_n)) - - num_k_tiles = cute.size(gX, mode=[3]) - num_k_blocks = self.tile_k // 16 - num_pair_tiles = cute.ceil_div(weight.shape[0], self.cluster_n * self.tile_n) - - root_max_addr = _map_dsmem_addr(storage.cluster_max.data_ptr(), 0) - root_sum_addr = _map_dsmem_addr(storage.cluster_sum.data_ptr(), 0) - root_target_addr = _map_dsmem_addr(storage.cluster_target.data_ptr(), 0) - root_weighted_addr = _map_dsmem_addr(storage.cluster_weighted.data_ptr(), 0) - - if warp_group < NUM_MMA_WARP_GROUPS: - # ================= MMA warp groups (WG0 + WG1) ================= - # Register budget is raised BEFORE the accumulator is materialised - # so the epilogue warp group never carries its live range. - cute.arch.setmaxregister_increase(self.mma_registers) - - mma_wg_layout = cute.make_layout(NUM_MMA_WARP_GROUPS, stride=128) - thr_mma = tiled_mma.get_slice(mma_wg_layout(Int32(warp_group))) - tCrX = tiled_mma.make_fragment_A(thr_mma.partition_A(sX)) - tCrW = tiled_mma.make_fragment_B(thr_mma.partition_B(sW)) - accum = cute.make_rmem_tensor( - thr_mma.partition_shape_C((self.tile_m, self.tile_n)), - Float32, - ) - n_accum = cute.size(accum) - row_base_in_tile = warp_group * 64 + local_warp * 16 + lane // 4 - col_base_in_tile = (lane % 4) * 2 - - read_index = Int32(0) - read_phase = Int32(0) - rel_index = Int32(0) - rel_phase = Int32(0) - load_index = Int32(0) - # Producer states start at phase 1 (``make_pipeline_state(Producer)``) - # so the first ``producer_acquire`` on a freshly-initialised empty - # barrier completes immediately. - load_phase = Int32(1) - logit_w_index = Int32(0) - logit_w_phase = Int32(1) - - while pid_m < num_m_tiles: - cute.arch.sync_threads() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - tXgX_m = tXgX[(None, pid_m, None, 0)] - load_k = Int32(0) - load_t = Int32(0) - - # ---- prologue: S complete TMA stages before any WGMMA ------ - if warp == 0: - prologue_stages = cutlass.min(Int32(self.stages), Int32(num_k_tiles) * Int32(num_pair_tiles)) - for _ in range(prologue_stages): - pstate = pipeline.PipelineState(self.stages, Int32(0), load_index, load_phase) - mainloop.producer_acquire(pstate) - bar = mainloop.producer_get_barrier(pstate) - cute.copy( - tma_x, - tXgX_m[(None, load_k)], - tXsX[(None, load_index)], - tma_bar_ptr=bar, - mcast_mask=0, - ) - tWgW_n = tWgW[(None, load_t * self.cluster_n + cluster_rank, None, 0)] - cute.copy( - tma_w, - tWgW_n[(None, load_k)], - tWsW[(None, load_index)], - tma_bar_ptr=bar, - ) - mainloop.producer_commit(pstate) - pstate.advance() - load_index = pstate.index - load_phase = pstate.phase - load_k += 1 - if load_k == num_k_tiles: - load_k = Int32(0) - load_t += 1 - - for _ in range(num_pair_tiles): - accum.fill(0.0) - tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, True) - - # ---- k = 0: no release yet, nothing to refill ---------- - read_state = pipeline.PipelineState(self.stages, Int32(0), read_index, read_phase) - mainloop.consumer_wait(read_state) - cute.nvgpu.warpgroup.fence() - for k_block in cutlass.range_constexpr(num_k_blocks): - coord = (None, None, k_block, read_index) - cute.gemm(tiled_mma, accum, tCrX[coord], tCrW[coord], accum) - cute.nvgpu.warpgroup.commit_group() - read_state.advance() - read_index = read_state.index - read_phase = read_state.phase - - # ---- steady state -------------------------------------- - for _ in range(1, num_k_tiles): - read_state = pipeline.PipelineState(self.stages, Int32(0), read_index, read_phase) - mainloop.consumer_wait(read_state) - cute.nvgpu.warpgroup.fence() - for k_block in cutlass.range_constexpr(num_k_blocks): - coord = (None, None, k_block, read_index) - cute.gemm(tiled_mma, accum, tCrX[coord], tCrW[coord], accum) - cute.nvgpu.warpgroup.commit_group() - # Retire only the previous group: the group just issued - # stays IN FLIGHT across the release + TMA refill below. - cute.nvgpu.warpgroup.wait_group(1) - rstate = pipeline.PipelineState(self.stages, Int32(0), rel_index, rel_phase) - mainloop.consumer_release(rstate) - rstate.advance() - rel_index = rstate.index - rel_phase = rstate.phase - if warp == 0: - if load_t < num_pair_tiles: - pstate = pipeline.PipelineState(self.stages, Int32(0), load_index, load_phase) - mainloop.producer_acquire(pstate) - bar = mainloop.producer_get_barrier(pstate) - cute.copy( - tma_x, - tXgX_m[(None, load_k)], - tXsX[(None, load_index)], - tma_bar_ptr=bar, - mcast_mask=0, - ) - tWgW_n = tWgW[(None, load_t * self.cluster_n + cluster_rank, None, 0)] - cute.copy( - tma_w, - tWgW_n[(None, load_k)], - tWsW[(None, load_index)], - tma_bar_ptr=bar, - ) - mainloop.producer_commit(pstate) - pstate.advance() - load_index = pstate.index - load_phase = pstate.phase - load_k += 1 - if load_k == num_k_tiles: - load_k = Int32(0) - load_t += 1 - read_state.advance() - read_index = read_state.index - read_phase = read_state.phase - - # ---- tile tail: the only WGMMA drain ------------------- - cute.nvgpu.warpgroup.wait_group(0) - rstate = pipeline.PipelineState(self.stages, Int32(0), rel_index, rel_phase) - mainloop.consumer_release(rstate) - rstate.advance() - rel_index = rstate.index - rel_phase = rstate.phase - if warp == 0: - if load_t < num_pair_tiles: - pstate = pipeline.PipelineState(self.stages, Int32(0), load_index, load_phase) - mainloop.producer_acquire(pstate) - bar = mainloop.producer_get_barrier(pstate) - cute.copy( - tma_x, - tXgX_m[(None, load_k)], - tXsX[(None, load_index)], - tma_bar_ptr=bar, - mcast_mask=0, - ) - tWgW_n = tWgW[(None, load_t * self.cluster_n + cluster_rank, None, 0)] - cute.copy( - tma_w, - tWgW_n[(None, load_k)], - tWsW[(None, load_index)], - tma_bar_ptr=bar, - ) - mainloop.producer_commit(pstate) - pstate.advance() - load_index = pstate.index - load_phase = pstate.phase - load_k += 1 - if load_k == num_k_tiles: - load_k = Int32(0) - load_t += 1 - - # ---- stage the FP32 accumulator as 16-bit SMEM ----- - wstate = pipeline.PipelineState(1, Int32(0), logit_w_index, logit_w_phase) - logits_pipeline.producer_acquire(wstate) - for i in cutlass.range_constexpr(n_accum): - row = row_base_in_tile + ((i % 4) // 2) * 8 - col = col_base_in_tile + (i // 4) * 8 + (i % 2) - sLogits[row, col] = self.logit_dtype(accum[i]) - cute.arch.fence_proxy("async.shared", space="cta") - cute.arch.fence_acq_rel_cta() - self.mma_barrier.arrive_and_wait() - if tid == 0: - logits_pipeline.producer_commit(wstate) - wstate.advance() - logit_w_index = wstate.index - logit_w_phase = wstate.phase - - cute.arch.sync_threads() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - pid_m += m_tile_stride - - if warp == 0: - ptail = pipeline.PipelineState(self.stages, Int32(0), load_index, load_phase) - mainloop.producer_tail(ptail) - else: - # ==================== epilogue warp group ====================== - cute.arch.setmaxregister_decrease(self.epi_registers) - - row = local_tid - chunk = cute.make_rmem_tensor((self.chunk_n,), Float32) - logit_r_index = Int32(0) - logit_r_phase = Int32(0) - - while pid_m < num_m_tiles: - if cluster_rank == 0: - cluster_max[row] = Float32(NEG_INF_F32) - cluster_sum[row] = Float32(0.0) - cluster_target[row] = Float32(0.0) - if cutlass.const_expr(self.return_entropy): - cluster_weighted[row] = Float32(0.0) - cute.arch.sync_threads() - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - run_max = Float32(NEG_INF_F32) - run_sum = Float32(0.0) - run_target = Float32(0.0) - run_weighted = Float32(0.0) - global_row = pid_m * self.tile_m + row - tgt_tile = Int32(-1) - tgt_col = Int32(-1) - if global_row < target.shape[0]: - tcol = Int32(target[global_row]) - if tcol >= 0: - tgt_tile = tcol // self.tile_n - tgt_col = tcol % self.tile_n - - for pair_tile in range(num_pair_tiles): - rstate = pipeline.PipelineState(1, Int32(0), logit_r_index, logit_r_phase) - logits_pipeline.consumer_wait(rstate) - cute.arch.fence_acq_rel_cta() - - n_tile = pair_tile * self.cluster_n + cluster_rank - n_base = n_tile * self.tile_n - valid_cols = Int32(weight.shape[0]) - n_base - - # Single pass: fold each chunk straight into the - # running online-softmax state. - for ci in cutlass.range_constexpr(self.tile_n // self.chunk_n): - c0 = ci * self.chunk_n - chunk_max = Float32(MASK_F32) - # One dynamic bound check per CHUNK instead of one - # per column: at the benchmark shape every chunk is - # fully in range, so the masked path is never taken - # and the hot loop is a pure LDS + max sequence. - if c0 + self.chunk_n <= valid_cols: - for j in cutlass.range_constexpr(self.chunk_n): - v = Float32(sLogits[row, c0 + j]) * inverse_temperature - chunk[j] = v - chunk_max = _fmax(chunk_max, v) - else: - for j in cutlass.range_constexpr(self.chunk_n): - v = Float32(sLogits[row, c0 + j]) * inverse_temperature - if c0 + j >= valid_cols: - v = Float32(MASK_F32) - chunk[j] = v - chunk_max = _fmax(chunk_max, v) - new_max = _fmax(run_max, chunk_max) - acc = Float32(0.0) - weighted = Float32(0.0) - for j in cutlass.range_constexpr(self.chunk_n): - exp_value = cute.math.exp2( - (chunk[j] - new_max) * LOG2_E, - fastmath=self.fast_math, - ) - acc += exp_value - if cutlass.const_expr(self.return_entropy): - weighted += exp_value * chunk[j] - old_scale = cute.math.exp2( - (run_max - new_max) * LOG2_E, - fastmath=self.fast_math, - ) - run_sum = run_sum * old_scale + acc - if cutlass.const_expr(self.return_entropy): - run_weighted = run_weighted * old_scale + weighted - run_max = new_max - - if tgt_tile == n_tile: - run_target += Float32(sLogits[row, tgt_col]) * inverse_temperature - - self.epi_barrier.arrive_and_wait() - if local_tid == 0: - logits_pipeline.consumer_release(rstate) - rstate.advance() - logit_r_index = rstate.index - logit_r_phase = rstate.phase - - cute.arch.sync_threads() - - _dsmem_atomic_fmax(root_max_addr + row * 4, run_max) - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - global_max = _dsmem_load_f32(root_max_addr + row * 4) - scaled_sum = run_sum * cute.math.exp2( - (run_max - global_max) * LOG2_E, - fastmath=self.fast_math, - ) - _dsmem_atomic_add_f32(root_sum_addr + row * 4, scaled_sum) - _dsmem_atomic_add_f32(root_target_addr + row * 4, run_target) - if cutlass.const_expr(self.return_entropy): - scaled_weighted = run_weighted * cute.math.exp2( - (run_max - global_max) * LOG2_E, - fastmath=self.fast_math, - ) - _dsmem_atomic_add_f32(root_weighted_addr + row * 4, scaled_weighted) - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - - if cluster_rank == 0: - if global_row < target.shape[0]: - row_lse = cluster_max[row] + cute.math.log2(cluster_sum[row], fastmath=self.fast_math) * LN2 - lse[global_row] = row_lse - if target[global_row] == ignore_index: - token_loss[global_row] = Float32(0.0) - if cutlass.const_expr(self.return_entropy): - entropy[global_row] = BFloat16(0.0) - else: - token_loss[global_row] = row_lse - cluster_target[row] - if cutlass.const_expr(self.return_entropy): - entropy[global_row] = BFloat16(row_lse - cluster_weighted[row] / cluster_sum[row]) - - cute.arch.cluster_arrive() - cute.arch.cluster_wait() - pid_m += m_tile_stride - - -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, - ) - - -def scaled_ce_forward( - _input, - weight, - target, - temperature=1.0, - ignore_index=-100, - return_entropy=False, - config: ScaledCEForwardConfig = None, -): - """M-outer (self-TMA) forward. Returns per-token ``(nll[M], lse[M])``.""" - _validate(_input, weight, target) - if config is None: - config = ScaledCEForwardConfig() - if config.return_entropy != return_entropy: - config = ScaledCEForwardConfig(**{**config.__dict__, "return_entropy": return_entropy}) - if config.tile_m != 128: - raise ValueError("tile_m must be 128 (one epilogue thread per row)") - if config.tile_n % 8 != 0 or config.tile_n % config.chunk_n != 0: - raise ValueError("tile_n must be a multiple of 8 and of chunk_n") - if config.smem_bytes() > HOPPER_MAX_SMEM_BYTES: - raise ValueError( - f"config needs ~{config.smem_bytes()} B of shared memory, " - f"which exceeds the {HOPPER_MAX_SMEM_BYTES} B Hopper opt-in limit" - ) - if config.register_total() > 64512: - raise ValueError(f"register budget {config.register_total()} exceeds the usable 64512-register budget") - - target = target.contiguous() - x_padded, weight_padded, _ = _pad_hidden(_input, weight, config.tile_k) - bt = x_padded.shape[0] - - lse = torch.empty(bt, device=_input.device, dtype=torch.float32) - token_loss = torch.empty(bt, device=_input.device, dtype=torch.float32) - entropy = torch.empty(bt if return_entropy else 1, device=_input.device, dtype=torch.bfloat16) - x_cute = to_cute_tensor(x_padded.unsqueeze(-1), leading_dim=1, assumed_align=16) - weight_cute = to_cute_tensor(weight_padded.unsqueeze(-1), leading_dim=1, assumed_align=16) - target_cute = to_cute_tensor(target, assumed_align=8) - lse_cute = to_cute_tensor(lse, assumed_align=4) - token_loss_cute = to_cute_tensor(token_loss, assumed_align=4) - entropy_cute = to_cute_tensor(entropy, assumed_align=2) - stream = _cute_stream() - inverse_temperature = Float32(1.0 / temperature) - max_active_clusters = _max_active_clusters(config.cluster_n) - - key = ("scaled_ce_forward", config, max_active_clusters) - compiled = _compile_cache.get(key) - if compiled is None: - compiled = cute.compile( - _ScaledCEForwardSM90(config), - x_cute, - weight_cute, - target_cute, - lse_cute, - token_loss_cute, - entropy_cute, - inverse_temperature, - Int32(ignore_index), - max_active_clusters, - stream, - ) - _compile_cache[key] = compiled - compiled( - x_cute, - weight_cute, - target_cute, - lse_cute, - token_loss_cute, - entropy_cute, - inverse_temperature, - Int32(ignore_index), - stream, - ) - - return token_loss, entropy if return_entropy else None, lse diff --git a/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_utils_sm90.py b/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_utils_sm90.py new file mode 100644 index 00000000..25e06c91 --- /dev/null +++ b/src/liger_kernel/ops/cutedsl/ops/_fused_scaled_cross_entropy_utils_sm90.py @@ -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, + ) diff --git a/src/liger_kernel/ops/cutedsl/ops/fused_scaled_cross_entropy_sm90.py b/src/liger_kernel/ops/cutedsl/ops/fused_scaled_cross_entropy_sm90.py index f1d86279..ef77d180 100644 --- a/src/liger_kernel/ops/cutedsl/ops/fused_scaled_cross_entropy_sm90.py +++ b/src/liger_kernel/ops/cutedsl/ops/fused_scaled_cross_entropy_sm90.py @@ -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, diff --git a/test/cutedsl/test_fused_scaled_cross_entropy_forward_fragment_sm90.py b/test/cutedsl/test_fused_scaled_cross_entropy_forward_fragment_sm90.py new file mode 100644 index 00000000..09fe6b36 --- /dev/null +++ b/test/cutedsl/test_fused_scaled_cross_entropy_forward_fragment_sm90.py @@ -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]) diff --git a/test/cutedsl/test_fused_scaled_cross_entropy_sm90.py b/test/cutedsl/test_fused_scaled_cross_entropy_sm90.py index 009ec494..538fdd3b 100644 --- a/test/cutedsl/test_fused_scaled_cross_entropy_sm90.py +++ b/test/cutedsl/test_fused_scaled_cross_entropy_sm90.py @@ -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():