Skip to content

[None][feat] Support per-buffer reuse policies for attention and SSM - #4

Merged
jiaganc merged 11 commits into
jiaganc:codex/layer-reuse-policyfrom
VALLIS-NERIA:feat/ssm-buffer-reuse-policy
Oct 28, 2026
Merged

jiaganc merged 11 commits into
jiaganc:codex/layer-reuse-policyfrom
VALLIS-NERIA:feat/ssm-buffer-reuse-policy

Conversation

@VALLIS-NERIA

@VALLIS-NERIA VALLIS-NERIA commented Oct 27, 2026 •

Copy link
Copy Markdown

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=None inherits 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.

  • Add separate get_buffer_group_id(layer_id, role) and buffer_grouping APIs. Keep LayerGroupId. Document get_layer_group_id(layer_id) and layer_grouping as deprecated without emitting warnings; preserve their uniform-layer results and reject ambiguous split-layer queries.
  • Existing Python consumers retain their layer-based APIs and require one lifecycle per layer. Adoption of buffer-based lookup in wrappers, sharing, connectors and disaggregation is deferred until those consumers support split buffers.
  • Replace persistent LayerAttr maps with BufferAttr.scratchPagesPerBlock and direct per-lifecycle rational scratch-utilization accounting. Preserve coalescing within lifecycle/size/coalescing-group boundaries, expansion and allocation rounding.
  • Retain per-buffer policy inheritance when filtering zero-sized buffers. PRIVATE buffers remain request-local; OPTIONAL snapshots are explicitly selected and unselected snapshots count as neither hits nor misses.
  • Register lifecycles only for actual buffers. Remove the layer-only lifecycle factory and validate each layer independently of its buffer count.

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 of 6f5596bd7b, 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

  • Description explains the behavior, rationale and consumer limitations.
  • Changes follow coding guidelines; commits have DCO sign-off and pass pre-commit hooks.
  • Behavioral regression tests cover the new paths.
  • Public stubs and developer documentation updated.
  • No high-level LLM API, dependency or ownership changes.
  • Final affected-test results verified.

Kept in draft; CI has not been triggered.

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
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>
@jiaganc jiaganc changed the title [None][feat] Support buffer-level reuse policy for SSM layers [None][feat] Support per-buffer reuse policies for attention and SSM Oct 28, 2026
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
@jiaganc
jiaganc merged commit 698b619 into jiaganc:codex/layer-reuse-policy Oct 28, 2026
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.

2 participants