Skip to content

Float8StaticActivationFloat8WeightConfig using keepdim=True#3759

Merged
jerryzh168 merged 1 commit into
mainfrom
gh/jerryzh168/31/head
Jan 29, 2026
Merged

Float8StaticActivationFloat8WeightConfig using keepdim=True#3759
jerryzh168 merged 1 commit into
mainfrom
gh/jerryzh168/31/head

Conversation

@jerryzh168

@jerryzh168 jerryzh168 commented Jan 29, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Summary:

  • Pass keepdim=True to all AffineQuantizedMinMaxObserver instances in Float8StaticActivationFloat8WeightConfig (prototype feature) transform to preserve scale tensor dimensions
  • Remove manual scale reshaping workarounds after calculate_qparams() calls since keepdim=True now handles this
  • Extend test_static_quant_with_output_quantization with parametrize to test both 2D and 3D inputs

Test Plan:

  • Run test_static_quant_with_output_quantization with both input shapes:
    • 2D: (4, 64)
    • 3D: (2, 4, 64)
  • Verify existing float8 static quantization tests still pass
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128
pip install --pre mslk --index-url https://download.pytorch.org/whl/nightly/cu128
pytest test/prototype/test_prototype_float8_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:

Summary:
- Pass keepdim=True to all AffineQuantizedMinMaxObserver instances in Float8StaticActivationFloat8WeightConfig transform to preserve scale tensor dimensions
- Remove manual scale reshaping workarounds after calculate_qparams() calls since keepdim=True now handles this
- Extend test_static_quant_with_output_quantization with parametrize to test both 2D and 3D inputs

Test Plan:
- Run test_static_quant_with_output_quantization with both input shapes:
  - 2D: (4, 64)
  - 3D: (2, 4, 64)
- Verify existing float8 static quantization tests still pass

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jan 29, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

✅ No Failures

As of commit 3176409 with merge base c6c34d4 (image):
💚 Looks good so far! There are no failures yet. 💚

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

jerryzh168 added a commit that referenced this pull request Jan 29, 2026
Summary:
- Pass keepdim=True to all AffineQuantizedMinMaxObserver instances in Float8StaticActivationFloat8WeightConfig transform to preserve scale tensor dimensions
- Remove manual scale reshaping workarounds after calculate_qparams() calls since keepdim=True now handles this
- Extend test_static_quant_with_output_quantization with parametrize to test both 2D and 3D inputs

Test Plan:
- Run test_static_quant_with_output_quantization with both input shapes:
  - 2D: (4, 64)
  - 3D: (2, 4, 64)
- Verify existing float8 static quantization tests still pass

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 9b69d32
Pull Request resolved: #3759
@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 Jan 29, 2026
@jerryzh168 jerryzh168 added module: not user facing Use this tag if you don't want this PR to show up in release notes topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) labels Jan 29, 2026
@jerryzh168 jerryzh168 changed the base branch from gh/jerryzh168/31/base to main January 29, 2026 17:53
@jerryzh168 jerryzh168 merged commit 35fc480 into main Jan 29, 2026
37 of 39 checks passed
andrewor14 added a commit that referenced this pull request Jan 29, 2026
**Summary:**

```
pytest test/test_ops_rowwise_scaled_linear_sparse_cutlass.py
```

This gives

```
E       RuntimeError: The size of tensor a (4) must match the size of tensor b (65) at non-singleton dimension 1
```

Seems like it's caused by #3759.
We just need to update the test.

**Test Plan:**

```
pytest test/test_ops_rowwise_scaled_linear_sparse_cutlass.py
```
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 topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants