repository·main·Indexed 19 days ago
https://github.com/meta-pytorch/kernelagentAn autonomous multi-agent system for the automated synthesis and optimization of GPU kernels. KernelAgent transforms PyTorch programs into verified, high-performance Triton kernels using LLM-assisted refactoring and hardware-guided optimization loops. The repository includes KernelAgent-Oink, which provides optimized custom operators for Blackwell GPUs (SM100/SM103), such as RMSNorm and fused Add + RMSNorm, with support for vLLM integration and PyTorch custom ops.
TMA is an NVIDIA GPU hardware feature that accelerates memory transfers for tensor operations. It replaces traditional pointer-based memory access with tensor descriptors that describe the entire tensor layout. This allows the hardware to optimize memory transfers automatically.
Key Benefits:
The Fuser pipeline writes all intermediate artifacts to a run directory under .fuse/<run_id>/. This includes:
orchestrator/code.py.tgz: The fused PyTorch refactor.subgraphs.json: Shape-specialized subgraph descriptions.kernels_out/<subgraph_id>/*: Individual KernelAgent sessions for each subgraph.compose_out/composed_kernel.py: The final composed Triton program with a self-test.compose_out/summary.json: Composition metadata.KernelAgent-Oink kernels currently in the repository.The optimization pipeline iteratively improves a verified Triton kernel using a loop of:
The loop terminates when the kernel reaches $\ge 95%$ Speed-of-Light (SOL) or performance converges.
Opt Worker Components are high-level, thin wrappers used by the OptimizationWorker. They wrap low-level utilities from the kernel_perf_agent package to provide specialized functionality for the optimization process, specifically:
OptimizationWorker lifecycle.When extending or debugging the optimization process, note that these components act as an abstraction layer over the core kernel_perf_agent implementation.
The repository is organized into several functional modules:
triton_kernel_agent/: The core logic, including agents, worker managers, provider adapters, and prompt templates.triton_kernel_agent/opt_worker_component/: The optimization pipeline, comprising the profiler, benchmarker, bottleneck analyzer, and orchestrator.kernel_perf_agent/kernel_opt: Utilities for roofline analysis, hardware specifications, and benchmarking.Fuser/: Orchestration pipeline, auto-router, CLIs, and Gradio UIs.triton_kernel_agent/templates/: Jinja templates used for prompting the TritonKernelAgent.examples/: Sample problems and prompt snippets.tests/: Unit tests.e2e_test.py: An example end-to-end kernel generation harness.scripts/: CLI entry points (e.g., Triton UI, autoroute coverage runners) and profiling/benchmark tooling.If you want to inspect existing outputs from the KernelAgent pipeline without running it yourself, you can browse curated artifact repositories:
subgraphs.json), per-subgraph Triton kernels, composed end-to-end Triton programs, and verification logs. View ArtifactsIf using Intel GPUs (Arc, Data Center, or integrated Xe), install the XPU-specific PyTorch build. You must be on Linux with appropriate Intel GPU drivers.
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/xpuHost-side TMA implementation involves creating TensorDescriptor objects on the CPU and passing them as arguments to the Triton kernel. This approach replaces manual pointer arithmetic with descriptor-based loading and storing.
Workflow:
triton.tools.tensor_descriptor.TensorDescriptor to define the tensor, its shape, strides, and the block size for TMA operations.@triton.jit kernel, use the .load([offset_coordinates]) and .store([offset_coordinates], value) methods on the descriptor objects to perform memory operations.from triton.tools.tensor_descriptor import TensorDescriptor
def matmul_with_tma(a, b, c, kernel, grid, BLOCK_SIZE_M, BLOCK_SIZE_K, BLOCK_SIZE_N, num_pid_n):
# Create TMA descriptors on host
a_desc = TensorDescriptor(
a, # the tensor
a.shape, # tensor shape
a.stride(), # tensor strides
[BLOCK_SIZE_M, BLOCK_SIZE_K] # block size for TMA operations
)
b_desc = TensorDescriptor(
b,
b.shape,
b.stride(),
[BLOCK_SIZE_K, BLOCK_SIZE_N]
)
c_desc = TensorDescriptor(
c,
c.shape,
c.stride(),
[BLOCK_SIZE_M, BLOCK_SIZE_N]
)
# Pass descriptors to kernel
kernel[grid](a_desc, b_desc, c_desc, ...)
@triton.jit
def matmul_kernel(a_desc, b_desc, c_desc, ...):
pid = tl.program_id(axis=0)
pid_m = pid // num_pid_n
pid_n = pid % num_pid_n
# Load using TMA descriptors
a = a_desc.load([pid_m * BLOCK_SIZE_M, 0])
b = b_desc.load([0, pid_n * BLOCK_SIZE_N])
# Compute
accumulator = tl.dot(a, b)
# Store using TMA descriptor
c_desc.store([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N], accumulator)You can run a comprehensive benchmark suite that includes the Quack-suite and DeepSeek-V3 (DSv3) workloads, saving all results to a timestamped directory. To include DeepSeek-V4-Flash (DSv4) workloads, use the --include-dsv4 flag.
# Run full Quack-suite + DSv3
conda run -n cute bash -lc 'PYTHONNOUSERSITE=1 CUTE_DSL_ARCH=sm_103 \
python benchmarks/readme/run_sm100_suite.py --dtype bf16'
# Include DeepSeek-V4-Flash norm workloads
conda run -n cute bash -lc 'PYTHONNOUSERSITE=1 CUTE_DSL_ARCH=sm_103 \
python benchmarks/readme/run_sm100_suite.py --dtype bf16 --include-dsv4 \
--out-dir /tmp/oink_sm103_suite_bf16_current'To run SM10x (GB200 / GB300 / Blackwell) microbenchmarks for Oink CuTeDSL kernels, ensure you have a GPU with torch.cuda.get_device_capability()[0] == 10.
Recommended environment variables:
PYTORCH_ALLOC_CONF=expandable_segments:TrueCUTE_DSL_ARCH=sm_103 (for GB300 / SM103)CUTE_DSL_ARCH=sm_100a (for GB200/B200 / SM100 historical runs)To create a pinned GB300 / SM103 benchmark environment using Conda, follow these steps:
conda create -y -n cute python=3.12
conda run -n cute python -m pip install --upgrade pip setuptools wheel packaging ninja
conda run -n cute python -m pip install --upgrade --index-url https://download.pytorch.org/whl/cu130 torch
conda run -n cute python -m pip install 'nvidia-cutlass-dsl==4.4.2' cuda-python triton matplotlib pytest pytest-cov
conda run -n cute python -m pip install -e '.[bench]'
conda run -n cute python -m pip install 'git+https://github.com/Dao-AILab/quack.git' # optional comparison baselineYou can execute the project's unit test suite using pytest:
pytest -v