diff --git a/README.md b/README.md index ac14e23b..740e2e1d 100644 --- a/README.md +++ b/README.md @@ -209,6 +209,28 @@ LIGER_KERNEL_IMPL=cutile python your_script.py `LIGER_KERNEL_IMPL` selects an opt-in implementation registered with Liger (currently `cutile` and `cutedsl`). Selecting one on an unsupported device, or without the required dependencies installed, raises an error. +#### cuTile Fused Linear Cross Entropy + +With `LIGER_KERNEL_IMPL=cutile`, `LigerFusedLinearCrossEntropyLoss` routes to a cuTile Fused Linear Cross Entropy kernel for BF16 `input`/`weight` on Hopper (SM90) and Blackwell (SM100). No API change is needed — the default route just works: + +```bash +LIGER_KERNEL_IMPL=cutile python your_script.py +``` + +```python +from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss + +# Same call as the default backend — there is no chunk_size argument. +loss_fn = LigerFusedLinearCrossEntropyLoss(reduction="mean") +loss = loss_fn(lm_head_weight, hidden_states, labels) +``` + +Implementation notes: + +- Token chunking uses the **exact same default geometry as the Triton backend** (`inc = ceil(V / (16·H))`, `chunk = min(N, next_pow2(ceil(N / inc)))`); there is no public `chunk_size` knob. +- All three GEMMs — projection, input-gradient (`dX`), and weight-gradient (`dW`) — run through PyTorch/cuBLAS. `dW` is accumulated across token chunks in FP32 (`mm` then `addmm`) and the combined upstream/mean scale is applied in FP32 before a single BF16 cast, so small gradients are preserved. +- Supports BF16 `input` and `weight` with `mean`/`sum` reductions; FP32 accumulation is the default and only accumulation policy. It is **not** a full feature/precision match for the Triton kernel (e.g. no bias, class weights, z-loss, label smoothing, softcap, or token scaling). + ### Enable CuTe DSL Backend CuTe DSL is the optional, CUDA-only Python DSL shipped with NVIDIA CUTLASS (`import cutlass.cute`), targeting Hopper (SM90) and Blackwell (SM100/SM110). After installing the `cutedsl` extra, enable it explicitly: diff --git a/src/liger_kernel/ops/cutile/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/cutile/ops/fused_linear_cross_entropy.py index 644cf424..fa5c7420 100644 --- a/src/liger_kernel/ops/cutile/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/cutile/ops/fused_linear_cross_entropy.py @@ -2,94 +2,84 @@ # # SPDX-License-Identifier: MIT -"""Hopper-optimized cuTile fused linear cross entropy. +"""cuTile fused linear cross entropy for Hopper (SM90) and Blackwell (SM100). -The projection and input-gradient GEMMs use PyTorch's Hopper-tuned BLAS path. -cuTile computes partitioned CE statistics, overwrites logits with dZ, and -computes dW with an FP32-accumulating 128x256x64 MMA schedule. +All three GEMMs -- the projection, the input-gradient (dX), and the +weight-gradient (dW) -- use PyTorch's tuned cuBLAS path. cuTile computes +partitioned CE statistics, overwrites a reusable ``[min(N, C), V]`` logits +buffer with dZ, and applies the final FP32 dW scale-and-cast. dW is accumulated +across token chunks in FP32 by cuBLAS ``mm`` (first chunk) then ``addmm`` +(subsequent chunks), so the transient logits workspace is +``O(min(N, C) * V)`` instead of the full ``O(N * V)`` materialization; the +retained gradients (raw BF16 dX and the FP32 dW accumulator) add ``O(N*H + V*H)``. + +The upstream/mean scale is applied out-of-place in backward: the FP32 dW +accumulator is retained until backward and cast to the weight dtype exactly once, +after multiplying by the combined ``upstream * mean_normalizer`` factor. Casting +after (not before) that multiply preserves the numeric range of small gradients +that would otherwise flush to zero in BF16. + +Two low-level interfaces are exported. ``fused_linear_cross_entropy_forward`` / +``fused_linear_cross_entropy_backward`` retain the original unchunked contract +(the forward returns the full ``[N, V]`` dZ and the 7-argument backward runs the +dW matmul). The token-chunked path used by the autograd Function lives in +``chunked_fused_linear_cross_entropy_forward`` / +``chunked_fused_linear_cross_entropy_backward``. """ +import operator + from typing import Optional import cuda.tile as ct import torch from liger_kernel.ops.cutile.ops.utils import _next_power_of_2 +from liger_kernel.ops.fused_linear_cross_entropy import _get_chunk_size from liger_kernel.ops.utils import amp_custom_bwd from liger_kernel.ops.utils import amp_custom_fwd +from liger_kernel.ops.utils import compare_version ConstBool = ct.Constant[bool] ConstInt = ct.Constant[int] LOG2E = 1.4426950408889634 -DW_TILE_M = 128 -DW_TILE_N = 128 -DW_TILE_K = 64 +SCALE_CAST_TILE_M = 128 +SCALE_CAST_TILE_N = 128 CE_BLOCK_SIZE = 2048 LOGITS_STATS_BLOCK_SIZE = 4096 MAX_STATS_BLOCK_SIZE = 1024 +# The FP32 dW accumulation uses the ``torch.mm`` / ``torch.addmm`` +# ``out_dtype=torch.float32`` overload (BF16 operands -> FP32 accumulator) added +# in torch 2.8. On older torch the operands are upcast to FP32 first; see +# ``_accumulate_dw_fp32``. This is a runtime capability check only -- it does not +# raise the package floor. +_ADDMM_SUPPORTS_OUT_DTYPE = compare_version("torch", operator.ge, "2.8.0") -@ct.kernel(num_worker_warps=ct.ByTarget(sm_90=8), opt_level=3) -def _matmul_tn_kernel( - a, - b, - output, - grad_output, + +@ct.kernel(num_worker_warps=ct.ByTarget(sm_90=8, default=8), opt_level=3) +def _scale_cast_kernel( + src, + dst, + scale, TILE_M: ConstInt, TILE_N: ConstInt, - TILE_K: ConstInt, ): - """Compute ``output = (a.T @ b) * grad_output``.""" + """Scale an FP32 ``src`` tile by a device scalar and cast into ``dst.dtype``. + + Used to turn the retained FP32 dW accumulator into the final weight-dtype + gradient without allocating an O(V*H) FP32 temporary: the multiply by the + combined ``upstream * mean_normalizer`` scalar happens in FP32 and the BF16 + cast is applied exactly once, per element, on store. + """ pid_m = ct.bid(0) pid_n = ct.bid(1) - a_view = a.tiled_view((TILE_K, TILE_M), padding_mode=ct.PaddingMode.ZERO) - b_view = b.tiled_view((TILE_K, TILE_N), padding_mode=ct.PaddingMode.ZERO) - acc = ct.zeros((TILE_M, TILE_N), dtype=ct.float32) - - for pid_k in range(a_view.num_tiles(0)): - a_tile = a_view.load((pid_k, pid_m), latency=10) - b_tile = b_view.load((pid_k, pid_n), latency=10) - acc = ct.mma(ct.transpose(a_tile), b_tile, acc) - - scale = ct.astype(ct.gather(grad_output, (), check_bounds=False), ct.float32) - ct.store(output, (pid_m, pid_n), ct.astype(acc * scale, output.dtype), latency=1) - - -@ct.kernel(num_worker_warps=ct.ByTarget(sm_90=8), opt_level=3) -def _matmul_tn_n4_kernel( - a, - b, - output, - grad_output, - TILE_M: ConstInt, - TILE_N: ConstInt, - TILE_K: ConstInt, - LOAD_LATENCY: ConstInt, -): - """Compute four adjacent dW column tiles while sharing each dZ load.""" - pid_n = ct.bid(0) * 4 - pid_m = ct.bid(1) - a_view = a.tiled_view((TILE_K, TILE_M), padding_mode=ct.PaddingMode.ZERO) - b_view = b.tiled_view((TILE_K, TILE_N), padding_mode=ct.PaddingMode.ZERO) - output_view = output.tiled_view((TILE_M, TILE_N), padding_mode=ct.PaddingMode.ZERO) - acc0 = ct.zeros((TILE_M, TILE_N), dtype=ct.float32) - acc1 = ct.zeros((TILE_M, TILE_N), dtype=ct.float32) - acc2 = ct.zeros((TILE_M, TILE_N), dtype=ct.float32) - acc3 = ct.zeros((TILE_M, TILE_N), dtype=ct.float32) - - for pid_k in range(a_view.num_tiles(0)): - a_tile = ct.transpose(a_view.load((pid_k, pid_m), latency=LOAD_LATENCY)) - acc0 = ct.mma(a_tile, b_view.load((pid_k, pid_n), latency=LOAD_LATENCY), acc0) - acc1 = ct.mma(a_tile, b_view.load((pid_k, pid_n + 1), latency=LOAD_LATENCY), acc1) - acc2 = ct.mma(a_tile, b_view.load((pid_k, pid_n + 2), latency=LOAD_LATENCY), acc2) - acc3 = ct.mma(a_tile, b_view.load((pid_k, pid_n + 3), latency=LOAD_LATENCY), acc3) - - scale = ct.astype(ct.gather(grad_output, (), check_bounds=False), ct.float32) - output_view.store((pid_m, pid_n), ct.astype(acc0 * scale, output.dtype), latency=1) - output_view.store((pid_m, pid_n + 1), ct.astype(acc1 * scale, output.dtype), latency=1) - output_view.store((pid_m, pid_n + 2), ct.astype(acc2 * scale, output.dtype), latency=1) - output_view.store((pid_m, pid_n + 3), ct.astype(acc3 * scale, output.dtype), latency=1) + src_view = src.tiled_view((TILE_M, TILE_N), padding_mode=ct.PaddingMode.ZERO) + dst_view = dst.tiled_view((TILE_M, TILE_N), padding_mode=ct.PaddingMode.ZERO) + tile = ct.astype(src_view.load((pid_m, pid_n), latency=2), ct.float32) + tile_scale = ct.astype(ct.gather(scale, (), check_bounds=False), ct.float32) + dst_view.store((pid_m, pid_n), ct.astype(tile * tile_scale, dst.dtype), latency=1) @ct.kernel(occupancy=8) @@ -177,36 +167,36 @@ def _fused_cross_entropy_dz_kernel( ct.scatter(logits, (row, cols), ct.astype(gradient, logits.dtype), check_bounds=True) -def _launch_matmul_tn( - a: torch.Tensor, - b: torch.Tensor, - output: torch.Tensor, - grad_output: torch.Tensor, +def _accumulate_dw_fp32( + dweight_accum: torch.Tensor, + dz: torch.Tensor, + x: torch.Tensor, + accumulate: bool, ) -> None: - if output.shape[1] % 256 == 0: - _launch_cutile( - output.device, - ( - output.shape[1] // 256, - (output.shape[0] + DW_TILE_M - 1) // DW_TILE_M, - 1, - ), - _matmul_tn_n4_kernel, - (a, b, output, grad_output, DW_TILE_M, 64, DW_TILE_K, 5), - ) - return + """Accumulate ``dZ.T @ X`` into the FP32 ``dweight_accum`` buffer via cuBLAS. - tile_n = min(256, _next_power_of_2(output.shape[1])) - _launch_cutile( - output.device, - ( - (output.shape[0] + DW_TILE_M - 1) // DW_TILE_M, - (output.shape[1] + tile_n - 1) // tile_n, - 1, - ), - _matmul_tn_kernel, - (a, b, output, grad_output, DW_TILE_M, tile_n, DW_TILE_K), - ) + ``dz`` is ``[rows, V]`` and ``x`` is ``[rows, H]``; the contribution + ``dZ.T @ X`` is ``[V, H]``. The first token chunk (``accumulate`` is + ``False``) overwrites the buffer with ``torch.mm``; subsequent chunks add + with ``torch.addmm``. Accumulation is always FP32: on torch>=2.8 the BF16 + operands drive the ``out_dtype=torch.float32`` overload directly; on older + torch (no overload) only the bounded BF16 operands are upcast to FP32 before + the same mm/addmm into the FP32 buffer. This single helper is shared by the + chunked and legacy paths so their dW math is identical. + """ + dz_t = dz.t() # [V, rows] view + if _ADDMM_SUPPORTS_OUT_DTYPE: + if accumulate: + torch.addmm(dweight_accum, dz_t, x, out_dtype=torch.float32, out=dweight_accum) + else: + torch.mm(dz_t, x, out_dtype=torch.float32, out=dweight_accum) + else: + dz_t_fp32 = dz_t.to(torch.float32) + x_fp32 = x.to(torch.float32) + if accumulate: + torch.addmm(dweight_accum, dz_t_fp32, x_fp32, out=dweight_accum) + else: + torch.mm(dz_t_fp32, x_fp32, out=dweight_accum) def _launch_cutile(device: torch.device, grid, kernel, args) -> None: @@ -214,6 +204,20 @@ def _launch_cutile(device: torch.device, grid, kernel, args) -> None: ct.launch(torch.cuda.current_stream(device), grid, kernel, args) +def _launch_scale_cast(src: torch.Tensor, dst: torch.Tensor, scale: torch.Tensor) -> None: + """Compute ``dst = (src * scale).astype(dst.dtype)`` without an FP32 temporary.""" + _launch_cutile( + dst.device, + ( + (src.shape[0] + SCALE_CAST_TILE_M - 1) // SCALE_CAST_TILE_M, + (src.shape[1] + SCALE_CAST_TILE_N - 1) // SCALE_CAST_TILE_N, + 1, + ), + _scale_cast_kernel, + (src, dst, scale, SCALE_CAST_TILE_M, SCALE_CAST_TILE_N), + ) + + def _reject_unsupported( bias, ce_weight, @@ -233,14 +237,34 @@ def _reject_unsupported( "label smoothing": bool(label_smoothing), "softcap": softcap is not None, "return_z_loss": return_z_loss, - "accum_dtype": accum_dtype is not None, "token scaling": use_token_scaling, "return_token_accuracy": return_token_accuracy, "return_predicted_tokens": return_predicted_tokens, } enabled = [name for name, value in unsupported.items() if value] if enabled: - raise NotImplementedError(f"cuTile SM90 FLCE does not support: {', '.join(enabled)}") + raise NotImplementedError(f"cuTile FLCE does not support: {', '.join(enabled)}") + # FP32 accumulation is the default and only supported accumulation dtype. + if accum_dtype is not None and accum_dtype != torch.float32: + raise NotImplementedError( + f"cuTile FLCE only supports FP32 accumulation (accum_dtype=torch.float32 or None), got {accum_dtype!r}" + ) + + +# Impl names that identify this cuTile kernel itself; accepted so callers that +# route by name can target it explicitly. These are self-identity only -- this +# Function does not dispatch to alternative CE implementations. +_SELF_CE_IMPLS = (None, "cutile", "nvidia-cutile") + + +def _validate_ce_dispatch(ce_impl, ce_mode) -> None: + if ce_impl not in _SELF_CE_IMPLS: + raise NotImplementedError( + f"cuTile FLCE does not dispatch to ce_impl={ce_impl!r}; " + f"expected one of {tuple(name for name in _SELF_CE_IMPLS)}" + ) + if ce_mode is not None: + raise NotImplementedError(f"cuTile FLCE only supports the default CE mode (ce_mode=None), got {ce_mode!r}") def _validate_inputs( @@ -249,14 +273,17 @@ def _validate_inputs( target: torch.Tensor, reduction: str, ) -> None: - if not torch.cuda.is_available() or torch.cuda.get_device_capability(_input.device) != (9, 0): - raise RuntimeError("cuTile FLCE requires a Hopper (compute capability 9.0) GPU") if _input.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: raise TypeError("cuTile FLCE supports BF16 input and weight only") if target.dtype != torch.int64: raise TypeError("target must be an int64 tensor") if _input.ndim != 2 or weight.ndim != 2 or target.ndim != 1: raise ValueError("expected input[M, H], weight[V, H], and target[M]") + if 0 in _input.shape or 0 in weight.shape or target.shape[0] == 0: + raise ValueError( + f"cuTile FLCE does not support empty tensors: input {tuple(_input.shape)}, " + f"weight {tuple(weight.shape)}, target {tuple(target.shape)}" + ) if _input.shape[0] != target.shape[0] or _input.shape[1] != weight.shape[1]: raise ValueError( f"incompatible input {tuple(_input.shape)}, weight {tuple(weight.shape)}, and target {tuple(target.shape)}" @@ -265,8 +292,18 @@ def _validate_inputs( raise ValueError(f"reduction must be 'mean' or 'sum', got {reduction!r}") if _input.device != weight.device or _input.device != target.device: raise ValueError("input, weight, and target must be on the same CUDA device") + # Reject CPU tensors before touching CUDA device-capability queries. + if _input.device.type != "cuda" or weight.device.type != "cuda" or target.device.type != "cuda": + raise RuntimeError("cuTile FLCE requires CUDA (Hopper 9.0 or Blackwell SM100 10.0) tensors") + if not torch.cuda.is_available(): + raise RuntimeError("cuTile FLCE requires an available CUDA device") if not _input.is_contiguous() or not weight.is_contiguous() or not target.is_contiguous(): raise ValueError("cuTile FLCE requires contiguous input, weight, and target tensors") + capability = torch.cuda.get_device_capability(_input.device) + if capability not in ((9, 0), (10, 0)): + raise RuntimeError( + f"cuTile FLCE requires a Hopper (compute capability 9.0) or Blackwell SM100 (10.0) GPU, got {capability}" + ) num_stats_partitions = (weight.shape[0] + LOGITS_STATS_BLOCK_SIZE - 1) // LOGITS_STATS_BLOCK_SIZE if num_stats_partitions > MAX_STATS_BLOCK_SIZE: raise NotImplementedError( @@ -274,6 +311,95 @@ def _validate_inputs( ) +def _ce_setup(_input, weight, target, reduction, ignore_index): + """Shared prologue: derive shapes, validate targets, and build the scales. + + Returns the geometry the CE kernel needs plus ``loss_scale`` (passed to the + kernel for the mean-reduced per-row loss) and ``gradient_scale`` (the global + mean normalizer, ``1`` for ``sum``) applied to the gradients in backward. + """ + tokens, hidden_size = _input.shape + vocab_size = weight.shape[0] + num_vocab_tiles = (vocab_size + LOGITS_STATS_BLOCK_SIZE - 1) // LOGITS_STATS_BLOCK_SIZE + stats_block_size = _next_power_of_2(num_vocab_tiles) + needs_dz = _input.requires_grad or weight.requires_grad + reduction_mean = reduction == "mean" + + target_mask = target != ignore_index + valid_target = ~target_mask | ((target >= 0) & (target < vocab_size)) + torch._assert_async(valid_target.all(), f"target values must be in [0, {vocab_size}) or equal ignore_index") + # CE normalization uses the GLOBAL count of non-ignored tokens across all chunks. + loss_scale = ( + target_mask.sum().clamp_min(1).to(torch.float32).reciprocal() + if reduction_mean + else torch.empty((), dtype=torch.float32, device=_input.device) + ) + gradient_scale = loss_scale if reduction_mean else torch.ones((), dtype=torch.float32, device=_input.device) + return ( + tokens, + hidden_size, + vocab_size, + num_vocab_tiles, + stats_block_size, + needs_dz, + reduction_mean, + loss_scale, + gradient_scale, + ) + + +def _project_and_dz( + input_chunk, + weight, + logits_chunk, + target_chunk, + loss_chunk, + loss_scale, + partial_max, + partial_sum, + completion_count, + vocab_size, + num_vocab_tiles, + stats_block_size, + ignore_index, + needs_dz, + reduction_mean, +): + """Project a token chunk to logits and overwrite it in place with dZ. + + ``torch.mm`` fills ``logits_chunk`` with the BF16 projection; the fused CE + kernel then writes the per-token loss and (when ``needs_dz``) replaces the + logits with the softmax-minus-one-hot gradient dZ. Shared by the legacy + unchunked path (one full-batch call) and the chunked path (per chunk). + """ + rows = logits_chunk.shape[0] + torch.mm(input_chunk, weight.t(), out=logits_chunk) + # Counters must be reset before every partition reduction. + completion_count.zero_() + _launch_cutile( + logits_chunk.device, + (num_vocab_tiles * rows, 1, 1), + _fused_cross_entropy_dz_kernel, + ( + logits_chunk, + target_chunk, + loss_chunk, + loss_scale, + partial_max, + partial_sum, + completion_count, + vocab_size, + num_vocab_tiles, + ignore_index, + LOGITS_STATS_BLOCK_SIZE, + CE_BLOCK_SIZE, + stats_block_size, + needs_dz, + reduction_mean, + ), + ) + + def fused_linear_cross_entropy_forward( _input, weight, @@ -291,6 +417,15 @@ def fused_linear_cross_entropy_forward( return_token_accuracy=False, return_predicted_tokens=False, ): + """Legacy unchunked forward. + + Materializes the full ``[N, V]`` logits buffer, overwrites it with dZ, and + returns that dZ as the third element (``grad_logits``). The paired + :func:`fused_linear_cross_entropy_backward` consumes ``grad_logits`` to run + the dW matmul. This is the original low-level contract, retained for + compatibility; the autograd Function uses the bounded + :func:`chunked_fused_linear_cross_entropy_forward` instead. + """ _reject_unsupported( bias, ce_weight, @@ -305,47 +440,40 @@ def fused_linear_cross_entropy_forward( ) _validate_inputs(_input, weight, target, reduction) - tokens, hidden_size = _input.shape - vocab_size = weight.shape[0] - num_vocab_tiles = (vocab_size + LOGITS_STATS_BLOCK_SIZE - 1) // LOGITS_STATS_BLOCK_SIZE - stats_block_size = _next_power_of_2(num_vocab_tiles) - needs_dz = _input.requires_grad or weight.requires_grad + ( + tokens, + hidden_size, + vocab_size, + num_vocab_tiles, + stats_block_size, + needs_dz, + reduction_mean, + loss_scale, + gradient_scale, + ) = _ce_setup(_input, weight, target, reduction, ignore_index) + logits = torch.empty((tokens, vocab_size), dtype=torch.bfloat16, device=_input.device) loss_1d = torch.empty(tokens, dtype=torch.float32, device=_input.device) partial_max = torch.empty((tokens, num_vocab_tiles), dtype=torch.float32, device=_input.device) partial_sum = torch.empty_like(partial_max) completion_count = torch.zeros(tokens, dtype=torch.int32, device=_input.device) - target_mask = target != ignore_index - valid_target = ~target_mask | ((target >= 0) & (target < vocab_size)) - torch._assert_async(valid_target.all(), f"target values must be in [0, {vocab_size}) or equal ignore_index") - loss_scale = ( - target_mask.sum().clamp_min(1).to(torch.float32).reciprocal() - if reduction == "mean" - else torch.empty((), dtype=torch.float32, device=_input.device) - ) - torch.mm(_input, weight.t(), out=logits) - _launch_cutile( - logits.device, - (num_vocab_tiles * tokens, 1, 1), - _fused_cross_entropy_dz_kernel, - ( - logits, - target, - loss_1d, - loss_scale, - partial_max, - partial_sum, - completion_count, - vocab_size, - num_vocab_tiles, - ignore_index, - LOGITS_STATS_BLOCK_SIZE, - CE_BLOCK_SIZE, - stats_block_size, - needs_dz, - reduction == "mean", - ), + _project_and_dz( + _input, + weight, + logits, + target, + loss_1d, + loss_scale, + partial_max, + partial_sum, + completion_count, + vocab_size, + num_vocab_tiles, + stats_block_size, + ignore_index, + needs_dz, + reduction_mean, ) grad_input = None @@ -354,7 +482,8 @@ def fused_linear_cross_entropy_forward( torch.mm(logits, weight, out=grad_input) loss = loss_1d.sum() - gradient_scale = loss_scale if reduction == "mean" else torch.ones((), dtype=torch.float32, device=_input.device) + # Third element is the full [N, V] dZ (legacy contract): the paired backward + # runs the dW matmul from it. Scaling is deferred to backward. return loss, grad_input, logits if weight.requires_grad else None, gradient_scale @@ -367,19 +496,191 @@ def fused_linear_cross_entropy_backward( weight_shape: torch.Size, weight_dtype: torch.dtype, ): + """Legacy 7-argument backward paired with the unchunked forward. + + Applies ``upstream * gradient_scale`` to dX. dW is computed in FP32 by cuBLAS + (``dZ.T @ X`` accumulated into an FP32 buffer), then the combined + ``upstream * gradient_scale`` factor is applied in FP32 and the result cast to + ``weight_dtype`` exactly once by the fused scale-and-cast kernel. Casting + after (not before) the upstream multiply preserves small gradients that would + flush to zero if rounded to BF16 first. All operations are out-of-place, so + the saved tensors are not mutated across repeated backward calls. + """ total_scale = grad_output * gradient_scale + grad_input_out = None if grad_input is not None: - grad_input = grad_input * total_scale + grad_input_out = grad_input * total_scale grad_weight = None if grad_logits is not None: + dweight_accum = torch.empty(weight_shape, dtype=torch.float32, device=grad_logits.device) + _accumulate_dw_fp32(dweight_accum, grad_logits, _input, accumulate=False) + combined_scale = total_scale.to(torch.float32) grad_weight = torch.empty(weight_shape, dtype=weight_dtype, device=grad_logits.device) - _launch_matmul_tn(grad_logits, _input, grad_weight, total_scale) + _launch_scale_cast(dweight_accum, grad_weight, combined_scale) - return grad_input, grad_weight + return grad_input_out, grad_weight + + +def chunked_fused_linear_cross_entropy_forward( + _input, + weight, + target, + ce_weight=None, + bias=None, + ignore_index=-100, + lse_square_scale=0.0, + label_smoothing=0.0, + reduction="mean", + softcap=None, + return_z_loss=False, + accum_dtype=None, + use_token_scaling=False, + return_token_accuracy=False, + return_predicted_tokens=False, +): + """Token-chunked forward used by the autograd Function. + + Processes tokens in row chunks so the transient logits/dZ workspace is + ``O(min(N, C) * V)``, where the token-chunk size ``C`` comes from the shared + :func:`liger_kernel.ops.fused_linear_cross_entropy._get_chunk_size` helper -- + the EXACT same default geometry as the Triton FLCE backend, with no public + ``chunk_size`` knob. dX is written as the raw BF16 projection and dW is + accumulated in an FP32 buffer; NEITHER is scaled here. The returned third + element is the raw FP32 dW accumulator (``O(V*H)``), which + :func:`chunked_fused_linear_cross_entropy_backward` scales and casts to the + weight dtype exactly once -- deferring the BF16 cast preserves small + gradients that would flush to zero if cast per-chunk in the forward. + """ + _reject_unsupported( + bias, + ce_weight, + lse_square_scale, + label_smoothing, + softcap, + return_z_loss, + accum_dtype, + use_token_scaling, + return_token_accuracy, + return_predicted_tokens, + ) + _validate_inputs(_input, weight, target, reduction) + + ( + tokens, + hidden_size, + vocab_size, + num_vocab_tiles, + stats_block_size, + needs_dz, + reduction_mean, + loss_scale, + gradient_scale, + ) = _ce_setup(_input, weight, target, reduction, ignore_index) + + chunk = _get_chunk_size(tokens, hidden_size, vocab_size) + buffer_rows = min(tokens, chunk) + + # Reusable per-chunk logits/dZ buffer of shape [min(N, C), V] and the + # partitioned CE statistics/counter scratch of shape [min(N, C), num_partitions]. + logits = torch.empty((buffer_rows, vocab_size), dtype=torch.bfloat16, device=_input.device) + loss_1d = torch.empty(tokens, dtype=torch.float32, device=_input.device) + partial_max = torch.empty((buffer_rows, num_vocab_tiles), dtype=torch.float32, device=_input.device) + partial_sum = torch.empty_like(partial_max) + completion_count = torch.zeros(buffer_rows, dtype=torch.int32, device=_input.device) + + grad_input = None + if _input.requires_grad: + grad_input = torch.empty((tokens, hidden_size), dtype=_input.dtype, device=_input.device) + + # FP32 dW accumulator reused across chunks. It is retained UNSCALED for + # backward; the mean normalizer and upstream scalar are combined and the + # single BF16 cast is applied there, never per chunk. + dweight_accum = None + if weight.requires_grad: + dweight_accum = torch.empty((vocab_size, hidden_size), dtype=torch.float32, device=_input.device) + + for chunk_index, start in enumerate(range(0, tokens, chunk)): + stop = min(start + chunk, tokens) + rows = stop - start + logits_chunk = logits[:rows] + input_chunk = _input[start:stop] + + _project_and_dz( + input_chunk, + weight, + logits_chunk, + target[start:stop], + loss_1d[start:stop], + loss_scale, + partial_max[:rows], + partial_sum[:rows], + completion_count[:rows], + vocab_size, + num_vocab_tiles, + stats_block_size, + ignore_index, + needs_dz, + reduction_mean, + ) + + # logits_chunk now holds dZ. dX is scaled AFTER the low-precision GEMM + # (in backward), matching the baseline; store the raw projection here. + if grad_input is not None: + torch.mm(logits_chunk, weight, out=grad_input[start:stop]) + + # Accumulate raw dZ.T @ X into the FP32 dW buffer via cuBLAS: overwrite + # (mm) on the first chunk, add (addmm) on subsequent chunks. + if dweight_accum is not None: + _accumulate_dw_fp32(dweight_accum, logits_chunk, input_chunk, chunk_index > 0) + + loss = loss_1d.sum() + # Retain the RAW FP32 dW accumulator (scaled + cast in backward) so small + # gradients are not lost to an early BF16 cast. + return loss, grad_input, dweight_accum, gradient_scale + + +def chunked_fused_linear_cross_entropy_backward( + grad_output: torch.Tensor, + grad_input: Optional[torch.Tensor], + dweight_accum: Optional[torch.Tensor], + gradient_scale: torch.Tensor, + weight_dtype: torch.dtype, +): + """Backward for the chunked forward. + + dX is the raw BF16 projection, scaled out-of-place by + ``upstream * gradient_scale``. dW is the retained FP32 accumulator: the + combined ``upstream * gradient_scale`` factor is applied in FP32 and the + result cast to ``weight_dtype`` exactly once by a fused scale-and-cast + kernel (no O(V*H) FP32 temporary, and small gradients survive the cast). + Every output is a fresh tensor, so the saved gradients are never mutated -- + repeated backward with a different upstream (including ``1``) stays correct. + """ + grad_input_out = None + if grad_input is not None: + grad_input_out = grad_input * (grad_output * gradient_scale) + + grad_weight_out = None + if dweight_accum is not None: + combined_scale = (grad_output * gradient_scale).to(torch.float32) + grad_weight_out = torch.empty(dweight_accum.shape, dtype=weight_dtype, device=dweight_accum.device) + _launch_scale_cast(dweight_accum, grad_weight_out, combined_scale) + + return grad_input_out, grad_weight_out class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): + # This Function computes dX/dW during forward with token chunking. The + # transient logits/dZ workspace is O(min(N, C) * V) instead of the full + # O(N * V); the retained gradients add O(N*H + V*H). The token-chunk size C + # is the shared default Triton geometry (_get_chunk_size); there is no public + # chunk_size knob. + # ce_impl / ce_mode are accepted only as self-identity placeholders (see + # _validate_ce_dispatch); this Function does NOT dispatch to alternative CE + # implementations, so it does not advertise inner-impl dispatch. + supports_inner_impl_dispatch = False + @staticmethod @amp_custom_fwd def forward( @@ -399,8 +700,11 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): use_token_scaling=False, return_token_accuracy=False, return_predicted_tokens=False, + ce_impl=None, + ce_mode=None, ): - loss, grad_input, grad_logits, gradient_scale = fused_linear_cross_entropy_forward( + _validate_ce_dispatch(ce_impl, ce_mode) + loss, grad_input, dweight_accum, gradient_scale = chunked_fused_linear_cross_entropy_forward( _input=_input, weight=weight, target=target, @@ -418,13 +722,10 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): return_predicted_tokens=return_predicted_tokens, ) - ctx.save_for_backward( - grad_input, - grad_logits, - _input.detach() if weight.requires_grad else None, - gradient_scale, - ) - ctx.weight_shape = weight.shape + # Retain the raw BF16 dX, the raw FP32 dW accumulator, and the scale. + # Neither the full [N, V] dZ nor the input activations are saved. The + # target weight dtype is tracked so backward can cast dW once. + ctx.save_for_backward(grad_input, dweight_accum, gradient_scale) ctx.weight_dtype = weight.dtype return loss, None, None, None @@ -432,37 +733,39 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function): @amp_custom_bwd def backward(ctx, grad_output, grad_output2, grad_output3, grad_output4): del grad_output2, grad_output3, grad_output4 - grad_input, grad_logits, _input, gradient_scale = ctx.saved_tensors - grad_input, grad_weight = fused_linear_cross_entropy_backward( + grad_input, dweight_accum, gradient_scale = ctx.saved_tensors + grad_input, grad_weight = chunked_fused_linear_cross_entropy_backward( grad_output, grad_input, - grad_logits, - _input, + dweight_accum, gradient_scale, - ctx.weight_shape, ctx.weight_dtype, ) return ( grad_input, grad_weight, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, + None, # target + None, # bias + None, # ce_weight + None, # ignore_index + None, # lse_square_scale + None, # label_smoothing + None, # reduction + None, # softcap + None, # return_z_loss + None, # accum_dtype + None, # use_token_scaling + None, # return_token_accuracy + None, # return_predicted_tokens + None, # ce_impl + None, # ce_mode ) __all__ = [ "LigerFusedLinearCrossEntropyFunction", + "chunked_fused_linear_cross_entropy_backward", + "chunked_fused_linear_cross_entropy_forward", "fused_linear_cross_entropy_backward", "fused_linear_cross_entropy_forward", ] diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index 1655fcd3..7db7ac92 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py @@ -40,6 +40,19 @@ _TORCH_VERSION = Version(torch.__version__.split("+")[0]) _ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0") +def _get_chunk_size(BT, H, V): + """Token-chunk size shared by the Triton and cuTile FLCE backends. + + Widens the transient logits budget to ``_CHUNK_MEM_CONST x BT x H`` (C=1 is a + memory floor, not a perf target): ``inc_factor = cdiv(V, C*H)``, then the + chunk is the next power of two of ``cdiv(BT, inc_factor)``, clamped to ``BT`` + so a single chunk covers ``BT`` when the budget allows (and when ``V <= C*H``). + """ + inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H) + chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) # (BT + inc_factor - 1) // inc_factor + return min(chunk_size, BT) # a single chunk covers BT when the budget allows; never exceed BT + + def fused_linear_cross_entropy_forward( _input, weight, @@ -89,9 +102,7 @@ def fused_linear_cross_entropy_forward( V = weight.shape[0] # widen the transient logits budget to C x BT x H (C=1 is a memory floor, not a perf target) - inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H) - chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) # (BT + inc_factor - 1) // inc_factor - chunk_size = min(chunk_size, BT) # a single chunk covers BT when the budget allows; never exceed BT + chunk_size = _get_chunk_size(BT, H, V) num_chunks = triton.cdiv(BT, chunk_size) # (BT + chunk_size - 1) // chunk_size grad_input = torch.zeros_like(_input, device=device) diff --git a/test/cutile/test_fused_linear_cross_entropy_sm90.py b/test/cutile/test_fused_linear_cross_entropy_sm90.py index d4b40bce..d1a74014 100644 --- a/test/cutile/test_fused_linear_cross_entropy_sm90.py +++ b/test/cutile/test_fused_linear_cross_entropy_sm90.py @@ -2,25 +2,67 @@ import pytest import torch import torch.nn.functional as F +from torch.utils._python_dispatch import TorchDispatchMode + from test.utils import assert_verbose_allclose from test.utils import set_seed pytest.importorskip("cuda.tile") +import liger_kernel.ops.cutile.ops.fused_linear_cross_entropy as flce_mod +import liger_kernel.ops.fused_linear_cross_entropy as triton_flce_mod + from liger_kernel.ops.cutile.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction +from liger_kernel.ops.cutile.ops.fused_linear_cross_entropy import chunked_fused_linear_cross_entropy_backward +from liger_kernel.ops.cutile.ops.fused_linear_cross_entropy import chunked_fused_linear_cross_entropy_forward +from liger_kernel.ops.cutile.ops.fused_linear_cross_entropy import fused_linear_cross_entropy_backward +from liger_kernel.ops.cutile.ops.fused_linear_cross_entropy import fused_linear_cross_entropy_forward +from liger_kernel.ops.fused_linear_cross_entropy import _get_chunk_size + +# The cuTile FLCE kernels support both Hopper (SM90) and Blackwell (SM100). +_SUPPORTED_CAPABILITIES = ((9, 0), (10, 0)) pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="cuTile FLCE requires CUDA"), pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0), - reason="cuTile FLCE requires Hopper", + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in _SUPPORTED_CAPABILITIES, + reason="cuTile FLCE requires Hopper (SM90) or Blackwell (SM100)", ), ] set_seed() +def _apply(x, weight, target, reduction="mean", ignore_index=-100): + """Invoke the Function with the 17-argument module-style contract. + + There is NO public ``chunk_size`` argument: the token-chunk geometry is the + shared default Triton policy (:func:`_get_chunk_size`). ``ce_impl`` / ``ce_mode`` + are self-identity placeholders and default to ``None``. + """ + return LigerFusedLinearCrossEntropyFunction.apply( + x, + weight, + target, + None, # bias + None, # ce_weight + ignore_index, + 0.0, # lse_square_scale + 0.0, # label_smoothing + reduction, + None, # softcap + False, # return_z_loss + None, # accum_dtype + False, # use_token_scaling + False, # return_token_accuracy + False, # return_predicted_tokens + None, # ce_impl + None, # ce_mode + ) + + def _reference(x, weight, target, ignore_index, reduction): + # Baseline: BF16 F.linear, then FP32 cross entropy (matches the kernel's math). return F.cross_entropy( F.linear(x, weight).float(), target, @@ -29,10 +71,106 @@ def _reference(x, weight, target, ignore_index, reduction): ) -@pytest.mark.parametrize("shape", [(128, 128, 256), (129, 96, 384), (32, 64, 8192)]) +def _tolerances(reduction): + atol = 5e-3 if reduction == "mean" else 5e-2 + return atol, 5e-2 + + +def _relative_l2(actual, expected): + return (torch.norm((actual - expected).float()) / torch.norm(expected.float()).clamp_min(1e-12)).item() + + +def _ref_dw_fp32(a, b): + """Independent FP32 ``a @ b`` mirroring the production dW accumulation. + + Tracks ``flce_mod._ADDMM_SUPPORTS_OUT_DTYPE``: on torch>=2.8 the BF16 operands + drive the ``out_dtype=torch.float32`` overload directly; on older torch (no + overload) the bounded BF16 operands are upcast to FP32 and fed to the plain + ``mm``. Both are numerically identical -- a BF16*BF16 product is exact in FP32 + -- so callers keep atol/rtol at 0. + """ + if flce_mod._ADDMM_SUPPORTS_OUT_DTYPE: + return torch.mm(a, b, out_dtype=torch.float32) + return torch.mm(a.to(torch.float32), b.to(torch.float32)) + + +def _expected_tiling(tokens, hidden_size, vocab_size): + """Row counts of every token chunk under the shared default geometry.""" + chunk = _get_chunk_size(tokens, hidden_size, vocab_size) + return [min(chunk, tokens - start) for start in range(0, tokens, chunk)] + + +# --------------------------------------------------------------------------- +# End-to-end correctness against the BF16-linear / FP32-CE reference on shapes +# that resolve to a SINGLE token chunk under the shared default geometry +# (V <= 16*H). One chunk means the projection GEMM tiling matches the reference, +# so the original tight elementwise tolerances apply. Shapes exercise a range of +# H (256-aligned and non-aligned), odd V, and > 4096 vocab (multiple partitions). +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "shape", + [ + (128, 128, 256), # H=128 + (129, 96, 384), # H=96, tail token + (32, 512, 8192), # V > 4096 -> multiple vocab partitions, still one chunk + (300, 256, 512), # 256-aligned H + (200, 128, 257), # odd vocab size + ], +) @pytest.mark.parametrize("reduction", ["mean", "sum"]) def test_correctness(shape, reduction): tokens, hidden_size, vocab_size = shape + # These shapes must be single-chunk so the tight tolerances below are valid. + assert len(_expected_tiling(tokens, hidden_size, vocab_size)) == 1 + ignore_index = -100 + upstream = torch.tensor(0.7, device="cuda") + x_data = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + weight_data = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) + target = torch.randint(vocab_size, (tokens,), device="cuda") + target[: max(1, tokens // 8)] = ignore_index # uneven ignored tokens + + x_ref = x_data.clone().requires_grad_(True) + weight_ref = weight_data.clone().requires_grad_(True) + loss_ref = _reference(x_ref, weight_ref, target, ignore_index, reduction) + loss_ref.backward(upstream) + + x = x_data.clone().requires_grad_(True) + weight = weight_data.clone().requires_grad_(True) + loss, _, _, _ = _apply(x, weight, target, reduction, ignore_index) + loss.backward(upstream) + + atol, rtol = _tolerances(reduction) + assert_verbose_allclose(loss_ref, loss, atol=atol, rtol=rtol) + assert_verbose_allclose(x_ref.grad, x.grad, atol=atol, rtol=rtol) + assert_verbose_allclose(weight_ref.grad, weight.grad, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# Multi-chunk + tail correctness against the reference on shapes whose DEFAULT +# geometry (V > 16*H) is genuinely multi-chunk -- no chunk_size knob is used. +# +# Multi-chunk re-tiles the projection GEMM, so the chunked BF16 logits differ +# from the full-batch reference by cuBLAS tiling variance. Under `mean` this is +# normalized away (tight elementwise tolerance holds); under `sum` the +# unnormalized near-zero gradient elements can diverge by O(0.1) purely from GEMM +# tiling, so correctness is asserted as a small aggregate relative-L2 error (a +# structural chunking bug would blow the whole tensor up, not just near-zeros). +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "shape", + [ + (513, 64, 4096), # inc=4 -> chunk 256 -> [256, 256, 1] + (200, 96, 8192), # H=96 (non-aligned), multi vocab partition + (250, 128, 8192), # V > 4096 -> multiple vocab partitions + (300, 64, 4096), # 3 chunks + (129, 16, 8192), # inc=32 -> chunk 8 -> 17 tiny chunks + tail + ], +) +@pytest.mark.parametrize("reduction", ["mean", "sum"]) +def test_multichunk_correctness(shape, reduction): + tokens, hidden_size, vocab_size = shape + # These shapes must be genuinely multi-chunk under the default geometry. + assert len(_expected_tiling(tokens, hidden_size, vocab_size)) >= 2 ignore_index = -100 upstream = torch.tensor(0.7, device="cuda") x_data = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) @@ -47,55 +185,205 @@ def test_correctness(shape, reduction): x = x_data.clone().requires_grad_(True) weight = weight_data.clone().requires_grad_(True) - loss, _, _, _ = LigerFusedLinearCrossEntropyFunction.apply( - x, - weight, - target, - None, - None, - ignore_index, - 0.0, - 0.0, - reduction, - ) + loss, _, _, _ = _apply(x, weight, target, reduction, ignore_index) loss.backward(upstream) - atol = 5e-3 if reduction == "mean" else 5e-2 - rtol = 5e-2 - assert_verbose_allclose(loss_ref, loss, atol=atol, rtol=rtol) - assert_verbose_allclose(x_ref.grad, x.grad, atol=atol, rtol=rtol) - assert_verbose_allclose(weight_ref.grad, weight.grad, atol=atol, rtol=rtol) + if reduction == "mean": + atol, rtol = _tolerances(reduction) + assert_verbose_allclose(loss_ref, loss, atol=atol, rtol=rtol) + assert_verbose_allclose(x_ref.grad, x.grad, atol=atol, rtol=rtol) + assert_verbose_allclose(weight_ref.grad, weight.grad, atol=atol, rtol=rtol) + else: + assert abs(loss.item() - loss_ref.item()) / abs(loss_ref.item()) < 2e-2 + assert _relative_l2(x.grad, x_ref.grad) < 2e-2 + assert _relative_l2(weight.grad, weight_ref.grad) < 2e-2 -@pytest.mark.parametrize("requires_input_grad, requires_weight_grad", [(True, False), (False, True)]) -def test_independent_gradient_requirements(requires_input_grad, requires_weight_grad): - x = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=requires_input_grad) - weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=requires_weight_grad) - target = torch.randint(256, (128,), device="cuda") +# --------------------------------------------------------------------------- +# Chunking invariance: a forced multi-chunk run (the documented C=1 budget floor +# monkeypatched onto the shared _CHUNK_MEM_CONST) must match the single-chunk run +# the same shape gets at the default C=16 budget. Loss is partition-invariant; +# gradients differ only by GEMM/addmm reduction order (tight tolerance). +# --------------------------------------------------------------------------- +def test_chunking_matches_single_chunk(monkeypatch): + tokens, hidden_size, vocab_size = 130, 128, 512 + ignore_index = -100 + upstream = torch.tensor(1.3, device="cuda") + x_data = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + weight_data = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) + target = torch.randint(vocab_size, (tokens,), device="cuda") + target[:9] = ignore_index - loss = LigerFusedLinearCrossEntropyFunction.apply(x, weight, target)[0] + # Default budget is a single chunk for this shape... + assert len(_expected_tiling(tokens, hidden_size, vocab_size)) == 1 + + def run(): + x = x_data.clone().requires_grad_(True) + weight = weight_data.clone().requires_grad_(True) + loss = _apply(x, weight, target, "mean", ignore_index)[0] + loss.backward(upstream) + return loss.detach(), x.grad, weight.grad + + loss_full, gx_full, gw_full = run() + + # ...forcing the C=1 memory floor re-tiles the SAME shape into many chunks. + monkeypatch.setattr(triton_flce_mod, "_CHUNK_MEM_CONST", 1) + assert len(_expected_tiling(tokens, hidden_size, vocab_size)) >= 2 + loss_c, gx_c, gw_c = run() + + assert_verbose_allclose(loss_full, loss_c, atol=1e-4, rtol=1e-3) + assert_verbose_allclose(gx_full, gx_c, atol=1e-3, rtol=1e-2) + assert_verbose_allclose(gw_full, gw_c, atol=1e-3, rtol=1e-2) + + +# --------------------------------------------------------------------------- +# Global-mean normalization: a middle chunk that is entirely ignored still +# normalizes by the GLOBAL non-ignored count. Shape (192,64,4096) is 3 chunks of +# 64 under the default geometry; the middle chunk [64:128] is fully ignored. +# --------------------------------------------------------------------------- +def test_global_mean_with_all_ignored_middle_chunk(): + tokens, hidden_size, vocab_size = 192, 64, 4096 + ignore_index = -100 + assert _expected_tiling(tokens, hidden_size, vocab_size) == [64, 64, 64] + x_data = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + weight_data = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) + target = torch.randint(vocab_size, (tokens,), device="cuda") + target[64:128] = ignore_index # middle chunk fully ignored + + x_ref = x_data.clone().requires_grad_(True) + weight_ref = weight_data.clone().requires_grad_(True) + loss_ref = _reference(x_ref, weight_ref, target, ignore_index, "mean") + loss_ref.backward() + + x = x_data.clone().requires_grad_(True) + weight = weight_data.clone().requires_grad_(True) + loss = _apply(x, weight, target, "mean", ignore_index)[0] loss.backward() - assert (x.grad is not None) == requires_input_grad - assert (weight.grad is not None) == requires_weight_grad + assert_verbose_allclose(loss_ref, loss, atol=5e-3, rtol=5e-2) + assert_verbose_allclose(x_ref.grad, x.grad, atol=5e-3, rtol=5e-2) + assert_verbose_allclose(weight_ref.grad, weight.grad, atol=5e-3, rtol=5e-2) -def test_rejects_unsupported_features(): - x = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) - weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) - target = torch.randint(256, (128,), device="cuda") - bias = torch.zeros(256, device="cuda", dtype=torch.bfloat16) +def test_all_ignored_entire_batch(): + # Multi-chunk shape so the all-ignored path is exercised across chunks. + tokens, hidden_size, vocab_size = 96, 64, 4096 + ignore_index = -100 + assert len(_expected_tiling(tokens, hidden_size, vocab_size)) >= 2 + x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.full((tokens,), ignore_index, device="cuda") - with pytest.raises(NotImplementedError, match="bias"): - LigerFusedLinearCrossEntropyFunction.apply(x, weight, target, bias) + loss = _apply(x, weight, target, "mean", ignore_index)[0] + loss.backward() + + assert loss.item() == 0.0 + assert torch.count_nonzero(x.grad) == 0 + assert torch.count_nonzero(weight.grad) == 0 + + +# --------------------------------------------------------------------------- +# Independent gradient requirements: forward-only / input-only / weight-only / +# both, on a multi-chunk shape. Chunked forward must honor each combination. +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "requires_input_grad, requires_weight_grad", + [(False, False), (True, False), (False, True), (True, True)], +) +def test_independent_gradient_requirements(requires_input_grad, requires_weight_grad): + torch.manual_seed(0) + tokens, hidden_size, vocab_size = 200, 64, 4096 + ignore_index = -100 + upstream = torch.tensor(-0.7, device="cuda") + tiling = _expected_tiling(tokens, hidden_size, vocab_size) + assert len(tiling) >= 2 and tiling[-1] == 8 # multi-chunk with an 8-token tail + + x_data = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + weight_data = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) + target = torch.randint(vocab_size, (tokens,), device="cuda") + # Ignored targets INSIDE a chunk exercise the global-mean normalizer. + target[20:37] = ignore_index + target[::13] = ignore_index + + x = x_data.clone().requires_grad_(requires_input_grad) + weight = weight_data.clone().requires_grad_(requires_weight_grad) + loss = _apply(x, weight, target, "mean", ignore_index)[0] + assert torch.isfinite(loss) + + # Same-chunk oracle on detached leaves -- the loss and each requested grad + # match tightly regardless of which gradients were requested. + loss_oracle, dx_oracle, dw_oracle = _same_chunk_oracle(x_data, weight_data, target, "mean", upstream, ignore_index) + assert_verbose_allclose(loss, loss_oracle, atol=1e-2, rtol=1e-3) + + if requires_input_grad or requires_weight_grad: + loss.backward(upstream) + assert (x.grad is not None) == requires_input_grad + assert (weight.grad is not None) == requires_weight_grad + if requires_input_grad: + assert_verbose_allclose(x.grad, dx_oracle, atol=2e-3, rtol=1e-3) + if requires_weight_grad: + assert_verbose_allclose(weight.grad, dw_oracle, atol=2e-3, rtol=1e-3) + else: + assert loss.grad_fn is None # no autograd graph when neither operand needs grad + + # Low-level chunked forward on leaves carrying the ORIGINAL requires flags: + # only the requested gradient buffers are materialized (with the right + # dtype/shape); the unrequested ones stay None. Run under no_grad on the + # leaves themselves -- the forward keys off tensor.requires_grad directly. + xl = x_data.clone().requires_grad_(requires_input_grad) + wl = weight_data.clone().requires_grad_(requires_weight_grad) + with torch.no_grad(): + ll_loss, grad_input, dweight_accum, gradient_scale = chunked_fused_linear_cross_entropy_forward( + xl, wl, target, reduction="mean", ignore_index=ignore_index + ) + assert_verbose_allclose(ll_loss, loss_oracle, atol=1e-2, rtol=1e-3) + if requires_input_grad: + assert grad_input is not None + assert grad_input.dtype == torch.bfloat16 + assert tuple(grad_input.shape) == (tokens, hidden_size) + else: + assert grad_input is None + if requires_weight_grad: + assert dweight_accum is not None + assert dweight_accum.dtype == torch.float32 + assert tuple(dweight_accum.shape) == (vocab_size, hidden_size) + else: + assert dweight_accum is None + assert gradient_scale.ndim == 0 + + +# --------------------------------------------------------------------------- +# Repeated backward with a *different* non-unit upstream: the retained gradients +# must not be mutated, and gradients must scale linearly. +# --------------------------------------------------------------------------- +def test_repeated_backward_different_upstream(): + tokens, hidden_size, vocab_size = 200, 64, 4096 + assert len(_expected_tiling(tokens, hidden_size, vocab_size)) >= 2 + x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab_size, (tokens,), device="cuda") + target[:10] = -100 + loss = _apply(x, weight, target, "mean", -100)[0] + + loss.backward(torch.tensor(1.0, device="cuda"), retain_graph=True) + gx1 = x.grad.clone() + gw1 = weight.grad.clone() + x.grad = None + weight.grad = None + loss.backward(torch.tensor(2.5, device="cuda"), retain_graph=True) + + assert_verbose_allclose(x.grad, 2.5 * gx1, atol=1e-3, rtol=1e-3) + assert_verbose_allclose(weight.grad, 2.5 * gw1, atol=1e-3, rtol=1e-3) def test_retained_backward_preserves_saved_gradient(): - x = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) - weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) - target = torch.randint(256, (128,), device="cuda") + tokens, hidden_size, vocab_size = 128, 64, 4096 + assert len(_expected_tiling(tokens, hidden_size, vocab_size)) >= 2 + x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab_size, (tokens,), device="cuda") upstream = torch.tensor(0.7, device="cuda") - loss = LigerFusedLinearCrossEntropyFunction.apply(x, weight, target)[0] + loss = _apply(x, weight, target, "mean", -100)[0] loss.backward(upstream, retain_graph=True) first_x_grad = x.grad.clone() @@ -108,18 +396,298 @@ def test_retained_backward_preserves_saved_gradient(): assert_verbose_allclose(first_weight_grad, weight.grad, atol=0.0, rtol=0.0) +# --------------------------------------------------------------------------- +# Memory contract: the forward retains only [N, H] dX and [V, H] dW (plus the +# scalar scale) -- never a full [N, V] logits/dZ tensor -- and the reusable +# logits buffer keeps peak memory below the full [N, V] materialization. +# --------------------------------------------------------------------------- +def test_retains_only_final_gradients_no_full_logits(): + tokens, hidden_size, vocab_size = 256, 128, 4096 + x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab_size, (tokens,), device="cuda") + + with torch.no_grad(): + _, grad_input, grad_weight, gradient_scale = chunked_fused_linear_cross_entropy_forward(x, weight, target) + + assert tuple(grad_input.shape) == (tokens, hidden_size) + assert tuple(grad_weight.shape) == (vocab_size, hidden_size) + # The retained dW accumulator is FP32 (scaled + cast to weight dtype in backward). + assert grad_weight.dtype == torch.float32 + assert gradient_scale.ndim == 0 + # No retained tensor has the full [N, V] footprint. + assert grad_input.numel() < tokens * vocab_size + assert grad_weight.numel() < tokens * vocab_size + + +def test_reusable_logits_buffer_bounds_peak_memory(): + # V > 16*H so the default geometry is genuinely multi-chunk (buffer_rows < N). + tokens, hidden_size, vocab_size = 4096, 128, 8192 + tiling = _expected_tiling(tokens, hidden_size, vocab_size) + assert len(tiling) >= 2 and tiling[0] < tokens + x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab_size, (tokens,), device="cuda") + + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + loss = _apply(x, weight, target, "mean", -100)[0] + loss.backward() + torch.cuda.synchronize() + peak_delta = torch.cuda.max_memory_allocated() - base + + full_logits_bytes = tokens * vocab_size * 2 # bf16 [N, V] + # A full [N, V] materialization would dominate; the chunked buffer must stay + # well under it (chunked logits is buffer_rows * V * 2 bytes). + assert peak_delta < full_logits_bytes + + +# --------------------------------------------------------------------------- +# Shared token-chunk geometry: cuTile reuses the EXACT Triton default policy. +# The imported helper is the same object the Triton backend uses, and it honors +# inc=cdiv(V, 16*H), chunk=next_pow2(cdiv(N, inc)) clamped to N (never a public +# chunk_size knob, never the removed CuTe <=1024 heuristic). +# --------------------------------------------------------------------------- +def test_cutile_reuses_triton_chunk_helper(): + # cuTile imports the Triton helper -- no second constant, no divergent policy. + assert flce_mod._get_chunk_size is triton_flce_mod._get_chunk_size + assert not hasattr(flce_mod, "_resolve_chunk_size") + assert not hasattr(flce_mod, "_largest_power_of_2_le") + + +@pytest.mark.parametrize( + "tokens, hidden_size, vocab_size, expected_chunk, expected_num_chunks", + [ + (8192, 4096, 128256, 4096, 2), # llama_3 reference: 4096 rows, 2 chunks (NOT 1024) + (200, 512, 256, 200, 1), # V <= 16*H -> single chunk == N + (128, 128, 256, 128, 1), # V <= 16*H -> single chunk == N + (513, 64, 4096, 256, 3), # odd N tail: [256, 256, 1] + (129, 16, 8192, 8, 17), # next_pow2 clamp -> tiny 8-row chunks + tail + (1, 16, 8192, 1, 1), # one token + ], +) +def test_chunk_geometry_table(tokens, hidden_size, vocab_size, expected_chunk, expected_num_chunks): + chunk = _get_chunk_size(tokens, hidden_size, vocab_size) + assert chunk == expected_chunk + tiling = _expected_tiling(tokens, hidden_size, vocab_size) + assert len(tiling) == expected_num_chunks + # Only the final chunk may be a short tail; all others are full width. + assert tiling[:-1] == [expected_chunk] * (expected_num_chunks - 1) + assert tiling[-1] == tokens - expected_chunk * (expected_num_chunks - 1) + + +class _ProjectionTracer(TorchDispatchMode): + """Record the row count (M) of every projection GEMM (arg1 == weight.t()).""" + + def __init__(self, hidden_size, vocab_size): + self._hv = (hidden_size, vocab_size) + self.rows = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + out = func(*args, **kwargs) + if str(func).startswith("aten.mm") and len(args) >= 2 and tuple(args[1].shape) == self._hv: + self.rows.append(int(args[0].shape[0])) + return out + + +@pytest.mark.parametrize( + "shape", + [ + (250, 128, 8192), # [64, 64, 64, 58] + (513, 64, 4096), # [256, 256, 1] + (200, 512, 256), # single chunk [200] + ], +) +def test_projection_tiling_matches_triton(shape): + """The cuTile projection GEMM tiling equals what the Triton backend uses. + + Both backends chunk tokens by the shared :func:`_get_chunk_size`. Trace the + per-chunk projection ``x_chunk @ weight.t()`` M-dims from each backend and + assert they are identical to each other and to the expected tiling. + """ + tokens, hidden_size, vocab_size = shape + x = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab_size, (tokens,), device="cuda") + expected = _expected_tiling(tokens, hidden_size, vocab_size) + + with torch.no_grad(), _ProjectionTracer(hidden_size, vocab_size) as ct_trace: + chunked_fused_linear_cross_entropy_forward(x, weight, target, reduction="mean") + + with torch.no_grad(), _ProjectionTracer(hidden_size, vocab_size) as tr_trace: + triton_flce_mod.fused_linear_cross_entropy_forward(x, weight, target, reduction="mean") + + assert ct_trace.rows == expected + assert tr_trace.rows == expected + assert ct_trace.rows == tr_trace.rows + + +# --------------------------------------------------------------------------- +# Argument arities: the Function accepts the direct 3-arg legacy apply, the +# 15-arg functional contract, and the 17-arg module contract (ce_impl/ce_mode +# placeholders). There is no chunk_size parameter and no supports_chunk_size +# marker; ce_impl/ce_mode do not enable real inner dispatch. +# --------------------------------------------------------------------------- +def test_no_chunk_size_public_api(): + import inspect + + assert not hasattr(LigerFusedLinearCrossEntropyFunction, "supports_chunk_size") + assert LigerFusedLinearCrossEntropyFunction.supports_inner_impl_dispatch is False + params = list(inspect.signature(LigerFusedLinearCrossEntropyFunction.forward).parameters) + assert "chunk_size" not in params + # ctx + 17 declared args (15 base + ce_impl + ce_mode). + assert len(params) == 18 + + +def test_apply_arg_arities(): + x_data = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16) + weight_data = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16) + target = torch.randint(256, (64,), device="cuda") + + # 3-arg legacy apply (minimal positional contract). + loss3 = LigerFusedLinearCrossEntropyFunction.apply( + x_data.clone().requires_grad_(True), weight_data.clone().requires_grad_(True), target + )[0] + assert torch.isfinite(loss3) + + # 15-arg functional contract (what liger_fused_linear_cross_entropy passes + # when supports_inner_impl_dispatch is False). + loss15 = LigerFusedLinearCrossEntropyFunction.apply( + x_data.clone().requires_grad_(True), + weight_data.clone().requires_grad_(True), + target, + None, + None, + -100, + 0.0, + 0.0, + "mean", + None, + False, + None, + False, + False, + False, + )[0] + assert torch.isfinite(loss15) + + # 17-arg module contract (what the nn.Module passes: + ce_impl=None, ce_mode=None). + x = x_data.clone().requires_grad_(True) + weight = weight_data.clone().requires_grad_(True) + loss17 = _apply(x, weight, target, "mean", -100)[0] + loss17.backward() + assert torch.isfinite(loss17) + assert x.grad is not None and weight.grad is not None + + +# --------------------------------------------------------------------------- +# Feature / dispatch rejection. +# --------------------------------------------------------------------------- +def test_rejects_unsupported_features(): + x = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(256, (128,), device="cuda") + bias = torch.zeros(256, device="cuda", dtype=torch.bfloat16) + + with pytest.raises(NotImplementedError, match="bias"): + LigerFusedLinearCrossEntropyFunction.apply(x, weight, target, bias) + + +def test_accum_dtype_fp32_accepted_others_rejected(): + x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(256, (64,), device="cuda") + + # FP32 accumulation is explicitly allowed (it is the default policy). + loss = LigerFusedLinearCrossEntropyFunction.apply( + x, weight, target, None, None, -100, 0.0, 0.0, "mean", None, False, torch.float32 + )[0] + assert torch.isfinite(loss) + + with pytest.raises(NotImplementedError, match="accum"): + LigerFusedLinearCrossEntropyFunction.apply( + x, weight, target, None, None, -100, 0.0, 0.0, "mean", None, False, torch.float16 + ) + + +def _apply_with_dispatch(x, weight, target, ce_impl, ce_mode): + return LigerFusedLinearCrossEntropyFunction.apply( + x, + weight, + target, + None, + None, + -100, + 0.0, + 0.0, + "mean", + None, + False, + None, + False, + False, + False, + ce_impl, + ce_mode, + ) + + +@pytest.mark.parametrize("ce_impl", [None, "cutile", "nvidia-cutile"]) +def test_ce_impl_self_identity_accepted(ce_impl): + x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(256, (64,), device="cuda") + loss = _apply_with_dispatch(x, weight, target, ce_impl=ce_impl, ce_mode=None)[0] + assert torch.isfinite(loss) + + +def test_ce_impl_and_mode_incompatible_rejected(): + x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(256, (64,), device="cuda") + + with pytest.raises(NotImplementedError): + _apply_with_dispatch(x, weight, target, ce_impl="triton", ce_mode=None) + with pytest.raises(NotImplementedError): + _apply_with_dispatch(x, weight, target, ce_impl=None, ce_mode="chunked") + + +# --------------------------------------------------------------------------- +# Input validation: empty dims and CPU tensors are caught before any CUDA +# capability query. +# --------------------------------------------------------------------------- +def test_rejects_empty_tensors(): + weight = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16) + empty_x = torch.randn(0, 128, device="cuda", dtype=torch.bfloat16) + empty_target = torch.randint(256, (0,), device="cuda") + with pytest.raises(ValueError, match="empty"): + with torch.no_grad(): + fused_linear_cross_entropy_forward(empty_x, weight, empty_target) + + +def test_rejects_cpu_tensors(): + x = torch.randn(64, 128, dtype=torch.bfloat16) + weight = torch.randn(256, 128, dtype=torch.bfloat16) + target = torch.randint(256, (64,)) + with pytest.raises(RuntimeError, match="CUDA"): + with torch.no_grad(): + fused_linear_cross_entropy_forward(x, weight, target) + + def test_large_token_grid(): tokens = 65536 x = torch.randn(tokens, 16, device="cuda", dtype=torch.bfloat16) weight = torch.randn(128, 16, device="cuda", dtype=torch.bfloat16) target = torch.randint(128, (tokens,), device="cuda") - loss = LigerFusedLinearCrossEntropyFunction.apply(x, weight, target)[0] + loss = _apply(x, weight, target)[0] assert torch.isfinite(loss) -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two Hopper GPUs") +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two GPUs") def test_noncurrent_cuda_device(): torch.cuda.set_device(0) device = torch.device("cuda:1") @@ -127,9 +695,380 @@ def test_noncurrent_cuda_device(): weight = torch.randn(256, 128, device=device, dtype=torch.bfloat16, requires_grad=True) target = torch.randint(256, (128,), device=device) - loss = LigerFusedLinearCrossEntropyFunction.apply(x, weight, target)[0] + loss = _apply(x, weight, target)[0] loss.backward() assert torch.isfinite(loss) assert x.grad is not None assert weight.grad is not None + + +# --------------------------------------------------------------------------- +# dW small-gradient preservation: the FP32 dW accumulator is retained and the +# combined upstream/mean scale is applied in FP32 BEFORE the single BF16 cast. +# An early cast in the forward (before the upstream scale) flushes recoverable +# small gradients to zero. +# --------------------------------------------------------------------------- +def test_small_dw_gradient_survives_bf16_cast(): + # dW element magnitude ~ (1/V) * 1e-37 is below the BF16 subnormal floor, + # but after the go=65536 upstream scale it is representable again. + x = torch.full((1, 64), 1e-37, dtype=torch.bfloat16, device="cuda", requires_grad=True) + weight = torch.zeros((4096, 64), dtype=torch.bfloat16, device="cuda", requires_grad=True) + target = torch.tensor([0], device="cuda") + upstream = torch.tensor(65536.0, device="cuda") + + loss = _apply(x, weight, target, "mean", -100)[0] + loss.backward(upstream) + + # Independent FP32 Torch reference: dZ = softmax(0) with the target row + # decremented; the raw dW is the FP32 cuBLAS product dZ.T @ X, then the + # combined (upstream * mean-normalizer, here 1/1) scale is applied in FP32 + # and cast to BF16 exactly once -- matching the retained-accumulator path. + dz = torch.full((1, 4096), 1.0 / 4096, dtype=torch.bfloat16, device="cuda") + dz[0, 0] -= 1.0 + dw_fp32 = _ref_dw_fp32(dz.t(), x.detach()) # FP32 dW = dZ.T @ X + ref = (dw_fp32 * upstream.to(torch.float32)).to(torch.bfloat16) # FP32 scale, single BF16 cast + + assert weight.grad[1, 0].item() != 0.0 + assert_verbose_allclose(weight.grad, ref, atol=0.0, rtol=0.0) + + +# --------------------------------------------------------------------------- +# Strict SAME-CHUNK oracle: reproduce the kernel math per chunk with the exact +# same GEMM tiling, DERIVED from the shared default geometry (not a chunk_size +# argument), then assert tight per-element tolerances. A structural chunking / +# scaling bug would move whole elements far beyond these tolerances. An +# independent full-batch PyTorch anchor is retained too. +# --------------------------------------------------------------------------- +def _same_chunk_oracle(x_data, weight_data, target, reduction, upstream, ignore_index=-100): + n_tokens, hidden = x_data.shape + vocab = weight_data.shape[0] + device = x_data.device + # Derive the chunk size from the shared policy -- identical to the kernel. + chunk_size = _get_chunk_size(n_tokens, hidden, vocab) + valid = target != ignore_index + global_scale = 1.0 / valid.sum().clamp_min(1).float() if reduction == "mean" else torch.tensor(1.0, device=device) + loss = torch.zeros((), dtype=torch.float32, device=device) + dx = torch.zeros((n_tokens, hidden), dtype=torch.bfloat16, device=device) + dw_fp32 = torch.zeros((vocab, hidden), dtype=torch.float32, device=device) + for start in range(0, n_tokens, chunk_size): + stop = min(start + chunk_size, n_tokens) + xc = x_data[start:stop] + tc = target[start:stop] + z = torch.mm(xc, weight_data.t()) # BF16 projection, same tiling as the kernel + zf = z.float() + ce = F.cross_entropy(zf, tc, ignore_index=ignore_index, reduction="none") + loss += (ce * global_scale).sum() + dz = torch.softmax(zf, dim=-1) + rows = torch.arange(stop - start, device=device) + valid_c = tc != ignore_index + dz[rows[valid_c], tc[valid_c]] -= 1.0 + dz[~valid_c] = 0.0 + dz = dz.to(torch.bfloat16) + dx[start:stop] = torch.mm(dz, weight_data) # raw BF16 dX + dw_fp32.add_(_ref_dw_fp32(dz.t(), xc)) # FP32-accumulated dW + combined = upstream * global_scale + dx_final = dx * combined # BF16 raw dX scaled by upstream * mean-normalizer + dw_final = (dw_fp32 * combined).to(torch.bfloat16) # FP32 scale, single BF16 cast + return loss, dx_final, dw_final + + +@pytest.mark.parametrize("seed", [7, 314]) +@pytest.mark.parametrize("hidden", [96, 256]) # non-aligned + 256-aligned H tails +@pytest.mark.parametrize("reduction", ["mean", "sum"]) +# force_fallback drives the pre-2.8 dW path (plain mm/addmm on FP32-upcast BF16 +# operands). The oracle's ``_ref_dw_fp32`` tracks the same flag, so both paths +# stay bit-identical and the SAME tight tolerances apply -- forced multi-chunk +# numerical coverage of the older-torch accumulation with no tolerance widening. +@pytest.mark.parametrize("force_fallback", [False, True]) +def test_same_chunk_oracle_regression(seed, hidden, reduction, force_fallback, monkeypatch): + if force_fallback: + monkeypatch.setattr(flce_mod, "_ADDMM_SUPPORTS_OUT_DTYPE", False) + torch.manual_seed(seed) + n_tokens, vocab = 65, 4113 # odd V, tail chunk under the default geometry + ignore_index = -100 + upstream = torch.tensor(-0.7, device="cuda") + x_data = torch.randn(n_tokens, hidden, device="cuda", dtype=torch.bfloat16) + weight_data = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16) + target = torch.randint(vocab, (n_tokens,), device="cuda") + target[17:34] = ignore_index + target[::7] = ignore_index + # These shapes are multi-chunk with a tail under the shared geometry. + assert len(_expected_tiling(n_tokens, hidden, vocab)) >= 2 + + x = x_data.clone().requires_grad_(True) + weight = weight_data.clone().requires_grad_(True) + loss = _apply(x, weight, target, reduction, ignore_index)[0] + loss.backward(upstream) + + loss_oracle, dx_oracle, dw_oracle = _same_chunk_oracle( + x_data, weight_data, target, reduction, upstream, ignore_index + ) + + # Strict per-element match against the same-tiling oracle (dX is bit-exact; + # dW differs only by FP32 accumulation order, far below these tolerances). + assert_verbose_allclose(loss, loss_oracle, atol=1e-2, rtol=1e-3) + assert_verbose_allclose(x.grad, dx_oracle, atol=2e-3, rtol=1e-3) + assert_verbose_allclose(weight.grad, dw_oracle, atol=2e-3, rtol=1e-3) + + # Independent full-batch PyTorch anchor (aggregate relative-L2). + x_ref = x_data.clone().requires_grad_(True) + weight_ref = weight_data.clone().requires_grad_(True) + loss_ref = _reference(x_ref, weight_ref, target, ignore_index, reduction) + loss_ref.backward(upstream) + assert _relative_l2(x.grad, x_ref.grad) < 2e-2 + assert _relative_l2(weight.grad, weight_ref.grad) < 2e-2 + + +# --------------------------------------------------------------------------- +# Persistent retained-graph alias safety: a first backward (and mutating + +# clearing the produced .grad) must not corrupt the saved FP32 dW accumulator or +# raw BF16 dX. A second backward with a different upstream must match a FRESH +# forward+backward at that upstream EXACTLY -- proving the scale/cast is +# recomputed from the retained FP32 dW (not a double-rounded rescale of the first +# BF16 result). Uses a multi-chunk shape so cross-chunk accumulation is retained. +# --------------------------------------------------------------------------- +def test_retained_backward_alias_safety(): + torch.manual_seed(11) + n_tokens, hidden, vocab = 128, 64, 4096 + ignore_index = -100 + assert len(_expected_tiling(n_tokens, hidden, vocab)) >= 2 + x_data = torch.randn(n_tokens, hidden, device="cuda", dtype=torch.bfloat16) + weight_data = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16) + target = torch.randint(vocab, (n_tokens,), device="cuda") + target[:9] = ignore_index + + x = x_data.clone().requires_grad_(True) + weight = weight_data.clone().requires_grad_(True) + loss = _apply(x, weight, target, "mean", ignore_index)[0] + + loss.backward(torch.tensor(1.0, device="cuda"), retain_graph=True) + # Mutate then clear the produced gradients (must not feed back into saved state). + x.grad.fill_(123.0) + weight.grad.fill_(-456.0) + x.grad = None + weight.grad = None + loss.backward(torch.tensor(-1.3, device="cuda"), retain_graph=True) + gx_second = x.grad.clone() + gw_second = weight.grad.clone() + + # Fresh forward+backward at the same -1.3 upstream. + x_fresh = x_data.clone().requires_grad_(True) + weight_fresh = weight_data.clone().requires_grad_(True) + loss_fresh = _apply(x_fresh, weight_fresh, target, "mean", ignore_index)[0] + loss_fresh.backward(torch.tensor(-1.3, device="cuda")) + + assert_verbose_allclose(gx_second, x_fresh.grad, atol=0.0, rtol=0.0) + assert_verbose_allclose(gw_second, weight_fresh.grad, atol=0.0, rtol=0.0) + + # Directly snapshot the retained low-level tensors and confirm two backward + # passes with different upstreams leave them untouched. + xg = x_data.clone().requires_grad_(True) + wg = weight_data.clone().requires_grad_(True) + with torch.no_grad(): + _, grad_input, dweight_accum, gradient_scale = chunked_fused_linear_cross_entropy_forward(xg, wg, target) + grad_input_snap = grad_input.clone() + dweight_snap = dweight_accum.clone() + chunked_fused_linear_cross_entropy_backward( + torch.tensor(1.0, device="cuda"), grad_input, dweight_accum, gradient_scale, torch.bfloat16 + ) + chunked_fused_linear_cross_entropy_backward( + torch.tensor(-1.3, device="cuda"), grad_input, dweight_accum, gradient_scale, torch.bfloat16 + ) + assert_verbose_allclose(grad_input, grad_input_snap, atol=0.0, rtol=0.0) + assert_verbose_allclose(dweight_accum, dweight_snap, atol=0.0, rtol=0.0) + + +# --------------------------------------------------------------------------- +# Legacy low-level contract: the unchunked helper returns the full [N, V] dZ as +# its third element, and the paired 7-argument backward consumes it to run the +# dW matmul. This guards the original exported interface against silent contract +# drift (the chunked path uses the separate chunked_* helpers, no C argument). +# --------------------------------------------------------------------------- +def test_legacy_lowlevel_forward_backward_contract(): + torch.manual_seed(0) + n_tokens, hidden, vocab = 128, 128, 257 # odd V exercises a non-256-aligned dW tail + ignore_index = -100 + upstream = torch.tensor(-0.7, device="cuda") + x = torch.randn(n_tokens, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab, (n_tokens,), device="cuda") + target[:8] = ignore_index + + with torch.no_grad(): + loss, grad_input, grad_logits, gradient_scale = fused_linear_cross_entropy_forward( + x, weight, target, reduction="mean" + ) + + # Third element is the full [N, V] dZ (legacy contract), not [V, H] dW. + assert tuple(grad_logits.shape) == (n_tokens, vocab) + assert grad_logits.dtype == torch.bfloat16 + assert gradient_scale.ndim == 0 + + grad_x, grad_weight = fused_linear_cross_entropy_backward( + upstream, + grad_input, + grad_logits, + x.detach(), + gradient_scale, + weight.shape, + weight.dtype, + ) + assert tuple(grad_x.shape) == (n_tokens, hidden) + assert grad_x.dtype == torch.bfloat16 + assert tuple(grad_weight.shape) == (vocab, hidden) + assert grad_weight.dtype == weight.dtype + + # Cross-check the legacy gradients against the BF16-linear / FP32-CE anchor. + x_ref = x.detach().clone().requires_grad_(True) + weight_ref = weight.detach().clone().requires_grad_(True) + loss_ref = _reference(x_ref, weight_ref, target, ignore_index, "mean") + loss_ref.backward(upstream) + assert abs(loss.sum().item() - loss_ref.item()) < 5e-3 + assert _relative_l2(grad_x, x_ref.grad) < 2e-2 + assert _relative_l2(grad_weight, weight_ref.grad) < 2e-2 + + +# --------------------------------------------------------------------------- +# GEMM ownership trace: the dW gradient runs through PyTorch/cuBLAS, not a custom +# cuTile MMA kernel. The chunked path must emit exactly one FP32-out ``aten::mm`` +# (first chunk) and one FP32-out ``aten::addmm`` per subsequent chunk; projection +# and dX use the plain (BF16-out) ``aten::mm.out`` overload. The removed cuTile +# matmul kernels must no longer exist on the module. +# --------------------------------------------------------------------------- +class _AtenOpRecorder(TorchDispatchMode): + """Record every aten op name the path issues, plus dW mm/addmm out dtypes.""" + + def __init__(self): + self.ops = [] + self.dw_out_dtypes = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + name = str(func) + out = func(*args, **kwargs) + self.ops.append(name) + if name in ("aten.mm.dtype_out", "aten.addmm.dtype_out"): + result = out[0] if isinstance(out, (tuple, list)) else out + self.dw_out_dtypes.append(result.dtype) + return out + + +def test_no_custom_matmul_kernels_remain(): + # The dW GEMM is cuBLAS now; the old cuTile MMA kernels/launcher are gone. + for removed in ("_matmul_tn_kernel", "_matmul_tn_n4_kernel", "_launch_matmul_tn"): + assert not hasattr(flce_mod, removed), f"{removed} should have been removed" + # The CE / scale-cast cuTile kernels are retained. + assert hasattr(flce_mod, "_scale_cast_kernel") + assert hasattr(flce_mod, "_fused_cross_entropy_dz_kernel") + + +def test_chunked_dw_uses_mm_then_addmm_fp32_out(): + torch.manual_seed(3) + # V > 16*H so the default geometry is multi-chunk: [16, 16, 16, 16, 6] -> 5 chunks. + n_tokens, hidden, vocab = 70, 96, 8192 + x = torch.randn(n_tokens, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab, (n_tokens,), device="cuda") + + n_chunks = len(_expected_tiling(n_tokens, hidden, vocab)) + assert n_chunks >= 2 + with _AtenOpRecorder() as rec, torch.no_grad(): + chunked_fused_linear_cross_entropy_forward(x, weight, target) + + # First dW chunk -> one FP32-out mm; every subsequent chunk -> FP32-out addmm. + if flce_mod._ADDMM_SUPPORTS_OUT_DTYPE: + assert rec.ops.count("aten.mm.dtype_out") == 1 + assert rec.ops.count("aten.addmm.dtype_out") == n_chunks - 1 + assert rec.dw_out_dtypes == [torch.float32] * n_chunks + # dW never accumulated straight into a BF16 tensor via a plain addmm. + assert "aten.addmm.out" not in rec.ops + assert "aten.addmm.default" not in rec.ops + else: + # Pre-2.8 fallback: no ``.dtype_out`` overload. dW overwrites once via a + # plain mm.out then accumulates via plain addmm.out into the FP32 buffer. + # addmm is issued ONLY by the dW accumulation, so its count is chunks - 1. + assert "aten.mm.dtype_out" not in rec.ops + assert "aten.addmm.dtype_out" not in rec.ops + assert rec.ops.count("aten.addmm.out") == n_chunks - 1 + assert rec.dw_out_dtypes == [] + assert "aten.addmm.default" not in rec.ops + + +def test_legacy_backward_dw_uses_fp32_mm_nonunit_upstream(): + torch.manual_seed(5) + n_tokens, hidden, vocab = 96, 128, 300 + upstream = torch.tensor(-1.7, device="cuda") # non-unit upstream + x = torch.randn(n_tokens, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab, (n_tokens,), device="cuda") + + with torch.no_grad(): + _, grad_input, grad_logits, gradient_scale = fused_linear_cross_entropy_forward( + x, weight, target, reduction="mean" + ) + + with _AtenOpRecorder() as rec: + _, grad_weight = fused_linear_cross_entropy_backward( + upstream, grad_input, grad_logits, x.detach(), gradient_scale, weight.shape, weight.dtype + ) + + # Legacy dW is a single FP32-out mm (unchunked); no addmm accumulation. + if flce_mod._ADDMM_SUPPORTS_OUT_DTYPE: + assert rec.ops.count("aten.mm.dtype_out") == 1 + assert "aten.addmm.dtype_out" not in rec.ops + assert rec.dw_out_dtypes == [torch.float32] + else: + # Pre-2.8 fallback: the single dW GEMM is a plain mm.out into the FP32 + # buffer -- no ``.dtype_out`` overload and no addmm accumulation. + assert "aten.mm.dtype_out" not in rec.ops + assert "aten.addmm.dtype_out" not in rec.ops + assert rec.dw_out_dtypes == [] + assert grad_weight.dtype == weight.dtype + + +# --------------------------------------------------------------------------- +# Pre-2.8 fallback: when the ``out_dtype`` mm/addmm overload is unavailable the +# helper must upcast the bounded operands to FP32 and use the plain mm/addmm (no +# ``.dtype_out`` op). Small gradients must still survive the deferred cast. +# --------------------------------------------------------------------------- +def test_pre_2_8_fallback_avoids_out_dtype_overload(monkeypatch): + monkeypatch.setattr(flce_mod, "_ADDMM_SUPPORTS_OUT_DTYPE", False) + torch.manual_seed(9) + # V > 16*H -> multi-chunk under the default geometry: [16, 16, 16, 2] -> 4 chunks. + n_tokens, hidden, vocab = 50, 64, 4096 + x = torch.randn(n_tokens, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + target = torch.randint(vocab, (n_tokens,), device="cuda") + + n_chunks = len(_expected_tiling(n_tokens, hidden, vocab)) + assert n_chunks >= 2 + with _AtenOpRecorder() as rec, torch.no_grad(): + chunked_fused_linear_cross_entropy_forward(x, weight, target) + + # No out_dtype overload is used on the fallback path. + assert "aten.mm.dtype_out" not in rec.ops + assert "aten.addmm.dtype_out" not in rec.ops + # dW still overwrites once (plain mm.out) then accumulates (plain addmm.out). + # addmm is issued ONLY by the dW accumulation, so its count is chunks - 1. + assert rec.ops.count("aten.addmm.out") == n_chunks - 1 + + +def test_pre_2_8_fallback_small_gradient_survives(monkeypatch): + monkeypatch.setattr(flce_mod, "_ADDMM_SUPPORTS_OUT_DTYPE", False) + x = torch.full((1, 64), 1e-37, dtype=torch.bfloat16, device="cuda", requires_grad=True) + weight = torch.zeros((4096, 64), dtype=torch.bfloat16, device="cuda", requires_grad=True) + target = torch.tensor([0], device="cuda") + upstream = torch.tensor(65536.0, device="cuda") + + loss = _apply(x, weight, target, "mean", -100)[0] + loss.backward(upstream) + + # The FP32 fallback preserves the sub-BF16 gradient just like the overload path. + dz = torch.full((1, 4096), 1.0 / 4096, dtype=torch.bfloat16, device="cuda") + dz[0, 0] -= 1.0 + dw_fp32 = torch.mm(dz.t().to(torch.float32), x.detach().to(torch.float32)) # bounded-operand upcast + ref = (dw_fp32 * upstream.to(torch.float32)).to(torch.bfloat16) + + assert weight.grad[1, 0].item() != 0.0 + assert_verbose_allclose(weight.grad, ref, atol=0.0, rtol=0.0)