[Attention] Add UltraQuant 4-bit KV cache backend (FlyDSL D=256) (#57057)

Signed-off-by: Aditi Rana <aditi.rana@amd.com>

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})

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

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)

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

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"
+    )

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()}"
+        )

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",

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
 

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,

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),
             )
 

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,
 }

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

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)
+        )

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

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

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}")

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)

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)

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,
+    )

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,
+    )

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