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

[PerryLink]

Summary

Six modules under tile_kernels/moe/ (16 allocation sites in total, covering the five entry
points named in the issue plus normalize_weight, which carries the byte-identical pattern)
allocate their output tensors with an unqualified device='cuda' and then launch the kernel
without a device guard. device='cuda' resolves to the current CUDA device, not to the device
of the input tensor, so when the input lives on a non-current device (for example input on
cuda:1 while torch.cuda.current_device() is 0) the outputs are allocated on the wrong device
and compilation/launch state is derived from the wrong device.

This PR derives the device from the input tensor, allocates and launches under
torch.cuda.device(device), and makes the device index part of the cache key of the two
device-property helpers in tile_kernels/config.py. A second commit extends the guard to five
sibling entry points that allocated correctly but still compiled and launched unguarded, and
fixes the no-argument paths of the property cache.

Fixes #28.

Affected allocation sites (default branch at 36d9e45)

API	Site
get_fused_mapping	tile_kernels/moe/get_fused_mapping_kernel.py:208-215 (8 tensors)
expand_to_fused	tile_kernels/moe/expand_to_fused_kernel.py:126
expand_to_fused_with_sf	tile_kernels/moe/expand_to_fused_kernel.py:192-193
reduce_fused	tile_kernels/moe/reduce_fused_kernel.py:115
group_count	tile_kernels/moe/group_count_kernel.py:66
aux_fi	tile_kernels/moe/aux_fi_kernel.py:69
normalize_weight	tile_kernels/moe/normalize_weight_kernel.py:64-65

normalize_weight is not listed in the issue but carries the same pattern in the same package, so
it is fixed here too.

Newly guarded compile/launch sites (second commit)

These five entry points already allocated with device=<input>.device or torch.empty_like, so
they had no wrong-device allocation, but their kernel was built and launched with no device guard.
Line numbers below are for the version in this PR.

API	Site
inplace_unique_group_indices	tile_kernels/moe/inplace_unique_group_indices_kernel.py:66-74 (also get_num_sms(device.index))
topk_gate	tile_kernels/moe/topk_gate_kernel.py:86-92
top2_sum_gate	tile_kernels/moe/top2_sum_gate_kernel.py:408-422
mask_indices_by_tp	tile_kernels/moe/mask_indices_by_tp_kernel.py:68-76
topk_sum_and_topk_group_idx	tile_kernels/moe/topk_sum_and_topk_group_idx_kernel.py:95-104

The primary input chosen per entry point: group_indices for inplace_unique_group_indices,
scores for topk_gate (its only tensor input), logits for top2_sum_gate (first argument;
bias, mask, fix_routing_mask, to_physical_map, logical_count and unmapped_topk_idx are
per-token/per-expert companions of logits, and both outputs are already allocated with
device=logits.device), indices for mask_indices_by_tp, and scores for
topk_sum_and_topk_group_idx (its only tensor input, and the device its output is already
allocated from).

In topk_sum_and_topk_group_idx the num_tokens == 0 early return sits between the build and the
launch. It is kept verbatim and stays inside the guard block, so the guard covers the build and the
launch on the non-empty path and the empty path returns the empty tensor while still exiting the
context manager; no line was reordered and no condition was rewritten.

Mechanism

1. torch.empty(..., device='cuda') binds to torch.cuda.current_device(). With a cuda:1
   input and cuda:0 current, every output lands on cuda:0.
2. get_num_sms() read torch.cuda.get_device_properties(torch.cuda.current_device()) behind a
   zero-argument functools.lru_cache(maxsize=None), so the SM count was resolved once, for
   whichever device happened to be current on the first call, and then reused for every later call
   on every other device. This value sizes a real tensor (num_experts_per_sm) and a compile-time
   kernel argument (num_sms grid size), not just a heuristic.
3. TileLang resolves the compilation target and the launch stream from the ambient current device,
   so even a correctly allocated output can be paired with a kernel built/launched for another
   device.

Fix

• New helper get_device_guard(device) in tile_kernels/utils.py: returns
  torch.cuda.device(device) for a CUDA device and contextlib.nullcontext() otherwise. The
  guard is entered before the kernel is built and before the outputs are allocated and launched,
  so allocation, compilation and launch all observe the input device as the current device.
• device = <primary input>.device in every affected entry point; every automatically created
  tensor now takes device=device.
