DiffusionGemmaSampler.cs

// Copyright (c) Zhongkai Fu. All rights reserved.
// https://github.com/zhongkaifu/TensorSharp
//
// This file is part of TensorSharp.
//
// TensorSharp is licensed under the BSD-3-Clause license found in the LICENSE file in the root directory of this source tree.
//
// TensorSharp is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the BSD-3-Clause License for more details.
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Numerics.Tensors;
using System.Threading;
using System.Threading.Tasks;

namespace TensorSharp.Models
{
    /// <summary>Hyper-parameters for the DiffusionGemma EntropyBound denoising sampler.
    /// Defaults mirror the reference (llama.cpp diffusion-gemma) values.</summary>
    public sealed class DiffusionEbParams
    {
        public int MaxDenoisingSteps = 48;
        public float TMin = 0.4f;            // temperature at the last step
        public float TMax = 0.8f;            // temperature at the first step
        public float EntropyBound = 0.1f;    // accept lowest-entropy tokens within this MI bound
        public int StabilityThreshold = 1;   // # of steps the argmax canvas must hold to count stable
        public float ConfidenceThreshold = 0.005f; // stop once mean canvas entropy drops below this
        public int Seed = 0;
        public int MaxBlocks = 1;            // block-autoregressive blocks (each = canvas_length tokens)
    }

    /// <summary>
    /// The DiffusionGemma "EntropyBound" denoising sampler.
    ///
    /// Each block of <c>canvas_length</c> tokens is generated by iterative denoising:
    ///  - The canvas is initialized with random tokens (NOT a mask token).
    ///  - Each step runs a bidirectional forward and, per canvas position, computes the argmax,
    ///    the Shannon entropy of softmax(logits/t), and one multinomial sample.
    ///  - The lowest-entropy positions (cumulative mutual-information &lt;= entropy_bound) are accepted
    ///    (kept = sampled token); the rest are re-noised with fresh random tokens.
    ///  - The emitted/output canvas is always the deterministic argmax canvas.
    ///  - Generation stops early when the argmax canvas is stable for <c>StabilityThreshold</c> steps
    ///    and the mean entropy is below <c>ConfidenceThreshold</c>.
    ///
    /// Multiple blocks are generated autoregressively: a finished block is committed to the prompt
    /// prefix and the next block is denoised, until an end-of-generation token or a repetition loop.
    /// </summary>
    public sealed class DiffusionGemmaSampler
    {
        private readonly DiffusionGemmaModel _model;
        private readonly int _vocab;
        private readonly int _canvasLength;

        // How to decode concurrent requests' canvases each step:
        //  - default (false): run the fast FUSED single-canvas kernel once per live sequence (time-slice).
        //    On a GPU already compute-saturated by one canvas (the case for this 128-expert MoE at
        //    canvas=256), this is the fastest aggregate path — the fused kernel is ~3x faster per canvas
        //    than the per-op true-batched forward, and true batching can't beat a compute-bound GPU.
        //  - DIFFUSION_BATCHED_FORWARD=1: fuse all live canvases into ONE per-op batched forward
        //    (DecodeCanvasBatched). Correct, and a win only when a single canvas leaves the GPU
        //    under-utilised (small canvas / large GPU); a loss when already compute-bound.
        private readonly bool _useBatchedForward = Environment.GetEnvironmentVariable("DIFFUSION_BATCHED_FORWARD") == "1";

        // One-time notice when the on-device sampling kernel rejects a step and the
        // block continues on the host logits path (full [vocab,C] logits cross PCIe
        // plus a host CPU sweep per step — slower).
        private bool _deviceSampleFallbackWarned;

        private void WarnDeviceSampleFallback()
        {
            if (_deviceSampleFallbackWarned) return;
            _deviceSampleFallbackWarned = true;
            Console.WriteLine("  [device-sample] on-device sampling kernel rejected the layout; " +
                "sampling moves to the host logits path for the rest of the block (full logits readback " +
                "per step, slower). Reported once.");
        }

        public DiffusionGemmaSampler(DiffusionGemmaModel model)
        {
            _model = model;
            _vocab = model.VocabSize;
            _canvasLength = model.CanvasLength;
        }

        /// <summary>Generate a full response (block-autoregressive) for the given prompt tokens.
        /// <paramref name="blockStepCallback"/> is invoked after every denoising step with
        /// (blockIndex, step, totalSteps, previewTokens) where previewTokens is the committed
        /// response so far plus the current block's best-guess (argmax) canvas — useful for a live
        /// "denoising" preview in a UI. <paramref name="ct"/> stops generation between steps.</summary>
        public List<int> Generate(int[] promptTokens, DiffusionEbParams p,
            Action<int, int, int, int[]> blockStepCallback = null,
            CancellationToken ct = default)
        {
            var prefix = new List<int>(promptTokens);
            var response = new List<int>();

            int blocks = Math.Max(1, p.MaxBlocks);
            for (int b = 0; b < blocks; b++)
            {
                if (ct.IsCancellationRequested) break;
                int bIdx = b;
                int[] canvas = DenoiseBlock(prefix.ToArray(), p,
                    blockStepCallback == null ? null : (step, total, argmax) =>
                    {
                        var preview = new int[response.Count + argmax.Length];
                        response.CopyTo(preview, 0);
                        argmax.CopyTo(preview, response.Count);
                        blockStepCallback(bIdx, step, total, preview);
                    },
                    ct);

                int cut = TrimCanvas(canvas, canvas.Length);
                for (int i = 0; i < cut; i++) response.Add(canvas[i]);

                if (cut < canvas.Length) break;            // end token or repetition loop -> done
                for (int i = 0; i < cut; i++) prefix.Add(canvas[i]);
            }
            return response;
        }

        /// <summary>Denoise a single block of <c>canvas_length</c> tokens. Returns the argmax canvas.
        /// <paramref name="stepCallback"/> receives (step, totalSteps, argmaxCanvasSnapshot) each step.</summary>
        public int[] DenoiseBlock(int[] promptTokens, DiffusionEbParams p,
            Action<int, int, int[]> stepCallback = null, CancellationToken ct = default)
        {
            int P = promptTokens.Length;
            int C = _canvasLength;
            int S = Math.Max(1, p.MaxDenoisingSteps);
            int vocab = _vocab;

            var rng = new DeterministicRng((ulong)p.Seed);

            int[] currentCanvas = new int[C];
            for (int i = 0; i < C; i++) currentCanvas[i] = rng.NextInt(vocab);

            // Self-conditioning reads the PREVIOUS step's logits. The model reads scBuffer at the START of
            // the forward (before it overwrites/produces the new logits), so we can hand it the previous
            // step's returned logits array directly instead of copying 268 MB into a side buffer each step.
            // Both paths return a reusable buffer (fused: _fusedLogitsBuffer, per-op: _canvasLogits) that
            // is only overwritten at the END of the forward — after SC has read it — so the alias is safe.
            float[] scBuffer = null;
            int[] argmaxCanvas = new int[C];
            int[] prevArgmax = new int[C];
            for (int i = 0; i < C; i++) prevArgmax[i] = -1;

            float[] entropy = new float[C];
            int[] denoiser = new int[C];
            int[] order = new int[C];
            float[] u = new float[C];
            int[] renoise = new int[C];

            // Prompt-KV caching: prefill the prompt's per-layer K/V once, then each step decodes only the
            // canvas (reading the cached prompt K/V). Falls back to the unified [prompt|canvas] forward.
            bool usePkv = _model.SupportsPromptKvCache;
            int[] tokens = null;
            if (usePkv)
            {
                _model.PrefillPrompt(promptTokens);
            }
            else
            {
                tokens = new int[P + C];
                Array.Copy(promptTokens, tokens, P);
            }

            // On-device sampling (CUDA): the lm_head kernel returns argmax/entropy/sampled/top-K directly,
            // so the full [vocab,C] logits never cross PCIe and no host full-vocab CPU sweep runs. Needs the
            // prompt-KV decode path (its canvas hidden stays on-device for the sampling tail). The per-step
            // top-K (double-buffered) feeds the NEXT step's self-conditioning, replacing the host logits.
            bool useDeviceSample = usePkv && _model.SupportsDeviceSampling;
            int K = _model.SelfCondTopK;
            int[] dArg = null, dSamp = null;
            int[][] topTok = null; float[][] topPrb = null; int dbuf = 0;
            if (useDeviceSample)
            {
                dArg = new int[C]; dSamp = new int[C];
                int kk = Math.Max(1, K);
                topTok = new[] { new int[C * kk], new int[C * kk] };
                topPrb = new[] { new float[C * kk], new float[C * kk] };
            }

            float prevTempInv = 1f;
            int held = 0;
            bool finished = false;

            for (int curStep = S; curStep >= 1 && !finished && !ct.IsCancellationRequested; curStep--)
            {
                int stepIdx = S - curStep;
                float t = p.TMin + (p.TMax - p.TMin) * ((float)curStep / S);
                float tempInv = 1f / t;
                float scUse = stepIdx == 0 ? 0f : 1f;

                if (useDeviceSample)
                {
                    // pre-draw step randomness (u then renoise per position) so the rng order matches the
                    // host sampler; u drives the on-device multinomial, renoise the rejected positions.
                    for (int pos = 0; pos < C; pos++) { u[pos] = rng.NextFloat(); renoise[pos] = rng.NextInt(vocab); }
                    int[] scTok = stepIdx == 0 ? null : topTok[1 - dbuf];
                    float[] scPrb = stepIdx == 0 ? null : topPrb[1 - dbuf];
                    bool ok = _model.DecodeCanvasSampled(currentCanvas, scTok, scPrb, scUse, tempInv, u, K,
                        dArg, entropy, dSamp, topTok[dbuf], topPrb[dbuf]);
                    if (ok)
                    {
                        finished = DenoiseStepFromDevice(dArg, entropy, dSamp, renoise, p,
                            currentCanvas, argmaxCanvas, prevArgmax, ref held, order);
                        dbuf ^= 1;
                        prevTempInv = tempInv;
                        stepCallback?.Invoke(stepIdx, S, (int[])argmaxCanvas.Clone());
                        continue;
                    }
                    // device path rejected the layout: fall back to the host logits path for the rest of the
                    // block (self-conditioning restarts from the next step, scBuffer null this step).
                    useDeviceSample = false;
                    WarnDeviceSampleFallback();
                }

                float[] logits;
                if (usePkv)
                {
                    logits = _model.DecodeCanvas(currentCanvas, scBuffer, scUse, prevTempInv);
                }
                else
                {
                    for (int i = 0; i < C; i++) tokens[P + i] = currentCanvas[i];
                    logits = _model.ForwardCanvas(tokens, P, scBuffer, scUse, prevTempInv);
                }

                // per-position sampling + accept + renoise + adaptive-stop (shared with the batched path)
                finished = DenoiseStep(logits, tempInv, rng, p, currentCanvas, argmaxCanvas, prevArgmax,
                    ref held, entropy, denoiser, order, u, renoise);
                prevTempInv = tempInv;

                // hand this step's logits to the next step's self-conditioning (no copy; see scBuffer note)
                if (_model.SelfConditioningEnabled) scBuffer = logits;

                stepCallback?.Invoke(stepIdx, S, (int[])argmaxCanvas.Clone());
            }

            return (int[])argmaxCanvas.Clone();
        }

        /// <summary>One denoising step's host-side work, shared by the single-request <see cref="DenoiseBlock"/>
        /// and the batched <see cref="RunBlockBatched"/> paths so they are numerically identical: per canvas
        /// position compute argmax + Shannon entropy of softmax(logits*tempInv) + one multinomial sample,
        /// accept the lowest-entropy positions within the MI bound (kept = sampled token, rest = fresh random),
        /// and report the adaptive-stop verdict (argmax stable for StabilityThreshold steps AND mean entropy
        /// below ConfidenceThreshold). Updates <paramref name="currentCanvas"/> (next step's input),
        /// <paramref name="argmaxCanvas"/> (the emitted best-guess), <paramref name="prevArgmax"/> and
        /// <paramref name="held"/>. The scratch arrays are caller-owned to keep the hot loop allocation-light.</summary>
        private bool DenoiseStep(float[] logits, float tempInv, DeterministicRng rng, DiffusionEbParams p,
            int[] currentCanvas, int[] argmaxCanvas, int[] prevArgmax, ref int held,
            float[] entropy, int[] denoiser, int[] order, float[] u, int[] renoise)
        {
            int C = _canvasLength;
            int vocab = _vocab;

            // pre-draw step randomness (single-threaded) for reproducibility
            for (int pos = 0; pos < C; pos++)
            {
                u[pos] = rng.NextFloat();
                renoise[pos] = rng.NextInt(vocab);
            }

            // per position: argmax + entropy of softmax(logits*tempInv) + one multinomial sample.
            // SIMD-vectorized via TensorPrimitives (exp/sum/dot), with per-worker scratch rows:
            //   s = logits*tempInv;  e = exp(s - max(s));  Z = Σe
            //   H = -Σ p ln p  with  p = e/Z  and  ln p = (s - m) - ln Z
            //     = ln Z + m - (Σ e·s)/Z
            //   sample = first v with cumulative Σe >= u*Z (the same CDF-inverse draw as the scalar
            //   reference, over the same token ordering).
            Parallel.For(0, C,
                () => (scaled: new float[vocab], exps: new float[vocab]),
                (pos, _, scratch) =>
                {
                    var row = logits.AsSpan(pos * vocab, vocab);
                    var s = scratch.scaled.AsSpan();
                    var e = scratch.exps.AsSpan();
                    TensorPrimitives.Multiply(row, tempInv, s);
                    int amax = TensorPrimitives.IndexOfMax<float>(s);
                    float m = s[amax];
                    TensorPrimitives.Subtract(s, m, e);
                    TensorPrimitives.Exp(e, e);
                    float Z = TensorPrimitives.Sum<float>(e);
                    float sz = TensorPrimitives.Dot<float>(e, s);
                    entropy[pos] = MathF.Log(Z) + m - sz / Z;

                    float target = u[pos] * Z;
                    float cum = 0f;
                    int sampled = vocab - 1;
                    for (int v = 0; v < vocab; v++)
                    {
                        cum += e[v];
                        if (cum >= target) { sampled = v; break; }
                    }
                    argmaxCanvas[pos] = amax;
                    denoiser[pos] = sampled;
                    return scratch;
                },
                _ => { });

            // accept lowest-entropy positions within the MI bound (sum of strictly-earlier entropies <= bound)
            for (int i = 0; i < C; i++) order[i] = i;
            Array.Sort(order, (a, bb) => entropy[a].CompareTo(entropy[bb]));
            var accepted = new bool[C];
            double cumE = 0.0;
            for (int kk = 0; kk < C; kk++)
            {
                int pos = order[kk];
                cumE += entropy[pos];
                if (cumE - entropy[pos] <= p.EntropyBound) accepted[pos] = true;
            }

            // renoise: accepted -> sampled token, rest -> fresh random; output = argmax canvas
            float entropySum = 0f;
            for (int pos = 0; pos < C; pos++)
            {
                currentCanvas[pos] = accepted[pos] ? denoiser[pos] : renoise[pos];
                entropySum += entropy[pos];
            }

            // adaptive stop: argmax stable for StabilityThreshold steps AND confident (low mean entropy)
            bool same = true;
            for (int i = 0; i < C; i++) if (prevArgmax[i] != argmaxCanvas[i]) { same = false; break; }
            held = same ? held + 1 : 0;
            bool confident = (entropySum / C) < p.ConfidenceThreshold;
            Array.Copy(argmaxCanvas, prevArgmax, C);
            return held >= p.StabilityThreshold && confident;
        }

        /// <summary>The host half of a denoising step when the per-position argmax/entropy/multinomial sample
        /// were computed ON-DEVICE (the GPU sampler). Identical accept/renoise/adaptive-stop logic to the back
        /// half of <see cref="DenoiseStep"/> — only the full-vocab per-position work is gone (done in the
        /// kernel). <paramref name="argmaxIn"/>/<paramref name="entropyIn"/>/<paramref name="sampledIn"/> are
        /// the device outputs; <paramref name="renoise"/> the host-drawn fresh tokens for rejected positions.</summary>
        private bool DenoiseStepFromDevice(int[] argmaxIn, float[] entropyIn, int[] sampledIn, int[] renoise,
            DiffusionEbParams p, int[] currentCanvas, int[] argmaxCanvas, int[] prevArgmax, ref int held, int[] order)
        {
            int C = _canvasLength;
            Array.Copy(argmaxIn, argmaxCanvas, C);   // emitted best-guess = device argmax

            // accept lowest-entropy positions within the MI bound (sum of strictly-earlier entropies <= bound)
            for (int i = 0; i < C; i++) order[i] = i;
            Array.Sort(order, (a, bb) => entropyIn[a].CompareTo(entropyIn[bb]));
            var accepted = new bool[C];
            double cumE = 0.0;
            for (int kk = 0; kk < C; kk++)
            {
                int pos = order[kk];
                cumE += entropyIn[pos];
                if (cumE - entropyIn[pos] <= p.EntropyBound) accepted[pos] = true;
            }

            // renoise: accepted -> device multinomial sample, rest -> fresh random; output = argmax canvas
            float entropySum = 0f;
            for (int pos = 0; pos < C; pos++)
            {
                currentCanvas[pos] = accepted[pos] ? sampledIn[pos] : renoise[pos];
                entropySum += entropyIn[pos];
            }

            bool same = true;
            for (int i = 0; i < C; i++) if (prevArgmax[i] != argmaxCanvas[i]) { same = false; break; }
            held = same ? held + 1 : 0;
            bool confident = (entropySum / C) < p.ConfidenceThreshold;
            Array.Copy(argmaxCanvas, prevArgmax, C);
            return held >= p.StabilityThreshold && confident;
        }

        // ===================================================================================
        //  Batched (multi-request) generation — the parallel-request throughput path.
        //  The server scheduler keeps a set of in-flight requests as DiffusionSeqRun objects and drives
        //  them one BLOCK at a time via RunBlockBatched, which prefills every active prompt then denoises
        //  all canvases together (one batched model forward per step). Sequences converge / end / hit their
        //  block budget independently; the scheduler admits new requests and retires finished ones between
        //  blocks (block-granular continuous batching).
        // ===================================================================================

