[DO NOT MERGE] Integration: #1371 (gfx1150 ROCM_RDNA_ATTN) + #1401 (skinny GEMM spills)

CMakeLists.txt

@@ -1576,12 +1576,12 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
       "csrc/rocm/gdn_chunked.cu")
   endif()
 
-  # RDNA3.5 attention kernels (gfx1151 only).
+  # RDNA3.5 attention kernels (gfx1150, gfx1151, gfx1152, gfx1153).
   set(VLLM_ROCM_HAS_RDNA_ATTN OFF)
-  if("gfx1151" IN_LIST VLLM_GPU_ARCHES)
+  include(cmake/rocm_rdna_attn.cmake)
+  rdna_attn_targets(VLLM_RDNA_ATTN_TARGETS)
+  if(VLLM_RDNA_ATTN_TARGETS)
     set(VLLM_ROCM_HAS_RDNA_ATTN ON)
-    include(cmake/rocm_rdna_attn.cmake)
-    rdna_attn_target(_rocm_C_rdna_attn "gfx1151")
     list(APPEND VLLM_ROCM_EXT_SRC "csrc/rocm/rdna_registry.cu")
   endif()
 
@@ -1634,7 +1634,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GDN_CHUNKED)
   endif()
   if(VLLM_ROCM_HAS_RDNA_ATTN)
-    target_link_libraries(_rocm_C PRIVATE _rocm_C_rdna_attn)
+    target_link_libraries(_rocm_C PRIVATE ${VLLM_RDNA_ATTN_TARGETS})
     target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_RDNA_ATTN)
   endif()
 endif()

benchmarks/attention_benchmarks/runner.py

@@ -506,7 +506,12 @@ def _run_single_benchmark(
     """
     total_q = q_list[0].shape[0]
     v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
-    out = torch.empty(total_q, config.num_q_heads, v_dim, device=device, dtype=dtype)
+    # One output per layer, like the inputs: a single buffer rewritten by every
+    # layer would stay in the cache and its writes would not reach DRAM.
+    out_list = [
+        torch.empty(total_q, config.num_q_heads, v_dim, device=device, dtype=dtype)
+        for _ in range(config.num_layers)
+    ]
 
     def benchmark_fn():
         for i in range(config.num_layers):
@@ -517,7 +522,7 @@ def benchmark_fn():
                 v_list[i],
                 cache_list[i],
                 attn_metadata,
-                output=out,
+                output=out_list[i],
             )
 
     if config.ncu_profile:

benchmarks/kernels/rocm_rdna_hip_attn/README.md

@@ -1,6 +1,7 @@
 # RDNA3.5 attention: adding, tuning and checking a configuration
 
-`ROCM_RDNA_ATTN` (gfx1151, selected by default there) runs decodes
+`ROCM_RDNA_ATTN` (gfx1150..gfx1153, selected by default on gfx1150 and gfx1151)
+runs decodes
 of up to 8 query tokens per sequence on
 `csrc/rocm/rdna_decode_attn.cu` and everything else on
 `csrc/rocm/rdna_prefill_attn.cu`. Its tables, in `csrc/rocm/`:
@@ -12,14 +13,16 @@ of up to 8 query tokens per sequence on
 | `rdna_prefill_configs.h` | prefill builds, one per step of M | one sequence |
 | `rdna_prefill_batch_configs.h` | per batch cell, a prefill build and head group | a batch of sequences |
 
-Calls without a build run on Triton.
+Calls without a build run on Triton. Each file has a table per GPU family:
+gfx1151's (1536 VGPRs a SIMD lane) and gfx1150's (1024), which gfx1152 and
+gfx1153 also use. The tools tune and write the tables of the GPU they run on.
 
 A *configuration* is `(Hq, Hkv, D, window, page size)`, in fp16 and bf16.
 
 ## Requirements
 
-- A gfx1151 GPU and a vLLM build of this tree; run from the tree root with
-  `PYTHONPATH=$PWD`.
+- A gfx1150..gfx1153 GPU and a vLLM build of this tree for it; run from the
+  tree root with `PYTHONPATH=$PWD`.
 - `D` in 64, 96, 128, 256, 512, `Hq % Hkv == 0`, and a page size that is a
   multiple of 16 (the one vLLM uses for the model, e.g. 544 for Qwen3.5-0.8B).
 - Nothing else running on the GPU.
@@ -78,7 +81,11 @@ PYTHONPATH=$PWD .venv/bin/python \
 ```
 
 An entry serves M from its `min_query_len` up to the next entry's. Steps start
-at 16, 32, ..., 4096; the first entry also serves 9 to 15 tokens.
+at 16, 32, ..., 8192 (vLLM's default chunk); the first entry also serves 9 to 15
+tokens. `--steps 8192` tunes only that step; `--apply` keeps the others. A step's
+winner loses at most `--prefill-loss-limit` (2%) to the current entry at any
+cell; `--reselect decisions.jsonl --decisions out.jsonl` chooses stored
+decisions again from their A/B samples, without timing.
 
 ## Batched prefill
 

benchmarks/kernels/rocm_rdna_hip_attn/tools/jit.py

@@ -31,7 +31,14 @@
 # Not the configs headers: their entries do not change a build.
 _SOURCES = (_CSRC / "rdna_decode_attn.cu", *_COMMON)
 _PREFILL_SOURCES = (_CSRC / "rdna_prefill_attn.cu", *_COMMON)
-_ARCH = os.environ.get("RDNA_JIT_ARCH", "gfx1151")
+
+
+@functools.cache
+def _arch() -> str:
+    """The target: RDNA_JIT_ARCH, else the current device's."""
+    arch = os.environ.get("RDNA_JIT_ARCH")
+    return arch or torch.cuda.get_device_properties(0).gcnArchName.split(":")[0]
+
 
 _WRAPPER = """#include "rdna_decode_attn.cu"
 using namespace rdna;
@@ -201,7 +208,7 @@ def _build(v: R.KernelVariant | R.PrefillVariant) -> Path:
         *clang,
         "-x",
         "hip",
-        f"--offload-arch={_ARCH}",
+        f"--offload-arch={_arch()}",
         "-fno-gpu-rdc",
         "-O3",
         "-std=c++20",
@@ -407,17 +414,24 @@ def failure(v: R.KernelVariant | R.PrefillVariant) -> str | None:
     return _failed.get(v)
 
 
-def precompile(variants) -> None:
-    """Build variants in parallel and load them."""
+def precompile(variants) -> int:
+    """Build variants in parallel and load them, sealed or not; the number
+    built."""
+    global _sealed
     todo = [
         dataclasses.replace(v, min_query_len=R.PREFILL_MIN_M)
         if isinstance(v, R.PrefillVariant)
         else v
         for v in variants
     ]
     todo = [v for v in dict.fromkeys(todo) if v not in _loaded]
-    with ThreadPoolExecutor(max_workers=max(1, (os.cpu_count() or 8) - 2)) as pool:
-        list(pool.map(load, todo))
+    sealed, _sealed = _sealed, False
+    try:
+        with ThreadPoolExecutor(max_workers=max(1, (os.cpu_count() or 8) - 2)) as pool:
+            list(pool.map(load, todo))
+    finally:
+        _sealed = sealed
+    return len(todo)
 
 
 def seal() -> None: