flagos-ai / flagos-ai/KernelGen
autotune: Code generator uses incorrect Triton math API paths, causing wasted iterations and final failure for gelu
- 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.