Faster inference by caching the KV projections

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)

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,

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))

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))

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