From be5121b452197fbab12a274755c4687017e80998 Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Tue, 20 Jan 2026 19:58:13 -0800 Subject: [PATCH] Fix NSA indexer in the nightly test (#17452) --- test/registered/kernels/test_nsa_indexer.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/test/registered/kernels/test_nsa_indexer.py b/test/registered/kernels/test_nsa_indexer.py index 87d11439b..212aea1d0 100644 --- a/test/registered/kernels/test_nsa_indexer.py +++ b/test/registered/kernels/test_nsa_indexer.py @@ -79,6 +79,22 @@ class MockIndexerMetadata(BaseIndexerMetadata): ) return page_table + def get_page_table_1(self) -> torch.Tensor: + """Return: (batch_size, num_blocks) int32, page table with page size 1.""" + # Create a simple page table for testing with page size 1 + max_seq_len = max(self.seq_lens) + num_blocks = max_seq_len # Page size 1 means num_blocks == max_seq_len + page_table = torch.zeros( + (self.batch_size, num_blocks), dtype=torch.int32, device=self.device + ) + for i in range(self.batch_size): + # Simple linear mapping: block i maps to page i + num_blocks_needed = self.seq_lens[i] + page_table[i, :num_blocks_needed] = torch.arange( + num_blocks_needed, device=self.device + ) + return page_table + def get_seqlens_expanded(self) -> torch.Tensor: """Return: (sum_extend_seq_len,) int32 tensor""" # For extend mode, each new token attends to progressively more tokens