HSTU backward error

[FindHao]

There is an error when I run the following command to test hstu attention ops' backward kernel.

% python hstu_attention_bench.py --bench-forward=False
python: ../../../lib/Tools/LinearLayout.cpp:565: LinearLayout mlir::triton::LinearLayout::reshapeOuts(ArrayRef<std::pair<StringAttr, int32_t>>) const: Assertion `getTotalOutDimSize() == std::accumulate( newOutDims.begin(), newOutDims.end(), 1, [&](int32_t acc, auto &outDim) { return acc * outDim.second; })' failed.
[1]    1368375 IOT instruction (core dumped)  python hstu_attention_bench.py --bench-forward=False

The tested pytorch is the latest nightly version 2.7.0.dev20250303+cu126 installed via pip.

The error changed to the following after I manually installed the newest triton, downloading its main branch and running pip install -e python.

% python hstu_attention_bench.py --bench-forward=False
python: /home/yhao/ptd/triton/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp:1154: mlir::triton::LinearLayout mlir::triton::gpu::{anonymous}::chooseStMatrixLayoutLeadingOffset(mlir::MLIRContext*, mlir::RankedTensorType, int): Assertion `instrN >= numColsPerChunk && "Each chunk is filled in with a single warp"' failed.
...

It is exactly the same with the issue 5609 reported in triton repo.

[xuzhao9]

cc @manman-ren @htyu