Simon Layton 7c6c5d04fe Add scaled_grouped_mm_v2 and python API (#165154) ...

Summary:

* Add `torch._scaled_grouped_mm_v2` with more functionality and
  extensibility for future formats
* Add `torch.nn.functional.scaled_grouped_mm` as public entrypoint
* Test both original and v2 functionality

Test Plan:

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

Reviewers:

Subscribers:

Tasks:

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

Pull Request resolved: https://github.com/pytorch/pytorch/pull/165154
Approved by: https://github.com/drisspg, https://github.com/danielvegamyhre

2025-10-15 17:47:23 +00:00

..

__init__.py

…

HasDecompTest.test_aten_core_operators.expect

[CUDA][cuBLAS] Aten GEMM overload for FP32 output from FP16/BF16 inputs (#150812)

2025-04-18 01:53:26 +00:00

HasDecompTest.test_has_decomposition.expect

Add scaled_grouped_mm_v2 and python API (#165154)

2025-10-15 17:47:23 +00:00

TestAutograd.test_function-x_grad_desc.expect

…

TestAutograd.test_function-y_grad_desc.expect

…

TestFXAPIBackwardCompatibility.test_class_member_back_compat-fx_backcompat_class_members.expect

Add output_node util function to fx.Graph (#139770)

2024-11-07 18:54:59 +00:00

TestFXAPIBackwardCompatibility.test_function_back_compat-fx_backcompat_function_signatures.expect

[fx] Allow customization of submod name in split graph (#164035)

2025-10-01 16:26:14 +00:00

TestJit.test_cu_escaped_number.expect

…

TestJit.test_import_method.expect

…

TestJit.test_non_ascii_string.expect

…

TestJit.test_pretty_printer-empty_float_list_test.expect

…

TestJit.test_pretty_printer-empty_int_list_test.expect

…

TestJit.test_pretty_printer-if_one.expect

…

TestJit.test_pretty_printer-if_test.expect

…

TestJit.test_pretty_printer-loop_use_test.expect

…

TestJit.test_pretty_printer-print_weird_test.expect

…

TestJit.test_pretty_printer-python_op_name_test.expect

…

TestJit.test_pretty_printer-while_if_test.expect

…

TestJit.test_pretty_printer-while_test.expect

…

TestPytorchExportModes.test_aten_fallback.expect

Fix ONNX ATen fallback for non-caffe2 engines

2022-04-14 23:18:45 +00:00

TestPytorchExportModes.test_onnx_aten.expect

Fix ONNX ATen fallback for non-caffe2 engines

2022-04-14 23:18:45 +00:00

TestScript.test_annot_ast_mypy_fn.expect

…

TestScript.test_annot_ast_mypy_method.expect

…

TestScript.test_annot_ast_py3_fn.expect

…

TestScript.test_annot_ast_py3_method.expect

…

TestScript.test_annot_string_mypy_fn.expect

…

TestScript.test_annot_string_mypy_method.expect

…

TestScript.test_annot_string_py3_fn.expect

…

TestScript.test_annot_string_py3_method.expect

…

TestScript.test_annotated_script_fn.expect

Make string serialization of C++ FunctionSchema consistent with torchgen.model.FunctionSchema

2022-05-24 19:39:26 +00:00

TestScript.test_annotated_script_method.expect

…

TestScript.test_format-stdout.expect

…

TestScript.test_listconstruct_erasure.expect

Fix ONNX ATen fallback for non-caffe2 engines

2022-04-14 23:18:45 +00:00

TestScript.test_parser_type_annotations_comment.expect

…

TestScript.test_parser_type_annotations.expect

…

TestScript.test_print-stdout.expect

Remove set_default_dtype calls from jit and ops tests (#105072)

2023-07-15 03:18:33 +00:00

TestScript.test_python_frontend_py2.expect

…

TestScript.test_python_frontend_py3.expect

…

TestScript.test_python_frontend.expect

…

TestScript.test_string_print-stdout.expect

…

TestScript.test_torch_dot_tensor_annotation.expect

…

TestSparseCompressedCPU.test_print_SparseBSC_cpu.expect

Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)

2022-11-30 02:13:33 +00:00

TestSparseCompressedCPU.test_print_SparseBSR_cpu.expect

Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)

2022-11-30 02:13:33 +00:00

TestSparseCompressedCPU.test_print_SparseCSC_cpu.expect

Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)

2022-11-30 02:13:33 +00:00

TestSparseCompressedCPU.test_print_SparseCSR_cpu.expect

Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)

2022-11-30 02:13:33 +00:00

TestSparseCompressedCUDA.test_print_SparseBSC_cuda.expect

Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)

2022-11-30 02:13:33 +00:00

TestSparseCompressedCUDA.test_print_SparseBSR_cuda.expect

Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)

