diff --git a/python/sglang/srt/layers/attention/mamba/ops/layernorm_gated.py b/python/sglang/srt/layers/attention/mamba/ops/layernorm_gated.py index 88b27eb5d..524e8d399 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/layernorm_gated.py +++ b/python/sglang/srt/layers/attention/mamba/ops/layernorm_gated.py @@ -119,7 +119,7 @@ def _layer_norm_fwd( # heuristics for number of warps num_warps = min(max(BLOCK_N // 256, 1), 8) grid = (M, ngroups) - with torch.cuda.device(x.device.index): + with torch.get_device_module(x.device).device(x.device.index): _layer_norm_fwd_1pass_kernel[grid]( x, out, diff --git a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py index f238d51b4..c89a4f86b 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +++ b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py @@ -427,7 +427,7 @@ def selective_state_update( else (0, 0) ) - with torch.cuda.device(x.device.index): + with torch.get_device_module(x.device).device(x.device.index): _selective_scan_update_kernel[grid]( state, x, diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssd_bmm.py b/python/sglang/srt/layers/attention/mamba/ops/ssd_bmm.py index 667d34afa..ae9803a83 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/ssd_bmm.py +++ b/python/sglang/srt/layers/attention/mamba/ops/ssd_bmm.py @@ -179,7 +179,7 @@ def _bmm_chunk_fwd(a, b, chunk_size, seq_idx=None, causal=False, output_dtype=No batch, nchunks if not has_groups else nchunks * ngroups, ) - with torch.cuda.device(a.device.index): + with torch.get_device_module(a.device).device(a.device.index): _bmm_chunk_fwd_kernel[grid]( a, b, diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py b/python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py index 2dd583800..162d859d4 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py +++ b/python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py @@ -460,7 +460,7 @@ def _chunk_cumsum_fwd( nchunks, triton.cdiv(nheads, META["BLOCK_SIZE_H"]), ) - with torch.cuda.device(dt.device.index): + with torch.get_device_module(dt.device).device(dt.device.index): _chunk_cumsum_fwd_kernel[grid_chunk_cs]( dt, A, @@ -520,7 +520,7 @@ def _chunk_state_fwd( batch * nchunks, nheads, ) - with torch.cuda.device(x.device.index): + with torch.get_device_module(x.device).device(x.device.index): _chunk_state_fwd_kernel[grid]( x, B, @@ -596,7 +596,7 @@ def chunk_state_varlen( batch, nheads, ) - with torch.cuda.device(x.device.index): + with torch.get_device_module(x.device).device(x.device.index): _chunk_state_varlen_kernel[grid]( x, B, diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py b/python/sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py index 5e8c32385..d448a1d5c 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py +++ b/python/sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py @@ -214,7 +214,7 @@ def _state_passing_fwd( (batch, nheads, dim), device=states.device, dtype=torch.float32 ) grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE"]), batch, nheads) - with torch.cuda.device(states.device.index): + with torch.get_device_module(states.device).device(states.device.index): _state_passing_fwd_kernel[grid]( states, out, diff --git a/test/registered/layers/mamba/test_mamba2_mixer.py b/test/registered/layers/mamba/test_mamba2_mixer.py index 9ff5060ba..41519f3be 100644 --- a/test/registered/layers/mamba/test_mamba2_mixer.py +++ b/test/registered/layers/mamba/test_mamba2_mixer.py @@ -13,6 +13,7 @@ from sglang.srt.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, ) +from sglang.srt.utils import get_device, get_device_count from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=32, suite="stage-b-test-2-gpu-large") @@ -35,14 +36,14 @@ def test_mixer2_gated_norm_multi_gpu( seq_len: int, hidden_size_n_groups: tuple[int, int], dtype: torch.dtype, - device: str = "cuda", + device: str = get_device(), ): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") assert ( - torch.cuda.device_count() >= NUM_GPUS - ), f"This test requires at least {NUM_GPUS} GPUs, but only {torch.cuda.device_count()} available" + get_device_count() >= NUM_GPUS + ), f"This test requires at least {NUM_GPUS} GPUs, but only {get_device_count()} available" hidden_size, n_groups = hidden_size_n_groups num_processes = NUM_GPUS @@ -79,8 +80,8 @@ def mixer2_gated_norm_tensor_parallel( ): torch.manual_seed(0) - device = torch.device(f"cuda:{local_rank}") - torch.cuda.set_device(device) + device = torch.device(get_device(local_rank)) + torch.get_device_module(device).set_device(device) torch.set_default_device(device) torch.set_default_dtype(dtype) diff --git a/test/registered/layers/mamba/test_mamba_ssm.py b/test/registered/layers/mamba/test_mamba_ssm.py index 6bd7d5734..8af1705b2 100644 --- a/test/registered/layers/mamba/test_mamba_ssm.py +++ b/test/registered/layers/mamba/test_mamba_ssm.py @@ -13,6 +13,7 @@ from einops import rearrange, repeat from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID from sglang.srt.layers.attention.mamba.ops import selective_state_update +from sglang.srt.utils import get_device def selective_state_update_ref( @@ -92,10 +93,9 @@ def selective_state_update_ref( @pytest.mark.parametrize("dstate", [16, 32, 64]) @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) def test_selective_state_update(dim, dstate, has_z, itype): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") - - device = "cuda" + device = get_device() + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2) if itype == torch.bfloat16: @@ -136,10 +136,9 @@ def test_selective_state_update(dim, dstate, has_z, itype): def test_selective_state_update_with_batch_indices( with_padding, dim, dstate, has_z, itype ): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") - - device = "cuda" + device = get_device() + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2) if itype == torch.bfloat16: rtol, atol = 1e-1, 1e-1 @@ -229,10 +228,9 @@ def test_selective_state_update_with_batch_indices( def test_selective_state_update_with_heads_with_batch_indices( dim, dstate, ngroups, has_z, tie_hdim, itype ): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") - - device = "cuda" + device = get_device() + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 3e-2) if itype == torch.bfloat16: rtol, atol = 1e-1, 1e-1 diff --git a/test/registered/layers/mamba/test_mamba_ssm_ssd.py b/test/registered/layers/mamba/test_mamba_ssm_ssd.py index 3b9670d7e..ec1b5c2a1 100644 --- a/test/registered/layers/mamba/test_mamba_ssm_ssd.py +++ b/test/registered/layers/mamba/test_mamba_ssm_ssd.py @@ -14,6 +14,7 @@ from einops import rearrange, repeat from sglang.srt.layers.attention.mamba.mamba2_metadata import Mamba2Metadata from sglang.srt.layers.attention.mamba.ops import mamba_chunk_scan_combined +from sglang.srt.utils import get_device from sglang.srt.utils.common import is_hip from sglang.utils import is_in_ci @@ -99,10 +100,12 @@ def ssd_minimal_discrete( return Y, final_state -def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device="cuda"): +def generate_random_inputs(batch_size, seqlen, n_heads, d_head, itype, device=None): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") + if device is None: + device = get_device() + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") torch.manual_seed(0) A = -torch.exp(torch.rand(n_heads, dtype=itype, device=device)) @@ -125,7 +128,7 @@ def generate_continuous_batched_examples( n_heads, d_head, itype, - device="cuda", + device=None, return_naive_ref=True, ): @@ -138,8 +141,10 @@ def generate_continuous_batched_examples( # generate the full-length example A, dt, X, B, C = generate_random_inputs( - num_examples, full_length, n_heads, d_head, itype + num_examples, full_length, n_heads, d_head, itype, device ) + # Capture the resolved device from the tensors + device = X.device if return_naive_ref: Y_min, final_state_min = ssd_minimal_discrete( @@ -227,8 +232,9 @@ if is_in_ci(): @pytest.mark.parametrize("d_head", SINGLE_DHEAD) @pytest.mark.parametrize("seq_len_chunk_size", SINGLE_SEQ_LEN_CHUNK_SIZE) def test_mamba_chunk_scan_single_example(d_head, n_heads, seq_len_chunk_size, itype): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") + device = get_device() + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") # this tests the kernels on a single example (no batching) @@ -319,8 +325,9 @@ if is_in_ci(): ], ) def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases, itype): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") + device = get_device() + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") # this test with multiple examples in a continuous batch # (i.e. chunked prefill) @@ -398,8 +405,9 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases, ], ) def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") + device = get_device() + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") # This test verifies the correctness of the chunked prefill implementation # in the mamba2 ssd kernels, by comparing concatenation (in the sequence @@ -632,8 +640,9 @@ def test_mamba_chunk_scan_intermediate_states( seq_len_chunk_size, itype, ): - if not torch.cuda.is_available(): - pytest.skip("CUDA device not available") + device = get_device() + if device not in ["cuda", "xpu"]: + pytest.skip("Test only supports CUDA and XPU devices") if itype == torch.bfloat16: atol, rtol = 5e-2, 5e-2