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

[Refactor] rename set_index_k_and_scale_buffer to set_index_k_scale_b… (#12956)

Browse Source
This commit is contained in:
Edwin Gao
2025-11-09 23:35:19 -05:00
committed by GitHub
parent 583bb1804e
commit d1be60c3c5
2 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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,

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

@@ -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,