@tryhamster/gerbil 1.12.0 → 1.13.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 +25 -0
- package/dist/{architectures-BH_z6k9d.mjs → architectures-DHwj9AQD.mjs} +37 -27
- package/dist/architectures-DHwj9AQD.mjs.map +1 -0
- package/dist/cli.mjs +8 -8
- package/dist/cli.mjs.map +1 -1
- package/dist/{defaults-DfGx4d1m.mjs → defaults-B0aQZJTM.mjs} +3 -2
- package/dist/defaults-B0aQZJTM.mjs.map +1 -0
- 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-Bw6do78d.mjs +4 -0
- package/dist/{gerbil-DF009Waa.d.mts → gerbil-CD_skiL3.d.mts} +2 -2
- package/dist/{gerbil-DF009Waa.d.mts.map → gerbil-CD_skiL3.d.mts.map} +1 -1
- package/dist/{gerbil-CpS3P240.mjs → gerbil-DAz_a4Jh.mjs} +22 -13
- package/dist/{gerbil-CpS3P240.mjs.map → gerbil-DAz_a4Jh.mjs.map} +1 -1
- package/dist/gpu/hooks.d.mts +1 -1
- package/dist/gpu/hooks.mjs +1 -1
- package/dist/gpu/index.d.mts +2 -2
- package/dist/gpu/index.mjs +5 -5
- package/dist/{gpu-DRFhiv4R.mjs → gpu-CU_Mldk0.mjs} +1275 -195
- package/dist/gpu-CU_Mldk0.mjs.map +1 -0
- package/dist/index-B3tjyDJI.d.mts.map +1 -1
- package/dist/{index-CNYoTRgr.d.mts → index-Fj2XkP-o.d.mts} +374 -3
- package/dist/index-Fj2XkP-o.d.mts.map +1 -0
- package/dist/index.d.mts +2 -2
- package/dist/index.mjs +6 -6
- 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-DAbWO8VS.mjs → mcp-CdMvLQ_9.mjs} +3 -3
- package/dist/{mcp-DAbWO8VS.mjs.map → mcp-CdMvLQ_9.mjs.map} +1 -1
- package/dist/{moonshine-stt-B1kV5n1c.mjs → moonshine-stt-COJeK2Zb.mjs} +778 -3
- package/dist/moonshine-stt-COJeK2Zb.mjs.map +1 -0
- package/dist/moonshine-stt-Dp4j5JEB.mjs +4 -0
- package/dist/{one-liner-CmP9ktUn.mjs → one-liner-CMiGNWJ7.mjs} +2 -2
- package/dist/{one-liner-CmP9ktUn.mjs.map → one-liner-CMiGNWJ7.mjs.map} +1 -1
- package/dist/repl-BhaLCJFb.mjs +9 -0
- package/dist/skills/index.d.mts +2 -2
- package/dist/skills/index.mjs +3 -3
- package/dist/{skills-BDOEbHSx.mjs → skills-BXybFWlG.mjs} +2 -2
- package/dist/{skills-BDOEbHSx.mjs.map → skills-BXybFWlG.mjs.map} +1 -1
- package/dist/tune/index.mjs +1 -1
- package/package.json +1 -1
- package/dist/architectures-BH_z6k9d.mjs.map +0 -1
- package/dist/defaults-DfGx4d1m.mjs.map +0 -1
- package/dist/gerbil-CZvoFo0T.mjs +0 -4
- package/dist/gpu-DRFhiv4R.mjs.map +0 -1
- package/dist/index-CNYoTRgr.d.mts.map +0 -1
- package/dist/moonshine-stt-B1kV5n1c.mjs.map +0 -1
- package/dist/moonshine-stt-DA1WuiZb.mjs +0 -4
- package/dist/repl-17NVeIaJ.mjs +0 -9
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
import { F as CANONICAL_KEYS, I as DTYPE_BYTES, N as DEFAULT_GROUP_SIZE, P as quantizeInt4, R as createDefaultHFKeyMapper, _ as moonshineEncoderFrames, g as generateMoonshineEncoderGraph, h as generateMoonshineDecoderGraph, j as nanoCodecWeightMap, t as generateGraph, v as parseMoonshineConfig, w as foldNanoCodecWeightNorm } from "./architectures-
|
|
1
|
+
import { F as CANONICAL_KEYS, I as DTYPE_BYTES, N as DEFAULT_GROUP_SIZE, P as quantizeInt4, R as createDefaultHFKeyMapper, _ as moonshineEncoderFrames, g as generateMoonshineEncoderGraph, h as generateMoonshineDecoderGraph, j as nanoCodecWeightMap, t as generateGraph, v as parseMoonshineConfig, w as foldNanoCodecWeightNorm } from "./architectures-DHwj9AQD.mjs";
|
|
2
2
|
|
|
3
3
|
//#region src/gpu/device.ts
|
|
4
4
|
const COPY_SRC = 4;
|
|
@@ -2096,6 +2096,103 @@ fn main(
|
|
|
2096
2096
|
}
|
|
2097
2097
|
}
|
|
2098
2098
|
`;
|
|
2099
|
+
const WGSL_QUAD_MATVEC_INT4 = `\
|
|
2100
|
+
const N_TILE: u32 = 16u;
|
|
2101
|
+
const K_THREADS: u32 = 16u;
|
|
2102
|
+
|
|
2103
|
+
struct Params {
|
|
2104
|
+
K: u32,
|
|
2105
|
+
N: u32,
|
|
2106
|
+
group_size: u32,
|
|
2107
|
+
n0: u32,
|
|
2108
|
+
n1: u32,
|
|
2109
|
+
n2: u32,
|
|
2110
|
+
_pad0: u32,
|
|
2111
|
+
_pad1: u32,
|
|
2112
|
+
}
|
|
2113
|
+
|
|
2114
|
+
@group(0) @binding(0) var<storage, read> A: array<vec4f>;
|
|
2115
|
+
@group(0) @binding(1) var<storage, read> B_q: array<vec4u>;
|
|
2116
|
+
@group(0) @binding(2) var<storage, read> scales: array<f32>;
|
|
2117
|
+
@group(0) @binding(3) var<storage, read> zeros: array<f32>;
|
|
2118
|
+
@group(0) @binding(4) var<storage, read_write> out0: array<f32>;
|
|
2119
|
+
@group(0) @binding(5) var<storage, read_write> out1: array<f32>;
|
|
2120
|
+
@group(0) @binding(6) var<storage, read_write> out2: array<f32>;
|
|
2121
|
+
@group(0) @binding(7) var<storage, read_write> out3: array<f32>;
|
|
2122
|
+
@group(0) @binding(8) var<storage, read> params: Params;
|
|
2123
|
+
|
|
2124
|
+
var<workgroup> shared_sums: array<f32, 256>;
|
|
2125
|
+
|
|
2126
|
+
@compute @workgroup_size(256)
|
|
2127
|
+
fn main(
|
|
2128
|
+
@builtin(local_invocation_id) lid: vec3u,
|
|
2129
|
+
@builtin(workgroup_id) wid: vec3u,
|
|
2130
|
+
) {
|
|
2131
|
+
let tid = lid.x;
|
|
2132
|
+
let n_idx = tid / K_THREADS;
|
|
2133
|
+
let k_tid = tid % K_THREADS;
|
|
2134
|
+
let col = wid.x * N_TILE + n_idx;
|
|
2135
|
+
let K = params.K;
|
|
2136
|
+
let k_vec = K / 32u;
|
|
2137
|
+
let groups_per_col = K / params.group_size;
|
|
2138
|
+
let packed_per_group = params.group_size / 8u;
|
|
2139
|
+
|
|
2140
|
+
var sum: f32 = 0.0;
|
|
2141
|
+
if (col < params.N) {
|
|
2142
|
+
let b_off_v = col * k_vec;
|
|
2143
|
+
let col_g_base = col * groups_per_col;
|
|
2144
|
+
|
|
2145
|
+
var v = k_tid;
|
|
2146
|
+
while (v < k_vec) {
|
|
2147
|
+
let bq = B_q[b_off_v + v];
|
|
2148
|
+
let g = (v * 4u) / packed_per_group;
|
|
2149
|
+
let scale = scales[col_g_base + g];
|
|
2150
|
+
let zero = zeros[col_g_base + g];
|
|
2151
|
+
let a_base = v * 8u;
|
|
2152
|
+
|
|
2153
|
+
for (var j: u32 = 0u; j < 4u; j++) {
|
|
2154
|
+
let packed = bq[j];
|
|
2155
|
+
let n0v = f32(packed & 0xFu);
|
|
2156
|
+
let n1v = f32((packed >> 4u) & 0xFu);
|
|
2157
|
+
let n2v = f32((packed >> 8u) & 0xFu);
|
|
2158
|
+
let n3v = f32((packed >> 12u) & 0xFu);
|
|
2159
|
+
let n4v = f32((packed >> 16u) & 0xFu);
|
|
2160
|
+
let n5v = f32((packed >> 20u) & 0xFu);
|
|
2161
|
+
let n6v = f32((packed >> 24u) & 0xFu);
|
|
2162
|
+
let n7v = f32((packed >> 28u) & 0xFu);
|
|
2163
|
+
|
|
2164
|
+
let b0 = vec4f(n0v - zero, n1v - zero, n2v - zero, n3v - zero) * scale;
|
|
2165
|
+
let b1 = vec4f(n4v - zero, n5v - zero, n6v - zero, n7v - zero) * scale;
|
|
2166
|
+
|
|
2167
|
+
sum += dot(A[a_base + j * 2u], b0) + dot(A[a_base + j * 2u + 1u], b1);
|
|
2168
|
+
}
|
|
2169
|
+
|
|
2170
|
+
v += K_THREADS;
|
|
2171
|
+
}
|
|
2172
|
+
}
|
|
2173
|
+
|
|
2174
|
+
shared_sums[tid] = sum;
|
|
2175
|
+
workgroupBarrier();
|
|
2176
|
+
|
|
2177
|
+
if (k_tid == 0u && col < params.N) {
|
|
2178
|
+
let base = n_idx * K_THREADS;
|
|
2179
|
+
var total: f32 = shared_sums[base];
|
|
2180
|
+
for (var i: u32 = 1u; i < K_THREADS; i++) {
|
|
2181
|
+
total += shared_sums[base + i];
|
|
2182
|
+
}
|
|
2183
|
+
// Route the column to its original destination buffer.
|
|
2184
|
+
if (col < params.n0) {
|
|
2185
|
+
out0[col] = total;
|
|
2186
|
+
} else if (col < params.n1) {
|
|
2187
|
+
out1[col - params.n0] = total;
|
|
2188
|
+
} else if (col < params.n2) {
|
|
2189
|
+
out2[col - params.n1] = total;
|
|
2190
|
+
} else {
|
|
2191
|
+
out3[col - params.n2] = total;
|
|
2192
|
+
}
|
|
2193
|
+
}
|
|
2194
|
+
}
|
|
2195
|
+
`;
|
|
2099
2196
|
const WGSL_GATED_MATVEC_INT4 = `\
|
|
2100
2197
|
const N_TILE: u32 = 8u;
|
|
2101
2198
|
const K_THREADS: u32 = 16u;
|
|
@@ -4331,6 +4428,133 @@ fn main(
|
|
|
4331
4428
|
}
|
|
4332
4429
|
}
|
|
4333
4430
|
`;
|
|
4431
|
+
const WGSL_NORM_MATVEC_INT4 = `\
|
|
4432
|
+
const N_TILE: u32 = 16u;
|
|
4433
|
+
const K_THREADS: u32 = 16u;
|
|
4434
|
+
|
|
4435
|
+
struct Params {
|
|
4436
|
+
K: u32,
|
|
4437
|
+
N: u32,
|
|
4438
|
+
group_size: u32,
|
|
4439
|
+
eps_bits: u32,
|
|
4440
|
+
}
|
|
4441
|
+
|
|
4442
|
+
@group(0) @binding(0) var<storage, read> A_a: array<vec4f>;
|
|
4443
|
+
@group(0) @binding(1) var<storage, read> A_b: array<vec4f>;
|
|
4444
|
+
@group(0) @binding(2) var<storage, read> W_norm: array<vec4f>;
|
|
4445
|
+
@group(0) @binding(3) var<storage, read> B_q: array<vec4u>;
|
|
4446
|
+
@group(0) @binding(4) var<storage, read> scales: array<f32>;
|
|
4447
|
+
@group(0) @binding(5) var<storage, read> zeros: array<f32>;
|
|
4448
|
+
@group(0) @binding(6) var<storage, read_write> C: array<f32>;
|
|
4449
|
+
@group(0) @binding(7) var<storage, read_write> residual_out: array<vec4f>;
|
|
4450
|
+
@group(0) @binding(8) var<storage, read_write> norm_out: array<vec4f>;
|
|
4451
|
+
@group(0) @binding(9) var<storage, read> params: Params;
|
|
4452
|
+
|
|
4453
|
+
var<workgroup> shared_sq: array<f32, 256>;
|
|
4454
|
+
var<workgroup> shared_sums: array<f32, 256>;
|
|
4455
|
+
|
|
4456
|
+
@compute @workgroup_size(256)
|
|
4457
|
+
fn main(
|
|
4458
|
+
@builtin(local_invocation_id) lid: vec3u,
|
|
4459
|
+
@builtin(workgroup_id) wid: vec3u,
|
|
4460
|
+
) {
|
|
4461
|
+
let tid = lid.x;
|
|
4462
|
+
let K = params.K;
|
|
4463
|
+
let k_vec4 = K / 4u;
|
|
4464
|
+
|
|
4465
|
+
// ── Prologue: residual sum + RMS (all threads, uniform control flow) ──
|
|
4466
|
+
var local_sq: f32 = 0.0;
|
|
4467
|
+
var v4 = tid;
|
|
4468
|
+
while (v4 < k_vec4) {
|
|
4469
|
+
let x = A_a[v4] + A_b[v4];
|
|
4470
|
+
if (wid.x == 0u) { residual_out[v4] = x; }
|
|
4471
|
+
local_sq += dot(x, x);
|
|
4472
|
+
v4 += 256u;
|
|
4473
|
+
}
|
|
4474
|
+
shared_sq[tid] = local_sq;
|
|
4475
|
+
workgroupBarrier();
|
|
4476
|
+
var stride: u32 = 128u;
|
|
4477
|
+
while (stride > 0u) {
|
|
4478
|
+
if (tid < stride) {
|
|
4479
|
+
shared_sq[tid] += shared_sq[tid + stride];
|
|
4480
|
+
}
|
|
4481
|
+
workgroupBarrier();
|
|
4482
|
+
stride = stride / 2u;
|
|
4483
|
+
}
|
|
4484
|
+
let eps = bitcast<f32>(params.eps_bits);
|
|
4485
|
+
let inv_rms = 1.0 / sqrt(shared_sq[0] / f32(K) + eps);
|
|
4486
|
+
workgroupBarrier();
|
|
4487
|
+
|
|
4488
|
+
// Workgroup 0 materializes the full normed vector — it stays live for the
|
|
4489
|
+
// sibling a_proj/b_proj/z_proj projections that read the same layernorm output.
|
|
4490
|
+
if (wid.x == 0u) {
|
|
4491
|
+
var w4 = tid;
|
|
4492
|
+
while (w4 < k_vec4) {
|
|
4493
|
+
norm_out[w4] = (A_a[w4] + A_b[w4]) * inv_rms * W_norm[w4];
|
|
4494
|
+
w4 += 256u;
|
|
4495
|
+
}
|
|
4496
|
+
}
|
|
4497
|
+
|
|
4498
|
+
// ── Main GEMV: C[col] = normed · dequant(B[col,:]) ──
|
|
4499
|
+
let n_idx = tid / K_THREADS;
|
|
4500
|
+
let k_tid = tid % K_THREADS;
|
|
4501
|
+
let col = wid.x * N_TILE + n_idx;
|
|
4502
|
+
let k_vec = K / 32u;
|
|
4503
|
+
let groups_per_col = K / params.group_size;
|
|
4504
|
+
let packed_per_group = params.group_size / 8u;
|
|
4505
|
+
|
|
4506
|
+
var sum: f32 = 0.0;
|
|
4507
|
+
if (col < params.N) {
|
|
4508
|
+
let b_off_v = col * k_vec;
|
|
4509
|
+
let col_g_base = col * groups_per_col;
|
|
4510
|
+
|
|
4511
|
+
var v = k_tid;
|
|
4512
|
+
while (v < k_vec) {
|
|
4513
|
+
let bq = B_q[b_off_v + v];
|
|
4514
|
+
let g = (v * 4u) / packed_per_group;
|
|
4515
|
+
let scale = scales[col_g_base + g];
|
|
4516
|
+
let zero = zeros[col_g_base + g];
|
|
4517
|
+
let a_base = v * 8u;
|
|
4518
|
+
|
|
4519
|
+
for (var j: u32 = 0u; j < 4u; j++) {
|
|
4520
|
+
let packed = bq[j];
|
|
4521
|
+
let n0 = f32(packed & 0xFu);
|
|
4522
|
+
let n1 = f32((packed >> 4u) & 0xFu);
|
|
4523
|
+
let n2 = f32((packed >> 8u) & 0xFu);
|
|
4524
|
+
let n3 = f32((packed >> 12u) & 0xFu);
|
|
4525
|
+
let n4 = f32((packed >> 16u) & 0xFu);
|
|
4526
|
+
let n5 = f32((packed >> 20u) & 0xFu);
|
|
4527
|
+
let n6 = f32((packed >> 24u) & 0xFu);
|
|
4528
|
+
let n7 = f32((packed >> 28u) & 0xFu);
|
|
4529
|
+
|
|
4530
|
+
let b0 = vec4f(n0 - zero, n1 - zero, n2 - zero, n3 - zero) * scale;
|
|
4531
|
+
let b1 = vec4f(n4 - zero, n5 - zero, n6 - zero, n7 - zero) * scale;
|
|
4532
|
+
|
|
4533
|
+
let i0 = a_base + j * 2u;
|
|
4534
|
+
let i1 = a_base + j * 2u + 1u;
|
|
4535
|
+
let a0 = (A_a[i0] + A_b[i0]) * inv_rms * W_norm[i0];
|
|
4536
|
+
let a1 = (A_a[i1] + A_b[i1]) * inv_rms * W_norm[i1];
|
|
4537
|
+
|
|
4538
|
+
sum += dot(a0, b0) + dot(a1, b1);
|
|
4539
|
+
}
|
|
4540
|
+
|
|
4541
|
+
v += K_THREADS;
|
|
4542
|
+
}
|
|
4543
|
+
}
|
|
4544
|
+
|
|
4545
|
+
shared_sums[tid] = sum;
|
|
4546
|
+
workgroupBarrier();
|
|
4547
|
+
|
|
4548
|
+
if (k_tid == 0u && col < params.N) {
|
|
4549
|
+
let base = n_idx * K_THREADS;
|
|
4550
|
+
var total: f32 = shared_sums[base];
|
|
4551
|
+
for (var i: u32 = 1u; i < K_THREADS; i++) {
|
|
4552
|
+
total += shared_sums[base + i];
|
|
4553
|
+
}
|
|
4554
|
+
C[col] = total;
|
|
4555
|
+
}
|
|
4556
|
+
}
|
|
4557
|
+
`;
|
|
4334
4558
|
const WGSL_DUAL_MATVEC_INT4 = `\
|
|
4335
4559
|
const N_TILE: u32 = 16u;
|
|
4336
4560
|
const K_THREADS: u32 = 16u;
|
|
@@ -6870,6 +7094,47 @@ const MATVEC_INT4_SPEC = {
|
|
|
6870
7094
|
]);
|
|
6871
7095
|
}
|
|
6872
7096
|
};
|
|
7097
|
+
const QUAD_MATVEC_INT4_SPEC = {
|
|
7098
|
+
shaderCode: WGSL_QUAD_MATVEC_INT4,
|
|
7099
|
+
entryPoint: "main",
|
|
7100
|
+
bindings: [
|
|
7101
|
+
{ type: "storage-read" },
|
|
7102
|
+
{ type: "storage-read" },
|
|
7103
|
+
{ type: "storage-read" },
|
|
7104
|
+
{ type: "storage-read" },
|
|
7105
|
+
{ type: "storage-read-write" },
|
|
7106
|
+
{ type: "storage-read-write" },
|
|
7107
|
+
{ type: "storage-read-write" },
|
|
7108
|
+
{ type: "storage-read-write" },
|
|
7109
|
+
{ type: "uniform" }
|
|
7110
|
+
],
|
|
7111
|
+
getDispatchSize(op) {
|
|
7112
|
+
const N = op.attributes.N;
|
|
7113
|
+
return [
|
|
7114
|
+
cdiv$1(N, 16),
|
|
7115
|
+
1,
|
|
7116
|
+
1
|
|
7117
|
+
];
|
|
7118
|
+
},
|
|
7119
|
+
buildParams(op) {
|
|
7120
|
+
const K = op.attributes.K;
|
|
7121
|
+
const N = op.attributes.N;
|
|
7122
|
+
const group_size = op.attributes.group_size;
|
|
7123
|
+
const n0 = op.attributes.n0;
|
|
7124
|
+
const n1 = op.attributes.n1;
|
|
7125
|
+
const n2 = op.attributes.n2;
|
|
7126
|
+
return buildUniformBuffer$1([
|
|
7127
|
+
K,
|
|
7128
|
+
N,
|
|
7129
|
+
group_size,
|
|
7130
|
+
n0,
|
|
7131
|
+
n1,
|
|
7132
|
+
n2,
|
|
7133
|
+
0,
|
|
7134
|
+
0
|
|
7135
|
+
]);
|
|
7136
|
+
}
|
|
7137
|
+
};
|
|
6873
7138
|
function deriveMatvecGeometry(nTile, kThreads) {
|
|
6874
7139
|
return {
|
|
6875
7140
|
...MATVEC_INT4_SPEC,
|
|
@@ -7084,6 +7349,41 @@ const NORM_SWIGLU_MATVEC_INT4_SPEC = {
|
|
|
7084
7349
|
]);
|
|
7085
7350
|
}
|
|
7086
7351
|
};
|
|
7352
|
+
const NORM_MATVEC_INT4_SPEC = {
|
|
7353
|
+
shaderCode: WGSL_NORM_MATVEC_INT4,
|
|
7354
|
+
entryPoint: "main",
|
|
7355
|
+
bindings: [
|
|
7356
|
+
{ type: "storage-read" },
|
|
7357
|
+
{ type: "storage-read" },
|
|
7358
|
+
{ type: "storage-read" },
|
|
7359
|
+
{ type: "storage-read" },
|
|
7360
|
+
{ type: "storage-read" },
|
|
7361
|
+
{ type: "storage-read" },
|
|
7362
|
+
{ type: "storage-read-write" },
|
|
7363
|
+
{ type: "storage-read-write" },
|
|
7364
|
+
{ type: "storage-read-write" },
|
|
7365
|
+
{ type: "uniform" }
|
|
7366
|
+
],
|
|
7367
|
+
getDispatchSize(op) {
|
|
7368
|
+
const N = op.attributes.N;
|
|
7369
|
+
return [
|
|
7370
|
+
cdiv$1(N, 16),
|
|
7371
|
+
1,
|
|
7372
|
+
1
|
|
7373
|
+
];
|
|
7374
|
+
},
|
|
7375
|
+
buildParams(op) {
|
|
7376
|
+
const K = op.attributes.K;
|
|
7377
|
+
const N = op.attributes.N;
|
|
7378
|
+
const group_size = op.attributes.group_size;
|
|
7379
|
+
return buildUniformBuffer$1([
|
|
7380
|
+
K,
|
|
7381
|
+
N,
|
|
7382
|
+
group_size,
|
|
7383
|
+
f32BitsToU32$1(op.attributes.norm_eps ?? 1e-6)
|
|
7384
|
+
]);
|
|
7385
|
+
}
|
|
7386
|
+
};
|
|
7087
7387
|
const DUAL_MATVEC_INT4_SPEC = {
|
|
7088
7388
|
shaderCode: WGSL_DUAL_MATVEC_INT4,
|
|
7089
7389
|
entryPoint: "main",
|
|
@@ -11246,6 +11546,11 @@ const BATCH_LORA_NONE = 4294967295;
|
|
|
11246
11546
|
var Executor = class Executor {
|
|
11247
11547
|
ctx;
|
|
11248
11548
|
graph;
|
|
11549
|
+
/** Output tensor names of this executor's graph (read-only; used by non-LM
|
|
11550
|
+
* codec executors that need to locate the decoded-audio output tensor). */
|
|
11551
|
+
get graphOutputs() {
|
|
11552
|
+
return this.graph.outputs;
|
|
11553
|
+
}
|
|
11249
11554
|
weightBuffers = /* @__PURE__ */ new Map();
|
|
11250
11555
|
activationBuffers = /* @__PURE__ */ new Map();
|
|
11251
11556
|
ssmStateBuffers = /* @__PURE__ */ new Map();
|
|
@@ -11310,6 +11615,13 @@ var Executor = class Executor {
|
|
|
11310
11615
|
dispatchEntries = [];
|
|
11311
11616
|
/** Dispatch entries for decode (M=1): uses K-parallel matvec kernels. */
|
|
11312
11617
|
decodeEntries = [];
|
|
11618
|
+
/**
|
|
11619
|
+
* Column-concatenated INT4 weight buffers built by the Mamba-front fusion
|
|
11620
|
+
* (qkv|a|b|z packed into one q/scales/zeros set per layer). Retained here so
|
|
11621
|
+
* they outlive the fusion pass and are not garbage-collected while the fused
|
|
11622
|
+
* bind groups still reference them.
|
|
11623
|
+
*/
|
|
11624
|
+
mambaFrontConcatBuffers = [];
|
|
11313
11625
|
/** Per-node decode entries captured BEFORE fusion (the overlay base). */
|
|
11314
11626
|
decodeEntriesUnfused = [];
|
|
11315
11627
|
/** True while a runtime adapter is applied. */
|
|
@@ -11594,6 +11906,7 @@ var Executor = class Executor {
|
|
|
11594
11906
|
this.decodeEntriesUnfused = this.decodeEntries.slice();
|
|
11595
11907
|
this.fuseNormSwiGLUDecodeEntries();
|
|
11596
11908
|
this.fuseSwiGLUDecodeEntries();
|
|
11909
|
+
this.fuseMambaFrontDecodeEntries();
|
|
11597
11910
|
this.fuseDualMatVecDecodeEntries();
|
|
11598
11911
|
this.fuseDualKVCacheAppendEntries();
|
|
11599
11912
|
this.fuseGatedOProjDecodeEntries();
|
|
@@ -11602,6 +11915,7 @@ var Executor = class Executor {
|
|
|
11602
11915
|
this.fuseDualRMSNormDecodeEntries();
|
|
11603
11916
|
this.fuseConvStateUpdateDecodeEntries();
|
|
11604
11917
|
this.fuseAddRMSNormDecodeEntries();
|
|
11918
|
+
this.fuseNormMatVecDecodeEntries();
|
|
11605
11919
|
console.log(`[executor] decode: ${this.decodeEntries.length} dispatches/token`);
|
|
11606
11920
|
if (this.batchSize > 0) this.buildBatchEntries();
|
|
11607
11921
|
const logitsBuffer = this.getBuffer("logits");
|
|
@@ -11631,6 +11945,7 @@ var Executor = class Executor {
|
|
|
11631
11945
|
size: this.uniformStagingCapacity,
|
|
11632
11946
|
usage: 12
|
|
11633
11947
|
});
|
|
11948
|
+
if (globalThis.process?.env?.GERBIL_DUMP_GRAPH === "1") globalThis.__GERBIL_DECODE_GRAPH__ = this.dumpDecodeGraph();
|
|
11634
11949
|
}
|
|
11635
11950
|
/** True while a runtime LoRA overlay is applied. */
|
|
11636
11951
|
get hasRuntimeAdapter() {
|
|
@@ -12041,6 +12356,97 @@ var Executor = class Executor {
|
|
|
12041
12356
|
return this.ctx.limits.maxStorageBuffersPerShaderStage;
|
|
12042
12357
|
}
|
|
12043
12358
|
/**
|
|
12359
|
+
* Build a machine-readable dump of the post-fusion decode entry list (the
|
|
12360
|
+
* exact per-token dispatch sequence), with per-tensor pool-buffer identities
|
|
12361
|
+
* so pooled-aliasing hazards are visible offline. Debug/analysis only — never
|
|
12362
|
+
* called on any hot path; reached via the GERBIL_DUMP_GRAPH hook at the end
|
|
12363
|
+
* of {@link initBindGroups} or directly by tooling.
|
|
12364
|
+
*/
|
|
12365
|
+
dumpDecodeGraph(dumpSeqPos = 64) {
|
|
12366
|
+
const bufferLabels = /* @__PURE__ */ new Map();
|
|
12367
|
+
const bufferKinds = /* @__PURE__ */ new Map();
|
|
12368
|
+
const register = (map, kind) => {
|
|
12369
|
+
for (const [name, buf] of map) if (!bufferLabels.has(buf)) {
|
|
12370
|
+
bufferLabels.set(buf, buf.label || `${kind}_${name}`);
|
|
12371
|
+
bufferKinds.set(buf, kind);
|
|
12372
|
+
}
|
|
12373
|
+
};
|
|
12374
|
+
register(this.weightBuffers, "weight");
|
|
12375
|
+
register(this.activationBuffers, "activation");
|
|
12376
|
+
register(this.ssmStateBuffers, "ssm-state");
|
|
12377
|
+
register(this.kvCacheBuffers, "kv-cache");
|
|
12378
|
+
const resolved = this.resolveShapes(1, dumpSeqPos);
|
|
12379
|
+
const runtimeContext = {
|
|
12380
|
+
seqPos: dumpSeqPos,
|
|
12381
|
+
tableBase: 0
|
|
12382
|
+
};
|
|
12383
|
+
const tensorRef = (tensor) => {
|
|
12384
|
+
const desc = this.graph.tensors[tensor];
|
|
12385
|
+
const buf = this.getBuffer(tensor);
|
|
12386
|
+
return {
|
|
12387
|
+
tensor,
|
|
12388
|
+
shape: desc ? [...desc.shape] : [],
|
|
12389
|
+
resolvedShape: resolved[tensor] ?? null,
|
|
12390
|
+
buffer: buf ? bufferLabels.get(buf) ?? buf.label ?? null : null,
|
|
12391
|
+
bufferBytes: buf ? buf.size : null
|
|
12392
|
+
};
|
|
12393
|
+
};
|
|
12394
|
+
const nodeById = new Map(this.graph.nodes.map((n) => [n.id, n]));
|
|
12395
|
+
const dumpNode = (id) => {
|
|
12396
|
+
const node = nodeById.get(id);
|
|
12397
|
+
if (!node) return null;
|
|
12398
|
+
const attributes = {};
|
|
12399
|
+
for (const [k, v] of Object.entries(node.attributes ?? {})) if (typeof v === "number" || typeof v === "string" || typeof v === "boolean") attributes[k] = v;
|
|
12400
|
+
return {
|
|
12401
|
+
id,
|
|
12402
|
+
opType: node.opType,
|
|
12403
|
+
inputs: node.inputs.map(tensorRef),
|
|
12404
|
+
outputs: node.outputs.map(tensorRef),
|
|
12405
|
+
attributes
|
|
12406
|
+
};
|
|
12407
|
+
};
|
|
12408
|
+
const entries = this.decodeEntries.map((entry, index) => {
|
|
12409
|
+
let dispatchSize = null;
|
|
12410
|
+
try {
|
|
12411
|
+
dispatchSize = entry.spec.getDispatchSize(entry.node, resolved, runtimeContext);
|
|
12412
|
+
} catch {}
|
|
12413
|
+
let paramsU32 = [];
|
|
12414
|
+
let paramsBytes = 0;
|
|
12415
|
+
try {
|
|
12416
|
+
const params = entry.spec.buildParams(entry.node, resolved, runtimeContext);
|
|
12417
|
+
paramsBytes = params.byteLength;
|
|
12418
|
+
paramsU32 = [...new Uint32Array(params.slice(0, Math.floor(params.byteLength / 4) * 4))];
|
|
12419
|
+
} catch {}
|
|
12420
|
+
const bindings = entry.spec.bindings ?? [];
|
|
12421
|
+
return {
|
|
12422
|
+
index,
|
|
12423
|
+
nodeId: entry.nodeId,
|
|
12424
|
+
opType: entry.node.opType,
|
|
12425
|
+
fused: entry.fusedNodeIds !== void 0,
|
|
12426
|
+
fusedNodeIds: entry.fusedNodeIds ? [...entry.fusedNodeIds] : null,
|
|
12427
|
+
storageBindings: bindings.filter((b) => b.type !== "uniform").length,
|
|
12428
|
+
totalBindings: bindings.length,
|
|
12429
|
+
dispatchSize,
|
|
12430
|
+
paramsBytes,
|
|
12431
|
+
paramsU32,
|
|
12432
|
+
bindBufferLabels: entry.buffers ? entry.buffers.map((b) => bufferLabels.get(b.buffer) ?? b.buffer.label ?? "uniform-or-unregistered") : null,
|
|
12433
|
+
nodes: this.coveredNodeIds(entry).map(dumpNode).filter((n) => n !== null)
|
|
12434
|
+
};
|
|
12435
|
+
});
|
|
12436
|
+
return {
|
|
12437
|
+
dispatchesPerToken: this.decodeEntries.length,
|
|
12438
|
+
unfusedDispatchesPerToken: this.decodeEntriesUnfused.length,
|
|
12439
|
+
maxStorageBuffersPerShaderStage: this.ctx.limits.maxStorageBuffersPerShaderStage,
|
|
12440
|
+
dumpSeqPos,
|
|
12441
|
+
buffers: [...bufferLabels.entries()].map(([buf, label]) => ({
|
|
12442
|
+
label,
|
|
12443
|
+
bytes: buf.size,
|
|
12444
|
+
kind: bufferKinds.get(buf) ?? "unknown"
|
|
12445
|
+
})),
|
|
12446
|
+
entries
|
|
12447
|
+
};
|
|
12448
|
+
}
|
|
12449
|
+
/**
|
|
12044
12450
|
* Profile ONE real decode step: times the actual `decodeEntries` (the kernels
|
|
12045
12451
|
* the pipelined greedy benchmark runs) with per-dispatch timestamps. Timing is
|
|
12046
12452
|
* token-independent, so pass any valid id; runs un-pipelined with a synchronous
|
|
@@ -14242,6 +14648,111 @@ var Executor = class Executor {
|
|
|
14242
14648
|
* when the device can't bind 11 storage buffers. Windows it skips fall
|
|
14243
14649
|
* through to fuseSwiGLUDecodeEntries.
|
|
14244
14650
|
*/
|
|
14651
|
+
/**
|
|
14652
|
+
* Fuse [ResidualRMSNorm → Mamba qkv MatVecInt4] → ONE NormMatVec dispatch.
|
|
14653
|
+
* The kernel recomputes the residual sum + RMS in a prologue (workgroup 0
|
|
14654
|
+
* persists BOTH the residual sum and the full normed vector — the latter is
|
|
14655
|
+
* live for the sibling a_proj/b_proj/z_proj), then the GEMV normalizes the
|
|
14656
|
+
* activation on the fly. Saves the ResidualRMSNorm dispatch (−1 per Mamba
|
|
14657
|
+
* layer, ×17). 10 storage bindings — fits the Dawn/Metal 10-binding adapter.
|
|
14658
|
+
* INT4-only.
|
|
14659
|
+
*
|
|
14660
|
+
* OPT-IN, DEFAULT OFF (GERBIL_FUSE_NORM_QKV=1). Round r28 (2026-07-25) measured
|
|
14661
|
+
* this a −4.6% REGRESSION on M4 Max/Metal: the qkv GEMV is the widest projection
|
|
14662
|
+
* (N=6144 → 384 workgroups), so recomputing the 1024-wide residual+RMS reduction
|
|
14663
|
+
* per workgroup + normalizing activations on the fly costs far more ALU than the
|
|
14664
|
+
* single saved ResidualRMSNorm dispatch returns on a bandwidth-plateaued backend.
|
|
14665
|
+
* The −17 dispatch/token reduction still targets Vulkan/WebKit (dispatch-count-
|
|
14666
|
+
* bound) — validate on the kernel-lab A5000 rig before enabling anywhere.
|
|
14667
|
+
*/
|
|
14668
|
+
fuseNormMatVecDecodeEntries() {
|
|
14669
|
+
const maxBindings = this.ctx.limits.maxStorageBuffersPerShaderStage;
|
|
14670
|
+
const dbg = process.env?.GERBIL_FUSE_DEBUG != null;
|
|
14671
|
+
if (process.env?.GERBIL_FUSE_NORM_QKV == null || maxBindings < 10) {
|
|
14672
|
+
if (dbg) console.log(`[norm-qkv-debug] disabled (maxBindings=${maxBindings})`);
|
|
14673
|
+
return;
|
|
14674
|
+
}
|
|
14675
|
+
if (process.env?.GERBIL_MATVEC_GEO != null) return;
|
|
14676
|
+
let fusionCount = 0;
|
|
14677
|
+
for (let i = 0; i < this.decodeEntries.length - 1; i++) {
|
|
14678
|
+
const e0 = this.decodeEntries[i];
|
|
14679
|
+
const e1 = this.decodeEntries[i + 1];
|
|
14680
|
+
if (e0.node.opType !== "ResidualRMSNorm") continue;
|
|
14681
|
+
if (e1.spec !== MATVEC_INT4_SPEC || e1.node.opType !== "MatMulInt4") continue;
|
|
14682
|
+
const normOut = e0.node.outputs[1];
|
|
14683
|
+
if (e1.node.inputs[0] !== normOut) {
|
|
14684
|
+
if (dbg) console.log(`[norm-qkv-debug] i=${i} reject=norm-input e1in=${e1.node.inputs[0]} normOut=${normOut}`);
|
|
14685
|
+
continue;
|
|
14686
|
+
}
|
|
14687
|
+
const aBuf = this.getBuffer(e0.node.inputs[0]);
|
|
14688
|
+
const bBuf = this.getBuffer(e0.node.inputs[1]);
|
|
14689
|
+
const wBuf = this.getBuffer(e0.node.inputs[2]);
|
|
14690
|
+
const qBuf = this.getBuffer(e1.node.inputs[1]);
|
|
14691
|
+
const sBuf = this.getBuffer(e1.node.inputs[2]);
|
|
14692
|
+
const zBuf = this.getBuffer(e1.node.inputs[3]);
|
|
14693
|
+
const outBuf = this.getBuffer(e1.node.outputs[0]);
|
|
14694
|
+
const residualBuf = this.getBuffer(e0.node.outputs[0]);
|
|
14695
|
+
const normBuf = this.getBuffer(normOut);
|
|
14696
|
+
if (!(aBuf && bBuf && wBuf && qBuf && sBuf && zBuf && outBuf && residualBuf && normBuf)) {
|
|
14697
|
+
if (dbg) console.log(`[norm-qkv-debug] i=${i} reject=missing-buffer`);
|
|
14698
|
+
continue;
|
|
14699
|
+
}
|
|
14700
|
+
const reads = [
|
|
14701
|
+
aBuf,
|
|
14702
|
+
bBuf,
|
|
14703
|
+
wBuf,
|
|
14704
|
+
qBuf,
|
|
14705
|
+
sBuf,
|
|
14706
|
+
zBuf
|
|
14707
|
+
];
|
|
14708
|
+
const writes = [
|
|
14709
|
+
outBuf,
|
|
14710
|
+
residualBuf,
|
|
14711
|
+
normBuf
|
|
14712
|
+
];
|
|
14713
|
+
let hazard = false;
|
|
14714
|
+
for (const w of writes) if (reads.includes(w)) {
|
|
14715
|
+
hazard = true;
|
|
14716
|
+
break;
|
|
14717
|
+
}
|
|
14718
|
+
if (hazard || outBuf === residualBuf || outBuf === normBuf || residualBuf === normBuf) {
|
|
14719
|
+
if (dbg) console.log(`[norm-qkv-debug] i=${i} reject=aliasing`);
|
|
14720
|
+
continue;
|
|
14721
|
+
}
|
|
14722
|
+
e1.node.attributes.norm_eps = e0.node.attributes.eps ?? 1e-6;
|
|
14723
|
+
const fusedSpec = NORM_MATVEC_INT4_SPEC;
|
|
14724
|
+
const fusedPipeline = getOrCreatePipeline(this.ctx, `norm_mv_${e1.nodeId}`, fusedSpec.shaderCode, fusedSpec.entryPoint);
|
|
14725
|
+
const fusedUniformData = fusedSpec.buildParams(e1.node, {}, { seqPos: 0 });
|
|
14726
|
+
const fusedUniform = createUniformBuffer(this.ctx, `uniform_norm_mv_${e1.nodeId}`, fusedUniformData);
|
|
14727
|
+
const bufferEntries = [
|
|
14728
|
+
{ buffer: aBuf },
|
|
14729
|
+
{ buffer: bBuf },
|
|
14730
|
+
{ buffer: wBuf },
|
|
14731
|
+
{ buffer: qBuf },
|
|
14732
|
+
{ buffer: sBuf },
|
|
14733
|
+
{ buffer: zBuf },
|
|
14734
|
+
{ buffer: outBuf },
|
|
14735
|
+
{ buffer: residualBuf },
|
|
14736
|
+
{ buffer: normBuf },
|
|
14737
|
+
{ buffer: fusedUniform }
|
|
14738
|
+
];
|
|
14739
|
+
const fusedBindGroup = createBindGroup(this.ctx, fusedPipeline, bufferEntries, `bg_norm_mv_${e1.nodeId}`);
|
|
14740
|
+
const fusedEntry = {
|
|
14741
|
+
nodeId: `fused_norm_qkv_${e1.nodeId}`,
|
|
14742
|
+
fusedNodeIds: [...this.coveredNodeIds(e0), ...this.coveredNodeIds(e1)],
|
|
14743
|
+
node: e1.node,
|
|
14744
|
+
spec: fusedSpec,
|
|
14745
|
+
pipeline: fusedPipeline,
|
|
14746
|
+
bindGroup: fusedBindGroup,
|
|
14747
|
+
uniformBuffer: fusedUniform,
|
|
14748
|
+
lastParamsBytes: null,
|
|
14749
|
+
lastDispatchSize: null
|
|
14750
|
+
};
|
|
14751
|
+
this.decodeEntries.splice(i, 2, fusedEntry);
|
|
14752
|
+
fusionCount++;
|
|
14753
|
+
}
|
|
14754
|
+
if (fusionCount > 0) console.log(`[executor] Fused ${fusionCount} norm+qkv Mamba blocks (saved ${fusionCount} dispatches)`);
|
|
14755
|
+
}
|
|
14245
14756
|
fuseNormSwiGLUDecodeEntries() {
|
|
14246
14757
|
const maxBindings = this.ctx.limits.maxStorageBuffersPerShaderStage;
|
|
14247
14758
|
if (process.env?.GERBIL_FUSE_DEBUG != null) {
|
|
@@ -14420,6 +14931,158 @@ var Executor = class Executor {
|
|
|
14420
14931
|
if (fusionCount > 0) console.log(`[executor] Fused ${fusionCount} SwiGLU MLP blocks (saved ${fusionCount * 2} dispatches)`);
|
|
14421
14932
|
}
|
|
14422
14933
|
/**
|
|
14934
|
+
* Fuse the Mamba-front projections — qkv (MatVec), a (MatVec), b (MatVec),
|
|
14935
|
+
* z (MatVec) — into ONE QuadMatVecInt4 dispatch per Mamba layer. All four
|
|
14936
|
+
* read the SAME `input_layernorm` output and share K and group_size, so their
|
|
14937
|
+
* INT4 weight matrices (stored column-major, one column contiguous) are
|
|
14938
|
+
* column-concatenated into a single q/scales/zeros buffer by a plain
|
|
14939
|
+
* byte-append (a GPU buffer-to-buffer copy at init). The fused kernel runs a
|
|
14940
|
+
* single GEMV over the concatenated N and routes each output column back to
|
|
14941
|
+
* its ORIGINAL destination buffer, so downstream consumers (conv1d reads qkv,
|
|
14942
|
+
* MambaSSM reads a/b, the gated out_proj reads z) are untouched.
|
|
14943
|
+
*
|
|
14944
|
+
* Numerically identical to the three separate dispatches (same dequant, same
|
|
14945
|
+
* K-parallel reduction, same per-column sum — the concatenated weights are the
|
|
14946
|
+
* same bytes), so this is token-exact and WebKit-safe. It removes 2
|
|
14947
|
+
* dispatches per Mamba layer (3 front dispatches → 1) plus one input read and
|
|
14948
|
+
* two weight-read passes.
|
|
14949
|
+
*
|
|
14950
|
+
* Runs BEFORE fuseDualMatVecDecodeEntries so the a/b pair is still two
|
|
14951
|
+
* standalone MatVec entries here (not yet a DualMatVec). Restricted to the
|
|
14952
|
+
* Mamba front by weight-name (`.in_proj_qkv/a/b/z.`), so the attention front
|
|
14953
|
+
* (q/gate/k/v) is never matched. Bindings: A + q/s/z + 4 outs + params = 9
|
|
14954
|
+
* (fits Metal's 10-binding wall). Gate: default-on when the adapter offers ≥9
|
|
14955
|
+
* storage bindings; opt-out via GERBIL_FUSE_MAMBA_FRONT=0.
|
|
14956
|
+
*/
|
|
14957
|
+
fuseMambaFrontDecodeEntries() {
|
|
14958
|
+
const maxBindings = this.ctx.limits.maxStorageBuffersPerShaderStage;
|
|
14959
|
+
const dbg = process.env?.GERBIL_FUSE_DEBUG != null;
|
|
14960
|
+
if (process.env?.GERBIL_FUSE_MAMBA_FRONT === "0" || maxBindings < 9) {
|
|
14961
|
+
if (dbg) console.log(`[mamba-front-debug] disabled (flag=${process.env?.GERBIL_FUSE_MAMBA_FRONT ?? "unset"} maxBindings=${maxBindings})`);
|
|
14962
|
+
return;
|
|
14963
|
+
}
|
|
14964
|
+
const isFrontProj = (entry, marker) => {
|
|
14965
|
+
if (entry.node.opType !== "MatMulInt4") return false;
|
|
14966
|
+
if (entry.fusedNodeIds && entry.fusedNodeIds.length > 1) return false;
|
|
14967
|
+
const wq = entry.node.inputs[1];
|
|
14968
|
+
return typeof wq === "string" && wq.includes(marker);
|
|
14969
|
+
};
|
|
14970
|
+
let fusionCount = 0;
|
|
14971
|
+
for (let i = 0; i < this.decodeEntries.length - 3; i++) {
|
|
14972
|
+
const eQkv = this.decodeEntries[i];
|
|
14973
|
+
const eA = this.decodeEntries[i + 1];
|
|
14974
|
+
const eB = this.decodeEntries[i + 2];
|
|
14975
|
+
const eZ = this.decodeEntries[i + 3];
|
|
14976
|
+
if (!isFrontProj(eQkv, ".in_proj_qkv.") || !isFrontProj(eA, ".in_proj_a.") || !isFrontProj(eB, ".in_proj_b.") || !isFrontProj(eZ, ".in_proj_z.")) continue;
|
|
14977
|
+
const inputTensor = eQkv.node.inputs[0];
|
|
14978
|
+
if (eA.node.inputs[0] !== inputTensor || eB.node.inputs[0] !== inputTensor || eZ.node.inputs[0] !== inputTensor) continue;
|
|
14979
|
+
const K = eQkv.node.attributes.K;
|
|
14980
|
+
const groupSize = eQkv.node.attributes.group_size;
|
|
14981
|
+
if (![
|
|
14982
|
+
eA,
|
|
14983
|
+
eB,
|
|
14984
|
+
eZ
|
|
14985
|
+
].every((e) => e.node.attributes.K === K && e.node.attributes.group_size === groupSize)) continue;
|
|
14986
|
+
const nQkv = eQkv.node.attributes.N;
|
|
14987
|
+
const nA = eA.node.attributes.N;
|
|
14988
|
+
const nB = eB.node.attributes.N;
|
|
14989
|
+
const nZ = eZ.node.attributes.N;
|
|
14990
|
+
const inputBuf = this.getBuffer(inputTensor);
|
|
14991
|
+
const outQkv = this.getBuffer(eQkv.node.outputs[0]);
|
|
14992
|
+
const outA = this.getBuffer(eA.node.outputs[0]);
|
|
14993
|
+
const outB = this.getBuffer(eB.node.outputs[0]);
|
|
14994
|
+
const outZ = this.getBuffer(eZ.node.outputs[0]);
|
|
14995
|
+
if (!(inputBuf && outQkv && outA && outB && outZ)) continue;
|
|
14996
|
+
const outs = [
|
|
14997
|
+
outQkv,
|
|
14998
|
+
outA,
|
|
14999
|
+
outB,
|
|
15000
|
+
outZ
|
|
15001
|
+
];
|
|
15002
|
+
if (outs.includes(inputBuf)) continue;
|
|
15003
|
+
if (new Set(outs).size !== outs.length) continue;
|
|
15004
|
+
const concat = (label, tensors) => {
|
|
15005
|
+
const bufs = tensors.map((t) => this.getBuffer(t));
|
|
15006
|
+
const total = bufs.reduce((acc, b) => acc + b.size, 0);
|
|
15007
|
+
const dst = createStorageBuffer(this.ctx, label, total);
|
|
15008
|
+
const enc = this.ctx.device.createCommandEncoder();
|
|
15009
|
+
let off = 0;
|
|
15010
|
+
for (const b of bufs) {
|
|
15011
|
+
enc.copyBufferToBuffer(b, 0, dst, off, b.size);
|
|
15012
|
+
off += b.size;
|
|
15013
|
+
}
|
|
15014
|
+
this.ctx.device.queue.submit([enc.finish()]);
|
|
15015
|
+
this.mambaFrontConcatBuffers.push(dst);
|
|
15016
|
+
return dst;
|
|
15017
|
+
};
|
|
15018
|
+
const concatQ = concat(`mamba_front_q_${eQkv.nodeId}`, [
|
|
15019
|
+
eQkv.node.inputs[1],
|
|
15020
|
+
eA.node.inputs[1],
|
|
15021
|
+
eB.node.inputs[1],
|
|
15022
|
+
eZ.node.inputs[1]
|
|
15023
|
+
]);
|
|
15024
|
+
const concatS = concat(`mamba_front_s_${eQkv.nodeId}`, [
|
|
15025
|
+
eQkv.node.inputs[2],
|
|
15026
|
+
eA.node.inputs[2],
|
|
15027
|
+
eB.node.inputs[2],
|
|
15028
|
+
eZ.node.inputs[2]
|
|
15029
|
+
]);
|
|
15030
|
+
const concatZ = concat(`mamba_front_z_${eQkv.nodeId}`, [
|
|
15031
|
+
eQkv.node.inputs[3],
|
|
15032
|
+
eA.node.inputs[3],
|
|
15033
|
+
eB.node.inputs[3],
|
|
15034
|
+
eZ.node.inputs[3]
|
|
15035
|
+
]);
|
|
15036
|
+
const fusedNode = {
|
|
15037
|
+
...eQkv.node,
|
|
15038
|
+
attributes: {
|
|
15039
|
+
...eQkv.node.attributes,
|
|
15040
|
+
K,
|
|
15041
|
+
N: nQkv + nA + nB + nZ,
|
|
15042
|
+
group_size: groupSize,
|
|
15043
|
+
n0: nQkv,
|
|
15044
|
+
n1: nQkv + nA,
|
|
15045
|
+
n2: nQkv + nA + nB
|
|
15046
|
+
}
|
|
15047
|
+
};
|
|
15048
|
+
const fusedSpec = QUAD_MATVEC_INT4_SPEC;
|
|
15049
|
+
const fusedPipeline = getOrCreatePipeline(this.ctx, `quad_mv_${eQkv.nodeId}`, fusedSpec.shaderCode, fusedSpec.entryPoint);
|
|
15050
|
+
const fusedUniform = createUniformBuffer(this.ctx, `uniform_quad_mv_${eQkv.nodeId}`, fusedSpec.buildParams(fusedNode, {}, { seqPos: 0 }));
|
|
15051
|
+
const bufferEntries = [
|
|
15052
|
+
{ buffer: inputBuf },
|
|
15053
|
+
{ buffer: concatQ },
|
|
15054
|
+
{ buffer: concatS },
|
|
15055
|
+
{ buffer: concatZ },
|
|
15056
|
+
{ buffer: outQkv },
|
|
15057
|
+
{ buffer: outA },
|
|
15058
|
+
{ buffer: outB },
|
|
15059
|
+
{ buffer: outZ },
|
|
15060
|
+
{ buffer: fusedUniform }
|
|
15061
|
+
];
|
|
15062
|
+
const fusedBindGroup = createBindGroup(this.ctx, fusedPipeline, bufferEntries, `bg_quad_mv_${eQkv.nodeId}`);
|
|
15063
|
+
const fusedEntry = {
|
|
15064
|
+
nodeId: `fused_mamba_front_${eQkv.nodeId}`,
|
|
15065
|
+
fusedNodeIds: [
|
|
15066
|
+
...this.coveredNodeIds(eQkv),
|
|
15067
|
+
...this.coveredNodeIds(eA),
|
|
15068
|
+
...this.coveredNodeIds(eB),
|
|
15069
|
+
...this.coveredNodeIds(eZ)
|
|
15070
|
+
],
|
|
15071
|
+
node: fusedNode,
|
|
15072
|
+
spec: fusedSpec,
|
|
15073
|
+
pipeline: fusedPipeline,
|
|
15074
|
+
bindGroup: fusedBindGroup,
|
|
15075
|
+
uniformBuffer: fusedUniform,
|
|
15076
|
+
lastParamsBytes: null,
|
|
15077
|
+
lastDispatchSize: null
|
|
15078
|
+
};
|
|
15079
|
+
this.decodeEntries.splice(i, 4, fusedEntry);
|
|
15080
|
+
fusionCount++;
|
|
15081
|
+
}
|
|
15082
|
+
if (fusionCount > 0) console.log(`[executor] Fused ${fusionCount} mamba-front qkv+a+b+z projections (saved ${fusionCount * 2} dispatches)`);
|
|
15083
|
+
else if (dbg) console.log("[mamba-front-debug] no matching qkv/a/b/z windows found");
|
|
15084
|
+
}
|
|
15085
|
+
/**
|
|
14423
15086
|
* Fuse two adjacent INT4 projections that share the same input activation and
|
|
14424
15087
|
* the same K/N (e.g. q_proj+gate_proj and k_proj+v_proj in full-attention
|
|
14425
15088
|
* decode) into a single DualMatVecInt4 dispatch. Reads the shared input vector
|
|
@@ -17964,6 +18627,118 @@ function quantizeKaniBackbone(graph, weights) {
|
|
|
17964
18627
|
* the embedding/LM-head are left untouched (the graph keeps them f32). Idempotent.
|
|
17965
18628
|
*/
|
|
17966
18629
|
const quantizeBackboneInt4 = quantizeKaniBackbone;
|
|
18630
|
+
/** Default NanoCodec checkpoint Gepard uses (NeMo 22 kHz, 1.89 kbps, 21.5 fps). */
|
|
18631
|
+
const GEPARD_NANOCODEC_REPO = "mlx-community/nemo-nano-codec-22khz-1.89kbps-21.5fps";
|
|
18632
|
+
/** Overlay key prefixes that run host-side (kept f32 in a separate map). */
|
|
18633
|
+
const GEPARD_OVERLAY_RE = /^(audio_embeddings\.|audio_embed_proj\.|audio_embed_scale$|codebook_heads\.|stop_head\.)/;
|
|
18634
|
+
/** Voice-cloning / training-only tensors, skipped until cloning is wired. */
|
|
18635
|
+
const GEPARD_SKIP_RE = /^(ref_compressor\.|null_prefix$|supcon_head\.)/;
|
|
18636
|
+
async function loadGepardTTS(options) {
|
|
18637
|
+
const revision = options.revision ?? "main";
|
|
18638
|
+
const repo = options.repo ?? "nineninesix/gepard-1.0";
|
|
18639
|
+
const codecRepo = options.codecRepo ?? GEPARD_NANOCODEC_REPO;
|
|
18640
|
+
const baseURL = await resolveModelBaseURL(repo, revision);
|
|
18641
|
+
const codecURL = await resolveModelBaseURL(codecRepo, "main");
|
|
18642
|
+
await initFs();
|
|
18643
|
+
options.onProgress?.(0, 100, "Fetching config + tokenizer…");
|
|
18644
|
+
const [rawConfig, gepardConfig, tokenizerJSON] = await Promise.all([
|
|
18645
|
+
fetchJSON(baseURL, "config.json", options.hfToken, options.cacheDir),
|
|
18646
|
+
fetchJSON(baseURL, "gepard_config.json", options.hfToken, options.cacheDir),
|
|
18647
|
+
fetchJSON(baseURL, "tokenizer.json", options.hfToken, options.cacheDir)
|
|
18648
|
+
]);
|
|
18649
|
+
const tokenizerConfig = await fetchJSON(baseURL, "tokenizer_config.json", options.hfToken, options.cacheDir).catch(() => null);
|
|
18650
|
+
const tokenizer = Tokenizer.fromJSON(tokenizerJSON, tokenizerConfig);
|
|
18651
|
+
options.onProgress?.(8, 100, "Downloading Gepard backbone…");
|
|
18652
|
+
const mainBuf = await fetchBinary(baseURL, "model.safetensors", options.hfToken, (loaded, total) => options.onProgress?.(8 + loaded / (total || 1) * 55, 100, "backbone weights"), options.cacheDir);
|
|
18653
|
+
const mainFile = parseSafetensorsHeader(mainBuf);
|
|
18654
|
+
const backboneWeights = /* @__PURE__ */ new Map();
|
|
18655
|
+
const overlayWeights = /* @__PURE__ */ new Map();
|
|
18656
|
+
let embedTokensBF16 = null;
|
|
18657
|
+
for (const entry of mainFile.entries) {
|
|
18658
|
+
if (GEPARD_SKIP_RE.test(entry.name)) continue;
|
|
18659
|
+
if (GEPARD_OVERLAY_RE.test(entry.name)) {
|
|
18660
|
+
overlayWeights.set(entry.name, {
|
|
18661
|
+
data: toF32(entry, mainBuf, mainFile),
|
|
18662
|
+
shape: entry.shape
|
|
18663
|
+
});
|
|
18664
|
+
continue;
|
|
18665
|
+
}
|
|
18666
|
+
if (!entry.name.startsWith("model.")) continue;
|
|
18667
|
+
const key = entry.name.slice(6);
|
|
18668
|
+
if (key === "embed_tokens.weight") {
|
|
18669
|
+
if (entry.dtype !== "BF16") throw new Error(`Gepard embed_tokens.weight is ${entry.dtype}; expected BF16.`);
|
|
18670
|
+
const view = getTensorData(mainBuf, mainFile, entry);
|
|
18671
|
+
embedTokensBF16 = new Uint16Array(view.buffer, view.byteOffset, view.byteLength / 2);
|
|
18672
|
+
continue;
|
|
18673
|
+
}
|
|
18674
|
+
backboneWeights.set(key, {
|
|
18675
|
+
data: toF32(entry, mainBuf, mainFile),
|
|
18676
|
+
shape: entry.shape
|
|
18677
|
+
});
|
|
18678
|
+
}
|
|
18679
|
+
if (!embedTokensBF16) throw new Error("Gepard checkpoint has no model.embed_tokens.weight.");
|
|
18680
|
+
const numLayers = rawConfig.num_hidden_layers;
|
|
18681
|
+
const numHeads = rawConfig.num_attention_heads;
|
|
18682
|
+
const headDim = rawConfig.head_dim ?? Math.floor(rawConfig.hidden_size / numHeads);
|
|
18683
|
+
for (let i = 0; i < numLayers; i++) {
|
|
18684
|
+
const qProjKey = `layers.${i}.self_attn.q_proj.weight`;
|
|
18685
|
+
const attnGateKey = `layers.${i}.self_attn.attn_gate.weight`;
|
|
18686
|
+
const fused = backboneWeights.get(qProjKey);
|
|
18687
|
+
if (!fused || backboneWeights.has(attnGateKey)) continue;
|
|
18688
|
+
const [fusedRows, cols] = fused.shape;
|
|
18689
|
+
if (fusedRows !== numHeads * 2 * headDim) continue;
|
|
18690
|
+
const src = fused.data;
|
|
18691
|
+
const halfRows = fusedRows / 2;
|
|
18692
|
+
const qData = new Float32Array(halfRows * cols);
|
|
18693
|
+
const gateData = new Float32Array(halfRows * cols);
|
|
18694
|
+
const headBlock = 2 * headDim;
|
|
18695
|
+
for (let h = 0; h < numHeads; h++) {
|
|
18696
|
+
const srcBase = h * headBlock * cols;
|
|
18697
|
+
const dstBase = h * headDim * cols;
|
|
18698
|
+
qData.set(src.subarray(srcBase, srcBase + headDim * cols), dstBase);
|
|
18699
|
+
gateData.set(src.subarray(srcBase + headDim * cols, srcBase + headBlock * cols), dstBase);
|
|
18700
|
+
}
|
|
18701
|
+
backboneWeights.set(qProjKey, {
|
|
18702
|
+
data: qData,
|
|
18703
|
+
shape: [halfRows, cols]
|
|
18704
|
+
});
|
|
18705
|
+
backboneWeights.set(attnGateKey, {
|
|
18706
|
+
data: gateData,
|
|
18707
|
+
shape: [halfRows, cols]
|
|
18708
|
+
});
|
|
18709
|
+
}
|
|
18710
|
+
const normKeys = ["norm.weight"];
|
|
18711
|
+
for (let i = 0; i < numLayers; i++) normKeys.push(`layers.${i}.input_layernorm.weight`, `layers.${i}.post_attention_layernorm.weight`, `layers.${i}.self_attn.q_norm.weight`, `layers.${i}.self_attn.k_norm.weight`);
|
|
18712
|
+
for (const key of normKeys) {
|
|
18713
|
+
const w = backboneWeights.get(key);
|
|
18714
|
+
if (!w) throw new Error(`Gepard backbone is missing norm weight ${key}.`);
|
|
18715
|
+
const data = w.data;
|
|
18716
|
+
for (let j = 0; j < data.length; j++) data[j] += 1;
|
|
18717
|
+
}
|
|
18718
|
+
options.onProgress?.(65, 100, "Downloading NanoCodec decoder…");
|
|
18719
|
+
const codecCacheDir = options.cacheDir ? `${options.cacheDir}/nanocodec-21_5fps` : void 0;
|
|
18720
|
+
const codecBuf = await fetchBinary(codecURL, "model.safetensors", options.hfToken, (loaded, total) => options.onProgress?.(65 + loaded / (total || 1) * 30, 100, "codec weights"), codecCacheDir);
|
|
18721
|
+
const codecFile = parseSafetensorsHeader(codecBuf);
|
|
18722
|
+
const codecRaw = /* @__PURE__ */ new Map();
|
|
18723
|
+
for (const entry of codecFile.entries) {
|
|
18724
|
+
if (!entry.name.startsWith("audio_decoder.")) continue;
|
|
18725
|
+
codecRaw.set(entry.name, {
|
|
18726
|
+
data: toF32(entry, codecBuf, codecFile),
|
|
18727
|
+
shape: entry.shape
|
|
18728
|
+
});
|
|
18729
|
+
}
|
|
18730
|
+
const codecWeights = buildKaniCodecWeights(codecRaw);
|
|
18731
|
+
options.onProgress?.(100, 100, "Gepard TTS loaded.");
|
|
18732
|
+
return {
|
|
18733
|
+
backboneWeights,
|
|
18734
|
+
overlayWeights,
|
|
18735
|
+
embedTokensBF16,
|
|
18736
|
+
codecWeights,
|
|
18737
|
+
tokenizer,
|
|
18738
|
+
rawConfig,
|
|
18739
|
+
gepardConfig
|
|
18740
|
+
};
|
|
18741
|
+
}
|
|
17967
18742
|
const OUTETTS_DEFAULT_REPO = "OuteAI/OuteTTS-1.0-0.6B";
|
|
17968
18743
|
/** Parse a folded DAC safetensors buffer into a canonical-name weight map (F32 only). */
|
|
17969
18744
|
function parseFoldedDacSafetensors(buf) {
|
|
@@ -18510,5 +19285,5 @@ function selectGraphWeights(graph, weights) {
|
|
|
18510
19285
|
}
|
|
18511
19286
|
|
|
18512
19287
|
//#endregion
|
|
18513
|
-
export {
|
|
18514
|
-
//# sourceMappingURL=moonshine-stt-
|
|
19288
|
+
export { createStorageBuffer as C, initGPU as D, getOrCreatePipeline as E, verifyGPU as O, createBindGroup as S, destroyBuffers as T, fetchAdapter as _, loadKaniTTS as a, MATMUL_BIAS_F16C_SPEC as b, loadOuteTTS as c, quantizeKaniBackbone as d, remapPrunedToken as f, buildLoRADeltas as g, applyLoRAToStore as h, loadGepardTTS as i, loadParlerTTS as l, Executor as m, MoonshineEncoderExecutor as n, loadModel as o, Tokenizer as p, createKeyMapperForArch as r, loadMoonshine as s, MoonshineSTT as t, quantizeBackboneInt4 as u, splitFusedQGateDelta as v, createUniformBuffer as w, clearPipelineCache as x, KERNEL_REGISTRY as y };
|
|
19289
|
+
//# sourceMappingURL=moonshine-stt-COJeK2Zb.mjs.map
|