From cfce99789636474b582e5033197e6e9603d38bb8 Mon Sep 17 00:00:00 2001
From: PerryLink <255665900+PerryLink@users.noreply.github.com>
Date: Thu, 17 Sep 2026 21:04:45 +0800
Subject: [PATCH 1/2] [BugFix][Multi-GPU] Allocate and launch MoE outputs on
 the input device

---
 tile_kernels/config.py                       | 22 ++++--
 tile_kernels/moe/aux_fi_kernel.py            | 17 +++--
 tile_kernels/moe/expand_to_fused_kernel.py   | 76 +++++++++++---------
 tile_kernels/moe/get_fused_mapping_kernel.py | 53 +++++++-------
 tile_kernels/moe/group_count_kernel.py       | 16 +++--
 tile_kernels/moe/normalize_weight_kernel.py  | 21 ++++--
 tile_kernels/moe/reduce_fused_kernel.py      | 42 +++++++----
 tile_kernels/utils.py                        | 17 +++++
 8 files changed, 167 insertions(+), 97 deletions(-)

diff --git a/tile_kernels/config.py b/tile_kernels/config.py
index 036dc04..28e6d7e 100644
--- a/tile_kernels/config.py
+++ b/tile_kernels/config.py
@@ -1,12 +1,18 @@
 import functools
+from typing import Optional
+
 import torch
 
 _num_sms = 0
 
 
 @functools.lru_cache(maxsize=None)
-def get_device_num_sms() -> int:
-    prop = torch.cuda.get_device_properties(torch.cuda.current_device())
+def get_device_num_sms(device_index: Optional[int] = None) -> int:
+    # The device index is part of the cache key: properties are looked up on the
+    # requested device and a second device never reuses the first device's value.
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    prop = torch.cuda.get_device_properties(device_index)
     return prop.multi_processor_count
 
 
@@ -16,14 +22,18 @@ def set_num_sms(num_sms: int) -> None:
     _num_sms = num_sms
 
 
-def get_num_sms() -> int:
+def get_num_sms(device_index: Optional[int] = None) -> int:
+    # `_num_sms` set through `set_num_sms` is a process-global override; it wins
+    # over the per-device value.
     global _num_sms
     if _num_sms == 0:
-        return get_device_num_sms()
+        return get_device_num_sms(device_index)
     return _num_sms
 
 
 @functools.lru_cache(maxsize=None)
-def get_max_smem_per_sm() -> int:
-    prop = torch.cuda.get_device_properties(torch.cuda.current_device())
+def get_max_smem_per_sm(device_index: Optional[int] = None) -> int:
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    prop = torch.cuda.get_device_properties(device_index)
     return prop.shared_memory_per_multiprocessor
diff --git a/tile_kernels/moe/aux_fi_kernel.py b/tile_kernels/moe/aux_fi_kernel.py
index 6607ac8..769e851 100644
--- a/tile_kernels/moe/aux_fi_kernel.py
+++ b/tile_kernels/moe/aux_fi_kernel.py
@@ -2,7 +2,7 @@
 import torch
 import tilelang
 from tilelang import language as T
-from tile_kernels.utils import align
+from tile_kernels.utils import align, get_device_guard
 from tile_kernels.config import get_num_sms
 
 
@@ -61,12 +61,17 @@ def aux_fi(topk_idx: torch.Tensor, num_experts: int, num_aux_topk: int) -> torch
     assert topk_idx.dim() == 2 and topk_idx.is_contiguous()
 
     num_topk = topk_idx.shape[1]
-    kernel = get_aux_fi_kernel(num_topk, num_experts, get_num_sms())
 
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
+    # Allocate and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    device = topk_idx.device
+    with get_device_guard(device):
+        kernel = get_aux_fi_kernel(num_topk, num_experts, get_num_sms(device.index))
 
