Add CUTLASS-based row-wise scaled sparse FP8 kernel

[alexsamardzic]

No description provided.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 6 New Failures

As of commit e1c9bd0 with merge base 867a91f ():

NEW FAILURES - The following jobs have failed:

• Build Linux Wheels / build / upload / manywheel-py3_9-cuda11_8 (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_ao__3.9_cu118_x86_64
• Build M1 Wheels / pytorch/ao / upload / wheel-py3_9-cpu (gh)
  The process '/usr/bin/git' failed with exit code 1
• Code Analysis with Ruff / build (3.9) (gh)
  Process completed with exit code 1.
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 7ad36df1d297152d3619a0f1ef80d7942cae376b0b3c3f5333c4eace574c1d1f /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 53d019a7b85003fae4486dcb64b6ea0d81c58efdefc3e9b81f9d33da0ccc11e0 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  RuntimeError: Command docker exec -t 3bb10a35ba6a87d098fab08439a3df2155e30447f61bc272b11407cb0216471f /exec failed with exit code 1

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

[alexsamardzic]

The kernel is ready and passes smoke test.

Remaining tasks:

• Write a converter to SM90 sparse semi-structured format
• Validate the kernel on proper test inputs
• Check compile warnings about the performance loss, write the benchmark
• Write Python-side code: sparsify/quantize method, Llama generator extension, etc.
• Provide that kernel is built with SM90a flags when torchao detects H100 card as SM90
• Further unify CUDA code with rowwise_scaled_linear_cutlass code
• Implement a meaningful config selection

@cpuhrsch @drisspg

[jcaip]

FYI: I think we have a variant of this in here:

ao/torchao/sparsity/utils.py

Line 26 in cc6244c

def create_semi_structured_tensor(r, c, dtype):