On August 1st, the 9th Meet AI Compiler Technical Salon, hosted by HyperAI, was held in Beijing. This event focused on the latest advancements in AI compilation technology, with several experts from industry and research institutions sharing their insights on programming languages, operator development, compilation optimization, and inference execution, showcasing the collaborative evolution of AI compilers from high-level language expression to hardware execution.

in,Zhu Honglin, a high-performance computing engineer at Tencent, shared his team's algorithm and operator optimization practices for low-complexity matrix multiplication in a presentation titled "FalconGEMM: Surpassing Hardware Peaks with Lower-Complexity Matrix Multiplication".

Faced with the challenge that mature operator libraries such as cuBLAS have pushed matrix multiplication performance to near hardware peak and that the traditional kernel-level optimization space is increasingly narrowing, the team started from the algorithm complexity again. Based on low-complexity matrix multiplication algorithms such as Strassen and AlphaTensor, they built a unified LCMA (Low-Complexity Matrix Algorithms) framework. By combining QDSL, operator fusion, Persistent Kernel, fine-grained scheduling and Cost Model, they transformed the theoretical advantage of "reducing the number of multiplications" into actual performance gains on GPUs.

In the FP16 and BF16 tests of NVIDIA H20, FalconGEMM outperforms cuBLAS on a large number of matrix shapes, with peak performance improvements of approximately 10%-16%, while maintaining numerical accuracy that is basically consistent with standard matrix multiplication in the language model benchmark.

HyperAI has compiled and summarized the shared content without altering its original meaning.

Follow the WeChat official account "HyperAI" and reply with the keyword "" in the background.0801 AI CompilerYou can obtain the authorized speaker's presentation PPT by clicking "...".

Starting with Strassen, we re-examine the optimization space for matrix multiplication.

Matrix multiplication is one of the most important fundamental operators in deep learning, and it usually accounts for the majority of the computational time of the model. Software stacks such as CUDA, MKL, and cuBLAS have been optimized for many years, and in many scenarios, the performance of a single GEMM Kernel is already very close to the hardware peak. This means that if we continue to make local optimizations only at the instruction, pipeline, and memory access levels, the room for improvement is becoming increasingly limited. Therefore, the team turned its attention back to the classic Strassen algorithm.

The Strassen algorithm was proposed by Volker Strassen in 1969. For the most basic 2×2 matrix multiplication, the traditional method requires 8 multiplications, while Strassen, by recombinizing the input matrix, only needs to perform 7 multiplications, and then uses additional addition and subtraction operations to restore the final result, which is equivalent to reducing the amount of multiplication computation by 1/8. If the operands are just scalars, this trade-off is not worthwhile; however, when the operands become submatrices, the complexity difference between matrix addition (O(N²)) and matrix multiplication (O(N³)) becomes significant.This makes "doing one less matrix multiplication and several more matrix additions" begin to have practical value.

If the Strassen algorithm is used recursively, the number of multiplications can be further reduced. For example, a 4×4 block matrix multiplication traditionally requires 64 block multiplications, while two layers of Strassen only require 49. However, increasing the number of recursive layers also introduces more additions, data organization, and memory access overhead.Therefore, practical systems often use only a limited number of layers, seeking a balance between computational reduction and additional overhead.

In 2022, DeepMind's AlphaTensor further expanded this algorithmic space. It transforms matrix multiplication into a Tensor Decomposition problem and uses reinforcement learning to search for decomposition methods with lower rank, demonstrating that in addition to the classic Strassen, there may be a large number of different low-complexity matrix multiplication algorithms under different M, N, and K shapes.

However, there is still a practical problem in moving from algorithm discovery to engineering applications:If each low-complexity algorithm requires a separate handwritten GPU kernel, the development and maintenance costs would obviously be too high.To address this, the team abstracted these algorithms into LCMA (Low-Complexity Matrix Algorithms), which uniformly describes which sub-blocks in the input matrix need to be pre-combined, how many matrix multiplications are actually performed, and how the intermediate results are finally combined into the output matrix. The corresponding implementation is then automatically generated using Codegen.

Therefore, the question arises from "how to implement a Strassen Kernel".The focus shifted to "how to build a unified framework that can support a variety of low-complexity matrix algorithms while maintaining high performance".

At the same time, low-complexity algorithms must also address the issue of numerical precision. While Strassen multiplication is algebraically equivalent to standard matrix multiplication, floating-point operations do not strictly adhere to associativity, and changes in the order of computation can introduce additional rounding errors. Therefore, LCMA, while pursuing performance, also needs to control error propagation in low-precision computations.

