Skip to content

Commit e881974

Browse files
zhoux77899zhoux77899
authored
[Refactor][Ops] Refactor multi-modal encoder attention sequence length handling (vllm-project#12084)
### What this PR does / why we need it? This PR unifies FIA `actual_seq_lengths` handling for ViT encoder ACL graph across eager, capture, and replay. Previously, `AscendMMEncoderAttention` and `encoder_acl_graph` built sequence lengths independently, which could produce incorrect results for padded `cu_seqlens` and budget-padding trailing zeros. Fix vllm-project#10549 and vllm-project#10824 Key changes: - Introduce shared helper `maybe_compute_actual_seq_lengths()` in `encoder_acl_graph.py` to convert `cu_seqlens` → host `list[int]`; in graph mode filter invalid endpoints and align `actual_seq_lengths[-1] == num_query_tokens`; when Q≠KV, scale KV endpoints by `kv/q`. - Update `_maybe_compute_cu_seqlens()` for `None` fallback, CPU normalization, and capture-time uniform `cu_seqlens`. - Route eager/capture/replay through the shared helper; simplify replay context to only `cu_seqlens_cpu`. - Remove redundant per-layer metadata (`vit_layer_idx`, `fullatt_block_indexes`, `sequence_lengths`, `cu_window_seqlens`) and use `self.scale` instead of `scale_value`. ```mermaid flowchart TB direction TB Before["Before: duplicated paths"] --> E1[cu_seqlens / sequence_lengths / cu_window_seqlens] E1 --> E2[AscendMMEncoderAttention<br/>local convert + per-layer metadata] E2 --> E3[_pad_actual_seq_lengths_for_fia] Before --> R1[cu_seqlens_cpu / cu_window_seqlens_cpu<br/>sequence_lengths_cpu + fullatt_block_indexes] R1 --> R2[encoder_acl_graph<br/>_maybe_compute_actual_seq_lengths] R2 --> R3[_pad_actual_seq_lengths_for_fia] E3 --> FIA1[FIA] R3 --> FIA1 After["After: shared pipeline"] --> IN[cu_seqlens] IN --> BUILD[maybe_compute_actual_seq_lengths<br/>filter + budget align] BUILD --> EAGER[eager / capture] BUILD --> REPLAY[replay via cu_seqlens_cpu context] EAGER --> FIA2[FIA actual_seq_lengths] REPLAY --> FIA2 ``` ### Does this PR introduce _any_ user-facing change? N/A. ### How was this patch tested? - The following UT added to verify the eager and capture paths as well as the sequence length computation. - `tests/ut/worker/test_encoder_acl_graph.py` - `tests/ut/ops/test_mm_encoder_attention.py` - Also test on `Qwen/Qwen3.6-27B`. ```diff | dataset | version | metric | mode | vllm-api-stream-chat | |----- | ----- | ----- | ----- | -----| - | videomme | 2f2602 | [Overall][overall] | gen | 67.78 | + | videomme | 2f2602 | [Overall][overall] | gen | 68.85 | ``` - vLLM version: v0.24.0 - vLLM main: vllm-project/vllm@85c09e9 --------- Signed-off-by: zhoux77899 <zhouxiang100@huawei.com>
1 parent 20c7cb8 commit e881974

5 files changed

Lines changed: 474 additions & 147 deletions

File tree

‎tests/ut/conftest.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,8 @@
107107
sys.modules["torch_npu"]._npu_reshape_and_cache = MagicMock() # type: ignore[attr-defined]
108108
sys.modules["torch_npu"].npu_scatter_pa_kv_cache = MagicMock() # type: ignore[attr-defined]
109109
sys.modules["torch_npu"].npu_gather_pa_kv_cache = MagicMock() # type: ignore[attr-defined]
110+
sys.modules["torch_npu"].npu_fused_infer_attention_score = MagicMock() # type: ignore[attr-defined]
111+
sys.modules["torch_npu"]._npu_fused_infer_attention_score_get_max_workspace = MagicMock() # type: ignore[attr-defined]
110112
sys.modules["torch_npu"].npu_moe_gating_top_k_softmax = MagicMock() # type: ignore[attr-defined]
111113
sys.modules["torch_npu"].npu_quant_matmul = MagicMock() # type: ignore[attr-defined]
112114
sys.modules["torch_npu"].npu_rms_norm = MagicMock() # type: ignore[attr-defined]

‎tests/ut/ops/test_mm_encoder_attention.py‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
from typing import Any
2+
from unittest.mock import MagicMock, patch
3+
4+
import torch
5+
from vllm.config import CompilationConfig, VllmConfig
6+
from vllm.config.vllm import get_cached_compilation_config
7+
8+
from tests.ut.base import TestBase
9+
from vllm_ascend.ops.mm_encoder_attention import (
10+
MAX_PAD_SIZE,
11+
AscendMMEncoderAttention,
12+
)
13+
from vllm_ascend.worker import encoder_acl_graph
14+
from vllm_ascend.worker.encoder_acl_graph import (
15+
get_encoder_graph_params,
16+
set_encoder_forward_context,
17+
set_encoder_graph_params,
18+
)
19+
20+
21+
class FIAMockMixin(TestBase):
22+
captured: dict[str, Any]
23+
24+
def _install_vllm_config_mock(self):
25+
mock_vllm_config = MagicMock(spec=VllmConfig)
26+
mock_vllm_config.compilation_config = CompilationConfig()
27+
patcher = patch(
28+
"vllm.config.vllm.get_current_vllm_config",
29+
return_value=mock_vllm_config,
30+
)
31+
patcher.start()
32+
self.addCleanup(patcher.stop)
33+
get_cached_compilation_config.cache_clear()
34+
self.addCleanup(get_cached_compilation_config.cache_clear)
35+
36+
def _make_layer(self, num_heads=4, num_kv_heads=4, head_size=72, scale=None):
37+
return AscendMMEncoderAttention(
38+
num_heads=num_heads,
39+
head_size=head_size,
40+
scale=scale,
41+
num_kv_heads=num_kv_heads,
42+
)
43+
44+
def _fake_fia(self, **kwargs):
45+
self.captured = {
46+
"mode": "functional",
47+
"q_shape": kwargs["query"].shape,
48+
"input_layout": kwargs["input_layout"],
49+
"actual_seq_lengths": kwargs["actual_seq_lengths"],
50+
}
51+
return torch.zeros_like(kwargs["query"]), None
52+
53+
def _fake_fia_out(self, *, workspace, out, **kwargs):
54+
self.captured = {"mode": "out", "softmax_lse": out[1]}
55+
out[0].zero_()
56+
57+
def _install_fia_mocks(self, *, capture: bool):
58+
self.captured = {}
59+
mock_fia = MagicMock(side_effect=self._fake_fia)
60+
mock_fia.out = self._fake_fia_out
61+
62+
patch_targets: list[tuple[str, Any]] = [
63+
(
64+
"vllm_ascend.ops.mm_encoder_attention.torch_npu.npu_fused_infer_attention_score",
65+
mock_fia,
66+
),
67+
(
68+
"vllm_ascend.ops.mm_encoder_attention.torch_npu._npu_fused_infer_attention_score_get_max_workspace",
69+
MagicMock(return_value=torch.zeros(1)),
70+
),
71+
]
72+
if capture:
73+
self.mock_graph_begin = MagicMock()
74+
self.mock_graph_end = MagicMock(return_value=42)
75+
mock_event = MagicMock()
76+
patch_targets.extend(
77+
[
78+
(
79+
"vllm_ascend.ops.mm_encoder_attention.weak_ref_tensors",
80+
lambda tensors: tensors,
81+
),
82+
(
83+
"vllm_ascend.ops.mm_encoder_attention.torch_npu.npu.current_stream",
84+
MagicMock(return_value=MagicMock()),
85+
),
86+
(
87+
"vllm_ascend.ops.mm_encoder_attention.torch.npu.ExternalEvent",
88+
MagicMock(return_value=mock_event),
89+
),
90+
(
91+
"vllm_ascend.ops.mm_encoder_attention.torch.npu.graph_task_group_begin",
92+
self.mock_graph_begin,
93+
),
94+
(
95+
"vllm_ascend.ops.mm_encoder_attention.torch.npu.graph_task_group_end",
96+
self.mock_graph_end,
97+
),
98+
]
99+
)
100+
101+
for target, replacement in patch_targets:
102+
patcher = patch(target, replacement)
103+
patcher.start()
104+
self.addCleanup(patcher.stop)
105+
106+
107+
class TestAscendMMEncoderAttentionEager(FIAMockMixin):
108+
def setUp(self):
109+
self._install_vllm_config_mock()
110+
self._install_fia_mocks(capture=False)
111+
112+
def test_forward_oot_basic(self):
113+
layer = self._make_layer(num_heads=4, num_kv_heads=4, head_size=128)
114+
bsz, q_len = 2, 4
115+
query = torch.randn(bsz, q_len, layer.num_heads * layer.head_size)
116+
key = query.clone()
117+
value = query.clone()
118+
cu_seqlens = torch.arange(0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32)
119+
120+
out = layer.forward_oot(query, key, value, cu_seqlens=cu_seqlens)
121+
122+
self.assertEqual(out.shape, (bsz, q_len, layer.num_heads * layer.head_size))
123+
self.assertEqual(self.captured["mode"], "functional")
124+
self.assertEqual(self.captured["input_layout"], "TND")
125+
126+
def test_forward_oot_seqlens(self):
127+
layer = self._make_layer(num_heads=4, num_kv_heads=4, head_size=72)
128+
seq_lens = [3, 7, 2]
129+
cu_seqlens = torch.tensor([0, 3, 10, 12], dtype=torch.int32, device="cpu")
130+
max_q_len = max(seq_lens)
131+
query = torch.randn(len(seq_lens), max_q_len, layer.num_heads, 72, dtype=torch.bfloat16)
132+
key = torch.randn_like(query)
133+
value = torch.randn_like(query)
134+
135+
out = layer.forward_oot(query, key, value, cu_seqlens=cu_seqlens)
136+
137+
self.assertEqual(out.shape, query.shape)
138+
self.assertEqual(self.captured["actual_seq_lengths"], [3, 10, 12])
139+
self.assertEqual(self.captured["q_shape"], (len(seq_lens) * max_q_len, 4, MAX_PAD_SIZE))
140+
141+
142+
class TestAscendMMEncoderAttentionCapture(FIAMockMixin):
143+
def setUp(self):
144+
self._install_vllm_config_mock()
145+
set_encoder_graph_params([2048])
146+
self._install_fia_mocks(capture=True)
147+
148+
def tearDown(self):
149+
encoder_acl_graph._encoder_graph_params = None
150+
encoder_acl_graph._reset_encoder_forward_context()
151+
152+
def test_forward_oot_basic(self):
153+
layer = self._make_layer(num_heads=4, num_kv_heads=4, head_size=72)
154+
bsz, q_len = 2, 4
155+
query = torch.randn(bsz, q_len, layer.num_heads, 72, dtype=torch.bfloat16)
156+
key = torch.randn_like(query)
157+
value = torch.randn_like(query)
158+
cu_seqlens = torch.arange(0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32)
159+
160+
with set_encoder_forward_context(2048, True):
161+
layer.forward_oot(query, key, value, cu_seqlens=cu_seqlens)
162+
163+
params = get_encoder_graph_params()
164+
self.assertIsNotNone(params)
165+
self.assertEqual(len(params.attn_params[2048]), 1)
166+
self.assertEqual(len(params.handles[2048]), 1)
167+
self.assertEqual(self.captured["mode"], "out")
168+
self.mock_graph_begin.assert_called_once()
169+
self.mock_graph_end.assert_called_once()
170+
171+
def test_forward_oot_seqlens(self):
172+
layer = self._make_layer(num_heads=4, num_kv_heads=4, head_size=72)
173+
seq_lens = [3, 7, 2]
174+
cu_seqlens = torch.tensor([0, 3, 10, 12], dtype=torch.int32, device="cpu")
175+
max_q_len = max(seq_lens)
176+
query = torch.randn(len(seq_lens), max_q_len, layer.num_heads, 72, dtype=torch.bfloat16)
177+
key = torch.randn_like(query)
178+
value = torch.randn_like(query)
179+
180+
captured_lengths: list[Any] = []
181+
182+
def capture_workspace(**kwargs):
183+
captured_lengths.append(kwargs.get("actual_seq_lengths"))
184+
return torch.zeros(1)
185+
186+
with (
187+
patch(
188+
"vllm_ascend.ops.mm_encoder_attention.torch_npu._npu_fused_infer_attention_score_get_max_workspace",
189+
side_effect=capture_workspace,
190+
),
191+
set_encoder_forward_context(2048, True),
192+
):
193+
layer.forward_oot(query, key, value, cu_seqlens=cu_seqlens)
194+
195+
self.assertEqual(captured_lengths[-1], [7, 14, 21])

‎tests/ut/worker/test_encoder_acl_graph.py‎

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
1+
from types import SimpleNamespace
2+
from unittest.mock import MagicMock, patch
3+
4+
import pytest
5+
import torch
6+
from vllm.config import CompilationConfig, VllmConfig
7+
8+
from vllm_ascend.worker import encoder_acl_graph
9+
from vllm_ascend.worker.encoder_acl_graph import (
10+
EncoderAclGraphManager,
11+
get_encoder_forward_context,
12+
get_encoder_graph_params,
13+
maybe_compute_actual_seq_lengths,
14+
set_encoder_graph_params,
15+
update_encoder_graph_params,
16+
)
17+
18+
19+
def _reset_encoder_acl_graph_state() -> None:
20+
encoder_acl_graph._encoder_graph_params = None
21+
encoder_acl_graph._reset_encoder_forward_context()
22+
23+
24+
@pytest.fixture(autouse=True)
25+
def _reset_state():
26+
_reset_encoder_acl_graph_state()
27+
yield
28+
_reset_encoder_acl_graph_state()
29+
30+
31+
@pytest.mark.parametrize(
32+
"cu_seqlens, num_tokens, expected",
33+
[
34+
(torch.tensor([0, 4, 16], dtype=torch.int32), 8, [4, 16]),
35+
],
36+
)
37+
def test_maybe_compute_actual_seq_lengths_eager(cu_seqlens, num_tokens, expected):
38+
actual_q, actual_kv = maybe_compute_actual_seq_lengths(
39+
cu_seqlens,
40+
num_tokens,
41+
num_tokens,
42+
cudagraph_mm_encoder=False,
43+
)
44+
assert actual_q == expected
45+
assert actual_kv == expected
46+
47+
48+
def test_maybe_compute_actual_seq_lengths_eager_unequal_q_kv():
49+
"""Molmo-style uniform cross-attention: scale KV endpoints by kv/q ratio."""
50+
cu_seqlens = torch.tensor([0, 1, 2], dtype=torch.int32)
51+
actual_q, actual_kv = maybe_compute_actual_seq_lengths(
52+
cu_seqlens,
53+
2,
54+
8,
55+
cudagraph_mm_encoder=False,
56+
)
57+
assert actual_q == [1, 2]
58+
assert actual_kv == [4, 8]
59+
60+
61+
@pytest.mark.parametrize(
62+
"cu_seqlens, num_tokens, expected",
63+
[
64+
(torch.tensor([0, 4, 8], dtype=torch.int32), 8, [4, 8]),
65+
(torch.tensor([0, 4, 16], dtype=torch.int32), 8, [4, 8]),
66+
],
67+
)
68+
def test_maybe_compute_actual_seq_lengths_graph(cu_seqlens, num_tokens, expected):
69+
actual_q, actual_kv = maybe_compute_actual_seq_lengths(
70+
cu_seqlens,
71+
num_tokens,
72+
num_tokens,
73+
cudagraph_mm_encoder=True,
74+
)
75+
assert actual_q == expected
76+
assert actual_kv == expected
77+
78+
79+
def test_update_encoder_graph_params_cu_seqlens():
80+
set_encoder_graph_params([2048])
81+
params = get_encoder_graph_params()
82+
query = MagicMock()
83+
query.shape = [8, 4, 72]
84+
key = MagicMock()
85+
key.shape = [8, 4, 72]
86+
packed = (
87+
query,
88+
key,
89+
MagicMock(),
90+
None,
91+
None,
92+
128,
93+
4,
94+
4,
95+
0.125,
96+
MagicMock(),
97+
MagicMock(),
98+
)
99+
params.handles[2048] = [1]
100+
params.events[2048] = [MagicMock()]
101+
params.attn_params[2048] = [packed]
102+
params.workspaces[2048] = MagicMock()
103+
104+
ctx = get_encoder_forward_context()
105+
ctx.cu_seqlens_cpu = torch.tensor([0, 4, 8], dtype=torch.int32)
106+
107+
captured = {}
108+
109+
def fake_out(**kwargs):
110+
captured["actual_seq_lengths"] = kwargs["actual_seq_lengths"]
111+
112+
fake_fia = SimpleNamespace(out=fake_out)
113+
with (
114+
patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.stream"),
115+
patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.graph_task_update_begin"),
116+
patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.graph_task_update_end"),
117+
patch(
118+
"vllm_ascend.worker.encoder_acl_graph.torch_npu.npu_fused_infer_attention_score",
119+
fake_fia,
120+
),
121+
):
122+
update_encoder_graph_params(MagicMock(), 2048)
123+
124+
assert captured["actual_seq_lengths"] == [4, 8]
125+
126+
127+
def _make_manager():
128+
vllm_config = MagicMock(spec=VllmConfig)
129+
vllm_config.compilation_config = CompilationConfig()
130+
mm_config = MagicMock()
131+
mm_config.get_limit_per_prompt.return_value = 0
132+
mm_config.mm_encoder_tp_mode = "tensor"
133+
vllm_config.model_config = MagicMock()
134+
vllm_config.model_config.multimodal_config = mm_config
135+
vllm_config.parallel_config = MagicMock()
136+
vllm_config.parallel_config.tensor_parallel_size = 1
137+
138+
model = MagicMock()
139+
model.get_encoder_cudagraph_config.return_value = MagicMock(
140+
modalities=["image"],
141+
buffer_keys=["cu_seqlens"],
142+
out_hidden_size=64,
143+
enable_dual_path_graph=False,
144+
padding_logics={},
145+
max_frames_per_video=1,
146+
)
147+
model.get_encoder_cudagraph_budget_range.return_value = (64, 2048)
148+
return EncoderAclGraphManager(vllm_config, "npu", "bfloat16", model), model
149+
150+
151+
def test_capture_graph_params():
152+
mgr, _ = _make_manager()
153+
mgr.token_budgets = [2048]
154+
155+
with patch("vllm.v1.worker.encoder_cudagraph.EncoderCudaGraphManager.capture", return_value=None):
156+
mgr.capture()
157+
158+
params = get_encoder_graph_params()
159+
assert params is not None
160+
assert 2048 in params.events
161+
162+
163+
def test_capture_budget_graph_npu():
164+
mgr, model = _make_manager()
165+
mgr.max_batch_size = 2
166+
mgr.max_frames_per_batch = 0
167+
capture_values = {"cu_seqlens": torch.zeros(3, dtype=torch.int32)}
168+
model.prepare_encoder_cudagraph_capture_inputs.return_value = MagicMock(
169+
values=capture_values,
170+
)
171+
model.encoder_cudagraph_forward.return_value = torch.zeros(2, 64)
172+
173+
fake_graph = MagicMock()
174+
with (
175+
patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.NPUGraph", return_value=fake_graph),
176+
patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.graph"),
177+
patch(
178+
"vllm_ascend.worker.encoder_acl_graph.weak_ref_tensors",
179+
side_effect=lambda tensors: tensors,
180+
),
181+
):
182+
mgr._capture_budget_graph(2048)
183+
184+
graph_meta = mgr._get_graph_set("default")[2048]
185+
assert graph_meta.graph is fake_graph
186+
assert graph_meta.input_buffers is capture_values

0 commit comments

Comments
 (0)