fix: fallback to triton for attention-sink models (flashinfer unsupported) (#23139)
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com> Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
co-authored by
Liangsheng Yin
hnyls2002
parent
6c2714f5ae
commit
e3782d04d2
@@ -381,6 +381,8 @@ class ModelConfig:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
self.has_attention_sinks = self._detect_attention_sinks()
|
||||||
|
|
||||||
self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [
|
self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [
|
||||||
"MiMoV2FlashForCausalLM",
|
"MiMoV2FlashForCausalLM",
|
||||||
"MiMoV2MTP",
|
"MiMoV2MTP",
|
||||||
@@ -388,6 +390,25 @@ class ModelConfig:
|
|||||||
"Gemma4ForConditionalGeneration",
|
"Gemma4ForConditionalGeneration",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
def _detect_attention_sinks(self) -> bool:
|
||||||
|
"""Check whether the model uses learned attention sinks.
|
||||||
|
|
||||||
|
Attention sinks are per-head scalars added to the softmax denominator
|
||||||
|
to compensate for evicted KV-cache entries under sliding-window
|
||||||
|
attention. Not every hybrid-SWA model uses them.
|
||||||
|
"""
|
||||||
|
archs = self.hf_config.architectures or []
|
||||||
|
# GptOss always creates sinks unconditionally.
|
||||||
|
if "GptOssForCausalLM" in archs:
|
||||||
|
return True
|
||||||
|
|
||||||
|
# MiMoV2 creates sinks only when the config flags are set.
|
||||||
|
if any(a in archs for a in ("MiMoV2FlashForCausalLM", "MiMoV2MTP")):
|
||||||
|
return getattr(
|
||||||
|
self.hf_text_config, "add_swa_attention_sink_bias", False
|
||||||
|
) or getattr(self.hf_text_config, "add_full_attention_sink_bias", False)
|
||||||
|
return False
|
||||||
|
|
||||||
def _derive_context_length(self, context_length: int):
|
def _derive_context_length(self, context_length: int):
|
||||||
is_draft_model = self.is_draft_model
|
is_draft_model = self.is_draft_model
|
||||||
derived_context_len = get_context_length(self.hf_text_config)
|
derived_context_len = get_context_length(self.hf_text_config)
|
||||||
|
|||||||
@@ -2422,7 +2422,10 @@ class ServerArgs:
|
|||||||
elif is_mps():
|
elif is_mps():
|
||||||
return "torch_native"
|
return "torch_native"
|
||||||
else:
|
else:
|
||||||
return "flashinfer" if is_flashinfer_available() else "triton"
|
# FlashInfer does not support attention sinks.
|
||||||
|
if is_flashinfer_available() and not model_config.has_attention_sinks:
|
||||||
|
return "flashinfer"
|
||||||
|
return "triton"
|
||||||
else:
|
else:
|
||||||
# MLA architecture
|
# MLA architecture
|
||||||
if is_hopper_with_cuda_12_3():
|
if is_hopper_with_cuda_12_3():
|
||||||
@@ -3161,9 +3164,14 @@ class ServerArgs:
|
|||||||
|
|
||||||
# If decode backend is implicit, pick a safe backend without changing io backend.
|
# If decode backend is implicit, pick a safe backend without changing io backend.
|
||||||
if not self.use_mla_backend():
|
if not self.use_mla_backend():
|
||||||
self.decode_attention_backend = (
|
# FlashInfer does not support attention sinks.
|
||||||
"flashinfer" if is_flashinfer_available() else "triton"
|
if (
|
||||||
)
|
is_flashinfer_available()
|
||||||
|
and not self.get_model_config().has_attention_sinks
|
||||||
|
):
|
||||||
|
self.decode_attention_backend = "flashinfer"
|
||||||
|
else:
|
||||||
|
self.decode_attention_backend = "triton"
|
||||||
else:
|
else:
|
||||||
self.decode_attention_backend = (
|
self.decode_attention_backend = (
|
||||||
"flashinfer" if is_sm100_supported() else "triton"
|
"flashinfer" if is_sm100_supported() else "triton"
|
||||||
|
|||||||
Reference in New Issue
Block a user