Class CustomQuantizeDynamicCache

src/utils/cache.py:174–334

自定义量化动态缓存类 扩展标准的QuantoQuantizedCache，添加额外的元数据存储和查询功能

        return super().update_router_kcache(key_states, layer_idx)

class CustomQuantizeDynamicCache(QuantoQuantizedCache):
    """
    自定义量化动态缓存类
    扩展标准的QuantoQuantizedCache，添加额外的元数据存储和查询功能
    """
    def __init__(self, cache_config):
        super().__init__(cache_config)
        self.cache_config = cache_config
        self.cache_kwargs = {}
        self.group_cache = {}
        self.meta = {}

    def record_kwargs(self, layer_idx, kwargs):
        """
        记录层的元数据信息

        Args:
            layer_idx: 层索引
            kwargs: 包含路由层信息的字典
        """
        if layer_idx in self.cache_kwargs:
            self.cache_kwargs[layer_idx].update(kwargs)
        else:
            self.cache_kwargs[layer_idx] = kwargs

    def get_layer_length(self):
        return len(self.cache_kwargs)

    def clear_kvcache(self):
        self._quantized_key_cache = []
        self._quantized_value_cache = []
        self.key_cache = []
        self.value_cache = []

    def get_kvcache(self, layer_idx):
        """
        获取指定层的KV缓存（反量化后）

        Args:
            layer_idx: 层索引

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (key_cache, value_cache)
        """
        dequant_key = self._dequantize(self._quantized_key_cache[layer_idx])
        dequant_value = self._dequantize(self._quantized_value_cache[layer_idx])
        return dequant_key, dequant_value

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs=None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        更新缓存