cuda: optimize Bonsai low-bit kernels for GB10

[pooyakhosravi]

Overview

This PR improves the CUDA execution path for Bonsai Q1_0 and ternary Q2_0/PQ2_0 models on NVIDIA GB10 (compute capability 12.1). It targets the low-bit matrix kernels and the surrounding elementwise and recurrent operations that dominate Bonsai inference.

The implementation stays inside the existing CUDA backend and uses GB10-specific dispatch wherever kernel geometry or resource tradeoffs are device-specific. It does not change GGUF formats or public APIs.

The main changes are:

• Fuse SWIGLU directly into Q8 activation quantization, avoiding an intermediate activation write/read and kernel launch.
• Fuse residual add, RMS normalization, and scaling patterns used by Bonsai graphs.
• Reuse a quantized Q8 activation across adjacent low-bit projections.
• Use cp.async.cg to overlap Q1_0, Q2_0, and PQ2_0 MMQ activation staging with computation.
• Prefetch packed Q1_0 J128 MMQ X tiles one K iteration ahead in 32-byte L2 sectors, overlapping global-memory dependency latency with current-tile work.
• Add GB10 scheduling for Gated Delta Net and precompute exp(g) for long prompts.
• Add vectorized scale, RMS-128, and transpose-concat specializations.
• Specialize fused and non-fused MMVQ kernels at compile time and prefetch packed low-bit weights into L2 during decode.
• Preserve Prism reduction geometry for arithmetic-sensitive shapes. In particular, the PQ2_0 6144x2048 down projection retains the reference geometry after a strict KL gate detected a reduction-order difference.
• Guard fused graphs against unsafe memory aliasing, cache their environment switches, precompute shared-Q8 consumer counts, and centralize Q8 MMQ buffer sizing.

The Q1_0 J128 prefetch is bounded by kb0_stop and compiled only for that type and tile width. It is a non-faulting cache hint: it does not write memory, change synchronization, alter addressing, or change the reduction order. The second asynchronous Y buffer remains equivalent because k00 advances only X-tile indices inside vec_dot; Y is indexed locally from the supplied base.

No benchmark reports or Markdown files are included in the commits.

Performance

Hardware: NVIDIA GB10, CUDA compute capability 12.1, full GPU offload. The GPU had no other compute processes during measurement.

Baseline: untouched Prism e311ed38f (build 10660). Final PR: 037953ae6 (build 10667). The intervening commits on the target prism branch only update policy files, so the CUDA source baseline is unchanged.

Every row below comes from one contemporaneous baseline-to-final experiment. Each result uses four independent alternating baseline/final process pairs, with five timed repetitions per process. Raw throughput errors are SEM across process means. The paired change uses a two-sided Student-t 95% confidence interval.

Model	Test	Prism tok/s	Final PR tok/s	Paired change (95% CI)
Bonsai 27B Q1_0	pp512	1032.114 +/- 1.799	1255.485 +/- 2.617	+21.643% +/- 0.928%
Bonsai 27B Q1_0	tg128	46.117 +/- 0.062	49.458 +/- 0.042	+7.245% +/- 0.170%
Ternary-Bonsai 27B PQ2_0	pp512	998.853 +/- 2.008	1171.533 +/- 0.959	+17.289% +/- 0.679%
Ternary-Bonsai 27B PQ2_0	tg128	29.942 +/- 0.015	30.668 +/- 0.014	+2.423% +/- 0.126%
Ternary-Bonsai 27B Q2_0 group 64	pp512	997.788 +/- 1.775	1165.948 +/- 0.084	+16.854% +/- 0.665%
Ternary-Bonsai 27B Q2_0 group 64	tg128	28.672 +/- 0.015	29.548 +/- 0.009	+3.056% +/- 0.271%
Bonsai 8B Q1_0	pp512	3980.046 +/- 4.366	4587.406 +/- 12.120	+15.262% +/- 1.371%
Bonsai 8B Q1_0	tg128	165.985 +/- 0.141	170.814 +/- 0.108	+2.910% +/- 0.192%
Ternary-Bonsai 8B PQ2_0	pp512	3996.699 +/- 5.562	4488.991 +/- 10.204	+12.317% +/- 0.470%
Ternary-Bonsai 8B PQ2_0	tg128	104.043 +/- 0.059	104.840 +/- 0.011	+0.766% +/- 0.167%
Bonsai 4B Q1_0	pp512	6172.892 +/- 10.152	7105.418 +/- 9.995	+15.108% +/- 0.800%
Bonsai 4B Q1_0	tg128	245.588 +/- 0.089	249.762 +/- 0.166	+1.700% +/- 0.158%
Ternary-Bonsai 4B PQ2_0	pp512	6221.004 +/- 15.366	6975.845 +/- 23.717	+12.133% +/- 0.696%
Ternary-Bonsai 4B PQ2_0	tg128	169.079 +/- 0.087	171.380 +/- 0.141	+1.362% +/- 0.395%
Bonsai 1.7B Q1_0	pp512	13778.072 +/- 34.029	15711.379 +/- 26.149	+14.033% +/- 0.856%
Bonsai 1.7B Q1_0	tg128	457.105 +/- 0.118	466.897 +/- 0.297	+2.142% +/- 0.145%
Ternary-Bonsai 1.7B PQ2_0	pp512	14136.859 +/- 26.124	15794.467 +/- 30.432	+11.726% +/- 0.423%
Ternary-Bonsai 1.7B PQ2_0	tg128	338.635 +/- 0.367	347.525 +/- 0.143	+2.626% +/- 0.415%

The group-64 row uses the official Ternary-Bonsai-27B-Q2_g64.gguf artifact, not the deprecated unsuffixed file.

DSpark control

The final Q1_0 27B target was also exercised end to end with its Q4_0 DFlash sidecar using draft-dspark, draft maximum 4, temperature 0, and a deterministic 256-token code prompt. Twelve independent alternating isolation pairs produced exactly 332 drafted and 177 accepted tokens in every run.

The J128 cache hint is intentionally a prefill optimization and does not measurably change DSpark: 82.551 +/- 0.117 versus 82.455 +/- 0.128 tok/s SEM, paired -0.114% +/- 0.416% at 95% confidence. The PR's ordinary target-decode improvements are represented by the tg128 rows above; no additional DSpark gain is claimed for the J128 prefetch.

Numerical validation

Final-head logits were compared with compressed reference logits generated by untouched Prism. Each model was tested in both modes:

• Batched prefill: four chunks of 512 tokens, -c 512 -b 2048 -ub 512.
• Sequential evaluation: two chunks of 128 tokens, -c 128 -b 1 -ub 1.

All tested models retain 100% top-token agreement. Mean KL remains at the compressed-reference numerical floor.

Model	Prefill PPL	Sequential PPL	Worst mean KL	Maximum KL	Same top token
Bonsai 8B Q1_0	12.428705	15.906751	approximately 0.000000	0.000055	100%
Ternary-Bonsai 8B PQ2_0	10.311228	14.188520	approximately 0.000000	0.000050	100%
Bonsai 4B Q1_0	15.156022	18.300087	-0.000001	0.000064	100%
Ternary-Bonsai 4B PQ2_0	11.378429	16.317651	0.000001	0.000046	100%
Bonsai 1.7B Q1_0	21.215227	22.931751	approximately 0.000000	0.000054	100%
Ternary-Bonsai 1.7B PQ2_0	16.801363	21.477474	0.000001	0.000055	100%

The 27B gates remain at the same floor:

• Q1_0: mean KL 0.000000, maximum KL 0.000052, RMS probability delta at most 0.001%, 100% same top token, and PPL ratio 1.000000.
• PQ2_0: mean KL approximately 0.000000, maximum KL 0.000050, 100% same top token.
• Official Q2_0 group 64: mean KL approximately 0.000000, maximum KL 0.000050, 100% same top token.

The final Q1 path was additionally isolated against the immediately preceding arithmetic-identical head across 27B/8B/4B/1.7B. Across batched and sequential modes, mean KL was effectively zero, maximum KL was at most 0.000060, RMS probability delta was at most 0.001%, and same-top agreement was 100%. Baseline self-comparisons exactly reproduced the small saved-logit PPL ratios on smaller models, identifying them as serialization/reload floor rather than incremental loss.

During development, a wider PQ2_0 schedule for the Ternary-Bonsai 1.7B down projection measured mean KL 0.000236 and maximum KL 0.002956. The submitted shape guard restores Prism reduction geometry for that projection. Its final sequential result is mean KL -0.000001, maximum KL 0.000039, and 100% top-token agreement while retaining the measured tg128 gain.

Profiling and roofline analysis

Nsight Systems attributes approximately 66-67% of Q1_0 pp512 CUDA time to the ordinary low-bit MMQ family, approximately 10% to Gated Delta Net, and approximately 6% to fused SWIGLU-to-Q8 quantization.

A final-head Nsight Compute profile of mul_mat_q<Q1_0,128,false> on 27B pp512 reports:

• 727.71 us for the sampled kernel launch.
• 40.87% of peak SM/issue throughput and 35.14% tensor-pipe activity.
• 37.96% memory-pipeline, 42.03% L1/TEX, and 29.16% L2 throughput.
• 254 registers per thread and 76.29 KiB dynamic shared memory per CTA.
• 16.67% theoretical and 16.84% achieved occupancy: one resident CTA per SM.
• Zero local-memory loads and stores, so there is no register spilling.
• 57.47% no-eligible-warp time, with long scoreboard the leading stall at 0.88 cycles per issued instruction.

The profile confirms that the kernel is latency/resource limited rather than at raw DRAM bandwidth. The next-tile sector hints raise cache activity while overlapping long-scoreboard latency; the remaining one-CTA residency is constrained by both registers and shared memory.

There is no single meaningful roofline percentage for this mixed packed-integer/tensor/FP32 kernel:

• The relevant measured hot-kernel utilization is 40.87% of SM issue peak and 35.14% of tensor-pipe peak.
• Nsight's generic FP32 chart reports 13% of FP32 peak, but FP32 is not the governing operation roof for packed Q1 MMQ.
• If the final MMQ family alone scaled ideally from its measured issue utilization to 100%, while all other kernels stayed fixed, Amdahl's law gives an approximate 2,046 tok/s pp512 ceiling. The measured 1,255 tok/s is about 61% of that limited engineering roof.
• NVIDIA specifies 273 GB/s memory bandwidth for GB10. Dividing it by the 3,792,459,776-byte Q1_0 model gives an optimistic ordinary-decode ceiling of 71.98 tok/s. The final 49.458 tok/s is about 68.7% of that ceiling, before accounting for KV and activation traffic.

Reaching 3,000 pp512 tok/s would require approximately 2.39x over the final result. Raising only MMQ to its measured issue roof tops out near 2.05k tok/s, so 3k requires broader changes to tiling/resource footprint, fusion, and non-MMQ kernels rather than another small cache hint.

A representative final Gated Delta Net profile reports 82.37% compute/memory SOL, 94.42% L1 throughput, 83.33% theoretical occupancy, 64.33% achieved occupancy, and 48 registers per thread.

Tests

• test-backend-ops test -b CUDA0 -o MUL_MAT: 1238/1238 passed.
• CUDA GATED_DELTA_NET: 46 cases passed.
• CUDA CONCAT: 192 cases passed.
• CUDA RMS_NORM, SCALE, and GLU backend tests passed.
• git diff --check passes.
• llama-bench, llama-perplexity, and llama-server build successfully.

Requirements

• I have read and agree with the contributing guidelines.
• AI usage disclosure: YES. AI assisted with implementation exploration, mechanical code generation, profiling, benchmark execution, numerical-quality validation, and regression isolation. I manually reviewed the submitted changes and accept responsibility for understanding, explaining, and maintaining them.

[Copilot]

Pull request overview

Optimizes Bonsai low-bit CUDA inference on NVIDIA GB10 without changing public APIs or model formats.

Changes:

• Adds fused normalization, SWIGLU quantization, and shared Q8 activation paths.
• Introduces GB10-specific MMQ/MMVQ staging, scheduling, and prefetching.
• Optimizes recurrent, scale, RMS, and concat kernels.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
tests/test-backend-ops.cpp	Adds a Gated Delta Net test case.
ggml/src/ggml-cuda/scale.cu	Adds vectorized GB10 scaling.
ggml/src/ggml-cuda/quantize.cuh	Declares fused Q8 quantizers.
ggml/src/ggml-cuda/quantize.cu	Implements SWIGLU and RMS Q8 fusion.
ggml/src/ggml-cuda/norm.cuh	Declares fused normalization operations.
ggml/src/ggml-cuda/norm.cu	Adds GB10 normalization kernels and fusion.
ggml/src/ggml-cuda/mmvq.cu	Specializes decode kernels and prefetching.
ggml/src/ggml-cuda/mmq.cuh	Adds GB10 configuration and async staging.
ggml/src/ggml-cuda/mmq.cu	Adds reusable Q8 and dual-projection execution.
ggml/src/ggml-cuda/mmq-vec-dot.cuh	Removes trailing whitespace.
ggml/src/ggml-cuda/mmq-config-blackwell.cuh	Adds the GB10 MMQ configuration entry point.
ggml/src/ggml-cuda/ggml-cuda.cu	Adds graph fusion and shared-Q8 dispatch.
ggml/src/ggml-cuda/gated_delta_net.cu	Adds GB10 scheduling and exponential precomputation.
ggml/src/ggml-cuda/concat.cu	Adds transpose-concat specialization.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[khosravipasha]

Reviewed in detail and tested for regressions on non-GB10 CUDA (L40S, sm_89 Ada, CUDA 12.8, build of this branch at 282ed52 vs prism e311ed3). Approving; one small suggestion below.

Code review. The isolation discipline here is exemplary: every behavioral change (the three graph fusions, the cp.async double-buffered y-tiles in MMQ, the MMVQ prefetch, the GDN cols-per-warp scheduling, the RMS-128/scale/concat specializations) is gated on GGML_CUDA_CC_DGX_SPARK host-side and __CUDA_ARCH__ device-side, with env kill-switches for the fusions. The numerics care shows too: the reduction-order comment in the fused RMS quantizer and the KL-gated shape guard on the down projection are exactly the right instincts for these models. The PR description's methodology (paired process alternation, CIs, KL gates) is a model for perf submissions.

L40S regression check.

Backend op suites all pass: MUL_MAT 1238/1238, GATED_DELTA_NET 37/37, RMS_NORM 51/51, CONCAT 177/177, ADD 99/99, SCALE 4/4.

Interleaved A/B llama-bench (r=5, two alternating process pairs per side):

Model	PP512 base -> PR	TG128 base -> PR
Bonsai-27B Q1_0	2936 -> 2921 (-0.5%)	107.4 -> 109.4 (+1.8%)
Ternary-27B PQ2_0	2997 -> 2980 (-0.6%)	74.3 -> 74.3 (flat)
Ternary-27B Q2_0 g64	2995 -> 2976 (-0.65%)	71.3 -> 71.0 (-0.4%)
Ternary-1.7B PQ2_0	33783 -> 33959 (+0.5%)	611 -> 621 (+1.7%)

Numerics on sm_89: 8-chunk wikitext PPL is bit-identical between builds (11.2139 +/- 0.72003 both), greedy 128-token outputs are bit-identical on both 27B models, and the draft-dspark path runs correctly on this branch (131 t/s vs 74 baseline at n-max 4). The Q1_0/1.7B decode gains look like a free bonus from has_gate becoming compile-time dispatch.

One suggestion (non-blocking). In the three fusion condition chains, getenv() runs before the cc == GGML_CUDA_CC_DGX_SPARK test, so every ADD/MUL_MAT node on every architecture pays a libc getenv per graph build; we suspect this explains the small but consistent 0.5-0.65% prefill dip we measure on the 27Bs. Putting the cc test first (it is a cached struct read) or caching the env lookups in a static should make non-GB10 devices pay nothing.

Good to go from our side otherwise; the GB10 numbers speak for themselves.

Converting to a comment for now: test results stand, but we want to scope the minimal change set before formally approving.

[bri-prism]

The measurement work here is the best I have seen on this repo and I want to say so before anything else. Paired process-level A/B with confidence intervals over process means, a KL gate on every model, and you caught your own regression: the wider PQ2_0 schedule on the 1.7B down projection at mean KL 2.36e-4, then guarded it back to the reference geometry while keeping most of the gain. That is exactly the discipline that stops a reduction-order change from shipping as a speedup.

One thing I would like changed before this lands, and a few smaller notes.

The two new ADD-based fusions skip the aliasing check. Your SWIGLU fusion correctly calls ggml_cuda_check_fusion_memory_ranges alongside ggml_can_fuse_subgraph. The dual-RMS fusion (5 nodes, outputs i, i+3, i+4) and the add-plus-rms-plus-mul fusion (3 nodes, outputs i, i+2) only call ggml_can_fuse_subgraph. Nearly every other fusion in ggml-cuda.cu pairs the two. That helper exists to refuse fusion when a fused output overlaps an input that is not an elided intermediate, and both of these write two live tensors while reading the residual, so they are exactly the shape it guards. Your KL numbers say it does not bite on the Bonsai graphs on GB10, which I believe, but the allocator picks the buffers and a different graph shape could alias silently into wrong numerics rather than a crash. It looks like one line each.

Three getenv calls sit on the per-node path and are evaluated before the device gate. GGML_CUDA_GB10_VIRTUAL_RMS_Q8 and GGML_CUDA_GB10_SHARED_Q8 in the node loop, GGML_CUDA_GB10_DUAL_RMS_Q8 in try_fuse. Because && short circuits left to right and the getenv is the first term, every CUDA device pays it per node per graph build, not just GB10. Caching each in a function-local static would keep the escape hatch and put the cost back where it belongs. The cc == GGML_CUDA_CC_DGX_SPARK checks that follow are correctly placed, so nothing else leaks onto other hardware.

