From 7107abb5d183b8390f2105096176722e685e6a34 Mon Sep 17 00:00:00 2001 From: debashishc Date: Sun, 8 Feb 2026 08:49:47 -0500 Subject: [PATCH] softmax_online: align issue 38 flow and simplify backend wiring --- bench/benchmark_online_softmax.py | 23 +- bench/run.py | 9 +- forge_cute_py/kernels/softmax_online.py | 333 +++++++++++++++++++++++ forge_cute_py/ops/__init__.py | 19 +- forge_cute_py/ops/softmax_online.py | 338 +++++++++++++----------- forge_cute_py/ref/softmax_online.py | 4 +- tests/test_softmax_online.py | 183 ++++++++----- 7 files changed, 663 insertions(+), 246 deletions(-) create mode 100644 forge_cute_py/kernels/softmax_online.py diff --git a/bench/benchmark_online_softmax.py b/bench/benchmark_online_softmax.py index 97e53ed..ed1e006 100644 --- a/bench/benchmark_online_softmax.py +++ b/bench/benchmark_online_softmax.py @@ -1,11 +1,16 @@ """Benchmark softmax_online op against torch.softmax and torch.compile(torch.softmax).""" import argparse -import os import torch -from forge_cute_py.ops.softmax_online import softmax_fwd, softmax_bwd +from forge_cute_py.ops.softmax_online import ( + get_softmax_online_backend, + list_softmax_online_backends, + set_softmax_online_backend, + softmax_bwd, + softmax_fwd, +) from forge_cute_py.util.bench import do_bench, estimate_bandwidth, summarize_times SHORT_M = [128, 512, 2048, 8192] @@ -26,6 +31,7 @@ def parse_str_list(s: str) -> list[str]: def main(): + backend_choices = list_softmax_online_backends() parser = argparse.ArgumentParser(description="Benchmark softmax_online op") parser.add_argument( "--long", action="store_true", help="Use long-N benchmark suite (small M, large N)" @@ -35,14 +41,9 @@ def main(): parser.add_argument("--dtypes", type=parse_str_list, default=DEFAULT_DTYPES) parser.add_argument("--warmup", type=int, default=20) parser.add_argument("--iterations", type=int, default=100) - parser.add_argument( - "--impl", - choices=["auto", "ref", "kernel"], - default="auto", - help="softmax_online backend mode (FORGE_SOFTMAX_IMPL)", - ) + parser.add_argument("--backend", choices=backend_choices, default="ref") args = parser.parse_args() - os.environ["FORGE_SOFTMAX_IMPL"] = args.impl + set_softmax_online_backend(args.backend) if args.m_sizes is None: args.m_sizes = LONG_M if args.long else SHORT_M @@ -54,7 +55,9 @@ def main(): gpu_name = torch.cuda.get_device_name(0) suite = "long" if args.long else "short" - print(f"softmax_online benchmarks [{suite}] ({gpu_name}) [impl={args.impl}]") + print( + f"softmax_online benchmarks [{suite}] ({gpu_name}) [backend={get_softmax_online_backend()}]" + ) print() header = ( diff --git a/bench/run.py b/bench/run.py index 1c1e1ec..2a48f88 100644 --- a/bench/run.py +++ b/bench/run.py @@ -1,6 +1,5 @@ import argparse import json -import os import sys from pathlib import Path @@ -153,14 +152,18 @@ def fn(): "reason": str(exc), } except Exception as exc: - impl = os.getenv("FORGE_SOFTMAX_IMPL", "auto") + backend = ( + ops.get_softmax_online_backend() + if hasattr(ops, "get_softmax_online_backend") + else "unknown" + ) return { "status": "skipped", "op": op_name, "shape": shape, "dtype": str(dtype).replace("torch.", ""), "dim": dim, - "reason": f"softmax_online failed (impl={impl}): {exc}", + "reason": f"softmax_online failed (backend={backend}): {exc}", } return { "status": "ok", diff --git a/forge_cute_py/kernels/softmax_online.py b/forge_cute_py/kernels/softmax_online.py new file mode 100644 index 0000000..c8e3139 --- /dev/null +++ b/forge_cute_py/kernels/softmax_online.py @@ -0,0 +1,333 @@ +"""PR #26 placeholder online-softmax kernels. + +This module is carried as a non-final reference implementation for contributors. +It is not treated as the production online-softmax kernel path yet. + +Credit: kernel implementation originated from Jonah Samost's PR #26 work. +""" + +import torch +import cutlass.cute as cute +import cuda.bindings.driver as cuda +from cutlass import BFloat16, Float16, Float32 +from cutlass.cute.runtime import from_dlpack + + +class SoftmaxOnlineBackward: + def __init__(self, dtype, N: int): + self.dtype = dtype + self.num_warps = 4 + self.bits_read = 128 + self.vec_load_size = self.bits_read // dtype.width + self.warp_size = 32 + self.threads_per_block = self.num_warps * self.warp_size + self.N = N # N is static at compile time, M is dynamic + + @cute.jit + def __call__(self, dY: cute.Tensor, y: cute.Tensor, dx: cute.Tensor, stream=None): + blocks_over_N = cute.ceil_div(self.N, self.vec_load_size * self.warp_size) + tiler_mn = ( # full covering tile + self.num_warps, + self.vec_load_size * self.warp_size * blocks_over_N, + ) + + copy_op = cute.nvgpu.CopyUniversalOp() + copy_atom = cute.make_copy_atom(copy_op, self.dtype, num_bits_per_copy=self.bits_read) + + thr_layout = cute.make_ordered_layout( + (self.num_warps, self.warp_size), + order=(1, 0), # cols move faster + ) + val_layout = cute.make_layout((1, self.vec_load_size)) + tiled_copy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout) + + blocks = cute.ceil_div(dY.shape[0], self.num_warps) + self.kernel(dY, y, dx, tiler_mn, tiled_copy).launch( + grid=(blocks, 1, 1), block=(self.threads_per_block, 1, 1), stream=stream + ) + + # type hints are not optional!!!! + @cute.kernel + def kernel( + self, + dY: cute.Tensor, + y: cute.Tensor, + dX: cute.Tensor, + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + ): + # Compute gradient (numerically stable) + # dot_product = (dy * y).sum(dim=dim, keepdim=True) + # result = y * (dy - dot_product) + # dx.copy_(result) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + dy_tile = cute.local_tile(dY, tiler_mn, (bidx, 0)) + y_tile = cute.local_tile(y, tiler_mn, (bidx, 0)) + dx_tile = cute.local_tile(dX, tiler_mn, (bidx, 0)) + + tidxSlice = tiled_copy.get_slice(tidx) + + dy_idx = tidxSlice.partition_S(dy_tile) + y_idx = tidxSlice.partition_S(y_tile) + dx_idx = tidxSlice.partition_D(dx_tile) + + dy_regs = cute.make_rmem_tensor_like(dy_idx) + y_regs = cute.make_rmem_tensor_like(y_idx) + + cute.autovec_copy(dy_idx, dy_regs) + cute.autovec_copy(y_idx, y_regs) + + dy_data = dy_regs.load() + y_data = y_regs.load() + + tidx_local_dp = dy_data * y_data + tidx_local_sum = tidx_local_dp.reduce( + cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0 + ) + row_dp_sum = cute.arch.warp_reduction_sum(tidx_local_sum) + + result = y_data * (dy_data - row_dp_sum) + dy_regs.store(result) + cute.autovec_copy(dy_regs, dx_idx) + + +class SoftmaxOnlineLoop: + def __init__(self, dtype): + self.dtype = dtype + self.num_warps = 1 + self.threads_per_block = self.num_warps * 32 + self.NEG_INF = Float32(float("-inf")) + + @cute.jit + def __call__(self, gInput: cute.Tensor, gOutput: cute.Tensor, stream: cuda.CUstream = None): + M, N = gInput.shape + thr_layout = cute.make_layout((self.threads_per_block,), stride=(1,)) + val_layout = cute.make_layout((1,), stride=(1,)) + tiler_mn_1d, tv_layout = cute.make_layout_tv(thr_layout, val_layout) + tiler_mn = (1, tiler_mn_1d[0]) + gX = cute.zipped_divide(gInput, tiler_mn) + gY = cute.zipped_divide(gOutput, tiler_mn) + + self.kernel(gX, gY, N).launch( + grid=(cute.size(gX, mode=[1, 0]), 1, 1), # RestM + block=(cute.size(tv_layout, mode=[0, 0])), # threads per block + stream=stream, + ) + + @cute.kernel + def kernel(self, gInput, gOutput, N): + maxValue = self.NEG_INF + sumValue = Float32(0.0) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + bdimx, _, _ = cute.arch.block_dim() + (TileM, TileN), (RestM, RestN) = gInput.shape + + for i in range(RestN): + idx = i * bdimx + tidx + value = Float32(gInput[(0, tidx), (bidx, i)]) if idx < N else self.NEG_INF + + curMax = cute.arch.warp_reduction_max(value) + prevMax = maxValue + maxValue = cute.arch.fmax(prevMax, curMax) + + scale = cute.math.exp(prevMax - maxValue) + scale_data = cute.math.exp(value - maxValue) + curSum = cute.arch.warp_reduction_sum(scale_data) + sumValue = sumValue * scale + curSum + + for i in range(RestN): + idx = i * bdimx + tidx + if idx < N: + value = gInput[(0, tidx), (bidx, i)].to(self.dtype) + data = cute.math.exp(value - maxValue) / sumValue + gOutput[(0, tidx), (bidx, i)] = data.to(self.dtype) + + +class SoftmaxOnline: + def __init__(self, dtype, N: int): + self.dtype = dtype + self.num_warps = 1 + self.threads_per_block = self.num_warps * 32 + self.NEG_INF = Float32(float("-inf")) + self.N = N + + self.bits_read = 128 + self.vec_load_size = self.bits_read // self.dtype.width + self.threads_per_row = 32 + self.num_warps = 4 + self.num_threads = self.num_warps * self.threads_per_row + + @cute.jit + def __call__(self, gInput: cute.Tensor, gOutput: cute.Tensor, stream: cuda.CUstream = None): + + blocks_vector_N = cute.ceil_div(self.N, self.bits_read // self.dtype.width) + blocks_over_N = cute.ceil_div(blocks_vector_N, self.threads_per_row) + tiler_mn = ( + self.num_warps, + self.vec_load_size * blocks_over_N * self.threads_per_row, + ) # [4, ~N] + + copy_op = cute.nvgpu.CopyUniversalOp() + copy_atom = cute.make_copy_atom(copy_op, self.dtype, num_bits_per_copy=self.bits_read) + thr_layout = cute.make_ordered_layout( + (self.num_warps, self.threads_per_row), + order=(1, 0), + ) + val_layout = cute.make_layout((1, self.vec_load_size)) + tiled_copy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout) + + blocks = cute.ceil_div(gInput.shape[0], tiler_mn[0]) + self.kernel(gInput, gOutput, tiler_mn, tiled_copy).launch( + grid=(blocks, 1, 1), block=(self.num_threads, 1, 1), stream=stream + ) + + # type hints are not optional!!!! + @cute.kernel + def kernel( + self, + gInput: cute.Tensor, + gOutput: cute.Tensor, + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + gX = cute.local_tile(gInput, tiler_mn, (bidx, 0)) + gY = cute.local_tile(gOutput, tiler_mn, (bidx, 0)) + # this thread is response for vectorized loads, striding 4 * 32 across the row + tidxSlice = tiled_copy.get_slice(tidx) + tidxIndices = tidxSlice.partition_S(gX) + tidxRegs = cute.make_rmem_tensor_like(tidxIndices) + cute.autovec_copy(tidxIndices, tidxRegs) + + tidxValues = tidxRegs.load() + tidLocalMax = tidxValues.reduce( + cute.ReductionOp.MAX, init_val=self.NEG_INF, reduction_profile=0 + ) + rowMax = cute.arch.warp_reduction_max(tidLocalMax) + tidScaledLocalSum = cute.math.exp(tidxValues - rowMax) + tidLocalSum = tidScaledLocalSum.reduce( + cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0 + ) + rowSum = cute.arch.warp_reduction_sum(tidLocalSum) + + writeValues = cute.math.exp(tidxValues - rowMax) / rowSum + tidxRegs.store(writeValues) + tidxOutIndices = tidxSlice.partition_D(gY) + cute.autovec_copy(tidxRegs, tidxOutIndices) + + +def benchmark(loopless=True): + import time + + dim = -1 + M, N = 4096, 768 + dtype = torch.float32 + dtype_map = { + torch.float16: Float16, + torch.float32: Float32, + torch.bfloat16: BFloat16, + } + cute_dtype = dtype_map[dtype] + + x = torch.randn(M, N, device="cuda", dtype=dtype) + output = torch.zeros_like(x) + + if loopless: + dx = x + dy = output + m = cute.sym_int() + input_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, N), stride_order=(1, 0)) + output_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, N), stride_order=(1, 0)) + softmax = SoftmaxOnline(dtype_map[dtype], N) + fn = cute.compile( + softmax, + input_cute, + output_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + fn(x, output) + else: + dx = from_dlpack(x, enable_tvm_ffi=True) + dy = from_dlpack(output, enable_tvm_ffi=True) + softmax = SoftmaxOnlineLoop(dtype_map[dtype]) + fn = cute.compile( + softmax, + dx, + dy, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + fn(dx, dy) + + print("Correctness check:") + expected = torch.nn.functional.softmax(x, dim=-1) + is_close = torch.allclose(output, expected, rtol=1e-3, atol=1e-3) + print(f" dim=-1: {'✓ PASS' if is_close else '✗ FAIL'}") + if not is_close: + max_diff = (output - expected).abs().max().item() + print(f" max diff: {max_diff}") + + print("\nBenchmarks:") + + # Warmup + for _ in range(10): + fn(dx, dy) + torch.cuda.synchronize() + + # Benchmark our softmax + start = time.perf_counter() + for _ in range(100): + fn(dx, dy) + torch.cuda.synchronize() + print(f" softmax_online dim=-1: {(time.perf_counter() - start) * 10:.3f} ms") + + # Compare to PyTorch + start = time.perf_counter() + for _ in range(100): + _ = torch.nn.functional.softmax(x, dim=-1) + torch.cuda.synchronize() + print(f" torch.softmax dim=-1: {(time.perf_counter() - start) * 10:.3f} ms") + + +def bench_back(): + dim = -1 + M, N = 256, 256 + dtype = torch.float32 + dtype_map = { + torch.float16: Float16, + torch.float32: Float32, + torch.bfloat16: BFloat16, + } + cute_dtype = dtype_map[dtype] + dy = torch.randn(M, N, device="cuda", dtype=dtype) + y = torch.randn(M, N, device="cuda", dtype=dtype) + dx = torch.randn(M, N, device="cuda", dtype=dtype) + + m = cute.sym_int() + dy_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, N), stride_order=(1, 0)) + y_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, N), stride_order=(1, 0)) + dx_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, N), stride_order=(1, 0)) + softmax = SoftmaxOnlineBackward(dtype_map[dtype], N) + fn = cute.compile( + softmax, + dy_cute, + y_cute, + dx_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + fn(dy, y, dx) + + +# benchmark(loopless=False) +# benchmark(loopless=True) + +# bench_back() diff --git a/forge_cute_py/ops/__init__.py b/forge_cute_py/ops/__init__.py index bcf0ee9..0ff9ba1 100644 --- a/forge_cute_py/ops/__init__.py +++ b/forge_cute_py/ops/__init__.py @@ -1,6 +1,21 @@ from forge_cute_py.ops.copy_transpose import copy_transpose from forge_cute_py.ops.reduce import reduce from forge_cute_py.ops.reduce_sum import reduce_sum -from forge_cute_py.ops.softmax_online import softmax_online +from forge_cute_py.ops.softmax_online import ( + get_softmax_online_backend, + list_softmax_online_backends, + register_softmax_online_backend, + set_softmax_online_backend, + softmax_online, +) -__all__ = ["copy_transpose", "reduce", "reduce_sum", "softmax_online"] +__all__ = [ + "copy_transpose", + "reduce", + "reduce_sum", + "softmax_online", + "register_softmax_online_backend", + "set_softmax_online_backend", + "get_softmax_online_backend", + "list_softmax_online_backends", +] diff --git a/forge_cute_py/ops/softmax_online.py b/forge_cute_py/ops/softmax_online.py index 889a21e..cd592c2 100644 --- a/forge_cute_py/ops/softmax_online.py +++ b/forge_cute_py/ops/softmax_online.py @@ -1,160 +1,33 @@ +"""Online softmax op with backend registry.""" + from __future__ import annotations -import importlib -import os -from types import ModuleType +from dataclasses import dataclass from typing import Callable import torch SoftmaxForwardImpl = Callable[[torch.Tensor, int], torch.Tensor] SoftmaxBackwardImpl = Callable[[torch.Tensor, torch.Tensor, int], torch.Tensor] - -_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32, torch.float64) -_SUPPORTED_IMPL_MODES = {"auto", "ref", "kernel"} - - -def _normalize_dim(dim: int, ndim: int) -> int: - dim = dim if dim >= 0 else ndim + dim - if dim not in (0, 1): - raise ValueError(f"softmax_online expects dim in {{-1, 0, 1}} for 2D tensors, got {dim}") - return dim - - -def _validate_impl_mode() -> str: - impl = os.getenv("FORGE_SOFTMAX_IMPL", "auto").strip().lower() - if impl not in _SUPPORTED_IMPL_MODES: - choices = ", ".join(sorted(_SUPPORTED_IMPL_MODES)) - raise ValueError(f"FORGE_SOFTMAX_IMPL must be one of {{{choices}}}, got '{impl}'") - return impl - - -def _load_kernel_module() -> tuple[ModuleType | None, str | None]: - module_name = "forge_cute_py.kernels.softmax_online" - try: - return importlib.import_module(module_name), None - except ModuleNotFoundError as exc: - if exc.name == module_name: - return None, f"module '{module_name}' was not found" - if exc.name is None: - return None, f"failed importing '{module_name}': {exc}" - return None, f"dependency '{exc.name}' missing while importing '{module_name}'" - except Exception as exc: # pragma: no cover - defensive runtime diagnostics - return None, f"failed importing '{module_name}': {exc}" - - -def _resolve_kernel_forward(module: ModuleType) -> SoftmaxForwardImpl: - for attr in ("softmax_fwd", "softmax_online"): - fn = getattr(module, attr, None) - if callable(fn): - return fn - raise AttributeError( - "kernel module must define a callable 'softmax_fwd(x, dim)' or 'softmax_online(x, dim)'" - ) - - -def _resolve_kernel_backward(module: ModuleType) -> SoftmaxBackwardImpl | None: - fn = getattr(module, "softmax_bwd", None) - if callable(fn): - return fn - return None - - -def _call_forward(fn: SoftmaxForwardImpl, x: torch.Tensor, dim: int) -> torch.Tensor: - try: - return fn(x, dim=dim) - except TypeError: - return fn(x, dim) - - -def _call_backward( - fn: SoftmaxBackwardImpl, - dy: torch.Tensor, - y: torch.Tensor, - dim: int, -) -> torch.Tensor: - try: - return fn(dy, y, dim=dim) - except TypeError: - return fn(dy, y, dim) - - -def _reference_softmax_forward(x: torch.Tensor, dim: int) -> torch.Tensor: - from forge_cute_py.ref import softmax_online as softmax_online_ref - - return softmax_online_ref(x, dim=dim) +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) -def _reference_softmax_backward(dy: torch.Tensor, y: torch.Tensor, dim: int) -> torch.Tensor: - dot_product = (dy * y).sum(dim=dim, keepdim=True) - return y * (dy - dot_product) +@dataclass(frozen=True) +class _SoftmaxBackend: + forward_impl: SoftmaxForwardImpl + backward_impl: SoftmaxBackwardImpl | None = None -def _forward_impl(x: torch.Tensor, dim: int) -> torch.Tensor: - impl = _validate_impl_mode() - if impl == "ref": - return _reference_softmax_forward(x, dim) +_softmax_online_backends: dict[str, _SoftmaxBackend] = {} +_active_softmax_online_backend = "ref" +_kernel_fwd_compile_cache: dict[tuple[object, int], object] = {} +_kernel_bwd_compile_cache: dict[tuple[object, int], object] = {} - module, reason = _load_kernel_module() - if module is None: - if impl == "auto": - return _reference_softmax_forward(x, dim) - raise NotImplementedError( - "FORGE_SOFTMAX_IMPL=kernel requested, but softmax kernel is unavailable: " - f"{reason}. Add forge_cute_py/kernels/softmax_online.py with softmax_fwd()." - ) - - try: - kernel_forward = _resolve_kernel_forward(module) - except AttributeError as exc: - if impl == "auto": - return _reference_softmax_forward(x, dim) - raise NotImplementedError( - "FORGE_SOFTMAX_IMPL=kernel requested, but softmax kernel forward entry point is " - f"incomplete: {exc}" - ) from exc - - try: - return _call_forward(kernel_forward, x, dim) - except NotImplementedError as exc: - if impl == "auto": - return _reference_softmax_forward(x, dim) - raise NotImplementedError( - "FORGE_SOFTMAX_IMPL=kernel requested, but softmax kernel forward is not implemented." - ) from exc - - -def _backward_impl(dy: torch.Tensor, y: torch.Tensor, dim: int) -> torch.Tensor: - impl = _validate_impl_mode() - if impl == "ref": - return _reference_softmax_backward(dy, y, dim) - - module, reason = _load_kernel_module() - if module is None: - if impl == "auto": - return _reference_softmax_backward(dy, y, dim) - raise NotImplementedError( - "FORGE_SOFTMAX_IMPL=kernel requested, but softmax kernel is unavailable: " - f"{reason}. Add forge_cute_py/kernels/softmax_online.py with softmax_bwd()." - ) - - kernel_backward = _resolve_kernel_backward(module) - if kernel_backward is None: - if impl == "auto": - return _reference_softmax_backward(dy, y, dim) - raise NotImplementedError( - "FORGE_SOFTMAX_IMPL=kernel requested, but 'softmax_bwd' is missing in " - "forge_cute_py.kernels.softmax_online." - ) - try: - return _call_backward(kernel_backward, dy, y, dim) - except NotImplementedError as exc: - if impl == "auto": - return _reference_softmax_backward(dy, y, dim) - raise NotImplementedError( - "FORGE_SOFTMAX_IMPL=kernel requested, but softmax kernel backward is not implemented." - ) from exc +def _normalize_dim(dim: int, ndim: int) -> int: + if dim != -1: + raise ValueError(f"softmax_online expects dim=-1 (row-wise) for 2D tensors, got {dim}") + return -1 def _ensure_forward_inputs(x: torch.Tensor, out: torch.Tensor, dim: int) -> int: @@ -187,25 +60,154 @@ def _ensure_backward_inputs(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor, return _normalize_dim(dim, dy.ndim) +def _reference_softmax_forward(x: torch.Tensor, dim: int) -> torch.Tensor: + from forge_cute_py.ref import softmax_online as softmax_online_ref + + return softmax_online_ref(x, dim=dim) + + +def _reference_softmax_backward(dy: torch.Tensor, y: torch.Tensor, dim: int) -> torch.Tensor: + dot_product = (dy * y).sum(dim=dim, keepdim=True) + return y * (dy - dot_product) + + +def _kernel_forward(x: torch.Tensor, dim: int) -> torch.Tensor: + del dim + if x.shape[1] % 32 != 0: + raise ValueError(f"Inner dimension N must be a multiple of 32, got {x.shape[1]}") + + import cutlass.cute as cute + from cutlass import BFloat16, Float16, Float32 + from forge_cute_py.kernels.softmax_online import SoftmaxOnline + + dtype_map = { + torch.float16: Float16, + torch.float32: Float32, + torch.bfloat16: BFloat16, + } + cute_dtype = dtype_map[x.dtype] + compile_key = (cute_dtype, x.shape[1]) + + if compile_key not in _kernel_fwd_compile_cache: + m = cute.sym_int() + n = x.shape[1] + input_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, n), stride_order=(1, 0)) + output_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, n), stride_order=(1, 0)) + _kernel_fwd_compile_cache[compile_key] = cute.compile( + SoftmaxOnline(cute_dtype, n), + input_cute, + output_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + out = torch.empty_like(x) + _kernel_fwd_compile_cache[compile_key](x, out) + return out + + +def _kernel_backward(dy: torch.Tensor, y: torch.Tensor, dim: int) -> torch.Tensor: + del dim + if dy.shape[1] % 32 != 0: + raise ValueError(f"Inner dimension N must be a multiple of 32, got {dy.shape[1]}") + + import cutlass.cute as cute + from cutlass import BFloat16, Float16, Float32 + from forge_cute_py.kernels.softmax_online import SoftmaxOnlineBackward + + dtype_map = { + torch.float16: Float16, + torch.float32: Float32, + torch.bfloat16: BFloat16, + } + cute_dtype = dtype_map[dy.dtype] + compile_key = (cute_dtype, dy.shape[1]) + + if compile_key not in _kernel_bwd_compile_cache: + m = cute.sym_int() + n = dy.shape[1] + dy_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, n), stride_order=(1, 0)) + y_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, n), stride_order=(1, 0)) + dx_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (m, n), stride_order=(1, 0)) + _kernel_bwd_compile_cache[compile_key] = cute.compile( + SoftmaxOnlineBackward(cute_dtype, n), + dy_cute, + y_cute, + dx_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + dx = torch.empty_like(dy) + _kernel_bwd_compile_cache[compile_key](dy, y, dx) + return dx + + +def register_softmax_online_backend( + name: str, + forward_impl: SoftmaxForwardImpl, + backward_impl: SoftmaxBackwardImpl | None = None, + *, + overwrite: bool = False, +) -> None: + if not name: + raise ValueError("Backend name must be a non-empty string") + if not callable(forward_impl): + raise TypeError("forward_impl must be callable") + if backward_impl is not None and not callable(backward_impl): + raise TypeError("backward_impl must be callable when provided") + if name in _softmax_online_backends and not overwrite: + raise ValueError(f"Backend '{name}' already exists. Pass overwrite=True to replace it.") + _softmax_online_backends[name] = _SoftmaxBackend( + forward_impl=forward_impl, backward_impl=backward_impl + ) + + +def list_softmax_online_backends() -> tuple[str, ...]: + return tuple(_softmax_online_backends.keys()) + + +def get_softmax_online_backend() -> str: + return _active_softmax_online_backend + + +def set_softmax_online_backend(name: str) -> None: + global _active_softmax_online_backend + if name not in _softmax_online_backends: + choices = ", ".join(list_softmax_online_backends()) or "" + raise ValueError(f"Unknown softmax_online backend '{name}'. Available: {choices}") + _active_softmax_online_backend = name + + +def _current_softmax_backend() -> _SoftmaxBackend: + return _softmax_online_backends[_active_softmax_online_backend] + + @torch.library.custom_op("forge_cute_py::_softmax_fwd", mutates_args={"out"}) def _softmax_fwd(x: torch.Tensor, out: torch.Tensor, dim: int = -1) -> None: """Softmax forward pass.""" dim = _ensure_forward_inputs(x, out, dim) - result = _forward_impl(x, dim) + backend = _current_softmax_backend() + result = backend.forward_impl(x, dim) if result.shape != x.shape: raise ValueError( - f"softmax forward produced invalid shape {result.shape}, expected {x.shape}" + f"Backend '{_active_softmax_online_backend}' returned invalid shape {result.shape}, " + f"expected {x.shape}." ) if result.dtype != x.dtype: raise ValueError( - f"softmax forward produced invalid dtype {result.dtype}, expected {x.dtype}" + f"Backend '{_active_softmax_online_backend}' returned invalid dtype {result.dtype}, " + f"expected {x.dtype}." ) if result.device != x.device: - raise ValueError(f"softmax forward produced output on {result.device}, expected {x.device}") + raise ValueError( + f"Backend '{_active_softmax_online_backend}' returned output on {result.device}, " + f"expected {x.device}." + ) out.copy_(result) -_softmax_fwd.compile_cache = {} +_softmax_fwd.compile_cache = _kernel_fwd_compile_cache def softmax_fwd(x: torch.Tensor, dim: int = -1) -> torch.Tensor: @@ -219,23 +221,31 @@ def softmax_fwd(x: torch.Tensor, dim: int = -1) -> torch.Tensor: def _softmax_backward(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor, dim: int = -1) -> None: """Softmax backward pass.""" dim = _ensure_backward_inputs(dy, y, dx, dim) - result = _backward_impl(dy, y, dim) + backend = _current_softmax_backend() + if backend.backward_impl is None: + result = _reference_softmax_backward(dy, y, dim) + else: + result = backend.backward_impl(dy, y, dim) + if result.shape != dy.shape: raise ValueError( - f"softmax backward produced invalid shape {result.shape}, expected {dy.shape}" + f"Backend '{_active_softmax_online_backend}' returned invalid grad shape {result.shape}, " + f"expected {dy.shape}." ) if result.dtype != dy.dtype: raise ValueError( - f"softmax backward produced invalid dtype {result.dtype}, expected {dy.dtype}" + f"Backend '{_active_softmax_online_backend}' returned invalid grad dtype {result.dtype}, " + f"expected {dy.dtype}." ) if result.device != dy.device: raise ValueError( - f"softmax backward produced output on {result.device}, expected {dy.device}" + f"Backend '{_active_softmax_online_backend}' returned grad on {result.device}, " + f"expected {dy.device}." ) dx.copy_(result) -_softmax_backward.compile_cache = {} +_softmax_backward.compile_cache = _kernel_bwd_compile_cache def softmax_bwd(dy: torch.Tensor, y: torch.Tensor, dim: int = -1) -> torch.Tensor: @@ -261,11 +271,19 @@ def backward(ctx, dy): def softmax_online(x: torch.Tensor, dim: int = -1) -> torch.Tensor: - """Online softmax with autograd support. - - Backend selection is controlled by FORGE_SOFTMAX_IMPL: - - auto (default): try kernel first, fallback to reference - - ref: force reference implementation - - kernel: require kernel implementation (raise if unavailable) - """ + """Online softmax with automatic differentiation support.""" return SoftmaxOnlineFunction.apply(x, dim) + + +register_softmax_online_backend( + "ref", + _reference_softmax_forward, + _reference_softmax_backward, + overwrite=True, +) +register_softmax_online_backend( + "kernel", + _kernel_forward, + _kernel_backward, + overwrite=True, +) diff --git a/forge_cute_py/ref/softmax_online.py b/forge_cute_py/ref/softmax_online.py index 4f8994a..9b6b14c 100644 --- a/forge_cute_py/ref/softmax_online.py +++ b/forge_cute_py/ref/softmax_online.py @@ -4,7 +4,7 @@ def softmax_online(x: torch.Tensor, dim: int = -1) -> torch.Tensor: if x.ndim != 2: raise ValueError("softmax_online expects a 2D tensor") - if dim not in (-1, 0, 1): - raise ValueError("softmax_online expects dim in {-1, 0, 1} for 2D tensors") + if dim != -1: + raise ValueError(f"softmax_online expects dim=-1 (row-wise) for 2D tensors, got {dim}") x_dtype = x.dtype return x.float().softmax(dim=dim).to(x_dtype) diff --git a/tests/test_softmax_online.py b/tests/test_softmax_online.py index 487d512..7443c34 100644 --- a/tests/test_softmax_online.py +++ b/tests/test_softmax_online.py @@ -1,34 +1,31 @@ -import importlib - import pytest import torch -from forge_cute_py.ops import softmax_online +from forge_cute_py.ops import ( + get_softmax_online_backend, + list_softmax_online_backends, + register_softmax_online_backend, + set_softmax_online_backend, + softmax_online, +) from forge_cute_py.ref import softmax_online as ref_softmax_online -softmax_online_ops = importlib.import_module("forge_cute_py.ops.softmax_online") - @pytest.fixture(autouse=True) -def _reset_softmax_impl_env(monkeypatch): - monkeypatch.delenv("FORGE_SOFTMAX_IMPL", raising=False) - - -def _patch_missing_kernel_module(monkeypatch): - original_import_module = softmax_online_ops.importlib.import_module +def _restore_softmax_backend(): + previous = get_softmax_online_backend() + try: + set_softmax_online_backend("ref") + yield + finally: + set_softmax_online_backend(previous) - def missing_kernel_module(name, *args, **kwargs): - if name == "forge_cute_py.kernels.softmax_online": - exc = ModuleNotFoundError(f"No module named '{name}'") - exc.name = name - raise exc - return original_import_module(name, *args, **kwargs) - monkeypatch.setattr(softmax_online_ops.importlib, "import_module", missing_kernel_module) +dims = [-1] @pytest.mark.parametrize("shape", [(4, 8), (2, 128)]) -@pytest.mark.parametrize("dim", [-1, 0, 1]) +@pytest.mark.parametrize("dim", dims) @pytest.mark.parametrize( "dtype, atol, rtol", [ @@ -46,8 +43,20 @@ def test_softmax_online_correctness(shape, dim, dtype, atol, rtol): assert torch.isfinite(y).all() +def test_softmax_online_rejects_dim0(): + x = torch.randn(4, 8, device="cuda", dtype=torch.float16) + with pytest.raises(ValueError, match="expects dim=-1"): + softmax_online(x, dim=0) + + +def test_softmax_online_rejects_dim1(): + x = torch.randn(4, 8, device="cuda", dtype=torch.float16) + with pytest.raises(ValueError, match="expects dim=-1"): + softmax_online(x, dim=1) + + @pytest.mark.parametrize("shape", [(4, 8), (2, 128)]) -@pytest.mark.parametrize("dim", [-1, 0, 1]) +@pytest.mark.parametrize("dim", dims) @pytest.mark.parametrize( "dtype, atol, rtol", [ @@ -77,53 +86,6 @@ def test_softmax_online_torch_compile(shape, dim, dtype, atol, rtol): torch.testing.assert_close(y, y_ref, atol=atol, rtol=rtol) -def test_softmax_online_auto_falls_back_to_ref_when_kernel_missing(monkeypatch): - monkeypatch.setenv("FORGE_SOFTMAX_IMPL", "auto") - _patch_missing_kernel_module(monkeypatch) - - x = torch.randn(4, 8, device="cuda", dtype=torch.float16) - y = softmax_online(x, dim=-1) - y_ref = ref_softmax_online(x, dim=-1) - torch.testing.assert_close(y, y_ref, atol=1e-2, rtol=1e-2) - - -def test_softmax_online_kernel_mode_requires_kernel(monkeypatch): - monkeypatch.setenv("FORGE_SOFTMAX_IMPL", "kernel") - _patch_missing_kernel_module(monkeypatch) - - x = torch.randn(4, 8, device="cuda", dtype=torch.float16) - with pytest.raises(NotImplementedError, match="FORGE_SOFTMAX_IMPL=kernel"): - softmax_online(x, dim=-1) - - -def test_softmax_online_ref_mode_skips_kernel_probe(monkeypatch): - import_called = {"value": False} - original_import_module = softmax_online_ops.importlib.import_module - - def tracking_import(name, *args, **kwargs): - if name == "forge_cute_py.kernels.softmax_online": - import_called["value"] = True - raise AssertionError("Kernel module should not be imported in ref mode") - return original_import_module(name, *args, **kwargs) - - monkeypatch.setenv("FORGE_SOFTMAX_IMPL", "ref") - monkeypatch.setattr(softmax_online_ops.importlib, "import_module", tracking_import) - - x = torch.randn(4, 8, device="cuda", dtype=torch.float16) - y = softmax_online(x, dim=-1) - y_ref = ref_softmax_online(x, dim=-1) - - assert import_called["value"] is False - torch.testing.assert_close(y, y_ref, atol=1e-2, rtol=1e-2) - - -def test_softmax_online_rejects_invalid_impl_mode(monkeypatch): - monkeypatch.setenv("FORGE_SOFTMAX_IMPL", "unknown") - x = torch.randn(4, 8, device="cuda", dtype=torch.float16) - with pytest.raises(ValueError, match="FORGE_SOFTMAX_IMPL"): - softmax_online(x, dim=-1) - - @pytest.mark.parametrize("input_dtype", [torch.float16, torch.float32]) def test_softmax_online_properties(input_dtype): x = torch.randn(16, 256, device="cuda", dtype=input_dtype) @@ -160,7 +122,7 @@ def test_softmax_online_extreme_values(input_dtype): @pytest.mark.parametrize("shape", [(4, 8), (16, 128), (32, 256)]) -@pytest.mark.parametrize("dim", [-1, 0, 1]) +@pytest.mark.parametrize("dim", dims) @pytest.mark.parametrize( "dtype, atol, rtol", [ @@ -172,6 +134,7 @@ def test_softmax_online_extreme_values(input_dtype): def test_softmax_online_backward(shape, dim, dtype, atol, rtol): """Test backward pass against PyTorch reference.""" # Create inputs with gradients enabled (scale by 0.1 to avoid overflow) + x = (0.1 * torch.randn(*shape, device="cuda", dtype=dtype)).requires_grad_(True) x_ref = x.detach().clone().requires_grad_(True) @@ -189,7 +152,7 @@ def test_softmax_online_backward(shape, dim, dtype, atol, rtol): @pytest.mark.parametrize("shape", [(4, 8), (16, 128)]) -@pytest.mark.parametrize("dim", [-1, 1]) +@pytest.mark.parametrize("dim", dims) @pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) def test_softmax_online_backward_torch_compile(shape, dim, dtype): """Test backward pass works with torch.compile.""" @@ -246,3 +209,85 @@ def test_softmax_online_gradient_properties(): # Row sums should be approximately zero (softmax Jacobian property) row_sums = dx.sum(dim=-1) torch.testing.assert_close(row_sums, torch.zeros_like(row_sums), atol=1e-6, rtol=1e-6) + + +def test_softmax_online_backend_registry_exposes_expected_backends(): + backends = list_softmax_online_backends() + assert "ref" in backends + assert "kernel" in backends + + +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float16, 1e-2, 1e-2), + (torch.float32, 1e-4, 1e-4), + ], +) +def test_softmax_online_kernel_backend_forward_matches_ref(dtype, atol, rtol): + set_softmax_online_backend("kernel") + + x = (0.1 * torch.randn(16, 128, device="cuda", dtype=dtype)).requires_grad_(True) + y = softmax_online(x, dim=-1) + y_ref = ref_softmax_online(x, dim=-1) + torch.testing.assert_close(y, y_ref, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float32, 1e-4, 1e-4), + ], +) +def test_softmax_online_kernel_backend_backward_matches_ref(dtype, atol, rtol): + set_softmax_online_backend("kernel") + + x = (0.1 * torch.randn(16, 128, device="cuda", dtype=dtype)).requires_grad_(True) + x_ref = x.detach().clone().requires_grad_(True) + + y = softmax_online(x, dim=-1) + y_ref = ref_softmax_online(x_ref, dim=-1) + dy = torch.randn_like(y) + + torch.cuda.synchronize() + (dx,) = torch.autograd.grad(y, x, grad_outputs=dy) + (dx_ref,) = torch.autograd.grad(y_ref, x_ref, grad_outputs=dy) + torch.testing.assert_close(dx, dx_ref, atol=atol, rtol=rtol) + + +def test_softmax_online_custom_backend_forward(): + def custom_ref_forward(x: torch.Tensor, dim: int) -> torch.Tensor: + return ref_softmax_online(x, dim=dim) + + register_softmax_online_backend("test_ref_forward", custom_ref_forward, overwrite=True) + set_softmax_online_backend("test_ref_forward") + + x = torch.randn(16, 128, device="cuda", dtype=torch.float16) + y = softmax_online(x, dim=-1) + y_ref = ref_softmax_online(x, dim=-1) + torch.testing.assert_close(y, y_ref, atol=1e-2, rtol=1e-2) + + +def test_softmax_online_custom_backend_backward_fallback(): + def forward_only(x: torch.Tensor, dim: int) -> torch.Tensor: + return ref_softmax_online(x, dim=dim) + + register_softmax_online_backend("test_fwd_only", forward_only, overwrite=True) + set_softmax_online_backend("test_fwd_only") + + x = (0.1 * torch.randn(16, 128, device="cuda", dtype=torch.float32)).requires_grad_(True) + x_ref = x.detach().clone().requires_grad_(True) + + y = softmax_online(x, dim=-1) + y_ref = ref_softmax_online(x_ref, dim=-1) + + dy = torch.randn_like(y) + torch.cuda.synchronize() + (dx,) = torch.autograd.grad(y, x, grad_outputs=dy) + (dx_ref,) = torch.autograd.grad(y_ref, x_ref, grad_outputs=dy) + torch.testing.assert_close(dx, dx_ref, atol=1e-4, rtol=1e-4) + + +def test_softmax_online_unknown_backend_raises(): + with pytest.raises(ValueError, match="Unknown softmax_online backend"): + set_softmax_online_backend("missing_backend")