diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py index 52a7a9daf028..de3a90ac105f 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py @@ -1,2 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .cache_manager import DeepseekV4CacheManager +from .deepseek_v4 import DeepseekV4AttentionType + +__all__ = ["DeepseekV4AttentionType", "DeepseekV4CacheManager"] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py new file mode 100644 index 000000000000..bf3cff3da7b9 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -0,0 +1,880 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import defaultdict +from typing import Dict, List, Optional, Tuple + +import torch + +from tensorrt_llm._torch.pyexecutor import llm_request +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import GPU_LEVEL, KVCacheManagerV2, Role +from tensorrt_llm._utils import ( + TensorWrapper, + convert_to_torch_tensor, + get_size_in_bytes, + prefer_pinned, +) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime import ModelConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BatchDesc, + BufferConfig, + GpuCacheTierConfig, + HostCacheTierConfig, + KVCacheDesc, + LayerId, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy + +from .compressor import KVCacheDtype +from .deepseek_v4 import ( + DEEPSEEK_V4_SPARSE_RATIO, + DeepseekV4AttentionType, + compress_ratio_has_attention, + get_attn_dim, + get_token_bytes, + is_compress_layer, + is_overlap_compressor, + is_sparse_layer, +) + + +def _estimate_bytes_per_token( + head_dim: int, + index_head_dim: int, + compress_ratios: List[int], + has_fp8_kv_cache, + attn_types: set[DeepseekV4AttentionType] | None = None, + indexer_k_dtype: str = "fp8", +) -> int: + total_bytes = 0 + for ratio in compress_ratios: + for attn in DeepseekV4AttentionType: + if attn_types is not None and attn not in attn_types: + continue + if compress_ratio_has_attention(ratio, attn): + total_bytes += _get_attn_bytes_per_token( + head_dim, + index_head_dim, + ratio, + attn, + has_fp8_kv_cache, + indexer_k_dtype=indexer_k_dtype, + ) + return total_bytes + + +def _get_attn_bytes_per_token( + head_dim: int, + index_head_dim: int, + compress_ratio: int, + attn_type: DeepseekV4AttentionType, + has_fp8_kv_cache: bool, + indexer_k_dtype: str = "fp8", +) -> int: + token_bytes = get_token_bytes( + head_dim, + index_head_dim, + compress_ratio, + attn_type, + has_fp8_kv_cache, + indexer_k_dtype=indexer_k_dtype, + ) + if attn_type in [DeepseekV4AttentionType.COMPRESS, DeepseekV4AttentionType.INDEXER_COMPRESS]: + token_bytes //= compress_ratio + return token_bytes + + +class DeepseekV4CacheManager(KVCacheManagerV2): + fixed_size_attention = { + DeepseekV4AttentionType.SWA, + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + } + # This tensor is for compatibility with AttentionOp, it only contains swa attention. + # kv_cache_pool_pointers contains pool pointers swa pool, shape: [1, 2] + # It assume the KVCacheManagerPy has only one pool for swa attention. + # The second column is always 0. + kv_cache_pool_pointers: torch.Tensor + # This tensor is for compatibility with AttentionOp, it only contains swa attention. + # kv_cache_pool_mapping contains pool id and layer offset for each layer's swa attention, + # shape: [num_local_layers, 2] + kv_cache_pool_mapping: torch.Tensor + # The block size of the (indexer) compressed cache. + # For other attention types, block size is tokens_per_block. + compressed_block_sizes: List[int] + + def __init__( + self, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: int = 1, + max_batch_size: int, + max_beam_width: int = 1, + tokens_per_block: int, + max_seq_len: int, + vocab_size: int, + mapping: Mapping, + dtype: DataType = DataType.BF16, + compressor_dtype: DataType = DataType.FLOAT, + sparse_attn_config: DeepSeekV4SparseAttentionConfig, + max_input_len: Optional[int] = None, + max_num_tokens: Optional[int] = None, + **kwargs, + ) -> None: + # DeepSeek-V4 specific attributes initialization + assert kv_cache_type == CacheTypeCpp.SELFKONLY, "DeepSeek-V4 only supports SELFKONLY" + assert num_kv_heads == 1, "DeepSeek-V4 only supports num_kv_heads == 1" + assert len(sparse_attn_config.compress_ratios) >= num_layers, ( + "The length of compress ratios must be >= the number of layers" + ) + assert dtype in [DataType.BF16, DataType.FP8], ( + f"Unsupported dtype: {dtype}, only support BF16 and FP8" + ) + assert compressor_dtype == DataType.FLOAT, ( + f"Unsupported compressor dtype: {compressor_dtype}, only support FP32/TF32" + ) + + assert tokens_per_block in [128, 256], ( + f"DeepseekV4CacheManager requires tokens_per_block in [128, 256], got {tokens_per_block}. " + f"Set kv_cache_config.tokens_per_block to 128 or 256." + ) + + self.index_head_dim = sparse_attn_config.index_head_dim + self._compress_ratios = sparse_attn_config.compress_ratios + # When MTP is enabled, enlarge the sliding window sizes by + # max_draft_len so that rewinding rejected draft tokens can still + # reach the KV entries that would otherwise have been evicted by the + # sliding-window policy. + spec_config = kwargs.get("spec_config", None) + self._max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 + self._swa_window_size = sparse_attn_config.window_size + self._compressor_dtype = compressor_dtype + # If MTP is enabled, append compress ratios for MTP virtual layers. + # MTP adds (max_draft_len - 1) extra layers that mirror the last real + # layer's attention pattern. Only NEW entries are appended; existing + # per-layer ratios are never modified, so a real layer with ratio==1 + # stays SWA-only. + if self._max_draft_len > 0: + self._compress_ratios = self._compress_ratios + [self._compress_ratios[-1]] * ( + self._max_draft_len - 1 + ) + self.compressed_block_sizes = [tokens_per_block // ratio for ratio in self._compress_ratios] + + self._init_indexer_dtype(sparse_attn_config) + + # _build_cache_config() needs them to build constraints + self._max_input_len = max_input_len + self._max_num_tokens = max_num_tokens + + # General initialization + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + max_batch_size=max_batch_size, + max_beam_width=max_beam_width, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + vocab_size=vocab_size, + mapping=mapping, + dtype=dtype, + **kwargs, + ) + self.is_vswa = True # DeepSeek-V4 must has VSWA + + # DeepSeek-V4 expects cache of all layers with the same attention type and compress ratio + # to be in the same pool and have the same scale. + self._assert_layer_pool_scale() + + # For DeepSeek-V4 Attention, the base pointer for SWA pool + # Use first PP layer instead of hardcoded 0 for pipeline parallelism. + first_pp_layer = self.pp_layers[0] + self.swa_pool_ptr = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[first_pp_layer, DeepseekV4AttentionType.SWA], Role.KEY + ) + + self.compress_pool_ptrs = {} + # Find first PP layer with each compress ratio for pool pointer lookup. + pp_compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] + if 4 in pp_compress_ratios: # indexer compressor + first_layer_with_4 = self.pp_layers[pp_compress_ratios.index(4)] + self.compress_pool_ptrs[4] = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[first_layer_with_4, DeepseekV4AttentionType.COMPRESS], + Role.KEY, + ) + if 128 in pp_compress_ratios: # compressor + first_layer_with_128 = self.pp_layers[pp_compress_ratios.index(128)] + self.compress_pool_ptrs[128] = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[ + first_layer_with_128, DeepseekV4AttentionType.COMPRESS + ], + Role.KEY, + ) + # Use pinned staging buffer to avoid pageable H2D memcpy + max_num_sequences = max_batch_size * mapping.pp_size + self._host_block_offsets_staging = torch.empty( + (max_num_sequences + 1) * max_beam_width, + 2, # key and value + self.max_blocks_per_seq, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + + def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> torch.Tensor: + """ + Get the buffers for a specific layer and attention type. + + Args: + layer_idx: The layer index + attn_type: The attention type + + Returns: + The buffer tensor (shape: [num_blocks, tokens_per_block, attn_dim]) + For blockwise FP8 layers, shape is [num_blocks, tokens_per_block, attn_dim + scale_size] + """ + layer_id = self._layer_attn_to_layer_id[(layer_idx, attn_type)] + addr = self.impl.get_mem_pool_base_address(layer_id, Role.KEY) + + block_size = self.tokens_per_block + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + block_size = self.compressed_block_sizes[layer_idx] + + attn_dim = get_attn_dim( + self.head_dim, self.index_head_dim, self._compress_ratios[layer_idx], attn_type + ) + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + # Indexer always pack data + per-block scales into the same row. + dim_per_token = self._indexer_data_size + self._indexer_scale_size + else: + dim_per_token = attn_dim + + shape = ( + self.impl.get_page_index_upper_bound(layer_id, Role.KEY), + block_size, + dim_per_token, + ) + + dtype = self.dtype + # (indexer) compressor state and score use compressor_dtype + if attn_type in [ + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ]: + dtype = self._compressor_dtype + elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + dtype = self._indexer_dtype + + return convert_to_torch_tensor(TensorWrapper(addr, dtype, shape)) + + def _get_window_size(self, compress_ratio: int, attn_type: DeepseekV4AttentionType) -> int: + if attn_type == DeepseekV4AttentionType.SWA: + base_window_size = self._swa_window_size + elif attn_type in self.fixed_size_attention: + state_factor = 2 if is_overlap_compressor(compress_ratio) else 1 + base_window_size = state_factor * compress_ratio + else: + raise ValueError(f"Unsupported fixed-size attention type: {attn_type}") + return base_window_size + self._max_draft_len + + def _build_pool_mapping_tensors(self) -> Tuple[torch.Tensor, torch.Tensor]: + first_pp_layer = self.pp_layers[0] + swa_bytes_per_block = self._get_attn_bytes_per_block( + DeepseekV4AttentionType.SWA, first_pp_layer + ) + swa_pool_ptr = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[first_pp_layer, DeepseekV4AttentionType.SWA], Role.KEY + ) + + def _get_layer_offset(pp_layer: int) -> int: + buffer_ptr = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[pp_layer, DeepseekV4AttentionType.SWA], Role.KEY + ) + return (buffer_ptr - swa_pool_ptr) // swa_bytes_per_block + + # Tensors for compatibility with AttentionOp, only contains swa attention. + # Assume the SWA of all layers share the same pool. + # shape: [1, 2] + kv_cache_pool_pointers = torch.tensor( + [[swa_pool_ptr, 0]], + dtype=torch.int64, + device="cpu", + pin_memory=prefer_pinned(), + ) + # shape: [num_local_layers, 2] + kv_cache_pool_mapping = torch.tensor( + [[0, _get_layer_offset(pp_layer)] for pp_layer in self.pp_layers], + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + return kv_cache_pool_pointers, kv_cache_pool_mapping + + def get_cache_indices( + self, + request_id: int, + layer_idx: int, + attn_type: DeepseekV4AttentionType, + ) -> List[int]: + """ + Get the cache block indices for a batch of requests at a specific layer and attention type. + + Args: + request_id: The request id + layer_idx: The layer index + attn_type: The attention type + + Returns: + The cache block indices, shape (max_blocks_per_seq,) + """ + layer_id = self._layer_attn_to_layer_id[(layer_idx, attn_type)] + pool_id = self.layer_to_pool_mapping_dict[layer_id] + base_indices = self.kv_cache_map[request_id].get_base_page_indices(pool_id).tolist() + converter = self.impl.get_page_index_converter(layer_id, Role.KEY) + return converter(base_indices) + + def _get_cache_quota(self, max_tokens: int) -> int: + quota = int(max_tokens * self.get_cache_bytes_per_token()) + # Add extra quota to ensure sufficient space for small max_tokens cases. + quota += len(DeepseekV4AttentionType) * (2 << 20) + return quota + + def _build_cache_config( + self, + kv_cache_config: KvCacheConfig, + *, + tokens_per_block: int, + vocab_size: int | None, + cache_tiers: List[GpuCacheTierConfig | HostCacheTierConfig], + ) -> KVCacheManagerConfigPy: + """ + Create the cache manager config for DeepSeek-V4. + """ + layers: List[AttentionLayerConfig] = [] + layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {} + + def _add_layer( + layer_idx: int, attn_type: DeepseekV4AttentionType, sliding_window_size: int | None + ): + nonlocal layers, layer_attn_to_layer_id + layer_id = LayerId(len(layers)) + # update the mapping from layer index and attention type to layer id + layer_attn_to_layer_id[layer_idx, attn_type] = layer_id + # add the layer to the layers list + layer_config = AttentionLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig( + role=Role.KEY, size=self._get_attn_bytes_per_block(attn_type, layer_idx) + ) + ], + sliding_window_size=sliding_window_size, + num_sink_tokens=None, + ) + layers.append(layer_config) + + # create the layer config for DeepSeek-V4 + for layer in self.pp_layers: + compress_ratio = self._compress_ratios[layer] + is_compress = is_compress_layer(compress_ratio) + is_sparse = is_sparse_layer(compress_ratio) + + # sliding window attention pool + _add_layer( + layer, + DeepseekV4AttentionType.SWA, + self._get_window_size(compress_ratio, DeepseekV4AttentionType.SWA), + ) + + if is_compress: + # compressed attention pool + _add_layer(layer, DeepseekV4AttentionType.COMPRESS, None) + # compressor state, managed as a sliding window attention cache, + # including compressor kv states and compressor score states. + # Add max_draft_len so rewind after rejected draft tokens can + # still reach past states. + compressor_window = self._get_window_size( + compress_ratio, DeepseekV4AttentionType.COMPRESSOR_STATE + ) + _add_layer(layer, DeepseekV4AttentionType.COMPRESSOR_STATE, compressor_window) + _add_layer(layer, DeepseekV4AttentionType.COMPRESSOR_SCORE, compressor_window) + + # sparse attention layer has indexer + if is_sparse: + # indexer kv cache pool, dim is indexer_head_dim + _add_layer(layer, DeepseekV4AttentionType.INDEXER_COMPRESS, None) + # indexer has its own compressor, so a separate compressor state + # similarly, indexer compressor state is managed as a sliding window attention cache + indexer_compressor_window = self._get_window_size( + compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE + ) + _add_layer( + layer, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + indexer_compressor_window, + ) + _add_layer( + layer, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + indexer_compressor_window, + ) + # the mapping from layer index and attention type to layer id + self._layer_attn_to_layer_id = layer_attn_to_layer_id + # number of layers in the KVCacheManagerPy + self._num_manager_layers = len(layers) + + # Build constraints and typical_step for better pool ratio. + max_batch_size = self.max_batch_size + max_seq_len = self.max_seq_len + max_num_tokens = self._max_num_tokens + max_draft_len = self._max_draft_len + + # For aggregated serving in large batch size: + # Use 1 context request + (max_batch_size - 1) generation requests as + # the typical step. An all-generation typical_step over-provisions the + # compressed-cache pool at the expense of the SWA pool, starving the + # SWA pool and artificially capping the achievable batch size. + ctx_capacity = max_num_tokens if max_num_tokens is not None else max_seq_len + typical_step = BatchDesc( + kv_caches=[ + KVCacheDesc(capacity=ctx_capacity, history_length=0), + ] + + [KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - max_draft_len - 1)] + * (max_batch_size - 1), + ) + + constraints = [] + # Constraint 1: cuda graph generation warmup — one decode request that has + # accumulated to the tail of max_seq_len. Using history_length=max_seq_len-1 + # (instead of 0) lets SWA / SSM pools collapse to their windowed working set, + # while full-cache pools still need max_seq_len/tokens_per_block blocks + # because they don't age. + constraints.append( + BatchDesc([KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1)]) + ) + + # Constraint 2: general / chunked-prefill warmup — one fresh context request + # at max_num_tokens (the per-iteration token budget). + if max_num_tokens is not None: + constraints.append(BatchDesc([KVCacheDesc(capacity=max_num_tokens, history_length=0)])) + + return KVCacheManagerConfigPy( + tokens_per_block=tokens_per_block, + vocab_size=vocab_size, + cache_tiers=cache_tiers, + max_util_for_resume=kv_cache_config.max_util_for_resume, + layers=layers, + typical_step=typical_step, + constraints=constraints, + enable_stats=self.enable_stats, + ) + + def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None: + # Indexer compressor cache layout. Two modes are supported: + # - "fp8" (FP8 blockwise): 1 byte per value + 1 fp32 scale per 128 + # values. + # - "fp4" (MXFP4 blockwise): ½ byte per value (two FP4 codes packed + # per byte) + 1 ue8m0 byte per 32 values. At index_head_dim=128 + # this halves the per-token indexer-K footprint vs FP8. + self._indexer_k_dtype = sparse_attn_config.indexer_k_dtype + if self._indexer_k_dtype == "fp8": + self._indexer_cache_dtype = KVCacheDtype.FP8_BLOCKWISE + self._indexer_dtype = DataType.FP8 + self.quant_block_size = 128 + self._indexer_data_size = self.index_head_dim + self._indexer_scale_size = get_size_in_bytes( + self.index_head_dim // self.quant_block_size, DataType.FLOAT + ) + elif self._indexer_k_dtype == "fp4": + assert self.index_head_dim == 128, ( + f"FP4 indexer K cache requires index_head_dim=128, got {self.index_head_dim}." + ) + self._indexer_cache_dtype = KVCacheDtype.MXFP4_BLOCKWISE + # Pool dtype is uint8 because PyTorch can't allocate float4 + # backing storage; downstream consumers reinterpret these + # raw bytes as packed E2M1 + UE8M0 exponents. + self._indexer_dtype = DataType.UINT8 + self.quant_block_size = 32 + # Two E2M1 codes pack into one byte → half the data footprint. + self._indexer_data_size = self.index_head_dim // 2 + # 1 UE8M0 byte per 32-element block. + self._indexer_scale_size = self.index_head_dim // self.quant_block_size + else: + raise ValueError( + f"Unsupported indexer_k_dtype " + f"{sparse_attn_config.indexer_k_dtype!r}; expected " + "'fp8' or 'fp4'." + ) + # FP4 indexer flag mirrors `DSACacheManager.use_fp4` so the shared + # base Indexer can branch without knowing the V4-specific enum. + self.use_fp4 = self._indexer_k_dtype == "fp4" + assert self.index_head_dim % self.quant_block_size == 0, ( + f"indexer_head_dim {self.index_head_dim} must be divisible by {self.quant_block_size}" + ) + + def _assert_layer_pool_scale(self) -> None: + attn_ratio_to_pool_id = defaultdict[DeepseekV4AttentionType, dict[int, int]](lambda: {}) + attn_ratio_to_scale = defaultdict[DeepseekV4AttentionType, dict[int, int]](lambda: {}) + + comb = [ + (attn_type, layer_idx) + for attn_type in DeepseekV4AttentionType + for layer_idx in self.pp_layers + if compress_ratio_has_attention(self._compress_ratios[layer_idx], attn_type) + ] + for attn_type, layer_idx in comb: + compress_ratio = self._compress_ratios[layer_idx] + layer_id = self._layer_attn_to_layer_id[layer_idx, attn_type] + pool_id = self.layer_to_pool_mapping_dict[layer_id] + converter = self.impl.get_page_index_converter(layer_id, Role.KEY) + scale = converter.scale + + # check if the pool id is consistent + if compress_ratio in attn_ratio_to_pool_id[attn_type]: + other_pool_id = attn_ratio_to_pool_id[attn_type][compress_ratio] + assert other_pool_id == pool_id, ( + f"Layer {layer_idx} with compress ratio {compress_ratio}, " + f"its attention type {attn_type.name} has pool id {pool_id}, " + f"but another layer with the same compress ratio and attention type has pool id {other_pool_id}." + "DeepSeek-V4 expects they share the same pool." + ) + else: + attn_ratio_to_pool_id[attn_type][compress_ratio] = pool_id + + # check if the scale is consistent + if compress_ratio in attn_ratio_to_scale[attn_type]: + other_scale = attn_ratio_to_scale[attn_type][compress_ratio] + assert other_scale == scale, ( + f"Layer {layer_idx} with compress ratio {compress_ratio}, " + f"its attention type {attn_type.name} has scale {scale}, " + f"but another layer with the same compress ratio and attention type has scale {other_scale}." + "DeepSeek-V4 expects they share the same scale." + ) + else: + attn_ratio_to_scale[attn_type][compress_ratio] = scale + + # check if all swa attentions are in the same pool and have the same scale + swa_pool_ids = set(attn_ratio_to_pool_id[DeepseekV4AttentionType.SWA].values()) + swa_scales = set(attn_ratio_to_scale[DeepseekV4AttentionType.SWA].values()) + assert len(swa_pool_ids) == 1, "All swa attentions must be in the same pool" + assert len(swa_scales) == 1, "All swa attentions must have the same scale" + + # Ensure all compress ratios have SWA entries, not just PP-local ones. + # The attention metadata uses compress_ratio=1 as a hardcoded SWA key, + # but with pipeline parallelism a PP stage may not have any layers with + # ratio 1. Since all SWA layers share the same pool, we populate entries + # for every compress ratio in the model. + swa_pool_id = next(iter(swa_pool_ids)) + swa_scale = next(iter(swa_scales)) + for ratio in set(self._compress_ratios): + if ratio not in attn_ratio_to_pool_id[DeepseekV4AttentionType.SWA]: + attn_ratio_to_pool_id[DeepseekV4AttentionType.SWA][ratio] = swa_pool_id + attn_ratio_to_scale[DeepseekV4AttentionType.SWA][ratio] = swa_scale + + self._attn_ratio_to_pool_id = attn_ratio_to_pool_id + self._attn_ratio_to_scale = attn_ratio_to_scale + + def _get_attn_bytes_per_block( + self, + attn_type: DeepseekV4AttentionType, + layer_idx: int, + ) -> int: + """ + Get the cache bytes per token for a specific attention type and layer. + """ + has_fp8_kv_cache = self.dtype == DataType.FP8 + token_bytes = get_token_bytes( + self.head_dim, + self.index_head_dim, + self._compress_ratios[layer_idx], + attn_type, + has_fp8_kv_cache, + indexer_k_dtype=self._indexer_k_dtype, + ) + + block_size = self.tokens_per_block + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + block_size = self.compressed_block_sizes[layer_idx] + + return token_bytes * block_size + + def get_cache_bytes_per_token(self) -> int: + """Get the average cache bytes per token for DeepSeek-V4.""" + has_fp8_kv_cache = self.dtype == DataType.FP8 + compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] + return _estimate_bytes_per_token( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + indexer_k_dtype=self._indexer_k_dtype, + ) + + def get_max_resource_count(self) -> int: + # Keep scheduler capacity tied to physical GPU KV quota in bytes. + return int(self.impl.get_quota(GPU_LEVEL)) + + def _is_context_request(self, request: llm_request.LlmRequest) -> bool: + if getattr(request, "is_context_init_state", False): + return True + return getattr(request, "state", None) == llm_request.LlmRequestState.CONTEXT_INIT + + def _is_generation_request(self, request: llm_request.LlmRequest) -> bool: + if ( + getattr(request, "is_generation_in_progress_state", False) + or getattr(request, "is_generation_to_complete_state", False) + or getattr(request, "is_disagg_generation_init_state", False) + ): + return True + return getattr(request, "state", None) in ( + llm_request.LlmRequestState.GENERATION_IN_PROGRESS, + llm_request.LlmRequestState.GENERATION_TO_COMPLETE, + ) + + def _get_context_bytes(self, request: llm_request.LlmRequest) -> int: + prompt_len = max(0, getattr(request, "prompt_len", request.orig_prompt_len)) + total_tokens = prompt_len + self.num_extra_kv_tokens + return total_tokens * self.get_cache_bytes_per_token() + + def _get_generation_bytes(self, request: llm_request.LlmRequest) -> int: + prompt_len = max(0, getattr(request, "prompt_len", request.orig_prompt_len)) + max_new_tokens = max(0, request.max_new_tokens) + total_tokens = prompt_len + max_new_tokens + self.num_extra_kv_tokens + has_fp8_kv_cache = self.dtype == DataType.FP8 + total_bytes = 0 + for layer in self.pp_layers: + compress_ratio = self._compress_ratios[layer] + for attn_type in DeepseekV4AttentionType: + if not compress_ratio_has_attention(compress_ratio, attn_type): + continue + token_bytes = _get_attn_bytes_per_token( + self.head_dim, + self.index_head_dim, + compress_ratio, + attn_type, + has_fp8_kv_cache, + indexer_k_dtype=self._indexer_k_dtype, + ) + attn_tokens = total_tokens + if attn_type in self.fixed_size_attention: + attn_tokens = self._get_window_size(compress_ratio, attn_type) + total_bytes += attn_tokens * token_bytes + return total_bytes + + def get_needed_resource_to_completion(self, request: llm_request.LlmRequest) -> int: + if self._is_generation_request(request): + return self._get_generation_bytes(request) + if self._is_context_request(request): + return self._get_context_bytes(request) + raise ValueError(f"Unsupported request state: {request.state}") + + def get_layer_bytes_per_token( + self, + local_layer_idx: int, + data_role: Role, + ) -> int: + raise NotImplementedError( + "DeepSeek-V4 doesn't support get_layer_bytes_per_token, use _get_attn_bytes_per_block" + ) + + def get_indexer_k_cache_buffers(self, layer_idx: int) -> torch.Tensor: + """ + Get the buffers for the indexer k cache for a specific layer. + """ + buffer = self.get_buffers(layer_idx, DeepseekV4AttentionType.INDEXER_COMPRESS).unsqueeze(2) + return buffer.view(torch.uint8) + + def get_batch_indexer_k_cache_indices(self, request_ids: List[int]) -> List[List[int]]: + """ + Get the indices for the indexer k cache for a batch of requests. + """ + return self.get_batch_attn_offset( + request_ids, + # use beam_width=1 and num_contexts=0 since we don't support beam search + 1, + 0, + len(request_ids), + DeepseekV4AttentionType.INDEXER_COMPRESS, + DEEPSEEK_V4_SPARSE_RATIO, + ).tolist() + + def copy_batch_block_offsets( + self, + dst_tensor: torch.Tensor, + request_ids: List[int], + beam_width: int, + num_contexts: int, + num_seqs: int, + ) -> None: + """For compatibility with AttentionOp, copy only the SWA block offsets.""" + offsets = self.get_batch_attn_offset( + request_ids, + beam_width, + num_contexts, + num_seqs, + DeepseekV4AttentionType.SWA, + # all compress ratios have SWA attention and they are in the same pool + self._compress_ratios[self.pp_layers[0]], + ) + self._host_block_offsets_staging[:num_seqs, :, :] = offsets[:, None, :] + dst_tensor[0, :num_seqs, :, :].copy_( + self._host_block_offsets_staging[:num_seqs, :, :], non_blocking=True + ) + + def get_batch_attn_offset( + self, + request_ids: List[int], + beam_width: int, + num_contexts: int, + num_seqs: int, + attn_type: DeepseekV4AttentionType, + compress_ratio: int, + ) -> torch.Tensor: + """ + Get the block offsets for a specific attention type for a batch of requests. + + Args: + request_ids: The request ids + beam_width: The beam width + num_contexts: The number of context requests + num_seqs: The number of sequence requests + attn_type: The attention type + compress_ratio: The compress ratio. Used for non-SWA attention types. + + Returns: + The block offsets, shape (num_seqs, max_blocks_per_seq) + """ + assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" + assert attn_type == DeepseekV4AttentionType.SWA or compress_ratio is not None, ( + "compress_ratio must be provided for non-SWA attention types" + ) + + copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + assert copy_idx.shape[0] == num_seqs + + pool_id = self._attn_ratio_to_pool_id[attn_type][compress_ratio] + scale = self._attn_ratio_to_scale[attn_type][compress_ratio] + offsets = self.host_kv_cache_block_offsets[pool_id, copy_idx, 0] * scale + offsets[offsets == -scale] = -1 + return offsets + + def get_batch_block_offsets( + self, + request_ids: List[int], + num_contexts: int, + attention_type_set: set, + ) -> Dict[Tuple[int, "DeepseekV4AttentionType"], torch.Tensor]: + """Get block offsets for all attention types in a single call. + + Calls get_copy_index once and deduplicates offset computation by + (pool_id, scale) to avoid redundant work. + + Args: + request_ids: The request ids. + num_contexts: The number of context requests. + attention_type_set: Set of (compress_ratio, attention_type) tuples. + + Returns: + Dict mapping (compress_ratio, attention_type) -> offset tensor. + """ + copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, 1) + + offset_cache = {} # (pool_id, scale) -> offsets tensor + result = {} + for compress_ratio, attention_type in attention_type_set: + pool_id = self._attn_ratio_to_pool_id[attention_type][compress_ratio] + scale = self._attn_ratio_to_scale[attention_type][compress_ratio] + cache_key = (pool_id, scale) + if cache_key not in offset_cache: + offsets = self.host_kv_cache_block_offsets[pool_id, copy_idx, 0] * scale + offsets[offsets == -scale] = -1 + offset_cache[cache_key] = offsets + result[(compress_ratio, attention_type)] = offset_cache[cache_key] + return result + + @staticmethod + def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwargs): + config = model_config.pretrained_config + head_dim = config.kv_lora_rank + config.qk_rope_head_dim + index_head_dim = model_config.sparse_attention_config.index_head_dim + pp_layers = mapping.pp_layers(model_config.get_num_attention_layers()) + compress_ratios = [ + model_config.sparse_attention_config.compress_ratios[layer] for layer in pp_layers + ] + quant_config = model_config.quant_config + if quant_config is not None: + has_fp8_kv_cache = quant_config.quant_mode.has_fp8_kv_cache() + else: + has_fp8_kv_cache = False + indexer_k_dtype = model_config.sparse_attention_config.indexer_k_dtype + return _estimate_bytes_per_token( + head_dim, + index_head_dim, + compress_ratios, + has_fp8_kv_cache, + indexer_k_dtype=indexer_k_dtype, + ) + + def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: + some_checks_unavailable = False + has_invalid_values = torch.tensor( + [False], dtype=torch.bool, device=torch.cuda.current_device() + ) + pool_handled = set() + + # Handle each layer from start to end to traverse the whole KV cache. + for (layer, attn), layer_id in self._layer_attn_to_layer_id.items(): + pool_id = self.layer_to_pool_mapping_dict[layer_id] + if pool_id in pool_handled: + continue + buffer = self.get_buffers(layer, attn) + # process in chunks of 256 pages to avoid OoM + for i in range(0, buffer.shape[0], 256): + buffer_slice = buffer[i : i + 256] + try: + has_invalid_values.logical_or_(torch.isnan(buffer_slice).any()) + has_invalid_values.logical_or_(torch.isinf(buffer_slice).any()) + except NotImplementedError: + some_checks_unavailable = True + if fill_with_zero: + buffer.zero_() + pool_handled.add(pool_id) + torch.cuda.synchronize() + + if some_checks_unavailable: + logger.warning( + "`torch.isnan` or `torch.isinf` is not implemented for current kv cache dtype, " + "related checks are skipped" + ) + return bool(has_invalid_values) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index dca52c504535..ed80c7958a52 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -1,21 +1,119 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. from enum import Enum +DEEPSEEK_V4_SPARSE_RATIO = 4 +DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO = 4 + class DeepseekV4AttentionType(Enum): - # Attention types backed by per-layer sliding-window cache state. SWA = 0 - COMPRESSOR_KV = 1 - COMPRESSOR_SCORE = 2 - INDEXER_COMPRESSOR_KV = 3 - INDEXER_COMPRESSOR_SCORE = 4 - - # Attention types backed by ratio-shared compressed cache state. - COMPRESS = 5 - INDEXER_COMPRESS = 6 - - # Backward-compatible names used by the standalone compressor primitive. - COMPRESSOR_STATE = COMPRESSOR_KV - INDEXER_COMPRESSOR_STATE = INDEXER_COMPRESSOR_KV + COMPRESS = 1 + COMPRESSOR_STATE = 2 + COMPRESSOR_SCORE = 3 + INDEXER_COMPRESS = 4 + INDEXER_COMPRESSOR_STATE = 5 + INDEXER_COMPRESSOR_SCORE = 6 + + +def is_overlap_compressor(compress_ratio: int) -> bool: + return compress_ratio == DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO + + +def is_sparse_layer(compress_ratio: int) -> bool: + return compress_ratio == DEEPSEEK_V4_SPARSE_RATIO + + +def is_compress_layer(compress_ratio: int) -> bool: + return compress_ratio > 1 + + +def compress_ratio_has_attention(compress_ratio: int, attn_type: DeepseekV4AttentionType) -> bool: + is_sparse = is_sparse_layer(compress_ratio) + is_compress = is_compress_layer(compress_ratio) + + if attn_type == DeepseekV4AttentionType.SWA: + return True + if attn_type == DeepseekV4AttentionType.COMPRESS: + return is_compress + if attn_type == DeepseekV4AttentionType.COMPRESSOR_STATE: + return is_compress + if attn_type == DeepseekV4AttentionType.COMPRESSOR_SCORE: + return is_compress + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + return is_sparse + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE: + return is_sparse + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: + return is_sparse + raise ValueError(f"Unsupported DeepSeek-V4 attention type: {attn_type}") + + +def get_attn_dim( + head_dim: int, index_head_dim: int, compress_ratio: int, attn_type: DeepseekV4AttentionType +) -> int: + state_factor = 2 if is_overlap_compressor(compress_ratio) else 1 + if attn_type == DeepseekV4AttentionType.SWA: + return head_dim + if attn_type == DeepseekV4AttentionType.COMPRESS: + return head_dim + if attn_type == DeepseekV4AttentionType.COMPRESSOR_STATE: + return state_factor * head_dim + if attn_type == DeepseekV4AttentionType.COMPRESSOR_SCORE: + return state_factor * head_dim + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + return index_head_dim + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE: + return state_factor * index_head_dim + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: + return state_factor * index_head_dim + raise ValueError(f"Unsupported DeepSeek-V4 attention type: {attn_type}") + + +def get_token_bytes( + head_dim: int, + index_head_dim: int, + compress_ratio: int, + attn_type: DeepseekV4AttentionType, + has_fp8_kv_cache: bool, + indexer_k_dtype: str = "fp8", +) -> int: + if not compress_ratio_has_attention(compress_ratio, attn_type): + raise ValueError( + f"Layer with compress ratio {compress_ratio} does not have attention type {attn_type}" + ) + + attn_dim = get_attn_dim(head_dim, index_head_dim, compress_ratio, attn_type) + + dtype_bytes = 1 if has_fp8_kv_cache else 2 + if attn_type in [ + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ]: + dtype_bytes = 4 + + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + if indexer_k_dtype == "fp8": + return attn_dim + index_head_dim // 128 * 4 + if indexer_k_dtype == "fp4": + return index_head_dim // 2 + index_head_dim // 32 + raise ValueError( + f"Unsupported indexer_k_dtype {indexer_k_dtype!r}; expected 'fp8' or 'fp4'." + ) + + return attn_dim * dtype_bytes diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index d62c16f36355..78251d1ae978 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -10,7 +10,8 @@ CacheTransceiverConfig, CalibConfig, CapacitySchedulerPolicy, ContextChunkingPolicy, CudaGraphConfig, DecodeCudaGraphConfig, - DeepSeekSparseAttentionConfig, DFlashDecodingConfig, + DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, DFlashDecodingConfig, DraftTargetDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, ExtendedRuntimePerfKnobConfig, @@ -83,6 +84,7 @@ 'RocketSparseAttentionConfig', 'ReorderRequestPolicyConfig', 'DeepSeekSparseAttentionConfig', + 'DeepSeekV4SparseAttentionConfig', 'MiniMaxM3SparseAttentionConfig', 'SchedulingParams', 'SkipSoftmaxAttentionConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 73bcc1936192..1fabaefb1c2d 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -792,6 +792,53 @@ def _value(name: str, default=None): ) +class DeepSeekV4SparseAttentionConfig(DeepSeekSparseAttentionConfig): + """Configuration for DeepSeek-V4 Sparse Attention.""" + + algorithm: Literal["deepseek_v4"] = "deepseek_v4" + index_head_dim: Optional[int] = Field( + default=128, + description="The dimension of the DeepSeek-V4 indexer heads.") + skip_indexer_for_short_seqs: bool = Field( + default=False, + description= + "Whether to skip the MQA and Top-K in the indexer for short sequences.") + compress_ratios: List[int] = Field( + default_factory=lambda: [1, 1, 4, 128, 4, 128, 4], + description="The compress ratios of each layer. DeepSeek-V4 uses 0 " + "for uncompressed/SWA-only layers; the LLM API config normalizes " + "0 to 1, while checkpoint-facing semantics remain unchanged.") + window_size: int = Field( + default=128, + description="The sliding window size in tokens for SWA layers.") + index_topk: Optional[int] = Field(default=512, + description="The top-k for the indexer.") + + @field_validator("index_head_dim") + @classmethod + def validate_index_head_dim(cls, index_head_dim): + if index_head_dim is None: + raise ValueError( + "index_head_dim is required for DeepSeek-V4 sparse attention.") + return index_head_dim + + @field_validator("compress_ratios") + @classmethod + def normalize_compress_ratios(cls, compress_ratios): + if not compress_ratios: + raise ValueError("compress_ratios must not be empty.") + if any(ratio < 0 for ratio in compress_ratios): + raise ValueError("compress_ratios must be non-negative.") + return [1 if ratio == 0 else ratio for ratio in compress_ratios] + + def supports_backend(self, backend: str) -> bool: + return backend == "pytorch" + + def needs_separate_short_long_cuda_graphs(self) -> bool: + # DeepSeek-V4 does not support short/long CUDA graph separation. + return False + + class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig): """Configuration for skip softmax attention.""" algorithm: Literal["skip_softmax"] = Field(default="skip_softmax") @@ -2907,6 +2954,7 @@ def supports_backend(self, backend: str) -> bool: Union[ RocketSparseAttentionConfig, DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, SkipSoftmaxAttentionConfig, MiniMaxM3SparseAttentionConfig, ], diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py new file mode 100644 index 000000000000..b9e47a021eb7 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py @@ -0,0 +1,2694 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests comparing Compressor with RefCompressor.""" + +import math +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from tensorrt_llm._torch.attention_backend.interface import ( + MLAParams, + PositionalEmbeddingParams, + PositionEmbeddingType, + RotaryScalingType, +) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import ( + Compressor, + KVCacheDtype, +) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( + DeepseekV4AttentionType, +) +from tensorrt_llm._torch.modules.rotary_embedding import RopeParams +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping + + +def _hadamard_transform(x: torch.Tensor, scale: float) -> torch.Tensor: + """Pure-Python Walsh-Hadamard transform (no external dependency). + + Matches the butterfly implementation in the CUDA postProcessScatterKernel. + """ + n = x.shape[-1] + assert n & (n - 1) == 0, "Last dim must be a power of 2" + y = x.float().clone() + stride = 1 + while stride < n: + idx = torch.arange(n, device=x.device) + lo_mask = (idx & stride) == 0 + hi_idx = idx ^ stride + a = y[..., lo_mask].clone() + b = y[..., hi_idx[lo_mask]].clone() + y[..., lo_mask] = a + b + y[..., hi_idx[lo_mask]] = a - b + stride <<= 1 + return (y * scale).to(x.dtype) + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + """Hadamard rotation matching the CUDA kernel (no fast_hadamard_transform needed).""" + return _hadamard_transform(x, scale=x.size(-1) ** -0.5) + + +# ============================================================================ +# Dummy Metadata for Testing +# ============================================================================ + + +class DummyAttentionMetadata: + """Dummy attention metadata for testing Compressor.forward.""" + + def __init__( + self, + num_contexts: int, + num_generations: int, + num_ctx_tokens: int, + num_tokens: int, + kv_cache_manager: DeepseekV4CacheManager, + block_tables: dict, + cu_seq_lens: dict, + cu_new_comp_kv: dict, + compressed_position_ids: dict, + compressed_kv_lens: dict, + past_kv_lens: dict, + new_comp_kv_lens_cuda: dict, + num_total_compressed_tokens: dict, + max_ctx_compressed_tokens: dict, + slot_mapping_fp8: torch.Tensor = None, + slot_mapping_scale: torch.Tensor = None, + compressed_mask_cuda: dict = None, + ): + self.num_contexts = num_contexts + self.num_generations = num_generations + self.num_ctx_tokens = num_ctx_tokens + self.num_tokens = num_tokens + self.kv_cache_manager = kv_cache_manager + self.block_tables = block_tables + self.cu_seq_lens_cuda = cu_seq_lens + self.cu_new_comp_kv_cuda = cu_new_comp_kv + self.compressed_position_ids_cuda = compressed_position_ids + self.compressed_kv_lens_cuda = compressed_kv_lens + self.past_kv_lens_cuda = past_kv_lens + self.slot_mapping_fp8 = slot_mapping_fp8 + self.slot_mapping_scale = slot_mapping_scale + self.new_comp_kv_lens_cuda = new_comp_kv_lens_cuda + self.num_total_compressed_tokens = num_total_compressed_tokens + self.max_ctx_compressed_tokens = max_ctx_compressed_tokens + self.compressed_mask_cuda = compressed_mask_cuda + self.num_gen_tokens_per_seq = 0 # Set by caller + self.kv_lens_cuda_runtime = None # Set by caller + self.cached_token_lens_cuda = None # Set by caller + + +# ============================================================================ +# Reference Implementation (DO NOT MODIFY) +# ============================================================================ + + +@dataclass +class ModelArgs: + """Model arguments for Compressor.""" + + max_batch_size: int = 16 + max_seq_len: int = 4096 + dim: int = 4096 + head_dim: int = 512 + rope_head_dim: int = 64 + norm_eps: float = 1e-6 + compress_ratios: Tuple[int, ...] = (1, 1, 4, 128, 4, 128, 4) + + +class RMSNorm(nn.Module): + """Root Mean Square Layer Normalization.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return (self.weight * x).to(dtype) + + +class Linear(nn.Module): + """Simple linear layer (fp32 weights for reference).""" + + def __init__(self, in_features: int, out_features: int, dtype=None): + super().__init__() + self.weight = nn.Parameter( + torch.empty(out_features, in_features, dtype=dtype or torch.float32) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.linear(x, self.weight) + + +def apply_rotary_emb( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Apply rotary positional embeddings.""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +class RefCompressor(nn.Module): + """Reference Compressor implementation for testing.""" + + def __init__( + self, args: ModelArgs, compress_ratio: int = 4, head_dim: int = 512, rotate: bool = False + ): + super().__init__() + self.dim = args.dim + self.head_dim = head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = head_dim - args.rope_head_dim + self.compress_ratio = compress_ratio + self.overlap = compress_ratio == 4 + self.rotate = rotate + coff = 1 + self.overlap + + self.ape = nn.Parameter( + torch.empty(compress_ratio, coff * self.head_dim, dtype=torch.float32) + ) + self.wkv = Linear(self.dim, coff * self.head_dim, dtype=torch.float32) + self.wgate = Linear(self.dim, coff * self.head_dim, dtype=torch.float32) + self.norm = RMSNorm(self.head_dim, args.norm_eps) + self.kv_cache = None + self.register_buffer( + "kv_state", + torch.zeros( + args.max_batch_size, + coff * compress_ratio, + coff * self.head_dim, + dtype=torch.float32, + ), + persistent=False, + ) + self.register_buffer( + "score_state", + torch.full( + (args.max_batch_size, coff * compress_ratio, coff * self.head_dim), + float("-inf"), + dtype=torch.float32, + ), + persistent=False, + ) + + def overlap_transform(self, tensor: torch.Tensor, value=0): + b, s, _, _ = tensor.size() + ratio, d = self.compress_ratio, self.head_dim + new_tensor = tensor.new_full((b, s, 2 * ratio, d), value) + new_tensor[:, :, ratio:] = tensor[:, :, :, d:] + new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d] + return new_tensor + + def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor): + assert self.kv_cache is not None + bsz, seqlen, _ = x.size() + ratio, overlap, d = self.compress_ratio, self.overlap, self.head_dim + dtype = x.dtype + x = x.float() + kv = self.wkv(x) + score = self.wgate(x) + if start_pos == 0: + should_compress = seqlen >= ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + freqs_cis = freqs_cis[:cutoff:ratio] + offset = ratio if overlap else 0 + if overlap and cutoff >= ratio: + self.kv_state[:bsz, :ratio] = kv[:, cutoff - ratio : cutoff] + self.score_state[:bsz, :ratio] = score[:, cutoff - ratio : cutoff] + self.ape + if remainder > 0: + kv, self.kv_state[:bsz, offset : offset + remainder] = kv.split( + [cutoff, remainder], dim=1 + ) + self.score_state[:bsz, offset : offset + remainder] = ( + score[:, cutoff:] + self.ape[:remainder] + ) + score = score[:, :cutoff] + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + self.ape + if overlap: + kv = self.overlap_transform(kv, 0) + score = self.overlap_transform(score, float("-inf")) + kv = (kv * score.softmax(dim=2)).sum(dim=2) + else: + # Handles any seqlen >= 1 with start_pos > 0 (decode or chunked prefill). + # freqs_cis must be the FULL precomputed array so output freqs can be + # indexed at absolute first-token-of-window positions (win_first = pos+1-ratio). + outputs = [] + output_freqs = [] + for t in range(seqlen): + pos = start_pos + t + kv_t = kv[:, t] + sc_t = score[:, t] + self.ape[pos % ratio] + if overlap: + self.kv_state[:bsz, ratio + pos % ratio] = kv_t + self.score_state[:bsz, ratio + pos % ratio] = sc_t + if (pos + 1) % ratio == 0: + kv_state = torch.cat( + [self.kv_state[:bsz, :ratio, :d], self.kv_state[:bsz, ratio:, d:]], + dim=1, + ) + score_state = torch.cat( + [ + self.score_state[:bsz, :ratio, :d], + self.score_state[:bsz, ratio:, d:], + ], + dim=1, + ) + comp = (kv_state * score_state.softmax(dim=1)).sum(dim=1, keepdim=True) + self.kv_state[:bsz, :ratio] = self.kv_state[:bsz, ratio:] + self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:] + outputs.append(comp) + output_freqs.append(freqs_cis[pos + 1 - ratio : pos + 2 - ratio]) + else: + self.kv_state[:bsz, pos % ratio] = kv_t + self.score_state[:bsz, pos % ratio] = sc_t + if (pos + 1) % ratio == 0: + comp = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum( + dim=1, keepdim=True + ) + outputs.append(comp) + output_freqs.append(freqs_cis[pos + 1 - ratio : pos + 2 - ratio]) + + if not outputs: + return None + should_compress = True + kv = torch.cat(outputs, dim=1) + freqs_cis = torch.cat(output_freqs, dim=0) + + if not should_compress: + return + # Match the kernel's hand-off: the kernel's kv_comp buffer is bf16, + # so we bf16-truncate the compression result here too. Then upcast + # back to fp32 for the postprocess so RMSNorm/RoPE/Hadamard run at + # full fp32 precision -- matching the kernel which keeps activations + # in fp32 registers throughout postprocess (no V4-Pro fake-quant bf16 + # round-trips between norm / rope / hadamard). Returns fp32 so the + # downstream QDQ reference sees the same full-precision values the + # kernel feeds into its QDQ on registers; only the cache write + # truncates to bf16. + kv = kv.to(dtype).float() + kv = self.norm(kv) + apply_rotary_emb(kv[..., -self.rope_head_dim :], freqs_cis) + if self.rotate: + kv = rotate_activation(kv) + if start_pos == 0: + self.kv_cache[:bsz, : seqlen // ratio] = kv.to(dtype) + else: + first_abs = start_pos // ratio + n_out = kv.size(1) + self.kv_cache[:bsz, first_abs : first_abs + n_out] = kv.to(dtype) + return kv + + +# ============================================================================ +# Test Configuration & Helpers +# ============================================================================ + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +DIM, HEAD_DIM, ROPE_DIM = 4096, 512, 64 +INDEX_HEAD_DIM = 128 # Fixed head_dim for indexer (INDEXER_COMPRESS) +MAX_BATCH, MAX_SEQ, PAGE_SIZE = 16, 4096, 128 +ORI_SEQ_LEN = 65536 +ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW = 40000.0, 4, 32, 1 + + +def precompute_freqs_cis( + dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow +) -> torch.Tensor: + """Precompute rotary embeddings.""" + + def find_correction_dim(num_rotations, dim, base, max_seq_len): + return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_seq_len): + low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len)) + high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len)) + return max(low, 0), min(high, dim - 1) + + def linear_ramp_factor(min, max, dim): + if min == max: + max += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if seqlen > original_seq_len: + low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_seq_len) + smooth = 1 - linear_ramp_factor(low, high, dim // 2) + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return freqs_cis + + +def assert_similar(t1: Optional[torch.Tensor], t2: Optional[torch.Tensor], name: str = "Output"): + """Assert tensors are similar (cosine sim >= 0.999).""" + if t1 is None and t2 is None: + return + assert t1 is not None and t2 is not None, f"{name}: One is None" + assert t1.shape == t2.shape, f"{name}: Shape mismatch {t1.shape} vs {t2.shape}" + t1, t2 = t1.float().flatten(), t2.float().flatten() + cos_sim = F.cosine_similarity(t1.unsqueeze(0), t2.unsqueeze(0)).item() + # Also check magnitude to avoid scaled-but-equal-direction false positives + max_diff = (t1 - t2).abs().max().item() + scale = max(t1.abs().max().item(), t2.abs().max().item(), 1e-3) + rel_err = max_diff / scale + assert cos_sim >= 0.999, f"{name}: cos_sim={cos_sim:.6f}" + assert rel_err <= 5e-2, f"{name}: rel_err={rel_err:.6f}, max_diff={max_diff:.6f}" + + +def dequantize_blockwise_fp8( + kv_fp8: torch.Tensor, kv_scale: torch.Tensor, block_size: int = 128 +) -> torch.Tensor: + """Dequantize blockwise FP8 data. + + Args: + kv_fp8: [num_tokens, head_dim] FP8 tensor + kv_scale: [num_tokens, num_scale_blocks] scale factors (one per 128 elements) + + Returns: + Dequantized float tensor + """ + num_tokens, head_dim = kv_fp8.shape + num_blocks = (head_dim + block_size - 1) // block_size + + kv_float = kv_fp8.float() + kv_dequant = torch.zeros_like(kv_float) + + for b in range(num_blocks): + start = b * block_size + end = min(start + block_size, head_dim) + kv_dequant[:, start:end] = kv_float[:, start:end] * kv_scale[:, b : b + 1] + + return kv_dequant + + +def assert_fp8_similar( + fp8_result: tuple, + ref_result: torch.Tensor, + kv_cache_dtype: str, + name: str = "FP8 Output", +): + """Assert FP8 result matches reference after dequantization. + + Uses FP8-appropriate tolerances: cos_sim >= 0.99, rel_err <= 10% + """ + kv_fp8, scale = fp8_result + + if kv_cache_dtype in ("fp8_blockwise"): + kv_dequant = dequantize_blockwise_fp8(kv_fp8, scale) + else: # fp8_pertensor + kv_dequant = kv_fp8.float() * scale.item() + + ref_float = ref_result.float().flatten() + dequant_flat = kv_dequant.flatten() + + # Cosine similarity (>= 0.99 for FP8) + cos_sim = F.cosine_similarity(dequant_flat.unsqueeze(0), ref_float.unsqueeze(0)).item() + assert cos_sim >= 0.99, f"{name}: cos_sim={cos_sim:.4f} < 0.99" + + # Relative error (<= 10% for FP8) + max_diff = (dequant_flat - ref_float).abs().max().item() + scale_val = max(ref_float.abs().max().item(), 1e-3) + rel_err = max_diff / scale_val + assert rel_err <= 0.1, f"{name}: rel_err={rel_err:.4f} > 0.1" + + +def read_paged_cache_tokens( + kv_cache: torch.Tensor, + block_offsets: torch.Tensor, + batch_idx: int, + num_tokens: int, + tokens_per_block: int, +) -> torch.Tensor: + """Materialize paged compressed cache for a batch into a contiguous view. + + Args: + kv_cache: Cache buffer with shape [num_blocks, tokens_per_block, head_dim] + block_offsets: Block offset table with shape [num_seqs, max_blocks] + batch_idx: Index of the batch to read + num_tokens: Number of tokens to read + tokens_per_block: Tokens per cache block + + Returns: + Tensor containing the read values with shape [num_tokens, head_dim] + """ + blocks_needed = (num_tokens + tokens_per_block - 1) // tokens_per_block + # Extract block indices for this batch + block_indices = block_offsets[batch_idx, :blocks_needed].tolist() + # Read all blocks at once and reshape + _, _, dim_per_token = kv_cache.shape + return kv_cache[block_indices].reshape(-1, dim_per_token)[:num_tokens] + + +def build_fp8_golden_cache( + kv_fp8: torch.Tensor, + kv_scale: torch.Tensor, + cache_shape: tuple, + block_offsets: torch.Tensor, + batch: int, + num_compressed: int, + tokens_per_block: int, + kv_cache_dtype: str, + head_dim: int = HEAD_DIM, +) -> torch.Tensor: + """Build Python golden reference cache from compressor's FP8 output. + + This validates the cache scatter operation by manually placing the compressor's + FP8 output into the expected cache positions. Quantization correctness is + validated separately via assert_fp8_similar. + + Args: + kv_fp8: Compressor's FP8 output [total_tokens, head_dim] + kv_scale: Compressor's scale output + cache_shape: Shape of the cache tensor + block_offsets: Block offset table [num_seqs, max_blocks] + batch: Batch size + num_compressed: Compressed tokens per batch + tokens_per_block: Tokens per cache block + kv_cache_dtype: "fp8_blockwise" or "fp8_pertensor" + head_dim: Head dimension (default HEAD_DIM=512, use INDEX_HEAD_DIM=128 for indexer) + + Returns: + golden_cache: Expected cache tensor for comparison + """ + total_comp_tokens = batch * num_compressed + golden_cache = torch.zeros(cache_shape, device=kv_fp8.device, dtype=torch.uint8) + + # Convert FP8 to bytes + kv_fp8_bytes = kv_fp8.contiguous().view(torch.uint8).view(total_comp_tokens, head_dim) + + if kv_cache_dtype in ("fp8_blockwise"): + # Blockwise: non-interleaved layout per block: + # [k0, k1, ..., kN, scale0, scale1, ..., scaleN] + num_scale_blocks = (head_dim + 127) // 128 + scale_size = num_scale_blocks * 4 + kv_scale_bytes = ( + kv_scale.flatten().contiguous().view(torch.uint8).view(total_comp_tokens, scale_size) + ) + + # Flatten to 2D [num_blocks, block_stride] for flat byte indexing + num_blocks = cache_shape[0] + golden_flat = golden_cache.view(num_blocks, -1) + + # Scatter into cache using non-interleaved offsets + global_token = 0 + for b in range(batch): + for i in range(num_compressed): + block_idx = i // tokens_per_block + pos_in_block = i % tokens_per_block + block_id = int(block_offsets[b, block_idx].item()) + + # FP8 data in first section of block + fp8_start = pos_in_block * head_dim + golden_flat[block_id, fp8_start : fp8_start + head_dim] = kv_fp8_bytes[global_token] + # Scale data in second section of block + scale_start = tokens_per_block * head_dim + pos_in_block * scale_size + golden_flat[block_id, scale_start : scale_start + scale_size] = kv_scale_bytes[ + global_token + ] + global_token += 1 + + # Reshape back to original shape + golden_cache = golden_flat.view(cache_shape) + else: + # Per-tensor: cache layout is [num_blocks, tokens_per_block, head_dim] + # Same as blockwise but without scale bytes + global_token = 0 + for b in range(batch): + for i in range(num_compressed): + block_idx = i // tokens_per_block + pos_in_block = i % tokens_per_block + block_id = int(block_offsets[b, block_idx].item()) + + golden_cache[block_id, pos_in_block, :head_dim] = kv_fp8_bytes[global_token] + global_token += 1 + + return golden_cache + + +def assert_fp8_cache_match( + kernel_cache: torch.Tensor, + golden_cache: torch.Tensor, + kv_cache_dtype: str, + name: str = "FP8 Cache", + head_dim: int = HEAD_DIM, +): + """Assert kernel cache matches Python golden reference.""" + # Convert both to uint8 bytes for comparison (kernel_cache may be Float8_e4m3fn) + kernel_bytes = kernel_cache.view(torch.uint8) + golden_bytes = golden_cache.view(torch.uint8) + + if torch.equal(kernel_bytes, golden_bytes): + return # Perfect match + + diff_mask = kernel_bytes != golden_bytes + num_diffs = diff_mask.sum().item() + total_bytes = kernel_bytes.numel() + + # Build detailed error message + msg = ( + f"{name}: {num_diffs}/{total_bytes} byte differences ({100 * num_diffs / total_bytes:.4f}%)" + ) + + if kv_cache_dtype in ("fp8_blockwise"): + # Reshape to original layout for detailed analysis + kernel_reshaped = kernel_bytes.view(kernel_cache.shape) + golden_reshaped = golden_bytes.view(golden_cache.shape) + fp8_diffs = ( + (kernel_reshaped[:, :, :head_dim] != golden_reshaped[:, :, :head_dim]).sum().item() + ) + scale_diffs = ( + (kernel_reshaped[:, :, head_dim:] != golden_reshaped[:, :, head_dim:]).sum().item() + ) + + msg += f" [FP8: {fp8_diffs}, Scale: {scale_diffs}]" + + # Show first few differences + diff_indices = torch.nonzero(diff_mask.view(-1))[:5] + for idx in diff_indices: + flat_idx = idx.item() + msg += f"\n Byte {flat_idx}: kernel={kernel_cache.view(-1)[flat_idx].item()}, " + msg += f"golden={golden_cache.view(-1)[flat_idx].item()}" + + raise AssertionError(msg) + + +def run_ref_segmented_forward( + ref: RefCompressor, + tokens: torch.Tensor, + freqs_cis: torch.Tensor, + segments: list[tuple[int, int]], +) -> Optional[torch.Tensor]: + """Run ref forward per segment, concatenating non-None outputs. + + For start_pos == 0 a freq slice is passed (forward indexes from 0). + For start_pos > 0 the FULL freqs_cis array is passed so forward() can + index at absolute win_first positions. + """ + outputs = [] + cursor = 0 + for start_pos, seg_len in segments: + seg_tokens = tokens[:, cursor : cursor + seg_len] + if start_pos > 0: + out = ref(seg_tokens, start_pos, freqs_cis) + else: + out = ref(seg_tokens, start_pos, freqs_cis[start_pos : start_pos + seg_len]) + if out is not None: + outputs.append(out) + cursor += seg_len + assert cursor <= tokens.size(1), "Segment lengths exceed provided tokens" + if not outputs: + return None + return torch.cat(outputs, dim=1) + + +class CompressorWrapper: + """Wrapper around Compressor to manage caches and provide a simpler test interface.""" + + # Class-level constants for DeepseekV4CacheManager + WINDOW_SIZE = 128 + VOCAB_SIZE = 129280 + + def __init__( + self, + compress_ratio: int = 4, + rotate: bool = False, + layer_idx: int = 0, + kv_cache_dtype: str = "default", + is_indexer: bool = False, + ): + self.compress_ratio = compress_ratio + self.overlap = compress_ratio == 4 + self.layer_idx = layer_idx + self.kv_cache_dtype = kv_cache_dtype + self.is_indexer = is_indexer + + # Create MLAParams + # For indexer mode, use INDEX_HEAD_DIM instead of HEAD_DIM + target_head_dim = INDEX_HEAD_DIM if is_indexer else HEAD_DIM + mla_params = MLAParams( + hidden_size=DIM, + qk_rope_head_dim=ROPE_DIM, + qk_nope_head_dim=target_head_dim - ROPE_DIM, + ) + + # Create RoPE - use no scaling to match precompute_freqs_cis behavior + # (precompute_freqs_cis only applies yarn scaling when seqlen > ORI_SEQ_LEN) + rope_params = RopeParams( + dim=ROPE_DIM, + theta=ROPE_THETA, + max_positions=4096, + beta_fast=BETA_FAST, + beta_slow=BETA_SLOW, + scale=1.0, # No scaling + mscale=1.0, + mscale_all_dim=1.0, + original_max_positions=ORI_SEQ_LEN, + scale_type=RotaryScalingType.none, # Match precompute_freqs_cis (no yarn for pos < ORI_SEQ_LEN) + ) + + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, # Basic RoPE without yarn + rope=rope_params, + is_neox=False, + ) + + # Create the Compressor + self.compressor = Compressor( + mla_params=mla_params, + layer_idx=layer_idx, + compress_ratio=compress_ratio, + norm_eps=1e-6, + skip_create_weights_in_init=False, + pos_embd_params=pos_embd_params, + dtype=DTYPE, + kv_cache_dtype=kv_cache_dtype, + is_indexer=is_indexer, + rotate_activation=rotate, + ).to(DEVICE) + + # Create DeepseekV4CacheManager + self.cache_manager = self._create_deepseek_v4_cache_manager(compress_ratio) + # COMPRESS cache has tokens_per_block from cache manager's compressed_block_sizes + self.tokens_per_block = self.cache_manager.compressed_block_sizes[self.layer_idx] + + # Track active requests for the cache manager + self.active_requests: Dict[int, LlmRequest] = {} + self.next_request_id = 0 + + # Store reference to kv_cache for test compatibility (read from cache manager) + self._update_kv_cache_reference() + + # Create block_offsets for test compatibility (will be updated per forward) + max_compressed = MAX_SEQ // compress_ratio + max_comp_blocks = (max_compressed + self.tokens_per_block - 1) // self.tokens_per_block + self.block_offsets = torch.zeros( + MAX_BATCH, max_comp_blocks, device=DEVICE, dtype=torch.int32 + ) + + def cleanup(self): + """Free all active requests and shut down the cache manager. + + Must be called before the wrapper is discarded, otherwise the cache + manager's destructor will fail with ResourceBusyError and GPU memory + will leak across tests. + """ + if not hasattr(self, "cache_manager"): + return + for req in self.active_requests.values(): + self.cache_manager.free_resources(req) + self.active_requests.clear() + self.cache_manager.shutdown() + + def __del__(self): + try: + self.cleanup() + except Exception: + pass + + def _create_deepseek_v4_cache_manager(self, compress_ratio: int) -> DeepseekV4CacheManager: + """Create a DeepseekV4CacheManager for testing.""" + # Single layer with the given compress ratio + compress_ratios = [compress_ratio] + + # Create sparse attention config + sparse_attn_config = DeepSeekV4SparseAttentionConfig( + index_head_dim=INDEX_HEAD_DIM, + window_size=self.WINDOW_SIZE, + compress_ratios=compress_ratios, + ) + + # Create KV cache config + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + max_tokens=MAX_SEQ * MAX_BATCH, + event_buffer_max_size=0, + ) + + # Create mapping (single GPU, no parallelism) + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + if self.kv_cache_dtype in ["fp8_pertensor", "fp8_blockwise"]: + cache_dtype = DataType.FP8 + else: + cache_dtype = DataType.BF16 + + # Create cache manager + cache_manager = DeepseekV4CacheManager( + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=len(compress_ratios), + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=PAGE_SIZE, + max_seq_len=MAX_SEQ, + max_batch_size=MAX_BATCH, + max_input_len=MAX_SEQ, + mapping=mapping, + dtype=cache_dtype, + compressor_dtype=DataType.FLOAT, # State caches always use FP32 + vocab_size=self.VOCAB_SIZE, + max_num_tokens=MAX_SEQ + MAX_BATCH, + sparse_attn_config=sparse_attn_config, + ) + + return cache_manager + + def _update_kv_cache_reference(self): + """Update the kv_cache reference from cache manager for test compatibility.""" + compress_type = ( + DeepseekV4AttentionType.INDEXER_COMPRESS + if self.is_indexer + else DeepseekV4AttentionType.COMPRESS + ) + self.kv_cache = self.cache_manager.get_buffers(self.layer_idx, compress_type) + + def _create_request(self, request_id: int, prompt_len: int) -> LlmRequest: + """Helper to create a test LlmRequest (following test_deepseek_v4_cache_manager pattern). + + Args: + request_id: Unique request identifier + prompt_len: Prompt length (number of tokens) + + Returns: + LlmRequest instance + """ + input_tokens = list(range(prompt_len)) + request = LlmRequest( + request_id=request_id, + max_new_tokens=1024, + input_tokens=input_tokens, + sampling_config=SamplingConfig(), + is_streaming=False, + ) + return request + + def _prepare_requests_for_batch( + self, + bsz: int, + seq_lens: torch.Tensor, + start_pos: torch.Tensor, + is_prefill: torch.Tensor, + batch_indices: torch.Tensor = None, + ) -> Tuple[List[LlmRequest], ScheduledRequests]: + """Prepare requests for a batch using the KVCacheV2 scheduler allocation flow. + + Args: + batch_indices: Optional tensor mapping batch positions to external batch indices. + If provided, requests are stored/retrieved using these indices. + If None, sequential indices (0, 1, 2...) are used. + + Returns: + Tuple of (requests list, scheduled_batch) for later update_resources call + """ + requests = [] + context_requests = [] + generation_requests = [] + + # Use batch_indices if provided, otherwise use sequential indices + if batch_indices is not None: + ext_indices = [int(batch_indices[b].item()) for b in range(bsz)] + else: + ext_indices = list(range(bsz)) + + # Separate prefill and generation indices. + prefill_indices = [] + gen_indices = [] + for b in range(bsz): + if is_prefill[b]: + prefill_indices.append(b) + else: + gen_indices.append(b) + + # Handle prefill requests, including chunked prefill with start_pos > 0. + for b in prefill_indices: + ext_idx = ext_indices[b] + req = self.active_requests.get(ext_idx) + total_prompt_len = int((start_pos[b] + seq_lens[b]).item()) + chunk_size = int(seq_lens[b].item()) + if req is None: + assert int(start_pos[b].item()) == 0, ( + "Chunked prefill requests must reuse an existing request created by " + "the initial context chunk" + ) + req = self._create_request(self.next_request_id, total_prompt_len) + self.active_requests[ext_idx] = req + self.next_request_id += 1 + req.state = LlmRequestState.CONTEXT_INIT + req.context_current_position = int(start_pos[b].item()) + req.prompt_len = total_prompt_len + req.py_prompt_len = total_prompt_len + req.context_chunk_size = chunk_size + req.py_draft_tokens = [] + context_requests.append(req) + + # Handle generation requests - reuse existing requests or create new ones + for b in gen_indices: + ext_idx = ext_indices[b] + req = self.active_requests.get(ext_idx) + req_seq_len = int(seq_lens[b].item()) + if req is None: + # Need to create a new request for generation (prefill was done in previous call) + pos = int(start_pos[b].item()) + req = self._create_request(self.next_request_id, pos) + req.state = LlmRequestState.GENERATION_IN_PROGRESS + req.context_current_position = pos + req.add_new_token(pos, 0) # Simulate having processed tokens + self.active_requests[ext_idx] = req + self.next_request_id += 1 + else: + # Existing request from previous prefill - mark as generation + req.state = LlmRequestState.GENERATION_IN_PROGRESS + # Scheduler v2 allocates 1 sampled token plus any draft tokens. + # Mirror that contract so multi-token generation reserves enough KV slots. + req.py_draft_tokens = [0] * max(req_seq_len - 1, 0) + generation_requests.append(req) + + # Build final request list in batch order + for b in range(bsz): + ext_idx = ext_indices[b] + requests.append(self.active_requests[ext_idx]) + + # Build scheduled batch and allocate buffers like KVCacheV2Scheduler. + scheduled_batch = ScheduledRequests() + for req in context_requests: + scheduled_batch.append_context_request(req) + scheduled_batch.generation_requests = generation_requests + for req in context_requests: + assert self.cache_manager.prepare_context(req), ( + f"Failed to prepare context for request {req.py_request_id}" + ) + assert self.cache_manager.resize_context(req, req.context_chunk_size), ( + f"Failed to resize context for request {req.py_request_id}" + ) + for req in generation_requests: + assert self.cache_manager.try_allocate_generation(req), ( + f"Failed to allocate generation KV cache for request {req.py_request_id}" + ) + + return requests, scheduled_batch + + def _get_block_table_for_request( + self, + req: LlmRequest, + attn_type: DeepseekV4AttentionType, + ) -> torch.Tensor: + """Get block table for a request and attention type.""" + page_indices = self.cache_manager.get_cache_indices( + request_id=req.py_request_id, + layer_idx=self.layer_idx, + attn_type=attn_type, + ) + return torch.tensor(page_indices, dtype=torch.int32, device=DEVICE) + + def forward( + self, + x: torch.Tensor, + start_pos: int | torch.Tensor, + freqs_cis: torch.Tensor, + batch_indices: torch.Tensor = None, + seq_lens: torch.Tensor = None, + *, + is_prefill: torch.Tensor | None = None, + ): + """Wrapper forward that matches the reference Compressor interface. + + Supports mixed prefill+decode when seq_lens and start_pos tensor are provided. + """ + ratio = self.compress_ratio + + def normalize_is_prefill( + prefill: torch.Tensor | None, default: torch.Tensor + ) -> torch.Tensor: + if prefill is None: + prefill_tensor = default + else: + assert isinstance(prefill, torch.Tensor), "is_prefill must be a torch.Tensor" + prefill_tensor = prefill.to(device=DEVICE, dtype=torch.bool) + if prefill_tensor.ndim == 0: + prefill_tensor = prefill_tensor.expand(bsz) + assert prefill_tensor.shape == (bsz,), ( + f"is_prefill must have shape ({bsz},), got {tuple(prefill_tensor.shape)}" + ) + return prefill_tensor + + # Handle variable-length sequences + if seq_lens is not None: + # Mixed batch mode with variable-length sequences + seq_lens = seq_lens.to(torch.int32) + bsz = seq_lens.size(0) + if isinstance(start_pos, torch.Tensor): + start_pos_tensor = start_pos.to(torch.int32) + else: + start_pos_tensor = torch.full((bsz,), start_pos, dtype=torch.int32, device=DEVICE) + + # Flatten input tokens + if x.ndim == 3: + x_flat = x.view(-1, DIM) + else: + x_flat = x + total_tokens = int(seq_lens.sum().item()) + x_flat = x_flat[:total_tokens] + + is_prefill_tensor = normalize_is_prefill( + is_prefill, (start_pos_tensor == 0) | (seq_lens > 1) + ) + + # Determine which sequences are context (prefill) vs generation (decode). + is_context = is_prefill_tensor + num_contexts = int(is_context.sum().item()) + num_generations = bsz - num_contexts + + # Reorder: contexts first, then generations + ctx_indices = torch.where(is_context)[0] + gen_indices = torch.where(~is_context)[0] + reorder_indices = torch.cat([ctx_indices, gen_indices]) + + seq_lens_reordered = seq_lens[reorder_indices] + start_pos_reordered = start_pos_tensor[reorder_indices] + + # Compute token offsets for reordering + cu_seq_original = torch.zeros(bsz + 1, dtype=torch.int32, device=DEVICE) + cu_seq_original[1:] = seq_lens.cumsum(0) + + # Reorder tokens: contexts first, then generations + token_indices = [] + for idx in reorder_indices: + start = cu_seq_original[idx].item() + end = cu_seq_original[idx + 1].item() + token_indices.extend(range(start, end)) + x_flat = x_flat[token_indices] + + num_ctx_tokens = ( + int(seq_lens_reordered[:num_contexts].sum().item()) if num_contexts > 0 else 0 + ) + num_gen_tokens = ( + int(seq_lens_reordered[num_contexts:].sum().item()) if num_generations > 0 else 0 + ) + seq_lens = seq_lens_reordered + past_kv_lens = start_pos_reordered + is_prefill = is_context[reorder_indices] + + # Use reorder_indices for request mapping (maps reordered position to original batch index) + batch_indices_for_requests = reorder_indices + start_pos_for_kv = past_kv_lens + + # Get number of compressed tokens + num_comp_per_seq = (past_kv_lens + seq_lens) // ratio - past_kv_lens // ratio + num_ctx_compressed_tokens = int(num_comp_per_seq[:num_contexts].sum().item()) + num_gen_compressed_tokens = int(num_comp_per_seq[num_contexts:].sum().item()) + else: + # Original single-mode logic + bsz, seqlen, _ = x.size() + x_flat = x.reshape(-1, DIM) + if isinstance(start_pos, torch.Tensor): + past_kv_lens = start_pos.to(torch.int32).to(DEVICE) + if past_kv_lens.ndim == 0: + past_kv_lens = past_kv_lens.expand(bsz) + else: + past_kv_lens = torch.full((bsz,), start_pos, dtype=torch.int32, device=DEVICE) + + is_prefill = normalize_is_prefill( + is_prefill, torch.full((bsz,), seqlen > 1, dtype=torch.bool, device=DEVICE) + ) + if not torch.all(is_prefill == is_prefill[0]): + raise ValueError( + "single-shape forward requires uniform is_prefill across the batch" + ) + is_prefill_value = bool(is_prefill[0].item()) + + if not is_prefill_value: + # Decode mode + num_contexts = 0 + num_generations = bsz + num_ctx_tokens = 0 + num_gen_tokens = bsz * seqlen + seq_lens = torch.full((bsz,), seqlen, dtype=torch.int32, device=DEVICE) + num_ctx_compressed_tokens = 0 + kv_lens_local = past_kv_lens + seq_lens + num_gen_compressed_tokens = int( + (kv_lens_local // ratio - past_kv_lens // ratio).sum().item() + ) + else: + # Prefill mode (may be chunked when start_pos > 0) + num_contexts = bsz + num_generations = 0 + num_ctx_tokens = bsz * seqlen + num_gen_tokens = 0 + seq_lens = torch.full((bsz,), seqlen, dtype=torch.int32, device=DEVICE) + kv_lens_local = past_kv_lens + seq_lens + num_ctx_compressed_tokens = int( + (kv_lens_local // ratio - past_kv_lens // ratio).sum().item() + ) + num_gen_compressed_tokens = 0 + # Use batch_indices for request mapping if provided + batch_indices_for_requests = batch_indices + start_pos_for_kv = past_kv_lens + + # Prepare requests for the cache manager + requests, scheduled_batch = self._prepare_requests_for_batch( + bsz, seq_lens, start_pos_for_kv, is_prefill, batch_indices_for_requests + ) + + cu_seq_lens = torch.zeros(bsz + 1, dtype=torch.int32, device=DEVICE) + cu_seq_lens[1:] = seq_lens.cumsum(0) + num_total_compressed_tokens = num_ctx_compressed_tokens + num_gen_compressed_tokens + + # Compute KV lengths (past + current) per sequence + kv_lens = past_kv_lens + seq_lens + + # Compute number of compressed outputs per batch using absolute-aligned formula. + # kv_len // ratio - past // ratio works for all cases: + # fresh prefill (past=0): kv_len // ratio + # chunked prefill (past>0, seqlen>1): correct window boundary counting + # generation (past>0): fires when chunk boundary is crossed + num_comp = (past_kv_lens + seq_lens) // ratio - past_kv_lens // ratio + + cu_new_comp_kv = torch.zeros(bsz + 1, dtype=torch.int32, device=DEVICE) + cu_new_comp_kv[1:] = num_comp.cumsum(0) + max_ctx_comp_kv_lens = num_comp[:num_contexts].max().item() if num_contexts > 0 else 0 + + # Create position IDs for compressed outputs. + # Always use first-token-of-window convention: position = (base_chunk + c) * ratio. + # This matches RefCompressor which uses win_first = pos + 1 - ratio for all cases. + position_ids = torch.zeros(num_total_compressed_tokens, dtype=torch.int32, device=DEVICE) + offset = 0 + for b in range(bsz): + n_out = num_comp[b].item() + sp = past_kv_lens[b].item() + base_chunk = sp // ratio + for c in range(n_out): + position_ids[offset + c] = (base_chunk + c) * ratio + offset += n_out + + # Determine attention types based on is_indexer + if self.is_indexer: + compress_type = DeepseekV4AttentionType.INDEXER_COMPRESS + state_type = DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE + score_type = DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE + else: + compress_type = DeepseekV4AttentionType.COMPRESS + state_type = DeepseekV4AttentionType.COMPRESSOR_STATE + score_type = DeepseekV4AttentionType.COMPRESSOR_SCORE + + # Build block_tables dict keyed by DeepseekV4AttentionType using cache manager + block_table_compress_list = [] + block_table_kv_state_list = [] + block_table_score_state_list = [] + + for b, req in enumerate(requests): + block_table_compress_list.append(self._get_block_table_for_request(req, compress_type)) + block_table_kv_state_list.append(self._get_block_table_for_request(req, state_type)) + block_table_score_state_list.append(self._get_block_table_for_request(req, score_type)) + + # Pad and stack block tables to handle variable-length block indices + max_blocks_compress = max(bt.size(0) for bt in block_table_compress_list) + block_table_compress = torch.zeros( + bsz, max_blocks_compress, dtype=torch.int32, device=DEVICE + ) + for b in range(bsz): + bt = block_table_compress_list[b] + block_table_compress[b, : bt.size(0)] = bt + + max_blocks_state = max(bt.size(0) for bt in block_table_kv_state_list) + block_table_kv_state = torch.zeros(bsz, max_blocks_state, dtype=torch.int32, device=DEVICE) + block_table_score_state = torch.zeros( + bsz, max_blocks_state, dtype=torch.int32, device=DEVICE + ) + for b in range(bsz): + bt_kv = block_table_kv_state_list[b] + bt_score = block_table_score_state_list[b] + block_table_kv_state[b, : bt_kv.size(0)] = bt_kv + block_table_score_state[b, : bt_score.size(0)] = bt_score + + # Update block_offsets for test compatibility + self.block_offsets = block_table_compress + + block_tables = { + (ratio, compress_type): block_table_compress, + (ratio, state_type): block_table_kv_state, + (ratio, score_type): block_table_score_state, + } + + # Both prefill and decode kernels use absolute token positions for the + # state cache, so pass the absolute kv_lens directly. + # Build dicts keyed by compress_ratio + cu_new_comp_kv_dict = {ratio: cu_new_comp_kv} + compressed_position_ids_dict = {ratio: position_ids} + compressed_kv_lens_dict = {ratio: kv_lens} + past_kv_lens_dict = {ratio: past_kv_lens // ratio} + new_comp_kv_lens_cuda_dict = {ratio: num_comp} + num_total_compressed_tokens_dict = {ratio: num_total_compressed_tokens} + max_ctx_compressed_tokens_dict = {ratio: max_ctx_comp_kv_lens} + + # Build per-token compressed_mask + compressed_mask_tokens = torch.zeros( + num_total_compressed_tokens, dtype=torch.bool, device=DEVICE + ) + offset = 0 + for b in range(bsz): + n = int(num_comp[b].item()) + compressed_mask_tokens[offset : offset + n] = True + slot_size = int(cu_new_comp_kv[b + 1].item() - cu_new_comp_kv[b].item()) + offset += slot_size + compressed_mask_cuda_dict = {ratio: compressed_mask_tokens} + + # Build attention metadata using DeepseekV4CacheManager + metadata = DummyAttentionMetadata( + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_ctx_tokens, + num_tokens=num_ctx_tokens + num_gen_tokens, + kv_cache_manager=self.cache_manager, + block_tables=block_tables, + cu_seq_lens=cu_seq_lens, + cu_new_comp_kv=cu_new_comp_kv_dict, + compressed_position_ids=compressed_position_ids_dict, + compressed_kv_lens=compressed_kv_lens_dict, + past_kv_lens=past_kv_lens_dict, + new_comp_kv_lens_cuda=new_comp_kv_lens_cuda_dict, + num_total_compressed_tokens=num_total_compressed_tokens_dict, + max_ctx_compressed_tokens=max_ctx_compressed_tokens_dict, + compressed_mask_cuda=compressed_mask_cuda_dict, + ) + # kv_lens_cuda_runtime: [num_seqs] total KV length per sequence (past + current) + metadata.kv_lens_cuda_runtime = kv_lens + metadata.cached_token_lens_cuda = past_kv_lens + metadata.num_gen_tokens_per_seq = ( + num_gen_tokens // num_generations if num_generations > 0 else 0 + ) + + # Update kv_cache reference for test compatibility + self._update_kv_cache_reference() + + # Call the compressor forward + result = self.compressor(x=x_flat, metadata=metadata) + + # Update request state and call update_resources after processing + for b, req in enumerate(requests): + token_count = int((start_pos_for_kv[b] + seq_lens[b]).item()) + if is_prefill[b]: + req.context_current_position = token_count + # Call add_new_token for BOTH prefill and generation requests. + req.add_new_token(token_count, 0) + self.cache_manager.update_resources(scheduled_batch) + + # Compressor.forward() returns (kv_comp, scale) tuple. + # For FP8 blockwise (indexer), scale is non-None → return directly. + if isinstance(result, tuple): + kv_comp, scale = result + if scale is not None: + return kv_comp, scale + else: + kv_comp = result + total_outputs = cu_new_comp_kv[-1].item() + if total_outputs == 0: + return None + + # Fused scatter writes postprocessed data to kv_cache but returns raw + # kv_comp. Apply postprocessing inline for test comparison with reference. + if kv_comp is not None and self.kv_cache_dtype == "fp8_pertensor": + # Read FP8 data directly from cache (written by fused kernel) to + # ensure golden cache comparison matches exactly. + all_fp8 = [] + for b in range(bsz): + n = cu_new_comp_kv[b + 1].item() - cu_new_comp_kv[b].item() + if n > 0: + tokens = read_paged_cache_tokens( + self.kv_cache, self.block_offsets, b, n, self.tokens_per_block + ) + all_fp8.append(tokens) + if all_fp8: + kv_fp8 = torch.cat(all_fp8, dim=0).view(torch.float8_e4m3fn) + kv_scale = torch.ones(1, dtype=torch.float32, device=kv_comp.device) + return kv_fp8, kv_scale + return None + elif kv_comp is not None and self.kv_cache_dtype == "default": + # Replay the kernel's fp32-throughout postprocess (RMSNorm + RoPE + # + optional Hadamard) so `out_comp` matches the kernel's cache + # values byte-for-byte. + nope_dim = self.compressor.nope_head_dim + rope_dim = self.compressor.rope_head_dim + kv_proc = kv_comp.clone().float() + var = kv_proc.pow(2).mean(-1, keepdim=True) + kv_proc = kv_proc * torch.rsqrt(var + self.compressor.norm.variance_epsilon) + kv_proc = kv_proc * self.compressor.norm.weight.float() + + pos_ids = metadata.compressed_position_ids_cuda[self.compress_ratio][: kv_proc.shape[0]] + cos_sin = self.compressor.rotary_emb.rotary_cos_sin.float() + half_rope = rope_dim // 2 + cos_v = cos_sin[pos_ids.long(), 0, :] + sin_v = cos_sin[pos_ids.long(), 1, :] + xn = kv_proc[:, :nope_dim] + xp = kv_proc[:, nope_dim:].view(-1, half_rope, 2) + x_even, x_odd = xp[..., 0], xp[..., 1] + xp = torch.stack( + [x_even * cos_v - x_odd * sin_v, x_odd * cos_v + x_even * sin_v], dim=-1 + ).view(-1, rope_dim) + kv_proc = torch.cat([xn, xp], dim=-1) + + if self.compressor.rotate_activation: + kv_proc = rotate_activation(kv_proc) + kv_comp = kv_proc.to(kv_comp.dtype) + + # Split packed output back to per-batch + outputs = [] + for b in range(bsz): + start = cu_new_comp_kv[b].item() + end = cu_new_comp_kv[b + 1].item() + if end > start: + outputs.append(kv_comp[start:end]) + else: + outputs.append(None) + + # If all batches have the same number of outputs, stack them + non_none_outputs = [o for o in outputs if o is not None] + if len(non_none_outputs) == 0: + return None + elif all(o is not None and o.size(0) == outputs[0].size(0) for o in outputs): + return torch.stack(outputs, dim=0) + else: + # Concatenate all non-None outputs along dim=0, then unsqueeze for batch dim + # This handles mixed batches with different compression counts + return torch.cat(non_none_outputs, dim=0).unsqueeze(0) + + def reset_state(self): + """Reset cache manager state for new sequences.""" + # Shutdown existing cache manager to release resources properly + if hasattr(self, "cache_manager") and self.cache_manager is not None: + self.cache_manager.shutdown() + + # Clear active requests and reset request ID counter + self.active_requests.clear() + self.next_request_id = 0 + + # Recreate the cache manager to reset all caches + self.cache_manager = self._create_deepseek_v4_cache_manager(self.compress_ratio) + self._update_kv_cache_reference() + + +def setup_compressors( + compress_ratio: int = 4, + rotate: bool = False, + kv_cache_dtype: str = "default", +): + """Create synced RefCompressor + Compressor with all caches initialized. + + Args: + compress_ratio: Compression ratio (default 4) + rotate: Whether to apply Hadamard rotation + kv_cache_dtype: Cache dtype preset - "default", "fp8_pertensor", + or "fp8_blockwise" + + Returns: + ref: RefCompressor (bf16 reference) + comp: CompressorWrapper (with specified kv_cache_dtype) + """ + args = ModelArgs() + overlap = compress_ratio == 4 + + # For indexer modes with ratio=4, use INDEXER_COMPRESS type and INDEX_HEAD_DIM=128. + # INDEXER_COMPRESS is only registered for sparse layers and uses INDEX_HEAD_DIM=128 + is_indexer = kv_cache_dtype == "fp8_blockwise" + + # Use appropriate head_dim based on is_indexer + target_head_dim = INDEX_HEAD_DIM if is_indexer else HEAD_DIM + + # Reference compressor (bf16) + ref = RefCompressor(args, compress_ratio, target_head_dim, rotate).to(DEVICE) + ref.ape.data.normal_(0, 0.02) + ref.wkv.weight.data.normal_(0, 0.02) + ref.wgate.weight.data.normal_(0, 0.02) + ref.kv_cache = torch.zeros( + MAX_BATCH, MAX_SEQ // compress_ratio, target_head_dim, device=DEVICE, dtype=DTYPE + ) + + # Compressor wrapper with specified kv_cache_dtype + comp = CompressorWrapper( + compress_ratio, rotate, kv_cache_dtype=kv_cache_dtype, is_indexer=is_indexer + ) + + # Copy weights from ref to compressor + coff = 2 if overlap else 1 + comp.compressor.wkv_gate.weight.data[: coff * target_head_dim] = ref.wkv.weight.data.clone() + comp.compressor.wkv_gate.weight.data[coff * target_head_dim :] = ref.wgate.weight.data.clone() + comp.compressor.ape.data.copy_(ref.ape.data) + comp.compressor.norm.weight.data.copy_(ref.norm.weight.data) + + return ref, comp + + +@pytest.fixture(autouse=True) +def seed(): + """Seed RNG for reproducibility and ensure GPU cleanup between tests.""" + torch.manual_seed(42) + torch.cuda.manual_seed(42) + yield + # Force garbage collection so CompressorWrapper.__del__ → cleanup() runs + # before the next test allocates new cache managers. + import gc + + gc.collect() + torch.cuda.empty_cache() + + +# ============================================================================ +# Tests +# ============================================================================ + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 130, 4), + (1, 2, 4), + (1, 64, 4), + (2, 256, 4), + (4, 128, 4), # batch/seqlen variations + (1, 256, 128), + (2, 512, 128), + (16, 234, 128), # ratio=128 coverage + ], +) +def test_prefill(batch, seqlen, ratio): + """Test prefill mode.""" + # rotate=True required: Compressor unconditionally applies Hadamard rotation + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + out_comp = comp.forward(x, 0, freqs) + + assert_similar(out_ref, out_comp) + if out_ref is not None: + num_tokens = out_ref.size(1) + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, out_ref[b : b + 1], f"Prefill ref cache[{b}]") + assert_similar(cached_comp, out_comp[b : b + 1], f"Prefill comp cache[{b}]") + assert_similar(cached_ref, cached_comp, f"Prefill cache parity[{b}]") + + +@pytest.mark.parametrize( + "prefill,steps,batch,ratio", + [ + (128, 8, 1, 4), + (128, 8, 2, 4), + (128, 24, 1, 4), + (128, 4, 1, 128), + ], +) +def test_decode(prefill, steps, batch, ratio): + """Test prefill + decode.""" + # rotate=True required: Compressor unconditionally applies Hadamard rotation + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + # Prefill + x = torch.randn(batch, prefill, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + assert_similar(ref(x, 0, freqs[:prefill]), comp.forward(x, 0, freqs[:prefill]), "Prefill") + + # Decode + for step in range(steps): + pos = prefill + step + x = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + out_ref = ref(x, pos, freqs) + out_comp = comp.forward(x, pos, freqs[pos : pos + 1]) + assert_similar(out_ref, out_comp, f"Decode[{step}]") + if out_ref is not None: + num_tokens = pos // ratio + 1 + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar( + cached_ref[:, -1:], out_ref[b : b + 1], f"Decode cache ref[{b}] step{step}" + ) + assert_similar( + cached_comp[:, -1:], + out_comp[b : b + 1], + f"Decode cache comp[{b}] step{step}", + ) + assert_similar(cached_ref, cached_comp, f"Decode cache parity[{b}] step{step}") + + +def test_varlen_batch(): + """Test variable-length prefill batch, compare with reference.""" + seq_lens = [64, 96, 128] + ratio = 4 + + # Test each sequence independently + for i, slen in enumerate(seq_lens): + # rotate=True required: Compressor unconditionally applies Hadamard rotation + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x = torch.randn(1, slen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + # Reset ref's state for each independent sequence + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + out_ref = ref(x, 0, freqs[:slen]) + out_comp = comp.forward(x, 0, freqs[:slen]) + + assert_similar(out_ref, out_comp, f"Varlen seq{i}") + + +def test_mixed_batch(): + """Test mixed context + generation requests in a single forward call. + + Simulates a realistic mixed batch with: + - 1 context request: 8 tokens, start_pos=0 + - 1 generation request: 1 token, start_pos=127 (triggers compression at 128) + + Generation requests have seqlen=1 and require pre-populated state. + """ + ratio = 4 + # rotate=True required: Compressor unconditionally applies Hadamard rotation + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + # Context request: 8 tokens at start_pos=0 → 2 compressed outputs + ctx_len = 8 + x_ctx = torch.randn(1, ctx_len, DIM, device=DEVICE, dtype=DTYPE) + + # Generation request: 1 token at start_pos=127 → triggers compression at 128 + # (127 + 1) % 4 == 0, so compression is triggered + gen_start_pos = 127 + x_gen = torch.randn(1, 1, DIM, device=DEVICE, dtype=DTYPE) + + # For ref: need to pre-populate state by running prefill of gen_start_pos tokens + # This simulates the generation request having processed gen_start_pos tokens already + x_gen_prefill = torch.randn(1, gen_start_pos, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + # === Reference: run each request separately === + + # Context request on ref (batch 0) + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + out_ref_ctx = ref(x_ctx, 0, freqs[:ctx_len]) + + # Generation request on ref (batch 0, but independent - reset state first) + # Pre-populate state by running prefill of gen_start_pos tokens + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + _ = ref(x_gen_prefill, 0, freqs[:gen_start_pos]) # Sets up state + # Now run the decode token + out_ref_gen = ref(x_gen, gen_start_pos, freqs) + + # Collect non-None outputs + ref_outputs = [o for o in [out_ref_ctx, out_ref_gen] if o is not None] + out_ref = torch.cat(ref_outputs, dim=1) if ref_outputs else None + + # === Compressor: single forward with mixed batch === + # Flatten tokens: [ctx_tokens, gen_token] + x_flat = torch.cat([x_ctx.squeeze(0), x_gen.squeeze(0)], dim=0) + + # Sequence lengths and start positions for each request + seq_lens = torch.tensor([ctx_len, 1], dtype=torch.int32, device=DEVICE) + start_pos_tensor = torch.tensor([0, gen_start_pos], dtype=torch.int32, device=DEVICE) + + # Pre-populate compressor's paged state for generation request (batch idx 1) + # by running prefill on batch 1 first + comp.reset_state() + comp.forward( + x_gen_prefill, + 0, + freqs[:gen_start_pos], + batch_indices=torch.tensor([1], device=DEVICE, dtype=torch.int32), + ) + + # Now run mixed batch + out_comp = comp.forward(x_flat, start_pos_tensor, freqs, seq_lens=seq_lens) + + if out_ref is None: + assert out_comp is None, "Mixed batch: expected no compression output" + else: + assert_similar(out_ref, out_comp, "Mixed batch context+generation single forward") + + +class _FakeCompressorCacheManager: + def __init__(self, head_dim: int, tokens_per_block: int = 4): + self.tokens_per_block = tokens_per_block + self.compressed_block_sizes = {0: tokens_per_block} + self._buffer = torch.empty(1, tokens_per_block * head_dim, device=DEVICE, dtype=DTYPE) + + def get_buffers(self, layer_idx, attn_type): + return self._buffer + + +def _create_small_compressor(kv_cache_dtype: str, is_indexer: bool) -> Compressor: + head_dim = 128 + rope_dim = 64 + mla_params = MLAParams( + hidden_size=16, + qk_rope_head_dim=rope_dim, + qk_nope_head_dim=head_dim - rope_dim, + ) + rope_params = RopeParams( + dim=rope_dim, + theta=ROPE_THETA, + max_positions=16, + original_max_positions=16, + scale_type=RotaryScalingType.none, + ) + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=rope_params, + is_neox=False, + ) + return Compressor( + mla_params=mla_params, + layer_idx=0, + compress_ratio=4, + norm_eps=1e-6, + skip_create_weights_in_init=False, + pos_embd_params=pos_embd_params, + dtype=DTYPE, + kv_cache_dtype=kv_cache_dtype, + is_indexer=is_indexer, + rotate_activation=True, + ).to(DEVICE) + + +def _create_minimal_metadata(compressor: Compressor, total_compressed_tokens: int = 1): + ratio = compressor.compress_ratio + bsz = 1 + block_table = torch.zeros(bsz, 1, device=DEVICE, dtype=torch.int32) + block_tables = {(ratio, attn_type): block_table for attn_type in DeepseekV4AttentionType} + metadata = DummyAttentionMetadata( + num_contexts=1, + num_generations=0, + num_ctx_tokens=ratio, + num_tokens=ratio, + kv_cache_manager=_FakeCompressorCacheManager(compressor.head_dim), + block_tables=block_tables, + cu_seq_lens=torch.tensor([0, ratio], device=DEVICE, dtype=torch.int32), + cu_new_comp_kv={ + ratio: torch.tensor([0, total_compressed_tokens], device=DEVICE, dtype=torch.int32) + }, + compressed_position_ids={ + ratio: torch.zeros(total_compressed_tokens, device=DEVICE, dtype=torch.int32) + }, + compressed_kv_lens={ + ratio: torch.tensor([total_compressed_tokens], device=DEVICE, dtype=torch.int32) + }, + past_kv_lens={ratio: torch.zeros(bsz, device=DEVICE, dtype=torch.int32)}, + new_comp_kv_lens_cuda={ + ratio: torch.tensor([total_compressed_tokens], device=DEVICE, dtype=torch.int32) + }, + num_total_compressed_tokens={ratio: total_compressed_tokens}, + max_ctx_compressed_tokens={ratio: total_compressed_tokens}, + compressed_mask_cuda={ + ratio: torch.ones(total_compressed_tokens, device=DEVICE, dtype=torch.bool) + }, + ) + metadata.kv_lens_cuda_runtime = torch.tensor([ratio], device=DEVICE, dtype=torch.int32) + metadata.cached_token_lens_cuda = torch.zeros(bsz, device=DEVICE, dtype=torch.int32) + return metadata + + +def test_compressor_wkv_gate_uses_checkpoint_dtype(): + compressor = _create_small_compressor(kv_cache_dtype="default", is_indexer=False) + + assert compressor.wkv_gate.weight.dtype == DTYPE + + +def _run_compressor_with_fake_postprocess(monkeypatch, kv_cache_dtype: str, is_indexer: bool): + compressor = _create_small_compressor(kv_cache_dtype, is_indexer) + metadata = _create_minimal_metadata(compressor) + seen = {} + + def fake_prefill_reduction(*args): + seen["kv_score_dtype"] = args[0].dtype + kv_comp = args[6] + kv_comp.fill_(0.25) + + def fake_paged_kv_compress(*args): + raise AssertionError("generation compression path should not run") + + def fake_postprocess_scatter( + kv_comp, + kv_out, + rms_weight, + rms_eps, + rotary_cos_sin, + position_ids, + nope_head_dim, + rope_head_dim, + kv_cache, + num_comp_tokens, + cu_new_comp_kv, + start_pos, + block_table, + compressed_mask, + tokens_per_block, + cache_dtype, + rotate_activation, + quant_output, + scale_output, + ): + seen["kv_out"] = kv_out + seen["quant_output"] = quant_output + seen["scale_output"] = scale_output + seen["cache_dtype"] = cache_dtype + if kv_out is not None: + kv_out.fill_(0.5) + if quant_output is not None: + quant_output.fill_(0x38) + if scale_output.dtype.is_floating_point: + scale_output.fill_(2.0) + else: + scale_output.fill_(0x7F) + + monkeypatch.setattr( + torch.ops.trtllm, + "compressor_prefill_reduction", + fake_prefill_reduction, + raising=False, + ) + monkeypatch.setattr( + torch.ops.trtllm, + "compressor_paged_kv_compress", + fake_paged_kv_compress, + raising=False, + ) + monkeypatch.setattr( + torch.ops.trtllm, + "compressor_postprocess_scatter", + fake_postprocess_scatter, + raising=False, + ) + + x = torch.randn(compressor.compress_ratio, 16, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + output = compressor(x, metadata) + return output, seen + + +def test_main_compressor_does_not_materialize_postprocess_output(monkeypatch): + """The fused kernel writes main compressor output directly to paged cache. + + Re-materializing kv_out here and doing a Python scatter afterwards defeats + the fused postprocess/scatter kernel and regresses end-to-end performance. + """ + output, seen = _run_compressor_with_fake_postprocess( + monkeypatch, kv_cache_dtype="default", is_indexer=False + ) + + assert seen["cache_dtype"] == int(KVCacheDtype.NONE) + assert seen["kv_score_dtype"] == torch.bfloat16 + assert seen["kv_out"] is None + assert seen["quant_output"] is None + assert seen["scale_output"] is None + assert output[0].shape == (1, 128) + assert output[1] is None + + +@pytest.mark.parametrize( + "kv_cache_dtype,cache_dtype,expected_dtype,expected_quant_shape,expected_scale_shape", + [ + ( + "fp8_blockwise", + KVCacheDtype.FP8_BLOCKWISE, + torch.float8_e4m3fn, + (1, 128), + (1, 1), + ), + ( + "mxfp4", + KVCacheDtype.MXFP4_BLOCKWISE, + torch.float4_e2m1fn_x2, + (1, 64), + (1, 4), + ), + ], +) +def test_indexer_returns_fused_quant_outputs( + monkeypatch, + kv_cache_dtype, + cache_dtype, + expected_dtype, + expected_quant_shape, + expected_scale_shape, +): + output, seen = _run_compressor_with_fake_postprocess( + monkeypatch, kv_cache_dtype=kv_cache_dtype, is_indexer=True + ) + + quant_output, scale_output = output + assert seen["cache_dtype"] == int(cache_dtype) + assert seen["kv_out"] is None + assert quant_output.dtype == expected_dtype + assert quant_output.shape == expected_quant_shape + assert scale_output.shape == expected_scale_shape + assert torch.equal(quant_output.view(torch.uint8), torch.full_like(seen["quant_output"], 0x38)) + if scale_output.dtype.is_floating_point: + assert torch.equal(scale_output, torch.full_like(scale_output, 2.0)) + else: + assert torch.equal(scale_output, torch.full_like(scale_output, 0x7F)) + + +# ============================================================================ +# FP8 Blockwise Quantization Tests +# ============================================================================ + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 64, 4), + (4, 256, 4), + (16, 256, 4), + (16, 512, 4), + ], +) +def test_fp8_blockwise_compressor(batch, seqlen, ratio): + """Test FP8 blockwise Compressor against RefCompressor (bf16 reference). + + FP8 blockwise uses INDEXER_COMPRESS cache type which has INDEX_HEAD_DIM=128. + This test validates FP8 quantization and cache scatter for the indexer path. + """ + num_compressed = seqlen // ratio + + ref, comp = setup_compressors(ratio, rotate=True, kv_cache_dtype="fp8_blockwise") + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + out_comp = comp.forward(x, 0, freqs) + + assert isinstance(out_comp, tuple), f"Expected tuple, got {type(out_comp)}" + kv_fp8, kv_scale = out_comp + + # Verify shapes - INDEXER_COMPRESS uses INDEX_HEAD_DIM=128 + total_comp_tokens = batch * num_compressed + num_scale_blocks = (INDEX_HEAD_DIM + 127) // 128 + assert kv_fp8.shape == (total_comp_tokens, INDEX_HEAD_DIM), ( + f"Expected shape {(total_comp_tokens, INDEX_HEAD_DIM)}, got {kv_fp8.shape}" + ) + assert kv_scale.shape == (total_comp_tokens, num_scale_blocks), ( + f"Expected scale shape {(total_comp_tokens, num_scale_blocks)}, got {kv_scale.shape}" + ) + + # Verify scales are positive and valid + assert (kv_scale > 0).all(), "All scales should be positive" + + # Compare dequantized FP8 with RefCompressor output (both use INDEX_HEAD_DIM=128) + assert_fp8_similar( + out_comp, out_ref.view(-1, INDEX_HEAD_DIM), "fp8_blockwise", "FP8 vs RefCompressor" + ) + + # Verify cache scatter + golden_cache = build_fp8_golden_cache( + kv_fp8, + kv_scale, + comp.kv_cache.shape, + comp.block_offsets, + batch, + num_compressed, + comp.tokens_per_block, + "fp8_blockwise", + head_dim=INDEX_HEAD_DIM, + ) + assert_fp8_cache_match( + comp.kv_cache, + golden_cache, + "fp8_blockwise", + "Blockwise cache layout", + head_dim=INDEX_HEAD_DIM, + ) + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 256, 4), + (1, 256, 128), + (2, 512, 128), # ratio=128 coverage + ], +) +def test_fp8_pertensor_compressor(batch, seqlen, ratio): + """Test per-tensor FP8 Compressor against RefCompressor (bf16 reference).""" + num_compressed = seqlen // ratio + + ref, comp = setup_compressors(ratio, rotate=True, kv_cache_dtype="fp8_pertensor") + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + out_comp = comp.forward(x, 0, freqs) + + assert isinstance(out_comp, tuple), f"Expected tuple, got {type(out_comp)}" + kv_fp8, kv_scale = out_comp + + # Verify shapes + total_comp_tokens = batch * num_compressed + assert kv_fp8.shape == (total_comp_tokens, HEAD_DIM) + assert kv_scale.numel() == 1, f"Per-tensor scale should be scalar, got {kv_scale.shape}" + assert kv_scale.item() > 0, "Scale should be positive" + + # Compare dequantized FP8 with reference + assert_fp8_similar( + out_comp, out_ref.view(-1, HEAD_DIM), "fp8_pertensor", "FP8 vs RefCompressor" + ) + + # Verify scale is fixed at 1.0 (compressor uses static scale following trtllm.py convention) + assert kv_scale.item() == 1.0, ( + f"Per-tensor scale should be fixed at 1.0, got {kv_scale.item():.6f}" + ) + + # Verify cache scatter (cache is FP8 dtype, same as compressor output) + golden_cache = build_fp8_golden_cache( + kv_fp8, + kv_scale, + comp.kv_cache.shape, + comp.block_offsets, + batch, + num_compressed, + comp.tokens_per_block, + "fp8_pertensor", + ) + assert_fp8_cache_match(comp.kv_cache, golden_cache, "fp8_pertensor", "Per-tensor cache layout") + + +# ============================================================================ +# Fused Kernel Tests (RMSNorm + RoPE + Hadamard + Scatter) +# +# The fused CUDA kernel matches the reference model.py pipeline: RMSNorm, +# RoPE, and Hadamard all in bf16 precision (via toBf16 round-trips in CUDA). +# This means fused and unfused paths should produce nearly identical results. +# ============================================================================ + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 130, 4), + (1, 64, 4), + (2, 256, 4), + (4, 128, 4), + (1, 256, 128), + (2, 512, 128), + ], +) +def test_fused_prefill(batch, seqlen, ratio): + """Test fused prefill cache against unfused bf16 reference cache. + + Both paths now use bf16 Hadamard so results should match closely. + """ + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + comp.forward(x, 0, freqs) + + if out_ref is not None: + num_tokens = out_ref.size(1) + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"Fused prefill cache[{b}]") + + +@pytest.mark.parametrize( + "prefill,steps,batch,ratio", + [ + (128, 8, 1, 4), + (128, 8, 2, 4), + (128, 24, 1, 4), + (128, 4, 1, 128), + ], +) +def test_fused_decode(prefill, steps, batch, ratio): + """Test fused prefill + decode cache against bf16 reference.""" + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x = torch.randn(batch, prefill, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + ref(x, 0, freqs[:prefill]) + comp.forward(x, 0, freqs[:prefill]) + + for step in range(steps): + pos = prefill + step + x = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + out_ref = ref(x, pos, freqs) + comp.forward(x, pos, freqs[pos : pos + 1]) + if out_ref is not None: + num_tokens = pos // ratio + 1 + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"Fused decode cache[{b}] step{step}") + + +def test_fused_mixed_batch(): + """Test fused kernel with mixed context + generation batch.""" + ratio = 4 + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + ctx_len = 8 + x_ctx = torch.randn(1, ctx_len, DIM, device=DEVICE, dtype=DTYPE) + gen_start_pos = 127 + x_gen = torch.randn(1, 1, DIM, device=DEVICE, dtype=DTYPE) + x_gen_prefill = torch.randn(1, gen_start_pos, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + ref(x_ctx, 0, freqs[:ctx_len]) + + # Save ctx cache before gen overwrites positions 0..1 + num_ctx_comp = ctx_len // ratio + cached_ref_ctx = ref.kv_cache[0:1, :num_ctx_comp].clone() + + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + _ = ref(x_gen_prefill, 0, freqs[:gen_start_pos]) + ref(x_gen, gen_start_pos, freqs) + + x_flat = torch.cat([x_ctx.squeeze(0), x_gen.squeeze(0)], dim=0) + seq_lens = torch.tensor([ctx_len, 1], dtype=torch.int32, device=DEVICE) + start_pos_tensor = torch.tensor([0, gen_start_pos], dtype=torch.int32, device=DEVICE) + + comp.reset_state() + comp.forward( + x_gen_prefill, + 0, + freqs[:gen_start_pos], + batch_indices=torch.tensor([1], device=DEVICE, dtype=torch.int32), + ) + comp.forward(x_flat, start_pos_tensor, freqs, seq_lens=seq_lens) + + # Compare context request cache (first 2 compressed tokens) + if num_ctx_comp > 0: + cached_ref = cached_ref_ctx + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, 0, num_ctx_comp, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, "Fused mixed ctx cache") + + +# ============================================================================ +# Tests with rotate_activation=False (no Hadamard transform) +# ============================================================================ + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 130, 4), + (4, 128, 4), + (1, 256, 128), + (2, 512, 128), + ], +) +def test_prefill_no_rotate(batch, seqlen, ratio): + """Test prefill mode with rotate_activation=False (Hadamard skipped).""" + ref, comp = setup_compressors(ratio, rotate=False) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + out_comp = comp.forward(x, 0, freqs) + + assert_similar(out_ref, out_comp) + if out_ref is not None: + num_tokens = out_ref.size(1) + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"Prefill no-rotate cache[{b}]") + + +@pytest.mark.parametrize( + "prefill,steps,batch,ratio", + [ + (128, 8, 1, 4), + (128, 8, 2, 4), + (128, 24, 1, 4), + (128, 4, 1, 128), + ], +) +def test_decode_no_rotate(prefill, steps, batch, ratio): + """Test prefill + decode with rotate_activation=False.""" + ref, comp = setup_compressors(ratio, rotate=False) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x = torch.randn(batch, prefill, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + assert_similar( + ref(x, 0, freqs[:prefill]), comp.forward(x, 0, freqs[:prefill]), "Prefill no-rotate" + ) + + for step in range(steps): + pos = prefill + step + x = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + out_ref = ref(x, pos, freqs) + out_comp = comp.forward(x, pos, freqs[pos : pos + 1]) + assert_similar(out_ref, out_comp, f"Decode no-rotate[{step}]") + if out_ref is not None: + num_tokens = pos // ratio + 1 + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar( + cached_ref, cached_comp, f"Decode no-rotate cache[{b}] step{step}" + ) + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 130, 4), + (2, 256, 4), + (1, 256, 128), + (2, 512, 128), + ], +) +def test_fused_prefill_no_rotate(batch, seqlen, ratio): + """Test fused prefill cache with rotate_activation=False.""" + ref, comp = setup_compressors(ratio, rotate=False) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + comp.forward(x, 0, freqs) + + if out_ref is not None: + num_tokens = out_ref.size(1) + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"Fused prefill no-rotate cache[{b}]") + + +@pytest.mark.parametrize( + "prefill,steps,batch,ratio", + [ + (128, 8, 1, 4), + (128, 8, 2, 4), + (128, 24, 1, 4), + (128, 4, 1, 128), + ], +) +def test_fused_decode_no_rotate(prefill, steps, batch, ratio): + """Test fused prefill + decode cache with rotate_activation=False.""" + ref, comp = setup_compressors(ratio, rotate=False) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x = torch.randn(batch, prefill, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + ref(x, 0, freqs[:prefill]) + comp.forward(x, 0, freqs[:prefill]) + + for step in range(steps): + pos = prefill + step + x = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + out_ref = ref(x, pos, freqs) + comp.forward(x, pos, freqs[pos : pos + 1]) + if out_ref is not None: + num_tokens = pos // ratio + 1 + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar( + cached_ref, cached_comp, f"Fused decode no-rotate cache[{b}] step{step}" + ) + + +@pytest.mark.parametrize( + "prefill_len, decode_steps, batch, ratio, rotate", + [ + # ratio=4 (overlap mode) + (1, 8, 1, 4, True), + (2, 8, 2, 4, True), + (3, 8, 1, 4, True), + (1, 4, 1, 4, False), + (3, 12, 2, 4, False), + # ratio=128 (non-overlap): these should pass + (64, 128, 1, 128, True), # half-ratio prefill, decode rest + (1, 128, 1, 128, True), # minimal prefill + (127, 4, 1, 128, True), # one token short of compression at prefill + ], +) +def test_short_prefill_then_decode(prefill_len, decode_steps, batch, ratio, rotate): + """Test prefill with fewer tokens than compress_ratio, then decode until compression. + + When prefill_len < compress_ratio, the prefill produces no compressed outputs; + all tokens are saved as remainder state. Subsequent decode tokens accumulate + in the state until compress_ratio is reached, at which point compression fires. + This tests the state handoff from prefill remainder to decode accumulation. + """ + ref, comp = setup_compressors(ratio, rotate=rotate) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x_prefill = torch.randn(batch, prefill_len, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + # Prefill: should produce no compressed tokens (prefill_len < ratio) + out_ref = ref(x_prefill, 0, freqs[:prefill_len]) + out_comp = comp.forward( + x_prefill, + 0, + freqs[:prefill_len], + is_prefill=torch.ones(batch, dtype=torch.bool, device=DEVICE), + ) + + if prefill_len < ratio: + # No compression expected from prefill + assert out_ref is None, ( + f"Ref should produce no output for prefill_len={prefill_len} < ratio={ratio}" + ) + + # Decode: step one token at a time + for step in range(decode_steps): + pos = prefill_len + step + x_decode = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + out_ref = ref(x_decode, pos, freqs) + out_comp = comp.forward(x_decode, pos, freqs[pos : pos + 1]) + + should_compress = (pos + 1) % ratio == 0 + if should_compress: + assert out_ref is not None, f"Ref should compress at step {step} (pos={pos})" + assert_similar(out_ref, out_comp, f"Short prefill decode step {step}") + + # Verify cache parity + num_comp_tokens = (pos + 1) // ratio + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_comp_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, + comp.block_offsets, + b, + num_comp_tokens, + comp.tokens_per_block, + ).unsqueeze(0) + assert_similar( + cached_ref, + cached_comp, + f"Short prefill cache parity[{b}] step{step}", + ) + else: + assert out_ref is None, f"Ref should NOT compress at step {step} (pos={pos})" + finally: + comp.cleanup() + + +MTP_DECODE_CASES = [ + # No prior compressed output: decode spans the first compression boundary. + (1, 8, 1, 4, 2), + (3, 8, 1, 4, 3), + # Prefill ends exactly on or just after a compression boundary. + (4, 8, 1, 4, 2), + (5, 8, 1, 4, 3), + # Large absolute positions: exercise the same boundary cases after many windows. + (127, 8, 1, 4, 2), + (128, 8, 2, 4, 3), +] + + +@pytest.mark.parametrize("prefill_len, decode_steps, batch, ratio, next_n", MTP_DECODE_CASES) +def test_mtp_decode_overlap(prefill_len, decode_steps, batch, ratio, next_n): + """MTP decode: ref(combined seqlen=n) == ref(chunked seqlen=1 each). + + Verifies that RefCompressor called once with all next_n tokens produces + the same output as calling it one token at a time — i.e., the else-branch + for-loop is equivalent to sequential single-token calls. + """ + ref_combined, comp = setup_compressors(ratio, rotate=True) + try: + # Build ref_chunked with identical weights + args = ModelArgs( + dim=DIM, + head_dim=HEAD_DIM, + rope_head_dim=ROPE_DIM, + max_seq_len=MAX_SEQ, + max_batch_size=MAX_BATCH, + ) + ref_chunked = RefCompressor(args, compress_ratio=ratio, head_dim=HEAD_DIM, rotate=True).to( + DEVICE + ) + ref_chunked.wkv.weight.data.copy_(ref_combined.wkv.weight.data) + ref_chunked.wgate.weight.data.copy_(ref_combined.wgate.weight.data) + ref_chunked.ape.data.copy_(ref_combined.ape.data) + ref_chunked.norm.weight.data.copy_(ref_combined.norm.weight.data) + ref_chunked.kv_cache = torch.zeros_like(ref_combined.kv_cache) + + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x_prefill = torch.randn(batch, prefill_len, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + ref_combined(x_prefill, 0, freqs[:prefill_len]) + ref_chunked(x_prefill, 0, freqs[:prefill_len]) + + step = 0 + while step < decode_steps: + n = min(next_n, decode_steps - step) + pos = prefill_len + step + x_decode = torch.randn(batch, n, DIM, device=DEVICE, dtype=DTYPE) + + # Combined: one forward call with all n tokens + out_combined = ref_combined(x_decode, pos, freqs) + + # Chunked: one token at a time + last_chunked_out = None + for t in range(n): + out_t = ref_chunked(x_decode[:, t : t + 1], pos + t, freqs) + if out_t is not None: + last_chunked_out = out_t + + if out_combined is not None: + assert last_chunked_out is not None, ( + f"MTP step {step}: chunked produced no output but combined compressed" + ) + assert_similar(out_combined, last_chunked_out, f"MTP decode step {step}") + for b in range(batch): + assert_similar( + ref_combined.kv_cache[b : b + 1], + ref_chunked.kv_cache[b : b + 1], + f"MTP cache[{b}] step {step}", + ) + else: + assert last_chunked_out is None, ( + f"MTP step {step}: chunked produced unexpected output" + ) + + step += n + finally: + comp.cleanup() + + +@pytest.mark.parametrize("prefill_len, decode_steps, batch, ratio, next_n", MTP_DECODE_CASES) +def test_mtp_decode_overlap_module(prefill_len, decode_steps, batch, ratio, next_n): + """MTP decode through CompressorWrapper requires explicit generation mode.""" + ref, comp = setup_compressors(ratio, rotate=True) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x_prefill = torch.randn(batch, prefill_len, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + ref(x_prefill, 0, freqs[:prefill_len]) + comp.forward( + x_prefill, + 0, + freqs[:prefill_len], + is_prefill=torch.ones(batch, dtype=torch.bool, device=DEVICE), + ) + + step = 0 + while step < decode_steps: + n = min(next_n, decode_steps - step) + pos = prefill_len + step + x_decode = torch.randn(batch, n, DIM, device=DEVICE, dtype=DTYPE) + + out_ref = ref(x_decode, pos, freqs) + out_comp = comp.forward( + x_decode, + pos, + freqs, + is_prefill=torch.zeros(batch, dtype=torch.bool, device=DEVICE), + ) + + if out_ref is None: + assert out_comp is None, ( + f"MTP module step {step}: expected no compressed output" + ) + else: + assert out_comp is not None, ( + f"MTP module step {step}: wrapper produced no output for generation" + ) + assert_similar(out_ref, out_comp, f"MTP module decode step {step}") + + step += n + finally: + comp.cleanup() + + +@pytest.mark.parametrize( + "batch, ratio, rotate", + [ + (1, 4, True), + (2, 4, True), + (1, 4, False), + (1, 128, True), + ], +) +def test_prefill_exact_ratio(batch, ratio, rotate): + """Prefill with seqlen == compress_ratio: exactly 1 full chunk, no remainder tokens. + + In overlap mode (ratio=4) the single output has no predecessor chunk so the + overlap first-half should be zero-weighted. In non-overlap mode (ratio=128) + this is a standard single-chunk prefill. Both output and cache are verified. + """ + ref, comp = setup_compressors(ratio, rotate=rotate) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x = torch.randn(batch, ratio, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs[:ratio]) + out_comp = comp.forward(x, 0, freqs[:ratio]) + + assert out_ref is not None, "Expected compression for seqlen==ratio" + assert out_comp is not None, "Expected compression for seqlen==ratio" + assert_similar(out_ref, out_comp, "exact_ratio prefill output") + + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :1] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, 1, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"exact_ratio cache[{b}]") + finally: + comp.cleanup() + + +@pytest.mark.parametrize( + "seq_lens_list, ratio, rotate", + [ + ([3, 8, 5], 4, True), # short / long / medium — varied output counts + ([1, 4, 7], 4, False), # minimal, exact-ratio, one-beyond-exact + ([64, 3, 128], 128, True), # ratio=128, mixed short / long + ], +) +def test_mixed_seqlen_contexts(seq_lens_list, ratio, rotate): + """Prefill batch with variable seqlens: some < ratio (no compressed output), some >= ratio. + + Verifies that zero-output sequences do not corrupt the cu_new_comp_kv layout + or the compressed token buffer for neighbouring sequences. Each sequence is + also run through the reference independently so that the paged cache content + can be compared token-by-token. + """ + batch = len(seq_lens_list) + max_sl = max(seq_lens_list) + + ref, comp = setup_compressors(ratio, rotate=rotate) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + xs = [torch.randn(1, sl, DIM, device=DEVICE, dtype=DTYPE) for sl in seq_lens_list] + + with torch.no_grad(): + # ---- Reference: run each sequence independently, save compressed cache ---- + ref_caches = {} + for i, sl in enumerate(seq_lens_list): + n_out = sl // ratio + if n_out > 0: + # Reset ref state so independent runs don't interfere + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + ref(xs[i], 0, freqs[:sl]) + # ref writes to kv_cache[0]; save before next independent run + ref_caches[i] = ref.kv_cache[0, :n_out].clone() + + # ---- Wrapper: run all sequences in one variable-length prefill call ---- + seq_lens_t = torch.tensor(seq_lens_list, dtype=torch.int32, device=DEVICE) + start_pos_t = torch.zeros(batch, dtype=torch.int32, device=DEVICE) + + # Build flat (non-padded) 2D token tensor. A padded 3D tensor would + # interleave padding zeros with real tokens after view(-1, DIM), so we + # concatenate actual token rows directly instead. + x_flat_input = torch.cat([xs[i][0] for i in range(batch)], dim=0) + + comp.forward(x_flat_input, start_pos_t, freqs[:max_sl], seq_lens=seq_lens_t) + + # ---- Compare per-sequence compressed caches ---- + for rank, sl in enumerate(seq_lens_list): + n_out = sl // ratio + if n_out > 0 and rank in ref_caches: + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, rank, n_out, comp.tokens_per_block + ) + assert_similar( + ref_caches[rank].unsqueeze(0), + cached_comp.unsqueeze(0), + f"mixed_seqlen cache seq[{rank}] sl={sl}", + ) + + # ---- Decode from remainder state for sequences with leftover tokens ---- + # Find the first sequence that has a non-zero remainder and run it to + # the next compression point; this validates state hand-off. + for rank, sl in enumerate(seq_lens_list): + remainder = sl % ratio + if remainder == 0: + continue + steps_to_compress = ratio - remainder + for step in range(steps_to_compress): + pos = sl + step + x_dec = torch.randn(1, 1, DIM, device=DEVICE, dtype=DTYPE) + pos_tensor = torch.tensor([pos], dtype=torch.int32, device=DEVICE) + seq_lens_dec = torch.ones(1, dtype=torch.int32, device=DEVICE) + out = comp.forward( + x_dec, pos_tensor, freqs[pos : pos + 1], seq_lens=seq_lens_dec + ) + should_compress = (pos + 1) % ratio == 0 + if should_compress: + assert out is not None, ( + f"mixed_seqlen seq[{rank}]: expected compression at pos={pos}" + ) + else: + assert out is None, ( + f"mixed_seqlen seq[{rank}]: unexpected output at pos={pos}" + ) + break # Only test the first remainder sequence for brevity + finally: + comp.cleanup() + + +@pytest.mark.parametrize( + "total_seqlen, split_pos, batch, ratio, rotate", + [ + # overlap mode (ratio=4): aligned split + (128, 20, 1, 4, True), + (128, 48, 2, 4, True), + (128, 20, 1, 4, False), + # overlap mode: unaligned split (split_pos % ratio != 0) + (128, 5, 1, 4, True), + (128, 6, 1, 4, True), + # non-overlap mode (ratio=128): aligned split + (256, 128, 1, 128, True), + # non-overlap mode: unaligned split + (256, 50, 1, 128, True), + ], +) +def test_chunked_prefill_ref(total_seqlen, split_pos, batch, ratio, rotate): + """Validate RefCompressor full prefill == RefCompressor two-step (initial + chunked). + + Pure-Python reference test — no CUDA kernels or cache managers. + Ensures the RefCompressor's sequential token-by-token path produces + identical kv_cache contents as the bulk prefill path. + """ + import copy + + args = ModelArgs() + + ref_full = RefCompressor(args, ratio, HEAD_DIM, rotate).to(DEVICE) + ref_full.ape.data.normal_(0, 0.02) + ref_full.wkv.weight.data.normal_(0, 0.02) + ref_full.wgate.weight.data.normal_(0, 0.02) + ref_full.kv_cache = torch.zeros( + MAX_BATCH, MAX_SEQ // ratio, HEAD_DIM, device=DEVICE, dtype=DTYPE + ) + + ref_split = copy.deepcopy(ref_full) + + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x = torch.randn(batch, total_seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_full = ref_full(x, 0, freqs[:total_seqlen]) + out_1 = ref_split(x[:, :split_pos], 0, freqs[:split_pos]) + out_2 = ref_split(x[:, split_pos:], split_pos, freqs) + + parts = [p for p in [out_1, out_2] if p is not None] + if out_full is None: + assert len(parts) == 0, "Split produced output but full did not" + return + assert len(parts) > 0, "Full produced output but split did not" + combined = torch.cat(parts, dim=1) + + assert_similar(out_full, combined, "Chunked prefill ref: full vs split") + + n_comp_full = out_full.size(1) + cache_full = ref_full.kv_cache[:batch, :n_comp_full] + cache_split = ref_split.kv_cache[:batch, :n_comp_full] + assert_similar(cache_full, cache_split, "Chunked prefill ref: cache parity") + + +@pytest.mark.parametrize( + "total_seqlen, split_pos, batch, ratio", + [ + (128, 20, 1, 4), + (128, 48, 2, 4), + (256, 128, 1, 128), + ], +) +def test_chunked_prefill_module(total_seqlen, split_pos, batch, ratio): + """Validate CompressorWrapper two-step matches RefCompressor full prefill. + + Step 1: RefCompressor processes the full sequence in one call (ground truth). + Step 2: CompressorWrapper processes the sequence in two calls (initial + chunked). + The concatenated CompressorWrapper outputs must match the RefCompressor output. + """ + ref, comp = setup_compressors(ratio, rotate=True) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x = torch.randn(batch, total_seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs[:total_seqlen]) + + with torch.no_grad(): + out_1 = comp.forward(x[:, :split_pos], 0, freqs[:split_pos]) + out_2 = comp.forward( + x[:, split_pos:], + split_pos, + freqs[:total_seqlen], + is_prefill=torch.ones(batch, dtype=torch.bool, device=DEVICE), + ) + + parts = [p for p in [out_1, out_2] if p is not None] + if out_ref is None: + assert len(parts) == 0, "Comp produced output but ref did not" + return + assert len(parts) > 0, "Ref produced output but comp did not" + combined = ( + torch.cat(parts, dim=1) + if all(p.dim() == out_ref.dim() for p in parts) + else torch.cat(parts, dim=0).unsqueeze(0) + ) + + assert_similar(out_ref, combined, "Chunked prefill module: ref vs comp") + finally: + comp.cleanup() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py new file mode 100644 index 000000000000..42a5b8b76a09 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -0,0 +1,949 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from typing import Dict, List, Optional, Tuple + +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( + DeepseekV4AttentionType, +) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._utils import binding_to_torch_dtype +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX + +_RequestCache = Dict[ + Tuple[int, DeepseekV4AttentionType], # (layer index, attention type) + Tuple[torch.Tensor, torch.Tensor | None], # (values tensor, scales tensor) +] + + +def test_cache_size_estimation_uses_model_attention_layer_count(): + class FakeModelConfig: + sparse_attention_config = SimpleNamespace( + index_head_dim=128, + compress_ratios=[1, 4, 1, 128], + indexer_k_dtype="fp8", + ) + pretrained_config = SimpleNamespace( + kv_lora_rank=512, + qk_rope_head_dim=64, + ) + quant_config = None + + def get_num_attention_layers(self) -> int: + return len(self.sparse_attention_config.compress_ratios) + + size_per_token = DeepseekV4CacheManager.get_cache_size_per_token( + FakeModelConfig(), + Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + is_disagg=True, + ) + + assert size_per_token > 0 + + +def _view_fp8_as_uint8(buffer: torch.Tensor) -> torch.Tensor: + """View an FP8 buffer as uint8. Non-FP8 buffers are returned as-is.""" + if buffer.dtype == torch.float8_e4m3fn: + return buffer.view(torch.uint8) + return buffer + + +@skip_pre_blackwell +@pytest.mark.skip_less_device_memory(80000) +class TestDeepseekV4CacheManager: + # deepseek_v4 specific param + head_dim = 512 + index_head_dim = 128 + window_size = 128 + vocab_size = 129280 + sparse_layer_ratio = 4 + overlap_compress_layer_ratio = 4 + + # indexer quantization config + indexer_dtype = DataType.FP8 + indexer_scale_dtype = DataType.FLOAT + indexer_quant_block_size = 128 + + # cache manager specific param + tokens_per_block = 128 + + def _is_compress_layer(self, compress_ratio: int) -> bool: + """Check if a layer uses compression based on its compress ratio. + + Args: + compress_ratio: The compression ratio for the layer + + Returns: + True if the layer uses compression (ratio > 1) + """ + return compress_ratio > 1 + + def _is_sparse_layer(self, compress_ratio: int) -> bool: + """Check if a layer uses sparse attention based on its compress ratio. + + Args: + compress_ratio: The compression ratio for the layer + + Returns: + True if the layer uses sparse attention (ratio == 4) + """ + return compress_ratio == self.sparse_layer_ratio + + def _is_overlap_compressor(self, compress_ratio: int) -> bool: + """Check if a layer uses overlap compressor based on its compress ratio. + + Args: + compress_ratio: The compression ratio for the layer + + Returns: + True if the layer uses overlap compressor (ratio == 4) + """ + return compress_ratio == self.overlap_compress_layer_ratio + + def _get_window_size(self, compress_ratio: int, attn_type: DeepseekV4AttentionType) -> int: + """Get the window size for a layer based on its compress ratio and attention type. + + Args: + compress_ratio: The compression ratio for the layer + attn_type: The attention type + + Returns: + The window size for the layer + """ + state_factor = 2 if self._is_overlap_compressor(compress_ratio) else 1 + if attn_type == DeepseekV4AttentionType.SWA: + return self.window_size + elif attn_type in [ + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ]: + return state_factor * compress_ratio + elif attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + return None + + def _create_deepseek_v4_cache_manager( + self, + tokens_per_block: int, + max_batch_size: int, + max_seq_len: int, + compress_ratios: List[int], + dtype: DataType, + compressor_dtype: DataType, + max_input_len: Optional[int] = None, + ) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]: + """Helper to create a DeepseekV4CacheManager for testing.""" + + # Create sparse attention config + sparse_attn_config = DeepSeekV4SparseAttentionConfig( + index_head_dim=self.index_head_dim, + window_size=self.window_size, + compress_ratios=compress_ratios, + ) + + # Create KV cache config + if max_input_len is None: + max_input_len = max_seq_len + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + max_tokens=max_seq_len * max_batch_size, + event_buffer_max_size=0, + ) + + # Create mapping (single GPU, no parallelism) + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + # Create cache manager + cache_manager = DeepseekV4CacheManager( + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=len(compress_ratios), + num_kv_heads=1, + head_dim=self.head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + max_input_len=max_input_len, + mapping=mapping, + dtype=dtype, + compressor_dtype=compressor_dtype, + vocab_size=self.vocab_size, + max_num_tokens=max_batch_size * (max_input_len + 1), + sparse_attn_config=sparse_attn_config, + ) + + return cache_manager, sparse_attn_config + + def _create_request(self, request_id: int, prompt_len: int) -> LlmRequest: + """Helper to create a test LlmRequest. + + Args: + request_id: Unique request identifier + prompt_len: Prompt length (number of tokens) + + Returns: + LlmRequest instance + """ + input_tokens = list(range(prompt_len)) + request = LlmRequest( + request_id=request_id, + max_new_tokens=1024, + input_tokens=input_tokens, + sampling_config=SamplingConfig(), + is_streaming=False, + ) + + return request + + def _rand_tensor( + self, + shape: Tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + if dtype in (torch.uint8, torch.float8_e4m3fn): + # Use uint8 for both uint8 and FP8 (same 1-byte layout) + return torch.randint(0, 255, shape, dtype=torch.uint8, device=device) + else: + return torch.randn(shape, dtype=dtype, device=device) * 1000.0 + + def _create_random_cache( + self, + seq_len: int, + head_dim: int, + sparse_attn_config: DeepSeekV4SparseAttentionConfig, + dtype: torch.dtype, + compressor_dtype: torch.dtype, + device: torch.device | None = None, + ) -> _RequestCache: + """Helper to create random cache values for all layers and attention types. + + Args: + seq_len: Sequence length + head_dim: Head dimension for regular attention + sparse_attn_config: Sparse attention configuration + + Returns: + Dictionary mapping (layer_idx, attn_type) to (values, scales) tuples. + scales is None for non-quantized attention types. + """ + device = device or torch.device("cuda") + cache: _RequestCache = {} + + for layer, ratio in enumerate(sparse_attn_config.compress_ratios): + is_overlap = self._is_overlap_compressor(ratio) + + cache[layer, DeepseekV4AttentionType.SWA] = ( + self._rand_tensor((seq_len, head_dim), dtype, device), + None, + ) + + if self._is_compress_layer(ratio): + compressor_dim = 2 * head_dim if is_overlap else head_dim + cache[layer, DeepseekV4AttentionType.COMPRESS] = ( + self._rand_tensor((seq_len // ratio, head_dim), dtype, device), + None, + ) + cache[layer, DeepseekV4AttentionType.COMPRESSOR_STATE] = ( + self._rand_tensor((seq_len, compressor_dim), compressor_dtype, device), + None, + ) + cache[layer, DeepseekV4AttentionType.COMPRESSOR_SCORE] = ( + self._rand_tensor((seq_len, compressor_dim), compressor_dtype, device), + None, + ) + + if self._is_sparse_layer(ratio): + # indexer kv cache is blockwise FP8 quantized + indexer_dim = sparse_attn_config.index_head_dim + indexer_num_tokens = seq_len // ratio + num_scales = indexer_dim // self.indexer_quant_block_size + indexer_values = self._rand_tensor( + (indexer_num_tokens, indexer_dim), torch.uint8, device + ) + indexer_scales = self._rand_tensor( + (indexer_num_tokens, num_scales), torch.float32, device + ) + cache[layer, DeepseekV4AttentionType.INDEXER_COMPRESS] = ( + indexer_values, + indexer_scales, + ) + + indexer_compressor_dim = 2 * indexer_dim if is_overlap else indexer_dim + cache[layer, DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE] = ( + self._rand_tensor((seq_len, indexer_compressor_dim), compressor_dtype, device), + None, + ) + cache[layer, DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE] = ( + self._rand_tensor((seq_len, indexer_compressor_dim), compressor_dtype, device), + None, + ) + + return cache + + def _split_blockwise_buffer(self, buffer: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Split a blockwise FP8 quantized buffer into value and scale buffers. + + Args: + buffer: The blockwise FP8 quantized buffer (shape: [num_blocks, tokens_per_block, bytes_per_token]) + + Returns: + Tuple of (value_buffer, scale_buffer) + """ + num_blocks, tokens_per_block, bytes_per_token = buffer.shape + bytes_per_block = bytes_per_token * tokens_per_block + + # Get value buffer + value_shape = (num_blocks, tokens_per_block, self.index_head_dim) + value_stride = (bytes_per_block, self.index_head_dim, 1) + value_buffer = buffer.as_strided(value_shape, value_stride, 0).view(torch.uint8) + + # Get scale buffer + scale_dim = self.index_head_dim // self.indexer_quant_block_size + scale_bytes = scale_dim * 4 # float32 = 4 bytes + scale_shape = (num_blocks, tokens_per_block, scale_bytes) + scale_stride = (bytes_per_block, scale_bytes, 1) + scale_offset = self.index_head_dim * tokens_per_block + scale_buffer = buffer.as_strided(scale_shape, scale_stride, scale_offset).view( + torch.float32 + ) + + return value_buffer, scale_buffer + + def _prefill_write_paged_cache( + self, + buffer: torch.Tensor, + block_indices: List[int], + values: torch.Tensor, + ) -> None: + """Write context values to a paged cache buffer. + + Args: + buffer: The cache buffer to write to (shape: [num_blocks, tokens_per_block, dim_per_token]) + block_indices: List of block indices to write to + values: Values to write (shape: [seq_len, dim_per_token]) + """ + assert buffer.size(2) == values.size(1), f"{buffer.size(2)=} != {values.size(1)=}" + tokens_per_block = buffer.size(1) + seq_len, dim_per_token = values.shape + + num_blocks = (seq_len + tokens_per_block - 1) // tokens_per_block + assert all(idx != BAD_PAGE_INDEX for idx in block_indices[:num_blocks]), ( + f"{block_indices[:num_blocks]=} contains BAD_PAGE_INDEX" + ) + + if seq_len % tokens_per_block != 0: + # pad the values to the nearest multiple of tokens_per_block + pad_len = tokens_per_block - (seq_len % tokens_per_block) + values = torch.cat( + [ + values, + self._rand_tensor((pad_len, dim_per_token), values.dtype, values.device), + ], + dim=0, + ) + + values_blocks = values.reshape(num_blocks, tokens_per_block, dim_per_token) + buffer[block_indices[:num_blocks]] = values_blocks + + def _decode_write_paged_cache( + self, + buffer: torch.Tensor, + block_indices: List[int], + token_idx: int, + value: torch.Tensor, + ) -> None: + """Simulate the decode phrase. Write one new token to the cache. + + Args: + buffer: The cache buffer to write to (shape: [num_blocks, tokens_per_block, dim_per_token]) + block_indices: List of block indices to write to + token_idx: Index of the new token to write + value: Value to write (shape: [dim_per_token]) + """ + assert buffer.size(2) == value.size(0), f"{buffer.size(2)=} != {value.size(0)=}" + num_blocks, tokens_per_block, _ = buffer.shape + + block_idx = token_idx // tokens_per_block + block_offset = token_idx % tokens_per_block + assert block_idx < num_blocks, f"{block_idx=} >= {num_blocks=}" + assert block_indices[block_idx] != BAD_PAGE_INDEX, ( + f"{block_indices[block_idx]=} == BAD_PAGE_INDEX" + ) + + buffer[block_indices[block_idx], block_offset] = value + + def _read_paged_cache( + self, buffer: torch.Tensor, block_indices: List[int], seq_len: int, window_size: int | None + ) -> torch.Tensor: + """Read values from a paged cache buffer. + + Args: + buffer: The cache buffer to read from (shape: [num_blocks, tokens_per_block, dim_per_token]) + block_indices: List of block/page indices to read from + seq_len: Sequence length + window_size: sliding window size to read from the cache + + Returns: + Tensor containing the read values (shape: [seq_len, dim_per_token] or [window_size, dim_per_token] + if window_size is given and seq_len > window_size) + """ + _, tokens_per_block, dim_per_token = buffer.shape + + # check if all blocks within the window are valid + end_block_idx = (seq_len + tokens_per_block - 1) // tokens_per_block + if window_size is not None: + start_block_idx = (seq_len - window_size + tokens_per_block - 1) // tokens_per_block + else: + start_block_idx = 0 + assert all(idx != BAD_PAGE_INDEX for idx in block_indices[start_block_idx:end_block_idx]), ( + f"{block_indices[start_block_idx:end_block_idx]=} contains BAD_PAGE_INDEX" + ) + + # read values from the cache + values = buffer[block_indices].reshape(-1, dim_per_token)[:seq_len] + if window_size is not None and seq_len > window_size: + values = values[-window_size:] + return values + + def _write_request_prefill( + self, + req: LlmRequest, + prompt_len: int, + cache_manager: DeepseekV4CacheManager, + cache_values: _RequestCache, + ) -> None: + """Write cache values for a request to the cache manager. + + Args: + req: The request to write cache for + prompt_len: Prompt length + cache_manager: The cache manager instance + cache_values: Request's cache to write + """ + compress_ratios = cache_manager._compress_ratios + for (layer_idx, attn_type), (values, scales) in cache_values.items(): + page_indices = cache_manager.get_batch_attn_offset( + [req.py_request_id], + beam_width=1, + num_contexts=1, + num_seqs=1, + attn_type=attn_type, + compress_ratio=compress_ratios[layer_idx], + ).squeeze(0) + + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + seq_len = prompt_len // compress_ratios[layer_idx] + else: + seq_len = prompt_len + + buffer = _view_fp8_as_uint8(cache_manager.get_buffers(layer_idx, attn_type)) + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + # indexer compress is blockwise FP8 quantized + values_buffer, scales_buffer = self._split_blockwise_buffer(buffer) + self._prefill_write_paged_cache( + buffer=values_buffer, + block_indices=page_indices, + values=values[:seq_len], + ) + self._prefill_write_paged_cache( + buffer=scales_buffer, + block_indices=page_indices, + values=scales[:seq_len], + ) + else: + self._prefill_write_paged_cache( + buffer=buffer, + block_indices=page_indices, + values=values[:seq_len], + ) + + def _write_request_decode( + self, + req: LlmRequest, + token_idx: int, + cache_manager: DeepseekV4CacheManager, + cache_values: _RequestCache, + ) -> None: + """Simulate the decode phrase. Write one new token to the cache. + + Args: + req: The request to write cache for + token_idx: Index of the new token to write + cache_manager: The cache manager instance + cache_values: Request's cache to write + """ + compress_ratios = cache_manager._compress_ratios + for (layer_idx, attn_type), (values, scales) in cache_values.items(): + block_indices = cache_manager.get_batch_attn_offset( + [req.py_request_id], + beam_width=1, + num_contexts=1, + num_seqs=1, + attn_type=attn_type, + compress_ratio=compress_ratios[layer_idx], + ).squeeze(0) + + compressed_token_idx = token_idx + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + if (token_idx + 1) % compress_ratios[layer_idx] != 0: + # skip if current token will not trigger compression + continue + compressed_token_idx = token_idx // compress_ratios[layer_idx] + + buffer = _view_fp8_as_uint8(cache_manager.get_buffers(layer_idx, attn_type)) + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + values_buffer, scales_buffer = self._split_blockwise_buffer(buffer) + self._decode_write_paged_cache( + buffer=values_buffer, + block_indices=block_indices, + token_idx=compressed_token_idx, + value=values[compressed_token_idx], + ) + self._decode_write_paged_cache( + buffer=scales_buffer, + block_indices=block_indices, + token_idx=compressed_token_idx, + value=scales[compressed_token_idx], + ) + else: + self._decode_write_paged_cache( + buffer=buffer, + block_indices=block_indices, + token_idx=compressed_token_idx, + value=values[compressed_token_idx], + ) + + def _read_request( + self, + req: LlmRequest, + seq_len: int, + cache_manager: DeepseekV4CacheManager, + compress_ratios: List[int], + ) -> _RequestCache: + """Read cache values for a request from the cache manager. + + Args: + req: The request to read cache for + seq_len: Sequence length + cache_manager: The cache manager instance + compress_ratios: Compression ratios for each layer + + Returns: + Request's cache + """ + cache_values: _RequestCache = {} + for layer, ratio in enumerate(compress_ratios): + attn_types = [DeepseekV4AttentionType.SWA] + if self._is_compress_layer(ratio): + attn_types.extend( + [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + ] + ) + if self._is_sparse_layer(ratio): + attn_types.extend( + [ + DeepseekV4AttentionType.INDEXER_COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ] + ) + + # read cache values for each attention type + for attn_type in attn_types: + page_indices = cache_manager.get_batch_attn_offset( + [req.py_request_id], + beam_width=1, + num_contexts=1, + num_seqs=1, + attn_type=attn_type, + compress_ratio=ratio, + ).squeeze(0) + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + attn_len = seq_len // ratio + else: + attn_len = seq_len + window_size = self._get_window_size(ratio, attn_type) + + buffer = _view_fp8_as_uint8(cache_manager.get_buffers(layer, attn_type)) + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + values_buffer, scales_buffer = self._split_blockwise_buffer(buffer) + values = self._read_paged_cache( + buffer=values_buffer, + block_indices=page_indices, + seq_len=attn_len, + window_size=window_size, + ) + scales = self._read_paged_cache( + buffer=scales_buffer, + block_indices=page_indices, + seq_len=attn_len, + window_size=window_size, + ) + else: + values = self._read_paged_cache( + buffer=buffer, + block_indices=page_indices, + seq_len=attn_len, + window_size=window_size, + ) + scales = None + + cache_values[layer, attn_type] = (values, scales) + + return cache_values + + def _assert_cache_equal( + self, seq_len: int, compress_ratios: List[int], expect: _RequestCache, actual: _RequestCache + ) -> None: + """Assert that two cache dictionaries contain equal values. + + Args: + seq_len: Sequence length + compress_ratios: Compression ratios for each layer + expected: Expected cache values + actual: Actual cache values read from cache manager + """ + # Check that keys match + assert set(expect.keys()) == set(actual.keys()), ( + f"Cache keys don't match. Expected: {set(expect.keys())}, Actual: {set(actual.keys())}" + ) + + # Check each tensor value + for layer_idx, attn_type in expect.keys(): + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + attn_len = seq_len // compress_ratios[layer_idx] + else: + attn_len = seq_len + + expect_values, expect_scales = expect[layer_idx, attn_type] + actual_values, actual_scales = actual[layer_idx, attn_type] + + # Slice to attention length + expect_values = expect_values[:attn_len] + if expect_scales is not None: + expect_scales = expect_scales[:attn_len] + + # Apply window size if applicable + window_size = self._get_window_size(compress_ratios[layer_idx], attn_type) + if window_size is not None: + expect_values = expect_values[-window_size:] + if expect_scales is not None: + expect_scales = expect_scales[-window_size:] + + # Assert values match + torch.testing.assert_close( + actual_values, + expect_values, + rtol=1e-5, + atol=1e-5, + msg=f"Mismatch for layer {layer_idx}, attention type {attn_type.name} (values)", + ) + + # Assert scales match (both should be None or both should be tensors) + if expect_scales is None: + assert actual_scales is None, ( + f"Expected no scales for layer {layer_idx}, attention type {attn_type.name}, " + f"but got scales with shape {actual_scales.shape}" + ) + else: + assert actual_scales is not None, ( + f"Expected scales for layer {layer_idx}, attention type {attn_type.name}, " + f"but got None" + ) + torch.testing.assert_close( + actual_scales, + expect_scales, + rtol=1e-5, + atol=1e-5, + msg=f"Mismatch for layer {layer_idx}, attention type {attn_type.name} (scales)", + ) + + def test_indexer_cache_layout_default(self): + """Indexer compressor cache: FP8 blockwise (128 fp8 + per-128 fp32 scale).""" + cache_manager, _ = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=1, + max_seq_len=512, + compress_ratios=[4], + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + ) + try: + buffer = cache_manager.get_buffers(0, DeepseekV4AttentionType.INDEXER_COMPRESS) + assert buffer.dtype == torch.float8_e4m3fn + assert buffer.shape[-1] == 128 + 4 + assert cache_manager.quant_block_size == 128 + finally: + cache_manager.shutdown() + + @pytest.mark.parametrize("compress_ratios", [[1, 4, 128]]) + @pytest.mark.parametrize( + "dtype,compressor_dtype", [(DataType.BF16, DataType.FLOAT), (DataType.FP8, DataType.FLOAT)] + ) + @pytest.mark.parametrize("prompt_lens", [[512, 128, 160], [1024, 2048, 4096]]) + @pytest.mark.parametrize("num_generation_steps", [2, 100]) + def test_write_read_cache( + self, + compress_ratios: List[int], + prompt_lens: List[int], + num_generation_steps: int, + dtype: DataType, + compressor_dtype: DataType, + ): + max_batch_size = len(prompt_lens) + max_seq_len = max(prompt_lens) + num_generation_steps + 1 + max_input_len = max(prompt_lens) + # Create cache manager and sparse attention config + cache_manager, sparse_attn_config = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + compress_ratios=compress_ratios, + dtype=dtype, + compressor_dtype=compressor_dtype, + max_input_len=max_input_len, + ) + + # Create requests and their cache values + requests = list[LlmRequest]() + try: + cache_values = dict[int, _RequestCache]() + for req_id, prompt_len in enumerate(prompt_lens): + req = self._create_request(req_id, prompt_len) + requests.append(req) + + # Generate random cache values for this request + cache_values[req_id] = self._create_random_cache( + seq_len=prompt_len + num_generation_steps + 1, + head_dim=self.head_dim, + sparse_attn_config=sparse_attn_config, + dtype=binding_to_torch_dtype(dtype), + compressor_dtype=binding_to_torch_dtype(compressor_dtype), + ) + + # Simulate the prefill phrase + scheduled_batch = ScheduledRequests() + scheduled_batch.context_requests_last_chunk = requests + for req in requests: + cache_manager.prepare_context(req) + cache_manager.resize_context(req, req.context_chunk_size) + + # Write context to cache + for req in requests: + self._write_request_prefill( + req=req, + prompt_len=prompt_lens[req.py_request_id], + cache_manager=cache_manager, + cache_values=cache_values[req.py_request_id], + ) + + # Update requests state and call update_resources + for req in requests: + req.context_current_position = prompt_lens[req.py_request_id] + req.add_new_token(prompt_lens[req.py_request_id], 0) + cache_manager.update_resources(scheduled_batch) + + # Read context from cache and verify + for req in requests: + actual_cache_values = self._read_request( + req=req, + seq_len=prompt_lens[req.py_request_id], + cache_manager=cache_manager, + compress_ratios=compress_ratios, + ) + self._assert_cache_equal( + seq_len=prompt_lens[req.py_request_id], + compress_ratios=compress_ratios, + expect=cache_values[req.py_request_id], + actual=actual_cache_values, + ) + + # Simulate the decode phrase + for i in range(num_generation_steps): + seq_lens = [prompt_len + i + 1 for prompt_len in prompt_lens] + scheduled_batch = ScheduledRequests() + scheduled_batch.generation_requests = requests + for req in requests: + cache_manager.try_allocate_generation(req) + + # Write new token to cache + for req in requests: + self._write_request_decode( + req=req, + token_idx=seq_lens[req.py_request_id] - 1, + cache_manager=cache_manager, + cache_values=cache_values[req.py_request_id], + ) + + # Read context from cache and verify + for req in requests: + actual_cache_values = self._read_request( + req=req, + seq_len=seq_lens[req.py_request_id], + cache_manager=cache_manager, + compress_ratios=compress_ratios, + ) + self._assert_cache_equal( + seq_len=seq_lens[req.py_request_id], + compress_ratios=compress_ratios, + expect=cache_values[req.py_request_id], + actual=actual_cache_values, + ) + + for req in requests: + req.add_new_token(seq_lens[req.py_request_id], 0) + cache_manager.update_resources(scheduled_batch) + finally: + try: + for req in requests: + cache_manager.free_resources(req) + finally: + cache_manager.shutdown() + + @pytest.mark.parametrize("compress_ratios", [[1, 4, 128]]) + @pytest.mark.parametrize( + "dtype,compressor_dtype", [(DataType.BF16, DataType.FLOAT), (DataType.FP8, DataType.FLOAT)] + ) + def test_kv_cache_pool_mapping( + self, compress_ratios: List[int], dtype: DataType, compressor_dtype: DataType + ): + # Create cache manager and sparse attention config + num_layers = len(compress_ratios) + cache_manager, _ = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=4, + max_seq_len=1024, + compress_ratios=compress_ratios, + dtype=dtype, + compressor_dtype=compressor_dtype, + ) + + try: + kv_cache_pool_mapping = cache_manager.kv_cache_pool_mapping + assert kv_cache_pool_mapping.shape == (num_layers, 2) + + assert torch.all(kv_cache_pool_mapping[:, 0] != -1), ( + "all layers should have swa attention pool" + ) + assert torch.all(kv_cache_pool_mapping[:, 1] >= 0), ( + "buffer pointer offset should be non-negative" + ) + assert torch.all(kv_cache_pool_mapping[:, 0] == kv_cache_pool_mapping[0, 0]), ( + "all layers should have the same pool_id" + ) + finally: + cache_manager.shutdown() + + @pytest.mark.parametrize("compress_ratios", [[1, 4, 128]]) + @pytest.mark.parametrize( + "dtype,compressor_dtype", [(DataType.BF16, DataType.FLOAT), (DataType.FP8, DataType.FLOAT)] + ) + @pytest.mark.parametrize("invalid", [False, True]) + @pytest.mark.parametrize("fill_with_zero", [False, True]) + def test_check_invalid_values_in_kv_cache( + self, + compress_ratios: List[int], + dtype: DataType, + compressor_dtype: DataType, + invalid: bool, + fill_with_zero: bool, + ): + """Test invalid value detection and optional zero-fill behavior in KV cache.""" + cache_manager, _ = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=4, + max_seq_len=1024, + compress_ratios=compress_ratios, + dtype=dtype, + compressor_dtype=compressor_dtype, + ) + + needs_invalid_cleanup = False + try: + # Fresh cache (zero-initialized) should have no invalid values + result = cache_manager.check_invalid_values_in_kv_cache() + assert not result, "Fresh cache should have no invalid values" + + if invalid: + # Inject invalid into a float buffer so NaN/Inf checks are supported. + layer_idx = next(i for i, ratio in enumerate(compress_ratios) if ratio > 1) + buffer = cache_manager.get_buffers( + layer_idx, DeepseekV4AttentionType.COMPRESSOR_STATE + ) + buffer[0, 0, 0] = torch.nan + needs_invalid_cleanup = True + + result = cache_manager.check_invalid_values_in_kv_cache(fill_with_zero=fill_with_zero) + if invalid and fill_with_zero: + needs_invalid_cleanup = False + assert result == invalid, ( + f"Expected invalid={invalid} from check_invalid_values_in_kv_cache, got {result}" + ) + + # Verify whether invalid values remain after the check. + post_check = cache_manager.check_invalid_values_in_kv_cache() + expected_post_check = invalid and not fill_with_zero + assert post_check == expected_post_check, ( + f"Expected post-check invalid={expected_post_check}, got {post_check}" + ) + + if expected_post_check: + # Cleanup for shutdown path when zero-fill wasn't requested above. + cache_manager.check_invalid_values_in_kv_cache(fill_with_zero=True) + needs_invalid_cleanup = False + finally: + try: + if needs_invalid_cleanup: + cache_manager.check_invalid_values_in_kv_cache(fill_with_zero=True) + finally: + cache_manager.shutdown() + + +if __name__ == "__main__": + tester = TestDeepseekV4CacheManager() + print("=== FP8, prompt_lens=[1024, 2048, 4096], steps=100 ===") + tester.test_write_read_cache([1, 4, 128], [1024, 2048, 4096], 100, DataType.FP8, DataType.FLOAT) + print("Test passed") diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index a415fdebcb98..f467252b3018 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -228,7 +228,7 @@ methods: default: null status: prototype sparse_attention_config: - annotation: Union[tensorrt_llm.llmapi.llm_args.RocketSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.SkipSoftmaxAttentionConfig, tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig, NoneType] + annotation: Union[tensorrt_llm.llmapi.llm_args.RocketSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekV4SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.SkipSoftmaxAttentionConfig, tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig, NoneType] default: null status: prototype otlp_traces_endpoint: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 41df9229ecc8..6888c11c6de5 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -32,6 +32,7 @@ CudaGraphConfig, DecodeCudaGraphConfig, DecodingBaseConfig, + DeepSeekV4SparseAttentionConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, @@ -2685,3 +2686,16 @@ def test_ckpt_sparse_attention_config_can_be_passed_directly(self): assert params.threshold_scale_factor_prefill == pytest.approx( 100.0 * math.exp(5.0 * 0.5)) + + +class TestDeepSeekV4SparseAttentionConfig: + + def test_zero_compress_ratios_are_normalized(self): + config = DeepSeekV4SparseAttentionConfig(compress_ratios=[0, 4, 128]) + + assert config.compress_ratios == [1, 4, 128] + + @pytest.mark.parametrize("compress_ratios", [[], [-1, 4, 128]]) + def test_invalid_compress_ratios_raise(self, compress_ratios): + with pytest.raises(ValidationError, match="compress_ratios"): + DeepSeekV4SparseAttentionConfig(compress_ratios=compress_ratios)