Skip to content

Conversation

@crcrpar
Copy link
Collaborator

@crcrpar crcrpar commented Nov 18, 2025

@crcrpar crcrpar requested a review from Copilot November 18, 2025 03:52
Copilot finished reviewing on behalf of crcrpar November 18, 2025 03:54
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull Request Overview

This PR adds support for torch.nn.functional.scaled_grouped_mm, a PyTorch function for scaled grouped matrix multiplication.

Key changes:

  • Implements the scaled_grouped_mm torchsymbol with input validation and shape inference
  • Adds three comprehensive test cases covering 2D×2D and 2D×3D tensor combinations with different scaling types
  • Registers the operation in the torch executor with appropriate availability checking

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File Description
thunder/torch/init.py Adds scaled_grouped_mm function with shape validation, dtype checking, and output shape computation
thunder/tests/test_ops.py Adds test cases for tensorwise and blockwise scaling scenarios with FP8 and MXFP8 dtypes
thunder/executors/torchex.py Registers scaled_grouped_mm operation and implements checker function for executor

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@crcrpar crcrpar marked this pull request as draft November 18, 2025 04:04
@crcrpar crcrpar requested a review from Copilot November 18, 2025 07:07
Copilot finished reviewing on behalf of crcrpar November 18, 2025 07:10
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull Request Overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

thunder/tests/test_ops.py:1

  • Both test_scaled_grouped_mm_3d2d_rowwise and test_scaled_grouped_mm_2d3d_rowwise test the same 2D @ 3D case (mat_a is 2D, mat_b after transpose is 3D). There is no test coverage for the 3D @ 2D case where mat_a would be 3D with shape (groups, m, k) and mat_b would be 2D with shape (k, n). Consider adding a test for this case or modifying one of the existing tests to cover it.
from collections.abc import Callable

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@crcrpar crcrpar requested a review from Copilot November 18, 2025 07:27
Copilot finished reviewing on behalf of crcrpar November 18, 2025 07:31
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull Request Overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@crcrpar crcrpar force-pushed the crpa/scaled_grouped_mm branch from 74c59d3 to f428a90 Compare November 18, 2025 07:37
@crcrpar crcrpar marked this pull request as ready for review November 18, 2025 11:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant