From 846ba3c62c223eec91d5d75fef7d06445fc024db Mon Sep 17 00:00:00 2001 From: alisonshao <54658187+alisonshao@users.noreply.github.com> Date: Tue, 25 Nov 2025 19:52:08 -0800 Subject: [PATCH] Fix nightly test failures: NSA indexer dtype and CPP radix cache init (#13958) --- python/sglang/srt/layers/attention/nsa/nsa_indexer.py | 7 ++++--- python/sglang/srt/mem_cache/radix_cache_cpp.py | 2 +- 2 files changed, 5 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 f5db3d7a3..ebc22ea71 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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( diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index e8f187b46..8bb08351d 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -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()