Skip to content

hook up real nvfp4 grouped_gemm#4316

Merged
vkuzo merged 34 commits intomainfrom
gh/vkuzo/254/head
Apr 23, 2026
Merged

hook up real nvfp4 grouped_gemm#4316
vkuzo merged 34 commits intomainfrom
gh/vkuzo/254/head

Conversation

@vkuzo
Copy link
Copy Markdown
Contributor

@vkuzo vkuzo commented Apr 22, 2026

Summary:

for nvfp4 + grouped_mm:

  1. hooks up mslk's to_nvfp4 quantization kernel for token group
    activations
  2. hooks up the nvfp4 recipe of torch._scaled_grouped_mm

still WIP

Test Plan:

pytest test/prototype/mx_formats/test_inference_workflow.py -s -k grouped_mm

vkuzo added 21 commits April 20, 2026 20:52
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
@vkuzo
Copy link
Copy Markdown
Contributor Author

vkuzo commented Apr 22, 2026

@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Apr 22, 2026

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 1 New Failure, 8 Pending

As of commit b0697ac with merge base bfa7a94 (image):

NEW FAILURE - The following job has failed:

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 22, 2026
not torch_version_at_least("2.8.0"),
reason="torch.compile requires PyTorch 2.8+",
)
def test_nvfp4_quantize_3d_param_similar_to_vllm(self):
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this isn't used so deleting

@vkuzo vkuzo added the module: not user facing Use this tag if you don't want this PR to show up in release notes label Apr 22, 2026
vkuzo added 4 commits April 23, 2026 09:23
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Apr 23, 2026
Summary:

for nvfp4 + grouped_mm:
1. hooks up mslk's to_nvfp4 quantization kernel for token group
   activations
2. hooks up the nvfp4 recipe of `torch._scaled_grouped_mm`

still WIP

Test Plan:

```
pytest test/prototype/mx_formats/test_inference_workflow.py -s -k grouped_mm
```
ghstack-source-id: 44d2809
ghstack-comment-id: 4299981698
Pull-Request: #4316
vkuzo added 3 commits April 23, 2026 09:24
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Apr 23, 2026
Summary:

for nvfp4 + grouped_mm:
1. hooks up mslk's to_nvfp4 quantization kernel for token group
   activations
2. hooks up the nvfp4 recipe of `torch._scaled_grouped_mm`

still WIP

Test Plan:

```
pytest test/prototype/mx_formats/test_inference_workflow.py -s -k grouped_mm
```
ghstack-source-id: 44d2809
ghstack-comment-id: 4299981698
Pull-Request: #4316
vkuzo added 2 commits April 23, 2026 09:25
[ghstack-poisoned]
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Apr 23, 2026
Summary:

for nvfp4 + grouped_mm:
1. hooks up mslk's to_nvfp4 quantization kernel for token group
   activations
2. hooks up the nvfp4 recipe of `torch._scaled_grouped_mm`

still WIP

Test Plan:

```
pytest test/prototype/mx_formats/test_inference_workflow.py -s -k grouped_mm
```
ghstack-source-id: 44d2809
ghstack-comment-id: 4299981698
Pull-Request: #4316
vkuzo added 2 commits April 23, 2026 09:27
[ghstack-poisoned]
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Apr 23, 2026
Summary:

for nvfp4 + grouped_mm:
1. hooks up mslk's to_nvfp4 quantization kernel for token group
   activations
2. hooks up the nvfp4 recipe of `torch._scaled_grouped_mm`

still WIP

Test Plan:

```
pytest test/prototype/mx_formats/test_inference_workflow.py -s -k grouped_mm
```
ghstack-source-id: 8d9fd1f
ghstack-comment-id: 4299981698
Pull-Request: #4316
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Apr 23, 2026
Summary:

for nvfp4 + grouped_mm:
1. hooks up mslk's to_nvfp4 quantization kernel for token group
   activations
2. hooks up the nvfp4 recipe of `torch._scaled_grouped_mm`

still WIP

Test Plan:

```
pytest test/prototype/mx_formats/test_inference_workflow.py -s -k grouped_mm
```
ghstack-source-id: cae5aae
ghstack-comment-id: 4299981698
Pull-Request: #4316
[ghstack-poisoned]
@vkuzo vkuzo changed the base branch from gh/vkuzo/253/head to main April 23, 2026 17:22
@vkuzo vkuzo merged commit e45f934 into main Apr 23, 2026
50 of 53 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: not user facing Use this tag if you don't want this PR to show up in release notes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants