Skip to content

Isolate MTP head metadata and decoder configuration - #2708

Draft
Tianlei Wu (tianleiwu) wants to merge 3 commits into
mainfrom
tlwu/20261012/ga01-head-metadata
Draft

Tianlei Wu (tianleiwu) wants to merge 3 commits into
mainfrom
tlwu/20261012/ga01-head-metadata

Conversation

@tianleiwu

@tianleiwu Tianlei Wu (tianleiwu) commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

Description

Keeps the MTP head decoder's configuration and state metadata independent from the target decoder, and resolves the disabled-indexer capacity from the head session's actual metadata instead of inheriting it from the target. This prevents a target-only setting (selected-logits indices, compact-state/indexer capacity) from leaking into the head and producing mis-sized graph inputs when target and head differ.

Stacked on #2714. Base is kvaishnavi/qwen-38-flash-run (head 62b34328); the diff shows only this change. Retarget to main once #2713 and #2714 merge.

Summary of Changes

Head/target isolation

File Change
src/config.cpp CreateMtpDecoderConfig clears the target-only logits_indices input and always resets the head's state_update_capacity to 0.
src/models/model.cpp Resolves the head decoder's disabled-indexer capacity from the head session's metadata before validating the session.
src/models/model_state_manifest.{h,cpp} Adds ResolveDisabledIndexerCapacity, which reads the capacity from actual head session metadata and rejects malformed or mismatched metadata without mutating state.

Tests and docs

File Change
test/cpp/engine/model_state_manifest_tests.cpp Malformed metadata (no mutation), inconsistent head layers, and enabled-target mismatch rejection.
test/cpp/engine/mtp_config_tests.cpp Six cross-capacity cases (target/head capacity 4/4, 8/4, 8/8 x indexed/non-indexed target logits); the existing one-layer fixed-indexer projection test now also asserts the head's cleared logits_indices and zero state-update capacity.
docs/paged_attention_engine.md Documents the head isolation behavior.

Intentionally excluded

Offload-specific clearing, lookup/hybrid policy, capture/debug instrumentation and exporter changes are not included.

Testing

  • CPU-only Release build (-DUSE_CUDA=OFF -DENABLE_PYTHON=OFF) on top of Add build and runtime support for Qwen-3.8 Flash Next #2714's head: 917 pass, 0 fail, 3 skipped vs 913 pass / 0 fail on the unmodified Add build and runtime support for Qwen-3.8 Flash Next #2714 head (+4 new tests).
  • Mutation checks: reverting the src/config.cpp change fails ProjectsOneLayerFixedIndexerStateForTheHead and ResolvesDisabledHeadIndexerCapacityIndependentOfTarget; making ResolveDisabledIndexerCapacity a no-op fails both manifest rejection tests and the MTP resolve test; resolving in place instead of on a copy fails the no-mutation test. EnabledIndexerCapacityMismatchRemainsRejected is a regression guard that passes either way, and the model.cpp call site has no unit test because it needs a real session.
  • git clang-format (20.1.8, the pinned version) and git diff --check are clean. Not run: CUDA, Python, guidance and telemetry surfaces.
  • Separate evidence from a private integration of the same logic (not this branch): real Qwen MTP1 head-only CUDA graph replay at query widths 1 and 2 with eager vs captured tokens identical.
  • No throughput, quality or default-behavior claim is made.

Motivation and Context

Found while qualifying head-only CUDA graph capture for Qwen MTP: a head that shares the target's selected-logits/indexer configuration can be planned with the wrong capacity. Making the head's own session metadata the source of truth removes that dependency. Rebased from the earlier qmoe-p0-base lineage onto #2714; #2714 removed engine.max_graph_query_tokens_per_request, so the assertions on it were dropped.

Checklist

  • Tests added/updated
  • No breaking changes (main configuration defaults verified unchanged)
  • Documentation updated

kunal-vaishnavi and others added 3 commits October 12, 2026 20:30
Split model builders, loaders, quantization, export tooling, builder tests, and builder documentation from kvaishnavi/qwen-38-flash (bf8e32c).
Stack runtime configuration, build integration, Engine and generator execution, Engram and hybrid/indexer state management, runtime tests, and runtime documentation on kvaishnavi/qwen-38-flash-create. Preserve committed changes from bf8e32c.
Keep the MTP head decoder's configuration independent from the target decoder: clear the target-only selected-logits input and state-update capacity when projecting the head config, and resolve the disabled indexer capacity from the head session's actual metadata instead of inheriting it from the target. A target-only setting can no longer leak into the head and produce mis-sized graph inputs when target and head differ.

Adds ModelStateManifest::ResolveDisabledIndexerCapacity (rejects malformed or mismatched metadata without mutating state), unit tests and documentation.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@tianleiwu
Tianlei Wu (tianleiwu) force-pushed the tlwu/20261012/ga01-head-metadata branch from 7c2240e to 61dafbd Compare October 12, 2026 23:09
@tianleiwu
Tianlei Wu (tianleiwu) changed the base branch from tlwu/20261012/qmoe-p0-base to kvaishnavi/qwen-38-flash-run October 12, 2026 23:10
Base automatically changed from kvaishnavi/qwen-38-flash-run to main October 30, 2026 00:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants