Add algorithm search to FP8 sparse linear benchmark (#4432)

[gyllstromk]

Summary:
What: Adds algorithm search support (--search-alg) to the FP8 sparse linear benchmark. Threads a new alg_id parameter through the quantization config (Float8DynamicActivationFloat8WeightConfig), the sparse tensor class, and down into the _cslt_sparse_mm kernel call. When --search-alg is passed, the benchmark calls _cslt_sparse_mm_search to find the best algorithm for each shape and benchmarks with that algorithm to report the speedup.

Why: hipSPARSELt supports multiple algorithms for sparse matmul, and the default (alg_id=0) isn't always the fastest. This lets us find the best algorithm for a given shape and measure the performance benefit.

Differential Revision: D102683062

Pulled By: gyllstromk

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 8f60b68 with merge base 918b32d ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

@gyllstromk has exported this pull request. If you are a Meta employee, you can view the originating Diff in D102683062.

[meta-codesync]

@gyllstromk has imported this pull request. If you are a Meta employee, you can view this in D102683062.

[vkuzo]

@claude review this PR

[claude]

Claude finished @vkuzo's task in 2m 0s —— View job

---

Review of PR #4432

• Read changed files and understand the diff
• Review code for correctness, style, and potential issues
• Post review feedback

---

Summary

Clean, well-structured change that threads an alg_id parameter from the config layer down through the sparse tensor class to the _cslt_sparse_mm kernel. The code is consistent with the existing patterns in this file.

Issues

1. Backward-compatibility risk with optional_tensor_attribute_names (minor)

alg_id is added to optional_tensor_attribute_names at line 72 of float8_sparse_2x4_1d_data_1d_metadata_tensor.py. The TorchAOBaseTensor.__setstate__ fallback (in torchao/utils.py:650-656) sets missing optional attributes to None when deserializing old checkpoints. Since alg_id is typed as int (not Optional[int]), a model saved before this change would deserialize with self.alg_id = None, which would then be passed as alg_id=None to torch._cslt_sparse_mm.

This is likely low-risk since this is a new ROCm-only feature and no production checkpoints would have been saved with this tensor type yet, but worth being aware of.

2. alg_id silently ignored for non-sparse packing formats (nit)

The alg_id field is on Float8DynamicActivationFloat8WeightConfig (line 1172 of quant_api.py) but only honored for SPARSE_1D_DATA_1D_METADATA. If a user sets alg_id=5 with PLAIN or SPARSE_2D_DATA_2D_METADATA packing, the value is silently dropped. Consider documenting that this field is only relevant for SPARSE_1D_DATA_1D_METADATA, or adding a warning/assertion.

Looks Good

• The argument ordering in __new__ and __init__ correctly matches tensor_data_names + tensor_attribute_names + optional_tensor_attribute_names, which is required by TorchAOBaseTensor._to_copy.
• The @custom_op signature and its register_fake are kept in sync.
• Default of alg_id=0 preserves existing behavior for all callers not explicitly passing the parameter.
• The dispatch in aten.linear.default correctly reads weight_tensor.alg_id and passes it through.

Verdict

LGTM with the minor comments above. The implementation is correct and follows the existing patterns well.

---

[claude]

Claude finished @vkuzo's task in 1m 23s —— View job

---

Review of PR #4432

• Read changed files and understand the diff
• Review code for correctness, style, and potential issues
• Post review feedback

---

Summary

Clean change that threads an alg_id parameter from Float8DynamicActivationFloat8WeightConfig through Float8Sparse2x4_1DData1DMetadataTensor and down into the _cslt_sparse_mm kernel call. The implementation follows existing patterns well and preserves backward compatibility via alg_id=0 defaults throughout.

Issues

1. Deserialization of old checkpoints sets alg_id=None instead of 0 (minor)

alg_id is listed in optional_tensor_attribute_names (line 72). The TorchAOBaseTensor.__setstate__ fallback (torchao/utils.py:650-656) sets missing optional attributes to None when deserializing old checkpoints. Since the code passes weight_tensor.alg_id directly to torch._cslt_sparse_mm (line 293), a model saved before this change would pass alg_id=None to the kernel.

This is likely low risk since this tensor type is new and ROCm-only, but if you want to be defensive, consider defaulting to 0 in the dispatch:

alg_id=weight_tensor.alg_id or 0,

2. alg_id silently ignored for non-sparse packing formats (nit)

The alg_id field on Float8DynamicActivationFloat8WeightConfig (line 1172) is only threaded through for SPARSE_1D_DATA_1D_METADATA (line 1278). If a user sets alg_id=5 with PLAIN or SPARSE_2D_DATA_2D_METADATA, the value is silently dropped. Consider either documenting this in the field's docstring or adding a warning in __post_init__ when alg_id != 0 and packing_format != SPARSE_1D_DATA_1D_METADATA.

Looks Good

• Argument ordering in __new__ and __init__ correctly matches tensor_data_names + tensor_attribute_names + optional_tensor_attribute_names, required by TorchAOBaseTensor._to_copy.
• The @custom_op signature and its register_fake are kept in sync.
• Default alg_id=0 preserves existing behavior for all callers.
• The dispatch in aten.linear.default correctly reads weight_tensor.alg_id and passes it through.

Verdict

LGTM with the minor comments above. The implementation is correct and follows existing patterns well.

---