[Feature] Support Flashinfer fmha on Blackwell (#6930)
This commit is contained in:
@@ -29,6 +29,7 @@ from sglang.srt.layers.attention.flashinfer_backend import (
|
||||
create_flashinfer_kv_indices_triton,
|
||||
)
|
||||
from sglang.srt.layers.dp_attention import get_attention_tp_size
|
||||
from sglang.srt.layers.utils import is_sm100_supported
|
||||
from sglang.srt.managers.schedule_batch import global_server_args_dict
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
|
||||
@@ -108,8 +109,11 @@ class FlashInferMLAAttnBackend(AttentionBackend):
|
||||
else:
|
||||
self.q_indptr_decode = q_indptr_decode_buf
|
||||
|
||||
fmha_backend = "auto"
|
||||
if is_sm100_supported():
|
||||
fmha_backend = "cutlass"
|
||||
self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
|
||||
self.workspace_buffer, "NHD"
|
||||
self.workspace_buffer, "NHD", backend=fmha_backend
|
||||
)
|
||||
|
||||
if not self.skip_prefill:
|
||||
|
||||
Reference in New Issue
Block a user