45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
|
|
||
|
|
"""
|
||
|
|
MagCache state management.
|
||
|
|
|
||
|
|
This module contains the MagCacheState class which manages the internal state
|
||
|
|
for MagCache caching logic, including residuals, accumulated metrics, and step tracking.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
class MagCacheState:
|
||
|
|
"""State management for MagCache caching logic."""
|
||
|
|
|
||
|
|
def __init__(self) -> None:
|
||
|
|
"""Initialize empty MagCache state."""
|
||
|
|
self.previous_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None
|
||
|
|
self.head_block_input: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None
|
||
|
|
self.should_compute: bool = True
|
||
|
|
self.accumulated_ratio: float = 1.0
|
||
|
|
self.accumulated_err: float = 0.0
|
||
|
|
self.accumulated_steps: int = 0
|
||
|
|
self.step_index: int = 0
|
||
|
|
self.calibration_ratios: list[float] = []
|
||
|
|
self.norm_ratios: list[float] = []
|
||
|
|
self.norm_stds: list[float] = []
|
||
|
|
self.cos_dises: list[float] = []
|
||
|
|
self._is_first_step: bool = True
|
||
|
|
|
||
|
|
def reset(self) -> None:
|
||
|
|
"""Reset all state variables for a new inference run."""
|
||
|
|
self.previous_residual = None
|
||
|
|
self.head_block_input = None
|
||
|
|
self.should_compute = True
|
||
|
|
self.accumulated_ratio = 1.0
|
||
|
|
self.accumulated_err = 0.0
|
||
|
|
self.accumulated_steps = 0
|
||
|
|
self.step_index = 0
|
||
|
|
self.calibration_ratios = []
|
||
|
|
self.norm_ratios = []
|
||
|
|
self.norm_stds = []
|
||
|
|
self.cos_dises = []
|
||
|
|
self._is_first_step = True
|