Flash Next on CUDA: image and video input on two ranks, concurrent and serial engines

src/tensorfold/families/qwen4_exp/cuda/decode.py

@@ -194,7 +194,7 @@ def twin(self) -> "Engine":
     def forward(self, tokens: Sequence[int]) -> torch.Tensor:
         """A decode step's forward (a CUDA graph when enabled): logits [R, V]."""
 
-        if self.graphs is not None:
+        if self.graphs is not None and self.st.image_positions is None:    # graphs are captured for text rows
             return self.graphs.forward(tokens)
         return forward(self.w, self.st, self.buf, tokens)
 
@@ -245,7 +245,7 @@ def sample_draft(self, logits: torch.Tensor, position: int, sampling: Sampling |
         return int(tok), float(np.exp(float(got[hit[0]]) - float(got[k]))) if len(hit) else 0.0
 
     def mtp_forward(self, next_tokens: Sequence[int], streams: torch.Tensor) -> torch.Tensor:
-        if self.graphs is not None:
+        if self.graphs is not None and self.st.image_positions is None:
             return self.graphs.mtp_forward(next_tokens, streams)
         return mtp_forward(self.w, self.st, self.mbuf, next_tokens, streams)
 

src/tensorfold/families/qwen4_exp/cuda/engine.py

@@ -72,8 +72,6 @@ def __init__(self, model_dir: Path, *, depth: int = DEPTH, confidence: float = C
         if (exl3 or quant_method(read_config(model_dir)) == "modelopt") and tp != 1:
             raise ValueError(f"{'EXL3 packs' if exl3 else 'NVFP4 checkpoints'} of Flash Next run on one GPU: drop --tp "
                              "2, or serve the MLX checkpoint (TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP) on two")
-        if vision and (streams < 2 or tp != 1):
-            raise ValueError("image input on Flash Next runs on one GPU with --parallel 2 or more")
         if exl3 and ple_on_ssd:
             raise ValueError("--ple-on-ssd reads the MLX checkpoint's n-gram tables; an EXL3 pack maps its own table "
                              "from its file, so drop --ple-on-ssd")
@@ -136,7 +134,7 @@ def __init__(self, model_dir: Path, *, depth: int = DEPTH, confidence: float = C
                   f"planned peak {peak:.2f} GiB at the admitted window", flush=True)
         self.max_len = self.capacity_plan["cache_slots"]
         if tp == 2:
-            self._same_settings(torch, ids)
+            self._same_settings(torch, ids, vision)
         build_kernels(exl3=exl3, nvfp4=not exl3 and quant_method(read_config(model_dir)) == "modelopt",
                       solo=streams == 1 or (graphs and mtp))
         from concurrent.futures import wait
@@ -165,15 +163,20 @@ def __init__(self, model_dir: Path, *, depth: int = DEPTH, confidence: float = C
         from tensorfold.cuda.markers import resume_points
 
         self.points = resume_points(model_dir)          # a prompt's message starts to keep states at, or None
-        if vision:
+        if vision and rank == 0:
             from tensorfold.vision.qwen_cuda import QwenCudaVision
 
             self.vision = QwenCudaVision(model_dir, torch.device("cuda", 0),
                                          allow_urls=vision_urls)
             torch.cuda.empty_cache()
             print(f"[tensorfold] vision: image{' and video' if self.vision.videos else ''} input, a "
                   f"{self.vision.weight_bytes / 2**30:.2f} GiB tower with {vision_workspace() / 2**30:.2f} GiB of "
-                  f"workspace reserved{'; https URLs allowed' if vision_urls else ''}", flush=True)
+                  f"workspace reserved{'; https URLs allowed' if vision_urls else ''}"
+                  f"{'; rank 1 receives each request\'s image features from this rank' if tp == 2 else ''}",
+                  flush=True)
+        elif vision:                                   # rank 1: no tower; rank 0's features arrive with each admission
+            print("[tensorfold] vision: image input on two ranks; this rank attaches the features rank 0 encodes",
+                  flush=True)
         # ``streams`` > 1: up to that many requests decoded together, every stream's chain in one forward
         self.concurrent = streams > 1
         self.refuses_structured_output = (
@@ -264,7 +267,7 @@ def __init__(self, model_dir: Path, *, depth: int = DEPTH, confidence: float = C
               f"decode graphs captured; idle prompt pieces {self.prefill_rows} rows; "
               f"prompt kernels warmed in {warm_s:.1f}s", flush=True)
 
-    def _same_settings(self, torch, ids) -> None:
+    def _same_settings(self, torch, ids, vision: bool = False) -> None:
         """Both ranks must decode with the same rule, context, draft vocabulary and KV cache, or they would fall out of step: refuse to start otherwise."""
 
         from .kvcache import BITS_OF
@@ -273,14 +276,15 @@ def _same_settings(self, torch, ids) -> None:
         mine = torch.tensor([self.depth, round(self.confidence * 1e6), self.max_len, self.streams,
                              int(self.graphs_enabled),
                              len(ids) if ids is not None else -1, total, BITS_OF[self.kv_dtype],
-                             self.prefill_rows, int(prompt_precision.fp8())], dtype=torch.int64, device="cuda")
+                             self.prefill_rows, int(bool(vision)), int(prompt_precision.fp8())],
+                            dtype=torch.int64, device="cuda")
         both = torch.empty((2 * mine.numel(),), dtype=torch.int64, device="cuda")
         self.comm.all_gather(mine, both)
         both = both.view(2, -1).cpu()
         prompt_precision.same_on_ranks(int(both[0, -1]), int(both[1, -1]))
         if not torch.equal(both[0], both[1]):
             raise RuntimeError(f"the two ranks were started with different settings (drafts, confidence, context, "
-                               f"parallel streams, graphs, draft vocabulary, KV cache, prompt rows): "
+                               f"parallel streams, graphs, draft vocabulary, KV cache, prompt rows, --vision): "
                                f"rank 0 {both[0].tolist()}, rank 1 {both[1].tolist()}")
 
     def _key(self, n: int) -> str:
@@ -296,13 +300,13 @@ def shutdown(self) -> None:
             self.comm.store.set(self._key(self.served), json.dumps({"stop": True}))
 
     def _share(self, prompt: list[int], max_tokens: int, sampling, draft: bool, cached: int, constraint=None,
-               stop_eos: bool = True) -> tuple:
+               stop_eos: bool = True, images: dict | None = None) -> tuple:
         from tensorfold.engine.grammar import pack
 
         points = getattr(self, "points", None)
         body = {"prompt": prompt, "max_tokens": max_tokens, "draft": bool(draft), "cached": int(cached),
-                "points": list(points(prompt)) if draft and points is not None else [],
-                "stop_eos": bool(stop_eos),
+                "points": list(points(prompt)) if draft and points is not None and images is None else [],
+                "stop_eos": bool(stop_eos), "images": images,       # an image prompt's rows, offset and feature shape
                 "sampling": None if sampling is None else [int(sampling.seed), float(sampling.temperature),
                                                            int(sampling.top_k), float(sampling.top_p),
                                                            float(sampling.min_p)],
@@ -338,7 +342,7 @@ def _unpack(text: str) -> tuple | None:
         s = body["sampling"]
         return (body["prompt"], body["max_tokens"], None if s is None else Sampling(*s),
                 body["draft"], body["cached"], body.get("grammar") or [], body.get("stop_eos", True),
-                body.get("points", []))
+                body.get("images"), body.get("points", []))      # the points stay last (older followers)
 
     @property
     def context_window(self) -> int:
@@ -378,7 +382,7 @@ def supports_logprobs(self) -> bool:
         return self.tp == 1
 
     def _serial(self, prompt: list[int], max_tokens: int, sampling, on_tokens, constraint=None,
-                stop_eos: bool = True, probabilities=None) -> dict[str, Any]:
+                stop_eos: bool = True, probabilities=None, vision=None) -> dict[str, Any]:
         """One token a round from a fresh prefill in the serial engine's own state (no drafts, no kept states)."""
 
         import torch
@@ -388,7 +392,8 @@ def _serial(self, prompt: list[int], max_tokens: int, sampling, on_tokens, const
         if self.serial is None:
             self.serial = self.e.twin()
         t0 = time.perf_counter()
-        first = prefill(self.serial, prompt, sampling, mtp=False, constraint=constraint, probabilities=probabilities)
+        first = prefill(self.serial, prompt, sampling, mtp=False, constraint=constraint, probabilities=probabilities,
+                        vision=vision)
         torch.cuda.synchronize()
         stats: dict[str, Any] = {"prefill_s": round(time.perf_counter() - t0, 4), "cached": 0, "drafts": False}
         if (on_tokens is not None and on_tokens([first])) or (stop_eos and first in self.eos) or max_tokens <= 1:
@@ -399,12 +404,14 @@ def _serial(self, prompt: list[int], max_tokens: int, sampling, on_tokens, const
         return stats
 
     def _decode(self, prompt: list[int], max_tokens: int, sampling, on_tokens, hit, constraint=None,
-                stop_eos: bool = True, probabilities=None, points=None) -> dict[str, Any]:
+                stop_eos: bool = True, probabilities=None, points=None, vision=None) -> dict[str, Any]:
         import torch
 
         from .decode import entry_end, mtp_decode, prefill, serial_decode
 
         t0 = time.perf_counter()
+        if vision is not None:
+            hit = None                                  # an image prompt prefills from its start
         self._start_from(hit)
         from tensorfold.cuda.markers import MIN_GAP
 
@@ -415,10 +422,15 @@ def _decode(self, prompt: list[int], max_tokens: int, sampling, on_tokens, hit,
         stops = [p for p in points if cached + MIN_GAP <= p < end]
         def keep(p, snap, tail):
             self._remember(list(prompt[:p]), {"state": snap, "tail": tail})
-        first = prefill(self.e, prompt, sampling, resume=hit[1] if hit else None, constraint=constraint,
-                        probabilities=probabilities, keep_at=end, stops=stops, keep=keep)
-        # the state one token before the prompt's end, so the same prompt or a next turn resumes from it
-        self._remember(list(prompt[:end]), self.e.kept)
+        if vision is not None:
+            # image rows and positions on the state; no kept snapshot (text only), the rounds run eagerly
+            # (Engine.forward: the captured graphs know no image positions), as the concurrent decoder does
+            first = prefill(self.e, prompt, sampling, constraint=constraint, probabilities=probabilities, vision=vision)
+        else:
+            first = prefill(self.e, prompt, sampling, resume=hit[1] if hit else None, constraint=constraint,
+                            probabilities=probabilities, keep_at=end, stops=stops, keep=keep)
+            # the state one token before the prompt's end, so the same prompt or a next turn resumes from it
+            self._remember(list(prompt[:end]), self.e.kept)
         torch.cuda.synchronize()
         stats: dict[str, Any] = {"prefill_s": round(time.perf_counter() - t0, 4), "cached": len(hit[0]) if hit else 0,
                                  "drafts": True}
@@ -461,18 +473,27 @@ def generate(self, prompt: list[int], max_tokens: int, sampling,
                                          **grammar, **({"background": True} if background else {}),
                                          probabilities=probabilities,
                                          **({"vision": vision} if vision is not None else {}))
-        hit = self._resume(prompt) if draft else None
+        hit = self._resume(prompt) if draft and vision is None else None
         points = None
+        # the serial engine's image prompts: rank 0 encodes before the request is shared (a refused image fails
+        # here alone), and after sharing sends the features to rank 1, which attaches the same tensors
+        encoded = self.vision.encode(vision, prompt) if vision is not None else None
         if self.tp == 2:                     # rank 0 decodes exactly what it hands rank 1
-            prompt, max_tokens, sampling, draft, _, _, stop_eos, points = self._share(
-                prompt, max_tokens, sampling, draft, len(hit[0]) if hit else 0, constraint, stop_eos)
+            from .vision_ranks import describe, exchange
+
+            prompt, max_tokens, sampling, draft, _, _, stop_eos, images, points = self._share(
+                prompt, max_tokens, sampling, draft, len(hit[0]) if hit else 0, constraint, stop_eos,
+                images=describe(encoded))
             self.served += 1
+            if images is not None:
+                encoded = exchange(self.comm, 0, len(prompt), images, encoded, hidden=int(self.w.cfg.hidden))
             emit = on_tokens
             on_tokens = lambda new: (emit(new), False)[1]       # noqa: E731  both ranks decode to the end
         if not draft:
-            return self._serial(prompt, max_tokens, sampling, on_tokens, constraint, stop_eos, probabilities=probabilities)
+            return self._serial(prompt, max_tokens, sampling, on_tokens, constraint, stop_eos,
+                                probabilities=probabilities, vision=encoded)
         return self._decode(prompt, max_tokens, sampling, on_tokens, hit, constraint, stop_eos,
-                            probabilities=probabilities, points=points)
+                            probabilities=probabilities, points=points, vision=encoded)
 
     def score_labels(self, prompt_ids, label_ids) -> tuple[list[float], float]:
         """Score the prompt's final row in a fresh state without adding a kept decision prefix."""
@@ -537,7 +558,9 @@ def follow(self) -> None:
             request = self._receive()
             if request is None:
                 return
-            prompt, max_tokens, sampling, draft, cached, packed, stop_eos, points = request
+            prompt, max_tokens, sampling, draft, cached, packed, stop_eos = request[:7]
+            # eight fields from a leader without image support, nine with the image description before the points
+            images, points = (None, request[7]) if len(request) == 8 else request[7:9]
             constraint = None
             if packed:                                      # the request's grammar, compiled here as on rank 0
                 from tensorfold.engine import grammar
@@ -551,9 +574,15 @@ def follow(self) -> None:
                 if hit is None:
                     raise RuntimeError(f"rank 1 has no kept state for the {cached} tokens rank 0 resumes from")
             try:
+                encoded = None
+                if images is not None:                      # rank 0's image features, attached as it attaches them
+                    from .vision_ranks import exchange
+
+                    encoded = exchange(self.comm, 1, len(prompt), images, hidden=int(self.w.cfg.hidden))
                 if draft:
-                    self._decode(prompt, max_tokens, sampling, None, hit, constraint, stop_eos, points=points)
+                    self._decode(prompt, max_tokens, sampling, None, hit, constraint, stop_eos, points=points,
+                                 vision=encoded)
                 else:
-                    self._serial(prompt, max_tokens, sampling, None, constraint, stop_eos)
+                    self._serial(prompt, max_tokens, sampling, None, constraint, stop_eos, vision=encoded)
             except ValueError as exc:                       # rank 0 raised at the same point on the same input
                 print(f"[tensorfold] request {self.served} failed on both ranks: {exc}", flush=True)

src/tensorfold/families/qwen4_exp/cuda/image_rows.py

@@ -24,13 +24,20 @@ def rope_axis(pos, ROPE, DELTA, length, index, MODE: tl.constexpr, S1: tl.conste
 
 
 def begin(engine, stream, tower) -> None:
-    """Encode an image request on its admitted slot before the normal prompt passes."""
+    """Encode an image request on its admitted slot before the normal prompt passes.
 
-    if stream.vision is not None:
+    On two ranks the stream already carries the features rank 0 encoded (``multi_tp._share_vision``): both ranks
+    attach the same tensors, and rank 1 needs no tower."""
+
+    if stream.vision is None:
+        return
+    encoded = stream.vision if hasattr(stream.vision, "features") else None
+    if encoded is None:
         if tower is None:
             raise ValueError("image inputs require starting this server with --vision")
-        attach(engine.st, tower.encode(stream.vision, stream.prompt), len(stream.prompt))
-        stream.vision = None
+        encoded = tower.encode(stream.vision, stream.prompt)
+    attach(engine.st, encoded, len(stream.prompt))
+    stream.vision = None
 
 
 def attach(st, encoded, length: int) -> None:

src/tensorfold/families/qwen4_exp/cuda/multi.py

@@ -250,10 +250,17 @@ def admit(self, s: Stream, told=None) -> None:
         s.count = max(1, min(s.count, room))
         if any(x.waiting for x in self.streams.values()):
             raise NoRoom("streams already wait for memory; a new request waits until one finishes")
-        if self.w.comm is not None and any(x is not None for x in (s.constraint, s.vision, s.probabilities)):
+        if self.w.comm is not None and any(x is not None for x in (s.constraint, s.probabilities)):
             raise ValueError("concurrent Flash Next on two ranks serves text without grammars or logprobs")
         t0 = time.perf_counter()
         if self.w.comm is not None:
+            if s.vision is not None and self.link is not None and not hasattr(s.vision, "features"):
+                # rank 0 encodes the images before rank 1 hears of the request: an image the tower refuses
+                # (too large, a video without its frontend) fails here and no message reaches rank 1; a request
+                # admitted again after waiting for memory keeps the features it already has
+                if self.vision is None:
+                    raise ValueError("image inputs require starting this server with --vision")
+                s.vision = self.vision.encode(s.vision, s.prompt)
             st, resume, s.cached = self._prepare_admission(s, told)
         else:
             st, resume, s.cached = self._slot_for(list(s.prompt), s.draft and s.vision is None)

src/tensorfold/families/qwen4_exp/cuda/multi_plan.py

@@ -19,6 +19,7 @@ class Shadow:
     def __init__(self, st, index: int, actions: list) -> None:
         self.source, self.index, self.actions = st, index, actions
         self.capacity, self.limit, self.pos, self.mtp_len = st.capacity, st.limit, st.pos, st.mtp_len
+        self.images = getattr(st, "image_positions", None) is not None     # an image prompt's rotary positions
 
     def cache_bytes(self, rows=None):
         return self.source.cache_bytes(self.capacity if rows is None else rows)
@@ -86,7 +87,8 @@ def layout(p) -> dict:
 
 def admission(dec, s) -> dict:
     p = view(dec)
-    st, resume, cached = p._slot_for(list(s.prompt), s.draft)
+    text = getattr(s, "vision", None) is None            # image prompts reuse no kept prefix (as on one GPU)
+    st, resume, cached = p._slot_for(list(s.prompt), s.draft and text)
     source = next((st.index for ids, st, snap, _ in p.kept
                    if resume is not None and snap is resume["state"] and ids == list(s.prompt[:cached])), None)
     error = None
@@ -96,7 +98,7 @@ def admission(dec, s) -> dict:
         else:
             p._remember(list(s.prompt[:cached]), st, resume["state"], resume["tail"])
         error = f"a {len(s.prompt)}-token prompt waits for memory until a live stream finishes"
-    markers = p.points(s.prompt) if s.draft and p.points is not None else []
+    markers = p.points(s.prompt) if s.draft and text and p.points is not None else []   # image prompts keep none
     points = sorted({n for n in markers if cached + MIN_GAP <= n < entry_end(s.prompt)})
     return {**layout(p), "slot": st.index, "cached": cached, "resume_slot": source, "error": error, "points": points}
 
@@ -107,7 +109,7 @@ def round_plan(dec) -> dict:
     live = [s for s in p.streams.values() if not s.done and not s.waiting]
     solo = None
     if (p.solo_on and not ended and not p.filling and len(p.streams) == 1
-            and len(live) == 1 and live[0].draft and live[0].constraint is None):
+            and len(live) == 1 and live[0].draft and live[0].constraint is None and not live[0].st.images):
         s = live[0]
         if s.st is not p.solo.st:
             p._move_to_solo(s)