[HotFix]Fix dtype mismatch in nsa indexer on AMD device (#17518)

This commit is contained in:
Baizhou Zhang
2026-01-22 12:17:02 +08:00
committed by GitHub
parent d6e2b88288
commit 3373545b9f

View File

@@ -186,7 +186,7 @@ class Indexer(MultiPlatformOp):
self.hidden_size,
self.n_heads,
bias=False,
params_dtype=torch.bfloat16,
params_dtype=torch.bfloat16 if _is_cuda else torch.float32,
prefix=add_prefix("weights_proj", prefix),
)
self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)