Skip to content

Improve Gemma4 long-context correctness and Metal decode performance - #473

Open
timkaye11 wants to merge 3 commits into
antflydb:mainfrom
timkaye11:gemma4-long-context
Open

Improve Gemma4 long-context correctness and Metal decode performance#473
timkaye11 wants to merge 3 commits into
antflydb:mainfrom
timkaye11:gemma4-long-context

Conversation

@timkaye11

Copy link
Copy Markdown
Contributor

[codex]

Summary

This PR improves Gemma4 E2B/E4B long-context inference

  • Fixes mixed sliding-window/global-attention KV behavior, bounded chunked prefill, paging, cache growth, and ownership for long prompts.
  • Replaces the fixed 2,048-token request cap with model-aware limits and route-aware KV admission that reserve output capacity and fail conservatively outside verified paged Metal paths.
  • Applies Gemma4 partial-rotary RoPE theta exactly once across native and Metal execution while preserving long-context MTP and pipelined-frame correctness.
  • Adds a bounded, fail-safe Metal split-GQA decode path for supported Gemma4 shapes, with telemetry, a kill switch, and automatic fallback.
  • Hardens shared generation admission with bounded speculative width, exact Qwen3.5 image-token accounting, expanded-context revalidation, and consistent benchmark ignore_eos handling.
  • Keeps short-context, unsupported shapes, and non-paged fallbacks on their existing conservative paths.

Benchmarks

Local, unattested Apple Metal measurements use ReleaseFast builds and fresh processes. The comparator is the branch's exact main fork point; decode excludes prefill.

Workload Main This PR Result
Gemma4 E2B, 470+64 34.50 tok/s 39.98 tok/s +15.9% decode; 6.7% lower total time
Gemma4 E2B, 2,010+64 34.78 tok/s 49.84 tok/s +43.3% decode; 4.1% lower total time
Gemma4 E4B, 444+64 21.26 tok/s 21.43 tok/s Performance parity; 0.8% lower total time
Gemma4 E4B, 2,002+300 18.85 tok/s 26.93 tok/s +42.9% decode; 21.1% lower total time

An additional E2B 2,065+300 long-output check used one warmup and five measured fresh processes:

TTFT Prefill Decode Generation Total
7,946 ms 259.88 tok/s 49.64 tok/s 13,990 ms 14,647 ms

Validation

  • Short E2B/E4B controls produced identical token IDs versus main.
  • Long E2B and E4B outputs were token-exact with the optimized Metal route enabled or disabled.
  • E4B 2K+300 recorded all 12,600 expected split-GQA dispatches with zero fallbacks.
  • Focused admission, bounded-prefill, generation-route, and memory-estimator tests passed 7/7.
  • Default Metal, -Dmetal=false, and ReleaseFast Metal builds passed.
  • The E2B 2,065+300 run completed deterministically at the requested 300-token length.
  • OpenAPI assembly is current, including the supported speculative-width bounds.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant