liger-kernels: add device context manager

liger-kernels/torch-ext/liger_kernels/cross_entropy.py

@@ -11,6 +11,8 @@
 from .utils import is_hip
 from .utils import infer_device
 from .utils import is_npu_available
+from .utils import device_context
+
 
 if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
     try:
@@ -372,44 +374,45 @@ def cross_entropy_forward(
     if target.stride(-1) != 1:
         target = target.contiguous()
 
-    # Here we use a trick to store X_ptr gradient in X_ptr so we can save memory
-    liger_cross_entropy_kernel[(n_rows,)](
-        X_ptr=_input,
-        X_stride=_input.stride(-2),
-        Y_ptr=target,
-        Y_stride=target.stride(-1),  # always 1
-        weight_ptr=weight,  # dummy if None
-        loss_ptr=loss_1d,
-        z_loss_ptr=z_loss_1d,
-        loss_stride=loss_1d.stride(-1),  # always 1
-        token_accuracy_ptr=token_accuracy_1d,
-        token_accuracy_stride=token_accuracy_1d.stride(-1)
-        if return_token_accuracy
-        else 0,  # always 1 if accuracy is enabled
-        predicted_tokens_ptr=predicted_tokens_1d,
-        predicted_tokens_stride=predicted_tokens_1d.stride(-1)
-        if return_predicted_tokens
-        else 0,  # always 1 if predicted tokens is enabled
-        n_cols=V,
-        n_non_ignore=n_non_ignore,
-        sum_non_ignore_weight=sum_non_ignore_weight,
-        ignore_index=ignore_index,
-        weight_sum=weight_sum,
-        lse_square_scale=lse_square_scale,
-        label_smoothing=label_smoothing,
-        reduction=reduction,
-        softcap=softcap,
-        RETURN_Z_LOSS=return_z_loss,
-        RETURN_TOKEN_ACCURACY=return_token_accuracy,
-        RETURN_PREDICTED_TOKENS=return_predicted_tokens,
-        BLOCK_SIZE=BLOCK_SIZE,
-        HAS_WEIGHT=True if weight is not None else False,
-        HAS_SOFTCAPPING=True if softcap is not None else False,
-        HAS_GRADIENTS=_input.requires_grad,
-        # TODO: 32 seems to give the best performance
-        # Performance is quite sensitive to num_warps
-        num_warps=32 if not is_hip() else 16,
-    )
+    with device_context(_input.device):
+        # Here we use a trick to store X_ptr gradient in X_ptr so we can save memory
+        liger_cross_entropy_kernel[(n_rows,)](
+            X_ptr=_input,
+            X_stride=_input.stride(-2),
+            Y_ptr=target,
+            Y_stride=target.stride(-1),  # always 1
+            weight_ptr=weight,  # dummy if None
+            loss_ptr=loss_1d,
+            z_loss_ptr=z_loss_1d,
+            loss_stride=loss_1d.stride(-1),  # always 1
+            token_accuracy_ptr=token_accuracy_1d,
+            token_accuracy_stride=token_accuracy_1d.stride(-1)
+            if return_token_accuracy
+            else 0,  # always 1 if accuracy is enabled
+            predicted_tokens_ptr=predicted_tokens_1d,
+            predicted_tokens_stride=predicted_tokens_1d.stride(-1)
+            if return_predicted_tokens
+            else 0,  # always 1 if predicted tokens is enabled
+            n_cols=V,
+            n_non_ignore=n_non_ignore,
+            sum_non_ignore_weight=sum_non_ignore_weight,
+            ignore_index=ignore_index,
+            weight_sum=weight_sum,
+            lse_square_scale=lse_square_scale,
+            label_smoothing=label_smoothing,
+            reduction=reduction,
+            softcap=softcap,
+            RETURN_Z_LOSS=return_z_loss,
+            RETURN_TOKEN_ACCURACY=return_token_accuracy,
+            RETURN_PREDICTED_TOKENS=return_predicted_tokens,
+            BLOCK_SIZE=BLOCK_SIZE,
+            HAS_WEIGHT=True if weight is not None else False,
+            HAS_SOFTCAPPING=True if softcap is not None else False,
+            HAS_GRADIENTS=_input.requires_grad,
+            # TODO: 32 seems to give the best performance
+            # Performance is quite sensitive to num_warps
+            num_warps=32 if not is_hip() else 16,
+        )
 
     if reduction == "none":
         loss = loss_1d
@@ -437,18 +440,19 @@ def cross_entropy_backward(_input, grad_output):
     # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
     # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
     else:
-        BT, V = _input.shape
-        n_rows = BT
-        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-
-        element_mul_kernel[(n_rows,)](
-            _input,
-            _input.stride(-2),
-            grad_output,
-            V,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=32 if not is_hip() else 16,
-        )
+        with device_context(_input.device):
+            BT, V = _input.shape
+            n_rows = BT
+            BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+
+            element_mul_kernel[(n_rows,)](
+                _input,
+                _input.stride(-2),
+                grad_output,
+                V,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=32 if not is_hip() else 16,
+            )
 
     return _input
 

liger-kernels/torch-ext/liger_kernels/dyt.py

@@ -9,6 +9,8 @@
 from .utils import get_npu_core_count
 from .utils import infer_device
 from .utils import is_npu_available
+from .utils import device_context
+
 
 if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
     try:
@@ -107,16 +109,17 @@ def liger_dyt_fwd(x, alpha, gamma, beta):
 
     y = torch.empty_like(x)
 
-    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), M)
-    _dyt_fwd_kernel[grid](
-        x,
-        y,
-        alpha,
-        gamma,
-        beta,
-        HAVE_BETA,
-        N,
-    )
+    with device_context(x.device):
+        grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), M)
+        _dyt_fwd_kernel[grid](
+            x,
+            y,
+            alpha,
+            gamma,
+            beta,
+            HAVE_BETA,
+            N,
+        )
     return y.view(input_shape)
 
 
