Skip to content

feat(gemma4): local decoder layer and its HF golden oracle - #875

Merged
FeathBow merged 1 commit into
pegainfer-project:mainfrom
FeathBow:feat/gemma4-local-layer
Aug 14, 2026
Merged

feat(gemma4): local decoder layer and its HF golden oracle#875
FeathBow merged 1 commit into
pegainfer-project:mainfrom
FeathBow:feat/gemma4-local-layer

Conversation

@FeathBow

Copy link
Copy Markdown
Collaborator

Description

Closes #874

The line's first real execution path: one local (sliding-attention) decoder layer, forward only, no KV cache — K and V stay contiguous into single_prefill, which is exact sliding attention for any prompt the 1024-token window never truncates. That equivalence only holds from position zero and inside the window, so the layer rejects start_pos > 0 and window-exceeding lengths instead of mis-computing them.

The load-bearing constants are the ones the issue calls out: attention is unscaled (scaling = 1.0), V takes a weightless RMS norm on local layers too (expressed as the plain-w norm with a ones weight — no new kernel), layer_scalar multiplies the layer output after both residual adds, all four norm sites are norm-then-add, and RoPE rotates the full 256-wide head at the sliding theta. The probe layers come from parsing the layer map, and the oracle asserts that parse against the fixture metadata's own parse before any numeric comparison.

attention routes through the hd256 single_prefill instantiation, whose contiguous cache is HND, so K and V are reassembled per head with existing copy ops (the width-generic-looking NHD entry validates head_dim against 128); the RoPE tables come from the shared core precompute at full head width (the identity case of the frequency/rotary split); layer_scalar is read at load as the host f32 it ends up as; the oracle re-verifies the checkpoint against the fixture's pinned hashes before comparing (full-file for the config files, header-only for the 22 GiB weights file). The layer-map parser and its tests run in the featureless build; the fixture plumbing and the oracle are test-only.

Test Env

Single GPU (sm_89, x86_64), CUDA 12.9, against the pinned 12B checkpoint.

Verification

  • Four comparisons (single/short x first/last sliding layer) against the declared tolerance of 0.4 absolute + 2% relative; non-finite values fail outright:
    • Both one-token comparisons are bitwise exact (max_abs = 0).
    • The nine-token comparisons show scattered rounding noise only: max_abs 0.1875 (layer 0) and 0.25 (layer 46), zero elements over tolerance, token 0 exact in both — the unscaled-attention signature (no rsqrt damping, so one-ulp bf16 GEMM differences shift softmax weights), deterministic across runs, not the whole-block scramble a layout defect produces.
  • Negative controls on the full checkpoint: comparing against each probe's input instead of its output turns all four comparisons red at max_abs 27-56; pinning sm_scale back to rsqrt(256) decisively fails both nine-token comparisons (max_abs 13-22) while the one-token comparisons stay exact — softmax over one key is scale-invariant, which is why the multi-token probe exists.
  • Memory: the 12B BF16 load resides ~22.2 GiB on the 48 GiB card; the --ignored gates run serially.

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Documentation update

Signed-off-by: Feathbow <feathbow@gmail.com>
@FeathBow
FeathBow merged commit feb5c7c into pegainfer-project:main Aug 14, 2026
13 checks passed
@FeathBow
FeathBow deleted the feat/gemma4-local-layer branch August 14, 2026 17:14
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.

gemma4: one local decoder layer, forward only, against the HF golden probes

1 participant