tt-model push qwen3-coder-next (container)

README.md

@@ -6,7 +6,6 @@ tags:
 - p300
 - p300x2
 - tt-model-cache
-- tt-model-catalog
 - tt-model-container
 - vllm-plugin
 pipeline_tag: text-generation
@@ -78,6 +77,13 @@ Point an OpenAI-compatible client at `http://127.0.0.1:20000` (or the port `tt-m
 (`--enable-auto-tool-choice` is on). Without sampling parameters a request uses the checkpoint's defaults
 (temperature 1.0, top_p 0.95, top_k 40).
 
+### Prefix caching
+On by default (every profile, together with chunked prefill). When a prompt starts with the same 2,048 or more tokens as a recent one (a document sent
+again, the next turn of a chat, a long shared system prompt), the server resumes from the last 2,048-token boundary
+the two share instead of prefilling from the start: vLLM reuses the attention KV blocks, and the server restores the
+Gated DeltaNet state it saved at that boundary (32 snapshots on four dies, 16 on two; each snapshot 27 MiB per die on four dies, 0.84 GiB per die in all, taken from the KV pool, least recently used evicted). A prompt seen for the first time
+runs exactly the prefill it ran without caching. Turn it off with `tt-model serve ... -- --no-enable-prefix-caching`.
+
 ### Profiles
 `tt-model serve raahemnabeel/qwen3-coder-next --profile <name>`:
 - `p300x2` (default) / `p150x4` - four dies: the faster profile.
@@ -216,6 +222,21 @@ Concurrency, 1024 input / 256 output tokens: v1 -> current
 | 16 | 64 | 4,161 -> 4,161 | 97.7 -> 31.0 | 10.2 -> 32.3 | 140.9 -> **338.8** (+140 %) | 29.07 -> 12.06 |
 | 32 | 128 | 7,135 -> 7,126 | 154.6 -> 65.3 | 6.5 -> 15.3 | 175.9 -> **344.0** (+96 %) | 46.55 -> 23.89 |
 
+### Prefix caching
+Measured 2026-10-28 on p300x2, one user, temperature 0: the same build served with caching off and on. TTFT of a
+prompt (Python source from this repo) sent a second time, and of each turn of a chat that opens with an 8k-token code
+document and resends the conversation so far (as every chat client does):
+
+| Request | caching off | caching on |
+| --- | ---: | ---: |
+| 7,160-token prompt, sent again | 1.79 s | 0.37 s (4.8x) |
+| 22,815-token prompt, sent again | 5.81 s | 0.36 s (16x) |
+| chat turns 2-6 (7.6k-8.2k tokens) | 1.94-2.08 s | 0.52-0.68 s |
+
+Prompts seen for the first time are unchanged: no-hit rows at 128 / 2,048 / 8,192 input tokens (1 user) and 128/128 at
+32 users, TTFT and TPOT within 1 % (1 user 128/128: TTFT 171 / 174 ms, TPOT 5.59 / 5.61 ms). Two dies (p300): a 25,113-token prompt sent again 7.70 -> 0.53 s, chat turns 2-4 0.62-0.63 s (turn 1, uncached, 2.35 s). Greedy output with caching on matched caching off on every resent prompt and on
+9 of 10 chat turns; the one that differs (a thinking-on turn) diverges at a near-tie token (the reused KV blocks were computed in a different prefill pass).
+
 ## Limitations
 
 - Two dies (`p300`, `p150x2`) run without the fabric layer boundary (its persistent receive buffers do not fit next to
@@ -224,6 +245,9 @@ Concurrency, 1024 input / 256 output tokens: v1 -> current
   32 rather than moved to host sampling.
 - Agentic benchmarks (SWE-Bench Verified, LiveCodeBench), which Qwen publishes for this checkpoint, were not run; the
   evals ran on the p300x2 profile only.
+- Prefix caching resumes only at 2,048-token boundaries: a prompt sharing fewer than 2,048 tokens with an earlier one
+  gets no benefit, and the shared tokens past the last boundary are prefilled again. Requests with images are not
+  cached.
 - The p150x4 and p150x2 profiles run the same meshes as p300x2 / p300 but were not run on p150 boards.
 
 ## Related packages
@@ -240,9 +264,9 @@ The exact sources the image was built from — `code/` in this repo is byte-iden
 
 | component | built from |
 | --- | --- |
-| tt-metal | [`4233e67db986d6cea6c0e98fd7d4593fc501c98e`](https://github.com/rnabeelTT/tt-metal/commit/4233e67db986d6cea6c0e98fd7d4593fc501c98e) |
+| tt-metal | [`de8e0fb2ba6f77a63d501f4009de1ab540ae96b5`](https://github.com/rnabeelTT/tt-metal/commit/de8e0fb2ba6f77a63d501f4009de1ab540ae96b5) |
 | vLLM | [`v0.26.0`](https://github.com/vllm-project/vllm/releases/tag/v0.26.0) |
 | vllm-tt-plugin | a local checkout — commit not published |
-| `code/` digest | `a62de654cd6c783c` (sha256, first 16 hex digits) |
-| built | 2026-10-27T19:31:06+00:00 by tt-model 0.1.0 |
+| `code/` digest | `e17bebd3bcea0db1` (sha256, first 16 hex digits) |
+| built | 2026-10-28T16:39:18+00:00 by tt-model 0.1.0 |
 

code/models/autoports/qwen_qwen3_coder_next/tt/generator.py

@@ -660,7 +660,9 @@ class Qwen35MoeGenerator(ReadinessGenerator):
         t = ttnn.to_torch(lg, mesh_composer=ttnn.ConcatMeshToTensor(self.mesh_device, dim=0))
         return t.float().reshape(-1, self.args.vocab_size)[:1].reshape(1, 1, -1)
 
-    def _prefill_single(self, tokens, page_table, segment=None, want_sharded=False, start=0, vision=None):
+    def _prefill_single(
+        self, tokens, page_table, segment=None, want_sharded=False, start=0, vision=None, cuts=(), on_cut=None
+    ):
         """B=1 prefill of any logical length through the model's masked fixed-bucket path (pads each segment to one of
         the pre-warmed bucket lengths, masks the GDN state to the exact valid length): full 2048-token segments carry the
         GDN/KV state forward through ``chunk_start``; the tail runs at the smallest bucket >= its length. Returns host
@@ -668,9 +670,15 @@ class Qwen35MoeGenerator(ReadinessGenerator):
         vocab-sharded logits [1,1,32,vocab/TP] of the last segment)`` for on-device sampling (row 0)."""
         segment = segment or int(getattr(self.model, "PREFILL_SEGMENT", 2048))
         T = tokens.shape[1]
-        # `start` (chunked prefill continuation): ``tokens`` are positions start .. start+T-1 of the request; the GDN
-        # state is reset by the bucket only at chunk_start 0, so a continuation runs on the state restored by the caller
-        assert start % segment == 0, f"chunk start {start} must be a multiple of the {segment}-token segment"
+        # `start` (chunked prefill continuation / prefix-cache hit): ``tokens`` are positions start .. start+T-1 of the
+        # request; the GDN state is reset by the bucket only at chunk_start 0, so a continuation runs on the state
+        # restored by the caller. Segments run [start, start + segment), ... and also end at every position in ``cuts``
+        # (prefix caching: ``on_cut(pos)`` snapshots the GDN state there); every segment start stays block aligned.
+        block = self.block_size
+        assert start % block == 0, f"chunk start {start} must be a multiple of the {block}-token KV block"
+        cuts = sorted(int(c) for c in cuts if start < int(c) <= start + T)
+        assert all(c % block == 0 for c in cuts), f"cut points {cuts} must be block aligned"
+        bounds = sorted({*range(start, start + T, segment), *(c for c in cuts if c < start + T)}) + [start + T]
         # multimodal request: ``vision`` = (packed image rows on device [n, H/TP], image_grid_thw). The M-RoPE table is
         # staged once over the WHOLE prompt (the segments slice it by position) and every segment splices its own
         # image rows (row offset = placeholders before the segment). Chunked continuations of image prompts are not
@@ -684,9 +692,9 @@ class Qwen35MoeGenerator(ReadinessGenerator):
             self.model._build_request_rope(tokens, vision_tokens)
             vkw = {"rope_staged": True}
         logits, sh = None, None
-        for cs in range(start, start + T, segment):
-            seg = tokens[:, cs - start : min(cs - start + segment, T)]
-            last = cs + segment >= start + T
+        for cs, ce in zip(bounds[:-1], bounds[1:]):
+            seg = tokens[:, cs - start : ce - start]
+            last = ce == start + T
             if vision_tokens is not None:
                 vkw["vision_tokens"] = vision_tokens
                 vkw["vis_row_offset"] = self.model._vis_row_offset_for(tokens, cs)
@@ -696,9 +704,65 @@ class Qwen35MoeGenerator(ReadinessGenerator):
                 )
             else:
                 lg = self.model.prefill_masked_bucket(seg, page_table, actual_len=seg.shape[1], chunk_start=cs, **vkw)
+            if on_cut is not None and ce in cuts:
+                on_cut(ce)
             logits = self._host_logits(lg)
         return (logits, sh) if want_sharded else logits
 
+    # ------------------------------------------------------------------ prefix caching (GDN state snapshots)
+    _STATE_BYTES = {"BFLOAT16": 2, "FLOAT32": 4, "BFLOAT8_B": 1.0625, "BFLOAT4_B": 0.5625}
+
+    def alloc_prefix_state_pool(self, n_slots):
+        """Device pool of ``n_slots`` GDN state snapshots for vLLM prefix caching (the plugin's ``prefix_state_cache``):
+        one slot = the recurrent state + conv taps of every GDN layer in the B = 1 prefill layout. The plugin's index
+        chooses the slots; :meth:`prefill_serving` saves into and restores from them. Allocated with the KV cache (before
+        any trace capture) and warmed with one save / restore. Returns the bytes per device of one slot."""
+        prev = self.model._bind_gdn_prefill_scratch() if self.max_batch_size > 1 else None
+        try:
+            dn_layers = [layer.attention for layer in self.model.layers if not layer.is_full_attention]
+            self._prefix_pool = [
+                (
+                    [ttnn.clone(dn.rec_state) for dn in dn_layers],
+                    [[ttnn.clone(c) for c in dn.conv_states] for dn in dn_layers],
+                )
+                for _ in range(int(n_slots))
+            ]
+            self._prefix_save(0)
+            self._prefix_restore(0)
+        finally:
+            if prev is not None:
+                self.model._unbind_gdn_prefill_scratch(prev)
+        recs, convs = self._prefix_pool[0]
+        tensors = recs + [t for taps in convs for t in taps]
+        return int(
+            sum(
+                math.prod(t.padded_shape) * self._STATE_BYTES.get(str(t.dtype).split(".")[-1].upper(), 2)
+                for t in tensors
+            )
+        )
+
+    def _prefix_save(self, slot):
+        """GDN state of the bound B = 1 prefill scratch (B = 1 build: the decode buffers) -> pool slot."""
+        recs, convs = self._prefix_pool[slot]
+        dn_layers = [layer.attention for layer in self.model.layers if not layer.is_full_attention]
+        for dn, rec, taps in zip(dn_layers, recs, convs):
+            ttnn.copy(dn.rec_state, rec)
+            for c, t in zip(dn.conv_states, taps):
+                ttnn.copy(c, t)
+
+    def _prefix_restore(self, slot):
+        """Pool slot -> the bound B = 1 prefill scratch, then the chunk-outer conv left context from the taps (as
+        :meth:`_pending_restore`)."""
+        from models.autoports.qwen_qwen3_coder_next.tt.mtp_rows import rebuild_conv_carry
+
+        recs, convs = self._prefix_pool[slot]
+        dn_layers = [layer.attention for layer in self.model.layers if not layer.is_full_attention]
+        for dn, rec, taps in zip(dn_layers, recs, convs):
+            ttnn.copy(rec, dn.rec_state)
+            for t, c in zip(taps, dn.conv_states):
+                ttnn.copy(t, c)
+            rebuild_conv_carry(dn)
+
     # ------------------------------------------------------------------ chunked prefill (turn 7)
     def _pending_key(self, pt_row):
         return int(
@@ -736,7 +800,7 @@ class Qwen35MoeGenerator(ReadinessGenerator):
             for t in taps:
                 ttnn.deallocate(t)
 
-    def _batched_groups(self, lens, starts, finals, vision, pt):
+    def _batched_groups(self, lens, starts, finals, vision, pt, keys=None, single=()):
         """Groups (lists of request indices, same bucket, 2..B each) of the requests the batched pass can admit: final
         text prompts starting at 0 whose length fits a batched bucket. Stale carried state on a reused first block is
         dropped here (as the single path does)."""
@@ -747,11 +811,11 @@ class Qwen35MoeGenerator(ReadinessGenerator):
 
         by_bucket = {}
         for u in range(len(lens)):
-            if not (finals[u] and starts[u] == 0 and (vision is None or vision[u] is None)):
+            if not (finals[u] and starts[u] == 0 and (vision is None or vision[u] is None)) or u in single:
                 continue
             if lens[u] > MAX_BATCHED_BUCKET:
                 continue
-            key = self._pending_key(pt[u])
+            key = keys[u] if keys is not None else self._pending_key(pt[u])
             if key in self._pending_prefill:
                 self._pending_free(self._pending_prefill.pop(key))
             by_bucket.setdefault(self.model._mask_bucket_for(lens[u]), []).append(u)
@@ -764,7 +828,17 @@ class Qwen35MoeGenerator(ReadinessGenerator):
         return groups
 
     def prefill_serving(
-        self, tokens, page_table, prompt_lens, slots, sampler=None, start_pos=None, final=None, vision=None, mtp=None
+        self,
+        tokens,
+        page_table,
+        prompt_lens,
+        slots,
+        sampler=None,
+        start_pos=None,
+        final=None,
+        vision=None,
+        mtp=None,
+        prefix_state=None,
     ):
         """vLLM prefill: ``tokens`` [N, max_T] host ids (row u valid up to ``prompt_lens[u]``), ``page_table`` [N, width]
         (row u = request u's blocks), ``slots`` = the decode slot of each request. Runs each request through the warmed
@@ -772,6 +846,10 @@ class Qwen35MoeGenerator(ReadinessGenerator):
         the prefill runs on the model's B=1 scratch). ``sampler(logits_sharded, u)`` -> host token (the adapter applies
         the request's sampling params to the on-device sampler and samples row 0 of the [1,1,32,vocab/TP] logits while
         the bucket output is still live); without it returns the host logits ``[N, 1, vocab]``.
+        ``prefix_state`` (vLLM prefix caching, the plugin's ``prefix_state_cache``): per request ``{"req_id", "restore":
+        (pool slot, pos) | None, "saves": [(pos, pool slot), ...]}`` - a hit restores the GDN state saved at ``pos`` (==
+        its start) and runs only the uncached suffix; every save position becomes a segment end whose GDN state is
+        copied into its pool slot. Every save must happen: the plugin marks the slot valid once this step returns.
         Returns ``[N, 1]`` int64 tokens (sampler) or ``[N, 1, vocab]`` float logits."""
         N = tokens.shape[0]
         lens = [int(prompt_lens[u]) for u in range(N)]  # chunked prefill: this chunk's END per request
@@ -785,10 +863,16 @@ class Qwen35MoeGenerator(ReadinessGenerator):
         A = int(getattr(self.model, "PREFILL_SEGMENT", 2048))
         V = int(self.args.vocab_size)
         outs_by_u = {}
+        ps = list(prefix_state) if prefix_state is not None else [None] * N
+        restores = [p.get("restore") if p else None for p in ps]
+        saves = [sorted((int(q), int(k)) for q, k in (p.get("saves") or ())) if p else [] for p in ps]
+        # carried chunked-prefill state is keyed by request when the plugin names it (prefix caching shares first blocks)
+        keys = [p["req_id"] if p and p.get("req_id") is not None else self._pending_key(pt[u]) for u, p in enumerate(ps)]
+        single = {u for u in range(N) if restores[u] is not None or saves[u]}
         # batched short-prompt admissions first: one pass per group of 2..4 same-bucket final text prompts; their GDN
         # state is written from the B-row scratch, their MTP head rows are skipped (MtpServing.note_batched marks the
         # slot: plain decode even when it ends up alone)
-        for g in self._batched_groups(lens, starts, finals, vision, pt):
+        for g in self._batched_groups(lens, starts, finals, vision, pt, keys=keys, single=single):
             g_lens = [lens[u] for u in g]
             g_slots = [slots[u] for u in g]
             lg, sh, hid = self.batched_prefill.prefill(
@@ -813,8 +897,13 @@ class Qwen35MoeGenerator(ReadinessGenerator):
             # request's own processed end s0 (the plugin's s may lie inside a segment: its tokens are resent with the next
             # chunk), carries the GDN state in device clones keyed by the request's first block, and returns zero logits;
             # the final chunk runs to e, samples and writes the decode slot as a whole prefill does.
-            key = self._pending_key(pt[u])
+            key = keys[u]
             pending = self._pending_prefill.get(key)
+            restore = restores[u]
+            if restore is not None:
+                assert pending is None and int(restore[1]) == starts[u] > 0, (
+                    f"request {u}: prefix-cache restore {restore} at start {starts[u]} (carried state {pending is not None})"
+                )
             if pending is not None and starts[u] == 0:
                 # a request starting from scratch on a first block that still carries another request's state (that
                 # request was aborted / preempted mid-prefill and vLLM reused its blocks): drop the stale entry
@@ -823,15 +912,21 @@ class Qwen35MoeGenerator(ReadinessGenerator):
             s0 = pending[0] if pending is not None else 0
             if pending is None:
                 assert (
-                    starts[u] == 0 or starts[u] % A == 0
+                    starts[u] == 0 or starts[u] % A == 0 or restore is not None
                 ), f"request {u}: chunk start {starts[u]} without carried state"
                 s0 = 0 if starts[u] == 0 else starts[u]
             E = lens[u] if finals[u] else (lens[u] // A) * A
+            # an intermediate chunk runs at least to its last save
+            E = max(E, max((q for q, _ in saves[u]), default=0))
+            assert all(s0 < q <= E for q, _ in saves[u]), f"request {u}: saves {saves[u]} outside [{s0}, {E}]"
+            cut_slot = dict(saves[u])
             batched = self.max_batch_size > 1
             prev = self.model._bind_gdn_prefill_scratch() if batched else None
             try:
                 if pending is not None:
                     self._pending_restore(key)
+                elif restore is not None:
+                    self._prefix_restore(int(restore[0]))
                 elif batched:
                     self.model._reset_gdn_state_for_new_sequence()
                 vis_u = vision[u] if vision is not None else None
@@ -842,6 +937,8 @@ class Qwen35MoeGenerator(ReadinessGenerator):
                         want_sharded=finals[u],
                         start=s0,
                         vision=vis_u,
+                        cuts=list(cut_slot),
+                        on_cut=lambda q: self._prefix_save(cut_slot[q]),
                     )
                     lg, sh = res if finals[u] else (res, None)
                 if finals[u]:

code/models/autoports/qwen_qwen3_coder_next/tt/generator_vllm.py

@@ -145,6 +145,22 @@ if os.environ.get("QWEN35_CHUNKED_PREFILL", "0") == "1":
 
 _MTP_K = 0  # Qwen3-Coder-Next ships no mtp.* weights
 _VISION = False  # text-only checkpoint
+
+# vLLM prefix caching (plugin prefix_state.py): QWEN35_PREFIX_STATE_SLOTS GDN state snapshots on device (~27 MiB per die on p300x2 each)
+# let a request resume from a cached KV prefix; 0 = off. A snapshot is saved at every prefill segment boundary and at the
+# last block boundary of each prompt of at least QWEN35_PREFIX_STATE_MIN tokens. Not wired with MTP.
+_PREFILL_SEGMENT = int(os.environ.get("QWEN35_PREFILL_SEGMENT", "2048"))  # the model's PREFILL_SEGMENT
+# dies the GDN state is sharded over (MESH_DEVICE): a snapshot per die is 4 / dies times the p300x2 size, so the default
+# slot count scales with the mesh (32 on four dies, 16 on two, 8 on one: the same ~0.7 GiB pool per die)
+_PREFIX_DEVICES = {"P150": 1, "N150": 1, "P100": 1, "P300": 2, "P150X2": 2}.get(
+    os.environ.get("MESH_DEVICE", "P300x2").strip().upper(), 4
+)
+_PREFIX_STATE_SLOTS = int(os.environ.get("QWEN35_PREFIX_STATE_SLOTS", str(8 * _PREFIX_DEVICES))) if _MTP_K == 0 else 0
+_PREFIX_STATE_MIN = int(os.environ.get("QWEN35_PREFIX_STATE_MIN", "1024"))
+# QWEN35_PREFIX_STATE_PROMPT_END=1 also snapshots each prompt's last block boundary: hits reach the end of the previous
+# turn, but every cold prompt pays one more prefill pass (+57..133 ms on Ornith p300x2) and its segment split (numerics)
+# changes. Off: snapshots only at segment boundaries, cold prefill unchanged and a hit is bit-identical to a cold run.
+_PREFIX_STATE_PROMPT_END = os.environ.get("QWEN35_PREFIX_STATE_PROMPT_END", "0") == "1"
 _VISION_MAX_IMAGES = 0
 
 
@@ -163,7 +179,19 @@ class Qwen3NextForCausalLM(Qwen36ForCausalLM):
     # is implemented (Qwen35MoeGenerator.decode_forward(read_from_device=False) -> read_decode_output -> process_...).
     # Overlap stays off until the page-growth refresh of the plugin revision is proven (doc/vllm_integration/README.md).
     model_capabilities = {
-        "supports_prefix_caching": False,
+        "supports_prefix_caching": _PREFIX_STATE_SLOTS > 0,
+        **(
+            {
+                "prefix_state_cache": {
+                    "slots": _PREFIX_STATE_SLOTS,
+                    "segment": _PREFILL_SEGMENT,
+                    "min_tokens": _PREFIX_STATE_MIN,
+                    "prompt_end": _PREFIX_STATE_PROMPT_END,
+                }
+            }
+            if _PREFIX_STATE_SLOTS > 0
+            else {}
+        ),
         "supports_async_decode": True,
         # Not read by this plugin revision (5961e48): overlap = vLLM's async scheduling (on by default, "Asynchronous
         # scheduling is enabled" in every served log) + supports_async_decode; the plugin's steady-decode fast path then
@@ -329,6 +357,8 @@ class Qwen3NextForCausalLM(Qwen36ForCausalLM):
     # and templates). The masked-bucket prefill programs allocate transients while a request runs (the expert-sorted
     # dispatch of the 2048-token bucket asked for one 512 MiB buffer during its trace capture) -> 1.25 GiB headroom.
     _PER_SLOT_GIB = 36.5 / 1024
+    # one prefix-cache GDN snapshot per die, p300x2 size scaled to the mesh (allocate_kv_cache checks the real size)
+    _PREFIX_SLOT_GIB = 28.0 / 1024 * 4 / _PREFIX_DEVICES
     _TRANSIENT_HEADROOM_GIB = 1.25
     _MTP_HEAD_GIB = 1.5  # QWEN35_MTP: head weights + step buffers + slack for the head K/V (also budgeted as a 13th KV layer; the first full server saw two 1017 MB head K/V buffers - the log line below reports the real size)
     _KV_LAYERS = 12  # attention layers holding paged K/V
@@ -356,6 +386,7 @@ class Qwen3NextForCausalLM(Qwen36ForCausalLM):
         c = contract.get(f"serving_tp{int(num_devices)}") or contract["datatype_sweep_stage8"]
         B = int(max_num_seqs or 1)
         nonkv = float(c["non_kv_footprint_gib"]) + (B - 1) * cls._PER_SLOT_GIB
+        nonkv += _PREFIX_STATE_SLOTS * cls._PREFIX_SLOT_GIB
         kv_bytes_per_token = int(c["kv_bytes_per_token_per_device"])
         if _MTP_K > 0:
             # the MTP head: one more attention layer's paged K/V over the same blocks (13 instead of 12 KV layers) plus
@@ -394,6 +425,13 @@ class Qwen3NextForCausalLM(Qwen36ForCausalLM):
         kv = self._gen.model.allocate_kv_caches(shape, self._gen.kv_dtype, batch_size=self._gen.max_batch_size)
         width = min(math.ceil(self._context / block), num_blocks)  # == the plugin's max_num_blocks_per_req
         self._gen.bind_kv_cache(kv, num_blocks, width)
+        if _PREFIX_STATE_SLOTS > 0:
+            slot_bytes = self._gen.alloc_prefix_state_pool(_PREFIX_STATE_SLOTS)
+            logger.info(
+                f"prefix cache: {_PREFIX_STATE_SLOTS} GDN state snapshots x {slot_bytes / 2**20:.1f} MiB per device "
+                f"(budgeted {self._PREFIX_SLOT_GIB * 1024:.1f} MiB)"
+            )
+            assert slot_bytes <= self._PREFIX_SLOT_GIB * 2**30, "prefix-cache snapshot larger than its KV-pool budget"
         if self._mtp is not None:
             self._mtp.allocate_kv(shape, self._gen.kv_dtype)
             hk = self._mtp.head.kv_cache[0]
@@ -472,8 +510,10 @@ class Qwen3NextForCausalLM(Qwen36ForCausalLM):
         N = tokens.shape[0]
         start_pos = kwargs.get("start_pos")
         starts = [int(v) for v in torch.as_tensor(start_pos).reshape(-1).tolist()] if start_pos is not None else [0] * N
-        if any(starts) and not self.model_capabilities["supports_chunked_prefill"]:
-            raise NotImplementedError("chunked / cached-prefix prefill is not supported (model_capabilities)")
+        prefix_state = kwargs.get("prefix_state")
+        restored = [bool(p and p.get("restore")) for p in prefix_state] if prefix_state is not None else [False] * N
+        if any(s and not r for s, r in zip(starts, restored)) and not self.model_capabilities["supports_chunked_prefill"]:
+            raise NotImplementedError("chunked prefill is not supported (model_capabilities)")
         # which rows end their prompt in this chunk: the plugin's intermediate mask (added to the kwargs by the chunked
         # prefill shim when QWEN35_CHUNKED_PREFILL=1); without it every row is final - the plugin never chunks a prompt
         # when chunked prefill is off. (A former heuristic "device sampling params present <=> final" marked every
@@ -525,6 +565,8 @@ class Qwen3NextForCausalLM(Qwen36ForCausalLM):
             f"Prefilling {N} request(s) into slots {slots} (chunk ends {plens}, starts {starts}, final {finals})"
         )
         chunked = dict(start_pos=starts, final=finals) if (any(starts) or not all(finals)) else {}
+        if prefix_state is not None:
+            chunked["prefix_state"] = prefix_state
         if self._mtp is not None:
             if chunked:
                 raise NotImplementedError(
@@ -549,7 +591,7 @@ class Qwen3NextForCausalLM(Qwen36ForCausalLM):
                 self._mtp.end_prefill(tok0[0], plens[0])
             self._live_slots |= set(slots)
             return out
-        if vision is not None and chunked:
+        if vision is not None and (any(starts) or not all(finals)):
             raise NotImplementedError("image prompts are prefilled whole (no chunked prefill)")
         out = self._gen.prefill_serving(tokens, pt, plens, slots, sampler=sampler, vision=vision, **chunked)
         self._live_slots |= set(slots)