ROCm / triton

Development repository for the Triton language and compiler
MIT License
80 stars 22 forks source link

fix: replace if/else statement with tl.where #545

Open Sara-KS opened 3 months ago

Sara-KS commented 3 months ago

While running the flash attention benchmark in a python environment, I encountered the following error is_mqa = hq != hk off_h_k = off_h_q % hk if is_mqa else off_h_q ^ Triton does not support if expressions (ternary operators) with dynamic conditions, useifstatements instead

Version information ROCm: 5.7.3 python: 3.11 Triton: pytorch-triton-rocm: 2.2.0 System: MI250 (gfx90a)

Solution: Replaced the if/else statement with tl.where function. With this replacement, the benchmark runs without error message.