# AMX, AVX512, and AVX2 Backends in kt-kernel: Key Differences Explained

> Discover the key differences between AMX, AVX512, and AVX2 backends in kt-kernel. Understand how each targets different CPU instruction sets for optimized performance in your AI applications.

- Repository: [kvcache.ai/ktransformers](https://github.com/kvcache-ai/ktransformers)
- Tags: deep-dive
- Published: 2026-07-20

---

**The AMX, AVX512, and AVX2 backends in kt-kernel share the same high-level MoE algorithm but target different CPU instruction sets—AMX leverages 1024-bit matrix tiles with AVX512-weighted sums on Sapphire Rapids, AVX2 provides a 256-bit SIMD baseline for Haswell+, and AVX512 is exposed as a distinct compile-time flag in the build system.**

kt-kernel is the CPU inference engine inside the `kvcache-ai/ktransformers` repository. It specializes **Mixture-of-Experts** (MoE) dense-and-mixture forward passes into three instruction-set tiers: **AMX**, **AVX512**, and **AVX2**. While the expert routing and gating logic remain identical across builds, the low-level GEMM kernels, activation helpers, and vector widths diverge sharply depending on which CMake target is enabled.

## Instruction Sets and Register Width

The architectural split begins with the hardware primitives each backend uses for matrix math.

- **AMX**: Built on Intel Advanced Matrix Extensions, the core GEMM kernels use the AMX tile API (`_tile_*` instructions) operating on 1024-bit × 1024-bit matrix tiles. As implemented in [`kt-kernel/operators/amx/moe_base.hpp`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/operators/amx/moe_base.hpp), the final weighted-sum reduction still compiles with **AVX512** intrinsics (`__m512` registers), making the AMX path a hybrid of tile acceleration and 512-bit vector math.

- **AVX512**: The CMake build system in [`kt-kernel/setup.py`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/setup.py) (lines 8–19) exposes `-DKTRANSFORMERS_CPU_USE_AVX512=ON` (mapped to `CPUINFER_ENABLE_AVX512=ON`) as a distinct backend flag. The provided source analysis highlights AVX512 primarily through its register width inside the AMX backend’s accumulation loops rather than through a separately documented base class.

- **AVX2**: Defined in [`kt-kernel/operators/avx2/moe_base.hpp`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/operators/avx2/moe_base.hpp), this backend relies on 256-bit AVX2 registers (`__m256`) and standard AVX2 intrinsics (`_mm256_*`). It performs no tile configuration and runs pure blocked SIMD-FMA.

## Source Files and Class Hierarchy

Each backend is rooted in a C++ CRTP base template that dispatches to static helper namespaces.

- **`AMX_MOE_BASE`** lives in [`kt-kernel/operators/amx/moe_base.hpp`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/operators/amx/moe_base.hpp) under the signature `template<class T, class Derived> class AMX_MOE_BASE`.
- **`AVX2_MOE_BASE`** lives in [`kt-kernel/operators/avx2/moe_base.hpp`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/operators/avx2/moe_base.hpp) under the signature `template<class T, class Derived> class AVX2_MOE_BASE`.

Both classes expose the same surface API—`load_weights`, `forward`, and `warm_up`—but route to `amx::*` or `avx2::*` kernel implementations.

## Matrix Multiplication Kernels

The heaviest workload in MoE inference is the gate-up and down projection GEMM. The backends differ fundamentally here.

- **AMX GEMM**: Calls `amx::do_gate_up_gemm` and `amx::do_down_gemm`. These routines load quantized FP16/INT8/INT4 weights into AMX tiles, execute packed multiply-accumulate operations via the tile engine, and spill results back to memory.
- **AVX2 GEMM**: Calls `avx2::do_gate_up_gemm` and `avx2::do_down_gemm`. These are handwritten blocked kernels that use 256-bit vector loads and `_mm256_fmadd_ps` for the FMA loop.

## Activation and Weighted-Sum Paths

After the GEMM phases, each backend applies its own activation and final weighted-sum reduction.

- **AMX activation**: Uses `amx::act_fn`, which processes data with `__m512` vectors and stores back to **BF16**.
- **AVX2 activation**: Uses `avx2::act_fn`, which processes eight BF16 elements at a time inside `__m256` registers.

For the expert-weighted reduction:

- **AMX weighted-sum**: Converts via `avx512_32xbf16_to_32xfp32` into `__m512`, executes FMADD, and converts back through `avx512_32xfp32_to_32xbf16`.
- **AVX2 weighted-sum**: Uses `avx2::load_16xbf16_to_2x8xfp32` to unpack into two `__m256` lanes, applies FMADD, and stores via `avx2::store_fp32_to_bf16`.

## Tile Configuration and Memory Pools

AMX hardware requires tile state management even when the base class abstracts it.

- Inside `AMX_MOE_BASE`, calls to `T::config()` are retained for API compatibility but are effectively no-ops; real tile configuration happens inside derived kernels such as `amx::act_fn`.
- The AMX allocator reserves larger tile-oriented buffers (`gate_up_ba_pool_`, `gate_bc_pool_`, etc.) shaped for the tile engine.
- The AVX2 backend allocates the same logical pool structure, but the buffers are consumed by 256-bit SIMD kernels rather than tile hardware.

## Compile-Time Selection and Python API

Backend selection is static: it happens at compile time via CMake flags parsed in [`kt-kernel/setup.py`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/setup.py), not at runtime.

- Enable AMX: `-DKTRANSFORMERS_CPU_USE_AMX=ON` (environment variable `CPUINFER_ENABLE_AMX=ON`)
- Enable AVX512: `-DKTRANSFORMERS_CPU_USE_AVX512=ON` (environment variable `CPUINFER_ENABLE_AVX512=ON`)
- Enable AVX2: `-DKTRANSFORMERS_CPU_USE_AVX2=ON` (environment variable `CPUINFER_ENABLE_AVX2=ON`)

The compiled wheel exports Python wrappers such as `AMXMoEWrapper` and `AVX2MoEWrapper`. Both accept the same config dictionary and expose identical methods.

```python

# AMX path: requires Sapphire Rapids or newer

import os
os.environ["CPUINFER_ENABLE_AMX"] = "ON"

from kt_kernel import AMXMoEWrapper

config = {
    "expert_num": 8,
    "hidden_size": 2048,
    "intermediate_size": 5120,
    "max_len": 1024,
    "num_experts_per_tok": 2,
    "swiglu_alpha": 0.0,
    "swiglu_limit": 0.0,
}
moe = AMXMoEWrapper(config)
moe.load_weights("path/to/weights")
output = moe.forward(input_tensor, expert_ids, weights)

```

```python

# AVX2 path: fallback for Haswell+ CPUs

import os
os.environ["CPUINFER_ENABLE_AMX"] = "OFF"
os.environ["CPUINFER_ENABLE_AVX2"] = "ON"

from kt_kernel import AVX2MoEWrapper

config = {
    "expert_num": 8,
    "hidden_size": 2048,
    "intermediate_size": 5120,
    "max_len": 1024,
    "num_experts_per_tok": 2,
}
moe = AVX2MoEWrapper(config)
moe.load_weights("path/to/weights")
output = moe.forward(input_tensor, expert_ids, weights)

```

## Performance Targets and Hardware Compatibility

The three backends map directly to Intel CPU generations.

- **AMX**: Targets **Sapphire Rapids** (2023) and newer Xeon CPUs. The tile engine delivers roughly **2×–3× speed-up** for large-batch INT8/INT4 MoE workloads compared to SIMD-only paths.
- **AVX512**: Represents the 512-bit vector tier selectable at compile time; inside the AMX backend it powers the final BF16-to-FP32 reduction loops.
- **AVX2**: Targets **Haswell and newer** (2013+). It provides the widest compatibility but lacks both tile-level parallelism and full 512-bit accumulation.

## Summary

- kt-kernel provides three CPU instruction-set backends—**AMX**, **AVX512**, and **AVX2**—unified under the same MoE algorithm but separated by low-level specialization.
- The **AMX** backend in [`kt-kernel/operators/amx/moe_base.hpp`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/operators/amx/moe_base.hpp) uses 1024-bit AMX tiles for GEMM and AVX512 (`__m512`) intrinsics for the weighted-sum step.
- The **AVX2** backend in [`kt-kernel/operators/avx2/moe_base.hpp`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/operators/avx2/moe_base.hpp) uses 256-bit SIMD (`__m256`) for GEMM, activation, and reduction.
- Build selection is controlled in [`kt-kernel/setup.py`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/setup.py) (lines 8–19) via CMake flags mapped to `CPUINFER_ENABLE_AMX`, `CPUINFER_ENABLE_AVX512`, and `CPUINFER_ENABLE_AVX2`.
- **AMX** requires Sapphire Rapids; **AVX2** runs on Haswell+; **AVX512** is available as a compile-time target and is integral to the AMX reduction path.

## Frequently Asked Questions

### Does kt-kernel support runtime switching between AMX and AVX2?

No. The backend is baked into the compiled extension module. You select the desired target via CMake flags in [`kt-kernel/setup.py`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/setup.py) before building, and the resulting wheel exports the corresponding Python wrapper (e.g., `AMXMoEWrapper` or `AVX2MoEWrapper`). The `CPUINFER_ENABLE_*` environment variables ensure the correct symbols are active at import time.

### Why does the AMX backend use AVX512 intrinsics if AMX has its own tile engine?

The AMX tile engine (`_tile_*` instructions) accelerates the core GEMM kernels (`amx::do_gate_up_gemm`, `amx::do_down_gemm`), but the activation helper `amx::act_fn` and the final weighted-sum reduction operate on 512-bit vector registers (`__m512`). This hybrid design is visible in [`kt-kernel/operators/amx/moe_base.hpp`](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/operators/amx/moe_base.hpp), where `avx512_32xbf16_to_32xfp32` feeds the accumulation before `avx512_32xfp32_to_32xbf16` stores the result.

### What is the minimum CPU generation required for each backend?

- **AVX2**: Intel Haswell or newer (2013+), plus compatible AMD Zen processors.
- **AVX512**: Intel Skylake-X, Ice Lake, Cooper Lake, or Sapphire Rapids server cores.
- **AMX**: Intel Sapphire Rapids (2023) Xeon processors only.

### Are the numerical results identical across AMX, AVX512, and AVX2?

Yes. The repository’s test suite—including files such as [`test_moe_amx_accuracy_int4.py`](https://github.com/kvcache-ai/ktransformers/blob/main/test_moe_amx_accuracy_int4.py) and [`test_moe_avx2_accuracy_fp8.py`](https://github.com/kvcache-ai/ktransformers/blob/main/test_moe_avx2_accuracy_fp8.py)—verifies parity between backends. Equivalent blocking strategies and the same high-level MoE routing ensure that output tensors match regardless of whether GEMM runs on AMX tiles or AVX2 SIMD lanes.