fix(core): take every kernel uniform from cachedUniform so none outlives its executor

packages/runntime/src/core/kernels/attention/attn.ts

@@ -1,5 +1,6 @@
 import tgpu, { d, std } from 'typegpu';
 import type { TgpuRoot } from 'typegpu';
+import { cachedBindGroup, cachedUniform } from '../../gpu/dispatchCache.ts';
 import {
   flatIndex,
   type F32Buffer,
@@ -209,12 +210,14 @@ export function attnHandle(
   },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const dims = root
-    .createBuffer(Dims, { qLen: args.qLen, kvLen: args.kvLen, qPosOffset: args.qPosOffset })
-    .$usage('uniform');
+  const dims = cachedUniform(root, Dims, {
+    qLen: args.qLen,
+    kvLen: args.kvLen,
+    qPosOffset: args.qPosOffset,
+  });
   const filler = f32Filler(buffers, elem);
   const { dummyF32: _dummy, ...bound } = buffers;
-  const bindGroup = root.createBindGroup(attnLayoutFor(elem), {
+  const bindGroup = cachedBindGroup(root, attnLayoutFor(elem), {
     ...bound,
     sinks: buffers.sinks ?? filler,
     segs: buffers.segs ?? filler,
@@ -700,12 +703,14 @@ export function attnRowsHandle(
   },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const dims = root
-    .createBuffer(Dims, { qLen: args.qLen, kvLen: args.kvLen, qPosOffset: args.qPosOffset })
-    .$usage('uniform');
+  const dims = cachedUniform(root, Dims, {
+    qLen: args.qLen,
+    kvLen: args.kvLen,
+    qPosOffset: args.qPosOffset,
+  });
   const filler = f32Filler(buffers, elem);
   const { dummyF32: _dummy, ...bound } = buffers;
-  const bindGroup = root.createBindGroup(attnLayoutFor(elem), {
+  const bindGroup = cachedBindGroup(root, attnLayoutFor(elem), {
     ...bound,
     sinks: buffers.sinks ?? filler,
     segs: buffers.segs ?? filler,
@@ -746,15 +751,13 @@ export function attnRowsSplitHandles(
   },
   elem: Elem = F32_ELEM,
 ): KernelHandle[] {
-  const splitDims = root
-    .createBuffer(SplitDims, {
-      qLen: args.qLen,
-      kvLen: args.kvLen,
-      qPosOffset: args.qPosOffset,
-      chunks: args.chunks,
-    })
-    .$usage('uniform');
-  const splitBind = root.createBindGroup(attnSplitLayoutFor(elem), {
+  const splitDims = cachedUniform(root, SplitDims, {
+    qLen: args.qLen,
+    kvLen: args.kvLen,
+    qPosOffset: args.qPosOffset,
+    chunks: args.chunks,
+  });
+  const splitBind = cachedBindGroup(root, attnSplitLayoutFor(elem), {
     q: buffers.q,
     k: buffers.k,
     v: buffers.v,
@@ -763,10 +766,8 @@ export function attnRowsSplitHandles(
     partials: buffers.partials,
     dims: splitDims,
   });
-  const mergeDims = root
-    .createBuffer(MergeDims, { qLen: args.qLen, chunks: args.chunks })
-    .$usage('uniform');
-  const mergeBind = root.createBindGroup(attnMergeLayoutFor(elem), {
+  const mergeDims = cachedUniform(root, MergeDims, { qLen: args.qLen, chunks: args.chunks });
+  const mergeBind = cachedBindGroup(root, attnMergeLayoutFor(elem), {
     partials: buffers.partials,
     out: buffers.out,
     dims: mergeDims,

packages/runntime/src/core/kernels/attention/attnDecode.ts

@@ -1,5 +1,6 @@
 import tgpu, { d, std } from 'typegpu';
 import type { TgpuRoot } from 'typegpu';
+import { cachedBindGroup, cachedUniform } from '../../gpu/dispatchCache.ts';
 import {
   type F32Buffer,
   f32Filler,
@@ -207,14 +208,17 @@ export function attnDecodeHandle(
   },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const dims = root
-    .createBuffer(Dims, { qLen: args.qLen, kvLen: args.kvLen, qPosOffset: args.qPosOffset })
-    .$usage('uniform');
-  const filler = f32Filler(buffers, elem);
-  const { dummyF32: _dummy, ...bound } = buffers;
-  const bindGroup = root.createBindGroup(attnDecodeLayoutFor(elem), {
-    ...bound,
-    sinks: buffers.sinks ?? filler,
+  const dims = cachedUniform(root, Dims, {
+    qLen: args.qLen,
+    kvLen: args.kvLen,
+    qPosOffset: args.qPosOffset,
+  });
+  const bindGroup = cachedBindGroup(root, attnDecodeLayoutFor(elem), {
+    q: buffers.q,
+    k: buffers.k,
+    v: buffers.v,
+    out: buffers.out,
+    sinks: buffers.sinks ?? f32Filler(buffers, elem),
     dims,
   });
   // One workgroup per (q row, q head), not one thread.
@@ -451,23 +455,21 @@ export function attnDecodeSplitHandles(
   },
   elem: Elem = F32_ELEM,
 ): [KernelHandle, KernelHandle] {
-  const dims = root
-    .createBuffer(SplitDims, {
-      kvLen: args.kvLen,
-      qPosOffset: args.qPosOffset,
-      chunks: args.chunks,
-    })
-    .$usage('uniform');
-  const partialBind = root.createBindGroup(attnDecodeSplitLayoutFor(elem), {
+  const dims = cachedUniform(root, SplitDims, {
+    kvLen: args.kvLen,
+    qPosOffset: args.qPosOffset,
+    chunks: args.chunks,
+  });
+  const partialBind = cachedBindGroup(root, attnDecodeSplitLayoutFor(elem), {
     q: buffers.q,
     k: buffers.k,
     v: buffers.v,
     sinks: buffers.sinks ?? f32Filler(buffers, elem),
     partials: buffers.partials,
     dims,
   });
-  const mergeDims = root.createBuffer(MergeDims, { chunks: args.chunks }).$usage('uniform');
-  const mergeBind = root.createBindGroup(attnDecodeMergeLayoutFor(elem), {
+  const mergeDims = cachedUniform(root, MergeDims, { chunks: args.chunks });
+  const mergeBind = cachedBindGroup(root, attnDecodeMergeLayoutFor(elem), {
     partials: buffers.partials,
     out: buffers.out,
     dims: mergeDims,

packages/runntime/src/core/kernels/conv/conv1d.ts

@@ -1,5 +1,6 @@
 import tgpu, { d } from 'typegpu';
 import type { TgpuRoot } from 'typegpu';
+import { cachedBindGroup, cachedUniform } from '../../gpu/dispatchCache.ts';
 import {
   type FloatBuffer,
   flatIndex,
@@ -198,9 +199,9 @@ export function conv1dHandle(
   buffers: { x: FloatBuffer; w: FloatBuffer; out: FloatBuffer; bias?: FloatBuffer },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const dims = root.createBuffer(Dims, { tIn: args.tIn, tOut: args.tOut }).$usage('uniform');
+  const dims = cachedUniform(root, Dims, { tIn: args.tIn, tOut: args.tOut });
   // An absent bias aliases `x`; hasBias=0 pipelines never read it.
-  const bindGroup = root.createBindGroup(conv1dVariant(elem).layout, {
+  const bindGroup = cachedBindGroup(root, conv1dVariant(elem).layout, {
     ...buffers,
     bias: buffers.bias ?? buffers.x,
     dims,

packages/runntime/src/core/kernels/elementwise/swigluChunk.ts

@@ -1,5 +1,6 @@
 import tgpu, { d, std } from 'typegpu';
 import type { TgpuRoot } from 'typegpu';
+import { cachedBindGroup, cachedUniform } from '../../gpu/dispatchCache.ts';
 import {
   type FloatBuffer,
   flatIndex,
@@ -66,7 +67,7 @@ export function swigluChunkHandle(
   buffers: { x: FloatBuffer; out: FloatBuffer },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const dims = root.createBuffer(Dims, { total: args.total, half: args.half }).$usage('uniform');
-  const bindGroup = root.createBindGroup(swigluChunkVariant(elem).layout, { ...buffers, dims });
+  const dims = cachedUniform(root, Dims, { total: args.total, half: args.half });
+  const bindGroup = cachedBindGroup(root, swigluChunkVariant(elem).layout, { ...buffers, dims });
   return makeHandle(pipeline, 'swigluChunk', bindGroup, Math.ceil(args.total / WORKGROUP_SIZE));
 }

packages/runntime/src/core/kernels/matmul/argmaxDot.ts

@@ -1,5 +1,6 @@
 import tgpu, { d, std } from 'typegpu';
 import type { TgpuRoot } from 'typegpu';
+import { cachedBindGroup, cachedUniform } from '../../gpu/dispatchCache.ts';
 import { type F32Buffer, type FloatBuffer, type KernelHandle, makeHandle } from '../common.ts';
 import { type Elem, F32_ELEM } from '../elem.ts';
 
@@ -163,8 +164,8 @@ export function argmaxDotPartialHandle(
   buffers: { w: FloatBuffer; x: FloatBuffer; partials: F32Buffer },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const dims = root.createBuffer(Dims, { rows }).$usage('uniform');
-  const bindGroup = root.createBindGroup(argmaxDotVariant(elem).layout, { ...buffers, dims });
+  const dims = cachedUniform(root, Dims, { rows });
+  const bindGroup = cachedBindGroup(root, argmaxDotVariant(elem).layout, { ...buffers, dims });
   return makeHandle(pipeline, 'argmaxDot', bindGroup, ARGMAX_SPLIT);
 }
 

packages/runntime/src/core/kernels/reduce/layerNorm.ts

@@ -1,5 +1,6 @@
 import tgpu, { d, std } from 'typegpu';
 import type { TgpuRoot } from 'typegpu';
+import { cachedBindGroup, cachedUniform } from '../../gpu/dispatchCache.ts';
 import {
   type FloatBuffer,
   flatWorkgroupId,
@@ -140,10 +141,13 @@ export function layerNormHandle(
   buffers: { x: FloatBuffer; weight: FloatBuffer; bias: FloatBuffer; out: FloatBuffer },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const dims = root
-    .createBuffer(Dims, { rows: args.rows, cols: args.cols, eps: args.eps, hasBias: args.hasBias })
-    .$usage('uniform');
-  const bindGroup = root.createBindGroup(layerNormVariant(elem).layout, { ...buffers, dims });
+  const dims = cachedUniform(root, Dims, {
+    rows: args.rows,
+    cols: args.cols,
+    eps: args.eps,
+    hasBias: args.hasBias,
+  });
+  const bindGroup = cachedBindGroup(root, layerNormVariant(elem).layout, { ...buffers, dims });
   // One workgroup per row — not one thread.
   return makeHandle(pipeline, 'layerNorm', bindGroup, args.rows);
 }

packages/runntime/src/core/kernels/shape/rows.ts

@@ -243,8 +243,8 @@ export function gatherRowsFromHandle(
   buffers: { x: FloatBuffer; idx: F32Buffer; out: FloatBuffer },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const meta = root.createBuffer(GatherFromMeta, args).$usage('uniform');
-  const bindGroup = root.createBindGroup(gatherRowsFromVariant(elem).layout, { ...buffers, meta });
+  const meta = cachedUniform(root, GatherFromMeta, args);
+  const bindGroup = cachedBindGroup(root, gatherRowsFromVariant(elem).layout, { ...buffers, meta });
   return makeHandle(pipeline, 'gatherRowsFrom', bindGroup, Math.ceil(args.total / WORKGROUP_SIZE));
 }
 
@@ -255,8 +255,8 @@ export function writeRowsHandle(
   buffers: { src: FloatBuffer; dst: FloatBuffer },
   elem: Elem = F32_ELEM,
 ): KernelHandle {
-  const meta = root.createBuffer(WriteMeta, args).$usage('uniform');
-  const bindGroup = root.createBindGroup(writeRowsVariant(elem).layout, { ...buffers, meta });
+  const meta = cachedUniform(root, WriteMeta, args);
+  const bindGroup = cachedBindGroup(root, writeRowsVariant(elem).layout, { ...buffers, meta });
   return makeHandle(pipeline, 'writeRows', bindGroup, Math.ceil(args.total / WORKGROUP_SIZE));
 }
 

packages/runntime/test/core/gpu/eval.gpu.test.ts

@@ -72,8 +72,6 @@ interface GraphCase {
   leftAfterCapture?: string[];
 }
 
-/** A uniform buffer a kernel makes itself instead of through cachedUniform. */
-const uniform = (bytes: number) => `${bytes} B, usage 76`;
 /** The 3-row gatherRows index list. */
 const INDEX_LIST = '12 B, usage 140';
 
@@ -97,9 +95,6 @@ const GRAPHS: GraphCase[] = [
   {
     name: 'layerNorm without a bias',
     build: ({ x, w }) => layerNorm(add(x, 1), w),
-    // dispose() keeps the layerNorm uniform.
-    leftOnDispose: [uniform(16)],
-    leftAfterCapture: [uniform(16)],
   },
   {
     name: 'softmax',
@@ -128,54 +123,33 @@ const GRAPHS: GraphCase[] = [
     needs: 'subgroups',
     // rowsSubgroupSplit: qLen > 1, f32, subgroups, headDim % 16 != 0
     build: ({ x, m }) => sdpa(x, m, m, TWO_HEAD_ATTENTION),
-    // dispose() keeps the split and merge uniforms.
-    leftOnDispose: [uniform(16), uniform(8)],
-    leftAfterCapture: [uniform(16), uniform(8)],
   },
   {
     name: 'decode attention',
     // decode: qLen == 1, kvLen <= ATTN_DECODE_CHUNK
     build: ({ w, m }) => sdpa(w, m, m, FULL_ATTENTION),
-    // dispose() keeps the attnDecode uniform.
-    leftOnDispose: [uniform(12)],
-    leftAfterCapture: [uniform(12)],
   },
   {
     name: 'conv1d',
     build: ({ x, m }) => conv1d(x, m, { kernelSize: 1, stride: 1 }),
-    // dispose() keeps the conv1d uniform.
-    leftOnDispose: [uniform(8)],
-    leftAfterCapture: [uniform(8)],
   },
   {
     name: 'swigluChunk',
     build: ({ x }) => swigluChunk(add(x, 1)),
-    // dispose() keeps the swigluChunk uniform.
-    leftOnDispose: [uniform(8)],
-    leftAfterCapture: [uniform(8)],
   },
   {
     name: 'argmaxDot',
     build: ({ w, m }) => argmaxDot(m, w),
-    // dispose() keeps the argmaxDot uniform.
-    leftOnDispose: [uniform(4)],
-    leftAfterCapture: [uniform(4)],
   },
   {
     name: 'gatherRowsFrom',
     build: ({ x, idx }) => gatherRowsFrom(add(x, 1), idx),
-    // dispose() keeps the gatherRowsFrom uniform.
-    leftOnDispose: [uniform(8)],
-    leftAfterCapture: [uniform(8)],
   },
   {
     name: 'writeRows',
     // Writes into x in place, as a KV-cache append does, then reads it: a
     // released writeRows result would hand x's buffer to the pool.
     build: ({ x, w }) => tanh(writeRows(x, add(w, 1), 2)),
-    // dispose() keeps the writeRows uniform.
-    leftOnDispose: [uniform(8)],
-    leftAfterCapture: [uniform(8)],
   },
 ];
 

packages/runntime/test/zoo/leaks.gpu.test.ts

@@ -101,11 +101,6 @@ const CASES: LeakCase[] = [
         },
       };
     },
-    // dispose() keeps the kernels' uniforms.
-    leftOnDispose: [
-      ...Array<string>(2).fill('8 B, usage 76'),
-      ...Array<string>(5).fill('16 B, usage 76'),
-    ],
   },
   {
     name: 'MiniLM, a batch',
@@ -128,12 +123,7 @@ const CASES: LeakCase[] = [
         },
       };
     },
-    // dispose() also keeps the kernels' uniforms.
-    leftOnDispose: [
-      ...MINILM_BATCH_UPLOADS,
-      ...Array<string>(5).fill('16 B, usage 76'),
-      MINILM_BATCH_RESULT,
-    ],
+    leftOnDispose: [...MINILM_BATCH_UPLOADS, MINILM_BATCH_RESULT],
     leftPerCall: [...MINILM_BATCH_UPLOADS, MINILM_BATCH_RESULT],
   },
   {