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"]