From 1f4363ceb5e34e9977275a77a67abfc84b5d81c8 Mon Sep 17 00:00:00 2001 From: Pekka Heikura Date: Sat, 13 Jun 2026 14:28:11 +0300 Subject: [PATCH 1/3] perf(cuda): dtype-aware auto-context for Gemma 4 (#220) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dense full-GPU Gemma path's auto-context (-g -1, no -c) came from CudaForwardPass.EstimateMaxContext, which priced KV at fp32 unconditionally — so --kv-type bf16/q8_0 bought NO extra context. On a 4070 Ti 12 GB, Gemma 4 12B was pinned to 1770 tokens (model max 262144) for fp32 AND bf16 AND q8_0. EstimateMaxContext already had a Gemma SWA/KV-share/per-layer-headdim branch but hardcoded `sizeof(float)` (and used NumKvHeads, no pow2 round-up), so it was dtype-blind and could drift from the real allocator. Replace its hand-rolled math with a binary search against EstimateKvCacheBytes(hp, ctx, kvDType) — the single source of truth the ctor's allocation and TierPlanner.SolveGpuCtxForPerLayerKv both use (dtype + SWA ring + KV-share skip + per-layer KV heads + pow2 round-up). The core is factored into the pure, GPU/GGUF-free CudaForwardPass.SolveMaxCtxForKv (mirrors the ResolveKvDType factoring) and the call site passes the requested KV dtype (ResolveConfiguredKvDType). Uniform-attention models keep the flat fp32 formula unchanged. Result (4070 Ti 12 GB, gemma-4-12b-it-qat-q4_0, -g -1): fp32 1770 -> 2048 (allocator-exact), bf16 -> 4096 (2x), q8_0 -> 30840 (17x). q8_0 exceeds 4x because SWA layers stop growing past their ring cap, so beyond it only the global layers consume KV per token. Verified the q8_0 30840-ctx run constructs + decodes coherently (53 t/s, no OOM); Qwen3-8B dense auto-ctx unchanged (12080). Tests: 3 GGUF-free unit tests for SolveMaxCtxForKv (dtype response in the linear regime: bf16 = 2x fp32, q8_0 ~3.76x; allocator-maximal on the SWA shape; uniform models ignore dtype). Existing KV-dtype parity + Gemma coherence + SWA ring-wrap tests green (they pass explicit ctx, unaffected). Release build clean under TreatWarningsAsErrors + AOT analyzers. Co-Authored-By: Claude Opus 4.8 --- src/SharpInference.Engine/CudaForwardPass.cs | 105 +++++++++--------- .../CudaForwardPassKvDtypeTests.cs | 94 ++++++++++++++++ 2 files changed, 145 insertions(+), 54 deletions(-) diff --git a/src/SharpInference.Engine/CudaForwardPass.cs b/src/SharpInference.Engine/CudaForwardPass.cs index cae74ec..570582f 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,57 @@ 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; - // 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 (cap <= floorCtx || EstimateKvCacheBytes(hp, floorCtx, kvDType) > availableKvBytes) + return Math.Min(floorCtx, cap); + 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..bd4c6bc 100644 --- a/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs +++ b/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs @@ -818,6 +818,100 @@ 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); + } + /// /// 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 From 1f7317104db6479016ce3d0d4358521f939557d5 Mon Sep 17 00:00:00 2001 From: Pekka Heikura Date: Sat, 13 Jun 2026 14:52:39 +0300 Subject: [PATCH 2/3] =?UTF-8?q?test(cuda):=20harden=20SolveMaxCtxForKv=20c?= =?UTF-8?q?overage=20=E2=80=94=20SWA-saturation,=20floor/cap,=20mixed=20he?= =?UTF-8?q?ad=5Fdim=20(#220=20review)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review-cycle follow-up addressing the pr-test-analyzer's two must-fix gaps: - SolveMaxCtxForKv_SwaSaturation_DtypeGainExceedsWidthRatio: the headline #220 claim (past the SWA ring cap, narrowing the KV dtype gains MORE than the bare width ratio because only the global layers grow) had zero coverage — the linear-regime test uses isSwa all-false. 1 global + 5 SWA layers, budget sized past the ring, asserts bf16 > 2× fp32 and q8 > 4× fp32 (relational, not magic numbers) + q8 allocator-maximal. - SolveMaxCtxForKv_ClampsToFloorAndModelMax: the floor early-return / cap clamp on the per-layer branch (the sibling SolveGpuCtxForPerLayerKv floor test doesn't transitively cover it — different signature). budget=1 → 512; ContextLength 256 → 256; huge budget → model max. - SolveMaxCtxForKv_MixedPerLayerHeadDim_IsAllocatorMaximal: distinct per-layer head_dim, allocator-maximal — guards the per-layer pricing not collapsing to hp.HeadDim. All 6 SolveMaxCtxForKv tests green. Test-only. Co-Authored-By: Claude Opus 4.8 --- .../CudaForwardPassKvDtypeTests.cs | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs b/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs index bd4c6bc..389e880 100644 --- a/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs +++ b/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs @@ -912,6 +912,99 @@ public void SolveMaxCtxForKv_UniformModel_IgnoresDtype() 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 (Math.Min(512, cap)) even with a huge budget. + var tinyCap = hp with { ContextLength = 256 }; + Assert.Equal(256, CudaForwardPass.SolveMaxCtxForKv(tinyCap, long.MaxValue, DType.BFloat16)); + + // 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 From 95b5d27e4c0fda10cdeff9f94a56cc7193dfc809 Mon Sep 17 00:00:00 2001 From: Pekka Heikura Date: Sat, 13 Jun 2026 14:54:52 +0300 Subject: [PATCH 3/3] fix(cuda): guard SolveMaxCtxForKv against ContextLength < floor (review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit gemini-code-assist (HIGH) flagged that the uniform branch's Math.Clamp(maxCtx, floorCtx, cap) throws ArgumentException when hp.ContextLength < 512 (min > max). Pre-existing on master (the refactor preserved it verbatim), but the SWA branch already handled cap <= floor gracefully — an ugly asymmetry. Hoist a single `if (cap <= floorCtx) return cap;` guard to the top so both branches are safe and consistent, and simplify the SWA early-return (cap > floorCtx is now guaranteed). Added a uniform-shape cap<512 case to SolveMaxCtxForKv_ClampsToFloorAndModelMax (would have thrown before). 6/6 SolveMaxCtxForKv tests green. Co-Authored-By: Claude Opus 4.8 --- src/SharpInference.Engine/CudaForwardPass.cs | 10 ++++++++-- .../CudaForwardPassKvDtypeTests.cs | 10 +++++++--- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/src/SharpInference.Engine/CudaForwardPass.cs b/src/SharpInference.Engine/CudaForwardPass.cs index 570582f..c7806be 100644 --- a/src/SharpInference.Engine/CudaForwardPass.cs +++ b/src/SharpInference.Engine/CudaForwardPass.cs @@ -4103,12 +4103,18 @@ internal static int SolveMaxCtxForKv(ModelHyperparams hp, long availableKvBytes, 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; + if (hp.LayerHeadDim is not null && hp.IsSwaLayer is not null) { // EstimateKvCacheBytes is monotonic non-decreasing in ctx, so an upper-bound binary // search converges. Floor 512 mirrors the uniform clamp below. - if (cap <= floorCtx || EstimateKvCacheBytes(hp, floorCtx, kvDType) > availableKvBytes) - return Math.Min(floorCtx, cap); + if (EstimateKvCacheBytes(hp, floorCtx, kvDType) > availableKvBytes) + return floorCtx; int lo = floorCtx, hi = cap; while (lo < hi) { diff --git a/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs b/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs index 389e880..4587fe4 100644 --- a/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs +++ b/tests/SharpInference.Tests.ForwardPass/CudaForwardPassKvDtypeTests.cs @@ -968,9 +968,13 @@ public void SolveMaxCtxForKv_ClampsToFloorAndModelMax() // 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 (Math.Min(512, cap)) even with a huge budget. - var tinyCap = hp with { ContextLength = 256 }; - Assert.Equal(256, CudaForwardPass.SolveMaxCtxForKv(tinyCap, long.MaxValue, 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));