diff --git a/src/SharpInference.Engine/CudaForwardPass.cs b/src/SharpInference.Engine/CudaForwardPass.cs
index cae74ec..c7806be 100644
--- a/src/SharpInference.Engine/CudaForwardPass.cs
+++ b/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(
///
/// VRAM-based context-length estimator: take the KV-cache budget from
- /// 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.
+ /// and find the largest context whose KV cache fits,
+ /// at the element width the forward pass will actually allocate.
+ /// For gemma4-style models (per-layer head_dim / SWA rings / KV-share aliasing) this binary-
+ /// searches against — 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 NumLayers × kvDim × maxCtx fp32 formula.
///
- 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);
+
+ ///
+ /// Pure (GPU-free, GGUF-free) core of : the largest context
+ /// whose KV cache fits at element width
+ /// . Factored out for unit testing (cf. ).
+ /// Gemma 4-style models (per-layer head_dim + SWA pattern) binary-search against
+ /// — the allocator-exact arithmetic (dtype + SWA ring +
+ /// KV-share skip + per-layer KV heads + pow2 round-up) the ctor reserves and
+ /// TierPlanner.SolveGpuCtxForPerLayerKv 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.
+ ///
+ 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);
}
///
diff --git a/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs b/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs
index ec4bd59..4587fe4 100644
--- a/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs
+++ b/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.
+
+ ///
+ /// (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.
+ ///
+ [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);
+ }
+
+ ///
+ /// (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.
+ ///
+ [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.");
+ }
+
+ ///
+ /// (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.
+ ///
+ [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);
+ }
+
+ ///
+ /// (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.
+ ///
+ [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);
+ }
+
+ ///
+ /// (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).
+ ///
+ [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));
+ }
+
+ ///
+ /// (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).
+ ///
+ [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.");
+ }
+
///
/// 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