Sync main: low-VRAM live test results (624x360)

b9fc164 verified 6 days ago

13.8 kB

from __future__ import annotations

import os
import time
from typing import Any, Callable, Optional

import torch

from xvideo.models.dit.rope import apply_rotary_emb
from xvideo.lowvram import tight_vram_enabled
from xvideo.models.dit import kv_entry_from_tensors, kv_entry_tensors, set_text_kv_active

GRAPH_ENV = "JOYOMNI_CUDA_GRAPH"
TEXT_KV_ENV = "JOYOMNI_TEXT_KV"

# The captured graph's KV pool is laid out as exactly [sink, prev1(, ref)]:
# an eager window [0, k-1, k] of this length is the only shape it can serve.
GRAPH_WINDOW_CHUNKS = 3


def graph_env_enabled() -> bool:
    if tight_vram_enabled() or os.environ.get("JOYOMNI_STREAM_DIT", "").strip().lower() in {
        "1", "true", "yes", "on"
    }:
        # Graph static buffers need the weights resident; streaming keeps them on host.
        return False
    return os.environ.get(GRAPH_ENV, "1").lower() in {"1", "true", "yes", "on"}


def text_kv_env_enabled() -> bool:
    return os.environ.get(TEXT_KV_ENV, "0").lower() in {"1", "true", "yes", "on"}


class StreamingGraphRunner:
    def __init__(
        self,
        transformer,
        *,
        chunk_tokens: int,
        ref_tokens: int,
        latent_shape: tuple[int, ...],
        txt_len: int,
        text_kv_by_step: Optional[list[dict[int, tuple[torch.Tensor, torch.Tensor]]]] = None,
        max_temporal_ids: int,
        pos_freqs: tuple[torch.Tensor, torch.Tensor],
        device: torch.device,
        dtype: torch.dtype,
        autocast_ctx: Callable[[], Any],
    ):
        self.transformer = transformer
        self.device = device
        self.dtype = dtype
        self.autocast_ctx = autocast_ctx
        self.num_layers = len(transformer.double_blocks)
        heads = int(transformer.config.heads_num)
        head_dim = int(transformer.config.hidden_size) // heads

        self.chunk_tokens = chunk_tokens
        self.ref_tokens = int(ref_tokens)
        self.txt_len = int(txt_len)
        # Static per-step text K/V buffers: the concatenation shape is baked into
        # the graph, the values are refreshed by bind_session() per session.
        self.text_kv_by_step: list[dict[int, tuple[torch.Tensor, torch.Tensor]]] = [
            {
                layer_idx: (
                    torch.zeros_like(key, device=device, dtype=dtype),
                    torch.zeros_like(value, device=device, dtype=dtype),
                )
                for layer_idx, (key, value) in cache.items()
            }
            for cache in (text_kv_by_step or ())
        ]
        self.text_kv_enabled = bool(self.text_kv_by_step)
        self.max_temporal_ids = int(max_temporal_ids)
        self.sink_off = 0
        self.prev1_off = chunk_tokens
        self.ref_off = 2 * chunk_tokens
        self.pool_len = 2 * chunk_tokens + self.ref_tokens

        L = self.num_layers
        self.pool_k = torch.zeros(L, 1, self.pool_len, heads, head_dim, device=device, dtype=dtype)
        self.pool_v = torch.zeros_like(self.pool_k)
        self.stage_k = torch.zeros(L, 1, chunk_tokens, heads, head_dim, device=device, dtype=dtype)
        self.stage_v = torch.zeros_like(self.stage_k)
        self.pub_k = torch.zeros(2, L, 1, chunk_tokens, heads, head_dim, device=device, dtype=dtype)
        self.pub_v = torch.zeros_like(self.pub_k)
        self._pub_parity = 0

        cos, sin = pos_freqs
        ct = chunk_tokens
        # Rope table for cached-KV positions 0..mti-1: row p = freqs of one chunk at temporal position p.
        self.pos_cos = cos.squeeze(0).view(self.max_temporal_ids, ct, -1).contiguous()
        self.pos_sin = sin.squeeze(0).view(self.max_temporal_ids, ct, -1).contiguous()
        # Writable buffer read by run_commit inside the graph; row is swapped per chunk via set_positions.
        self.commit_cos = self.pos_cos[self.max_temporal_ids - 1].clone().unsqueeze(0)
        self.commit_sin = self.pos_sin[self.max_temporal_ids - 1].clone().unsqueeze(0)

        self.in_latent = torch.zeros(latent_shape, device=device, dtype=dtype)
        self.in_timestep = torch.zeros((latent_shape[0],), device=device, dtype=torch.float32)
        self.in_ref_latent = torch.zeros(latent_shape, device=device, dtype=dtype)
        self.in_store_latent = torch.zeros(latent_shape, device=device, dtype=dtype)
        self.in_store_timestep = torch.zeros((latent_shape[0],), device=device, dtype=dtype)
        self.in_prompt = torch.zeros((1, self.txt_len, int(transformer.config.text_states_dim)),
                                     device=device, dtype=dtype)
        self.in_current_ids = torch.full((1, latent_shape[2]), self.max_temporal_ids,
                                         device=device, dtype=torch.long)

        self.full_graph: Optional[torch.cuda.CUDAGraph] = None
        self.in_noise = torch.zeros(latent_shape, device=device, dtype=dtype)
        self.out_latents: Optional[torch.Tensor] = None
        self.pool_dirty = True
        self.ready = False
        self.bound_token: Optional[int] = None
        self.last_graph_chunk = -(10 ** 9)

    def static_bytes(self) -> int:
        """Bytes held by the graph's static buffers (KV pool, staging, I/O)."""
        tensors = [
            self.pool_k, self.pool_v, self.stage_k, self.stage_v, self.pub_k, self.pub_v,
            self.in_latent, self.in_timestep, self.in_ref_latent, self.in_store_latent,
            self.in_store_timestep, self.in_prompt, self.in_current_ids, self.in_noise,
            self.out_latents, self.pos_cos, self.pos_sin, self.commit_cos, self.commit_sin,
        ]
        for cache in self.text_kv_by_step:
            for key, value in cache.values():
                tensors.extend((key, value))
        return sum(t.numel() * t.element_size() for t in tensors if isinstance(t, torch.Tensor))

    def _pool_assembler(self, layer_idx, *, device, dtype, cached_freqs_cis=None):
        return self.pool_k[layer_idx], self.pool_v[layer_idx]

    def _stage_writer(self, layer_idx, key, value):
        self.stage_k[layer_idx].copy_(key)
        self.stage_v[layer_idx].copy_(value)

    def capture(self, *, timesteps: torch.Tensor, sigmas: torch.Tensor) -> None:
        tf = self.transformer
        t0 = time.time()
        mem_pool = torch.cuda.graph_pool_handle()
        torch.cuda.synchronize(self.device)

        def run_denoise():
            with self.autocast_ctx():
                with tf.cache_context("cond"):
                    return tf(
                        hidden_states=self.in_latent,
                        timestep=self.in_timestep,
                        encoder_hidden_states=self.in_prompt,
                        ref_video_latent=self.in_ref_latent,
                        current_temporal_ids=self.in_current_ids,
                        cached_temporal_ids=None,
                        kv_cache_mode="reuse",
                        kv_cache_scope="cond",
                        kv_cache_chunk_id=None,
                        kv_cache_selected_chunk_ids=[],
                        kv_cache_pre_rope=True,
                        skip_text_stream=self.text_kv_enabled,
                    )[0]

        def run_store():
            with self.autocast_ctx():
                with tf.cache_context("cond"):
                    tf(
                        hidden_states=self.in_store_latent,
                        timestep=self.in_store_timestep,
                        encoder_hidden_states=self.in_prompt,
                        current_temporal_ids=self.in_current_ids,
                        cached_temporal_ids=None,
                        kv_cache_mode="store",
                        kv_cache_scope="cond",
                        kv_cache_chunk_id=None,
                        kv_cache_selected_chunk_ids=[],
                        kv_cache_pre_rope=True,
                        skip_text_stream=True,
                    )

        def run_commit():
            seg = slice(self.prev1_off, self.prev1_off + self.chunk_tokens)
            for li in range(self.num_layers):
                roped = apply_rotary_emb(self.stage_k[li], (self.commit_cos, self.commit_sin))
                self.pool_k[li, :, seg].copy_(roped)
                self.pool_v[li, :, seg].copy_(self.stage_v[li])

        ts_vals = [float(t) for t in timesteps]
        dt_vals = [float(sigmas[i + 1] - sigmas[i]) for i in range(len(ts_vals))]

        def run_full():
            lat32 = self.in_noise.to(torch.float32)
            tf._graph_kv_assembler = self._pool_assembler
            tf._graph_kv_writer = None
            for step_i, (t_val, dt) in enumerate(zip(ts_vals, dt_vals)):
                self.in_latent.copy_(lat32.to(self.dtype))
                self.in_timestep.fill_(t_val)
                set_text_kv_active(
                    self.text_kv_by_step[step_i] if self.text_kv_enabled else None
                )
                pred = run_denoise()
                lat32 = lat32 + pred.to(torch.float32) * dt
            self.in_store_latent.copy_(lat32.to(self.dtype))
            tf._graph_kv_assembler = None
            tf._graph_kv_writer = self._stage_writer
            # The store pass skips the text stream and must not see the cache.
            set_text_kv_active(None)
            try:
                run_store()
            finally:
                tf._graph_kv_writer = None
            run_commit()
            return lat32

        try:
            run_full()
            torch.cuda.synchronize(self.device)
            g = torch.cuda.CUDAGraph()
            with torch.cuda.graph(g, pool=mem_pool, capture_error_mode="thread_local"):
                self.out_latents = run_full()
            self.full_graph = g
        finally:
            tf._graph_kv_assembler = None
            tf._graph_kv_writer = None
            set_text_kv_active(None)

        torch.cuda.synchronize(self.device)
        self.ready = True
        static_gib = sum(t.numel() * t.element_size() for t in (
            self.pool_k, self.pool_v, self.stage_k, self.stage_v, self.pub_k, self.pub_v)) / 2 ** 30
        print(f"#####[GRAPH] captured whole-chunk graph in {time.time() - t0:.1f}s "
              f"(static bufs {static_gib:.2f} GiB, txt_len={self.txt_len}, "
              f"ref_tokens={self.ref_tokens}, mti={self.max_temporal_ids})", flush=True)

    def bind_session(
        self,
        token: int,
        prompt_embeds: torch.Tensor,
        text_kv_by_step: Optional[list[dict[int, tuple[torch.Tensor, torch.Tensor]]]] = None,
    ) -> None:
        self.in_prompt.copy_(prompt_embeds.to(dtype=self.dtype))
        for dst_cache, src_cache in zip(self.text_kv_by_step, text_kv_by_step or ()):
            for layer_idx, (key, value) in src_cache.items():
                dst = dst_cache.get(layer_idx)
                if dst is None:
                    continue
                dst[0].copy_(key.to(device=self.device, dtype=self.dtype))
                dst[1].copy_(value.to(device=self.device, dtype=self.dtype))
        self.bound_token = token
        self.pool_dirty = True
        self.last_graph_chunk = -(10 ** 9)

    def _pos_freqs(self, pos: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.pos_cos[pos].unsqueeze(0), self.pos_sin[pos].unsqueeze(0)

    def set_positions(self, chunk_idx: int) -> None:
        """Point the graph's position-dependent buffers at chunk_idx's temporal layout.
        Mirrors _gather_window_temporal_ids for the window [0, k-1, k]:
        current sits at min(k, mti), the committed KV (prev1 of chunk k+1) at min(k, mti-1).
        """
        mti = self.max_temporal_ids
        self.in_current_ids.fill_(min(chunk_idx, mti))
        commit_pos = min(chunk_idx, mti - 1)
        self.commit_cos.copy_(self.pos_cos[commit_pos].unsqueeze(0))
        self.commit_sin.copy_(self.pos_sin[commit_pos].unsqueeze(0))

    def _seed_segment(self, chunk_store: dict, off: int, length: int,
                      freqs: Optional[tuple[torch.Tensor, torch.Tensor]]) -> None:
        for li in range(self.num_layers):
            entry = chunk_store[li]
            k, v = kv_entry_tensors(entry, device=self.device, dtype=self.dtype)
            if entry.get("pre_rope", False):
                if freqs is None:
                    raise RuntimeError("graph seed: pre-rope entry in a no-rope segment")
                k = apply_rotary_emb(k, freqs)
            self.pool_k[li, :, off:off + length].copy_(k)
            self.pool_v[li, :, off:off + length].copy_(v)

    def seed_from_cache(self, scope_store: dict, *, sink_mem_id: int,
                        prev1_mem_id: int, prev1_pos: int, ref_mem_id: Optional[int]) -> None:
        for mem_id, off, freqs in (
            (sink_mem_id, self.sink_off, self._pos_freqs(0)),
            (prev1_mem_id, self.prev1_off, self._pos_freqs(prev1_pos)),
        ):
            store = scope_store.get(mem_id)
            if store is None:
                raise RuntimeError(f"graph seed: python cache missing mem_id={mem_id}")
            self._seed_segment(store, off, self.chunk_tokens, freqs)
        if self.ref_tokens > 0:
            store = scope_store.get(ref_mem_id)
            if store is None:
                raise RuntimeError(f"graph seed: python cache missing ref mem_id={ref_mem_id}")
            self._seed_segment(store, self.ref_off, self.ref_tokens, None)
        self.pool_dirty = False

    def publish_to_cache(self, scope_store: dict, mem_id: int) -> None:
        p = self._pub_parity
        self._pub_parity ^= 1
        self.pub_k[p].copy_(self.stage_k)
        self.pub_v[p].copy_(self.stage_v)
        chunk_store = {}
        for li in range(self.num_layers):
            chunk_store[li] = kv_entry_from_tensors(
                self.pub_k[p, li], self.pub_v[p, li], pre_rope=True
            )
        scope_store[mem_id] = chunk_store

	from __future__ import annotations

	import os
	import time
	from typing import Any, Callable, Optional

	import torch

	from xvideo.models.dit.rope import apply_rotary_emb
	from xvideo.lowvram import tight_vram_enabled
	from xvideo.models.dit import kv_entry_from_tensors, kv_entry_tensors, set_text_kv_active

	GRAPH_ENV = "JOYOMNI_CUDA_GRAPH"
	TEXT_KV_ENV = "JOYOMNI_TEXT_KV"

	# The captured graph's KV pool is laid out as exactly [sink, prev1(, ref)]:
	# an eager window [0, k-1, k] of this length is the only shape it can serve.
	GRAPH_WINDOW_CHUNKS = 3


	def graph_env_enabled() -> bool:
	if tight_vram_enabled() or os.environ.get("JOYOMNI_STREAM_DIT", "").strip().lower() in {
	"1", "true", "yes", "on"
	}:
	# Graph static buffers need the weights resident; streaming keeps them on host.
	return False
	return os.environ.get(GRAPH_ENV, "1").lower() in {"1", "true", "yes", "on"}


	def text_kv_env_enabled() -> bool:
	return os.environ.get(TEXT_KV_ENV, "0").lower() in {"1", "true", "yes", "on"}


	class StreamingGraphRunner:
	def __init__(
	self,
	transformer,
	*,
	chunk_tokens: int,
	ref_tokens: int,
	latent_shape: tuple[int, ...],
	txt_len: int,
	text_kv_by_step: Optional[list[dict[int, tuple[torch.Tensor, torch.Tensor]]]] = None,
	max_temporal_ids: int,
	pos_freqs: tuple[torch.Tensor, torch.Tensor],
	device: torch.device,
	dtype: torch.dtype,
	autocast_ctx: Callable[[], Any],
	):
	self.transformer = transformer
	self.device = device
	self.dtype = dtype
	self.autocast_ctx = autocast_ctx
	self.num_layers = len(transformer.double_blocks)
	heads = int(transformer.config.heads_num)
	head_dim = int(transformer.config.hidden_size) // heads

	self.chunk_tokens = chunk_tokens
	self.ref_tokens = int(ref_tokens)
	self.txt_len = int(txt_len)
	# Static per-step text K/V buffers: the concatenation shape is baked into
	# the graph, the values are refreshed by bind_session() per session.
	self.text_kv_by_step: list[dict[int, tuple[torch.Tensor, torch.Tensor]]] = [
	{
	layer_idx: (
	torch.zeros_like(key, device=device, dtype=dtype),
	torch.zeros_like(value, device=device, dtype=dtype),
	)
	for layer_idx, (key, value) in cache.items()
	}
	for cache in (text_kv_by_step or ())
	]
	self.text_kv_enabled = bool(self.text_kv_by_step)
	self.max_temporal_ids = int(max_temporal_ids)
	self.sink_off = 0
	self.prev1_off = chunk_tokens
	self.ref_off = 2 * chunk_tokens
	self.pool_len = 2 * chunk_tokens + self.ref_tokens

	L = self.num_layers
	self.pool_k = torch.zeros(L, 1, self.pool_len, heads, head_dim, device=device, dtype=dtype)
	self.pool_v = torch.zeros_like(self.pool_k)
	self.stage_k = torch.zeros(L, 1, chunk_tokens, heads, head_dim, device=device, dtype=dtype)
	self.stage_v = torch.zeros_like(self.stage_k)
	self.pub_k = torch.zeros(2, L, 1, chunk_tokens, heads, head_dim, device=device, dtype=dtype)
	self.pub_v = torch.zeros_like(self.pub_k)
	self._pub_parity = 0

	cos, sin = pos_freqs
	ct = chunk_tokens
	# Rope table for cached-KV positions 0..mti-1: row p = freqs of one chunk at temporal position p.
	self.pos_cos = cos.squeeze(0).view(self.max_temporal_ids, ct, -1).contiguous()
	self.pos_sin = sin.squeeze(0).view(self.max_temporal_ids, ct, -1).contiguous()
	# Writable buffer read by run_commit inside the graph; row is swapped per chunk via set_positions.
	self.commit_cos = self.pos_cos[self.max_temporal_ids - 1].clone().unsqueeze(0)
	self.commit_sin = self.pos_sin[self.max_temporal_ids - 1].clone().unsqueeze(0)

	self.in_latent = torch.zeros(latent_shape, device=device, dtype=dtype)
	self.in_timestep = torch.zeros((latent_shape[0],), device=device, dtype=torch.float32)
	self.in_ref_latent = torch.zeros(latent_shape, device=device, dtype=dtype)
	self.in_store_latent = torch.zeros(latent_shape, device=device, dtype=dtype)
	self.in_store_timestep = torch.zeros((latent_shape[0],), device=device, dtype=dtype)
	self.in_prompt = torch.zeros((1, self.txt_len, int(transformer.config.text_states_dim)),
	device=device, dtype=dtype)
	self.in_current_ids = torch.full((1, latent_shape[2]), self.max_temporal_ids,
	device=device, dtype=torch.long)

	self.full_graph: Optional[torch.cuda.CUDAGraph] = None
	self.in_noise = torch.zeros(latent_shape, device=device, dtype=dtype)
	self.out_latents: Optional[torch.Tensor] = None
	self.pool_dirty = True
	self.ready = False
	self.bound_token: Optional[int] = None
	self.last_graph_chunk = -(10 ** 9)

	def static_bytes(self) -> int:
	"""Bytes held by the graph's static buffers (KV pool, staging, I/O)."""
	tensors = [
	self.pool_k, self.pool_v, self.stage_k, self.stage_v, self.pub_k, self.pub_v,
	self.in_latent, self.in_timestep, self.in_ref_latent, self.in_store_latent,
	self.in_store_timestep, self.in_prompt, self.in_current_ids, self.in_noise,
	self.out_latents, self.pos_cos, self.pos_sin, self.commit_cos, self.commit_sin,
	]
	for cache in self.text_kv_by_step:
	for key, value in cache.values():
	tensors.extend((key, value))
	return sum(t.numel() * t.element_size() for t in tensors if isinstance(t, torch.Tensor))

	def _pool_assembler(self, layer_idx, *, device, dtype, cached_freqs_cis=None):
	return self.pool_k[layer_idx], self.pool_v[layer_idx]

	def _stage_writer(self, layer_idx, key, value):
	self.stage_k[layer_idx].copy_(key)
	self.stage_v[layer_idx].copy_(value)

	def capture(self, *, timesteps: torch.Tensor, sigmas: torch.Tensor) -> None:
	tf = self.transformer
	t0 = time.time()
	mem_pool = torch.cuda.graph_pool_handle()
	torch.cuda.synchronize(self.device)

	def run_denoise():
	with self.autocast_ctx():
	with tf.cache_context("cond"):
	return tf(
	hidden_states=self.in_latent,
	timestep=self.in_timestep,
	encoder_hidden_states=self.in_prompt,
	ref_video_latent=self.in_ref_latent,
	current_temporal_ids=self.in_current_ids,
	cached_temporal_ids=None,
	kv_cache_mode="reuse",
	kv_cache_scope="cond",
	kv_cache_chunk_id=None,
	kv_cache_selected_chunk_ids=[],
	kv_cache_pre_rope=True,
	skip_text_stream=self.text_kv_enabled,
	)[0]

	def run_store():
	with self.autocast_ctx():
	with tf.cache_context("cond"):
	tf(
	hidden_states=self.in_store_latent,
	timestep=self.in_store_timestep,
	encoder_hidden_states=self.in_prompt,
	current_temporal_ids=self.in_current_ids,
	cached_temporal_ids=None,
	kv_cache_mode="store",
	kv_cache_scope="cond",
	kv_cache_chunk_id=None,
	kv_cache_selected_chunk_ids=[],
	kv_cache_pre_rope=True,
	skip_text_stream=True,
	)

	def run_commit():
	seg = slice(self.prev1_off, self.prev1_off + self.chunk_tokens)
	for li in range(self.num_layers):
	roped = apply_rotary_emb(self.stage_k[li], (self.commit_cos, self.commit_sin))
	self.pool_k[li, :, seg].copy_(roped)
	self.pool_v[li, :, seg].copy_(self.stage_v[li])

	ts_vals = [float(t) for t in timesteps]
	dt_vals = [float(sigmas[i + 1] - sigmas[i]) for i in range(len(ts_vals))]

	def run_full():
	lat32 = self.in_noise.to(torch.float32)
	tf._graph_kv_assembler = self._pool_assembler
	tf._graph_kv_writer = None
	for step_i, (t_val, dt) in enumerate(zip(ts_vals, dt_vals)):
	self.in_latent.copy_(lat32.to(self.dtype))
	self.in_timestep.fill_(t_val)
	set_text_kv_active(
	self.text_kv_by_step[step_i] if self.text_kv_enabled else None
	)
	pred = run_denoise()
	lat32 = lat32 + pred.to(torch.float32) * dt
	self.in_store_latent.copy_(lat32.to(self.dtype))
	tf._graph_kv_assembler = None
	tf._graph_kv_writer = self._stage_writer
	# The store pass skips the text stream and must not see the cache.
	set_text_kv_active(None)
	try:
	run_store()
	finally:
	tf._graph_kv_writer = None
	run_commit()
	return lat32

	try:
	run_full()
	torch.cuda.synchronize(self.device)
	g = torch.cuda.CUDAGraph()
	with torch.cuda.graph(g, pool=mem_pool, capture_error_mode="thread_local"):
	self.out_latents = run_full()
	self.full_graph = g
	finally:
	tf._graph_kv_assembler = None
	tf._graph_kv_writer = None
	set_text_kv_active(None)

	torch.cuda.synchronize(self.device)
	self.ready = True
	static_gib = sum(t.numel() * t.element_size() for t in (
	self.pool_k, self.pool_v, self.stage_k, self.stage_v, self.pub_k, self.pub_v)) / 2 ** 30
	print(f"#####[GRAPH] captured whole-chunk graph in {time.time() - t0:.1f}s "
	f"(static bufs {static_gib:.2f} GiB, txt_len={self.txt_len}, "
	f"ref_tokens={self.ref_tokens}, mti={self.max_temporal_ids})", flush=True)

	def bind_session(
	self,
	token: int,
	prompt_embeds: torch.Tensor,
	text_kv_by_step: Optional[list[dict[int, tuple[torch.Tensor, torch.Tensor]]]] = None,
	) -> None:
	self.in_prompt.copy_(prompt_embeds.to(dtype=self.dtype))
	for dst_cache, src_cache in zip(self.text_kv_by_step, text_kv_by_step or ()):
	for layer_idx, (key, value) in src_cache.items():
	dst = dst_cache.get(layer_idx)
	if dst is None:
	continue
	dst[0].copy_(key.to(device=self.device, dtype=self.dtype))
	dst[1].copy_(value.to(device=self.device, dtype=self.dtype))
	self.bound_token = token
	self.pool_dirty = True
	self.last_graph_chunk = -(10 ** 9)

	def _pos_freqs(self, pos: int) -> tuple[torch.Tensor, torch.Tensor]:
	return self.pos_cos[pos].unsqueeze(0), self.pos_sin[pos].unsqueeze(0)

	def set_positions(self, chunk_idx: int) -> None:
	"""Point the graph's position-dependent buffers at chunk_idx's temporal layout.

	Mirrors _gather_window_temporal_ids for the window [0, k-1, k]:
	current sits at min(k, mti), the committed KV (prev1 of chunk k+1) at min(k, mti-1).
	"""
	mti = self.max_temporal_ids
	self.in_current_ids.fill_(min(chunk_idx, mti))
	commit_pos = min(chunk_idx, mti - 1)
	self.commit_cos.copy_(self.pos_cos[commit_pos].unsqueeze(0))
	self.commit_sin.copy_(self.pos_sin[commit_pos].unsqueeze(0))

	def _seed_segment(self, chunk_store: dict, off: int, length: int,
	freqs: Optional[tuple[torch.Tensor, torch.Tensor]]) -> None:
	for li in range(self.num_layers):
	entry = chunk_store[li]
	k, v = kv_entry_tensors(entry, device=self.device, dtype=self.dtype)
	if entry.get("pre_rope", False):
	if freqs is None:
	raise RuntimeError("graph seed: pre-rope entry in a no-rope segment")
	k = apply_rotary_emb(k, freqs)
	self.pool_k[li, :, off:off + length].copy_(k)
	self.pool_v[li, :, off:off + length].copy_(v)

	def seed_from_cache(self, scope_store: dict, *, sink_mem_id: int,
	prev1_mem_id: int, prev1_pos: int, ref_mem_id: Optional[int]) -> None:
	for mem_id, off, freqs in (
	(sink_mem_id, self.sink_off, self._pos_freqs(0)),
	(prev1_mem_id, self.prev1_off, self._pos_freqs(prev1_pos)),
	):
	store = scope_store.get(mem_id)
	if store is None:
	raise RuntimeError(f"graph seed: python cache missing mem_id={mem_id}")
	self._seed_segment(store, off, self.chunk_tokens, freqs)
	if self.ref_tokens > 0:
	store = scope_store.get(ref_mem_id)
	if store is None:
	raise RuntimeError(f"graph seed: python cache missing ref mem_id={ref_mem_id}")
	self._seed_segment(store, self.ref_off, self.ref_tokens, None)
	self.pool_dirty = False

	def publish_to_cache(self, scope_store: dict, mem_id: int) -> None:
	p = self._pub_parity
	self._pub_parity ^= 1
	self.pub_k[p].copy_(self.stage_k)
	self.pub_v[p].copy_(self.stage_v)
	chunk_store = {}
	for li in range(self.num_layers):
	chunk_store[li] = kv_entry_from_tensors(
	self.pub_k[p, li], self.pub_v[p, li], pre_rope=True
	)
	scope_store[mem_id] = chunk_store