support ds4

lightllm/common/basemodel/attention/create_utils.py

@@ -130,21 +130,25 @@ def get_mla_decode_att_backend_class(index=0, priority_list: list = ["flashinfer
         return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list)
 
 
-def get_nsa_prefill_att_backend_class(index=0, priority_list: list = ["flashmla_sparse"]) -> BaseAttBackend:
+def get_nsa_prefill_att_backend_class(
+    index=0, priority_list: list = ["flashmla_sparse"], backend_map=nsa_data_type_to_backend
+) -> BaseAttBackend:
     args = get_env_start_args()
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_prefill_att_backend[index]
     if backend_str != "auto":
-        return nsa_data_type_to_backend[llm_dtype][backend_str]
+        return backend_map[llm_dtype][backend_str]
     else:
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=nsa_data_type_to_backend, priority_list=priority_list)
+        return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)
 
 
-def get_nsa_decode_att_backend_class(index=0, priority_list: list = ["flashmla_sparse"]) -> BaseAttBackend:
+def get_nsa_decode_att_backend_class(
+    index=0, priority_list: list = ["flashmla_sparse"], backend_map=nsa_data_type_to_backend
+) -> BaseAttBackend:
     args = get_env_start_args()
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_decode_att_backend[index]
     if backend_str != "auto":
-        return nsa_data_type_to_backend[llm_dtype][backend_str]
+        return backend_map[llm_dtype][backend_str]
     else:
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=nsa_data_type_to_backend, priority_list=priority_list)
+        return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)

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

