[helion] [regression] Triton update introduced "unknown extension 'SPV_KHR_bfloat16'"

[Stonepia]

Describe the bug

When testing on helion's test, we got an error message:

$cd helion
$python -m pytest -sv test/test_dot.py  -k test_addmm_multiple_small_dims

InvalidModule: Invalid SPIR-V module: input SPIR-V module uses unknown extension 'SPV_KHR_bfloat16'

Reproducer

No Need to install helion. Just run the reproducer is enough.

# Run triton reproducer
cd SPV_KHR_bfloat16_repro
python SPV_KHR_bfloat16_repro.py

Attached files

SPV_KHR_bfloat16_repro.zip

Also attach the log of:

MLIR_ENABLE_DUMP=1 python repro.py

success_triton_3.5.log

failure_triton_3.6.log

Environment details

Workable Triton: pytorch-triton-xpu-3.5.0+git1b0418a9
Failure Triton: 3.6.0+gitd3a28a60 (nightly whl for wheels-pytorch-py3.12-20251126)

Torch: 2.10.0.dev20251123+xpu

Testing machine: PVC
uname -a : Linux 984fee015d7d.jf.intel.com 6.8.0-60-generic #63-Ubuntu SMP PREEMPT_DYNAMIC Tue Apr 15 19:04:15 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux

[anmyachev]

The use of some features is determined at build time, so if you build in one environment and run in another, it won't work:

intel-xpu-backend-for-triton/setup.py

Line 499 in ba1d008

cmake_args.append(f"-DAGAMA_VERSION={agama_version}")

@whitneywhtsang it looks like this solution is worse than we thought.

[whitneywhtsang]

I can reproduce the good case with the latest Triton.
@Stonepia What driver is it using? pure LTS2, or new IGC?