diff --git a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py index 77447e8f..ef89d328 100644 --- a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py +++ b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py @@ -170,8 +170,8 @@ def test_llama_qwen3_eagle_correctness( "max_model_len": 128, }, compilation_config=CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - cudagraph_capture_sizes=[12])) as llm: + cudagraph_mode="FULL", + cudagraph_capture_sizes=[5, 12])) as llm: spec_outputs = llm.generate(example_prompts, sampling_params) cleanup_dist_env_and_memory() del llm diff --git a/vllm_ascend/spec_decode/eagle_proposer.py b/vllm_ascend/spec_decode/eagle_proposer.py index b8e0c2f9..b9d4e3bd 100644 --- a/vllm_ascend/spec_decode/eagle_proposer.py +++ b/vllm_ascend/spec_decode/eagle_proposer.py @@ -315,10 +315,11 @@ class AscendEagleProposer(EagleProposer): aclgraph_runtime_mode = CUDAGraphMode.NONE if aclgraph_runtime_mode == CUDAGraphMode.FULL and len(self.runner.attn_groups) > 0: num_computed_tokens_cpu = self.runner.input_batch.num_computed_tokens_cpu_tensor[:num_reqs] - self.query_start_loc.cpu[: num_reqs + 1] = torch.tensor( - [0] + self.runner.actual_seq_lengths_q[:num_reqs], device="cpu", dtype=torch.int32 - ) + + # num_reqs is already the padded version + self.query_start_loc.cpu[: num_reqs + 1].copy_(self.runner.query_start_loc.cpu[: num_reqs + 1]) self.query_start_loc.copy_to_gpu() + common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1], diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 8000dbca..a102d0b1 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -420,9 +420,6 @@ class NPUModelRunner(GPUModelRunner): assert isinstance(self.drafter, AscendEagleProposer) self.use_aux_hidden_state_outputs = self.drafter.eagle3_use_aux_hidden_state self.rejection_sampler = RejectionSampler(self.sampler) - self.actual_seq_lengths_q = list( - range(self.decode_token_per_req, self.max_num_tokens + 1, self.decode_token_per_req) - ) self.discard_request_indices = self._make_buffer(self.max_num_reqs, dtype=torch.int64) self.num_discarded_requests = 0