With a unified algorithm description in place, the next step was to find a suitable GPU implementation. The team tried CUDA, Triton, TiLang, and QDSL. CUDA offered the strongest hardware control, but when faced with a large number of different LCMA algorithms, registers, shared memory, and intermediate summation structures all required specific adjustments, resulting in high expansion and maintenance costs.

Triton achieves near-CUDA performance in basic Strassen scenarios, but when the algorithm scales to larger block structures, it needs precise reuse of register buffers across multiple intermediate computations, making Triton prone to additional spills. TiLang offers more flexibility in register and shared memory control, but its performance in team tests was still approximately 51-101 TP3T lower than Triton. For basic Strassen, with a theoretical gain of only 12.51 TP3T, this loss is significant enough to erode the algorithm's gains.

final,The team chose QDSL as the main implementation backend for FalconGEMM. QDSL offers development granularity close to CUDA, while also possessing code generation capabilities and supporting embedded PTX. This facilitates the migration of existing high-performance implementations and is suitable for batch code generation based on different LCMA descriptions, providing greater flexibility for subsequent integration and customized optimization.

From LCMA to FalconGEMM, transferring algorithmic gains to GPUs

The most straightforward Strassen GPU implementation can be divided into several steps: combining submatrices A and B to generate 7 new input pairs; executing 7 batched GEMMs; and finally combining the 7 sets of intermediate results into the final matrix C. Compared to a regular GEMM, the truly computationally intensive matrix multiplication part is only 7/8 of the original.Therefore, as long as the extra time spent on pre- and post-processing is less than 1/8 of the saved computational load, there is an overall opportunity to gain benefits.

The team first tested on the NVIDIA H20. The H20, with its high memory bandwidth and relatively low peak computation, is well-suited for this approach of "increasing some data processing in exchange for reduced computational load." At matrix sizes of approximately 2048³ and above, the basic implementation already showed stable gains. However, at smaller shapes, the proportion of input combination, intermediate result write-back, and output combination increases rapidly, easily consuming the saved computational load.

Therefore, the focus of subsequent optimization shifted from GEMM itself to intermediate memory access. The most direct approach is operator fusion.Try to keep intermediate results on the chip as much as possible, rather than repeatedly writing them back to Global Memory.However, combining input A/B directly into GEMM is not suitable because the same sub-block may be used by multiple SMs, easily leading to duplicate loading and summation. In contrast, post-processing fusion of Batched GEMM and Combine H is more feasible.

The real challenge lies in the fact that Strassen's seven intermediate results contribute to the final four output submatrices in different ways. If H is used as the parallel unit, multiple SMs may simultaneously write back to the same C, leading to severe atomic conflicts; if C is used as the parallel unit, some H will be repeatedly calculated by different SMs. Both approaches negate the benefits of reducing multiplication.

The team ultimately abandoned organizing tasks based on Strassen's intermediate results. Instead, they grouped tasks according to the spatial coordinates of the matrix: seven multiplication tiles at the same position in seven batched GEMMs were grouped into a single group and executed on the same SM. This way, after a group completed its calculations, the results could be directly accumulated to the final C on-chip, eliminating the need to write intermediate results back to Global Memory and avoiding significant cross-SM write conflicts.

This fusion approach significantly reduces the additional memory accesses introduced by Strassen, but the larger group granularity leads to load imbalance. For example, with a 4096³ matrix multiplication, coarse-grained scheduling may cause approximately 211 TP3T of additional wave waste, even exceeding the 12.51 TP3T computation reduction of Strassen itself.

to this end,The team drew inspiration from Stream-K to split a Group into two SMs for execution when necessary.The scheduling layer still uses Group as the basic unit, but the actual execution can be further refined to Tile, thereby reducing the idle time of tail SM and improving hardware utilization while retaining the advantages of Group-level data reuse.

However, after resolving load balancing, a new problem arose: L2 cache thrashing. After the groups were split, different types of intermediate multiplications might be mixed within the same wave, and the accessed data was independent, leading to a significant decrease in L2 hit rate. At the same time, GEMM was already heavily utilizing Tensor Cores, and when memory access pressure was nearing full load, the H20 would hit its power limit. In actual testing, the core frequency dropped from approximately 1.8 GHz to 1.6 GHz, resulting in a decrease in computational performance, and some of the benefits of fusion were once again offset.

To address L2 cache thrashing, the team further adjusted the order of the split groups, ensuring that intermediate results of the same type were processed within the same wave, with mixed processing only occurring in a few tail waves. This preserved the load balancing provided by fine-grained scheduling while restoring better L2 data locality, ultimately eliminating the significant frequency reduction issue.

