Research map / Large language models / Agents, code and math
QiMeng-GEMM: Automatically Generating High-Performance Matrix Multiplication Code by Exploiting Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes QiMeng-GEMM, a meta-prompt plus tree-search auto-tuning framework that lets LLMs iteratively generate high-performance GEMM code on RISC-V CPU and NVIDIA GPUs.
Problem
LLMs can produce functionally correct simple code but not high-performance GEMM code, because they do not understand diverse hardware architectures (RISC-V CPU, NVIDIA GPU) and prompts do not unlock their optimization ability.
Approach
Meta-prompts built from five GEMM optimization primitives (tiling, reordering, vectorization, layout, pipeline), each combined with platform-specific hints and code skeletons from hint repositories, plus a beam-style tree search over sequences of primitives where LLM-generated code is evaluated and the top-k are kept at each step.
Key results
- On C910 CPU with GPT-4o, QiMeng-GEMM reaches up to 113.67x speedup over vanilla prompts and 211% of OpenBLAS (10.23 vs 4.85 GFLOPS at size 4096).
- On RTX 4070, the generated code reaches up to 9.65x over vanilla prompts and 115% of cuBLAS (10.55 vs 9.14 TFLOPS at size 768).
- Development time drops from 24 days (junior coder) or 5 days (senior coder) on A100 to 10 minutes, a reported cost reduction of over 240x versus a senior coder.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace a vanilla prompt with QiMeng-GEMM | GPT-4o GEMM code generation · C910 CPU, matrix dimension 4096 | GFLOPS: 0.09 -> 10.23 | Helped |
| Replace a vanilla prompt with QiMeng-GEMM | GPT-4o GEMM code generation · RTX 4070, matrix dimension 8192 | TFLOPS: 1.48 -> 14.28 | Helped |
| Replace a vanilla prompt with QiMeng-GEMM | GPT-4o GEMM code generation · A100, matrix dimension 16384 | TFLOPS: 4.73 -> 18.55 | Helped |
| Use QiMeng-GEMM instead of OpenBLAS 0.3.27 | GEMM on C910 CPU, matrix dimension 4096 · C910 CPU, matrix dimension 4096 | GFLOPS: 4.85 -> 10.23 | Helped |
| Use QiMeng-GEMM instead of cuBLAS 12.1 | GEMM on RTX 4070, matrix dimension 8192 · RTX 4070, matrix dimension 8192 | TFLOPS: 14.25 -> 14.28 | Helped |
| Use QiMeng-GEMM instead of cuBLAS 12.1 | GEMM on A100, matrix dimension 16384 · A100, matrix dimension 16384 | TFLOPS: 19.22 -> 18.55 | Hurt |
| Add the platform-specific hint for block tiling (ablation) | QiMeng-GEMM on A100 GPU · A100 GPU | TFLOPS: 4.74 -> 4.78 | No clear effect |
| Add the platform-specific hint for thread tiling (ablation) | QiMeng-GEMM on A100 GPU · A100 GPU | TFLOPS: 4.31 -> 14.59 | Helped |
| Add the platform-specific hint for double buffering (ablation) | QiMeng-GEMM on A100 GPU · A100 GPU | TFLOPS: 16.72 -> 18.51 | Helped |
| Search optimization-primitive combinations using auto-tuning (ablation) | QiMeng-GEMM on C910 CPU · C910 CPU | performance relative to OpenBLAS: 100% | Helped |
Surprising findings
- Without platform-specific hints, LLMs could not generate correctly executable code for warp tiling, thread tiling, or float4 vectorization on A100.
- On A100, Codestral-22B with QiMeng-GEMM gave only about 1.0x over its vanilla prompt, while larger models gained several-fold.
Limitations
- Gains depend on the LLM; smaller open models such as Codestral-22B show little or no improvement on GPU.
- Results are limited to three hardware platforms and a set of matrix dimensions, and the paper says more architectures are future work.
Experimental setup
Datasets: matrix dimensions from deep learning models such as Llama2-7b and Mistral-7B
Benchmarks: GEMM performance on XuanTie C910 RISC-V CPU, NVIDIA RTX 4070, and A100 GPU across matrix sizes
Metrics: GFLOPS TFLOPS speedup over vanilla prompt development time
Baselines: vanilla prompt CoT prompt OpenBLAS cuBLAS Junior Coder Senior Coder
Builds on: chain-of-thought prompting
Closest papers by meaning
- LibAMM: Empirical Insights into Approximate Computing for Accelerating Matrix Multiplication (NeurIPS 2024 · similarity 0.89)
- CodeGEMM: A Codebook-Centric Approach to Efficient GEMM in Quantized LLMs (NeurIPS 2025 · similarity 0.86)
- LUT-GEMM: Quantized Matrix Multiplication based on LUTs for Efficient Inference in Large-Scale Generative Language Models (ICLR 2024 · similarity 0.84)
- ShiftAddLLM: Accelerating Pretrained LLMs via Post-Training Multiplication-Less Reparameterization (NeurIPS 2024 · similarity 0.84)
- HOT: Hadamard-based Optimized Training (CVPR 2025 · similarity 0.84)
- IM-Unpack: Training and Inference with Arbitrarily Low Precision Integers (ICML 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.