Repository navigation
[None][feat] Support per-buffer reuse policies for attention and SSM - #4
Merged
jiaganc merged 11 commits intoOct 28, 2026
Conversation
Add an optional per-buffer reuse policy. A BufferConfig without one uses its layer's policy. An SSM buffer may set a different policy and then belongs to the SSM lifecycle of that policy while keeping its layer ID, so one SSM layer can span several layer groups. Attention buffers must use their layer's policy. Lifecycle registration and the storage layout are built per buffer, and per-layer storage attributes are keyed by (layer, lifecycle), so page index converters use the buffer's own lifecycle. Add getLayerGroupId(layerId, role) and its get_layer_group_id(layer_id, role) binding. The one-argument overload throws when a layer's buffers span several groups; layerGrouping lists such a layer under each group. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
…By Agent) _remove_zero_size_buffers rebuilt SSM layer configs without their reuse policy, which turned PRIVATE and OPTIONAL SSM layers into REQUIRED ones whenever any layer had a zero-size buffer. Forward the layer policy; the retained buffers keep their own policies. Describe how the layer-group helpers handle an SSM layer listed under several groups and cover them with a split layer. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
jiaganc
marked this pull request as draft
October 28, 2026 05:23
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
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.
Description
Builds on NVIDIA#21088 (
8083cb24e0). Buffers in one model layer can now use different reuse policies without inventing virtual layer IDs.BufferConfig.reuse_policy=Noneinherits the existing layer policy; an explicit override applies to attention or SSM buffers. Lifecycle deduplication remains based on semantic properties, so overrides introduce at most three policy variants per attention window/sink/sparsity combination and three SSM groups.get_buffer_group_id(layer_id, role)andbuffer_groupingAPIs. KeepLayerGroupId. Documentget_layer_group_id(layer_id)andlayer_groupingas deprecated without emitting warnings; preserve their uniform-layer results and reject ambiguous split-layer queries.LayerAttrmaps withBufferAttr.scratchPagesPerBlockand direct per-lifecycle rational scratch-utilization accounting. Preserve coalescing within lifecycle/size/coalescing-group boundaries, expansion and allocation rounding.Consumers must supply independent indices for independent policies. Existing Python consumers reject split layers through the legacy native API's ambiguity checks. Current Mamba kernels share indices between recurrent state and replay buffers; separate indices are deferred to a follow-up PR. This PR does not turn existing Mamba replay buffers PRIVATE by default. Custom consumers can use the native per-buffer APIs for independent buffers. No high-level LLM API configuration changes.
The native scratch converter handles unequal buffer membership, including a layer block that crosses coalesced slots. Existing attention scratch-table kernels require whole layer blocks per slot; custom consumers must use the native converter or provide compatible index generation for crossing layouts. Kernel changes and adoption of these layouts in Python consumers are outside this PR. OPTIONAL attention reuse retains its existing block-aligned sliding-window endpoint requirements.
Test Coverage
Library and bindings compiled successfully with the lifecycle factory cleanup at
c1dd7a9a1e. Pre-commit checks passed. Validation of6f5596bd7b, including the Python consumer rollback, started on shared B200-02 (umbriel-b200-034). All 9 C++ configuration tests passed before the SSH connection dropped during the next suite. The host is currently unreachable; remaining results are unverified, pending recovery.Direct native API coverage includes uniform and ambiguous lookup, inheritance, attention and SSM OPTIONAL selection at full/partial endpoints, unselected snapshot statistics, private state snapshots/beams/offload, and scratch address separation with unequal role membership and expansion. The validation run includes all 12 KVCM V2 C++ binaries with invariant checks enabled and the full Python KVCM V2 suite plus existing wrapper, Mamba, CSA2, staging/sharing, connector and disaggregation tests. Python tests use normal mode because the existing V1 KVCacheIndex sentinel fails unrelated debug assertions.
PR Checklist
Kept in draft; CI has not been triggered.