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

[DeepseekV32] Enable flashmla_prefill kernel with fp8 kvcache (#11655)

Browse Source 
Signed-off-by: Hao Lu <14827759+hlu1@users.noreply.github.com>
This commit is contained in:
hlu1
2025-10-27 23:11:48 -07:00
committed by GitHub
parent 83b2240074
commit 81a632ace6
4 changed files with 367 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -135,7 +135,16 @@ GRAMMAR_BACKEND_CHOICES = ["xgrammar", "outlines", "llguidance", "none"]
DETERMINISTIC_ATTENTION_BACKEND_CHOICES = ["flashinfer", "fa3", "triton"]
NSA_CHOICES = ["flashmla_sparse", "flashmla_kv", "fa3", "tilelang", "aiter"]
DEFAULT_LORA_EVICTION_POLICY = "lru"
NSA_CHOICES = [
"flashmla_sparse",
"flashmla_kv",
"flashmla_auto",
"fa3",
"tilelang",
"aiter",
]
RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu"]
@@ -1022,16 +1031,30 @@ class ServerArgs:
import torch
major, _ = torch.cuda.get_device_capability()
if major >= 10:
self.kv_cache_dtype = "fp8_e4m3"
logger.warning("Setting KV cache dtype to fp8.")
if self.kv_cache_dtype == "auto":
self.kv_cache_dtype = "fp8_e4m3" if major >= 10 else "bfloat16"
logger.warning(
f"Setting KV cache dtype to {self.kv_cache_dtype} for DeepSeek NSA."
)
if self.kv_cache_dtype == "bf16":
self.kv_cache_dtype = "bfloat16"
assert self.kv_cache_dtype in [
"bfloat16",
"fp8_e4m3",
], "DeepSeek NSA only supports bf16/bfloat16 or fp8_e4m3 kv_cache_dtype"
if self.kv_cache_dtype == "fp8_e4m3":
self.nsa_prefill_backend = "flashmla_kv"
# flashmla_auto dispatches to flashmla_sparse/flashmla_kv based on hardware and heuristics
self.nsa_prefill_backend = "flashmla_auto"
self.nsa_decode_backend = "flashmla_kv"
logger.warning(
"Setting NSA backend to flashmla_kv for FP8 KV Cache."
"Setting NSA backend to flashmla_auto for prefill and flashmla_kv for decode for FP8 KV Cache."
)
else:
# set prefill/decode backends for Blackwell. The default settings are for Hopper.
if major >= 10:
self.nsa_prefill_backend = "flashmla_sparse"
self.nsa_decode_backend = "flashmla_sparse"
# Logging env vars for NSA
from sglang.srt.layers.attention.nsa.utils import (