@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/glm52/moe.ts

@@ -0,0 +1,223 @@
+// Exact GLM-5.2 noaux_tc routing and shared+routed MoE composition seams.
+//
+// The 256-way route decision is a deliberately bounded host correctness path:
+// it preserves Colibri's float32 sigmoid, correction-bias ranking, and strict
+// lower-ID tie behavior. Expert tensor arithmetic remains on MLX.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import {
+  buildExpertBatchUnion,
+  type ExpertUnionWave,
+} from "../../artifacts/experts/residency";
+import {
+  routeTrueTopKF32,
+  swiGluF32,
+  type Glm52Route,
+  type NumericVector,
+  type SwiGluWeights,
+} from "./reference";
+
+export interface Glm52MoeRouterConfig {
+  readonly topK: number;
+  readonly normalize: boolean;
+  readonly routedScale: number;
+}
+
+export interface Glm52RoutedExpertOutput<T> {
+  readonly expertId: number;
+  readonly output: T;
+}
+
+export interface Glm52MoeExpertJob {
+  readonly expertId: number;
+  readonly rows: Int32Array;
+  readonly ranks: Int32Array;
+  readonly weights: Float32Array;
+}
+
+export interface Glm52MoeBatchWave {
+  readonly jobs: readonly Glm52MoeExpertJob[];
+}
+
+export interface Glm52MoeBatchPlan {
+  readonly routes: readonly Glm52Route[];
+  readonly waves: readonly Glm52MoeBatchWave[];
+  readonly uniqueExperts: number;
+}
+
+function f32(value: number): number {
+  return Math.fround(value);
+}
+
+function shapeEquals(left: MlxArray, right: MlxArray): boolean {
+  const a = left.shape;
+  const b = right.shape;
+  return a.length === b.length && a.every((dimension, index) => dimension === b[index]);
+}
+
+/** Exact, bounded host route selection for the GLM noaux_tc router. */
+export function routeGlm52MoeF32(
+  logits: NumericVector,
+  correctionBias: NumericVector,
+  config: Glm52MoeRouterConfig,
+): Glm52Route {
+  return routeTrueTopKF32(
+    logits,
+    correctionBias,
+    config.topK,
+    config.normalize,
+    config.routedScale,
+  );
+}
+
+/**
+ * Route flattened [batch,tokens] rows in order, then form the exact first-seen
+ * expert union consumed by the global residency working bank.
+ */
+export function planGlm52MoeBatchF32(
+  logitsByRow: readonly NumericVector[],
+  correctionBias: NumericVector,
+  config: Glm52MoeRouterConfig,
+  maxUniquePerWave = 64,
+): Glm52MoeBatchPlan {
+  const routes = logitsByRow.map((logits) =>
+    routeGlm52MoeF32(logits, correctionBias, config));
+  const union: readonly ExpertUnionWave[] =
+    buildExpertBatchUnion(routes, maxUniquePerWave);
+  let uniqueExperts = 0;
+  const waves = union.map((wave) => ({
+    jobs: wave.entries.map((entry) => {
+      uniqueExperts++;
+      const rows = new Int32Array(entry.consumers.length);
+      const ranks = new Int32Array(entry.consumers.length);
+      const weights = new Float32Array(entry.consumers.length);
+      for (let index = 0; index < entry.consumers.length; index++) {
+        const consumer = entry.consumers[index]!;
+        rows[index] = consumer.row;
+        ranks[index] = consumer.rank;
+        weights[index] = routes[consumer.row]!.executionWeights[consumer.rank]!;
+      }
+      return {
+        expertId: entry.expertId,
+        rows,
+        ranks,
+        weights,
+      };
+    }),
+  }));
+  return { routes, waves, uniqueExperts };
+}
+
+function validateRouteOutputs<T>(
+  route: Glm52Route,
+  routed: ReadonlyArray<Glm52RoutedExpertOutput<T>>,
+): void {
+  if (routed.length !== route.indices.length)
+    throw new Error(`routed output count ${routed.length} != route count ${route.indices.length}`);
+  for (let rank = 0; rank < routed.length; rank++) {
+    if (routed[rank]!.expertId !== route.indices[rank]) {
+      throw new Error(
+        `routed output ${rank} expert ${routed[rank]!.expertId} != route ${route.indices[rank]}`,
+      );
+    }
+  }
+}
+
+/**
+ * Float32 oracle for the composition seam. Routed outputs are accumulated in
+ * route order with raw-sigmoid-derived weights; the shared output is added
+ * once, unweighted, after the routed sum.
+ */
+export function composeGlm52MoeOutputsF32(
+  route: Glm52Route,
+  routed: ReadonlyArray<Glm52RoutedExpertOutput<NumericVector>>,
+  shared: NumericVector | null,
+): Float32Array {
+  validateRouteOutputs(route, routed);
+  if (routed.length === 0 && shared === null)
+    throw new Error("GLM MoE composition needs a routed or shared output");
+
+  const width = routed[0]?.output.length ?? shared!.length;
+  if (!Number.isSafeInteger(width) || width < 1)
+    throw new Error("GLM MoE output width must be positive");
+  const out = new Float32Array(width);
+  for (let rank = 0; rank < routed.length; rank++) {
+    const value = routed[rank]!.output;
+    if (value.length !== width) throw new Error("routed expert output dimensions differ");
+    const weight = route.executionWeights[rank]!;
+    for (let index = 0; index < width; index++)
+      out[index] = f32(out[index]! + f32(weight * value[index]!));
+  }
+  if (shared !== null) {
+    if (shared.length !== width) throw new Error("shared expert output dimension differs");
+    for (let index = 0; index < width; index++)
+      out[index] = f32(out[index]! + shared[index]!);
+  }
+  return out;
+}
+
+/**
+ * MLX composition path. Inputs remain caller-owned; returned output is owned by
+ * the caller. The API intentionally accepts already-computed shared output so
+ * the model can issue it in parallel with expert residency/I/O.
+ */
+export function composeGlm52MoeOutputsMlx(
+  route: Glm52Route,
+  routed: ReadonlyArray<Glm52RoutedExpertOutput<MlxArray>>,
+  shared: MlxArray | null,
+): MlxArray {
+  validateRouteOutputs(route, routed);
+  if (routed.length === 0 && shared === null)
+    throw new Error("GLM MoE composition needs a routed or shared output");
+
+  const exemplar = routed[0]?.output ?? shared!;
+  for (const contribution of routed)
+    if (!shapeEquals(contribution.output, exemplar))
+      throw new Error("routed expert output dimensions differ");
+  if (shared !== null && !shapeEquals(shared, exemplar))
+    throw new Error("shared expert output dimension differs");
+
+  let out: MlxArray | null = null;
+  for (let rank = 0; rank < routed.length; rank++) {
+    const scaled = ops.mulScalar(routed[rank]!.output, route.executionWeights[rank]!);
+    if (out === null) {
+      out = scaled;
+    } else {
+      const next = ops.add(out, scaled);
+      out.dispose();
+      scaled.dispose();
+      out = next;
+    }
+  }
+  if (shared !== null) {
+    if (out === null) {
+      // Return a new owned graph value, never an alias of the caller's input.
+      out = ops.mulScalar(shared, 1);
+    } else {
+      const next = ops.add(out, shared);
+      out.dispose();
+      out = next;
+    }
+  }
+  return out!;
+}
+
+/**
+ * Tiny/reference execution seam: the caller resolves streamed experts by ID;
+ * shared weights are passed independently so their work can later be issued
+ * before routed expert loads complete.
+ */
+export function executeGlm52MoeReferenceF32(
+  input: NumericVector,
+  route: Glm52Route,
+  routedExpert: (expertId: number) => SwiGluWeights,
+  sharedExpert: SwiGluWeights | null,
+): Float32Array {
+  const routed = route.indices.map((expertId) => ({
+    expertId,
+    output: swiGluF32(input, routedExpert(expertId)),
+  }));
+  const shared = sharedExpert === null ? null : swiGluF32(input, sharedExpert);
+  return composeGlm52MoeOutputsF32(route, routed, shared);
+}

package/src/models/glm52/mtp.ts

@@ -0,0 +1,60 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { NativeMtpHead } from "../../contracts/mlx/drafter";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Glm52Config } from "../../artifacts/glm52-config";
+import type { Glm52WeightSource } from "../../contracts/mlx/glm52-weights";
+import type { LatentAttentionCache } from "../../contracts/mlx/cache";
+import { GLM52_PHASE_READS, rmsNormF32Mlx, type Glm52PhaseReads } from "./mla";
+
+/** What the MTP graph reads of the GLM-5.2 model it shares weights with. */
+export interface Glm52MtpHost {
+  readonly glmConfig: Glm52Config;
+  readonly weights: Glm52WeightSource;
+  logitsFromHidden(hidden: MlxArray): MlxArray;
+}
+
+/** The MTP decoder layer (the layer after the target's last). */
+export interface Glm52MtpLayer {
+  forwardAsync(input: MlxArray, cache: LatentAttentionCache, dsaState: null, reads: Glm52PhaseReads): Promise<MlxArray>;
+}
+
+/** Native MTP numerical graph. Tokens, anchor hidden and compressed state may
+ * have one or several rows; sampling and state retention belong to callers. */
+export class Glm52MtpGraph implements NativeMtpHead {
+  constructor(readonly model: Glm52MtpHost, readonly layer: Glm52MtpLayer) {}
+
+  get hiddenSize(): number { return this.model.glmConfig.hiddenSize; }
+
+  get cache(): NativeMtpHead["cache"] {
+    const config = this.model.glmConfig;
+    return { kvLoraRank: config.kvLoraRank, ropeHeadDim: config.qkRopeHeadDim, maxTokens: config.maxPositionEmbeddings };
+  }
+
+  /** One token per sequence (a draft step): the MTP layer reads its compressed state absorbed. */
+  decode(ids: MlxArray, hidden: MlxArray, cache: LatentAttentionCache): Promise<MlxArray> {
+    return this.#forward(ids, hidden, cache, GLM52_PHASE_READS.oneToken);
+  }
+
+  /** Two or more tokens per sequence (the accepted window a commit replays): reconstructed. */
+  window(ids: MlxArray, hidden: MlxArray, cache: LatentAttentionCache): Promise<MlxArray> {
+    return this.#forward(ids, hidden, cache, GLM52_PHASE_READS.window);
+  }
+
+  async #forward(ids: MlxArray, hidden: MlxArray, cache: LatentAttentionCache, reads: Glm52PhaseReads): Promise<MlxArray> {
+    const { glmConfig: config, weights } = this.model;
+    const prefix = `model.layers.${config.numHiddenLayers}`;
+    using embedded = weights.embedding(ids, "model.embed_tokens.weight", config.vocabSize, config.hiddenSize);
+    using embeddedNorm = rmsNormF32Mlx(embedded, weights.tensor(`${prefix}.enorm.weight`), config.rmsNormEps);
+    using hiddenNorm = rmsNormF32Mlx(hidden, weights.tensor(`${prefix}.hnorm.weight`), config.rmsNormEps);
+    using joined = ops.concatAxis([embeddedNorm, hiddenNorm], 2);
+    using projected = weights.linear(joined, `${prefix}.eh_proj.weight`, config.hiddenSize, 2 * config.hiddenSize);
+    return await this.layer.forwardAsync(projected, cache, null, reads);
+  }
+
+  project(hidden: MlxArray): MlxArray {
+    const { glmConfig: config, weights } = this.model;
+    using normalized = rmsNormF32Mlx(hidden,
+      weights.tensor(`model.layers.${config.numHiddenLayers}.shared_head.norm.weight`), config.rmsNormEps);
+    return this.model.logitsFromHidden(normalized);
+  }
+}