flagos-ai / flagos-ai/KernelGen

autotune: Code generator uses incorrect Triton math API paths, causing wasted iterations and final failure for gelu

Open
#78 2 comments 0 reactions 0 assignees View on GitHub
bug flagos2.1-rc2 P1
Dominant language
Python
Stars
80
Forks
13
PR merge metrics
No merged PRs in 30d

Description

## Summary

When running `autotune_kernel` on the `gelu` operator, the code generator does not know the correct Triton API path for `tanh`, wasting 2 rounds on wrong APIs before finding the correct one in round 3 — but the final result still fails accuracy.

## Steps to Reproduce

```
autotune_kernel(operator_name="gelu", target_speedup=1, max_rounds=3)
```

## Actual Behavior

| Round | Generated Code | Error |
|:---:|------|------|
| v1 | `tl.math.tanh(x)` | `AttributeError: module 'triton.language.math' has no attribute 'tanh'` |
| v2 | `tl.libdevice.tanh(x)` | `AttributeError: module 'triton.language' has no attribute 'libdevice'` |
| v3 | `from triton.language.extra import libdevice` + `libdevice.tanh(x)` | Compiled but precision fails (`max_diff ≈ 4.7e-4`), only 5/12 tests pass |

Additionally, the initial code generation also failed once (`初始代码生成失败`), meaning 4 attempts in total with no successful result.

## Expected Behavior

The code generator should use the correct API from round 1:
```python
from triton.language.extra import libdevice
result = libdevice.tanh(x)
```

## Root Cause Analysis

1. **Triton `math.py` (17 functions, NO tanh)**: `exp`, `exp2`, `log`, `log2`, `cos`, `sin`, `sqrt`, `sqrt_rn`, `rsqrt`, `abs`, `fdiv`, `div_rn`, `erf`, `floor`, `ceil`, `fma`, `umulhi`
- Source: https://github.com/triton-lang/triton/blob/main/python/triton/language/math.py

2. **Triton `extra/libdevice.py` (198 functions, INCLUDES tanh)**: Contains `tanh`, `fast_tanhf`, `sinh`, `cosh`, `atan`, `asin`, `acos`, and many more.
- Source: https://github.com/triton-lang/triton/blob/main/python/triton/language/extra/libdevice.py

3. **KernelGen generation prompts** (`kernelgen-generate.md`, `kernelgen-optimize.md`, `SKILL.md`) contain **zero guidance** on Triton math API paths. The LLM relies entirely on its training data, which has inaccurate knowledge of the `tanh` API location in Triton.

## Impact Scope

Any operator that depends on `tl.math.tanh`, `tl.libdevice.*`, or other math functions with ambiguous API paths (e.g., gelu, tanh, sigmoid, and other activation functions) may encounter similar issues.

## Suggested Fix

1. **Short-term**: Add explicit Triton math API path guidance to the generation/optimization prompts — documenting which functions are in `tl.math.*` vs `triton.language.extra.libdevice.*`
2. **Medium-term**: Add an API availability check in the pre-check phase before compilation/testing, so API-missing errors are caught early without consuming iteration rounds

## Test Environment

- Tool: `autotune_kernel`
- Operator: `gelu`
- Parameters: `target_speedup=1`, `max_rounds=3`
- Date: 2026-06-04

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.