Skip to content

gpu, gpu-gemma-embedding2, gpu-gates: the EmbeddingGemma 2 encoder from the common kernels, with its end-to-end gate - #299

Merged
midagedev merged 11 commits into
midagedev:mainfrom
juliankang4:dgx-spark-emb2-body
Oct 16, 2026
Merged

midagedev merged 11 commits into
midagedev:mainfrom
juliankang4:dgx-spark-emb2-body

Conversation

@juliankang4

@juliankang4 juliankang4 commented Oct 15, 2026 •

Copy link
Copy Markdown
Collaborator

Part of #59 (plan and PR order in #190), E5: the EmbeddingGemma 2 encoder, crates/gpu-gemma-embedding2, and its end-to-end gate against mainline llama.cpp's reference set.

What and why

The encoder adds no device code of its own. One call runs bloomery-gpu's kernels over an input's ids in llama.cpp's graph order (src/models/gemma-embedding2.cpp) and returns the 768-wide row of every position and their mean. With it, program_of returns a program for gemma-embedding2, and the gate just gate-gpu-gemma-embedding2-e2e holds every layer of it to mainline's text set.

Release notes: Added: the EmbeddingGemma 2 encoder on the card (gpu-gemma-embedding2) and its end-to-end gate against llama.cpp mainline; no server seat reads it yet.

The PR sits on main and needs no other open PR. It reads what the earlier steps of #190 put there: the reader (#238), the gemma4 tokenizer's id set (#241), the reference sets and their families (#260), and the head-512 prefill entries, the prefill window and the GELU-gated entry (#233, #234, #237). The /v1/embeddings seat is the next step, E6.

Common owners

The rows of docs/contrib/new-model.md's common-owner table, cut to what this PR touches:

Owner This PR Why
Placement plan (placement::plan_routed) does not attach The model has no routed expert. The whole encoder is one fixed byte figure, card_bytes(split, ctx); the serve seat (E6 of #190) places it.
Kernel table (placement/src/kernels.rs) already attached Every q8_0 matrix of the encoder goes through site::gemm on the table's Q8_0 entry (Gemm32Weight::Q8_0Plane for a dense projection), and per_layer_model_proj through the table's F32 entry, enqueue_f32_tile. The load refuses a matrix of another type by name. One part differs from #190: see the paragraph below the table.
Card expert leg does not attach A dense model: no expert.
Residency does not attach No expert; every resident weight stays on the card.
Expert tier does not attach No expert.
Prompt walk (runtime/src/prompt.rs) does not attach Every position attends both ways, so one input is one call over all its positions, not a walk in pieces.
Resident slots does not attach Nothing stays between inputs: the call rewrites its cache rows each time.
Serve seat does not attach here The /v1/embeddings seat is E6 of #190.
Timing does not attach here Through that seat, as for the others.

per_layer_model_proj does not load through Weights as one Bf16AsF32 matrix, as #190 planned. The encoder decodes it on the host and uploads each layer's 512 rows as their own f32 tensor, so the f32 tile writes that layer's slice of the product contiguous, each value the whole projection's (lib.rs's ple_rows). As one matrix, each position's product would hold the 24 layers' slices side by side, and every layer's norm would read a strided slice. The bytes are the same as Bf16AsF32's, four a value, and card_bytes counts them.

What the PR adds

By contract, with the commits that carry each part:

  1. gpu: head_norm_neox_append_256_whole. The sliding layers turn a head of 256 values whole (NEOX pairs (i, i + 128)) and read their queries from plain query rows, as the global layers do at head 512. The new entry is head_norm_neox_append_512's body (norm_partial_append) at HEAD = ROT = 256 over the f16 planes, 128 threads a head. enqueue_head_norm_neox_append_256_whole takes the 512 entry's PartialNeoxArgs and refuses a short buffer by name through partial_launch. It is an unsafe fn: its # Safety section says the positions must be distinct, because two tokens at one position write the same cache rows from two blocks. Its rows are pinned in the three ptx-shapes tables.

  2. model: gemma-embedding2's CARD_CTX The header's context_length is 262,144, but the model card names an input length of 8,192. The file does not carry that number, so the reader holds it as one constant, hparams::CARD_CTX, beside the keys it reads. The encoder's load bound reads it, and the /v1/embeddings seat (E6 of model: EmbeddingGemma and EmbeddingGemma 2 on a side seat #190) will read the same constant.

  3. models, model, placement, gpu-gates: the encoder parts are needs. The reader (model, models: gemma-embedding2's header reader, with an EncoderSpec on ModelSpec #238) put the encoder's own parts under the layer program item. They are now needs of their own. EncoderSpec gains four bools, one per fact of llama.cpp's graph that no other spec field carries:

    • embed_scale: the token rows are multiplied by √hidden before the first layer.
    • value_norm: each value head is RMS-normed with no gain.
    • post_norms: the attention and the feed-forward outputs are RMS-normed with a gain before each residual add.
    • layer_scale: each layer's output is multiplied by a trained scalar.

    The gemma-embedding (v1) spec sets the same four: embed_scale and post_norms on, value_norm and layer_scale off.

    needs() lists them with the per-layer input and the pooled output. Program::GemmaEmbedding2Body covers the shapes the encoder runs: the head-256 flash with a first key a row, the head-512 flash over every key, both whole NEOX turns, the GELU GLU at 2048, the per-layer input at 512 and the 768-wide output. Each AVAILABLE row names the function that runs it. Its matrices are pinned to q8_0, and its token table and per-layer projection to the host's decoded rows. It runs no routed expert, so leg_arms is None, routed_rule puts nothing on a card, and the routed-legs gate skips its sets as "no routed expert". The meta test now wants no coverage item for the file.

  4. gpu-gemma-embedding2, gpu, tools: the encoder. The new crate:

    • Encoder::open(split, ctx, gpu) loads the file for calls of up to ctx positions. It refuses by name a context of 0 or past CARD_CTX or the file's context_length, and a hidden width that is not a positive multiple of the f32 tile's line. It runs the program's coverage check, so it also refuses any part the program does not run (a partial rope, a head other than 256 or 512, a window on a head-512 layer, a group the flash pack does not divide, a width other than 2048, 512 or 768). A tensor of a type it does not read is refused too (q8_0 for a matrix, f32 for a vector, a host-decoded type for the token table and per_layer_model_proj). So is a load whose resident tensors, per-layer projection, rope tables, quantizer and route scratch, f32 scratch, cache planes and constant rows add up to more than card_bytes.
    • card_bytes(split, ctx) is the figure a plan reserves: the resident tensors in their device formats, the per-layer projection as f32, the two rope tables, the fixed scratch, and the rows each call uploads (the token rows, one pair of key bound rows for each distinct reach, the output rows).
    • Encoder::pass(ids, taps) runs the encoder once over ids. The token rows are decoded on the host and scaled by √hidden. Each layer runs three steps in order, each its own method (attention, glu, per_layer). The attention holds the q8_0 GEMMs, the q and k norms with their gains and the v norm without one, the whole NEOX turn, and the bidirectional flash: head 256 with a first key a row for the sliding layers, head 512 over every key for the global layers. The GELU GLU follows, then the per-layer input with the layer's output scale. The output norm and the output GEMM of every position come last, and the host takes their mean (Pass::pooled), summed in f64. An empty input, one past the load's context and an id past the vocabulary are refused by name, and so is an output value that is not finite, by its position. The key bounds are uploaded once a reach, not once a layer. With taps, it reads back the intermediate tensors the gate compares.
    • A sliding layer's reach is the shared encoder reader's EncoderSpec::reach, which the encoder and the gate both read. Each row's key bounds have one owner in this model's reader, hparams::key_bounds(reach, t, n). A host test holds key_bounds to llama.cpp's symmetric mask, |i - j| <= reach, for every length up to 12 and every reach up to 13.
    • Two card tests run in the crate (--include-ignored). hw_per_layer_projection_scale_near_epsilon builds a synthetic 64-wide model whose per-layer projection has a mean square near the RMS epsilon, and holds the inp_per_layer tap to an f64 host recomputation within 1e-5, so a missing 1/√hidden scale fails it. hw_scratch_census_reads_buffer_lengths holds the f32 scratch census to the buffers' lengths.
    • tools/check-arch.sh counts crates/gpu-gemma-embedding2/ as the architecture's own crate: a crate name spells the arch/gemma_embedding2/ directory's underscore as a hyphen. just check and just lint build the gates with the new gemma-embedding2 feature, and the host nightly list in tools/recipes.py leaves the new card crate out as it does the others.
  5. gpu-gates, refset, tools: gate_gemma_embedding2_e2e The end-to-end gate (just gate-gpu-gemma-embedding2-e2e), described under "What I ran". It opens the text set and its CPU twin through refset's text-gemma-embedding2 families and Emb2Set::open, which now name the gate as their consumer. It reads each case's ids with refset's read_ids, the parser Text::ids uses. The sha256 helper it calls moves from serve_client (behind deepseek41) to the crate root, and gate_glm5next_serve imports it from there. tools/gate-paths.tsv names the gates that reach these files.

    The recipe runs four steps in order: the encoder crate's tests with the card tests, the symmetric key mask test, the gate's host tests, and the gate on the card. The mask test runs before the tap bands, so a wrong key bound fails on its own. The gate also requires an l_out tap for each of the layers.

Code shape

  • No second copy. The encoder calls the common kernels and adds none. The mean over positions is Pass::pooled, the one mean on this path. The gate reads the sets through refset's families and reader, and its sha256 is the crate's one helper. The sliding reach is read from the shared EncoderSpec::reach, not computed again.
  • Model code holds only architecture facts. The new crate, the reader's CARD_CTX and key_bounds are this model's; reach stays the shared encoder reader's. A common owner gained arms of the existing kind (Program::GemmaEmbedding2Body with its AVAILABLE rows, its leg_arms and routed_rule arms, routed_legs.rs's family name and skip reason) and EncoderSpec's four bools above.
  • Dead code is gone. fault.rs's, oracle/mod.rs's and the reader module's "no program runs it" texts for gemma-embedding2 now name the encoder and its gate's sets.
  • No silent failure. Every refusal above names its cause, a fault word raised during a call ends it with Encoder::pass named, and a non-finite output row is an error. In the gate, a row is within its band only when its distance and the twin's distance are both finite.
  • Comments state what is true now.
  • No relaxed gate. tools/unsafe-ratchet.txt moves two rows, each with PIN(...) lines: gpu by one block (the new entry's call into norm_partial_append) and one fn (the launcher); and gpu-gemma-embedding2 by two blocks, gpu-gemma-embedding2 2 0 0 (the encoder's calls of the head-512 and the whole-turn launchers, each with its SAFETY line). The gate's TWIN and FLOOR are new; FLOOR carries a PIN(...) line with the distances it was set over.
  • cargo fmt and cargo clippy are clean on what I touched; the lint count did not go up. The numbers are under "What I ran".

Duplication found or removed (path:line): serve_client.rs's sha256 moves to the crate root, where the new gate and gate_glm5next_serve call it. Other copies are left as they are: gate_qwen3_serve.rs:3081, gate_clef_hidden.rs:239, gate_qwen35_media.rs:286, shared/qwen3moe_taps.rs:75 and, through hf::fetch::digest, gate_serve_start.rs:367. head_norm_neox_append_256_whole is the 512 entry's body at another head width, through the same generic core.

What I ran

The branch is now rebased onto main a688accb as head 3e8da8dc. The rebase changed no line of the encoder, its kernel entry, gate or tests; in the first three commits only the match arms beside main's new Bolt and nemotron arms in accepts.rs, coverage.rs and routed_legs.rs moved. At 3e8da8dc, cargo fmt --all -- --check, the host tests, just gate-gpu-gemma-embedding2-head512 and the Mac mac-static ran again with the same results as below (the Mac lint step now counts 59 warnings, as main a688accb does on the same Mac). Every other result below is from head f001e25 on base d66ef09. The card runs used a DGX Spark and the host checks a second one of the same type, both in the project's dev container (bloomery-dev:cuda13.3.1, image sha256:853c20d59e68), CUDA 13.3.1, ptxas 13.3.73. A check that is red on this head and on the base alike says so.

Check On this head On the base
just gate-gpu-gemma-embedding2-e2e exit 0, 88 of 88 cases PASS no such gate
just gate-gpu-gemma-embedding2-head512 exit 0 no such gate
cargo fmt --all -- --check rc 0 not run
cargo clippy on main's just lint feature set rc 0, 39 warning lines rc 0, the same 39 lines
lint step of ./tools/mac-static.sh upstream/main (Mac), with gemma-embedding2 rc 0, 41 warnings, the ratchet's 41 not run
host tests (model, models, tokenizer, placement, levers, refset, gates library with and without gemma-embedding2) all rc 0 not run
gemma_embedding2_meta with the model file pins_differ=0 not run
just contrib-check --tier t0 84 checks, 1 red: check-recipes the same 1 red
the 11 tools/check-*.sh scripts all rc 0 not run
just affected gates 182 gates, 37 exit 0, 145 exit non-zero the same 145 exit non-zero with the same codes
ptx-compare against the base, 7 binaries the new entry added in 6, the 7th identical not applicable
just gate-ptx-spill (cc 12.1) violations=11 FAIL in each core binary the same 11
./tools/mac-static.sh upstream/main, other steps 1 red: check-recipes the same 1 red (main fc47ee6)

The sections below give the command and output behind each row.

The end-to-end gate. The gate checks the file at refset's MODEL path against its pinned sha256 and opens the encoder at CARD_CTX (8,192) positions. It runs every text case of mainline's card set, with each case's ids read from the tokenizer oracle's gemma4 set (ref-tokenizer-gemma4) and held to the count and digest the set's # case line records. It first checks that the CPU twin's # case lines are the card set's. A tap's distance from the card set is the relative L2 distance, and its band is 2.5 times the CPU twin's own distance, and at least 1e-2.

On a DGX Spark, against the sets mainline dumped on that card type (NVIDIA GB10, 595.84, 12.1, the model file below):

$ BLOOMERY_ARCH=sm_121 BLOOMERY_BOX_LOCAL=1 just gate-gpu-gemma-embedding2-e2e
test tests::hw_scratch_census_reads_buffer_lengths ... ok
test pass::tests::hw_per_layer_projection_scale_near_epsilon ... ok
test result: ok. 5 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.36s
test arch::gemma_embedding2::spec::tests::key_bounds_are_the_symmetric_mask ... ok
test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 378 filtered out; finished in 0.00s
test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s
gate_gemma_embedding2_e2e: device NVIDIA GB10: /models/embeddinggemma-2/embeddinggemma-2-Q8_0.gguf at 8192 positions, card figure 911742976 B; 88 cases; band 2.5 x the CPU twin, at least 1e-2
(t) the window: the longest case holds 1734 ids, a sliding layer reaches 512 on each side PASS
(t) cases/66: 1734 ids, inp_scaled bits PASS; inp_per_layer 1.589e-3 (twin 2.410e-7, band 1.000e-2); l_out worst layer 23 2.277e-2 (twin 2.118e-2, band 5.294e-2); result_norm 6.554e-2 (twin 6.137e-2, band 1.534e-1); embd 6.445e-2 (twin 5.995e-2, band 1.499e-1); pooled 1.047e-2 (twin 1.022e-2, band 2.556e-2); 375.4 ms (runtime value) PASS
(d) cases/66: 1734 ids, a rerun's rows bit for bit PASS
(r) no ids: refused by name PASS
(r) one past the context: refused by name PASS
(r) an id past the vocabulary: refused by name PASS
(r) cases/66: after the refusals the rows bit for bit PASS
(r) a load at 8193 positions: refused by name PASS
gpu-gate.sh: gate_gemma_embedding2_e2e ran 4 s (exit 0)

The first test result is the encoder crate's five tests, the two card tests among them. The second is the symmetric mask test. The third is the gate's two host tests, the CPU case identity test and the window case test. All 88 case lines end PASS; one is shown above. In every case inp_scaled is bit for bit the set's.

The band's measured distances. On the DGX Spark, inp_per_layer is 1.283e-3 to 1.594e-3 over the 88 cases, against twins of 1.047e-7 to 2.410e-7. So FLOOR sets its band in every case. Four cases have their worst layer on FLOOR too: cases/02, cases/34, cases/59 and cases/60, whose worst is l_out 9.046e-3 at layer 0.

FAIL-first. Thirteen temporary edits, each run through the gate's recipe on the DGX Spark and reverted, and the recipe run again (rc 0 after every one):

Edit What stops it rc
the v norm skipped cases/01 layer 0 l_out 2.276e-1 past its band 1.061e-2 1
the attention's post norm skipped before the residual add cases/01 layer 0 l_out 6.892e-1 past its band 1
key_bounds gives every key on a sliding layer key_bounds_are_the_symmetric_mask 101
the per-layer input's 1/√hidden scale dropped hw_per_layer_projection_scale_near_epsilon, row 0 of layer 0 at 0.9919 against 0.6987 101
the token rows not scaled by √hidden the same test, its mean square at 1.490e-8, below the band around the epsilon 101
the vocabulary check dropped (r) an id past the vocabulary: refused by name FAIL 1
every layer's output scale set to 1 cases/01 layer 0 l_out 1.560e0 past its band 1
key_bounds one key short on the right key_bounds_are_the_symmetric_mask 101
the card's key counts one short on a sliding layer cases/01 layer 0 l_out 3.356e-1 past its band 1
the CPU twin's case lines not compared a_cpu_case_of_another_digest_is_refused_by_name 101
ValueNorm dropped from the needs coverage_names_the_encoder_parts 101
the coverage refusal in Encoder::open dropped a_shape_outside_the_program_is_refused_by_name 101
the context bound n > ctx made n >= ctx the same near-epsilon test, which passes a 2-position input at a context of 2 101

Each edit failed where the table says, and each revert passed. The ValueNorm edit ran bloomery-model's test coverage_names_the_encoder_parts alone, because the gate's recipe does not run it.

Static and host checks. On a DGX Spark:

  • cargo fmt --all -- --check is rc 0.
  • cargo clippy with the feature set of main's just lint is rc 0 at the head and at the base, with warning_lines=39 and the same warning set at both. That set does not include the gates' gemma-embedding2 feature, which just lint now builds.
  • The host test runs of bloomery-model, bloomery-models, bloomery-tokenizer, bloomery-placement, bloomery-levers, bloomery-refset and the gates' library with and without gemma-embedding2 all pass: unit-tests aggregate rc=0. The model's library alone is test result: ok. 375 passed; 0 failed; 4 ignored.
  • gemma_embedding2_meta with the model file (--ignored) passes with pins_differ=0.
  • The 11 tools/check-*.sh scripts are rc 0: check-scripts aggregate rc=0.
  • just contrib-check --tier t0 ends DONE tier=t0 total=84 red=1 wall=321s red: check-recipes. The one red is --v41-queue runs under BLOOMERY_BOX_READONLY=1 and syncs nothing: rc 1 (want 0). The base gives DONE tier=t0 total=84 red=1 wall=320s red: check-recipes with the same line.

On a Mac, ./tools/mac-static.sh upstream/main runs the just lint recipe's command, with the gemma-embedding2 feature: the lint step is rc 0 and reads ^warning: 41 = the ratchet 41. The run's one red step is red steps: check-recipes (rc 1), from the ptx-compare self-test (33 cases, 1 failed). Main fc47ee62 on the same Mac gives the same single red step and the same self-test failure.

The kernel and PTX checks.

  • just gate-gpu-gemma-embedding2-head512 on a DGX Spark: gate_gemma_embedding2_head512 ran 8 s (exit 0).
  • ptx-compare of the new entry against main: in six of the seven binaries the new entry shows added 1: head_norm_neox_append_256_whole, removed none and changed none. The seventh, gate_qwen3vl_tower, is identical. No existing entry changed.
  • ptxas -arch=sm_121 -v on the scanned module gives head_norm_neox_append_256_whole 40 registers with no stack, spill stores or spill loads, beside head_norm_neox_append_512 at 38.
  • just gate-ptx-spill on a DGX Spark (cc 12.1) gives violations=11 FAIL for generate_ds41 and for gate_e2e, at the head and at the base. The 11 are the same at both and are not this PR's: embed_rows_q5_0, gemm_q5_0p, gqa_flash_seg_g5, iq2_xxs_gemv_sel, iq2_xxs_rows, mamba2_norm_rows, mamba2_ring_tail, mamba2_step, mamba2_steps, nvfp4_gemv and relu_sqr, unpinned in sm_121-cc12.1.tsv on main, each with spill=0 jit_local=0. The head has one more entry and one more pinned row than the base in each (entries=317 pinned=306 against entries=316 pinned=305 for generate_ds41; entries=250 pinned=239 against entries=249 pinned=238 for gate_e2e): the new entry, pinned. The vision binary is entries=16 pinned=16 nonzero=none violations=0 PASS.

The just affected gates. The recipe selects 182 gates for this change. gates-affected: 182 gates, ok 37, same as base 145, differs 0, unproven 0. The 37 that exit 0 include gate-gpu-gemma-embedding2-e2e, gate-gpu-gemma-embedding2-head512 and gate-gemma-embedding2-meta. The other 145 exit non-zero on this DGX Spark with the same code on the base (1 for 92, 101 for 33, 64 for 18 and 66 for 2), so they show nothing about this head and count as not run.

The cc 8.6 rows. No cc 8.6 card was at hand, so the new entry's rows in ptx-shapes.tsv and sm_86-cc12.1.tsv are ptxas 13.3.73's sm_86 frame.

Not run here.

  • A cc 8.6 card (above).
  • The gate on an RTX 5090 machine at this head.
  • just contrib-check --tier t1, which needs an x86_64 Linux host.
  • The 145 just affected gates listed above.

What I ran into

None of this blocks this pull request.

  • inp_per_layer and the shortest cases' l_out sit on FLOOR, not on the twin rule. The per-layer input's distance from the card set is 5.9e3 to 1.4e4 times the CPU twin's in every case [derived], so FLOOR = 1e-2 sets its band. The suspected cause is the reference's per_layer_model_proj, which is BF16 in the Q8_0 file: llama.cpp rounds the input of that matrix product to BF16, where the encoder decodes the matrix to f32 on the host and keeps the input in f32 [inferred, not run]. A check would round inp_scaled to BF16 once on the host before the f32 tile and see whether inp_per_layer falls to the twin's distance. If it does, the options are to round the same way, so that the twin rule holds that tap, or to keep f32 and keep FLOOR.
  • The ids set has no refset family. The gate opens ref-tokenizer-gemma4 by name, reads its ids with read_ids and holds each file to the digest the text set records. It does not check that set's MANIFEST trailer, build or vocabulary. I left out a family row because the set's MANIFEST records different md5 sums for llama-tokenize and libllama on a DGX Spark and on an RTX 5090 machine at the same mainline revision 6868a52cb: be4957f2c2f96704054e0557c48b9b09 against 2de0937303f64c46c35959c91b54a500, and 10eff38604798775e975958c13b3b385 against 044b3a8b828a9f31508ba850f9846169. A family pins one build, so one row would refuse one of the two machines' sets. Our default is the digest check above; a family can come once the build pin has a rule for per-machine binaries.
  • The helper scan of tools/ptx-spill-check.sh asks for an RTX 3090. Run directly on a DGX Spark, its JIT half stops with the driver JIT (oxart_jit) failed: gpu-gate.sh: the 3090 lookup failed at the head and at the base, for all three binaries, so it is not a kernel result. The recipe just gate-ptx-spill above is the cc 12.1 result.

Hardware (if the change is for a machine)

  • Two DGX Sparks (aarch64), the same type: NVIDIA GB10, 595.84, 12.1 from nvidia-smi, the project's dev container bloomery-dev:cuda13.3.1 (image sha256:853c20d59e68), CUDA 13.3.1, ptxas 13.3.73.
  • The model file: ggml-org/embeddinggemma-2-GGUF embeddinggemma-2-Q8_0.gguf (sha256 2188ac1d...), the file the reference sets were dumped from.

AI help

AI 코딩 에이전트가 코드를 작성한 후 설명을 썼고, 테스트를 마쳤습니다. diff 리뷰와 PR 초안 리뷰는 제가 직접 했습니다.

@github-actions github-actions Bot added gpu The card engine: crates/gpu*, crates/app host Host crates: qdot, threads, engram, engram-lab, runtime, sampler, levers, decision, vision, refset model The CPU engine and model files: crates/model, models, hf, gguf, tokenizer, jinja needs-gpu Changes code a card must check, and carries no gate run on a card placement Placement planning: crates/placement tools Tools, recipes, CI and the toolchain: tools/, justfile, .github/, Cargo labels Oct 15, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: MUST_FIX: Remove the single-use tensor wrapper and its unused metadata result.

Must fix

1. Inline the single-use tensor lookup

crates/gpu-gemma-embedding2/src/lib.rs:504 introduces tensor for one caller, output_scale, which discards its TensorInfo result at line 567. This is a kind B unnecessary wrapper under the fewest-elements rule.
With one caller, removing the definition and unused type plumbing needs no API migration. Delete tensor, remove TensorInfo from the import, and replace output_scale with:

fn output_scale(split: &Split, l: usize) -> Result<f32, GpuError> {
    let name = format!("blk.{l}.{}", names::LAYER_OUTPUT_SCALE);
    let missing = || GpuError::Tensor {
        what: OPEN,
        name: name.clone(),
        need: "a tensor of the file",
    };
    let (shard, t) = split.find(&name).ok_or_else(missing)?;
    let data = split.shard(shard).ok_or_else(missing)?.data(t)?;
    let b: [u8; 4] = data
        .get(..4)
        .and_then(|b| b.try_into().ok())
        .ok_or_else(|| GpuError::Shape {
            what: OPEN,
            detail: format!("{name} holds {} bytes, not one f32", data.len()),
        })?;
    Ok(f32::from_le_bytes(b))
}

Non-blocking

⚠ Accept the exact context limit: crates/gpu-gates/src/bin/gate_gemma_embedding2_e2e.rs:254 checks the rejected side only. Changing n > self.ctx to n >= self.ctx at pass.rs:402 survives the reported cases and refusals.
Add enc.pass(&vec![c.ids[0]; enc.ctx()], false)?; before the existing rerun. This catches the off-by-one using the current clause.

Optimization headroom

  • G, Per-launch and per-sync fixed cost: crates/gpu-gemma-embedding2/src/pass.rs:479: reuse the idle feed-forward scratch for output, reading back only the live output prefix. Its width is 2048; the output needs 768, and the projection accepts a larger buffer. This removes one device allocation per input and 3,072 zeroed bytes per token [derived: 4 × 768], without another buffer field.

Checked

Duplicates: 79 added definitions searched; no duplicate owner found. The one dryer pair consists of thin entries sharing norm_partial_append.
Complexity: new fallible pipelines are pass 26 (inputs/layers/output), open 23 (load/constants/census), attention 24 (projections/rope/flash), per_layer 18 (projection/taps/residual), and gate run 20 (load/cases/refusals). Scratch::new 24 constructs buffers. needs 41→48 adds encoder branches; Need::fmt 74→80 adds harmless match arms.
Silent failure: five fallback/conversion sites followed. Positive layer counts, validated tensor shapes and nonempty cases block the defaults; failed conversions become named errors. Device faults and non-finite output reach errors.
Tests, source mutation audit: check_program at encoder lib.rs:253 returning success fails the shape test; dropping RoPE at line 244 fails the byte figure; wrapping arithmetic at lines 60/66 fails overflow assertions; omitting x at line 351 fails the census. A key edge shortened at hparams.rs:279, dropped ValueNorm at need.rs:430, skipped case equality at gate line 82, and a missing program in program_of fail the mask, needs, identity and metadata assertions respectively. The PR reports 11/12 encoder mutants caught; scale removal survives. The context-boundary mutant above also survives.
Gates: new FLOOR and increased unsafe counts have dated PIN reasons; no existing numeric band was relaxed.
Escape catalog: C1–C4 checked; none matched the changed shell code.

Notes

GPU builds and gates were not rerun in this review. The PR carries card evidence and gate output.

Card check: gpu-evidence. This pull request changes crates/gpu-gates/Cargo.toml, crates/gpu-gates/src/bin/gate_gemma_embedding2_e2e.rs, crates/gpu-gates/src/bin/gate_glm5next_serve.rs … and carries a gate run on the author's card, which counts as verified.

🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)

@github-actions github-actions Bot added gpu-evidence Verified on the author's card: the pull request carries the gate run and its output and removed needs-gpu Changes code a card must check, and carries no gate run on a card labels Oct 15, 2026
… over plain query rows at head 256 (EmbeddingGemma 2's sliding layers)

The sliding layers turn a head of 256 values whole (NEOX pairs
(i, i + 128)) and read their queries from plain query rows, as the
global layers do at head 512. The entry is norm_partial_append at
HEAD = ROT = 256 over the f16 planes, 128 threads a head, and its enqueue
takes the 512 entry's PartialNeoxArgs and refuses a short buffer by name
through partial_launch. Its rows are pinned in the three ptx-shapes
tables at spill 0 and jit_local 0; the cc 8.6 rows' jit_local is
ptxas's sm_86 stack frame.
The header's context_length (262144) is not the input length the model
was trained to: the model card names 8,192. The file does not carry that
fact, so the reader holds it as one cited constant beside the keys it
reads, and a consumer that bounds an input reads it from here instead of
writing the number.
… are needs, and its program is the card encoder

The encoder carries four parts a decoder's layer does not: the token
rows' scale, the gainless value norm, the norms after both blocks and the
layer's trained scalar. EncoderSpec names them, needs() lists them with
the per-layer input and the pooled output, and the coverage rows of
Program::GemmaEmbedding2Body cover the shapes the encoder runs: the head
256 flash a first key a row, the head 512 flash over every key, both
whole NEOX turns, the GELU GLU at 2048, the per-layer input at 512 and
the 768-wide output. Its matrices are pinned to q8_0, its token table and
per-layer projection to the host's decoded rows.
…loomery-gpu's kernels

One pass over an input's ids in llama.cpp's graph order, with no device
code of its own: the token rows decoded on the host and scaled by the
root of the hidden width; per layer the q8_0 GEMMs over a one-expert
dense route, the q and k norms with their gains and the gainless v norm,
the whole NEOX turn and the bidirectional flash (head 256 with a first
key a row from the layer's reach, head 512 over every key), the GELU GLU,
the norms after both blocks, the per-layer input (its projection read as
f32 on the card, scaled and normed, gating its own block) and the
layer's scalar; then the output norm, the output GEMM of every position
and their mean on the host.

card_bytes is the figure a plan reserves and the one the open holds
itself to. A load's context is at most the reader's CARD_CTX (the model
card's input length) and the file's context_length.

check-arch counts the crate as the architecture's own: a crate name
spells the arch directory's underscores as hyphens.
…line's text set and its CPU twin

Every text case of the card set through one pass at CARD_CTX positions:
inp_scaled bit for bit, the other taps, the rows and their mean within
TWIN times the CPU twin's distance and at least FLOOR; the longest case
rerun bit for bit; a pass of no ids, past the load and past the
vocabulary, and a load past CARD_CTX, each refused by name. The
gemma-embedding2 feature turns the encoder on in gpu-gates for this bin.
…ner, and its coverage rows name their owners

Hparams::reach(l) is a sliding layer's reach, half the window, and
hparams::key_bounds(reach, t, n) is row t's keys among n positions: every
key j with |t - j| <= reach, as llama.cpp's symmetric mask keeps them. The
spec reads reach; a host test holds key_bounds to the mask at every length
up to 12 and every reach up to 13, so a bound one key short or long on
either side is red. The pooled row's coverage names gpu-gemma-embedding2's
mean, the per-layer projection's rule reads names::PER_LAYER_MODEL_PROJ,
and the metadata gate opens refset's MODEL, the file the sets are dumped
from. The routed-legs gate skips a gemma-embedding2 set for no routed
expert.
…fe, and a pass reads its key bounds once a reach and refuses a non-finite row

enqueue_head_norm_neox_append_256_whole is an unsafe fn: its Safety
section names the distinct positions the caller keeps, since two tokens at
one position write the same cache rows from two blocks. Encoder::pass
vouches for it with Consts::pos, 0..ctx.

Encoder::pass runs each layer as three steps (the attention, the GELU GLU,
the per-layer input and the layer's output scale), each its own method. It
uploads one pair of key bound rows for each distinct reach, from
key_bounds and narrowed through try_from, and card_bytes counts that many
pairs and the output rows, which each pass now holds at its own n
positions and reads back whole. An output value that is not finite is an
error naming its position. The ratchet moves gpu by one fn and
gpu-gemma-embedding2 by one block, each with its PIN line; check-arch's
crate-name line is in the file's language.
serve_client sits behind deepseek41, so a gate without that feature could
not call its file digest. The helper is unchanged, and gate_glm5next_serve
imports it from the root.
…rough refset's family and reader, holds the window and both distances finite

The gate opens the card set and its CPU twin through
refset::arch::gemma_embedding2's TEXT and TEXT_CPU families and
Emb2Set::open (the set's kind, device and the model's path and sha256),
in place of its own copy of the family. It runs the file at refset's
MODEL, held to MODEL_SHA256 before the card, so --model goes; the case ids
come from the tokenizer oracle's gemma4 set, as the dumper reads them.
bloomery_levers::at_main runs first.

A row is within its band only when its distance and the twin's are both
finite: rel makes a NaN or a length mismatch infinite, and an infinite
twin made an infinite band. Every l_out line prints its twin distance. A
set with no case longer than a sliding layer's reach plus one fails by
name, since no row's keys would stop short of the input's ends. FLOOR's
comment is a PIN line with the distances it was set over.
github-actions[bot]
github-actions Bot previously approved these changes Oct 16, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: CLEAN: No must-fix findings remain in the validated review.

Checked

Duplicates: The SHA helper moves from serve_client to the gpu-gates library. The new head-256 kernel entry reuses norm_partial_append. Dryer reports zero pairs in its searched crates; gpu and gpu-gates were excluded from that scan.
Complexity: Diagnostics flag the new encoder and gate pipelines. needs adds encoder branches; Need::fmt adds match arms.
Silent failure: Load contexts and pass lengths are checked. Invalid token IDs, conversion failures, device faults, and non-finite output rows return errors. Missing window coverage and failed cases keep the gate red.
Tests: Assertions cover projection scaling, symmetric key bounds, scratch byte counts, coverage, encoder needs, CPU case identities, and window boundaries. The PR reports FAIL-first mutation runs; these were not rerun during validation.
Gates: The new numeric floor and unsafe-ratchet increases carry dated PIN comments and reasons. The PR reports unchanged existing PTX entries.
Adversarial checks: The gate checks empty and oversized calls, invalid IDs, repeated calls, and a shorter call after a full-context call. Encoder::pass requires exclusive mutable access.
Escape catalog: No C1–C4 pattern matches the changed script code.

Notes

Validation was static; repository tests and CUDA gates were not rerun. The PR supplies gate output, both commit hashes, the model path, and NVIDIA GB10 card evidence. Its cc 8.6 evidence is a compiler resource scan.

Card check: gpu-evidence. This pull request changes crates/gpu-gates/Cargo.toml, crates/gpu-gates/src/bin/gate_gemma_embedding2_e2e.rs, crates/gpu-gates/src/bin/gate_glm5next_serve.rs … and carries a gate run on the author's card, which counts as verified.

🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: CLEAN: Input bounds, device faults, and non-finite results reach named errors. No must-fix defect was found.

Optimization headroom

  • G. Per-launch and per-sync fixed cost: crates/gpu-gemma-embedding2/src/pass.rs:336. Explore scaling the per-layer projection weights during loading and removing this separate scale launch. The opportunity is 24 launches per input, or 24/n per token [derived, 24 layers and n input positions]. Validate the changed rounding with the near-epsilon test and end-to-end gate, then time an A/B.

Checked

Duplicates: 81 added or re-exposed item declarations checked against existing owners; no new duplicate found. Dryer reported zero pairs. The normalization entry reuses its generic core; hashing and ID parsing reuse existing implementations.

Complexity: new Encoder::pass (26), Scratch::new (24), Step::attention (24), Encoder::open (23), gate::run (22), and Step::per_layer (18) are load or launch sequences with error propagation; attention, GLU, and per-layer stages are separate. needs (+7) adds encoder branches; Need::fmt (+6) adds match arms.

Silent failure: five default/conversion sites followed. Header guards block empty geometry; failed conversions return named errors; an insufficient window case fails the gate. Optional taps and failed-case reporting do not hide errors.

Tests: mutation review found assertions against dropped norms/scales/coverage/identity checks, key-edge and context off-by-ones, omitted scratch, and wrapping byte counts. The PR reports 13 FAIL-first runs. Independent mask and scratch arithmetic checks passed during review.

Gates: new FLOOR and unsafe-counter pins carry dated reasons. No existing numeric band was relaxed. The feature-gated binary declares required-features.

Escape catalog: no added code matches C1–C4. Boundary scenarios checked include a full-context pass, a shorter pass after it, invalid IDs, unsupported shapes, and non-finite outputs.

Notes

Dryer did not scan crates/gpu or crates/gpu-gates; both received manual owner searches.

Build and GPU tests were not rerun during this source review. The supplied GB10 end-to-end output predates the rebase; the PR reports the head-512 gate and host checks rerun at 3e8da8d.

Card check: gpu-evidence. This pull request changes crates/gpu-gates/Cargo.toml, crates/gpu-gates/src/bin/gate_gemma_embedding2_e2e.rs, crates/gpu-gates/src/bin/gate_glm5next_serve.rs … and carries a gate run on the author's card, which counts as verified.

🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)

@juliankang4

Copy link
Copy Markdown
Collaborator Author

The review is done and the contribution check passed, so I'm merging this.

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

Labels

gpu The card engine: crates/gpu*, crates/app gpu-evidence Verified on the author's card: the pull request carries the gate run and its output gpu-verified host Host crates: qdot, threads, engram, engram-lab, runtime, sampler, levers, decision, vision, refset model The CPU engine and model files: crates/model, models, hf, gguf, tokenizer, jinja placement Placement planning: crates/placement tools Tools, recipes, CI and the toolchain: tools/, justfile, .github/, Cargo

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants