gpu: three of the four remaining backend-layer gaps - #4
Merged
Merged
Conversation
…ines Every other fused op already had a gpu::generic composition to fall back to; these two eager training ops still dropped straight to the host -- adam_step's own comment even said so. Both now compose from the same tier-0 primitives the rest of gpu::generic uses (broadcast subtract and exp plus a one-hot scatter for xent_bwd; a handful of affines, adds and a divide reused through three scratch buffers for adam_step), so a backend with no fused kernel keeps the update on the device instead of paying a round trip. No launch in either composition reads and writes the same buffer, the WebGPU constraint the rest of gpu::generic already works around. adam_step's chain folds an affine into each of its producing mul and sqrt launches (14 launches down to 12); the algebra is exact under the shared op's *scale+offset epilogue, not an approximation. Verified against the fused kernel on the host reference backend and native Metal, under WebGPU via Deno (the one backend that actually lacks both kernels, so this is the composition genuinely running), and by a clean CUDA host-side trace diff (new test launches only, no existing kernel call changed). Mutations of both compositions were each caught by the existing parity test.
kop carried two families no other backend implemented -- Metal's ten GEMM tile ids and its eighteen attention-kernel variants (by D and whether the KV cache is bf16) -- kept there only because bind_/pso_ took a kop directly. Metal now caches any (id, name) pair through one pso_id_, so its own gtile and attn_op enums use it the same way kop does, offset into disjoint ranges (gtile negative, attn_op further negative) so the shared cache never confuses one family's pipeline for another's -- a static_assert now checks the gap directly rather than leaving it to a comment. CUDA's one f32 GEMM fallback used to borrow kop::sgemm32 as an arbitrary cache key -- its cached_(slot, name) helper already didn't need a kop at all, so it gets its own named slot instead. docs/backends.md folds the CUDA wave plan's design reasoning into the existing paragraph that already says a backend's own kernel choice stays with the backend, rather than repeating it as an open gap. Verified on native Metal (117 cases, 11640 assertions with real device access, plus check_qwen's greedy tokens against the numpy oracle) and gpu_host's no-GPU reference build, under WebGPU via Deno, and by a clean CUDA host-side trace diff (no existing kernel call changed). Two collision mutations (gtile's key losing its offset, then attn_op's) each hung the real device outright rather than silently misbehaving -- about as decisive a catch as a test gets.
cuda.h and webgpu.h each kept their own copy of the same shape: a struct pairing a host buffer with a device handle and a residency state, a map from the native handle to it, and a size-keyed free list alloc()/release() recycle through. Only the device handle's type differed (CUdeviceptr, wgpu::Buffer). gpu::mirror_table<Handle> (gpu_abi.h) holds that shape once; each backend instantiates it for its own handle type and keeps doing its own copying (before_kernel_, sync_to_host), which is where the two drivers actually diverge. mirror_table::insert moves its Handle parameter into the entry rather than copying it, so a ref-counted handle (wgpu::Buffer) picks up only one AddRef per allocation instead of two. Verified on native Metal (117 cases, 11640 assertions with real device access, plus check_qwen's greedy tokens against the numpy oracle) and gpu_host's no-GPU reference build, under WebGPU via Deno, and by a clean CUDA host-side trace diff against the pre-refactor commit (no kernel call changed). A mutation that let take() hand out an already-live pooled buffer twice was caught hard by the WebGPU suite (32 of 117 cases failed).
5 tasks
6 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Three of the four items left after the backend-layer redesign (PR #2, #3):
gpu::genericcomposition to fall back to; these two eager training ops still dropped straight to the host. Both now compose from the same tier-0 primitives the rest ofgpu::genericuses, so a backend with no fused kernel keeps the update on the device instead of paying a round trip.kopcarried two families no other backend implemented (Metal's ten GEMM tile ids, its eighteen attention-kernel variants), kept there only becausebind_/pso_took akopdirectly. Metal now caches any(id, name)pair through onepso_id_, so its owngtileandattn_openums use it the same waykopdoes, offset into disjoint ranges so the shared cache never confuses one family's pipeline for another's.gpu_abi.h.cuda.handwebgpu.heach kept their own copy of the same shape: a struct pairing a host buffer with a device handle and a residency state, a map from the native handle to it, and a size-keyed free listalloc()/release()recycle through.gpu::mirror_table<Handle>holds that shape once; each backend instantiates it for its own handle type and keeps doing its own copying, which is where the two drivers actually diverge.The CUDA f32 GEMM wave plan (originally the fourth item) turned out to have no Metal counterpart to share it with —
docs/backends.mdnow says so directly instead of listing it as an open gap.The remaining item — rewriting WebGPU's WGSL kernels to the canonical ABI and dropping
marshal_— is out of scope here; it is the largest and riskiest of the four and will be its own PR.Test plan
ctest— 117/117 cases, 11640 assertions with real device accesstensorlib_check_qwen— greedy tokens match the numpy oraclegpu_hostno-GPU reference build (-DTENSORLIB_HOST_GPU=ON) — 117/117test/wasm/build.sh tests && deno run --allow-all test/wasm/deno_run.js) — 117/117, both--gpuand--autotools/cuda_trace/compare.sh) — identical kernel calls before/after each change