Skip to content

Add CUDA backend for Laguna S 2.1#594

Draft
devteapot wants to merge 16 commits into
antirez:laguna-s2.1from
devteapot:laguna-s2.1-cuda
Draft

Add CUDA backend for Laguna S 2.1#594
devteapot wants to merge 16 commits into
antirez:laguna-s2.1from
devteapot:laguna-s2.1-cuda

Conversation

@devteapot

@devteapot devteapot commented Jul 23, 2026

Copy link
Copy Markdown

Summary

  • add native CUDA inference support for Laguna S 2.1 while keeping the shared Metal/CUDA model graph as the scheduler
  • rebase onto laguna-s2.1 commit 3b0ec5f and support all released weight recipes: mixed routed Q2_K/Q3_K, revised Q8_0 signal/all-Q4_K routed, and legacy F16/Q4_K/Q6_K
  • implement the canonical Q3_K scale/mask decode with DP4A dot products, reuse the existing Q2_K arithmetic, and route both through Laguna-specific decode and expert-sorted prefill/down kernels
  • preserve revised-checkpoint accuracy with bounded Q8-to-F16 caching plus a streaming dequant/cuBLAS fallback for multi-token projections
  • implement Laguna embedding, per-head RMSNorm/YaRN RoPE, gated GQA attention, F16 ring KV cache, and routed/shared MoE for decode and prefill
  • support Poolside's official BF16 Laguna DFlash drafter on CUDA, including six target-layer captures, feature fusion/KV injection, causal block drafting, and batched target verification
  • add optional capability-gated Blackwell decode-attention and DFlash 9:1 GQA tuning with portable fallbacks and rollback switches
  • add Blackwell warp-per-query-head target prefill, 16-row K/V staging, and exact Q4_K INT8 tensor-core gate/up/down kernels, including 16-pair model-specific MoE tiles

Motivation

Laguna S 2.1 already had model loading and a Metal path, but no end-to-end NVIDIA backend. The source branch now also accepts a 44.95 GiB mixed checkpoint whose routed layers 1..20 use Q2_K and layers 21..47 use Q3_K. CUDA previously had generic Q2_K arithmetic but did not deliberately route this Laguna layout, and it had no Q3_K implementation.

This PR gives all three Laguna recipes an accurate CUDA path and adds target-verified DFlash acceleration for workloads with sufficient draft acceptance.

Impact

Users with a supported NVIDIA GPU can run:

  • the 44.95 GiB mixed Q2_K/Q3_K checkpoint on 64 GiB systems
  • the pinned 68,248,759,648-byte (63.56 GiB) Q4_K_M checkpoint on systems with at least 96 GiB
  • the earlier 70.01 GiB F16/Q4_K/Q6_K checkpoint

The optional official DFlash BF16 GGUF adds about 2.08 GiB of weights plus graph buffers and works with either the mixed or Q4 target. Raw decode remains the default. Laguna remains single-GPU/full-residency: SSD streaming, distributed inference, ROCm inference, and tensor parallelism are not enabled.

GB10 benchmark: checkpoints with and without DFlash

Current end-to-end results on the NVIDIA GB10 using this PR's native Blackwell build at commit 37abbc5:

make cuda-spark -j2
./ds4-bench -m <checkpoint> --cuda \
  --prompt-file ds4.c --ctx-start 2048 --ctx-max 2048 \
  --gen-tokens 256
# Add for DFlash:
# --mtp gguf/laguna-s-2.1-DFlash-BF16.gguf --mtp-draft 15

The build emits native sm_121 code. No rollback environment variables were set, so these runs use the Blackwell target-prefill, raw-decode, and DFlash kernels. Each value is the median of three repeated warm runs after normalizing unified-memory pressure; ranges show process-to-process spread. Model startup is outside the measured interval. Generation is greedy, and steady generation excludes the first decode/speculative cycle.

