perf(glm): skip unused NextN prefill computation

lightllm/models/glm5_next/indexer.py

@@ -41,7 +41,7 @@ def select_topk_indices(self, logits, q_pool_seq_len, topk_pool_indices):
             selected = torch.topk(logits, topk_pool_indices.shape[1], dim=-1, sorted=True).indices
             topk_pool_indices.copy_(selected.masked_fill(selected >= q_pool_seq_len[:, None], -1))
 
-    def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight):
+    def write_k_pools(self, hidden_states, infer_state, att_state, layer_weight):
         k = layer_weight.k_norm_(layer_weight.wk_proj_.mm(hidden_states), eps=self.rms_norm_eps)
         gate = layer_weight.index_kpool_compress_gate.mm(hidden_states)
         k_and_gate = torch.cat((k, gate), -1)
@@ -71,6 +71,10 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig
             max_q_seq_len=infer_state.max_q_seq_len,
             b_mtp_index=None if infer_state.is_prefill else infer_state.b_mtp_index,
         )
+        return indexer_k
+
+    def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight):
+        indexer_k = self.write_k_pools(hidden_states, infer_state, att_state, layer_weight)
 
         if infer_state.max_kv_seq_len <= self.index_topk:
             return expand_topk(

lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py

@@ -129,7 +129,7 @@ def _moe_ffn(self, input, infer_state, layer_weight) -> torch.Tensor:
 
         return output.view(num_tokens, hidden_dim)
 
-    def _get_qkv(self, input, infer_state, layer_weight):
+    def _project_sparse_qkv_a(self, input, infer_state, layer_weight):
         if self.is_linear_attention_layer:
             raise AssertionError("KDA projections use _kda_projections")
 
@@ -139,17 +139,24 @@ def _get_qkv(self, input, infer_state, layer_weight):
             input = infer_state._all_to_all_unbalance_get(data=input)
 
         q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split([self.q_lora_rank, self.kv_lora_rank], dim=-1)
-        q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_)
-        infer_state.get_topk_indices_params = {"hidden_states": input, "q_lora": q}
-        q = layer_weight.q_b_proj_.mm(q).view(-1, self.tp_q_head_num_, self.qk_nope_head_dim)
+        return input, q, cache_kv
+
+    def _normalize_sparse_kv(self, cache_kv, layer_weight):
         cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank)
         rmsnorm_forward(
             cache_kv[:, :, : self.kv_lora_rank],
             weight=layer_weight.kv_a_layernorm_.weight,
             eps=self.eps_,
             out=cache_kv[:, :, : self.kv_lora_rank],
         )
-        return q, cache_kv
+        return cache_kv
+
+    def _get_qkv(self, input, infer_state, layer_weight):
+        input, q, cache_kv = self._project_sparse_qkv_a(input, infer_state, layer_weight)
+        q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_)
+        infer_state.get_topk_indices_params = {"hidden_states": input, "q_lora": q}
+        q = layer_weight.q_b_proj_.mm(q).view(-1, self.tp_q_head_num_, self.qk_nope_head_dim)
+        return q, self._normalize_sparse_kv(cache_kv, layer_weight)
 
     def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None):
         q = layer_weight.k_b_proj_.bmm(q.transpose(0, 1)).transpose(0, 1).contiguous()

lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py

@@ -0,0 +1,11 @@
+from lightllm.common.basemodel.batch_objs import PostLayerOutput
+from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector
+from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
+
+
+class Glm5NextMTPPostLayerInfer(LlamaPostLayerInfer):
+    def token_forward(self, input_embeddings, infer_state, layer_weight):
+        if infer_state.is_prefill and input_embeddings.shape[0] == 0:
+            infer_state.hidden_collector = NoopHiddenCollector()
+            return PostLayerOutput(logits=input_embeddings.new_empty((0,)))
+        return super().token_forward(input_embeddings, infer_state, layer_weight)

lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py

@@ -0,0 +1,24 @@
+from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer
+
+
+class Glm5NextMTPTransformerLayerInfer(Glm5NextTransformerLayerInfer):
+    def context_forward(self, input_embeddings, infer_state, layer_weight):
+        if not getattr(infer_state, "prefill_kv_only", False):
+            return super().context_forward(input_embeddings, infer_state, layer_weight)
+
+        normed = self._att_norm(input_embeddings, infer_state, layer_weight)
+        hidden, _, kv = self._project_sparse_qkv_a(normed, infer_state, layer_weight)
+        infer_state.mtp_draft_input_hiddens = None
+        self._post_cache_kv(self._normalize_sparse_kv(kv, layer_weight), infer_state, layer_weight)
+        self.indexer.write_k_pools(hidden, infer_state, infer_state.prefill_att_state, layer_weight)
+        # BaseModel pads empty ranks to at least one row. An empty result marks
+        # the discarded draft output for our post layer without retaining hiddens.
+        return input_embeddings.new_empty((0, self.embed_dim_))
+
+    def overlap_tpsp_context_forward(
+        self, input_embeddings, input_embeddings1, infer_state, infer_state1, layer_weight
+    ):
+        return (
+            self.context_forward(input_embeddings, infer_state, layer_weight),
+            self.context_forward(input_embeddings1, infer_state1, layer_weight),
+        )

