Upload triton_v2.py with huggingface_hub

triton_v2.py

@@ -0,0 +1,419 @@
+#!/usr/bin/env python3
+"""
+Triton-fused Chunked Sparse Backward Pass — v2.
+
+Fixes from review:
+  1. Bias folded into dW kernel (kills the uncoalesced column-striding bias kernel)
+  2. block_ptr / TMA for dW and dX loads (hardware-accelerated 2D tile fetch)
+  3. No autotune (fixed config to eliminate compilation overhead + divergence risk)
+
+Benchmarks v1 (manual ptrs + separate bias) vs v2 (block_ptr + fused bias)
+vs Python-loop baseline vs Dense.
+"""
+
+import math, os, time
+import torch, torch.nn as nn, torch.nn.functional as F
+import triton, triton.language as tl
+
+# ═══════════════════════════════════════════════════════════════════
+# V2 KERNELS — block_ptr + fused bias
+# ═══════════════════════════════════════════════════════════════════
+
+# Fixed tile sizes — no autotune. CS=64 means one N-block covers the whole chunk.
+# BM=64: token tile for the M-reduction loop
+# BK=64: tile along d_in
+# BN=64: tile along chunk (== CS for chunk_size=64, so 1 block per chunk)
+
+@triton.jit
+def _v2_sparse_bwd_dW_kernel(
+    X_ptr, dY_ptr, dW_ptr, dB_ptr, chunk_ids_ptr,
+    M, d_in, d_out, num_active,
+    stride_xm, stride_xk,
+    stride_dym, stride_dyn,
+    stride_dwn, stride_dwk,
+    HAS_BIAS: tl.constexpr,
+    CS: tl.constexpr,
+    BK: tl.constexpr,
+    BM: tl.constexpr,
+):
+    """
+    Each program computes one [CS, BK] tile of dW for one active chunk,
+    plus the [CS] bias slice if HAS_BIAS.
+
+    Grid: (num_active, ceil(d_in / BK))
+    Since CS fits in one tile (CS==64, BN==CS), pid0 == chunk index directly.
+    """
+    chunk_linear_id = tl.program_id(0)
+    k_block_id = tl.program_id(1)
+
+    if chunk_linear_id >= num_active:
+        return
+
+    chunk_idx = tl.load(chunk_ids_ptr + chunk_linear_id)
+    chunk_start = chunk_idx * CS
+
+    k_offset = k_block_id * BK
+
+    # Block pointer for dY transposed: we want dY.T[chunk_cols, :] = shape (CS, M)
+    # dY is (M, d_out) row-major. Transposed view: shape=(d_out, M), strides=(stride_dyn, stride_dym)
+    dy_block_ptr = tl.make_block_ptr(
+        base=dY_ptr,
+        shape=(d_out, M),
+        strides=(stride_dyn, stride_dym),
+        offsets=(chunk_start, 0),
+        block_shape=(CS, BM),
+        order=(1, 0),
+    )
+
+    # Block pointer for X: shape (M, d_in), reading (BM, BK) tiles
+    x_block_ptr = tl.make_block_ptr(
+        base=X_ptr,
+        shape=(M, d_in),
+        strides=(stride_xm, stride_xk),
+        offsets=(0, k_offset),
+        block_shape=(BM, BK),
+        order=(1, 0),
+    )
+
+    # Accumulators
+    acc_dw = tl.zeros((CS, BK), dtype=tl.float32)
+    # Bias accumulator: only on the first k-block to avoid redundant work
+    compute_bias = HAS_BIAS and (k_block_id == 0)
+    acc_db = tl.zeros((CS,), dtype=tl.float32)
+
+    # Reduction over M
+    for m_start in range(0, M, BM):
+        dy_t = tl.load(dy_block_ptr, boundary_check=(0, 1))   # (CS, BM)
+        x = tl.load(x_block_ptr, boundary_check=(0, 1))       # (BM, BK)
+
+        # dW += dY.T @ X  ->  (CS, BM) @ (BM, BK) = (CS, BK)
+        acc_dw = tl.dot(dy_t, x, acc=acc_dw)
+
+        # Bias: sum over M dimension of dY chunk columns
+        # dy_t is (CS, BM) = transposed chunk. Sum along dim=1 = sum over tokens.
+        if compute_bias:
+            acc_db += tl.sum(dy_t, axis=1)
+
+        dy_block_ptr = tl.advance(dy_block_ptr, (0, BM))
+        x_block_ptr = tl.advance(x_block_ptr, (BM, 0))
+
+    # Store dW tile: dW[chunk_start:chunk_start+CS, k_offset:k_offset+BK]
+    dw_block_ptr = tl.make_block_ptr(
+        base=dW_ptr,
+        shape=(d_out, d_in),
+        strides=(stride_dwn, stride_dwk),
+        offsets=(chunk_start, k_offset),
+        block_shape=(CS, BK),
+        order=(1, 0),
+    )
+    tl.store(dw_block_ptr, acc_dw.to(dW_ptr.dtype.element_ty), boundary_check=(0, 1))
+
+    # Store bias (only from k_block_id == 0)
+    if compute_bias:
+        rn = chunk_start + tl.arange(0, CS)
+        n_mask = rn < d_out
+        tl.store(dB_ptr + rn, acc_db.to(dB_ptr.dtype.element_ty), mask=n_mask)
+
+
+def v2_sparse_bwd_dW(X, dY, active_chunks, chunk_size, d_out, bias=True):
+    """Fused dW + dBias via block_ptr kernel."""
+    M, d_in = X.shape
+    num_active = active_chunks.shape[0]
+    CS = chunk_size
+
+    dW = torch.zeros(d_out, d_in, device=X.device, dtype=X.dtype)
+    dB = torch.zeros(d_out, device=X.device, dtype=X.dtype) if bias else None
+
+    if num_active == 0:
+        return dW, dB
+
+    chunk_ids = active_chunks.to(torch.int32).contiguous()
+
+    BK = 64
+    BM = 64
+
+    grid = (num_active, triton.cdiv(d_in, BK))
+
+    _v2_sparse_bwd_dW_kernel[grid](
+        X, dY, dW, dB if bias else X,  # dummy ptr if no bias
+        chunk_ids,
+        M, d_in, d_out, num_active,
+        X.stride(0), X.stride(1),
+        dY.stride(0), dY.stride(1),
+        dW.stride(0), dW.stride(1),
+        HAS_BIAS=bias,
+        CS=CS, BK=BK, BM=BM,
+    )
+    return dW, dB
+
+
+# ── V2 dX kernel with block_ptr ──
+
+@triton.jit
+def _v2_sparse_bwd_dX_kernel(
+    dY_ptr, W_ptr, dX_ptr, chunk_ids_ptr,
+    M, d_in, d_out, num_active,
+    stride_dym, stride_dyn,
+    stride_wn, stride_wk,
+    stride_dxm, stride_dxk,
+    CS: tl.constexpr,
+    BM: tl.constexpr,
+    BK: tl.constexpr,
+):
+    """
+    Each program computes one [BM, BK] tile of dX by accumulating over active chunks.
+    Grid: (ceil(M/BM), ceil(d_in/BK))
+    """
+    pid_m = tl.program_id(0)
+    pid_k = tl.program_id(1)
+
+    m_offset = pid_m * BM
+    k_offset = pid_k * BK
+
+    acc = tl.zeros((BM, BK), dtype=tl.float32)
+
+    for i in range(num_active):
+        chunk_idx = tl.load(chunk_ids_ptr + i)
+        chunk_start = chunk_idx * CS
+
+        # dY tile: (BM, CS) at [m_offset, chunk_start]
+        dy_block_ptr = tl.make_block_ptr(
+            base=dY_ptr,
+            shape=(M, d_out),
+            strides=(stride_dym, stride_dyn),
+            offsets=(m_offset, chunk_start),
+            block_shape=(BM, CS),
+            order=(1, 0),
+        )
+
+        # W tile: (CS, BK) at [chunk_start, k_offset]
+        w_block_ptr = tl.make_block_ptr(
+            base=W_ptr,
+            shape=(d_out, d_in),
+            strides=(stride_wn, stride_wk),
+            offsets=(chunk_start, k_offset),
+            block_shape=(CS, BK),
+            order=(1, 0),
+        )
+
+        dy = tl.load(dy_block_ptr, boundary_check=(0, 1))  # (BM, CS)
+        w = tl.load(w_block_ptr, boundary_check=(0, 1))    # (CS, BK)
+
+        # dY @ W -> (BM, BK)
+        acc = tl.dot(dy, w, acc=acc)
+
+    # Store dX tile
+    dx_block_ptr = tl.make_block_ptr(
+        base=dX_ptr,
+        shape=(M, d_in),
+        strides=(stride_dxm, stride_dxk),
+        offsets=(m_offset, k_offset),
+        block_shape=(BM, BK),
+        order=(1, 0),
+    )
+    tl.store(dx_block_ptr, acc.to(dX_ptr.dtype.element_ty), boundary_check=(0, 1))
+
+
+def v2_sparse_bwd_dX(dY, W, active_chunks, chunk_size, M, d_in):
+    """Fused dX via block_ptr kernel."""
+    num_active = active_chunks.shape[0]
+    d_out = dY.shape[1]
+    CS = chunk_size
+
+    dX = torch.zeros(M, d_in, device=dY.device, dtype=dY.dtype)
+    if num_active == 0:
+        return dX
+
+    chunk_ids = active_chunks.to(torch.int32).contiguous()
+
+    BM = 64
+    BK = 64
+
+    grid = (triton.cdiv(M, BM), triton.cdiv(d_in, BK))
+
+    _v2_sparse_bwd_dX_kernel[grid](
+        dY, W, dX, chunk_ids,
+        M, d_in, d_out, num_active,
+        dY.stride(0), dY.stride(1),
+        W.stride(0), W.stride(1),
+        dX.stride(0), dX.stride(1),
+        CS=CS, BM=BM, BK=BK,
+    )
+    return dX
+
+
+# ═══════════════════════════════════════════════════════════════════
+# V1 KERNELS (old, for comparison) — import from triton_sparse.py
+# ═══════════════════════════════════════════════════════════════════
+
+from triton_sparse import (
+    sparse_bwd_dW as v1_sparse_bwd_dW,
+    sparse_bwd_dX as v1_sparse_bwd_dX,
+    sparse_bwd_dbias as v1_sparse_bwd_dbias,
+)
+
+
+# ═══════════════════════════════════════════════════════════════════
+# CORRECTNESS TEST
+# ═══════════════════════════════════════════════════════════════════
+
+def test_correctness():
+    print("V2 Correctness Tests")
+    print("=" * 60)
+    device = "cuda"
+    torch.manual_seed(42)
+
+    for d_in, d_out, cs in [(512, 2048, 64), (1024, 4096, 64), (256, 1024, 64)]:
+        M = 2048
+        n_chunks = d_out // cs
+        n_active = max(1, int(0.1 * n_chunks))
+        active = torch.randperm(n_chunks, device=device)[:n_active].sort().values
+
+        x = torch.randn(M, d_in, device=device)
+        w = torch.randn(d_out, d_in, device=device)
+        gy = torch.randn(M, d_out, device=device)
+
+        # Reference
+        ref_dw = torch.zeros_like(w)
+        ref_db = torch.zeros(d_out, device=device)
+        for c in active.tolist():
+            s, e = c * cs, (c + 1) * cs
+            ref_dw[s:e] = gy[:, s:e].t() @ x
+            ref_db[s:e] = gy[:, s:e].sum(0)
+
+        ref_dx = torch.zeros_like(x)
+        for c in active.tolist():
+            s, e = c * cs, (c + 1) * cs
+            ref_dx += gy[:, s:e] @ w[s:e]
+
+        # V2
+        v2_dw, v2_db = v2_sparse_bwd_dW(x, gy, active, cs, d_out, bias=True)
+        v2_dx = v2_sparse_bwd_dX(gy, w, active, cs, M, d_in)
+
+        dw_err = (v2_dw - ref_dw).abs().max().item()
+        db_err = (v2_db - ref_db).abs().max().item()
+        dx_err = (v2_dx - ref_dx).abs().max().item()
+
+        ok = dw_err < 0.01 and db_err < 0.01 and dx_err < 0.01
+        print(f"  {'✓' if ok else '✗'} d_in={d_in} d_out={d_out} cs={cs}: dW={dw_err:.6f} dB={db_err:.6f} dX={dx_err:.6f}")
+
+    print()
+
+
+# ═══════════════════════════════════════════════════════════════════
+# BENCHMARK
+# ═══════════════════════════════════════════════════════════════════
+
+def benchmark():
+    print("=" * 100)
+    print("BENCHMARK: Dense vs PyLoop vs V1-Triton vs V2-Triton (block_ptr + fused bias)")
+    print("=" * 100)
+    device = "cuda"
+    B, T = 8, 256
+    M = B * T
+    cs = 64
+    af = 0.10
+    warmup = 20
+    iters = 100
+
+    print(f"\nM={M}, chunk_size={cs}, active_frac={af}, {iters} iters after {warmup} warmup")
+    print(f"{'d':>5} | {'ffn':>5} | {'act':>3} | {'Dense':>9} | {'PyLoop':>9} | {'V1-Tri':>9} | {'V2-Tri':>9} | {'V2/Dense':>9} | {'V2/PyLoop':>9} | {'V2/V1':>9}")
+    print("-" * 105)
+
+    for d_in in [256, 512, 768, 1024, 1536, 2048]:
+        d_out = 4 * d_in
+        nc = d_out // cs
+        na = max(1, int(af * nc))
+        active = torch.randperm(nc, device=device)[:na].sort().values
+
+        x = torch.randn(M, d_in, device=device)
+        w = torch.randn(d_out, d_in, device=device)
+        gy = torch.randn(M, d_out, device=device)
+
+        def dense():
+            return gy.t() @ x, gy @ w, gy.sum(0)
+
+        def pyloop():
+            dw = torch.zeros_like(w); db = torch.zeros(d_out, device=device)
+            dx = gy @ w
+            for c in active.tolist():
+                s, e = c*cs, (c+1)*cs
+                dw[s:e] = gy[:, s:e].t() @ x
+                db[s:e] = gy[:, s:e].sum(0)
+            return dw, dx, db
+
+        def v1_tri():
+            dw = v1_sparse_bwd_dW(x, gy, active, cs, d_out)
+            dx = gy @ w
+            db = v1_sparse_bwd_dbias(gy, active, cs, d_out)
+            return dw, dx, db
+
+        def v2_tri():
+            dw, db = v2_sparse_bwd_dW(x, gy, active, cs, d_out, bias=True)
+            dx = gy @ w
+            return dw, dx, db
+
+        # Warmup all
+        for _ in range(warmup):
+            dense(); pyloop(); v1_tri(); v2_tri()
+        torch.cuda.synchronize()
+
+        times = {}
+        for name, fn in [("dense", dense), ("pyloop", pyloop), ("v1", v1_tri), ("v2", v2_tri)]:
+            torch.cuda.synchronize(); t0 = time.perf_counter()
+            for _ in range(iters): fn()
+            torch.cuda.synchronize()
+            times[name] = (time.perf_counter() - t0) / iters
+
+        td, tp, t1, t2 = times["dense"], times["pyloop"], times["v1"], times["v2"]
+        print(f"{d_in:>5} | {d_out:>5} | {na:>3} | {td*1000:>8.2f}ms | {tp*1000:>8.2f}ms | {t1*1000:>8.2f}ms | {t2*1000:>8.2f}ms | {td/t2:>8.2f}x | {tp/t2:>8.2f}x | {t1/t2:>8.2f}x")
+
+    # Sparse dX comparison: V1 vs V2
+    print(f"\n{'='*80}")
+    print("Sparse dX (both dW+dX sparse): V1 vs V2")
+    print(f"{'d':>5} | {'Dense':>9} | {'V1-all':>9} | {'V2-all':>9} | {'V2/Dense':>9}")
+    print("-" * 55)
+
+    for d_in in [512, 1024, 2048]:
+        d_out = 4 * d_in; nc = d_out // cs; na = max(1, int(af * nc))
+        active = torch.randperm(nc, device=device)[:na].sort().values
+        x = torch.randn(M, d_in, device=device)
+        w = torch.randn(d_out, d_in, device=device)
+        gy = torch.randn(M, d_out, device=device)
+
+        def dense_all():
+            return gy.t() @ x, gy @ w
+        def v1_all():
+            return v1_sparse_bwd_dW(x, gy, active, cs, d_out), v1_sparse_bwd_dX(gy, w, active, cs, M, d_in)
+        def v2_all():
+            dw, _ = v2_sparse_bwd_dW(x, gy, active, cs, d_out, bias=False)
+            return dw, v2_sparse_bwd_dX(gy, w, active, cs, M, d_in)
+
+        for _ in range(warmup): dense_all(); v1_all(); v2_all()
+        torch.cuda.synchronize()
+
+        for name, fn, store in [("dense", dense_all, "td"), ("v1", v1_all, "t1"), ("v2", v2_all, "t2")]:
+            torch.cuda.synchronize(); t0 = time.perf_counter()
+            for _ in range(iters): fn()
+            torch.cuda.synchronize()
+            locals()[store] = (time.perf_counter() - t0) / iters
+
+        # Need to read them back since locals() trick doesn't work cleanly
+        torch.cuda.synchronize(); t0 = time.perf_counter()
+        for _ in range(iters): dense_all()
+        torch.cuda.synchronize(); td = (time.perf_counter() - t0) / iters
+
+        torch.cuda.synchronize(); t0 = time.perf_counter()
+        for _ in range(iters): v1_all()
+        torch.cuda.synchronize(); t1 = (time.perf_counter() - t0) / iters
+
+        torch.cuda.synchronize(); t0 = time.perf_counter()
+        for _ in range(iters): v2_all()
+        torch.cuda.synchronize(); t2 = (time.perf_counter() - t0) / iters
+
+        print(f"{d_in:>5} | {td*1000:>8.2f}ms | {t1*1000:>8.2f}ms | {t2*1000:>8.2f}ms | {td/t2:>8.2f}x")
+
+
+if __name__ == "__main__":
+    test_correctness()
+    benchmark()