checkpoint mode prefill median (range) generation median (range) steady generation generation vs raw
revised Q4_K_M, 63.56 GiB raw CUDA 749.92 t/s (748.39–753.85) 21.67 t/s (21.67–21.67) 21.80 t/s
revised Q4_K_M, 63.56 GiB DFlash, depth 15 652.34 t/s (647.75–653.97) 45.96 t/s (45.83–46.25) 47.29 t/s +112.1%
mixed Q2_K/Q3_K, 44.94 GiB raw CUDA 376.19 t/s (376.12–377.97) 25.19 t/s (25.10–25.22) 25.37 t/s
mixed Q2_K/Q3_K, 44.94 GiB DFlash, depth 15 363.02 t/s (362.57–363.32) 45.00 t/s (44.98–49.29) 46.09 t/s +78.6%

These paired results measure prefill and generation from the same binary and workload. Enabling DFlash reduces prefill throughput by 13.0% on revised Q4 and 3.5% on mixed Q2/Q3 because the capture path is active during prompt ingestion; DFlash accelerates generation, not prefill. Its generation gain remains workload- and acceptance-dependent, which is why it stays explicitly opt-in.

The revised Q4 target has the fastest prefill because its exact INT8 tensor-core path uses 16-pair gate/up and chunked-down tiles. The smaller mixed checkpoint uses the Q2_K/Q3_K DP4A path and is faster for raw generation, but not for prefill.

On DFlash, the Blackwell 9:1 GQA block shares each K/V fetch across the official drafter's nine query heads per KV head. DS4_CUDA_DFLASH_NO_BLACKWELL=1 selects the portable kernel. For target prefill, DS4_CUDA_LAGUNA_NO_WARP_GQA_PREFILL=1 selects portable attention and DS4_CUDA_LAGUNA_NO_Q4_MMA=1 disables exact Q4_K tensor-core routing. DS4_CUDA_LAGUNA_NO_Q4_MMA_TILE16=1 and DS4_CUDA_LAGUNA_NO_Q4_MMA_DOWN_TILE16=1 retain the eight-pair rollback paths. DS4_CUDA_LAGUNA_NO_BLACKWELL_SPLIT16=1 disables only the wider raw-decode attention schedule.

Validation

  • make cuda-spark -j2
  • make cpu -j2
  • make cuda-regression CUDA_ARCH=native
  • portable sm_75 Laguna compile and regression on GB10
  • tests/cuda_laguna_smoke, including patterned Q3_K scale/mask coverage, Q2_K/Q3_K decode and expert-sorted batch MoE, real 72:8 Blackwell GQA coverage, and 128-token Q4_K tensor-core/rollback equivalence at 1e-6
  • tests/cuda_long_context_smoke
  • full official mixed checkpoint: 2K prefill, sustained 256-token raw generation, and sustained 256-token portable/Blackwell DFlash generation
  • patch-md validate laguna-s21-cuda --ready

The broader model-dependent tool-call/logprob fixtures remain sensitive to which Laguna checkpoint is linked. The focused Laguna CUDA builds, numerical regressions, and real-model runtime checks pass.

carlid added 6 commits July 24, 2026 17:00
- Tile warp GQA prefill and add A/B controls
- Select tile16 MoE kernels only for compatible PTX
- Extend smoke tests for tiled attention equivalence
- Add an environment toggle to isolate the down-projection tile
- Document the new CUDA tuning control
- Fuse Q/K/V/gate Q8 projections and single-token Q/K norm/RoPE launches
- Add portable fallbacks, controls, documentation, and CUDA smoke coverage
- Calibrate target decode cost per session and adjust depth based on profitability
- Add DS4_DFLASH_ADAPTIVE override and adaptive scheduling diagnostics
- Document adaptive depth behavior and configuration
- Preserve the target model mapping on CUDA
- Commit only accepted DFlash prefixes and roll back rejected SWA rows
- Add transient-cache coverage to CUDA smoke tests
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant