diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py index 41ce7fe9e..1cdf65b91 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py +++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 7ec4dd60f..cf362e712 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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( diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 6e5062c88..a5ad78b89 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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( diff --git a/test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py b/test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py deleted file mode 100644 index 75f83f6db..000000000 --- a/test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py +++ /dev/null @@ -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") diff --git a/test/manual/layers/attention/nsa/test_index_buf_accessor.py b/test/manual/layers/attention/nsa/test_index_buf_accessor.py index 4e38b859f..5e17e1859 100644 --- a/test/manual/layers/attention/nsa/test_index_buf_accessor.py +++ b/test/manual/layers/attention/nsa/test_index_buf_accessor.py @@ -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)