diff --git a/src/tensorfold/families/qwen4_exp/cuda/decode.py b/src/tensorfold/families/qwen4_exp/cuda/decode.py index ea765aa03e..3286a8cddf 100644 --- a/src/tensorfold/families/qwen4_exp/cuda/decode.py +++ b/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) diff --git a/src/tensorfold/families/qwen4_exp/cuda/engine.py b/src/tensorfold/families/qwen4_exp/cuda/engine.py index 002ff1fb6e..9cb91645a3 100644 --- a/src/tensorfold/families/qwen4_exp/cuda/engine.py +++ b/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,7 +163,7 @@ 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), @@ -173,7 +171,12 @@ def __init__(self, model_dir: Path, *, depth: int = DEPTH, confidence: float = C 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) diff --git a/src/tensorfold/families/qwen4_exp/cuda/image_rows.py b/src/tensorfold/families/qwen4_exp/cuda/image_rows.py index 73e8278153..2f4992c029 100644 --- a/src/tensorfold/families/qwen4_exp/cuda/image_rows.py +++ b/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: diff --git a/src/tensorfold/families/qwen4_exp/cuda/multi.py b/src/tensorfold/families/qwen4_exp/cuda/multi.py index e4326b4693..a19abc4784 100644 --- a/src/tensorfold/families/qwen4_exp/cuda/multi.py +++ b/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) diff --git a/src/tensorfold/families/qwen4_exp/cuda/multi_plan.py b/src/tensorfold/families/qwen4_exp/cuda/multi_plan.py index 0c83983898..7270a0eac5 100644 --- a/src/tensorfold/families/qwen4_exp/cuda/multi_plan.py +++ b/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) diff --git a/src/tensorfold/families/qwen4_exp/cuda/multi_tp.py b/src/tensorfold/families/qwen4_exp/cuda/multi_tp.py index 5dcc6b96f2..f9c5574c81 100644 --- a/src/tensorfold/families/qwen4_exp/cuda/multi_tp.py +++ b/src/tensorfold/families/qwen4_exp/cuda/multi_tp.py @@ -108,7 +108,8 @@ class OutOfStep(RuntimeError): def shape(dec) -> list: """The logical state both ranks must share; cache contents are checked when a prefix is selected.""" - return [dec.next_id, [[s.capacity, s.limit, s.pos, s.mtp_len, list(s.cur), s.kv_dtype] for s in dec.slots], + return [dec.next_id, [[s.capacity, s.limit, s.pos, s.mtp_len, list(s.cur), s.kv_dtype, + getattr(s, "image_positions", None) is not None] for s in dec.slots], [[s.sid, dec._index(s.st), len(s.out), s.out[-1] if s.out else None, list(s.drafts), bool(s.done), bool(s.waiting), s.background, s.draft, s.stop_eos, s.count, _pack(s.sampling)] for s in [*dec.streams.values(), *dec.filling]], @@ -168,11 +169,17 @@ def _prepare_admission(self, s, told): resume = next(({"state": k[2], "tail": k[3]} for k in self.kept if self._index(k[1]) == source and len(k[0]) == cached and s.prompt[:cached] == k[0]), None) valid = 0 <= slot < len(self.slots) and (cached == 0 or resume is not None) + # an image prompt: rank 0 holds the features it encoded, rank 1 the metadata rank 0 sent; both agree on + # the rows, the decode offset and the feature shape before the tensors cross (``_share_vision``) + vision = getattr(s, "vision", None) + images = None if vision is None else vision if isinstance(vision, dict) else { + "rows": [int(r) for r in vision.rows], "delta": int(vision.rope_delta), + "shape": [int(n) for n in vision.features.shape]} if self.link is not None: self.link.send(["admit", self.next_id, list(s.prompt), s.count, _pack(s.sampling), bool(s.draft), - bool(s.stop_eos), bool(s.background), plan]) + bool(s.stop_eos), bool(s.background), plan, images]) valid, fits = self._agree("admission", [shape(self), plan, list(s.prompt), s.count, _pack(s.sampling), - bool(s.draft), bool(s.stop_eos), bool(s.background)], + bool(s.draft), bool(s.stop_eos), bool(s.background), images], (valid, ready(self, plan))) if not valid: raise OutOfStep("the agreed prefix is not available on both ranks") @@ -180,12 +187,30 @@ def _prepare_admission(self, s, told): if not self.streams and not self.filling: raise ValueError("the prompt cannot fit available memory; shorten it or lower the server context") raise NoRoom("the proposed prompt cache growth does not fit on both ranks") + if images is not None: # both ranks reach this point or neither does + s.vision = self._share_vision(s, images) apply(self, plan) if plan["error"]: raise NoRoom(plan["error"]) s.stops = list(plan["points"]) return self.slots[slot], resume, cached + def _share_vision(self, s, images: dict): + """Rank 0 sends the image features and rotary positions it encoded; rank 1 receives the same tensors. + + Point to point over the model communicator, inside the admission's lockstep (after the agreed plan, + before any round collective). Rank 1 attaches them exactly as rank 0 does, so every layer sees identical + hidden states for the image rows; it has no tower of its own.""" + + from .vision_ranks import exchange + + leader = self.link is not None + try: + return exchange(self.w.comm, 0 if leader else 1, len(s.prompt), images, s.vision if leader else None, + hidden=int(self.w.cfg.hidden)) + except ValueError as exc: + raise OutOfStep(str(exc)) from exc + def _prepare_round(self, told): from .multi_plan import apply, ready, round_plan @@ -212,9 +237,10 @@ def follow(self, link: Link) -> None: if op is None or op[0] == "stop": return if op[0] == "admit": - sid, prompt, count, smp, drafted, stop, background, plan = op[1:] + sid, prompt, count, smp, drafted, stop, background, plan = op[1:9] s = Stream(prompt, count, _unpack(smp), draft=drafted, stop_eos=stop) s.background = background + s.vision = op[9] if len(op) > 9 else None # an image prompt's rows, delta and feature shape try: self.admit(s, told=plan) except (ValueError, NoRoom, OutOfStep) as exc: diff --git a/src/tensorfold/families/qwen4_exp/cuda/vision_ranks.py b/src/tensorfold/families/qwen4_exp/cuda/vision_ranks.py new file mode 100644 index 0000000000..5e7c4796bc --- /dev/null +++ b/src/tensorfold/families/qwen4_exp/cuda/vision_ranks.py @@ -0,0 +1,47 @@ +"""Image features on two ranks: rank 0 encodes a request's images once, rank 1 receives the same tensors. + +Both ranks must feed identical hidden states into their halves of every layer. Text does that by embedding the +same ids; images only if both ranks hold the same feature rows, so they are computed once and sent, not recomputed. +``describe`` is what rank 0 tells rank 1 with the request (JSON), ``exchange`` moves the tensors point to point once +both ranks agreed to serve it. The concurrent decoder (``multi_tp``) and the serial engine (``engine``) share this.""" + +from __future__ import annotations + +import torch + + +def describe(encoded) -> dict | None: + """The rows, decode offset and feature shape of an encoded image prompt, as the request message carries them.""" + + if encoded is None: + return None + return {"rows": [int(r) for r in encoded.rows], "delta": int(encoded.rope_delta), + "shape": [int(n) for n in encoded.features.shape]} + + +def exchange(comm, rank: int, length: int, images: dict, encoded=None, *, hidden: int): + """Rank 0 sends ``encoded``'s bf16 features and int32 positions; rank 1 receives them into tensors of the told shape. + + Both ranks validate the same description against the same prompt length, so a malformed one fails on both.""" + + from tensorfold.cuda.comm import exchange as swap + from tensorfold.vision.qwen_cuda import EncodedVision + + rows, delta, shape = images["rows"], int(images["delta"]), tuple(int(n) for n in images["shape"]) + if (len(shape) != 2 or shape[0] != len(rows) or shape[1] != int(hidden) or not rows + or sorted(set(rows)) != list(rows) or rows[0] < 0 or rows[-1] >= length): + raise ValueError("the image rows rank 0 encoded do not fit the prompt it admitted") + if rank == 0: + if encoded is None: + raise ValueError("rank 0 shares an image prompt it has not encoded") + features = encoded.features.to(dtype=torch.bfloat16).contiguous() + positions = encoded.positions.to(dtype=torch.int32).contiguous() + if tuple(features.shape) != shape or tuple(positions.shape) != (3, length): + raise ValueError("the image features rank 0 encoded do not match what it told rank 1") + swap(comm, [features, positions], [features.new_empty((0,)), positions.new_empty((0,))], 1) + return EncodedVision(tuple(rows), features, positions, delta) + device = torch.device("cuda", 0) + features = torch.empty(shape, dtype=torch.bfloat16, device=device) + positions = torch.empty((3, length), dtype=torch.int32, device=device) + swap(comm, [features.new_empty((0,)), positions.new_empty((0,))], [features, positions], 0) + return EncodedVision(tuple(rows), features, positions, delta) diff --git a/tests/test_flashnext_vision_two_ranks_host.py b/tests/test_flashnext_vision_two_ranks_host.py new file mode 100644 index 0000000000..9d14ebf005 --- /dev/null +++ b/tests/test_flashnext_vision_two_ranks_host.py @@ -0,0 +1,269 @@ +"""Image prompts on two ranks: rank 0 encodes, the admission carries the rows, rank 1 attaches the same features.""" + +import ast +import hashlib +import json +from pathlib import Path +import threading +from types import SimpleNamespace +from typing import Callable + +import pytest + + +def _methods(): + root = Path(__file__).resolve().parents[1] / 'src/tensorfold' + module = ast.parse((root / 'families/qwen4_exp/cuda/multi_tp.py').read_text()) + owner = next(node for node in module.body if isinstance(node, ast.ClassDef) and node.name == 'TwoRanks') + methods = [node for node in owner.body if isinstance(node, ast.FunctionDef) + and node.name in ('_agree', '_prepare_admission')] + for method in methods: + method.body = [node for node in method.body if not isinstance(node, ast.ImportFrom)] + body = ast.ClassDef(name='Ranks', bases=[], keywords=[], body=methods, decorator_list=[], type_params=[]) + capacity = ast.parse((root / 'cuda/capacity.py').read_text()) + gather = next(node for node in capacity.body if isinstance(node, ast.FunctionDef) and node.name == 'gather_ints') + return ast.fix_missing_locations(ast.Module(body=[gather, body], type_ignores=[])) + + +def _ranks(visions, expected): + """Run rank 0 (with ``visions[0]``) and rank 1 (with ``visions[1]``) through one admission; the outcomes.""" + + barrier = threading.Barrier(2) + values, results, sent, shared = {}, [None, None], [], {} + plan = {'slot': 0, 'cached': 0, 'resume_slot': None, 'error': None, 'points': [], 'need': 7} + + class OutOfStep(RuntimeError): + pass + + class NoRoom(RuntimeError): + pass + + class Tensor: + def __init__(self, data): + self.data, self.rows = list(data), None + + def view(self, rows, unused): + self.rows = rows + return self + + def tolist(self): + width = len(self.data) // self.rows + return [self.data[start:start + width] for start in range(0, len(self.data), width)] + + torch = SimpleNamespace(int64=None, tensor=lambda data, **kwargs: Tensor(data), + empty=lambda shape, **kwargs: Tensor([0] * shape[0])) + namespace = {'hashlib': hashlib, 'json': json, 'torch': torch, 'Callable': Callable, + 'OutOfStep': OutOfStep, 'NoRoom': NoRoom, + 'shape': lambda dec: ['shared-state'], '_pack': lambda sampling: None, + 'admission': lambda dec, stream: dict(plan), + 'ready': lambda dec, proposed: True, + 'apply': lambda dec, proposed: setattr(dec, 'applied', True)} + + def gather(rank, send, receive): + values[rank] = list(send.data) + barrier.wait(timeout=2) + receive.data = [*values[0], *values[1]] + barrier.wait(timeout=2) + exec(compile(_methods(), 'actual-tp-planning-methods', 'exec'), namespace) + + def run(rank): + decoder = namespace['Ranks']() + decoder.w = SimpleNamespace(comm=SimpleNamespace(all_gather=lambda send, receive: gather(rank, send, receive))) + decoder.next_id, decoder.link = 0, None + if rank == 0: + decoder.link = SimpleNamespace(send=lambda op: sent.append(op)) + decoder._share_vision = lambda s, images: shared.setdefault(rank, ('features', images)) + decoder.slots, decoder.kept, decoder.streams, decoder.filling = [object()], [], {}, [] + stream = SimpleNamespace(prompt=[1, 2, 3, 4], count=4, sampling=None, draft=True, + stop_eos=False, background=False, vision=visions[rank]) + try: + decoder._prepare_admission(stream, None if rank == 0 else dict(plan)) + results[rank] = ('admit', stream.vision) + except Exception as error: + results[rank] = (type(error).__name__, None) + + threads = [threading.Thread(target=run, args=(rank,)) for rank in (0, 1)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(5) + assert not thread.is_alive() + assert [r[0] for r in results] == [expected, expected] + return results, sent, shared + + +def test_text_admissions_carry_no_image_metadata_and_share_nothing(): + results, sent, shared = _ranks([None, None], 'admit') + assert sent[0][0] == 'admit' and sent[0][-1] is None and len(sent[0]) == 10 + assert shared == {} and [r[1] for r in results] == [None, None] + + +def test_rank_0_sends_the_rows_it_encoded_and_both_ranks_share_the_same_tensors(): + encoded = SimpleNamespace(rows=(1, 2), rope_delta=-1, features=SimpleNamespace(shape=(2, 2560))) + meta = {'rows': [1, 2], 'delta': -1, 'shape': [2, 2560]} + results, sent, shared = _ranks([encoded, json.loads(json.dumps(meta))], 'admit') + assert sent[0][-1] == meta # what the follower is told, JSON-exact + assert shared == {0: ('features', meta), 1: ('features', meta)} + assert [r[1] for r in results] == [('features', meta)] * 2 # the stream now carries the shared features + + +def test_an_image_admission_one_rank_does_not_know_about_falls_out_of_step(): + encoded = SimpleNamespace(rows=(1, 2), rope_delta=-1, features=SimpleNamespace(shape=(2, 2560))) + _, _, shared = _ranks([encoded, None], 'OutOfStep') + assert shared == {} # the digest differs before any tensor crosses + + +def test_begin_attaches_features_another_rank_encoded_without_a_tower(): + torch = pytest.importorskip('torch') + from tensorfold.families.qwen4_exp.cuda import image_rows + + class State: + def __init__(self): + self.image_positions = self.image_rows = self.image_features = None + self.rope_delta = 0 + + def set_rope_delta(self, delta): + self.rope_delta = delta + + st = State() + encoded = SimpleNamespace(rows=(1, 2), features=torch.zeros((2, 8)), rope_delta=-1, + positions=torch.zeros((3, 4), dtype=torch.int32)) + stream = SimpleNamespace(prompt=[5, 6, 7, 8], vision=encoded) + image_rows.begin(SimpleNamespace(st=st), stream, None) + assert stream.vision is None and st.image_rows == (1, 2) and st.rope_delta == -1 + assert st.image_features is encoded.features and tuple(st.image_positions.shape) == (4, 3) + with pytest.raises(ValueError, match='--vision'): # a prepared prompt still needs the tower + image_rows.begin(SimpleNamespace(st=State()), SimpleNamespace(prompt=[5], vision=object()), None) + + +@pytest.fixture +def planning(): + pytest.importorskip('torch') + from tests.test_cuda_geometry import allocations as geometry # noqa: F401 (patched allocation sizes) + import tests.test_flashnext_round_plan as plans + from tensorfold.families.qwen4_exp.cuda.multi import MultiDecoder + from tensorfold.families.qwen4_exp.cuda.multi_plan import admission, round_plan + + plans.MultiDecoder = MultiDecoder # what its autouse fixture would set + return plans.decoder, admission, round_plan + + +def test_image_prompts_reuse_no_prefix_keep_no_points_and_stay_off_the_graph_slot(planning): + from tensorfold.cuda.streams import Stream + + decoder, admission, round_plan = planning + d = decoder() + d.points = lambda prompt: [8, 16] + kept_prompt = [7] * 30 + d.slots[0].pos, d.slots[0].mtp_len = 30, 29 + d.free = d.slots[1:] + d.kept = [(kept_prompt, d.slots[0], {'mtp_len': 29}, None)] + text = admission(d, Stream(kept_prompt + [9] * 10, 8)) + assert text['cached'] == 30 and text['points'] == [] # the text prompt resumes the kept prefix + image = admission(d, Stream(kept_prompt + [9] * 10, 8, vision={'rows': [3], 'delta': 0, 'shape': [1, 8]})) + assert image['cached'] == 0 and image['slot'] != 0 and image['points'] == [] + d = decoder() + live = Stream([7], 12, sid=0, st=d.slots[0], out=[3], drafts=[4]) + live.st.pos, live.st.image_positions = 255, object() + d.streams, d.free = {0: live}, d.slots[1:] + d.solo, d.solo_on = SimpleNamespace(st=d.slots[1]), True + assert round_plan(d)['solo'] is None # the graph slot's captures know no image rows + live.st.image_positions = None + assert round_plan(d)['solo'] == 0 + + +# --- the serial engine (no --parallel): image prompts on one GPU and on two ranks --- + +def test_serial_request_message_carries_the_image_description_and_keeps_no_points(): + pytest.importorskip('torch') + from tensorfold.families.qwen4_exp.cuda.engine import FlashNextEngine + + class Store: + def __init__(self): + self.kv = {} + + def set(self, key, value): + self.kv[key] = value + + e = object.__new__(FlashNextEngine) + e.comm, e.served, e.points = SimpleNamespace(store=Store()), 3, lambda prompt: [2, 4] + text = e._share([1, 2, 3, 4, 5], 8, None, True, 0) + assert text[7] is None and text[-1] == [2, 4] and len(text) == 9 # the points stay the last field + images = {'rows': [1, 2], 'delta': -1, 'shape': [2, 2560]} + shared = e._share([1, 2, 3, 4, 5], 8, None, True, 0, images=images) + assert shared[7] == images and shared[-1] == [] # an image prompt keeps no prompt states + assert json.loads(e.comm.store.kv['tensorfold/flashnext/request/3'])['images'] == images + assert FlashNextEngine._unpack(json.dumps({'prompt': [1], 'max_tokens': 1, 'draft': True, 'cached': 0, + 'sampling': None}))[7] is None # an older leader's message + follower = object.__new__(FlashNextEngine) + follower.served, follower.cache, follower.multi, follower.points = 0, [], None, None + requests = iter([([1, 2, 3], 4, None, True, 0, [], True, [2]), # eight fields: no images + ([1, 2, 3], 4, None, True, 0, [], True, None, [2]), None]) # nine: images before points + follower._receive = lambda: next(requests) + seen = [] + follower._decode = lambda *args, **kwargs: seen.append((kwargs['points'], kwargs['vision'])) + follower.follow() + assert seen == [([2], None), ([2], None)] and follower.served == 2 + + +def test_describe_and_exchange_validate_the_same_description_on_both_ranks(): + torch = pytest.importorskip('torch') + from tensorfold.families.qwen4_exp.cuda import vision_ranks + + encoded = SimpleNamespace(rows=(1, 2), rope_delta=-1, features=torch.zeros((2, 16)), + positions=torch.zeros((3, 4), dtype=torch.int32)) + images = vision_ranks.describe(encoded) + assert images == {'rows': [1, 2], 'delta': -1, 'shape': [2, 16]} and vision_ranks.describe(None) is None + for bad in ({**images, 'rows': [2, 1]}, {**images, 'rows': [1, 4]}, {**images, 'shape': [3, 16]}, + {**images, 'shape': [2, 8]}): + with pytest.raises(ValueError, match='do not fit'): + vision_ranks.exchange(None, 1, 4, bad, hidden=16) + with pytest.raises(ValueError, match='has not encoded'): + vision_ranks.exchange(None, 0, 4, images, None, hidden=16) + moved = [] + comm = SimpleNamespace(rank=0, world=2, exchange=lambda sends, recvs, peer: moved.append((len(sends), peer))) + out = vision_ranks.exchange(comm, 0, 4, images, encoded, hidden=16) + assert moved == [(2, 1)] and out.rows == (1, 2) and out.rope_delta == -1 + assert out.features.dtype == torch.bfloat16 and tuple(out.positions.shape) == (3, 4) + + +def test_serial_engine_decodes_an_image_state_eagerly_and_keeps_no_snapshot(monkeypatch): + pytest.importorskip('torch') + from tensorfold.families.qwen4_exp.cuda import decode, engine as engine_module + from tensorfold.families.qwen4_exp.cuda.engine import FlashNextEngine + + calls = [] + eng = object.__new__(decode.Engine) + eng.w, eng.st, eng.buf, eng.mbuf = object(), SimpleNamespace(image_positions=None), object(), object() + eng.graphs = SimpleNamespace(forward=lambda tokens: calls.append('graph'), + mtp_forward=lambda nxt, streams: calls.append('graph-mtp')) + monkeypatch.setattr(decode, 'forward', lambda w, st, b, tokens: calls.append('eager')) + monkeypatch.setattr(decode, 'mtp_forward', lambda w, st, b, nxt, streams: calls.append('eager-mtp')) + eng.forward([1]) + eng.mtp_forward([1], None) + eng.st.image_positions = object() + eng.forward([1]) + eng.mtp_forward([1], None) + assert calls == ['graph', 'graph-mtp', 'eager', 'eager-mtp'] + + seen = {} + + def prefill(e, prompt, sampling, **kwargs): + seen.update(kwargs) + return 7 + + monkeypatch.setattr(decode, 'prefill', prefill) + monkeypatch.setattr(engine_module.time, 'perf_counter', lambda: 0.0) + import torch + monkeypatch.setattr(torch.cuda, 'synchronize', lambda: None) + fe = object.__new__(FlashNextEngine) + fe.e, fe.cache, fe.eos, fe.depth, fe.points = SimpleNamespace(kept={'state': 's', 'tail': None}), [([1], {})], (), 3, None + kept = ([1], {'state': 'old'}) + fe.cache = [kept] + stats = fe._decode([1, 2, 3], 1, None, lambda new: True, kept, vision=SimpleNamespace(rows=(1,))) + assert seen == {'constraint': None, 'probabilities': None, 'vision': seen['vision']} + assert fe.cache == [] and stats['cached'] == 0 # from the start, nothing kept, nothing resumed + seen.clear() + fe.cache = [kept] + fe._decode([1, 2, 3], 1, None, lambda new: True, kept) + assert 'vision' not in seen and seen['resume'] == {'state': 'old'} and fe.cache[-1][0] == [1, 2]