stats_manager

Statistics manager for sparse attention modules.

Classes

SparseAttentionStatsManager

Centralized statistics manager for sparse attention.
class SparseAttentionStatsManager

Bases: object

Centralized statistics manager for sparse attention.

This class is the single source of truth for all statistics collection in sparse attention modules. It handles both runtime aggregation and per-sample calibration statistics.

Design principles: - Single responsibility: only stats management - No computation: receives pre-computed stats from methods - Optional: can be None if stats collection disabled - Zero overhead when disabled

__init__(module_name, enabled=True)

Initialize stats manager.

Parameters:

  • module_name (str) – Name of the module this manager is attached to

  • enabled (bool) – Whether stats collection is enabled

collect(stats)

Collect statistics from a single forward pass.

Parameters:

stats (dict) – Dictionary containing statistics from method computation. Expected keys: sparsity, phase, total_blocks, sparse_blocks, sample_length (optional)

get_calibration_stats(phase=None)

Get per-sample calibration statistics, optionally filtered by phase.

Note: Returns historical stats collected while calibration_mode was enabled. Stats remain accessible even after calibration_mode is disabled. New stats are only collected when calibration_mode is True.

Parameters:

phase (str | None) – Optional phase to filter by (‘prefill’ or ‘decode’). If None, returns all stats.

Returns:

List of per-sample statistics dictionaries. Empty list if no stats were collected or no stats match the phase.

Return type:

list[dict]

get_summary()

Get aggregated statistics summary.

Returns:

Dictionary with module name, total calls, average sparsity, and phase distribution.

Return type:

dict

reset()

Reset all statistics to initial state.

set_calibration_mode(enabled, reset_history=True)

Enable or disable calibration mode.

In calibration mode, per-sample statistics are stored for detailed analysis. Otherwise, only aggregated stats are maintained.

Parameters:

  • enabled (bool) – Whether to enable calibration mode

  • reset_history (bool) – Whether to clear per_sample_stats when enabling