minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Fix] Fix backend selection after flashinfer version update (#18364)

Browse Source
This commit is contained in:
DarkSharpness
2026-02-08 11:20:41 +08:00
committed by GitHub
parent 00248d85c7
commit 8e2e835c2f
1 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashinfer_backend.py
+9 -3
View File
@@ -121,6 +121,8 @@ class FlashInferAttnBackend(AttentionBackend):
init_new_workspace: bool = False, init_new_workspace: bool = False,
): ):
super().__init__() super().__init__()
self.prefill_backend = "fa2"
self.decode_backend = "fa2"
# Store multi-item scoring delimiter for efficient access # Store multi-item scoring delimiter for efficient access
self.multi_item_scoring_delimiter = ( self.multi_item_scoring_delimiter = (
@@ -264,19 +266,21 @@ class FlashInferAttnBackend(AttentionBackend):
BatchPrefillWithPagedKVCacheWrapper( BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer, self.workspace_buffer,
"NHD", "NHD",
backend="fa2", backend=self.prefill_backend,
) )
) )
self.prefill_wrappers_verify.append( self.prefill_wrappers_verify.append(
BatchPrefillWithPagedKVCacheWrapper( BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer, self.workspace_buffer,
"NHD", "NHD",
backend=self.prefill_backend,
) )
) )
self.decode_wrappers.append( self.decode_wrappers.append(
BatchDecodeWithPagedKVCacheWrapper( BatchDecodeWithPagedKVCacheWrapper(
self.workspace_buffer, self.workspace_buffer,
"NHD", "NHD",
backend=self.decode_backend,
use_tensor_cores=self.decode_use_tensor_cores, use_tensor_cores=self.decode_use_tensor_cores,
) )
) )
@@ -555,6 +559,7 @@ class FlashInferAttnBackend(AttentionBackend):
BatchDecodeWithPagedKVCacheWrapper( BatchDecodeWithPagedKVCacheWrapper(
self.workspace_buffer, self.workspace_buffer,
"NHD", "NHD",
backend=self.decode_backend,
use_cuda_graph=True, use_cuda_graph=True,
use_tensor_cores=self.decode_use_tensor_cores, use_tensor_cores=self.decode_use_tensor_cores,
paged_kv_indptr_buffer=self.kv_indptr[i][: num_tokens + 1], paged_kv_indptr_buffer=self.kv_indptr[i][: num_tokens + 1],
@@ -590,6 +595,7 @@ class FlashInferAttnBackend(AttentionBackend):
self.workspace_buffer, self.workspace_buffer,
"NHD", "NHD",
use_cuda_graph=True, use_cuda_graph=True,
backend=self.prefill_backend,
qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1], qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1],
paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1], paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1],
paged_kv_indices_buf=self.cuda_graph_kv_indices[i], paged_kv_indices_buf=self.cuda_graph_kv_indices[i],
@@ -619,7 +625,7 @@ class FlashInferAttnBackend(AttentionBackend):
BatchPrefillWithPagedKVCacheWrapper( BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer, self.workspace_buffer,
"NHD", "NHD",
backend="fa2", backend=self.prefill_backend,
use_cuda_graph=True, use_cuda_graph=True,
qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1], qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1],
paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1], paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1],
@@ -649,7 +655,7 @@ class FlashInferAttnBackend(AttentionBackend):
BatchPrefillWithPagedKVCacheWrapper( BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer, self.workspace_buffer,
"NHD", "NHD",
backend="fa2", backend=self.prefill_backend,
use_cuda_graph=True, use_cuda_graph=True,
qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1], qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1],
paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1], paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1],