2022-11-30 02:13:33 +00:00

TestSparseCompressedCUDA.test_print_SparseCSC_cuda.expect

Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)

2022-11-30 02:13:33 +00:00

TestSparseCompressedCUDA.test_print_SparseCSR_cuda.expect

Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)

2022-11-30 02:13:33 +00:00

TestSparseCPU.test_print_coalesced_cpu_float64.expect

…

TestSparseCPU.test_print_uncoalesced_cpu_float64.expect

…

TestSparseCUDA.test_print_coalesced_cuda_float64.expect

…

TestSparseCUDA.test_print_uncoalesced_cuda_float64.expect

…

TestSparseMeta.test_print_meta_SparseBSC_float64.expect

Add sparse compressed meta tensor support (#120707)

2024-03-01 13:28:47 +00:00

TestSparseMeta.test_print_meta_SparseBSR_float64.expect

Add sparse compressed meta tensor support (#120707)

2024-03-01 13:28:47 +00:00

TestSparseMeta.test_print_meta_SparseCOO_float64.expect

Add sparse compressed meta tensor support (#120707)

2024-03-01 13:28:47 +00:00

TestSparseMeta.test_print_meta_SparseCSC_float64.expect

Add sparse compressed meta tensor support (#120707)

2024-03-01 13:28:47 +00:00

TestSparseMeta.test_print_meta_SparseCSR_float64.expect

Add sparse compressed meta tensor support (#120707)

2024-03-01 13:28:47 +00:00

TestSparseMPS.test_print_coalesced_mps_float32.expect

[MPS] Move sparsemps testing from test_mps to test_sparse (#161852)

2025-09-02 19:04:11 +00:00

TestSparseMPS.test_print_uncoalesced_mps_float32.expect

[MPS] Move sparsemps testing from test_mps to test_sparse (#161852)

2025-09-02 19:04:11 +00:00

TestTensorBoard.test_audio.expect

Bump protobuf version and refactor tensorboard tests (#154244)

2025-05-25 00:50:07 +00:00

TestTensorBoard.test_caffe2_simple_cnnmodel.expect

…

TestTensorBoard.test_caffe2_simple_model.expect

…

TestTensorBoard.test_histogram_auto.expect

Bump protobuf version and refactor tensorboard tests (#154244)

2025-05-25 00:50:07 +00:00

TestTensorBoard.test_histogram_doane.expect

Bump protobuf version and refactor tensorboard tests (#154244)

2025-05-25 00:50:07 +00:00

TestTensorBoard.test_histogram_fd.expect

Bump protobuf version and refactor tensorboard tests (#154244)

2025-05-25 00:50:07 +00:00

TestTensorBoard.test_hparams_bool.expect

…

TestTensorBoard.test_hparams_number.expect

…

TestTensorBoard.test_hparams_string.expect

…

TestTensorBoard.test_image_with_3_channel_batched.expect

Avoid overflow in tensorboard image summary (#90423)

2022-12-08 08:31:52 +00:00

TestTensorBoard.test_image_with_boxes.expect

Avoid overflow in tensorboard image summary (#90423)

2022-12-08 08:31:52 +00:00

TestTensorBoard.test_image_with_one_channel_batched.expect

Avoid overflow in tensorboard image summary (#90423)

2022-12-08 08:31:52 +00:00

TestTensorBoard.test_image_with_one_channel.expect

Avoid overflow in tensorboard image summary (#90423)

2022-12-08 08:31:52 +00:00

TestTensorBoard.test_image_without_channel.expect

Avoid overflow in tensorboard image summary (#90423)

2022-12-08 08:31:52 +00:00

TestTensorBoard.test_mesh.expect

Bump protobuf version and refactor tensorboard tests (#154244)

2025-05-25 00:50:07 +00:00

TestTensorBoard.test_nested_nn_squential.expect

Remove some unnecessary python functional wrappers (#61608)

2022-02-01 16:59:26 +00:00

TestTensorBoard.test_pr_curve_raw.expect

…

TestTensorBoard.test_pr_curve.expect

…

TestTensorBoard.test_pytorch_graph.expect

…

TestTensorBoard.test_scalar_new_style.expect

Bump protobuf version and refactor tensorboard tests (#154244)

2025-05-25 00:50:07 +00:00

TestTensorBoard.test_text.expect

…

TestTensorBoard.test_video.expect

…

TestTorch.test_is_nonzero-empty.expect

…

TestTorch.test_is_nonzero-multiple.expect

…

TestTorch.test_print-non_contiguous.expect

…