diff --git a/needle/model/architecture.py b/needle/model/architecture.py
index c62cc4d8..c370d357 100644
--- a/needle/model/architecture.py
+++ b/needle/model/architecture.py
@@ -88,35 +88,85 @@ class MultiHeadAttention(nn.Module):
     num_layers: int
     dtype: jnp.dtype = jnp.bfloat16
     rope_keys_only: bool = False
+    decode: bool = False
+    cache_len: int = 0
+    cross_cache: bool = False
 
     @nn.compact
-    def __call__(self, q_input, kv_input, mask=None, rope=None):
+    def __call__(self, q_input, kv_input, mask=None, rope=None, cache_pos=None, cross_prefill=True):
         head_dim = self.d_model // self.num_heads
         kv_dim = self.num_kv_heads * head_dim
         B = q_input.shape[0]
 
+        # Query always depends on the (changing) decoder hidden state, so it is
+        # projected on every step.
         q = nn.Dense(self.d_model, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="q_proj")(q_input)
-        k = nn.Dense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="k_proj")(kv_input)
-        v = nn.Dense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="v_proj")(kv_input)
-
         q = q.reshape(B, -1, self.num_heads, head_dim).transpose(0, 2, 1, 3)
-        k = k.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)
-        v = v.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)
-
         q = ZCRMSNorm(dtype=self.dtype, name="q_norm")(q)
-        k = ZCRMSNorm(dtype=self.dtype, name="k_norm")(k)
+        if rope is not None and not self.rope_keys_only:
+            cos, sin = rope
+            q = apply_rope(q, cos, sin)
+
+        if self.cross_cache:
+            # Static cross-attention K/V cache: kv_input (encoder output) is
+            # constant across decode steps, so project K/V once when
+            # cross_prefill=True and reuse them on every subsequent step.
+            cached_key = self.variable(
+                "cache", "cross_k",
+                lambda: jnp.zeros((B, self.num_kv_heads, kv_input.shape[1], head_dim), self.dtype),
+            )
+            cached_value = self.variable(
+                "cache", "cross_v",
+                lambda: jnp.zeros((B, self.num_kv_heads, kv_input.shape[1], head_dim), self.dtype),
+            )
+            if cross_prefill:
+                k = nn.Dense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="k_proj")(kv_input)
+                v = nn.Dense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="v_proj")(kv_input)
+                k = k.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)
+                v = v.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)
+                k = ZCRMSNorm(dtype=self.dtype, name="k_norm")(k)
+                cached_key.value = k
+                cached_value.value = v
+            k = cached_key.value
+            v = cached_value.value
+        else:
+            k = nn.Dense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="k_proj")(kv_input)
+            v = nn.Dense(kv_dim, dtype=self.dtype, use_bias=False, kernel_init=default_init(), name="v_proj")(kv_input)
+            k = k.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)
+            v = v.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)
+            k = ZCRMSNorm(dtype=self.dtype, name="k_norm")(k)
+
+            if rope is not None:
+                cos, sin = rope
+                k = apply_rope(k, cos, sin)
+
+            if self.decode:
+                # Incremental single-token step: append this token's (rope'd) K/V to a
+                # fixed-size cache and attend against the whole prefix. Cache is stored
+                # in float32 to match the non-cached path's precision exactly.
+                cached_key = self.variable(
+                    "cache", "cached_key",
+                    lambda: jnp.zeros((B, self.num_kv_heads, self.cache_len, head_dim), jnp.float32),
+                )
+                cached_value = self.variable(
+                    "cache", "cached_value",
+                    lambda: jnp.zeros((B, self.num_kv_heads, self.cache_len, head_dim), jnp.float32),
+                )
+                zero = jnp.array(0, jnp.int32)
+                cp = jnp.asarray(cache_pos, jnp.int32)
+                idx = (zero, zero, cp, zero)
+                new_k = jax.lax.dynamic_update_slice(cached_key.value, k.astype(jnp.float32), idx)
+                new_v = jax.lax.dynamic_update_slice(cached_value.value, v.astype(jnp.float32), idx)
+                cached_key.value = new_k
+                cached_value.value = new_v
+                k = new_k
+                v = new_v
 
         repeats = self.num_heads // self.num_kv_heads
         if repeats > 1:
             k = jnp.repeat(k, repeats, axis=1)
             v = jnp.repeat(v, repeats, axis=1)
 