It is worth mentioning thatThe key foundation for these scheduling optimizations is the Persistent Kernel.Unlike a regular kernel where the Task Advisor (CTA) exits after completing a block, a Persistent Kernel allows the CTA to reside on the Streaming Service (SM) for an extended period, continuously accepting subsequent tasks. This gives developers more flexibility in controlling the execution order of Groups and Tiles and enabling on-chip resource reuse. Task splitting, scheduling rearrangement, and caching optimization can therefore also be completed within the same kernel.

From cache reordering to Cost Model, peak performance improvement 10%-16%

After merging, load balancing, and cache rearrangement, FalconGEMM is able to unleash the computational advantages of LCMA on more shapes. However, LCMA is not superior to ordinary GEMM in all cases:Its essence remains the same: trading extra data processing for fewer multiplication calculations.If the original GEMM is already limited by memory access, further reducing computation will not bring sufficient benefits; low-complexity algorithms are only more advantageous when computational density is high.

therefore,The team further designed a Roofline-like Cost Model to determine when to use LCMA and which of the various LCMAs to choose.Since the goal is to "choose the right algorithm" rather than to accurately predict execution time, the model mainly analyzes the computational and memory access costs of different schemes, and estimates the computational/memory access bottleneck range of the target GPU by combining its computing power and bandwidth.

In this model, the reduced number of multiplications in the low-complexity algorithm corresponds to computational gains, while additional data combinations and repeated memory accesses constitute new memory overhead; the aforementioned fusion optimization further reduces this memory access cost. Therefore, FalconGEMM can determine the benefit boundaries of conventional GEMM and different LCMA schemes based on different M, N, and K shapes, and automatically select a more suitable implementation.

By leveraging QDSL's code generation capabilities, the entire framework ultimately forms a relatively complete execution flow: First, a corresponding fused persistent kernel is generated based on the LCMA description, reducing intermediate memory accesses through fusion; then, the Cost Model selects an appropriate matrix multiplication algorithm for the specific shape; finally, QDSL automatically generates and compiles the target code. In this way, LCMA is no longer just a fixed Strassen implementation, but forms an algorithm space that can be dynamically selected based on workload.

Performance testing was primarily conducted on the NVIDIA H20. Results show that in various low-precision matrix multiplication scenarios,FalconGEMM outperforms cuBLAS on a large number of shapes, with peak performance improvements of approximately 10%–16%.After completing the group splitting and cache rearrangement, the large shape can stably obtain the computational benefits brought by the low-complexity algorithm, the performance of the small shape is also improved, and the frequency reduction caused by the power wall triggered by L2 cache jitter is avoided.

Based on the results of the Cost Model selection, the model was able to select the better-performing implementation for most shapes that cross the LCMA return threshold, indicating that the "algorithm selection + fusion kernel" approach can effectively cover matrix multiplication scenarios with different computational densities.

Besides performance, numerical accuracy is also a crucial issue that FalconGEMM must verify. The team encountered significant errors in early low-precision experiments, primarily due to changes in the order of floating-point addition. For example, A+C+B−C is algebraically equal to A+B, but in finite-precision floating-point calculations, they are not necessarily strictly equal.

Further analysis revealed that the truly significant errors primarily stemmed from low-precision casts, rather than the FP32 accumulation itself. In common FP16/BF16 input scenarios, matrix multiplication is typically performed first using FP32 accumulation, then converted back to lower precision; if intermediate results are frequently cast, the FP32 mantissa information is continuously discarded.

The fusion approach actually alleviates this problem. WGMMA's output maintains FP32 precision, while FalconGEMM directly combines and accumulates the final C values on-chip using FP32, only to cast them back to the target precision after the computation is complete. Compared to repeatedly writing back low-precision intermediate results between multiple independent kernels, this method reduces one or more precision conversions, allowing errors caused by changes in computation order to remain more in the lower bits of FP32.

In the language model benchmark, the final scores obtained using FalconGEMM and standard matrix multiplication are almost identical, with only very slight differences, indicating that the current implementation does not lead to a significant decrease in model accuracy.

In the next phase, the team plans to continue in two directions: first, to further integrate Combine A/B by adjusting the order of Batch Groups and K Loops to further reduce memory access for intermediate input results; second, to extend LCMA to Attention. Flash Attention also has a high computation-to-memory ratio, and if low-complexity matrix algorithms can be further combined with its block decomposition and pipeline, it may also bring new performance potential.

From Strassen and AlphaTensor to LCMA and FalconGEMM, the significance of this work goes beyond simply making the already highly optimized GEMM a few percentage points faster. It offers another approach: when the kernel itself is already approaching the hardware limits, performance optimization can not only continue to dig deeper into instructions and pipelines, but also find new room for improvement in algorithm complexity. Then, through compilation, fusion, and scheduling, the theoretical computational reduction can be truly transformed into runtime gains.