Jesse Cai a8e354a9a0 [sparse][semi-structured] enable fp32 support, separate sparse and dense constraints (#115550) ...

Summary:

Both cuSPASRELt and CUTLASS support 1:2 semi-structured sparsity for
fp32, which this PR enables.(thanks @alexsamardzic).

Furthermore, this PR also updates the sparse_config to take into account
the different shape constraints for sparse and dense matrices.

Technically, cuSPARSELt supports smaller sparse matrix constraints as it
seens to pad to the CUTLASS constraints under the hood. However, in
practice small sparse matrices are not commonly used and we care more
about the dense constraints for LLM inference.

For now, we keep the CUTLASS constraints in place for both cuSPARSELt
and CUTLASS tensors

This PR also reconnects the _FUSE_TRANSPOSE flag for cuSPARSELt tensors.

Test Plan:
```
python test/test_sparse_semi_structured.py
```

Reviewers:

Subscribers:

Tasks:

Tags:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/115550
Approved by: https://github.com/cpuhrsch

2023-12-15 02:28:17 +00:00

..

__init__.py

Fixed docstring errors in _fuser.py, _state.py, __init__.py, _freeze.py, _async.py, _recursive.py, _tensorboard_vis.py, _trace.py, _await.py, _check.py, _serialization.py, _script.py, annotations.py, _monkeytype_config.py (#113371)

2023-11-12 03:19:02 +00:00

_semi_structured_conversions.py

[sparse][semi-structured] enable fp32 support, separate sparse and dense constraints (#115550)

2023-12-15 02:28:17 +00:00

_triton_ops_meta.py

Add instructions for generating optimal Triton kernel parameters of bsr_dense_addmm (#115504)

2023-12-12 16:44:51 +00:00

_triton_ops.py

Add tune_bsr_dense_addmm as an API to find optimal triton kernel parameters for bsr_dense_addmm (#115499)

2023-12-12 16:44:51 +00:00

semi_structured.py

[sparse][semi-structured] enable fp32 support, separate sparse and dense constraints (#115550)

2023-12-15 02:28:17 +00:00