A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 06:00 UTC
[Bugfix] [Tests] Enforce `out` tensor device in `kernel/moe/test_cutedsl_moe.py` (#39644)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
Apr 12, 2026 at 20:08 UTC
620e8924d9c6b2a0b1d49ac0dcf2588fffcbe390
1 file changed
+4
-2
tests/kernels/moe/test_cutedsl_moe.py
+4
-2
@@ -142,7 +142,9 @@ def prepare_inputs(
142
# Initialize the hidden_states_3d with ones instead of empty to avoid nan
143
# issue.
144
hidden_states_3d = torch.ones(
145
- (num_experts, max(masked_m), hidden_states.shape[1]), dtype=hidden_states.dtype
145
+ (num_experts, max(masked_m), hidden_states.shape[1]),
146
+ dtype=hidden_states.dtype,
147
+ device=hidden_states.device,
148
)
149
for i in range(num_experts):
150
hidden_states_3d[i, : masked_m[i], :] = hidden_states[topk_idx.view(-1) == i]
@@ -426,7 +428,7 @@ def test_flashinfer_cutedsl_moe_masked(
428
w1_alpha = 1.0 / (input_global_scale * w1_global_scale)
429
w2_alpha = 1.0 / (a2_global_scale * w2_global_scale)
430
429
- out = torch.empty_like(hidden_states_3d)
431
+ out = torch.empty_like(hidden_states_3d, device=hidden_states.device)
432
# Note: the 1st dim shouldn't be bs
433
wk = torch.empty(
434
num_experts,