minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fixes hardcoded "cuda" device references in unit tests to use a dynamic device selection (#12761)

Browse Source
This commit is contained in:
Kalyan Kumar
2025-11-07 11:38:35 +08:00
committed by GitHub
parent 0ee5ab5a74
commit 3b1cc466c0
1 changed files with 88 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/test_triton_attention_kernels.py
+88 -81
View File
@@ -18,6 +18,7 @@ from sglang.srt.layers.attention.triton_ops.extend_attention import (
from sglang.srt.layers.attention.triton_ops.prefill_attention import ( from sglang.srt.layers.attention.triton_ops.prefill_attention import (
context_attention_fwd, context_attention_fwd,
) )
from sglang.srt.utils import get_device
from sglang.test.test_utils import CustomTestCase 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): def _test_extend_attention_once(self, B, N_CTX, H_Q, H_KV, D):
dtype = torch.bfloat16 dtype = torch.bfloat16
device = get_device()
b_seq_len_prefix = torch.randint( 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( 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_seq_len = b_seq_len_prefix + b_seq_len_extend
max_len_in_batch = torch.max(b_seq_len, 0)[0].item() max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
b_req_idx = torch.arange(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="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[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) 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_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
kv_indices = torch.zeros( 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): for i in range(B):
@@ -142,15 +144,15 @@ class TestTritonAttention(CustomTestCase):
total_token_num = torch.sum(b_seq_len).item() total_token_num = torch.sum(b_seq_len).item()
extend_token_num = torch.sum(b_seq_len_extend).item() extend_token_num = torch.sum(b_seq_len_extend).item()
k_buffer = torch.empty( 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) ).normal_(mean=0.1, std=0.2)
v_buffer = torch.empty( 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) ).normal_(mean=0.1, std=0.2)
k_extend = torch.empty((extend_token_num, H_KV, 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="cuda") 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="cuda") q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
for i in range(B): for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] 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] 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 extend_start_in_buffer:extend_end_in_buffer
] ]
q_extend[extend_start:extend_end] = torch.empty( 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) ).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( 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( 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 b_seq_len_extend = b_seq_len - b_seq_len_prefix
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item() 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) qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
custom_mask = None custom_mask = None
@@ -200,9 +202,9 @@ class TestTritonAttention(CustomTestCase):
b_seq_mask_len = b_seq_len_extend * b_seq_len b_seq_mask_len = b_seq_len_extend * b_seq_len
custom_mask = torch.ones( 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) mask_indptr[1 : B + 1] = torch.cumsum(b_seq_mask_len[:B], dim=0)
for i in range(B): for i in range(B):
causal_mask = ( causal_mask = (
@@ -261,24 +263,25 @@ class TestTritonAttention(CustomTestCase):
self, B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE self, B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE
): ):
dtype = torch.bfloat16 dtype = torch.bfloat16
device = get_device()
b_seq_len_prefix = torch.randint( 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( 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_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[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) 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_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
kv_indices = torch.zeros( 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): for i in range(B):
@@ -289,15 +292,15 @@ class TestTritonAttention(CustomTestCase):
total_token_num = torch.sum(b_seq_len).item() total_token_num = torch.sum(b_seq_len).item()
extend_token_num = torch.sum(b_seq_len_extend).item() extend_token_num = torch.sum(b_seq_len_extend).item()
k_buffer = torch.empty( 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) ).normal_(mean=0.1, std=0.2)
v_buffer = torch.empty( 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) ).normal_(mean=0.1, std=0.2)
k_extend = torch.empty((extend_token_num, H_KV, 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="cuda") 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="cuda") q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
for i in range(B): for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] 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] 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 extend_start_in_buffer:extend_end_in_buffer
] ]
q_extend[extend_start:extend_end] = torch.empty( 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) ).normal_(mean=0.1, std=0.2)
o_extend_triton = torch.empty( 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( 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 b_seq_len_extend = b_seq_len - b_seq_len_prefix
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item() 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) qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
extend_attention_fwd( extend_attention_fwd(
@@ -368,19 +371,20 @@ class TestTritonAttention(CustomTestCase):
def _test_context_attention_once(self, head_dim, is_causal): def _test_context_attention_once(self, head_dim, is_causal):
# Set up a simple test case # Set up a simple test case
device = get_device()
num_heads = 4 num_heads = 4
seq_lens = [8, 12] seq_lens = [8, 12]
max_seq_len = max(seq_lens) max_seq_len = max(seq_lens)
# Create random input tensors # Create random input tensors
q = torch.randn(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="cuda") k = torch.randn(sum(seq_lens), num_heads, head_dim, device=device)
v = torch.randn(sum(seq_lens), num_heads, head_dim, device="cuda") v = torch.randn(sum(seq_lens), num_heads, head_dim, device=device)
o = torch.zeros(sum(seq_lens), num_heads, head_dim, device="cuda") o = torch.zeros(sum(seq_lens), num_heads, head_dim, device=device)
# Create b_start_loc and b_seq_len tensors # Create b_start_loc and b_seq_len tensors
b_start_loc = torch.tensor([0, seq_lens[0]], device="cuda") b_start_loc = torch.tensor([0, seq_lens[0]], device=device)
b_seq_len = torch.tensor(seq_lens, device="cuda") b_seq_len = torch.tensor(seq_lens, device=device)
context_attention_fwd( context_attention_fwd(
q, k, v, o, b_start_loc, b_seq_len, max_seq_len, is_causal=is_causal 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) self._test_context_attention_once(dim, is_causal)
def _test_decode_attention_once(self, B, H_Q, H_KV, D): def _test_decode_attention_once(self, B, H_Q, H_KV, D):
device = get_device()
dtype = torch.bfloat16 dtype = torch.bfloat16
seq_len = 10 # This represents the number of tokens already in the sequence seq_len = 10 # This represents the number of tokens already in the sequence
total_tokens = B * seq_len total_tokens = B * seq_len
sm_scale = 1.0 / (D**0.5) sm_scale = 1.0 / (D**0.5)
max_kv_splits = 8 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 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 and v_buffer represent all previous tokens
k_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="cuda") v_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
# o will have the same shape as q # 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_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( attn_logits = torch.empty(
(B, H_Q, max_kv_splits, D), (B, H_Q, max_kv_splits, D),
dtype=torch.float32, dtype=torch.float32,
device="cuda", device=device,
) )
attn_lse = torch.empty( attn_lse = torch.empty(
(B, H_Q, max_kv_splits), (B, H_Q, max_kv_splits),
dtype=torch.float32, dtype=torch.float32,
device="cuda", device=device,
) )
decode_attention_fwd( 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): def _test_grouped_decode_attention_once(self, B, S, H_Q, H_KV, D, D_V):
dtype = torch.bfloat16 dtype = torch.bfloat16
device = get_device()
seq_len = S # This represents the number of tokens already in the sequence seq_len = S # This represents the number of tokens already in the sequence
total_tokens = B * seq_len total_tokens = B * seq_len
sm_scale = 1.0 / (D**0.5) sm_scale = 1.0 / (D**0.5)
max_kv_splits = 8 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 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 and v_buffer represent all previous tokens
k_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_V, dtype=dtype, device="cuda") v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device=device)
# o will have the same shape as q # o will have the same shape as q
o = 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="cuda") 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_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( attn_logits = torch.empty(
(B, H_Q, max_kv_splits, D_V), (B, H_Q, max_kv_splits, D_V),
dtype=torch.float32, dtype=torch.float32,
device="cuda", device=device,
) )
attn_lse = torch.empty( attn_lse = torch.empty(
(B, H_Q, max_kv_splits), (B, H_Q, max_kv_splits),
dtype=torch.float32, dtype=torch.float32,
device="cuda", device=device,
) )
decode_attention_fwd_normal( decode_attention_fwd_normal(
@@ -529,12 +535,12 @@ class TestTritonAttention(CustomTestCase):
attn_logits1 = torch.empty( attn_logits1 = torch.empty(
(B, H_Q, max_kv_splits, D_V), (B, H_Q, max_kv_splits, D_V),
dtype=torch.float32, dtype=torch.float32,
device="cuda", device=device,
) )
attn_lse1 = torch.empty( attn_lse1 = torch.empty(
(B, H_Q, max_kv_splits, D_V), (B, H_Q, max_kv_splits, D_V),
dtype=torch.float32, dtype=torch.float32,
device="cuda", device=device,
) )
decode_attention_fwd_grouped( 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): 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.""" """Test that unified kernel produces same results as 2-stage kernel."""
dtype = torch.bfloat16 dtype = torch.bfloat16
device = get_device()
b_seq_len_prefix = torch.randint( 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( 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_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[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) b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
# Setup prefix KV indices # 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_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
kv_indices = torch.zeros( 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): for i in range(B):
@@ -605,15 +612,15 @@ class TestTritonAttention(CustomTestCase):
total_token_num = torch.sum(b_seq_len).item() total_token_num = torch.sum(b_seq_len).item()
extend_token_num = torch.sum(b_seq_len_extend).item() extend_token_num = torch.sum(b_seq_len_extend).item()
k_buffer = torch.empty( 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) ).normal_(mean=0.1, std=0.2)
v_buffer = torch.empty( 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) ).normal_(mean=0.1, std=0.2)
k_extend = torch.empty((extend_token_num, H_KV, 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="cuda") 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="cuda") q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
for i in range(B): for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] 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 extend_start_in_buffer:extend_end_in_buffer
] ]
q_extend[extend_start:extend_end] = torch.empty( 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) ).normal_(mean=0.1, std=0.2)
# Setup for extend attention # Setup for extend attention
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item() 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) qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
# Run 2-stage kernel # 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( extend_attention_fwd(
q_extend, q_extend,
k_extend, k_extend,
@@ -658,9 +665,9 @@ class TestTritonAttention(CustomTestCase):
total_token_num - extend_token_num, total_token_num - extend_token_num,
total_token_num, total_token_num,
dtype=torch.int64, 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) extend_start_loc[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
unified_kv_indptr, unified_kv_indices, prefix_lens = build_unified_kv_indices( unified_kv_indptr, unified_kv_indices, prefix_lens = build_unified_kv_indices(
@@ -673,7 +680,7 @@ class TestTritonAttention(CustomTestCase):
) )
# Run unified kernel # 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( extend_attention_fwd_unified(
q_extend, q_extend,
o_unified, o_unified,
@@ -716,7 +723,7 @@ class TestTritonAttention(CustomTestCase):
"""Test build_unified_kv_indices correctness.""" """Test build_unified_kv_indices correctness."""
B = 4 B = 4
dtype = torch.int64 dtype = torch.int64
device = "cuda" device = get_device()
# Setup test data # Setup test data
prefix_lens = torch.tensor([10, 20, 15, 25], dtype=torch.int32, device=device) prefix_lens = torch.tensor([10, 20, 15, 25], dtype=torch.int32, device=device)