Add backward pass for softmax kernel

stack-info: PR: #744, branch: karthickai/stack/4

Author: karthickai

examples/softmax.py

@@ -11,6 +11,8 @@
 # %%
 from __future__ import annotations
 
+from typing import Any
+
 import torch
 
 import helion
@@ -89,6 +91,69 @@ def softmax_two_pass(x: torch.Tensor) -> torch.Tensor:
     return out
 
 
+@helion.kernel()
+def softmax_bwd(
+    grad_output: torch.Tensor, softmax_output: torch.Tensor
+) -> torch.Tensor:
+    """
+    Helion kernel implementing softmax backward pass.
+
+    dy/dx = softmax_output * (grad_output - sum(softmax_output * grad_output))
+
+    Args:
+        grad_output (torch.Tensor): Gradient from downstream layers of shape [m, n]
+        softmax_output (torch.Tensor): Output from forward softmax pass of shape [m, n]
+
+    Returns:
+        torch.Tensor: Gradient with respect to input of shape [m, n]
+    """
+    m, n = grad_output.size()
+    grad_input = torch.empty_like(grad_output)
+
+    block_size_m = hl.register_block_size(m)
+    block_size_n = hl.register_block_size(n)
+
+    for tile_m in hl.tile(m, block_size=block_size_m):
+        sum_per_row = hl.zeros([tile_m], dtype=torch.float32)
+        for tile_n in hl.tile(n, block_size=block_size_n):
+            sum_per_row += torch.sum(
+                softmax_output[tile_m, tile_n] * grad_output[tile_m, tile_n], dim=1
+            )
+        for tile_n in hl.tile(n, block_size=block_size_n):
+            grad_input[tile_m, tile_n] = softmax_output[tile_m, tile_n] * (
+                grad_output[tile_m, tile_n] - sum_per_row[:, None]
+            )
+
+    return grad_input
+
+
+class SoftmaxFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,  # noqa: ANN401
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        y = softmax_two_pass(x)
+        ctx.save_for_backward(y)
+        return y
+
+    @staticmethod
+    def backward(  # type: ignore[override]
+        ctx: Any,  # noqa: ANN401
+        grad_output: torch.Tensor,
+    ) -> tuple[torch.Tensor | None]:
+        (softmax_output,) = ctx.saved_tensors
+        grad_x = softmax_bwd(grad_output, softmax_output)
+        return (grad_x,)
+
+
+def softmax_fwd_bwd(
+    x: torch.Tensor,
+) -> torch.Tensor:
+    """Softmax with forward + backward support."""
+    return SoftmaxFunction.apply(x)  # type: ignore[no-any-return]
+
+
 # %%
 def check(m: int, n: int) -> None:
     """
@@ -105,6 +170,17 @@ def check(m: int, n: int) -> None:
     }
     run_example(kernels, lambda x: torch.nn.functional.softmax(x, dim=1), (x,))
 
+    print("\n\n=== Forward + Backward Pass Test ===")
+    x_grad = torch.randn([m, n], device="cuda", dtype=torch.float16, requires_grad=True)
+    run_example(
+        softmax_fwd_bwd,
+        torch.nn.functional.softmax,
+        (x_grad,),
+        rtol=1e-3,
+        atol=1e-3,
+        bwd=True,
+    )
+
 
 # %%
 def main() -> None:

test/test_examples.expected

@@ -3420,6 +3420,67 @@ def softmax(x: torch.Tensor, *, _launcher=_default_launcher):
     _launcher(_helion_softmax, (n,), x, out, out.size(0), out.size(1), x.size(0), x.size(1), out.stride(0), out.stride(1), x.stride(0), x.stride(1), _m, _RDIM_SIZE_1, 1, num_warps=4, num_stages=1)
     return out
 
+--- assertExpectedJournal(TestExamples.test_softmax_bwd)
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from helion.runtime import default_launcher as _default_launcher
+
+@triton.jit
+def _helion_softmax_bwd(softmax_output, grad_output, grad_input, grad_input_stride_0, grad_input_stride_1, grad_output_stride_0, grad_output_stride_1, softmax_output_stride_0, softmax_output_stride_1, m, n, _BLOCK_SIZE_0: tl.constexpr, _BLOCK_SIZE_1: tl.constexpr):
+    pid_0 = tl.program_id(0)
+    offset_0 = pid_0 * _BLOCK_SIZE_0
+    indices_0 = (offset_0 + tl.arange(0, _BLOCK_SIZE_0)).to(tl.int32)
+    mask_0 = indices_0 < m
+    sum_per_row = tl.full([_BLOCK_SIZE_0], 0.0, tl.float32)
+    for offset_2 in tl.range(0, n.to(tl.int32), _BLOCK_SIZE_1):
+        indices_2 = offset_2 + tl.arange(0, _BLOCK_SIZE_1).to(tl.int32)
+        mask_1 = indices_2 < n
+        sum_per_row_copy = sum_per_row
+        sum_per_row_copy_0 = sum_per_row_copy
+        load = tl.load(softmax_output + (indices_0[:, None] * softmax_output_stride_0 + indices_2[None, :] * softmax_output_stride_1), mask_0[:, None] & mask_1[None, :], other=0)
+        load_1 = tl.load(grad_output + (indices_0[:, None] * grad_output_stride_0 + indices_2[None, :] * grad_output_stride_1), mask_0[:, None] & mask_1[None, :], other=0)
+        v_0 = load * load_1
+        sum_1 = tl.cast(tl.sum(v_0, 1), tl.float16)
+        v_1 = tl.cast(sum_1, tl.float32)
+        sum_per_row = sum_per_row_copy_0 + v_1
+    for offset_2 in tl.range(0, n.to(tl.int32), _BLOCK_SIZE_1):
+        indices_2 = offset_2 + tl.arange(0, _BLOCK_SIZE_1).to(tl.int32)
+        mask_2 = indices_2 < n
+        sum_per_row_copy_1 = sum_per_row
+        sum_per_row_copy_1_0 = sum_per_row_copy_1
+        load_2 = tl.load(softmax_output + (indices_0[:, None] * softmax_output_stride_0 + indices_2[None, :] * softmax_output_stride_1), mask_0[:, None] & mask_2[None, :], other=0)
+        load_3 = tl.load(grad_output + (indices_0[:, None] * grad_output_stride_0 + indices_2[None, :] * grad_output_stride_1), mask_0[:, None] & mask_2[None, :], other=0)
+        subscript = sum_per_row_copy_1_0[:, None]
+        v_3 = tl.cast(load_3, tl.float32)
+        v_4 = v_3 - subscript
+        v_5 = tl.cast(load_2, tl.float32)
+        v_6 = v_5 * v_4
+        v_7 = tl.cast(v_6, tl.float16)
+        tl.store(grad_input + (indices_0[:, None] * grad_input_stride_0 + indices_2[None, :] * grad_input_stride_1), v_7, mask_0[:, None] & mask_2[None, :])
+
+def softmax_bwd(grad_output: torch.Tensor, softmax_output: torch.Tensor, *, _launcher=_default_launcher):
+    """
+    Helion kernel implementing softmax backward pass.
+
+    dy/dx = softmax_output * (grad_output - sum(softmax_output * grad_output))
+
+    Args:
+        grad_output (torch.Tensor): Gradient from downstream layers of shape [m, n]
+        softmax_output (torch.Tensor): Output from forward softmax pass of shape [m, n]
+
+    Returns:
+        torch.Tensor: Gradient with respect to input of shape [m, n]
+    """
+    m, n = grad_output.size()
+    grad_input = torch.empty_like(grad_output)
+    _BLOCK_SIZE_0 = 32
+    _BLOCK_SIZE_1 = 32
+    _launcher(_helion_softmax_bwd, (triton.cdiv(m, _BLOCK_SIZE_0),), softmax_output, grad_output, grad_input, grad_input.stride(0), grad_input.stride(1), grad_output.stride(0), grad_output.stride(1), softmax_output.stride(0), softmax_output.stride(1), m, n, _BLOCK_SIZE_0, _BLOCK_SIZE_1, num_warps=4, num_stages=3)
+    return grad_input
+
 --- assertExpectedJournal(TestExamples.test_softmax_decomposed)
 from __future__ import annotations
 

test/test_examples.py

@@ -881,6 +881,32 @@ def test_layernorm_bwd_dx(self):
             )
         )
 
+    def test_softmax_bwd(self):
+        m, n = 2048, 2048
+        x = torch.randn([m, n], device=DEVICE, dtype=torch.float16, requires_grad=True)
+        grad_out = torch.randn([m, n], device=DEVICE, dtype=torch.float16)
+
+        from examples.softmax import softmax_two_pass
+
+        config = helion.Config(block_size=[128, 128], num_warps=4, num_stages=3)
+        configured_kernel = helion.kernel(softmax_two_pass.fn, config=config)
+        y = configured_kernel(x)
+
+        x_torch = x.detach().clone().requires_grad_(True)
+        y_torch = torch.nn.functional.softmax(x_torch, dim=-1)
+        y_torch.backward(grad_out)
+
+        self.assertExpectedJournal(
+            check_example(
+                "softmax",
+                (grad_out, y),
+                x_torch.grad,
+                fn_name="softmax_bwd",
+                rtol=1e-3,
+                atol=1e-3,
+            )
+        )
+
     def test_layernorm_without_bias(self):
         x = -2.3 + 0.5 * torch.randn([32, 64], device=DEVICE, dtype=torch.float16)
         weight = torch.randn([64], device=DEVICE, dtype=torch.float16)