vllm.model_executor.layers.attention.attn_capture ¶

class AttentionCapture:
    """Per-worker attention capture state.

    Instantiated once per ModelRunner. Buffers Q vectors during inference
    and computes attention scores at request completion time.
    """

    def __init__(self, config: CaptureConfig, model_config=None):
        self.config = config
        self.model_config = model_config
        # Ephemeral: (layer_idx, slot_id) -> Q tensor, cleared after capture().
        self.q_buffer: dict[tuple[int, int], torch.Tensor] = {}
        # Persistent across requests while KV slot lives.
        # Allows full [T,H,T] capture with prefix caching (cache hits reuse initial Q).
        # FIFO-capped by _q_cache_max to limit memory.
        self.q_cache: dict[tuple[int, int], torch.Tensor] = {}
        self._q_cache_max: int = 32768  # ~128MB at float16 128d 16h
        self._layer_idx_cache: dict[str, int] = {}
        self.runtime_enabled_this_step = False
        self.capture_slots: set[int] | None = None

    def buffer_query(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        attn_metadata,
        layer_name: str,
    ) -> None:
        """
        Buffer Q tokens at attention-computation time.
        K is NOT buffered — it is read from KV cache at capture time
        (request completion).
        """
        if not self.config.enabled or attn_metadata is None:
            return

        layer_idx = self.extract_layer_idx(layer_name)
        if layer_idx < 0:
            return

        # Extract slot_ids using attn_metadata from the gpu_model_runner process
        slot_ids = attn_metadata.slot_mapping
        if query.shape[0] != slot_ids.shape[0]:
            return

        # Detaching the query tensor for buffering. Severing CUDA computation trace
        # Query tensor stays in the dict for all requests
        try:
            query_cpu = query.detach().cpu().clone()
        except Exception:
            return  # This error may occur due to gpu memory access issues

        capture_slots = self.capture_slots
        for i in range(query.shape[0]):
            slot_id = slot_ids[i].item()

            # Buffer queries for only set requests at _update_attn_capture_slots()
            if slot_id < 0 or (capture_slots and slot_id not in capture_slots):
                continue

            # cast to float16 for uniform buffer dtype; no-op if already float16
            self.q_buffer[(layer_idx, slot_id)] = query_cpu[i].to(torch.float16)

    def capture(
        self,
        req_state,
        block_size: int,
        kv_caches: list,
        prefix: str | None = None,
    ) -> None:
        """Capture attention scores for a completed request.

        Called after request completion. Computes Q*K attention on GPU
        and writes results to shared memory.
        """
        req_id = req_state.req_id
        layers = resolve_target_layers(req_state, self.config.layers)
        snapshots: list[dict] = []

        for layer_idx in layers:
            # Collect slots that have Q in either the ephemeral buffer or
            # the persistent q_cache (for prefix-cached tokens).
            buf_slots = [sid for (li, sid) in self.q_buffer if li == layer_idx]
            cached_slots = [sid for (li, sid) in self.q_cache if li == layer_idx]
            q_slots = set(buf_slots) | set(cached_slots)
            if not q_slots or not req_state.block_ids:
                continue
            # Find which block group contains buffered slots for this layer
            grp_idx = None
            for gi, block_list in enumerate(req_state.block_ids):
                if not block_list:
                    continue
                if q_slots & slots_from_blocks(block_list, block_size):
                    grp_idx = gi
                    break
            if grp_idx is None:
                continue

            # From the current request block group,
            # order the slot_ids for alignment with QK^T and deduplication.
            slots = ordered_slots_for_group(
                req_state.block_ids[grp_idx], req_state.num_tokens, block_size
            )
            if not slots:
                continue
            # Collect Q tensors from buffer (falling back to q_cache for
            # prefix-cached tokens that were not re-computed this request).
            q_list, q_sids, tok_idx = self._collect_query(layer_idx, slots)
            if not q_list:
                continue

            # Persist Q vectors so future requests can reuse them on cache hit.
            # Evict oldest entries if cap is reached (simple FIFO approximation).
            for sid, q in zip(q_sids, q_list):
                if len(self.q_cache) >= self._q_cache_max:
                    self.q_cache.pop(next(iter(self.q_cache)))
                self.q_cache[(layer_idx, sid)] = q

            # NOTE(jehyun): This class receives the KV Cache from the worker.
            # Reads the KV Cache right after the request is finished
            # before the KV cache is freed.
            # Extract K for ALL slots (not just q_sids) so the key axis covers
            # the full sequence, including prefix-cached tokens.
            kv_idx = layer_idx if layer_idx < len(kv_caches) else grp_idx
            if not kv_caches or kv_idx is None or kv_idx >= len(kv_caches):
                continue
            k_raw = extract_k_from_kv_cache(kv_caches[kv_idx], slots)
            k_list = [k_raw[i] for i in range(k_raw.shape[0])]
            if not k_list:
                continue

            # Build tok_idx covering all slots; q_list entries map to q_sids
            # positions within the full slot list for the QK^T alignment.
            sid_to_pos = {sid: pos for pos, sid in enumerate(slots)}
            q_tok_idx = [sid_to_pos[sid] for sid in q_sids]

            # Checking for alignment before QK Calculation for safety issues
            q_tok_idx, q_list, _ = self._filter_compatible_qk(
                q_tok_idx, q_list, [k_list[i] for i in q_tok_idx]
            )
            if not q_list:
                continue
            tok_idx = q_tok_idx

            # Trim K to the same token positions as Q so shapes always match.
            k_list = [k_list[i] for i in q_tok_idx]
            # Compute Q*K^T: Q=[T, H, d], K=[T, H, d] → [T, H, T]
            q_t, k_t = torch.stack(q_list), torch.stack(k_list)
            if k_t.is_cuda and not q_t.is_cuda:
                q_t = q_t.to(k_t.device)
            scale = 1.0 / (q_t.shape[2] ** 0.5)
            attn = compute_qk_attention(q_t, k_t, scale)
            if attn is None:
                continue

            # Apply prefix slice
            if prefix:
                parts = prefix.split(":")
                q_start = int(parts[0]) if parts[0] else 0
                q_end = int(parts[1]) if len(parts) > 1 and parts[1] else None
                attn = attn[q_start:q_end, :, :]
                tok_idx = tok_idx[q_start:q_end]

            # Create meta_data for tokens, later used by clients for 1:1 matching
            # between attn_scores and token_idx
            tmeta = build_token_meta(req_state, tok_idx, ordered_slots_len=len(slots))
            snapshots.append(encode_snapshot(attn, layer_idx, tmeta))

            # Remove processed slots from q_buffer to prevent
            # cross-request contamination.
            for sid in set(slots):
                self.q_buffer.pop((layer_idx, sid), None)

        # NOTE(jehyun): Write snapshots to shared memory for cross-process delivery.
        # The output_processor (main process) reads them via _shm_read().
        if snapshots:
            _shm_write(req_id, snapshots)

    def cleanup_request_buffers(
        self,
        block_ids: list[list[int]],
        block_size: int,
    ) -> None:
        """
        Remove buffered Q vectors for a finished request.
        Called for ALL requests to prevent stale data leaking into future requests.
        q_cache is evicted only when the underlying KV block is reclaimed.
        """
        if not block_ids:
            return
        remove: set[int] = set()
        for block_list in block_ids:
            remove |= slots_from_blocks(block_list, block_size)
        if self.q_buffer:
            for k in [k for k in self.q_buffer if k[1] in remove]:
                del self.q_buffer[k]

    def extract_layer_idx(self, layer_name: str) -> int:
        """Parse layer index from attention layer name, with caching."""
        v = self._layer_idx_cache.get(layer_name)
        if v is not None:
            return v

        for pat in _LAYER_PATTERNS:
            m = pat.search(layer_name)
            if not m:
                continue
            idx = int(m.group(1))
            self._layer_idx_cache[layer_name] = idx
            return idx

        self._layer_idx_cache[layer_name] = -1
        return -1

    def _collect_query(
        self,
        layer_idx: int,
        slots: list[int],
    ) -> tuple[list[torch.Tensor], list[int], list[int]]:
        """
        Collect Q tensors from buffer in deterministic slot order.
        Falls back to q_cache for prefix-cached tokens whose Q was not
        re-computed this request (prefix caching ON).
        Returns (q_list, q_slot_ids, tok_idx).
        """
        q_list, q_sids, tok_idx = [], [], []
        for si, sid in enumerate(slots):
            key = (layer_idx, sid)
            q = self.q_buffer.get(key)
            if q is None:
                q = self.q_cache.get(key)
            if q is None:
                continue
            q_list.append(q)
            q_sids.append(sid)
            tok_idx.append(si)
        return q_list, q_sids, tok_idx

    @staticmethod
    def _filter_compatible_qk(
        tok_idx: list[int],
        q_list: list[torch.Tensor],
        k_list: list[torch.Tensor],
    ) -> tuple[list[int], list[torch.Tensor], list[torch.Tensor]]:
        """Filter Q/K pairs to those with matching shape and device."""
        if not q_list or not k_list:
            return [], [], []
        q0, k0 = q_list[0], k_list[0]
        triples = [
            (idx, q, k)
            for idx, q, k in zip(tok_idx, q_list, k_list)
            if _sanity_check(q, q0) and _sanity_check(k, k0)
        ]
        if not triples:
            return [], [], []
        return (
            [t[0] for t in triples],
            [t[1] for t in triples],
            [t[2] for t in triples],
        )

@dataclass
class CaptureConfig:
    """Configuration for attention capture."""

    enabled: bool = False
    layers: set[int] | None = None