[BugFix][Multi-GPU] Allocate and launch MoE outputs on the input device

tile_kernels/config.py

@@ -1,29 +1,53 @@
 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: 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
 
 
-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())
-    return prop.shared_memory_per_multiprocessor
+def get_max_smem_per_sm(device_index: Optional[int] = None) -> int:
+    return _get_max_smem_per_sm(_resolve_device_index(device_index))

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

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

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())
 

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

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)

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

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)