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

144
python/sglang/srt/layers/attention/nsa/dequant_k_cache.py
View File

@@ -22,6 +22,10 @@ def _dequantize_k_cache_slow(
De-quantize the k-cache
"""
assert dv % tile_size == 0
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_tiles = dv // tile_size
num_blocks, block_size, h_k, _ = quant_k_cache.shape
assert h_k == 1
@@ -45,8 +49,10 @@ def _dequantize_k_cache_slow(
cur_nope * cur_scales
)
result = result.view(num_blocks, block_size, 1, d)
return result
if original_ndim == 3:
return result.view(num_blocks, 1, -1)
else:
return result.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast_wrapped(
@@ -54,7 +60,10 @@ def _dequantize_k_cache_fast_wrapped(
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
# TODO the final API may be 2D instead of 4D, thus we convert them here
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_blocks, block_size, _, dim_quant = quant_k_cache.shape
assert dv == 512
assert dim_quant == 656
@@ -63,7 +72,10 @@ def _dequantize_k_cache_fast_wrapped(
output = _dequantize_k_cache_fast(quant_k_cache)
return output.view(num_blocks, block_size, 1, -1)
if original_ndim == 3:
return output.view(num_blocks, 1, -1)
else:
return output.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
@@ -85,7 +97,6 @@ def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
NUM_NOPE_BLOCKS = dim_nope // group_size
input_nope_q = quant_k_cache[:, :dim_nope]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
@@ -102,7 +113,7 @@ def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
@@ -159,5 +170,126 @@ def _dequantize_k_cache_fast_kernel(
tl.store(dst_ptr, data, mask=mask)
def dequantize_k_cache_paged(
quant_k_cache: torch.Tensor,
page_table_1_flattened: torch.Tensor,
group_size: int = 128,
) -> torch.Tensor:
"""
De-quantize the k-cache with paged layout
Args:
quant_k_cache: [total_num_tokens, 1, dim_quant] or [num_blocks, block_size, 1, dim_quant], the quantized k-cache in paged layout
page_table_1_flattened: [num_tokens], the flattened page_table_1 with the page indices in each requests concatenated together
Returns:
output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
"""
dim_quant = quant_k_cache.shape[-1]
assert (
dim_quant == 656
), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
quant_k_cache = quant_k_cache.view((-1, dim_quant))
total_num_tokens, _ = quant_k_cache.shape
num_tokens = page_table_1_flattened.shape[0]
assert num_tokens <= total_num_tokens
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size # 512 // 128 = 4
output = torch.empty(
(num_tokens, 1, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
# cdiv(512 + 64, 128) = 5
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
# [:, 512:512+4*4] = [:, 512:528]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
# [:, 528:]
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_paged_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
page_table_1_flattened,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_paged_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
page_table_1_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id_paged * input_nope_q_stride_0 + offs_q
ptr_s = (
input_nope_s_ptr
+ token_id_paged * input_nope_s_stride_0
+ effective_block_id
)
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id_paged * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
raise Exception("UT is in quant_k_cache.py")

56
python/sglang/srt/layers/attention/nsa/quant_k_cache.py
View File

@@ -206,6 +206,8 @@ def _quantize_k_cache_fast_kernel(
if __name__ == "__main__":
import dequant_k_cache
for num_blocks, block_size in [
(1, 1),
(10, 64),
@@ -217,21 +219,9 @@ if __name__ == "__main__":
dtype=torch.bfloat16,
device="cuda",
)
# temp debug
# input_k_cache = (576 - torch.arange(num_blocks * block_size * 1 * dim_nope_and_rope, device="cuda")).to(torch.bfloat16).reshape(num_blocks, block_size, 1, dim_nope_and_rope)
ref_quant = _quantize_k_cache_slow(input_k_cache)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
# print(f"{input_k_cache=}")
# print(f"{ref_quant=}")
# print(f"{actual_quant=}")
# print(f"{ref_quant == actual_quant=}")
# print(f"{actual_quant.to(torch.float32) - ref_quant.to(torch.float32)=}")
# print(f"{ref_quant.view(torch.bfloat16)=}")
# print(f"{actual_quant.view(torch.bfloat16)=}")
# assert torch.all(ref_quant == actual_quant)
import dequant_k_cache
ref_ref_dequant = dequant_k_cache._dequantize_k_cache_slow(ref_quant)
ref_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(ref_quant)
@@ -252,4 +242,46 @@ if __name__ == "__main__":
ref_ref_dequant, actual_actual_dequant, atol=0.2, rtol=0.2
)
# test dequant_k_cache_paged
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
actual_dequant_paged = dequant_k_cache.dequantize_k_cache_paged(
actual_quant, page_table_1
).reshape(actual_actual_dequant.shape)
print(f"{torch.mean(actual_actual_dequant - actual_dequant_paged)=}")
torch.testing.assert_close(
ref_ref_dequant, actual_dequant_paged, atol=0.2, rtol=0.2
)
print("Passed")
print("Do benchmark...")
for num_blocks, block_size in [
(1, 64),
(64, 64),
(128, 64),
(256, 64),
(512, 64),
(1024, 64),
(2048, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
def run_ans():
return dequant_k_cache.dequantize_k_cache_paged(actual_quant, page_table_1)
ans_time: float = triton.testing.do_bench(run_ans, warmup=10, rep=20) / 1000 # type: ignore
print(f"seq_kv: {num_blocks * block_size}, time: {ans_time * 1e6: 4.0f} us")

176
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -1,12 +1,14 @@
from __future__ import annotations
from dataclasses import dataclass
from enum import IntEnum, auto
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, TypeAlias
import torch
from sglang.srt.configs.model_config import get_nsa_index_topk, is_deepseek_nsa
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged
from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.attention.nsa.quant_k_cache import quantize_k_cache
from sglang.srt.layers.attention.nsa.transform_index import (
@@ -98,11 +100,27 @@ class NSAMetadata:
nsa_max_seqlen_q: Literal[1] = 1 # always 1 for decode, variable for extend
flashmla_metadata: Optional[NSAFlashMLAMetadata] = None
# The sum of sequence lengths for key, prefill only
seq_lens_sum: Optional[int] = None
# The flattened 1D page table with shape (seq_lens_sum,), prefill only
# this table is always with page_size = 1
page_table_1_flattened: Optional[torch.Tensor] = None
# The offset of topk indices in ragged kv, prefill only
# shape: (seq_lens_sum,)
topk_indices_offset: Optional[torch.Tensor] = None
class TopkTransformMethod(IntEnum):
# Transform topk indices to indices to the page table (page_size = 1)
PAGED = auto()
# Transform topk indices to indices to ragged kv (non-paged)
RAGGED = auto()
@dataclass(frozen=True)
class NSAIndexerMetadata(BaseIndexerMetadata):
attn_metadata: NSAMetadata
topk_transform_method: TopkTransformMethod
def get_seqlens_int32(self) -> torch.Tensor:
return self.attn_metadata.cache_seqlens_int32
@@ -118,23 +136,36 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
logits: torch.Tensor,
topk: int,
) -> torch.Tensor:
from sgl_kernel import fast_topk_transform_fused, fast_topk_v2
from sgl_kernel import (
fast_topk_transform_fused,
fast_topk_transform_ragged_fused,
fast_topk_v2,
)
if not NSA_FUSE_TOPK:
return fast_topk_v2(logits, self.get_seqlens_expanded(), topk)
# NOTE(dark): if fused, we return a transformed page table directly
return fast_topk_transform_fused(
score=logits,
lengths=self.get_seqlens_expanded(),
page_table_size_1=self.attn_metadata.page_table_1,
cu_seqlens_q=self.attn_metadata.cu_seqlens_q,
topk=topk,
)
elif self.topk_transform_method == TopkTransformMethod.PAGED:
# NOTE(dark): if fused, we return a transformed page table directly
return fast_topk_transform_fused(
score=logits,
lengths=self.get_seqlens_expanded(),
page_table_size_1=self.attn_metadata.page_table_1,
cu_seqlens_q=self.attn_metadata.cu_seqlens_q,
topk=topk,
)
elif self.topk_transform_method == TopkTransformMethod.RAGGED:
return fast_topk_transform_ragged_fused(
score=logits,
lengths=self.get_seqlens_expanded(),
topk_indices_offset=self.attn_metadata.topk_indices_offset,
topk=topk,
)
else:
assert False, f"Unsupported {self.topk_transform_method = }"
def compute_cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor:
assert seqlens.dtype == torch.int32 and seqlens.is_cuda
assert seqlens.dtype == torch.int32
return torch.nn.functional.pad(
torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)
)
@@ -181,6 +212,7 @@ class NativeSparseAttnBackend(AttentionBackend):
global NSA_PREFILL_IMPL, NSA_DECODE_IMPL
NSA_PREFILL_IMPL = model_runner.server_args.nsa_prefill_backend
NSA_DECODE_IMPL = model_runner.server_args.nsa_decode_backend
self.enable_auto_select_prefill_impl = NSA_PREFILL_IMPL == "flashmla_auto"
self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32)
@@ -231,10 +263,16 @@ class NativeSparseAttnBackend(AttentionBackend):
cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32)
assert forward_batch.seq_lens_cpu is not None
max_seqlen_k = int(forward_batch.seq_lens_cpu.max().item() + draft_token_num)
# [b, max_seqlen_k]
page_table = forward_batch.req_to_token_pool.req_to_token[
forward_batch.req_pool_indices, :max_seqlen_k
]
page_table_1_flattened = None
topk_indices_offset = None
self.set_nsa_prefill_impl(forward_batch)
topk_transform_method = self.get_topk_transform_method()
if forward_batch.forward_mode.is_decode_or_idle():
extend_seq_lens_cpu = [1] * batch_size
max_seqlen_q = 1
@@ -295,6 +333,7 @@ class NativeSparseAttnBackend(AttentionBackend):
else:
max_seqlen_q = max_seqlen_k
cu_seqlens_q = cu_seqlens_k
seqlens_expanded = torch.cat(
[
torch.arange(
@@ -310,6 +349,24 @@ class NativeSparseAttnBackend(AttentionBackend):
)
]
)
if topk_transform_method == TopkTransformMethod.RAGGED:
page_table_1_flattened = torch.cat(
[
page_table[i, :kv_len]
for i, kv_len in enumerate(
forward_batch.seq_lens_cpu.tolist(),
)
]
)
assert (
page_table_1_flattened.shape[0] == forward_batch.seq_lens_sum
), f"{page_table_1_flattened.shape[0] = } must be the same as {forward_batch.seq_lens_sum = }"
topk_indices_offset = torch.repeat_interleave(
cu_seqlens_k[:-1],
forward_batch.extend_seq_lens,
)
else:
assert False, f"Unsupported {forward_batch.forward_mode = }"
@@ -328,7 +385,9 @@ class NativeSparseAttnBackend(AttentionBackend):
max_seq_len_k=max_seqlen_k,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
seq_lens_sum=forward_batch.seq_lens_sum,
page_table_1=page_table,
page_table_1_flattened=page_table_1_flattened,
flashmla_metadata=(
self._compute_flashmla_metadata(
cache_seqlens=nsa_cache_seqlens_int32,
@@ -344,6 +403,7 @@ class NativeSparseAttnBackend(AttentionBackend):
nsa_extend_seq_lens_list=extend_seq_lens_cpu,
real_page_table=self._transform_table_1_to_real(page_table),
nsa_max_seqlen_q=1,
topk_indices_offset=topk_indices_offset,
)
self.forward_metadata = metadata
@@ -396,6 +456,8 @@ class NativeSparseAttnBackend(AttentionBackend):
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
):
self.set_nsa_prefill_impl(forward_batch=None)
"""Initialize forward metadata for capturing CUDA graph."""
if forward_mode.is_decode_or_idle():
# Normal Decode
@@ -586,6 +648,8 @@ class NativeSparseAttnBackend(AttentionBackend):
"""Initialize forward metadata for replaying CUDA graph."""
assert seq_lens_cpu is not None
self.set_nsa_prefill_impl(forward_batch=None)
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
req_pool_indices = req_pool_indices[:bs]
@@ -780,17 +844,31 @@ class NativeSparseAttnBackend(AttentionBackend):
q_rope = q_all[:, :, layer.v_head_dim :]
# NOTE(dark): here, we use page size = 1
topk_transform_method = self.get_topk_transform_method()
if NSA_FUSE_TOPK:
page_table_1 = topk_indices
else:
assert metadata.nsa_extend_seq_lens_list is not None
page_table_1 = transform_index_page_table_prefill(
page_table=metadata.page_table_1,
topk_indices=topk_indices,
extend_lens_cpu=metadata.nsa_extend_seq_lens_list,
page_size=1,
)
if topk_transform_method == TopkTransformMethod.RAGGED:
topk_indices_offset = metadata.topk_indices_offset
assert topk_indices_offset is not None
mask = topk_indices != -1
topk_indices_offset = (
topk_indices_offset.unsqueeze(1)
if topk_indices_offset.ndim == 1
else topk_indices_offset
)
topk_indices = torch.where(
mask, topk_indices + topk_indices_offset, topk_indices
)
elif topk_transform_method == TopkTransformMethod.PAGED:
assert metadata.nsa_extend_seq_lens_list is not None
page_table_1 = transform_index_page_table_prefill(
page_table=metadata.page_table_1,
topk_indices=topk_indices,
extend_lens_cpu=metadata.nsa_extend_seq_lens_list,
page_size=1,
)
if NSA_PREFILL_IMPL == "tilelang":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
@@ -804,6 +882,22 @@ class NativeSparseAttnBackend(AttentionBackend):
elif NSA_PREFILL_IMPL == "flashmla_sparse":
if q_rope is not None:
q_all = torch.cat([q_nope, q_rope], dim=-1)
# NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8 has no effect here,
# because the flashmla_sparse kernel doesn't support fp8 compute
if topk_transform_method == TopkTransformMethod.RAGGED:
if any(forward_batch.extend_prefix_lens_cpu):
page_table_1_flattened = (
self.forward_metadata.page_table_1_flattened
)
assert page_table_1_flattened is not None
kv_cache = dequantize_k_cache_paged(
kv_cache, page_table_1_flattened
)
else:
kv_cache = torch.cat([k, k_rope], dim=-1)
page_table_1 = topk_indices
return self._forward_flashmla_sparse(
q_all=q_all,
kv_cache=kv_cache,
@@ -1121,10 +1215,52 @@ class NativeSparseAttnBackend(AttentionBackend):
"""Get the fill value for sequence length in CUDA graph."""
return 1
def set_nsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None) -> str:
from sglang.srt.utils import is_blackwell
global NSA_PREFILL_IMPL
if self.enable_auto_select_prefill_impl:
if self.nsa_kv_cache_store_fp8:
if (
# TODO(hlu1): enable MTP
is_blackwell()
and forward_batch is not None
and forward_batch.forward_mode.is_extend()
and forward_batch.spec_algorithm.is_none()
):
total_kv_tokens = forward_batch.seq_lens_sum
total_q_tokens = forward_batch.extend_num_tokens
# Heuristic based on benchmarking flashmla_kv vs flashmla_sparse + dequantize_k_cache_paged
if total_kv_tokens < total_q_tokens * 512:
NSA_PREFILL_IMPL = "flashmla_sparse"
return
NSA_PREFILL_IMPL = "flashmla_kv"
else:
# bf16 kv cache
NSA_PREFILL_IMPL = "flashmla_sparse"
def get_topk_transform_method(self) -> TopkTransformMethod:
"""
NSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel.
This method is used to select the topk transform method which can be fused or unfused.
"""
if (
# disable for MTP
self.nsa_kv_cache_store_fp8
and NSA_PREFILL_IMPL == "flashmla_sparse"
):
topk_transform_method = TopkTransformMethod.RAGGED
else:
topk_transform_method = TopkTransformMethod.PAGED
return topk_transform_method
def get_indexer_metadata(
self, layer_id: int, forward_batch: ForwardBatch
) -> NSAIndexerMetadata:
return NSAIndexerMetadata(attn_metadata=self.forward_metadata)
return NSAIndexerMetadata(
attn_metadata=self.forward_metadata,
topk_transform_method=self.get_topk_transform_method(),
)
def _compute_flashmla_metadata(self, cache_seqlens: torch.Tensor, seq_len_q: int):
from flash_mla import get_mla_metadata

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 (