@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / de42abb366
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 07:30 UTC

[CI] Heavy refactoring of Voxtral multimodal audio model tests (#34294)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Feb 13, 2026 at 22:04 UTC de42abb366032519bca073e057331ead6270e09f
11 files changed +350 -70
requirements/rocm-test.txt
+2
@@ -96,3 +96,5 @@ albumentations==1.4.6
96 transformers==4.57.3
97 # Pin HF Hub version
98 huggingface-hub==0.36.2
99 +# Pin Mistral Common
100 +mistral-common[image,audio]==1.9.1
tests/conftest.py
-2
@@ -419,7 +419,6 @@ class HfRunner:
419 self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
420 AutoTokenizer.from_pretrained(
421 model_name,
422 - dtype=dtype,
422 trust_remote_code=trust_remote_code,
423 )
424 )
@@ -430,7 +429,6 @@ class HfRunner:
429
430 self.processor = AutoProcessor.from_pretrained(
431 model_name,
433 - dtype=dtype,
432 trust_remote_code=trust_remote_code,
433 )
434 if skip_tokenizer_init:
tests/models/multimodal/generation/test_voxtral.py
+138 -43
@@ -4,16 +4,18 @@
4 import json
5
6 import pytest
7 -import pytest_asyncio
7 from mistral_common.audio import Audio
8 from mistral_common.protocol.instruct.chunk import AudioChunk, RawAudio, TextChunk
9 from mistral_common.protocol.instruct.messages import UserMessage
10 +from transformers import VoxtralForConditionalGeneration
11
12 from vllm.tokenizers.mistral import MistralTokenizer
13
14 from ....conftest import AudioTestAssets
15 from ....utils import RemoteOpenAIServer
16 +from ...utils import check_logprobs_close
17 from .test_ultravox import MULTI_AUDIO_PROMPT, run_multi_audio_test
18 +from .vlm_utils import model_utils
19
20 MODEL_NAME = "mistralai/Voxtral-Mini-3B-2507"
21 MISTRAL_FORMAT_ARGS = [
@@ -26,40 +28,21 @@ MISTRAL_FORMAT_ARGS = [
28 ]
29
30
29 -@pytest.fixture()
30 -def server(request, audio_assets: AudioTestAssets):
31 - args = [
32 - "--enforce-eager",
33 - "--limit-mm-per-prompt",
34 - json.dumps({"audio": len(audio_assets)}),
35 - ] + MISTRAL_FORMAT_ARGS
36 -
37 - with RemoteOpenAIServer(
38 - MODEL_NAME, args, env_dict={"VLLM_AUDIO_FETCH_TIMEOUT": "30"}
39 - ) as remote_server:
40 - yield remote_server
41 -
42 -
43 -@pytest_asyncio.fixture
44 -async def client(server):
45 - async with server.get_async_client() as async_client:
46 - yield async_client
47 -
48 -
49 -def _get_prompt(audio_assets, question):
31 +def _get_prompt(audio_assets: AudioTestAssets, question: str) -> list[int]:
32 + """Build a token-ID prompt via mistral_common for vLLM offline inference."""
33 tokenizer = MistralTokenizer.from_pretrained(MODEL_NAME)
34
35 audios = [
53 - Audio.from_file(str(audio_assets[i].get_local_path()), strict=False)
54 - for i in range(len(audio_assets))
36 + Audio.from_file(str(asset.get_local_path()), strict=False)
37 + for asset in audio_assets
38 ]
39 audio_chunks = [
40 AudioChunk(input_audio=RawAudio.from_audio(audio)) for audio in audios
41 ]
42
60 - text_chunk = TextChunk(text=question)
61 - messages = [UserMessage(content=[*audio_chunks, text_chunk]).to_openai()]
62 -
43 + messages = [
44 + UserMessage(content=[*audio_chunks, TextChunk(text=question)]).to_openai()
45 + ]
46 return tokenizer.apply_chat_template(messages=messages)
47
48
@@ -77,7 +60,7 @@ def test_models_with_multiple_audios(
60 vllm_prompt = _get_prompt(audio_assets, MULTI_AUDIO_PROMPT)
61 run_multi_audio_test(
62 vllm_runner,
80 - [(vllm_prompt, [audio.audio_and_sample_rate for audio in audio_assets])],
63 + [(vllm_prompt, [a.audio_and_sample_rate for a in audio_assets])], # type: ignore[list-item]
64 MODEL_NAME,
65 dtype=dtype,
66 max_tokens=max_tokens,
@@ -86,30 +69,142 @@ def test_models_with_multiple_audios(
69 )
70
71
89 -@pytest.mark.asyncio
90 -async def test_online_serving(client, audio_assets: AudioTestAssets):
91 - """Exercises online serving with/without chunked prefill enabled."""
72 +def test_online_serving(vllm_runner, audio_assets: AudioTestAssets):
73 + """Two-layer accuracy and serving validation using Mistral format.
74 +
75 + 1. Offline vLLM greedy output (runs first to avoid CUDA fork issues
76 + with multiprocessing - see vlm_utils/core.py).
77 + 2. Online OpenAI-compatible API output must match offline — validates
78 + that the serving path (chat template, audio encoding, tokenization)
79 + does not corrupt anything.
80 +
81 + Steps run sequentially so each releases the GPU before the next starts.
82 + """
83
93 - def asset_to_chunk(asset):
84 + question = f"What's happening in these {len(audio_assets)} audio clips?"
85 + max_tokens = 10
86 + audio_data = [asset.audio_and_sample_rate for asset in audio_assets]
87 +
88 + vllm_prompt = _get_prompt(audio_assets, question)
89 + with vllm_runner(
90 + MODEL_NAME,
91 + dtype="half",
92 + enforce_eager=True,
93 + tokenizer_mode="mistral",
94 + config_format="mistral",
95 + load_format="mistral",
96 + limit_mm_per_prompt={"audio": len(audio_assets)},
97 + ) as vllm_model:
98 + offline_outputs = vllm_model.generate_greedy(
99 + [vllm_prompt],
100 + max_tokens,
101 + audios=[audio_data],
102 + )
103 +
104 + offline_text = offline_outputs[0][1]
105 + assert offline_text, "Offline vLLM inference produced empty output"
106 +
107 + def _asset_to_openai_chunk(asset):
108 audio = Audio.from_file(str(asset.get_local_path()), strict=False)
109 audio.format = "wav"
96 - audio_dict = AudioChunk.from_audio(audio).to_openai()
97 - return audio_dict
110 + return AudioChunk.from_audio(audio).to_openai()
111
99 - audio_chunks = [asset_to_chunk(asset) for asset in audio_assets]
100 - text = f"What's happening in these {len(audio_assets)} audio clips?"
112 messages = [
113 {
114 "role": "user",
104 - "content": [*audio_chunks, {"type": "text", "text": text}],
115 + "content": [
116 + *[_asset_to_openai_chunk(a) for a in audio_assets],
117 + {"type": "text", "text": question},
118 + ],
119 }
120 ]
121
108 - chat_completion = await client.chat.completions.create(
109 - model=MODEL_NAME, messages=messages, max_tokens=10
110 - )
122 + server_args = [
123 + "--enforce-eager",
124 + "--limit-mm-per-prompt",
125 + json.dumps({"audio": len(audio_assets)}),
126 + *MISTRAL_FORMAT_ARGS,
127 + ]
128
112 - assert len(chat_completion.choices) == 1
113 - choice = chat_completion.choices[0]
114 - assert choice.message.content == "In the first audio clip, you hear a brief"
129 + with RemoteOpenAIServer(
130 + MODEL_NAME,
131 + server_args,
132 + env_dict={"VLLM_AUDIO_FETCH_TIMEOUT": "30"},
133 + ) as remote_server:
134 + client = remote_server.get_client()
135 + completion = client.chat.completions.create(
136 + model=MODEL_NAME,
137 + messages=messages,
138 + max_tokens=max_tokens,
139 + temperature=0,
140 + )
141 +
142 + assert len(completion.choices) == 1
143 + choice = completion.choices[0]
144 assert choice.finish_reason == "length"
145 + assert choice.message.content == offline_text, (
146 + f"Online serving output does not match offline inference.\n"
147 + f" Online: {choice.message.content!r}\n"
148 + f" Offline: {offline_text!r}"
149 + )
150 +
151 +
152 +def test_hf_reference(hf_runner, vllm_runner, audio_assets: AudioTestAssets):
153 + """Compare vLLM Mistral-format output against HF Transformers reference.
154 +
155 + Instead of requiring an exact text match (which is brittle across
156 + attention backends), we compare per-token logprobs using the standard
157 + check_logprobs_close helper: when tokens diverge at a position, each
158 + runner's chosen token must appear in the other's top-k logprobs.
159 +
160 + Marked xfail(strict=False) so remaining edge-case mismatches
161 + don't block CI.
162 + """
163 + question = f"What's happening in these {len(audio_assets)} audio clips?"
164 + max_tokens = 10
165 + num_logprobs = 5
166 + audio_data = [asset.audio_and_sample_rate for asset in audio_assets]
167 +
168 + vllm_prompt = _get_prompt(audio_assets, question)
169 + with vllm_runner(
170 + MODEL_NAME,
171 + dtype="half",
172 + enforce_eager=True,
173 + tokenizer_mode="mistral",
174 + config_format="mistral",
175 + load_format="mistral",
176 + limit_mm_per_prompt={"audio": len(audio_assets)},
177 + ) as vllm_model:
178 + vllm_outputs = vllm_model.generate_greedy_logprobs(
179 + [vllm_prompt],
180 + max_tokens,
181 + num_logprobs,
182 + audios=[audio_data],
183 + )
184 + assert vllm_outputs[0][1], "vLLM inference produced empty output"
185 +
186 + with hf_runner(
187 + MODEL_NAME,
188 + dtype="half",
189 + auto_cls=VoxtralForConditionalGeneration,
190 + ) as hf_model:
191 + hf_model = model_utils.voxtral_patch_hf_runner(hf_model)
192 + hf_outputs = hf_model.generate_greedy_logprobs_limit(
193 + [question],
194 + max_tokens,
195 + num_logprobs,
196 + audios=[audio_data],
197 + )
198 + assert hf_outputs[0][1], "HF Transformers produced empty output"
199 +
200 + print(
201 + f"HF Reference Comparison\n"
202 + f" vLLM: {vllm_outputs[0][1]!r}\n"
203 + f" HF: {hf_outputs[0][1]!r}"
204 + )
205 + check_logprobs_close(
206 + outputs_0_lst=vllm_outputs,
207 + outputs_1_lst=hf_outputs,
208 + name_0="vllm",
209 + name_1="hf",
210 + )
tests/models/multimodal/generation/test_voxtral_realtime.py
+6 -2
@@ -10,6 +10,7 @@ from mistral_common.protocol.transcription.request import (
10 TranscriptionRequest,
11 )
12 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
13 +from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy
14
15 from vllm import LLM, EngineArgs, SamplingParams
16 from vllm.assets.audio import AudioAsset
@@ -26,7 +27,7 @@ ENGINE_CONFIG = dict(
27 load_format="mistral",
28 tokenizer_mode="mistral",
29 enforce_eager=True,
29 - gpu_memory_utilization=0.4,
30 + gpu_memory_utilization=0.9,
31 )
32
33
@@ -148,6 +149,9 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine)
149
150 output_tokens_list.append(output_tokens)
151
151 - texts = [tokenizer.decode(output_tokens) for output_tokens in output_tokens_list]
152 + texts = [
153 + tokenizer.decode(output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE)
154 + for output_tokens in output_tokens_list
155 + ]
156 texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
157 assert texts == EXPECTED_TEXT
tests/models/multimodal/generation/vlm_utils/model_utils.py
+88
@@ -1215,3 +1215,91 @@ def tarsier_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
1215 hf_processor.patch_size = vision_encoder_info.get_patch_size()
1216
1217 return hf_model
1218 +
1219 +
1220 +def voxtral_patch_hf_runner(hf_model: "HfRunner") -> "HfRunner":
1221 + """Patch HfRunner for Voxtral's conversation-based processor.
1222 +
1223 + Two issues in HfRunner require patching:
1224 +
1225 + 1. VoxtralProcessor requires ``apply_chat_template()`` with conversation
1226 + dicts (accepting ``url``, ``path``, or ``base64`` audio) rather than
1227 + the standard ``processor(text=, audio=, sampling_rate=)`` interface.
1228 + 2. HfRunner.get_inputs cannot handle multi-audio per prompt because it
1229 + mis-unpacks ``[(arr1, sr1), (arr2, sr2)]`` via a ``len == 2`` check.
1230 +
1231 + We override ``get_inputs`` to build conversation dicts and call
1232 + ``apply_chat_template`` directly, bypassing both issues. We also wrap
1233 + ``model.generate`` to strip prompt tokens before decoding, since
1234 + HfRunner.generate calls batch_decode on the full sequence (prompt +
1235 + generated).
1236 + """
1237 +
1238 + import base64
1239 + import io
1240 +
1241 + import soundfile as sf
1242 +
1243 + processor = hf_model.processor
1244 +
1245 + def _audio_to_base64(audio_array, sample_rate: int) -> str:
1246 + """Encode a numpy audio array as a base64 WAV string."""
1247 + buf = io.BytesIO()
1248 + sf.write(buf, audio_array, int(sample_rate), format="WAV")
1249 + return base64.b64encode(buf.getvalue()).decode("ascii")
1250 +
1251 + def patched_get_inputs(prompts, images=None, videos=None, audios=None, **kwargs):
1252 + all_inputs = []
1253 + for i, prompt in enumerate(prompts):
1254 + content: list[dict] = []
1255 +
1256 + if audios is not None and audios[i] is not None:
1257 + items = audios[i]
1258 + if not isinstance(items, list):
1259 + items = [items]
1260 + for item in items:
1261 + if isinstance(item, (list, tuple)) and len(item) == 2:
1262 + arr, sr = item
1263 + else:
1264 + arr, sr = item, 16_000
1265 + content.append(
1266 + {
1267 + "type": "audio",
1268 + "base64": _audio_to_base64(arr, sr),
1269 + }
1270 + )
1271 +
1272 + content.append({"type": "text", "text": prompt})
1273 +
1274 + inputs = processor.apply_chat_template(
1275 + [{"role": "user", "content": content}]
1276 + )
1277 + if hasattr(inputs, "to"):
1278 + inputs = inputs.to(dtype=hf_model.dtype)
1279 + all_inputs.append(inputs)
1280 +
1281 + return all_inputs
1282 +
1283 + _orig_generate = hf_model.model.generate
1284 +
1285 + def patched_generate(*args, **kwargs):
1286 + """Strip prompt tokens so only generated tokens are decoded."""
1287 + input_ids = kwargs.get("input_ids")
1288 + if input_ids is None and args:
1289 + input_ids = args[0]
1290 + prompt_len = input_ids.shape[1] if input_ids is not None else 0
1291 +
1292 + output = _orig_generate(*args, **kwargs)
1293 + if prompt_len:
1294 + if isinstance(output, torch.Tensor):
1295 + output = output[:, prompt_len:]
1296 + else:
1297 + # GenerateDecoderOnlyOutput - trim sequences but preserve
1298 + # scores/logits so generate_greedy_logprobs_limit can
1299 + # extract per-token logprobs.
1300 + output.sequences = output.sequences[:, prompt_len:]
1301 + return output
1302 +
1303 + hf_model.get_inputs = patched_get_inputs # type: ignore[method-assign, assignment]
1304 + hf_model.model.generate = patched_generate # type: ignore[method-assign]
1305 + return hf_model
tests/models/multimodal/processing/test_common.py
+34 -14
@@ -184,22 +184,42 @@ def get_text_token_prompts(
184 text_prompt: str | None
185 token_prompt: list[int]
186 if isinstance(tokenizer, MistralTokenizer):
187 - images = parsed_data.get("image", [])
188 - request = ChatCompletionRequest(
189 - messages=[
190 - UserMessage(
191 - content=[
192 - TextChunk(text=""),
193 - *(ImageChunk(image=image) for image in images),
194 - ]
195 - ),
196 - ]
187 + # ChatCompletionRequest only supports ImageChunk natively;
188 + # for other modalities (e.g. audio), fall back to the model's
189 + # own dummy inputs builder which knows the right placeholders.
190 + has_non_image = any(
191 + k != "image" and count > 0 for k, count in mm_counts.items()
192 )
198 - res = tokenizer.mistral.encode_chat_completion(request)
193
200 - # Mistral does not support decode_tokens with skip_special_tokens=False
201 - text_prompt = None
202 - token_prompt = res.tokens
194 + if has_non_image:
195 + inputs = dummy_inputs.get_dummy_processor_inputs(
196 + model_config.max_model_len,
197 + mm_counts,
198 + )
199 + text_prompt = None
200 + token_prompt = (
201 + inputs.prompt
202 + if isinstance(inputs.prompt, list)
203 + else tokenizer.encode(inputs.prompt, add_special_tokens=False)
204 + )
205 + else:
206 + images = parsed_data.get("image", [])
207 + request = ChatCompletionRequest(
208 + messages=[
209 + UserMessage(
210 + content=[
211 + TextChunk(text=""),
212 + *(ImageChunk(image=image) for image in images),
213 + ]
214 + ),
215 + ]
216 + )
217 + res = tokenizer.mistral.encode_chat_completion(request)
218 +
219 + # Mistral does not support decode_tokens with
220 + # skip_special_tokens=False
221 + text_prompt = None
222 + token_prompt = res.tokens
223 else:
224 inputs = dummy_inputs.get_dummy_processor_inputs(
225 model_config.max_model_len,
vllm/model_executor/models/voxtral.py
+28
@@ -291,6 +291,34 @@ class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo])
291 # skip validation here
292 ...
293
294 + def _apply_hf_processor_mm_only(
295 + self,
296 + mm_items: MultiModalDataItems,
297 + hf_processor_mm_kwargs: Mapping[str, object],
298 + tokenization_kwargs: Mapping[str, object],
299 + ) -> BatchFeature:
300 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
301 + processor_data, passthrough_data = self._get_hf_mm_data(mm_items)
302 + audios = processor_data.get("audios", [])
303 + if not isinstance(audios, list):
304 + audios = [audios]
305 +
306 + audio_config = processor._audio_processor.audio_config
307 + audio_tensors: list[torch.Tensor] = []
308 + for audio in audios:
309 + audio = np.asarray(audio, dtype=np.float32).ravel()
310 + if not audio_config.is_streaming:
311 + audio = processor._audio_processor.pad(
312 + audio,
313 + processor.sampling_rate,
314 + audio_config.is_streaming,
315 + )
316 + audio_tensors.append(torch.tensor(audio))
317 +
318 + result = BatchFeature({"audio_arrays": audio_tensors} if audio_tensors else {})
319 + result.update(passthrough_data)
320 + return result
321 +
322 def _get_prompt_updates(
323 self,
324 mm_items: MultiModalDataItems,
vllm/model_executor/models/whisper_causal.py
+44 -4
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import copy
4 import functools
5 +import logging
6 import math
7 from dataclasses import replace
8 from functools import partial
@@ -30,11 +31,20 @@ from vllm.v1.attention.backend import (
31 subclass_attention_backend_with_overrides,
32 )
33 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
34 +
35 +try:
36 + from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
37 +except ImportError:
38 + AiterFlashAttentionBackend = None
39 +from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
40 +from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
41 from vllm.v1.attention.selector import get_attn_backend
42 from vllm.v1.kv_cache_interface import AttentionSpec
43
44 from .utils import make_layers
45
46 +logger = logging.getLogger(__name__)
47 +
48 CausalRMSNorm = partial(RMSNorm, eps=1e-5)
49
50
@@ -122,6 +132,13 @@ def create_whisper_attention_backend_with_block_pooling(
132 num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size,
133 )
134 super().__init__(kv_cache_spec, layer_names, vllm_config, device)
135 + # Override model_config-derived values with the actual
136 + # encoder values from kv_cache_spec
137 + self.num_heads_kv = kv_cache_spec.num_kv_heads
138 + self.headdim = kv_cache_spec.head_size
139 + # num_heads_q for the encoder is the same as num_kv_heads
140 + # (no GQA in whisper encoder)
141 + self.num_heads_q = kv_cache_spec.num_kv_heads
142
143 def build(
144 self,
@@ -192,13 +209,36 @@ def create_whisper_attention_backend_with_block_pooling(
209 output_block_scale,
210 )
211
195 - if not issubclass(underlying_attn_backend, FlashAttentionBackend):
212 + _SUPPORTED_BACKENDS = tuple(
213 + b
214 + for b in (
215 + AiterFlashAttentionBackend,
216 + FlashAttentionBackend,
217 + RocmAttentionBackend,
218 + TritonAttentionBackend,
219 + )
220 + if b is not None
221 + )
222 +
223 + if not issubclass(underlying_attn_backend, _SUPPORTED_BACKENDS):
224 raise NotImplementedError(
225 f"{underlying_attn_backend} is not yet supported."
226 "Contributions to support more backends are much "
227 "appreciated."
228 )
229
230 + if not issubclass(underlying_attn_backend, FlashAttentionBackend):
231 + logger.info(
232 + "Using %s for Whisper causal attention with block pooling. "
233 + "This backend was recently enabled for this model. "
234 + "If you encounter any accuracy or performance issues, "
235 + "please open an issue at "
236 + "https://github.com/vllm-project/vllm/issues "
237 + "with the [ROCm] tag so it can be triaged by the "
238 + "appropriate team.",
239 + underlying_attn_backend.get_name(),
240 + )
241 +
242 attn_backend = subclass_attention_backend_with_overrides(
243 name_prefix=prefix,
244 attention_backend_cls=underlying_attn_backend,
@@ -209,14 +249,14 @@ def create_whisper_attention_backend_with_block_pooling(
249 block_size,
250 num_kv_heads,
251 head_size,
212 - cache_dtype_str: (
213 - 2,
252 + cache_dtype_str: underlying_attn_backend.get_kv_cache_shape(
253 num_blocks,
254 # we stretch each block by `block_pool_size`
255 block_size * block_pool_size,
256 num_kv_heads // block_pool_size,
257 head_size,
219 - ), # TODO: generalize to other backends
258 + cache_dtype_str,
259 + ),
260 "forward_includes_kv_cache_update": True,
261 },
262 )
vllm/reasoning/mistral_reasoning_parser.py
+2 -2
@@ -43,8 +43,8 @@ class MistralReasoningParser(BaseThinkingReasoningParser):
43 "constructor during construction."
44 )
45
46 - self.start_token_id = tokenizer.tokenizer.get_control_token(self.start_token)
47 - self.end_token_id = tokenizer.tokenizer.get_control_token(self.end_token)
46 + self.start_token_id = tokenizer.tokenizer.get_special_token(self.start_token)
47 + self.end_token_id = tokenizer.tokenizer.get_special_token(self.end_token)
48
49 if self.start_token_id is None or self.end_token_id is None:
50 raise RuntimeError(
vllm/tokenizers/mistral.py
+1 -1
@@ -517,7 +517,7 @@ class MistralTokenizer(TokenizerLike):
517 return [self.tokenizer.id_to_piece(token_id) for token_id in ids]
518
519 non_skip_special_tokens_ids = {
520 - self.tokenizer.get_control_token(SpecialTokens.tool_calls),
520 + self.tokenizer.get_special_token(SpecialTokens.tool_calls),
521 }
522 if isinstance(self.instruct, InstructTokenizerV13):
523 if self.instruct.BEGIN_THINK:
vllm/v1/attention/backends/rocm_aiter_fa.py
+7 -2
@@ -425,8 +425,13 @@ class AiterFlashAttentionMetadataBuilder(
425
426 sliding_window_configs: set[tuple[int, int] | None] = set()
427 layers = get_layers_from_vllm_config(self.vllm_config, Attention)
428 - for layer in layers.values():
429 - assert isinstance(layer.impl, AiterFlashAttentionImpl)
428 + for name, layer in layers.items():
429 + if name not in layer_names:
430 + continue
431 + assert isinstance(layer.impl, AiterFlashAttentionImpl), (
432 + "Aiter Flash Attention Metadata Builder can only be used "
433 + "with Aiter Flash Attention Impl."
434 + )
435 sliding_window_configs.add(layer.impl.sliding_window)
436
437 while len(sliding_window_configs) > 0: