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

Fix nightly test failure: NSA indexer dtype (#14017)

Browse Source
This commit is contained in:
alisonshao
2025-11-26 18:27:51 -08:00
committed by GitHub
parent b12c9e5c0a
commit d941a3befa
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
test/nightly/test_nsa_indexer.py
View File

@@ -248,12 +248,15 @@ class TestNSAIndexer(CustomTestCase):
# Move indexer to CUDA device
indexer = indexer.to(device=self.device)
# Convert linear layer weights to bfloat16 (but preserve LayerNorm's float32)
# Convert linear layer weights to bfloat16 (but preserve LayerNorm's float32
# and weights_proj's float32 - it uses params_dtype=torch.float32 in production)
# Need to recursively convert LinearBase submodules (like ReplicatedLinear)
for name, module in indexer.named_modules():
# Check for LinearBase (parent of ReplicatedLinear) but exclude LayerNorm
# Also exclude weights_proj which uses float32 params in production
if isinstance(module, LinearBase) and not isinstance(module, LayerNorm):
module.to(dtype=self.dtype)
if "weights_proj" not in name:
module.to(dtype=self.dtype)
return indexer