diff --git a/README.md b/README.md index fe5f6ed..ef27228 100644 --- a/README.md +++ b/README.md @@ -51,13 +51,19 @@ coherent (`scripts/bench-all.ps1`); top-1 parity vs llama.cpp b8585 verified on | Qwen3.6-35B-A3B-MTP (GDN+MoE) | (same) | 22 GB | **CUDA** `-g -1 --no-thinking` (hybrid) | **65.0** | **22.9** | Requires `SHARPI_CPU_MOE=1`: 30 GDN + 10 attn + shared expert on GPU, routed experts CPU mmap. 100% acceptance. Fused GDN scan + batched SDPA (#114-B/#118), bit-identical, grows with ctx | | Carnice (Qwen3.6-35B-A3B-MTP finetune) | [mudler](https://huggingface.co/mudler/Carnice-Qwen3.6-MoE-35B-A3B-APEX-MTP-GGUF) | 17 GB | **CUDA** `-g -1 --no-thinking` (hybrid) | **43.6** | **25.0** | agentic finetune of 35B-A3B-MTP; 77% acceptance (`bench-carnice.ps1` — the default prompt 1-token-EOSes on this terser tune). APEX mixed-precision (Q3_K + Q8_0 experts); Q8_KS per-32 int dots auto-enable at load (#99/#101/#107), +4.6% decode at ~4× tighter parity vs plain Q8_K (`SHARPI_Q3K_Q8K=0`/`SHARPI_Q8_0_Q8K=0` to disable). Fused GDN scan + wave SDPA (#114-B/#118) bit-identical past 4096 | | Gemma 4 E4B-it Q8 | [unsloth](https://huggingface.co/unsloth/gemma-4-E4B-it-GGUF) | 8 GB | CPU | 4.9 | 5.0 | dense 42-layer gemma4: per-layer head_dim (256 SWA / 512 global), dual-RoPE, KV-share tail (18 layers), 5:1 SWA:global, logit softcap 30, PLE-256 injection (~4.2 GB mmap-resident) | -| Gemma 4 E4B-it Q8 | (same) | 8 GB | **CUDA** `-g -1 -c 2048` | **109.4** | **49.1** | all 42 layers fit at `-c 2048`. KV-share alias + SWA/global split per layer; PLE projections (~215 MB) upload at construction. Batched-trunk prefill (#136, bit-identical; `SHARPI_BATCHED_PREFILL=0` → 53.4 t/s to bisect) collapses the per-position attention/FFN/PLE launches into batched GEMM-N + batched SWA/full attention (≈2.05× prefill). Launch-bound decode also gains (44→49) from the fused PLE projection, dropped Q-prescale, and fused Q/K HeadNorm | +| Gemma 4 E4B-it Q8 | (same) | 8 GB | **CUDA** `-g -1 -c 2048` | **2853** | **59** | all 42 layers fit at `-c 2048`. KV-share alias + SWA/global split per layer; PLE projections (~215 MB) upload at construction. **Prefill (#141):** int8 **tensor-core MMQ** matmul (`mma.m16n8k32.s8`, each Q8_0 weight read once as int8 — beats the dequant→fp16→cuBLAS GEMM, drops its fp16 HBM temp; `SHARPI_PREFILL_MMQ=0` reverts) + a memory-efficient **flash-attention** prefill (shared K/V tiles + online softmax + half2 fp16x2 QK dot — replaces the scalar O(n²) per-query attention that re-streamed each query's K/V window up to ~512×; `SHARPI_PREFILL_FLASH=0` reverts) + a batched Q8_0 embedding lookup. **~1.8× at ~1K ctx (1564→2853), ~2.05× at 1.8K** — profiling showed *attention*, not the matmul, was the dominant prefill cost at realistic prompt lengths, so the win grows with prompt length. **Decode (#142):** dp4a/Q8_1 int8 matvec (`SHARPI_Q80_DP4A=0` to bisect) + CUDA-graph capture/replay default-on (`SHARPI_CUDA_GRAPH=0` to bisect). All prefill/decode fast paths are argmax-stable vs the fp32 path, not bit-exact. Remaining gap to llama.cpp (~8000 prefill / ~78 decode): full tensor-core flash at d=512 + decode matvec work | | Gemma 4 E4B-it Q8 | (same) | 8 GB | **CUDA** `-g 22 -c 2048` (hybrid) | 6.6 | 6.8 | 22 GPU + 20 CPU layers. `-g ≤ 22` required so the CPU shared-KV tail can read its own-KV source layers; CPU dense-FFN dominates decode (bandwidth-bound). `SHARPI_CUDA_PROFILE=1` for per-phase breakdown | _Numbers re-measured across every on-disk row at ~1K ctx so the prefill column is comparable; per-issue before/after figures in the notes are historical. Llama-4 Scout and Qwen3-Coder Vulkan-hybrid keep their prior values (not re-runnable on the bench machine)._ +**Recommended sampling for Gemma 4 E4B-it:** `--temp 1.0 --top-k 64 --top-p 0.95 --min-p 0` +(the Gemma 3/4 family defaults). Gemma 4 E4B-it is **not** a reasoning model, so the CLI now +defaults `enable_thinking=false` for it automatically (no `--no-thinking` needed) — otherwise the chat +template renders a `` block the model wasn't trained to fill and the output degenerates. Greedy +(`--temp 0`) is not recommended for it either; use the sampling values above. + `--backend auto` (default) picks CUDA when available, sizing the GPU/CPU split from VRAM via TierPlanner; falls through to Vulkan only when CUDA isn't present. For hybrid `qwen35moe` models the CUDA backend keeps attention KV, the GDN layers, and the shared expert in VRAM; routed-expert dispatch auto-selects between diff --git a/src/SharpInference.Cli/RunCommand.cs b/src/SharpInference.Cli/RunCommand.cs index 6e83a69..b90d2fb 100644 --- a/src/SharpInference.Cli/RunCommand.cs +++ b/src/SharpInference.Cli/RunCommand.cs @@ -230,9 +230,20 @@ protected override int Execute(CommandContext context, Settings settings, Cancel s_endThinkTokenId = endChannelId; } + // Gemma 4 E4B-it brackets a <|channel>thought block in its chat template but is + // NOT a reasoning model — rendering enable_thinking=true makes it try to fill a + // think section it wasn't trained for and the output degenerates. Default thinking + // OFF for Gemma 4 (its recommended config); --no-thinking still forces it off for + // every other model. Pass --temp 1.0 --top-k 64 --top-p 0.95 for Gemma 4. + bool modelDefaultsThinkingOff = s_arch == "gemma4"; + s_noThinking = settings.NoThinking || modelDefaultsThinkingOff; + if (modelDefaultsThinkingOff && !settings.NoThinking) + AnsiConsole.MarkupLine("[dim]Gemma 4 is not a reasoning model — defaulting to --no-thinking " + + "(recommended: --temp 1.0 --top-k 64 --top-p 0.95).[/]"); + // Greedy on a reasoning model tends to "wait, but actually" itself into infinite // loops; --no-thinking sidesteps the issue since the model won't reason at all. - if (s_thinkTokenId > 0 && settings.Temperature == 0f && !settings.NoThinking) + if (s_thinkTokenId > 0 && settings.Temperature == 0f && !s_noThinking) { AnsiConsole.MarkupLine("[yellow]Warning:[/] Greedy decoding (--temp 0) on a reasoning model often produces"); AnsiConsole.MarkupLine("infinite \"wait, but actually\" loops. Consider [yellow]--temp 0.6 --top-p 0.95 --top-k 20[/]."); @@ -668,7 +679,7 @@ private static int RunSpeculativeSinglePrompt(Settings s, ForwardPass target, ForwardPass draft, GgufTokenizer tok, SamplingParams sp) { - var prompt = FormatPrompt(s.Prompt!, s.SystemPrompt, enableThinking: !s.NoThinking); + var prompt = FormatPrompt(s.Prompt!, s.SystemPrompt, enableThinking: !s_noThinking); var tokens = tok.Encode(prompt); if (!s.NoDisplayPrompt) @@ -726,7 +737,7 @@ private static int RunSpeculativeInteractive(Settings s, if (input is null or "/exit" or "/quit") break; if (string.IsNullOrWhiteSpace(input)) continue; - var prompt = FormatPrompt(input, s.SystemPrompt, enableThinking: !s.NoThinking); + var prompt = FormatPrompt(input, s.SystemPrompt, enableThinking: !s_noThinking); var tokens = tok.Encode(prompt); target.Cache.Reset(); @@ -775,7 +786,7 @@ private static int RunSinglePrompt(Settings s, GgufTokenizer tok, SamplingParams sp, Random rng, IForwardPass? mtpFwd) { - var prompt = FormatPrompt(s.Prompt!, s.SystemPrompt, enableThinking: !s.NoThinking); + var prompt = FormatPrompt(s.Prompt!, s.SystemPrompt, enableThinking: !s_noThinking); var tokens = tok.Encode(prompt); // SHARPI_RAW_PROMPT bypasses the chat template, so we need to add BOS @@ -808,7 +819,7 @@ private static int RunSinglePrompt(Settings s, // MTP head AND sampling is greedy AND the user disabled thinking on the chat // template (--no-thinking) AND sp.SpecType permits. SHARPI_DISABLE_MTP=1 is // a back-compat off-switch that wins. - bool useMtp = ResolveCliMtp(mtpFwd, sp, s.NoThinking, out string? mtpReject); + bool useMtp = ResolveCliMtp(mtpFwd, sp, s_noThinking, out string? mtpReject); if (mtpReject != null) { AnsiConsole.MarkupLine($"[red]Error:[/] {Markup.Escape(mtpReject)}"); @@ -956,7 +967,7 @@ private static int RunInteractive(Settings s, if (input is null or "/exit" or "/quit") break; if (string.IsNullOrWhiteSpace(input)) continue; - var prompt = FormatPrompt(input, s.SystemPrompt, enableThinking: !s.NoThinking); + var prompt = FormatPrompt(input, s.SystemPrompt, enableThinking: !s_noThinking); var tokens = tok.Encode(prompt); resetCache(); @@ -1072,6 +1083,9 @@ private static bool EmitToken(int next, GgufTokenizer tok, Utf8StreamDecoder str } private static string s_arch = "qwen2"; // set during model load + // Effective "thinking off" state: --no-thinking OR a model whose recommended config + // disables reasoning (Gemma 4 E4B-it is not a reasoning model). Set during model load. + private static bool s_noThinking; private static int s_thinkTokenId = -1; // token for any model using the / special-token convention private static int s_endThinkTokenId = -1; // token for any model using the / special-token convention private static JinjaChatTemplate? s_jinja; // parsed from GGUF tokenizer.chat_template diff --git a/src/SharpInference.Cuda/CudaBackend.cs b/src/SharpInference.Cuda/CudaBackend.cs index 0ba2814..2644c1e 100644 --- a/src/SharpInference.Cuda/CudaBackend.cs +++ b/src/SharpInference.Cuda/CudaBackend.cs @@ -126,6 +126,7 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis private nint _embedLookupQ4KKernel; private nint _embedLookupQ5KKernel; private nint _embedLookupQ80Kernel; + private nint _embedLookupQ80BatchedKernel; private nint _matvecF32Kernel; private nint _matvecQ4KKernel; private nint _matvecQ5KKernel; @@ -134,6 +135,8 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis // the GPU. Without this, Q8_0 weights would dequant to F32 on upload and // blow out VRAM ~2.1×. Geometry mirrors Q5_K/Q6_K (8 rows/block × 32 thr/row). private nint _matvecQ80Kernel; + // Issue #142: dp4a/Q8_1 decode matvec (quantize activation to int8, __dp4a dot). + private nint _matvecQ80Dp4aKernel; // Issue #43: N=2 (two-input, two-output) variants — read each weight row // once and accumulate into two outputs. Used by MTP BatchForward2's // on-GPU dense FFN to halve weight-bandwidth cost per output. @@ -151,6 +154,13 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis private nint _matvecQ5KGemmNKernel; private nint _matvecQ6KGemmNKernel; private nint _matvecQ80GemmNKernel; + // Issue #141: compute-bound prefill GEMM — dequant Q8_0 weight + convert + // activations to fp16, then one cublasGemmEx (weight read once per batch). + private nint _dequantQ80F16Kernel; + private nint _f32ToF16Kernel; + // Issue #141 (MMQ): int8 tensor-core Q8_0×Q8_1 matmul — weight read once as + // int8, no fp16 HBM round-trip, m16n8k32 s8 mma. Replaces the dequant→GEMM path. + private nint _mmqQ80Kernel; // Issue #111: batched trunk elementwise/norm kernels (one launch over N tokens). private nint _rmsNormBatchedKernel; private nint _headNormBatchedKernel; @@ -200,6 +210,10 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis private nint _fullSeqAttentionBf16Kernel; private nint _fullSeqAttentionGlobalKernel; private nint _fullSeqAttentionGlobalBf16Kernel; + // Issue #141 (attention): memory-efficient flash-attention prefill (shared K/V + // tiles reused across a query tile + online softmax) replacing the scalar + // full_seq / swa_batched kernels' O(n²) global K/V re-reads. + private nint _flashAttnPrefillKernel; // Grow-only global score scratch for the wave-based >4096 batched-query SDPA // (issue #118). Sized W × num_heads × score_stride floats; W is chosen so this @@ -220,6 +234,12 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis // N contiguous q81 rows). Grow-only, same policy as _q81Buf. private nint _q81BatchBuf; private nuint _q81BatchBufSize; + // Issue #141: fp16 scratch for the prefill GEMM — dequantized weight and + // converted activations. Grow-only; sized to the largest trunk matmul. + private nint _gemmWf16Buf; + private nuint _gemmWf16Size; + private nint _gemmAf16Buf; + private nuint _gemmAf16Size; // Tracks dtype per tensor handle so MatMul can dispatch to the right matvec variant // (Q4_K / Q5_K / Q6_K / F32). Norm/bias weights upload as F32; quantized weight bytes @@ -232,6 +252,16 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis public bool SupportsGpuDequant => false; + /// + /// Issue #142: route Q8_0 decode matvecs through the dp4a/Q8_1 kernel + /// () instead of the fp32-decode kernel. + /// Default on (SHARPI_Q80_DP4A); the dp4a path quantizes the activation + /// to int8 so it is argmax-stable, not bit-exact to the fp32 matvec. Settable so + /// bit-parity oracles can pin both sides to the fp32 path. + /// + public bool Q80Dp4aEnabled { get; set; } = + Environment.GetEnvironmentVariable("SHARPI_Q80_DP4A") != "0"; + /// Total VRAM on the active CUDA device, in bytes. Queried once at backend creation. public ulong VramBytes { @@ -1472,6 +1502,13 @@ public void MatMul(Tensor output, Tensor matrix, Tensor vector, DType weightDTyp DispatchMatVecQ4K(wPtr, xPtr, yPtr, rows, cols); return; } + // Issue #142: Q8_0 decode matvec via dp4a/Q8_1 (cols % 32 == 0 — every LLM + // hidden dim qualifies). Falls back to the fp32-decode kernel otherwise. + if (weightDType == DType.Q8_0 && Q80Dp4aEnabled && (cols & 31) == 0) + { + DispatchMatVecQ80Dp4a(wPtr, xPtr, yPtr, rows, cols); + return; + } int pRows = rows, pCols = cols; nint* args = stackalloc nint[5] @@ -1650,6 +1687,197 @@ public void MatMulBatched(Tensor outputAll, Tensor matrix, Tensor inputAll, int MatMulBatched(outputAll, matrix, inputAll, nTok, dtype); } + /// + /// Compute-bound batched matmul for prefill (issue #141). Dequantizes the + /// Q8_0 [rows×cols] to an fp16 scratch once, + /// converts the activation vectors to fp16, then runs a + /// single cublasGemmEx (fp16×fp16 → fp32, fp32 accumulate). The + /// matvec GEMM-N re-streams the weight matrix once + /// per token (memory-bound); this reads each weight once per ~nTok batch, so on + /// a compute-rich GPU prefill becomes compute-bound — the ~70× pp512 gap vs + /// llama.cpp the matvec path could never close. + /// + /// NOT bit-exact to the matvec path: the weight value d*q and + /// each activation are rounded to fp16 before the tensor-core multiply (fp32 + /// accumulation). Result tracks the fp32 path to fp tolerance (argmax-stable), + /// not byte-for-byte. Callers that need byte-parity (GDN/MTP draft verify) must + /// use . Q8_0 weights only — other dtypes throw. + /// + /// Layout matches : is + /// token-major [nTok × cols], is + /// [nTok × rows]. + /// + public void MatMulBatchedGemm(Tensor outputAll, Tensor matrix, Tensor inputAll, + int nTok, DType weightDType) + { + EnsureImageKernels(); + if (!_imageKernelsAvailable) + throw new NotSupportedException("NVRTC kernels are not available on this system."); + if (weightDType != DType.Q8_0) + throw new NotSupportedException( + $"CUDA MatMulBatchedGemm: weight dtype {weightDType} not supported (Q8_0 only)."); + if (nTok <= 0) + throw new ArgumentOutOfRangeException(nameof(nTok), nTok, "nTok must be > 0."); + if (outputAll.ElementCount % nTok != 0 || inputAll.ElementCount % nTok != 0) + throw new ArgumentException( + $"MatMulBatchedGemm: outputAll ({outputAll.ElementCount}) and inputAll " + + $"({inputAll.ElementCount}) element counts must be divisible by nTok ({nTok})."); + + int rows = (int)(outputAll.ElementCount / nTok); + int cols = (int)(inputAll.ElementCount / nTok); + if ((cols & 31) != 0) + throw new InvalidOperationException( + $"CUDA MatMulBatchedGemm requires cols % 32 == 0 (got {cols})."); + + nint wPtr = GetDevPtr(matrix); + nint xPtr = GetDevPtr(inputAll); + nint yPtr = GetDevPtr(outputAll); + + EnsureGemmWf16((nuint)((long)rows * cols * 2L)); + EnsureGemmAf16((nuint)((long)nTok * cols * 2L)); + + // 1) Dequant Q8_0 weight → fp16 (one block per row). + { + nint wp = wPtr, op = _gemmWf16Buf; + int pRows = rows, pCols = cols; + nint* args = stackalloc nint[4] { (nint)(&wp), (nint)(&op), (nint)(&pRows), (nint)(&pCols) }; + int r = NvrtcInterop.LaunchKernel(_dequantQ80F16Kernel, (uint)rows, 1, 1, + 256, 1, 1, 0, _stream, args, null); + if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(dequant_q8_0_to_f16) failed: {r}"); + } + + // 2) Convert activations fp32 → fp16. + { + nint ip = xPtr, op = _gemmAf16Buf; + int n = nTok * cols; + nint* args = stackalloc nint[3] { (nint)(&ip), (nint)(&op), (nint)(&n) }; + uint grid = (uint)((n + 255) / 256); + int r = NvrtcInterop.LaunchKernel(_f32ToF16Kernel, grid, 1, 1, + 256, 1, 1, 0, _stream, args, null); + if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(f32_to_f16) failed: {r}"); + } + + // 3) GemmEx: C[nTok×rows] f32 = A[nTok×cols] f16 × B[rows×cols] f16ᵀ, fp32 accum. + // Row-major via the col-major transpose identity (mirrors Sgemm): + // row-major C=A·Bᵀ ≡ col-major Cᵀ = B·Aᵀ. + { + float alpha = 1f, beta = 0f; + int M = nTok, K = cols, N = rows; + int status = CuBlasInterop.GemmEx( + _handle, + CuBlasInterop.OpT, CuBlasInterop.OpN, + N, M, K, + ref alpha, + _gemmWf16Buf, CuBlasInterop.CUDA_R_16F, K, + _gemmAf16Buf, CuBlasInterop.CUDA_R_16F, K, + ref beta, + yPtr, CuBlasInterop.CUDA_R_32F, N, + CuBlasInterop.CUBLAS_COMPUTE_32F, + CuBlasInterop.CUBLAS_GEMM_DEFAULT); + if (status != 0) + throw new InvalidOperationException($"cublasGemmEx (prefill GEMM) failed: {status}"); + } + } + + /// + /// Issue #141 (MMQ): int8 tensor-core batched matmul for Q8_0 weights. + /// C[nTok×rows] f32 = X[nTok×cols] · Wᵀ where W is Q8_0 [rows×cols]. The input is + /// quantized to Q8_1 (per-32-block int8 + fp16 scale) and multiplied by the Q8_0 + /// weight via the m16n8k32 s8 mma — the weight is read once as int8, with no fp16 + /// dequant temp written to HBM (the cost that capped ). + /// Argmax-stable, not bit-exact (both operands are int8-quantized). + /// + public void MatMulBatchedMmq(Tensor outputAll, Tensor matrix, Tensor inputAll, + int nTok, DType weightDType) + { + EnsureImageKernels(); + if (!_imageKernelsAvailable) + throw new NotSupportedException("NVRTC kernels are not available on this system."); + if (weightDType != DType.Q8_0) + throw new NotSupportedException( + $"CUDA MatMulBatchedMmq: weight dtype {weightDType} not supported (Q8_0 only)."); + if (nTok <= 0) + throw new ArgumentOutOfRangeException(nameof(nTok), nTok, "nTok must be > 0."); + if (outputAll.ElementCount % nTok != 0 || inputAll.ElementCount % nTok != 0) + throw new ArgumentException( + $"MatMulBatchedMmq: outputAll ({outputAll.ElementCount}) and inputAll " + + $"({inputAll.ElementCount}) element counts must be divisible by nTok ({nTok})."); + + int rows = (int)(outputAll.ElementCount / nTok); + int cols = (int)(inputAll.ElementCount / nTok); + if ((cols & 31) != 0) + throw new InvalidOperationException( + $"CUDA MatMulBatchedMmq requires cols % 32 == 0 (got {cols})."); + + nint wPtr = GetDevPtr(matrix); + nint xPtr = GetDevPtr(inputAll); + nint yPtr = GetDevPtr(outputAll); + + // 1) Quantize activations [nTok×cols] f32 → contiguous Q8_1 (36 B/block). The + // per-block quantize is independent, so a single launch over nTok×subBlocks + // is bit-identical to per-token quantization (mirrors DispatchMatVecQ4KBatched). + int subBlocks = cols / 32; + long totalSub = (long)subBlocks * nTok; + EnsureQ81BatchBuf((nuint)(totalSub * 36L)); + { + nint qIn = xPtr, qOut = _q81BatchBuf; + int qN = (int)((long)cols * nTok); + if ((long)cols * nTok > int.MaxValue) + throw new InvalidOperationException( + $"MatMulBatchedMmq: cols*nTok ({(long)cols * nTok}) exceeds int range."); + nint* args = stackalloc nint[3] { (nint)(&qIn), (nint)(&qOut), (nint)(&qN) }; + int rq = NvrtcInterop.LaunchKernel(_quantizeQ81Kernel, (uint)totalSub, 1, 1, + 32, 1, 1, 0, _stream, args, null); + if (rq != 0) throw new InvalidOperationException($"cuLaunchKernel(quantize_q8_1 mmq) failed: {rq}"); + } + + // 2) int8 mma MMQ: grid = ((rows+63)/64, (nTok+127)/128), 256 threads/block, + // each block a 64×128 output tile (8 warps × 8 m16n8k32 mma per K-block). + { + nint q81 = _q81BatchBuf; + int pRows = rows, pCols = cols, pN = nTok; + nint* args = stackalloc nint[6] + { + (nint)(&wPtr), (nint)(&q81), (nint)(&yPtr), + (nint)(&pRows), (nint)(&pCols), (nint)(&pN) + }; + uint gx = (uint)((rows + 63) / 64), gy = (uint)((nTok + 127) / 128); + int rm = NvrtcInterop.LaunchKernel(_mmqQ80Kernel, gx, gy, 1, + 256, 1, 1, 0, _stream, args, null); + if (rm != 0) throw new InvalidOperationException($"cuLaunchKernel(mmq_q8_0) failed: {rm}"); + } + } + + private void EnsureGemmWf16(nuint required) + { + if (_gemmWf16Buf != nint.Zero && _gemmWf16Size >= required) return; + if (_gemmWf16Buf != nint.Zero) + { + CuBlasInterop.CudaFree(_gemmWf16Buf); + _gemmWf16Buf = nint.Zero; + _gemmWf16Size = 0; + } + nuint newSize = (required + 0xffffu) & ~(nuint)0xffffu; + int r = CuBlasInterop.CudaMalloc(out _gemmWf16Buf, newSize); + if (r != 0) throw new InvalidOperationException($"cudaMalloc(gemm wf16, {newSize} B) failed: {r}"); + _gemmWf16Size = newSize; + } + + private void EnsureGemmAf16(nuint required) + { + if (_gemmAf16Buf != nint.Zero && _gemmAf16Size >= required) return; + if (_gemmAf16Buf != nint.Zero) + { + CuBlasInterop.CudaFree(_gemmAf16Buf); + _gemmAf16Buf = nint.Zero; + _gemmAf16Size = 0; + } + nuint newSize = (required + 0xffffu) & ~(nuint)0xffffu; + int r = CuBlasInterop.CudaMalloc(out _gemmAf16Buf, newSize); + if (r != 0) throw new InvalidOperationException($"cudaMalloc(gemm af16, {newSize} B) failed: {r}"); + _gemmAf16Size = newSize; + } + /// /// Q4_K batched GEMM-N: quantizes all input vectors into a /// single contiguous Q8_1 scratch (one launch over nTok × subBlocks blocks — @@ -1826,6 +2054,64 @@ private void DispatchMatVecQ4K(nint wPtr, nint xPtr, nint yPtr, int rows, int co } } + /// + /// Q8_0 decode matvec via dp4a (issue #142): quantize the input vector to Q8_1 + /// (32-element sub-blocks), then dispatch llm_matvec_q8_0_dp4a (1 row / + /// block, MATVEC_Q80_NWARPS warps). The int8·int8 dp4a inner product replaces + /// the per-element int8→float decode of the fp32 matvec, cutting instruction + /// count on the bandwidth-bound decode path. Requires cols % 32 == 0. + /// + /// + /// Pre-grow the Q8_1 input-quantization scratch to hold a -wide + /// vector (issue #142). Call before a CUDA-graph capture region that contains a dp4a + /// matvec: capture forbids cudaMalloc, so the buffer must already be at its max + /// size. No-op if already large enough. + /// + public void EnsureQ81Scratch(int cols) + { + if (cols <= 0) return; + int subBlocks = (cols + 31) / 32; + EnsureQ81Buf((nuint)((long)subBlocks * 36L)); + } + + private void DispatchMatVecQ80Dp4a(nint wPtr, nint xPtr, nint yPtr, int rows, int cols) + { + if ((cols & 31) != 0) + throw new InvalidOperationException( + $"CUDA matvec_q8_0_dp4a requires cols % 32 == 0 (got {cols})."); + + int subBlocks = cols / 32; + nuint q81Bytes = (nuint)((long)subBlocks * 36L); + EnsureQ81Buf(q81Bytes); + + // Quantize input → Q8_1 (32 threads per sub-block). + { + nint qInPtr = xPtr; + nint qOutPtr = _q81Buf; + int qN = cols; + nint* args = stackalloc nint[3] { (nint)(&qInPtr), (nint)(&qOutPtr), (nint)(&qN) }; + int rq = NvrtcInterop.LaunchKernel( + _quantizeQ81Kernel, (uint)subBlocks, 1, 1, + 32, 1, 1, 0, _stream, args, null); + if (rq != 0) throw new InvalidOperationException($"cuLaunchKernel(quantize_q8_1 for q8_0) failed: {rq}"); + } + + // Cooperative dp4a matvec: 1 row/block, MATVEC_Q80_NWARPS warps × 32 threads. + { + nint q81Ptr = _q81Buf; + int pRows = rows, pCols = cols; + nint* args = stackalloc nint[5] + { + (nint)(&wPtr), (nint)(&q81Ptr), (nint)(&yPtr), + (nint)(&pRows), (nint)(&pCols) + }; + int rm = NvrtcInterop.LaunchKernel( + _matvecQ80Dp4aKernel, (uint)rows, 1, 1, + 32, 8, 1, 0, _stream, args, null); + if (rm != 0) throw new InvalidOperationException($"cuLaunchKernel(matvec_q8_0_dp4a) failed: {rm}"); + } + } + private void EnsureQ81Buf(nuint required) { if (_q81Buf != nint.Zero && _q81BufSize >= required) return; @@ -2612,6 +2898,55 @@ public void AttentionSwaBatched(Tensor qAll, Tensor kCache, Tensor vCache, Tenso if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(attention_swa_batched) failed: {r}"); } + /// + /// Issue #141 (attention): memory-efficient flash-attention prefill. Replaces + /// (global, =0) and + /// (sliding window, windowSize>0) for the + /// fp32-KV batched-trunk prefill. A block handles a tile of 8 queries of one head + /// and streams K/V through shared-memory tiles with an online softmax, so each + /// key is read from global once per 8 queries instead of once per query — cutting + /// the scalar kernels' O(n²) (SWA: up to ~512×) redundant K/V traffic. GQA, causal, + /// optional sliding window, per-layer . Matches the scalar + /// kernels to fp tolerance (online softmax), not bit-exact. + /// + public void FlashAttentionPrefill(Tensor qAll, Tensor kCache, Tensor vCache, Tensor outAll, + int numHeads, int numKvHeads, int headDim, + int startPos, int windowSize, int maxSeqLen, int nTok, float attnScale = -1f) + { + EnsureImageKernels(); + if (!_imageKernelsAvailable) + throw new NotSupportedException("NVRTC kernels are not available."); + if (nTok <= 0) + throw new ArgumentOutOfRangeException(nameof(nTok), nTok, "nTok must be > 0."); + if (headDim > 512) + throw new NotSupportedException( + $"FlashAttentionPrefill supports head_dim ≤ 512 (16 dims/lane); got {headDim}."); + + // Pick the streaming K-tile so the per-key shared tile (K fp16 = headDim*2 B + + // V fp32 = headDim*4 B = 6*headDim B) fits a 48 KB budget (no >48 KB opt-in). + const int sharedBudget = 48 * 1024; + int ktTile = sharedBudget / (6 * headDim); + if (ktTile < 1) ktTile = 1; + if (ktTile > 32) ktTile = 32; + uint sharedBytes = (uint)(6 * ktTile * headDim); + + nint qP = GetDevPtr(qAll), kP = GetDevPtr(kCache), vP = GetDevPtr(vCache), oP = GetDevPtr(outAll); + int pNH = numHeads, pNKV = numKvHeads, pHD = headDim; + int pSP = startPos, pWS = windowSize, pMSL = maxSeqLen, pN = nTok, pKT = ktTile; + float pScale = attnScale; + nint* args = stackalloc nint[13] + { + (nint)(&qP), (nint)(&kP), (nint)(&vP), (nint)(&oP), + (nint)(&pNH), (nint)(&pNKV), (nint)(&pHD), + (nint)(&pSP), (nint)(&pWS), (nint)(&pMSL), (nint)(&pN), (nint)(&pKT), (nint)(&pScale) + }; + const int faQt = 16; // FA_QT in the kernel (warps/block = K/V reuse factor) + uint gy = (uint)((nTok + faQt - 1) / faQt); + int r = NvrtcInterop.LaunchKernel(_flashAttnPrefillKernel, (uint)numHeads, gy, 1, + (uint)(faQt * 32), 1, 1, sharedBytes, _stream, args, null); + if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(flash_attn_prefill) failed: {r}"); + } + /// /// Bf16-store variant of . Inputs stay fp32; the K/V /// cache tensors must be -allocated (half the @@ -3236,6 +3571,35 @@ public void EmbedLookupQ8_0(Tensor embTable, Tensor output, int tokenId, int emb if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(embed_lookup_q8_0) failed: {r}"); } + /// + /// Batched Q8_0 embedding lookup: writes all token rows into + /// ([nTok × embDim]) in one launch, reading the per-token + /// ids from the device buffer (nTok int32). Collapses the + /// prefill's per-token + device copy (2·N host launches) + /// into a single grid.x = nTok launch. Bit-identical to the per-token path. + /// + public void EmbedLookupQ8_0Batched(Tensor embTable, Tensor outputAll, Tensor tokenIds, int nTok, int embDim) + { + EnsureImageKernels(); + if (!_imageKernelsAvailable) + throw new NotSupportedException("NVRTC kernels are not available."); + if ((embDim & 0xff) != 0) + throw new ArgumentException($"EmbedLookupQ8_0Batched requires embDim to be a multiple of 256 (got {embDim})."); + if (nTok <= 0) + throw new ArgumentOutOfRangeException(nameof(nTok), nTok, "nTok must be > 0."); + + nint tP = GetDevPtr(embTable); + nint oP = GetDevPtr(outputAll); + nint idP = GetDevPtr(tokenIds); + int pN = nTok, pE = embDim; + nint* args = stackalloc nint[5] + { + (nint)(&tP), (nint)(&oP), (nint)(&idP), (nint)(&pN), (nint)(&pE) + }; + int r = NvrtcInterop.LaunchKernel(_embedLookupQ80BatchedKernel, (uint)nTok, 1, 1, 256, 1, 1, 0, _stream, args, null); + if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(embed_lookup_q8_0_batched) failed: {r}"); + } + /// Set every element of to zero. /// /// The kernel writes fp32-sized lanes; for sub-fp32 dtypes (e.g. BFloat16) @@ -3849,12 +4213,12 @@ private void ForceEagerJit() _snapKvScoreKernel, _snapKvScoreBf16Kernel, _kvCompactKernel, _kvCompactBf16Kernel, _embedLookupF32Kernel, _embedLookupQ4KKernel, _embedLookupQ5KKernel, - _embedLookupQ80Kernel, + _embedLookupQ80Kernel, _embedLookupQ80BatchedKernel, _matvecF32Kernel, _matvecQ4KKernel, _matvecQ5KKernel, _matvecQ6KKernel, _matvecQ80Kernel, _matvecF32N2Kernel, _matvecQ4KN2Kernel, _matvecQ5KN2Kernel, _matvecQ6KN2Kernel, _matvecF32GemmNKernel, _matvecQ4KGemmNKernel, _matvecQ5KGemmNKernel, _matvecQ6KGemmNKernel, - _matvecQ80GemmNKernel, + _matvecQ80GemmNKernel, _mmqQ80Kernel, _rmsNormBatchedKernel, _headNormBatchedKernel, _headNormQkKernel, _headNormQkBatchedKernel, _splitQgBatchedKernel, _ropeNeoxPartialBatchedKernel, _attentionKernel, _attentionBf16Kernel, _attentionSwaKernel, _attentionSwaBatchedKernel, @@ -3869,6 +4233,7 @@ private void ForceEagerJit() _kvAppendBatchedKernel, _kvAppendBatchedBf16Kernel, _fullSeqAttentionKernel, _fullSeqAttentionBf16Kernel, _fullSeqAttentionGlobalKernel, _fullSeqAttentionGlobalBf16Kernel, + _flashAttnPrefillKernel, ]; foreach (nint k in kernels) { @@ -3914,11 +4279,13 @@ private void LoadKernelFunctions() _embedLookupQ4KKernel = GetKernelFunc("llm_embed_lookup_q4k"); _embedLookupQ5KKernel = GetKernelFunc("llm_embed_lookup_q5k"); _embedLookupQ80Kernel = GetKernelFunc("llm_embed_lookup_q8_0"); + _embedLookupQ80BatchedKernel = GetKernelFunc("llm_embed_lookup_q8_0_batched"); _matvecF32Kernel = GetKernelFunc("llm_matvec_f32"); _matvecQ4KKernel = GetKernelFunc("llm_matvec_q4k"); _matvecQ5KKernel = GetKernelFunc("llm_matvec_q5k"); _matvecQ6KKernel = GetKernelFunc("llm_matvec_q6k"); _matvecQ80Kernel = GetKernelFunc("llm_matvec_q8_0"); + _matvecQ80Dp4aKernel = GetKernelFunc("llm_matvec_q8_0_dp4a"); _matvecF32N2Kernel = GetKernelFunc("llm_matvec_f32_n2"); _matvecQ4KN2Kernel = GetKernelFunc("llm_matvec_q4k_n2"); _matvecQ5KN2Kernel = GetKernelFunc("llm_matvec_q5k_n2"); @@ -3928,6 +4295,10 @@ private void LoadKernelFunctions() _matvecQ5KGemmNKernel = GetKernelFunc("llm_matvec_q5k_gemm_n"); _matvecQ6KGemmNKernel = GetKernelFunc("llm_matvec_q6k_gemm_n"); _matvecQ80GemmNKernel = GetKernelFunc("llm_matvec_q8_0_gemm_n"); + _dequantQ80F16Kernel = GetKernelFunc("llm_dequant_q8_0_to_f16"); + _f32ToF16Kernel = GetKernelFunc("llm_f32_to_f16"); + _mmqQ80Kernel = GetKernelFunc("llm_mmq_q8_0"); + _flashAttnPrefillKernel = GetKernelFunc("llm_flash_attn_prefill_f32"); _rmsNormBatchedKernel = GetKernelFunc("llm_rmsnorm_batched"); _headNormBatchedKernel = GetKernelFunc("llm_head_norm_batched"); _headNormQkKernel = GetKernelFunc("llm_head_norm_qk"); @@ -4383,6 +4754,18 @@ public void Dispose() _q81BatchBuf = nint.Zero; _q81BatchBufSize = 0; } + if (_gemmWf16Buf != nint.Zero) + { + CuBlasInterop.CudaFree(_gemmWf16Buf); + _gemmWf16Buf = nint.Zero; + _gemmWf16Size = 0; + } + if (_gemmAf16Buf != nint.Zero) + { + CuBlasInterop.CudaFree(_gemmAf16Buf); + _gemmAf16Buf = nint.Zero; + _gemmAf16Size = 0; + } if (_waveScratchBuf != nint.Zero) { CuBlasInterop.CudaFree(_waveScratchBuf); diff --git a/src/SharpInference.Cuda/CudaTextKernels.cs b/src/SharpInference.Cuda/CudaTextKernels.cs index 3f2868a..2b74e70 100644 --- a/src/SharpInference.Cuda/CudaTextKernels.cs +++ b/src/SharpInference.Cuda/CudaTextKernels.cs @@ -32,6 +32,35 @@ __device__ __forceinline__ float sharpi_fp16_to_fp32(unsigned int h) return result; } +// ── half2 (fp16x2) pack / fma / unpack via inline PTX ────────────────────── +// NVRTC compiles this source without cuda_fp16.h, so the __half2 intrinsics are +// unavailable; these wrap the raw PTX. A packed fp16x2 lives in one unsigned int +// (low half = first element). Used by the flash-attention QK dot to do 2 multiply- +// accumulates per instruction (FP16x2 path, ~2× the fp32 FMA rate) — the scores +// tolerate fp16-rounded inputs (argmax-stable), and the running sum over only a few +// pairs per lane keeps fp16 accumulation error negligible (final reduce is fp32). +__device__ __forceinline__ unsigned int sharpi_f32x2_to_f16x2(float a, float b) +{ + unsigned int r; + asm(""{ .reg .b16 lo, hi; cvt.rn.f16.f32 lo, %1; cvt.rn.f16.f32 hi, %2; mov.b32 %0, {lo, hi}; }"" + : ""=r""(r) : ""f""(a), ""f""(b)); + return r; +} +__device__ __forceinline__ unsigned int sharpi_hfma2(unsigned int a, unsigned int b, unsigned int acc) +{ + unsigned int r; + asm(""fma.rn.f16x2 %0, %1, %2, %3;"" : ""=r""(r) : ""r""(a), ""r""(b), ""r""(acc)); + return r; +} +// Sum the low and high fp16 halves of a packed fp16x2 into one fp32. +__device__ __forceinline__ float sharpi_f16x2_sum(unsigned int p) +{ + float lo, hi; + asm(""{ .reg .b16 l, h; mov.b32 {l, h}, %2; cvt.f32.f16 %0, l; cvt.f32.f16 %1, h; }"" + : ""=f""(lo), ""=f""(hi) : ""r""(p)); + return lo + hi; +} + // Warp-level sum reduction over 32 lanes (full-warp mask). __device__ __forceinline__ float sharpi_warp_reduce_sum(float v) { @@ -90,6 +119,21 @@ __device__ __forceinline__ int sharpi_int8_at(const unsigned int* __restrict__ b return v >= 128 ? v - 256 : v; } +// Read a 4-byte little-endian word from a uint32-stride buffer at an arbitrary +// (possibly 2-byte-aligned) byte offset B. Q8_0 qs start at a 2-byte offset inside +// each 34-byte block, so a raw int load would be misaligned; this assembles the +// word from the two covering aligned uints via funnelshift (1 extra load only when +// B is not 4-aligned). Mirrors llama.cpp's get_int_b2 (vecdotq.cuh). +__device__ __forceinline__ unsigned int sharpi_uint_at(const unsigned int* __restrict__ buf, long B) +{ + long aB = B & ~3L; + unsigned int sh = (unsigned int)(B & 3L) * 8u; + unsigned int lo = buf[aB >> 2]; + if (sh == 0u) return lo; + unsigned int hi = buf[(aB >> 2) + 1]; + return __funnelshift_r(lo, hi, sh); +} + // ── RmsNorm ──────────────────────────────────────────────────────────────── // output[i] = input[i] / rms * weight[i], rms = sqrt(mean(input^2) + eps). // 1 block of 256 threads. Push-equivalent: (n, eps). @@ -785,6 +829,46 @@ __device__ __forceinline__ int sharpi_int8_at(const unsigned int* __restrict__ b } } +// Batched Q8_0 embedding lookup: one block per query token (grid.x = n_tok), +// reading token_ids[blockIdx.x] and writing row blockIdx.x of output. Collapses +// the prefill's per-token EmbedLookupQ8_0 + copy (2·N host launches) into one +// launch — the per-token body is identical to llm_embed_lookup_q8_0. +extern ""C"" __global__ void llm_embed_lookup_q8_0_batched( + const unsigned char* __restrict__ emb_data, + float* __restrict__ output, // [n_tok * emb_dim] + const int* __restrict__ token_ids, // [n_tok] + int n_tok, int emb_dim) +{ + __shared__ unsigned char blk[272]; + unsigned int tid = threadIdx.x; + int i = (int)blockIdx.x; + if (i >= n_tok) return; + + int token_id = token_ids[i]; + float* out_row = output + (long)i * emb_dim; + int num_blocks = emb_dim >> 5; + long bytes_per_row = (long)num_blocks * 34L; + long row_byte_base = (long)token_id * bytes_per_row; + + int outer_iters = emb_dim >> 8; + for (int outer = 0; outer < outer_iters; outer++) { + long base_byte = row_byte_base + (long)(outer * 8) * 34L; + if (tid < 272) blk[tid] = emb_data[base_byte + tid]; + if (tid < 16) blk[256 + tid] = emb_data[base_byte + 256 + tid]; + __syncthreads(); + + unsigned int block_in_outer = tid >> 5; + unsigned int lane = tid & 31u; + unsigned int block_off = block_in_outer * 34u; + unsigned int d_bits = (unsigned int)blk[block_off] + | ((unsigned int)blk[block_off + 1u] << 8); + float d = sharpi_fp16_to_fp32(d_bits); + int q = (int)(signed char)blk[block_off + 2u + lane]; + out_row[outer * 256 + (int)tid] = d * (float)q; + __syncthreads(); + } +} + // ── MatVec F32 ───────────────────────────────────────────────────────────── // 256 threads/block, 8 rows/block, 32 threads/row → warp reduce. // One grid dim x covers ceil(rows/8) blocks. @@ -1215,6 +1299,236 @@ __device__ __forceinline__ int sharpi_int8_at(const unsigned int* __restrict__ b if (lane == 0) output[row] = result; } +// ── MatVec Q8_0 — __dp4a / Q8_1 path (issue #142) ───────────────────────── +// Decode matvec mirroring llama.cpp's mul_mat_vec_q8_0_q8_1. The input vector is +// pre-quantized to Q8_1 (36-byte sub-blocks: fp16 d at [0:2], 32 int8 at [4:36]), +// so each 4-int8 inner product is one __dp4a instruction instead of four +// int8→float converts + fp32 FMAs — far fewer instructions per byte of weight, +// pushing the (already memory-coalesced) Q8_0 matvec closer to the HBM ceiling. +// +// One output row per block; MATVEC_Q80_NWARPS warps cooperate. Within a warp the +// 32 lanes split into 4 groups of 8: group g = lane>>3 owns one Q8_0 block of the +// warp's 4-block stripe; sub-lane s = lane&7 handles int32 word s (4 int8) of that +// block via one __dp4a. The 8 sub-lane partials are summed (shfl_xor over the +// group of 8) and scaled by d_w·d_a, then accumulated across the stripe. +// +// Q8_0 block = 34 bytes (fp16 d + 32 int8); qs is only 2-byte aligned, so the +// 4-int8 weight words are assembled with __funnelshift_r from two aligned uint +// loads (the activation Q8_1 qs at +4 is naturally 4-aligned). +#define MATVEC_Q80_NWARPS 8 +extern ""C"" __global__ void llm_matvec_q8_0_dp4a( + const unsigned int* __restrict__ weights, + const unsigned char* __restrict__ y_q81, + float* __restrict__ output, + int rows, int cols) +{ + int row = (int)blockIdx.x; + if (row >= rows) return; + int warp_id = (int)threadIdx.y; // 0..NWARPS-1 + int lane = (int)threadIdx.x; // 0..31 + int grp = lane >> 3; // 0..3 block within the warp's 4-block stripe + int sub = lane & 7; // 0..7 int32 word within the block + + int num_blocks = cols >> 5; // cols / 32 + long row_base_bytes = (long)row * (long)num_blocks * 34L; + + float acc = 0.f; + + for (int block0 = warp_id * 4; block0 < num_blocks; block0 += MATVEC_Q80_NWARPS * 4) { + int block = block0 + grp; + float part = 0.f; + if (block < num_blocks) { + long b0 = row_base_bytes + (long)block * 34L; + unsigned int dlo = sharpi_byte_at(weights, b0); + unsigned int dhi = sharpi_byte_at(weights, b0 + 1); + float dw = sharpi_fp16_to_fp32(dlo | (dhi << 8)); + + // This sub-lane's 4 weight int8 = qs[sub*4 .. sub*4+4) at byte b0+2+sub*4. + long wb = b0 + 2 + (long)sub * 4; + long aligned = wb & ~3L; + unsigned int shift = (unsigned int)(wb & 3L) * 8u; + unsigned int w_lo = weights[aligned >> 2]; + int wq; + if (shift == 0u) wq = (int)w_lo; + else { + unsigned int w_hi = weights[(aligned >> 2) + 1]; + wq = (int)__funnelshift_r(w_lo, w_hi, shift); + } + + // Activation Q8_1 sub-block = block; d at base (low 16), 32 int8 at +4. + long ab = (long)block * 36L; + unsigned int d_bits = (*reinterpret_cast(y_q81 + ab)) & 0xffffu; + float da = sharpi_fp16_to_fp32(d_bits); + int aq = *reinterpret_cast(y_q81 + ab + 4 + (long)sub * 4); + + int dot = __dp4a(wq, aq, 0); + part = dw * da * (float)dot; + } + // Sum the 8 sub-lanes within each aligned group of 8. + part += __shfl_xor_sync(0xffffffffu, part, 4); + part += __shfl_xor_sync(0xffffffffu, part, 2); + part += __shfl_xor_sync(0xffffffffu, part, 1); + if (sub == 0) acc += part; + } + + // Group leaders (sub==0: lanes 0,8,16,24) hold per-stripe sums; reduce across + // the 4 groups and all warps via shared memory. + __shared__ float warp_acc[MATVEC_Q80_NWARPS][4]; + if (sub == 0) warp_acc[warp_id][grp] = acc; + __syncthreads(); + if (warp_id == 0 && lane == 0) { + float s = 0.f; + #pragma unroll + for (int w = 0; w < MATVEC_Q80_NWARPS; w++) + for (int g = 0; g < 4; g++) s += warp_acc[w][g]; + output[row] = s; + } +} + +// ── Q8_0 × Q8_1 int8 tensor-core MMQ (issue #141 prefill) ────────────────── +// Replaces the dequant→fp16→cuBLAS GEMM round-trip with a direct int8 tensor-core +// multiply: each Q8_0 weight is read once as int8 (no fp16 weight temp written to +// HBM) and fed to the m16n8k32 s8 mma. Activations are pre-quantized to the +// 36-byte/block Q8_1 layout (fp16 d at [0:2], 32 int8 at [4:36]) — the same buffer +// the dp4a decode matvec uses. result[r,t] = Σ_block ( int32(W_blk·Y_blk) · d_w · d_a ). +// Q8_0 is symmetric (scale only, no min), so — like llama.cpp's D4 layout — there +// is NO sum/bias correction term; the activation `s` field is never read. +// +// The mma fragment register layout follows the PTX m16n8k32.row.col spec exactly: +// groupID = lane>>2 (0..7), tig = lane&3 (0..3) +// A(16×32 s8): a0=W[grp][tig*4..], a1=W[grp+8][..], a2=W[grp][16+tig*4..], a3=W[grp+8][..] +// B(32×8 s8): b0=Y[tig*4..][grp], b1=Y[16+tig*4..][grp] (col=token=grp) +// C(16×8 s32): c0=[grp][tig*2], c1=[grp][tig*2+1], c2=[grp+8][tig*2], c3=[grp+8][tig*2+1] +// +// Shared-tiled: a 256-thread block computes a 64(row)×128(token) output tile, looping +// K in 32-wide Q8_0 blocks. Each K-block stages the weight + activation sub-tiles into +// shared once; 8 warps (4 row × 2 col) then run 8 m16n8k32 s8 mma's each (one per +// 8-token N-tile), accumulating per-block-scaled products in fp32 registers. Every +// weight row is read once per (row-block, K-block) and reused across all 128 tokens in +// the tile — the weight-read-once property the dequant->fp16->cuBLAS path bought with a +// 2x fp16 HBM temp, here without the temp and on int8 tensor cores (2x fp16 TC peak). +// The wide 128-token tile halves the weight re-read factor (nTok/MMQ_BN) vs a 64-token +// tile, which matters because the Q8_0 qs 2-byte misalignment taxes each weight word. +#define MMQ_BM 64 +#define MMQ_BN 128 +extern ""C"" __global__ void llm_mmq_q8_0( + const unsigned int* __restrict__ weights, // Q8_0 [rows × cols], 34 B/block + const unsigned char* __restrict__ y_q81, // Q8_1 [n_tok × cols], 36 B/block + float* __restrict__ output, // [n_tok × rows] fp32 + int rows, int cols, int n_tok) +{ + __shared__ int sW[MMQ_BM * 8]; // 64 weight rows × 8 int32 (32 int8) + __shared__ float sWd[MMQ_BM]; // 64 weight block-scales + __shared__ int sY[MMQ_BN * 8]; // 128 tokens × 8 int32 acts + __shared__ float sYd[MMQ_BN]; // 128 act block-scales + + int row_block = (int)blockIdx.x * MMQ_BM; + int tok_block = (int)blockIdx.y * MMQ_BN; + int nb = cols >> 5; + + int tid = (int)threadIdx.x; // 0..255 + int warp = tid >> 5; // 0..7 + int lane = tid & 31; + int grp = lane >> 2; // 0..7 + int tig = lane & 3; // 0..3 + int wr = warp & 3; // 0..3 row-group → rows [wr*16 : +16] + int wc = warp >> 2; // 0..1 col-group → tokens [wc*64 : +64] + int mrow0 = wr * 16; + + float acc[8][4]; // 8 N-tiles × 4 C registers, fp32 + #pragma unroll + for (int n = 0; n < 8; n++) { acc[n][0] = acc[n][1] = acc[n][2] = acc[n][3] = 0.f; } + + // Register-prefetch double-buffer: each thread stages 2 weight words (li=tid, + // tid+256) and 4 act words (li=tid,+256,+512,+768) plus, for tid in range, one + // weight/act block-scale. The next K-tile's global loads are issued into these + // registers while the current tile's mma's run, so global latency hides behind + // compute instead of stalling the per-K-block barrier (cp.async can't be used — + // the Q8_0 qs is only 2-byte aligned). Macro loads tile `KB` into the rX regs. + unsigned int rW0, rW1, rY0, rY1, rY2, rY3; + float rWd, rYd; + #define MMQ_LOAD_TILE(KB) do { \ + int gw0 = row_block + (tid >> 3); \ + rW0 = (gw0 < rows) ? sharpi_uint_at(weights, ((long)gw0 * nb + (KB)) * 34L + 2 + (long)(tid & 7) * 4) : 0u; \ + int gw1 = row_block + ((tid + 256) >> 3); \ + rW1 = (gw1 < rows) ? sharpi_uint_at(weights, ((long)gw1 * nb + (KB)) * 34L + 2 + (long)((tid + 256) & 7) * 4) : 0u; \ + if (tid < MMQ_BM) { long wb = ((long)(row_block + tid) * nb + (KB)) * 34L; \ + rWd = (row_block + tid < rows) ? sharpi_fp16_to_fp32(sharpi_byte_at(weights, wb) | (sharpi_byte_at(weights, wb + 1) << 8)) : 0.f; } \ + int gy0 = tok_block + (tid >> 3); \ + rY0 = (gy0 < n_tok) ? *reinterpret_cast(y_q81 + ((long)gy0 * nb + (KB)) * 36L + 4 + (long)(tid & 7) * 4) : 0u; \ + int gy1 = tok_block + ((tid + 256) >> 3); \ + rY1 = (gy1 < n_tok) ? *reinterpret_cast(y_q81 + ((long)gy1 * nb + (KB)) * 36L + 4 + (long)((tid + 256) & 7) * 4) : 0u; \ + int gy2 = tok_block + ((tid + 512) >> 3); \ + rY2 = (gy2 < n_tok) ? *reinterpret_cast(y_q81 + ((long)gy2 * nb + (KB)) * 36L + 4 + (long)((tid + 512) & 7) * 4) : 0u; \ + int gy3 = tok_block + ((tid + 768) >> 3); \ + rY3 = (gy3 < n_tok) ? *reinterpret_cast(y_q81 + ((long)gy3 * nb + (KB)) * 36L + 4 + (long)((tid + 768) & 7) * 4) : 0u; \ + if (tid < MMQ_BN) { int gt = tok_block + tid; \ + rYd = (gt < n_tok) ? sharpi_fp16_to_fp32((*reinterpret_cast(y_q81 + ((long)gt * nb + (KB)) * 36L)) & 0xffffu) : 0.f; } \ + } while (0) + + MMQ_LOAD_TILE(0); + + for (int kb = 0; kb < nb; kb++) { + // Publish the prefetched tile to shared. + sW[tid] = (int)rW0; sW[tid + 256] = (int)rW1; + if (tid < MMQ_BM) sWd[tid] = rWd; + sY[tid] = (int)rY0; sY[tid + 256] = (int)rY1; sY[tid + 512] = (int)rY2; sY[tid + 768] = (int)rY3; + if (tid < MMQ_BN) sYd[tid] = rYd; + __syncthreads(); + + // Issue the next tile's global loads (in flight during the mma's below). + if (kb + 1 < nb) MMQ_LOAD_TILE(kb + 1); + + // A fragment for this warp's 16-row tile (read once, reused over 8 N-tiles). + int a0 = sW[(mrow0 + grp) * 8 + tig]; + int a1 = sW[(mrow0 + grp + 8) * 8 + tig]; + int a2 = sW[(mrow0 + grp) * 8 + tig + 4]; + int a3 = sW[(mrow0 + grp + 8) * 8 + tig + 4]; + float dwA = sWd[mrow0 + grp]; + float dwB = sWd[mrow0 + grp + 8]; + + #pragma unroll + for (int nt = 0; nt < 8; nt++) { + int ncol0 = wc * 64 + nt * 8; + int b0 = sY[(ncol0 + grp) * 8 + tig]; + int b1 = sY[(ncol0 + grp) * 8 + tig + 4]; + float daC0 = sYd[ncol0 + tig * 2]; + float daC1 = sYd[ncol0 + tig * 2 + 1]; + int c0 = 0, c1 = 0, c2 = 0, c3 = 0; + asm( + ""mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "" + ""{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"" + : ""+r""(c0), ""+r""(c1), ""+r""(c2), ""+r""(c3) + : ""r""(a0), ""r""(a1), ""r""(a2), ""r""(a3), ""r""(b0), ""r""(b1)); + acc[nt][0] += (float)c0 * dwA * daC0; + acc[nt][1] += (float)c1 * dwA * daC1; + acc[nt][2] += (float)c2 * dwB * daC0; + acc[nt][3] += (float)c3 * dwB * daC1; + } + __syncthreads(); + } + + int rowA = row_block + mrow0 + grp; + int rowB = rowA + 8; + #pragma unroll + for (int nt = 0; nt < 8; nt++) { + int ncol0 = tok_block + wc * 64 + nt * 8; + int tokC0 = ncol0 + tig * 2; + int tokC1 = ncol0 + tig * 2 + 1; + if (rowA < rows) { + if (tokC0 < n_tok) output[(long)tokC0 * rows + rowA] = acc[nt][0]; + if (tokC1 < n_tok) output[(long)tokC1 * rows + rowA] = acc[nt][1]; + } + if (rowB < rows) { + if (tokC0 < n_tok) output[(long)tokC0 * rows + rowB] = acc[nt][2]; + if (tokC1 < n_tok) output[(long)tokC1 * rows + rowB] = acc[nt][3]; + } + } +} +#undef MMQ_LOAD_TILE +#undef MMQ_BM +#undef MMQ_BN + // Q8_0 GEMM-N: N tokens through one weight matrix. grid=( (rows+7)/8, n_tok ). // Input [n_tok, cols] and output [n_tok, rows] are offset by token; the per-row // accumulation + warp reduce is identical to llm_matvec_q8_0, so this is @@ -1253,6 +1567,48 @@ __device__ __forceinline__ int sharpi_int8_at(const unsigned int* __restrict__ b if (lane == 0) output[(long)token * (long)rows + row] = result; } +// ── Q8_0 → FP16 dequant for cuBLAS prefill GEMM (issue #141) ─────────────── +// Dequantizes a Q8_0-packed weight matrix [rows × cols] into a row-major fp16 +// matrix [rows × cols]. One block per row (256 threads), each thread strides +// over the row's columns. Element (row, c): super-block b = c >> 5, lane = +// c & 31; the per-block fp16 scale d lives at byte (row*nb + b)*34, the int8 +// quant at +2+lane. Stored value d*q is rounded to fp16 — this is the only +// lossy step vs the fp32 matvec (which keeps d*q*x in fp32), and it's what lets +// the prefill GEMM read each weight once per batch instead of once per token. +extern ""C"" __global__ void llm_dequant_q8_0_to_f16( + const unsigned int* __restrict__ weights, + unsigned short* __restrict__ out, // [rows * cols] fp16 + int rows, int cols) +{ + int row = (int)blockIdx.x; + if (row >= rows) return; + int num_blocks = cols >> 5; + long row_base_bytes = (long)row * (long)num_blocks * 34L; + long out_row = (long)row * (long)cols; + for (int c = (int)threadIdx.x; c < cols; c += (int)blockDim.x) { + int block = c >> 5; + int lane = c & 31; + long b0 = row_base_bytes + (long)block * 34L; + unsigned int dlo = sharpi_byte_at(weights, b0 + 0); + unsigned int dhi = sharpi_byte_at(weights, b0 + 1); + float d = sharpi_fp16_to_fp32(dlo | (dhi << 8)); + int q = sharpi_int8_at(weights, b0 + 2 + (long)lane); + out[out_row + c] = (unsigned short)sharpi_fp32_to_fp16(d * (float)q); + } +} + +// ── FP32 → FP16 elementwise convert (issue #141) ────────────────────────── +// Converts the prefill activation batch [n] fp32 → fp16 so it can feed the +// cuBLAS fp16 GEMM alongside the dequantized weights. +extern ""C"" __global__ void llm_f32_to_f16( + const float* __restrict__ in, + unsigned short* __restrict__ out, + int n) +{ + int i = (int)(blockIdx.x * blockDim.x + threadIdx.x); + if (i < n) out[i] = (unsigned short)sharpi_fp32_to_fp16(in[i]); +} + // ── MatVec F32 — N=2 variant (issue #43) ────────────────────────────────── // Two input vectors, two output vectors, single weight read. Mirrors the // llm_matvec_f32 launch geometry (8 rows × 32 threads/row). Wins over two @@ -2696,6 +3052,153 @@ __device__ __forceinline__ int sharpi_int8_at(const unsigned int* __restrict__ b } } +// ── Flash-attention prefill (issue #141 attention) ───────────────────────── +// Memory-efficient batched SDPA replacing the scalar llm_full_seq_attention / +// llm_attention_swa_batched (one 256-thread block PER query, each re-reading its +// whole K/V range from global → O(n²) global traffic; for SWA, adjacent queries' +// 512-wide windows overlap ~99%, so K/V is re-streamed up to ~512×). Here a block +// handles a TILE of FA_QT queries of one head and streams K/V in shared-memory +// tiles of `kt_tile` keys, so each key is read from global once per FA_QT queries +// (FA_QT× less traffic) and the softmax runs online (running max/sum + rescaled +// output accumulator, FlashAttention-2 style) — no n²-sized score buffer. +// +// One warp = one query. Lane L owns head dims {L, L+32, …}; qreg/oreg hold up to +// 512/32 = 16 dims/lane. The QK dot is a warp reduce (score is then warp-uniform, +// so the causal/window mask `continue` and the online-softmax update never diverge +// within a warp). GQA via kv_head = h/(num_heads/num_kv_heads); per-layer head_dim; +// causal (key ≤ start_pos+qi) + optional sliding window (window_size>0). Matches the +// scalar kernels to fp tolerance (online softmax reassociates the same sum), not +// bit-exact. Dynamic shared = 2*kt_tile*head_dim floats (sized by the host to fit). +// FA_QT warps/block = queries sharing each K/V tile load (the reuse factor). +#define FA_QT 16 +extern ""C"" __global__ void llm_flash_attn_prefill_f32( + const float* __restrict__ q_all, // [n_tok, num_heads*head_dim] + const float* __restrict__ k_cache, + const float* __restrict__ v_cache, + float* __restrict__ out_all, // [n_tok, num_heads*head_dim] + int num_heads, int num_kv_heads, int head_dim, + int start_pos, int window_size, int max_seq_len, int n_tok, + int kt_tile, float attn_scale) +{ + // Dynamic shared: K as fp16 (half2-packed, hd/2 uints/key) then V as fp32. K is + // read with the half2 QK dot (fp16-rounded inputs — argmax-stable); V stays fp32 + // for the exact scalar PV. Each lane owns dim-PAIRS pi = lane+32*p (dims 2·pi, + // 2·pi+1) so the shared half2 loads stay coalesced. host sizes 6·kt_tile·head_dim B. + extern __shared__ unsigned int fa_smem[]; + int hd2 = head_dim >> 1; // pairs per key + unsigned int* sKh = fa_smem; // [kt_tile * hd2] fp16x2 + float* sV = (float*)(fa_smem + kt_tile * hd2); // [kt_tile * head_dim] fp32 + + int h = (int)blockIdx.x; + int warp = (int)(threadIdx.x >> 5); // 0..FA_QT-1 (query within the tile) + int lane = (int)(threadIdx.x & 31); + int tid = (int)threadIdx.x; + int qi = (int)blockIdx.y * FA_QT + warp; + + int kv_head = h / (num_heads / num_kv_heads); + int kv_dim = num_kv_heads * head_dim; + float scale = (attn_scale > 0.f) ? attn_scale : rsqrtf((float)head_dim); + int q_dim = num_heads * head_dim; + int ndj2 = (hd2 + 31) >> 5; // dim-pairs per lane (≤8) + + bool active = (qi < n_tok); + int qpos = start_pos + qi; + int win_end = qpos + 1; // keys [.., qpos] + int win_start = (window_size > 0) ? (win_end - window_size) : 0; + if (win_start < 0) win_start = 0; + + unsigned int qh2[8]; // this query's dim-pairs, fp16x2-packed + float oreg[16]; // PV accumulator, fp32, 2 per pair + #pragma unroll + for (int p = 0; p < 8; p++) { qh2[p] = 0u; oreg[2 * p] = 0.f; oreg[2 * p + 1] = 0.f; } + if (active) { + long q_base = (long)qi * q_dim + (long)h * head_dim; + for (int p = 0; p < ndj2; p++) { + int pi = lane + 32 * p; + if (pi < hd2) qh2[p] = sharpi_f32x2_to_f16x2(q_all[q_base + 2 * pi], q_all[q_base + 2 * pi + 1]); + } + } + float m_run = sharpi_neg_inf(); + float l_run = 0.f; + + // Union key range over the FA_QT queries in this block. + int last_qi = (int)blockIdx.y * FA_QT + (FA_QT - 1); + if (last_qi > n_tok - 1) last_qi = n_tok - 1; + int blk_key_end = (start_pos + last_qi) + 1; + int first_qpos = start_pos + (int)blockIdx.y * FA_QT; + int blk_key_start = (window_size > 0) ? (first_qpos + 1 - window_size) : 0; + if (blk_key_start < 0) blk_key_start = 0; + + for (int kt0 = blk_key_start; kt0 < blk_key_end; kt0 += kt_tile) { + int tile_keys = blk_key_end - kt0; + if (tile_keys > kt_tile) tile_keys = kt_tile; + + // Stage K (fp32 global → fp16x2 shared, one pair/thread-step). + for (int idx = tid; idx < kt_tile * hd2; idx += (int)blockDim.x) { + int kk = idx / hd2, pr = idx - kk * hd2; + unsigned int kh = 0u; + if (kk < tile_keys) { + long off = (long)(kt0 + kk) * kv_dim + (long)kv_head * head_dim + 2 * pr; + kh = sharpi_f32x2_to_f16x2(k_cache[off], k_cache[off + 1]); + } + sKh[idx] = kh; + } + // Stage V (fp32 → fp32 shared). + for (int idx = tid; idx < kt_tile * head_dim; idx += (int)blockDim.x) { + int kk = idx / head_dim, d = idx - kk * head_dim; + sV[idx] = (kk < tile_keys) + ? v_cache[(long)(kt0 + kk) * kv_dim + (long)kv_head * head_dim + d] + : 0.f; + } + __syncthreads(); + + if (active) { + for (int kk = 0; kk < tile_keys; kk++) { + int abs_t = kt0 + kk; + if (abs_t >= win_end || abs_t < win_start) continue; // warp-uniform + unsigned int acc = 0u; + for (int p = 0; p < ndj2; p++) { + int pi = lane + 32 * p; + if (pi < hd2) acc = sharpi_hfma2(qh2[p], sKh[kk * hd2 + pi], acc); + } + float part = sharpi_f16x2_sum(acc); + part += __shfl_xor_sync(0xffffffffu, part, 16); + part += __shfl_xor_sync(0xffffffffu, part, 8); + part += __shfl_xor_sync(0xffffffffu, part, 4); + part += __shfl_xor_sync(0xffffffffu, part, 2); + part += __shfl_xor_sync(0xffffffffu, part, 1); + float score = part * scale; + float m_new = fmaxf(m_run, score); + float alpha = __expf(m_run - m_new); + float pw = __expf(score - m_new); + l_run = l_run * alpha + pw; + for (int p = 0; p < ndj2; p++) { + int pi = lane + 32 * p; + float v0 = (pi < hd2) ? sV[kk * head_dim + 2 * pi] : 0.f; + float v1 = (pi < hd2) ? sV[kk * head_dim + 2 * pi + 1] : 0.f; + oreg[2 * p] = oreg[2 * p] * alpha + pw * v0; + oreg[2 * p + 1] = oreg[2 * p + 1] * alpha + pw * v1; + } + m_run = m_new; + } + } + __syncthreads(); + } + + if (active) { + float inv = (l_run > 0.f) ? (1.f / l_run) : 0.f; + long o_base = (long)qi * q_dim + (long)h * head_dim; + for (int p = 0; p < ndj2; p++) { + int pi = lane + 32 * p; + if (pi < hd2) { + out_all[o_base + 2 * pi] = oreg[2 * p] * inv; + out_all[o_base + 2 * pi + 1] = oreg[2 * p + 1] * inv; + } + } + } +} +#undef FA_QT + // Batched sliding-window attention over N query tokens (Gemma 4 SWA layers in // batched-trunk prefill). Grid = (num_heads, n_tok); query token i sits at // absolute position start_pos+i and attends [max(0,pos+1-window), pos+1). The diff --git a/src/SharpInference.Engine/CudaForwardPass.cs b/src/SharpInference.Engine/CudaForwardPass.cs index 9bb1ccb..7e79c46 100644 --- a/src/SharpInference.Engine/CudaForwardPass.cs +++ b/src/SharpInference.Engine/CudaForwardPass.cs @@ -169,6 +169,35 @@ public sealed unsafe class CudaForwardPass : IForwardPass /// batched path against the bit-exact per-token loop on one model instance. /// public bool BatchedPrefillEnabled { get; set; } + /// + /// Issue #141: route Q8_0 trunk matmuls in the batched prefill through the + /// compute-bound cuBLAS GEMM () + /// instead of the memory-bound matvec GEMM-N. Default on + /// (SHARPI_PREFILL_GEMM); settable so tests can A/B the GEMM path + /// against the bit-exact matvec path on one model instance. Non-Q8_0 trunk + /// weights always fall back to the matvec GEMM-N regardless. + /// + public bool PrefillGemmEnabled { get; set; } + /// + /// Issue #141 (MMQ): route Q8_0 trunk matmuls through the int8 tensor-core MMQ + /// kernel () instead of the + /// dequant→fp16→cuBLAS GEMM (). Reads each Q8_0 + /// weight once as int8 with no fp16 HBM temp, on int8 tensor cores. Takes + /// precedence over when both are set. Default on + /// (SHARPI_PREFILL_MMQ=0 reverts to the GEMM path). Argmax-stable, not + /// bit-exact (both operands int8-quantized), like the GEMM path it replaces. + /// + public bool PrefillMmqEnabled { get; set; } + /// + /// Issue #141 (attention): route the batched-trunk prefill attention through the + /// memory-efficient flash kernel (, + /// shared K/V tiles + online softmax) instead of the scalar per-query + /// / + /// (which re-read each query's whole K/V range from global — O(n²), the dominant + /// prefill cost). fp32 KV only. Argmax-stable, not bit-exact (online softmax). + /// Default on (SHARPI_PREFILL_FLASH=0 reverts to the scalar kernels). + /// + public bool PrefillFlashAttnEnabled { get; set; } private int _bpCapacity; // current N the scratch is sized for (0 = none) private Tensor? _bpHidden, _bpResidual, _bpNorm; // [N × embDim] private Tensor? _bpQ, _bpAttnOut; // [N × numHeads*maxHeadDim] @@ -265,6 +294,19 @@ public CudaForwardPass(GgufModel model, CudaBackend gpu, ModelHyperparams hp, // bit-exact per-token loop (useful for A/B and parity debugging). BatchedPrefillEnabled = Environment.GetEnvironmentVariable("SHARPI_BATCHED_PREFILL") != "0"; + PrefillGemmEnabled = + Environment.GetEnvironmentVariable("SHARPI_PREFILL_GEMM") != "0"; + // Issue #141 (MMQ): default on — the int8 tensor-core MMQ beats the + // dequant→fp16→cuBLAS GEMM matmul (~316ms vs ~332ms over a 1848-tok Gemma 4 + // prefill) and drops the fp16 weight HBM temp. SHARPI_PREFILL_MMQ=0 reverts to + // the cuBLAS GEMM path. Gated under PrefillGemmEnabled (the compute-bound path). + PrefillMmqEnabled = + Environment.GetEnvironmentVariable("SHARPI_PREFILL_MMQ") != "0"; + // Issue #141 (attention): default on — the flash kernel cuts batched-prefill + // attention ~929ms→~411ms (2.26×) at N=1848, lifting Gemma 4 prefill + // ~1389→~2180 t/s. SHARPI_PREFILL_FLASH=0 reverts to the scalar kernels. + PrefillFlashAttnEnabled = + Environment.GetEnvironmentVariable("SHARPI_PREFILL_FLASH") != "0"; _maxHeadDim = _headDim; if (hp.LayerHeadDim is { } lhdMax) for (int i = 0; i < hp.NumLayers; i++) @@ -776,6 +818,10 @@ private void BenchMatVec() // Profiling state (only used when SHARPI_CUDA_PROFILE is set). private static readonly bool s_profile = Environment.GetEnvironmentVariable("SHARPI_CUDA_PROFILE") == "1"; + private static readonly bool s_prefillProfile = + Environment.GetEnvironmentVariable("SHARPI_PREFILL_PROFILE") == "1"; + private readonly System.Diagnostics.Stopwatch _profSw = new(); + private double _profAttnMs; private readonly double[] _phaseMs = new double[10]; private readonly long[] _phaseCount = new long[10]; private const int PH_EMBED = 0, PH_QKV = 1, PH_ROPE_QKN = 2, PH_KV_ATTN = 3, @@ -786,11 +832,15 @@ private void BenchMatVec() // CUDA Graph decode (issue #136): the Gemma 4 layer + output region has static // topology across decode tokens (only `position` varies), so capture it once on the // first decode token and replay per token — collapsing ~1k host launches/token into - // one cuGraphLaunch + a handful of node-param updates. Off by default; opt in with - // SHARPI_CUDA_GRAPH=1 or the UseCudaGraph setter. Falls back to direct launches on - // any capture/replay failure. + // one cuGraphLaunch + a handful of node-param updates. Falls back to direct launches + // on any capture/replay failure. + // + // Default ON (issue #142): after the #137 launch fusions and the dp4a decode matvec, + // graphs measure +9–10% at both low and ~1K context on the all-GPU Gemma 4 path (the + // earlier short-context regression that kept #136 default-off is gone). SHARPI_CUDA_GRAPH=0 + // reverts to direct launches. private bool _useCudaGraph = - Environment.GetEnvironmentVariable("SHARPI_CUDA_GRAPH") == "1"; + Environment.GetEnvironmentVariable("SHARPI_CUDA_GRAPH") != "0"; private bool _graphCaptured; // Number of KV entries SnapKV physically dropped when it compacted the cache // (= N - K at eviction, 0 otherwise). Decode indexes the compacted cache by the @@ -1580,6 +1630,13 @@ public ReadOnlySpan Prefill(IReadOnlyList tokens, int startPos = 0) int N = tokens.Count; LastPrefillWasBatched = false; + // Issue #142: pre-grow the dp4a Q8_1 input scratch to the widest decode + // matvec (FFN down: cols = intermDim) before any CUDA-graph capture on the + // first decode token — capture forbids cudaMalloc, so the buffer must + // already be at max size. Harmless when the dp4a path is disabled. + if (_isGemma4Like) + _gpu.EnsureQ81Scratch(Math.Max(_embDim, _intermDim)); + // SnapKV (issue #59) gating: only run eviction when this is a fresh // prefill (startPos==0), the effective budget is positive (env-set or // VRAM-scaled auto), the prompt is long enough that eviction would drop @@ -1665,8 +1722,39 @@ private bool IsGemma4BatchedPrefillSupported() return BatchableWeight(_wOutput); } - private void GpuMatMulBatched(Tensor outAll, Tensor weights, Tensor inAll, int n) => - _gpu.MatMulBatched(outAll, weights, inAll, n, WDType(weights)); + private void GpuMatMulBatched(Tensor outAll, Tensor weights, Tensor inAll, int n) + { + var dt = WDType(weights); + // Issue #141: route the trunk matmuls through compute-bound cuBLAS GEMM + // (each weight read once per batch) instead of the memory-bound matvec + // GEMM-N (weight re-streamed per token). Q8_0 weights dequant to fp16 + // first; F32 weights (PLE projections) feed Sgemm's TF32 path directly. + // ~argmax-stable, not byte-exact — fine for Gemma 4 prefill, which has no + // MTP/GDN byte-parity oracle. Other dtypes keep the matvec GEMM-N. + if (!PrefillGemmEnabled) + { + _gpu.MatMulBatched(outAll, weights, inAll, n, dt); + return; + } + if (dt == DType.Q8_0) + { + if (PrefillMmqEnabled) + _gpu.MatMulBatchedMmq(outAll, weights, inAll, n, dt); + else + _gpu.MatMulBatchedGemm(outAll, weights, inAll, n, dt); + } + else if (dt == DType.Float32) + { + // C[n×rows] = A[n×cols] · B[rows×cols]ᵀ → Sgemm(C, A, B, M=n, K=cols, N=rows). + int rows = (int)(outAll.ElementCount / n); + int cols = (int)(inAll.ElementCount / n); + _gpu.Sgemm(outAll, inAll, weights, n, cols, rows); + } + else + { + _gpu.MatMulBatched(outAll, weights, inAll, n, dt); + } + } /// (Re)allocate the batched-trunk scratch for a prompt of length N. private void EnsureBatchedTrunkScratch(int n) @@ -1720,23 +1808,54 @@ private ReadOnlySpan PrefillBatchedTrunkGemma4(IReadOnlyList tokens, EnsureBatchedTrunkScratch(N); int embDim = _embDim; - // 1. Embed every token into _bpHidden, then batched embedding scale. - for (int i = 0; i < N; i++) + bool profile = s_prefillProfile; + var swp = profile ? System.Diagnostics.Stopwatch.StartNew() : null; + long tEmbed = 0, tPle = 0, tLayers = 0; + + // 1. Embed every token into _bpHidden, then batched embedding scale. The Q8_0 + // table (Gemma 4) goes through a single batched launch instead of 2·N host- + // driven EmbedLookup + copy launches (bit-identical); other dtypes keep the loop. + var embDType = _embIsQuantized + ? _weightDTypes.GetValueOrDefault(_gpuEmbedding.Handle, DType.Q4_K) + : DType.Float32; + if (_embIsQuantized && embDType == DType.Q8_0) + { + int[] ids = tokens as int[] ?? System.Linq.Enumerable.ToArray(tokens); + var idTensor = _gpu.UploadRaw( + System.Runtime.InteropServices.MemoryMarshal.AsBytes(ids), + TensorShape.D1(N), DType.Float32); + _gpu.EmbedLookupQ8_0Batched(_gpuEmbedding, _bpHidden!, idTensor, N, embDim); + _gpu.Free(idTensor); + } + else { - EmbedTokenGpu(tokens[i]); // writes _hidden - _gpu.CopyDeviceRegion(_bpHidden!, (long)i * embDim * sizeof(float), - _hidden, 0, (long)embDim * sizeof(float)); + for (int i = 0; i < N; i++) + { + EmbedTokenGpu(tokens[i]); // writes _hidden + _gpu.CopyDeviceRegion(_bpHidden!, (long)i * embDim * sizeof(float), + _hidden, 0, (long)embDim * sizeof(float)); + } } if (_hp.EmbeddingScale != 1f) _gpu.ScaleInPlace(_bpHidden!, _hp.EmbeddingScale); + if (profile) { _gpu.Synchronize(); tEmbed = swp!.ElapsedMilliseconds; } // 2. PLE pre-pass batched (builds _bpProjAll = [N × L*pleWidth]). if (_hp.HasPerLayerTokenEmbd) BuildPerLayerProjectionsBatched(tokens); + if (profile) { _gpu.Synchronize(); tPle = swp!.ElapsedMilliseconds; } // 3. Transformer layers, batched across N. for (int layer = 0; layer < _hp.NumLayers; layer++) GpuLayerBatchedTrunkGemma4(layer, N, startPos); + if (profile) + { + _gpu.Synchronize(); + tLayers = swp!.ElapsedMilliseconds; + Console.Error.WriteLine($"[prefill-profile] N={N} embed={tEmbed}ms ple={tPle - tEmbed}ms " + + $"layers={tLayers - tPle}ms (attn={_profAttnMs:F0}ms)"); + _profAttnMs = 0; + } // 4. Final norm + output projection on the last token only. var lastHidden = _gpu.View(_bpHidden!, (long)(N - 1) * embDim, embDim); @@ -1780,16 +1899,21 @@ private void BuildPerLayerProjectionsBatched(IReadOnlyList tokens) * DTypeInfo.BytesPerBlock(pleRef.DType); // CPU dequant of each token's PLE row into the [N × stackedDim] host buffer. + // Parallelized across tokens (each row is independent): for a long prompt + // this serial gather+dequant of N×stackedDim Q8_0 elements was ~30% of the + // whole batched prefill (issue #141 profiling). var host = _bpPleRowHostAll!; - for (int i = 0; i < N; i++) + byte* basePtr = pleRef.DataPtr; + var dtype = pleRef.DType; + System.Threading.Tasks.Parallel.For(0, N, i => { - byte* rowPtr = pleRef.DataPtr + (long)tokens[i] * bytesPerRow; + byte* rowPtr = basePtr + (long)tokens[i] * bytesPerRow; var dst = new Span(host).Slice(i * stackedDim, stackedDim); - if (pleRef.DType == DType.Float32) + if (dtype == DType.Float32) new ReadOnlySpan((float*)rowPtr, stackedDim).CopyTo(dst); else - Dequantize.ToFloat32(new ReadOnlySpan(rowPtr, bytesPerRow), dst, pleRef.DType, stackedDim); - } + Dequantize.ToFloat32(new ReadOnlySpan(rowPtr, bytesPerRow), dst, dtype, stackedDim); + }); _gpu.UploadInto(_bpPleRowAll!, host); _gpu.ScaleInPlace(_bpPleRowAll!, MathF.Sqrt(_pleWidth)); @@ -1865,13 +1989,18 @@ private void GpuLayerBatchedTrunkGemma4(int layer, int N, int startPos) int effLayerCtx = (_hp.IsSwaLayer is { } swaEff && swaEff[effLayer] && window > 0) ? Math.Min(_maxSeqLen, window) : _maxSeqLen; + if (s_prefillProfile) { _gpu.Synchronize(); _profSw.Restart(); } // Gemma 4: attention_scale = 1.0, passed explicitly (kernel skips its rsqrtf). - if (isSwa) + if (PrefillFlashAttnEnabled) + _gpu.FlashAttentionPrefill(qAll, _gpuKCache[effLayer], _gpuVCache[effLayer], attnAll, + _numHeads, _numKvHeads, layerHd, startPos, isSwa ? window : 0, effLayerCtx, N, attnScale: 1f); + else if (isSwa) _gpu.AttentionSwaBatched(qAll, _gpuKCache[effLayer], _gpuVCache[effLayer], attnAll, _numHeads, _numKvHeads, layerHd, startPos, window, effLayerCtx, N, attnScale: 1f); else _gpu.AttentionBatched(qAll, _gpuKCache[effLayer], _gpuVCache[effLayer], attnAll, _numHeads, _numKvHeads, layerHd, startPos, effLayerCtx, N, attnScale: 1f); + if (s_prefillProfile) { _gpu.Synchronize(); _profAttnMs += _profSw.Elapsed.TotalMilliseconds; } GpuMatMulBatched(_bpHidden!, _wo[layer], attnAll, N); if (_wPostAttnNorm is not null) diff --git a/tests/SharpInference.Tests.ForwardPass/CudaFlashAttnTests.cs b/tests/SharpInference.Tests.ForwardPass/CudaFlashAttnTests.cs new file mode 100644 index 0000000..4cfe77b --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/CudaFlashAttnTests.cs @@ -0,0 +1,94 @@ +using SharpInference.Core; +using SharpInference.Cuda; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Issue #141 (attention): parity for the memory-efficient flash-attention prefill +/// (, kernel +/// llm_flash_attn_prefill_f32) against the validated scalar batched kernels +/// it replaces — (global) and +/// (sliding window). Random Q/K/V are +/// run through both; the flash kernel's online softmax reassociates the same sum and +/// rounds K to fp16 for the half2 QK dot, so it tracks the reference to fp tolerance +/// (~fp16-relative), not bit-exactly. Exercises GQA, both head_dims Gemma 4 uses +/// (256 SWA / 512 global), causal masking, windowing, and a partial final query tile. +/// +/// Silent no-op on hosts without CUDA, matching the other Cuda* test files. +/// +public sealed unsafe class CudaFlashAttnTests +{ + private static CudaBackend? TryCreate() + { + if (!CudaBackend.IsAvailable()) return null; + try { return CudaBackend.Create(); } + catch { return null; } + } + + [Fact] + public void FlashAttentionPrefill_MatchesScalarBatched() + { + using var gpu = TryCreate(); + if (gpu is null) return; + + // (numHeads, numKvHeads, headDim, window, nTok). window=0 → global (full causal). + (int nh, int nkv, int hd, int win, int nTok)[] cfgs = + { + (8, 2, 256, 0, 200), // global, SWA head_dim, partial last tile (200%8=0 → exact; ok) + (8, 2, 512, 0, 173), // global, global head_dim, partial last tile (173%8≠0) + (8, 2, 256, 64, 200), // sliding window 64 < nTok (windowing exercised) + (8, 2, 512, 96, 130), // sliding window 96, global head_dim + (4, 4, 128, 0, 64), // MHA (no GQA), small head_dim + }; + + foreach (var (nh, nkv, hd, win, nTok) in cfgs) + { + var rng = new Random(20260605 + nh * 7 + hd * 13 + win * 17 + nTok); + int qDim = nh * hd, kvDim = nkv * hd; + + var q = new float[(long)nTok * qDim]; + for (int i = 0; i < q.Length; i++) q[i] = (float)(rng.NextDouble() * 2 - 1); + // K/V cache filled for positions [0, nTok) (prefill startPos=0). + var k = new float[(long)nTok * kvDim]; + var v = new float[(long)nTok * kvDim]; + for (int i = 0; i < k.Length; i++) { k[i] = (float)(rng.NextDouble() * 2 - 1); v[i] = (float)(rng.NextDouble() * 2 - 1); } + + var gq = gpu.Upload(q, TensorShape.D1(q.Length)); + var gk = gpu.Upload(k, TensorShape.D1(k.Length)); + var gv = gpu.Upload(v, TensorShape.D1(v.Length)); + var gRef = gpu.Allocate(TensorShape.D1(q.Length)); + var gFla = gpu.Allocate(TensorShape.D1(q.Length)); + + if (win == 0) + gpu.AttentionBatched(gq, gk, gv, gRef, nh, nkv, hd, startPos: 0, maxSeqLen: nTok, nTok: nTok); + else + gpu.AttentionSwaBatched(gq, gk, gv, gRef, nh, nkv, hd, startPos: 0, windowSize: win, maxSeqLen: nTok, nTok: nTok); + gpu.FlashAttentionPrefill(gq, gk, gv, gFla, nh, nkv, hd, startPos: 0, windowSize: win, maxSeqLen: nTok, nTok: nTok); + gpu.Synchronize(); + + var outRef = new float[q.Length]; + var outFla = new float[q.Length]; + gpu.Download(gRef, outRef); + gpu.Download(gFla, outFla); + gpu.Free(gq); gpu.Free(gk); gpu.Free(gv); gpu.Free(gRef); gpu.Free(gFla); + + double sumSq = 0; + for (int i = 0; i < outRef.Length; i++) sumSq += (double)outRef[i] * outRef[i]; + float rms = (float)Math.Sqrt(sumSq / outRef.Length) + 1e-9f; + + float maxAbs = 0; + int mismatches = 0; + for (int i = 0; i < outRef.Length; i++) + { + float diff = MathF.Abs(outFla[i] - outRef[i]); + maxAbs = MathF.Max(maxAbs, diff); + // K is fp16-rounded for the half2 QK dot → ~fp16-relative score error. + if (diff > 5e-3f * rms) mismatches++; + } + Console.WriteLine( + $"Flash nh={nh} nkv={nkv} hd={hd} win={win} nTok={nTok}: maxAbs={maxAbs:E2} rms={rms:E2} mismatches={mismatches}/{outRef.Length}"); + Assert.True(mismatches == 0, + $"Flash attention diverged from scalar reference: {mismatches}/{outRef.Length} beyond 5e-3·rms ({rms:E3}), maxAbs={maxAbs:E3} (nh={nh} hd={hd} win={win} nTok={nTok})."); + } + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/CudaMatMulBatchedTests.cs b/tests/SharpInference.Tests.ForwardPass/CudaMatMulBatchedTests.cs index a1eef0d..a8aaa71 100644 --- a/tests/SharpInference.Tests.ForwardPass/CudaMatMulBatchedTests.cs +++ b/tests/SharpInference.Tests.ForwardPass/CudaMatMulBatchedTests.cs @@ -261,6 +261,10 @@ public void MatMulBatched_Q8_0_BitwiseMatchesSequential() { using var gpu = TryCreate(); if (gpu is null) return; + // The batched matvec GEMM-N path is fp32; pin the sequential MatMul + // reference to the fp32 kernel too (the default dp4a path quantizes the + // activation to int8 — argmax-stable, not bit-exact). Issue #142. + gpu.Q80Dp4aEnabled = false; foreach ((int rows, int cols, int nTok) in new[] { (8, 256, 2), (33, 512, 5), (64, 1024, 17), (130, 2048, 40) }) diff --git a/tests/SharpInference.Tests.ForwardPass/CudaMmqQ8_0Tests.cs b/tests/SharpInference.Tests.ForwardPass/CudaMmqQ8_0Tests.cs new file mode 100644 index 0000000..4f68f0c --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/CudaMmqQ8_0Tests.cs @@ -0,0 +1,113 @@ +using SharpInference.Core; +using SharpInference.Cpu; +using SharpInference.Cuda; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Issue #141 (MMQ): parity for the int8 tensor-core Q8_0×Q8_1 batched matmul +/// (, kernel llm_mmq_q8_0). +/// A small Q8_0 weight matrix [rows×cols] and an fp32 activation batch [nTok×cols] +/// are multiplied both on the GPU (MMQ) and on the CPU (, +/// fp32 reference). MMQ quantizes the activation to int8 (Q8_1) before the int8 mma, +/// so — like the dp4a decode matvec — it tracks the fp32 reference to a loose +/// per-row-RMS tolerance rather than bit-exactly. The int32 dot itself is exact; the +/// only error is the ~Q8 activation quantization. +/// +/// Silent no-op on hosts without CUDA, matching the other Cuda* test files. +/// +public sealed unsafe class CudaMmqQ8_0Tests +{ + private static CudaBackend? TryCreate() + { + if (!CudaBackend.IsAvailable()) return null; + try { return CudaBackend.Create(); } + catch { return null; } + } + + private static ushort HalfToUshort(Half h) => + BitConverter.ToUInt16(BitConverter.GetBytes(h), 0); + + /// Build × Q8_0 (34 B/block). + private static byte[] BuildQ8_0Matrix(int rows, int cols, Random rng) + { + int blocksPerRow = cols / 32; + int bytesPerRow = blocksPerRow * 34; + var bytes = new byte[rows * bytesPerRow]; + for (int r = 0; r < rows; r++) + for (int b = 0; b < blocksPerRow; b++) + { + int off = r * bytesPerRow + b * 34; + float d = (float)(rng.NextDouble() * 0.09 + 0.01); + ushort dHalf = HalfToUshort((Half)d); + bytes[off + 0] = (byte)(dHalf & 0xFF); + bytes[off + 1] = (byte)(dHalf >> 8); + for (int i = 0; i < 32; i++) + bytes[off + 2 + i] = (byte)(sbyte)(rng.Next(255) - 127); + } + return bytes; + } + + [Fact] + public void MatMulBatchedMmq_Q8_0_TracksCpuReference() + { + using var gpu = TryCreate(); + if (gpu is null) return; + + // Cover: square, non-multiple-of-16 rows + non-multiple-of-8 tokens (partial + // tile guards), and a wide single-block-token batch. + foreach ((int rows, int cols, int nTok) in new[] + { (256, 256, 8), (1024, 512, 12), (33, 256, 5), (128, 2560, 64) }) + { + var rng = new Random(20260605 + rows * 31 + cols * 7 + nTok); + byte[] weightBytes = BuildQ8_0Matrix(rows, cols, rng); + + var acts = new float[nTok * cols]; + for (int i = 0; i < acts.Length; i++) + acts[i] = (float)(rng.NextDouble() * 2 - 1); + + // CPU reference: out[t*rows + r] = Σ W[r]·acts[t] (fp32, exact-byte Q8_0). + int bytesPerRow = (cols / 32) * 34; + var cpuOut = new float[nTok * rows]; + fixed (byte* wPtr = weightBytes) + fixed (float* aPtr = acts) + { + for (int t = 0; t < nTok; t++) + for (int r = 0; r < rows; r++) + cpuOut[t * rows + r] = SimdKernels.DotQ8_0(wPtr + r * bytesPerRow, aPtr + t * cols, cols); + } + + var gpuW = gpu.UploadRaw(weightBytes, TensorShape.D1(weightBytes.Length), DType.Q8_0); + var gpuX = gpu.Upload(acts, TensorShape.D1(acts.Length)); + var gpuY = gpu.Allocate(TensorShape.D1((long)nTok * rows)); + + gpu.MatMulBatchedMmq(gpuY, gpuW, gpuX, nTok, DType.Q8_0); + gpu.Synchronize(); + + var gpuOut = new float[nTok * rows]; + gpu.Download(gpuY, gpuOut); + gpu.Free(gpuW); + gpu.Free(gpuX); + gpu.Free(gpuY); + + // Per-row magnitude scale for the relative bound (dot of ±1 acts over + // `cols` int8 weights has stddev ~ sqrt(cols)). + double sumSq = 0; + for (int i = 0; i < cpuOut.Length; i++) sumSq += (double)cpuOut[i] * cpuOut[i]; + float refRms = (float)Math.Sqrt(sumSq / cpuOut.Length); + + int mismatches = 0; + float maxAbs = 0; + for (int i = 0; i < cpuOut.Length; i++) + { + float diff = MathF.Abs(gpuOut[i] - cpuOut[i]); + maxAbs = MathF.Max(maxAbs, diff); + if (diff > 0.02f * refRms) mismatches++; + } + Console.WriteLine( + $"MMQ rows={rows} cols={cols} nTok={nTok}: maxAbs={maxAbs:E2} refRms={refRms:E2} mismatches={mismatches}/{cpuOut.Length}"); + Assert.True(mismatches <= cpuOut.Length / 100 + 1, + $"MMQ Q8_0 drifted from fp32 reference: {mismatches}/{cpuOut.Length} beyond 2% of RMS ({refRms:E3}), maxAbs={maxAbs:E3}."); + } + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/CudaQ8_0Tests.cs b/tests/SharpInference.Tests.ForwardPass/CudaQ8_0Tests.cs index fa949e9..b696872 100644 --- a/tests/SharpInference.Tests.ForwardPass/CudaQ8_0Tests.cs +++ b/tests/SharpInference.Tests.ForwardPass/CudaQ8_0Tests.cs @@ -66,6 +66,11 @@ public void MatVec_Q8_0_MatchesCpuReference() { using var gpu = TryCreate(); if (gpu is null) return; + // Validate the fp32-decode kernel (llm_matvec_q8_0): exact-byte Q8_0 weights + // dotted with fp32 activations, so only fp16-d rounding + reduction order + // differ from the CPU reference. Pin off the dp4a path (issue #142), which + // quantizes the activation to int8 and is covered by the looser test below. + gpu.Q80Dp4aEnabled = false; foreach ((int rows, int cols) in new[] { (256, 256), (1024, 1024), (33, 512) }) { @@ -127,6 +132,71 @@ public void MatVec_Q8_0_MatchesCpuReference() } } + /// + /// Issue #142: the dp4a Q8_0 matvec (llm_matvec_q8_0_dp4a) quantizes the + /// activation to int8 (Q8_1) before the int8·int8 dp4a dot — exactly llama.cpp's + /// decode matvec. That introduces ~Q8 activation-quant error (~1%), so it tracks + /// the fp32 CPU reference to a loose relative tolerance, not 1e-3. The aggregate + /// dot must still be accurate enough to be argmax-stable, which this bounds. + /// + [Fact] + public void MatVec_Q8_0_Dp4a_TracksCpuReference() + { + using var gpu = TryCreate(); + if (gpu is null) return; + gpu.Q80Dp4aEnabled = true; + + foreach ((int rows, int cols) in new[] { (256, 256), (1024, 1024), (64, 2560) }) + { + var rng = new Random(20260605 + rows * 31 + cols); + byte[] weightBytes = BuildQ8_0Matrix(rows, cols, rng); + + var input = new float[cols]; + for (int i = 0; i < cols; i++) + input[i] = (float)(rng.NextDouble() * 2 - 1); + + int bytesPerRow = (cols / 32) * 34; + var cpuOutput = new float[rows]; + fixed (byte* wPtr = weightBytes) + fixed (float* iPtr = input) + { + for (int r = 0; r < rows; r++) + cpuOutput[r] = SimdKernels.DotQ8_0(wPtr + r * bytesPerRow, iPtr, cols); + } + + var gpuWeights = gpu.UploadRaw(weightBytes, TensorShape.D1(weightBytes.Length), DType.Q8_0); + var gpuInput = gpu.Upload(input, TensorShape.D1(cols)); + var gpuOutput = gpu.Allocate(TensorShape.D1(rows)); + gpu.MatMul(gpuOutput, gpuWeights, gpuInput, DType.Q8_0); + gpu.Synchronize(); + var gpuResult = new float[rows]; + gpu.Download(gpuOutput, gpuResult); + gpu.Free(gpuWeights); + gpu.Free(gpuInput); + gpu.Free(gpuOutput); + + // Per-row magnitude scale for a relative bound (the dot of random ±1 + // activations over `cols` int8 weights has stddev ~ sqrt(cols)). + float refRms = 0f; + for (int r = 0; r < rows; r++) refRms += cpuOutput[r] * cpuOutput[r]; + refRms = MathF.Sqrt(refRms / rows); + + int mismatches = 0; + float maxAbs = 0; + for (int r = 0; r < rows; r++) + { + float diff = MathF.Abs(gpuResult[r] - cpuOutput[r]); + maxAbs = MathF.Max(maxAbs, diff); + // Q8 activation quant: per-element error ~ scale/2; allow 2% of the + // typical row magnitude as the absolute envelope. + if (diff > 0.02f * refRms) mismatches++; + } + Console.WriteLine($"MatVecQ8_0-dp4a rows={rows} cols={cols}: maxAbs={maxAbs:E2} refRms={refRms:E2} mismatches={mismatches}/{rows}"); + Assert.True(mismatches <= rows / 100 + 1, + $"dp4a Q8_0 matvec drifted from fp32 reference: {mismatches}/{rows} rows beyond 2% of row RMS ({refRms:E3}), maxAbs={maxAbs:E3}."); + } + } + [Fact] public void EmbedLookup_Q8_0_MatchesCpuReference() { diff --git a/tests/SharpInference.Tests.ForwardPass/Gemma4CudaBatchedPrefillTests.cs b/tests/SharpInference.Tests.ForwardPass/Gemma4CudaBatchedPrefillTests.cs index b92439f..8eef8fa 100644 --- a/tests/SharpInference.Tests.ForwardPass/Gemma4CudaBatchedPrefillTests.cs +++ b/tests/SharpInference.Tests.ForwardPass/Gemma4CudaBatchedPrefillTests.cs @@ -5,17 +5,21 @@ namespace SharpInference.Tests.ForwardPass; /// -/// Issue #136: the all-GPU batched-trunk prefill for Gemma 4 must produce the -/// same post-prompt logits as the bit-exact per-token -/// loop. Every batched primitive it uses (Q8_0/Q5_K/… GEMM-N, RoPEWithFactorsBatched, -/// AttentionSwaBatched, AttentionBatched, HeadNormBatched, RmsNormBatched, -/// KvAppendBatched, strided GeluTanhMul) is individually proven bit-identical to its -/// per-token form, so the assembled trunk should match bit-for-bit; the test allows a -/// hair of tolerance only to stay robust to driver-level reassociation. +/// Issue #141: the all-GPU batched-trunk prefill for Gemma 4 routes its Q8_0 trunk +/// matmuls through the compute-bound cuBLAS GEMM (), +/// which dequantizes the weight and rounds activations to fp16 before a tensor-core +/// multiply. That is not bit-exact to the per-token fp32 matvec loop, so this +/// oracle asserts argmax-stability plus a tolerance envelope rather than the +/// >95%-bit-exact match the pre-#141 matvec GEMM-N path satisfied. +/// +/// A second case () +/// keeps the strict bit-exact oracle for the matvec GEMM-N path (PrefillGemmEnabled=false), +/// so the #136 batched primitives stay individually verified. /// /// Single model instance, toggled via -/// with a ResetCache between runs (two 8 GB instances would not co-reside). -/// Silent-skips when CUDA or the GGUF is absent. +/// / with a ResetCache between +/// runs (two 8 GB instances would not co-reside). Silent-skips when CUDA or the GGUF +/// is absent. /// public sealed class Gemma4CudaBatchedPrefillTests { @@ -55,7 +59,7 @@ private static int Argmax(ReadOnlySpan logits) } [Fact] - public void Gemma4_E4B_BatchedPrefill_MatchesSequential() + public void Gemma4_E4B_BatchedPrefill_GemmMatchesSequential() { using var gpu = TryCreate(); if (gpu is null) return; @@ -73,45 +77,210 @@ public void Gemma4_E4B_BatchedPrefill_MatchesSequential() using var fwd = new CudaForwardPass(model, gpu, hp, maxContextLength: 512); - // Batched (default on). + // Batched cuBLAS-GEMM prefill (#141). Pin MMQ + flash off so this isolates the + // dequant→fp16→cuBLAS GEMM path (each has its own oracle below). fwd.BatchedPrefillEnabled = true; + fwd.PrefillGemmEnabled = true; + fwd.PrefillMmqEnabled = false; + fwd.PrefillFlashAttnEnabled = false; var batched = fwd.Prefill(tokens).ToArray(); Assert.True(fwd.LastPrefillWasBatched, "Batched-trunk prefill did not engage — check IsGemma4BatchedPrefillSupported gating."); - // Sequential per-token loop on the same instance. + // Sequential per-token loop on the same instance (the fp32 reference). + fwd.ResetCache(); + fwd.BatchedPrefillEnabled = false; + var sequential = fwd.Prefill(tokens).ToArray(); + Assert.False(fwd.LastPrefillWasBatched); + + Assert.Equal(sequential.Length, batched.Length); + + // The decisive parity signal: argmax must match the fp32 reference. The + // GEMM rounds weights + activations to fp16, so logits track to fp tolerance + // (post-softcap, |logit| ≲ FinalLogitSoftcap), not bit-for-bit. + Assert.Equal(Argmax(sequential), Argmax(batched)); + + float maxAbs = 0f; + for (int i = 0; i < sequential.Length; i++) + maxAbs = MathF.Max(maxAbs, MathF.Abs(sequential[i] - batched[i])); + // fp16 trunk over 42 layers + softcap: a few tenths of a logit is expected, + // a whole-number divergence would signal a real wiring bug. + Assert.True(maxAbs < 1.0f, + $"GEMM-batched vs sequential logits diverged beyond fp16 tolerance: maxAbs={maxAbs}."); + + // Top-5 set should be stable under fp16 rounding (order may shuffle below #1). + var seqTop = TopKSet(sequential, 5); + var batTop = TopKSet(batched, 5); + int overlap = 0; + foreach (var t in batTop) if (seqTop.Contains(t)) overlap++; + Assert.True(overlap >= 4, + $"GEMM-batched top-5 overlaps the fp32 reference in only {overlap}/5 slots; " + + "fp16 GEMM is diverging more than rounding explains."); + } + + [Fact] + public void Gemma4_E4B_BatchedPrefill_MmqMatchesSequential() + { + using var gpu = TryCreate(); + if (gpu is null) return; + var path = FindModelPath(); + if (path is null) return; + + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata, model); + Assert.NotNull(hp.LayerHeadDim); + Assert.True(hp.HasPerLayerTokenEmbd); + + var tokens = new int[] { 2, 651, 6037, 576, 6081, 603, 1234, 4567, 8901, 222, 333, 444 }; + + using var fwd = new CudaForwardPass(model, gpu, hp, maxContextLength: 512); + + // Int8 tensor-core MMQ prefill (#141, default on). Like the cuBLAS GEMM path + // it int8-quantizes both operands, so it is argmax-stable to the fp32 per-token + // reference, not bit-exact. + fwd.BatchedPrefillEnabled = true; + fwd.PrefillGemmEnabled = true; + fwd.PrefillMmqEnabled = true; + fwd.PrefillFlashAttnEnabled = false; // isolate the MMQ matmul (flash has its own oracle) + var batched = fwd.Prefill(tokens).ToArray(); + Assert.True(fwd.LastPrefillWasBatched, + "Batched-trunk MMQ prefill did not engage — check IsGemma4BatchedPrefillSupported gating."); + + fwd.ResetCache(); + fwd.BatchedPrefillEnabled = false; + var sequential = fwd.Prefill(tokens).ToArray(); + Assert.False(fwd.LastPrefillWasBatched); + + Assert.Equal(sequential.Length, batched.Length); + Assert.Equal(Argmax(sequential), Argmax(batched)); + + float maxAbs = 0f; + for (int i = 0; i < sequential.Length; i++) + maxAbs = MathF.Max(maxAbs, MathF.Abs(sequential[i] - batched[i])); + Assert.True(maxAbs < 1.0f, + $"MMQ-batched vs sequential logits diverged beyond int8 tolerance: maxAbs={maxAbs}."); + + var seqTop = TopKSet(sequential, 5); + var batTop = TopKSet(batched, 5); + int overlap = 0; + foreach (var t in batTop) if (seqTop.Contains(t)) overlap++; + Assert.True(overlap >= 4, + $"MMQ-batched top-5 overlaps the fp32 reference in only {overlap}/5 slots."); + } + + [Fact] + public void Gemma4_E4B_BatchedPrefill_FlashAttnMatchesSequential() + { + using var gpu = TryCreate(); + if (gpu is null) return; + var path = FindModelPath(); + if (path is null) return; + + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata, model); + Assert.NotNull(hp.LayerHeadDim); + Assert.True(hp.HasPerLayerTokenEmbd); + + var tokens = new int[] { 2, 651, 6037, 576, 6081, 603, 1234, 4567, 8901, 222, 333, 444 }; + + using var fwd = new CudaForwardPass(model, gpu, hp, maxContextLength: 512); + + // Flash-attention prefill (#141, default on). Its online softmax reassociates + // the score sum, so it is argmax-stable to the fp32 per-token reference, not + // bit-exact. Exercises both the SWA-windowed and global attention layers. + fwd.BatchedPrefillEnabled = true; + fwd.PrefillFlashAttnEnabled = true; + var batched = fwd.Prefill(tokens).ToArray(); + Assert.True(fwd.LastPrefillWasBatched); + fwd.ResetCache(); fwd.BatchedPrefillEnabled = false; var sequential = fwd.Prefill(tokens).ToArray(); Assert.False(fwd.LastPrefillWasBatched); Assert.Equal(sequential.Length, batched.Length); + Assert.Equal(Argmax(sequential), Argmax(batched)); + + float maxAbs = 0f; + for (int i = 0; i < sequential.Length; i++) + maxAbs = MathF.Max(maxAbs, MathF.Abs(sequential[i] - batched[i])); + Assert.True(maxAbs < 1.0f, + $"Flash-attn prefill vs sequential logits diverged beyond fp tolerance: maxAbs={maxAbs}."); + + var seqTop = TopKSet(sequential, 5); + var batTop = TopKSet(batched, 5); + int overlap = 0; + foreach (var t in batTop) if (seqTop.Contains(t)) overlap++; + Assert.True(overlap >= 4, + $"Flash-attn prefill top-5 overlaps the fp32 reference in only {overlap}/5 slots."); + } + + [Fact] + public void Gemma4_E4B_BatchedPrefill_GemmOff_MatchesSequentialBitExact() + { + using var gpu = TryCreate(); + if (gpu is null) return; + var path = FindModelPath(); + if (path is null) return; + + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata, model); + Assert.NotNull(hp.LayerHeadDim); + Assert.True(hp.HasPerLayerTokenEmbd); - // Argmax must match exactly; logits must match to a hair. + var tokens = new int[] { 2, 651, 6037, 576, 6081, 603, 1234, 4567, 8901, 222, 333, 444 }; + + // Pin Q8_0 matvec to the fp32-decode kernel on both sides so the matvec + // GEMM-N batched path can be compared bit-exactly to the per-token loop + // (the default dp4a path quantizes activations and is only argmax-stable). + gpu.Q80Dp4aEnabled = false; + + using var fwd = new CudaForwardPass(model, gpu, hp, maxContextLength: 512); + + // Batched matvec GEMM-N path (#136), bit-exact to per-token by construction. + // Flash attention pinned off: its online softmax reassociates the score sum + // and is only argmax-stable, which would break this bit-exact oracle (it + // verifies the batched matvec/norm/rope primitives, not the attention method). + fwd.BatchedPrefillEnabled = true; + fwd.PrefillGemmEnabled = false; + fwd.PrefillFlashAttnEnabled = false; + var batched = fwd.Prefill(tokens).ToArray(); + Assert.True(fwd.LastPrefillWasBatched); + + fwd.ResetCache(); + fwd.BatchedPrefillEnabled = false; + var sequential = fwd.Prefill(tokens).ToArray(); + Assert.False(fwd.LastPrefillWasBatched); + + Assert.Equal(sequential.Length, batched.Length); Assert.Equal(Argmax(sequential), Argmax(batched)); - float maxAbs = 0f, maxRel = 0f; + float maxAbs = 0f; int exact = 0; for (int i = 0; i < sequential.Length; i++) { if (BitConverter.SingleToInt32Bits(sequential[i]) == BitConverter.SingleToInt32Bits(batched[i])) exact++; - float d = MathF.Abs(sequential[i] - batched[i]); - maxAbs = MathF.Max(maxAbs, d); - maxRel = MathF.Max(maxRel, d / (MathF.Abs(sequential[i]) + 1e-6f)); + maxAbs = MathF.Max(maxAbs, MathF.Abs(sequential[i] - batched[i])); } Assert.True(maxAbs < 1e-2f, - $"Batched vs sequential logits diverged: maxAbs={maxAbs}, maxRel={maxRel}, " + - $"bit-exact {exact}/{sequential.Length}."); - - // Every batched primitive is individually bit-identical to its per-token form, - // so the assembled trunk should be too — assert a high bit-exact fraction (not - // just the loose maxAbs envelope) so a future single-kernel reassociation that - // stays under 1e-2 still trips this oracle. The maxAbs check above remains the - // net for any driver-level reassociation across the differing grid shapes. + $"matvec-batched vs sequential logits diverged: maxAbs={maxAbs}, bit-exact {exact}/{sequential.Length}."); + + // Every batched primitive is individually bit-identical to its per-token form. float exactFrac = (float)exact / sequential.Length; Assert.True(exactFrac > 0.95f, - $"Batched vs sequential only {exact}/{sequential.Length} ({exactFrac:P1}) bit-exact " + + $"matvec-batched vs sequential only {exact}/{sequential.Length} ({exactFrac:P1}) bit-exact " + $"(expected >95%); maxAbs={maxAbs}."); } + + private static HashSet TopKSet(ReadOnlySpan logits, int k) + { + var idx = new int[logits.Length]; + for (int i = 0; i < idx.Length; i++) idx[i] = i; + var arr = logits.ToArray(); + Array.Sort(idx, (a, b) => arr[b].CompareTo(arr[a])); + var set = new HashSet(); + for (int i = 0; i < k && i < idx.Length; i++) set.Add(idx[i]); + return set; + } }