add backward support

Author: xsank

fastvideo-kernel/benchmarks/bench_fused_compress_topk.py

@@ -13,6 +13,8 @@
 
 Reports per-kernel latency (ms), speedup, and numerical accuracy (max abs error,
 cosine similarity for compress; mask match rate for topk).
+
+Also benchmarks backward pass of compress (fused Triton bwd kernel vs. PyTorch autograd).
 """
 
 from __future__ import annotations
@@ -116,14 +118,13 @@ def accuracy_topk(ref_mask: torch.Tensor, test_mask: torch.Tensor) -> dict:
 # Benchmark runner
 # ---------------------------------------------------------------------------
 
-def bench_compress(
+def bench_compress_fwd(
     B: int, H: int, seq_len: int, D: int, block_elements: int,
     dtype: torch.dtype, warmup: int, rep: int,
 ) -> None:
     num_blocks = seq_len // block_elements
     x = torch.randn(B, H, seq_len, D, dtype=dtype, device="cuda")
     vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device="cuda")
-    # Make a few blocks partially filled to exercise variable block sizes
     if num_blocks > 4:
         vbs[1] = block_elements - 2
         vbs[-2] = block_elements - 5
@@ -138,7 +139,58 @@ def bench_compress(
     new_ms = do_bench(lambda: fused_block_mean(x, vbs, block_elements), warmup=warmup, rep=rep)
 
     speedup = old_ms / new_ms if new_ms > 0 else float("inf")
-    print(f"  compress  | old: {old_ms:8.3f} ms | new: {new_ms:8.3f} ms | speedup: {speedup:5.2f}x "
+    print(f"  compress fwd | old: {old_ms:8.3f} ms | new: {new_ms:8.3f} ms | speedup: {speedup:5.2f}x "
+          f"| max_abs_err: {acc['max_abs_err']:.2e} | cos_sim: {acc['cosine_sim']:.8f}")
+
+
+def bench_compress_bwd(
+    B: int, H: int, seq_len: int, D: int, block_elements: int,
+    dtype: torch.dtype, warmup: int, rep: int,
+) -> None:
+    num_blocks = seq_len // block_elements
+    vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device="cuda")
+    if num_blocks > 4:
+        vbs[1] = block_elements - 2
+        vbs[-2] = block_elements - 5
+
+    # --- Accuracy: compare gradients ---
+    x_old = torch.randn(B, H, seq_len, D, dtype=dtype, device="cuda", requires_grad=True)
+    grad_out = torch.randn(B, H, num_blocks, D, dtype=dtype, device="cuda")
+
+    out_old = pytorch_block_mean(x_old, vbs, block_elements)
+    out_old.backward(grad_out)
+    grad_ref = x_old.grad.clone()
+
+    x_new = x_old.detach().clone().requires_grad_(True)
+    out_new = fused_block_mean(x_new, vbs, block_elements)
+    out_new.backward(grad_out)
+    grad_fused = x_new.grad.clone()
+
+    acc = accuracy_compress(grad_ref, grad_fused)
+
+    # --- Latency: isolate backward-only via retain_graph ---
+    x_o = x_old.detach().clone().requires_grad_(True)
+    out_o = pytorch_block_mean(x_o, vbs, block_elements)
+    loss_o = (out_o * grad_out).sum()
+    for _ in range(warmup):
+        torch.autograd.grad(loss_o, x_o, retain_graph=True)
+    old_ms = do_bench(
+        lambda: torch.autograd.grad(loss_o, x_o, retain_graph=True),
+        warmup=0, rep=rep,
+    )
+
+    x_n = x_old.detach().clone().requires_grad_(True)
+    out_n = fused_block_mean(x_n, vbs, block_elements)
+    loss_n = (out_n * grad_out).sum()
+    for _ in range(warmup):
+        torch.autograd.grad(loss_n, x_n, retain_graph=True)
+    new_ms = do_bench(
+        lambda: torch.autograd.grad(loss_n, x_n, retain_graph=True),
+        warmup=0, rep=rep,
+    )
+
+    speedup = old_ms / new_ms if new_ms > 0 else float("inf")
+    print(f"  compress bwd | old: {old_ms:8.3f} ms | new: {new_ms:8.3f} ms | speedup: {speedup:5.2f}x "
           f"| max_abs_err: {acc['max_abs_err']:.2e} | cos_sim: {acc['cosine_sim']:.8f}")
 
 
@@ -188,7 +240,8 @@ def main() -> None:
         print(f"seq_len={seq_len}, num_blocks={num_blocks}, topk={topk}")
         print("-" * 100)
 
-        bench_compress(B, H, seq_len, D, block_elements, dtype, args.warmup, args.rep)
+        bench_compress_fwd(B, H, seq_len, D, block_elements, dtype, args.warmup, args.rep)
+        bench_compress_bwd(B, H, seq_len, D, block_elements, dtype, args.warmup, args.rep)
         bench_topk(B, H, num_blocks, topk, dtype, args.warmup, args.rep)
 
 

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/fused_compress_topk.py

@@ -58,21 +58,75 @@ def _fused_block_mean_kernel(
     tl.store(out_base, acc.to(tl.bfloat16))
 
 
-def fused_block_mean(
-    x: torch.Tensor,
+@triton.jit
+def _fused_block_mean_bwd_kernel(
+    GradOut_ptr,
+    GradX_ptr,
+    VBS_ptr,
+    stride_go_bh,
+    stride_go_blk,
+    stride_gx_bh,
+    stride_gx_seq,
+    num_blocks,
+    BLOCK_ELEMENTS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+):
+    """Backward of block mean: broadcast grad_out / vbs to each token in the block.
+
+    Mirrors the forward kernel: one program per (block, bh).
+    GradOut is [B*H, num_blocks, HEAD_DIM].
+    GradX  is [B*H, num_blocks*BLOCK_ELEMENTS, HEAD_DIM].
+    """
+    block_idx = tl.program_id(0)
+    bh_idx = tl.program_id(1)
+
+    if block_idx >= num_blocks:
+        return
+
+    vbs = tl.load(VBS_ptr + block_idx).to(tl.float32)
+
+    go_base = GradOut_ptr + bh_idx * stride_go_bh + block_idx * stride_go_blk
+    dim_offsets = tl.arange(0, HEAD_DIM)
+    grad_val = tl.load(go_base + dim_offsets).to(tl.float32) / vbs
+
+    gx_base = GradX_ptr + bh_idx * stride_gx_bh + block_idx * BLOCK_ELEMENTS * stride_gx_seq
+    grad_out_cast = grad_val.to(tl.bfloat16)
+    for i in range(BLOCK_ELEMENTS):
+        tl.store(gx_base + i * stride_gx_seq + dim_offsets, grad_out_cast)
+
+
+def _fused_block_mean_bwd(
+    grad_output: torch.Tensor,
     variable_block_sizes: torch.Tensor,
     block_elements: int,
 ) -> torch.Tensor:
-    """Compute block-wise mean with fp32 accumulation, fused in one kernel.
+    B, H, num_blocks, D = grad_output.shape
+    seq_len = num_blocks * block_elements
 
-    Args:
-        x: [B, H, seq_len, D] in bf16
-        variable_block_sizes: [num_blocks] number of valid tokens per block
-        block_elements: tokens per block (e.g. 64)
+    grad_x = torch.empty(B, H, seq_len, D, dtype=grad_output.dtype, device=grad_output.device)
 
-    Returns:
-        [B, H, num_blocks, D] in bf16
-    """
+    go_flat = grad_output.contiguous().view(B * H, num_blocks, D)
+    gx_flat = grad_x.view(B * H, seq_len, D)
+
+    grid = (num_blocks, B * H)
+
+    _fused_block_mean_bwd_kernel[grid](
+        go_flat, gx_flat, variable_block_sizes,
+        go_flat.stride(0), go_flat.stride(1),
+        gx_flat.stride(0), gx_flat.stride(1),
+        num_blocks,
+        BLOCK_ELEMENTS=block_elements,
+        HEAD_DIM=D,
+    )
+
+    return grad_x
+
+
+def _fused_block_mean_fwd(
+    x: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    block_elements: int,
+) -> torch.Tensor:
     B, H, seq_len, D = x.shape
     num_blocks = seq_len // block_elements
     assert seq_len % block_elements == 0
@@ -97,6 +151,42 @@ def fused_block_mean(
     return out
 
 
+class _FusedBlockMeanAutograd(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, variable_block_sizes, block_elements):
+        ctx.save_for_backward(variable_block_sizes)
+        ctx.block_elements = block_elements
+        return _fused_block_mean_fwd(x, variable_block_sizes, block_elements)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        variable_block_sizes, = ctx.saved_tensors
+        block_elements = ctx.block_elements
+        return _fused_block_mean_bwd(grad_output, variable_block_sizes, block_elements), None, None
+
+
+def fused_block_mean(
+    x: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    block_elements: int,
+) -> torch.Tensor:
+    """Compute block-wise mean with fp32 accumulation, fused in one kernel.
+
+    Forward: fused Triton kernel (bf16 read → fp32 accumulate → div → bf16 write).
+    Backward: broadcasts grad_output / vbs back to each token position.
+
+    Args:
+        x: [B, H, seq_len, D] in bf16
+        variable_block_sizes: [num_blocks] number of valid tokens per block
+        block_elements: tokens per block (e.g. 64)
+
+    Returns:
+        [B, H, num_blocks, D] in bf16
+    """
+    return _FusedBlockMeanAutograd.apply(x, variable_block_sizes, block_elements)
+
+
 @triton.jit
 def _fused_topk_mask_kernel(
     Scores_ptr,