@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 44cf3f0466

[XPU][Tests] Make tests device-agnostic (#51968)

Signed-off-by: pmanczak <pawel.manczak@intel.com>

pmanczak committed Aug 20, 2026 at 11:25 UTC 44cf3f046604188729572e22cabd747357098c2a

12 files changed +148 -111

tests/conftest.py

+6 -7

@@ -284,19 +284,18 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool):

284 def workspace_init():

285 """Initialize the workspace manager for tests that need it.

286

287 - This fixture initializes the workspace manager with a CUDA device

288 - if available, and resets it after the test completes. Tests that

289 - create a full vLLM engine should NOT use this fixture as the engine

290 - will initialize the workspace manager itself.

287 + This fixture initializes the workspace manager with the current

288 + platform's accelerator device if available, and resets it after the test

289 + completes. Tests that create a full vLLM engine should NOT use this

290 + fixture as the engine will initialize the workspace manager itself.

291 """

292 from vllm.v1.worker.workspace import (

293 init_workspace_manager,

294 reset_workspace_manager,

295 )

296

297 - if torch.cuda.is_available():

298 - device = torch.device("cuda:0")

299 - init_workspace_manager(device)

297 + if torch.accelerator.is_available():

298 + init_workspace_manager(torch.device(0))

299 yield

300 reset_workspace_manager()

301

tests/kernels/attention/test_merge_attn_states.py

+46 -42

@@ -5,7 +5,7 @@ import pytest

5 import torch

7 from vllm._custom_ops import (

8 - merge_attn_states as merge_attn_states_cuda,

8 + merge_attn_states as merge_attn_states_native,

9 )

10 from vllm._custom_ops import (

11 scaled_fp8_quant,

@@ -15,11 +15,13 @@ from vllm.v1.attention.ops.triton_merge_attn_states import (

15 merge_attn_states as merge_attn_states_triton,

16 )

18 +DEVICE = current_platform.device_type

19 +

20 pytestmark = [

21 pytest.mark.skip_global_cleanup,

22 pytest.mark.skipif(

21 - not current_platform.is_cuda_alike(),

22 - reason="merge_attn_states kernels require CUDA or ROCm.",

23 + not (current_platform.is_cuda_alike() or current_platform.is_xpu()),

24 + reason="merge_attn_states kernels require CUDA, ROCm or XPU.",

25 ),

26 ]

@@ -81,20 +83,22 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]

83 all_case_info: list[tuple] = []

84 -@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])

86 +@pytest.mark.parametrize(

87 + "merge_fn", [merge_attn_states_native, merge_attn_states_triton]

88 +)

89 @pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])

90 def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:

91 """When a token is empty on both sides (both LSE -inf), the 0/0 softmax

92 scales must not surface as NaN in the merged output."""

93 num_tokens, num_heads, head_size = 6, 8, 128

94 prefix_output = torch.zeros(

91 - num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype

95 + num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype

96 )

93 - prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")

97 + prefix_lse = torch.randn(num_heads, num_tokens, device=DEVICE)

98 suffix_output = torch.zeros(

95 - num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype

99 + num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype

100 )

97 - suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")

101 + suffix_lse = torch.randn(num_heads, num_tokens, device=DEVICE)

102

103 # Tokens 2 and 3 are empty on both sides.

104 empty = slice(2, 4)

@@ -111,7 +115,7 @@ def generate_markdown_table():

115 global all_case_info

116 table_header = (

117 "| tokens | heads | headsize | dtype "

119 )

120 table_separator = "| --- | --- | --- | --- | --- | --- | --- | --- | --- |"

121

@@ -132,7 +136,7 @@ def generate_markdown_table():

136 device,

137 avg_time_torch_kernel,

138 avg_time_triton_kernel,

135 - avg_time_cuda_kernel,

139 + avg_time_native_kernel,

140 performance_improved,

141 ) = info

142 dtype = shortly_dtype(dtype)

@@ -141,7 +145,7 @@ def generate_markdown_table():

145 f"| {num_tokens} | {num_heads} | {head_size} "

146 f"| {dtype} | {device} | {avg_time_torch_kernel:.5f}ms "

147 f"| {avg_time_triton_kernel:.5f}ms "

144 - f"| {avg_time_cuda_kernel:.5f}ms "

148 + f"| {avg_time_native_kernel:.5f}ms "

