Skip to content

kernel_generation.md

Latest commit

 

History

History
183 lines (131 loc) · 6.92 KB

kernel_generation.md

File metadata and controls

183 lines (131 loc) · 6.92 KB

中文 | English

Kernel Generation with LLMs

Problem Statement

给定一个 kernel 的功能描述 (e.g., "write a fused GELU + bias add Triton kernel for fp16 inputs")，让 LLM 生成高性能 GPU kernel 代码。

挑战

  1. Correctness: 生成的 kernel 必须在数值上正确
  2. Performance: 不仅要正确，还要快 (vs cuBLAS/cuDNN baselines)
  3. Edge cases: 处理非对齐 shape、masking、不同 dtype
  4. Hardware awareness: 不同 GPU 架构有不同的最优策略

KernelBench

KernelBench 是 Stanford Scaling Intelligence Lab 发布的 LLM kernel generation benchmark。

设计

  • 大量 kernel 题目: 从 PyTorch 参考实现生成 CUDA/DSL kernel；规模随版本更新（官方 README 划分为 4 个 Level，例如 Level 1/2 各约 100 题、Level 3 约 50 题，另有 Level 4 面向 Hugging Face 整模）
  • 4 个难度级别（与 KernelBench README 一致）:
    • Level 1: 单算子 / 单 kernel（如 Conv、GEMM、LayerNorm）
    • Level 2: 简单融合模式（如 Conv+Bias+ReLU、Matmul+Scale+Sigmoid）
    • Level 3: 完整模型架构端到端（如 MobileNet、MiniGPT）
    • Level 4: Hugging Face 上的整模架构优化（「Level Hugging Face」）
  • 评估指标:
    • Correctness: 与 PyTorch 参考实现的数值一致性
    • Performance: 相对于 PyTorch (cuBLAS/cuDNN) 的加速比

结果 (截至 2025)

Model Level 1 Pass Rate Level 2 Pass Rate Level 3 Pass Rate
GPT-4o ~50% ~25% ~10%
Claude 3.5 Sonnet ~45% ~20% ~8%
Deepseek-V3 ~40% ~18% ~5%
Open-source 7B ~15% ~5% ~1%

注: 数据为近似值且多来自早期对 Level 1–3 的报告；Level 4 较新，具体结果因 prompt 和评测条件而异。

关键发现

  1. LLM 在简单 kernel (elementwise, reduction) 上表现不错
  2. 复杂 kernel (tiled GEMM, attention) 生成正确代码的成功率显著下降
  3. 性能优化（接近 cuBLAS 水平）仍然困难
  4. 多轮迭代 (generate → test → fix) 显著提高成功率

NVIDIA KernelLLM

NVIDIA 于 2025 年发布的专门针对 GPU kernel generation 的 LLM:

方法

  1. 数据集构建: 从 CUTLASS, Triton, FlashAttention 等开源项目收集 kernel 代码
  2. Fine-tuning: 在 kernel 代码 + 性能数据上微调 LLM
  3. Instruction format: natural language description → kernel code
  4. Iterative refinement: 结合 compiler feedback 和 profiling 数据迭代

能力

  • 生成 Triton kernel (主要目标)
  • 生成 CUDA kernel
  • 分析现有 kernel 的性能瓶颈
  • 建议优化策略

Triton Kernel CodeGen 实践

Prompt Engineering for Kernel Generation

有效的 prompt 应该包含：

1. 功能描述: "Write a Triton kernel for batched matrix multiplication"
2. 输入/输出规格: "A: [B, M, K] fp16, B: [B, K, N] fp16, C: [B, M, N] fp32"
3. 性能约束: "Target 80%+ of cuBLAS on H100"
4. 硬件信息: "NVIDIA H100, 80GB HBM3, SM90"
5. 优化要求: "Use software pipelining with num_stages=3"

Few-Shot Example Template

# Prompt: Generate a Triton kernel for fused bias + GELU activation
# Input: x [M, N] fp16, bias [N] fp16
# Output: y [M, N] fp16
# y = GELU(x + bias)

@triton.jit
def fused_bias_gelu_kernel(
    x_ptr, bias_ptr, y_ptr,
    M, N,
    stride_xm, stride_xn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)

    x = tl.load(x_ptr + offs_m[:, None] * stride_xm + offs_n[None, :] * stride_xn, mask=mask)
    bias = tl.load(bias_ptr + offs_n, mask=offs_n < N)

    x = x + bias[None, :]

    # Approximate GELU
    y = 0.5 * x * (1.0 + tl.math.tanh(0.7978845608 * (x + 0.044715 * x * x * x)))

    tl.store(y_ptr + offs_m[:, None] * stride_xm + offs_n[None, :] * stride_xn, y, mask=mask)

Common Failure Modes

  1. Shape/stride 错误: LLM 经常混淆 row-major / col-major stride
  2. Mask 遗漏: 边界条件处理不当导致 OOB 访问
  3. 数据类型: fp16 accumulation 精度问题
  4. 过度简化: 生成正确但性能差的 naive 实现
  5. 幻觉 API: 使用不存在的 tl.xxx API

Evaluation Framework

评估 kernel generation 的标准流程:

def evaluate_generated_kernel(kernel_fn, ref_fn, input_shapes, dtype):
    # 1. Correctness check
    inputs = [torch.randn(s, dtype=dtype, device="cuda") for s in input_shapes]
    ref_output = ref_fn(*inputs)
    gen_output = kernel_fn(*inputs)
    assert torch.allclose(ref_output, gen_output, rtol=1e-2, atol=1e-3)

    # 2. Performance benchmark
    ref_time = triton.testing.do_bench(lambda: ref_fn(*inputs))
    gen_time = triton.testing.do_bench(lambda: kernel_fn(*inputs))
    speedup = ref_time / gen_time

    return {"correct": True, "speedup": speedup}

HF Upskills: 从交互 Trace 蒸馏 Kernel Skills

Hugging Face 的 upskills 项目展示了一种从 LLM 交互过程中蒸馏"结构化专家知识包"(SKILL.md) 的方法:

流程

  1. 用强模型（teacher，如 Claude）完成 kernel 开发任务 + 记录完整 trace
  2. 从 trace 自动生成 skill 文档（SKILL.md）
  3. 自动生成 test cases
  4. 在小模型（student）上评估 skill 提升（skill lift）

Skill 的本质

Skill = 从 teacher 模型的成功交互中蒸馏出来的结构化专家知识包，包含:

  • Architecture-specific memory access patterns
  • Vectorization strategies
  • Warp shuffle reductions

这意味着 skill + agent = 端到端 CUDA kernel 工程系统:

  • HF kernels 提供 skill 基础设施
  • Upskills 提供 skill 生成方法论
  • Agent 负责在运行时应用 skills

新学习范式

AI 时代学习源码不再是"读懂它"，而是"拆掉它、重建它、打败它":

  1. 执行优先 (Execution First): 不先读代码，先跑起来 + 改输入 + 看输出
  2. 改写 (Rewrite): 用 AI 把复杂系统压缩成最小版本（100 行以内）
  3. 对抗 (Adversarial): 构造 agent 失败的 case，挑战其极限
  4. 抽象 (Abstraction): 提炼 optimization loop、skill 表示、feedback 机制

Reference