From 09d92d09c189d4f1675b246cff00cab34e1ef822 Mon Sep 17 00:00:00 2001 From: Harry Robertson <69441195+Harry25R@users.noreply.github.com> Date: Sat, 22 Aug 2026 02:54:01 +1000 Subject: [PATCH 1/2] perf: cache conditioning K/V projections during flow sampling The conditioning tensor is constant across ODE solver steps, so its key/value projections in every cross-attention block can be computed once and reused. Adds run_fast_flow / OptimizedFlow and a fast=True flag on FlowPipeline; attention blocks accept an optional precomputed kv tuple. Output is identical to the default path. --- src/phoenix/helpers/fast_sampler.py | 107 ++++++++++++++++++++++++++++ src/phoenix/helpers/inference.py | 10 ++- src/phoenix/models/flow_llama3.py | 30 +++++--- src/phoenix/models/flow_simple.py | 30 +++++--- 4 files changed, 156 insertions(+), 21 deletions(-) create mode 100644 src/phoenix/helpers/fast_sampler.py diff --git a/src/phoenix/helpers/fast_sampler.py b/src/phoenix/helpers/fast_sampler.py new file mode 100644 index 0000000..e08fd99 --- /dev/null +++ b/src/phoenix/helpers/fast_sampler.py @@ -0,0 +1,107 @@ +""" +Fast flow inference that caches the conditioning K/V projections. + +The conditioning tensor `c` is constant across ODE steps, so its key/value +projections in every cross-attention block can be computed once and reused, +instead of being recomputed at every solver step. +""" + +import torch +from zuko.utils import odeint + +# ------------------------------------------------------------------------------- + + +class OptimizedFlow: + """ + Wraps a flow transformer model, caching the conditioning K/V projections. + + Parameters + ---------- + flow_model + Flow transformer model (``flow_llama3`` or ``flow_simple``) to wrap. + """ + + def __init__(self, flow_model): + self.m = flow_model + self._kv = None + + @torch.no_grad() + def prep(self, c): + """ + Precompute the K/V projections of the conditioning tensor `c`. + + `c` is constant in the ODE, so its K/V projections can be cached once. + """ + m = self.m + self._kv = None # drop any previous cache + + if c.dim() == 4: # if c is an image, extract image features + c = m.vision_forward(c) + + # mirror the model's conditioning path dtype (fp32) + pc = torch.arange(c.size(1), device=c.device) + c = m.c_projection(m.c_norm(c)) # c is the feature vector + c = c + m.pc_embedding(pc) + for layer in m.layers: # conditioning blocks + c = layer(c) + + c_bsz, c_sql, _ = c.shape + kv = [] + + # mirror the attention blocks' dtype (bf16) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + for block in m.blocks: + attn = block.xattn + xk = attn.weight_k(c) + xv = attn.weight_v(c) + xk = xk.view(c_bsz, c_sql, attn.n_heads, attn.d_head) + xv = xv.view(c_bsz, c_sql, attn.n_heads, attn.d_head) + xk = attn.norm_k(xk) + kv.append((xk, xv)) + + self._kv = kv # c isn't needed anymore + + @torch.no_grad() + def velocity(self, t_scalar, x, y=None, device=None): + """Evaluate the velocity field at time `t_scalar`, reusing the cached K/V.""" + m = self.m + + t = torch.full((x.shape[0],), t_scalar, device=device or x.device) + + px = torch.arange(x.size(1), device=x.device) + x = m.x_projection(x) + x = x + m.px_embedding(px) + + t = m.t_embedding(t) + if m.cfg.n_classes > 0 and y is not None: + t = t + m.y_embedding(y, m.training) + + if self._kv is None: + raise RuntimeError("call prep(c) before velocity()") + for block, kv in zip(m.blocks, self._kv, strict=True): + x = block(x, t, None, kv) + + return m.head(x, t) + + +# ------------------------------------------------------------------------------- + + +@torch.no_grad() +def run_fast_flow(flow_model, x_0, t_0, t_1, c, y, atol, rtol, device="cpu"): + """ + Integrate the velocity field, caching the conditioning K/V projections. + + Drop-in replacement for ``run_flow`` that computes the K/V projections of + `c` once instead of recomputing them at every ODE solver step. + """ + flow_model.eval() + phi = flow_model.parameters() + opt = OptimizedFlow(flow_model) + opt.prep(c) + + def f(t: float, x: torch.Tensor): + return opt.velocity(t, x, y=y, device=device) + + return odeint(f, x_0, t_0, t_1, phi=phi, atol=atol, rtol=rtol) diff --git a/src/phoenix/helpers/inference.py b/src/phoenix/helpers/inference.py index a251602..2cd1253 100644 --- a/src/phoenix/helpers/inference.py +++ b/src/phoenix/helpers/inference.py @@ -11,6 +11,8 @@ from tqdm import tqdm from zuko.utils import odeint +from phoenix.helpers.fast_sampler import run_fast_flow + # ------------------------------------------------------------------------------- @@ -77,6 +79,9 @@ class FlowPipeline: Absolute tolerance passed to the ODE solver. rtol Relative tolerance passed to the ODE solver. + fast + Whether to sample with ``run_fast_flow``, which caches the conditioning + K/V projections across ODE steps, instead of ``run_flow``. """ def __init__( @@ -87,12 +92,14 @@ def __init__( t_1: float = 1.0, atol: float = 1e-1, rtol: float = 1e-1, + fast: bool = False, ): self.model = model self.t_0 = t_0 self.t_1 = t_1 self.atol = atol self.rtol = rtol + self.fast = fast if stats is None: raise ValueError("FlowPipeline requires `stats` with 'mean' and 'std' entries") self.mean, self.std = stats["mean"], stats["std"] @@ -121,6 +128,7 @@ def __call__(self, gene_list: list, dataloader: DataLoader): """ self.model.eval() device = self.device + sampler = run_fast_flow if self.fast else run_flow pred_list, coords_list = [], [] for batch in tqdm(dataloader, desc="Flow sampling"): @@ -131,7 +139,7 @@ def __call__(self, gene_list: list, dataloader: DataLoader): feats = self.model.vision_forward(image) # type: ignore[operator] noise = torch.randn(image.size(0), len(gene_list), 1, device=device) - gex_pred = run_flow( + gex_pred = sampler( flow_model=self.model, x_0=noise.float(), t_0=self.t_0, diff --git a/src/phoenix/models/flow_llama3.py b/src/phoenix/models/flow_llama3.py index a54e461..47bf8ae 100644 --- a/src/phoenix/models/flow_llama3.py +++ b/src/phoenix/models/flow_llama3.py @@ -337,7 +337,7 @@ def __init__( self.dropout = nn.Dropout(proj_drop) @torch.autocast(device_type="cuda", dtype=torch.bfloat16) - def forward(self, x: Tensor, c: Tensor | None = None): + def forward(self, x: Tensor, c: Tensor | None = None, kv: tuple | None = None): """ Apply attention, using `x` as query and `c` (or `x`, if `c` is `None`) as key/value. @@ -348,6 +348,9 @@ def forward(self, x: Tensor, c: Tensor | None = None): c Optional key/value input, shape ``(batch, tokens_c, d_cross)``; falls back to self-attention on `x` when `None`. + kv + Optional precomputed ``(key, value)`` projections, each of shape + ``(batch, tokens_c, n_heads, d_head)``; when given, `c` is ignored. Returns ------- @@ -360,15 +363,19 @@ def forward(self, x: Tensor, c: Tensor | None = None): c_bsz, c_sql, _ = x.shape xq = self.weight_q(x) # inputs serve as query - xk = self.weight_k(c if c is not None else x) - xv = self.weight_v(c if c is not None else x) - xq = xq.view(x_bsz, x_sql, self.n_heads, self.d_head) - xk = xk.view(c_bsz, c_sql, self.n_heads, self.d_head) - xv = xv.view(c_bsz, c_sql, self.n_heads, self.d_head) - xq = self.norm_q(xq) - xk = self.norm_k(xk) + + if kv is not None: + xk, xv = kv # precomputed by the caller + if xk.shape[0] != x_bsz: + raise ValueError(f"cached k/v batch {xk.shape[0]} != query batch {x_bsz}") + else: + xk = self.weight_k(c if c is not None else x) + xv = self.weight_v(c if c is not None else x) + xk = xk.view(c_bsz, c_sql, self.n_heads, self.d_head) + xv = xv.view(c_bsz, c_sql, self.n_heads, self.d_head) + xk = self.norm_k(xk) xo = flash_attn_func( xq, @@ -524,7 +531,7 @@ def __init__(self, cfg): ), ) - def forward(self, x: Tensor, t: Tensor, c: Tensor | None = None): + def forward(self, x: Tensor, t: Tensor, c: Tensor | None = None, kv: tuple | None = None): """ Apply one flow-matching transformer block. @@ -537,6 +544,9 @@ def forward(self, x: Tensor, t: Tensor, c: Tensor | None = None): ``(batch, d_model)``. c Conditioning tokens for cross-attention, shape ``(batch, tokens_c, d_cross)``. + kv + Optional precomputed ``(key, value)`` projections for the cross-attention; + when given, `c` is ignored. Returns ------- @@ -549,7 +559,7 @@ def forward(self, x: Tensor, t: Tensor, c: Tensor | None = None): r = self.zattn(modulate(self.norm_1(x), shift_zattn, scale_zattn)) x = x + gate_zattn.unsqueeze(1) * r - r = self.xattn(self.norm_2(x), c) + r = self.xattn(self.norm_2(x), c, kv) x = x + r r = self.mlp(modulate(self.norm_3(x), shift_mlp, scale_mlp)) diff --git a/src/phoenix/models/flow_simple.py b/src/phoenix/models/flow_simple.py index 9f116c7..f33104b 100644 --- a/src/phoenix/models/flow_simple.py +++ b/src/phoenix/models/flow_simple.py @@ -419,7 +419,7 @@ def __init__( self.dropout = nn.Dropout(proj_drop) @torch.autocast(device_type="cuda", dtype=torch.bfloat16) - def forward(self, x: Tensor, c: Tensor | None = None): + def forward(self, x: Tensor, c: Tensor | None = None, kv: tuple | None = None): """ Apply attention, using `x` as query and `c` (or `x`, if `c` is `None`) as key/value. @@ -430,6 +430,9 @@ def forward(self, x: Tensor, c: Tensor | None = None): c Optional key/value input, shape ``(batch, tokens_c, d_cross)``; falls back to self-attention on `x` when `None`. + kv + Optional precomputed ``(key, value)`` projections, each of shape + ``(batch, tokens_c, n_heads, d_head)``; when given, `c` is ignored. Returns ------- @@ -442,15 +445,19 @@ def forward(self, x: Tensor, c: Tensor | None = None): c_bsz, c_sql, _ = x.shape xq = self.weight_q(x) # inputs serve as query - xk = self.weight_k(c if c is not None else x) - xv = self.weight_v(c if c is not None else x) - xq = xq.view(x_bsz, x_sql, self.n_heads, self.d_head) - xk = xk.view(c_bsz, c_sql, self.n_heads, self.d_head) - xv = xv.view(c_bsz, c_sql, self.n_heads, self.d_head) - xq = self.norm_q(xq) - xk = self.norm_k(xk) + + if kv is not None: + xk, xv = kv # precomputed by the caller + if xk.shape[0] != x_bsz: + raise ValueError(f"cached k/v batch {xk.shape[0]} != query batch {x_bsz}") + else: + xk = self.weight_k(c if c is not None else x) + xv = self.weight_v(c if c is not None else x) + xk = xk.view(c_bsz, c_sql, self.n_heads, self.d_head) + xv = xv.view(c_bsz, c_sql, self.n_heads, self.d_head) + xk = self.norm_k(xk) xq = xq.transpose(1, 2) xk = xk.transpose(1, 2) @@ -608,7 +615,7 @@ def __init__(self, cfg): ), ) - def forward(self, x: Tensor, t: Tensor, c: Tensor | None = None): + def forward(self, x: Tensor, t: Tensor, c: Tensor | None = None, kv: tuple | None = None): """ Apply one flow-matching transformer block. @@ -621,6 +628,9 @@ def forward(self, x: Tensor, t: Tensor, c: Tensor | None = None): ``(batch, d_model)``. c Conditioning tokens for cross-attention, shape ``(batch, tokens_c, d_cross)``. + kv + Optional precomputed ``(key, value)`` projections for the cross-attention; + when given, `c` is ignored. Returns ------- @@ -633,7 +643,7 @@ def forward(self, x: Tensor, t: Tensor, c: Tensor | None = None): r = self.zattn(modulate(self.norm_1(x), shift_zattn, scale_zattn)) x = x + gate_zattn.unsqueeze(1) * r - r = self.xattn(self.norm_2(x), c) + r = self.xattn(self.norm_2(x), c, kv) x = x + r r = self.mlp(modulate(self.norm_3(x), shift_mlp, scale_mlp)) From 63c9e601a1baa2fbd3576f258dd4f15ff35f6a4f Mon Sep 17 00:00:00 2001 From: Harry Robertson <69441195+Harry25R@users.noreply.github.com> Date: Sat, 22 Aug 2026 03:20:05 +1000 Subject: [PATCH 2/2] test: cover the fast K/V-caching sampler on CPU --- tests/test_inference.py | 51 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tests/test_inference.py b/tests/test_inference.py index d0506f6..d50f416 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -87,3 +87,54 @@ def test_flow_pipeline_runs_on_cpu(pipeline_model): assert gex_pred.shape == (n_samples, n_genes) assert np.isfinite(gex_pred).all() assert sum(len(c) for c in coords_list) == n_samples + + +def test_run_fast_flow_matches_run_flow(tiny_model): + """The K/V-caching sampler must produce the same output as run_flow.""" + from phoenix.helpers.fast_sampler import run_fast_flow + + torch.manual_seed(0) + x_0 = torch.randn(1, 6, 4) + c = torch.randn(1, 5, 16) + + slow = run_flow(tiny_model, x_0, t_0=0.0, t_1=1.0, c=c, y=None, atol=1e-1, rtol=1e-1) + fast = run_fast_flow(tiny_model, x_0, t_0=0.0, t_1=1.0, c=c, y=None, atol=1e-1, rtol=1e-1) + + torch.testing.assert_close(fast, slow) + + +def test_optimized_flow_requires_prep(tiny_model): + from phoenix.helpers.fast_sampler import OptimizedFlow + + opt = OptimizedFlow(tiny_model) + with pytest.raises(RuntimeError, match="prep"): + opt.velocity(0.5, torch.randn(1, 6, 4)) + + +def test_fast_attention_rejects_mismatched_kv_batch(tiny_model): + from phoenix.helpers.fast_sampler import OptimizedFlow + + opt = OptimizedFlow(tiny_model) + opt.prep(torch.randn(1, 5, 16)) + with pytest.raises(ValueError, match="batch"): + opt.velocity(0.5, torch.randn(2, 6, 4)) + + +def test_flow_pipeline_fast_runs_on_cpu(pipeline_model): + """FlowPipeline(fast=True) must run end-to-end without a GPU.""" + from torch.utils.data import DataLoader, TensorDataset + + torch.manual_seed(0) + n_samples, n_genes = 4, 3 + feats = torch.randn(n_samples, 5, 16) + coords = torch.zeros(n_samples, 2) + loader = DataLoader(TensorDataset(feats, coords), batch_size=2) + + stats = {"mean": np.zeros(n_genes), "std": np.ones(n_genes)} + pipeline = FlowPipeline(model=pipeline_model, stats=stats, atol=1e-1, rtol=1e-1, fast=True) + + gex_pred, coords_list = pipeline(["A", "B", "C"], loader) + + assert gex_pred.shape == (n_samples, n_genes) + assert np.isfinite(gex_pred).all() + assert sum(len(c) for c in coords_list) == n_samples