chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Check KV4 compatibility with attention backends and add KV4 support to the attention_backend doc (#14467)

Browse Source 
Signed-off-by: Ho-Ren (Jack) Chuang <horenchuang@bytedance.com>
This commit is contained in:
Ho-Ren (Jack) Chuang
2025-12-11 19:00:53 -08:00
committed by GitHub
parent b62e7e3bfe
commit 10146af099
2 changed files with 108 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
python/sglang/srt/server_args.py
View File

@@ -656,6 +656,7 @@ class ServerArgs:
# Set kernel backends.
self._handle_sampling_backend()
self._handle_attention_backend_compatibility()
self._handle_kv4_compatibility()
self._handle_page_size()
self._handle_amd_specifics()
self._handle_grammar_backend()
@@ -1555,6 +1556,84 @@ class ServerArgs:
self.enable_mixed_chunk = False
self.disable_radix_cache = True
def _handle_kv4_compatibility(self):
"""Check FP4 KV cache compatibility with the attention backend"""
if self.kv_cache_dtype != "fp4_e2m1":
return
use_mla_backend = self.use_mla_backend()
# self.attention_backend didn't overwrite self.prefill/decode_attention_backend yet
self.prefill_attention_backend_str, self.decode_attention_backend_str = (
self.get_attention_backends()
)
if is_cuda():
if (
self.prefill_attention_backend_str != self.decode_attention_backend_str
and self.prefill_attention_backend_str != "fa4"
): # Take care of prefill=fa4 later
logger.warning(
f"Attention: Using KV4 with PREFILL = {self.prefill_attention_backend_str} "
f"and DECODE = {self.decode_attention_backend_str}. "
f"Compatibility issues are unlikely, but may occur in rare edge cases."
)
else:
if self.prefill_attention_backend_str == "fa4":
if use_mla_backend: # FA4 + MLA
KV4_FA4_MLA_BACKEND_CHOICES = [
"cutlass_mla",
"flashinfer",
"trtllm_mla",
]
assert (
self.decode_attention_backend_str
in KV4_FA4_MLA_BACKEND_CHOICES
), (
f"KV4 FA4 MLA expects decode_attention_backend to be one of "
f"{KV4_FA4_MLA_BACKEND_CHOICES}, but got {self.decode_attention_backend_str}"
)
else: # FA4 + MHA
KV4_FA4_MHA_BACKEND_CHOICES = [
"triton",
"torch_native",
"flex_attention",
]
assert (
self.decode_attention_backend_str
in KV4_FA4_MHA_BACKEND_CHOICES
), (
f"KV4 FA4 MHA expects decode_attention_backend to be one of "
f"{KV4_FA4_MHA_BACKEND_CHOICES}, but got {self.decode_attention_backend_str}"
)
else:
if use_mla_backend: # !FA4 + MLA
KV4_ATTENTION_MLA_BACKEND_CHOICES = [
"cutlass_mla",
"flashinfer",
"trtllm_mla",
]
assert (
self.attention_backend in KV4_ATTENTION_MLA_BACKEND_CHOICES
), (
f"KV4 MLA expects attention_backend to be one of "
f"{KV4_ATTENTION_MLA_BACKEND_CHOICES}, but got {self.attention_backend}"
)
else: # !FA4 + MHA
KV4_ATTENTION_MHA_BACKEND_CHOICES = [
"triton",
"torch_native",
"flex_attention",
"trtllm_mha",
]
assert (
self.attention_backend in KV4_ATTENTION_MHA_BACKEND_CHOICES
), (
f"KV4 MHA expects attention_backend to be one of "
f"{KV4_ATTENTION_MHA_BACKEND_CHOICES}, but got {self.attention_backend}"
)
else:
raise RuntimeError("KV4 is not tested on non-CUDA platforms.")
def _handle_page_size(self):
if self.page_size is None:
self.page_size = 1