[ROCm][Attention] Wide MHA decode attention kernel for gfx1151

CMakeLists.txt

@@ -1306,6 +1306,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/rocm/skinny_gemms_w8a8/instantiate_n4.cu"
     "csrc/rocm/skinny_gemms_w8a8/instantiate_n5.cu"
     "csrc/rocm/moe_gemm_w4a16_wmma.cu"
+    "csrc/rocm/rdna35_causal_mha_attn.cu"
     "csrc/rocm/attention.cu")
 
   set(VLLM_ROCM_FLAGS ${VLLM_GPU_FLAGS})

csrc/rocm/ops.h

@@ -142,6 +142,20 @@ void gdn_chunked(torch::Tensor& q, torch::Tensor& k, torch::Tensor& v,
                  torch::Tensor& cu_seqlens, torch::Tensor& out,
                  torch::Tensor& final_state, double scale);
 
+// Causal MHA decode attention for RDNA3.5 (defined in
+// rdna35_causal_mha_attn.cu; real body is gfx11-only, stub elsewhere). Mutates
+// out and the three partial buffers in place. Callers gate the shape (Python
+// rdna35_causal_mha_attn.can_run), so an unsupported one raises via TORCH_CHECK
+// rather than leaving out untouched.
+void rdna35_causal_mha_attn(torch::Tensor& out, torch::Tensor& query,
+                            torch::Tensor& key_cache,
+                            torch::Tensor& value_cache,
+                            torch::Tensor& block_table, torch::Tensor& seq_lens,
+                            torch::Tensor& partial_out,
+                            torch::Tensor& partial_max,
+                            torch::Tensor& partial_sum, double scale,
+                            double softcap);
+
 void paged_attention(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
     torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,

csrc/rocm/torch_bindings.cpp

@@ -184,6 +184,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
       "int n_valid_tokens, int top_k, int block_m, int num_blocks) -> ()");
   rocm_ops.impl("moe_gemm_w4a16", torch::kCUDA, &moe_gemm_w4a16);
 
+  // Wide decode attention for MHA (num_heads == num_kv_heads) on AMD RDNA3
+  // (gfx11). Always registered; the kernel body is gfx11-only (stub elsewhere)
+  // and Python gates calls on on_gfx1151(). Mutates out and the three fp32
+  // partial buffers in place; an unsupported shape raises via TORCH_CHECK
+  // (callers gate via the Python rdna35_causal_mha_attn.can_run predicate).
+  rocm_ops.def(
+      "rdna35_causal_mha_attn(Tensor! out, Tensor query, Tensor key_cache, "
+      "Tensor value_cache, Tensor block_table, Tensor seq_lens, "
+      "Tensor! partial_out, Tensor! partial_max, Tensor! partial_sum, "
+      "float scale, float softcap) -> ()");
+  rocm_ops.impl("rdna35_causal_mha_attn", torch::kCUDA,
+                &rdna35_causal_mha_attn);
+
   // Custom attention op
   // Compute the attention between an input query and the cached
   // keys/values using PagedAttention.