[jason-fxz]

_multi_row_decode_launch halves BLOCK_N while the Q and K/V tiles exceed the opt-in shared memory, and skips that check when the budget is 0. Torch reads shared_memory_per_block_optin only on CUDA, so on ROCm the budget is always 0 and multi-row decode (#900) can pick tiles over the 64 KiB LDS. Reported by @TashaSkyUp in #897 (comment): the BF16 head_dim 256 case of test_decode_triton_attention_multi_row_matches_reference needs 65,792 B on gfx1201.

On ROCm the multi-row launch now takes its budget from shared_memory_per_block (the LDS size per workgroup), as qsa/attend.py already does. Single-row decode, extend and all CUDA launches still read the opt-in value and do not change.

Not run on AMD hardware. I compiled the wrapper's real stage-1 launches for gfx1201 with Triton 3.8 on a CUDA box (64 CUs, opt-in smem 0 as on ROCm): the multi-row tests plus a sweep of 10 head shapes x bf16/fp8/int8 x 1-8 rows x bs 1-64, with and without a window, 469 kernels.

gfx1201 LDS, multi-row kernels	main	this PR
over 64 KiB	4 (D512 bf16, 2 rows, GQA 4: 131,072 B)	0
exactly 64 KiB	34 (D256 / D512 bf16)	0 (largest 36,864 B)
single-row and quantized launches		unchanged

The exactly-64-KiB bf16 tiles move to half the BLOCK_N too: the failing test case is one of them (65,536 B in this cross-compile, 65,792 B on the ROCm toolchain in the report). @zihaomu @TashaSkyUp, could you run pytest tests/kernels/test_triton_attention.py -k multi_row on gfx1201 with this branch?

CUDA: pytest tests/kernels/test_triton_attention.py -k "multi_row or decode_rows or decode_triton" passes on an RTX PRO 6000 (sm_120).