        /// <summary>Run ONE block of denoising over every active sequence in lockstep. Prefills each
        /// sequence's current prefix, denoises all canvases together (batched forward per step), and commits
        /// each sequence's trimmed block tokens via <see cref="DiffusionSeqRun.CommitBlock"/>. A sequence that
        /// converges early (or whose request is cancelled) freezes its canvas and is dropped from the batch
        /// for the remaining steps so it doesn't waste GPU work on the slower sequences.</summary>
        public void RunBlockBatched(IReadOnlyList<DiffusionSeqRun> active, CancellationToken stopToken = default)
        {
            int A = active.Count;
            if (A == 0) return;
            int C = _canvasLength;
            int vocab = _vocab;

            var seqs = new DiffusionSeqState[A];
            var canvas = new int[A][];
            var argmaxCanvas = new int[A][];
            var prevArgmax = new int[A][];
            var scBuffer = new float[A][];
            var rng = new DeterministicRng[A];
            var held = new int[A];
            var finished = new bool[A];
            var prevTempInv = new float[A];
            int Smax = 1;

            // Prompt-KV caching (device-glue backends): prefill each sequence's prefix K/V once, then each
            // step decodes only its canvas. CPU backends have no prompt-KV store, so each step runs the
            // unified [prefix|canvas] forward per sequence instead — the same fallback DenoiseBlock uses.
            bool usePkv = _model.SupportsPromptKvCache;
            var unifiedTokens = usePkv ? null : new int[A][];   // per-seq [prefix|canvas] buffer
            var promptLen = usePkv ? null : new int[A];

            // On-device sampling (CUDA, model fits): per-sequence the lm_head kernel returns argmax/entropy/
            // sampled/top-K directly, so no full [vocab,C] logits cross PCIe. Each sequence keeps its own
            // double-buffered top-K (for self-conditioning), argmax/sampled, mirroring DenoiseBlock.
            bool useDeviceSample = usePkv && _model.SupportsDeviceSampling;
            int K = _model.SelfCondTopK;
            int[][][] dTopTok = useDeviceSample ? new int[A][][] : null;
            float[][][] dTopPrb = useDeviceSample ? new float[A][][] : null;
            int[] dBuf = useDeviceSample ? new int[A] : null;
            int[][] dArg = useDeviceSample ? new int[A][] : null;
            int[][] dSamp = useDeviceSample ? new int[A][] : null;

            // per-sequence prefill + block init
            for (int a = 0; a < A; a++)
            {
                var run = active[a];
                seqs[a] = run.State;
                if (usePkv)
                {
                    _model.PrefillSeq(run.State, run.Prefix.ToArray());
                }
                else
                {
                    int P = run.Prefix.Count;
                    promptLen[a] = P;
                    unifiedTokens[a] = new int[P + C];
                    run.Prefix.CopyTo(unifiedTokens[a], 0);
                }

                int S = Math.Max(1, run.Params.MaxDenoisingSteps);
                if (S > Smax) Smax = S;
                rng[a] = new DeterministicRng((ulong)run.Params.Seed);
                canvas[a] = new int[C];
                for (int i = 0; i < C; i++) canvas[a][i] = rng[a].NextInt(vocab);
                argmaxCanvas[a] = new int[C];
                prevArgmax[a] = new int[C];
                for (int i = 0; i < C; i++) prevArgmax[a][i] = -1;
                // device sampling carries SC via the top-K buffers, so the 268 MB host logits buffer is unused
                scBuffer[a] = (_model.SelfConditioningEnabled && !useDeviceSample) ? new float[(long)C * vocab] : null;
                prevTempInv[a] = 1f;
                if (useDeviceSample)
                {
                    int kk = Math.Max(1, K);
                    dTopTok[a] = new[] { new int[C * kk], new int[C * kk] };
                    dTopPrb[a] = new[] { new float[C * kk], new float[C * kk] };
                    dArg[a] = new int[C];
                    dSamp[a] = new int[C];
                }
            }

            // caller-owned per-step scratch (one set; the per-position work is over a single canvas at a time)
            float[] entropy = new float[C];
            int[] denoiser = new int[C];
            int[] order = new int[C];
            float[] u = new float[C];
            int[] renoise = new int[C];

            // indices of sequences still denoising this block
            var live = new List<int>(A);
            for (int a = 0; a < A; a++) live.Add(a);

            for (int stepIdx = 0; stepIdx < Smax && live.Count > 0 && !stopToken.IsCancellationRequested; stepIdx++)
            {
                // drop cancelled requests from the live set before computing
                for (int j = live.Count - 1; j >= 0; j--)
                {
                    int a = live[j];
                    if (active[a].Ct.IsCancellationRequested) { finished[a] = true; live.RemoveAt(j); }
                }
                int L = live.Count;
                if (L == 0) break;

                if (_useBatchedForward && usePkv && L > 1)
                {
                    // EXPERIMENTAL true-batched forward: all live canvases through one per-op forward. A win
                    // only when a single canvas under-utilises the GPU; a loss when compute-bound (this model).
                    var subSeqs = new DiffusionSeqState[L];
                    var subCanvas = new int[L][];
                    var subScPrev = new float[L][];
                    var subScUse = new float[L];
                    var subPrevTempInv = new float[L];
                    var subTempInv = new float[L];
                    for (int j = 0; j < L; j++)
                    {
                        int a = live[j];
                        var run = active[a];
                        int S = Math.Max(1, run.Params.MaxDenoisingSteps);
                        int curStep = S - stepIdx;
                        subTempInv[j] = 1f / (run.Params.TMin + (run.Params.TMax - run.Params.TMin) * ((float)curStep / S));
                        subScUse[j] = stepIdx == 0 ? 0f : 1f;
                        subScPrev[j] = scBuffer[a];
                        subSeqs[j] = seqs[a];
                        subCanvas[j] = canvas[a];
                        subPrevTempInv[j] = prevTempInv[a];
                    }
                    float[][] logits = _model.DecodeCanvasBatched(subSeqs, subCanvas, subScPrev, subScUse, subPrevTempInv);
                    for (int j = 0; j < L; j++)
                    {
                        int a = live[j];
                        var run = active[a];
                        if (scBuffer[a] != null) Array.Copy(logits[j], scBuffer[a], (long)C * vocab);
                        bool seqFinished = DenoiseStep(logits[j], subTempInv[j], rng[a], run.Params,
                            canvas[a], argmaxCanvas[a], prevArgmax[a], ref held[a], entropy, denoiser, order, u, renoise);
                        prevTempInv[a] = subTempInv[j];
                        run.EmitPreview(stepIdx, Math.Max(1, run.Params.MaxDenoisingSteps), argmaxCanvas[a]);
                        if (seqFinished) finished[a] = true;
                    }
                }
                else
                {
                    // DEFAULT: decode each live sequence with the fast FUSED single-canvas kernel, then sample
                    // immediately (the kernel returns the model's shared readback buffer, so its logits must be
                    // consumed before the next sequence's decode overwrites it). Time-slices the GPU across the
                    // active requests at full fused speed. Without prompt-KV (CPU backends) each sequence runs
                    // the unified [prefix|canvas] forward instead, which has the same shared-buffer contract.
                    for (int j = 0; j < L; j++)
                    {
                        int a = live[j];
                        var run = active[a];
                        int S = Math.Max(1, run.Params.MaxDenoisingSteps);
                        int curStep = S - stepIdx;                  // mirrors single-seq: curStep counts S..1
                        float tempInv = 1f / (run.Params.TMin + (run.Params.TMax - run.Params.TMin) * ((float)curStep / S));
                        float scUse = stepIdx == 0 ? 0f : 1f;

                        if (useDeviceSample)
                        {
                            for (int pos = 0; pos < C; pos++) { u[pos] = rng[a].NextFloat(); renoise[pos] = rng[a].NextInt(vocab); }
                            int[] scTok = stepIdx == 0 ? null : dTopTok[a][1 - dBuf[a]];
                            float[] scPrb = stepIdx == 0 ? null : dTopPrb[a][1 - dBuf[a]];
                            bool ok = _model.DecodeCanvasSampledSeq(seqs[a], canvas[a], scTok, scPrb, scUse, tempInv, u, K,
                                dArg[a], entropy, dSamp[a], dTopTok[a][dBuf[a]], dTopPrb[a][dBuf[a]]);
                            if (ok)
                            {
                                bool seqDone = DenoiseStepFromDevice(dArg[a], entropy, dSamp[a], renoise, run.Params,
                                    canvas[a], argmaxCanvas[a], prevArgmax[a], ref held[a], order);
                                dBuf[a] ^= 1;
                                prevTempInv[a] = tempInv;
                                run.EmitPreview(stepIdx, S, argmaxCanvas[a]);
                                if (seqDone) finished[a] = true;
                                continue;
                            }
                            useDeviceSample = false;   // kernel rejected: fall back to the host path for the rest
                            WarnDeviceSampleFallback();
                        }

                        float[] lg;
                        if (usePkv)
                        {
                            lg = _model.DecodeCanvasSeq(seqs[a], canvas[a], scBuffer[a], scUse, prevTempInv[a]);
                        }
                        else
                        {
                            Array.Copy(canvas[a], 0, unifiedTokens[a], promptLen[a], C);
                            lg = _model.ForwardCanvas(unifiedTokens[a], promptLen[a], scBuffer[a], scUse, prevTempInv[a]);
                        }
                        if (scBuffer[a] != null) Array.Copy(lg, scBuffer[a], (long)C * vocab);
                        bool seqFinished = DenoiseStep(lg, tempInv, rng[a], run.Params,
                            canvas[a], argmaxCanvas[a], prevArgmax[a], ref held[a], entropy, denoiser, order, u, renoise);
                        prevTempInv[a] = tempInv;
                        run.EmitPreview(stepIdx, S, argmaxCanvas[a]);
                        if (seqFinished) finished[a] = true;
                    }
                }

                // retire converged sequences from the batch for the remaining steps
                live.RemoveAll(a => finished[a]);
            }

            // commit each sequence's trimmed block
            for (int a = 0; a < A; a++)
            {
                int cut = TrimCanvas(argmaxCanvas[a], C);
                var blockTokens = new List<int>(cut);
                for (int i = 0; i < cut; i++) blockTokens.Add(argmaxCanvas[a][i]);
                active[a].CommitBlock(blockTokens, ended: cut < C);
            }
        }

