Skip to content

[ROCm] Add scaled_mm v2 support.#165528

Closed
jagadish-amd wants to merge 3 commits intopytorch:mainfrom
jagadish-amd:add_scaled_mm_rocm
Closed

[ROCm] Add scaled_mm v2 support.#165528
jagadish-amd wants to merge 3 commits intopytorch:mainfrom
jagadish-amd:add_scaled_mm_rocm

Conversation

@jagadish-amd
Copy link
Contributor

@jagadish-amd jagadish-amd commented Oct 15, 2025 •
edited
Loading

Add mx fp4 support in Blas.cpp.
Updated the scale_kernel_dispatch array and ScaledGemmImplementation enum to include MXFP4 support.
Modify the tests under test_scaled_matmul_cuda accordingly.

PYTORCH_TEST_WITH_ROCM=1 python test/test_scaled_matmul_cuda.py -v -k test_blockwise
115 test passed.

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @naromero77amd

@jagadish-amd
Add scaled_mm v2 support.
45ea717 
Add mx fp4 support in Blas.cpp.
Modify the tests under test_scaled_matmul_cuda accordingly.

PYTORCH_TEST_WITH_ROCM=1 python test/test_scaled_matmul_cuda.py -v -k test_blockwise
115 test passed.

Signed-off-by: Jagadish Krishnamoorthy <jagadish.krishnamoorthy@amd.com>
@pytorch-bot
Copy link

pytorch-bot bot commented Oct 15, 2025 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/165528

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 05ddd83 with merge base 53f9ae0 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@jagadish-amd
Copy link
Contributor Author

jagadish-amd commented Oct 15, 2025

cc @jeffdaily

@jagadish-amd jagadish-amd changed the title Add scaled_mm v2 support. [ROCm] Add scaled_mm v2 support. Oct 15, 2025
@pytorch-bot pytorch-bot bot added the module: rocm AMD GPU support for Pytorch label Oct 15, 2025
@jagadish-amd
Copy link
Contributor Author

jagadish-amd commented Oct 15, 2025

@pytorchbot label "topic: not user facing"

@pytorch-bot pytorch-bot bot added the topic: not user facing topic category label Oct 15, 2025
@jeffdaily jeffdaily added ciflow/rocm-mi355 Trigger "default" config CI on ROCm MI355 runners ciflow/rocm Trigger "default" config CI on ROCm ciflow/rocm-mi300 Trigger "default" config CI on ROCm MI300 labels Oct 15, 2025
@jeffdaily jeffdaily marked this pull request as ready for review October 15, 2025 22:27
@jeffdaily
Copy link
Collaborator

jeffdaily commented Oct 15, 2025

@pytorchbot merge

@pytorch-bot
Copy link

pytorch-bot bot commented Oct 15, 2025

Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge. Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt. If you think this is a mistake, please contact PyTorch Dev Infra.

@jeffdaily jeffdaily added release notes: rocm mandatorylabel and removed topic: not user facing topic category labels Oct 15, 2025
@jeffdaily
Copy link
Collaborator

jeffdaily commented Oct 15, 2025

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 15, 2025
@pytorchmergebot
Copy link
Collaborator

pytorchmergebot commented Oct 15, 2025

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@pytorch-bot pytorch-bot bot removed ciflow/trunk Trigger trunk jobs on your pull request ciflow/rocm Trigger "default" config CI on ROCm ciflow/rocm-mi300 Trigger "default" config CI on ROCm MI300 ciflow/rocm-mi355 Trigger "default" config CI on ROCm MI355 runners labels Oct 16, 2025
@jeffdaily jeffdaily added ciflow/rocm Trigger "default" config CI on ROCm ciflow/rocm-mi300 Trigger "default" config CI on ROCm MI300 ciflow/rocm-mi355 Trigger "default" config CI on ROCm MI355 runners labels Oct 16, 2025
@jeffdaily
Copy link
Collaborator

jeffdaily commented Oct 16, 2025

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 16, 2025
@pytorchmergebot
Copy link
Collaborator

pytorchmergebot commented Oct 16, 2025

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@slayton58 slayton58 mentioned this pull request Oct 17, 2025
Closed
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Fix B200 test fails in scaled_mm
99214a0 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

@jagadish-amd @jeffdaily @drisspg

Subscribers:

@Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlayton@meta.com>

[ghstack-poisoned]
@slayton58 slayton58 mentioned this pull request Oct 17, 2025
Closed
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Fix B200 test fails in scaled_mm
27ef1ea 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

jagadish-amd jeffdaily drisspg

Subscribers:

Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

ghstack-source-id: 3bb8060
Pull Request resolved: #165746
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Update base for Update on "Fix B200 test fails in scaled_mm"
010bfe0 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

jagadish-amd jeffdaily drisspg

Subscribers:

Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

[ghstack-poisoned]
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Update on "Fix B200 test fails in scaled_mm"
42553e4 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

jagadish-amd jeffdaily drisspg

Subscribers:

Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

[ghstack-poisoned]
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Fix B200 test fails in scaled_mm
6e823c6 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

jagadish-amd jeffdaily drisspg

Subscribers:

Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

ghstack-source-id: 01149f8
Pull Request resolved: #165746
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Fix B200 test fails in scaled_mm
b823d96 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

jagadish-amd jeffdaily drisspg

Subscribers:

Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

[ghstack-poisoned]
@slayton58 slayton58 mentioned this pull request Oct 17, 2025
Closed
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Fix B200 test fails in scaled_mm
9a402b3 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

jagadish-amd jeffdaily drisspg

Subscribers:

Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

ghstack-source-id: 01149f8
Pull Request resolved: #165747
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Update on "Fix B200 test fails in scaled_mm"
b0962bb 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

jagadish-amd jeffdaily drisspg

Subscribers:

Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

[ghstack-poisoned]
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Fix B200 test fails in scaled_mm
41f3ba9 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

@jagadish-amd @jeffdaily @drisspg

Subscribers:

@Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

ghstack-source-id: 3e93657
Pull Request resolved: #165747

Signed-off-by: Simon Layton <simonlayton@meta.com>
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Update on "Fix B200 test fails in scaled_mm"
7461e85 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

jagadish-amd jeffdaily drisspg

Subscribers:

Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

[ghstack-poisoned]
slayton58 added a commit that referenced this pull request Oct 17, 2025
@slayton58
Fix B200 test fails in scaled_mm
63b7b83 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

@jagadish-amd @jeffdaily @drisspg

Subscribers:

@Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>

ghstack-source-id: 7f834a5
Pull Request resolved: #165747

Signed-off-by: Simon Layton <simonlayton@meta.com>
pytorchmergebot pushed a commit that referenced this pull request Oct 17, 2025
@slayton58 @pytorchmergebot
Fix B200 test fails in scaled_mm (#165747)
39e0a83 
Summary:

PR #165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the scale/swizzle inference code to prevent this.

Fixes #165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

@jagadish-amd @jeffdaily @drisspg

Subscribers:

@Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>
Pull Request resolved: #165747
Approved by: https://github.com/eqy, https://github.com/drisspg, https://github.com/jeffdaily
Chao1Han pushed a commit to Chao1Han/pytorch that referenced this pull request Oct 21, 2025
@jagadish-amd @Chao1Han
[ROCm] Add scaled_mm v2 support. (pytorch#165528)
33b3c5d 
Add mx fp4 support in Blas.cpp.
Updated the scale_kernel_dispatch array and ScaledGemmImplementation enum to include MXFP4 support.
Modify the tests under test_scaled_matmul_cuda accordingly.

PYTORCH_TEST_WITH_ROCM=1 python test/test_scaled_matmul_cuda.py -v -k test_blockwise
115 test passed.

Pull Request resolved: pytorch#165528
Approved by: https://github.com/jeffdaily
Chao1Han pushed a commit to Chao1Han/pytorch that referenced this pull request Oct 21, 2025
@slayton58 @Chao1Han
Fix B200 test fails in scaled_mm (pytorch#165747)
7d03511 
Summary:

PR pytorch#165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the scale/swizzle inference code to prevent this.

Fixes pytorch#165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

@jagadish-amd @jeffdaily @drisspg

Subscribers:

@Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>
Pull Request resolved: pytorch#165747
Approved by: https://github.com/eqy, https://github.com/drisspg, https://github.com/jeffdaily
zhudada0120 pushed a commit to zhudada0120/pytorch that referenced this pull request Oct 22, 2025
@jagadish-amd @zhudada0120
[ROCm] Add scaled_mm v2 support. (pytorch#165528)
79c62b8 
Add mx fp4 support in Blas.cpp.
Updated the scale_kernel_dispatch array and ScaledGemmImplementation enum to include MXFP4 support.
Modify the tests under test_scaled_matmul_cuda accordingly.

PYTORCH_TEST_WITH_ROCM=1 python test/test_scaled_matmul_cuda.py -v -k test_blockwise
115 test passed.

Pull Request resolved: pytorch#165528
Approved by: https://github.com/jeffdaily
zhudada0120 pushed a commit to zhudada0120/pytorch that referenced this pull request Oct 22, 2025
@slayton58 @zhudada0120
Fix B200 test fails in scaled_mm (pytorch#165747)
dc90e1f 
Summary:

PR pytorch#165528 changes some scale/swizzle inference behavior in scaled_mm
tests - mxfp8 tests on Blackwell can get incorrectly classified,
resulting in failures.

Fix the scale/swizzle inference code to prevent this.

Fixes pytorch#165743

Test Plan:

```
pytest -svv test/test_scaled_matmul_cuda.py
```

Reviewers:

@jagadish-amd @jeffdaily @drisspg

Subscribers:

@Aidyn-A

Tasks:

Tags:
Signed-off-by: Simon Layton <simonlaytonmeta.com>
Pull Request resolved: pytorch#165747
Approved by: https://github.com/eqy, https://github.com/drisspg, https://github.com/jeffdaily
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jeffdaily jeffdaily jeffdaily approved these changes

@eqy eqy Awaiting requested review from eqy

@syed-ahmed syed-ahmed Awaiting requested review from syed-ahmed

@Aidyn-A Aidyn-A Awaiting requested review from Aidyn-A

Assignees

No one assigned

Labels

ciflow/rocm Trigger "default" config CI on ROCm ciflow/rocm-mi300 Trigger "default" config CI on ROCm MI300 ciflow/rocm-mi355 Trigger "default" config CI on ROCm MI355 runners ciflow/trunk Trigger trunk jobs on your pull request Merged module: rocm AMD GPU support for Pytorch open source release notes: rocm mandatorylabel

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@jagadish-amd @jeffdaily @pytorchmergebot @pytorchbot