@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/input/media-fetch.ts

@@ -0,0 +1,439 @@
+// Remote media (image_url / audio_url) fetching with an SSRF/resource-
+// exhaustion policy. A chat request's URL is attacker-controlled input to
+// a server that may hold credentials for its LAN (the classic target:
+// cloud metadata at 169.254.169.254, router admin pages, other local
+// services), so by default the destination of every hop — the original
+// URL and each redirect — must be a public http(s) host, the whole fetch
+// is wall-clock bounded, and the response is size-capped.
+//
+// Callers may allow private/loopback/link-local destinations with the policy's
+// `allowPrivate` option. Scheme/timeout/size limits still apply.
+//
+// DNS names are resolved and every resolved address is checked before the
+// fetch. The request then dials one of those checked IP literals while
+// preserving the logical Host header and HTTPS SNI, so the transport cannot
+// re-resolve to a private address between validation and connect.
+//
+// This module must stay free of mlx imports so packages/inference/tests/input/media-fetch.test.ts
+// runs model-free.
+
+import { lookup } from "node:dns/promises";
+import { runtimeValue } from "../runtime/config";
+
+export interface MediaFetchPolicy {
+  /** Permit private/loopback/link-local destinations (LAN hosts). */
+  allowPrivate: boolean;
+  /** Optional caller-specific explanation when private destinations fail. */
+  privateDestinationHint?: string;
+  /** Wall-clock budget for the whole fetch — every redirect hop plus the
+   *  body read share one AbortSignal.timeout. */
+  timeoutMs: number;
+  /** Response body cap in bytes (checked against Content-Length up front
+   *  AND while streaming, since the header is optional/spoofable). */
+  maxBytes: number;
+  maxRedirects: number;
+}
+
+export type RestrictedFetch = (
+  input: string | URL,
+  init?: BunFetchRequestInit,
+) => Promise<Response>;
+
+export type ResolveHost = (
+  hostname: string,
+) => Promise<ReadonlyArray<{ address: string }>>;
+
+export interface RestrictedFetchDependencies {
+  fetch: RestrictedFetch;
+  resolve: ResolveHost;
+}
+
+export interface RestrictedFetchOptions {
+  /** Optional caller cancellation, combined with the policy timeout. */
+  signal?: AbortSignal;
+  /** Request headers applied on every validated redirect hop. */
+  headers?: Record<string, string>;
+  /** Test seam for deterministic DNS/redirect/body behavior. */
+  dependencies?: Partial<RestrictedFetchDependencies>;
+  /** Media callers reject non-2xx before reading; web_fetch returns the body. */
+  rejectHttpErrors?: boolean;
+}
+
+export interface RestrictedFetchResult {
+  ok: boolean;
+  status: number;
+  finalUrl: string;
+  contentType: string;
+  bytes: Uint8Array;
+}
+
+export function defaultMediaFetchPolicy(): MediaFetchPolicy {
+  return {
+    allowPrivate: runtimeValue("MLX_BUN_ALLOW_PRIVATE_MEDIA") === "1",
+    // 10 s / 64 MB: well above what a 30 s audio clip (~5 MB of 16 kHz
+    // WAV, less compressed) or any supported image needs.
+    timeoutMs: 10_000,
+    maxBytes: 64 * 1024 * 1024,
+    maxRedirects: 5,
+  };
+}
+
+const PRIVATE_HINT =
+  "private/loopback hosts are blocked by default; set MediaFetchPolicy.allowPrivate to true to fetch from LAN hosts";
+
+function privateHint(policy: MediaFetchPolicy): string {
+  return policy.privateDestinationHint ?? PRIVATE_HINT;
+}
+
+function parseIpv4(s: string): number[] | null {
+  const parts = s.split(".");
+  if (parts.length !== 4) return null;
+  const out: number[] = [];
+  for (const p of parts) {
+    if (!/^\d{1,3}$/.test(p)) return null;
+    const n = Number(p);
+    if (n > 255) return null;
+    out.push(n);
+  }
+  return out;
+}
+
+function blockedIpv4(o: number[]): boolean {
+  const [a, b] = o as [number, number];
+  if (a === 0 || a === 10 || a === 127) return true; // this-net, 10/8, loopback
+  if (a === 100 && b >= 64 && b <= 127) return true; // CGNAT 100.64/10
+  if (a === 169 && b === 254) return true; // link-local (cloud metadata lives here)
+  if (a === 172 && b >= 16 && b <= 31) return true; // 172.16/12
+  if (a === 192 && b === 168) return true; // 192.168/16
+  if (a >= 224) return true; // multicast + reserved + broadcast
+  return false;
+}
+
+/** Expand an IPv6 literal (with `::` compression, an optional zone index,
+ *  and an optional dotted-quad tail like ::ffff:127.0.0.1) into its 8
+ *  16-bit groups. null = not valid IPv6. */
+function parseIpv6(addr: string): number[] | null {
+  let s = addr;
+  const pct = s.indexOf("%");
+  if (pct !== -1) s = s.slice(0, pct);
+  if (s.includes(".")) {
+    // rewrite the dotted-quad tail as two hex groups
+    const i = s.lastIndexOf(":");
+    const v4 = parseIpv4(s.slice(i + 1));
+    if (!v4) return null;
+    s =
+      s.slice(0, i + 1) +
+      (((v4[0]! << 8) | v4[1]!).toString(16)) + ":" +
+      (((v4[2]! << 8) | v4[3]!).toString(16));
+  }
+  const halves = s.split("::");
+  if (halves.length > 2) return null;
+  const toGroups = (part: string): number[] | null => {
+    if (part === "") return [];
+    const out: number[] = [];
+    for (const g of part.split(":")) {
+      if (!/^[0-9a-fA-F]{1,4}$/.test(g)) return null;
+      out.push(parseInt(g, 16));
+    }
+    return out;
+  };
+  const head = toGroups(halves[0]!);
+  if (head === null) return null;
+  if (halves.length === 1) return head.length === 8 ? head : null;
+  const tail = toGroups(halves[1]!);
+  if (tail === null) return null;
+  if (head.length + tail.length > 7) return null;
+  return [...head, ...new Array(8 - head.length - tail.length).fill(0), ...tail];
+}
+
+function blockedIpv6(g: number[]): boolean {
+  if (g.every((x) => x === 0)) return true; // :: unspecified
+  if (g.slice(0, 7).every((x) => x === 0) && g[7] === 1) return true; // ::1
+  if ((g[0]! & 0xffc0) === 0xfe80) return true; // link-local fe80::/10
+  if ((g[0]! & 0xfe00) === 0xfc00) return true; // ULA fc00::/7
+  if ((g[0]! & 0xff00) === 0xff00) return true; // multicast ff00::/8
+  // v4-mapped (::ffff:0:0/96) and NAT64 (64:ff9b::/96) embed an IPv4
+  // address in the low 32 bits — judge the embedded address.
+  const v4Embedded =
+    (g.slice(0, 5).every((x) => x === 0) && g[5] === 0xffff) ||
+    (g[0] === 0x64 && g[1] === 0xff9b && g.slice(2, 6).every((x) => x === 0));
+  if (v4Embedded)
+    return blockedIpv4([g[6]! >> 8, g[6]! & 0xff, g[7]! >> 8, g[7]! & 0xff]);
+  return false;
+}
+
+/** True when `addr` is a private/loopback/link-local/reserved IP literal
+ *  — or unparseable as an IP at all (fail closed; callers only pass
+ *  things believed to be addresses). Accepts bracketed IPv6. */
+export function isBlockedAddress(addr: string): boolean {
+  let s = addr;
+  if (s.startsWith("[") && s.endsWith("]")) s = s.slice(1, -1);
+  const v4 = parseIpv4(s);
+  if (v4) return blockedIpv4(v4);
+  const v6 = parseIpv6(s);
+  if (v6) return blockedIpv6(v6);
+  return true;
+}
+
+/** Sync per-hop destination check: scheme + hostname-literal policy.
+ *  Returns a rejection reason or null. DNS names pass here and get their
+ *  resolved addresses checked in fetchMediaBytes (async). Exported pure
+ *  so the policy is testable without any network. */
+export function checkMediaUrl(u: URL, policy: MediaFetchPolicy): string | null {
+  if (u.protocol !== "http:" && u.protocol !== "https:")
+    return `unsupported url scheme "${u.protocol.replace(/:$/, "")}" (only http/https)`;
+  const host =
+    u.hostname.startsWith("[") && u.hostname.endsWith("]")
+      ? u.hostname.slice(1, -1)
+      : u.hostname;
+  if (!host) return "missing host";
+  if (policy.allowPrivate) return null;
+  const lower = host.toLowerCase();
+  if (lower === "localhost" || lower.endsWith(".localhost"))
+    return `destination "${host}" — ${privateHint(policy)}`;
+  // WHATWG URL parsing already canonicalized numeric IPv4 forms
+  // (http://2130706433/, http://0x7f000001/ → 127.0.0.1).
+  if (parseIpv4(lower) || parseIpv6(lower)) {
+    if (isBlockedAddress(lower))
+      return `destination "${host}" — ${privateHint(policy)}`;
+  }
+  return null;
+}
+
+export async function assertResolvesPublic(
+  hostname: string,
+  kind: string,
+  policy: MediaFetchPolicy,
+  resolve: ResolveHost = async (host) =>
+    await lookup(host, { all: true, verbatim: true }),
+): Promise<void> {
+  if (policy.allowPrivate) return;
+  const host =
+    hostname.startsWith("[") && hostname.endsWith("]")
+      ? hostname.slice(1, -1)
+      : hostname;
+  if (parseIpv4(host) || parseIpv6(host)) return; // literal — checked sync
+  let addrs: { address: string }[];
+  try {
+    addrs = [...await resolve(host)];
+  } catch {
+    throw new Error(`${kind} fetch failed: could not resolve host "${host}"`);
+  }
+  for (const a of addrs) {
+    if (isBlockedAddress(a.address))
+      throw new Error(
+        `${kind} url rejected: host "${host}" resolves to ${a.address} — ${privateHint(policy)}`,
+      );
+  }
+}
+
+async function resolveFetchAddresses(
+  hostname: string,
+  kind: string,
+  policy: MediaFetchPolicy,
+  resolve: ResolveHost,
+): Promise<string[]> {
+  const host =
+    hostname.startsWith("[") && hostname.endsWith("]")
+      ? hostname.slice(1, -1)
+      : hostname;
+  if (parseIpv4(host) || parseIpv6(host)) return [host];
+  if (policy.allowPrivate) return [];
+  let addrs: { address: string }[];
+  try {
+    addrs = [...await resolve(host)];
+  } catch {
+    throw new Error(`${kind} fetch failed: could not resolve host "${host}"`);
+  }
+  if (addrs.length === 0)
+    throw new Error(`${kind} fetch failed: host "${host}" resolved to no addresses`);
+  for (const a of addrs) {
+    if (isBlockedAddress(a.address))
+      throw new Error(
+        `${kind} url rejected: host "${host}" resolves to ${a.address} — ${privateHint(policy)}`,
+      );
+  }
+  return addrs.map((a) => a.address);
+}
+
+/**
+ * Fetch an http(s) response through the shared public-destination policy.
+ * Every redirect is handled manually and revalidated, caller cancellation
+ * shares one signal with the wall-clock timeout, and the body is streamed
+ * into a byte-capped buffer.
+ */
+export async function fetchRestrictedHttpBytes(
+  url: string,
+  kind: string,
+  policy: MediaFetchPolicy,
+  options: RestrictedFetchOptions = {},
+): Promise<RestrictedFetchResult> {
+  let current: URL;
+  try {
+    current = new URL(url);
+  } catch {
+    throw new Error(`unsupported ${kind} url scheme: ${url.slice(0, 16)}`);
+  }
+
+  const fetchImpl = options.dependencies?.fetch ??
+    ((input: string | URL, init?: BunFetchRequestInit) => fetch(input, init));
+  const resolve = options.dependencies?.resolve ??
+    (async (host: string) => await lookup(host, { all: true, verbatim: true }));
+  const timeout = AbortSignal.timeout(policy.timeoutMs);
+  const signal = options.signal
+    ? AbortSignal.any([options.signal, timeout])
+    : timeout;
+  const capMb = Math.round(policy.maxBytes / (1024 * 1024));
+  let overCap = false;
+
+  try {
+    for (let hop = 0; ; hop++) {
+      const reason = checkMediaUrl(current, policy);
+      if (reason) throw new Error(`${kind} url rejected: ${reason}`);
+      const addresses = await resolveFetchAddresses(
+        current.hostname,
+        kind,
+        policy,
+        resolve,
+      );
+      const logicalHostname = current.hostname.replace(/^\[|\]$/g, "");
+      const connectUrl = new URL(current);
+      if (!policy.allowPrivate && addresses.length > 0) {
+        const address = addresses[0]!;
+        connectUrl.hostname = address.includes(":") ? `[${address}]` : address;
+      }
+      const headers = new Headers(options.headers);
+      if (connectUrl.hostname !== current.hostname && !headers.has("host"))
+        headers.set("host", current.host);
+      const init: BunFetchRequestInit = {
+        redirect: "manual",
+        signal,
+        headers,
+      };
+      if (current.protocol === "https:" && connectUrl.hostname !== current.hostname) {
+        init.tls = { serverName: logicalHostname };
+      }
+      const res = await fetchImpl(connectUrl, init);
+      if (res.status >= 300 && res.status < 400) {
+        const loc = res.headers.get("location");
+        if (!loc)
+          throw new Error(`${kind} fetch failed: ${res.status} redirect without location`);
+        if (hop >= policy.maxRedirects)
+          throw new Error(`${kind} fetch failed: more than ${policy.maxRedirects} redirects`);
+        try {
+          await res.body?.cancel();
+        } catch {}
+        try {
+          current = new URL(loc, current);
+        } catch {
+          throw new Error(`${kind} fetch failed: malformed redirect location`);
+        }
+        continue;
+      }
+
+      if (options.rejectHttpErrors && !res.ok) {
+        try {
+          await res.body?.cancel();
+        } catch {}
+        throw new Error(`${kind} fetch failed: ${res.status} ${current}`);
+      }
+
+      const declared = Number(res.headers.get("content-length") ?? "0");
+      if (Number.isFinite(declared) && declared > policy.maxBytes) {
+        overCap = true;
+        try {
+          await res.body?.cancel();
+        } catch {}
+        throw new Error(`${kind} response exceeds the ${capMb} MB limit`);
+      }
+
+      let bytes: Uint8Array;
+      if (!res.body) {
+        bytes = new Uint8Array(await res.arrayBuffer());
+        if (bytes.byteLength > policy.maxBytes) {
+          overCap = true;
+          throw new Error(`${kind} response exceeds the ${capMb} MB limit`);
+        }
+      } else {
+        const reader = res.body.getReader();
+        const chunks: Uint8Array[] = [];
+        let total = 0;
+        for (;;) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          total += value.byteLength;
+          if (total > policy.maxBytes) {
+            overCap = true;
+            try {
+              await reader.cancel();
+            } catch {}
+            throw new Error(`${kind} response exceeds the ${capMb} MB limit`);
+          }
+          chunks.push(value);
+        }
+        bytes = new Uint8Array(total);
+        let offset = 0;
+        for (const chunk of chunks) {
+          bytes.set(chunk, offset);
+          offset += chunk.byteLength;
+        }
+      }
+
+      return {
+        ok: res.ok,
+        status: res.status,
+        finalUrl: current.href,
+        contentType: res.headers.get("content-type") ?? "",
+        bytes,
+      };
+    }
+  } catch (error) {
+    if (timeout.aborted && !options.signal?.aborted && !overCap)
+      throw new Error(
+        `${kind} fetch timed out after ${policy.timeoutMs} ms: ${url.slice(0, 200)}`,
+      );
+    throw error;
+  }
+}
+
+/** Fetch image/audio bytes for an OpenAI-style image_url/audio_url part.
+ *  data: URLs decode locally (no policy — the bytes came in the request);
+ *  http(s) URLs go through the destination policy above, with every
+ *  redirect hop re-validated, one wall-clock timeout, and a streaming
+ *  size cap. Throws Error with a client-presentable message — callers
+ *  surface it as a 400 (the prompt-build catch in
+ *  apps/mlx-bun/src/server/chat-stage.ts). */
+/** Video clips are legitimately larger than images/audio: same SSRF guard
+ *  and timeout discipline, quadruple the body cap (a 30 s 1080p H.264 clip
+ *  runs tens of MB; frame sampling truncates long clips server-side). */
+export function videoMediaFetchPolicy(): MediaFetchPolicy {
+  return { ...defaultMediaFetchPolicy(), maxBytes: 256 * 1024 * 1024 };
+}
+
+export async function fetchMediaBytes(
+  url: string,
+  kind: "image" | "audio" | "video",
+  policy: MediaFetchPolicy = kind === "video"
+    ? videoMediaFetchPolicy()
+    : defaultMediaFetchPolicy(),
+): Promise<Uint8Array> {
+  if (url.startsWith("data:")) {
+    const comma = url.indexOf(",");
+    if (comma === -1) throw new Error("malformed data: URL");
+    const meta = url.slice(0, comma);
+    const body = url.slice(comma + 1);
+    if (!meta.includes("base64")) throw new Error("data: URL must be base64");
+    // The same body cap as the HTTP path — inline payloads must not become
+    // the uncapped route (base64 inflates 4/3, so check pre-decode too).
+    if (body.length * 0.75 > policy.maxBytes)
+      throw new Error(
+        `${kind} data: URL exceeds the ${Math.round(policy.maxBytes / 1024 / 1024)} MB cap`,
+      );
+    return Uint8Array.from(Buffer.from(body, "base64"));
+  }
+  return (
+    await fetchRestrictedHttpBytes(url, kind, policy, {
+      rejectHttpErrors: true,
+    })
+  ).bytes;
+}

package/src/input/packed-mixed.ts

@@ -0,0 +1,43 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../contracts/mlx/cache";
+import type { MlxPhaseGroup } from "../contracts/mlx/graph";
+import { mapPackedTokens } from "./token-groups";
+
+/** Group geometry and ownership for written packed operations. The graph
+ * supplies each layer's direct calls; this scope selects no numerical operation. */
+export class PackedMixedScope {
+  readonly #hidden:MlxArray[]=[];
+  readonly #bounded:boolean;
+  constructor(readonly groups:readonly MlxPhaseGroup[],embed:(ids:MlxArray)=>MlxArray) {
+    this.#bounded=groups.some(group=>group.phase==="window" && group.rows>=5);
+    try { for (const group of groups) this.#hidden.push(embed(group.ids)); }
+    catch (error) { this[Symbol.dispose](); throw error; }
+  }
+  accept(layer:number,attention:(hidden:MlxArray,cache:Cache,phase:"decode"|"window")=>MlxArray,
+    mlp:(hidden:MlxArray)=>MlxArray):void {
+    const mids:MlxArray[]=[];
+    try {
+      for (const [row,group] of this.groups.entries()) mids.push(attention(this.#hidden[row]!,group.state[layer]!,group.phase==="decode"?"decode":"window"));
+      const outputs=mapPackedTokens(mids,mlp);
+      for (const [row,output] of outputs.entries()) { this.#hidden[row]!.dispose(); this.#hidden[row]=output; }
+      if (this.#bounded) {
+        const caches=this.groups.map(group=>group.state[layer]!);
+        const state=caches.map(cache=>cache.state());
+        try { ops.evalAll([...this.#hidden,...state.flat()]); }
+        finally { for (const [row,cache] of caches.entries()) if (cache.stateNeedsDispose) for (const array of state[row]!) array.dispose(); }
+      }
+      for (const [row,group] of this.groups.entries()) group.capture?.(layer,this.#hidden[row]!);
+    } finally { for (const mid of mids) mid.dispose(); }
+  }
+  finish(layer:number,norm:(hidden:MlxArray)=>MlxArray):MlxArray[] {
+    const result:MlxArray[]=[];
+    try {
+      for (const [row,hidden] of this.#hidden.entries()) {
+        const output=norm(hidden); result.push(output); this.groups[row]!.capture?.(layer,output);
+      }
+      return result;
+    } catch (error) { for (const output of result) output.dispose(); throw error; }
+  }
+  [Symbol.dispose]():void { for (const hidden of this.#hidden.splice(0)) hidden.dispose(); }
+}

package/src/input/special-tokens.ts

@@ -0,0 +1,75 @@
+// Special tokens a model's generated text or media prompts depend on. The model
+// profile declares them as token TEXT; the tokenizer in use resolves the ids, so
+// no consumer carries a family's vocabulary ids.
+
+import type { AudioTokenIds, VisionTokenIds } from "./vision/prompt";
+import type { LoadedTokenizer } from "./tokenizer";
+
+type Encoder = Pick<LoadedTokenizer, "encode">;
+
+/** The id of `text` when the tokenizer encodes it as exactly one token. */
+export function singleTokenId(tokenizer: Encoder, text: string): number | null {
+  const ids = tokenizer.encode(text, false);
+  return ids.length === 1 ? ids[0]! : null;
+}
+
+/** Markers that delimit tool calls and the reasoning channel in generated
+ * text. They are special tokens the content decoder strips, so a router must
+ * see them as ids. */
+export interface SentinelTokenTexts {
+  readonly toolCallStart: string;
+  readonly toolCallEnd: string;
+  readonly channelStart: string;
+  readonly channelEnd: string;
+}
+export interface SentinelTokens {
+  readonly toolCallStart: number;
+  readonly toolCallEnd: number;
+  readonly channelStart: number;
+  readonly channelEnd: number;
+}
+
+/** Resolve declared sentinel texts against a tokenizer. A declared marker that
+ * is not one token there means the declaration does not fit this tokenizer. */
+export function resolveSentinelTokens(tokenizer: Encoder, texts: SentinelTokenTexts): SentinelTokens {
+  const id = (text: string): number => {
+    const found = singleTokenId(tokenizer, text);
+    if (found === null) throw new Error(`declared sentinel ${JSON.stringify(text)} is not a single token in this tokenizer`);
+    return found;
+  };
+  return { toolCallStart: id(texts.toolCallStart), toolCallEnd: id(texts.toolCallEnd),
+    channelStart: id(texts.channelStart), channelEnd: id(texts.channelEnd) };
+}
+
+/** Soft-token markers of a graph's tower prompts, by role. */
+export interface MediaTokenTexts {
+  readonly vision: { readonly image: string; readonly begin: string; readonly end: string };
+  readonly audio: { readonly audio: string; readonly begin: string; readonly end: string };
+}
+
+const configId = (raw: Record<string, unknown>, key: string): number | undefined =>
+  typeof raw[key] === "number" ? raw[key] as number : undefined;
+
+/** Image soft-token ids: the checkpoint's config wins, the declared token text
+ * fills what it omits. Null when neither supplies all three. */
+export function resolveVisionTokenIds(
+  raw: Record<string, unknown>, tokenizer: Encoder, texts: MediaTokenTexts["vision"] | null,
+): VisionTokenIds | null {
+  const from = (key: string, text: string | undefined) =>
+    configId(raw, key) ?? (text === undefined ? null : singleTokenId(tokenizer, text));
+  const imageTokenId = from("image_token_id", texts?.image), boiTokenId = from("boi_token_id", texts?.begin),
+    eoiTokenId = from("eoi_token_id", texts?.end);
+  return imageTokenId == null || boiTokenId == null || eoiTokenId == null ? null : { imageTokenId, boiTokenId, eoiTokenId };
+}
+
+/** Audio soft-token ids, for a checkpoint that declares `audio_config`; same precedence as vision. */
+export function resolveAudioTokenIds(
+  raw: Record<string, unknown>, tokenizer: Encoder, texts: MediaTokenTexts["audio"] | null,
+): AudioTokenIds | null {
+  if (!raw.audio_config) return null;
+  const from = (key: string, text: string | undefined) =>
+    configId(raw, key) ?? (text === undefined ? null : singleTokenId(tokenizer, text));
+  const audioTokenId = from("audio_token_id", texts?.audio), boaTokenId = from("boa_token_id", texts?.begin),
+    eoaTokenId = from("eoa_token_id", texts?.end);
+  return audioTokenId == null || boaTokenId == null || eoaTokenId == null ? null : { audioTokenId, boaTokenId, eoaTokenId };
+}

package/src/input/token-groups.ts

@@ -0,0 +1,49 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { MlxPhaseGroup } from "../contracts/mlx/graph";
+
+/** Apply a tokenwise block over the groups of a mixed forward: the groups that
+ * may share it packed into one call of the class the caller named (`packed`),
+ * a verify group (which keeps its projection geometry) and every group when
+ * `pack` is off on its own rows, at its own class. */
+export function mapTokenGroups(
+  groups: readonly MlxPhaseGroup[], inputs: readonly MlxArray[],
+  operation: (hidden: MlxArray, width: number) => MlxArray, packed: number, pack = true,
+): MlxArray[] {
+  const results: MlxArray[] = [];
+  const packedRows = groups.flatMap((group, row) => pack && group.phase !== "verify" ? [row] : []);
+  try {
+    const outputs = packedRows.length === 1 ? [operation(inputs[packedRows[0]!]!, groups[packedRows[0]!]!.rows)]
+      : packedRows.length ? mapPackedTokens(packedRows.map(row => inputs[row]!), hidden => operation(hidden, packed)) : [];
+    for (const [index, row] of packedRows.entries()) results[row] = outputs[index]!;
+    for (let row = 0; row < groups.length; row++) results[row] ??= operation(inputs[row]!, groups[row]!.rows);
+    return results;
+  } catch (error) { for (const result of results) result?.dispose(); throw error; }
+}
+
+/** Pack real tokens for a tokenwise operation, then restore each group's
+ * geometry. Attention and recurrent operations must run outside this port.
+ * No padding, host readback, or evaluation is introduced here. */
+export function mapPackedTokens(
+  inputs: readonly MlxArray[], operation: (packed: MlxArray) => MlxArray,
+): MlxArray[] {
+  if (inputs.length === 1) return [operation(inputs[0]!)];
+  const flat = inputs.map(input => ops.reshape(input, [1, -1, input.shape.at(-1)!]));
+  let packed: MlxArray;
+  try { packed = ops.concatAxis(flat, 1); }
+  finally { for (const input of flat) input.dispose(); }
+  using joined = packed;
+  using output = operation(joined);
+  const results: MlxArray[] = [];
+  let position = 0;
+  try {
+    for (const input of inputs) {
+      const [batch, length] = input.shape;
+      const end = position + batch! * length!;
+      using part = output.slice([0, position, 0], [1, end, output.shape[2]!]);
+      results.push(ops.reshape(part, [batch!, length!, output.shape[2]!]));
+      position = end;
+    }
+    return results;
+  } catch (error) { for (const result of results) result.dispose(); throw error; }
+}