xx

Author: wangxiaoxin-sherie

vllm_ascend/worker/model_runner_v1.py

@@ -2271,8 +2271,10 @@ def _build_dummy_attn_metadata(
             self.seq_lens_np[:num_reqs] = seq_lens
             self.seq_lens_np[num_reqs:] = 0
 
-            self.query_start_loc[:num_reqs + 1] = torch.arange(num_reqs + 1)
-            self.query_start_loc_cpu[:num_reqs + 1] = torch.arange(num_reqs + 1)
+            cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens)
+
+            self.query_start_loc[1:num_reqs + 1] = torch.Tensor(cu_num_tokens)
+            self.query_start_loc_cpu[1:num_reqs + 1] = torch.Tensor(cu_num_tokens)
 
             num_computed_tokens_cpu = (
                 self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs])
@@ -2427,6 +2429,7 @@ def _dummy_run(
             max_query_len=max_query_len,
             aclgraph_runtime_mode=aclgraph_runtime_mode,
             force_attention=force_attention,
+            num_scheduled_tokens=num_scheduled_tokens,
         )
 
         if not self.in_profile_run and self.dynamic_eplb: