Technical Summary: Xe-Forge

Problem Statement

Porting deep learning algorithms to new hardware accelerators, specifically Intel GPUs, requires developers to manually apply repetitive low-level optimizations (quantization, memory coalescing, tile tuning, architecture-specific workarounds) to every Triton kernel. This manual effort creates a significant bottleneck: while optimization patterns are consistent, the hardware constraints (warp counts, GRF modes, SLM sizing) vary across devices and are often absent from Large Language Model (LLM) training data. Existing LLM-based systems primarily focus on generating kernels from scratch for NVIDIA or AMD hardware, leaving a critical gap in optimizing existing functionally correct Triton kernels for Intel's specific architecture (Arc/Arc Pro series).

Methodology

The authors present Xe-Forge, a multi-stage, LLM-powered pipeline designed to automate the optimization of existing Triton kernels for Intel GPUs. The system operates through the following core components:

1. Multi-Stage Pipeline Architecture

Xe-Forge processes a functionally correct Triton kernel through up to nine sequential optimization stages. An LLM-based planner determines the execution order based on detected issues, subject to hard dependency constraints (e.g., algorithmic restructuring must precede dtype conversion). The stages include:

1. Analysis: Detects optimization opportunities.
2. Algorithmic Optimization: Restructures computation to reduce FLOPs and memory access (e.g., eliminating redundant GEMMs).
3. Discovery: Applies open-ended, novel optimizations not covered by named stages.
4. Dtype Fix: Converts unnecessary precision (e.g., float64 to float32/16).
5. Fusion: Merges separate kernels or elementwise chains.
6. Memory Access: Optimizes coalescing, prefetching, and data layout.
7. Block Pointers: Replaces manual pointer arithmetic with tl.make_block_ptr.
8. Persistent Kernel: Restructures loops for cross-iteration data reuse.
9. GPU-Specific Tuning: Applies Intel-specific constraints (warp counts, GRF modes, tile sizing, GROUP_SIZE_M swizzling).
10. Autotuning: Generates a @triton.autotune configuration grid.

Stages with no associated issues are skipped to reduce cost and code churn.

2. Chain-of-Verification-and-Refinement (CoVeR)

Each stage is executed by a CoVeR agent (implemented in DSPy) that iteratively generates candidates, validates them, and refines based on feedback. The agent operates in a loop:

• Generate: Proposes code changes.
• Verify: Runs a four-level cascade:
  1. Syntax: Python AST parsing.
  2. Structure: Checks for hardware constraints (e.g., power-of-two warps, valid GRF modes).
  3. Correctness: Compares outputs against a PyTorch reference using torch.allclose.
  4. Performance: Benchmarks on real hardware; rejects if slower than the baseline.
• Refine: If verification fails, error diagnostics are fed back into the context for the next iteration.

3. Curated Knowledge Base

To address the lack of Intel-specific data in LLM training corpora, Xe-Forge utilizes a machine-readable YAML knowledge base. This encodes hard constraints (e.g., valid warp counts, SLM capacity limits) and optimization patterns (before/after code transformations) specific to Intel Xe architectures. This ensures the model stays within architecturally valid bounds.

4. Validation Infrastructure (AI Bench)

The system relies on AI Bench, a standalone framework that decouples measurement from optimization logic. It ensures reproducible timing (using hardware-event timing with L2 cache flushing) and correctness checking. Crucially, the pipeline enforces a kernel harness separation: the LLM can only modify the kernel body, while the Model wrapper, input generation, and execution path remain fixed. This prevents adversarial optimizations where the model bypasses the benchmark logic.

Key Contributions

1. Triton-to-Triton Optimization Formulation: Unlike existing PyTorch-to-Triton generation approaches, Xe-Forge optimizes existing kernels while preserving semantics.
2. First Intel GPU Targeting System: The first LLM-based system targeting Intel GPUs, featuring a hardware-aware knowledge base for constraints absent from training data.
3. Multi-Stage CoVeR Pipeline: A decomposed approach where each stage has dedicated patterns and verification criteria, balancing flexibility with reproducibility.
4. AI Bench Framework: A benchmarking tool that decouples execution/timing from optimization logic, ensuring fair comparison across backends.
5. Adversarial Mitigation: A design that fixes the execution harness to prevent the LLM from "cheating" the benchmark by short-circuiting the model.

Experimental Results

The system was evaluated on 97 Level-2 KernelBench kernels and Flash Attention configurations on the Intel Arc Pro B70 GPU.

• Overall Performance: Achieved a 1.17× geometric mean speedup over PyTorch eager, with 67% of kernels showing improvement.
• Extreme Speedups: Nine kernels exceeded 5× speedup, with the highest reaching 82× (Matmul_Min_Subtract). These gains were driven by algorithmic restructuring that eliminated redundant computation, effectively reducing the FLOP count below the original accounting.
• Flash Attention: Achieved 2× to 13.3× speedups across all 16 tested configurations (varying model sizes, sequence lengths, and batch sizes) without regression. Notably, the system enabled 32k sequence length inference on a 32GB GPU by utilizing persistent kernels and SLM tiling, which was infeasible with the naive kernel.
• Correctness: 100% of optimized kernels passed correctness checks against the PyTorch reference.
• Roofline Analysis: For compute-bound kernels (GEMM/MatMul), the pipeline lifted performance to near-peak efficiency. For bandwidth-bound kernels, performance was generally competitive, though some showed mild regressions (0.5–0.8×) due to memory throughput limits.

Significance and Claims

The paper claims that structured domain knowledge combined with hardware-in-the-loop verification is more effective than model scale alone for hardware-specific kernel optimization.

• Bridging the Knowledge Gap: By encoding Intel-specific constraints (warp counts, GRF modes) in a curated knowledge base, the system compensates for the LLM's lack of training data on these specific architectures.
• Systematic Elimination of Porting Effort: Xe-Forge demonstrates that repetitive porting work can be automated, systematically removing the bottleneck that currently gates algorithm deployment on new accelerators.
• Robustness: The separation of mutable kernel code from immutable harness code effectively prevents the class of "evasion" failures seen in other generation systems, where models bypass constraints to pass correctness checks.

The authors conclude that while the pipeline is not perfect (a minority of kernels show mild regressions, and novel algorithmic patterns remain challenging), it successfully automates the optimization of standard operators and fused compositions that constitute the bulk of production workloads.