Skip to content

[MLAS] Integrate prototype decomposable QGEMM#28745

Draft
JonathanC-ARM wants to merge 3 commits into
microsoft:mainfrom
JonathanC-ARM:jonclo01/decomposable_gemm_integration
Draft

[MLAS] Integrate prototype decomposable QGEMM#28745
JonathanC-ARM wants to merge 3 commits into
microsoft:mainfrom
JonathanC-ARM:jonclo01/decomposable_gemm_integration

Conversation

@JonathanC-ARM

Copy link
Copy Markdown
Contributor

Summary

This change integrates Arm® KleidiAI™ prototype decomposable QGEMM kernels into the MLAS quantized GEMM path, enabling backend-selected acceleration for quantized CPU operators on supported Arm SME2
platforms.

The patch adds a new MLAS QGEMM override hook, threads MLAS_BACKEND_KERNEL_SELECTOR_CONFIG through quantized GEMM packing and execution call sites, and implements KleidiAI-backed {u8,i8} x {u8,i8} QGEMM
paths for both int32 accumulation output and fused float ScaleBias output.

Key Changes

  • Updates the KleidiAI dependency to a prototype Arm GitLab commit containing decomposable GEMM kernel support.
  • Adds an MLAS QGEMM backend override path alongside the existing SGEMM and DynamicQGEMM override mechanisms.
  • Extends MlasGemmBatch, MlasGemm, and MlasGemmPackB to accept backend kernel selector configuration.
  • Implements KleidiAI SME2 QGEMM support for:
    • u8 x u8 -> i32
    • u8 x i8 -> i32
    • i8 x u8 -> i32
    • i8 x i8 -> i32
    • Corresponding fused ScaleBias -> f32 variants
  • Adds packed-B support for the KleidiAI QGEMM path by appending a backend-specific packed representation after the base MLAS packed buffer.
  • Adds output processor introspection so backend overrides can detect and fuse MLAS_QGEMM_SCALE_BIAS_OUTPUT_PROCESSOR without RTTI.
  • Propagates backend selector configuration through quantized CPU operators including:
    • MatMulInteger
    • QLinearMatMul
    • QLinearConv
    • ConvInteger
    • quantized attention
    • dynamic quantized matmul/LSTM
    • RNN helper GEMM paths
  • Restores packed weight size metadata when shared prepacked buffers are reused by QAttention and QLinearConv.
  • Adds a 1D unsigned ConvInteger im2col fast path.
  • Optimizes single-GEMM DynamicQuantizeMatMul execution by calling MlasDynamicQGemm directly instead of the batch API.

Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>
Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>
@hariharans29 hariharans29 changed the title Integrate prototype decomposable QGEMM [MLAS] Integrate prototype decomposable QGEMM Jun 1, 2026
Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>
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