"""Tiled continuation prefill for the TurboQuant backend. During chunked prefill vLLM's TurboQuant backend takes one of two branches: q_len <= 128 the decode kernel walks the packed cache on CUDA cores; q_len > 128 ``_continuation_prefill`` dequantises the cached context to fp16 and runs FlashAttention on tensor cores. Measured on an RTX 5090 with this model's shapes, the second branch is 58x faster per query-token-layer: a literal 1,010,000-token prefill costs about 16 hours through the decode kernel and about 17 minutes through FlashAttention. The fast branch is nevertheless unusable at a million tokens, because it materialises two full-length fp16 buffers - the dequantisation workspace and a concatenated K/V - which come to roughly 7.7 GiB and do not fit beside the weights and the KV cache. Attention decomposes exactly over disjoint key ranges when partial results are merged by log-sum-exp, so the cached context can be walked one tile at a time. The arithmetic stays on tensor cores while the buffers shrink to a single tile. Measured on an RTX 5090 with the real dequantisation kernel at 1,009,664 cached tokens and a 1024-token chunk: 0.170-0.174 ns per query-token-layer whatever the tile, i.e. under 2 % spread, while peak memory moves from 4.64 GiB at a 128K tile to 1.77 GiB at 32K. Speed does not pay for the smaller tile, so 32K is the default. Against the one-shot form the result matches to 1.5e-5 absolute - fp16 epsilon - and a sliced block table dequantises bit-identically to the corresponding slice of a full dequantisation. Projected from those measurements, a literal 1,010,000-token prefill takes about 23 minutes instead of about 16 hours. Nothing about the cache format, the codec, or the model changes; this only alters the order in which cached keys and values are consumed. """ from __future__ import annotations import math from typing import Any import torch REPAIR_ID = "ZENIT_TILED_CONTINUATION_PREFILL_V1" DEFAULT_TILE_TOKENS = 32_768 def _round_to_blocks(tokens: int, block_size: int) -> int: return math.ceil(tokens / block_size) * block_size # One tile is in flight at a time across the whole model, so these buffers are # shared by every attention layer. Keying them per layer would hold sixteen # copies of the same scratch space and defeat the point. _TILE_BUFFERS: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} _ROTATION_BUFFERS: dict[tuple, torch.Tensor] = {} def _tile_buffers( tile: int, heads: int, head_dim: int, dtype: torch.dtype, device: Any ) -> tuple[torch.Tensor, torch.Tensor]: """Two persistent (tile, heads, head_dim) buffers in the model's dtype.""" key = (tile, heads, head_dim, dtype, str(device)) cached = _TILE_BUFFERS.get(key) if cached is None: shape = (tile, heads, head_dim) cached = ( torch.empty(shape, dtype=dtype, device=device), torch.empty(shape, dtype=dtype, device=device), ) _TILE_BUFFERS[key] = cached return cached def _rotation_buffer(rows: int, head_dim: int, device: Any) -> torch.Tensor: """Persistent destination for the inverse key rotation, in fp16.""" key = (rows, head_dim, str(device)) cached = _ROTATION_BUFFERS.get(key) if cached is None: cached = torch.empty((rows, head_dim), dtype=torch.float16, device=device) _ROTATION_BUFFERS[key] = cached return cached def _cu_seqlens(impl: Any, device: Any) -> tuple[torch.Tensor, torch.Tensor]: """The two cumulative-length tensors flash_attn needs, kept per layer. Allocating them per call would be wasteful; sharing them between layers would be wrong, since two layers could be in flight with different chunk lengths. One pair per attention implementation is both correct and free. """ cached = getattr(impl, "_zenit_cu_seqlens", None) if cached is None or cached[0].device != device: cached = ( torch.zeros(2, device=device, dtype=torch.int32), torch.zeros(2, device=device, dtype=torch.int32), ) impl._zenit_cu_seqlens = cached return cached def install_tiled_continuation_prefill(tile_tokens: int = DEFAULT_TILE_TOKENS) -> bool: """Replace the one-shot continuation prefill with a tiled one, once.""" from vllm.v1.attention.backends.turboquant_attn import ( TurboQuantAttentionImpl, TurboQuantMetadataBuilder, ) if getattr(TurboQuantAttentionImpl, "_zenit_tiled_prefill", False): return False original = TurboQuantAttentionImpl._continuation_prefill def tiled_continuation_prefill( self, layer: Any, query: torch.Tensor, # (q_len, Hq, D) key_chunk: torch.Tensor, # (q_len, Hk, D) val_chunk: torch.Tensor, # (q_len, Hk, D) kv_cache: torch.Tensor, block_table: torch.Tensor, # (1, max_num_blocks) cached_len: int, seq_len: int, Pi: torch.Tensor, centroids: torch.Tensor, ) -> torch.Tensor: from vllm.v1.attention.backends.turboquant_attn import ( _tq_full_dequant_kv, _use_fp8_e4b15, ) from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.worker.workspace import current_workspace_manager try: from vllm.vllm_flash_attn import flash_attn_varlen_func except ImportError: # pragma: no cover - environment without FA return original( self, layer, query, key_chunk, val_chunk, kv_cache, block_table, cached_len, seq_len, Pi, centroids, ) q_len, Hq, D = query.shape Hk = key_chunk.shape[1] device = query.device block_size = kv_cache.shape[1] tile = max(block_size, _round_to_blocks(tile_tokens, block_size)) # A context that already fits in one tile gains nothing from tiling. if cached_len <= tile: return original( self, layer, query, key_chunk, val_chunk, kv_cache, block_table, cached_len, seq_len, Pi, centroids, ) alloc = tile k_buf, v_buf = current_workspace_manager().get_simultaneous( ((1, Hk, alloc, D), torch.float16), ((1, Hk, alloc, D), torch.float16), ) # The dequant workspace is fp16; the model runs in bfloat16. Converting # with ``.to()`` inside the loop would allocate a fresh tile-sized tensor # on every iteration - about 33 MiB each, four times per tile. On a card # that is 98% full, that churn makes the caching allocator fall back to # cudaFree/cudaMalloc, which synchronises the device and stalls the whole # pipeline. Persistent buffers, filled with copy_, remove the churn. qdtype = query.dtype rotate = not self.tq_config.key_fp8 if qdtype != torch.float16: k_tile_out, v_tile_out = _tile_buffers(alloc, Hk, D, qdtype, device) else: k_tile_out = v_tile_out = None if rotate: rot_buf = _rotation_buffer(alloc * Hk, D, device) # Cumulative sequence lengths for flash_attn. Two rules, both learned # the hard way and both stated in vLLM's own code: keep the tensors # across calls, and write to them with a *slice* assignment. A plain # ``cu_k[1] = span`` is an indexed assignment, which synchronises the # host with the device; ``cu_k[1:2] = span`` lowers to fill_ and does # not. The tile loop performs this write once per tile - about fifteen # times per call at a million tokens against four at 262K - so a sync # here drains the pipeline on every tile and gets worse with context. cu_q, cu_k = _cu_seqlens(self, device) cu_q[1:2] = q_len # Ping-pong buffers: merging reads one pair and writes the other, so no # allocation happens inside the tile loop. prefix_out: torch.Tensor | None = None prefix_lse: torch.Tensor | None = None pong_out = torch.empty_like(query) pong_lse = torch.empty(Hq, q_len, device=device, dtype=torch.float32) for start in range(0, cached_len, tile): stop = min(start + tile, cached_len) span = stop - start span_alloc = _round_to_blocks(span, block_size) first_block = start // block_size last_block = math.ceil(stop / block_size) tile_table = block_table[:, first_block:last_block].contiguous() k_tile = k_buf[:, :, :span_alloc, :] v_tile = v_buf[:, :, :span_alloc, :] _tq_full_dequant_kv[(span_alloc, Hk)]( kv_cache, tile_table, centroids, k_tile, v_tile, k_tile.stride(0), k_tile.stride(1), k_tile.stride(2), v_tile.stride(0), v_tile.stride(1), v_tile.stride(2), kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), tile_table.stride(0), HEAD_DIM=D, BLOCK_SIZE=block_size, NUM_KV_HEADS=Hk, MSE_BYTES=self._mse_bytes, KPS=self.tq_config.key_packed_size, VQB=self.tq_config.effective_value_quant_bits, VAL_DATA_BYTES=self._val_data_bytes, MSE_BITS=self.tq_config.key_mse_bits, KEY_FP8=1 if self.tq_config.key_fp8 else 0, BLOCK_D=1 << (D - 1).bit_length(), NORM_CORRECTION=1 if self.tq_config.norm_correction else 0, FP8_E4B15=_use_fp8_e4b15(device.index or 0), num_warps=4, ) if rotate: rotated = rot_buf[: span * Hk] torch.mm( k_tile[0, :, :span, :].reshape(-1, D), layer._tq_Pi_half, out=rotated, ) k_source = rotated.reshape(Hk, span, D).transpose(0, 1) else: k_source = k_tile[0, :, :span, :].transpose(0, 1) v_source = v_tile[0, :, :span, :].transpose(0, 1) if k_tile_out is None: k_span, v_span = k_source, v_source else: k_span = k_tile_out[:span] v_span = v_tile_out[:span] k_span.copy_(k_source) v_span.copy_(v_source) cu_k[1:2] = span # Every cached token precedes every query token, so no mask here. out_i, lse_i = flash_attn_varlen_func( q=query, k=k_span, v=v_span, cu_seqlens_q=cu_q, cu_seqlens_k=cu_k, max_seqlen_q=q_len, max_seqlen_k=span, softmax_scale=self.scale, causal=False, return_softmax_lse=True, ) if prefix_out is None: prefix_out, prefix_lse = out_i, lse_i else: merge_attn_states(pong_out, prefix_out, prefix_lse, out_i, lse_i, output_lse=pong_lse) prefix_out, pong_out = pong_out, prefix_out prefix_lse, pong_lse = pong_lse, prefix_lse # The chunk's own tokens are causal against each other. cu_k[1:2] = q_len suffix_out, suffix_lse = flash_attn_varlen_func( q=query, k=key_chunk, v=val_chunk, cu_seqlens_q=cu_q, cu_seqlens_k=cu_k, max_seqlen_q=q_len, max_seqlen_k=q_len, softmax_scale=self.scale, causal=True, return_softmax_lse=True, ) final = torch.empty_like(query) merge_attn_states(final, prefix_out, prefix_lse, suffix_out, suffix_lse) return final TurboQuantAttentionImpl._continuation_prefill = tiled_continuation_prefill TurboQuantAttentionImpl._zenit_tiled_prefill = True TurboQuantAttentionImpl._zenit_tile_tokens = tile_tokens # The builder reserves a continuation workspace sized to the whole context. # It derives that size from ``model_config.max_model_len`` and uses the value # for nothing else, so presenting a tile-sized length for the duration of the # call reserves a tile-sized buffer without duplicating any of its logic. original_reserve = TurboQuantMetadataBuilder._reserve_workspace def reserve_one_tile(self) -> None: model_config = self.vllm_config.model_config real_len = model_config.max_model_len if real_len is None or real_len <= tile_tokens + 1: return original_reserve(self) try: model_config.max_model_len = tile_tokens + 1 return original_reserve(self) finally: model_config.max_model_len = real_len TurboQuantMetadataBuilder._reserve_workspace = reserve_one_tile TurboQuantMetadataBuilder._zenit_tiled_workspace = True return True __all__ = ["DEFAULT_TILE_TOKENS", "REPAIR_ID", "install_tiled_continuation_prefill"]