feat(hybrid): add opt-in shared-prefix attention and Mamba execution

[jalbericiola]

GRPO batches contain several completions with identical prompts. This adds an opt-in packed star/forest representation so the hybrid backbone computes a prompt once per packed group and processes each suffix independently.

Attention reuses prompt keys/values while each suffix remains causally isolated. Mamba shares the aligned prompt head, forks recurrent state into independent branches, and handles the residual prompt tail and convolution halo per branch. Backward accumulates branch contributions into the shared prefix. Ragged layouts, TP/CP transport, activation recomputation, logical-token router accounting and dense MTP compatibility are part of the same execution contract. MTP prefix sharing is not included.

The existing dense path remains the default. Supported combinations must be checked explicitly; unsupported layouts and training modes fail before execution.

Related design discussion: #7386. That proposal covers exact agent-call trees. This draft focuses on repeated GRPO prompts, packed star/forest execution, aligned Mamba state forking and dense MTP compatibility. The layout/API and numerical acceptance criteria need maintainer review before merge.

Numerical status

Within-implementation logprob/training agreement has been qualified for the production configuration. A separate controlled dense/shared backbone comparison still showed 13.9305% relative L2 difference against approximately 1% repeat variation. That gap remains open. Treat this as an experimental draft until the cross-implementation contract and current-main GPU coverage are resolved.

Observed one-seed end-to-end quality is supporting evidence, not proof of gradient equivalence.

Validation and review status

18 isolated CPU contracts pass and23 exported Python files compile. Five upstream API methods are AST-identical and no new undefined-name/redefinition/syntax or direct global process-group reads were found. Current-main distributed GPU and determinism-registry gates remain open.

The adapted patch applies to the pinned current-main base and reproduces candidate file hashes exactly. This remains a draft pending the stated runtime/CI gates.

Final local checks

The native tools/autoformat.sh workflow, Black, isort, pylint, Ruff and git diff --check pass. Formatter changes preserve executable AST behavior. Mypy remains advisory and is not clean; these results do not establish native distributed runtime or GPU qualification. Deterministic replay coverage, layer-ownership factoring and the numerical acceptance contract remain draft merge gates.

Diff size

Incremental contribution: 5,847 added lines, including 130 added test lines (2.22%); 0 test lines removed. These are physical diff lines in test directories, including comments and blank lines, not test coverage.

Commits are cryptographically signed and locally verified. GitHub signing-key registration is pending, so GitHub may label the signature unverified.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.