Guard Triton launches by tensor device

src/liger_kernel/ops/geglu.py

@@ -6,6 +6,7 @@
 
 from liger_kernel.ops.utils import calculate_settings
 from liger_kernel.ops.utils import compare_version
+from liger_kernel.ops.utils import device_context
 from liger_kernel.ops.utils import ensure_contiguous
 from liger_kernel.utils import is_npu_available
 
@@ -94,15 +95,16 @@ def geglu_forward(a, b):
 
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
 
-    _geglu_tanh_forward_kernel[(n_rows,)](
-        a,
-        b,
-        c,
-        c.stride(-2),
-        n_cols=n_cols,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=num_warps,
-    )
+    with device_context(a.device):
+        _geglu_tanh_forward_kernel[(n_rows,)](
+            a,
+            b,
+            c,
+            c.stride(-2),
+            n_cols=n_cols,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
     return a, b, c.view(*ori_shape)
 
 
@@ -114,15 +116,16 @@ def geglu_backward(a, b, dc):
 
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
 
-    _geglu_tanh_backward_kernel[(n_rows,)](
-        dc,
-        a,
-        b,
-        dc.stride(-2),
-        n_cols=n_cols,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=num_warps,
-    )
+    with device_context(a.device):
+        _geglu_tanh_backward_kernel[(n_rows,)](
+            dc,
+            a,
+            b,
+            dc.stride(-2),
+            n_cols=n_cols,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
 
     return a.view(*ori_shape), b.view(*ori_shape)
 

src/liger_kernel/ops/rms_norm.py

@@ -19,6 +19,7 @@
 
 from liger_kernel.ops.utils import calculate_settings
 from liger_kernel.ops.utils import compare_version
+from liger_kernel.ops.utils import device_context
 from liger_kernel.ops.utils import ensure_contiguous
 from liger_kernel.ops.utils import get_npu_core_count
 from liger_kernel.ops.utils import set_large_grf_mode
@@ -450,47 +451,48 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
     kernel_args = {}
     if X.device.type == "xpu":
         set_large_grf_mode(kernel_args)
-    if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode:
-        _rms_norm_forward_kernel[(n_rows,)](
-            Y,
-            Y.stride(0),
-            X,
-            X.stride(0),
-            W,
-            W.stride(0) if elementwise_affine else 0,
-            RSTD,
-            RSTD.stride(0),
-            n_cols,
-            eps,
-            offset,
-            casting_mode,
-            elementwise_affine=elementwise_affine,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-            **kernel_args,  # XPU-specific optimization
-        )
-    else:
-        BLOCK_ROW = 16
-        kernel_args["BLOCK_ROW"] = BLOCK_ROW
-        _block_rms_norm_forward_kernel[(triton.cdiv(n_rows, BLOCK_ROW),)](
-            Y,
-            Y.stride(0),
-            X,
-            X.stride(0),
-            W,
-            W.stride(0) if elementwise_affine else 0,
-            RSTD,
-            RSTD.stride(0),
-            n_rows,
-            n_cols,
-            eps,
-            offset,
-            casting_mode,
-            elementwise_affine=elementwise_affine,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-            **kernel_args,  # XPU-specific optimization
-        )
+    with device_context(X.device):
+        if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode:
+            _rms_norm_forward_kernel[(n_rows,)](
+                Y,
+                Y.stride(0),
+                X,
+                X.stride(0),
+                W,
+                W.stride(0) if elementwise_affine else 0,
+                RSTD,
+                RSTD.stride(0),
+                n_cols,
+                eps,
+                offset,
+                casting_mode,
+                elementwise_affine=elementwise_affine,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=num_warps,
+                **kernel_args,  # XPU-specific optimization
+            )
+        else:
+            BLOCK_ROW = 16
+            kernel_args["BLOCK_ROW"] = BLOCK_ROW
+            _block_rms_norm_forward_kernel[(triton.cdiv(n_rows, BLOCK_ROW),)](
+                Y,
+                Y.stride(0),
+                X,
+                X.stride(0),
+                W,
+                W.stride(0) if elementwise_affine else 0,
+                RSTD,
+                RSTD.stride(0),
+                n_rows,
+                n_cols,
+                eps,
+                offset,
+                casting_mode,
+                elementwise_affine=elementwise_affine,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=num_warps,
+                **kernel_args,  # XPU-specific optimization
+            )
     return Y.view(*shape), X, RSTD, BLOCK_SIZE, num_warps, casting_mode
 
 
@@ -531,57 +533,58 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     if X.device.type == "xpu":
         set_large_grf_mode(kernel_args)
 
-    if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode:
-        _rms_norm_backward_kernel[grid](
-            dY,
-            dY.stride(0),
-            dX,
-            dX.stride(0),
-            X,
-            X.stride(0),
-            torch_to_triton_dtype[X.dtype],
-            W,
-            W.stride(0) if elementwise_affine else 0,
-            RSTD,
-            RSTD.stride(0),
-            _dW,
-            _dW.stride(0) if elementwise_affine else 0,
-            n_rows,
-            n_cols,
-            offset,
-            rows_per_program,
-            casting_mode,
-            elementwise_affine=elementwise_affine,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-            **kernel_args,  # XPU-specific optimization
-        )
-    else:
-        BLOCK_ROW = 16
-        kernel_args["BLOCK_ROW"] = BLOCK_ROW
-        _block_rms_norm_backward_kernel[grid](
-            dY,
-            dY.stride(0),
-            dX,
-            dX.stride(0),
-            X,
-            X.stride(0),
-            torch_to_triton_dtype[X.dtype],
-            W,
-            W.stride(0) if elementwise_affine else 0,
-            RSTD,
-            RSTD.stride(0),
-            _dW,
-            _dW.stride(0) if elementwise_affine else 0,
-            n_rows,
-            n_cols,
-            offset,
-            casting_mode,
-            elementwise_affine=elementwise_affine,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-            **kernel_args,  # XPU-specific optimization
-        )
+    with device_context(X.device):
+        if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode:
+            _rms_norm_backward_kernel[grid](
+                dY,
+                dY.stride(0),
+                dX,
+                dX.stride(0),
+                X,
+                X.stride(0),
+                torch_to_triton_dtype[X.dtype],
+                W,
+                W.stride(0) if elementwise_affine else 0,
+                RSTD,
+                RSTD.stride(0),
+                _dW,
+                _dW.stride(0) if elementwise_affine else 0,
+                n_rows,
+                n_cols,
+                offset,
+                rows_per_program,
+                casting_mode,
+                elementwise_affine=elementwise_affine,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=num_warps,
+                **kernel_args,  # XPU-specific optimization
+            )
+        else:
+            BLOCK_ROW = 16
+            kernel_args["BLOCK_ROW"] = BLOCK_ROW
+            _block_rms_norm_backward_kernel[grid](
+                dY,
+                dY.stride(0),
+                dX,
+                dX.stride(0),
+                X,
+                X.stride(0),
+                torch_to_triton_dtype[X.dtype],
+                W,
+                W.stride(0) if elementwise_affine else 0,
+                RSTD,
+                RSTD.stride(0),
+                _dW,
+                _dW.stride(0) if elementwise_affine else 0,
+                n_rows,
+                n_cols,
+                offset,
+                casting_mode,
+                elementwise_affine=elementwise_affine,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=num_warps,
+                **kernel_args,  # XPU-specific optimization
+            )
     dX = dX.view(*shape)
 
     if elementwise_affine:

src/liger_kernel/ops/swiglu.py

@@ -3,6 +3,7 @@
 import triton.language as tl
 
 from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import device_context
 from liger_kernel.ops.utils import ensure_contiguous
 from liger_kernel.utils import infer_device_arch
 
@@ -150,30 +151,32 @@ def swiglu_forward(a, b, gate_multiplier: float = 1.0):
         # Blackwell (B200), wide rows: column-tiled 2D grid for higher SM occupancy.
         BLOCK_SIZE, num_warps = _swiglu_tile_settings(n_cols)
         grid = (n_rows, triton.cdiv(n_cols, BLOCK_SIZE))
-        _swiglu_forward_kernel_tiled[grid](
+        with device_context(a.device):
+            _swiglu_forward_kernel_tiled[grid](
+                a,
+                b,
+                c,
+                c.stride(-2),
+                float(gate_multiplier),
+                n_cols,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=num_warps,
+            )
+        return a, b, c.view(*ori_shape)
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    with device_context(a.device):
+        _swiglu_forward_kernel[(n_rows,)](
             a,
             b,
             c,
             c.stride(-2),
             float(gate_multiplier),
-            n_cols,
+            n_cols=n_cols,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
         )
-        return a, b, c.view(*ori_shape)
-
-    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
-
-    _swiglu_forward_kernel[(n_rows,)](
-        a,
-        b,
-        c,
-        c.stride(-2),
-        float(gate_multiplier),
-        n_cols=n_cols,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=num_warps,
-    )
     return a, b, c.view(*ori_shape)
 
 
@@ -187,30 +190,32 @@ def swiglu_backward(a, b, dc, gate_multiplier: float = 1.0):
         # Blackwell (B200), wide rows: column-tiled 2D grid for higher SM occupancy.
         BLOCK_SIZE, num_warps = _swiglu_tile_settings(n_cols)
         grid = (n_rows, triton.cdiv(n_cols, BLOCK_SIZE))
-        _swiglu_backward_kernel_tiled[grid](
+        with device_context(a.device):
+            _swiglu_backward_kernel_tiled[grid](
+                dc,
+                a,
+                b,
+                dc.stride(-2),
+                float(gate_multiplier),
+                n_cols,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=num_warps,
+            )
+        return a.view(*ori_shape), b.view(*ori_shape)
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    with device_context(a.device):
+        _swiglu_backward_kernel[(n_rows,)](
             dc,
             a,
             b,
             dc.stride(-2),
             float(gate_multiplier),
-            n_cols,
+            n_cols=n_cols,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
         )
-        return a.view(*ori_shape), b.view(*ori_shape)
-
-    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
-
-    _swiglu_backward_kernel[(n_rows,)](
-        dc,
-        a,
-        b,
-        dc.stride(-2),
-        float(gate_multiplier),
-        n_cols=n_cols,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=num_warps,
-    )
     return a.view(*ori_shape), b.view(*ori_shape)
 
 

src/liger_kernel/ops/utils.py

@@ -14,6 +14,7 @@
 import importlib
 import operator
 
+from contextlib import contextmanager
 from typing import Callable
 
 import torch
@@ -29,6 +30,17 @@ def is_hip() -> bool:
     return torch.version.hip is not None
 
 
+@contextmanager
+def device_context(device):
+    device = torch.device(device)
+    backend = getattr(torch, device.type, None)
+    if backend is not None and hasattr(backend, "device"):
+        with backend.device(device):
+            yield
+    else:
+        yield
+
+
 def ensure_contiguous(fn):
     @functools.wraps(fn)
     def wrapper(ctx, *args, **kwargs):