Skip to content

FA4 Inference - #4186

Merged
wdykas merged 23 commits into
NVIDIA:mainfrom
wdykas:fa4-inference
Apr 18, 2026
Merged

FA4 Inference#4186
wdykas merged 23 commits into
NVIDIA:mainfrom
wdykas:fa4-inference

Conversation

@wdykas

@wdykas wdykas commented Apr 7, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Adding FA4 for 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 the @mcore-oncall.

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 the @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"

  1. When your PR is ready, click Ready for Review.
  2. 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.

For MRs into `dev` branch The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

@copy-pr-bot

copy-pr-bot Bot commented Apr 7, 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.

@wdykas

wdykas commented Apr 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test db6fe05

@svcnvidia-nemo-ci svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 7, 2026
@wdykas

wdykas commented Apr 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test ba8f399

@wdykas

wdykas commented Apr 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 481be87

@wdykas

wdykas commented Apr 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 481be87

@wdykas wdykas removed the Run tests label Apr 7, 2026
@wdykas

wdykas commented Apr 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 439e670

@wdykas
wdykas marked this pull request as ready for review April 8, 2026 14:58
@wdykas
wdykas requested review from a team as code owners April 8, 2026 14:58
@wdykas

wdykas commented Apr 8, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test bab7ba2

@svcnvidia-nemo-ci
svcnvidia-nemo-ci requested a review from a team April 8, 2026 14:59
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Apr 8, 2026
sidsingh-nvidia
sidsingh-nvidia reviewed Apr 8, 2026
View reviewed changes
Comment thread megatron/core/transformer/attention.py
@wdykas

wdykas commented Apr 17, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c68d142

@wdykas
wdykas added this pull request to the merge queue Apr 17, 2026
@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Apr 17, 2026
@wdykas
wdykas added this pull request to the merge queue Apr 17, 2026
@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Apr 17, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Apr 17, 2026
@wdykas
wdykas enabled auto-merge April 17, 2026 22:18
@wdykas

wdykas commented Apr 17, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 2bfbabe2bfbabe

@bbuschkaemper bbuschkaemper mentioned this pull request Apr 18, 2026
Closed
5 tasks
@wdykas

wdykas commented Apr 18, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test ff8a941ff8a941

@wdykas

wdykas commented Apr 18, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test ff8a941

@wdykas
wdykas added this pull request to the merge queue Apr 18, 2026
@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Apr 18, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit 76ac7c2 Apr 18, 2026
183 of 184 checks passed
@wdykas
wdykas deleted the fa4-inference branch April 18, 2026 21:08
Victarry pushed a commit to yanring/Megatron-LM that referenced this pull request Apr 20, 2026
Dennis Liu
Merge remote-tracking branch 'origin/main' into pr3431/fix-conflicts
9d4c1bc 
* origin/main: (286 commits)
  Rename MambaModel/MambaStack to HybridModel/HybridStack (NVIDIA#4099)
  Fix Megatron initialization with extra_args_provider (NVIDIA#4327)
  Fix RL to once again work with --skip-train (NVIDIA#4249)
  Add activation logging and tokens per expert logging (NVIDIA#3842)
  Make param_index_map always use unpacked (full numel) offsets (NVIDIA#4328)
  FA4 Inference (NVIDIA#4186)
  Fix RL reward due to stop token (NVIDIA#4096)
  cp: Fix UT timeout (NVIDIA#4310) (NVIDIA#4373)
  feat(ckpt): add --async-ckpt-use-cpu-shm argument (NVIDIA#4355)
  Update copy-pr-bot.yaml [skip ci]
  Docs: improve docstrings and comments in example training loop (NVIDIA#4041)
  Add QK layernorm support for dot-product attention in MambaModel (NVIDIA#4067)
  Fix bug with non-partial rollouts (NVIDIA#3964)
  [docs] ci: use parent-relative json_url for version picker (NVIDIA#4367)
  Add tables and histogram for RL staleness (NVIDIA#4097)
  Port DeepSeek Sparse Attention to `MambaModel` (NVIDIA#3553)
  docs: bump versions1.json to 0.17.0 (latest) (NVIDIA#4360)
  Fix potential coredump issue that occurs when saving a checkpoint (NVIDIA#1871)
  ci(gb200): add 1-node mr-github functional test variants (NVIDIA#4334)
  fix: wait for async P2P send before deallocating output tensor (NVIDIA#4047)
  ...

# Conflicts:
#	megatron/core/transformer/cuda_graphs.py
yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026
@wdykas William Dykas
root root
FA4 Inference (NVIDIA#4186)
bab54b7 
Co-authored-by: William Dykas <wdykas@oci-hsg-cs-001-vscode-03.cm.cluster>
Co-authored-by: root <root@nvl72065-T16.cm.cluster>
Co-authored-by: root <root@nvl72163-T17.cm.cluster>
yhgalaxy pushed a commit to yhgalaxy/Megatron-LM that referenced this pull request Jun 17, 2026
@wdykas William Dykas
root root @yhgalaxy
FA4 Inference (NVIDIA#4186)
4a7af36 
Co-authored-by: William Dykas <wdykas@oci-hsg-cs-001-vscode-03.cm.cluster>
Co-authored-by: root <root@nvl72065-T16.cm.cluster>
Co-authored-by: root <root@nvl72163-T17.cm.cluster>
Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
jon-barker pushed a commit to jon-barker/Megatron-LM that referenced this pull request Jul 10, 2026
@wdykas William Dykas
root root Jon Barker
FA4 Inference (NVIDIA#4186)
3c8f230 
Co-authored-by: William Dykas <wdykas@oci-hsg-cs-001-vscode-03.cm.cluster>
Co-authored-by: root <root@nvl72065-T16.cm.cluster>
Co-authored-by: root <root@nvl72163-T17.cm.cluster>
Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
@wdykas William Dykas
root root @terminator123
FA4 Inference (NVIDIA#4186)
842461f 
Co-authored-by: William Dykas <wdykas@oci-hsg-cs-001-vscode-03.cm.cluster>
Co-authored-by: root <root@nvl72065-T16.cm.cluster>
Co-authored-by: root <root@nvl72163-T17.cm.cluster>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@santhnm2 santhnm2 santhnm2 approved these changes
@kvareddy kvareddy kvareddy approved these changes
@shanmugamr1992 shanmugamr1992 shanmugamr1992 approved these changes
@sidsingh-nvidia sidsingh-nvidia sidsingh-nvidia approved these changes

Assignees

No one assigned

Labels

Approved All necessary approvals have been made complexity: medium Run functional tests

Projects

None yet

Milestone

Core 0.16

Development

Successfully merging this pull request may close these issues.

6 participants

@wdykas @svcnvidia-nemo-ci @santhnm2 @kvareddy @shanmugamr1992 @sidsingh-nvidia