        /// <summary>Trim a denoised canvas at the first end-of-generation token, or at the onset of a
        /// repetition loop (a token recurring at stride 1-2 for >= 6 steps).</summary>
        private int TrimCanvas(int[] canvas, int n)
        {
            int cut = n;
            for (int i = 0; i < n; i++)
            {
                if (_model.Tokenizer.IsEos(canvas[i])) { cut = i; break; }
            }
            for (int i = 0; i + 1 < cut; i++)
            {
                bool loop = false;
                for (int stride = 1; stride <= 2 && !loop; stride++)
                {
                    int reps = 0;
                    for (int j = i; j + stride < n && canvas[j] == canvas[j + stride]; j += stride) reps++;
                    if (reps >= 6) loop = true;
                }
                if (loop) { cut = i; break; }
            }
            return cut;
        }

        /// <summary>Small deterministic PRNG (splitmix64 + xorshift) so generation is reproducible
        /// for a given seed independent of the BCL's Random implementation.</summary>
        private sealed class DeterministicRng
        {
            private ulong _state;
            public DeterministicRng(ulong seed)
            {
                // splitmix64 seeding
                _state = seed == 0 ? 0x9E3779B97F4A7C15UL : seed;
            }

            private ulong NextU64()
            {
                ulong z = (_state += 0x9E3779B97F4A7C15UL);
                z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9UL;
                z = (z ^ (z >> 27)) * 0x94D049BB133111EBUL;
                return z ^ (z >> 31);
            }

            public float NextFloat()
            {
                // 24-bit mantissa in [0,1)
                return (NextU64() >> 40) * (1.0f / 16777216.0f);
            }

            public int NextInt(int bound)
            {
                return (int)(NextU64() % (ulong)bound);
            }
        }
    }

    /// <summary>One in-flight request in a batched diffusion generation. Carries the request's params, its
    /// per-sequence model K/V state, the growing prefix (prompt + committed blocks) fed to each block's
    /// prefill, the committed response tokens, and a preview callback. The scheduler drives a set of these
    /// through <see cref="DiffusionGemmaSampler.RunBlockBatched"/> one block at a time; <see cref="Done"/>
    /// flips true when the request ends (end token / repetition) or exhausts its block budget.</summary>
    public sealed class DiffusionSeqRun
    {
        public DiffusionEbParams Params { get; }
        public DiffusionSeqState State { get; }
        public List<int> Prefix { get; }
        public List<int> Response { get; }
        public int BlockIndex { get; private set; }
        public bool Done { get; private set; }
        public CancellationToken Ct { get; }

        private readonly Action<DiffusionSeqRun, int, int, int[]> _onPreview;   // (run, step, totalSteps, previewTokens)

        public DiffusionSeqRun(int[] promptTokens, DiffusionEbParams p, DiffusionSeqState state,
            CancellationToken ct, Action<DiffusionSeqRun, int, int, int[]> onPreview)
        {
            Params = p;
            State = state;
            Ct = ct;
            _onPreview = onPreview;
            Prefix = new List<int>(promptTokens);
            Response = new List<int>();
        }

        /// <summary>Report the current best-guess (committed response + this block's argmax canvas) to the
        /// preview consumer. Whole-message "replace" semantics, matching the single-request preview.</summary>
        internal void EmitPreview(int step, int totalSteps, int[] argmaxCanvas)
        {
            if (_onPreview == null) return;
            var preview = new int[Response.Count + argmaxCanvas.Length];
            Response.CopyTo(preview, 0);
            argmaxCanvas.CopyTo(preview, Response.Count);
            _onPreview(this, step, totalSteps, preview);
        }

        /// <summary>Commit this block's trimmed tokens. If the block ended (end token / repetition) or this
        /// was the last permitted block, the request is done; otherwise the tokens extend the prefix and the
        /// next block is generated.</summary>
        internal void CommitBlock(List<int> blockTokens, bool ended)
        {
            Response.AddRange(blockTokens);
            if (ended || BlockIndex + 1 >= Math.Max(1, Params.MaxBlocks))
                Done = true;
            else
            {
                Prefix.AddRange(blockTokens);
                BlockIndex++;
            }
        }
    }
}