• config.get_device_num_sms / config.get_max_smem_per_sm / config.get_num_sms keep their
  public signatures but dispatch to private @functools.lru_cache functions keyed by a concrete
  int device index. device_index=None is resolved with torch.cuda.current_device() before
  the cached call, so the key is never None and zero-argument callers (including
  set_num_sms() and tile_kernels/testing/generator.py::generate_num_sms()) are per-device
  correct rather than pinned to the first current device. The MoE callers pass device.index.
• Tensor inputs that must share the primary input's device are validated with an assertion
  (expand_to_fused*), and in reduce_fused the optional preallocated out is validated instead
  of being moved implicitly, as the issue suggests. In the second commit the optional sf tensor
  is validated the same way, alongside token_topk_to_pos, topk_weights and x_sf.

Kernel-cache caveat / needs two GPUs

The guard makes the input device current around allocation, compilation and launch. It does not
give each device its own compiled kernel, because TileLang's JIT wrapper caches per factory
argument tuple. What is established from the source:

1. The in-memory cache key carries no device component. tilelang/jit/__init__.py holds
   self._kernel_cache: dict[tuple, Kernel] = {} and fills it from
   key, kernel_args = self.func.parse_args(*args, **kwargs) followed by
   if kernel is None: kernel = self.compile(...). That parse_args
   (tilelang/language/eager/builder.py:1577-1590) returns (bound.p1_key, p2_key), where
   p1_key comes from self._argument_binder.bind(args, kwargs) (the factory arguments) and
   p2_key from tir_temp._parse_phase2_key(**bound.tensor_args, **bound.compile_kwargs)
   (tensor/compile arguments). Nothing in that tuple identifies a CUDA device. So on a host with
   identical GPUs, a second call with the same factory arguments hits the kernel object that was
   compiled while the first device was current.
2. The on-disk cache key includes target (hence the architecture) but no explicit device index:
   tilelang/cache/kernel_cache.py:269-284 builds
   key_data = {"func": sha256(func_binary), "out_idx": ..., "args_repr": ..., "target": str(target), ...}.
   Same-architecture devices therefore share a disk entry; different-architecture devices do not.
3. The launch path itself is device-aware: tilelang/jit/adapter/base.py:72-95 resolves the
   stream from the current device at call time (get_stream(current_device()) /
   torch.device("cuda", current_device())).

What is therefore not established, and is the reason this PR no longer claims that
"allocation, compilation and launch all agree on the input device": whether launching a kernel
object that was compiled or loaded while device A was current, now with device B current, succeeds
or fails. That is a runtime question about per-context modules, not a cache-key question, and
answering it needs two GPUs.

Consequence for the code in this PR, stated precisely:

• For factories whose arguments are device-independent, the guard alone cannot produce a second
  kernel. topk_gate is the clearest case: get_topk_gate_kernel(num_experts, num_topk) has no
  device-derived argument, so a second device reuses the first device's kernel object.
• For get_fused_mapping and inplace_unique_group_indices, num_sms is a factory argument and
  is now device-specific (get_num_sms(device.index)), so devices with different SM counts get
  separate kernels, while identical GPUs still share one.

What a reviewer with a 2-GPU machine should test (this is the open item):

1. With cuda:0 current, call each of the twelve guarded entry points on cuda:0 inputs, so the
   first kernel object for each factory argument tuple is built there.
2. Then set device 1 current (torch.cuda.set_device(1)) and call the same entry points with the
   same shapes and dtypes but cuda:1 inputs.
3. Record per entry point: whether the call raises, whether the returned tensors are on cuda:1,
   whether the result matches the cuda:0 result numerically, and whether
   torch.cuda.synchronize() succeeds on both devices afterwards.
4. Repeat with two GPUs of different architecture (or at least different SM counts) to separate the
   in-memory cache case (1) from the disk cache case (2).

Redesigning that cache (device in the factory arguments, a target/device-scoped key, or a
per-device kernel registry) is deliberately out of scope for this PR: it changes TileLang, not the
call sites, and its correct shape depends on the answer to the two-GPU question above. It is
offered as a follow-up.

Why it is safe

• Single-GPU path: x.device is exactly the device device='cuda' would have resolved to,
  torch.cuda.device(x.device) is a no-op when the device is already current, and
  get_num_sms(x.device.index) returns the same cached integer as get_num_sms(). Allocation,
  layout, dtype, shapes, kernel arguments and return values are unchanged.
• The set_num_sms() override keeps its exact previous semantics: a process-global override that
  wins over the per-device value. Its bound is still checked against the current device only, and
  the value is not re-validated per device when it is later used; that limitation is now stated in
  a comment next to the function rather than being implicit.
