Skip to content

Commit 6238004

Browse files
committed
Update test_attention_v1.py
1 parent 77cac33 commit 6238004

File tree

1 file changed

+3
-4
lines changed

1 file changed

+3
-4
lines changed

tests/ut/attention/test_attention_v1.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -429,10 +429,9 @@ def test_forward_decode_only(self, mock_paged_attention,
429429
mock_paged_attention.assert_called_once()
430430
assert output.shape == (10, 8 * 64)
431431

432-
@patch('torch_npu._npu_reshape_and_cache')
433-
@patch('torch_npu._npu_paged_attention')
434432
@patch('torch_npu.npu_fused_infer_attention_score')
435-
def test_forward_decode_only_swa(self, mock_paged_attention,
433+
@patch('torch_npu._npu_reshape_and_cache')
434+
def test_forward_decode_only_swa(self, mock_fused_infer_attention_score,
436435
mock_npu_reshape_and_cache):
437436
"""Test forward pass in DecodeOnly state"""
438437
query = torch.randn(10, 8 * 64)
@@ -455,7 +454,7 @@ def test_forward_decode_only_swa(self, mock_paged_attention,
455454
metadata,
456455
trace_flag=False)
457456

458-
mock_paged_attention.assert_called_once()
457+
mock_fused_infer_attention_score.assert_called_once()
459458
assert output.shape == (10, 8 * 64)
460459

461460
@patch('vllm_ascend.attention.attention_v1.is_310p', return_value=False)

0 commit comments

Comments
 (0)