Skip to content

feat(spec): GPU draft sampled speculative decoding for Gemma 4 (#178) #295

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pekkah merged 2 commits into from
Jun 18, 2026
Merged

feat(spec): GPU draft sampled speculative decoding for Gemma 4 (#178) #295

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2428240
feat(spec): GPU draft sampled speculative decoding for Gemma 4 (#178)
pekkah Jun 18, 2026
4b4cb4d
perf(spec): address Gemini review — drop redundant distribution passes
pekkah Jun 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 35 additions & 14 deletions src/SharpInference.Cli/RunCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand All @@ -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)
{
Expand Down Expand Up @@ -831,17 +847,17 @@ 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
{
using var draftCpuBackend = new CpuBackend();
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)
Expand Down Expand Up @@ -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);
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
34 changes: 24 additions & 10 deletions src/SharpInference.Engine/CudaForwardPass.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4599,16 +4599,22 @@ internal float[][] BatchForwardMulti(int[] tokens, int[] positions, CudaSequence
// ── Speculative-decode batched verify (issue #207) ──────────────────────────────────

/// <summary>
/// Whether <see cref="BatchVerify"/> can run: the dense batched-decode configuration
/// (<see cref="DenseBatchedDecodeSupported"/> — non-MoE, non-Gemma-4, no TurboQuant, no
/// final-logit softcap, GEMM-N-batchable weights) with an uncompacted cache. Unlike
/// <see cref="SupportsContinuousBatching"/>, a CONFIGURED SnapKV budget does not disable
/// verify — only an actual prefill-time eviction does (then physical slot != logical
/// Whether <see cref="BatchVerify"/> can run: a batched-decode-capable configuration with an
/// uncompacted cache — either the dense path (<see cref="DenseBatchedDecodeSupported"/> —
/// non-MoE, non-Gemma-4, no final-logit softcap) OR the Gemma-4 path
/// (<see cref="Gemma4BatchedDecodeSupported"/> — 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 <see cref="SupportsContinuousBatching"/>, 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 <c>_kvEvictedCount</c> is
/// always 0 there and the guard is a no-op.)
/// </summary>
public bool SupportsBatchVerify => _kvEvictedCount == 0 && DenseBatchedDecodeSupported();
public bool SupportsBatchVerify =>
_kvEvictedCount == 0 && (DenseBatchedDecodeSupported() || Gemma4BatchedDecodeSupported());

/// <summary>
/// Batched k-token verify for single-user speculative decoding (issue #207): one packed
Expand All @@ -4629,15 +4635,23 @@ internal float[][] BatchForwardMulti(int[] tokens, int[] positions, CudaSequence
/// the cache; the caller rewinds rejected tokens via <see cref="TruncateTo"/>. Issues
/// direct launches only — the per-token decode CUDA graph (owned-cache pointers) stays
/// valid for the surrounding Forward steps.
/// <para>Gemma-4 (issue #178): <see cref="RunBatchedTrunk"/> dispatches the same packed pass
/// to <see cref="RunBatchedTrunkGemma4"/>, 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 <see cref="Forward"/> the k==1 shortcut
/// uses.</para>
/// </summary>
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<float[]>();
if (startPos < 0 || startPos + k > _maxSeqLen)
Expand Down
139 changes: 139 additions & 0 deletions src/SharpInference.Engine/Sampler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,103 @@ public static int Sample(ReadOnlySpan<float> logits, SamplingParams p, Random? r
return SampleFromDistribution(probs, rng);
}

/// <summary>
/// Build the full-vocabulary, post-filter sampling distribution for <paramref name="logits"/>
/// under <paramref name="p"/> into <paramref name="probs"/> (length must equal
/// <c>logits.Length</c>; entries outside the kept support are 0 and the kept entries sum to 1).
/// Applies the SAME pipeline as <see cref="Sample"/>'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 <see cref="Sample"/>.
///
/// 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 <paramref name="p"/> for the
/// min(1, p/q) accept ratio to be meaningful. On Temperature ≤ 0 writes a one-hot at the
/// greedy argmax.
/// </summary>
public static void BuildFilteredDistribution(ReadOnlySpan<float> logits, SamplingParams p, Span<float> 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);
}

/// <summary>
/// Sample like <see cref="Sample"/> but also materialise the full post-filter distribution
/// into <paramref name="probs"/> (see <see cref="BuildFilteredDistribution"/>) and return the
/// drawn token, so <c>probs[token]</c> 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.
/// </summary>
public static int SampleWithDistribution(ReadOnlySpan<float> logits, SamplingParams p, Span<float> 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);
}
Comment on lines +172 to +185

@gemini-code-assist gemini-code-assist Bot Jun 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In SampleWithDistribution, when p.Temperature <= 0f, we can avoid calling BuildFilteredDistribution entirely. This avoids a redundant call to Greedy(logits) and the overhead of copying/clearing arrays. We also introduce a public SampleFromProbs helper to allow sampling from a pre-computed distribution directly.

    public static int SampleWithDistribution(ReadOnlySpan<float> logits, SamplingParams p, Span<float> probs, Random? rng = null)
    {
        if (p.Temperature <= 0f)
        {
            int argmax = Greedy(logits);
            probs.Clear();
            probs[argmax] = 1f;
            return argmax;
        }
        BuildFilteredDistribution(logits, p, probs);
        rng ??= Random.Shared;
        return SampleFromDistribution(probs, rng);
    }

    /// <summary>
    /// Sample a token from a pre-computed probability distribution.
    /// </summary>
    public static int SampleFromProbs(ReadOnlySpan<float> probs, Random rng)
    {
        return SampleFromDistribution(probs, rng);
    }


/// <summary>
/// Draw a token from a pre-computed normalized distribution (e.g. one already built by
/// <see cref="BuildFilteredDistribution"/>), 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.
/// </summary>
public static int SampleFromProbs(ReadOnlySpan<float> probs, Random rng)
{
ArgumentNullException.ThrowIfNull(rng);
return SampleFromDistribution(probs, rng);
}

/// <summary>
/// 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
Expand Down Expand Up @@ -367,6 +464,48 @@ private static int SampleFromDistribution(ReadOnlySpan<float> probs, Random rng)
return probs.Length - 1;
}

/// <summary>
/// Sample a token proportional to the residual max(0, p[x] − q[x]) of two full-vocabulary
/// distributions built with the SAME <see cref="SamplingParams"/> (see
/// <see cref="BuildFilteredDistribution"/>). 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 <paramref name="p"/>. If the residual mass is ≈0 (numerical /
/// q already dominated p on the kept support), falls back to sampling from <paramref name="p"/>.
/// </summary>
public static int ResampleResidual(ReadOnlySpan<float> p, ReadOnlySpan<float> 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);
}
Comment on lines +492 to +507

@gemini-code-assist gemini-code-assist Bot Jun 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In ResampleResidual, we can optimize the loop to avoid the third pass over the arrays. By tracking the last index with a positive residual during the second pass, we can return it directly if the cumulative sum check is not met due to floating-point inaccuracies. This reduces the worst-case passes from 3 to 2.

        float target = (float)rng.NextDouble() * sum;
        float cum = 0f;
        int lastPositiveIdx = -1;
        for (int i = 0; i < p.Length; i++)
        {
            float r = p[i] - q[i];
            if (r > 0f)
            {
                lastPositiveIdx = i;
                cum += r;
                if (target <= cum) return i;
            }
        }
        if (lastPositiveIdx >= 0) return lastPositiveIdx;
        return SampleFromDistribution(p, rng);


/// <summary>
/// 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.
Expand Down
Loading