fixes for device and localheaddim

sgl-project · Feb 2, 2025 · 80621ab · 80621ab
1 parent e0e7589
commit 80621ab
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
@@ -673,10 +673,10 @@ def forward_absorb_fused_mla(
         attn_logits = torch.empty(
                              (
                                 forward_batch.batch_size,
-                                num_local_heads,
+                                self.num_local_heads,
                                 num_kv_split,
                                 self.kv_lora_rank + 1,
-                             ), dtype=torch.float32, device=self.device)
+                             ), dtype=torch.float32, device=q.device)
 
         decode_attention_fwd_normal(
                         q,