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

6128f7cf · Yineng Zhang · GitHub · a2486eb5 · 6128f7cf
Unverified Commit 6128f7cf authored Dec 08, 2024 by Yineng Zhang Committed by GitHub Dec 08, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

python/sglang/srt/layers/attention/flashinfer_backend.py python/sglang/srt/layers/attention/flashinfer_backend.py +2 -0

No files found.
--- a/python/sglang/srt/layers/attention/flashinfer_backend.py
+++ b/python/sglang/srt/layers/attention/flashinfer_backend.py
@@ -678,6 +678,7 @@ class FlashInferIndicesUpdaterPrefill:
                self.num_qo_heads,
                self.num_kv_heads,
                self.head_dim,
+                q_data_type=self.q_data_type,
            )

        # cached part
@@ -691,6 +692,7 @@ class FlashInferIndicesUpdaterPrefill:
            self.num_kv_heads,
            self.head_dim,
            1,
+            q_data_type=self.q_data_type,
        )