[diffusion] feat: support quant-videogen prq kv-cache quantization (memory-saving) for causal-dit (#32581)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/pyproject.toml

@@ -123,6 +123,14 @@ diffusion = [
   "xatlas",
 ]
 
+diffusion-qvg = [
+  # quant-videogen 0.1.0 pins an incompatible Torch version, so it is
+  # installed separately with --no-deps; this extra provides its missing
+  # import-time dependencies.
+  "loguru>=0.7",
+  "termcolor>=2.3",
+]
+
 ray = [
   "ray[default]>=2.55.1",
 ]

python/sglang/multimodal_gen/configs/quantization/qvg_kv.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CLI-facing configuration for Quant-VideoGen PRQ KV-cache quantization.
+
+Mirrors the SRT `--kv-cache-dtype` pattern: the on/off + tuning knobs live on a
+typed config object carried by ServerArgs (see `kv_cache_quant_config`), instead
+of a pile of raw environment variables.
+
+Defaults match the tuned per-chunk setting
+(kmeans 1 stage, 128 centroids, block 64, symmetric, 2 iters, recent 1,
+per-chunk sink) — i.e. `--kv-cache-quant int4` alone reproduces it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+_BITS = {"off": None, "none": None, "bf16": None, "int4": 4, "int2": 2}
+
+
+def _parse_bits(val: str | None) -> int | None:
+    if val is None:
+        return None
+    key = str(val).strip().lower()
+    if key not in _BITS:
+        raise ValueError(f"kv-cache-quant must be one of {list(_BITS)}, got {val!r}")
+    return _BITS[key]
+
+
+@dataclass
+class QVGKVQuantArgs:
+    """PRQ (multi-stage k-means) KV-cache quantization settings.
+
+    ``bits is None`` means quantization is OFF (plain bf16 cache). Defaults
+    reproduce the tuned per-chunk config from offline sweeps.
+    """
+
+    bits: int | None = None  # None => off; 2 or 4 (master switch)
+    centroids: int = 128  # k-means centroids per stage
+    block_size: int = 64  # residual scale block size
+    stages: int = 1  # PRQ k-means stages
+    kmeans_iters: int = 2  # k-means iterations
+    asymmetric: bool = False  # KIVI-style asymmetric residual quant
+    keep_recent_chunks: int = 1  # completed chunks kept bf16 (recency guard)
+    sink: bool = True  # quantize the attention sink too
+    sink_keep_chunks: int = 0  # leading sink chunks kept bf16 forever
+
+    @property
+    def enabled(self) -> bool:
+        return self.bits is not None
+
+    def validate(self) -> QVGKVQuantArgs:
+        if self.bits not in (None, 2, 4):
+            raise ValueError(f"kv-cache-quant bits must be 2 or 4, got {self.bits}")
+        if self.centroids <= 0:
+            raise ValueError("kv-cache-quant-centroids must be > 0")
+        if self.block_size <= 0:
+            raise ValueError("kv-cache-quant-block-size must be > 0")
+        if self.stages <= 0:
+            raise ValueError("kv-cache-quant-stages must be > 0")
+        if self.kmeans_iters <= 0:
+            raise ValueError("kv-cache-quant-iters must be > 0")
+        if self.keep_recent_chunks < 0 or self.sink_keep_chunks < 0:
+            raise ValueError("keep_recent_chunks / sink_keep_chunks must be >= 0")
+        return self
+
+    def describe(self) -> str:
+        if not self.enabled:
+            return "off"
+        return (
+            f"int{self.bits} centroids={self.centroids} block={self.block_size} "
+            f"stages={self.stages} iters={self.kmeans_iters} "
+            f"asym={self.asymmetric} recent={self.keep_recent_chunks} "
+            f"sink={self.sink} sink_keep={self.sink_keep_chunks}"
+        )
+
+    @classmethod
+    def from_dict(cls, kwargs: dict) -> QVGKVQuantArgs:
+        """Build from flat CLI kwargs (dest names ``kv_cache_quant*``)."""
+        master = kwargs.get("kv_cache_quant")
+        if master is None:
+            return cls()
+        inst = cls(bits=_parse_bits(master))
+        if kwargs.get("kv_cache_quant_centroids") is not None:
+            inst.centroids = kwargs["kv_cache_quant_centroids"]
+        if kwargs.get("kv_cache_quant_block_size") is not None:
+            inst.block_size = kwargs["kv_cache_quant_block_size"]
+        if kwargs.get("kv_cache_quant_stages") is not None:
+            inst.stages = kwargs["kv_cache_quant_stages"]
+        if kwargs.get("kv_cache_quant_iters") is not None:
+            inst.kmeans_iters = kwargs["kv_cache_quant_iters"]
+        if kwargs.get("kv_cache_quant_keep_recent") is not None:
+            inst.keep_recent_chunks = kwargs["kv_cache_quant_keep_recent"]
+        if kwargs.get("kv_cache_quant_sink_keep") is not None:
+            inst.sink_keep_chunks = kwargs["kv_cache_quant_sink_keep"]
+        if kwargs.get("kv_cache_quant_asymmetric") is not None:
+            inst.asymmetric = kwargs["kv_cache_quant_asymmetric"]
+        if kwargs.get("kv_cache_quant_sink") is not None:
+            inst.sink = kwargs["kv_cache_quant_sink"]
+        inst.sink = bool(inst.sink)
+        inst.asymmetric = bool(inst.asymmetric)
+        return inst.validate()

python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py

@@ -42,6 +42,10 @@ class CausalSelfAttentionKVCache:
         if self.attention_window_size == 0:
             self.attention_window_size = self.cache_size
 
+    @property
+    def num_cache_heads(self) -> int:
+        return self.k.shape[2]
+
     def reset_indices(self) -> None:
         self.global_end_index.zero_()
         self.local_end_index.zero_()

python/sglang/multimodal_gen/runtime/layers/kvcache/qvg_packed_cache.py

@@ -0,0 +1,401 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Quant-VideoGen KV cache with PRQ-packed storage for completed chunks.
+
+Storage model (mirrors Quant-VideoGen's ChunkedKVCache, fitted to SGLang's
+``update_and_get_attention_kv`` contract):
+
+  * The retained window is split into frame/chunk-aligned *segments* in global
+    token order. Each segment is either BF16 (`k`/`v` tensors resident) or
+    PRQ-packed (`packed_k`/`packed_v` dicts resident, BF16 freed).
+  * The current (still-denoising) chunk and the newest ``keep_recent_chunks``
+    completed chunks stay BF16 (rewritten each denoise step / attended cleanly).
+  * Older completed segments are PRQ-packed once and their BF16 freed -> the
+    resident footprint drops to ~(sink + recent) BF16 + packed tail.
+  * On read the visible window is reconstructed densely on the fly (dequantize
+    packed segments + cat BF16 ones) and returned to attention; that transient
+    dense tensor is freed after the layer's attention, so only ONE layer is
+    dense at a time vs. all ``num_layers`` resident dense windows before.
+
+Scope: the LingBot realtime causal path only (sliding window + sink, chunk-
+aligned writes, optional ulysses head-slice, ``recent_window_tokens`` None or a
+non-negative int). Unsupported base-class features raise NotImplementedError
+rather than silently diverge.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cache
+
+import torch
+
+from sglang.multimodal_gen.configs.quantization.qvg_kv import QVGKVQuantArgs
+from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import (
+    CausalAttentionKVView,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@cache
+def _qvg_functions():
+    try:
+        from quant_videogen.functions import (
+            triton_prq_dequantize_tensor,
+            triton_prq_quantize_tensor,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "Quant-VideoGen KV-cache quantization requires its optional "
+            "runtime dependencies. Install them with: "
+            "pip install 'sglang[diffusion-qvg]' && "
+            "pip install --no-deps quant-videogen==0.1.0."
+        ) from e
+    return triton_prq_quantize_tensor, triton_prq_dequantize_tensor
+
+
+@dataclass
+class _Segment:
+    g0: int  # global start token (inclusive)
+    g1: int  # global end token (exclusive)
+    is_sink: bool  # sink segments are never evicted
+    k: torch.Tensor | None = None
+    v: torch.Tensor | None = None
+    packed_k: dict | None = None
+    packed_v: dict | None = None
+
+    @property
+    def packed(self) -> bool:
+        return self.packed_k is not None
+
+    def nbytes(self) -> int:
+        if self.packed:
+            return _packed_nbytes(self.packed_k) + _packed_nbytes(self.packed_v)
+        return (
+            self.k.numel() * self.k.element_size()
+            + self.v.numel() * self.v.element_size()
+        )
+
+
+def _packed_nbytes(packed: dict) -> int:
+    total = 0
+    for key in ("centroids_list", "cluster_ids_list"):
+        for t in packed.get(key) or []:
+            total += t.numel() * t.element_size()
+    for key in ("residual_quant", "scales", "zeros", "residual", "scale_factor"):
+        t = packed.get(key)
+        if isinstance(t, torch.Tensor):
+            total += t.numel() * t.element_size()
+    return total
+
+
+class QVGPackedCausalKVCache:
+    """Chunk-segmented causal KV cache with PRQ-packed cold segments."""
+
+    def __init__(
+        self,
+        *,
+        batch_size: int,
+        cache_size: int,
+        num_heads: int,
+        head_dim: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        global_end_index: torch.Tensor,
+        local_end_index: torch.Tensor,
+        use_int_indices: bool = False,
+        sink_tokens: int = 0,
+        attention_window_size: int | None = None,
+        quant_args: QVGKVQuantArgs,
+    ) -> None:
+        self.batch_size = batch_size
+        self.cache_size = cache_size
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.dtype = dtype
+        self.device = device
+        self.sink_tokens = sink_tokens
+        self.global_sink_tokens = 0
+        self.attention_window_size = attention_window_size or cache_size
+        self.q = quant_args
+        # kept for API compatibility with the dense cache (indices are not read
+        # by consumers, but reset/patterns touch them)
+        self.global_end_index = global_end_index
+        self.local_end_index = local_end_index
+        self.global_end_index_int = 0 if use_int_indices else None
+        self.local_end_index_int = 0 if use_int_indices else None
+
+        self._segments: list[_Segment] = []  # completed, global-ordered
+        self._cur: _Segment | None = None  # current (mutable) chunk
+        self._global_end = 0
+        self._chunk_tokens = 0  # inferred from first advance
+
+    # ------------------------------------------------------------------ api
+    def reset_indices(self) -> None:
+        self._segments = []
+        self._cur = None
+        self._global_end = 0
+        if self.global_end_index_int is not None:
+            self.global_end_index_int = 0
+            self.local_end_index_int = 0
+        self.global_end_index.zero_()
+        self.local_end_index.zero_()
+
+    def can_direct_current_attention(self, num_new_tokens: int) -> bool:
+        return (
+            self.sink_tokens == 0
+            and self.cache_size == num_new_tokens
+            and self.attention_window_size == num_new_tokens
+        )
+
+    @property
+    def num_cache_heads(self) -> int:
+        return self.num_heads
+
+    def pin_current_chunk(self, current_num_tokens: int) -> None:
+        raise NotImplementedError(
+            "QVGPackedCausalKVCache does not support pinned-sink (longlive2); "
+            "packed KV quant is scoped to the LingBot realtime path."
+        )
+
+    def resident_nbytes(self) -> int:
+        total = sum(s.nbytes() for s in self._segments)
+        if self._cur is not None:
+            total += self._cur.nbytes()
+        return total
+
+    # -------------------------------------------------------------- helpers
+    def _new_bf16_segment(self, g0: int, g1: int, is_sink: bool) -> _Segment:
+        n = g1 - g0
+        k = torch.zeros(
+            self.batch_size,
+            n,
+            self.num_heads,
+            self.head_dim,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        v = torch.zeros_like(k)
+        return _Segment(g0=g0, g1=g1, is_sink=is_sink, k=k, v=v)
+
+    def _write(self, seg: _Segment, key, value, head_slice) -> None:
+        if head_slice is None:
+            seg.k.copy_(key)
+            seg.v.copy_(value)
+        else:
+            seg.k[:, :, head_slice, :] = key
+            seg.v[:, :, head_slice, :] = value
+
+    def _pack(self, seg: _Segment) -> None:
+        if seg.packed or seg.k is None:
+            return
+        triton_prq_quantize_tensor, _ = _qvg_functions()
+
+        def q(x):
+            xb = x.permute(0, 2, 1, 3).contiguous()  # [B,S,H,D]->[B,H,S,D]
+            devices = [xb.device] if xb.is_cuda else []
+            with torch.random.fork_rng(devices=devices):
+                torch.manual_seed(1234)
+                return triton_prq_quantize_tensor(
+                    xb,
+                    num_stages=self.q.stages,
+                    num_clusters=self.q.centroids,
+                    block_size=self.q.block_size,
+                    max_iters=self.q.kmeans_iters,
+                    quantize_fn=lambda _t: self.q.bits,
+                    asymmetric=self.q.asymmetric,
+                )
+
+        seg.packed_k = q(seg.k)
+        seg.packed_v = q(seg.v)
+        seg.k = None
+        seg.v = None
+        logger.info_once(f"Using QVG packed KV cache: {self.q.describe()}")
+
+    def _dequant(self, packed: dict) -> torch.Tensor:
+        _, triton_prq_dequantize_tensor = _qvg_functions()
+        return triton_prq_dequantize_tensor(
+            packed, self.q.block_size, self.q.bits, output_dtype=self.dtype
+        )  # [B,H,S,D]
+
+    def _all_segments(self) -> list[_Segment]:
+        segs = list(self._segments)
+        if self._cur is not None:
+            segs.append(self._cur)
+        return segs
+
+    def _sink_end(self) -> int:
+        return min(self.sink_tokens, self._global_end)
+
+    def _tail_start(self) -> int:
+        """Global start of the rolling recent tail (sink occupies its own
+        budget at the window front, matching the dense cache's roll)."""
+        sink_end = self._sink_end()
+        recent_budget = max(0, self.attention_window_size - sink_end)
+        return max(sink_end, self._global_end - recent_budget)
+
+    def _pack_and_evict(self) -> None:
+        """Pack completed segments older than the recency guard; drop segments
+        that have slid entirely out of the window (sink is never evicted)."""
+        tail_start = self._tail_start()
+
+        # eviction: drop non-sink segments fully left of the rolling tail
+        kept = []
+        for s in self._segments:
+            if not s.is_sink and s.g1 <= tail_start:
+                continue
+            kept.append(s)
+        self._segments = kept
+
+        if not self.q.enabled:
+            return
+        # recency guard: keep the newest `keep_recent_chunks` completed
+        # non-sink chunks in bf16; pack everything older.
+        recent = self.q.keep_recent_chunks
+        nonsink = [s for s in self._segments if not s.is_sink]
+        cutoff_idx = len(nonsink) - recent
+        for i, s in enumerate(nonsink):
+            if i < cutoff_idx:
+                self._pack(s)
+        # sink packing policy
+        if self.q.sink:
+            sink_keep_tokens = self.q.sink_keep_chunks * max(1, self._chunk_tokens)
+            for s in self._segments:
+                if s.is_sink and s.g0 >= sink_keep_tokens:
+                    # only pack sink chunks past the protected prefix, and only
+                    # once they are no longer the current recency-recent region
+                    self._pack(s)
+
+    # -------------------------------------------------------------- contract
+    def update_and_get_attention_kv(
+        self,
+        *,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        current_chunk_start: int,
+        cache_head_start: int | None = None,
+        recent_window_tokens: int | None = None,
+        debug_name: str = "QVG packed KV cache",
+    ) -> CausalAttentionKVView:
+        num_new = key.shape[1]
+        num_input_heads = key.shape[2]
+        head_slice = None
+        if num_input_heads != self.num_heads:
+            if cache_head_start is None:
+                raise ValueError(
+                    f"{debug_name}: cache_head_start required for head slice"
+                )
+            head_slice = slice(cache_head_start, cache_head_start + num_input_heads)
+        cend = current_chunk_start + num_new
+
+        if self._cur is not None and current_chunk_start == self._cur.g0:
+            # rewrite current chunk in place (denoise step)
+            if cend != self._cur.g1:
+                raise NotImplementedError(
+                    f"{debug_name}: current-chunk rewrite size changed"
+                )
+            self._write(self._cur, key, value, head_slice)
+        elif current_chunk_start == self._global_end:
+            # advance: finalize current chunk, start a new one
+            if self._cur is not None:
+                self._segments.append(self._cur)
+            is_sink = current_chunk_start < self.sink_tokens
+            if self._chunk_tokens == 0:
+                self._chunk_tokens = num_new
+            self._cur = self._new_bf16_segment(current_chunk_start, cend, is_sink)
+            self._write(self._cur, key, value, head_slice)
+            self._global_end = cend
+            self._pack_and_evict()
+        else:
+            raise NotImplementedError(
+                f"{debug_name}: non-sequential write current_start="
+                f"{current_chunk_start} global_end={self._global_end} "
+                f"cur={None if self._cur is None else self._cur.g0}"
+            )
+
+        local_end = min(self._global_end, self.cache_size)
+        if self.global_end_index_int is not None:
+            self.global_end_index_int = self._global_end
+            self.local_end_index_int = local_end
+        else:
+            self.global_end_index.fill_(self._global_end)
+            self.local_end_index.fill_(local_end)
+
+        vk, vv = self._reconstruct(current_chunk_start, recent_window_tokens)
+        return CausalAttentionKVView(
+            k=vk,
+            v=vv,
+            local_start_index=0,
+            local_end_index=num_new,
+            visible_local_end=min(self._global_end, self.cache_size),
+            visible_global_end=self._global_end,
+        )
+
+    def _reconstruct(
+        self,
+        current_chunk_start: int,
+        recent_window_tokens: int | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Dense visible window = sink prefix ++ rolling recent tail, matching
+        the dense cache's [sink | rolled-recent] buffer content."""
+        sink_end = self._sink_end()
+        if recent_window_tokens is None:
+            tail_start = self._tail_start()
+        else:
+            if recent_window_tokens < 0:
+                raise ValueError("recent_window_tokens must be >= 0 or None")
+            tail_start = max(sink_end, current_chunk_start - recent_window_tokens)
+
+        if tail_start <= sink_end:
+            ranges = [(0, self._global_end)]
+        else:
+            ranges = [(0, sink_end), (tail_start, self._global_end)]
+
+        visible_segments: list[tuple[_Segment, int, int]] = []
+        visible_tokens = 0
+        for g_lo, g_hi in ranges:
+            for seg in self._all_segments():
+                a = max(g_lo, seg.g0)
+                b = min(g_hi, seg.g1)
+                if b <= a:
+                    continue
+                visible_segments.append((seg, a - seg.g0, b - seg.g0))
+                visible_tokens += b - a
+
+        if len(visible_segments) == 1 and not visible_segments[0][0].packed:
+            seg, i0, i1 = visible_segments[0]
+            return seg.k[:, i0:i1], seg.v[:, i0:i1]
+
+        output_shape = (
+            self.batch_size,
+            visible_tokens,
+            self.num_heads,
+            self.head_dim,
+        )
+        vk = torch.empty(output_shape, dtype=self.dtype, device=self.device)
+        vv = torch.empty_like(vk)
+        output_start = 0
+
+        # Dequantize one tensor at a time so reconstruction needs only the
+        # final dense view plus one segment-sized temporary.
+        for seg, i0, i1 in visible_segments:
+            output_end = output_start + i1 - i0
+            if seg.packed:
+                dequantized = self._dequant(seg.packed_k)
+                vk[:, output_start:output_end].copy_(
+                    dequantized[:, :, i0:i1].permute(0, 2, 1, 3)
+                )
+                del dequantized
+
+                dequantized = self._dequant(seg.packed_v)
+                vv[:, output_start:output_end].copy_(
+                    dequantized[:, :, i0:i1].permute(0, 2, 1, 3)
+                )
+                del dequantized
+            else:
+                vk[:, output_start:output_end].copy_(seg.k[:, i0:i1])
+                vv[:, output_start:output_end].copy_(seg.v[:, i0:i1])
+            output_start = output_end
+
+        return vk, vv

python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py

@@ -12,6 +12,9 @@ from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import
     CausalSelfAttentionKVCache,
     CrossAttentionKVCache,
 )
+from sglang.multimodal_gen.runtime.layers.kvcache.qvg_packed_cache import (
+    QVGPackedCausalKVCache,
+)
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
     get_or_create_request_scheduler,
@@ -46,6 +49,7 @@ logger = init_logger(__name__)
 CAUSAL_BLOCK_PROMPTS_KEY = "causal_block_prompts"
 CAUSAL_SCENE_CUT_MASK_KEY = "causal_scene_cut_mask"
 CAUSAL_SHOT_INDICES_KEY = "causal_shot_indices"
+CausalKVCache = CausalSelfAttentionKVCache | QVGPackedCausalKVCache
 
 
 def expand_causal_block_prompts(
@@ -134,7 +138,7 @@ class CausalDMDCachePolicy:
 class CausalDMDRealtimeCacheContext:
     cache_state: RealtimeCausalDiTState
     persist_state: bool
-    kv_cache: list[CausalSelfAttentionKVCache]
+    kv_cache: list[CausalKVCache]
     crossattn_cache: list[CrossAttentionKVCache]
     current_start_frame: int
     chunk_idx: int
@@ -415,6 +419,17 @@ class CausalDMDDenoisingStage(DenoisingStage):
                 raise ValueError("realtime_causal_kv_cache_num_frames must be positive")
             self.sliding_window_num_frames = int(kv_cache_num_frames)
 
+        if (
+            server_args.kv_cache_quant_config.enabled
+            and not self._supports_qvg_kv_cache_quantization()
+        ):
+            raise ValueError(
+                f"{type(self).__name__} does not support QVG KV-cache quantization"
+            )
+
+    def _supports_qvg_kv_cache_quantization(self) -> bool:
+        return False
+
     def _causal_sequence_shard_enabled(self, batch: Req) -> bool:
         return False
 
@@ -487,8 +502,8 @@ class CausalDMDDenoisingStage(DenoisingStage):
             or crossattn_cache is None
             or len(causal_kv_cache) != self.num_transformer_blocks
             or len(crossattn_cache) != self.num_transformer_blocks
-            or causal_kv_cache[0].k.shape[1] != policy.expected_cache_tokens
-            or causal_kv_cache[0].k.shape[2] != policy.num_attention_heads
+            or causal_kv_cache[0].cache_size != policy.expected_cache_tokens
+            or causal_kv_cache[0].num_cache_heads != policy.num_attention_heads
             or causal_kv_cache[0].sink_tokens != policy.expected_sink_tokens
         )
 
@@ -1111,46 +1126,27 @@ class CausalDMDDenoisingStage(DenoisingStage):
         global_sink_tokens: int = 0,
         attention_window_size: int | None = None,
         allow_growth: bool = False,
-    ) -> list[CausalSelfAttentionKVCache]:
-        causal_kv_cache = []
-        int_index = 0 if use_int_indices else None
+    ) -> list[CausalKVCache]:
         if attention_window_size is None:
             attention_window_size = kv_cache_size
-        for _ in range(self.num_transformer_blocks):
-            causal_kv_cache.append(
-                CausalSelfAttentionKVCache(
-                    k=torch.zeros(
-                        [
-                            batch_size,
-                            kv_cache_size,
-                            num_attention_heads,
-                            attention_head_dim,
-                        ],
-                        dtype=dtype,
-                        device=device,
-                    ),
-                    v=torch.zeros(
-                        [
-                            batch_size,
-                            kv_cache_size,
-                            num_attention_heads,
-                            attention_head_dim,
-                        ],
-                        dtype=dtype,
-                        device=device,
-                    ),
-                    global_end_index=torch.zeros(1, dtype=torch.long, device=device),
-                    local_end_index=torch.zeros(1, dtype=torch.long, device=device),
-                    global_end_index_int=int_index,
-                    local_end_index_int=int_index,
-                    cache_size=kv_cache_size,
-                    sink_tokens=sink_tokens,
-                    global_sink_tokens=global_sink_tokens,
-                    attention_window_size=attention_window_size,
-                    allow_growth=allow_growth,
-                )
+        int_index = 0 if use_int_indices else None
+        shape = [batch_size, kv_cache_size, num_attention_heads, attention_head_dim]
+        return [
+            CausalSelfAttentionKVCache(
+                k=torch.zeros(shape, dtype=dtype, device=device),
+                v=torch.zeros(shape, dtype=dtype, device=device),
+                global_end_index=torch.zeros(1, dtype=torch.long, device=device),
+                local_end_index=torch.zeros(1, dtype=torch.long, device=device),
+                global_end_index_int=int_index,
+                local_end_index_int=int_index,
+                cache_size=kv_cache_size,
+                sink_tokens=sink_tokens,
+                global_sink_tokens=global_sink_tokens,
+                attention_window_size=attention_window_size,
+                allow_growth=allow_growth,
             )
-        return causal_kv_cache
+            for _ in range(self.num_transformer_blocks)
+        ]
 
     @torch.no_grad()
     def forward(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py

@@ -12,6 +12,9 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_ring_parallel_world_size,
     get_ulysses_parallel_world_size,
 )
+from sglang.multimodal_gen.runtime.layers.kvcache.qvg_packed_cache import (
+    QVGPackedCausalKVCache,
+)
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import (
@@ -19,6 +22,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import
     CausalDMDDenoisingStage,
     CausalDMDForwardContext,
     CausalDMDRealtimeCacheContext,
+    CausalKVCache,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.constants import (
     LINGBOT_C2WS_PLUCKER_EMB_CACHE,
@@ -48,6 +52,63 @@ class LingBotWorldCausalDMDDenoisingStage(CausalDMDDenoisingStage):
     Each call processes one chunk (num_frames_per_block frames).
     """
 
+    def _supports_qvg_kv_cache_quantization(self) -> bool:
+        return True
+
+    def _allocate_causal_kv_cache(
+        self,
+        *,
+        batch_size: int,
+        kv_cache_size: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dtype: torch.dtype,
+        device,
+        use_int_indices: bool = False,
+        sink_tokens: int = 0,
+        global_sink_tokens: int = 0,
+        attention_window_size: int | None = None,
+        allow_growth: bool = False,
+    ) -> list[CausalKVCache]:
+        if not self._kv_quant_args.enabled:
+            return super()._allocate_causal_kv_cache(
+                batch_size=batch_size,
+                kv_cache_size=kv_cache_size,
+                num_attention_heads=num_attention_heads,
+                attention_head_dim=attention_head_dim,
+                dtype=dtype,
+                device=device,
+                use_int_indices=use_int_indices,
+                sink_tokens=sink_tokens,
+                global_sink_tokens=global_sink_tokens,
+                attention_window_size=attention_window_size,
+                allow_growth=allow_growth,
+            )
+        if global_sink_tokens or allow_growth:
+            raise NotImplementedError(
+                "QVG packed KV cache supports only the LingBot realtime "
+                "sliding-window and sink path"
+            )
+        if attention_window_size is None:
+            attention_window_size = kv_cache_size
+        return [
+            QVGPackedCausalKVCache(
+                batch_size=batch_size,
+                cache_size=kv_cache_size,
+                num_heads=num_attention_heads,
+                head_dim=attention_head_dim,
+                dtype=dtype,
+                device=device,
+                use_int_indices=use_int_indices,
+                global_end_index=torch.zeros(1, dtype=torch.long, device=device),
+                local_end_index=torch.zeros(1, dtype=torch.long, device=device),
+                sink_tokens=sink_tokens,
+                attention_window_size=attention_window_size,
+                quant_args=self._kv_quant_args,
+            )
+            for _ in range(self.num_transformer_blocks)
+        ]
+
     def _get_causal_kv_cache_size(
         self,
         *,
@@ -141,6 +202,7 @@ class LingBotWorldCausalDMDDenoisingStage(CausalDMDDenoisingStage):
     ) -> None:
         self._reset_causal_cache_config_defaults()
         super()._apply_causal_cache_overrides(batch, server_args)
+        self._kv_quant_args = server_args.kv_cache_quant_config
         self._sync_interactive_kv_cache_window(server_args)
 
     def _reset_causal_cache_config_defaults(self) -> None:

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -26,6 +26,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
     is_ltx23_native_variant,
 )
 from sglang.multimodal_gen.configs.quantization.nunchaku import NunchakuSVDQuantArgs
+from sglang.multimodal_gen.configs.quantization.qvg_kv import QVGKVQuantArgs
 from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
     NunchakuConfig,
@@ -358,6 +359,12 @@ class ServerArgs(DisaggServerArgsMixin):
         default_factory=NunchakuSVDQuantArgs, repr=False
     )
 
+    # KV-cache quantization (Quant-VideoGen PRQ). Off by default; mirrors the
+    # SRT --kv-cache-dtype pattern (typed config, not a pile of env vars).
+    kv_cache_quant_config: QVGKVQuantArgs = field(
+        default_factory=QVGKVQuantArgs, repr=False
+    )
+
     # Master port for distributed inference
     master_port: int = 30005
 
@@ -1814,6 +1821,67 @@ class ServerArgs(DisaggServerArgsMixin):
             help="Disable autocast for denoising loop and vae decoding in pipeline sampling",
         )
 
+        # KV-cache quantization (Quant-VideoGen PRQ)
+        parser.add_argument(
+            "--kv-cache-quant",
+            type=str,
+            default=None,
+            choices=["off", "int4", "int2"],
+            help="Enable Quant-VideoGen PRQ KV-cache quantization (off|int4|int2). "
+            "Defaults reproduce the tuned per-chunk config (stages=1, "
+            "centroids=128, block=64, symmetric, iters=2, recent=1, "
+            "per-chunk sink).",
+        )
+        parser.add_argument(
+            "--kv-cache-quant-centroids",
+            type=int,
+            default=None,
+            help="PRQ k-means centroids per stage (default 128).",
+        )
+        parser.add_argument(
+            "--kv-cache-quant-block-size",
+            type=int,
+            default=None,
+            help="PRQ residual scale block size (default 64).",
+        )
+        parser.add_argument(
+            "--kv-cache-quant-stages",
+            type=int,
+            default=None,
+            help="PRQ k-means stages (default 1).",
+        )
+        parser.add_argument(
+            "--kv-cache-quant-iters",
+            type=int,
+            default=None,
+            help="PRQ k-means iterations (default 2).",
+        )
+        parser.add_argument(
+            "--kv-cache-quant-asymmetric",
+            action="store_true",
+            default=None,
+            help="Use KIVI-style asymmetric residual quantization.",
+        )
+        parser.add_argument(
+            "--kv-cache-quant-keep-recent",
+            type=int,
+            default=None,
+            help="Completed chunks kept bf16 before quantizing (default 1).",
+        )
+        parser.add_argument(
+            "--kv-cache-quant-sink",
+            type=int,
+            default=None,
+            choices=[0, 1],
+            help="Quantize the attention sink too (1, default) " "or keep it bf16 (0).",
+        )
+        parser.add_argument(
+            "--kv-cache-quant-sink-keep",
+            type=int,
+            default=None,
+            help="Leading sink chunks kept bf16 forever (default 0).",
+        )
+
         # quantization
         parser.add_argument(
             "--quantization",
@@ -2311,6 +2379,19 @@ class ServerArgs(DisaggServerArgsMixin):
             elif attr == "nunchaku_config":
                 nunchaku_config = NunchakuSVDQuantArgs.from_dict(kwargs)
                 server_args_kwargs["nunchaku_config"] = nunchaku_config
+            elif attr == "kv_cache_quant_config":
+                kv_quant_config = kwargs.get("kv_cache_quant_config")
+                if kv_quant_config is None:
+                    kv_quant_config = QVGKVQuantArgs.from_dict(kwargs)
+                elif isinstance(kv_quant_config, dict):
+                    kv_quant_config = QVGKVQuantArgs(**kv_quant_config).validate()
+                elif isinstance(kv_quant_config, QVGKVQuantArgs):
+                    kv_quant_config.validate()
+                else:
+                    raise TypeError(
+                        "kv_cache_quant_config must be QVGKVQuantArgs or a dict"
+                    )
+                server_args_kwargs["kv_cache_quant_config"] = kv_quant_config
             elif attr in kwargs:
                 server_args_kwargs[attr] = kwargs[attr]
 

python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py

@@ -7,6 +7,7 @@ import torch
 from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import (
     LingBotWorldCausalDMDConfig,
 )
+from sglang.multimodal_gen.configs.quantization.qvg_kv import QVGKVQuantArgs
 from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import (
     CausalSelfAttentionKVCache,
     CrossAttentionKVCache,
@@ -72,12 +73,16 @@ def test_lingbot_realtime_cache_config_overrides_checkpoint_defaults():
     stage.sink_size = 9
     stage.sliding_window_num_frames = 18
     stage.num_token_per_frame = 10
+    stage.transformer = SimpleNamespace(
+        config=SimpleNamespace(arch_config=SimpleNamespace())
+    )
 
     server_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             realtime_causal_sink_size=3,
             realtime_causal_kv_cache_num_frames=45,
-        )
+        ),
     )
     stage._apply_causal_cache_overrides(SimpleNamespace(), server_args)
 
@@ -93,16 +98,20 @@ def test_lingbot_realtime_cache_config_uses_request_overrides():
     stage.sink_size = 9
     stage.sliding_window_num_frames = 18
     stage.num_token_per_frame = 10
+    stage.transformer = SimpleNamespace(
+        config=SimpleNamespace(arch_config=SimpleNamespace())
+    )
 
     batch = SimpleNamespace(
         realtime_causal_sink_size=4,
         realtime_causal_kv_cache_num_frames=12,
     )
     server_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             realtime_causal_sink_size=3,
             realtime_causal_kv_cache_num_frames=45,
-        )
+        ),
     )
     stage._apply_causal_cache_overrides(batch, server_args)
 
@@ -120,6 +129,7 @@ def test_lingbot_realtime_attention_cache_rolls_with_sink_window():
     stage.num_token_per_frame = 1
     stage.num_frames_per_block = 3
     stage.sliding_window_num_frames = 6
+    stage._kv_quant_args = QVGKVQuantArgs()
     stage.transformer = SimpleNamespace(
         num_attention_heads=1,
         attention_head_dim=1,
@@ -279,12 +289,13 @@ def test_lingbot_interactive_kv_window_samples_base_moving_and_still(monkeypatch
     stage.sliding_window_num_frames = 18
     stage.transformer = SimpleNamespace(num_attention_heads=1)
     server_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             interactive_kv_window_enable=True,
             interactive_kv_moving_window=12,
             interactive_kv_still_window=3,
             interactive_kv_still_chunks=2,
-        )
+        ),
     )
     cache_state = RealtimeCausalDiTState()
 
@@ -323,6 +334,7 @@ def test_lingbot_interactive_kv_window_none_disables_moving_window(monkeypatch):
     stage.sliding_window_num_frames = 18
     stage.transformer = SimpleNamespace(num_attention_heads=1)
     server_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             realtime_causal_sink_size=9,
             realtime_causal_kv_cache_num_frames=18,
@@ -330,7 +342,7 @@ def test_lingbot_interactive_kv_window_none_disables_moving_window(monkeypatch):
             interactive_kv_moving_window=None,
             interactive_kv_still_window=3,
             interactive_kv_still_chunks=2,
-        )
+        ),
     )
     cache_state = RealtimeCausalDiTState()
     batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]})
@@ -353,6 +365,7 @@ def test_lingbot_interactive_kv_window_zero_is_valid_moving_window(monkeypatch):
     stage.sliding_window_num_frames = 18
     stage.transformer = SimpleNamespace(num_attention_heads=1)
     server_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             realtime_causal_sink_size=9,
             realtime_causal_kv_cache_num_frames=18,
@@ -360,7 +373,7 @@ def test_lingbot_interactive_kv_window_zero_is_valid_moving_window(monkeypatch):
             interactive_kv_moving_window=0,
             interactive_kv_still_window=None,
             interactive_kv_still_chunks=2,
-        )
+        ),
     )
     cache_state = RealtimeCausalDiTState()
     batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]})
@@ -385,6 +398,7 @@ def test_lingbot_interactive_kv_window_updates_total_window_for_moving_default(
     stage.sliding_window_num_frames = 18
     stage.transformer = SimpleNamespace(num_attention_heads=1)
     server_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             realtime_causal_sink_size=9,
             realtime_causal_kv_cache_num_frames=18,
@@ -392,7 +406,7 @@ def test_lingbot_interactive_kv_window_updates_total_window_for_moving_default(
             interactive_kv_moving_window=12,
             interactive_kv_still_window=3,
             interactive_kv_still_chunks=2,
-        )
+        ),
     )
     batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]})
 
@@ -424,6 +438,7 @@ def test_lingbot_interactive_kv_window_resets_stage_window_between_requests(
         ),
     )
     dynamic_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             realtime_causal_sink_size=9,
             realtime_causal_kv_cache_num_frames=18,
@@ -431,15 +446,16 @@ def test_lingbot_interactive_kv_window_resets_stage_window_between_requests(
             interactive_kv_moving_window=12,
             interactive_kv_still_window=3,
             interactive_kv_still_chunks=2,
-        )
+        ),
     )
     disabled_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             interactive_kv_window_enable=False,
             interactive_kv_moving_window=12,
             interactive_kv_still_window=3,
             interactive_kv_still_chunks=2,
-        )
+        ),
     )
 
     dynamic_batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"]]})
@@ -469,6 +485,7 @@ def test_lingbot_interactive_kv_window_default_disabled(monkeypatch):
     stage.sliding_window_num_frames = 18
     stage.transformer = SimpleNamespace(num_attention_heads=1)
     server_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             realtime_causal_sink_size=9,
             realtime_causal_kv_cache_num_frames=18,
@@ -476,7 +493,7 @@ def test_lingbot_interactive_kv_window_default_disabled(monkeypatch):
             interactive_kv_moving_window=12,
             interactive_kv_still_window=3,
             interactive_kv_still_chunks=2,
-        )
+        ),
     )
     cache_state = RealtimeCausalDiTState()
     batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]})
@@ -501,6 +518,7 @@ def test_lingbot_interactive_kv_window_env_can_enable_default(monkeypatch):
     stage.sliding_window_num_frames = 18
     stage.transformer = SimpleNamespace(num_attention_heads=1)
     server_args = SimpleNamespace(
+        kv_cache_quant_config=QVGKVQuantArgs(),
         pipeline_config=SimpleNamespace(
             realtime_causal_sink_size=9,
             realtime_causal_kv_cache_num_frames=18,
@@ -508,7 +526,7 @@ def test_lingbot_interactive_kv_window_env_can_enable_default(monkeypatch):
             interactive_kv_moving_window=12,
             interactive_kv_still_window=3,
             interactive_kv_still_chunks=2,
-        )
+        ),
     )
     cache_state = RealtimeCausalDiTState()
     batch = SimpleNamespace(condition_inputs={"camera_actions": [["w"], [], []]})
@@ -531,6 +549,7 @@ def test_lingbot_interactive_kv_window_allocates_expected_cache_size():
     stage.num_token_per_frame = 10
     stage.num_frames_per_block = 3
     stage.sliding_window_num_frames = 18
+    stage._kv_quant_args = QVGKVQuantArgs()
     stage.transformer = SimpleNamespace(num_attention_heads=1, attention_head_dim=1)
     policy = CausalDMDCachePolicy(
         sequence_shard_enabled=False,

python/sglang/multimodal_gen/test/unit/test_qvg_packed_kv.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for Quant-VideoGen packed KV-cache (QVGPackedCausalKVCache).
+
+- config parsing + defaults: CPU, no deps.
+- packed-storage equivalence vs the dense cache (quant disabled -> bf16
+  segments): CPU-only, exercises the segment / sliding-window / eviction /
+  reconstruction logic bit-exactly without needing quant-videogen or a GPU.
+- quant memory saving + reconstruction: requires CUDA + quant-videogen.
+"""
+
+import importlib.util
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.multimodal_gen.configs.quantization.qvg_kv import QVGKVQuantArgs
+from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import (
+    CausalSelfAttentionKVCache,
+)
+from sglang.multimodal_gen.runtime.layers.kvcache.qvg_packed_cache import (
+    QVGPackedCausalKVCache,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import (
+    CausalDMDDenoisingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.lingbot_world_causal_denoising import (
+    LingBotWorldCausalDMDDenoisingStage,
+)
+
+_HAS_QVG = importlib.util.find_spec("quant_videogen") is not None
+_HAS_CUDA = torch.cuda.is_available()
+
+
+def _base(B, W, H, D, sink, dev):
+    return CausalSelfAttentionKVCache(
+        k=torch.zeros(B, W, H, D, device=dev),
+        v=torch.zeros(B, W, H, D, device=dev),
+        global_end_index=torch.zeros(1, dtype=torch.long, device=dev),
+        local_end_index=torch.zeros(1, dtype=torch.long, device=dev),
+        cache_size=W,
+        sink_tokens=sink,
+        attention_window_size=W,
+    )
+
+
+def _packed(B, W, H, D, sink, dev, quant, use_int_indices=False):
+    return QVGPackedCausalKVCache(
+        batch_size=B,
+        cache_size=W,
+        num_heads=H,
+        head_dim=D,
+        dtype=torch.float32,
+        device=torch.device(dev),
+        use_int_indices=use_int_indices,
+        global_end_index=torch.zeros(1, dtype=torch.long, device=dev),
+        local_end_index=torch.zeros(1, dtype=torch.long, device=dev),
+        sink_tokens=sink,
+        attention_window_size=W,
+        quant_args=quant,
+    )
+
+
+def _replay(cache, chunks, C, H, D, B, nsteps, rwt, data, dev):
+    views = []
+    for ci in range(chunks):
+        g = ci * C
+        for step in range(nsteps):
+            k, v = data[(ci, step)]
+            view = cache.update_and_get_attention_kv(
+                key=k, value=v, current_chunk_start=g, recent_window_tokens=rwt
+            )
+        views.append((view.k.clone(), view.v.clone()))
+    return views
+
+
+class TestQVGKVQuantArgs(unittest.TestCase):
+    def test_off_by_default(self):
+        self.assertFalse(QVGKVQuantArgs().enabled)
+        self.assertFalse(QVGKVQuantArgs.from_dict({}).enabled)
+
+    def test_master_flag_defaults(self):
+        a = QVGKVQuantArgs.from_dict({"kv_cache_quant": "int4"})
+        self.assertTrue(a.enabled)
+        self.assertEqual(
+            (
+                a.bits,
+                a.stages,
+                a.centroids,
+                a.block_size,
+                a.kmeans_iters,
+                a.asymmetric,
+                a.keep_recent_chunks,
+                a.sink,
+                a.sink_keep_chunks,
+            ),
+            (4, 1, 128, 64, 2, False, 1, True, 0),
+        )
+
+    def test_overrides_and_validation(self):
+        a = QVGKVQuantArgs.from_dict(
+            {
+                "kv_cache_quant": "int2",
+                "kv_cache_quant_stages": 3,
+                "kv_cache_quant_sink": 0,
+            }
+        )
+        self.assertEqual((a.bits, a.stages, a.sink), (2, 3, False))
+        with self.assertRaises(ValueError):
+            QVGKVQuantArgs.from_dict({"kv_cache_quant": "int8"})
+
+
+class TestPackedStorageEquivalence(unittest.TestCase):
+    """quant disabled -> packed cache stores bf16 segments; must match the
+    dense cache bit-for-bit across the sliding-window + eviction lifecycle."""
+
+    def _run(self, rwt):
+        B, H, D, C, sink, chunks, nsteps = 1, 2, 4, 8, 8, 8, 3
+        W = C * 5
+        dev = "cpu"
+        torch.manual_seed(0)
+        data = {
+            (ci, s): (torch.randn(B, C, H, D), torch.randn(B, C, H, D))
+            for ci in range(chunks)
+            for s in range(nsteps)
+        }
+        base = _base(B, W, H, D, sink, dev)
+        packed = _packed(B, W, H, D, sink, dev, QVGKVQuantArgs())  # disabled
+        ob = _replay(base, chunks, C, H, D, B, nsteps, rwt, data, dev)
+        op = _replay(packed, chunks, C, H, D, B, nsteps, rwt, data, dev)
+        for i, ((bk, bv), (pk, pv)) in enumerate(zip(ob, op)):
+            self.assertEqual(bk.shape, pk.shape, f"chunk {i} shape")
+            self.assertTrue(torch.equal(bk, pk), f"chunk {i} K mismatch")
+            self.assertTrue(torch.equal(bv, pv), f"chunk {i} V mismatch")
+
+    def test_full_window(self):
+        self._run(rwt=None)
+
+    def test_recent_window_selection(self):
+        self._run(rwt=16)
+
+    def test_host_index_cursors_reset_without_device_updates(self):
+        cache = _packed(
+            B=1,
+            W=16,
+            H=2,
+            D=4,
+            sink=0,
+            dev="cpu",
+            quant=QVGKVQuantArgs(),
+            use_int_indices=True,
+        )
+        cache.update_and_get_attention_kv(
+            key=torch.ones(1, 4, 2, 4),
+            value=torch.ones(1, 4, 2, 4),
+            current_chunk_start=0,
+        )
+        self.assertEqual(cache.global_end_index_int, 4)
+        self.assertEqual(cache.local_end_index_int, 4)
+        self.assertEqual(int(cache.global_end_index.item()), 0)
+        self.assertEqual(int(cache.local_end_index.item()), 0)
+
+        cache.reset_indices()
+        self.assertEqual(cache.global_end_index_int, 0)
+        self.assertEqual(cache.local_end_index_int, 0)
+
+    def test_quantization_is_rejected_by_unsupported_causal_stages(self):
+        stage = CausalDMDDenoisingStage.__new__(CausalDMDDenoisingStage)
+        stage.sink_size = 0
+        stage.sliding_window_num_frames = 1
+
+        with self.assertRaisesRegex(ValueError, "does not support QVG"):
+            stage._apply_causal_cache_overrides(
+                SimpleNamespace(),
+                SimpleNamespace(
+                    pipeline_config=SimpleNamespace(),
+                    kv_cache_quant_config=QVGKVQuantArgs(bits=4),
+                ),
+            )
+
+    def test_quantization_is_available_for_lingbot(self):
+        stage = LingBotWorldCausalDMDDenoisingStage.__new__(
+            LingBotWorldCausalDMDDenoisingStage
+        )
+        stage._kv_quant_args = QVGKVQuantArgs(bits=4)
+        stage.num_transformer_blocks = 1
+
+        caches = stage._allocate_causal_kv_cache(
+            batch_size=1,
+            kv_cache_size=8,
+            num_attention_heads=1,
+            attention_head_dim=4,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+        )
+
+        self.assertIsInstance(caches[0], QVGPackedCausalKVCache)
+
+
+@unittest.skipUnless(
+    _HAS_CUDA and _HAS_QVG,
+    "needs CUDA + quant-videogen",
+)
+class TestPackedQuantMemory(unittest.TestCase):
+    def test_saves_memory_and_reconstructs(self):
+        B, H, D, C, sink, chunks, nsteps = 1, 8, 128, 512, 512, 12, 2
+        W = C * 6
+        dev = "cuda"
+        q = QVGKVQuantArgs.from_dict({"kv_cache_quant": "int4"})
+        torch.manual_seed(1)
+        data = {
+            (ci, s): (
+                torch.randn(B, C, H, D, device=dev),
+                torch.randn(B, C, H, D, device=dev),
+            )
+            for ci in range(chunks)
+            for s in range(nsteps)
+        }
+        base = _base(B, W, H, D, sink, dev)
+        packed = _packed(B, W, H, D, sink, dev, q)
+        ob = _replay(base, chunks, C, H, D, B, nsteps, None, data, dev)
+        op = _replay(packed, chunks, C, H, D, B, nsteps, None, data, dev)
+        dense_bytes = base.k.numel() * base.k.element_size() * 2
+        self.assertLess(
+            packed.resident_nbytes(),
+            dense_bytes * 0.7,
+            "packed cache should save >30% resident memory",
+        )
+        bk, _ = ob[-1]
+        pk, _ = op[-1]
+        self.assertEqual(bk.shape, pk.shape)
+        mse = ((bk - pk) ** 2).mean().item()
+        self.assertLess(mse, 1.0, "int4 reconstruction error too large")
+
+
+if __name__ == "__main__":
+    unittest.main()