From 3b1cc466c01cf46b8b32cc3b1f68494858d1c63e Mon Sep 17 00:00:00 2001 From: Kalyan Kumar Date: Fri, 7 Nov 2025 09:08:35 +0530 Subject: [PATCH] fixes hardcoded "cuda" device references in unit tests to use a dynamic device selection (#12761) --- test/srt/test_triton_attention_kernels.py | 169 +++++++++++----------- 1 file changed, 88 insertions(+), 81 deletions(-) diff --git a/test/srt/test_triton_attention_kernels.py b/test/srt/test_triton_attention_kernels.py index 16c107006..5a7fb5472 100644 --- a/test/srt/test_triton_attention_kernels.py +++ b/test/srt/test_triton_attention_kernels.py @@ -18,6 +18,7 @@ from sglang.srt.layers.attention.triton_ops.extend_attention import ( from sglang.srt.layers.attention.triton_ops.prefill_attention import ( context_attention_fwd, ) +from sglang.srt.utils import get_device from sglang.test.test_utils import CustomTestCase @@ -112,26 +113,27 @@ class TestTritonAttention(CustomTestCase): def _test_extend_attention_once(self, B, N_CTX, H_Q, H_KV, D): dtype = torch.bfloat16 + device = get_device() b_seq_len_prefix = torch.randint( - 1, N_CTX // 2, (B,), dtype=torch.int32, device="cuda" + 1, N_CTX // 2, (B,), dtype=torch.int32, device=device ) b_seq_len_extend = torch.randint( - 1, N_CTX // 2, (B,), dtype=torch.int32, device="cuda" + 1, N_CTX // 2, (B,), dtype=torch.int32, device=device ) b_seq_len = b_seq_len_prefix + b_seq_len_extend max_len_in_batch = torch.max(b_seq_len, 0)[0].item() - b_req_idx = torch.arange(B, dtype=torch.int32, device="cuda") - b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda") + b_req_idx = torch.arange(B, dtype=torch.int32, device=device) + b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device) b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0) - b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device="cuda") + b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device) b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) - kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0) kv_indices = torch.zeros( - (b_seq_len_prefix.sum().item(),), dtype=torch.int32, device="cuda" + (b_seq_len_prefix.sum().item(),), dtype=torch.int32, device=device ) for i in range(B): @@ -142,15 +144,15 @@ class TestTritonAttention(CustomTestCase): total_token_num = torch.sum(b_seq_len).item() extend_token_num = torch.sum(b_seq_len_extend).item() k_buffer = torch.empty( - (total_token_num, H_KV, D), dtype=dtype, device="cuda" + (total_token_num, H_KV, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) v_buffer = torch.empty( - (total_token_num, H_KV, D), dtype=dtype, device="cuda" + (total_token_num, H_KV, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) - k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda") - v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda") - q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda") + k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) for i in range(B): extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] @@ -163,20 +165,20 @@ class TestTritonAttention(CustomTestCase): extend_start_in_buffer:extend_end_in_buffer ] q_extend[extend_start:extend_end] = torch.empty( - (b_seq_len_extend[i], H_Q, D), dtype=dtype, device="cuda" + (b_seq_len_extend[i], H_Q, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) - o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda") + o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) o_extend_mask = torch.empty( - (extend_token_num, H_Q, D), dtype=dtype, device="cuda" + (extend_token_num, H_Q, D), dtype=dtype, device=device ) o_redundant = torch.empty( - (extend_token_num, H_Q, D), dtype=dtype, device="cuda" + (extend_token_num, H_Q, D), dtype=dtype, device=device ) b_seq_len_extend = b_seq_len - b_seq_len_prefix max_len_extend = torch.max(b_seq_len_extend, 0)[0].item() - qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0) custom_mask = None @@ -200,9 +202,9 @@ class TestTritonAttention(CustomTestCase): b_seq_mask_len = b_seq_len_extend * b_seq_len custom_mask = torch.ones( - (b_seq_mask_len.sum().item(),), dtype=torch.bool, device="cuda" + (b_seq_mask_len.sum().item(),), dtype=torch.bool, device=device ) - mask_indptr = torch.zeros((B + 1,), dtype=torch.int64, device="cuda") + mask_indptr = torch.zeros((B + 1,), dtype=torch.int64, device=device) mask_indptr[1 : B + 1] = torch.cumsum(b_seq_mask_len[:B], dim=0) for i in range(B): causal_mask = ( @@ -261,24 +263,25 @@ class TestTritonAttention(CustomTestCase): self, B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE ): dtype = torch.bfloat16 + device = get_device() b_seq_len_prefix = torch.randint( - 1, N_CTX // 2, (B,), dtype=torch.int32, device="cuda" + 1, N_CTX // 2, (B,), dtype=torch.int32, device=device ) b_seq_len_extend = torch.randint( - 1, N_CTX // 2, (B,), dtype=torch.int32, device="cuda" + 1, N_CTX // 2, (B,), dtype=torch.int32, device=device ) b_seq_len = b_seq_len_prefix + b_seq_len_extend - b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda") + b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device) b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0) - b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device="cuda") + b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device) b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) - kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0) kv_indices = torch.zeros( - (b_seq_len_prefix.sum().item(),), dtype=torch.int32, device="cuda" + (b_seq_len_prefix.sum().item(),), dtype=torch.int32, device=device ) for i in range(B): @@ -289,15 +292,15 @@ class TestTritonAttention(CustomTestCase): total_token_num = torch.sum(b_seq_len).item() extend_token_num = torch.sum(b_seq_len_extend).item() k_buffer = torch.empty( - (total_token_num, H_KV, D), dtype=dtype, device="cuda" + (total_token_num, H_KV, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) v_buffer = torch.empty( - (total_token_num, H_KV, D), dtype=dtype, device="cuda" + (total_token_num, H_KV, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) - k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda") - v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda") - q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda") + k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) for i in range(B): extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] @@ -310,19 +313,19 @@ class TestTritonAttention(CustomTestCase): extend_start_in_buffer:extend_end_in_buffer ] q_extend[extend_start:extend_end] = torch.empty( - (b_seq_len_extend[i], H_Q, D), dtype=dtype, device="cuda" + (b_seq_len_extend[i], H_Q, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) o_extend_triton = torch.empty( - (extend_token_num, H_Q, D), dtype=dtype, device="cuda" + (extend_token_num, H_Q, D), dtype=dtype, device=device ) o_extend_torch = torch.empty( - (extend_token_num, H_Q, D), dtype=dtype, device="cuda" + (extend_token_num, H_Q, D), dtype=dtype, device=device ) b_seq_len_extend = b_seq_len - b_seq_len_prefix max_len_extend = torch.max(b_seq_len_extend, 0)[0].item() - qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0) extend_attention_fwd( @@ -368,19 +371,20 @@ class TestTritonAttention(CustomTestCase): def _test_context_attention_once(self, head_dim, is_causal): # Set up a simple test case + device = get_device() num_heads = 4 seq_lens = [8, 12] max_seq_len = max(seq_lens) # Create random input tensors - q = torch.randn(sum(seq_lens), num_heads, head_dim, device="cuda") - k = torch.randn(sum(seq_lens), num_heads, head_dim, device="cuda") - v = torch.randn(sum(seq_lens), num_heads, head_dim, device="cuda") - o = torch.zeros(sum(seq_lens), num_heads, head_dim, device="cuda") + q = torch.randn(sum(seq_lens), num_heads, head_dim, device=device) + k = torch.randn(sum(seq_lens), num_heads, head_dim, device=device) + v = torch.randn(sum(seq_lens), num_heads, head_dim, device=device) + o = torch.zeros(sum(seq_lens), num_heads, head_dim, device=device) # Create b_start_loc and b_seq_len tensors - b_start_loc = torch.tensor([0, seq_lens[0]], device="cuda") - b_seq_len = torch.tensor(seq_lens, device="cuda") + b_start_loc = torch.tensor([0, seq_lens[0]], device=device) + b_seq_len = torch.tensor(seq_lens, device=device) context_attention_fwd( q, k, v, o, b_start_loc, b_seq_len, max_seq_len, is_causal=is_causal @@ -413,38 +417,39 @@ class TestTritonAttention(CustomTestCase): self._test_context_attention_once(dim, is_causal) def _test_decode_attention_once(self, B, H_Q, H_KV, D): + device = get_device() dtype = torch.bfloat16 seq_len = 10 # This represents the number of tokens already in the sequence total_tokens = B * seq_len sm_scale = 1.0 / (D**0.5) max_kv_splits = 8 - num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device="cuda") + num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device=device) # q represents the new token being generated, one per batch - q = torch.randn(B, H_Q, D, dtype=dtype, device="cuda") + q = torch.randn(B, H_Q, D, dtype=dtype, device=device) # k_buffer and v_buffer represent all previous tokens - k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") - v_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") + k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device) + v_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device) # o will have the same shape as q - o = torch.zeros(B, H_Q, D, dtype=dtype, device="cuda") + o = torch.zeros(B, H_Q, D, dtype=dtype, device=device) - b_seq_len = torch.full((B,), seq_len, device="cuda") + b_seq_len = torch.full((B,), seq_len, device=device) - kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len[:B], dim=0) - kv_indices = torch.arange(total_tokens, device="cuda") + kv_indices = torch.arange(total_tokens, device=device) attn_logits = torch.empty( (B, H_Q, max_kv_splits, D), dtype=torch.float32, - device="cuda", + device=device, ) attn_lse = torch.empty( (B, H_Q, max_kv_splits), dtype=torch.float32, - device="cuda", + device=device, ) decode_attention_fwd( @@ -478,38 +483,39 @@ class TestTritonAttention(CustomTestCase): def _test_grouped_decode_attention_once(self, B, S, H_Q, H_KV, D, D_V): dtype = torch.bfloat16 + device = get_device() seq_len = S # This represents the number of tokens already in the sequence total_tokens = B * seq_len sm_scale = 1.0 / (D**0.5) max_kv_splits = 8 - num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device="cuda") + num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device=device) # q represents the new token being generated, one per batch - q = torch.randn(B, H_Q, D, dtype=dtype, device="cuda") + q = torch.randn(B, H_Q, D, dtype=dtype, device=device) # k_buffer and v_buffer represent all previous tokens - k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") - v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device="cuda") + k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device) + v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device=device) # o will have the same shape as q - o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda") - o_grouped = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda") + o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device) + o_grouped = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device) - b_seq_len = torch.full((B,), seq_len, device="cuda") + b_seq_len = torch.full((B,), seq_len, device=device) - kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len[:B], dim=0) - kv_indices = torch.arange(total_tokens, device="cuda") + kv_indices = torch.arange(total_tokens, device=device) attn_logits = torch.empty( (B, H_Q, max_kv_splits, D_V), dtype=torch.float32, - device="cuda", + device=device, ) attn_lse = torch.empty( (B, H_Q, max_kv_splits), dtype=torch.float32, - device="cuda", + device=device, ) decode_attention_fwd_normal( @@ -529,12 +535,12 @@ class TestTritonAttention(CustomTestCase): attn_logits1 = torch.empty( (B, H_Q, max_kv_splits, D_V), dtype=torch.float32, - device="cuda", + device=device, ) attn_lse1 = torch.empty( (B, H_Q, max_kv_splits, D_V), dtype=torch.float32, - device="cuda", + device=device, ) decode_attention_fwd_grouped( @@ -576,25 +582,26 @@ class TestTritonAttention(CustomTestCase): def _test_extend_attention_unified_vs_regular_once(self, B, N_CTX, H_Q, H_KV, D): """Test that unified kernel produces same results as 2-stage kernel.""" dtype = torch.bfloat16 + device = get_device() b_seq_len_prefix = torch.randint( - 1, N_CTX // 2, (B,), dtype=torch.int32, device="cuda" + 1, N_CTX // 2, (B,), dtype=torch.int32, device=device ) b_seq_len_extend = torch.randint( - 1, N_CTX // 2, (B,), dtype=torch.int32, device="cuda" + 1, N_CTX // 2, (B,), dtype=torch.int32, device=device ) b_seq_len = b_seq_len_prefix + b_seq_len_extend - b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda") + b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device) b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0) - b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device="cuda") + b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device) b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) # Setup prefix KV indices - kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0) kv_indices = torch.zeros( - (b_seq_len_prefix.sum().item(),), dtype=torch.int64, device="cuda" + (b_seq_len_prefix.sum().item(),), dtype=torch.int64, device=device ) for i in range(B): @@ -605,15 +612,15 @@ class TestTritonAttention(CustomTestCase): total_token_num = torch.sum(b_seq_len).item() extend_token_num = torch.sum(b_seq_len_extend).item() k_buffer = torch.empty( - (total_token_num, H_KV, D), dtype=dtype, device="cuda" + (total_token_num, H_KV, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) v_buffer = torch.empty( - (total_token_num, H_KV, D), dtype=dtype, device="cuda" + (total_token_num, H_KV, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) - k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda") - v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda") - q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda") + k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) for i in range(B): extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] @@ -627,16 +634,16 @@ class TestTritonAttention(CustomTestCase): extend_start_in_buffer:extend_end_in_buffer ] q_extend[extend_start:extend_end] = torch.empty( - (b_seq_len_extend[i], H_Q, D), dtype=dtype, device="cuda" + (b_seq_len_extend[i], H_Q, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) # Setup for extend attention max_len_extend = torch.max(b_seq_len_extend, 0)[0].item() - qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") + qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0) # Run 2-stage kernel - o_regular = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda") + o_regular = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) extend_attention_fwd( q_extend, k_extend, @@ -658,9 +665,9 @@ class TestTritonAttention(CustomTestCase): total_token_num - extend_token_num, total_token_num, dtype=torch.int64, - device="cuda", + device=device, ) - extend_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda") + extend_start_loc = torch.zeros((B,), dtype=torch.int32, device=device) extend_start_loc[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) unified_kv_indptr, unified_kv_indices, prefix_lens = build_unified_kv_indices( @@ -673,7 +680,7 @@ class TestTritonAttention(CustomTestCase): ) # Run unified kernel - o_unified = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda") + o_unified = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) extend_attention_fwd_unified( q_extend, o_unified, @@ -716,7 +723,7 @@ class TestTritonAttention(CustomTestCase): """Test build_unified_kv_indices correctness.""" B = 4 dtype = torch.int64 - device = "cuda" + device = get_device() # Setup test data prefix_lens = torch.tensor([10, 20, 15, 25], dtype=torch.int32, device=device)