Skip to content

Serving: read cached DeltaNet recurrent states by bank index - #258

Open
DorBernsohn wants to merge 2 commits into
jaredpalmer:mainfrom
DorBernsohn:perf/indexed-deltanet-state-bank
Open

DorBernsohn wants to merge 2 commits into
jaredpalmer:mainfrom
DorBernsohn:perf/indexed-deltanet-state-bank

Conversation

@DorBernsohn

@DorBernsohn DorBernsohn commented Oct 24, 2026 •

Copy link
Copy Markdown

Each question row currently gathers the document's fp32 DeltaNet recurrent state with index_select, then copies it into a row buffer before running the scan. Questions sharing a document repeat those copies. This reads the recurrent initial state directly from the CUDA bank using each row's bank-entry index instead.

The indexed kernel is the FLA 0.5.2 state scan with its initial-state address changed. The index stride is explicit because entries are a column of the graph upload buffer. Question intermediates remain row-local, and final-state writes stay disabled, so repeated indices cannot advance or overwrite the document's cached state. State prefill and unfused paths keep their existing calculation. Convolution and attention gathers remain. The vendored kernel retains its MIT license, including in built wheels.

Addresses the indexed DeltaNet state-bank item in #107. This does not close the umbrella issue.

Measurements and limits

The baseline was measured before implementation at 2074cc44256200bb37540103fa4572901faa10f3. The final measured core sources match 8e339ae34979678fc154366f1d1f06a50cb49e2b by SHA-256; the later commit adds reports, the throughput investigation helper, and license packaging metadata.

H100 80GB, driver 610.57.04, CUDA 12.8, Torch 2.8.0, Transformers 5.17.0, FLA/fla-core 0.5.2, Triton 3.8.0, causal-conv1d 1.7.0. Checkpoint jaredpalmer/kev-4b@6cfce5c2fa4b4bd64026336ab649c5ca78857d52, merged bf16, fused layers and CUDA graphs. Dependency requirements and uv.lock are unchanged.

The controlled comparison runs both paths on the same loaded weights, encoded inputs and physical GPU. It warms three calls, captures the known graph shapes, then measures 15 calls with synchronized CUDA events and wall clocks. It covers shared and distinct documents, padded counts, longer branches, and cold/mixed/warm prefixes. Two exploratory trials reverse measurement order, followed by confirmation on the final committed core sources. Capture/warmup costs are recorded separately.

Final controlled warm batch, 16 documents × 6 questions Gathered Indexed
Median GPU-event model latency 106.050 ms 91.979 ms
Wall-time throughput 150.8 requests/s 173.9 requests/s
Maximum absolute probability difference — 0
Answer flips — 0/504 across all seven cases

That is a 13.3% latency reduction for this fixed batch. The two earlier trials measured 12.9% and 12.5%; longer branches reduced the gain to about 4.5%. The isolated original recurrent gather/copy measured 5.367 ms over 24 layers at 96 rows and sixteen bank entries. Removing that operation still leaves the scan's initial-state read, so its isolated timing is not a predicted speedup.

Mixed serving throughput is inconclusive. The original full serving benchmark improves six-question short-state traffic at 64 clients from 112.8 to 120.1 requests/s, but mixed decision-v7 falls from 128.3 to 121.2. A within-container ABBA investigation gives gathered 122.12/120.01 and indexed 120.07/130.86 requests/s at 64 clients. New graph shapes remain during those second passes. The separate-container regression does not repeat consistently, and this does not establish a universal serving-throughput gain. Both favorable and unfavorable results are summarized in the report. Only the final controlled benchmark JSON is committed; it contains timing samples, parity results, package versions and the three changed core modules' source hashes. Workspace URLs and run identifiers are excluded.

Row buffers are still allocated. Peak-memory differences follow measurement order, so no memory-saving claim is made. Performance was measured only on Kev-4B/H100; Kev-0.8B has the existing weight-backed parity test, and 9B/27B and other GPUs remain unmeasured. The copied scan must be checked when the FLA pin changes.

Full report, raw samples, environment, exact source hashes, and reproduction commands. The minimal launcher uses modal_app.py's image/functions without registering unrelated secret-dependent mirroring functions. Example:

uvx --python 3.13 --from 'modal==1.5.5' --with 'cbor2<6' modal run runs/indexed-delta-4b/launch.py --kind compare --name indexed-compare-NEW --flags '--reps 15'

The correctness criterion was set before candidate results: bit-identical probabilities against the gathered serving calculation, zero flips, and an unchanged bank after question reads. The full serving benchmark's reference and isolation blocks also match the original baseline exactly. Their existing fused/bf16 versus fp32 differences are not presented as candidate-versus-baseline differences.

Test plan

  • Required fast suites on Linux CPU: 527 passed, 26 skipped. Full tracked fixtures were mounted; optional CUDA-only FLA/convolution packages were removed only from that disposable CPU test container.
  • New CUDA tests on H100: 8 passed. fp32/bf16 inputs, lengths 1/17/64/129, row counts 1/3/8/17, repeated/strided/changed indices across graph replays, and unchanged bank contents.
  • Existing tests/test_model.py::test_cuda_graphs_match_eager on H100: 1 passed, pinned Kev-0.8B. Covers state prefill, shared prefixes, cache hits/misses, large/padded row batches, long states and eager fallback.
  • Controlled model parity: all 504 question outputs bit-identical in each of three comparison runs, zero flips.
  • Original serving benchmark repeated before/after with --n 8 --reps 5 --isolation; all captures succeed and reference/isolation blocks match.
  • Final source-convention checks: 20 passed. Structural review completed; canonical device, environment and JSON helpers are reused.
  • uv lock --check, Python compilation, and git diff --check pass. Wheel builds and includes both the indexed kernel and its MIT license.

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