@@ -0,0 +1,186 @@
+import dataclasses
+from typing import TYPE_CHECKING
+
+import torch
+
+from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState
+
+if TYPE_CHECKING:
+    from lightllm.common.basemodel.infer_struct import InferStateInfo
+
+
+# The current FlashMLA MODEL1 binary only instantiates these Q-head counts.
+_SUPPORTED_Q_HEADS = (64, 128)
+
+
+def get_dsv4_flashmla_padded_q_heads(q_head_num: int) -> int:
+    for supported_head_num in _SUPPORTED_Q_HEADS:
+        if q_head_num <= supported_head_num:
+            return supported_head_num
+    raise ValueError(f"FlashMLA does not support {q_head_num} local Q heads; supported counts: {_SUPPORTED_Q_HEADS}")
+
+
+def _view_cache(buffer: torch.Tensor, page_size: int) -> torch.Tensor:
+    from lightllm.common.kv_cache_mem_manager.deepseek4_mem_manager import DSV4_MLA_BYTES_PER_TOKEN
+
+    byte_num = page_size * DSV4_MLA_BYTES_PER_TOKEN
+    return buffer[:, :byte_num].view(buffer.shape[0], page_size, 1, DSV4_MLA_BYTES_PER_TOKEN)
+
+
+def _get_vllm_flashmla():
+    from vllm.v1.attention.ops import flashmla
+
+    return flashmla
+
+
+class DeepseekV4FlashMlaFp8SparseAttBackend(BaseAttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+        self.real_q_head_num = model.config["num_attention_heads"] // model.tp_world_size_
+        self.padded_q_head_num = get_dsv4_flashmla_padded_q_heads(self.real_q_head_num)
+        self.compress_ratios = tuple(dict.fromkeys(model.config["compress_ratios"]))
+
+    def _flashmla_att(
+        self,
+        q: torch.Tensor,
+        packed_kv: torch.Tensor,
+        mem_manager,
+        nsa_dict: dict,
+        sched_meta,
+        flash_mla,
+        flashmla_out: torch.Tensor = None,
+    ) -> torch.Tensor:
+        from lightllm.common.kv_cache_mem_manager.deepseek4_mem_manager import (
+            DSV4_C128_PAGE_SIZE,
+            DSV4_C4_PAGE_SIZE,
+            DSV4_SWA_PAGE_SIZE,
+        )
+
+        ratio = nsa_dict["compress_ratio"]
+        extra_cache = None
+        if ratio == 4:
+            extra_page_size = DSV4_C4_PAGE_SIZE
+        elif ratio == 128:
+            extra_page_size = DSV4_C128_PAGE_SIZE
+        elif ratio != 0:
+            raise ValueError(f"unsupported DeepSeek-V4 compress ratio: {ratio}")
+        if ratio:
+            buffer = mem_manager.get_compressed_kv_buffer(nsa_dict["layer_index"])
+            extra_cache = _view_cache(buffer, extra_page_size)
+
+        kwargs = dict(
+            q=q.unsqueeze(1),
+            k_cache=_view_cache(packed_kv, DSV4_SWA_PAGE_SIZE),
+            block_table=None,
+            cache_seqlens=None,
+            head_dim_v=nsa_dict["head_dim_v"],
+            tile_scheduler_metadata=sched_meta,
+            num_splits=None,
+            softmax_scale=nsa_dict["softmax_scale"],
+            causal=False,
+            is_fp8_kvcache=True,
+            indices=nsa_dict["swa_indices"],
+            attn_sink=nsa_dict["attn_sink"],
+            topk_length=nsa_dict["swa_lengths"],
+            extra_k_cache=extra_cache,
+            extra_indices_in_kvcache=nsa_dict.get("extra_indices"),
+            extra_topk_length=nsa_dict.get("extra_lengths"),
+        )
+        if flashmla_out is not None:
+            kwargs["out"] = flashmla_out
+        full_out, _ = flash_mla.flash_mla_with_kvcache(**kwargs)
+        return full_out[:, 0, : self.real_q_head_num, :]
+
+    def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "_PrefillAttState":
+        return _PrefillAttState(backend=self, infer_state=infer_state)
+
+    def create_att_decode_state(self, infer_state: "InferStateInfo") -> "_DecodeAttState":
+        return _DecodeAttState(backend=self, infer_state=infer_state)
+
+
+@dataclasses.dataclass
+class _PrefillAttState(BasePrefillAttState):
+    flashmla_sched_meta: dict = None
+    flash_mla: object = None
+
+    def init_state(self):
+        self.flash_mla = _get_vllm_flashmla()
+        self.flashmla_sched_meta = {}
+
+    def _get_sched_meta(self, compress_ratio: int):
+        if compress_ratio not in self.flashmla_sched_meta:
+            self.flashmla_sched_meta[compress_ratio] = self.flash_mla.get_mla_metadata()[0]
+        return self.flashmla_sched_meta[compress_ratio]
+
+    def prefill_att(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        att_control: AttControl = AttControl(),
+        alloc_func=torch.empty,
+        *,
+        out: torch.Tensor = None,
+    ) -> 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"
+        nsa_dict = att_control.nsa_prefill_dict
+        if out is None:
+            out = alloc_func(
+                (q.shape[0], self.backend.real_q_head_num, nsa_dict["head_dim_v"]),
+                dtype=q.dtype,
+                device=q.device,
+            )
+        full_out = self.infer_state.dsv4_workspace.flashmla_prefill_full_out[: q.shape[0]]
+        out.copy_(
+            self.backend._flashmla_att(
+                q,
+                k,
+                self.infer_state.mem_manager,
+                nsa_dict,
+                self._get_sched_meta(nsa_dict["compress_ratio"]),
+                self.flash_mla,
+                flashmla_out=full_out,
+            )
+        )
+        return out
+
+
+@dataclasses.dataclass
+class _DecodeAttState(BaseDecodeAttState):
+    flashmla_sched_meta: dict = None
+    flash_mla: object = None
+
+    def init_state(self):
+        self.reset_sched_meta_for_capture()
+
+    def reset_sched_meta_for_capture(self):
+        # FlashMLA lazily binds extra-cache geometry, so ratios cannot share one sched-meta object.
+        self.flash_mla = _get_vllm_flashmla()
+        self.flashmla_sched_meta = {
+            ratio: self.flash_mla.get_mla_metadata()[0] for ratio in self.backend.compress_ratios
+        }
+
+    def decode_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_decode, "nsa_decode must be True for NSA decode attention"
+        assert att_control.nsa_decode_dict is not None, "nsa_decode_dict is required"
+        nsa_dict = att_control.nsa_decode_dict
+        real_out = self.backend._flashmla_att(
+            q,
+            k,
+            self.infer_state.mem_manager,
+            nsa_dict,
+            self.flashmla_sched_meta[nsa_dict["compress_ratio"]],
+            self.flash_mla,
+        )
+        return real_out.contiguous()
+
+
+DSV4_NSA_BACKENDS = {"fp8kv_dsa": {"flashmla_sparse": DeepseekV4FlashMlaFp8SparseAttBackend}}

lightllm/common/basemodel/basemodel.py

@@ -42,6 +42,8 @@
 
 
 class TpPartBaseModel:
+    is_mtp_draft_model = False
+
     # weight class
     pre_and_post_weight_class = None
     transformer_weight_class = None
@@ -554,8 +556,10 @@ def _decode(
         else:
             infer_batch_size = model_input.batch_size
 
-        if self.graph is not None and self.graph.can_run(
-            batch_size=infer_batch_size, max_len_in_batch=model_input.max_kv_seq_len
+        if (
+            self.graph is not None
+            and not self.is_mtp_draft_model
+            and self.graph.can_run(batch_size=infer_batch_size, max_len_in_batch=model_input.max_kv_seq_len)
         ):
             infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size)
             model_input = self._create_padded_decode_model_input(
@@ -600,6 +604,7 @@ def _decode(
     def _context_forward(self, infer_state: InferStateInfo):
 
         input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
+        infer_state.mtp_draft_input_hiddens = None
         if self.args.enable_dp_prefill_balance:
             assert not self.args.enable_prefill_cudagraph, "not support now"
             infer_state.prepare_prefill_dp_balance()
@@ -645,14 +650,20 @@ def prefill_func(input_tensors, infer_state):
             last_input_embs = infer_state._all_to_all_unbalance_get(data=last_input_embs)
 
         predict_logits = self.post_infer.token_forward(last_input_embs, infer_state, self.pre_post_weight)
+        mtp_main_output_hiddens = None
+        if isinstance(predict_logits, tuple):
+            predict_logits, mtp_main_output_hiddens = predict_logits
         model_output = ModelOutput(logits=predict_logits)
 
         # 特殊模型特殊模式的额外输出
         if self.is_mtp_mode:
-            input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
-            if infer_state.need_dp_prefill_balance:
-                input_embs = infer_state._all_to_all_unbalance_get(data=input_embs)
-            model_output.mtp_main_output_hiddens = input_embs.contiguous()
+            if mtp_main_output_hiddens is not None:
+                model_output.mtp_main_output_hiddens = mtp_main_output_hiddens.contiguous()
+            else:
+                input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+                if infer_state.need_dp_prefill_balance:
+                    input_embs = infer_state._all_to_all_unbalance_get(data=input_embs)
+                model_output.mtp_main_output_hiddens = input_embs.contiguous()
 
         # 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候
         # 该调用没有实际意义
@@ -664,23 +675,30 @@ def _token_forward(self, infer_state: InferStateInfo):
         input_ids = infer_state.input_ids
         cuda_input_ids = input_ids
         input_embs = self.pre_infer.token_forward(cuda_input_ids, infer_state, self.pre_post_weight)
+        infer_state.mtp_draft_input_hiddens = None
         input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state)
 
         for i in range(self.layers_num):
             layer = self.layers_infer[i]
             input_embs: torch.Tensor = layer.token_forward(input_embs, infer_state, self.trans_layers_weight[i])
 
         last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
-        predict_logits: torch.Tensor = self.post_infer.token_forward(
+        predict_logits = self.post_infer.token_forward(
             last_input_embs, infer_state=infer_state, layer_weight=self.pre_post_weight
         )
+        mtp_main_output_hiddens = None
+        if isinstance(predict_logits, tuple):
+            predict_logits, mtp_main_output_hiddens = predict_logits
 
         model_output = ModelOutput(logits=predict_logits.contiguous())
 
         # 特殊模型特殊模式的额外输出
         if self.is_mtp_mode:
-            input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
-            model_output.mtp_main_output_hiddens = input_embs.contiguous()
+            if mtp_main_output_hiddens is not None:
+                model_output.mtp_main_output_hiddens = mtp_main_output_hiddens.contiguous()
+            else:
+                input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+                model_output.mtp_main_output_hiddens = input_embs.contiguous()
 
         # 在 cuda graph 模式下，输出需要转为 no ref tensor, 加强mem pool 的复用，降低显存的使用。
         if infer_state.is_cuda_graph:
@@ -917,18 +935,31 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state
             last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight
         )
         g_cache_manager.cache_env_out()
+        mtp_main_output_hiddens = None
+        mtp_main_output_hiddens1 = None
+        if isinstance(predict_logits, tuple):
+            predict_logits, mtp_main_output_hiddens = predict_logits
+        if isinstance(predict_logits1, tuple):
+            predict_logits1, mtp_main_output_hiddens1 = predict_logits1
 
         model_output = ModelOutput(logits=predict_logits.contiguous())
         model_output1 = ModelOutput(logits=predict_logits1.contiguous())
 
         if self.is_mtp_mode:
-            input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
-            input_embs1 = self.pre_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
-            if infer_state.need_dp_prefill_balance:
-                input_embs = infer_state._all_to_all_unbalance_get(data=input_embs)
-                input_embs1 = infer_state1._all_to_all_unbalance_get(data=input_embs1)
-            model_output.mtp_main_output_hiddens = input_embs.contiguous()
-            model_output1.mtp_main_output_hiddens = input_embs1.contiguous()
+            if mtp_main_output_hiddens is not None:
+                model_output.mtp_main_output_hiddens = mtp_main_output_hiddens.contiguous()
+            else:
+                input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+                if infer_state.need_dp_prefill_balance:
+                    input_embs = infer_state._all_to_all_unbalance_get(data=input_embs)
+                model_output.mtp_main_output_hiddens = input_embs.contiguous()
+            if mtp_main_output_hiddens1 is not None:
+                model_output1.mtp_main_output_hiddens = mtp_main_output_hiddens1.contiguous()
+            else:
+                input_embs1 = self.pre_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
+                if infer_state.need_dp_prefill_balance:
+                    input_embs1 = infer_state1._all_to_all_unbalance_get(data=input_embs1)
+                model_output1.mtp_main_output_hiddens = input_embs1.contiguous()
 
         return model_output, model_output1
 
@@ -955,15 +986,27 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1:
         predict_logits, predict_logits1 = self.post_infer.overlap_tpsp_token_forward(
             last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight
         )
+        mtp_main_output_hiddens = None
+        mtp_main_output_hiddens1 = None
+        if isinstance(predict_logits, tuple):
+            predict_logits, mtp_main_output_hiddens = predict_logits
+        if isinstance(predict_logits1, tuple):
+            predict_logits1, mtp_main_output_hiddens1 = predict_logits1
 
         model_output = ModelOutput(logits=predict_logits.contiguous())
         model_output1 = ModelOutput(logits=predict_logits1.contiguous())
 
         if self.is_mtp_mode:
-            input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
-            input_embs1 = self.pre_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
-            model_output.mtp_main_output_hiddens = input_embs.contiguous()
-            model_output1.mtp_main_output_hiddens = input_embs1.contiguous()
+            if mtp_main_output_hiddens is not None:
+                model_output.mtp_main_output_hiddens = mtp_main_output_hiddens.contiguous()
+            else:
+                input_embs = self.pre_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+                model_output.mtp_main_output_hiddens = input_embs.contiguous()
+            if mtp_main_output_hiddens1 is not None:
+                model_output1.mtp_main_output_hiddens = mtp_main_output_hiddens1.contiguous()
+            else:
+                input_embs1 = self.pre_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
+                model_output1.mtp_main_output_hiddens = input_embs1.contiguous()
 
         if infer_state.is_cuda_graph:
             model_output.to_no_ref_tensor()
@@ -1171,13 +1214,7 @@ def _init_padded_req(self):
     def _gen_special_model_input(self, token_num: int):
         special_model_input = {}
 
-        is_mtp_draft_model = (
-            "Deepseek3MTPModel" in str(self.__class__)
-            or "Qwen3MOEMTPModel" in str(self.__class__)
-            or "MistralMTPModel" in str(self.__class__)
-            or "Glm4MoeLiteMTPModel" in str(self.__class__)
-        )
-        if is_mtp_draft_model:
+        if self.is_mtp_draft_model:
             special_model_input["mtp_draft_input_hiddens"] = torch.randn(
                 token_num, self.config["hidden_size"], dtype=self.data_type, device="cuda"
             )