@@ -139,8 +142,9 @@ def liger_dyt_bwd(dy, x, alpha, gamma, beta):
     db = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device) if HAVE_BETA else None
     dx = torch.empty_like(dy)
 
-    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), NUM_SMS)
-    _dyt_bwd_kernel[grid](dy, dx, da, dg, db, x, alpha, gamma, HAVE_BETA, M, N)
+    with device_context(x.device):
+        grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), NUM_SMS)
+        _dyt_bwd_kernel[grid](dy, dx, da, dg, db, x, alpha, gamma, HAVE_BETA, M, N)
     if HAVE_BETA:
         db = db.sum(0).to(x.dtype)
     dg = dg.sum(0).to(gamma.dtype)

liger-kernels/torch-ext/liger_kernels/fused_linear_cross_entropy.py

@@ -7,6 +7,7 @@
 from .utils import element_mul_kernel
 from .utils import is_hip
 from .utils import infer_device
+from .utils import device_context
 
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
@@ -147,42 +148,43 @@ def fused_linear_cross_entropy_forward(
         logits_chunk = logits_chunk.contiguous()
         target_chunk = target_chunk.contiguous()
 
-        # Here we calculate the gradient of logits_chunk in place so we can save memory.
-        liger_cross_entropy_kernel[(n_rows,)](
-            X_ptr=logits_chunk,
-            X_stride=logits_chunk.stride(-2),
-            Y_ptr=target_chunk,
-            Y_stride=target_chunk.stride(-1),  # always 1
-            weight_ptr=ce_weight,
-            loss_ptr=loss_1d_slice,
-            z_loss_ptr=z_loss_1d_slice,
-            loss_stride=loss_1d_slice.stride(-1),  # always 1
-            token_accuracy_ptr=token_accuracy_1d_slice,
-            token_accuracy_stride=token_accuracy_1d_slice.stride(-1)
-            if return_token_accuracy
-            else 0,  # always 1 if accuracy is enabled
-            predicted_tokens_ptr=predicted_tokens_1d_slice,
-            predicted_tokens_stride=predicted_tokens_1d_slice.stride(-1)
-            if return_predicted_tokens
-            else 0,  # always 1 if predicted tokens is enabled
-            n_cols=V,
-            n_non_ignore=total_n_non_ignore,
-            sum_non_ignore_weight=total_sum_non_ignore_ce_weight,
-            weight_sum=ce_weight_sum,
-            ignore_index=ignore_index,
-            lse_square_scale=lse_square_scale,
-            label_smoothing=label_smoothing,
-            reduction=reduction,
-            softcap=softcap,
-            RETURN_Z_LOSS=return_z_loss,
-            RETURN_TOKEN_ACCURACY=return_token_accuracy,
-            RETURN_PREDICTED_TOKENS=return_predicted_tokens,
-            HAS_WEIGHT=True if ce_weight is not None else False,
-            HAS_SOFTCAPPING=True if softcap is not None else False,
-            HAS_GRADIENTS=input_requires_grad,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=32 if not is_hip() else 16,
-        )
+        with device_context(device):
+            # Here we calculate the gradient of logits_chunk in place so we can save memory.
+            liger_cross_entropy_kernel[(n_rows,)](
+                X_ptr=logits_chunk,
+                X_stride=logits_chunk.stride(-2),
+                Y_ptr=target_chunk,
+                Y_stride=target_chunk.stride(-1),  # always 1
+                weight_ptr=ce_weight,
+                loss_ptr=loss_1d_slice,
+                z_loss_ptr=z_loss_1d_slice,
+                loss_stride=loss_1d_slice.stride(-1),  # always 1
+                token_accuracy_ptr=token_accuracy_1d_slice,
+                token_accuracy_stride=token_accuracy_1d_slice.stride(-1)
+                if return_token_accuracy
+                else 0,  # always 1 if accuracy is enabled
+                predicted_tokens_ptr=predicted_tokens_1d_slice,
+                predicted_tokens_stride=predicted_tokens_1d_slice.stride(-1)
+                if return_predicted_tokens
+                else 0,  # always 1 if predicted tokens is enabled
+                n_cols=V,
+                n_non_ignore=total_n_non_ignore,
+                sum_non_ignore_weight=total_sum_non_ignore_ce_weight,
+                weight_sum=ce_weight_sum,
+                ignore_index=ignore_index,
+                lse_square_scale=lse_square_scale,
+                label_smoothing=label_smoothing,
+                reduction=reduction,
+                softcap=softcap,
+                RETURN_Z_LOSS=return_z_loss,
+                RETURN_TOKEN_ACCURACY=return_token_accuracy,
+                RETURN_PREDICTED_TOKENS=return_predicted_tokens,
+                HAS_WEIGHT=True if ce_weight is not None else False,
+                HAS_SOFTCAPPING=True if softcap is not None else False,
+                HAS_GRADIENTS=input_requires_grad,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=32 if not is_hip() else 16,
+            )
 
         # Apply token scaling if requested
         if use_token_scaling:
@@ -247,47 +249,48 @@ def fused_linear_cross_entropy_forward(
 def fused_linear_cross_entropy_backward(grad_output, grad_input, grad_weight, grad_bias):
     # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
     if not torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
-        # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
-        # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
-        BT, H = grad_input.shape
-        n_rows = BT
-        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
-
-        element_mul_kernel[(n_rows,)](
-            grad_input,
-            grad_input.stride(-2),
-            grad_output,
-            H,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=32 if not is_hip() else 16,
-        )
-
-        # handle grad_weight
-        if grad_weight is not None:
-            V, H = grad_weight.shape
-            n_rows = V
+        with device_context(grad_input.device):
+            # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
+            # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
+            BT, H = grad_input.shape
+            n_rows = BT
+            BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
 
             element_mul_kernel[(n_rows,)](
-                grad_weight,
-                grad_weight.stride(-2),
+                grad_input,
+                grad_input.stride(-2),
                 grad_output,
                 H,
                 BLOCK_SIZE=BLOCK_SIZE,
                 num_warps=32 if not is_hip() else 16,
             )
 
-        if grad_bias is not None:
-            V = grad_bias.shape[0]
-            n_rows = V
-
-            element_mul_kernel[(n_rows,)](
-                grad_bias,
-                grad_bias.stride(-1),
-                grad_output,
-                1,
-                BLOCK_SIZE=BLOCK_SIZE,
-                num_warps=32 if not is_hip() else 16,
-            )
+            # handle grad_weight
+            if grad_weight is not None:
+                V, H = grad_weight.shape
+                n_rows = V
+
+                element_mul_kernel[(n_rows,)](
+                    grad_weight,
+                    grad_weight.stride(-2),
+                    grad_output,
+                    H,
+                    BLOCK_SIZE=BLOCK_SIZE,
+                    num_warps=32 if not is_hip() else 16,
+                )
+
+            if grad_bias is not None:
+                V = grad_bias.shape[0]
+                n_rows = V
+
+                element_mul_kernel[(n_rows,)](
+                    grad_bias,
+                    grad_bias.stride(-1),
+                    grad_output,
+                    1,
+                    BLOCK_SIZE=BLOCK_SIZE,
+                    num_warps=32 if not is_hip() else 16,
+                )
     return grad_input, grad_weight, grad_bias