diff --git a/python/sglang/srt/debug_utils/dump_comparator.py b/python/sglang/srt/debug_utils/dump_comparator.py index 6f5a3397d..e052ef7a1 100644 --- a/python/sglang/srt/debug_utils/dump_comparator.py +++ b/python/sglang/srt/debug_utils/dump_comparator.py @@ -15,6 +15,7 @@ from typing import Callable, List, Optional import torch from sglang.srt.debug_utils.dumper import get_truncated_value +from sglang.srt.utils import get_device def main(args): @@ -259,7 +260,7 @@ def _load_object(path): if not isinstance(x, torch.Tensor): print(f"Skip load {path} since {type(x)=} is not a Tensor ({x=})") return None - return x.cuda() + return x.to(get_device()) def _comparison_preprocessor(x_baseline, x_target, name): diff --git a/test/registered/debug_utils/test_tensor_dump_forward_hook.py b/test/registered/debug_utils/test_tensor_dump_forward_hook.py index 4ebbfd17d..00dee4adf 100644 --- a/test/registered/debug_utils/test_tensor_dump_forward_hook.py +++ b/test/registered/debug_utils/test_tensor_dump_forward_hook.py @@ -7,6 +7,7 @@ from sglang.srt.debug_utils.tensor_dump_forward_hook import ( register_forward_hook_for_model, ) from sglang.srt.distributed.parallel_state import ( + get_default_distributed_backend, init_distributed_environment, initialize_model_parallel, ) @@ -14,7 +15,7 @@ from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import LinearBase from sglang.srt.models.qwen2 import Qwen2MLP from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler -from sglang.srt.utils import add_prefix +from sglang.srt.utils import add_prefix, get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( @@ -78,8 +79,10 @@ def init_weights(module): def test_model_forward_dump(tmp_path): set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) + device = get_device() + backend = get_default_distributed_backend(device) init_distributed_environment( - backend="nccl", + backend=backend, world_size=1, rank=0, local_rank=0, @@ -88,14 +91,14 @@ def test_model_forward_dump(tmp_path): initialize_model_parallel() model = MockCausalLM() model.apply(init_weights) - model = model.cuda().bfloat16() + model = model.to(device=device, dtype=torch.bfloat16) dumper = register_forward_hook_for_model( model, tmp_path / "sglang_dump", [0], 0, 0, 0 ) dir_path = dumper.get_dump_dir() inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01 - result = model(inp.cuda()) + result = model(inp.to(device)) data = torch.load(f"{dir_path}/Pass00000.pt") assert "model.layernorm" in data assert "model.mlp.down_proj" in data diff --git a/test/registered/kernels/test_fused_topk_deepseek.py b/test/registered/kernels/test_fused_topk_deepseek.py index beed115cc..ea7347a3c 100644 --- a/test/registered/kernels/test_fused_topk_deepseek.py +++ b/test/registered/kernels/test_fused_topk_deepseek.py @@ -4,6 +4,7 @@ import pytest import torch from sglang.srt.layers.moe.topk import biased_grouped_topk_gpu, biased_grouped_topk_impl +from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=2, suite="nightly-1-gpu", nightly=True) @@ -29,11 +30,12 @@ def test_fused_topk_deepseek(seq_length, params, apply_routed_scaling_factor_on_ """ num_experts, num_expert_group, topk_group, topk = params dtype = torch.float32 + device = get_device() torch.manual_seed(seq_length) - hidden_states = torch.randn(seq_length, 128, dtype=dtype, device="cuda") - gating_output = torch.randn(seq_length, num_experts, dtype=dtype, device="cuda") - correction_bias = torch.randn(num_experts, dtype=dtype, device="cuda") + hidden_states = torch.randn(seq_length, 128, dtype=dtype, device=device) + gating_output = torch.randn(seq_length, num_experts, dtype=dtype, device=device) + correction_bias = torch.randn(num_experts, dtype=dtype, device=device) routed_scaling_factor = 2.5 if apply_routed_scaling_factor_on_output else None @@ -71,8 +73,8 @@ def test_fused_topk_deepseek(seq_length, params, apply_routed_scaling_factor_on_ sum_check = torch.allclose(output_sum, ref_output_sum, rtol=1e-03, atol=1e-04) # Check 2: Scatter-based comparison with allowance for tie-breaking - res = torch.zeros(seq_length, num_experts, dtype=torch.float32, device="cuda") - ref = torch.zeros(seq_length, num_experts, dtype=torch.float32, device="cuda") + res = torch.zeros(seq_length, num_experts, dtype=torch.float32, device=device) + ref = torch.zeros(seq_length, num_experts, dtype=torch.float32, device=device) res.scatter_(1, indices.long(), output) ref.scatter_(1, ref_indices.long(), ref_output) diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 47b565678..5bbe92b5f 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -39,6 +39,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( ) from sglang.srt.mem_cache.mamba_radix_cache import TreeNode as MambaTreeNode from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode +from sglang.srt.utils import get_device # Test constants DEFAULT_PAGE_SIZE = 4 @@ -774,7 +775,7 @@ class TestRadixCache(unittest.TestCase): base_prefix_len = 10000 suffix_len = 100 - torch_allocated_before = torch.cuda.memory_allocated() + torch_allocated_before = torch.get_device_module().memory_allocated() # build dataset with common prefix common_prefix = [ @@ -784,7 +785,7 @@ class TestRadixCache(unittest.TestCase): suffix = [random.randint(1, vocab_size - 1) for _ in range(suffix_len)] seq = common_prefix + suffix keys.append(seq) - values.append(torch.zeros(len(seq), device="cuda", dtype=torch.int32)) + values.append(torch.zeros(len(seq), device=get_device(), dtype=torch.int32)) cache: RadixCache = RadixCache.create_simulated() @@ -793,7 +794,9 @@ class TestRadixCache(unittest.TestCase): del values - torch_allocated = torch.cuda.memory_allocated() - torch_allocated_before + torch_allocated = ( + torch.get_device_module().memory_allocated() - torch_allocated_before + ) cache_size_bytes = cache.total_size() * 4 print(f"\nCache size (MB): {cache_size_bytes / (1024 * 1024)}") print(f"Torch allocated (MB): {torch_allocated / (1024 * 1024)}")