From 3e80309e60e245ce16eac38b0721738c7c45bcc0 Mon Sep 17 00:00:00 2001 From: XUE TONGYAO <268319265+Alex-stack-hub@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:28:28 +0800 Subject: [PATCH] feat(spec_decode): add minimal Gemma4 MTP support Signed-off-by: XUE TONGYAO <268319265+Alex-stack-hub@users.noreply.github.com> --- .github/workflows/scripts/test_config.yaml | 3 + tests/ut/ops/test_rotary_embedding.py | 18 ++ tests/ut/spec_decode/test_gemma4_proposer.py | 196 ++++++++++++++++++ vllm_ascend/ops/rotary_embedding.py | 24 ++- .../platform/patch_speculative_config.py | 3 + vllm_ascend/spec_decode/__init__.py | 3 + vllm_ascend/spec_decode/gemma4_proposer.py | 116 +++++++++++ vllm_ascend/spec_decode/llm_base_proposer.py | 37 +++- vllm_ascend/worker/model_runner_v1.py | 34 ++- 9 files changed, 418 insertions(+), 16 deletions(-) create mode 100644 tests/ut/spec_decode/test_gemma4_proposer.py create mode 100644 vllm_ascend/spec_decode/gemma4_proposer.py diff --git a/.github/workflows/scripts/test_config.yaml b/.github/workflows/scripts/test_config.yaml index b02cfcf79a5e..a4b133a4c81e 100644 --- a/.github/workflows/scripts/test_config.yaml +++ b/.github/workflows/scripts/test_config.yaml @@ -467,6 +467,8 @@ source_file_dependencies: - vllm_ascend/spec_decode/llm_base_proposer.py - vllm_ascend/spec_decode/eagle_proposer.py + - vllm_ascend/spec_decode/gemma4_proposer.py + - vllm_ascend/patch/platform/patch_speculative_config.py - vllm_ascend/models/llama_eagle3_vwn.py tests: - tests/e2e/pull_request/one_card/spec_decode/test_eagle.py @@ -474,6 +476,7 @@ - tests/e2e/pull_request/two_card/spec_decode/test_spec_decode.py - tests/e2e/pull_request/four_card/spec_decode/test_mtp_qwen3_next.py - tests/ut/spec_decode/test_speculators_vwn_eagle3.py + - tests/ut/spec_decode/test_gemma4_proposer.py - name: spec_decode_dspark optional: false diff --git a/tests/ut/ops/test_rotary_embedding.py b/tests/ut/ops/test_rotary_embedding.py index 2a6ead2483c1..4f599f3ca6e8 100644 --- a/tests/ut/ops/test_rotary_embedding.py +++ b/tests/ut/ops/test_rotary_embedding.py @@ -160,6 +160,24 @@ def test_basic_call_delegates_to_npu_op(self, mock_get_forward_context, mock_npu ) assert result is expected_output + @patch("vllm_ascend.ops.rotary_embedding.is_forward_context_available", return_value=False) + @patch("vllm_ascend.ops.rotary_embedding.rope_forward_oot") + def test_q_only_uses_throwaway_key(self, mock_rope, _mock_is_ctx, make_embedding): + emb = make_embedding() + positions, query, _ = _make_tensors() + expected_query = torch.randn_like(query) + mock_rope.return_value = expected_query, torch.empty_like(query) + + with patch("vllm_ascend.ops.rotary_embedding.HAS_TRITON", False): + result = emb.forward_oot(positions, query, None) + + assert result[0] is expected_query + assert result[1] is None + dummy_key = mock_rope.call_args.args[2] + assert dummy_key.shape == (query.shape[0], HEAD_SIZE) + assert dummy_key.dtype == query.dtype + assert dummy_key.device == query.device + @patch("torch.ops.vllm.npu_rotary_embedding") @patch("vllm_ascend.ascend_forward_context.get_forward_context") def test_neox_style_override_true(self, mock_get_forward_context, mock_npu_op, make_embedding): diff --git a/tests/ut/spec_decode/test_gemma4_proposer.py b/tests/ut/spec_decode/test_gemma4_proposer.py new file mode 100644 index 000000000000..42a57b9921c7 --- /dev/null +++ b/tests/ut/spec_decode/test_gemma4_proposer.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import torch + +import vllm_ascend.spec_decode as spec_decode +from vllm_ascend.attention.attention_v1 import AscendAttentionState +from vllm_ascend.patch.platform import patch_speculative_config +from vllm_ascend.spec_decode.gemma4_proposer import AscendGemma4Proposer +from vllm_ascend.spec_decode.llm_base_proposer import AscendSpecDecodeBaseProposer + + +def test_routes_gemma4_mtp_to_ascend_proposer(): + speculative_config = MagicMock() + speculative_config.use_gemma4_mtp.return_value = True + vllm_config = SimpleNamespace(speculative_config=speculative_config) + expected = object() + + with patch.object( + spec_decode, + "AscendGemma4Proposer", + return_value=expected, + ) as proposer_cls: + result = spec_decode.get_spec_decode_method( + "mtp", + vllm_config, + device="npu", + runner=object(), + ) + + assert result is expected + proposer_cls.assert_called_once() + + +def test_gemma_config_override_delegates_to_vllm(monkeypatch): + hf_config = SimpleNamespace( + architectures=["Gemma4ForConditionalGeneration"], + model_type="gemma4_assistant", + ) + expected = object() + original_override = MagicMock(return_value=expected) + monkeypatch.setattr( + patch_speculative_config, + "_orig_hf_config_override", + original_override, + ) + + result = patch_speculative_config.hf_config_override(hf_config) + + assert result is expected + original_override.assert_called_once_with(hf_config) + + +def test_sync_kv_sharing_target_to_impl(): + proposer = AscendGemma4Proposer.__new__(AscendGemma4Proposer) + proposer.vllm_config = MagicMock() + proposer._draft_attn_layer_names = {"draft.attn"} + + impl = SimpleNamespace(kv_sharing_target_layer_name=None) + attn = SimpleNamespace( + impl=impl, + kv_sharing_target_layer_name="target.attn", + ) + with patch( + "vllm_ascend.spec_decode.gemma4_proposer.get_layers_from_vllm_config", + return_value={"draft.attn": attn}, + ): + proposer._sync_kv_sharing_target_to_impl() + + assert impl.kv_sharing_target_layer_name == "target.attn" + + +def test_keeps_draft_lm_head(): + proposer = AscendGemma4Proposer.__new__(AscendGemma4Proposer) + draft_lm_head = object() + proposer.model = SimpleNamespace(lm_head=draft_lm_head) + proposer.method = "mtp" + proposer.use_cuda_graph = False + proposer.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(is_deepseek_mla=False), + compilation_config=SimpleNamespace( + cudagraph_mode=SimpleNamespace( + has_full_cudagraphs=lambda: False, + ) + ), + ) + + proposer._maybe_share_lm_head(SimpleNamespace(lm_head=object())) + + assert proposer.model.lm_head is draft_lm_head + + +def test_build_draft_attn_metadata_uses_per_group_block_tables(): + proposer = AscendGemma4Proposer.__new__(AscendGemma4Proposer) + block_tables = { + 0: torch.arange(12).view(3, 4), + 1: torch.arange(12, 24).view(3, 4), + } + proposer._per_group_block_tables = block_tables + proposer.runner = SimpleNamespace(get_model=MagicMock(return_value=object())) + metadata = [ + SimpleNamespace(attn_state=None, causal=True), + SimpleNamespace(attn_state=None, causal=False, attn_mask=object()), + ] + builders = [MagicMock(), MagicMock()] + for builder, group_metadata in zip(builders, metadata): + builder.build.return_value = group_metadata + proposer.draft_attn_groups = [ + SimpleNamespace( + kv_cache_group_id=gid, + layer_names=[f"draft.attn.{gid}"], + get_metadata_builder=MagicMock(return_value=builders[gid]), + ) + for gid in range(2) + ] + common_metadata = SimpleNamespace( + num_reqs=2, + block_table_tensor=torch.zeros(2, 4), + ) + + multi_steps, first_metadata = proposer.build_draft_attn_metadata( + common_metadata, + num_input_tokens=2, + num_actual_tokens=2, + ) + + assert first_metadata is metadata[0] + assert multi_steps == [ + { + "draft.attn.0": metadata[0], + "draft.attn.1": metadata[1], + } + ] + for gid, builder in enumerate(builders): + group_common_metadata = builder.build.call_args.args[1] + assert group_common_metadata is not common_metadata + assert torch.equal( + group_common_metadata.block_table_tensor, + block_tables[gid][:2], + ) + assert metadata[gid].attn_state == AscendAttentionState.SpecDecoding + assert metadata[1].attn_mask is None + + graph_metadata = [object(), object()] + for builder, graph_item in zip(builders, graph_metadata): + builder.build_for_graph_capture.return_value = graph_item + graph_result = proposer._build_multi_group_graph_capture_metadata( + common_metadata, + draft_index=0, + ) + + assert graph_result == { + "draft.attn.0": graph_metadata[0], + "draft.attn.1": graph_metadata[1], + } + for gid, builder in enumerate(builders): + call_args = builder.build_for_graph_capture.call_args.args + assert torch.equal(call_args[0].block_table_tensor, block_tables[gid][:2]) + assert call_args[1] == AscendAttentionState.SpecDecoding + + +def test_attn_update_uses_only_active_group_block_table(): + proposer = AscendGemma4Proposer.__new__(AscendGemma4Proposer) + proposer._per_group_block_tables = {1: torch.arange(12).view(3, 4)} + common_metadata = SimpleNamespace( + num_reqs=2, + block_table_tensor=torch.zeros(2, 4), + ) + attn_group = SimpleNamespace(kv_cache_group_id=1) + expected = object() + + with patch.object( + AscendSpecDecodeBaseProposer, + "attn_update_stack_num_spec_norm", + return_value=expected, + ) as base_update: + result = proposer.attn_update_stack_num_spec_norm( + 1, + common_metadata, + 2, + 2, + torch.tensor([3, 4]), + "none", + attn_group=attn_group, + ) + + assert result is expected + group_common_metadata = base_update.call_args.args[1] + assert group_common_metadata is not common_metadata + assert torch.equal( + group_common_metadata.block_table_tensor, + proposer._per_group_block_tables[1][:2], + ) + assert base_update.call_args.kwargs["attn_group"] is attn_group diff --git a/vllm_ascend/ops/rotary_embedding.py b/vllm_ascend/ops/rotary_embedding.py index 646e449cef69..8a7948532d30 100644 --- a/vllm_ascend/ops/rotary_embedding.py +++ b/vllm_ascend/ops/rotary_embedding.py @@ -235,7 +235,7 @@ def forward_oot( self, positions: torch.Tensor, query: torch.Tensor, - key: torch.Tensor, + key: torch.Tensor | None, offsets: torch.Tensor | None = None, is_neox_style_override: bool | None = None, ): @@ -246,6 +246,28 @@ def forward_oot( flash_comm_v1_enabled = _EXTRA_CTX.flash_comm_v1_enabled if is_forward_context_available() else False if is_draft_model and self.use_mtp and flash_comm_v1_enabled: positions = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(positions.contiguous(), True) + if key is None: + # Gemma4 MTP reads K/V from the target cache. Reuse the regular + # rotary implementation with a throwaway key buffer. + dummy_key = ( + torch.empty(query.shape[0], 0, self.head_size, dtype=query.dtype, device=query.device) + if HAS_TRITON + else torch.empty( + (query.shape[0], 1, self.head_size) if query.ndim == 3 else (query.shape[0], self.head_size), + dtype=query.dtype, + device=query.device, + ) + ) + query, _ = rope_forward_oot( + positions, + query, + dummy_key, + self.cos_sin_cache, + self.head_size, + self.rotary_dim, + is_neox_style, + ) + return query, None return torch.ops.vllm.npu_rotary_embedding( positions, query, key, self.cos_sin_cache, self.head_size, self.rotary_dim, is_neox_style ) diff --git a/vllm_ascend/patch/platform/patch_speculative_config.py b/vllm_ascend/patch/platform/patch_speculative_config.py index 97d772d9d47b..14ecbc86a6d0 100644 --- a/vllm_ascend/patch/platform/patch_speculative_config.py +++ b/vllm_ascend/patch/platform/patch_speculative_config.py @@ -4,6 +4,7 @@ from vllm.utils.import_utils import LazyLoader _orig_post_init = SpeculativeConfig.__post_init__ +_orig_hf_config_override = SpeculativeConfig.hf_config_override if TYPE_CHECKING: import vllm.model_executor.layers.quantization as me_quant @@ -16,6 +17,8 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: initial_architecture = hf_config.architectures[0] + if hf_config.model_type in ("gemma4_assistant", "gemma4_unified_assistant"): + return _orig_hf_config_override(hf_config) if hf_config.model_type in ("deepseek_v3", "deepseek_v32", "deepseek_v4", "glm_moe_dsa"): target_model_type = hf_config.model_type hf_config.model_type = "deepseek_mtp" diff --git a/vllm_ascend/spec_decode/__init__.py b/vllm_ascend/spec_decode/__init__.py index 0863adfd71fa..da4c24107cfb 100644 --- a/vllm_ascend/spec_decode/__init__.py +++ b/vllm_ascend/spec_decode/__init__.py @@ -25,6 +25,7 @@ from vllm_ascend.spec_decode.extract_hidden_states_proposer import ( AscendExtractHiddenStatesProposer, ) +from vllm_ascend.spec_decode.gemma4_proposer import AscendGemma4Proposer from vllm_ascend.spec_decode.medusa_proposer import AscendMedusaProposer from vllm_ascend.spec_decode.ngram_proposer import AscendNgramProposer from vllm_ascend.spec_decode.ngram_proposer_npu import AscendNgramProposerNPU @@ -43,6 +44,8 @@ def get_spec_decode_method(method, vllm_config, device, runner): return AscendMedusaProposer(vllm_config, device) elif method == "dspark": return AscendDSparkProposer(vllm_config, device, runner) + elif method == "mtp" and vllm_config.speculative_config.use_gemma4_mtp(): + return AscendGemma4Proposer(vllm_config, device, runner) elif method in ("eagle", "eagle3", "mtp"): speculative_config = vllm_config.speculative_config if speculative_config is not None and speculative_config.use_step3p5_mtp(): diff --git a/vllm_ascend/spec_decode/gemma4_proposer.py b/vllm_ascend/spec_decode/gemma4_proposer.py new file mode 100644 index 000000000000..508f60455ff6 --- /dev/null +++ b/vllm_ascend/spec_decode/gemma4_proposer.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Gemma4 MTP proposer for Ascend NPUs.""" + +import copy + +from vllm.config import get_layers_from_vllm_config +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.v1.spec_decode.gemma4 import Gemma4Proposer + +from vllm_ascend.attention.attention_v1 import AscendAttentionState +from vllm_ascend.spec_decode.llm_base_proposer import AscendSpecDecodeBaseProposer + + +class AscendGemma4Proposer(Gemma4Proposer, AscendSpecDecodeBaseProposer): + """Reuse vLLM's Gemma4 proposer with Ascend execution support.""" + + def _setup_centroids_cuda_graphs(self) -> None: + # Centroid sampling runs eagerly on NPU; ACL graph capture is handled by + # AscendSpecDecodeBaseProposer. + self._centroids_sizes: list[int] = [] + + def _maybe_share_lm_head(self, target_language_model) -> None: + # Gemma4 keeps its draft-dimension lm_head. The Ascend implementation + # still needs to install the ACL graph wrapper. + AscendSpecDecodeBaseProposer._maybe_share_lm_head(self, target_language_model) + + def load_model(self, target_model) -> None: + super().load_model(target_model) + self.supports_mm_inputs = False + self._sync_kv_sharing_target_to_impl() + + def _sync_kv_sharing_target_to_impl(self) -> None: + """Propagate late-bound KV-sharing targets to Ascend backends.""" + attn_layers = get_layers_from_vllm_config( + self.vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + ) + for layer_name in self._draft_attn_layer_names: + attn = attn_layers[layer_name] + target = getattr(attn, "kv_sharing_target_layer_name", None) + impl = getattr(attn, "impl", None) + if target is not None and impl is not None: + impl.kv_sharing_target_layer_name = target + + def _get_group_common_attn_metadata(self, attn_group, common_attn_metadata): + block_table = self._per_group_block_tables.get(attn_group.kv_cache_group_id) + if block_table is None: + return common_attn_metadata + group_metadata = copy.copy(common_attn_metadata) + group_metadata.block_table_tensor = block_table[: common_attn_metadata.num_reqs] + return group_metadata + + def _build_multi_group_graph_capture_metadata(self, common_attn_metadata, draft_index): + per_layer_attn_metadata = {} + for attn_group in self.draft_attn_groups: + group_metadata = self._get_group_common_attn_metadata( + attn_group, + common_attn_metadata, + ) + attn_metadata = attn_group.get_metadata_builder().build_for_graph_capture( + group_metadata, + AscendAttentionState.SpecDecoding, + ) + for layer_name in attn_group.layer_names: + per_layer_attn_metadata[layer_name] = attn_metadata + return per_layer_attn_metadata + + def _get_attn_metadata_layer_names(self, attn_group): + return attn_group.layer_names + + def attn_update_stack_num_spec_norm( + self, + draft_index, + old_common_metadata, + *args, + attn_group=None, + **kwargs, + ): + assert attn_group is not None + group_metadata = self._get_group_common_attn_metadata( + attn_group, + old_common_metadata, + ) + kwargs["attn_group"] = attn_group + return super().attn_update_stack_num_spec_norm( + draft_index, + group_metadata, + *args, + **kwargs, + ) + + def build_draft_attn_metadata( + self, + common_attn_metadata, + num_input_tokens, + num_actual_tokens, + ): + per_layer_attn_metadata = {} + for attn_group in self.draft_attn_groups: + group_metadata = self._get_group_common_attn_metadata( + attn_group, + common_attn_metadata, + ) + attn_metadata = attn_group.get_metadata_builder().build( + 0, + group_metadata, + self.runner.get_model(), + ) + attn_metadata.attn_state = AscendAttentionState.SpecDecoding + if hasattr(attn_metadata, "causal") and not attn_metadata.causal: + attn_metadata.attn_mask = None + for layer_name in attn_group.layer_names: + per_layer_attn_metadata[layer_name] = attn_metadata + + attn_metadata = per_layer_attn_metadata[self.draft_attn_groups[0].layer_names[0]] + return [per_layer_attn_metadata], attn_metadata diff --git a/vllm_ascend/spec_decode/llm_base_proposer.py b/vllm_ascend/spec_decode/llm_base_proposer.py index 39c5f65a1c54..d78a6796afdd 100644 --- a/vllm_ascend/spec_decode/llm_base_proposer.py +++ b/vllm_ascend/spec_decode/llm_base_proposer.py @@ -337,6 +337,8 @@ def load_model(self, model: nn.Module) -> None: "Qwen3VLMoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", + "Gemma4ForConditionalGeneration", + "Gemma4UnifiedForConditionalGeneration", "Step3p7ForConditionalGeneration", ]: self.model.config.image_token_index = model.config.image_token_id @@ -535,6 +537,12 @@ def shallow_copy_metadata(self, attn_metadata): # when update. So we can use the shallow copy. return copy.copy(attn_metadata) + def _build_multi_group_graph_capture_metadata(self, common_attn_metadata, draft_index): + return None + + def _get_attn_metadata_layer_names(self, attn_group): + return self.attn_layer_names + @torch.inference_mode() def dummy_run( self, @@ -644,6 +652,13 @@ def dummy_run( if self.dcp_size > 1 and draft_index > 0: assert self.block_table_tensor_clone is not None, "block_table_tensor_clone is not init" common_attn_metadata.block_table_tensor = self.block_table_tensor_clone[:num_reqs] + per_layer_attn_metadata = self._build_multi_group_graph_capture_metadata( + common_attn_metadata, + draft_index, + ) + if per_layer_attn_metadata is not None: + multi_steps_attn_metadata.append(per_layer_attn_metadata) + continue if not self.use_compress or draft_index == 0: attn_metadata_eagle = builder.build_for_graph_capture( common_attn_metadata, @@ -966,7 +981,7 @@ def _propose( **draft_cp_kwargs, attn_group=attn_group, ) - for layer_name in self.attn_layer_names: + for layer_name in self._get_attn_metadata_layer_names(attn_group): per_layer_attn_metadata[layer_name] = attn_metadata multi_steps_attn_metadata.append(per_layer_attn_metadata) @@ -1200,7 +1215,8 @@ def _run_merged_draft( # cast to int32 is crucial when eagle model is compiled. # tensor.argmax() returns int64 by default. input_ids = draft_token_ids_tensor[draft_index] - positions += 1 + if not getattr(self, "constant_draft_positions", False): + positions += 1 # NOTE(woosuk): We should handle the case where the draft model # generates tokens beyond the max model length. Since it is complex @@ -1542,8 +1558,9 @@ def attn_update_stack_num_spec_norm( common_attn_metadata.graph_pad_size = -1 common_attn_metadata.num_input_tokens = input_batch_size - # The loop part - used_update_positions += 1 + advance_draft_positions = not getattr(self, "constant_draft_positions", False) + if advance_draft_positions: + used_update_positions += 1 # Clone the data so that when calculating the data at position 2 and position 3 # in the merged graph, it does not affect position 1 @@ -1578,21 +1595,25 @@ def attn_update_stack_num_spec_norm( # operations in case they are modified in next step's `prepare_input` # of main model. # Increment the sequence lengths. - common_attn_metadata.seq_lens[:batch_size] += 1 + if advance_draft_positions: + common_attn_metadata.seq_lens[:batch_size] += 1 # For the requests that exceed the max model length, we set the # sequence length to 1 to minimize their overheads in attention. exceeds_mask = common_attn_metadata.seq_lens[:batch_size] > self.max_model_len common_attn_metadata.seq_lens[:batch_size].masked_fill_(exceeds_mask, 1) if common_attn_metadata.seq_lens_cpu is not None: - common_attn_metadata.seq_lens_cpu[:batch_size] = common_attn_metadata.seq_lens_cpu[:batch_size] + 1 + if advance_draft_positions: + common_attn_metadata.seq_lens_cpu[:batch_size] = common_attn_metadata.seq_lens_cpu[:batch_size] + 1 exceeds_mask_cpu = common_attn_metadata.seq_lens_cpu[:batch_size] > self.max_model_len common_attn_metadata.seq_lens_cpu[:batch_size].masked_fill_(exceeds_mask_cpu, 1) if common_attn_metadata._seq_lens_cpu is not None: - common_attn_metadata._seq_lens_cpu[:batch_size] = common_attn_metadata._seq_lens_cpu[:batch_size] + 1 + if advance_draft_positions: + common_attn_metadata._seq_lens_cpu[:batch_size] = common_attn_metadata._seq_lens_cpu[:batch_size] + 1 exceeds_mask_internal_cpu = common_attn_metadata._seq_lens_cpu[:batch_size] > self.max_model_len common_attn_metadata._seq_lens_cpu[:batch_size].masked_fill_(exceeds_mask_internal_cpu, 1) if common_attn_metadata.num_computed_tokens_cpu is not None: - common_attn_metadata.num_computed_tokens_cpu[:batch_size] += 1 + if advance_draft_positions: + common_attn_metadata.num_computed_tokens_cpu[:batch_size] += 1 if self.uses_mrope: common_attn_metadata.positions[:batch_size].copy_(clamped_positions[0]) else: diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index b0e532de95ea..a6142cec8df6 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -140,6 +140,7 @@ from vllm_ascend.spec_decode.extract_hidden_states_proposer import ( AscendExtractHiddenStatesProposer, ) +from vllm_ascend.spec_decode.gemma4_proposer import AscendGemma4Proposer from vllm_ascend.spec_decode.medusa_proposer import AscendMedusaProposer from vllm_ascend.spec_decode.ngram_proposer import AscendNgramProposer from vllm_ascend.spec_decode.ngram_proposer_npu import AscendNgramProposerNPU @@ -568,6 +569,7 @@ def _set_up_drafter(self): | AscendDraftModelProposer | AscendDflashProposer | AscendDSparkProposer + | AscendGemma4Proposer | AscendSuffixDecodingProposer | AscendMedusaProposer | AscendExtractHiddenStatesProposer @@ -3019,9 +3021,14 @@ def _build_attn_group_metadata( # build per-step attention metadata for the active MTP layer. self.drafter.set_per_group_attn_metadata( kv_cache_gid, cm.block_table_tensor, cm.slot_mapping) + elif self.speculative_config and isinstance(self.drafter, AscendGemma4Proposer): + self.drafter.set_per_group_block_table(kv_cache_gid, cm.block_table_tensor) if self.speculative_config and spec_decode_common_attn_metadata is None: - if isinstance(self.drafter, AscendEagleProposer | AscendDraftModelProposer | AscendDflashProposer - | AscendDSparkProposer): + if isinstance( + self.drafter, + AscendEagleProposer | AscendGemma4Proposer + | AscendDraftModelProposer | AscendDflashProposer | AscendDSparkProposer, + ): if self.drafter.attn_layer_names[0] in kv_cache_group.layer_names: spec_decode_common_attn_metadata = cm else: @@ -3586,11 +3593,24 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: ): assert isinstance( self.drafter, - AscendEagleProposer | AscendDflashProposer | AscendDSparkProposer | AscendDraftModelProposer, + AscendEagleProposer + | AscendGemma4Proposer + | AscendDflashProposer + | AscendDSparkProposer + | AscendDraftModelProposer, + ) + kernel_block_sizes = ( + self.kernel_block_sizes[0] + if isinstance(self.kernel_block_sizes, list) + else self.kernel_block_sizes ) - block_size = (self.kernel_block_sizes[0] if isinstance( - self.kernel_block_sizes, list) else self.kernel_block_sizes) - self.drafter.initialize_attn_backend(kv_cache_config, block_size) + if isinstance(self.drafter, AscendGemma4Proposer): + assert isinstance(self.kernel_block_sizes, list) + kernel_block_sizes = [ + sizes if isinstance(sizes, int) else sizes[0] + for sizes in self.kernel_block_sizes + ] + self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes) if has_kv_transfer_group(): get_kv_transfer_group().register_kv_caches(kv_caches) @@ -4697,7 +4717,7 @@ def _check_and_update_cudagraph_mode( ): assert isinstance( self.drafter, - AscendEagleProposer | AscendDflashProposer | AscendExtractHiddenStatesProposer, + AscendEagleProposer | AscendDflashProposer | AscendExtractHiddenStatesProposer | AscendGemma4Proposer, ) self.drafter.initialize_cudagraph_keys(cudagraph_mode)