Skip to content

Commit f731be8

Browse files
authored
[Feature][MRV2] GQA C8 modelrunnerv2 adaptation (vllm-project#12981)
### What this PR does / why we need it? This PR fixes a bug in `modelrunnerv1` and `modelrunnerv2` caused by inconsistent `attn_state` values ​​during graph mode compilation. ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? - vLLM version: v0.25.1 - vLLM main: vllm-project/vllm@fe784ff --------- Signed-off-by: pichangping <1337510399@qq.com>
1 parent 6fc3e33 commit f731be8

1 file changed

Lines changed: 7 additions & 4 deletions

File tree

vllm_ascend/attention/attention_v1.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1714,11 +1714,14 @@ def forward(
17141714
attn_output = self._forward_encoder_attention(query, key, value, attn_metadata, output)
17151715
output[:num_tokens] = attn_output[:num_tokens]
17161716
return output
1717+
1718+
# When `modelrunnerv2` compiles the graph, the value of `attn_metadata.attn_state` is `None`;
1719+
# therefore, the graph-mode condition needs to be evaluated earlier.
1720+
if _EXTRA_CTX.capturing:
1721+
attn_output, num_tokens = self.full_graph_fia(query, key, value, attn_metadata, output, layer)
1722+
output[:num_tokens] = attn_output[:num_tokens]
1723+
return output
17171724
if attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
1718-
if _EXTRA_CTX.capturing:
1719-
attn_output, num_tokens = self.full_graph_fia(query, key, value, attn_metadata, output, layer)
1720-
output[:num_tokens] = attn_output[:num_tokens]
1721-
return output
17221725
return self._forward_c8_decode(query, attn_metadata, output, layer)
17231726
elif attn_metadata.attn_state == AscendAttentionState.ChunkedPrefill:
17241727
return self._forward_c8_chunked_prefill(query, float_key, float_value, attn_metadata, output, layer)

0 commit comments

Comments
 (0)