Repository navigation
feat(dynamo): alias multi-token KV cache writes so multi-profile works with static cache - #4749
Open
cehongwang wants to merge 2 commits into
Open
cehongwang wants to merge 2 commits into
cehongwang wants to merge 2 commits into
Conversation
The KV fast path for `aten.index_copy` accepted only a single-position write (`src.shape[2] == 1`), which made multi-optimization-profile engines mutually exclusive with an aliased static cache. A dynamic query axis -- the thing that makes a prefill profile mean anything -- leaves the write extent symbolic, so the validator rejected, the converter fell back to scatter, and aliasing was lost. Because `_buffer_lifting` had already erased the `copy_` write-back on the promise of aliasing, the cache then silently stopped updating: cosine similarity against eager decayed 1.00 -> 0.77 -> 0.52 -> 0.15, no error raised. TensorRT has taken a multi-token update since 10.15.1 (`[B, N, S_new, H]`, with kLINEAR defined as `out[i, :, writeIndices[i] + s, :] = update[i, :, s, :]`), so the restriction was ours, not the layer's. Accept any contiguous run of write positions, static or dynamic. Contiguity is load-bearing rather than cosmetic: kLINEAR writes a block from one start, so a scattered index of (5, 1, 9) would land at 5, 6, 7 -- silently, with the cache aliased. Proving it has to survive lowering. `constant_fold`, inside `post_lowering`, replaces a static-length `arange` with a frozen parameter, and this validator is consulted on both sides: by `_buffer_lifting` before, and by the partitioner after. Matching only the `arange` node answers True before and False after, and a write whose `copy_` was dropped as engine-aliased then aliases nothing, which `assert_predicted_kv_aliased` turns into a hard error. So contiguity is also read off a materialised constant index by value, which additionally unlocks aliasing for a fixed-length prefill. Also: * `writeIndices` must be `[batch]` holding only the start position. It was passed the whole index tensor, which is correct only at length one. * `emit_kv_cache_update_layer` now declines unless the update agrees with the cache on batch/heads/head_dim and those, plus maxSequenceLength, are static. An extent static in FX metadata can still reach the network as -1, and emitting anyway fails the build from inside TensorRT's validate() with an error naming the layer and nothing about where the shape went. * `index_copy_fallback` handles dynamic shapes, reusing the shape-aware expand the `slice_scatter` fallback already uses. Once the validator claims a node under `supports_dynamic_shapes=True` the registry commits it to this converter, so the in-function fallback is the only recourse left; raising there would abort a compile that previously succeeded via a Torch fallback. * The validator requires a rank-1 index. aten accepts a 0-d one, which reached `add_kv_cache_update` as a bare `nbDims == 1` check naming nothing. Two incidental fixes were needed to get the touched file past its own pre-commit hooks, both latent on main: an unused `type: ignore` (replaced with a narrowing isinstance, so it holds across mypy versions) and a `typos` false positive on "mis-evaluate". Verified on `google/gemma-3-1b-it` (26 layers) with a prefill profile [1,16] and a decode profile [1,1]: one engine, no Torch fallbacks, 52 aliased `kv_cache_update` bindings, all 52 caches mutated in place, and top-1 tokens matching eager across an 8-token prefill then decode at positions 8..11.
…e validator HF's static cache holds cumulative_length as a [1] tensor, so a multi-position write's index is arange(n) + <[1] tensor>. _is_uniform_shift only accepted 0-d, which rejected every Gemma-3 cache write under a dynamic prefill profile: no aliasing, the copy_ already dropped, and decode silently read a stale cache.
github-actions
Bot
added
component: tests
This branch has not been deployed
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
2 participants
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.
Description
Multi-optimization profiles and the aliased static KV cache could not be used together. When an engine had a dynamic prefill profile next to the decode profile, every cache write had a symbolic length.
_index_copy_kv_eligiblerejected those writes because it requiredsrc.shape[2] == 1. The write then went to a Torch fallback,lift_mutated_buffershad already dropped itscopy_, and decode read a stale cache without raising an error. On mini Gemma-3 this gave 3 engines, 0 aliased caches, and cosine similarity against eager falling from 0.52 to −0.29 within four decode steps.IKVCacheUpdateLayeralready supports multi-token and dynamic-length updates (updateis[B, N, S_new, H]), so the single-token limit came from Torch-TensorRT, not TensorRT. This PR removes it:arangewith step 1, optionally shifted by a scalar, 0-d, or[1]offset (HF'scumulative_lengthis[1]), or a folded constant holding an ascending run. It reads that constant becauseconstant_foldrewrites thearange, and the validator runs on both sides ofpost_lowering. Scattered or strided indices still fall back, since the layer writes one contiguous block.writeIndicesis now the start position only (index[0]), andsupports_dynamic_shapes=True.index_copy_fallbacknow handles dynamic extents instead of raising.emit_kv_cache_update_layerdeclines when batch, heads, or head_dim don't match the cache or aren't static in the network.Verification
gemma-3-1b-itwith a decode[1,1]profile and a prefill[1,16]profile gives 1 engine, 0 fallbacks, 52/52 aliasedkv_cache_updatebindings, and 52/52 caches updated in place. Top-1 matches eager on an 8-token prefill and on decode at positions 8–11, with cosine similarity ≥ 0.9999. Single-profile decode is unchanged.103 tests pass across
test_index_copy_kv,test_aliased_io(including the cudagraphs and serialization variants), andtest_buffer_lifting. The new tests cover dynamic and static multi-position writes, both operand orders, a[1]offset, and the fallbacks for scattered, strided, 0-d-index, and declined-layer writes.Type of change
Checklist: