[Correctness][Multi-GPU] MoE APIs allocate and launch on the current CUDA device instead of the input device

[GenTang]

Summary

Several public MoE APIs allocate outputs using an unqualified device="cuda" instead of the device of their input tensors.

When an input tensor is on a non-current CUDA device, for example:

• input tensor: cuda:1
• current CUDA device: cuda:0

the outputs are allocated on cuda:0. Kernel compilation and launch also occur without an explicit device guard, which can result in program failure.

Affected APIs include at least:

• get_fused_mapping
• expand_to_fused
• reduce_fused
• group_count
• aux_fi

For example, get_fused_mapping() currently allocates all output tensors with device="cuda":

pos_to_expert = torch.empty(..., device="cuda")
pos_to_token = torch.empty(..., device="cuda")
pos_to_token_topk = torch.empty(..., device="cuda")
token_topk_to_pos = torch.empty(..., device="cuda")
expert_start = torch.empty(..., device="cuda")
expert_end = torch.empty(..., device="cuda")
num_tokens_per_expert = torch.empty(..., device="cuda")
num_experts_per_sm = torch.empty(..., device="cuda")

Relevant code: get_fused_mapping output allocation

Root cause

1. Unqualified CUDA allocations

In PyTorch, device="cuda" resolves to the current CUDA device rather than the device of an existing input tensor.

Therefore, the following is not device-safe: out = torch.empty(shape, device="cuda") when the input may be on a non-current device. The output should instead be allocated using the relevant input device: out = torch.empty(shape, device=input.device)

2. Missing device guards around compilation and launch

Even after allocating outputs on the correct device, JIT compilation and kernel launch must run with the input device active as current CUDA device. By code, this is as follow:

device = input.device

with torch.cuda.device(device):
    kernel = get_kernel(...)
    kernel(input, output, ...)

Without this guard, device-specific compilation state and launch state can still be derived from the previously current device.

3. Device-property caches do not include the device index

The following functions are zero-argument lru_cache functions:

@functools.lru_cache(maxsize=None)
def get_device_num_sms() -> int:
    prop = torch.cuda.get_device_properties(torch.cuda.current_device())
    return prop.multi_processor_count

@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

Relevant code: config.py

The cache key is always the empty argument tuple. After the first call, switching to another CUDA device can still return the first device's SM count or shared-memory limit.

This is observable when the process uses GPUs with different properties.

Minimal reproduction

1. Multi-GPU bug

The reproduction of bug requires at least two GPU devices. The following uses group_count() since it has a small input and a single output:

import torch
from tile_kernels.moe import group_count

if torch.cuda.device_count() < 2:
    raise RuntimeError("This reproduction requires two visible CUDA devices")

# Keep cuda:0 as the current device.
torch.cuda.set_device(0)

# Allocate the input on cuda:1 without leaving cuda:1 current.
with torch.cuda.device(1):
    group_idx = torch.tensor(
        [[0, 1], [1, 0]],
        dtype=torch.int64,
        device="cuda:1",
    ).contiguous()

assert torch.cuda.current_device() == 0
assert group_idx.device == torch.device("cuda:1")

out = group_count(group_idx, num_groups=2)

# Surface asynchronous launch errors.
for device_index in (0, 1):
    torch.cuda.synchronize(device_index)

print("current device:", torch.cuda.current_device())
print("input device:", group_idx.device)
print("output device:", out.device)

assert out.device == group_idx.device

Run with synchronous CUDA error reporting:

CUDA_VISIBLE_DEVICES=0,1 CUDA_LAUNCH_BLOCKING=1 python repro.py

On the current implementation, the call will fail.

2. Device-property cache reproduction

This reproduction requires two visible devices with different SM counts:

import torch
import tile_kernels.config as config

if torch.cuda.device_count() < 2:
    raise RuntimeError("This reproduction requires two visible CUDA devices")

config.get_device_num_sms.cache_clear()

torch.cuda.set_device(0)
sm0 = config.get_device_num_sms()
actual0 = torch.cuda.get_device_properties(0).multi_processor_count

torch.cuda.set_device(1)
sm1 = config.get_device_num_sms()
actual1 = torch.cuda.get_device_properties(1).multi_processor_count

print("cuda:0:", sm0, actual0)
print("cuda:1:", sm1, actual1)

assert sm0 == actual0
assert sm1 == actual1

If the two devices have different SM counts, sm1 incorrectly retains the value returned for device 0. On identical GPUs, the cache is still device-insensitive, but the problem is numerically hidden because both devices happen to have the same property value.

Why existing tests did not catch this

Existing tests generally create inputs using an unqualified device="cuda" and run them on the current CUDA device. As a result, the input device, output device, compilation device, and launch device are normally all the same.

Proposed fix

1. Output allocation and device guard

Derive the device from the primary input:

device = input.device

if device.type != "cuda":
    raise ValueError(f"input must be a CUDA tensor, got {device}")

with torch.cuda.device(device):
    out = torch.empty(..., device=device)
    kernel = get_kernel(...)
    kernel(input, out, ...)

For APIs with multiple tensor inputs, validate their device relationship before entering the guard:

if mapping.device != input.device:
    raise ValueError(
        f"input and mapping must be on the same device, "
        f"got {input.device} and {mapping.device}"
    )

Optional preallocated outputs should be validated rather than moved implicitly.

2. Device-aware property caches

Make the device index part of the cache key(Callers can resolve the relevant device index from their primary input):

@functools.lru_cache(maxsize=None)
def get_device_num_sms(device_index: int) -> int:
    return torch.cuda.get_device_properties(
        device_index
    ).multi_processor_count


@functools.lru_cache(maxsize=None)
def get_max_smem_per_sm(device_index: int) -> int:
    return torch.cuda.get_device_properties(
        device_index
    ).shared_memory_per_multiprocessor

If set_num_sms() is intended to override device properties, its currently process-global _num_sms state may also need to be made device-specific or explicitly documented as a global override.

Suggested test coverage

Most of the fix can be tested without specialized multi-GPU hardware:

1. Mock output allocation and verify that every automatically created tensor receives input.device.
2. Mock CUDA device contexts and verify that kernel construction and invocation occur under the input device.
3. Mock two devices with different SM/shared-memory properties and verify independent cached results.
4. Run all existing single-GPU correctness tests.

An optional two-GPU integration test should additionally verify:

torch.cuda.set_device(0)
input = make_input(device="cuda:1")
output = api(input)
assert output.device == input.device

A heterogeneous pair with different SM is only required to reproduce the device-property cache issue on real hardware; deterministic mocked tests can cover this behavior in regular CI.