# SPDX-License-Identifier: Apache-2.0 # Format-dispatched geometry now lives in the ``kv_format`` package: each # ``EngineKVFormat`` has a :class:`KVFormatSpec` (geometry accessors) and # detection is split per engine in ``kv_format.detection``. The public # functions below are a thin, backwards-compatible facade that delegates # to ``get_spec`` / ``get_spec_class`` / ``detect_format`` so existing # callers keep working unchanged. # mypy: disable-error-code="union-attr,call-overload" # Standard from collections.abc import Hashable, Sequence from typing import TYPE_CHECKING, Optional, Union # Third Party import torch # First Party from lmcache.logging import init_logger from lmcache.python_ops_fallback import set_shape_desc_dtype from lmcache.utils import EngineType, lmcache_deprecate from lmcache.v1.config import LMCacheEngineConfig from lmcache.v1.gpu_connector.kv_format import ( concrete_shape, describe_shape, detect_format, get_spec, get_spec_class, ) from lmcache.v1.gpu_connector.kv_format.types import DiscoverableKVCache, LayoutHints if TYPE_CHECKING: # First Party from lmcache.v1.gpu_connector.gpu_connectors import GPUConnectorInterface # First Party import lmcache.c_ops as lmc_ops logger = init_logger(__name__) def assert_contiguous(tensor: torch.Tensor) -> None: """Assert that *tensor* has a contiguous physical layout with zero offset. LMCache transfer kernels assume logical and physical views match for coalesced memory accesses. Used at boundaries where we receive a tensor we can't or shouldn't permute (e.g. raw CUDA-IPC reconstruction in :class:`~lmcache.v1.platform.cuda.ipc_wrapper.RawCudaIPCWrapper`). Raises: ValueError: If *tensor* has a nonzero storage offset, or is non-contiguous. """ if tensor.storage_offset() != 0: raise ValueError(f"expected storage_offset 0, got {tensor.storage_offset()}") if not tensor.is_contiguous(): raise ValueError("tensor is not contiguous") def need_gpu_interm_buffer(lmcache_config: LMCacheEngineConfig): """ Check if the GPU Connector needs to create an intermediate buffer on the GPU """ if lmcache_config.enable_pd: return False else: return True def assert_layerwise_gpu_connector(gpu_connector: "GPUConnectorInterface"): """ Assert that a GPU Connector is a layerwise connector. """ # Import at runtime to avoid circular dependency # First Party from lmcache.v1.gpu_connector import gpu_connectors, xpu_connectors valid_connectors = ( gpu_connectors.VLLMPagedMemLayerwiseGPUConnector, gpu_connectors.VLLMBufferLayerwiseGPUConnector, gpu_connectors.SGLangLayerwiseGPUConnector, xpu_connectors.VLLMPagedMemLayerwiseXPUConnector, xpu_connectors.VLLMBufferLayerwiseXPUConnector, xpu_connectors.SGLangLayerwiseXPUConnector, ) assert isinstance(gpu_connector, valid_connectors) def is_mla(engine_kv_format: "lmc_ops.EngineKVFormat") -> bool: """Return ``True`` for a Multi-head Latent Attention (MLA) layout.""" return lmc_ops.is_mla(engine_kv_format) def get_engine_kv_shape_description(engine_kv_format: "lmc_ops.EngineKVFormat") -> str: """Return a human-readable symbolic shape legend for the Engine KV format. Uses short names matching the ``EngineKVFormat`` enum convention: NB=num_blocks, NL=num_layers, BS=block_size, NH=num_heads, HS=head_size, PBS=page_buffer_size (NB*BS). """ try: return describe_shape(engine_kv_format) except KeyError: return f"Unknown ({engine_kv_format})" def get_attention_backend(engine_kv_format: "lmc_ops.EngineKVFormat") -> str: """Return a representative attention-backend label for the format. Diagnostic only. A format may be produced by several (engine, attention-backend) combinations; the spec lists them in ``attention_backends`` and this returns the first (canonical) one. """ try: backends = get_spec_class(engine_kv_format).attention_backends except ValueError: return f"Unknown ({engine_kv_format})" return backends[0] if backends else f"Unknown ({engine_kv_format})" def get_concrete_engine_kv_shape( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat" ) -> str: """Return the shape with actual numeric values substituted. For example, instead of ``NL x [2, NB, BS, NH, HS]`` this returns ``80 x [2, 2048, 128, 8, 128]``. """ try: return get_spec(kv_caches, engine_kv_format).concrete_shape_str() except ValueError: return f"Unknown ({engine_kv_format})" def get_concrete_engine_kv_shape_from_shape_desc( shape_desc: "lmc_ops.PageBufferShapeDesc", engine_kv_format: "lmc_ops.EngineKVFormat", ) -> str: """Return the concrete shape for a single kernel group's ``shape_desc``. Like :func:`get_concrete_engine_kv_shape`, but the numeric values are read from a per-group :class:`PageBufferShapeDesc` rather than from the whole ``kv_caches`` structure. This makes the result *group-accurate*: ``shape_desc.nl`` is the layer count of the group (not the model total), so each kernel group of a hybrid model reports its own shape. For example, instead of ``NL x [2, NB, BS, NH, HS]`` this returns ``80 x [2, 2048, 128, 8, 128]``. Args: shape_desc: The kernel group's shape descriptor. Sizes are read from its ``nl``/``nb``/``bs``/``nh``/``hs`` fields; fused page-buffer (``NBBS``) formats use ``nb * bs``. engine_kv_format: The format whose member name is the shape template. Returns: The shape string with numeric values substituted, or ``"Unknown ()"`` for an unrecognised format. """ sizes = { "NB": shape_desc.nb, "NL": shape_desc.nl, "BS": shape_desc.bs, "NH": shape_desc.nh, "HS": shape_desc.hs, "PBS": shape_desc.nb * shape_desc.bs, } try: return concrete_shape(engine_kv_format, lambda label: sizes[label]) except KeyError: return f"Unknown ({engine_kv_format})" def legible_print_engine_kv_format(engine_kv_format: "lmc_ops.EngineKVFormat"): """ Print the Engine KV Format in a legible way """ shape = get_engine_kv_shape_description(engine_kv_format) backend = get_attention_backend(engine_kv_format) if shape.startswith("Unknown"): logger.warning(f"Unknown Engine KV Format: {engine_kv_format}") else: logger.info("Engine KV Format: %s", shape) logger.info("Currently used by:\n - %s", backend) def normalize_kv_and_discover_format( kv_caches: DiscoverableKVCache, serving_engine: EngineType, layout_hints: "LayoutHints | None" = None, ) -> tuple["lmc_ops.EngineKVFormat", DiscoverableKVCache]: """Normalize ``kv_caches`` into canonical form and discover its Engine KV format. Thin wrapper over :func:`lmcache.v1.gpu_connector.kv_format.detect_format`; see that function for the full contract. Args: kv_caches: The KV cache tensors (possibly nested lists of tensors). serving_engine: Which serving engine produced the caches. layout_hints: See :class:`LayoutHints`. Returns: ``(engine_kv_format, normalized_kv_caches)``. Callers must use the returned tensor structure for subsequent operations -- it shares storage with the input but may be a permuted view. """ return detect_format(kv_caches, serving_engine, layout_hints) def normalize_and_discover_per_layer_formats( kv_caches: "DiscoverableKVCache", layer_index_groups: "Sequence[Sequence[int]]", serving_engine: EngineType, layout_hints: "LayoutHints | None" = None, ) -> "tuple[DiscoverableKVCache, list[lmc_ops.EngineKVFormat]]": """Normalize the KV caches and return one Engine KV format per layer. Reports each layer's own format, so models whose layers do not all share one format -- e.g. a K+V main cache (``kv_size=2``) alongside a key-only MLA index cache (``kv_size=1``) -- get a correct per-layer format rather than a single model-wide one. Args: kv_caches: The registered KV caches: a per-layer list, or a single fused tensor for cross-layer formats. layer_index_groups: Layer indices of each engine group (one inner sequence per group). Empty means a single non-hybrid group. serving_engine: Which serving engine produced the caches. layout_hints: See :class:`LayoutHints`. Returns: ``(normalized_kv_caches, engine_kv_formats)``: the canonical KV cache structure and one format per layer (length equals the layer count), ready for :func:`lmcache.v1.kv_layer_groups.group_layers_by_identity`. """ # Detect the whole structure once. A format that isn't a per-layer list (a # cross-layer tensor, or a K/V-split) is single-format -- return it whole. whole_format, whole_normalized = detect_format( kv_caches, serving_engine, layout_hints ) if not lmc_ops.is_layer_list(whole_format): return whole_normalized, [whole_format] * get_num_layers( whole_normalized, whole_format ) # Per-layer list: re-detect per engine group, split by tensor shape so a group # that mixes layouts gets the right format per layer. groups = layer_index_groups or [range(len(kv_caches))] detected: dict[int, tuple[DiscoverableKVCache, "lmc_ops.EngineKVFormat"]] = {} for indices in groups: layers_by_shape: dict[Hashable, list[int]] = {} for i in indices: shape = getattr(kv_caches[i], "shape", None) key = tuple(shape) if shape is not None else None layers_by_shape.setdefault(key, []).append(i) for same_shape_indices in layers_by_shape.values(): fmt, normalized = detect_format( [kv_caches[i] for i in same_shape_indices], serving_engine, layout_hints, ) for sub_idx, layer_idx in enumerate(same_shape_indices): detected[layer_idx] = (normalized[sub_idx], fmt) # A layer in no group (cross-layer KV sharing) keeps its own tensor and is # skipped downstream; give it any detected format so every layer has one. fallback_format = next(fmt for _, fmt in detected.values()) normalized_per_layer = [ detected[i][0] if i in detected else kv_caches[i] for i in range(len(kv_caches)) ] engine_kv_formats = [ detected[i][1] if i in detected else fallback_format for i in range(len(kv_caches)) ] return normalized_per_layer, engine_kv_formats def get_num_layers( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat" ) -> int: """Return the number of layers from ``kv_caches``.""" return get_spec(kv_caches, engine_kv_format).num_layers() def get_num_blocks( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat" ) -> int: """Return the number of blocks from ``kv_caches``. Raises: ValueError: For NBBS-fused formats with no separate block axis. """ return get_spec(kv_caches, engine_kv_format).num_blocks() def get_block_size( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat", layer_idx: int = 0, ) -> int: """Return the block size (tokens per block) for layer ``layer_idx``. ``layer_idx`` is honoured only for per-layer formats where BS may differ across layers (e.g. mixed-compression MLA pools). For cross-layer formats BS is shared across layers and ``layer_idx`` is ignored. Raises: ValueError: For NBBS-fused formats with no separate block axis. """ return get_spec(kv_caches, engine_kv_format).block_size(layer_idx) @lmcache_deprecate( "page_buffer_size is only used by the legacy non-MP (in-process) connectors; " "the MP transfer path reads geometry from a per-group PageBufferShapeDesc instead" ) def get_page_buffer_size( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat" ) -> int: """Return the page buffer size (num_blocks * block_size) from ``kv_caches``.""" return get_spec(kv_caches, engine_kv_format).page_buffer_size() def get_num_heads( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat", layer_idx: int = 0, ) -> int: """Return the number of heads for a layer (defaults to layer 0).""" return get_spec(kv_caches, engine_kv_format).num_heads(layer_idx) def get_hidden_dim_size( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat", layer_idx: int = 0, ) -> int: """Return the hidden dimension for a layer (defaults to layer 0).""" return get_spec(kv_caches, engine_kv_format).hidden_dim(layer_idx) def get_head_size( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat", layer_idx: int = 0, ) -> int: """Return the head size for a layer (defaults to layer 0).""" return get_spec(kv_caches, engine_kv_format).head_size(layer_idx) def get_tokens_per_layer( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat" ) -> int: """Return the number of tokens per layer (num_blocks * block_size).""" return get_spec(kv_caches, engine_kv_format).tokens_per_layer() def get_elements_per_layer( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat" ) -> int: """Return the number of elements per layer (both K and V for non-MLA).""" return get_spec(kv_caches, engine_kv_format).elements_per_layer() def get_dtype( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat", layer_idx: int = 0, ) -> torch.dtype: """Return the dtype for a layer (defaults to layer 0).""" return get_spec(kv_caches, engine_kv_format).dtype(layer_idx) def get_group_data_ptrs( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat", layer_indices: list[int], ) -> list[int]: """Return device pointers for a group of layers in kernel-expected order. See :meth:`KVFormatSpec.data_ptrs` for the per-format pointer-array shape (per-layer list, SGLang K-then-V, or single cross-layer base). Args: kv_caches: Full kv_caches structure. engine_kv_format: Format returned by :func:`normalize_kv_and_discover_format`. layer_indices: 0-based layer indices in the group, in kernel order. Returns: Device pointers (int), in kernel-expected order. Raises: ValueError: If *engine_kv_format* is not recognized. """ return get_spec(kv_caches, engine_kv_format).data_ptrs(layer_indices) def assert_is_vllm_flash_attn_or_flash_infer( engine_kv_format: "lmc_ops.EngineKVFormat", ): """ Ensure that we have an Engine KV Cache Format that is either vLLM's flash attention or flash infer. """ assert engine_kv_format in ( lmc_ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS, lmc_ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS, lmc_ops.EngineKVFormat.NL_X_TWO_NB_NH_BS_HS, lmc_ops.EngineKVFormat.NL_X_NB_TWO_NH_BS_HS, # Blocks-first fused K/V (vLLM CPU): a per-layer non-MLA layout that # shares this transfer path even though it is not literally flash-*. lmc_ops.EngineKVFormat.NL_X_NB_NH_BS_TWO_HS, ) def assert_is_vllm_mla_or_flash_attn_or_flash_infer( engine_kv_format: "lmc_ops.EngineKVFormat", ) -> None: """ Ensure that we have an Engine KV Cache Format that is either vLLM's MLA, flash attention, or flash infer. Accepted formats: - ``NL_X_TWO_NB_BS_NH_HS`` (flash attention, NHD) - ``NL_X_NB_TWO_BS_NH_HS`` (flash infer, NHD) - ``NL_X_TWO_NB_NH_BS_HS`` (flash attention, HND) - ``NL_X_NB_TWO_NH_BS_HS`` (flash infer, HND) - ``NL_X_NB_BS_HS`` (MLA) Raises: AssertionError: If *engine_kv_format* is not one of the accepted formats. """ assert engine_kv_format in ( lmc_ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS, lmc_ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS, lmc_ops.EngineKVFormat.NL_X_TWO_NB_NH_BS_HS, lmc_ops.EngineKVFormat.NL_X_NB_TWO_NH_BS_HS, lmc_ops.EngineKVFormat.NL_X_NB_BS_HS, ) def get_device(kv_caches: DiscoverableKVCache) -> torch.device: """Return the device of the KV cache tensors. Descends into any list nesting until a tensor is found; assumes all tensors in *kv_caches* live on the same device (true for every current :class:`EngineKVFormat`). """ probe: DiscoverableKVCache = kv_caches while isinstance(probe, list): probe = probe[0] return probe.device # Formats whose per-layer tensor dim-0 is the *block* axis AND for # which we currently support dim-0 padding (e.g. DeepSeek V4 # compressor / indexer caches sharing a KV pool with larger attn # groups). Today only the MLA layout (``NL_X_NB_BS_HS``, kv_size==1) # is exercised by real mixed-compression workloads. # # ``NL_X_NB_TWO_BS_NH_HS`` *could* in principle also be the block # axis on dim-0, but no real serving engine emits a padded layout of # that format yet, and supporting it would require: (a) deciding # (without a ground-truth example) which axis carries the padding # -- NB boundary vs K<->V offset -- and (b) a coordinated change in # ``attempt_permute_to_contiguous_view`` to let interior-dim padding # through for that one format. Rather than ship an unverifiable code # path, we keep ``NL_X_NB_TWO_BS_NH_HS`` out of this set, which means # any padded tensor of that format will fail loudly via the # non-block-axis dim-0-padding check below. Revisit and add a # properly-tested branch when a concrete use case lands. _BLOCK_AXIS_FORMATS: frozenset = frozenset( { lmc_ops.EngineKVFormat.NL_X_NB_BS_HS, } ) def resolve_block_stride_and_log_layout( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat", layer_idx: int, group_idx: int, ) -> Optional[int]: """Resolve the per-block stride for a KV layer group and log its layout. Single entry point for :class:`KVLayerGroupsManager` to obtain the ``block_stride_elems`` value for :class:`PageBufferShapeDesc` and emit a one-shot layout audit line. All ``EngineKVFormat``-aware reasoning is kept here so callers never touch a "representative KV cache" tensor. * Block-axis formats (:data:`_BLOCK_AXIS_FORMATS`): ``stride(0)`` is the per-block step and is returned as-is. A value larger than the tight stride indicates dim-0 padding (e.g. DeepSeek V4 compressor caches sharing a KV pool with larger attn groups). * Other formats: dim-0 is not the block axis, so ``None`` is returned and ``shape_desc`` falls back to the tight stride. Any dim-0 padding in such formats is rejected with ``ValueError`` since downstream kernels cannot honour it. Args: kv_caches: Full KV cache structure (already normalised). engine_kv_format: Format of ``kv_caches``. layer_idx: 0-based layer index used as the layout probe. group_idx: 0-based group index, used only for logging. Returns: ``stride(0)`` for block-axis formats; ``None`` otherwise. Raises: ValueError: Non-block-axis format carries dim-0 padding. """ def _pick_layout_probe_tensor() -> torch.Tensor: # Layout probe only (shape/stride/storage_offset/dtype); not # the K/V slice fed to the transfer kernel. # - Cross-layer formats: ``kv_caches`` is the single backing # tensor packing all layers along dim-1; indexing dim-0 # (= NB) would yield a per-block slice, so return the whole # tensor (its ``stride(0)`` is the authoritative per-NB step). # - SGL MHA: outer list is K/V (length 2), inner list is # per-layer; K & V share shape/stride by construction. # - Other formats: ``kv_caches`` is already a per-layer list. if lmc_ops.is_cross_layer(engine_kv_format): if not isinstance(kv_caches, torch.Tensor): raise TypeError( "Cross-layer EngineKVFormat expects a single backing " f"torch.Tensor, got {type(kv_caches).__name__}." ) return kv_caches if lmc_ops.is_kv_list(engine_kv_format): return kv_caches[0][layer_idx] # type: ignore[index,return-value] return kv_caches[layer_idx] # type: ignore[index,return-value] rep = _pick_layout_probe_tensor() block_stride_elems: Optional[int] if engine_kv_format in _BLOCK_AXIS_FORMATS and rep.ndim > 0: block_stride_elems = int(rep.stride(0)) else: # Non-block-axis format: detect forbidden dim-0 padding. if rep.ndim >= 2: tight_dim0 = 1 for d in range(1, rep.ndim): tight_dim0 *= int(rep.shape[d]) padding = int(rep.stride(0)) - tight_dim0 if padding > 0: raise ValueError( "resolve_block_stride_and_log_layout: group's probe " f"tensor has dim-0 padding ({padding} elements per " f"block) but engine_kv_format={engine_kv_format!r} is not " "a supported dim-0-padded format (only " "NL_X_NB_BS_HS is); downstream transfer kernels " "cannot honour this padding and would read/write " "wrong bytes. " f"layer_idx={layer_idx}, shape={tuple(rep.shape)}, " f"stride={tuple(rep.stride())}, " f"tight_stride0={tight_dim0}, " f"storage_offset={int(rep.storage_offset())}, " f"dtype={rep.dtype}." ) block_stride_elems = None # Best-effort layout audit log; the log line itself must not raise. shape = tuple(rep.shape) stride = tuple(rep.stride()) try: inner = 1 for s in shape[1:]: inner *= int(s) padding_per_block = stride[0] - inner if stride else 0 except Exception: padding_per_block = -1 try: storage_nbytes = rep.untyped_storage().nbytes() except Exception: storage_nbytes = -1 logger.info( "Group %d first-layer tensor: layer_idx=%d shape=%s " "stride=%s is_contiguous=%s dtype=%s device=%s " "storage_offset=%d numel=%d storage_nbytes=%d " "padding_per_block=%d", group_idx, layer_idx, shape, stride, rep.is_contiguous(), rep.dtype, rep.device, rep.storage_offset(), rep.numel(), storage_nbytes, padding_per_block, ) return block_stride_elems def make_page_buffer_shape_desc( kv_caches: DiscoverableKVCache, engine_kv_format: "lmc_ops.EngineKVFormat", layer_idx: int, num_layers_in_group: int, num_blocks: int, block_size: int, block_stride_elems: Optional[int] = None, ) -> "lmc_ops.PageBufferShapeDesc": """Build a :class:`PageBufferShapeDesc` from a representative layer. Args: kv_caches: Full kv_caches structure. engine_kv_format: Format returned by :func:`normalize_kv_and_discover_format`. layer_idx: 0-based index of the representative layer. num_layers_in_group: Number of layers in the group (``nl``). num_blocks: Number of paged blocks (``nb``). block_size: Tokens per block (``bs``). block_stride_elems: Physical per-block stride in *elements* (= ``tensor.stride(0)`` of the representative layer). Pass the real value whenever the group's KV pool may be dim-0-padded (e.g. DeepSeek V4 compressor/indexer caches sharing a row width with a larger group in the same pool); otherwise downstream transfer kernels will skip into padding and corrupt data. Leave as ``None`` for unpadded pools -- the kernel's ``per_block_stride()`` fallback (block_stride_elems <= 0) will reconstruct the tight stride from ``kv_size`` and ``scalars_per_block`` itself, so we don't duplicate that arithmetic on the Python side. Returns: A populated ``PageBufferShapeDesc``. """ desc = lmc_ops.PageBufferShapeDesc() desc.kv_size = 1 if is_mla(engine_kv_format) else 2 desc.nl = num_layers_in_group desc.nb = num_blocks desc.bs = block_size desc.nh = ( 1 if is_mla(engine_kv_format) else get_num_heads(kv_caches, engine_kv_format, layer_idx) ) desc.hs = get_head_size(kv_caches, engine_kv_format, layer_idx) dtype = get_dtype(kv_caches, engine_kv_format, layer_idx) desc.element_size = dtype.itemsize # The C++ PageBufferShapeDesc has no ``dtype`` field, but the pure-Python # CPU fallback does -- and needs it to disambiguate float16 vs bfloat16 # (both have itemsize 2, so element_size alone is not enough). Best-effort. set_shape_desc_dtype(desc, dtype) resolved_stride = int(block_stride_elems) if block_stride_elems else 0 desc.block_stride_elems = resolved_stride return desc def _split_token2d_kv(token2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Accepts either: - [2, T, D] - [T, 2, D] Returns: - k_tok: [T, D] - v_tok: [T, D] """ if token2d.dim() != 3: raise ValueError(f"Expected token2d dim=3, got {token2d.shape}") if token2d.shape[0] == 2: # [2, T, D] return token2d[0], token2d[1] if token2d.shape[1] == 2: # [T, 2, D] return token2d[:, 0, :], token2d[:, 1, :] raise ValueError(f"Unrecognized token2d layout: {token2d.shape}") def _get_head_size_view( kv_cache_layer: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]], *, use_mla: bool, engine_kv_format: Optional["lmc_ops.EngineKVFormat"] = None, ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: """ Returns flattened views for index_copy/index_select. If engine_kv_format is provided, use it to interpret tensor layout explicitly. If not provided, fall back to current structural behavior: - MLA: expects Tensor [P, B, HS] - Non-MLA: expects either * Tensor [2, P, B, NH, HS] OR * (k, v) tuple each [P, B, NH, HS] (and also supports [P, 2, B, NH, HS] as a safe extension) """ # ------------------------- # MLA # ------------------------- if use_mla: if not isinstance(kv_cache_layer, torch.Tensor): raise ValueError("MLA expects kv_cache_layer as Tensor") if kv_cache_layer.dim() != 3: raise ValueError(f"MLA expects 3D [P,B,HS], got {kv_cache_layer.shape}") p, b, hs = kv_cache_layer.shape return kv_cache_layer.view(p * b, hs) # ------------------------- # non-MLA (K/V) # ------------------------- # If already provided (k, v) in canonical per-layer form, no format needed. if not isinstance(kv_cache_layer, torch.Tensor): k, v = kv_cache_layer if k.dim() != 4 or v.dim() != 4: raise ValueError(f"Expected (k,v) 4D [P,B,NH,HS], got {k.shape}, {v.shape}") p, b, nh, hs = k.shape if v.shape != (p, b, nh, hs): raise ValueError(f"k/v shape mismatch: {k.shape} vs {v.shape}") return k.view(p * b, nh * hs), v.view(p * b, nh * hs) t = kv_cache_layer if t.dim() != 5: raise ValueError(f"Expected 5D tensor for non-MLA, got {t.shape}") # If we have the format enum, decode explicitly. if engine_kv_format is not None: if engine_kv_format == lmc_ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS: # per-layer: [2, NB, BS, NH, HS] if t.shape[0] != 2: raise ValueError( f"{engine_kv_format} expects [2,NB,BS,NH,HS], got {t.shape}" ) k, v = t[0], t[1] # [NB,BS,NH,HS] elif engine_kv_format == lmc_ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS: # per-layer: [NB, 2, BS, NH, HS] if t.shape[1] != 2: raise ValueError( f"{engine_kv_format} expects [NB,2,BS,NH,HS], got {t.shape}" ) k, v = t[:, 0], t[:, 1] # [NB,BS,NH,HS] else: # Other formats are either MLA-only or require upstream normalization. raise NotImplementedError( f"engine_kv_format={engine_kv_format} not supported in non-MLA " "path here. Normalize to (k,v) tuple [NB,BS,NH,HS] per-layer " "before calling." ) else: # No enum available: Assumed [2,P,B,H,D] (or [2,NB,BS,NH,HS] per-layer). # Also accept [P,2,B,H,D] (or [NB,2,BS,NH,HS]) to be more robust. if t.shape[0] == 2: k, v = t[0], t[1] elif t.shape[1] == 2: k, v = t[:, 0], t[:, 1] else: raise ValueError( f"engine_kv_format is None and tensor does not look like stacked KV. " f"Expected axis0==2 or axis1==2, got {t.shape}" ) if k.dim() != 4 or v.dim() != 4: raise ValueError(f"Expected k/v 4D [NB,BS,NH,HS], got {k.shape}, {v.shape}") nb, bs, nh, hs = k.shape if v.shape != (nb, bs, nh, hs): raise ValueError(f"k/v shape mismatch after decode: {k.shape} vs {v.shape}") return k.view(nb * bs, nh * hs), v.view(nb * bs, nh * hs)