mlx-bun 0.4.0 → 1.0.0

package/src/model/profile.ts

@@ -1,443 +0,0 @@
-import type { ModelConfig } from "../config";
-import { configFingerprint } from "./fingerprint";
-import {
-  isDiffusionGemmaConfig,
-  isGlm52Config,
-  isMiniCPM5Config,
-  isQwen35Config,
-  isQwen3Config,
-  isQwen3MoeConfig,
-} from "./support";
-import { GENERATED } from "./generated";
-import { GENERIC_MODEL_TYPES, genericArgsFor, remapModelType } from "./universal/archs";
-
-export const ENGINE_CAPABILITIES = Object.freeze([
-  "autoregressive",
-  "colibri-container",
-  "diffusion",
-  "diffusion-gemma-graph",
-  "gemma4-graph",
-  "generated-graph",
-  "glm5.2-graph",
-  "minicpm5-graph",
-  "mixed-precision-kv",
-  "native-mtp",
-  "qwen3-graph",
-  "qwen3-moe-graph",
-  "qwen3.5-graph",
-  "recurrent-state",
-  "safetensors",
-  "streamed-experts",
-  "universal-dense-graph",
-  "vision-sidecar",
-] as const);
-
-export type EngineCapability = typeof ENGINE_CAPABILITIES[number];
-export type FidelityTier = "l1" | "l2" | "l3";
-
-export type FidelityTarget =
-  | Readonly<{ tier: "l1"; oracle: "mlx-lm"; claim: "bit-exact" }>
-  | Readonly<{ tier: "l2"; oracle: "mlx-optiq"; claim: "bit-exact" }>
-  | Readonly<{ tier: "l3"; oracle: null; claim: "measured" }>;
-
-export type ModelLoader = "safetensors" | "colibri";
-export type ModelGraph =
-  | "gemma4"
-  | "minicpm5"
-  | "qwen3.5"
-  | "qwen3"
-  | "qwen3-moe"
-  | "diffusion-gemma"
-  | "glm5.2"
-  | "universal-dense";
-export type GenerationLoop = "autoregressive" | "diffusion";
-export type ModelSpecialization = "artifact" | "dedicated" | "generated" | "generic";
-
-export interface ModelExecutionComposition {
-  readonly loader: ModelLoader;
-  readonly graph: ModelGraph;
-  readonly loop: GenerationLoop;
-  readonly specialization: ModelSpecialization;
-  /** Engine-owned code registration. Model files never supply executable code.
-   * Omit to retain the existing graph implementation. */
-  readonly implementation?: string;
-}
-
-/** A profile declares construction only. Request methods such as MTP, KV
- * schemes, adapters, grammar, and sampling are resolved independently and
- * cannot be rewritten by profile selection. */
-export interface ModelProfile {
-  readonly id: string;
-  /** Exact external identity. Family profiles omit this field. */
-  readonly artifactFingerprint?: string;
-  /** Structural guard for an exact artifact declaration. */
-  readonly configFingerprint?: string;
-  readonly fidelity: FidelityTarget;
-  readonly requiredCapabilities: readonly EngineCapability[];
-  readonly execution: ModelExecutionComposition;
-}
-
-export interface ArtifactModelProfile extends ModelProfile {
-  readonly artifactFingerprint: string;
-  readonly configFingerprint: string;
-}
-
-export interface ModelArtifactIdentity {
-  readonly fingerprint: string | null;
-  readonly configFingerprint: string;
-}
-
-export interface ResolvedModelProfile {
-  readonly profile: ModelProfile;
-  readonly artifact: ModelArtifactIdentity;
-  readonly exactArtifact: boolean;
-}
-
-export interface ResolveModelProfileOptions {
-  /** Explicit identity for a non-Hugging-Face artifact. The caller owns its
-   * provenance; mlx-bun never mistakes a local path for a content identity. */
-  readonly artifactFingerprint?: string | null;
-  /** Additional exact declarations composed with the shipped profiles. */
-  readonly artifactProfiles?: readonly ArtifactModelProfile[];
-  readonly engineCapabilities?: readonly EngineCapability[];
-}
-
-const L1: FidelityTarget = Object.freeze({
-  tier: "l1", oracle: "mlx-lm", claim: "bit-exact",
-});
-const L2: FidelityTarget = Object.freeze({
-  tier: "l2", oracle: "mlx-optiq", claim: "bit-exact",
-});
-const L3: FidelityTarget = Object.freeze({
-  tier: "l3", oracle: null, claim: "measured",
-});
-
-function freezeProfile<T extends ModelProfile>(profile: T): T {
-  Object.freeze(profile.fidelity);
-  Object.freeze(profile.requiredCapabilities);
-  Object.freeze(profile.execution);
-  return Object.freeze(profile);
-}
-
-function snapshotProfile(profile: ModelProfile): ModelProfile {
-  return freezeProfile({
-    ...profile,
-    fidelity: { ...profile.fidelity },
-    requiredCapabilities: [...profile.requiredCapabilities],
-    execution: { ...profile.execution },
-  });
-}
-
-/** Exact artifact declarations already backed by this repository's parity or
- * measured evidence. A new revision gets the family profile until it earns a
- * declaration of its own. */
-export const BUILTIN_ARTIFACT_PROFILES: readonly ArtifactModelProfile[] = Object.freeze([
-  freezeProfile({
-    id: "qwen3.8-27b-optiq-4bit",
-    artifactFingerprint:
-      "hf:mlx-community/Qwen3.8-27B-OptiQ-4bit@b04599de95d7a9bfbd7f208d347c0f10d9432a42",
-    configFingerprint: "50117975cb405944",
-    fidelity: L1,
-    requiredCapabilities: Object.freeze([
-      "safetensors", "autoregressive", "qwen3.5-graph", "recurrent-state", "vision-sidecar",
-    ]),
-    execution: Object.freeze({
-      loader: "safetensors",
-      graph: "qwen3.5",
-      loop: "autoregressive",
-      specialization: "artifact",
-    }),
-  }),
-  freezeProfile({
-    id: "glm5.2-colibri-int4-int8-mtp",
-    artifactFingerprint:
-      "hf:mateogrgic/GLM-5.2-colibri-int4-with-int8-mtp@3cc8db99b1b13fc79325d987ba3c1c430766b3b8",
-    configFingerprint: "8b49d0941e18f1b9",
-    fidelity: L3,
-    requiredCapabilities: Object.freeze([
-      "colibri-container", "autoregressive", "glm5.2-graph", "streamed-experts",
-    ]),
-    execution: Object.freeze({
-      loader: "colibri",
-      graph: "glm5.2",
-      loop: "autoregressive",
-      specialization: "artifact",
-    }),
-  }),
-]);
-
-const FAMILY_PROFILES = {
-  gemma4Generated: freezeProfile({
-    id: "gemma4-generated",
-    fidelity: L1,
-    requiredCapabilities: ["safetensors", "autoregressive", "gemma4-graph", "generated-graph"],
-    execution: {
-      loader: "safetensors", graph: "gemma4", loop: "autoregressive", specialization: "generated",
-    },
-  }),
-  gemma4: freezeProfile({
-    id: "gemma4-dedicated",
-    fidelity: L1,
-    requiredCapabilities: ["safetensors", "autoregressive", "gemma4-graph"],
-    execution: {
-      loader: "safetensors", graph: "gemma4", loop: "autoregressive", specialization: "dedicated",
-    },
-  }),
-  minicpm5: freezeProfile({
-    id: "minicpm5-dedicated",
-    fidelity: L1,
-    requiredCapabilities: ["safetensors", "autoregressive", "minicpm5-graph"],
-    execution: {
-      loader: "safetensors", graph: "minicpm5", loop: "autoregressive", specialization: "dedicated",
-    },
-  }),
-  qwen35: freezeProfile({
-    id: "qwen3.5-dedicated",
-    fidelity: L1,
-    requiredCapabilities: ["safetensors", "autoregressive", "qwen3.5-graph", "recurrent-state"],
-    execution: {
-      loader: "safetensors", graph: "qwen3.5", loop: "autoregressive", specialization: "dedicated",
-    },
-  }),
-  qwen3: freezeProfile({
-    id: "qwen3-dedicated",
-    fidelity: L1,
-    requiredCapabilities: ["safetensors", "autoregressive", "qwen3-graph"],
-    execution: {
-      loader: "safetensors", graph: "qwen3", loop: "autoregressive", specialization: "dedicated",
-    },
-  }),
-  qwen3Moe: freezeProfile({
-    id: "qwen3-moe-dedicated",
-    fidelity: L1,
-    requiredCapabilities: ["safetensors", "autoregressive", "qwen3-moe-graph"],
-    execution: {
-      loader: "safetensors", graph: "qwen3-moe", loop: "autoregressive", specialization: "dedicated",
-    },
-  }),
-  diffusionGemma: freezeProfile({
-    id: "diffusion-gemma-dedicated",
-    fidelity: L2,
-    requiredCapabilities: ["safetensors", "diffusion", "diffusion-gemma-graph"],
-    execution: {
-      loader: "safetensors", graph: "diffusion-gemma", loop: "diffusion", specialization: "dedicated",
-    },
-  }),
-  glm52: freezeProfile({
-    id: "glm5.2-colibri",
-    fidelity: L3,
-    requiredCapabilities: ["colibri-container", "autoregressive", "glm5.2-graph", "streamed-experts"],
-    execution: {
-      loader: "colibri", graph: "glm5.2", loop: "autoregressive", specialization: "dedicated",
-    },
-  }),
-  universal: freezeProfile({
-    id: "universal-dense",
-    fidelity: L1,
-    requiredCapabilities: ["safetensors", "autoregressive", "universal-dense-graph"],
-    execution: {
-      loader: "safetensors", graph: "universal-dense", loop: "autoregressive", specialization: "generic",
-    },
-  }),
-} satisfies Record<string, ModelProfile>;
-
-/** Return a stable external fingerprint only when the directory is an exact HF
- * snapshot revision. Relocation does not change the identity. Arbitrary local
- * directories and mutable aliases intentionally return null. */
-export function externalArtifactFingerprint(modelDir: string): string | null {
-  const match = modelDir.replaceAll("\\", "/").match(
-    /(?:^|\/)models--([^/]+)\/snapshots\/([0-9a-f]{40,64})(?:\/|$)/i,
-  );
-  if (!match) return null;
-  const encodedRepo = match[1]!;
-  const separator = encodedRepo.indexOf("--");
-  if (separator <= 0 || separator === encodedRepo.length - 2) return null;
-  const repo = `${encodedRepo.slice(0, separator)}/${encodedRepo.slice(separator + 2)}`;
-  return `hf:${repo}@${match[2]!.toLowerCase()}`;
-}
-
-function familyProfile(config: ModelConfig, fingerprint: string): ModelProfile {
-  if (isGlm52Config(config)) return FAMILY_PROFILES.glm52;
-  if (isDiffusionGemmaConfig(config)) return FAMILY_PROFILES.diffusionGemma;
-  if (isMiniCPM5Config(config)) return FAMILY_PROFILES.minicpm5;
-  if (isQwen35Config(config)) return FAMILY_PROFILES.qwen35;
-  if (isQwen3MoeConfig(config)) return FAMILY_PROFILES.qwen3Moe;
-  if (isQwen3Config(config)) return FAMILY_PROFILES.qwen3;
-  if (config.modelType.startsWith("gemma4"))
-    return GENERATED.has(fingerprint) ? FAMILY_PROFILES.gemma4Generated : FAMILY_PROFILES.gemma4;
-  if (genericArgsFor(config)) return FAMILY_PROFILES.universal;
-
-  const arch = remapModelType(config.modelType);
-  throw new Error(
-    `unsupported model_type "${config.modelType}"` +
-    (arch !== config.modelType ? ` (mlx-lm remaps it to "${arch}")` : "") +
-    ` — targeted: gemma4*, diffusion_gemma, glm_moe_dsa, qwen3_5, qwen3, qwen3_moe, MiniCPM5;` +
-    ` generic (Tier-0): ${[...GENERIC_MODEL_TYPES].sort().join(", ")}`,
-  );
-}
-
-interface GraphMetadata {
-  readonly accepts: (config: ModelConfig) => boolean;
-  readonly capabilities: readonly EngineCapability[];
-}
-
-const GRAPH_METADATA: Readonly<Record<ModelGraph, GraphMetadata>> = Object.freeze({
-  "gemma4": {
-    accepts: (config) => config.modelType.startsWith("gemma4"),
-    capabilities: ["gemma4-graph"],
-  },
-  "minicpm5": { accepts: isMiniCPM5Config, capabilities: ["minicpm5-graph"] },
-  "qwen3.5": {
-    accepts: isQwen35Config,
-    capabilities: ["qwen3.5-graph", "recurrent-state"],
-  },
-  "qwen3": { accepts: isQwen3Config, capabilities: ["qwen3-graph"] },
-  "qwen3-moe": { accepts: isQwen3MoeConfig, capabilities: ["qwen3-moe-graph"] },
-  "diffusion-gemma": {
-    accepts: isDiffusionGemmaConfig,
-    capabilities: ["diffusion-gemma-graph"],
-  },
-  "glm5.2": {
-    accepts: isGlm52Config,
-    capabilities: ["glm5.2-graph", "streamed-experts"],
-  },
-  "universal-dense": {
-    accepts: (config) => {
-      try { return genericArgsFor(config) !== null; } catch { return false; }
-    },
-    capabilities: ["universal-dense-graph"],
-  },
-});
-
-function graphAccepts(profile: ModelProfile, config: ModelConfig): boolean {
-  return GRAPH_METADATA[profile.execution.graph].accepts(config);
-}
-
-function executionCapabilities(profile: ModelProfile): EngineCapability[] {
-  const required: EngineCapability[] = [profile.execution.loader === "colibri"
-    ? "colibri-container"
-    : "safetensors"];
-  if (profile.execution.loop === "diffusion") required.push("diffusion");
-  else required.push("autoregressive");
-  required.push(...GRAPH_METADATA[profile.execution.graph].capabilities);
-  if (profile.execution.specialization === "generated") required.push("generated-graph");
-  return required;
-}
-
-function validateProfile(profile: ModelProfile): void {
-  if (!profile.id) throw new Error("model profile id must not be empty");
-  if (profile.execution.implementation !== undefined &&
-      (!profile.execution.implementation.trim() || !profile.artifactFingerprint ||
-       profile.execution.specialization !== "artifact"))
-    throw new Error(`model profile ${profile.id} must bind a named implementation to an exact artifact`);
-  if ((profile.artifactFingerprint === undefined) !== (profile.configFingerprint === undefined))
-    throw new Error(
-      `model profile ${profile.id} must declare artifactFingerprint and configFingerprint together`,
-    );
-  const fidelityOk =
-    (profile.fidelity.tier === "l1" && profile.fidelity.oracle === "mlx-lm" &&
-      profile.fidelity.claim === "bit-exact") ||
-    (profile.fidelity.tier === "l2" && profile.fidelity.oracle === "mlx-optiq" &&
-      profile.fidelity.claim === "bit-exact") ||
-    (profile.fidelity.tier === "l3" && profile.fidelity.oracle === null &&
-      profile.fidelity.claim === "measured");
-  if (!fidelityOk) throw new Error(`model profile ${profile.id} has an invalid fidelity contract`);
-  const expectedLoader = profile.execution.graph === "glm5.2" ? "colibri" : "safetensors";
-  const expectedLoop = profile.execution.graph === "diffusion-gemma"
-    ? "diffusion"
-    : "autoregressive";
-  if (profile.execution.loader !== expectedLoader || profile.execution.loop !== expectedLoop)
-    throw new Error(
-      `model profile ${profile.id} has an invalid execution composition for ` +
-      `${profile.execution.graph}`,
-    );
-  const declared = new Set(profile.requiredCapabilities);
-  const underdeclared = executionCapabilities(profile).filter((capability) => !declared.has(capability));
-  if (underdeclared.length)
-    throw new Error(
-      `model profile ${profile.id} does not declare execution capabilities: ` +
-      `${underdeclared.join(", ")}`,
-    );
-}
-
-export function resolveModelProfile(
-  config: ModelConfig,
-  options: ResolveModelProfileOptions = {},
-): ResolvedModelProfile {
-  const fingerprint = configFingerprint(config);
-  const artifactFingerprint = options.artifactFingerprint === undefined
-    ? externalArtifactFingerprint(config.modelDir)
-    : options.artifactFingerprint;
-  const artifactProfiles = options.artifactProfiles
-    ? [...BUILTIN_ARTIFACT_PROFILES, ...options.artifactProfiles]
-    : BUILTIN_ARTIFACT_PROFILES;
-  const capabilities = new Set(options.engineCapabilities ?? ENGINE_CAPABILITIES);
-
-  let profile: ModelProfile | undefined;
-  let exactArtifact = false;
-  if (artifactFingerprint) {
-    const matches = artifactProfiles.filter((entry) => {
-      validateProfile(entry);
-      return entry.artifactFingerprint === artifactFingerprint;
-    });
-    if (matches.length > 1)
-      throw new Error(`multiple model profiles declare artifact ${artifactFingerprint}`);
-    profile = matches[0];
-    if (profile) {
-      exactArtifact = true;
-      if (profile.configFingerprint !== fingerprint)
-        throw new Error(
-          `model profile ${profile.id} matched artifact ${artifactFingerprint}, but config fingerprint ` +
-          `${fingerprint} != declared ${profile.configFingerprint}; refusing to fall back`,
-        );
-      if (!graphAccepts(profile, config))
-        throw new Error(
-          `model profile ${profile.id} selects ${profile.execution.graph}, which is incompatible with ` +
-          `${config.modelType}; refusing to fall back`,
-        );
-    }
-  }
-  profile ??= familyProfile(config, fingerprint);
-  validateProfile(profile);
-  profile = snapshotProfile(profile);
-
-  const missing = profile.requiredCapabilities.filter((capability) => !capabilities.has(capability));
-  if (missing.length)
-    throw new Error(
-      `model profile ${profile.id} requires missing engine capabilities: ${missing.join(", ")}; ` +
-      `refusing to fall back`,
-    );
-
-  return Object.freeze({
-    profile,
-    artifact: Object.freeze({ fingerprint: artifactFingerprint, configFingerprint: fingerprint }),
-    exactArtifact,
-  });
-}
-
-/** Guard the factory seam when a caller passes a previously resolved profile. */
-export function assertResolvedModelProfile(
-  config: ModelConfig,
-  resolved: ResolvedModelProfile,
-): void {
-  validateProfile(resolved.profile);
-  const fingerprint = configFingerprint(config);
-  if (resolved.artifact.configFingerprint !== fingerprint)
-    throw new Error(
-      `model profile ${resolved.profile.id} was resolved for config ` +
-      `${resolved.artifact.configFingerprint}, not ${fingerprint}`,
-    );
-  if (!graphAccepts(resolved.profile, config))
-    throw new Error(
-      `model profile ${resolved.profile.id} selects ${resolved.profile.execution.graph}, ` +
-      `which is incompatible with ${config.modelType}`,
-    );
-  if (resolved.exactArtifact !== (resolved.profile.artifactFingerprint !== undefined) ||
-      (resolved.exactArtifact &&
-       (resolved.profile.artifactFingerprint !== resolved.artifact.fingerprint ||
-        resolved.profile.configFingerprint !== fingerprint)))
-    throw new Error(
-      `model profile ${resolved.profile.id} does not match artifact ` +
-      `${resolved.artifact.fingerprint ?? "<unidentified>"}`,
-    );
-}

package/src/model/qwen-append.ts

@@ -1,12 +0,0 @@
-// MLX 0.32.2, mlx/backend/metal/scaled_dot_product_attention.cpp (MIT):
-// https://github.com/ml-explore/mlx/blob/v0.32.2/mlx/backend/metal/scaled_dot_product_attention.cpp
-// For the qualified Qwen GQA-6/D256 graph on applegpu_g16s, one-pass/vector
-// selection and two-pass block counts change after these inclusive KV lengths.
-// A committed span must stay within one arithmetic regime to match M=1.
-const ATTENTION_BOUNDARIES = [1023, 1024, 8192, 32768, 65536];
-
-export function qwenAppendChunkSize(cacheTokens: number): number {
-  for (const end of ATTENTION_BOUNDARIES)
-    if (cacheTokens < end) return Math.min(4, end - cacheTokens);
-  return 4;
-}

package/src/model/qwen-conv.ts

@@ -1,52 +0,0 @@
-// Experimental Qwen GDN convolution. Its two outputs are evaluated together:
-// consuming the activation also materializes the small, independently owned
-// state tail, without keeping a concatenated prefill buffer alive for it.
-// Arithmetic follows MLX 0.31.2 depthwise_conv_1d and unary sigmoid (MIT,
-// Copyright Apple Inc.), preserving the bf16 intermediate rounding.
-import type { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import { MetalKernel } from "../mlx/metal-kernel";
-
-/** Borrow all inputs; return owned activation and an independent state tail. */
-export type QwenConvolution = (x: MlxArray, state: MlxArray, weight: MlxArray) => [MlxArray, MlxArray];
-
-let kernel: MetalKernel | null = null;
-export const fusedQwenConvolution: QwenConvolution = (x, state, weight) => {
-  const [B, M, D] = x.shape as [number, number, number];
-  const K = weight.shape[1]!;
-  if (x.shape.length !== 3 || weight.shape.length !== 3 || weight.shape[0] !== D ||
-      weight.shape[2] !== 1 || K !== 4 || state.shape.join() !== [B, K - 1, D].join() ||
-      ![Dtype.bfloat16, Dtype.float32].includes(x.dtype) || state.dtype !== x.dtype || weight.dtype !== x.dtype)
-    throw new Error("fused Qwen convolution requires matching bf16/f32 inputs, [B,M,D], [B,3,D], [D,4,1]");
-  kernel ??= new MetalKernel({ name: "qwen_causal_conv_silu_tail",
-    inputNames: ["state", "x", "w"], outputNames: ["out", "tail"],
-    source: String.raw`
-      const uint c = thread_position_in_grid.x;
-      const uint t = thread_position_in_grid.y;
-      const uint b = thread_position_in_grid.z;
-      if (c >= D) return;
-      if (t < M) {
-        float acc = 0.0f;
-        for (uint i = 0; i < K; ++i) {
-          const uint row = t + i;
-          const T value = row < K-1 ? state[(b*(K-1)+row)*D+c] : x[(b*M+row-(K-1))*D+c];
-          acc += float(value) * float(w[c*K+i]);
-        }
-        const T conv = T(acc);
-        // Preserve the pinned MLX sigmoid's typed operations and bf16
-        // intermediate rounding. Promoting this expression to f32 differs.
-        const auto y = 1 / (1 + metal::exp(metal::abs(conv)));
-        const T sig = conv < 0 ? y : 1-y;
-        out[(b*M+t)*D+c] = conv * sig;
-      }
-      if (t < K-1) {
-        const uint row = M + t;
-        tail[(b*(K-1)+t)*D+c] = row < K-1 ? state[(b*(K-1)+row)*D+c] : x[(b*M+row-(K-1))*D+c];
-      }
-    ` });
-  return kernel.apply([state, x, weight], {
-    outputs: [{ shape: [B, M, D], dtype: x.dtype }, { shape: [B, K - 1, D], dtype: x.dtype }],
-    grid: [D, Math.max(M, K - 1), B], threadGroup: [256, 1, 1],
-    templateDtypes: { T: x.dtype }, templateInts: { D, M, K },
-  }) as [MlxArray, MlxArray];
-};

package/src/model/qwen3-delta-state.ts

@@ -1,74 +0,0 @@
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import { MetalKernel } from "../mlx/metal-kernel";
-import * as ops from "../mlx/ops";
-import { computeG } from "./qwen3-delta";
-
-// The same recurrent arithmetic as gated_delta_step, with no query/output
-// calculation. Each row replays only its retained prefix of the verify window.
-const SOURCE = String.raw`
-  const int T = Tin[0];
-  auto n = thread_position_in_grid.z;
-  auto b_idx = n / Hv;
-  auto hv_idx = n % Hv;
-  auto hk_idx = hv_idx / (Hv / Hk);
-  constexpr int n_per_t = Dk / 32;
-  auto dk_idx = thread_position_in_threadgroup.x;
-  auto dv_idx = thread_position_in_grid.y;
-  auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
-  auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
-  auto g_ = g + b_idx * T * Hv;
-  auto beta_ = beta + b_idx * T * Hv;
-  auto i_state = state_in + (n * Dv + dv_idx) * Dk;
-  auto o_state = state_out + (n * Dv + dv_idx) * Dk;
-  float state[n_per_t];
-  for (int i = 0; i < n_per_t; ++i)
-    state[i] = static_cast<float>(i_state[n_per_t * dk_idx + i]);
-  const int keep = lengths[b_idx];
-  for (int t = 0; t < keep; ++t) {
-    float kv_mem = 0.0f;
-    for (int i = 0; i < n_per_t; ++i) {
-      auto s_idx = n_per_t * dk_idx + i;
-      state[i] = state[i] * g_[hv_idx];
-      kv_mem += state[i] * k_[s_idx];
-    }
-    kv_mem = simd_sum(kv_mem);
-    auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
-    for (int i = 0; i < n_per_t; ++i) {
-      auto s_idx = n_per_t * dk_idx + i;
-      state[i] = state[i] + k_[s_idx] * delta;
-    }
-    k_ += Hk * Dk;
-    v_ += Hv * Dv;
-    g_ += Hv;
-    beta_ += Hv;
-  }
-  for (int i = 0; i < n_per_t; ++i)
-    o_state[n_per_t * dk_idx + i] = state[i];
-`;
-
-let kernel: MetalKernel | undefined;
-
-/** State-only prefix replay. Inputs keep their original [B,T,...] strides;
- * lengths are request-owned accepted counts, never scheduling decisions. */
-export function gatedDeltaState(
-  k: MlxArray, v: MlxArray, a: MlxArray, b: MlxArray,
-  aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null,
-  lengths: readonly number[],
-): MlxArray {
-  const [B, T, Hk, Dk] = k.shape as [number, number, number, number];
-  const [, , Hv, Dv] = v.shape as [number, number, number, number];
-  using beta = ops.sigmoid(b);
-  using g = computeG(aLog, a, dtBias);
-  using zero = state ? null : ops.zeros([B, Hv, Dv, Dk], Dtype.float32);
-  using t = ops.fromInt32([T], [1]);
-  using kept = ops.fromInt32([...lengths], [B]);
-  kernel ??= new MetalKernel({ name: "gated_delta_prefix_state",
-    inputNames: ["k", "v", "g", "beta", "state_in", "Tin", "lengths"],
-    outputNames: ["state_out"], source: SOURCE, ensureRowContiguous: true });
-  return kernel.apply([k, v, g, beta, state ?? zero!, t, kept], {
-    outputs: [{ shape: [B, Hv, Dv, Dk], dtype: Dtype.float32 }],
-    grid: [32, Dv, B * Hv], threadGroup: [32, 4, 1],
-    templateInts: { Dk, Dv, Hk, Hv },
-  })[0]!;
-}