[Bugfix][V1][TurboQuant] Reserve workspace before CUDA graph capture (#44053)
Signed-off-by: Guipeng Zhang <zhangguipeng23z@ict.ac.cn> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Guipeng Zhang committed
Jun 23, 2026 at 06:47 UTC
183b5f27eafa734117605932e969de47fd1324a6
2 files changed
+163
tests/quantization/test_turboquant.py
+124
@@ -6,6 +6,7 @@ Run: .venv/bin/python -m pytest tests/quantization/test_turboquant.py -v
6
"""
7
8
import math
9
+from types import SimpleNamespace
10
11
import pytest
12
import torch
@@ -276,6 +277,129 @@ class TestHybridAttentionIndices:
277
assert _get_full_attention_layer_indices(mc) == []
278
279
280
+class TestTurboQuantWorkspaceReservation:
281
+ @staticmethod
282
+ def _fake_vllm_config(
283
+ *,
284
+ max_num_seqs: int = 16,
285
+ max_num_batched_tokens: int = 4096,
286
+ enable_chunked_prefill: bool = True,
287
+ max_model_len: int = 8192,
288
+ dtype: torch.dtype = torch.float16,
289
+ max_num_kv_splits: int = 4,
290
+ ):
291
+ return SimpleNamespace(
292
+ scheduler_config=SimpleNamespace(
293
+ max_num_seqs=max_num_seqs,
294
+ max_num_batched_tokens=max_num_batched_tokens,
295
+ enable_chunked_prefill=enable_chunked_prefill,
296
+ ),
297
+ model_config=SimpleNamespace(
298
+ max_model_len=max_model_len,
299
+ dtype=dtype,
300
+ get_num_attention_heads=lambda parallel_config: 8,
301
+ ),
302
+ parallel_config=SimpleNamespace(
303
+ tensor_parallel_size=2,
304
+ decode_context_parallel_size=1,
305
+ ),
306
+ attention_config=SimpleNamespace(
307
+ tq_max_kv_splits_for_cuda_graph=max_num_kv_splits
308
+ ),
309
+ )
310
+
311
+ @staticmethod
312
+ def _fake_kv_cache_spec():
313
+ from vllm.v1.kv_cache_interface import TQFullAttentionSpec
314
+
315
+ return TQFullAttentionSpec(
316
+ block_size=32,
317
+ num_kv_heads=4,
318
+ head_size=128,
319
+ head_size_v=128,
320
+ dtype=torch.uint8,
321
+ tq_slot_size=102,
322
+ )
323
+
324
+ def test_metadata_builder_reserves_decode_and_continuation_prefill_workspace(
325
+ self, monkeypatch
326
+ ):
327
+ from vllm.v1.attention.backends import turboquant_attn
328
+
329
+ calls = []
330
+
331
+ class FakeWorkspaceManager:
332
+ def get_simultaneous(self, *shapes_and_dtypes):
333
+ calls.append(shapes_and_dtypes)
334
+
335
+ monkeypatch.setattr(
336
+ turboquant_attn,
337
+ "current_workspace_manager",
338
+ lambda: FakeWorkspaceManager(),
339
+ )
340
+ monkeypatch.setattr(
341
+ turboquant_attn,
342
+ "is_workspace_manager_initialized",
343
+ lambda: True,
344
+ )
345
+
346
+ turboquant_attn.TurboQuantMetadataBuilder(
347
+ kv_cache_spec=self._fake_kv_cache_spec(),
348
+ layer_names=["layers.0.self_attn.attn"],
349
+ vllm_config=self._fake_vllm_config(),
350
+ device=torch.device("cuda"),
351
+ )
352
+
353
+ assert calls == [
354
+ (
355
+ ((16, 8, 4, 129), torch.float32),
356
+ ((16, 8, 128), torch.float16),
357
+ ((16, 8), torch.float32),
358
+ ),
359
+ (
360
+ ((1, 4, 8192, 128), torch.float16),
361
+ ((1, 4, 8192, 128), torch.float16),
362
+ ),
363
+ ]
364
+
365
+ def test_metadata_builder_skips_continuation_prefill_when_disabled(
366
+ self, monkeypatch
367
+ ):
368
+ from vllm.v1.attention.backends import turboquant_attn
369
+
370
+ calls = []
371
+
372
+ class FakeWorkspaceManager:
373
+ def get_simultaneous(self, *shapes_and_dtypes):
374
+ calls.append(shapes_and_dtypes)
375
+
376
+ monkeypatch.setattr(
377
+ turboquant_attn,
378
+ "current_workspace_manager",
379
+ lambda: FakeWorkspaceManager(),
380
+ )
381
+ monkeypatch.setattr(
382
+ turboquant_attn,
383
+ "is_workspace_manager_initialized",
384
+ lambda: True,
385
+ )
386
+
387
+ turboquant_attn.TurboQuantMetadataBuilder(
388
+ kv_cache_spec=self._fake_kv_cache_spec(),
389
+ layer_names=["layers.0.self_attn.attn"],
390
+ vllm_config=self._fake_vllm_config(enable_chunked_prefill=False),
391
+ device=torch.device("cuda"),
392
+ )
393
+
394
+ assert calls == [
395
+ (
396
+ ((16, 8, 4, 129), torch.float32),
397
+ ((16, 8, 128), torch.float16),
398
+ ((16, 8), torch.float32),
399
+ )
400
+ ]
401
+
402
+
403
# ============================================================================
404
# Centroids tests (CPU-only)
405
# ============================================================================
vllm/v1/attention/backends/turboquant_attn.py
+39
@@ -30,6 +30,7 @@ from vllm.model_executor.layers.quantization.turboquant.centroids import (
30
get_centroids,
31
)
32
from vllm.triton_utils import triton
33
+from vllm.utils.math_utils import round_up
34
from vllm.v1.attention.backend import (
35
AttentionBackend,
36
AttentionCGSupport,
@@ -201,6 +202,44 @@ class TurboQuantMetadataBuilder(AttentionMetadataBuilder[TurboQuantMetadata]):
202
def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
203
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
204
self._init_reorder_batch_threshold(1, supports_spec_as_decode=False)
205
+ self._reserve_workspace()
206
+
207
+ def _reserve_workspace(self) -> None:
208
+ if not is_workspace_manager_initialized():
209
+ return
210
+
211
+ scheduler_config = self.vllm_config.scheduler_config
212
+ model_config = self.vllm_config.model_config
213
+ parallel_config = self.vllm_config.parallel_config
214
+
215
+ max_num_reqs = scheduler_config.max_num_seqs
216
+ num_heads = model_config.get_num_attention_heads(parallel_config)
217
+ num_kv_heads = self.kv_cache_spec.num_kv_heads
218
+ head_size = self.kv_cache_spec.head_size
219
+ max_num_splits = (
220
+ self.vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph
221
+ )
222
+
223
+ current_workspace_manager().get_simultaneous(
224
+ ((max_num_reqs, num_heads, max_num_splits, head_size + 1), torch.float32),
225
+ ((max_num_reqs, num_heads, head_size), model_config.dtype),
226
+ ((max_num_reqs, num_heads), torch.float32),
227
+ )
228
+
229
+ reserve_continuation_prefill = (
230
+ scheduler_config.enable_chunked_prefill
231
+ and scheduler_config.max_num_batched_tokens > _CONTINUATION_DECODE_THRESHOLD
232
+ )
233
+ if not reserve_continuation_prefill:
234
+ return
235
+
236
+ max_cached_len = max(0, model_config.max_model_len - 1)
237
+ alloc_len = round_up(max_cached_len, self.kv_cache_spec.block_size)
238
+ cache_buf_shape = (1, num_kv_heads, alloc_len, head_size)
239
+ current_workspace_manager().get_simultaneous(
240
+ (cache_buf_shape, torch.float16),
241
+ (cache_buf_shape, torch.float16),
242
+ )
243
244
def build_for_cudagraph_capture(
245
self, common_attn_metadata: CommonAttentionMetadata