mlx-bun 0.5.0 → 1.0.0

package/src/model/trellis-wide-prefill.ts

@@ -1,94 +0,0 @@
-// Packed Trellis prefill with MLX 0.32.2 GemvWide arithmetic.
-// Ported from ml-explore/mlx v0.32.2, mlx/backend/metal/kernels/gemv.h
-// (MIT). Decode weights to bf16 inside the native dot/unroll/shuffle order;
-// share each decoded row across every input vector in this narrow regime.
-import { MetalKernel } from "../mlx/metal-kernel";
-import { type MlxArray } from "../mlx/array";
-import { Dtype, deviceArchitecture } from "../mlx/ffi";
-import { type TrellisGeometry } from "./trellis-linear";
-
-/** Matches the bundled runtime's native reduction on M3 and newer GPUs. */
-export function nativeTrellisWidePrefill(m: number): boolean {
-  return m >= 5 && m <= 15 &&
-    Number(/^applegpu_[a-z](\d{2})/.exec(deviceArchitecture())?.[1] ?? 0) >= 15;
-}
-
-export function wideTrellisPrefillEligible(g: TrellisGeometry, m: number, dtype: Dtype): boolean {
-  return dtype === Dtype.bfloat16 && nativeTrellisWidePrefill(m) &&
-    g.axis === 1 && g.T === 256 && g.L === 12 && [2, 3, 4].includes(g.k) &&
-    g.inFeatures === 5120 && g.outFeatures === 17408 && !g.blockInterleave;
-}
-
-let kernel: MetalKernel | undefined;
-function wideKernel() {
-  return kernel ??= new MetalKernel({ name: 'mlx_bun_trellis_wide_prefill',
-    inputNames: ['x', 'codes', 'scales'], outputNames: ['out'], source: String.raw`
-    constexpr uint unroll=8, rows_per_simd=32/KL, simdgroups=KL/8;
-    const uint lane=thread_index_in_simdgroup, sg=simdgroup_index_in_threadgroup;
-    const uint kl=lane%KL, row=threadgroup_position_in_grid.y*4+sg*rows_per_simd+lane/KL;
-    const uint vec0=threadgroup_position_in_grid.x*SM;
-    constexpr uint wpb=BT*BITS/32, row_words=KC*BITS/32;
-    const float scale=float(scales[row]);
-    constexpr uint n_v4=KC/4, n_main=n_v4-n_v4%(KL*unroll);
-    float result[SM]={0};
-    for(uint base=0;base<n_main;base+=KL*unroll) {
-      float4 wf[unroll];
-      #pragma clang loop unroll(full)
-      for(uint i=0;i<unroll;i++) {
-        #pragma clang loop unroll(full)
-        for(uint c=0;c<4;c++) {
-          const uint k=(base+i*KL+kl)*4+c, block=k/BT, t=k%BT;
-          const device uint* bw=codes+(ulong)row*row_words+block*wpb;
-          const uint pos=(BT-1-t)*BITS, wi=pos>>5, off=pos&31u;
-          uint win=bw[wi]>>off;
-          if(off+L>32)win|=bw[wi+1==wpb?0:wi+1]<<(32-off);
-          win&=(1u<<L)-1;
-          const uint z=win*34038481u+76625530u;
-          const uint p=(z&0x00FF00FFu)+((z>>8)&0x00FF00FFu);
-          const int code=int((p&0xFFFFu)+(p>>16))-510;
-          const float reciprocal=1.0f/147.800537109375f;
-          const float approx=float(code)*reciprocal;
-          const float residual=metal::fma(-approx,147.800537109375f,float(code));
-          const float value=metal::fma(residual,reciprocal,approx);
-          wf[i][c]=float(T(value*scale));
-        }
-      }
-      #pragma clang loop unroll(full)
-      for(uint v=0;v<SM;v++) {
-        const uint m=min(vec0+v,uint(M-1));
-        float acc=0;
-        #pragma clang loop unroll(full)
-        for(uint i=0;i<unroll;i++) {
-          const ulong offset=(ulong)m*KC+(base+i*KL+kl)*4;
-          const float4 xv=float4(float(x[offset]),float(x[offset+1]),float(x[offset+2]),float(x[offset+3]));
-          acc+=metal::dot(wf[i],xv);
-        }
-        result[v]+=acc;
-      }
-    }
-    #pragma clang loop unroll(full)
-    for(uint v=0;v<SM;v++) {
-      #pragma clang loop unroll(full)
-      for(ushort off=KL/2;off>=1;off>>=1) result[v]+=simd_shuffle_down(result[v],off);
-      if(kl==0 && vec0+v<M) out[(ulong)(vec0+v)*N+row]=T(result[v]);
-    }
-  ` });
-}
-/** Requires the caller to prove aligned, row-contiguous x, as Qwen does
- * with its RMSNorm output. Preserve native fallback when that proof is absent. */
-export function wideTrellisPrefill(x: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
-  const [m, k] = x.shape;
-  if (x.ndim !== 2 || x.dtype !== Dtype.bfloat16 || !m || m < 5 || m > 15 ||
-      k !== 5120 || g.inFeatures !== k || g.axis !== 1 || g.T !== 256 || g.L !== 12 ||
-      ![2, 3, 4].includes(g.k) || g.blockInterleave || g.rows < 4 || g.rows % 4 ||
-      g.outFeatures !== g.rows || g.cols !== k || codes.dtype !== Dtype.uint32 ||
-      codes.ndim !== 2 || codes.shape[0] !== g.rows || codes.shape[1] !== k * g.k / 32 || scales.size !== g.rows)
-    throw new Error("unsupported wide trellis prefill geometry, dtype or layout");
-  const kl = m <= 5 ? 32 : 16;
-  return wideKernel().apply([x, codes, scales], {
-    outputs: [{ shape: [m, g.outFeatures], dtype: x.dtype }],
-    grid: [32, g.outFeatures / 4 * (kl / 8), 1], threadGroup: [32, kl / 8, 1],
-    templateDtypes: { T: x.dtype },
-    templateInts: { M: m, SM: m, KL: kl, KC: k, N: g.outFeatures, BT: g.T, L: g.L, BITS: g.k },
-  })[0]!;
-}

package/src/model/turboquant-codec.ts

@@ -1,89 +0,0 @@
-import { runtimeConfig, type RuntimeConfig } from "../runtime-config";
-import type { MlxArray } from "../mlx/array";
-import * as tq from "../mlx/turboquant-ops";
-import type { KvCodec } from "../backends/mlx/kv-codec";
-
-/** Resolve once when a cache is created; storage copies retain this policy. */
-export function turboQuantFusedDecode(config: RuntimeConfig = runtimeConfig()): boolean {
-  return config.flag("MLX_BUN_TURBOQUANT_FUSED_DECODE", true);
-}
-
-/** One TurboQuant-encoded (K, V) storage tuple — the 5 arrays kv-store.ts
- *  and toQuantized/fromKVCache pass around. Not `ops.QuantizedTensor`
- *  (mlx's affine int4/int8 scheme): this is the rotation + Lloyd-Max
- *  layout (docs/design/turboquant.md), asymmetric-affine for keys,
- *  FWHT+Lloyd-Max for values. Field order is the kv-store.ts tensor-slot
- *  contract — do not reorder without updating snapshotCache/loadKvCache. */
-export interface TurboQuantTensor {
-  kIdx: MlxArray;
-  kScales: MlxArray;
-  kZeros: MlxArray;
-  vPacked: MlxArray;
-  vScales: MlxArray;
-}
-
-export const disposeTurboQuant = (t: TurboQuantTensor): void => {
-  t.kIdx.dispose();
-  t.kScales.dispose();
-  t.kZeros.dispose();
-  t.vPacked.dispose();
-  t.vScales.dispose();
-};
-
-/** Numerical representation only. Storage owns capacity, positions, row
- * membership and persistence; this codec retains the existing kernel choices. */
-export class TurboQuantCodec implements KvCodec<TurboQuantTensor> {
-  constructor(readonly kBits: number, readonly vBits: number, readonly fusedDecode: boolean) {}
-
-  /** Encode k/v [B,H,L,D] into a fresh TurboQuantTensor (packed at rest). */
-  encode(k: MlxArray, v: MlxArray): TurboQuantTensor {
-    const signed = this.kBits === 8;
-    const kEnc = tq.encodeKeys(k, this.kBits, signed);
-    const kIdxPacked = this.kBits < 8 ? tq.packBits(kEnc.indices, this.kBits) : kEnc.indices;
-    if (kIdxPacked !== kEnc.indices) kEnc.indices.dispose();
-    const vEnc = tq.encodeValues(v, this.vBits);
-    const vIdxPacked = this.vBits < 8 ? tq.packBits(vEnc.indices, this.vBits) : vEnc.indices;
-    if (vIdxPacked !== vEnc.indices) vEnc.indices.dispose();
-    return {
-      kIdx: kIdxPacked, kScales: kEnc.scales, kZeros: kEnc.zeros,
-      vPacked: vIdxPacked, vScales: vEnc.scales,
-    };
-  }
-
-  /** Decode a TurboQuantTensor's active [B,H,upTo,*] window back to
-   *  bf16 [B,H,upTo,headDim] (k, v). Caller owns the returned arrays.
-   *  deferV leaves V in the rotated domain (decodeValuesRotated) — the
-   *  caller un-rotates its attention output instead (tq.unrotateValues). */
-  decode(t: TurboQuantTensor, upTo: number, headDim: number, deferV = false): [MlxArray, MlxArray] {
-    const cut = (a: MlxArray): MlxArray => {
-      const [B, H, , D] = a.shape as [number, number, number, number];
-      return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
-    };
-    if (this.fusedDecode) {
-      const inputs = [cut(t.kIdx), cut(t.kScales), cut(t.kZeros), cut(t.vPacked), cut(t.vScales)] as const;
-      try {
-        const decoded = tq.tryDecodePackedKv(inputs, this.kBits, this.vBits, headDim, deferV);
-        if (decoded) return decoded;
-      } finally { for (const input of inputs) input.dispose(); }
-    }
-    const kIdxCut = cut(t.kIdx);
-    const kIdxUnpacked = this.kBits < 8 ? tq.unpackBits(kIdxCut, this.kBits, headDim) : kIdxCut;
-    if (kIdxUnpacked !== kIdxCut) kIdxCut.dispose();
-    const kScalesCut = cut(t.kScales);
-    const kZerosCut = cut(t.kZeros);
-    const k = tq.decodeKeys(kIdxUnpacked, kScalesCut, kZerosCut);
-    for (const a of [kIdxUnpacked, kScalesCut, kZerosCut]) a.dispose();
-
-    const vPackedCut = cut(t.vPacked);
-    const vIdxUnpacked = this.vBits < 8 ? tq.unpackBits(vPackedCut, this.vBits, headDim) : vPackedCut;
-    if (vIdxUnpacked !== vPackedCut) vPackedCut.dispose();
-    const vScalesCut = cut(t.vScales);
-    const v = deferV
-      ? tq.decodeValuesRotated(vIdxUnpacked, vScalesCut, this.vBits)
-      : tq.decodeValues(vIdxUnpacked, vScalesCut, this.vBits);
-    for (const a of [vIdxUnpacked, vScalesCut]) a.dispose();
-
-    return [k, v];
-  }
-
-}

