diff --git a/src/SharpInference.Cli/RunCommand.cs b/src/SharpInference.Cli/RunCommand.cs index c41a2d44..14b70e5b 100644 --- a/src/SharpInference.Cli/RunCommand.cs +++ b/src/SharpInference.Cli/RunCommand.cs @@ -757,18 +757,34 @@ protected override int Execute(CommandContext context, Settings settings, Cancel if (settings.DraftModelPath is not null || settings.DraftLookup) { bool cudaSpecTarget = gpuFwd is CudaForwardPass { SupportsBatchVerify: true }; + // Sampled speculative decoding (issue #178): temp>0 now drives distribution-preserving + // spec sampling on the model-draft path (greedy at temp 0 stays byte-stable). Gated to + // model drafts (lookup proposals expose no q), to non-penalized/-biased sampling (draft + // and target must agree on the distribution), and bypassable via SHARPI_SPEC_SAMPLE=0. + bool sampledSpec = settings.Temperature > 0f; + bool specSampleDisabled = Environment.GetEnvironmentVariable("SHARPI_SPEC_SAMPLE") == "0"; + bool hasPenalty = sp.RepetitionPenalty != 1f && sp.PreviousTokens is { Count: > 0 }; + bool hasBias = sp.LogitBias is { Count: > 0 }; if (settings.DraftModelPath is not null && settings.DraftLookup) { AnsiConsole.MarkupLine("[red]Error:[/] --draft-model and --draft-lookup are mutually exclusive."); return 1; } - if (settings.Temperature > 0f) + if (nGpuLayers != 0 && !cudaSpecTarget) { - AnsiConsole.MarkupLine("[yellow]Warning:[/] Speculative decoding requires greedy sampling (--temp 0). Falling back to normal generation."); + AnsiConsole.MarkupLine("[yellow]Warning:[/] Speculative decoding requires pure CPU (-g 0) or full CUDA offload of a dense or Gemma-4 model. Falling back to normal generation."); } - else if (nGpuLayers != 0 && !cudaSpecTarget) + else if (sampledSpec && settings.DraftLookup) { - AnsiConsole.MarkupLine("[yellow]Warning:[/] Speculative decoding requires pure CPU (-g 0) or full CUDA offload of a dense model. Falling back to normal generation."); + AnsiConsole.MarkupLine("[yellow]Warning:[/] --draft-lookup supports greedy (--temp 0) only; sampled speculative decoding needs --draft-model. Falling back to normal generation."); + } + else if (sampledSpec && specSampleDisabled) + { + AnsiConsole.MarkupLine("[yellow]Note:[/] SHARPI_SPEC_SAMPLE=0 — sampled speculative decoding disabled; using normal sampled generation."); + } + else if (sampledSpec && (hasPenalty || hasBias)) + { + AnsiConsole.MarkupLine("[yellow]Warning:[/] sampled speculative decoding does not yet support --repeat-penalty / logit bias (draft and target must share the same distribution); falling back to normal generation."); } else if (settings.DraftLookup) { @@ -780,8 +796,8 @@ protected override int Execute(CommandContext context, Settings settings, Cancel IForwardPass lookupTarget = cudaSpecTarget ? (CudaForwardPass)gpuFwd! : fwd!; AnsiConsole.MarkupLine($"[dim]Speculative decoding: prompt-lookup (n-gram) drafting | Lookahead k={settings.SpecLookahead}[/]"); if (settings.Prompt is not null) - return RunSpeculativeSinglePrompt(settings, lookupTarget, null, tokenizer, sp); - return RunSpeculativeInteractive(settings, lookupTarget, null, tokenizer, sp); + return RunSpeculativeSinglePrompt(settings, lookupTarget, null, tokenizer, sp, rng); + return RunSpeculativeInteractive(settings, lookupTarget, null, tokenizer, sp, rng); } catch (Exception ex) { @@ -831,8 +847,8 @@ protected override int Execute(CommandContext context, Settings settings, Cancel using var draftFwd = new CudaForwardPass(draftModel, draftCuda, draftHp, draftCtx); AnsiConsole.MarkupLine($"[dim]Draft model: {draftHp.NumLayers}L, {draftHp.EmbeddingDim}d ([green]CUDA[/]) | Lookahead k={settings.SpecLookahead}[/]"); if (settings.Prompt is not null) - return RunSpeculativeSinglePrompt(settings, target, draftFwd, tokenizer, sp); - return RunSpeculativeInteractive(settings, target, draftFwd, tokenizer, sp); + return RunSpeculativeSinglePrompt(settings, target, draftFwd, tokenizer, sp, rng); + return RunSpeculativeInteractive(settings, target, draftFwd, tokenizer, sp, rng); } else { @@ -840,8 +856,8 @@ protected override int Execute(CommandContext context, Settings settings, Cancel using var draftFwd = new ForwardPass(draftModel, draftCpuBackend, draftHp); AnsiConsole.MarkupLine($"[dim]Draft model: {draftHp.NumLayers}L, {draftHp.EmbeddingDim}d ([blue]CPU[/]) | Lookahead k={settings.SpecLookahead}[/]"); if (settings.Prompt is not null) - return RunSpeculativeSinglePrompt(settings, fwd!, draftFwd, tokenizer, sp); - return RunSpeculativeInteractive(settings, fwd!, draftFwd, tokenizer, sp); + return RunSpeculativeSinglePrompt(settings, fwd!, draftFwd, tokenizer, sp, rng); + return RunSpeculativeInteractive(settings, fwd!, draftFwd, tokenizer, sp, rng); } } catch (Exception ex) @@ -905,7 +921,7 @@ private static bool SpecWindowExhausted(int promptTokens, private static int RunSpeculativeSinglePrompt(Settings s, IForwardPass target, IForwardPass? draft, - GgufTokenizer tok, SamplingParams sp) + GgufTokenizer tok, SamplingParams sp, Random rng) { var prompt = FormatPrompt(s.Prompt!, s.SystemPrompt, enableThinking: !s_noThinking); var tokens = tok.Encode(prompt); @@ -930,7 +946,10 @@ private static int RunSpeculativeSinglePrompt(Settings s, SpeculativeDecoder spec; if (draft is not null) { - spec = new SpeculativeDecoder(target, draft, s.SpecLookahead); + // temp>0 → sampled (distribution-preserving) accept; temp 0 → greedy (byte-stable). + spec = sp.Temperature > 0f + ? new SpeculativeDecoder(target, draft, sp, rng, s.SpecLookahead) + : new SpeculativeDecoder(target, draft, s.SpecLookahead); spec.Initialize(tokens.Count, targetLogits, draftLogits); } else @@ -974,11 +993,13 @@ private static int RunSpeculativeSinglePrompt(Settings s, private static int RunSpeculativeInteractive(Settings s, IForwardPass target, IForwardPass? draft, - GgufTokenizer tok, SamplingParams sp) + GgufTokenizer tok, SamplingParams sp, Random rng) { AnsiConsole.MarkupLine("[green]Interactive chat (speculative decoding).[/] Type a message, or [yellow]/exit[/] to quit.\n"); var spec = draft is not null - ? new SpeculativeDecoder(target, draft, s.SpecLookahead) + ? (sp.Temperature > 0f + ? new SpeculativeDecoder(target, draft, sp, rng, s.SpecLookahead) + : new SpeculativeDecoder(target, draft, s.SpecLookahead)) : new SpeculativeDecoder(target, new PromptLookupDraft(), s.SpecLookahead); while (true) diff --git a/src/SharpInference.Engine/CudaForwardPass.cs b/src/SharpInference.Engine/CudaForwardPass.cs index 68b8f514..99e36f90 100644 --- a/src/SharpInference.Engine/CudaForwardPass.cs +++ b/src/SharpInference.Engine/CudaForwardPass.cs @@ -4599,16 +4599,22 @@ internal float[][] BatchForwardMulti(int[] tokens, int[] positions, CudaSequence // ── Speculative-decode batched verify (issue #207) ────────────────────────────────── /// - /// Whether can run: the dense batched-decode configuration - /// ( — non-MoE, non-Gemma-4, no TurboQuant, no - /// final-logit softcap, GEMM-N-batchable weights) with an uncompacted cache. Unlike - /// , a CONFIGURED SnapKV budget does not disable - /// verify — only an actual prefill-time eviction does (then physical slot != logical + /// Whether can run: a batched-decode-capable configuration with an + /// uncompacted cache — either the dense path ( — + /// non-MoE, non-Gemma-4, no final-logit softcap) OR the Gemma-4 path + /// ( — per-layer head_dim, SWA rings, shared-KV, + /// k_eq_v, PLE, sandwich norms, and the final softcap the Gemma-4 finisher applies; issue + /// #178 GPU draft speculation). Both exclude TurboQuant and require GEMM-N-batchable weights. + /// Unlike , a CONFIGURED SnapKV budget does not + /// disable verify — only an actual prefill-time eviction does (then physical slot != logical /// position and the batched kernels would mis-index). Dynamic: flips false after such a /// prefill, so the speculative decoder (which re-checks per step) degrades to sequential - /// verify — the same once-evicted gating the GDN passes use (#130). + /// verify — the same once-evicted gating the GDN passes use (#130). (For Gemma-4 the SnapKV + /// budget is structurally off — the constructor forces it — so _kvEvictedCount is + /// always 0 there and the guard is a no-op.) /// - public bool SupportsBatchVerify => _kvEvictedCount == 0 && DenseBatchedDecodeSupported(); + public bool SupportsBatchVerify => + _kvEvictedCount == 0 && (DenseBatchedDecodeSupported() || Gemma4BatchedDecodeSupported()); /// /// Batched k-token verify for single-user speculative decoding (issue #207): one packed @@ -4629,15 +4635,23 @@ internal float[][] BatchForwardMulti(int[] tokens, int[] positions, CudaSequence /// the cache; the caller rewinds rejected tokens via . Issues /// direct launches only — the per-token decode CUDA graph (owned-cache pointers) stays /// valid for the surrounding Forward steps. + /// Gemma-4 (issue #178): dispatches the same packed pass + /// to , whose per-sequence attention loop appends each + /// row's K/V into the shared owned cache then attends in ascending row order — the same + /// append-then-attend causality. The k contiguous positions write distinct SWA ring slots + /// (k ≪ window), shared-KV layers read the source layer's already-filled cache, and the + /// finisher applies the Gemma-4 softcap; so the returned logits are the model's softcapped + /// logits, consistent with the single-token Gemma-4 the k==1 shortcut + /// uses. /// public float[][] BatchVerify(int[] tokens, int startPos) { ArgumentNullException.ThrowIfNull(tokens); if (!SupportsBatchVerify) throw new NotSupportedException( - "BatchVerify requires the dense batching-capable configuration (no MoE / " + - "Gemma-4 / TurboQuant / SnapKV / softcap, GEMM-N-batchable weights) and an " + - "uncompacted cache. Check SupportsBatchVerify before calling."); + "BatchVerify requires a batched-decode-capable configuration (dense or Gemma-4, " + + "no MoE / TurboQuant, GEMM-N-batchable weights) with an uncompacted cache. " + + "Check SupportsBatchVerify before calling."); int k = tokens.Length; if (k == 0) return Array.Empty(); if (startPos < 0 || startPos + k > _maxSeqLen) diff --git a/src/SharpInference.Engine/Sampler.cs b/src/SharpInference.Engine/Sampler.cs index 5bacbf4c..823459ee 100644 --- a/src/SharpInference.Engine/Sampler.cs +++ b/src/SharpInference.Engine/Sampler.cs @@ -99,6 +99,103 @@ public static int Sample(ReadOnlySpan logits, SamplingParams p, Random? r return SampleFromDistribution(probs, rng); } + /// + /// Build the full-vocabulary, post-filter sampling distribution for + /// under into (length must equal + /// logits.Length; entries outside the kept support are 0 and the kept entries sum to 1). + /// Applies the SAME pipeline as 's slow path — logit bias, repetition + /// penalty, temperature, softmax, top-k (+ renormalise), min-p, top-p, normalise — so a token + /// drawn from this distribution is distributed identically to . + /// + /// Used by distribution-preserving speculative sampling (issue #178), which needs the draft + /// proposal probability q(x) and the residual max(0, p − q) over a SHARED support: the draft + /// and target distributions MUST be built with the same for the + /// min(1, p/q) accept ratio to be meaningful. On Temperature ≤ 0 writes a one-hot at the + /// greedy argmax. + /// + public static void BuildFilteredDistribution(ReadOnlySpan logits, SamplingParams p, Span probs) + { + int vocabSize = logits.Length; + if (probs.Length != vocabSize) + throw new ArgumentException( + $"probs length ({probs.Length}) must equal logits length ({vocabSize}).", nameof(probs)); + + if (p.Temperature <= 0f) + { + probs.Clear(); + probs[Greedy(logits)] = 1f; + return; + } + + logits.CopyTo(probs); + + // Logit bias (additive, before temperature) — mirrors Sample's slow path. + if (p.LogitBias is { Count: > 0 }) + foreach (var (id, bias) in p.LogitBias) + if ((uint)id < (uint)vocabSize) + probs[id] += bias; + + // Repetition penalty (in logit space, before temperature). + if (p.RepetitionPenalty != 1.0f && p.PreviousTokens is { Count: > 0 }) + foreach (int id in p.PreviousTokens) + if ((uint)id < (uint)vocabSize) + probs[id] = probs[id] > 0f ? probs[id] / p.RepetitionPenalty : probs[id] * p.RepetitionPenalty; + + if (p.Temperature != 1.0f) + { + float invTemp = 1.0f / p.Temperature; + for (int i = 0; i < vocabSize; i++) + probs[i] *= invTemp; + } + + Softmax(probs); + + if (p.TopK > 0 && p.TopK < vocabSize) + { + ApplyTopK(probs, p.TopK); + Normalize(probs); + } + if (p.MinP > 0f) + ApplyMinP(probs, p.MinP); + if (p.TopP < 1.0f && p.TopP > 0f) + ApplyTopP(probs, p.TopP); + Normalize(probs); + } + + /// + /// Sample like but also materialise the full post-filter distribution + /// into (see ) and return the + /// drawn token, so probs[token] is its proposal probability q(token). On Temperature ≤ 0 + /// returns the greedy argmax. Distribution-preserving speculative sampling (issue #178) uses + /// this to propose draft tokens while retaining q for the accept/residual test. + /// + public static int SampleWithDistribution(ReadOnlySpan logits, SamplingParams p, Span probs, Random? rng = null) + { + if (p.Temperature <= 0f) + { + // One-hot at the greedy argmax — no filter pipeline, no RNG draw. + int argmax = Greedy(logits); + probs.Clear(); + probs[argmax] = 1f; + return argmax; + } + BuildFilteredDistribution(logits, p, probs); + rng ??= Random.Shared; + return SampleFromDistribution(probs, rng); + } + + /// + /// Draw a token from a pre-computed normalized distribution (e.g. one already built by + /// ), avoiding a redundant rebuild. Used by the + /// speculative decoder's looser-accept reject branch, where the target distribution at the + /// rejection position is already in hand. + /// + public static int SampleFromProbs(ReadOnlySpan probs, Random rng) + { + ArgumentNullException.ThrowIfNull(rng); + return SampleFromDistribution(probs, rng); + } + /// /// Top-k-first sampling fast path (no logit bias). Selects the top-k logits in one /// O(vocab) pass, then applies temperature, softmax, min-p, and top-p over only those @@ -367,6 +464,48 @@ private static int SampleFromDistribution(ReadOnlySpan probs, Random rng) return probs.Length - 1; } + /// + /// Sample a token proportional to the residual max(0, p[x] − q[x]) of two full-vocabulary + /// distributions built with the SAME (see + /// ). This is the rejection correction of speculative + /// sampling (Leviathan et al. / Chen et al.): drawing the corrected token from the residual + /// after a draft proposal is rejected makes the overall emitted-token distribution identical + /// to sampling directly from . If the residual mass is ≈0 (numerical / + /// q already dominated p on the kept support), falls back to sampling from . + /// + public static int ResampleResidual(ReadOnlySpan p, ReadOnlySpan q, Random rng) + { + ArgumentNullException.ThrowIfNull(rng); + if (p.Length != q.Length) + throw new ArgumentException( + $"p and q must have equal length ({p.Length} vs {q.Length}).", nameof(q)); + + float sum = 0f; + for (int i = 0; i < p.Length; i++) + { + float r = p[i] - q[i]; + if (r > 0f) sum += r; + } + if (sum <= 0f || float.IsNaN(sum)) + return SampleFromDistribution(p, rng); // degenerate: empty residual → fall back to p + + float target = (float)rng.NextDouble() * sum; + float cum = 0f; + int lastPositive = -1; + for (int i = 0; i < p.Length; i++) + { + float r = p[i] - q[i]; + if (r > 0f) + { + lastPositive = i; + cum += r; + if (target <= cum) return i; + } + } + // Rounding fallback: the last token with positive residual (tracked above, no extra pass). + return lastPositive >= 0 ? lastPositive : SampleFromDistribution(p, rng); + } + /// /// Find the k-th largest value in the span (1-indexed: k=1 is the largest). /// Simple O(n*k) selection — fine for Phase 1 correctness. diff --git a/src/SharpInference.Engine/SpeculativeDecoder.cs b/src/SharpInference.Engine/SpeculativeDecoder.cs index da622fe5..21e021d2 100644 --- a/src/SharpInference.Engine/SpeculativeDecoder.cs +++ b/src/SharpInference.Engine/SpeculativeDecoder.cs @@ -3,17 +3,31 @@ namespace SharpInference.Engine; /// -/// Speculative decoding (greedy): a small draft model proposes tokens which the target model -/// verifies via a batched forward pass, accepting each where they agree and correcting at the -/// first divergence. +/// Speculative decoding: a small draft model proposes tokens which the target model verifies +/// via a batched forward pass, accepting each where they agree and correcting at the first +/// divergence. /// -/// Each step packs the CERTAIN next token (argmax of the saved target logits) together with -/// k−1 draft proposals into ONE batched target pass (the llama.cpp formulation): the batch -/// yields both the verification logits and the next step's saved logits, so the target runs -/// exactly one batched pass per step — no separate correction-commit forward. On memory-bound -/// decode paths the batched pass costs ~1–2× a single forward (issue #194/#207 weight -/// amortization), so the speedup is ≈ E[tokens/step] / (cost_batch/cost_forward + draft -/// overhead), with E[tokens/step] = 1 + E[accepted of k−1] for per-token acceptance α. +/// Each step packs the CERTAIN next token together with k−1 draft proposals into ONE batched +/// target pass (the llama.cpp formulation): the batch yields both the verification logits and +/// the next step's certain token, so the target runs exactly one batched pass per step — no +/// separate correction-commit forward. On memory-bound decode paths the batched pass costs +/// ~1–2× a single forward (issue #194/#207 weight amortization), so the speedup is +/// ≈ E[tokens/step] / (cost_batch/cost_forward + draft overhead), with +/// E[tokens/step] = 1 + E[accepted of k−1] for per-token acceptance α. +/// +/// Two accept modes (issue #178): +/// +/// Greedy (temp ≤ 0, or the non-sampling ctor): the certain token is the argmax +/// of the saved target logits and a draft token is accepted iff the target's argmax +/// matches it. Byte-stable vs non-spec greedy. +/// Sampled (temp > 0, the sampling ctor): distribution-preserving speculative +/// sampling (Leviathan/Chen) — draft tokens are SAMPLED with proposal probability q, +/// accepted with prob min(1, p/q) against the target distribution p, and a rejection +/// resamples the correction from the residual max(0, p − q); a full accept samples a +/// bonus from the last verify position. The emitted-token distribution is identical to +/// direct target sampling. The opt-in --spec-draft-p-min (p ∈ (0,1)) selects a +/// looser, distribution-diverging accept (parity with the MTP path's pMin rule). +/// /// /// Both target and draft must share the same tokenizer (same vocab size). /// Note: does NOT take ownership of the forward pass instances. @@ -26,13 +40,29 @@ public sealed class SpeculativeDecoder private readonly bool _batchVerify; private int _lookahead; + // Sampled-accept mode (issue #178). _sampling is non-null only for temp > 0 + a draft + // model; otherwise the greedy path runs (byte-stable). _trueSpecSampling selects the + // distribution-preserving rule (default) vs the looser --spec-draft-p-min accept. + private readonly SamplingParams? _sampling; + private readonly Random? _rng; + private readonly bool _trueSpecSampling; + // Generation state. Invariant at step boundaries: both caches hold exactly _nextPos - // positions and _savedTargetLogits are the target's logits after the token at - // _nextPos−1 (so argmax(_savedTargetLogits) is the next emitted token, by greedy - // construction). The draft's own last logits are not part of the state — each step's - // first proposal requires forwarding the certain token through the draft anyway. + // positions. In GREEDY mode _savedTargetLogits are the target's logits after the token at + // _nextPos−1 (so argmax(_savedTargetLogits) is the next emitted token). In SAMPLED mode the + // next certain token was already drawn last step (the deferred correction/bonus) and is held + // in _savedSampledToken — re-sampling from logits would consume an extra RNG draw and break + // determinism. The draft's own last logits are not part of the state — each step's first + // proposal requires forwarding the certain token through the draft anyway. private int _nextPos; private float[] _savedTargetLogits; + private int _savedSampledToken; // sampled mode: next step's certain token + + // Sampled-mode scratch (lazily sized to vocab; _draftDists to the step's k). _pDist holds the + // target's filtered distribution at the verify position; _draftDists[i] retains the draft's + // full proposal distribution for token i so a rejection can resample from the residual. + private float[]? _pDist; + private float[][]? _draftDists; // Acceptance statistics private long _totalAccepted; @@ -97,6 +127,32 @@ public SpeculativeDecoder(IForwardPass target, IForwardPass draft, int lookahead _savedTargetLogits = new float[target.VocabSize]; } + /// + /// Sampled speculative decoding (issue #178): like the model-draft ctor, but verification + /// uses distribution-preserving speculative sampling when has + /// Temperature > 0. The SAME params drive the draft proposals + /// and the target accept test (so the min(1, p/q) ratio is well-defined), and the SAME + /// instance must be threaded from the caller for determinism. With + /// Temperature ≤ 0 this is exactly the greedy ctor (byte-stable). --spec-draft-p-min + /// in (0,1) opts into the looser, distribution-diverging accept. + /// + public SpeculativeDecoder(IForwardPass target, IForwardPass draft, SamplingParams sampling, Random rng, int lookahead = 4) + : this(target, draft, lookahead) + { + ArgumentNullException.ThrowIfNull(sampling); + ArgumentNullException.ThrowIfNull(rng); + if (sampling.Temperature > 0f) + { + _sampling = sampling; + _rng = rng; + // Default = true distribution-preserving sampling. SpecDraftPMin ∈ (0,1) → looser, + // distribution-diverging accept (same threshold SHAPE as MtpDecoder's pMin rule, but + // tested against the FILTERED target prob here — see StepSampled). 1.0 (default) / ≤0 + // → strict spec sampling. + _trueSpecSampling = !(sampling.SpecDraftPMin > 0f && sampling.SpecDraftPMin < 1f); + } + } + /// Adaptive lookahead: increase/decrease based on recent acceptance rate. public int Lookahead { @@ -135,6 +191,13 @@ public void Initialize(int prefillLength, ReadOnlySpan targetLogits, Read DraftMs = 0; VerifyMs = 0; CommitMs = 0; + if (_sampling is not null) + { + // Draw the first certain token from the prefill tail's target distribution; every + // later step inherits its certain token as the prior step's correction/bonus. + EnsureSampledScratch(_lookahead); + _savedSampledToken = Sampler.SampleWithDistribution(_savedTargetLogits, _sampling, _pDist!, _rng!); + } } /// @@ -166,7 +229,7 @@ public void Decode(int maxTokens, ReadOnlySpan stopTokenIds, Action em { int remaining = maxTokens - generated; int k = Math.Min(_lookahead, remaining); - int[] emitted = Step(k); + int[] emitted = _sampling is null ? Step(k) : StepSampled(k); foreach (int token in emitted) { @@ -277,6 +340,129 @@ private int[] Step(int k) return emitted; } + /// + /// One sampled speculative step (issue #178). Same folded structure as , + /// but the certain token is SAMPLED (drawn last step as the deferred correction/bonus), draft + /// proposals are sampled with proposal probability q, and acceptance is the distribution- + /// preserving rule: accept draft token i with prob min(1, p_i/q_i); on first rejection draw + /// the correction from the residual max(0, p_i − q_i); on full accept draw a bonus from the + /// last verify position. The correction/bonus is deferred to the next step's certain token + /// (so the target runs exactly one batched pass per step, like greedy). Returns the emitted + /// token array (the certain token + accepted proposals). + /// + private int[] StepSampled(int k) + { + int P = _nextPos; + var sampling = _sampling!; + var rng = _rng!; + EnsureSampledScratch(k); + var pDist = _pDist!; + var draftDists = _draftDists!; + + // tokens[0] is CERTAIN — it was sampled last step (the deferred correction/bonus), so it + // is NOT re-drawn here (re-sampling would consume an extra RNG draw and break determinism). + var tokens = new int[k]; + tokens[0] = _savedSampledToken; + + // ── Draft phase ────────────────────────────────────────────────────────── + // Each proposal is SAMPLED from the draft's filtered distribution; the full distribution + // is retained in draftDists[i] so a rejection at i can resample from the residual. + _phaseSw.Restart(); + for (int i = 1; i < k; i++) + { + var draftLogits = _draft!.Forward(tokens[i - 1], P + i - 1); + tokens[i] = Sampler.SampleWithDistribution(draftLogits, sampling, draftDists[i], rng); + } + DraftMs += _phaseSw.Elapsed.TotalMilliseconds; + + // ── Target batch-verify ────────────────────────────────────────────────── + // batch[i] = target logits AFTER tokens[i]; target cache advances to P + k. + _phaseSw.Restart(); + float[][] batch = BatchVerifyTarget(tokens, P); + VerifyMs += _phaseSw.Elapsed.TotalMilliseconds; + + // ── Speculative-sampling accept / reject ────────────────────────────────── + int accepted = 0; + int correction = -1; // residual/bonus token, deferred to next step's tokens[0] + for (int i = 1; i < k; i++) + { + // p_i = target distribution after tokens[i-1] (predicts the slot tokens[i] sits in). + Sampler.BuildFilteredDistribution(batch[i - 1], sampling, pDist); + float px = pDist[tokens[i]]; + bool accept; + if (_trueSpecSampling) + { + float qx = draftDists[i][tokens[i]]; + float a = qx > 0f ? MathF.Min(1f, px / qx) : (px > 0f ? 1f : 0f); + accept = rng.NextDouble() < a; + } + else + { + // Looser opt-in (--spec-draft-p-min): accept iff the target's prob of the draft + // token ≥ pMin OR the draft token is the target's own argmax. Diverges from the + // target distribution. Same threshold shape as MtpDecoder.AcceptDraft, but px is + // the FILTERED target prob (temp/top-k/top-p/min-p applied) — consistent with the + // p used by the strict path — whereas MtpDecoder thresholds a raw temp-1 softmax, + // so a given pMin is not numerically identical across the two paths. + accept = px >= sampling.SpecDraftPMin || tokens[i] == Sampler.Greedy(batch[i - 1]); + } + if (accept) { accepted++; continue; } + + // Reject at i: correction from the residual at this position (true sampling) or a + // fresh target sample (looser pMin mode — already off-distribution). pDist already + // holds BuildFilteredDistribution(batch[i-1]) from the accept test above, so the + // looser path samples it directly rather than rebuilding it. + correction = _trueSpecSampling + ? Sampler.ResampleResidual(pDist, draftDists[i], rng) + : Sampler.SampleFromProbs(pDist, rng); + break; + } + if (correction < 0) + { + // All k−1 drafts accepted: bonus token sampled from the final verify position. + correction = Sampler.SampleWithDistribution(batch[k - 1], sampling, pDist, rng); + } + + _totalAccepted += accepted; + _totalEmitted += accepted + 1; + + // ── Roll caches back; defer the correction/bonus to next step ───────────── + _phaseSw.Restart(); + int newPos = P + 1 + accepted; + _target.TruncateTo(newPos); + if (accepted == k - 1) + // Full accept: the draft never forwarded tokens[k-1] (its cache is at P+k-1). Sync it + // so the next chain starts at newPos. (Identical to Step's full-accept branch.) + _draft!.Forward(tokens[^1], P + k - 1); + else + _draft!.TruncateTo(newPos); + CommitMs += _phaseSw.Elapsed.TotalMilliseconds; + + // ── Update state ────────────────────────────────────────────────────────── + _nextPos = newPos; + _savedSampledToken = correction; // emitted next step as its certain token + + // Emit tokens[0..accepted]; the correction/bonus rides into the next step (folded form). + var emitted = new int[accepted + 1]; + for (int i = 0; i <= accepted; i++) emitted[i] = tokens[i]; + return emitted; + } + + /// Lazily (re)allocate the sampled-mode scratch: the target distribution buffer and + /// per-draft-position distribution buffers, each vocab-sized, growing _draftDists to k. + private void EnsureSampledScratch(int k) + { + int v = _target.VocabSize; + _pDist ??= new float[v]; + if (_draftDists is null || _draftDists.Length < k) + { + var old = _draftDists; + _draftDists = new float[k][]; + for (int i = 0; i < k; i++) + _draftDists[i] = old is not null && i < old.Length ? old[i] : new float[v]; + } + } + /// /// Batch-verify draft tokens with the target model. Targets that report /// (CPU , dense diff --git a/tests/SharpInference.Tests.ForwardPass/CudaSpecBatchVerifyGemma4Tests.cs b/tests/SharpInference.Tests.ForwardPass/CudaSpecBatchVerifyGemma4Tests.cs new file mode 100644 index 00000000..a7a84477 --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/CudaSpecBatchVerifyGemma4Tests.cs @@ -0,0 +1,361 @@ +using SharpInference.Core; +using SharpInference.Cpu; +using SharpInference.Engine; +using SharpInference.Cuda; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Issue #178: single-user speculative-decode on the +/// Gemma-4 path. now admits Gemma-4 +/// (it routed only the dense gate before), so a GPU draft can batch-verify on the 12B/E4B target. +/// The packed pass dispatches through RunBatchedTrunkGemma4, whose per-sequence attention +/// loop appends each of the k rows' K/V into the SHARED owned cache then attends in ascending row +/// order — the same append-then-attend causality the dense ragged path documents, but exercising +/// per-layer head_dim, SWA rings, the shared-KV tail, k_eq_v, PLE, sandwich norms, and the final +/// softcap. +/// +/// Correctness contract (argmax-stable class, like ): +/// the Gemma-4 batched decode routes matmuls through cuBLAS GEMM (fp16), so BatchVerify is +/// argmax-stable — NOT bit-exact — vs the per-token fp32 ForwardGemma4 loop. Asserted with +/// the maxAbs/top-5 tolerances of the dense . +/// +/// The ring-boundary oracle is the case the dense test (non-SWA Qwen3) never reaches and the one +/// the plan flagged as highest-risk: it prefills PAST the SWA ring (window + headroom) so the +/// verify operates on a wrapped ring where physical slot != logical position — the common case at +/// the ≥32K context this feature targets. Silent-skips when CUDA or the GGUF is absent; mirrors +/// . +/// +public sealed class CudaSpecBatchVerifyGemma4Tests +{ + // E4B Q8_0 exercises the richest Gemma-4 geometry (per-layer head_dim 256, SWA rings, the + // 18-layer shared-KV tail, PLE) and its Q8_0 weights are GEMM-N-batchable, so it reports + // SupportsBatchVerify. Falls back to a 12B Q4_K_M (the issue's headline model) if present. + private static readonly string[] TargetCandidates = + { + "gemma-4-E4B-it-Q8_0.gguf", + "gemma-4-12B-it-qat-Q4_K_M.gguf", + "gemma4-12b-q4km.gguf", + }; + // Optional small same-vocab draft for the e2e oracle (skipped if absent). + private static readonly string[] DraftCandidates = + { + "gemma-3-1b-it-Q8_0.gguf", + "gemma-4-E2B-it-Q8_0.gguf", + }; + + private static CudaBackend? TryCreate() + { + if (!CudaBackend.IsAvailable()) return null; + try { return CudaBackend.Create(); } + catch { return null; } + } + + // SnapKV pinned off (it is structurally off for Gemma-4 anyway): keeps the oracle + // machine-independent, matching the other Gemma-4 CUDA test fixtures. + private static CudaForwardPass NewFwd(GgufModel model, CudaBackend gpu, ModelHyperparams hp, + int ctx, string? kvDtype = null) + { + var prevSnap = Environment.GetEnvironmentVariable("SHARPI_SNAPKV_BUDGET"); + var prevKv = Environment.GetEnvironmentVariable("SHARPI_KV_DTYPE"); + Environment.SetEnvironmentVariable("SHARPI_SNAPKV_BUDGET", "0"); + if (kvDtype is not null) Environment.SetEnvironmentVariable("SHARPI_KV_DTYPE", kvDtype); + try { return new CudaForwardPass(model, gpu, hp, maxContextLength: ctx); } + finally + { + Environment.SetEnvironmentVariable("SHARPI_SNAPKV_BUDGET", prevSnap); + Environment.SetEnvironmentVariable("SHARPI_KV_DTYPE", prevKv); + } + } + + private static string? FindFirst(string[] candidates) + { + foreach (var file in candidates) + { + string[] absolute = { $@"E:\models\{file}", $@"C:\p\sharpi\models\{file}" }; + foreach (var p in absolute) + if (File.Exists(p)) return p; + var dir = Directory.GetCurrentDirectory(); + for (int i = 0; i < 8; i++) + { + var p = Path.Combine(dir, "models", file); + if (File.Exists(p)) return p; + var parent = Directory.GetParent(dir); + if (parent is null) break; + dir = parent.FullName; + } + } + return null; + } + + private static int Argmax(ReadOnlySpan logits) + { + int best = 0; + float bestVal = logits[0]; + for (int i = 1; i < logits.Length; i++) + if (logits[i] > bestVal) { bestVal = logits[i]; best = i; } + return best; + } + + private static HashSet TopKSet(ReadOnlySpan logits, int k) + { + var idx = new int[logits.Length]; + for (int i = 0; i < idx.Length; i++) idx[i] = i; + var arr = logits.ToArray(); + Array.Sort(idx, (a, b) => arr[b].CompareTo(arr[a])); + var set = new HashSet(); + for (int i = 0; i < k && i < idx.Length; i++) set.Add(idx[i]); + return set; + } + + private static (float maxAbs, int overlap) Compare(float[] reference, float[] candidate) + { + Assert.Equal(reference.Length, candidate.Length); + float maxAbs = 0f; + for (int i = 0; i < reference.Length; i++) + maxAbs = MathF.Max(maxAbs, MathF.Abs(reference[i] - candidate[i])); + var refTop = TopKSet(reference, 5); + var candTop = TopKSet(candidate, 5); + int overlap = 0; + foreach (var t in candTop) if (refTop.Contains(t)) overlap++; + return (maxAbs, overlap); + } + + // Argmax parity tolerant of an fp16-GEMM near-tie flip — accepted ONLY when the reference's + // top-2 are within tieEps (mirrors Gemma4CudaBatchForwardMultiTests.AssertArgmaxOrNearTie). + private static void AssertArgmaxOrNearTie(float[] reference, float[] candidate, float tieEps, string label) + { + int rArg = Argmax(reference), cArg = Argmax(candidate); + if (rArg == cArg) return; + float gap = MathF.Abs(reference[rArg] - reference[cArg]); + Assert.True(gap < tieEps, + $"{label}: batched argmax {cArg} != sequential {rArg}, NOT a near-tie (reference gap {gap:F3} ≥ {tieEps:F1}) " + + "— a real wiring divergence (per-layer geometry / SWA ring / shared-KV / PLE), not fp16 rounding."); + } + + // Real Gemma-4 token ids (BOS=2 + natural mid-vocab subwords), matching the activation regime + // the established Gemma-4 batched oracles assert under. Natural tokens for the long ring prefill. + private static readonly int[] GemmaPrompt = { 2, 651, 6037, 576, 6081, 603, 1234, 4567, 8901, 222 }; + private static int[] NaturalTokens(int count) + { + var t = new int[count]; + for (int i = 0; i < count; i++) t[i] = i == 0 ? 2 : 200 + (i * 37) % 8000; + return t; + } + + // The k-row packed verify batches more rows through the fp16 cuBLAS GEMM than the N=2 decode + // the sibling oracles' maxAbs<1.0 bound was calibrated for, so absolute logit divergence is + // modestly larger (~1.1 at k=4/6 on the ±softcap range). Argmax-or-near-tie + top-5 overlap are + // the real correctness contract; maxAbs is a coarse divergence guard at the fp16-GEMM scale. + private static void AssertParity(CudaForwardPass fwd, int[] prompt, int k, string label, float maxAbsTol = 1.5f) + { + fwd.ResetCache(); + var prefillLogits = fwd.Prefill(prompt); + int P = prompt.Length; + + // Greedy-chain k tokens so the verified positions carry realistic activations. + var tokens = new int[k]; + tokens[0] = Argmax(prefillLogits); + var reference = new float[k][]; + for (int i = 0; i < k; i++) + { + var logits = fwd.Forward(tokens[i], P + i); + reference[i] = logits.ToArray(); + if (i + 1 < k) tokens[i + 1] = Argmax(logits); + } + + // Soft rewind (stale K/V must be overwritten) and batch-verify. + fwd.TruncateTo(P); + float[][] batch = fwd.BatchVerify(tokens, P); + + Assert.Equal(k, batch.Length); + for (int i = 0; i < k; i++) + { + var (maxAbs, overlap) = Compare(reference[i], batch[i]); + AssertArgmaxOrNearTie(reference[i], batch[i], tieEps: 0.5f, $"{label} pos {i}"); + Assert.True(overlap >= 4, + $"{label} pos {i}: batched top-5 overlaps sequential in {overlap}/5 (maxAbs={maxAbs})."); + Assert.True(maxAbs < maxAbsTol, + $"{label} pos {i}: batched vs sequential diverged: maxAbs={maxAbs} (tol {maxAbsTol})."); + } + } + + [Theory] + [InlineData(4)] + [InlineData(6)] // not a capacity-stamped WS size — exercises pad-to-capacity dispatch + public void Gemma4_BatchVerify_MatchesSequentialForward(int k) + { + using var gpu = TryCreate(); + if (gpu is null) return; + var path = FindFirst(TargetCandidates); + if (path is null) return; + + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata, model); + Assert.NotNull(hp.LayerHeadDim); // Gemma-4 marker + + using var fwd = NewFwd(model, gpu, hp, ctx: 512); + Assert.True(fwd.SupportsBatchVerify, + "A GEMM-N-batchable Gemma-4 model must report SupportsBatchVerify on the CUDA path (#178)."); + + AssertParity(fwd, GemmaPrompt, k, "Gemma4 verify"); + } + + /// + /// q8_0 KV variant — the issue's headline config (q8 KV frees the VRAM the draft needs). The + /// quantized ring is lossy, so the maxAbs tolerance (not exact equality) carries it; argmax + /// must stay stable, the same contract the q8 KV decode path holds elsewhere. + /// + [Fact] + public void Gemma4_BatchVerify_Q8Kv_MatchesSequentialForward() + { + using var gpu = TryCreate(); + if (gpu is null) return; + var path = FindFirst(TargetCandidates); + if (path is null) return; + + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata, model); + using var fwd = NewFwd(model, gpu, hp, ctx: 512, kvDtype: "q8_0"); + if (!fwd.SupportsBatchVerify) return; // q8 KV geometry unsupported on this build → skip + + // q8 KV is lossy → a slightly looser maxAbs than the fp32-KV path; argmax stays stable. + AssertParity(fwd, GemmaPrompt, k: 4, "Gemma4 q8-KV verify", maxAbsTol: 2.0f); + } + + /// + /// Ring-wrap oracle (the highest-risk Gemma-4 case): prefill PAST the SWA ring + /// (window + headroom, ≈5K) so the verify writes ring slots whose physical index != + /// logical position — the common case at ≥32K context. The batched packed verify must still + /// match k sequential forwards at the wrapped offset. Heavy (multi-thousand-token prefill); + /// model-gated so it only runs locally on a GPU with the GGUF. + /// + [Fact] + public void Gemma4_BatchVerify_AcrossSwaRingBoundary() + { + using var gpu = TryCreate(); + if (gpu is null) return; + var path = FindFirst(TargetCandidates); + if (path is null) return; + + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata, model); + if (hp.SlidingWindowSize <= 0) return; // no SWA layers → nothing to wrap + + // SwaRingSize = min(ctx, window + SwaRingHeadroom>=4096). Pick a ctx comfortably above + // window+4096 and a prefill length past the ring so the SWA cache has wrapped. + int ring = hp.SlidingWindowSize + 4096; + int ctx = ring + 3072; + int prefillLen = ring + 512; + + using var fwd = NewFwd(model, gpu, hp, ctx: ctx); + if (!fwd.SupportsBatchVerify) return; + if (prefillLen + 8 >= fwd.MaxSeqLen) return; // model can't seat the wrap → skip + + // Long synthetic context → looser maxAbs (deeper fp16 accumulation over thousands of + // positions); the point is argmax-stable correctness on the WRAPPED ring. + AssertParity(fwd, NaturalTokens(prefillLen), k: 4, "Gemma4 ring-wrap verify", maxAbsTol: 3.0f); + } + + /// + /// Rollback oracle: verify [t0, junk, junk, junk], accept only t0, TruncateTo(P+1), commit the + /// correction t1. Post-rollback logits must match the sequential trajectory that never saw the + /// rejected tokens — catches stale SWA ring-slot leaks past the truncation point. + /// + [Fact] + public void Gemma4_BatchVerify_TruncateAndCommit_MatchesSequential() + { + using var gpu = TryCreate(); + if (gpu is null) return; + var path = FindFirst(TargetCandidates); + if (path is null) return; + + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata, model); + using var fwd = NewFwd(model, gpu, hp, ctx: 512); + Assert.True(fwd.SupportsBatchVerify); + + var prompt = GemmaPrompt; + fwd.ResetCache(); + var prefillLogits = fwd.Prefill(prompt); + int P = prompt.Length; + int t0 = Argmax(prefillLogits); + + // Sequential reference trajectory: accept t0 → t1, then commit t1. + float[] afterT0 = fwd.Forward(t0, P).ToArray(); + int t1 = Argmax(afterT0); + float[] reference = fwd.Forward(t1, P + 1).ToArray(); + + // Spec-step shape: rewind, verify [t0, junk, junk, junk], accept only t0, commit t1. + fwd.TruncateTo(P); + int junk = (t0 + 7919) % hp.VocabSize; + float[][] batch = fwd.BatchVerify([t0, junk, junk, junk], P); + AssertArgmaxOrNearTie(afterT0, batch[0], tieEps: 0.5f, "after-t0"); // verify[0] still picks t1 + + fwd.TruncateTo(P + 1); + float[] committed = fwd.Forward(t1, P + 1).ToArray(); + + var (maxAbs, overlap) = Compare(reference, committed); + AssertArgmaxOrNearTie(reference, committed, tieEps: 0.5f, "post-rollback commit"); + Assert.True(overlap >= 4, $"Post-rollback top-5 overlap {overlap}/5 (maxAbs={maxAbs})."); + Assert.True(maxAbs < 1.0f, $"Post-rollback diverged: maxAbs={maxAbs}."); + } + + /// + /// E2E greedy parity: SpeculativeDecoder with a CUDA Gemma-4 target + a small same-vocab CPU + /// draft must emit the target's own non-spec greedy continuation — the spec invariant (the + /// draft only proposes; every emitted token is the target's argmax). Gemma-4 BatchVerify is + /// argmax-stable (not bit-exact), so a divergence means a real verify/rollback bug or an + /// FP-borderline argmax flip — investigate before weakening. Skips without the draft GGUF. + /// + [Fact] + public void Gemma4_SpecDecode_GreedyParity_E2E() + { + using var gpu = TryCreate(); + if (gpu is null) return; + var targetPath = FindFirst(TargetCandidates); + var draftPath = FindFirst(DraftCandidates); + if (targetPath is null || draftPath is null) return; + + const int DecodeTokens = 32; + + using var targetModel = GgufModel.Open(targetPath); + var targetHp = ModelHyperparams.FromGgufMetadata(targetModel.Metadata, targetModel); + using var target = NewFwd(targetModel, gpu, targetHp, ctx: 512); + Assert.True(target.SupportsBatchVerify); + + using var draftModel = GgufModel.Open(draftPath); + var draftHp = ModelHyperparams.FromGgufMetadata(draftModel.Metadata, draftModel); + if (targetHp.VocabSize != draftHp.VocabSize) return; // different tokenizer → not a draft + + var prompt = GemmaPrompt; + + // Non-spec greedy baseline on the CUDA target. + target.ResetCache(); + var logits = target.Prefill(prompt); + int P = prompt.Length; + var baseline = new List(); + int tok = Argmax(logits); + for (int i = 0; i < DecodeTokens; i++) + { + baseline.Add(tok); + logits = target.Forward(tok, P + i); + tok = Argmax(logits); + } + + using var cpu = new CpuBackend(); + using var draft = new SharpInference.Engine.ForwardPass(draftModel, cpu, draftHp); + + target.ResetCache(); + var targetLogits = target.Prefill(prompt).ToArray(); + var draftLogits = draft.Prefill(prompt).ToArray(); + + var spec = new SpeculativeDecoder(target, draft, lookahead: 4); + spec.Initialize(P, targetLogits, draftLogits); + + var emitted = new List(); + spec.Decode(DecodeTokens, [], emitted.Add); + + Assert.Equal(baseline, emitted); + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/SamplerTests.cs b/tests/SharpInference.Tests.ForwardPass/SamplerTests.cs index 48974977..2c561c8c 100644 --- a/tests/SharpInference.Tests.ForwardPass/SamplerTests.cs +++ b/tests/SharpInference.Tests.ForwardPass/SamplerTests.cs @@ -328,4 +328,85 @@ public void SampleTopK_TopKThenTopP_NucleusTakenAfterTopK() var reach = ReachableSet(logits, p, seed: 2, trials: 500); Assert.True(reach.IsSubsetOf(new[] { 0 }), $"got {string.Join(",", reach)}"); } + + // ── Distribution helpers for speculative sampling (issue #178) ─────────────────── + + [Fact] + public void BuildFilteredDistribution_NoFilters_MatchesSoftmax() + { + float[] logits = [1f, 2f, 3f, 0f]; + var p = new SamplingParams { Temperature = 1f }; // no top-k/top-p/min-p + var probs = new float[logits.Length]; + Sampler.BuildFilteredDistribution(logits, p, probs); + + // Reference softmax. + double max = 3.0, sum = 0; + var expected = new double[logits.Length]; + for (int i = 0; i < logits.Length; i++) { expected[i] = Math.Exp(logits[i] - max); sum += expected[i]; } + for (int i = 0; i < logits.Length; i++) + Assert.Equal(expected[i] / sum, probs[i], 5); + Assert.Equal(1f, probs.Sum(), 5); + } + + [Fact] + public void BuildFilteredDistribution_TempZero_OneHotAtArgmax() + { + float[] logits = [1f, 9f, 3f, 2f]; + var probs = new float[logits.Length]; + Sampler.BuildFilteredDistribution(logits, new SamplingParams { Temperature = 0f }, probs); + Assert.Equal(1f, probs[1]); + Assert.Equal(0f, probs[0] + probs[2] + probs[3]); + } + + [Fact] + public void BuildFilteredDistribution_TopK_ZerosOutsideTopKAndSumsToOne() + { + float[] logits = [5f, 4f, 3f, 2f, 1f]; + var probs = new float[logits.Length]; + Sampler.BuildFilteredDistribution(logits, new SamplingParams { Temperature = 1f, TopK = 2 }, probs); + Assert.True(probs[0] > 0f && probs[1] > 0f); + Assert.Equal(0f, probs[2] + probs[3] + probs[4]); // outside top-2 zeroed + Assert.Equal(1f, probs.Sum(), 5); + } + + [Fact] + public void SampleWithDistribution_FillsProbsAndReturnsDrawnToken() + { + // Dominant token 1 → with low temperature the draw is token 1 and probs[1] ≈ 1. + float[] logits = [0f, 12f, 0f, 0f]; + var probs = new float[logits.Length]; + int tok = Sampler.SampleWithDistribution(logits, new SamplingParams { Temperature = 0.1f }, probs, new Random(7)); + Assert.Equal(1, tok); + Assert.True(probs[1] > 0.99f); + Assert.Equal(1f, probs.Sum(), 5); + } + + [Fact] + public void ResampleResidual_ConcentratedResidual_AlwaysReturnsResidualToken() + { + // p mass on {1,2}; q takes all of token 1 → residual = {2}. Every draw must be token 2. + float[] p = [0f, 0.5f, 0.5f, 0f]; + float[] q = [0f, 1.0f, 0.0f, 0f]; + var rng = new Random(3); + for (int i = 0; i < 200; i++) + Assert.Equal(2, Sampler.ResampleResidual(p, q, rng)); + } + + [Fact] + public void ResampleResidual_EmptyResidual_FallsBackToP() + { + // q dominates p everywhere on the support → residual is empty → fall back to sampling p. + float[] p = [0.25f, 0.25f, 0.25f, 0.25f]; + float[] q = [0.25f, 0.25f, 0.25f, 0.25f]; + var rng = new Random(5); + int tok = Sampler.ResampleResidual(p, q, rng); + Assert.InRange(tok, 0, p.Length - 1); // a valid token from p, not a sentinel + } + + [Fact] + public void ResampleResidual_MismatchedLengths_Throws() + { + Assert.Throws(() => + Sampler.ResampleResidual(new float[4], new float[3], new Random(1))); + } } diff --git a/tests/SharpInference.Tests.ForwardPass/SpeculativeSamplingTests.cs b/tests/SharpInference.Tests.ForwardPass/SpeculativeSamplingTests.cs new file mode 100644 index 00000000..36e10656 --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/SpeculativeSamplingTests.cs @@ -0,0 +1,145 @@ +using SharpInference.Core; +using SharpInference.Engine; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Sampled (temp > 0) speculative decoding — issue #178. The core invariant is +/// distribution preservation: with the distribution-preserving accept rule (draft sampled with +/// proposal prob q, accepted with min(1, p/q), rejection resampled from the residual max(0, p−q), +/// full accept drawn from the last verify position), the emitted-token distribution is identical +/// to sampling directly from the target — REGARDLESS of the draft's distribution. +/// +/// Tested with a position/token-independent fixed-distribution mock: the target always returns the +/// same distribution p and the draft a deliberately different q, so the aggregate histogram of all +/// emitted tokens must converge to softmax(p). Also covers RNG determinism and that temp ≤ 0 still +/// routes to the byte-stable greedy path. +/// +public sealed class SpeculativeSamplingTests +{ + // Two deliberately different fixed distributions over a small vocab. + private static readonly float[] PTarget = [0.2f, 2.5f, 0.1f, 1.0f, 0.3f, 1.8f, 0.0f, 0.5f]; + private static readonly float[] QDraft = [2.0f, 0.1f, 1.5f, 0.2f, 1.0f, 0.0f, 2.2f, 0.3f]; + + [Fact] + public void SampledSpec_PreservesTargetDistribution() + { + var target = new FixedDistForwardPass(PTarget, supportsBatchVerify: true); + var draft = new FixedDistForwardPass(QDraft, supportsBatchVerify: false); + var spec = new SpeculativeDecoder(target, draft, new SamplingParams { Temperature = 1f }, new Random(12345), lookahead: 4); + spec.Initialize(1, PTarget); + + const int n = 40000; + var counts = new int[PTarget.Length]; + spec.Decode(n, [], t => counts[t]++); + + var expected = Softmax(PTarget); + double maxDev = 0; + for (int i = 0; i < PTarget.Length; i++) + maxDev = Math.Max(maxDev, Math.Abs((double)counts[i] / n - expected[i])); + + // ~8σ headroom at N=40000 (max prob ≈ 0.45 → σ ≈ 0.0025) yet far below the ≥0.1 + // deviation a broken accept rule (emitting ≈q, or a biased p/q mix) would produce. + Assert.True(maxDev < 0.02, $"emitted histogram deviates from softmax(target) by {maxDev:F4} (> 0.02)"); + // The distributions differ, so some drafts are rejected and some accepted. + Assert.InRange(spec.AcceptanceRate, 0.01f, 0.99f); + } + + [Fact] + public void SampledSpec_PMinLooserAccept_StillProducesValidTokens() + { + // --spec-draft-p-min in (0,1) opts into the looser accept; not distribution-preserving, + // but must still emit valid in-vocab tokens and accept more than the strict rule would. + var target = new FixedDistForwardPass(PTarget, supportsBatchVerify: true); + var draft = new FixedDistForwardPass(QDraft, supportsBatchVerify: false); + var spec = new SpeculativeDecoder(target, draft, + new SamplingParams { Temperature = 1f, SpecDraftPMin = 0.05f }, new Random(7), lookahead: 4); + spec.Initialize(1, PTarget); + + var emitted = new List(); + spec.Decode(500, [], emitted.Add); + Assert.Equal(500, emitted.Count); + Assert.All(emitted, t => Assert.InRange(t, 0, PTarget.Length - 1)); + } + + [Fact] + public void SampledSpec_SameSeed_Deterministic() + => Assert.Equal(RunSampled(seed: 999, n: 200), RunSampled(seed: 999, n: 200)); + + [Fact] + public void SampledSpec_DifferentSeed_Differs() + => Assert.NotEqual(RunSampled(seed: 1, n: 200), RunSampled(seed: 2, n: 200)); + + [Fact] + public void SamplingCtor_TempZero_RoutesToGreedy() + { + // Temp ≤ 0 leaves _sampling null → the greedy path runs (byte-stable). With a fixed + // target distribution, greedy emits argmax(p) at every position; the draft never matches. + var target = new FixedDistForwardPass(PTarget, supportsBatchVerify: true); + var draft = new FixedDistForwardPass(QDraft, supportsBatchVerify: false); + var spec = new SpeculativeDecoder(target, draft, new SamplingParams { Temperature = 0f }, new Random(1), lookahead: 4); + spec.Initialize(1, PTarget); + + var emitted = new List(); + spec.Decode(20, [], emitted.Add); + int argmax = Sampler.Greedy(PTarget); + Assert.All(emitted, t => Assert.Equal(argmax, t)); + } + + private static List RunSampled(int seed, int n) + { + var target = new FixedDistForwardPass(PTarget, supportsBatchVerify: true); + var draft = new FixedDistForwardPass(QDraft, supportsBatchVerify: false); + var spec = new SpeculativeDecoder(target, draft, new SamplingParams { Temperature = 1f }, new Random(seed), lookahead: 4); + spec.Initialize(1, PTarget); + var emitted = new List(); + spec.Decode(n, [], emitted.Add); + return emitted; + } + + private static double[] Softmax(float[] logits) + { + double max = logits.Max(); + double sum = 0; + var e = new double[logits.Length]; + for (int i = 0; i < logits.Length; i++) { e[i] = Math.Exp(logits[i] - max); sum += e[i]; } + for (int i = 0; i < logits.Length; i++) e[i] /= sum; + return e; + } + + /// + /// Position/token-independent mock: and always + /// return the same fixed logits, so every verified position has the same distribution and the + /// aggregate emitted histogram is directly comparable to that distribution. + /// + private sealed class FixedDistForwardPass : IForwardPass + { + private readonly float[] _logits; + private readonly bool _bv; + + public FixedDistForwardPass(float[] logits, bool supportsBatchVerify) + { + _logits = logits; + _bv = supportsBatchVerify; + } + + public int VocabSize => _logits.Length; + public int MaxSeqLen => 1 << 20; // mock ignores positions; allow long runs + public bool SupportsPartialRewind => true; + public bool SupportsBatchVerify => _bv; + + public ReadOnlySpan Forward(int token, int position) => _logits; + + public float[][] BatchVerify(int[] tokens, int startPos) + { + var r = new float[tokens.Length][]; + for (int i = 0; i < tokens.Length; i++) r[i] = (float[])_logits.Clone(); + return r; + } + + public ReadOnlySpan Prefill(IReadOnlyList tokens, int startPos = 0) => _logits; + public void TruncateTo(int length) { } + public void ResetCache() { } + public void Dispose() { } + } +}