gb10_shared_q8_consumer_count scans the whole graph per call. It is called once in the guard and again when the entry is constructed, plus once per entry at the top of each while pass, which makes it quadratic in node count on GB10 graph builds. It is correctly behind the cc check so no other device sees it. Worth a single prepass building a map from src1 and type to a count, if it ever shows up on a prefill profile.

A duplicated invariant worth collapsing. ggml_cuda_mul_mat_q_q8_size hardcodes sizeof(block_q8_1_mmq)/QK8_1_MMQ while the allocation it has to match uses y_block_size/y_values_per_block, and ggml_cuda_mul_mat_q_fused_two writes the same ratio a third way as sizeof(block_q8_1)/QK8_1. All three are 1.125 bytes per value today so this is not a bug, and the !use_native_fp4 assert keeps the FP4 case out. But an external buffer that is silently too small is device heap corruption, so I would rather see one shared helper than three expressions that happen to agree.

One question rather than a finding. In the async_buffer_y path the second vec_dot now reads tile_y_next, a separate buffer filled from by1, at offset MMQ_TILE_NE_K, where the original reads the same tile_y at that offset. The synchronisation looks right to me, cp_async_wait_all is followed by __syncthreads() on both paths, so visibility is covered. I could not convince myself from the diff alone that the addressing is equivalent, and your KL results say it is. Could you say a sentence about why, so the next person reading it does not have to re-derive it?

Scope note on this review: I have no GB10 here, so none of this was run. Everything above is read off the diff and the prism sources, and your performance and KL tables are unverified by me. The code only activates on GGML_CUDA_CC_DGX_SPARK, so the only point that touches other hardware is the getenv one.

[pooyakhosravi]

Addressed the review feedback in fff11bb.

Changes:

• Added ggml_cuda_check_fusion_memory_ranges guards to the dual-projection and add/RMS/MUL fusions, and to the virtual-RMS variant of the same graph shape.
• Cached the three GB10 fusion environment switches in function-local statics.
• Replaced repeated whole-graph shared-Q8 consumer scans with one prepass.
• Consolidated Q8 MMQ buffer sizing behind one helper and used the MMQ block layout consistently for the fused stride.
• Documented why the second async Y buffer retains the original logical K offset.
• Guarded padded SWIGLU gate loads, corrected the GDN pool allocation units, moved RMS128 dispatch out of NORM and into RMS_NORM, and removed the duplicate GDN test.
• Retained persistent shared-Q8 and row-scale allocations in one global allocation-order container, then release them in exact reverse order for the stack-like VMM pool.

Validation on the GB10:

• Clean full CUDA build.
• Targeted CUDA NORM/RMS_NORM/GATED_DELTA_NET suite: 137/137 supported cases passed.
• End-to-end 27B llama-bench, 3 samples each:
  • Q1_0: pp512 1204.9 ± 18.3 tok/s; tg32 49.14 ± 0.61 tok/s.
  • PQ2_0: pp512 1182.8 ± 22.5 tok/s; tg32 30.58 ± 0.34 tok/s.
  • Official Q2_g64: pp512 1175.6 ± 24.5 tok/s; tg32 29.53 ± 0.11 tok/s.
• Strict Wikitext-2 4x512 comparison against the Prism references:
  • Q1_0: mean/max KL 0.000000 / 0.000046, RMS Δp 0.001%, same-top 100%.
  • PQ2_0: mean/max KL 0.000001 / 0.000057, RMS Δp 0.001%, same-top 100%.
  • Q2_g64: mean/max KL 0.000001 / 0.000057, RMS Δp 0.001%, same-top 100%.

The quality results remain at the compressed-reference floor.

[pooyakhosravi]

Follow-up 6012190 sharpens the async-Y addressing comment: k00 advances only X-tile indices inside vec_dot; Y is indexed locally from the supplied base. Both paths stage by1 at local Y offset zero, so tile_y_next changes only the shared-memory storage address while MMQ_TILE_NE_K selects the second half of X.

[pooyakhosravi]

Added the evaluated Q1/J128 next-tile L2 sector prefetch in 037953ae6. The PR description now includes the corrected matched-head benchmarks, all-size quality gates, DSpark control, and Nsight roofline/scheduler analysis. Headline: Q1 prefill improves 3.46-3.80%; decode and DSpark are neutral; mean KL is effectively zero with 100% top-token agreement. No benchmark or Markdown artifacts were committed.