-        if rope is not None:
-            cos, sin = rope
-            if not self.rope_keys_only:
-                q = apply_rope(q, cos, sin)
-            k = apply_rope(k, cos, sin)
-
         scale = jnp.sqrt(jnp.float32(head_dim))
         attn_weights = jnp.matmul(q, k.transpose(0, 1, 3, 2)) / scale
 
@@ -244,22 +294,31 @@ class DecoderBlock(nn.Module):
     activation: str = "drelu"
     dropout_rate: float = 0.0
     no_feedforward: bool = True
+    decode: bool = False
+    cache_len: int = 0
+    use_cross_cache: bool = True
 
     @nn.compact
-    def __call__(self, x, encoder_out, self_mask=None, cross_mask=None, rope=None, ffn_mask=None, deterministic=True):
+    def __call__(self, x, encoder_out, self_mask=None, cross_mask=None, rope=None, ffn_mask=None, deterministic=True, cache_pos=None, cross_prefill=True):
         self_gate = nn.sigmoid(self.param("self_attn_gate", jinit.zeros, ())).astype(self.dtype)
         residual = x
         x = ZCRMSNorm(dtype=self.dtype)(x)
-        x = MultiHeadAttention(self.num_heads, self.num_kv_heads, self.d_model, self.num_layers, self.dtype, name="self_attn")(
-            x, x, mask=self_mask, rope=rope
+        x = MultiHeadAttention(
+            self.num_heads, self.num_kv_heads, self.d_model, self.num_layers, self.dtype,
+            decode=self.decode, cache_len=self.cache_len, name="self_attn",
+        )(
+            x, x, mask=self_mask, rope=rope, cache_pos=cache_pos
         )
         x = residual + self_gate * nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
 
         cross_gate = nn.sigmoid(self.param("cross_attn_gate", jinit.zeros, ())).astype(self.dtype)
         residual = x
         x = ZCRMSNorm(dtype=self.dtype)(x)
-        x = MultiHeadAttention(self.num_heads, self.num_kv_heads, self.d_model, self.num_layers, self.dtype, name="cross_attn")(
-            x, encoder_out, mask=cross_mask
+        x = MultiHeadAttention(
+            self.num_heads, self.num_kv_heads, self.d_model, self.num_layers, self.dtype,
+            cross_cache=(self.decode and self.use_cross_cache), name="cross_attn",
+        )(
+            x, encoder_out, mask=cross_mask, cross_prefill=cross_prefill
         )
         x = residual + cross_gate * nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
 
@@ -275,7 +334,7 @@ def __call__(self, x, encoder_out, self_mask=None, cross_mask=None, rope=None, f
 
 
 class _DecoderScanBody(nn.Module):
-    """Wraps DecoderBlock for nn.scan: carry = (x, encoder_out, self_mask, cross_mask, rope, ffn_mask)."""
+    """Wraps DecoderBlock for nn.scan: carry = (x, encoder_out, self_mask, cross_mask, rope, ffn_mask, cache_pos)."""
     num_heads: int
     num_kv_heads: int
     d_model: int
@@ -286,38 +345,47 @@ class _DecoderScanBody(nn.Module):
     dropout_rate: float = 0.0
     no_feedforward: bool = True
     deterministic: bool = True
+    decode: bool = False
+    cache_len: int = 0
+    cross_prefill: bool = True
+    use_cross_cache: bool = True
 
     @nn.compact
     def __call__(self, carry, _):
-        x, encoder_out, self_mask, cross_mask, rope, ffn_mask = carry
+        x, encoder_out, self_mask, cross_mask, rope, ffn_mask, cache_pos = carry
         x = DecoderBlock(
             self.num_heads, self.num_kv_heads, self.d_model, self.d_ff,
             self.num_layers, self.dtype, self.activation, self.dropout_rate,
-            self.no_feedforward,
-        )(x, encoder_out, self_mask, cross_mask, rope, ffn_mask, self.deterministic)
-        return (x, encoder_out, self_mask, cross_mask, rope, ffn_mask), None
+            self.no_feedforward, self.decode, self.cache_len, self.use_cross_cache,
+        )(x, encoder_out, self_mask, cross_mask, rope, ffn_mask, self.deterministic, cache_pos, self.cross_prefill)
+        return (x, encoder_out, self_mask, cross_mask, rope, ffn_mask, cache_pos), None
 
 
 class Decoder(nn.Module):
     config: TransformerConfig
 
     @nn.compact
-    def __call__(self, x, encoder_out, self_mask=None, cross_mask=None, rope=None, ffn_mask=None, deterministic=True):
+    def __call__(self, x, encoder_out, self_mask=None, cross_mask=None, rope=None, ffn_mask=None, deterministic=True, decode=False, cache_pos=None, cache_len=0, cross_prefill=True, use_cross_cache=True):
         cfg = self.config
         dt = cfg.jax_dtype
         x = x.astype(dt)
 
+        # In decode mode the KV cache ('cache' collection) is scanned per layer.
+        variable_axes = {"params": 0, "cache": 0} if decode else {"params": 0}
+        if cache_pos is None:
+            cache_pos = jnp.array(0, jnp.int32)
+
         ScanBlock = nn.scan(
             nn.remat(_DecoderScanBody),
-            variable_axes={"params": 0},
+            variable_axes=variable_axes,
             split_rngs={"params": True, "dropout": True},
             length=cfg.num_decoder_layers,
         )
-        (x, _, _, _, _, _), _ = ScanBlock(
+        (x, _, _, _, _, _, _), _ = ScanBlock(
             cfg.num_heads, cfg.num_kv_heads, cfg.d_model, cfg.d_ff,
             cfg.total_layers, dt, cfg.activation, cfg.dropout_rate,
-            cfg.no_feedforward, deterministic, name="layers",
-        )((x, encoder_out, self_mask, cross_mask, rope, ffn_mask), None)
+            cfg.no_feedforward, deterministic, decode, cache_len, cross_prefill, use_cross_cache, name="layers",
+        )((x, encoder_out, self_mask, cross_mask, rope, ffn_mask, cache_pos), None)
 
         x = ZCRMSNorm(dtype=dt)(x)
         return x
@@ -363,6 +431,31 @@ def decode(self, tgt, encoder_out, self_mask=None, cross_mask=None, deterministi
         logits = x.astype(jnp.float32) @ self.embedding.embedding.T
         return logits
 
+    def decode_step(self, tgt_token, encoder_out, self_mask, cross_mask, rope, cache_pos, cache_len, cross_prefill=True, use_cross_cache=True):
+        """Single-token incremental decode with KV cache. Returns logits (B, vocab).
+
+        tgt_token: (B, 1) int tokens for the current step.
+        self_mask: (B|1, 1, 1, cache_len) bool mask over cached key positions.
+        cross_mask: (B, 1, 1, T_enc) encoder padding mask.
+        rope: (cos, sin) for the single current position.
+        cache_pos: int32 scalar absolute position of this token.
+        cache_len: static int cache length.
+        cross_prefill: if True, project and store the (constant) cross-attention
+            K/V into the cache; if False, reuse the cached cross K/V. Set True on
+            the first (prefill) step and False on every subsequent step. Only
+            meaningful when use_cross_cache is True.
+        use_cross_cache: if True, cache the static cross-attention K/V; if False,
+            re-project them every step (original behavior).
+        """
+        x = self.embedding(tgt_token) * self.embed_scale
+        x = self.decoder(
+            x, encoder_out, self_mask=self_mask, cross_mask=cross_mask, rope=rope,
+            deterministic=True, decode=True, cache_pos=cache_pos, cache_len=cache_len,
+            cross_prefill=cross_prefill, use_cross_cache=use_cross_cache,
+        )
+        logits = x.astype(jnp.float32) @ self.embedding.embedding.T
+        return logits[:, 0, :]
+
     def _mean_pool(self, encoder_out, enc_mask):
         """Mean-pool encoder output over non-padded positions. Returns (B, d_model)."""
         if enc_mask is not None:
diff --git a/needle/model/run.py b/needle/model/run.py
index 954022e5..1911baf4 100644
--- a/needle/model/run.py
+++ b/needle/model/run.py
@@ -14,6 +14,7 @@
     TransformerConfig,
     make_causal_mask,
     make_padding_mask,
+    precompute_rope_freqs,
 )
 
 
@@ -81,6 +82,150 @@ def decode_step(params, dec_buffer, encoder_out, cross_mask):
     return _decode_fn_cache[key]
 
 
+_step_fn_cache = {}
+
+
+def _default_host_cross_cache():
+    """Default cross-attention caching policy for the host-driven decode loop.
+
+    The static cross-attention K/V cache regressed the GPU host-loop path (the
+    per-token host sync dominates and the prefill/gather overhead isn't hidden),
+    but helped on CPU. So default it off on GPU and on elsewhere. Callers can
+    override explicitly via use_cross_cache=True/False.
+    """
+    try:
+        return jax.default_backend() != "gpu"
+    except Exception:
+        return True
+
+
+def _get_step_fn(model, cache_len, use_cross_cache=True):
+    """Return a JIT-compiled single-token KV-cached decode step.
+
+    The compiled function takes the current cache and returns (logits, new_cache).
+    cache_len is baked in (static) so the cache buffers have a fixed shape.
+    use_cross_cache controls whether the static cross-attention K/V are reused
+    (True) or re-projected every step (False); it is part of the cache key since
+    it changes the compiled graph and cache structure.
+    """
+    key = (id(model), cache_len, bool(use_cross_cache))
+    if key not in _step_fn_cache:
+
+        @jax.jit
+        def step_fn(params, cache, tgt_token, encoder_out, cross_mask, self_mask, cos_p, sin_p, cache_pos):
+            # cross_prefill=False: reuse the cross-attention K/V cached at prefill.
+            logits, mutated = model.apply(
+                {"params": params, "cache": cache},
+                tgt_token, encoder_out, self_mask, cross_mask, (cos_p, sin_p),
+                cache_pos, cache_len, False, use_cross_cache,
+                method="decode_step", mutable=["cache"],
+            )
+            return logits, mutated["cache"]
+
+        _step_fn_cache[key] = step_fn
+    return _step_fn_cache[key]
+
+
+def _init_cache_step(model, params, cache_len, tgt_token, encoder_out, cross_mask, self_mask, cos_p, sin_p, cache_pos, use_cross_cache=True):
+    """First (prefill) decode step: allocates the KV cache (including the static
+    cross-attention K/V when use_cross_cache is True) and returns (logits, cache)."""
+    logits, mutated = model.apply(
+        {"params": params},
+        tgt_token, encoder_out, self_mask, cross_mask, (cos_p, sin_p),
+        cache_pos, cache_len, True, use_cross_cache,
+        method="decode_step", mutable=["cache"],
+    )
+    return logits, mutated["cache"]
+
+
+def _rope_tables(config, cache_len):
+    head_dim = config.d_model // config.num_heads
+    return precompute_rope_freqs(head_dim, cache_len, config.rope_theta)
+
+
+_prefill_fn_cache = {}
+_ondevice_fn_cache = {}
+
+
+def _get_prefill_fn(model, cache_len):
+    """Return a jitted prefill step (cross_prefill=True).
+
+    Allocates the KV cache (self-attention buffers + static cross-attention K/V),
+    runs the first decode step, and returns (next_token (B,), cache). The static
+    cross K/V stored here are reused by the on-device loop without recomputation.
+    """
+    key = (id(model), cache_len)
+    if key not in _prefill_fn_cache:
+
+        @jax.jit
+        def prefill_fn(params, first_token, encoder_out, cross_mask, cos0, sin0, eos_id):
+            self_mask = (jnp.arange(cache_len) <= 0)[None, None, None, :]
+            logits, mutated = model.apply(
+                {"params": params},
+                first_token, encoder_out, self_mask, cross_mask, (cos0, sin0),
+                jnp.array(0, jnp.int32), cache_len, True, True,
+                method="decode_step", mutable=["cache"],
+            )
+            nxt = jnp.argmax(logits, axis=-1).astype(jnp.int32)  # (B,)
+            return nxt, mutated["cache"]
+
+        _prefill_fn_cache[key] = prefill_fn
+    return _prefill_fn_cache[key]
+
+
+def _get_ondevice_fn(model, cache_len):
+    """Return a jitted, fully on-device greedy decode loop (no per-token host sync).
+
+    Runs the whole autoregressive loop inside lax.while_loop with early stop when
+    every sequence has emitted EOS. Assumes the cache has been prefilled (so cross
+    K/V are already stored); each step reuses them (cross_prefill=False). Greedy
+    argmax only (no constrained decoding, which requires host-side automaton
+    state). Returns (gen_tokens, lengths).
+    """
+    key = (id(model), cache_len)
+    if key not in _ondevice_fn_cache:
+        max_steps = cache_len - 1
+
+        @jax.jit
+        def loop_fn(params, cache, tok0, encoder_out, cross_mask, cos_full, sin_full, eos_id):
+            B = tok0.shape[0]
+            # tok0 is the token produced by the prefill step (stored at column 0).
+            gen0 = jnp.zeros((B, max_steps), jnp.int32).at[:, 0].set(tok0)
+            fin0 = tok0 == eos_id
+            len0 = jnp.where(fin0, 0, max_steps)
+
+            def cond(state):
+                pos, cur, cache, gen, fin, length = state
+                return jnp.logical_and(pos < max_steps, jnp.logical_not(jnp.all(fin)))
+
+            def body(state):
+                pos, cur, cache, gen, fin, length = state
+                cos_p = jax.lax.dynamic_slice_in_dim(cos_full, pos, 1, axis=0)
+                sin_p = jax.lax.dynamic_slice_in_dim(sin_full, pos, 1, axis=0)
+                self_mask = (jnp.arange(cache_len) <= pos)[None, None, None, :]
+                logits, mutated = model.apply(
+                    {"params": params, "cache": cache},
+                    cur, encoder_out, self_mask, cross_mask, (cos_p, sin_p),
+                    pos, cache_len, False, True,
+                    method="decode_step", mutable=["cache"],
+                )
+                nxt = jnp.argmax(logits, axis=-1).astype(jnp.int32)  # (B,)
+                is_eos = nxt == eos_id
+                newly = jnp.logical_and(jnp.logical_not(fin), is_eos)
+                # Record token count before the first EOS (EOS itself is not emitted).
+                length = jnp.where(newly, pos, length)
+                gen = gen.at[:, pos].set(nxt)
+                fin = jnp.logical_or(fin, is_eos)
+                return (pos + 1, nxt[:, None], mutated["cache"], gen, fin, length)
+
+            state = (jnp.array(1, jnp.int32), tok0[:, None], cache, gen0, fin0, len0)
+            _, _, _, gen, _, length = jax.lax.while_loop(cond, body, state)
+            return gen, length
+
+        _ondevice_fn_cache[key] = loop_fn
+    return _ondevice_fn_cache[key]
+
+
 def load_checkpoint(path):
     with open(path, "rb") as f:
         data = pickle.load(f)
@@ -103,12 +248,30 @@ def _build_encoder_input(tokenizer, query, tools, max_enc_len=DEFAULT_MAX_ENC_LE
     return q_toks + [tools_sep_id] + t_toks
 
 
-def generate(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, seed=0, stream=True, task_token_id=None, normalize=True, constrained=True):
+def generate(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, seed=0, stream=True, task_token_id=None, normalize=True, constrained=True, use_cache=True, use_cross_cache=None):
     """Generate tool-call output.
 
     Encoder: [query_tokens..., <tools>, tools_tokens...] truncated to max_enc_len.
     Decoder: prefilled with [EOS], model predicts <tool_call> first, then answer tokens.
+
+    use_cache=True uses the incremental KV-cached decode path (single-token steps);
+    use_cache=False uses the original full-buffer re-decode path.
+    use_cross_cache toggles the static cross-attention K/V cache on the host loop
+    (None=auto: off on GPU, on elsewhere); the on-device path always uses it.
     """
+    if use_cache and not constrained and not stream:
+        # Fully on-device single-token loop (no per-token host sync).
+        return generate_batch_ondevice(
+            model, params, tokenizer, [query], [tools], max_gen_len=max_gen_len,
+            max_enc_len=max_enc_len, normalize=normalize,
+        )[0]
+    if use_cache:
+        return generate_cached(
+            model, params, tokenizer, query, tools=tools, max_gen_len=max_gen_len,
+            max_enc_len=max_enc_len, seed=seed, stream=stream, task_token_id=task_token_id,
+            normalize=normalize, constrained=constrained, use_cross_cache=use_cross_cache,
+        )
+
     name_map = {}
     if normalize:
         tools, name_map = normalize_tools(tools)
@@ -179,14 +342,122 @@ def generate(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MA
     return result
 
 
-def generate_batch(model, params, tokenizer, queries, tools_list, max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, normalize=True, constrained=True):
+def generate_cached(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, seed=0, stream=True, task_token_id=None, normalize=True, constrained=True, use_cross_cache=None):
+    """Incremental KV-cached single-example generation. Token-for-token equivalent
+    to generate(..., use_cache=False) but O(T^2) instead of O(T^3) work.
+
+    use_cross_cache=None auto-selects (off on GPU, on elsewhere); pass True/False
+    to force the static cross-attention K/V cache on or off."""
+    if use_cross_cache is None:
+        use_cross_cache = _default_host_cross_cache()
+    name_map = {}
+    if normalize:
+        tools, name_map = normalize_tools(tools)
+
+    enc_tokens = _build_encoder_input(tokenizer, query, tools, max_enc_len)
+    enc_input = jnp.array([enc_tokens])
+
+    pad_id = tokenizer.pad_token_id
+    eos_id = tokenizer.eos_token_id
+
+    src_mask = make_padding_mask(enc_input, pad_id)
+    encoder_out, enc_mask = model.apply(
+        {"params": params}, enc_input, src_mask=src_mask, method="encode"
+    )
+
+    cache_len = max_gen_len
+    cos_full, sin_full = _rope_tables(model.config, cache_len)
+    step_fn = _get_step_fn(model, cache_len, use_cross_cache)
+
+    constrained_decoder = None
+    if constrained:
+        from .constrained import build_constrained_decoder
+        constrained_decoder = build_constrained_decoder([tools], tokenizer)
+
+    if stream:
+        sys.stdout.write("\n")
+        sys.stdout.flush()
+
+    generated_tokens = []
+    current_token = jnp.array([[eos_id]], dtype=jnp.int32)
+    cache = None
+
+    for i in range(0, max_gen_len - 1):
+        cos_p = cos_full[i:i + 1]
+        sin_p = sin_full[i:i + 1]
+        self_mask = (jnp.arange(cache_len) <= i)[None, None, None, :]
+        pos = jnp.array(i, dtype=jnp.int32)
+
+        if cache is None:
+            logits, cache = _init_cache_step(
+                model, params, cache_len, current_token, encoder_out, enc_mask,
+                self_mask, cos_p, sin_p, pos, use_cross_cache,
+            )
+        else:
+            logits, cache = step_fn(
+                params, cache, current_token, encoder_out, enc_mask,
+                self_mask, cos_p, sin_p, pos,
+            )
+
+        next_logits = logits[0]
+
+        if constrained_decoder and constrained_decoder.is_active(0):
+            logits_np = np.array(next_logits)
+            logits_np = constrained_decoder.constrain_logits(logits_np, 0)
+            next_token = int(np.argmax(logits_np))
+        else:
+            next_token = int(jnp.argmax(next_logits))
+
+        if constrained_decoder:
+            constrained_decoder.update(0, next_token)
+
+        if next_token == eos_id:
+            break
+
+        generated_tokens.append(next_token)
+        current_token = jnp.array([[next_token]], dtype=jnp.int32)
+
+        if stream:
+            sys.stdout.write(tokenizer.decode([next_token]))
+            sys.stdout.flush()
+
+    if stream:
+        sys.stdout.write("\n")
+
+    result = tokenizer.decode(generated_tokens)
+    if result.startswith("<tool_call>"):
+        result = result[len("<tool_call>"):]
+    if normalize and name_map:
+        result = restore_tool_names(result, name_map)
+    return result
+
+
+def generate_batch(model, params, tokenizer, queries, tools_list, max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, normalize=True, constrained=True, use_cache=True, use_cross_cache=None):
     """Batch-generate tool-call outputs for multiple examples at once.
 
     Encoder: [query_tokens..., <tools>, tools_tokens...] per example, truncated to max_enc_len.
     Decoder: prefilled with [EOS], model predicts <tool_call> first, then answer tokens.
 
     Returns a list of decoded strings, one per example.
+
+    use_cross_cache toggles the static cross-attention K/V cache on the host loop
+    (None=auto: off on GPU, on elsewhere); the on-device path always uses it.
     """
+    if use_cache and not constrained:
+        # Fully on-device loop (no per-token host sync). Constrained decoding
+        # cannot use this path (host-side automaton state), so it falls through
+        # to the host-driven cached loop below.
+        return generate_batch_ondevice(
+            model, params, tokenizer, queries, tools_list, max_gen_len=max_gen_len,
+            max_enc_len=max_enc_len, normalize=normalize,
+        )
+    if use_cache:
+        return generate_batch_cached(
+            model, params, tokenizer, queries, tools_list, max_gen_len=max_gen_len,
+            max_enc_len=max_enc_len, normalize=normalize, constrained=constrained,
+            use_cross_cache=use_cross_cache,
+        )
+
     name_maps = []
     if normalize:
         normed_tools = []
@@ -265,6 +536,165 @@ def generate_batch(model, params, tokenizer, queries, tools_list, max_gen_len=DE
     return results
 
 
+def generate_batch_cached(model, params, tokenizer, queries, tools_list, max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, normalize=True, constrained=True, use_cross_cache=None):
+    """Incremental KV-cached batched generation. Token-for-token equivalent to
+    generate_batch(..., use_cache=False).
+
+    use_cross_cache=None auto-selects (off on GPU, on elsewhere); pass True/False
+    to force the static cross-attention K/V cache on or off."""
+    if use_cross_cache is None:
+        use_cross_cache = _default_host_cross_cache()
+    name_maps = []
+    if normalize:
+        normed_tools = []
+        for t in tools_list:
+            nt, nm = normalize_tools(t)
+            normed_tools.append(nt)
+            name_maps.append(nm)
+        tools_list = normed_tools
+
+    B = len(queries)
+    pad_id = tokenizer.pad_token_id
+    eos_id = tokenizer.eos_token_id
+
+    enc_token_lists = []
+    for q, t in zip(queries, tools_list):
+        enc_token_lists.append(_build_encoder_input(tokenizer, q, t, max_enc_len))
+    max_enc = max(len(toks) for toks in enc_token_lists)
+    enc_input = np.full((B, max_enc), pad_id, dtype=np.int32)
+    for i, toks in enumerate(enc_token_lists):
+        enc_input[i, :len(toks)] = toks
+    enc_input = jnp.array(enc_input)
+    src_mask = make_padding_mask(enc_input, pad_id)
+
+    encoder_out, enc_mask = model.apply(
+        {"params": params}, enc_input, src_mask=src_mask, method="encode"
+    )
+
+    cache_len = max_gen_len
+    cos_full, sin_full = _rope_tables(model.config, cache_len)
+    step_fn = _get_step_fn(model, cache_len, use_cross_cache)
+
+    constrained_decoder = None
+    if constrained:
+        from .constrained import build_constrained_decoder
+        constrained_decoder = build_constrained_decoder(tools_list, tokenizer)
+
+    finished = [False] * B
+    gen_tokens = [[] for _ in range(B)]
+    current_token = np.full((B, 1), eos_id, dtype=np.int32)
+    cache = None
+
+    for pos_i in range(0, max_gen_len - 1):
+        cos_p = cos_full[pos_i:pos_i + 1]
+        sin_p = sin_full[pos_i:pos_i + 1]
+        self_mask = (jnp.arange(cache_len) <= pos_i)[None, None, None, :]
+        pos = jnp.array(pos_i, dtype=jnp.int32)
+        tok_in = jnp.array(current_token)
+
+        if cache is None:
+            logits, cache = _init_cache_step(
+                model, params, cache_len, tok_in, encoder_out, enc_mask,
+                self_mask, cos_p, sin_p, pos, use_cross_cache,
+            )
+        else:
+            logits, cache = step_fn(
+                params, cache, tok_in, encoder_out, enc_mask,
+                self_mask, cos_p, sin_p, pos,
+            )
+
+        for i in range(B):
+            if finished[i]:
+                continue
+            if constrained_decoder and constrained_decoder.is_active(i):
+                logits_np = np.array(logits[i])
+                logits_np = constrained_decoder.constrain_logits(logits_np, i)
+                next_token = int(np.argmax(logits_np))
+            else:
+                next_token = int(jnp.argmax(logits[i]))
+            if constrained_decoder:
+                constrained_decoder.update(i, next_token)
+            if next_token == eos_id:
+                finished[i] = True
+                continue
+            gen_tokens[i].append(next_token)
+            current_token[i, 0] = next_token
+
+        if all(finished):
+            break
+
+    results = []
+    for i in range(B):
+        text = tokenizer.decode(gen_tokens[i])
+        if text.startswith("<tool_call>"):
+            text = text[len("<tool_call>"):]
+        results.append(text)
+    if normalize and name_maps:
+        results = [restore_tool_names(r, nm) for r, nm in zip(results, name_maps)]
+    return results
+
+
+def generate_batch_ondevice(model, params, tokenizer, queries, tools_list, max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, normalize=True):
+    """Fully on-device KV-cached batched generation (greedy, unconstrained).
+
+    The autoregressive loop runs inside a single jitted lax.while_loop, so there
+    is no per-token host sync (only one device->host transfer at the end). This
+    removes the launch-bound overhead that hurts small-batch GPU decoding.
+    Token-for-token equivalent to generate_batch(..., use_cache=True) with
+    constrained=False.
+    """
+    name_maps = []
+    if normalize:
+        normed_tools = []
+        for t in tools_list:
+            nt, nm = normalize_tools(t)
+            normed_tools.append(nt)
+            name_maps.append(nm)
+        tools_list = normed_tools
+
+    B = len(queries)
+    pad_id = tokenizer.pad_token_id
+    eos_id = tokenizer.eos_token_id
+
+    enc_token_lists = [_build_encoder_input(tokenizer, q, t, max_enc_len) for q, t in zip(queries, tools_list)]
+    max_enc = max(len(toks) for toks in enc_token_lists)
+    enc_input = np.full((B, max_enc), pad_id, dtype=np.int32)
+    for i, toks in enumerate(enc_token_lists):
+        enc_input[i, :len(toks)] = toks
+    enc_input = jnp.array(enc_input)
+    src_mask = make_padding_mask(enc_input, pad_id)
+
+    encoder_out, enc_mask = model.apply(
+        {"params": params}, enc_input, src_mask=src_mask, method="encode"
+    )
+
+    cache_len = max_gen_len
+    cos_full, sin_full = _rope_tables(model.config, cache_len)
+    first_token = jnp.full((B, 1), eos_id, dtype=jnp.int32)
+    eos = jnp.array(eos_id, dtype=jnp.int32)
+
+    # Prefill step: allocates the cache and projects the static cross-attention
+    # K/V once. The on-device loop then reuses them on every subsequent step.
+    prefill_fn = _get_prefill_fn(model, cache_len)
+    loop_fn = _get_ondevice_fn(model, cache_len)
+
+    tok0, cache = prefill_fn(params, first_token, encoder_out, enc_mask, cos_full[0:1], sin_full[0:1], eos)
+    gen, length = loop_fn(params, cache, tok0, encoder_out, enc_mask, cos_full, sin_full, eos)
+    gen = np.array(gen)          # single device -> host transfer
+    length = np.array(length)
+
+    results = []
+    for i in range(B):
+        toks = gen[i, :int(length[i])].tolist()
+        text = tokenizer.decode(toks)
+        if text.startswith("<tool_call>"):
+            text = text[len("<tool_call>"):]
+        results.append(text)
+    if normalize and name_maps:
+        results = [restore_tool_names(r, nm) for r, nm in zip(results, name_maps)]
+    return results
+
+
 def main(args):
     print(f"Loading checkpoint: {args.checkpoint}")
     params, config = load_checkpoint(args.checkpoint)