Enhance MM Kernel Performance and Coverage for Specific Input Scenarios

[kiddyjinjin]

PR Category

Operator

Type of Change

Performance Optimization

Description

This PR introduces improvements to the matrix multiplication (MM) kernels to address specific scenarios and optimize performance across different input patterns and sizes.

Key Updates:

1. Support Column-Major Input for Matrix B in Tests and Benchmarks

• Added support to test and benchmark scenarios where the MM input matrix B is column-major.
• Ensured correctness and performance consistency across different input layouts.

2. Expanded Tuning Configurations for Large-K Scenarios

• Introduced additional tuning configs to better handle large-K situations.
• Achieved latency-speedup improvements from 0.5x to 0.85x under large-K conditions.

3. Optimized Kernels for Small M and N Inputs

• Added new kernels specifically optimized for very small M and N inputs.
• These kernels fully utilize L2 cache for faster memory access and reduced latency.

4. Two-Stage Kernel for Large-K Scenarios

• Implemented a split-K two-stage kernel approach:
  • Stage 1: A split-K kernel computes partial results and stores them in an intermediate [SPLIT_K, M, N] matrix.
  • Stage 2: A merging kernel combines these intermediate results into the final [M, N] output matrix.
• This approach improves computational efficiency and resource utilization for large-K scenarios.

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

Large-K Performance with Matrix B as Column-Major

1. Large-K with Matrix B as Column-Major (float16)

• Base Performance:
  
• Current Performance:
  

Small M and N Matrix

1. Small M and N (float16)

• Base Performance:
  
• Current Performance:
• b_row_major
  
• b_column_major
  

1. Small M and N (float32)

• Base Performance:
  
• Current Performance:
• b_row_major
  
• b_column_major
  

1. Small M and N (bfloat16)

• Base Performance:
  

• Current Performance:

• b_row_major
  

• b_column_major