[Bugfix] Honor STEP token pooling in DispatchPooler.for_seq_cls (#55307)
Signed-off-by: weihan <130119942+krsish@users.noreply.github.com>
weihan committed
Sep 8, 2026 at 12:13 UTC
9ea8f3ffc354901b740f0b31988900897b7221d7
3 files changed
+73
-47
examples/pooling/token_classify/forced_alignment_online.py
+24
-45
@@ -4,20 +4,25 @@
4
# https://github.com/QwenLM/Qwen3-ASR
5
6
"""
7
-Online forced alignment example using Qwen3-ForcedAligner-0.6B.
7
+Online forced alignment example using Qwen3-ForcedAligner-0.6B-hf.
8
9
Forced alignment takes audio and reference text as input and produces
10
word-level timestamps. The model predicts a time bin at each <timestamp>
11
token position; multiplying by ``timestamp_segment_time`` gives milliseconds.
12
13
+This example uses STEP pooling to return logits only for <timestamp> tokens.
14
+The ``step_tag_id`` value 151705 is this checkpoint's ``timestamp_token_id``
15
+from its config.json.
16
+
17
Start the server with:
18
15
- vllm serve Qwen/Qwen3-ForcedAligner-0.6B \\
16
- --runner pooling \\
17
- --enforce-eager \\
18
- --trust-request-chat-template \\
19
- --hf-overrides \\
20
- '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'
19
+ vllm serve Qwen/Qwen3-ForcedAligner-0.6B-hf \
20
+ --runner pooling \
21
+ --convert classify \
22
+ --pooler-config \
23
+ '{"tok_pooling_type":"STEP","step_tag_id":151705,"use_activation":false}' \
24
+ --enforce-eager \
25
+ --trust-request-chat-template
26
27
Then run:
28
@@ -108,7 +113,7 @@ def parse_response(response: requests.Response) -> dict[str, Any]:
113
return result
114
115
111
-def load_timestamp_config(model: str) -> tuple[int, float]:
116
+def load_timestamp_segment_time(model: str) -> float:
117
model_path = Path(model)
118
config_path = (
119
model_path / "config.json"
@@ -119,7 +124,7 @@ def load_timestamp_config(model: str) -> tuple[int, float]:
124
with config_path.open() as f:
125
config = json.load(f)
126
122
- return config["timestamp_token_id"], config["timestamp_segment_time"]
127
+ return config["timestamp_segment_time"]
128
129
130
def parse_args():
@@ -129,7 +134,7 @@ def parse_args():
134
parser.add_argument(
135
"--model",
136
type=str,
132
- default="Qwen/Qwen3-ForcedAligner-0.6B",
137
+ default="Qwen/Qwen3-ForcedAligner-0.6B-hf",
138
)
139
parser.add_argument(
140
"--audio-path",
@@ -147,8 +152,6 @@ def parse_args():
152
153
154
def main(args):
150
- from transformers import AutoTokenizer
151
-
155
api_url = f"http://{args.host}:{args.port}/pooling"
156
prompt = build_prompt(args.words)
157
audio_uri = (
@@ -161,46 +164,22 @@ def main(args):
164
pooling_response = post_http_request(payload=payload, api_url=api_url)
165
result = parse_response(pooling_response)
166
164
- tokenizer = AutoTokenizer.from_pretrained(args.model)
165
- timestamp_token_id, timestamp_segment_time = load_timestamp_config(args.model)
166
-
167
+ timestamp_segment_time = load_timestamp_segment_time(args.model)
168
output = result["data"][0]
169
logits = torch.tensor(output["data"])
170
predictions = logits.argmax(dim=-1)
170
- token_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
171
- audio_pad_token_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
172
-
173
- usage = result.get("usage") or {}
174
- prompt_tokens = usage.get("prompt_tokens")
175
- if prompt_tokens is not None and prompt_tokens != len(predictions):
176
- raise RuntimeError(
177
- "The response length does not match the reported prompt token count."
178
- )
171
180
- try:
181
- audio_pad_index = token_ids.index(audio_pad_token_id)
182
- except ValueError as exc:
183
- raise RuntimeError("The prompt does not contain the audio pad token.") from exc
184
-
185
- audio_token_shift = len(predictions) - len(token_ids)
186
- if audio_token_shift < 0:
172
+ expected_timestamps = len(args.words) * 2
173
+ if len(predictions) != expected_timestamps:
174
raise RuntimeError(
188
- "The response is shorter than the locally tokenized prompt. "
189
- "Check that the server was started with --trust-request-chat-template."
190
- )
191
-
192
- ts_predictions = []
193
- for i, token_id in enumerate(token_ids):
194
- if token_id != timestamp_token_id:
195
- continue
196
-
197
- prediction_index = i + audio_token_shift if i > audio_pad_index else i
198
- ts_predictions.append(
199
- predictions[prediction_index].item() * timestamp_segment_time
175
+ f"Expected {expected_timestamps} timestamp predictions, "
176
+ f"but received {len(predictions)}. Check that the server was started "
177
+ "with STEP pooling."
178
)
179
202
- if len(ts_predictions) < len(args.words) * 2:
203
- raise RuntimeError("The model did not return enough timestamp predictions.")
180
+ ts_predictions = [
181
+ prediction.item() * timestamp_segment_time for prediction in predictions
182
+ ]
183
184
for i, word in enumerate(args.words):
185
start_ms = ts_predictions[i * 2]
tests/model_executor/layers/test_pooler_methods.py
+48
@@ -3,12 +3,15 @@
3
"""Unit tests for sequence and token pooling methods and their factories."""
4
5
from dataclasses import dataclass
6
+from types import SimpleNamespace
7
from unittest.mock import patch
8
9
import numpy as np
10
import pytest
11
import torch
12
+from transformers import PretrainedConfig
13
14
+from vllm.config import PoolerConfig, set_current_vllm_config
15
from vllm.model_executor.layers.pooler import PoolingParamsUpdate
16
from vllm.model_executor.layers.pooler.seqwise.methods import (
17
CLSPool,
@@ -16,6 +19,7 @@ from vllm.model_executor.layers.pooler.seqwise.methods import (
19
MeanPool,
20
get_seq_pooling_method,
21
)
22
+from vllm.model_executor.layers.pooler.special import DispatchPooler
23
from vllm.model_executor.layers.pooler.tokwise.methods import (
24
AllPool,
25
StepPool,
@@ -300,6 +304,50 @@ class _FakeVllmConfig:
304
scheduler_config: _FakeSchedulerConfig
305
306
307
+@pytest.mark.parametrize("tok_pooling_type", ["ALL", "STEP"])
308
+def test_dispatch_seq_cls_honors_token_pooling_type(tok_pooling_type):
309
+ """Filter STEP rows before classification; preserve all rows for resolved ALL."""
310
+ pooler_config = PoolerConfig(
311
+ seq_pooling_type="LAST", tok_pooling_type=tok_pooling_type
312
+ )
313
+ config = SimpleNamespace(
314
+ scheduler_config=_FakeSchedulerConfig(),
315
+ model_config=SimpleNamespace(
316
+ pooler_config=pooler_config,
317
+ head_dtype=None,
318
+ hf_config=PretrainedConfig(num_labels=3),
319
+ ),
320
+ )
321
+ classifier = torch.nn.Linear(4, 3)
322
+ with set_current_vllm_config(config):
323
+ pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=classifier)
324
+
325
+ token_ids = [[2, 99, 3, 99], [4, 99], [5, 6]]
326
+ params = [
327
+ PoolingParams(
328
+ task="token_classify",
329
+ use_activation=False,
330
+ step_tag_id=99 if tok_pooling_type == "STEP" else None,
331
+ )
332
+ for _ in token_ids
333
+ ]
334
+ metadata = _make_metadata([4, 2, 2], token_ids=token_ids, pooling_params=params)
335
+ hidden = torch.arange(32, dtype=torch.float32).reshape(8, 4)
336
+ outputs = pooler(hidden, metadata)
337
+
338
+ expected_rows = (
339
+ [[1, 3], [5], []]
340
+ if tok_pooling_type == "STEP"
341
+ else [[0, 1, 2, 3], [4, 5], [6, 7]]
342
+ )
343
+ assert len(outputs) == len(expected_rows)
344
+ for output, rows in zip(outputs, expected_rows):
345
+ torch.testing.assert_close(output, classifier(hidden[rows]))
346
+ assert pooler.get_pooling_updates("token_classify").requires_token_ids == (
347
+ tok_pooling_type == "STEP"
348
+ )
349
+
350
+
351
class TestAllPool:
352
@staticmethod
353
def _make_all_pool(
vllm/model_executor/layers/pooler/special.py
+1
-2
@@ -19,7 +19,7 @@ from .seqwise import (
19
pooler_for_classify,
20
pooler_for_embed,
21
)
22
-from .tokwise import AllPool, pooler_for_token_classify, pooler_for_token_embed
22
+from .tokwise import pooler_for_token_classify, pooler_for_token_embed
23
24
25
class DispatchPooler(Pooler):
@@ -46,7 +46,6 @@ class DispatchPooler(Pooler):
46
{
47
"token_classify": pooler_for_token_classify(
48
pooler_config,
49
- pooling=AllPool(),
49
classifier=classifier,
50
),
51
"classify": pooler_for_classify(