• Existing multi-GPU callers that worked by accident (input on the current device, or identical
  GPUs) are unaffected; identical GPUs still hit the same cached SM count because the cache key is
  the device index and the values are equal.
• Multi-GPU callers that were broken now get outputs on the input device, matching what
  tile_kernels/torch/*.py reference implementations already do.

Changes since the first commit

• tile_kernels/moe/inplace_unique_group_indices_kernel.py: device = group_indices.device,
  get_num_sms(device.index) instead of get_num_sms(), and the kernel build plus launch wrapped
  in get_device_guard(device).
• tile_kernels/moe/topk_gate_kernel.py: build plus launch wrapped in
  get_device_guard(scores.device). No co-location assertion was added: this module has no
  device-assertion style to follow, and scores is its only tensor input.
• tile_kernels/moe/top2_sum_gate_kernel.py: build plus launch wrapped in
  get_device_guard(logits.device) (see the primary-input note above).
• tile_kernels/moe/mask_indices_by_tp_kernel.py: build plus launch wrapped in
  get_device_guard(indices.device).
• tile_kernels/moe/topk_sum_and_topk_group_idx_kernel.py: build plus launch wrapped in
  get_device_guard(scores.device). The num_tokens == 0 early return stays between them inside
  the guard block (see the note above); no line was reordered.
• tile_kernels/config.py: the cached helpers are now private functions keyed by a concrete
  int; None is resolved to torch.cuda.current_device() before the cached call, so the
  no-argument callers named in the review (set_num_sms(),
  tile_kernels/testing/generator.py::generate_num_sms()) are per-device correct. Public
  signatures and the set_num_sms override semantics are unchanged.
• tile_kernels/moe/reduce_fused_kernel.py: the optional sf tensor (a kernel argument when
  fp8_format == 'e4m3') is now validated against the input device like the other tensors.

Validation

• Second commit: py_compile on all seven changed files: passed (Python 3.12.14), compiled into a
  scratch directory so the payload files are untouched.
• Second commit: git apply --check and git apply of the follow-up patch onto the head commit
  cfce997: passed, and the applied tree is byte-identical (SHA-256) to the tree that was syntax
  checked; the patch touches exactly those seven files and every other file in the head tree is
  byte-identical before and after.
• Second commit: each changed file was compared against the GitHub blob it was edited from (SHA-1
  of blob <len>\0<bytes> equals the tree entry), so the diff base is the reviewed revision and
  not a local approximation.
• Every changed line is within the repo's line-length = 150 ruff limit; the files are LF-only,
  UTF-8 without BOM, have no trailing whitespace and keep their trailing newline; no reformatting.
• Not validated: no CUDA-capable environment was used, so no kernel was compiled or launched and
  no two-GPU reproduction was executed. The device-guard placement inside TileLang's JIT wrapper,
  and the kernel-cache behaviour in the section above, are argued from the source, not measured.
  The behaviour of set_num_sms() as a global override under simultaneous multi-device use is
  left as-is and only documented.

Out of scope

The same unqualified device='cuda' pattern exists outside moe/
(tile_kernels/quant/per_channel_cast_fused_kernel.py:198-199,
tile_kernels/quant/swiglu_forward_and_per_token_cast_kernel.py:252,
tile_kernels/transpose/batched_transpose_kernel.py:115, and the
device: torch.device = 'cuda' default of tile_kernels/quant/common.py::alloc_scaling_factors,
which seven quant call sites rely on). Those are the same class of bug but a separate change set;
this PR stays inside moe/ plus the shared property cache the issue points at.
tile_kernels/testing/generator.py (six device='cuda' sites, test-data generation for the
current device) and the test modules are also unchanged.

Also out of scope, and offered as a follow-up: the TileLang kernel-cache redesign described in the
caveat above.

Prepared with AI assistance; reviewed before submission.

[ds-review-bot]

🔵 suggestion: set_num_sms() still validates num_sms against get_device_num_sms() with no argument, i.e. against the current device only. On hosts with heterogeneous GPUs the process-global override may be accepted based on one device and then applied to another with fewer SMs. This is pre-existing behaviour and the PR documents the override as intentionally global, so it is not a blocker; consider a follow-up that validates against the minimum SM count over visible devices, or makes the override per-device.

_{🤖 v5}

[ds-review-bot]

🟡 warning: The device guard only fixes the first call per kernel-factory argument tuple. The get_*_kernel(...) factories are decorated with @tilelang.jit, which caches the compiled JITKernel by factory arguments. The first call compiles/loads the module on whichever device is current at that moment; every subsequent call with the same argument tuple returns that cached kernel regardless of the device set by get_device_guard. So calling e.g. get_fused_mapping with identical args on cuda:0 and then cuda:1 will reuse the cuda:0 module and the second launch either errors or targets the wrong device. This directly contradicts the PR statement that "allocation, compilation and launch all agree on the input device". Please verify with a 2-GPU test that invokes each entry point again on a non-current device after a first call on the current device, and if it fails, include the device in the factory arguments/cache key (or maintain a per-device kernel cache).

_{🤖 v4}

[ds-review-bot]

🟡 warning: Making device_index part of the lru_cache key only helps callers that pass it. When device_index is None, the resolved torch.cuda.current_device() is not recorded in the key, so the first current device's multi_processor_count stays cached under (None,) and is returned even after the current device changes. set_num_sms() (line 21) and tile_kernels/testing/generator.py::generate_num_sms() still call get_device_num_sms() with no argument, so the original multi-GPU SM-count bug survives on those paths, and set_num_sms's assertion can be checked against the wrong device. Consider resolving the device first and dispatching as get_device_num_sms(torch.cuda.current_device()) (so the concrete index is cached), or passing the device index from those callers.

_{🤖 v4}

[ds-review-bot]

🔵 suggestion: The new device validation covers token_topk_to_pos, topk_weights, x_sf and the optional out, but omits sf, which is also a kernel argument when fp8_format == 'e4m3' (passed at line 146). A cross-device sf therefore still slips past validation and fails at launch. Add if sf is not None: assert sf.device == device, ... alongside the other checks.

_{🤖 v4}

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

未发现本次变更引入的明确缺陷。8 个文件通过语法检查，模拟检查覆盖了 7 个入口的设备选择与恢复；当前环境缺少 PyTorch 和 TileLang，未验证真实多 GPU 编译和执行。

v5

Reviewed commit cfce997 ([BugFix][Multi-GPU] Allocate and launch MoE outputs on the input device). Verdict: approve. The change matches its description: all 16 device='cuda' allocation sites across six tile_kernels/moe/ modules now derive device from the primary input tensor, and kernel construction, output allocation and launch are wrapped in the new get_device_guard(device) context manager in tile_kernels/utils.py (torch.cuda.device(device) for CUDA, nullcontext() otherwise; accepts a torch.device and is a no-op when that device is already current, so the single-GPU path is unchanged). config.get_device_num_sms / get_max_smem_per_sm / get_num_sms gain an optional device_index that becomes part of the lru_cache key, fixing the stale-SM-count problem; MoE callers pass device.index, while all pre-existing callers (engram, mhc ops, quant, tests, __init__.py re-exports) keep the zero-argument form and are behaviourally unchanged. In get_fused_mapping the guard correctly covers get_num_sms(device.index), all eight allocations (including num_experts_per_sm, whose shape depends on the per-device SM count) and the launch. reduce_fused validates the preallocated out and secondary inputs (token_topk_to_pos/topk_weights/x_sf) with descriptive assertions rather than moving them, and expand_to_fused* assert co-location of secondary inputs, as the issue suggested; the out_sf shape selection, slice and .T handling are unchanged. Verification: git grep device='cuda' confirms nothing remains under moe/ (remaining hits are exactly the out-of-scope files listed in the PR plus testing/generator.py); no changed line exceeds the 150-column ruff limit. Caveat: only static review was possible in this environment (no Python/CUDA execution), consistent with the PR's own 'Not validated' note. One non-blocking suggestion left on set_num_sms.

v4

The change is directionally correct: the new get_device_guard helper plus device = <primary_input>.device makes the six touched MoE entry points allocate their outputs on the input tensor's device, and config.get_device_num_sms/get_max_smem_per_sm/get_num_sms now accept a device_index that participates in the cache key. However, the patch is incomplete and its central guarantee is not established. (1) Only six of the moe/ entry points are guarded; several others (topk_gate, top2_sum_gate, topk_sum_and_topk_group_idx, mask_indices_by_tp, and especially inplace_unique_group_indices) still compile/launch on the ambient current device, and inplace_unique_group_indices also sizes its grid with get_num_sms() for the wrong device, so the stated bug ("allocate and launch ... on the input device") is only partially fixed. (2) The guard only helps if TileLang recompiles per device; if its @tilelang.jit cache is keyed by factory arguments (the usual behaviour), a second device reuses the first device's compiled kernel and the guard is a no-op after the first call. The PR itself admits this is "argued from the code, not measured". (3) The device_index=None path still caches the first current device's SM count, so set_num_sms() and testing/generator.generate_num_sms() remain device-unsafe. (4) reduce_fused validates every tensor argument except sf. None of the existing tests exercise a non-current CUDA device, so all of these gaps are untested.

Files reviewed: 8
Issues found: 🟡 5 warning | 🔵 4 suggestion
Inline comments posted: 4
General comments (无法定位到 diff): 5

---

📍 未定位到 diff 的评论

🟡 warning tile_kernels/moe/inplace_unique_group_indices_kernel.py:L63: This entry point is declared out of scope because it "allocates nothing", but the issue explicitly covers launch and this is the one remaining moe/ path that also derives kernel parameters from the wrong device: get_inplace_unique_group_indices_kernel(..., get_num_sms()) uses the ambient device's SM count (sizing grid_x = num_sms * 2), and the kernel is built/launched at line 69 with no device guard. With group_indices on cuda:1 and torch.cuda.current_device() == 0 it uses device 0's SM count and launches on device 0. It should follow the same pattern as the fixed sites: device = group_indices.device, get_num_sms(device.index), and wrap build/launch in get_device_guard(device). _{🤖 v4}

🟡 warning tile_kernels/moe/top2_sum_gate_kernel.py:L406: top2_sum_gate allocates topk_idx/topk_weights from logits.device (good) but still builds and launches get_top2_sum_gate_kernel on the ambient current device. By the PR's own mechanism #3, TileLang resolves the target device at compile/launch time, so with logits on cuda:1 while the current device is 0 this kernel is compiled/launched for cuda:0. The entry point should be wrapped in get_device_guard(logits.device) like the six fixed sites (and ideally the co-located inputs validated). _{🤖 v4}

🟡 warning tile_kernels/moe/topk_gate_kernel.py:L84: Same inconsistency: the output is correctly allocated from scores.device at line 80, but get_topk_gate_kernel(...) is built and launched (line 89) without a get_device_guard(scores.device), so a scores tensor on a non-current device is compiled/launched against the current device. If the guard is required for the six fixed entry points, it is required here too. _{🤖 v4}

🔵 suggestion tile_kernels/moe/mask_indices_by_tp_kernel.py:L64: mask_indices_by_tp uses torch.empty_like(indices) (device-correct) but builds/launches the kernel at lines 64/71 without a device guard, so the same non-current-device defect as the fixed entry points remains. Please wrap build and launch in get_device_guard(indices.device) for consistency with the stated fix. _{🤖 v4}

🔵 suggestion tile_kernels/moe/topk_sum_and_topk_group_idx_kernel.py:L93: topk_sum_and_topk_group_idx allocates topk_group_idx on scores.device but builds/launches the kernel at lines 93/101 without a device guard, leaving the ambient-device compile/launch issue in place for a non-current scores tensor. Apply the same guard used in the six fixed entry points. _{🤖 v4}

[PerryLink]

Thanks for the review.

Fixed in the follow-up commit:

• All five unguarded entry points you named (inplace_unique_group_indices, topk_gate,
  top2_sum_gate, mask_indices_by_tp, topk_sum_and_topk_group_idx) now build and launch inside
  get_device_guard(<primary input>.device); inplace_unique_group_indices also passes
  get_num_sms(device.index). In topk_sum_and_topk_group_idx the num_tokens == 0 early return
  stays between them inside the guard block, so build and launch are both covered.
• config.py: device_index=None is resolved to torch.cuda.current_device() before the cached
  call, so the key is always a concrete device index and zero-argument callers (set_num_sms(),
  generate_num_sms()) are per-device correct. Public signatures unchanged.
• reduce_fused: sf is now validated alongside the other tensors.

Deliberately not changed:

• The TileLang kernel-factory cache: its key is the factory argument tuple with no device, so a
  second identical-GPU device can receive the first device's kernel object. Whether that launch
  succeeds needs two GPUs to answer; fixing it changes TileLang, not the call sites, so the PR body
  states what is established and offers it as a follow-up.
• set_num_sms() keeps its process-global semantics; the current-device-only bound is now
  documented in a comment.

One thing we could not do here: we have no CUDA environment, so the two-GPU run described in the new
"Kernel-cache caveat / needs two GPUs" section of the PR body is still open. Its result decides
whether the factory cache needs a device component, so if a maintainer or reviewer can run it we
will follow up with whatever it shows.