diff --git a/README.md b/README.md index ef272286..7fcaae65 100644 --- a/README.md +++ b/README.md @@ -51,7 +51,7 @@ 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` | **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 -1 -c 2048` | **3698** | **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 **tensor-core flash-attention** prefill (#146/#147): both QK^T and P·V on the mma cores (`mma.m16n8k16.f16`), multi-warp **d-split** so the O tile stays register-resident — replaces the scalar O(n²) per-query attention (which re-streamed each query's K/V window up to ~512×) and is **+27% at ~1K / +40% at 1.8K** over the earlier half2 flash kernel (`SHARPI_PREFILL_FLASH_TC=0` reverts to half2, `=…_FLASH=0` to scalar) + a **SoA Q8_0 weight repack** (#149): all Q8_0 readers (MMQ, dp4a, fp32 matvec, GEMM-N, dequant) read the quants 16-byte-aligned with the fp16 scales split out, killing the `qs` 2-byte-misalignment funnelshift tax — **+10-12% prefill, bit-identical**; `SHARPI_MMQ_SOA=0` reverts) + a batched Q8_0 embedding lookup. **109→3698 at ~1K ctx, →4240 at 1.8K** — profiling showed *attention*, then the matmul inner-loop efficiency, were the dominant prefill costs at realistic prompt lengths. **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 (the SoA repack is bit-identical). Remaining gap to llama.cpp (~8475 prefill / ~78 decode): cp.async-pipelined MMQ on the SoA layout + 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 diff --git a/scripts/bench-gemma4-flashtc-146.ps1 b/scripts/bench-gemma4-flashtc-146.ps1 new file mode 100644 index 00000000..75334dee --- /dev/null +++ b/scripts/bench-gemma4-flashtc-146.ps1 @@ -0,0 +1,55 @@ +# A/B: half2 flash prefill (default) vs tensor-core flash prefill (#146, +# SHARPI_PREFILL_FLASH_TC=1). Same ~1K prompt + warm-discard as bench-gemma4-136.ps1, +# all-GPU (-g -1). Reports prefill t/s for each so the TC kernel can be judged +# against the half2 baseline (re-baselined at 2589 t/s on merged master). +param( + [string]$ModelDir = "E:\models", + [int]$Sections = 6 +) +$ErrorActionPreference = "Continue" +$model = Join-Path $ModelDir "gemma-4-E4B-it-Q8_0.gguf" + +$para = @( +"Modern large language model inference is dominated by memory bandwidth rather than raw compute.", +"Each decode step streams the full weight matrix for every layer, so quantization formats like Q4_K and Q5_K trade a small accuracy loss for a large reduction in bytes moved per token.", +"Mixture-of-experts models complicate this: only a handful of the hundreds of experts fire per token, and which fire varies token to token, defeating simple weight caching.", +"Gated DeltaNet layers replace quadratic attention with a linear recurrent state update, bounding per-token cost as context grows, at the price of a strictly sequential scan.", +"Hybrid placement keeps the attention and recurrent trunk on the accelerator while streaming routed-expert weights from host memory, overlapping the two so neither stalls." +) -join " " +$sb = [System.Text.StringBuilder]::new() +[void]$sb.Append("Read the following engineering notes and then write a concise technical summary.`n`n") +for ($i = 1; $i -le $Sections; $i++) { [void]$sb.Append("Section $i. $para`n`n") } +[void]$sb.Append("Summarize the main performance trade-offs across the sections above.") +$prompt = $sb.ToString() + +$cudaArgs = @("-g","-1","--backend","cuda","--ctx-size","2048") + +Write-Host "--- warming model ---" -ForegroundColor DarkGray +$null = .\scripts\bench-textgen.ps1 -Model $model -Tag "ftc-warm" -NTokens 8 -Prompt "Hello, world." -TimeoutSec 900 -ExtraArgs $cudaArgs + +$rows = @() +# variant: half2 | tc1 (single-warp #146) | tc2 (multi-warp #147) +function Run($tag, $variant) { + try { + Remove-Item env:SHARPI_PREFILL_FLASH_TC -ErrorAction SilentlyContinue + Remove-Item env:SHARPI_PREFILL_FLASH_TC1 -ErrorAction SilentlyContinue + if ($variant -eq "tc1") { $env:SHARPI_PREFILL_FLASH_TC = "1"; $env:SHARPI_PREFILL_FLASH_TC1 = "1" } + elseif ($variant -eq "tc2") { $env:SHARPI_PREFILL_FLASH_TC = "1" } + $null = .\scripts\bench-textgen.ps1 -Model $model -Tag "$tag-w" -NTokens 60 -Prompt $prompt -TimeoutSec 900 -ExtraArgs $cudaArgs + $r = .\scripts\bench-textgen.ps1 -Model $model -Tag $tag -NTokens 60 -Prompt $prompt -TimeoutSec 900 -ExtraArgs $cudaArgs + $script:rows += [PSCustomObject]@{ Tag=$tag; PrefTok=$r.PrefillTok; PrefillTps=$r.PrefillTps; DecodeTps=$r.DecodeTps; Wall=$r.WallSec; TO=$r.TimedOut; Sample=$r.Sample } + Write-Host (" {0,-18} pref={1,7} t/s dec={2,6} t/s ({3} tok)" -f $tag,$r.PrefillTps,$r.DecodeTps,$r.PrefillTok) -ForegroundColor Green + } finally { + Remove-Item env:SHARPI_PREFILL_FLASH_TC -ErrorAction SilentlyContinue + Remove-Item env:SHARPI_PREFILL_FLASH_TC1 -ErrorAction SilentlyContinue + } +} + +Run "flash-half2" "half2" +Run "flash-tc1" "tc1" +Run "flash-tc2" "tc2" + +Write-Host "" +Write-Host "=== Gemma 4 flash half2 vs TC (#146, ~1K ctx, warm) ===" -ForegroundColor Cyan +$rows | Format-Table Tag,PrefTok,PrefillTps,DecodeTps,Wall,TO -AutoSize +$rows | ForEach-Object { Write-Host (" {0}: {1}" -f $_.Tag, $_.Sample) -ForegroundColor DarkGray } diff --git a/scripts/bench-gemma4-mmq-soa-149.ps1 b/scripts/bench-gemma4-mmq-soa-149.ps1 new file mode 100644 index 00000000..e82a333d --- /dev/null +++ b/scripts/bench-gemma4-mmq-soa-149.ps1 @@ -0,0 +1,35 @@ +# A/B: Q8_0 SoA weight layout (#149, SHARPI_MMQ_SOA) off vs on, both with the TC +# flash default. Same ~1K/1.8K prompt + warm-discard, all-GPU (-g -1). +param([string]$ModelDir = "E:\models", [int]$Sections = 13) +$ErrorActionPreference = "Continue" +$model = Join-Path $ModelDir "gemma-4-E4B-it-Q8_0.gguf" +$para = @( +"Modern large language model inference is dominated by memory bandwidth rather than raw compute.", +"Each decode step streams the full weight matrix for every layer, so quantization formats like Q4_K and Q5_K trade a small accuracy loss for a large reduction in bytes moved per token.", +"Mixture-of-experts models complicate this: only a handful of the hundreds of experts fire per token, and which fire varies token to token, defeating simple weight caching.", +"Gated DeltaNet layers replace quadratic attention with a linear recurrent state update, bounding per-token cost as context grows, at the price of a strictly sequential scan.", +"Hybrid placement keeps the attention and recurrent trunk on the accelerator while streaming routed-expert weights from host memory, overlapping the two so neither stalls." +) -join " " +$sb = [System.Text.StringBuilder]::new() +[void]$sb.Append("Read the following engineering notes and then write a concise technical summary.`n`n") +for ($i = 1; $i -le $Sections; $i++) { [void]$sb.Append("Section $i. $para`n`n") } +[void]$sb.Append("Summarize the main performance trade-offs across the sections above.") +$prompt = $sb.ToString() +$cudaArgs = @("-g","-1","--backend","cuda","--ctx-size","2048") + +Write-Host "--- warming model ---" -ForegroundColor DarkGray +$null = .\scripts\bench-textgen.ps1 -Model $model -Tag "soa-warm" -NTokens 8 -Prompt "Hello, world." -TimeoutSec 900 -ExtraArgs $cudaArgs +$rows = @() +function Run($tag, $on) { + try { + if ($on) { $env:SHARPI_MMQ_SOA = "1" } else { Remove-Item env:SHARPI_MMQ_SOA -ErrorAction SilentlyContinue } + $null = .\scripts\bench-textgen.ps1 -Model $model -Tag "$tag-w" -NTokens 60 -Prompt $prompt -TimeoutSec 900 -ExtraArgs $cudaArgs + $r = .\scripts\bench-textgen.ps1 -Model $model -Tag $tag -NTokens 60 -Prompt $prompt -TimeoutSec 900 -ExtraArgs $cudaArgs + $script:rows += [PSCustomObject]@{ Tag=$tag; PrefTok=$r.PrefillTok; PrefillTps=$r.PrefillTps; DecodeTps=$r.DecodeTps; TO=$r.TimedOut } + Write-Host (" {0,-14} pref={1,7} t/s dec={2,6} t/s ({3} tok)" -f $tag,$r.PrefillTps,$r.DecodeTps,$r.PrefillTok) -ForegroundColor Green + } finally { Remove-Item env:SHARPI_MMQ_SOA -ErrorAction SilentlyContinue } +} +Run "soa-off" $false +Run "soa-on" $true +Write-Host "`n=== Gemma 4 MMQ SoA #149 (~$([int]($Sections*158))ctx, warm) ===" -ForegroundColor Cyan +$rows | Format-Table Tag,PrefTok,PrefillTps,DecodeTps,TO -AutoSize diff --git a/scripts/bench-gemma4-rebaseline.ps1 b/scripts/bench-gemma4-rebaseline.ps1 new file mode 100644 index 00000000..499987d9 --- /dev/null +++ b/scripts/bench-gemma4-rebaseline.ps1 @@ -0,0 +1,42 @@ +# Focused re-baseline of the all-GPU Gemma 4 prefill/decode on merged master, +# before the #146 full-TC-flash work. Same ~1K prompt + warm-discard methodology +# as bench-gemma4-136.ps1, but only the all-GPU (-g -1) cell, run twice (keep 2nd). +param( + [string]$ModelDir = "E:\models" +) +$ErrorActionPreference = "Continue" +$model = Join-Path $ModelDir "gemma-4-E4B-it-Q8_0.gguf" + +# Identical ~1K-token prompt to bench-allrows-1k.ps1 / bench-gemma4-136.ps1. +$para = @( +"Modern large language model inference is dominated by memory bandwidth rather than raw compute.", +"Each decode step streams the full weight matrix for every layer, so quantization formats like Q4_K and Q5_K trade a small accuracy loss for a large reduction in bytes moved per token.", +"Mixture-of-experts models complicate this: only a handful of the hundreds of experts fire per token, and which fire varies token to token, defeating simple weight caching.", +"Gated DeltaNet layers replace quadratic attention with a linear recurrent state update, bounding per-token cost as context grows, at the price of a strictly sequential scan.", +"Hybrid placement keeps the attention and recurrent trunk on the accelerator while streaming routed-expert weights from host memory, overlapping the two so neither stalls." +) -join " " +$sb = [System.Text.StringBuilder]::new() +[void]$sb.Append("Read the following engineering notes and then write a concise technical summary.`n`n") +for ($i = 1; $i -le 6; $i++) { [void]$sb.Append("Section $i. $para`n`n") } +[void]$sb.Append("Summarize the main performance trade-offs across the sections above.") +$prompt = $sb.ToString() + +$cudaArgs = @("-g","-1","--backend","cuda","--ctx-size","2048") + +Write-Host "--- warming model ---" -ForegroundColor DarkGray +$null = .\scripts\bench-textgen.ps1 -Model $model -Tag "g4-rb-warm" -NTokens 8 -Prompt "Hello, world." -TimeoutSec 900 -ExtraArgs $cudaArgs + +# Run twice, keep the second (warm). +$null = .\scripts\bench-textgen.ps1 -Model $model -Tag "g4-rb-w" -NTokens 60 -Prompt $prompt -TimeoutSec 900 -ExtraArgs $cudaArgs +$r = .\scripts\bench-textgen.ps1 -Model $model -Tag "g4-rb" -NTokens 60 -Prompt $prompt -TimeoutSec 900 -ExtraArgs $cudaArgs + +Write-Host "" +Write-Host "=== Gemma 4 all-GPU re-baseline (merged master, ~1K ctx, warm) ===" -ForegroundColor Cyan +[PSCustomObject]@{ + PrefillTok = $r.PrefillTok + PrefillTps = $r.PrefillTps + DecodeTps = $r.DecodeTps + WallSec = $r.WallSec + TimedOut = $r.TimedOut + Sample = $r.Sample +} | Format-List diff --git a/src/SharpInference.Cuda/CudaBackend.cs b/src/SharpInference.Cuda/CudaBackend.cs index 2644c1e3..fdebf042 100644 --- a/src/SharpInference.Cuda/CudaBackend.cs +++ b/src/SharpInference.Cuda/CudaBackend.cs @@ -137,6 +137,13 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis private nint _matvecQ80Kernel; // Issue #142: dp4a/Q8_1 decode matvec (quantize activation to int8, __dp4a dot). private nint _matvecQ80Dp4aKernel; + // Issue #149: SoA-layout dp4a decode matvec + the interleaved→SoA repack kernel. + private nint _matvecQ80Dp4aSoaKernel; + private nint _q80RepackSoaKernel; + // Issue #149: SoA variants of the remaining Q8_0 readers (fp32 matvec / GEMM-N / dequant). + private nint _matvecQ80SoaKernel; + private nint _matvecQ80GemmNSoaKernel; + private nint _dequantQ80F16SoaKernel; // 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. @@ -161,6 +168,9 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis // 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 #149: SoA-layout MMQ (quants 16B-aligned, scales separate) — kills the + // Q8_0 qs 2-byte-misalignment funnelshift tax in the weight load. + private nint _mmqQ80SoaKernel; // Issue #111: batched trunk elementwise/norm kernels (one launch over N tokens). private nint _rmsNormBatchedKernel; private nint _headNormBatchedKernel; @@ -215,6 +225,18 @@ public sealed unsafe class CudaBackend : IComputeBackend, IImageOpsBackend, IDis // full_seq / swa_batched kernels' O(n²) global K/V re-reads. private nint _flashAttnPrefillKernel; + // Issue #146: single-warp mma.sync m16n8k16 fragment-layout validation harness + // (de-risks the TC flash-attention fragment packing). Test-only. + private nint _mmaTestM16N8K16Kernel; + + // Issue #146: tensor-core flash-attention prefill (QK^T + P·V on the mma cores, + // shared-memory O). The compute-bound successor to the half2 flash kernel. + private nint _flashAttnPrefillTcKernel; + + // Issue #147: multi-warp / d-split TC flash (register-resident O, ~10× the + // single-warp occupancy). Default TC path when head_dim % 64 == 0. + private nint _flashAttnPrefillTc2Kernel; + // 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 // stays under a bounded budget. Freed in Dispose. @@ -509,6 +531,9 @@ public Tensor View(Tensor parent, long elemOffset, long elemCount, DType dtype = public void Free(Tensor tensor) { + // #149: drop any SoA-layout mark (harmless no-op for non-repacked handles) so + // the set doesn't grow across model load/free cycles. + _soaHandles.TryRemove(tensor.Handle, out _); if (_viewHandles.TryRemove(tensor.Handle, out _)) { // Non-owning view: drop the handle registration only; the parent owns @@ -1081,6 +1106,27 @@ public Tensor UploadRaw(ReadOnlySpan data, TensorShape shape, DType dtype, return tensor; } + /// Allocate an uninitialized device buffer of bytes + /// (issue #149 SoA repack destination). Mirrors 's allocation + /// without the host→device copy. The shape is a 1-D byte count for bookkeeping. + public Tensor AllocateRawBytes(long bytes, DType dtype, bool exact = false) + { + nuint byteSize = (nuint)bytes; + nuint allocSize = exact ? byteSize : GpuBufferPool.RoundUp(byteSize); + nint devPtr = exact ? nint.Zero : _pool.Rent(allocSize); + if (devPtr == nint.Zero) + { + int status = CuBlasInterop.CudaMalloc(out devPtr, allocSize); + if (status != 0) + throw new InvalidOperationException($"cudaMalloc failed: {status}"); + } + var handle = (nint)Interlocked.Increment(ref _nextHandle); + _devPtrs[handle] = (devPtr, allocSize); + if (exact) _exactHandles[handle] = 0; + _tensorDTypes[handle] = dtype; + return new Tensor(TensorShape.D1(bytes), dtype, handle); + } + // ── Async background upload path (issue #78) ────────────────────────── // // Mirrors the Vulkan UploadBackground entry point: dispatches the host→device @@ -1506,7 +1552,7 @@ public void MatMul(Tensor output, Tensor matrix, Tensor vector, DType weightDTyp // 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); + DispatchMatVecQ80Dp4a(wPtr, xPtr, yPtr, rows, cols, soa: _soaHandles.ContainsKey(matrix.Handle)); return; } @@ -1517,11 +1563,12 @@ public void MatMul(Tensor output, Tensor matrix, Tensor vector, DType weightDTyp (nint)(&pRows), (nint)(&pCols) }; + bool soa = weightDType == DType.Q8_0 && _soaHandles.ContainsKey(matrix.Handle); nint kernel = weightDType switch { DType.Q5_K => _matvecQ5KKernel, DType.Q6_K => _matvecQ6KKernel, - DType.Q8_0 => _matvecQ80Kernel, + DType.Q8_0 => soa ? _matvecQ80SoaKernel : _matvecQ80Kernel, DType.Float32 => _matvecF32Kernel, _ => throw new NotSupportedException($"CUDA MatMul: weight dtype {weightDType} not supported (expected Q4_K, Q5_K, Q6_K, Q8_0, or Float32)."), }; @@ -1654,11 +1701,12 @@ public void MatMulBatched(Tensor outputAll, Tensor matrix, Tensor inputAll, { // All take F32 input; the Q5_K/Q6_K/Q8_0 kernels decode the weight per // element. Same (rows+7)/8 × nTok geometry across all four. + bool soa = weightDType == DType.Q8_0 && _soaHandles.ContainsKey(matrix.Handle); nint kernel = weightDType switch { DType.Q6_K => _matvecQ6KGemmNKernel, DType.Q5_K => _matvecQ5KGemmNKernel, - DType.Q8_0 => _matvecQ80GemmNKernel, + DType.Q8_0 => soa ? _matvecQ80GemmNSoaKernel : _matvecQ80GemmNKernel, _ => _matvecF32GemmNKernel, }; int pRows = rows, pCols = cols, pN = nTok; @@ -1736,12 +1784,13 @@ public void MatMulBatchedGemm(Tensor outputAll, Tensor matrix, Tensor inputAll, EnsureGemmWf16((nuint)((long)rows * cols * 2L)); EnsureGemmAf16((nuint)((long)nTok * cols * 2L)); - // 1) Dequant Q8_0 weight → fp16 (one block per row). + // 1) Dequant Q8_0 weight → fp16 (one block per row). #149: SoA-aware. { nint wp = wPtr, op = _gemmWf16Buf; int pRows = rows, pCols = cols; + nint dqKern = _soaHandles.ContainsKey(matrix.Handle) ? _dequantQ80F16SoaKernel : _dequantQ80F16Kernel; nint* args = stackalloc nint[4] { (nint)(&wp), (nint)(&op), (nint)(&pRows), (nint)(&pCols) }; - int r = NvrtcInterop.LaunchKernel(_dequantQ80F16Kernel, (uint)rows, 1, 1, + int r = NvrtcInterop.LaunchKernel(dqKern, (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}"); } @@ -1842,12 +1891,108 @@ public void MatMulBatchedMmq(Tensor outputAll, Tensor matrix, Tensor inputAll, (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, + // #149: if this weight was repacked SoA, use the aligned-load kernel (same args). + nint kern = _soaHandles.ContainsKey(matrix.Handle) ? _mmqQ80SoaKernel : _mmqQ80Kernel; + int rm = NvrtcInterop.LaunchKernel(kern, gx, gy, 1, 256, 1, 1, 0, _stream, args, null); if (rm != 0) throw new InvalidOperationException($"cuLaunchKernel(mmq_q8_0) failed: {rm}"); } } + /// + /// Issue #149: SoA-layout variant of . Identical + /// math/output, but is a single buffer holding the + /// Q8_0 weight repacked struct-of-arrays — [quants rows*cols B][scales rows*nb + /// fp16] — so the 32 quants/block are contiguous & 16-byte aligned and the + /// weight load uses plain aligned uint reads instead of the funnelshift the + /// interleaved 34-byte block forces. The host splits the buffer at byte rows*cols + /// into the quant and scale pointers. Bit-identical to . + /// + public void MatMulBatchedMmqSoa(Tensor outputAll, Tensor soaWeight, Tensor inputAll, int nTok) + { + EnsureImageKernels(); + if (!_imageKernelsAvailable) + throw new NotSupportedException("NVRTC kernels are not available on this system."); + 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( + $"MatMulBatchedMmqSoa: 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 MatMulBatchedMmqSoa requires cols % 32 == 0 (got {cols})."); + + nint wPtr = GetDevPtr(soaWeight); + nint xPtr = GetDevPtr(inputAll); + nint yPtr = GetDevPtr(outputAll); + + int subBlocks = cols / 32; + long totalSub = (long)subBlocks * nTok; + EnsureQ81BatchBuf((nuint)(totalSub * 36L)); + { + nint qIn = xPtr, qOut = _q81BatchBuf; + long ce = (long)cols * nTok; + if (ce > int.MaxValue) + throw new InvalidOperationException($"MatMulBatchedMmqSoa: cols*nTok ({ce}) exceeds int range."); + int qN = (int)ce; + 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-soa) failed: {rq}"); + } + { + 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(_mmqQ80SoaKernel, gx, gy, 1, + 256, 1, 1, 0, _stream, args, null); + if (rm != 0) throw new InvalidOperationException($"cuLaunchKernel(mmq_q8_0_soa) failed: {rm}"); + } + } + + /// Issue #149: device handles of Q8_0 weights repacked into the SoA layout + /// (see ). and the decode + /// matvec auto-route to the aligned-load SoA kernels for these. + private readonly ConcurrentDictionary _soaHandles = new(); + + /// + /// Issue #149: allocate a new buffer the same size as the interleaved Q8_0 + /// [rows×nb×34 B], repack it into the SoA layout + /// [quants rows*cols B][scales rows*nb fp16] on the GPU, free , + /// and mark the new handle so the matmul/matvec dispatch uses the SoA kernels. + /// + public Tensor RepackQ8_0Soa(Tensor src, int rows, int cols) + { + EnsureImageKernels(); + if (!_imageKernelsAvailable) + throw new NotSupportedException("NVRTC kernels are not available on this system."); + if ((cols & 31) != 0) + throw new InvalidOperationException($"RepackQ8_0Soa requires cols % 32 == 0 (got {cols})."); + + long bytes = (long)rows * (cols / 32) * 34L; // == rows*cols + rows*nb*2 + var dst = AllocateRawBytes(bytes, DType.Q8_0, exact: true); + nint sPtr = GetDevPtr(src), dPtr = GetDevPtr(dst); + int pRows = rows, pCols = cols; + nint* args = stackalloc nint[4] { (nint)(&sPtr), (nint)(&dPtr), (nint)(&pRows), (nint)(&pCols) }; + long totalBlocks = (long)rows * (cols / 32); + uint grid = (uint)((totalBlocks + 255) / 256); + int r = NvrtcInterop.LaunchKernel(_q80RepackSoaKernel, grid, 1, 1, 256, 1, 1, 0, _stream, args, null); + if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(q8_0_repack_soa) failed: {r}"); + Synchronize(); + Free(src); + _soaHandles[dst.Handle] = 0; + return dst; + } + private void EnsureGemmWf16(nuint required) { if (_gemmWf16Buf != nint.Zero && _gemmWf16Size >= required) return; @@ -2074,7 +2219,7 @@ public void EnsureQ81Scratch(int cols) EnsureQ81Buf((nuint)((long)subBlocks * 36L)); } - private void DispatchMatVecQ80Dp4a(nint wPtr, nint xPtr, nint yPtr, int rows, int cols) + private void DispatchMatVecQ80Dp4a(nint wPtr, nint xPtr, nint yPtr, int rows, int cols, bool soa = false) { if ((cols & 31) != 0) throw new InvalidOperationException( @@ -2106,7 +2251,7 @@ private void DispatchMatVecQ80Dp4a(nint wPtr, nint xPtr, nint yPtr, int rows, in (nint)(&pRows), (nint)(&pCols) }; int rm = NvrtcInterop.LaunchKernel( - _matvecQ80Dp4aKernel, (uint)rows, 1, 1, + soa ? _matvecQ80Dp4aSoaKernel : _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}"); } @@ -2947,6 +3092,106 @@ public void FlashAttentionPrefill(Tensor qAll, Tensor kCache, Tensor vCache, Ten if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(flash_attn_prefill) failed: {r}"); } + /// + /// Issue #146 (test-only): one tensor-core mma.sync.m16n8k16.f32.f16.f16.f32 + /// computing [16×8] = [16×16] · + /// [16×8] on a single warp. Validates the A/B/C fragment + /// layouts used by the TC flash kernel; inputs are fp32 (rounded to fp16 inside), + /// so the result tracks an fp32 reference to fp16 tolerance, not bit-exactly. + /// is row-major [16*16], is K-major + /// [16*8] (bIn[k*8+n] = B[k][n]), is row-major [16*8]. + /// + public void MmaTestM16N8K16(Tensor aIn, Tensor bIn, Tensor cOut) + { + EnsureImageKernels(); + if (!_imageKernelsAvailable) + throw new NotSupportedException("NVRTC kernels are not available."); + + nint aP = GetDevPtr(aIn), bP = GetDevPtr(bIn), cP = GetDevPtr(cOut); + nint* args = stackalloc nint[3] { (nint)(&aP), (nint)(&bP), (nint)(&cP) }; + int r = NvrtcInterop.LaunchKernel(_mmaTestM16N8K16Kernel, 1, 1, 1, + 32, 1, 1, 0, _stream, args, null); + if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(mma_test_m16n8k16) failed: {r}"); + } + + /// + /// Issue #146: tensor-core flash-attention prefill. Drop-in for + /// — same args/semantics — but runs both + /// QK^T and P·V on the mma cores (one warp per 16-query tile, online softmax, + /// O accumulated in shared fp32). Requires % 16 == 0 + /// and ≤ 512 (shared = 16·headDim·6 B = 48 KB at d=512). Matches the scalar + /// kernels to fp tolerance (fp16 Q/K/V/P + online softmax), not bit-exact. + /// + public void FlashAttentionPrefillTc(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 || (headDim & 15) != 0) + throw new NotSupportedException( + $"FlashAttentionPrefillTc requires head_dim ≤ 512 and a multiple of 16; got {headDim}."); + + uint sharedBytes = (uint)(16 * headDim * 6); // O fp32 (×4) + K/V fp16 (×2) + + 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; + float pScale = attnScale; + nint* args = stackalloc nint[12] + { + (nint)(&qP), (nint)(&kP), (nint)(&vP), (nint)(&oP), + (nint)(&pNH), (nint)(&pNKV), (nint)(&pHD), + (nint)(&pSP), (nint)(&pWS), (nint)(&pMSL), (nint)(&pN), (nint)(&pScale) + }; + uint gy = (uint)((nTok + 15) / 16); + int r = NvrtcInterop.LaunchKernel(_flashAttnPrefillTcKernel, (uint)numHeads, gy, 1, + 32, 1, 1, sharedBytes, _stream, args, null); + if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(flash_attn_prefill_tc) failed: {r}"); + } + + /// + /// Issue #147: multi-warp / d-split tensor-core flash-attention prefill — same + /// args/semantics as but W=4 warps cooperate + /// on each 16-query tile with the head dim split across them, so O is register- + /// resident (no shared-O rescale) and occupancy rises ~10×. Requires + /// % 64 == 0 (W·16) and ≤ 512. Argmax-stable, not bit-exact. + /// + public void FlashAttentionPrefillTc2(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 || (headDim & 63) != 0) + throw new NotSupportedException( + $"FlashAttentionPrefillTc2 requires head_dim ≤ 512 and a multiple of 64 (W·16); got {headDim}."); + + const int w = 4; + uint sharedBytes = (uint)(16 * headDim * 2 + w * 256 * 4); // K/V fp16 + S scratch fp32 + + 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; + float pScale = attnScale; + nint* args = stackalloc nint[12] + { + (nint)(&qP), (nint)(&kP), (nint)(&vP), (nint)(&oP), + (nint)(&pNH), (nint)(&pNKV), (nint)(&pHD), + (nint)(&pSP), (nint)(&pWS), (nint)(&pMSL), (nint)(&pN), (nint)(&pScale) + }; + uint gy = (uint)((nTok + 15) / 16); + int r = NvrtcInterop.LaunchKernel(_flashAttnPrefillTc2Kernel, (uint)numHeads, gy, 1, + (uint)(w * 32), 1, 1, sharedBytes, _stream, args, null); + if (r != 0) throw new InvalidOperationException($"cuLaunchKernel(flash_attn_prefill_tc2) failed: {r}"); + } + /// /// Bf16-store variant of . Inputs stay fp32; the K/V /// cache tensors must be -allocated (half the @@ -4218,7 +4463,9 @@ private void ForceEagerJit() _matvecQ80Kernel, _matvecF32N2Kernel, _matvecQ4KN2Kernel, _matvecQ5KN2Kernel, _matvecQ6KN2Kernel, _matvecF32GemmNKernel, _matvecQ4KGemmNKernel, _matvecQ5KGemmNKernel, _matvecQ6KGemmNKernel, - _matvecQ80GemmNKernel, _mmqQ80Kernel, + _matvecQ80GemmNKernel, _mmqQ80Kernel, _mmqQ80SoaKernel, + _matvecQ80Dp4aSoaKernel, _q80RepackSoaKernel, + _matvecQ80SoaKernel, _matvecQ80GemmNSoaKernel, _dequantQ80F16SoaKernel, _rmsNormBatchedKernel, _headNormBatchedKernel, _headNormQkKernel, _headNormQkBatchedKernel, _splitQgBatchedKernel, _ropeNeoxPartialBatchedKernel, _attentionKernel, _attentionBf16Kernel, _attentionSwaKernel, _attentionSwaBatchedKernel, @@ -4233,7 +4480,8 @@ private void ForceEagerJit() _kvAppendBatchedKernel, _kvAppendBatchedBf16Kernel, _fullSeqAttentionKernel, _fullSeqAttentionBf16Kernel, _fullSeqAttentionGlobalKernel, _fullSeqAttentionGlobalBf16Kernel, - _flashAttnPrefillKernel, + _flashAttnPrefillKernel, _mmaTestM16N8K16Kernel, _flashAttnPrefillTcKernel, + _flashAttnPrefillTc2Kernel, ]; foreach (nint k in kernels) { @@ -4298,7 +4546,16 @@ private void LoadKernelFunctions() _dequantQ80F16Kernel = GetKernelFunc("llm_dequant_q8_0_to_f16"); _f32ToF16Kernel = GetKernelFunc("llm_f32_to_f16"); _mmqQ80Kernel = GetKernelFunc("llm_mmq_q8_0"); + _mmqQ80SoaKernel = GetKernelFunc("llm_mmq_q8_0_soa"); + _matvecQ80Dp4aSoaKernel = GetKernelFunc("llm_matvec_q8_0_dp4a_soa"); + _q80RepackSoaKernel = GetKernelFunc("llm_q8_0_repack_soa"); + _matvecQ80SoaKernel = GetKernelFunc("llm_matvec_q8_0_soa"); + _matvecQ80GemmNSoaKernel = GetKernelFunc("llm_matvec_q8_0_gemm_n_soa"); + _dequantQ80F16SoaKernel = GetKernelFunc("llm_dequant_q8_0_to_f16_soa"); _flashAttnPrefillKernel = GetKernelFunc("llm_flash_attn_prefill_f32"); + _mmaTestM16N8K16Kernel = GetKernelFunc("llm_mma_test_m16n8k16_f32"); + _flashAttnPrefillTcKernel = GetKernelFunc("llm_flash_attn_prefill_tc"); + _flashAttnPrefillTc2Kernel = GetKernelFunc("llm_flash_attn_prefill_tc2"); _rmsNormBatchedKernel = GetKernelFunc("llm_rmsnorm_batched"); _headNormBatchedKernel = GetKernelFunc("llm_head_norm_batched"); _headNormQkKernel = GetKernelFunc("llm_head_norm_qk"); @@ -4722,6 +4979,7 @@ public void Dispose() foreach (var entry in _devPtrs.Values) CuBlasInterop.CudaFree(entry.devPtr); _devPtrs.Clear(); + _soaHandles.Clear(); // #149 _pool.Dispose(); diff --git a/src/SharpInference.Cuda/CudaTextKernels.cs b/src/SharpInference.Cuda/CudaTextKernels.cs index 2b74e70e..d44a7eb2 100644 --- a/src/SharpInference.Cuda/CudaTextKernels.cs +++ b/src/SharpInference.Cuda/CudaTextKernels.cs @@ -1299,6 +1299,39 @@ __device__ __forceinline__ unsigned int sharpi_uint_at(const unsigned int* __res if (lane == 0) output[row] = result; } +// SoA-layout fp32 matvec (issue #149): bit-identical to llm_matvec_q8_0 but reads the +// Q8_0 weight from the SoA buffer [quants rows*cols B][scales rows*nb fp16] — aligned +// quant bytes, one aligned fp16 scale per block, no funnelshift. +extern ""C"" __global__ void llm_matvec_q8_0_soa( + const unsigned int* __restrict__ weights, + const float* __restrict__ input, + float* __restrict__ output, + int rows, int cols) +{ + const int N_ROWS = 8; + const int THREADS_PER_ROW = 32; + unsigned int tid = threadIdx.x; + int row_in_wg = (int)tid / THREADS_PER_ROW; + int lane = (int)tid & (THREADS_PER_ROW - 1); + int row = (int)blockIdx.x * N_ROWS + row_in_wg; + if (row >= rows) return; + + int num_blocks = cols >> 5; + long qrow = (long)row * cols; + const unsigned short* ws = (const unsigned short*)((const char*)weights + (long)rows * cols); + long srow = (long)row * num_blocks; + + float acc = 0.f; + for (int block = 0; block < num_blocks; block++) { + float d = sharpi_fp16_to_fp32(ws[srow + block]); + int q = sharpi_int8_at(weights, qrow + (long)block * 32 + lane); + float x = input[block * 32 + lane]; + acc += d * (float)q * x; + } + float result = sharpi_warp_reduce_sum(acc); + 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]), @@ -1385,6 +1418,84 @@ __device__ __forceinline__ unsigned int sharpi_uint_at(const unsigned int* __res } } +// SoA-layout dp4a decode matvec (issue #149). Same as llm_matvec_q8_0_dp4a but the +// Q8_0 weight is in the SoA buffer [quants rows*cols B][scales rows*nb fp16], so the +// quant words are plain aligned loads (no funnelshift) and the scale is one aligned +// fp16 read. Bit-identical to the AoS dp4a. ws locates the scale region at byte rows*cols. +extern ""C"" __global__ void llm_matvec_q8_0_dp4a_soa( + const unsigned int* __restrict__ weights, // SoA: [quants][scales] + 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; + int lane = (int)threadIdx.x; + int grp = lane >> 3; + int sub = lane & 7; + + int num_blocks = cols >> 5; + long qrow = (long)row * (cols >> 2); // uint index of this row's quants + const unsigned short* ws = (const unsigned short*)((const char*)weights + (long)rows * cols); + long srow = (long)row * num_blocks; // ushort index of this row's scales + + 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) { + float dw = sharpi_fp16_to_fp32(ws[srow + block]); + int wq = (int)weights[qrow + (long)block * 8 + sub]; // aligned, no funnelshift + + 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; + } + 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; + } + + __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; + } +} + +// One-time repack of an interleaved Q8_0 weight [rows × nb × 34 B] into the SoA +// buffer [quants rows*cols B][scales rows*nb fp16] (issue #149). One thread per +// 32-int8 block: copies the 2-byte fp16 scale to the scale region and the 32 quants +// to the (16-byte-aligned) quant region. Runs once at weight upload. +extern ""C"" __global__ void llm_q8_0_repack_soa( + const unsigned char* __restrict__ src, // interleaved, 34 B/block + unsigned char* __restrict__ dst, // SoA [quants][scales] + int rows, int cols) +{ + long blk = (long)blockIdx.x * blockDim.x + threadIdx.x; + int nb = cols >> 5; + long total = (long)rows * nb; + if (blk >= total) return; + long srcOff = blk * 34L; + long qDst = blk * 32L; // quants region + long sDst = (long)rows * cols + blk * 2L; // scales region + dst[sDst] = src[srcOff]; + dst[sDst + 1] = src[srcOff + 1]; + #pragma unroll + for (int i = 0; i < 32; i++) dst[qDst + i] = src[srcOff + 2 + i]; +} + // ── 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 @@ -1529,6 +1640,134 @@ __device__ __forceinline__ unsigned int sharpi_uint_at(const unsigned int* __res #undef MMQ_BM #undef MMQ_BN +// ── SoA-layout int8 MMQ (issue #149) ─────────────────────────────────────── +// Same tiling/mma as llm_mmq_q8_0, but the Q8_0 weights are pre-repacked into a +// struct-of-arrays layout: `qw` holds the 32 int8 quants per block contiguous and +// 16-byte aligned (8 uints/block, row-major by (row*nb+kb)), and `ws` holds the +// fp16 block scales separately. This kills the AoS tax: in the interleaved 34-byte +// layout the qs start at a 2-byte offset, so every weight word costs an extra load +// + __funnelshift (sharpi_uint_at); here every weight word is a plain aligned load. +// The roofline probe put the interleaved MMQ at 23-34% of int8 TC peak, inner-loop- +// bound on exactly this. Activations (Q8_1) are already 4-aligned, so only the +// weight-load path changes. Bit-identical to llm_mmq_q8_0 given a faithful repack. +#define MMQ_BM 64 +#define MMQ_BN 128 +extern ""C"" __global__ void llm_mmq_q8_0_soa( + const unsigned int* __restrict__ weights, // SoA buffer: [quants rows*cols B][scales rows*nb fp16] + 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) +{ + // Single SoA buffer; split into the quant and scale views (host stores both + // contiguous so the kernel signature matches llm_mmq_q8_0 — the dispatch just + // swaps the kernel function pointer based on whether the weight was repacked). + const unsigned int* qw = weights; + const unsigned short* ws = (const unsigned short*)((const char*)weights + (long)rows * cols); + + __shared__ int sW[MMQ_BM * 8]; + __shared__ float sWd[MMQ_BM]; + __shared__ int sY[MMQ_BN * 8]; + __shared__ float sYd[MMQ_BN]; + + 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; + int warp = tid >> 5; + int lane = tid & 31; + int grp = lane >> 2; + int tig = lane & 3; + int wr = warp & 3; + int wc = warp >> 2; + int mrow0 = wr * 16; + + float acc[8][4]; + #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, SoA weights: aligned uint loads (no funnelshift). + unsigned int rW0, rW1, rY0, rY1, rY2, rY3; + float rWd, rYd; + #define MMQ_LOAD_TILE_SOA(KB) do { \ + int gw0 = row_block + (tid >> 3); \ + rW0 = (gw0 < rows) ? qw[((long)gw0 * nb + (KB)) * 8L + (tid & 7)] : 0u; \ + int gw1 = row_block + ((tid + 256) >> 3); \ + rW1 = (gw1 < rows) ? qw[((long)gw1 * nb + (KB)) * 8L + ((tid + 256) & 7)] : 0u; \ + if (tid < MMQ_BM) { int gw = row_block + tid; \ + rWd = (gw < rows) ? sharpi_fp16_to_fp32(ws[(long)gw * nb + (KB)]) : 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_SOA(0); + + for (int kb = 0; kb < nb; kb++) { + 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(); + + if (kb + 1 < nb) MMQ_LOAD_TILE_SOA(kb + 1); + + 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_SOA +#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 @@ -1567,6 +1806,40 @@ __device__ __forceinline__ unsigned int sharpi_uint_at(const unsigned int* __res if (lane == 0) output[(long)token * (long)rows + row] = result; } +// SoA-layout GEMM-N (issue #149): bit-identical to llm_matvec_q8_0_gemm_n, aligned +// SoA weight reads. +extern ""C"" __global__ void llm_matvec_q8_0_gemm_n_soa( + const unsigned int* __restrict__ weights, + const float* __restrict__ input, + float* __restrict__ output, + int rows, int cols, int n_tok) +{ + const int N_ROWS = 8; + const int THREADS_PER_ROW = 32; + unsigned int tid = threadIdx.x; + int row_in_wg = (int)tid / THREADS_PER_ROW; + int lane = (int)tid & (THREADS_PER_ROW - 1); + int row = (int)blockIdx.x * N_ROWS + row_in_wg; + int token = (int)blockIdx.y; + if (row >= rows || token >= n_tok) return; + + int num_blocks = cols >> 5; + long qrow = (long)row * cols; + const unsigned short* ws = (const unsigned short*)((const char*)weights + (long)rows * cols); + long srow = (long)row * num_blocks; + const float* in = input + (long)token * (long)cols; + + float acc = 0.f; + for (int block = 0; block < num_blocks; block++) { + float d = sharpi_fp16_to_fp32(ws[srow + block]); + int q = sharpi_int8_at(weights, qrow + (long)block * 32 + lane); + float x = in[block * 32 + lane]; + acc += d * (float)q * x; + } + float result = sharpi_warp_reduce_sum(acc); + 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 @@ -1597,6 +1870,29 @@ __device__ __forceinline__ unsigned int sharpi_uint_at(const unsigned int* __res } } +// SoA-layout dequant (issue #149): bit-identical to llm_dequant_q8_0_to_f16, aligned +// SoA weight reads. +extern ""C"" __global__ void llm_dequant_q8_0_to_f16_soa( + const unsigned int* __restrict__ weights, + unsigned short* __restrict__ out, + int rows, int cols) +{ + int row = (int)blockIdx.x; + if (row >= rows) return; + int num_blocks = cols >> 5; + long qrow = (long)row * cols; + const unsigned short* ws = (const unsigned short*)((const char*)weights + (long)rows * cols); + long srow = (long)row * num_blocks; + 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; + float d = sharpi_fp16_to_fp32(ws[srow + block]); + int q = sharpi_int8_at(weights, qrow + (long)block * 32 + 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. @@ -3052,6 +3348,483 @@ __device__ __forceinline__ unsigned int sharpi_uint_at(const unsigned int* __res } } +// ── mma.sync m16n8k16 fragment-layout validation (issue #146) ────────────── +// Single-warp test harness for the int8/fp16 tensor-core building block used by +// the TC flash-attention prefill. One block / 32 threads computes C[16x8] = +// A[16x16] · B[16x8] on the tensor cores (fp16 multiplicands, fp32 accumulate) +// and writes C back in (row,col) order so a host unit test can compare against a +// CPU fp32 matmul. Validating the A/B/C fragment→lane→register maps in isolation +// (PTX ISA m16n8k16 .f16/.f32 tables) de-risks the full kernel — a wrong mapping +// silently produces garbage. a_in is row-major [16*16], b_in is row-major K-major +// [16*8] (b_in[k*8+n] = B[k][n]), c_out is row-major [16*8]. +extern ""C"" __global__ void llm_mma_test_m16n8k16_f32( + const float* __restrict__ a_in, // [16*16] row-major A + const float* __restrict__ b_in, // [16*8] K-major B (b_in[k*8+n]) + float* __restrict__ c_out) // [16*8] row-major C = A·B +{ + int lane = (int)(threadIdx.x & 31); + int gid = lane >> 2; // groupID 0..7 + int tig = lane & 3; // threadID_in_group 0..3 + + // A fragment (16x16 fp16, row-major): 4 regs, each a column-pair {2c,2c+1}. + unsigned int a0 = sharpi_f32x2_to_f16x2(a_in[(gid ) * 16 + (2 * tig )], a_in[(gid ) * 16 + (2 * tig + 1)]); + unsigned int a1 = sharpi_f32x2_to_f16x2(a_in[(gid + 8) * 16 + (2 * tig )], a_in[(gid + 8) * 16 + (2 * tig + 1)]); + unsigned int a2 = sharpi_f32x2_to_f16x2(a_in[(gid ) * 16 + (2 * tig + 8)], a_in[(gid ) * 16 + (2 * tig + 9)]); + unsigned int a3 = sharpi_f32x2_to_f16x2(a_in[(gid + 8) * 16 + (2 * tig + 8)], a_in[(gid + 8) * 16 + (2 * tig + 9)]); + + // B fragment (16x8 fp16, col-major): 2 regs, each a K-row-pair of column `gid`. + unsigned int b0 = sharpi_f32x2_to_f16x2(b_in[(2 * tig ) * 8 + gid], b_in[(2 * tig + 1) * 8 + gid]); + unsigned int b1 = sharpi_f32x2_to_f16x2(b_in[(2 * tig + 8) * 8 + gid], b_in[(2 * tig + 9) * 8 + gid]); + + float c0 = 0.f, c1 = 0.f, c2 = 0.f, c3 = 0.f; + asm volatile( + ""mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "" + ""{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"" + : ""+f""(c0), ""+f""(c1), ""+f""(c2), ""+f""(c3) + : ""r""(a0), ""r""(a1), ""r""(a2), ""r""(a3), ""r""(b0), ""r""(b1)); + + // C fragment (16x8 fp32): c0=(gid,2tig) c1=(gid,2tig+1) c2=(gid+8,2tig) c3=(gid+8,2tig+1). + c_out[(gid ) * 8 + (2 * tig )] = c0; + c_out[(gid ) * 8 + (2 * tig + 1)] = c1; + c_out[(gid + 8) * 8 + (2 * tig )] = c2; + c_out[(gid + 8) * 8 + (2 * tig + 1)] = c3; +} + +// ── Tensor-core flash-attention prefill (issue #146) ─────────────────────── +// Full TC version of llm_flash_attn_prefill_f32: both QK^T and P·V run on the +// tensor cores (mma.sync m16n8k16, fp16 multiplicands / fp32 accumulate). One +// warp owns a 16-query tile and streams the keys in 16-key tiles with an online +// softmax. The elegant part: the QK^T score C-fragment is reused *directly* as +// the P·V A-fragment — the key index is QK^T's N-column AND P·V's contraction +// dim, and the m16n8k16 C and A fragment layouts coincide on (row, 2·tig), so no +// transpose / shared round-trip is needed for P. O[16×head_dim] is too large for +// registers (head_dim/8 × 4 = 256 regs/lane at d=512), so it lives in shared fp32 +// and is rescaled in place per key-tile. K and V time-share one 16×head_dim fp16 +// shared buffer (K is fully consumed by QK^T before V is staged for P·V), so the +// whole kernel fits 16·head_dim·(4+2) = 48 KB at d=512. Requires head_dim%16==0. +// Matches the scalar kernels to fp tolerance (online softmax + fp16 Q/K/V/P), not +// bit-exact. GQA, causal, optional sliding window, per-layer head_dim. +__device__ __forceinline__ float fatc_mask(float s, int qpos, int abs_k, int window_size) +{ + bool ok = (abs_k <= qpos) && (window_size <= 0 || abs_k >= qpos + 1 - window_size); + return ok ? s : sharpi_neg_inf(); +} +#define FATC_KT 16 +extern ""C"" __global__ void llm_flash_attn_prefill_tc( + 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, float attn_scale) +{ + extern __shared__ unsigned int fatc_smem[]; + float* sO = (float*)fatc_smem; // [16 * head_dim] fp32 + unsigned short* sKV = (unsigned short*)(sO + 16 * head_dim); // [16 * head_dim] fp16 + + int lane = (int)(threadIdx.x & 31); + int gid = lane >> 2; // 0..7 (query-row base; also the N-col 0..7) + int tig = lane & 3; // 0..3 + int h = (int)blockIdx.x; + int qtile0 = (int)blockIdx.y * 16; + + int kv_head = h / (num_heads / num_kv_heads); + int kv_dim = num_kv_heads * head_dim; + int q_dim = num_heads * head_dim; + float scale = (attn_scale > 0.f) ? attn_scale : rsqrtf((float)head_dim); + int nchunk = head_dim >> 4; // d / 16 (QK^T contract chunks) + int ndtile = head_dim >> 3; // d / 8 (P·V head-dim N-tiles) + + int q0 = qtile0 + gid; // this lane's two query rows + int q1 = qtile0 + gid + 8; + int qpos0 = start_pos + q0; + int qpos1 = start_pos + q1; + + for (int idx = lane; idx < 16 * head_dim; idx += 32) sO[idx] = 0.f; + + float m0 = sharpi_neg_inf(), l0 = 0.f; // running max/sum, query row gid + float m1 = sharpi_neg_inf(), l1 = 0.f; // query row gid+8 + + // Union key range over the 16 queries in this tile (causal + optional window). + int last_q = qtile0 + 15; if (last_q > n_tok - 1) last_q = n_tok - 1; + int key_end = start_pos + last_q + 1; + int first_qpos = start_pos + qtile0; + int key_start = (window_size > 0) ? (first_qpos + 1 - window_size) : 0; + if (key_start < 0) key_start = 0; + key_start = (key_start / FATC_KT) * FATC_KT; // align down to a key tile + __syncthreads(); + + for (int kt0 = key_start; kt0 < key_end; kt0 += FATC_KT) { + // Stage K-tile [16 keys × head_dim] → sKV (fp16). Out-of-range keys load 0 + // (later masked to -inf by the causal bound). + for (int idx = lane; idx < FATC_KT * head_dim; idx += 32) { + int kk = idx / head_dim, d = idx - kk * head_dim; + int abs_k = kt0 + kk; + float kv = (abs_k < key_end) + ? k_cache[(long)abs_k * kv_dim + (long)kv_head * head_dim + d] : 0.f; + sKV[idx] = (unsigned short)sharpi_fp32_to_fp16(kv); + } + __syncthreads(); + + // ── QK^T → S[16q × 16k], two N=8 key sub-tiles (nt0: keys 0-7, nt1: 8-15) ── + float s0[4] = {0.f,0.f,0.f,0.f}; // nt0 C-frag c0..c3 + float s1[4] = {0.f,0.f,0.f,0.f}; // nt1 C-frag + for (int dc = 0; dc < nchunk; dc++) { + int d0 = dc * 16; + long qb = (long)h * head_dim + d0; + // A = Q frag: rows {gid→q0, gid+8→q1}, cols {2tig,2tig+1,2tig+8,2tig+9}. + float qa0l = (q0 < n_tok) ? q_all[(long)q0 * q_dim + qb + 2*tig ] : 0.f; + float qa0h = (q0 < n_tok) ? q_all[(long)q0 * q_dim + qb + 2*tig + 1] : 0.f; + float qa1l = (q1 < n_tok) ? q_all[(long)q1 * q_dim + qb + 2*tig ] : 0.f; + float qa1h = (q1 < n_tok) ? q_all[(long)q1 * q_dim + qb + 2*tig + 1] : 0.f; + float qa2l = (q0 < n_tok) ? q_all[(long)q0 * q_dim + qb + 2*tig + 8] : 0.f; + float qa2h = (q0 < n_tok) ? q_all[(long)q0 * q_dim + qb + 2*tig + 9] : 0.f; + float qa3l = (q1 < n_tok) ? q_all[(long)q1 * q_dim + qb + 2*tig + 8] : 0.f; + float qa3h = (q1 < n_tok) ? q_all[(long)q1 * q_dim + qb + 2*tig + 9] : 0.f; + unsigned int a0 = sharpi_f32x2_to_f16x2(qa0l, qa0h); + unsigned int a1 = sharpi_f32x2_to_f16x2(qa1l, qa1h); + unsigned int a2 = sharpi_f32x2_to_f16x2(qa2l, qa2h); + unsigned int a3 = sharpi_f32x2_to_f16x2(qa3l, qa3h); + // B = K frag (col-major; col = key within sub-tile, rows = contract d). + #pragma unroll + for (int nt = 0; nt < 2; nt++) { + int kbase = (nt * 8 + gid) * head_dim + d0; // key (nt*8+gid) + unsigned int b0 = ((unsigned int)sKV[kbase + 2*tig + 1] << 16) | (unsigned int)sKV[kbase + 2*tig ]; + unsigned int b1 = ((unsigned int)sKV[kbase + 2*tig + 9] << 16) | (unsigned int)sKV[kbase + 2*tig + 8]; + float* s = nt ? s1 : s0; + asm volatile( + ""mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "" + ""{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"" + : ""+f""(s[0]), ""+f""(s[1]), ""+f""(s[2]), ""+f""(s[3]) + : ""r""(a0), ""r""(a1), ""r""(a2), ""r""(a3), ""r""(b0), ""r""(b1)); + } + } + + // Scale + causal/window mask. C-frag: s0[c0]=S[q0,kt0+2tig], s0[c1]=S[q0,kt0+2tig+1], + // s0[c2]=S[q1,kt0+2tig], s0[c3]=S[q1,kt0+2tig+1]; s1 keys shifted by +8. + s0[0] = fatc_mask(s0[0]*scale, qpos0, kt0 + 2*tig, window_size); + s0[1] = fatc_mask(s0[1]*scale, qpos0, kt0 + 2*tig + 1, window_size); + s0[2] = fatc_mask(s0[2]*scale, qpos1, kt0 + 2*tig, window_size); + s0[3] = fatc_mask(s0[3]*scale, qpos1, kt0 + 2*tig + 1, window_size); + s1[0] = fatc_mask(s1[0]*scale, qpos0, kt0 + 2*tig + 8, window_size); + s1[1] = fatc_mask(s1[1]*scale, qpos0, kt0 + 2*tig + 9, window_size); + s1[2] = fatc_mask(s1[2]*scale, qpos1, kt0 + 2*tig + 8, window_size); + s1[3] = fatc_mask(s1[3]*scale, qpos1, kt0 + 2*tig + 9, window_size); + + // Online softmax, per query row. Row gid uses {s0[0],s0[1],s1[0],s1[1]} (this + // lane's 4 keys) reduced across the 4 lanes of the group (tig). Row gid+8 uses + // {s0[2],s0[3],s1[2],s1[3]}. + float tmax0 = fmaxf(fmaxf(s0[0], s0[1]), fmaxf(s1[0], s1[1])); + float tmax1 = fmaxf(fmaxf(s0[2], s0[3]), fmaxf(s1[2], s1[3])); + tmax0 = fmaxf(tmax0, __shfl_xor_sync(0xffffffffu, tmax0, 1)); + tmax0 = fmaxf(tmax0, __shfl_xor_sync(0xffffffffu, tmax0, 2)); + tmax1 = fmaxf(tmax1, __shfl_xor_sync(0xffffffffu, tmax1, 1)); + tmax1 = fmaxf(tmax1, __shfl_xor_sync(0xffffffffu, tmax1, 2)); + float mnew0 = fmaxf(m0, tmax0); + float mnew1 = fmaxf(m1, tmax1); + // A query row whose every key in this tile is masked (sliding window: a later + // query's early tiles fall entirely outside its window) keeps tmax = -inf; if + // no valid key has been seen yet mnew is also -inf, so exp(m-mnew)=exp(-inf+inf) + // is NaN. Guard: when mnew is -inf, skip the update (alpha=1 leaves O=0 intact, + // probabilities are 0). Once any valid key exists, mnew is finite and a masked + // score (-inf) safely yields exp(-inf - finite) = 0. + bool ok0 = mnew0 > sharpi_neg_inf(); + bool ok1 = mnew1 > sharpi_neg_inf(); + float alpha0 = ok0 ? __expf(m0 - mnew0) : 1.f; + float alpha1 = ok1 ? __expf(m1 - mnew1) : 1.f; + // Probabilities (unnormalized). Masked (-inf) or no-valid-key → 0. + float p0[4], p1[4]; + p0[0] = ok0 ? __expf(s0[0] - mnew0) : 0.f; p0[1] = ok0 ? __expf(s0[1] - mnew0) : 0.f; + p0[2] = ok1 ? __expf(s0[2] - mnew1) : 0.f; p0[3] = ok1 ? __expf(s0[3] - mnew1) : 0.f; + p1[0] = ok0 ? __expf(s1[0] - mnew0) : 0.f; p1[1] = ok0 ? __expf(s1[1] - mnew0) : 0.f; + p1[2] = ok1 ? __expf(s1[2] - mnew1) : 0.f; p1[3] = ok1 ? __expf(s1[3] - mnew1) : 0.f; + float lt0 = p0[0] + p0[1] + p1[0] + p1[1]; + float lt1 = p0[2] + p0[3] + p1[2] + p1[3]; + lt0 += __shfl_xor_sync(0xffffffffu, lt0, 1); lt0 += __shfl_xor_sync(0xffffffffu, lt0, 2); + lt1 += __shfl_xor_sync(0xffffffffu, lt1, 1); lt1 += __shfl_xor_sync(0xffffffffu, lt1, 2); + l0 = l0 * alpha0 + lt0; + l1 = l1 * alpha1 + lt1; + m0 = mnew0; m1 = mnew1; + + // P A-fragment (fp16), reusing the score layout directly (no transpose): + // a0=(gid; keys 2tig,2tig+1) a1=(gid+8; …) a2=(gid; keys 8+2tig,…) a3=(gid+8; …). + unsigned int pa0 = sharpi_f32x2_to_f16x2(p0[0], p0[1]); + unsigned int pa1 = sharpi_f32x2_to_f16x2(p0[2], p0[3]); + unsigned int pa2 = sharpi_f32x2_to_f16x2(p1[0], p1[1]); + unsigned int pa3 = sharpi_f32x2_to_f16x2(p1[2], p1[3]); + + // Stage V over the (now-consumed) K buffer. + __syncthreads(); + for (int idx = lane; idx < FATC_KT * head_dim; idx += 32) { + int kk = idx / head_dim, d = idx - kk * head_dim; + int abs_k = kt0 + kk; + float vv = (abs_k < key_end) + ? v_cache[(long)abs_k * kv_dim + (long)kv_head * head_dim + d] : 0.f; + sKV[idx] = (unsigned short)sharpi_fp32_to_fp16(vv); + } + __syncthreads(); + + // ── P·V → O[16q × head_dim], rescaling shared O in place per N=8 d-tile ── + for (int dt = 0; dt < ndtile; dt++) { + int dbase = dt * 8 + gid; // V col = head-dim index dt*8+gid + unsigned int b0 = ((unsigned int)sKV[(2*tig + 1) * head_dim + dbase] << 16) | (unsigned int)sKV[(2*tig ) * head_dim + dbase]; + unsigned int b1 = ((unsigned int)sKV[(2*tig + 9) * head_dim + dbase] << 16) | (unsigned int)sKV[(2*tig + 8) * head_dim + dbase]; + float o0 = 0.f, o1 = 0.f, o2 = 0.f, o3 = 0.f; + asm volatile( + ""mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "" + ""{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"" + : ""+f""(o0), ""+f""(o1), ""+f""(o2), ""+f""(o3) + : ""r""(pa0), ""r""(pa1), ""r""(pa2), ""r""(pa3), ""r""(b0), ""r""(b1)); + // o0=O[gid][dt*8+2tig] o1=O[gid][dt*8+2tig+1] o2=O[gid+8][…] o3=O[gid+8][…+1]. + int c0 = gid * head_dim + dt*8 + 2*tig; + int c1 = c0 + 1; + int c2 = (gid + 8) * head_dim + dt*8 + 2*tig; + int c3 = c2 + 1; + sO[c0] = sO[c0] * alpha0 + o0; + sO[c1] = sO[c1] * alpha0 + o1; + sO[c2] = sO[c2] * alpha1 + o2; + sO[c3] = sO[c3] * alpha1 + o3; + } + __syncthreads(); + } + + // Normalize and write out. The 4 lanes of a group share (gid, l0, l1) so they + // cooperatively stride the head dim by 4 (tig). + float inv0 = (l0 > 0.f) ? (1.f / l0) : 0.f; + float inv1 = (l1 > 0.f) ? (1.f / l1) : 0.f; + for (int d = tig; d < head_dim; d += 4) { + if (q0 < n_tok) out_all[(long)q0 * q_dim + (long)h * head_dim + d] = sO[gid * head_dim + d] * inv0; + if (q1 < n_tok) out_all[(long)q1 * q_dim + (long)h * head_dim + d] = sO[(gid + 8) * head_dim + d] * inv1; + } +} +#undef FATC_KT + +// ── Multi-warp / d-split tensor-core flash-attention prefill (issue #147) ─── +// Fixes the single-warp llm_flash_attn_prefill_tc occupancy limit (1 warp/block + +// 48 KB shared → ~2 warps/SM). Here a block is W warps that cooperate on ONE +// 16-query tile, splitting the head dim: warp w owns output columns [w·dW, …) with +// dW = head_dim/W, so O[16×dW] stays REGISTER-resident (16×128 = 64 regs/lane at +// d=512,W=4) instead of in shared — no per-key-tile shared-O rescale traffic, and +// the freed shared lets occupancy rise to ~16-20 warps/SM. Each warp computes a +// PARTIAL QK^T over its d-slice; the partials are summed across warps through a +// small shared S buffer ([W×16×16] fp32), after which every warp holds the full +// reduced score tile in its C-fragment and proceeds exactly like the single-warp +// kernel (in-warp softmax, no-transpose score→P, P·V for its d-slice). Requires +// head_dim % (W·16) == 0. Shared = 16·head_dim·2 (K/V fp16) + W·256·4 (S) B. +#define FATC2_W 4 +#define FATC2_KT 16 +extern ""C"" __global__ void llm_flash_attn_prefill_tc2( + const float* __restrict__ q_all, + const float* __restrict__ k_cache, + const float* __restrict__ v_cache, + float* __restrict__ out_all, + int num_heads, int num_kv_heads, int head_dim, + int start_pos, int window_size, int max_seq_len, int n_tok, float attn_scale) +{ + extern __shared__ unsigned int fatc2_smem[]; + unsigned short* sKV = (unsigned short*)fatc2_smem; // [16 * head_dim] fp16 + float* sS = (float*)(sKV + 16 * head_dim); // [W * 16 * 16] fp32 + + int tid = (int)threadIdx.x; + int warp = tid >> 5; // 0..W-1 (d-slice owner) + int lane = tid & 31; + int gid = lane >> 2; // 0..7 + int tig = lane & 3; // 0..3 + int h = (int)blockIdx.x; + int qtile0 = (int)blockIdx.y * 16; + + int kv_head = h / (num_heads / num_kv_heads); + int kv_dim = num_kv_heads * head_dim; + int q_dim = num_heads * head_dim; + float scale = (attn_scale > 0.f) ? attn_scale : rsqrtf((float)head_dim); + int dW = head_dim / FATC2_W; // this warp's head-dim slice width + int d_off = warp * dW; // first dim this warp owns + int nchunk = dW >> 4; // dW / 16 (QK^T contract chunks for the slice) + int ndt = dW >> 3; // dW / 8 (P·V N-tiles for the slice) + + int q0 = qtile0 + gid; + int q1 = qtile0 + gid + 8; + int qpos0 = start_pos + q0; + int qpos1 = start_pos + q1; + + // Register-resident O for this warp's d-slice: ndt N-tiles × 4 (c0..c3) fp32. + float oacc[64]; + #pragma unroll + for (int i = 0; i < 64; i++) oacc[i] = 0.f; + + float m0 = sharpi_neg_inf(), l0 = 0.f; + float m1 = sharpi_neg_inf(), l1 = 0.f; + + int last_q = qtile0 + 15; if (last_q > n_tok - 1) last_q = n_tok - 1; + int key_end = start_pos + last_q + 1; + int first_qpos = start_pos + qtile0; + int key_start = (window_size > 0) ? (first_qpos + 1 - window_size) : 0; + if (key_start < 0) key_start = 0; + key_start = (key_start / FATC2_KT) * FATC2_KT; + __syncthreads(); + + for (int kt0 = key_start; kt0 < key_end; kt0 += FATC2_KT) { + // Stage K-tile [16 × head_dim] → sKV (all warps cooperate over the full tile). + for (int idx = tid; idx < FATC2_KT * head_dim; idx += FATC2_W * 32) { + int kk = idx / head_dim, d = idx - kk * head_dim; + int abs_k = kt0 + kk; + float kv = (abs_k < key_end) + ? k_cache[(long)abs_k * kv_dim + (long)kv_head * head_dim + d] : 0.f; + sKV[idx] = (unsigned short)sharpi_fp32_to_fp16(kv); + } + __syncthreads(); + + // ── Partial QK^T over this warp's d-slice → s0/s1 (C-frag, 2 key sub-tiles) ── + float s0[4] = {0.f,0.f,0.f,0.f}; + float s1[4] = {0.f,0.f,0.f,0.f}; + for (int dc = 0; dc < nchunk; dc++) { + int d0 = d_off + dc * 16; // absolute head dim of this chunk + long qb = (long)h * head_dim + d0; + float qa0l = (q0 < n_tok) ? q_all[(long)q0 * q_dim + qb + 2*tig ] : 0.f; + float qa0h = (q0 < n_tok) ? q_all[(long)q0 * q_dim + qb + 2*tig + 1] : 0.f; + float qa1l = (q1 < n_tok) ? q_all[(long)q1 * q_dim + qb + 2*tig ] : 0.f; + float qa1h = (q1 < n_tok) ? q_all[(long)q1 * q_dim + qb + 2*tig + 1] : 0.f; + float qa2l = (q0 < n_tok) ? q_all[(long)q0 * q_dim + qb + 2*tig + 8] : 0.f; + float qa2h = (q0 < n_tok) ? q_all[(long)q0 * q_dim + qb + 2*tig + 9] : 0.f; + float qa3l = (q1 < n_tok) ? q_all[(long)q1 * q_dim + qb + 2*tig + 8] : 0.f; + float qa3h = (q1 < n_tok) ? q_all[(long)q1 * q_dim + qb + 2*tig + 9] : 0.f; + unsigned int a0 = sharpi_f32x2_to_f16x2(qa0l, qa0h); + unsigned int a1 = sharpi_f32x2_to_f16x2(qa1l, qa1h); + unsigned int a2 = sharpi_f32x2_to_f16x2(qa2l, qa2h); + unsigned int a3 = sharpi_f32x2_to_f16x2(qa3l, qa3h); + #pragma unroll + for (int nt = 0; nt < 2; nt++) { + int kbase = (nt * 8 + gid) * head_dim + d0; + unsigned int b0 = ((unsigned int)sKV[kbase + 2*tig + 1] << 16) | (unsigned int)sKV[kbase + 2*tig ]; + unsigned int b1 = ((unsigned int)sKV[kbase + 2*tig + 9] << 16) | (unsigned int)sKV[kbase + 2*tig + 8]; + float* s = nt ? s1 : s0; + asm volatile( + ""mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "" + ""{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"" + : ""+f""(s[0]), ""+f""(s[1]), ""+f""(s[2]), ""+f""(s[3]) + : ""r""(a0), ""r""(a1), ""r""(a2), ""r""(a3), ""r""(b0), ""r""(b1)); + } + } + + // ── Reduce partial scores across warps via shared, then read the full S back ── + // Each warp writes its C-frag to sS[warp][q][k]; lane (gid,tig) of every warp + // owns the same (q,k) cells, so after the barrier each lane sums over warps. + float* sw = sS + warp * 256; + sw[(gid ) * 16 + (2*tig )] = s0[0]; + sw[(gid ) * 16 + (2*tig + 1)] = s0[1]; + sw[(gid + 8) * 16 + (2*tig )] = s0[2]; + sw[(gid + 8) * 16 + (2*tig + 1)] = s0[3]; + sw[(gid ) * 16 + (8 + 2*tig )] = s1[0]; + sw[(gid ) * 16 + (8 + 2*tig + 1)] = s1[1]; + sw[(gid + 8) * 16 + (8 + 2*tig )] = s1[2]; + sw[(gid + 8) * 16 + (8 + 2*tig + 1)] = s1[3]; + __syncthreads(); + // Sum the W partials (start from this warp's own value, add the others). + #pragma unroll + for (int ww = 1; ww < FATC2_W; ww++) { + int o = ((warp + ww) % FATC2_W) * 256; + s0[0] += sS[o + (gid )*16 + 2*tig ]; + s0[1] += sS[o + (gid )*16 + 2*tig + 1]; + s0[2] += sS[o + (gid + 8)*16 + 2*tig ]; + s0[3] += sS[o + (gid + 8)*16 + 2*tig + 1]; + s1[0] += sS[o + (gid )*16 + 8 + 2*tig ]; + s1[1] += sS[o + (gid )*16 + 8 + 2*tig + 1]; + s1[2] += sS[o + (gid + 8)*16 + 8 + 2*tig ]; + s1[3] += sS[o + (gid + 8)*16 + 8 + 2*tig + 1]; + } + + // Scale + causal/window mask (full reduced scores now). + s0[0] = fatc_mask(s0[0]*scale, qpos0, kt0 + 2*tig, window_size); + s0[1] = fatc_mask(s0[1]*scale, qpos0, kt0 + 2*tig + 1, window_size); + s0[2] = fatc_mask(s0[2]*scale, qpos1, kt0 + 2*tig, window_size); + s0[3] = fatc_mask(s0[3]*scale, qpos1, kt0 + 2*tig + 1, window_size); + s1[0] = fatc_mask(s1[0]*scale, qpos0, kt0 + 2*tig + 8, window_size); + s1[1] = fatc_mask(s1[1]*scale, qpos0, kt0 + 2*tig + 9, window_size); + s1[2] = fatc_mask(s1[2]*scale, qpos1, kt0 + 2*tig + 8, window_size); + s1[3] = fatc_mask(s1[3]*scale, qpos1, kt0 + 2*tig + 9, window_size); + + // Online softmax (per query row; reduce the 4 keys across the group's 4 lanes). + float tmax0 = fmaxf(fmaxf(s0[0], s0[1]), fmaxf(s1[0], s1[1])); + float tmax1 = fmaxf(fmaxf(s0[2], s0[3]), fmaxf(s1[2], s1[3])); + tmax0 = fmaxf(tmax0, __shfl_xor_sync(0xffffffffu, tmax0, 1)); + tmax0 = fmaxf(tmax0, __shfl_xor_sync(0xffffffffu, tmax0, 2)); + tmax1 = fmaxf(tmax1, __shfl_xor_sync(0xffffffffu, tmax1, 1)); + tmax1 = fmaxf(tmax1, __shfl_xor_sync(0xffffffffu, tmax1, 2)); + float mnew0 = fmaxf(m0, tmax0); + float mnew1 = fmaxf(m1, tmax1); + bool ok0 = mnew0 > sharpi_neg_inf(); + bool ok1 = mnew1 > sharpi_neg_inf(); + float alpha0 = ok0 ? __expf(m0 - mnew0) : 1.f; + float alpha1 = ok1 ? __expf(m1 - mnew1) : 1.f; + float p0[4], p1[4]; + p0[0] = ok0 ? __expf(s0[0] - mnew0) : 0.f; p0[1] = ok0 ? __expf(s0[1] - mnew0) : 0.f; + p0[2] = ok1 ? __expf(s0[2] - mnew1) : 0.f; p0[3] = ok1 ? __expf(s0[3] - mnew1) : 0.f; + p1[0] = ok0 ? __expf(s1[0] - mnew0) : 0.f; p1[1] = ok0 ? __expf(s1[1] - mnew0) : 0.f; + p1[2] = ok1 ? __expf(s1[2] - mnew1) : 0.f; p1[3] = ok1 ? __expf(s1[3] - mnew1) : 0.f; + float lt0 = p0[0] + p0[1] + p1[0] + p1[1]; + float lt1 = p0[2] + p0[3] + p1[2] + p1[3]; + lt0 += __shfl_xor_sync(0xffffffffu, lt0, 1); lt0 += __shfl_xor_sync(0xffffffffu, lt0, 2); + lt1 += __shfl_xor_sync(0xffffffffu, lt1, 1); lt1 += __shfl_xor_sync(0xffffffffu, lt1, 2); + l0 = l0 * alpha0 + lt0; + l1 = l1 * alpha1 + lt1; + m0 = mnew0; m1 = mnew1; + + unsigned int pa0 = sharpi_f32x2_to_f16x2(p0[0], p0[1]); + unsigned int pa1 = sharpi_f32x2_to_f16x2(p0[2], p0[3]); + unsigned int pa2 = sharpi_f32x2_to_f16x2(p1[0], p1[1]); + unsigned int pa3 = sharpi_f32x2_to_f16x2(p1[2], p1[3]); + + // Stage V over the consumed K buffer. + __syncthreads(); + for (int idx = tid; idx < FATC2_KT * head_dim; idx += FATC2_W * 32) { + int kk = idx / head_dim, d = idx - kk * head_dim; + int abs_k = kt0 + kk; + float vv = (abs_k < key_end) + ? v_cache[(long)abs_k * kv_dim + (long)kv_head * head_dim + d] : 0.f; + sKV[idx] = (unsigned short)sharpi_fp32_to_fp16(vv); + } + __syncthreads(); + + // ── P·V over this warp's d-slice → register O, rescaled in place ── + for (int dt = 0; dt < ndt; dt++) { + int dcol = d_off + dt * 8 + gid; // absolute head dim (V col) + unsigned int b0 = ((unsigned int)sKV[(2*tig + 1) * head_dim + dcol] << 16) | (unsigned int)sKV[(2*tig ) * head_dim + dcol]; + unsigned int b1 = ((unsigned int)sKV[(2*tig + 9) * head_dim + dcol] << 16) | (unsigned int)sKV[(2*tig + 8) * head_dim + dcol]; + float o0 = 0.f, o1 = 0.f, o2 = 0.f, o3 = 0.f; + asm volatile( + ""mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "" + ""{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"" + : ""+f""(o0), ""+f""(o1), ""+f""(o2), ""+f""(o3) + : ""r""(pa0), ""r""(pa1), ""r""(pa2), ""r""(pa3), ""r""(b0), ""r""(b1)); + int base = dt * 4; + oacc[base + 0] = oacc[base + 0] * alpha0 + o0; // O[gid][dcol_2tig] + oacc[base + 1] = oacc[base + 1] * alpha0 + o1; + oacc[base + 2] = oacc[base + 2] * alpha1 + o2; // O[gid+8] + oacc[base + 3] = oacc[base + 3] * alpha1 + o3; + } + __syncthreads(); + } + + // Normalize and write this warp's d-slice from registers. oacc[dt] cells map to + // O[q=gid/gid+8][d_off + dt*8 + {2tig,2tig+1}]. + float inv0 = (l0 > 0.f) ? (1.f / l0) : 0.f; + float inv1 = (l1 > 0.f) ? (1.f / l1) : 0.f; + for (int dt = 0; dt < ndt; dt++) { + int d = d_off + dt * 8 + 2*tig; + int base = dt * 4; + if (q0 < n_tok) { + out_all[(long)q0 * q_dim + (long)h * head_dim + d ] = oacc[base + 0] * inv0; + out_all[(long)q0 * q_dim + (long)h * head_dim + d + 1] = oacc[base + 1] * inv0; + } + if (q1 < n_tok) { + out_all[(long)q1 * q_dim + (long)h * head_dim + d ] = oacc[base + 2] * inv1; + out_all[(long)q1 * q_dim + (long)h * head_dim + d + 1] = oacc[base + 3] * inv1; + } + } +} +#undef FATC2_W +#undef FATC2_KT + // ── 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 diff --git a/src/SharpInference.Engine/CudaForwardPass.cs b/src/SharpInference.Engine/CudaForwardPass.cs index 7e79c465..92b686e7 100644 --- a/src/SharpInference.Engine/CudaForwardPass.cs +++ b/src/SharpInference.Engine/CudaForwardPass.cs @@ -198,6 +198,18 @@ public sealed unsafe class CudaForwardPass : IForwardPass /// Default on (SHARPI_PREFILL_FLASH=0 reverts to the scalar kernels). /// public bool PrefillFlashAttnEnabled { get; set; } + + /// + /// Issues #146/#147: use the tensor-core flash-attention prefill (QK^T + P·V on the + /// mma cores) instead of the half2 kernel. Takes + /// precedence within the flash path. Requires head_dim % 16 == 0 (Gemma 4: 256/512); + /// the multi-warp #147 kernel (head_dim % 64 == 0) is +27-40% over half2, the + /// single-warp #146 fallback +5%. Default on (SHARPI_PREFILL_FLASH_TC=0 reverts + /// to half2). Argmax-stable, not bit-exact (fp16 Q/K/V/P + online softmax). + /// + public bool PrefillFlashTcEnabled { get; set; } + private bool _forceFlashTc1; // #147 A/B: pin the single-warp TC kernel + private readonly bool _mmqSoa; // #149: repack 2-D Q8_0 weights to SoA at upload 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] @@ -270,11 +282,18 @@ public sealed unsafe class CudaForwardPass : IForwardPass public CudaForwardPass(GgufModel model, CudaBackend gpu, ModelHyperparams hp, int maxContextLength = 0, - bool enableTurboQuant = false, int tqFp32Window = 256, int tqBits = 3) + bool enableTurboQuant = false, int tqFp32Window = 256, int tqBits = 3, + bool? mmqSoa = null) { _model = model; _gpu = gpu; _hp = hp; + // Issue #149: repack 2-D Q8_0 weights into the SoA layout at upload so all the + // Q8_0 readers (prefill MMQ, decode dp4a, fp32 matvec, GEMM-N, dequant) use + // aligned loads instead of the qs-misalignment funnelshift — +10-12% prefill, + // bit-identical. The backend auto-routes per repacked handle. Default on + // (SHARPI_MMQ_SOA=0 reverts). + _mmqSoa = mmqSoa ?? (Environment.GetEnvironmentVariable("SHARPI_MMQ_SOA") != "0"); _tqEnabled = enableTurboQuant; _tqBits = enableTurboQuant ? tqBits : 0; @@ -307,6 +326,15 @@ public CudaForwardPass(GgufModel model, CudaBackend gpu, ModelHyperparams hp, // ~1389→~2180 t/s. SHARPI_PREFILL_FLASH=0 reverts to the scalar kernels. PrefillFlashAttnEnabled = Environment.GetEnvironmentVariable("SHARPI_PREFILL_FLASH") != "0"; + // Issues #146/#147: tensor-core flash prefill — default on (the #147 multi-warp + // kernel is +27-40% over half2 on Gemma 4 at d=512). SHARPI_PREFILL_FLASH_TC=0 + // reverts to the half2 kernel. + PrefillFlashTcEnabled = + Environment.GetEnvironmentVariable("SHARPI_PREFILL_FLASH_TC") != "0" + && (_headDim & 15) == 0; + // Issue #147 A/B: force the single-warp TC kernel even where #147 would apply. + _forceFlashTc1 = + Environment.GetEnvironmentVariable("SHARPI_PREFILL_FLASH_TC1") == "1"; _maxHeadDim = _headDim; if (hp.LayerHeadDim is { } lhdMax) for (int i = 0; i < hp.NumLayers; i++) @@ -1991,7 +2019,18 @@ private void GpuLayerBatchedTrunkGemma4(int layer, int N, int startPos) if (s_prefillProfile) { _gpu.Synchronize(); _profSw.Restart(); } // Gemma 4: attention_scale = 1.0, passed explicitly (kernel skips its rsqrtf). - if (PrefillFlashAttnEnabled) + if (PrefillFlashTcEnabled && (layerHd & 15) == 0) + { + // #147 multi-warp/d-split when head_dim is a multiple of 64 (W·16); else the + // #146 single-warp kernel. SHARPI_PREFILL_FLASH_TC1=1 forces single-warp (A/B). + if (!_forceFlashTc1 && (layerHd & 63) == 0) + _gpu.FlashAttentionPrefillTc2(qAll, _gpuKCache[effLayer], _gpuVCache[effLayer], attnAll, + _numHeads, _numKvHeads, layerHd, startPos, isSwa ? window : 0, effLayerCtx, N, attnScale: 1f); + else + _gpu.FlashAttentionPrefillTc(qAll, _gpuKCache[effLayer], _gpuVCache[effLayer], attnAll, + _numHeads, _numKvHeads, layerHd, startPos, isSwa ? window : 0, effLayerCtx, N, attnScale: 1f); + } + else if (PrefillFlashAttnEnabled) _gpu.FlashAttentionPrefill(qAll, _gpuKCache[effLayer], _gpuVCache[effLayer], attnAll, _numHeads, _numKvHeads, layerHd, startPos, isSwa ? window : 0, effLayerCtx, N, attnScale: 1f); else if (isSwa) @@ -2384,6 +2423,14 @@ private Tensor UploadWeight(string name) else if (info.DType == DType.Q4_K || info.DType == DType.Q6_K || info.DType == DType.Q8_0) { result = _gpu.UploadRaw(data, TensorShape.D1(data.Length), info.DType, exact: true); + // #149: repack 2-D Q8_0 GEMM weights (norms/biases are 1-D; embedding uploads + // elsewhere) into the SoA layout. Dimensions are GGUF ne order: [cols, rows]. + if (_mmqSoa && info.DType == DType.Q8_0 && info.NDimensions == 2) + { + int cols = (int)info.Dimensions[0]; + int rows = (int)info.Dimensions[1]; + result = _gpu.RepackQ8_0Soa(result, rows, cols); + } _weightDTypes[result.Handle] = info.DType; } else diff --git a/tests/SharpInference.Tests.ForwardPass/CudaFlashAttnTcTests.cs b/tests/SharpInference.Tests.ForwardPass/CudaFlashAttnTcTests.cs new file mode 100644 index 00000000..0be568fc --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/CudaFlashAttnTcTests.cs @@ -0,0 +1,112 @@ +using SharpInference.Core; +using SharpInference.Cuda; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Issue #146: parity for the tensor-core flash-attention prefill +/// (, kernel +/// llm_flash_attn_prefill_tc) against the validated scalar batched kernels +/// ( / ). +/// The TC kernel rounds Q, K, V *and* the softmax probabilities P to fp16 for the +/// mma multiplicands (the half2 kernel kept V fp32), so it tracks the reference to a +/// looser fp16 tolerance. Same config matrix as the half2 flash test: GQA, both +/// Gemma 4 head_dims (256 SWA / 512 global), causal, windowing, partial last tile. +/// +/// Silent no-op on hosts without CUDA, matching the other Cuda* test files. +/// +public sealed unsafe class CudaFlashAttnTcTests +{ + private static CudaBackend? TryCreate() + { + if (!CudaBackend.IsAvailable()) return null; + try { return CudaBackend.Create(); } + catch { return null; } + } + + // (numHeads, numKvHeads, headDim, window, nTok). window=0 → global (full causal). + // All head_dims are multiples of 64 so the same matrix exercises tc2 (#147, W·16=64). + private static (int nh, int nkv, int hd, int win, int nTok)[] Configs() => new[] + { + (8, 2, 256, 0, 200), // global, SWA head_dim + (8, 2, 512, 0, 173), // global, global head_dim, partial last tile + (8, 2, 256, 64, 200), // sliding window 64 < nTok + (8, 2, 512, 96, 130), // sliding window 96, global head_dim + (4, 4, 128, 0, 64), // MHA (no GQA), small head_dim + }; + + [Fact] + public void FlashAttentionPrefillTc_MatchesScalarBatched() // #146 single-warp + { + using var gpu = TryCreate(); + if (gpu is null) return; + RunParity(gpu, tc2: false, label: "TC1"); + } + + [Fact] + public void FlashAttentionPrefillTc2_MatchesScalarBatched() // #147 multi-warp/d-split + { + using var gpu = TryCreate(); + if (gpu is null) return; + RunParity(gpu, tc2: true, label: "TC2"); + } + + private static void RunParity(CudaBackend gpu, bool tc2, string label) + { + foreach (var (nh, nkv, hd, win, nTok) in Configs()) + { + var rng = new Random(20260606 + 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); + 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 gTc = 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); + if (tc2) + gpu.FlashAttentionPrefillTc2(gq, gk, gv, gTc, nh, nkv, hd, startPos: 0, windowSize: win, maxSeqLen: nTok, nTok: nTok); + else + gpu.FlashAttentionPrefillTc(gq, gk, gv, gTc, nh, nkv, hd, startPos: 0, windowSize: win, maxSeqLen: nTok, nTok: nTok); + gpu.Synchronize(); + + var outRef = new float[q.Length]; + var outTc = new float[q.Length]; + gpu.Download(gRef, outRef); + gpu.Download(gTc, outTc); + gpu.Free(gq); gpu.Free(gk); gpu.Free(gv); gpu.Free(gRef); gpu.Free(gTc); + + 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(outTc[i] - outRef[i]); + // NaN/Inf must fail (a masked-tile online-softmax bug produces NaN, and + // NaN > threshold is false — count it explicitly so it can't slip through). + if (!float.IsFinite(outTc[i])) { mismatches++; continue; } + maxAbs = MathF.Max(maxAbs, diff); + // Q,K,V,P all fp16-rounded for the mma → ~fp16-relative attention error. + if (diff > 2e-2f * rms) mismatches++; + } + Console.WriteLine( + $"Flash{label} nh={nh} nkv={nkv} hd={hd} win={win} nTok={nTok}: maxAbs={maxAbs:E2} rms={rms:E2} mismatches={mismatches}/{outRef.Length}"); + // Allow a tiny tail of outliers from fp16 P·V accumulation; the bulk must match. + Assert.True(mismatches <= outRef.Length / 200 + 1, + $"{label} flash attention diverged from scalar reference: {mismatches}/{outRef.Length} beyond 2e-2·rms ({rms:E3}), maxAbs={maxAbs:E3} (nh={nh} hd={hd} win={win} nTok={nTok})."); + } + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/CudaMmaPrimitiveTests.cs b/tests/SharpInference.Tests.ForwardPass/CudaMmaPrimitiveTests.cs new file mode 100644 index 00000000..76ee74a2 --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/CudaMmaPrimitiveTests.cs @@ -0,0 +1,93 @@ +using SharpInference.Core; +using SharpInference.Cuda; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Issue #146: validates the mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 +/// tensor-core building block (kernel llm_mma_test_m16n8k16_f32, +/// ) in isolation before it is used by the +/// TC flash-attention prefill. A known A[16×16] and B[16×8] are multiplied on the +/// tensor cores (fp16 multiplicands, fp32 accumulate) and compared against a CPU +/// fp16-rounded reference. A wrong A/B/C fragment→lane→register map silently +/// produces garbage, so this is the go/no-go gate for the fragment layouts. +/// +/// Silent no-op on hosts without CUDA, matching the other Cuda* test files. +/// +public sealed unsafe class CudaMmaPrimitiveTests +{ + private static CudaBackend? TryCreate() + { + if (!CudaBackend.IsAvailable()) return null; + try { return CudaBackend.Create(); } + catch { return null; } + } + + [Fact] + public void MmaM16N8K16_F32_MatchesCpuReference() + { + using var gpu = TryCreate(); + if (gpu is null) return; + + const int M = 16, K = 16, N = 8; + var rng = new Random(20260606); + + // A[16×16] row-major, B[16×8] K-major (b[k*N+n] = B[k][n]). Keep magnitudes + // modest so fp16 rounding of the inputs stays the only meaningful error. + var a = new float[M * K]; + var b = new float[K * N]; + for (int i = 0; i < a.Length; i++) a[i] = (float)(rng.NextDouble() * 2 - 1); + for (int i = 0; i < b.Length; i++) b[i] = (float)(rng.NextDouble() * 2 - 1); + + // CPU reference with inputs rounded to fp16 (the tensor core rounds A,B to + // fp16 before the multiply; accumulation is fp32 on both sides). + var cRef = new float[M * N]; + for (int i = 0; i < M; i++) + for (int n = 0; n < N; n++) + { + float acc = 0f; + for (int k = 0; k < K; k++) + acc += (float)(Half)a[i * K + k] * (float)(Half)b[k * N + n]; + cRef[i * N + n] = acc; + } + + var gpuA = gpu.Upload(a, TensorShape.D1(a.Length)); + var gpuB = gpu.Upload(b, TensorShape.D1(b.Length)); + var gpuC = gpu.Allocate(TensorShape.D1(M * N)); + + gpu.MmaTestM16N8K16(gpuA, gpuB, gpuC); + gpu.Synchronize(); + + var cGpu = new float[M * N]; + gpu.Download(gpuC, cGpu); + gpu.Free(gpuA); + gpu.Free(gpuB); + gpu.Free(gpuC); + + // Per-element magnitude ~ sqrt(K) for ±1 inputs. fp16 mantissa is ~2^-11, so + // the dominant error is the input rounding; a few 1e-3 absolute is expected. + float maxAbs = 0f; + int mismatches = 0; + for (int i = 0; i < cRef.Length; i++) + { + float diff = MathF.Abs(cGpu[i] - cRef[i]); + maxAbs = MathF.Max(maxAbs, diff); + if (diff > 1e-2f) mismatches++; + } + Console.WriteLine($"mma m16n8k16: maxAbs={maxAbs:E3} mismatches={mismatches}/{cRef.Length}"); + // Dump the matrices on failure so a fragment-layout bug is diagnosable. + if (mismatches > 0) + { + Console.WriteLine(" row | gpu vs ref"); + for (int i = 0; i < M; i++) + { + var line = new System.Text.StringBuilder($" q{i,2}: "); + for (int n = 0; n < N; n++) + line.Append($"{cGpu[i * N + n],8:F3}/{cRef[i * N + n],-8:F3} "); + Console.WriteLine(line.ToString()); + } + } + Assert.True(mismatches == 0, + $"mma.sync m16n8k16 fragment layout produced wrong results: {mismatches}/{cRef.Length} elements off, maxAbs={maxAbs:E3}."); + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/CudaMmqRooflineProbe.cs b/tests/SharpInference.Tests.ForwardPass/CudaMmqRooflineProbe.cs new file mode 100644 index 00000000..e67554d6 --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/CudaMmqRooflineProbe.cs @@ -0,0 +1,90 @@ +using System.Diagnostics; +using SharpInference.Core; +using SharpInference.Cuda; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Issue #141/#145 roofline probe (not a correctness test): times the int8 MMQ +/// GEMM () at an FFN-representative prefill +/// shape and reports the achieved int8 TOPS, so the gap to the RTX 4070 Ti's ~160 +/// TOPS dense int8 tensor-core peak tells us how much headroom a pipelined MMQ +/// rewrite would have. Profiling (#146/#147) showed the matmul/FFN GEMMs are now the +/// dominant prefill cost; this quantifies whether MMQ is compute-saturated or not. +/// +/// Run explicitly: --filter FullyQualifiedName~CudaMmqRooflineProbe. Silent no-op +/// without CUDA. Always asserts true — it only prints the measurement. +/// +public sealed unsafe class CudaMmqRooflineProbe +{ + 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); + + private static byte[] BuildQ8_0Matrix(int rows, int cols, Random rng) + { + int blocksPerRow = cols / 32, 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; + ushort dHalf = HalfToUshort((Half)(float)(rng.NextDouble() * 0.09 + 0.01)); + bytes[off] = (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 Mmq_Q8_0_AchievedTops_AtFfnShape() + { + using var gpu = TryCreate(); + if (gpu is null) return; + + // FFN-ish prefill GEMMs for a ~Gemma-4-E4B layer (rows=out, cols=in, N=tokens). + (int rows, int cols, int nTok, string what)[] shapes = + { + (8192, 2048, 1024, "ffn-gate/up"), + (2048, 8192, 1024, "ffn-down"), + (6144, 2048, 1024, "qkv"), + }; + + var rng = new Random(20260606); + foreach (var (rows, cols, nTok, what) in shapes) + { + 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); + + 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)); + + // Warm up (incl. first-call NVRTC/quantize buffer alloc), then time. + for (int i = 0; i < 5; i++) gpu.MatMulBatchedMmq(gpuY, gpuW, gpuX, nTok, DType.Q8_0); + gpu.Synchronize(); + + const int iters = 50; + var sw = Stopwatch.StartNew(); + for (int i = 0; i < iters; i++) gpu.MatMulBatchedMmq(gpuY, gpuW, gpuX, nTok, DType.Q8_0); + gpu.Synchronize(); + sw.Stop(); + + gpu.Free(gpuW); gpu.Free(gpuX); gpu.Free(gpuY); + + double msPer = sw.Elapsed.TotalMilliseconds / iters; + double macs = (double)rows * cols * nTok; // multiply-accumulates + double tops = (2.0 * macs) / (msPer * 1e-3) / 1e12; // ×2 (MAC=mul+add) + double peak = 160.0; // ~RTX 4070 Ti dense int8 TC TOPS + Console.WriteLine( + $"MMQ {what,-12} [{rows}×{cols}]·[{cols}×{nTok}]: {msPer:F3} ms/call {tops:F1} int8 TOPS ({100*tops/peak:F0}% of ~{peak:F0} peak)"); + } + Assert.True(true); + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/CudaMmqSoaE2ETests.cs b/tests/SharpInference.Tests.ForwardPass/CudaMmqSoaE2ETests.cs new file mode 100644 index 00000000..9e3ea8dd --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/CudaMmqSoaE2ETests.cs @@ -0,0 +1,96 @@ +using SharpInference.Core; +using SharpInference.Cuda; +using SharpInference.Engine; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Issue #149: end-to-end check that repacking the Gemma 4 Q8_0 weights into the SoA +/// layout (mmqSoa: true) produces bit-identical logits to the interleaved +/// layout. The SoA prefill MMQ and decode dp4a kernels are each bit-identical to their +/// AoS counterparts (per ), so the whole forward pass — +/// batched prefill (MMQ) + greedy decode (dp4a) — must match to the bit. Two 8 GB +/// instances cannot co-reside, so this runs SoA-off, disposes, then SoA-on, and compares. +/// +/// Silent no-op without CUDA or the GGUF. +/// +public sealed class CudaMmqSoaE2ETests +{ + private const string ModelFile = "gemma-4-E4B-it-Q8_0.gguf"; + + private static CudaBackend? TryCreate() + { + if (!CudaBackend.IsAvailable()) return null; + try { return CudaBackend.Create(); } + catch { return null; } + } + + private static string? FindModelPath() + { + string[] absolute = { $@"E:\models\{ModelFile}", $@"C:\p\sharpi\models\{ModelFile}" }; + foreach (var p in absolute) + if (File.Exists(p)) return p; + return null; + } + + private static int Argmax(ReadOnlySpan v) + { + int best = 0; float bv = v[0]; + for (int i = 1; i < v.Length; i++) if (v[i] > bv) { bv = v[i]; best = i; } + return best; + } + + // Prefill the prompt, then greedily decode `nDecode` tokens, returning the logit + // array captured at prefill and at each decode step. + private static List RunForward(GgufModel model, CudaBackend gpu, ModelHyperparams hp, bool soa, + int[] tokens, int nDecode) + { + using var fwd = new CudaForwardPass(model, gpu, hp, maxContextLength: 512, mmqSoa: soa); + var captures = new List(); + + var logits = fwd.Prefill(tokens).ToArray(); + captures.Add(logits); + int pos = tokens.Length; + for (int i = 0; i < nDecode; i++) + { + int next = Argmax(logits); + logits = fwd.Forward(next, pos).ToArray(); + captures.Add(logits); + pos++; + } + return captures; + } + + [Fact] + public void Gemma4_E4B_MmqSoa_BitIdenticalToInterleaved() + { + 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); + + var tokens = new int[] { 2, 651, 6037, 576, 6081, 603, 1234, 4567, 8901, 222, 333, 444 }; + const int nDecode = 4; + + // Interleaved (default) first, then SoA — one instance at a time (8 GB each). + var refRun = RunForward(model, gpu, hp, soa: false, tokens, nDecode); + var soaRun = RunForward(model, gpu, hp, soa: true, tokens, nDecode); + + Assert.Equal(refRun.Count, soaRun.Count); + float maxAbs = 0f; + for (int step = 0; step < refRun.Count; step++) + { + var a = refRun[step]; + var b = soaRun[step]; + Assert.Equal(a.Length, b.Length); + for (int i = 0; i < a.Length; i++) + maxAbs = MathF.Max(maxAbs, MathF.Abs(a[i] - b[i])); + } + Console.WriteLine($"Gemma4 MMQ-SoA e2e: {refRun.Count} steps, maxAbs(SoA−AoS)={maxAbs:E3}"); + Assert.True(maxAbs == 0f, + $"SoA weight layout changed the Gemma 4 forward output (expected bit-identical): maxAbs={maxAbs:E3}."); + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/CudaMmqSoaTests.cs b/tests/SharpInference.Tests.ForwardPass/CudaMmqSoaTests.cs new file mode 100644 index 00000000..576c4de9 --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/CudaMmqSoaTests.cs @@ -0,0 +1,221 @@ +using System.Diagnostics; +using SharpInference.Core; +using SharpInference.Cuda; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Issue #149: the SoA-layout int8 MMQ (, +/// kernel llm_mmq_q8_0_soa) repacks Q8_0 weights so the 32 quants/block are +/// contiguous & 16-byte aligned and the fp16 scales are separate, eliminating the +/// 2-byte-misalignment __funnelshift the interleaved 34-byte block forces on +/// every weight word. This (a) checks it is bit-identical to the interleaved +/// (same int8 mma + scale math, only the +/// weight read differs) and (b) times both at an FFN-shaped prefill GEMM at the +/// real prefill nTok to measure the funnelshift-elimination speedup. +/// +/// Silent no-op without CUDA. +/// +public sealed unsafe class CudaMmqSoaTests +{ + 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); + + private static byte[] BuildQ8_0Matrix(int rows, int cols, Random rng) + { + int nb = cols / 32, bytesPerRow = nb * 34; + var bytes = new byte[rows * bytesPerRow]; + for (int r = 0; r < rows; r++) + for (int b = 0; b < nb; b++) + { + int off = r * bytesPerRow + b * 34; + ushort dHalf = HalfToUshort((Half)(float)(rng.NextDouble() * 0.09 + 0.01)); + bytes[off] = (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; + } + + /// Repack interleaved Q8_0 [34 B/block] → a single SoA buffer + /// [quants rows*cols B][scales rows*nb fp16] (matches CudaBackend's split at rows*cols). + private static byte[] RepackToSoA(byte[] interleaved, int rows, int cols) + { + int nb = cols / 32, bytesPerRow = nb * 34; + var soa = new byte[rows * cols + rows * nb * 2]; + int scaleBase = rows * cols; + for (int r = 0; r < rows; r++) + for (int b = 0; b < nb; b++) + { + int off = r * bytesPerRow + b * 34; + int blk = r * nb + b; + soa[scaleBase + blk * 2] = interleaved[off]; soa[scaleBase + blk * 2 + 1] = interleaved[off + 1]; + Array.Copy(interleaved, off + 2, soa, blk * 32, 32); + } + return soa; + } + + [Fact] + public void MatMulBatchedMmqSoa_BitIdenticalToInterleaved() + { + using var gpu = TryCreate(); + if (gpu is null) return; + + foreach ((int rows, int cols, int nTok) in new[] + { (256, 256, 8), (1024, 512, 12), (128, 2560, 64), (2048, 8192, 256), (6144, 2048, 200) }) + { + var rng = new Random(20260606 + rows * 31 + cols * 7 + nTok); + byte[] interleaved = BuildQ8_0Matrix(rows, cols, rng); + byte[] soa = RepackToSoA(interleaved, rows, cols); + + var acts = new float[nTok * cols]; + for (int i = 0; i < acts.Length; i++) acts[i] = (float)(rng.NextDouble() * 2 - 1); + + var gW = gpu.UploadRaw(interleaved, TensorShape.D1(interleaved.Length), DType.Q8_0); + var gSoa = gpu.UploadRaw(soa, TensorShape.D1(soa.Length), DType.Q8_0); + var gX = gpu.Upload(acts, TensorShape.D1(acts.Length)); + var gYi = gpu.Allocate(TensorShape.D1((long)nTok * rows)); + var gYs = gpu.Allocate(TensorShape.D1((long)nTok * rows)); + + gpu.MatMulBatchedMmq(gYi, gW, gX, nTok, DType.Q8_0); + gpu.MatMulBatchedMmqSoa(gYs, gSoa, gX, nTok); + gpu.Synchronize(); + + var yi = new float[nTok * rows]; + var ys = new float[nTok * rows]; + gpu.Download(gYi, yi); + gpu.Download(gYs, ys); + gpu.Free(gW); gpu.Free(gSoa); gpu.Free(gX); gpu.Free(gYi); gpu.Free(gYs); + + int diffs = 0; float maxAbs = 0; + for (int i = 0; i < yi.Length; i++) + { + float d = MathF.Abs(yi[i] - ys[i]); + maxAbs = MathF.Max(maxAbs, d); + if (d != 0f) diffs++; + } + Console.WriteLine($"MMQ-SoA rows={rows} cols={cols} nTok={nTok}: maxAbs={maxAbs:E2} diffs={diffs}/{yi.Length}"); + Assert.True(maxAbs == 0f, + $"SoA MMQ not bit-identical to interleaved: {diffs}/{yi.Length} differ, maxAbs={maxAbs:E3} (rows={rows} cols={cols} nTok={nTok})."); + } + } + + [Fact] + public void MmqSoa_Vs_Interleaved_Speed_AtRealPrefillNtok() + { + using var gpu = TryCreate(); + if (gpu is null) return; + + // FFN/qkv prefill GEMMs at a REALISTIC prefill nTok (~2K prompt batched at once), + // not the 1024 the earlier roofline probe used (which mismeasured occupancy). + (int rows, int cols, int nTok, string what)[] shapes = + { + (8192, 2048, 2048, "ffn-gate/up"), + (2048, 8192, 2048, "ffn-down"), + (6144, 2048, 2048, "qkv"), + }; + var rng = new Random(20260606); + foreach (var (rows, cols, nTok, what) in shapes) + { + byte[] interleaved = BuildQ8_0Matrix(rows, cols, rng); + byte[] soa = RepackToSoA(interleaved, rows, cols); + var acts = new float[nTok * cols]; + for (int i = 0; i < acts.Length; i++) acts[i] = (float)(rng.NextDouble() * 2 - 1); + + var gW = gpu.UploadRaw(interleaved, TensorShape.D1(interleaved.Length), DType.Q8_0); + var gSoa = gpu.UploadRaw(soa, TensorShape.D1(soa.Length), DType.Q8_0); + var gX = gpu.Upload(acts, TensorShape.D1(acts.Length)); + var gY = gpu.Allocate(TensorShape.D1((long)nTok * rows)); + + for (int i = 0; i < 5; i++) gpu.MatMulBatchedMmq(gY, gW, gX, nTok, DType.Q8_0); + gpu.Synchronize(); + const int iters = 50; + var sw = Stopwatch.StartNew(); + for (int i = 0; i < iters; i++) gpu.MatMulBatchedMmq(gY, gW, gX, nTok, DType.Q8_0); + gpu.Synchronize(); sw.Stop(); + double msAos = sw.Elapsed.TotalMilliseconds / iters; + + for (int i = 0; i < 5; i++) gpu.MatMulBatchedMmqSoa(gY, gSoa, gX, nTok); + gpu.Synchronize(); + sw.Restart(); + for (int i = 0; i < iters; i++) gpu.MatMulBatchedMmqSoa(gY, gSoa, gX, nTok); + gpu.Synchronize(); sw.Stop(); + double msSoa = sw.Elapsed.TotalMilliseconds / iters; + + gpu.Free(gW); gpu.Free(gSoa); gpu.Free(gX); gpu.Free(gY); + + double macs = (double)rows * cols * nTok; + double topsAos = 2.0 * macs / (msAos * 1e-3) / 1e12; + double topsSoa = 2.0 * macs / (msSoa * 1e-3) / 1e12; + Console.WriteLine( + $"MMQ-SoA {what,-12} nTok={nTok}: AoS {msAos:F3}ms ({topsAos:F1} TOPS) → SoA {msSoa:F3}ms ({topsSoa:F1} TOPS) {100*(msAos-msSoa)/msAos:+0.0;-0.0}%"); + } + Assert.True(true); + } + + /// + /// Exercises the production (GPU repack) and + /// the auto-routing through ALL five Q8_0 readers — dp4a + fp32 decode matvec, GEMM-N, + /// dequant→GEMM, and MMQ — by repacking on the GPU and asserting each reader's output + /// is bit-identical to the interleaved kernel. GGUF-free, runs on any CUDA box, so it + /// is the durable regression net the model-level oracles (bench-machine only) lack. + /// + [Fact] + public void GpuRepack_AllSoaReaders_BitIdenticalToInterleaved() + { + using var gpu = TryCreate(); + if (gpu is null) return; + + foreach ((int rows, int cols) in new[] { (256, 256), (1024, 512), (512, 2560) }) + { + var rng = new Random(20260607 + rows * 7 + cols); + byte[] interleaved = BuildQ8_0Matrix(rows, cols, rng); + const int nTok = 20; + + var vec = new float[cols]; + for (int i = 0; i < vec.Length; i++) vec[i] = (float)(rng.NextDouble() * 2 - 1); + var input = new float[(long)nTok * cols]; + for (int i = 0; i < input.Length; i++) input[i] = (float)(rng.NextDouble() * 2 - 1); + var gVec = gpu.Upload(vec, TensorShape.D1(cols)); + var gIn = gpu.Upload(input, TensorShape.D1(input.Length)); + + // Run `reader` on an interleaved weight and on a GPU-repacked SoA weight (the + // backend auto-routes the SoA handle to the aligned-load kernel) → bit-identical. + void Check(string label, int outElems, Action reader) + { + var gW = gpu.UploadRaw(interleaved, TensorShape.D1(interleaved.Length), DType.Q8_0); + var gWr = gpu.UploadRaw(interleaved, TensorShape.D1(interleaved.Length), DType.Q8_0); + var gSoa = gpu.RepackQ8_0Soa(gWr, rows, cols); // frees gWr, marks gSoa SoA + var gA = gpu.Allocate(TensorShape.D1(outElems)); + var gB = gpu.Allocate(TensorShape.D1(outElems)); + reader(gW, gA); // interleaved + reader(gSoa, gB); // SoA (auto-routed) + gpu.Synchronize(); + var a = new float[outElems]; var b = new float[outElems]; + gpu.Download(gA, a); gpu.Download(gB, b); + gpu.Free(gW); gpu.Free(gSoa); gpu.Free(gA); gpu.Free(gB); + float maxAbs = 0; int diffs = 0; + for (int i = 0; i < outElems; i++) { float d = MathF.Abs(a[i] - b[i]); maxAbs = MathF.Max(maxAbs, d); if (d != 0f) diffs++; } + Console.WriteLine($"SoA-reader {label,-16} rows={rows} cols={cols}: maxAbs={maxAbs:E2} diffs={diffs}/{outElems}"); + Assert.True(maxAbs == 0f, + $"GPU-repacked SoA reader '{label}' not bit-identical to interleaved: {diffs}/{outElems} differ, maxAbs={maxAbs:E3} (rows={rows} cols={cols})."); + } + + gpu.Q80Dp4aEnabled = true; + Check("decode-dp4a", rows, (w, o) => gpu.MatMul(o, w, gVec, DType.Q8_0)); + gpu.Q80Dp4aEnabled = false; + Check("decode-fp32", rows, (w, o) => gpu.MatMul(o, w, gVec, DType.Q8_0)); + gpu.Q80Dp4aEnabled = true; + Check("gemm-n", nTok * rows, (w, o) => gpu.MatMulBatched(o, w, gIn, nTok, DType.Q8_0)); + Check("dequant-gemm", nTok * rows, (w, o) => gpu.MatMulBatchedGemm(o, w, gIn, nTok, DType.Q8_0)); + Check("mmq", nTok * rows, (w, o) => gpu.MatMulBatchedMmq(o, w, gIn, nTok, DType.Q8_0)); + + gpu.Free(gVec); gpu.Free(gIn); + } + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/Gemma4CudaBatchedPrefillTests.cs b/tests/SharpInference.Tests.ForwardPass/Gemma4CudaBatchedPrefillTests.cs index 8eef8fa4..d9259fb6 100644 --- a/tests/SharpInference.Tests.ForwardPass/Gemma4CudaBatchedPrefillTests.cs +++ b/tests/SharpInference.Tests.ForwardPass/Gemma4CudaBatchedPrefillTests.cs @@ -83,6 +83,7 @@ public void Gemma4_E4B_BatchedPrefill_GemmMatchesSequential() fwd.PrefillGemmEnabled = true; fwd.PrefillMmqEnabled = false; fwd.PrefillFlashAttnEnabled = false; + fwd.PrefillFlashTcEnabled = false; var batched = fwd.Prefill(tokens).ToArray(); Assert.True(fwd.LastPrefillWasBatched, "Batched-trunk prefill did not engage — check IsGemma4BatchedPrefillSupported gating."); @@ -142,6 +143,7 @@ public void Gemma4_E4B_BatchedPrefill_MmqMatchesSequential() fwd.PrefillGemmEnabled = true; fwd.PrefillMmqEnabled = true; fwd.PrefillFlashAttnEnabled = false; // isolate the MMQ matmul (flash has its own oracle) + fwd.PrefillFlashTcEnabled = false; var batched = fwd.Prefill(tokens).ToArray(); Assert.True(fwd.LastPrefillWasBatched, "Batched-trunk MMQ prefill did not engage — check IsGemma4BatchedPrefillSupported gating."); @@ -168,6 +170,54 @@ public void Gemma4_E4B_BatchedPrefill_MmqMatchesSequential() $"MMQ-batched top-5 overlaps the fp32 reference in only {overlap}/5 slots."); } + [Fact] + public void Gemma4_E4B_BatchedPrefill_FlashTcMatchesSequential() + { + 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); + + // Tensor-core flash-attention prefill (#146). Like the half2 flash it is + // argmax-stable (online softmax + fp16 Q/K/V/P), not bit-exact. End-to-end + // full-model check on top of the per-kernel CudaFlashAttnTcTests parity, over + // the real per-layer head_dim mix (256 SWA / 512 global) and KV-share tail. + fwd.BatchedPrefillEnabled = true; + fwd.PrefillFlashTcEnabled = 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, + $"TC flash prefill vs sequential logits diverged beyond fp tolerance: maxAbs={maxAbs}."); + + var seqTop5 = TopKSet(sequential, 5); + var batTop5 = TopKSet(batched, 5); + int ov = 0; + foreach (var t in batTop5) if (seqTop5.Contains(t)) ov++; + Assert.True(ov >= 4, + $"TC flash prefill top-5 overlaps the fp32 reference in only {ov}/5 slots."); + } + [Fact] public void Gemma4_E4B_BatchedPrefill_FlashAttnMatchesSequential() { @@ -190,6 +240,7 @@ public void Gemma4_E4B_BatchedPrefill_FlashAttnMatchesSequential() // bit-exact. Exercises both the SWA-windowed and global attention layers. fwd.BatchedPrefillEnabled = true; fwd.PrefillFlashAttnEnabled = true; + fwd.PrefillFlashTcEnabled = false; // this oracle targets the half2 kernel; TC has its own var batched = fwd.Prefill(tokens).ToArray(); Assert.True(fwd.LastPrefillWasBatched); @@ -244,6 +295,7 @@ public void Gemma4_E4B_BatchedPrefill_GemmOff_MatchesSequentialBitExact() fwd.BatchedPrefillEnabled = true; fwd.PrefillGemmEnabled = false; fwd.PrefillFlashAttnEnabled = false; + fwd.PrefillFlashTcEnabled = false; var batched = fwd.Prefill(tokens).ToArray(); Assert.True(fwd.LastPrefillWasBatched);