lightllm/models/glm5_next_mtp/model.py

@@ -5,10 +5,10 @@
 from lightllm.common.basemodel import TpPartBaseModel
 from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend
 from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend
-from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer
 from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight
-from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
+from .layer_infer.post_layer_infer import Glm5NextMTPPostLayerInfer
 from .layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer
+from .layer_infer.transformer_layer_infer import Glm5NextMTPTransformerLayerInfer
 from .layer_weights.pre_and_post_layer_weight import Glm5NextMTPPreAndPostLayerWeight
 
 
@@ -24,14 +24,28 @@ class Glm5NextMTPModel(TpPartBaseModel):
     pre_and_post_weight_class = Glm5NextMTPPreAndPostLayerWeight
     transformer_weight_class = Glm5NextTransformerLayerWeight
     pre_layer_infer_class = Glm5NextMTPPreLayerInfer
-    post_layer_infer_class = LlamaPostLayerInfer
-    transformer_layer_infer_class = Glm5NextTransformerLayerInfer
+    post_layer_infer_class = Glm5NextMTPPostLayerInfer
+    transformer_layer_infer_class = Glm5NextMTPTransformerLayerInfer
 
     def __init__(self, kvargs):
         self.main_model = kvargs.pop("main_model")
         self.mtp_previous_draft_models = kvargs.pop("mtp_previous_draft_models")
         super().__init__(kvargs)
 
+    def _supports_prefill_kv_only(self):
+        args = self.args
+        # EAGLE reuses one draft model and discards its prefill output. Vanilla
+        # only does so for a single chained draft; later chains consume the
+        # preceding draft output.
+        return args.mtp_mode == "eagle_with_att" or (args.mtp_mode == "vanilla_with_att" and args.mtp_step == 1)
+
+    def _create_inferstate(self, model_input, microbatch_index=0):
+        infer_state = super()._create_inferstate(model_input, microbatch_index=microbatch_index)
+        # Keep the mode-specific marker on this model's infer state so the
+        # generic BaseModel pipeline remains unchanged.
+        infer_state.prefill_kv_only = infer_state.is_prefill and self._supports_prefill_kv_only()
+        return infer_state
+
     def _init_config(self):
         # Reuse the validated target configuration without changing its mHC layout.
         self.config = self.normalize_model_config(deepcopy(self.main_model.config))

lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/vanilla_with_att.py

@@ -47,14 +47,16 @@ def fill_draft_model_kv_state_overlap(
         ]
         draft_token_ids = list(target_next_token_ids)
 
-        for draft_model in self.backend.draft_models:
+        for draft_index, draft_model in enumerate(self.backend.draft_models):
             for batch_index, model_input in enumerate(model_inputs):
                 model_inputs[batch_index] = self._prepare_mtp_prefill_inputs(
                     model_input=model_input,
                     b_next_token_ids=draft_token_ids[batch_index],
                     mtp_draft_input_hiddens=draft_hiddens[batch_index],
                 )
             draft_outputs = draft_model._microbatch_overlap_prefill_cuda(*model_inputs)
+            if draft_index + 1 == len(self.backend.draft_models):
+                break
             for batch_index, draft_output in enumerate(draft_outputs):
                 draft_hiddens[batch_index] = draft_output.mtp_collector.spec_hidden
                 draft_token_ids[batch_index] = self.backend._gen_argmax_token_ids(draft_output)

lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py

@@ -30,13 +30,17 @@ def fill_draft_model_kv_state(
         draft_input = copy.copy(target_model_input)
         draft_hidden = target_model_output.mtp_collector.spec_hidden
         draft_token_ids = target_next_token_ids
-        for draft_model in self.backend.draft_models:
+        for draft_index, draft_model in enumerate(self.backend.draft_models):
             draft_input = self._prepare_mtp_prefill_inputs(
                 model_input=draft_input,
                 b_next_token_ids=draft_token_ids,
                 mtp_draft_input_hiddens=draft_hidden,
             )
             draft_output = draft_model.forward(draft_input)
+            # The final chained draft only needs its KV cache. Its output is
+            # consumed by no subsequent draft model.
+            if draft_index + 1 == len(self.backend.draft_models):
+                break
             draft_hidden = draft_output.mtp_collector.spec_hidden
             draft_token_ids = self.backend._gen_argmax_token_ids(draft_output)