JCNTH opened a new pull request, #20428: URL: https://github.com/apache/tvm/pull/20428
**Stacked on the compute-pass PR (#20427); only the last commit is this PR. Rebased onto `main` once that one lands.** Follow-up to the review of #18871, which asked for bind group caching to live in `CacheState` with the key computation and invalidation spelled out. During LLM decode the memory planner hands each kernel the same buffers every token, so nearly every launch binds a combination that was bound before; creating a bind group is a round trip to the GPU process. - **Key** (`CacheState.computeBindGroupKey`): an exact string, not a hash, so distinct bindings cannot collide: shader id (fixes the layout and the uniform binding size), ids of the bound buffers in binding order (buffer arguments are bound whole), id of the uniform pool buffer. - **Ids** come from `CacheState.allocUid()` when a shader or `GPUBuffer` is created and are never reused. GPU pointers are not used because the buffer table recycles them after a free. Because the cache owns the id space, keys stay unique if several contexts share one `CacheState`. - **Id storage**: a buffer's id lives at its pointer's index next to the `bufferTable` assignment, is cleared on free, and a launch asserts every id is defined, so a missing id can never alias two buffers in a key. - **Invalidation**: none needed. Bind groups are immutable, and an entry is reachable only with the ids of the buffers it references, which stop occurring once one is destroyed. Dead entries hold no GPU memory and leave by LRU eviction. `dispose()` clears the cache. - **Sizing**: the key includes the uniform pool buffer, which is tied to the launch's position in the batch, so the cache must hold one repeating unit of work (a decode step: ~184 launches for a 1B model, ~900 for 80 layers); below that every lookup misses and evicts. Default 2048, constructor argument. - POD values are not in the key; they live in the uniform buffer's contents. Tests: reuse across flushes, distinct buffers and shaders, recycled GPU pointer, a freed buffer after the table has grown, defined ids in every key, replaced uniform pool buffer, LRU bound, two contexts sharing one `CacheState` and dispose. lint/tsc clean; jest 30 passed. ## Measurement (shared by the four runtime PRs) Real WebLLM (0.2.85, built from source) against tvmjs from each commit; Llama-3.2-1B-Instruct-q4f16_1, Apple M4 Max, Chrome Canary; same model lib, weights and prompt, temperature 0; one engine per arm, arms interleaved within each round with rotating order; 12 rounds × 256 decode tokens; "paired" is the median per-round ratio to `main`. Generated text identical in all arms. Logits through the WebLLM API are bit-exact against `main` for the whole stack. | tvmjs | passes / token | bind groups / token | uploads / token | decode tok/s | paired | |---|---|---|---|---|---| | main | 184 | 184 | 199 | 210.0 | 1.000 | | + shared compute pass | 5 | 184 | 199 | 218.1 | 1.026 | | + bind group cache | 5 | 4 | 199 | 222.5 | 1.062 | | + skip unchanged uploads | 5 | 4 | 54 | 237.2 | 1.119 | | + periodic submission (separate run, vs its own main at 206.8) | | | | 257.1 | 1.277 | Repeat runs of the first three: 1.118×, 1.137×, 1.136×; 1.04–1.08× when decoding after a 2750-token prompt; prefill neutral (0.97–1.04×). The saving is a fixed cost per launch, so larger models see less: Llama-3.1-8B 1.07× (54.5 → 58.9 tok/s) for the first three. One adapter tested. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
