feat: fp8 dsa support

Author: 钮圣虓

docker/Dockerfile

@@ -4,6 +4,7 @@ FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu22.04
 ARG PYTHON_VERSION=3.10
 ARG MAMBA_VERSION=24.7.1-0
 ARG VLLM_VERSION=0.16.0
+ARG FLASH_MLA_REF=47c35a7
 ARG TARGETPLATFORM
 ARG ENABLE_DEEPEP=1
 ARG ENABLE_NIXL=1
@@ -45,6 +46,11 @@ COPY ./requirements.txt /lightllm/requirements.txt
 RUN pip install -U pip
 RUN pip install -r /lightllm/requirements.txt --no-cache-dir
 RUN pip install --no-cache-dir vllm==${VLLM_VERSION}
+RUN git clone https://github.com/deepseek-ai/FlashMLA.git /root/FlashMLA && \
+    cd /root/FlashMLA && \
+    git checkout ${FLASH_MLA_REF} && \
+    git submodule update --init --recursive && \
+    FLASH_MLA_DISABLE_SM100=1 pip install --no-cache-dir .
 
 RUN apt-get update && apt-get install -y libnuma-dev && rm -rf /var/lib/apt/lists/*
 

lightllm/common/basemodel/attention/__init__.py

@@ -12,6 +12,7 @@
 
 # NSA backend
 from .nsa.flashmla_sparse import NsaFlashMlaSparseAttBackend
+from .nsa.fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend
 
 from .create_utils import (
     get_prefill_att_backend_class,

lightllm/common/basemodel/attention/create_utils.py

@@ -15,6 +15,7 @@
 from .flashinfer.fp import FlashInferAttBackend
 from .flashinfer.mla import MlaFlashInferAttBackend
 from .nsa.flashmla_sparse import NsaFlashMlaSparseAttBackend
+from .nsa.fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend
 
 logger = init_logger(__name__)
 
@@ -56,6 +57,9 @@
         "flashmla_sparse": NsaFlashMlaSparseAttBackend,
         # Future backends: "fa3", "tilelang", "aiter"
     },
+    "fp8kv_dsa": {
+        "flashmla_sparse": NsaFlashMlaFp8SparseAttBackend,
+    },
 }
 
 

lightllm/common/basemodel/attention/nsa/__init__.py

@@ -5,9 +5,17 @@
     NsaFlashMlaSparsePrefillAttState,
     NsaFlashMlaSparseDecodeAttState,
 )
+from .fp8_flashmla_sparse import (
+    NsaFlashMlaFp8SparseAttBackend,
+    NsaFlashMlaFp8SparsePrefillAttState,
+    NsaFlashMlaFp8SparseDecodeAttState,
+)
 
 __all__ = [
     "NsaFlashMlaSparseAttBackend",
     "NsaFlashMlaSparsePrefillAttState",
     "NsaFlashMlaSparseDecodeAttState",
+    "NsaFlashMlaFp8SparseAttBackend",
+    "NsaFlashMlaFp8SparsePrefillAttState",
+    "NsaFlashMlaFp8SparseDecodeAttState",
 ]

lightllm/common/basemodel/attention/nsa/flashmla_sparse.py

@@ -165,7 +165,7 @@ def _nsa_decode_att(
         from sgl_kernel.flash_attn import flash_attn_with_kvcache
 
         nsa_dict = att_control.nsa_decode_dict
-        topk_indices = nsa_dict["topk_indices"]
+        topk_mem_indices = nsa_dict["topk_mem_indices"]
         softmax_scale = nsa_dict["softmax_scale"]
         kv_lora_rank = nsa_dict["kv_lora_rank"]
         qk_rope_head_dim = nsa_dict["qk_rope_head_dim"]
@@ -181,7 +181,7 @@ def _nsa_decode_att(
             k_cache=k_rope,
             v_cache=kv_nope,
             qv=q_nope,
-            page_table=topk_indices,
+            page_table=topk_mem_indices,
             cache_seqlens=self.nsa_cache_seqlens,
             cu_seqlens_q=self.infer_state.b1_cu_q_seq_len,
             cu_seqlens_k_new=self.nsa_cu_seqlens_k_new,

lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py

@@ -0,0 +1,196 @@
+import dataclasses
+import torch
+from typing import TYPE_CHECKING, Tuple
+
+from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState
+from lightllm.utils.dist_utils import get_current_device_id
+
+if TYPE_CHECKING:
+    from lightllm.common.basemodel.infer_struct import InferStateInfo
+
+
+class NsaFlashMlaFp8SparseAttBackend(BaseAttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+        device = get_current_device_id()
+        self.ragged_mem_buffers = [
+            torch.empty(model.graph_max_batch_size * model.max_seq_length, dtype=torch.int32, device=device)
+            for _ in range(2)
+        ]
+
+    def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaFp8SparsePrefillAttState":
+        return NsaFlashMlaFp8SparsePrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaFp8SparseDecodeAttState":
+        return NsaFlashMlaFp8SparseDecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class NsaFlashMlaFp8SparsePrefillAttState(BasePrefillAttState):
+    ks: torch.Tensor = None
+    ke: torch.Tensor = None
+    lengths: torch.Tensor = None
+    ragged_mem_index: torch.Tensor = None
+
+    def init_state(self):
+        self.backend: NsaFlashMlaFp8SparseAttBackend = self.backend
+        self.ragged_mem_index = torch.empty(
+            self.infer_state.total_token_num,
+            dtype=torch.int32,
+            device=get_current_device_id(),
+        )
+        from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke
+
+        self.ks, self.ke, self.lengths = gen_nsa_ks_ke(
+            b_seq_len=self.infer_state.b_seq_len,
+            b_q_seq_len=self.infer_state.b_q_seq_len,
+            b_req_idx=self.infer_state.b_req_idx,
+            req_to_token_index=self.infer_state.req_manager.req_to_token_indexs,
+            q_token_num=self.infer_state.total_token_num - self.infer_state.prefix_total_token_num,
+            ragged_mem_index=self.ragged_mem_index,
+            hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID,
+        )
+        return
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert att_control.nsa_prefill, "nsa_prefill must be True for NSA prefill attention"
+        assert att_control.nsa_prefill_dict is not None, "nsa_prefill_dict is required"
+        return self._nsa_prefill_att(q=q, packed_kv=k, att_control=att_control)
+
+    def _nsa_prefill_att(
+        self,
+        q: torch.Tensor,
+        packed_kv: torch.Tensor,
+        att_control: AttControl,
+    ) -> torch.Tensor:
+        import flash_mla
+
+        nsa_dict = att_control.nsa_prefill_dict
+        topk_indices = nsa_dict["topk_indices"]
+        softmax_scale = nsa_dict["softmax_scale"]
+        kv_lora_rank = nsa_dict["kv_lora_rank"]
+        topk_mem_indices = nsa_dict["topk_mem_indices"]
+        prefill_cache_kv = nsa_dict["prefill_cache_kv"]
+
+        if self.infer_state.prefix_total_token_num > 0:
+            kv, topk_indices = self.infer_state.mem_manager.get_prefill_kv_cache_and_remap_indices(
+                packed_kv=packed_kv,
+                topk_indices=topk_mem_indices,
+                prefill_mem_index=self.infer_state.mem_index,
+                prefill_cache_kv=prefill_cache_kv,
+            )
+        else:
+            kv = prefill_cache_kv
+
+        if topk_indices.ndim == 2:
+            topk_indices = topk_indices.unsqueeze(1)
+
+        mla_out, _, _ = flash_mla.flash_mla_sparse_fwd(
+            q=q,
+            kv=kv,
+            indices=topk_indices,
+            sm_scale=softmax_scale,
+            d_v=kv_lora_rank,
+        )
+        return mla_out
+
+
+@dataclasses.dataclass
+class NsaFlashMlaFp8SparseDecodeAttState(BaseDecodeAttState):
+    ks: torch.Tensor = None
+    ke: torch.Tensor = None
+    lengths: torch.Tensor = None
+    ragged_mem_index: torch.Tensor = None
+    flashmla_sched_meta: object = None
+
+    def init_state(self):
+        self.backend: NsaFlashMlaFp8SparseAttBackend = self.backend
+        model = self.backend.model
+        use_cuda_graph = (
+            self.infer_state.batch_size <= model.graph_max_batch_size
+            and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch
+        )
+
+        if use_cuda_graph:
+            self.ragged_mem_index = self.backend.ragged_mem_buffers[self.infer_state.microbatch_index]
+        else:
+            self.ragged_mem_index = torch.empty(
+                self.infer_state.total_token_num,
+                dtype=torch.int32,
+                device=get_current_device_id(),
+            )
+
+        from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke
+
+        self.ks, self.ke, self.lengths = gen_nsa_ks_ke(
+            b_seq_len=self.infer_state.b_seq_len,
+            b_q_seq_len=self.infer_state.b_q_seq_len,
+            b_req_idx=self.infer_state.b_req_idx,
+            req_to_token_index=self.infer_state.req_manager.req_to_token_indexs,
+            q_token_num=self.infer_state.b_seq_len.shape[0],
+            ragged_mem_index=self.ragged_mem_index,
+            hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID,
+        )
+        import flash_mla
+
+        self.flashmla_sched_meta, _ = flash_mla.get_mla_metadata()
+        return
+
+    def decode_att(
+        self,
+        q: Tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+    ) -> torch.Tensor:
+        assert att_control.nsa_decode, "nsa_decode must be True for NSA decode attention"
+        assert att_control.nsa_decode_dict is not None, "nsa_decode_dict is required"
+        return self._nsa_decode_att(q=q, packed_kv=k, att_control=att_control)
+
+    def _nsa_decode_att(
+        self,
+        q: Tuple[torch.Tensor, torch.Tensor],
+        packed_kv: torch.Tensor,
+        att_control: AttControl,
+    ) -> torch.Tensor:
+        import flash_mla
+
+        nsa_dict = att_control.nsa_decode_dict
+        topk_mem_indices = nsa_dict["topk_mem_indices"]
+        softmax_scale = nsa_dict["softmax_scale"]
+        kv_lora_rank = nsa_dict["kv_lora_rank"]
+
+        if topk_mem_indices.ndim == 2:
+            topk_mem_indices = topk_mem_indices.unsqueeze(1)
+        assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1"
+
+        q_nope, q_rope = q
+        q_all = torch.cat([q_nope, q_rope], dim=-1).unsqueeze(1).contiguous()
+        kv = torch.as_strided(
+            packed_kv,
+            size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]),
+            stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)),
+        )
+
+        o_tensor, _ = flash_mla.flash_mla_with_kvcache(
+            q=q_all,
+            k_cache=kv,
+            block_table=None,
+            cache_seqlens=None,
+            head_dim_v=kv_lora_rank,
+            tile_scheduler_metadata=self.flashmla_sched_meta,
+            num_splits=None,
+            softmax_scale=softmax_scale,
+            causal=False,
+            is_fp8_kvcache=True,
+            indices=topk_mem_indices,
+        )
+        return o_tensor[:, 0, :, :]  # [b, 1, h, d] -> [b, h, d]

lightllm/common/kv_cache_mem_manager/__init__.py

@@ -3,6 +3,7 @@
 from .ppl_int4kv_mem_manager import PPLINT4KVMemoryManager
 from .deepseek2_mem_manager import Deepseek2MemoryManager
 from .deepseek3_2mem_manager import Deepseek3_2MemoryManager
+from .fp8_per_token_group_quant_deepseek3_2mem_manager import FP8PerTokenGroupQuantDeepseek3_2MemoryManager
 from .fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager
 from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager
 
@@ -13,6 +14,7 @@
     "PPLINT8KVMemoryManager",
     "Deepseek2MemoryManager",
     "Deepseek3_2MemoryManager",
+    "FP8PerTokenGroupQuantDeepseek3_2MemoryManager",
     "FP8StaticPerHeadQuantMemManager",
     "FP8StaticPerTensorQuantMemManager",
 ]

lightllm/common/kv_cache_mem_manager/deepseek3_2mem_manager.py

@@ -34,3 +34,6 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv:
     def get_att_input_params(self, layer_index: int) -> Any:
         kv = self.kv_buffer[layer_index][:, :, : (self.head_dim - (144 // 2))]
         return kv
+
+    def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor:
+        return self.kv_buffer[layer_index].view(dtype=torch.uint8)[:, :, -132:]

lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py

@@ -0,0 +1,83 @@
+import torch
+from typing import Any
+
+from .deepseek2_mem_manager import Deepseek2MemoryManager
+
+
+class FP8PerTokenGroupQuantDeepseek3_2MemoryManager(Deepseek2MemoryManager):
+    kv_nope_dim = 512
+    kv_rope_dim = 64
+    # 576 = 512 + 64
+    kv_head_dim = kv_nope_dim + kv_rope_dim
+
+    quant_group_size = 128
+    # 4 = 512 / 128
+    quant_group_num = kv_nope_dim // quant_group_size
+    # 4 * 4 = quant_group_num * fp32
+    # 64 * 2 = kv_rope_dim * bfloat16
+    # 656 bytes = 512 + (4 * 4) + (64 * 2)
+    flashmla_bytes_per_token = kv_nope_dim + quant_group_num * 4 + kv_rope_dim * 2
+
+    indexer_head_dim = 128
+    # 128 + 4 = indexer_head_dim + fp32
+    # 132 bytes = 128 + 4
+    indexer_bytes_per_token = indexer_head_dim + 4
+
+    # 16-byte 对齐，满足FlashMLA的对齐要求
+    alignment = 16
+    total_bytes_per_token = (
+        (flashmla_bytes_per_token + indexer_bytes_per_token + alignment - 1) // alignment * alignment
+    )
+
+    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
+        assert head_num == 1, "DeepSeek-V3.2 DSA FP8 path expects MQA-style head_num == 1"
+        self.prefill_dtype = dtype
+        super().__init__(size, torch.uint8, head_num, self.total_bytes_per_token, layer_num, always_copy, mem_fraction)
+
+    def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor):
+        from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_kv_flashmla_fp8 import (
+            destindex_copy_kv_flashmla_fp8,
+        )
+
+        rope_dim = 64
+        kv_lora_rank = kv.shape[2] - rope_dim
+        assert kv_lora_rank == 512, f"Expected kv_lora_rank=512, got {kv_lora_rank}"
+
+        o_nope = self.kv_buffer[layer_index][:, :, :512].view(torch.float8_e4m3fn)
+        o_scale = self.kv_buffer[layer_index][:, :, 512:528].view(torch.float32)
+        o_rope = self.kv_buffer[layer_index][:, :, 528 : self.flashmla_bytes_per_token].view(torch.bfloat16)
+        destindex_copy_kv_flashmla_fp8(
+            kv[:, :, :kv_lora_rank],
+            kv[:, :, kv_lora_rank:],
+            mem_index,
+            o_nope,
+            o_scale,
+            o_rope,
+        )
+
+    def get_att_input_params(self, layer_index: int) -> Any:
+        return self.kv_buffer[layer_index][:, :, : self.flashmla_bytes_per_token]
+
+    def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor:
+        begin = self.flashmla_bytes_per_token
+        end = begin + self.indexer_bytes_per_token
+        return self.kv_buffer[layer_index][:, :, begin:end]
+
+    def get_prefill_kv_cache_and_remap_indices(
+        self,
+        packed_kv: torch.Tensor,
+        topk_indices: torch.Tensor,
+        prefill_mem_index: torch.Tensor,
+        prefill_cache_kv: torch.Tensor,
+    ):
+        from lightllm.models.deepseek3_2.triton_kernel.prefill_compact_kv_flashmla_fp8 import (
+            get_prefill_kv_cache_and_remap_indices_triton,
+        )
+
+        return get_prefill_kv_cache_and_remap_indices_triton(
+            packed_kv=packed_kv,
+            topk_mem_indices=topk_indices,
+            prefill_mem_index=prefill_mem_index,
+            prefill_cache_kv=prefill_cache_kv,
+            prefill_dtype=self.prefill_dtype,
+        )