From 00b6c4c565a8732c131e1f452d4a17803bb0842b Mon Sep 17 00:00:00 2001 From: JartX Date: Mon, 26 Oct 2026 05:35:00 +0100 Subject: [PATCH] [ROCm] Octave KV (`ROCM_OCTAVE`): native 3-bit KV cache backend for AMD GPUs (#59774) Signed-off-by: JartX Signed-off-by: tjtanaa Co-authored-by: tjtanaa --- .buildkite/test-amd.yaml | 19 + .buildkite/test_areas/lm_eval.yaml | 21 + CMakeLists.txt | 21 + benchmarks/attention_benchmarks/benchmark.py | 15 +- benchmarks/attention_benchmarks/common.py | 4 +- benchmarks/attention_benchmarks/runner.py | 47 +- csrc/attention/octave_arch.cuh | 36 + csrc/attention/octave_attn.cu | 1180 +++++++++++++++++ csrc/attention/octave_decode_mfma_cdna.cu | 719 ++++++++++ csrc/attention/octave_decode_wmma_rdna3.cu | 687 ++++++++++ csrc/attention/octave_format.cuh | 78 ++ csrc/attention/octave_prefill_cdna.cu | 666 ++++++++++ csrc/attention/octave_prefill_rdna3.cu | 809 +++++++++++ csrc/attention/octave_torch_bindings.cpp | 82 ++ docs/design/rocm_octave.md | 197 +++ .../Qwen3.5-4B-Octave-k3v3-compact.yaml | 5 + .../gsm8k/configs/Qwen3.5-4B-Octave-k3v3.yaml | 5 + .../gsm8k/configs/Qwen3.5-4B-Octave-k3v4.yaml | 5 + tests/evals/gsm8k/configs/models-octave.txt | 3 + tests/kernels/attention/test_rocm_octave.py | 475 +++++++ tests/quantization/test_per_token_kv_cache.py | 13 + vllm/config/cache.py | 3 + vllm/models/qwen4_exp/amd/qsa.py | 212 ++- vllm/platforms/interface.py | 11 +- vllm/platforms/rocm.py | 9 + vllm/utils/torch_utils.py | 3 + vllm/v1/attention/backends/registry.py | 3 + .../v1/attention/backends/rocm_octave_attn.py | 496 +++++++ vllm/v1/attention/ops/rocm_octave.py | 390 ++++++ vllm/v1/kv_cache_interface.py | 15 + 30 files changed, 6193 insertions(+), 36 deletions(-) create mode 100644 csrc/attention/octave_arch.cuh create mode 100644 csrc/attention/octave_attn.cu create mode 100644 csrc/attention/octave_decode_mfma_cdna.cu create mode 100644 csrc/attention/octave_decode_wmma_rdna3.cu create mode 100644 csrc/attention/octave_format.cuh create mode 100644 csrc/attention/octave_prefill_cdna.cu create mode 100644 csrc/attention/octave_prefill_rdna3.cu create mode 100644 csrc/attention/octave_torch_bindings.cpp create mode 100644 docs/design/rocm_octave.md create mode 100644 tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3-compact.yaml create mode 100644 tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3.yaml create mode 100644 tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v4.yaml create mode 100644 tests/evals/gsm8k/configs/models-octave.txt create mode 100644 tests/kernels/attention/test_rocm_octave.py create mode 100644 vllm/v1/attention/backends/rocm_octave_attn.py create mode 100644 vllm/v1/attention/ops/rocm_octave.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 85383b36fbf..8ba09f55ae7 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3709,6 +3709,25 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt +- label: ":amd: (MI300) LM Eval Octave KV Cache" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - tests/evals/gsm8k/ + - tests/utils.py + - tests/conftest.py + - vllm/v1/attention/backends/rocm_octave_attn.py + - vllm/v1/attention/ops/rocm_octave.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-octave.txt + - label: ":amd: (MI355) LM Eval Spec Decode (Fixed-length FP8)" # TBD timeout_in_minutes: 145 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index a3298f2c5b3..3855cbcfdc1 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -1121,6 +1121,27 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt +# Octave kernels are built only for gfx942 and gfx1100, so this runs on MI300. +- label: ":amd: (MI300) LM Eval Octave KV Cache" + key: amd-lm-eval-octave-mi300 + dind: false + device: mi300_1 + num_devices: 1 + timeout_in_minutes: 45 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - tests/evals/gsm8k/ + - tests/utils.py + - tests/conftest.py + - vllm/v1/attention/backends/rocm_octave_attn.py + - vllm/v1/attention/ops/rocm_octave.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-octave.txt + - label: ":nvidia: (H100) [Nightly] GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-2xh100 timeout_in_minutes: 35 diff --git a/CMakeLists.txt b/CMakeLists.txt index bac0426cd21..7694e6f8bc4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -368,6 +368,20 @@ if(VLLM_GPU_LANG STREQUAL "HIP") set(VLLM_EXT_SRC "csrc/rocm/quickreduce/torch_bindings.cpp" "csrc/rocm/quickreduce/custom_quickreduce.cu") + # Octave KV cache (store, decode, rotation, prefill): only built for the + # validated architectures, gfx1100 (RDNA3) and gfx942 (MI300). + set(VLLM_OCTAVE_KV OFF) + if(VLLM_GPU_ARCHES MATCHES "gfx942|gfx1100") + set(VLLM_OCTAVE_KV ON) + list(APPEND VLLM_EXT_SRC + "csrc/attention/octave_torch_bindings.cpp" + "csrc/attention/octave_attn.cu" + "csrc/attention/octave_decode_wmma_rdna3.cu" + "csrc/attention/octave_prefill_rdna3.cu" + "csrc/attention/octave_prefill_cdna.cu" + "csrc/attention/octave_decode_mfma_cdna.cu" + ) + endif() message(STATUS "Enabling C extension.") define_extension_target( @@ -379,6 +393,13 @@ if(VLLM_GPU_LANG STREQUAL "HIP") ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) + if(VLLM_OCTAVE_KV) + # MFMA accumulators in VGPRs: in AGPRs every online-softmax rescale of the + # output copies it through 64 extra VGPRs and halves the occupancy. + set_source_files_properties( + "${CMAKE_CURRENT_BINARY_DIR}/csrc/attention/octave_decode_mfma_cdna.hip" + PROPERTIES COMPILE_FLAGS "-mllvm -amdgpu-mfma-vgpr-form") + endif() endif() # _C HIP endif if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 2242e231bb1..c42ebe895b3 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -32,6 +32,7 @@ import subprocess import sys from dataclasses import replace from pathlib import Path +from typing import get_args import yaml from rich.console import Console @@ -50,6 +51,7 @@ from common import ( is_mla_backend, ) +from vllm.config.cache import CacheDType from vllm.v1.worker.workspace import init_workspace_manager @@ -548,6 +550,11 @@ def main(): # Benchmark settings parser.add_argument("--device", default="cuda:0", help="Device") + parser.add_argument( + "--model", + default="meta-llama/Meta-Llama-3-8B", + help="Model whose config the backend is built with", + ) parser.add_argument( "--warmup-ms", type=int, @@ -561,8 +568,8 @@ def main(): parser.add_argument( "--kv-cache-dtype", default="auto", - choices=["auto", "fp8"], - 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, diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index b6fd26c8034..0af316d59a5 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -284,8 +284,10 @@ class BenchmarkConfig: torch_profile_iters: int = 3 warmup_ms: int | None = None - # "auto" or "fp8" + # Any CacheDType: "auto", "fp8", "int8_per_token_head", ... kv_cache_dtype: str = "auto" + # Model whose config the backend sees (dtype, rotary dims, ...) + model: str = "meta-llama/Meta-Llama-3-8B" # MLA-specific prefill_backend: str | None = None diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index a1b7d045050..abb8858abe6 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -36,6 +36,7 @@ from vllm.config import ( set_current_vllm_config, ) from vllm.platforms import current_platform +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.attention.backends.utils import ( CommonAttentionMetadata, get_supported_kv_cache_layouts, @@ -48,6 +49,7 @@ from vllm.v1.kv_cache_interface import ( compute_layer_kv_cache_shape_bytes, compute_layout_strides, create_kv_cache_views, + get_kv_quant_mode, ) # ============================================================================ @@ -142,15 +144,16 @@ def _build_common_attn_metadata( def _create_vllm_config( config: BenchmarkConfig, max_num_blocks: int, + max_model_len: int = 0, ) -> VllmConfig: """Create a VllmConfig for benchmarking with mock model methods.""" model_config = ModelConfig( - model="meta-llama/Meta-Llama-3-8B", - 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 diff --git a/csrc/attention/octave_arch.cuh b/csrc/attention/octave_arch.cuh new file mode 100644 index 00000000000..536f36b2594 --- /dev/null +++ b/csrc/attention/octave_arch.cuh @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Per-architecture primitives of the Octave kernels. Every AMD target vLLM +// builds compiles the portable tier; faster tiers are selected at runtime by +// the host code (see docs/design/rocm_octave.md, "Architecture contract"). +// +// int8 dot4: v_dot4_i32_iu8 on gfx11/gfx12, v_dot4_i32_i8 on gfx9/gfx10.3 +// decode, fast tier: octave_decode_wmma_rdna3.cu (WMMA, gfx11) +// prefill: octave_prefill_rdna3.cu (WMMA, gfx11) +// decode, portable: octave_attn.cu decode_kernel (dot4, every target) +// +// The portable kernels are written for 32-lane logical waves: shuffles never +// cross 32 lanes, so they are correct on wave64 (CDNA) as well. + +#pragma once + +#include + +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 diff --git a/csrc/attention/octave_attn.cu b/csrc/attention/octave_attn.cu new file mode 100644 index 00000000000..b4056a1662a --- /dev/null +++ b/csrc/attention/octave_attn.cu @@ -0,0 +1,1180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave KV cache kernels, head_size 256 with 64 RoPE dims. Portable across +// the AMD targets vLLM builds (octave_arch.cuh); the WMMA decode tier for +// gfx11 lives in octave_decode_wmma_rdna3.cu. +// Format: vllm/v1/attention/ops/rocm_octave.py (the PyTorch reference there +// is the source of truth for the byte layout). +// +// octave_cache_store quantize + pack K/V into slots (one wave per slot). +// octave_decode split-KV attention per query group; K dot products +// with int8 dot4 on the codebook values of the packed +// codes, the query quantized to int8 once per block. +// Consecutive query tokens of one request (MTP verify) +// share the KV read. +// octave_rotate the store-time rotations, for the prefill path. +// +// Everything is in the rotated space: K and V are sign-flipped and +// Hadamard-rotated at store time; the query is rotated like K and the +// attention output is rotated back in the reduce kernel. + +#include +#include +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #include "octave_arch.cuh" + #include "octave_format.cuh" + +namespace octave { + +constexpr int DPL = D / 32; // dims per lane when a wave owns one slot + +template +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(half x) { + return __half2float(x); +} +template +__device__ __forceinline__ T from_f(float x) { + return (T)x; +} +template <> +__device__ __forceinline__ half from_f(float x) { + return __float2half(x); +} + +__device__ __forceinline__ float half_bits_to_f(uint32_t bits16) { + return __half2float(__ushort_as_half((unsigned short)bits16)); +} +__device__ __forceinline__ uint16_t f_to_half_bits(float x) { + return __half_as_ushort(__float2half(x)); +} + +template +__device__ __forceinline__ float xor_sum(float v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v += __shfl_xor(v, m); + return v; +} +template +__device__ __forceinline__ float xor_max(float v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v = fmaxf(v, __shfl_xor(v, m)); + return v; +} +template +__device__ __forceinline__ uint32_t xor_or(uint32_t v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v |= (uint32_t)__shfl_xor((int)v, m); + return v; +} + +// Walsh-Hadamard butterflies over DPL consecutive dims per lane, then across +// lanes with xor distance 1..LANE_MASK. Unnormalized. +template +__device__ __forceinline__ void fwht_wave(float (&x)[DPL], int lane) { + #pragma unroll + for (int h = 1; h < DPL; h <<= 1) { + #pragma unroll + for (int i = 0; i < DPL; ++i) { + if ((i & h) == 0) { + float a = x[i], b = x[i + h]; + x[i] = a + b; + x[i + h] = a - b; + } + } + } + #pragma unroll + for (int m = 1; m <= LANE_MASK; m <<= 1) { + const bool hi = lane & m; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + float o = __shfl_xor(x[i], m); + x[i] = hi ? o - x[i] : x[i] + o; + } + } +} + +__device__ __forceinline__ float sign_of(const int* bits, int i) { + return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f; +} + +// Rotate a lane's 8 dims like K (64-dim blocks) or like V (one 256-dim +// block). +__device__ __forceinline__ void rotate_k(float (&x)[DPL], int lane, + const int* signs) { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= sign_of(signs, lane * DPL + i); + fwht_wave<4>(x, lane); + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= 0.125f; +} + +__device__ __forceinline__ void rotate_full(float (&x)[DPL], int lane, + const int* signs) { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= sign_of(signs, lane * DPL + i); + fwht_wave<16>(x, lane); + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= 0.0625f; +} + +// Midpoints of the int8 codebooks (rocm_octave.LUT), in codebook units. +__device__ constexpr float kThr3[7] = {-103.f, -62.f, -29.5f, 0.f, + 29.5f, 62.f, 103.f}; +__device__ constexpr float kThr4[15] = {-110.5f, -84.5f, -66.f, -50.5f, -36.5f, + -24.f, -12.f, 0.f, 12.f, 24.f, + 36.5f, 50.5f, 66.f, 84.5f, 110.5f}; + +// Nearest codebook entry on the group RMS (4- or 3-bit per lane), codes in +// [0, 2^bits). Returns the stored scale, times LUT_ONE: the one that makes +// x . x_hat = |x|^2, so quantization does not shrink scores or outputs. +template +__device__ __forceinline__ float quantize_group(const float (&x)[DPL], + int (&codes)[DPL], + float group_dims, bool four) { + float ss = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) ss += x[i] * x[i]; + ss = xor_sum(ss); + const float to_lut = + (four ? 46.0f : 59.0f) / fmaxf(sqrtf(ss / group_dims), 1e-12f); + float num = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + const float y = x[i] * to_lut; + int c = 0; + #pragma unroll + for (int k = 0; k < 15; ++k) + c += (four ? y > kThr4[k] : (k < 7 && y > kThr3[k])) ? 1 : 0; + codes[i] = c; + num += x[i] * (float)(int8_t)(four ? lut4((uint32_t)c) : lut3((uint32_t)c)); + } + num = xor_sum(num); + return num > 0.f ? ss / num * LUT_ONE : 0.f; +} + +// 4-bit: one word per 8 dims (lane-local). 3-bit: this lane's contribution to +// the 2-bit plane word (16 dims, 2 lanes) and the 1-bit plane word (32 dims, +// 4 lanes); callers OR across lanes. +__device__ __forceinline__ uint32_t pack4(const int (&c)[DPL]) { + uint32_t w = 0; + #pragma unroll + for (int i = 0; i < DPL; ++i) + w |= (uint32_t)c[i] << (8 * (i & 3) + 4 * (i >> 2)); + return w; +} +__device__ __forceinline__ void pack3(const int (&c)[DPL], int idx, + uint32_t& lo, uint32_t& hi) { + lo = 0; + hi = 0; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + const int e16 = (idx & 1) * DPL + i; // dim within the 16-dim lo word + const int e32 = (idx & 3) * DPL + i; // dim within the 32-dim hi word + lo |= (uint32_t)(c[i] & 3) << (8 * (e16 & 3) + 2 * (e16 >> 2)); + hi |= (uint32_t)(c[i] >> 2) << (8 * (e32 & 3) + (e32 >> 2)); + } +} + +// --------------------------------------------------------------------------- +// Store: grid (num_tokens, num_kv_heads), block 32. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(32) + store_kernel(const KT* __restrict__ key, const KT* __restrict__ value, + uint8_t* __restrict__ cache, + const int64_t* __restrict__ slot_mapping, + const int* __restrict__ k_signs, + const int* __restrict__ v_signs, int block_size, int64_t sk0, + int64_t sk1, int64_t sv0, int64_t sv1, int64_t scb, + int64_t sch, int64_t sct) { + using F = Format; + const int t = blockIdx.x; + const int h = blockIdx.y; + const int lane = threadIdx.x; + const int64_t slot = slot_mapping[t]; + if (slot < 0) return; + uint8_t* dst = + cache + (slot / block_size) * scb + h * sch + (slot % block_size) * sct; + + float k[DPL], v[DPL]; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + k[i] = to_f(key[t * sk0 + h * sk1 + lane * DPL + i]); + v[i] = to_f(value[t * sv0 + h * sv1 + lane * DPL + i]); + } + + // K, block layout: rotate, then quantize per 64-dim block (8 lanes): the + // RoPE block at 4 bits, the NoPE blocks at 3. Compact layout: rotate over + // all dims like V, 3 bits, one scale. + const bool rope_lane = !KC && lane < R / DPL; + int kc[DPL]; + float k_scale; + if constexpr (KC) { + rotate_full(k, lane, k_signs); + k_scale = quantize_group<16>(k, kc, (float)D, false); + } else { + rotate_k(k, lane, k_signs); + k_scale = quantize_group<4>(k, kc, 64.f, rope_lane); + } + const int nl = lane - F::KR / DPL; // 3-bit lane index (valid when >= 0) + { + uint32_t lo, hi; + pack3(kc, nl & 3, lo, hi); + lo = xor_or<1>(lo); + hi = xor_or<2>(hi); + if (rope_lane) { + *reinterpret_cast(dst + 4 * lane) = pack4(kc); + } else { + if ((nl & 1) == 0) + *reinterpret_cast(dst + F::OFF_KN + 4 * (nl >> 1)) = lo; + if ((nl & 3) == 0) + *reinterpret_cast(dst + F::OFF_KN + F::KN / 4 + + 4 * (nl >> 2)) = hi; + } + } + + // V: rotate all 256 dims, one scale. + rotate_full(v, lane, v_signs); + int vc[DPL]; + const float v_scale = quantize_group<16>(v, vc, (float)D, VB == 4); + if constexpr (VB == 4) { + *reinterpret_cast(dst + F::OFF_V + 4 * lane) = pack4(vc); + } else { + uint32_t lo, hi; + pack3(vc, lane, lo, hi); + lo = xor_or<1>(lo); + hi = xor_or<2>(hi); + if ((lane & 1) == 0) + *reinterpret_cast(dst + F::OFF_V + 4 * (lane >> 1)) = lo; + if ((lane & 3) == 0) + *reinterpret_cast(dst + F::OFF_V + D / 4 + 4 * (lane >> 2)) = + hi; + } + + uint16_t* sc = reinterpret_cast(dst + F::OFF_SC); + if (KC ? lane == 0 : (lane & 7) == 0) sc[lane >> 3] = f_to_half_bits(k_scale); + if (lane == 0) sc[F::NUM_KSC] = f_to_half_bits(v_scale); +} + +// --------------------------------------------------------------------------- +// Decode stage 1: grid (num_q, num_kv_heads, num_splits), block 128. +// A block serves up to QG consecutive query tokens of one request (only the +// first token of each group of a run launches work) and all G query heads of +// its KV head: ROWS = QG * G attention rows share every K/V load. +// --------------------------------------------------------------------------- +constexpr int WAVES = 4; +constexpr int CHUNK = 32 * WAVES; + +template +__global__ void __launch_bounds__(128) + decode_kernel(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_log2, + int num_q, int block_size, int max_blocks, int num_reqs, + int num_phys_blocks, int num_splits, int64_t sq0, int64_t sq1, + int64_t scb, int64_t sch, int64_t sct, int64_t smo, + int64_t smh, int64_t sms) { + using F = Format; + constexpr int ROWS = QG * G; + constexpr int QW = D / 4; // int8 query words per row + + __shared__ int q8[ROWS][QW]; + __shared__ float q_bs[ROWS][NKB]; // query scale per K block + __shared__ float S[ROWS][CHUNK]; + __shared__ float m_s[ROWS], l_s[ROWS], alpha_s[ROWS]; + __shared__ const uint8_t* tok_ptr[CHUNK]; + __shared__ float tok_vs[CHUNK]; + + const int qi = blockIdx.x; + const int kvh = blockIdx.y; + const int si = blockIdx.z; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int wave = tid >> 5; + + int req = q_to_req[qi]; + const bool req_ok = req >= 0 && req < num_reqs; + int nq = 1; + if constexpr (QG > 1) { + if (req_ok) { + int s = qi; + while (s > 0 && qi - s < 256 && q_to_req[s - 1] == req) --s; + if ((qi - s) % QG != 0) return; // served by the group leader + while (nq < QG && qi + nq < num_q && q_to_req[qi + nq] == req) ++nq; + } + } + const int max_kv = max_blocks * block_size; + int klen[QG]; + int kv_max = 0; + #pragma unroll + for (int j = 0; j < QG; ++j) { + int kl = (req_ok && j < nq) ? q_to_klen[qi + j] : 0; + kl = min(max(kl, 0), max_kv); + klen[j] = kl; + kv_max = max(kv_max, kl); + } + if (!req_ok) req = 0; + + const int tps = (kv_max + num_splits - 1) / num_splits; + const int start = si * tps; + const int end = min(start + tps, kv_max); + + auto write_empty = [&]() { + for (int r = wave; r < ROWS; r += WAVES) { + const int j = r / G, g = r % G; + if (j >= nq) continue; + float* op = mid_o + (qi + j) * smo + (kvh * G + g) * smh + si * sms; + for (int d = lane; d < D; d += 32) op[d] = 0.f; + if (lane == 0) { + op[D] = -INFINITY; + op[D + 1] = 0.f; + } + } + }; + if (start >= end) { + write_empty(); + return; + } + + // Query: rotate like K, quantize to int8 with one scale per 64-dim block. + for (int r = wave; r < ROWS; r += WAVES) { + const int j = r / G, g = r % G; + float x[DPL]; + if (j < nq) { + const QT* qp = Q + (qi + j) * sq0 + (kvh * G + g) * sq1 + lane * DPL; + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] = to_f(qp[i]); + } else { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] = 0.f; + } + if constexpr (KC) + rotate_full(x, lane, k_signs); + else + rotate_k(x, lane, k_signs); + float amax = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) amax = fmaxf(amax, fabsf(x[i])); + const float s = fmaxf(xor_max<4>(amax), 1e-20f) / 127.0f; + uint32_t w[2] = {0, 0}; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + int q = __float2int_rn(x[i] / s); + q = min(max(q, -127), 127); + w[i >> 2] |= (uint32_t)(q & 0xFF) << (8 * (i & 3)); + } + q8[r][2 * lane] = (int)w[0]; + q8[r][2 * lane + 1] = (int)w[1]; + if ((lane & 7) == 0) q_bs[r][lane >> 3] = s; + } + __syncthreads(); + + if (tid < ROWS) { + m_s[tid] = -INFINITY; + l_s[tid] = 0.f; + } + float o0[ROWS], o1[ROWS]; // this lane's two V dims (rotated space) + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + o0[r] = 0.f; + o1[r] = 0.f; + } + const int my_d = wave * 64 + 2 * lane; // first of the two V dims owned + + for (int base = start; base < end; base += CHUNK) { + // ---- scores: lane = token ---- + const int tok = base + wave * 32 + lane; + const bool valid = tok < end; + if (valid) { + const int lb = tok / block_size; + int pb = block_table[req * max_blocks + lb]; + if (pb < 0 || pb >= num_phys_blocks) pb = 0; + const uint8_t* p = + cache + pb * scb + kvh * sch + (int64_t)(tok - lb * block_size) * sct; + tok_ptr[wave * 32 + lane] = p; + + constexpr int KW = F::OFF_V / 4; // K code words + uint32_t kw[KW]; + const uint32_t* p4 = reinterpret_cast(p); + #pragma unroll + for (int w = 0; w < KW; ++w) kw[w] = p4[w]; + const uint32_t* kn = kw + F::OFF_KN / 4; + const uint32_t* ps = reinterpret_cast(p + F::OFF_SC); + float ks[NKB]; + if constexpr (KC) { + const uint32_t u = ps[0]; + #pragma unroll + for (int b = 0; b < NKB; ++b) ks[b] = half_bits_to_f(u & 0xFFFF); + tok_vs[wave * 32 + lane] = half_bits_to_f(u >> 16) * (1.0f / LUT_ONE); + } else { + const uint32_t u0 = ps[0], u1 = ps[1]; + ks[0] = half_bits_to_f(u0 & 0xFFFF); + ks[1] = half_bits_to_f(u0 >> 16); + ks[2] = half_bits_to_f(u1 & 0xFFFF); + ks[3] = half_bits_to_f(u1 >> 16); + tok_vs[wave * 32 + lane] = + half_bits_to_f(ps[2] & 0xFFFF) * (1.0f / LUT_ONE); + } + + // Rows two at a time: unrolling all of them kept too much live and + // spilled to scratch on gfx11. + #pragma unroll 2 + for (int r = 0; r < ROWS; ++r) { + const int* qr = q8[r]; + float sk = 0.f; + if constexpr (F::KR > 0) { + int dr = 0; + #pragma unroll + for (int u = 0; u < F::KR / 8; ++u) { + const uint32_t w = kw[u]; + dr = dot4_i8(qr[2 * u], (int)lut4(w & 0x0F0F0F0Fu), dr); + dr = dot4_i8(qr[2 * u + 1], (int)lut4((w >> 4) & 0x0F0F0F0Fu), dr); + } + sk = q_bs[r][0] * ks[0] * (float)dr; + } + #pragma unroll + for (int b = 0; b < F::KN / 64; ++b) { + constexpr int B0 = F::KR / 64; // first 3-bit block + int dn = 0; + #pragma unroll + for (int v = 0; v < 4; ++v) { // 2-bit plane word: 16 dims + const uint32_t lo = kn[4 * b + v]; + const uint32_t hi = kn[F::KN / 16 + 2 * b + (v >> 1)]; + #pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t c = + ((lo >> (2 * i)) & 0x03030303u) | + (((hi >> (4 * (v & 1) + i)) & 0x01010101u) << 2); + dn = + dot4_i8(qr[F::KR / 4 + 16 * b + 4 * v + i], (int)lut3(c), dn); + } + } + sk += q_bs[r][B0 + b] * ks[B0 + b] * (float)dn; + } + const float sc = sk * (sm_scale_log2 / LUT_ONE); + int kl = klen[0]; // klen by constant index only (no scratch) + #pragma unroll + for (int j = 1; j < QG; ++j) kl = r / G == j ? klen[j] : kl; + S[r][wave * 32 + lane] = tok < kl ? sc : -INFINITY; + } + } else { + #pragma unroll + for (int r = 0; r < ROWS; ++r) S[r][wave * 32 + lane] = -INFINITY; + } + __syncthreads(); + + // ---- online softmax: each wave owns rows r = wave, wave + WAVES, ... ---- + for (int r = wave; r < ROWS; r += WAVES) { + float x[CHUNK / 32]; + float mx = -INFINITY; + #pragma unroll + for (int c = 0; c < CHUNK / 32; ++c) { + x[c] = S[r][32 * c + lane]; + mx = fmaxf(mx, x[c]); + } + mx = xor_max<16>(mx); + const float m_old = m_s[r]; + const float m_new = fmaxf(m_old, mx); + float ps = 0.f; + #pragma unroll + for (int c = 0; c < CHUNK / 32; ++c) { + const float p = x[c] == -INFINITY ? 0.f : exp2f(x[c] - m_new); + S[r][32 * c + lane] = p; + ps += p; + } + ps = xor_sum<16>(ps); + if (lane == 0) { + const float a = m_new == -INFINITY ? 1.f : exp2f(m_old - m_new); + alpha_s[r] = a; + l_s[r] = l_s[r] * a + ps; + m_s[r] = m_new; + } + } + __syncthreads(); + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float a = alpha_s[r]; + o0[r] *= a; + o1[r] *= a; + } + + // ---- P @ V: wave owns 64 dims, lane owns 2 ---- + const int n_valid = min(CHUNK, end - base); + for (int t = 0; t < n_valid; ++t) { + const uint8_t* p = tok_ptr[t]; + uint32_t c; // codes of dims my_d, my_d + 1 in bytes 0, 1 + if constexpr (VB == 4) { + const uint32_t w = + *reinterpret_cast(p + F::OFF_V + 4 * (my_d >> 3)); + const int e = my_d & 7; // even + c = ((w >> (8 * (e & 3) + 4 * (e >> 2))) & 15) | + (((w >> (8 * ((e + 1) & 3) + 4 * ((e + 1) >> 2))) & 15) << 8); + } else { + const uint32_t lo = + *reinterpret_cast(p + F::OFF_V + 4 * (my_d >> 4)); + const uint32_t hi = *reinterpret_cast( + p + F::OFF_V + D / 4 + 4 * (my_d >> 5)); + const int e16 = my_d & 15, e32 = my_d & 31; + const int f16 = e16 + 1, f32 = e32 + 1; + c = (((lo >> (8 * (e16 & 3) + 2 * (e16 >> 2))) & 3) | + (((hi >> (8 * (e32 & 3) + (e32 >> 2))) & 1) << 2)) | + ((((lo >> (8 * (f16 & 3) + 2 * (f16 >> 2))) & 3) | + (((hi >> (8 * (f32 & 3) + (f32 >> 2))) & 1) << 2)) + << 8); + } + const uint32_t val = lut(c); + const float vs = tok_vs[t]; + const float c0 = (float)(int8_t)(val & 0xFF) * vs; + const float c1 = (float)(int8_t)((val >> 8) & 0xFF) * vs; + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float pr = S[r][t]; + o0[r] += pr * c0; + o1[r] += pr * c1; + } + } + __syncthreads(); + } + + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const int j = r / G, g = r % G; + if (j >= nq) continue; + float* op = mid_o + (qi + j) * smo + (kvh * G + g) * smh + si * sms; + op[my_d] = o0[r]; + op[my_d + 1] = o1[r]; + if (tid == 0) { + op[D] = m_s[r]; + op[D + 1] = l_s[r]; + } + } +} + +// --------------------------------------------------------------------------- +// Decode stage 2: combine splits, rotate V back. grid (num_q, num_q_heads), +// block 256. Up to 1024 splits: their weights at once, then 8 waves take the +// splits round-robin (lane = 8 dims) and skip the ones that weigh zero. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(256) + reduce_kernel(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) { + __shared__ float sw[1024], swm[8], swl[8]; + __shared__ float so[8][D]; + const int qi = blockIdx.x, hi = blockIdx.y; + const int tid = threadIdx.x, w = tid >> 5, lane = tid & 31; + const float* base = mid_o + qi * smo + hi * smh; + + float ms[4], ls[4]; + float m = -INFINITY; + #pragma unroll + for (int k = 0; k < 4; ++k) { + const int s = tid + 256 * k; + ms[k] = s < num_splits ? base[s * sms + D] : -INFINITY; + ls[k] = s < num_splits ? base[s * sms + D + 1] : 0.0f; + m = fmaxf(m, ms[k]); + } + m = xor_max<16>(m); + if (lane == 0) swm[w] = m; + __syncthreads(); + float mg = swm[0]; + #pragma unroll + for (int i = 1; i < 8; ++i) mg = fmaxf(mg, swm[i]); + float l = 0.0f; + #pragma unroll + for (int k = 0; k < 4; ++k) { + const float a = ms[k] == -INFINITY ? 0.0f : exp2f(ms[k] - mg); + if (tid + 256 * k < 1024) sw[tid + 256 * k] = a; + l += ls[k] * a; + } + l = xor_sum<16>(l); + if (lane == 0) swl[w] = l; + __syncthreads(); + + float o[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int s = w; s < num_splits; s += 8) { + const float as = sw[s]; + if (as == 0.0f) continue; + const float4* v = (const float4*)(base + s * sms + lane * 8); + const float4 x = v[0], y = v[1]; + o[0] += x.x * as; + o[1] += x.y * as; + o[2] += x.z * as; + o[3] += x.w * as; + o[4] += y.x * as; + o[5] += y.y * as; + o[6] += y.z * as; + o[7] += y.w * as; + } + #pragma unroll + for (int d = 0; d < 8; ++d) so[w][lane * 8 + d] = o[d]; + __syncthreads(); + float lg = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) lg += swl[i]; + float acc = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) acc += so[i][tid]; + __syncthreads(); + // Inverse rotation of V over the 256 dims (so[0] reused as scratch). + float* x = so[0]; + x[tid] = acc; + __syncthreads(); + for (int hh = 1; hh < D; hh <<= 1) { + float a = 0.f, b = 0.f; + const bool lo = (tid & hh) == 0; + if (lo) { + a = x[tid]; + b = x[tid + hh]; + } + __syncthreads(); + if (lo) { + x[tid] = a + b; + x[tid + hh] = a - b; + } + __syncthreads(); + } + const float inv_l = lg > 0.f ? 1.0f / lg : 0.f; + out[qi * soo + hi * soh + tid] = + from_f(x[tid] * 0.0625f * sign_of(v_signs, tid) * inv_l); +} + +// --------------------------------------------------------------------------- +// In-place rotation of [T, H, 256] vectors for the prefill path, like K +// (signs, then 64-dim Hadamard blocks) or like V (one 256-dim block); +// `inverse` applies the transform then the signs. grid (T, H), block 32. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(32) + rotate_kernel(T* __restrict__ x, const int* __restrict__ signs, + bool k_layout, bool inverse, int64_t s0, int64_t s1) { + const int lane = threadIdx.x; + T* p = x + blockIdx.x * s0 + blockIdx.y * s1 + lane * DPL; + float v[DPL]; + #pragma unroll + for (int i = 0; i < DPL; ++i) v[i] = to_f(p[i]); + if (!inverse) { + #pragma unroll + for (int i = 0; i < DPL; ++i) v[i] *= sign_of(signs, lane * DPL + i); + } + float scale; + if (k_layout) { + fwht_wave<4>(v, lane); + scale = 0.125f; + } else { + fwht_wave<16>(v, lane); + scale = 0.0625f; + } + #pragma unroll + for (int i = 0; i < DPL; ++i) { + float y = v[i] * scale; + if (inverse) y *= sign_of(signs, lane * DPL + i); + p[i] = from_f(y); + } +} + +// --------------------------------------------------------------------------- +// Exact window (sparse decode). Direct-mapped fp16 shadow tables keyed by the +// cache slot: entry slot & mask holds that token's exact K and V ([hkv, 2D]) +// and is valid while tags[entry] == slot. The recent table takes every stored +// token, the sink table the first kSinks positions of a sequence. A lost +// entry only means the token is read from the packed cache instead. +// --------------------------------------------------------------------------- +constexpr int kWindow = 64; +constexpr int kSinks = 4; +constexpr int kMaxWindowHits = 128; + +// One wave per token. The per-entry lock keeps two tokens that share an +// entry in the same launch from leaving a mixed row behind. +template +__global__ void __launch_bounds__(32) + window_store_kernel(const T* __restrict__ key, const T* __restrict__ value, + const int64_t* __restrict__ slot_mapping, + const int* __restrict__ positions, T* __restrict__ wtab, + int* __restrict__ wtags, int* __restrict__ wlocks, + int wmask, T* __restrict__ stab, + int* __restrict__ stags, int* __restrict__ slocks, + int smask, int hkv, int64_t sk0, int64_t sk1, + int64_t sv0, int64_t sv1) { + const int i = blockIdx.x, lane = threadIdx.x; + const int64_t slot = slot_mapping[i]; + if (slot < 0) return; + const int pos = positions[i]; + auto put = [&](T* tab, int* tags, int* locks, int e) { + if (lane == 0) { + while (atomicCAS(&locks[e], 0, 1) != 0) { + } + atomicExch(&tags[e], -1); + } + __threadfence(); + for (int h = 0; h < hkv; ++h) { + T* row = tab + ((int64_t)e * hkv + h) * (2 * D); + #pragma unroll + for (int d = 0; d < DPL; ++d) { + row[lane * DPL + d] = key[i * sk0 + h * sk1 + lane * DPL + d]; + row[D + lane * DPL + d] = value[i * sv0 + h * sv1 + lane * DPL + d]; + } + } + __threadfence(); + if (lane == 0) { + atomicExch(&tags[e], (int)slot); + __threadfence(); + atomicExch(&locks[e], 0); + } + }; + put(wtab, wtags, wlocks, (int)(slot & wmask)); + if (pos >= 0 && pos < kSinks) put(stab, stags, slocks, (int)(slot & smask)); +} + +// Attention of one query row over its selected positions that sit in the +// exact window, written as one more split of mid_o (in the rotated V space, +// like the packed-cache splits). grid (num_q, hkv), block 256. +template +__global__ void __launch_bounds__(256) window_attn_kernel( + const T* __restrict__ query, const int* __restrict__ positions, + const int* __restrict__ q_to_req, const int* __restrict__ block_table, + const int* __restrict__ indices, const T* __restrict__ wtab, + const int* __restrict__ wtags, int wmask, const T* __restrict__ stab, + const int* __restrict__ stags, int smask, float* __restrict__ mid_o, + const int* __restrict__ v_signs, float sm_scale, int split, int topk, + int64_t sidx, int block_size, int max_blocks, int num_reqs, + int num_phys_blocks, int hq, int hkv, int64_t sq0, int64_t sq1, int64_t smo, + int64_t smh, int64_t sms) { + __shared__ int s_entry[kMaxWindowHits]; + __shared__ float s_score[8][kMaxWindowHits]; + __shared__ int s_count; + const int qi = blockIdx.x, kvh = blockIdx.y; + const int tid = threadIdx.x, w = tid >> 5, lane = tid & 31; + const int hpk = hq / hkv; + if (tid == 0) s_count = 0; + __syncthreads(); + const int req = q_to_req[qi]; + const int p = positions[qi]; + if (req >= 0 && req < num_reqs) { + for (int j = tid; j < topk; j += 256) { + const int t = indices[qi * sidx + j]; + const int d = p - t; + if (t < 0 || d < 0 || (d >= kWindow && t >= kSinks)) continue; + const int lb = t / block_size; + if (lb >= max_blocks) continue; + const int pb = block_table[req * max_blocks + lb]; + if (pb < 0 || pb >= num_phys_blocks) continue; + const int slot = pb * block_size + (t - lb * block_size); + int e = -1; + if (d < kWindow && wtags[slot & wmask] == slot) + e = slot & wmask; + else if (t < kSinks && stags[slot & smask] == slot) + e = (slot & smask) | (1 << 30); + if (e < 0) continue; + const int k = atomicAdd(&s_count, 1); + if (k < kMaxWindowHits) s_entry[k] = e; + } + } + __syncthreads(); + const int n = min(s_count, kMaxWindowHits); + const float sl2 = sm_scale * 1.4426950408889634f; + for (int g = w; g < hpk; g += 8) { + const int head = kvh * hpk + g; + float q[DPL]; + #pragma unroll + for (int d = 0; d < DPL; ++d) + q[d] = to_f(query[qi * sq0 + head * sq1 + lane * DPL + d]); + auto row_of = [&](int e) { + const T* tab = (e >> 30) ? stab : wtab; + return tab + ((int64_t)(e & ((1 << 30) - 1)) * hkv + kvh) * (2 * D); + }; + float m = -INFINITY; + for (int i = 0; i < n; ++i) { + const T* r = row_of(s_entry[i]); + float acc = 0.f; + #pragma unroll + for (int d = 0; d < DPL; ++d) acc += q[d] * to_f(r[lane * DPL + d]); + acc = xor_sum<16>(acc) * sl2; + if (lane == 0) s_score[w][i] = acc; + m = fmaxf(m, acc); + } + __builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront"); + __builtin_amdgcn_wave_barrier(); + __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront"); + float l = 0.f, o[DPL] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int i = 0; i < n; ++i) { + const float pr = exp2f(s_score[w][i] - m); + l += pr; + const T* r = row_of(s_entry[i]) + D; + #pragma unroll + for (int d = 0; d < DPL; ++d) o[d] += pr * to_f(r[lane * DPL + d]); + } + rotate_full(o, lane, v_signs); + float* op = mid_o + qi * smo + head * smh + (int64_t)split * sms; + #pragma unroll + for (int d = 0; d < DPL; ++d) op[lane * DPL + d] = o[d]; + if (lane == 0) { + op[D] = n > 0 ? m : -INFINITY; + op[D + 1] = l; + } + } +} + +} // namespace octave + +// --------------------------------------------------------------------------- +// Host launchers +// --------------------------------------------------------------------------- +using namespace octave; + +// fmt: V fmt (3 or 4), plus kCompactFlag for compact K (3-bit V only). +int octave_slot_bytes(int64_t fmt) { + switch (fmt) { + case 3: + return Format<3>::SLOT; + case 4: + return Format<4>::SLOT; + case 3 + kCompactFlag: + return Format<3, true>::SLOT; + } + TORCH_CHECK(false, "octave: unknown format id ", fmt); + return 0; +} + +static void check_cache(const torch::Tensor& cache, int64_t fmt) { + TORCH_CHECK(cache.dtype() == at::kByte && cache.dim() == 4, + "octave: cache must be uint8 [blocks, heads, block_size, slot]"); + const int slot = octave_slot_bytes(fmt); + TORCH_CHECK(cache.size(3) == slot && cache.stride(3) == 1, + "octave: slot size mismatch, expected ", slot, " got ", + cache.size(3)); +} + +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) { + const int n = slot_mapping.size(0); + if (n == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(key.size(-1) == D && key.stride(-1) == 1 && value.stride(-1) == 1, + "octave: head_size must be 256 with contiguous last dim"); + TORCH_CHECK(slot_mapping.dtype() == at::kLong); + const int hkv = cache.size(1); + const at::cuda::OptionalCUDAGuard guard(device_of(key)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + dim3 grid(n, hkv); + #define SQ_STORE(B, KC, T) \ + store_kernel<<>>( \ + (const T*)key.data_ptr(), (const T*)value.data_ptr(), \ + (uint8_t*)cache.data_ptr(), slot_mapping.data_ptr(), \ + k_signs.data_ptr(), v_signs.data_ptr(), cache.size(2), \ + key.stride(0), key.stride(1), value.stride(0), value.stride(1), \ + cache.stride(0), cache.stride(1), cache.stride(2)) + const bool bf = key.dtype() == at::kBFloat16; + if (fmt == 4) { + if (bf) + SQ_STORE(4, false, __hip_bfloat16); + else + SQ_STORE(4, false, half); + } else if (fmt == 3) { + if (bf) + SQ_STORE(3, false, __hip_bfloat16); + else + SQ_STORE(3, false, half); + } else { + if (bf) + SQ_STORE(3, true, __hip_bfloat16); + else + SQ_STORE(3, true, half); + } + #undef SQ_STORE +} + +template +static void launch_decode(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, float sm_scale_log2, int ns, + hipStream_t stream) { + const int num_q = query.size(0); + dim3 grid(num_q, cache.size(1), ns); + decode_kernel<<>>( + (const QT*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), + block_table.data_ptr(), q_to_req.data_ptr(), + q_to_klen.data_ptr(), mid_o.data_ptr(), + k_signs.data_ptr(), sm_scale_log2, num_q, cache.size(2), + block_table.size(1), block_table.size(0), cache.size(0), ns, + 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)); +} + +template +static void dispatch_decode(int g, int qg, 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, + float sl2, int ns, hipStream_t stream) { + #define SQ_DEC(GG, QQ) \ + if (g == GG && qg == QQ) { \ + launch_decode(query, cache, block_table, q_to_req, \ + q_to_klen, mid_o, k_signs, sl2, ns, \ + stream); \ + return; \ + } + SQ_DEC(1, 1); + SQ_DEC(1, 4); + SQ_DEC(2, 1); + SQ_DEC(2, 4); + SQ_DEC(4, 1); + SQ_DEC(4, 4); + SQ_DEC(6, 1); + SQ_DEC(6, 4); + SQ_DEC(8, 1); + SQ_DEC(8, 4); + #undef SQ_DEC + TORCH_CHECK(false, "octave_decode: unsupported GQA group ", g); +} + +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& indices, + const int* qpos, const int* wtags, int wmask, + const int* stags, int smask, int64_t query_group); + +bool octave_reduce_mfma(torch::Tensor out, torch::Tensor mid_o, + torch::Tensor v_signs, int num_splits, + cudaStream_t stream); + +static void launch_reduce(torch::Tensor out, torch::Tensor mid_o, + torch::Tensor v_signs, int num_splits, + cudaStream_t stream) { + static const bool cdna3 = + std::string(at::cuda::getCurrentDeviceProperties()->gcnArchName) + .rfind("gfx942", 0) == 0; + if (cdna3 && octave_reduce_mfma(out, mid_o, v_signs, num_splits, stream)) + return; + dim3 grid(out.size(0), out.size(1)); + if (out.dtype() == at::kBFloat16) + reduce_kernel<__hip_bfloat16><<>>( + mid_o.data_ptr(), (__hip_bfloat16*)out.data_ptr(), + v_signs.data_ptr(), num_splits, mid_o.stride(0), mid_o.stride(1), + mid_o.stride(2), out.stride(0), out.stride(1)); + else + reduce_kernel<<>>( + mid_o.data_ptr(), (half*)out.data_ptr(), v_signs.data_ptr(), + num_splits, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), + out.stride(0), out.stride(1)); +} + +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) { + const int num_q = query.size(0); + if (num_q == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(query.size(2) == D && query.stride(2) == 1); + TORCH_CHECK(out.dtype() == query.dtype()); + TORCH_CHECK(mid_o.size(2) >= num_kv_splits && mid_o.size(3) >= D + 2); + TORCH_CHECK(num_kv_splits <= 1024 && mid_o.stride(3) == 1); + const int hq = query.size(1), hkv = cache.size(1); + TORCH_CHECK(hq % hkv == 0); + const int g = hq / hkv; + const int qg = query_group > 1 ? 4 : 1; + int ns = (int)num_kv_splits; + const float sl2 = (float)sm_scale * 1.4426950408889634f; + const at::cuda::OptionalCUDAGuard guard(device_of(query)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + const bool bf = query.dtype() == at::kBFloat16; + // Fast tier when the device and shapes allow it, else the portable kernel. + const int wmma_splits = + use_wmma ? octave_decode_wmma(query, cache, block_table, q_to_req, + q_to_klen, mid_o, k_signs, sm_scale, + num_kv_splits, fmt, std::nullopt, nullptr, + nullptr, 0, nullptr, 0, qg) + : 0; + if (wmma_splits > 0) { + ns = wmma_splits; + } else { + #define SQ_DISPATCH(B, KC) \ + if (bf) \ + dispatch_decode(g, qg, query, cache, block_table, \ + q_to_req, q_to_klen, mid_o, \ + k_signs, sl2, ns, stream); \ + else \ + dispatch_decode(g, qg, query, cache, block_table, q_to_req, \ + q_to_klen, mid_o, k_signs, sl2, ns, stream) + if (fmt == 4) { + SQ_DISPATCH(4, false); + } else if (fmt == 3) { + SQ_DISPATCH(3, false); + } else { + SQ_DISPATCH(3, true); + } + #undef SQ_DISPATCH + } + launch_reduce(out, mid_o, v_signs, ns, stream); +} + +// Sparse decode: query row i attends to the logical positions indices[i, :] +// of request q_to_req[i] (negative entries are skipped). WMMA tier only. +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& positions, + const std::optional& wtab, + const std::optional& wtags, + const std::optional& stab, + const std::optional& stags) { + const int num_q = query.size(0); + if (num_q == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(query.size(2) == D && query.stride(2) == 1); + TORCH_CHECK(out.dtype() == query.dtype()); + TORCH_CHECK(indices.dim() == 2 && indices.size(0) >= num_q && + indices.stride(1) == 1 && indices.dtype() == at::kInt); + TORCH_CHECK(mid_o.size(2) >= num_kv_splits && mid_o.size(3) >= D + 2); + TORCH_CHECK(num_kv_splits <= 1024 && mid_o.stride(3) == 1); + const at::cuda::OptionalCUDAGuard guard(device_of(query)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + // With the exact window, the packed-cache splits leave the last split of + // mid_o to octave_window_attn. + const bool win = positions.has_value(); + TORCH_CHECK(!win || (wtab && wtags && stab && stags && num_kv_splits >= 2)); + const int wmask = win ? (int)wtags->size(0) - 1 : 0; + const int smask = win ? (int)stags->size(0) - 1 : 0; + TORCH_CHECK( + !win || (((wmask + 1) & wmask) == 0 && ((smask + 1) & smask) == 0), + "octave_decode_sparse: shadow tables must be powers of two"); + const int ns = octave_decode_wmma( + query, cache, block_table, q_to_req, q_to_req, mid_o, k_signs, sm_scale, + win ? num_kv_splits - 1 : num_kv_splits, fmt, indices, + win ? positions->data_ptr() : nullptr, + win ? wtags->data_ptr() : nullptr, wmask, + win ? stags->data_ptr() : nullptr, smask, 4); + TORCH_CHECK(ns > 0, + "octave_decode_sparse: needs the WMMA tier (gfx1100, head size " + "256, GQA group <= 8)"); + if (!win) { + launch_reduce(out, mid_o, v_signs, ns, stream); + return; + } + const int hq = query.size(1), hkv = cache.size(1); + TORCH_CHECK(hq / hkv <= 8 && wtab->size(1) == hkv && wtab->size(2) == 2 * D); + dim3 grid(num_q, hkv); + #define SQ_WIN(T) \ + window_attn_kernel<<>>( \ + (const T*)query.data_ptr(), positions->data_ptr(), \ + q_to_req.data_ptr(), block_table.data_ptr(), \ + indices.data_ptr(), (const T*)wtab->data_ptr(), \ + wtags->data_ptr(), wmask, (const T*)stab->data_ptr(), \ + stags->data_ptr(), smask, mid_o.data_ptr(), \ + v_signs.data_ptr(), (float)sm_scale, ns, (int)indices.size(1), \ + indices.stride(0), cache.size(2), block_table.size(1), \ + block_table.size(0), cache.size(0), hq, hkv, query.stride(0), \ + query.stride(1), mid_o.stride(0), mid_o.stride(1), mid_o.stride(2)) + if (query.dtype() == at::kBFloat16) + SQ_WIN(__hip_bfloat16); + else + SQ_WIN(half); + #undef SQ_WIN + launch_reduce(out, mid_o, v_signs, ns + 1, stream); +} + +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) { + const int n = slot_mapping.size(0); + if (n == 0) return; + TORCH_CHECK(slot_mapping.dtype() == at::kLong && + positions.dtype() == at::kInt); + TORCH_CHECK(wtab.dtype() == key.dtype() && stab.dtype() == key.dtype()); + const int hkv = wtab.size(1); + const int wmask = (int)wtags.size(0) - 1, smask = (int)stags.size(0) - 1; + const at::cuda::OptionalCUDAGuard guard(device_of(key)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + #define SQ_WSTORE(T) \ + window_store_kernel<<>>( \ + (const T*)key.data_ptr(), (const T*)value.data_ptr(), \ + slot_mapping.data_ptr(), positions.data_ptr(), \ + (T*)wtab.data_ptr(), wtags.data_ptr(), wlocks.data_ptr(), \ + wmask, (T*)stab.data_ptr(), stags.data_ptr(), \ + slocks.data_ptr(), smask, hkv, key.stride(0), key.stride(1), \ + value.stride(0), value.stride(1)) + if (key.dtype() == at::kBFloat16) + SQ_WSTORE(__hip_bfloat16); + else + SQ_WSTORE(half); + #undef SQ_WSTORE +} + +void octave_rotate(torch::Tensor x, torch::Tensor signs, bool k_layout, + bool inverse) { + if (x.numel() == 0) return; + TORCH_CHECK(x.dim() == 3 && x.size(2) == D && x.stride(2) == 1, + "octave_rotate: x must be [T, H, 256] with contiguous last dim"); + const at::cuda::OptionalCUDAGuard guard(device_of(x)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + dim3 grid(x.size(0), x.size(1)); + if (x.dtype() == at::kBFloat16) + rotate_kernel<__hip_bfloat16><<>>( + (__hip_bfloat16*)x.data_ptr(), signs.data_ptr(), k_layout, inverse, + x.stride(0), x.stride(1)); + else + rotate_kernel + <<>>((half*)x.data_ptr(), signs.data_ptr(), + k_layout, inverse, x.stride(0), x.stride(1)); +} + +#else +void octave_cache_store(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, int64_t) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_decode(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, double, int64_t, int64_t, int64_t, bool) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_decode_sparse(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, double, + int64_t, int64_t, const std::optional&, + const std::optional&, + const std::optional&, + const std::optional&, + const std::optional&) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_window_store(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_rotate(torch::Tensor, torch::Tensor, bool, bool) { + TORCH_CHECK(false, "octave requires ROCm"); +} +#endif diff --git a/csrc/attention/octave_decode_mfma_cdna.cu b/csrc/attention/octave_decode_mfma_cdna.cu new file mode 100644 index 00000000000..07f899600c6 --- /dev/null +++ b/csrc/attention/octave_decode_mfma_cdna.cu @@ -0,0 +1,719 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave decode on MFMA for CDNA3 (gfx942). Same scheme as the WMMA decode +// (octave_decode_wmma_rdna3.cu): independent waves, one (split, 16-query row +// tile) per wave, transposed orientation, codebook bytes fed to fp16 as +// x + 1152 and debiased afterwards. With 16x16x16 MFMA on wave64 (lane l +// holds column l % 16 and rows 4 (l / 16) .. + 3 of the result): +// +// QK S^T[token][q] = K x Q^T. A lane owns one token (l % 16) and, in the +// four k-steps of each 64-dim block, the 16 contiguous dims +// 16 (l / 16) .. + 15 of that block, so its codes are one contiguous +// run of the slot. One accumulator per K scale. +// PV O^T[dim][q] = V^T x P^T. Output row r of fragment c is dim 16 r + c, +// so a lane reads the 16 contiguous dims 16 (l % 16) .. + 15 of tokens +// 4 (l / 16) .. + 3. P^T is the S^T accumulator, one query per lane. + +#include +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #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 +__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 +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(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 +__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; + 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(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(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(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 +__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><<>>( + mid_o.data_ptr(), (__hip_bfloat16*)out.data_ptr(), + v_signs.data_ptr(), num_splits, mid_o.stride(0), mid_o.stride(1), + mid_o.stride(2), out.stride(0), out.stride(1)); + else + reduce_mfma<<>>( + mid_o.data_ptr(), (half*)out.data_ptr(), v_signs.data_ptr(), + 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(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({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<<>>( \ + (const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), \ + block_table.data_ptr(), q_to_req.data_ptr(), \ + q_to_klen.data_ptr(), mid_o.data_ptr(), \ + k_signs.data_ptr(), (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 diff --git a/csrc/attention/octave_decode_wmma_rdna3.cu b/csrc/attention/octave_decode_wmma_rdna3.cu new file mode 100644 index 00000000000..6f4c6537a08 --- /dev/null +++ b/csrc/attention/octave_decode_wmma_rdna3.cu @@ -0,0 +1,687 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave decode on WMMA for RDNA3 (gfx11). Independent waves, one (split, +// 16-query row tile) per wave, transposed orientation, and lane pairs sharing +// A fragments through a DPP swap. The products: +// +// QK S^T[token][q] = K x Q^T over 16 k-steps of 16 dims: steps 0-3 the +// RoPE block (4-bit codes), 4-15 the NoPE blocks (3-bit codes). Codes +// are mapped to their int8 codebook values and fed as x + 1152. Each +// 64-dim block has its own per-token scale, so the steps accumulate +// into four tiles that are scaled and debiased in the softmax. Compact +// K: 16 steps of 3-bit codes and one scale, so one tile. +// PV O^T[dim][q] = V^T x P^T with the V codebook values fed as x + 1152; +// the per-token V scale is folded into P, the bias removed with sum(P). +// +// The query is rotated like K while it is staged into LDS; the output stays +// rotated and octave_reduce rotates it back. + +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #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 +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(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 +__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; + 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(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& 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(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() : 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 \ + <<>>( \ + (const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), \ + block_table.data_ptr(), q_to_req.data_ptr(), \ + q_to_klen.data_ptr(), mid_o.data_ptr(), \ + k_signs.data_ptr(), (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 diff --git a/csrc/attention/octave_format.cuh b/csrc/attention/octave_format.cuh new file mode 100644 index 00000000000..faf8a884864 --- /dev/null +++ b/csrc/attention/octave_format.cuh @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Byte layout of one Octave slot (one token, one KV head). The PyTorch +// reference in vllm/v1/attention/ops/rocm_octave.py is the source of truth. +// +// [0, OFF_KN) K RoPE block (dims 0-63), rotated, 4-bit codes +// [OFF_KN, OFF_V) K NoPE blocks (dims 64-255), rotated, 3-bit codes +// [OFF_V, OFF_SC) V, rotated, VBITS-bit codes +// [OFF_SC, ...) fp16 scales: K block 0..3, V +// +// Compact K (KC): K rotated over all 256 dims like V, 3-bit codes for every +// dim (OFF_KN = 0), one K scale; slots pad to 4 bytes instead of 8. +// +// 4-bit codes: dims 8i..8i+7 in one dword, low nibbles = dims 0-3, high +// nibbles = dims 4-7. 3-bit codes: a 2-bit plane (16 dims per dword, dims +// 4k..4k+3 at bit 2k of bytes 0-3) followed by a 1-bit plane (32 dims per +// dword, dims 4k..4k+3 at bit k of bytes 0-3). A code c of a block with +// scale s stands for LUT[c] * s / LUT_ONE, LUT being the int8 Lloyd-Max +// codebook of its width. + +#pragma once + +#include + +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 +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 +__device__ __forceinline__ uint32_t lut(uint32_t codes) { + if constexpr (BITS == 4) + return lut4(codes); + else + return lut3(codes); +} + +} // namespace octave diff --git a/csrc/attention/octave_prefill_cdna.cu b/csrc/attention/octave_prefill_cdna.cu new file mode 100644 index 00000000000..d1638219b79 --- /dev/null +++ b/csrc/attention/octave_prefill_cdna.cu @@ -0,0 +1,666 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave prefill for CDNA3 (gfx942), head size 256, fp16. Same contract as +// octave_prefill_rdna3.cu: reads the packed cache directly, everything in the +// rotated space, phase 1 writes (O, m, l) partials over the cached prefix and +// phase 2 runs causal attention over the chunk and merges. +// +// Both phases compute transposed tiles so that every MFMA result is already in +// the layout the next MFMA consumes (16x16 MFMA, wave64: lane l holds column +// l % 16 and rows 4 (l / 16) .. 4 (l / 16) + 3): +// S^T[token][q] = K x Q^T (phase 1: int8 16x16x32, phase 2: fp16) +// O^T[dim][q] += V^T x P^T (fp16 16x16x16) +// so P^T comes straight out of the S^T accumulator, one query per lane. +// +// Phase 1: one block = 128 query vectors (token, head) of one KV head, 8 waves +// of 16 query vectors x 256 dims. 256 threads unpack the next 16-token tile +// into LDS while the current one is consumed. + +#include +#include + +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + + #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 +__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 +__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; + 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(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<<>>( \ + (const __half*)q.data_ptr(), cache.data_ptr(), \ + block_table.data_ptr(), cu_seqlens_q.data_ptr(), \ + seq_lens.data_ptr(), 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<<>>( + (__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(), seq_lens.data_ptr(), 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 diff --git a/csrc/attention/octave_prefill_rdna3.cu b/csrc/attention/octave_prefill_rdna3.cu new file mode 100644 index 00000000000..3566bf62626 --- /dev/null +++ b/csrc/attention/octave_prefill_rdna3.cu @@ -0,0 +1,809 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave prefill for RDNA3 (gfx11), head size 256, fp16. Reads the packed +// cache directly; no intermediate buffers. Everything is in the rotated space +// (the caller rotates Q/K/V of the chunk and rotates the output back). +// +// Phase 1 (prefix_attn): attention of the chunk's queries over the cached +// prefix, split over the prefix length; writes (O, m, l) partials. +// One block = 128 query vectors (token, head) of one KV head, so the GQA +// heads share every K/V tile. 16 compute waves in pairs: wave g of a pair +// holds Q (int8) and O for head dims [128 g, 128 g + 128). 8 loader waves +// unpack 16-token tiles into double-buffered LDS, one tile ahead. +// QK S^T[token][q] = K x Q^T with int8 WMMA. K rows hold the int8 codebook +// value of every dim. Each wave's 128 dims are two 64-dim blocks with +// their own per-token scale; the wave scales both and the pair +// exchanges the float partial through LDS. +// PV O^T[dim][q] = V^T x P^T with V as its exact fp16 codebook values; the +// V scale is folded into P. +// Phase 2 (chunk_attn): causal attention over the chunk's own K/V (fp16), +// merged with the phase-1 partials; writes the output. + +#include +#include +#include +#include +#include + +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + + #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 +__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 +__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; + 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(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> kept; + std::lock_guard lock(mtx); + std::vector& 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<<>>( \ + (const __half*)q.data_ptr(), cache.data_ptr(), \ + block_table.data_ptr(), cu_seqlens_q.data_ptr(), \ + seq_lens.data_ptr(), 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<<>>( + (__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(), seq_lens.data_ptr(), 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 diff --git a/csrc/attention/octave_torch_bindings.cpp b/csrc/attention/octave_torch_bindings.cpp new file mode 100644 index 00000000000..9446606df85 --- /dev/null +++ b/csrc/attention/octave_torch_bindings.cpp @@ -0,0 +1,82 @@ +// Op registration for the Octave KV cache kernels (legacy _C extension, full +// torch API). Built only for gfx1100 and gfx942; see CMakeLists.txt. +#include + +#include +#include + +#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& positions, + const std::optional& wtab, + const std::optional& wtags, + const std::optional& stab, + const std::optional& 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); +} diff --git a/docs/design/rocm_octave.md b/docs/design/rocm_octave.md new file mode 100644 index 00000000000..3248419608c --- /dev/null +++ b/docs/design/rocm_octave.md @@ -0,0 +1,197 @@ +# Octave KV: a 3-bit KV cache for ROCm + +`ROCM_OCTAVE` is an attention backend for AMD GPUs that keeps the KV cache of +full-attention layers at 3-4 bits per value, with native kernels (HIP today). It is +selected with `--kv-cache-dtype octave_k3v3_compact`, `octave_k3v3` or +`octave_k3v4` (the backend is picked from the cache dtype; +`--attention-backend ROCM_OCTAVE` also works). The name comes from the 3-bit +codes: each one indexes an eight-level codebook. + +## Why it exists + +Long-context serving on AMD GPUs is bounded by KV-cache memory long before it +is bounded by compute. The existing options on ROCm were: + +- fp16/bf16 KV (`TRITON_ATTN`, `ROCM_ATTN`): 16 bits per value. +- fp8 / int8 per-token-head: 8 bits per value. +- `TURBOQUANT`: 3-4 bits per value, but its kernels are Triton (store and + decode) and its prefill goes through FlashAttention; it does not support the + fused multi-step MTP draft. + +Octave KV targets the same bit budget as TurboQuant with HIP kernels written for +the AMD matrix instructions (WMMA on RDNA3, MFMA on CDNA as a contract slot), +reading the packed cache directly in both decode and prefill, and with +multi-token verification (MTP / speculative decoding) as a first-class case. + +## Supported models + +The format exploits *partial rotary embeddings*: models whose attention heads +apply RoPE to a fraction of the head dims. Today the kernels are built for +head size 256 with 64 rotary dims (`partial_rotary_factor = 0.25`), the +attention shape of the Qwen3-Next / Qwen3.5 / Qwen3.6 / Qwen3.8 family (dense +and MoE, e.g. Qwen3.8-27B, Qwen3.6-35B-A3B, Qwen3.5-122B-A10B). Like the MLA +backends, which serve only models with multi-head latent attention, Octave KV is +tied to an attention shape, not to a model. + +Models with full RoPE (Llama, Qwen3, Mistral, ...) or other head sizes are +rejected at startup with an explicit error. Sliding window, ALiBi, attention +sinks, soft-capping and encoder attention are not supported. + +The decode kernels cover GQA groups of 1, 2, 4, 6 or 8 query heads per KV +head on each GPU. Under tensor parallelism the group is counted per GPU, so +models with 32 query heads and 2 KV heads (Qwen3.5-122B-A10B) need TP ≥ 4. + +## Format + +One slot per (token, KV head). V is always rotated as one 256-dim block +(random signs, then a Walsh-Hadamard transform) and gets 3 or 4 bits per dim. +K has two layouts: + +- **Block K** (`octave_k3v3`, `octave_k3v4`): four 64-dim blocks, each rotated + on its own. The first block holds exactly the 64 rotary dims and gets 4 bits + per dim, the three NoPE blocks 3 bits; each block has its own scale. +- **Compact K** (`octave_k3v3_compact`): rotated over all 256 dims like V, + 3 bits per dim, one scale. The rotation spreads the rotary dims over the + whole head, so they need no extra bits; slots pad to 4 bytes instead of 8. + +| Bytes (block K) | Content | +| --- | --- | +| `[0, 32)` | K RoPE block, 4-bit codes | +| `[32, 104)` | K NoPE blocks, 3-bit codes | +| next `256·v/8` | V, `v`-bit codes | +| next 10 | fp16 scales (K blocks 0-3, V) | + +| Bytes (compact K) | Content | +| --- | --- | +| `[0, 96)` | K, 3-bit codes | +| `[96, 192)` | V, 3-bit codes | +| `[192, 196)` | fp16 scales (K, V) | + +| Cache dtype | Slot bytes | Bits per value | +| --- | --- | --- | +| `octave_k3v3_compact` | 196 | 3.06 | +| `octave_k3v3` | 216 | 3.38 | +| `octave_k3v4` | 248 | 3.88 | + +For reference, `turboquant_3bit_nc` uses 198 bytes and `turboquant_k3v4_nc` +230. + +After the rotation the coordinates are close to Gaussian, so each code +indexes the Lloyd-Max codebook of N(0, 1) for its width, stored as int8 +(`LUT` in the reference): a code `c` of a block with scale `s` stands for +`LUT[c] · s / 64`. The scale is chosen so that `x · x̂ = |x|²`. A +least-squares fit minimizes the reconstruction error but shrinks every +reconstructed vector (by about 3% at 3 bits), which flattens the softmax and +scales down the attention output; at 380k tokens it cost 2-3× more NLL than +this scale, at the same size. The quantizer is dynamic: no calibration, no per-model +dictionaries. The kernels map four codes to their int8 values with one +`v_perm_b32` (two for 4-bit codes), so the codebook costs nothing in +bandwidth and the dot products stay integer. The byte layout's source of +truth is the PyTorch reference in `vllm/v1/attention/ops/rocm_octave.py`. + +Why these layouts: measured on real Q/K/V of Qwen3.8-27B (32k tokens, every +full-attention layer, attention-output error), both K layouts sit on the same +bytes-versus-error curve, about 25-30% below TurboQuant at equal size; block +K reaches the 4-bit rotary block, compact K the smallest slot. Rotating the +rotary block alone at 3 bits does not work (one layer reaches 32% error); +rotating it with the rest of the head does. Moving bits between K and V or +weighting K channels by query energy did not help at a fixed size. The int8 +codebook loses nothing against the floating-point one. + +Attention runs in the rotated space: the query is rotated like K and the +output is rotated back once per query. The rotations are orthogonal, so the +scores are those of the unrotated vectors. + +## Kernels and data flow + +| Step | Kernel | File | +| --- | --- | --- | +| Cache write | `octave_cache_store`: one wave per slot, rotate + quantize + pack | `csrc/attention/octave_attn.cu` | +| Decode / MTP verify (≤128 query tokens per request) | `octave_decode`: split-KV; WMMA tier on gfx11, portable dot4 tier elsewhere; consecutive query tokens of a request share each K/V read | `octave_decode_wmma_rdna3.cu`, `octave_attn.cu` | +| Split reduce + inverse V rotation | `reduce_kernel` | `octave_attn.cu` | +| Prefill | `octave_prefill`: phase 1 reads the cached prefix from the packed cache (int8 WMMA for QK with per-segment scales, fp16 WMMA for PV), phase 2 attends to the chunk's own unquantized K/V and merges | `octave_prefill_rdna3.cu` | +| Q/K/V chunk rotation for prefill | `octave_rotate` | `octave_attn.cu` | +| Codebooks (`v_perm_b32` tables) and byte layout | — | `octave_format.cuh` | + +The metadata builder builds, on the device and without host syncs, the +per-query request and causal-length maps the kernels index by. They are what +makes MTP verification batches (several query tokens per request, each with +its own causal length) work without special cases, and they give cudagraph +padding queries a K length of 0. + +## Architecture contract + +Kernels are native: HIP or FlyDSL, never Triton. There is no fallback: an +architecture without a required kernel is rejected at startup with an error +naming what is missing. The kernels in this tree are HIP; a FlyDSL kernel +(for example a CDNA prefill) must keep the same signature and semantics, pass +the same tests, and report a clear error when `flydsl` is not installed. + +Octave is enabled only on RDNA3 and CDNA3, validated on gfx1100 and gfx942 +(MI300). The backend rejects any other family at selection +(`supports_compute_capability`, via `get_rdna_version()` and +`get_cdna_version()`), and the kernel tests skip elsewhere. CMake builds the +Octave sources and registers the ops only when the target list includes +gfx1100 or gfx942; other targets in the same build compile the +architecture-specific kernels as empty stubs, and the kernel dispatch reports +any other RDNA3 chip as not validated. + +| Piece | gfx1100 (RDNA3) | gfx942 (MI300) | +| --- | --- | --- | +| int8 dot4 (`octave_arch.cuh`) | `v_dot4_i32_iu8` | `v_dot4_i32_i8` | +| Codebook lookup (`octave_format.cuh`) | `v_perm_b32` | `v_perm_b32` | +| Cache store, reduce, rotate | portable | portable (decode reduce: gfx942 kernel) | +| Decode | WMMA | MFMA (`v_mfma_i32_16x16x32_i8`) | +| Prefill | WMMA | MFMA | + +The portable dot4 decode stays as the reference tier the tests compare the +native tiers against. + +To add an architecture: + +1. Implement `octave_prefill` for it, in HIP or FlyDSL (same signature and semantics as + `octave_prefill_rdna3.cu`: phase-1 partials over the packed prefix, phase-2 + chunk attention and merge, rotated space in and out). +2. Add a decode tier, or validate the portable dot4 kernel there. +3. Add the architecture to the CMake guard, to `_kernel_tier` and + `supports_compute_capability` in + `vllm/v1/attention/backends/rocm_octave_attn.py`, and to the test gate. +4. Run `tests/kernels/attention/test_rocm_octave.py`: every kernel is checked + against the PyTorch reference, and `test_decode_matches_reference` runs + both the fast and the portable decode tier. + +## Results + +All measurements on 4× RX 7900 XTX (gfx1100), Qwen3.8-27B W4A16, tensor +parallel 4, MTP k=3, 8 GB of KV cache per GPU. + +### Capacity + +| KV cache | Tokens that fit | +| --- | --- | +| fp16 | TBD | +| `turboquant_k3v4_nc` | 1,790,625 | +| `octave_k4v4` | 1,402k | +| `octave_k3v3` | 1,679k | + +### Speed + +TBD: `vllm bench serve` (random prompts, no prefix reuse) TTFT and TPOT at +8k / 32k / 100k / 380k input, and per-layer kernel times from +`benchmarks/attention_benchmarks`. + +### Quality + +TBD: GSM8K (`tests/evals/gsm8k/gsm8k_eval.py`, 1319 questions, 5-shot, +temperature 0.6) and next-token NLL over a 380k-token document. + +## Reproducing + +```bash +vllm serve Qwen/Qwen3.8-27B --kv-cache-dtype octave_k3v4 ... +pytest tests/kernels/attention/test_rocm_octave.py +python benchmarks/attention_benchmarks/benchmark.py \ + --backends ROCM_OCTAVE --kv-cache-dtype octave_k3v4 \ + --model Qwen/Qwen3.8-27B --head-dim 256 --num-q-heads 6 --num-kv-heads 1 \ + --batch-specs q1s100k q4s100k q2ks100k +``` diff --git a/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3-compact.yaml b/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3-compact.yaml new file mode 100644 index 00000000000..c69e3805d8e --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3-compact.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3.5-4B" +accuracy_threshold: 0.76 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype octave_k3v3_compact --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3.yaml b/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3.yaml new file mode 100644 index 00000000000..e903284c8f8 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v3.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3.5-4B" +accuracy_threshold: 0.77 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype octave_k3v3 --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v4.yaml b/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v4.yaml new file mode 100644 index 00000000000..5f53f8dd40c --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3.5-4B-Octave-k3v4.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3.5-4B" +accuracy_threshold: 0.78 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype octave_k3v4 --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/models-octave.txt b/tests/evals/gsm8k/configs/models-octave.txt new file mode 100644 index 00000000000..5669b9ce76a --- /dev/null +++ b/tests/evals/gsm8k/configs/models-octave.txt @@ -0,0 +1,3 @@ +Qwen3.5-4B-Octave-k3v4.yaml +Qwen3.5-4B-Octave-k3v3.yaml +Qwen3.5-4B-Octave-k3v3-compact.yaml diff --git a/tests/kernels/attention/test_rocm_octave.py b/tests/kernels/attention/test_rocm_octave.py new file mode 100644 index 00000000000..27c2a77d42e --- /dev/null +++ b/tests/kernels/attention/test_rocm_octave.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Octave HIP kernels against the PyTorch reference in +vllm/v1/attention/ops/rocm_octave.py.""" + +import math + +import pytest +import torch + +from tests.v1.attention.utils import dense_kv_cache_views +from vllm.platforms import current_platform +from vllm.v1.attention.backends.rocm_octave_attn import RocmOctaveAttentionBackend +from vllm.v1.attention.ops import rocm_octave as sq +from vllm.v1.kv_cache_interface import FullAttentionSpec, get_kv_quant_mode + + +def _arch() -> str | None: + if not current_platform.is_rocm(): + return None + from vllm.platforms.rocm import get_cdna_version, get_rdna_version + + if get_rdna_version() == 3: + return "rdna3" + return "cdna3" if get_cdna_version() == 3 else None + + +pytestmark = pytest.mark.skipif( + _arch() is None or not hasattr(torch.ops._C, "octave_decode"), + reason="Octave kernels need RDNA3 or CDNA3", +) + +HEAD, ROPE, BLOCK = 256, 64, 16 +CACHE_DTYPES = ["octave_k3v4", "octave_k3v3", "octave_k3v3_compact"] + + +def _format(cache_dtype: str) -> sq.OctaveFormat: + return sq.OctaveFormat.from_cache_dtype(cache_dtype, HEAD, ROPE) + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a.float() - b.float()).norm() / b.float().norm()).item() + + +def _token_slots(cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor: + """(blocks, H, block_size, slot) cache -> (T, H, slot) for flat slots.""" + flat = cache.permute(0, 2, 1, 3).reshape(-1, cache.shape[1], cache.shape[3]) + return flat[slots] + + +def _fill(cache, block_ids, k, v, k_signs, v_signs, fmt): + n = k.shape[0] + slots = ( + block_ids.long().repeat_interleave(BLOCK)[:n] * BLOCK + + torch.arange(n, device=k.device) % BLOCK + ) + torch.ops._C.octave_cache_store( + k, v, cache, slots, k_signs, v_signs, fmt.kernel_code + ) + return slots + + +@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES) +def test_store_matches_reference(cache_dtype): + torch.manual_seed(0) + dev, fmt = "cuda", _format(cache_dtype) + ksg = vsg = sq.sign_bits(HEAD).to(dev) + t, hkv = 1024, 2 + k = torch.randn(t, hkv, HEAD, device=dev) + k[..., 100] *= 8 # an outlier channel is spread by the rotation + k = k.bfloat16() + v = torch.randn(t, hkv, HEAD, device=dev).bfloat16() + nblk = t // BLOCK + 4 + cache = torch.zeros(nblk, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev) + slots = _fill(cache, torch.randperm(nblk, device=dev), k, v, ksg, vsg, fmt) + + k_got, v_got = sq.reference_dequantize(_token_slots(cache, slots), fmt) + ref = sq.reference_quantize(k.float(), v.float(), fmt) + k_ref, v_ref = sq.reference_dequantize(ref, fmt) + # Same codes up to rounding ties from a different summation order. + assert _rel(k_got, k_ref) < 1e-2 and _rel(v_got, v_ref) < 1e-2 + # Lloyd-Max on Gaussian data: 3-bit codes (4-bit for block K's RoPE + # block) keep ~19% of the norm as error, 4-bit ~10%. + assert _rel(k_got, k) < 0.2 + assert _rel(v_got, v) < (0.12 if fmt.v_bits == 4 else 0.22) + + +@pytest.mark.parametrize( + "cache_dtype,slot_bytes", + [("octave_k3v4", 248), ("octave_k3v3", 216), ("octave_k3v3_compact", 196)], +) +def test_slot_bytes_match_kernel_layout(cache_dtype, slot_bytes): + """The cache is sized from the Python format; the kernels index it with + octave::Format in csrc/attention/octave_format.cuh.""" + assert _format(cache_dtype).slot_bytes == slot_bytes + + +@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES) +def test_store_skips_padding_slots(cache_dtype): + """Tokens with slot -1 (cudagraph padding) leave every byte unchanged.""" + torch.manual_seed(0) + dev, fmt = "cuda", _format(cache_dtype) + ksg = vsg = sq.sign_bits(HEAD).to(dev) + t, hkv = 8, 2 + k = torch.randn(t, hkv, HEAD, device=dev).bfloat16() + v = torch.randn(t, hkv, HEAD, device=dev).bfloat16() + cache = torch.randint( + 0, 256, (2, hkv, BLOCK, fmt.slot_bytes), dtype=torch.uint8, device=dev + ) + before = cache.clone() + slots = torch.full((t,), -1, dtype=torch.long, device=dev) + slots[3] = 21 + torch.ops._C.octave_cache_store(k, v, cache, slots, ksg, vsg, fmt.kernel_code) + written = torch.zeros_like(cache, dtype=torch.bool) + written[1, :, 5] = True + assert torch.equal(cache[~written], before[~written]) + assert not torch.equal(cache[written], before[written]) + + +@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES) +@pytest.mark.parametrize("query_group", [1, 4]) +@pytest.mark.parametrize("num_splits", [1, 7, 64]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("use_wmma", [True, False]) +def test_decode_matches_reference( + cache_dtype, query_group, num_splits, dtype, use_wmma +): + """Covers MTP verification: up to 4 consecutive query tokens per request, + each with its own causal length, plus a padded query with no request. + ``use_wmma=False`` forces the portable kernel every architecture runs.""" + torch.manual_seed(1) + dev, fmt = "cuda", _format(cache_dtype) + ksg = vsg = sq.sign_bits(HEAD).to(dev) + hkv, group = 1, 6 + lens = [1, 16, 37, 700, 3000] + cache = torch.zeros(512, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev) + max_blocks = max(math.ceil(n / BLOCK) for n in lens) + block_table = torch.zeros(len(lens), max_blocks, dtype=torch.int32, device=dev) + free = torch.randperm(512, device=dev) + q_to_req, q_to_klen, used = [], [], 0 + for i, n in enumerate(lens): + nb = math.ceil(n / BLOCK) + block_table[i, :nb] = free[used : used + nb].int() + used += nb + k = torch.randn(n, hkv, HEAD, device=dev).bfloat16() + v = torch.randn(n, hkv, HEAD, device=dev).bfloat16() + _fill(cache, block_table[i, :nb], k, v, ksg, vsg, fmt) + nq = min(query_group, n) + for j in range(nq): + q_to_req.append(i) + q_to_klen.append(n - nq + 1 + j) + q_to_req.append(len(lens) + 3) # cudagraph padding: no request + q_to_klen.append(0) + q_to_req = torch.tensor(q_to_req, dtype=torch.int32, device=dev) + q_to_klen = torch.tensor(q_to_klen, dtype=torch.int32, device=dev) + + num_q = q_to_req.numel() + q = (torch.randn(num_q, hkv * group, HEAD, device=dev) * 2).to(dtype) + out = torch.zeros_like(q) + mid = torch.empty( + num_q, hkv * group, num_splits, HEAD + 2, dtype=torch.float32, device=dev + ) + torch.ops._C.octave_decode( + out, + q, + cache, + block_table, + q_to_req, + q_to_klen, + mid, + ksg, + vsg, + 1 / 16, + num_splits, + fmt.kernel_code, + query_group, + use_wmma, + ) + ref = sq.reference_attention( + q.float(), + cache, + block_table, + q_to_req.clamp(max=len(lens) - 1), + q_to_klen, + fmt, + 1 / 16, + ) + # Only the kernel's query precision (fp16 or int8) differs. + assert _rel(out[:-1], ref[:-1]) < 2e-2 + assert out[-1].abs().max().item() == 0 + + +@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES) +@pytest.mark.parametrize("query_lens", [[300], [129, 1, 517]]) +def test_prefill_matches_reference(cache_dtype, query_lens): + """Chunked prefill: each request's chunk attends causally to itself + (unquantized K/V) and to its cached prefix (packed cache), in the rotated + space the backend runs it in.""" + torch.manual_seed(3) + dev, fmt = "cuda", _format(cache_dtype) + ksg = vsg = sq.sign_bits(HEAD).to(dev) + hkv, group = 1, 6 + ctx_lens = [2500, 0, 4100][: len(query_lens)] + seq_lens = [c + n for c, n in zip(ctx_lens, query_lens)] + max_blocks = max(math.ceil(n / BLOCK) for n in seq_lens) + nblk = sum(math.ceil(n / BLOCK) for n in seq_lens) + 8 + cache = torch.zeros(nblk, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev) + block_table = torch.zeros(len(seq_lens), max_blocks, dtype=torch.int32, device=dev) + free, used = torch.randperm(nblk, device=dev), 0 + ks, vs = [], [] + for i, n in enumerate(seq_lens): + nb = math.ceil(n / BLOCK) + block_table[i, :nb] = free[used : used + nb].int() + used += nb + k = torch.randn(n, hkv, HEAD, device=dev).half() + v = torch.randn(n, hkv, HEAD, device=dev).half() + _fill(cache, block_table[i, :nb], k, v, ksg, vsg, fmt) + ks.append(k) + vs.append(v) + + qsl = torch.tensor([0, *query_lens], device=dev).cumsum(0).int() + total = int(qsl[-1]) + q = (torch.randn(total, hkv * group, HEAD, device=dev) * 2).half() + k_new = torch.cat([k[c:] for k, c in zip(ks, ctx_lens)]) + v_new = torch.cat([v[c:] for v, c in zip(vs, ctx_lens)]) + q_rot, k_rot, v_rot = q.clone(), k_new.clone(), v_new.clone() + torch.ops._C.octave_rotate(q_rot, ksg, not fmt.compact, False) + torch.ops._C.octave_rotate(k_rot, ksg, not fmt.compact, False) + torch.ops._C.octave_rotate(v_rot, vsg, False, False) + out = torch.empty_like(q_rot) + torch.ops._C.octave_prefill( + out, + q_rot, + k_rot, + v_rot, + cache, + block_table, + qsl, + torch.tensor(seq_lens, dtype=torch.int32, device=dev), + max(query_lens), + 1 / 16, + fmt.kernel_code, + ) + torch.ops._C.octave_rotate(out, vsg, False, True) + + ref = torch.empty(total, hkv * group, HEAD, device=dev) + for i, (c, n) in enumerate(zip(ctx_lens, query_lens)): + k_all, v_all = ks[i][c:].float(), vs[i][c:].float() + if c > 0: + pos = torch.arange(c, device=dev) + slots = cache[block_table[i, pos // BLOCK].long(), :, pos % BLOCK] + k_pre, v_pre = sq.reference_dequantize(slots, fmt) + k_all, v_all = torch.cat([k_pre, k_all]), torch.cat([v_pre, v_all]) + k_all = k_all.repeat_interleave(group, 1) + v_all = v_all.repeat_interleave(group, 1) + qi = q[int(qsl[i]) : int(qsl[i + 1])].float() + s = torch.einsum("qhd,thd->hqt", qi, k_all) / 16 + mask = ( + torch.arange(c + n, device=dev)[None] + > (c + torch.arange(n, device=dev))[:, None] + ) + p = torch.softmax(s.masked_fill(mask, float("-inf")), -1) + ref[int(qsl[i]) : int(qsl[i + 1])] = torch.einsum("hqt,thd->qhd", p, v_all) + # Only the kernel's int8 query and fp16 P/V rounding differ. + assert _rel(out, ref) < 2e-2 + + +def _decode(cache, q, block_table, q_to_req, q_to_klen, signs, fmt, use_wmma): + out = torch.zeros_like(q) + mid = torch.empty(*q.shape[:2], 8, HEAD + 2, dtype=torch.float32, device=q.device) + torch.ops._C.octave_decode( + out, + q, + cache, + block_table, + q_to_req, + q_to_klen, + mid, + signs, + signs, + 1 / 16, + 8, + fmt.kernel_code, + 1, + use_wmma, + ) + return out + + +def _prefill(cache, q, k, v, block_table, qsl, seq_lens, signs, fmt): + out = torch.empty_like(q) + torch.ops._C.octave_prefill( + out, + q, + k, + v, + cache, + block_table, + qsl, + seq_lens, + int((qsl[1:] - qsl[:-1]).max()), + 1 / 16, + fmt.kernel_code, + ) + return out + + +@pytest.mark.parametrize( + "layout", RocmOctaveAttentionBackend.supported_kv_cache_layouts() +) +def test_kernels_follow_declared_cache_layouts(layout): + """Store, decode and prefill on the cache view vLLM allocates for each + declared layout match a contiguous cache bit for bit. Two KV heads, so the + head and token strides differ between layouts.""" + torch.manual_seed(4) + dev, cache_dtype = "cuda", "octave_k3v3" + fmt, hkv, group, nblk = _format(cache_dtype), 2, 6, 80 + signs = sq.sign_bits(HEAD).to(dev) + sq.register_rope_dim(HEAD, ROPE) + spec = RocmOctaveAttentionBackend.customize_spec( + FullAttentionSpec( + block_size=BLOCK, + num_kv_heads=hkv, + head_size=HEAD, + dtype=torch.uint8, + kv_quant_mode=get_kv_quant_mode(cache_dtype), + ) + ) + raw = torch.zeros(nblk * spec.page_size_bytes, dtype=torch.int8, device=dev) + cache = dense_kv_cache_views(raw, spec, nblk, 1, layout)[0] + dense = torch.zeros(nblk, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev) + assert cache.shape == dense.shape and cache.dtype == dense.dtype + + seq_lens, query_lens = [700, 300], [100, 300] + max_blocks = max(math.ceil(n / BLOCK) for n in seq_lens) + block_table = torch.zeros(len(seq_lens), max_blocks, dtype=torch.int32, device=dev) + free, used, ks, vs = torch.randperm(nblk, device=dev), 0, [], [] + for i, n in enumerate(seq_lens): + nb = math.ceil(n / BLOCK) + block_table[i, :nb] = free[used : used + nb].int() + used += nb + ks.append(torch.randn(n, hkv, HEAD, device=dev).half()) + vs.append(torch.randn(n, hkv, HEAD, device=dev).half()) + for c in (cache, dense): + _fill(c, block_table[i, :nb], ks[i], vs[i], signs, signs, fmt) + assert torch.equal(cache, dense) + + q_to_req = torch.arange(len(seq_lens), dtype=torch.int32, device=dev) + q_to_klen = torch.tensor(seq_lens, dtype=torch.int32, device=dev) + q = torch.randn(len(seq_lens), hkv * group, HEAD, device=dev).half() + for use_wmma in (True, False): + out, ref = ( + _decode(c, q, block_table, q_to_req, q_to_klen, signs, fmt, use_wmma) + for c in (cache, dense) + ) + assert out.abs().sum() > 0 + assert torch.equal(out, ref) + + qsl = torch.tensor([0, *query_lens], device=dev).cumsum(0).int() + q = torch.randn(int(qsl[-1]), hkv * group, HEAD, device=dev).half() + ctx = [n - m for n, m in zip(seq_lens, query_lens)] + k = torch.cat([x[c:] for x, c in zip(ks, ctx)]) + v = torch.cat([x[c:] for x, c in zip(vs, ctx)]) + out, ref = ( + _prefill(c, q, k, v, block_table, qsl, q_to_klen, signs, fmt) + for c in (cache, dense) + ) + assert torch.equal(out, ref) + + +def _window_tables(entries, hkv, dev): + return ( + torch.zeros(entries, hkv, 2 * HEAD, dtype=torch.float16, device=dev), + torch.full((entries,), -1, dtype=torch.int32, device=dev), + torch.zeros(entries, dtype=torch.int32, device=dev), + ) + + +@pytest.mark.skipif(_arch() != "rdna3", reason="sparse decode is RDNA3-only") +@pytest.mark.parametrize("cache_dtype", CACHE_DTYPES) +@pytest.mark.parametrize("hkv,group", [(1, 6), (2, 6)]) +@pytest.mark.parametrize("window", ["off", "on", "evicted"]) +def test_sparse_decode_matches_reference(cache_dtype, hkv, group, window): + """Sparse (top-k) decode as qwen4_exp QSA runs it: each query reads the + logical positions in its index row, -1 entries are holes. With the window + on, the first 4 and the last 64 positions come from the fp16 tables; + an entry whose tag was overwritten falls back to the packed cache.""" + torch.manual_seed(4) + dev, fmt = "cuda", _format(cache_dtype) + signs = sq.sign_bits(HEAD).to(dev) + lens, max_blocks, nblk = [1500, 2600, 777], 192, 1024 + cache = torch.zeros(nblk, hkv, BLOCK, fmt.slot_bytes, dtype=torch.uint8, device=dev) + block_table = ( + torch.randperm(nblk, device=dev)[: len(lens) * max_blocks] + .view(len(lens), max_blocks) + .int() + ) + wtab, wtags, wlocks = _window_tables(4096, hkv, dev) + stab, stags, slocks = _window_tables(256, hkv, dev) + ks, vs = [], [] + for r, n in enumerate(lens): + k = torch.randn(n, hkv, HEAD, device=dev).half() + v = torch.randn_like(k) + pos = torch.arange(n, device=dev) + slots = block_table[r, pos // BLOCK].long() * BLOCK + pos % BLOCK + torch.ops._C.octave_cache_store( + k, v, cache, slots, signs, signs, fmt.kernel_code + ) + torch.ops._C.octave_window_store( + k, v, slots, pos.int(), wtab, wtags, wlocks, stab, stags, slocks + ) + ks.append(k) + vs.append(v) + if window == "evicted": + wtags[::2] = -7 + stags[::2] = -7 + + q_to_req = torch.tensor([0, 1, 1, 2, 0, 1, 2], dtype=torch.int32, device=dev) + num_q, topk = q_to_req.numel(), 512 + qpos = torch.tensor( + [lens[r] - 1 - 37 * i for i, r in enumerate(q_to_req.tolist())], + dtype=torch.int32, + device=dev, + ) + indices = torch.full((num_q, topk), -1, dtype=torch.int32, device=dev) + for i in range(num_q): + n = int(qpos[i]) + 1 + cand = torch.cat( + [ + torch.arange(max(0, n - 80), n, device=dev), + torch.arange(4, device=dev), + torch.randperm(n, device=dev)[:300], + ] + ).unique() + indices[i, : cand.numel()] = cand.int() + indices[i] = indices[i][torch.randperm(topk, device=dev)] + + q = torch.randn(num_q, hkv * group, HEAD, device=dev).half() + out = torch.empty_like(q) + mid = torch.empty(num_q, hkv * group, 64, HEAD + 2, dtype=torch.float32, device=dev) + scale = HEAD**-0.5 + window_args = () if window == "off" else (qpos, wtab, wtags, stab, stags) + torch.ops._C.octave_decode_sparse( + out, + q, + cache, + block_table, + q_to_req, + indices, + mid, + signs, + signs, + scale, + 64, + fmt.kernel_code, + *window_args, + ) + + ref = torch.zeros(num_q, hkv * group, HEAD, device=dev) + for i in range(num_q): + r, p = int(q_to_req[i]), int(qpos[i]) + sel = indices[i][indices[i] >= 0].long() + blocks = block_table[r, sel // BLOCK].long() + k, v = sq.reference_dequantize(cache[blocks, :, sel % BLOCK], fmt) + if window != "off": + slot = blocks * BLOCK + sel % BLOCK + exact = ((p - sel) < 64) & (wtags[slot % 4096].long() == slot) + exact |= (sel < 4) & (stags[slot % 256].long() == slot) + k[exact] = ks[r][sel[exact]].float() + v[exact] = vs[r][sel[exact]].float() + k = k.repeat_interleave(group, 1) + v = v.repeat_interleave(group, 1) + p_attn = torch.softmax(torch.einsum("hd,thd->ht", q[i].float(), k) * scale, -1) + ref[i] = torch.einsum("ht,thd->hd", p_attn, v) + assert _rel(out, ref) < 2e-2 diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index 4d619321182..e66a577737b 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -163,6 +163,19 @@ class TestIsQuantizedKvCache: assert get_kv_quant_mode("fp8_per_token_head") == KVQuantMode.FP8_PER_TOKEN_HEAD + @pytest.mark.parametrize( + "cache_dtype", ["octave_k3v4", "octave_k3v3", "octave_k3v3_compact"] + ) + def test_octave(self, cache_dtype): + from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype + from vllm.v1.kv_cache_interface import get_kv_quant_mode + + mode = get_kv_quant_mode(cache_dtype) + assert mode.name.lower() == cache_dtype + assert mode.is_octave and not mode.is_per_token_head + assert is_quantized_kv_cache(cache_dtype) + assert kv_cache_dtype_str_to_dtype(cache_dtype, None) == torch.uint8 + # =========================================================================== # 2. Triton per-token-head kernel (reshape-and-cache) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index b540b1ec8f1..626825dba8c 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -55,6 +55,9 @@ CacheDType = Literal[ "turboquant_4bit_nc", "turboquant_k3v4_nc", "turboquant_3bit_nc", + "octave_k3v4", + "octave_k3v3", + "octave_k3v3_compact", "ultraquant_4bit", "int4_per_token_head", "int8_per_token_head", diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index 709fd0b08ef..6a246dd31ec 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -10,7 +10,7 @@ import torch from torch import nn from transformers import Qwen4ExpTextConfig -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context @@ -46,28 +46,93 @@ from vllm.v1.attention.backends.flash_attn import ( FlashAttentionMetadata, FlashAttentionMetadataBuilder, ) +from vllm.v1.attention.ops import rocm_octave from vllm.v1.kv_cache_interface import ( + AttentionSpec, FullAttentionSpec, KVCacheSpec, get_kv_quant_mode, ) +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) from ..common.qsa_cache import QSAForwardMetadata from . import model from .indexer_qsa import QSAIndexer +def _octave_mid_o_capacity(vllm_config: VllmConfig) -> int: + """Query rows times splits the Octave split-KV partials hold: every split + up to the cudagraph capture size, and one split per row of a full chunk.""" + capture = vllm_config.compilation_config.max_cudagraph_capture_size or 4 + return max( + capture * vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph, + 2 * vllm_config.scheduler_config.max_num_batched_tokens, + ) + + +# Exact window of the Octave QSA cache: fp16 shadow tables (entries, powers of +# two) of every stored token and of each sequence's leading positions; the +# kernels keep the last 64 and the first 4 selected positions exact. +_OCTAVE_WINDOW_ENTRIES = 16384 +_OCTAVE_SINK_ENTRIES = 2048 + + class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder): """Flash metadata supporting uniform decode and target-verify graphs.""" _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + # Layers run one after another, so they share one set of Octave + # split-KV partials; reserve it before the workspace is locked. + if kv_cache_spec.kv_quant_mode.is_octave and is_workspace_manager_initialized(): + num_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + numel = ( + _octave_mid_o_capacity(vllm_config) + * num_heads + * (kv_cache_spec.head_size + 2) + ) + current_workspace_manager().get_simultaneous(((numel,), torch.float32)) + + +_OCTAVE_CACHE_DTYPES: tuple[CacheDType, ...] = ( + "octave_k3v4", + "octave_k3v3", + "octave_k3v3_compact", +) +_QSA_CACHE_DTYPES: tuple[CacheDType, ...] = ("auto", "bfloat16", *_OCTAVE_CACHE_DTYPES) + class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): - """FullAttentionSpec backend used by the merged QSA owner.""" + """FullAttentionSpec backend used by the merged QSA owner. + + With an Octave ``kv_cache_dtype`` the main K/V cache holds one packed + Octave slot per (token, KV head) and sparse attention runs the Octave + sparse decode kernel over the selected positions. + """ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = list(_QSA_CACHE_DTYPES) + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + from vllm.v1.attention.backends.rocm_octave_attn import ( + RocmOctaveAttentionBackend, + ) + + return RocmOctaveAttentionBackend.customize_spec(spec) @staticmethod def get_name() -> str: @@ -109,9 +174,106 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): raise NotImplementedError( "Qwen4Exp QSA does not support decode context parallelism" ) - if self.kv_cache_dtype not in ("auto", "bfloat16"): - 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, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 65f165ad3ad..de899f7a1bf 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -869,7 +869,8 @@ class Platform: kv_quant_mode=get_kv_quant_mode(cache_dtype), ) # The backend owns its packing - return backend_cls.customize_spec(spec).page_size_bytes + with set_current_vllm_config(vllm_config): + return backend_cls.customize_spec(spec).page_size_bytes primary_dtype = ( STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype] @@ -1035,9 +1036,11 @@ class Platform: dtype=kv_cache_dtype, kv_quant_mode=kv_quant_mode, ) - attn_page_size_1_token = backend_cls.customize_spec( - 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( diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 0ce608bd729..34727fe0490 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -376,6 +376,14 @@ def get_cdna_version() -> int: return 0 +def get_rdna_version() -> int: + if on_rdna4(): + return 4 + if on_gfx11(): + return 3 + return 0 + + @cache def get_rocm_version() -> tuple[int, ...] | None: """Return the installed ROCm release as (major, minor, patch), or None.""" @@ -521,6 +529,7 @@ def _get_backend_priorities( backends.insert(0, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN) backends.append(AttentionBackendEnum.TRITON_ATTN) backends.append(AttentionBackendEnum.TURBOQUANT) + backends.append(AttentionBackendEnum.ROCM_OCTAVE) backends.append(AttentionBackendEnum.ULTRAQUANT) return backends diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 43e1efd7392..7e889e244ad 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -57,6 +57,9 @@ STR_DTYPE_TO_TORCH_DTYPE = { "turboquant_4bit_nc": torch.uint8, "turboquant_k3v4_nc": torch.uint8, "turboquant_3bit_nc": torch.uint8, + "octave_k3v4": torch.uint8, + "octave_k3v3": torch.uint8, + "octave_k3v3_compact": torch.uint8, "ultraquant_4bit": torch.uint8, "nvfp4": torch.uint8, "nvfp4_4over6": torch.uint8, diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 0499c469028..093a536cd4e 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -160,6 +160,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): CPU_MLA = "vllm.v1.attention.backends.mla.cpu_mla.CPUMLABackend" AMX_MLA = "vllm.v1.attention.backends.mla.amx_mla.AMXMLABackend" TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend" + ROCM_OCTAVE = ( + "vllm.v1.attention.backends.rocm_octave_attn.RocmOctaveAttentionBackend" + ) ULTRAQUANT = "vllm.v1.attention.backends.ultraquant_attn.UltraQuantAttentionBackend" # Placeholder for third-party/custom backends - must be registered before use # set to None to avoid alias with other backend, whose value is an empty string diff --git a/vllm/v1/attention/backends/rocm_octave_attn.py b/vllm/v1/attention/backends/rocm_octave_attn.py new file mode 100644 index 00000000000..3af17b4aaf6 --- /dev/null +++ b/vllm/v1/attention/backends/rocm_octave_attn.py @@ -0,0 +1,496 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Octave attention backend for AMD GPUs (native kernels: HIP or FlyDSL). + +KV cache format and reference: ``vllm/v1/attention/ops/rocm_octave.py``. +Kernels: ``csrc/attention/octave_attn.cu`` and the per-architecture tiers in +``_kernel_tier``; RDNA3 and CDNA3, validated on gfx1100 and gfx942. See +``docs/design/rocm_octave.md``. + +Decode, MTP verification and short continuation chunks (every request with +at most ``_DECODE_MAX_QUERY_LEN`` query tokens) run the split-KV decode +kernel over the packed cache; consecutive query tokens of one request share +each K/V read. Longer prefills run the prefill kernel, which reads the cached +prefix from the packed cache and the current chunk's K/V unquantized. +""" + +import contextlib +import math +from dataclasses import dataclass, replace +from typing import ClassVar + +import torch + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.cache import CacheDType +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionImpl, + AttentionLayer, + AttentionMetadataBuilder, + AttentionType, + CommonAttentionMetadata, + MultipleOf, +) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.attention.ops import rocm_octave as sq +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec +from vllm.v1.kv_cache_layout import KVCacheLayout +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + + +# Native kernel tier per architecture family; any other family is rejected at +# selection (supports_compute_capability) until it is validated. +def _kernel_tier() -> str: + from vllm.platforms.rocm import get_cdna_version, get_rdna_version + + if get_rdna_version() == 3: + return "wmma" + if get_cdna_version() == 3: + return "mfma" + raise NotImplementedError( + "ROCM_OCTAVE supports RDNA3 and CDNA3 only; " + "see the architecture contract in docs/design/rocm_octave.md" + ) + + +_DECODE_MAX_QUERY_LEN = 128 +_KERNEL_HEAD_SIZE = 256 +_KERNEL_ROPE_DIM = 64 + + +def _mid_o_shape( + num_q: int, + num_heads: int, + max_splits: int, + capture_size: int, + split_cap: int | None = None, +) -> tuple[tuple[int, int, int, int], int]: + """Split-KV partials for num_q queries: every split up to the cudagraph + capture size, fewer splits past it so the buffer never outgrows that. + With split_cap, small batches spend the same buffer on up to split_cap + splits (a wide GPU needs them to fill its CUs with one query).""" + budget = capture_size * max_splits + splits = max_splits if split_cap is None else min(split_cap, budget // num_q) + if num_q > capture_size: + splits = max(1, min(max_splits, budget // num_q)) + return (num_q, num_heads, splits, _KERNEL_HEAD_SIZE + 2), splits + + +def _max_capture_size(vllm_config: VllmConfig) -> int: + return vllm_config.compilation_config.max_cudagraph_capture_size or 4 + + +class RocmOctaveAttentionBackend(AttentionBackend): + accept_output_buffer: bool = True + forward_includes_kv_cache_update: bool = False + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "octave_k3v4", + "octave_k3v3", + "octave_k3v3_compact", + ] + + @staticmethod + def get_name() -> str: + return "ROCM_OCTAVE" + + @staticmethod + def get_impl_cls() -> type["RocmOctaveAttentionImpl"]: + return RocmOctaveAttentionImpl + + @staticmethod + def get_builder_cls() -> type["RocmOctaveMetadataBuilder"]: + return RocmOctaveMetadataBuilder + + @staticmethod + def get_supported_kernel_block_sizes( + kv_cache_spec: "KVCacheSpec | None" = None, + ) -> list[int | MultipleOf]: + return [MultipleOf(16)] + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + from vllm.platforms.rocm import get_cdna_version, get_rdna_version + + # DeviceCapability is unreliable on ROCm; use the gcnArchName. + return current_platform.is_rocm() and ( + get_rdna_version() == 3 or get_cdna_version() == 3 + ) + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [_KERNEL_HEAD_SIZE] + + @classmethod + def supports_attn_type(cls, attn_type: str) -> bool: + return attn_type == AttentionType.DECODER + + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + # The kernels take block/head/token strides; the slot is contiguous. + # The MFMA decode loads 16 consecutive tokens of a head as one tile. + if current_platform.is_rocm(): + from vllm.platforms.rocm import get_cdna_version + + if get_cdna_version() == 3: + return (KVCacheLayout.LBHNC,) + return (KVCacheLayout.LBHNC, KVCacheLayout.LBNHC) + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + """One packed K+V slot per (token, KV head).""" + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_octave: + return spec + fmt = sq.OctaveFormat.from_cache_dtype( + spec.kv_quant_mode.name.lower(), + spec.head_size, + sq.registered_rope_dim(spec.head_size), + ) + return replace(spec, state_content_bytes=fmt.slot_bytes) + + +@dataclass +class RocmOctaveMetadata: + num_actual_tokens: int + max_query_len: int + query_start_loc: torch.Tensor + seq_lens: torch.Tensor + block_table: torch.Tensor + slot_mapping: torch.Tensor + # Decode requests (one query token) come first in the batch. + num_decodes: int + num_decode_tokens: int + # Per query token: its request and its causal K length (0 for padding). + q_to_req: torch.Tensor + q_to_klen: torch.Tensor + + +class RocmOctaveMetadataBuilder(AttentionMetadataBuilder[RocmOctaveMetadata]): + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + supports_draft_decode_metadata_update = True + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self._init_reorder_batch_threshold(1, supports_spec_as_decode=False) + # Persistent so captured graphs keep reading the current contents. + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self._q_to_req = torch.empty(max_tokens, dtype=torch.int32, device=device) + self._q_to_klen = torch.empty(max_tokens, dtype=torch.int32, device=device) + # Layers run one after another, so they share one set of split-KV + # partials; reserve it before the workspace is locked. + if is_workspace_manager_initialized(): + num_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + capture_size = _max_capture_size(vllm_config) + shape, _ = _mid_o_shape( + capture_size, + num_heads, + vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph, + capture_size, + ) + current_workspace_manager().get_simultaneous((shape, torch.float32)) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> RocmOctaveMetadata: + cm = common_attn_metadata + n = cm.num_actual_tokens + num_decodes, _, num_decode_tokens, _ = split_decodes_and_prefills(cm) + md = RocmOctaveMetadata( + num_actual_tokens=n, + max_query_len=cm.max_query_len, + query_start_loc=cm.query_start_loc, + seq_lens=cm.seq_lens, + block_table=cm.block_table_tensor, + slot_mapping=cm.slot_mapping, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + q_to_req=self._q_to_req[:n], + q_to_klen=self._q_to_klen[:n], + ) + self._fill_query_maps(md) + return md + + def update_draft_decode_metadata(self, metadata: RocmOctaveMetadata) -> None: + # Draft steps advance seq_lens in place; the K lengths follow it. + self._fill_query_maps(metadata) + + @staticmethod + def _fill_query_maps(md: RocmOctaveMetadata) -> None: + """On the device, without host syncs (also inside graph capture). + Queries past the last request (cudagraph padding) get K length 0.""" + n = md.num_actual_tokens + num_reqs = md.query_start_loc.shape[0] - 1 + if n == 0 or num_reqs == 0: + return + idx = torch.arange(n, dtype=torch.int32, device=md.q_to_req.device) + qsl = md.query_start_loc.to(torch.int32) + req = torch.searchsorted(qsl[1:], idx, right=True).to(torch.int32) + r = req.clamp(max=num_reqs - 1).long() + q_lens = qsl[1:] - qsl[:-1] + klen = md.seq_lens.to(torch.int32)[r] - q_lens[r] + (idx - qsl[r]) + 1 + md.q_to_req.copy_(req) + md.q_to_klen.copy_(torch.where(req < num_reqs, klen, 0)) + + +class RocmOctaveAttentionImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int | None = None, + alibi_slopes: list[float] | None = None, + sliding_window: int | None = None, + kv_cache_dtype: str = "auto", + logits_soft_cap: float | None = None, + attn_type: str = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + sinks: torch.Tensor | None = None, + **kwargs, + ) -> None: + if alibi_slopes is not None or sinks is not None: + raise NotImplementedError("ROCM_OCTAVE does not support ALiBi or sinks") + if sliding_window is not None or logits_soft_cap: + raise NotImplementedError( + "ROCM_OCTAVE does not support sliding window or soft cap" + ) + if attn_type != AttentionType.DECODER: + raise NotImplementedError("ROCM_OCTAVE supports decoder attention only") + + self.num_heads = num_heads + self.head_size = head_size + self.scale = float(scale) + self.num_kv_heads = num_kv_heads or num_heads + self.kv_cache_dtype = kv_cache_dtype + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name + + vllm_config = get_current_vllm_config() + rope_dim = sq.rope_dim_from_hf_config( + vllm_config.model_config.hf_text_config, head_size + ) + sq.register_rope_dim(head_size, rope_dim) + self.fmt = sq.OctaveFormat.from_cache_dtype(kv_cache_dtype, head_size, rope_dim) + if not ( + current_platform.is_rocm() + and head_size == _KERNEL_HEAD_SIZE + and rope_dim == _KERNEL_ROPE_DIM + and hasattr(torch.ops._C, "octave_decode") + ): + raise NotImplementedError( + "ROCM_OCTAVE kernels need ROCm, head_size=256 and 64 rotary " + f"dims; got head_size={head_size}, rotary dims={rope_dim}" + ) + # The MFMA decode takes up to 256 splits (octave_decode_mfma_cdna.cu). + self._split_cap = 256 if _kernel_tier() == "mfma" else None + + self.max_num_kv_splits = ( + vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph + ) + self._max_capture_size = _max_capture_size(vllm_config) + self._k_signs: torch.Tensor | None = None + self._v_signs: torch.Tensor | None = None + self._mid_o: torch.Tensor | None = None + + def _signs(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + if self._k_signs is None or self._k_signs.device != device: + self._k_signs = sq.sign_bits(self.head_size).to(device) + self._v_signs = self._k_signs + assert self._v_signs is not None + return self._k_signs, self._v_signs + + def _mid_o_buffer( + self, num_q: int, device: torch.device + ) -> tuple[torch.Tensor, int]: + """Split-KV partials, from the workspace every layer shares. Without a + workspace, a buffer of this layer's own, allocated once at the + cudagraph capture size since captured graphs hold its address.""" + shape, splits = _mid_o_shape( + num_q, + self.num_heads, + self.max_num_kv_splits, + self._max_capture_size, + self._split_cap, + ) + if is_workspace_manager_initialized(): + (buf,) = current_workspace_manager().get_simultaneous( + (shape, torch.float32) + ) + return buf, splits + if self._mid_o is None: + full, _ = _mid_o_shape( + self._max_capture_size, + self.num_heads, + self.max_num_kv_splits, + self._max_capture_size, + ) + self._mid_o = torch.empty(full, dtype=torch.float32, device=device) + if math.prod(shape) <= self._mid_o.numel(): + return self._mid_o.view(-1)[: math.prod(shape)].view(shape), splits + return torch.empty(shape, dtype=torch.float32, device=device), splits + + def do_kv_cache_update( + self, + layer: AttentionLayer, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + if slot_mapping.numel() == 0: + return + k_signs, v_signs = self._signs(key.device) + n = slot_mapping.shape[0] + torch.ops._C.octave_cache_store( + key[:n], + value[:n], + kv_cache, + slot_mapping, + k_signs, + v_signs, + self.fmt.kernel_code, + ) + + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: RocmOctaveMetadata, + output: torch.Tensor | None = None, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + assert output is not None + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError("ROCM_OCTAVE does not fuse output quant") + if attn_metadata is None: + return output.fill_(0) + + n = attn_metadata.num_actual_tokens + out = output[:n].view(n, self.num_heads, self.head_size) + q = query[:n].view(n, self.num_heads, self.head_size) + + if attn_metadata.max_query_len <= _DECODE_MAX_QUERY_LEN: + self._decode( + out, + q, + kv_cache, + attn_metadata, + attn_metadata.q_to_req, + attn_metadata.q_to_klen, + ) + return output + + num_dt = attn_metadata.num_decode_tokens + if num_dt > 0: + q_to_req, q_to_klen = attn_metadata.q_to_req, attn_metadata.q_to_klen + self._decode( + out[:num_dt], + q[:num_dt], + kv_cache, + attn_metadata, + q_to_req[:num_dt], + q_to_klen[:num_dt], + ) + self._prefill(out, q, key, value, kv_cache, attn_metadata) + return output + + def _decode( + self, + out: torch.Tensor, + q: torch.Tensor, + kv_cache: torch.Tensor, + md: RocmOctaveMetadata, + q_to_req: torch.Tensor, + q_to_klen: torch.Tensor, + ) -> None: + num_q = q.shape[0] + if num_q == 0: + return + k_signs, v_signs = self._signs(q.device) + mid_o, splits = self._mid_o_buffer(num_q, q.device) + # Async scheduling can free and reuse an input's memory before this + # kernel reads it; tell the allocator these are in use on the stream. + stream = torch.cuda.current_stream() + for t in (q, out, md.block_table, q_to_req, q_to_klen, mid_o): + with contextlib.suppress(Exception): + t.record_stream(stream) + torch.ops._C.octave_decode( + out, + q, + kv_cache, + md.block_table, + q_to_req, + q_to_klen, + mid_o, + k_signs, + v_signs, + self.scale, + splits, + self.fmt.kernel_code, + 4 if md.max_query_len > 1 else 1, + True, # native tier; False runs the portable reference kernel + ) + + def _prefill( + self, + out: torch.Tensor, + q: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + md: RocmOctaveMetadata, + ) -> None: + """Requests after the decodes, in the rotated space, without host + syncs: one call covers them all, reading the cached prefixes straight + from the packed cache and the chunks' own K/V unquantized.""" + nd, ndt, n = md.num_decodes, md.num_decode_tokens, md.num_actual_tokens + num_reqs = md.query_start_loc.shape[0] - 1 + if num_reqs - nd <= 0 or ndt >= n: + return + k_signs, v_signs = self._signs(q.device) + # The kernel runs in fp16; the rotated copies are made in fp16. + q_rot = q[ndt:n].to(torch.float16, copy=True) + k_rot = key[ndt:n].to(torch.float16, copy=True) + v_rot = value[ndt:n].to(torch.float16, copy=True) + k_blocks = not self.fmt.compact + torch.ops._C.octave_rotate(q_rot, k_signs, k_blocks, False) + torch.ops._C.octave_rotate(k_rot, k_signs, k_blocks, False) + torch.ops._C.octave_rotate(v_rot, v_signs, False, False) + o_rot = torch.empty_like(q_rot) + torch.ops._C.octave_prefill( + o_rot, + q_rot, + k_rot, + v_rot, + kv_cache, + md.block_table[nd:num_reqs], + (md.query_start_loc[nd:] - ndt).to(torch.int32), + md.seq_lens[nd:num_reqs].to(torch.int32), + md.max_query_len, + self.scale, + self.fmt.kernel_code, + ) + torch.ops._C.octave_rotate(o_rot, v_signs, False, True) + out[ndt:n] = o_rot diff --git a/vllm/v1/attention/ops/rocm_octave.py b/vllm/v1/attention/ops/rocm_octave.py new file mode 100644 index 00000000000..628c99592d0 --- /dev/null +++ b/vllm/v1/attention/ops/rocm_octave.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Octave KV-cache format: layout, rotations and a PyTorch reference. + +Each (token, KV head) is stored in one byte slot. V is always quantized +whole; K has two layouts: + +* Block K (``octave_k3v4``, ``octave_k3v3``): random sign flips plus a + Walsh-Hadamard transform over each 64-dim block. The RoPE blocks get 4-bit + codes, the NoPE blocks 3-bit codes, and every block has its own scale. +* Compact K (``octave_k3v3_compact``): sign flips plus one Hadamard transform + over all dims, 3-bit codes, one scale. Spreading the RoPE dims over the + whole head lets them share the 3-bit budget. +* V: sign flips plus one Hadamard transform over all dims, ``v_bits`` codes. + +The rotated coordinates are close to Gaussian, so each code indexes a +Lloyd-Max codebook for N(0, 1), stored as int8 (``LUT``): a code ``c`` of a +block with scale ``s`` stands for ``LUT[c] * s / LUT_ONE``. The scale makes +``x . x_hat = |x|^2``: a least-squares fit would shrink every reconstructed +vector, which flattens the softmax and scales down the attention output. +Everything is computed on the fly; there is no calibration. + +Scores are computed in the rotated space (the query is rotated like K) and +the attention output is rotated back once per query head. + +Slot layout (bytes):: + + [k_rope codes][k_nope codes][v codes][fp16 scales][pad to 8] + +Compact K has no 4-bit part and pads to 4 bytes. + +4-bit codes: each little-endian uint32 covers 8 dims; byte j holds dim j in +its low nibble and dim j + 4 in its high nibble, so ``w & 0x0F0F0F0F`` and +``(w >> 4) & 0x0F0F0F0F`` give dims 0..3 and 4..7 as bytes. + +3-bit codes: a 2-bit plane (uint32 per 16 dims; ``(w >> 2i) & 0x03030303`` +gives dims 4i..4i+3) followed by a 1-bit plane (uint32 per 32 dims; +``(w >> i) & 0x01010101`` gives dims 4i..4i+3). code = lo + 4 * hi. + +Scales (fp16): one per 64-dim K block (one in total for compact K), then V. +""" + +import functools +import math +from dataclasses import dataclass + +import torch + +HADAMARD_BLOCK = 64 +_SIGN_SEED = 0x5E17 +ROPE_BITS = 4 +NOPE_BITS = 3 +# Lloyd-Max codebooks for N(0, 1) as int8: LUT[bits][c] / LUT_MULT[bits] +# approximates centroid c (the kernels map codes through the same tables). +LUT = { + 3: (-127, -79, -45, -14, 14, 45, 79, 127), + 4: (-126, -95, -74, -58, -43, -30, -18, -6, 6, 18, 30, 43, 58, 74, 95, 126), +} +LUT_MULT = {3: 59.0, 4: 46.0} +# Fixed point of the stored scales: a code stands for LUT[c] * scale / LUT_ONE, +# which keeps small blocks' scales out of the fp16 subnormal range. +LUT_ONE = 64.0 + +_rope_dim_by_head_size: dict[int, int] = {} + + +def register_rope_dim(head_size: int, rope_dim: int) -> None: + prev = _rope_dim_by_head_size.setdefault(head_size, rope_dim) + if prev != rope_dim: + raise ValueError( + f"Octave: conflicting rotary dims {prev} and {rope_dim} for " + f"head_size={head_size}" + ) + + +def rope_dim_from_hf_config(hf_text_config, head_size: int) -> int: + factor = getattr(hf_text_config, "partial_rotary_factor", None) + if factor is None: + rope_params = getattr(hf_text_config, "rope_parameters", None) or {} + factor = rope_params.get("partial_rotary_factor", 1.0) + return int(head_size * factor) + + +def registered_rope_dim(head_size: int) -> int: + if head_size not in _rope_dim_by_head_size: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None or vllm_config.model_config is None: + raise RuntimeError( + "Octave needs the model's rotary dim before sizing the KV cache" + ) + register_rope_dim( + head_size, + rope_dim_from_hf_config(vllm_config.model_config.hf_text_config, head_size), + ) + return _rope_dim_by_head_size[head_size] + + +@dataclass(frozen=True) +class OctaveFormat: + head_size: int + rope_dim: int + v_bits: int + compact: bool = False + + def __post_init__(self): + d, r = self.head_size, self.rope_dim + if self.v_bits not in (3, 4): + raise ValueError(f"Octave supports 3- or 4-bit V, got {self.v_bits}") + if self.compact and d & (d - 1): + raise ValueError( + f"Compact Octave K needs a power-of-two head_size, got {d}" + ) + if r % HADAMARD_BLOCK or d % HADAMARD_BLOCK: + raise ValueError( + f"Octave needs head_size and rope_dim in multiples of " + f"{HADAMARD_BLOCK}; got head_size={d}, rope_dim={r}" + ) + + @classmethod + def from_cache_dtype(cls, cache_dtype: str, head_size: int, rope_dim: int): + v_bits, compact = { + "octave_k3v4": (4, False), + "octave_k3v3": (3, False), + "octave_k3v3_compact": (3, True), + }[cache_dtype] + return cls(head_size, rope_dim, v_bits, compact) + + @property + def kernel_code(self) -> int: + """Format id the kernels take: V bits, plus 16 for compact K.""" + return self.v_bits + (16 if self.compact else 0) + + @property + def k_rotation_block(self) -> int: + return self.head_size if self.compact else HADAMARD_BLOCK + + @property + def k_wide_dim(self) -> int: + """K dims stored with ROPE_BITS codes (the rest use NOPE_BITS).""" + return 0 if self.compact else self.rope_dim + + @property + def nope_dim(self) -> int: + return self.head_size - self.rope_dim + + @property + def num_k_blocks(self) -> int: + return self.head_size // HADAMARD_BLOCK + + @property + def num_rope_blocks(self) -> int: + return self.rope_dim // HADAMARD_BLOCK + + @property + def off_k_nope(self) -> int: + return self.k_wide_dim * ROPE_BITS // 8 + + @property + def off_v(self) -> int: + return self.off_k_nope + (self.head_size - self.k_wide_dim) * NOPE_BITS // 8 + + @property + def off_scales(self) -> int: + return self.off_v + self.head_size * self.v_bits // 8 + + @property + def num_k_scales(self) -> int: + return 1 if self.compact else self.num_k_blocks + + @property + def num_scales(self) -> int: + return self.num_k_scales + 1 + + @property + def slot_bytes(self) -> int: + raw = self.off_scales + 2 * self.num_scales + align = 4 if self.compact else 8 + return (raw + align - 1) // align * align + + +@functools.cache +def _signs_cpu(n: int) -> torch.Tensor: + g = torch.Generator().manual_seed(_SIGN_SEED) + return torch.randint(0, 2, (n,), generator=g, dtype=torch.int64) * -2 + 1 + + +@functools.cache +def signs(n: int, device_str: str) -> torch.Tensor: + """Fixed ±1 pattern applied before the Hadamard transform.""" + return _signs_cpu(n).to(device=device_str, dtype=torch.float32) + + +def sign_bits(n: int) -> torch.Tensor: + """The sign pattern as int32 words, bit i set where the sign is -1.""" + neg = (_signs_cpu(n) < 0).to(torch.int64).view(-1, 32) + words = (neg << torch.arange(32)).sum(-1) + return torch.where(words >= 2**31, words - 2**32, words).to(torch.int32) + + +def fwht_blocks(x: torch.Tensor, block: int = HADAMARD_BLOCK) -> torch.Tensor: + """Orthonormal Walsh-Hadamard transform over consecutive blocks of the + last dim.""" + shape = x.shape + y = x.reshape(-1, shape[-1] // block, block).float() + h = 1 + while h < block: + y = y.view(*y.shape[:-1], block // (2 * h), 2, h) + a, b = y[..., 0, :], y[..., 1, :] + y = torch.stack((a + b, a - b), dim=-2).view(*y.shape[:-3], block) + h *= 2 + return (y * (1.0 / math.sqrt(block))).view(shape) + + +def rotate( + x: torch.Tensor, sgn: torch.Tensor, block: int = HADAMARD_BLOCK +) -> torch.Tensor: + return fwht_blocks(x.float() * sgn, block) + + +def unrotate( + y: torch.Tensor, sgn: torch.Tensor, block: int = HADAMARD_BLOCK +) -> torch.Tensor: + return fwht_blocks(y.float(), block) * sgn + + +def v_block(head_size: int) -> int: + """V is rotated as one block when head_size is a power of two.""" + if head_size & (head_size - 1) == 0: + return head_size + return HADAMARD_BLOCK + + +@functools.cache +def _lut(bits: int, device_str: str) -> torch.Tensor: + return torch.tensor(LUT[bits], dtype=torch.float32, device=device_str) + + +def _quantize_lut(x: torch.Tensor, bits: int): + """Nearest Lloyd-Max code on the block RMS; the scale keeps + ``x . x_hat = |x|^2``. + + Returns (codes uint8 in [0, 2**bits), scale) with + x ≈ LUT[codes] * scale / LUT_ONE. + """ + lut = _lut(bits, str(x.device)) + ss = x.pow(2).sum(-1, keepdim=True) + rms = (ss / x.shape[-1]).sqrt().clamp_min(1e-12) + y = x / rms * LUT_MULT[bits] + codes = torch.bucketize(y, (lut[1:] + lut[:-1]) / 2) + dot = (x * lut[codes]).sum(-1, keepdim=True) + scale = torch.where(dot > 0, ss / dot.clamp_min(1e-30), 0.0) * LUT_ONE + return codes.to(torch.uint8), scale + + +def _pack(codes: torch.Tensor, bits: int) -> torch.Tensor: + """codes: (..., n) uint8 -> (..., n * bits / 8) uint8, layout above.""" + c = codes.to(torch.int64) + if bits == 4: + g = c.view(*c.shape[:-1], -1, 2, 4) # (word, half, byte) + b = g[..., 0, :] | (g[..., 1, :] << 4) + return b.reshape(*c.shape[:-1], -1).to(torch.uint8) + lo = (c & 3).view(*c.shape[:-1], -1, 4, 4) # (word, i, byte) + lo_b = (lo << (2 * torch.arange(4, device=c.device)).view(4, 1)).sum(-2) + hi = (c >> 2).view(*c.shape[:-1], -1, 8, 4) + hi_b = (hi << torch.arange(8, device=c.device).view(8, 1)).sum(-2) + return torch.cat( + (lo_b.reshape(*c.shape[:-1], -1), hi_b.reshape(*c.shape[:-1], -1)), -1 + ).to(torch.uint8) + + +def _unpack(packed: torch.Tensor, n: int, bits: int) -> torch.Tensor: + p = packed.to(torch.int64) + if bits == 4: + g = p.view(*p.shape[:-1], -1, 4) + c = torch.stack((g & 15, g >> 4), dim=-2) + return c.reshape(*p.shape[:-1], n) + lo_bytes = n // 4 + lo = p[..., :lo_bytes].view(*p.shape[:-1], -1, 1, 4) + lo = (lo >> (2 * torch.arange(4, device=p.device)).view(4, 1)) & 3 + hi = p[..., lo_bytes:].view(*p.shape[:-1], -1, 1, 4) + hi = (hi >> torch.arange(8, device=p.device).view(8, 1)) & 1 + return lo.reshape(*p.shape[:-1], n) + 4 * hi.reshape(*p.shape[:-1], n) + + +def reference_quantize( + key: torch.Tensor, value: torch.Tensor, fmt: OctaveFormat +) -> torch.Tensor: + """(T, H, D) K and V -> (T, H, slot_bytes) uint8 slots.""" + t, h, d = key.shape + r, dev = fmt.rope_dim, key.device + out = torch.zeros(t, h, fmt.slot_bytes, dtype=torch.uint8, device=dev) + nkb = fmt.num_k_blocks + + if fmt.compact: + k = rotate(key, signs(d, str(dev)), d) + k_codes, k_scale = _quantize_lut(k, NOPE_BITS) + out[..., : fmt.off_v] = _pack(k_codes, NOPE_BITS) + else: + k = rotate(key, signs(d, str(dev))).view(t, h, nkb, HADAMARD_BLOCK) + nr = fmt.num_rope_blocks + kr_codes, kr_scale = _quantize_lut(k[:, :, :nr], ROPE_BITS) + kn_codes, kn_scale = _quantize_lut(k[:, :, nr:], NOPE_BITS) + out[..., : fmt.off_k_nope] = _pack(kr_codes.view(t, h, r), ROPE_BITS) + out[..., fmt.off_k_nope : fmt.off_v] = _pack( + kn_codes.view(t, h, fmt.nope_dim), NOPE_BITS + ) + k_scale = torch.cat((kr_scale[..., 0], kn_scale[..., 0]), -1) + + vr = rotate(value, signs(d, str(dev)), v_block(d)) + v_codes, v_scale = _quantize_lut(vr, fmt.v_bits) + out[..., fmt.off_v : fmt.off_scales] = _pack(v_codes, fmt.v_bits) + + scales = torch.cat((k_scale, v_scale), -1) + sc = scales.to(torch.float16).view(torch.uint8).view(t, h, -1) + out[..., fmt.off_scales : fmt.off_scales + sc.shape[-1]] = sc + return out + + +def reference_dequantize( + slots: torch.Tensor, fmt: OctaveFormat, rotated: bool = False +) -> tuple[torch.Tensor, torch.Tensor]: + """(T, H, slot) slots -> fp32 K and V of shape (T, H, D). + + With ``rotated=True`` K and V stay in the rotated space, which is what + the kernels consume. + """ + t, h, _ = slots.shape + r, d, dev = fmt.rope_dim, fmt.head_size, slots.device + sc = ( + slots[..., fmt.off_scales : fmt.off_scales + 2 * fmt.num_scales] + .contiguous() + .view(torch.float16) + .float() + / LUT_ONE + ) + if fmt.compact: + k = _unpack(slots[..., : fmt.off_v], d, NOPE_BITS) + k = _lut(NOPE_BITS, str(dev))[k] * sc[..., :1] + else: + kr = _unpack(slots[..., : fmt.off_k_nope], r, ROPE_BITS) + kn = _unpack(slots[..., fmt.off_k_nope : fmt.off_v], fmt.nope_dim, NOPE_BITS) + k = torch.cat( + (_lut(ROPE_BITS, str(dev))[kr], _lut(NOPE_BITS, str(dev))[kn]), -1 + ).view(t, h, fmt.num_k_blocks, HADAMARD_BLOCK) + k = (k * sc[..., : fmt.num_k_blocks, None]).view(t, h, d) + + v = _unpack(slots[..., fmt.off_v : fmt.off_scales], d, fmt.v_bits) + v = _lut(fmt.v_bits, str(dev))[v] * sc[..., -1:] + if not rotated: + k = unrotate(k, signs(d, str(dev)), fmt.k_rotation_block) + v = unrotate(v, signs(d, str(dev)), v_block(d)) + return k, v + + +def reference_attention( + query: torch.Tensor, + cache: torch.Tensor, + block_table: torch.Tensor, + q_to_req: torch.Tensor, + q_to_klen: torch.Tensor, + fmt: OctaveFormat, + scale: float, +) -> torch.Tensor: + """Per-query-token attention over a Octave paged cache. + + query: (Q, Hq, D). cache: (num_blocks, Hkv, block_size, slot) uint8. + Query token i attends to the first ``q_to_klen[i]`` tokens of request + ``q_to_req[i]``. + """ + num_q, hq, d = query.shape + _, hkv, block_size, _ = cache.shape + group = hq // hkv + out = torch.zeros(num_q, hq, d, dtype=torch.float32, device=query.device) + for i in range(num_q): + klen = int(q_to_klen[i]) + if klen <= 0: + continue + req = int(q_to_req[i]) + pos = torch.arange(klen, device=cache.device) + blocks = block_table[req, pos // block_size].long() + slots = cache[blocks, :, pos % block_size] # (klen, Hkv, slot) + k, v = reference_dequantize(slots, fmt) + k = k.repeat_interleave(group, dim=1) + v = v.repeat_interleave(group, dim=1) + s = torch.einsum("hd,thd->ht", query[i].float(), k) * scale + p = torch.softmax(s, dim=-1) + out[i] = torch.einsum("ht,thd->hd", p, v) + return out diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 8ff16735452..67c9c329203 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -57,6 +57,10 @@ class KVQuantMode(IntEnum): TURBOQUANT_3BIT_NC = 9 NVFP4_DS_MLA = 10 # opaque-bytes NVFP4 DS-MLA layouts (FlashMLA sparse) ULTRAQUANT_4BIT = 11 # packed FP4 values + UE8M0 group scales + # Hadamard-rotated Lloyd-Max codes, packed K+V per slot (ROCm). + OCTAVE_K3V4 = 12 + OCTAVE_K3V3 = 13 + OCTAVE_K3V3_COMPACT = 14 # Backend self-manages kernel dispatch. Reads as quantized but matches no # generic kernel mode (is_per_token_head / is_nvfp4 / is_turboquant all @@ -78,6 +82,15 @@ class KVQuantMode(IntEnum): """True for NVFP4 packed quantization mode.""" return self == KVQuantMode.NVFP4 + @property + def is_octave(self) -> bool: + """True for any Octave quantization mode.""" + return self in ( + KVQuantMode.OCTAVE_K3V4, + KVQuantMode.OCTAVE_K3V3, + KVQuantMode.OCTAVE_K3V3_COMPACT, + ) + @property def is_turboquant(self) -> bool: """True for any turboquant quantization mode.""" @@ -127,6 +140,8 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: return KVQuantMode.NVFP4 if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"): return KVQuantMode[kv_cache_dtype.upper()] + if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("octave_"): + return KVQuantMode[kv_cache_dtype.upper()] if kv_cache_dtype == "ultraquant_4bit": return KVQuantMode.ULTRAQUANT_4BIT if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):