[Feature] implement eagle spec decoding for model runner v2 (#5840)

### What this PR does / why we need it?
this pr implement eagle spec decoding for model runner v2, please see
RFC https://github.com/vllm-project/vllm-ascend/issues/5208

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?
vLLM version: v0.13.0

---------

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>

vllm_ascend/patch/__init__.py

@@ -174,7 +174,8 @@
 #
 # ** 6. File: worker/patch_triton.py**
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-#   1. `vllm.model_executor.layers.mamba.ops`, `vllm.model_executor.layers.fla.ops`
+#   1. `vllm.model_executor.layers.mamba.ops`, `vllm.model_executor.layers.fla.ops`,
+#      `vllm.v1.worker.gpu.sample.gumbel.gumbel_sample`
 #    Why:
 #       triton ops in vLLM perform not good on NPU. And there is no dispatch mechanism for triton ops.
 #    How：
@@ -263,3 +264,15 @@
 #    Future Plan:
 #       Remove this patch when vLLM support these operators.
 #
+# ** 12. File: worker/patch_v2_eagle.py**
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#   1. `vllm.v1.worker.gpu.spec_decode.eagle.EagleSpeculator.propose`
+#    Why:
+#       `propose` method use torch.gather, but the gather operator will
+#       pollute the arguments passed to it. the bug is reported to huawei
+#       CANN team, but not fixed yet.
+#    How：
+#       clone the out attribute ahead of gather to avoid the bug.
+#    Future Plan:
+#       Remove this patch when cann fix the gather bug.
+#

vllm_ascend/patch/worker/__init__.py

@@ -32,3 +32,4 @@ import vllm_ascend.patch.worker.patch_qwen3_next  # noqa
 import vllm_ascend.patch.worker.patch_qwen3_next_mtp  # noqa
 import vllm_ascend.patch.worker.patch_rejection_sampler  # noqa
 import vllm_ascend.patch.worker.patch_qwen3_next  # noqa
+import vllm_ascend.patch.worker.patch_v2_egale  # noqa

vllm_ascend/patch/worker/patch_triton.py

@@ -1,4 +1,5 @@
 import vllm.model_executor.layers.mamba.ops.causal_conv1d
+import vllm.v1.worker.gpu.sample.gumbel
 
 from vllm_ascend.ops.triton.fla.chunk import chunk_gated_delta_rule
 from vllm_ascend.ops.triton.fla.layernorm_guard import LayerNormFn
@@ -6,9 +7,12 @@ from vllm_ascend.ops.triton.fla.sigmoid_gating import \
     fused_recurrent_gated_delta_rule_fwd_kernel
 from vllm_ascend.ops.triton.mamba.causal_conv1d import (
     causal_conv1d_fn, causal_conv1d_update_npu)
+from vllm_ascend.worker.v2.sample.gumbel import \
+    gumbel_sample as ascend_gumbel_sample
 
 vllm.model_executor.layers.mamba.ops.causal_conv1d.causal_conv1d_update = causal_conv1d_update_npu
 vllm.model_executor.layers.mamba.ops.causal_conv1d.causal_conv1d_fn = causal_conv1d_fn
 vllm.model_executor.layers.fla.ops.fused_recurrent.fused_recurrent_gated_delta_rule_fwd_kernel = fused_recurrent_gated_delta_rule_fwd_kernel
 vllm.model_executor.layers.fla.ops.layernorm_guard.LayerNormFn = LayerNormFn
 vllm.model_executor.layers.fla.ops.chunk_gated_delta_rule = chunk_gated_delta_rule
+vllm.v1.worker.gpu.sample.gumbel.gumbel_sample = ascend_gumbel_sample

vllm_ascend/patch/worker/patch_v2_egale.py

@@ -0,0 +1,166 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/sample/spec_decode/eagle.py
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+import numpy as np
+import torch
+import vllm
+from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
+from vllm.v1.worker.gpu.spec_decode.eagle import (prepare_eagle_decode,
+                                                  prepare_eagle_inputs)
+
+from vllm_ascend.worker.v2.attn_utils import build_attn_metadata
+
+
+@torch.inference_mode()
+def propose(
+    self,
+    input_batch: InputBatch,
+    sampling_metadata: SamplingMetadata,
+    # [num_tokens, hidden_size]
+    last_hidden_states: torch.Tensor,
+    # num_layers x [num_tokens, hidden_size]
+    aux_hidden_states: list[torch.Tensor] | None,
+    # [num_reqs]
+    num_sampled: torch.Tensor,
+    # [num_reqs]
+    num_rejected: torch.Tensor,
+    # [num_reqs]
+    last_sampled: torch.Tensor,
+    # [num_reqs]
+    next_prefill_tokens: torch.Tensor,
+) -> torch.Tensor:
+    # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
+    # number of rejected tokens, we maintain the size of eagle's input_ids and
+    # hidden_states the same as the target model's. This means, we pad each
+    # request's query length to include any rejected positions. By doing so,
+    # we can also reuse the attention metadata (e.g., query_start_loc,
+    # seq_lens) of the target model.
+    if aux_hidden_states:
+        assert self.method == "eagle3"
+        hidden_states = self.model.combine_hidden_states(
+            torch.cat(aux_hidden_states, dim=-1))
+    else:
+        hidden_states = last_hidden_states
+    num_tokens = input_batch.num_tokens_after_padding
+    self.hidden_states[:num_tokens] = hidden_states
+
+    # Get the input ids and last token indices for the speculator.
+    last_token_indices = prepare_eagle_inputs(
+        self.input_buffers,
+        input_batch,
+        num_sampled,
+        num_rejected,
+        last_sampled,
+        next_prefill_tokens,
+    )
+
+    # Prefill: Run the eagle speculator with eager mode.
+    # TODO(woosuk): Support CUDA graph for prefill.
+    last_hidden_states, hidden_states = self.run_model(
+        num_tokens,
+        input_batch.attn_metadata,
+        num_tokens_across_dp=None,  # FIXME
+    )
+    sample_hidden_states = last_hidden_states[last_token_indices]
+    logits = self.model.compute_logits(sample_hidden_states)
+
+    num_reqs = input_batch.num_reqs
+    cu_num_logits = input_batch.cu_num_logits[:num_reqs]
+    # NOTE(woosuk): For draft sampling, we only consider the temperature
+    # and ignore the other sampling parameters such as top_k and top_p,
+    # for simplicity and performance.
+    # While this may slightly degrade the acceptance rate, it does not
+    # affect the output distribution after rejection sampling.
+    # NOTE(Ronald1995): torch.gather will pollute the cache such as self.input_buffers.positions
+    # the bug is reported to huawei CANN team, but not fixed yet.
+    # So we clone the tensors before calling torch.gather to avoid the issue.
+    temperature = self.temperature[:num_reqs].clone()
+    seeds = self.seeds[:num_reqs].clone()
+    pos = self.input_buffers.positions[:num_reqs].clone()
+    # Gather the values and copy them to the pre-allocated buffers.
+    torch.gather(sampling_metadata.temperature,
+                 0,
+                 cu_num_logits,
+                 out=temperature)
+    torch.gather(sampling_metadata.seeds, 0, cu_num_logits, out=seeds)
+    torch.gather(input_batch.positions, 0, last_token_indices, out=pos)
+    # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
+    # used for draft and target sampling.
+    draft_tokens = gumbel_sample(logits,
+                                 temperature,
+                                 seeds,
+                                 pos + 1,
+                                 apply_temperature=True)
+    if self.num_speculative_steps == 1:
+        # Early exit.
+        return draft_tokens.view(-1, 1)
+
+    # Save the draft tokens for the first step.
+    self.draft_tokens[:num_reqs, 0] = draft_tokens
+    # Prepare the inputs for the decode steps.
+    prepare_eagle_decode(
+        draft_tokens,
+        hidden_states,
+        last_token_indices,
+        input_batch.seq_lens,
+        num_rejected,
+        self.input_buffers,
+        self.hidden_states,
+        self.max_model_len,
+        self.max_num_reqs,
+    )
+    query_start_loc = self.input_buffers.query_start_loc
+    query_start_loc_gpu = query_start_loc.gpu[:num_reqs + 1]
+    slot_mappings = self.block_tables.compute_slot_mappings(
+        query_start_loc_gpu, pos)
+
+    cudagraph_size = self.cudagraph_manager.get_cudagraph_size(num_reqs)
+    if cudagraph_size is not None:
+        # Run CUDA graph.
+        self.cudagraph_manager.run(cudagraph_size)
+        return self.draft_tokens[:num_reqs]
+
+    # Run eager mode.
+    query_start_loc.np[:num_reqs + 1] = np.arange(num_reqs + 1)
+    query_start_loc_cpu = query_start_loc.cpu[:num_reqs + 1]
+    # HACK(woosuk)
+    seq_lens_np = np.full(num_reqs, self.max_model_len, dtype=np.int32)
+    block_tables = [x[:num_reqs] for x in self.block_tables.input_block_tables]
+
+    # FIXME(woosuk): This is UNSAFE!!
+    attn_metadata = build_attn_metadata(
+        attn_metadata_builders=self.attn_metadata_builders,
+        num_reqs=num_reqs,
+        num_tokens=num_reqs,
+        query_start_loc_gpu=query_start_loc_gpu,
+        query_start_loc_cpu=query_start_loc_cpu,
+        seq_lens=self.input_buffers.seq_lens[:num_reqs],
+        seq_lens_np=seq_lens_np,
+        num_computed_tokens_cpu=None,  # FIXME
+        block_tables=block_tables,
+        slot_mappings=slot_mappings,
+        kv_cache_config=self.kv_cache_config,
+    )
+    self.generate_draft(num_reqs, attn_metadata,
+                        num_tokens_across_dp=None)  # FIXME
+    return self.draft_tokens[:num_reqs]
+
+
+vllm.v1.worker.gpu.spec_decode.eagle.EagleSpeculator.propose = propose