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

[AMD] fix amd ci dpskv32 (#17432)

Browse Source 
Co-authored-by: michaelzhang-ai <michaelzhang.ai@users.noreply.github.com>
This commit is contained in:
YC Tseng
2026-01-21 20:34:24 -08:00
committed by GitHub
co-authored by michaelzhang-ai
parent fafa171529
commit 17807caf82
5 changed files with 16 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+4
View File
@@ -220,6 +220,8 @@ class Indexer(MultiPlatformOp):
@torch.compile(dynamic=True) if not _is_hip else lambda f: f
def _project_and_scale_head_gates(self, x: torch.Tensor):
if _is_hip:
x = x.to(self.weights_proj.weight.dtype)
weights, _ = self.weights_proj(x)
weights = weights.float()
weights = weights * self.n_heads**-0.5
@@ -227,6 +229,8 @@ class Indexer(MultiPlatformOp):
@torch.compile(dynamic=True) if not _is_hip else lambda f: f
def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
if _is_hip:
x = x.to(self.weights_proj.weight.dtype)
weights, _ = self.weights_proj(x)
weights = weights.float()
weights = weights * self.n_heads**-0.5