-    out = torch.zeros(num_experts, dtype=torch.float32, device='cuda')
-    kernel(topk_idx, out, num_aux_topk)
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
+
+        out = torch.zeros(num_experts, dtype=torch.float32, device=device)
+        kernel(topk_idx, out, num_aux_topk)
 
     return out
diff --git a/tile_kernels/moe/expand_to_fused_kernel.py b/tile_kernels/moe/expand_to_fused_kernel.py
index 92a18c4..402707a 100644
--- a/tile_kernels/moe/expand_to_fused_kernel.py
+++ b/tile_kernels/moe/expand_to_fused_kernel.py
@@ -4,7 +4,7 @@
 from tilelang import language as T
 from typing import Optional
 
-from tile_kernels.utils import align, ceil_div
+from tile_kernels.utils import align, ceil_div, get_device_guard
 from tile_kernels.quant.common import QuantTensor
 
 
@@ -112,20 +112,24 @@ def expand_to_fused(x: torch.Tensor, token_topk_to_pos: torch.Tensor, pos_to_exp
     num_expanded_tokens = pos_to_expert.shape[0]
     assert num_tokens == num_tokens_
 
-    kernel = get_expand_to_fused_kernel(
-        hidden,
-        num_topk,
-        None, None, None,
-        T.dtype(x.dtype),
-        T.dtype(x.dtype),
-    )
+    device = x.device
+    assert token_topk_to_pos.device == device and pos_to_expert.device == device
 
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
+    with get_device_guard(device):
+        kernel = get_expand_to_fused_kernel(
+            hidden,
+            num_topk,
+            None, None, None,
+            T.dtype(x.dtype),
+            T.dtype(x.dtype),
+        )
 
-    out = torch.empty((num_expanded_tokens, hidden), dtype=x.dtype, device='cuda')
-    if num_tokens > 0:
-        kernel(x, None, out, None, token_topk_to_pos, pos_to_expert)
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
+
+        out = torch.empty((num_expanded_tokens, hidden), dtype=x.dtype, device=device)
+        if num_tokens > 0:
+            kernel(x, None, out, None, token_topk_to_pos, pos_to_expert)
 
     return out
 
@@ -176,25 +180,33 @@ def expand_to_fused_with_sf(
 
     assert hidden_sf == x_sf.shape[1]
 
-    kernel = get_expand_to_fused_kernel(
-        hidden,
-        num_topk,
-        num_per_channels,
-        use_tma_aligned_col_major_sf,
-        use_packed_ue8m0,
-        T.dtype(x.dtype),
-        T.dtype(x_sf.dtype),
-    )
-
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
-
-    out = torch.empty((num_expanded_tokens, hidden), dtype=x.dtype, device='cuda')
-    out_sf = torch.empty((hidden_sf, num_expanded_sf_tokens) if use_tma_aligned_col_major_sf else (num_expanded_tokens, hidden_sf), dtype=x_sf.dtype, device='cuda')
-    out_sf = out_sf[:, :num_expanded_tokens] if use_tma_aligned_col_major_sf else out_sf
-
-    if num_tokens > 0:
-        kernel(x, x_sf, out, out_sf, token_topk_to_pos, pos_to_expert)
+    device = x.device
+    assert x_sf.device == device and token_topk_to_pos.device == device and pos_to_expert.device == device
+
+    with get_device_guard(device):
+        kernel = get_expand_to_fused_kernel(
+            hidden,
+            num_topk,
+            num_per_channels,
+            use_tma_aligned_col_major_sf,
+            use_packed_ue8m0,
+            T.dtype(x.dtype),
+            T.dtype(x_sf.dtype),
+        )
+
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
+
+        out = torch.empty((num_expanded_tokens, hidden), dtype=x.dtype, device=device)
+        out_sf = torch.empty(
+            (hidden_sf, num_expanded_sf_tokens) if use_tma_aligned_col_major_sf else (num_expanded_tokens, hidden_sf),
+            dtype=x_sf.dtype,
+            device=device,
+        )
+        out_sf = out_sf[:, :num_expanded_tokens] if use_tma_aligned_col_major_sf else out_sf
+
+        if num_tokens > 0:
+            kernel(x, x_sf, out, out_sf, token_topk_to_pos, pos_to_expert)
     out_sf = out_sf.T if use_tma_aligned_col_major_sf else out_sf
 
     return out, out_sf
diff --git a/tile_kernels/moe/get_fused_mapping_kernel.py b/tile_kernels/moe/get_fused_mapping_kernel.py
index 998189e..31dc0f4 100644
--- a/tile_kernels/moe/get_fused_mapping_kernel.py
+++ b/tile_kernels/moe/get_fused_mapping_kernel.py
@@ -4,7 +4,7 @@
 from tilelang import language as T
 from tile_kernels.config import get_num_sms
 
-from tile_kernels.utils import align
+from tile_kernels.utils import align, get_device_guard
 
 
 @T.macro
@@ -203,30 +203,33 @@ def get_fused_mapping(
         should_sync = True
         num_expanded_tokens = (num_tokens * num_topk + (alignment - 1) * num_experts) // alignment * alignment
 
-    # Allocate output
-    num_sms = get_num_sms()
-    pos_to_expert = torch.empty((num_expanded_tokens, ), dtype=torch.int32, device='cuda')
-    pos_to_token = torch.empty((num_expanded_tokens, ), dtype=torch.int32, device='cuda')
-    pos_to_token_topk = torch.empty((num_expanded_tokens, ), dtype=torch.int32, device='cuda')
-    token_topk_to_pos = torch.empty((num_tokens, num_topk), dtype=torch.int32, device='cuda')
-    expert_start = torch.empty((num_experts, ), dtype=torch.int32, device='cuda')
-    expert_end = torch.empty((num_experts, ), dtype=torch.int32, device='cuda')
-    num_tokens_per_expert = torch.empty((num_experts, ), dtype=torch.int32, device='cuda')
-    num_experts_per_sm = torch.empty((num_sms, num_experts), dtype=torch.int32, device='cuda')
-
-    # Get kernel and launch
-    mapping_kernel = get_get_fused_mapping_kernel(num_experts, num_topk, alignment, num_sms)
-    mapping_kernel(
-        topk_idx,
-        pos_to_expert,
-        pos_to_token,
-        pos_to_token_topk,
-        token_topk_to_pos,
-        expert_start,
-        expert_end,
-        num_tokens_per_expert,
-        num_experts_per_sm,
-    )
+    # Allocate output and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    device = topk_idx.device
+    with get_device_guard(device):
+        num_sms = get_num_sms(device.index)
+        pos_to_expert = torch.empty((num_expanded_tokens, ), dtype=torch.int32, device=device)
+        pos_to_token = torch.empty((num_expanded_tokens, ), dtype=torch.int32, device=device)
+        pos_to_token_topk = torch.empty((num_expanded_tokens, ), dtype=torch.int32, device=device)
+        token_topk_to_pos = torch.empty((num_tokens, num_topk), dtype=torch.int32, device=device)
+        expert_start = torch.empty((num_experts, ), dtype=torch.int32, device=device)
+        expert_end = torch.empty((num_experts, ), dtype=torch.int32, device=device)
+        num_tokens_per_expert = torch.empty((num_experts, ), dtype=torch.int32, device=device)
+        num_experts_per_sm = torch.empty((num_sms, num_experts), dtype=torch.int32, device=device)
+
+        # Get kernel and launch
+        mapping_kernel = get_get_fused_mapping_kernel(num_experts, num_topk, alignment, num_sms)
+        mapping_kernel(
+            topk_idx,
+            pos_to_expert,
+            pos_to_token,
+            pos_to_token_topk,
+            token_topk_to_pos,
+            expert_start,
+            expert_end,
+            num_tokens_per_expert,
+            num_experts_per_sm,
+        )
     if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
         print(mapping_kernel.get_kernel_source())
 
diff --git a/tile_kernels/moe/group_count_kernel.py b/tile_kernels/moe/group_count_kernel.py
index bcf91bf..4dd4ceb 100644
--- a/tile_kernels/moe/group_count_kernel.py
+++ b/tile_kernels/moe/group_count_kernel.py
@@ -4,7 +4,7 @@
 from tilelang import language as T
 
 from tile_kernels.config import get_num_sms
-from tile_kernels.utils import align
+from tile_kernels.utils import align, get_device_guard
 
 
 @tilelang.jit(
@@ -58,12 +58,16 @@ def group_count(group_idx: torch.Tensor, num_groups: int) -> torch.Tensor:
     """
     assert group_idx.dim() == 2 and group_idx.is_contiguous()
 
-    kernel = get_group_count_kernel(group_idx.shape[1], num_groups, get_num_sms())
+    # Allocate and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    device = group_idx.device
+    with get_device_guard(device):
+        kernel = get_group_count_kernel(group_idx.shape[1], num_groups, get_num_sms(device.index))
 
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
 
-    out = torch.zeros(num_groups, dtype=torch.int32, device='cuda')
-    kernel(group_idx, out)
+        out = torch.zeros(num_groups, dtype=torch.int32, device=device)
+        kernel(group_idx, out)
 
     return out
diff --git a/tile_kernels/moe/normalize_weight_kernel.py b/tile_kernels/moe/normalize_weight_kernel.py
index 236faa1..8016a5b 100644
--- a/tile_kernels/moe/normalize_weight_kernel.py
+++ b/tile_kernels/moe/normalize_weight_kernel.py
@@ -3,6 +3,8 @@
 import tilelang
 from tilelang import language as T
 
+from tile_kernels.utils import get_device_guard
+
 
 @tilelang.jit(
     pass_configs={
@@ -56,15 +58,20 @@ def normalize_weight(topk_weights: torch.Tensor) -> tuple[torch.Tensor, torch.Te
     assert topk_weights.dtype == torch.float32
 
     num_tokens, num_topk = topk_weights.shape
-    kernel = get_normalize_weight_kernel(num_topk)
 
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
+    # Allocate and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    device = topk_weights.device
+    with get_device_guard(device):
+        kernel = get_normalize_weight_kernel(num_topk)
+
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
 
-    denominator = torch.empty((num_tokens,), dtype=torch.float32, device='cuda')
-    normalized_weights = torch.empty((num_tokens, num_topk), dtype=torch.float32, device='cuda')
+        denominator = torch.empty((num_tokens,), dtype=torch.float32, device=device)
+        normalized_weights = torch.empty((num_tokens, num_topk), dtype=torch.float32, device=device)
 
-    if num_tokens > 0:
-        kernel(topk_weights, denominator, normalized_weights)
+        if num_tokens > 0:
+            kernel(topk_weights, denominator, normalized_weights)
 
     return (denominator, normalized_weights)
diff --git a/tile_kernels/moe/reduce_fused_kernel.py b/tile_kernels/moe/reduce_fused_kernel.py
index 970139f..c08245c 100644
--- a/tile_kernels/moe/reduce_fused_kernel.py
+++ b/tile_kernels/moe/reduce_fused_kernel.py
@@ -4,6 +4,7 @@
 from tilelang import language as T
 from typing import Optional, Union
 from tile_kernels.quant.common import *
+from tile_kernels.utils import get_device_guard
 
 
 @tilelang.jit(
@@ -108,29 +109,40 @@ def reduce_fused(
     else:
         assert sf is None, 'Only FP8 output supports sf.'
 
+    # Allocate (when needed) and launch on the device of the input tensor, which is
+    # not necessarily the current CUDA device.
+    device = x.device
+    assert token_topk_to_pos.device == device, f'token_topk_to_pos is on {token_topk_to_pos.device}, expected {device}'
+    if topk_weights is not None:
+        assert topk_weights.device == device, f'topk_weights is on {topk_weights.device}, expected {device}'
+    if x_sf is not None:
+        assert x_sf.device == device, f'x_sf is on {x_sf.device}, expected {device}'
+
     if out is not None:
         num_tokens_, hidden_ = out.shape
         assert num_tokens == num_tokens_ and hidden == hidden_
+        assert out.device == device, f'out is on {out.device}, expected {device}'
     else:
-        out = torch.empty((num_tokens, hidden), dtype=out_dtype, device='cuda')
+        out = torch.empty((num_tokens, hidden), dtype=out_dtype, device=device)
 
     if x_sf is not None:
         num_expanded_tokens_ = x_sf.shape[0]
         assert num_expanded_tokens == num_expanded_tokens_
 
-    kernel = get_reduce_fused_kernel(
-        hidden,
-        num_topk,
-        T.dtype(in_dtype),
-        T.dtype(out_dtype),
-        sf is not None,
-        topk_weights is not None,
-        x_sf is not None,
-    )
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
-
-    if num_tokens > 0:
-        kernel(x, topk_weights, token_topk_to_pos, out, sf, x_sf)
+    with get_device_guard(device):
+        kernel = get_reduce_fused_kernel(
+            hidden,
+            num_topk,
+            T.dtype(in_dtype),
+            T.dtype(out_dtype),
+            sf is not None,
+            topk_weights is not None,
+            x_sf is not None,
+        )
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
+
+        if num_tokens > 0:
+            kernel(x, topk_weights, token_topk_to_pos, out, sf, x_sf)
 
     return out
diff --git a/tile_kernels/utils.py b/tile_kernels/utils.py
index 06cc4b0..9b93c51 100644
--- a/tile_kernels/utils.py
+++ b/tile_kernels/utils.py
@@ -1,3 +1,20 @@
+import contextlib
+
+import torch
+
+
+def get_device_guard(device: torch.device):
+    """Return a context manager that makes `device` the current CUDA device.
+
+    Needed by the MoE entry points because output allocation, JIT compilation and
+    kernel launch all resolve `cuda` against the current device, not against the
+    device the input tensors already live on.
+    """
+    if device.type == 'cuda':
+        return torch.cuda.device(device)
+    return contextlib.nullcontext()
+
+
 def ceil_div(x: int, y: int) -> int:
     return (x + y - 1) // y
 

From ae28577768b39b200403b37813021a7c263919ab Mon Sep 17 00:00:00 2001
From: PerryLink <255665900+PerryLink@users.noreply.github.com>
Date: Fri, 18 Sep 2026 12:49:06 +0800
Subject: [PATCH 2/2] Guard the remaining moe/ launch paths and make the
 SM-count cache device-specific

Follow-up to the review on #30:

- inplace_unique_group_indices: derive the device from group_indices, size the
  grid from that device's SM count and guard build + launch. It was the one
  remaining moe/ path that also fed a wrong-device value into the kernel.
- topk_gate, top2_sum_gate, mask_indices_by_tp, topk_sum_and_topk_group_idx: wrap
  kernel build and launch in get_device_guard so compilation and launch happen
  with the input device current, matching the six entry points fixed in the
  first commit.
- config: resolve the current device index before the cached call, so the
  no-argument path is cached per concrete device instead of under (None,) and
  set_num_sms()/generate_num_sms() are no longer device-unsafe.
- reduce_fused: validate sf's device alongside the other kernel arguments.
---
 tile_kernels/config.py                        | 34 +++++++++++-----
 .../inplace_unique_group_indices_kernel.py    | 17 +++++---
 tile_kernels/moe/mask_indices_by_tp_kernel.py | 17 +++++---
 tile_kernels/moe/reduce_fused_kernel.py       |  2 +
 tile_kernels/moe/top2_sum_gate_kernel.py      | 40 ++++++++++---------
 tile_kernels/moe/topk_gate_kernel.py          | 13 +++---
 .../moe/topk_sum_and_topk_group_idx_kernel.py | 20 ++++++----
 7 files changed, 90 insertions(+), 53 deletions(-)

diff --git a/tile_kernels/config.py b/tile_kernels/config.py
index 28e6d7e..1b17321 100644
--- a/tile_kernels/config.py
+++ b/tile_kernels/config.py
@@ -7,16 +7,34 @@
 
 
 @functools.lru_cache(maxsize=None)
-def get_device_num_sms(device_index: Optional[int] = None) -> int:
-    # The device index is part of the cache key: properties are looked up on the
-    # requested device and a second device never reuses the first device's value.
-    if device_index is None:
-        device_index = torch.cuda.current_device()
+def _get_device_num_sms(device_index: int) -> int:
     prop = torch.cuda.get_device_properties(device_index)
     return prop.multi_processor_count
 
 
+@functools.lru_cache(maxsize=None)
+def _get_max_smem_per_sm(device_index: int) -> int:
+    prop = torch.cuda.get_device_properties(device_index)
+    return prop.shared_memory_per_multiprocessor
+
+
+def _resolve_device_index(device_index: Optional[int]) -> int:
+    # Resolve before the cached call: caching under `None` would pin whichever device
+    # happened to be current on the first call and reuse that value on every other device.
+    if device_index is None:
+        return torch.cuda.current_device()
+    return device_index
+
+
+def get_device_num_sms(device_index: Optional[int] = None) -> int:
+    # The cached key is always a concrete device index, so a second device never reuses
+    # the first device's value, whether or not the caller passed an index.
+    return _get_device_num_sms(_resolve_device_index(device_index))
+
+
 def set_num_sms(num_sms: int) -> None:
+    # The override below is process-global, so the bound is checked against the current
+    # device only; the value is not re-validated per device when it is later used.
     global _num_sms
     assert 0 < num_sms <= get_device_num_sms()
     _num_sms = num_sms
@@ -31,9 +49,5 @@ def get_num_sms(device_index: Optional[int] = None) -> int:
     return _num_sms
 
 
-@functools.lru_cache(maxsize=None)
 def get_max_smem_per_sm(device_index: Optional[int] = None) -> int:
-    if device_index is None:
-        device_index = torch.cuda.current_device()
-    prop = torch.cuda.get_device_properties(device_index)
-    return prop.shared_memory_per_multiprocessor
+    return _get_max_smem_per_sm(_resolve_device_index(device_index))
diff --git a/tile_kernels/moe/inplace_unique_group_indices_kernel.py b/tile_kernels/moe/inplace_unique_group_indices_kernel.py
index 0222689..cc0c4b3 100644
--- a/tile_kernels/moe/inplace_unique_group_indices_kernel.py
+++ b/tile_kernels/moe/inplace_unique_group_indices_kernel.py
@@ -4,7 +4,7 @@
 from tilelang import language as T
 
 from tile_kernels.config import get_num_sms
-from tile_kernels.utils import align
+from tile_kernels.utils import align, get_device_guard
 
 
 @tilelang.jit(
@@ -60,10 +60,15 @@ def inplace_unique_group_indices(group_indices: torch.Tensor, num_groups: int) -
 
     num_topk = group_indices.shape[1]
     num_groups_aligned = align(num_groups, 64)
-    kernel = get_inplace_unique_group_indices_kernel(num_topk, num_groups_aligned, get_num_sms())
 
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
+    # Compile and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    device = group_indices.device
+    with get_device_guard(device):
+        kernel = get_inplace_unique_group_indices_kernel(num_topk, num_groups_aligned, get_num_sms(device.index))
 
-    if group_indices.shape[0] > 0:
-        kernel(group_indices)
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
+
+        if group_indices.shape[0] > 0:
+            kernel(group_indices)
diff --git a/tile_kernels/moe/mask_indices_by_tp_kernel.py b/tile_kernels/moe/mask_indices_by_tp_kernel.py
index f30d2e1..fce9df4 100644
--- a/tile_kernels/moe/mask_indices_by_tp_kernel.py
+++ b/tile_kernels/moe/mask_indices_by_tp_kernel.py
@@ -3,6 +3,8 @@
 import tilelang
 from tilelang import language as T
 
+from tile_kernels.utils import get_device_guard
+
 
 @tilelang.jit(
     pass_configs={
@@ -61,13 +63,16 @@ def mask_indices_by_tp(indices: torch.Tensor, n: int, num_ep_ranks: int, tp_rank
     num_topk = indices.shape[1]
     per_gpu = n // num_ep_ranks
     per_dp = num_tp_ranks * per_gpu
-    kernel = get_mask_indices_by_tp_kernel(num_topk, T.dtype(indices.dtype))
+    # Compile and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    with get_device_guard(indices.device):
+        kernel = get_mask_indices_by_tp_kernel(num_topk, T.dtype(indices.dtype))
 
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
 
-    masked_indices = torch.empty_like(indices)
-    if indices.shape[0] > 0:
-        kernel(indices, masked_indices, per_gpu, per_dp, num_tp_ranks, tp_rank)
+        masked_indices = torch.empty_like(indices)
+        if indices.shape[0] > 0:
+            kernel(indices, masked_indices, per_gpu, per_dp, num_tp_ranks, tp_rank)
 
     return masked_indices
diff --git a/tile_kernels/moe/reduce_fused_kernel.py b/tile_kernels/moe/reduce_fused_kernel.py
index c08245c..4e95b73 100644
--- a/tile_kernels/moe/reduce_fused_kernel.py
+++ b/tile_kernels/moe/reduce_fused_kernel.py
@@ -117,6 +117,8 @@ def reduce_fused(
         assert topk_weights.device == device, f'topk_weights is on {topk_weights.device}, expected {device}'
     if x_sf is not None:
         assert x_sf.device == device, f'x_sf is on {x_sf.device}, expected {device}'
+    if sf is not None:
+        assert sf.device == device, f'sf is on {sf.device}, expected {device}'
 
     if out is not None:
         num_tokens_, hidden_ = out.shape
diff --git a/tile_kernels/moe/top2_sum_gate_kernel.py b/tile_kernels/moe/top2_sum_gate_kernel.py
index 768b52c..ce36938 100644
--- a/tile_kernels/moe/top2_sum_gate_kernel.py
+++ b/tile_kernels/moe/top2_sum_gate_kernel.py
@@ -4,7 +4,7 @@
 from typing import Optional
 import os
 
-from tile_kernels.utils import align, ceil_div
+from tile_kernels.utils import align, ceil_div, get_device_guard
 from tile_kernels.moe.scoring import ScoringFunc, softplus
 from tile_kernels.moe.common import get_topk_group_idx
 
@@ -403,22 +403,26 @@ def top2_sum_gate(
         assert fix_routing_mask.dtype == torch.bool
         assert fix_routing_mask.dim() == 1 and fix_routing_mask.size(0) == num_tokens
 
-    kernel = get_top2_sum_gate_kernel(
-        ScoringFunc.from_str(scoring_func).value,
-        num_topk,
-        num_topk_groups, num_groups,
-        num_routed_experts,
-        mask is not None, fix_routing_mask is not None,
-        unmapped_topk_idx is not None, to_physical_map is not None,
-    )  # fmt: off
-
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
-
-    kernel(logits, bias,
-           mask, fix_routing_mask, to_physical_map, logical_count,
-           topk_idx, unmapped_topk_idx, topk_weights,
-           num_extra_experts, routed_scaling_factor,
-           ep_rank, num_ep_ranks, tp_rank, num_tp_ranks)  # fmt: off
+    # Compile and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    device = logits.device
+    with get_device_guard(device):
+        kernel = get_top2_sum_gate_kernel(
+            ScoringFunc.from_str(scoring_func).value,
+            num_topk,
+            num_topk_groups, num_groups,
+            num_routed_experts,
+            mask is not None, fix_routing_mask is not None,
+            unmapped_topk_idx is not None, to_physical_map is not None,
+        )  # fmt: off
+
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
+
+        kernel(logits, bias,
+               mask, fix_routing_mask, to_physical_map, logical_count,
+               topk_idx, unmapped_topk_idx, topk_weights,
+               num_extra_experts, routed_scaling_factor,
+               ep_rank, num_ep_ranks, tp_rank, num_tp_ranks)  # fmt: off
 
     return topk_idx, topk_weights
diff --git a/tile_kernels/moe/topk_gate_kernel.py b/tile_kernels/moe/topk_gate_kernel.py
index 1b2ce99..206e090 100644
--- a/tile_kernels/moe/topk_gate_kernel.py
+++ b/tile_kernels/moe/topk_gate_kernel.py
@@ -4,7 +4,7 @@
 import torch
 from tilelang import language as T
 
-from tile_kernels.utils import align
+from tile_kernels.utils import align, get_device_guard
 
 
 @tilelang.jit(
@@ -81,10 +81,13 @@ def topk_gate(scores: torch.Tensor, num_topk: int) -> torch.Tensor:
     if num_tokens == 0:
         return topk_idx
 
-    kernel = get_topk_gate_kernel(num_experts, num_topk)
+    # Compile and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    with get_device_guard(scores.device):
+        kernel = get_topk_gate_kernel(num_experts, num_topk)
 
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
 
-    kernel(scores, topk_idx)
+        kernel(scores, topk_idx)
     return topk_idx
diff --git a/tile_kernels/moe/topk_sum_and_topk_group_idx_kernel.py b/tile_kernels/moe/topk_sum_and_topk_group_idx_kernel.py
index 77bb4f7..110782e 100644
--- a/tile_kernels/moe/topk_sum_and_topk_group_idx_kernel.py
+++ b/tile_kernels/moe/topk_sum_and_topk_group_idx_kernel.py
@@ -5,7 +5,7 @@
 from tilelang import language as T
 
 from tile_kernels.moe.common import get_topk_group_idx
-from tile_kernels.utils import align
+from tile_kernels.utils import align, get_device_guard
 
 
 @tilelang.jit(
@@ -90,13 +90,17 @@ def topk_sum_and_topk_group_idx(scores: torch.Tensor, num_topk_sum: int, num_top
     num_tokens, num_groups, num_experts_per_group = scores.shape
     assert num_topk_sum <= num_experts_per_group and num_topk_sum in (1, 2) and num_topk_groups <= num_groups
 
-    kernel = get_topk_sum_and_topk_group_idx_kernel(num_groups, num_experts_per_group, num_topk_groups, num_topk_sum)
-    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
-        print(kernel.get_kernel_source())
+    # Compile and launch on the device of the input tensor, which is not
+    # necessarily the current CUDA device.
+    with get_device_guard(scores.device):
+        kernel = get_topk_sum_and_topk_group_idx_kernel(num_groups, num_experts_per_group, num_topk_groups, num_topk_sum)
+        if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+            print(kernel.get_kernel_source())
 
-    topk_group_idx = torch.empty(num_tokens, num_topk_groups, dtype=torch.int64, device=scores.device)
-    if num_tokens == 0:
-        return topk_group_idx
+        topk_group_idx = torch.empty(num_tokens, num_topk_groups, dtype=torch.int64, device=scores.device)
+        if num_tokens == 0:
+            return topk_group_idx
+
+        kernel(scores.view(num_tokens, -1), topk_group_idx)
 
-    kernel(scores.view(num_tokens, -1), topk_group_idx)
     return topk_group_idx