Skip to content

dsv41 metal: the KV compressor and Engram - #1713

Closed
gilby wants to merge 2 commits into
ashhart:mainfrom
gilby:dsv41-metal-pr/09-compressor-engram
Closed

gilby wants to merge 2 commits into
ashhart:mainfrom
gilby:dsv41-metal-pr/09-compressor-engram

Conversation

@gilby

@gilby gilby commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

The KV compressor (compress.metal) and Engram (engram.metal: FP8 row dequantisation, projection, gate, residual add) on Metal, registered with the kernel contract.

Main's shared files change only by one-line hooks; this part's records, suites and limits go into the family's own files (dsv41_metal_records.zig, metal_dsv41_main.zig, metal_dsv41_limits.zig).

Part 9 of 18 of the stack for #1676; stacks on dsv41-metal-pr/08-indexer-topk-route (part 8, #1712), on main with parts 1 to 5 (#1705 to #1709) merged.

Receipt

  • Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision: TensorFold 1a2c89a on main ad98b60 (native Zig 0.17.0; no MLX or PyTorch); Apple M3 Ultra (60-core GPU), 96 GB, macOS, for the kernel and device tests; checkpoint deepseek-ai/DeepSeek-V4.1-Flash-1017.
  • Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false: n/a: kernels only; the kernel contract checks each launch geometry's bits against the others and against float64.
  • Decode speed before and after, from tools/bench_openai.py: measured on the PR that turns serving on.
  • Prompt speed before and after: the cold first request's ttft_s from tools/bench_cache.py --require-cold on a long prompt, or "n/a" with the reason when the change cannot reach prompt processing: measured on the PR that turns serving on.
  • Tests run, with pass, fail and skip counts: zig build test tf-kernel-contract -Dcpu=apple_m1: exit 0, 621/621 steps succeeded; kernel contract, tf-kernel-contract --only dsv41 (Metal): SUMMARY 65 kernels: 65 pass gates 1-4, 0 fail, 0 known findings; 129 controls: 129 red as planted, 0 missed; 0 limit problems; gates 118459 ms, limits 157 ms; python3 tools/zig/lean_check.py: 5068 files, 424 problems, the same count as main (5027 files, 424). Zig's summary counts only the test binaries it ran in that invocation, so the step count is the comparable figure.
  • Not run, and why: CUDA: no CUDA code changes; Metal on M1, M2, M4 and M5: one M3 Ultra only (the contract's limits check covers the M1 to M5 pipeline limits).

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
  • Prompt processing is no slower than the last release.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thanks @gilby. Rebased and clean: 58 kernels and 107 controls red, with the Engram path checked against its own float64 reference. This lands after #1712.

@gilby
gilby force-pushed the dsv41-metal-pr/09-compressor-engram branch 2 times, most recently from a31c325 to 48b5050 Compare November 2, 2026 07:49
@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @gilby: nothing owed from you here. Every kernel in this part is gated against float64 with planted controls, and the push is a clean rebase. We're running tf-kernel-contract --only dsv41 on our M5 Ultra over the whole stack, then it lands in stack order after #1710 (#1712 has one small ask for route_select controls).

@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Our Metal kernel contract over the stack (tf-kernel-contract --only dsv41, head #1713, M5 Ultra): 65 kernels pass gates 1-4, 0 fail, 0 known findings; 121 controls, 121 red as planted, 0 missed; 0 limit problems. That clears the device side for #1711 to #1713. They land in order after #1710, with #1712's one ask (route_select controls) first.

gilby added 2 commits November 3, 2026 10:35
index.metal scores FP4 queries against FP4 keys in one stated order with causal and block masks, topk512.metal
keeps each row's top 512, cand.metal picks the candidate blocks as a byte mask, and route.metal takes the top-k of
sqrt(softplus(logit)) + bias for the backbone and the drafter. Each is registered with the contract.
compress.metal pools the ratio-2 and ratio-4 compressed KV and round-trips it through FP4; engram.metal dequantises
the FP8 Engram rows, projects them, gates them and adds them to the residual. Registered with the contract against
float64 and bit-equal to the phase-1 glue where they replace it.
@gilby
gilby force-pushed the dsv41-metal-pr/09-compressor-engram branch from 1a2c89a to d9cf9de Compare November 3, 2026 09:39
@ashhart

ashhart commented Nov 3, 2026

Copy link
Copy Markdown
Owner

Thanks @gilby. This part landed on main with #1719's merge (cfcbecb), which carries #1712-#1719 as your own commits in order; closing it as included. The same Metal contract, full test and DeepSeek device suite cleared the whole stack.

@ashhart ashhart closed this Nov 3, 2026
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.

2 participants