149 f"| {performance_improved:.4f}x |"

150 )

151

@@ -171,7 +175,7 @@ def test_merge_attn_states(

175 output_scale = None

176 if use_fp8:

177 output_dtype = current_platform.fp8_dtype()

174 - output_scale = torch.tensor([0.05], dtype=torch.float32, device="cuda")

178 + output_scale = torch.tensor([0.05], dtype=torch.float32, device=DEVICE)

179

180 print(

181 f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, "

@@ -182,8 +186,8 @@ def test_merge_attn_states(

186 )

187

188 # prefix_lse and suffix_lse contain inf and normal values

185 - prefix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device="cuda")

186 - suffix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device="cuda")

189 + prefix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device=DEVICE)

190 + suffix_lse = torch.randn(NUM_HEADS, NUM_TOKENS, dtype=torch.float32, device=DEVICE)

191

192 # Generate boolean masks

193 mask_prefix = torch.rand(NUM_HEADS, NUM_TOKENS) < 0.1

@@ -199,16 +203,16 @@ def test_merge_attn_states(

203 # Other input tensors (need to be initialized but

204 # no actual calculation needed)

205 output = torch.zeros(

202 - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"

206 + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device=DEVICE

207 )

208 output_lse = torch.zeros(

205 - (NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device="cuda"

209 + (NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device=DEVICE

210 )

211 prefix_output = torch.randn(

208 - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda"

212 + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device=DEVICE

213 )

214 suffix_output = torch.randn(

211 - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda"

215 + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device=DEVICE

216 )

217

218 warmup_times = 2

@@ -293,19 +297,19 @@ def test_merge_attn_states(

297

298 avg_time_triton_kernel = total_time_triton_kernel / repeat_times

299

296 - # 2. Run the CUDA kernel

297 - total_time_cuda_kernel = 0

298 - output_cuda = output.clone()

299 - output_lse_cuda = output_lse.clone()

300 + # 2. Run the native (CUDA/XPU) kernel

301 + total_time_native_kernel = 0

302 + output_native = output.clone()

303 + output_lse_native = output_lse.clone()

304

305 for _ in range(warmup_times):

302 - merge_attn_states_cuda(

303 - output_cuda,

306 + merge_attn_states_native(

307 + output_native,

308 prefix_output,

309 prefix_lse,

310 suffix_output,

311 suffix_lse,

308 - output_lse_cuda,

312 + output_lse_native,

313 prefill_tokens_with_context,

314 output_scale,

315 )

@@ -313,28 +317,28 @@ def test_merge_attn_states(

317

318 for _ in range(repeat_times):

319 start.record()

316 - merge_attn_states_cuda(

317 - output_cuda,

320 + merge_attn_states_native(

321 + output_native,

322 prefix_output,

323 prefix_lse,

324 suffix_output,

325 suffix_lse,

322 - output_lse_cuda,

326 + output_lse_native,

327 prefill_tokens_with_context,

328 output_scale,

329 )

330 end.record()

331 torch.accelerator.synchronize()

328 - total_time_cuda_kernel += start.elapsed_time(end)

332 + total_time_native_kernel += start.elapsed_time(end)

333

330 - avg_time_cuda_kernel = total_time_cuda_kernel / repeat_times

334 + avg_time_native_kernel = total_time_native_kernel / repeat_times

335

336 # 3. Performance compare

333 - performance_improved = avg_time_triton_kernel / avg_time_cuda_kernel

337 + performance_improved = avg_time_triton_kernel / avg_time_native_kernel

338 print(f" Torch time: {avg_time_torch_kernel:.6f}ms")

339 print(f"Triton time: {avg_time_triton_kernel:.6f}ms")

340 print(

337 - f" CUDA time: {avg_time_cuda_kernel:.6f}ms, "

341 + f"Native time: {avg_time_native_kernel:.6f}ms, "

342 f"Performance: {performance_improved:.5f}x"

343 )

344 print("-" * 100)

@@ -362,12 +366,12 @@ def test_merge_attn_states(

366 return max_diff

367

368 # Use Triton output as reference because we want to replace

365 - # the Triton kernel with custom CUDA kernel for merge attn

369 + # the Triton kernel with the custom native kernel for merge attn

370 # states operation.

371 output_ref = output_ref_triton

372 output_lse_ref = output_lse_ref_triton

373 torch.testing.assert_close(

370 - output_cuda.float() * scale,

374 + output_native.float() * scale,

375 output_ref.float() * scale,

376 atol=atol,

377 rtol=rtol,

@@ -379,19 +383,19 @@ def test_merge_attn_states(

383 )

384 _diff = diff(output_ref.float() * scale, output_torch.float() * scale)

385 print(f"(Triton vs Torch) : {_diff}")

382 - _diff = diff(output_torch.float() * scale, output_cuda.float() * scale)

383 - print(f" (CUDA vs Torch) : {_diff}")

384 - _diff = diff(output_ref.float() * scale, output_cuda.float() * scale)

385 - print(f" (CUDA vs Triton): {_diff}")

386 + _diff = diff(output_torch.float() * scale, output_native.float() * scale)

387 + print(f"(Native vs Torch) : {_diff}")

388 + _diff = diff(output_ref.float() * scale, output_native.float() * scale)

389 + print(f"(Native vs Triton): {_diff}")

390 print("-" * 100)

391

392 torch.testing.assert_close(

389 - output_lse_cuda.float(), output_lse_ref.float(), atol=atol, rtol=rtol

393 + output_lse_native.float(), output_lse_ref.float(), atol=atol, rtol=rtol

394 )

395 print("Output LSE all match, max abs diff:")

396 print(f"(Triton vs Torch) : {diff(output_lse_torch, output_lse_ref)}")

393 - print(f" (CUDA vs Torch) : {diff(output_lse_torch, output_lse_cuda)}")

394 - print(f" (CUDA vs Triton): {diff(output_lse_ref, output_lse_cuda)}")

397 + print(f"(Native vs Torch) : {diff(output_lse_torch, output_lse_native)}")

398 + print(f"(Native vs Triton): {diff(output_lse_ref, output_lse_native)}")

399 print("-" * 100)

400

401 print(

@@ -410,7 +414,7 @@ def test_merge_attn_states(

414 device,

415 avg_time_torch_kernel,

416 avg_time_triton_kernel,

413 - avg_time_cuda_kernel,

417 + avg_time_native_kernel,

418 performance_improved,

419 )

420 )

tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py

+6 -2

@@ -36,6 +36,7 @@ from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_outpu

36 from vllm.model_executor.layers.mamba.ops.ssd_combined import (

37 mamba_chunk_scan_combined_varlen,

38 )

39 +from vllm.platforms import current_platform

40 from vllm.utils.torch_utils import set_random_seed

41 from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata

@@ -68,7 +69,7 @@ def _run_prefill_decode_equivalence(

69 flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill

70 and decode apply the softplus/bias preprocessing consistently. ``state_dtype``

71 is the recurrent-state precision; ``act_dtype`` the activation/buffer one."""

71 - device = "cuda"

72 + device = current_platform.device_type

73 rtol, atol = _prefill_tolerances(act_dtype)

74 set_random_seed(seed)

@@ -203,7 +204,10 @@ def _run_prefill_decode_equivalence(

204 )

205

206

206 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

207 +@pytest.mark.skipif(

208 + not (current_platform.is_cuda_alike() or current_platform.is_xpu()),

209 + reason="Mamba2 ReplaySSM kernels require a CUDA-alike or XPU device.",

210 +)

211 @pytest.mark.parametrize(

212 "precision",

213 # fp32 state is the default; bf16/fp16 are reduced-footprint configs. fp16

tests/kernels/moe/test_batched_moe.py

+28 -23

@@ -23,6 +23,13 @@ from vllm.platforms import current_platform

23 from vllm.triton_utils import tl

24 from vllm.utils.torch_utils import set_random_seed

26 +DEVICE = current_platform.device_type

27 +

28 +pytestmark = pytest.mark.skipif(

29 + not (current_platform.is_cuda_alike() or current_platform.is_xpu()),

30 + reason="Triton MoE kernels require CUDA/ROCm/XPU.",

31 +)

32 +

33 MNK_FACTORS = [

34 (1, 128, 128),

35 (1, 512, 512),

@@ -71,19 +78,19 @@ class BatchedMMTensors:

78 A = (

79 torch.randn(

80 (config.num_experts, config.max_tokens_per_expert, config.K),

74 - device="cuda",

81 + device=DEVICE,

82 dtype=config.in_dtype,

83 )

84 / 10

85 )

86 B = torch.randn(

87 (config.num_experts, config.N, config.K),

81 - device="cuda",

88 + device=DEVICE,

89 dtype=config.in_dtype,

90 )

91 C = torch.zeros(

92 (config.num_experts, config.max_tokens_per_expert, config.N),

86 - device="cuda",

93 + device=DEVICE,

94 dtype=config.out_dtype,

95 )

@@ -91,7 +98,7 @@ class BatchedMMTensors:

98 low=0,

99 high=config.max_tokens_per_expert,

100 size=(config.num_experts,),

94 - device="cuda",

101 + device=DEVICE,

102 dtype=torch.int32,

103 )

104

@@ -120,8 +127,10 @@ def test_batched_mm(

127

128 use_fp8_w8a8 = dtype == torch.float8_e4m3fn

129

123 - if (dtype == torch.float8_e4m3fn) and not current_platform.has_device_capability(

124 - 89

130 + if (

131 + dtype == torch.float8_e4m3fn

132 + and current_platform.is_cuda_alike()

133 + and not current_platform.has_device_capability(89)

134 ):

135 pytest.skip(

136 "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"

@@ -144,7 +153,7 @@ def test_batched_mm(

153 low=0,

154 high=max_tokens_per_expert,

155 size=(num_experts,),

147 - device="cuda",

156 + device=DEVICE,

157 dtype=torch.int32,

158 )

159

@@ -169,9 +178,9 @@ def test_batched_mm(

178 )

179

180 out_shape = (num_experts, max_tokens_per_expert, N)

172 - test_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda")

173 - ref_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda")

174 - q_ref_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda")

181 + test_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE)

182 + ref_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE)

183 + q_ref_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE)

184

185 compute_tl_dtype = {

186 torch.float16: tl.float16,

@@ -257,8 +266,10 @@ def test_fused_moe_batched_experts(

266

267 use_fp8_w8a8 = dtype == torch.float8_e4m3fn

268

260 - if (dtype == torch.float8_e4m3fn) and not current_platform.has_device_capability(

261 - 89

269 + if (

270 + dtype == torch.float8_e4m3fn

271 + and current_platform.is_cuda_alike()

272 + and not current_platform.has_device_capability(89)

273 ):

274 pytest.skip(

275 "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"

@@ -273,8 +284,8 @@ def test_fused_moe_batched_experts(

284 if per_act_token_quant and block_shape is not None:

285 pytest.skip("Skip illegal quantization test.")

286

276 - a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10

277 - score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)

287 + a = torch.randn((m, k), device=DEVICE, dtype=torch.bfloat16) / 10

288 + score = torch.randn((m, e), device=DEVICE, dtype=torch.bfloat16)

289

290 if dtype.itemsize == 1:

291 act_dtype = torch.bfloat16

@@ -294,8 +305,8 @@ def test_fused_moe_batched_experts(

305 )

306

307 if input_scales and quant_dtype is not None:

297 - a1_scale = torch.tensor(1, device="cuda", dtype=torch.float32)

298 - a2_scale = torch.tensor(1, device="cuda", dtype=torch.float32)

308 + a1_scale = torch.tensor(1, device=DEVICE, dtype=torch.float32)

309 + a2_scale = torch.tensor(1, device=DEVICE, dtype=torch.float32)

310 else:

311 a1_scale = None

312 a2_scale = None

@@ -489,10 +500,7 @@ def test_batched_triton_experts_supports_current_device():

500 BatchedTritonExperts,

501 )

502

492 - if current_platform.is_xpu() or current_platform.is_cuda_alike():

493 - assert BatchedTritonExperts._supports_current_device()

494 - else:

495 - pytest.skip("No GPU device available")

503 + assert BatchedTritonExperts._supports_current_device()

504

505

506 @pytest.mark.parametrize("m,n,k,e,topk", [(32, 512, 512, 8, 2), (45, 1024, 128, 8, 1)])

@@ -500,9 +508,6 @@ def test_batched_experts_end_to_end(m, n, k, e, topk):

508 """End-to-end BatchedTritonExperts via the reference (no-comms)

509 BatchedPrepareAndFinalize, validated against a torch reference. Exercises

510 the device-enablement path."""

503 - if not (current_platform.is_xpu() or current_platform.is_cuda_alike()):

504 - pytest.skip("No GPU device available")

505 -

511 from vllm.v1.worker.workspace import init_workspace_manager

512

513 set_random_seed(7)

tests/kernels/moe/test_count_expert_num_tokens.py

+10 -2

@@ -10,6 +10,14 @@ import pytest

10 import torch

12 from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens

13 +from vllm.platforms import current_platform

14 +

15 +DEVICE = current_platform.device_type

16 +

17 +pytestmark = pytest.mark.skipif(

18 + not (current_platform.is_cuda_alike() or current_platform.is_xpu()),

19 + reason="Triton MoE kernels require CUDA/ROCm/XPU.",

20 +)

23 @dataclasses.dataclass

@@ -88,9 +96,9 @@ def do_test_compute_expert_num_tokens(

96 (num_local_experts), device="cpu", dtype=torch.int32

97 )

98 ref_impl(tt_rank, ref_expert_num_tokens)

91 - ref_expert_num_tokens = ref_expert_num_tokens.to("cuda")

99 + ref_expert_num_tokens = ref_expert_num_tokens.to(DEVICE)

100

93 - tt_rank.to_device("cuda")

101 + tt_rank.to_device(DEVICE)

102 # Test with expert_map

103 triton_expert_num_tokens_w_emap = count_expert_num_tokens(

104 tt_rank.topk_ids, num_local_experts, tt_rank.expert_map

tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py

+4 -3

@@ -13,6 +13,7 @@ from vllm.triton_utils import triton

13 from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used

14 from vllm.utils.torch_utils import set_random_seed

16 +DEVICE = current_platform.device_type

17 FLOAT8_DTYPE = torch.float8_e4m3fn

18 GROUP_SIZE = 128

@@ -61,7 +62,7 @@ def reference_quant(x: torch.Tensor, use_ue8m0: bool):

63 def reference(x: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]:

64 T, N = x.size()

64 - ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device="cuda")

65 + ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device=DEVICE)

66 torch.ops._C.silu_and_mul(ref_act_out, x)

67 return reference_quant(ref_act_out, use_ue8m0)

@@ -93,7 +94,7 @@ def reference_with_clamp(

94 def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int):

95 set_random_seed(42)

96 - input = torch.rand((T, N), dtype=torch.bfloat16, device="cuda")

97 + input = torch.rand((T, N), dtype=torch.bfloat16, device=DEVICE)

99 use_ue8m0 = is_deep_gemm_e8m0_used()

100

@@ -122,7 +123,7 @@ def test_silu_mul_fp8_quant_deep_gemm_clamp(T: int, N: int, clamp_limit: float):

123 # Use a wide distribution so values routinely exceed both clamp limits and

124 # the clamp branch is actually exercised (uniform [0, 1) inputs would never

125 # trigger it).

125 - input = torch.randn((T, N), dtype=torch.bfloat16, device="cuda") * 8.0

126 + input = torch.randn((T, N), dtype=torch.bfloat16, device=DEVICE) * 8.0

127

128 use_ue8m0 = is_deep_gemm_e8m0_used()

129

tests/kernels/moe/test_triton_moe_no_act_mul.py

+15 -13

@@ -23,6 +23,20 @@ from vllm.model_executor.layers.fused_moe.config import (

23 from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts

24 from vllm.platforms import current_platform

26 +DEVICE = current_platform.device_type

27 +

28 +pytestmark = [

29 + pytest.mark.skipif(

30 + not (current_platform.is_cuda_alike() or current_platform.is_xpu()),

31 + reason="Triton MoE kernels require a CUDA-alike or XPU device.",

32 + ),

33 + pytest.mark.skipif(

34 + current_platform.is_cuda_alike()

35 + and not current_platform.has_device_capability(80),

36 + reason="Triton MoE kernels require compute capability >= 8.0 on CUDA/ROCm.",

37 + ),

38 +]

39 +

40 # Test parameters

41 M_SIZES = [1, 16, 64]

42 N_SIZES = [128, 256]

@@ -86,7 +100,7 @@ def make_test_tensors(

100 num_experts: int,

101 topk: int,

102 dtype: torch.dtype = torch.bfloat16,

89 - device: str = "cuda",

103 + device: str = DEVICE,

104 ):

105 """Create test tensors for MoE with non-gated activation.

106

@@ -106,10 +120,6 @@ def make_test_tensors(

120 return hidden_states, w1, w2, topk_weights, topk_ids

121

122

109 -@pytest.mark.skipif(

110 - not current_platform.has_device_capability(80),

111 - reason="Requires compute capability >= 8.0",

112 -)

123 @pytest.mark.parametrize("m", M_SIZES)

124 @pytest.mark.parametrize("n", N_SIZES)

125 @pytest.mark.parametrize("k", K_SIZES)

@@ -192,10 +202,6 @@ def test_triton_experts_no_mul_activation(

202 assert output.abs().sum() > 0, "Output is all zeros"

203

204

195 -@pytest.mark.skipif(

196 - not current_platform.has_device_capability(80),

197 - reason="Requires compute capability >= 8.0",

198 -)

205 @torch.inference_mode()

206 def test_workspace_shapes_no_mul_vs_gated():

207 """Test that workspace shapes differ correctly between gated and non-gated."""

@@ -233,10 +239,6 @@ def test_workspace_shapes_no_mul_vs_gated():

239 assert out_no_mul == out_gated == (M, K)

240

241

236 -@pytest.mark.skipif(

237 - not current_platform.has_device_capability(80),

238 - reason="Requires compute capability >= 8.0",

239 -)

242 @torch.inference_mode()

243 def test_adjust_n_for_activation():

244 """Test the adjust_N_for_activation method."""

tests/kernels/moe/utils.py

+9 -7

@@ -40,6 +40,8 @@ from vllm.platforms import current_platform

40 from vllm.utils.deep_gemm import per_block_cast_to_fp8

41 from vllm.utils.math_utils import round_up

43 +DEVICE = current_platform.device_type

44 +

46 def shuffle_weight(w: torch.Tensor) -> torch.Tensor:

47 """Fold weights to adjacent locations for Triton MoE / SwiGLU kernel layout."""

@@ -80,7 +82,7 @@ def make_dummy_moe_config(

82 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

83 activation=activation,

84 in_dtype=in_dtype,

83 - device="cuda",

85 + device=DEVICE,

86 routing_method=RoutingMethodType.TopK,

87 max_num_tokens=max_num_tokens,

88 )

@@ -238,7 +240,7 @@ def make_quantized_test_activations(

240 block_shape: list[int] | None = None,

241 per_act_token_quant: bool = False,

242 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

241 - a = torch.randn((E, m, k), device="cuda", dtype=in_dtype) / 10

243 + a = torch.randn((E, m, k), device=DEVICE, dtype=in_dtype) / 10

244 a_q = a

245 a_scale = None

246

@@ -369,7 +371,7 @@ def make_test_weight(

371 block_shape: list[int] | None = None,

372 per_out_ch_quant: bool = False,

373 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:

372 - w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15

374 + w_16 = torch.randn((e, rows, cols), device=DEVICE, dtype=in_dtype) / 15

375

376 if quant_dtype is not None:

377 w, w_s, w_gs = moe_quantize_weights(

@@ -449,8 +451,8 @@ def make_test_quant_config(

451 a1_gscale: torch.Tensor | None = None

452 a2_gscale: torch.Tensor | None = None

453 if quant_dtype == "nvfp4":

452 - a1_gscale = torch.ones((e,), device="cuda", dtype=torch.float32)

453 - a2_gscale = torch.ones((e,), device="cuda", dtype=torch.float32)

454 + a1_gscale = torch.ones((e,), device=DEVICE, dtype=torch.float32)

455 + a2_gscale = torch.ones((e,), device=DEVICE, dtype=torch.float32)

456 a1_scale = a1_gscale

457 a2_scale = a2_gscale

458 else:

@@ -551,8 +553,8 @@ def make_naive_shared_experts(

553 K: int,

554 in_dtype: torch.dtype = torch.bfloat16,

555 ) -> torch.nn.Module:

554 - w1 = torch.randn((K, N * 2), device="cuda", dtype=in_dtype) / 15

555 - w2 = torch.randn((N, K), device="cuda", dtype=in_dtype) / 15

556 + w1 = torch.randn((K, N * 2), device=DEVICE, dtype=in_dtype) / 15

557 + w2 = torch.randn((N, K), device=DEVICE, dtype=in_dtype) / 15

558 return TestMLP(w1, w2, out_dtype=in_dtype)

559

560

tests/kernels/quant_utils.py

+2 -1

@@ -13,10 +13,11 @@ from vllm.utils.deep_gemm import _ceil_to_ue8m0, is_deep_gemm_e8m0_used

13 from vllm.utils.math_utils import round_up

15 FP8_DTYPE = current_platform.fp8_dtype()

16 +DEVICE = current_platform.device_type

19 def as_float32_tensor(x: float | torch.Tensor) -> torch.Tensor:

19 - return torch.as_tensor(x, dtype=torch.float32, device="cuda")

20 + return torch.as_tensor(x, dtype=torch.float32, device=DEVICE)

23 def ref_dynamic_per_token_quant(

tests/kernels/quantization/test_per_token_group_quant.py

+8 -4

@@ -18,19 +18,20 @@ from vllm.platforms import current_platform

18 @pytest.mark.parametrize("scale_ue8m0", [False, True])

19 @pytest.mark.parametrize("group_size", [64, 128])

20 @pytest.mark.skipif(

21 - not current_platform.is_cuda_alike(), reason="Only test on CUDA/ROCm."

21 + not (current_platform.is_cuda_alike() or current_platform.is_xpu()),

22 + reason="Only test on CUDA/ROCm/XPU.",

23 )

24 def test_per_token_group_quant_fp8(

25 shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int

26 ):

26 - device = "cuda"

27 + device = current_platform.device_type

29 torch.manual_seed(42)

30 num_tokens, hidden_dim = shape

32 x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8

33 - # cuda path

34 + # native kernel path

35 out_q, scale = fp8_utils.per_token_group_quant_fp8(

36 x,

37 group_size,

@@ -40,7 +41,10 @@ def test_per_token_group_quant_fp8(

41 )

43 # triton ref

43 - with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):

44 + with (

45 + patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False),

46 + patch("vllm.platforms.current_platform.is_xpu", return_value=False),

47 + ):

48 ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(

49 x,

50 group_size,

tests/kernels/test_fused_recurrent_packed_decode.py

+10 -4

@@ -4,13 +4,20 @@

4 import pytest

5 import torch

7 +from vllm.platforms import current_platform

8 from vllm.third_party.flash_linear_attention.ops import (

9 fused_recurrent_gated_delta_rule,

10 fused_recurrent_gated_delta_rule_packed_decode,

11 )

13 +DEVICE = current_platform.device_type

14 +

15 +pytestmark = pytest.mark.skipif(

16 + not (current_platform.is_cuda_alike() or current_platform.is_xpu()),

17 + reason="Gated delta rule Triton kernels require a CUDA-alike or XPU device.",

18 +)

19 +

13 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

21 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])

22 @pytest.mark.parametrize("strided_mixed_qkv", [False, True])

23 def test_fused_recurrent_packed_decode_matches_reference(

@@ -26,7 +33,7 @@ def test_fused_recurrent_packed_decode_matches_reference(

33 V = 128

34 qkv_dim = 2 * (H * K) + (HV * V)

29 - device = torch.device("cuda")

36 + device = torch.device(DEVICE)

38 if strided_mixed_qkv:

39 # Simulate a packed view into a larger projection buffer:

@@ -102,10 +109,9 @@ def test_fused_recurrent_packed_decode_matches_reference(

109 torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)

110

111

105 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

112 def test_packed_decode_supports_large_batch_head_grid():

113 B, H, HV, K, V = 1024, 8, 64, 1, 1

108 - device = torch.device("cuda")

114 + device = torch.device(DEVICE)

115 gates = torch.empty((B, HV), device=device)

116 params = torch.empty((HV,), device=device)

117 out = torch.empty((B, 1, HV, V), device=device)

vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py

+4 -3

@@ -37,9 +37,10 @@ from vllm.triton_utils.allocation import set_triton_allocator

39 def _is_capturing_or_compiling() -> bool:

40 - # torch.cuda.is_current_stream_capturing() is unavailable on non-CUDA (XPU) torch.

41 - return torch.compiler.is_compiling() or (

42 - current_platform.is_cuda_alike() and torch.cuda.is_current_stream_capturing()

40 + # The accelerator API dispatches to the active backend.

41 + return (

42 + torch.compiler.is_compiling()

43 + or torch.accelerator.current_stream().is_capturing()

44 )