@tryhamster/gerbil 1.11.3 → 1.12.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +79 -0
- package/dist/{architectures-DmZMEFsA.mjs → architectures-BH_z6k9d.mjs} +14 -12
- package/dist/architectures-BH_z6k9d.mjs.map +1 -0
- package/dist/browser/index.d.ts.map +1 -1
- package/dist/browser/index.js +11 -0
- package/dist/browser/index.js.map +1 -1
- package/dist/cli.mjs +8 -8
- package/dist/cli.mjs.map +1 -1
- package/dist/frameworks/express.mjs +1 -1
- package/dist/frameworks/fastify.mjs +1 -1
- package/dist/frameworks/hono.mjs +1 -1
- package/dist/frameworks/next.d.mts +2 -2
- package/dist/frameworks/next.mjs +1 -1
- package/dist/frameworks/trpc.mjs +1 -1
- package/dist/gerbil-CZvoFo0T.mjs +4 -0
- package/dist/{gerbil-CYVmU8sQ.mjs → gerbil-CpS3P240.mjs} +13 -2
- package/dist/gerbil-CpS3P240.mjs.map +1 -0
- package/dist/{gerbil-Cu7YwBBz.d.mts → gerbil-DF009Waa.d.mts} +2 -2
- package/dist/{gerbil-Cu7YwBBz.d.mts.map → gerbil-DF009Waa.d.mts.map} +1 -1
- package/dist/gpu/hooks.d.mts +1 -1
- package/dist/gpu/index.d.mts +2 -2
- package/dist/gpu/index.mjs +4 -4
- package/dist/{gpu-EyIH7Qfc.mjs → gpu-DRFhiv4R.mjs} +836 -140
- package/dist/gpu-DRFhiv4R.mjs.map +1 -0
- package/dist/{index-DeKYzJgJ.d.mts → index-CNYoTRgr.d.mts} +1198 -550
- package/dist/index-CNYoTRgr.d.mts.map +1 -0
- package/dist/index.d.mts +3 -3
- package/dist/index.d.mts.map +1 -1
- package/dist/index.mjs +5 -5
- package/dist/index.mjs.map +1 -1
- package/dist/integrations/ai-sdk.mjs +1 -1
- package/dist/integrations/langchain.mjs +1 -1
- package/dist/integrations/llamaindex.mjs +1 -1
- package/dist/integrations/mcp.d.mts +2 -2
- package/dist/integrations/mcp.mjs +4 -4
- package/dist/{mcp-CAsD7eCj.mjs → mcp-DAbWO8VS.mjs} +3 -3
- package/dist/{mcp-CAsD7eCj.mjs.map → mcp-DAbWO8VS.mjs.map} +1 -1
- package/dist/{moonshine-stt-BXoZaHJE.mjs → moonshine-stt-B1kV5n1c.mjs} +5251 -1678
- package/dist/moonshine-stt-B1kV5n1c.mjs.map +1 -0
- package/dist/moonshine-stt-DA1WuiZb.mjs +4 -0
- package/dist/{one-liner-ppgw4jHH.mjs → one-liner-CmP9ktUn.mjs} +2 -2
- package/dist/{one-liner-ppgw4jHH.mjs.map → one-liner-CmP9ktUn.mjs.map} +1 -1
- package/dist/repl-17NVeIaJ.mjs +9 -0
- package/dist/skills/index.d.mts +2 -2
- package/dist/skills/index.mjs +3 -3
- package/dist/{skills-CQa1Gshd.mjs → skills-BDOEbHSx.mjs} +2 -2
- package/dist/{skills-CQa1Gshd.mjs.map → skills-BDOEbHSx.mjs.map} +1 -1
- package/dist/tune/index.d.mts.map +1 -1
- package/dist/tune/index.mjs +1 -1
- package/package.json +1 -1
- package/dist/architectures-DmZMEFsA.mjs.map +0 -1
- package/dist/gerbil-CTefAwKp.mjs +0 -4
- package/dist/gerbil-CYVmU8sQ.mjs.map +0 -1
- package/dist/gpu-EyIH7Qfc.mjs.map +0 -1
- package/dist/index-DeKYzJgJ.d.mts.map +0 -1
- package/dist/moonshine-stt-BXoZaHJE.mjs.map +0 -1
- package/dist/moonshine-stt-DZVnKgPO.mjs +0 -4
- package/dist/repl-C0Ew7_Z-.mjs +0 -9
|
@@ -814,524 +814,6 @@ declare function quantizeKaniBackbone(graph: ModelGraph, weights: Map<string, {
|
|
|
814
814
|
shape: number[];
|
|
815
815
|
}>): void;
|
|
816
816
|
//#endregion
|
|
817
|
-
//#region src/gpu/sampler.d.ts
|
|
818
|
-
/**
|
|
819
|
-
* CPU-side token sampling from logits.
|
|
820
|
-
*
|
|
821
|
-
* Applies temperature, top-k, and top-p (nucleus) filtering,
|
|
822
|
-
* then samples from the resulting probability distribution.
|
|
823
|
-
*
|
|
824
|
-
* Uses typed arrays and min-heap for zero-allocation top-K selection.
|
|
825
|
-
* For vocab_size ~152K with topK=50, this avoids creating 152K JS tuples.
|
|
826
|
-
*/
|
|
827
|
-
interface SamplingParams {
|
|
828
|
-
temperature?: number;
|
|
829
|
-
topK?: number;
|
|
830
|
-
topP?: number;
|
|
831
|
-
repetitionPenalty?: number;
|
|
832
|
-
}
|
|
833
|
-
//#endregion
|
|
834
|
-
//#region src/gpu/vision-preprocess.d.ts
|
|
835
|
-
/**
|
|
836
|
-
* Host-side vision preprocessing for the Qwen3.5 ViT.
|
|
837
|
-
*
|
|
838
|
-
* The learned position embeddings (bilinear-interpolated over the patch grid)
|
|
839
|
-
* and the 2D rotary cos/sin tables are functions of the image grid (t, h, w)
|
|
840
|
-
* ONLY — not of the model weights or pixel values. They are cheap to compute on
|
|
841
|
-
* the CPU and fed to the GPU graph as input activations, keeping the graph to
|
|
842
|
-
* weight-dependent math while staying byte-identical to HF transformers.
|
|
843
|
-
*
|
|
844
|
-
* Ports (verified against transformers 5.12 vision_utils.py + modeling_qwen3_5):
|
|
845
|
-
* - get_vision_bilinear_indices_and_weights → buildPosEmbeds()
|
|
846
|
-
* - get_vision_position_ids → buildPositionIds()
|
|
847
|
-
* - Qwen3_5VisionRotaryEmbedding → buildRotaryCosSin()
|
|
848
|
-
*
|
|
849
|
-
* Patch ordering: both the pos-embed gather and the position ids reorder patches
|
|
850
|
-
* into spatial_merge_size×spatial_merge_size groups, matching the HF image
|
|
851
|
-
* processor's output ordering, so the merger's [N,h]→[N/u,h*u] reshape lines up.
|
|
852
|
-
*/
|
|
853
|
-
interface VisionGridConfig {
|
|
854
|
-
hiddenSize: number;
|
|
855
|
-
numHeads: number;
|
|
856
|
-
numPositionEmbeddings: number;
|
|
857
|
-
spatialMergeSize: number;
|
|
858
|
-
ropeTheta?: number;
|
|
859
|
-
}
|
|
860
|
-
interface VisionPositionTensors {
|
|
861
|
-
/** [N, hidden_size] bilinear-interpolated learned position embeddings. */
|
|
862
|
-
posEmbeds: Float32Array;
|
|
863
|
-
/** [N, head_dim] rotary cosines. */
|
|
864
|
-
cos: Float32Array;
|
|
865
|
-
/** [N, head_dim] rotary sines. */
|
|
866
|
-
sin: Float32Array;
|
|
867
|
-
numPatches: number;
|
|
868
|
-
}
|
|
869
|
-
/**
|
|
870
|
-
* Build bilinear-interpolated learned position embeddings [N, hidden].
|
|
871
|
-
* posEmbedTable is the raw pos_embed.weight [num_position_embeddings, hidden].
|
|
872
|
-
*/
|
|
873
|
-
declare function buildPosEmbeds(gridTHW: [number, number, number], posEmbedTable: Float32Array, cfg: VisionGridConfig): Float32Array;
|
|
874
|
-
/**
|
|
875
|
-
* Build the reordered (row, col) position ids [N, 2] for rotary, matching
|
|
876
|
-
* get_vision_position_ids.
|
|
877
|
-
*/
|
|
878
|
-
declare function buildPositionIds(gridTHW: [number, number, number], merge: number): Int32Array;
|
|
879
|
-
/**
|
|
880
|
-
* Build rotary cos/sin tables [N, head_dim] from position ids, matching
|
|
881
|
-
* Qwen3_5VisionRotaryEmbedding + the cat((rotary, rotary)) in VisionModel.forward.
|
|
882
|
-
*
|
|
883
|
-
* rotary_pos_emb(position_ids) = (position_ids[..,None] * inv_freq).flatten(1)
|
|
884
|
-
* where inv_freq has length (head_dim/2)/2 = head_dim/4, computed over dim=head_dim/2.
|
|
885
|
-
* For each token the two position components (h, w) each produce head_dim/4 freqs,
|
|
886
|
-
* concatenated → head_dim/2, then duplicated → head_dim for cos/sin.
|
|
887
|
-
*/
|
|
888
|
-
declare function buildRotaryCosSin(positionIds: Int32Array, headDim: number, theta?: number): {
|
|
889
|
-
cos: Float32Array;
|
|
890
|
-
sin: Float32Array;
|
|
891
|
-
numPatches: number;
|
|
892
|
-
};
|
|
893
|
-
/**
|
|
894
|
-
* Build all host position tensors for a single image grid in one call.
|
|
895
|
-
*/
|
|
896
|
-
declare function buildVisionPositionTensors(gridTHW: [number, number, number], posEmbedTable: Float32Array, cfg: VisionGridConfig): VisionPositionTensors;
|
|
897
|
-
interface Gemma4VisionGridConfig {
|
|
898
|
-
hiddenSize: number;
|
|
899
|
-
numHeads: number;
|
|
900
|
-
headDim: number;
|
|
901
|
-
ropeTheta: number;
|
|
902
|
-
poolingKernelSize: number;
|
|
903
|
-
}
|
|
904
|
-
interface Gemma4VisionPositionTensors {
|
|
905
|
-
/** [N, hidden] axial position embeddings (table[0][x] + table[1][y]). */
|
|
906
|
-
posEmbeds: Float32Array;
|
|
907
|
-
/** [N, headDim] axial rotary cosines. */
|
|
908
|
-
cos: Float32Array;
|
|
909
|
-
/** [N, headDim] axial rotary sines. */
|
|
910
|
-
sin: Float32Array;
|
|
911
|
-
/** [Np, N] average-pooling matrix (1/k² in-cell, 0 elsewhere). */
|
|
912
|
-
poolMatrix: Float32Array;
|
|
913
|
-
/** number of patches N (= gridH*gridW). */
|
|
914
|
-
numPatches: number;
|
|
915
|
-
/** number of pooled (soft) tokens Np (= ceil(gridH/k)*ceil(gridW/k)). */
|
|
916
|
-
numPooled: number;
|
|
917
|
-
}
|
|
918
|
-
/**
|
|
919
|
-
* Build axial learned position embeddings [N, hidden] from the [2, posSize, hidden]
|
|
920
|
-
* table: pos[p] = table[0][x_p] + table[1][y_p]. Direct lookup, no interpolation
|
|
921
|
-
* (HF F.embedding on clamped positions).
|
|
922
|
-
*/
|
|
923
|
-
declare function buildGemma4PosEmbeds(gridH: number, gridW: number, posEmbedTable: Float32Array,
|
|
924
|
-
// [2, posSize, hidden] flattened
|
|
925
|
-
hidden: number, posSize: number): Float32Array;
|
|
926
|
-
/**
|
|
927
|
-
* Build the 2D axial rotary cos/sin tables [N, headDim].
|
|
928
|
-
* spatial_dim = headDim / 2; inv_freq[j] = 1/theta^((2j)/spatial_dim), j in [0, spatial_dim/2)
|
|
929
|
-
* per spatial dim: f = pos * inv_freq (spatial_dim/2 values); emb = cat(f, f) (spatial_dim values)
|
|
930
|
-
* cos/sin = cat([emb_x, emb_y]) → headDim values, layout [fx,fx,fy,fy].
|
|
931
|
-
* Applied with the global-half rotate_half kernel (ApplyRotaryEmb), which computes
|
|
932
|
-
* out = x*cos + rotate_half(x)*sin element-wise — exact for this layout.
|
|
933
|
-
*/
|
|
934
|
-
declare function buildGemma4RotaryCosSin(gridH: number, gridW: number, headDim: number, theta: number): {
|
|
935
|
-
cos: Float32Array;
|
|
936
|
-
sin: Float32Array;
|
|
937
|
-
};
|
|
938
|
-
/**
|
|
939
|
-
* Build the [Np, N] average-pooling matrix for k×k spatial pooling over the real
|
|
940
|
-
* (unpadded) grid, matching modeling_gemma4's kernel_idxs/one_hot pooling:
|
|
941
|
-
* cell(p) = floor(x_p/k) + ceil(gridW/k) * floor(y_p/k)
|
|
942
|
-
* poolMatrix[cell, p] = 1/k² (so pooled = poolMatrix @ hidden = mean over the k×k block)
|
|
943
|
-
* Np = ceil(gridH/k) * ceil(gridW/k). Each pooled cell averages exactly the patches
|
|
944
|
-
* that fall in it (edge cells with fewer than k² patches still divide by k², matching
|
|
945
|
-
* HF's fixed 1/k² normalization).
|
|
946
|
-
*/
|
|
947
|
-
declare function buildGemma4PoolMatrix(gridH: number, gridW: number, k: number): {
|
|
948
|
-
poolMatrix: Float32Array;
|
|
949
|
-
numPooled: number;
|
|
950
|
-
};
|
|
951
|
-
/**
|
|
952
|
-
* Build all Gemma 4 vision host tensors for one image grid in one call.
|
|
953
|
-
* `posEmbedTable` is the raw [2, posSize, hidden] flattened table.
|
|
954
|
-
*/
|
|
955
|
-
declare function buildGemma4VisionPositionTensors(gridH: number, gridW: number, posEmbedTable: Float32Array, posSize: number, cfg: Gemma4VisionGridConfig): Gemma4VisionPositionTensors;
|
|
956
|
-
/** Gemma 4 image processor config (from processor_config.json). */
|
|
957
|
-
declare const GEMMA4_IMAGE_PROCESSOR: ImageProcessorConfig;
|
|
958
|
-
interface Gemma4PreprocessedImage {
|
|
959
|
-
/** Flattened patches [N, 3·patch²] row-major (row-major patch grid). */
|
|
960
|
-
patches: Float32Array;
|
|
961
|
-
/** Patch grid (gridH, gridW). */
|
|
962
|
-
gridHW: [number, number];
|
|
963
|
-
}
|
|
964
|
-
/**
|
|
965
|
-
* Preprocess a decoded RGB image for the Gemma 4 ViT: aspect-preserving resize so
|
|
966
|
-
* the patch grid is ≤ max_soft_tokens·k² patches and H,W divisible by k·patch,
|
|
967
|
-
* rescale ×1/255 (no normalize), patchify row-major into [N, 3·16·16].
|
|
968
|
-
*
|
|
969
|
-
* @param pixels row-major HWC RGB (0..255), length width*height*3.
|
|
970
|
-
*/
|
|
971
|
-
declare function preprocessImageGemma4(pixels: Float32Array | Uint8ClampedArray | Uint8Array, width: number, height: number, maxSoftTokens?: number, poolingKernelSize?: number, patchSize?: number): Gemma4PreprocessedImage;
|
|
972
|
-
interface ImageProcessorConfig {
|
|
973
|
-
patchSize: number;
|
|
974
|
-
temporalPatchSize: number;
|
|
975
|
-
mergeSize: number;
|
|
976
|
-
imageMean: [number, number, number];
|
|
977
|
-
imageStd: [number, number, number];
|
|
978
|
-
/** rescale factor applied to raw 0..255 pixels before normalization (1/255). */
|
|
979
|
-
rescaleFactor: number;
|
|
980
|
-
/** min total pixels after resize (shortest_edge). */
|
|
981
|
-
minPixels: number;
|
|
982
|
-
/** max total pixels after resize (longest_edge). */
|
|
983
|
-
maxPixels: number;
|
|
984
|
-
}
|
|
985
|
-
declare const QWEN3_5_IMAGE_PROCESSOR: ImageProcessorConfig;
|
|
986
|
-
interface PreprocessedImage {
|
|
987
|
-
/** Flattened patches [N, 1536] in the spatial-merge order encodeImage expects. */
|
|
988
|
-
patches: Float32Array;
|
|
989
|
-
/** (t, h, w) patch grid. t=1 for a single image, h/w in patch units. */
|
|
990
|
-
gridTHW: [number, number, number];
|
|
991
|
-
}
|
|
992
|
-
/**
|
|
993
|
-
* Qwen2-VL smart-resize: round H and W to multiples of factor=patch*merge,
|
|
994
|
-
* keeping aspect ratio and clamping the total pixel budget to [minPixels, maxPixels].
|
|
995
|
-
* Matches transformers.models.qwen2_vl.image_processing.smart_resize.
|
|
996
|
-
*/
|
|
997
|
-
declare function smartResize(height: number, width: number, factor: number, minPixels: number, maxPixels: number): [number, number];
|
|
998
|
-
/**
|
|
999
|
-
* Preprocess a decoded RGB image into the [N, 1536] patch tensor + grid_thw that
|
|
1000
|
-
* `encodeImage()` expects, matching the HF Qwen2-VL image processor:
|
|
1001
|
-
* smart_resize → rescale (×1/255) → normalize → temporal-pair (×temporal_patch_size)
|
|
1002
|
-
* → patchify into spatial_merge×spatial_merge blocks → flatten to [N, C·T·P·P].
|
|
1003
|
-
*
|
|
1004
|
-
* @param pixels row-major HWC RGB, length width*height*3. Values 0..255 (default)
|
|
1005
|
-
* or already 0..1 if `rescaled` is true.
|
|
1006
|
-
* @param width source pixel width
|
|
1007
|
-
* @param height source pixel height
|
|
1008
|
-
*/
|
|
1009
|
-
declare function preprocessImage(pixels: Float32Array | Uint8ClampedArray | Uint8Array, width: number, height: number, cfg?: ImageProcessorConfig, rescaled?: boolean): PreprocessedImage;
|
|
1010
|
-
declare function buildMRoPEPositionIds(inputIds: Int32Array | Uint32Array | number[], imageGrids: Array<[number, number, number]>, imageTokenId: number, mergeSize: number): Int32Array;
|
|
1011
|
-
/**
|
|
1012
|
-
* Per-pair frequency→dimension assignment for interleaved M-RoPE, matching
|
|
1013
|
-
* Qwen3_5TextRotaryEmbedding.apply_interleaved_mrope. For pair index i in
|
|
1014
|
-
* [0, sum(section)) the position component is section-cyclic: T,H,W,T,H,W,...
|
|
1015
|
-
* but each component capped at its section count. Returns an array of length
|
|
1016
|
-
* (rope_dim/2) with values 0=T, 1=H, 2=W.
|
|
1017
|
-
*/
|
|
1018
|
-
declare function mropeFreqDims(mropeSection: [number, number, number]): Int32Array;
|
|
1019
|
-
/**
|
|
1020
|
-
* Build the interleaved-M-RoPE cos/sin tables [seq, rope_dim] from 3D position
|
|
1021
|
-
* ids, matching Qwen3_5TextRotaryEmbedding.forward:
|
|
1022
|
-
* freqs[d][i] = pos[d] * inv_freq[i], inv_freq[i] = 1/theta^(2i/rope_dim)
|
|
1023
|
-
* freq[i] picks component mropeFreqDims[i]; emb = cat(freqs, freqs).
|
|
1024
|
-
* cos/sin have length seq*rope_dim. For text-only (all 3 pos rows equal) this
|
|
1025
|
-
* reduces exactly to standard 1D partial RoPE.
|
|
1026
|
-
*
|
|
1027
|
-
* @param positionIds3 [3, seq] as produced by buildMRoPEPositionIds.
|
|
1028
|
-
* @param ropeDim number of rotated dims per head (head_dim * partial_factor).
|
|
1029
|
-
*/
|
|
1030
|
-
declare function buildMRoPECosSin(positionIds3: Int32Array, seq: number, ropeDim: number, theta: number, mropeSection: [number, number, number]): {
|
|
1031
|
-
cos: Float32Array;
|
|
1032
|
-
sin: Float32Array;
|
|
1033
|
-
};
|
|
1034
|
-
//#endregion
|
|
1035
|
-
//#region src/gpu/defaults.d.ts
|
|
1036
|
-
/**
|
|
1037
|
-
* Default model per capability. Kept in its own tiny module (no heavy imports)
|
|
1038
|
-
* so the React hooks can resolve defaults without statically pulling in the GPU
|
|
1039
|
-
* engine — they import the engine dynamically to stay light.
|
|
1040
|
-
*/
|
|
1041
|
-
declare const DEFAULT_MODELS: {
|
|
1042
|
-
/** Text generation (also the vision-capable checkpoint). */
|
|
1043
|
-
readonly text: "mlx-community/Qwen3.5-0.8B-4bit";
|
|
1044
|
-
/** Image understanding — same checkpoint, vision tower built on demand. */
|
|
1045
|
-
readonly vision: "mlx-community/Qwen3.5-0.8B-4bit";
|
|
1046
|
-
/** Text embeddings. */
|
|
1047
|
-
readonly embedding: "mlx-community/embeddinggemma-300m-4bit";
|
|
1048
|
-
/** Text-to-speech (Apache-2.0; LFM2-350M backbone + NanoCodec). Small/fast default. */
|
|
1049
|
-
readonly tts: "nineninesix/kani-tts-450m-0.2-ft";
|
|
1050
|
-
/**
|
|
1051
|
-
* Selectable many-voices TTS (OuteTTS-1.0-0.6B: Qwen3-0.6B codec-LM + DAC.speech).
|
|
1052
|
-
* Pass this as the `repo`/`model` to engine.speak()/useTTS to route to OuteTTS.
|
|
1053
|
-
*/
|
|
1054
|
-
readonly ttsOute: "OuteAI/OuteTTS-1.0-0.6B";
|
|
1055
|
-
/**
|
|
1056
|
-
* Selectable describe-your-voice TTS (Parler-TTS-mini-v1: Flan-T5 encoder +
|
|
1057
|
-
* cross-attending decoder LM + dac_44khz). Pass this as the `repo`/`model` to
|
|
1058
|
-
* engine.speak()/useTTS with a `describeVoice` natural-language description.
|
|
1059
|
-
*/
|
|
1060
|
-
readonly ttsParler: "parler-tts/parler-tts-mini-v1";
|
|
1061
|
-
/** Speech-to-text. */
|
|
1062
|
-
readonly stt: "UsefulSensors/moonshine-base";
|
|
1063
|
-
};
|
|
1064
|
-
declare const OUTETTS_ASSETS: {
|
|
1065
|
-
/** Folded DAC.speech.v1.0 encoder+decoder weights (prep-dac-speech.py output). */
|
|
1066
|
-
readonly codecWeightsURL: "https://pub-0522ecec25fc412bb35a523f3e3f63c6.r2.dev/tryhamster/gerbil-outetts-assets/dac-speech-folded.safetensors";
|
|
1067
|
-
/** Base URL for the preset speaker JSONs (`{name}.json`). */
|
|
1068
|
-
readonly speakersBaseURL: "https://pub-0522ecec25fc412bb35a523f3e3f63c6.r2.dev/tryhamster/gerbil-outetts-assets/speakers";
|
|
1069
|
-
/** The default preset voice. */
|
|
1070
|
-
readonly defaultVoice: "en-female-1-neutral";
|
|
1071
|
-
};
|
|
1072
|
-
/**
|
|
1073
|
-
* Bundled OuteTTS preset voices.
|
|
1074
|
-
*
|
|
1075
|
-
* Provenance (honest labelling):
|
|
1076
|
-
* - `en-female-1-neutral` is the ONE AUTHENTIC OuteAI voice — the OuteTTS v3 bundled
|
|
1077
|
-
* default speaker (default_speakers/en-female-1-neutral), copied verbatim.
|
|
1078
|
-
* - `gerbil-voice-2..6` are DERIVED voices: distinct timbres produced by running
|
|
1079
|
-
* pitch/tempo-varied renderings of that authentic clip back through the REAL
|
|
1080
|
-
* DAC.speech encode + prosody pipeline (scripts/engine/stage-outetts-assets.py).
|
|
1081
|
-
* They are genuine DAC-encoded OuteSpeakers, but they are NOT OuteAI's own
|
|
1082
|
-
* `en-male-*` voices, so they are not labelled as such. (Earlier builds mislabelled
|
|
1083
|
-
* them as `en-male-2-neutral` etc.; renamed for honesty.)
|
|
1084
|
-
*/
|
|
1085
|
-
declare const OUTETTS_PRESET_VOICES: readonly ["en-female-1-neutral", "gerbil-voice-2", "gerbil-voice-3", "gerbil-voice-4", "gerbil-voice-5", "gerbil-voice-6"];
|
|
1086
|
-
/**
|
|
1087
|
-
* True when `repo`/`architecture` identify a text-to-speech checkpoint that must
|
|
1088
|
-
* run through a dedicated native TTS engine (KaniTTS, OuteTTS, or Parler-TTS)
|
|
1089
|
-
* rather than the single-graph base text loader. Mirrors the routing used by
|
|
1090
|
-
* `WebGPUEngine.speak()` / `Gerbil.ensureNativeTTSEngine`:
|
|
1091
|
-
* - Parler → encoder-decoder (`ParlerTTSForConditionalGeneration`)
|
|
1092
|
-
* - OuteTTS → Qwen3-0.6B codec-LM
|
|
1093
|
-
* - Kani → conv-backbone codec-LM (`KaniTTS2ForCausalLM`, or the Apache
|
|
1094
|
-
* default which is a plain `Lfm2ForCausalLM` — hence the repo-name
|
|
1095
|
-
* gate; the architecture string alone does not identify it).
|
|
1096
|
-
* At `create()` time only the repo is known (the architecture is discovered
|
|
1097
|
-
* during load), so the repo-name signals below cover all three DEFAULT_MODELS
|
|
1098
|
-
* and conventionally-named TTS repos; the architecture arg is honored when a
|
|
1099
|
-
* caller already has it.
|
|
1100
|
-
*/
|
|
1101
|
-
declare function isTTSRepo(repo?: string, architecture?: string): boolean;
|
|
1102
|
-
/** Resolve the model repo for a set of options, falling back to the defaults. */
|
|
1103
|
-
declare function resolveDefaultRepo(opts: {
|
|
1104
|
-
repo?: string;
|
|
1105
|
-
embedding?: boolean;
|
|
1106
|
-
enableVision?: boolean;
|
|
1107
|
-
}): string;
|
|
1108
|
-
//#endregion
|
|
1109
|
-
//#region src/gpu/architectures/gemma4_vision.d.ts
|
|
1110
|
-
interface Gemma4VisionGraphInfo {
|
|
1111
|
-
hiddenSize: number;
|
|
1112
|
-
numHeads: number;
|
|
1113
|
-
headDim: number;
|
|
1114
|
-
depth: number;
|
|
1115
|
-
intermediateSize: number;
|
|
1116
|
-
textHidden: number;
|
|
1117
|
-
patchSize: number;
|
|
1118
|
-
patchDim: number;
|
|
1119
|
-
poolingKernelSize: number;
|
|
1120
|
-
ropeTheta: number;
|
|
1121
|
-
rmsNormEps: number;
|
|
1122
|
-
}
|
|
1123
|
-
/**
|
|
1124
|
-
* Resolve the Gemma 4 vision dims from a raw HF config. Accepts either the
|
|
1125
|
-
* top-level config (reads `.vision_config` + `.text_config.hidden_size`) or a
|
|
1126
|
-
* bare vision_config (then `textHidden` falls back to the projector row count if
|
|
1127
|
-
* present, else hidden). Family-general — no E2B constants.
|
|
1128
|
-
*/
|
|
1129
|
-
declare function resolveGemma4VisionInfo(rawConfig: Record<string, unknown>): Gemma4VisionGraphInfo;
|
|
1130
|
-
/**
|
|
1131
|
-
* Dequantize an MLX affine-int4 weight to a plain f32 [rows, cols] matrix.
|
|
1132
|
-
* MLX packs 8 int4 values per u32 (low-nibble first); each group of `groupSize`
|
|
1133
|
-
* columns shares one scale + bias: w[r,c] = scale[r, c/gs] * q + bias[r, c/gs].
|
|
1134
|
-
* Used for the Gemma 4 multimodal projector (`embed_vision.embedding_projection`)
|
|
1135
|
-
* in MLX-4bit checkpoints, where (unlike the BF16 ViT body) the projector is int4.
|
|
1136
|
-
*/
|
|
1137
|
-
declare function dequantizeMLXProjection(packed: Uint32Array, scales: Float32Array, biases: Float32Array, rows: number, cols: number, groupSize: number): Float32Array;
|
|
1138
|
-
/**
|
|
1139
|
-
* If the Gemma 4 multimodal projector arrived as an MLX affine-int4 triplet
|
|
1140
|
-
* (`embed_vision.embedding_projection.{weight(U32), scales, biases}`), dequantize
|
|
1141
|
-
* it in-place to a plain f32 `embed_vision.embedding_projection.weight` and drop
|
|
1142
|
-
* the scales/biases, so the vision graph's plain MatMul on the projector works for
|
|
1143
|
-
* MLX-4bit checkpoints too. No-op for BF16 (HF) checkpoints (weight already f32).
|
|
1144
|
-
*/
|
|
1145
|
-
declare function dequantizeGemma4VisionProjection(weights: Map<string, {
|
|
1146
|
-
data: ArrayBufferView;
|
|
1147
|
-
shape: number[];
|
|
1148
|
-
}>, groupSize: number, rows: number, cols: number): void;
|
|
1149
|
-
/**
|
|
1150
|
-
* Patch the ClippedMatMul nodes of a Gemma 4 vision graph with the calibrated clip
|
|
1151
|
-
* scalars from the checkpoint (Gemma4ClippableLinear's per-tensor input/output
|
|
1152
|
-
* min/max buffers), then drop those scalar tensors from the weights map so the
|
|
1153
|
-
* vision executor doesn't try to upload them as GPU buffers. Call BEFORE
|
|
1154
|
-
* VisionExecutor.uploadWeights(). Missing scalars default to ±inf (clip = identity),
|
|
1155
|
-
* so a checkpoint without calibration still loads.
|
|
1156
|
-
*/
|
|
1157
|
-
declare function patchGemma4VisionClips(graph: ModelGraph, weights: Map<string, {
|
|
1158
|
-
data: ArrayBufferView;
|
|
1159
|
-
shape: number[];
|
|
1160
|
-
}>): void;
|
|
1161
|
-
/**
|
|
1162
|
-
* Build the Gemma 4 ViT graph. Shaped by symbolic "N" (number of patches, runtime)
|
|
1163
|
-
* and "Np" (number of pooled tokens = ceil(grid_h/k)·ceil(grid_w/k), runtime),
|
|
1164
|
-
* resolved from input tensor dims — like the Qwen ViT's "N"/"Nm".
|
|
1165
|
-
*/
|
|
1166
|
-
declare function generateGemma4VisionGraph(rawConfig: Record<string, unknown>): ModelGraph;
|
|
1167
|
-
//#endregion
|
|
1168
|
-
//#region src/gpu/architectures/kani_tts.d.ts
|
|
1169
|
-
/** Parsed KaniTTS2 backbone config (the LFM2 dims + the TTS-specific fields). */
|
|
1170
|
-
interface KaniConfig {
|
|
1171
|
-
textVocabSize: number;
|
|
1172
|
-
vocabSize: number;
|
|
1173
|
-
tokensPerFrame: number;
|
|
1174
|
-
audioStep: number;
|
|
1175
|
-
useLearnableRope: boolean;
|
|
1176
|
-
alphaMin: number;
|
|
1177
|
-
alphaMax: number;
|
|
1178
|
-
speakerEmbDim: number;
|
|
1179
|
-
audioTokensStart: number;
|
|
1180
|
-
startOfSpeech: number;
|
|
1181
|
-
endOfSpeech: number;
|
|
1182
|
-
codebookSize: number;
|
|
1183
|
-
/**
|
|
1184
|
-
* Whether the backbone uses KaniTTS2's frame-level position IDs (the 4 audio
|
|
1185
|
-
* tokens of a frame share one logical position) + learnable per-layer RoPE.
|
|
1186
|
-
* The `KaniTTS2ForCausalLM` checkpoints (e.g. kani-tts-2-en) do; the plain
|
|
1187
|
-
* `Lfm2ForCausalLM` family (e.g. the Apache kani-tts-450m-0.2-ft / kani-tts-370m)
|
|
1188
|
-
* does NOT — those run standard sequential positions with α=1 plain RoPE.
|
|
1189
|
-
* Detected from the presence of `use_learnable_rope`/`tokens_per_frame`/the
|
|
1190
|
-
* `KaniTTS2ForCausalLM` architecture marker in config.json.
|
|
1191
|
-
*/
|
|
1192
|
-
frameLevelPositions: boolean;
|
|
1193
|
-
}
|
|
1194
|
-
declare function parseKaniConfig(rawConfig: Record<string, unknown>): KaniConfig;
|
|
1195
|
-
/**
|
|
1196
|
-
* Convert the model's flat audio-token stream (the IDs between start/end-of-speech)
|
|
1197
|
-
* into NanoCodec codes [NUM_GROUPS, T]. Mirrors NemoAudioPlayer.get_nano_codes:
|
|
1198
|
-
* reshape [-1, 4]; codes[:,c] -= codebook_size*c; codes -= audio_tokens_start;
|
|
1199
|
-
* transpose → [4, T]. Returns a Uint32Array laid out group-major ([g*T + t]).
|
|
1200
|
-
*/
|
|
1201
|
-
declare function audioTokensToCodes(audioTokenIds: number[], cfg?: KaniConfig): {
|
|
1202
|
-
codes: Uint32Array;
|
|
1203
|
-
numFrames: number;
|
|
1204
|
-
};
|
|
1205
|
-
interface NanoCodecGraphOptions {
|
|
1206
|
-
/** Number of audio frames T (the code grid width). PCM length = T * 1764. */
|
|
1207
|
-
numFrames: number;
|
|
1208
|
-
}
|
|
1209
|
-
declare function generateNanoCodecDecoderGraph(opts: NanoCodecGraphOptions): ModelGraph;
|
|
1210
|
-
/**
|
|
1211
|
-
* Generate the KaniTTS codec-LM backbone graph (LFM2-350M body, full-vocab logits,
|
|
1212
|
-
* per-layer learnable MRoPE). Mirrors generateLfm2Graph block-for-block; the only
|
|
1213
|
-
* difference is that each attention layer rotates Q/K with the MRoPE op fed a
|
|
1214
|
-
* per-layer host cos/sin table instead of the position-counter RoPE op.
|
|
1215
|
-
*
|
|
1216
|
-
* dtype:
|
|
1217
|
-
* • undefined / "f32" — full precision (the checkpoint is bf16/f32 → f32).
|
|
1218
|
-
* • "q4" — weight-only group-wise INT4 (W4A16) for the LFM2 linear projections
|
|
1219
|
-
* (q/k/v/o, gate/up/down, short-conv in/out_proj) via MatMulInt4. The tied
|
|
1220
|
-
* embedding + LM head stay at f32 on purpose: decoder-LM heads are the most
|
|
1221
|
-
* quantization-sensitive part of a codec-LM and the embed table is the single
|
|
1222
|
-
* buffer shared by both, so quantizing it would degrade audio quality for no
|
|
1223
|
-
* load-size win that the int4 projections do not already capture. Backbone
|
|
1224
|
-
* drops from ~1.5 GB (f32) to ~0.37 GB. The NanoCodec decoder is a separate
|
|
1225
|
-
* graph and always runs f32.
|
|
1226
|
-
*/
|
|
1227
|
-
declare function generateKaniTtsGraph(rawConfig: Record<string, unknown>, dtype?: GraphDType, groupSize?: number, kvDtype?: KVDType): ModelGraph;
|
|
1228
|
-
//#endregion
|
|
1229
|
-
//#region src/gpu/architectures/moonshine.d.ts
|
|
1230
|
-
/**
|
|
1231
|
-
* Historical record of the executor-side work the keystone left open; all items
|
|
1232
|
-
* are now implemented (see the STATUS note above). Retained as exported
|
|
1233
|
-
* documentation of the dependency surface.
|
|
1234
|
-
*/
|
|
1235
|
-
declare const MOONSHINE_REMAINING_WORK: readonly ["DONE: Transpose op kernel (conv output [C,L] → encoder input [L,C]).", "DONE: GroupNorm(num_groups=1) over conv channels (weight+bias).", "DONE: Tanh elementwise op (conv1 activation).", "DONE: Interleaved-RoPE variant (Moonshine rotates adjacent dim pairs 2p/2p+1).", "DONE: No-bias LayerNorm variant (Moonshine norms are weight-only).", "DONE: Per-utterance conv graph regeneration from the concrete sample count.", "DONE: Dual-graph executor — encoder once → frozen cross-attn K/V → AR decode.", "DONE: engine.transcribe(pcm) host path (MoonshineSTT): conv → encode → AR decode → detokenize."];
|
|
1236
|
-
interface MoonshineDims {
|
|
1237
|
-
hidden_size: number;
|
|
1238
|
-
enc_layers: number;
|
|
1239
|
-
dec_layers: number;
|
|
1240
|
-
num_heads: number;
|
|
1241
|
-
num_kv_heads: number;
|
|
1242
|
-
head_dim: number;
|
|
1243
|
-
rotary_dim: number;
|
|
1244
|
-
intermediate_size: number;
|
|
1245
|
-
vocab_size: number;
|
|
1246
|
-
rope_base: number;
|
|
1247
|
-
context_length: number;
|
|
1248
|
-
ln_eps: number;
|
|
1249
|
-
}
|
|
1250
|
-
/** Pull and derive Moonshine dimensions from the raw HF config. */
|
|
1251
|
-
declare function parseMoonshineConfig(raw: Record<string, unknown>): MoonshineDims;
|
|
1252
|
-
/** Number of encoder frames produced by the conv frontend for n_samples PCM. */
|
|
1253
|
-
declare function moonshineEncoderFrames(nSamples: number): number;
|
|
1254
|
-
/**
|
|
1255
|
-
* Encoder graph: raw-waveform conv frontend + bidirectional transformer.
|
|
1256
|
-
* @param nSamples concrete PCM sample count (the conv frontend is length-static).
|
|
1257
|
-
* The output tensor "encoder_out" is [T_frames, hidden] and is consumed (after
|
|
1258
|
-
* per-layer K/V projection) as the frozen K/V for the decoder's cross-attention.
|
|
1259
|
-
*/
|
|
1260
|
-
declare function generateMoonshineEncoderGraph(raw: Record<string, unknown>, nSamples: number): ModelGraph;
|
|
1261
|
-
/**
|
|
1262
|
-
* Decoder graph: AR transformer with causal self-attn (KV-cache) + cross-attn to
|
|
1263
|
-
* the frozen encoder output. Built for a single decode step (T=1). The encoder K/V
|
|
1264
|
-
* are supplied as graph inputs "enc_k_layer{i}" / "enc_v_layer{i}" — the host
|
|
1265
|
-
* pre-projects the encoder output through each layer's encoder_attn.k_proj/v_proj
|
|
1266
|
-
* ONCE and binds them frozen for the whole decode (the CrossAttention contract).
|
|
1267
|
-
*/
|
|
1268
|
-
declare function generateMoonshineDecoderGraph(raw: Record<string, unknown>, sEnc: number): ModelGraph;
|
|
1269
|
-
//#endregion
|
|
1270
|
-
//#region src/gpu/architectures/outetts.d.ts
|
|
1271
|
-
/** Parsed OuteTTS backbone config (the Qwen3 dims). Reuses the generic Qwen path. */
|
|
1272
|
-
interface OuteTtsConfig {
|
|
1273
|
-
hidden_size: number;
|
|
1274
|
-
num_layers: number;
|
|
1275
|
-
num_heads: number;
|
|
1276
|
-
num_kv_heads: number;
|
|
1277
|
-
head_dim: number;
|
|
1278
|
-
intermediate_size: number;
|
|
1279
|
-
vocab_size: number;
|
|
1280
|
-
rope_base: number;
|
|
1281
|
-
rms_norm_eps: number;
|
|
1282
|
-
max_position_embeddings: number;
|
|
1283
|
-
}
|
|
1284
|
-
declare function parseOuteTtsConfig(rawConfig: Record<string, unknown>): OuteTtsConfig;
|
|
1285
|
-
/**
|
|
1286
|
-
* Convert the flat generated token stream (the run BETWEEN <|audio_start|> and
|
|
1287
|
-
* <|audio_end|>) into a DAC code grid [NUM_CODEBOOKS, T] (codebook-major [c*T + t]).
|
|
1288
|
-
*
|
|
1289
|
-
* Mirrors outetts v3 `extract_audio_from_tokens`: filter the stream into two ordered
|
|
1290
|
-
* lists by membership of the c1 / c2 ranges, recover the raw code
|
|
1291
|
-
* (id − base), drop the reserved index 1024, and truncate to the shorter list. The
|
|
1292
|
-
* model emits c1,c2 interleaved per frame, but the parse is membership-based (it does
|
|
1293
|
-
* not assume strict pairing) so a stray token cannot desync the grid.
|
|
1294
|
-
*/
|
|
1295
|
-
declare function audioTokensToDacCodes(tokenIds: number[]): {
|
|
1296
|
-
codes: Uint32Array;
|
|
1297
|
-
numFrames: number;
|
|
1298
|
-
};
|
|
1299
|
-
interface DacSpeechGraphOptions {
|
|
1300
|
-
/** Number of audio frames T (the code-grid width). */
|
|
1301
|
-
numFrames: number;
|
|
1302
|
-
/**
|
|
1303
|
-
* Number of RVQ codebooks. Default 2 (DAC.speech.v1.0). Parler-TTS's
|
|
1304
|
-
* dac_44khz_8kbps uses 9 — the RVQDequant kernel loops over NCB so the only
|
|
1305
|
-
* change is the codebook/out_proj tensor first-dim and the codes grid height.
|
|
1306
|
-
*/
|
|
1307
|
-
numCodebooks?: number;
|
|
1308
|
-
/**
|
|
1309
|
-
* Decoder upsampling rates (the DecoderBlock strides). Default [8,5,4,2]
|
|
1310
|
-
* (DAC.speech.v1.0). Parler's dac_44khz uses [8,8,4,2] (hop 512). Channels
|
|
1311
|
-
* still halve per stage from `decoderDim`.
|
|
1312
|
-
*/
|
|
1313
|
-
decoderRates?: readonly number[];
|
|
1314
|
-
/** First-conv output channels (decoder_dim). Default 1536. */
|
|
1315
|
-
decoderDim?: number;
|
|
1316
|
-
/** Quantizer out_proj output / first-conv input dim (latent_dim). Default 1024. */
|
|
1317
|
-
latentDim?: number;
|
|
1318
|
-
}
|
|
1319
|
-
declare function generateDacSpeechDecoderGraph(opts: DacSpeechGraphOptions): ModelGraph;
|
|
1320
|
-
/**
|
|
1321
|
-
* Final decoded PCM length (samples) for a T-frame code grid, via the DAC stage chain.
|
|
1322
|
-
* `rates` defaults to DAC.speech.v1.0 ([8,5,4,2]); pass [8,8,4,2] for Parler's dac_44khz.
|
|
1323
|
-
*/
|
|
1324
|
-
declare function dacOutputLength(numFrames: number, rates?: readonly number[]): number;
|
|
1325
|
-
declare function generateOuteTtsBackboneGraph(rawConfig: Record<string, unknown>, dtype?: GraphDType, groupSize?: number, kvDtype?: KVDType): ModelGraph;
|
|
1326
|
-
//#endregion
|
|
1327
|
-
//#region src/gpu/architectures/qwen3_5_vision.d.ts
|
|
1328
|
-
/**
|
|
1329
|
-
* Build the ViT graph. The graph is shaped by symbolic "N" (number of patches),
|
|
1330
|
-
* resolved at run time from the input tensor's first dim — exactly like the LM's
|
|
1331
|
-
* symbolic "T".
|
|
1332
|
-
*/
|
|
1333
|
-
declare function generateQwen3_5VisionGraph(rawConfig: Record<string, unknown>): ModelGraph;
|
|
1334
|
-
//#endregion
|
|
1335
817
|
//#region src/gpu/executor.d.ts
|
|
1336
818
|
/**
|
|
1337
819
|
* Safari/Metal workaround: shader variant alternation.
|
|
@@ -1399,6 +881,17 @@ declare class Executor {
|
|
|
1399
881
|
private argmaxReadback;
|
|
1400
882
|
/** Readback ring for pipelined greedy decode (created lazily). */
|
|
1401
883
|
private decodeReadbacks;
|
|
884
|
+
/** Window size K (GERBIL_DECODE_WINDOW). 0 = off. */
|
|
885
|
+
readonly decodeWindowK: number;
|
|
886
|
+
/** Readback ring for decode windows (4·K bytes per slot; created lazily). */
|
|
887
|
+
private windowReadbacks;
|
|
888
|
+
/**
|
|
889
|
+
* Per-step uniform/bind-group variants for the few seq-pos-dependent decode
|
|
890
|
+
* entries (RoPE position, KV-append offset, attention S). Keyed by nodeId;
|
|
891
|
+
* index k holds the variant used at window step k (1..K-1 — step 0 uses the
|
|
892
|
+
* entry's base uniform). Created lazily, rewritten per window pre-submit.
|
|
893
|
+
*/
|
|
894
|
+
private windowVariants;
|
|
1402
895
|
/**
|
|
1403
896
|
* Staging buffer for uniform param updates.
|
|
1404
897
|
* Safari/Metal has weaker visibility guarantees for queue.writeBuffer() to
|
|
@@ -1431,6 +924,48 @@ declare class Executor {
|
|
|
1431
924
|
private kvMode;
|
|
1432
925
|
private seqPos;
|
|
1433
926
|
private webkitGroupSize;
|
|
927
|
+
private readonly pagedKV;
|
|
928
|
+
private kvBlockManager;
|
|
929
|
+
/** Lockstep batch width (GERBIL_BATCH). 0 = batching off. */
|
|
930
|
+
readonly batchSize: number;
|
|
931
|
+
/** Block-table row stride — set for any paged mode. */
|
|
932
|
+
private maxBlocksPerSeq;
|
|
933
|
+
/** Block-table row the single-sequence paths target (batch prefill slot). */
|
|
934
|
+
private activeSeqSlot;
|
|
935
|
+
/** Batched dispatch entries (built once in initBindGroups when batch on). */
|
|
936
|
+
private batchEntries;
|
|
937
|
+
/** Per-slot SSM/conv state pools: slot b's state at offset b*slotBytes. */
|
|
938
|
+
private batchSsmPools;
|
|
939
|
+
/** Batch-widened tail buffers (SliceLastRow output + logits at [N, ...]). */
|
|
940
|
+
private batchTailBuffers;
|
|
941
|
+
private batchPositionsBuffer;
|
|
942
|
+
private batchSeqLensBuffer;
|
|
943
|
+
private batchSampledBuffer;
|
|
944
|
+
private batchReadbacks;
|
|
945
|
+
private batchArgmaxEntry;
|
|
946
|
+
/** JS-side authoritative per-slot token counts (positions source). */
|
|
947
|
+
private batchSeqLens;
|
|
948
|
+
/** Scratch for per-step positions/seq_lens uploads. */
|
|
949
|
+
private batchPositionsScratch;
|
|
950
|
+
private batchSeqLensScratch;
|
|
951
|
+
/** Cached batch-resolved shapes (static: T == batchSize every step). */
|
|
952
|
+
private batchShapes;
|
|
953
|
+
/** True once the batched-argmax uniform has been written this session. */
|
|
954
|
+
private batchArgmaxParamsWritten;
|
|
955
|
+
/** True when GERBIL_BATCH_LORA=1 rode in on an active batch path. */
|
|
956
|
+
readonly batchLoraEnabled: boolean;
|
|
957
|
+
/** Registered adapters (expanded deltas), indexed by adapter id. */
|
|
958
|
+
private batchAdapterDeltas;
|
|
959
|
+
/** Per-lane adapter ids (u32; 0xFFFFFFFF = bare base). */
|
|
960
|
+
private batchAdapterIdsBuffer;
|
|
961
|
+
/** GPU buffers owned by the batch-LoRA overlay (factors, maps, tmp, uniforms). */
|
|
962
|
+
private batchLoraBuffers;
|
|
963
|
+
/** Batch entries snapshotted BEFORE the batch fusion passes (overlay base). */
|
|
964
|
+
private batchEntriesUnfused;
|
|
965
|
+
/** Batch entries with the per-lane LoRA dispatches spliced in. */
|
|
966
|
+
private batchEntriesLoRA;
|
|
967
|
+
/** True while at least one lane carries an adapter (set by setBatchLaneAdapters). */
|
|
968
|
+
private batchLoraActive;
|
|
1434
969
|
private profileEnabled;
|
|
1435
970
|
private readonly profileData;
|
|
1436
971
|
private querySet;
|
|
@@ -1573,8 +1108,15 @@ declare class Executor {
|
|
|
1573
1108
|
private buildSelectiveLoraDecode;
|
|
1574
1109
|
/**
|
|
1575
1110
|
* Run a forward pass. Uses matvec kernels for M=1 (decode), tiled for M>1 (prefill).
|
|
1111
|
+
*
|
|
1112
|
+
* `opts.readLogits: false` skips the logits copy + readback entirely (the
|
|
1113
|
+
* returned logits array is empty) and returns right after submission — used
|
|
1114
|
+
* by the scheduler's chunked prefill, where only the FINAL chunk's logits
|
|
1115
|
+
* are ever sampled. seqPos still advances by T.
|
|
1576
1116
|
*/
|
|
1577
|
-
forward(inputIds: Uint32Array
|
|
1117
|
+
forward(inputIds: Uint32Array, opts?: {
|
|
1118
|
+
readLogits?: boolean;
|
|
1119
|
+
}): Promise<ForwardResult>;
|
|
1578
1120
|
/**
|
|
1579
1121
|
* Profiling variant of the desktop dispatch path: one compute pass per dispatch,
|
|
1580
1122
|
* each bracketed by timestamp queries, so we get per-op GPU time. Accumulates
|
|
@@ -1605,6 +1147,14 @@ declare class Executor {
|
|
|
1605
1147
|
* tiny dispatch) is intentionally excluded — it is not a hotspot target.
|
|
1606
1148
|
*/
|
|
1607
1149
|
profileDecodeStep(tokenId: number): Promise<void>;
|
|
1150
|
+
/**
|
|
1151
|
+
* Profile ONE batched decode step (GERBIL_BATCH): times each batchEntries
|
|
1152
|
+
* dispatch with timestamp queries, accumulating GPU ns per opType into the
|
|
1153
|
+
* same profile as profileDecodeStep. Measurement-only (un-pipelined, CPU
|
|
1154
|
+
* tokens). Call after prefilling the batch slots; advances every slot by
|
|
1155
|
+
* one position like a real step.
|
|
1156
|
+
*/
|
|
1157
|
+
profileBatchDecodeStep(tokens: Uint32Array<ArrayBuffer>): Promise<void>;
|
|
1608
1158
|
/**
|
|
1609
1159
|
* Run a single forward pass over `inputIds` and read back the L2-normalized
|
|
1610
1160
|
* embedding vector. Requires an embedding graph (one whose output tensor is
|
|
@@ -1663,6 +1213,44 @@ declare class Executor {
|
|
|
1663
1213
|
submitGreedyDecodeStep(tokenId: number | null, slot: number): void;
|
|
1664
1214
|
/** Read back the token produced by the pipelined step that used `slot`. */
|
|
1665
1215
|
readDecodeToken(slot: number): Promise<number>;
|
|
1216
|
+
/**
|
|
1217
|
+
* K-step greedy decode window (GERBIL_DECODE_WINDOW=K — dispatch-collapse
|
|
1218
|
+
* campaign Lever 1; Dawn only).
|
|
1219
|
+
*
|
|
1220
|
+
* Records `steps` full decode forwards + argmax into ONE command buffer and
|
|
1221
|
+
* submits once. Between recorded steps the previous step's argmax result is
|
|
1222
|
+
* copied into input_ids ON THE GPU (the same 4-byte chain the pipelined path
|
|
1223
|
+
* uses across submits), so no CPU sync exists inside the window. Each step's
|
|
1224
|
+
* winning token id is copied into the slot's window readback at offset 4·k;
|
|
1225
|
+
* {@link readDecodeWindow} maps it once per window.
|
|
1226
|
+
*
|
|
1227
|
+
* Seq-pos-dependent entries (RoPE position, KV-append offset, attention S —
|
|
1228
|
+
* ~3 per full-attention layer; Mamba layers carry state and have none) get
|
|
1229
|
+
* lazily-built per-step uniform/bind-group variants, ALL written with
|
|
1230
|
+
* queue.writeBuffer BEFORE the submit (queue-ordered, so safe with a prior
|
|
1231
|
+
* window still in flight). Params are built against per-step resolved shapes
|
|
1232
|
+
* (L_max tensors resolve to seqPos+T — attention S comes from the KV shape).
|
|
1233
|
+
* Dispatch grids are encoded ONCE per window at the worst case (last step's
|
|
1234
|
+
* shapes, pos+K-1): decode grids are step-invariant in practice, and any
|
|
1235
|
+
* L_max-derived grid over-dispatches for early steps, which is safe because
|
|
1236
|
+
* kernels bound their work from uniforms.
|
|
1237
|
+
*
|
|
1238
|
+
* EOS is only observable at window granularity: up to K-1 tokens beyond EOS
|
|
1239
|
+
* are computed and discarded by the caller. KV overshoot is harmless (seqPos
|
|
1240
|
+
* governs reads) but SSM/conv state is NOT rewindable — continuing the same
|
|
1241
|
+
* context past an intra-window EOS requires a re-prefill from that point.
|
|
1242
|
+
*
|
|
1243
|
+
* Evidence base: measured FLAT on M4 Max/Metal (results.jsonl r13-E2 /
|
|
1244
|
+
* r14-E2v2 — depth-2 pipelining already hides submit+readback there);
|
|
1245
|
+
* targeted at Dawn/Vulkan where the submit roundtrip is 52µs and decode is
|
|
1246
|
+
* dispatch/driver-bound (A5000 gate, arXiv:2604.02344). See
|
|
1247
|
+
* docs/design/dispatch-collapse-campaign.md §3.
|
|
1248
|
+
*/
|
|
1249
|
+
submitGreedyDecodeWindow(firstToken: number | null, steps: number, slot: number): void;
|
|
1250
|
+
/** Get (or lazily create) entry's per-step uniform/bind-group variant. */
|
|
1251
|
+
private getWindowVariant;
|
|
1252
|
+
/** Read back the `steps` token ids produced by the window that used `slot`. */
|
|
1253
|
+
readDecodeWindow(slot: number, steps: number): Promise<Uint32Array>;
|
|
1666
1254
|
reset(): void;
|
|
1667
1255
|
/**
|
|
1668
1256
|
* Diagnostic: dispatch ONLY the first kernel (EmbeddingInt4) using the
|
|
@@ -1744,12 +1332,235 @@ declare class Executor {
|
|
|
1744
1332
|
private allocateActivationBuffers;
|
|
1745
1333
|
private allocateSSMStateBuffers;
|
|
1746
1334
|
private allocateKVCacheBuffers;
|
|
1335
|
+
/** f16 SSM-state kernels are selected on Dawn with f16 support (not WebKit). */
|
|
1336
|
+
private get usesF16Ssm();
|
|
1337
|
+
/**
|
|
1338
|
+
* Batched-decode state (GERBIL_BATCH=N): per-slot SSM/conv state pools, the
|
|
1339
|
+
* per-row positions/seq_lens buffers, the sampled-token buffer + readback
|
|
1340
|
+
* ring, and batch-widened buffers for the lm-head tail (SliceLastRow output
|
|
1341
|
+
* and logits become [N, ...]).
|
|
1342
|
+
*
|
|
1343
|
+
* SSM pool slot stride: the f16 MambaSSM kernel indexes its state buffer in
|
|
1344
|
+
* f16 ELEMENTS (the buffer is f32-sized with the upper half unused), so its
|
|
1345
|
+
* pool slot stride is elems*2 bytes; conv state (f32 kernels) uses the full
|
|
1346
|
+
* f32 byte size. The batched kernels compute `slot * elems` in their own
|
|
1347
|
+
* element space, so the pool stride must match each kernel's element width.
|
|
1348
|
+
*/
|
|
1349
|
+
private allocateBatchState;
|
|
1747
1350
|
private resolveShapes;
|
|
1748
1351
|
private getBuffer;
|
|
1352
|
+
/**
|
|
1353
|
+
* Batch-path buffer resolution: lm-head tail tensors use their [B, ...]
|
|
1354
|
+
* widened buffers, ssm_state tensors bind the per-slot pool, everything
|
|
1355
|
+
* else falls through to the regular buffer maps.
|
|
1356
|
+
*/
|
|
1357
|
+
private getBatchBuffer;
|
|
1358
|
+
/**
|
|
1359
|
+
* True when this node's KV traffic goes through the paged path: the flag is
|
|
1360
|
+
* on AND the node actually touches a kv_cache tensor (KVCacheAppend writes
|
|
1361
|
+
* one; self-attention reads one as its K input). SSM/conv state and
|
|
1362
|
+
* ViT/cross-attention over plain activations never match — on hybrid models
|
|
1363
|
+
* only the attention layers page, exactly as the design requires.
|
|
1364
|
+
*/
|
|
1365
|
+
private usesPagedKV;
|
|
1366
|
+
/**
|
|
1367
|
+
* Paged KV only: make sure blocks are mapped for logical positions
|
|
1368
|
+
* [0, seqPos + T) BEFORE this forward's dispatches are submitted.
|
|
1369
|
+
* queue.writeBuffer is queue-ordered, so the (rare — once per 16 tokens)
|
|
1370
|
+
* table upload always lands ahead of the submits that read it, including in
|
|
1371
|
+
* the pipelined decode path with PIPELINE_DEPTH steps in flight.
|
|
1372
|
+
*/
|
|
1373
|
+
private preparePagedKV;
|
|
1374
|
+
/** Runtime context for batched dispatches (positions live in buffers). */
|
|
1375
|
+
private batchRuntimeCtx;
|
|
1376
|
+
/**
|
|
1377
|
+
* Shapes for the batched step: every "T" dimension is the batch width, and
|
|
1378
|
+
* the lm-head tail tensors (statically [1, ...] in the graph) widen to
|
|
1379
|
+
* [B, ...] — every row IS a last row on the batch path.
|
|
1380
|
+
*/
|
|
1381
|
+
private resolveBatchShapes;
|
|
1382
|
+
/**
|
|
1383
|
+
* Build the batched dispatch-entry list: one lockstep decode step for all
|
|
1384
|
+
* batchSize rows. Row-independent ops (embedding, norms, elementwise) keep
|
|
1385
|
+
* their prefill kernels at T=B; linear layers run the batched INT4 GEMM;
|
|
1386
|
+
* every position/state-dependent op swaps to its batch-indexed variant
|
|
1387
|
+
* (paged attention + append via per-row block-table rows, RoPE via the
|
|
1388
|
+
* positions buffer, conv/SSM via the per-slot state pool).
|
|
1389
|
+
*/
|
|
1390
|
+
private buildBatchEntries;
|
|
1391
|
+
/**
|
|
1392
|
+
* Batched-path fusion driver (Phase 4). Ports the single-sequence decode
|
|
1393
|
+
* fusions onto `batchEntries`: identical math (the fused kernels reproduce
|
|
1394
|
+
* the unfused chain operation-for-operation — exact-tier GEMM reductions
|
|
1395
|
+
* keep the V2 order, elementwise prologues/epilogues keep the registry
|
|
1396
|
+
* kernels' formulas), strictly fewer dispatches. `GERBIL_BATCH_FUSE`
|
|
1397
|
+
* selects a subset for A/B measurement: unset/"all" = every pass,
|
|
1398
|
+
* "none" = none, or a comma list of pass names.
|
|
1399
|
+
*/
|
|
1400
|
+
private fuseBatchDecodeEntries;
|
|
1401
|
+
/**
|
|
1402
|
+
* Batch port of fuseSwiGLUDecodeEntries: gate MatMulInt4 + up MatMulInt4 +
|
|
1403
|
+
* elementwise SwiGLU → ONE fused batched GEMM (saves 2 dispatches per MLP
|
|
1404
|
+
* block and the gate/up activation round-trips).
|
|
1405
|
+
*/
|
|
1406
|
+
private fuseBatchSwiGLUMLPEntries;
|
|
1407
|
+
/**
|
|
1408
|
+
* Batch port of fuseDualMatVecDecodeEntries: two adjacent MatMulInt4s
|
|
1409
|
+
* sharing one input and identical K/N/group_size (q+gate, k+v) → ONE dual
|
|
1410
|
+
* batched GEMM dispatch (wid.y selects the projection — per-projection
|
|
1411
|
+
* math unchanged).
|
|
1412
|
+
*/
|
|
1413
|
+
private fuseBatchDualMatVecEntries;
|
|
1414
|
+
/**
|
|
1415
|
+
* Batch port of fuseDualKVCacheAppendEntries: adjacent K+V paged appends →
|
|
1416
|
+
* ONE BATCHED_PAGED_DUAL_KV_APPEND dispatch (pure memcpys — identical
|
|
1417
|
+
* writes, one dispatch per attention layer saved).
|
|
1418
|
+
*/
|
|
1419
|
+
private fuseBatchDualKVAppendEntries;
|
|
1420
|
+
/**
|
|
1421
|
+
* Batch port of fuseGatedOProjDecodeEntries ("sigmoid": SigmoidGate →
|
|
1422
|
+
* o_proj) and fuseSwiGLUGatedProjDecodeEntries ("swiglu": Mamba SwiGLU →
|
|
1423
|
+
* out_proj): the elementwise gate is recomputed inside the GEMM's input
|
|
1424
|
+
* load with the registry kernel's exact formula — one dispatch per layer
|
|
1425
|
+
* saved plus the gated-activation round-trip.
|
|
1426
|
+
*/
|
|
1427
|
+
private fuseBatchGatedProjEntries;
|
|
1428
|
+
/**
|
|
1429
|
+
* Batch port of fuseDualRMSNormRoPEDecodeEntries: [q_norm, k_norm, RoPE]
|
|
1430
|
+
* → ONE BATCHED_DUAL_RMSNORM_ROPE dispatch (per-row rotation position from
|
|
1431
|
+
* the positions buffer). Saves 2 dispatches per full-attention layer.
|
|
1432
|
+
*/
|
|
1433
|
+
private fuseBatchDualRMSNormRoPEEntries;
|
|
1434
|
+
/**
|
|
1435
|
+
* Reset ALL batch slots for a new lockstep generation: frees every
|
|
1436
|
+
* sequence's KV blocks, zeroes the single-sequence SSM state (the prefill
|
|
1437
|
+
* staging area), and clears per-slot lengths. State-pool slots need no
|
|
1438
|
+
* zeroing — finishBatchSlot overwrites a slot in full before it is read.
|
|
1439
|
+
*/
|
|
1440
|
+
resetBatch(): void;
|
|
1441
|
+
/**
|
|
1442
|
+
* Route the single-sequence forward() path at batch slot `slot`: paged KV
|
|
1443
|
+
* writes/reads use the slot's block-table row (RuntimeContext.tableBase) and
|
|
1444
|
+
* SSM state accumulates in the singleton staging buffers (zeroed).
|
|
1445
|
+
* Call once per prompt, then forward(promptIds), then finishBatchSlot().
|
|
1446
|
+
*/
|
|
1447
|
+
beginBatchSlot(slot: number): void;
|
|
1448
|
+
/**
|
|
1449
|
+
* Adopt the just-prefilled sequence into `slot`: records its length, copies
|
|
1450
|
+
* the singleton SSM/conv state into the slot's pool rows (KV already lives
|
|
1451
|
+
* in the slot's block-table row via tableBase), and zeroes the singleton
|
|
1452
|
+
* staging state for the next prefill.
|
|
1453
|
+
*/
|
|
1454
|
+
finishBatchSlot(slot: number): void;
|
|
1455
|
+
/** Per-slot token counts (prompt + generated) on the batch path. */
|
|
1456
|
+
get batchSlotLengths(): Uint32Array;
|
|
1457
|
+
/**
|
|
1458
|
+
* Submit one lockstep batched decode step (pipelined — does not await).
|
|
1459
|
+
* `tokens` feeds input_ids on the first step; pass null afterwards to chain
|
|
1460
|
+
* the previous step's sampled ids into input_ids ON the GPU. Sampled ids
|
|
1461
|
+
* land in the `slot` readback (read one step behind via readBatchTokens).
|
|
1462
|
+
*/
|
|
1463
|
+
submitBatchDecodeStep(tokens: Uint32Array<ArrayBuffer> | null, slot: number): void;
|
|
1464
|
+
/** Read the sampled ids produced by the batched step that used `slot`. */
|
|
1465
|
+
readBatchTokens(slot: number): Promise<Uint32Array>;
|
|
1466
|
+
/** Batch entries in effect for the current lane-adapter state. */
|
|
1467
|
+
private get activeBatchEntries();
|
|
1468
|
+
/** Number of adapters registered for per-lane batched decode. */
|
|
1469
|
+
get batchAdapterCount(): number;
|
|
1470
|
+
/**
|
|
1471
|
+
* Register a LoRA adapter for per-lane batched decode and return its integer
|
|
1472
|
+
* adapter id (assign lanes via {@link setBatchLaneAdapters}). Factors are
|
|
1473
|
+
* packed into per-adapted-node GPU buffers ONCE — every lane running this
|
|
1474
|
+
* adapter reads the same factors. Fused query+gate deltas are expanded with
|
|
1475
|
+
* the same deinterleave the single-sequence overlay uses.
|
|
1476
|
+
*
|
|
1477
|
+
* Phase-1 constraint: every adapter adapting a given linear must share that
|
|
1478
|
+
* node's LoRA rank (the packed factor arena is fixed-stride; mixed ranks are
|
|
1479
|
+
* deferred).
|
|
1480
|
+
*/
|
|
1481
|
+
registerBatchAdapter(deltas: LoRADelta[]): number;
|
|
1482
|
+
/** Drop every registered batch adapter and free the overlay's GPU buffers. */
|
|
1483
|
+
clearBatchAdapters(): void;
|
|
1484
|
+
/**
|
|
1485
|
+
* Assign each batch lane its adapter: `ids[b]` is a registered adapter id or
|
|
1486
|
+
* -1 for the bare base. Uploads the per-lane adapter_ids buffer (the
|
|
1487
|
+
* positions/seq_lens per-row pattern) and switches the step's entry list to
|
|
1488
|
+
* the LoRA overlay iff any lane carries an adapter — an all-base assignment
|
|
1489
|
+
* runs the untouched plain batch entries at zero extra cost.
|
|
1490
|
+
*/
|
|
1491
|
+
setBatchLaneAdapters(ids: ArrayLike<number>): void;
|
|
1492
|
+
/** Free the batch-LoRA overlay's GPU buffers (factors, maps, tmp, uniforms). */
|
|
1493
|
+
private destroyBatchLoraBuffers;
|
|
1494
|
+
/**
|
|
1495
|
+
* (Re)build the per-lane LoRA overlay: pack every registered adapter's
|
|
1496
|
+
* factors per adapted node, and derive `batchEntriesLoRA` from the fused
|
|
1497
|
+
* batch entry list with SELECTIVE unfusing — every fused group that covers
|
|
1498
|
+
* no adapted linear stays exactly as the plain batch path runs it; groups
|
|
1499
|
+
* covering an adapted node fall back to their pre-fusion per-node entries
|
|
1500
|
+
* (token-exact by the Phase 2/4 derivation discipline) with the per-lane
|
|
1501
|
+
* `[down, up]` gather pair spliced right after each adapted node, so the
|
|
1502
|
+
* corrected output is live before any consumer reads it.
|
|
1503
|
+
*/
|
|
1504
|
+
private rebuildBatchLoraOverlay;
|
|
1505
|
+
/** Total physical KV blocks in the paged pool (admission-control bound). */
|
|
1506
|
+
get kvTotalBlocks(): number;
|
|
1507
|
+
/** KV blocks currently unallocated. */
|
|
1508
|
+
get kvFreeBlocks(): number;
|
|
1509
|
+
/** Tokens per KV block (admission-control block-count math). */
|
|
1510
|
+
get kvBlockSizeTokens(): number;
|
|
1511
|
+
/** Blocks a single sequence may map (== ceil(maxSeqLen / blockSize)). */
|
|
1512
|
+
get kvMaxBlocksPerSeq(): number;
|
|
1513
|
+
/** Effective context length (scheduler admission math). */
|
|
1514
|
+
get maxSequenceLength(): number;
|
|
1515
|
+
/**
|
|
1516
|
+
* Retire a batch lane: free its KV blocks and zero its length so masked
|
|
1517
|
+
* steps treat the slot as empty (seq_lens[slot] = 0). The SSM pool row
|
|
1518
|
+
* needs no cleanup — re-admission overwrites it in full (finishBatchSlot).
|
|
1519
|
+
*/
|
|
1520
|
+
retireBatchSlot(slot: number): void;
|
|
1521
|
+
/**
|
|
1522
|
+
* Seed `slot`'s next input token (the first token sampled from its prefill
|
|
1523
|
+
* logits). Written into the sampled-ids buffer so the next masked step's
|
|
1524
|
+
* on-GPU chain copy propagates it into input_ids. queue.writeBuffer is
|
|
1525
|
+
* queue-ordered: it lands AFTER any in-flight step's argmax write to the
|
|
1526
|
+
* same word and BEFORE the next submitted step's chain copy — correct even
|
|
1527
|
+
* with PIPELINE_DEPTH steps in flight.
|
|
1528
|
+
*/
|
|
1529
|
+
injectBatchToken(slot: number, token: number): void;
|
|
1530
|
+
/**
|
|
1531
|
+
* Submit one MASKED batched decode step (the scheduler's step; pipelined —
|
|
1532
|
+
* does not await). Only lanes with `active[b] != 0` advance: inactive lanes
|
|
1533
|
+
* get positions=0 / seq_lens=0, which masks their KV appends (kernel guard)
|
|
1534
|
+
* and clamps their attention window to nothing; their other row-local ops
|
|
1535
|
+
* compute garbage that stays confined to their rows/state slots and is
|
|
1536
|
+
* fully overwritten on re-admission.
|
|
1537
|
+
*
|
|
1538
|
+
* The step's dispatch width is dynamic: kernels take their batch dimension
|
|
1539
|
+
* from the runtime context, so only rows [0, maxActiveSlot] are dispatched
|
|
1540
|
+
* (GEMM b-tiles may cover a few dead rows past the last active one —
|
|
1541
|
+
* in-bounds stale writes that are never read). Input tokens always chain
|
|
1542
|
+
* from the sampled-ids buffer on the GPU; newly admitted lanes are seeded
|
|
1543
|
+
* via injectBatchToken. Sampled ids land in the `slot` readback (read one
|
|
1544
|
+
* step behind via readBatchTokens).
|
|
1545
|
+
*/
|
|
1546
|
+
submitBatchDecodeStepMasked(active: Uint8Array, slot: number): void;
|
|
1547
|
+
/** Write the batched-argmax uniform once per batch session (invariant). */
|
|
1548
|
+
private writeBatchArgmaxParams;
|
|
1549
|
+
private requireBatch;
|
|
1749
1550
|
/**
|
|
1750
1551
|
* Detect gate_proj + up_proj + SwiGLU patterns in decode entries and replace
|
|
1751
1552
|
* with a single fused SwiGLUMatVec dispatch. Saves 2 dispatches per MLP block.
|
|
1752
1553
|
*/
|
|
1554
|
+
/**
|
|
1555
|
+
* Fuse [ResidualRMSNorm, gate MatMulInt4, up MatMulInt4, SwiGLU] → ONE
|
|
1556
|
+
* NormSwiGLUMatVec dispatch. The kernel computes the residual sum + RMS in a
|
|
1557
|
+
* prologue (workgroup 0 persists the residual sum for the downstream residual
|
|
1558
|
+
* chain) and normalizes the activation on the fly, so the norm dispatch AND
|
|
1559
|
+
* the normed tensor's write+read roundtrip disappear. INT4-only; falls back
|
|
1560
|
+
* when the device can't bind 11 storage buffers. Windows it skips fall
|
|
1561
|
+
* through to fuseSwiGLUDecodeEntries.
|
|
1562
|
+
*/
|
|
1563
|
+
private fuseNormSwiGLUDecodeEntries;
|
|
1753
1564
|
private fuseSwiGLUDecodeEntries;
|
|
1754
1565
|
/**
|
|
1755
1566
|
* Fuse two adjacent INT4 projections that share the same input activation and
|
|
@@ -1808,45 +1619,797 @@ declare class Executor {
|
|
|
1808
1619
|
*/
|
|
1809
1620
|
private fuseSwiGLUGatedProjDecodeEntries;
|
|
1810
1621
|
/**
|
|
1811
|
-
* Fuse two adjacent per-row RMSNorms sharing hidden_size + eps into a single
|
|
1812
|
-
* DualRMSNorm dispatch (e.g. the per-head q_norm and k_norm in full-attention
|
|
1813
|
-
* decode). One workgroup still handles one row; the fused grid just spans both
|
|
1814
|
-
* inputs' rows, so each row's reduction is unchanged — numerically identical to
|
|
1815
|
-
* two separate RMSNorm dispatches. One round-trip saved per fused pair.
|
|
1816
|
-
*
|
|
1817
|
-
* WebKit risk: low — same single-workgroup reduction as the proven RMSNorm
|
|
1818
|
-
* kernel, only the row→input routing is added.
|
|
1622
|
+
* Fuse two adjacent per-row RMSNorms sharing hidden_size + eps into a single
|
|
1623
|
+
* DualRMSNorm dispatch (e.g. the per-head q_norm and k_norm in full-attention
|
|
1624
|
+
* decode). One workgroup still handles one row; the fused grid just spans both
|
|
1625
|
+
* inputs' rows, so each row's reduction is unchanged — numerically identical to
|
|
1626
|
+
* two separate RMSNorm dispatches. One round-trip saved per fused pair.
|
|
1627
|
+
*
|
|
1628
|
+
* WebKit risk: low — same single-workgroup reduction as the proven RMSNorm
|
|
1629
|
+
* kernel, only the row→input routing is added.
|
|
1630
|
+
*/
|
|
1631
|
+
/**
|
|
1632
|
+
* Fuse [CausalConv1dSiLU|CausalConv1dGated → ConvStateUpdate] into one
|
|
1633
|
+
* dispatch (decode entries only). At T=1 the conv read and the rolling-state
|
|
1634
|
+
* write for a channel live in the same thread, so the fused kernel is
|
|
1635
|
+
* race-free; prefill (T>1) keeps the two-dispatch path.
|
|
1636
|
+
*/
|
|
1637
|
+
private fuseConvStateUpdateDecodeEntries;
|
|
1638
|
+
/**
|
|
1639
|
+
* Fuse [per-head q_norm RMSNorm, k_norm RMSNorm, in-place RoPE] into ONE
|
|
1640
|
+
* DualRMSNormRoPE dispatch. Numerically identical (norm then rotate, same
|
|
1641
|
+
* math); requires head_dim <= 256 and RoPE operating in-place on exactly the
|
|
1642
|
+
* two norm outputs. Saves 2 dispatches per full-attention layer — a mobile
|
|
1643
|
+
* (WebKit submit-per-dispatch) win, desktop-neutral.
|
|
1644
|
+
*/
|
|
1645
|
+
/**
|
|
1646
|
+
* Fuse [Add → RMSNorm] into a single ResidualRMSNorm dispatch: the add's sum
|
|
1647
|
+
* is still materialized (sum_output — downstream residual chain reads it) and
|
|
1648
|
+
* the norm result is written in the same dispatch. Numerically identical.
|
|
1649
|
+
*/
|
|
1650
|
+
private fuseAddRMSNormDecodeEntries;
|
|
1651
|
+
private fuseDualRMSNormRoPEDecodeEntries;
|
|
1652
|
+
private fuseDualRMSNormDecodeEntries;
|
|
1653
|
+
/**
|
|
1654
|
+
* Gather buffer entries for a bind group, matching the kernel spec's binding layout.
|
|
1655
|
+
* Uses the pre-allocated inputIdsBuffer for the "input_ids" tensor.
|
|
1656
|
+
*/
|
|
1657
|
+
private gatherBuffers;
|
|
1658
|
+
/** Lazily allocate a scratch storage buffer at least `minBytes` large. */
|
|
1659
|
+
private getBindingScratchBuffer;
|
|
1660
|
+
}
|
|
1661
|
+
//#endregion
|
|
1662
|
+
//#region src/gpu/batch/scheduler.d.ts
|
|
1663
|
+
/** Options for a single scheduled request. */
|
|
1664
|
+
interface ServeOptions {
|
|
1665
|
+
/** Max new tokens (clamped to context; default 128). */
|
|
1666
|
+
maxTokens?: number;
|
|
1667
|
+
stopSequences?: string[];
|
|
1668
|
+
systemPrompt?: string;
|
|
1669
|
+
/** Greedy-only in Phase 5: temperature must be 0 (or omitted). */
|
|
1670
|
+
sampling?: {
|
|
1671
|
+
temperature?: number;
|
|
1672
|
+
};
|
|
1673
|
+
/**
|
|
1674
|
+
* Registered adapter name this request decodes with (SGMV — requires
|
|
1675
|
+
* `GERBIL_BATCH_LORA=1` and a prior `registerAdapter` call on the engine);
|
|
1676
|
+
* undefined = bare base. Requests with different adapters (or none) decode
|
|
1677
|
+
* concurrently over the shared base, each token-exact vs a solo run with
|
|
1678
|
+
* that adapter loaded.
|
|
1679
|
+
*/
|
|
1680
|
+
adapter?: string;
|
|
1681
|
+
/** Streamed per-token callback (decoded text piece). */
|
|
1682
|
+
onToken?: (token: string, meta: {
|
|
1683
|
+
requestId: number;
|
|
1684
|
+
tokenId: number;
|
|
1685
|
+
generated: number;
|
|
1686
|
+
}) => void;
|
|
1687
|
+
}
|
|
1688
|
+
/** Per-request result + serving timings (all performance.now() ms). */
|
|
1689
|
+
interface ServeResult {
|
|
1690
|
+
requestId: number;
|
|
1691
|
+
text: string;
|
|
1692
|
+
tokenIds: number[];
|
|
1693
|
+
tokensGenerated: number;
|
|
1694
|
+
finishReason: "eos" | "max_tokens" | "stop_sequence";
|
|
1695
|
+
/** The adapter name the request decoded with (echoed through), if any. */
|
|
1696
|
+
adapter?: string;
|
|
1697
|
+
/** submit() call time. */
|
|
1698
|
+
submittedAt: number;
|
|
1699
|
+
/** Admission into a batch slot (prefill start). */
|
|
1700
|
+
admittedAt: number;
|
|
1701
|
+
/** First token available (end of prefill). */
|
|
1702
|
+
firstTokenAt: number;
|
|
1703
|
+
finishedAt: number;
|
|
1704
|
+
/** Time-to-first-token from submission (queue wait + prefill). */
|
|
1705
|
+
ttftMs: number;
|
|
1706
|
+
e2eMs: number;
|
|
1707
|
+
}
|
|
1708
|
+
/** Wiring the engine hands the scheduler (see WebGPUEngine.createScheduler). */
|
|
1709
|
+
interface SchedulerHost {
|
|
1710
|
+
executor: Executor;
|
|
1711
|
+
encodeChat(messages: ChatMessage[]): number[];
|
|
1712
|
+
decodeToken(id: number): string;
|
|
1713
|
+
eosTokenId: number | null;
|
|
1714
|
+
eotTokenId: number | null;
|
|
1715
|
+
/** Vocab-prune remap for CPU-sampled tokens (identity when unpruned). */
|
|
1716
|
+
remapToken(id: number): number;
|
|
1717
|
+
/**
|
|
1718
|
+
* Resolve a registered batch-adapter name to its executor adapter id +
|
|
1719
|
+
* parsed deltas. Must throw a clear error for unregistered names (surfaced
|
|
1720
|
+
* at submit() time, before the request is queued).
|
|
1721
|
+
*/
|
|
1722
|
+
resolveAdapter(name: string): {
|
|
1723
|
+
id: number;
|
|
1724
|
+
deltas: LoRADelta[];
|
|
1725
|
+
};
|
|
1726
|
+
/** Serialize against generate()/embed() on the same engine. */
|
|
1727
|
+
acquireLock(): Promise<() => void>;
|
|
1728
|
+
}
|
|
1729
|
+
interface SchedulerStats {
|
|
1730
|
+
submitted: number;
|
|
1731
|
+
completed: number;
|
|
1732
|
+
stepsSubmitted: number;
|
|
1733
|
+
prefillChunks: number;
|
|
1734
|
+
tokensGenerated: number;
|
|
1735
|
+
/** Wall-clock ms the serving loop was live (work outstanding). */
|
|
1736
|
+
busyMs: number;
|
|
1737
|
+
peakActiveLanes: number;
|
|
1738
|
+
}
|
|
1739
|
+
declare class RequestScheduler {
|
|
1740
|
+
private readonly host;
|
|
1741
|
+
private readonly batchSize;
|
|
1742
|
+
private readonly maxSeqLen;
|
|
1743
|
+
private readonly blockSize;
|
|
1744
|
+
private readonly prefillChunk;
|
|
1745
|
+
/** Free KV blocks kept unreserved as safety headroom. */
|
|
1746
|
+
private readonly blockHeadroom;
|
|
1747
|
+
private nextId;
|
|
1748
|
+
private queue;
|
|
1749
|
+
/** slot → occupying request (prefilling or decoding). */
|
|
1750
|
+
private slots;
|
|
1751
|
+
/** Sum of blocksReserved across live (admitted, unfinished) requests. */
|
|
1752
|
+
private reservedBlocks;
|
|
1753
|
+
private prefilling;
|
|
1754
|
+
private inFlight;
|
|
1755
|
+
/**
|
|
1756
|
+
* The slot→adapter-id assignment last uploaded to the executor (null =
|
|
1757
|
+
* nothing uploaded this batch session). Steps only re-upload when the
|
|
1758
|
+
* assignment changes (admissions/retirements), and the very first upload is
|
|
1759
|
+
* skipped entirely while every lane is bare base — an adapter-free serving
|
|
1760
|
+
* session never touches the LoRA machinery.
|
|
1761
|
+
*/
|
|
1762
|
+
private lastLaneAdapterIds;
|
|
1763
|
+
private stepCounter;
|
|
1764
|
+
private loop;
|
|
1765
|
+
private closed;
|
|
1766
|
+
private stats;
|
|
1767
|
+
constructor(host: SchedulerHost, options?: {
|
|
1768
|
+
prefillChunkTokens?: number;
|
|
1769
|
+
});
|
|
1770
|
+
/**
|
|
1771
|
+
* Submit a request; resolves with the full result when it finishes.
|
|
1772
|
+
* Tokens stream through options.onToken as they are decoded.
|
|
1773
|
+
*/
|
|
1774
|
+
submit(prompt: string | ChatMessage[], options?: ServeOptions): Promise<ServeResult>;
|
|
1775
|
+
/** Resolves when every submitted request has completed. */
|
|
1776
|
+
drain(): Promise<void>;
|
|
1777
|
+
/** Stop accepting requests; resolves when in-flight work drains. */
|
|
1778
|
+
close(): Promise<void>;
|
|
1779
|
+
getStats(): SchedulerStats;
|
|
1780
|
+
private pump;
|
|
1781
|
+
private hasWork;
|
|
1782
|
+
private activeMask;
|
|
1783
|
+
/** The serving loop. Holds the engine generation lock while live. */
|
|
1784
|
+
private run;
|
|
1785
|
+
/**
|
|
1786
|
+
* Upload the slot→adapter assignment for the step about to be submitted,
|
|
1787
|
+
* when it differs from the last uploaded one. Called immediately before
|
|
1788
|
+
* every decode-step submit: `queue.writeBuffer` is queue-ordered, so steps
|
|
1789
|
+
* already in flight keep the assignment they were submitted under, and a
|
|
1790
|
+
* retiring lane's slot can never leak its adapter into the slot's next
|
|
1791
|
+
* occupant — the next occupant's assignment is re-derived from `slots`
|
|
1792
|
+
* before its first decode step. Masked (free/prefilling) lanes are pinned
|
|
1793
|
+
* to bare base; their output is discarded regardless.
|
|
1794
|
+
*/
|
|
1795
|
+
private syncLaneAdapters;
|
|
1796
|
+
/**
|
|
1797
|
+
* Admission: move the next queued request into a free slot when (a) a slot
|
|
1798
|
+
* is free, (b) no other prefill is staged (the single-sequence prefill path
|
|
1799
|
+
* stages SSM state in singleton buffers — one at a time), and (c) its
|
|
1800
|
+
* worst-case block need fits the unreserved KV pool with headroom.
|
|
1801
|
+
*/
|
|
1802
|
+
private admit;
|
|
1803
|
+
/**
|
|
1804
|
+
* Run one bounded prefill chunk for the staged request on the
|
|
1805
|
+
* single-sequence path (into its slot's block-table row via tableBase).
|
|
1806
|
+
* Intermediate chunks are submit-only; the final chunk reads logits, adopts
|
|
1807
|
+
* the singleton SSM state into the slot's pool row, seeds the lane's first
|
|
1808
|
+
* token, and flips the lane to decoding.
|
|
1809
|
+
*/
|
|
1810
|
+
private prefillChunkStep;
|
|
1811
|
+
/** Read the oldest in-flight step and deliver its tokens. */
|
|
1812
|
+
private readStep;
|
|
1813
|
+
/** Append one generated token; finish + retire the lane when terminal. */
|
|
1814
|
+
private consumeToken;
|
|
1815
|
+
/** Retire the lane immediately: free KV blocks, open the slot, resolve. */
|
|
1816
|
+
private finish;
|
|
1817
|
+
}
|
|
1818
|
+
//#endregion
|
|
1819
|
+
//#region src/gpu/batch/swarm.d.ts
|
|
1820
|
+
/** Per-request options a swarm sub-task may set (the member picks the adapter). */
|
|
1821
|
+
type SwarmTaskOptions = Omit<ServeOptions, "adapter">;
|
|
1822
|
+
/** One planned sub-task: which member decodes which prompt. */
|
|
1823
|
+
interface SwarmTask {
|
|
1824
|
+
/** Member name (a key of the swarm's `members`). */
|
|
1825
|
+
member: string;
|
|
1826
|
+
prompt: string | ChatMessage[];
|
|
1827
|
+
/** Per-task overrides merged over the run's options. */
|
|
1828
|
+
options?: SwarmTaskOptions;
|
|
1829
|
+
}
|
|
1830
|
+
/** One member's completed sub-task (plan order is preserved). */
|
|
1831
|
+
interface SwarmMemberResult {
|
|
1832
|
+
member: string;
|
|
1833
|
+
result: ServeResult;
|
|
1834
|
+
}
|
|
1835
|
+
/**
|
|
1836
|
+
* Swarm configuration.
|
|
1837
|
+
*
|
|
1838
|
+
* @template R The reduced result type (defaults to the member-keyed map the
|
|
1839
|
+
* default reduce produces).
|
|
1840
|
+
*/
|
|
1841
|
+
interface SwarmOptions<R = Record<string, ServeResult>> {
|
|
1842
|
+
/**
|
|
1843
|
+
* Member name → adapter source (`hf:owner/repo`, URL, or `file:./dir`), or
|
|
1844
|
+
* null for a member that decodes on the bare base. Every non-null source is
|
|
1845
|
+
* registered on the engine (factors packed into shared GPU buffers once).
|
|
1846
|
+
*/
|
|
1847
|
+
members: Record<string, string | null>;
|
|
1848
|
+
/**
|
|
1849
|
+
* input → sub-tasks routed to members. Default: broadcast the input to
|
|
1850
|
+
* every member (one sub-task each, declaration order).
|
|
1851
|
+
*/
|
|
1852
|
+
plan?: (input: string, members: string[]) => SwarmTask[] | Promise<SwarmTask[]>;
|
|
1853
|
+
/**
|
|
1854
|
+
* Joined sub-results → final answer. Default: a member-keyed map of
|
|
1855
|
+
* ServeResults (a duplicated member's later tasks get `name#2`, `name#3`…).
|
|
1856
|
+
*/
|
|
1857
|
+
reduce?: (results: SwarmMemberResult[]) => R | Promise<R>;
|
|
1858
|
+
}
|
|
1859
|
+
/**
|
|
1860
|
+
* A swarm handle: registered members + plan/reduce over the engine's
|
|
1861
|
+
* scheduler. Create via {@link WebGPUEngine.createSwarm}.
|
|
1862
|
+
*/
|
|
1863
|
+
declare class Swarm<R = Record<string, ServeResult>> {
|
|
1864
|
+
private readonly scheduler;
|
|
1865
|
+
private readonly memberSources;
|
|
1866
|
+
private readonly plan;
|
|
1867
|
+
private readonly reduce;
|
|
1868
|
+
constructor(scheduler: RequestScheduler, options: SwarmOptions<R>);
|
|
1869
|
+
/** Member names, declaration order. */
|
|
1870
|
+
get members(): string[];
|
|
1871
|
+
/**
|
|
1872
|
+
* Run one input through the swarm: plan → concurrent fan-out (one batched
|
|
1873
|
+
* pass across the scheduler's lanes) → reduce.
|
|
1874
|
+
*
|
|
1875
|
+
* @param input The task input handed to the plan (default plan: broadcast
|
|
1876
|
+
* it to every member as the prompt).
|
|
1877
|
+
* @param options Per-request options applied to every sub-task (a task's
|
|
1878
|
+
* own `options` win field-by-field).
|
|
1879
|
+
*/
|
|
1880
|
+
run(input: string, options?: SwarmTaskOptions): Promise<R>;
|
|
1881
|
+
/**
|
|
1882
|
+
* Mode-A access: decode one prompt with a single member (no plan/reduce).
|
|
1883
|
+
* Still rides the shared scheduler, so it batches with any concurrent work.
|
|
1884
|
+
*/
|
|
1885
|
+
generate(member: string, prompt: string | ChatMessage[], options?: SwarmTaskOptions): Promise<ServeResult>;
|
|
1886
|
+
private submitFor;
|
|
1887
|
+
}
|
|
1888
|
+
//#endregion
|
|
1889
|
+
//#region src/gpu/sampler.d.ts
|
|
1890
|
+
/**
|
|
1891
|
+
* CPU-side token sampling from logits.
|
|
1892
|
+
*
|
|
1893
|
+
* Applies temperature, top-k, and top-p (nucleus) filtering,
|
|
1894
|
+
* then samples from the resulting probability distribution.
|
|
1895
|
+
*
|
|
1896
|
+
* Uses typed arrays and min-heap for zero-allocation top-K selection.
|
|
1897
|
+
* For vocab_size ~152K with topK=50, this avoids creating 152K JS tuples.
|
|
1898
|
+
*/
|
|
1899
|
+
interface SamplingParams {
|
|
1900
|
+
temperature?: number;
|
|
1901
|
+
topK?: number;
|
|
1902
|
+
topP?: number;
|
|
1903
|
+
repetitionPenalty?: number;
|
|
1904
|
+
/**
|
|
1905
|
+
* Seed for the sampling RNG. When set, sampled decoding is reproducible:
|
|
1906
|
+
* the same params object draws a deterministic random stream (seeded on its
|
|
1907
|
+
* first use — pass a fresh options object per generation for run-to-run
|
|
1908
|
+
* reproducibility). Unset: Math.random() (non-deterministic, the default).
|
|
1909
|
+
*/
|
|
1910
|
+
seed?: number;
|
|
1911
|
+
}
|
|
1912
|
+
//#endregion
|
|
1913
|
+
//#region src/gpu/vision-preprocess.d.ts
|
|
1914
|
+
/**
|
|
1915
|
+
* Host-side vision preprocessing for the Qwen3.5 ViT.
|
|
1916
|
+
*
|
|
1917
|
+
* The learned position embeddings (bilinear-interpolated over the patch grid)
|
|
1918
|
+
* and the 2D rotary cos/sin tables are functions of the image grid (t, h, w)
|
|
1919
|
+
* ONLY — not of the model weights or pixel values. They are cheap to compute on
|
|
1920
|
+
* the CPU and fed to the GPU graph as input activations, keeping the graph to
|
|
1921
|
+
* weight-dependent math while staying byte-identical to HF transformers.
|
|
1922
|
+
*
|
|
1923
|
+
* Ports (verified against transformers 5.12 vision_utils.py + modeling_qwen3_5):
|
|
1924
|
+
* - get_vision_bilinear_indices_and_weights → buildPosEmbeds()
|
|
1925
|
+
* - get_vision_position_ids → buildPositionIds()
|
|
1926
|
+
* - Qwen3_5VisionRotaryEmbedding → buildRotaryCosSin()
|
|
1927
|
+
*
|
|
1928
|
+
* Patch ordering: both the pos-embed gather and the position ids reorder patches
|
|
1929
|
+
* into spatial_merge_size×spatial_merge_size groups, matching the HF image
|
|
1930
|
+
* processor's output ordering, so the merger's [N,h]→[N/u,h*u] reshape lines up.
|
|
1931
|
+
*/
|
|
1932
|
+
interface VisionGridConfig {
|
|
1933
|
+
hiddenSize: number;
|
|
1934
|
+
numHeads: number;
|
|
1935
|
+
numPositionEmbeddings: number;
|
|
1936
|
+
spatialMergeSize: number;
|
|
1937
|
+
ropeTheta?: number;
|
|
1938
|
+
}
|
|
1939
|
+
interface VisionPositionTensors {
|
|
1940
|
+
/** [N, hidden_size] bilinear-interpolated learned position embeddings. */
|
|
1941
|
+
posEmbeds: Float32Array;
|
|
1942
|
+
/** [N, head_dim] rotary cosines. */
|
|
1943
|
+
cos: Float32Array;
|
|
1944
|
+
/** [N, head_dim] rotary sines. */
|
|
1945
|
+
sin: Float32Array;
|
|
1946
|
+
numPatches: number;
|
|
1947
|
+
}
|
|
1948
|
+
/**
|
|
1949
|
+
* Build bilinear-interpolated learned position embeddings [N, hidden].
|
|
1950
|
+
* posEmbedTable is the raw pos_embed.weight [num_position_embeddings, hidden].
|
|
1951
|
+
*/
|
|
1952
|
+
declare function buildPosEmbeds(gridTHW: [number, number, number], posEmbedTable: Float32Array, cfg: VisionGridConfig): Float32Array;
|
|
1953
|
+
/**
|
|
1954
|
+
* Build the reordered (row, col) position ids [N, 2] for rotary, matching
|
|
1955
|
+
* get_vision_position_ids.
|
|
1956
|
+
*/
|
|
1957
|
+
declare function buildPositionIds(gridTHW: [number, number, number], merge: number): Int32Array;
|
|
1958
|
+
/**
|
|
1959
|
+
* Build rotary cos/sin tables [N, head_dim] from position ids, matching
|
|
1960
|
+
* Qwen3_5VisionRotaryEmbedding + the cat((rotary, rotary)) in VisionModel.forward.
|
|
1961
|
+
*
|
|
1962
|
+
* rotary_pos_emb(position_ids) = (position_ids[..,None] * inv_freq).flatten(1)
|
|
1963
|
+
* where inv_freq has length (head_dim/2)/2 = head_dim/4, computed over dim=head_dim/2.
|
|
1964
|
+
* For each token the two position components (h, w) each produce head_dim/4 freqs,
|
|
1965
|
+
* concatenated → head_dim/2, then duplicated → head_dim for cos/sin.
|
|
1966
|
+
*/
|
|
1967
|
+
declare function buildRotaryCosSin(positionIds: Int32Array, headDim: number, theta?: number): {
|
|
1968
|
+
cos: Float32Array;
|
|
1969
|
+
sin: Float32Array;
|
|
1970
|
+
numPatches: number;
|
|
1971
|
+
};
|
|
1972
|
+
/**
|
|
1973
|
+
* Build all host position tensors for a single image grid in one call.
|
|
1974
|
+
*/
|
|
1975
|
+
declare function buildVisionPositionTensors(gridTHW: [number, number, number], posEmbedTable: Float32Array, cfg: VisionGridConfig): VisionPositionTensors;
|
|
1976
|
+
interface Gemma4VisionGridConfig {
|
|
1977
|
+
hiddenSize: number;
|
|
1978
|
+
numHeads: number;
|
|
1979
|
+
headDim: number;
|
|
1980
|
+
ropeTheta: number;
|
|
1981
|
+
poolingKernelSize: number;
|
|
1982
|
+
}
|
|
1983
|
+
interface Gemma4VisionPositionTensors {
|
|
1984
|
+
/** [N, hidden] axial position embeddings (table[0][x] + table[1][y]). */
|
|
1985
|
+
posEmbeds: Float32Array;
|
|
1986
|
+
/** [N, headDim] axial rotary cosines. */
|
|
1987
|
+
cos: Float32Array;
|
|
1988
|
+
/** [N, headDim] axial rotary sines. */
|
|
1989
|
+
sin: Float32Array;
|
|
1990
|
+
/** [Np, N] average-pooling matrix (1/k² in-cell, 0 elsewhere). */
|
|
1991
|
+
poolMatrix: Float32Array;
|
|
1992
|
+
/** number of patches N (= gridH*gridW). */
|
|
1993
|
+
numPatches: number;
|
|
1994
|
+
/** number of pooled (soft) tokens Np (= ceil(gridH/k)*ceil(gridW/k)). */
|
|
1995
|
+
numPooled: number;
|
|
1996
|
+
}
|
|
1997
|
+
/**
|
|
1998
|
+
* Build axial learned position embeddings [N, hidden] from the [2, posSize, hidden]
|
|
1999
|
+
* table: pos[p] = table[0][x_p] + table[1][y_p]. Direct lookup, no interpolation
|
|
2000
|
+
* (HF F.embedding on clamped positions).
|
|
2001
|
+
*/
|
|
2002
|
+
declare function buildGemma4PosEmbeds(gridH: number, gridW: number, posEmbedTable: Float32Array,
|
|
2003
|
+
// [2, posSize, hidden] flattened
|
|
2004
|
+
hidden: number, posSize: number): Float32Array;
|
|
2005
|
+
/**
|
|
2006
|
+
* Build the 2D axial rotary cos/sin tables [N, headDim].
|
|
2007
|
+
* spatial_dim = headDim / 2; inv_freq[j] = 1/theta^((2j)/spatial_dim), j in [0, spatial_dim/2)
|
|
2008
|
+
* per spatial dim: f = pos * inv_freq (spatial_dim/2 values); emb = cat(f, f) (spatial_dim values)
|
|
2009
|
+
* cos/sin = cat([emb_x, emb_y]) → headDim values, layout [fx,fx,fy,fy].
|
|
2010
|
+
* Applied with the global-half rotate_half kernel (ApplyRotaryEmb), which computes
|
|
2011
|
+
* out = x*cos + rotate_half(x)*sin element-wise — exact for this layout.
|
|
2012
|
+
*/
|
|
2013
|
+
declare function buildGemma4RotaryCosSin(gridH: number, gridW: number, headDim: number, theta: number): {
|
|
2014
|
+
cos: Float32Array;
|
|
2015
|
+
sin: Float32Array;
|
|
2016
|
+
};
|
|
2017
|
+
/**
|
|
2018
|
+
* Build the [Np, N] average-pooling matrix for k×k spatial pooling over the real
|
|
2019
|
+
* (unpadded) grid, matching modeling_gemma4's kernel_idxs/one_hot pooling:
|
|
2020
|
+
* cell(p) = floor(x_p/k) + ceil(gridW/k) * floor(y_p/k)
|
|
2021
|
+
* poolMatrix[cell, p] = 1/k² (so pooled = poolMatrix @ hidden = mean over the k×k block)
|
|
2022
|
+
* Np = ceil(gridH/k) * ceil(gridW/k). Each pooled cell averages exactly the patches
|
|
2023
|
+
* that fall in it (edge cells with fewer than k² patches still divide by k², matching
|
|
2024
|
+
* HF's fixed 1/k² normalization).
|
|
2025
|
+
*/
|
|
2026
|
+
declare function buildGemma4PoolMatrix(gridH: number, gridW: number, k: number): {
|
|
2027
|
+
poolMatrix: Float32Array;
|
|
2028
|
+
numPooled: number;
|
|
2029
|
+
};
|
|
2030
|
+
/**
|
|
2031
|
+
* Build all Gemma 4 vision host tensors for one image grid in one call.
|
|
2032
|
+
* `posEmbedTable` is the raw [2, posSize, hidden] flattened table.
|
|
2033
|
+
*/
|
|
2034
|
+
declare function buildGemma4VisionPositionTensors(gridH: number, gridW: number, posEmbedTable: Float32Array, posSize: number, cfg: Gemma4VisionGridConfig): Gemma4VisionPositionTensors;
|
|
2035
|
+
/** Gemma 4 image processor config (from processor_config.json). */
|
|
2036
|
+
declare const GEMMA4_IMAGE_PROCESSOR: ImageProcessorConfig;
|
|
2037
|
+
interface Gemma4PreprocessedImage {
|
|
2038
|
+
/** Flattened patches [N, 3·patch²] row-major (row-major patch grid). */
|
|
2039
|
+
patches: Float32Array;
|
|
2040
|
+
/** Patch grid (gridH, gridW). */
|
|
2041
|
+
gridHW: [number, number];
|
|
2042
|
+
}
|
|
2043
|
+
/**
|
|
2044
|
+
* Preprocess a decoded RGB image for the Gemma 4 ViT: aspect-preserving resize so
|
|
2045
|
+
* the patch grid is ≤ max_soft_tokens·k² patches and H,W divisible by k·patch,
|
|
2046
|
+
* rescale ×1/255 (no normalize), patchify row-major into [N, 3·16·16].
|
|
2047
|
+
*
|
|
2048
|
+
* @param pixels row-major HWC RGB (0..255), length width*height*3.
|
|
2049
|
+
*/
|
|
2050
|
+
declare function preprocessImageGemma4(pixels: Float32Array | Uint8ClampedArray | Uint8Array, width: number, height: number, maxSoftTokens?: number, poolingKernelSize?: number, patchSize?: number): Gemma4PreprocessedImage;
|
|
2051
|
+
interface ImageProcessorConfig {
|
|
2052
|
+
patchSize: number;
|
|
2053
|
+
temporalPatchSize: number;
|
|
2054
|
+
mergeSize: number;
|
|
2055
|
+
imageMean: [number, number, number];
|
|
2056
|
+
imageStd: [number, number, number];
|
|
2057
|
+
/** rescale factor applied to raw 0..255 pixels before normalization (1/255). */
|
|
2058
|
+
rescaleFactor: number;
|
|
2059
|
+
/** min total pixels after resize (shortest_edge). */
|
|
2060
|
+
minPixels: number;
|
|
2061
|
+
/** max total pixels after resize (longest_edge). */
|
|
2062
|
+
maxPixels: number;
|
|
2063
|
+
}
|
|
2064
|
+
declare const QWEN3_5_IMAGE_PROCESSOR: ImageProcessorConfig;
|
|
2065
|
+
interface PreprocessedImage {
|
|
2066
|
+
/** Flattened patches [N, 1536] in the spatial-merge order encodeImage expects. */
|
|
2067
|
+
patches: Float32Array;
|
|
2068
|
+
/** (t, h, w) patch grid. t=1 for a single image, h/w in patch units. */
|
|
2069
|
+
gridTHW: [number, number, number];
|
|
2070
|
+
}
|
|
2071
|
+
/**
|
|
2072
|
+
* Qwen2-VL smart-resize: round H and W to multiples of factor=patch*merge,
|
|
2073
|
+
* keeping aspect ratio and clamping the total pixel budget to [minPixels, maxPixels].
|
|
2074
|
+
* Matches transformers.models.qwen2_vl.image_processing.smart_resize.
|
|
2075
|
+
*/
|
|
2076
|
+
declare function smartResize(height: number, width: number, factor: number, minPixels: number, maxPixels: number): [number, number];
|
|
2077
|
+
/**
|
|
2078
|
+
* Preprocess a decoded RGB image into the [N, 1536] patch tensor + grid_thw that
|
|
2079
|
+
* `encodeImage()` expects, matching the HF Qwen2-VL image processor:
|
|
2080
|
+
* smart_resize → rescale (×1/255) → normalize → temporal-pair (×temporal_patch_size)
|
|
2081
|
+
* → patchify into spatial_merge×spatial_merge blocks → flatten to [N, C·T·P·P].
|
|
2082
|
+
*
|
|
2083
|
+
* @param pixels row-major HWC RGB, length width*height*3. Values 0..255 (default)
|
|
2084
|
+
* or already 0..1 if `rescaled` is true.
|
|
2085
|
+
* @param width source pixel width
|
|
2086
|
+
* @param height source pixel height
|
|
2087
|
+
*/
|
|
2088
|
+
declare function preprocessImage(pixels: Float32Array | Uint8ClampedArray | Uint8Array, width: number, height: number, cfg?: ImageProcessorConfig, rescaled?: boolean): PreprocessedImage;
|
|
2089
|
+
declare function buildMRoPEPositionIds(inputIds: Int32Array | Uint32Array | number[], imageGrids: Array<[number, number, number]>, imageTokenId: number, mergeSize: number): Int32Array;
|
|
2090
|
+
/**
|
|
2091
|
+
* Per-pair frequency→dimension assignment for interleaved M-RoPE, matching
|
|
2092
|
+
* Qwen3_5TextRotaryEmbedding.apply_interleaved_mrope. For pair index i in
|
|
2093
|
+
* [0, sum(section)) the position component is section-cyclic: T,H,W,T,H,W,...
|
|
2094
|
+
* but each component capped at its section count. Returns an array of length
|
|
2095
|
+
* (rope_dim/2) with values 0=T, 1=H, 2=W.
|
|
2096
|
+
*/
|
|
2097
|
+
declare function mropeFreqDims(mropeSection: [number, number, number]): Int32Array;
|
|
2098
|
+
/**
|
|
2099
|
+
* Build the interleaved-M-RoPE cos/sin tables [seq, rope_dim] from 3D position
|
|
2100
|
+
* ids, matching Qwen3_5TextRotaryEmbedding.forward:
|
|
2101
|
+
* freqs[d][i] = pos[d] * inv_freq[i], inv_freq[i] = 1/theta^(2i/rope_dim)
|
|
2102
|
+
* freq[i] picks component mropeFreqDims[i]; emb = cat(freqs, freqs).
|
|
2103
|
+
* cos/sin have length seq*rope_dim. For text-only (all 3 pos rows equal) this
|
|
2104
|
+
* reduces exactly to standard 1D partial RoPE.
|
|
2105
|
+
*
|
|
2106
|
+
* @param positionIds3 [3, seq] as produced by buildMRoPEPositionIds.
|
|
2107
|
+
* @param ropeDim number of rotated dims per head (head_dim * partial_factor).
|
|
2108
|
+
*/
|
|
2109
|
+
declare function buildMRoPECosSin(positionIds3: Int32Array, seq: number, ropeDim: number, theta: number, mropeSection: [number, number, number]): {
|
|
2110
|
+
cos: Float32Array;
|
|
2111
|
+
sin: Float32Array;
|
|
2112
|
+
};
|
|
2113
|
+
//#endregion
|
|
2114
|
+
//#region src/gpu/defaults.d.ts
|
|
2115
|
+
/**
|
|
2116
|
+
* Default model per capability. Kept in its own tiny module (no heavy imports)
|
|
2117
|
+
* so the React hooks can resolve defaults without statically pulling in the GPU
|
|
2118
|
+
* engine — they import the engine dynamically to stay light.
|
|
2119
|
+
*/
|
|
2120
|
+
declare const DEFAULT_MODELS: {
|
|
2121
|
+
/** Text generation (also the vision-capable checkpoint). */
|
|
2122
|
+
readonly text: "mlx-community/Qwen3.5-0.8B-4bit";
|
|
2123
|
+
/** Image understanding — same checkpoint, vision tower built on demand. */
|
|
2124
|
+
readonly vision: "mlx-community/Qwen3.5-0.8B-4bit";
|
|
2125
|
+
/** Text embeddings. */
|
|
2126
|
+
readonly embedding: "mlx-community/embeddinggemma-300m-4bit";
|
|
2127
|
+
/** Text-to-speech (Apache-2.0; LFM2-350M backbone + NanoCodec). Small/fast default. */
|
|
2128
|
+
readonly tts: "nineninesix/kani-tts-450m-0.2-ft";
|
|
2129
|
+
/**
|
|
2130
|
+
* Selectable many-voices TTS (OuteTTS-1.0-0.6B: Qwen3-0.6B codec-LM + DAC.speech).
|
|
2131
|
+
* Pass this as the `repo`/`model` to engine.speak()/useTTS to route to OuteTTS.
|
|
1819
2132
|
*/
|
|
2133
|
+
readonly ttsOute: "OuteAI/OuteTTS-1.0-0.6B";
|
|
1820
2134
|
/**
|
|
1821
|
-
*
|
|
1822
|
-
*
|
|
1823
|
-
*
|
|
1824
|
-
* race-free; prefill (T>1) keeps the two-dispatch path.
|
|
2135
|
+
* Selectable describe-your-voice TTS (Parler-TTS-mini-v1: Flan-T5 encoder +
|
|
2136
|
+
* cross-attending decoder LM + dac_44khz). Pass this as the `repo`/`model` to
|
|
2137
|
+
* engine.speak()/useTTS with a `describeVoice` natural-language description.
|
|
1825
2138
|
*/
|
|
1826
|
-
|
|
2139
|
+
readonly ttsParler: "parler-tts/parler-tts-mini-v1";
|
|
2140
|
+
/** Speech-to-text. */
|
|
2141
|
+
readonly stt: "UsefulSensors/moonshine-base";
|
|
2142
|
+
};
|
|
2143
|
+
declare const OUTETTS_ASSETS: {
|
|
2144
|
+
/** Folded DAC.speech.v1.0 encoder+decoder weights (prep-dac-speech.py output). */
|
|
2145
|
+
readonly codecWeightsURL: "https://pub-0522ecec25fc412bb35a523f3e3f63c6.r2.dev/tryhamster/gerbil-outetts-assets/dac-speech-folded.safetensors";
|
|
2146
|
+
/** Base URL for the preset speaker JSONs (`{name}.json`). */
|
|
2147
|
+
readonly speakersBaseURL: "https://pub-0522ecec25fc412bb35a523f3e3f63c6.r2.dev/tryhamster/gerbil-outetts-assets/speakers";
|
|
2148
|
+
/** The default preset voice. */
|
|
2149
|
+
readonly defaultVoice: "en-female-1-neutral";
|
|
2150
|
+
};
|
|
2151
|
+
/**
|
|
2152
|
+
* Bundled OuteTTS preset voices.
|
|
2153
|
+
*
|
|
2154
|
+
* Provenance (honest labelling):
|
|
2155
|
+
* - `en-female-1-neutral` is the ONE AUTHENTIC OuteAI voice — the OuteTTS v3 bundled
|
|
2156
|
+
* default speaker (default_speakers/en-female-1-neutral), copied verbatim.
|
|
2157
|
+
* - `gerbil-voice-2..6` are DERIVED voices: distinct timbres produced by running
|
|
2158
|
+
* pitch/tempo-varied renderings of that authentic clip back through the REAL
|
|
2159
|
+
* DAC.speech encode + prosody pipeline (scripts/engine/stage-outetts-assets.py).
|
|
2160
|
+
* They are genuine DAC-encoded OuteSpeakers, but they are NOT OuteAI's own
|
|
2161
|
+
* `en-male-*` voices, so they are not labelled as such. (Earlier builds mislabelled
|
|
2162
|
+
* them as `en-male-2-neutral` etc.; renamed for honesty.)
|
|
2163
|
+
*/
|
|
2164
|
+
declare const OUTETTS_PRESET_VOICES: readonly ["en-female-1-neutral", "gerbil-voice-2", "gerbil-voice-3", "gerbil-voice-4", "gerbil-voice-5", "gerbil-voice-6"];
|
|
2165
|
+
/**
|
|
2166
|
+
* True when `repo`/`architecture` identify a text-to-speech checkpoint that must
|
|
2167
|
+
* run through a dedicated native TTS engine (KaniTTS, OuteTTS, or Parler-TTS)
|
|
2168
|
+
* rather than the single-graph base text loader. Mirrors the routing used by
|
|
2169
|
+
* `WebGPUEngine.speak()` / `Gerbil.ensureNativeTTSEngine`:
|
|
2170
|
+
* - Parler → encoder-decoder (`ParlerTTSForConditionalGeneration`)
|
|
2171
|
+
* - OuteTTS → Qwen3-0.6B codec-LM
|
|
2172
|
+
* - Kani → conv-backbone codec-LM (`KaniTTS2ForCausalLM`, or the Apache
|
|
2173
|
+
* default which is a plain `Lfm2ForCausalLM` — hence the repo-name
|
|
2174
|
+
* gate; the architecture string alone does not identify it).
|
|
2175
|
+
* At `create()` time only the repo is known (the architecture is discovered
|
|
2176
|
+
* during load), so the repo-name signals below cover all three DEFAULT_MODELS
|
|
2177
|
+
* and conventionally-named TTS repos; the architecture arg is honored when a
|
|
2178
|
+
* caller already has it.
|
|
2179
|
+
*/
|
|
2180
|
+
declare function isTTSRepo(repo?: string, architecture?: string): boolean;
|
|
2181
|
+
/** Resolve the model repo for a set of options, falling back to the defaults. */
|
|
2182
|
+
declare function resolveDefaultRepo(opts: {
|
|
2183
|
+
repo?: string;
|
|
2184
|
+
embedding?: boolean;
|
|
2185
|
+
enableVision?: boolean;
|
|
2186
|
+
}): string;
|
|
2187
|
+
//#endregion
|
|
2188
|
+
//#region src/gpu/architectures/gemma4_vision.d.ts
|
|
2189
|
+
interface Gemma4VisionGraphInfo {
|
|
2190
|
+
hiddenSize: number;
|
|
2191
|
+
numHeads: number;
|
|
2192
|
+
headDim: number;
|
|
2193
|
+
depth: number;
|
|
2194
|
+
intermediateSize: number;
|
|
2195
|
+
textHidden: number;
|
|
2196
|
+
patchSize: number;
|
|
2197
|
+
patchDim: number;
|
|
2198
|
+
poolingKernelSize: number;
|
|
2199
|
+
ropeTheta: number;
|
|
2200
|
+
rmsNormEps: number;
|
|
2201
|
+
}
|
|
2202
|
+
/**
|
|
2203
|
+
* Resolve the Gemma 4 vision dims from a raw HF config. Accepts either the
|
|
2204
|
+
* top-level config (reads `.vision_config` + `.text_config.hidden_size`) or a
|
|
2205
|
+
* bare vision_config (then `textHidden` falls back to the projector row count if
|
|
2206
|
+
* present, else hidden). Family-general — no E2B constants.
|
|
2207
|
+
*/
|
|
2208
|
+
declare function resolveGemma4VisionInfo(rawConfig: Record<string, unknown>): Gemma4VisionGraphInfo;
|
|
2209
|
+
/**
|
|
2210
|
+
* Dequantize an MLX affine-int4 weight to a plain f32 [rows, cols] matrix.
|
|
2211
|
+
* MLX packs 8 int4 values per u32 (low-nibble first); each group of `groupSize`
|
|
2212
|
+
* columns shares one scale + bias: w[r,c] = scale[r, c/gs] * q + bias[r, c/gs].
|
|
2213
|
+
* Used for the Gemma 4 multimodal projector (`embed_vision.embedding_projection`)
|
|
2214
|
+
* in MLX-4bit checkpoints, where (unlike the BF16 ViT body) the projector is int4.
|
|
2215
|
+
*/
|
|
2216
|
+
declare function dequantizeMLXProjection(packed: Uint32Array, scales: Float32Array, biases: Float32Array, rows: number, cols: number, groupSize: number): Float32Array;
|
|
2217
|
+
/**
|
|
2218
|
+
* If the Gemma 4 multimodal projector arrived as an MLX affine-int4 triplet
|
|
2219
|
+
* (`embed_vision.embedding_projection.{weight(U32), scales, biases}`), dequantize
|
|
2220
|
+
* it in-place to a plain f32 `embed_vision.embedding_projection.weight` and drop
|
|
2221
|
+
* the scales/biases, so the vision graph's plain MatMul on the projector works for
|
|
2222
|
+
* MLX-4bit checkpoints too. No-op for BF16 (HF) checkpoints (weight already f32).
|
|
2223
|
+
*/
|
|
2224
|
+
declare function dequantizeGemma4VisionProjection(weights: Map<string, {
|
|
2225
|
+
data: ArrayBufferView;
|
|
2226
|
+
shape: number[];
|
|
2227
|
+
}>, groupSize: number, rows: number, cols: number): void;
|
|
2228
|
+
/**
|
|
2229
|
+
* Patch the ClippedMatMul nodes of a Gemma 4 vision graph with the calibrated clip
|
|
2230
|
+
* scalars from the checkpoint (Gemma4ClippableLinear's per-tensor input/output
|
|
2231
|
+
* min/max buffers), then drop those scalar tensors from the weights map so the
|
|
2232
|
+
* vision executor doesn't try to upload them as GPU buffers. Call BEFORE
|
|
2233
|
+
* VisionExecutor.uploadWeights(). Missing scalars default to ±inf (clip = identity),
|
|
2234
|
+
* so a checkpoint without calibration still loads.
|
|
2235
|
+
*/
|
|
2236
|
+
declare function patchGemma4VisionClips(graph: ModelGraph, weights: Map<string, {
|
|
2237
|
+
data: ArrayBufferView;
|
|
2238
|
+
shape: number[];
|
|
2239
|
+
}>): void;
|
|
2240
|
+
/**
|
|
2241
|
+
* Build the Gemma 4 ViT graph. Shaped by symbolic "N" (number of patches, runtime)
|
|
2242
|
+
* and "Np" (number of pooled tokens = ceil(grid_h/k)·ceil(grid_w/k), runtime),
|
|
2243
|
+
* resolved from input tensor dims — like the Qwen ViT's "N"/"Nm".
|
|
2244
|
+
*/
|
|
2245
|
+
declare function generateGemma4VisionGraph(rawConfig: Record<string, unknown>): ModelGraph;
|
|
2246
|
+
//#endregion
|
|
2247
|
+
//#region src/gpu/architectures/kani_tts.d.ts
|
|
2248
|
+
/** Parsed KaniTTS2 backbone config (the LFM2 dims + the TTS-specific fields). */
|
|
2249
|
+
interface KaniConfig {
|
|
2250
|
+
textVocabSize: number;
|
|
2251
|
+
vocabSize: number;
|
|
2252
|
+
tokensPerFrame: number;
|
|
2253
|
+
audioStep: number;
|
|
2254
|
+
useLearnableRope: boolean;
|
|
2255
|
+
alphaMin: number;
|
|
2256
|
+
alphaMax: number;
|
|
2257
|
+
speakerEmbDim: number;
|
|
2258
|
+
audioTokensStart: number;
|
|
2259
|
+
startOfSpeech: number;
|
|
2260
|
+
endOfSpeech: number;
|
|
2261
|
+
codebookSize: number;
|
|
1827
2262
|
/**
|
|
1828
|
-
*
|
|
1829
|
-
*
|
|
1830
|
-
*
|
|
1831
|
-
*
|
|
1832
|
-
*
|
|
2263
|
+
* Whether the backbone uses KaniTTS2's frame-level position IDs (the 4 audio
|
|
2264
|
+
* tokens of a frame share one logical position) + learnable per-layer RoPE.
|
|
2265
|
+
* The `KaniTTS2ForCausalLM` checkpoints (e.g. kani-tts-2-en) do; the plain
|
|
2266
|
+
* `Lfm2ForCausalLM` family (e.g. the Apache kani-tts-450m-0.2-ft / kani-tts-370m)
|
|
2267
|
+
* does NOT — those run standard sequential positions with α=1 plain RoPE.
|
|
2268
|
+
* Detected from the presence of `use_learnable_rope`/`tokens_per_frame`/the
|
|
2269
|
+
* `KaniTTS2ForCausalLM` architecture marker in config.json.
|
|
1833
2270
|
*/
|
|
2271
|
+
frameLevelPositions: boolean;
|
|
2272
|
+
}
|
|
2273
|
+
declare function parseKaniConfig(rawConfig: Record<string, unknown>): KaniConfig;
|
|
2274
|
+
/**
|
|
2275
|
+
* Convert the model's flat audio-token stream (the IDs between start/end-of-speech)
|
|
2276
|
+
* into NanoCodec codes [NUM_GROUPS, T]. Mirrors NemoAudioPlayer.get_nano_codes:
|
|
2277
|
+
* reshape [-1, 4]; codes[:,c] -= codebook_size*c; codes -= audio_tokens_start;
|
|
2278
|
+
* transpose → [4, T]. Returns a Uint32Array laid out group-major ([g*T + t]).
|
|
2279
|
+
*/
|
|
2280
|
+
declare function audioTokensToCodes(audioTokenIds: number[], cfg?: KaniConfig): {
|
|
2281
|
+
codes: Uint32Array;
|
|
2282
|
+
numFrames: number;
|
|
2283
|
+
};
|
|
2284
|
+
interface NanoCodecGraphOptions {
|
|
2285
|
+
/** Number of audio frames T (the code grid width). PCM length = T * 1764. */
|
|
2286
|
+
numFrames: number;
|
|
2287
|
+
}
|
|
2288
|
+
declare function generateNanoCodecDecoderGraph(opts: NanoCodecGraphOptions): ModelGraph;
|
|
2289
|
+
/**
|
|
2290
|
+
* Generate the KaniTTS codec-LM backbone graph (LFM2-350M body, full-vocab logits,
|
|
2291
|
+
* per-layer learnable MRoPE). Mirrors generateLfm2Graph block-for-block; the only
|
|
2292
|
+
* difference is that each attention layer rotates Q/K with the MRoPE op fed a
|
|
2293
|
+
* per-layer host cos/sin table instead of the position-counter RoPE op.
|
|
2294
|
+
*
|
|
2295
|
+
* dtype:
|
|
2296
|
+
* • undefined / "f32" — full precision (the checkpoint is bf16/f32 → f32).
|
|
2297
|
+
* • "q4" — weight-only group-wise INT4 (W4A16) for the LFM2 linear projections
|
|
2298
|
+
* (q/k/v/o, gate/up/down, short-conv in/out_proj) via MatMulInt4. The tied
|
|
2299
|
+
* embedding + LM head stay at f32 on purpose: decoder-LM heads are the most
|
|
2300
|
+
* quantization-sensitive part of a codec-LM and the embed table is the single
|
|
2301
|
+
* buffer shared by both, so quantizing it would degrade audio quality for no
|
|
2302
|
+
* load-size win that the int4 projections do not already capture. Backbone
|
|
2303
|
+
* drops from ~1.5 GB (f32) to ~0.37 GB. The NanoCodec decoder is a separate
|
|
2304
|
+
* graph and always runs f32.
|
|
2305
|
+
*/
|
|
2306
|
+
declare function generateKaniTtsGraph(rawConfig: Record<string, unknown>, dtype?: GraphDType, groupSize?: number, kvDtype?: KVDType): ModelGraph;
|
|
2307
|
+
//#endregion
|
|
2308
|
+
//#region src/gpu/architectures/moonshine.d.ts
|
|
2309
|
+
/**
|
|
2310
|
+
* Historical record of the executor-side work the keystone left open; all items
|
|
2311
|
+
* are now implemented (see the STATUS note above). Retained as exported
|
|
2312
|
+
* documentation of the dependency surface.
|
|
2313
|
+
*/
|
|
2314
|
+
declare const MOONSHINE_REMAINING_WORK: readonly ["DONE: Transpose op kernel (conv output [C,L] → encoder input [L,C]).", "DONE: GroupNorm(num_groups=1) over conv channels (weight+bias).", "DONE: Tanh elementwise op (conv1 activation).", "DONE: Interleaved-RoPE variant (Moonshine rotates adjacent dim pairs 2p/2p+1).", "DONE: No-bias LayerNorm variant (Moonshine norms are weight-only).", "DONE: Per-utterance conv graph regeneration from the concrete sample count.", "DONE: Dual-graph executor — encoder once → frozen cross-attn K/V → AR decode.", "DONE: engine.transcribe(pcm) host path (MoonshineSTT): conv → encode → AR decode → detokenize."];
|
|
2315
|
+
interface MoonshineDims {
|
|
2316
|
+
hidden_size: number;
|
|
2317
|
+
enc_layers: number;
|
|
2318
|
+
dec_layers: number;
|
|
2319
|
+
num_heads: number;
|
|
2320
|
+
num_kv_heads: number;
|
|
2321
|
+
head_dim: number;
|
|
2322
|
+
rotary_dim: number;
|
|
2323
|
+
intermediate_size: number;
|
|
2324
|
+
vocab_size: number;
|
|
2325
|
+
rope_base: number;
|
|
2326
|
+
context_length: number;
|
|
2327
|
+
ln_eps: number;
|
|
2328
|
+
}
|
|
2329
|
+
/** Pull and derive Moonshine dimensions from the raw HF config. */
|
|
2330
|
+
declare function parseMoonshineConfig(raw: Record<string, unknown>): MoonshineDims;
|
|
2331
|
+
/** Number of encoder frames produced by the conv frontend for n_samples PCM. */
|
|
2332
|
+
declare function moonshineEncoderFrames(nSamples: number): number;
|
|
2333
|
+
/**
|
|
2334
|
+
* Encoder graph: raw-waveform conv frontend + bidirectional transformer.
|
|
2335
|
+
* @param nSamples concrete PCM sample count (the conv frontend is length-static).
|
|
2336
|
+
* The output tensor "encoder_out" is [T_frames, hidden] and is consumed (after
|
|
2337
|
+
* per-layer K/V projection) as the frozen K/V for the decoder's cross-attention.
|
|
2338
|
+
*/
|
|
2339
|
+
declare function generateMoonshineEncoderGraph(raw: Record<string, unknown>, nSamples: number): ModelGraph;
|
|
2340
|
+
/**
|
|
2341
|
+
* Decoder graph: AR transformer with causal self-attn (KV-cache) + cross-attn to
|
|
2342
|
+
* the frozen encoder output. Built for a single decode step (T=1). The encoder K/V
|
|
2343
|
+
* are supplied as graph inputs "enc_k_layer{i}" / "enc_v_layer{i}" — the host
|
|
2344
|
+
* pre-projects the encoder output through each layer's encoder_attn.k_proj/v_proj
|
|
2345
|
+
* ONCE and binds them frozen for the whole decode (the CrossAttention contract).
|
|
2346
|
+
*/
|
|
2347
|
+
declare function generateMoonshineDecoderGraph(raw: Record<string, unknown>, sEnc: number): ModelGraph;
|
|
2348
|
+
//#endregion
|
|
2349
|
+
//#region src/gpu/architectures/outetts.d.ts
|
|
2350
|
+
/** Parsed OuteTTS backbone config (the Qwen3 dims). Reuses the generic Qwen path. */
|
|
2351
|
+
interface OuteTtsConfig {
|
|
2352
|
+
hidden_size: number;
|
|
2353
|
+
num_layers: number;
|
|
2354
|
+
num_heads: number;
|
|
2355
|
+
num_kv_heads: number;
|
|
2356
|
+
head_dim: number;
|
|
2357
|
+
intermediate_size: number;
|
|
2358
|
+
vocab_size: number;
|
|
2359
|
+
rope_base: number;
|
|
2360
|
+
rms_norm_eps: number;
|
|
2361
|
+
max_position_embeddings: number;
|
|
2362
|
+
}
|
|
2363
|
+
declare function parseOuteTtsConfig(rawConfig: Record<string, unknown>): OuteTtsConfig;
|
|
2364
|
+
/**
|
|
2365
|
+
* Convert the flat generated token stream (the run BETWEEN <|audio_start|> and
|
|
2366
|
+
* <|audio_end|>) into a DAC code grid [NUM_CODEBOOKS, T] (codebook-major [c*T + t]).
|
|
2367
|
+
*
|
|
2368
|
+
* Mirrors outetts v3 `extract_audio_from_tokens`: filter the stream into two ordered
|
|
2369
|
+
* lists by membership of the c1 / c2 ranges, recover the raw code
|
|
2370
|
+
* (id − base), drop the reserved index 1024, and truncate to the shorter list. The
|
|
2371
|
+
* model emits c1,c2 interleaved per frame, but the parse is membership-based (it does
|
|
2372
|
+
* not assume strict pairing) so a stray token cannot desync the grid.
|
|
2373
|
+
*/
|
|
2374
|
+
declare function audioTokensToDacCodes(tokenIds: number[]): {
|
|
2375
|
+
codes: Uint32Array;
|
|
2376
|
+
numFrames: number;
|
|
2377
|
+
};
|
|
2378
|
+
interface DacSpeechGraphOptions {
|
|
2379
|
+
/** Number of audio frames T (the code-grid width). */
|
|
2380
|
+
numFrames: number;
|
|
1834
2381
|
/**
|
|
1835
|
-
*
|
|
1836
|
-
*
|
|
1837
|
-
* the
|
|
2382
|
+
* Number of RVQ codebooks. Default 2 (DAC.speech.v1.0). Parler-TTS's
|
|
2383
|
+
* dac_44khz_8kbps uses 9 — the RVQDequant kernel loops over NCB so the only
|
|
2384
|
+
* change is the codebook/out_proj tensor first-dim and the codes grid height.
|
|
1838
2385
|
*/
|
|
1839
|
-
|
|
1840
|
-
private fuseDualRMSNormRoPEDecodeEntries;
|
|
1841
|
-
private fuseDualRMSNormDecodeEntries;
|
|
2386
|
+
numCodebooks?: number;
|
|
1842
2387
|
/**
|
|
1843
|
-
*
|
|
1844
|
-
*
|
|
2388
|
+
* Decoder upsampling rates (the DecoderBlock strides). Default [8,5,4,2]
|
|
2389
|
+
* (DAC.speech.v1.0). Parler's dac_44khz uses [8,8,4,2] (hop 512). Channels
|
|
2390
|
+
* still halve per stage from `decoderDim`.
|
|
1845
2391
|
*/
|
|
1846
|
-
|
|
1847
|
-
/**
|
|
1848
|
-
|
|
2392
|
+
decoderRates?: readonly number[];
|
|
2393
|
+
/** First-conv output channels (decoder_dim). Default 1536. */
|
|
2394
|
+
decoderDim?: number;
|
|
2395
|
+
/** Quantizer out_proj output / first-conv input dim (latent_dim). Default 1024. */
|
|
2396
|
+
latentDim?: number;
|
|
1849
2397
|
}
|
|
2398
|
+
declare function generateDacSpeechDecoderGraph(opts: DacSpeechGraphOptions): ModelGraph;
|
|
2399
|
+
/**
|
|
2400
|
+
* Final decoded PCM length (samples) for a T-frame code grid, via the DAC stage chain.
|
|
2401
|
+
* `rates` defaults to DAC.speech.v1.0 ([8,5,4,2]); pass [8,8,4,2] for Parler's dac_44khz.
|
|
2402
|
+
*/
|
|
2403
|
+
declare function dacOutputLength(numFrames: number, rates?: readonly number[]): number;
|
|
2404
|
+
declare function generateOuteTtsBackboneGraph(rawConfig: Record<string, unknown>, dtype?: GraphDType, groupSize?: number, kvDtype?: KVDType): ModelGraph;
|
|
2405
|
+
//#endregion
|
|
2406
|
+
//#region src/gpu/architectures/qwen3_5_vision.d.ts
|
|
2407
|
+
/**
|
|
2408
|
+
* Build the ViT graph. The graph is shaped by symbolic "N" (number of patches),
|
|
2409
|
+
* resolved at run time from the input tensor's first dim — exactly like the LM's
|
|
2410
|
+
* symbolic "T".
|
|
2411
|
+
*/
|
|
2412
|
+
declare function generateQwen3_5VisionGraph(rawConfig: Record<string, unknown>): ModelGraph;
|
|
1850
2413
|
//#endregion
|
|
1851
2414
|
//#region src/gpu/moonshine-executor.d.ts
|
|
1852
2415
|
interface EncoderResult {
|
|
@@ -2252,6 +2815,16 @@ interface GenerateResult {
|
|
|
2252
2815
|
totalTime: number;
|
|
2253
2816
|
/** Why generation stopped. */
|
|
2254
2817
|
finishReason: "eos" | "max_tokens" | "stop_sequence";
|
|
2818
|
+
/**
|
|
2819
|
+
* Decode-only wall time in ms (excludes prefill). Set by the batched path
|
|
2820
|
+
* (generateBatch) so aggregate decode throughput is measurable directly.
|
|
2821
|
+
*/
|
|
2822
|
+
decodeTime?: number;
|
|
2823
|
+
/**
|
|
2824
|
+
* Raw generated token ids, in order. Used by parity harnesses (e.g. the
|
|
2825
|
+
* paged-KV token-exactness gate) to compare runs without decode ambiguity.
|
|
2826
|
+
*/
|
|
2827
|
+
tokenIds?: number[];
|
|
2255
2828
|
/** Thinking content if model produced it (future). */
|
|
2256
2829
|
thinking?: string;
|
|
2257
2830
|
}
|
|
@@ -2364,8 +2937,12 @@ declare class WebGPUEngine {
|
|
|
2364
2937
|
private _outeVoice;
|
|
2365
2938
|
/** Source of the runtime LoRA adapter currently applied on the static base, if any. */
|
|
2366
2939
|
private _currentAdapter;
|
|
2940
|
+
/** Named adapters registered for per-lane batched decode (GERBIL_BATCH_LORA). */
|
|
2941
|
+
private _batchAdapters;
|
|
2367
2942
|
/** Lazily-created Parler-TTS engine (Flan-T5 encoder + decoder LM + dac_44khz). */
|
|
2368
2943
|
private _parlerTTS;
|
|
2944
|
+
/** Lazily-created continuous-batching scheduler (GERBIL_BATCH, Phase 5). */
|
|
2945
|
+
private _scheduler;
|
|
2369
2946
|
/**
|
|
2370
2947
|
* WebKit group-size probe state. When true, a candidate group size is being
|
|
2371
2948
|
* tried this page-load and must be promoted (or capped) after the FIRST
|
|
@@ -2423,6 +3000,17 @@ declare class WebGPUEngine {
|
|
|
2423
3000
|
loadAdapter(source: string | null): Promise<void>;
|
|
2424
3001
|
/** The runtime LoRA adapter currently applied on the base, or null. */
|
|
2425
3002
|
getAdapter(): string | null;
|
|
3003
|
+
/**
|
|
3004
|
+
* Register a named LoRA adapter for per-lane batched decode (SGMV phase 1,
|
|
3005
|
+
* `GERBIL_BATCH=N` + `GERBIL_BATCH_LORA=1`). The adapter's factors are
|
|
3006
|
+
* fetched and packed into the executor's shared factor buffers ONCE; any
|
|
3007
|
+
* number of batch lanes can then run it concurrently by passing its name in
|
|
3008
|
+
* `generateBatch`'s per-lane `adapters` option. Registering the same name
|
|
3009
|
+
* twice is a no-op.
|
|
3010
|
+
*/
|
|
3011
|
+
registerAdapter(name: string, source: string): Promise<void>;
|
|
3012
|
+
/** Names of the adapters registered for per-lane batched decode. */
|
|
3013
|
+
getRegisteredAdapters(): string[];
|
|
2426
3014
|
/**
|
|
2427
3015
|
* Write a coarse crash-phase breadcrumb that survives a GPU-process kill / page
|
|
2428
3016
|
* reload. The iPad harness reads `localStorage["gerbil-crash-phase"]` after a
|
|
@@ -2501,6 +3089,66 @@ declare class WebGPUEngine {
|
|
|
2501
3089
|
* the same instance awaits this one (see {@link _genQueueTail}).
|
|
2502
3090
|
*/
|
|
2503
3091
|
generate(prompt: string | ChatMessage[], options?: GenerateOptions): Promise<GenerateResult>;
|
|
3092
|
+
/**
|
|
3093
|
+
* Batched lockstep generation (GERBIL_BATCH=N — Phase 2 of the
|
|
3094
|
+
* continuous-batching campaign, Dawn only).
|
|
3095
|
+
*
|
|
3096
|
+
* Prefills each prompt sequentially into its batch slot (KV into the slot's
|
|
3097
|
+
* block-table row, SSM/conv state adopted into the slot's state-pool rows),
|
|
3098
|
+
* then decodes all N sequences in lockstep: one batched dispatch stream per
|
|
3099
|
+
* step produces one token per row. Greedy-only in Phase 2 (per-row GPU
|
|
3100
|
+
* argmax); rows that hit EOS/stop keep decoding in lockstep with their
|
|
3101
|
+
* output discarded until every row has finished — rows never interact, so
|
|
3102
|
+
* per-row streams are token-exact vs single-sequence generation.
|
|
3103
|
+
*
|
|
3104
|
+
* Requires exactly `executor.batchSize` prompts (fixed lockstep batch; the
|
|
3105
|
+
* Phase 3 scheduler will lift this).
|
|
3106
|
+
*/
|
|
3107
|
+
generateBatch(prompts: Array<string | ChatMessage[]>, options?: GenerateOptions & {
|
|
3108
|
+
maxTokensPerRow?: number[];
|
|
3109
|
+
/**
|
|
3110
|
+
* Per-lane adapter names (SGMV phase 1 — requires GERBIL_BATCH_LORA=1
|
|
3111
|
+
* and prior {@link registerAdapter} calls). `adapters[i]` selects the
|
|
3112
|
+
* adapter lane i decodes with; null/undefined = bare base. Lanes with
|
|
3113
|
+
* different adapters (or none) decode concurrently over the shared base.
|
|
3114
|
+
*/
|
|
3115
|
+
adapters?: Array<string | null | undefined>;
|
|
3116
|
+
}): Promise<GenerateResult[]>;
|
|
3117
|
+
/**
|
|
3118
|
+
* Continuous-batching scheduler (GERBIL_BATCH=N — Phase 5 of the batching
|
|
3119
|
+
* campaign, Dawn only). Returns the engine's request scheduler (created on
|
|
3120
|
+
* first call): submit many requests with independent prompts/limits and the
|
|
3121
|
+
* scheduler multiplexes them over the batch lanes with dynamic admission,
|
|
3122
|
+
* immediate lane retirement, and chunked prefill. Greedy-only; per-request
|
|
3123
|
+
* output is token-exact vs single-sequence generate().
|
|
3124
|
+
*/
|
|
3125
|
+
createScheduler(): RequestScheduler;
|
|
3126
|
+
/**
|
|
3127
|
+
* Create a swarm — the MoTA mode-B primitive (docs/research/slm-swarms.md):
|
|
3128
|
+
* named members (LoRA adapters on this engine's shared base, or null for the
|
|
3129
|
+
* bare base), an optional plan that fans an input out into per-member
|
|
3130
|
+
* sub-tasks, and an optional reduce that joins the results. `swarm.run()`
|
|
3131
|
+
* decodes every sub-task CONCURRENTLY through the continuous-batching
|
|
3132
|
+
* scheduler with per-request adapters — one batched pass, so the fan-out
|
|
3133
|
+
* costs roughly the slowest member, not the sum.
|
|
3134
|
+
*
|
|
3135
|
+
* Every non-null member source is registered via {@link registerAdapter}
|
|
3136
|
+
* (downloads in parallel; factors packed into shared GPU buffers once).
|
|
3137
|
+
* Requires `GERBIL_BATCH=N` (N >= 2) + `GERBIL_BATCH_LORA=1` on the Dawn
|
|
3138
|
+
* path; greedy-only like the scheduler underneath.
|
|
3139
|
+
*/
|
|
3140
|
+
createSwarm<R = Record<string, ServeResult>>(options: SwarmOptions<R>): Promise<Swarm<R>>;
|
|
3141
|
+
/**
|
|
3142
|
+
* Resolve the end-of-turn stop token id, or null if the model has none.
|
|
3143
|
+
*
|
|
3144
|
+
* Chat models like Gemma 4 end an assistant turn with a dedicated end-of-turn
|
|
3145
|
+
* token (`<turn|>`, id 106) rather than the generic `<eos>`. Every text-decode
|
|
3146
|
+
* loop must treat it as a stop token or the model runs on past its answer and
|
|
3147
|
+
* repeats until it hits `maxTokens`. Resolved from the tokenizer's added tokens
|
|
3148
|
+
* so it is null for models without a turn delimiter (Qwen, LFM2, …), leaving
|
|
3149
|
+
* their decode behaviour unchanged.
|
|
3150
|
+
*/
|
|
3151
|
+
private resolveEndOfTurnId;
|
|
2504
3152
|
private _generateLocked;
|
|
2505
3153
|
/**
|
|
2506
3154
|
* Inline autocomplete: continue `prefix` with a brief, single-line continuation.
|
|
@@ -2834,5 +3482,5 @@ declare class WebGPUEngine {
|
|
|
2834
3482
|
private _acquireGenLock;
|
|
2835
3483
|
}
|
|
2836
3484
|
//#endregion
|
|
2837
|
-
export {
|
|
2838
|
-
//# sourceMappingURL=index-
|
|
3485
|
+
export { Gemma4VisionGridConfig as $, GraphDType as $t, generateOuteTtsBackboneGraph as A, Executor as At, parseKaniConfig as B, AdapterSource as Bt, TranscribeOptions as C, SwarmOptions as Ct, audioTokensToDacCodes as D, SchedulerStats as Dt, generateQwen3_5VisionGraph as E, RequestScheduler as Et, moonshineEncoderFrames as F, loadMoonshine as Ft, patchGemma4VisionClips as G, splitFusedQGateDelta as Gt, dequantizeGemma4VisionProjection as H, applyLoRAToStore as Ht, parseMoonshineConfig as I, quantizeKaniBackbone as It, OUTETTS_ASSETS as J, KaniTTSOptions as Jt, resolveGemma4VisionInfo as K, KaniDType as Kt, audioTokensToCodes as L, ChatMessage as Lt, MOONSHINE_REMAINING_WORK as M, LoadedMoonshine as Mt, generateMoonshineDecoderGraph as N, loadKaniTTS as Nt, dacOutputLength as O, ServeOptions as Ot, generateMoonshineEncoderGraph as P, loadModel as Pt, GEMMA4_IMAGE_PROCESSOR as Q, initGPU as Qt, generateKaniTtsGraph as R, Tokenizer as Rt, MoonshineSTTOptions as S, SwarmMemberResult as St, MoonshineEncoderExecutor as T, SwarmTaskOptions as Tt, dequantizeMLXProjection as U, buildLoRADeltas as Ut, Gemma4VisionGraphInfo as V, LoRADelta as Vt, generateGemma4VisionGraph as W, fetchAdapter as Wt, isTTSRepo as X, SpeakResult as Xt, OUTETTS_PRESET_VOICES as Y, SpeakOptions as Yt, resolveDefaultRepo as Z, GPUDiagnosticResult as Zt, ParlerSpeakOptions as _, preprocessImage as _t, GenerateObjectOptions as a, VisionPositionTensors as at, ParlerTTSOptions as b, SamplingParams as bt, GenerateResult as c, buildGemma4RotaryCosSin as ct, ObjectSchema as d, buildMRoPEPositionIds as dt, KVDType as en, Gemma4VisionPositionTensors as et, ObjectValidator as f, buildPosEmbeds as ft, VisionInputs as g, mropeFreqDims as gt, VisionExecutor as h, buildVisionPositionTensors as ht, EncodeImageResult as i, createDefaultHFKeyMapper as in, VisionGridConfig as it, parseOuteTtsConfig as j, LoadedKaniTTS as jt, generateDacSpeechDecoderGraph as k, ServeResult as kt, IntegrityCheckEntry as l, buildGemma4VisionPositionTensors as lt, WebGPUEngineOptions as m, buildRotaryCosSin as mt, AgentTool as n, ModelArchConfig as nn, PreprocessedImage as nt, GenerateObjectResult as o, buildGemma4PoolMatrix as ot, WebGPUEngine as p, buildPositionIds as pt, DEFAULT_MODELS as q, KaniTTS as qt, EmbedOptions as r, ModelCapabilities as rn, QWEN3_5_IMAGE_PROCESSOR as rt, GenerateOptions as s, buildGemma4PosEmbeds as st, AgentStep as t, KvMode as tn, ImageProcessorConfig as tt, IntegrityCheckResult as u, buildMRoPECosSin as ut, ParlerSpeakResult as v, preprocessImageGemma4 as vt, TranscribeResult as w, SwarmTask as wt, MoonshineSTT as x, Swarm as xt, ParlerTTS as y, smartResize as yt, generateNanoCodecDecoderGraph as z, AdapterConfig as zt };
|
|
3486
|
+
//# sourceMappingURL=index-CNYoTRgr.d.mts.map
|