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..0b9324a9d4 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,14 @@ 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 _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/layer_infer/post_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py new file mode 100644 index 0000000000..4b1196ea8d --- /dev/null +++ b/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) diff --git a/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..29e55ebaed --- /dev/null +++ b/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), + ) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 6c5f80a414..005ef29586 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/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)) diff --git a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/vanilla_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/vanilla_with_att.py index ca02970ba5..966df0b822 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/vanilla_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/vanilla_with_att.py @@ -47,7 +47,7 @@ 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, @@ -55,6 +55,8 @@ def fill_draft_model_kv_state_overlap( 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) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py index 7ad136f9e3..4383c58e2b 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py +++ b/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) 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..5b60fece98 --- /dev/null +++ b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py @@ -0,0 +1,278 @@ +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.attention_state.full_att.packed_mla.packed_mla_state import ( + PackedMLAAttentionConfig, + PackedMLAFullAttentionState, + PackedMLALogicalConfig, +) +from lightllm.common.basemodel.batch_objs import ModelInput, PostLayerOutput +from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer +from lightllm.models.glm5_next_mtp.layer_infer.post_layer_infer import Glm5NextMTPPostLayerInfer +from lightllm.models.glm5_next_mtp.layer_infer.transformer_layer_infer import Glm5NextMTPTransformerLayerInfer +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer +from lightllm.server.core.objs.start_args_type import StartArgs + + +def make_model(): + model = Glm5NextMTPModel.__new__(Glm5NextMTPModel) + model.args = StartArgs() + 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", + [ + ("mtp_mode", "eagle_no_att"), + ], +) +def test_unsupported_execution_uses_full_forward(field, value): + model = make_model() + assert model._supports_prefill_kv_only() + setattr(model.args, field, value) + assert not model._supports_prefill_kv_only() + + +@pytest.mark.parametrize("field,value", [("mtp_mode", "eagle_no_att")]) +def test_unsupported_layout_uses_full_forward(field, value): + model = make_model() + setattr(model.args, field, value) + assert not model._supports_prefill_kv_only() + + +def test_eagle_multistep_and_vanilla_multistep_capabilities(): + model = make_model() + model.args.mtp_step = 3 + assert model._supports_prefill_kv_only() + model.args.mtp_mode = "vanilla_with_att" + assert not model._supports_prefill_kv_only() + + +@pytest.mark.parametrize("kv_type", ["fp8kv_dsa", "bf16"]) +@pytest.mark.parametrize("tp_world_size", [1, 4]) +def test_cache_only_capability_does_not_depend_on_kv_format_or_tp(kv_type, tp_world_size): + model = make_model() + model.args.llm_kv_type = kv_type + model.tp_world_size_ = tp_world_size + assert model._supports_prefill_kv_only() + + +@pytest.mark.parametrize("tokens", [0, 1, 8192]) +@pytest.mark.parametrize("mode", ["eagle_with_att", "vanilla_with_att"]) +def test_forward_uses_common_prefill(monkeypatch, tokens, mode): + from lightllm.common.basemodel import basemodel + + model = make_model() + model.args.mtp_mode = mode + model.mtp_previous_draft_models = [] + model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=3) + model.att_state_manager = object() + model.hidden_collector_prototype = FinalHiddenCollector() + model.return_all_prompt_logics = False + model.prefill_att_backend = model.prefill_att_backend1 = None + model._select_mem_indexes = lambda model_input: torch.zeros_like(model_input.input_ids) + batch = int(tokens > 0) + model_input = ModelInput( + is_prefill=True, + batch_size=batch, + total_token_num=tokens, + max_q_seq_len=tokens, + max_kv_seq_len=tokens, + max_cache_len=0, + input_ids=torch.zeros(tokens, dtype=torch.int64), + b_req_idx=torch.zeros(batch, dtype=torch.int32), + b_mtp_index=torch.zeros(batch, dtype=torch.int32), + b_seq_len=torch.full((batch,), tokens, dtype=torch.int32), + b_ready_cache_len=torch.zeros(batch, dtype=torch.int32), + b_prefill_start_loc=torch.zeros(batch, dtype=torch.int32), + b_is_decode_req=torch.zeros(batch, dtype=torch.bool), + b_prefill_has_output_cpu=[True] * batch, + multimodal_params=[{"images": [], "audios": []} for _ in range(batch)], + mtp_draft_input_hiddens=torch.ones(tokens, 2), + ) + model.pre_infer = SimpleNamespace( + context_forward=Mock(side_effect=lambda ids, *_: torch.zeros(ids.shape[0], 2)), + _tpsp_sp_split=lambda input, **_: input, + ) + model.post_infer = Glm5NextMTPPostLayerInfer.__new__(Glm5NextMTPPostLayerInfer) + model.post_infer.tp_world_size_ = 1 + model.post_infer._tpsp_allgather = Mock(side_effect=lambda input, **_: input) + post_forward = Mock(side_effect=lambda hidden, state, _: PostLayerOutput(torch.ones(state.batch_size, 3))) + monkeypatch.setattr(LlamaPostLayerInfer, "token_forward", post_forward) + model.pre_post_weight = object() + layer = Glm5NextMTPTransformerLayerInfer.__new__(Glm5NextMTPTransformerLayerInfer) + layer.is_linear_attention_layer = layer.use_mhc = False + layer.layer_num_, layer.num_hidden_layers = 45, 1 + layer.embed_dim_ = 2 + layer._att_norm = layer._ffn_norm = lambda hidden, *_: hidden + layer.context_attention_forward = Mock(side_effect=lambda hidden, *_: torch.zeros_like(hidden)) + layer._ffn = Mock(side_effect=lambda hidden, *_: torch.zeros_like(hidden)) + layer._project_sparse_qkv_a = lambda hidden, *_: (hidden, None, hidden) + layer._normalize_sparse_kv = lambda kv, _: kv + layer.indexer = SimpleNamespace(write_k_pools=Mock()) + writer = Mock() + layer._post_cache_kv = writer + model.layers_infer = [layer] + allocator = SimpleNamespace(cache_env_in=Mock(), cache_env_out=Mock()) + monkeypatch.setattr(basemodel, "g_cache_manager", allocator) + monkeypatch.setattr(basemodel.dist_group_manager, "get_group", Mock()) + monkeypatch.setattr(basemodel.dist_group_manager, "clear_deepep_buffer", Mock()) + monkeypatch.setattr(ModelInput, "to_cuda", lambda self: self.check_input()) + monkeypatch.setattr(InferStateInfo, "init_some_extra_state", Mock()) + event = Mock() + monkeypatch.setattr(torch.cuda, "Event", Mock(return_value=event)) + output = model.forward(model_input) + assert output.prefill_mem_indexes_ready_event is event + assert output.logits.numel() == 0 + assert output.mtp_collector.spec_hidden is None + ids, infer, _ = model.pre_infer.context_forward.call_args.args + assert ids.shape == (max(1, tokens),) + assert infer.b_req_idx.tolist() == ([0] if tokens else [model.req_manager.HOLD_REQUEST_ID]) + if mode == "eagle_with_att": + assert infer.mtp_draft_input_hiddens is None + writer.assert_called_once() + layer.context_attention_forward.assert_not_called() + layer._ffn.assert_not_called() + post_forward.assert_not_called() + placeholder = model.post_infer._tpsp_allgather.call_args.kwargs["input"] + assert placeholder.shape == (0, layer.embed_dim_) + assert placeholder.untyped_storage().nbytes() == 0 + layer.indexer.write_k_pools.assert_called_once() + else: + writer.assert_called_once() + layer.context_attention_forward.assert_not_called() + layer._ffn.assert_not_called() + post_forward.assert_not_called() + 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("kv_type", ["fp8kv_dsa", "bf16"]) +@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, kv_type, 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 = Glm5NextMTPTransformerLayerInfer.__new__(Glm5NextMTPTransformerLayerInfer) + 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"): + if kv_type == "fp8kv_dsa": + full = FP8E8M0PackedMLAState(FP8E8M0PackedMLAConfig(logical_layer_groups=((layer_id,),))) + else: + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + PackedMLALogicalConfig( + logical_layer_groups=((layer_id,),), + latent_dim=512, + format="packed_bf16", + indexer_dim=128, + ) + ), + device="cuda", + ) + 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, + prefill_kv_only=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() + placeholder = layer.context_forward(hidden, infer, weights) + assert placeholder.shape == (0, layer.embed_dim_) + assert placeholder.untyped_storage().nbytes() == 0 + 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])