Autotune with fusion panics for output position out of bounds

[laggui]

Describe the bug

Fusion autotune panics for reduce:

2026-05-29T17:21:58.972579Z  INFO burn_train::learner::supervised::strategies::single::epoch: Executing training step for epoch 1
2026-05-29T17:21:58.972630Z  INFO bunsen_preview_chat_dataloader::arrow::parquet: Loading Parquet: "/media/Data/nanochat/shard_00016.parquet"
2026-05-29T17:21:58.997351Z  INFO burn_train::learner::supervised::strategies::single::epoch: Iteration 1
2026-05-29T17:21:59.001445Z  INFO cubecl_runtime::tune::tune_cache: Load autotune cache ...
2026-05-29T17:21:59.001544Z  INFO cubecl_runtime::tune::tune_cache: Loaded 0 autotune cached entries
2026-05-29T17:21:59.001556Z  INFO cubecl_runtime::tune::tuner: Tuning FusedReduceAutotuneKey - ReduceKey: ReduceAutotuneKey { elem_input: Float(F32), elem_output: Float(F32), elem_acc: Float(F32), axis_is_contiguous: true, vector_size: 1024, vector_count: 4096 }, FuseNumReads: 1, FuseNumWrites: 8, FuseNumOps: 8
2026-05-29T17:21:59.136427Z  INFO cubecl_runtime::tune::tune_cache: Load autotune cache ...
2026-05-29T17:21:59.136496Z  INFO cubecl_runtime::tune::tune_cache: Loaded 0 autotune cached entries
2026-05-29T17:21:59.136503Z  INFO cubecl_runtime::tune::tuner: Tuning ReduceAutotuneKey - ElemInput: Float(F32), ElemOutput: Float(F32), ElemAcc: Float(F32), AxisIsContiguous: true, VectorSize: 1024, VectorCount: 4096
2026-05-29T17:21:59.668041Z ERROR burn_train::learner::application_logger: PANIC => panicked at /home/crutcher/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-cubecl-fusion-0.21.0/src/engine/codegen/ir.rs:500:63:
index out of bounds: the len is 1 but the index is 1

Somehow the registered output position is out of bounds in the global args?

To reproduce

• repo: https://github.com/zspacelabs/bunsen
• cmd: cargo run --release -p train-chat --features cuda -- --dataset-dir ~/Data/nanochat --batch-size 2 --shards 0..20

you have to pick a data dir you can drop some shards in.

Expected behavior
Should executed successfully.

Additional context
Originally reported on discord.

[0xAlcibiades]

I've got an independent repro on this.

Repro: a Mamba-3 SSM language model in Burn 0.21

Identical panic:

thread 'DSD-0-0' panicked at burn-cubecl-fusion-0.21.0/src/engine/codegen/ir.rs:500:63:
index out of bounds: the len is 1 but the index is 1

• In our runs the panic follows a FusedMatmulAutotuneKey , not a FusedReduceAutotuneKey. Same ir.rs:500 output-position OOB, so the off-by-one in the registered output position looks general to fused-op codegen, rather than reduce-only. Can reproduce with just cuda, cuda + autotune, cuda + fusion, or cuda + autotune + fusion, so interestingly exists in the general cuda path.

What avoids it: the same model on --features wgpu / metal trains fine. So it's specific to cubecl-cuda fusion.

Env: burn 0.21.0, cubecl 0.10.0; NVIDIA RTX PRO 6000 Blackwell, aarch64 Linux, driver 595.71.05 / CUDA 13.2.

#4878 - may relate to this, i.e. the ir is producing a zero element tensor traversing :500:63 and it's not otherwise guarded.

[laggui]

Ahh thanks for the additional info!

• Can reproduce with just cuda, cuda + autotune, cuda + fusion, or cuda + autotune + fusion, so interestingly exists in the general cuda path.

The panicking code is strictly related to fusion, and requires the fusion feature. But it's enabled by default, so if default features are not explicitly disabled "just cuda" will still use cuda + fusion.

Will try to investigate this issue today.