Open
Fix segfault in MatMulNBits prepacking when weights/scales are parent-graph initializers in subgraph#31141
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
… initializers in subgraph Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>
Copilot
AI
changed the title
[WIP] Fix segfault in MatMulNBits prepacking with parent graph initializers
Fix segfault in MatMulNBits prepacking when weights/scales are parent-graph initializers in subgraph
Jul 29, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes a session-initialization crash when a MatMulNBits node lives inside a control-flow subgraph (e.g., If branch) but its quantized weights/scales are initializers in an enclosing (parent) graph scope. The core change ensures kernel construction (and prepacking) can resolve those parent-scope constant initializers when building OpKernelInfo.
Changes:
- Add
SessionState::GetConstantInitializedTensorsForKernelCreation()and subgraph-only bookkeeping to expose a kernel-creation constant map augmented with parent-scope initializers. - Build the augmented constant map for subgraph
SessionStateinstances duringFinalizeSessionStateImplbefore kernel creation. - Use the augmented map in
KernelRegistryManager::CreateKernel, plus add a regression test for the reportedIf-subgraph topology.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/matmul_4bits_test.cc | Adds regression test covering MatMulNBits inside If with parent-scope Bq/Bs initializers and accuracy_level=4. |
| onnxruntime/core/framework/session_state.h | Declares new kernel-creation constant map accessor and stores subgraph augmented-map state. |
| onnxruntime/core/framework/session_state.cc | Implements the accessor and builds the subgraph augmented constant map prior to CreateKernels(). |
| onnxruntime/core/framework/kernel_registry_manager.cc | Switches kernel creation to use the new kernel-creation constant tensor map. |
…s post-finalization After all subgraphs are finalized, erase the parent-scope OrtValue copies from outer_scope_augmented_constant_tensors_ (tracked in outer_scope_parent_only_indices_). This allows the parent's constant_initialized_tensors_ to free underlying tensor buffers once all prepack use counts reach zero, avoiding the memory regression the reviewer identified where copied OrtValues would hold extra shared_ptr references preventing memory release after PrepackConstantInitializedTensors. The copies are retained until after the subgraph finalization loop so that nested (grandchild) session states can still read them via GetConstantInitializedTensorsForKernelCreation() when building their own augmented maps.
tianleiwu
marked this pull request as ready for review
July 30, 2026 05:42
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.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
When a
MatMulNBitsnode lives inside a subgraph (e.g.Ifbranch) and its quantized weightsBandscalesare initializers of the enclosing graph,TryGetConstantInputfails to resolve those tensors during kernel construction andPrePack. On ARM64 withaccuracy_level=4(KleidiAI path), this causes anullptrdereference → segfault at session initialization.Root cause:
KernelRegistryManager::CreateKernelbuildsOpKernelInfousingsession_state.GetConstantInitializedTensors(), which only contains the current subgraph's constants. Parent-scope initializers are absent, soTryGetConstantInput(scales)silently returns null, andSQ4BitGemmPackQuantBDataAndBlkSumassertsQuantBScaleBegin != nullptr.Changes:
session_state.h— addsGetConstantInitializedTensorsForKernelCreation()method + two private fields (outer_scope_augmented_constant_tensors_,outer_scope_augmented_map_built_).session_state.cc— inFinalizeSessionStateImpl, betweenCleanInitializedTensorsFromGraph()andCreateKernels(), builds an augmented constant-tensor map for subgraph session states by iteratingparent_node->ImplicitInputDefs()and re-indexing each parent-scope constant into the current subgraph'sOrtValueNameIdxMap. The parent's own augmented map is consulted (already built, since parent finalizes before child), so arbitrary nesting depth is handled correctly. This map is intentionally not used byPrepackConstantInitializedTensorsto avoid double-prepacking outer-scope tensors.kernel_registry_manager.cc— one-line change:GetConstantInitializedTensors()→GetConstantInitializedTensorsForKernelCreation()inCreateKernel.matmul_4bits_test.cc— regression testMatMulNBits.SubgraphParentScopeInitializersthat constructs the exact topology from the bug report (If node, MatMulNBits in both branches, B/scales as parent-graph initializers,accuracy_level=4) and asserts session initialization and inference succeed.Motivation and Context
On ARM64, ORT segfaults during session initialization when loading a model with a
MatMulNBitsnode inside a subgraph whose quantized weights and scales are initializers of the parent graph withaccuracy_level=4. Reported against v1.28.0 on macOS/ARM64 and reproducible with any model following this pattern (e.g. decoder models with quantized weights shared acrossIfbranches).MatMulNBitsprepacking when weights and scales are initializers of parent graph #31137