perf(cuda): dtype-aware auto-context for Gemma 4 (#220)

src/SharpInference.Engine/CudaForwardPass.cs

@@ -566,7 +566,14 @@ public CudaForwardPass(GgufModel model, CudaBackend gpu, ModelHyperparams hp,
         else if (_tqEnabled)
             _maxSeqLen = EstimateMaxContextTq(model, gpu, hp, tqFp32Window, tqBits);
         else
-            _maxSeqLen = EstimateMaxContext(model, gpu, hp);
+            // Size the auto-context for the KV dtype the operator requested (#220): a bf16/q8_0
+            // KV store fits 2×/4× the positions of fp32, but EstimateMaxContext previously priced
+            // fp32 unconditionally, so --kv-type was silently ignored for auto-context. Pass the
+            // *requested* dtype (pre-auto-narrow): an explicit narrowed choice should expand the
+            // window, and the fp32 default still yields the fp32-fit context (the auto-narrow
+            // below never fires for an fp32-sized auto-context, since fp32 fits there by
+            // construction).
+            _maxSeqLen = EstimateMaxContext(model, gpu, hp, ResolveConfiguredKvDType());
         // The KV-append/attention kernels index the cache at `pos % _maxSeqLen` (the ring
         // modulo, identity for full caches), so a zero context — e.g. a malformed GGUF with
         // context_length=0 reached via an explicit ctx-size — would be an in-kernel
@@ -4068,67 +4075,63 @@ internal static DType ResolveKvDType(
 
     /// <summary>
     /// VRAM-based context-length estimator: take the KV-cache budget from
-    /// <see cref="EstimateAvailableKvVram"/> and divide what's left between K and V caches
-    /// (each FP32, [maxSeqLen, kvDim] per layer), with per-layer / SWA-ring sizing for
-    /// gemma4-style models.
+    /// <see cref="EstimateAvailableKvVram"/> and find the largest context whose KV cache fits,
+    /// at the element width <paramref name="kvDType"/> the forward pass will actually allocate.
+    /// For gemma4-style models (per-layer head_dim / SWA rings / KV-share aliasing) this binary-
+    /// searches against <see cref="EstimateKvCacheBytes"/> — the same allocator-exact arithmetic
+    /// the constructor reserves — so bf16/q8_0 correctly buy ~2×/4× the positions of fp32 (#220).
+    /// Uniform-attention models keep the flat <c>NumLayers × kvDim × maxCtx</c> fp32 formula.
     /// </summary>
-    public static int EstimateMaxContext(GgufModel model, CudaBackend gpu, ModelHyperparams hp)
+    public static int EstimateMaxContext(
+        GgufModel model, CudaBackend gpu, ModelHyperparams hp, DType kvDType = DType.Float32)
+        => SolveMaxCtxForKv(hp, EstimateAvailableKvVram(model, gpu, hp), kvDType);
+
+    /// <summary>
+    /// Pure (GPU-free, GGUF-free) core of <see cref="EstimateMaxContext"/>: the largest context
+    /// whose KV cache fits <paramref name="availableKvBytes"/> at element width
+    /// <paramref name="kvDType"/>. Factored out for unit testing (cf. <see cref="ResolveKvDType"/>).
+    /// <para>Gemma 4-style models (per-layer head_dim + SWA pattern) binary-search against
+    /// <see cref="EstimateKvCacheBytes"/> — the allocator-exact arithmetic (dtype + SWA ring +
+    /// KV-share skip + per-layer KV heads + pow2 round-up) the ctor reserves and
+    /// <c>TierPlanner.SolveGpuCtxForPerLayerKv</c> uses — so the estimate can't drift from what is
+    /// actually allocated and bf16/q8_0 correctly buy more context (#220). Because SWA layers
+    /// stop growing past their ring cap, the gain over fp32 exceeds the bare width ratio once the
+    /// context clears that cap. Uniform-attention models keep the flat fp32 formula unchanged.</para>
+    /// </summary>
+    internal static int SolveMaxCtxForKv(ModelHyperparams hp, long availableKvBytes, DType kvDType)
     {
-        long available = EstimateAvailableKvVram(model, gpu, hp);
-        int headDim = hp.HeadDim;
+        const int floorCtx = 512;
+        int cap = hp.ContextLength;
+
+        // A model whose context is at/below the floor clamps to the cap (and avoids the
+        // Math.Clamp(_, 512, cap) below throwing when cap < 512). Mirrors the floor convention:
+        // return the small ctx and let the ctor's allocation fail loudly if even that won't fit.
+        if (cap <= floorCtx)
+            return cap;
 
-        // Gemma 4 per-layer head-dim path: each layer's K/V buffer takes its own
-        // head_dim, and SWA layers cap at SlidingWindowSize regardless of the
-        // global context window. Solve for the largest maxCtx s.t. the summed
-        // per-layer bytes still fit in `available`. Without this branch the
-        // non-gemma4 formula (NumLayers × headDim × maxCtx) wildly under- or
-        // over-counts depending on which side of the head-dim mix dominates.
-        if (hp.LayerHeadDim is { } lhd && hp.IsSwaLayer is { } swa)
+        if (hp.LayerHeadDim is not null && hp.IsSwaLayer is not null)
         {
-            int swaWindow = hp.SlidingWindowSize > 0 ? hp.SlidingWindowSize : int.MaxValue;
-            // SWA layers are sized as a ring of window + SwaRingHeadroom positions (issue
-            // #162), so the cap for the per-token byte formula is the ring size, not the
-            // bare window. Guard against overflow when swaWindow is "unbounded".
-            long swaCap = swaWindow == int.MaxValue ? long.MaxValue : (long)swaWindow + SwaRingHeadroom;
-            long globalKvDimPerToken = 0;
-            long swaKvDimPerToken    = 0;
-            for (int i = 0; i < hp.NumLayers; i++)
+            // EstimateKvCacheBytes is monotonic non-decreasing in ctx, so an upper-bound binary
+            // search converges. Floor 512 mirrors the uniform clamp below.
+            if (EstimateKvCacheBytes(hp, floorCtx, kvDType) > availableKvBytes)
+                return floorCtx;
+            int lo = floorCtx, hi = cap;
+            while (lo < hi)
             {
-                // KV-share layers don't allocate their own pages (the source layer
-                // already counted). Skip from both buckets.
-                if (hp.KvSourceLayer is { } ksl && ksl[i] >= 0) continue;
-                long layerKvDim = 2L * hp.NumKvHeads * lhd[i] * sizeof(float);
-                if (swa[i]) swaKvDimPerToken    += layerKvDim;
-                else        globalKvDimPerToken += layerKvDim;
-            }
-            // For a given maxCtx C: bytes = globalKvDimPerToken * C
-            //                            + swaKvDimPerToken    * min(C, swaCap)
-            // Solve for the largest C ≤ hp.ContextLength that fits in `available`.
-            // Branch on whether C ≤ swaCap:
-            //   if C ≤ swaCap: bytes = (global+swa) * C
-            //   else:          bytes = global * C + swa * swaCap
-            long globalPlusSwa = globalKvDimPerToken + swaKvDimPerToken;
-            int candA = globalPlusSwa > 0 ? (int)(available / globalPlusSwa) : int.MaxValue;
-            int maxCtxL;
-            if (candA <= swaCap)
-            {
-                maxCtxL = candA;
-            }
-            else
-            {
-                // swaCap is finite here (candA ≤ long.MaxValue always takes the branch
-                // above when swaCap is unbounded), so swaKvDimPerToken * swaCap is safe.
-                long remain = available - swaKvDimPerToken * swaCap;
-                int candB = globalKvDimPerToken > 0 && remain > 0
-                    ? (int)(remain / globalKvDimPerToken) : 0;
-                maxCtxL = (int)Math.Max(swaCap, candB);
+                int mid = lo + (hi - lo + 1) / 2;
+                if (EstimateKvCacheBytes(hp, mid, kvDType) <= availableKvBytes)
+                    lo = mid;
+                else
+                    hi = mid - 1;
             }
-            return Math.Clamp(maxCtxL, 512, hp.ContextLength);
+            return lo;
         }
 
-        long bytesPerToken = 2L * hp.NumLayers * hp.NumKvHeads * headDim * sizeof(float);
-        int maxCtx = (int)(available / bytesPerToken);
-        return Math.Clamp(maxCtx, 512, hp.ContextLength);
+        // Uniform-attention models: unchanged flat fp32 formula (#220 is scoped to the
+        // SWA/per-layer Gemma path; dtype-aware sizing for uniform models is out of scope).
+        long bytesPerToken = 2L * hp.NumLayers * hp.NumKvHeads * hp.HeadDim * sizeof(float);
+        int maxCtx = (int)(availableKvBytes / bytesPerToken);
+        return Math.Clamp(maxCtx, floorCtx, cap);
     }
 
     /// <summary>

tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs

@@ -818,6 +818,197 @@ public void SolveGpuCtxForPerLayerKv_BudgetBelowFloor_ReturnsFloor()
             hp, autoCtxCap: 256, vramBudget: long.MaxValue, dtype, gpuLayers: 3));
     }
 
+    // ── Issue #220: dense auto-context is dtype-aware (CudaForwardPass.SolveMaxCtxForKv) ──
+    // The full-GPU Gemma path's auto-context comes from EstimateMaxContext → SolveMaxCtxForKv,
+    // which previously priced fp32 unconditionally, so --kv-type bf16/q8_0 bought NO extra
+    // context (observed 1770 flat across all three dtypes on a 12 GB card). The fix binary-
+    // searches the largest ctx whose EstimateKvCacheBytes(.., kvDType) fits — the same
+    // allocator-exact arithmetic the ctor reserves — so narrowed KV expands the window. These
+    // pin: (1) the dtype response in a linear all-global regime (clean ratios), (2) the
+    // allocator-maximal contract on the SWA shape, (3) uniform-attention models are UNCHANGED.
+
+    /// <summary>
+    /// (1) Dtype response, linear regime. An all-global per-layer model (LayerHeadDim set,
+    /// IsSwaLayer all false → the per-layer binary-search branch, but no SWA-ring capping) is
+    /// linear in ctx, so a budget sized to exactly fit 4096-ctx fp32 yields a clean dtype
+    /// progression: bf16 doubles the ctx (half the per-element width), q8_0 ~3.76× it (its
+    /// 34-byte/32-elem blocks fall short of a clean 4×). Dims chosen so fp32/bf16 buffers land
+    /// exactly on power-of-two pool buckets (kvDim 1024 × power-of-two ctx), isolating the ratio
+    /// from round-up noise.
+    /// </summary>
+    [Fact]
+    public void SolveMaxCtxForKv_RespondsToDtype_LinearRegime()
+    {
+        var hp = Gemma4ShapedHp(
+            layerHeadDim: [128, 128, 128, 128],
+            layerKvHeads: [8, 8, 8, 8],          // kvDim = 1024
+            isSwa: [false, false, false, false], // all global → linear in ctx, no SWA cap
+            kvSource: [-1, -1, -1, -1],
+            slidingWindow: 4096);                // irrelevant (no SWA layer)
+
+        // Budget = exactly the fp32 footprint at 4096 ctx (each [4096×1024] fp32 buffer = 2^24).
+        long budget = CudaForwardPass.EstimateKvCacheBytes(hp, 4096, DType.Float32);
+
+        int fp32 = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.Float32);
+        int bf16 = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.BFloat16);
+        int q8   = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.Q8_0);
+
+        Assert.Equal(4096, fp32);
+        Assert.Equal(8192, bf16);                 // exactly 2× — half the element width
+        Assert.Equal(2 * fp32, bf16);
+        Assert.True(q8 > bf16 && q8 >= 15000,     // ~3.76× (q8_0's 34/32 block overhead < clean 4×)
+            $"q8_0 ctx {q8} should be well past 2× fp32 ({fp32}) — narrowed KV must expand the window.");
+        // Allocator-exact: the chosen q8_0 ctx fits the budget and the next step up does not.
+        Assert.True(CudaForwardPass.EstimateKvCacheBytes(hp, q8, DType.Q8_0) <= budget);
+        Assert.True(CudaForwardPass.EstimateKvCacheBytes(hp, q8 + 1, DType.Q8_0) > budget);
+    }
+
+    /// <summary>
+    /// (2) Allocator-maximal on the SWA shape (the #220 contract). Sizing the budget to a
+    /// reference context's bf16 footprint, the solver must fit (never under-reserve → no
+    /// runtime OOM), return at least the reference, and be maximal (the next ctx overflows).
+    /// Uses the gemma4 shape (per-layer head_dim + SWA ring + KV-share aliasing) so the search
+    /// runs against the real per-layer allocator arithmetic.
+    /// </summary>
+    [Fact]
+    public void SolveMaxCtxForKv_SwaShape_IsAllocatorMaximal()
+    {
+        const int refCtx = 8192;
+        var dtype = DType.BFloat16;
+        var hp = Gemma4ShapedHp(
+            layerHeadDim: [256, 256, 256],
+            layerKvHeads: [8, 8, 8],
+            isSwa: [false, true, false],
+            kvSource: [-1, -1, 0],   // layer 2 aliases layer 0
+            slidingWindow: 1024);
+
+        long budget = CudaForwardPass.EstimateKvCacheBytes(hp, refCtx, dtype);
+        int got = CudaForwardPass.SolveMaxCtxForKv(hp, budget, dtype);
+
+        Assert.True(CudaForwardPass.EstimateKvCacheBytes(hp, got, dtype) <= budget,
+            $"solved ctx {got} over-reserves vs budget {budget} — would OOM at runtime (#220).");
+        Assert.True(got >= refCtx, $"solved ctx {got} below the reference {refCtx}, which fits exactly.");
+        Assert.True(got == hp.ContextLength ||
+            CudaForwardPass.EstimateKvCacheBytes(hp, got + 1, dtype) > budget,
+            $"ctx {got + 1} also fits budget {budget} — solver did not return the LARGEST fitting context.");
+    }
+
+    /// <summary>
+    /// (3) Uniform-attention models are UNCHANGED: a flat model (no LayerHeadDim / IsSwaLayer)
+    /// keeps the fp32 formula regardless of the requested KV dtype, so bf16/q8_0 do NOT alter
+    /// its auto-context. This is the #220 acceptance "no change for uniform-attention models"
+    /// guard — the dtype-aware sizing is scoped to the SWA/per-layer Gemma path only.
+    /// </summary>
+    [Fact]
+    public void SolveMaxCtxForKv_UniformModel_IgnoresDtype()
+    {
+        var hp = FlatHp(numLayers: 8, numKvHeads: 8, headDim: 128, ctx: 131072);
+        long budget = 256L * 1024 * 1024;
+        int fp32 = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.Float32);
+        int bf16 = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.BFloat16);
+        int q8   = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.Q8_0);
+        Assert.True(fp32 > 512, $"sanity: expected a mid-range ctx, got {fp32}.");
+        Assert.Equal(fp32, bf16);   // dtype ignored for uniform models — no change vs pre-#220
+        Assert.Equal(fp32, q8);
+    }
+
+    /// <summary>
+    /// (1b) The headline #220 claim: once the context clears the SWA ring cap, the SWA layers
+    /// stop growing, so a narrower KV dtype's freed budget flows entirely into the (few) global
+    /// layers — gaining MORE than the bare width ratio (2×/4×). Shape: 1 global + 5 SWA layers
+    /// with a 512 window (ring = 4608); the budget is sized to an fp32 context (8192) already
+    /// PAST that ring, so all SWA layers are capped for every dtype. Relational asserts (not
+    /// magic numbers) so the test is robust to the exact pow2-bucket arithmetic.
+    /// </summary>
+    [Fact]
+    public void SolveMaxCtxForKv_SwaSaturation_DtypeGainExceedsWidthRatio()
+    {
+        var hp = Gemma4ShapedHp(
+            layerHeadDim: [256, 256, 256, 256, 256, 256],
+            layerKvHeads: [8, 8, 8, 8, 8, 8],                 // kvDim = 2048
+            isSwa: [false, true, true, true, true, true],      // 1 global, 5 SWA → SWA dominates
+            kvSource: [-1, -1, -1, -1, -1, -1],
+            slidingWindow: 512);                               // ring = min(ctx, 512+4096) = 4608
+
+        // Budget = fp32 footprint at ctx 8192 (> 4608 ring → SWA layers capped for all dtypes).
+        long budget = CudaForwardPass.EstimateKvCacheBytes(hp, 8192, DType.Float32);
+
+        int fp32 = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.Float32);
+        int bf16 = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.BFloat16);
+        int q8   = CudaForwardPass.SolveMaxCtxForKv(hp, budget, DType.Q8_0);
+
+        Assert.True(q8 > bf16 && bf16 > fp32, $"monotonic dtype response expected; got fp32={fp32} bf16={bf16} q8={q8}.");
+        // Super-linear: past the SWA cap only the global layers grow, so narrowing beats the
+        // width ratio. (Asserting strict > the ratio, with margin from the 5:1 SWA:global mix.)
+        Assert.True(bf16 > 2 * fp32, $"bf16 ctx {bf16} should exceed 2× fp32 ({fp32}) once SWA layers are capped.");
+        Assert.True(q8 > 4 * fp32, $"q8_0 ctx {q8} should exceed 4× fp32 ({fp32}) once SWA layers are capped.");
+        // Allocator-maximal for the narrowest dtype (fits + next step overflows or hits the cap).
+        Assert.True(CudaForwardPass.EstimateKvCacheBytes(hp, q8, DType.Q8_0) <= budget);
+        Assert.True(q8 == hp.ContextLength ||
+            CudaForwardPass.EstimateKvCacheBytes(hp, q8 + 1, DType.Q8_0) > budget);
+    }
+
+    /// <summary>
+    /// (1c) Floor and cap clamps on the per-layer branch (distinct from the sibling
+    /// SolveGpuCtxForPerLayerKv floor test — different signature: the cap here is hp.ContextLength).
+    /// A budget too small for even 512-ctx returns the floor (the alloc then fails loudly, not a
+    /// silently-smaller context); a model whose ContextLength is below 512 clamps to that; and a
+    /// huge budget clamps UP to the model max (not unbounded).
+    /// </summary>
+    [Fact]
+    public void SolveMaxCtxForKv_ClampsToFloorAndModelMax()
+    {
+        var hp = Gemma4ShapedHp(
+            layerHeadDim: [256, 256, 256],
+            layerKvHeads: [8, 8, 8],
+            isSwa: [false, true, false],
+            kvSource: [-1, -1, 0],
+            slidingWindow: 1024);
+
+        // Budget = 1 byte → can't hold even a 512-ctx cache → floor (512).
+        Assert.Equal(512, CudaForwardPass.SolveMaxCtxForKv(hp, 1, DType.BFloat16));
+
+        // ContextLength below the floor → clamp to the cap even with a huge budget (must NOT throw
+        // — the SWA branch and the uniform branch both route through the cap<=floor guard, else the
+        // uniform Math.Clamp(_, 512, cap) would throw ArgumentException for cap < 512).
+        var tinyCapSwa = hp with { ContextLength = 256 };
+        Assert.Equal(256, CudaForwardPass.SolveMaxCtxForKv(tinyCapSwa, long.MaxValue, DType.BFloat16));
+        var tinyCapUniform = FlatHp(numLayers: 8, numKvHeads: 8, headDim: 128, ctx: 256);
+        Assert.Equal(256, CudaForwardPass.SolveMaxCtxForKv(tinyCapUniform, long.MaxValue, DType.Float32));
+
+        // Huge budget → clamp UP to the model max, not beyond.
+        Assert.Equal(hp.ContextLength, CudaForwardPass.SolveMaxCtxForKv(hp, long.MaxValue, DType.Q8_0));
+    }
+
+    /// <summary>
+    /// (1d) Mixed per-layer head_dim: the solver must price each layer at its own head_dim (via
+    /// EstimateKvCacheBytes), not collapse to hp.HeadDim or layer 0. A shape with distinct
+    /// per-layer dims, asserted allocator-maximal (the contract that depends on the per-layer
+    /// arithmetic being exact).
+    /// </summary>
+    [Fact]
+    public void SolveMaxCtxForKv_MixedPerLayerHeadDim_IsAllocatorMaximal()
+    {
+        const int refCtx = 8192;
+        var dtype = DType.Q8_0;
+        var hp = Gemma4ShapedHp(
+            layerHeadDim: [256, 128, 256, 128],   // mixed per-layer head_dim
+            layerKvHeads: [8, 8, 8, 8],
+            isSwa: [false, true, false, true],
+            kvSource: [-1, -1, -1, -1],
+            slidingWindow: 1024);
+
+        long budget = CudaForwardPass.EstimateKvCacheBytes(hp, refCtx, dtype);
+        int got = CudaForwardPass.SolveMaxCtxForKv(hp, budget, dtype);
+
+        Assert.True(CudaForwardPass.EstimateKvCacheBytes(hp, got, dtype) <= budget,
+            $"mixed-headdim solved ctx {got} over-reserves vs budget {budget}.");
+        Assert.True(got >= refCtx, $"solved ctx {got} below the reference {refCtx}, which fits exactly.");
+        Assert.True(got == hp.ContextLength ||
+            CudaForwardPass.EstimateKvCacheBytes(hp, got + 1, dtype) > budget,
+            $"ctx {got + 1} also fits — not the largest fitting context for the mixed-headdim shape.");
+    }
+
     /// <summary>
     /// Q8KvGeometrySupported returns false when ANY single (non-aliased) layer violates the
     /// %32 rule — not just when all do. A mixed set with one bad layer must fail, matching