Unfuse forward/backward and use lock

Signed-off-by: Tcc0403 <[email protected]>

Author: Tcc0403

src/liger_kernel/ops/helion/fused_linear_cross_entropy.py

@@ -1,12 +1,35 @@
+import math
+
 import helion
 import helion.language as hl
 import torch
 
 # Best config for default llama3Config(hidden_size=4096, vocab_size=32000) with 4096 bs*seqlen input
 config = helion.Config(block_sizes=[32, 32, 256], indexing=['tensor_descriptor', 'tensor_descriptor', 'pointer', 'tensor_descriptor', 'pointer', 'pointer', 'tensor_descriptor', 'pointer', 'tensor_descriptor'], load_eviction_policies=['first', '', 'first', 'last', '', '', 'last', 'first'], num_stages=5, num_warps=4, pid_type='flat', range_flattens=[None, True, False], range_multi_buffers=[None, True, False], range_num_stages=[0, 0, 0], range_unroll_factors=[0, 1, 1], range_warp_specializes=[])
 
-@helion.kernel(config=config, ignore_warnings=[helion.exc.TensorOperationInWrapper])
-def fused_linear_cross_entropy_fwd_bwd(
+def helion_lock_acquire(lock_ptr, lock_index):
+    hl.inline_triton(
+        """
+        while tl.atomic_cas({0} + {1}, 0, 1, sem="acquire") == 1:
+            pass
+        """,
+        args=(lock_ptr, lock_index),
+        output_like=None,
+    )
+
+def helion_lock_release(lock_ptr, lock_index):
+    hl.inline_triton(
+        """
+        tl.atomic_xchg({0} + {1}, 0, sem="release")
+        """,
+        args=(lock_ptr, lock_index),
+        output_like=None,
+    )
+
+
+# @helion.kernel(config=config, ignore_warnings=[helion.exc.TensorOperationInWrapper])
+@helion.kernel(autotune_effort="none", ignore_warnings=[helion.exc.TensorOperationInWrapper])
+def fused_linear_cross_entropy_fwd(
     x: torch.Tensor,
     weight: torch.Tensor,
     target: torch.Tensor,
@@ -31,15 +54,15 @@ def fused_linear_cross_entropy_fwd_bwd(
     block_size_v = hl.register_block_size(V)
 
     nll = torch.zeros(BT, device=x.device, dtype=torch.float32)
-    grad_x = torch.zeros_like(x, dtype=torch.float32)
-    grad_w = torch.zeros_like(weight, dtype=torch.float32)
+    lse = torch.zeros(BT, device=x.device, dtype=torch.float32)
 
     # May be useful for splitting fwd and bwd
     # lse = torch.full((BT,), fill_value=-torch.inf, device=x.device, dtype=torch.float32)
     # neg_target_logits = torch.zeros(BT, device=x.device, dtype=torch.float32)
 
     n_non_ignore = (target != ignore_index).sum().unsqueeze(0)
 
+    # forward
     for tile_bt in hl.tile(BT, block_size=block_size_bt):
         m_i = hl.zeros([tile_bt], dtype=torch.float32) - float("inf")
         d_i = hl.zeros([tile_bt], dtype=torch.float32)
@@ -75,50 +98,86 @@ def fused_linear_cross_entropy_fwd_bwd(
             nll_tile /= n_non_ignore_value
 
         nll[tile_bt] = nll_tile
-        # gradients computation
-        for tile_v in hl.tile(V, block_size=block_size_v):
-            # Restore logits
-            acc = hl.zeros([tile_bt, tile_v], dtype=torch.float32)
-            for tile_h in hl.tile(H, block_size=block_size_h):
-                x_tile = x[tile_bt, tile_h]
-                weight_tile = weight[tile_v, tile_h]
-                acc = hl.dot(x_tile, weight_tile.T, acc=acc, out_dtype=torch.float32)
-
-            # softmax(x_i) = exp(x_i) / sum(exp(x_i))
-            #              = exp(x_i) / log(exp(sum(x_i)))
-            #              = exp(x_i) / lse = exp(x_i - lse)
-            grad_logits_tile = torch.exp(acc - lse_tile[:, None])
-            offset = tile_v.index.unsqueeze(0)  # [1, tile_v]
-            mask = target_indices == offset  # [tile_bt, tile_v]
-            grad_logits_tile = grad_logits_tile - mask.float()
-            # handle out of bound values in grad_logits_tile
-            grad_logits_tile = grad_logits_tile * ((tile_bt.index < BT)[:, None] & (tile_v.index < V)[None, :])
-
-            if reduction == "mean":
-                grad_logits_tile /= n_non_ignore_value
-
-            for tile_h in hl.tile(H, block_size=block_size_h):
-                # grad_x = grad_logits @ weight
-                rhs_tile = weight[tile_v, tile_h]
-                partial_grad_x = hl.dot(grad_logits_tile, rhs_tile, out_dtype=torch.float32)
-                hl.atomic_add(grad_x, [tile_bt, tile_h], partial_grad_x)
-                # grad_w = grad_logits.T[tile_v, tile_bt] @ x[tile_bt, tile_h]
-                rhs_tile = x[tile_bt, tile_h]
-                partial_grad_w = hl.dot(grad_logits_tile.T, rhs_tile, out_dtype=torch.float32)
-                hl.atomic_add(grad_w, [tile_v, tile_h], partial_grad_w)
-
+        lse[tile_bt] = lse_tile
+    
     if reduction != "none":
         loss = nll.sum() 
     else:
         loss = nll
+    
+    return loss, lse
+
+@helion.kernel(autotune_effort="none", ignore_warnings=[helion.exc.TensorOperationInWrapper])
+def fused_linear_cross_entropy_bwd(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    target: torch.Tensor,
+    lse: torch.Tensor,
+    ignore_index: int = -100,
+    reduction: str = "mean",
+):
+    BT, H = x.size()
+    V = weight.size(0)
+    block_size_bt = hl.register_block_size(BT)
+    block_size_h = hl.register_block_size(H)
+    block_size_v = hl.register_block_size(V)
+    grad_x = torch.zeros_like(x, dtype=torch.float32)
+    grad_w = torch.zeros_like(weight, dtype=torch.float32)
+    n_non_ignore = (target != ignore_index).sum().unsqueeze(0)
+
+    num_block_bt = (BT + block_size_bt - 1)//block_size_bt
+    num_block_h = (H + block_size_h - 1)//block_size_h
+    num_block_v = (V + block_size_v - 1)//block_size_v
+    grad_x_lock = torch.zeros((num_block_bt, num_block_h), dtype=torch.int32, device=x.device)
+    grad_w_lock = torch.zeros((num_block_v, num_block_h), dtype=torch.int32, device=x.device)
+    # backward
+    for tile_bt, tile_v in hl.tile([BT, V], block_size=(block_size_bt, block_size_v)):
+        # Restore logits
+        acc2 = hl.zeros([tile_bt, tile_v], dtype=torch.float32)
+        for tile_h in hl.tile(H, block_size=block_size_h):
+            x_tile = x[tile_bt, tile_h]
+            weight_tile = weight[tile_v, tile_h]
+            acc2 = hl.dot(x_tile, weight_tile.T, acc=acc2, out_dtype=torch.float32)
+
+        # softmax(x_i) = exp(x_i) / sum(exp(x_i))
+        #              = exp(x_i) / log(exp(sum(x_i)))
+        #              = exp(x_i) / lse = exp(x_i - lse)
+        lse_tile = lse[tile_bt]
+        target_indices = target[tile_bt].unsqueeze(1)  # [tile_bt, 1]
+        if reduction == "mean":
+            n_non_ignore_value = hl.load(n_non_ignore, [0])
+
+        grad_logits_tile = torch.exp(acc2 - lse_tile[:, None])
+        offset = tile_v.index.unsqueeze(0)  # [1, tile_v]
+        mask = target_indices == offset  # [tile_bt, tile_v]
+        grad_logits_tile = grad_logits_tile - mask.float()
+        # handle out of bound values in grad_logits_tile
+        grad_logits_tile = grad_logits_tile * ((tile_bt.index < BT)[:, None] & (tile_v.index < V)[None, :])
+
+        if reduction == "mean":
+            grad_logits_tile /= n_non_ignore_value
+
+        for tile_h in hl.tile(H, block_size=block_size_h):
+            # grad_x = grad_logits @ weight
+            rhs_tile = weight[tile_v, tile_h]
+            partial_grad_x = hl.dot(grad_logits_tile, rhs_tile, out_dtype=torch.float32)
+            helion_lock_acquire(grad_x_lock, tile_bt.id * num_block_h + tile_h.id)
+            grad_x[tile_bt, tile_h] += partial_grad_x
+            helion_lock_release(grad_x_lock, tile_bt.id * num_block_h + tile_h.id)
+            # hl.atomic_add(grad_x, [tile_bt, tile_h], partial_grad_x)
+        
+        # for tile_h in hl.tile(H, block_size=block_size_h):
+            # grad_w = grad_logits.T[tile_v, tile_bt] @ x[tile_bt, tile_h]
+            rhs_tile = x[tile_bt, tile_h]
+            partial_grad_w = hl.dot(grad_logits_tile.T, rhs_tile, out_dtype=torch.float32)
+            helion_lock_acquire(grad_w_lock, tile_v.id * num_block_h + tile_h.id)
+            grad_w[tile_v, tile_h] += partial_grad_w
+            helion_lock_release(grad_w_lock, tile_v.id * num_block_h + tile_h.id)
+            # hl.atomic_add(grad_w, [tile_v, tile_h], partial_grad_w)
+
 
-    # return format is not determined yet
-    return loss, dict(
-        {
-            "grad_x": grad_x,
-            "grad_w": grad_w,
-        }
-    )
+
+    return grad_x, grad_w
 
 
 class LigerFusedLinearCrossEntropyHelionFunction(torch.autograd.Function):
@@ -131,20 +190,30 @@ def forward(
         ignore_index=-100,
         reduction="mean",
     ):
-        loss, aux_output = fused_linear_cross_entropy_fwd_bwd(
+        loss, lse = fused_linear_cross_entropy_fwd(
             _input,
             weight,
             target,
             ignore_index,
             reduction,
         )
-        ctx.save_for_backward(aux_output["grad_x"], aux_output["grad_w"])
+        ctx.ignore_index = ignore_index
+        ctx.reduction = reduction
+        ctx.save_for_backward(_input, lse)
         return loss
 
     @staticmethod
     def backward(ctx, grad_output):
         assert grad_output.ndim == 0, "token_scaling is not supported. grad_output must be a scalar"
-        grad_input, grad_weight = ctx.saved_tensors
+        _input, lse = ctx.saved_tensors
+        grad_input, grad_weight = fused_linear_cross_entropy_bwd(
+            _input,
+            weight,
+            target,
+            lse,
+            ctx.ignore_index,
+            ctx.reduction,
+        )
         return grad_input * grad_output, grad_weight * grad_output, None, None, None
 
 
@@ -250,6 +319,7 @@ def forward(self, x, target):
     from helion._testing import run_example
     from functools import partial
 
+
     def fwd_bwd_fn(input, target, fn):
         loss = fn(input, target)
         loss.backward()