Skip to content

refactor(zoo): move the Moonshine and privacy-filter state dict hooks onto the shared weight helpers - #78

Merged
chmjkb merged 12 commits into
mainfrom
@chmjkb/hooks-moonshine-privacy-filter
Oct 30, 2026
Merged

chmjkb merged 12 commits into
mainfrom
@chmjkb/hooks-moonshine-privacy-filter

Conversation

@chmjkb

@chmjkb chmjkb commented Oct 27, 2026

Copy link
Copy Markdown
Collaborator

Description

Moves the Moonshine, Moonshine streaming and privacy-filter state dict hooks onto the shared weight helpers, and adds reshapeWeight.

  • reshapeWeight: changes the shape only. No copy, no conversion; the tensor keeps its dtype and byteLength. Mismatched element counts throw CHECKPOINT_MISMATCH.
  • Moonshine (+16 / −61): q/k/v fusion and cross-attention q/kv use concatRows + scaleWeight. stackRows is gone.
  • Moonshine streaming (+7 / −8): gets the same q/k/v change through the shared fuseQkv. gamma + 1 stays f32 on purpose (rounding it to f16 would lose precision on devices without shader-f16), as do exp(log_k) and the all-ones CMVN weight.
  • privacy-filter (+17 / −68): its private lazyF32/concatRows/column split are gone. Gate/up is splitCols + reshapeWeight, down_proj is reshapeWeight. The quantized path is untouched: quantized files skip this hook.
  • The F16 safetensors builder used by the weight and hook tests is now one shared helper, test/core/weights/checkpoint.ts.

Names, shapes and byteLengths reaching the models are unchanged, so the hosted files load as before. F16 checkpoints now stay f16 through these hooks.

Introduces a breaking change?

  • Yes
  • No

Type of change

  • Bug fix (change which fixes an issue)
  • New feature (change which adds functionality)
  • Documentation update (improves or adds clarity to existing documentation)
  • Performance (kernel tuning, graph building, anything related to performance)
  • Other (chores, tests, code style improvements etc.)

Testing instructions

  1. pnpm typecheck && pnpm lint && pnpm --filter runntime test:unit (923 pass)
  2. New tests: test/zoo/{moonshine,moonshine-streaming,privacy-filter}/stateDictHooks.test.ts, each on F32 and F16 synthetic checkpoints, values against hand-worked math. reshapeWeight cases in test/core/weights/transforms.test.ts.
  3. pnpm --filter runntime test:gpu: 1439 passed, 16 expected failures.

Related issues

Part of #70. Stacked on #76.

Checklist

  • I have performed a self-review of my code
  • I have commented my code, particularly in hard-to-understand areas
  • I have updated the documentation accordingly
  • My changes generate no new warnings

Additional notes

The only hosted privacy-filter file is int8, so its changed float path only runs when someone loads the official HF checkpoint.

@chmjkb chmjkb added core Engine: graph ops, dispatch, buffer pool, weight loading zoo Task runners and model loaders in src/zoo labels Oct 27, 2026
@chmjkb
chmjkb force-pushed the @chmjkb/hooks-moonshine-privacy-filter branch from a6e580c to c2f4e66 Compare October 28, 2026 09:16
@chmjkb
chmjkb force-pushed the @chmjkb/linear-f16-load branch 2 times, most recently from 756048a to b4448d4 Compare October 28, 2026 10:49
@chmjkb
chmjkb force-pushed the @chmjkb/hooks-moonshine-privacy-filter branch from c2f4e66 to 1930710 Compare October 28, 2026 10:49
@chmjkb
chmjkb force-pushed the @chmjkb/linear-f16-load branch from b4448d4 to 1505b3e Compare October 28, 2026 10:57
@chmjkb
chmjkb force-pushed the @chmjkb/hooks-moonshine-privacy-filter branch from 1930710 to e4659b0 Compare October 28, 2026 10:58
@chmjkb
chmjkb force-pushed the @chmjkb/linear-f16-load branch from 0e802e6 to 75b704f Compare October 30, 2026 11:26
@chmjkb
chmjkb force-pushed the @chmjkb/hooks-moonshine-privacy-filter branch from e4659b0 to 52294eb Compare October 30, 2026 11:26
Base automatically changed from @chmjkb/linear-f16-load to main October 30, 2026 11:40
@chmjkb
chmjkb force-pushed the @chmjkb/hooks-moonshine-privacy-filter branch from 52294eb to 5ada244 Compare October 30, 2026 12:47
@chmjkb
chmjkb merged commit b204178 into main Oct 30, 2026
2 checks passed
@chmjkb
chmjkb deleted the @chmjkb/hooks-moonshine-privacy-filter branch October 30, 2026 12:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Engine: graph ops, dispatch, buffer pool, weight loading zoo Task runners and model loaders in src/zoo

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants