fix: fix bugs in the ep v2 and mega moe for sm100

lightllm/common/basemodel/attention/create_utils.py

@@ -141,21 +141,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,172 @@
+import dataclasses
+from typing import TYPE_CHECKING
+
+import torch
+from vllm.v1.attention.ops import flashmla
+
+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)
+
+
+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,
+        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, _ = flashmla.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
+
+    def init_state(self):
+        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] = flashmla.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"]),
+                flashmla_out=full_out,
+            )
+        )
+        return out
+
+
+@dataclasses.dataclass
+class _DecodeAttState(BaseDecodeAttState):
+    flashmla_sched_meta: dict = 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.flashmla_sched_meta = {ratio: flashmla.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"]],
+        )
+        return real_out.contiguous()
+
+
+DSV4_NSA_BACKENDS = {"fp8kv_dsa": {"flashmla_sparse": DeepseekV4FlashMlaFp8SparseAttBackend}}

lightllm/common/basemodel/basemodel.py

@@ -49,6 +49,8 @@
 
 
 class TpPartBaseModel:
+    is_mtp_draft_model = False
+
     # weight class
     pre_and_post_weight_class = None
     transformer_weight_class = None
@@ -168,7 +170,13 @@ def _verify_params(self):
         return
 
     def _init_quant(self):
-        self.quant_cfg = Quantcfg(self.config, self.quant_type, self.quant_cfg_path, self.expert_dtype)
+        self.quant_cfg = Quantcfg(
+            self.config,
+            self.quant_type,
+            self.quant_cfg_path,
+            self.expert_dtype,
+            enable_ep_moe=self.args.enable_ep_moe,
+        )
         logger.info(f"Initial quantization. " f"The default quantization method is {self.quant_cfg.quant_type}")
 
     def _init_weights(self, start_layer_index=0):
@@ -618,8 +626,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(
@@ -665,6 +675,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()
@@ -710,14 +721,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, prompt_logics=infer_state.prompt_logics)
 
         # 特殊模型特殊模式的额外输出
         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做资源清理，没有启用的时候
         # 该调用没有实际意义
@@ -729,23 +746,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:
@@ -982,18 +1006,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(), prompt_logics=infer_state.prompt_logics)
         model_output1 = ModelOutput(logits=predict_logits1.contiguous(), prompt_logics=infer_state1.prompt_logics)
 
         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
 
@@ -1020,15 +1057,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()