[ROCm][gfx1151] ROCM_ATTN: HIP decode attention on gfx1151

CMakeLists.txt

@@ -1591,6 +1591,17 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
       "csrc/rocm/gdn_chunked.cu")
   endif()
 
+  # RDNA3.5 decode attention: the variants of csrc/rocm/rdna35_attn/variants.def.
+  # Their code is emitted for gfx1151 only, where they were tuned; other shapes
+  # and devices are served by Triton.
+  set(VLLM_ROCM_HAS_RDNA35_ATTN OFF)
+  if(VLLM_GPU_ARCHES MATCHES "gfx1151")
+    set(VLLM_ROCM_HAS_RDNA35_ATTN ON)
+    include(cmake/rdna35_attn.cmake)
+    rdna35_attn_sources(VLLM_RDNA35_ATTN_SRC)
+    list(APPEND VLLM_ROCM_EXT_SRC ${VLLM_RDNA35_ATTN_SRC})
+  endif()
+
   define_extension_target(
     _rocm_C
     DESTINATION vllm
@@ -1639,6 +1650,9 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
   if(VLLM_ROCM_HAS_GDN_CHUNKED)
     target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GDN_CHUNKED)
   endif()
+  if(VLLM_ROCM_HAS_RDNA35_ATTN)
+    target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_RDNA35_ATTN)
+  endif()
 endif()
 
 # Must run after the last HIP `define_extension_target` so every extension

benchmarks/attention_benchmarks/README.md

@@ -21,6 +21,41 @@ python benchmark.py \
     --output-csv results.csv
 ```
 
+## Cache residency
+
+`benchmark_fn` walks one KV cache per layer and divides the result by the layer
+count, so the layer loop doubles as a rotation over disjoint KV regions. Whether
+a measurement reads DRAM or a resident last-level cache therefore depends on
+`num_layers * kv_bytes_per_layer`, not on one layer's cache.
+
+The default of 10 layers is enough for most shapes. It stops being enough when
+the product falls below the cache, which on a Strix Halo means a 32 MiB MALL
+that CUDA-graph replay never flushes. Measured on that board, TRITON_ATTN
+decode, with and without the option:
+
+| working set (10 layers) | cells | median bias | worst |
+| --- | --- | --- | --- |
+| under 32 MiB | 6 | +35% | +64% |
+| over 32 MiB | 18 | +0% | +3% |
+
+The split is clean: every cell that fits the MALL was reporting a time too low
+to be sustainable, and every cell that does not was already fine. Small models
+at short contexts are the ones affected — `32x8x128` at 512 tokens reads 20 MB
+across ten layers and came out 64% fast.
+
+`--min-working-set-mb` raises `--num-layers` until the KV working set reaches
+the given size. It only ever raises, so shapes already past the cache are
+untouched and their timings do not move.
+
+```bash
+# 20 MB across the default 10 layers: inside the MALL.
+python benchmark.py --backends TRITON_ATTN --batch-specs q1s512 \
+    --num-q-heads 32 --num-kv-heads 8 --head-dim 128      # 11 us, cached
+python benchmark.py --backends TRITON_ATTN --batch-specs q1s512 \
+    --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \
+    --min-working-set-mb 96                               # 18 us, DRAM
+```
+
 ## Simplified Batch Specification Grammar
 
 Express workloads concisely using query length and sequence length:

benchmarks/attention_benchmarks/benchmark.py

@@ -35,6 +35,7 @@
 from dataclasses import replace
 from pathlib import Path
 
+import torch
 import yaml
 from rich.console import Console
 from tqdm import tqdm
@@ -688,6 +689,25 @@ def main():
 
     # Model config
     parser.add_argument("--num-layers", type=int, default=10, help="Number of layers")
+    parser.add_argument(
+        "--dtype",
+        choices=["float16", "bfloat16"],
+        default="float16",
+        help="Q/K/V and output dtype.",
+    )
+    parser.add_argument(
+        "--min-working-set-mb",
+        type=float,
+        default=None,
+        help=(
+            "Grow --num-layers until the KV working set reaches this many MB. "
+            "The layer loop walks one KV cache per layer, so it doubles as a "
+            "rotation over disjoint KV; without this a small model at a short "
+            "context can stay resident in a large last-level cache and report "
+            "bandwidth the memory system cannot sustain. Suggested value on "
+            "Strix Halo (32 MiB MALL): 96."
+        ),
+    )
     parser.add_argument("--head-dim", type=int, default=128, help="Head dimension")
     parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads")
     parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads")
@@ -1087,6 +1107,8 @@ def main():
                         backend=decode_backend,
                         batch_spec=spec,
                         num_layers=args.num_layers,
+                        dtype=getattr(torch, args.dtype),
+                        min_working_set_mb=args.min_working_set_mb,
                         head_dim=args.head_dim,
                         num_q_heads=args.num_q_heads,
                         num_kv_heads=args.num_kv_heads,
@@ -1157,6 +1179,8 @@ def main():
                         backend=backend,  # Will be overridden later
                         batch_spec=batch_spec,
                         num_layers=args.num_layers,
+                        dtype=getattr(torch, args.dtype),
+                        min_working_set_mb=args.min_working_set_mb,
                         head_dim=args.head_dim,
                         num_q_heads=args.num_q_heads,
                         num_kv_heads=args.num_kv_heads,
@@ -1426,6 +1450,8 @@ def main():
                             backend=f"{backend}_{variant_label}",
                             batch_spec=spec,
                             num_layers=args.num_layers,
+                            dtype=getattr(torch, args.dtype),
+                            min_working_set_mb=args.min_working_set_mb,
                             head_dim=args.head_dim,
                             num_q_heads=args.num_q_heads,
                             num_kv_heads=args.num_kv_heads,
@@ -1500,6 +1526,8 @@ def main():
         # Model parameter sweep
         base_config_args = {
             "num_layers": args.num_layers,
+            "dtype": getattr(torch, args.dtype),
+            "min_working_set_mb": args.min_working_set_mb,
             "head_dim": args.head_dim,
             "v_head_dim": args.v_head_dim,
             "num_q_heads": args.num_q_heads,
@@ -1532,6 +1560,8 @@ def main():
         # Unified parameter sweep
         base_config_args = {
             "num_layers": args.num_layers,
+            "dtype": getattr(torch, args.dtype),
+            "min_working_set_mb": args.min_working_set_mb,
             "head_dim": args.head_dim,
             "v_head_dim": args.v_head_dim,
             "num_q_heads": args.num_q_heads,
@@ -1573,6 +1603,8 @@ def main():
                             backend=backend,
                             batch_spec=spec,
                             num_layers=args.num_layers,
+                            dtype=getattr(torch, args.dtype),
+                            min_working_set_mb=args.min_working_set_mb,
                             head_dim=args.head_dim,
                             v_head_dim=getattr(args, "v_head_dim", None),
                             num_q_heads=args.num_q_heads,
@@ -1649,6 +1681,8 @@ def main():
                             backend=decode_backend,
                             batch_spec=spec,
                             num_layers=args.num_layers,
+                            dtype=getattr(torch, args.dtype),
+                            min_working_set_mb=args.min_working_set_mb,
                             head_dim=args.head_dim,
                             num_q_heads=args.num_q_heads,
                             num_kv_heads=args.num_kv_heads,

benchmarks/attention_benchmarks/common.py

@@ -38,6 +38,30 @@ def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
         return (0, 0, 0)
 
 
+def layers_for_working_set(
+    num_layers: int,
+    kv_bytes_per_layer: int,
+    min_working_set_mb: float | None,
+) -> int:
+    """Raise the layer count until the KV working set outgrows the cache.
+
+    ``benchmark_fn`` walks one KV cache per layer and the result is divided by
+    the layer count, so the layer loop doubles as a rotation over disjoint KV
+    regions.  Whether a measurement reads DRAM or a last-level cache therefore
+    depends on ``num_layers * kv_bytes_per_layer``, not on the size of a single
+    layer's cache.
+
+    Left at the default this is a no-op.  Set it on parts with a large
+    last-level cache (for example the 32 MiB MALL on Strix Halo), where a small
+    model at a short context otherwise stays resident across replays and
+    reports bandwidth the memory system cannot sustain.
+    """
+    if not min_working_set_mb or kv_bytes_per_layer <= 0:
+        return num_layers
+    needed = int(min_working_set_mb * 1024**2)
+    return max(num_layers, -(-needed // kv_bytes_per_layer))
+
+
 def run_do_bench(
     benchmark_fn,
     use_cuda_graphs: bool,
@@ -284,10 +308,17 @@ class BenchmarkConfig:
     torch_profile_dir: str | None = None
     torch_profile_iters: int = 3
     warmup_ms: int | None = None
+    # Grow num_layers until the KV working set reaches this, so the timed
+    # kernel reads DRAM instead of a resident last-level cache.
+    min_working_set_mb: float | None = None
 
     # "auto" or "fp8"
     kv_cache_dtype: str = "auto"
 
+    # Sliding window in keys (None: full attention).  Set on the impl and the
+    # KV cache spec, so every backend sees the same windowed layer.
+    sliding_window: int | None = None
+
     # MLA-specific
     prefill_backend: str | None = None
     kv_lora_rank: int | None = None