[WIP] Move float8 cutlass sparse layout to Float8SemiSparseTensor #3182

bbeckca · 2025-10-15T22:29:55Z

Moving float8 cutlass sparse layout into its own class:
https://github.com/pytorch/ao/blob/main/torchao/dtypes/floatx/cutlass_semi_sparse_layout.py

cc @jerryzh168 @jcaip @danielvegamyhre

Summary: Moving float8 cutlass sparse layout into its own class: https://github.com/pytorch/ao/blob/main/torchao/dtypes/floatx/cutlass_semi_sparse_layout.py Differential Revision: D84467190

Signed-off-by: Benji Beck <benjibeck@meta.com>

pytorch-bot · 2025-10-15T22:29:59Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 960bc91 with merge base 30082cb ():

NEW FAILURE - The following job has failed:

PR Label Check / Check PR Labels (gh)
Process completed with exit code 1.

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

bbeckca · 2025-10-15T22:38:20Z

torchao/quantization/quantize_/common/packing_format.py

    needed for the rest of the system to understand the specific format that's adopted.
    """
    OPAQUE = "opaque"
+    # todo: add semi-sparse


@jerryzh168 It seems we may want to add a packing format for sparse. Wondering if there's a preference between adding it here or in a separate file (similar to int4) for float8?

Do we need packing format if we have a separate config? It looks like packing format is mostly to support different Int4WeightOnlyConfig kernel options (tinygemm, sparse marlin, etc).

Good point, I noticed that we seem to replace the dense weight with quantized semi-sparse in the transform Would it make more sense to integrate Float8SemiSparseTensor here rather than gating with packing-format as I proposed previously? cc @jerryzh168

Yup I think that the transform should call your subclass.

jcaip · 2025-10-16T14:30:46Z

test/quantization/quantize_/workflows/float8/test_float8_semi_sparse.py

+from torchao.testing.utils import skip_if_rocm
+from torchao.utils import torch_version_at_least
+
+BF16_ACT_CONFIG = Float8WeightOnlyConfig(


I don't think this config makes sense, it's not something we support. From what I understand this is a bf16 a + fp8 sparse weight? We only have kernel support for fp8xfp8 +2:4 sparse matmul, no support for mixed input dtypes currently.

You're right, it seems I should be mirroring test_fp8_cutlass_sparse (from test_sparse_api.py) instead
with the difference being using the new flag/config which exposes the tensor subclass being added?

I think Float8DynamicActivationFloat8SemiSparseWeightConfig should eventually resolve to your subclass.
But I would like to hold off on that until

We're out of the QRT period (cc @RandySheriff just FYI on refactor plans)

We have the same functionality (addmm/mm support)

Sounds good. In that case, should we sequence the changes as follows?

Land Float8SemiSparseTensor with linear support

Add mm / addmm ops for feature parity

Integrate into Float8DynamicActivationFloat8SemiSparseWeightConfig after QRT

@jcaip @jerryzh168 Mind confirming that you're onboard with this direction?

Yeah, that sounds good to me.

jcaip · 2025-10-16T14:32:21Z

torchao/quantization/quantize_/workflows/float8/float8_semi_sparse_tensor.py

+implements_torch_function = Float8SemiSparseTensor.implements_torch_function
+
+
+@implements(aten.linear.default)


We'll also need to make sure mm and addmm are supported ops as well. The arg order is different from linear but it should be the same logic.

Sounds good, I'm onboard with that. Mind if I add those ops in a follow-up diff after this lands?

Yup that's fine with me :)

Signed-off-by: Benji Beck <benjibeck@meta.com>

bbeckca · 2025-10-27T17:35:24Z

test/quantization/quantize_/workflows/float8/test_float8_semi_sparse.py

+            ((2, 32, 128), 256, 128),
+        ],
+    )
+    def test_sparse_vs_dense_fp8(self, sizes):


@jcaip Updated testing to follow style of [test_sparse_apu.py](https://fburl.com/18n157bf. For now I'm omitting any config related changes until QRT, however this diff does include all ops (linear, addmm, mm) so that integration can be done as follow up.

Could I get feedback on the construction of the implementations and test before adding similar for addmm, mm and adding polish?

jcaip

cc @bbeckca I think the implementation and test look good, left a couple of comments.

jcaip · 2025-10-28T13:19:33Z

test/quantization/quantize_/workflows/float8/test_float8_semi_sparse.py

+        )
+        dense_output = torch.nn.functional.linear(input_fp8, weight_fp8, linear.bias)
+
+        weight_sparse_fp8 = Float8SemiSparseTensor.from_hp(linear.weight.data, [1, K])


nit: .detach() instead of .data?

jcaip · 2025-10-28T13:20:56Z

torchao/quantization/quantize_/workflows/float8/float8_semi_sparse_tensor.py

+
+
+class Float8SemiSparseTensor(TorchAOBaseTensor):
+    tensor_data_names = ["sparse", "meta", "scale"]


nit: we should use [ compressed_values, metadata ] instead of sparse and meta here.

jcaip · 2025-10-28T13:23:27Z

torchao/quantization/quantize_/workflows/float8/float8_semi_sparse_tensor.py

+implements_torch_function = Float8SemiSparseTensor.implements_torch_function
+
+
+@implements(aten.t.default)


Why do you have to implement transpose? Transpose on a sparse matrix is kind of tricky, we should probably throw an ValueError if it's called.

jcaip · 2025-10-28T13:24:05Z

torchao/quantization/quantize_/workflows/float8/float8_semi_sparse_tensor.py

+        Float8Tensor,
+    )
+
+    if isinstance(input_tensor, Float8Tensor):


nit: what does this conditional do?

jcaip · 2025-10-28T13:27:17Z

torchao/quantization/quantize_/workflows/float8/float8_semi_sparse_tensor.py

+    else:  # aten.mm.default
+        input_tensor, weight_tensor = args
+        bias = None
+


I think you may need to do some transpose trickery here to support mm and addmm,

my understanding is that linear(x, w, bias) will return xW^t + bias so for mm / addmm you need to pass in a transposed weight

bbeckca added 2 commits October 15, 2025 12:41

Add CutlassSemiSparseFp8Tensor

9b2cfe8

Summary: Moving float8 cutlass sparse layout into its own class: https://github.com/pytorch/ao/blob/main/torchao/dtypes/floatx/cutlass_semi_sparse_layout.py Differential Revision: D84467190

Implement packing and linear

fc80e43

Signed-off-by: Benji Beck <benjibeck@meta.com>

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 Oct 15, 2025

bbeckca commented Oct 15, 2025

View reviewed changes

jcaip reviewed Oct 16, 2025

View reviewed changes

bbeckca added 3 commits October 26, 2025 12:27

Add test for fp8 semi-sparse vs. dense

4199d06

Signed-off-by: Benji Beck <benjibeck@meta.com>

Clean up comments

450d4f6

Signed-off-by: Benji Beck <benjibeck@meta.com>

Apply ruff formatting fix

960bc91

Signed-off-by: Benji Beck <benjibeck@meta.com>

bbeckca force-pushed the fp8-sparse branch from dec85f5 to 960bc91 Compare October 26, 2025 19:35

bbeckca commented Oct 27, 2025

View reviewed changes

jcaip reviewed Oct 28, 2025

View reviewed changes

		implements_torch_function = Float8SemiSparseTensor.implements_torch_function


		@implements(aten.linear.default)



		class Float8SemiSparseTensor(TorchAOBaseTensor):
		tensor_data_names = ["sparse", "meta", "scale"]

		implements_torch_function = Float8SemiSparseTensor.implements_torch_function


		@implements(aten.t.default)

Uh oh!

[WIP] Move float8 cutlass sparse layout to Float8SemiSparseTensor #3182

Are you sure you want to change the base?

[WIP] Move float8 cutlass sparse layout to Float8SemiSparseTensor #3182

Conversation

bbeckca commented Oct 15, 2025

Uh oh!

pytorch-bot bot commented Oct 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 1 New Failure

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jcaip Oct 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jcaip left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot bot commented Oct 15, 2025 •

edited

Loading

jcaip Oct 16, 2025 •

edited

Loading