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

[deepseekv3.2] fix get_k_and_s_triton kenel for 128K seqlen case bug (#19319)

Browse Source 
Co-authored-by: abing <wangbingjia.wbj@alibaba-inc.com>
This commit is contained in:
BingjiaWang
2026-03-12 03:56:33 +08:00
committed by GitHub
parent e6a6cd1f0c
commit 006bd44cf9
5 changed files with 380 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

191
test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py Normal file
View File

@@ -0,0 +1,191 @@
import torch
from sglang.srt.layers.attention.nsa.index_buf_accessor import (
_get_k_and_s_triton_kernel,
)
def golden_torch_gen(
seq_len_tensor: torch.Tensor,
buffer_indexer: torch.Tensor,
buffer: torch.Tensor,
index_head_dim,
page_size,
):
dim_split = page_size * index_head_dim
torch_k_out = buffer[:, 0:dim_split]
torch_s_out = buffer[:, dim_split:]
torch_k_out = torch_k_out.reshape(-1, 128)
torch_s_out = torch_s_out.reshape(-1, 4)
batch = seq_len_tensor.shape[0]
index_list = []
for i in range(batch):
seq_len = seq_len_tensor[i].item()
buffer_index_ = buffer_indexer[i]
align_seq_len = ((seq_len + page_size - 1) / page_size) * page_size
needed_block_num = int((seq_len + page_size - 1) / page_size)
for j in range(needed_block_num):
block_idx = buffer_index_[j].item()
start_idx = block_idx * page_size
end_idx = 0
if j == (needed_block_num - 1):
end_idx = block_idx * page_size + (
seq_len - (needed_block_num - 1) * page_size
)
else:
end_idx = (block_idx + 1) * page_size
index_tensor = (
torch.arange(start=start_idx, end=end_idx, step=1)
.type(torch.int32)
.cuda()
)
index_list.append(index_tensor)
index_list_ = torch.cat(index_list, dim=0)
torch_k_out = torch.index_select(torch_k_out, dim=0, index=index_list_)
torch_s_out = torch.index_select(torch_s_out, dim=0, index=index_list_)
return torch_k_out, torch_s_out
def get_k_and_s_triton():
index_head_dim = 128
page_size = 64
num_page = 128
s_offset_in_page = page_size * index_head_dim
seq_len_tensor = torch.tensor(
[256, 267, 215, 32, 129], dtype=torch.int64, device="cuda"
) # 4 + 5 + 3 + 1 + 3 block
buffer_indexer = torch.tensor(
[
[1, 2, 3, 4, 0],
[7, 6, 5, 8, 9],
[10, 11, 12, 0, 0],
[13, 0, 0, 0, 0],
[14, 15, 16, 0, 0],
],
dtype=torch.int32,
device="cuda",
)
seq_len_sum = seq_len_tensor.sum()
batch = seq_len_tensor.shape[0]
triton_k_out = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device="cuda"
)
triton_s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device="cuda")
buffer = torch.randint(
0,
num_page,
(num_page, page_size * index_head_dim + page_size * 4),
device="cuda",
).type(torch.uint8)
_, buf_numel_per_page = buffer.shape
_, page_indice_batch_offset = buffer_indexer.shape
max_seq_len = seq_len_tensor.max().item()
BLOCK_SIZE = 256
BLOCK_SIZE_K = 128
num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
grid = (batch, num_token_blocks, num_k_threads)
seq_num_pow2 = 1
while seq_num_pow2 < batch:
seq_num_pow2 *= 2
# acc test =====================
_get_k_and_s_triton_kernel[grid](
buf_ptr=buffer,
page_indices_ptr=buffer_indexer,
k_out_ptr=triton_k_out,
s_out_ptr=triton_s_out,
seq_len_ptr=seq_len_tensor,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
torch_k_out, torch_s_out = golden_torch_gen(
seq_len_tensor=seq_len_tensor,
buffer_indexer=buffer_indexer,
buffer=buffer,
index_head_dim=index_head_dim,
page_size=page_size,
)
torch.testing.assert_close(
triton_k_out, torch_k_out, rtol=0, atol=0, msg="k outputs differ!"
)
torch.testing.assert_close(
triton_s_out, torch_s_out, rtol=0, atol=0, msg="s outputs differ!"
)
print("_get_k_and_s_triton_kernel test pass")
# perf test =====================
import time
torch.cuda.synchronize()
for _ in range(10):
_get_k_and_s_triton_kernel[grid](
buf_ptr=buffer,
page_indices_ptr=buffer_indexer,
k_out_ptr=triton_k_out,
s_out_ptr=triton_s_out,
seq_len_ptr=seq_len_tensor,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
torch.cuda.synchronize()
start_time = time.perf_counter()
_get_k_and_s_triton_kernel[grid](
buf_ptr=buffer,
page_indices_ptr=buffer_indexer,
k_out_ptr=triton_k_out,
s_out_ptr=triton_s_out,
seq_len_ptr=seq_len_tensor,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
end_time = time.perf_counter()
print(
f"_get_k_and_s_triton_kernel triton kernel infer time is {((end_time-start_time)*1000):.4f} ms\n"
)
if __name__ == "__main__":
if not torch.cuda.is_available():
print("CUDA not available. Skipping tests.")
exit(0)
print("Start test cases...\n")
get_k_and_s_triton()
print("End test cases...\n")

55
test/manual/layers/attention/nsa/test_index_buf_accessor.py
View File

@@ -264,6 +264,7 @@ class TestGetKAndS:
# Ensure seq_len doesn't exceed available pages
max_seq_len = num_pages * page_size
seq_len = min(seq_len, max_seq_len)
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
# Create mock pool
pool = MockNSATokenToKVPool(
@@ -283,13 +284,16 @@ class TestGetKAndS:
page_indices = torch.randint(
0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
)
page_indices_ = page_indices.unsqueeze(0)
# Run baseline: separate torch_fast calls
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
# Run fused Triton implementation
k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices)
k_triton, s_triton = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
# Verify shapes
assert k_torch.shape == (seq_len, index_head_dim)
@@ -320,6 +324,7 @@ class TestGetKAndS:
index_head_dim = 128
num_pages = 10
seq_len = 320 # 5 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
@@ -328,13 +333,16 @@ class TestGetKAndS:
# Sequential page indices [0, 1, 2, 3, 4]
page_indices = torch.arange(5, dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Baseline
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
# Fused
k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices)
k_triton, s_triton = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
@@ -346,6 +354,7 @@ class TestGetKAndS:
index_head_dim = 128
num_pages = 5
seq_len = 192 # 3 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
@@ -354,13 +363,16 @@ class TestGetKAndS:
# Repeated page indices [2, 2, 2]
page_indices = torch.full((3,), 2, dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Baseline
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
# Fused
k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices)
k_triton, s_triton = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
@@ -372,6 +384,7 @@ class TestGetKAndS:
index_head_dim = 128
num_pages = 5
seq_len = 100 # Not a multiple of 64
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
@@ -380,13 +393,16 @@ class TestGetKAndS:
num_pages_needed = (seq_len + page_size - 1) // page_size
page_indices = torch.arange(num_pages_needed, dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Baseline
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
# Fused
k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices)
k_triton, s_triton = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
# Should handle partial pages correctly
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
@@ -404,12 +420,14 @@ class TestEdgeCases:
index_head_dim = 128
num_pages = 2
seq_len = 1
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
page_indices = torch.tensor([0], dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Test GetK
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
@@ -422,7 +440,9 @@ class TestEdgeCases:
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices)
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
@@ -433,12 +453,14 @@ class TestEdgeCases:
index_head_dim = 128
num_pages = 5
seq_len = 192 # Exactly 3 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
page_indices = torch.arange(3, dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Test GetK
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
@@ -451,7 +473,9 @@ class TestEdgeCases:
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices)
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
@@ -462,6 +486,7 @@ class TestEdgeCases:
index_head_dim = 128
num_pages = 100
seq_len = 4096 # 64 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
@@ -472,6 +497,7 @@ class TestEdgeCases:
page_indices = torch.randint(
0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
)
page_indices_ = page_indices.unsqueeze(0)
# Test GetK
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
@@ -484,7 +510,9 @@ class TestEdgeCases:
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices)
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
@@ -532,14 +560,23 @@ if __name__ == "__main__":
print("✓ GetS tests passed\n")
# Test GetKAndS
print("Testing GetKAndS...")
print("Testing GetKAndS SeqLen=256...")
test_get_k_and_s = TestGetKAndS()
test_get_k_and_s.test_get_k_and_s_correctness(
num_pages=4, seq_len=256, page_size=64, index_head_dim=128
)
test_get_k_and_s.test_get_k_and_s_sequential_pages()
test_get_k_and_s.test_get_k_and_s_partial_page()
print("✓ GetKAndS tests passed\n")
print("✓ GetKAndS SeqLen=256 tests passed\n")
print("Testing GetKAndS SeqLen=128K...")
test_get_k_and_s = TestGetKAndS()
test_get_k_and_s.test_get_k_and_s_correctness(
num_pages=2048, seq_len=131072, page_size=64, index_head_dim=128
)
test_get_k_and_s.test_get_k_and_s_sequential_pages()
test_get_k_and_s.test_get_k_and_s_partial_page()
print("✓ GetKAndS SeqLen=128K tests passed\n")
# Test edge cases
print("Testing edge cases...")