[ROCm] Octave KV (ROCM_OCTAVE): native 3-bit KV cache backend for AMD GPUs (#59774)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>

.buildkite/test-amd.yaml

@@ -3709,6 +3709,25 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt
 
+- label: ":amd: (MI300) LM Eval Octave KV Cache"
+  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - tests/evals/gsm8k/
+  - tests/utils.py
+  - tests/conftest.py
+  - vllm/v1/attention/backends/rocm_octave_attn.py
+  - vllm/v1/attention/ops/rocm_octave.py
+  - vllm/v1/attention/selector.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-octave.txt
+
 - label: ":amd: (MI355) LM Eval Spec Decode (Fixed-length FP8)" # TBD
   timeout_in_minutes: 145
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

.buildkite/test_areas/lm_eval.yaml

@@ -1121,6 +1121,27 @@ steps:
       commands:
       - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt
 
+# Octave kernels are built only for gfx942 and gfx1100, so this runs on MI300.
+- label: ":amd: (MI300) LM Eval Octave KV Cache"
+  key: amd-lm-eval-octave-mi300
+  dind: false
+  device: mi300_1
+  num_devices: 1
+  timeout_in_minutes: 45
+  working_dir: /vllm-workspace/tests
+  depends_on:
+  - image-build-amd
+  source_file_dependencies:
+  - tests/evals/gsm8k/
+  - tests/utils.py
+  - tests/conftest.py
+  - vllm/v1/attention/backends/rocm_octave_attn.py
+  - vllm/v1/attention/ops/rocm_octave.py
+  - vllm/v1/attention/selector.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-octave.txt
+
 - label: ":nvidia: (H100) [Nightly] GPQA Eval (GPT-OSS)"
   key: gpqa-eval-gpt-oss-2xh100
   timeout_in_minutes: 35

CMakeLists.txt

@@ -368,6 +368,20 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
   set(VLLM_EXT_SRC
     "csrc/rocm/quickreduce/torch_bindings.cpp"
     "csrc/rocm/quickreduce/custom_quickreduce.cu")
+  # Octave KV cache (store, decode, rotation, prefill): only built for the
+  # validated architectures, gfx1100 (RDNA3) and gfx942 (MI300).
+  set(VLLM_OCTAVE_KV OFF)
+  if(VLLM_GPU_ARCHES MATCHES "gfx942|gfx1100")
+    set(VLLM_OCTAVE_KV ON)
+    list(APPEND VLLM_EXT_SRC
+      "csrc/attention/octave_torch_bindings.cpp"
+      "csrc/attention/octave_attn.cu"
+      "csrc/attention/octave_decode_wmma_rdna3.cu"
+      "csrc/attention/octave_prefill_rdna3.cu"
+      "csrc/attention/octave_prefill_cdna.cu"
+      "csrc/attention/octave_decode_mfma_cdna.cu"
+    )
+  endif()
 
   message(STATUS "Enabling C extension.")
   define_extension_target(
@@ -379,6 +393,13 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     ARCHITECTURES ${VLLM_GPU_ARCHES}
     USE_SABI 3
     WITH_SOABI)
+  if(VLLM_OCTAVE_KV)
+    # MFMA accumulators in VGPRs: in AGPRs every online-softmax rescale of the
+    # output copies it through 64 extra VGPRs and halves the occupancy.
+    set_source_files_properties(
+      "${CMAKE_CURRENT_BINARY_DIR}/csrc/attention/octave_decode_mfma_cdna.hip"
+      PROPERTIES COMPILE_FLAGS "-mllvm -amdgpu-mfma-vgpr-form")
+  endif()
 endif() # _C HIP endif
 
 if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

benchmarks/attention_benchmarks/benchmark.py

@@ -32,6 +32,7 @@ import subprocess
 import sys
 from dataclasses import replace
 from pathlib import Path
+from typing import get_args
 
 import yaml
 from rich.console import Console
@@ -50,6 +51,7 @@ from common import (
     is_mla_backend,
 )
 
+from vllm.config.cache import CacheDType
 from vllm.v1.worker.workspace import init_workspace_manager
 
 
@@ -548,6 +550,11 @@ def main():
 
     # Benchmark settings
     parser.add_argument("--device", default="cuda:0", help="Device")
+    parser.add_argument(
+        "--model",
+        default="meta-llama/Meta-Llama-3-8B",
+        help="Model whose config the backend is built with",
+    )
     parser.add_argument(
         "--warmup-ms",
         type=int,
@@ -561,8 +568,8 @@ def main():
     parser.add_argument(
         "--kv-cache-dtype",
         default="auto",
-        choices=["auto", "fp8"],
+        choices=get_args(CacheDType),
-        help="KV cache dtype: auto or fp8",
+        help="KV cache dtype",
     )
     parser.add_argument(
         "--cuda-graphs",
@@ -894,6 +901,7 @@ def main():
                         device=args.device,
                         profile_memory=args.profile_memory,
                         kv_cache_dtype=args.kv_cache_dtype,
+                        model=args.model,
                         use_cuda_graphs=args.cuda_graphs,
                         prefill_backend="fa4",
                     )
@@ -964,6 +972,7 @@ def main():
                         device=args.device,
                         profile_memory=args.profile_memory,
                         kv_cache_dtype=args.kv_cache_dtype,
+                        model=args.model,
                         use_cuda_graphs=args.cuda_graphs,
                         ncu_profile=args.ncu_profile,
                         warmup_ms=args.warmup_ms,
@@ -1210,6 +1219,7 @@ def main():
                             device=args.device,
                             max_model_len=getattr(args, "max_model_len", None),
                             kv_cache_dtype=args.kv_cache_dtype,
+                            model=args.model,
                             profile_memory=args.profile_memory,
                             use_cuda_graphs=args.cuda_graphs,
                             ncu_profile=args.ncu_profile,
@@ -1355,6 +1365,7 @@ def main():
                             device=args.device,
                             profile_memory=args.profile_memory,
                             kv_cache_dtype=args.kv_cache_dtype,
+                            model=args.model,
                             use_cuda_graphs=args.cuda_graphs,
                             ncu_profile=args.ncu_profile,
                             warmup_ms=args.warmup_ms,

benchmarks/attention_benchmarks/common.py

@@ -284,8 +284,10 @@ class BenchmarkConfig:
     torch_profile_iters: int = 3
     warmup_ms: int | None = None
 
-    # "auto" or "fp8"
+    # Any CacheDType: "auto", "fp8", "int8_per_token_head", ...
     kv_cache_dtype: str = "auto"
+    # Model whose config the backend sees (dtype, rotary dims, ...)
+    model: str = "meta-llama/Meta-Llama-3-8B"
 
     # MLA-specific
     prefill_backend: str | None = None

benchmarks/attention_benchmarks/runner.py

@@ -36,6 +36,7 @@ from vllm.config import (
     set_current_vllm_config,
 )
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.attention.backends.utils import (
     CommonAttentionMetadata,
     get_supported_kv_cache_layouts,
@@ -48,6 +49,7 @@ from vllm.v1.kv_cache_interface import (
     compute_layer_kv_cache_shape_bytes,
     compute_layout_strides,
     create_kv_cache_views,
+    get_kv_quant_mode,
 )
 
 # ============================================================================
@@ -142,15 +144,16 @@ def _build_common_attn_metadata(
 def _create_vllm_config(
     config: BenchmarkConfig,
     max_num_blocks: int,
+    max_model_len: int = 0,
 ) -> VllmConfig:
     """Create a VllmConfig for benchmarking with mock model methods."""
     model_config = ModelConfig(
-        model="meta-llama/Meta-Llama-3-8B",
+        model=config.model,
-        tokenizer="meta-llama/Meta-Llama-3-8B",
+        tokenizer=config.model,
         trust_remote_code=False,
         dtype="auto",  # Use model's native dtype
         seed=0,
-        max_model_len=1024,
+        max_model_len=max(1024, max_model_len),
     )
 
     cache_config = CacheConfig(
@@ -164,7 +167,7 @@ def _create_vllm_config(
     scheduler_config = SchedulerConfig(
         max_num_seqs=256,
         max_num_batched_tokens=8192,
-        max_model_len=8192,
+        max_model_len=max(8192, max_model_len),
         is_encoder_decoder=False,
         enable_chunked_prefill=True,
     )
@@ -342,6 +345,22 @@ def _create_input_tensors(
     return q_list, k_list, v_list
 
 
+def _kv_cache_spec(config: BenchmarkConfig, dtype: torch.dtype) -> FullAttentionSpec:
+    if config.kv_cache_dtype == "auto":
+        cache_dtype = dtype
+    elif config.kv_cache_dtype.startswith("fp8"):
+        cache_dtype = current_platform.fp8_dtype()
+    else:
+        cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[config.kv_cache_dtype]
+    return FullAttentionSpec(
+        block_size=config.block_size,
+        num_kv_heads=config.num_kv_heads,
+        head_size=config.head_dim,
+        dtype=cache_dtype,
+        kv_quant_mode=get_kv_quant_mode(config.kv_cache_dtype),
+    )
+
+
 def _create_kv_cache(
     config: BenchmarkConfig,
     max_num_blocks: int,
@@ -351,16 +370,7 @@ def _create_kv_cache(
     layout: KVCacheLayout,
 ) -> list:
     """Create KV cache tensors for all layers using the standard allocator."""
-    if config.kv_cache_dtype.startswith("fp8"):
+    spec = _kv_cache_spec(config, dtype)
-        cache_dtype = current_platform.fp8_dtype()
-    else:
-        cache_dtype = dtype
-    spec = FullAttentionSpec(
-        block_size=config.block_size,
-        num_kv_heads=config.num_kv_heads,
-        head_size=config.head_dim,
-        dtype=cache_dtype,
-    )
     # Apply the backend's page customization, as the worker does for the real spec.
     spec = backend_class.customize_spec(spec)
     total_bytes = (
@@ -492,7 +502,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
     # Suppress vLLM logs during setup to reduce spam
     with log_warnings_and_errors_only():
         # Create vllm_config first - uses model's native dtype via "auto"
-        vllm_config = _create_vllm_config(config, max_num_blocks)
+        vllm_config = _create_vllm_config(config, max_num_blocks, max_kv)
         dtype = vllm_config.model_config.dtype
 
         # Wrap everything in set_current_vllm_config context
@@ -510,12 +520,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 q_lens, kv_lens, config.block_size, device
             )
 
-            kv_cache_spec = FullAttentionSpec(
+            kv_cache_spec = backend_class.customize_spec(_kv_cache_spec(config, dtype))
-                block_size=config.block_size,
-                num_kv_heads=config.num_kv_heads,
-                head_size=config.head_dim,
-                dtype=dtype,
-            )
 
             builder = _create_metadata_builder(
                 backend_class, kv_cache_spec, vllm_config, device, config.backend

csrc/attention/octave_arch.cuh

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Per-architecture primitives of the Octave kernels. Every AMD target vLLM
+// builds compiles the portable tier; faster tiers are selected at runtime by
+// the host code (see docs/design/rocm_octave.md, "Architecture contract").
+//
+//   int8 dot4: v_dot4_i32_iu8 on gfx11/gfx12, v_dot4_i32_i8 on gfx9/gfx10.3
+//   decode, fast tier: octave_decode_wmma_rdna3.cu (WMMA, gfx11)
+//   prefill: octave_prefill_rdna3.cu (WMMA, gfx11)
+//   decode, portable: octave_attn.cu decode_kernel (dot4, every target)
+//
+// The portable kernels are written for 32-lane logical waves: shuffles never
+// cross 32 lanes, so they are correct on wave64 (CDNA) as well.
+
+#pragma once
+
+#include <cstdint>
+
+namespace octave {
+
+// sum_i a.i8[i] * b.i8[i] + c, signed bytes on both sides.
+__device__ __forceinline__ int dot4_i8(int a, int b, int c) {
+#if defined(__GFX11__) || defined(__GFX12__)
+  return __builtin_amdgcn_sudot4(true, a, true, b, c, false);
+#elif defined(__gfx942__)
+  return __builtin_amdgcn_sdot4(a, b, c, false);
+#else
+  #pragma unroll
+  for (int i = 0; i < 4; ++i)
+    c += (int)(int8_t)(a >> (8 * i)) * (int)(int8_t)(b >> (8 * i));
+  return c;
+#endif
+}
+
+}  // namespace octave

csrc/attention/octave_decode_mfma_cdna.cu

@@ -0,0 +1,719 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Octave decode on MFMA for CDNA3 (gfx942). Same scheme as the WMMA decode
+// (octave_decode_wmma_rdna3.cu): independent waves, one (split, 16-query row
+// tile) per wave, transposed orientation, codebook bytes fed to fp16 as
+// x + 1152 and debiased afterwards. With 16x16x16 MFMA on wave64 (lane l
+// holds column l % 16 and rows 4 (l / 16) .. + 3 of the result):
+//
+//   QK  S^T[token][q] = K x Q^T. A lane owns one token (l % 16) and, in the
+//       four k-steps of each 64-dim block, the 16 contiguous dims
+//       16 (l / 16) .. + 15 of that block, so its codes are one contiguous
+//       run of the slot. One accumulator per K scale.
+//   PV  O^T[dim][q] = V^T x P^T. Output row r of fragment c is dim 16 r + c,
+//       so a lane reads the 16 contiguous dims 16 (l % 16) .. + 15 of tokens
+//       4 (l / 16) .. + 3. P^T is the S^T accumulator, one query per lane.
+
+#include <algorithm>
+#include <cstdint>
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#if defined(USE_ROCM)
+  #include <hip/hip_runtime.h>
+  #include <hip/hip_bf16.h>
+  #include <hip/hip_fp16.h>
+
+  #include "octave_format.cuh"
+
+  #if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx942__)
+    #define OCTAVE_MFMA_STUB
+  #endif
+
+namespace octave_mfma {
+
+typedef _Float16 h4v __attribute__((ext_vector_type(4)));
+typedef _Float16 h8v __attribute__((ext_vector_type(8)));
+typedef _Float16 h2v __attribute__((ext_vector_type(2)));
+typedef float f4v __attribute__((ext_vector_type(4)));
+typedef int i4v __attribute__((ext_vector_type(4)));
+typedef const int __attribute__((address_space(4))) cint;
+
+using octave::D;
+using octave::Format;
+using octave::LUT_ONE;
+[[maybe_unused]] constexpr int QROW = D + 8;
+[[maybe_unused]] constexpr float PSCALE = 256.0f;
+[[maybe_unused]] constexpr float BIAS = 1152.0f;  // byte x fed as x + 1152
+[[maybe_unused]] constexpr uint32_t X = 0x80808080u;
+// Bytes of a 16-token tile (slot <= 256).
+[[maybe_unused]] constexpr int TILE_MAX = 4096;
+
+// Bytes b of a dword -> fp16 1024 + b, exact.
+__device__ __forceinline__ h4v bytes_to_h4(uint32_t w) {
+  const uint32_t lo = __builtin_amdgcn_perm(0x64646464u, w, 0x07010700u);
+  const uint32_t hi = __builtin_amdgcn_perm(0x64646464u, w, 0x07030702u);
+  return __builtin_bit_cast(h4v, (uint2){lo, hi});
+}
+
+template <bool BIASED>
+__device__ __forceinline__ uint32_t lut4_fast(uint32_t codes) {
+  constexpr uint32_t B = BIASED ? X : 0u;
+  const uint32_t lo = __builtin_amdgcn_perm(
+      octave::LUT4_1 ^ B, octave::LUT4_0 ^ B, codes & 0x07070707u);
+  const uint32_t hi = __builtin_amdgcn_perm(
+      octave::LUT4_3 ^ B, octave::LUT4_2 ^ B, codes & 0x07070707u);
+  const uint32_t t = codes & 0x08080808u;
+  const uint32_t m = (t << 5) - (t >> 3);
+  return (hi & m) | (lo & ~m);
+}
+
+// 4 dwords of 4 tokens each (same 4 dims) -> 4 dwords of 4 dims each.
+__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2,
+                                    uint32_t a3, uint32_t o[4]) {
+  const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u);
+  const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u);
+  const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u);
+  const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u);
+  o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u);
+  o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u);
+  o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u);
+  o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u);
+}
+
+template <typename T>
+__device__ __forceinline__ float to_f(T x) {
+  return (float)x;
+}
+template <>
+__device__ __forceinline__ float to_f<half>(half x) {
+  return __half2float(x);
+}
+
+__device__ __forceinline__ float sign_of(const int* bits, int i) {
+  return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f;
+}
+__device__ __forceinline__ float lo16(uint32_t x) {
+  return __half2float(__ushort_as_half((unsigned short)(x & 0xFFFF)));
+}
+__device__ __forceinline__ float hi16(uint32_t x) {
+  return __half2float(__ushort_as_half((unsigned short)(x >> 16)));
+}
+// Over the four lanes of a query (l, l ^ 16, l ^ 32, l ^ 48).
+__device__ __forceinline__ float qmax4(float v) {
+  v = fmaxf(v, __shfl_xor(v, 16));
+  return fmaxf(v, __shfl_xor(v, 32));
+}
+__device__ __forceinline__ float qsum4(float v) {
+  v += __shfl_xor(v, 16);
+  return v + __shfl_xor(v, 32);
+}
+
+// QG: query tokens per block (4 for multi-token verification, else 1).
+template <int VB, bool KC, int NW, int QG, typename QT>
+__global__ __launch_bounds__(64 * NW) void decode_mfma(
+    const QT* __restrict__ Q, const uint8_t* __restrict__ cache,
+    const int* __restrict__ block_table, const int* __restrict__ q_to_req,
+    const int* __restrict__ q_to_klen, float* __restrict__ mid_o,
+    const int* __restrict__ k_signs, float sm_scale, int num_q, int num_q_heads,
+    int num_kv_heads, int block_size, int max_blocks, int num_reqs,
+    int num_phys_blocks, int num_splits, int min_tps, int num_row_tiles,
+    int64_t sq0, int64_t sq1, int64_t scb, int64_t sch, int64_t sct,
+    int64_t smo, int64_t smh, int64_t sms) {
+  #ifndef OCTAVE_MFMA_STUB
+  using F = Format<VB, KC>;
+  constexpr int NG = KC ? 1 : 4;  // QK accumulators (one per K scale)
+  const int grp = blockIdx.x;
+  const int kvh = blockIdx.y;
+  const int tid = threadIdx.x;
+  const int w = __builtin_amdgcn_readfirstlane(tid >> 6), lane = tid & 63;
+  const int col = lane & 15, gq = lane >> 4;
+  const int rt = num_row_tiles == 2 ? (w & 1) : 0;
+  const int splits_per_block = NW / num_row_tiles;
+  const int si =
+      blockIdx.z * splits_per_block + (num_row_tiles == 2 ? w >> 1 : w);
+  const int hpk = num_q_heads / num_kv_heads;
+  const int max_kv = max_blocks * block_size;
+  const float sm_log2 = sm_scale * 1.4426950408889634f / LUT_ONE;
+
+  // The staged queries, then (once they are in registers) one 4 KB tile
+  // buffer per wave.
+  constexpr int SMEM =
+      32 * QROW * 2 > NW * TILE_MAX ? 32 * QROW * 2 : NW * TILE_MAX;
+  __shared__ __attribute__((aligned(16))) uint8_t smem[SMEM];
+  auto sQ = reinterpret_cast<_Float16 (*)[QROW]>(smem);
+
+  // Stage the group's queries rotated like K: each 32-lane half owns one
+  // query vector, 8 dims per lane.
+  {
+    const int l32 = lane & 31;
+    for (int v = 2 * w + (lane >> 5); v < 32; v += 2 * NW) {
+      const int r = v / hpk, g = v % hpk;
+      const int qi = grp * QG + r;
+      float x[8];
+      const bool live = v < QG * hpk && qi < num_q;
+    #pragma unroll
+      for (int i = 0; i < 8; ++i)
+        x[i] = live ? to_f<QT>(Q[(int64_t)qi * sq0 +
+                                 (int64_t)(kvh * hpk + g) * sq1 + l32 * 8 + i])
+                    : 0.0f;
+    #pragma unroll
+      for (int i = 0; i < 8; ++i) x[i] *= sign_of(k_signs, l32 * 8 + i);
+    #pragma unroll
+      for (int hh = 1; hh < 8; hh <<= 1) {
+    #pragma unroll
+        for (int i = 0; i < 8; ++i)
+          if ((i & hh) == 0) {
+            const float a = x[i], b = x[i + hh];
+            x[i] = a + b;
+            x[i + hh] = a - b;
+          }
+      }
+    #pragma unroll
+      for (int m = 1; m <= (KC ? 16 : 4); m <<= 1) {
+        const bool hi = l32 & m;
+    #pragma unroll
+        for (int i = 0; i < 8; ++i) {
+          const float o = __shfl_xor(x[i], m);
+          x[i] = hi ? o - x[i] : x[i] + o;
+        }
+      }
+    #pragma unroll
+      for (int i = 0; i < 8; ++i)
+        sQ[v][l32 * 8 + i] = (_Float16)(x[i] * (KC ? 0.0625f : 0.125f));
+    }
+  }
+  __syncthreads();
+
+  int rreq[QG], rlen[QG];
+    #pragma unroll
+  for (int r = 0; r < QG; ++r) {
+    const int qi = grp * QG + r;
+    int rq = -1, ln = 0;
+    if (qi < num_q) {
+      rq = q_to_req[qi];
+      ln = q_to_klen[qi];
+      if (rq < 0 || rq >= num_reqs) {
+        rq = 0;
+        ln = 0;
+      }
+      ln = min(max(ln, 0), max_kv);
+    }
+    rreq[r] = __builtin_amdgcn_readfirstlane(rq);
+    rlen[r] = __builtin_amdgcn_readfirstlane(ln);
+  }
+
+  // This lane's query (column) as int8 B fragments for the i8 MFMA: step
+  // 2 b + h holds dims 64 b + 16 gq + 8 h .. + 7 (the dot product does not
+  // care how dims map to k, only that K uses the same map). Two int8 halves,
+  // q ~ qsc (hi + lo / 256), keep ~16 bits of the query.
+  const int qv = rt * 16 + col;
+  long Qh[8], Ql[8];
+  float qsc;
+  {
+    float x[64];
+    float qmax = 1e-20f;
+    #pragma unroll
+    for (int i = 0; i < 64; ++i) {
+      x[i] = (float)sQ[qv][64 * (i >> 4) + 16 * gq + (i & 15)];
+      qmax = fmaxf(qmax, fabsf(x[i]));
+    }
+    qmax = qmax4(qmax);
+    const float inv = 127.0f / qmax;
+    qsc = qmax * (1.0f / 127.0f);
+    #pragma unroll
+    for (int st = 0; st < 8; ++st) {
+      uint64_t wh = 0, wl = 0;
+    #pragma unroll
+      for (int e = 0; e < 8; ++e) {
+        const float y = x[8 * st + e] * inv;
+        const float h = rintf(y);
+        const int lo = max(-127, min(127, (int)rintf((y - h) * 256.0f)));
+        wh |= (uint64_t)(uint8_t)(int8_t)(int)h << (8 * e);
+        wl |= (uint64_t)(uint8_t)(int8_t)lo << (8 * e);
+      }
+      Qh[st] = (long)wh;
+      Ql[st] = (long)wl;
+    }
+  }
+  const float sm_q = sm_log2 * qsc * (1.0f / 256.0f);
+  __syncthreads();
+  uint8_t* lt = smem + w * TILE_MAX;
+
+  // rreq/rlen are only ever indexed by constants (unrolled selections): a
+  // dynamic index would put them in scratch memory.
+  auto pick = [&](const int* a, int i) {
+    int x = a[0];
+    #pragma unroll
+    for (int r = 1; r < QG; ++r) x = i == r ? a[r] : x;
+    return x;
+  };
+  int seg0 = 0;
+  while (seg0 < QG && pick(rreq, seg0) >= 0) {
+    const int req = pick(rreq, seg0);
+    int seg1 = seg0 + 1;
+    while (seg1 < QG && pick(rreq, seg1) == req) ++seg1;
+    int Lmax = 0;
+    #pragma unroll
+    for (int r = 0; r < QG; ++r)
+      if (r >= seg0 && r < seg1) Lmax = max(Lmax, rlen[r]);
+    const int qv_lo = seg0 * hpk, qv_hi = seg1 * hpk;
+    int tps = (Lmax + num_splits - 1) / num_splits;
+    if (tps < min_tps) tps = min_tps;
+    tps = (tps + 15) & ~15;
+
+    const bool tile_live = qv_hi > rt * 16 && qv_lo < rt * 16 + 16;
+    const bool qlive = qv >= qv_lo && qv < qv_hi;
+    if (blockIdx.z * splits_per_block * tps >= Lmax) {
+      if (tile_live && si < num_splits && gq == 0 && qlive) {
+        float* op = mid_o + (int64_t)(grp * QG + qv / hpk) * smo +
+                    (int64_t)(kvh * hpk + qv % hpk) * smh + (int64_t)si * sms;
+        op[D] = -INFINITY;
+        op[D + 1] = 0.0f;
+      }
+      seg0 = seg1;
+      continue;
+    }
+
+    if (tile_live && si < num_splits) {
+      int qlen = 0;
+    #pragma unroll
+      for (int r = 0; r < QG; ++r)
+        if (qlive && qv / hpk == r) qlen = rlen[r];
+
+      const int start = si * tps;
+      const int end = min(start + tps, Lmax);
+      float m = -INFINITY, l = 0.0f, cp = 0.0f;
+      f4v O[16];
+    #pragma unroll
+      for (int c = 0; c < 16; ++c) O[c] = (f4v){0, 0, 0, 0};
+
+      const uint8_t* cache_h = cache + (int64_t)kvh * sch;
+      auto tile_ptr = [&](int pos) {
+        const int lb = pos / block_size;
+        int pb = ((cint*)block_table)[req * max_blocks + lb];
+        if (pb < 0 || pb >= num_phys_blocks) pb = 0;
+        return cache_h + (int64_t)pb * scb +
+               (int64_t)(pos - lb * block_size) * sct;
+      };
+
+      // Raw K of token col: per 64-dim block, the codes of dims
+      // 16 gq .. + 15. 4-bit: 8 bytes. 3-bit: 2-bit plane word gq and the
+      // 16 bits of the 1-bit plane. Scales of tokens 4 gq .. + 3.
+      constexpr int NB3 = KC ? 4 : 3;  // 3-bit blocks
+      uint2 kr = {0, 0};
+      uint2 kn[NB3];
+      uint32_t sa[4], sb[4], sv[4];
+      // A 16-token tile is one contiguous run of a block: each lane fetches
+      // four 16-byte chunks of the next tile into registers, then stores
+      // them to its wave's LDS buffer, where the fragments are read.
+      constexpr int tb = 16 * F::SLOT;
+      uint4 raw[TILE_MAX / 1024], raw2[TILE_MAX / 1024];
+      auto fetch = [&](uint4* dst, const uint8_t* tile) {
+    #pragma unroll
+        for (int j = 0; j < TILE_MAX / 1024; ++j) {
+          const int o = 16 * (lane + 64 * j);
+          dst[j] = o < tb ? *(const uint4*)(tile + o) : make_uint4(0, 0, 0, 0);
+        }
+      };
+      auto stash = [&]() {
+    #pragma unroll
+        for (int j = 0; j < TILE_MAX / 1024; ++j) {
+          const int o = 16 * (lane + 64 * j);
+          if (o < tb) *(uint4*)(lt + o) = raw[j];
+        }
+      };
+      auto load_k = [&](const uint8_t* tile) {
+        const uint8_t* pk = tile + (int64_t)col * F::SLOT;
+        if constexpr (!KC) kr = *(const uint2*)(pk + 8 * gq);
+    #pragma unroll
+        for (int b = 0; b < NB3; ++b) {
+          kn[b].x = *(const uint32_t*)(pk + F::OFF_KN + 16 * b + 4 * gq);
+          kn[b].y = *(const uint32_t*)(pk + F::OFF_KN + F::KN / 4 + 8 * b +
+                                       4 * (gq >> 1)) >>
+                    (4 * (gq & 1));
+        }
+    #pragma unroll
+        for (int e = 0; e < 4; ++e) {
+          const uint8_t* ps =
+              tile + (int64_t)(4 * gq + e) * F::SLOT + F::OFF_SC;
+          if constexpr (KC) {
+            sa[e] = *(const uint32_t*)ps;
+            sv[e] = sa[e] >> 16;
+          } else {
+            const uint2 u = *(const uint2*)ps;
+            sa[e] = u.x;
+            sb[e] = u.y;
+            sv[e] = *(const uint16_t*)(ps + 8);
+          }
+        }
+      };
+      // Raw V of tokens 4 gq .. + 3, dims 16 col .. + 15.
+      uint2 vr[4];
+      auto load_v = [&](const uint8_t* tile) {
+    #pragma unroll
+        for (int e = 0; e < 4; ++e) {
+          const uint8_t* p = tile + (int64_t)(4 * gq + e) * F::SLOT + F::OFF_V;
+          if constexpr (VB == 4) {
+            vr[e] = *(const uint2*)(p + 8 * col);
+          } else {
+            vr[e].x = *(const uint32_t*)(p + 4 * col);
+            vr[e].y = *(const uint32_t*)(p + D / 4 + 4 * (col >> 1)) >>
+                      (4 * (col & 1));
+          }
+        }
+      };
+      // Codebook bytes of dims 4 s .. + 3 of a 16-dim run: int8 for K, biased
+      // (x + 128, the bias folded into the tables) for V.
+      auto codes4 = [&](uint2 raw, int s) -> uint32_t {
+        const uint32_t wd = s < 2 ? raw.x : raw.y;
+        return (wd >> (4 * (s & 1))) & 0x0F0F0F0Fu;
+      };
+      auto codes3 = [&](uint2 raw, int s) -> uint32_t {
+        return ((raw.x >> (2 * s)) & 0x03030303u) |
+               (((raw.y >> s) & 0x01010101u) << 2);
+      };
+      auto k_bytes = [&](int b, int s) -> uint32_t {
+        return (!KC && b == 0) ? lut4_fast<false>(codes4(kr, s))
+                               : octave::lut3(codes3(kn[KC ? b : b - 1], s));
+      };
+      auto v_bytes = [&](uint2 raw, int s) -> uint32_t {
+        if constexpr (VB == 4)
+          return lut4_fast<true>(codes4(raw, s));
+        else
+          return __builtin_amdgcn_perm(octave::LUT3_HI ^ X, octave::LUT3_LO ^ X,
+                                       codes3(raw, s));
+      };
+
+      if (start < end) fetch(raw, tile_ptr(start));
+      if (start + 16 < end) fetch(raw2, tile_ptr(start + 16));
+
+      for (int base = start; base < end; base += 16) {
+        stash();
+        load_k(lt);
+    #pragma unroll
+        for (int j = 0; j < TILE_MAX / 1024; ++j) raw[j] = raw2[j];
+        if (base + 32 < end) fetch(raw2, tile_ptr(base + 32));
+
+        i4v Sh[NG], Sl[NG];
+    #pragma unroll
+        for (int g = 0; g < NG; ++g) Sh[g] = Sl[g] = (i4v){0, 0, 0, 0};
+    #pragma unroll
+        for (int b = 0; b < 4; ++b) {
+          const int g = KC ? 0 : b;
+    #pragma unroll
+          for (int h = 0; h < 2; ++h) {
+            const long a = (long)(((uint64_t)k_bytes(b, 2 * h + 1) << 32) |
+                                  k_bytes(b, 2 * h));
+            Sh[g] = __builtin_amdgcn_mfma_i32_16x16x32_i8(a, Qh[2 * b + h],
+                                                          Sh[g], 0, 0, 0);
+            Sl[g] = __builtin_amdgcn_mfma_i32_16x16x32_i8(a, Ql[2 * b + h],
+                                                          Sl[g], 0, 0, 0);
+          }
+        }
+        auto dot = [&](int g, int e) {
+          return (float)(Sh[g][e] * 256 + Sl[g][e]);
+        };
+
+        load_v(lt);
+
+        float sc[4];
+        float mloc = -INFINITY;
+    #pragma unroll
+        for (int e = 0; e < 4; ++e) {
+          const int t = base + 4 * gq + e;
+          const bool ok = t < end && t < qlen;
+          float v;
+          if constexpr (KC) {
+            v = lo16(sa[e]) * dot(0, e);
+          } else {
+            v = lo16(sa[e]) * dot(0, e) + hi16(sa[e]) * dot(1, e) +
+                lo16(sb[e]) * dot(2, e) + hi16(sb[e]) * dot(3, e);
+          }
+          sc[e] = ok ? v * sm_q : -INFINITY;
+          mloc = fmaxf(mloc, sc[e]);
+        }
+        const float mn = fmaxf(m, qmax4(mloc));
+        float alpha = 1.0f;
+        if (mn > m)
+          alpha = (m == -INFINITY) ? 0.0f : __builtin_amdgcn_exp2f(m - mn);
+        m = mn;
+        h4v pb;
+        float psum = 0.0f, lsum = 0.0f;
+    #pragma unroll
+        for (int e = 0; e < 4; ++e) {
+          const float p =
+              sc[e] == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(sc[e] - mn);
+          lsum += p;
+          const _Float16 hp = (_Float16)(p * lo16(sv[e]) * PSCALE);
+          psum += (float)hp;
+          pb[e] = hp;
+        }
+        l = l * alpha + lsum;
+        cp = cp * alpha + psum;
+        if (__ballot(alpha != 1.0f)) {
+    #pragma unroll
+          for (int c = 0; c < 16; ++c) O[c] *= alpha;
+        }
+
+          // A of fragment 4 s + d: dim 16 col + 4 s + d of tokens 4 gq .. + 3.
+          // One perm per token pair transposes the biased bytes and widens
+          // them to fp16 [b, 0x00] (selector 0x0C), i.e. the exact denormal
+          // b * 2^-24.
+    #pragma unroll
+        for (int s = 0; s < 4; ++s) {
+          uint32_t a[4];
+    #pragma unroll
+          for (int e = 0; e < 4; ++e) a[e] = v_bytes(vr[e], s);
+    #pragma unroll
+          for (int d = 0; d < 4; ++d) {
+            constexpr uint32_t kSel = 0x0C000C00u;
+            const uint32_t sel = kSel | ((4u + d) << 16) | (uint32_t)d;
+            const h4v A = __builtin_bit_cast(
+                h4v, (uint2){__builtin_amdgcn_perm(a[1], a[0], sel),
+                             __builtin_amdgcn_perm(a[3], a[2], sel)});
+            O[4 * s + d] = __builtin_amdgcn_mfma_f32_16x16x16f16(
+                A, pb, O[4 * s + d], 0, 0, 0);
+          }
+        }
+      }
+
+      l = qsum4(l);
+      cp = qsum4(cp);
+      if (qlive) {
+        float* op = mid_o + (int64_t)(grp * QG + qv / hpk) * smo +
+                    (int64_t)(kvh * hpk + qv % hpk) * smh + (int64_t)si * sms;
+        const bool any = m != -INFINITY;
+        // O holds sum hp (x + 128) 2^-24.
+        const float corr = 128.0f * 0x1p-24f * cp;
+        constexpr float kOut = 0x1p24f / (PSCALE * LUT_ONE);
+    #pragma unroll
+        for (int e = 0; e < 4; ++e) {
+          float* dp = op + 16 * (4 * gq + e);
+    #pragma unroll
+          for (int c = 0; c < 16; c += 4) {
+            float4 v;
+            v.x = any ? (O[c][e] - corr) * kOut : 0.0f;
+            v.y = any ? (O[c + 1][e] - corr) * kOut : 0.0f;
+            v.z = any ? (O[c + 2][e] - corr) * kOut : 0.0f;
+            v.w = any ? (O[c + 3][e] - corr) * kOut : 0.0f;
+            *(float4*)(dp + c) = v;
+          }
+        }
+        if (gq == 0) {
+          op[D] = m;
+          op[D + 1] = l;
+        }
+      }
+    }
+    seg0 = seg1;
+  }
+  #endif  // OCTAVE_MFMA_STUB
+}
+
+// Split combine for gfx942: grid (num_q, num_q_heads), 16 waves. Thread t
+// weighs split t; wave w sums splits w, w + 16, ... four at a time over
+// dims 4 lane .. + 3; wave 0 adds the 16 partials and undoes the 256-dim V
+// rotation in registers (two in-lane stages, six shuffle stages).
+template <typename OutT>
+__global__ __launch_bounds__(1024) void reduce_mfma(
+    const float* __restrict__ mid_o, OutT* __restrict__ out,
+    const int* __restrict__ v_signs, int num_splits, int64_t smo, int64_t smh,
+    int64_t sms, int64_t soo, int64_t soh) {
+  #ifndef OCTAVE_MFMA_STUB
+  __shared__ float sw[1024], swm[16], swl[16];
+  __shared__ __attribute__((aligned(16))) float so[16][D];
+  const int qi = blockIdx.x, hi = blockIdx.y;
+  const int tid = threadIdx.x, w = tid >> 6, lane = tid & 63;
+  const float* base = mid_o + qi * smo + hi * smh;
+
+  const float ms = tid < num_splits ? base[tid * sms + D] : -INFINITY;
+  const float ls = tid < num_splits ? base[tid * sms + D + 1] : 0.0f;
+  float m = ms;
+    #pragma unroll
+  for (int k = 1; k < 64; k <<= 1) m = fmaxf(m, __shfl_xor(m, k));
+  if (lane == 0) swm[w] = m;
+  __syncthreads();
+  float mg = swm[0];
+    #pragma unroll
+  for (int i = 1; i < 16; ++i) mg = fmaxf(mg, swm[i]);
+  const float a = ms == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(ms - mg);
+  sw[tid] = a;
+  float l = ls * a;
+    #pragma unroll
+  for (int k = 1; k < 64; k <<= 1) l += __shfl_xor(l, k);
+  if (lane == 0) swl[w] = l;
+  __syncthreads();
+
+  float4 o = {0, 0, 0, 0};
+  for (int s0 = w; s0 < num_splits; s0 += 64) {
+    float4 x[4];
+    float as[4];
+    #pragma unroll
+    for (int u = 0; u < 4; ++u) {
+      const int s = s0 + 16 * u;
+      as[u] = s < num_splits ? sw[s] : 0.0f;
+      x[u] = s < num_splits ? *(const float4*)(base + s * sms + 4 * lane)
+                            : make_float4(0, 0, 0, 0);
+    }
+    #pragma unroll
+    for (int u = 0; u < 4; ++u) {
+      // Splits without tokens weigh zero; their O was never written.
+      o.x += as[u] == 0.0f ? 0.0f : x[u].x * as[u];
+      o.y += as[u] == 0.0f ? 0.0f : x[u].y * as[u];
+      o.z += as[u] == 0.0f ? 0.0f : x[u].z * as[u];
+      o.w += as[u] == 0.0f ? 0.0f : x[u].w * as[u];
+    }
+  }
+  *(float4*)&so[w][4 * lane] = o;
+  __syncthreads();
+  if (w != 0) return;
+  float x[4] = {0, 0, 0, 0};
+    #pragma unroll
+  for (int i = 0; i < 16; ++i) {
+    const float4 v = *(const float4*)&so[i][4 * lane];
+    x[0] += v.x;
+    x[1] += v.y;
+    x[2] += v.z;
+    x[3] += v.w;
+  }
+  float lg = 0.0f;
+    #pragma unroll
+  for (int i = 0; i < 16; ++i) lg += swl[i];
+  {
+    const float a0 = x[0] + x[1], a1 = x[0] - x[1];
+    const float a2 = x[2] + x[3], a3 = x[2] - x[3];
+    x[0] = a0 + a2;
+    x[1] = a1 + a3;
+    x[2] = a0 - a2;
+    x[3] = a1 - a3;
+  }
+    #pragma unroll
+  for (int mk = 1; mk < 64; mk <<= 1) {
+    const bool hi_lane = lane & mk;
+    #pragma unroll
+    for (int i = 0; i < 4; ++i) {
+      const float p = __shfl_xor(x[i], mk);
+      x[i] = hi_lane ? p - x[i] : x[i] + p;
+    }
+  }
+  const float inv_l = lg > 0.f ? 0.0625f / lg : 0.f;
+  OutT* op = out + qi * soo + hi * soh + 4 * lane;
+    #pragma unroll
+  for (int i = 0; i < 4; ++i)
+    op[i] = (OutT)(x[i] * sign_of(v_signs, 4 * lane + i) * inv_l);
+  #endif  // OCTAVE_MFMA_STUB
+}
+
+}  // namespace octave_mfma
+
+// Split combine on gfx942; false when the shape is not covered.
+bool octave_reduce_mfma(torch::Tensor out, torch::Tensor mid_o,
+                        torch::Tensor v_signs, int num_splits,
+                        cudaStream_t stream) {
+  using namespace octave_mfma;
+  // 1024 threads per (query, head) pay off while those blocks alone cannot
+  // fill the GPU; big batches keep the 256-thread kernel.
+  static const int max_blocks =
+      2 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  if (num_splits > 1024 || out.size(0) * out.size(1) > max_blocks ||
+      out.size(2) != D || out.stride(2) != 1 || mid_o.stride(3) != 1)
+    return false;
+  dim3 grid(out.size(0), out.size(1));
+  if (out.dtype() == at::kBFloat16)
+    reduce_mfma<__hip_bfloat16><<<grid, 1024, 0, stream>>>(
+        mid_o.data_ptr<float>(), (__hip_bfloat16*)out.data_ptr(),
+        v_signs.data_ptr<int>(), num_splits, mid_o.stride(0), mid_o.stride(1),
+        mid_o.stride(2), out.stride(0), out.stride(1));
+  else
+    reduce_mfma<half><<<grid, 1024, 0, stream>>>(
+        mid_o.data_ptr<float>(), (half*)out.data_ptr(), v_signs.data_ptr<int>(),
+        num_splits, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2),
+        out.stride(0), out.stride(1));
+  return true;
+}
+
+constexpr int kMfmaSplits = 256;
+constexpr int kMfmaMinTps = 32;
+constexpr int kMfmaNw = 4;
+
+// Launches the MFMA decode when the shape is covered; returns the number of
+// splits it wrote, or 0 when the caller must use its own kernel. Dense only.
+int octave_decode_mfma(torch::Tensor query, torch::Tensor cache,
+                       torch::Tensor block_table, torch::Tensor q_to_req,
+                       torch::Tensor q_to_klen, torch::Tensor mid_o,
+                       torch::Tensor k_signs, double sm_scale,
+                       int64_t num_kv_splits, int64_t fmt,
+                       int64_t query_group) {
+  using namespace octave_mfma;
+  const int num_q = query.size(0);
+  const int num_q_heads = query.size(1);
+  const int num_kv_heads = cache.size(1);
+  const int hpk = num_q_heads / num_kv_heads;
+  const int QG = query_group > 1 ? 4 : 1;
+  const int slot = fmt == 4   ? octave::Format<4, false>::SLOT
+                   : fmt == 3 ? octave::Format<3, false>::SLOT
+                              : octave::Format<3, true>::SLOT;
+  const bool eligible =
+      query.size(2) == D && query.stride(2) == 1 && hpk * QG <= 32 &&
+      cache.size(2) % 16 == 0 &&
+      cache.stride(2) % (fmt & octave::kCompactFlag ? 4 : 8) == 0 &&
+      cache.stride(2) == slot && cache.stride(1) % 16 == 0 &&
+      cache.stride(0) % 16 == 0 &&
+      reinterpret_cast<uintptr_t>(cache.data_ptr()) % 16 == 0 &&
+      num_kv_splits >= 1;
+  if (!eligible) return 0;
+  const int nrt = std::min(num_q, QG) * hpk > 16 ? 2 : 1;
+  // Enough splits for ~two resident waves per SIMD, no more: every extra
+  // split is another partial for the reduction to read.
+  static const int target_waves =
+      8 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  const int groups = std::max(1, (num_q + QG - 1) / QG) * num_kv_heads * nrt;
+  const int ns =
+      std::min<int>({kMfmaSplits, (int)num_kv_splits,
+                     std::max(1, (target_waves + groups - 1) / groups)});
+  if (num_q == 0) return ns;
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int spb = kMfmaNw / nrt;
+  dim3 grid((num_q + QG - 1) / QG, num_kv_heads, (ns + spb - 1) / spb);
+  #define SQM_G(B, KC, G, T)                                                   \
+    decode_mfma<B, KC, kMfmaNw, G, T><<<grid, 64 * kMfmaNw, 0, stream>>>(      \
+        (const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(),          \
+        block_table.data_ptr<int>(), q_to_req.data_ptr<int>(),                 \
+        q_to_klen.data_ptr<int>(), mid_o.data_ptr<float>(),                    \
+        k_signs.data_ptr<int>(), (float)sm_scale, num_q, num_q_heads,          \
+        num_kv_heads, cache.size(2), block_table.size(1), block_table.size(0), \
+        cache.size(0), ns, kMfmaMinTps, nrt, query.stride(0), query.stride(1), \
+        cache.stride(0), cache.stride(1), cache.stride(2), mid_o.stride(0),    \
+        mid_o.stride(1), mid_o.stride(2))
+  #define SQM(B, KC, T)     \
+    do {                    \
+      if (QG == 4)          \
+        SQM_G(B, KC, 4, T); \
+      else                  \
+        SQM_G(B, KC, 1, T); \
+    } while (0)
+  const bool bf = query.dtype() == at::kBFloat16;
+  if (fmt == 4) {
+    if (bf)
+      SQM(4, false, __hip_bfloat16);
+    else
+      SQM(4, false, half);
+  } else if (fmt == 3) {
+    if (bf)
+      SQM(3, false, __hip_bfloat16);
+    else
+      SQM(3, false, half);
+  } else {
+    if (bf)
+      SQM(3, true, __hip_bfloat16);
+    else
+      SQM(3, true, half);
+  }
+  #undef SQM
+  #undef SQM_G
+  return ns;
+}
+
+#endif  // USE_ROCM

csrc/attention/octave_decode_wmma_rdna3.cu

@@ -0,0 +1,687 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Octave decode on WMMA for RDNA3 (gfx11). Independent waves, one (split,
+// 16-query row tile) per wave, transposed orientation, and lane pairs sharing
+// A fragments through a DPP swap. The products:
+//
+//   QK  S^T[token][q] = K x Q^T over 16 k-steps of 16 dims: steps 0-3 the
+//       RoPE block (4-bit codes), 4-15 the NoPE blocks (3-bit codes). Codes
+//       are mapped to their int8 codebook values and fed as x + 1152. Each
+//       64-dim block has its own per-token scale, so the steps accumulate
+//       into four tiles that are scaled and debiased in the softmax. Compact
+//       K: 16 steps of 3-bit codes and one scale, so one tile.
+//   PV  O^T[dim][q] = V^T x P^T with the V codebook values fed as x + 1152;
+//       the per-token V scale is folded into P, the bias removed with sum(P).
+//
+// The query is rotated like K while it is staged into LDS; the output stays
+// rotated and octave_reduce rotates it back.
+
+#include <cstdint>
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#if defined(USE_ROCM)
+  #include <hip/hip_runtime.h>
+  #include <hip/hip_bf16.h>
+  #include <hip/hip_fp16.h>
+
+  #include "octave_format.cuh"
+
+  #if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx1100__)
+    #define OCTAVE_WMMA_STUB
+  #endif
+
+namespace octave_wmma {
+
+typedef _Float16 h16v __attribute__((ext_vector_type(16)));
+typedef _Float16 h8v __attribute__((ext_vector_type(8)));
+typedef _Float16 h2v __attribute__((ext_vector_type(2)));
+typedef float f8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u8v __attribute__((ext_vector_type(8)));
+typedef const int __attribute__((address_space(4))) cint;
+
+using octave::D;
+using octave::Format;
+using octave::LUT_ONE;
+using octave::R;
+constexpr int QG = 4;  // query tokens per block
+// Exact recent and leading positions (sparse window mode).
+[[maybe_unused]] constexpr int kWindow = 64;
+[[maybe_unused]] constexpr int kSinks = 4;
+[[maybe_unused]] constexpr int QROW = D + 8;  // padded LDS row (bank spread)
+[[maybe_unused]] constexpr float PSCALE = 256.0f;
+[[maybe_unused]] constexpr float BIAS = 1152.0f;  // int8 byte x fed as x + 1152
+[[maybe_unused]] constexpr uint32_t X = 0x80808080u;
+
+__device__ __forceinline__ uint32_t swap1(uint32_t v) {
+  return (uint32_t)__builtin_amdgcn_update_dpp(0, (int)v, 0xB1, 0xF, 0xF,
+                                               false);
+}
+__device__ __forceinline__ uint32_t xhalf(uint32_t v) {
+  return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xFEDCBA98u, false,
+                                      false);
+}
+__device__ __forceinline__ float xhalf(float v) {
+  return __uint_as_float(xhalf(__float_as_uint(v)));
+}
+
+// Bytes b -> fp16 1024 + b, exact.
+__device__ __forceinline__ h16v bytes_to_h16(uint32_t w0, uint32_t w1,
+                                             uint32_t w2, uint32_t w3) {
+  const uint32_t ws[4] = {w0, w1, w2, w3};
+  u8v r;
+  #pragma unroll
+  for (int i = 0; i < 4; ++i) {
+    r[2 * i] = __builtin_amdgcn_perm(0x64646464u, ws[i], 0x07010700u);
+    r[2 * i + 1] = __builtin_amdgcn_perm(0x64646464u, ws[i], 0x07030702u);
+  }
+  return __builtin_bit_cast(h16v, r);
+}
+
+__device__ __forceinline__ h16v ld16(const _Float16* p) {
+  const h8v lo = *(const h8v*)p;
+  const h8v hi = *(const h8v*)(p + 8);
+  return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
+                                 12, 13, 14, 15);
+}
+
+// 4 dwords of 4 tokens each (same 4 dims) -> 4 dwords of 4 dims each.
+__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2,
+                                    uint32_t a3, uint32_t o[4]) {
+  const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u);
+  const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u);
+  const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u);
+  const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u);
+  o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u);
+  o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u);
+  o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u);
+  o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u);
+}
+
+template <typename T>
+__device__ __forceinline__ float to_f(T x) {
+  return (float)x;
+}
+template <>
+__device__ __forceinline__ float to_f<half>(half x) {
+  return __half2float(x);
+}
+
+__device__ __forceinline__ float sign_of(const int* bits, int i) {
+  return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f;
+}
+
+// SP: sparse mode. Row qi attends to the logical positions idx[qi, 0..topk)
+// (negative entries are skipped) instead of its first q_to_klen[qi] tokens,
+// and every row is its own segment since rows do not share positions.
+template <int VB, bool KC, int NSB, typename QT, bool SP>
+__global__ __launch_bounds__(64 * NSB) void decode_wmma(
+    const QT* __restrict__ Q, const uint8_t* __restrict__ cache,
+    const int* __restrict__ block_table, const int* __restrict__ q_to_req,
+    const int* __restrict__ q_to_klen, float* __restrict__ mid_o,
+    const int* __restrict__ k_signs, float sm_scale, int num_q, int num_q_heads,
+    int num_kv_heads, int block_size, int max_blocks, int num_reqs,
+    int num_phys_blocks, int num_splits, int min_tps, int num_row_tiles,
+    int64_t sq0, int64_t sq1, int64_t scb, int64_t sch, int64_t sct,
+    int64_t smo, int64_t smh, int64_t sms, const int* __restrict__ idx,
+    int64_t sidx, int topk, const int* __restrict__ qpos,
+    const int* __restrict__ wtags, int wmask, const int* __restrict__ stags,
+    int smask) {
+  #ifndef OCTAVE_WMMA_STUB
+  using F = Format<VB, KC>;
+  constexpr int NG = KC ? 1 : 4;  // QK tiles (one per K scale)
+  const int grp = blockIdx.x;
+  const int kvh = blockIdx.y;
+  const int tid = threadIdx.x;
+  const int w = __builtin_amdgcn_readfirstlane(tid >> 5), lane = tid & 31;
+  // Two row tiles when the group can hold more than 16 query vectors;
+  // otherwise every wave takes its own split.
+  const int rt = num_row_tiles == 2 ? (w & 1) : 0;
+  const int splits_per_block = 2 * NSB / num_row_tiles;
+  const int si =
+      blockIdx.z * splits_per_block + (num_row_tiles == 2 ? w >> 1 : w);
+  const int j = lane & 15, h = lane >> 4;
+  const bool useful = (j & 1) == h;
+  const int rr = useful ? j : (j ^ 1);
+  const int hpk = num_q_heads / num_kv_heads;
+  const int max_kv = max_blocks * block_size;
+  const float sm_log2 = sm_scale * 1.4426950408889634f / LUT_ONE;
+
+  __shared__ __attribute__((aligned(16))) _Float16 sQ[32][QROW];
+
+  // Stage the group's queries rotated like K: lane owns 8 dims.
+  for (int v = w; v < 32; v += 2 * NSB) {
+    const int r = v / hpk, g = v % hpk;
+    const int qi = grp * QG + r;
+    float x[8];
+    const bool live = v < QG * hpk && qi < num_q;
+    #pragma unroll
+    for (int i = 0; i < 8; ++i)
+      x[i] = live ? to_f<QT>(Q[(int64_t)qi * sq0 +
+                               (int64_t)(kvh * hpk + g) * sq1 + lane * 8 + i])
+                  : 0.0f;
+    #pragma unroll
+    for (int i = 0; i < 8; ++i) x[i] *= sign_of(k_signs, lane * 8 + i);
+    #pragma unroll
+    for (int hh = 1; hh < 8; hh <<= 1) {
+    #pragma unroll
+      for (int i = 0; i < 8; ++i)
+        if ((i & hh) == 0) {
+          const float a = x[i], b = x[i + hh];
+          x[i] = a + b;
+          x[i + hh] = a - b;
+        }
+    }
+    #pragma unroll
+    for (int m = 1; m <= (KC ? 16 : 4); m <<= 1) {
+      const bool hi = lane & m;
+    #pragma unroll
+      for (int i = 0; i < 8; ++i) {
+        const float o = __shfl_xor(x[i], m);
+        x[i] = hi ? o - x[i] : x[i] + o;
+      }
+    }
+    #pragma unroll
+    for (int i = 0; i < 8; ++i)
+      sQ[v][lane * 8 + i] = (_Float16)(x[i] * (KC ? 0.0625f : 0.125f));
+  }
+  __syncthreads();
+
+  int rreq[QG], rlen[QG];
+    #pragma unroll
+  for (int r = 0; r < QG; ++r) {
+    const int qi = grp * QG + r;
+    int rq = -1, ln = 0;
+    if (qi < num_q) {
+      rq = q_to_req[qi];
+      ln = SP ? topk : q_to_klen[qi];
+      if (rq < 0 || rq >= num_reqs) {
+        rq = 0;
+        ln = 0;
+      }
+      ln = min(max(ln, 0), SP ? topk : max_kv);
+    }
+    rreq[r] = __builtin_amdgcn_readfirstlane(rq);
+    rlen[r] = __builtin_amdgcn_readfirstlane(ln);
+  }
+
+  // Per-segment sums of this lane's query (debiasing): half h sums dims
+  // [128 h, 128 h + 128) as two 64-dim segments; the halves swap.
+  const int qv = rt * 16 + j;
+  float bq[4];
+  {
+    const h2v one = {(_Float16)1.0f, (_Float16)1.0f};
+    float s0 = 0.f, s1 = 0.f;
+    #pragma unroll
+    for (int s = 0; s < 8; ++s) {
+      const h8v a = *(const h8v*)&sQ[qv][128 * h + 8 * s];
+      const h8v b = *(const h8v*)&sQ[qv][128 * h + 64 + 8 * s];
+    #pragma unroll
+      for (int i = 0; i < 8; i += 2) {
+        s0 = __builtin_amdgcn_fdot2((h2v){a[i], a[i + 1]}, one, s0, false);
+        s1 = __builtin_amdgcn_fdot2((h2v){b[i], b[i + 1]}, one, s1, false);
+      }
+    }
+    const float o0 = xhalf(s0), o1 = xhalf(s1);
+    bq[0] = (h ? o0 : s0) * BIAS;
+    bq[1] = (h ? o1 : s1) * BIAS;
+    bq[2] = (h ? s0 : o0) * BIAS;
+    bq[3] = (h ? s1 : o1) * BIAS;
+    if constexpr (KC) bq[0] += bq[1] + bq[2] + bq[3];
+  }
+
+  int seg0 = 0;
+  while (seg0 < QG && rreq[seg0] >= 0) {
+    int seg1 = seg0 + 1;
+    if constexpr (!SP)
+      while (seg1 < QG && rreq[seg1] == rreq[seg0]) ++seg1;
+    const int req = rreq[seg0];
+    int Lmax = 0;
+    for (int r = seg0; r < seg1; ++r) Lmax = max(Lmax, rlen[r]);
+    const int qv_lo = seg0 * hpk, qv_hi = seg1 * hpk;
+    int tps = (Lmax + num_splits - 1) / num_splits;
+    if (tps < min_tps) tps = min_tps;
+    tps = (tps + 15) & ~15;
+
+    const bool tile_live = qv_hi > rt * 16 && qv_lo < rt * 16 + 16;
+    const bool qlive = qv >= qv_lo && qv < qv_hi;
+    if (blockIdx.z * splits_per_block * tps >= Lmax) {
+      if (tile_live && si < num_splits && h == 0 && qlive) {
+        int qx = qv;
+        asm volatile("" : "+v"(qx));
+        float* op = mid_o + (int64_t)(grp * QG + qx / hpk) * smo +
+                    (int64_t)(kvh * hpk + qx % hpk) * smh + (int64_t)si * sms;
+        op[D] = -INFINITY;
+        op[D + 1] = 0.0f;
+      }
+      seg0 = seg1;
+      continue;
+    }
+
+    if (tile_live && si < num_splits) {
+      int qlen = 0;
+    #pragma unroll
+      for (int r = 0; r < QG; ++r)
+        if (qlive && qv / hpk == r) qlen = rlen[r];
+
+      const int start = si * tps;
+      const int end = min(start + tps, Lmax);
+      float m = -INFINITY, l = 0.0f, cp = 0.0f;
+      f8v O[16];
+    #pragma unroll
+      for (int c = 0; c < 16; ++c) O[c] = (f8v){0, 0, 0, 0, 0, 0, 0, 0};
+
+      const uint8_t* cache_h = cache + (int64_t)kvh * sch;
+      auto pickw = [&](uint32_t v, int t) {
+        const uint32_t ev = __builtin_amdgcn_readlane(v, t);
+        const uint32_t od = __builtin_amdgcn_readlane(v, t + 1);
+        return h ? od : ev;
+      };
+      auto tile_ptr = [&](int pos) {
+        const int lb = pos / block_size;
+        int pb = ((cint*)block_table)[req * max_blocks + lb];
+        if (pb < 0 || pb >= num_phys_blocks) pb = 0;
+        return cache_h + (int64_t)pb * scb +
+               (int64_t)(pos - lb * block_size) * sct;
+      };
+      // Sparse mode: slot of the pos-th selected position of this row, and
+      // whether that position maps to a live cache slot. Dead entries read
+      // slot 0 of block 0 and are masked out of the softmax.
+      const int* irow = SP ? idx + (int64_t)(grp * QG + seg0) * sidx : nullptr;
+      auto sel_block = [&](int pos, int& off) {
+        const int t = irow[min(pos, topk - 1)];
+        const int lb = max(t, 0) / block_size;
+        off = max(t, 0) - lb * block_size;
+        const int pb = (t >= 0 && lb < max_blocks)
+                           ? ((cint*)block_table)[req * max_blocks + lb]
+                           : -1;
+        return (pb >= 0 && pb < num_phys_blocks) ? pb : -1;
+      };
+      // Exact-window tokens (attended by octave_window_attn instead): the
+      // last kWindow positions before the query, and the first kSinks.
+      const int rpos = SP && wtags ? qpos[grp * QG + seg0] : 0;
+      auto in_window = [&](int pos, int pb, int off) {
+        const int t = irow[min(pos, topk - 1)];
+        const int d = rpos - t;
+        const int slot = pb * block_size + off;
+        return (d >= 0 && d < kWindow && wtags[slot & wmask] == slot) ||
+               (t >= 0 && t < kSinks && stags[slot & smask] == slot);
+      };
+      auto sel_ptr = [&](int pos) {
+        int off;
+        const int pb = sel_block(pos, off);
+        return pb < 0 ? cache_h
+                      : cache_h + (int64_t)pb * scb + (int64_t)off * sct;
+      };
+
+      // Raw K loads, one tile ahead. Each lane of a pair owns 8 of the 16
+      // k-steps of token rr: NR4 4-bit steps (8 bytes of codes each), then
+      // 3-bit steps (2-bit plane word, 1-bit plane word pre-shifted). Block
+      // K: the useful lane's slots are steps {0,1,4..9}, its partner's
+      // {2,3,10..15}. Compact K: {0..7} and {8..15}. The partner's reach the
+      // useful lane through the DPP swap.
+      constexpr int NR4 = F::KR / 32;  // 4-bit slots per lane
+      constexpr int NR3 = 8 - NR4;     // 3-bit slots per lane
+      uint2 kr_raw[NR4 > 0 ? NR4 : 1];
+      uint2 kn_raw[NR3];
+      // Scales of token j (lane j), raw fp16 pairs. Block K: {k0, k1},
+      // {k2, k3}, {v}. Compact K: {k, v}.
+      uint32_t sc_a = 0, sc_b = 0, sc_v = 0;
+      auto load_k = [&](const uint8_t* tile, int base, float dep) {
+        uint32_t off = (uint32_t)(rr * (int)sct);
+        asm volatile("" : "+v"(off) : "v"(dep));
+        const uint8_t* pk = SP ? sel_ptr(base + rr) : tile + off;
+        const int chunk0 = useful ? 0 : NR3;
+        if constexpr (NR4 > 0) {
+          const int rope0 = useful ? 0 : NR4;
+    #pragma unroll
+          for (int k = 0; k < NR4; ++k)
+            kr_raw[k] = *(const uint2*)(pk + 8 * (rope0 + k));
+        }
+    #pragma unroll
+        for (int k = 0; k < NR3; ++k) {
+          const int c = chunk0 + k;
+          kn_raw[k].x = *(const uint32_t*)(pk + F::OFF_KN + 4 * c);
+          kn_raw[k].y =
+              *(const uint32_t*)(pk + F::OFF_KN + F::KN / 4 + 4 * (c >> 1)) >>
+              (4 * (c & 1));
+        }
+        const uint8_t* ps =
+            (SP ? sel_ptr(base + j) : tile + (int64_t)j * sct) + F::OFF_SC;
+        if constexpr (KC) {
+          sc_a = *(const uint32_t*)ps;
+          sc_v = sc_a >> 16;
+        } else {
+          const uint2 u = *(const uint2*)ps;
+          sc_a = u.x;
+          sc_b = u.y;
+          sc_v = *(const uint16_t*)(ps + 8);
+        }
+      };
+      // Raw V loads: lane rr's 16 dims of the even (useful) or odd tokens.
+      uint2 vraw[8];
+      auto load_v = [&](const uint8_t* tile, int base, float dep) {
+        uint32_t off = (uint32_t)((useful ? 0 : 1) * (int)sct);
+        asm volatile("" : "+v"(off) : "v"(dep));
+        const uint8_t* pv = tile + off + F::OFF_V;
+    #pragma unroll
+        for (int i = 0; i < 8; ++i) {
+          const uint8_t* p =
+              SP ? sel_ptr(base + 2 * i + (useful ? 0 : 1)) + F::OFF_V
+                 : pv + (int64_t)(2 * i) * sct;
+          if constexpr (VB == 4) {
+            vraw[i] = *(const uint2*)(p + 8 * rr);
+          } else {
+            vraw[i].x = *(const uint32_t*)(p + 4 * rr);
+            vraw[i].y = *(const uint32_t*)(p + D / 4 + 4 * (rr >> 1));
+          }
+        }
+      };
+      // Biased codebook bytes (x + 128) of a 16-dim step from its raw words.
+      auto rope_bytes = [&](uint2 raw, uint32_t o[4]) {
+        o[0] = octave::lut4(raw.x & 0x0F0F0F0Fu) ^ X;
+        o[1] = octave::lut4((raw.x >> 4) & 0x0F0F0F0Fu) ^ X;
+        o[2] = octave::lut4(raw.y & 0x0F0F0F0Fu) ^ X;
+        o[3] = octave::lut4((raw.y >> 4) & 0x0F0F0F0Fu) ^ X;
+      };
+      auto nope_bytes = [&](uint2 raw, uint32_t o[4]) {
+    #pragma unroll
+        for (int i = 0; i < 4; ++i)
+          o[i] = octave::lut3(((raw.x >> (2 * i)) & 0x03030303u) |
+                              (((raw.y >> i) & 0x01010101u) << 2)) ^
+                 X;
+      };
+
+      // V of a tile is loaded after its QK (softmax hides it): holding it
+      // through the QK, next to O, the four S tiles and the raw K, spilled.
+      const uint8_t* tcur = nullptr;
+      if (start < end) {
+        tcur = SP ? nullptr : tile_ptr(start);
+        load_k(tcur, start, 0.f);
+      }
+
+      for (int base = start; base < end; base += 16) {
+        auto lo16 = [](uint32_t x) {
+          return __half2float(__ushort_as_half((unsigned short)(x & 0xFFFF)));
+        };
+        auto hi16 = [](uint32_t x) {
+          return __half2float(__ushort_as_half((unsigned short)(x >> 16)));
+        };
+        const uint32_t sa_c = sc_a, sb_c = sc_b, sv_c = sc_v;
+
+        // QK. Iteration k feeds the useful lane's slot k and, through the
+        // swap, the partner's slot k.
+        f8v S[NG];
+    #pragma unroll
+        for (int g = 0; g < NG; ++g) S[g] = (f8v){0, 0, 0, 0, 0, 0, 0, 0};
+        int qoff = qv * QROW;
+        asm volatile("" : "+v"(qoff));
+        constexpr int kStepOwn[8] = {0, 1, 4, 5, 6, 7, 8, 9};
+        constexpr int kStepPar[8] = {2, 3, 10, 11, 12, 13, 14, 15};
+    #pragma unroll
+        for (int k = 0; k < 8; ++k) {
+          uint32_t wd[4];
+          if (k < NR4)
+            rope_bytes(kr_raw[k < NR4 ? k : 0], wd);
+          else
+            nope_bytes(kn_raw[k - NR4], wd);
+          const int so = KC ? k : kStepOwn[k], sp = KC ? 8 + k : kStepPar[k];
+          const int go = KC ? 0 : so / 4, gp = KC ? 0 : sp / 4;
+          // One A and one B fragment live at a time. Each read waits for the
+          // previous WMMA: without the dependence the compiler hoists all Q
+          // fragments and spills.
+          asm volatile("" : "+v"(qoff) : "v"(S[gp][0]));
+          {
+            const h16v b = ld16(&sQ[0][0] + qoff + 16 * so);
+            const h16v a = bytes_to_h16(wd[0], wd[1], wd[2], wd[3]);
+            S[go] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, S[go]);
+          }
+          __builtin_amdgcn_sched_barrier(0);
+          asm volatile("" : "+v"(qoff) : "v"(S[go][0]));
+          {
+            const h16v b = ld16(&sQ[0][0] + qoff + 16 * sp);
+            const h16v a = bytes_to_h16(swap1(wd[0]), swap1(wd[1]),
+                                        swap1(wd[2]), swap1(wd[3]));
+            S[gp] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, S[gp]);
+          }
+          __builtin_amdgcn_sched_barrier(0);
+        }
+
+        load_v(tcur, base, S[0][0]);
+        // Next tile's K and scales now: softmax and PV hide them.
+        const bool next = base + 16 < end;
+        const uint8_t* tn = next && !SP ? tile_ptr(base + 16) : nullptr;
+        if (next) load_k(tn, base + 16, S[0][0]);
+
+        float sc[8];
+        float mloc = -INFINITY;
+    #pragma unroll
+        for (int e = 0; e < 8; ++e) {
+          const int t = base + 2 * e + h;
+          bool ok = t < end && t < qlen;
+          if constexpr (SP) {
+            int off;
+            const int pb = sel_block(t, off);
+            ok = ok && pb >= 0 && !(wtags && in_window(t, pb, off));
+          }
+          const uint32_t a = pickw(sa_c, 2 * e);
+          float v;
+          if constexpr (KC) {
+            v = lo16(a) * (S[0][e] - bq[0]);
+          } else {
+            const uint32_t b = pickw(sb_c, 2 * e);
+            v = lo16(a) * (S[0][e] - bq[0]) + hi16(a) * (S[1][e] - bq[1]) +
+                lo16(b) * (S[2][e] - bq[2]) + hi16(b) * (S[3][e] - bq[3]);
+          }
+          sc[e] = ok ? v * sm_log2 : -INFINITY;
+          mloc = fmaxf(mloc, sc[e]);
+        }
+        const float mn = fmaxf(m, fmaxf(mloc, xhalf(mloc)));
+        float alpha = 1.0f;
+        if (mn > m)
+          alpha = (m == -INFINITY) ? 0.0f : __builtin_amdgcn_exp2f(m - mn);
+        m = mn;
+        uint32_t pd[4];
+        float psum = 0.0f, lsum = 0.0f;
+    #pragma unroll
+        for (int e = 0; e < 8; e += 2) {
+          const float v0 = lo16(pickw(sv_c, 2 * e));
+          const float v1 = lo16(pickw(sv_c, 2 * e + 2));
+          const float p0 =
+              sc[e] == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(sc[e] - mn);
+          const float p1 = sc[e + 1] == -INFINITY
+                               ? 0.0f
+                               : __builtin_amdgcn_exp2f(sc[e + 1] - mn);
+          lsum += p0 + p1;
+          const _Float16 h0 = (_Float16)(p0 * v0 * PSCALE);
+          const _Float16 h1 = (_Float16)(p1 * v1 * PSCALE);
+          psum += (float)h0 + (float)h1;
+          pd[e >> 1] = __builtin_bit_cast(uint32_t, (h2v){h0, h1});
+        }
+        l = l * alpha + lsum;
+        cp = cp * alpha + psum;
+        if (__builtin_amdgcn_ballot_w32(alpha != 1.0f)) {
+    #pragma unroll
+          for (int c = 0; c < 16; ++c) O[c] *= alpha;
+        }
+        u8v pf;
+    #pragma unroll
+        for (int i = 0; i < 4; ++i) {
+          const uint32_t o = xhalf(pd[i]);
+          pf[i] = h ? o : pd[i];
+          pf[4 + i] = h ? pd[i] : o;
+        }
+        const h16v pb16 = __builtin_bit_cast(h16v, pf);
+
+        // Biased codebook bytes of dims 4q..4q+3 of this lane's chunk,
+        // token i.
+        auto vq = [&](int i, int q) -> uint32_t {
+          if constexpr (VB == 4) {
+            const uint32_t w = q < 2 ? vraw[i].x : vraw[i].y;
+            return octave::lut4((w >> (4 * (q & 1))) & 0x0F0F0F0Fu) ^ X;
+          } else {
+            const uint32_t hi = vraw[i].y >> (4 * (rr & 1));
+            return octave::lut3(((vraw[i].x >> (2 * q)) & 0x03030303u) |
+                                (((hi >> q) & 0x01010101u) << 2)) ^
+                   X;
+          }
+        };
+    #pragma unroll
+        for (int q = 0; q < 4; ++q) {
+          uint32_t t0[4], t1[4];
+          tr4(vq(0, q), vq(1, q), vq(2, q), vq(3, q), t0);
+          tr4(vq(4, q), vq(5, q), vq(6, q), vq(7, q), t1);
+    #pragma unroll
+          for (int b = 0; b < 4; ++b) {
+            const int c = 4 * q + b;
+            uint32_t o0 = t0[b], o1 = t1[b];
+            if (c >= 2)
+              asm volatile("" : "+v"(o0), "+v"(o1) : "v"(O[c - 2][0]));
+            const h16v a = bytes_to_h16(o0, o1, swap1(o0), swap1(o1));
+            O[c] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, pb16, O[c]);
+          }
+        }
+        tcur = tn;
+      }
+
+      l += xhalf(l);
+      cp += xhalf(cp);
+      if (qlive) {
+        int qx = qv;
+        asm volatile("" : "+v"(qx));
+        float* op = mid_o + (int64_t)(grp * QG + qx / hpk) * smo +
+                    (int64_t)(kvh * hpk + qx % hpk) * smh + (int64_t)si * sms;
+        const bool any = m != -INFINITY;
+        const float corr = BIAS * cp;
+        constexpr float kOut = 1.0f / (PSCALE * LUT_ONE);
+    #pragma unroll
+        for (int e = 0; e < 8; ++e) {
+          float* dp = op + 16 * (2 * e + h);
+    #pragma unroll
+          for (int c = 0; c < 16; c += 4) {
+            float4 v;
+            v.x = any ? (O[c][e] - corr) * kOut : 0.0f;
+            v.y = any ? (O[c + 1][e] - corr) * kOut : 0.0f;
+            v.z = any ? (O[c + 2][e] - corr) * kOut : 0.0f;
+            v.w = any ? (O[c + 3][e] - corr) * kOut : 0.0f;
+            *(float4*)(dp + c) = v;
+          }
+        }
+        if (h == 0) {
+          op[D] = m;
+          op[D + 1] = l;
+        }
+      }
+    }
+    seg0 = seg1;
+  }
+  #endif  // OCTAVE_WMMA_STUB
+}
+
+}  // namespace octave_wmma
+
+// Splits per row group and minimum tokens per split, as in the int8 kernel.
+constexpr int kSqWmmaSplits = 192;
+constexpr int kSqWmmaMinTps = 64;
+constexpr int kSqWmmaNsb = 2;
+
+int octave_decode_mfma(torch::Tensor query, torch::Tensor cache,
+                       torch::Tensor block_table, torch::Tensor q_to_req,
+                       torch::Tensor q_to_klen, torch::Tensor mid_o,
+                       torch::Tensor k_signs, double sm_scale,
+                       int64_t num_kv_splits, int64_t fmt, int64_t query_group);
+
+// Launches the WMMA decode when the shape is covered; returns the number of
+// splits it wrote (for the reduce), or 0 when the caller must use its own
+// kernel.
+// With `indices` ([num_q, topk] int32 logical positions) the sparse variant
+// runs and q_to_klen is ignored.
+int octave_decode_wmma(torch::Tensor query, torch::Tensor cache,
+                       torch::Tensor block_table, torch::Tensor q_to_req,
+                       torch::Tensor q_to_klen, torch::Tensor mid_o,
+                       torch::Tensor k_signs, double sm_scale,
+                       int64_t num_kv_splits, int64_t fmt,
+                       const std::optional<torch::Tensor>& indices,
+                       const int* qpos, const int* wtags, int wmask,
+                       const int* stags, int smask, int64_t query_group) {
+  using namespace octave_wmma;
+  // 0: gfx1100 (this file), 1: gfx942 (octave_decode_mfma_cdna.cu, dense
+  // only).
+  static const int arch = [] {
+    const std::string name =
+        at::cuda::getCurrentDeviceProperties()->gcnArchName;
+    if (name.rfind("gfx1100", 0) == 0) return 0;
+    if (name.rfind("gfx942", 0) == 0) return 1;
+    return -1;
+  }();
+  if (arch == 1)
+    return indices.has_value()
+               ? 0
+               : octave_decode_mfma(query, cache, block_table, q_to_req,
+                                    q_to_klen, mid_o, k_signs, sm_scale,
+                                    num_kv_splits, fmt, query_group);
+  const bool arch_ok = arch == 0;
+  const int num_q = query.size(0);
+  const int num_q_heads = query.size(1);
+  const int num_kv_heads = cache.size(1);
+  const int hpk = num_q_heads / num_kv_heads;
+  const bool eligible =
+      arch_ok && query.size(2) == D && query.stride(2) == 1 && hpk * QG <= 32 &&
+      cache.size(2) % 16 == 0 &&
+      cache.stride(2) % (fmt & octave::kCompactFlag ? 4 : 8) == 0 &&
+      num_kv_splits >= 1;
+  if (!eligible) return 0;
+  const int ns = std::min<int>(kSqWmmaSplits, (int)num_kv_splits);
+  if (num_q == 0) return ns;
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const bool sp = indices.has_value();
+  const int* idx = sp ? indices->data_ptr<int>() : nullptr;
+  const int64_t sidx = sp ? indices->stride(0) : 0;
+  const int topk = sp ? (int)indices->size(1) : 0;
+  const int nrt = std::min(num_q, QG) * hpk > 16 ? 2 : 1;
+  const int spb = 2 * kSqWmmaNsb / nrt;
+  dim3 grid((num_q + QG - 1) / QG, num_kv_heads, (ns + spb - 1) / spb);
+  #define SQW_SP(B, KC, T, SP)                                              \
+    decode_wmma<B, KC, kSqWmmaNsb, T, SP>                                   \
+        <<<grid, 64 * kSqWmmaNsb, 0, stream>>>(                             \
+            (const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(),   \
+            block_table.data_ptr<int>(), q_to_req.data_ptr<int>(),          \
+            q_to_klen.data_ptr<int>(), mid_o.data_ptr<float>(),             \
+            k_signs.data_ptr<int>(), (float)sm_scale, num_q, num_q_heads,   \
+            num_kv_heads, cache.size(2), block_table.size(1),               \
+            block_table.size(0), cache.size(0), ns, kSqWmmaMinTps, nrt,     \
+            query.stride(0), query.stride(1), cache.stride(0),              \
+            cache.stride(1), cache.stride(2), mid_o.stride(0),              \
+            mid_o.stride(1), mid_o.stride(2), idx, sidx, topk, qpos, wtags, \
+            wmask, stags, smask)
+  #define SQW(B, KC, T)          \
+    do {                         \
+      if (sp)                    \
+        SQW_SP(B, KC, T, true);  \
+      else                       \
+        SQW_SP(B, KC, T, false); \
+    } while (0)
+  const bool bf = query.dtype() == at::kBFloat16;
+  if (fmt == 4) {
+    if (bf)
+      SQW(4, false, __hip_bfloat16);
+    else
+      SQW(4, false, half);
+  } else if (fmt == 3) {
+    if (bf)
+      SQW(3, false, __hip_bfloat16);
+    else
+      SQW(3, false, half);
+  } else {
+    if (bf)
+      SQW(3, true, __hip_bfloat16);
+    else
+      SQW(3, true, half);
+  }
+  #undef SQW
+  #undef SQW_SP
+  return ns;
+}
+
+#endif  // USE_ROCM

csrc/attention/octave_format.cuh

@@ -0,0 +1,78 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Byte layout of one Octave slot (one token, one KV head). The PyTorch
+// reference in vllm/v1/attention/ops/rocm_octave.py is the source of truth.
+//
+//   [0, OFF_KN)       K RoPE block (dims 0-63), rotated, 4-bit codes
+//   [OFF_KN, OFF_V)   K NoPE blocks (dims 64-255), rotated, 3-bit codes
+//   [OFF_V, OFF_SC)   V, rotated, VBITS-bit codes
+//   [OFF_SC, ...)     fp16 scales: K block 0..3, V
+//
+// Compact K (KC): K rotated over all 256 dims like V, 3-bit codes for every
+// dim (OFF_KN = 0), one K scale; slots pad to 4 bytes instead of 8.
+//
+// 4-bit codes: dims 8i..8i+7 in one dword, low nibbles = dims 0-3, high
+// nibbles = dims 4-7. 3-bit codes: a 2-bit plane (16 dims per dword, dims
+// 4k..4k+3 at bit 2k of bytes 0-3) followed by a 1-bit plane (32 dims per
+// dword, dims 4k..4k+3 at bit k of bytes 0-3). A code c of a block with
+// scale s stands for LUT[c] * s / LUT_ONE, LUT being the int8 Lloyd-Max
+// codebook of its width.
+
+#pragma once
+
+#include <cstdint>
+
+namespace octave {
+
+constexpr int D = 256;       // head size
+constexpr int R = 64;        // RoPE dims (K block 0)
+constexpr int N = D - R;     // NoPE dims (K blocks 1-3)
+constexpr int NKB = D / 64;  // K Hadamard blocks
+constexpr float LUT_ONE = 64.0f;
+
+template <int VBITS, bool KC = false>
+struct Format {
+  static constexpr int KR = KC ? 0 : R;  // K dims with 4-bit codes
+  static constexpr int KN = D - KR;      // K dims with 3-bit codes
+  static constexpr int OFF_KN = KR * 4 / 8;
+  static constexpr int OFF_V = OFF_KN + KN * 3 / 8;
+  static constexpr int OFF_SC = OFF_V + D * VBITS / 8;
+  static constexpr int NUM_KSC = KC ? 1 : NKB;
+  static constexpr int NUM_SC = NUM_KSC + 1;
+  static constexpr int ALIGN = KC ? 4 : 8;
+  static constexpr int SLOT = (OFF_SC + 2 * NUM_SC + ALIGN - 1) / ALIGN * ALIGN;
+};
+
+// Format id passed by the host: V bits, plus 16 for compact K.
+constexpr int kCompactFlag = 16;
+
+// int8 Lloyd-Max codebooks as perm tables: byte c of the table is LUT[c].
+//   3-bit: -127 -79 -45 -14 14 45 79 127
+//   4-bit: -126 -95 -74 -58 -43 -30 -18 -6 6 18 30 43 58 74 95 126
+constexpr uint32_t LUT3_LO = 0xF2D3B181u, LUT3_HI = 0x7F4F2D0Eu;
+constexpr uint32_t LUT4_0 = 0xC6B6A182u, LUT4_1 = 0xFAEEE2D5u;
+constexpr uint32_t LUT4_2 = 0x2B1E1206u, LUT4_3 = 0x7E5F4A3Au;
+
+// Four 3-bit codes (one per byte) -> their four int8 codebook values.
+__device__ __forceinline__ uint32_t lut3(uint32_t codes) {
+  return __builtin_amdgcn_perm(LUT3_HI, LUT3_LO, codes);
+}
+// Four 4-bit codes (one per byte) -> their four int8 codebook values.
+__device__ __forceinline__ uint32_t lut4(uint32_t codes) {
+  const uint32_t lo =
+      __builtin_amdgcn_perm(LUT4_1, LUT4_0, codes & 0x07070707u);
+  const uint32_t hi =
+      __builtin_amdgcn_perm(LUT4_3, LUT4_2, codes & 0x07070707u);
+  const uint32_t m = ((codes >> 3) & 0x01010101u) * 0xFFu;
+  return (hi & m) | (lo & ~m);
+}
+template <int BITS>
+__device__ __forceinline__ uint32_t lut(uint32_t codes) {
+  if constexpr (BITS == 4)
+    return lut4(codes);
+  else
+    return lut3(codes);
+}
+
+}  // namespace octave

csrc/attention/octave_prefill_cdna.cu

@@ -0,0 +1,666 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Octave prefill for CDNA3 (gfx942), head size 256, fp16. Same contract as
+// octave_prefill_rdna3.cu: reads the packed cache directly, everything in the
+// rotated space, phase 1 writes (O, m, l) partials over the cached prefix and
+// phase 2 runs causal attention over the chunk and merges.
+//
+// Both phases compute transposed tiles so that every MFMA result is already in
+// the layout the next MFMA consumes (16x16 MFMA, wave64: lane l holds column
+// l % 16 and rows 4 (l / 16) .. 4 (l / 16) + 3):
+//   S^T[token][q] = K x Q^T      (phase 1: int8 16x16x32, phase 2: fp16)
+//   O^T[dim][q]  += V^T x P^T    (fp16 16x16x16)
+// so P^T comes straight out of the S^T accumulator, one query per lane.
+//
+// Phase 1: one block = 128 query vectors (token, head) of one KV head, 8 waves
+// of 16 query vectors x 256 dims. 256 threads unpack the next 16-token tile
+// into LDS while the current one is consumed.
+
+#include <cstdint>
+#include <algorithm>
+
+#include <torch/all.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#if defined(USE_ROCM)
+  #include <hip/hip_runtime.h>
+  #include <hip/hip_fp16.h>
+
+  #include "octave_format.cuh"
+
+  #if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx942__)
+    #define OCTAVE_CDNA_STUB
+  #endif
+
+namespace octave_cdna {
+
+using octave::D;
+using octave::Format;
+using octave::LUT_ONE;
+
+typedef _Float16 h4v __attribute__((ext_vector_type(4)));
+typedef _Float16 h2v __attribute__((ext_vector_type(2)));
+typedef float f4v __attribute__((ext_vector_type(4)));
+typedef int i4v __attribute__((ext_vector_type(4)));
+typedef uint32_t u4v __attribute__((ext_vector_type(4)));
+typedef uint32_t u2v __attribute__((ext_vector_type(2)));
+
+constexpr int KT = 16;                       // tokens per tile
+[[maybe_unused]] constexpr int FR = D / 16;  // 16-dim output fragments
+// Both phases run 8 waves (128 query vectors of one KV head) per block: every
+// block reads the whole K/V it attends to, so more query vectors per block
+// means less of that work per query.
+constexpr int NW1 = 8;
+constexpr int NT1 = NW1 * 64;
+constexpr int QB1 = NW1 * 16;
+constexpr int KT1 = 32;  // phase-1 tile: two 16-token sub-tiles
+[[maybe_unused]] constexpr int VST1 = KT1 + 4;
+[[maybe_unused]] constexpr int KROW = D + 16;  // bytes per int8 K row
+[[maybe_unused]] constexpr int VST = KT + 4;   // halves per V^T row
+[[maybe_unused]] constexpr int HROW = D + 8;   // halves per fp16 K row
+// Phase 2 LDS rows, picked for conflict-free banks: a K row is 2 dwords
+// mod 32, and V^T rows of 40 halves with the 4-token group XORed by
+// (row >> 3) & 7 spread both the transposed stores and the MFMA reads.
+[[maybe_unused]] constexpr int KROW2 = D + 4;
+[[maybe_unused]] constexpr int VST2 = KT1 + 8;
+[[maybe_unused]] constexpr float PSCALE = 4096.0f;
+
+constexpr int kMinSplitTokens = 2048;
+__device__ __forceinline__ int used_splits(int ctx, int num_splits) {
+  if (ctx <= 0) return 0;
+  return min(num_splits, (ctx + kMinSplitTokens - 1) / kMinSplitTokens);
+}
+__device__ __forceinline__ int split_len(int ctx, int used) {
+  return ((ctx + used - 1) / used + 31) / 32 * 32;
+}
+
+// Reductions over the four lanes of a query (l, l ^ 16, l ^ 32, l ^ 48).
+__device__ __forceinline__ float qmax4(float v) {
+  v = fmaxf(v, __shfl_xor(v, 16));
+  return fmaxf(v, __shfl_xor(v, 32));
+}
+__device__ __forceinline__ float qsum4(float v) {
+  v += __shfl_xor(v, 16);
+  return v + __shfl_xor(v, 32);
+}
+
+// Codes of dims 4k..4k+3 of a 64-dim block (see octave_prefill_rdna3.cu).
+template <int BITS>
+__device__ __forceinline__ uint32_t block_codes(const uint32_t* raw,
+                                                uint32_t hi, int k) {
+  if constexpr (BITS == 4) {
+    return (raw[k >> 1] >> (4 * (k & 1))) & 0x0F0F0F0Fu;
+  } else {
+    return ((raw[0] >> (2 * (k & 3))) & 0x03030303u) |
+           (((hi >> (k & 7)) & 0x01010101u) << 2);
+  }
+}
+
+// 4 dwords, one per token (same 4 dims) -> 4 dwords, one per dim (4 tokens).
+__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2,
+                                    uint32_t a3, uint32_t o[4]) {
+  const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u);
+  const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u);
+  const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u);
+  const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u);
+  o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u);
+  o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u);
+  o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u);
+  o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u);
+}
+
+// ===========================================================================
+// Phase 1: prefix
+// ===========================================================================
+template <int VB, bool KC>
+__global__ __launch_bounds__(NT1) void prefix_attn(
+    const __half* __restrict__ q, const uint8_t* __restrict__ cache,
+    const int* __restrict__ block_table, const int* __restrict__ cu_seqlens_q,
+    const int* __restrict__ seq_lens, float* __restrict__ ws, float sm_scale,
+    int num_q_heads, int num_kv_heads, int block_size, int max_blocks,
+    int num_splits, int total_q_tokens, int64_t sq_tok, int64_t sq_head,
+    int64_t scb, int64_t sch, int64_t sct) {
+  #ifndef OCTAVE_CDNA_STUB
+  using F = Format<VB, KC>;
+  const int seq = blockIdx.x, kvh = blockIdx.y;
+  const int rowtile = blockIdx.z / num_splits, split = blockIdx.z % num_splits;
+  const int tid = threadIdx.x;
+  const int w = tid >> 6, lane = tid & 63;
+  const int col = lane & 15, gq = lane >> 4;
+  const int hpk = num_q_heads / num_kv_heads;
+
+  const int q0 = cu_seqlens_q[seq];
+  const int qlen = cu_seqlens_q[seq + 1] - q0;
+  int ctx = seq_lens[seq] - qlen;
+  ctx = max(0, min(ctx, max_blocks * block_size));
+  const int nqv = qlen * hpk;
+  if (rowtile * QB1 >= nqv) return;
+
+  const int used = used_splits(ctx, num_splits);
+  if (split >= used) return;
+  const int k_len = split_len(ctx, used);
+  const int k_begin = min(ctx, split * k_len);
+  const int k_stop = min(ctx, k_begin + k_len);
+  const int ntiles = k_stop > k_begin ? (k_stop - k_begin + KT1 - 1) / KT1 : 0;
+
+  __shared__ __attribute__((aligned(16))) int8_t sK[2][KT1][KROW];
+  __shared__ __attribute__((aligned(16))) _Float16 sV[2][D * VST1];
+  __shared__ float sKs[2][4][KT1], sVs[2][KT1];
+
+  // ---- loader mapping. Thread half sub (256 threads each) loads the 16-token
+  // sub-tile sub of a tile; within it, K: dims 16 kp .. + 15 of block kq of
+  // token tk; V: dims vd + 4 vh .. + 3 of token quad vq.
+  const int sub = tid >> 8, lt = tid & 255;
+  const int tk = lt >> 4, kq = (lt >> 2) & 3, kp = lt & 3;
+  const int vu = lt >> 1, vh = lt & 1, vq = vu & 3, vd = (vu >> 2) * 8;
+  const bool k3 = KC || kq > 0;
+  const int b3 = KC ? kq : kq - 1;
+  const int koff0 = k3 ? F::OFF_KN + 16 * b3 + 4 * kp : 8 * kp;
+  const int koff1 =
+      k3 ? F::OFF_KN + F::KN / 4 + 8 * b3 + 4 * (kp >> 1) : 8 * kp + 4;
+  constexpr int VW = VB == 4 ? 1 : 2;
+  const int tlast = k_stop - 1;
+  const uint8_t* const head = cache + (int64_t)kvh * sch;
+  const int* const bt = block_table + seq * max_blocks;
+
+  struct Stg {
+    uint32_t k0, k1;
+    uint32_t v[4][VW];
+    uint32_t s0, s1;
+  };
+  // Sub-tiles start 16-aligned and sit in one cache block (block sizes are
+  // multiples of 16), so all 16 slots exist: tokens past k_stop are read as
+  // they are and masked (K to -inf, P to 0 by selection, never by product).
+  // Slot offsets inside a sub-tile are fixed per thread.
+  auto sub_start = [&](int i) {
+    return __builtin_amdgcn_readfirstlane(
+        min(k_begin + i * KT1 + 16 * sub, tlast & ~15));
+  };
+  const uint32_t ok0 = tk * (uint32_t)sct + koff0;
+  const uint32_t ok1 = tk * (uint32_t)sct + koff1;
+  uint32_t ov[4];
+    #pragma unroll
+  for (int r = 0; r < 4; ++r) {
+    ov[r] = (4 * vq + r) * (uint32_t)sct + F::OFF_V +
+            (VB == 4 ? vd / 2 : 4 * (vd >> 4));
+  }
+  const uint32_t os =
+      (lt & 15) * (uint32_t)sct + F::OFF_SC + (!KC && lt >= 16 ? 8 : 0);
+  auto fetch = [&](int i, Stg& S) {
+    const int start = sub_start(i);
+    const int pb = __builtin_amdgcn_readfirstlane(bt[start / block_size]);
+    const uint8_t* base = head + (int64_t)(pb < 0 ? 0 : pb) * scb +
+                          (int64_t)(start % block_size) * sct;
+    S.k0 = *(const uint32_t*)(base + ok0);
+    S.k1 = *(const uint32_t*)(base + ok1);
+    #pragma unroll
+    for (int r = 0; r < 4; ++r) {
+      S.v[r][0] = *(const uint32_t*)(base + ov[r]);
+      if constexpr (VB != 4)
+        S.v[r][1] = *(const uint32_t*)(base + ov[r] - 4 * (vd >> 4) + D / 4 +
+                                       4 * (vd >> 5));
+    }
+    if (lt < 32) {
+      if constexpr (KC) {
+        S.s0 = *(const uint32_t*)(base + os);
+      } else {
+        const u2v x = *(const u2v*)(base + os);
+        S.s0 = x[0];
+        S.s1 = x[1];
+      }
+    }
+  };
+  auto store = [&](const Stg& S, int bf) {
+    const int t0 = 16 * sub;
+    {
+      const uint32_t r4[2] = {S.k0, S.k1};
+      const uint32_t r3[1] = {S.k0};
+      const uint32_t hi = S.k1 >> (4 * (kp & 1));
+      u4v o;
+    #pragma unroll
+      for (int x = 0; x < 4; ++x)
+        o[x] = k3 ? octave::lut3(block_codes<3>(r3, hi, x))
+                  : octave::lut4(block_codes<4>(r4, 0, x));
+      *(u4v*)&sK[bf][t0 + tk][64 * kq + 16 * kp] = o;
+    }
+    uint32_t c[4];
+    #pragma unroll
+    for (int r = 0; r < 4; ++r) {
+      if constexpr (VB == 4) {
+        c[r] = (S.v[r][0] >> (4 * vh)) & 0x0F0F0F0Fu;
+      } else {
+        const int sl = 2 * ((vd & 15) >> 2) + 2 * vh;
+        const int sh = ((vd & 31) >> 2) + vh;
+        c[r] = ((S.v[r][0] >> sl) & 0x03030303u) |
+               (((S.v[r][1] >> sh) & 0x01010101u) << 2);
+      }
+      c[r] = octave::lut<VB>(c[r]) ^ 0x80808080u;
+    }
+    uint32_t t[4];
+    tr4(c[0], c[1], c[2], c[3], t);
+    // Codebook bytes fed as fp16 1024 + (x + 128), less 1152: exact x.
+    const h2v kofs = {(_Float16)1152.0f, (_Float16)1152.0f};
+    #pragma unroll
+    for (int d = 0; d < 4; ++d) {
+      const int row = vd + 4 * vh + d;
+      const h2v p01 = __builtin_bit_cast(
+          h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07010700u));
+      const h2v p23 = __builtin_bit_cast(
+          h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07030702u));
+      *(u2v*)&sV[bf][row * VST1 + t0 + 4 * vq] =
+          (u2v){__builtin_bit_cast(uint32_t, p01 - kofs),
+                __builtin_bit_cast(uint32_t, p23 - kofs)};
+    }
+    if (lt < 16) {
+      const h2v a = __builtin_bit_cast(h2v, S.s0);
+      sKs[bf][0][t0 + lt] = (float)a[0];
+      if constexpr (!KC) {
+        const h2v cc = __builtin_bit_cast(h2v, S.s1);
+        sKs[bf][1][t0 + lt] = (float)a[1];
+        sKs[bf][2][t0 + lt] = (float)cc[0];
+        sKs[bf][3][t0 + lt] = (float)cc[1];
+      }
+    } else if (lt < 32) {
+      sVs[bf][t0 + lt - 16] =
+          (float)__builtin_bit_cast(_Float16,
+                                    (uint16_t)(KC ? S.s0 >> 16 : S.s0)) *
+          (PSCALE / LUT_ONE);
+    }
+  };
+
+  // ---- this lane's query vector, int8 over dims 32 s + 8 gq .. + 7
+  const int qv = rowtile * QB1 + w * 16 + col;
+  const bool qlive = qv < nqv;
+  long Q[8];
+  float qsc;
+  {
+    const int qtok = qlive ? qv / hpk : 0;
+    const int qhead = kvh * hpk + (qlive ? qv % hpk : 0);
+    const __half* qrow =
+        q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)qhead * sq_head;
+    _Float16 x[8][8];
+    float qmax = 1e-8f;
+    #pragma unroll
+    for (int s = 0; s < 8; ++s) {
+      const u4v r = qlive ? *(const u4v*)(qrow + 32 * s + 8 * gq) : (u4v){};
+      __builtin_memcpy(x[s], &r, 16);
+    #pragma unroll
+      for (int e = 0; e < 8; ++e) qmax = fmaxf(qmax, fabsf((float)x[s][e]));
+    }
+    qmax = qmax4(qmax);
+    const float qinv = 127.0f / qmax;
+    qsc = qmax * (1.0f / 127.0f);
+    #pragma unroll
+    for (int s = 0; s < 8; ++s) {
+      uint64_t word = 0;
+    #pragma unroll
+      for (int e = 0; e < 8; ++e) {
+        int v = (int)lrintf((float)x[s][e] * qinv);
+        v = max(-127, min(127, v));
+        word |= ((uint64_t)(uint8_t)(int8_t)v) << (8 * e);
+      }
+      Q[s] = (long)word;
+    }
+  }
+
+  f4v O[FR];
+    #pragma unroll
+  for (int c = 0; c < FR; ++c) O[c] = (f4v){0, 0, 0, 0};
+  const float c1 = sm_scale * qsc * (1.4426950408889634f / LUT_ONE);
+  float m = -INFINITY, l = 0.f;
+
+  // Pipeline: iteration i computes tile i from LDS while the loads of tile
+  // i + 1 are in flight, then stores them into the other buffer.
+  Stg S;
+  if (ntiles > 0) {
+    fetch(0, S);
+    store(S, 0);
+  }
+  __syncthreads();
+  for (int i = 0; i < ntiles; ++i) {
+    const int bf = i & 1;
+    if (i + 1 < ntiles) fetch(i + 1, S);
+
+    // S^T = K x Q^T over both sub-tiles first (independent MFMA chains),
+    // then one online-softmax step over the 32 tokens, then PV.
+    f4v Sf[2];
+    #pragma unroll
+    for (int j = 0; j < 2; ++j) {
+      Sf[j] = (f4v){0, 0, 0, 0};
+      const int8_t* krow = &sK[bf][16 * j + col][8 * gq];
+      const int tq = 16 * j + 4 * gq;
+      if constexpr (KC) {
+        i4v Si = {0, 0, 0, 0};
+    #pragma unroll
+        for (int s = 0; s < 8; ++s)
+          Si = __builtin_amdgcn_mfma_i32_16x16x32_i8(
+              *(const long*)(krow + 32 * s), Q[s], Si, 0, 0, 0);
+    #pragma unroll
+        for (int e = 0; e < 4; ++e)
+          Sf[j][e] = sKs[bf][0][tq + e] * (float)Si[e];
+      } else {
+    #pragma unroll
+        for (int b = 0; b < 4; ++b) {
+          i4v Si = {0, 0, 0, 0};
+    #pragma unroll
+          for (int s = 2 * b; s < 2 * b + 2; ++s)
+            Si = __builtin_amdgcn_mfma_i32_16x16x32_i8(
+                *(const long*)(krow + 32 * s), Q[s], Si, 0, 0, 0);
+    #pragma unroll
+          for (int e = 0; e < 4; ++e)
+            Sf[j][e] += sKs[bf][b][tq + e] * (float)Si[e];
+        }
+      }
+    }
+
+    // Online softmax in log2 units; this lane holds tokens 16 j + 4 gq + e.
+    float sc[2][4];
+    float mloc = -INFINITY;
+    #pragma unroll
+    for (int j = 0; j < 2; ++j)
+    #pragma unroll
+      for (int e = 0; e < 4; ++e) {
+        const int t = k_begin + i * KT1 + 16 * j + 4 * gq + e;
+        sc[j][e] = t < k_stop ? Sf[j][e] * c1 : -INFINITY;
+        mloc = fmaxf(mloc, sc[j][e]);
+      }
+    const float mn = fmaxf(m, qmax4(mloc));
+    const float alpha = __builtin_amdgcn_exp2f(m - mn);
+    m = mn;
+    h4v pb[2];
+    float lsum = 0.f;
+    #pragma unroll
+    for (int j = 0; j < 2; ++j)
+    #pragma unroll
+      for (int e = 0; e < 4; ++e) {
+        const bool live = sc[j][e] != -INFINITY;
+        const float p = live ? __builtin_amdgcn_exp2f(sc[j][e] - mn) : 0.f;
+        lsum += p;
+        pb[j][e] =
+            live ? (_Float16)(p * sVs[bf][16 * j + 4 * gq + e]) : (_Float16)0.f;
+      }
+    l = l * alpha + lsum;
+    if (__ballot(alpha != 1.0f)) {
+    #pragma unroll
+      for (int c = 0; c < FR; ++c) O[c] *= alpha;
+    }
+
+    // O^T += V^T x P^T: A = V^T[dim 16 c + col][tokens 16 j + 4 gq .. + 3].
+    #pragma unroll
+    for (int j = 0; j < 2; ++j) {
+      const _Float16* vrow = &sV[bf][col * VST1 + 16 * j + 4 * gq];
+    #pragma unroll
+      for (int c = 0; c < FR; ++c)
+        O[c] = __builtin_amdgcn_mfma_f32_16x16x16f16(
+            *(const h4v*)(vrow + 16 * c * VST1), pb[j], O[c], 0, 0, 0);
+    }
+
+    if (i + 1 < ntiles) store(S, bf ^ 1);
+    __syncthreads();
+  }
+
+  // ---- partials for the phase-2 merge (natural-log max, unnormalized O)
+  l = qsum4(l);
+  if (!qlive) return;
+  const int qtok = qv / hpk, qhead = kvh * hpk + qv % hpk;
+  float* wp = ws + (int64_t)split * total_q_tokens * num_q_heads * (D + 2) +
+              ((int64_t)(q0 + qtok) * num_q_heads + qhead) * (D + 2);
+    #pragma unroll
+  for (int c = 0; c < FR; ++c)
+    #pragma unroll
+    for (int e = 0; e < 4; ++e)
+      wp[16 * c + 4 * gq + e] = O[c][e] * (1.0f / PSCALE);
+  if (gq == 0) {
+    wp[D] = m * 0.6931471805599453f;
+    wp[D + 1] = l;
+  }
+  #endif  // OCTAVE_CDNA_STUB
+}
+
+// ===========================================================================
+// Phase 2: the chunk's own K/V (fp16, causal), merged with the prefix
+// partials. Like phase 1, one block = 128 query vectors (token, head) of one
+// KV head, so the GQA heads share every K/V tile; 8 waves of 16 vectors.
+// ===========================================================================
+__global__ void __launch_bounds__(NT1)
+    chunk_attn(__half* __restrict__ out, const float* __restrict__ ws,
+               const __half* __restrict__ q, const __half* __restrict__ k,
+               const __half* __restrict__ v,
+               const int* __restrict__ cu_seqlens_q,
+               const int* __restrict__ seq_lens, int num_q_heads,
+               int num_kv_heads, int max_ctx, float sm_scale, int num_splits,
+               int total_q_tokens, int64_t sq_tok, int64_t sq_head,
+               int64_t sk_tok, int64_t sk_head, int64_t sv_tok, int64_t sv_head,
+               int64_t so_tok, int64_t so_head) {
+  #ifndef OCTAVE_CDNA_STUB
+  constexpr float LOG2E = 1.4426950408889634f;
+  const int seq = blockIdx.x, kvh = blockIdx.y, rowtile = blockIdx.z;
+  const int tid = threadIdx.x, w = tid >> 6, lane = tid & 63;
+  const int col = lane & 15, gq = lane >> 4;
+  const int hpk = num_q_heads / num_kv_heads;
+  const int q0 = cu_seqlens_q[seq];
+  const int qlen = cu_seqlens_q[seq + 1] - q0;
+  const int nqv = qlen * hpk;
+  if (rowtile * QB1 >= nqv) return;
+  const int ctx = max(0, min(seq_lens[seq] - qlen, max_ctx));
+  const int qv = rowtile * QB1 + w * 16 + col;
+  const bool qlive = qv < nqv;
+  const int qtok = qlive ? qv / hpk : 0;
+  const int head = kvh * hpk + (qlive ? qv % hpk : 0);
+  // Keys this block and this wave need (causal: up to their last token).
+  const int k_end = min(qlen, (min(nqv, (rowtile + 1) * QB1) - 1) / hpk + 1);
+  const int wave_first = rowtile * QB1 + w * 16;
+  const bool wave_live = wave_first < nqv;
+  const int wave_last_tok = min(nqv - 1, wave_first + 15) / hpk;
+
+  // One 32-token tile in LDS (two would not fit in 64 KB with fp16 K and V);
+  // the next tile waits in registers.
+  __shared__ __attribute__((aligned(16))) _Float16 sK[KT1 * KROW2];
+  __shared__ __attribute__((aligned(16))) _Float16 sV[D * VST2];
+
+  // Merge the prefix splits (log2 units from here on). O^T fragment c holds
+  // dims 16 c + 4 gq + e of this lane's vector; l stays partial per lane.
+  f4v acc[FR];
+    #pragma unroll
+  for (int c = 0; c < FR; ++c) acc[c] = (f4v){0, 0, 0, 0};
+  float m = -INFINITY, l = 0.f;
+  const int64_t split_stride = (int64_t)total_q_tokens * num_q_heads * (D + 2);
+  const int used = used_splits(ctx, num_splits);
+  if (qlive && used > 0) {
+    const float* w0 =
+        ws + ((int64_t)(q0 + qtok) * num_q_heads + head) * (D + 2);
+    for (int s = 0; s < used; ++s)
+      m = fmaxf(m, w0[s * split_stride + D] * LOG2E);
+    for (int s = 0; s < used; ++s) {
+      const float* wp = w0 + s * split_stride;
+      if (wp[D] == -INFINITY) continue;
+      const float wgt = __builtin_amdgcn_exp2f(wp[D] * LOG2E - m);
+      l += wp[D + 1] * wgt;
+    #pragma unroll
+      for (int c = 0; c < FR; ++c)
+    #pragma unroll
+        for (int e = 0; e < 4; ++e) acc[c][e] += wp[16 * c + 4 * gq + e] * wgt;
+    }
+    if (gq != 0) l = 0.f;
+  }
+
+  // Q^T operand: dims 16 s + 4 gq .. + 3 of this lane's vector.
+  h4v Qh[FR];
+  {
+    const __half* qrow =
+        q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)head * sq_head;
+    #pragma unroll
+    for (int s = 0; s < FR; ++s) Qh[s] = *(const h4v*)(qrow + 16 * s + 4 * gq);
+  }
+  const float c2 = sm_scale * LOG2E;
+
+  // Loaders. K: thread t moves dims 8 (t % 32) .. + 7 of tokens t / 32
+  // and t / 32 + 16 of the 32-token tile. V: dims 4 (t % 64) .. + 3 of
+  // tokens 4 (t / 64) .. + 3, transposed in registers into V^T rows. Rows
+  // past the chunk are zeros.
+  const int lt_tok = tid >> 5, lt_d = (tid & 31) * 8;
+  const int vt_d = (tid & 63) * 4, vt_g = tid >> 6;
+  const __half* kbase = k + (int64_t)kvh * sk_head + lt_d;
+  const __half* vbase = v + (int64_t)kvh * sv_head + vt_d;
+  u4v Kr[2];
+  u2v Vr[4];
+  auto fetch = [&](int n0) {
+    #pragma unroll
+    for (int h = 0; h < 2; ++h) {
+      const int t = n0 + lt_tok + 16 * h;
+      const bool ok = t < qlen;
+      const int64_t r = q0 + (ok ? t : 0);
+      Kr[h] = ok ? *(const u4v*)(kbase + r * sk_tok) : (u4v){0, 0, 0, 0};
+    }
+    #pragma unroll
+    for (int e = 0; e < 4; ++e) {
+      const int t = n0 + 4 * vt_g + e;
+      const bool ok = t < qlen;
+      const int64_t r = q0 + (ok ? t : 0);
+      Vr[e] = ok ? *(const u2v*)(vbase + r * sv_tok) : (u2v){0, 0};
+    }
+  };
+  auto store = [&]() {
+    #pragma unroll
+    for (int h = 0; h < 2; ++h) {
+      _Float16* kp = &sK[(lt_tok + 16 * h) * KROW2 + lt_d];
+      *(u2v*)kp = (u2v){Kr[h][0], Kr[h][1]};
+      *(u2v*)(kp + 4) = (u2v){Kr[h][2], Kr[h][3]};
+    }
+    #pragma unroll
+    for (int d = 0; d < 4; ++d) {
+      const uint32_t sel = (d & 1) ? 0x07060302u : 0x05040100u;
+      const uint32_t lo =
+          __builtin_amdgcn_perm(Vr[1][d >> 1], Vr[0][d >> 1], sel);
+      const uint32_t hi =
+          __builtin_amdgcn_perm(Vr[3][d >> 1], Vr[2][d >> 1], sel);
+      const int row = vt_d + d;
+      *(u2v*)&sV[row * VST2 + 4 * (vt_g ^ ((row >> 3) & 7))] = (u2v){lo, hi};
+    }
+  };
+
+  const int ntiles = (k_end + KT1 - 1) / KT1;
+  fetch(0);
+  for (int i = 0; i < ntiles; ++i) {
+    const int n0 = i * KT1;
+    store();
+    __syncthreads();
+    if (i + 1 < ntiles) fetch(n0 + KT1);
+    if (wave_live && n0 <= wave_last_tok) {
+      // S^T for both 16-token halves, one softmax step, then PV.
+      f4v sacc[2];
+    #pragma unroll
+      for (int j = 0; j < 2; ++j) {
+        sacc[j] = (f4v){0, 0, 0, 0};
+        const _Float16* krow = &sK[(16 * j + col) * KROW2 + 4 * gq];
+    #pragma unroll
+        for (int s = 0; s < FR; ++s)
+          sacc[j] = __builtin_amdgcn_mfma_f32_16x16x16f16(
+              *(const h4v*)(krow + 16 * s), Qh[s], sacc[j], 0, 0, 0);
+      }
+      float sc[2][4];
+      float mloc = -INFINITY;
+    #pragma unroll
+      for (int j = 0; j < 2; ++j)
+    #pragma unroll
+        for (int e = 0; e < 4; ++e) {
+          const int kabs = n0 + 16 * j + 4 * gq + e;
+          sc[j][e] = qlive && kabs <= qtok ? sacc[j][e] * c2 : -INFINITY;
+          mloc = fmaxf(mloc, sc[j][e]);
+        }
+      const float mn = fmaxf(m, qmax4(mloc));
+      const float al = mn == -INFINITY ? 1.f : __builtin_amdgcn_exp2f(m - mn);
+      h4v pb[2];
+      float lsum = 0.f;
+    #pragma unroll
+      for (int j = 0; j < 2; ++j)
+    #pragma unroll
+        for (int e = 0; e < 4; ++e) {
+          const float p = sc[j][e] == -INFINITY
+                              ? 0.f
+                              : __builtin_amdgcn_exp2f(sc[j][e] - mn);
+          lsum += p;
+          pb[j][e] = (_Float16)p;
+        }
+      m = mn;
+      l = l * al + lsum;
+      if (__ballot(al != 1.0f)) {
+    #pragma unroll
+        for (int c = 0; c < FR; ++c) acc[c] *= al;
+      }
+    #pragma unroll
+      for (int j = 0; j < 2; ++j) {
+          // Row 16 c + col; its group 4 j + gq sits at (4 j + gq) ^ (row >> 3).
+    #pragma unroll
+        for (int c = 0; c < FR; ++c) {
+          const int row = 16 * c + col;
+          const _Float16* vp =
+              &sV[row * VST2 + 4 * ((4 * j + gq) ^ ((row >> 3) & 7))];
+          acc[c] = __builtin_amdgcn_mfma_f32_16x16x16f16(*(const h4v*)vp, pb[j],
+                                                         acc[c], 0, 0, 0);
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  l = qsum4(l);
+  if (!qlive) return;
+  const float inv = 1.0f / (l + 1e-10f);
+  __half* orow = out + (int64_t)(q0 + qtok) * so_tok + (int64_t)head * so_head;
+    #pragma unroll
+  for (int c = 0; c < FR; ++c)
+    #pragma unroll
+    for (int e = 0; e < 4; ++e)
+      orow[16 * c + 4 * gq + e] = __float2half(acc[c][e] * inv);
+  #endif  // OCTAVE_CDNA_STUB
+}
+
+}  // namespace octave_cdna
+
+// Called by octave_prefill (octave_prefill_rdna3.cu) on gfx942; same
+// arguments, already validated there.
+void octave_prefill_cdna(torch::Tensor out, torch::Tensor q, torch::Tensor k,
+                         torch::Tensor v, torch::Tensor cache,
+                         torch::Tensor block_table, torch::Tensor cu_seqlens_q,
+                         torch::Tensor seq_lens, int64_t max_query_len,
+                         double sm_scale, int64_t fmt, float* ws,
+                         int num_splits) {
+  using namespace octave_cdna;
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int num_seqs = seq_lens.size(0);
+  const int hq = q.size(1), hkv = cache.size(1);
+  const int hpk = hq / hkv;
+  const int block_size = cache.size(2);
+  const int max_blocks = block_table.size(1);
+  const int total_q = q.size(0);
+  const int rowtiles = (int)((max_query_len * hpk + QB1 - 1) / QB1);
+
+  dim3 grid1(num_seqs, hkv, rowtiles * num_splits);
+  #define PREFIX(B, KC)                                                     \
+    prefix_attn<B, KC><<<grid1, NT1, 0, stream>>>(                          \
+        (const __half*)q.data_ptr(), cache.data_ptr<uint8_t>(),             \
+        block_table.data_ptr<int>(), cu_seqlens_q.data_ptr<int>(),          \
+        seq_lens.data_ptr<int>(), ws, (float)sm_scale, hq, hkv, block_size, \
+        max_blocks, num_splits, total_q, q.stride(0), q.stride(1),          \
+        cache.stride(0), cache.stride(1), cache.stride(2))
+  if (fmt == 4)
+    PREFIX(4, false);
+  else if (fmt == 3)
+    PREFIX(3, false);
+  else
+    PREFIX(3, true);
+  #undef PREFIX
+  dim3 grid2(num_seqs, hkv, rowtiles);
+  chunk_attn<<<grid2, NT1, 0, stream>>>(
+      (__half*)out.data_ptr(), ws, (const __half*)q.data_ptr(),
+      (const __half*)k.data_ptr(), (const __half*)v.data_ptr(),
+      cu_seqlens_q.data_ptr<int>(), seq_lens.data_ptr<int>(), hq, hkv,
+      max_blocks * block_size, (float)sm_scale, num_splits, total_q,
+      q.stride(0), q.stride(1), k.stride(0), k.stride(1), v.stride(0),
+      v.stride(1), out.stride(0), out.stride(1));
+}
+
+#endif  // USE_ROCM

csrc/attention/octave_prefill_rdna3.cu

@@ -0,0 +1,809 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Octave prefill for RDNA3 (gfx11), head size 256, fp16. Reads the packed
+// cache directly; no intermediate buffers. Everything is in the rotated space
+// (the caller rotates Q/K/V of the chunk and rotates the output back).
+//
+// Phase 1 (prefix_attn): attention of the chunk's queries over the cached
+// prefix, split over the prefix length; writes (O, m, l) partials.
+//   One block = 128 query vectors (token, head) of one KV head, so the GQA
+//   heads share every K/V tile. 16 compute waves in pairs: wave g of a pair
+//   holds Q (int8) and O for head dims [128 g, 128 g + 128). 8 loader waves
+//   unpack 16-token tiles into double-buffered LDS, one tile ahead.
+//   QK  S^T[token][q] = K x Q^T with int8 WMMA. K rows hold the int8 codebook
+//       value of every dim. Each wave's 128 dims are two 64-dim blocks with
+//       their own per-token scale; the wave scales both and the pair
+//       exchanges the float partial through LDS.
+//   PV  O^T[dim][q] = V^T x P^T with V as its exact fp16 codebook values; the
+//       V scale is folded into P.
+// Phase 2 (chunk_attn): causal attention over the chunk's own K/V (fp16),
+// merged with the phase-1 partials; writes the output.
+
+#include <cstdint>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+#include <algorithm>
+
+#include <torch/all.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#if defined(USE_ROCM)
+  #include <hip/hip_runtime.h>
+  #include <hip/hip_fp16.h>
+
+  #include "octave_format.cuh"
+
+  #if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx1100__)
+    #define OCTAVE_PREFILL_STUB
+  #endif
+
+namespace octave_pf {
+
+using octave::D;
+using octave::Format;
+using octave::LUT_ONE;
+
+typedef _Float16 h16v __attribute__((ext_vector_type(16)));
+typedef _Float16 h8v __attribute__((ext_vector_type(8)));
+typedef _Float16 h2v __attribute__((ext_vector_type(2)));
+typedef float f8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u4v __attribute__((ext_vector_type(4)));
+typedef uint32_t u2v __attribute__((ext_vector_type(2)));
+typedef int v4i_t __attribute__((ext_vector_type(4)));
+typedef int v8i_t __attribute__((ext_vector_type(8)));
+
+constexpr int KT = 16;  // tokens per tile
+
+// Splits of the cached prefix used by a sequence: one per kMinSplitTokens, at
+// most num_splits. Both phases must agree.
+constexpr int kMinSplitTokens = 2048;
+__device__ __forceinline__ int used_splits(int ctx, int num_splits) {
+  if (ctx <= 0) return 0;
+  return min(num_splits, (ctx + kMinSplitTokens - 1) / kMinSplitTokens);
+}
+__device__ __forceinline__ int split_len(int ctx, int used) {
+  return ((ctx + used - 1) / used + KT - 1) / KT * KT;
+}
+
+// ===========================================================================
+// Phase 1: prefix
+// ===========================================================================
+constexpr int NW = 8;  // compute wave pairs (16 query vectors each)
+constexpr int NL = 8;  // loader waves
+[[maybe_unused]] constexpr int KROW = D + 16;  // bytes per K row in LDS
+[[maybe_unused]] constexpr float PSCALE = 4096.0f;
+// The LDS read of PV step c waits for the WMMA of step c - DEPV, which keeps
+// the scheduler from hoisting every read (and spilling).
+[[maybe_unused]] constexpr int DEPV = 4;
+constexpr int NTHREADS = (2 * NW + NL) * 32;
+constexpr int VPAD = 16;  // halves after every 8 rows of V^T
+constexpr int HD = D / 2;
+[[maybe_unused]] constexpr int QSTEPS = HD / 16;
+[[maybe_unused]] constexpr int OFRAGS = HD / 16;
+
+__device__ __forceinline__ int vrow(int row) {
+  return row * KT + (row >> 3) * VPAD;
+}
+__device__ __forceinline__ uint32_t xhalf(uint32_t v) {
+  return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xFEDCBA98u, false,
+                                      false);
+}
+__device__ __forceinline__ float xhalf(float v) {
+  return __uint_as_float(xhalf(__float_as_uint(v)));
+}
+// 4 dwords, one per token (same 4 dims) -> 4 dwords, one per dim (4 tokens).
+__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2,
+                                    uint32_t a3, uint32_t o[4]) {
+  const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u);
+  const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u);
+  const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u);
+  const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u);
+  o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u);
+  o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u);
+  o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u);
+  o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u);
+}
+// LDS-only barrier: __syncthreads() would also wait for the loaders' global
+// prefetch.
+__device__ __forceinline__ void tile_barrier() {
+  asm volatile("s_waitcnt lgkmcnt(0)\n\ts_barrier" ::: "memory");
+}
+
+// Codes of dims 4k..4k+3 of a 64-dim block, one byte each. 4-bit: raw is
+// the block's 32 bytes. 3-bit: raw starts at the block's 16 bytes of the
+// 2-bit plane; hi0/hi1 are its two dwords of the 1-bit plane.
+template <int BITS>
+__device__ __forceinline__ uint32_t block_codes(const uint32_t* raw,
+                                                uint32_t hi0, uint32_t hi1,
+                                                int k) {
+  if constexpr (BITS == 4) {
+    return (raw[k >> 1] >> (4 * (k & 1))) & 0x0F0F0F0Fu;
+  } else {
+    const uint32_t hi = k < 8 ? hi0 : hi1;
+    return ((raw[k >> 2] >> (2 * (k & 3))) & 0x03030303u) |
+           (((hi >> (k & 7)) & 0x01010101u) << 2);
+  }
+}
+
+template <int VB, bool KC>
+__global__ __launch_bounds__(NTHREADS) void prefix_attn(
+    const __half* __restrict__ q, const uint8_t* __restrict__ cache,
+    const int* __restrict__ block_table, const int* __restrict__ cu_seqlens_q,
+    const int* __restrict__ seq_lens, float* __restrict__ ws, float sm_scale,
+    int num_q_heads, int num_kv_heads, int block_size, int max_blocks,
+    int num_splits, int total_q_tokens, int64_t sq_tok, int64_t sq_head,
+    int64_t scb, int64_t sch, int64_t sct) {
+  #ifndef OCTAVE_PREFILL_STUB
+  using F = Format<VB, KC>;
+  const int seq = blockIdx.x, kvh = blockIdx.y;
+  const int rowtile = blockIdx.z / num_splits, split = blockIdx.z % num_splits;
+  const int tid = threadIdx.x;
+  const int w = __builtin_amdgcn_readfirstlane(tid >> 5), lane = tid & 31;
+  const int j = lane & 15, h = lane >> 4;
+  const int pr = w >> 1, g = w & 1;
+  const int hpk = num_q_heads / num_kv_heads;
+  const bool loader = w >= 2 * NW;
+
+  const int q0 = cu_seqlens_q[seq];
+  const int qlen = cu_seqlens_q[seq + 1] - q0;
+  int ctx = seq_lens[seq] - qlen;
+  ctx = max(0, min(ctx, max_blocks * block_size));
+  const int nqv = qlen * hpk;
+  if (rowtile * NW * 16 >= nqv) return;
+
+  const int used = used_splits(ctx, num_splits);
+  if (split >= used) return;
+  const int k_len = split_len(ctx, used);
+  const int k_begin = min(ctx, split * k_len);
+  const int k_stop = min(ctx, k_begin + k_len);
+  const int ntiles = k_stop > k_begin ? (k_stop - k_begin + KT - 1) / KT : 0;
+
+  __shared__ __attribute__((aligned(16))) int8_t sK[2][KT][KROW];
+  __shared__
+      __attribute__((aligned(16))) _Float16 sV[2][D * KT + (D / 8) * VPAD];
+  __shared__ __attribute__((aligned(16))) u4v sS[2][2 * NW][2][32];
+  // K scales travel with K (block 0 = RoPE, 1-3 = NoPE), V scales with V.
+  __shared__ float sKs[2][4][KT], sVs[2][KT];
+
+  // ---- loader waves
+  if (loader) {
+    // 256 threads. K: thread lt unpacks dims 16 kp .. 16 kp + 15 of block kq
+    // of token tk. V: dims 4 vh .. 4 vh + 3 of the 8-dim group vd of token
+    // quad vq.
+    const int lt = tid - 2 * NW * 32;
+    const int tk = lt >> 4, kq = (lt >> 2) & 3, kp = lt & 3;
+    const int vu = lt >> 1, vh = lt & 1, vq = vu & 3, vd = (vu >> 2) * 8;
+    // Block sizes are multiples of KT, so a tile sits in one cache block.
+    // Tiles past the end read the last one; tokens past k_stop read the last
+    // token: their K is masked to -inf and their V meets p = 0.
+    const int tlast = k_stop - 1;
+    const uint8_t* const head = cache + (int64_t)kvh * sch;
+    const int* const bt = block_table + seq * max_blocks;
+    // Walk the tiles in order: block index and offset inside it, advanced by
+    // KT per tile and frozen on the last tile.
+    struct Pos {
+      int start, blk, off;
+    };
+    auto first = [&]() {
+      return Pos{k_begin, k_begin / block_size, k_begin % block_size};
+    };
+    auto next = [&](Pos p) {
+      if (p.start + KT > tlast) return p;
+      p.start += KT;
+      p.off += KT;
+      if (p.off == block_size) {
+        p.off = 0;
+        ++p.blk;
+      }
+      return p;
+    };
+    auto pbase = [&](Pos p, int pb) {
+      return head + (int64_t)(pb < 0 ? 0 : pb) * scb + (int64_t)p.off * sct;
+    };
+    auto tok = [&](int start, int r) {
+      return (int64_t)min(r, tlast - start) * sct;
+    };
+    // 4-bit: dwords 2 kp, 2 kp + 1 of the RoPE codes. 3-bit: dword kp of the
+    // block's 2-bit plane and dword kp / 2 of its 1-bit plane.
+    const bool k3 = KC || kq > 0;
+    const int b = KC ? kq : kq - 1;  // 3-bit block index
+    const int koff0 = k3 ? F::OFF_KN + 16 * b + 4 * kp : 8 * kp;
+    const int koff1 =
+        k3 ? F::OFF_KN + F::KN / 4 + 8 * b + 4 * (kp >> 1) : 8 * kp + 4;
+    constexpr int VW = VB == 4 ? 1 : 2;  // dwords of 8 V dims
+    struct Stg {
+      uint32_t k0, k1;
+      uint32_t v[4][VW];
+      uint32_t s0, s1;
+    };
+    auto fetch = [&](int TK, const uint8_t* bK, int TV, const uint8_t* bV,
+                     Stg& S) {
+      const uint8_t* pk = bK + tok(TK, tk);
+      S.k0 = *(const uint32_t*)(pk + koff0);
+      S.k1 = *(const uint32_t*)(pk + koff1);
+    #pragma unroll
+      for (int i = 0; i < 4; ++i) {
+        const uint8_t* pv = bV + tok(TV, 4 * vq + i) + F::OFF_V;
+        if constexpr (VB == 4) {
+          S.v[i][0] = *(const uint32_t*)(pv + vd / 2);
+        } else {
+          S.v[i][0] = *(const uint32_t*)(pv + 4 * (vd >> 4));
+          S.v[i][1] = *(const uint32_t*)(pv + D / 4 + 4 * (vd >> 5));
+        }
+      }
+      // Threads 0-15: the K scales of token lt; 16-31: the V scale of token
+      // lt - KT. Block K: two 8-byte loads inside the slot, V scale in the
+      // low half of s0. Compact K: one 4-byte load {k, v}.
+      if (lt < 2 * KT) {
+        if constexpr (KC) {
+          S.s0 = *(const uint32_t*)((lt < KT ? bK + tok(TK, lt)
+                                             : bV + tok(TV, lt - KT)) +
+                                    F::OFF_SC);
+        } else {
+          const uint8_t* ps = lt < KT ? bK + tok(TK, lt) + F::OFF_SC
+                                      : bV + tok(TV, lt - KT) + F::OFF_SC + 8;
+          const u2v x = *(const u2v*)ps;
+          S.s0 = x[0];
+          S.s1 = x[1];
+        }
+      }
+    };
+    auto store_k = [&](const Stg& S, int bf) {
+      const uint32_t r4[2] = {S.k0, S.k1};
+      const uint32_t r3[1] = {S.k0};
+      const uint32_t hi = S.k1 >> (4 * (kp & 1));
+      u4v o;
+    #pragma unroll
+      for (int x = 0; x < 4; ++x)
+        o[x] = k3 ? octave::lut3(block_codes<3>(r3, hi, hi, x))
+                  : octave::lut4(block_codes<4>(r4, 0, 0, x));
+      *(u4v*)&sK[bf][tk][64 * kq + 16 * kp] = o;
+      if (lt < KT) {
+        const h2v a = __builtin_bit_cast(h2v, S.s0);
+        if constexpr (KC) {
+          sKs[bf][0][lt] = (float)a[0];
+        } else {
+          const h2v c = __builtin_bit_cast(h2v, S.s1);
+          sKs[bf][0][lt] = (float)a[0];
+          sKs[bf][1][lt] = (float)a[1];
+          sKs[bf][2][lt] = (float)c[0];
+          sKs[bf][3][lt] = (float)c[1];
+        }
+      }
+    };
+    auto store_v = [&](const Stg& S, int bf) {
+      uint32_t c[4];  // [token] code bytes of dims vd + 4 vh ..
+    #pragma unroll
+      for (int i = 0; i < 4; ++i) {
+        if constexpr (VB == 4) {
+          c[i] = (S.v[i][0] >> (4 * vh)) & 0x0F0F0F0Fu;
+        } else {
+          const int sl = 2 * ((vd & 15) >> 2) + 2 * vh;
+          const int sh = ((vd & 31) >> 2) + vh;
+          c[i] = ((S.v[i][0] >> sl) & 0x03030303u) |
+                 (((S.v[i][1] >> sh) & 0x01010101u) << 2);
+        }
+        c[i] = octave::lut<VB>(c[i]) ^ 0x80808080u;
+      }
+      uint32_t t[4];
+      tr4(c[0], c[1], c[2], c[3], t);
+      // Codebook bytes fed as fp16 1024 + (x + 128), less 1152: exact x.
+      const h2v kofs = {(_Float16)1152.0f, (_Float16)1152.0f};
+    #pragma unroll
+      for (int d = 0; d < 4; ++d) {
+        const int row = vd + 4 * vh + d, sw = ((row >> 3) & 1) * 8;
+        const h2v p01 = __builtin_bit_cast(
+            h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07010700u));
+        const h2v p23 = __builtin_bit_cast(
+            h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07030702u));
+        *(u2v*)&sV[bf][vrow(row) + ((4 * vq) ^ sw)] =
+            (u2v){__builtin_bit_cast(uint32_t, p01 - kofs),
+                  __builtin_bit_cast(uint32_t, p23 - kofs)};
+      }
+      if (lt >= KT && lt < 2 * KT)
+        sVs[bf][lt - KT] = (float)__builtin_bit_cast(
+                               _Float16, (uint16_t)(KC ? S.s0 >> 16 : S.s0)) *
+                           (PSCALE / LUT_ONE);
+    };
+    // Pipeline: K of tile i + 1 and V of tile i are written in iteration i
+    // from registers; the block number of a tile is loaded two tiles ahead.
+    Stg S;
+    if (ntiles > 0) {
+      const Pos p0 = first(), p1 = next(p0), p2 = next(p1);
+      const uint8_t* b0 = pbase(p0, bt[p0.blk]);
+      const int q1 = bt[p1.blk];
+      fetch(p0.start, b0, p0.start, b0, S);
+      const int q2 = bt[p2.blk];
+      Pos p3 = next(p2);
+      int pn = bt[p3.blk];  // block of tile i + 3
+      const uint8_t* bv = pbase(p1, q1);
+      const uint8_t* bk = pbase(p2, q2);
+      int sv = p1.start, sk = p2.start;
+      store_k(S, 0);
+      fetch(p1.start, bv, p0.start, b0, S);
+      tile_barrier();
+      for (int i = 0; i < ntiles; ++i) {
+        store_k(S, (i + 1) & 1);
+        store_v(S, i & 1);
+        fetch(sk, bk, sv, bv, S);
+        bv = bk;
+        sv = sk;
+        bk = pbase(p3, pn);
+        sk = p3.start;
+        p3 = next(p3);
+        pn = bt[p3.blk];
+        tile_barrier();
+      }
+    } else {
+      tile_barrier();
+    }
+    tile_barrier();
+    return;
+  }
+
+  // ---- compute waves: this lane's query row, int8 over this wave's dims
+  const int qr = pr * 16 + j;
+  const int qv = rowtile * NW * 16 + qr;
+  const bool qlive = qv < nqv;
+  u4v Q[QSTEPS];
+  float qsc;
+  {
+    const int qtok = qlive ? qv / hpk : 0;
+    const int qhead = kvh * hpk + (qlive ? qv % hpk : 0);
+    const __half* qrow =
+        q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)qhead * sq_head;
+    float qmax = 1e-8f;
+    for (int i = 0; i < D / 8; ++i) {
+      const h8v x = qlive ? *(const h8v*)(qrow + 8 * i) : (h8v){};
+    #pragma unroll
+      for (int k = 0; k < 8; ++k) qmax = fmaxf(qmax, fabsf((float)x[k]));
+    }
+    const float qinv = 127.0f / qmax;
+    qsc = qmax * (1.0f / 127.0f);
+    #pragma unroll
+    for (int f = 0; f < QSTEPS; ++f) {
+      uint32_t wd[4];
+    #pragma unroll
+      for (int hh = 0; hh < 2; ++hh) {
+        const h8v x =
+            qlive ? *(const h8v*)(qrow + HD * g + 16 * f + 8 * hh) : (h8v){};
+    #pragma unroll
+        for (int d4 = 0; d4 < 2; ++d4) {
+          uint32_t word = 0;
+    #pragma unroll
+          for (int b = 0; b < 4; ++b) {
+            int v = (int)lrintf((float)x[4 * d4 + b] * qinv);
+            v = max(-127, min(127, v));
+            word |= ((uint32_t)(uint8_t)(int8_t)v) << (8 * b);
+          }
+          wd[2 * hh + d4] = word;
+        }
+      }
+      Q[f] = (u4v){wd[0], wd[1], wd[2], wd[3]};
+    }
+  }
+
+  f8v O[OFRAGS];
+    #pragma unroll
+  for (int c = 0; c < OFRAGS; ++c) O[c] = (f8v){0, 0, 0, 0, 0, 0, 0, 0};
+  // A operands are only read in lanes with (j % 2) == h; the other lane of
+  // each pair reads its partner's row, so the pair costs one LDS fetch.
+  const int rr = ((j & 1) == h) ? j : (j ^ 1);
+  const int sw = ((rr >> 3) & 1) * 8;
+  const float c1 = sm_scale * qsc * (1.4426950408889634f / LUT_ONE);
+  float m = -INFINITY, l = 0.f;
+  f8v Sp = {0, 0, 0, 0, 0, 0, 0, 0};  // this wave's partial S, tile i-1
+
+  tile_barrier();
+  for (int i = 0; i <= ntiles; ++i) {
+    f8v Sn = {0, 0, 0, 0, 0, 0, 0, 0};
+    if (i < ntiles) {
+      // QK over this wave's two 64-dim blocks, S^T = K x Q^T, 4 k-steps
+      // each; block s is scaled into Sn before block s + 1 starts, so one
+      // int accumulator is live at a time.
+      uint32_t koff = (uint32_t)(rr * KROW + HD * g);
+      if constexpr (KC) {  // one K scale: one accumulator for both blocks
+        v8i_t Si = {0, 0, 0, 0, 0, 0, 0, 0};
+    #pragma unroll
+        for (int f = 0; f < 8; ++f) {
+          const u4v a = *(const u4v*)(&sK[i & 1][0][0] + koff + 16 * f);
+          Si = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(
+              true, __builtin_bit_cast(v4i_t, a), true,
+              __builtin_bit_cast(v4i_t, Q[f]), Si, false);
+        }
+        const float* ks = &sKs[i & 1][0][0];
+    #pragma unroll
+        for (int e = 0; e < 8; ++e) Sn[e] = ks[2 * e + h] * (float)Si[e];
+      }
+    #pragma unroll
+      for (int sg = 0; sg < (KC ? 0 : 2); ++sg) {
+        v8i_t Si = {0, 0, 0, 0, 0, 0, 0, 0};
+    #pragma unroll
+        for (int f = 4 * sg; f < 4 * sg + 4; ++f) {
+          const u4v a = *(const u4v*)(&sK[i & 1][0][0] + koff + 16 * f);
+          Si = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(
+              true, __builtin_bit_cast(v4i_t, a), true,
+              __builtin_bit_cast(v4i_t, Q[f]), Si, false);
+        }
+        const float* ks = &sKs[i & 1][2 * g + sg][0];
+    #pragma unroll
+        for (int e = 0; e < 8; ++e) Sn[e] += ks[2 * e + h] * (float)Si[e];
+        if (sg == 0) asm volatile("" : "+v"(koff) : "v"(Sn[0]));
+      }
+      const u8v su = __builtin_bit_cast(u8v, Sn);
+      sS[i & 1][w][0][lane] = (u4v){su[0], su[1], su[2], su[3]};
+      sS[i & 1][w][1][lane] = (u4v){su[4], su[5], su[6], su[7]};
+    }
+    if (i > 0) {
+      const int bf = (i - 1) & 1, base = k_begin + (i - 1) * KT;
+      const u4v o0 = sS[bf][w ^ 1][0][lane], o1 = sS[bf][w ^ 1][1][lane];
+      const f8v S =
+          Sp + __builtin_bit_cast(f8v, (u8v){o0[0], o0[1], o0[2], o0[3], o1[0],
+                                             o1[1], o1[2], o1[3]});
+      // Scores in log2 units are S * c1 (c1 > 0). Only the last tile has
+      // keys past k_stop; exp2(-inf) = 0, and m = -inf only before the first
+      // tile, so nothing else needs a guard.
+      float s[8];
+    #pragma unroll
+      for (int e = 0; e < 8; ++e) s[e] = S[e];
+      if (base + KT > k_stop) {
+    #pragma unroll
+        for (int e = 0; e < 8; ++e)
+          if (base + 2 * e + h >= k_stop) s[e] = -INFINITY;
+      }
+      float mloc = fmaxf(fmaxf(fmaxf(s[0], s[1]), fmaxf(s[2], s[3])),
+                         fmaxf(fmaxf(s[4], s[5]), fmaxf(s[6], s[7])));
+      const float mn = fmaxf(m, fmaxf(mloc, xhalf(mloc)) * c1);
+      const float alpha = __builtin_amdgcn_exp2f(m - mn);
+      m = mn;
+      uint32_t pd[4];
+      float lsum = 0.f;
+    #pragma unroll
+      for (int e = 0; e < 8; e += 2) {
+        const float p0 = __builtin_amdgcn_exp2f(fmaf(s[e], c1, -mn));
+        const float p1 = __builtin_amdgcn_exp2f(fmaf(s[e + 1], c1, -mn));
+        lsum += p0 + p1;
+        const _Float16 h0 = (_Float16)(p0 * sVs[bf][2 * e + h]);
+        const _Float16 h1 = (_Float16)(p1 * sVs[bf][2 * e + 2 + h]);
+        pd[e >> 1] = __builtin_bit_cast(uint32_t, (h2v){h0, h1});
+      }
+      l = l * alpha + lsum;
+      if (__builtin_amdgcn_ballot_w32(alpha != 1.0f)) {
+    #pragma unroll
+        for (int c = 0; c < OFRAGS; ++c) O[c] *= alpha;
+      }
+      u8v pf;
+        // B operand in natural token order: lane half 0 holds the even tokens,
+        // half 1 the odd ones; interleave them per pair.
+    #pragma unroll
+      for (int k = 0; k < 4; ++k) {
+        const uint32_t o = xhalf(pd[k]);
+        const uint32_t ev = h ? o : pd[k], od = h ? pd[k] : o;
+        pf[2 * k] = __builtin_amdgcn_perm(od, ev, 0x05040100u);
+        pf[2 * k + 1] = __builtin_amdgcn_perm(od, ev, 0x07060302u);
+      }
+      const h16v pb16 = __builtin_bit_cast(h16v, pf);
+    #pragma unroll
+      for (int c = 0; c < OFRAGS; ++c) {
+        uint32_t voff = (uint32_t)vrow(HD * g + 16 * c + rr);
+        if (c >= DEPV) asm volatile("" : "+v"(voff) : "v"(O[c - DEPV][0]));
+        const _Float16* vr = &sV[bf][0] + voff;
+        const h8v lo = *(const h8v*)(vr + sw);
+        const h8v hi = *(const h8v*)(vr + (8 ^ sw));
+        const h16v a = __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7,
+                                               8, 9, 10, 11, 12, 13, 14, 15);
+        O[c] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, pb16, O[c]);
+      }
+    }
+    Sp = Sn;
+    tile_barrier();
+  }
+
+  // ---- partials for the phase-2 merge (natural-log max, unnormalized O)
+  l += xhalf(l);
+  if (!qlive) return;
+  const int qtok = qv / hpk, qhead = kvh * hpk + qv % hpk;
+  float* wp = ws + (int64_t)split * total_q_tokens * num_q_heads * (D + 2) +
+              ((int64_t)(q0 + qtok) * num_q_heads + qhead) * (D + 2) + HD * g;
+    #pragma unroll
+  for (int c = 0; c < OFRAGS; ++c)
+    #pragma unroll
+    for (int e = 0; e < 8; ++e)
+      wp[16 * c + 2 * e + h] = O[c][e] * (1.0f / PSCALE);
+  if (g == 0 && h == 0) {
+    wp[D] = m * 0.6931471805599453f;
+    wp[D + 1] = l;
+  }
+  #endif  // OCTAVE_PREFILL_STUB
+}
+
+// ===========================================================================
+// Phase 2: the chunk's own K/V (fp16, causal), merged with the prefix
+// partials. 8 waves, 16 query rows each, one query head per block.
+// ===========================================================================
+constexpr int CW = 8;        // waves
+constexpr int CM = 16 * CW;  // query rows per block
+[[maybe_unused]] constexpr int FRAGS = D / 16;
+[[maybe_unused]] constexpr int X = 8;  // fp16 per 16-byte chunk
+
+__device__ __forceinline__ float row16_max(float v) {
+  v = fmaxf(v, __shfl_xor(v, 1));
+  v = fmaxf(v, __shfl_xor(v, 2));
+  v = fmaxf(v, __shfl_xor(v, 4));
+  return fmaxf(v, __shfl_xor(v, 8));
+}
+__device__ __forceinline__ float row16_sum(float v) {
+  v += __shfl_xor(v, 1);
+  v += __shfl_xor(v, 2);
+  v += __shfl_xor(v, 4);
+  return v + __shfl_xor(v, 8);
+}
+
+__global__ void __launch_bounds__(D)
+    chunk_attn(__half* __restrict__ out, const float* __restrict__ ws,
+               const __half* __restrict__ q, const __half* __restrict__ k,
+               const __half* __restrict__ v,
+               const int* __restrict__ cu_seqlens_q,
+               const int* __restrict__ seq_lens, int num_q_heads,
+               int num_kv_heads, int max_ctx, float sm_scale, int num_splits,
+               int total_q_tokens, int64_t sq_tok, int64_t sq_head,
+               int64_t sk_tok, int64_t sk_head, int64_t sv_tok, int64_t sv_head,
+               int64_t so_tok, int64_t so_head) {
+  #ifndef OCTAVE_PREFILL_STUB
+  const int seq = blockIdx.x, head = blockIdx.y;
+  const int tid = threadIdx.x, wave = tid >> 5, lane = tid & 31;
+  const int lo = lane & 15, hi = lane >> 4;
+  const int q0 = cu_seqlens_q[seq];
+  const int qlen = cu_seqlens_q[seq + 1] - q0;
+  const int ctx = max(0, min(seq_lens[seq] - qlen, max_ctx));
+  const int tile0 = blockIdx.z * CM;
+  if (tile0 >= qlen) return;
+  const int wq0 = tile0 + wave * 16;
+  const int nq_wave = max(0, min(16, qlen - wq0));
+  const int kvh = head / (num_q_heads / num_kv_heads);
+
+  __shared__ __attribute__((aligned(16))) _Float16 sK[FRAGS * 2 * KT * X];
+  __shared__ __attribute__((aligned(16))) _Float16 sV[D * KT];
+  __shared__ __attribute__((aligned(16))) _Float16 sP[CW][16 * KT];
+
+  float m_st[8], l_st[8];
+  f8v acc[FRAGS];
+  // Merge the prefix splits (one weight per split; exp(0) = 1 with one).
+  const int64_t split_stride = (int64_t)total_q_tokens * num_q_heads * (D + 2);
+  const int used = used_splits(ctx, num_splits);
+    #pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    m_st[i] = -INFINITY;
+    l_st[i] = 0.f;
+  }
+    #pragma unroll
+  for (int f = 0; f < FRAGS; ++f) acc[f] = (f8v){0, 0, 0, 0, 0, 0, 0, 0};
+    #pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    const int qp = wq0 + 2 * i + hi;
+    if (qp >= qlen) continue;
+    const float* w0 = ws + ((int64_t)(q0 + qp) * num_q_heads + head) * (D + 2);
+    float mx = -INFINITY;
+    for (int s = 0; s < used; ++s) mx = fmaxf(mx, w0[s * split_stride + D]);
+    float lsum = 0.f;
+    for (int s = 0; s < used; ++s) {
+      const float* wp = w0 + s * split_stride;
+      if (wp[D] == -INFINITY) continue;
+      const float wgt = __expf(wp[D] - mx);
+      lsum += wp[D + 1] * wgt;
+    #pragma unroll
+      for (int f = 0; f < FRAGS; ++f) acc[f][i] += wp[16 * f + lo] * wgt;
+    }
+    m_st[i] = mx;
+    l_st[i] = lsum;
+  }
+
+  // Q rows are re-read from global (L1/L2) in the loop instead of being held
+  // in registers; out-of-range rows are masked anyway.
+  const __half* qrow = q + (int64_t)(q0 + min(wq0 + lo, qlen - 1)) * sq_tok +
+                       (int64_t)head * sq_head;
+  const int k_end = min(qlen, tile0 + min(CM, qlen - tile0));
+  _Float16* sPw = &sP[wave][0];
+  for (int n0 = 0; n0 < k_end; n0 += KT) {
+    {  // K tile: [16-dim chunk][token][8] fp16
+      const int tk = tid / (FRAGS), dc = (tid % FRAGS) * 2;
+      const bool ok = n0 + tk < qlen;
+      const __half* row =
+          k + (int64_t)(q0 + n0 + tk) * sk_tok + (int64_t)kvh * sk_head;
+    #pragma unroll
+      for (int d = 0; d < 2; ++d) {
+        int4 x = ok ? *(const int4*)(row + (dc + d) * X) : int4{0, 0, 0, 0};
+        *(int4*)&sK[(dc + d) * (KT * X) + tk * X] = x;
+      }
+    }
+    {  // V tile transposed: [dim][token]
+    #pragma unroll
+      for (int p = 0; p < 2; ++p) {
+        const int tk = tid / (D / KT),
+                  dbase = ((tid % (D / KT)) + p * (D / KT)) * 8;
+        const bool ok = n0 + tk < qlen;
+        int4 x = ok ? *(const int4*)(v + (int64_t)(q0 + n0 + tk) * sv_tok +
+                                     (int64_t)kvh * sv_head + dbase)
+                    : int4{0, 0, 0, 0};
+        _Float16 t8[8];
+        __builtin_memcpy(t8, &x, 16);
+    #pragma unroll
+        for (int e = 0; e < 8; ++e) sV[(dbase + e) * KT + tk] = t8[e];
+      }
+    }
+    __syncthreads();
+    if (nq_wave > 0 && n0 <= wq0 + nq_wave - 1) {
+      f8v s = {0, 0, 0, 0, 0, 0, 0, 0};
+    #pragma unroll 2
+      for (int f = 0; f < FRAGS; ++f) {
+        h16v b, a;
+        const int4 klo = *(const int4*)&sK[(2 * f) * (KT * X) + lo * X];
+        const int4 khi = *(const int4*)&sK[(2 * f + 1) * (KT * X) + lo * X];
+        __builtin_memcpy(&b, &klo, 16);
+        __builtin_memcpy(((char*)&b) + 16, &khi, 16);
+        __builtin_memcpy(&a, qrow + 16 * f, 32);
+        s = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, s);
+      }
+      const int kabs = n0 + lo;
+      float p[8];
+    #pragma unroll
+      for (int i = 0; i < 8; ++i) {
+        const int r = 2 * i + hi;
+        const bool keep = r < nq_wave && kabs < qlen && kabs <= wq0 + r;
+        const float sc = keep ? s[i] * sm_scale : -INFINITY;
+        const float mi = row16_max(sc);
+        const float mn = fmaxf(m_st[i], mi);
+        const float al = m_st[i] == -INFINITY ? 0.f : __expf(m_st[i] - mn);
+        p[i] = mn == -INFINITY ? 0.f : __expf(sc - mn);
+        l_st[i] = l_st[i] * al + row16_sum(p[i]);
+        m_st[i] = mn;
+    #pragma unroll
+        for (int f = 0; f < FRAGS; ++f) acc[f][i] *= al;
+        sPw[r * KT + lo] = (_Float16)p[i];
+      }
+      h16v pa;
+      const int4 plo = *(const int4*)&sPw[lo * KT];
+      const int4 phi = *(const int4*)&sPw[lo * KT + 8];
+      __builtin_memcpy(&pa, &plo, 16);
+      __builtin_memcpy(((char*)&pa) + 16, &phi, 16);
+    #pragma unroll
+      for (int f = 0; f < FRAGS; ++f) {
+        h16v vb;
+        const int4 vlo = *(const int4*)&sV[(16 * f + lo) * KT];
+        const int4 vhi = *(const int4*)&sV[(16 * f + lo) * KT + 8];
+        __builtin_memcpy(&vb, &vlo, 16);
+        __builtin_memcpy(((char*)&vb) + 16, &vhi, 16);
+        acc[f] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(pa, vb, acc[f]);
+      }
+    }
+    __syncthreads();
+  }
+
+    #pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    const int qp = wq0 + 2 * i + hi;
+    if (qp >= qlen) continue;
+    const float inv = 1.0f / (l_st[i] + 1e-10f);
+    __half* orow = out + (int64_t)(q0 + qp) * so_tok + (int64_t)head * so_head;
+    #pragma unroll
+    for (int f = 0; f < FRAGS; ++f)
+      orow[16 * f + lo] = __float2half(acc[f][i] * inv);
+  }
+  #endif  // OCTAVE_PREFILL_STUB
+}
+
+// Partials of phase 1, kept alive for the process: a graph that captured a
+// pointer must keep finding it valid. Grows geometrically, so it settles.
+float* workspace(int64_t need, const at::TensorOptions& opts) {
+  static std::mutex mtx;
+  static std::unordered_map<int, std::vector<at::Tensor>> kept;
+  std::lock_guard<std::mutex> lock(mtx);
+  std::vector<at::Tensor>& bufs = kept[opts.device().index()];
+  if (bufs.empty() || bufs.back().numel() < need) {
+    int64_t n = 8 << 20;
+    while (n < need) n *= 2;
+    bufs.push_back(at::empty({n}, opts));
+  }
+  return (float*)bufs.back().data_ptr();
+}
+
+}  // namespace octave_pf
+
+int octave_slot_bytes(int64_t fmt);
+void octave_prefill_cdna(torch::Tensor out, torch::Tensor q, torch::Tensor k,
+                         torch::Tensor v, torch::Tensor cache,
+                         torch::Tensor block_table, torch::Tensor cu_seqlens_q,
+                         torch::Tensor seq_lens, int64_t max_query_len,
+                         double sm_scale, int64_t fmt, float* ws,
+                         int num_splits);
+
+void octave_prefill(torch::Tensor out, torch::Tensor q, torch::Tensor k,
+                    torch::Tensor v, torch::Tensor cache,
+                    torch::Tensor block_table, torch::Tensor cu_seqlens_q,
+                    torch::Tensor seq_lens, int64_t max_query_len,
+                    double sm_scale, int64_t fmt) {
+  using namespace octave_pf;
+  TORCH_CHECK(q.dtype() == at::kHalf && k.dtype() == at::kHalf &&
+                  v.dtype() == at::kHalf && out.dtype() == at::kHalf,
+              "octave_prefill: fp16 only");
+  TORCH_CHECK(q.size(2) == D && q.stride(2) == 1 && k.stride(2) == 1 &&
+              v.stride(2) == 1 && out.stride(2) == 1);
+  TORCH_CHECK(cache.dtype() == at::kByte && cache.dim() == 4 &&
+              cache.stride(3) == 1);
+  TORCH_CHECK(cache.size(3) == octave_slot_bytes(fmt),
+              "octave_prefill: slot size mismatch");
+  const int align = fmt & octave::kCompactFlag ? 4 : 8;
+  TORCH_CHECK(cache.stride(2) % align == 0 && cache.stride(1) % align == 0 &&
+                  cache.stride(0) % align == 0,
+              "octave_prefill: misaligned slots");
+  // 0: gfx1100 (WMMA), 1: gfx942 (MFMA, octave_prefill_cdna.cu).
+  static const int arch = [] {
+    const std::string name =
+        at::cuda::getCurrentDeviceProperties()->gcnArchName;
+    if (name.rfind("gfx1100", 0) == 0) return 0;
+    if (name.rfind("gfx942", 0) == 0) return 1;
+    return -1;
+  }();
+  TORCH_CHECK(arch >= 0, "octave_prefill: gfx1100 or gfx942 only");
+  const int num_seqs = seq_lens.size(0);
+  if (num_seqs == 0 || q.size(0) == 0) return;
+  const at::cuda::OptionalCUDAGuard guard(device_of(q));
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int hq = q.size(1), hkv = cache.size(1);
+  TORCH_CHECK(hq % hkv == 0);
+  const int hpk = hq / hkv;
+  const int block_size = cache.size(2);
+  const int max_blocks = block_table.size(1);
+  TORCH_CHECK(block_table.stride(1) == 1 && block_table.stride(0) == max_blocks,
+              "octave_prefill: block_table rows must be contiguous");
+  const int total_q = q.size(0);
+
+  const int rowtiles = (int)((max_query_len * hpk + 127) / 128);
+  const int base_blocks = std::max(1, num_seqs * hkv * rowtiles);
+  // Blocks whose split has no keys exit at once, so ask for plenty.
+  const int num_splits =
+      std::min(16, std::max(1, (1024 + base_blocks - 1) / base_blocks));
+  float* ws = workspace((int64_t)num_splits * total_q * hq * (D + 2),
+                        q.options().dtype(at::kFloat));
+  if (arch == 1) {
+    octave_prefill_cdna(out, q, k, v, cache, block_table, cu_seqlens_q,
+                        seq_lens, max_query_len, sm_scale, fmt, ws, num_splits);
+    return;
+  }
+
+  dim3 grid1(num_seqs, hkv, rowtiles * num_splits);
+  #define PREFIX(B, KC)                                                     \
+    prefix_attn<B, KC><<<grid1, NTHREADS, 0, stream>>>(                     \
+        (const __half*)q.data_ptr(), cache.data_ptr<uint8_t>(),             \
+        block_table.data_ptr<int>(), cu_seqlens_q.data_ptr<int>(),          \
+        seq_lens.data_ptr<int>(), ws, (float)sm_scale, hq, hkv, block_size, \
+        max_blocks, num_splits, total_q, q.stride(0), q.stride(1),          \
+        cache.stride(0), cache.stride(1), cache.stride(2))
+  if (fmt == 4)
+    PREFIX(4, false);
+  else if (fmt == 3)
+    PREFIX(3, false);
+  else
+    PREFIX(3, true);
+  #undef PREFIX
+  dim3 grid2(num_seqs, hq, (max_query_len + CM - 1) / CM);
+  chunk_attn<<<grid2, D, 0, stream>>>(
+      (__half*)out.data_ptr(), ws, (const __half*)q.data_ptr(),
+      (const __half*)k.data_ptr(), (const __half*)v.data_ptr(),
+      cu_seqlens_q.data_ptr<int>(), seq_lens.data_ptr<int>(), hq, hkv,
+      max_blocks * block_size, (float)sm_scale, num_splits, total_q,
+      q.stride(0), q.stride(1), k.stride(0), k.stride(1), v.stride(0),
+      v.stride(1), out.stride(0), out.stride(1));
+}
+
+#else
+void octave_prefill(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+                    torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+                    int64_t, double, int64_t) {
+  TORCH_CHECK(false, "octave requires ROCm");
+}
+#endif

csrc/attention/octave_torch_bindings.cpp

@@ -0,0 +1,82 @@
+// Op registration for the Octave KV cache kernels (legacy _C extension, full
+// torch API). Built only for gfx1100 and gfx942; see CMakeLists.txt.
+#include <optional>
+
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include "core/registration.h"
+
+// fmt: V bits (3 or 4), plus 16 for compact K.
+void octave_cache_store(torch::Tensor key, torch::Tensor value,
+                        torch::Tensor cache, torch::Tensor slot_mapping,
+                        torch::Tensor k_signs, torch::Tensor v_signs,
+                        int64_t fmt);
+
+void octave_decode(torch::Tensor out, torch::Tensor query, torch::Tensor cache,
+                   torch::Tensor block_table, torch::Tensor q_to_req,
+                   torch::Tensor q_to_klen, torch::Tensor mid_o,
+                   torch::Tensor k_signs, torch::Tensor v_signs,
+                   double sm_scale, int64_t num_kv_splits, int64_t fmt,
+                   int64_t query_group, bool use_wmma);
+
+void octave_decode_sparse(torch::Tensor out, torch::Tensor query,
+                          torch::Tensor cache, torch::Tensor block_table,
+                          torch::Tensor q_to_req, torch::Tensor indices,
+                          torch::Tensor mid_o, torch::Tensor k_signs,
+                          torch::Tensor v_signs, double sm_scale,
+                          int64_t num_kv_splits, int64_t fmt,
+                          const std::optional<torch::Tensor>& positions,
+                          const std::optional<torch::Tensor>& wtab,
+                          const std::optional<torch::Tensor>& wtags,
+                          const std::optional<torch::Tensor>& stab,
+                          const std::optional<torch::Tensor>& stags);
+
+void octave_window_store(torch::Tensor key, torch::Tensor value,
+                         torch::Tensor slot_mapping, torch::Tensor positions,
+                         torch::Tensor wtab, torch::Tensor wtags,
+                         torch::Tensor wlocks, torch::Tensor stab,
+                         torch::Tensor stags, torch::Tensor slocks);
+
+void octave_rotate(torch::Tensor x, torch::Tensor signs, bool k_layout,
+                   bool inverse);
+
+void octave_prefill(torch::Tensor out, torch::Tensor q, torch::Tensor k,
+                    torch::Tensor v, torch::Tensor cache,
+                    torch::Tensor block_table, torch::Tensor cu_seqlens_q,
+                    torch::Tensor seq_lens, int64_t max_query_len,
+                    double sm_scale, int64_t fmt);
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def(
+      "octave_cache_store(Tensor key, Tensor value, Tensor! cache, "
+      "Tensor slot_mapping, Tensor k_signs, Tensor v_signs, int fmt) -> ()");
+  ops.impl("octave_cache_store", torch::kCUDA, &octave_cache_store);
+  ops.def(
+      "octave_decode(Tensor! out, Tensor query, Tensor cache, "
+      "Tensor block_table, Tensor q_to_req, Tensor q_to_klen, "
+      "Tensor! mid_o, Tensor k_signs, Tensor v_signs, float sm_scale, "
+      "int num_kv_splits, int fmt, int query_group, bool use_wmma) -> ()");
+  ops.impl("octave_decode", torch::kCUDA, &octave_decode);
+  ops.def(
+      "octave_decode_sparse(Tensor! out, Tensor query, Tensor cache, "
+      "Tensor block_table, Tensor q_to_req, Tensor indices, Tensor! mid_o, "
+      "Tensor k_signs, Tensor v_signs, float sm_scale, int num_kv_splits, "
+      "int fmt, Tensor? positions=None, Tensor? wtab=None, "
+      "Tensor? wtags=None, Tensor? stab=None, Tensor? stags=None) -> ()");
+  ops.impl("octave_decode_sparse", torch::kCUDA, &octave_decode_sparse);
+  ops.def(
+      "octave_window_store(Tensor key, Tensor value, Tensor slot_mapping, "
+      "Tensor positions, Tensor! wtab, Tensor! wtags, Tensor! wlocks, "
+      "Tensor! stab, Tensor! stags, Tensor! slocks) -> ()");
+  ops.impl("octave_window_store", torch::kCUDA, &octave_window_store);
+  ops.def(
+      "octave_rotate(Tensor! x, Tensor signs, bool k_layout, bool inverse) "
+      "-> ()");
+  ops.impl("octave_rotate", torch::kCUDA, &octave_rotate);
+  ops.def(
+      "octave_prefill(Tensor! out, Tensor q, Tensor k, Tensor v, Tensor cache, "
+      "Tensor block_table, Tensor cu_seqlens_q, Tensor seq_lens, "
+      "int max_query_len, float sm_scale, int fmt) -> ()");
+  ops.impl("octave_prefill", torch::kCUDA, &octave_prefill);
+}

docs/design/rocm_octave.md

@@ -0,0 +1,197 @@
+# Octave KV: a 3-bit KV cache for ROCm
+
+`ROCM_OCTAVE` is an attention backend for AMD GPUs that keeps the KV cache of
+full-attention layers at 3-4 bits per value, with native kernels (HIP today). It is
+selected with `--kv-cache-dtype octave_k3v3_compact`, `octave_k3v3` or
+`octave_k3v4` (the backend is picked from the cache dtype;
+`--attention-backend ROCM_OCTAVE` also works). The name comes from the 3-bit
+codes: each one indexes an eight-level codebook.
+
+## Why it exists
+
+Long-context serving on AMD GPUs is bounded by KV-cache memory long before it
+is bounded by compute. The existing options on ROCm were:
+
+- fp16/bf16 KV (`TRITON_ATTN`, `ROCM_ATTN`): 16 bits per value.
+- fp8 / int8 per-token-head: 8 bits per value.
+- `TURBOQUANT`: 3-4 bits per value, but its kernels are Triton (store and
+  decode) and its prefill goes through FlashAttention; it does not support the
+  fused multi-step MTP draft.
+
+Octave KV targets the same bit budget as TurboQuant with HIP kernels written for
+the AMD matrix instructions (WMMA on RDNA3, MFMA on CDNA as a contract slot),
+reading the packed cache directly in both decode and prefill, and with
+multi-token verification (MTP / speculative decoding) as a first-class case.
+
+## Supported models
+
+The format exploits *partial rotary embeddings*: models whose attention heads
+apply RoPE to a fraction of the head dims. Today the kernels are built for
+head size 256 with 64 rotary dims (`partial_rotary_factor = 0.25`), the
+attention shape of the Qwen3-Next / Qwen3.5 / Qwen3.6 / Qwen3.8 family (dense
+and MoE, e.g. Qwen3.8-27B, Qwen3.6-35B-A3B, Qwen3.5-122B-A10B). Like the MLA
+backends, which serve only models with multi-head latent attention, Octave KV is
+tied to an attention shape, not to a model.
+
+Models with full RoPE (Llama, Qwen3, Mistral, ...) or other head sizes are
+rejected at startup with an explicit error. Sliding window, ALiBi, attention
+sinks, soft-capping and encoder attention are not supported.
+
+The decode kernels cover GQA groups of 1, 2, 4, 6 or 8 query heads per KV
+head on each GPU. Under tensor parallelism the group is counted per GPU, so
+models with 32 query heads and 2 KV heads (Qwen3.5-122B-A10B) need TP ≥ 4.
+
+## Format
+
+One slot per (token, KV head). V is always rotated as one 256-dim block
+(random signs, then a Walsh-Hadamard transform) and gets 3 or 4 bits per dim.
+K has two layouts:
+
+- **Block K** (`octave_k3v3`, `octave_k3v4`): four 64-dim blocks, each rotated
+  on its own. The first block holds exactly the 64 rotary dims and gets 4 bits
+  per dim, the three NoPE blocks 3 bits; each block has its own scale.
+- **Compact K** (`octave_k3v3_compact`): rotated over all 256 dims like V,
+  3 bits per dim, one scale. The rotation spreads the rotary dims over the
+  whole head, so they need no extra bits; slots pad to 4 bytes instead of 8.
+
+| Bytes (block K) | Content |
+| --- | --- |
+| `[0, 32)` | K RoPE block, 4-bit codes |
+| `[32, 104)` | K NoPE blocks, 3-bit codes |
+| next `256·v/8` | V, `v`-bit codes |
+| next 10 | fp16 scales (K blocks 0-3, V) |
+
+| Bytes (compact K) | Content |
+| --- | --- |
+| `[0, 96)` | K, 3-bit codes |
+| `[96, 192)` | V, 3-bit codes |
+| `[192, 196)` | fp16 scales (K, V) |
+
+| Cache dtype | Slot bytes | Bits per value |
+| --- | --- | --- |
+| `octave_k3v3_compact` | 196 | 3.06 |
+| `octave_k3v3` | 216 | 3.38 |
+| `octave_k3v4` | 248 | 3.88 |
+
+For reference, `turboquant_3bit_nc` uses 198 bytes and `turboquant_k3v4_nc`
+230.
+
+After the rotation the coordinates are close to Gaussian, so each code
+indexes the Lloyd-Max codebook of N(0, 1) for its width, stored as int8
+(`LUT` in the reference): a code `c` of a block with scale `s` stands for
+`LUT[c] · s / 64`. The scale is chosen so that `x · x̂ = |x|²`. A
+least-squares fit minimizes the reconstruction error but shrinks every
+reconstructed vector (by about 3% at 3 bits), which flattens the softmax and
+scales down the attention output; at 380k tokens it cost 2-3× more NLL than
+this scale, at the same size. The quantizer is dynamic: no calibration, no per-model
+dictionaries. The kernels map four codes to their int8 values with one
+`v_perm_b32` (two for 4-bit codes), so the codebook costs nothing in
+bandwidth and the dot products stay integer. The byte layout's source of
+truth is the PyTorch reference in `vllm/v1/attention/ops/rocm_octave.py`.
+
+Why these layouts: measured on real Q/K/V of Qwen3.8-27B (32k tokens, every
+full-attention layer, attention-output error), both K layouts sit on the same
+bytes-versus-error curve, about 25-30% below TurboQuant at equal size; block
+K reaches the 4-bit rotary block, compact K the smallest slot. Rotating the
+rotary block alone at 3 bits does not work (one layer reaches 32% error);
+rotating it with the rest of the head does. Moving bits between K and V or
+weighting K channels by query energy did not help at a fixed size. The int8
+codebook loses nothing against the floating-point one.
+
+Attention runs in the rotated space: the query is rotated like K and the
+output is rotated back once per query. The rotations are orthogonal, so the
+scores are those of the unrotated vectors.
+
+## Kernels and data flow
+
+| Step | Kernel | File |
+| --- | --- | --- |
+| Cache write | `octave_cache_store`: one wave per slot, rotate + quantize + pack | `csrc/attention/octave_attn.cu` |
+| Decode / MTP verify (≤128 query tokens per request) | `octave_decode`: split-KV; WMMA tier on gfx11, portable dot4 tier elsewhere; consecutive query tokens of a request share each K/V read | `octave_decode_wmma_rdna3.cu`, `octave_attn.cu` |
+| Split reduce + inverse V rotation | `reduce_kernel` | `octave_attn.cu` |
+| Prefill | `octave_prefill`: phase 1 reads the cached prefix from the packed cache (int8 WMMA for QK with per-segment scales, fp16 WMMA for PV), phase 2 attends to the chunk's own unquantized K/V and merges | `octave_prefill_rdna3.cu` |
+| Q/K/V chunk rotation for prefill | `octave_rotate` | `octave_attn.cu` |
+| Codebooks (`v_perm_b32` tables) and byte layout | — | `octave_format.cuh` |
+
+The metadata builder builds, on the device and without host syncs, the
+per-query request and causal-length maps the kernels index by. They are what
+makes MTP verification batches (several query tokens per request, each with
+its own causal length) work without special cases, and they give cudagraph
+padding queries a K length of 0.
+
+## Architecture contract
+
+Kernels are native: HIP or FlyDSL, never Triton. There is no fallback: an
+architecture without a required kernel is rejected at startup with an error
+naming what is missing. The kernels in this tree are HIP; a FlyDSL kernel
+(for example a CDNA prefill) must keep the same signature and semantics, pass
+the same tests, and report a clear error when `flydsl` is not installed.
+
+Octave is enabled only on RDNA3 and CDNA3, validated on gfx1100 and gfx942
+(MI300). The backend rejects any other family at selection
+(`supports_compute_capability`, via `get_rdna_version()` and
+`get_cdna_version()`), and the kernel tests skip elsewhere. CMake builds the
+Octave sources and registers the ops only when the target list includes
+gfx1100 or gfx942; other targets in the same build compile the
+architecture-specific kernels as empty stubs, and the kernel dispatch reports
+any other RDNA3 chip as not validated.
+
+| Piece | gfx1100 (RDNA3) | gfx942 (MI300) |
+| --- | --- | --- |
+| int8 dot4 (`octave_arch.cuh`) | `v_dot4_i32_iu8` | `v_dot4_i32_i8` |
+| Codebook lookup (`octave_format.cuh`) | `v_perm_b32` | `v_perm_b32` |
+| Cache store, reduce, rotate | portable | portable (decode reduce: gfx942 kernel) |
+| Decode | WMMA | MFMA (`v_mfma_i32_16x16x32_i8`) |
+| Prefill | WMMA | MFMA |
+
+The portable dot4 decode stays as the reference tier the tests compare the
+native tiers against.
+
+To add an architecture:
+
+1. Implement `octave_prefill` for it, in HIP or FlyDSL (same signature and semantics as
+   `octave_prefill_rdna3.cu`: phase-1 partials over the packed prefix, phase-2
+   chunk attention and merge, rotated space in and out).
+2. Add a decode tier, or validate the portable dot4 kernel there.
+3. Add the architecture to the CMake guard, to `_kernel_tier` and
+   `supports_compute_capability` in
+   `vllm/v1/attention/backends/rocm_octave_attn.py`, and to the test gate.
+4. Run `tests/kernels/attention/test_rocm_octave.py`: every kernel is checked
+   against the PyTorch reference, and `test_decode_matches_reference` runs
+   both the fast and the portable decode tier.
+
+## Results
+
+All measurements on 4× RX 7900 XTX (gfx1100), Qwen3.8-27B W4A16, tensor
+parallel 4, MTP k=3, 8 GB of KV cache per GPU.
+
+### Capacity
+
+| KV cache | Tokens that fit |
+| --- | --- |
+| fp16 | TBD |
+| `turboquant_k3v4_nc` | 1,790,625 |
+| `octave_k4v4` | 1,402k |
+| `octave_k3v3` | 1,679k |
+
+### Speed
+
+TBD: `vllm bench serve` (random prompts, no prefix reuse) TTFT and TPOT at
+8k / 32k / 100k / 380k input, and per-layer kernel times from
+`benchmarks/attention_benchmarks`.
+
+### Quality
+
+TBD: GSM8K (`tests/evals/gsm8k/gsm8k_eval.py`, 1319 questions, 5-shot,
+temperature 0.6) and next-token NLL over a 380k-token document.
+
+## Reproducing
+
+```bash
+vllm serve Qwen/Qwen3.8-27B --kv-cache-dtype octave_k3v4 ...
+pytest tests/kernels/attention/test_rocm_octave.py
+python benchmarks/attention_benchmarks/benchmark.py \
+    --backends ROCM_OCTAVE --kv-cache-dtype octave_k3v4 \
+    --model Qwen/Qwen3.8-27B --head-dim 256 --num-q-heads 6 --num-kv-heads 1 \
+    --batch-specs q1s100k q4s100k q2ks100k
+```

tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3-compact.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3.5-4B"
+accuracy_threshold: 0.76
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype octave_k3v3_compact --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3.5-4B"
+accuracy_threshold: 0.77
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype octave_k3v3 --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v4.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3.5-4B"
+accuracy_threshold: 0.78
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype octave_k3v4 --max-model-len 4096"

tests/evals/gsm8k/configs/models-octave.txt

@@ -0,0 +1,3 @@
+Qwen3.5-4B-Octave-k3v4.yaml
+Qwen3.5-4B-Octave-k3v3.yaml
+Qwen3.5-4B-Octave-k3v3-compact.yaml

tests/kernels/attention/test_rocm_octave.py

@@ -0,0 +1,475 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Octave HIP kernels against the PyTorch reference in
+vllm/v1/attention/ops/rocm_octave.py."""
+
+import math
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import dense_kv_cache_views
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.rocm_octave_attn import RocmOctaveAttentionBackend
+from vllm.v1.attention.ops import rocm_octave as sq
+from vllm.v1.kv_cache_interface import FullAttentionSpec, get_kv_quant_mode
+
+
+def _arch() -> str | None:
+    if not current_platform.is_rocm():
+        return None
+    from vllm.platforms.rocm import get_cdna_version, get_rdna_version
+
+    if get_rdna_version() == 3:
+        return "rdna3"
+    return "cdna3" if get_cdna_version() == 3 else None
+
+
+pytestmark = pytest.mark.skipif(
+    _arch() is None or not hasattr(torch.ops._C, "octave_decode"),
+    reason="Octave kernels need RDNA3 or CDNA3",
+)
+
+HEAD, ROPE, BLOCK = 256, 64, 16
+CACHE_DTYPES = ["octave_k3v4", "octave_k3v3", "octave_k3v3_compact"]
+
+
+def _format(cache_dtype: str) -> sq.OctaveFormat:
+    return sq.OctaveFormat.from_cache_dtype(cache_dtype, HEAD, ROPE)
+
+
+def _rel(a: torch.Tensor, b: torch.Tensor) -> float:
+    return ((a.float() - b.float()).norm() / b.float().norm()).item()
+
+
+def _token_slots(cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor:
+    """(blocks, H, block_size, slot) cache -> (T, H, slot) for flat slots."""
+    flat = cache.permute(0, 2, 1, 3).reshape(-1, cache.shape[1], cache.shape[3])
+    return flat[slots]
+
+
+def _fill(cache, block_ids, k, v, k_signs, v_signs, fmt):
+    n = k.shape[0]
+    slots = (
+        block_ids.long().repeat_interleave(BLOCK)[:n] * BLOCK
+        + torch.arange(n, device=k.device) % BLOCK
+    )
+    torch.ops._C.octave_cache_store(
+        k, v, cache, slots, k_signs, v_signs, fmt.kernel_code
+    )
+    return slots
+
+
+@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES)
+def test_store_matches_reference(cache_dtype):
+    torch.manual_seed(0)
+    dev, fmt = "cuda", _format(cache_dtype)
+    ksg = vsg = sq.sign_bits(HEAD).to(dev)
+    t, hkv = 1024, 2
+    k = torch.randn(t, hkv, HEAD, device=dev)
+    k[..., 100] *= 8  # an outlier channel is spread by the rotation
+    k = k.bfloat16()
+    v = torch.randn(t, hkv, HEAD, device=dev).bfloat16()
+    nblk = t // BLOCK + 4
+    cache = torch.zeros(nblk, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev)
+    slots = _fill(cache, torch.randperm(nblk, device=dev), k, v, ksg, vsg, fmt)
+
+    k_got, v_got = sq.reference_dequantize(_token_slots(cache, slots), fmt)
+    ref = sq.reference_quantize(k.float(), v.float(), fmt)
+    k_ref, v_ref = sq.reference_dequantize(ref, fmt)
+    # Same codes up to rounding ties from a different summation order.
+    assert _rel(k_got, k_ref) < 1e-2 and _rel(v_got, v_ref) < 1e-2
+    # Lloyd-Max on Gaussian data: 3-bit codes (4-bit for block K's RoPE
+    # block) keep ~19% of the norm as error, 4-bit ~10%.
+    assert _rel(k_got, k) < 0.2
+    assert _rel(v_got, v) < (0.12 if fmt.v_bits == 4 else 0.22)
+
+
+@pytest.mark.parametrize(
+    "cache_dtype,slot_bytes",
+    [("octave_k3v4", 248), ("octave_k3v3", 216), ("octave_k3v3_compact", 196)],
+)
+def test_slot_bytes_match_kernel_layout(cache_dtype, slot_bytes):
+    """The cache is sized from the Python format; the kernels index it with
+    octave::Format in csrc/attention/octave_format.cuh."""
+    assert _format(cache_dtype).slot_bytes == slot_bytes
+
+
+@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES)
+def test_store_skips_padding_slots(cache_dtype):
+    """Tokens with slot -1 (cudagraph padding) leave every byte unchanged."""
+    torch.manual_seed(0)
+    dev, fmt = "cuda", _format(cache_dtype)
+    ksg = vsg = sq.sign_bits(HEAD).to(dev)
+    t, hkv = 8, 2
+    k = torch.randn(t, hkv, HEAD, device=dev).bfloat16()
+    v = torch.randn(t, hkv, HEAD, device=dev).bfloat16()
+    cache = torch.randint(
+        0, 256, (2, hkv, BLOCK, fmt.slot_bytes), dtype=torch.uint8, device=dev
+    )
+    before = cache.clone()
+    slots = torch.full((t,), -1, dtype=torch.long, device=dev)
+    slots[3] = 21
+    torch.ops._C.octave_cache_store(k, v, cache, slots, ksg, vsg, fmt.kernel_code)
+    written = torch.zeros_like(cache, dtype=torch.bool)
+    written[1, :, 5] = True
+    assert torch.equal(cache[~written], before[~written])
+    assert not torch.equal(cache[written], before[written])
+
+
+@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES)
+@pytest.mark.parametrize("query_group", [1, 4])
+@pytest.mark.parametrize("num_splits", [1, 7, 64])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_wmma", [True, False])
+def test_decode_matches_reference(
+    cache_dtype, query_group, num_splits, dtype, use_wmma
+):
+    """Covers MTP verification: up to 4 consecutive query tokens per request,
+    each with its own causal length, plus a padded query with no request.
+    ``use_wmma=False`` forces the portable kernel every architecture runs."""
+    torch.manual_seed(1)
+    dev, fmt = "cuda", _format(cache_dtype)
+    ksg = vsg = sq.sign_bits(HEAD).to(dev)
+    hkv, group = 1, 6
+    lens = [1, 16, 37, 700, 3000]
+    cache = torch.zeros(512, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev)
+    max_blocks = max(math.ceil(n / BLOCK) for n in lens)
+    block_table = torch.zeros(len(lens), max_blocks, dtype=torch.int32, device=dev)
+    free = torch.randperm(512, device=dev)
+    q_to_req, q_to_klen, used = [], [], 0
+    for i, n in enumerate(lens):
+        nb = math.ceil(n / BLOCK)
+        block_table[i, :nb] = free[used : used + nb].int()
+        used += nb
+        k = torch.randn(n, hkv, HEAD, device=dev).bfloat16()
+        v = torch.randn(n, hkv, HEAD, device=dev).bfloat16()
+        _fill(cache, block_table[i, :nb], k, v, ksg, vsg, fmt)
+        nq = min(query_group, n)
+        for j in range(nq):
+            q_to_req.append(i)
+            q_to_klen.append(n - nq + 1 + j)
+    q_to_req.append(len(lens) + 3)  # cudagraph padding: no request
+    q_to_klen.append(0)
+    q_to_req = torch.tensor(q_to_req, dtype=torch.int32, device=dev)
+    q_to_klen = torch.tensor(q_to_klen, dtype=torch.int32, device=dev)
+
+    num_q = q_to_req.numel()
+    q = (torch.randn(num_q, hkv * group, HEAD, device=dev) * 2).to(dtype)
+    out = torch.zeros_like(q)
+    mid = torch.empty(
+        num_q, hkv * group, num_splits, HEAD + 2, dtype=torch.float32, device=dev
+    )
+    torch.ops._C.octave_decode(
+        out,
+        q,
+        cache,
+        block_table,
+        q_to_req,
+        q_to_klen,
+        mid,
+        ksg,
+        vsg,
+        1 / 16,
+        num_splits,
+        fmt.kernel_code,
+        query_group,
+        use_wmma,
+    )
+    ref = sq.reference_attention(
+        q.float(),
+        cache,
+        block_table,
+        q_to_req.clamp(max=len(lens) - 1),
+        q_to_klen,
+        fmt,
+        1 / 16,
+    )
+    # Only the kernel's query precision (fp16 or int8) differs.
+    assert _rel(out[:-1], ref[:-1]) < 2e-2
+    assert out[-1].abs().max().item() == 0
+
+
+@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES)
+@pytest.mark.parametrize("query_lens", [[300], [129, 1, 517]])
+def test_prefill_matches_reference(cache_dtype, query_lens):
+    """Chunked prefill: each request's chunk attends causally to itself
+    (unquantized K/V) and to its cached prefix (packed cache), in the rotated
+    space the backend runs it in."""
+    torch.manual_seed(3)
+    dev, fmt = "cuda", _format(cache_dtype)
+    ksg = vsg = sq.sign_bits(HEAD).to(dev)
+    hkv, group = 1, 6
+    ctx_lens = [2500, 0, 4100][: len(query_lens)]
+    seq_lens = [c + n for c, n in zip(ctx_lens, query_lens)]
+    max_blocks = max(math.ceil(n / BLOCK) for n in seq_lens)
+    nblk = sum(math.ceil(n / BLOCK) for n in seq_lens) + 8
+    cache = torch.zeros(nblk, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev)
+    block_table = torch.zeros(len(seq_lens), max_blocks, dtype=torch.int32, device=dev)
+    free, used = torch.randperm(nblk, device=dev), 0
+    ks, vs = [], []
+    for i, n in enumerate(seq_lens):
+        nb = math.ceil(n / BLOCK)
+        block_table[i, :nb] = free[used : used + nb].int()
+        used += nb
+        k = torch.randn(n, hkv, HEAD, device=dev).half()
+        v = torch.randn(n, hkv, HEAD, device=dev).half()
+        _fill(cache, block_table[i, :nb], k, v, ksg, vsg, fmt)
+        ks.append(k)
+        vs.append(v)
+
+    qsl = torch.tensor([0, *query_lens], device=dev).cumsum(0).int()
+    total = int(qsl[-1])
+    q = (torch.randn(total, hkv * group, HEAD, device=dev) * 2).half()
+    k_new = torch.cat([k[c:] for k, c in zip(ks, ctx_lens)])
+    v_new = torch.cat([v[c:] for v, c in zip(vs, ctx_lens)])
+    q_rot, k_rot, v_rot = q.clone(), k_new.clone(), v_new.clone()
+    torch.ops._C.octave_rotate(q_rot, ksg, not fmt.compact, False)
+    torch.ops._C.octave_rotate(k_rot, ksg, not fmt.compact, False)
+    torch.ops._C.octave_rotate(v_rot, vsg, False, False)
+    out = torch.empty_like(q_rot)
+    torch.ops._C.octave_prefill(
+        out,
+        q_rot,
+        k_rot,
+        v_rot,
+        cache,
+        block_table,
+        qsl,
+        torch.tensor(seq_lens, dtype=torch.int32, device=dev),
+        max(query_lens),
+        1 / 16,
+        fmt.kernel_code,
+    )
+    torch.ops._C.octave_rotate(out, vsg, False, True)
+
+    ref = torch.empty(total, hkv * group, HEAD, device=dev)
+    for i, (c, n) in enumerate(zip(ctx_lens, query_lens)):
+        k_all, v_all = ks[i][c:].float(), vs[i][c:].float()
+        if c > 0:
+            pos = torch.arange(c, device=dev)
+            slots = cache[block_table[i, pos // BLOCK].long(), :, pos % BLOCK]
+            k_pre, v_pre = sq.reference_dequantize(slots, fmt)
+            k_all, v_all = torch.cat([k_pre, k_all]), torch.cat([v_pre, v_all])
+        k_all = k_all.repeat_interleave(group, 1)
+        v_all = v_all.repeat_interleave(group, 1)
+        qi = q[int(qsl[i]) : int(qsl[i + 1])].float()
+        s = torch.einsum("qhd,thd->hqt", qi, k_all) / 16
+        mask = (
+            torch.arange(c + n, device=dev)[None]
+            > (c + torch.arange(n, device=dev))[:, None]
+        )
+        p = torch.softmax(s.masked_fill(mask, float("-inf")), -1)
+        ref[int(qsl[i]) : int(qsl[i + 1])] = torch.einsum("hqt,thd->qhd", p, v_all)
+    # Only the kernel's int8 query and fp16 P/V rounding differ.
+    assert _rel(out, ref) < 2e-2
+
+
+def _decode(cache, q, block_table, q_to_req, q_to_klen, signs, fmt, use_wmma):
+    out = torch.zeros_like(q)
+    mid = torch.empty(*q.shape[:2], 8, HEAD + 2, dtype=torch.float32, device=q.device)
+    torch.ops._C.octave_decode(
+        out,
+        q,
+        cache,
+        block_table,
+        q_to_req,
+        q_to_klen,
+        mid,
+        signs,
+        signs,
+        1 / 16,
+        8,
+        fmt.kernel_code,
+        1,
+        use_wmma,
+    )
+    return out
+
+
+def _prefill(cache, q, k, v, block_table, qsl, seq_lens, signs, fmt):
+    out = torch.empty_like(q)
+    torch.ops._C.octave_prefill(
+        out,
+        q,
+        k,
+        v,
+        cache,
+        block_table,
+        qsl,
+        seq_lens,
+        int((qsl[1:] - qsl[:-1]).max()),
+        1 / 16,
+        fmt.kernel_code,
+    )
+    return out
+
+
+@pytest.mark.parametrize(
+    "layout", RocmOctaveAttentionBackend.supported_kv_cache_layouts()
+)
+def test_kernels_follow_declared_cache_layouts(layout):
+    """Store, decode and prefill on the cache view vLLM allocates for each
+    declared layout match a contiguous cache bit for bit. Two KV heads, so the
+    head and token strides differ between layouts."""
+    torch.manual_seed(4)
+    dev, cache_dtype = "cuda", "octave_k3v3"
+    fmt, hkv, group, nblk = _format(cache_dtype), 2, 6, 80
+    signs = sq.sign_bits(HEAD).to(dev)
+    sq.register_rope_dim(HEAD, ROPE)
+    spec = RocmOctaveAttentionBackend.customize_spec(
+        FullAttentionSpec(
+            block_size=BLOCK,
+            num_kv_heads=hkv,
+            head_size=HEAD,
+            dtype=torch.uint8,
+            kv_quant_mode=get_kv_quant_mode(cache_dtype),
+        )
+    )
+    raw = torch.zeros(nblk * spec.page_size_bytes, dtype=torch.int8, device=dev)
+    cache = dense_kv_cache_views(raw, spec, nblk, 1, layout)[0]
+    dense = torch.zeros(nblk, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev)
+    assert cache.shape == dense.shape and cache.dtype == dense.dtype
+
+    seq_lens, query_lens = [700, 300], [100, 300]
+    max_blocks = max(math.ceil(n / BLOCK) for n in seq_lens)
+    block_table = torch.zeros(len(seq_lens), max_blocks, dtype=torch.int32, device=dev)
+    free, used, ks, vs = torch.randperm(nblk, device=dev), 0, [], []
+    for i, n in enumerate(seq_lens):
+        nb = math.ceil(n / BLOCK)
+        block_table[i, :nb] = free[used : used + nb].int()
+        used += nb
+        ks.append(torch.randn(n, hkv, HEAD, device=dev).half())
+        vs.append(torch.randn(n, hkv, HEAD, device=dev).half())
+        for c in (cache, dense):
+            _fill(c, block_table[i, :nb], ks[i], vs[i], signs, signs, fmt)
+    assert torch.equal(cache, dense)
+
+    q_to_req = torch.arange(len(seq_lens), dtype=torch.int32, device=dev)
+    q_to_klen = torch.tensor(seq_lens, dtype=torch.int32, device=dev)
+    q = torch.randn(len(seq_lens), hkv * group, HEAD, device=dev).half()
+    for use_wmma in (True, False):
+        out, ref = (
+            _decode(c, q, block_table, q_to_req, q_to_klen, signs, fmt, use_wmma)
+            for c in (cache, dense)
+        )
+        assert out.abs().sum() > 0
+        assert torch.equal(out, ref)
+
+    qsl = torch.tensor([0, *query_lens], device=dev).cumsum(0).int()
+    q = torch.randn(int(qsl[-1]), hkv * group, HEAD, device=dev).half()
+    ctx = [n - m for n, m in zip(seq_lens, query_lens)]
+    k = torch.cat([x[c:] for x, c in zip(ks, ctx)])
+    v = torch.cat([x[c:] for x, c in zip(vs, ctx)])
+    out, ref = (
+        _prefill(c, q, k, v, block_table, qsl, q_to_klen, signs, fmt)
+        for c in (cache, dense)
+    )
+    assert torch.equal(out, ref)
+
+
+def _window_tables(entries, hkv, dev):
+    return (
+        torch.zeros(entries, hkv, 2 * HEAD, dtype=torch.float16, device=dev),
+        torch.full((entries,), -1, dtype=torch.int32, device=dev),
+        torch.zeros(entries, dtype=torch.int32, device=dev),
+    )
+
+
+@pytest.mark.skipif(_arch() != "rdna3", reason="sparse decode is RDNA3-only")
+@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES)
+@pytest.mark.parametrize("hkv,group", [(1, 6), (2, 6)])
+@pytest.mark.parametrize("window", ["off", "on", "evicted"])
+def test_sparse_decode_matches_reference(cache_dtype, hkv, group, window):
+    """Sparse (top-k) decode as qwen4_exp QSA runs it: each query reads the
+    logical positions in its index row, -1 entries are holes. With the window
+    on, the first 4 and the last 64 positions come from the fp16 tables;
+    an entry whose tag was overwritten falls back to the packed cache."""
+    torch.manual_seed(4)
+    dev, fmt = "cuda", _format(cache_dtype)
+    signs = sq.sign_bits(HEAD).to(dev)
+    lens, max_blocks, nblk = [1500, 2600, 777], 192, 1024
+    cache = torch.zeros(nblk, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev)
+    block_table = (
+        torch.randperm(nblk, device=dev)[: len(lens) * max_blocks]
+        .view(len(lens), max_blocks)
+        .int()
+    )
+    wtab, wtags, wlocks = _window_tables(4096, hkv, dev)
+    stab, stags, slocks = _window_tables(256, hkv, dev)
+    ks, vs = [], []
+    for r, n in enumerate(lens):
+        k = torch.randn(n, hkv, HEAD, device=dev).half()
+        v = torch.randn_like(k)
+        pos = torch.arange(n, device=dev)
+        slots = block_table[r, pos // BLOCK].long() * BLOCK + pos % BLOCK
+        torch.ops._C.octave_cache_store(
+            k, v, cache, slots, signs, signs, fmt.kernel_code
+        )
+        torch.ops._C.octave_window_store(
+            k, v, slots, pos.int(), wtab, wtags, wlocks, stab, stags, slocks
+        )
+        ks.append(k)
+        vs.append(v)
+    if window == "evicted":
+        wtags[::2] = -7
+        stags[::2] = -7
+
+    q_to_req = torch.tensor([0, 1, 1, 2, 0, 1, 2], dtype=torch.int32, device=dev)
+    num_q, topk = q_to_req.numel(), 512
+    qpos = torch.tensor(
+        [lens[r] - 1 - 37 * i for i, r in enumerate(q_to_req.tolist())],
+        dtype=torch.int32,
+        device=dev,
+    )
+    indices = torch.full((num_q, topk), -1, dtype=torch.int32, device=dev)
+    for i in range(num_q):
+        n = int(qpos[i]) + 1
+        cand = torch.cat(
+            [
+                torch.arange(max(0, n - 80), n, device=dev),
+                torch.arange(4, device=dev),
+                torch.randperm(n, device=dev)[:300],
+            ]
+        ).unique()
+        indices[i, : cand.numel()] = cand.int()
+        indices[i] = indices[i][torch.randperm(topk, device=dev)]
+
+    q = torch.randn(num_q, hkv * group, HEAD, device=dev).half()
+    out = torch.empty_like(q)
+    mid = torch.empty(num_q, hkv * group, 64, HEAD + 2, dtype=torch.float32, device=dev)
+    scale = HEAD**-0.5
+    window_args = () if window == "off" else (qpos, wtab, wtags, stab, stags)
+    torch.ops._C.octave_decode_sparse(
+        out,
+        q,
+        cache,
+        block_table,
+        q_to_req,
+        indices,
+        mid,
+        signs,
+        signs,
+        scale,
+        64,
+        fmt.kernel_code,
+        *window_args,
+    )
+
+    ref = torch.zeros(num_q, hkv * group, HEAD, device=dev)
+    for i in range(num_q):
+        r, p = int(q_to_req[i]), int(qpos[i])
+        sel = indices[i][indices[i] >= 0].long()
+        blocks = block_table[r, sel // BLOCK].long()
+        k, v = sq.reference_dequantize(cache[blocks, :, sel % BLOCK], fmt)
+        if window != "off":
+            slot = blocks * BLOCK + sel % BLOCK
+            exact = ((p - sel) < 64) & (wtags[slot % 4096].long() == slot)
+            exact |= (sel < 4) & (stags[slot % 256].long() == slot)
+            k[exact] = ks[r][sel[exact]].float()
+            v[exact] = vs[r][sel[exact]].float()
+        k = k.repeat_interleave(group, 1)
+        v = v.repeat_interleave(group, 1)
+        p_attn = torch.softmax(torch.einsum("hd,thd->ht", q[i].float(), k) * scale, -1)
+        ref[i] = torch.einsum("ht,thd->hd", p_attn, v)
+    assert _rel(out, ref) < 2e-2

tests/quantization/test_per_token_kv_cache.py

@@ -163,6 +163,19 @@ class TestIsQuantizedKvCache:
 
         assert get_kv_quant_mode("fp8_per_token_head") == KVQuantMode.FP8_PER_TOKEN_HEAD
 
+    @pytest.mark.parametrize(
+        "cache_dtype", ["octave_k3v4", "octave_k3v3", "octave_k3v3_compact"]
+    )
+    def test_octave(self, cache_dtype):
+        from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
+        from vllm.v1.kv_cache_interface import get_kv_quant_mode
+
+        mode = get_kv_quant_mode(cache_dtype)
+        assert mode.name.lower() == cache_dtype
+        assert mode.is_octave and not mode.is_per_token_head
+        assert is_quantized_kv_cache(cache_dtype)
+        assert kv_cache_dtype_str_to_dtype(cache_dtype, None) == torch.uint8
+
 
 # ===========================================================================
 # 2. Triton per-token-head kernel (reshape-and-cache)

vllm/config/cache.py

@@ -55,6 +55,9 @@ CacheDType = Literal[
     "turboquant_4bit_nc",
     "turboquant_k3v4_nc",
     "turboquant_3bit_nc",
+    "octave_k3v4",
+    "octave_k3v3",
+    "octave_k3v3_compact",
     "ultraquant_4bit",
     "int4_per_token_head",
     "int8_per_token_head",

vllm/models/qwen4_exp/amd/qsa.py

@@ -10,7 +10,7 @@ import torch
 from torch import nn
 from transformers import Qwen4ExpTextConfig
 
-from vllm.config import VllmConfig
+from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.config.cache import CacheDType
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.forward_context import get_forward_context
@@ -46,28 +46,93 @@ from vllm.v1.attention.backends.flash_attn import (
     FlashAttentionMetadata,
     FlashAttentionMetadataBuilder,
 )
+from vllm.v1.attention.ops import rocm_octave
 from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
     FullAttentionSpec,
     KVCacheSpec,
     get_kv_quant_mode,
 )
+from vllm.v1.worker.workspace import (
+    current_workspace_manager,
+    is_workspace_manager_initialized,
+)
 
 from ..common.qsa_cache import QSAForwardMetadata
 from . import model
 from .indexer_qsa import QSAIndexer
 
 
+def _octave_mid_o_capacity(vllm_config: VllmConfig) -> int:
+    """Query rows times splits the Octave split-KV partials hold: every split
+    up to the cudagraph capture size, and one split per row of a full chunk."""
+    capture = vllm_config.compilation_config.max_cudagraph_capture_size or 4
+    return max(
+        capture * vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph,
+        2 * vllm_config.scheduler_config.max_num_batched_tokens,
+    )
+
+
+# Exact window of the Octave QSA cache: fp16 shadow tables (entries, powers of
+# two) of every stored token and of each sequence's leading positions; the
+# kernels keep the last 64 and the first 4 selected positions exact.
+_OCTAVE_WINDOW_ENTRIES = 16384
+_OCTAVE_SINK_ENTRIES = 2048
+
+
 class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
     """Flash metadata supporting uniform decode and target-verify graphs."""
 
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
 
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        # Layers run one after another, so they share one set of Octave
+        # split-KV partials; reserve it before the workspace is locked.
+        if kv_cache_spec.kv_quant_mode.is_octave and is_workspace_manager_initialized():
+            num_heads = vllm_config.model_config.get_num_attention_heads(
+                vllm_config.parallel_config
+            )
+            numel = (
+                _octave_mid_o_capacity(vllm_config)
+                * num_heads
+                * (kv_cache_spec.head_size + 2)
+            )
+            current_workspace_manager().get_simultaneous(((numel,), torch.float32))
+
+
+_OCTAVE_CACHE_DTYPES: tuple[CacheDType, ...] = (
+    "octave_k3v4",
+    "octave_k3v3",
+    "octave_k3v3_compact",
+)
+_QSA_CACHE_DTYPES: tuple[CacheDType, ...] = ("auto", "bfloat16", *_OCTAVE_CACHE_DTYPES)
+
 
 class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
-    """FullAttentionSpec backend used by the merged QSA owner."""
+    """FullAttentionSpec backend used by the merged QSA owner.
+
+    With an Octave ``kv_cache_dtype`` the main K/V cache holds one packed
+    Octave slot per (token, KV head) and sparse attention runs the Octave
+    sparse decode kernel over the selected positions.
+    """
 
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
-    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = list(_QSA_CACHE_DTYPES)
+
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        from vllm.v1.attention.backends.rocm_octave_attn import (
+            RocmOctaveAttentionBackend,
+        )
+
+        return RocmOctaveAttentionBackend.customize_spec(spec)
 
     @staticmethod
     def get_name() -> str:
@@ -109,9 +174,106 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
             raise NotImplementedError(
                 "Qwen4Exp QSA does not support decode context parallelism"
             )
-        if self.kv_cache_dtype not in ("auto", "bfloat16"):
+        if self.kv_cache_dtype not in _QSA_CACHE_DTYPES:
-            raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
+            raise NotImplementedError(
+                "Qwen4Exp QSA requires a BF16 or Octave main KV cache"
+            )
         self.supports_quant_query_input = False
+        self.octave_fmt: rocm_octave.OctaveFormat | None = None
+        if self.kv_cache_dtype in _OCTAVE_CACHE_DTYPES:
+            self.octave_fmt = rocm_octave.OctaveFormat.from_cache_dtype(
+                self.kv_cache_dtype,
+                self.head_size,
+                rocm_octave.registered_rope_dim(self.head_size),
+            )
+            self._octave_signs: torch.Tensor | None = None
+            self._octave_mid_o: torch.Tensor | None = None
+            vllm_config = get_current_vllm_config()
+            self._octave_splits = (
+                vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph
+            )
+            self._octave_capacity = _octave_mid_o_capacity(vllm_config)
+            self._octave_positions: torch.Tensor | None = None
+            device = torch.device("cuda", torch.accelerator.current_device_index())
+            dtype = vllm_config.model_config.dtype
+            row = (self.num_kv_heads, 2 * self.head_size)
+
+            def table(entries: int) -> tuple[torch.Tensor, ...]:
+                return (
+                    torch.zeros((entries, *row), dtype=dtype, device=device),
+                    torch.full((entries,), -1, dtype=torch.int32, device=device),
+                    torch.zeros(entries, dtype=torch.int32, device=device),
+                )
+
+            self._wtab, self._wtags, self._wlocks = table(_OCTAVE_WINDOW_ENTRIES)
+            self._stab, self._stags, self._slocks = table(_OCTAVE_SINK_ENTRIES)
+
+    def _octave_buffers(
+        self, num_q: int, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor, int]:
+        """Sign bits and split-KV partials, from the workspace every layer
+        shares (without one, a buffer allocated once, since captured graphs
+        hold its address). Larger batches use fewer splits so they fit."""
+        if self._octave_signs is None or self._octave_signs.device != device:
+            self._octave_signs = rocm_octave.sign_bits(self.head_size).to(device)
+        capacity = self._octave_capacity
+        splits = max(2, min(self._octave_splits, capacity // max(num_q, 1)))
+        row = self.head_size + 2
+        numel = capacity * self.num_heads * row
+        if is_workspace_manager_initialized():
+            (buf,) = current_workspace_manager().get_simultaneous(
+                ((numel,), torch.float32)
+            )
+        else:
+            if self._octave_mid_o is None:
+                self._octave_mid_o = torch.empty(
+                    numel, dtype=torch.float32, device=device
+                )
+            buf = self._octave_mid_o
+        mid_o = buf[: num_q * self.num_heads * splits * row]
+        return (
+            self._octave_signs,
+            mid_o.view(num_q, self.num_heads, splits, row),
+            splits,
+        )
+
+    def do_kv_cache_update(
+        self,
+        layer: torch.nn.Module,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ) -> None:
+        if self.octave_fmt is None:
+            super().do_kv_cache_update(layer, key, value, kv_cache, slot_mapping)
+            return
+        n = slot_mapping.shape[0]
+        if n == 0:
+            return
+        signs, _, _ = self._octave_buffers(1, key.device)
+        torch.ops._C.octave_cache_store(
+            key[:n],
+            value[:n],
+            kv_cache,
+            slot_mapping,
+            signs,
+            signs,
+            self.octave_fmt.kernel_code,
+        )
+        assert self._octave_positions is not None
+        torch.ops._C.octave_window_store(
+            key[:n],
+            value[:n],
+            slot_mapping,
+            self._octave_positions[:n],
+            self._wtab,
+            self._wtags,
+            self._wlocks,
+            self._stab,
+            self._stags,
+            self._slocks,
+        )
 
     def forward_qsa(
         self,
@@ -144,6 +306,29 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
             raise RuntimeError("QSA owner did not provide its top-k buffer")
         logical_indices = topk_buffer[:num_tokens]
         token_to_req = token_to_req[:num_tokens]
+        if self.octave_fmt is not None:
+            assert self._octave_positions is not None
+            signs, mid_o, splits = self._octave_buffers(num_tokens, query.device)
+            torch.ops._C.octave_decode_sparse(
+                output[:num_tokens],
+                query[:num_tokens],
+                kv_cache,
+                attn_metadata.block_table,
+                token_to_req,
+                logical_indices,
+                mid_o,
+                signs,
+                signs,
+                self.scale,
+                splits,
+                self.octave_fmt.kernel_code,
+                self._octave_positions[:num_tokens],
+                self._wtab,
+                self._wtags,
+                self._stab,
+                self._stags,
+            )
+            return output
         key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
         key_cache = canonicalize_singleton_dim_strides(key_cache)
         value_cache = canonicalize_singleton_dim_strides(value_cache)
@@ -186,8 +371,10 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
             raise ValueError("Qwen4Exp QSA requires a paged KV cache")
         if model_config.dtype != torch.bfloat16:
             raise NotImplementedError("Qwen4Exp QSA currently requires BF16")
-        if cache_config.cache_dtype not in ("auto", "bfloat16"):
+        if cache_config.cache_dtype not in _QSA_CACHE_DTYPES:
-            raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
+            raise NotImplementedError(
+                "Qwen4Exp QSA requires a BF16 or Octave main KV cache"
+            )
         if getattr(quant_config, "kv_cache_scheme", None) is not None:
             raise NotImplementedError("Qwen4Exp QSA does not support KV quantization")
         parallel_config = vllm_config.parallel_config
@@ -269,7 +456,10 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
         self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
             self.kv_cache_dtype, model_config
         )
-        if self.kv_cache_torch_dtype != torch.bfloat16:
+        if (
+            self.kv_cache_dtype not in _OCTAVE_CACHE_DTYPES
+            and self.kv_cache_torch_dtype != torch.bfloat16
+        ):
             raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage")
         self.kv_sharing_target_layer_name = None
         self.kv_cache = torch.tensor([])
@@ -296,10 +486,13 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
             quant_config=quant_config,
             prefix=f"{prefix}.indexer",
         )
+        # The fused prepare writes raw BF16 K/V; an Octave cache is packed by
+        # do_kv_cache_update instead.
         self.use_fused_qsa_prepare = (
             self.attn_output_gate
             and self.indexer.use_fused_pre_indexer
             and self.head_dim & (self.head_dim - 1) == 0
+            and self.kv_cache_dtype not in _OCTAVE_CACHE_DTYPES
         )
         max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
         self.register_buffer(
@@ -379,6 +572,9 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
         impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)
         if main_outputs is None:
             assert key is not None and value is not None
+            if impl.octave_fmt is not None:
+                flat = positions[0] if positions.dim() == 2 else positions
+                impl._octave_positions = flat.to(torch.int32)
             impl.do_kv_cache_update(
                 self,
                 key,

vllm/platforms/interface.py

@@ -869,7 +869,8 @@ class Platform:
                 kv_quant_mode=get_kv_quant_mode(cache_dtype),
             )
             # The backend owns its packing
-            return backend_cls.customize_spec(spec).page_size_bytes
+            with set_current_vllm_config(vllm_config):
+                return backend_cls.customize_spec(spec).page_size_bytes
 
         primary_dtype = (
             STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
@@ -1035,9 +1036,11 @@ class Platform:
                 dtype=kv_cache_dtype,
                 kv_quant_mode=kv_quant_mode,
             )
-            attn_page_size_1_token = backend_cls.customize_spec(
+            # customize_spec may read the model config (e.g. Octave's rotary dim).
-                attn_spec
+            with set_current_vllm_config(vllm_config):
-            ).page_size_bytes
+                attn_page_size_1_token = backend_cls.customize_spec(
+                    attn_spec
+                ).page_size_bytes
 
         # Compute mamba page size
         model_cls, _ = ModelRegistry.resolve_model_cls(

vllm/platforms/rocm.py

@@ -376,6 +376,14 @@ def get_cdna_version() -> int:
     return 0
 
 
+def get_rdna_version() -> int:
+    if on_rdna4():
+        return 4
+    if on_gfx11():
+        return 3
+    return 0
+
+
 @cache
 def get_rocm_version() -> tuple[int, ...] | None:
     """Return the installed ROCm release as (major, minor, patch), or None."""
@@ -521,6 +529,7 @@ def _get_backend_priorities(
         backends.insert(0, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
     backends.append(AttentionBackendEnum.TRITON_ATTN)
     backends.append(AttentionBackendEnum.TURBOQUANT)
+    backends.append(AttentionBackendEnum.ROCM_OCTAVE)
     backends.append(AttentionBackendEnum.ULTRAQUANT)
 
     return backends

vllm/utils/torch_utils.py

@@ -57,6 +57,9 @@ STR_DTYPE_TO_TORCH_DTYPE = {
     "turboquant_4bit_nc": torch.uint8,
     "turboquant_k3v4_nc": torch.uint8,
     "turboquant_3bit_nc": torch.uint8,
+    "octave_k3v4": torch.uint8,
+    "octave_k3v3": torch.uint8,
+    "octave_k3v3_compact": torch.uint8,
     "ultraquant_4bit": torch.uint8,
     "nvfp4": torch.uint8,
     "nvfp4_4over6": torch.uint8,

vllm/v1/attention/backends/registry.py

@@ -160,6 +160,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     CPU_MLA = "vllm.v1.attention.backends.mla.cpu_mla.CPUMLABackend"
     AMX_MLA = "vllm.v1.attention.backends.mla.amx_mla.AMXMLABackend"
     TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend"
+    ROCM_OCTAVE = (
+        "vllm.v1.attention.backends.rocm_octave_attn.RocmOctaveAttentionBackend"
+    )
     ULTRAQUANT = "vllm.v1.attention.backends.ultraquant_attn.UltraQuantAttentionBackend"
     # Placeholder for third-party/custom backends - must be registered before use
     # set to None to avoid alias with other backend, whose value is an empty string

vllm/v1/attention/backends/rocm_octave_attn.py

@@ -0,0 +1,496 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Octave attention backend for AMD GPUs (native kernels: HIP or FlyDSL).
+
+KV cache format and reference: ``vllm/v1/attention/ops/rocm_octave.py``.
+Kernels: ``csrc/attention/octave_attn.cu`` and the per-architecture tiers in
+``_kernel_tier``; RDNA3 and CDNA3, validated on gfx1100 and gfx942. See
+``docs/design/rocm_octave.md``.
+
+Decode, MTP verification and short continuation chunks (every request with
+at most ``_DECODE_MAX_QUERY_LEN`` query tokens) run the split-KV decode
+kernel over the packed cache; consecutive query tokens of one request share
+each K/V read. Longer prefills run the prefill kernel, which reads the cached
+prefix from the packed cache and the current chunk's K/V unquantized.
+"""
+
+import contextlib
+import math
+from dataclasses import dataclass, replace
+from typing import ClassVar
+
+import torch
+
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config.cache import CacheDType
+from vllm.platforms import current_platform
+from vllm.platforms.interface import DeviceCapability
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionImpl,
+    AttentionLayer,
+    AttentionMetadataBuilder,
+    AttentionType,
+    CommonAttentionMetadata,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+from vllm.v1.attention.ops import rocm_octave as sq
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
+from vllm.v1.kv_cache_layout import KVCacheLayout
+from vllm.v1.worker.workspace import (
+    current_workspace_manager,
+    is_workspace_manager_initialized,
+)
+
+
+# Native kernel tier per architecture family; any other family is rejected at
+# selection (supports_compute_capability) until it is validated.
+def _kernel_tier() -> str:
+    from vllm.platforms.rocm import get_cdna_version, get_rdna_version
+
+    if get_rdna_version() == 3:
+        return "wmma"
+    if get_cdna_version() == 3:
+        return "mfma"
+    raise NotImplementedError(
+        "ROCM_OCTAVE supports RDNA3 and CDNA3 only; "
+        "see the architecture contract in docs/design/rocm_octave.md"
+    )
+
+
+_DECODE_MAX_QUERY_LEN = 128
+_KERNEL_HEAD_SIZE = 256
+_KERNEL_ROPE_DIM = 64
+
+
+def _mid_o_shape(
+    num_q: int,
+    num_heads: int,
+    max_splits: int,
+    capture_size: int,
+    split_cap: int | None = None,
+) -> tuple[tuple[int, int, int, int], int]:
+    """Split-KV partials for num_q queries: every split up to the cudagraph
+    capture size, fewer splits past it so the buffer never outgrows that.
+    With split_cap, small batches spend the same buffer on up to split_cap
+    splits (a wide GPU needs them to fill its CUs with one query)."""
+    budget = capture_size * max_splits
+    splits = max_splits if split_cap is None else min(split_cap, budget // num_q)
+    if num_q > capture_size:
+        splits = max(1, min(max_splits, budget // num_q))
+    return (num_q, num_heads, splits, _KERNEL_HEAD_SIZE + 2), splits
+
+
+def _max_capture_size(vllm_config: VllmConfig) -> int:
+    return vllm_config.compilation_config.max_cudagraph_capture_size or 4
+
+
+class RocmOctaveAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+    forward_includes_kv_cache_update: bool = False
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "octave_k3v4",
+        "octave_k3v3",
+        "octave_k3v3_compact",
+    ]
+
+    @staticmethod
+    def get_name() -> str:
+        return "ROCM_OCTAVE"
+
+    @staticmethod
+    def get_impl_cls() -> type["RocmOctaveAttentionImpl"]:
+        return RocmOctaveAttentionImpl
+
+    @staticmethod
+    def get_builder_cls() -> type["RocmOctaveMetadataBuilder"]:
+        return RocmOctaveMetadataBuilder
+
+    @staticmethod
+    def get_supported_kernel_block_sizes(
+        kv_cache_spec: "KVCacheSpec | None" = None,
+    ) -> list[int | MultipleOf]:
+        return [MultipleOf(16)]
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        from vllm.platforms.rocm import get_cdna_version, get_rdna_version
+
+        # DeviceCapability is unreliable on ROCm; use the gcnArchName.
+        return current_platform.is_rocm() and (
+            get_rdna_version() == 3 or get_cdna_version() == 3
+        )
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [_KERNEL_HEAD_SIZE]
+
+    @classmethod
+    def supports_attn_type(cls, attn_type: str) -> bool:
+        return attn_type == AttentionType.DECODER
+
+    @classmethod
+    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
+        # The kernels take block/head/token strides; the slot is contiguous.
+        # The MFMA decode loads 16 consecutive tokens of a head as one tile.
+        if current_platform.is_rocm():
+            from vllm.platforms.rocm import get_cdna_version
+
+            if get_cdna_version() == 3:
+                return (KVCacheLayout.LBHNC,)
+        return (KVCacheLayout.LBHNC, KVCacheLayout.LBNHC)
+
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        """One packed K+V slot per (token, KV head)."""
+        if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_octave:
+            return spec
+        fmt = sq.OctaveFormat.from_cache_dtype(
+            spec.kv_quant_mode.name.lower(),
+            spec.head_size,
+            sq.registered_rope_dim(spec.head_size),
+        )
+        return replace(spec, state_content_bytes=fmt.slot_bytes)
+
+
+@dataclass
+class RocmOctaveMetadata:
+    num_actual_tokens: int
+    max_query_len: int
+    query_start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    block_table: torch.Tensor
+    slot_mapping: torch.Tensor
+    # Decode requests (one query token) come first in the batch.
+    num_decodes: int
+    num_decode_tokens: int
+    # Per query token: its request and its causal K length (0 for padding).
+    q_to_req: torch.Tensor
+    q_to_klen: torch.Tensor
+
+
+class RocmOctaveMetadataBuilder(AttentionMetadataBuilder[RocmOctaveMetadata]):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    supports_draft_decode_metadata_update = True
+
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ):
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self._init_reorder_batch_threshold(1, supports_spec_as_decode=False)
+        # Persistent so captured graphs keep reading the current contents.
+        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        self._q_to_req = torch.empty(max_tokens, dtype=torch.int32, device=device)
+        self._q_to_klen = torch.empty(max_tokens, dtype=torch.int32, device=device)
+        # Layers run one after another, so they share one set of split-KV
+        # partials; reserve it before the workspace is locked.
+        if is_workspace_manager_initialized():
+            num_heads = vllm_config.model_config.get_num_attention_heads(
+                vllm_config.parallel_config
+            )
+            capture_size = _max_capture_size(vllm_config)
+            shape, _ = _mid_o_shape(
+                capture_size,
+                num_heads,
+                vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph,
+                capture_size,
+            )
+            current_workspace_manager().get_simultaneous((shape, torch.float32))
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> RocmOctaveMetadata:
+        cm = common_attn_metadata
+        n = cm.num_actual_tokens
+        num_decodes, _, num_decode_tokens, _ = split_decodes_and_prefills(cm)
+        md = RocmOctaveMetadata(
+            num_actual_tokens=n,
+            max_query_len=cm.max_query_len,
+            query_start_loc=cm.query_start_loc,
+            seq_lens=cm.seq_lens,
+            block_table=cm.block_table_tensor,
+            slot_mapping=cm.slot_mapping,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            q_to_req=self._q_to_req[:n],
+            q_to_klen=self._q_to_klen[:n],
+        )
+        self._fill_query_maps(md)
+        return md
+
+    def update_draft_decode_metadata(self, metadata: RocmOctaveMetadata) -> None:
+        # Draft steps advance seq_lens in place; the K lengths follow it.
+        self._fill_query_maps(metadata)
+
+    @staticmethod
+    def _fill_query_maps(md: RocmOctaveMetadata) -> None:
+        """On the device, without host syncs (also inside graph capture).
+        Queries past the last request (cudagraph padding) get K length 0."""
+        n = md.num_actual_tokens
+        num_reqs = md.query_start_loc.shape[0] - 1
+        if n == 0 or num_reqs == 0:
+            return
+        idx = torch.arange(n, dtype=torch.int32, device=md.q_to_req.device)
+        qsl = md.query_start_loc.to(torch.int32)
+        req = torch.searchsorted(qsl[1:], idx, right=True).to(torch.int32)
+        r = req.clamp(max=num_reqs - 1).long()
+        q_lens = qsl[1:] - qsl[:-1]
+        klen = md.seq_lens.to(torch.int32)[r] - q_lens[r] + (idx - qsl[r]) + 1
+        md.q_to_req.copy_(req)
+        md.q_to_klen.copy_(torch.where(req < num_reqs, klen, 0))
+
+
+class RocmOctaveAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int | None = None,
+        alibi_slopes: list[float] | None = None,
+        sliding_window: int | None = None,
+        kv_cache_dtype: str = "auto",
+        logits_soft_cap: float | None = None,
+        attn_type: str = AttentionType.DECODER,
+        kv_sharing_target_layer_name: str | None = None,
+        sinks: torch.Tensor | None = None,
+        **kwargs,
+    ) -> None:
+        if alibi_slopes is not None or sinks is not None:
+            raise NotImplementedError("ROCM_OCTAVE does not support ALiBi or sinks")
+        if sliding_window is not None or logits_soft_cap:
+            raise NotImplementedError(
+                "ROCM_OCTAVE does not support sliding window or soft cap"
+            )
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError("ROCM_OCTAVE supports decoder attention only")
+
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.kv_cache_dtype = kv_cache_dtype
+        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
+
+        vllm_config = get_current_vllm_config()
+        rope_dim = sq.rope_dim_from_hf_config(
+            vllm_config.model_config.hf_text_config, head_size
+        )
+        sq.register_rope_dim(head_size, rope_dim)
+        self.fmt = sq.OctaveFormat.from_cache_dtype(kv_cache_dtype, head_size, rope_dim)
+        if not (
+            current_platform.is_rocm()
+            and head_size == _KERNEL_HEAD_SIZE
+            and rope_dim == _KERNEL_ROPE_DIM
+            and hasattr(torch.ops._C, "octave_decode")
+        ):
+            raise NotImplementedError(
+                "ROCM_OCTAVE kernels need ROCm, head_size=256 and 64 rotary "
+                f"dims; got head_size={head_size}, rotary dims={rope_dim}"
+            )
+        # The MFMA decode takes up to 256 splits (octave_decode_mfma_cdna.cu).
+        self._split_cap = 256 if _kernel_tier() == "mfma" else None
+
+        self.max_num_kv_splits = (
+            vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph
+        )
+        self._max_capture_size = _max_capture_size(vllm_config)
+        self._k_signs: torch.Tensor | None = None
+        self._v_signs: torch.Tensor | None = None
+        self._mid_o: torch.Tensor | None = None
+
+    def _signs(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        if self._k_signs is None or self._k_signs.device != device:
+            self._k_signs = sq.sign_bits(self.head_size).to(device)
+            self._v_signs = self._k_signs
+        assert self._v_signs is not None
+        return self._k_signs, self._v_signs
+
+    def _mid_o_buffer(
+        self, num_q: int, device: torch.device
+    ) -> tuple[torch.Tensor, int]:
+        """Split-KV partials, from the workspace every layer shares. Without a
+        workspace, a buffer of this layer's own, allocated once at the
+        cudagraph capture size since captured graphs hold its address."""
+        shape, splits = _mid_o_shape(
+            num_q,
+            self.num_heads,
+            self.max_num_kv_splits,
+            self._max_capture_size,
+            self._split_cap,
+        )
+        if is_workspace_manager_initialized():
+            (buf,) = current_workspace_manager().get_simultaneous(
+                (shape, torch.float32)
+            )
+            return buf, splits
+        if self._mid_o is None:
+            full, _ = _mid_o_shape(
+                self._max_capture_size,
+                self.num_heads,
+                self.max_num_kv_splits,
+                self._max_capture_size,
+            )
+            self._mid_o = torch.empty(full, dtype=torch.float32, device=device)
+        if math.prod(shape) <= self._mid_o.numel():
+            return self._mid_o.view(-1)[: math.prod(shape)].view(shape), splits
+        return torch.empty(shape, dtype=torch.float32, device=device), splits
+
+    def do_kv_cache_update(
+        self,
+        layer: AttentionLayer,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ) -> None:
+        if slot_mapping.numel() == 0:
+            return
+        k_signs, v_signs = self._signs(key.device)
+        n = slot_mapping.shape[0]
+        torch.ops._C.octave_cache_store(
+            key[:n],
+            value[:n],
+            kv_cache,
+            slot_mapping,
+            k_signs,
+            v_signs,
+            self.fmt.kernel_code,
+        )
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: RocmOctaveMetadata,
+        output: torch.Tensor | None = None,
+        output_scale: torch.Tensor | None = None,
+        output_block_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        assert output is not None
+        if output_scale is not None or output_block_scale is not None:
+            raise NotImplementedError("ROCM_OCTAVE does not fuse output quant")
+        if attn_metadata is None:
+            return output.fill_(0)
+
+        n = attn_metadata.num_actual_tokens
+        out = output[:n].view(n, self.num_heads, self.head_size)
+        q = query[:n].view(n, self.num_heads, self.head_size)
+
+        if attn_metadata.max_query_len <= _DECODE_MAX_QUERY_LEN:
+            self._decode(
+                out,
+                q,
+                kv_cache,
+                attn_metadata,
+                attn_metadata.q_to_req,
+                attn_metadata.q_to_klen,
+            )
+            return output
+
+        num_dt = attn_metadata.num_decode_tokens
+        if num_dt > 0:
+            q_to_req, q_to_klen = attn_metadata.q_to_req, attn_metadata.q_to_klen
+            self._decode(
+                out[:num_dt],
+                q[:num_dt],
+                kv_cache,
+                attn_metadata,
+                q_to_req[:num_dt],
+                q_to_klen[:num_dt],
+            )
+        self._prefill(out, q, key, value, kv_cache, attn_metadata)
+        return output
+
+    def _decode(
+        self,
+        out: torch.Tensor,
+        q: torch.Tensor,
+        kv_cache: torch.Tensor,
+        md: RocmOctaveMetadata,
+        q_to_req: torch.Tensor,
+        q_to_klen: torch.Tensor,
+    ) -> None:
+        num_q = q.shape[0]
+        if num_q == 0:
+            return
+        k_signs, v_signs = self._signs(q.device)
+        mid_o, splits = self._mid_o_buffer(num_q, q.device)
+        # Async scheduling can free and reuse an input's memory before this
+        # kernel reads it; tell the allocator these are in use on the stream.
+        stream = torch.cuda.current_stream()
+        for t in (q, out, md.block_table, q_to_req, q_to_klen, mid_o):
+            with contextlib.suppress(Exception):
+                t.record_stream(stream)
+        torch.ops._C.octave_decode(
+            out,
+            q,
+            kv_cache,
+            md.block_table,
+            q_to_req,
+            q_to_klen,
+            mid_o,
+            k_signs,
+            v_signs,
+            self.scale,
+            splits,
+            self.fmt.kernel_code,
+            4 if md.max_query_len > 1 else 1,
+            True,  # native tier; False runs the portable reference kernel
+        )
+
+    def _prefill(
+        self,
+        out: torch.Tensor,
+        q: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        md: RocmOctaveMetadata,
+    ) -> None:
+        """Requests after the decodes, in the rotated space, without host
+        syncs: one call covers them all, reading the cached prefixes straight
+        from the packed cache and the chunks' own K/V unquantized."""
+        nd, ndt, n = md.num_decodes, md.num_decode_tokens, md.num_actual_tokens
+        num_reqs = md.query_start_loc.shape[0] - 1
+        if num_reqs - nd <= 0 or ndt >= n:
+            return
+        k_signs, v_signs = self._signs(q.device)
+        # The kernel runs in fp16; the rotated copies are made in fp16.
+        q_rot = q[ndt:n].to(torch.float16, copy=True)
+        k_rot = key[ndt:n].to(torch.float16, copy=True)
+        v_rot = value[ndt:n].to(torch.float16, copy=True)
+        k_blocks = not self.fmt.compact
+        torch.ops._C.octave_rotate(q_rot, k_signs, k_blocks, False)
+        torch.ops._C.octave_rotate(k_rot, k_signs, k_blocks, False)
+        torch.ops._C.octave_rotate(v_rot, v_signs, False, False)
+        o_rot = torch.empty_like(q_rot)
+        torch.ops._C.octave_prefill(
+            o_rot,
+            q_rot,
+            k_rot,
+            v_rot,
+            kv_cache,
+            md.block_table[nd:num_reqs],
+            (md.query_start_loc[nd:] - ndt).to(torch.int32),
+            md.seq_lens[nd:num_reqs].to(torch.int32),
+            md.max_query_len,
+            self.scale,
+            self.fmt.kernel_code,
+        )
+        torch.ops._C.octave_rotate(o_rot, v_signs, False, True)
+        out[ndt:n] = o_rot

vllm/v1/attention/ops/rocm_octave.py

@@ -0,0 +1,390 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Octave KV-cache format: layout, rotations and a PyTorch reference.
+
+Each (token, KV head) is stored in one byte slot. V is always quantized
+whole; K has two layouts:
+
+* Block K (``octave_k3v4``, ``octave_k3v3``): random sign flips plus a
+  Walsh-Hadamard transform over each 64-dim block. The RoPE blocks get 4-bit
+  codes, the NoPE blocks 3-bit codes, and every block has its own scale.
+* Compact K (``octave_k3v3_compact``): sign flips plus one Hadamard transform
+  over all dims, 3-bit codes, one scale. Spreading the RoPE dims over the
+  whole head lets them share the 3-bit budget.
+* V: sign flips plus one Hadamard transform over all dims, ``v_bits`` codes.
+
+The rotated coordinates are close to Gaussian, so each code indexes a
+Lloyd-Max codebook for N(0, 1), stored as int8 (``LUT``): a code ``c`` of a
+block with scale ``s`` stands for ``LUT[c] * s / LUT_ONE``. The scale makes
+``x . x_hat = |x|^2``: a least-squares fit would shrink every reconstructed
+vector, which flattens the softmax and scales down the attention output.
+Everything is computed on the fly; there is no calibration.
+
+Scores are computed in the rotated space (the query is rotated like K) and
+the attention output is rotated back once per query head.
+
+Slot layout (bytes)::
+
+    [k_rope codes][k_nope codes][v codes][fp16 scales][pad to 8]
+
+Compact K has no 4-bit part and pads to 4 bytes.
+
+4-bit codes: each little-endian uint32 covers 8 dims; byte j holds dim j in
+its low nibble and dim j + 4 in its high nibble, so ``w & 0x0F0F0F0F`` and
+``(w >> 4) & 0x0F0F0F0F`` give dims 0..3 and 4..7 as bytes.
+
+3-bit codes: a 2-bit plane (uint32 per 16 dims; ``(w >> 2i) & 0x03030303``
+gives dims 4i..4i+3) followed by a 1-bit plane (uint32 per 32 dims;
+``(w >> i) & 0x01010101`` gives dims 4i..4i+3). code = lo + 4 * hi.
+
+Scales (fp16): one per 64-dim K block (one in total for compact K), then V.
+"""
+
+import functools
+import math
+from dataclasses import dataclass
+
+import torch
+
+HADAMARD_BLOCK = 64
+_SIGN_SEED = 0x5E17
+ROPE_BITS = 4
+NOPE_BITS = 3
+# Lloyd-Max codebooks for N(0, 1) as int8: LUT[bits][c] / LUT_MULT[bits]
+# approximates centroid c (the kernels map codes through the same tables).
+LUT = {
+    3: (-127, -79, -45, -14, 14, 45, 79, 127),
+    4: (-126, -95, -74, -58, -43, -30, -18, -6, 6, 18, 30, 43, 58, 74, 95, 126),
+}
+LUT_MULT = {3: 59.0, 4: 46.0}
+# Fixed point of the stored scales: a code stands for LUT[c] * scale / LUT_ONE,
+# which keeps small blocks' scales out of the fp16 subnormal range.
+LUT_ONE = 64.0
+
+_rope_dim_by_head_size: dict[int, int] = {}
+
+
+def register_rope_dim(head_size: int, rope_dim: int) -> None:
+    prev = _rope_dim_by_head_size.setdefault(head_size, rope_dim)
+    if prev != rope_dim:
+        raise ValueError(
+            f"Octave: conflicting rotary dims {prev} and {rope_dim} for "
+            f"head_size={head_size}"
+        )
+
+
+def rope_dim_from_hf_config(hf_text_config, head_size: int) -> int:
+    factor = getattr(hf_text_config, "partial_rotary_factor", None)
+    if factor is None:
+        rope_params = getattr(hf_text_config, "rope_parameters", None) or {}
+        factor = rope_params.get("partial_rotary_factor", 1.0)
+    return int(head_size * factor)
+
+
+def registered_rope_dim(head_size: int) -> int:
+    if head_size not in _rope_dim_by_head_size:
+        from vllm.config import get_current_vllm_config_or_none
+
+        vllm_config = get_current_vllm_config_or_none()
+        if vllm_config is None or vllm_config.model_config is None:
+            raise RuntimeError(
+                "Octave needs the model's rotary dim before sizing the KV cache"
+            )
+        register_rope_dim(
+            head_size,
+            rope_dim_from_hf_config(vllm_config.model_config.hf_text_config, head_size),
+        )
+    return _rope_dim_by_head_size[head_size]
+
+
+@dataclass(frozen=True)
+class OctaveFormat:
+    head_size: int
+    rope_dim: int
+    v_bits: int
+    compact: bool = False
+
+    def __post_init__(self):
+        d, r = self.head_size, self.rope_dim
+        if self.v_bits not in (3, 4):
+            raise ValueError(f"Octave supports 3- or 4-bit V, got {self.v_bits}")
+        if self.compact and d & (d - 1):
+            raise ValueError(
+                f"Compact Octave K needs a power-of-two head_size, got {d}"
+            )
+        if r % HADAMARD_BLOCK or d % HADAMARD_BLOCK:
+            raise ValueError(
+                f"Octave needs head_size and rope_dim in multiples of "
+                f"{HADAMARD_BLOCK}; got head_size={d}, rope_dim={r}"
+            )
+
+    @classmethod
+    def from_cache_dtype(cls, cache_dtype: str, head_size: int, rope_dim: int):
+        v_bits, compact = {
+            "octave_k3v4": (4, False),
+            "octave_k3v3": (3, False),
+            "octave_k3v3_compact": (3, True),
+        }[cache_dtype]
+        return cls(head_size, rope_dim, v_bits, compact)
+
+    @property
+    def kernel_code(self) -> int:
+        """Format id the kernels take: V bits, plus 16 for compact K."""
+        return self.v_bits + (16 if self.compact else 0)
+
+    @property
+    def k_rotation_block(self) -> int:
+        return self.head_size if self.compact else HADAMARD_BLOCK
+
+    @property
+    def k_wide_dim(self) -> int:
+        """K dims stored with ROPE_BITS codes (the rest use NOPE_BITS)."""
+        return 0 if self.compact else self.rope_dim
+
+    @property
+    def nope_dim(self) -> int:
+        return self.head_size - self.rope_dim
+
+    @property
+    def num_k_blocks(self) -> int:
+        return self.head_size // HADAMARD_BLOCK
+
+    @property
+    def num_rope_blocks(self) -> int:
+        return self.rope_dim // HADAMARD_BLOCK
+
+    @property
+    def off_k_nope(self) -> int:
+        return self.k_wide_dim * ROPE_BITS // 8
+
+    @property
+    def off_v(self) -> int:
+        return self.off_k_nope + (self.head_size - self.k_wide_dim) * NOPE_BITS // 8
+
+    @property
+    def off_scales(self) -> int:
+        return self.off_v + self.head_size * self.v_bits // 8
+
+    @property
+    def num_k_scales(self) -> int:
+        return 1 if self.compact else self.num_k_blocks
+
+    @property
+    def num_scales(self) -> int:
+        return self.num_k_scales + 1
+
+    @property
+    def slot_bytes(self) -> int:
+        raw = self.off_scales + 2 * self.num_scales
+        align = 4 if self.compact else 8
+        return (raw + align - 1) // align * align
+
+
+@functools.cache
+def _signs_cpu(n: int) -> torch.Tensor:
+    g = torch.Generator().manual_seed(_SIGN_SEED)
+    return torch.randint(0, 2, (n,), generator=g, dtype=torch.int64) * -2 + 1
+
+
+@functools.cache
+def signs(n: int, device_str: str) -> torch.Tensor:
+    """Fixed ±1 pattern applied before the Hadamard transform."""
+    return _signs_cpu(n).to(device=device_str, dtype=torch.float32)
+
+
+def sign_bits(n: int) -> torch.Tensor:
+    """The sign pattern as int32 words, bit i set where the sign is -1."""
+    neg = (_signs_cpu(n) < 0).to(torch.int64).view(-1, 32)
+    words = (neg << torch.arange(32)).sum(-1)
+    return torch.where(words >= 2**31, words - 2**32, words).to(torch.int32)
+
+
+def fwht_blocks(x: torch.Tensor, block: int = HADAMARD_BLOCK) -> torch.Tensor:
+    """Orthonormal Walsh-Hadamard transform over consecutive blocks of the
+    last dim."""
+    shape = x.shape
+    y = x.reshape(-1, shape[-1] // block, block).float()
+    h = 1
+    while h < block:
+        y = y.view(*y.shape[:-1], block // (2 * h), 2, h)
+        a, b = y[..., 0, :], y[..., 1, :]
+        y = torch.stack((a + b, a - b), dim=-2).view(*y.shape[:-3], block)
+        h *= 2
+    return (y * (1.0 / math.sqrt(block))).view(shape)
+
+
+def rotate(
+    x: torch.Tensor, sgn: torch.Tensor, block: int = HADAMARD_BLOCK
+) -> torch.Tensor:
+    return fwht_blocks(x.float() * sgn, block)
+
+
+def unrotate(
+    y: torch.Tensor, sgn: torch.Tensor, block: int = HADAMARD_BLOCK
+) -> torch.Tensor:
+    return fwht_blocks(y.float(), block) * sgn
+
+
+def v_block(head_size: int) -> int:
+    """V is rotated as one block when head_size is a power of two."""
+    if head_size & (head_size - 1) == 0:
+        return head_size
+    return HADAMARD_BLOCK
+
+
+@functools.cache
+def _lut(bits: int, device_str: str) -> torch.Tensor:
+    return torch.tensor(LUT[bits], dtype=torch.float32, device=device_str)
+
+
+def _quantize_lut(x: torch.Tensor, bits: int):
+    """Nearest Lloyd-Max code on the block RMS; the scale keeps
+    ``x . x_hat = |x|^2``.
+
+    Returns (codes uint8 in [0, 2**bits), scale) with
+    x ≈ LUT[codes] * scale / LUT_ONE.
+    """
+    lut = _lut(bits, str(x.device))
+    ss = x.pow(2).sum(-1, keepdim=True)
+    rms = (ss / x.shape[-1]).sqrt().clamp_min(1e-12)
+    y = x / rms * LUT_MULT[bits]
+    codes = torch.bucketize(y, (lut[1:] + lut[:-1]) / 2)
+    dot = (x * lut[codes]).sum(-1, keepdim=True)
+    scale = torch.where(dot > 0, ss / dot.clamp_min(1e-30), 0.0) * LUT_ONE
+    return codes.to(torch.uint8), scale
+
+
+def _pack(codes: torch.Tensor, bits: int) -> torch.Tensor:
+    """codes: (..., n) uint8 -> (..., n * bits / 8) uint8, layout above."""
+    c = codes.to(torch.int64)
+    if bits == 4:
+        g = c.view(*c.shape[:-1], -1, 2, 4)  # (word, half, byte)
+        b = g[..., 0, :] | (g[..., 1, :] << 4)
+        return b.reshape(*c.shape[:-1], -1).to(torch.uint8)
+    lo = (c & 3).view(*c.shape[:-1], -1, 4, 4)  # (word, i, byte)
+    lo_b = (lo << (2 * torch.arange(4, device=c.device)).view(4, 1)).sum(-2)
+    hi = (c >> 2).view(*c.shape[:-1], -1, 8, 4)
+    hi_b = (hi << torch.arange(8, device=c.device).view(8, 1)).sum(-2)
+    return torch.cat(
+        (lo_b.reshape(*c.shape[:-1], -1), hi_b.reshape(*c.shape[:-1], -1)), -1
+    ).to(torch.uint8)
+
+
+def _unpack(packed: torch.Tensor, n: int, bits: int) -> torch.Tensor:
+    p = packed.to(torch.int64)
+    if bits == 4:
+        g = p.view(*p.shape[:-1], -1, 4)
+        c = torch.stack((g & 15, g >> 4), dim=-2)
+        return c.reshape(*p.shape[:-1], n)
+    lo_bytes = n // 4
+    lo = p[..., :lo_bytes].view(*p.shape[:-1], -1, 1, 4)
+    lo = (lo >> (2 * torch.arange(4, device=p.device)).view(4, 1)) & 3
+    hi = p[..., lo_bytes:].view(*p.shape[:-1], -1, 1, 4)
+    hi = (hi >> torch.arange(8, device=p.device).view(8, 1)) & 1
+    return lo.reshape(*p.shape[:-1], n) + 4 * hi.reshape(*p.shape[:-1], n)
+
+
+def reference_quantize(
+    key: torch.Tensor, value: torch.Tensor, fmt: OctaveFormat
+) -> torch.Tensor:
+    """(T, H, D) K and V -> (T, H, slot_bytes) uint8 slots."""
+    t, h, d = key.shape
+    r, dev = fmt.rope_dim, key.device
+    out = torch.zeros(t, h, fmt.slot_bytes, dtype=torch.uint8, device=dev)
+    nkb = fmt.num_k_blocks
+
+    if fmt.compact:
+        k = rotate(key, signs(d, str(dev)), d)
+        k_codes, k_scale = _quantize_lut(k, NOPE_BITS)
+        out[..., : fmt.off_v] = _pack(k_codes, NOPE_BITS)
+    else:
+        k = rotate(key, signs(d, str(dev))).view(t, h, nkb, HADAMARD_BLOCK)
+        nr = fmt.num_rope_blocks
+        kr_codes, kr_scale = _quantize_lut(k[:, :, :nr], ROPE_BITS)
+        kn_codes, kn_scale = _quantize_lut(k[:, :, nr:], NOPE_BITS)
+        out[..., : fmt.off_k_nope] = _pack(kr_codes.view(t, h, r), ROPE_BITS)
+        out[..., fmt.off_k_nope : fmt.off_v] = _pack(
+            kn_codes.view(t, h, fmt.nope_dim), NOPE_BITS
+        )
+        k_scale = torch.cat((kr_scale[..., 0], kn_scale[..., 0]), -1)
+
+    vr = rotate(value, signs(d, str(dev)), v_block(d))
+    v_codes, v_scale = _quantize_lut(vr, fmt.v_bits)
+    out[..., fmt.off_v : fmt.off_scales] = _pack(v_codes, fmt.v_bits)
+
+    scales = torch.cat((k_scale, v_scale), -1)
+    sc = scales.to(torch.float16).view(torch.uint8).view(t, h, -1)
+    out[..., fmt.off_scales : fmt.off_scales + sc.shape[-1]] = sc
+    return out
+
+
+def reference_dequantize(
+    slots: torch.Tensor, fmt: OctaveFormat, rotated: bool = False
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """(T, H, slot) slots -> fp32 K and V of shape (T, H, D).
+
+    With ``rotated=True`` K and V stay in the rotated space, which is what
+    the kernels consume.
+    """
+    t, h, _ = slots.shape
+    r, d, dev = fmt.rope_dim, fmt.head_size, slots.device
+    sc = (
+        slots[..., fmt.off_scales : fmt.off_scales + 2 * fmt.num_scales]
+        .contiguous()
+        .view(torch.float16)
+        .float()
+        / LUT_ONE
+    )
+    if fmt.compact:
+        k = _unpack(slots[..., : fmt.off_v], d, NOPE_BITS)
+        k = _lut(NOPE_BITS, str(dev))[k] * sc[..., :1]
+    else:
+        kr = _unpack(slots[..., : fmt.off_k_nope], r, ROPE_BITS)
+        kn = _unpack(slots[..., fmt.off_k_nope : fmt.off_v], fmt.nope_dim, NOPE_BITS)
+        k = torch.cat(
+            (_lut(ROPE_BITS, str(dev))[kr], _lut(NOPE_BITS, str(dev))[kn]), -1
+        ).view(t, h, fmt.num_k_blocks, HADAMARD_BLOCK)
+        k = (k * sc[..., : fmt.num_k_blocks, None]).view(t, h, d)
+
+    v = _unpack(slots[..., fmt.off_v : fmt.off_scales], d, fmt.v_bits)
+    v = _lut(fmt.v_bits, str(dev))[v] * sc[..., -1:]
+    if not rotated:
+        k = unrotate(k, signs(d, str(dev)), fmt.k_rotation_block)
+        v = unrotate(v, signs(d, str(dev)), v_block(d))
+    return k, v
+
+
+def reference_attention(
+    query: torch.Tensor,
+    cache: torch.Tensor,
+    block_table: torch.Tensor,
+    q_to_req: torch.Tensor,
+    q_to_klen: torch.Tensor,
+    fmt: OctaveFormat,
+    scale: float,
+) -> torch.Tensor:
+    """Per-query-token attention over a Octave paged cache.
+
+    query: (Q, Hq, D). cache: (num_blocks, Hkv, block_size, slot) uint8.
+    Query token i attends to the first ``q_to_klen[i]`` tokens of request
+    ``q_to_req[i]``.
+    """
+    num_q, hq, d = query.shape
+    _, hkv, block_size, _ = cache.shape
+    group = hq // hkv
+    out = torch.zeros(num_q, hq, d, dtype=torch.float32, device=query.device)
+    for i in range(num_q):
+        klen = int(q_to_klen[i])
+        if klen <= 0:
+            continue
+        req = int(q_to_req[i])
+        pos = torch.arange(klen, device=cache.device)
+        blocks = block_table[req, pos // block_size].long()
+        slots = cache[blocks, :, pos % block_size]  # (klen, Hkv, slot)
+        k, v = reference_dequantize(slots, fmt)
+        k = k.repeat_interleave(group, dim=1)
+        v = v.repeat_interleave(group, dim=1)
+        s = torch.einsum("hd,thd->ht", query[i].float(), k) * scale
+        p = torch.softmax(s, dim=-1)
+        out[i] = torch.einsum("ht,thd->hd", p, v)
+    return out

vllm/v1/kv_cache_interface.py

@@ -57,6 +57,10 @@ class KVQuantMode(IntEnum):
     TURBOQUANT_3BIT_NC = 9
     NVFP4_DS_MLA = 10  # opaque-bytes NVFP4 DS-MLA layouts (FlashMLA sparse)
     ULTRAQUANT_4BIT = 11  # packed FP4 values + UE8M0 group scales
+    # Hadamard-rotated Lloyd-Max codes, packed K+V per slot (ROCm).
+    OCTAVE_K3V4 = 12
+    OCTAVE_K3V3 = 13
+    OCTAVE_K3V3_COMPACT = 14
 
     # Backend self-manages kernel dispatch. Reads as quantized but matches no
     # generic kernel mode (is_per_token_head / is_nvfp4 / is_turboquant all
@@ -78,6 +82,15 @@ class KVQuantMode(IntEnum):
         """True for NVFP4 packed quantization mode."""
         return self == KVQuantMode.NVFP4
 
+    @property
+    def is_octave(self) -> bool:
+        """True for any Octave quantization mode."""
+        return self in (
+            KVQuantMode.OCTAVE_K3V4,
+            KVQuantMode.OCTAVE_K3V3,
+            KVQuantMode.OCTAVE_K3V3_COMPACT,
+        )
+
     @property
     def is_turboquant(self) -> bool:
         """True for any turboquant quantization mode."""
@@ -127,6 +140,8 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
         return KVQuantMode.NVFP4
     if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):
         return KVQuantMode[kv_cache_dtype.upper()]
+    if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("octave_"):
+        return KVQuantMode[kv_cache_dtype.upper()]
     if kv_cache_dtype == "ultraquant_4bit":
         return KVQuantMode.ULTRAQUANT_4BIT
     if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):