Spaces:
Running on Zero
Running on Zero
File size: 8,304 Bytes
ee69ab7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import torch
# These constants are deliberately tied to the released FlashVSR+ hybrid
# checkpoint and its 1536x2560 RTX 4090 reference run. The policy is not a
# generic PyTorch memory planner; it is a conservative guard for this stream
# inference shape.
REFERENCE_SPATIAL_BLOCKS = 240
REFERENCE_NO_OFFLOAD_PEAK_MIB = 17_394.0
REFERENCE_OFFLOAD_PEAK_MIB = 14_792.0
REFERENCE_STATIC_ALLOCATED_MIB = 10_000.0
TOTAL_MEMORY_SAFETY_FACTOR = 0.90
@dataclass(frozen=True)
class StreamOffloadDecision:
mode: str
enabled: bool
layers: tuple[int, ...]
spatial_blocks: int
free_memory_mib: float | None
total_memory_mib: float | None
estimated_no_offload_peak_mib: float
estimated_offload_peak_mib: float
reason: str
def metadata(self) -> dict[str, Any]:
return {
"cache_offload_mode": self.mode,
"cache_offload_enabled": self.enabled,
"cache_offload_layers": list(self.layers),
"cache_offload_spatial_blocks": self.spatial_blocks,
"cache_offload_free_memory_mib": self.free_memory_mib,
"cache_offload_total_memory_mib": self.total_memory_mib,
"cache_offload_estimated_no_offload_peak_mib": (
self.estimated_no_offload_peak_mib
),
"cache_offload_estimated_offload_peak_mib": (
self.estimated_offload_peak_mib
),
"cache_offload_decision_reason": self.reason,
}
def _memory_info(device: torch.device) -> tuple[float | None, float | None]:
if device.type != "cuda" or not torch.cuda.is_available():
return None, None
try:
free, total = torch.cuda.mem_get_info(device)
except RuntimeError:
return None, None
return free / 1024**2, total / 1024**2
def _original_layers(denoiser: torch.nn.Module, config) -> tuple[int, ...]:
blocks = list(getattr(denoiser, "blocks", ()))
if not blocks:
return ()
# SlidingLocalSelfAttention subclasses BlockSparseSelfAttention, so use
# the concrete class name rather than isinstance on the base class.
if getattr(config, "is_hybrid_sliding", False):
return tuple(
index
for index, block in enumerate(blocks)
if type(block.self_attn).__name__ != "SlidingLocalSelfAttention"
)
return tuple(range(len(blocks)))
def _estimate_peak(spatial_blocks: int) -> tuple[float, float]:
# The persistent original-layer cache scales approximately linearly with
# spatial blocks, while draft/workspace pressure grows more steeply. A
# squared ratio is intentionally conservative at the 4K boundary.
ratio = max(spatial_blocks, 1) / REFERENCE_SPATIAL_BLOCKS
estimated_no_offload = REFERENCE_OFFLOAD_PEAK_MIB + (
REFERENCE_NO_OFFLOAD_PEAK_MIB - REFERENCE_OFFLOAD_PEAK_MIB
) * ratio * ratio
estimated_offload = REFERENCE_OFFLOAD_PEAK_MIB * max(ratio, 1.0)
return estimated_no_offload, estimated_offload
def choose_stream_cache_offload(
config,
denoiser: torch.nn.Module,
video: torch.Tensor,
device: torch.device,
) -> StreamOffloadDecision:
"""Choose cache residency for the current request without running a probe.
Explicit layer indices remain a backwards-compatible force-on override,
except when ``never`` is explicitly requested. ``never`` and ``always``
are intended for debugging/benchmarking; normal inference uses ``auto``.
"""
height, width = video.shape[-2:]
spatial_blocks = (height // 16 // 8) * (width // 16 // 8)
free_mib, total_mib = _memory_info(device)
estimated_no, estimated_offload = _estimate_peak(spatial_blocks)
original_layers = _original_layers(denoiser, config)
is_base = not getattr(config, "is_hybrid_sliding", False)
if is_base:
# Conservative base estimate: retain the reference workspace floor,
# then budget every original layer's full K/V history separately.
# This is an estimate, not a calibrated base benchmark. Do not scale
# the hybrid five-layer cache delta as though it represented 30 layers.
cached_frames = int(config.sparse_kv_ratio) * config.sparse_window_size[0]
tokens = video.shape[0] * (height // 16) * (width // 16) * cached_frames
cache_mib = sum(
2 * tokens * block.self_attn.q.out_features
* block.self_attn.q.weight.element_size() / 1024**2
for block in getattr(denoiser, "blocks", ())
)
estimated_no = max(estimated_no, estimated_offload + cache_mib)
explicit_layers = tuple(getattr(config, "stream_cache_offload_layers", ()))
mode = getattr(config, "stream_cache_offload_mode", "auto")
if mode == "never":
enabled = False
reason = "forced never"
elif explicit_layers:
return StreamOffloadDecision(
mode="always",
enabled=True,
layers=explicit_layers,
spatial_blocks=spatial_blocks,
free_memory_mib=free_mib,
total_memory_mib=total_mib,
estimated_no_offload_peak_mib=estimated_no,
estimated_offload_peak_mib=estimated_offload,
reason="explicit stream_cache_offload_layers",
)
elif mode == "always":
enabled = True
reason = "forced always"
else:
reasons = []
if is_base and device.type == "cuda" and (free_mib is None or total_mib is None):
reasons.append("base: VRAM information unavailable; conservatively offload KV")
# Compare both residency plans with the current GPU capacity. The
# reference peak includes roughly REFERENCE_STATIC_ALLOCATED_MIB of
# model/static allocations; replace that part with the actual current
# process allocation when CUDA reports it. This keeps auto policy
# tied to the installed GPU instead of a resolution-only hard limit.
capacity_mib = None if total_mib is None else total_mib * TOTAL_MEMORY_SAFETY_FACTOR
if free_mib is not None and total_mib is not None:
allocated_mib = total_mib - free_mib
estimated_no = allocated_mib + max(
estimated_no - REFERENCE_STATIC_ALLOCATED_MIB, 0.0
)
estimated_offload = allocated_mib + max(
estimated_offload - REFERENCE_STATIC_ALLOCATED_MIB, 0.0
)
no_fits = capacity_mib is not None and estimated_no <= capacity_mib
offload_fits = capacity_mib is not None and estimated_offload <= capacity_mib
if capacity_mib is not None and not no_fits:
reasons.append(
f"estimated_no_offload_peak={estimated_no:.0f}MiB exceeds "
f"available VRAM ({capacity_mib:.0f}MiB)"
)
# If the no-offload estimate does not fit, offload is the only plan
# that can reduce persistent KV residency. The extrapolated
# offload estimate can be pessimistic at resolutions beyond the
# reference run, so do not disable the only mitigation merely because
# that estimate is also above the safety budget.
if capacity_mib is None:
enabled = bool(reasons)
else:
enabled = not no_fits
if capacity_mib is None and not reasons:
reason = "VRAM information unavailable; keep cache on GPU"
elif not no_fits and not reasons:
reason = "estimated peak fits available VRAM"
else:
reason = "; ".join(reasons)
if not offload_fits:
reason += (
f"; offload estimate also exceeds safety budget "
f"({estimated_offload:.0f}MiB), enabling as best effort"
)
return StreamOffloadDecision(
mode=mode,
enabled=enabled,
layers=original_layers if enabled else (),
spatial_blocks=spatial_blocks,
free_memory_mib=free_mib,
total_memory_mib=total_mib,
estimated_no_offload_peak_mib=estimated_no,
estimated_offload_peak_mib=estimated_offload,
reason=reason,
)
__all__ = ["StreamOffloadDecision", "choose_stream_cache_offload"]
|