diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index de2cacd5a7..b7e6b9b7f3 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/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( diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index be7a2fa781..50c9d46c0d 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/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,9 +139,9 @@ 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], @@ -149,7 +149,23 @@ def _get_qkv(self, input, infer_state, layer_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 write_draft_prefill_cache(self, input_embeddings, infer_state, layer_weight): + """Write the single NextN block's MLA KV, pooled keys and raw tail.""" + assert infer_state.is_prefill and infer_state.is_draft_model + assert not self.use_mhc and self.layer_num_ >= self.num_hidden_layers + normed = self._att_norm(input_embeddings, infer_state, layer_weight) + hidden, _, kv = self._project_sparse_qkv_a(normed, infer_state, layer_weight) + 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) 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() diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 6c5f80a414..caa3c16100 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -5,6 +5,7 @@ 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.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager 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 @@ -32,6 +33,48 @@ def __init__(self, kvargs): self.mtp_previous_draft_models = kvargs.pop("mtp_previous_draft_models") super().__init__(kvargs) + def _can_fill_prefill_kv_only(self): + args = self.args + # This decision must be uniform across EP ranks, including empty ranks. + # Keep graph, TP and mixed-mode execution on the normal forward path. + return ( + args.run_mode == "prefill" + and args.mtp_mode == "eagle_with_att" + and args.mtp_step == 1 + and args.llm_kv_type == "fp8kv_dsa" + and self.tp_world_size_ == 1 + and not args.enable_tpsp_mix_mode + and not args.enable_dp_prefill_balance + and not args.enable_prefill_decode_mixed + and not args.enable_prefill_microbatch_overlap + and self.prefill_graph is None + and self.layers_num == len(self.layers_infer) == len(self.trans_layers_weight) == 1 + and not self.layers_infer[0].is_linear_attention_layer + and not self.layers_infer[0].use_mhc + ) + + @torch.no_grad() + def fill_prefill_kv(self, model_input): + """Skip unused draft outputs when pure Prefill only needs the cache.""" + assert model_input.is_prefill + if not self._can_fill_prefill_kv_only(): + return self.forward(model_input) + + model_input.to_cuda() + model_input = self._create_padded_prefill_model_input( + model_input=model_input, new_handle_token_num=max(1, model_input.input_ids.shape[0]) + ) + infer_state = self._create_inferstate(model_input) + infer_state.init_some_extra_state(self) + infer_state.init_att_state() + input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) + infer_state.mtp_draft_input_hiddens = None + g_cache_manager.cache_env_in() + try: + self.layers_infer[0].write_draft_prefill_cache(input_embs, infer_state, self.trans_layers_weight[0]) + finally: + g_cache_manager.cache_env_out() + 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)) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 721af1c27d..b06a596f45 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -33,7 +33,8 @@ def fill_draft_model_kv_state( b_next_token_ids=target_next_token_ids, mtp_draft_input_hiddens=target_model_output.mtp_collector.spec_hidden, ) - self.backend.draft_models[0].forward(draft_input) + draft = self.backend.draft_models[0] + getattr(draft, "fill_prefill_kv", draft.forward)(draft_input) def propose_next( self, diff --git a/unit_tests/models/glm5_next/test_mtp_prefill_cache.py b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py new file mode 100644 index 0000000000..bbb87e7b7b --- /dev/null +++ b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py @@ -0,0 +1,189 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig +from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_state import ( + FP8E8M0PackedMLAConfig, + FP8E8M0PackedMLAState, +) +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.server.core.objs.start_args_type import StartArgs + + +def make_model(): + model = Glm5NextMTPModel.__new__(Glm5NextMTPModel) + model.args = StartArgs() + model.args.run_mode = "prefill" + model.args.mtp_mode = "eagle_with_att" + model.args.mtp_step = 1 + model.args.llm_kv_type = "fp8kv_dsa" + model.tp_world_size_ = model.layers_num = 1 + model.prefill_graph = None + model.layers_infer = [SimpleNamespace(is_linear_attention_layer=False, use_mhc=False)] + model.trans_layers_weight = [object()] + return model + + +@pytest.mark.parametrize( + "field,value", + [ + ("run_mode", "decode"), + ("run_mode", "normal"), + ("mtp_mode", "eagle_no_att"), + ("mtp_step", 3), + ("llm_kv_type", "none"), + ("enable_tpsp_mix_mode", True), + ("enable_dp_prefill_balance", True), + ("enable_prefill_decode_mixed", True), + ("enable_prefill_microbatch_overlap", True), + ], +) +def test_unsupported_execution_uses_full_forward(field, value): + model = make_model() + assert model._can_fill_prefill_kv_only() + setattr(model.args, field, value) + model.forward = Mock(return_value=object()) + model_input = SimpleNamespace(is_prefill=True) + assert model.fill_prefill_kv(model_input) is model.forward.return_value + model.forward.assert_called_once_with(model_input) + + +@pytest.mark.parametrize("field,value", [("tp_world_size_", 2), ("prefill_graph", object()), ("layers_num", 2)]) +def test_unsupported_layout_uses_full_forward(field, value): + model = make_model() + setattr(model, field, value) + assert not model._can_fill_prefill_kv_only() + + +@pytest.mark.parametrize("tokens", [0, 1, 8192]) +@pytest.mark.parametrize("fail_write", [False, True]) +def test_fill_pads_empty_ranks_and_releases_allocator(monkeypatch, tokens, fail_write): + from lightllm.models.glm5_next_mtp import model as model_module + + model = make_model() + model_input = SimpleNamespace(is_prefill=True, input_ids=torch.empty(tokens), to_cuda=Mock()) + infer = SimpleNamespace( + input_ids=object(), init_some_extra_state=Mock(), init_att_state=Mock(), mtp_draft_input_hiddens=object() + ) + model._create_padded_prefill_model_input = Mock(return_value=model_input) + model._create_inferstate = Mock(return_value=infer) + model.pre_infer = SimpleNamespace(context_forward=Mock(return_value=object())) + model.pre_post_weight = object() + writer = Mock(side_effect=RuntimeError("cache write failed") if fail_write else None) + model.layers_infer[0].write_draft_prefill_cache = writer + allocator = SimpleNamespace(cache_env_in=Mock(), cache_env_out=Mock()) + monkeypatch.setattr(model_module, "g_cache_manager", allocator) + if fail_write: + with pytest.raises(RuntimeError, match="cache write failed"): + model.fill_prefill_kv(model_input) + else: + assert model.fill_prefill_kv(model_input) is None + model._create_padded_prefill_model_input.assert_called_once_with( + model_input=model_input, new_handle_token_num=max(1, tokens) + ) + model.pre_infer.context_forward.assert_called_once_with(infer.input_ids, infer, model.pre_post_weight) + writer.assert_called_once_with(model.pre_infer.context_forward.return_value, infer, model.trans_layers_weight[0]) + assert infer.mtp_draft_input_hiddens is None + allocator.cache_env_in.assert_called_once() + allocator.cache_env_out.assert_called_once() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "requests", + [((3, 0, 1),), ((0, 0, 8192),), ((0, 131071, 129), (2, 1048573, 3)), ((1, 3, 1), (0, 5, 7))], +) +def test_cache_only_matches_full_prefill_cache_writes(monkeypatch, requests): + from lightllm.models.glm5_next import indexer as indexer_module + + torch.manual_seed(53) + layer_id, tokens = 45, sum(q for _, _, q in requests) + layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer) + layer.layer_num_, layer.num_hidden_layers = layer_id, 1 + layer.use_mhc = layer.is_linear_attention_layer = False + layer.tp_world_size_ = layer.tp_q_head_num_ = 1 + layer.embed_dim_ = layer.q_lora_rank = 64 + layer.kv_lora_rank = layer.qk_nope_head_dim = 512 + layer.eps_ = 1e-5 + layer.indexer = Glm5NextNsaInfer( + layer_id, dict(index_topk=2048, index_n_heads=1, index_head_dim=128, rms_norm_eps=1e-5, index_kpool=4), 1 + ) + + def projection(rows, cols): + weight = torch.randn(rows, cols, device="cuda", dtype=torch.bfloat16) * 0.1 + return SimpleNamespace(mm=Mock(side_effect=lambda x: x @ weight)) + + def norm(input, eps, **_): + return rmsnorm_forward(input, None, eps) + + weights = SimpleNamespace( + att_norm_weight_=norm, + qkv_a_proj_with_mqa_=projection(64, 576), + q_a_layernorm_=SimpleNamespace(weight=torch.ones(64, device="cuda", dtype=torch.bfloat16)), + kv_a_layernorm_=SimpleNamespace(weight=torch.ones(512, device="cuda", dtype=torch.bfloat16)), + q_b_proj_=projection(64, 512), + wk_proj_=projection(64, 128), + k_norm_=norm, + index_kpool_compress_gate=projection(64, 128), + index_kpool_compress_ape=SimpleNamespace(weight=torch.randn(4, 128, device="cuda", dtype=torch.bfloat16)), + ) + resources = AttentionStateManager(request_capacity=4) + stores = [] + for role in ("main", "draft0"): + full = FP8E8M0PackedMLAState(FP8E8M0PackedMLAConfig(logical_layer_groups=((layer_id,),))) + full.allocate_storage(2 * tokens + 2) + full.storage.fill_(0x5A) + tail = WindowCompressorState( + WindowCompressorStateConfig((layer_id,), 4, 4, 128, "indexer", torch.bfloat16), 4, mtp_step=1 + ) + tail.buffer.normal_() + resources.register_state(role, full) + resources.register_state(role, tail) + stores.extend((full.storage, tail.buffer)) + resources.finish_state_registration() + initial = [s.clone() for s in stores] + q_lengths = [q for _, _, q in requests] + lengths = torch.cat([torch.arange(past + 1, past + q + 1, device="cuda") for _, past, q in requests]).int() + infer = SimpleNamespace( + is_prefill=True, + is_draft_model=True, + need_dp_prefill_balance=False, + att_state_role="draft0", + att_state_manager=resources, + mem_index=2 * torch.arange(tokens, device="cuda", dtype=torch.int32) + 1, + b_req_idx=torch.tensor([r for r, _, _ in requests], device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([p + q for _, p, q in requests], device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=torch.tensor([0, *q_lengths], device="cuda", dtype=torch.int32).cumsum(0).int(), + max_q_seq_len=max(q_lengths), + max_kv_seq_len=2048, + prefill_att_state=SimpleNamespace(q_kv_seq_len=lengths, q_req_idx=None), + req_manager=SimpleNamespace(req_to_token_indexs=None), + ) + # Top-k outputs are discarded here; run the normal projections and cache writers. + monkeypatch.setattr(indexer_module, "expand_topk", Mock(return_value=(None, None))) + hidden = torch.randn(tokens, 64, device="cuda", dtype=torch.bfloat16) + _, kv = layer._get_qkv(layer._att_norm(hidden, infer, weights), infer, weights) + layer._post_cache_kv(kv, infer, weights) + layer.indexer._get_indices( + **infer.get_topk_indices_params, infer_state=infer, att_state=infer.prefill_att_state, layer_weight=weights + ) + expected = [s.clone() for s in stores] + assert not torch.equal(initial[2], expected[2]) + for store, before in zip(stores, initial): + store.copy_(before) + weights.q_b_proj_.mm.reset_mock() + layer.write_draft_prefill_cache(hidden, infer, weights) + weights.q_b_proj_.mm.assert_not_called() + for actual, reference in zip(stores, expected): + assert torch.equal(actual, reference) + for actual, before in zip(stores[:2], initial[:2]): + assert torch.equal(actual, before) + assert torch.equal(stores[2][:, ::2], initial[2][:, ::2]) + assert torch.equal(stores[3][:, 3], initial[3][:, 3]) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py index 365acb30b2..38e97defa1 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py @@ -73,7 +73,8 @@ def test_eagle_with_att_rejects_zero_draft_steps(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_eagle_with_att_prefill_builds_draft_kv_without_mutating_target(): +@pytest.mark.parametrize("cache_only", [False, True]) +def test_eagle_with_att_prefill_builds_draft_kv_without_mutating_target(cache_only): device = "cuda" original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) @@ -101,8 +102,12 @@ def forward(model_input): ) return ModelOutput(logits=torch.empty((0,), device=device)) + draft = SimpleNamespace(forward=forward) + if cache_only: + draft.fill_prefill_kv = forward + draft.forward = lambda _: pytest.fail("cache-only draft must not run full forward") proposer = EagleWithAttProposer( - backend=SimpleNamespace(draft_models=[SimpleNamespace(forward=forward)]), + backend=SimpleNamespace(draft_models=[draft]), enable_dynmaic_mtp=False, ) proposer.fill_draft_model_kv_state(