Skip to content

Commit d3d63d2

Browse files
committed
refact attn metadata build
Signed-off-by: weiguihua2 <[email protected]>
1 parent 8c74a86 commit d3d63d2

File tree

3 files changed

+3
-5
lines changed

3 files changed

+3
-5
lines changed

vllm_ascend/attention/attention_v1_torchair.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -175,8 +175,7 @@ def build_torchair_graph_dummy(
175175
self, common_attn_metadata: AscendCommonAttentionMetadata) -> AscendTorchairMetadata:
176176
device = self.device
177177
num_reqs = common_attn_metadata.num_reqs
178-
_, max_blocks = self.max_blocks
179-
block_table = torch.zeros((num_reqs, max_blocks),
178+
block_table = torch.zeros((num_reqs, self.max_blocks),
180179
dtype=torch.int32,
181180
device=device)
182181
block_table = self._get_graph_runner_block_tables(

vllm_ascend/attention/mla_v1.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -287,8 +287,7 @@ def build_torchair_graph_dummy(
287287
self, common_attn_metadata: AscendCommonAttentionMetadata,) -> AscendMLAMetadata:
288288
device = self.device
289289
num_reqs = common_attn_metadata.num_reqs
290-
_, max_blocks = self.max_blocks
291-
block_table = torch.zeros((num_reqs, max_blocks),
290+
block_table = torch.zeros((num_reqs, self.max_blocks),
292291
dtype=torch.int32,
293292
device=device)
294293
block_table = self._get_graph_runner_block_tables(

vllm_ascend/torchair/torchair_model_runner.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ def _build_attention_metadata(self, with_prefill, num_reqs, skip_attn):
7676
actual_seq_lengths_q=self.actual_seq_lengths_q,
7777
attn_mask=self.attn_mask,
7878
spec_attn_mask=self.spec_attn_mask,
79-
attn_state=self.attn_state,
79+
decode_token_per_req=self.decode_token_per_req,
8080
)
8181
attn_metadata = self.attn_metadata_builder.build_torchair_graph_dummy(common_attn_metadata)
8282
else:

0 commit comments

Comments
 (0)