From 385d86a6cf6f05051899b53d83ec55c63cccb379 Mon Sep 17 00:00:00 2001 From: aditi-amd Date: Tue, 6 Oct 2026 13:42:52 -0700 Subject: [PATCH] [Attention] Add UltraQuant 4-bit KV cache backend (FlyDSL D=256) (#57057) Signed-off-by: Aditi Rana --- tests/kernels/ultraquant/test_config.py | 36 + .../ultraquant/test_decode_vs_reference.py | 111 ++ .../ultraquant/test_dequant_vs_reference.py | 88 ++ .../ultraquant/test_flydsl_decode_d256.py | 149 +++ .../kernels/ultraquant/test_seam_contract.py | 106 ++ .../ultraquant/test_store_vs_reference.py | 149 +++ vllm/config/cache.py | 1 + vllm/engine/arg_utils.py | 18 +- vllm/platforms/interface.py | 34 +- vllm/platforms/rocm.py | 59 +- vllm/utils/torch_utils.py | 1 + vllm/v1/attention/backends/registry.py | 1 + vllm/v1/attention/backends/ultraquant_attn.py | 609 +++++++++ .../flydsl_kernels/ultraquant_decode_hd256.py | 983 +++++++++++++++ .../attention/ops/flydsl_ultraquant_decode.py | 834 +++++++++++++ vllm/v1/attention/ops/ultraquant/__init__.py | 19 + vllm/v1/attention/ops/ultraquant/format.py | 218 ++++ vllm/v1/attention/ops/ultraquant/reference.py | 349 ++++++ .../ops/ultraquant/triton_dequant.py | 190 +++ .../attention/ops/ultraquant/triton_store.py | 351 ++++++ .../ultraquant/triton_unified_attention.py | 1087 +++++++++++++++++ vllm/v1/kv_cache_interface.py | 8 + 22 files changed, 5377 insertions(+), 24 deletions(-) create mode 100644 tests/kernels/ultraquant/test_config.py create mode 100644 tests/kernels/ultraquant/test_decode_vs_reference.py create mode 100644 tests/kernels/ultraquant/test_dequant_vs_reference.py create mode 100644 tests/kernels/ultraquant/test_flydsl_decode_d256.py create mode 100644 tests/kernels/ultraquant/test_seam_contract.py create mode 100644 tests/kernels/ultraquant/test_store_vs_reference.py create mode 100644 vllm/v1/attention/backends/ultraquant_attn.py create mode 100644 vllm/v1/attention/ops/flydsl_kernels/ultraquant_decode_hd256.py create mode 100644 vllm/v1/attention/ops/flydsl_ultraquant_decode.py create mode 100644 vllm/v1/attention/ops/ultraquant/__init__.py create mode 100644 vllm/v1/attention/ops/ultraquant/format.py create mode 100644 vllm/v1/attention/ops/ultraquant/reference.py create mode 100644 vllm/v1/attention/ops/ultraquant/triton_dequant.py create mode 100644 vllm/v1/attention/ops/ultraquant/triton_store.py create mode 100644 vllm/v1/attention/ops/ultraquant/triton_unified_attention.py diff --git a/tests/kernels/ultraquant/test_config.py b/tests/kernels/ultraquant/test_config.py new file mode 100644 index 0000000000..14036b44e8 --- /dev/null +++ b/tests/kernels/ultraquant/test_config.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only integration tests for the UltraQuant cache dtype.""" + +from vllm.v1.attention.backends.ultraquant_attn import ( + UltraQuantAttentionBackend, +) +from vllm.v1.attention.ops.flydsl_ultraquant_decode import ( + ultraquant_flydsl_decode_eligible, +) + + +def test_ultraquant_prefers_block_size_64(): + # FlyDSL D=256 decode geometry requires a 64-token KV block. + assert UltraQuantAttentionBackend.get_preferred_block_size(16) == 64 + assert UltraQuantAttentionBackend.get_preferred_block_size(128) == 64 + + +def test_ultraquant_flydsl_decode_routing(): + """CPU check of FlyDSL vs unified-Triton decode eligibility.""" + eligible = dict( + head_size=256, + num_kv_groups=8, + has_sinks=False, + sliding_window=None, + flydsl_loaded=True, + ) + assert ultraquant_flydsl_decode_eligible(**eligible) + for gqa in (6, 8, 16): + assert ultraquant_flydsl_decode_eligible(**{**eligible, "num_kv_groups": gqa}) + + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "flydsl_loaded": False}) + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "head_size": 128}) + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "num_kv_groups": 4}) + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "has_sinks": True}) + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "sliding_window": 4096}) diff --git a/tests/kernels/ultraquant/test_decode_vs_reference.py b/tests/kernels/ultraquant/test_decode_vs_reference.py new file mode 100644 index 0000000000..9a8d431001 --- /dev/null +++ b/tests/kernels/ultraquant/test_decode_vs_reference.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the UltraQuant Triton unified-attention fallback.""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.ultraquant.format import slot_size +from vllm.v1.attention.ops.ultraquant.reference import reference_ultraquant_attention +from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store +from vllm.v1.attention.ops.ultraquant.triton_unified_attention import ( + ultraquant_unified_attention, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +def _build_paged_cache(key, value, block_size): + batch_size, seq_len, num_kv_heads, head_size = key.shape + blocks_per_seq = (seq_len + block_size - 1) // block_size + kv_cache = torch.zeros( + batch_size * blocks_per_seq, + block_size, + num_kv_heads, + slot_size(head_size), + dtype=torch.uint8, + device=key.device, + ) + block_table = torch.arange( + batch_size * blocks_per_seq, + dtype=torch.int32, + device=key.device, + ).view(batch_size, blocks_per_seq) + slot_mapping = torch.arange( + batch_size * blocks_per_seq * block_size, + dtype=torch.int64, + device=key.device, + ).view(batch_size, -1)[:, :seq_len] + ultraquant_store( + key.reshape(-1, num_kv_heads, head_size), + value.reshape(-1, num_kv_heads, head_size), + kv_cache, + slot_mapping.reshape(-1), + ) + return kv_cache, block_table + + +def _attention_quality(actual: torch.Tensor, expected: torch.Tensor) -> dict: + a = actual.float().reshape(-1) + b = expected.float().reshape(-1) + diff = (a - b).abs() + return { + "cos_sim": torch.nn.functional.cosine_similarity(a, b, dim=0).item(), + "mean_abs": diff.mean().item(), + "max_abs": diff.max().item(), + "norm_ratio": (a.norm() / b.norm()).item(), + } + + +def _unified_decode(query, kv_cache, block_table, seq_lens, scale, sinks=None): + batch_size = query.shape[0] + return ultraquant_unified_attention( + query=query, + kv_cache=kv_cache, + block_table=block_table, + seq_lens=seq_lens, + query_start_loc=torch.arange( + batch_size + 1, dtype=torch.int32, device=query.device + ), + scale=scale, + max_query_len=1, + max_seq_len=int(seq_lens.max().item()), + sinks=sinks, + ) + + +@pytest.mark.parametrize("query_group_size", [6, 8]) +def test_unified_decode_d256_matches_reference(query_group_size): + """The D=256 fallback covers both target Qwen GQA layouts.""" + device = torch.device("cuda") + set_random_seed(0xD256 + query_group_size) + + B, Hk, N, D = 2, 1, 256, 256 + Hq = query_group_size * Hk + block_size = 32 + dtype = torch.bfloat16 + scale = D**-0.5 + + query = torch.randn(B, Hq, D, dtype=dtype, device=device) + key = torch.randn(B, N, Hk, D, dtype=dtype, device=device) + value = torch.randn_like(key) + kv_cache, block_table = _build_paged_cache(key, value, block_size) + seq_lens = torch.full((B,), N, dtype=torch.int32, device=device) + + actual = _unified_decode(query, kv_cache, block_table, seq_lens, scale) + expected = reference_ultraquant_attention( + query=query, + key=key, + value=value, + scale=scale, + ) + + quality = _attention_quality(actual, expected) + assert quality["cos_sim"] > 0.98, quality + assert quality["mean_abs"] < 0.08, quality + assert quality["max_abs"] < 0.8, quality diff --git a/tests/kernels/ultraquant/test_dequant_vs_reference.py b/tests/kernels/ultraquant/test_dequant_vs_reference.py new file mode 100644 index 0000000000..db56f56e29 --- /dev/null +++ b/tests/kernels/ultraquant/test_dequant_vs_reference.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the UltraQuant full-KV dequant kernel (continuation prefill). + +The kernel reconstructs cached K/V from FP4 codes plus UE8M0 group scales. +Because dequant is an exact table lookup times a power of two, the result +must match the reference to within the output dtype's rounding. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.ultraquant.format import slot_size +from vllm.v1.attention.ops.ultraquant.reference import ( + ultraquant_dequant, + ultraquant_encode, +) +from vllm.v1.attention.ops.ultraquant.triton_dequant import ( + ultraquant_full_dequant_kv, +) +from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("head_dim", [128, 256]) +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +@pytest.mark.parametrize("cached_len", [32, 96]) +def test_full_dequant_matches_reference(dtype, head_dim, num_kv_heads, cached_len): + device = torch.device("cuda") + set_random_seed(0xD4) + + block_size = 32 + alloc_len = cached_len + num_blocks = alloc_len // block_size + 1 + + key = torch.randn(alloc_len, num_kv_heads, head_dim, dtype=dtype, device=device) + value = torch.randn(alloc_len, num_kv_heads, head_dim, dtype=dtype, device=device) + + kv_cache = torch.zeros( + num_blocks, + block_size, + num_kv_heads, + slot_size(head_dim), + dtype=torch.uint8, + device=device, + ) + # Leave block 0 unused so the block table exercises a non-identity mapping. + slot_mapping = torch.arange( + block_size, block_size + alloc_len, dtype=torch.int64, device=device + ) + ultraquant_store(key, value, kv_cache, slot_mapping) + + block_table = torch.arange( + 1, num_blocks, dtype=torch.int32, device=device + ).unsqueeze(0) + + k_out = torch.zeros( + 1, num_kv_heads, alloc_len, head_dim, dtype=dtype, device=device + ) + v_out = torch.zeros_like(k_out) + ultraquant_full_dequant_kv( + kv_cache=kv_cache, + block_table=block_table, + k_out=k_out, + v_out=v_out, + alloc_len=alloc_len, + ) + torch.accelerator.synchronize() + + # K is stored Hadamard-rotated, V unrotated; the kernel returns them as + # stored, so the reference must not inverse-rotate either. + ref_k = ultraquant_dequant(ultraquant_encode(key, rotate=True), head_dim) + ref_v = ultraquant_dequant(ultraquant_encode(value, rotate=False), head_dim) + + got_k = k_out[0].transpose(0, 1).to(torch.float32) + got_v = v_out[0].transpose(0, 1).to(torch.float32) + + # Reference is fp32; the kernel rounds to the output dtype once. + tol = 1e-2 if dtype is torch.bfloat16 else 1e-3 + torch.testing.assert_close(got_k, ref_k.to(torch.float32), rtol=tol, atol=tol) + torch.testing.assert_close(got_v, ref_v.to(torch.float32), rtol=tol, atol=tol) diff --git a/tests/kernels/ultraquant/test_flydsl_decode_d256.py b/tests/kernels/ultraquant/test_flydsl_decode_d256.py new file mode 100644 index 0000000000..33a1a8273f --- /dev/null +++ b/tests/kernels/ultraquant/test_flydsl_decode_d256.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the production UltraQuant D=256 FlyDSL decode.""" + +import pytest +import torch + +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.flydsl_ultraquant_decode import ( + flydsl_ultraquant_decode_attention, + is_ultraquant_flydsl_available, + ultraquant_flydsl_decode_eligible, +) +from vllm.v1.attention.ops.ultraquant.format import slot_size +from vllm.v1.attention.ops.ultraquant.reference import ( + hadamard_matrix, + reference_ultraquant_attention, +) +from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +def _build_paged_cache(key, value, block_size): + batch_size, seq_len, num_kv_heads, head_size = key.shape + blocks_per_seq = (seq_len + block_size - 1) // block_size + kv_cache = torch.zeros( + batch_size * blocks_per_seq, + block_size, + num_kv_heads, + slot_size(head_size), + dtype=torch.uint8, + device=key.device, + ) + block_table = torch.arange( + batch_size * blocks_per_seq, + dtype=torch.int32, + device=key.device, + ).view(batch_size, blocks_per_seq) + slot_mapping = torch.arange( + batch_size * blocks_per_seq * block_size, + dtype=torch.int64, + device=key.device, + ).view(batch_size, -1)[:, :seq_len] + ultraquant_store( + key.reshape(-1, num_kv_heads, head_size), + value.reshape(-1, num_kv_heads, head_size), + kv_cache, + slot_mapping.reshape(-1), + ) + return kv_cache, block_table + + +@pytest.mark.parametrize( + ("batch_size", "query_group_size", "block_size", "contiguous_query"), + [ + (1, 8, 32, True), + (8, 8, 64, True), + (9, 8, 256, True), + (1, 16, 32, True), + (1, 6, 16, True), + (1, 6, 32, True), + (2, 6, 64, True), + (1, 6, 128, True), + (1, 6, 256, True), + (1, 6, 32, False), + ], +) +def test_flydsl_decode_matches_reference( + batch_size, + query_group_size, + block_size, + contiguous_query, +): + if not ultraquant_flydsl_decode_eligible( + head_size=256, + num_kv_groups=query_group_size, + has_sinks=False, + sliding_window=None, + flydsl_loaded=is_ultraquant_flydsl_available(), + ): + pytest.skip("UltraQuant D=256 FlyDSL kernel is unavailable") + + set_random_seed(0x57A1D + batch_size + block_size + query_group_size) + device = torch.device("cuda") + dtype = torch.bfloat16 + num_kv_heads = 1 + head_size = 256 + seq_len = 2048 + scale = head_size**-0.5 + + key = torch.randn( + batch_size, + seq_len, + num_kv_heads, + head_size, + dtype=dtype, + device=device, + ) + value = torch.randn_like(key) + if contiguous_query: + qkv = torch.randn( + batch_size, + query_group_size + 2 * num_kv_heads, + head_size, + dtype=dtype, + device=device, + ) + query = qkv[:, :query_group_size] + else: + query_storage = torch.randn( + batch_size, + query_group_size, + head_size * 2, + dtype=dtype, + device=device, + ) + query = query_storage[..., ::2] + kv_cache, block_table = _build_paged_cache(key, value, block_size) + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) + hadamard = hadamard_matrix(head_size, device) + + actual = flydsl_ultraquant_decode_attention( + query=query, + kv_cache=kv_cache, + block_table=block_table, + seq_lens=seq_lens, + scale=scale, + PiT=hadamard, + max_seq_len=seq_len, + ) + expected = reference_ultraquant_attention( + query=query, + key=key, + value=value, + scale=scale, + ) + + actual_f32 = actual.float() + expected_f32 = expected.float() + cosine = torch.nn.functional.cosine_similarity( + actual_f32.flatten(), expected_f32.flatten(), dim=0 + ) + error = (actual_f32 - expected_f32).abs() + assert cosine > 0.995 + assert error.mean() < 0.05 + assert error.max() < 0.8 diff --git a/tests/kernels/ultraquant/test_seam_contract.py b/tests/kernels/ultraquant/test_seam_contract.py new file mode 100644 index 0000000000..a5f2116797 --- /dev/null +++ b/tests/kernels/ultraquant/test_seam_contract.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""UltraQuant inherits TurboQuant's forward/do_kv_cache_update orchestration. + +UltraQuant only overrides the store/decode/prefill seams, so a TurboQuant-side +change to a seam signature, to an inherited helper, or to the ``layer._tq_*`` +attributes ``forward`` reads will break UltraQuant at serve time without +touching any UltraQuant kernel test. These checks pin that contract. +""" + +import ast +import inspect +import textwrap + +from vllm.v1.attention.backends import ultraquant_attn +from vllm.v1.attention.backends.turboquant_attn import TurboQuantAttentionImpl +from vllm.v1.attention.backends.ultraquant_attn import UltraQuantAttentionImpl + +_SEAMS = ("_ensure_on_device", "_store_kv", "_decode_attention", "_prefill_attention") + + +def _method_ast(cls: type, name: str) -> ast.AST: + return ast.parse(textwrap.dedent(inspect.getsource(getattr(cls, name)))) + + +def _tq_attrs_read(cls: type, name: str) -> set[str]: + """``_tq_*`` attribute names read (not assigned) in ``cls.name``.""" + tree = _method_ast(cls, name) + assigned = { + t.attr + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + for t in node.targets + if isinstance(t, ast.Attribute) + } + read = { + node.attr + for node in ast.walk(tree) + if isinstance(node, ast.Attribute) and node.attr.startswith("_tq_") + } + return read - assigned + + +def _tq_attrs_written(cls: type, name: str) -> set[str]: + tree = _method_ast(cls, name) + return { + t.attr + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + for t in node.targets + if isinstance(t, ast.Attribute) and t.attr.startswith("_tq_") + } + + +def test_ultraquant_matches_turboquant_seam_contract(): + # 1. Each overridden seam accepts every parameter the base declares, so any + # call the inherited orchestration makes stays valid for UltraQuant. + for name in _SEAMS: + base = list( + inspect.signature(getattr(TurboQuantAttentionImpl, name)).parameters + ) + override = list( + inspect.signature(getattr(UltraQuantAttentionImpl, name)).parameters + ) + assert override[: len(base)] == base, ( + f"UltraQuantAttentionImpl.{name} no longer accepts the arguments " + f"TurboQuantAttentionImpl passes: base={base} override={override}" + ) + + # 2. Keyword arguments UltraQuant passes to *inherited* TurboQuant helpers + # must exist on those helpers (e.g. _flash_attn_varlen). + own = vars(UltraQuantAttentionImpl) + module = ast.parse(inspect.getsource(ultraquant_attn)) + for node in ast.walk(module): + if not isinstance(node, ast.Call): + continue + fn = node.func + if not ( + isinstance(fn, ast.Attribute) + and isinstance(fn.value, ast.Name) + and fn.value.id == "self" + ): + continue + target = getattr(TurboQuantAttentionImpl, fn.attr, None) + if target is None or fn.attr in own or not callable(target): + continue + params = inspect.signature(target).parameters + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + continue + for kw in node.keywords: + if kw.arg is None: + continue + assert kw.arg in params, ( + f"UltraQuant calls inherited {fn.attr}(..., {kw.arg}=...) but " + f"TurboQuantAttentionImpl.{fn.attr} accepts {list(params)}" + ) + + # 3. Every layer._tq_* attribute the inherited forward reads is set by + # UltraQuant's _ensure_on_device. + required = _tq_attrs_read(TurboQuantAttentionImpl, "forward") + provided = _tq_attrs_written(UltraQuantAttentionImpl, "_ensure_on_device") + assert required, "expected inherited forward to read layer._tq_* attributes" + assert required <= provided, ( + f"UltraQuant._ensure_on_device does not set {sorted(required - provided)}, " + f"which the inherited TurboQuant forward reads" + ) diff --git a/tests/kernels/ultraquant/test_store_vs_reference.py b/tests/kernels/ultraquant/test_store_vs_reference.py new file mode 100644 index 0000000000..2f02e9578c --- /dev/null +++ b/tests/kernels/ultraquant/test_store_vs_reference.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Triton UltraQuant store kernel. + +The Triton kernel must produce bit-identical codes and scale bytes vs +the PyTorch reference (``reference.ultraquant_encode``). +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.ultraquant.format import ( + k_codes_offset, + k_scales_offset, + n_groups, + slot_size, + v_codes_offset, + v_scales_offset, +) +from vllm.v1.attention.ops.ultraquant.reference import ultraquant_encode +from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +def _make_cache( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_dim: int, + device: torch.device, +) -> torch.Tensor: + slot_b = slot_size(head_dim) + return torch.zeros( + num_blocks, block_size, num_kv_heads, slot_b, dtype=torch.uint8, device=device + ) + + +def _read_codes_and_scales( + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + num_kv_heads: int, + head_dim: int, +): + """Pull (k_codes, k_scale_bytes, v_codes, v_scale_bytes) for each + (token, head) from the AoS slot layout. Scales come back as uint8 + (one E8M0 byte per group).""" + block_size = kv_cache.shape[1] + N = slot_mapping.shape[0] + Gk = n_groups(head_dim) + + k_codes_off = k_codes_offset(head_dim) + k_scales_off = k_scales_offset(head_dim) + v_codes_off = v_codes_offset(head_dim) + v_scales_off = v_scales_offset(head_dim) + codes_bytes = head_dim // 2 + + k_codes = torch.zeros( + N, num_kv_heads, codes_bytes, dtype=torch.uint8, device=kv_cache.device + ) + k_scale_bytes = torch.zeros( + N, num_kv_heads, Gk, dtype=torch.uint8, device=kv_cache.device + ) + v_codes = torch.zeros_like(k_codes) + v_scale_bytes = torch.zeros_like(k_scale_bytes) + + for ti in range(N): + slot = int(slot_mapping[ti].item()) + if slot < 0: + continue + blk = slot // block_size + off = slot % block_size + for h in range(num_kv_heads): + slot_bytes = kv_cache[blk, off, h] # [slot_size_aligned] uint8 + + k_codes[ti, h] = slot_bytes[k_codes_off : k_codes_off + codes_bytes] + k_scale_bytes[ti, h] = slot_bytes[k_scales_off : k_scales_off + Gk] + v_codes[ti, h] = slot_bytes[v_codes_off : v_codes_off + codes_bytes] + v_scale_bytes[ti, h] = slot_bytes[v_scales_off : v_scales_off + Gk] + + return k_codes, k_scale_bytes, v_codes, v_scale_bytes + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("head_dim", [128, 256]) +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +@pytest.mark.parametrize("num_tokens", [1, 16, 17]) +def test_store_matches_reference(dtype, head_dim, num_kv_heads, num_tokens): + """Triton store output (codes + E8M0 scale bytes) matches the reference + bit-exactly.""" + device = torch.device("cuda") + set_random_seed(0xF8) + + block_size = 32 + num_blocks = max(2, (num_tokens + block_size - 1) // block_size + 1) + + key = torch.randn(num_tokens, num_kv_heads, head_dim, dtype=dtype, device=device) + value = torch.randn(num_tokens, num_kv_heads, head_dim, dtype=dtype, device=device) + + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + if num_tokens >= 4: + slot_mapping[1] = -1 # skip-path test + + kv_cache = _make_cache(num_blocks, block_size, num_kv_heads, head_dim, device) + + ultraquant_store(key, value, kv_cache, slot_mapping) + torch.accelerator.synchronize() + + ref_k = ultraquant_encode(key, rotate=True) + ref_v = ultraquant_encode(value, rotate=False) + + k_codes, k_scale_bytes, v_codes, v_scale_bytes = _read_codes_and_scales( + kv_cache, slot_mapping, num_kv_heads, head_dim + ) + + if num_tokens >= 4: + assert (kv_cache[0, 1, :, :] == 0).all(), ( + "skipped slot's physical cache region was unexpectedly written" + ) + + for ti in range(num_tokens): + slot = int(slot_mapping[ti].item()) + if slot < 0: + continue + + assert torch.equal(k_codes[ti], ref_k.codes_packed[ti]), ( + f"K codes mismatch at token {ti}\n" + f"got: {k_codes[ti].cpu().numpy().tolist()[:16]}\n" + f"expected: {ref_k.codes_packed[ti].cpu().numpy().tolist()[:16]}" + ) + assert torch.equal(k_scale_bytes[ti], ref_k.scale_bytes[ti]), ( + f"K scale bytes mismatch at token {ti}\n" + f"got: {k_scale_bytes[ti].cpu().numpy().tolist()}\n" + f"expected: {ref_k.scale_bytes[ti].cpu().numpy().tolist()}" + ) + + assert torch.equal(v_codes[ti], ref_v.codes_packed[ti]), ( + f"V codes mismatch at token {ti}" + ) + assert torch.equal(v_scale_bytes[ti], ref_v.scale_bytes[ti]), ( + f"V scale bytes mismatch at token {ti}\n" + f"got: {v_scale_bytes[ti].cpu().numpy().tolist()}\n" + f"expected: {ref_v.scale_bytes[ti].cpu().numpy().tolist()}" + ) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 92b984c0fa..e6436f98f9 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -50,6 +50,7 @@ CacheDType = Literal[ "turboquant_4bit_nc", "turboquant_k3v4_nc", "turboquant_3bit_nc", + "ultraquant_4bit", "int4_per_token_head", "int8_per_token_head", "fp8_per_token_head", diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 98a0b7cd88..527b8956d0 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2271,7 +2271,13 @@ class EngineArgs: kv_offloading_backend=self.kv_offloading_backend, ) - if resolved_cache_dtype.startswith("turboquant_"): + # TurboQuant and UltraQuant both keep boundary attention layers at the + # native dtype, and compute those layers the same way. + uses_packed_kv_backend = ( + resolved_cache_dtype.startswith("turboquant_") + or resolved_cache_dtype == "ultraquant_4bit" + ) + if uses_packed_kv_backend: from vllm.model_executor.layers.quantization.turboquant.config import ( TurboQuantConfig, ) @@ -2681,14 +2687,18 @@ class EngineArgs: # TurboQuant requires FlashAttention 2 — FA3 boundary layers assert # FlashAttentionImpl which fails with TurboQuantAttentionImpl. - if resolved_cache_dtype.startswith("turboquant_") and ( + # UltraQuant subclasses that impl, so it inherits the same limit. + if uses_packed_kv_backend and ( attention_config.flash_attn_version is None or attention_config.flash_attn_version >= 3 ): logger.warning( - "TurboQuant is not yet compatible with FlashAttention >= 3. " + "%s is not yet compatible with FlashAttention >= 3. " "Overriding flash_attn_version to 2. To silence this " - "warning, pass --attention-config.flash_attn_version=2" + "warning, pass --attention-config.flash_attn_version=2", + "TurboQuant" + if resolved_cache_dtype.startswith("turboquant_") + else "UltraQuant", ) attention_config.flash_attn_version = 2 diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 57a7a82f21..f4583daba2 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -906,15 +906,29 @@ class Platform: cache_dtype_str=cache_config.cache_dtype, kv_quant_mode=kv_quant_mode, ).page_size_bytes - elif cache_config.cache_dtype.startswith("turboquant_"): - # TQ has a packed K|V layout; the standard FullAttentionSpec - # formula over-sizes it and trips unify_kv_cache_spec_page_size - # when all attention layers are TQ. With mixed skip+TQ the skip - # layers still use the standard layout — take max so mamba - # padding covers the largest actual page. - from vllm.v1.attention.backends.turboquant_attn import ( - TurboQuantAttentionBackend, - ) + elif ( + cache_config.cache_dtype.startswith("turboquant_") + or cache_config.cache_dtype == "ultraquant_4bit" + ): + # TurboQuant and UltraQuant have packed K|V layouts; the standard + # FullAttentionSpec formula over-sizes them and trips + # unify_kv_cache_spec_page_size when all attention layers are + # packed. With mixed skip+packed the skip layers still use the + # standard layout — take max so mamba padding covers the largest + # actual page. Each dtype packs via its own backend. + _pack_backend: type[AttentionBackend] + if cache_config.cache_dtype == "ultraquant_4bit": + from vllm.v1.attention.backends.ultraquant_attn import ( + UltraQuantAttentionBackend, + ) + + _pack_backend = UltraQuantAttentionBackend + else: + from vllm.v1.attention.backends.turboquant_attn import ( + TurboQuantAttentionBackend, + ) + + _pack_backend = TurboQuantAttentionBackend tq_spec = FullAttentionSpec( block_size=1, @@ -923,7 +937,7 @@ class Platform: dtype=kv_cache_dtype, kv_quant_mode=kv_quant_mode, ) - tq_page = TurboQuantAttentionBackend.customize_spec(tq_spec).page_size_bytes + tq_page = _pack_backend.customize_spec(tq_spec).page_size_bytes if cache_config.kv_cache_dtype_skip_layers: skip_page = FullAttentionSpec( block_size=1, diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 487cffa670..f394e56bf6 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -30,6 +30,7 @@ logger = init_logger(__name__) _KV_CACHE_DTYPE_REASON = "kv_cache_dtype not supported" _TURBOQUANT_LAYOUT_REASON = "no KV cache layout in common with TURBOQUANT" +_ULTRAQUANT_LAYOUT_REASON = "no KV cache layout in common with ULTRAQUANT" try: from amdsmi import ( @@ -505,6 +506,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.ULTRAQUANT) return backends @@ -517,6 +519,14 @@ def _uses_turboquant(vllm_config: "VllmConfig | None") -> bool: ) +def _uses_ultraquant(vllm_config: "VllmConfig | None") -> bool: + """Whether the run's KV cache dtype is the UltraQuant 4-bit format.""" + cache_config = getattr(vllm_config, "cache_config", None) + return ( + cache_config is not None and str(cache_config.cache_dtype) == "ultraquant_4bit" + ) + + def _shares_layout_with_turboquant(backend_class: type["AttentionBackend"]) -> bool: """Whether this backend reads a KV cache layout TURBOQUANT reads too. @@ -535,12 +545,29 @@ def _shares_layout_with_turboquant(backend_class: type["AttentionBackend"]) -> b return not set(layouts).isdisjoint(turboquant_layouts) +def _shares_layout_with_ultraquant(backend_class: type["AttentionBackend"]) -> bool: + """Whether this backend reads a KV cache layout ULTRAQUANT reads too. + + Same constraint as _shares_layout_with_turboquant: an ultraquant_4bit run + keeps its boundary layers at the native dtype, and one layout has to serve + the whole worker. + """ + layouts = backend_class.supported_kv_cache_layouts() + ultraquant_layouts = ( + AttentionBackendEnum.ULTRAQUANT.get_class().supported_kv_cache_layouts() + ) + if layouts is None or ultraquant_layouts is None: + return True + return not set(layouts).isdisjoint(ultraquant_layouts) + + def _get_invalid_reasons( backend_class: type["AttentionBackend"], device_capability: DeviceCapability, attn_selector_config: "AttentionSelectorConfig", *, is_turboquant_run: bool, + is_ultraquant_run: bool = False, ) -> list[str]: """Why this backend cannot serve the layer, empty when it can.""" invalid_reasons = backend_class.validate_configuration( @@ -553,6 +580,12 @@ def _get_invalid_reasons( and not _shares_layout_with_turboquant(backend_class) ): invalid_reasons = [_TURBOQUANT_LAYOUT_REASON] + if ( + not invalid_reasons + and is_ultraquant_run + and not _shares_layout_with_ultraquant(backend_class) + ): + invalid_reasons = [_ULTRAQUANT_LAYOUT_REASON] return invalid_reasons @@ -683,7 +716,7 @@ class RocmPlatform(Platform): @classmethod def get_attn_backend_cls( cls, - selected_backend: "AttentionBackendEnum", + selected_backend: "AttentionBackendEnum | None", attn_selector_config: "AttentionSelectorConfig", num_heads: int | None = None, ) -> str: @@ -695,13 +728,16 @@ class RocmPlatform(Platform): # Keep lazy: vllm.config imports current_platform during initialization. from vllm.config import get_current_vllm_config_or_none - is_turboquant_run = _uses_turboquant(get_current_vllm_config_or_none()) + vllm_config = get_current_vllm_config_or_none() + is_turboquant_run = _uses_turboquant(vllm_config) + is_ultraquant_run = _uses_ultraquant(vllm_config) try: sel_invalid_reasons = _get_invalid_reasons( selected_backend.get_class(), device_capability, attn_selector_config, is_turboquant_run=is_turboquant_run, + is_ultraquant_run=is_ultraquant_run, ) except ImportError: sel_invalid_reasons = ["ImportError"] @@ -711,21 +747,23 @@ class RocmPlatform(Platform): selected_backend.name, ) return selected_backend.get_path() - # Only tolerate the mismatch when turboquant is in play: boundary - # layers keep the native dtype while every other layer needs - # TURBOQUANT, so no single --attention-backend can serve every layer. - # For any other dtype the selection is genuinely invalid -> fail loud. + # Only tolerate the mismatch when turboquant or ultraquant is in play: + # boundary layers keep the native dtype while every other layer needs + # the packed backend, so no single --attention-backend can serve every + # layer. For any other dtype the selection is genuinely invalid -> fail + # loud. kv_dtype = attn_selector_config.kv_cache_dtype layer_is_turboquant = kv_dtype is not None and str(kv_dtype).startswith( "turboquant" ) - is_turboquant_fallback = ( - is_turboquant_run or layer_is_turboquant + is_packed_kv_fallback = ( + is_turboquant_run or layer_is_turboquant or is_ultraquant_run ) and sel_invalid_reasons in ( [_KV_CACHE_DTYPE_REASON], [_TURBOQUANT_LAYOUT_REASON], + [_ULTRAQUANT_LAYOUT_REASON], ) - if not is_turboquant_fallback: + if not is_packed_kv_fallback: raise ValueError( f"Selected backend {selected_backend} is not valid for " f"this configuration. Reason: {sel_invalid_reasons}" @@ -733,10 +771,11 @@ class RocmPlatform(Platform): # NOTE: pass a str (not the list) -- info_once hashes its args. logger.info_once( "Selected backend %s is incompatible with this layer (%s) of " - "the turboquant run; using the auto-selected per-layer backend. " + "the %s run; using the auto-selected per-layer backend. " "Reason: %s", selected_backend.name, attn_selector_config.attn_type, + "ultraquant" if is_ultraquant_run else "turboquant", str(sel_invalid_reasons), ) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 025941eeb6..dce5e84fce 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -50,6 +50,7 @@ STR_DTYPE_TO_TORCH_DTYPE = { "turboquant_4bit_nc": torch.uint8, "turboquant_k3v4_nc": torch.uint8, "turboquant_3bit_nc": 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 f28a331854..ad184f0913 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -152,6 +152,7 @@ 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" + 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 CUSTOM = None diff --git a/vllm/v1/attention/backends/ultraquant_attn.py b/vllm/v1/attention/backends/ultraquant_attn.py new file mode 100644 index 0000000000..b62b7a6830 --- /dev/null +++ b/vllm/v1/attention/backends/ultraquant_attn.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""UltraQuant 4-bit KV-cache attention backend. + +UltraQuant is a 4-bit KV-cache format (FP4 E2M1 codes + UE8M0 power-of-two +per-group-32 scales) with Hadamard-rotated keys. It reuses the TurboQuant +attention scaffolding (forward orchestration, metadata, cache-layout addressing, +CUDA-graph pre-warm) by subclassing, and overrides only the format-specific +seams: + + * store -> ultraquant_store (rotate K, pack FP4 + UE8M0 scales) + * decode -> FlyDSL D=256 scaled-MFMA kernel, Triton unified fallback + * continuation -> small chunk: Triton unified; large chunk: dequant + FA + +``UltraQuantAttentionBackend`` is a standalone backend with its own identity +(``ULTRAQUANT``) and dtype (``ultraquant_4bit``); the TurboQuant backend/impl +carry no UltraQuant knowledge. +""" + +import contextlib +import math +from dataclasses import replace +from typing import Any, ClassVar + +import torch +import torch.nn.functional as F + +from vllm.config import get_current_vllm_config +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.fa_utils import get_flash_attn_version +from vllm.v1.attention.backends.turboquant_attn import ( + _CONTINUATION_DECODE_THRESHOLD, + _HAS_FLASH_ATTN, + TurboQuantAttentionBackend, + TurboQuantAttentionImpl, + TurboQuantMetadata, + _build_hadamard, +) +from vllm.v1.attention.ops.flydsl_ultraquant_decode import ( + flydsl_ultraquant_decode_attention, + is_ultraquant_flydsl_available, + ultraquant_flydsl_decode_eligible, +) +from vllm.v1.attention.ops.ultraquant.format import slot_size +from vllm.v1.attention.ops.ultraquant.triton_unified_attention import ( + ultraquant_unified_attention, +) +from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + +logger = init_logger(__name__) + +# CK FlashAttention on ROCm has no compiled kernel above head_dim 256 and no +# sinks argument; those cases route through the Triton unified kernel instead. +_CK_MAX_HEAD_DIM = 256 + + +class UltraQuantAttentionBackend(TurboQuantAttentionBackend): + """Standalone backend for the ``ultraquant_4bit`` KV-cache dtype. + + Subclasses :class:`TurboQuantAttentionBackend` for the shared metadata and + cache-layout plumbing, and overrides identity, spec, impl, block size, and + capabilities so no UltraQuant knowledge lives in the TurboQuant backend. + """ + + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["ultraquant_4bit"] + # FlyDSL D=256 decode loads query as bf16; fp16 is not supported. Selection + # reads this from the backend class, not from the impl. + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + """UltraQuant packs K+V into one 4-bit FP4/UE8M0 slot per head.""" + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_ultraquant: + return spec + return replace(spec, state_content_bytes=slot_size(spec.head_size)) + + @staticmethod + def get_name() -> str: + return "ULTRAQUANT" + + @classmethod + def get_preferred_block_size(cls, default_block_size: int) -> int: + # FlyDSL D=256 decode geometry requires a 64-token KV block; pin it so + # the KV-cache manager pages stay dense and splittable into kernel blocks. + return 64 + + @staticmethod + def get_impl_cls() -> type["UltraQuantAttentionImpl"]: + return UltraQuantAttentionImpl + + @classmethod + def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: + return kv_cache_dtype == "ultraquant_4bit" + + @classmethod + def supports_sink(cls) -> bool: + # Decode and continuation prefill route sinks to the Triton unified + # kernel, but first-chunk prefill goes through flash_attn_varlen, which + # takes no sinks argument. Until that path is covered, reject sink + # models at selection rather than returning silently wrong output. + return False + + +class UltraQuantAttentionImpl(TurboQuantAttentionImpl): + """UltraQuant 4-bit KV-cache attention (FP4 E2M1 + UE8M0, rotated K). + + Overrides the init, on-device setup, store, decode, and prefill/continuation + seams of :class:`TurboQuantAttentionImpl`; the shared ``forward`` / + ``do_kv_cache_update`` orchestration and metadata handling are inherited. + """ + + # FlyDSL D=256 decode loads query as bf16; fp16 is not supported. + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + + 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, + **kwargs, + ): + self.num_heads = num_heads + self.head_size = head_size + self.scale = scale + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + + self.fa_version = get_flash_attn_version(head_size=head_size) + + vllm_config = get_current_vllm_config() + self.max_num_kv_splits = ( + vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph + ) + self.sliding_window = sliding_window + self.sinks = kwargs.get("sinks") + self._max_model_len = vllm_config.model_config.max_model_len + + self._use_ultraquant_flydsl = is_ultraquant_flydsl_available() + # FlyDSL decode eligibility is static per layer config; compute once. + self._flydsl_decode_eligible = ultraquant_flydsl_decode_eligible( + head_size=self.head_size, + num_kv_groups=self.num_kv_groups, + has_sinks=self.sinks is not None, + sliding_window=self.sliding_window, + flydsl_loaded=self._use_ultraquant_flydsl, + ) + + def _ensure_on_device(self, layer, device): + """Build the shared Hadamard rotation and pre-warm FlyDSL buffers. + + UltraQuant needs no Lloyd-Max centroids, but the inherited forward reads + ``layer._tq_centroids``, so an empty tensor is kept on the layer. + """ + if self._use_ultraquant_flydsl: + # CUDA-graph capture safety for the FlyDSL decode path on ROCm: + # pre-allocate index buffers and grow the workspace to its max size + # before capture so replay never lands on stale addresses. + _max_len = self._max_model_len + _already_ok = ( + hasattr(self, "_arange_cache") + and self._arange_cache.device.type == str(device).split(":")[0] + and self._arange_cache.shape[0] >= _max_len + 2 + ) + if not _already_ok: + self._arange_cache = torch.arange( + 0, _max_len + 2, device=device, dtype=torch.int32 + ) + if not hasattr(self, "_cu_2") or self._cu_2.device != torch.device(device): + self._cu_2 = torch.zeros(2, device=device, dtype=torch.int32) + if ( + is_workspace_manager_initialized() + and not current_workspace_manager().is_locked() + ): + B_max = self._max_capture_batch_size() + D = self.head_size + Hq = self.num_heads + S = self.max_num_kv_splits + _pre_warm_bytes = ( + B_max * Hq * (S * (D + 1) + D) * 4 # fp32 mid_o + fp32 lse + + B_max * Hq * D * 2 # query-dtype output (bf16 = 2 B) + + 512 # alignment padding + ) + with contextlib.suppress(AssertionError): + current_workspace_manager().get_simultaneous( + ((_pre_warm_bytes,), torch.uint8) + ) + + if not hasattr(layer, "_tq_cached"): + D = self.head_size + # Pure Hadamard: orthonormal + symmetric (H = H^T), enabling + # in-kernel butterfly fusion and trivial inverse for continuation. + H = _build_hadamard(D, str(device)) + layer._tq_PiT = H + layer._tq_Pi = H + layer._tq_Pi_half = H.to(torch.float16) + # Inherited TurboQuantAttentionImpl.forward reads _tq_centroids + # unconditionally; UltraQuant has no centroid table, so an empty + # tensor satisfies the read. + layer._tq_centroids = torch.empty(0, device=device, dtype=torch.float32) + layer._tq_cached = True + + def _store_kv( + self, + key: torch.Tensor, # (N, Hk, D) + value: torch.Tensor, # (N, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + slot_mapping: torch.Tensor, + layer: Any, + ): + """Rotate K by the Hadamard and pack K+V into FP4/UE8M0 slots.""" + from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + + ultraquant_store( + key, + value, + kv_cache, + slot_mapping, + PiT=layer._tq_PiT, + ) + + def _decode_attention( + self, + query: torch.Tensor, # (B, Hq, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + attn_metadata: TurboQuantMetadata, + Pi: torch.Tensor, + centroids: torch.Tensor, + PiT: torch.Tensor | None = None, + layer: torch.nn.Module | None = None, + ) -> torch.Tensor: + # Acquire shared decode scratch from the WorkspaceManager (UltraQuant + # only needs the output buffer; the mid_o/lse slots keep the triple + # allocation identical to the base decode for buffer reuse). + B = query.shape[0] + D = self.head_size + S = self.max_num_kv_splits + Hq = self.num_heads + output_buf = None + if is_workspace_manager_initialized(): + _mid_o, output_buf, _lse = current_workspace_manager().get_simultaneous( + ((B, Hq, S, D + 1), torch.float32), + ((B, Hq, D), query.dtype), + ((B, Hq), torch.float32), + ) + + if self._flydsl_decode_eligible: + return flydsl_ultraquant_decode_attention( + query=query, + kv_cache=kv_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + scale=self.scale, + PiT=PiT, + max_seq_len=attn_metadata.max_seq_len, + output_buf=output_buf, + buf_holder=layer, + max_num_kv_splits=self.max_num_kv_splits, + sinks=self.sinks, + ) + + logger.info_once( + "UltraQuant FlyDSL is ineligible for head_size=%s, GQA=%s, " + "sinks=%s, sliding_window=%s; using Triton fallback.", + self.head_size, + self.num_kv_groups, + self.sinks is not None, + self.sliding_window, + ) + return ultraquant_unified_attention( + query=query, + kv_cache=kv_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + query_start_loc=attn_metadata.query_start_loc, + scale=self.scale, + PiT=PiT, + output=output_buf, + max_query_len=1, + max_seq_len=attn_metadata.max_seq_len, + sinks=self.sinks, + sliding_window=self.sliding_window, + ) + + def _prefill_attention( + self, + query: torch.Tensor, # (N, Hq, D) + key: torch.Tensor, # (N, Hk, D) + value: torch.Tensor, # (N, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + attn_metadata: TurboQuantMetadata, + Pi: torch.Tensor, + centroids: torch.Tensor, + PiT: torch.Tensor | None = None, + layer: Any = None, + ) -> torch.Tensor: + N, Hq, D = query.shape + + # Fast path: flash_attn for first-chunk prefills (all K/V in batch). + # CK cannot serve head dims above 256, so those fall through to the + # per-request unified-attention path below. + if ( + _HAS_FLASH_ATTN + and attn_metadata.max_query_len == attn_metadata.max_seq_len + and D <= _CK_MAX_HEAD_DIM + ): + return self._flash_attn_varlen( + q=query, + k=key, + v=value, + cu_seqlens_q=attn_metadata.query_start_loc, + cu_seqlens_k=attn_metadata.query_start_loc, + max_seqlen_q=attn_metadata.max_query_len, + max_seqlen_k=attn_metadata.max_query_len, + ) + + # Continuation or no flash_attn: per-request attention. + Hk = key.shape[1] + use_gqa = Hk < Hq + query_start_loc = attn_metadata.query_start_loc + num_reqs = query_start_loc.shape[0] - 1 + + output = torch.zeros(N, Hq, D, device=query.device, dtype=query.dtype) + + if attn_metadata.query_start_loc_cpu is not None: + qsl = attn_metadata.query_start_loc_cpu.tolist() + else: + qsl = query_start_loc.tolist() + if attn_metadata.seq_lens_cpu is not None: + seq_lens_list = attn_metadata.seq_lens_cpu.tolist() + else: + seq_lens_list = attn_metadata.seq_lens.tolist() + + if not hasattr(self, "_cu_2"): + self._cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32) + _max_seq = attn_metadata.max_seq_len + _ac: torch.Tensor | None = getattr(self, "_arange_cache", None) + if _ac is None or _ac.shape[0] <= _max_seq: + _ac = torch.arange( + 0, _max_seq + 1, device=query.device, dtype=attn_metadata.seq_lens.dtype + ) + self._arange_cache = _ac + _arange_cache: torch.Tensor = _ac + + for i in range(num_reqs): + q_start = qsl[i] + q_end = qsl[i + 1] + q_len = q_end - q_start + if q_len <= 0: + continue + + seq_len = seq_lens_list[i] + q_seq = query[q_start:q_end] # (q_len, Hq, D) + k_seq = key[q_start:q_end] # (q_len, Hk, D) + v_seq = value[q_start:q_end] # (q_len, Hk, D) + + if q_len == seq_len: + # First-chunk prefill: all K/V are in the current batch. + if D > _CK_MAX_HEAD_DIM: + # Head dims CK cannot handle (Gemma 4 global layers, + # D=512). This kernel is generic over head size. + from vllm.v1.attention.ops.triton_unified_attention import ( + unified_attention, + ) + + out = torch.empty_like(q_seq) + cu_single = torch.tensor( + [0, q_len], dtype=torch.int32, device=query.device + ) + seq_lens_single = torch.tensor( + [q_len], dtype=torch.int32, device=query.device + ) + block_table_single = torch.zeros( + (1, 1), dtype=torch.int32, device=query.device + ) + unified_attention( + q=q_seq, + k=k_seq.unsqueeze(0), + v=v_seq.unsqueeze(0), + out=out, + cu_seqlens_q=cu_single, + max_seqlen_q=q_len, + seqused_k=seq_lens_single, + max_seqlen_k=q_len, + softmax_scale=self.scale, + causal=True, + window_size=(-1, -1), + block_table=block_table_single, + softcap=0.0, + q_descale=None, + k_descale=None, + v_descale=None, + sinks=None, + ) + elif _HAS_FLASH_ATTN: + self._cu_2[1:2] = q_len + cu = self._cu_2 + out = self._flash_attn_varlen( + q=q_seq, + k=k_seq, + v=v_seq, + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=q_len, + max_seqlen_k=q_len, + ) + else: + q_t = q_seq.transpose(0, 1).contiguous() + k_t = k_seq.transpose(0, 1).contiguous() + v_t = v_seq.transpose(0, 1).contiguous() + out = F.scaled_dot_product_attention( + q_t, + k_t, + v_t, + is_causal=True, + scale=self.scale, + enable_gqa=use_gqa, + ).transpose(0, 1) + output[q_start:q_end] = out.to(query.dtype) + else: + # Continuation chunk: prefix already stored (Hadamard-rotated + # FP4) by do_kv_cache_update. Small chunks reuse the Triton + # unified decode-shaped kernel; large chunks dequant the prefix + # once and run a dense flash-attention prefill. + cached_len = seq_len - q_len + if q_len <= _CONTINUATION_DECODE_THRESHOLD: + synth_seq_lens = _arange_cache[cached_len + 1 : seq_len + 1] + synth_bt = attn_metadata.block_table[i : i + 1].expand(q_len, -1) + out = ultraquant_unified_attention( + query=q_seq, + kv_cache=kv_cache, + block_table=synth_bt, + seq_lens=synth_seq_lens, + query_start_loc=_arange_cache[: q_len + 1], + scale=self.scale, + PiT=PiT, + max_query_len=1, + max_seq_len=seq_len, + sinks=self.sinks, + sliding_window=self.sliding_window, + ) + else: + out = self._ultraquant_continuation_prefill( + layer=layer, + query=q_seq, + key_chunk=k_seq, + val_chunk=v_seq, + kv_cache=kv_cache, + block_table=attn_metadata.block_table[i : i + 1], + cached_len=cached_len, + seq_len=seq_len, + PiT=PiT, + ) + output[q_start:q_end] = out.to(query.dtype) + + return output + + def _ultraquant_continuation_prefill( + self, + *, + layer: Any, + query: torch.Tensor, # (q_len, Hq, D) unrotated + key_chunk: torch.Tensor, # (q_len, Hk, D) unrotated + val_chunk: torch.Tensor, # (q_len, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + block_table: torch.Tensor, # (1, max_num_blocks) + cached_len: int, + seq_len: int, + PiT: torch.Tensor, # (D, D) fp32 Hadamard + ) -> torch.Tensor: + """Large continuation chunk: dequant the cached prefix, then flash_attn. + + The cache stores K already Hadamard-rotated, so only the current chunk's + K and the queries need rotating. Hadamard is symmetric orthonormal, so + rotating both sides leaves the attention scores unchanged. + """ + from vllm.v1.attention.ops.ultraquant.triton_dequant import ( + ultraquant_full_dequant_kv, + ) + + q_len, Hq, D = query.shape + Hk = key_chunk.shape[1] + device = query.device + qdtype = query.dtype + block_size = kv_cache.shape[1] + alloc_len = math.ceil(cached_len / block_size) * block_size + + # Dequant target, shared across layers by the workspace manager so long + # context does not pay one allocation per layer. + buf_shape = (1, Hk, alloc_len, D) + k_buf, v_buf = current_workspace_manager().get_simultaneous( + (buf_shape, qdtype), + (buf_shape, qdtype), + ) + ultraquant_full_dequant_kv( + kv_cache=kv_cache, + block_table=block_table, + k_out=k_buf, + v_out=v_buf, + alloc_len=alloc_len, + ) + + # Layer-cached prefix+chunk concatenation, grown once to worst case. + cap = block_table.shape[1] * block_size + k_full_buf = getattr(layer, "_uq_kfull_buf", None) + if ( + k_full_buf is None + or k_full_buf.shape[0] < cap + or k_full_buf.dtype != qdtype + ): + k_full_buf = torch.empty(cap, Hk, D, dtype=qdtype, device=device) + layer._uq_kfull_buf = k_full_buf + v_full_buf = getattr(layer, "_uq_vfull_buf", None) + if ( + v_full_buf is None + or v_full_buf.shape[0] < cap + or v_full_buf.dtype != qdtype + ): + v_full_buf = torch.empty(cap, Hk, D, dtype=qdtype, device=device) + layer._uq_vfull_buf = v_full_buf + k_full = k_full_buf[:seq_len] + v_full = v_full_buf[:seq_len] + + k_full[:cached_len] = k_buf[0, :, :cached_len, :].transpose(0, 1) + v_full[:cached_len] = v_buf[0, :, :cached_len, :].transpose(0, 1) + k_full[cached_len:] = (key_chunk.to(torch.float32) @ PiT).to(qdtype) + v_full[cached_len:] = val_chunk + q_rot = (query.to(torch.float32) @ PiT).to(qdtype) + + # Unequal q/k lengths with causal=True gives lower-right alignment: + # query i (absolute position cached_len + i) sees keys 0..cached_len+i. + # + # FA2 on ROCm has no sinks parameter, and CK cannot serve head dims + # above 256, so those cases route through the Triton unified kernel. + # Q/K are already in Hadamard-rotated space (q_rot, k_full) so the dot + # products match the original-space attention scores. + if self.sinks is not None or D > _CK_MAX_HEAD_DIM: + from vllm.v1.attention.ops.triton_unified_attention import ( + unified_attention, + ) + + out = torch.empty_like(q_rot) + cu_q = torch.tensor([0, q_len], dtype=torch.int32, device=device) + seqused_k = torch.tensor([seq_len], dtype=torch.int32, device=device) + block_table_single = torch.zeros((1, 1), dtype=torch.int32, device=device) + window = ( + (self.sliding_window - 1, 0) + if self.sliding_window and self.sliding_window > 0 + else (-1, -1) + ) + unified_attention( + q=q_rot, + k=k_full.unsqueeze(0), + v=v_full.unsqueeze(0), + out=out, + cu_seqlens_q=cu_q, + max_seqlen_q=q_len, + seqused_k=seqused_k, + max_seqlen_k=seq_len, + softmax_scale=self.scale, + causal=True, + window_size=window, + block_table=block_table_single, + softcap=0.0, + q_descale=None, + k_descale=None, + v_descale=None, + sinks=self.sinks, + ) + return out.to(qdtype) + if _HAS_FLASH_ATTN: + cu_q = torch.tensor([0, q_len], dtype=torch.int32, device=device) + cu_k = torch.tensor([0, seq_len], dtype=torch.int32, device=device) + return self._flash_attn_varlen( + q=q_rot, + k=k_full, + v=v_full, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=q_len, + max_seqlen_k=seq_len, + ) + q_t = q_rot.transpose(0, 1).unsqueeze(0) + k_t = k_full.transpose(0, 1).unsqueeze(0) + v_t = v_full.transpose(0, 1).unsqueeze(0) + if Hq > Hk: + k_t = k_t.expand(1, Hq, -1, -1) + v_t = v_t.expand(1, Hq, -1, -1) + return ( + F.scaled_dot_product_attention( + q_t, k_t, v_t, is_causal=True, scale=self.scale + ) + .squeeze(0) + .transpose(0, 1) + .to(qdtype) + ) diff --git a/vllm/v1/attention/ops/flydsl_kernels/ultraquant_decode_hd256.py b/vllm/v1/attention/ops/flydsl_kernels/ultraquant_decode_hd256.py new file mode 100644 index 0000000000..6a03f06f45 --- /dev/null +++ b/vllm/v1/attention/ops/flydsl_kernels/ultraquant_decode_hd256.py @@ -0,0 +1,983 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# UltraQuant D=256 decode (FlyDSL) — gfx950 / CDNA4 +# +# Production path only: scaled FP4×E4M3 QK MFMA, native V CVT, HW V-transpose, +# Q hoist, strided tile-groups, and in-kernel Walsh-Hadamard for GQA 6/8/16. +# +# AoS slot (D=256, group_size=32 → 8 groups, 272 B): +# [0:128) K FP4 codes; [128:136) K UE8M0; [136:264) V FP4; [264:272) V UE8M0 +# QK MFMA: A=K, B=Q → C[token, query]; PV MFMA: A=V_T, B=P → C[head_dim, query] + +from __future__ import annotations + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops, vector +from flydsl._mlir import ir +from flydsl._mlir.dialects import scf as _scf +from flydsl.expr import ( + arith, + gpu, + rocdl, +) +from flydsl.expr.primitive import const_expr, range_constexpr +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + + +def _vector_insert(value, dest, *, static_position): + """Insert a DSL scalar using FlyDSL 0.3's raw vector dialect binding.""" + return vector.insert( + vector.as_ir_value(value), + vector.as_ir_value(dest), + dynamic_position=[], + static_position=static_position, + ) + + +# === Constants (256-wide-head ultraquant decode profile) ======================= +HEAD_SIZE = 256 +KV_BLOCK_SIZE = 16 # default; overridable via build(kv_block_size=...) +TILE_SIZE = 16 # MFMA tile = 16 tokens +N_CENTROIDS = 16 # FP4 E2M1 value table (16 entries indexed by nibble) +QUERY_GROUP_SIZE = 16 # default for tests; overridable via build(...) +WARP_SIZE = 64 +BLOCK_THREADS = WARP_SIZE +KV_COMPUTE_BLOCK = 256 # 16 K-tiles × 16 tokens + +# ultraquant AoS slot layout (per (slot, head), D=256, group_size=32 => 8 groups) +FP8_GROUP_SIZE = 32 +N_GROUPS = HEAD_SIZE // FP8_GROUP_SIZE # 8 UE8M0 scale bytes per head +KEY_CODE_BYTES = HEAD_SIZE // 2 # 128 +VAL_CODE_BYTES = HEAD_SIZE // 2 # 128 +K_SCALES_OFFSET = KEY_CODE_BYTES # 128 +V_CODES_OFFSET = KEY_CODE_BYTES + N_GROUPS # 136 +V_SCALES_OFFSET = V_CODES_OFFSET + VAL_CODE_BYTES # 264 +SLOT_CONTENT_BYTES = V_SCALES_OFFSET + N_GROUPS # 272 + +# MFMA +MFMA_N = 16 +PV_N_CHUNKS = HEAD_SIZE // MFMA_N # 16 for HEAD_SIZE=256 + +# --- Per-lane data-movement geometry (generalized from the 128 kernel) ------ +# The K/V dequant lane assignment splits each token's packed code bytes across +# the 4 ``chunk_in_tok`` sub-lanes (lane % 4). Each lane processes +# ``KEY_CODE_BYTES // 4`` code bytes = 32 B at 256. buffer_load maxes at 16 B +# (dwordx4) so a lane issues ``HALVES``=2 loads; each 16-B half covers exactly +# one UE8M0 group of 32 head-dims. +# HEAD_SIZE=128: LANE_CODE_BYTES=16, HALVES=1, SUBCHUNK_HDIMS=32 +# HEAD_SIZE=256: LANE_CODE_BYTES=32, HALVES=2, SUBCHUNK_HDIMS=64 +LANE_CODE_BYTES = KEY_CODE_BYTES // 4 # 32 code bytes per lane per K-tile +HALVES = LANE_CODE_BYTES // 16 # 2 for HEAD_SIZE=256 +HALF_HDIMS = 32 # head-dims per half (= one group) +SUBCHUNK_HDIMS = HEAD_SIZE // 4 # head-dims per chunk_in_tok sub-lane (64) +# Q row-major LDS load geometry (STEP B): each lane loads 8 bf16 (=4 i32), +# so a row spans ``HEAD_SIZE // 8`` column-groups. +Q_GROUPS_PER_ROW = HEAD_SIZE // 8 # 32 for 256 +Q_ROW_SHIFT = Q_GROUPS_PER_ROW.bit_length() - 1 # 5 for 256 +Q_COL_MASK = Q_GROUPS_PER_ROW - 1 # 31 for 256 + +# LDS regions +CENTROID_LDS_BYTES = N_CENTROIDS * 4 # 64 + +# KV LDS rows are padded by 16 B to reduce bank conflicts while preserving +# ds_read/ds_write_b128 alignment. +KV_ROW_PAD_ELEMS = 8 +assert KV_ROW_PAD_ELEMS % 8 == 0, "KV row pad must be a multiple of 8 bf16 (16 B)" +KV_ROW_ELEMS = HEAD_SIZE + KV_ROW_PAD_ELEMS # 264 @ pad 8 +KV_ROW_BYTES = KV_ROW_ELEMS * 2 # 528 +KFP4_ROW_PAD_I32 = 4 +assert KFP4_ROW_PAD_I32 % 4 == 0, "FP4-K row pad must be a multiple of 4 i32 (16 B)" +KFP4_ROW_I32 = HEAD_SIZE // 8 + KFP4_ROW_PAD_I32 +KV_TILE_LDS_BYTES_PADDED = TILE_SIZE * KV_ROW_BYTES # 8448 (V dominates) +SCALE_LDS_BYTES = TILE_SIZE * N_GROUPS * 4 # 16*8*4 = 512 for 256 + +LOG2E = 1.4426950408889634 +NEG_INF_VAL = float("-inf") + + +def _vsplat_mul(vec, scalar): + s = scalar.ir_value() if hasattr(scalar, "ir_value") else scalar + return vec * vector.broadcast(T.f32x4, s) + + +allocator = None + + +def build_ultraquant_decode_hd256_module( + num_seqs: int, + num_kv_heads: int, + num_partitions: int, + padded_slot: int, + max_blocks_per_seq: int = 512, + softmax_scale: float | None = None, + query_group_size: int = QUERY_GROUP_SIZE, + kv_block_size: int = KV_BLOCK_SIZE, + tile_groups_per_partition: int = 1, + fuse_qrot: bool = True, + stride_q_seq: int | None = None, + stride_q_head: int | None = None, +): + """Build the production UltraQuant D=256 decode kernel. + + ``padded_slot`` is the per-(token, head) slot stride in bytes + (``kv_cache.shape[3]``, >= 272). Cache layout is + ``[num_blocks, block_size, num_kv_heads, padded_slot]`` uint8 AoS. + + QK is always scaled FP4×E4M3 MFMA; V always uses native CVT + HW + transpose; tile-groups are always strided. ``fuse_qrot`` selects the + in-kernel WHT when the query layout is eligible. + """ + # QG=6 (Qwen3.6 full-attention layers) reuses the 8-row WHT lane map with + # two inactive rows masked. STEP C may read rows QG..15 from adjacent LDS, + # but those MFMA rows are discarded by the `mfma_row < QG` output gate. + assert query_group_size in (6, 8, 16), ( + f"query_group_size must be 6, 8 or 16; got {query_group_size}" + ) + # head_dim=256 hybrid models (Qwen3.6/3.8) force a larger mamba-aligned KV + # block; the tiled inner loop (_TILES_PER_BLOCK = kv_block_size // TILE_SIZE) + # handles any multiple of TILE_SIZE, so accept 16/32/64/128/256. + assert kv_block_size in (16, 32, 64, 128, 256), ( + f"kv_block_size must be 16/32/64/128/256; got {kv_block_size}" + ) + assert kv_block_size % TILE_SIZE == 0 + assert int(tile_groups_per_partition) >= 1 + + TGPP = int(tile_groups_per_partition) + # Partition p owns tile-groups p, p+P, p+2P, ... This keeps the graph-fixed + # grid wide for short sequences; the reducer is agnostic to the mapping. + NUM_PARTS_C = int(num_partitions) + MFMA_SCALED_K = 128 + MFMA_ISSUES = HEAD_SIZE // MFMA_SCALED_K # 2 for HEAD_SIZE=256 + GRPS_PER_ISSUE = MFMA_SCALED_K // FP8_GROUP_SIZE # 4 + + QG = int(query_group_size) + QG_LOAD_ITERS = (QG * Q_GROUPS_PER_ROW) // WARP_SIZE # 8@QG16, 4@QG8, 3@QG6 + + _BLOCK_THREADS = WARP_SIZE + + # Fuse Q rotation (q_rot = Q @ PiT) into this kernel as an in-register WHT. + # The 8-row lane map masks its two inactive rows for GQA-6. + _FUSE_QROT = bool(fuse_qrot) + + # --- Cache policy on the KV code stream ---------------------------------- + # KV codes are read EXACTLY ONCE per decode (81.6 MB at B=4/seq=75k) and + # never reused, so caching them only evicts Q / block-table / scale lines + # that ARE reused. gfx9 buffer aux bits: bit0=sc0, bit1=nt, bit4=sc1. + # 0 = default (what we ship today: no hint at all) + # 2 = nt (non-temporal: stream, do not pollute) + # 18 = nt|sc1 (also bypass the far cache level) + _KV_AUX = 0 + + # --- Occupancy: size the Q LDS to the ACTUAL query group ------------------- + # STEP B only writes rows [0, QG) and STEP C/F only read those rows. For + # QG=8 the upper 4 KB is allocated-but-never-touched, so shrinking it to + # QG*HEAD_SIZE*2 is bitwise identical and frees LDS/workgroup (17 KB -> 13 KB + # at QG=8), lifting waves/CU (~3 -> ~5) since the kernel is LDS-occupancy + # bound. The launcher pads segm pools to QG=16 regardless, so this is purely + # an in-kernel LDS footprint reduction. + _Q_LDS_BYTES = QG * HEAD_SIZE * 2 + + _BS = int(kv_block_size) + _TILES_PER_BLOCK = _BS // TILE_SIZE + _PADDED_SLOT = int(padded_slot) + assert _PADDED_SLOT >= SLOT_CONTENT_BYTES, ( + f"padded_slot={_PADDED_SLOT} < {SLOT_CONTENT_BYTES} (ultraquant D=256 content)" + ) + assert _PADDED_SLOT % 4 == 0, f"padded_slot={_PADDED_SLOT} must be 4-byte aligned" + + arch = get_hip_arch() + + if softmax_scale is None: + softmax_scale = 1.0 / (HEAD_SIZE**0.5) + _qk_scale = float(softmax_scale) + + # --- Strides --- + _Hq = num_kv_heads * QG + # Defaults describe the pooled [B, Hq, D] q_rot buffer. When the rotation is + # fused in (STEP B'), Q is read straight from the model's query tensor, + # which is a SPLIT OF THE FUSED QKV PROJECTION -- so its row stride is + # (Hq + 2*Hk)*D, not Hq*D. Those strides are constant for a given model and + # TP rank, so the caller passes them in and they bake in like any other + # compile-time constant (they are part of the kernel cache key). + _stride_q_seq = int(stride_q_seq) if stride_q_seq is not None else _Hq * HEAD_SIZE + _stride_q_head = int(stride_q_head) if stride_q_head is not None else HEAD_SIZE + _stride_bt_seq = max_blocks_per_seq + + # ultraquant AoS cache: [num_blocks, block_size, num_kv_heads, padded_slot]. + _stride_cache_head = _PADDED_SLOT + _stride_cache_pos = num_kv_heads * _PADDED_SLOT + _stride_cache_block = _BS * _stride_cache_pos + + _stride_out_part = QG * HEAD_SIZE + _stride_out_head = num_partitions * QG * HEAD_SIZE + _stride_out_seq = num_kv_heads * num_partitions * QG * HEAD_SIZE + _stride_es_seq = num_kv_heads * num_partitions * QG + _stride_ml_seq = _stride_es_seq + + # --- LDS layout --- + # `allocator` is intentionally a BUILD-LOCAL captured by the kernel closure + # (a freevar), NOT a module global. The @flyc.jit drift guard + # (_check_globals_drift) snapshots every module global referenced by the JIT + # launcher's dependency tree (which includes this kernel) and ABORTS if one + # changes between compiles. When more than one partition variant is + # JIT-resident at once (e.g. a ragged/adaptive batch mixes num_partitions + # 256 and 512, or fixed cap=512 yields 294@75k and 512@131k), each build + # created a fresh SmemAllocator; as a reassigned MODULE GLOBAL that tripped + # "global 'allocator' changed since first compile" and crashed the worker. + # Keeping it local + a per-variant symbol name lets the variants coexist. + # The build-local is handed to the launcher via a module attribute set AFTER + # the kernel is built (see end of function); that assignment is safe because + # no @flyc.jit function references the module global anymore. + allocator = SmemAllocator( + None, + arch=arch, + global_sym_name=(f"ultraquant_hd256_smem_p{int(num_partitions)}"), + ) + centroid_off = 0 + allocator.ptr = CENTROID_LDS_BYTES + q_off = allocator.ptr + allocator.ptr += _Q_LDS_BYTES + kv_off = allocator.ptr + allocator.ptr += KV_TILE_LDS_BYTES_PADDED + scale_off = allocator.ptr + allocator.ptr += SCALE_LDS_BYTES + + @flyc.kernel + def ultraquant_decode_hd256_kernel( + out_ptr: fx.Tensor, + exp_sums_ptr: fx.Tensor, + max_logits_ptr: fx.Tensor, + query_ptr: fx.Tensor, + kv_cache_ptr: fx.Tensor, + centroids_ptr: fx.Tensor, + block_tables_ptr: fx.Tensor, + seq_lens_ptr: fx.Tensor, + ): + # ---- IDs --------------------------------------------------------- + tid = gpu.thread_idx.x + seq = gpu.block_idx.x + kv_h = gpu.block_idx.y + part = gpu.block_idx.z + lane = tid # 0..63 + mfma_row = lane & fx.Int32(15) + mfma_col_grp = lane >> fx.Int32(4) # 0..3, K-group dim + + # ---- Buffer resources ------------------------------------------- + q_rsrc = buffer_ops.create_buffer_resource(query_ptr, max_size=True) + bt_rsrc = buffer_ops.create_buffer_resource(block_tables_ptr, max_size=True) + sl_rsrc = buffer_ops.create_buffer_resource(seq_lens_ptr, max_size=True) + cent_rsrc = buffer_ops.create_buffer_resource(centroids_ptr, max_size=True) + out_rsrc = buffer_ops.create_buffer_resource(out_ptr, max_size=True) + es_rsrc = buffer_ops.create_buffer_resource(exp_sums_ptr, max_size=True) + ml_rsrc = buffer_ops.create_buffer_resource(max_logits_ptr, max_size=True) + + # ---- LDS pointers ----------------------------------------------- + base = allocator.get_base() + cent_lds = SmemPtr(base, centroid_off, T.f32, shape=(N_CENTROIDS,)) + q_lds_i32 = SmemPtr(base, q_off, T.i32, shape=(_Q_LDS_BYTES // 4,)).get() + q_lds_i64 = SmemPtr(base, q_off, T.i64, shape=(_Q_LDS_BYTES // 8,)).get() + # KV / scale views span ALL warps' private regions (W*stride); the + # per-warp base offset is folded into the element indices below via + # _kvi*/_sci (a literal 0 for W=1, so the single-warp path is unchanged). + kv_lds_i32 = SmemPtr( + base, kv_off, T.i32, shape=(KV_TILE_LDS_BYTES_PADDED // 4,) + ).get() + kv_lds_i64 = SmemPtr( + base, kv_off, T.i64, shape=(KV_TILE_LDS_BYTES_PADDED // 8,) + ).get() + scale_lds_i32 = SmemPtr(base, scale_off, T.i32, shape=(TILE_SIZE * N_GROUPS,)) + + def _kvi32(e): + return e + + def _kvi64(e): + return e + + def _sci(e): + return e + + # ---- Constants --------------------------------------------------- + c_sq = fx.Int32(_stride_q_seq) + c_qh = fx.Int32(_stride_q_head) + c_qg = fx.Int32(QG) + c_bt = fx.Int32(_stride_bt_seq) + c_stride_pos = fx.Int32(_stride_cache_pos) + c_stride_head = fx.Int32(_stride_cache_head) + c_kscale_off = fx.Int32(K_SCALES_OFFSET) # 128 + c_vcode_off = fx.Int32(V_CODES_OFFSET) # 136 + c_vscale_off = fx.Int32(V_SCALES_OFFSET) # 264 + c_w = fx.Int32(WARP_SIZE) + + NEG_INF = arith.constant(NEG_INF_VAL, type=T.f32) + ZERO_F = fx.Float32(0.0) + ONE_F = fx.Float32(1.0) + LOG2E_C = arith.constant(LOG2E, type=T.f32) + QK_SCALE = arith.constant(_qk_scale, type=T.f32) + + def _ival(v): + return v.ir_value() if hasattr(v, "ir_value") else v + + c_zero_i32 = arith.constant(0, type=T.i32) + + def _f32x8_to_fp8_i64(f): + # 8 f32 -> 8 E4M3 bytes packed as i64 (MFMA fp8 operand). + w0 = rocdl.cvt_pk_fp8_f32(T.i32, f[0], f[1], c_zero_i32, 0) + w0 = rocdl.cvt_pk_fp8_f32(T.i32, f[2], f[3], w0, 1) + w1 = rocdl.cvt_pk_fp8_f32(T.i32, f[4], f[5], c_zero_i32, 0) + w1 = rocdl.cvt_pk_fp8_f32(T.i32, f[6], f[7], w1, 1) + pv = vector.from_elements(T.vec(2, T.i32), [w0, w1]) + return vector.extract( + vector.bitcast(T.vec(1, T.i64), pv), static_position=[0] + ) + + def _bf16x8_to_fp8_i64(v_bf16): + f = [ + arith.extf(T.f32, vector.extract(v_bf16, static_position=[i])) + for i in range(8) + ] + return _f32x8_to_fp8_i64(f) + + # ===== STEP A: Load centroids → LDS (cooperative, race-safe) ===== + c_idx_safe = lane & fx.Int32(N_CENTROIDS - 1) + c_val = buffer_ops.buffer_load(cent_rsrc, c_idx_safe, vec_width=1, dtype=T.f32) + cent_lds.store(c_val, [arith.index_cast(T.index, c_idx_safe)]) + gpu.barrier() + + # ===== STEP B': in-kernel Q rotation (fused WHT) ================== + # The shipped path needs a SEPARATE prologue kernel to compute + # q_rot = Q @ PiT. That kernel is ~7 us at B=4 -- essentially ALL fixed + # per-dispatch cost (measured: 2.1M MACs, and sweeping its grid 4x moves + # it 0%). PiT is the orthonormal Sylvester Hadamard with a column + # permutation folded in on the host, so Q @ PiT is a WALSH-HADAMARD + # TRANSFORM: 8 butterfly stages (2048 add/sub per row) instead of a + # 256x256 matmul (65536 MACs) -- 32x less work, cheap enough to redo + # redundantly in every partition and delete the dispatch entirely. + # + # Verified offline: butterfly vs matmul differs by <=1.8e-6 in fp32, and + # after the E4M3 haircut the two agree EXACTLY (0/2048 elements differ), + # because e4m3 keeps only 3 mantissa bits. So this is bit-exact, not an + # approximation. + # + # Lane map: lane L owns row r = L>>3 and the 32 CONSECUTIVE head-dims + # d = (L&7)*32 + t. That puts butterfly stages 0..4 (strides 1..16) + # entirely inside one lane's registers, and leaves only stages 5..7 + # (strides 32/64/128 == chunk XOR 1/2/4) needing a cross-lane exchange. + if const_expr(_FUSE_QROT): + _wrow = lane >> fx.Int32(3) # 0..7 row within the group + _wchk = lane & fx.Int32(7) # 0..7 32-dim chunk + # inv(qperm) is a pure bit permutation of the head-dim index: + # dest = (d & 0x8F) | ((d>>2)&0x10) | ((d<<1)&0x60) + # which splits additively into a compile-time part in t and a + # per-lane part in the chunk index (verified exhaustively): + # ct(t) = (t & 0x0F) | ((t & 0x10) << 1) + # rt(c) = ((c&2)<<3) | ((c&1)<<6) | ((c&4)<<5) + _perm_base = ( + ((_wchk & fx.Int32(2)) << fx.Int32(3)) + | ((_wchk & fx.Int32(1)) << fx.Int32(6)) + | ((_wchk & fx.Int32(4)) << fx.Int32(5)) + ) + _inv_sqrtD = arith.constant(1.0 / math.sqrt(HEAD_SIZE), type=T.f32) + for _qit in range_constexpr((QG + 7) // 8): + _qrow = _wrow + fx.Int32(_qit * 8) + if _qrow < fx.Int32(QG): + _qb = ( + seq * c_sq + (kv_h * c_qg + _qrow) * c_qh + _wchk * fx.Int32(32) + ) + # ---- load this lane's 32 raw-Q head-dims (4 x dwordx4) ---- + _v = [] + for _j in range_constexpr(4): + _raw = buffer_ops.buffer_load( + q_rsrc, + (_qb + fx.Int32(_j * 8)) // fx.Int32(2), + vec_width=4, + dtype=T.i32, + ) + _rb = vector.bitcast(T.vec(8, T.bf16), _raw) + for _e in range_constexpr(8): + _v.append( + arith.extf( + T.f32, vector.extract(_rb, static_position=[_e]) + ) + ) + # ---- stages 0..4: in-register, no LDS, no cross-lane ---- + for _s in range_constexpr(5): + _st = 1 << _s + _nx = list(_v) + for _g in range_constexpr(32 // (2 * _st)): + for _k in range_constexpr(_st): + _i0 = _g * 2 * _st + _k + _i1 = _i0 + _st + _a, _b = _v[_i0], _v[_i1] + _nx[_i0] = _a + _b + _nx[_i1] = _a - _b + _v = _nx + # ---- stages 5..7: chunk XOR 1/2/4 via ds_swizzle ---- + # ds_swizzle_b32 offset (bit15=0, 32-lane groups): + # [4:0]=and_mask [9:5]=or_mask [14:10]=xor_mask + # pure XOR k -> and_mask=0x1F, or_mask=0, xor_mask=k. The + # partner is always within the same 32-lane group (k<=4). + for _s in range_constexpr(3): + _k = 1 << _s + _swz = arith.constant((_k << 10) | 0x1F, type=T.i32) + # High half of the pair computes (partner - mine), the low + # half (mine + partner). Fold that into a single sign + # computed ONCE per stage, so each element costs one + # multiply-add rather than a branch or a per-element select. + _chk_bit = _wchk & fx.Int32(_k) + _is_hi = arith.cmpi( + arith.CmpIPredicate.ne, + _chk_bit.ir_value() + if hasattr(_chk_bit, "ir_value") + else _chk_bit, + fx.Int32(0).ir_value(), + ) + _sf = arith.select( + _is_hi, + arith.constant(-1.0, type=T.f32), + arith.constant(1.0, type=T.f32), + ) + _nx = [] + for _t in range_constexpr(32): + _mi = arith.bitcast(T.i32, _v[_t]) + _pt = arith.bitcast( + T.f32, rocdl.ds_swizzle(T.i32, _mi, _swz) + ) + _nx.append(_v[_t] * _sf + _pt) + _v = _nx + # ---- 1/sqrt(D) normalization + E4M3 haircut -> bf16 ---- + # cvt_pk_fp8_f32 is the SAME instruction the QK operand build + # uses, so this reproduces the host haircut exactly. E4M3 is a + # subset of bf16, so the bf16 store below is lossless and every + # downstream consumer stays byte-identical. + _hc = [] + for _p in range_constexpr(16): + _w = rocdl.cvt_pk_fp8_f32( + T.i32, + _v[2 * _p] * _inv_sqrtD, + _v[2 * _p + 1] * _inv_sqrtD, + fx.Int32(0), + 0, + ) + _u = rocdl.cvt_pk_f32_fp8(T.vec(2, T.f32), _w, 0) + _hc.append( + arith.truncf( + T.bf16, vector.extract(_u, static_position=[0]) + ) + ) + _hc.append( + arith.truncf( + T.bf16, vector.extract(_u, static_position=[1]) + ) + ) + # ---- permuted store ---- + # ct(t) maps t=0..15 -> 0..15 and t=16..31 -> 32..47, i.e. TWO + # contiguous 16-element runs, so the scatter is still 4 aligned + # 16-byte vector stores rather than 32 scalar writes. + _qrow_elem = _qrow * fx.Int32(HEAD_SIZE) + for _half in range_constexpr(2): + for _q4 in range_constexpr(2): + _elems = [_hc[_half * 16 + _q4 * 8 + _e] for _e in range(8)] + _vec = vector.from_elements(T.vec(8, T.bf16), _elems) + _dst = ( + _qrow_elem + _perm_base + fx.Int32(_half * 32 + _q4 * 8) + ) + vector.store( + vector.bitcast(T.vec(4, T.i32), _vec), + q_lds_i32, + [arith.index_cast(T.index, _dst // fx.Int32(2))], + ) + gpu.barrier() + + # ===== STEP B: Load pre-rotated q_rot → row-major Q LDS ========== + for c in range_constexpr(0 if _FUSE_QROT else QG_LOAD_ITERS): + row_chunk = lane + fx.Int32(c * WARP_SIZE) + row = row_chunk >> fx.Int32(Q_ROW_SHIFT) # 0..QG-1 + col_b = row_chunk & fx.Int32(Q_COL_MASK) # 0..Q_GROUPS_PER_ROW-1 + col_elem = col_b * fx.Int32(8) # bf16 elem + q_off_elem = seq * c_sq + (kv_h * c_qg + row) * c_qh + col_elem + q_v = buffer_ops.buffer_load( + q_rsrc, + q_off_elem // fx.Int32(2), + vec_width=4, + dtype=T.i32, + ) + q_lds_byte = row * fx.Int32(HEAD_SIZE * 2) + col_elem * fx.Int32(2) + vector.store( + q_v, + q_lds_i32, + [arith.index_cast(T.index, q_lds_byte // fx.Int32(4))], + ) + gpu.barrier() + + # Scaled-MFMA Q operands (loop-invariant): one per MFMA issue. + q_op_hoisted = [] + for h in range_constexpr(MFMA_ISSUES): + _q_hoist_words = [] + for j in range_constexpr(4): + _q_hoist_idx = ( + mfma_row * fx.Int32(HEAD_SIZE * 2 // 8) + + (fx.Int32(h * GRPS_PER_ISSUE) + mfma_col_grp) * fx.Int32(8) + + fx.Int32(j * 2) + ) + _q_hoist_v = vector.load_op( + T.vec(2, T.i64), + q_lds_i64, + [arith.index_cast(T.index, _q_hoist_idx)], + ) + _q_hoist_words.append( + _bf16x8_to_fp8_i64(vector.bitcast(T.vec(8, T.bf16), _q_hoist_v)) + ) + q_op_hoisted.append( + vector.bitcast( + T.vec(8, T.i32), + vector.from_elements(T.vec(4, T.i64), _q_hoist_words), + ) + ) + + # ===== STEP D: Online softmax + PV state ========================= + running_max = NEG_INF + running_sum = ZERO_F + zero_v4 = arith.constant_vector(0.0, T.f32x4) + acc_pv = [zero_v4 for _ in range(PV_N_CHUNKS)] + + # ===== STEP E: Sequence-len + partition base ==================== + seq_len = buffer_ops.buffer_load(sl_rsrc, seq, vec_width=1, dtype=T.i32) + + # Per-K-tile dequant lane assignment: lane t -> token = t/4, + # chunk_in_tok = t%4 (each chunk = SUBCHUNK_HDIMS=64 head-dims = 2 groups). + tok_in_tile = lane >> fx.Int32(2) + chunk_in_tok = lane & fx.Int32(3) + # Both UE8M0 group bytes for this lane's chunk (groups 2c, 2c+1) live in + # the SAME 4-byte scale word at offset (chunk_in_tok>>1)*4 within the + # 8-byte scale region. Load once; extract per-half below. + c_scale_word_off = (chunk_in_tok >> fx.Int32(1)) * fx.Int32(4) + + # ===== STEP F: K-tile loop ======================================= + c_kcb = fx.Int32(KV_COMPUTE_BLOCK) + c_tgpp = fx.Int32(TGPP) + c_zero_i32 = fx.Int32(0) + c_one_i32 = fx.Int32(1) + c_nparts = fx.Int32(NUM_PARTS_C) + total_tgs = (seq_len + c_kcb - c_one_i32) // c_kcb + avail = total_tgs - part + in_range = avail > c_zero_i32 + trip_raw = (avail + c_nparts - c_one_i32) // c_nparts + trip_clamped = (trip_raw > c_tgpp).select(c_tgpp, trip_raw) + trip_or_zero = in_range.select(trip_clamped, c_zero_i32) + c_zero_idx = arith.constant(0, index=True) + c_one_idx = arith.constant(1, index=True) + trip_idx = arith.index_cast( + T.index, + trip_or_zero.ir_value() + if hasattr(trip_or_zero, "ir_value") + else trip_or_zero, + ) + _init_iter = [ + _ival(running_max), + _ival(running_sum), + *[_ival(p) for p in acc_pv], + ] + _for_op = _scf.ForOp( + c_zero_idx, + trip_idx, + c_one_idx, + _init_iter, + ) + _for_ip = ir.InsertionPoint(_for_op.body) + _for_ip.__enter__() + try: + tg_idx = _for_op.induction_variable + tg_i32 = fx.Int32(arith.index_cast(T.i32, tg_idx)) + partition_start = (part + tg_i32 * c_nparts) * c_kcb + bt_seq_base = seq * c_bt + (partition_start // fx.Int32(_BS)) + running_max = _for_op.inner_iter_args[0] + running_sum = _for_op.inner_iter_args[1] + acc_pv = list(_for_op.inner_iter_args[2:]) + + # ==== SOFTWARE PREFETCH ===================================== + # Emit ONLY the HBM buffer_loads for one tile (K codes, V codes, + # K/V UE8M0 scale words) and return the loaded register values. + # Emitting tile N+1's loads BEFORE tile N's dequant+MFMA lets the + # vmem fetches overlap compute: the compiler's s_waitcnt for these + # values lands at their first use one iteration later, hiding HBM + # latency (the decode kernel is ~85% HBM-stalled at low batch). + # This is bit-identical to the pre-prefetch path -- only the + # *emission order* of independent loads changes, never any + # arithmetic. All addressing locals (slot_base_byte, blk_rsrc, + # ...) are consumed only inside this closure; the dequant below + # needs just the four returned register values. + def _emit_tile_loads(n_tile): + # n_tile is a compile-time int on the single-warp path and a + # runtime fx.Int32 on the multi-warp path (each warp strides the + # 16 tiles). Compute the block/tile/slot geometry accordingly. + if const_expr(isinstance(n_tile, int)): + block_in_part = fx.Int32(n_tile // _TILES_PER_BLOCK) + tile_in_block = n_tile % _TILES_PER_BLOCK + tile_start_tok = partition_start + fx.Int32(n_tile * TILE_SIZE) + slot = fx.Int32(tile_in_block * TILE_SIZE) + tok_in_tile + else: + block_in_part = n_tile // fx.Int32(_TILES_PER_BLOCK) + tile_in_block = n_tile - block_in_part * fx.Int32(_TILES_PER_BLOCK) + tile_start_tok = partition_start + n_tile * fx.Int32(TILE_SIZE) + slot = tile_in_block * fx.Int32(TILE_SIZE) + tok_in_tile + tile_in_seq = tile_start_tok < seq_len + bt_off = bt_seq_base + block_in_part + bt_off_safe = tile_in_seq.select(bt_off, seq * c_bt) + phys_block = buffer_ops.buffer_load( + bt_rsrc, + bt_off_safe, + vec_width=1, + dtype=T.i32, + ) + # A buffer descriptor addresses with a 32-bit voffset, so + # ``phys_block * c_block`` wraps once the cache view exceeds + # 4 GiB. Fold the block base into the descriptor's 64-bit + # base pointer and keep only in-block offsets below. + blk_off_i64 = arith.extui(T.i64, phys_block) * fx.Int64( + _stride_cache_block + ) + blk_rsrc = buffer_ops.create_buffer_resource( + kv_cache_ptr, + max_size=True, + base_byte_offset=blk_off_i64, + ) + block_base = fx.Int32(0) + + slot_base_byte = block_base + slot * c_stride_pos + kv_h * c_stride_head + # K codes: LANE_CODE_BYTES (=32) at slot_base + chunk*32, in + # HALVES x 16-byte buffer_loads. + k_byte0 = slot_base_byte + chunk_in_tok * fx.Int32(LANE_CODE_BYTES) + _kpl = [] + for hf in range_constexpr(HALVES): + _kpl.append( + buffer_ops.buffer_load( + blk_rsrc, + (k_byte0 + fx.Int32(hf * 16)) // fx.Int32(4), + vec_width=4, + dtype=T.i32, + cache_modifier=_KV_AUX, + ) + ) + + # ---- V codes + K/V scale HBM loads ------------------------ + v_byte0 = ( + slot_base_byte + + c_vcode_off + + chunk_in_tok * fx.Int32(LANE_CODE_BYTES) + ) + _vpl = [] + for hf in range_constexpr(HALVES): + _vpl.append( + buffer_ops.buffer_load( + blk_rsrc, + (v_byte0 + fx.Int32(hf * 16)) // fx.Int32(4), + vec_width=4, + dtype=T.i32, + cache_modifier=_KV_AUX, + ) + ) + # UE8M0 group scale words (one 4-byte word covers this lane's + # two groups). scale = 2^(byte-127) = bitcast_f32(byte<<23); + # byte==0 -> +0.0 (zero sentinel). + _ksw = buffer_ops.buffer_load( + blk_rsrc, + (slot_base_byte + c_kscale_off + c_scale_word_off) // fx.Int32(4), + vec_width=1, + dtype=T.i32, + ) + _vsw = buffer_ops.buffer_load( + blk_rsrc, + (slot_base_byte + c_vscale_off + c_scale_word_off) // fx.Int32(4), + vec_width=1, + dtype=T.i32, + ) + return (_kpl, _vpl, _ksw, _vsw) + + # Depth-1 software prefetch: emit the next tile's HBM loads before + # consuming the current tile without changing arithmetic order. + _PF_DEPTH = 1 + # Prime the pipeline with the first _PF_DEPTH tiles' loads. + _pf_queue = [] + for _pi in range_constexpr(min(_PF_DEPTH, 16)): + _pf_queue.append(_emit_tile_loads(_pi)) + _tile_range = range_constexpr(16) + for _jt in _tile_range: + n_tile = _jt + _ntile_tok = fx.Int32(n_tile * TILE_SIZE) + (k_packed_list, v_packed_list, kscale_word, vscale_word) = ( + _pf_queue.pop(0) + ) + _pf_nxt = n_tile + _PF_DEPTH + if const_expr(_pf_nxt < 16): + _pf_queue.append(_emit_tile_loads(_pf_nxt)) + # Store raw FP4 codes (natural contiguous order, + # 16 B = 32 nibbles per UE8M0 group) straight to KV LDS as + # the scaled-MFMA A operand — no dequant here. + for hf in range_constexpr(HALVES): + grp = chunk_in_tok * fx.Int32(2) + fx.Int32(hf) + grp_shift = (grp & fx.Int32(3)) * fx.Int32(8) + kscale_byte = (kscale_word >> grp_shift) & fx.Int32(0xFF) + k_packed = k_packed_list[hf] + vector.store( + k_packed, + kv_lds_i32, + [ + arith.index_cast( + T.index, + _kvi32( + tok_in_tile * fx.Int32(KFP4_ROW_I32) + + grp * fx.Int32(4) + ), + ) + ], + ) + scale_lds_i32.store( + kscale_byte, + [ + arith.index_cast( + T.index, + _sci(tok_in_tile * fx.Int32(N_GROUPS) + grp), + ) + ], + ) + gpu.barrier() + # QK: MFMA_ISSUES native scaled MFMAs over K=128, + # chained through the accumulator. + # A = K raw FP4 codes (vec4 i32 = 16 B) per (token, group) + # B = Q E4M3 (vec8 i32 = 32 B) for the same group + # scaleA = per-(token,group) UE8M0 byte at op_sel byte 0 + # scaleB = 0x7F identity (Q carries no scale) + BLGP_E2M1 = 4 # A operand type code = fp4 e2m1 + CBSZ_E4M3 = 0 # B operand type code = fp8 e4m3 + IDENT = fx.Int32(0x7F) + qk_acc = zero_v4 + for h in range_constexpr(MFMA_ISSUES): + grp = fx.Int32(h * GRPS_PER_ISSUE) + mfma_col_grp + k_op = vector.load_op( + T.vec(4, T.i32), + kv_lds_i32, + [ + arith.index_cast( + T.index, + _kvi32( + mfma_row * fx.Int32(KFP4_ROW_I32) + + grp * fx.Int32(4) + ), + ) + ], + ) + q_op = q_op_hoisted[h] + scbyte = fx.Int32( + scale_lds_i32.load( + [ + arith.index_cast( + T.index, + _sci(mfma_row * fx.Int32(N_GROUPS) + grp), + ) + ] + ) + ) + kscale = fx.Int32(0x7F7F7F00) | scbyte + qk_acc = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), + [ + k_op, + q_op, + qk_acc, + BLGP_E2M1, + CBSZ_E4M3, + 0, + kscale, + 0, + IDENT, + ], + ) + + # Scale + mask out-of-context tokens. + qk_acc = _vsplat_mul(qk_acc, QK_SCALE) + for elem in range_constexpr(4): + kv_tok = ( + partition_start + + _ntile_tok + + mfma_col_grp * fx.Int32(4) + + fx.Int32(elem) + ) + in_b = kv_tok < seq_len + v = vector.extract(qk_acc, static_position=[elem]) + qk_acc = _vector_insert( + in_b.select(v, NEG_INF), + qk_acc, + static_position=[elem], + ) + + # FA2 online softmax: per-query-row reduce. + local_max = vector.reduction(T.f32, "maxnumf", qk_acc) + r1 = local_max.shuffle_xor(fx.Int32(16), c_w) + local_max = local_max.maximumf(r1) + r2 = local_max.shuffle_xor(fx.Int32(32), c_w) + tile_max = local_max.maximumf(r2) + + new_max = running_max.maximumf(tile_max) + max_diff = running_max - new_max + safe_diff = (running_max > NEG_INF).select(max_diff, ZERO_F) + scale = (safe_diff * LOG2E_C).exp2(fastmath=arith.FastMathFlags.fast) + running_sum = running_sum * scale + for h in range_constexpr(PV_N_CHUNKS): + acc_pv[h] = _vsplat_mul(acc_pv[h], scale) + running_max = new_max + + tile_sum = ZERO_F + for elem in range_constexpr(4): + s = vector.extract(qk_acc, static_position=[elem]) + d = s - new_max + d = (new_max > NEG_INF).select(d, NEG_INF) + p = (d * LOG2E_C).exp2(fastmath=arith.FastMathFlags.fast) + tile_sum = tile_sum + p + qk_acc = _vector_insert(p, qk_acc, static_position=[elem]) + + ts1 = tile_sum.shuffle_xor(fx.Int32(16), c_w) + tile_sum = tile_sum + ts1 + ts2 = tile_sum.shuffle_xor(fx.Int32(32), c_w) + tile_sum = tile_sum + ts2 + running_sum = running_sum + tile_sum + # V dequant → LDS [token][head_dim] (ROW-MAJOR, no transpose) + # written as ds_write_b128 (vec(2,i64)); read back via the + # ds_read_tr16_b64 HW transpose in the PV block below. + v_lds_elem_base = tok_in_tile * fx.Int32( + KV_ROW_ELEMS + ) + chunk_in_tok * fx.Int32(SUBCHUNK_HDIMS) + for hf in range_constexpr(HALVES): + v_packed = v_packed_list[hf] + half_hd = fx.Int32(hf * HALF_HDIMS) + grp_shift = ( + (chunk_in_tok * fx.Int32(2) + fx.Int32(hf)) & fx.Int32(3) + ) * fx.Int32(8) + vscale_byte = (vscale_word >> grp_shift) & fx.Int32(0xFF) + vscale_f32 = arith.bitcast( + T.f32, _ival(vscale_byte << fx.Int32(23)) + ) + for w in range_constexpr(4): + word_i32 = vector.extract(v_packed, static_position=[w]) + # Native CDNA4 scaled convert: 2-wide + # cvt_scalef32_pk_bf16_fp4. srcSel s picks byte s + # of the word (nibbles 2s, 2s+1) → 2 bf16 with the + # UE8M0 group scale fused. 4 cvt calls/word. + bf16_elems = [] + for s in range_constexpr(4): + v2 = rocdl.cvt_scalef32_pk_bf16_fp4( + T.vec(2, T.bf16), + _ival(word_i32), + _ival(vscale_f32), + int(s), + ) + bf16_elems.append(vector.extract(v2, static_position=[0])) + bf16_elems.append(vector.extract(v2, static_position=[1])) + v_bf16 = vector.from_elements(T.vec(8, T.bf16), bf16_elems) + v_i64 = vector.bitcast(T.vec(2, T.i64), v_bf16) + v_lds_i64_idx = ( + v_lds_elem_base + half_hd + fx.Int32(w * 8) + ) // fx.Int32(4) + vector.store( + v_i64, + kv_lds_i64, + [arith.index_cast(T.index, _kvi64(v_lds_i64_idx))], + ) + # HW V transpose cross-lane LDS fence (see 128 kernel). + rocdl.sched_barrier(0) + rocdl.s_waitcnt(0xC07F) + gpu.barrier() + + # ---- PV MFMA: A=V[head_dim, token], B=P (=qk_acc bf16) ----- + p_bf16 = arith.trunc_f(T.vec(4, T.bf16), qk_acc) + p_op = vector.bitcast(T.vec(4, T.i16), p_bf16) + # HW-transpose PV: V_lds row-major V[token][head_dim]. + token_idx = lane >> fx.Int32(2) + hd_sub = (lane & fx.Int32(3)) * fx.Int32(4) + v_lane_byte = ( + fx.Int32(kv_off) + + token_idx * fx.Int32(KV_ROW_BYTES) + + hd_sub * fx.Int32(2) + ) + for h in range_constexpr(PV_N_CHUNKS): + v_byte_off = v_lane_byte + fx.Int32(h * 32) + v_byte_i64 = fx.Int64(v_byte_off) + v_ptr = buffer_ops.create_llvm_ptr( + v_byte_i64, + address_space=3, + ) + v_op_raw = rocdl.ds_read_tr16_b64( + T.vec(4, T.i16), + v_ptr, + ).result + acc_pv[h] = rocdl.mfma_f32_16x16x16bf16_1k( + T.f32x4, [v_op_raw, p_op, acc_pv[h], 0, 0, 0] + ) + + _scf.YieldOp( + [ + _ival(running_max), + _ival(running_sum), + *[_ival(p) for p in acc_pv], + ] + ) + finally: + _for_ip.__exit__(None, None, None) + running_max = _for_op.results[0] + running_sum = _for_op.results[1] + acc_pv = list(_for_op.results[2:]) + + # ===== STEP G: Output =========================================== + safe_sum = (running_sum > ZERO_F).select(running_sum, ONE_F) + rcp = ONE_F / safe_sum + + c_os = fx.Int32(_stride_out_seq) + c_oh = fx.Int32(_stride_out_head) + c_op_ = fx.Int32(_stride_out_part) + out_base = seq * c_os + kv_h * c_oh + part * c_op_ + + valid_row_pred = arith.cmpi( + arith.CmpIPredicate.ult, + mfma_row.ir_value() if hasattr(mfma_row, "ir_value") else mfma_row, + arith.constant(QG, type=T.i32), + ) + _if = _scf.IfOp(valid_row_pred) + with ir.InsertionPoint(_if.then_block): + for h in range_constexpr(PV_N_CHUNKS): + pv_norm = _vsplat_mul(acc_pv[h], rcp) + pv_bf16 = arith.trunc_f(T.vec(4, T.bf16), pv_norm) + pv_i32x2 = vector.bitcast(T.vec(2, T.i32), pv_bf16) + head_dim_start = fx.Int32(h * 16) + mfma_col_grp * fx.Int32(4) + out_off_elem = ( + out_base + mfma_row * fx.Int32(HEAD_SIZE) + head_dim_start + ) + buffer_ops.buffer_store( + pv_i32x2, + out_rsrc, + out_off_elem * fx.Int32(2), + offset_is_bytes=True, + ) + + c_npq = fx.Int32(num_partitions * QG) + ml_off = ( + seq * fx.Int32(_stride_ml_seq) + kv_h * c_npq + part * c_qg + mfma_row + ) + es_off = ( + seq * fx.Int32(_stride_es_seq) + kv_h * c_npq + part * c_qg + mfma_row + ) + buffer_ops.buffer_store(running_max, ml_rsrc, ml_off) + buffer_ops.buffer_store(running_sum, es_rsrc, es_off) + _scf.YieldOp([]) + + # Hand the build-local allocator to the launcher. `_get_kernel` reads + # `kmod.allocator` immediately after this returns (before the next build can + # run), so each launcher captures its own variant's allocator. This does NOT + # reintroduce the drift bug: the kernel above closes over the LOCAL + # `allocator`, so no @flyc.jit function references module-global `allocator` + # and it is never snapshotted by _check_globals_drift. + globals()["allocator"] = allocator + return ultraquant_decode_hd256_kernel diff --git a/vllm/v1/attention/ops/flydsl_ultraquant_decode.py b/vllm/v1/attention/ops/flydsl_ultraquant_decode.py new file mode 100644 index 0000000000..7f8c720678 --- /dev/null +++ b/vllm/v1/attention/ops/flydsl_ultraquant_decode.py @@ -0,0 +1,834 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlyDSL launcher for the optimized UltraQuant 4-bit D=256 decode kernel. + +The production kernel uses scaled FP4×E4M3 QK MFMA, native V conversion, +query hoisting, in-kernel Walsh-Hadamard rotation for GQA-6/8/16, and strided +tile-group scheduling. These choices are fixed and are not environment-tunable. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from vllm.logger import init_logger +from vllm.triton_utils import triton + +logger = init_logger(__name__) + +_FLYDSL_AVAILABLE: bool | None = None +_ULTRAQUANT_MOD: Any = None +_FLYC: Any = None +_FX: Any = None +_TYPING_T: Any = None +_CC: Any = None +_IR: Any = None + + +def is_ultraquant_flydsl_available() -> bool: + """Return whether the gfx950 UltraQuant FlyDSL kernel can be loaded.""" + global _FLYDSL_AVAILABLE, _ULTRAQUANT_MOD + global _FLYC, _FX, _TYPING_T, _CC, _IR + if _FLYDSL_AVAILABLE is not None: + return _FLYDSL_AVAILABLE + try: + from vllm.platforms.rocm import on_gfx950 + + if not on_gfx950(): + _FLYDSL_AVAILABLE = False + return False + + import flydsl.compiler as flyc + import flydsl.expr as fx + from flydsl._mlir import ir + from flydsl.compiler.kernel_function import CompilationContext + from flydsl.expr.typing import T + + from vllm.v1.attention.ops.flydsl_kernels import ( + ultraquant_decode_hd256 as ultraquant_mod, + ) + + _FLYC = flyc + _FX = fx + _TYPING_T = T + _CC = CompilationContext + _IR = ir + _ULTRAQUANT_MOD = ultraquant_mod + _FLYDSL_AVAILABLE = True + logger.info_once("FlyDSL UltraQuant D=256 decode is available") + except Exception as ex: # noqa: BLE001 + _FLYDSL_AVAILABLE = False + logger.warning_once( + "FlyDSL UltraQuant decode is unavailable (%s); using Triton fallback.", + ex, + ) + return _FLYDSL_AVAILABLE + + +_ULTRAQUANT_FLYDSL_GQA = (6, 8, 16) + + +def ultraquant_flydsl_decode_eligible( + *, + head_size: int, + num_kv_groups: int, + has_sinks: bool, + sliding_window: int | None, + flydsl_loaded: bool, +) -> bool: + """Return whether decode should use the gfx950 FlyDSL kernel. + + Production FlyDSL covers D=256 and GQA in {6, 8, 16}. Sinks and sliding + window fall back to unified Triton. ``flydsl_loaded`` is the process-level + compiler/hardware probe so this policy can be unit-tested without a GPU. + """ + return ( + flydsl_loaded + and head_size == 256 + and num_kv_groups in _ULTRAQUANT_FLYDSL_GQA + and not has_sinks + and not (sliding_window and sliding_window > 0) + ) + + +# -- Kernel module cache ------------------------------------------------------- +# Compiling a FlyDSL kernel is expensive; cache by the constexprs that +# parameterize stride math: num_kv_heads, num_partitions, max_blocks_per_seq, +# scale, query_group_size, kv_block_size, padded_slot, hw_v_transpose, +# tile_groups_per_partition. +_KERN_CACHE: dict[tuple, Any] = {} + +_LOG_INVOKED_ONCE: bool = False + +# Per-device cached FP4 E2M1 value table (used as the kernel "centroids" LUT). +_FP4_LUT_CACHE: dict[tuple[str, torch.dtype], torch.Tensor] = {} + + +def _fp4_value_lut(device: torch.device) -> torch.Tensor: + """Return the fixed 16-entry FP4 E2M1 value table as fp32 [16] on device. + + Indexed by the stored 4-bit E2M1 code, so ``lut[nibble]`` is the dequant + value directly (matches ``format.FP4_BITS_TO_VALUE``). + """ + from vllm.v1.attention.ops.ultraquant.format import FP4_BITS_TO_VALUE + + key = (str(device), torch.float32) + t = _FP4_LUT_CACHE.get(key) + if t is None: + t = torch.tensor( + FP4_BITS_TO_VALUE, dtype=torch.float32, device=device + ).contiguous() + _FP4_LUT_CACHE[key] = t + return t + + # Per-device cached qperm head-dim permutation for scaled QK MFMA. + + +# qperm[phys] = the natural head-dim placed at physical operand position phys. +# It realigns the fixed FP4xFP8 scaled-MFMA contraction so that, with K fed +# native-contiguous, Q's columns pair correctly. Folded into PiT on the host +# (PiT[:, qperm]) so the Q rotation matmul emits already-permuted q_rot at zero +# extra runtime cost. Derived + validated in tmp_qk_scaled_mfma.py (contigK). +_QPERM_CACHE: dict[str, torch.Tensor] = {} + + +def _kmap_phys(group: int, s: int) -> int: + # Physical operand position of group's element s for the K=128 fp4/fp8 + # scaled MFMA (matches tmp_qk_scaled_mfma.py::_kmap). + return 32 * (s >> 4) + 64 * (group % 2) + 16 * (group // 2) + (s & 15) + + +def _qperm_index( + device: torch.device, D: int = 128, group_size: int = 32 +) -> torch.Tensor: + # _kmap_phys describes ONE K=128 contraction (4 groups of 32). A head dim + # wider than 128 is covered by ceil(D/128) back-to-back MFMA issues, each + # contracting its own independent 128-wide slice, so the same map repeats + # per slice: qperm = [qperm128, qperm128 + 128, ...]. Applying _kmap_phys + # with group >= 4 instead yields positions outside the slice and is not a + # permutation at all. + key = f"{device}:{D}:{group_size}" + t = _QPERM_CACHE.get(key) + if t is None: + k_per_mfma = 128 + assert D % k_per_mfma == 0, ( + f"scaled QK MFMA needs D to be a multiple of {k_per_mfma}, got {D}" + ) + groups_per_mfma = k_per_mfma // group_size + qperm = [0] * D + for half in range(D // k_per_mfma): + base = half * k_per_mfma + for g in range(groups_per_mfma): + for s in range(group_size): + qperm[base + _kmap_phys(g, s)] = base + g * group_size + s + assert sorted(qperm) == list(range(D)), "qperm is not a permutation" + t = torch.tensor(qperm, dtype=torch.long, device=device).contiguous() + _QPERM_CACHE[key] = t + return t + + +def _detect_max_capture_B() -> int: + try: + from vllm.config import get_current_vllm_config + + cfg = get_current_vllm_config() + candidates: list[int] = [] + sizes = cfg.compilation_config.cudagraph_capture_sizes + if sizes: + candidates.append(int(max(sizes))) + sched = getattr(cfg, "scheduler_config", None) + if sched is not None and getattr(sched, "max_num_seqs", None): + candidates.append(int(sched.max_num_seqs)) + if candidates: + return max(candidates) + except Exception: # noqa: BLE001 + pass + return 512 + + +class _SegmBufPool: + """Single-bucket buffer pool for segm_out/segm_max/segm_sum/output + the + pooled Q-rotation intermediates. Identical to the TQ pool except that + ultraquant needs an extra ``q_fp8`` (float8_e4m3fn) slot for the Q precision + haircut (``q_rot_fp32 -> fp8_e4m3 -> bf16``) so that no fresh allocation + lands in the HIP graph memory pool post-capture. + """ + + __slots__ = ("_bufs", "_max_B") + + def __init__(self) -> None: + self._bufs: dict[tuple, dict[str, torch.Tensor]] = {} + self._max_B: int | None = None + + def get( + self, + B: int, + Hk: int, + Hq: int, + num_partitions: int, + QG: int, + D: int, + device: torch.device, + q_dtype: torch.dtype, + ) -> dict[str, torch.Tensor]: + if self._max_B is None: + self._max_B = _detect_max_capture_B() + B_bucket = max(self._max_B, int(B)) + key = ( + int(Hk), + int(Hq), + int(num_partitions), + int(QG), + int(D), + str(device), + q_dtype, + ) + bufs = self._bufs.get(key) + if bufs is None or bufs["segm_out"].shape[0] < B_bucket: + if bufs is not None and bufs["segm_out"].shape[0] < B_bucket: + # First-time grow: warn so user knows cudagraphs may be + # invalidated. This should only happen during warmup or in + # always-eager configs (a batch larger than any captured + # cudagraph size). + logger.warning_once( + "FlyDSL UltraQuant buffer pool growing from %d " + "to %d (B=%d). If this happens AFTER cudagraph warmup the " + "captured graphs hold stale pointers and will GPU-fault.", + bufs["segm_out"].shape[0], + B_bucket, + B, + ) + bufs = { + "segm_out": torch.empty( + (B_bucket, Hk, num_partitions, QG, D), + dtype=torch.bfloat16, + device=device, + ), + "segm_max": torch.empty( + (B_bucket, Hk, num_partitions, QG), + dtype=torch.float32, + device=device, + ), + "segm_sum": torch.empty( + (B_bucket, Hk, num_partitions, QG), + dtype=torch.float32, + device=device, + ), + "output": torch.empty( + (B_bucket, Hq, D), + dtype=q_dtype, + device=device, + ), + "q_rot": torch.empty( + (B_bucket, Hq, D), + dtype=q_dtype, + device=device, + ), + "q_float": torch.empty( + (B_bucket, Hq, D), + dtype=torch.float32, + device=device, + ), + "q_rot_fp32": torch.empty( + (B_bucket, Hq, D), + dtype=torch.float32, + device=device, + ), + # ultraquant-specific: E4M3 haircut intermediate (pooled so the + # cast does not allocate a fresh tensor post-capture). + "q_fp8": torch.empty( + (B_bucket, Hq, D), + dtype=torch.float8_e4m3fn, + device=device, + ), + } + self._bufs[key] = bufs + self._max_B = B_bucket + logger.info_once( + "FlyDSL UltraQuant buffer pool allocated " + "shape=(Hk=%d, Hq=%d, P=%d, QG=%d, D=%d, dtype=%s) " + "B_bucket=%d. VRAM = %.1f MiB / shape.", + Hk, + Hq, + num_partitions, + QG, + D, + q_dtype, + B_bucket, + sum(t.numel() * t.element_size() for t in bufs.values()) / (1 << 20), + ) + return { + "segm_out": bufs["segm_out"][:B], + "segm_max": bufs["segm_max"][:B], + "segm_sum": bufs["segm_sum"][:B], + "output": bufs["output"][:B], + "q_rot": bufs["q_rot"][:B], + "q_float": bufs["q_float"][:B], + "q_rot_fp32": bufs["q_rot_fp32"][:B], + "q_fp8": bufs["q_fp8"][:B], + } + + +_SEGM_POOL = _SegmBufPool() + +_GET_KERNEL_STATS = {"hits": 0, "misses": 0, "build_total_s": 0.0} + + +def _get_kernel( + num_kv_heads: int, + num_partitions: int, + max_blocks_per_seq: int, + scale: float, + query_group_size: int, + kv_block_size: int, + padded_slot: int, + fuse_qrot: bool, + num_seqs_hint: int, + tile_groups_per_partition: int, + stride_q_seq: int, + stride_q_head: int, +): + key = ( + num_kv_heads, + num_partitions, + max_blocks_per_seq, + round(scale, 8), + query_group_size, + kv_block_size, + padded_slot, + fuse_qrot, + tile_groups_per_partition, + stride_q_seq, + stride_q_head, + ) + cached = _KERN_CACHE.get(key) + if cached is not None: + _GET_KERNEL_STATS["hits"] += 1 + return cached + + assert is_ultraquant_flydsl_available() + import time + + build_start = time.perf_counter() + kmod = _ULTRAQUANT_MOD + assert kmod is not None + kfn = kmod.build_ultraquant_decode_hd256_module( + num_seqs=num_seqs_hint, + num_kv_heads=num_kv_heads, + num_partitions=num_partitions, + padded_slot=padded_slot, + max_blocks_per_seq=max_blocks_per_seq, + softmax_scale=scale, + query_group_size=query_group_size, + kv_block_size=kv_block_size, + tile_groups_per_partition=tile_groups_per_partition, + fuse_qrot=fuse_qrot, + stride_q_seq=stride_q_seq, + stride_q_head=stride_q_head, + ) + allocator = kmod.allocator + flyc = _FLYC + fx = _FX + T = _TYPING_T + CompilationContext = _CC + ir_mod = _IR + + @flyc.jit + def _launch( + out, + es, + ml, + q, + kvc, + cents, + bt, + sl, + gx: fx.Int32, # type: ignore[name-defined] + gy: fx.Int32, # type: ignore[name-defined] + gz: fx.Int32, # type: ignore[name-defined] + stream: fx.Stream, # type: ignore[name-defined] + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir_mod.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + from flydsl.expr import arith + + grid_x = arith.index_cast(T.index, gx.ir_value()) + grid_y = arith.index_cast(T.index, gy.ir_value()) + grid_z = arith.index_cast(T.index, gz.ir_value()) + kfn(out, es, ml, q, kvc, cents, bt, sl).launch( + grid=(grid_x, grid_y, grid_z), + block=(kmod.BLOCK_THREADS, 1, 1), + stream=stream, + ) + + _KERN_CACHE[key] = _launch + build_time = time.perf_counter() - build_start + _GET_KERNEL_STATS["misses"] += 1 + _GET_KERNEL_STATS["build_total_s"] += build_time + logger.info( + "Built FlyDSL UltraQuant kernel #%d in %.2fs (total %.1fs), key=%s", + _GET_KERNEL_STATS["misses"], + build_time, + _GET_KERNEL_STATS["build_total_s"], + key, + ) + return _launch + + +# -- Partition reducer --------------------------------------------------------- +from vllm.triton_utils import tl # noqa: E402 + + +@triton.jit +def _reduce_partitions( + output_ptr, # [N, Hq, D] in OUT_DTYPE + segm_out_ptr, # [N, Hk, P, QG, D] bf16 + segm_max_ptr, # [N, Hk, P, QG] fp32 + segm_sum_ptr, # [N, Hk, P, QG] fp32 + out_stride_n: tl.int64, + out_stride_h: tl.int64, + NUM_KV_HEADS: tl.constexpr, + QG: tl.constexpr, + NUM_PARTS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, +): + n = tl.program_id(0) + hq = tl.program_id(1) + kv_h = hq // QG + qg = hq % QG + # Splitting the head-dim across programs is what gives this kernel enough + # parallelism to use the GPU: the reduction runs over PARTITIONS, so every + # head-dim column is independent. A (B, Hq) grid is only 32 workgroups at + # B=4 -- ~12% of a 256-CU part -- which is why it moved 8.4 MB at just + # ~995 GB/s. The partition reduction itself is untouched (same values, same + # order), so the result is bit-identical; only the columns each program + # owns change. The [P] max/sum vectors get re-read per D-block, which is + # 4 KB of duplicated traffic against a 262 KB tile. + + msum_base = n * (NUM_KV_HEADS * NUM_PARTS * QG) + kv_h * (NUM_PARTS * QG) + qg + so_base = ( + n * (NUM_KV_HEADS * NUM_PARTS * QG * HEAD_SIZE) + + kv_h * (NUM_PARTS * QG * HEAD_SIZE) + + qg * HEAD_SIZE + ) + + p_off = tl.arange(0, NUM_PARTS) + d_off = tl.program_id(2) * BLOCK_D + tl.arange(0, BLOCK_D) + + m_idx = msum_base + p_off * QG + seg_max = tl.load(segm_max_ptr + m_idx) + seg_sum = tl.load(segm_sum_ptr + m_idx) + + valid = seg_max > float("-inf") + overall_max = tl.max(tl.where(valid, seg_max, float("-inf"))) + + rescale = tl.where(valid, tl.exp(seg_max - overall_max), 0.0) + seg_sum_rescaled = seg_sum * rescale + overall_sum = tl.sum(seg_sum_rescaled) + + so_idx = so_base + p_off[:, None] * (QG * HEAD_SIZE) + d_off[None, :] + seg_out = tl.load(segm_out_ptr + so_idx).to(tl.float32) + weighted = seg_out * seg_sum_rescaled[:, None] + acc_sum = tl.sum(weighted, axis=0) + acc = tl.where(overall_sum > 0.0, acc_sum / overall_sum, 0.0) + + out_off = n * out_stride_n + hq * out_stride_h + d_off + tl.store(output_ptr + out_off, acc.to(output_ptr.dtype.element_ty)) + + +# -- Fused Q rotation + E4M3 haircut (Triton) --------------------------------- +# Replaces the 4-op torch prologue (copy bf16->fp32, mm by PiT, copy fp32->e4m3, +# copy e4m3->bf16) with ONE kernel. The win is kernel COUNT, not bytes: each of +# those small ops costs ~2-4us of fixed GPU-side cost here largely independent of +# element count (verified: sweeping the pool's B_bucket 512->1, i.e. 512x fewer +# elements in the copies, moved the floor 0%). Under CUDA graphs the host launch +# cost is already gone, so only per-kernel GPU cost remains and collapsing 4 +# kernels into 1 is the lever that is left. +@triton.jit +def _fused_q_rot_haircut( + q_ptr, # [B, Hq, D] query dtype (bf16/fp16), strided + pit_ptr, # [D, D] fp32 contiguous (already qperm-permuted) + out_ptr, # [B_bucket, Hq, D] out dtype, contiguous + M, # B * Hq (runtime) + q_stride_n: tl.int64, + q_stride_h: tl.int64, + HQ: tl.constexpr, + D: tl.constexpr, + HAIRCUT: tl.constexpr, # round through E4M3 (fp8-consuming QK MFMA paths) + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_d = tl.program_id(1) + m_off = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + m_mask = m_off < M + n = m_off // HQ + h = m_off % HQ + + # Splitting the OUTPUT columns across programs (not just the rows) is what + # gives this kernel parallelism at small batch: M = B*Hq is only 8 at B=1, + # so a row-only grid launches ONE workgroup and serializes the whole + # 256x256 rotation on a single CU. The K reduction is untouched by the + # split, so the result stays bit-identical. + d_off = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) + for k0 in range(0, D, BLOCK_K): + k_off = k0 + tl.arange(0, BLOCK_K) + q_idx = n[:, None] * q_stride_n + h[:, None] * q_stride_h + k_off[None, :] + q_blk = tl.load(q_ptr + q_idx, mask=m_mask[:, None], other=0.0) + p_blk = tl.load(pit_ptr + k_off[:, None] * D + d_off[None, :]) + acc = tl.dot(q_blk.to(tl.float32), p_blk, acc) + + if HAIRCUT: + # fp32 -> E4M3 -> fp32. Mirrors the torch path's + # q_rot_fp32 -> float8_e4m3fn -> bf16 round trip (e4m3 is a subset of + # bf16, so the second hop is exact and folds away here). + acc = acc.to(tl.float8e4nv).to(tl.float32) + + out_idx = m_off[:, None] * D + d_off[None, :] + tl.store(out_ptr + out_idx, acc.to(out_ptr.dtype.element_ty), mask=m_mask[:, None]) + + +# Prologue Q-rotation fusion (fallback for non-hd256 / in-kernel-WHT-off paths). +# Cleared the accuracy ladder (offline parity -> autoregressive drift -> paired +# GSM8K); pinned ON as part of the validated production config. +_FUSE_Q_ROT = True + +# Stronger form: do the rotation INSIDE the decode kernel as a Walsh-Hadamard +# butterfly and drop the prologue dispatch entirely. The prologue is ~7us at +# B=4 and that is almost all fixed per-dispatch cost (sweeping its grid 4x moves +# it 0%), so deleting the launch -- not optimizing it -- is the win. +# Default ON: verified bit-exact (torch.equal) against the prologue path across +# B={1,2,3,4,8,16} x seq={4k,8k,16k,32k,50021,75k}, including non-power-of-2 +# batch and seq. +_FUSE_QROT_INKERNEL = True + + +def _run_fused_q_rot(query, PiT_used, out, B, Hq, D, haircut): + """Launch the fused prologue. ``out`` is the pooled [B_bucket, Hq, D] buf.""" + M = B * Hq + BLOCK_M = 16 + BLOCK_K = 32 + BLOCK_D = 64 + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(D, BLOCK_D)) + _fused_q_rot_haircut[grid]( + query, + PiT_used, + out, + M, + query.stride(0), + query.stride(1), + HQ=int(Hq), + D=int(D), + HAIRCUT=bool(haircut), + BLOCK_M=BLOCK_M, + BLOCK_K=BLOCK_K, + BLOCK_D=BLOCK_D, + ) + + +# -- Public launcher ---------------------------------------------------------- +def flydsl_ultraquant_decode_attention( + query: torch.Tensor, + kv_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + scale: float, + PiT: torch.Tensor | None = None, + max_seq_len: int = 0, + output_buf: torch.Tensor | None = None, + buf_holder: Any = None, + max_num_kv_splits: int = 32, + sinks: torch.Tensor | None = None, +) -> torch.Tensor: + """Launch the optimized FlyDSL UltraQuant D=256 decode.""" + if not is_ultraquant_flydsl_available(): + raise RuntimeError( + "FlyDSL UltraQuant decode requires gfx950 and importable FlyDSL." + ) + if sinks is not None: + raise NotImplementedError("FlyDSL UltraQuant decode does not support sinks") + + B, Hq, D = query.shape + Hk = kv_cache.shape[2] + block_size = kv_cache.shape[1] + padded_slot = int(kv_cache.shape[3]) + QG = Hq // Hk + assert D == 256, f"UltraQuant FlyDSL requires head_dim=256, got {D}" + assert block_size in (16, 32, 64, 128, 256), ( + f"UltraQuant FlyDSL supports kv_block_size 16/32/64/128/256, got {block_size}" + ) + assert QG in _ULTRAQUANT_FLYDSL_GQA, ( + f"UltraQuant FlyDSL supports GQA factor " + f"{', '.join(map(str, _ULTRAQUANT_FLYDSL_GQA))}, got {QG}" + ) + + device = query.device + + # ---- PiT (Hadamard) — cache the contiguous fp32 form on the layer ----- + PiT_f32: torch.Tensor + if buf_holder is not None: + PiT_f32 = getattr(buf_holder, "_ultraquant_PiT_f32", None) + if PiT_f32 is None: + assert PiT is not None, ( + "UltraQuant launcher requires PiT (the Hadamard rotation)" + ) + PiT_f32 = PiT if PiT.dtype == torch.float32 else PiT.to(torch.float32) + if not PiT_f32.is_contiguous(): + PiT_f32 = PiT_f32.contiguous() + buf_holder._ultraquant_PiT_f32 = PiT_f32 + else: + assert PiT is not None, ( + "UltraQuant launcher requires PiT (the Hadamard rotation)" + ) + PiT_f32 = PiT if PiT.dtype == torch.float32 else PiT.to(torch.float32) + if not PiT_f32.is_contiguous(): + PiT_f32 = PiT_f32.contiguous() + + centroids_c = _fp4_value_lut(device) + + # ---- Partition count (FA2 split-KV) ----------------------------------- + kv_compute_block = _ULTRAQUANT_MOD.KV_COMPUTE_BLOCK + worst_case_max_seq_len = int(block_table.shape[1]) * int(block_size) + if max_seq_len <= 0: + max_seq_len = worst_case_max_seq_len + sizing_max_seq_len = worst_case_max_seq_len + + # Split-KV parallelism cap. The old default (32) silently serialized long + # contexts: once required_num_partitions exceeds the cap, TGPP grows and each + # workgroup walks MORE tile-groups instead of the grid getting wider. The + # decode kernel is latency-bound at ~1-2 waves/CU (grid = B*Hk*P workgroups + # of a single wavefront), so that serialization dominated long-context decode. + # Raising the cap to 256 keeps TGPP ~1 and measured (B=16, rocprof, per step + # decode+reduce+store): 16k 82.7->40.0us, 32k 152.0->55.1us, 64k 339.4->95.7us, + # 131k 671.3->179.0us (2.1x - 3.8x, growing with context). 512 regresses + # because _reduce_partitions cost grows with partition count. + # Cost: the segm pool scales with P (~8x vs P=32). + # BATCH-ADAPTIVE split-KV cap (low-concurrency occupancy fix). + # rocprof (B=4, 75k): the decode kernel is occupancy/latency-bound, not + # compute-bound -- OccupancyPercent 5.8%, MfmaUtil 4.5%, VALUBusy 14.5% + # (stalled ~85% on HBM), VALUUtilization 90% (no divergence). The grid is + # B*Hk*num_partitions single-wavefront workgroups; at B=4/Hk=1/P=256 that's + # only 1024 wg = 4/CU, too few to hide memory latency. Splitting MORE at low + # batch adds resident wavefronts and cuts ITL; at high batch the grid is + # already full so extra partitions only inflate _reduce_partitions + empty- + # partition overhead. Measured end-to-end (graphs on, TP=8, uniform 75k, + # median ITL vs kv8): cap512 vs cap256 -> C4 +9.0%->+4.9%, C8 +6.6%->+2.5% + # (better), C16 +2.8%->+4.3% (worse). So use 512 at low batch, 256 at high. + # This is only SAFE because the build-local-allocator fix in + # ultraquant_decode_hd256.py lets the two num_partitions variants (256 & 512) + # coexist across CUDA-graph batch buckets without the @flyc.jit + # "global 'allocator' changed since first compile" drift crash. + MAX_PARTITIONS = 512 if B <= 8 else 256 + required_num_partitions = ( + sizing_max_seq_len + kv_compute_block - 1 + ) // kv_compute_block + parallelism_floor = min(MAX_PARTITIONS, max(1, max_num_kv_splits)) + num_partitions_actual = max( + parallelism_floor, + min(MAX_PARTITIONS, required_num_partitions), + ) + # Cap partitions by how much work there actually IS to split. Nothing above + # forces a partition to hold a useful number of tokens -- max_num_kv_splits + # sets a parallelism FLOOR -- so a short context gets split into partitions + # of a handful of tokens each, where every partition is almost entirely + # fixed cost (and each one still writes a full QG x D partial for the + # reduce to read back). Measured sweet spot is ~130-147 tokens/partition, + # roughly flat in between; next_power_of_2 below rounds this target up, so + # the effective range lands at ~96-192 tokens. At B=4 this is worth 27.9 -> + # 24.4 us at seq=2048 and 27.0 -> 24.8 us at seq=16384, and is inert once + # the context is long enough to want every partition (seq >= ~64k). + _tok_per_part = 192 + if _tok_per_part > 0: + _useful_parts = max( + 1, (int(sizing_max_seq_len) + _tok_per_part - 1) // _tok_per_part + ) + num_partitions_actual = min(num_partitions_actual, _useful_parts) + num_partitions = max(2, triton.next_power_of_2(num_partitions_actual)) + _tgpp_required = max( + 1, + (required_num_partitions + num_partitions - 1) // num_partitions, + ) + tile_groups_per_partition = int(triton.next_power_of_2(_tgpp_required)) + + # ---- Pooled buffers --------------------------------------------------- + pool_bufs = _SEGM_POOL.get( + B, + Hk, + Hq, + num_partitions, + QG, + D, + device, + query.dtype, + ) + segm_out = pool_bufs["segm_out"] + segm_max = pool_bufs["segm_max"] + segm_sum = pool_bufs["segm_sum"] + if output_buf is None: + output = pool_bufs["output"] + else: + output = output_buf[:B] if output_buf.shape[0] != B else output_buf + + # Fold the scaled-MFMA head-dim permutation into PiT's columns so the + # WHT / prologue produces native operand order. Always apply the E4M3 + # haircut: scaled QK consumes Q as E4M3. + PiT_used = PiT_f32 + qperm_idx = _qperm_index(device, D=D) + if buf_holder is not None: + PiT_used = getattr(buf_holder, "_ultraquant_PiT_perm_f32", None) + if PiT_used is None: + PiT_used = PiT_f32.index_select(1, qperm_idx).contiguous() + buf_holder._ultraquant_PiT_perm_f32 = PiT_used + else: + PiT_used = PiT_f32.index_select(1, qperm_idx).contiguous() + _q_float = pool_bufs["q_float"] + _q_rot_f32 = pool_bufs["q_rot_fp32"] + _q_fp8 = pool_bufs["q_fp8"] + _q_rot_out = pool_bufs["q_rot"] + + _inkernel_qrot = ( + _FUSE_QROT_INKERNEL + and int(D) == 256 + and QG in _ULTRAQUANT_FLYDSL_GQA + and query.dim() == 3 + and query.stride(2) == 1 + and query.stride(1) == int(D) + ) + if _inkernel_qrot: + q_for_kernel = query + # The fused kernel indexes the head-dim contiguously; fall back if not. + elif _FUSE_Q_ROT and query.stride(-1) == 1: + # One kernel for the whole prologue instead of four ops. + _run_fused_q_rot(query, PiT_used, _q_rot_out, B, Hq, D, haircut=True) + else: + _q_float.copy_(query) + torch.mm( + _q_float.view(B * Hq, D), + PiT_used, + out=_q_rot_f32.view(B * Hq, D), + ) + _q_fp8.copy_(_q_rot_f32) # fp32 -> e4m3 + _q_rot_out.copy_(_q_fp8) # e4m3 -> bf16 + if not _inkernel_qrot: + q_for_kernel = _q_rot_out + + # ---- FlyDSL kernel launch -------------------------------------------- + max_bps = int(block_table.shape[1]) + launch = _get_kernel( + Hk, + num_partitions, + max_bps, + scale, + QG, + block_size, + padded_slot, + fuse_qrot=_inkernel_qrot, + num_seqs_hint=int(B), + tile_groups_per_partition=int(tile_groups_per_partition), + stride_q_seq=int(q_for_kernel.stride(0)), + stride_q_head=int(q_for_kernel.stride(1)), + ) + global _LOG_INVOKED_ONCE + if not _LOG_INVOKED_ONCE: + _LOG_INVOKED_ONCE = True + logger.info( + "FlyDSL UltraQuant launcher invoked: B=%d Hk=%d Hq=%d D=%d QG=%d " + "num_partitions=%d (actual=%d, cap=%d) TGPP=%d max_bps=%d " + "block_size=%d padded_slot=%d max_seq_len=%d " + "(coverage=%d tokens, worst_case=%d tokens)", + B, + Hk, + Hq, + D, + QG, + num_partitions, + num_partitions_actual, + MAX_PARTITIONS, + tile_groups_per_partition, + max_bps, + int(block_size), + padded_slot, + int(max_seq_len), + num_partitions * tile_groups_per_partition * kv_compute_block, + worst_case_max_seq_len, + ) + launch( + segm_out, + segm_sum, + segm_max, + q_for_kernel, + kv_cache, + centroids_c, + block_table, + seq_lens, + B, + Hk, + num_partitions, + torch.cuda.current_stream(), + ) + + # ---- Reduce partitions -> [B, Hq, D] --------------------------------- + # 64 splits the head-dim 4 ways, taking the grid from 32 to 128 workgroups. + # Bit-identical by construction (the partition reduction is per-column). + _red_block_d = min(int(D), 64) + _reduce_partitions[(B, Hq, triton.cdiv(int(D), _red_block_d))]( + output_ptr=output, + segm_out_ptr=segm_out, + segm_max_ptr=segm_max, + segm_sum_ptr=segm_sum, + out_stride_n=output.stride(0), + out_stride_h=output.stride(1), + NUM_KV_HEADS=Hk, + QG=QG, + NUM_PARTS=num_partitions, + HEAD_SIZE=D, + BLOCK_D=_red_block_d, + ) + return output diff --git a/vllm/v1/attention/ops/ultraquant/__init__.py b/vllm/v1/attention/ops/ultraquant/__init__.py new file mode 100644 index 0000000000..dbe17bfdc8 --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""UltraQuant 4-bit KV cache: FP4 E2M1 codes + UE8M0 group-of-32 scales. + +Production decode uses the FlyDSL D=256 kernel on gfx950, with Triton +unified attention as the fallback. Slot size is ``slot_size(head_dim)`` +(272 B at D=256). Format helpers live in ``format``; import them from +there, not this package root. +""" + +__all__ = ["ultraquant_store"] + + +def __getattr__(name: str): + if name == "ultraquant_store": + from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + + return ultraquant_store + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/v1/attention/ops/ultraquant/format.py b/vllm/v1/attention/ops/ultraquant/format.py new file mode 100644 index 0000000000..20977d528f --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/format.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FP4 codepoint + UE8M0 scale constants for the UltraQuant KV cache. + +K/V codes are FP4 E2M1. Per-group scales are UE8M0 (one byte, a power of +two): ``s = 2^round(log2(c · absmax))`` with ``c = 0.156``. Decode +consumes Q as FP8 E4M3 so scaled F8F6F4 MFMA can run natively on CDNA4. + +AoS slot layout (group_size=32), per (token, head): + bytes [0 .. D/2) : K codes (FP4 nibbles, 2/byte) + bytes [D/2 .. D/2 + Gk) : K scales (UE8M0, 1 byte × Gk groups) + bytes [D/2 + Gk .. D + Gk) : V codes + bytes [D + Gk .. D + 2*Gk) : V scales + D=256 → Gk=8 → 272 B/slot. D=128 → Gk=4 → 136 B/slot. + +No per-token norm-fold. No V rotation. K is Hadamard-rotated at store. +""" + +from __future__ import annotations + +import math + +# ── 15 unique FP4 E2M1 levels, sorted ascending ──── +FP4_LEVELS_SORTED: tuple[float, ...] = ( + -6.0, + -4.0, + -3.0, + -2.0, + -1.5, + -1.0, + -0.5, + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, +) +assert len(FP4_LEVELS_SORTED) == 15 + +# Midpoints between consecutive sorted levels (14 boundaries). +MIDPOINTS_SORTED: tuple[float, ...] = tuple( + (FP4_LEVELS_SORTED[i] + FP4_LEVELS_SORTED[i + 1]) / 2.0 + for i in range(len(FP4_LEVELS_SORTED) - 1) +) +assert len(MIDPOINTS_SORTED) == 14 + +# FP4 E2M1 bit-pattern decode (16 entries indexed by 4-bit code). +# Sign bit (MSB) | 2 exp bits | 1 mantissa bit (LSB). +FP4_BITS_TO_VALUE: tuple[float, ...] = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) +assert len(FP4_BITS_TO_VALUE) == 16 + +# Sorted index → FP4 E2M1 bit pattern (uint4). +SORTED_TO_BITS: tuple[int, ...] = ( + 0b1111, + 0b1110, + 0b1101, + 0b1100, + 0b1011, + 0b1010, + 0b1001, + 0b0000, + 0b0001, + 0b0010, + 0b0011, + 0b0100, + 0b0101, + 0b0110, + 0b0111, +) +assert len(SORTED_TO_BITS) == 15 +for _i, _bits in enumerate(SORTED_TO_BITS): + _v = FP4_BITS_TO_VALUE[_bits] + _v = 0.0 if _v == 0.0 else _v + assert _v == FP4_LEVELS_SORTED[_i] + +# ── Format constants ────────────────────────────────────────────────────── +GROUP_SIZE: int = 32 +"""Elements per group; one E8M0 scale per group. Fixed at 32 because the +AMD scaled F8F6F4 MFMA instruction also consumes one E8M0 scale per 32 +elements along K. Other group sizes would force a fallback to plain MFMA ++ accumulator-side scale multiply (no hardware fast-path).""" + +FP4_MAX: float = 6.0 +"""Largest representable FP4 magnitude.""" + +DEFAULT_CONSTANT_C: float = 0.156 +"""MSE-optimal scale constant; ``s_raw = c * absmax``, then UE8M0-snapped.""" + + +def get_constant_c() -> float: + """Return the fixed UltraQuant scale constant.""" + return DEFAULT_CONSTANT_C + + +def get_group_size() -> int: + """Return the fixed group size required by scaled MFMA.""" + return GROUP_SIZE + + +# ── Scale encoding ──────────────────────────────────────────────────────── +# Per-group scale `s = pow2_round(c · absmax)` stored as the E8M0 byte; +# codebook is the raw FP4 grid {0, ±0.5, ..., ±6}; dequant = code · s. + + +# ── UE8M0 helpers ────────────────────────────────────────────────────────── +# E8M0 = 8 exponent bits, 0 mantissa bits, no sign. The byte value `e` +# represents `2^(e - 127)` for e in [1, 254]. e=0 and e=255 are reserved +# (zero and NaN respectively in the OCP spec). We use e=0 as our zero +# sentinel for sink-zero groups. + +UE8M0_BIAS: int = 127 + +# fp32 minimum positive normal = 2^-126 ≈ 1.175e-38. Anything smaller +# clamps to the zero sentinel. +_UE8M0_MIN_EXP: int = -126 +_UE8M0_MAX_EXP: int = 127 + + +def ue8m0_encode(s: float) -> int: + """Snap a positive fp32 scale `s` to the nearest power of 2 and + encode as a UE8M0 byte. `s <= 0` encodes as 0 (zero sentinel).""" + if s <= 0.0 or not math.isfinite(s): + return 0 + exp = int(round(math.log2(s))) + if exp < _UE8M0_MIN_EXP: + return 0 + if exp > _UE8M0_MAX_EXP: + exp = _UE8M0_MAX_EXP + return exp + UE8M0_BIAS + + +def ue8m0_decode(byte: int) -> float: + """Decode a UE8M0 byte back to fp32. 0 → 0.0 (zero sentinel).""" + if byte == 0: + return 0.0 + return 2.0 ** (byte - UE8M0_BIAS) + + +# ── Per-slot AoS byte layout ────────────────────────────────────────────── +# Per slot, per head (group_size=GS=32, Gk = D/GS): +# bytes [0 .. D/2) : K codes (2 nibbles/byte) +# bytes [D/2 .. D/2 + Gk) : K scales (E8M0, 1 byte each) +# bytes [D/2 + Gk .. D + Gk) : V codes +# bytes [D + Gk .. D + 2*Gk) : V scales (E8M0, 1 byte each) +# D=256 → 272 B/slot. D=128 → 136 B/slot. + + +def k_codes_bytes(head_dim: int) -> int: + """Bytes for one head's packed K codes.""" + if head_dim % 2 != 0: + raise ValueError(f"head_dim must be even, got {head_dim}") + return head_dim // 2 + + +def n_groups(head_dim: int, group_size: int | None = None) -> int: + gs = group_size if group_size is not None else get_group_size() + if head_dim % gs != 0: + raise ValueError(f"head_dim={head_dim} must be a multiple of group_size={gs}") + return head_dim // gs + + +def k_scales_bytes(head_dim: int, group_size: int | None = None) -> int: + """Bytes for one head's K scales (one E8M0 byte per group).""" + return n_groups(head_dim, group_size) + + +def v_codes_bytes(head_dim: int) -> int: + return k_codes_bytes(head_dim) + + +def v_scales_bytes(head_dim: int, group_size: int | None = None) -> int: + return k_scales_bytes(head_dim, group_size) + + +def slot_size(head_dim: int, group_size: int | None = None) -> int: + """Bytes per (token, head) slot.""" + return ( + k_codes_bytes(head_dim) + + k_scales_bytes(head_dim, group_size) + + v_codes_bytes(head_dim) + + v_scales_bytes(head_dim, group_size) + ) + + +def k_codes_offset(head_dim: int, group_size: int | None = None) -> int: + return 0 + + +def k_scales_offset(head_dim: int, group_size: int | None = None) -> int: + return k_codes_bytes(head_dim) + + +def v_codes_offset(head_dim: int, group_size: int | None = None) -> int: + return k_codes_bytes(head_dim) + k_scales_bytes(head_dim, group_size) + + +def v_scales_offset(head_dim: int, group_size: int | None = None) -> int: + return v_codes_offset(head_dim, group_size) + v_codes_bytes(head_dim) diff --git a/vllm/v1/attention/ops/ultraquant/reference.py b/vllm/v1/attention/ops/ultraquant/reference.py new file mode 100644 index 0000000000..0d7dbe745b --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/reference.py @@ -0,0 +1,349 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""PyTorch reference for the UltraQuant KV cache format. + +Scale is UE8M0 (power of two): after ``s = c * absmax`` snap +``s = 2^round(log2(s))`` and store one byte. Q is Hadamard-rotated then +cast to FP8 E4M3 before the QK matmul, matching the kernel launcher. + +Same FP4 E2M1 grid. ``c = 0.156``. No per-token L2 norm. No V rotation. +K is rotated at encode; Q arrives pre-rotated at decode. + +This module is the bit-comparison ground truth for the kernels. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from vllm.v1.attention.ops.ultraquant.format import ( + FP4_LEVELS_SORTED, + GROUP_SIZE, + MIDPOINTS_SORTED, + SORTED_TO_BITS, + UE8M0_BIAS, + get_constant_c, +) + +_LEVELS_CACHE: dict[tuple[torch.device, torch.dtype], torch.Tensor] = {} +_MIDPOINTS_CACHE: dict[tuple[torch.device, torch.dtype], torch.Tensor] = {} +_SORTED_TO_BITS_CACHE: dict[torch.device, torch.Tensor] = {} + + +def _get_levels(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + key = (device, dtype) + t = _LEVELS_CACHE.get(key) + if t is None: + t = torch.tensor(FP4_LEVELS_SORTED, device=device, dtype=dtype) + _LEVELS_CACHE[key] = t + return t + + +def _get_midpoints(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + key = (device, dtype) + t = _MIDPOINTS_CACHE.get(key) + if t is None: + t = torch.tensor(MIDPOINTS_SORTED, device=device, dtype=dtype) + _MIDPOINTS_CACHE[key] = t + return t + + +def _get_sorted_to_bits(device: torch.device) -> torch.Tensor: + t = _SORTED_TO_BITS_CACHE.get(device) + if t is None: + t = torch.tensor(SORTED_TO_BITS, device=device, dtype=torch.uint8) + _SORTED_TO_BITS_CACHE[device] = t + return t + + +_HADAMARD_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.Tensor] = {} + + +def hadamard_matrix( + dim: int, device: torch.device, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """Sylvester Hadamard, normalised so H @ H.T = I.""" + if dim <= 0 or (dim & (dim - 1)) != 0: + raise ValueError(f"hadamard_matrix requires power-of-two dim, got {dim}") + key = (dim, device, dtype) + cached = _HADAMARD_CACHE.get(key) + if cached is not None: + return cached + H = torch.tensor([[1.0]], dtype=torch.float64) + while H.shape[0] < dim: + H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0) + H = (H / (dim**0.5)).to(device=device, dtype=dtype).contiguous() + _HADAMARD_CACHE[key] = H + return H + + +# ── Core quantization primitives ─────────────────────────────────────────── +def _snap_to_sorted_idx( + x_norm: torch.Tensor, *, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """Snap normalised values to nearest FP4 level, returning sorted index + in [0, 14]. Same semantic as `torch.bucketize` against midpoints + (reproduced by 14 `tl.where` in the kernel).""" + boundaries = _get_midpoints(x_norm.device, dtype) + x = x_norm.to(dtype) + # ROCm's bucketize uses int32 element indexing and raises + # hipErrorInvalidConfiguration once numel >= 2**31 (e.g. a [B, Hk, N, D] + # K tensor at B=128, N=16384, Hk=8, D=128 is exactly 2**31). Chunk the + # flattened input so each call stays under the limit (bit-identical output). + _INT32_LIMIT = 2**31 + if x.numel() < _INT32_LIMIT: + return torch.bucketize(x, boundaries) + flat = x.reshape(-1) + out = torch.empty_like(flat, dtype=torch.long) + chunk = _INT32_LIMIT // 2 + for i in range(0, flat.numel(), chunk): + out[i : i + chunk] = torch.bucketize(flat[i : i + chunk], boundaries) + return out.reshape(x.shape) + + +def _pack_nibbles_last_dim(codes: torch.Tensor) -> torch.Tensor: + if codes.shape[-1] % 2 != 0: + raise ValueError(f"pack_nibbles last dim must be even, got {codes.shape[-1]}") + lo = codes[..., 0::2] & 0xF + hi = codes[..., 1::2] & 0xF + return (lo | (hi << 4)).to(torch.uint8) + + +def _unpack_nibbles_last_dim(packed: torch.Tensor, out_dim: int) -> torch.Tensor: + if out_dim != 2 * packed.shape[-1]: + packed_dim = packed.shape[-1] + raise ValueError( + f"unpack_nibbles: out_dim={out_dim} must be 2*packed.last={2 * packed_dim}" + ) + p = packed.to(torch.int32) + lo = (p & 0xF).to(torch.uint8) + hi = ((p >> 4) & 0xF).to(torch.uint8) + out_shape = packed.shape[:-1] + (out_dim,) + out = torch.empty(out_shape, dtype=torch.uint8, device=packed.device) + out[..., 0::2] = lo + out[..., 1::2] = hi + return out + + +def _ue8m0_snap_tensor( + s_raw: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Snap a positive-scale tensor `s_raw` to a power of two and return + ``(s_snapped_fp32, byte_uint8)``. + + Matches the kernel: zero / non-finite inputs map to byte=0 / value=0.0. + Round-half-up via ``floor(log2(s) + 0.5)``, matching Triton's + ``tl.floor(... + 0.5)`` idiom. + """ + is_zero = (s_raw <= 0) | ~torch.isfinite(s_raw) + s_safe = torch.where(is_zero, torch.ones_like(s_raw), s_raw) + log2_s = torch.log2(s_safe) + exp = torch.floor(log2_s + 0.5).to(torch.int32) + # Clamp to UE8M0's representable exponent range (signed pre-bias). + exp = torch.clamp(exp, min=-126, max=127) + byte = (exp + UE8M0_BIAS).to(torch.int32) + byte = torch.where(is_zero, torch.zeros_like(byte), byte).to(torch.uint8) + s_snapped = torch.where( + is_zero, torch.zeros_like(s_raw), torch.exp2(exp.to(torch.float32)) + ) + return s_snapped, byte + + +@dataclass +class UltraQuantEncoded: + """Encoded ultraquant representation of one tensor along its last axis. + + Shapes assume input shape `[..., D]`: + - codes_packed: `[..., D // 2]` uint8 — 2 FP4 nibbles per byte + - scale_bytes: `[..., D // GROUP_SIZE]` uint8 — E8M0 byte per group + Dequant is ``code · 2^(byte - 127)`` (`c` is folded into the byte). + """ + + codes_packed: torch.Tensor # uint8, [..., D//2] + scale_bytes: torch.Tensor # uint8, [..., D//GROUP_SIZE] + + @property + def scales_fp32(self) -> torch.Tensor: + """Decode the E8M0 bytes back to fp32 (for dequant / inspection).""" + b = self.scale_bytes.to(torch.int32) + exp = b - UE8M0_BIAS + is_zero = b == 0 + s = torch.where( + is_zero, + torch.zeros_like(b, dtype=torch.float32), + torch.exp2(exp.to(torch.float32)), + ) + return s + + +def ultraquant_encode( + x: torch.Tensor, + *, + rotate: bool = True, + constant_c: float | None = None, +) -> UltraQuantEncoded: + """Encode `x` (shape [..., D]) to FP4 codes + E8M0 per-group-of-32 scales. + + Steps: + 1. Optional Hadamard rotation: `x_rot = x @ H.T` + 2. Reshape to `[..., G, GROUP_SIZE]`. + 3. `absmax = max(|x_rot|)` per group. + 4. `s_raw = c · absmax`; `s_snapped = 2^round(log2(s_raw))`. + Zero-amax groups encode byte=0. + 5. `sorted_idx = bucketize(x_g / s_snapped, midpoints)` ∈ [0, 14]. + 6. Remap sorted_idx → FP4 E2M1 bit pattern via SORTED_TO_BITS. + 7. Pack pairs of 4-bit codes into bytes. + """ + if x.shape[-1] % GROUP_SIZE != 0: + raise ValueError( + f"ultraquant_encode: last dim {x.shape[-1]} must be a multiple of " + f"GROUP_SIZE={GROUP_SIZE}" + ) + D = x.shape[-1] + G = D // GROUP_SIZE + c = constant_c if constant_c is not None else get_constant_c() + + x_f32 = x.to(torch.float32) + if rotate: + H = hadamard_matrix(D, x.device, torch.float32) + x_rot = x_f32 @ H.T + else: + x_rot = x_f32 + + x_g = x_rot.reshape(*x.shape[:-1], G, GROUP_SIZE) + absmax = x_g.abs().amax(dim=-1, keepdim=True) # [..., G, 1] + s_raw = absmax * c + s_snapped, scale_bytes = _ue8m0_snap_tensor(s_raw.squeeze(-1)) + s_snapped = s_snapped.unsqueeze(-1) # [..., G, 1] + s_for_div = torch.where(s_snapped == 0, torch.ones_like(s_snapped), s_snapped) + + sorted_idx = _snap_to_sorted_idx(x_g / s_for_div) + sorted_idx = torch.where( + (s_snapped == 0).expand_as(sorted_idx), + torch.full_like(sorted_idx, 7), # sorted idx 7 → +0.0 + sorted_idx, + ) + + sorted_to_bits = _get_sorted_to_bits(x.device) + fp4_bits = sorted_to_bits[sorted_idx] # [..., G, GS] + fp4_bits_flat = fp4_bits.reshape(*x.shape[:-1], D) + codes_packed = _pack_nibbles_last_dim(fp4_bits_flat) # [..., D//2] + + return UltraQuantEncoded( + codes_packed=codes_packed, + scale_bytes=scale_bytes, + ) + + +def ultraquant_dequant( + encoded: UltraQuantEncoded, + head_dim: int, +) -> torch.Tensor: + """Decode `encoded` to fp32 (no inverse rotation — result in the + rotated basis if rotation was applied at encode). + + Dequant is ``code · 2^(byte - 127)`` (`c` is already folded into the byte). + """ + G = head_dim // GROUP_SIZE + codes = _unpack_nibbles_last_dim(encoded.codes_packed, head_dim) + bits_to_val = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + device=encoded.codes_packed.device, + dtype=torch.float32, + ) + decoded = bits_to_val[codes.to(torch.long)] # [..., D] fp32 + decoded_g = decoded.reshape(*decoded.shape[:-1], G, GROUP_SIZE) + scales = encoded.scales_fp32.unsqueeze(-1) # [..., G, 1] + return (decoded_g * scales).reshape(*decoded.shape[:-1], head_dim) + + +def ultraquant_encode_decode( + x: torch.Tensor, + *, + rotate: bool = True, + constant_c: float | None = None, +) -> torch.Tensor: + """Full encode→decode round-trip in the rotated basis. Returns fp32.""" + enc = ultraquant_encode( + x, + rotate=rotate, + constant_c=constant_c, + ) + return ultraquant_dequant(enc, x.shape[-1]) + + +# ── Reference attention against the ultraquant encoded cache ────────────────── +def reference_ultraquant_attention( + query: torch.Tensor, # [B, Hq, D] bf16 — RAW (will be rotated + cast to FP8 here) + key: torch.Tensor, # [B, N, Hk, D] bf16 + value: torch.Tensor, # [B, N, Hk, D] bf16 + *, + scale: float, + sinks: torch.Tensor | None = None, + constant_c: float | None = None, +) -> torch.Tensor: + """End-to-end reference attention with ultraquant K/V encode→decode in + the rotated basis. Q is Hadamard-rotated AND cast through FP8 E4M3 + (scale = 1) before the QK matmul — mirrors what the launcher does. + + Output shape: [B, Hq, D] in query.dtype. + """ + B, Hq, D = query.shape + B2, N, Hk, D2 = key.shape + assert B == B2 and D == D2, (B, B2, D, D2) + assert value.shape == key.shape + assert Hq % Hk == 0, f"Hq={Hq} must be a multiple of Hk={Hk}" + g = Hq // Hk + + H = hadamard_matrix(D, query.device, torch.float32) + q_rot = query.to(torch.float32) @ H.T + # Round-trip Q through FP8 E4M3 with scale=1 — same precision haircut + # as the launcher (`q_rot.to(torch.float8_e4m3fn)`). + q_fp8 = q_rot.to(torch.float8_e4m3fn) + q_for_matmul = q_fp8.to(torch.float32) + + k_recon_rot = ultraquant_encode_decode( + key.transpose(1, 2).contiguous(), + rotate=True, + constant_c=constant_c, + ) + v_recon = ultraquant_encode_decode( + value.transpose(1, 2).contiguous(), + rotate=False, + constant_c=constant_c, + ) + + k_recon_rot = k_recon_rot.unsqueeze(2).expand(B, Hk, g, N, D).reshape(B, Hq, N, D) + v_recon = v_recon.unsqueeze(2).expand(B, Hk, g, N, D).reshape(B, Hq, N, D) + + qk = torch.einsum("bhd,bhnd->bhn", q_for_matmul, k_recon_rot) * scale # [B, Hq, N] + + if sinks is not None: + sink_log = sinks.to(torch.float32).reshape(1, Hq, 1).expand(B, Hq, 1) + qk_full = torch.cat([qk, sink_log], dim=-1) + p_full = torch.softmax(qk_full, dim=-1) + p = p_full[..., :N] + else: + p = torch.softmax(qk, dim=-1) + + out = torch.einsum("bhn,bhnd->bhd", p, v_recon) + return out.to(query.dtype) diff --git a/vllm/v1/attention/ops/ultraquant/triton_dequant.py b/vllm/v1/attention/ops/ultraquant/triton_dequant.py new file mode 100644 index 0000000000..9ff8d01bf4 --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/triton_dequant.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Full KV dequant for the UltraQuant cache format. + +Used by continuation prefill: when a chunk brings many new query tokens on +top of a long cached prefix, dequanting the prefix once and running a dense +prefill kernel beats replaying the decode kernel per query token. + +Reads FP4 codes plus UE8M0 group scales and writes K (Hadamard-rotated, as +stored) and V (raw) into pre-allocated fp16/bf16 buffers. +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.ultraquant.format import ( + FP4_BITS_TO_VALUE, + UE8M0_BIAS, + get_group_size, + k_scales_offset, + n_groups, + v_codes_offset, + v_scales_offset, +) +from vllm.v1.attention.ops.ultraquant.triton_store import _kv_cache_flat + +_FP4_DECODE_CACHE: dict[tuple[torch.device, torch.dtype], torch.Tensor] = {} + + +def _get_fp4_decode_table(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """16-entry FP4 E2M1 bit-pattern -> value table (cached per device).""" + key = (device, dtype) + t = _FP4_DECODE_CACHE.get(key) + if t is None: + t = torch.tensor(FP4_BITS_TO_VALUE, device=device, dtype=dtype).contiguous() + _FP4_DECODE_CACHE[key] = t + return t + + +@triton.jit +def _ultraquant_full_dequant_kv( + KV_cache_ptr, # uint8 view, flat + Block_table_ptr, # [B, max_num_blocks] int32 + Fp4_decode_ptr, # [16] FP4 bit pattern -> value + K_out_ptr, + V_out_ptr, + stride_ko_b: tl.int64, + stride_ko_h: tl.int64, + stride_ko_s: tl.int64, + stride_vo_b: tl.int64, + stride_vo_h: tl.int64, + stride_vo_s: tl.int64, + stride_cache_block: tl.int64, + stride_cache_pos: tl.int64, + stride_cache_head: tl.int64, + stride_bt_b: tl.int64, + HEAD_DIM: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + BLOCK_D: tl.constexpr, + OUT_BF16: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, +): + pos = tl.program_id(0) + bh = tl.program_id(1) + bid = bh // NUM_KV_HEADS + hid = bh % NUM_KV_HEADS + + page_idx = pos // BLOCK_SIZE + page_off = pos % BLOCK_SIZE + block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64) + slot_base = ( + block_num * stride_cache_block + + page_off * stride_cache_pos + + tl.cast(hid, tl.int64) * stride_cache_head + ) + + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + byte_idx = (d_offs // 2).to(tl.int32) + nibble_shift = ((d_offs % 2) * 4).to(tl.int32) + g_offs = tl.arange(0, N_GROUPS_C) + + out_dtype = tl.bfloat16 if OUT_BF16 else tl.float16 + + # K side: codes are packed nibbles, one UE8M0 byte per group of 32. + k_byte_raw = tl.load(KV_cache_ptr + slot_base + byte_idx, mask=d_mask, other=0).to( + tl.int32 + ) + k_codes = (k_byte_raw >> nibble_shift) & 0xF + k_dec = tl.load(Fp4_decode_ptr + k_codes).to(tl.float32) + k_dec = tl.where(d_mask, k_dec, 0.0) + + # UE8M0 scale: value = 2^(byte - bias); byte == 0 encodes exact zero. + k_scale_bytes = tl.load(KV_cache_ptr + slot_base + K_SCALES_OFFSET + g_offs).to( + tl.int32 + ) + k_scales = tl.where( + k_scale_bytes == 0, + 0.0, + tl.exp2(tl.cast(k_scale_bytes - UE8M0_BIAS_C, tl.float32)), + ) + + k_g = tl.reshape(k_dec, [N_GROUPS_C, GROUP_SIZE_C]) + k_recon = tl.reshape(k_g * k_scales[:, None], [BLOCK_D]) + k_recon = tl.where(d_mask, k_recon, 0.0) + + ko_base = bid * stride_ko_b + hid * stride_ko_h + pos * stride_ko_s + tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(out_dtype), mask=d_mask) + + # V side mirrors K but is stored unrotated. + v_byte_raw = tl.load( + KV_cache_ptr + slot_base + V_CODES_OFFSET + byte_idx, + mask=d_mask, + other=0, + ).to(tl.int32) + v_codes = (v_byte_raw >> nibble_shift) & 0xF + v_dec = tl.load(Fp4_decode_ptr + v_codes).to(tl.float32) + v_dec = tl.where(d_mask, v_dec, 0.0) + + v_scale_bytes = tl.load(KV_cache_ptr + slot_base + V_SCALES_OFFSET + g_offs).to( + tl.int32 + ) + v_scales = tl.where( + v_scale_bytes == 0, + 0.0, + tl.exp2(tl.cast(v_scale_bytes - UE8M0_BIAS_C, tl.float32)), + ) + + v_g = tl.reshape(v_dec, [N_GROUPS_C, GROUP_SIZE_C]) + v_recon = tl.reshape(v_g * v_scales[:, None], [BLOCK_D]) + v_recon = tl.where(d_mask, v_recon, 0.0) + + vo_base = bid * stride_vo_b + hid * stride_vo_h + pos * stride_vo_s + tl.store(V_out_ptr + vo_base + d_offs, v_recon.to(out_dtype), mask=d_mask) + + +def ultraquant_full_dequant_kv( + kv_cache: torch.Tensor, + block_table: torch.Tensor, + k_out: torch.Tensor, + v_out: torch.Tensor, + alloc_len: int, +) -> None: + """Dequant ``alloc_len`` cached positions into ``k_out`` / ``v_out``. + + ``k_out`` / ``v_out`` are ``[B, Hk, alloc_len, D]`` in fp16 or bf16. + """ + B = block_table.shape[0] + Hk = kv_cache.shape[2] + D = k_out.shape[3] + block_size = kv_cache.shape[1] + + gs = get_group_size() + _ultraquant_full_dequant_kv[(alloc_len, B * Hk)]( + _kv_cache_flat(kv_cache), + block_table, + _get_fp4_decode_table(kv_cache.device, k_out.dtype), + k_out, + v_out, + k_out.stride(0), + k_out.stride(1), + k_out.stride(2), + v_out.stride(0), + v_out.stride(1), + v_out.stride(2), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + block_table.stride(0), + HEAD_DIM=D, + BLOCK_SIZE=block_size, + NUM_KV_HEADS=Hk, + K_SCALES_OFFSET=k_scales_offset(D, gs), + V_CODES_OFFSET=v_codes_offset(D, gs), + V_SCALES_OFFSET=v_scales_offset(D, gs), + GROUP_SIZE_C=gs, + N_GROUPS_C=n_groups(D, gs), + BLOCK_D=triton.next_power_of_2(D), + OUT_BF16=1 if k_out.dtype == torch.bfloat16 else 0, + UE8M0_BIAS_C=UE8M0_BIAS, + num_warps=4, + ) diff --git a/vllm/v1/attention/ops/ultraquant/triton_store.py b/vllm/v1/attention/ops/ultraquant/triton_store.py new file mode 100644 index 0000000000..b43f4264e7 --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/triton_store.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton store kernel for the ultraquant KV cache format. + +Per-token, per-head, per-group-of-32 the kernel computes: + + s_raw = c * absmax(group) # c = 0.156 by default + s = 2^round(log2(s_raw)) # UE8M0 snap (power of 2) + code = quantize_to_fp4(group / s) # OCP E2M1, 4 bits/elt + +and stores ``code`` as packed nibbles + ``s`` as a single UE8M0 byte per +group. The output layout matches ``format.py`` (272 B/slot for D=256). + +K is Hadamard-rotated in-register; V is not. There is no per-token L2 +norm. The UE8M0 snap makes ``s`` a power of two so scaled F8F6F4 MFMA +can consume it on the read side. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.ultraquant.format import ( + MIDPOINTS_SORTED, + SORTED_TO_BITS, + UE8M0_BIAS, + get_constant_c, + get_group_size, + k_scales_offset, + n_groups, + slot_size, + v_codes_offset, + v_scales_offset, +) + + +def _kv_cache_flat(kv_cache: torch.Tensor) -> torch.Tensor: + """1-D uint8 view of the raw KV allocation (handles padded layouts).""" + if kv_cache.is_contiguous(): + return kv_cache.reshape(-1) + n = kv_cache.untyped_storage().nbytes() // kv_cache.element_size() + return torch.empty(0, dtype=kv_cache.dtype, device=kv_cache.device).set_( + kv_cache.untyped_storage(), 0, (n,) + ) + + +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16} + +# Built lazily on first launch so importing this module does not require triton. +_MIDPOINTS_C: Any = None + + +def _ensure_const_tables() -> None: + global _MIDPOINTS_C + if _MIDPOINTS_C is None: + _MIDPOINTS_C = tl.constexpr(tuple(float(m) for m in MIDPOINTS_SORTED)) + + +# ── Cached per-device constant tensors ───────────────────────────────────── +_MIDPOINTS_CACHE: dict[torch.device, torch.Tensor] = {} +_SORTED_TO_BITS_CACHE: dict[torch.device, torch.Tensor] = {} +_PIT_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.Tensor] = {} + + +def _get_midpoints_tensor(device: torch.device) -> torch.Tensor: + t = _MIDPOINTS_CACHE.get(device) + if t is None: + t = torch.tensor(MIDPOINTS_SORTED, device=device, dtype=torch.float32) + _MIDPOINTS_CACHE[device] = t + return t + + +def _get_sorted_to_bits_tensor(device: torch.device) -> torch.Tensor: + t = _SORTED_TO_BITS_CACHE.get(device) + if t is None: + t = torch.tensor(SORTED_TO_BITS, device=device, dtype=torch.int32) + _SORTED_TO_BITS_CACHE[device] = t + return t + + +def _get_hadamard( + dim: int, device: torch.device, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """Sylvester Hadamard in fp32 (cached, normalised).""" + key = (dim, device, dtype) + cached = _PIT_CACHE.get(key) + if cached is not None: + return cached + if dim <= 0 or (dim & (dim - 1)) != 0: + raise ValueError(f"ultraquant store requires power-of-two dim, got {dim}") + H = torch.tensor([[1.0]], dtype=torch.float64) + while H.shape[0] < dim: + H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0) + H = (H / (dim**0.5)).to(device=device, dtype=dtype).contiguous() + _PIT_CACHE[key] = H + return H + + +# ═══════════════════════════════════════════════════════════════════════════ +# Triton kernel +# ═══════════════════════════════════════════════════════════════════════════ + + +@triton.jit +def _ultraquant_store_kernel( + Key_ptr, # [N*H, D] in raw dtype (bf16 or fp16) + Value_ptr, # [N*H, D] in raw dtype + KV_cache_ptr, # raw uint8 view, flat + Slot_mapping_ptr, # [N] int64 + Midpoints_ptr, # [14] fp32 + SortedToBits_ptr, # [15] int32 + # Cache strides (in bytes) + stride_cache_block: tl.constexpr, + stride_cache_pos: tl.constexpr, + stride_cache_head: tl.constexpr, + # Dimensions + HEAD_DIM: tl.constexpr, + H: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + LOG2_D: tl.constexpr, + # Layout + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, + FP4_C: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, + # When 1, the FP4 midpoints and the sorted->bits remap are materialised as + # immediates instead of being read from `Midpoints_ptr`/`SortedToBits_ptr`. + # That removes 28 scalar global loads and two 256-lane divergent gathers + # from the critical path of a kernel whose runtime is pure latency (flat at + # ~22 us from batch 1 to batch 64). Numerically identical: the immediates + # are the same fp32 values, and the remap formula is exact (asserted in + # `format`, verified for all 15 entries). + CONST_TABLES: tl.constexpr = 1, +): + pid = tl.program_id(0) + token_idx = pid // H + head_idx = pid % H + + slot = tl.load(Slot_mapping_ptr + token_idx) + if slot < 0: + return + blk = slot // BLOCK_SIZE + off = slot % BLOCK_SIZE + slot_base = ( + blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head + ).to(tl.int64) + k_code_base = slot_base + k_scale_base = slot_base + K_SCALES_OFFSET + v_code_base = slot_base + V_CODES_OFFSET + v_scale_base = slot_base + V_SCALES_OFFSET + + base = pid * HEAD_DIM + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + shifts_4 = tl.arange(0, 2) * 4 # 0, 4 for nibble packing + g_offs = tl.arange(0, N_GROUPS_C) + + # ── K side ───────────────────────────────────────────────────────── + # 1. Load + Hadamard-rotate K via WHT butterfly. + k_raw = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32) + k_rot = tl.where(d_mask, k_raw, 0.0) + + INV_SQRT2: tl.constexpr = 0.7071067811865476 + for _stage in tl.static_range(LOG2_D): + _bmask = 1 << _stage + _pidx = (d_offs ^ _bmask).to(tl.int32) + _pval = tl.gather(k_rot, _pidx, 0) + _sign = tl.where((d_offs & _bmask) == 0, 1.0, -1.0).to(tl.float32) + k_rot = (_sign * k_rot + _pval) * INV_SQRT2 + k_rot = tl.where(d_mask, k_rot, 0.0) + + # 2. Per-group absmax. + k_g = tl.reshape(k_rot, [N_GROUPS_C, GROUP_SIZE_C]) + k_absmax = tl.max(tl.abs(k_g), axis=1) # [N_GROUPS] + k_is_zero = k_absmax == 0.0 + + # 3. UE8M0 snap: s_raw = c * absmax; exp = round(log2(s_raw)); + # divisor = 2^exp (c folded in). Zero groups write byte=0. + k_s_raw = k_absmax * tl.cast(FP4_C, tl.float32) + k_s_safe = tl.where(k_is_zero, 1.0, k_s_raw) + k_log2 = tl.log2(k_s_safe) + k_exp = tl.cast(tl.floor(k_log2 + 0.5), tl.int32) # round-half-up + k_s_snapped = tl.exp2(tl.cast(k_exp, tl.float32)) + k_s_div = tl.where(k_is_zero, 1.0, k_s_snapped) + + # 4. Normalize + snap to sorted FP4 idx via 14 sequential tl.where (`>` for + # bucketize(right=False) parity with the reference). + k_norm = k_g / k_s_div[:, None] + k_sorted = tl.zeros([N_GROUPS_C, GROUP_SIZE_C], dtype=tl.int32) + for i in tl.static_range(14): + mid = _MIDPOINTS_C.value[i] if CONST_TABLES else tl.load(Midpoints_ptr + i) + k_sorted += tl.where(k_norm > mid, 1, 0) + k_sorted = tl.where(k_is_zero[:, None], 7, k_sorted) + + # 5. Remap sorted idx → FP4 E2M1 bit pattern. + if CONST_TABLES: + k_bits = tl.where(k_sorted < 7, 15 - k_sorted, k_sorted - 7) + else: + k_bits = tl.load(SortedToBits_ptr + k_sorted) + + # 6. Pack pairs of 4-bit codes into uint8 bytes. + k_bits_flat = tl.reshape(k_bits, [HEAD_DIM]) + k_pairs = tl.reshape(k_bits_flat, [HEAD_DIM // 2, 2]) + k_packed = tl.sum((k_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8) + k_code_addrs = k_code_base + tl.arange(0, HEAD_DIM // 2) + tl.store(KV_cache_ptr + k_code_addrs, k_packed) + + # 7. Store K scales as E8M0 bytes (one per group). Zero-amax groups + # write byte=0 (zero sentinel). + k_byte = tl.where(k_is_zero, 0, k_exp + UE8M0_BIAS_C) + k_byte = tl.where((k_byte < 0) | (k_byte > 255), 0, k_byte).to(tl.uint8) + k_scale_addrs = k_scale_base + g_offs + tl.store(KV_cache_ptr + k_scale_addrs, k_byte) + + # ── V side (mirror K but skip rotation) ──────────────────────────── + v_raw = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32) + v_raw = tl.where(d_mask, v_raw, 0.0) + + v_g = tl.reshape(v_raw, [N_GROUPS_C, GROUP_SIZE_C]) + v_absmax = tl.max(tl.abs(v_g), axis=1) + v_is_zero = v_absmax == 0.0 + + v_s_raw = v_absmax * tl.cast(FP4_C, tl.float32) + v_s_safe = tl.where(v_is_zero, 1.0, v_s_raw) + v_log2 = tl.log2(v_s_safe) + v_exp = tl.cast(tl.floor(v_log2 + 0.5), tl.int32) + v_s_snapped = tl.exp2(tl.cast(v_exp, tl.float32)) + v_s_div = tl.where(v_is_zero, 1.0, v_s_snapped) + + v_norm = v_g / v_s_div[:, None] + v_sorted = tl.zeros([N_GROUPS_C, GROUP_SIZE_C], dtype=tl.int32) + for i in tl.static_range(14): + mid = _MIDPOINTS_C.value[i] if CONST_TABLES else tl.load(Midpoints_ptr + i) + v_sorted += tl.where(v_norm > mid, 1, 0) + v_sorted = tl.where(v_is_zero[:, None], 7, v_sorted) + + if CONST_TABLES: + v_bits = tl.where(v_sorted < 7, 15 - v_sorted, v_sorted - 7) + else: + v_bits = tl.load(SortedToBits_ptr + v_sorted) + v_bits_flat = tl.reshape(v_bits, [HEAD_DIM]) + v_pairs = tl.reshape(v_bits_flat, [HEAD_DIM // 2, 2]) + v_packed = tl.sum((v_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8) + v_code_addrs = v_code_base + tl.arange(0, HEAD_DIM // 2) + tl.store(KV_cache_ptr + v_code_addrs, v_packed) + + v_byte = tl.where(v_is_zero, 0, v_exp + UE8M0_BIAS_C) + v_byte = tl.where((v_byte < 0) | (v_byte > 255), 0, v_byte).to(tl.uint8) + v_scale_addrs = v_scale_base + g_offs + tl.store(KV_cache_ptr + v_scale_addrs, v_byte) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Launcher +# ═══════════════════════════════════════════════════════════════════════════ + + +def ultraquant_store( + key: torch.Tensor, # [N, H, D] bf16 or fp16 + value: torch.Tensor, # [N, H, D] bf16 or fp16 + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, slot_size_aligned] uint8 + slot_mapping: torch.Tensor, # [N] int64 + *, + PiT: torch.Tensor | None = None, # unused; kept for call-site compatibility + constant_c: float | None = None, +) -> None: + """Launch the ultraquant store kernel. Writes K/V into `kv_cache` + in-place at the slots specified by `slot_mapping`.""" + if key.dtype not in _SUPPORTED_DTYPES: + raise ValueError( + f"ultraquant_store: key.dtype must be one of " + f"{_SUPPORTED_DTYPES}, got {key.dtype}" + ) + if value.dtype != key.dtype: + raise ValueError( + f"ultraquant_store: key/value dtype mismatch ({key.dtype} vs {value.dtype})" + ) + if slot_mapping.dtype != torch.int64: + slot_mapping = slot_mapping.to(torch.int64) + + N, H, D = key.shape + NH = N * H + block_size = kv_cache.shape[1] + num_kv_heads = kv_cache.shape[2] + padded_slot = kv_cache.shape[3] + + expected_slot = slot_size(D) + if padded_slot < expected_slot: + raise ValueError( + f"ultraquant_store: kv_cache slot {padded_slot} < expected " + f"{expected_slot} for head_dim={D}" + ) + if num_kv_heads != H: + raise ValueError( + f"ultraquant_store: kv_cache num_kv_heads {num_kv_heads} != key heads {H}" + ) + + midpoints = _get_midpoints_tensor(key.device) + sorted_to_bits = _get_sorted_to_bits_tensor(key.device) + + k_flat = key.reshape(NH, D).contiguous() + v_flat = value.reshape(NH, D).contiguous() + + gs = get_group_size() + BLOCK_D = triton.next_power_of_2(D) + N_GROUPS_C = n_groups(D, gs) + K_SCALES_OFFSET = k_scales_offset(D, gs) + V_CODES_OFFSET = v_codes_offset(D, gs) + V_SCALES_OFFSET = v_scales_offset(D, gs) + + stride_block = kv_cache.stride(0) + stride_pos = kv_cache.stride(1) + stride_head = kv_cache.stride(2) + + c = constant_c if constant_c is not None else get_constant_c() + + _ensure_const_tables() + grid = (NH,) + _ultraquant_store_kernel[grid]( + k_flat, + v_flat, + _kv_cache_flat(kv_cache), + slot_mapping, + midpoints, + sorted_to_bits, + stride_cache_block=stride_block, + stride_cache_pos=stride_pos, + stride_cache_head=stride_head, + HEAD_DIM=D, + H=H, + BLOCK_SIZE=block_size, + BLOCK_D=BLOCK_D, + LOG2_D=int(D).bit_length() - 1, + GROUP_SIZE_C=gs, + N_GROUPS_C=N_GROUPS_C, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + FP4_C=c, + UE8M0_BIAS_C=UE8M0_BIAS, + CONST_TABLES=1, + num_warps=4, + num_stages=1, + ) diff --git a/vllm/v1/attention/ops/ultraquant/triton_unified_attention.py b/vllm/v1/attention/ops/ultraquant/triton_unified_attention.py new file mode 100644 index 0000000000..69ec744a07 --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/triton_unified_attention.py @@ -0,0 +1,1087 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unified Triton fallback for the UltraQuant 4-bit KV-cache format. + +K and V use packed FP4 E2M1 values with UE8M0 group-of-32 scales. QK uses +native scaled FP4×E4M3 MFMA on CDNA4. Production choices are fixed in code; +there are no environment-variable tuning switches. +""" + +from __future__ import annotations + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.triton_unified_attention import ( + reduce_segments, +) + +# Reuse the sequence-index search and vectorized partition reducer. +from vllm.v1.attention.ops.turboquant_soa.triton_turboquant_unified_attention import ( + _find_seq_idx, +) +from vllm.v1.attention.ops.ultraquant.format import ( + UE8M0_BIAS, + get_group_size, + k_scales_offset, + n_groups, + slot_size, + v_codes_offset, + v_scales_offset, +) +from vllm.v1.attention.ops.ultraquant.triton_store import _kv_cache_flat + +_is_hip = current_platform.is_rocm() + + +@triton.jit +def _ultraquant_fp4_decode_arith(codes): + """Arithmetic FP4 E2M1 decode, bit-exact to format.FP4_BITS_TO_VALUE.""" + mag = codes & 7 + e = mag >> 1 + m = mag & 1 + exp_field = 126 + e + mant_field = tl.where(e != 0, m, 0) << 22 + bits = (exp_field << 23) | mant_field + magval = bits.to(tl.float32, bitcast=True) + magval = tl.where(mag == 0, 0.0, magval) + signf = tl.where((codes & 8) != 0, -1.0, 1.0) + return magval * signf + + +@triton.jit +def _ultraquant_load_k_packed( + KV_cache_ptr, + data_bases, + k_scales_addrs, + d_half_offs, + half_mask, + tile_mask, + BLOCK_D: tl.constexpr, + HEAD_DIM: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + TILE_SIZE: tl.constexpr, + UNMASKED: tl.constexpr, +): + """Load packed FP4 K codes and UE8M0 scales for scaled QK MFMA.""" + addrs = data_bases[:, None] + d_half_offs[None, :] + if UNMASKED: + K_codes = tl.load(KV_cache_ptr + addrs, mask=half_mask[None, :], other=0) + else: + K_codes = tl.load( + KV_cache_ptr + addrs, + mask=tile_mask[:, None] & half_mask[None, :], + other=0, + ) + K_T_packed = tl.trans(K_codes) + grp = tl.arange(0, N_GROUPS_C) + scale_addrs = k_scales_addrs[:, None] + grp[None, :] + if UNMASKED: + K_scales = tl.load(KV_cache_ptr + scale_addrs) + else: + K_scales = tl.load(KV_cache_ptr + scale_addrs, mask=tile_mask[:, None], other=0) + _ = HEAD_DIM + _ = GROUP_SIZE_C + _ = BLOCK_D + return K_T_packed, K_scales + + +@triton.jit +def _ultraquant_load_v_tile( + KV_cache_ptr, + val_bases, + v_scales_addrs, + d_offs, + d_mask, + tile_mask, + OUT_DTYPE: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + TILE_SIZE: tl.constexpr, + UNMASKED: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, +): + """Load and dequant a V tile with arithmetic FP4 decode.""" + half_idx = d_offs // 2 + nibble_shift = (d_offs % 2) * 4 + addrs = val_bases[:, None] + half_idx[None, :] + if UNMASKED: + byte_raw = tl.load(KV_cache_ptr + addrs, mask=d_mask[None, :], other=0).to( + tl.int32 + ) + else: + byte_raw = tl.load( + KV_cache_ptr + addrs, + mask=tile_mask[:, None] & d_mask[None, :], + other=0, + ).to(tl.int32) + codes = (byte_raw >> nibble_shift[None, :]) & 0xF + grp = tl.arange(0, N_GROUPS_C) + scale_addrs = v_scales_addrs[:, None] + grp[None, :] + if UNMASKED: + scale_bytes = tl.load(KV_cache_ptr + scale_addrs).to(tl.int32) + else: + scale_bytes = tl.load( + KV_cache_ptr + scale_addrs, mask=tile_mask[:, None], other=0 + ).to(tl.int32) + scale_exp = scale_bytes - UE8M0_BIAS_C + fp4_vals = _ultraquant_fp4_decode_arith(codes) + scales = tl.where(scale_bytes == 0, 0.0, tl.exp2(tl.cast(scale_exp, tl.float32))) + V_g = tl.reshape(fp4_vals, [TILE_SIZE, N_GROUPS_C, GROUP_SIZE_C]) + return tl.reshape(V_g * scales[:, :, None], [TILE_SIZE, BLOCK_D]).to(OUT_DTYPE) + + +@triton.jit +def _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + d_offs, + d_mask, + tile_mask, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_M: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + TILE_SIZE: tl.constexpr, + UNMASKED: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, +): + V = _ultraquant_load_v_tile( + KV_cache_ptr, + val_bases, + v_scales_addrs, + d_offs, + d_mask, + tile_mask, + OUT_DTYPE=tl.bfloat16, + HEAD_DIM=HEAD_DIM, + BLOCK_D=BLOCK_D, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=UNMASKED, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + _ = BLOCK_M + return tl.dot(P.to(tl.bfloat16), V, out_dtype=tl.float32) + + +@triton.jit +def _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + tile_mask, + BLOCK_D: tl.constexpr, + HEAD_DIM: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + TILE_SIZE: tl.constexpr, + UNMASKED: tl.constexpr, +): + K_T_packed, K_scales = _ultraquant_load_k_packed( + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + tile_mask, + BLOCK_D=BLOCK_D, + HEAD_DIM=HEAD_DIM, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=UNMASKED, + ) + return tl.dot_scaled( + Q, + None, + "e4m3", + K_T_packed, + K_scales, + "e2m1", + out_dtype=tl.float32, + ) + + +@triton.jit +def _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos: tl.int64, + stride_cache_head: tl.int64, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, +): + data_bases = ( + block_base + + slot_within_block * stride_cache_pos + + tl.cast(kv_head_idx, tl.int64) * stride_cache_head + ) + return ( + data_bases, + data_bases + K_SCALES_OFFSET, + data_bases + V_CODES_OFFSET, + data_bases + V_SCALES_OFFSET, + ) + + +# =========================================================================== +# Unified 2D attention kernel (prefill / short-context decode) +# =========================================================================== + + +@triton.jit +def kernel_ultraquant_unified_attention_2d( + output_ptr, + query_ptr, # FP8 E4M3, already Hadamard-rotated + KV_cache_ptr, # uint8 view + block_tables_ptr, + seq_lens_ptr, + query_start_len_ptr, + sinks_ptr, + scale, + num_query_heads: tl.constexpr, + num_queries_per_kv: tl.constexpr, + block_table_stride: tl.int64, + query_stride_0: tl.int64, + query_stride_1: tl.int64, + output_stride_0: tl.int64, + output_stride_1: tl.int64, + stride_cache_block: tl.int64, + stride_cache_pos: tl.int64, + stride_cache_head: tl.int64, + BLOCK_SIZE: tl.constexpr, + TILE_SIZE: tl.constexpr, + HEAD_SIZE: tl.constexpr, + HEAD_SIZE_PADDED: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_M: tl.constexpr, + num_seqs: tl.int32, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, + USE_SINKS: tl.constexpr = 0, + SLIDING_WINDOW: tl.constexpr = 0, +): + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + + seq_idx = _find_seq_idx( + query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True + ) + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx + q_block_local_idx = q_block_global_idx - q_block_start_idx + + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1) + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + offs_d_half = tl.arange(0, HEAD_SIZE_PADDED // 2) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d[None, :] + ) + + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1) + half_mask = tl.where(offs_d_half * 2 < HEAD_SIZE, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=tl.zeros([], tl.float8e4nv), + ) + + block_table_offset = seq_idx * block_table_stride + + if USE_SINKS: + M = tl.load( + sinks_ptr + query_offset_1, + mask=query_mask_1, + other=float("-inf"), + ).to(tl.float32) + else: + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32) + + seq_len = tl.load(seq_lens_ptr + seq_idx) + context_len = seq_len - cur_batch_query_len + + max_seq_prefix_len = ( + context_len + + q_block_local_idx * BLOCK_Q + + (BLOCK_M - 1) // num_queries_per_kv + + 1 + ) + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) + num_tiles = tl.cdiv(max_seq_prefix_len, TILE_SIZE) + + tile_start = 0 + tile_end = num_tiles + if SLIDING_WINDOW > 0: + qpos_lo = q_block_local_idx * BLOCK_Q + qpos_hi = tl.minimum( + qpos_lo + (BLOCK_M - 1) // num_queries_per_kv, + cur_batch_query_len - 1, + ) + first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1 + last_allowed_key = context_len + qpos_hi + tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) + tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) + + query_abs_pos = context_len + query_pos[:, None] + dummy_tile_mask = tl.full([TILE_SIZE], 1, tl.int1) + + for j in range(tile_start, tile_end - 1): + seq_offset = j * TILE_SIZE + offs_t + if TILE_SIZE == BLOCK_SIZE: + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j).to( + tl.int64 + ) + slot_within_block = offs_t.to(tl.int64) + else: + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64) + block_base = physical_block_idx * stride_cache_block + data_bases, k_scales_addrs, val_bases, v_scales_addrs = _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos, + stride_cache_head, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + ) + + S = scale * _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + dummy_tile_mask, + BLOCK_D=HEAD_SIZE_PADDED, + HEAD_DIM=HEAD_SIZE, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=True, + ) + seq_mask = seq_offset[None, :] <= query_abs_pos + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ( + (query_abs_pos - seq_offset[None, :]) < SLIDING_WINDOW + ) + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, + S, + float("-inf"), + ) + m_j = tl.maximum(M, tl.max(S, axis=1)) + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + P = tl.exp(S - m_j[:, None]) + l_j = tl.sum(P, axis=1) + alpha = tl.exp(M - m_j) + acc = acc * alpha[:, None] + L = L * alpha + l_j + M = m_j + acc += _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + offs_d, + dim_mask, + dummy_tile_mask, + HEAD_DIM=HEAD_SIZE, + BLOCK_D=HEAD_SIZE_PADDED, + BLOCK_M=BLOCK_M, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=True, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + + # Tail tile (masked) + if tile_end > tile_start: + j = tile_end - 1 + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + if TILE_SIZE == BLOCK_SIZE: + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j).to( + tl.int64 + ) + slot_within_block = offs_t.to(tl.int64) + else: + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64) + block_base = physical_block_idx * stride_cache_block + data_bases, k_scales_addrs, val_bases, v_scales_addrs = _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos, + stride_cache_head, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + ) + + S = scale * _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + tile_mask, + BLOCK_D=HEAD_SIZE_PADDED, + HEAD_DIM=HEAD_SIZE, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=False, + ) + seq_mask = seq_offset[None, :] <= query_abs_pos + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ( + (query_abs_pos - seq_offset[None, :]) < SLIDING_WINDOW + ) + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, + S, + float("-inf"), + ) + m_j = tl.maximum(M, tl.max(S, axis=1)) + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + P = tl.exp(S - m_j[:, None]) + l_j = tl.sum(P, axis=1) + alpha = tl.exp(M - m_j) + acc = acc * alpha[:, None] + L = L * alpha + l_j + M = m_j + acc += _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + offs_d, + dim_mask, + tile_mask, + HEAD_DIM=HEAD_SIZE, + BLOCK_D=HEAD_SIZE_PADDED, + BLOCK_M=BLOCK_M, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=False, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + + acc = acc / L[:, None] + + output_offset = ( + query_offset_0[:, None] * output_stride_0 + + query_offset_1[:, None] * output_stride_1 + + offs_d[None, :] + ) + tl.store( + output_ptr + output_offset, + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + + +# =========================================================================== +# Unified 3D (split-KV) attention kernel +# =========================================================================== + + +@triton.jit +def kernel_ultraquant_unified_attention_3d( + segm_output_ptr, + segm_max_ptr, + segm_expsum_ptr, + query_ptr, # FP8 E4M3, already Hadamard-rotated + KV_cache_ptr, + block_tables_ptr, + seq_lens_ptr, + query_start_len_ptr, + sinks_ptr, + scale, + num_query_heads: tl.constexpr, + num_queries_per_kv: tl.constexpr, + block_table_stride: tl.int64, + query_stride_0: tl.int64, + query_stride_1: tl.int64, + stride_cache_block: tl.int64, + stride_cache_pos: tl.int64, + stride_cache_head: tl.int64, + BLOCK_SIZE: tl.constexpr, + TILE_SIZE: tl.constexpr, + HEAD_SIZE: tl.constexpr, + HEAD_SIZE_PADDED: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_M: tl.constexpr, + num_seqs: tl.int32, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, + USE_SINKS: tl.constexpr = 0, + SLIDING_WINDOW: tl.constexpr = 0, +): + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + segm_idx = tl.program_id(2) + + seq_idx = _find_seq_idx( + query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True + ) + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx + q_block_local_idx = q_block_global_idx - q_block_start_idx + + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1) + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + seq_len = tl.load(seq_lens_ptr + seq_idx) + tiles_per_segment = tl.cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len: + return + + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + offs_d_half = tl.arange(0, HEAD_SIZE_PADDED // 2) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d[None, :] + ) + + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1) + half_mask = tl.where(offs_d_half * 2 < HEAD_SIZE, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=tl.zeros([], tl.float8e4nv), + ) + + block_table_offset = seq_idx * block_table_stride + + if USE_SINKS and segm_idx == 0: + M = tl.load( + sinks_ptr + query_offset_1, + mask=query_mask_1, + other=float("-inf"), + ).to(tl.float32) + else: + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32) + + context_len = seq_len - cur_batch_query_len + max_seq_prefix_len = ( + context_len + + q_block_local_idx * BLOCK_Q + + (BLOCK_M - 1) // num_queries_per_kv + + 1 + ) + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) + num_tiles = tl.cdiv(max_seq_prefix_len, TILE_SIZE) + + tile_lo = segm_idx * tiles_per_segment + tile_hi = tl.minimum((segm_idx + 1) * tiles_per_segment, num_tiles) + + if SLIDING_WINDOW > 0: + qpos_lo = q_block_local_idx * BLOCK_Q + qpos_hi = tl.minimum( + qpos_lo + (BLOCK_M - 1) // num_queries_per_kv, + cur_batch_query_len - 1, + ) + first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1 + last_allowed_key = context_len + qpos_hi + swa_tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) + swa_tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) + tile_lo = tl.maximum(tile_lo, swa_tile_start) + tile_hi = tl.minimum(tile_hi, swa_tile_end) + + query_abs_pos = context_len + query_pos[:, None] + dummy_tile_mask = tl.full([TILE_SIZE], 1, tl.int1) + tail = tile_hi - 1 + + for j in range(tile_lo, tile_hi - 1): + seq_offset = j * TILE_SIZE + offs_t + if TILE_SIZE == BLOCK_SIZE: + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j).to( + tl.int64 + ) + slot_within_block = offs_t.to(tl.int64) + else: + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64) + block_base = physical_block_idx * stride_cache_block + data_bases, k_scales_addrs, val_bases, v_scales_addrs = _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos, + stride_cache_head, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + ) + + S = scale * _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + dummy_tile_mask, + BLOCK_D=HEAD_SIZE_PADDED, + HEAD_DIM=HEAD_SIZE, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=True, + ) + seq_mask = seq_offset[None, :] <= query_abs_pos + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ( + (query_abs_pos - seq_offset[None, :]) < SLIDING_WINDOW + ) + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, + S, + float("-inf"), + ) + m_j = tl.maximum(M, tl.max(S, axis=1)) + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + P = tl.exp(S - m_j[:, None]) + l_j = tl.sum(P, axis=1) + alpha = tl.exp(M - m_j) + acc = acc * alpha[:, None] + L = L * alpha + l_j + M = m_j + acc += _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + offs_d, + dim_mask, + dummy_tile_mask, + HEAD_DIM=HEAD_SIZE, + BLOCK_D=HEAD_SIZE_PADDED, + BLOCK_M=BLOCK_M, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=True, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + + # Tail + if tile_lo < tile_hi: + j = tail + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + if TILE_SIZE == BLOCK_SIZE: + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j).to( + tl.int64 + ) + slot_within_block = offs_t.to(tl.int64) + else: + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64) + block_base = physical_block_idx * stride_cache_block + data_bases, k_scales_addrs, val_bases, v_scales_addrs = _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos, + stride_cache_head, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + ) + + S = scale * _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + tile_mask, + BLOCK_D=HEAD_SIZE_PADDED, + HEAD_DIM=HEAD_SIZE, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=False, + ) + seq_mask = seq_offset[None, :] <= query_abs_pos + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ( + (query_abs_pos - seq_offset[None, :]) < SLIDING_WINDOW + ) + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, + S, + float("-inf"), + ) + m_j = tl.maximum(M, tl.max(S, axis=1)) + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + P = tl.exp(S - m_j[:, None]) + l_j = tl.sum(P, axis=1) + alpha = tl.exp(M - m_j) + acc = acc * alpha[:, None] + L = L * alpha + l_j + M = m_j + acc += _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + offs_d, + dim_mask, + tile_mask, + HEAD_DIM=HEAD_SIZE, + BLOCK_D=HEAD_SIZE_PADDED, + BLOCK_M=BLOCK_M, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=False, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + + # Write segment partials for stage-2 reduce. + segm_output_offset = ( + query_offset_0[:, None].to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + segm_idx * HEAD_SIZE_PADDED + + tl.arange(0, HEAD_SIZE_PADDED)[None, :] + ) + tl.store( + segm_output_ptr + segm_output_offset, + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + segm_offset = ( + query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_offset_1 * NUM_SEGMENTS_PER_SEQ + + segm_idx + ) + tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1) + tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1) + + +# =========================================================================== +# Launcher +# =========================================================================== + + +_HADAMARD_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.Tensor] = {} + + +def _get_pit(dim: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Sylvester Hadamard projection / sqrt(dim). PiT = PiT.T (symmetric).""" + key = (dim, device, dtype) + H = _HADAMARD_CACHE.get(key) + if H is None: + assert (dim & (dim - 1)) == 0, f"dim={dim} must be power of 2" + H = torch.tensor([[1.0]], dtype=torch.float32) + while H.shape[0] < dim: + H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0) + H = (H / (dim**0.5)).to(device=device, dtype=dtype).contiguous() + _HADAMARD_CACHE[key] = H + return H + + +def ultraquant_unified_attention( + query: torch.Tensor, # [num_tokens, Hq, D] fp16/bf16 (raw) + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 + block_table: torch.Tensor, # [num_seqs, max_num_blocks] int32 + seq_lens: torch.Tensor, # [num_seqs] int32 + query_start_loc: torch.Tensor, # [num_seqs+1] int32 + scale: float, + PiT: torch.Tensor | None = None, + output: torch.Tensor | None = None, + tile_size: int | None = None, + max_query_len: int | None = None, + max_seq_len: int | None = None, + num_kv_splits: int | None = None, + force_2d: bool = False, + sinks: torch.Tensor | None = None, + sliding_window: int | None = None, +) -> torch.Tensor: + """Launch unified UltraQuant attention with scaled F8F6F4 MFMA.""" + assert query.dim() == 3, f"query must be [N, Hq, D], got {query.shape}" + num_tokens, Hq, D = query.shape + Hk = kv_cache.shape[2] + block_size = kv_cache.shape[1] + padded_slot = kv_cache.shape[3] + kv_group_size = Hq // Hk + num_seqs = int(query_start_loc.shape[0] - 1) + device = query.device + + gs = get_group_size() + if padded_slot < slot_size(D, gs): + raise ValueError( + f"ultraquant_unified: cache slot {padded_slot} < expected " + f"{slot_size(D, gs)} for D={D} group_size={gs}" + ) + if Hq % Hk != 0: + raise ValueError(f"Hq={Hq} must be a multiple of Hk={Hk}") + + if PiT is None: + PiT = _get_pit(D, device, torch.float32) + elif PiT.dtype != torch.float32 or not PiT.is_contiguous(): + PiT = PiT.to(torch.float32).contiguous() + + q_rot = (query.float() @ PiT).contiguous() + q_for_kernel = q_rot.to(torch.float8_e4m3fn).contiguous() + + if sinks is not None: + sinks_f32 = sinks if sinks.dtype == torch.float32 else sinks.to(torch.float32) + if not sinks_f32.is_contiguous(): + sinks_f32 = sinks_f32.contiguous() + assert sinks_f32.numel() == Hq, ( + f"sinks must have shape [Hq={Hq}], got numel={sinks_f32.numel()}" + ) + use_sinks = True + else: + sinks_f32 = q_for_kernel # harmless dummy; never dereferenced when USE_SINKS=0 + use_sinks = False + + if output is None: + output = torch.empty_like(query) + + # BLOCK_M heuristic shared with the compressed-KV fallback kernels. + if max_query_len is not None: + is_prefill_like = max_query_len > 1 + else: + is_prefill_like = num_tokens > num_seqs + + if is_prefill_like: + BLOCK_M = max(128, triton.next_power_of_2(kv_group_size)) + else: + BLOCK_M = 16 if kv_group_size <= 16 else triton.next_power_of_2(kv_group_size) + BLOCK_Q = BLOCK_M // kv_group_size + + total_num_q_blocks = num_tokens // BLOCK_Q + num_seqs + + if tile_size is None: + tile_size = 32 if is_prefill_like else 16 + + num_stages_2d = 1 if _is_hip else 2 + num_stages_3d = 3 if _is_hip else 2 + + BLOCK_D = triton.next_power_of_2(D) + N_GROUPS_C = n_groups(D, gs) + K_SCALES_OFFSET = k_scales_offset(D, gs) + V_CODES_OFFSET = v_codes_offset(D, gs) + V_SCALES_OFFSET = v_scales_offset(D, gs) + + scale_for_kernel = float(scale) + + kv_flat = _kv_cache_flat(kv_cache) + + # Dispatch: 2D for prefill / chunked; 3D for pure decode with long KV. + if max_seq_len is None: + max_seq_len_hint = int(block_table.shape[1]) * int(block_size) + else: + max_seq_len_hint = int(max_seq_len) + use_3d = (not force_2d) and (not is_prefill_like) and max_seq_len_hint >= 1024 + + if not use_3d: + kernel_ultraquant_unified_attention_2d[(total_num_q_blocks, Hk)]( + output_ptr=output, + query_ptr=q_for_kernel, + KV_cache_ptr=kv_flat, + block_tables_ptr=block_table, + seq_lens_ptr=seq_lens, + query_start_len_ptr=query_start_loc, + sinks_ptr=sinks_f32, + scale=scale_for_kernel, + num_query_heads=Hq, + num_queries_per_kv=kv_group_size, + block_table_stride=block_table.stride(0), + query_stride_0=q_for_kernel.stride(0), + query_stride_1=q_for_kernel.stride(1), + output_stride_0=output.stride(0), + output_stride_1=output.stride(1), + stride_cache_block=kv_cache.stride(0), + stride_cache_pos=kv_cache.stride(1), + stride_cache_head=kv_cache.stride(2), + BLOCK_SIZE=block_size, + TILE_SIZE=tile_size, + HEAD_SIZE=D, + HEAD_SIZE_PADDED=BLOCK_D, + BLOCK_Q=BLOCK_Q, + BLOCK_M=BLOCK_M, + num_seqs=num_seqs, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + GROUP_SIZE_C=gs, + N_GROUPS_C=N_GROUPS_C, + UE8M0_BIAS_C=UE8M0_BIAS, + USE_SINKS=1 if use_sinks else 0, + SLIDING_WINDOW=int(sliding_window) + if sliding_window and sliding_window > 0 + else 0, + num_warps=4, + num_stages=num_stages_2d, + ) + return output + + # 3D split-KV path + if num_kv_splits is None: + num_kv_splits = 16 + if num_kv_splits < 1: + num_kv_splits = 1 + if num_kv_splits & (num_kv_splits - 1) != 0: + num_kv_splits = 1 << (num_kv_splits.bit_length() - 1) + max_possible_splits = max(1, (max_seq_len_hint + tile_size - 1) // tile_size) + num_segments = max(1, min(num_kv_splits, max_possible_splits)) + + segm_output = torch.empty( + (num_tokens, Hq, num_segments, BLOCK_D), + dtype=torch.float32, + device=device, + ) + segm_max = torch.empty( + (num_tokens, Hq, num_segments), + dtype=torch.float32, + device=device, + ) + segm_expsum = torch.empty( + (num_tokens, Hq, num_segments), + dtype=torch.float32, + device=device, + ) + + kernel_ultraquant_unified_attention_3d[(total_num_q_blocks, Hk, num_segments)]( + segm_output_ptr=segm_output, + segm_max_ptr=segm_max, + segm_expsum_ptr=segm_expsum, + query_ptr=q_for_kernel, + KV_cache_ptr=kv_flat, + block_tables_ptr=block_table, + seq_lens_ptr=seq_lens, + query_start_len_ptr=query_start_loc, + sinks_ptr=sinks_f32, + scale=scale_for_kernel, + num_query_heads=Hq, + num_queries_per_kv=kv_group_size, + block_table_stride=block_table.stride(0), + query_stride_0=q_for_kernel.stride(0), + query_stride_1=q_for_kernel.stride(1), + stride_cache_block=kv_cache.stride(0), + stride_cache_pos=kv_cache.stride(1), + stride_cache_head=kv_cache.stride(2), + BLOCK_SIZE=block_size, + TILE_SIZE=tile_size, + HEAD_SIZE=D, + HEAD_SIZE_PADDED=BLOCK_D, + BLOCK_Q=BLOCK_Q, + BLOCK_M=BLOCK_M, + num_seqs=num_seqs, + NUM_SEGMENTS_PER_SEQ=num_segments, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + GROUP_SIZE_C=gs, + N_GROUPS_C=N_GROUPS_C, + UE8M0_BIAS_C=UE8M0_BIAS, + USE_SINKS=1 if use_sinks else 0, + SLIDING_WINDOW=int(sliding_window) + if sliding_window and sliding_window > 0 + else 0, + num_warps=2, + num_stages=num_stages_3d, + ) + + # Reduce split-KV partials with the shared vectorized reducer. + reduce_segments[(num_tokens, Hq)]( + output_ptr=output, + segm_output_ptr=segm_output, + segm_max_ptr=segm_max, + segm_expsum_ptr=segm_expsum, + seq_lens_ptr=seq_lens, + num_seqs=num_seqs, + num_query_heads=Hq, + out_scale_inv=1.0, + output_stride_0=output.stride(0), + output_stride_1=output.stride(1), + block_table_stride=block_table.stride(0), + TILE_SIZE=tile_size, + HEAD_SIZE=D, + HEAD_SIZE_PADDED=BLOCK_D, + query_start_len_ptr=query_start_loc, + BLOCK_Q=BLOCK_Q, + NUM_SEGMENTS_PER_SEQ=num_segments, + USE_FP8=False, + ) + + return output diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 846e6bcad0..b6913003be 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -55,6 +55,7 @@ class KVQuantMode(IntEnum): TURBOQUANT_K3V4_NC = 8 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 @property def is_per_token_head(self) -> bool: @@ -80,6 +81,11 @@ class KVQuantMode(IntEnum): KVQuantMode.TURBOQUANT_3BIT_NC, ) + @property + def is_ultraquant(self) -> bool: + """True for the UltraQuant 4-bit KV-cache format.""" + return self == KVQuantMode.ULTRAQUANT_4BIT + def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`.""" @@ -98,6 +104,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 kv_cache_dtype == "ultraquant_4bit": + return KVQuantMode.ULTRAQUANT_4BIT if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"): return KVQuantMode.FP8_PER_TENSOR return KVQuantMode.NONE