Repository navigation
Serving: read cached DeltaNet recurrent states by bank index - #258
Open
DorBernsohn wants to merge 2 commits into
Open
DorBernsohn wants to merge 2 commits into
DorBernsohn wants to merge 2 commits into
Conversation
DorBernsohn
force-pushed
the
perf/indexed-deltanet-state-bank
branch
from
October 24, 2026 11:33
1de3f4a to
a4d6100
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 match8e339ae34979678fc154366f1d1f06a50cb49e2bby 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 anduv.lockare 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.
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: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
tests/test_model.py::test_cuda_graphs_match_eageron H100: 1 passed, pinned Kev-0.8B. Covers state prefill, shared prefixes, cache hits/misses, large/padded row batches, long states and eager fallback.--n 8 --reps 5 --isolation; all captures succeed and reference/isolation blocks match.uv lock --check, Python compilation, andgit diff --checkpass. Wheel builds and includes both the indexed kernel and its MIT license.