From d1be60c3c58b63c00ecdc3969949120b15a13e21 Mon Sep 17 00:00:00 2001 From: Edwin Gao <105094497+edwingao28@users.noreply.github.com> Date: Sun, 9 Nov 2025 23:35:19 -0500 Subject: [PATCH] =?UTF-8?q?[Refactor]=20rename=20set=5Findex=5Fk=5Fand=5Fs?= =?UTF-8?q?cale=5Fbuffer=20to=20set=5Findex=5Fk=5Fscale=5Fb=E2=80=A6=20(#1?= =?UTF-8?q?2956)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- python/sglang/srt/layers/attention/nsa/nsa_indexer.py | 4 ++-- python/sglang/srt/mem_cache/memory_pool.py | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index c8903d279..eed4c9230 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -418,7 +418,7 @@ class Indexer(CustomOp): if not forward_batch.out_cache_loc.is_contiguous(): forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous() - forward_batch.token_to_kv_pool.set_index_k_and_scale_buffer( + forward_batch.token_to_kv_pool.set_index_k_scale_buffer( layer_id=layer_id, loc=forward_batch.out_cache_loc, index_k=k_fp8, @@ -599,7 +599,7 @@ class Indexer(CustomOp): # k_scale_cache: (num_total_tokens + page_size, head_dim // block_size = 1) fp8_e4m3fn if not forward_batch.out_cache_loc.is_contiguous(): forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous() - forward_batch.token_to_kv_pool.set_index_k_and_scale_buffer( + forward_batch.token_to_kv_pool.set_index_k_scale_buffer( layer_id=layer_id, loc=forward_batch.out_cache_loc, index_k=k_fp8, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 244c0acd9..739e28943 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1640,8 +1640,7 @@ class NSATokenToKVPool(MLATokenToKVPool): self, buf, seq_len=seq_len, page_indices=page_indices ) - # TODO rename later (currently use diff name to avoid confusion) - def set_index_k_and_scale_buffer( + def set_index_k_scale_buffer( self, layer_id: int, loc: torch.Tensor,