add triton based Softmax and unit tests

tests/modules/layers/test_triton_softmax.py

@@ -0,0 +1,149 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+import pytest
+import torch
+
+from tests.test_utils import assert_expected, gpu_test, set_rng_seed
+from torchmultimodal.triton.layers.softmax import fused_softmax
+
+
+@pytest.fixture(autouse=True)
+def set_seed():
+ set_rng_seed(2020)
+
+
+@gpu_test()
+class TestForwardSoftMax:
+ def test_forward_2d_float32(
+ self,
+ ):
+ # float32
+ seq_len = 768
+
+ sample_constant_float32 = torch.ones(
+ (seq_len, seq_len), dtype=torch.float32, device="cuda"
+ )
+ sample_random_float32 = torch.randn_like(sample_constant_float32)
+
+ expected_out_constant32 = torch.softmax(sample_constant_float32, dim=1)
+ expected_out_random32 = torch.softmax(sample_random_float32, dim=1)
+
+ triton_out_c32 = fused_softmax(sample_constant_float32)
+ triton_out_random32 = fused_softmax(sample_random_float32)
+
+ assert_expected(triton_out_c32, expected_out_constant32)
+ assert_expected(triton_out_random32, expected_out_random32)
+
+ def test_forward_2d_bfloat16(
+ self,
+ ):
+ # bfloat16
+ seq_len = 2048
+ sample_constant_bf16 = torch.ones(
+ (seq_len, seq_len), dtype=torch.bfloat16, device="cuda"
+ )
+ sample_random_bf16 = torch.randn_like(sample_constant_bf16)
+
+ expected_out_c_bf16 = torch.softmax(sample_constant_bf16, dim=1)
+ expected_out_rand_bf16 = torch.softmax(sample_random_bf16, dim=1)
+
+ triton_out_c_bf16 = fused_softmax(sample_constant_bf16)
+ triton_out_rand_bf16 = fused_softmax(sample_random_bf16)
+
+ assert_expected(triton_out_c_bf16, expected_out_c_bf16)
+ assert_expected(triton_out_rand_bf16, expected_out_rand_bf16)
+
+ def test_forward_3d_bfloat16(
+ self,
+ ):
+ # bfloat16
+ seq_len = 2048
+ batch = 12
+
+ sample_constant_bf16 = torch.ones(
+ (batch, seq_len, seq_len), dtype=torch.bfloat16, device="cuda"
+ )
+ sample_random_bf16 = torch.randn_like(sample_constant_bf16)
+
+ expected_out_c_bf16 = torch.softmax(sample_constant_bf16, dim=1)
+ expected_out_rand_bf16 = torch.softmax(sample_random_bf16, dim=1)
+
+ triton_out_c_bf16 = fused_softmax(sample_constant_bf16)
+ triton_out_rand_bf16 = fused_softmax(sample_random_bf16)
+
+ assert_expected(triton_out_c_bf16, expected_out_c_bf16, rtol=1e-4, atol=1e-2)
+ assert_expected(
+ triton_out_rand_bf16, expected_out_rand_bf16, rtol=1e-4, atol=1e-2
+ )
+
+
+@gpu_test()
+class TestBackwardSoftMax:
+ def test_backward_2d(
+ self,
+ ):
+ seq_len = 1024
+
+ sample_constant_float32 = torch.ones(
+ (seq_len, seq_len), dtype=torch.float32, device="cuda", requires_grad=True
+ )
+ sample_random_float32 = torch.randn_like(
+ sample_constant_float32, requires_grad=True
+ )
+
+ expected_fwd_constant32 = torch.softmax(sample_constant_float32, dim=1)
+ expected_fwd_random32 = torch.softmax(sample_random_float32, dim=1)
+
+ triton_fwd_c32 = fused_softmax(sample_constant_float32)
+ triton_fwd_random32 = fused_softmax(sample_random_float32)
+
+ dout = torch.randn_like(sample_constant_float32)
+
+ expected_bwd_c32 = expected_fwd_constant32.backward(dout)
+ expected_bwd_r32 = expected_fwd_random32.backward(dout)
+
+ triton_bwd_c32 = triton_fwd_c32.backward(dout)
+ triton_bwd_r32 = triton_fwd_random32.backward(dout)
+
+ assert_expected(triton_bwd_c32, expected_bwd_c32)
+ assert_expected(triton_bwd_r32, expected_bwd_r32)
+
+ def test_bwd_3d(
+ self,
+ ):
+ seq_len = 2048
+ batch = 4
+
+ sample_constant_float32 = torch.ones(
+ (batch, seq_len, seq_len),
+ dtype=torch.float32,
+ device="cuda",
+ requires_grad=True,
+ )
+ sample_random_float32 = torch.randn_like(
+ sample_constant_float32, requires_grad=True
+ )
+
+ expected_fwd_constant32 = torch.softmax(sample_constant_float32, dim=1)
+ expected_fwd_random32 = torch.softmax(sample_random_float32, dim=1)
+
+ triton_fwd_c32 = fused_softmax(sample_constant_float32)
+ triton_fwd_random32 = fused_softmax(sample_random_float32)
+
+ dout = torch.randn_like(sample_constant_float32)
+
+ expected_bwd_c32 = expected_fwd_constant32.backward(dout)
+ expected_bwd_r32 = expected_fwd_random32.backward(dout)
+
+ triton_bwd_c32 = triton_fwd_c32.backward(dout)
+ triton_bwd_r32 = triton_fwd_random32.backward(dout)
+
+ assert_expected(triton_bwd_c32, expected_bwd_c32)
+ assert_expected(triton_bwd_r32, expected_bwd_r32)

torchmultimodal/triton/layers/softmax.py

@@ -0,0 +1,155 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# ---- Fused Softmax written in Triton ------
+# Extra Credits:
+# Triton Softmax Tutorial
+# LucidRains Triton_Transformers
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from torch import autograd
+
+
+def _get_num_warps(block_size: int) -> int:
+ num_warps = 4
+ if block_size > 2047:
+ num_warps = 8
+ if block_size > 4095:
+ num_warps = 16
+ return num_warps
+
+
+@triton.jit
+def _softmax_kernel_fwd(
+ output_ptr,
+ output_row_stride,
+ input_ptr,
+ input_row_stride,
+ n_cols,
+ block_size: tl.constexpr,
+) -> None:
+
+ # setup input location
+ row_index = tl.program_id(0)
+ input_row_ptr = input_ptr + (row_index * input_row_stride)
+ col_offsets = tl.arange(0, block_size)
+ input_ptrs = input_row_ptr + col_offsets
+ rw_mask = col_offsets < n_cols
+ row = tl.load(input_ptrs, mask=rw_mask, other=float("-inf"))
+
+ # safe softmax proper
+ safe_row = row - tl.max(row, axis=0)
+ numerator = tl.exp(safe_row)
+ denom = tl.sum(numerator, axis=0)
+ sm_out = numerator / denom
+
+ # write results to HBM
+ out_row_ptr = output_ptr + (row_index * output_row_stride)
+ out_row_ptrs = out_row_ptr + col_offsets
+ tl.store(out_row_ptrs, sm_out, mask=rw_mask)
+
+
+@triton.jit
+def _softmax_kernel_bwd(
+ output_ptr,
+ stride_output_row,
+ grad_ptr,
+ stride_grad_row,
+ input_ptr,
+ stride_input_row,
+ n_cols,
+ block_size: tl.constexpr,
+) -> None:
+ # setup input locations - need both grad and input access
+ row_index = tl.program_id(0)
+
+ input_row_ptr = input_ptr + (row_index * stride_input_row)
+ grad_row_ptr = grad_ptr + (row_index * stride_grad_row)
+
+ col_offsets = tl.arange(0, block_size)
+ rw_mask = col_offsets < n_cols
+
+ input_row_ptrs = input_row_ptr + col_offsets
+ grad_row_ptrs = grad_row_ptr + col_offsets
+
+ probs_row = tl.load(input_row_ptrs, mask=rw_mask, other=0)
+ grads_row = tl.load(grad_row_ptrs, mask=rw_mask, other=0)
+
+ # compute derivatives
+ dx = probs_row * grads_row
+ dsm_out = dx - probs_row * (tl.sum(dx, axis=0))
+
+ # write to HBM
+ output_row_ptr = output_ptr + (row_index * stride_output_row)
+ output_ptrs = output_row_ptr + col_offsets
+ tl.store(output_ptrs, dsm_out, mask=rw_mask)
+
+
+class TritonSoftmax(autograd.Function):
+ @staticmethod
+ def forward(ctx, x: torch.Tensor) -> torch.Tensor:
+ orig_shape = x.shape
+ x = x.view(-1, orig_shape[-1])
+ nrows, ncols = x.shape
+
+ block_size = triton.next_power_of_2(ncols)
+ num_warps = _get_num_warps(block_size)
+
+ res = torch.empty_like(x)
+ grid = (nrows,)
+
+ _softmax_kernel_fwd[grid](
+ res,
+ res.stride(0),
+ x,
+ x.stride(0),
+ ncols,
+ block_size=block_size,
+ num_warps=num_warps,
+ )
+
+ if x.requires_grad:
+ ctx.save_for_backward(res)
+
+ return res.view(*orig_shape)
+
+ @staticmethod
+ def backward(ctx, grad_probs: torch.Tensor) -> Tuple[torch.Tensor, None]:
+ orig_shape = grad_probs.shape
+ (probs,) = ctx.saved_tensors
+
+ grad_probs = grad_probs.view(-1, orig_shape[-1])
+ nrows, ncols = grad_probs.shape
+
+ block_size = triton.next_power_of_2(ncols)
+ num_warps = _get_num_warps(block_size)
+
+ dx = torch.empty_like(probs)
+ grid = (nrows,)
+
+ _softmax_kernel_bwd[grid](
+ dx,
+ dx.stride(0),
+ probs,
+ probs.stride(0),
+ grad_probs,
+ grad_probs.stride(0),
+ ncols,
+ block_size=block_size,
+ num_warps=num_warps,
+ )
+ return dx.view(*orig_shape), None
+
+
+fused_softmax = TritonSoftmax.apply