Repository navigation
feat(cuda): PQ2_0 (PrismML Bonsai ternary) CPU + CUDA kernels - #153
Merged
Merged
Conversation
Adds the PQ2_0 GGUF quant type (type 42) for PrismML's Bonsai-27B ternary model: CPU dequant + scalar GEMV/GEMM kernels, dispatch wiring in TransformerModel.cs, and CUDA dequant + GEMV kernels wired into CudaTransformerModel.cs's Project() dispatch (prefill dequant-then-cuBLAS path and decode GEMV path, mirroring the simple I2_S sites). Unlike I2_S's single per-tensor scale, PQ2_0's scale is per-128-element group (empirically verified against real Ternary-Bonsai-27B-Q2_0.gguf bytes and confirmed a second time via a real-GGUF dequant smoke test this session). The CUDA GEMV kernel reads x directly from global memory rather than staging it in shared memory like I2_S does, since real Bonsai FFN rows (k=17408) would overflow sm_86's 48KB static shared-memory cap. CPU kernel validated against a real Bonsai tensor from the actual GGUF file; CUDA kernel validated against the CPU reference on the real RTX 3060 at real Bonsai dims (hidden=5120, ffn=17408). Correctness-first scope only — no AVX2/AVX-VNNI SIMD tier, no CUDA fused/W2A8 variants, matching I2_S's original scoping. Vulkan kernels are a follow-up (tracked in the issue). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Adds support for PrismML Bonsai’s ternary GGUF quantization PQ2_0 (type 42) across CPU + CUDA backends, including byte-size accounting, CPU dequant + GEMV/GEMM kernels, CUDA dequant + GEMV kernels, and unit/smoke/parity tests. This fits into dotLLM’s quantized-weight execution paths by extending QuantizationType and wiring new dispatch points in both the CPU TransformerModel and CUDA CudaTransformerModel.
Changes:
- Introduces
QuantizationType.PQ2_0 = 42with correct byte-count/row-stride formulas (group-128 with per-group fp16 scale). - Adds CPU scalar dequant + GEMV/GEMM implementation and hooks it into
TransformerModel’s GEMV/GEMM dispatch (and fused-kernel gating). - Adds CUDA PQ2_0 dequant-to-F16 (prefill) and PQ2_0 GEMV kernels (decode), plus CPU-vs-GPU parity and real-GGUF smoke tests.
Reviewed changes
Copilot reviewed 13 out of 14 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/DotLLM.Tests.Unit/Cuda/CudaPQ2_0GemvTest.cs | Adds CPU-vs-GPU PQ2_0 GEMV parity test at Bonsai-like dimensions and a synthetic packer helper |
| tests/DotLLM.Tests.Unit/Cpu/Kernels/PQ2_0Tests.cs | Adds unit tests for PQ2_0 layout, byte accounting, dequant, GEMV, and GEMM correctness |
| tests/DotLLM.Tests.Unit/Cpu/Kernels/PQ2_0RealGgufSmokeTests.cs | Adds gated real-GGUF smoke test to validate dequant output sanity on real Bonsai tensors |
| src/DotLLM.Models/Architectures/TransformerModel.cs | Wires PQ2_0 into CPU GEMV/GEMM dispatch and excludes ternary types from fused decode path |
| src/DotLLM.Cuda/CudaTransformerModel.cs | Wires PQ2_0 into CUDA projection dispatch for both prefill dequant+HGEMM and decode GEMV |
| src/DotLLM.Cuda/CudaKernels.cs | Loads PQ2_0 PTX modules and exposes launch wrappers for PQ2_0 GEMV + dequant-to-F16 |
| src/DotLLM.Cpu/Kernels/MatMul.PQ2S.cs | Adds CPU scalar PQ2_0 GEMV/GEMM kernels (unpack row to float then TensorPrimitives.Dot) |
| src/DotLLM.Cpu/Kernels/Dequantize.cs | Adds PQ2_0 row stride accounting and scalar PQ2_0 dequant routine |
| src/DotLLM.Core/Configuration/QuantizationTypeExtensions.cs | Adds PQ2_0 ComputeByteCount implementation |
| src/DotLLM.Core/Configuration/QuantizationType.cs | Adds PQ2_0 enum member + detailed format documentation |
| native/ptx/dequant_pq2_0.ptx | Adds compiled PTX for GPU PQ2_0 dequant-to-F16 |
| native/kernels/pq2_0_gemv.cu | Adds CUDA PQ2_0 GEMV kernels (f32-in parity + f16-in production) |
| native/kernels/dequant_pq2_0.cu | Adds CUDA PQ2_0 dequant-to-F16 kernel for prefill HGEMM |
Comment on lines
+1239
to
+1253
| /// <summary>PQ2_0 ternary GEMV with FP16 activations/output — the production decode path.</summary> | ||
| public void LaunchPQ2_0GemvF16In(nint quantWeight, nint xF16, nint yF16, int n, int k, nint stream) | ||
| { | ||
| nint wArg = quantWeight, xArg = xF16, yArg = yF16; | ||
| int nArg = n, kArg = k; | ||
| void** args = stackalloc void*[] {&wArg, &xArg, &yArg, &nArg, &kArg}; | ||
|
|
||
| int warpsPerBlock = BlockSize / 32; | ||
| uint gridDim = (uint)Math.Min((n + warpsPerBlock - 1) / warpsPerBlock, MaxDequantGridSize); | ||
| if (gridDim == 0) gridDim = 1; | ||
|
|
||
| CudaDriverApi.cuLaunchKernel(_pq2_0GemvF16InFunc, | ||
| gridDim, 1, 1, BlockSize, 1, 1, | ||
| 0, stream, (nint)args, 0).ThrowOnError(); | ||
| } |
Comment on lines
+1255
to
+1274
| /// <summary> | ||
| /// Dequantizes a PQ2_0 (PrismML Bonsai ternary) weight matrix to dense FP16 on the GPU for | ||
| /// prefill GEMM. <c>dst[row·k + col] = (code(W[row,col]) - 1) · group_scale</c>. Unlike I2_S, | ||
| /// the scale is per-128-element-group (not per-tensor), so no tensor-tail offset is read — | ||
| /// <paramref name="src"/> points at the packed row-major PQ2_0 payload only. | ||
| /// </summary> | ||
| public void LaunchDequantPQ2_0ToF16(nint src, nint dst, int n, int k, nint stream) | ||
| { | ||
| nint srcArg = src, dstArg = dst; | ||
| int nArg = n, kArg = k; | ||
| void** args = stackalloc void*[] {&srcArg, &dstArg, &nArg, &kArg}; | ||
|
|
||
| long totalGroups = (long)n * (k / 128); | ||
| uint gridDim = (uint)Math.Min((totalGroups + BlockSize - 1) / BlockSize, MaxDequantGridSize); | ||
| if (gridDim == 0) gridDim = 1; | ||
|
|
||
| CudaDriverApi.cuLaunchKernel(_dequantPQ2_0F16Func, | ||
| gridDim, 1, 1, BlockSize, 1, 1, | ||
| 0, stream, (nint)args, 0).ThrowOnError(); | ||
| } |
Comment on lines
+126
to
+128
| Half scale = (Half)groupScales[r * groupsPerRow + g]; | ||
| BitConverter.GetBytes(BitConverter.HalfToUInt16Bits(scale)).CopyTo(buf, (int)groupBase); | ||
|
|
Comment on lines
+1224
to
+1237
| public void LaunchPQ2_0GemvF32In(nint quantWeight, nint xF32, nint yF32, int n, int k, nint stream) | ||
| { | ||
| nint wArg = quantWeight, xArg = xF32, yArg = yF32; | ||
| int nArg = n, kArg = k; | ||
| void** args = stackalloc void*[] {&wArg, &xArg, &yArg, &nArg, &kArg}; | ||
|
|
||
| int warpsPerBlock = BlockSize / 32; | ||
| uint gridDim = (uint)Math.Min((n + warpsPerBlock - 1) / warpsPerBlock, MaxDequantGridSize); | ||
| if (gridDim == 0) gridDim = 1; | ||
|
|
||
| CudaDriverApi.cuLaunchKernel(_pq2_0GemvF32InFunc, | ||
| gridDim, 1, 1, BlockSize, 1, 1, | ||
| 0, stream, (nint)args, 0).ThrowOnError(); | ||
| } |
jamesburton
added a commit
that referenced
this pull request
Jul 31, 2026
…tract (#153) Review follow-up. SiLu.Execute: `Overlaps(Span)` is true for ANY overlap, not just exact in-place aliasing, so the tiled aliased path silently corrupted results when input and result overlapped at a non-zero offset (a tile write clobbered input elements a later tile still had to read). Switched to `Overlaps(other, out int elementOffset)` and added a memmove-style shifted path: each tile is snapshotted into stack scratch and tiles are walked backwards when the destination is ahead of the source. FusedOps.SwiGLU: three-span aliasing can require conflicting traversal directions, so shifted overlaps are rejected with ArgumentException instead of being mis-computed; the documented contract is now exact aliasing only. Also made the `up` snapshot buffer a conditional stackalloc so the common non-aliased path keeps a smaller stack frame. Tests: 8 shifted-overlap SiLu cases (both directions, shifts smaller and larger than the 256-float tile) — 7 of 8 fail against the previous implementation — plus 4 SwiGLU shifted-overlap rejection cases. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
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
CudaTransformerModel.cs'sProject()dispatch.Ternary-Bonsai-27B-Q2_0.ggufbytes, then re-confirmed via a real-GGUF dequant smoke test and a CPU-vs-real-RTX-3060 CUDA GEMV parity test at real Bonsai dims (hidden=5120, ffn=17408).Partial PR for #151 — this covers the CPU + CUDA backends only; Vulkan is tracked separately and will land in a follow-up PR before #151 is fully closed.
Test plan
PQ2_0Tests(12/12) — CPU dequant + GEMV/GEMM unit testsPQ2_0RealGgufSmokeTests— dequant smoke test against the real Bonsai GGUF (gated,DOTLLM_BONSAI_PQ2_0_GGUF)CudaPQ2_0GemvTest— CPU-vs-GPU numeric parity on real RTX 3060 at real Bonsai dims (n=512, k=5120/17408)Closes part of #151 (Vulkan follow-up still open on the issue).
🤖 Generated with Claude Code