Use transpose_result=True for fp8_sparse_mm (#4433)

[gyllstromk]

Summary:
What: Replace result.t() + result.contiguous() with transpose_result=True in fp8_sparse_mm.

Why: Use native transpose through hipSPARSELt + torch._cslt_sparse_mm.

Reviewed By: ZihaoLiu0927

Differential Revision: D105060886

Pulled By: gyllstromk

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 14b126c with merge base c81eb79 ():
💚 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 D105060886.

[vkuzo]

@claude review this PR

[claude]

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

---

Review of PR #4433: Use transpose_result=True for fp8_sparse_mm

• Gather context and read changed files
• Analyze the changes
• Post review feedback

---

Overall: This PR looks good. The core change is clean and correct — using transpose_result=True in _cslt_sparse_mm avoids a separate .t() + .contiguous() which is a nice performance improvement by leveraging hipSPARSELt's native transpose support.

Changes Summary

This PR contains two commits:

1. 1406dda — Threads alg_id through the config → tensor → kernel path (from PR Add algorithm search to FP8 sparse linear benchmark (#4432) #4432)
2. 16522c3 — The titular change: replaces result.t() + result.contiguous() with transpose_result=True

Detailed Review

The transpose_result change (float8_sparse_2x4_1d_data_1d_metadata_tensor.py:211-222):

# Before:
result = torch._cslt_sparse_mm(packed, dense.t(), bias=bias, alpha=alpha, out_dtype=torch.float32, alg_id=alg_id)
result = result.t()
if to_pad:
    result = result.narrow(0, 0, batch)
return result.contiguous()

# After:
result = torch._cslt_sparse_mm(packed, dense.t(), bias=bias, alpha=alpha, out_dtype=torch.float32, alg_id=alg_id, transpose_result=True)
if to_pad:
    result = result.narrow(0, 0, batch)
return result

This is correct because:

• With transpose_result=True, hipSPARSELt natively produces the transposed result in contiguous memory, avoiding a separate transpose + copy.
• The .contiguous() removal is safe: when to_pad > 0, narrow(0, 0, batch) on a row-major (batch+pad, out_features) tensor produces a contiguous view (narrowing the leading dimension from contiguous storage preserves contiguity).

The alg_id threading — properly plumbed through all layers:

• Float8DynamicActivationFloat8WeightConfig.alg_id → Float8Sparse2x4_1DData1DMetadataTensor.from_hp() → stored as attribute → passed to fp8_sparse_mm custom op → passed to _cslt_sparse_mm
• Listed in optional_tensor_attribute_names for proper serialization via the base class machinery
• Fake tensor registration (_fp8_sparse_mm_fake) correctly updated with the new parameter

Minor Observations

1. No documentation for alg_id on the config class — Float8DynamicActivationFloat8WeightConfig has a docstring (line ~1130) but alg_id isn't described there. Since this is a ROCm/hipSPARSELt-specific parameter, a brief comment in the config class might help users understand what values are valid and when to use it. This is minor and non-blocking.

2. The PR title/description only mentions the transpose_result change but the diff also includes the alg_id additions (from commit 1406dda). This appears to be because commit 1406dda was originally from PR Add algorithm search to FP8 sparse linear benchmark (#4432) #4432 and is stacked here. Not a code issue, just noting for clarity.

LGTM — the changes are correct, well-structured, and improve performance by using the native transpose path.

---

[vkuzo]

looks like there are rebase conflicts, please resolve and we can review again