[roberteg16]

ROCM_RDNA_ATTN for the multimodal encoders on gfx1151: a non-causal build of the prefill kernel behind --mm-encoder-attn-backend ROCM_RDNA_ATTN (opt-in; the default does not change). Gemma 4's and Idefics2's vision towers now reach the selected mm-encoder backend instead of SDPA with a mask, which also helps TRITON_ATTN.

Changes

1. TRITON_ATTN mm-encoder wrapper: no host read under CUDA graph capture. The profiler label's segment lengths ([ROCm] Emit block-diagonal segment lengths in attention profiler scope #1372) read cu_seqlens on the host and broke encoder CUDA graph capture with more than one segment: today cudagraph_mm_encoder fails on Qwen3-VL and Gemma 3, and commits 6 and 7 would extend it to Gemma 4 and MiniCPM-V.

2. Head dims 32, 48, 72 and 80 in the decode and prefill kernels (D=72 padded to 80). They compile for gfx1151 and gfx1201; on gfx12, v_t 2 does not take D=48 or D=72. No causal rows.

3. Non-causal prefill builds: PrefillConfig.causal = 0. Each token sees every key of its segment. Q, K and V are read in place through token strides (a fused qkv needs no copy); K/V buffer loads end at the segment's last key. row_tiles = 2 lets a wave run two 16-row tiles sharing each K/V operand read (non-causal, gfx11 only). New op rdna_encoder_attn.

4. Non-causal rows and batch cells for 21 encoder shapes, gfx1151's tables only (the encoder backend is offered on gfx1151):

   D	Hq/Hkv
   32	4/4, 8/8, 16/16, 32/32
   48	32/32
   64	4/4, 6/6, 8/8, 12/12, 14/14, 16/16, 20/20, 24/24, 64/64
   72, 80	16/16
   128	4/4, 8/8, 12/12, 16/16, 24/24

5. MMEncoderAttention: all of a call's segments in one launch, planned on the host. Calls the kernels cannot serve (dtype, head size, GQA, missing builds) run on TRITON_ATTN. The launches follow the host's segment lengths, so cudagraph_mm_encoder with this backend is refused at startup. Qwen2.5-VL and Qwen3-VL accept the backend.

6. Gemma 4's vision tower on the mm-encoder backend: its towers come from transformers, so --mm-encoder-attn-backend never reached them. Each attention module gets an MMEncoderAttention behind a registered transformers attention function, given each image's patches, then its padding, as segments.

7. Idefics2's vision tower (MiniCPM-V/o): padded slices become segments instead of an additive mask that sent attention to SDPA. TORCH_SDPA and FLASHINFER keep the mask.

Results

TTFT of one-image requests on gfx1151 (Strix Halo), median of 12, a new random image per request, encoder without CUDA graphs. "Today" is what runs on gfx11 with TRITON_ATTN, the default on gfx1151 without flash_attn: SDPA with a mask for Gemma 4 and for MiniCPM-V-4 at 1024x800 (padded slices), TRITON_ATTN for the rest.

Model	Image	Today (ms)	ROCM_RDNA_ATTN (ms)	Change
MiniCPM-V-4	1024x800	468.2	386.2	−17.5 %
	336x336	87.4	85.6	−2.1 %
Gemma 3 4B (w4a16)	1024x800	322.2	301.4	−6.5 %
	336x336	309.9	289.7	−6.5 %
Gemma 4 E4B (AWQ)	1024x800	217.7	209.1	−4.0 %
	336x336	208.8	201.2	−3.6 %
Phi-3.5-vision	1024x800	299.6	294.4	−1.7 %
	336x336	286.7	282.1	−1.6 %
Qwen3-VL-4B (AWQ)	1024x800	486.0	486.2	0 %
	336x336	74.9	75.0	0 %

With TRITON_ATTN, commits 6 and 7 alone take MiniCPM-V-4 at 1024x800 from 468.2 to 403.3 ms and Gemma 4 from 217.7 to 211.8 ms. At D=64 TRITON_ATTN is close to these kernels (Qwen3-VL-4B's tower: 39.5 vs 37.3 ms of attention), so Qwen3-VL-4B does not move.

Kernel, sustained back-to-back: 16/16/{64,72,80} reach 77–94 % of the roofline at L ≥ 1200; the 18 other shapes 78–101 % at L ≥ 2048 (do_bench, which reads ~10–13 % above sustained). Shorter sequences are below 80 %.

Accuracy

MMStar, first 300 questions, TRITON_ATTN vs ROCM_RDNA_ATTN, greedy, no fallbacks (measured on an earlier revision of the integration, same kernels):

Model	Same answer	Accuracy (Triton / ours)
Phi-3.5-vision	299/300	54.0 / 53.7 %
Phi-4-multimodal	298/300	53.3 / 53.7 %
MiniCPM-V (custom)	298/300	62.0 / 62.0 %
MiniCPM-o (vision)	298/300	59.7 / 59.7 %
Gemma 4 E4B	298/300	56.7 / 56.7 %
Gemma 4 26B	293/300	58.7 / 58.3 %

On this revision, the first 20 MMStar questions give the same letter as today's backend in 19–20 of 20 for each model in the results table. Qwen3-Omni's audio encoder (20/20/64) transcribes librispeech samples correctly.

Not covered

• Encoder CUDA graphs with ROCM_RDNA_ATTN (refused at startup).
• PaliGemma2: not validated (gated checkpoint).
• MiniCPM-o's Whisper and Gemma 4's audio tower do not use MMEncoderAttention.
• gfx1150 and gfx1201 tables carry no non-causal rows yet.

Tests

pytest tests/kernels/attention/test_rocm_rdna_hip.py tests/kernels/attention/test_mha_attn.py tests/config/test_multimodal_config.py on gfx1151: 704 passed, 16 skipped.

No open PR covers encoder attention on gfx1151. AI-assisted (Claude Code); the submitter reviews every line.

🤖 Generated with Claude Code