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

Fix nightly test failures: NSA indexer dtype and CPP radix cache init (#13958)

Browse Source
This commit is contained in:
alisonshao
2025-11-25 19:52:08 -08:00
committed by GitHub
parent 18fb51583f
commit 846ba3c62c
2 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -169,9 +169,10 @@ class Indexer(CustomOp):
@torch.compile(dynamic=True)
def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
weights, _ = self.weights_proj(x.float())
weights = weights * self.n_heads**-0.5
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
# Keep x in original dtype (bfloat16) for projection, then convert to float32
weights, _ = self.weights_proj(x)
weights = weights.float() * self.n_heads**-0.5
weights = weights.unsqueeze(-1) * q_scale.float() * self.softmax_scale
return weights
def _get_q_k_bf16(

2
python/sglang/srt/mem_cache/radix_cache_cpp.py
View File

@@ -45,8 +45,8 @@ class RadixCacheCpp(BasePrefixCache):
self.write_through_threshold = (
1 if server_args.hicache_write_policy == "write_through" else 2
)
self.device = self.token_to_kv_pool_allocator.device
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
self.device = self.token_to_kv_pool_allocator.device
self.req_to_token_pool = params.req_to_token_pool
self.page_size = params.page_size
self.kv_cache = self.token_to_kv_pool_allocator.get_kvcache()