From 006bd44cf92064bdd32a96f150a1aa77c2eb7cde Mon Sep 17 00:00:00 2001 From: BingjiaWang Date: Thu, 12 Mar 2026 03:56:33 +0800 Subject: [PATCH] [deepseekv3.2] fix get_k_and_s_triton kenel for 128K seqlen case bug (#19319) Co-authored-by: abing --- .../attention/nsa/index_buf_accessor.py | 149 +++++++++----- .../srt/layers/attention/nsa/nsa_indexer.py | 55 ++--- python/sglang/srt/mem_cache/memory_pool.py | 11 +- .../nsa/test_get_k_scale_triton_kernel.py | 191 ++++++++++++++++++ .../attention/nsa/test_index_buf_accessor.py | 55 ++++- 5 files changed, 380 insertions(+), 81 deletions(-) create mode 100644 test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py 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 1cdf65b91..da8121fc7 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py +++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -167,11 +167,20 @@ class GetKAndS: @classmethod def triton( - cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor + cls, + pool: "NSATokenToKVPool", + buf: torch.Tensor, + page_indices: torch.Tensor, + seq_len_tensor: torch.Tensor, + seq_len_sum: int, + max_seq_len: int, ): """ 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 @@ -179,7 +188,9 @@ class GetKAndS: return _get_k_and_s_triton( buf=buf, page_indices=page_indices, - seq_len=seq_len, + seq_lens=seq_len_tensor, + seq_len_sum=seq_len_sum, + max_seq_len=max_seq_len, page_size=pool.page_size, index_head_dim=pool.index_head_dim, ) @@ -599,7 +610,9 @@ def _get_s_triton_kernel( def _get_k_and_s_triton( buf: torch.Tensor, page_indices: torch.Tensor, - seq_len: int, + seq_lens: torch.Tensor, + seq_len_sum: int, + max_seq_len: int, page_size: int, index_head_dim: int, ): @@ -609,33 +622,52 @@ 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_len: int, number of tokens to gather + :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 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 """ - num_pages, buf_numel_per_page = buf.shape - s_offset_in_page = page_size * index_head_dim # Scales start after K data - # 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) + 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) + + _, buf_numel_per_page = buf.shape + _, page_indice_batch_offset = page_indices.shape + s_offset_in_page = page_size * index_head_dim # Launch kernel with one thread per token - grid = (seq_len,) + 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 + + seq_num = seq_lens.shape[0] + grid = (seq_num, num_token_blocks, num_k_threads) + seq_num_pow2 = 1 + while seq_num_pow2 < seq_num: + seq_num_pow2 *= 2 + _get_k_and_s_triton_kernel[grid]( - 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, + 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, + BLOCK_SIZE=BLOCK_SIZE, + BLOCK_SIZE_K=BLOCK_SIZE_K, ) return k_out, s_out @@ -647,11 +679,14 @@ def _get_k_and_s_triton_kernel( page_indices_ptr, k_out_ptr, s_out_ptr, - seq_len: tl.constexpr, + seq_len_ptr, + seq_len_num_pow: 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: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): """ @@ -659,40 +694,62 @@ 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. """ - token_id = tl.program_id(0) + batch_id = tl.program_id(0) + block_token_start = tl.program_id(1) * BLOCK_SIZE + thread_idx = tl.program_id(2) - # Calculate which page and offset within page - page_idx = token_id // page_size - token_offset_in_page = token_id % page_size + # Define the token range within the block and the K dimension range handled by the thread. + token_ids_in_block = tl.arange(0, BLOCK_SIZE) + token_ids = block_token_start + token_ids_in_block + k_offsets = thread_idx * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) - # Load the page index from page_indices - page_index = tl.load(page_indices_ptr + page_idx) + seq_len = tl.load(seq_len_ptr + batch_id) + token_valid_mask = token_ids < seq_len - # ===== Load K data (128 bytes) ===== - # Calculate source offset for K in buf - k_src_base_offset = ( - page_index * buf_numel_per_page + token_offset_in_page * index_head_dim + 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) + batch_token_offset = tl.sum(prev_seq_lens) + + # Batch calculate the page index and in-page offset of each token. + page_idx = token_ids // page_size + token_offset_in_page = token_ids % page_size + page_indices_base = batch_id * page_indice_batch_offset + page_idx_valid_mask = page_idx < page_indice_batch_offset + page_index = tl.load( + page_indices_ptr + page_idx + page_indices_base, + mask=token_valid_mask & page_idx_valid_mask, ) - # Load 128 bytes (index_head_dim elements) - k_offsets = tl.arange(0, BLOCK_SIZE_K) - k_mask = k_offsets < index_head_dim - k_data = tl.load(buf_ptr + k_src_base_offset + k_offsets, mask=k_mask) + # ===== Load K data ===== + # The address calculation logic for K: page_index * total number of elements in a single page + K offset of the token within the page. + k_src_token_offset = token_offset_in_page * index_head_dim + k_src_base_offset = page_index * buf_numel_per_page + k_src_token_offset + + k_load_addr = buf_ptr + k_src_base_offset[:, None] + k_offsets[None, :] + k_dim_mask = k_offsets[None, :] < index_head_dim + k_mask = token_valid_mask[:, None] & k_dim_mask + + k_data = tl.load(k_load_addr, mask=k_mask, other=0) # Store K to output - k_dst_offset = token_id * index_head_dim - tl.store(k_out_ptr + k_dst_offset + k_offsets, k_data, mask=k_mask) + k_dst_token_offset = batch_token_offset + token_ids + k_dst_base_offset = k_dst_token_offset * index_head_dim + k_store_addr = k_out_ptr + k_dst_base_offset[:, None] + k_offsets[None, :] + tl.store(k_store_addr, k_data, mask=k_mask) - # ===== Load S data (4 bytes) ===== - # Calculate source offset for S in buf - s_src_base_offset = ( - page_index * buf_numel_per_page + s_offset_in_page + token_offset_in_page * 4 - ) + # ===== Load S data ===== + # The address calculation logic for S: page_index * total number of elements in a single page + starting offset of S within the page + offset of token within S in the page + s_src_token_offset = s_offset_in_page + token_offset_in_page * 4 + s_src_base_offset = page_index * buf_numel_per_page + s_src_token_offset - # Load 4 bytes (fp32 scale) s_offsets = tl.arange(0, 4) - s_data = tl.load(buf_ptr + s_src_base_offset + s_offsets) + s_load_addr = buf_ptr + s_src_base_offset[:, None] + s_offsets[None, :] + s_mask = token_valid_mask[:, None] & (s_offsets[None, :] < 4) + s_data = tl.load(s_load_addr, mask=s_mask, other=0) # Store S to output - s_dst_offset = token_id * 4 - tl.store(s_out_ptr + s_dst_offset + s_offsets, s_data) + s_dst_token_offset = batch_token_offset + token_ids + s_dst_base_offset = s_dst_token_offset * 4 + s_store_addr = s_out_ptr + s_dst_base_offset[:, None] + s_offsets[None, :] + tl.store(s_store_addr, s_data, mask=s_mask) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index d2c146f3f..675b6f9ed 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -491,6 +491,7 @@ class Indexer(MultiPlatformOp): def _get_topk_ragged( self, + enable_dual_stream: bool, forward_batch: ForwardBatch, layer_id: int, q_fp8: torch.Tensor, @@ -509,9 +510,11 @@ 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() @@ -526,38 +529,37 @@ 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 - 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 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fnuz) - else: - 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) 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, + ) + if _is_fp8_fnuz: + k_fp8 = k_fp8.view(torch.float8_e4m3fnuz) + else: + k_fp8 = k_fp8.view(torch.float8_e4m3fn) + + k_scale = k_scale.view(torch.float32).squeeze(-1) + kv_fp8 = (k_fp8, k_scale) + + # Check if we need to chunk to avoid OOM 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: @@ -1187,7 +1189,12 @@ class Indexer(MultiPlatformOp): return torch.cat([topk_result_prev, topk_result_next], dim=0) else: topk_result = self._get_topk_ragged( - forward_batch, layer_id, q_fp8, weights, metadata + enable_dual_stream, + 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 16b1410c3..476b3248d 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1860,8 +1860,10 @@ class NSATokenToKVPool(MLATokenToKVPool): def get_index_k_scale_buffer( self, layer_id: int, - seq_len: int, + seq_len_tensor: torch.Tensor, 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. @@ -1876,7 +1878,12 @@ class NSATokenToKVPool(MLATokenToKVPool): """ buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] return index_buf_accessor.GetKAndS.execute( - self, buf, seq_len=seq_len, page_indices=page_indices + self, + buf, + page_indices=page_indices, + seq_len_tensor=seq_len_tensor, + seq_len_sum=seq_len_sum, + max_seq_len=max_seq_len, ) 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 new file mode 100644 index 000000000..296567559 --- /dev/null +++ b/test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py @@ -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") 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 5e17e1859..49395263d 100644 --- a/test/manual/layers/attention/nsa/test_index_buf_accessor.py +++ b/test/manual/layers/attention/nsa/test_index_buf_accessor.py @@ -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...")