mlx-bun 0.0.8 → 0.0.9

package/src/model/steel-linear-kernel.ts

@@ -0,0 +1,96 @@
+// Plain quantized linear y = x @ Wᵀ as OUR dispatch, composed from the steel
+// building blocks already in src/train/steel-qmm-header.ts (BlockMMA + Block/
+// QuantizedBlockLoader — MLX's verbatim GEMM, the same code qmm_t_impl wraps).
+// This is the projection primitive for the controlled forward (q/k/v/o/down):
+// we own the dispatch and the eval boundary instead of emitting an ops.* node
+// into MLX's lazy graph. N (out) and H (in/K) are baked as template literals;
+// M comes from the x_shape helper. Output bit-matches the BlockMMA store epilogue
+// (TransformNone → cast to T), the same store qmm_t_impl uses.
+
+import { MlxArray } from "../mlx/array";
+import { MetalKernel } from "../mlx/metal-kernel";
+import { STEEL_QMM_HEADER } from "../train/steel-qmm-header";
+import * as ops from "../mlx/ops";
+import type { QuantSpec } from "../mlx/ops";
+
+const BM = 32, BN = 32, BK = 32, WM = 2, WN = 2;
+const TGN = WM * WN * 32;
+
+// One BM×BN output tile per threadgroup (grid.y = N-tiles, grid.z = M-tiles), the
+// K=H reduction looped inside — the qmm_t_impl body, dims via templateInts.
+const SOURCE = String.raw`
+  const uint M=x_shape[0]; constexpr int N=N_T, H=H_T;
+  constexpr int gs=GS_T, bts=BITS_T;
+  constexpr int pf=get_pack_factor<bts,8>(), bpp=get_bytes_per_pack<bts>();
+  constexpr int BKp=BK_T+16/sizeof(T);
+  threadgroup T Xs[BM_T*BKp];
+  threadgroup T Ws[BN_T*BKp];
+  using mma_t=mlx::steel::BlockMMA<T,T,BM_T,BN_T,BK_T,WM_T,WN_T,false,true,BKp,BKp>;
+  using lx_t=mlx::steel::BlockLoader<T,BM_T,BK_T,BKp,1,WM_T*WN_T*SIMD_SIZE>;
+  using lw_t=QuantizedBlockLoader<T,BN_T,BK_T,BKp,1,WM_T*WN_T*SIMD_SIZE,gs,bts>;
+  const uint sg=simdgroup_index_in_threadgroup, sl=thread_index_in_simdgroup;
+  const int Hh=(int)H, K_w=Hh*bpp/pf, K_g=Hh/gs;
+  const int m0=(int)threadgroup_position_in_grid.z*BM_T;
+  const int n0=(int)threadgroup_position_in_grid.y*BN_T;
+  const short num_els=(short)min(BM_T,(int)M-m0);
+  const short num_outs=(short)min(BN_T,(int)N-n0);
+
+  mma_t mma_op(sg,sl);
+  lx_t  lx(x + (int64_t)m0*Hh, Hh, Xs, sg, sl);
+  lw_t  lw((const device uint8_t*)w + (int64_t)n0*K_w, scales + (int64_t)n0*K_g, biases + (int64_t)n0*K_g, Hh, Ws, sg, sl);
+
+  for (int k=0; k<Hh; k+=BK_T){
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (num_els<BM_T) lx.load_safe(short2(BK_T,num_els)); else lx.load_unsafe();
+    if (num_outs<BN_T) lw.load_safe(short2(BK_T,num_outs)); else lw.load_unsafe();
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    mma_op.mma(Xs,Ws);
+    lx.next(); lw.next();
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  // Store straight to device y[m0.., n0..] via the BlockMMA epilogue (float accum
+  // → T, the SAME rounding MLX's quantized_matmul applies).
+  device T* yp = y + (int64_t)m0*(int64_t)N + n0;
+  if (num_els<BM_T || num_outs<BN_T) mma_op.store_result_safe(yp, N, short2(num_outs,num_els));
+  else mma_op.store_result(yp, N);
+`;
+
+const _kernels = new Map<string, MetalKernel>();
+function getKernel(gs: number, bits: number): MetalKernel {
+  const key = `${gs}:${bits}`;
+  let k = _kernels.get(key);
+  if (!k) {
+    k = new MetalKernel({
+      name: "mlx_bun_steel_linear",
+      inputNames: ["x", "w", "scales", "biases"],
+      outputNames: ["y"],
+      source: SOURCE,
+      header:
+        `${STEEL_QMM_HEADER}\n` +
+        `#define GS_T ${gs}\n#define BITS_T ${bits}\n` +
+        `#define BM_T ${BM}\n#define BN_T ${BN}\n#define BK_T ${BK}\n` +
+        `#define WM_T ${WM}\n#define WN_T ${WN}\n#define TGN_T ${TGN}\n`,
+      ensureRowContiguous: true,
+    });
+    _kernels.set(key, k);
+  }
+  return k;
+}
+
+/** y = x @ Wᵀ, quantized, as our own dispatch. `x` is [M, H] (bf16/f16); `w` is
+ *  the packed weight [N, H*bits/32] with per-group scales/biases [N, H/groupSize].
+ *  Returns [M, N] in x.dtype. H % BK == 0 and groupSize % BK == 0 required. */
+export function steelLinear(
+  x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray, spec: QuantSpec,
+): MlxArray {
+  const M = x.shape[0]!, H = x.shape[1]!, N = w.shape[0]!;
+  const [y] = getKernel(spec.groupSize, spec.bits).apply([x, w, scales, biases], {
+    outputShapeFn: (ins) => [{ shape: [ins[0]!.shape[0]!, N], dtype: x.dtype }],
+    grid: (ins) => [TGN, Math.ceil(N / BN), Math.ceil(ins[0]!.shape[0]! / BM)],
+    threadGroup: [TGN, 1, 1],
+    templateInts: { N_T: N, H_T: H },
+    templateDtypes: { T: x.dtype },
+  });
+  void M;
+  return y!;
+}

package/src/model/support.ts

@@ -1,4 +1,5 @@
 import type { ModelConfig } from "../config";
+import { GENERIC_MODEL_TYPES, genericArgsFor, remapModelType } from "./universal/archs";
 
 export function isMiniCPM5Config(config: ModelConfig): boolean {
   const t = config.text;
@@ -22,6 +23,13 @@ export function isQwen35Config(config: ModelConfig): boolean {
     config.text.fullAttentionInterval > 0;
 }
 
+/** Plain Qwen3 (model_type `qwen3`, Qwen3ForCausalLM) — a standard dense decoder
+ *  with per-head q/k norm and tied embeddings. Used here as the text-embedding
+ *  backbone (mlx-community Qwen3-Embedding-*). Distinct from the qwen3_5 hybrid. */
+export function isQwen3Config(config: ModelConfig): boolean {
+  return config.modelType === "qwen3";
+}
+
 /** Speculative-decoding drafters (e.g. `gemma4_assistant`) are companion
  *  artifacts to a target model — Q-only, centroid-head, no standalone LM
  *  head. They are never servable/selectable on their own (the spec path
@@ -38,17 +46,39 @@ export function isDiffusionGemmaConfig(config: ModelConfig): boolean {
   return config.modelType === "diffusion_gemma";
 }
 
-export function isSupportedModelRecord(modelType: string, repoId = ""): boolean {
-  if (isDrafterModelType(modelType)) return false;
-  if (modelType.startsWith("gemma4")) return true;
-  if (modelType === "diffusion_gemma") return true;
+/** Support tier for a registry record (docs/design/generic-model-support.md):
+ *  "targeted" = dedicated/generated forward + L2/L3 paths;
+ *  "generic"  = the Tier-0 universal module (L1 monolith only);
+ *  null       = unsupported. Generic never shadows targeted. */
+export function supportTier(modelType: string, repoId = ""): "targeted" | "generic" | null {
+  if (isDrafterModelType(modelType)) return null;
+  if (modelType.startsWith("gemma4")) return "targeted";
+  if (modelType === "diffusion_gemma") return "targeted";
   // qwen3_5 / qwen3_5_text (dense hybrid). The MoE variant is a separate type.
-  if (modelType === "qwen3_5" || modelType === "qwen3_5_text") return true;
-  return modelType === "llama" && repoId.toLowerCase().includes("minicpm5-1b-optiq-4bit");
+  if (modelType === "qwen3_5" || modelType === "qwen3_5_text") return "targeted";
+  if (modelType === "qwen3") return "targeted"; // plain Qwen3 (embedding backbone)
+  if (modelType === "llama" && repoId.toLowerCase().includes("minicpm5-1b-optiq-4bit"))
+    return "targeted";
+  // Tier-0 generic fallback: the universal-dense descriptor table.
+  if (GENERIC_MODEL_TYPES.has(remapModelType(modelType))) return "generic";
+  return null;
+}
+
+export function isSupportedModelRecord(modelType: string, repoId = ""): boolean {
+  return supportTier(modelType, repoId) !== null;
 }
 
 export function isSupportedModelConfig(config: ModelConfig): boolean {
   if (isDrafterModelType(config.modelType)) return false;
-  return config.modelType.startsWith("gemma4") || isMiniCPM5Config(config) ||
-    isQwen35Config(config) || isDiffusionGemmaConfig(config);
+  if (
+    config.modelType.startsWith("gemma4") || isMiniCPM5Config(config) ||
+    isQwen35Config(config) || isQwen3Config(config) || isDiffusionGemmaConfig(config)
+  ) return true;
+  // Generic tier: an arch descriptor exists for this model_type (and its
+  // config parses — a malformed config is unsupported, not a crash).
+  try {
+    return genericArgsFor(config) !== null;
+  } catch {
+    return false;
+  }
 }

package/src/model/universal/archs.ts

@@ -0,0 +1,332 @@
+// 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;
+}