mirrors/vllm
0
0 Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-10-28 09:45:26 +08:00
Code Activity

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

Browse Source 
Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
This commit is contained in:
JartXandtjtanaa authored and GitHub committed 2026-10-26 12:35:00 +08:00
1 parent d7595a7678
commit 00b6c4c565
30 files changed
+6193 -36
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/test-amd.yaml
+19
View File
@@ -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
+21
View File
@@ -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
+21
View File
@@ -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
+13 -2
View File
@@ -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"],
help="KV cache dtype: auto or fp8",
choices=get_args(CacheDType),
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
+3 -1
View File
@@ -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
+26 -21
View File
@@ -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",
tokenizer="meta-llama/Meta-Llama-3-8B",
model=config.model,
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"):
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,
)
spec = _kv_cache_spec(config, 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(
block_size=config.block_size,
num_kv_heads=config.num_kv_heads,
head_size=config.head_dim,
dtype=dtype,
)
kv_cache_spec = backend_class.customize_spec(_kv_cache_spec(config, dtype))
builder = _create_metadata_builder(
backend_class, kv_cache_spec, vllm_config, device, config.backend
csrc/attention/octave_arch.cuh
+36
View File
@@ -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_attn.cu
+1180
View File
File diff suppressed because it is too large. Load diff
csrc/attention/octave_decode_mfma_cdna.cu
+719
View File
@@ -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
+687
View File
@@ -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
+78
View File
@@ -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
+666
View File
@@ -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
+809
View File
@@ -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
+82
View File
@@ -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
+197
View File
@@ -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
+5
View File
@@ -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
+5
View File
@@ -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
+5
View File
@@ -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
+3
View File
@@ -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
+475
View File
@@ -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
+13
View File
@@ -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
+3
View File
@@ -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
+204 -8
View File
@@ -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"):
raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
if self.kv_cache_dtype not in _QSA_CACHE_DTYPES:
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"):
raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
if cache_config.cache_dtype not in _QSA_CACHE_DTYPES:
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
+7 -4
View File
@@ -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(
attn_spec
).page_size_bytes
# customize_spec may read the model config (e.g. Octave's rotary dim).
with set_current_vllm_config(vllm_config):
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
+9
View File
@@ -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
+3
View File
@@ -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
+3
View File
@@ -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
+496
View File
@@ -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
+390
View File
@@ -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
+15
View File
@@ -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"):