diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 19c6518..69cbba1 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -7,7 +7,7 @@ For quickstart usage, see `README.md`. v0.1 targets KernelHeim Weeks 0-2: - Week 0: copy/transpose -- Week 1: reductions (sum) variants +- Week 1: reductions (row-wise sum) - Week 2: single-pass online softmax Out of scope for v0: FlashAttention kernels, KV-cache decode, FP8, NCCL, C++ diff --git a/README.md b/README.md index ebb856f..7c62218 100644 --- a/README.md +++ b/README.md @@ -106,7 +106,7 @@ ncu --set full -o profiles/copy_transpose uv run python bench/benchmark_copy_tra | Op | Status | Variants | Notes | | --- | --- | --- | --- | | copy_transpose | Implemented | tile_size=16/32 | CuTe DSL kernel with tiled shared memory | -| reduce_sum | Stub (ref) | naive/improved/shfl | Uses PyTorch reference; kernel to be implemented | +| reduce_sum | Stub (ref) | row-wise | Uses PyTorch reference; kernel to be implemented | | softmax_online | Stub (ref) | single-pass | Uses PyTorch reference with autograd support; kernel to be implemented | --- @@ -174,7 +174,7 @@ For kernel development workflow and architecture details, see: ## Roadmap (v0.1 completion) * [x] Week 0 copy/transpose: end-to-end correctness + benchmark + profile scripts -* [ ] Week 1 reductions: multiple variants, correctness + benchmark coverage +* [ ] Week 1 reductions: row-wise sum, correctness + benchmark coverage * [ ] Week 2 online softmax: correctness + benchmark coverage + profiling notes * [ ] CI: run correctness on supported GPU runners; optional perf smoke checks diff --git a/ROADMAP.md b/ROADMAP.md index faf02c2..a1ec121 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -39,9 +39,9 @@ Target: Harness infrastructure + Week 0-2 kernel implementations aligned to Kern --- -## Week 1: Reductions (Sum) +## Week 1: Reductions (Row-wise Sum) -**Goal**: Multiple variants (naive → improved → shuffle) with correctness + benchmark coverage +**Goal**: Row-wise sum kernel with correctness + benchmark coverage ### Infrastructure - [x] Test infrastructure (using PyTorch reference) @@ -49,22 +49,12 @@ Target: Harness infrastructure + Week 0-2 kernel implementations aligned to Kern - [x] Ops registration and API design - [x] Benchmark integration -### Kernel Implementations -- [ ] **Naive variant**: Simple reduction without optimizations +### Kernel Implementation +- [ ] **Row-wise sum kernel** - [ ] CuTe DSL kernel implementation - [ ] Correctness tests vs PyTorch reference - [ ] Benchmark baseline -- [ ] **Improved variant**: Optimized reduction with shared memory - - [ ] CuTe DSL kernel implementation - - [ ] Correctness tests vs PyTorch reference - - [ ] Benchmark comparison vs naive - -- [ ] **Shuffle variant**: Warp-level shuffle reduction - - [ ] CuTe DSL kernel implementation - - [ ] Correctness tests vs PyTorch reference - - [ ] Benchmark comparison vs improved - - [ ] **Documentation**: Profiling notes and performance analysis **Status**: ⏳ Test infrastructure ready, kernels pending diff --git a/bench/benchmark_reduce_sum.py b/bench/benchmark_reduce_sum.py new file mode 100644 index 0000000..48cbdbb --- /dev/null +++ b/bench/benchmark_reduce_sum.py @@ -0,0 +1,98 @@ +"""Benchmark reduce_sum kernel variants against torch.sum.""" + +import argparse + +import torch + +from forge_cute_py.ops.reduce_sum import reduce_sum +from forge_cute_py.util.bench import do_bench, estimate_bandwidth, summarize_times + +DEFAULT_SIZES = [1024, 2048, 4096, 8192] +DEFAULT_DTYPES = ["float16", "bfloat16", "float32"] +DEFAULT_DIMS = [1] + + +def parse_int_list(s: str) -> list[int]: + return [int(x.strip()) for x in s.split(",")] + + +def parse_str_list(s: str) -> list[str]: + return [x.strip() for x in s.split(",")] + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark reduce_sum variants") + parser.add_argument("--sizes", type=parse_int_list, default=DEFAULT_SIZES) + parser.add_argument("--dtypes", type=parse_str_list, default=DEFAULT_DTYPES) + parser.add_argument("--dims", type=parse_int_list, default=DEFAULT_DIMS) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA required for benchmarking") + + gpu_name = torch.cuda.get_device_name(0) + print(f"reduce_sum benchmarks ({gpu_name})") + print() + + header = f"{'Size':>6} {'Dtype':<10} {'Dim':>3} {'Op':<12} {'p50 (ms)':>10} {'BW (GB/s)':>10} {'vs torch':>10}" + print(header) + print("-" * len(header)) + + for size in args.sizes: + for dtype_str in args.dtypes: + dtype = getattr(torch, dtype_str) + for dim in args.dims: + x = torch.randn(size, size, device="cuda", dtype=dtype) + + # Output shape determines output bytes + if dim == 1: + out_numel = size + else: + out_numel = size + input_bytes = x.numel() * x.element_size() + output_bytes = out_numel * x.element_size() + total_bytes = input_bytes + output_bytes + + # Benchmark torch.sum as reference + torch_fn = lambda: torch.sum(x, dim=dim) + torch_times = do_bench(torch_fn, warmup=args.warmup, rep=args.iterations) + torch_stats = summarize_times(torch_times) + torch_p50 = torch_stats["p50_ms"] + torch_bw = estimate_bandwidth(total_bytes, torch_p50) + print( + f"{size:>6} {dtype_str:<10} {dim:>3} {'torch.sum':<14} " + f"{torch_p50:>10.4f} {torch_bw:>10.2f} {1.0:>10.2f}x" + ) + + if dim not in (-1, 1): + raise ValueError(f"Unsupported dim={dim}; reduce_sum supports only -1 or 1") + + # Warm the JIT cache + try: + reduce_sum(x, dim=dim) + except Exception as e: + print( + f"{size:>6} {dtype_str:<10} {dim:>3} {'reduce_sum':<12} " + f"{'ERROR':>10} {'':>10} {'':>10} {e}" + ) + print() + continue + + fn = lambda d=dim: reduce_sum(x, dim=d) + times = do_bench(fn, warmup=args.warmup, rep=args.iterations) + stats = summarize_times(times) + p50 = stats["p50_ms"] + bw = estimate_bandwidth(total_bytes, p50) + ratio = p50 / torch_p50 if torch_p50 > 0 else float("inf") + print( + f"{size:>6} {dtype_str:<10} {dim:>3} {'reduce_sum':<12} " + f"{p50:>10.4f} {bw:>10.2f} {ratio:>10.2f}x" + ) + + print() + + +if __name__ == "__main__": + main() diff --git a/bench/run.py b/bench/run.py index c6255fa..7430841 100644 --- a/bench/run.py +++ b/bench/run.py @@ -84,16 +84,12 @@ def fn(): if not hasattr(ops, "reduce_sum"): return {"status": "skipped", "reason": "reduce_sum not available"} dim = case.get("dim", -1) - variant = case.get("variant", "shfl") x = torch.randn(*shape, device="cuda", dtype=dtype) def fn(): - return ops.reduce_sum(x, dim=dim, variant=variant) + return ops.reduce_sum(x, dim=dim) - try: - times = do_bench(fn, warmup=warmup, rep=iterations) - except NotImplementedError: - return {"status": "skipped", "reason": f"variant {variant} not implemented"} + times = do_bench(fn, warmup=warmup, rep=iterations) stats = summarize_times(times) bytes_moved = _estimate_bytes(op_name, shape, dtype, dim=dim) bw = estimate_bandwidth(bytes_moved, stats["p50_ms"]) @@ -103,7 +99,6 @@ def fn(): "shape": shape, "dtype": str(dtype).replace("torch.", ""), "dim": dim, - "variant": variant, "times_ms": stats, "bandwidth_gbps": bw, } diff --git a/bench/suites.yaml b/bench/suites.yaml index 39540c7..349faa1 100644 --- a/bench/suites.yaml +++ b/bench/suites.yaml @@ -11,7 +11,6 @@ suites: shape: [1024, 1024] dtype: float16 dim: -1 - variant: shfl - op: softmax_online shape: [256, 2048] dtype: float16 diff --git a/forge_cute_py/kernels/reduce_sum.py b/forge_cute_py/kernels/reduce_sum.py new file mode 100644 index 0000000..1bcf191 --- /dev/null +++ b/forge_cute_py/kernels/reduce_sum.py @@ -0,0 +1,157 @@ +""" +Row-wise sum reduction kernel using CuTe DSL (WIP). +""" + +import operator +from typing import Callable + +import cutlass +import cutlass.cute as cute +from cutlass import const_expr + + +@cute.jit +def block_reduce( + val: cute.Numeric, + op: Callable, + reduction_buffer: cute.Tensor, + init_val: cute.Numeric, +) -> cute.Numeric: + """Block-wide reduction using warp shuffles + shared memory cross-warp step.""" + lane_idx = cute.arch.lane_idx() + warp_idx = cute.arch.warp_idx() + num_warps = cute.size(reduction_buffer.shape[0]) + + if lane_idx == 0: + reduction_buffer[warp_idx] = val + cute.arch.barrier() + + block_reduce_val = init_val + if lane_idx < num_warps: + block_reduce_val = reduction_buffer[lane_idx] + return cute.arch.warp_reduction(block_reduce_val, op) + + +@cute.jit +def row_reduce( + x: cute.TensorSSA | cute.Numeric, + op: cute.ReductionOp, + threads_per_row: cutlass.Constexpr[int], + reduction_buffer: cute.Tensor, + init_val: cute.Numeric = 0.0, +) -> cute.Numeric: + """Thread + warp + block reduction for a single row.""" + if const_expr(isinstance(x, cute.TensorSSA)): + val = x.reduce(op, init_val=init_val, reduction_profile=0) + else: + val = x + warp_op = { + cute.ReductionOp.ADD: operator.add, + cute.ReductionOp.MAX: cute.arch.fmax, + cute.ReductionOp.MIN: min, + cute.ReductionOp.MUL: operator.mul, + }[op] + val = cute.arch.warp_reduction( + val, + warp_op, + threads_in_group=min(threads_per_row, cute.arch.WARP_SIZE), + ) + if const_expr(cute.size(reduction_buffer.shape[0]) > 1): + val = block_reduce(val, warp_op, reduction_buffer, init_val=init_val) + return val + + +class ReduceSumRow: + """Row-wise sum reduction (one block per row).""" + + def __init__(self, dtype: type, block_size: int = 256): + self.dtype = dtype + self.block_size = block_size + + @cute.jit + def __call__( + self, + x: cute.Tensor, + out: cute.Tensor, + ): + """ + Launch row-wise reduction. + + Args: + x: Input tensor of shape (M, N) + out: Output tensor of shape (M,) + stream: CUDA stream + """ + M, _ = x.shape + block = const_expr(self.block_size) + self.kernel(x, out).launch( + grid=[M, 1, 1], + block=[block, 1, 1], + ) + + @cute.kernel + def kernel( + self, + x: cute.Tensor, + out: cute.Tensor, + ): + """Row-wise sum reduction with vectorized loads and block reduction.""" + tidx, _, _ = cute.arch.thread_idx() + row_idx, _, _ = cute.arch.block_idx() + + threads_per_block = const_expr(self.block_size) + num_warps = threads_per_block // 32 + + smem = cutlass.utils.SmemAllocator() + warp_sums = smem.allocate_tensor( + cutlass.Float32, layout=cute.make_layout(num_warps), byte_alignment=16 + ) + + row = x[(row_idx, None)] + n_cols = row.shape[0] + acc = cutlass.Float32(0.0) + + vec_size = const_expr(128 // x.element_type.width) + tile_n = threads_per_block * vec_size + num_full_tiles = n_cols // tile_n + + sX = smem.allocate_tensor( + x.element_type, layout=cute.make_layout(tile_n), byte_alignment=16 + ) + copy_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + x.element_type, + num_bits_per_copy=128, + ) + thr_layout = cute.make_layout(threads_per_block) + val_layout = cute.make_layout(vec_size) + tiled_copy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout) + thr_copy = tiled_copy.get_slice(tidx) + + for tile in cutlass.range(0, num_full_tiles, 1): + gX = cute.local_tile(row, (tile_n,), (tile,)) + tXgX = thr_copy.partition_S(gX) + tXsX = thr_copy.partition_D(sX) + tXrX = cute.make_fragment_like(tXgX) + cute.copy(copy_atom, tXgX, tXsX) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + + cute.autovec_copy(tXsX, tXrX) + x_vec = tXrX.load().to(cutlass.Float32) + acc += x_vec.reduce(cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0) + + tail_start = num_full_tiles * tile_n + tidx + for col in cutlass.range(tail_start, n_cols, threads_per_block): + acc += row[col].to(cutlass.Float32) + + block_val = row_reduce( + acc, + cute.ReductionOp.ADD, + threads_per_row=threads_per_block, + reduction_buffer=warp_sums, + init_val=cutlass.Float32(0.0), + ) + if tidx == 0: + out[row_idx] = block_val.to(out.element_type) diff --git a/forge_cute_py/ops/reduce_sum.py b/forge_cute_py/ops/reduce_sum.py index d501eca..275af50 100644 --- a/forge_cute_py/ops/reduce_sum.py +++ b/forge_cute_py/ops/reduce_sum.py @@ -2,18 +2,17 @@ @torch.library.custom_op("forge_cute_py::_reduce_sum", mutates_args={"out"}) -def _reduce_sum(x: torch.Tensor, out: torch.Tensor, dim: int = -1, variant: str = "shfl") -> None: - """Row/column sum reduction (reference implementation stub). +def _reduce_sum(x: torch.Tensor, out: torch.Tensor, dim: int = -1) -> None: + """Row-wise sum reduction (reference implementation stub). Args: x: Input tensor of shape (M, N) out: Output tensor (mutated in-place) - dim: Dimension to reduce over (-1, 0, or 1) - variant: Reduction variant (naive, improved, shfl) - currently unused + dim: Dimension to reduce over (-1 or 1) """ assert x.dim() == 2, "reduce_sum expects a 2D tensor" assert x.is_cuda, f"reduce_sum is CUDA-only, got device={x.device}" - assert dim in (-1, 0, 1), f"reduce_sum expects dim in {{-1, 0, 1}}, got {dim}" + assert dim in (-1, 1), f"reduce_sum expects dim in {{-1, 1}}, got {dim}" assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], ( f"Unsupported dtype: {x.dtype}" ) @@ -22,7 +21,7 @@ def _reduce_sum(x: torch.Tensor, out: torch.Tensor, dim: int = -1, variant: str dim = dim if dim >= 0 else x.ndim + dim # For now, use reference implementation - # Future: call kernel implementation based on variant when available + # Future: call kernel implementation when available from forge_cute_py.ref import reduce_sum as reduce_sum_ref result = reduce_sum_ref(x, dim=dim) @@ -32,16 +31,15 @@ def _reduce_sum(x: torch.Tensor, out: torch.Tensor, dim: int = -1, variant: str _reduce_sum.compile_cache = {} -def reduce_sum(x: torch.Tensor, dim: int = -1, variant: str = "shfl") -> torch.Tensor: - """Row/column sum reduction. +def reduce_sum(x: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Row-wise sum reduction. Args: x: Input tensor of shape (M, N) - dim: Dimension to reduce over (-1 for last dim, 0 or 1) - variant: Reduction variant (naive, improved, shfl) - currently unused + dim: Dimension to reduce over (-1 or 1) Returns: - Reduced tensor of shape (M,) if dim=1 or (N,) if dim=0 + Reduced tensor of shape (M,) Examples: >>> x = torch.randn(32, 128, device='cuda', dtype=torch.float16) @@ -53,13 +51,11 @@ def reduce_sum(x: torch.Tensor, dim: int = -1, variant: str = "shfl") -> torch.T dim = dim if dim >= 0 else x.ndim + dim # Determine output shape - if dim == 0: - out_shape = (x.shape[1],) - elif dim == 1: + if dim == 1: out_shape = (x.shape[0],) else: raise ValueError(f"Invalid dim={dim} for 2D tensor") out = torch.empty(out_shape, dtype=x.dtype, device=x.device) - _reduce_sum(x, out, dim, variant) + _reduce_sum(x, out, dim) return out diff --git a/forge_cute_py/ref/reduce_sum.py b/forge_cute_py/ref/reduce_sum.py index 5b9c3ba..63c1fe5 100644 --- a/forge_cute_py/ref/reduce_sum.py +++ b/forge_cute_py/ref/reduce_sum.py @@ -1,10 +1,9 @@ import torch -def reduce_sum(x: torch.Tensor, dim: int = -1, variant: str = "shfl") -> torch.Tensor: +def reduce_sum(x: torch.Tensor, dim: int = -1) -> torch.Tensor: if x.ndim != 2: raise ValueError("reduce_sum expects a 2D tensor") - if dim not in (-1, 0, 1): - raise ValueError("reduce_sum expects dim in {-1, 0, 1} for 2D tensors") - x_dtype = x.dtype - return x.float().sum(dim=dim).to(x_dtype) + if dim not in (-1, 1): + raise ValueError("reduce_sum expects dim in {-1, 1} for 2D tensors") + return x.sum(dim=dim) diff --git a/tests/test_reduce_sum.py b/tests/test_reduce_sum.py index dd773b5..e1ab65c 100644 --- a/tests/test_reduce_sum.py +++ b/tests/test_reduce_sum.py @@ -9,7 +9,7 @@ "shape, dim", [ ((4, 8), -1), - ((8, 4), 0), + ((8, 4), 1), ], ) @pytest.mark.parametrize( @@ -20,13 +20,9 @@ (torch.bfloat16, 1e-2, 1e-2), ], ) -@pytest.mark.parametrize("variant", ["naive", "improved", "shfl"]) -def test_reduce_sum_correctness(shape, dim, dtype, atol, rtol, variant): +def test_reduce_sum_correctness(shape, dim, dtype, atol, rtol): x = torch.randn(*shape, device="cuda", dtype=dtype) - try: - y = reduce_sum(x, dim=dim, variant=variant) - except NotImplementedError: - pytest.skip(f"reduce_sum variant {variant} not implemented") + y = reduce_sum(x, dim=dim) y_ref = ref_reduce_sum(x, dim=dim) torch.testing.assert_close(y, y_ref, atol=atol, rtol=rtol) @@ -50,7 +46,5 @@ def test_reduce_sum_torch_compile(): y = compiled(x) except unsupported_exc as exc: pytest.skip(f"torch.compile unsupported for reduce_sum op: {exc}") - except NotImplementedError: - pytest.skip("reduce_sum shfl variant not implemented") y_ref = ref_reduce_sum(x, -1) torch.testing.assert_close(y, y_ref, atol=1e-2, rtol=1e-2)