@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2cc008e7b4

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 4, 15:00 UTC

[Attention][TurboQuant] Share dequant buffers, eliminate float16_copy (#40941)

Signed-off-by: Bhoomit Vasani <bhoomit.2010@gmail.com> Signed-off-by: Vasani Bhoomit <bhoomit.2010@gmail.com>

Bhoomit committed Apr 26, 2026 at 22:48 UTC 2cc008e7b491bb77f0caf4d27ad55a83f196114c

4 files changed +98 -92

vllm/model_executor/layers/attention/attention.py

-48

@@ -386,10 +386,6 @@ class Attention(nn.Module, AttentionLayerBase):

386 # Initialize KV cache quantization attributes

387 _init_kv_cache_quant(self, quant_config, prefix)

388

389 - # Initialize TurboQuant buffers (Pi, S, centroids) if tq cache dtype

390 - if kv_cache_dtype.startswith("turboquant_"):

391 - self._init_turboquant_buffers(kv_cache_dtype, head_size, prefix)

392 -

389 # for attn backends supporting query quantization

390 self.query_quant = None

391 if (

@@ -410,50 +406,6 @@ class Attention(nn.Module, AttentionLayerBase):

406 else GroupShape.PER_TENSOR,

407 )

408

413 - def _init_turboquant_buffers(

414 - self, cache_dtype: str, head_size: int, prefix: str

415 - ) -> None:

416 - """Initialize TurboQuant centroids for Lloyd-Max quantization."""

417 - from vllm.model_executor.layers.quantization.turboquant.centroids import (

418 - get_centroids,

419 - )

420 - from vllm.model_executor.layers.quantization.turboquant.config import (

421 - TurboQuantConfig,

422 - )

423 -

424 - tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size)

425 -

426 - self.register_buffer(

427 - "_tq_centroids",

428 - get_centroids(head_size, tq_config.centroid_bits),

429 - )

430 - self._tq_config = tq_config

431 -

432 - # Pre-allocate decode intermediate buffers so model.to(device) moves

433 - # them to GPU *before* the memory profiler runs. Without this the

434 - # profiler gives all free memory to KV cache blocks and the first

435 - # decode OOMs when these buffers are lazily allocated.

436 - _vllm_cfg = get_current_vllm_config()

437 - B = _vllm_cfg.scheduler_config.max_num_seqs

438 - Hq = self.num_heads

439 - S = _vllm_cfg.attention_config.tq_max_kv_splits_for_cuda_graph

440 - D = head_size

441 - self.register_buffer(

442 - "_tq_mid_o_buf",

443 - torch.empty(B, Hq, S, D + 1, dtype=torch.float32),

444 - persistent=False,

445 - )

446 - self.register_buffer(

447 - "_tq_output_buf",

448 - torch.empty(B, Hq, D, dtype=torch.float32),

449 - persistent=False,

450 - )

451 - self.register_buffer(

452 - "_tq_lse_buf",

453 - torch.empty(B, Hq, dtype=torch.float32),

454 - persistent=False,

455 - )

456 -

409 def forward(

410 self,

411 query: torch.Tensor,

vllm/v1/attention/backends/turboquant_attn.py

+83 -40

@@ -26,6 +26,9 @@ import torch.nn.functional as F

27 from vllm.config import get_current_vllm_config

28 from vllm.config.cache import CacheDType

29 +from vllm.model_executor.layers.quantization.turboquant.centroids import (

30 + get_centroids,

31 +)

32 from vllm.triton_utils import triton

33 from vllm.v1.attention.backend import (

34 AttentionBackend,

@@ -49,6 +52,10 @@ from vllm.v1.attention.ops.triton_turboquant_decode import (

52 triton_turboquant_decode_attention,

53 )

54 from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store

55 +from vllm.v1.worker.workspace import (

56 + current_workspace_manager,

57 + is_workspace_manager_initialized,

58 +)

60 _HAS_FLASH_ATTN = is_flash_attn_varlen_func_available()

61 if _HAS_FLASH_ATTN:

@@ -335,9 +342,15 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

342 H = _build_hadamard(D, str(device))

343 layer._tq_PiT = H

344 layer._tq_Pi = H

345 + # fp16 copy for rotation in continuation prefill path

346 + layer._tq_Pi_half = H.to(torch.float16)

347

339 - c = layer._tq_centroids.to(device=device, dtype=torch.float32)

340 - c_sorted, _ = c.sort()

348 + # Centroids for Lloyd-Max quantization.

349 + layer._tq_centroids = get_centroids(D, self.tq_config.centroid_bits).to(

350 + device=device, dtype=torch.float32

351 + )

352 +

353 + c_sorted, _ = layer._tq_centroids.sort()

354 layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2

355 layer._tq_cached = True

356

@@ -572,7 +585,17 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

585

586 # Pre-allocate cu_seqlens for single-request flash_attn calls

587 # to avoid per-request host→device tensor creation.

575 - _cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32)

588 + if not hasattr(self, "_cu_2"):

589 + self._cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32)

590 + # Cache arange on self (avoid per-call kernel launch).

591 + _max_seq = attn_metadata.max_seq_len

592 + _ac: torch.Tensor | None = getattr(self, "_arange_cache", None)

593 + if _ac is None or _ac.shape[0] <= _max_seq:

594 + _ac = torch.arange(

595 + 0, _max_seq + 1, device=query.device, dtype=attn_metadata.seq_lens.dtype

596 + )

597 + self._arange_cache = _ac

598 + _arange_cache: torch.Tensor = _ac

599

600 for i in range(num_reqs):

601 q_start = qsl[i]

@@ -589,8 +612,8 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

612 if q_len == seq_len:

613 # First-chunk prefill: all K/V are in the current batch.

614 if _HAS_FLASH_ATTN:

592 - _cu_2[1] = q_len

593 - cu = _cu_2

615 + self._cu_2[1] = q_len

616 + cu = self._cu_2

617 out = self._flash_attn_varlen(

618 q=q_seq,

619 k=k_seq,

@@ -622,12 +645,8 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

645 if q_len <= _CONTINUATION_DECODE_THRESHOLD:

646 # Fast path: treat each query as a decode request

647 # with incremental seq_lens for causal masking.

625 - synth_seq_lens = torch.arange(

626 - cached_len + 1,

627 - seq_len + 1,

628 - device=query.device,

629 - dtype=attn_metadata.seq_lens.dtype,

630 - )

648 + # Slice from pre-built arange (no kernel launch)

649 + synth_seq_lens = _arange_cache[cached_len + 1 : seq_len + 1]

650 synth_bt = attn_metadata.block_table[i : i + 1].expand(q_len, -1)

651 out = triton_turboquant_decode_attention(

652 query=q_seq,

@@ -695,16 +714,17 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

714 # Reuse cached buffers to avoid per-call allocation (~16MB at 8K).

715 alloc_len = math.ceil(cached_len / block_size) * block_size

716 buf_shape = (1, Hk, alloc_len, D)

698 - k_buf = getattr(layer, "_tq_k_dequant_buf", None)

699 - if k_buf is None or k_buf.shape[2] < alloc_len:

700 - k_buf = torch.empty(buf_shape, dtype=torch.float16, device=device)

701 - v_buf = torch.empty(buf_shape, dtype=torch.float16, device=device)

702 - layer._tq_k_dequant_buf = k_buf

703 - layer._tq_v_dequant_buf = v_buf

704 - else:

705 - v_buf = layer._tq_v_dequant_buf

706 - k_cached = k_buf[:, :, :alloc_len, :].zero_()

707 - v_cached = v_buf[:, :, :alloc_len, :].zero_()

717 + # Use WorkspaceManager for dequant buffers.

718 + # Shared across all layers — saves 60× memory at long context.

719 + # Required for CUDA Graph capture (per-layer growth incompatible with CG).

720 + k_buf, v_buf = current_workspace_manager().get_simultaneous(

721 + (buf_shape, torch.float16),

722 + (buf_shape, torch.float16),

723 + )

724 + # Skip .zero_() — kernel writes all positions up to cached_len,

725 + # and we only read [:cached_len] afterwards.

726 + k_cached = k_buf[:, :, :alloc_len, :]

727 + v_cached = v_buf[:, :, :alloc_len, :]

728

729 grid = (alloc_len, 1 * Hk)

730 _tq_full_dequant_kv[grid](

@@ -740,29 +760,41 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

760

761 # Inverse-rotate MSE keys back to original space

762 if not self.tq_config.key_fp8:

743 - k_flat = k_cached[0, :, :cached_len, :].reshape(-1, D).float()

744 - k_flat = k_flat @ Pi

745 - k_cached_trim = (

746 - k_flat.to(torch.float16).reshape(Hk, cached_len, D).transpose(0, 1)

747 - ) # (cached_len, Hk, D)

763 + # fp16 matmul for rotation (2× less bandwidth, uses fp16 tensor cores)

764 + Pi_half = layer._tq_Pi_half

765 + k_flat = k_cached[0, :, :cached_len, :].reshape(-1, D)

766 + k_flat = k_flat @ Pi_half

767 + k_cached_trim = k_flat.reshape(Hk, cached_len, D).transpose(

768 + 0, 1

769 + ) # (cached_len, Hk, D) — already fp16

770 else:

749 - k_cached_trim = (

750 - k_cached[0, :, :cached_len, :].transpose(0, 1).contiguous()

771 + k_cached_trim = k_cached[0, :, :cached_len, :].transpose(

772 + 0, 1

773 ) # (cached_len, Hk, D)

774

753 - v_cached_trim = (

754 - v_cached[0, :, :cached_len, :].transpose(0, 1).contiguous()

755 - ) # (cached_len, Hk, D)

775 + # Skip .contiguous() — the copy into k_full/v_full handles layout

776 + v_cached_trim = v_cached[0, :, :cached_len, :].transpose(0, 1)

777

778 # Concatenate cached + current chunk K/V (match query dtype)

779 + # Pre-allocate full K/V buffer, copy into slices (no cat alloc)

780 qdtype = query.dtype

759 - k_full = torch.cat([k_cached_trim.to(qdtype), key_chunk], dim=0)

760 - v_full = torch.cat([v_cached_trim.to(qdtype), val_chunk], dim=0)

781 + k_full = torch.empty(seq_len, Hk, D, dtype=qdtype, device=device)

782 + v_full = torch.empty(seq_len, Hk, D, dtype=qdtype, device=device)

783 + k_full[:cached_len] = k_cached_trim.to(qdtype)

784 + k_full[cached_len:] = key_chunk

785 + v_full[:cached_len] = v_cached_trim.to(qdtype)

786 + v_full[cached_len:] = val_chunk

787

788 # Attention: q_len queries attending to seq_len K/V with causal mask

789 if _HAS_FLASH_ATTN:

764 - cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32)

765 - cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32)

790 + # Reuse pre-allocated cu_seqlens (avoid host→device transfer)

791 + if not hasattr(self, "_cu_2_q"):

792 + self._cu_2_q = torch.zeros(2, device=device, dtype=torch.int32)

793 + self._cu_2_k = torch.zeros(2, device=device, dtype=torch.int32)

794 + self._cu_2_q[1] = q_len

795 + self._cu_2_k[1] = seq_len

796 + cu_seqlens_q = self._cu_2_q

797 + cu_seqlens_k = self._cu_2_k

798 return self._flash_attn_varlen(

799 q=query,

800 k=k_full,

@@ -805,12 +837,23 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

837 PiT: torch.Tensor | None = None,

838 layer: torch.nn.Module | None = None,

839 ) -> torch.Tensor:

808 - # Grab cached decode buffers from the layer (lazily allocated).

840 + # Acquire shared decode scratch buffers from WorkspaceManager.

841 + # Layers execute sequentially so one set of buffers is sufficient.

842 + # Falls back to kernel-internal allocation if workspace unavailable.

843 + B = query.shape[0]

844 + D = self.head_size

845 + S = self.max_num_kv_splits

846 + Hq = self.num_heads

847 mid_o_buf = output_buf = lse_buf = None

810 - if layer is not None:

811 - mid_o_buf = getattr(layer, "_tq_mid_o_buf", None)

812 - output_buf = getattr(layer, "_tq_output_buf", None)

813 - lse_buf = getattr(layer, "_tq_lse_buf", None)

848 + if is_workspace_manager_initialized():

849 + # output_buf in query dtype — matches the in-kernel fp16 cast in stage2.

850 + mid_o_buf, output_buf, lse_buf = (

851 + current_workspace_manager().get_simultaneous(

852 + ((B, Hq, S, D + 1), torch.float32),

853 + ((B, Hq, D), query.dtype),

854 + ((B, Hq), torch.float32),

855 + )

856 + )

857

858 result = triton_turboquant_decode_attention(

859 query=query,

vllm/v1/attention/ops/triton_decode_attention.py

+5 -1

@@ -551,6 +551,7 @@ def _fwd_kernel_stage2(

551 NUM_KV_SPLITS: tl.constexpr,

552 BLOCK_DV: tl.constexpr,

553 Lv: tl.constexpr,

554 + OUTPUT_FP16: tl.constexpr = 0,

555 ):

556 cur_batch = tl.program_id(0)

557 cur_head = tl.program_id(1)

@@ -587,9 +588,12 @@ def _fwd_kernel_stage2(

588 e_sum = e_sum * old_scale + exp_logic

589 e_max = n_e_max

590

591 + result = acc / e_sum

592 + if OUTPUT_FP16:

593 + result = result.to(tl.float16)

594 tl.store(

595 o + cur_batch * stride_obs + cur_head * stride_oh + offs_d,

592 - acc / e_sum,

596 + result,

597 mask=mask_d,

598 )

599 lse_val = e_max + tl.log(e_sum)

vllm/v1/attention/ops/triton_turboquant_decode.py

+10 -3

@@ -588,10 +588,16 @@ def triton_turboquant_decode_attention(

588 )

589

590 # Stage 2: Reduce across KV splits

591 - if output_buf is not None and output_buf.shape[0] >= B:

591 + # Output in query dtype — eliminates float16_copy kernel after stage2

592 + out_dtype = query.dtype

593 + if (

594 + output_buf is not None

595 + and output_buf.shape[0] >= B

596 + and output_buf.dtype == out_dtype

597 + ):

598 output = output_buf[:B, :Hq, :D]

599 else:

594 - output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)

600 + output = torch.empty(B, Hq, D, dtype=out_dtype, device=device)

601 if buf_holder is not None:

602 buf_holder._tq_output_buf = output

603 if lse_buf is not None and lse_buf.shape[0] >= B:

@@ -616,8 +622,9 @@ def triton_turboquant_decode_attention(

622 NUM_KV_SPLITS=NUM_KV_SPLITS,

623 BLOCK_DV=cfg["BLOCK_D"],

624 Lv=D,

625 + OUTPUT_FP16=1 if out_dtype == torch.float16 else 0,

626 num_warps=4,

627 num_stages=2,

628 )

629

623 - return output.to(query.dtype)

630 + return output # already in query dtype