From d98c764493a43ca92cc2ba4fb1c41d22997464df Mon Sep 17 00:00:00 2001 From: chmjkb Date: Fri, 23 Oct 2026 17:45:42 +0200 Subject: [PATCH 1/2] fix(core): take every kernel uniform from cachedUniform so none outlives its executor --- .../src/core/kernels/attention/attn.ts | 43 ++++++++++--------- .../src/core/kernels/attention/attnDecode.ts | 38 ++++++++-------- .../runntime/src/core/kernels/conv/conv1d.ts | 5 ++- .../core/kernels/elementwise/swigluChunk.ts | 5 ++- .../src/core/kernels/matmul/argmaxDot.ts | 5 ++- .../src/core/kernels/reduce/layerNorm.ts | 12 ++++-- .../runntime/src/core/kernels/shape/rows.ts | 8 ++-- 7 files changed, 63 insertions(+), 53 deletions(-) diff --git a/packages/runntime/src/core/kernels/attention/attn.ts b/packages/runntime/src/core/kernels/attention/attn.ts index c0adfbe..1c27577 100644 --- a/packages/runntime/src/core/kernels/attention/attn.ts +++ b/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, diff --git a/packages/runntime/src/core/kernels/attention/attnDecode.ts b/packages/runntime/src/core/kernels/attention/attnDecode.ts index 1dde5ba..e349837 100644 --- a/packages/runntime/src/core/kernels/attention/attnDecode.ts +++ b/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,14 +455,12 @@ 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, @@ -466,8 +468,8 @@ export function attnDecodeSplitHandles( 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, diff --git a/packages/runntime/src/core/kernels/conv/conv1d.ts b/packages/runntime/src/core/kernels/conv/conv1d.ts index 8c241e2..1b8c243 100644 --- a/packages/runntime/src/core/kernels/conv/conv1d.ts +++ b/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, diff --git a/packages/runntime/src/core/kernels/elementwise/swigluChunk.ts b/packages/runntime/src/core/kernels/elementwise/swigluChunk.ts index 48f4ed6..03fcdba 100644 --- a/packages/runntime/src/core/kernels/elementwise/swigluChunk.ts +++ b/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)); } diff --git a/packages/runntime/src/core/kernels/matmul/argmaxDot.ts b/packages/runntime/src/core/kernels/matmul/argmaxDot.ts index b6aab33..dae80a3 100644 --- a/packages/runntime/src/core/kernels/matmul/argmaxDot.ts +++ b/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); } diff --git a/packages/runntime/src/core/kernels/reduce/layerNorm.ts b/packages/runntime/src/core/kernels/reduce/layerNorm.ts index f8547b0..f664960 100644 --- a/packages/runntime/src/core/kernels/reduce/layerNorm.ts +++ b/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); } diff --git a/packages/runntime/src/core/kernels/shape/rows.ts b/packages/runntime/src/core/kernels/shape/rows.ts index afd99fb..ea60f1f 100644 --- a/packages/runntime/src/core/kernels/shape/rows.ts +++ b/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)); } From 13f36e117cbc6c1f92edb3d08055ae93322caa3b Mon Sep 17 00:00:00 2001 From: chmjkb Date: Fri, 23 Oct 2026 17:46:21 +0200 Subject: [PATCH 2/2] test(core): expect no kernel uniform to outlive its executor --- .../runntime/test/core/gpu/eval.gpu.test.ts | 26 ------------------- packages/runntime/test/zoo/leaks.gpu.test.ts | 12 +-------- 2 files changed, 1 insertion(+), 37 deletions(-) diff --git a/packages/runntime/test/core/gpu/eval.gpu.test.ts b/packages/runntime/test/core/gpu/eval.gpu.test.ts index ad3ab21..376975d 100644 --- a/packages/runntime/test/core/gpu/eval.gpu.test.ts +++ b/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)], }, ]; diff --git a/packages/runntime/test/zoo/leaks.gpu.test.ts b/packages/runntime/test/zoo/leaks.gpu.test.ts index 4220ce4..0337e61 100644 --- a/packages/runntime/test/zoo/leaks.gpu.test.ts +++ b/packages/runntime/test/zoo/leaks.gpu.test.ts @@ -101,11 +101,6 @@ const CASES: LeakCase[] = [ }, }; }, - // dispose() keeps the kernels' uniforms. - leftOnDispose: [ - ...Array(2).fill('8 B, usage 76'), - ...Array(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(5).fill('16 B, usage 76'), - MINILM_BATCH_RESULT, - ], + leftOnDispose: [...MINILM_BATCH_UPLOADS, MINILM_BATCH_RESULT], leftPerCall: [...MINILM_BATCH_UPLOADS, MINILM_BATCH_RESULT], }, {