@tryhamster/gerbil 1.6.0 → 1.6.1
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/dist/{architectures-Ol7yUmKf.mjs → architectures-CpmH4Gha.mjs} +86 -82
- package/dist/architectures-CpmH4Gha.mjs.map +1 -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-_68_JSGV.d.mts → gerbil--YDNklsY.d.mts} +2 -2
- package/dist/{gerbil-_68_JSGV.d.mts.map → gerbil--YDNklsY.d.mts.map} +1 -1
- package/dist/gerbil-BPpEvMYU.mjs +4 -0
- package/dist/{gerbil-DsgYlkgf.mjs → gerbil-D13QAlk2.mjs} +2 -2
- package/dist/{gerbil-DsgYlkgf.mjs.map → gerbil-D13QAlk2.mjs.map} +1 -1
- package/dist/gpu/hooks.d.mts +1 -1
- package/dist/gpu/index.d.mts +1 -1
- package/dist/gpu/index.mjs +3 -3
- package/dist/{gpu-DfUBMUGe.mjs → gpu-DtA28rES.mjs} +12 -8
- package/dist/{gpu-DfUBMUGe.mjs.map → gpu-DtA28rES.mjs.map} +1 -1
- package/dist/{index-BdbeBDTz.d.mts → index-CZVIHkaH.d.mts} +23 -2
- package/dist/index-CZVIHkaH.d.mts.map +1 -0
- package/dist/index.d.mts +2 -2
- package/dist/index.mjs +5 -5
- 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-D3XRaVAc.mjs → mcp-DuaxaIwt.mjs} +3 -3
- package/dist/{mcp-D3XRaVAc.mjs.map → mcp-DuaxaIwt.mjs.map} +1 -1
- package/dist/moonshine-stt-BkU9cx2d.mjs +4 -0
- package/dist/{moonshine-stt-bExSnd2o.mjs → moonshine-stt-BwX7NQ0O.mjs} +749 -2
- package/dist/moonshine-stt-BwX7NQ0O.mjs.map +1 -0
- package/dist/{one-liner-UvMt5g0G.mjs → one-liner-CMk_GZgS.mjs} +2 -2
- package/dist/{one-liner-UvMt5g0G.mjs.map → one-liner-CMk_GZgS.mjs.map} +1 -1
- package/dist/repl-CD4NDrku.mjs +9 -0
- package/dist/skills/index.d.mts +4 -4
- package/dist/skills/index.mjs +3 -3
- package/dist/{skills-BJUofoaY.mjs → skills-CJHhljKM.mjs} +2 -2
- package/dist/{skills-BJUofoaY.mjs.map → skills-CJHhljKM.mjs.map} +1 -1
- package/package.json +1 -1
- package/dist/architectures-Ol7yUmKf.mjs.map +0 -1
- package/dist/gerbil-CNyb67ca.mjs +0 -4
- package/dist/index-BdbeBDTz.d.mts.map +0 -1
- package/dist/moonshine-stt-B0dwDCgT.mjs +0 -4
- package/dist/moonshine-stt-bExSnd2o.mjs.map +0 -1
- package/dist/repl-B5xHAr9n.mjs +0 -9
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
import { C as quantizeInt4, D as createDefaultHFKeyMapper, S as DEFAULT_GROUP_SIZE, T as DTYPE_BYTES, a as generateMoonshineEncoderGraph, b as nanoCodecWeightMap, i as generateMoonshineDecoderGraph, o as moonshineEncoderFrames, p as foldNanoCodecWeightNorm, s as parseMoonshineConfig, t as generateGraph, w as CANONICAL_KEYS } from "./architectures-
|
|
1
|
+
import { C as quantizeInt4, D as createDefaultHFKeyMapper, S as DEFAULT_GROUP_SIZE, T as DTYPE_BYTES, a as generateMoonshineEncoderGraph, b as nanoCodecWeightMap, i as generateMoonshineDecoderGraph, o as moonshineEncoderFrames, p as foldNanoCodecWeightNorm, s as parseMoonshineConfig, t as generateGraph, w as CANONICAL_KEYS } from "./architectures-CpmH4Gha.mjs";
|
|
2
2
|
|
|
3
3
|
//#region src/gpu/device.ts
|
|
4
4
|
const COPY_SRC = 4;
|
|
@@ -2341,6 +2341,98 @@ fn main(
|
|
|
2341
2341
|
}
|
|
2342
2342
|
}
|
|
2343
2343
|
`;
|
|
2344
|
+
const WGSL_DUAL_RMSNORM_ROPE = `\
|
|
2345
|
+
struct Params {
|
|
2346
|
+
rows0: u32, // T * num_q_heads
|
|
2347
|
+
rows1: u32, // T * num_kv_heads
|
|
2348
|
+
head_dim: u32,
|
|
2349
|
+
eps_bits: u32,
|
|
2350
|
+
num_q_heads: u32,
|
|
2351
|
+
num_kv_heads: u32,
|
|
2352
|
+
rope_base_bits: u32,
|
|
2353
|
+
position_offset: u32,
|
|
2354
|
+
rope_half: u32,
|
|
2355
|
+
rope_denom: u32,
|
|
2356
|
+
rope_active_pairs: u32,
|
|
2357
|
+
_pad: u32,
|
|
2358
|
+
}
|
|
2359
|
+
|
|
2360
|
+
@group(0) @binding(0) var<storage, read> input0: array<f32>;
|
|
2361
|
+
@group(0) @binding(1) var<storage, read> weight0: array<f32>;
|
|
2362
|
+
@group(0) @binding(2) var<storage, read> input1: array<f32>;
|
|
2363
|
+
@group(0) @binding(3) var<storage, read> weight1: array<f32>;
|
|
2364
|
+
@group(0) @binding(4) var<storage, read_write> output0: array<f32>;
|
|
2365
|
+
@group(0) @binding(5) var<storage, read_write> output1: array<f32>;
|
|
2366
|
+
@group(0) @binding(6) var<storage, read> params: Params;
|
|
2367
|
+
|
|
2368
|
+
var<workgroup> shared_sum: array<f32, 256>;
|
|
2369
|
+
var<workgroup> shared_row: array<f32, 256>;
|
|
2370
|
+
|
|
2371
|
+
@compute @workgroup_size(256)
|
|
2372
|
+
fn main(
|
|
2373
|
+
@builtin(local_invocation_id) lid: vec3u,
|
|
2374
|
+
@builtin(workgroup_id) wid: vec3u,
|
|
2375
|
+
) {
|
|
2376
|
+
let wg = wid.x;
|
|
2377
|
+
if (wg >= params.rows0 + params.rows1) { return; }
|
|
2378
|
+
|
|
2379
|
+
let tid = lid.x;
|
|
2380
|
+
let hd = params.head_dim;
|
|
2381
|
+
let eps = bitcast<f32>(params.eps_bits);
|
|
2382
|
+
|
|
2383
|
+
let is1 = wg >= params.rows0;
|
|
2384
|
+
let row = select(wg, wg - params.rows0, is1);
|
|
2385
|
+
let row_offset = row * hd;
|
|
2386
|
+
let heads = select(params.num_q_heads, params.num_kv_heads, is1);
|
|
2387
|
+
let pos = row / heads;
|
|
2388
|
+
|
|
2389
|
+
// ── Phase 1: row sum of squares ──
|
|
2390
|
+
var val: f32 = 0.0;
|
|
2391
|
+
if (tid < hd) {
|
|
2392
|
+
if (is1) { val = input1[row_offset + tid]; } else { val = input0[row_offset + tid]; }
|
|
2393
|
+
}
|
|
2394
|
+
shared_sum[tid] = val * val;
|
|
2395
|
+
workgroupBarrier();
|
|
2396
|
+
var stride: u32 = 128u;
|
|
2397
|
+
while (stride > 0u) {
|
|
2398
|
+
if (tid < stride) { shared_sum[tid] += shared_sum[tid + stride]; }
|
|
2399
|
+
workgroupBarrier();
|
|
2400
|
+
stride = stride / 2u;
|
|
2401
|
+
}
|
|
2402
|
+
let rms = sqrt(shared_sum[0] / f32(hd) + eps);
|
|
2403
|
+
|
|
2404
|
+
// ── Phase 2: normalized row into shared memory ──
|
|
2405
|
+
if (tid < hd) {
|
|
2406
|
+
var w: f32;
|
|
2407
|
+
if (is1) { w = weight1[tid]; } else { w = weight0[tid]; }
|
|
2408
|
+
shared_row[tid] = (val / rms) * w;
|
|
2409
|
+
}
|
|
2410
|
+
workgroupBarrier();
|
|
2411
|
+
|
|
2412
|
+
// ── Phase 3: rotate_half rotation from shared pairs, write out ──
|
|
2413
|
+
if (tid < hd) {
|
|
2414
|
+
let rope_half = params.rope_half;
|
|
2415
|
+
var out_val = shared_row[tid];
|
|
2416
|
+
if (tid < 2u * rope_half) {
|
|
2417
|
+
let is_hi = tid >= rope_half;
|
|
2418
|
+
let pair_idx = select(tid, tid - rope_half, is_hi);
|
|
2419
|
+
var cos_val = 1.0;
|
|
2420
|
+
var sin_val = 0.0;
|
|
2421
|
+
if (pair_idx < params.rope_active_pairs) {
|
|
2422
|
+
let rope_base = bitcast<f32>(params.rope_base_bits);
|
|
2423
|
+
let position = f32(pos + params.position_offset);
|
|
2424
|
+
let freq = position / pow(rope_base, f32(2u * pair_idx) / f32(params.rope_denom));
|
|
2425
|
+
cos_val = cos(freq);
|
|
2426
|
+
sin_val = sin(freq);
|
|
2427
|
+
}
|
|
2428
|
+
let lo = shared_row[pair_idx];
|
|
2429
|
+
let hi = shared_row[pair_idx + rope_half];
|
|
2430
|
+
out_val = select(lo * cos_val - hi * sin_val, lo * sin_val + hi * cos_val, is_hi);
|
|
2431
|
+
}
|
|
2432
|
+
if (is1) { output1[row_offset + tid] = out_val; } else { output0[row_offset + tid] = out_val; }
|
|
2433
|
+
}
|
|
2434
|
+
}
|
|
2435
|
+
`;
|
|
2344
2436
|
const WGSL_LAYERNORM = `\
|
|
2345
2437
|
// Layer Normalization: output = ((x - mean) / sqrt(var + eps)) * weight + bias
|
|
2346
2438
|
|
|
@@ -4235,6 +4327,235 @@ fn main(@builtin(global_invocation_id) gid: vec3u) {
|
|
|
4235
4327
|
output[t * params.channels + c] = sum / (1.0 + exp(max(-sum, -80.0)));
|
|
4236
4328
|
}
|
|
4237
4329
|
`;
|
|
4330
|
+
const WGSL_CAUSAL_CONV1D_SILU_UPDATE = `\
|
|
4331
|
+
struct Params {
|
|
4332
|
+
seq_len: u32,
|
|
4333
|
+
channels: u32,
|
|
4334
|
+
kernel_size: u32,
|
|
4335
|
+
_pad: u32,
|
|
4336
|
+
}
|
|
4337
|
+
|
|
4338
|
+
@group(0) @binding(0) var<storage, read> input: array<f32>;
|
|
4339
|
+
@group(0) @binding(1) var<storage, read> weight: array<f32>;
|
|
4340
|
+
@group(0) @binding(2) var<storage, read_write> conv_state: array<f32>;
|
|
4341
|
+
@group(0) @binding(3) var<storage, read_write> output: array<f32>;
|
|
4342
|
+
@group(0) @binding(4) var<storage, read> params: Params;
|
|
4343
|
+
|
|
4344
|
+
@compute @workgroup_size(256)
|
|
4345
|
+
fn main(@builtin(global_invocation_id) gid: vec3u) {
|
|
4346
|
+
let idx = gid.x;
|
|
4347
|
+
let total = params.seq_len * params.channels;
|
|
4348
|
+
if (idx >= total) { return; }
|
|
4349
|
+
|
|
4350
|
+
let t = idx / params.channels;
|
|
4351
|
+
let c = idx % params.channels;
|
|
4352
|
+
let state_size = params.kernel_size - 1u;
|
|
4353
|
+
|
|
4354
|
+
// Snapshot this channel's rolling state (kernel_size <= 8) before any write.
|
|
4355
|
+
var olds: array<f32, 8>;
|
|
4356
|
+
for (var s: u32 = 0u; s < state_size; s = s + 1u) {
|
|
4357
|
+
olds[s] = conv_state[s * params.channels + c];
|
|
4358
|
+
}
|
|
4359
|
+
|
|
4360
|
+
var sum: f32 = 0.0;
|
|
4361
|
+
for (var k: u32 = 0u; k < params.kernel_size; k = k + 1u) {
|
|
4362
|
+
let src_t = i32(t) - i32(k);
|
|
4363
|
+
var in_val: f32 = 0.0;
|
|
4364
|
+
if (src_t >= 0) {
|
|
4365
|
+
in_val = input[u32(src_t) * params.channels + c];
|
|
4366
|
+
} else {
|
|
4367
|
+
let state_t = i32(state_size) + src_t;
|
|
4368
|
+
if (state_t >= 0) {
|
|
4369
|
+
in_val = olds[u32(state_t)];
|
|
4370
|
+
}
|
|
4371
|
+
}
|
|
4372
|
+
let w_val = weight[c * params.kernel_size + (params.kernel_size - 1u - k)];
|
|
4373
|
+
sum += in_val * w_val;
|
|
4374
|
+
}
|
|
4375
|
+
|
|
4376
|
+
// Fused SiLU: silu(x) = x / (1 + exp(-x))
|
|
4377
|
+
output[t * params.channels + c] = sum / (1.0 + exp(max(-sum, -80.0)));
|
|
4378
|
+
|
|
4379
|
+
// Rolling-state shift (decode T=1): new state = [olds[1..], input[0,c]].
|
|
4380
|
+
if (params.seq_len == 1u && t == 0u) {
|
|
4381
|
+
for (var s: u32 = 0u; s + 1u < state_size; s = s + 1u) {
|
|
4382
|
+
conv_state[s * params.channels + c] = olds[s + 1u];
|
|
4383
|
+
}
|
|
4384
|
+
if (state_size > 0u) {
|
|
4385
|
+
conv_state[(state_size - 1u) * params.channels + c] = input[c];
|
|
4386
|
+
}
|
|
4387
|
+
}
|
|
4388
|
+
}
|
|
4389
|
+
`;
|
|
4390
|
+
const WGSL_CAUSAL_CONV1D_GATED_UPDATE = `\
|
|
4391
|
+
struct Params {
|
|
4392
|
+
seq_len: u32,
|
|
4393
|
+
channels: u32,
|
|
4394
|
+
kernel_size: u32,
|
|
4395
|
+
_pad: u32,
|
|
4396
|
+
}
|
|
4397
|
+
|
|
4398
|
+
@group(0) @binding(0) var<storage, read> input: array<f32>;
|
|
4399
|
+
@group(0) @binding(1) var<storage, read> weight: array<f32>;
|
|
4400
|
+
@group(0) @binding(2) var<storage, read_write> conv_state: array<f32>;
|
|
4401
|
+
@group(0) @binding(3) var<storage, read> gate: array<f32>;
|
|
4402
|
+
@group(0) @binding(4) var<storage, read_write> output: array<f32>;
|
|
4403
|
+
@group(0) @binding(5) var<storage, read> params: Params;
|
|
4404
|
+
|
|
4405
|
+
@compute @workgroup_size(256)
|
|
4406
|
+
fn main(@builtin(global_invocation_id) gid: vec3u) {
|
|
4407
|
+
let idx = gid.x;
|
|
4408
|
+
let total = params.seq_len * params.channels;
|
|
4409
|
+
if (idx >= total) { return; }
|
|
4410
|
+
|
|
4411
|
+
let t = idx / params.channels;
|
|
4412
|
+
let c = idx % params.channels;
|
|
4413
|
+
let state_size = params.kernel_size - 1u;
|
|
4414
|
+
|
|
4415
|
+
var olds: array<f32, 8>;
|
|
4416
|
+
for (var s: u32 = 0u; s < state_size; s = s + 1u) {
|
|
4417
|
+
olds[s] = conv_state[s * params.channels + c];
|
|
4418
|
+
}
|
|
4419
|
+
|
|
4420
|
+
var sum: f32 = 0.0;
|
|
4421
|
+
for (var k: u32 = 0u; k < params.kernel_size; k = k + 1u) {
|
|
4422
|
+
let src_t = i32(t) - i32(k);
|
|
4423
|
+
var in_val: f32 = 0.0;
|
|
4424
|
+
if (src_t >= 0) {
|
|
4425
|
+
in_val = input[u32(src_t) * params.channels + c];
|
|
4426
|
+
} else {
|
|
4427
|
+
let state_t = i32(state_size) + src_t;
|
|
4428
|
+
if (state_t >= 0) {
|
|
4429
|
+
in_val = olds[u32(state_t)];
|
|
4430
|
+
}
|
|
4431
|
+
}
|
|
4432
|
+
let w_val = weight[c * params.kernel_size + (params.kernel_size - 1u - k)];
|
|
4433
|
+
sum += in_val * w_val;
|
|
4434
|
+
}
|
|
4435
|
+
|
|
4436
|
+
output[t * params.channels + c] = gate[idx] * sum;
|
|
4437
|
+
|
|
4438
|
+
if (params.seq_len == 1u && t == 0u) {
|
|
4439
|
+
for (var s: u32 = 0u; s + 1u < state_size; s = s + 1u) {
|
|
4440
|
+
conv_state[s * params.channels + c] = olds[s + 1u];
|
|
4441
|
+
}
|
|
4442
|
+
if (state_size > 0u) {
|
|
4443
|
+
conv_state[(state_size - 1u) * params.channels + c] = input[c];
|
|
4444
|
+
}
|
|
4445
|
+
}
|
|
4446
|
+
}
|
|
4447
|
+
`;
|
|
4448
|
+
const WGSL_CAUSAL_CONV1D_GATED_BCX = `\
|
|
4449
|
+
struct Params {
|
|
4450
|
+
seq_len: u32,
|
|
4451
|
+
channels: u32,
|
|
4452
|
+
kernel_size: u32,
|
|
4453
|
+
in_width: u32, // bcx row stride = 3 * channels
|
|
4454
|
+
}
|
|
4455
|
+
|
|
4456
|
+
@group(0) @binding(0) var<storage, read> bcx: array<f32>;
|
|
4457
|
+
@group(0) @binding(1) var<storage, read> weight: array<f32>;
|
|
4458
|
+
@group(0) @binding(2) var<storage, read_write> conv_state: array<f32>;
|
|
4459
|
+
@group(0) @binding(3) var<storage, read_write> output: array<f32>;
|
|
4460
|
+
@group(0) @binding(4) var<storage, read> params: Params;
|
|
4461
|
+
|
|
4462
|
+
fn bx_at(t: u32, c: u32) -> f32 {
|
|
4463
|
+
let row = t * params.in_width;
|
|
4464
|
+
return bcx[row + c] * bcx[row + 2u * params.channels + c];
|
|
4465
|
+
}
|
|
4466
|
+
|
|
4467
|
+
@compute @workgroup_size(256)
|
|
4468
|
+
fn main(@builtin(global_invocation_id) gid: vec3u) {
|
|
4469
|
+
let idx = gid.x;
|
|
4470
|
+
let total = params.seq_len * params.channels;
|
|
4471
|
+
if (idx >= total) { return; }
|
|
4472
|
+
|
|
4473
|
+
let t = idx / params.channels;
|
|
4474
|
+
let c = idx % params.channels;
|
|
4475
|
+
let state_size = params.kernel_size - 1u;
|
|
4476
|
+
|
|
4477
|
+
// Snapshot this channel's rolling state (kernel_size <= 8) before any write.
|
|
4478
|
+
var olds: array<f32, 8>;
|
|
4479
|
+
for (var s: u32 = 0u; s < state_size; s = s + 1u) {
|
|
4480
|
+
olds[s] = conv_state[s * params.channels + c];
|
|
4481
|
+
}
|
|
4482
|
+
|
|
4483
|
+
var sum: f32 = 0.0;
|
|
4484
|
+
for (var k: u32 = 0u; k < params.kernel_size; k = k + 1u) {
|
|
4485
|
+
let src_t = i32(t) - i32(k);
|
|
4486
|
+
var in_val: f32 = 0.0;
|
|
4487
|
+
if (src_t >= 0) {
|
|
4488
|
+
in_val = bx_at(u32(src_t), c);
|
|
4489
|
+
} else {
|
|
4490
|
+
let state_t = i32(state_size) + src_t;
|
|
4491
|
+
if (state_t >= 0) {
|
|
4492
|
+
in_val = olds[u32(state_t)];
|
|
4493
|
+
}
|
|
4494
|
+
}
|
|
4495
|
+
let w_val = weight[c * params.kernel_size + (params.kernel_size - 1u - k)];
|
|
4496
|
+
sum += in_val * w_val;
|
|
4497
|
+
}
|
|
4498
|
+
|
|
4499
|
+
// Fused multiplicative gate from the C block.
|
|
4500
|
+
output[t * params.channels + c] = bcx[t * params.in_width + params.channels + c] * sum;
|
|
4501
|
+
|
|
4502
|
+
// Rolling-state shift (decode T=1 only; prefill uses ConvStateUpdateBcx).
|
|
4503
|
+
if (params.seq_len == 1u && t == 0u) {
|
|
4504
|
+
for (var s: u32 = 0u; s + 1u < state_size; s = s + 1u) {
|
|
4505
|
+
conv_state[s * params.channels + c] = olds[s + 1u];
|
|
4506
|
+
}
|
|
4507
|
+
if (state_size > 0u) {
|
|
4508
|
+
conv_state[(state_size - 1u) * params.channels + c] = bx_at(0u, c);
|
|
4509
|
+
}
|
|
4510
|
+
}
|
|
4511
|
+
}
|
|
4512
|
+
`;
|
|
4513
|
+
const WGSL_CONV_STATE_UPDATE_BCX = `\
|
|
4514
|
+
// ConvStateUpdate reading raw bcx (computes B*x for the copied timesteps).
|
|
4515
|
+
struct Params {
|
|
4516
|
+
seq_len: u32,
|
|
4517
|
+
channels: u32,
|
|
4518
|
+
state_size: u32,
|
|
4519
|
+
in_width: u32,
|
|
4520
|
+
}
|
|
4521
|
+
|
|
4522
|
+
@group(0) @binding(0) var<storage, read> bcx: array<f32>;
|
|
4523
|
+
@group(0) @binding(1) var<storage, read_write> conv_state: array<f32>;
|
|
4524
|
+
@group(0) @binding(2) var<storage, read> params: Params;
|
|
4525
|
+
|
|
4526
|
+
fn bx_at(t: u32, c: u32) -> f32 {
|
|
4527
|
+
let row = t * params.in_width;
|
|
4528
|
+
return bcx[row + c] * bcx[row + 2u * params.channels + c];
|
|
4529
|
+
}
|
|
4530
|
+
|
|
4531
|
+
@compute @workgroup_size(256)
|
|
4532
|
+
fn main(@builtin(global_invocation_id) gid: vec3u) {
|
|
4533
|
+
let c = gid.x;
|
|
4534
|
+
if (c >= params.channels) { return; }
|
|
4535
|
+
let ss = params.state_size;
|
|
4536
|
+
|
|
4537
|
+
if (params.seq_len >= ss) {
|
|
4538
|
+
for (var s: u32 = 0u; s < ss; s = s + 1u) {
|
|
4539
|
+
conv_state[s * params.channels + c] = bx_at(params.seq_len - ss + s, c);
|
|
4540
|
+
}
|
|
4541
|
+
} else {
|
|
4542
|
+
var olds: array<f32, 8>;
|
|
4543
|
+
for (var s: u32 = 0u; s < ss; s = s + 1u) {
|
|
4544
|
+
olds[s] = conv_state[s * params.channels + c];
|
|
4545
|
+
}
|
|
4546
|
+
let shift = ss - params.seq_len;
|
|
4547
|
+
for (var s: u32 = 0u; s < ss; s = s + 1u) {
|
|
4548
|
+
var val: f32 = 0.0;
|
|
4549
|
+
if (s < shift) {
|
|
4550
|
+
val = olds[s + params.seq_len];
|
|
4551
|
+
} else {
|
|
4552
|
+
val = bx_at(s - shift, c);
|
|
4553
|
+
}
|
|
4554
|
+
conv_state[s * params.channels + c] = val;
|
|
4555
|
+
}
|
|
4556
|
+
}
|
|
4557
|
+
}
|
|
4558
|
+
`;
|
|
4238
4559
|
const WGSL_CAUSAL_CONV1D_GATED = `\
|
|
4239
4560
|
// Causal 1D depthwise convolution with a fused multiplicative output gate:
|
|
4240
4561
|
// output = gate * conv(input)
|
|
@@ -6461,6 +6782,62 @@ const embeddingInt4Spec = {
|
|
|
6461
6782
|
]);
|
|
6462
6783
|
}
|
|
6463
6784
|
};
|
|
6785
|
+
const DUAL_RMSNORM_ROPE_SPEC = {
|
|
6786
|
+
shaderCode: WGSL_DUAL_RMSNORM_ROPE,
|
|
6787
|
+
entryPoint: "main",
|
|
6788
|
+
bindings: [
|
|
6789
|
+
{ type: "storage-read" },
|
|
6790
|
+
{ type: "storage-read" },
|
|
6791
|
+
{ type: "storage-read" },
|
|
6792
|
+
{ type: "storage-read" },
|
|
6793
|
+
{ type: "storage-read-write" },
|
|
6794
|
+
{ type: "storage-read-write" },
|
|
6795
|
+
{ type: "uniform" }
|
|
6796
|
+
],
|
|
6797
|
+
getDispatchSize(op, resolvedShapes) {
|
|
6798
|
+
const num_q_heads = op.attributes.num_q_heads;
|
|
6799
|
+
const num_kv_heads = op.attributes.num_kv_heads;
|
|
6800
|
+
const seqRef = op.attributes.seq_len_tensor;
|
|
6801
|
+
let T;
|
|
6802
|
+
if (seqRef && resolvedShapes[seqRef]) T = resolvedShapes[seqRef][0];
|
|
6803
|
+
else T = op.attributes.seq_len ?? 1;
|
|
6804
|
+
return [
|
|
6805
|
+
T * (num_q_heads + num_kv_heads),
|
|
6806
|
+
1,
|
|
6807
|
+
1
|
|
6808
|
+
];
|
|
6809
|
+
},
|
|
6810
|
+
buildParams(op, resolvedShapes, context) {
|
|
6811
|
+
const head_dim = op.attributes.hidden_size;
|
|
6812
|
+
const eps = op.attributes.eps ?? 1e-6;
|
|
6813
|
+
const num_q_heads = op.attributes.num_q_heads;
|
|
6814
|
+
const num_kv_heads = op.attributes.num_kv_heads;
|
|
6815
|
+
const rope_base = op.attributes.rope_base;
|
|
6816
|
+
const rope_dim = op.attributes.rope_dim ?? head_dim;
|
|
6817
|
+
const rope_half = op.attributes.rope_half ?? rope_dim / 2;
|
|
6818
|
+
const rope_denom = op.attributes.rope_denom ?? rope_dim;
|
|
6819
|
+
const rope_active_pairs = op.attributes.rope_active_pairs ?? rope_dim / 2;
|
|
6820
|
+
const position_offset = context?.seqPos ?? op.attributes.position_offset ?? 0;
|
|
6821
|
+
const seqRef = op.attributes.seq_len_tensor;
|
|
6822
|
+
let T;
|
|
6823
|
+
if (seqRef && resolvedShapes[seqRef]) T = resolvedShapes[seqRef][0];
|
|
6824
|
+
else T = op.attributes.seq_len ?? 1;
|
|
6825
|
+
return buildUniformBuffer([
|
|
6826
|
+
T * num_q_heads,
|
|
6827
|
+
T * num_kv_heads,
|
|
6828
|
+
head_dim,
|
|
6829
|
+
f32BitsToU32(eps),
|
|
6830
|
+
num_q_heads,
|
|
6831
|
+
num_kv_heads,
|
|
6832
|
+
f32BitsToU32(rope_base),
|
|
6833
|
+
position_offset,
|
|
6834
|
+
rope_half,
|
|
6835
|
+
rope_denom,
|
|
6836
|
+
rope_active_pairs,
|
|
6837
|
+
0
|
|
6838
|
+
]);
|
|
6839
|
+
}
|
|
6840
|
+
};
|
|
6464
6841
|
const ropeSpec = {
|
|
6465
6842
|
shaderCode: WGSL_ROPE,
|
|
6466
6843
|
entryPoint: "main",
|
|
@@ -6965,6 +7342,134 @@ const causalConv1dSiluSpec = {
|
|
|
6965
7342
|
]);
|
|
6966
7343
|
}
|
|
6967
7344
|
};
|
|
7345
|
+
const convUpdateDispatch = {
|
|
7346
|
+
getDispatchSize(op, resolvedShapes) {
|
|
7347
|
+
const channels = op.attributes.channels;
|
|
7348
|
+
const ref = op.attributes.seq_len_tensor;
|
|
7349
|
+
let seq_len;
|
|
7350
|
+
if (ref && resolvedShapes[ref]) seq_len = resolvedShapes[ref][0];
|
|
7351
|
+
else seq_len = op.attributes.seq_len;
|
|
7352
|
+
return [
|
|
7353
|
+
cdiv(seq_len * channels, 256),
|
|
7354
|
+
1,
|
|
7355
|
+
1
|
|
7356
|
+
];
|
|
7357
|
+
},
|
|
7358
|
+
buildParams(op, resolvedShapes) {
|
|
7359
|
+
const channels = op.attributes.channels;
|
|
7360
|
+
const kernel_size = op.attributes.kernel_size;
|
|
7361
|
+
const ref = op.attributes.seq_len_tensor;
|
|
7362
|
+
let seq_len;
|
|
7363
|
+
if (ref && resolvedShapes[ref]) seq_len = resolvedShapes[ref][0];
|
|
7364
|
+
else seq_len = op.attributes.seq_len;
|
|
7365
|
+
return buildUniformBuffer([
|
|
7366
|
+
seq_len,
|
|
7367
|
+
channels,
|
|
7368
|
+
kernel_size,
|
|
7369
|
+
0
|
|
7370
|
+
]);
|
|
7371
|
+
}
|
|
7372
|
+
};
|
|
7373
|
+
const CAUSAL_CONV1D_SILU_UPDATE_SPEC = {
|
|
7374
|
+
shaderCode: WGSL_CAUSAL_CONV1D_SILU_UPDATE,
|
|
7375
|
+
entryPoint: "main",
|
|
7376
|
+
bindings: [
|
|
7377
|
+
{ type: "storage-read" },
|
|
7378
|
+
{ type: "storage-read" },
|
|
7379
|
+
{ type: "storage-read-write" },
|
|
7380
|
+
{ type: "storage-read-write" },
|
|
7381
|
+
{ type: "uniform" }
|
|
7382
|
+
],
|
|
7383
|
+
getDispatchSize: convUpdateDispatch.getDispatchSize,
|
|
7384
|
+
buildParams: convUpdateDispatch.buildParams
|
|
7385
|
+
};
|
|
7386
|
+
const CAUSAL_CONV1D_GATED_UPDATE_SPEC = {
|
|
7387
|
+
shaderCode: WGSL_CAUSAL_CONV1D_GATED_UPDATE,
|
|
7388
|
+
entryPoint: "main",
|
|
7389
|
+
bindings: [
|
|
7390
|
+
{ type: "storage-read" },
|
|
7391
|
+
{ type: "storage-read" },
|
|
7392
|
+
{ type: "storage-read-write" },
|
|
7393
|
+
{ type: "storage-read" },
|
|
7394
|
+
{ type: "storage-read-write" },
|
|
7395
|
+
{ type: "uniform" }
|
|
7396
|
+
],
|
|
7397
|
+
getDispatchSize: convUpdateDispatch.getDispatchSize,
|
|
7398
|
+
buildParams: convUpdateDispatch.buildParams
|
|
7399
|
+
};
|
|
7400
|
+
const bcxDispatch = (op, resolvedShapes) => {
|
|
7401
|
+
const channels = op.attributes.channels;
|
|
7402
|
+
const ref = op.attributes.seq_len_tensor;
|
|
7403
|
+
let seq_len;
|
|
7404
|
+
if (ref && resolvedShapes[ref]) seq_len = resolvedShapes[ref][0];
|
|
7405
|
+
else seq_len = op.attributes.seq_len ?? 1;
|
|
7406
|
+
return {
|
|
7407
|
+
channels,
|
|
7408
|
+
seq_len
|
|
7409
|
+
};
|
|
7410
|
+
};
|
|
7411
|
+
const CAUSAL_CONV1D_GATED_BCX_SPEC = {
|
|
7412
|
+
shaderCode: WGSL_CAUSAL_CONV1D_GATED_BCX,
|
|
7413
|
+
entryPoint: "main",
|
|
7414
|
+
bindings: [
|
|
7415
|
+
{ type: "storage-read" },
|
|
7416
|
+
{ type: "storage-read" },
|
|
7417
|
+
{ type: "storage-read-write" },
|
|
7418
|
+
{ type: "storage-read-write" },
|
|
7419
|
+
{ type: "uniform" }
|
|
7420
|
+
],
|
|
7421
|
+
getDispatchSize(op, resolvedShapes) {
|
|
7422
|
+
const { channels, seq_len } = bcxDispatch(op, resolvedShapes);
|
|
7423
|
+
return [
|
|
7424
|
+
cdiv(seq_len * channels, 256),
|
|
7425
|
+
1,
|
|
7426
|
+
1
|
|
7427
|
+
];
|
|
7428
|
+
},
|
|
7429
|
+
buildParams(op, resolvedShapes) {
|
|
7430
|
+
const { channels, seq_len } = bcxDispatch(op, resolvedShapes);
|
|
7431
|
+
const kernel_size = op.attributes.kernel_size;
|
|
7432
|
+
const in_width = op.attributes.in_width;
|
|
7433
|
+
return buildUniformBuffer([
|
|
7434
|
+
seq_len,
|
|
7435
|
+
channels,
|
|
7436
|
+
kernel_size,
|
|
7437
|
+
in_width
|
|
7438
|
+
]);
|
|
7439
|
+
}
|
|
7440
|
+
};
|
|
7441
|
+
const CONV_STATE_UPDATE_BCX_SPEC = {
|
|
7442
|
+
shaderCode: WGSL_CONV_STATE_UPDATE_BCX,
|
|
7443
|
+
entryPoint: "main",
|
|
7444
|
+
bindings: [
|
|
7445
|
+
{ type: "storage-read" },
|
|
7446
|
+
{ type: "storage-read-write" },
|
|
7447
|
+
{ type: "uniform" }
|
|
7448
|
+
],
|
|
7449
|
+
getDispatchSize(op) {
|
|
7450
|
+
const channels = op.attributes.channels;
|
|
7451
|
+
return [
|
|
7452
|
+
cdiv(channels, 256),
|
|
7453
|
+
1,
|
|
7454
|
+
1
|
|
7455
|
+
];
|
|
7456
|
+
},
|
|
7457
|
+
buildParams(op, resolvedShapes) {
|
|
7458
|
+
const channels = op.attributes.channels;
|
|
7459
|
+
const state_size = op.attributes.state_size;
|
|
7460
|
+
const in_width = op.attributes.in_width;
|
|
7461
|
+
const ref = op.attributes.T_tensor;
|
|
7462
|
+
let T;
|
|
7463
|
+
if (ref && resolvedShapes[ref]) T = resolvedShapes[ref][0];
|
|
7464
|
+
else T = op.attributes.T ?? 1;
|
|
7465
|
+
return buildUniformBuffer([
|
|
7466
|
+
T,
|
|
7467
|
+
channels,
|
|
7468
|
+
state_size,
|
|
7469
|
+
in_width
|
|
7470
|
+
]);
|
|
7471
|
+
}
|
|
7472
|
+
};
|
|
6968
7473
|
const causalConv1dGatedSpec = {
|
|
6969
7474
|
shaderCode: WGSL_CAUSAL_CONV1D_GATED,
|
|
6970
7475
|
entryPoint: "main",
|
|
@@ -8124,6 +8629,8 @@ const KERNEL_REGISTRY = {
|
|
|
8124
8629
|
CausalConv1d: causalConv1dSpec,
|
|
8125
8630
|
CausalConv1dSiLU: causalConv1dSiluSpec,
|
|
8126
8631
|
CausalConv1dGated: causalConv1dGatedSpec,
|
|
8632
|
+
CausalConv1dGatedBcx: CAUSAL_CONV1D_GATED_BCX_SPEC,
|
|
8633
|
+
ConvStateUpdateBcx: CONV_STATE_UPDATE_BCX_SPEC,
|
|
8127
8634
|
SigmoidGate: sigmoidGateSpec,
|
|
8128
8635
|
MambaSSM: mambaSSMSpec,
|
|
8129
8636
|
KVCacheAppend: kvCacheAppendSpec,
|
|
@@ -8406,7 +8913,10 @@ var Executor = class Executor {
|
|
|
8406
8913
|
this.fuseDualKVCacheAppendEntries();
|
|
8407
8914
|
this.fuseGatedOProjDecodeEntries();
|
|
8408
8915
|
this.fuseSwiGLUGatedProjDecodeEntries();
|
|
8916
|
+
this.fuseDualRMSNormRoPEDecodeEntries();
|
|
8409
8917
|
this.fuseDualRMSNormDecodeEntries();
|
|
8918
|
+
this.fuseConvStateUpdateDecodeEntries();
|
|
8919
|
+
this.fuseAddRMSNormDecodeEntries();
|
|
8410
8920
|
console.log(`[executor] decode: ${this.decodeEntries.length} dispatches/token`);
|
|
8411
8921
|
const logitsBuffer = this.getBuffer("logits");
|
|
8412
8922
|
if (logitsBuffer) {
|
|
@@ -9748,6 +10258,243 @@ var Executor = class Executor {
|
|
|
9748
10258
|
* WebKit risk: low — same single-workgroup reduction as the proven RMSNorm
|
|
9749
10259
|
* kernel, only the row→input routing is added.
|
|
9750
10260
|
*/
|
|
10261
|
+
/**
|
|
10262
|
+
* Fuse [CausalConv1dSiLU|CausalConv1dGated → ConvStateUpdate] into one
|
|
10263
|
+
* dispatch (decode entries only). At T=1 the conv read and the rolling-state
|
|
10264
|
+
* write for a channel live in the same thread, so the fused kernel is
|
|
10265
|
+
* race-free; prefill (T>1) keeps the two-dispatch path.
|
|
10266
|
+
*/
|
|
10267
|
+
fuseConvStateUpdateDecodeEntries() {
|
|
10268
|
+
let fusionCount = 0;
|
|
10269
|
+
for (let i = 0; i < this.decodeEntries.length - 1; i++) {
|
|
10270
|
+
const e1 = this.decodeEntries[i];
|
|
10271
|
+
const e2 = this.decodeEntries[i + 1];
|
|
10272
|
+
if (e1.node.opType === "CausalConv1dGatedBcx" && e2.node.opType === "ConvStateUpdateBcx" && e2.node.inputs[0] === e1.node.inputs[0] && e2.node.outputs[0] === e1.node.inputs[2] && e1.node.attributes.kernel_size <= 8) {
|
|
10273
|
+
this.decodeEntries.splice(i + 1, 1);
|
|
10274
|
+
fusionCount++;
|
|
10275
|
+
continue;
|
|
10276
|
+
}
|
|
10277
|
+
const isSilu = e1.node.opType === "CausalConv1dSiLU";
|
|
10278
|
+
const isGated = e1.node.opType === "CausalConv1dGated";
|
|
10279
|
+
if (!isSilu && !isGated) continue;
|
|
10280
|
+
if (e2.node.opType !== "ConvStateUpdate") continue;
|
|
10281
|
+
if (e2.node.inputs[0] !== e1.node.inputs[0]) continue;
|
|
10282
|
+
if (e2.node.outputs[0] !== e1.node.inputs[2]) continue;
|
|
10283
|
+
const kernelSize = e1.node.attributes.kernel_size;
|
|
10284
|
+
if (e2.node.attributes.state_size !== kernelSize - 1) continue;
|
|
10285
|
+
if (kernelSize > 8) continue;
|
|
10286
|
+
const inputBuf = this.getBuffer(e1.node.inputs[0]);
|
|
10287
|
+
const weightBuf = this.getBuffer(e1.node.inputs[1]);
|
|
10288
|
+
const stateBuf = this.getBuffer(e1.node.inputs[2]);
|
|
10289
|
+
const gateBuf = isGated ? this.getBuffer(e1.node.inputs[3]) : null;
|
|
10290
|
+
const outBuf = this.getBuffer(e1.node.outputs[0]);
|
|
10291
|
+
if (!inputBuf || !weightBuf || !stateBuf || !outBuf || isGated && !gateBuf) continue;
|
|
10292
|
+
if (outBuf === inputBuf || outBuf === stateBuf || stateBuf === inputBuf) continue;
|
|
10293
|
+
if (gateBuf && (gateBuf === outBuf || gateBuf === stateBuf)) continue;
|
|
10294
|
+
const fusedSpec = isSilu ? CAUSAL_CONV1D_SILU_UPDATE_SPEC : CAUSAL_CONV1D_GATED_UPDATE_SPEC;
|
|
10295
|
+
const fusedPipeline = getOrCreatePipeline(this.ctx, `conv_update_${e1.nodeId}`, fusedSpec.shaderCode, fusedSpec.entryPoint);
|
|
10296
|
+
const fusedUniformData = fusedSpec.buildParams(e1.node, {}, { seqPos: 0 });
|
|
10297
|
+
const fusedUniform = createUniformBuffer(this.ctx, `uniform_conv_update_${e1.nodeId}`, fusedUniformData);
|
|
10298
|
+
const bufferList = isSilu ? [
|
|
10299
|
+
{ buffer: inputBuf },
|
|
10300
|
+
{ buffer: weightBuf },
|
|
10301
|
+
{ buffer: stateBuf },
|
|
10302
|
+
{ buffer: outBuf },
|
|
10303
|
+
{ buffer: fusedUniform }
|
|
10304
|
+
] : [
|
|
10305
|
+
{ buffer: inputBuf },
|
|
10306
|
+
{ buffer: weightBuf },
|
|
10307
|
+
{ buffer: stateBuf },
|
|
10308
|
+
{ buffer: gateBuf },
|
|
10309
|
+
{ buffer: outBuf },
|
|
10310
|
+
{ buffer: fusedUniform }
|
|
10311
|
+
];
|
|
10312
|
+
const fusedBindGroup = createBindGroup(this.ctx, fusedPipeline, bufferList, `bg_conv_update_${e1.nodeId}`);
|
|
10313
|
+
const fusedEntry = {
|
|
10314
|
+
nodeId: `fused_conv_update_${e1.nodeId}`,
|
|
10315
|
+
node: e1.node,
|
|
10316
|
+
spec: fusedSpec,
|
|
10317
|
+
pipeline: fusedPipeline,
|
|
10318
|
+
bindGroup: fusedBindGroup,
|
|
10319
|
+
uniformBuffer: fusedUniform,
|
|
10320
|
+
lastParamsBytes: null,
|
|
10321
|
+
lastDispatchSize: null,
|
|
10322
|
+
buffers: bufferList
|
|
10323
|
+
};
|
|
10324
|
+
this.decodeEntries.splice(i, 2, fusedEntry);
|
|
10325
|
+
fusionCount++;
|
|
10326
|
+
}
|
|
10327
|
+
if (fusionCount > 0) console.log(`[executor] Fused ${fusionCount} conv+state-update pairs (saved ${fusionCount} dispatches)`);
|
|
10328
|
+
}
|
|
10329
|
+
/**
|
|
10330
|
+
* Fuse [per-head q_norm RMSNorm, k_norm RMSNorm, in-place RoPE] into ONE
|
|
10331
|
+
* DualRMSNormRoPE dispatch. Numerically identical (norm then rotate, same
|
|
10332
|
+
* math); requires head_dim <= 256 and RoPE operating in-place on exactly the
|
|
10333
|
+
* two norm outputs. Saves 2 dispatches per full-attention layer — a mobile
|
|
10334
|
+
* (WebKit submit-per-dispatch) win, desktop-neutral.
|
|
10335
|
+
*/
|
|
10336
|
+
/**
|
|
10337
|
+
* Fuse [Add → RMSNorm] into a single ResidualRMSNorm dispatch: the add's sum
|
|
10338
|
+
* is still materialized (sum_output — downstream residual chain reads it) and
|
|
10339
|
+
* the norm result is written in the same dispatch. Numerically identical.
|
|
10340
|
+
*/
|
|
10341
|
+
fuseAddRMSNormDecodeEntries() {
|
|
10342
|
+
let fusionCount = 0;
|
|
10343
|
+
for (let i = 0; i < this.decodeEntries.length - 1; i++) {
|
|
10344
|
+
const e0 = this.decodeEntries[i];
|
|
10345
|
+
const e1 = this.decodeEntries[i + 1];
|
|
10346
|
+
if (e0.node.opType !== "Add" || e1.node.opType !== "RMSNorm") continue;
|
|
10347
|
+
if (e1.node.inputs[0] !== e0.node.outputs[0]) continue;
|
|
10348
|
+
const a = e0.node.inputs[0];
|
|
10349
|
+
const b = e0.node.inputs[1];
|
|
10350
|
+
const sumOut = e0.node.outputs[0];
|
|
10351
|
+
const normW = e1.node.inputs[1];
|
|
10352
|
+
const normOut = e1.node.outputs[0];
|
|
10353
|
+
const aBuf = this.getBuffer(a);
|
|
10354
|
+
const bBuf = this.getBuffer(b);
|
|
10355
|
+
const sumBuf = this.getBuffer(sumOut);
|
|
10356
|
+
const wBuf = this.getBuffer(normW);
|
|
10357
|
+
const outBuf = this.getBuffer(normOut);
|
|
10358
|
+
if (!aBuf || !bBuf || !sumBuf || !wBuf || !outBuf) continue;
|
|
10359
|
+
const bufs = [
|
|
10360
|
+
aBuf,
|
|
10361
|
+
bBuf,
|
|
10362
|
+
wBuf,
|
|
10363
|
+
sumBuf,
|
|
10364
|
+
outBuf
|
|
10365
|
+
];
|
|
10366
|
+
if (new Set(bufs).size !== bufs.length) continue;
|
|
10367
|
+
const fusedNode = {
|
|
10368
|
+
id: `fused_addnorm_${e0.nodeId}`,
|
|
10369
|
+
opType: "ResidualRMSNorm",
|
|
10370
|
+
inputs: [
|
|
10371
|
+
a,
|
|
10372
|
+
b,
|
|
10373
|
+
normW
|
|
10374
|
+
],
|
|
10375
|
+
outputs: [sumOut, normOut],
|
|
10376
|
+
attributes: {
|
|
10377
|
+
hidden_size: e1.node.attributes.hidden_size,
|
|
10378
|
+
eps: e1.node.attributes.eps,
|
|
10379
|
+
seq_len_tensor: e1.node.attributes.seq_len_tensor ?? e0.node.attributes.count_tensor
|
|
10380
|
+
}
|
|
10381
|
+
};
|
|
10382
|
+
const fusedSpec = KERNEL_REGISTRY.ResidualRMSNorm;
|
|
10383
|
+
if (!fusedSpec) continue;
|
|
10384
|
+
const fusedPipeline = getOrCreatePipeline(this.ctx, `add_norm_${e0.nodeId}`, fusedSpec.shaderCode, fusedSpec.entryPoint);
|
|
10385
|
+
const dummyShapes = this.resolveShapes(1);
|
|
10386
|
+
const fusedUniformData = fusedSpec.buildParams(fusedNode, dummyShapes, { seqPos: 0 });
|
|
10387
|
+
const fusedUniform = createUniformBuffer(this.ctx, `uniform_add_norm_${e0.nodeId}`, fusedUniformData);
|
|
10388
|
+
const bufferEntries = [
|
|
10389
|
+
{ buffer: aBuf },
|
|
10390
|
+
{ buffer: bBuf },
|
|
10391
|
+
{ buffer: wBuf },
|
|
10392
|
+
{ buffer: sumBuf },
|
|
10393
|
+
{ buffer: outBuf },
|
|
10394
|
+
{ buffer: fusedUniform }
|
|
10395
|
+
];
|
|
10396
|
+
const fusedBindGroup = createBindGroup(this.ctx, fusedPipeline, bufferEntries, `bg_add_norm_${e0.nodeId}`);
|
|
10397
|
+
this.decodeEntries.splice(i, 2, {
|
|
10398
|
+
nodeId: fusedNode.id,
|
|
10399
|
+
node: fusedNode,
|
|
10400
|
+
spec: fusedSpec,
|
|
10401
|
+
pipeline: fusedPipeline,
|
|
10402
|
+
bindGroup: fusedBindGroup,
|
|
10403
|
+
uniformBuffer: fusedUniform,
|
|
10404
|
+
lastParamsBytes: null,
|
|
10405
|
+
lastDispatchSize: null,
|
|
10406
|
+
buffers: bufferEntries
|
|
10407
|
+
});
|
|
10408
|
+
fusionCount++;
|
|
10409
|
+
}
|
|
10410
|
+
if (fusionCount > 0) console.log(`[executor] Fused ${fusionCount} add+RMSNorm pairs (saved ${fusionCount} dispatches)`);
|
|
10411
|
+
}
|
|
10412
|
+
fuseDualRMSNormRoPEDecodeEntries() {
|
|
10413
|
+
let fusionCount = 0;
|
|
10414
|
+
for (let i = 0; i < this.decodeEntries.length - 2; i++) {
|
|
10415
|
+
const e0 = this.decodeEntries[i];
|
|
10416
|
+
const e1 = this.decodeEntries[i + 1];
|
|
10417
|
+
const e2 = this.decodeEntries[i + 2];
|
|
10418
|
+
if (e0.node.opType !== "RMSNorm" || e1.node.opType !== "RMSNorm") continue;
|
|
10419
|
+
if (e2.node.opType !== "RoPE") continue;
|
|
10420
|
+
if (e2.node.attributes.interleaved === true) continue;
|
|
10421
|
+
if (e0.node.attributes.hidden_size !== e1.node.attributes.hidden_size) continue;
|
|
10422
|
+
if (e0.node.attributes.eps !== e1.node.attributes.eps) continue;
|
|
10423
|
+
const headDim = e0.node.attributes.hidden_size;
|
|
10424
|
+
if (headDim > 256) continue;
|
|
10425
|
+
if (e2.node.attributes.head_dim !== headDim) continue;
|
|
10426
|
+
const out0 = e0.node.outputs[0];
|
|
10427
|
+
const out1 = e1.node.outputs[0];
|
|
10428
|
+
if (e2.node.inputs[0] !== out0 || e2.node.inputs[1] !== out1) continue;
|
|
10429
|
+
if (e2.node.outputs[0] !== out0 || e2.node.outputs[1] !== out1) continue;
|
|
10430
|
+
const in0 = e0.node.inputs[0];
|
|
10431
|
+
const w0 = e0.node.inputs[1];
|
|
10432
|
+
const in1 = e1.node.inputs[0];
|
|
10433
|
+
const w1 = e1.node.inputs[1];
|
|
10434
|
+
const in0Buf = this.getBuffer(in0);
|
|
10435
|
+
const w0Buf = this.getBuffer(w0);
|
|
10436
|
+
const out0Buf = this.getBuffer(out0);
|
|
10437
|
+
const in1Buf = this.getBuffer(in1);
|
|
10438
|
+
const w1Buf = this.getBuffer(w1);
|
|
10439
|
+
const out1Buf = this.getBuffer(out1);
|
|
10440
|
+
if (!in0Buf || !w0Buf || !out0Buf || !in1Buf || !w1Buf || !out1Buf) continue;
|
|
10441
|
+
if (out0Buf === out1Buf) continue;
|
|
10442
|
+
if (in0Buf === out1Buf || in1Buf === out0Buf) continue;
|
|
10443
|
+
const fusedNode = {
|
|
10444
|
+
id: `fused_dualnorm_rope_${e0.nodeId}`,
|
|
10445
|
+
opType: "RMSNorm",
|
|
10446
|
+
inputs: [
|
|
10447
|
+
in0,
|
|
10448
|
+
w0,
|
|
10449
|
+
in1,
|
|
10450
|
+
w1
|
|
10451
|
+
],
|
|
10452
|
+
outputs: [out0, out1],
|
|
10453
|
+
attributes: {
|
|
10454
|
+
hidden_size: headDim,
|
|
10455
|
+
eps: e0.node.attributes.eps,
|
|
10456
|
+
seq_len_tensor: e0.node.attributes.seq_len_tensor,
|
|
10457
|
+
seq_len: e0.node.attributes.seq_len,
|
|
10458
|
+
num_q_heads: e2.node.attributes.num_q_heads,
|
|
10459
|
+
num_kv_heads: e2.node.attributes.num_kv_heads,
|
|
10460
|
+
rope_base: e2.node.attributes.rope_base,
|
|
10461
|
+
rope_dim: e2.node.attributes.rope_dim,
|
|
10462
|
+
rope_half: e2.node.attributes.rope_half,
|
|
10463
|
+
rope_denom: e2.node.attributes.rope_denom,
|
|
10464
|
+
rope_active_pairs: e2.node.attributes.rope_active_pairs
|
|
10465
|
+
}
|
|
10466
|
+
};
|
|
10467
|
+
const fusedSpec = DUAL_RMSNORM_ROPE_SPEC;
|
|
10468
|
+
const fusedPipeline = getOrCreatePipeline(this.ctx, `dual_norm_rope_${e0.nodeId}`, fusedSpec.shaderCode, fusedSpec.entryPoint);
|
|
10469
|
+
const dummyShapes = this.resolveShapes(1);
|
|
10470
|
+
const fusedUniformData = fusedSpec.buildParams(fusedNode, dummyShapes, { seqPos: 0 });
|
|
10471
|
+
const fusedUniform = createUniformBuffer(this.ctx, `uniform_dual_norm_rope_${e0.nodeId}`, fusedUniformData);
|
|
10472
|
+
const bufferEntries = [
|
|
10473
|
+
{ buffer: in0Buf },
|
|
10474
|
+
{ buffer: w0Buf },
|
|
10475
|
+
{ buffer: in1Buf },
|
|
10476
|
+
{ buffer: w1Buf },
|
|
10477
|
+
{ buffer: out0Buf },
|
|
10478
|
+
{ buffer: out1Buf },
|
|
10479
|
+
{ buffer: fusedUniform }
|
|
10480
|
+
];
|
|
10481
|
+
const fusedBindGroup = createBindGroup(this.ctx, fusedPipeline, bufferEntries, `bg_dual_norm_rope_${e0.nodeId}`);
|
|
10482
|
+
const fusedEntry = {
|
|
10483
|
+
nodeId: fusedNode.id,
|
|
10484
|
+
node: fusedNode,
|
|
10485
|
+
spec: fusedSpec,
|
|
10486
|
+
pipeline: fusedPipeline,
|
|
10487
|
+
bindGroup: fusedBindGroup,
|
|
10488
|
+
uniformBuffer: fusedUniform,
|
|
10489
|
+
lastParamsBytes: null,
|
|
10490
|
+
lastDispatchSize: null,
|
|
10491
|
+
buffers: bufferEntries
|
|
10492
|
+
};
|
|
10493
|
+
this.decodeEntries.splice(i, 3, fusedEntry);
|
|
10494
|
+
fusionCount++;
|
|
10495
|
+
}
|
|
10496
|
+
if (fusionCount > 0) console.log(`[executor] Fused ${fusionCount} norm+norm+RoPE triples (saved ${fusionCount * 2} dispatches)`);
|
|
10497
|
+
}
|
|
9751
10498
|
fuseDualRMSNormDecodeEntries() {
|
|
9752
10499
|
let fusionCount = 0;
|
|
9753
10500
|
for (let i = 0; i < this.decodeEntries.length - 1; i++) {
|
|
@@ -13425,4 +14172,4 @@ function selectGraphWeights(graph, weights) {
|
|
|
13425
14172
|
|
|
13426
14173
|
//#endregion
|
|
13427
14174
|
export { verifyGPU as S, createStorageBuffer as _, loadMoonshine as a, getOrCreatePipeline as b, quantizeBackboneInt4 as c, Tokenizer as d, Executor as f, createBindGroup as g, clearPipelineCache as h, loadModel as i, quantizeKaniBackbone as l, MATMUL_BIAS_F16C_SPEC as m, MoonshineEncoderExecutor as n, loadOuteTTS as o, KERNEL_REGISTRY as p, loadKaniTTS as r, loadParlerTTS as s, MoonshineSTT as t, remapPrunedToken as u, createUniformBuffer as v, initGPU as x, destroyBuffers as y };
|
|
13428
|
-
//# sourceMappingURL=moonshine-stt-
|
|
14175
|
+
//# sourceMappingURL=moonshine-stt-BwX7NQ0O.mjs.map
|