diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/indexer_qsa.py index 4345763fd4..7770a38fa4 100644 --- a/vllm/models/qwen4_exp/amd/indexer_qsa.py +++ b/vllm/models/qwen4_exp/amd/indexer_qsa.py @@ -119,6 +119,15 @@ class QSAIndexer(nn.Module): eps=float(getattr(config, "rms_norm_eps", 1e-6)), ) + # QSAKeyStateCache rows now carry + # a tag word the NVIDIA pre-indexer writes and validates; this path + # neither writes nor checks it, so it would run the widened layout + # with no ring validation at all. Refuse it rather than share it. + raise NotImplementedError( + "Qwen4Exp QSA on the AMD path is not supported on this fork: the " + "ring-row tag contract is implemented for the NVIDIA " + "pre-indexer only (vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py)." + ) cache_config = vllm_config.cache_config cache_prefix = f"{prefix}." if prefix else "" self.raw_key_cache = QSAKeyStateCache( diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index 3b34917361..eb6be6562d 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -44,6 +44,42 @@ from vllm.v1.kv_cache_interface import ( MLAAttentionSpec, ) +# Ring-row tag contract. Every raw-key ring row is the key (bf16), padded to a +# whole int64, followed by a tail of four int64 words: +# (pos_t, pos_h, pos_w, tag), tag = logical_position + 1. +# The tail exists on every ring, with or without cached RoPE positions: the tag +# is what lets the reader prove the row it is about to pool was committed by +# this request for exactly that logical position. 0 is never a valid tag (so a +# cleared row can never validate) and the manager clears the tag word of every +# newly claimed ring block before its new owner's first step. +QSA_RING_TAIL_WORDS = 4 +QSA_RING_TAG_WORD = 3 +QSA_RING_BF16_PER_INT64 = 4 +QSA_RING_TAIL_ELEMENTS = QSA_RING_TAIL_WORDS * QSA_RING_BF16_PER_INT64 + + +def qsa_ring_row_tags(logical_positions: torch.Tensor) -> torch.Tensor: + """The tag the writer stores for each committed row (int64).""" + return logical_positions.to(torch.int64) + 1 + + +def qsa_ring_group_rows_valid( + found_tags: torch.Tensor, + source_positions: torch.Tensor, + historical: torch.Tensor, +) -> torch.Tensor: + """Reference for the pre-indexer's per-group check (torch, any device). + + ``found_tags``/``source_positions``/``historical`` are ``[groups, ratio]``: + the tag word read from each source row's ring slot, the row's logical + position, and whether the row must come from the ring (it precedes the + chunk start). Returns a ``[groups]`` bool: True when EVERY historical row + carries ``source_position + 1``. In-chunk rows are never checked. + """ + expected = qsa_ring_row_tags(source_positions) + bad = historical & (found_tags.to(torch.int64) != expected) + return ~bad.any(dim=-1) + def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor: """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows.""" @@ -868,37 +904,65 @@ class _QSAStateCache(nn.Module, AttentionLayerBase): class QSAKeyStateCache(_QSAStateCache): - """Raw BF16 key, optionally followed by exact int64 MRoPE positions.""" + """Raw BF16 key followed by the int64 ring tail (pos_t, pos_h, pos_w, tag). - _BF16_PER_INT64 = 4 + The position words are consumed only when ``cache_rope_positions``; the tag + word is always written and always validated.""" + + _BF16_PER_INT64 = QSA_RING_BF16_PER_INT64 _NUM_ROPE_AXES = 3 def __init__(self, *, cache_rope_positions: bool = False, **kwargs) -> None: key_head_size = int(kwargs.pop("head_size")) self.key_head_size = key_head_size self.cache_rope_positions = bool(cache_rope_positions) + # Element offset of the int64 tail (the name predates the tag word). self.rope_position_offset = ( (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64 ) * self._BF16_PER_INT64 - storage_head_size = key_head_size - if self.cache_rope_positions: - storage_head_size = self.rope_position_offset + ( - self._NUM_ROPE_AXES * self._BF16_PER_INT64 - ) + storage_head_size = self.rope_position_offset + QSA_RING_TAIL_ELEMENTS super().__init__(head_size=storage_head_size, **kwargs) + self._ring_tag_cache: torch.Tensor | None = None + self._ring_tail: torch.Tensor | None = None + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + super().bind_kv_cache(kv_cache) + qsa_cache = self.kv_cache + tail_end = self.rope_position_offset + QSA_RING_TAIL_ELEMENTS + tail = qsa_cache[..., self.rope_position_offset : tail_end].view(torch.int64) + assert tail.shape[-1] == QSA_RING_TAIL_WORDS + # Derived on access so `kv_cache` stays the only reference to the bound + # storage: clearing it (e.g. after CUDA graph memory profiling) must free + # the cache, or the freed block stays pinned in the allocator. Only the + # tag word is read outside a bind, so it is the one cached attribute. + self._ring_tail = tail + self._ring_tag_cache = tail[..., QSA_RING_TAG_WORD : QSA_RING_TAG_WORD + 1] - # Derived on access so `kv_cache` stays the only reference to the bound - # storage: clearing it (e.g. after CUDA graph memory profiling) must free - # the cache, or the freed block stays pinned in the allocator. @property def key_cache(self) -> torch.Tensor: return self.kv_cache[..., : self.key_head_size] @property def rope_position_cache(self) -> torch.Tensor | None: - if not self.cache_rope_positions: + if not self.cache_rope_positions or self._ring_tail is None: return None - return self.kv_cache[..., self.rope_position_offset :].view(torch.int64) + return self._ring_tail[..., : self._NUM_ROPE_AXES] + + @property + def ring_tag_cache(self) -> torch.Tensor | None: + return self._ring_tag_cache + + def clear_ring_row_tags(self, block_ids: torch.Tensor) -> None: + """Zero the tag word of every row of ``block_ids`` (ring-row tag claim clear). + + Only the tag word: a claimed ring block's keys and positions stay as + the previous occupant left them, and the zero tag is what makes every + such row fail validation until this request commits it.""" + if self._ring_tag_cache is None: + raise RuntimeError( + f"{self.prefix}: ring tag clear before the KV cache was bound" + ) + self._ring_tag_cache.index_fill_(0, block_ids, 0) def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # Hold the open group's committed keys plus every row a speculative @@ -936,6 +1000,9 @@ class QSACompressedKeyCache(_QSAStateCache): __all__ = [ + "QSA_RING_TAG_WORD", + "QSA_RING_TAIL_ELEMENTS", + "QSA_RING_TAIL_WORDS", "QSACompressedKeyCache", "QSAForwardMetadata", "QSAKeyStateCache", @@ -944,4 +1011,6 @@ __all__ = [ "canonical_qsa_rope_positions", "circular_qsa_slot_mapping", "compressed_qsa_slot_mapping", + "qsa_ring_group_rows_valid", + "qsa_ring_row_tags", ] diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py index c8ffa0dbfb..d1268a89e7 100644 --- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py @@ -20,6 +20,7 @@ from ..common.qsa_cache import ( QSAForwardMetadata, QSAKeyStateCache, canonical_qsa_rope_positions, + qsa_ring_row_tags, ) from .ops.qsa_prepare import qsa_prepare @@ -339,6 +340,7 @@ class QSAIndexer(nn.Module): main_slot_mapping=slot_mapping[:num_tokens], main_k_scale=attn._k_scale_float, main_v_scale=attn._v_scale_float, + layer_id=self.layer_id, ) else: # Unfused reference path @@ -405,6 +407,18 @@ class QSAIndexer(nn.Module): raw_metadata.slot_mapping, position_rows, ) + # The unfused writer commits the same tag word as the fused + # one, so both paths leave byte-identical rings. (This reference + # reader does not validate tags; the fused kernel does.) + ring_tag_cache = getattr(raw_key_state_cache, "ring_tag_cache", None) + if ring_tag_cache is not None: + qsa_store_cache_rows( + ring_tag_cache, + raw_metadata.slot_mapping, + qsa_ring_row_tags(raw_metadata.logical_positions[:num_tokens]).view( + -1, 1, 1 + ), + ) if self.skip_topk: if out is None: diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py index d5e809ee0e..f1ae4405a4 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py @@ -4,6 +4,10 @@ import torch +from vllm.models.qwen4_exp.common.qsa_cache import ( + QSA_RING_TAG_WORD, + QSA_RING_TAIL_ELEMENTS, +) from vllm.triton_utils import tl, triton from vllm.v1.worker import guard_telemetry @@ -12,6 +16,14 @@ from vllm.v1.worker import guard_telemetry # from a @triton.jit body -- an earlier build shipped one and every guarded kernel died # with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. _GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) +# Int64 word of the ring tail that carries `logical_position + 1`. +_RING_TAG_WORD = tl.constexpr(QSA_RING_TAG_WORD) +# Ring-row tag record packing and the encodable range of the found-tag slot. +_REC_POS_BITS = tl.constexpr(guard_telemetry.RING_RECORD_POS_BITS) +_REC_REQ_BITS = tl.constexpr(guard_telemetry.RING_RECORD_REQ_BITS) +_REC_LAYER_BITS = tl.constexpr(guard_telemetry.RING_RECORD_LAYER_BITS) +_GUARD_MAX_ENC = tl.constexpr(guard_telemetry.GUARD_MAX_ENCODABLE) +_GUARD_MIN_ENC = tl.constexpr(guard_telemetry.GUARD_MIN_ENCODABLE) @@ -184,6 +196,7 @@ def _store_rotated(dst, y, o1, o2, scale): "num_compressed_blocks", "num_k_work", "num_requests", + "layer_id", ] ) def _qsa_prepare_kernel( @@ -223,6 +236,8 @@ def _qsa_prepare_kernel( num_requests, # bounds-guard instrumentation guard_ctr, # bounds-guard instrumentation: persistent int64[3] ring_skip_ctr, # ring-skip telemetry, persistent int64[3], SITE_QSA_RING_SKIP + ring_uncommitted_ctr, # ring-row tag fix, persistent int64[3], SITE_QSA_RING_UNCOMMITTED + layer_id, # ring-row tag fix: named in the uncommitted-row record only HQ: tl.constexpr, D: tl.constexpr, TILE_T_Q: tl.constexpr, @@ -446,6 +461,64 @@ def _qsa_prepare_kernel( + safe_state_block * state_cache_stride_block + (source_positions % STATE_SIZE)[:, None] * state_cache_stride_token ) + # Ring-row tag fix: EVERY historical source row must carry the tag the + # writer stores for exactly this logical position, before anything + # reads its key or its position tail. A missing ring block (state + # table entry out of range) reads tag 0 and fails the same way -- + # before the ring-row tag fix it pooled zeros silently. In-chunk rows come from + # this launch's own k and are never checked. + historical = valid & ~source_in_chunk + source_tags = tl.load( + ( + state_cache_ptr + + safe_state_block * state_cache_stride_block + + (source_positions % STATE_SIZE) * state_cache_stride_token + + D + ).to(tl.pointer_type(tl.int64)) + + _RING_TAG_WORD, + mask=historical & state_block_valid, + other=0, + ) + row_bad = historical & (source_tags != source_positions.to(tl.int64) + 1) + num_bad_rows = tl.sum(row_bad.to(tl.int32), axis=0) + group_ok = num_bad_rows == 0 + if num_bad_rows > 0: + # One coherent record per site, first offender wins: the CAS + # claims slot 1 and only its winner writes slot 2, so both + # fields always describe the same row. + bad_position = tl.min( + tl.where(row_bad, source_positions.to(tl.int64), 1 << 62), + axis=0, + ) + bad_tag = tl.sum( + tl.where( + row_bad & (source_positions.to(tl.int64) == bad_position), + source_tags, + 0, + ), + axis=0, + ) + record = ( + ( + ((layer_id + 1).to(tl.int64) & ((1 << _REC_LAYER_BITS) - 1)) + << (_REC_REQ_BITS + _REC_POS_BITS) + ) + | ( + (request.to(tl.int64) & ((1 << _REC_REQ_BITS) - 1)) + << _REC_POS_BITS + ) + | (bad_position & ((1 << _REC_POS_BITS) - 1)) + ) + tl.atomic_add(ring_uncommitted_ctr, num_bad_rows.to(tl.int64)) + zero64 = bad_position * 0 + prev = tl.atomic_cas( + ring_uncommitted_ctr + 1, zero64, record + _GUARD_BIAS + ) + if prev == 0: + clamped_tag = tl.minimum( + tl.maximum(bad_tag, _GUARD_MIN_ENC), _GUARD_MAX_ENC + ) + tl.atomic_xchg(ring_uncommitted_ctr + 2, clamped_tag + _GUARD_BIAS) # Only the first completed group can cross the chunk boundary. Select # historical rows from the ring without issuing two masked loads. source_base = tl.where(source_in_chunk[:, None], current_base, cached_base) @@ -454,9 +527,11 @@ def _qsa_prepare_kernel( source_valid = tl.where( source_in_chunk, source_tokens_valid, state_block_valid ) + # A failed group loads nothing (the mask, not a tl.where on the + # result, is what suppresses the memory operation). source = tl.load( source_base + dims[None, :], - mask=valid & source_valid[:, None], + mask=valid & group_ok & source_valid[:, None], other=0.0, ).to(tl.float32) # Match the unfused path's BF16 pooled tensor before RMSNorm. @@ -504,7 +579,9 @@ def _qsa_prepare_kernel( + (first_position % STATE_SIZE) * state_cache_stride_token + D ).to(tl.pointer_type(tl.int64)) - load_cached_position = ~first_in_chunk & state_block_valid + # The first row's tag was validated above with the rest; + # a failed group never reads its tail (pos stays 0, in range). + load_cached_position = ~first_in_chunk & state_block_valid & group_ok cached_pos_t = tl.load(tail, mask=load_cached_position, other=0) cached_pos_h = tl.load(tail + 1, mask=load_cached_position, other=0) cached_pos_w = tl.load(tail + 2, mask=load_cached_position, other=0) @@ -532,13 +609,16 @@ def _qsa_prepare_kernel( ) compressed_block = (compressed_slot // COMP_PAGE_SIZE).to(tl.int64) compressed_row = compressed_slot % COMP_PAGE_SIZE + # A group built from an uncommitted row must never become + # usable compressed state. The slot stays unwritten; the site is + # fatal because the selection kernels will still read it. tl.store( compressed_cache_ptr + compressed_block * compressed_cache_stride_block + compressed_row * compressed_cache_stride_token + dims, tl.reshape(y, (D,)), - mask=valid, + mask=valid & group_ok, ) if work_in_request == 0: @@ -609,6 +689,20 @@ def _qsa_prepare_kernel( tl.store(tail, pos_t.to(tl.int64), mask=valid_slot) tl.store(tail + 1, pos_h.to(tl.int64), mask=valid_slot) tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot) + # Ring-row tag fix: the tag names the logical position this row now + # holds, taken from the authoritative per-token metadata (not + # re-derived from chunk arithmetic), stored after key + tail. + # Every committed row gets one, with or without RoPE positions. + token_logical = tl.load( + logical_positions_ptr + tl.maximum(token, 0), + mask=valid_slot, + other=-1, + ) + tl.store( + (state_row + D).to(tl.pointer_type(tl.int64)) + _RING_TAG_WORD, + token_logical.to(tl.int64) + 1, + mask=valid_slot, + ) def qsa_prepare( @@ -640,11 +734,16 @@ def qsa_prepare( main_slot_mapping: torch.Tensor, main_k_scale: float, main_v_scale: float, + layer_id: int = -1, ) -> tuple[torch.Tensor, torch.Tensor]: """Normalize Q, compress K, then update the circular raw state. Also prepares the main attention (QK-norm/RoPE, gate copy, K/V cache write) and returns its Q and gate. + + ``state_cache`` rows are ``[key D bf16 | 4 int64: pos_t, pos_h, pos_w, tag]`` + (``QSAKeyStateCache``, ring-row tag fix); the tail is present with or without + ``rope_pos_offset``. ``layer_id`` only labels the uncommitted-row record. """ num_tokens = q.shape[0] main_head_dim = main_kv_cache.shape[-1] // 2 @@ -669,6 +768,15 @@ def qsa_prepare( is_k_mrope = bool(mrope_section) cache_has_rope_pos = rope_pos_offset is not None assert rope_pos_offset is None or rope_pos_offset == head_dim + # The kernel addresses the tail at +head_dim and the tag at word 3 of + # it. A ring without the four-word tail (the older D or D+12 layouts) + # would put the tag store into the NEXT row, so refuse it outright. + assert state_cache.shape[-1] == head_dim + QSA_RING_TAIL_ELEMENTS, ( + f"QSA ring row width {state_cache.shape[-1]} != head_dim + " + f"{QSA_RING_TAIL_ELEMENTS} (key + pos_t/h/w + tag)" + ) + assert state_cache.dtype == torch.bfloat16 + assert state_cache.stride(1) % 4 == 0 and state_cache.stride(0) % 4 == 0 if is_2d_positions: assert positions.shape == (3, num_tokens) assert is_k_mrope @@ -727,6 +835,8 @@ def qsa_prepare( min(query_start_loc.numel() - 1, state_block_table.shape[0]), guard_telemetry.counter(guard_telemetry.SITE_QSA_PRE_INDEXER, q.device), guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_SKIP, q.device), + guard_telemetry.counter(guard_telemetry.SITE_QSA_RING_UNCOMMITTED, q.device), + int(layer_id), HQ=num_q_heads, D=head_dim, TILE_T_Q=TILE_T_Q, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 37ceb4e474..ad7bf18090 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -867,6 +867,21 @@ class KVCacheManager: ids.extend(mgr.take_new_block_ids()) return ids + def take_new_ring_block_ids(self) -> list[int]: + """Drain ring blocks claimed since the last call (ring-row tag clear). + Independent of KV zeroing: every claim is returned.""" + ids: list[int] = [] + for mgr in self.coordinator.single_type_managers: + ids.extend(mgr.take_new_ring_block_ids()) + return ids + + def get_ring_block_ids(self, request_id: str) -> list[int]: + """The request's ring (CircularBufferSpec) block ids, all groups.""" + ids: list[int] = [] + for mgr in self.coordinator.single_type_managers: + ids.extend(mgr.get_ring_block_ids(request_id)) + return ids + def record_blocks_for_zeroing(self, request_id: str, start_token: int) -> None: """Re-record the request's blocks from start_token onwards for zeroing, e.g. blocks a failed async KV load left unwritten. diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 9c42d91b38..46464fcb7d 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -301,6 +301,13 @@ class SchedulerOutput: # preventing stale NaN/data from corrupting attention or SSM computation. new_block_ids_to_zero: list[int] | None = None + # v2.2 ring-row tag fix: ring (CircularBufferSpec) blocks claimed during this + # scheduling step. Independent of new_block_ids_to_zero and of + # needs_kv_cache_zeroing: the worker clears only the rows' tag words, after + # the zeroer and before any copy or forward of this step, i.e. before the + # new owner's first read or write of the ring. + new_ring_block_ids_to_clear: list[int] | None = None + # CoW copies to apply after zeroing new blocks and before forward. kv_cache_block_copies: list[KVCacheBlockCopy] | None = None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index d8aa1b2d2f..e93c0c5e7c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -366,6 +366,11 @@ class Scheduler(SchedulerInterface): # Bind after construction so connectors can access the cache manager. if self.connector is not None: self.connector.bind_kv_cache_manager(self.kv_cache_manager) + # Ring blocks an async KV load is about to fill. Their tag clear + # is skipped for the same reason as zeroing (it could land after the + # out-of-band write and wipe the producer's committed tags). Tracked + # whatever needs_kv_cache_zeroing says, like the clear itself. + self._skip_ring_clear_block_ids: set[int] = set() self.use_pp = self.parallel_config.pipeline_parallel_size > 1 self.use_v2_model_runner = vllm_config.use_v2_model_runner @@ -1333,6 +1338,12 @@ class Scheduler(SchedulerInterface): request.num_computed_tokens = num_computed_tokens self._inflight_prefills.add(request) self._set_kv_fetch_stage(request, KV_FETCH_IN_PROGRESS) + # Same race for the ring's tag clear; the load + # brings the producer's rows and tags (or none, and the + # reader then fails closed on the missing history). + self._skip_ring_clear_block_ids.update( + self.kv_cache_manager.get_ring_block_ids(request.request_id) + ) skip_request(request_queue) continue @@ -1525,6 +1536,8 @@ class Scheduler(SchedulerInterface): finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=self._get_new_block_ids_to_zero(), + # Drained every step, never gated on KV zeroing. + new_ring_block_ids_to_clear=self._get_new_ring_block_ids_to_clear(), has_sync_kv_loads=has_sync_kv_loads, kv_cache_block_copies=pending_kv_cache_block_copies, kv_connector_block_state=kv_connector_block_state, @@ -1572,6 +1585,17 @@ class Scheduler(SchedulerInterface): ) -> KVConnectorMetadata: return connector.build_connector_meta(scheduler_output) + def _get_new_ring_block_ids_to_clear(self) -> list[int] | None: + """Ring claims of this step whose tag words the worker clears. + Unconditional (not under needs_kv_cache_zeroing); only blocks an async + KV load is about to write are left out.""" + ids = self.kv_cache_manager.take_new_ring_block_ids() + if self._skip_ring_clear_block_ids: + skip = self._skip_ring_clear_block_ids + ids = [b for b in ids if b not in skip] + skip.clear() + return ids or None + def _get_new_block_ids_to_zero(self) -> list[int] | None: # Drain new attention block ids every step so the manager-side list # does not grow unbounded; only kv-cache zeroing consumes them. diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 7f23a886ed..e4bdbc62e7 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -121,6 +121,12 @@ class SingleTypeKVCacheManager(ABC): and not isinstance(kv_cache_spec, CircularBufferSpec) ) self.new_block_ids: list[int] = [] + # v2.2 ring-row tag fix: ring blocks claimed since the last drain. The ring stays + # OUT of whole-block zeroing above; instead every claim is recorded + # UNCONDITIONALLY (never under needs_kv_cache_zeroing) so the worker + # clears the rows' tag words before the new owner's first step. Only + # CircularBufferManager appends here. + self.new_ring_block_ids: list[int] = [] # Mapping from request ID to blocks to track the blocks allocated # for each request, so that we can free the blocks when the request @@ -422,6 +428,17 @@ class SingleTypeKVCacheManager(ABC): self.new_block_ids = [] return ids + def take_new_ring_block_ids(self) -> list[int]: + """Drain ring blocks claimed since the last call; empty for + every manager but CircularBufferManager.""" + ids = self.new_ring_block_ids + self.new_ring_block_ids = [] + return ids + + def get_ring_block_ids(self, request_id: str) -> list[int]: + """The request's ring block ids; empty for non-ring managers.""" + return [] + def take_pending_cow_copies( self, ) -> list[tuple[KVCacheBlock, KVCacheBlock]]: @@ -1214,8 +1231,37 @@ class CircularBufferManager(FullAttentionManager): req_blocks.extend(new_blocks) if self._record_new_block_ids and record_for_zeroing: self.new_block_ids.extend(block.block_id for block in new_blocks) + # Unconditional -- not gated on needs_kv_cache_zeroing. + self.new_ring_block_ids.extend(block.block_id for block in new_blocks) return new_blocks + def take_new_ring_block_ids(self) -> list[int]: + """Drain the ring blocks claimed since the last call, keeping only + blocks a request of this manager still holds, in claim order, once. + + A block claimed and then freed within the same scheduling pass (the + request was preempted after its allocation) may already belong to + another group by the time the worker acts; clearing its tag words + there would write into that group's page, so it is dropped here. + """ + claimed = self.new_ring_block_ids + self.new_ring_block_ids = [] + if not claimed: + return [] + held = { + block.block_id for blocks in self.req_to_blocks.values() for block in blocks + } + out: list[int] = [] + seen: set[int] = set() + for block_id in claimed: + if block_id in held and block_id not in seen: + seen.add(block_id) + out.append(block_id) + return out + + def get_ring_block_ids(self, request_id: str) -> list[int]: + return [block.block_id for block in self.req_to_blocks.get(request_id, ())] + def get_num_blocks_to_allocate( self, request_id: str, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 374a1df853..9e0c5b77c7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -180,6 +180,7 @@ from vllm.v1.worker.gpu.ubatch_utils import ( maybe_build_ubatch_runner, ) from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin +from vllm.v1.worker.ring_tags import clear_ring_row_tags, ring_tag_layers from vllm.v1.worker.utils import ( KVBlockZeroer, clear_layer_kv_caches, @@ -1326,6 +1327,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): assert self.kv_block_zeroer is not None self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero) + # v2.2 ring-row tag fix: clear the tag words of newly claimed ring blocks. Not + # under the zeroing gate; after the zeroer, before copies and forward. + if scheduler_output.new_ring_block_ids_to_clear: + clear_ring_row_tags( + ring_tag_layers(self.compilation_config.static_forward_context), + scheduler_output.new_ring_block_ids_to_clear, + self.device, + ) + # Apply copy-on-write block copies for partial prefix-cache hits, after # zeroing new blocks and before the forward pass reads them. if scheduler_output.kv_cache_block_copies: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 03b11d8909..86774b8a60 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -214,6 +214,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin +from vllm.v1.worker.ring_tags import clear_ring_row_tags, ring_tag_layers from vllm.v1.worker.ubatch_utils import ( UBatchSlices, check_ubatch_thresholds, @@ -1218,6 +1219,14 @@ class GPUModelRunner( # stale NaN/data from corrupting attention or SSM computation. if scheduler_output.new_block_ids_to_zero: self._zero_block_ids(scheduler_output.new_block_ids_to_zero) + # v2.2 ring-row tag fix: clear the tag words of newly claimed ring blocks. Not + # under the zeroing gate; after the zeroer, before copies and forward. + if scheduler_output.new_ring_block_ids_to_clear: + clear_ring_row_tags( + ring_tag_layers(self.compilation_config.static_forward_context), + scheduler_output.new_ring_block_ids_to_clear, + self.device, + ) if scheduler_output.kv_cache_block_copies: copy_kv_cache_blocks_inplace( self.kv_caches, diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py index adca93049f..63c450f14d 100644 --- a/vllm/v1/worker/guard_telemetry.py +++ b/vllm/v1/worker/guard_telemetry.py @@ -146,7 +146,53 @@ SITE_QSA_ROPE_POS = 9 # The skip itself is the earlier behaviour; the count is the instrument # for the v2.2 ring-row reconciliation. Telemetry only: warn. SITE_QSA_RING_SKIP = 10 -NUM_SITES = 11 +# The pre-indexer found a historical raw-key ring row whose tag is +# not `logical_position + 1` -- a row this request never committed for that +# position (claim residue, a skipped suffix commit, a stale mapping). The +# group's pooled load, its RoPE-tail read and its compressed-slot store are all +# MASKED, so the compressed slot is left unwritten: durable damage, fatal. +# Unlike the other sites the three slots form ONE coherent record, claimed +# first-offender-wins with atomic_cas: slot 0 = rows flagged (all offenders), +# slot 1 = packed (layer + 1, batch request row, source logical position) of +# the first offender, slot 2 = the tag that offender's row actually held. +# Decode slot 1 with decode_ring_record(). +SITE_QSA_RING_UNCOMMITTED = 11 +NUM_SITES = 12 + +# Bit layout of the SITE_QSA_RING_UNCOMMITTED record (before GUARD_BIAS). +RING_RECORD_POS_BITS = 32 +RING_RECORD_REQ_BITS = 16 +RING_RECORD_LAYER_BITS = 14 + + +def encode_ring_record(layer_id: int, request_row: int, position: int) -> int: + """Python twin of the kernel's packing (tests; the kernel inlines it).""" + layer = (int(layer_id) + 1) & ((1 << RING_RECORD_LAYER_BITS) - 1) + row = int(request_row) & ((1 << RING_RECORD_REQ_BITS) - 1) + pos = int(position) & ((1 << RING_RECORD_POS_BITS) - 1) + return ( + (layer << (RING_RECORD_REQ_BITS + RING_RECORD_POS_BITS)) + | (row << RING_RECORD_POS_BITS) + | pos + ) + + +def decode_ring_record(biased: int) -> dict[str, int] | None: + """Unpack slot 1 of SITE_QSA_RING_UNCOMMITTED. ``layer`` is -1 when the + launcher was not given one. None = nothing recorded (or not yet visible).""" + value = decode_slot(biased) + if value is None: + return None + pos_mask = (1 << RING_RECORD_POS_BITS) - 1 + req_mask = (1 << RING_RECORD_REQ_BITS) - 1 + layer_mask = (1 << RING_RECORD_LAYER_BITS) - 1 + return { + "layer": ((value >> (RING_RECORD_REQ_BITS + RING_RECORD_POS_BITS)) & layer_mask) + - 1, + "request_row": (value >> RING_RECORD_POS_BITS) & req_mask, + "source_position": value & pos_mask, + } + SITE_NAMES = { SITE_MAMBA_COPY: "_copy_mamba_state_block (block-table row/column)", @@ -165,6 +211,10 @@ SITE_NAMES = { SITE_QSA_RING_SKIP: "_qsa_pre_indexer_kernel suffix commit SKIPPED " "(state slot < 0 or >= ring capacity; slot 1 = token index, " "slot 2 = state slot value)", + SITE_QSA_RING_UNCOMMITTED: "_qsa_pre_indexer_kernel historical ring row " + "UNCOMMITTED (tag != logical_position + 1; group load, RoPE tail and " + "compressed-slot store masked; slot 1 = packed first-offender record, " + "slot 2 = the tag its row held)", } @@ -373,6 +423,14 @@ _ALWAYS_FATAL_SITES: set[int] = set() # whatever VLLM_E2_GUARD_MODE says. (bounds-guard instrumentation, rope site) _ALWAYS_FATAL_SITES.add(SITE_QSA_ROPE_POS) +# An uncommitted ring row leaves the group's compressed slot +# UNWRITTEN (the store is masked, there is no correct value to store), and the +# selection kernels will read that slot on every later step. That is durable +# damage by the same argument as the rope site, so it stops the rank under +# warn too. The failure boundary is the rank, not the request: per-request +# containment (discard + quarantine + TP/PP coordination) is not implemented. +_ALWAYS_FATAL_SITES.add(SITE_QSA_RING_UNCOMMITTED) + def _normalize(device: torch.device) -> torch.device: """torch.device("cuda") and torch.device("cuda:0") are different dict keys. @@ -479,6 +537,19 @@ def _report_rows(device: torch.device, rows: list[list[int]]) -> None: _seen[(device, site)] = n_masked row = decode_slot(biased_row) val = decode_slot(biased_val) + if site == SITE_QSA_RING_UNCOMMITTED: + # One coherent first-offender record, not independent maxima. + record = decode_ring_record(biased_row) + msg = ( + f"ring-row tag guard masked {n_masked} uncommitted historical " + f"ring row(s) in {SITE_NAMES[site]} on {device}; first " + f"offender {record}, its row held tag {val} (expected " + f"source_position + 1). The group was not pooled and its " + f"compressed slot was NOT written." + ) + logger.warning("%s", msg) + _arm(msg) + continue msg = ( f"bounds guard masked {n_masked} out-of-range access(es) " f"in {SITE_NAMES.get(site, site)} on {device} " diff --git a/vllm/v1/worker/ring_tags.py b/vllm/v1/worker/ring_tags.py new file mode 100644 index 0000000000..ef0a8b7c4a --- /dev/null +++ b/vllm/v1/worker/ring_tags.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""v2.2 ring-row tag fix: worker side of the ring-row tag clear. + +The scheduler reports every newly claimed ring (``CircularBufferSpec``) block +in ``SchedulerOutput.new_ring_block_ids_to_clear``, whether or not KV zeroing +is on. Before the step's forward, the runner hands them to every layer that +keeps a tagged ring (``clear_ring_row_tags``; today ``QSAKeyStateCache``), +which zeroes only the tag word of each row. A zero tag never validates, so the +new owner cannot pool a row it did not commit itself, whatever the previous +occupant (or another group that held the block) left there. + +ORDER. The runners call this right after the KV zeroer and before the CoW +copies and the forward, on the same stream, inside the same execute_model as +the claim's first step. So: previous owner's work done (deferred free) -> +claim -> tag clear -> new owner's first read/write. +""" + +from collections.abc import Iterable, Mapping + +import torch + + +def ring_tag_layers(static_forward_context: Mapping[str, object]) -> list: + """Layers of this worker that keep a tagged ring.""" + return [ + layer + for layer in static_forward_context.values() + if callable(getattr(layer, "clear_ring_row_tags", None)) + ] + + +def clear_ring_row_tags( + layers: Iterable, + block_ids: list[int], + device: torch.device, +) -> int: + """Zero the tag word of every row of ``block_ids`` in every layer. + + Returns the number of layers cleared. Raises (fail loud) on a block id + outside a layer's pool: an index_fill_ with a bad index is a device-side + assert, far from its cause.""" + if not block_ids: + return 0 + layers = list(layers) + if not layers: + return 0 + lo, hi = min(block_ids), max(block_ids) + index = torch.tensor(block_ids, dtype=torch.long, device=device) + for layer in layers: + tags = getattr(layer, "ring_tag_cache", None) + if tags is None: + raise RuntimeError( + f"ring tag clear: layer {getattr(layer, 'prefix', layer)!r} has " + "no bound ring_tag_cache" + ) + if lo < 0 or hi >= tags.shape[0]: + raise ValueError( + f"ring tag clear: block ids [{lo}, {hi}] outside the pool of " + f"{tags.shape[0]} blocks ({getattr(layer, 'prefix', layer)!r})" + ) + layer.clear_ring_row_tags(index) + return len(layers)