Add unit tests for GDP inference by santhnm2 · Pull Request #6505 · NVIDIA/Megatron-LM

santhnm2 · 2026-08-13T05:15:47Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds unit tests for GDP inference.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com> Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Signed-off-by: Mikail Khona <mkhona@cw-dfw-cs-001-vscode-01.cm.cluster> Signed-off-by: Mikail Khona (NVIDIA) <mkhona@nvidia.com> Co-authored-by: Mikail Khona <mkhona@cw-dfw-cs-001-vscode-01.cm.cluster> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Give each householder copy its own checkpoint key in GatedDeltaProductMixer's sharded_state_dict so distributed checkpointing concatenates TP shards within a copy before the copies are merged. This keeps resharded (e.g. TP=2 -> TP=1) in_proj/conv1d tensors in the semantic [M0-all-ranks, M1-all-ranks, ...] order that the forward rearranges expect. Add _get_in_proj_checkpoint_split_layout and _get_conv_checkpoint_split_layout helpers (applied to both weight and bias) and a resharding unit test. Squashed from two commits on kezhik/dev-arch-mar2026: 465077d86347 Fix GDP checkpoint resharding across TP sizes 2759392384fc Fix GDP checkpoint resharding unit test Adapted to the renamed gated_delta_product.py module (was gated_delta_product_original_v4.py on the source branch); the test import was updated to match. Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Add a dynamic-batching inference path to GatedDeltaProductMixer that separates decode and prefill requests, runs each through the GDP kernels, and merges the results back into packed token order. Decode uses causal_conv1d_update plus the fused recurrent gated-delta-rule kernel with per-request state gathered and scattered through DynamicInferenceContext's slot-indexed caches; prefill runs a single variable-length chunk_gated_delta_product call. MVP scope excludes context parallelism, speculative decoding, chunked prefill, prefix caching, and CUDA-graph capture. Cherry-picked from kezhik/dev-arch-mar2026 commit 296b04884456 and adapted to this fork: - Applied onto the renamed gated_delta_product.py (was gated_delta_product_original_v4.py on the source branch). - Ported megatron/core/ssm/_packed_seq_helpers.py, a dependency the source branch already had but this fork lacked. Only check_fla_sequence_packing_support is used here, so the import is narrowed to that symbol to avoid unused imports. Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

…path Wire THD/packed-sequence support into GatedDeltaProductMixer's training and prefill forward path: build cu_seqlens and a per-token seq_idx from packed_seq_params (after the in_proj sequence-parallel all-gather and the context-parallel all-to-all), thread seq_idx through causal_conv1d_fn to reset convolution boundaries at document edges, and pass cu_seqlens to chunk_gated_delta_product. GDPContextParallel.pre_conv_ssm/post_conv_ssm now take packed_seq_params so the load-balancing undo/redo uses the packed layout. Add self.chunk_size and a causal_conv1d version check in __init__. Integrates NVIDIA-NeMo/nv-mistralai-megatron PR NVIDIA#122 (commit b8726f2edf4b, "IDM for GDP from internal gitlab") onto this branch: - Reconciled the forward() inference branch with the dynamic-inference path already on staging; NVIDIA#122's "no packed sequences during inference" assert now sits alongside the dynamic-batching dispatch. - Kept the _packed_seq_helpers.py port already added here (content-identical to NVIDIA#122's; only formatting differed) and dropped NVIDIA#122's duplicate copy. Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Packed sequence indices are built on every SSM layer forward. Reading cu_seqlens[-1] with .item() synchronizes the GPU and CPU even though the value is used only for validation. Remove the eager host read and let repeat_interleave reject invalid negative lengths. The number of output indices is already known from total_tokens, so pass it as output_size. This also avoids synchronizing to infer the output shape and keeps index construction on-device. Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Rename the GDP-specific submodule dataclass at its definition site and update its type annotations and call sites. This removes the alias required to distinguish it from the unrelated Mamba mixer submodule dataclass and reduces the risk of wrong imports. Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Record the new gdp_num_householder TransformerConfig field and its backward-compatible default of three in the Mamba MoE golden configuration. This keeps the config drift test aligned with the intentional GDP configuration surface. Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Gated Delta Product carried its own `_dynamic_inference` / `_ssm_decode` / `_ssm_prefill` trio, duplicating the request-level control flow that `SSMDynamicInferenceMixin` already owns for Mamba2: fetch the per-layer (conv_state, ssm_state) slabs, project, split the packed batch into decode and prefill partitions, run each through its kernels, merge back into packed token order, and project out. Subclass the mixin and implement only the two variant hooks. `ssm_decode` now takes the mixin's batch-first `[n, seq_len, proj_dim]` layout and rejects the speculative-decoding intermediate buffers explicitly rather than silently ignoring them; `ssm_prefill` reads its varlen metadata off the context instead of taking it as an argument list, and owns the chunked-prefill assertion, as the interface prescribes. Quarantine the static-batching path the same way MambaMixer does. Static decode moves out of the shared `forward` body into `_static_decode`, which delegates to `ssm_decode` with `batch_indices=None`; `forward` keeps only the training and static-prefill body. This removes the two `seqlen_offset > 0` branches that threaded static-batching bookkeeping through the training math. No functional change: `pre_conv_ssm` / `post_conv_ssm` are identity at cp_size == 1, which static decode already required. Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com> # Conflicts: # megatron/core/ssm/gated_delta_product.py

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

…tests

copy-pr-bot · 2026-08-13T05:15:50Z

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.

santhnm2 · 2026-08-13T05:36:33Z

/ok to test 02dc028

svcnvidia-nemo-ci · 2026-08-13T08:43:38Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31683334710

svcnvidia-nemo-ci · 2026-08-13T09:13:22Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31685592929

devnkong and others added 25 commits August 10, 2026 10:17

Add gated delta product (GDP) mixer and spec (NVIDIA#45)

d025a7f

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com> Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Tag GDP packed parameters for refit

02a7cf9

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Fail fast on invalid module specs

c3af950

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Fix GDP mixer module naming

e2c8855

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Update FLOPs calculation to account for GDP

9c4f6de

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

fix(optimizer): skip Muon for Mamba input projection

0e95267

Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

chore: add missing GDP test copyright header

f286a4c

Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

chore: apply PR formatting fixes

2feb52b

Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

refactor(ssm): rename packed sequence helper module

2fe04d6

Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

docs(ssm): describe gated delta product mixer

20d9d5b

Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

refactor(ssm): remove unused GDP constructor arguments

fe926f0

Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

fix(ssm): configure GDP Householder reflections

0d03d60

Signed-off-by: Kezhi Kong <kezhik@nvidia.com> Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Linting

6878aeb

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Merge with main

2483de0

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Merge remote-tracking branch 'upstream/main' into gdp_common_interface

b0280ab

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com> # Conflicts: # megatron/core/ssm/gated_delta_product.py

Add GDP inference unit tests

3b8a54f

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Merge remote-tracking branch 'upstream/main' into gdp_inference_unit_…

02dc028

…tests

santhnm2 marked this pull request as ready for review August 13, 2026 05:21

santhnm2 requested review from a team as code owners August 13, 2026 05:21

kvareddy approved these changes Aug 13, 2026

View reviewed changes

github-actions Bot added the complexity: low label Aug 13, 2026

Phlip79 approved these changes Aug 13, 2026

View reviewed changes

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Aug 13, 2026

santhnm2 enabled auto-merge August 13, 2026 05:36

copy-pr-bot Bot temporarily deployed to public August 13, 2026 05:37 Inactive

copy-pr-bot Bot deployed to test August 13, 2026 05:37 Active

copy-pr-bot Bot temporarily deployed to public August 13, 2026 05:40 Inactive

copy-pr-bot Bot temporarily deployed to public August 13, 2026 05:50 Inactive

santhnm2 added this pull request to the merge queue Aug 13, 2026

Merged via the queue into NVIDIA:main with commit fecb29f Aug 13, 2026
95 checks passed

Navigation Menu

Uh oh!

Add unit tests for GDP inference - #6505

Add unit tests for GDP inference#6505
santhnm2 merged 25 commits into
NVIDIA:mainfrom
santhnm2:gdp_inference_unit_tests

santhnm2 commented Aug 13, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 13, 2026

Uh oh!

santhnm2 commented Aug 13, 2026

Uh oh!

svcnvidia-nemo-ci commented Aug 13, 2026

Uh oh!

svcnvidia-nemo-ci commented Aug 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Uh oh!

Conversation

santhnm2 commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Aug 13, 2026

Uh oh!

santhnm2 commented Aug 13, 2026

Uh oh!

svcnvidia-nemo-ci commented Aug 13, 2026

Uh oh!

svcnvidia-nemo-ci commented Aug 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

santhnm2 commented Aug 13, 2026 •

edited

Loading