fix: specify dtype with begin_forward aka plan (#2404)

sgl-project · Dec 8, 2024 · 6128f7c · 6128f7c
1 parent a2486eb
commit 6128f7c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py
@@ -678,6 +678,7 @@ def call_begin_forward(
                 self.num_qo_heads,
                 self.num_kv_heads,
                 self.head_dim,
+                q_data_type=self.q_data_type,
             )
 
         # cached part
@@ -691,6 +692,7 @@ def call_begin_forward(
             self.num_kv_heads,
             self.head_dim,
             1,
+            q_data_type=self.q_data_type,
         )