minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Baizhou Zhang
2025-11-26 10:45:23 -08:00
committed by GitHub
parent 7e964b5198
commit 8a9b8b8457
2 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+3 -4
View File
@@ -169,10 +169,9 @@ class Indexer(CustomOp):
@torch.compile(dynamic=True) @torch.compile(dynamic=True)
def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor): def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
# Keep x in original dtype (bfloat16) for projection, then convert to float32 weights, _ = self.weights_proj(x.float())
weights, _ = self.weights_proj(x) weights = weights * self.n_heads**-0.5
weights = weights.float() * self.n_heads**-0.5 weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
weights = weights.unsqueeze(-1) * q_scale.float() * self.softmax_scale
return weights return weights
def _get_q_k_bf16( def _get_q_k_bf16(
python/sglang/srt/mem_cache/radix_cache_cpp.py
+1 -1
View File
@@ -45,8 +45,8 @@ class RadixCacheCpp(BasePrefixCache):
self.write_through_threshold = ( self.write_through_threshold = (
1 if server_args.hicache_write_policy == "write_through" else 2 1 if server_args.hicache_write_policy == "write_through" else 2
) )
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
self.device = self.token_to_kv_pool_allocator.device self.device = self.token_to_kv_pool_allocator.device
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
self.req_to_token_pool = params.req_to_token_pool self.req_to_token_pool = params.req_to_token_pool
self.page_size = params.page_size self.page_size = params.page_size
self.kv_cache = self.token_to_kv_pool_allocator.get_kvcache() self.kv_cache = self.token_to_kv_pool_allocator.get_kvcache()