Skip to content

[Dev] Rename signature FineGrain in Kernel with MMA as MFMA implementation will be introduced for rocm#289

Merged
LeiWang1999 merged 4 commits intomicrosoft:mainfrom
LeiWang1999:separate_mma
Feb 2, 2025
Merged

[Dev] Rename signature FineGrain in Kernel with MMA as MFMA implementation will be introduced for rocm#289
LeiWang1999 merged 4 commits intomicrosoft:mainfrom
LeiWang1999:separate_mma

Conversation

@LeiWang1999
Copy link
Contributor

This pull request includes significant changes to the matrix multiplication and dequantization scheduling code in the bitblas library. The primary changes involve replacing the existing TensorCore-based schedulers with new Tile Library and MMA-based schedulers. Additionally, several files have been renamed to reflect these changes.

Scheduler Updates:

  • Replaced MatmulBlockScheduler, MatmulFineGrainScheduler, and MatmulWeightPropagationScheduler with MatmulTileLibraryScheduler, MatmulMMAScheduler, and MatmulMMAWeightPropagationScheduler respectively in bitblas/ops/general_matmul/tilelang/dense/__init__.py.
  • Updated MatmulDequantizeBlockScheduler, MatmulDequantizeFineGrainedScheduler, and MatmulDequantizeWeightPropagationScheduler to MatmulDequantizeTileLibraryScheduler, MatmulDequantizeMMAScheduler, and MatmulDequantizeMMAWeightPropagationScheduler respectively in bitblas/ops/general_matmul/tilelang/dequantize/__init__.py.

Function and Class Updates:

  • Modified the is_int4_dtype function to use the new MMA-based schedulers in bitblas/ops/general_matmul/tilelang/dense/__init__.py. [1] [2] [3]
  • Updated MatmulScheduler class to use the new schedulers in bitblas/ops/general_matmul/tilelang/dense/matmul.py.
  • Renamed classes and updated their hints to reflect the new MMA-based schedulers in bitblas/ops/general_matmul/tilelang/dequantize/matmul_dequantize_mma.py. [1] [2] [3]

File Renaming:

  • Renamed matmul_dequantize_tensorcore_finegrained.py to matmul_dequantize_mma.py and updated its imports and class names accordingly. [1] [2]
  • Renamed matmul_dequantize_tensorcore_weight_transform.py to matmul_dequantize_mma_weight_transform.py and updated its imports and class names accordingly. [1] [2]
  • Renamed matmul_dequantize_tensorcore.py to matmul_dequantize_tile.py and updated its imports and class names accordingly. [1] [2] [3]

Test Updates:

  • Updated test files to use the new schedulers in testing/python/operators/test_general_matmul_tilelang_impl.py and testing/python/operators/test_general_matmul_tilelang_kernel.py. [1] [2]

@LeiWang1999 LeiWang1999 merged commit f2aeded into microsoft:main Feb 2, 2025
5 of 6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant