From d9cfd34fab27b606a101106ec7f335e12fc1e026 Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Fri, 10 Jul 2026 13:17:39 -0400 Subject: [PATCH] Efficient fixed-length KV-cache for recurrent Transformer sampling The recurrent Transformer sampling path was both incorrect and slow: the estimator re-fed the whole committed prefix each step while the module threaded a growing carry, so the KV-cache and absolute positions diverged (cache grew 0,1,2,4,7 instead of 0,1,2,3,4). It survived CI only because the sole integration test asserts isfinite(loss). This makes the recurrent Transformer carry correct and adds an opt-in in-place KV-cache: - TransformerDiscreteSequenceModel.init_carry(max_len=...) preallocates [B, n_heads, max_len, head_dim] K/V buffers + a write cursor; the block writes new K/V in place (narrow().copy_(), TorchScript-safe) instead of torch.cat. max_len=None keeps the existing grow-by-cat carry. - RecurrentDiscretePolicyEstimator gains use_kv_cache: False (default) is a corrected grad-bearing full-prefix baseline; True feeds only the newest token against the persistent in-place cache under no_grad (O(L^2) vs the O(L^3) re-forward). A fixed-length guard rejects ragged active batches; variable-length carry slicing is out of scope (see the Gap-2 TODO in prob_calculations.py). Both paths build the canonical [BOS, w0, ..., w_{t-1}] stream and are numerically identical to one full-sequence forward. Tests (testing/test_kv_cache.py): dense-equivalence across shapes, in-place preallocation, decode-position sensitivity, use_kv_cache OFF==ON==teacher-forced, and the fixed-length guard. Benchmark (tutorials/examples/benchmark_kv_cache.py) reports the sampling speedup, which widens with sequence length (~5x at L=64, CPU). Stateless policies and the Sampler API are unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/gfn/estimators.py | 197 ++++++++++--- src/gfn/utils/modules.py | 145 ++++++++-- testing/test_kv_cache.py | 238 ++++++++++++++++ tutorials/examples/benchmark_kv_cache.py | 344 +++++++++++++++++++++++ 4 files changed, 860 insertions(+), 64 deletions(-) create mode 100644 testing/test_kv_cache.py create mode 100644 tutorials/examples/benchmark_kv_cache.py diff --git a/src/gfn/estimators.py b/src/gfn/estimators.py index b816cd76..12ffeb9e 100644 --- a/src/gfn/estimators.py +++ b/src/gfn/estimators.py @@ -1239,12 +1239,20 @@ class RecurrentDiscretePolicyEstimator(RecurrentPolicyMixin, DiscretePolicyEstim per-step artifacts. Non-recurrent estimators should use the default PolicyMixin and the standard ``DiscretePolicyEstimator`` base class instead. + An optional, fixed-length-only performance mode (opt-in; the default leaves + behavior unchanged): + + - ``use_kv_cache=True`` decodes one token per step against a persistent, in-place + preallocated KV-cache, cutting rollout cost from O(L^3) to O(L^2). The in-place + buffers cannot carry gradients, so cached sampling runs under ``no_grad`` -- it is + a sampling-time optimization, ideal for inference/rollout. + Notes ----- - - Forward is intended for on-policy generation; off-policy evaluation over - entire trajectories typically requires different batching and masking. - ``init_carry`` is a hard requirement for compatibility with the recurrent PolicyMixin. + - The cached fast path assumes fixed-length trajectories; see the Gap-2 TODO in + ``gfn/utils/prob_calculations.py``. Attributes: module: The neural network module to use. @@ -1261,14 +1269,27 @@ def __init__( preprocessor: Preprocessor | None = None, is_backward: bool = False, debug: bool = False, + use_kv_cache: bool = False, + cache_max_len: int | None = None, ): """Initializes a RecurrentDiscretePolicyEstimator. + See the class docstring for what ``use_kv_cache`` does and why cached sampling + is ``no_grad``. + Args: module: The neural network module to use. n_actions: Total number of actions in the discrete environment. preprocessor: Preprocessor object that transforms states to tensors. + is_backward: Flag indicating whether this estimator is for backward policy. debug: If True, enables expensive validation checks. + use_kv_cache: Opt into the O(L^2) in-place KV-cache sampling fast path. + Cached sampling runs under ``no_grad``, so it is a sampling-time + optimization (ideal for inference/rollout). If False (default), sampling + uses the corrected grad-bearing full-prefix forward each step. + cache_max_len: KV-cache preallocation length; defaults to the module's + ``max_position_embeddings``. Rarely set by hand; must be at least the + env horizon + 1 (the +1 is the BOS token). """ if preprocessor is None: preprocessor = IdentityPreprocessor(output_dim=None) @@ -1280,49 +1301,71 @@ def __init__( is_backward=is_backward, debug=debug, ) + self.use_kv_cache = use_kv_cache + # A preallocated in-place cache is only available for modules that expose a + # bounded position range (e.g. TransformerDiscreteSequenceModel). Modules + # without one (e.g. RNNs, whose carry is already fixed-size) decode one token + # per step with their ordinary carry. + max_positions = getattr(module, "max_position_embeddings", None) + if cache_max_len is None: + cache_max_len = max_positions + # Validate the cache budget eagerly (against the module's position range) so a + # too-small value fails here, not deep inside a rollout with a confusing error. + if ( + cache_max_len is not None + and max_positions is not None + and cache_max_len > max_positions + ): + raise ValueError( + f"cache_max_len ({cache_max_len}) exceeds the module's " + f"max_position_embeddings ({max_positions})." + ) + self._cache_max_len = cache_max_len + + @property + def _bos_index(self) -> int: + """BOS token index (the sequence model reserves ``vocab_size`` for BOS).""" + return int(getattr(self.module, "vocab_size", self.n_actions - 1)) def forward( self, states: States, carry: dict[str, torch.Tensor], ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: - """Forward pass of the module. + """Per-step forward: next-action logits for the current rollout states. + + Both paths build the same canonical token stream ``[BOS, w_0, ..., w_{t-1}]`` + and take the logits at the last position, so they are numerically equivalent + to one teacher-forced forward over the full sequence: + + - ``use_kv_cache=False`` (default): re-embed the whole committed prefix with a + fresh carry each step (grad-bearing, O(L^3) over a rollout). This is the + corrected baseline. + - ``use_kv_cache=True``: feed only the newest committed token against the + persistent carry (in-place KV-cache for transformers), under ``no_grad``. Args: - states: The input states. - carry: The carry from the previous step. + states: The current (active) rollout states, ``(B, W)`` word indices with + ``-1`` padding. + carry: The recurrent carry threaded from the previous step. Returns: - The output of the module, as a tensor of shape (*batch_shape, output_dim). + ``(logits (B, n_actions), updated carry)``. """ - # Prepare integer token sequences without -1 padding and use a BOS index. - # We infer the active sequence length per row from (token != -1). tokens = states.tensor.long() + if tokens.dim() != 2: + raise ValueError( + "RecurrentDiscretePolicyEstimator expects 2D state tensors " + f"(batch, seq); got shape {tuple(tokens.shape)}." + ) + lengths = (states.tensor >= 0).sum(dim=-1) # committed words per row - # Replace padding (-1) with BOS index expected by the sequence model. - # RecurrentDiscreteSequenceModel reserves index == vocab_size for BOS. - bos_index = getattr(self.module, "vocab_size", self.n_actions - 1) - tokens = torch.where( - tokens < 0, torch.as_tensor(bos_index, device=tokens.device), tokens - ) - - # Determine a common prefix length across the (active) batch. - # Active rows in a rollout step share the same length; use max for safety. - # We still derive length from original states.tensor where -1 marks padding. - valid_mask = states.tensor >= 0 - if valid_mask.ndim == 1: - max_len = int(valid_mask.sum().item()) + if self.use_kv_cache: + logits, carry = self._forward_cached(tokens, lengths, carry) else: - max_len = int(valid_mask.sum(dim=-1).max().item()) - if max_len == 0: - max_len = 1 # Ensure at least BOS is processed - - # Trim to the common active prefix length and run the sequence model. - seq_input = tokens[..., :max_len] - logits, carry = self.module(seq_input, carry) + logits, carry = self._forward_full_prefix(tokens, lengths) - # Use the logits corresponding to the last processed token. - logits = logits[:, -1, :] # (b, n_actions) + logits = logits[:, -1, :] # (B, n_actions) if self.expected_output_dim is not None: assert logits.shape[-1] == self.expected_output_dim, ( @@ -1332,20 +1375,104 @@ def forward( return logits, carry - def init_carry( + def _canonical_sequence( + self, tokens: torch.Tensor, lengths: torch.Tensor + ) -> torch.Tensor: + """Build ``[BOS, w_0, ..., w_{L-1}]`` up to the longest committed length ``L``. + + Used by the (no-cache) full-prefix path and the teacher-forced loss recompute. + """ + bos = self._bos_index + max_len = int(lengths.max().item()) + # Replace any -1 padding with BOS; unlike the cached path this one has no + # fixed-length guard, so shorter rows can carry padding inside [:max_len]. + words = torch.where(tokens < 0, torch.full_like(tokens, bos), tokens)[ + :, :max_len + ] + bos_col = torch.full( + (tokens.shape[0], 1), bos, dtype=torch.long, device=tokens.device + ) + return torch.cat([bos_col, words], dim=1) # (B, L + 1) + + def _forward_full_prefix( self, - batch_size: int, - device: torch.device, + tokens: torch.Tensor, + lengths: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Grad-bearing baseline: full prefix + fresh carry each step.""" + seq = self._canonical_sequence(tokens, lengths) + fresh = self._module_init_carry(seq.shape[0], seq.device, preallocate=False) + logits, new_carry = self.module(seq, fresh) + return logits, new_carry + + def _forward_cached( + self, + tokens: torch.Tensor, + lengths: torch.Tensor, + carry: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Cached fast path: feed only the newest committed token, under ``no_grad``. + + Assumes fixed-length rollouts (all active rows share a committed length). The + always-on backstop is the module's batch-size carry check (a masked, unsliced + carry mismatches the active batch); the stronger equal-length / cursor-sync + checks are gated behind ``debug`` so the hot path avoids per-step device syncs. + """ + length0 = int(lengths[0].item()) if lengths.numel() > 0 else 0 + + if self.debug and lengths.numel() > 0: + if not bool(torch.all(lengths == length0)): + raise ValueError( + "use_kv_cache=True requires equal-length active trajectories " + "(fixed-length env). Variable-length carry slicing is out of scope " + "(see the Gap-2 TODO in gfn/utils/prob_calculations.py)." + ) + cursor = int(carry["cache_len"].item()) if "cache_len" in carry else 0 + if cursor != length0: + raise ValueError( + f"KV-cache cursor ({cursor}) is out of sync with the committed " + f"length ({length0}); the carry must be threaded unmodified " + "across steps." + ) + + # Feed the canonical stream [BOS, w_0, ..., w_{L-1}] one token per step: at a + # committed length of ``length0`` that is column ``length0 - 1``, or BOS when + # nothing is committed yet. + if length0 == 0: + step_tok = torch.full( + (tokens.shape[0], 1), + self._bos_index, + dtype=torch.long, + device=tokens.device, + ) + else: + step_tok = tokens[:, length0 - 1 : length0] + + with torch.no_grad(): + logits, new_carry = self.module(step_tok, carry) + return logits, new_carry + + def _module_init_carry( + self, batch_size: int, device: torch.device, preallocate: bool ) -> dict[str, torch.Tensor]: init_carry = getattr(self.module, "init_carry", None) if not callable(init_carry): raise NotImplementedError( "Module does not implement init_carry(batch_size, device)." ) - init_carry_fn = cast(Callable[[int, torch.device], Any], init_carry) - + init_carry_fn = cast(Callable[..., dict[str, torch.Tensor]], init_carry) + if preallocate and self._cache_max_len is not None: + return init_carry_fn(batch_size, device, max_len=self._cache_max_len) return init_carry_fn(batch_size, device) + def init_carry( + self, + batch_size: int, + device: torch.device, + ) -> dict[str, torch.Tensor]: + # Preallocate an in-place cache only for the cached fast path. + return self._module_init_carry(batch_size, device, preallocate=self.use_kv_cache) + class DiffusionPolicyEstimator(PolicyMixin, Estimator): """Base class for diffusion policy estimators.""" diff --git a/src/gfn/utils/modules.py b/src/gfn/utils/modules.py index 14e5fc61..82c68d4b 100644 --- a/src/gfn/utils/modules.py +++ b/src/gfn/utils/modules.py @@ -1284,7 +1284,27 @@ def forward( hidden: torch.Tensor, key_carry: torch.Tensor, value_carry: torch.Tensor, + cache_len: Optional[int] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Self-attention over the K/V carry, in one of two carry modes. + + Args: + hidden: (batch, timesteps, embed_dim) block input. + key_carry: The running key cache. + value_carry: The running value cache. + cache_len: Selects the carry mode. When ``None`` (default) the block runs + in **grow-by-cat** mode: the new K/V are concatenated onto the carry + and freshly allocated buffers are returned. This is autograd-safe and + used by the full-sequence / no-cache paths. When an int, it is the + write cursor for **in-place** mode: ``key_carry`` / ``value_carry`` are + preallocated buffers, the new K/V are written at + ``[cache_len:cache_len+T]`` without reallocation, and the *same* buffer + objects are returned. In-place mutation is only version-counter-safe + under ``no_grad``. + + Returns: + (block output, updated key carry, updated value carry). + """ batch, timesteps, _ = hidden.size() normed_hidden = self.norm1(hidden) @@ -1297,30 +1317,45 @@ def forward( k = k.view(batch, timesteps, self.num_heads, self.head_dim).transpose(1, 2) v = v.view(batch, timesteps, self.num_heads, self.head_dim).transpose(1, 2) - carry_length = key_carry.size(2) - updated_key_carry = torch.cat((key_carry, k), dim=2) - updated_value_carry = torch.cat((value_carry, v), dim=2) - - attn_scores = torch.matmul(q, updated_key_carry.transpose(-2, -1)) / math.sqrt( + if cache_len is None: + # Grow-by-cat mode (grad-safe): O(L^2) reallocation across a rollout. + prefix_length = key_carry.size(2) + updated_key_carry = torch.cat((key_carry, k), dim=2) + updated_value_carry = torch.cat((value_carry, v), dim=2) + keys = updated_key_carry + values = updated_value_carry + else: + # In-place mode: write the new K/V into the preallocated buffer at the + # current cursor without reallocating (narrow+copy_ is TorchScript-safe). + prefix_length = cache_len + end = cache_len + timesteps + key_carry.narrow(2, cache_len, timesteps).copy_(k) + value_carry.narrow(2, cache_len, timesteps).copy_(v) + updated_key_carry = key_carry + updated_value_carry = value_carry + keys = key_carry.narrow(2, 0, end) + values = value_carry.narrow(2, 0, end) + + attn_scores = torch.matmul(q, keys.transpose(-2, -1)) / math.sqrt( float(self.head_dim) ) - if timesteps > 1 or carry_length > 0: - total_kv_length = carry_length + timesteps + if timesteps > 1 or prefix_length > 0: + total_kv_length = prefix_length + timesteps kv_positions = torch.arange( total_kv_length, device=hidden.device, dtype=torch.long ) query_positions = torch.arange( timesteps, device=hidden.device, dtype=torch.long ).unsqueeze(1) - causal_mask = kv_positions.unsqueeze(0) <= (query_positions + carry_length) + causal_mask = kv_positions.unsqueeze(0) <= (query_positions + prefix_length) attn_scores = attn_scores.masked_fill( ~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf") ) attn_weights = torch.softmax(attn_scores, dim=-1) attn_weights = self.attn_dropout(attn_weights) - attn_output = torch.matmul(attn_weights, updated_value_carry) + attn_output = torch.matmul(attn_weights, values) attn_output = attn_output.transpose(1, 2).reshape( batch, timesteps, self.embed_dim ) @@ -1406,20 +1441,57 @@ def init_carry( self, batch_size: int, device: torch.device, + max_len: Optional[int] = None, ) -> dict[str, torch.Tensor]: + """Initialize the recurrent carry. + + Args: + batch_size: Number of parallel rollouts. + device: Device to allocate the carry on. + max_len: If ``None`` (default), allocates length-0 K/V buffers that are + grown by concatenation each step (grad-safe; O(L^2) reallocation). + If an int, preallocates fixed ``[batch, n_heads, max_len, head_dim]`` + K/V buffers plus a scalar ``cache_len`` write cursor, so the block + writes new K/V in place (O(L), ``no_grad`` only). ``max_len`` must be + at least the maximum number of tokens decoded in a rollout (the fixed + env horizon + 1 for the BOS token) and no larger than + ``max_position_embeddings``. + + Returns: + The initialized carry dict. + """ weight = self.token_embedding.weight carry: dict[str, torch.Tensor] = { "position": torch.zeros(batch_size, dtype=torch.long, device=device), } - empty_key = weight.new_empty(batch_size, self.num_heads, 0, self.head_dim).to( - device - ) - empty_value = weight.new_empty(batch_size, self.num_heads, 0, self.head_dim).to( - device - ) - for key_name, value_name in zip(self.key_names, self.value_names): - carry[key_name] = empty_key.clone() - carry[value_name] = empty_value.clone() + if max_len is None: + empty_key = weight.new_empty( + batch_size, self.num_heads, 0, self.head_dim + ).to(device) + empty_value = weight.new_empty( + batch_size, self.num_heads, 0, self.head_dim + ).to(device) + for key_name, value_name in zip(self.key_names, self.value_names): + carry[key_name] = empty_key.clone() + carry[value_name] = empty_value.clone() + else: + if max_len <= 0: + raise ValueError("max_len must be positive for a preallocated carry.") + if max_len > self.max_position_embeddings: + raise ValueError( + "max_len exceeds max_position_embeddings; increase the model's " + "max_position_embeddings to preallocate a carry of this length." + ) + # Scalar cursor into the preallocated buffers (shared across rows, which + # advance in lockstep for the fixed-length fast path). + carry["cache_len"] = torch.zeros((), dtype=torch.long, device=device) + for key_name, value_name in zip(self.key_names, self.value_names): + carry[key_name] = weight.new_empty( + batch_size, self.num_heads, max_len, self.head_dim + ).to(device) + carry[value_name] = weight.new_empty( + batch_size, self.num_heads, max_len, self.head_dim + ).to(device) return carry @@ -1462,6 +1534,13 @@ def forward( updated_carry: dict[str, torch.Tensor] = {} + # An int cursor selects the in-place preallocated carry path in each block; + # None keeps the grow-by-cat path. The cursor is shared by all layers since + # they decode the same number of tokens. + cache_len: Optional[int] = None + if "cache_len" in carry: + cache_len = int(carry["cache_len"].item()) + for idx, layer in enumerate(self.layers): key_name = self.key_names[idx] value_name = self.value_names[idx] @@ -1471,22 +1550,28 @@ def forward( ) key_carry = carry[key_name] value_carry = carry[value_name] - if key_carry.size(0) != batch or key_carry.size(1) != self.num_heads: + if ( + key_carry.size(0) != batch + or value_carry.size(0) != batch + or key_carry.size(1) != self.num_heads + or value_carry.size(1) != self.num_heads + or key_carry.size(-1) != self.head_dim + or value_carry.size(-1) != self.head_dim + or key_carry.device != device + or value_carry.device != device + ): raise ValueError( - "Key carry shape is incompatible with the provided tokens." + "Key/value carry is incompatible with the provided tokens " + "(batch, num_heads, head_dim, or device mismatch)." ) - if value_carry.size(0) != batch or value_carry.size(1) != self.num_heads: + # Actionable guard kept separate: overflow points the user at max_len. + if cache_len is not None and cache_len + timesteps > key_carry.size(2): raise ValueError( - "Value carry shape is incompatible with the provided tokens." + "Preallocated carry is too short for the decoded sequence; " + "increase max_len passed to init_carry." ) - if (key_carry.size(-1) != self.head_dim) or ( - value_carry.size(-1) != self.head_dim - ): - raise ValueError("Key/value carry head dimension mismatch detected.") - if key_carry.device != device or value_carry.device != device: - raise ValueError("Key/value carry tensors must share the input device.") hidden, updated_key_carry, updated_value_carry = layer( - hidden, key_carry, value_carry + hidden, key_carry, value_carry, cache_len ) updated_carry[key_name] = updated_key_carry updated_carry[value_name] = updated_value_carry @@ -1495,6 +1580,8 @@ def forward( logits = self.output_projection(hidden) updated_carry["position"] = positions + timesteps + if "cache_len" in carry: + updated_carry["cache_len"] = carry["cache_len"] + timesteps return logits, updated_carry @property diff --git a/testing/test_kv_cache.py b/testing/test_kv_cache.py new file mode 100644 index 00000000..21ffe9cb --- /dev/null +++ b/testing/test_kv_cache.py @@ -0,0 +1,238 @@ +"""Correctness gates for the fixed-length KV-cache sampling optimization. + +Covers two layers: + +- The transformer module's in-place preallocated carry (PR1a): incremental decode + equals a single full-sequence forward, the buffers are preallocated and mutated in + place, and the decode position is load-bearing. +- The recurrent estimator's ``use_kv_cache`` toggle (PR1b): the cached fast path and + the corrected full-prefix baseline produce identical per-step logits, both equal to + a teacher-forced forward, and the fixed-length guard fires on ragged batches. + +These are the gates the spec calls "load-bearing": they turn red on the previous +(grow-both) recurrent path and green on the corrected implementation. +""" + +from typing import Literal + +import pytest +import torch + +from gfn.estimators import RecurrentDiscretePolicyEstimator +from gfn.gym.bitSequence import BitSequence +from gfn.samplers import Sampler +from gfn.utils.modules import TransformerDiscreteSequenceModel + +ATOL = 1e-5 + + +def _make_model( + vocab_size: int = 7, + embedding_dim: int = 16, + num_heads: int = 2, + num_layers: int = 2, + max_position_embeddings: int = 32, + positional_embedding: Literal["learned", "sinusoidal"] = "learned", +) -> TransformerDiscreteSequenceModel: + model = TransformerDiscreteSequenceModel( + vocab_size=vocab_size, + embedding_dim=embedding_dim, + num_heads=num_heads, + ff_hidden_dim=2 * embedding_dim, + num_layers=num_layers, + max_position_embeddings=max_position_embeddings, + dropout=0.0, + positional_embedding=positional_embedding, + ) + model.eval() + return model + + +# -------------------------------------------------------------------------------- +# PR1a: module-level in-place preallocated carry +# -------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "batch, num_heads, num_layers, seq_len", + [(1, 1, 1, 4), (3, 2, 2, 6), (2, 4, 3, 9)], +) +@pytest.mark.parametrize("positional_embedding", ["learned", "sinusoidal"]) +def test_inplace_cache_matches_full_forward( + batch: int, + num_heads: int, + num_layers: int, + seq_len: int, + positional_embedding: Literal["learned", "sinusoidal"], +) -> None: + """Incremental in-place decode == a single full-sequence forward (dense gate).""" + vocab_size = 7 + model = _make_model( + vocab_size=vocab_size, + embedding_dim=4 * num_heads, + num_heads=num_heads, + num_layers=num_layers, + max_position_embeddings=seq_len + 2, + positional_embedding=positional_embedding, + ) + device = torch.device("cpu") + tokens = torch.randint(0, vocab_size, (batch, seq_len)) + + with torch.no_grad(): + carry = model.init_carry(batch, device) + logits_full, _ = model(tokens, carry) + + carry = model.init_carry(batch, device, max_len=seq_len) + step_logits = [] + for t in range(seq_len): + lo, carry = model(tokens[:, t : t + 1], carry) + step_logits.append(lo) + logits_inplace = torch.cat(step_logits, dim=1) + + assert torch.allclose(logits_full, logits_inplace, atol=ATOL) + + +def test_inplace_cache_is_preallocated_and_mutated_in_place() -> None: + """Buffers are horizon-length from step 0, the same object across steps, and the + cursor advances one slot per decoded token.""" + vocab_size, batch, seq_len = 7, 3, 6 + model = _make_model(vocab_size=vocab_size, max_position_embeddings=seq_len + 2) + device = torch.device("cpu") + tokens = torch.randint(0, vocab_size, (batch, seq_len)) + + with torch.no_grad(): + carry = model.init_carry(batch, device, max_len=seq_len) + key0_obj = carry["key_0"] + # Preallocated to the full horizon from the very first step. + assert carry["key_0"].size(2) == seq_len + assert carry["value_0"].size(2) == seq_len + assert int(carry["cache_len"].item()) == 0 + + for t in range(seq_len): + _, carry = model(tokens[:, t : t + 1], carry) + assert carry["key_0"] is key0_obj # same tensor object, mutated in place + assert carry["key_0"].size(2) == seq_len # never reallocated / grown + assert int(carry["cache_len"].item()) == t + 1 + + +def test_decode_position_is_load_bearing() -> None: + """Perturbing the decode position by +1 must break equivalence (guards abs-pos).""" + vocab_size, batch, seq_len = 7, 2, 5 + model = _make_model(vocab_size=vocab_size, max_position_embeddings=seq_len + 2) + device = torch.device("cpu") + tokens = torch.randint(0, vocab_size, (batch, seq_len)) + + with torch.no_grad(): + carry_ref = model.init_carry(batch, device, max_len=seq_len) + logits_ref, _ = model(tokens[:, 0:1], carry_ref) + + carry_bad = model.init_carry(batch, device, max_len=seq_len) + carry_bad["position"] = carry_bad["position"] + 1 # off-by-one + logits_bad, _ = model(tokens[:, 0:1], carry_bad) + + assert not torch.allclose(logits_ref, logits_bad, atol=ATOL) + + +# -------------------------------------------------------------------------------- +# PR1b: estimator-level use_kv_cache toggle +# -------------------------------------------------------------------------------- + + +def _env_and_model(word_size: int = 2, seq_size: int = 12): + env = BitSequence( + word_size=word_size, + seq_size=seq_size, + n_modes=2, + device_str="cpu", + seed=0, + debug=False, + ) + model = _make_model( + vocab_size=env.n_actions, + num_heads=2, + num_layers=2, + max_position_embeddings=env.words_per_seq + 4, + ) + return env, model + + +def _per_step_logits(env, est, words: torch.Tensor) -> torch.Tensor: + """Drive the estimator over reconstructed per-step states (0..L words committed).""" + batch = words.shape[0] + device = words.device + carry = est.init_carry(batch, device) + out = [] + with torch.no_grad(): + for t in range(env.words_per_seq + 1): + tensor = words.clone() + tensor[:, t:] = -1 + logits, carry = est(env.States(tensor), carry) + out.append(logits) + return torch.stack(out, dim=0) # (L+1, B, n_actions) + + +def test_estimator_cache_equivalence() -> None: + """use_kv_cache OFF == ON == a single teacher-forced forward, at the logit level.""" + env, model = _env_and_model() + batch, L = 4, env.words_per_seq + words = torch.randint(0, env.n_actions - 1, (batch, L)) + + est_off = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=False + ) + est_on = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=True + ) + logits_off = _per_step_logits(env, est_off, words) + logits_on = _per_step_logits(env, est_on, words) + + with torch.no_grad(): + bos = torch.full((batch, 1), env.n_actions, dtype=torch.long) + seq = torch.cat([bos, words], dim=1) + tf_logits, _ = model(seq, model.init_carry(batch, torch.device("cpu"))) + tf_logits = tf_logits.transpose(0, 1) # (L+1, B, n_actions) + + assert torch.allclose(logits_off, tf_logits, atol=ATOL) + assert torch.allclose(logits_on, tf_logits, atol=ATOL) + assert torch.allclose(logits_off, logits_on, atol=ATOL) + + +def test_estimator_fixed_length_guard() -> None: + """The cached path rejects ragged (variable-length) active batches (debug on).""" + env, model = _env_and_model() + # The equal-length check is gated behind debug (it forces a per-step device sync). + est = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=True, debug=True + ) + carry = est.init_carry(3, torch.device("cpu")) + ragged = torch.randint(0, env.n_actions - 1, (3, env.words_per_seq)) + ragged[0, 1:] = -1 # row 0 has a different committed length + with pytest.raises(ValueError, match="equal-length"): + est(env.States(ragged), carry) + + +def test_cache_max_len_validated_eagerly() -> None: + """A cache_max_len beyond the model's position range fails at construction.""" + env, model = _env_and_model() + with pytest.raises(ValueError, match="max_position_embeddings"): + RecurrentDiscretePolicyEstimator( + module=model, + n_actions=env.n_actions, + use_kv_cache=True, + cache_max_len=model.max_position_embeddings + 10, + ) + + +def test_sampling_runs_in_both_modes() -> None: + """Both toggle settings produce valid fixed-length trajectories end-to-end.""" + env, model = _env_and_model() + for use_cache in (False, True): + est = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=use_cache + ) + with torch.no_grad(): + trajs = Sampler(estimator=est).sample_trajectories( + env, n=5, save_logprobs=True, save_estimator_outputs=False + ) + # Fixed-length: every trajectory has the same number of steps. + assert trajs.max_length == env.words_per_seq + 1 diff --git a/tutorials/examples/benchmark_kv_cache.py b/tutorials/examples/benchmark_kv_cache.py new file mode 100644 index 00000000..ff5897cb --- /dev/null +++ b/tutorials/examples/benchmark_kv_cache.py @@ -0,0 +1,344 @@ +#!/usr/bin/env python +r"""Instructional benchmark: KV-caching for autoregressive GFlowNet sampling. + +Run me: + + python tutorials/examples/benchmark_kv_cache.py + python tutorials/examples/benchmark_kv_cache.py --seq_sizes 16 32 64 128 --plot out.png + +I sample the *same* transformer policy on ``BitSequence`` two ways and time both, so the +wall-clock difference is exactly the KV-cache. No GPU needed -- the win is a FLOP +reduction and shows up on CPU (and grows with sequence length). + +================================================================================ +WHY KV-CACHING HELPS +================================================================================ +A GFlowNet policy over sequences is *autoregressive*: to choose token ``t`` it runs a +transformer over the tokens chosen so far and reads the logits at the last position. + +Inside each attention layer, the query at position ``t`` attends to the keys/values +(K/V) of every earlier position ``0..t``. Crucially, the K/V of positions ``0..t-1`` +are a function of tokens that are *already fixed* -- they do not change as the sequence +grows. Yet naive sampling recomputes the whole prefix from scratch at every step: + + step t=3, NO cache: embed+attend [w0 w1 w2] <- recompute K/V for w0,w1 again + step t=3, cache: embed+attend [ w2] <- compute K/V for w2 only, + reuse cached K/V for w0,w1 + +So a length-``L`` rollout without a cache re-embeds an ever-growing prefix each step: +roughly ``sum_t O(t^2) = O(L^3)`` attention FLOPs. A cache stores each position's K/V +the first time it is computed and reuses it, so step ``t`` only computes the new +token's Q/K/V and attends it against the ``t`` cached keys: ``sum_t O(t) = O(L^2)``. +That is a ~``L``x speedup that *widens with sequence length* -- the effect this script +makes visible. + +================================================================================ +HOW torchgfn IMPLEMENTS IT +================================================================================ +The seam is the *recurrent carry* threaded through sampling by ``Sampler``: + + RecurrentDiscretePolicyEstimator.forward(states, carry) -> (logits, carry) + +The carry is a ``dict[str, Tensor]`` holding, per transformer layer, the running K/V +cache, plus a scalar write cursor and the absolute position. Two forward paths (both +build the same ``[BOS, w0, ..., w_{t-1}]`` token stream and are numerically identical +to one full-sequence forward): + + * ``use_kv_cache=False`` (default): re-embed the full committed prefix with a fresh + carry each step. Correct and grad-bearing, but O(L^3) -- the honest baseline. + + * ``use_kv_cache=True``: feed only the newest token; ``TransformerDiscreteSequenceModel`` + writes its K/V *in place* into a preallocated ``[B, n_heads, max_len, head_dim]`` + buffer at the cursor (no ``torch.cat``, no realloc) and attends against the filled + prefix. This is the O(L^2) fast path. + +The autograd catch: writing K/V in place mutates a buffer whose earlier slice was saved +for a previous step's backward, which invalidates the autograd graph. So the cached path +runs under ``torch.no_grad()`` -- which is fine, because *sampling only needs sampled +actions, not gradients*. This makes ``use_kv_cache=True`` ideal for inference/rollout; +to *train* with it, recompute the loss log-probs with a grad-bearing forward +(``recalculate_all_logprobs=True``). A companion "teacher-forced" parallel recompute +makes that efficient and turns this sampling win into an end-to-end training win. + +================================================================================ +WHAT THIS SCRIPT MEASURES +================================================================================ +For each sequence length it (1) asserts the two configs produce identical per-step +logits (``|Δ logits| ~ 1e-7`` -- proof it is the same model), then times pure rollout +under ``no_grad`` for each -- isolating the cache (O(L^3) vs O(L^2)). +""" + +from __future__ import annotations + +import argparse +import statistics +import time +from typing import Callable + +import torch + +from gfn.estimators import RecurrentDiscretePolicyEstimator +from gfn.gflownet import TBGFlowNet +from gfn.gym.bitSequence import BitSequence +from gfn.utils.common import set_seed +from gfn.utils.modules import TransformerDiscreteSequenceModel + + +def _sync(device: torch.device) -> None: + """Block until queued device work finishes so timings are accurate. + + GPU/MPS kernels launch asynchronously, so ``perf_counter`` would otherwise time + only the Python dispatch, not the compute. A no-op on CPU. + """ + if device.type == "cuda": + torch.cuda.synchronize() + elif device.type == "mps": + torch.mps.synchronize() + + +def build(args, seq_size: int, use_kv_cache: bool, device: torch.device): + """Construct ``(env, gflownet)`` for one benchmark configuration. + + This is the only place the KV-cache API appears -- everything else (sampler, loss, + optimizer) is the ordinary torchgfn stack, unchanged. + + Args: + args: Parsed CLI arguments (model/env hyperparameters). + seq_size: BitSequence length; with ``word_size=1`` this equals the number of + autoregressive decode steps per trajectory (the ``L`` in the docstring). + use_kv_cache: Selects the sampling path (see module docstring). + device: Torch device. + + Returns: + The environment and a ``TBGFlowNet`` wrapping a transformer forward policy. + """ + env = BitSequence( + word_size=args.word_size, + seq_size=seq_size, + n_modes=args.n_modes, + device_str=str(device), + seed=args.seed, + debug=False, + ) + # The KV-cache is preallocated to the model's position range, so it must cover the + # whole rollout: one slot per word plus the leading BOS token (+ a little headroom). + max_positions = env.words_per_seq + 2 + model = TransformerDiscreteSequenceModel( + vocab_size=env.n_actions, + embedding_dim=args.embedding_dim, + num_heads=args.num_heads, + ff_hidden_dim=args.ff_hidden_dim, + num_layers=args.num_layers, + max_position_embeddings=max_positions, + dropout=0.0, + ).to(device) + est = RecurrentDiscretePolicyEstimator( + module=model, + n_actions=env.n_actions, + is_backward=False, + # The single switch that turns on the in-place, feed-newest-token fast path. + use_kv_cache=use_kv_cache, + ) + # Tree-structured DAG (each string has a unique parent) => backward policy is + # constant, hence pb=None + constant_pb=True. Only the forward policy is timed. + gflownet = TBGFlowNet(pf=est, pb=None, init_logZ=0.0, constant_pb=True).to(device) + return env, gflownet + + +def _median_ms(fn: Callable[[], None], device: torch.device, warmup: int, iters: int): + """Return the median (min, max) wall-clock of ``fn`` in milliseconds. + + Warm-up iterations are discarded so one-time costs (lazy allocations, the first + preallocation of the cache buffers) do not pollute the measurement. + """ + for _ in range(warmup): + fn() + _sync(device) + samples = [] + for _ in range(iters): + _sync(device) + t0 = time.perf_counter() + fn() + _sync(device) + samples.append((time.perf_counter() - t0) * 1e3) + lo, hi = min(samples), max(samples) + return statistics.median(samples), lo, hi + + +def time_sampling(gflownet, env, batch: int, device, warmup, iters): + """Time pure rollout under ``no_grad`` -- the operation the cache directly speeds up. + + Both configs run under ``no_grad`` here, so this isolates the cache: identical model, + identical outputs, the only difference is O(L^3) re-forward vs O(L^2) cached decode. + """ + + def _step(): + with torch.no_grad(): + gflownet.sample_trajectories( + env, n=batch, save_logprobs=False, save_estimator_outputs=False + ) + + return _median_ms(_step, device, warmup, iters) + + +def check_equivalence(args, seq_size: int, device: torch.device) -> float: + """Prove the cache changes speed, not results, at the per-step logit level. + + We reconstruct the states at every decode step of a fixed batch and run the + estimator over them both ways (``use_kv_cache`` False and True). Because the cached + incremental decode equals the full-prefix forward, the logits must agree to + floating-point tolerance; a large gap would mean the cache computes the wrong thing. + Returns the max absolute logit difference. + """ + set_seed(args.seed) + env, gflownet = build(args, seq_size, use_kv_cache=False, device=device) + # Both estimators wrap the SAME module weights, so the only difference is the path. + model = gflownet.pf.module + est_off = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=False + ) + est_on = RecurrentDiscretePolicyEstimator( + module=model, n_actions=env.n_actions, use_kv_cache=True + ) + + batch, L = args.batch_size, env.words_per_seq + words = torch.randint(0, env.n_actions - 1, (batch, L), device=device) + + def per_step_logits(est) -> torch.Tensor: + carry = est.init_carry(batch, device) + out = [] + with torch.no_grad(): + for t in range(L + 1): # steps 0..L (L+1 decode steps) + tensor = words.clone() + tensor[:, t:] = -1 # only first t words committed + logits, carry = est(env.States(tensor), carry) + out.append(logits) + return torch.stack(out, dim=0) + + return (per_step_logits(est_off) - per_step_logits(est_on)).abs().max().item() + + +def main(args) -> None: + device = torch.device(args.device) + print( + f"device={device} batch={args.batch_size} " + f"model=(dim={args.embedding_dim}, heads={args.num_heads}, " + f"layers={args.num_layers}) word_size={args.word_size}\n" + ) + + # Columns: sampling base/cache + speedup, and the equivalence residual proving both + # columns describe the same model. + header = ( + f"{'seq_len':>7} | {'sample base':>12} {'sample cache':>13} " + f"{'speedup':>8} | {'|Δlogits|':>10}" + ) + print(header) + print("-" * len(header)) + + rows = [] + for seq_size in args.seq_sizes: + set_seed(args.seed) + + # Correctness gate first: never report a speedup for a model that changed. + max_delta = check_equivalence(args, seq_size, device) + assert ( + max_delta < 1e-4 + ), f"baseline vs cached logits differ by {max_delta} at seq_size={seq_size}" + + # Fresh models per config so neither warms the other's caches/allocator. + env_b, gfn_b = build(args, seq_size, use_kv_cache=False, device=device) + env_c, gfn_c = build(args, seq_size, use_kv_cache=True, device=device) + + s_base = time_sampling( + gfn_b, env_b, args.batch_size, device, args.warmup, args.iters + ) + s_cache = time_sampling( + gfn_c, env_c, args.batch_size, device, args.warmup, args.iters + ) + speedup = s_base[0] / s_cache[0] + rows.append((seq_size, s_base[0], s_cache[0], speedup)) + print( + f"{seq_size:>7} | {s_base[0]:>10.2f}ms {s_cache[0]:>11.2f}ms " + f"{speedup:>7.1f}x | {max_delta:>10.1e}" + ) + + print( + "\nBoth configs are numerically identical (see |Δlogits|); the speedup is the " + "KV-cache.\nIt isolates the sampling win (O(L^3)->O(L^2)) and widens with " + "seq_len -- that is the point.\nTo turn this into an end-to-end training " + "speedup, pair it with the teacher-forced loss recompute." + ) + + if args.plot: + _save_plot(rows, args.plot) + + +def _save_plot(rows, path: str) -> None: + """Optionally save a sampling-time / speedup figure (requires matplotlib).""" + try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ImportError: + print("matplotlib not available; skipping --plot") + return + + seq = [r[0] for r in rows] + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4)) + ax1.plot(seq, [r[1] for r in rows], "o-", label="baseline (no cache)") + ax1.plot(seq, [r[2] for r in rows], "s-", label="KV-cache") + ax1.set( + xlabel="sequence length (steps)", + ylabel="ms / rollout", + title="Pure sampling time", + ) + ax1.legend() + ax1.grid(alpha=0.3) + ax2.plot(seq, [r[3] for r in rows], "o-") + ax2.axhline(1.0, color="gray", ls="--", lw=0.8) + ax2.set( + xlabel="sequence length (steps)", + ylabel="speedup (x)", + title="KV-cache sampling speedup", + ) + ax2.grid(alpha=0.3) + fig.tight_layout() + fig.savefig(path, dpi=120) + print(f"saved plot to {path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Benchmark KV-cached vs baseline autoregressive GFlowNet sampling.", + ) + parser.add_argument("--device", type=str, default="cpu", help="cpu | mps | cuda") + parser.add_argument( + "--seq_sizes", + type=int, + nargs="+", + default=[8, 16, 32, 64], + help="BitSequence lengths to sweep (== rollout lengths with word_size=1)", + ) + parser.add_argument( + "--word_size", + type=int, + default=1, + help="Bits per action; 1 => one decode step per sequence bit", + ) + parser.add_argument("--n_modes", type=int, default=4) + parser.add_argument("--batch_size", type=int, default=32) + parser.add_argument("--embedding_dim", type=int, default=64) + parser.add_argument("--num_heads", type=int, default=4) + parser.add_argument("--ff_hidden_dim", type=int, default=128) + parser.add_argument("--num_layers", type=int, default=2) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iters", type=int, default=5) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--plot", + type=str, + default=None, + help="Optional path to save a speedup figure (PNG)", + ) + main(parser.parse_args())