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

Revert "optimize get_topk_ragged by fusing get k and k_scale triton kernel" (#18471)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-02-09 16:37:19 +08:00
committed by GitHub
parent 875ad6cf35
commit 615a02dcd4
5 changed files with 57 additions and 326 deletions
Download Patch File Download Diff File Expand all files Collapse all files

96
python/sglang/srt/layers/attention/nsa/index_buf_accessor.py
View File

@@ -167,20 +167,11 @@ class GetKAndS:
@classmethod
def triton(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering both K and S data from paged buffer in a single call.
:param page_indices: (num_pages,), int32/int64
:param seq_len_tensor: (num_pages,), int32/int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of all sequence len, int32
:return: tuple of (k_fp8, k_scale) where
k_fp8: (seq_len, index_head_dim), uint8
k_scale: (seq_len, 4), uint8
@@ -188,9 +179,7 @@ class GetKAndS:
return _get_k_and_s_triton(
buf=buf,
page_indices=page_indices,
seq_lens=seq_len_tensor,
seq_len_sum=seq_len_sum,
max_seq_len=max_seq_len,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
@@ -610,9 +599,7 @@ def _get_s_triton_kernel(
def _get_k_and_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
seq_len: int,
page_size: int,
index_head_dim: int,
):
@@ -622,44 +609,32 @@ def _get_k_and_s_triton(
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_lens: tensor of sequence lens, int64
:param seq_len_sum: sum of all sequence len, int32
:param seq_len_sum: max of sequence len, int32
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: tuple of (k_out, s_out) where
k_out: (seq_len, index_head_dim), uint8
s_out: (seq_len, 4), uint8
"""
# Allocate outputs
k_out = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device=buf.device)
num_pages, buf_numel_per_page = buf.shape
s_offset_in_page = page_size * index_head_dim # Scales start after K data
_, buf_numel_per_page = buf.shape
_, page_indice_batch_offset = page_indices.shape
s_offset_in_page = page_size * index_head_dim
# Allocate outputs
k_out = torch.empty((seq_len, index_head_dim), dtype=torch.uint8, device=buf.device)
s_out = torch.empty((seq_len, 4), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
seq_num = seq_lens.shape[0]
grid = (seq_num, max_seq_len)
seq_num_pow2 = 1
while seq_num_pow2 < seq_num:
seq_num_pow2 *= 2
grid = (seq_len,)
_get_k_and_s_triton_kernel[grid](
buf_ptr=buf,
page_indices_ptr=page_indices,
k_out_ptr=k_out,
s_out_ptr=s_out,
seq_len_ptr=seq_lens,
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,
buf,
page_indices,
k_out,
s_out,
seq_len,
page_size,
buf_numel_per_page,
index_head_dim,
s_offset_in_page,
BLOCK_SIZE_K=128,
)
@@ -672,13 +647,11 @@ def _get_k_and_s_triton_kernel(
page_indices_ptr,
k_out_ptr,
s_out_ptr,
seq_len_ptr,
seq_len_num_pow: tl.constexpr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
s_offset_in_page: tl.constexpr,
page_indice_batch_offset: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
):
"""
@@ -686,30 +659,14 @@ def _get_k_and_s_triton_kernel(
Each program handles one token (seq_len tokens total).
Loads 128 bytes (K) + 4 bytes (S) from the appropriate page.
"""
batch_id = tl.program_id(0)
token_id = tl.program_id(1)
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load batch id seq len from seq_len_ptr
seq_len = tl.load(seq_len_ptr + batch_id)
if token_id >= seq_len:
return
pre_batch_idx = tl.arange(0, seq_len_num_pow)
mask_pre_batch_idx = pre_batch_idx < batch_id
prev_seq_lens = tl.load(seq_len_ptr + pre_batch_idx, mask=mask_pre_batch_idx)
seq_len_offset = tl.sum(prev_seq_lens)
k_offset_batch = seq_len_offset * index_head_dim
s_offset_batch = seq_len_offset * 4
# Load the page index from page_indices
page_index = tl.load(
page_indices_ptr + page_idx + batch_id * page_indice_batch_offset
)
page_index = tl.load(page_indices_ptr + page_idx)
# ===== Load K data (128 bytes) =====
# Calculate source offset for K in buf
@@ -719,12 +676,12 @@ def _get_k_and_s_triton_kernel(
# Load 128 bytes (index_head_dim elements)
k_offsets = tl.arange(0, BLOCK_SIZE_K)
k_mask = (k_offsets < index_head_dim) & (token_id < seq_len)
k_mask = k_offsets < index_head_dim
k_data = tl.load(buf_ptr + k_src_base_offset + k_offsets, mask=k_mask)
# Store K to output
k_dst_offset = token_id * index_head_dim
tl.store(k_out_ptr + k_dst_offset + k_offsets + k_offset_batch, k_data, mask=k_mask)
tl.store(k_out_ptr + k_dst_offset + k_offsets, k_data, mask=k_mask)
# ===== Load S data (4 bytes) =====
# Calculate source offset for S in buf
@@ -734,9 +691,8 @@ def _get_k_and_s_triton_kernel(
# Load 4 bytes (fp32 scale)
s_offsets = tl.arange(0, 4)
s_mask = (s_offsets < 4) & (token_id < seq_len)
s_data = tl.load(buf_ptr + s_src_base_offset + s_offsets, mask=s_mask)
s_data = tl.load(buf_ptr + s_src_base_offset + s_offsets)
# Store S to output
s_dst_offset = token_id * 4
tl.store(s_out_ptr + s_dst_offset + s_offsets + s_offset_batch, s_data, mask=s_mask)
tl.store(s_out_ptr + s_dst_offset + s_offsets, s_data)

51
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -469,7 +469,6 @@ class Indexer(MultiPlatformOp):
def _get_topk_ragged(
self,
enable_dual_stream: bool,
forward_batch: ForwardBatch,
layer_id: int,
q_fp8: torch.Tensor,
@@ -488,11 +487,9 @@ class Indexer(MultiPlatformOp):
assert page_size == 64, "only support page size 64"
assert len(weights.shape) == 3
assert (
forward_batch.seq_lens_cpu is not None
and forward_batch.extend_seq_lens_cpu is not None
)
weights = weights.squeeze(-1)
k_fp8_list = []
k_scale_list = []
if _is_hip:
block_tables = metadata.get_page_table_1()
@@ -507,37 +504,38 @@ class Indexer(MultiPlatformOp):
batch_size = len(block_tables)
token_nums, _, _ = q_fp8.shape
device = q_fp8.device
topk_result = torch.full(
(token_nums, self.index_topk), -1, device=device, dtype=torch.int32
)
if batch_size == 0:
return topk_result
ks, ke = metadata.get_indexer_kvcache_range()
seq_len_sum = forward_batch.seq_lens_sum
max_seq_len = torch.max(forward_batch.seq_lens_cpu).item()
k_fp8, k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_buffer(
layer_id,
forward_batch.seq_lens,
block_tables,
seq_len_sum,
max_seq_len,
)
indexer_seq_lens_cpu = metadata.get_indexer_seq_len_cpu()
assert len(indexer_seq_lens_cpu) == batch_size
for i in range(batch_size):
seq_len = indexer_seq_lens_cpu[i].item()
assert isinstance(seq_len, int)
# Use fused Triton kernel to get both K and scale in a single call
k_fp8, k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_buffer(
layer_id,
seq_len,
block_tables[i],
)
k_fp8_list.append(k_fp8)
k_scale_list.append(k_scale)
if _is_fp8_fnuz:
k_fp8 = k_fp8.view(torch.float8_e4m3fnuz)
k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fnuz)
else:
k_fp8 = k_fp8.view(torch.float8_e4m3fn)
k_scale = k_scale.view(torch.float32).squeeze(-1)
k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1)
kv_fp8 = (k_fp8, k_scale)
# Check if we need to chunk to avoid OOM
ks, ke = metadata.get_indexer_kvcache_range()
seq_lens_expanded = metadata.get_seqlens_expanded()
token_to_batch_idx = metadata.get_token_to_batch_idx()
q_offset = ks.shape[0]
k_offset = k_fp8.shape[0]
# Check if we need to chunk to avoid OOM
need_chunk, free_mem = self._should_chunk_mqa_logits(q_offset, k_offset, device)
if not need_chunk:
@@ -1113,12 +1111,7 @@ class Indexer(MultiPlatformOp):
return torch.cat([topk_result_prev, topk_result_next], dim=0)
else:
topk_result = self._get_topk_ragged(
enable_dual_stream,
forward_batch,
layer_id,
q_fp8,
weights,
metadata,
forward_batch, layer_id, q_fp8, weights, metadata
)
else:
topk_result = self.forward_indexer(

11
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -1837,10 +1837,8 @@ class NSATokenToKVPool(MLATokenToKVPool):
def get_index_k_scale_buffer(
self,
layer_id: int,
seq_len_tensor: torch.Tensor,
seq_len: int,
page_indices: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
"""
Fused method to get both index K and scale data in a single call using Triton.
@@ -1855,12 +1853,7 @@ class NSATokenToKVPool(MLATokenToKVPool):
"""
buf = self.index_k_with_scale_buffer[layer_id - self.start_layer]
return index_buf_accessor.GetKAndS.execute(
self,
buf,
page_indices=page_indices,
seq_len_tensor=seq_len_tensor,
seq_len_sum=seq_len_sum,
max_seq_len=max_seq_len,
self, buf, seq_len=seq_len, page_indices=page_indices
)
def set_index_k_scale_buffer(

183
test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py
View File

@@ -1,183 +0,0 @@
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()
grid = (batch, max_seq_len)
BLOCK_SIZE = 128
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_K=BLOCK_SIZE,
)
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_K=BLOCK_SIZE,
)
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_K=BLOCK_SIZE,
)
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")

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

@@ -264,7 +264,6 @@ 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(
@@ -284,16 +283,13 @@ 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, page_indices_, seq_len_tensor, seq_len, seq_len
)
k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices)
# Verify shapes
assert k_torch.shape == (seq_len, index_head_dim)
@@ -324,7 +320,6 @@ 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
@@ -333,16 +328,13 @@ 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, page_indices_, seq_len_tensor, seq_len, seq_len
)
k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
@@ -354,7 +346,6 @@ 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
@@ -363,16 +354,13 @@ 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, page_indices_, seq_len_tensor, seq_len, seq_len
)
k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
@@ -384,7 +372,6 @@ 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
@@ -393,16 +380,13 @@ 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, page_indices_, seq_len_tensor, seq_len, seq_len
)
k_triton, s_triton = GetKAndS.triton(pool, buf, seq_len, page_indices)
# Should handle partial pages correctly
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
@@ -420,14 +404,12 @@ 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)
@@ -440,9 +422,7 @@ class TestEdgeCases:
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
@@ -453,14 +433,12 @@ 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)
@@ -473,9 +451,7 @@ class TestEdgeCases:
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
@@ -486,7 +462,6 @@ 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
@@ -497,7 +472,6 @@ 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)
@@ -510,9 +484,7 @@ class TestEdgeCases:
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
k_triton2, s_triton2 = GetKAndS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)