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

[pekkah]

Summary

Fixes #220. On a 12 GB card, Gemma 4 12B's auto-context (-g -1, no -c) was pinned at 1770 tokens for fp32 AND bf16 AND q8_0 — --kv-type bought no extra context, despite the model supporting 262144.

The dense full-GPU path's auto-context comes from CudaForwardPass.EstimateMaxContext, which 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. (TierPlanner.RecommendedCtxSize is computed but discarded for the full-GPU case; the live source is EstimateMaxContext.)

Fix

Replace EstimateMaxContext's hand-rolled SWA 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); 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)

--kv-type	before	after
fp32	1770	2048 (allocator-exact)
bf16	1770	4096 (2×)
q8_0	1770	30840 (17×)

q8_0 exceeds 4× 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 new GGUF-free unit tests for SolveMaxCtxForKv (CudaForwardPassKvDtypeTests):

• dtype response in the linear regime — bf16 = 2× fp32, q8_0 ~3.76× (q8_0's 34/32-byte blocks fall short of a clean 4×);
• allocator-maximal on the SWA shape (fits, never under-reserves, maximal);
• uniform models ignore dtype (the "no change for uniform models" guard).

Existing KV-dtype parity + Gemma coherence + SWA ring-wrap tests stay green (they pass explicit ctx, so are unaffected). Release build clean under TreatWarningsAsErrors + AOT analyzers.

Follow-up (separate)

The auto value (30840) is still ~8× below what actually fits (full 256K runs with -c), because EstimateAvailableKvVram reserves a blunt max(VRAM/3, 2GB). That's tracked in #228 (gated on a spill-cliff A/B) — deliberately out of scope here.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request updates the auto-context length estimation in CudaForwardPass to be aware of the requested KV data type (kvDType) for Gemma 4-style models, allowing narrower data types like bf16 and q8_0 to correctly expand the context window. Uniform-attention models remain unchanged, and comprehensive unit tests have been added to verify these behaviors. The reviewer identified a potential ArgumentException in SolveMaxCtxForKv when hp.ContextLength is less than floorCtx (512) due to invalid arguments in Math.Clamp, and provided a code suggestion to handle this case early.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If hp.ContextLength (cap) is less than floorCtx (512), Math.Clamp(maxCtx, floorCtx, cap) will throw an ArgumentException because the minimum value (512) is greater than the maximum value (cap).

To prevent this, we can check if cap <= floorCtx at the beginning of the method and return cap immediately. This also simplifies the SWA branch and avoids potential division-by-zero issues if bytesPerToken is zero.

    internal static int SolveMaxCtxForKv(ModelHyperparams hp, long availableKvBytes, DType kvDType)
    {
        const int floorCtx = 512;
        int cap = hp.ContextLength;

        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 (EstimateKvCacheBytes(hp, floorCtx, kvDType) > availableKvBytes)
                return floorCtx;
            int lo = floorCtx, hi = cap;
            while (lo < hi)
            {
                int mid = lo + (hi - lo + 1) / 2;
                if (EstimateKvCacheBytes(hp, mid, kvDType) <= availableKvBytes)
                    lo = mid;
                else
                    hi = mid - 1;
            }
            return lo;
        }

        // 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);
        if (bytesPerToken == 0)
            return floorCtx;
        int maxCtx = (int)(availableKvBytes / bytesPerToken);
        return Math.Clamp(maxCtx, floorCtx, cap);
    }

[pekkah]

Review cycle complete

Ran three independent review lenses (correctness, silent-failure/OOM, test-coverage) plus the automated gemini-code-assist review.

Correctness + silent-failure: no high-confidence issues. Confirmed the binary search is allocator-exact and strictly safer than master on OOM — master's old Gemma formula summed raw bytes with no pow2 round-up and could hand back a context that cudaMalloc-fails; the new search rounds each buffer exactly like the pool. Requested-vs-actual KV dtype can't diverge to under-reserve (auto-narrow only narrows, never widens). No integer overflow (long arithmetic; the (int) cast is only on the unchanged uniform path).

Addressed findings:

• gemini (HIGH): Math.Clamp(_, 512, cap) in the uniform branch threw ArgumentException for ContextLength < 512. Fixed with a top-level if (cap <= floorCtx) return cap; guard (95b5d27) — both branches now safe + consistent; added a uniform-shape cap<512 test.
• pr-test-analyzer (2× must-fix): the SWA-saturation super-linear claim and the floor/cap clamp had no coverage. Added SolveMaxCtxForKv_SwaSaturation_DtypeGainExceedsWidthRatio, _ClampsToFloorAndModelMax, _MixedPerLayerHeadDim_IsAllocatorMaximal (1f73171).

Final: 6/6 SolveMaxCtxForKv unit tests green, existing KV-dtype/coherence/ring tests green, Release build clean under TreatWarningsAsErrors + AOT analyzers. Merging.