Skip to content

fix: initialize dense local CUDA graph state - #6167

Open
nvegesna-netizen wants to merge 1 commit into
NVIDIA:mainfrom
nvegesna-netizen:fix/local-cudagraph-dense-init
Open

fix: initialize dense local CUDA graph state#6167
nvegesna-netizen wants to merge 1 commit into
NVIDIA:mainfrom
nvegesna-netizen:fix/local-cudagraph-dense-init

Conversation

@nvegesna-netizen

@nvegesna-netizen nvegesna-netizen commented Jul 31, 2026

Copy link
Copy Markdown

Problem

GraphableMegatronModule can create a local CUDA graph manager from its constructor. For an MLP-only graph scope, that path reads is_moe_layer before TransformerLayer initializes it, so dense layers fail during construction.

Fix

Initialize the dense default before the base constructor while preserving the value already set by MoETransformerLayer. TransformerLayer still derives the final value from the constructed MLP later in initialization.

Validation

  • Added a focused unit regression that constructs a dense TransformerBlock with local MLP CUDA graph scope and verifies every layer receives a graph manager.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Jul 31, 2026

Copy link
Copy Markdown

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.

@nvegesna-netizen
nvegesna-netizen marked this pull request as ready for review July 31, 2026 02:16
@nvegesna-netizen
nvegesna-netizen requested review from a team as code owners July 31, 2026 02:16
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.

1 participant