package/src/model/universal/archs.ts

@@ -1,332 +0,0 @@
-// Arch descriptor table for the Tier-0 universal dense module
-// (docs/design/generic-model-support.md §3.1). One entry per mlx-lm
-// model_type; each parse function is the ~15–25-line transcription of that
-// arch's ModelArgs (defaults included) + structural deltas, read from the
-// RAW config.json exactly like the oracle's dataclass does.
-//
-// Support is this table's keys (+ MODEL_REMAPPING aliases) — an explicit,
-// declared surface. No config sniffing. A model_type absent here is
-// unsupported until a descriptor + green parity manifest entry exist.
-
-import type { ModelConfig } from "../../config";
-import type { RopeScalingConfig } from "./rope";
-
-/** mlx_lm/utils.py MODEL_REMAPPING, vendored (free aliases). Only entries
- *  whose TARGET has a descriptor below actually load; the rest are kept so
- *  the reject message can name the real arch. */
-export const MODEL_REMAPPING: Record<string, string> = {
-  mistral: "llama",
-  llava: "mistral3",
-  "phi-msft": "phixtral",
-  falcon_mamba: "mamba",
-  joyai_llm_flash: "deepseek_v3",
-  kimi_k2: "deepseek_v3",
-  qwen2_5_vl: "qwen2_vl",
-  minimax_m2: "minimax",
-  iquestcoder: "llama",
-};
-
-export function remapModelType(modelType: string): string {
-  return MODEL_REMAPPING[modelType] ?? modelType;
-}
-
-/** Normalized args for the universal dense forward — the union of the
- *  launch archs' ModelArgs, one value per knob (no optionals left). */
-export interface UniversalArgs {
-  modelType: string;
-  hiddenSize: number;
-  numHiddenLayers: number;
-  intermediateSize: number;
-  numHeads: number;
-  numKvHeads: number;
-  headDim: number;
-  normEps: number;
-  vocabSize: number;
-  ropeTheta: number;
-  ropeTraditional: boolean;
-  ropeScaling: RopeScalingConfig | null;
-  maxPositionEmbeddings: number | null;
-  /** q/k/v projection bias; o and mlp separately (qwen2: qkv only). */
-  qkvBias: boolean;
-  oBias: boolean;
-  mlpBias: boolean;
-  tieWordEmbeddings: boolean;
-  /** null ⇒ every layer full attention. */
-  layerTypes: string[] | null;
-  slidingWindow: number | null;
-  // ---- block structure ----
-  /** pre  = llama (norm → sublayer → residual)
-   *  post = olmo2 (sublayer → norm → residual)
-   *  gemma2 = sandwich (input/post-attn + pre/post-ffn norms)
-   *  glm4 = post-sublayer norms with pre-norms too */
-  block: "pre" | "post" | "gemma2" | "glm4";
-  norm: "rmsnorm" | "rmsnorm_plus_one" | "layernorm";
-  /** qwen3 = per-head RMSNorm over head_dim (before transpose);
-   *  olmo2 = full-width RMSNorm on the flat projection (before reshape). */
-  qkNorm: "none" | "head" | "full";
-  mlp: "swiglu" | "geglu" | "geglu_approx" | "fused_swiglu" | "gelu_mlp";
-  /** phi3: one fused qkv_proj, split as activations (mx.split — exactly
-   *  what mlx-lm does; the quantized matrix stays fused). */
-  fusedQkv: boolean;
-  attnScale: number;
-  /** gemma2 attention-score softcap ⇒ manual q·kᵀ→tanh→softmax path. */
-  attnLogitSoftcap: number | null;
-  finalLogitSoftcap: number | null;
-  /** granite: logits / logits_scaling. */
-  logitsDivisor: number | null;
-  /** gemma: hidden**0.5; granite: embedding_multiplier. */
-  embedMultiplier: number | null;
-  /** granite residual_multiplier. */
-  residualMultiplier: number | null;
-  rope: "initialize" | "plain" | "phi3";
-  /** rope dims = floor(headDim * partialRotaryFactor) (glm4, phi3). */
-  partialRotaryFactor: number;
-  /** phi3 SuScaledRoPE original context (TOP-LEVEL config key). */
-  originalMaxPositionEmbeddings: number;
-  /** smollm3: per-layer 1 = rope, 0 = NoPE. null ⇒ rope everywhere. */
-  noRopeLayers: number[] | null;
-  /** gemma2: create_attention_mask(..., return_array=True). */
-  maskArray: boolean;
-  /** sanitize: tensors mlx-lm drops on load (audit allowances). */
-  dropWeightPatterns: RegExp[];
-}
-
-type Raw = Record<string, any>;
-
-const INV_FREQ = /self_attn\.rotary_emb\.inv_freq/;
-
-/** Common fields most archs share; per-arch parsers override the deltas. */
-function common(raw: Raw, modelType: string): UniversalArgs {
-  const heads = raw.num_attention_heads as number;
-  return {
-    modelType,
-    hiddenSize: raw.hidden_size,
-    numHiddenLayers: raw.num_hidden_layers,
-    intermediateSize: raw.intermediate_size,
-    numHeads: heads,
-    numKvHeads: raw.num_key_value_heads ?? heads,
-    headDim: Math.floor(raw.hidden_size / heads),
-    normEps: raw.rms_norm_eps,
-    vocabSize: raw.vocab_size,
-    ropeTheta: raw.rope_theta ?? 10000,
-    ropeTraditional: raw.rope_traditional ?? false,
-    ropeScaling: (raw.rope_scaling as RopeScalingConfig | undefined) ?? null,
-    maxPositionEmbeddings: raw.max_position_embeddings ?? null,
-    qkvBias: false,
-    oBias: false,
-    mlpBias: false,
-    tieWordEmbeddings: raw.tie_word_embeddings ?? true,
-    layerTypes: null,
-    slidingWindow: null,
-    block: "pre",
-    norm: "rmsnorm",
-    qkNorm: "none",
-    mlp: "swiglu",
-    fusedQkv: false,
-    attnScale: Math.pow(Math.floor(raw.hidden_size / heads), -0.5),
-    attnLogitSoftcap: null,
-    finalLogitSoftcap: null,
-    logitsDivisor: null,
-    embedMultiplier: null,
-    residualMultiplier: null,
-    rope: "initialize",
-    partialRotaryFactor: 1.0,
-    originalMaxPositionEmbeddings: 4096,
-    noRopeLayers: null,
-    maskArray: false,
-    dropWeightPatterns: [],
-  };
-}
-
-function parseLlama(raw: Raw, modelType = "llama"): UniversalArgs {
-  const a = common(raw, modelType);
-  a.headDim = raw.head_dim ?? a.headDim;
-  a.attnScale = Math.pow(a.headDim, -0.5);
-  a.qkvBias = raw.attention_bias ?? false;
-  a.oBias = raw.attention_bias ?? false; // llama.py wires o_proj bias to the same flag
-  a.mlpBias = raw.mlp_bias ?? false;
-  a.layerTypes = (raw.layer_types as string[] | undefined) ?? null;
-  a.slidingWindow = raw.sliding_window ?? null;
-  a.dropWeightPatterns = [INV_FREQ];
-  if (a.tieWordEmbeddings) a.dropWeightPatterns.push(/^lm_head\.weight$/);
-  return a;
-}
-
-export const ARCHS: Record<string, (raw: Raw) => UniversalArgs> = {
-  /** llama (+ mistral / iquestcoder via remap): the reference dense block. */
-  llama: (raw) => parseLlama(raw),
-
-  /** smollm3: llama with NoPE layers (every no_rope_layer_interval-th). */
-  smollm3: (raw) => {
-    const a = parseLlama(raw, "smollm3");
-    const interval = raw.no_rope_layer_interval ?? 4;
-    const list = raw.no_rope_layers as number[] | undefined;
-    if (list && list.length !== a.numHiddenLayers)
-      throw new Error("`no_rope_layers` length mismatch");
-    a.noRopeLayers = list ??
-      Array.from({ length: a.numHiddenLayers }, (_, i) => ((i + 1) % interval !== 0 ? 1 : 0));
-    return a;
-  },
-
-  /** qwen2: llama block + hardcoded q/k/v bias (o unbiased), theta 1e6. */
-  qwen2: (raw) => {
-    const a = common(raw, "qwen2");
-    a.ropeTheta = raw.rope_theta ?? 1000000;
-    a.maxPositionEmbeddings = raw.max_position_embeddings ?? 32768;
-    a.qkvBias = true;
-    a.dropWeightPatterns = [INV_FREQ];
-    if (a.tieWordEmbeddings) a.dropWeightPatterns.push(/^lm_head\.weight$/);
-    return a;
-  },
-
-  /** qwen3: per-head q/k RMSNorm over head_dim BEFORE transpose+rope. */
-  qwen3: (raw) => {
-    const a = common(raw, "qwen3");
-    a.headDim = raw.head_dim; // required in qwen3 ModelArgs
-    a.attnScale = Math.pow(a.headDim, -0.5);
-    a.qkNorm = "head";
-    a.tieWordEmbeddings = raw.tie_word_embeddings ?? false;
-    if (a.tieWordEmbeddings) a.dropWeightPatterns = [/^lm_head\.weight$/];
-    return a;
-  },
-
-  /** gemma (v1): (1+w) RMSNorm, √hidden embed scale, PRECISE-gelu gated
-   *  MLP, always-tied head, plain rope (rope_scaling ignored). */
-  gemma: (raw) => {
-    const a = common(raw, "gemma");
-    a.headDim = raw.head_dim; // required
-    a.attnScale = Math.pow(a.headDim, -0.5);
-    a.norm = "rmsnorm_plus_one";
-    a.mlp = "geglu";
-    a.embedMultiplier = Math.pow(raw.hidden_size, 0.5);
-    a.tieWordEmbeddings = true; // as_linear unconditionally
-    a.rope = "plain";
-    a.ropeScaling = null;
-    return a;
-  },
-
-  /** gemma2: gemma + sandwich norms, attention-score softcap (manual
-   *  attention), query_pre_attn_scalar scale, final logit softcap,
-   *  gelu_approx MLP, always-array mask. mlx-lm runs every layer full
-   *  attention (no sliding window in its gemma2) — we match the oracle. */
-  gemma2: (raw) => {
-    const a = common(raw, "gemma2");
-    a.headDim = raw.head_dim; // required
-    a.norm = "rmsnorm_plus_one";
-    a.block = "gemma2";
-    a.mlp = "geglu_approx";
-    a.embedMultiplier = Math.pow(raw.hidden_size, 0.5);
-    a.tieWordEmbeddings = true;
-    a.rope = "plain";
-    a.ropeScaling = null;
-    a.attnScale = 1.0 / Math.pow(raw.query_pre_attn_scalar ?? 144.0, 0.5);
-    a.attnLogitSoftcap = raw.attn_logit_softcapping ?? 50.0;
-    a.finalLogitSoftcap = raw.final_logit_softcapping ?? 30.0;
-    a.maskArray = true;
-    return a;
-  },
-
-  /** phi3: fused qkv + fused gate_up (split as ACTIVATIONS, like mlx-lm),
-   *  longrope/su (or linear) rope on partial dims, untied head. */
-  phi3: (raw) => {
-    const a = common(raw, "phi3");
-    a.tieWordEmbeddings = raw.tie_word_embeddings ?? false;
-    a.fusedQkv = true;
-    a.mlp = "fused_swiglu";
-    a.rope = "phi3";
-    a.partialRotaryFactor = raw.partial_rotary_factor ?? 1.0;
-    a.maxPositionEmbeddings = raw.max_position_embeddings ?? 131072;
-    a.originalMaxPositionEmbeddings = raw.original_max_position_embeddings ?? 4096;
-    // phi3.py __post_init__: rope_scaling must carry {long_factor, type};
-    // only longrope/su/linear supported, anything else → scaling off.
-    if (a.ropeScaling) {
-      for (const key of ["long_factor", "type"])
-        if (!(key in a.ropeScaling))
-          throw new Error(`rope_scaling must contain keys {'long_factor', 'type'}`);
-      const t = String(a.ropeScaling.type);
-      if (!["longrope", "su", "linear"].includes(t)) a.ropeScaling = null;
-    }
-    return a;
-  },
-
-  /** olmo2: POST-norm block (norm on sublayer output), full-width q/k
-   *  RMSNorm on the flat projections. */
-  olmo2: (raw) => {
-    const a = common(raw, "olmo2");
-    a.headDim = raw.head_dim ?? a.headDim;
-    a.attnScale = Math.pow(a.headDim, -0.5);
-    a.block = "post";
-    a.qkNorm = "full";
-    a.qkvBias = raw.attention_bias ?? false;
-    a.oBias = raw.attention_bias ?? false;
-    a.mlpBias = raw.mlp_bias ?? false;
-    a.dropWeightPatterns = [INV_FREQ];
-    return a;
-  },
-
-  /** glm4: pre+post sublayer norms (post_self_attn / post_mlp), fused
-   *  gate_up MLP, PARTIAL rope with traditional=True default, always a
-   *  separate lm_head. */
-  glm4: (raw) => {
-    const a = common(raw, "glm4");
-    a.headDim = raw.head_dim ?? a.headDim;
-    a.attnScale = Math.pow(a.headDim, -0.5);
-    a.block = "glm4";
-    a.mlp = "fused_swiglu";
-    a.qkvBias = raw.attention_bias;
-    a.rope = "plain";
-    a.ropeScaling = null;
-    a.ropeTraditional = raw.rope_traditional ?? true;
-    a.partialRotaryFactor = raw.partial_rotary_factor;
-    a.maxPositionEmbeddings = raw.max_position_embeddings ?? 32768;
-    a.tieWordEmbeddings = false; // glm4.py constructs lm_head unconditionally
-    return a;
-  },
-
-  /** granite: llama block scaled everywhere — attention_multiplier as the
-   *  SDPA scale, embedding/residual multipliers, logits / logits_scaling. */
-  granite: (raw) => {
-    const a = common(raw, "granite");
-    a.attnScale = raw.attention_multiplier;
-    a.embedMultiplier = raw.embedding_multiplier;
-    a.residualMultiplier = raw.residual_multiplier;
-    a.logitsDivisor = raw.logits_scaling;
-    a.qkvBias = raw.attention_bias;
-    a.oBias = raw.attention_bias;
-    a.mlpBias = raw.mlp_bias;
-    a.ropeTraditional = false; // granite.py hardcodes False
-    return a;
-  },
-
-  /** starcoder2: LayerNorm (bias affine), bias on ALL projections, plain
-   *  c_fc→gelu(precise)→c_proj MLP, theta 1e5. */
-  starcoder2: (raw) => {
-    const a = common(raw, "starcoder2");
-    a.normEps = raw.norm_epsilon ?? 1e-5;
-    a.vocabSize = raw.vocab_size ?? 49152;
-    a.ropeTheta = raw.rope_theta ?? 100000;
-    a.norm = "layernorm";
-    a.mlp = "gelu_mlp";
-    a.qkvBias = true;
-    a.oBias = true;
-    a.mlpBias = true;
-    a.rope = "plain";
-    a.ropeScaling = null;
-    return a;
-  },
-};
-
-/** The declared generic support surface (post-remap model_types). */
-export const GENERIC_MODEL_TYPES: ReadonlySet<string> = new Set(Object.keys(ARCHS));
-
-/** Descriptor lookup for a loaded config; null ⇒ not generically supported. */
-export function genericArgsFor(config: ModelConfig): UniversalArgs | null {
-  const arch = remapModelType(config.modelType);
-  const parse = ARCHS[arch];
-  if (!parse) return null;
-  const raw = (config.raw.text_config ?? config.raw) as Raw;
-  const args = parse(raw);
-  // keep the ORIGINAL model_type visible in errors/labels
-  args.modelType = config.modelType;
-  return args;
-}