Pull requests / #646

#646 perf(cuda,decode): zero-doorbell resident verify graph, sub-warp expert packing & shared-mem staging (+39-72% tok/s)

closed · @stuchapin909 · 0 comentários · No GitHub

BenchmarksMulti-GPUNVIDIA / CUDAModels & quants

Descrição

## Summary

This PR contributes a set of **100% bit-exact** CUDA kernel, verify graph, and multi-GPU pipeline optimizations for speculative decode (Verifier + MtpDrafter), profiled and verified on dual NVIDIA RTX 3090s (sm_86, 48 GB VRAM) running **Swift 1.5 IQ2_XS** at **256K context (K8V4)**.

All optimizations preserve exact IEEE-754 floating-point summation trees and produce bit-identical token trajectories and MTP draft acceptance counts at 	emperature=0.0.

---

## Key Optimizations

### 1. Routed Expert CUDA Kernels (src/kernels/cuda/iq_kernels.cu)
- **Sub-Warp Row Packing in 
ative_down_multi_kernel (
_ff == 640):**
  - **IQ4_NL (TD == 20,  = 40$ chunks/row):** Previously, 32 lanes/row left 24 of 32 lanes (75%) idle on the second pass (k = 32..39) and launched 320 blocks/group. Packing **4 rows per warp (8 lanes/row, 32 rows/block)** via 
ow_dot_40_sub8 achieves 100% active lanes, cuts launched blocks and __shared__ activation staging by **4x** (80 blocks/group), and matches the 32-lane warp_sum addition tree bit-for-bit via __fadd_rn.
  - **Q2_0 (TD == 42,  = 20$ chunks/row):** Previously left lanes 20..31 (37.5% of every warp) permanently idle. Packing **2 rows per warp (16 lanes/row, 16 rows/block)** via 
ow_dot_20_sub16 cuts launched blocks/staging by **2x** (160 blocks/group) with bit-for-bit parity.
- **Exact-$ (
 == 1, 
 == 2) Group Specialization:** Added 	emplate<int TY, int NC, bool EXACT_N = false> to 
ow_dot_multi, 
ow_dot_40_sub8, and 
ow_dot_20_sub16. Since ~80-86% of routed expert groups in a speculative window have  = 1$ and ~10% have  = 2$, dispatching <TY, 1, true> and <TY, 2, true> removes inner-loop predication (if (c < n)) and eliminates 75% of warp_sum reductions on single-token groups.
- **__shared__ Memory Staging for Codebooks & Down Activations:**
  - Staged IQ1_M (iq1s_grid_gpu), IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, and IQ3_S lookup grids into __shared__ memory (stage_iq_grid) in mmvq_multi_kernel and 
ative_gu_multi_kernel to eliminate __constant__ cache bank conflicts across divergent warp lanes.
  - Staged the group's  \le 4$ Q8_1 activation columns into __shared__ memory (s_hq_buf, stride 9 words = bank-conflict-free) in 
ative_down_multi_kernel.
  - Moved if (blockIdx.y >= *n_groups) return; above __shared__ staging so unused group blocks exit immediately.

### 2. Zero-Doorbell Multi-Layer Verify Graph (src/core/verify.cpp)
- When ll_resident_ is true (100% of routed experts resident in VRAM), captures the entire multi-layer verify pass per GPU into a single back-to-back CUDA graph (ull_gr_), eliminating 48 per-layer CPU-GPU doorbell round-trips (GPU-reach wait: 0.00 ms, per-layer host: 0.00 ms).
- Eliminated 48 redundant 32_to_bf16_bulk kernel launches per window on sh_stream when shared_expert_native_bf16_enabled() is active.

### 3. Shared Expert, PLE, GDN, Fused GR & MTP (shared_expert.cu, ple.cu, erify_kernels.cu, used_gr.cu, mtp.cpp)
- **Shared Expert (shared_expert.cu):** Added shared_gu_1t_kernel (4 warps/block with __shared__ Q8_1 input staging for =1$) and __shared__ Q8_1 staging in shared_down_kernel.
- **Batched Layer-1 PLE Projections (ple.cu, erify.cpp):** Batched ple_key (52.4 MB BF16) and ple_value (13.1 MB BF16) across all $ window tokens via ple_block_projected / f16_gemv_fp32_mmvf_multi, reading 65.5 MB of BF16 weights once per window instead of $ times.
- **GDN & Mixers (erify_kernels.cu):** Added single-launch gdn_conv_multi_kernel, state-only gdn_step_commit_kernel (192 blocks, skipping unused output projections during state commit), 48-block gdn_ab_multi_kernel, single-warp hc_write_multi_kernel, and multi-token head_mix_multi_kernel.
- **Register-Light gr_down_v3 (used_gr.cu):** Moved weight loads into the dot loop after __syncthreads() and added EXACT_T = true specializations for  \le 6$ with 100% shared-memory carveout.
- **Fused MTP Projections & Prefetch Overlap (mtp.cpp, 
gram.cpp):** Fused contiguous q_proj + q_idx and gate + up projections in MtpDrafter and overlapped NVMe/mmap PLE row prefetches with draft steps.

---

## Benchmark Results (Dual RTX 3090 sm_86, Swift 1.5 IQ2_XS, 256K K8V4, --spec 4, 	emp=0)

Every prompt was verified to produce the **exact same token sequence, window count, and draft acceptance count** (125/133, 141/185, 189/242) before and after:

| Workload | Baseline (ms/win) | Optimized (ms/win) | Baseline (	ok/s) | Optimized Engine (	ok/s) | Optimized HTTP Client (	ok/s) | Speedup |
| :--- | :---: | :---: | :---: | :---: | :---: | :---: |
| **LRU Cache (Code)** *(47 win, 3.60 tok/win)* | 30.69 ms | **18.13 ms** | 115.4 tok/s | **198.4 tok/s** | 187.8 tok/s | **+71.9%** |
| **Skip-List Reasoning** *(159 win, 1.89 tok/win)* | 19.70 ms | **13.78 ms** | 95.8 tok/s | **136.9 tok/s** | 133.9 tok/s | **+42.9%** |
| **Technical Prose** *(113 win, 2.65 tok/win)* | 22.62 ms | **16.43 ms** | 117.4 tok/s | **162.9 tok/s** | 159.0 tok/s | **+38.8%** |

No site

Links install, modelos, releases.