test_fused_qk_norm_rope_cache_quant.py

# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

import copy
import torch
from torch import Tensor
import aiter
from aiter.test_common import checkAllclose, perftest, benchmark
from aiter.utility.dtypes import get_dtype_fp8
from aiter.utility import dtypes
import argparse
import pandas as pd


def rms_norm_forward(x: Tensor, weight: Tensor, eps: float):
    input_dtype = x.dtype
    variance = x.float().pow(2).mean(-1, keepdim=True)
    x = x * torch.rsqrt(variance + eps)
    x = x.to(input_dtype)
    return weight * x


def rms_norm_diffusers_forward(x: Tensor, weight: Tensor, eps: float):
    input_dtype = x.dtype
    variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
    hidden_states = x * torch.rsqrt(variance + eps)

    if weight is not None:
        if weight.dtype in [torch.float16, torch.bfloat16]:
            hidden_states = hidden_states.to(weight.dtype)
        hidden_states = hidden_states * weight
    else:
        hidden_states = hidden_states.to(input_dtype)

    return hidden_states


def apply_interleaved_rope(x: torch.Tensor, mrope_section: list[int]) -> torch.Tensor:
    """Apply interleaved MRoPE to 3D rotary embeddings.
    Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
    interleaved [THTHWHTHW...TT], preserving frequency continuity.
    """
    x_t = x[0].clone()
    x_t[..., 1 : mrope_section[1] * 3 : 3] = x[1, ..., 1 : mrope_section[1] * 3 : 3]
    x_t[..., 2 : mrope_section[2] * 3 : 3] = x[2, ..., 2 : mrope_section[2] * 3 : 3]
    return x_t


def apply_rotary_emb_torch(
    x: Tensor,
    cos: Tensor,
    sin: Tensor,
    is_neox_style: bool,
) -> Tensor:
    cos = cos.unsqueeze(-2).to(x.dtype)
    sin = sin.unsqueeze(-2).to(x.dtype)
    if is_neox_style:
        x1, x2 = torch.chunk(x, 2, dim=-1)
    else:
        x1 = x[..., ::2]
        x2 = x[..., 1::2]
    o1 = x1 * cos - x2 * sin
    o2 = x2 * cos + x1 * sin
    if is_neox_style:
        return torch.cat((o1, o2), dim=-1)
    else:
        return torch.stack((o1, o2), dim=-1).flatten(-2)


def apply_rotary_emb_diffusers(
    x: Tensor,
    cos: Tensor,
    sin: Tensor,
    is_neox_style: bool,
) -> Tensor:
    """Diffusers / qwen-image-edit reference rope: cos/sin stay fp32, x is upcast
    to fp32 for the multiply, output is cast back to x's original dtype.
    Mirrors the semantics of `_apply_rope_complex` (complex multiply in fp32,
    `.to(original_dtype)` on the result)."""
    out_dtype = x.dtype
    cos = cos.unsqueeze(-2).float()
    sin = sin.unsqueeze(-2).float()
    x = x.float()
    if is_neox_style:
        x1, x2 = torch.chunk(x, 2, dim=-1)
    else:
        x1 = x[..., ::2]
        x2 = x[..., 1::2]
    o1 = x1 * cos - x2 * sin
    o2 = x2 * cos + x1 * sin
    if is_neox_style:
        out = torch.cat((o1, o2), dim=-1)
    else:
        out = torch.stack((o1, o2), dim=-1).flatten(-2)
    return out.to(out_dtype)


def apply_rotary_emb_dispatch(
    x: Tensor,
    cos: Tensor,
    sin: Tensor,
    is_neox_style: bool,
    rotary_dim: int = 0,
) -> Tensor:
    """
    Args:
        x: [num_tokens, num_heads, head_size]
        cos: [num_tokens, rotary_dim // 2]
        sin: [num_tokens, rotary_dim // 2]
        is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
            positional embeddings.
        rotary_dim: 0 means full rotary; otherwise only the first rotary_dim
            channels are rotated.
    """
    head_size = x.shape[-1]
    rotary_dim_ = rotary_dim if rotary_dim > 0 else head_size
    if rotary_dim_ < head_size:
        x_rot = apply_rotary_emb_torch(x[..., :rotary_dim_], cos, sin, is_neox_style)
        return torch.cat((x_rot, x[..., rotary_dim_:]), dim=-1)
    return apply_rotary_emb_torch(x, cos, sin, is_neox_style)


def split_qkv(
    qkv: Tensor,
    num_heads_q: int,
    num_heads_k: int,
    num_heads_v: int,
    head_size: int,
) -> tuple[Tensor, Tensor, Tensor]:
    q_size = num_heads_q * head_size
    k_size = num_heads_k * head_size
    v_size = num_heads_v * head_size
    qkv_2d = qkv.view(qkv.shape[0], q_size + k_size + v_size)
    return (
        qkv_2d[:, :q_size],
        qkv_2d[:, q_size : q_size + k_size],
        qkv_2d[:, q_size + k_size :],
    )


def clone_qkv_inputs(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    num_heads_q: int,
    num_heads_k: int,
    num_heads_v: int,
    head_size: int,
) -> tuple[Tensor, Tensor, Tensor]:
    # deepcopy preserves the original split-view aliasing/strides, so perf/warmup
    # iterations do not mutate the caller's tensors in-place.
    return copy.deepcopy((q, k, v))


@perftest()
def run_torch_qk_norm_rope_cache_quant_shuffle(
    qkv: Tensor,  # contiguous (num_tokens * (num_heads_q + num_heads_k + num_heads_v) * head_size)
    qw: Tensor,  #  contiguous (head_size)
    kw: Tensor,  #  contiguous (head_size)
    cos_sin: Tensor,  # contiguous (max_positions * rotary_dim), rotary_dim <= head_size
    positions: Tensor,  # contiguous (3 * num_tokens) or (num_tokens)
    num_tokens: int,
    num_heads_q: int,
    num_heads_k: int,
    num_heads_v: int,
    head_size: int,
    is_neox_style: bool,
    eps: float,
    k_cache: Tensor,  # [num_blocks, num_heads_k, head_size // x, page_size, x]
    v_cache: Tensor,  # [num_blocks, num_heads_v, head_size, page_size]
    k_scale: Tensor,  # [num_blocks, page_size]
    v_scale: Tensor,  # [num_blocks, page_size]
    slot_mapping: Tensor,
    kv_cache_dtype: str,
):
    q, k, v = split_qkv(qkv, num_heads_q, num_heads_k, num_heads_v, head_size)

    q_by_head = q.view(num_tokens, num_heads_q, head_size)
    q_by_head = rms_norm_forward(q_by_head, qw, eps)
    q = q_by_head.view(q.shape)

    k_by_head = k.view(num_tokens, num_heads_k, head_size)
    k_by_head = rms_norm_forward(k_by_head, kw, eps)
    k = k_by_head.view(k.shape)

    rotary_dim = cos_sin.shape[-1]
    cos_sin = cos_sin.view(cos_sin.shape[0], rotary_dim)
    cos_sin = cos_sin[positions]
    cos, sin = cos_sin.chunk(2, dim=-1)

    q_shape = q.shape
    q = q.view(num_tokens, -1, head_size)
    q = apply_rotary_emb_dispatch(q, cos, sin, is_neox_style, rotary_dim)
    q = q.reshape(q_shape)

    k_shape = k.shape
    k = k.view(num_tokens, -1, head_size)
    k = apply_rotary_emb_dispatch(k, cos, sin, is_neox_style, rotary_dim)

    v = v.view(num_tokens, -1, head_size)

    from aiter import reshape_and_cache_with_pertoken_quant, reshape_and_cache

    if kv_cache_dtype == "auto":
        reshape_and_cache(
            k,
            v,
            k_cache,
            v_cache,
            slot_mapping,
            kv_cache_dtype,
            None,
            None,
            asm_layout=True,
        )
    else:
        reshape_and_cache_with_pertoken_quant(
            k, v, k_cache, v_cache, k_scale, v_scale, slot_mapping, asm_layout=True
        )

    k = k.reshape(k_shape)
    v = v.reshape(k_shape)
    return q, k, v, k_cache, v_cache


@perftest()
def run_aiter_qk_norm_rope_cache_quant_shuffle(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    qw: Tensor,  #  contiguous (head_size)
    kw: Tensor,  #  contiguous (head_size)
    cos_sin: Tensor,  # contiguous (max_positions * head_size)
    positions: Tensor,  # contiguous (3 * num_tokens)
    num_heads_q: int,
    num_heads_k: int,
    num_heads_v: int,
    head_size: int,
    is_neox_style: bool,
    eps: float,
    k_cache: Tensor,
    v_cache: Tensor,
    slot_mapping: Tensor,
    kv_cache_dtype: str,
    k_scale: Tensor,
    v_scale: Tensor,
):
    q, k, v = clone_qkv_inputs(
        q, k, v, num_heads_q, num_heads_k, num_heads_v, head_size
    )
    aiter.fused_qk_norm_rope_cache_quant_shuffle(
        q,
        k,
        v,
        num_heads_q=num_heads_q,
        num_heads_k=num_heads_k,
        num_heads_v=num_heads_v,
        head_dim=head_size,
        eps=eps,
        qw=qw,
        kw=kw,
        cos_sin_cache=cos_sin,
        is_neox_style=is_neox_style,
        pos_ids=positions,
        k_cache=k_cache,
        v_cache=v_cache,
        slot_mapping=slot_mapping,
        kv_cache_dtype=kv_cache_dtype,
        k_scale=k_scale,
        v_scale=v_scale,
    )
    return q, k, v, k_cache, v_cache


@benchmark()
def test_shuffle_contiguous_inputs_match_split_views():
    """Contiguous q/k/v inputs should match non-contiguous split views."""
    dtype = torch.bfloat16
    num_tokens = 11
    num_heads_q = num_heads_k = num_heads_v = 2
    head_size = 128
    is_neox_style = False
    eps = 1e-6
    kv_cache_dtype = "auto"
    num_blocks = 4
    page_size = 16
    max_positions = 4096

    k_cache = torch.randn(
        [num_blocks, page_size, num_heads_k, head_size],
        dtype=dtype,
        device="cuda",
    )
    v_cache = torch.randn(
        [num_blocks, page_size, num_heads_v, head_size],
        dtype=dtype,
        device="cuda",
    )
    x = 16 // k_cache.element_size()
    k_cache = (
        k_cache.view([num_blocks, page_size, num_heads_k, head_size // x, x])
        .permute(0, 2, 3, 1, 4)
        .contiguous()
    )
    v_cache = v_cache.permute(0, 2, 3, 1).contiguous()
    slot_mapping = torch.randperm(num_tokens, dtype=torch.int64, device="cuda")
    k_scale = torch.zeros(
        [num_blocks, num_heads_k, page_size], dtype=torch.float32, device="cuda"
    )
    v_scale = torch.zeros(
        [num_blocks, num_heads_v, page_size], dtype=torch.float32, device="cuda"
    )
    qkv = torch.randn(
        (num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )
    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_positions, head_size), dtype=dtype, device="cuda")
    positions = torch.randint(
        0, max_positions, (num_tokens,), dtype=torch.int64, device="cuda"
    )

    k_cache_c = k_cache.clone()
    v_cache_c = v_cache.clone()
    k_scale_c = k_scale.clone()
    v_scale_c = v_scale.clone()
    qkv_c = qkv.clone()

    q1_src, k1_src, v1_src = split_qkv(
        qkv.clone(), num_heads_q, num_heads_k, num_heads_v, head_size
    )
    q1_src = q1_src.contiguous()
    k1_src = k1_src.contiguous()
    v1_src = v1_src.contiguous()
    q2_src, k2_src, v2_src = split_qkv(
        qkv_c, num_heads_q, num_heads_k, num_heads_v, head_size
    )
    assert not q2_src.is_contiguous()
    assert not k2_src.is_contiguous()
    assert not v2_src.is_contiguous()

    (q1, k1, v1, kc1, vc1), _ = run_aiter_qk_norm_rope_cache_quant_shuffle(
        q1_src,
        k1_src,
        v1_src,
        qw,
        kw,
        cos_sin,
        positions,
        num_heads_q,
        num_heads_k,
        num_heads_v,
        head_size,
        is_neox_style,
        eps,
        k_cache,
        v_cache,
        slot_mapping,
        kv_cache_dtype,
        k_scale,
        v_scale,
    )

    (q2, k2, v2, kc2, vc2), _ = run_aiter_qk_norm_rope_cache_quant_shuffle(
        q2_src,
        k2_src,
        v2_src,
        qw,
        kw,
        cos_sin,
        positions,
        num_heads_q,
        num_heads_k,
        num_heads_v,
        head_size,
        is_neox_style,
        eps,
        k_cache_c,
        v_cache_c,
        slot_mapping,
        kv_cache_dtype,
        k_scale_c,
        v_scale_c,
    )

    checkAllclose(q1, q2, msg="q contiguous vs split-view", rtol=1e-2, atol=0.05)
    checkAllclose(k1, k2, msg="k contiguous vs split-view", rtol=1e-2, atol=0.05)
    checkAllclose(v1, v2, msg="v contiguous vs split-view", rtol=1e-2, atol=0.05)
    checkAllclose(kc1.float(), kc2.float(), msg="k_cache", rtol=1e-2, atol=0.05)
    checkAllclose(vc1.float(), vc2.float(), msg="v_cache", rtol=1e-2, atol=0.05)
    checkAllclose(k_scale, k_scale_c, msg="k_scale", rtol=1e-2, atol=0.05)
    checkAllclose(v_scale, v_scale_c, msg="v_scale", rtol=1e-2, atol=0.05)


@benchmark()
def test_shuffle_noncontiguous_split_views_match_contiguous_inputs():
    """Non-contiguous q/k/v split views should match equivalent contiguous inputs."""
    dtype = torch.bfloat16
    num_tokens = 11
    num_heads_q = num_heads_k = num_heads_v = 2
    head_size = 128
    is_neox_style = False
    eps = 1e-6
    kv_cache_dtype = "auto"
    num_blocks = 4
    page_size = 16
    max_positions = 4096

    k_cache = torch.randn(
        [num_blocks, page_size, num_heads_k, head_size],
        dtype=dtype,
        device="cuda",
    )
    v_cache = torch.randn(
        [num_blocks, page_size, num_heads_v, head_size],
        dtype=dtype,
        device="cuda",
    )
    x = 16 // k_cache.element_size()
    k_cache = (
        k_cache.view([num_blocks, page_size, num_heads_k, head_size // x, x])
        .permute(0, 2, 3, 1, 4)
        .contiguous()
    )
    v_cache = v_cache.permute(0, 2, 3, 1).contiguous()
    slot_mapping = torch.randperm(num_tokens, dtype=torch.int64, device="cuda")
    k_scale = torch.zeros(
        [num_blocks, num_heads_k, page_size], dtype=torch.float32, device="cuda"
    )
    v_scale = torch.zeros(
        [num_blocks, num_heads_v, page_size], dtype=torch.float32, device="cuda"
    )
    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_positions, head_size), dtype=dtype, device="cuda")
    positions = torch.randint(
        0, max_positions, (num_tokens,), dtype=torch.int64, device="cuda"
    )

    qs = num_heads_q * head_size
    ks = num_heads_k * head_size
    vs = num_heads_v * head_size
    total = qs + ks + vs
    flat = torch.randn((num_tokens, total), dtype=dtype, device="cuda")

    buf_ref = flat.clone()
    k_cache_ref = k_cache.clone()
    v_cache_ref = v_cache.clone()
    k_scale_ref = k_scale.clone()
    v_scale_ref = v_scale.clone()
    q_ref_src, k_ref_src, v_ref_src = split_qkv(
        buf_ref, num_heads_q, num_heads_k, num_heads_v, head_size
    )
    q_ref_src = q_ref_src.contiguous()
    k_ref_src = k_ref_src.contiguous()
    v_ref_src = v_ref_src.contiguous()
    (q_ref, k_ref, v_ref, kc_ref, vc_ref), _ = (
        run_aiter_qk_norm_rope_cache_quant_shuffle(
            q_ref_src,
            k_ref_src,
            v_ref_src,
            qw,
            kw,
            cos_sin,
            positions,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache_ref,
            v_cache_ref,
            slot_mapping,
            kv_cache_dtype,
            k_scale_ref,
            v_scale_ref,
        )
    )

    buf = flat.clone()
    k_cache_a = k_cache.clone()
    v_cache_a = v_cache.clone()
    k_scale_a = k_scale.clone()
    v_scale_a = v_scale.clone()
    qkv_2d = buf.view(num_tokens, total)
    q_nc = qkv_2d[:, :qs]
    k_nc = qkv_2d[:, qs : qs + ks]
    v_nc = qkv_2d[:, qs + ks :]
    assert not q_nc.is_contiguous()
    assert not k_nc.is_contiguous()
    assert not v_nc.is_contiguous()

    aiter.fused_qk_norm_rope_cache_quant_shuffle(
        q_nc,
        k_nc,
        v_nc,
        num_heads_q=num_heads_q,
        num_heads_k=num_heads_k,
        num_heads_v=num_heads_v,
        head_dim=head_size,
        eps=eps,
        qw=qw,
        kw=kw,
        cos_sin_cache=cos_sin,
        is_neox_style=is_neox_style,
        pos_ids=positions,
        k_cache=k_cache_a,
        v_cache=v_cache_a,
        slot_mapping=slot_mapping,
        kv_cache_dtype=kv_cache_dtype,
        k_scale=k_scale_a,
        v_scale=v_scale_a,
    )

    checkAllclose(q_nc, q_ref, msg="q split-view vs contiguous", rtol=1e-2, atol=0.05)
    checkAllclose(k_nc, k_ref, msg="k split-view vs contiguous", rtol=1e-2, atol=0.05)
    checkAllclose(v_nc, v_ref, msg="v split-view vs contiguous", rtol=1e-2, atol=0.05)
    checkAllclose(
        k_cache_a.float(), kc_ref.float(), msg="k_cache", rtol=1e-2, atol=0.05
    )
    checkAllclose(
        v_cache_a.float(), vc_ref.float(), msg="v_cache", rtol=1e-2, atol=0.05
    )
    checkAllclose(k_scale_a, k_scale_ref, msg="k_scale", rtol=1e-2, atol=0.05)
    checkAllclose(v_scale_a, v_scale_ref, msg="v_scale", rtol=1e-2, atol=0.05)


@benchmark()
def test_shuffle_3d_inputs_match_2d_inputs():
    """The operator should produce the same result for equivalent 2D and 3D q/k/v inputs."""
    dtype = torch.bfloat16
    num_tokens = 9
    num_heads_q = num_heads_k = num_heads_v = 2
    head_size = 128
    is_neox_style = False
    eps = 1e-6
    kv_cache_dtype = "auto"
    num_blocks = 3
    page_size = 16
    max_positions = 2048

    def make_kv_caches():
        kc = torch.randn(
            [num_blocks, page_size, num_heads_k, head_size],
            dtype=dtype,
            device="cuda",
        )
        vc = torch.randn(
            [num_blocks, page_size, num_heads_v, head_size],
            dtype=dtype,
            device="cuda",
        )
        x_ = 16 // kc.element_size()
        kc = (
            kc.view([num_blocks, page_size, num_heads_k, head_size // x_, x_])
            .permute(0, 2, 3, 1, 4)
            .contiguous()
        )
        vc = vc.permute(0, 2, 3, 1).contiguous()
        return kc, vc

    slot_mapping = torch.randperm(num_tokens, dtype=torch.int64, device="cuda")
    qkv = torch.randn(
        (num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )
    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_positions, head_size), dtype=dtype, device="cuda")
    positions = torch.randint(
        0, max_positions, (num_tokens,), dtype=torch.int64, device="cuda"
    )

    qs = int(num_heads_q * head_size)
    ks = int(num_heads_k * head_size)
    vs = int(num_heads_v * head_size)
    flat = qkv.view(num_tokens, qs + ks + vs)
    q_src = flat[:, :qs].contiguous()
    k_src = flat[:, qs : qs + ks].contiguous()
    v_src = flat[:, qs + ks :].contiguous()

    k_cache_init, v_cache_init = make_kv_caches()
    k_cache_a, v_cache_a = k_cache_init.clone(), v_cache_init.clone()
    k_scale_a = torch.zeros(
        [num_blocks, num_heads_k, page_size], dtype=torch.float32, device="cuda"
    )
    v_scale_a = torch.zeros(
        [num_blocks, num_heads_v, page_size], dtype=torch.float32, device="cuda"
    )
    q_a, k_a, v_a = q_src.clone(), k_src.clone(), v_src.clone()
    aiter.fused_qk_norm_rope_cache_quant_shuffle(
        q_a,
        k_a,
        v_a,
        num_heads_q=num_heads_q,
        num_heads_k=num_heads_k,
        num_heads_v=num_heads_v,
        head_dim=head_size,
        eps=eps,
        qw=qw,
        kw=kw,
        cos_sin_cache=cos_sin,
        is_neox_style=is_neox_style,
        pos_ids=positions,
        k_cache=k_cache_a,
        v_cache=v_cache_a,
        slot_mapping=slot_mapping,
        kv_cache_dtype=kv_cache_dtype,
        k_scale=k_scale_a,
        v_scale=v_scale_a,
    )

    k_cache_b, v_cache_b = k_cache_init.clone(), v_cache_init.clone()
    k_scale_b = torch.zeros_like(k_scale_a)
    v_scale_b = torch.zeros_like(v_scale_a)
    q_b = q_src.clone().view(num_tokens, num_heads_q, head_size)
    k_b = k_src.clone().view(num_tokens, num_heads_k, head_size)
    v_b = v_src.clone().view(num_tokens, num_heads_v, head_size)
    aiter.fused_qk_norm_rope_cache_quant_shuffle(
        q_b,
        k_b,
        v_b,
        num_heads_q=num_heads_q,
        num_heads_k=num_heads_k,
        num_heads_v=num_heads_v,
        head_dim=head_size,
        eps=eps,
        qw=qw,
        kw=kw,
        cos_sin_cache=cos_sin,
        is_neox_style=is_neox_style,
        pos_ids=positions,
        k_cache=k_cache_b,
        v_cache=v_cache_b,
        slot_mapping=slot_mapping,
        kv_cache_dtype=kv_cache_dtype,
        k_scale=k_scale_b,
        v_scale=v_scale_b,
    )

    checkAllclose(q_a, q_b.view_as(q_a), msg="q 2d vs 3d", rtol=1e-2, atol=0.05)
    checkAllclose(k_a, k_b.view_as(k_a), msg="k 2d vs 3d", rtol=1e-2, atol=0.05)
    checkAllclose(v_a, v_b.view_as(v_a), msg="v 2d vs 3d", rtol=1e-2, atol=0.05)
    checkAllclose(
        k_cache_a.float(), k_cache_b.float(), msg="k_cache", rtol=1e-2, atol=0.05
    )
    checkAllclose(
        v_cache_a.float(), v_cache_b.float(), msg="v_cache", rtol=1e-2, atol=0.05
    )
    checkAllclose(k_scale_a, k_scale_b, msg="k_scale", rtol=1e-2, atol=0.05)
    checkAllclose(v_scale_a, v_scale_b, msg="v_scale", rtol=1e-2, atol=0.05)


@perftest(num_iters=2)
def run_torch_qk_norm_rope_cache_block_quant_shuffle(
    qkv: Tensor,  # contiguous (num_tokens * (num_heads_q + num_heads_k + num_heads_v) * head_size)
    qw: Tensor,  #  contiguous (head_size)
    kw: Tensor,  #  contiguous (head_size)
    cos_sin: Tensor,  # contiguous (max_positions * head_size)
    positions: Tensor,  # contiguous (3 * num_tokens) or (num_tokens)
    num_tokens: int,
    num_heads_q: int,
    num_heads_k: int,
    num_heads_v: int,
    head_size: int,
    is_neox_style: bool,
    eps: float,
    k_cache: Tensor,  # [num_blocks, num_heads_k, head_size // x, block_size, x]
    v_cache: Tensor,  # [num_blocks, num_heads_v, block_size // x, head_size, x]
    k_scale: Tensor,  # [num_blocks, num_kv_heads]
    v_scale: Tensor,  # [num_blocks, num_kv_heads]
    slot_mapping: Tensor,
    kv_cache_dtype: str,
):
    q_size = num_heads_q * head_size
    k_size = num_heads_k * head_size
    v_size = num_heads_v * head_size
    qkv = qkv.view(num_tokens, q_size + k_size + v_size)
    q, k, v = qkv.split([q_size, k_size, v_size], dim=-1)

    q_by_head = q.view(num_tokens, num_heads_q, head_size)
    q_by_head = rms_norm_forward(q_by_head, qw, eps)
    q = q_by_head.view(q.shape)

    k_by_head = k.view(num_tokens, num_heads_k, head_size)
    k_by_head = rms_norm_forward(k_by_head, kw, eps)
    k = k_by_head.view(k.shape)

    cos_sin = cos_sin.view(cos_sin.shape[0], head_size)
    cos_sin = cos_sin[positions]
    cos, sin = cos_sin.chunk(2, dim=-1)

    q_shape = q.shape
    q = q.view(num_tokens, -1, head_size)
    q = apply_rotary_emb_dispatch(q, cos, sin, is_neox_style)
    q = q.reshape(q_shape)

    k_shape = k.shape
    k = k.view(num_tokens, -1, head_size)
    k = apply_rotary_emb_dispatch(k, cos, sin, is_neox_style)

    v = v.view(num_tokens, -1, head_size)

    from aiter import reshape_and_cache

    if kv_cache_dtype == "auto":
        reshape_and_cache(
            k,
            v,
            k_cache,
            v_cache,
            slot_mapping,
            kv_cache_dtype,
            None,
            None,
            asm_layout=True,
        )
    else:
        # Block quant ref using pertoken_quant (same approach as test_kvcache_blockscale.py)
        # k_cache: [num_blocks, num_heads_k, head_size // x, block_size, x]
        # v_cache: [num_blocks, num_heads_v, block_size // x, head_size, x]
        num_blocks = k_cache.shape[0]
        block_size = k_cache.shape[-2]  # page_size
        x_val = k_cache.shape[-1]
        cache_dtype = k_cache.dtype

        # Step 1: Unflatten k_cache to [num_blocks, block_size, num_heads_k, head_size]
        # and DEQUANTIZE using old scale (multiply raw fp8 values by scale)
        # k_cache: [num_blocks, num_heads_k, head_size//x, block_size, x]
        #       -> permute(0, 3, 1, 2, 4) -> [num_blocks, block_size, num_heads_k, head_size//x, x]
        #       -> view [num_blocks, block_size, num_heads_k, head_size]
        k_cache_flat = (
            k_cache.float()
            .permute(0, 3, 1, 2, 4)
            .contiguous()
            .view(num_blocks, block_size, num_heads_k, head_size)
        )
        # Dequantize: k_cache_flat *= k_scale[num_blocks, num_heads_k] (broadcast over block_size, head_size)
        k_cache_flat = k_cache_flat * k_scale.view(num_blocks, 1, num_heads_k, 1)
        k_cache_flat = k_cache_flat.view(-1, num_heads_k, head_size)

        # v_cache: [num_blocks, num_heads_v, block_size//x, head_size, x]
        #       -> permute(0, 2, 4, 1, 3) -> [num_blocks, block_size//x, x, num_heads_v, head_size]
        #       -> view [num_blocks, block_size, num_heads_v, head_size]
        v_cache_flat = (
            v_cache.float()
            .permute(0, 2, 4, 1, 3)
            .contiguous()
            .view(num_blocks, block_size, num_heads_v, head_size)
        )
        # Dequantize: v_cache_flat *= v_scale[num_blocks, num_heads_v]
        v_cache_flat = v_cache_flat * v_scale.view(num_blocks, 1, num_heads_v, 1)
        v_cache_flat = v_cache_flat.view(-1, num_heads_v, head_size)

        # Step 2: Scatter K/V into dequantized cache
        k_flat = k.view(-1, num_heads_k, head_size)
        v_flat = v.view(-1, num_heads_v, head_size)
        k_cache_flat[slot_mapping] = k_flat.float()
        v_cache_flat[slot_mapping] = v_flat.float()

        # Step 3: Reshape to [num_blocks, num_heads, block_size*head_size] and pertoken_quant
        k_cache_for_quant = (
            k_cache_flat.view(num_blocks, block_size, num_heads_k, head_size)
            .permute(0, 2, 1, 3)
            .contiguous()
            .view(num_blocks, num_heads_k, -1)
        )
        k_cache_q, k_scale_new = aiter.pertoken_quant(
            k_cache_for_quant,
            scale_dtype=torch.float32,
            quant_dtype=cache_dtype,
        )
        k_scale_new = k_scale_new.view(num_blocks, num_heads_k)

        v_cache_for_quant = (
            v_cache_flat.view(num_blocks, block_size, num_heads_v, head_size)
            .permute(0, 2, 1, 3)
            .contiguous()
            .view(num_blocks, num_heads_v, -1)
        )
        v_cache_q, v_scale_new = aiter.pertoken_quant(
            v_cache_for_quant,
            scale_dtype=torch.float32,
            quant_dtype=cache_dtype,
        )
        v_scale_new = v_scale_new.view(num_blocks, num_heads_v)

        # Step 4: Reshape back to tiled layout
        # k_cache_q: [num_blocks, num_heads_k, block_size*head_size]
        #         -> view [num_blocks, num_heads_k, block_size, head_size//x, x]
        #         -> permute(0, 1, 3, 2, 4) -> [num_blocks, num_heads_k, head_size//x, block_size, x]
        k_cache.copy_(
            k_cache_q.view(
                num_blocks, num_heads_k, block_size, head_size // x_val, x_val
            )
            .permute(0, 1, 3, 2, 4)
            .contiguous()
        )

        # v_cache_q: [num_blocks, num_heads_v, block_size*head_size]
        #         -> view [num_blocks, num_heads_v, block_size, head_size]
        #         -> view [num_blocks, num_heads_v, block_size//x, x, head_size]
        #         -> permute(0, 1, 2, 4, 3) -> [num_blocks, num_heads_v, block_size//x, head_size, x]
        v_cache.copy_(
            v_cache_q.view(
                num_blocks, num_heads_v, block_size // x_val, x_val, head_size
            )
            .permute(0, 1, 2, 4, 3)
            .contiguous()
        )

        k_scale.copy_(k_scale_new)
        v_scale.copy_(v_scale_new)

    k = k.reshape(k_shape)
    v = v.reshape(k_shape)
    return q, k, v, k_cache, v_cache


@perftest(num_iters=31, num_rotate_args=31)
def run_aiter_qk_norm_rope_cache_block_quant_shuffle(
    qkv: Tensor,  # contiguous (num_tokens * (num_heads_q + num_heads_k + num_heads_v) * head_size)
    qw: Tensor,  #  contiguous (head_size)
    kw: Tensor,  #  contiguous (head_size)
    cos_sin: Tensor,  # contiguous (max_positions * head_size)
    positions: Tensor,  # contiguous (3 * num_tokens)
    num_tokens: int,
    num_heads_q: int,
    num_heads_k: int,
    num_heads_v: int,
    head_size: int,
    is_neox_style: bool,
    eps: float,
    k_cache: Tensor,
    v_cache: Tensor,
    slot_mapping: Tensor,
    cu_q_len: Tensor,
    kv_cache_dtype: str,
    k_scale: Tensor,
    v_scale: Tensor,
    max_tokens_per_batch: int = 0,
):
    qkv = qkv.clone()  # inplace op

    aiter.fused_qk_norm_rope_cache_block_quant_shuffle(
        qkv,
        num_heads_q,
        num_heads_k,
        num_heads_v,
        head_size,
        eps,
        qw,
        kw,
        cos_sin,
        is_neox_style,
        positions,
        k_cache,
        v_cache,
        slot_mapping,
        cu_q_len,
        kv_cache_dtype,
        k_scale,
        v_scale,
        max_tokens_per_batch,
    )

    q_size = num_heads_q * head_size
    k_size = num_heads_k * head_size
    v_size = num_heads_v * head_size

    qkv = qkv.view(num_tokens, q_size + k_size + v_size)
    q, k, v = qkv.split([q_size, k_size, v_size], dim=-1)
    return q, k, v, k_cache, v_cache


@benchmark()
def test_qk_norm_rope_cache_quant(
    dtype,
    num_tokens,
    num_heads_q,
    num_heads_k,
    num_heads_v,
    head_size,
    is_neox_style,
    eps,
    kv_cache_dtype,
    num_blocks,
    page_size,
    max_positions: int = 10000,
    rotary_dim: int = 0,
):
    # Construct tensors inside the function
    if kv_cache_dtype == "fp8_e4m3":
        cache_dtype = get_dtype_fp8()
    else:
        cache_dtype = dtype
    rotary_dim_ = rotary_dim if rotary_dim > 0 else head_size
    assert rotary_dim_ <= head_size
    assert rotary_dim_ % 2 == 0

    k_cache = torch.randn(
        [num_blocks, page_size, num_heads_k, head_size],
        dtype=dtype,
        device="cuda",
    ).to(cache_dtype)
    v_cache = torch.randn(
        [num_blocks, page_size, num_heads_v, head_size],
        dtype=dtype,
        device="cuda",
    ).to(cache_dtype)

    x = 16 // k_cache.element_size()
    k_cache = (
        k_cache.view([num_blocks, page_size, num_heads_k, head_size // x, x])
        .permute(0, 2, 3, 1, 4)
        .contiguous()
    )
    v_cache = v_cache.permute(0, 2, 3, 1).contiguous()

    slot_mapping = torch.randperm(num_tokens, dtype=torch.int64, device="cuda")
    k_scale = torch.zeros(
        [num_blocks, num_heads_k, page_size], dtype=torch.float32, device="cuda"
    )
    v_scale = torch.zeros(
        [num_blocks, num_heads_v, page_size], dtype=torch.float32, device="cuda"
    )
    qkv = torch.randn(
        (num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )
    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_positions, rotary_dim_), dtype=dtype, device="cuda")
    pos_shape = (num_tokens,)
    positions = torch.randint(
        0, max_positions, pos_shape, dtype=torch.int64, device="cuda"
    )
    k_scale_ref = k_scale.clone()
    v_scale_ref = v_scale.clone()

    (q_ref, k_ref, v_ref, k_cache_ref, v_cache_ref), avg_torch = (
        run_torch_qk_norm_rope_cache_quant_shuffle(
            qkv,
            qw,
            kw,
            cos_sin,
            positions,
            num_tokens,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache,
            v_cache,
            k_scale_ref,
            v_scale_ref,
            slot_mapping,
            kv_cache_dtype,
        )
    )
    q_in, k_in, v_in = split_qkv(
        qkv.clone(), num_heads_q, num_heads_k, num_heads_v, head_size
    )
    (q, k, v, k_cache, v_cache), avg_cu = run_aiter_qk_norm_rope_cache_quant_shuffle(
        q_in,
        k_in,
        v_in,
        qw,
        kw,
        cos_sin,
        positions,
        num_heads_q,
        num_heads_k,
        num_heads_v,
        head_size,
        is_neox_style,
        eps,
        k_cache,
        v_cache,
        slot_mapping,
        kv_cache_dtype,
        k_scale,
        v_scale,
    )

    info = f"dtype:{dtype}, num_tokens:{num_tokens}, num_heads_q:{num_heads_q}, num_heads_k:{num_heads_k}, num_heads_v:{num_heads_v}, head_size:{head_size}, is_neox_style:{is_neox_style}"
    if rotary_dim > 0:
        info += f", rotary_dim:{rotary_dim_}"
    msg = f"[perf] === {info} === torch avg: {avg_torch:<8.2f} us, cu avg: {avg_cu:<8.2f} us, uplift: {avg_torch / avg_cu - 1:<5.1%}"
    checkAllclose(q_ref, q, msg="q", rtol=1e-2, atol=0.05)
    checkAllclose(k_ref, k, msg="k", rtol=1e-2, atol=0.05)
    checkAllclose(v_ref, v, msg=msg, rtol=1e-2, atol=0.05)
    checkAllclose(
        k_cache_ref.float(), k_cache.float(), msg="k_cache", rtol=1e-2, atol=0.05
    )
    checkAllclose(
        v_cache_ref.float(), v_cache.float(), msg="v_cache", rtol=1e-2, atol=0.05
    )
    checkAllclose(k_scale_ref, k_scale, msg="k_scale", rtol=1e-2, atol=0.05)
    checkAllclose(v_scale_ref, v_scale, msg="v_scale", rtol=1e-2, atol=0.05)
    ret = {}
    ret["fused_qk_us"] = avg_cu
    ret["unfused_us"] = avg_torch
    ret["rotary_dim"] = rotary_dim_
    ret["aiter_bw(TB/s)"] = (
        num_tokens
        * (num_heads_k + num_heads_v + num_heads_q)
        * head_size
        * (torch.finfo(dtype).bits // 8)
        + num_tokens * num_heads_q * head_size * (torch.finfo(dtype).bits // 8)
        + num_tokens * num_heads_k * head_size * (torch.finfo(cache_dtype).bits // 8)
        + num_tokens * num_heads_v * head_size * (torch.finfo(cache_dtype).bits // 8)
    ) / (avg_cu * 1e6)
    return ret


@benchmark()
def test_qk_norm_rope_cache_quant_v_shuffle_layout(
    dtype,
    num_tokens,
    num_heads_q,
    num_heads_k,
    num_heads_v,
    head_size,
    is_neox_style,
    eps,
    kv_cache_dtype,
    num_blocks,
    page_size,
    max_positions: int = 10000,
):
    """5D v_cache shuffle: host must use k_cache page_size, not v.size(-1).

    Two aiter runs with independent 5D shuffle buffers must match (q/k/v, caches, scales).
    """
    if kv_cache_dtype == "fp8_e4m3":
        cache_dtype = get_dtype_fp8()
    else:
        cache_dtype = dtype

    k0 = torch.zeros(
        [num_blocks, page_size, num_heads_k, head_size],
        dtype=dtype,
        device="cuda",
    ).to(cache_dtype)
    x = 16 // k0.element_size()
    k_cache = (
        k0.view([num_blocks, page_size, num_heads_k, head_size // x, x])
        .permute(0, 2, 3, 1, 4)
        .contiguous()
    )
    v_cache_5d = (
        torch.zeros(
            [num_blocks, page_size, num_heads_v, head_size],
            dtype=dtype,
            device="cuda",
        )
        .to(cache_dtype)
        .view(num_blocks, page_size // x, num_heads_v, head_size, x)
        .permute(0, 2, 1, 3, 4)
        .contiguous()
    )

    slot_mapping = torch.randperm(num_tokens, dtype=torch.int64, device="cuda")
    k_scale = torch.zeros(
        [num_blocks, num_heads_k, page_size], dtype=torch.float32, device="cuda"
    )
    v_scale = torch.zeros(
        [num_blocks, num_heads_v, page_size], dtype=torch.float32, device="cuda"
    )
    qkv = torch.randn(
        (num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )
    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_positions, head_size), dtype=dtype, device="cuda")
    positions = torch.randint(
        0, max_positions, (num_tokens,), dtype=torch.int64, device="cuda"
    )

    k_scale_a = k_scale.clone()
    v_scale_a = v_scale.clone()
    k_scale_b = k_scale.clone()
    v_scale_b = v_scale.clone()

    q_a_in, k_a_in, v_a_in = split_qkv(
        qkv.clone(), num_heads_q, num_heads_k, num_heads_v, head_size
    )
    q_b_in, k_b_in, v_b_in = split_qkv(
        qkv.clone(), num_heads_q, num_heads_k, num_heads_v, head_size
    )

    (q_a, k_a, v_a, k_cache_a, v_cache_a), avg_a = (
        run_aiter_qk_norm_rope_cache_quant_shuffle(
            q_a_in,
            k_a_in,
            v_a_in,
            qw,
            kw,
            cos_sin,
            positions,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache.clone(),
            v_cache_5d.clone(),
            slot_mapping,
            kv_cache_dtype,
            k_scale_a,
            v_scale_a,
        )
    )

    (q_b, k_b, v_b, k_cache_b, v_cache_b), avg_b = (
        run_aiter_qk_norm_rope_cache_quant_shuffle(
            q_b_in,
            k_b_in,
            v_b_in,
            qw,
            kw,
            cos_sin,
            positions,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache.clone(),
            v_cache_5d.clone(),
            slot_mapping,
            kv_cache_dtype,
            k_scale_b,
            v_scale_b,
        )
    )

    slots_edit = torch.unique(slot_mapping // page_size)
    cache_rtol = 5e-2 if kv_cache_dtype == "fp8_e4m3" else 1e-2
    cache_atol = 0.05

    info = (
        f"v_shuffle dtype:{dtype}, tok:{num_tokens}, Hq:{num_heads_q}, Hkv:{num_heads_k}, "
        f"D:{head_size}, neox:{is_neox_style}, kvd:{kv_cache_dtype}"
    )
    msg = f"[perf] === {info} === run_a {avg_a:.2f} us, run_b {avg_b:.2f} us"
    checkAllclose(q_a, q_b, msg="q " + msg, rtol=1e-2, atol=0.05)
    checkAllclose(k_a, k_b, msg="k " + msg, rtol=1e-2, atol=0.05)
    checkAllclose(v_a, v_b, msg="v " + msg, rtol=1e-2, atol=0.05)
    checkAllclose(
        k_cache_a.float()[slots_edit],
        k_cache_b.float()[slots_edit],
        msg="k_cache 5D " + msg,
        rtol=cache_rtol,
        atol=cache_atol,
    )
    checkAllclose(
        v_cache_a.float()[slots_edit],
        v_cache_b.float()[slots_edit],
        msg="v_cache 5D shuffle " + msg,
        rtol=cache_rtol,
        atol=cache_atol,
    )
    checkAllclose(k_scale_a, k_scale_b, msg="k_scale", rtol=1e-2, atol=0.05)
    checkAllclose(v_scale_a, v_scale_b, msg="v_scale", rtol=1e-2, atol=0.05)
    return {"v_shuffle_us_a": avg_a, "v_shuffle_us_b": avg_b}


@perftest()
def run_torch_qk_norm_rope_2way(
    q0: Tensor,  # contiguous (batch_size * num_tokens0 * num_heads_q * head_size)
    k0: Tensor,  # contiguous (batch_size * num_tokens0 * num_heads_k * head_size)
    q1: Tensor,  # contiguous (batch_size * num_tokens1 * num_heads_q * head_size)
    k1: Tensor,  # contiguous (batch_size * num_tokens1 * num_heads_k * head_size)
    w_q0: Tensor,  # contiguous (head_size)
    w_k0: Tensor,  # contiguous (head_size)
    w_q1: Tensor,  # contiguous (head_size)
    w_k1: Tensor,  # contiguous (head_size)
    cos_sin0: Tensor,  # contiguous (num_tokens0 * head_size)
    cos_sin1: Tensor,  # contiguous (num_tokens1 * head_size)
    batch_size: int,
    num_tokens0: int,
    num_tokens1: int,
    num_heads_q: int,
    num_heads_k: int,
    head_size: int,
    is_interleaved: bool,
    eps: float,
):
    is_neox_style = not is_interleaved
    q0_shape = q0.shape
    k0_shape = k0.shape
    q1_shape = q1.shape
    k1_shape = k1.shape
    q0_by_head = rms_norm_forward(
        q0.view(batch_size, num_tokens0, num_heads_q, head_size), w_q0, eps
    )
    k0_by_head = rms_norm_forward(
        k0.view(batch_size, num_tokens0, num_heads_k, head_size), w_k0, eps
    )
    q1_by_head = rms_norm_forward(
        q1.view(batch_size, num_tokens1, num_heads_q, head_size), w_q1, eps
    )
    k1_by_head = rms_norm_forward(
        k1.view(batch_size, num_tokens1, num_heads_k, head_size), w_k1, eps
    )
    cos_sin0 = cos_sin0.view(num_tokens0, head_size)
    cos_sin1 = cos_sin1.view(num_tokens1, head_size)
    cos0, sin0 = cos_sin0.chunk(2, dim=-1)
    cos1, sin1 = cos_sin1.chunk(2, dim=-1)
    q0 = apply_rotary_emb_torch(q0_by_head, cos0, sin0, is_neox_style)
    k0 = apply_rotary_emb_torch(k0_by_head, cos0, sin0, is_neox_style)
    q1 = apply_rotary_emb_torch(q1_by_head, cos1, sin1, is_neox_style)
    k1 = apply_rotary_emb_torch(k1_by_head, cos1, sin1, is_neox_style)
    q0 = q0.reshape(q0_shape)
    k0 = k0.reshape(k0_shape)
    q1 = q1.reshape(q1_shape)
    k1 = k1.reshape(k1_shape)
    q01 = torch.cat([q0, q1], dim=1)
    k01 = torch.cat([k0, k1], dim=1)
    return q01, k01


@perftest()
def run_fused_qk_norm_rope_2way(
    q0: Tensor,  # contiguous (batch_size * num_tokens0 * num_heads_q * head_size)
    k0: Tensor,  # contiguous (batch_size * num_tokens0 * num_heads_k * head_size)
    q1: Tensor,  # contiguous (batch_size * num_tokens1 * num_heads_q * head_size)
    k1: Tensor,  # contiguous (batch_size * num_tokens1 * num_heads_k * head_size)
    w_q0: Tensor,  # contiguous (head_size)
    w_k0: Tensor,  # contiguous (head_size)
    w_q1: Tensor,  # contiguous (head_size)
    w_k1: Tensor,  # contiguous (head_size)
    cos_sin0: Tensor,  # contiguous (num_tokens0 * head_size)
    cos_sin1: Tensor,  # contiguous (num_tokens1 * head_size)
    batch_size: int,
    num_tokens0: int,
    num_tokens1: int,
    num_heads_q: int,
    num_heads_k: int,
    head_size: int,
    is_interleaved: bool,
    eps: float,
):
    q01 = torch.empty(
        (batch_size, num_tokens0 + num_tokens1, num_heads_q, head_size),
        dtype=q0.dtype,
        device=q0.device,
    )
    k01 = torch.empty(
        (batch_size, num_tokens0 + num_tokens1, num_heads_k, head_size),
        dtype=k0.dtype,
        device=k0.device,
    )
    aiter.fused_qk_norm_rope_2way(
        q0,
        k0,
        q1,
        k1,
        w_q0,
        w_k0,
        w_q1,
        w_k1,
        cos_sin0,
        cos_sin1,
        batch_size,
        num_tokens0,
        num_tokens1,
        num_heads_q,
        num_heads_k,
        head_size,
        is_interleaved,
        eps,
        q01,
        k01,
    )
    return q01, k01


@benchmark()
def test_qk_norm_rope_2way(
    dtype,
    batch_size,
    num_tokens0,
    num_tokens1,
    num_heads_q,
    num_heads_k,
    head_size,
    is_interleaved,
    eps=1e-6,
):
    q0 = torch.randn(
        (batch_size, num_tokens0, num_heads_q, head_size),
        dtype=dtype,
        device="cuda",
    )
    k0 = torch.randn(
        (batch_size, num_tokens0, num_heads_k, head_size),
        dtype=dtype,
        device="cuda",
    )
    q1 = torch.randn(
        (batch_size, num_tokens1, num_heads_q, head_size),
        dtype=dtype,
        device="cuda",
    )
    k1 = torch.randn(
        (batch_size, num_tokens1, num_heads_k, head_size),
        dtype=dtype,
        device="cuda",
    )
    w_q0 = torch.randn(head_size, dtype=dtype, device="cuda")
    w_k0 = torch.randn(head_size, dtype=dtype, device="cuda")
    w_q1 = torch.randn(head_size, dtype=dtype, device="cuda")
    w_k1 = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin0 = torch.randn(
        (num_tokens0, head_size),
        dtype=dtype,
        device="cuda",
    )
    cos_sin1 = torch.randn(
        (num_tokens1, head_size),
        dtype=dtype,
        device="cuda",
    )
    (q01_ref, k01_ref), avg_torch = run_torch_qk_norm_rope_2way(
        q0,
        k0,
        q1,
        k1,
        w_q0,
        w_k0,
        w_q1,
        w_k1,
        cos_sin0,
        cos_sin1,
        batch_size,
        num_tokens0,
        num_tokens1,
        num_heads_q,
        num_heads_k,
        head_size,
        is_interleaved,
        eps,
    )
    (q01, k01), avg_cu = run_fused_qk_norm_rope_2way(
        q0,
        k0,
        q1,
        k1,
        w_q0,
        w_k0,
        w_q1,
        w_k1,
        cos_sin0,
        cos_sin1,
        batch_size,
        num_tokens0,
        num_tokens1,
        num_heads_q,
        num_heads_k,
        head_size,
        is_interleaved,
        eps,
    )

    info = f"dtype:{dtype}, batch_size:{batch_size}, num_tokens0:{num_tokens0}, num_tokens1:{num_tokens1}, num_heads_q:{num_heads_q}, num_heads_k:{num_heads_k}"
    info += f", head_size:{head_size}, is_interleaved:{is_interleaved}, eps:{eps}"
    msg = f"[perf] === {info} === torch avg: {avg_torch:<8.2f} us, cu avg: {avg_cu:<8.2f} us, uplift: {avg_torch/avg_cu-1:<5.1%}"
    checkAllclose(q01_ref, q01, msg="q01", rtol=1e-2, atol=0.05)
    checkAllclose(k01_ref, k01, msg="k01", rtol=1e-2, atol=0.05)
    print(msg, flush=True)

    ret = {}
    ret["dtype"] = dtype
    ret["batch_size"] = batch_size
    ret["num_tokens0"] = num_tokens0
    ret["num_tokens1"] = num_tokens1
    ret["num_heads_q"] = num_heads_q
    ret["num_heads_k"] = num_heads_k
    ret["head_size"] = head_size
    ret["is_interleaved"] = "1" if is_interleaved else "0"
    ret["avg_torch"] = avg_torch
    ret["avg_cu"] = avg_cu
    ret["speedup"] = avg_torch / avg_cu
    return ret


@perftest()
def run_torch_qk_norm_rope_1way(
    q: Tensor,  # contiguous (batch_size * num_tokens * num_heads_q * head_size)
    k: Tensor,  # contiguous (batch_size * num_tokens * num_heads_k * head_size)
    w_q: Tensor,  # contiguous (head_size)
    w_k: Tensor,  # contiguous (head_size)
    cos_sin: Tensor,  # contiguous (num_tokens * head_size)
    batch_size: int,
    num_tokens: int,
    num_heads_q: int,
    num_heads_k: int,
    head_size: int,
    is_interleaved: bool,
    eps: float,
):
    is_neox_style = not is_interleaved
    q_shape = q.shape
    k_shape = k.shape
    q_by_head = rms_norm_diffusers_forward(
        q.view(batch_size, num_tokens, num_heads_q, head_size), w_q, eps
    )
    k_by_head = rms_norm_diffusers_forward(
        k.view(batch_size, num_tokens, num_heads_k, head_size), w_k, eps
    )
    # cos_sin must arrive as fp32 — diffusers / qwen-image-edit reference
    # passes the complex freqs in fp32 to keep the rope multiply precision.
    assert (
        cos_sin.dtype == torch.float32
    ), f"cos_sin must be fp32 to match the diffusers reference, got {cos_sin.dtype}"
    cos_sin = cos_sin.view(num_tokens, head_size)
    cos, sin = cos_sin.chunk(2, dim=-1)
    q = apply_rotary_emb_diffusers(q_by_head, cos, sin, is_neox_style)
    k = apply_rotary_emb_diffusers(k_by_head, cos, sin, is_neox_style)
    q = q.reshape(q_shape)
    k = k.reshape(k_shape)
    return q, k


@perftest()
def run_fused_qk_norm_rope_1way(
    q: Tensor,  # contiguous (batch_size * num_tokens * num_heads_q * head_size)
    k: Tensor,  # contiguous (batch_size * num_tokens * num_heads_k * head_size)
    w_q: Tensor,  # contiguous (head_size)
    w_k: Tensor,  # contiguous (head_size)
    cos_sin: Tensor,  # contiguous (num_tokens * head_size)
    batch_size: int,
    num_tokens: int,
    num_heads_q: int,
    num_heads_k: int,
    head_size: int,
    is_interleaved: bool,
    eps: float,
):
    out_q = torch.empty(
        (batch_size, num_tokens, num_heads_q, head_size),
        dtype=q.dtype,
        device=q.device,
    )
    out_k = torch.empty(
        (batch_size, num_tokens, num_heads_k, head_size),
        dtype=k.dtype,
        device=k.device,
    )
    aiter.fused_qk_norm_rope_1way(
        q,
        k,
        w_q,
        w_k,
        cos_sin,
        batch_size,
        num_tokens,
        num_heads_q,
        num_heads_k,
        head_size,
        is_interleaved,
        eps,
        out_q,
        out_k,
    )
    return out_q, out_k


@benchmark()
def test_qk_norm_rope_1way(
    dtype,
    batch_size,
    num_tokens,
    num_heads_q,
    num_heads_k,
    head_size,
    is_interleaved,
    eps=1e-6,
):
    q = torch.randn(
        (batch_size, num_tokens, num_heads_q, head_size),
        dtype=dtype,
        device="cuda",
    )
    k = torch.randn(
        (batch_size, num_tokens, num_heads_k, head_size),
        dtype=dtype,
        device="cuda",
    )
    w_q = torch.randn(head_size, dtype=dtype, device="cuda")
    w_k = torch.randn(head_size, dtype=dtype, device="cuda")
    # cos_sin is fp32 to match the kernel's new dtype contract (kernel will
    # TORCH_CHECK; diffusers reference also expects fp32).
    cos_sin = torch.randn(
        (num_tokens, head_size),
        dtype=torch.float32,
        device="cuda",
    )
    (q_ref, k_ref), avg_torch = run_torch_qk_norm_rope_1way(
        q,
        k,
        w_q,
        w_k,
        cos_sin,
        batch_size,
        num_tokens,
        num_heads_q,
        num_heads_k,
        head_size,
        is_interleaved,
        eps,
    )
    (q_out, k_out), avg_cu = run_fused_qk_norm_rope_1way(
        q,
        k,
        w_q,
        w_k,
        cos_sin,
        batch_size,
        num_tokens,
        num_heads_q,
        num_heads_k,
        head_size,
        is_interleaved,
        eps,
    )

    info = f"dtype:{dtype}, batch_size:{batch_size}, num_tokens:{num_tokens}, num_heads_q:{num_heads_q}, num_heads_k:{num_heads_k}"
    info += f", head_size:{head_size}, is_interleaved:{is_interleaved}, eps:{eps}"
    msg = f"[perf] === {info} === torch avg: {avg_torch:<8.2f} us, cu avg: {avg_cu:<8.2f} us, uplift: {avg_torch/avg_cu-1:<5.1%}"
    checkAllclose(q_ref, q_out, msg="q", rtol=1e-2, atol=0.05)
    checkAllclose(k_ref, k_out, msg="k", rtol=1e-2, atol=0.05)
    print(msg, flush=True)

    ret = {}
    ret["dtype"] = dtype
    ret["batch_size"] = batch_size
    ret["num_tokens"] = num_tokens
    ret["num_heads_q"] = num_heads_q
    ret["num_heads_k"] = num_heads_k
    ret["head_size"] = head_size
    ret["is_interleaved"] = "1" if is_interleaved else "0"
    ret["avg_torch"] = avg_torch
    ret["avg_cu"] = avg_cu
    ret["speedup"] = avg_torch / avg_cu
    return ret


def minimax_rms_norm_forward(x: Tensor, weight: Tensor, eps: float) -> Tensor:
    input_dtype = x.dtype
    variance = x.float().pow(2).mean(-1, keepdim=True)
    x = x.float() * torch.rsqrt(variance + eps) * weight.float()
    return x.to(input_dtype)


@perftest()
def run_torch_minimax_qk_norm_rope(
    qkv: Tensor,
    q_weight: Tensor,
    k_weight: Tensor,
    cos_sin: Tensor,
    positions: Tensor,
    num_tokens: int,
    num_heads_q: int,
    num_heads_k: int,
    head_size: int,
    rotary_dim: int,
    is_neox_style: bool,
    eps: float,
):
    q_size = num_heads_q * head_size
    kv_size = num_heads_k * head_size
    q, k, v = qkv.view(num_tokens, q_size + 2 * kv_size).split(
        [q_size, kv_size, kv_size], dim=-1
    )

    q = minimax_rms_norm_forward(q, q_weight, eps)
    k = minimax_rms_norm_forward(k, k_weight, eps)

    cos_sin = cos_sin.view(cos_sin.shape[0], rotary_dim)[positions]
    cos, sin = cos_sin.chunk(2, dim=-1)
    q = apply_rotary_emb_dispatch(
        q.view(num_tokens, num_heads_q, head_size),
        cos,
        sin,
        is_neox_style,
        rotary_dim,
    ).reshape(num_tokens, q_size)
    k = apply_rotary_emb_dispatch(
        k.view(num_tokens, num_heads_k, head_size),
        cos,
        sin,
        is_neox_style,
        rotary_dim,
    ).reshape(num_tokens, kv_size)
    return q, k, v


@perftest()
def run_aiter_minimax_qk_norm_rope(
    qkv: Tensor,
    q_weight: Tensor,
    k_weight: Tensor,
    cos_sin: Tensor,
    positions: Tensor,
    num_heads_q: int,
    num_heads_k: int,
    head_size: int,
    rotary_dim: int,
    is_neox_style: bool,
    eps: float,
):
    num_tokens = qkv.size(0)
    q_out = torch.empty(
        (num_tokens, num_heads_q * head_size), dtype=qkv.dtype, device=qkv.device
    )
    k_out = torch.empty(
        (num_tokens, num_heads_k * head_size), dtype=qkv.dtype, device=qkv.device
    )
    v_out = torch.empty_like(k_out)
    q_ret, k_ret, v_ret = aiter.minimax_qk_norm_rope(
        qkv,
        q_weight,
        k_weight,
        cos_sin,
        positions,
        num_heads_q=num_heads_q,
        num_heads_k=num_heads_k,
        head_dim=head_size,
        rotary_dim=rotary_dim,
        eps=eps,
        is_neox_style=is_neox_style,
        q_out=q_out,
        k_out=k_out,
        v_out=v_out,
    )
    assert q_ret.data_ptr() == q_out.data_ptr()
    assert k_ret.data_ptr() == k_out.data_ptr()
    assert v_ret.data_ptr() == v_out.data_ptr()
    return q_ret, k_ret, v_ret


@benchmark()
def test_minimax_qk_norm_rope(
    dtype,
    num_tokens,
    num_heads_q,
    num_heads_k,
    head_size,
    rotary_dim,
    is_neox_style,
    eps=1e-6,
):
    torch.manual_seed(0)
    q_size = num_heads_q * head_size
    kv_size = num_heads_k * head_size
    max_positions = max(num_tokens * 2, 128)
    qkv = torch.randn(
        (num_tokens, q_size + 2 * kv_size),
        dtype=dtype,
        device="cuda",
    )
    q_weight = torch.randn(q_size, dtype=torch.float32, device="cuda")
    k_weight = torch.randn(kv_size, dtype=torch.float32, device="cuda")
    cos_sin = torch.randn(
        (max_positions, rotary_dim),
        dtype=dtype,
        device="cuda",
    )
    positions = torch.randint(
        0, max_positions, (num_tokens,), dtype=torch.int64, device="cuda"
    )

    (q_ref, k_ref, v_ref), avg_torch = run_torch_minimax_qk_norm_rope(
        qkv,
        q_weight,
        k_weight,
        cos_sin,
        positions,
        num_tokens,
        num_heads_q,
        num_heads_k,
        head_size,
        rotary_dim,
        is_neox_style,
        eps,
    )
    (q_out, k_out, v_out), avg_cu = run_aiter_minimax_qk_norm_rope(
        qkv,
        q_weight,
        k_weight,
        cos_sin,
        positions,
        num_heads_q,
        num_heads_k,
        head_size,
        rotary_dim,
        is_neox_style,
        eps,
    )

    info = f"dtype:{dtype}, num_tokens:{num_tokens}, num_heads_q:{num_heads_q}, num_heads_k:{num_heads_k}"
    info += f", head_size:{head_size}, rotary_dim:{rotary_dim}, is_neox_style:{is_neox_style}, eps:{eps}"
    msg = f"[perf] === {info} === torch avg: {avg_torch:<8.2f} us, cu avg: {avg_cu:<8.2f} us, uplift: {avg_torch/avg_cu-1:<5.1%}"
    checkAllclose(q_ref, q_out, msg="minimax q", rtol=1e-2, atol=0.05)
    checkAllclose(k_ref, k_out, msg="minimax k", rtol=1e-2, atol=0.05)
    checkAllclose(v_ref, v_out, msg="minimax v", rtol=1e-2, atol=0.05)
    print(msg, flush=True)

    ret = {}
    ret["dtype"] = dtype
    ret["num_tokens"] = num_tokens
    ret["num_heads_q"] = num_heads_q
    ret["num_heads_k"] = num_heads_k
    ret["head_size"] = head_size
    ret["rotary_dim"] = rotary_dim
    ret["is_neox_style"] = "1" if is_neox_style else "0"
    ret["avg_torch"] = avg_torch
    ret["avg_cu"] = avg_cu
    ret["speedup"] = avg_torch / avg_cu
    return ret


@benchmark()
def test_qk_norm_rope_cache_block_quant(
    dtype,
    num_tokens,
    num_heads_q,
    num_heads_k,
    num_heads_v,
    head_size,
    num_blocks,
    page_size,
    is_neox_style,
    eps,
    kv_cache_dtype,
    batch=1,
    decode_tokens_per_batch=1,
):
    torch.manual_seed(0)
    # Initialize cache to zeros so unused slots have 0 -> ref pertoken_quant
    # matches kernel (kernel only writes pages with tokens; randn would pollute scale)
    if kv_cache_dtype == "fp8_e4m3":
        cache_dtype = get_dtype_fp8()
    else:
        cache_dtype = dtype
    # Zeros init: unused slots stay 0 so ref's pertoken_quant max = kernel's block_max
    k_cache = torch.zeros(
        [num_blocks, page_size, num_heads_k, head_size],
        dtype=dtype,
        device="cuda",
    ).to(cache_dtype)
    v_cache = torch.zeros(
        [num_blocks, page_size, num_heads_v, head_size],
        dtype=dtype,
        device="cuda",
    ).to(cache_dtype)

    # Check for NaN values in k_cache and v_cache
    if torch.isnan(k_cache).any():
        aiter.logger.warning(f"k_cache contains NaN values! dtype={cache_dtype}")
    if torch.isnan(v_cache).any():
        aiter.logger.warning(f"v_cache contains NaN values! dtype={cache_dtype}")

    # slot_mapping built after cu_q_len (see below)
    x = 16 // k_cache.element_size()
    k_cache = (
        k_cache.view(
            [
                num_blocks,
                page_size,
                num_heads_k,
                head_size // x,
                x,
            ]
        )
        .permute(0, 2, 3, 1, 4)
        .contiguous()
    )
    # Value cache [num_blocks, num_kv_heads, block_size // x, head_size, x]
    v_cache = (
        v_cache.view(
            [
                num_blocks,
                page_size // x,
                num_heads_v,
                head_size,
                x,
            ]
        )
        .permute(0, 2, 1, 3, 4)
        .contiguous()
    )
    batch_size = batch
    base_len = num_tokens // batch_size
    remainder = num_tokens % batch_size
    seq_lens = [base_len + 1] * remainder + [base_len] * (batch_size - remainder)
    total_len = sum(seq_lens)
    if total_len > num_tokens:
        seq_lens[-1] -= total_len - num_tokens
    elif total_len < num_tokens:
        seq_lens[-1] += num_tokens - total_len
    max_tpb = max(seq_lens)
    #
    cu_q_len = torch.zeros(batch_size + 1, dtype=torch.int64, device="cuda")

    cu_q_len[0] = 0
    for i in range(batch_size):
        cu_q_len[i + 1] = cu_q_len[i] + seq_lens[i]
    #
    assert (
        cu_q_len[-1].item() == num_tokens
    ), f"cu_q_len[-1]={cu_q_len[-1].item()} != num_tokens={num_tokens}"
    #
    # slot_mapping: each batch maps to disjoint blocks (no cross-batch block sharing)
    slot_start_per_batch = []
    next_slot = 0
    for i in range(batch_size):
        slot_start_per_batch.append(next_slot)
        blocks_needed = (seq_lens[i] + page_size - 1) // page_size
        next_slot += blocks_needed * page_size
    prefill_slot_end = next_slot  # chunk/decode must allocate after this
    slot_mapping = torch.zeros(num_tokens, dtype=torch.int64, device="cuda")
    for i in range(batch_size):
        start = cu_q_len[i].item()
        end = cu_q_len[i + 1].item()
        slot_mapping[start:end] = torch.arange(
            slot_start_per_batch[i],
            slot_start_per_batch[i] + (end - start),
            dtype=torch.int64,
            device="cuda",
        )
    qkv = torch.randn(
        (num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )

    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_positions, head_size), dtype=dtype, device="cuda")
    pos_shape = (num_tokens,)
    positions = torch.randint(
        0, max_positions, pos_shape, dtype=torch.int64, device="cuda"
    )
    k_scale = torch.zeros(
        [num_blocks, num_heads_k],
        dtype=torch.float32,
        device="cuda",
    )
    v_scale = torch.zeros(
        [num_blocks, num_heads_v],
        dtype=torch.float32,
        device="cuda",
    )
    k_scale_ref = k_scale.clone()
    v_scale_ref = v_scale.clone()
    ## Use separate caches so ref and aiter each write to their own (avoid self-compare)
    k_cache_ref = k_cache.clone()
    v_cache_ref = v_cache.clone()

    (q_ref, k_ref, v_ref, k_cache_ref, v_cache_ref), avg_torch = (
        run_torch_qk_norm_rope_cache_block_quant_shuffle(
            qkv,
            qw,
            kw,
            cos_sin,
            positions,
            num_tokens,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache_ref,
            v_cache_ref,
            k_scale_ref,
            v_scale_ref,
            slot_mapping,
            kv_cache_dtype,
        )
    )
    (q, k, v, k_cache, v_cache), avg_cu = (
        run_aiter_qk_norm_rope_cache_block_quant_shuffle(
            qkv,
            qw,
            kw,
            cos_sin,
            positions,
            num_tokens,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache,
            v_cache,
            slot_mapping,
            cu_q_len,
            kv_cache_dtype,
            k_scale,
            v_scale,
            max_tokens_per_batch=max_tpb,
        )
    )

    info = f"dtype:{dtype}, batch:{batch}, num_tokens:{num_tokens}, num_heads_q:{num_heads_q}, num_heads_k:{num_heads_k}, num_heads_v:{num_heads_v}, head_size:{head_size}, is_neox_style:{is_neox_style}"
    msg = f"[perf] === {info} === torch avg: {avg_torch:<8.2f} us, cu avg: {avg_cu:<8.2f} us, uplift: {avg_torch / avg_cu - 1:<5.1%}"
    checkAllclose(q_ref[:64,], q[:64,], msg="prefill q", rtol=1e-2, atol=0.05)
    checkAllclose(k_ref, k, msg="prefill k", rtol=1e-2, atol=0.05)
    checkAllclose(v_ref, v, msg=msg, rtol=1e-2, atol=0.05)
    # Only check pages that have actual token data (via slot_mapping)
    page_size = k_cache.shape[
        -2
    ]  # k_cache: [num_blocks, num_kv_heads, head_size//x, page_size, x]
    slots_edit = torch.unique(slot_mapping // page_size)
    checkAllclose(
        k_cache_ref.float()[slots_edit],
        k_cache.float()[slots_edit],
        msg="prefill k_cache",
        rtol=5e-2,
        atol=0.05,
    )
    checkAllclose(
        v_cache_ref.float()[slots_edit],
        v_cache.float()[slots_edit],
        msg="prefill v_cache",
        rtol=5e-2,
        atol=0.05,
    )

    checkAllclose(
        k_scale_ref[slots_edit],
        k_scale[slots_edit],
        msg="prefill k_scale",
        rtol=1e-2,
        atol=0.01,
    )
    checkAllclose(
        v_scale_ref[slots_edit],
        v_scale[slots_edit],
        msg="prefill v_scale",
        rtol=1e-2,
        atol=0.01,
    )

    ret = {}
    ret["fused_qk_us"] = avg_cu
    # ret["k_cache_err"] = k_cache_err
    # ret["v_cache_err"] = v_cache_err
    # ========== chunk-prefill part ==========
    # Chunk: (page_size-1) tokens per batch, K=0.001 + kw=1 -> k_scale_global small.
    # Decode fills last slot per block -> k_scale_val > k_scale_global triggers requantization.
    # batch_size = cu_q_len.size(0) - 1
    page_size = k_cache.shape[
        -2
    ]  # k_cache: [num_blocks, num_kv_heads, head_size//x, page_size, x]
    chunk_left_ctx_lens = page_size - 1
    chunk_total_tokens = batch_size * chunk_left_ctx_lens
    #
    chunk_qkv = torch.randn(
        (chunk_total_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )
    q_size, k_size, v_size = (
        num_heads_q * head_size,
        num_heads_k * head_size,
        num_heads_v * head_size,
    )
    ## to test requant in worst case
    # chunk_qkv[:, q_size : q_size + k_size] = 1e-6  # K ~0 -> k_scale_global ~2e-6
    # chunk_qkv[:, q_size + k_size : q_size + k_size + v_size] = (
    #   1  # V ~0 -> v_scale_global ~2e-6
    # )
    kw = torch.ones(
        head_size, dtype=dtype, device="cuda"
    )  # kw=1 so k_scale_global small
    # slot_mapping: each batch fills slots [0..page_size-2] of its chunk block
    chunk_slot_mapping = (
        (
            prefill_slot_end
            + torch.arange(batch_size, device="cuda").unsqueeze(1) * page_size
            + torch.arange(chunk_left_ctx_lens, device="cuda").unsqueeze(0)
        )
        .reshape(-1)
        .to(torch.int64)
    )
    chunk_cu_q_len = torch.zeros(batch_size + 1, dtype=torch.int64, device="cuda")
    for i in range(batch_size):
        chunk_cu_q_len[i + 1] = chunk_cu_q_len[i] + chunk_left_ctx_lens
    chunk_positions = torch.randint(
        0, max_positions, (chunk_total_tokens,), dtype=torch.int64, device="cuda"
    )

    k_scale_chunk_ref = k_scale_ref.clone()
    v_scale_chunk_ref = v_scale_ref.clone()
    k_scale_chunk = k_scale.clone()
    v_scale_chunk = v_scale.clone()
    #
    (
        q_chunk_ref,
        k_chunk_ref,
        v_chunk_ref,
        k_cache_ref,
        v_cache_ref,
    ), avg_torch_chunk = run_torch_qk_norm_rope_cache_block_quant_shuffle(
        chunk_qkv,
        qw,
        kw,
        cos_sin,
        chunk_positions,
        chunk_total_tokens,
        num_heads_q,
        num_heads_k,
        num_heads_v,
        head_size,
        is_neox_style,
        eps,
        k_cache_ref,
        v_cache_ref,
        k_scale_chunk_ref,
        v_scale_chunk_ref,
        chunk_slot_mapping,
        kv_cache_dtype,
    )
    (q_chunk, k_chunk, v_chunk, k_cache, v_cache), avg_cu_chunk = (
        run_aiter_qk_norm_rope_cache_block_quant_shuffle(
            chunk_qkv,
            qw,
            kw,
            cos_sin,
            chunk_positions,
            chunk_total_tokens,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache,
            v_cache,
            chunk_slot_mapping,
            chunk_cu_q_len,
            kv_cache_dtype,
            k_scale_chunk,
            v_scale_chunk,
            max_tokens_per_batch=chunk_left_ctx_lens,
        )
    )
    #
    print(
        f"chunk-prefill: torch avg: {avg_torch_chunk:.2f} us, cu avg: {avg_cu_chunk:.2f} us"
    )
    checkAllclose(q_chunk_ref, q_chunk, msg="chunk q", rtol=1e-2, atol=0.05)
    checkAllclose(k_chunk_ref, k_chunk, msg="chunk k", rtol=1e-2, atol=0.05)
    checkAllclose(v_chunk_ref, v_chunk, msg="chunk v", rtol=1e-2, atol=0.05)
    # Combine prefill + chunk slots to check all pages with data
    all_slots_so_far = torch.cat([slot_mapping, chunk_slot_mapping])
    chunk_slots_edit = torch.unique(all_slots_so_far // page_size)
    chunk_k_cache_err = checkAllclose(
        k_cache_ref.float()[chunk_slots_edit],
        k_cache.float()[chunk_slots_edit],
        msg="chunk k_cache",
        rtol=5e-2,
        atol=0.05,
    )
    chunk_v_cache_err = checkAllclose(
        v_cache_ref.float()[chunk_slots_edit],
        v_cache.float()[chunk_slots_edit],
        msg="chunk v_cache",
        rtol=5e-2,
        atol=0.05,
    )
    checkAllclose(
        k_scale_chunk_ref[chunk_slots_edit],
        k_scale_chunk[chunk_slots_edit],
        msg="chunk k_scale",
        rtol=1e-2,
        atol=0.01,
    )
    checkAllclose(
        v_scale_chunk_ref[chunk_slots_edit],
        v_scale_chunk[chunk_slots_edit],
        msg="chunk v_scale",
        rtol=1e-2,
        atol=0.01,
    )
    # ret["chunk_fused_qk_us"] = avg_cu_chunk
    # ret["chunk_k_cache_err"] = chunk_k_cache_err
    # ret["chunk_v_cache_err"] = chunk_v_cache_err
    #
    # ========== decode test 1: 1 token per batch (append after chunk's last slot) ==========
    dtpb = decode_tokens_per_batch
    decode1_total = batch_size
    decode1_slot_mapping = torch.tensor(
        [
            int(chunk_slot_mapping[(bsID + 1) * chunk_left_ctx_lens - 1].item()) + 1
            for bsID in range(batch_size)
        ],
        dtype=torch.int64,
        device="cuda",
    )
    decode1_qkv = torch.randn(
        (decode1_total, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )
    decode1_qkv[:, q_size : q_size + k_size] = torch.randn(
        (decode1_total, k_size), dtype=dtype, device="cuda"
    )
    decode1_qkv[:, q_size + k_size : q_size + k_size + v_size] = (
        torch.randn((decode1_total, v_size), dtype=dtype, device="cuda") * 2.0 + 2.0
    )
    decode1_cu_q_len = torch.zeros(batch_size + 1, dtype=torch.int64, device="cuda")
    for i in range(batch_size):
        decode1_cu_q_len[i + 1] = decode1_cu_q_len[i] + 1
    decode1_positions = torch.randint(
        0, max_positions, (decode1_total,), dtype=torch.int64, device="cuda"
    )
    k_scale_d1_ref = k_scale_chunk_ref.clone()
    v_scale_d1_ref = v_scale_chunk_ref.clone()
    k_scale_d1 = k_scale_chunk.clone()
    v_scale_d1 = v_scale_chunk.clone()

    print(
        f"decode1: 1 tok/batch, total={decode1_total}"
    )  # , slots={decode1_slot_mapping}
    (q_d1_ref, k_d1_ref, v_d1_ref, k_cache_ref, v_cache_ref), avg_torch_d1 = (
        run_torch_qk_norm_rope_cache_block_quant_shuffle(
            decode1_qkv,
            qw,
            kw,
            cos_sin,
            decode1_positions,
            decode1_total,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache_ref,
            v_cache_ref,
            k_scale_d1_ref,
            v_scale_d1_ref,
            decode1_slot_mapping,
            kv_cache_dtype,
        )
    )
    (q_d1, k_d1, v_d1, k_cache, v_cache), avg_cu_d1 = (
        run_aiter_qk_norm_rope_cache_block_quant_shuffle(
            decode1_qkv,
            qw,
            kw,
            cos_sin,
            decode1_positions,
            decode1_total,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache,
            v_cache,
            decode1_slot_mapping,
            decode1_cu_q_len,
            kv_cache_dtype,
            k_scale_d1,
            v_scale_d1,
            max_tokens_per_batch=1,
        )
    )
    print(f"decode1: torch avg: {avg_torch_d1:.2f} us, cu avg: {avg_cu_d1:.2f} us")
    checkAllclose(q_d1_ref, q_d1, msg="decode1 q", rtol=1e-2, atol=0.05)
    checkAllclose(k_d1_ref, k_d1, msg="decode1 k", rtol=1e-2, atol=0.05)
    checkAllclose(v_d1_ref, v_d1, msg="decode1 v", rtol=1e-2, atol=0.05)
    all_slots_d1 = torch.cat([slot_mapping, chunk_slot_mapping, decode1_slot_mapping])
    d1_pages = torch.unique(all_slots_d1 // page_size)
    d1_k_err = checkAllclose(
        k_cache_ref.float()[d1_pages],
        k_cache.float()[d1_pages],
        msg="decode1 k_cache",
        rtol=5e-2,
        atol=0.05,
    )
    d1_v_err = checkAllclose(
        v_cache_ref.float()[d1_pages],
        v_cache.float()[d1_pages],
        msg="decode1 v_cache",
        rtol=5e-2,
        atol=0.05,
    )
    checkAllclose(
        k_scale_d1_ref[d1_pages],
        k_scale_d1[d1_pages],
        msg="decode1 k_scale",
        rtol=1e-2,
        atol=0.05,
    )
    checkAllclose(
        v_scale_d1_ref[d1_pages],
        v_scale_d1[d1_pages],
        msg="decode1 v_scale",
        rtol=1e-2,
        atol=0.05,
    )
    ret["decode1_fused_qk_us"] = avg_cu_d1
    ret["decode1_k_cache_err"] = d1_k_err
    ret["decode1_v_cache_err"] = d1_v_err

    # ========== decode test 2: dtpb tokens per batch (cross-page, only when dtpb > 1) ==========
    if dtpb > 1:
        decode2_total = batch_size * dtpb
        last_used_slot = int(chunk_slot_mapping[-1].item())
        decode2_page_base = (last_used_slot + page_size) // page_size * page_size
        num_blocks * page_size
        pages_needed = batch_size * 2 + (decode2_page_base // page_size)
        assert (
            pages_needed <= num_blocks
        ), f"decode2 needs {pages_needed} pages but num_blocks={num_blocks}. Increase -b."
        decode2_slots = []
        for bsID in range(batch_size):
            start_slot = (
                decode2_page_base + bsID * 2 * page_size + (page_size - 1)
            )  # (dtpb // 2)
            for t in range(dtpb):
                decode2_slots.append(start_slot + t)
        decode2_slot_mapping = torch.tensor(
            decode2_slots, dtype=torch.int64, device="cuda"
        )
        decode2_qkv = torch.randn(
            (decode2_total, (num_heads_q + num_heads_k + num_heads_v) * head_size),
            dtype=dtype,
            device="cuda",
        )
        decode2_qkv[:, q_size : q_size + k_size] = torch.randn(
            (decode2_total, k_size), dtype=dtype, device="cuda"
        )
        decode2_qkv[:, q_size + k_size : q_size + k_size + v_size] = (
            torch.randn((decode2_total, v_size), dtype=dtype, device="cuda") * 2.0 + 2.0
        )
        decode2_cu_q_len = torch.zeros(batch_size + 1, dtype=torch.int64, device="cuda")
        for i in range(batch_size):
            decode2_cu_q_len[i + 1] = decode2_cu_q_len[i] + dtpb
        decode2_positions = torch.randint(
            0, max_positions, (decode2_total,), dtype=torch.int64, device="cuda"
        )
        k_scale_d2_ref = k_scale_d1_ref.clone()
        v_scale_d2_ref = v_scale_d1_ref.clone()
        k_scale_d2 = k_scale_d1.clone()
        v_scale_d2 = v_scale_d1.clone()

        print(
            f"decode2: {dtpb} tok/batch (cross-page), total={decode2_total}"
        )  # , slots={decode2_slot_mapping}
        (q_d2_ref, k_d2_ref, v_d2_ref, k_cache_ref, v_cache_ref), avg_torch_d2 = (
            run_torch_qk_norm_rope_cache_block_quant_shuffle(
                decode2_qkv,
                qw,
                kw,
                cos_sin,
                decode2_positions,
                decode2_total,
                num_heads_q,
                num_heads_k,
                num_heads_v,
                head_size,
                is_neox_style,
                eps,
                k_cache_ref,
                v_cache_ref,
                k_scale_d2_ref,
                v_scale_d2_ref,
                decode2_slot_mapping,
                kv_cache_dtype,
            )
        )
        (q_d2, k_d2, v_d2, k_cache, v_cache), avg_cu_d2 = (
            run_aiter_qk_norm_rope_cache_block_quant_shuffle(
                decode2_qkv,
                qw,
                kw,
                cos_sin,
                decode2_positions,
                decode2_total,
                num_heads_q,
                num_heads_k,
                num_heads_v,
                head_size,
                is_neox_style,
                eps,
                k_cache,
                v_cache,
                decode2_slot_mapping,
                decode2_cu_q_len,
                kv_cache_dtype,
                k_scale_d2,
                v_scale_d2,
                max_tokens_per_batch=dtpb,
            )
        )
        print(f"decode2: torch avg: {avg_torch_d2:.2f} us, cu avg: {avg_cu_d2:.2f} us")
        checkAllclose(q_d2_ref, q_d2, msg="decode2 q", rtol=1e-2, atol=0.05)
        checkAllclose(k_d2_ref, k_d2, msg="decode2 k", rtol=1e-2, atol=0.05)
        checkAllclose(v_d2_ref, v_d2, msg="decode2 v", rtol=1e-2, atol=0.05)
        all_slots_d2 = torch.cat([all_slots_d1, decode2_slot_mapping])
        d2_pages = torch.unique(all_slots_d2 // page_size)
        checkAllclose(
            k_cache_ref.float()[d2_pages],
            k_cache.float()[d2_pages],
            msg="decode2 k_cache",
            rtol=5e-2,
            atol=0.05,
        )
        checkAllclose(
            v_cache_ref.float()[d2_pages],
            v_cache.float()[d2_pages],
            msg="decode2 v_cache",
            rtol=5e-2,
            atol=0.05,
        )
        checkAllclose(
            k_scale_d2_ref[d2_pages],
            k_scale_d2[d2_pages],
            msg="decode2 k_scale",
            rtol=1e-2,
            atol=0.05,
        )
        checkAllclose(
            v_scale_d2_ref[d2_pages],
            v_scale_d2[d2_pages],
            msg="decode2 v_scale",
            rtol=1e-2,
            atol=0.05,
        )
        ret["decode2_fused_qk_us"] = avg_cu_d2
        # ret["decode2_k_cache_err"] = d2_k_err
        # ret["decode2_v_cache_err"] = d2_v_err

    # ========== decode test 3: dtpb tokens per batch from block_offset=0 (fresh page) ==========
    if dtpb > 1:
        decode3_total = batch_size * dtpb

        decode3_page_base = decode2_page_base + batch_size * 2 * page_size
        pages_needed_d3 = batch_size + (decode3_page_base // page_size)
        assert (
            pages_needed_d3 <= num_blocks
        ), f"decode3 needs {pages_needed_d3} pages but num_blocks={num_blocks}. Increase -b."
        decode3_slots = []
        for bsID in range(batch_size):
            base_slot = decode3_page_base + bsID * page_size
            for t in range(dtpb):
                decode3_slots.append(base_slot + t)
        decode3_slot_mapping = torch.tensor(
            decode3_slots, dtype=torch.int64, device="cuda"
        )
        decode3_qkv = torch.randn(
            (decode3_total, (num_heads_q + num_heads_k + num_heads_v) * head_size),
            dtype=dtype,
            device="cuda",
        )
        decode3_qkv[:, q_size : q_size + k_size] = torch.randn(
            (decode3_total, k_size), dtype=dtype, device="cuda"
        )
        decode3_qkv[:, q_size + k_size : q_size + k_size + v_size] = (
            torch.randn((decode3_total, v_size), dtype=dtype, device="cuda") * 2.0 + 2.0
        )
        decode3_cu_q_len = torch.zeros(batch_size + 1, dtype=torch.int64, device="cuda")
        for i in range(batch_size):
            decode3_cu_q_len[i + 1] = decode3_cu_q_len[i] + dtpb
        decode3_positions = torch.randint(
            0, max_positions, (decode3_total,), dtype=torch.int64, device="cuda"
        )
        k_scale_d3_ref = k_scale_d1_ref.clone()
        v_scale_d3_ref = v_scale_d1_ref.clone()
        k_scale_d3 = k_scale_d1.clone()
        v_scale_d3 = v_scale_d1.clone()

        print(
            f"decode3: {dtpb} tok/batch (offset=0, fresh page), total={decode3_total}, slots={decode3_slot_mapping[:16]}..."
        )
        (q_d3_ref, k_d3_ref, v_d3_ref, k_cache_ref, v_cache_ref), avg_torch_d3 = (
            run_torch_qk_norm_rope_cache_block_quant_shuffle(
                decode3_qkv,
                qw,
                kw,
                cos_sin,
                decode3_positions,
                decode3_total,
                num_heads_q,
                num_heads_k,
                num_heads_v,
                head_size,
                is_neox_style,
                eps,
                k_cache_ref,
                v_cache_ref,
                k_scale_d3_ref,
                v_scale_d3_ref,
                decode3_slot_mapping,
                kv_cache_dtype,
            )
        )
        (q_d3, k_d3, v_d3, k_cache, v_cache), avg_cu_d3 = (
            run_aiter_qk_norm_rope_cache_block_quant_shuffle(
                decode3_qkv,
                qw,
                kw,
                cos_sin,
                decode3_positions,
                decode3_total,
                num_heads_q,
                num_heads_k,
                num_heads_v,
                head_size,
                is_neox_style,
                eps,
                k_cache,
                v_cache,
                decode3_slot_mapping,
                decode3_cu_q_len,
                kv_cache_dtype,
                k_scale_d3,
                v_scale_d3,
                max_tokens_per_batch=dtpb,
            )
        )
        print(f"decode3: torch avg: {avg_torch_d3:.2f} us, cu avg: {avg_cu_d3:.2f} us")
        checkAllclose(q_d3_ref, q_d3, msg="decode3 q", rtol=1e-2, atol=0.05)
        checkAllclose(k_d3_ref, k_d3, msg="decode3 k", rtol=1e-2, atol=0.05)
        checkAllclose(v_d3_ref, v_d3, msg="decode3 v", rtol=1e-2, atol=0.05)
        all_slots_d3 = torch.cat([all_slots_d1, decode3_slot_mapping])
        d3_pages = torch.unique(all_slots_d3 // page_size)
        d3_k_err = checkAllclose(
            k_cache_ref.float()[d3_pages],
            k_cache.float()[d3_pages],
            msg="decode3 k_cache",
            rtol=5e-2,
            atol=0.05,
        )
        d3_v_err = checkAllclose(
            v_cache_ref.float()[d3_pages],
            v_cache.float()[d3_pages],
            msg="decode3 v_cache",
            rtol=5e-2,
            atol=0.05,
        )
        checkAllclose(
            k_scale_d3_ref[d3_pages],
            k_scale_d3[d3_pages],
            msg="decode3 k_scale",
            rtol=1e-2,
            atol=0.05,
        )
        checkAllclose(
            v_scale_d3_ref[d3_pages],
            v_scale_d3[d3_pages],
            msg="decode3 v_scale",
            rtol=1e-2,
            atol=0.05,
        )
        ret["decode3_fused_qk_us"] = avg_cu_d3
        ret["decode3_k_cache_err"] = d3_k_err
        ret["decode3_v_cache_err"] = d3_v_err

    return ret


def test_mixed_prefill_decode_block_quant(
    dtype,
    num_heads_q,
    num_heads_k,
    num_heads_v,
    head_size,
    num_blocks,
    page_size,
    is_neox_style,
    eps,
    kv_cache_dtype,
    num_decode_batches=120,
    num_prefill_batches=8,
    prefill_seq_len=100,
):
    """Test mixed prefill/decode: some batches have 1 token (decode),
    others have many tokens (prefill). Verifies that the general path
    (non-decode) correctly handles non-uniform batch distributions
    where avg tokens_per_batch < page_size but max > page_size."""
    torch.manual_seed(0)
    batch_size = num_decode_batches + num_prefill_batches
    seq_lens = [1] * num_decode_batches + [prefill_seq_len] * num_prefill_batches
    num_tokens = sum(seq_lens)
    max_tpb = max(seq_lens)
    avg_tpb = (num_tokens + batch_size - 1) // batch_size

    print("\n=== Mixed prefill/decode test ===")
    print(f"  batch_size={batch_size}, num_tokens={num_tokens}")
    print(
        f"  seq_lens: {num_decode_batches}x1 (decode) + {num_prefill_batches}x{prefill_seq_len} (prefill)"
    )
    print(f"  avg_tpb={avg_tpb}, max_tpb={max_tpb}, page_size={page_size}")
    print(
        f"  is_decode would be: avg<ps={avg_tpb < page_size}, max<ps={max_tpb < page_size}"
    )

    if kv_cache_dtype == "fp8_e4m3":
        cache_dtype = get_dtype_fp8()
    else:
        cache_dtype = dtype

    k_cache = torch.zeros(
        [num_blocks, page_size, num_heads_k, head_size],
        dtype=dtype,
        device="cuda",
    ).to(cache_dtype)
    v_cache = torch.zeros(
        [num_blocks, page_size, num_heads_v, head_size],
        dtype=dtype,
        device="cuda",
    ).to(cache_dtype)

    x = 16 // k_cache.element_size()
    k_cache = (
        k_cache.view([num_blocks, page_size, num_heads_k, head_size // x, x])
        .permute(0, 2, 3, 1, 4)
        .contiguous()
    )
    v_cache = (
        v_cache.view([num_blocks, page_size // x, num_heads_v, head_size, x])
        .permute(0, 2, 1, 3, 4)
        .contiguous()
    )

    cu_q_len = torch.zeros(batch_size + 1, dtype=torch.int64, device="cuda")
    for i in range(batch_size):
        cu_q_len[i + 1] = cu_q_len[i] + seq_lens[i]
    assert cu_q_len[-1].item() == num_tokens

    slot_start_per_batch = []
    next_slot = 0
    for i in range(batch_size):
        slot_start_per_batch.append(next_slot)
        blocks_needed = (seq_lens[i] + page_size - 1) // page_size
        next_slot += blocks_needed * page_size
    assert (
        next_slot <= num_blocks * page_size
    ), f"Need {next_slot // page_size} pages but num_blocks={num_blocks}. Increase -b."

    slot_mapping = torch.zeros(num_tokens, dtype=torch.int64, device="cuda")
    for i in range(batch_size):
        start = cu_q_len[i].item()
        end = cu_q_len[i + 1].item()
        slot_mapping[start:end] = torch.arange(
            slot_start_per_batch[i],
            slot_start_per_batch[i] + (end - start),
            dtype=torch.int64,
            device="cuda",
        )

    qkv = torch.randn(
        (num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )
    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_positions, head_size), dtype=dtype, device="cuda")
    positions = torch.randint(
        0, max_positions, (num_tokens,), dtype=torch.int64, device="cuda"
    )

    k_scale = torch.zeros([num_blocks, num_heads_k], dtype=torch.float32, device="cuda")
    v_scale = torch.zeros([num_blocks, num_heads_v], dtype=torch.float32, device="cuda")
    k_scale_ref = k_scale.clone()
    v_scale_ref = v_scale.clone()
    k_cache_ref = k_cache.clone()
    v_cache_ref = v_cache.clone()

    (q_ref, k_ref, v_ref, k_cache_ref, v_cache_ref), avg_torch = (
        run_torch_qk_norm_rope_cache_block_quant_shuffle(
            qkv,
            qw,
            kw,
            cos_sin,
            positions,
            num_tokens,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache_ref,
            v_cache_ref,
            k_scale_ref,
            v_scale_ref,
            slot_mapping,
            kv_cache_dtype,
        )
    )
    (q, k, v, k_cache, v_cache), avg_cu = (
        run_aiter_qk_norm_rope_cache_block_quant_shuffle(
            qkv,
            qw,
            kw,
            cos_sin,
            positions,
            num_tokens,
            num_heads_q,
            num_heads_k,
            num_heads_v,
            head_size,
            is_neox_style,
            eps,
            k_cache,
            v_cache,
            slot_mapping,
            cu_q_len,
            kv_cache_dtype,
            k_scale,
            v_scale,
            max_tokens_per_batch=max_tpb,
        )
    )

    info = (
        f"mixed_prefill_decode: batch={batch_size}, decode={num_decode_batches}x1, "
        f"prefill={num_prefill_batches}x{prefill_seq_len}, "
        f"num_tokens={num_tokens}, max_tpb={max_tpb}"
    )
    msg = (
        f"[perf] === {info} === torch avg: {avg_torch:<8.2f} us, "
        f"cu avg: {avg_cu:<8.2f} us, uplift: {avg_torch / avg_cu - 1:<5.1%}"
    )
    print(msg)

    checkAllclose(q_ref, q, msg="mixed q", rtol=1e-2, atol=0.05)
    checkAllclose(k_ref, k, msg="mixed k", rtol=1e-2, atol=0.05)
    checkAllclose(v_ref, v, msg="mixed v", rtol=1e-2, atol=0.05)

    slots_edit = torch.unique(slot_mapping // page_size)
    checkAllclose(
        k_cache_ref.float()[slots_edit],
        k_cache.float()[slots_edit],
        msg="mixed k_cache",
        rtol=5e-2,
        atol=0.05,
    )
    checkAllclose(
        v_cache_ref.float()[slots_edit],
        v_cache.float()[slots_edit],
        msg="mixed v_cache",
        rtol=5e-2,
        atol=0.05,
    )
    checkAllclose(
        k_scale_ref[slots_edit],
        k_scale[slots_edit],
        msg="mixed k_scale",
        rtol=1e-2,
        atol=0.01,
    )
    checkAllclose(
        v_scale_ref[slots_edit],
        v_scale[slots_edit],
        msg="mixed v_scale",
        rtol=1e-2,
        atol=0.01,
    )
    print("  PASSED: mixed prefill/decode correctness verified")
    return {"mixed_fused_qk_us": avg_cu, "mixed_unfused_us": avg_torch}


def apply_partial_rotary_emb(
    x: Tensor, cos: Tensor, sin: Tensor, rotary_dim: int, is_neox_style: bool
) -> Tensor:
    """Apply RoPE only to the first rotary_dim elements, pass through the rest."""
    x_rot = x[..., :rotary_dim]
    x_pass = x[..., rotary_dim:]
    x_rot = apply_rotary_emb_torch(x_rot, cos, sin, is_neox_style)
    return torch.cat((x_rot, x_pass), dim=-1)


def ref_partial_rotary_pts_quant(
    qkv: Tensor,
    qw: Tensor,
    kw: Tensor,
    cos_sin: Tensor,
    positions: Tensor,
    num_tokens: int,
    num_heads_q: int,
    num_heads_k: int,
    num_heads_v: int,
    head_size: int,
    rotary_dim: int,
    is_neox_style: bool,
    eps: float,
):
    """Reference implementation: RMSNorm + partial rotary RoPE."""
    q_size = num_heads_q * head_size
    k_size = num_heads_k * head_size
    v_size = num_heads_v * head_size
    qkv_flat = qkv.view(num_tokens, q_size + k_size + v_size)
    q, k, v = qkv_flat.split([q_size, k_size, v_size], dim=-1)

    q = rms_norm_forward(q.view(num_tokens, num_heads_q, head_size), qw, eps)
    k = rms_norm_forward(k.view(num_tokens, num_heads_k, head_size), kw, eps)
    v = v.view(num_tokens, num_heads_v, head_size)

    indexed = cos_sin[positions]
    cos, sin = indexed.chunk(2, dim=-1)

    q = apply_partial_rotary_emb(q, cos, sin, rotary_dim, is_neox_style)
    k = apply_partial_rotary_emb(k, cos, sin, rotary_dim, is_neox_style)
    return q, k, v


def test_partial_rotary_pts_quant(
    dtype,
    num_tokens,
    num_heads_q,
    num_heads_k,
    num_heads_v,
    head_size,
    rotary_dim,
    is_neox_style,
    eps=1e-6,
):
    max_pos = 4096
    num_slots = num_tokens + 64

    qkv = torch.randn(
        (num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_size),
        dtype=dtype,
        device="cuda",
    )
    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_pos, rotary_dim), dtype=dtype, device="cuda")
    positions = torch.randint(
        0, max_pos, (num_tokens,), dtype=torch.int64, device="cuda"
    )

    k_cache = torch.zeros(
        (num_slots, num_heads_k, head_size), dtype=dtype, device="cuda"
    )
    v_cache = torch.zeros(
        (num_slots, num_heads_v, head_size), dtype=dtype, device="cuda"
    )
    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device="cuda")
    per_tensor_k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
    per_tensor_v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")

    q_out = torch.empty(
        (num_tokens, num_heads_q, head_size), dtype=dtype, device="cuda"
    )
    k_out = torch.empty(
        (num_tokens, num_heads_k, head_size), dtype=dtype, device="cuda"
    )
    v_out = torch.empty(
        (num_tokens, num_heads_v, head_size), dtype=dtype, device="cuda"
    )

    q_ref, k_ref, v_ref = ref_partial_rotary_pts_quant(
        qkv,
        qw,
        kw,
        cos_sin,
        positions,
        num_tokens,
        num_heads_q,
        num_heads_k,
        num_heads_v,
        head_size,
        rotary_dim,
        is_neox_style,
        eps,
    )

    aiter.fused_qk_norm_rope_cache_pts_quant_shuffle(
        qkv.clone(),
        qw,
        kw,
        cos_sin,
        positions,
        num_tokens,
        num_heads_q,
        num_heads_k,
        num_heads_v,
        head_size,
        is_neox_style,
        eps,
        q_out,
        k_cache,
        v_cache,
        slot_mapping,
        per_tensor_k_scale,
        per_tensor_v_scale,
        k_out,
        v_out,
        True,
        False,
        0,
        0,
        rotary_dim,
    )

    tag = (
        f"partial_rotary dtype={dtype}, tokens={num_tokens}, "
        f"Hq={num_heads_q}, Hk={num_heads_k}, D={head_size}, "
        f"rotary_dim={rotary_dim}, neox={is_neox_style}"
    )
    checkAllclose(
        q_ref.reshape(num_tokens, -1),
        q_out.reshape(num_tokens, -1),
        msg=f"q  {tag}",
        rtol=1e-2,
        atol=0.05,
    )
    checkAllclose(
        k_ref.reshape(num_tokens, -1),
        k_out.reshape(num_tokens, -1),
        msg=f"k  {tag}",
        rtol=1e-2,
        atol=0.05,
    )
    checkAllclose(
        v_ref.reshape(num_tokens, -1),
        v_out.reshape(num_tokens, -1),
        msg=f"v  {tag}",
        rtol=1e-2,
        atol=0.05,
    )
    print(f"[PASS] {tag}", flush=True)
    return {
        "dtype": str(dtype),
        "num_tokens": num_tokens,
        "num_heads_q": num_heads_q,
        "num_heads_k": num_heads_k,
        "head_size": head_size,
        "rotary_dim": rotary_dim,
        "is_neox_style": "1" if is_neox_style else "0",
        "status": "PASS",
    }


def test_pts_quant_shuffle_block_layout_parity(
    dtype,
    num_tokens,
    num_heads_q,
    num_heads_kv,
    head_size,
    rotary_dim,
    is_neox_style,
    block_size=16,
    cache_dtype=None,
    eps=1e-6,
    use_shuffle_layout=True,
):
    """Parity: the pts write must give identical KV cache for the original [2, num_blocks, ...] and new [num_blocks, 2, ...] (unbind(1)) paged layouts, for both use_shuffle_layout=True and False (both honor the cache's per-block stride)."""
    cache_dtype = cache_dtype or dtype  # None => auto (cache dtype == qkv dtype)
    x = (
        16 // torch.empty(0, dtype=cache_dtype).element_size()
    )  # 8 (bf16/fp16), 16 (fp8)
    assert block_size % x == 0, f"block_size {block_size} must be a multiple of x={x}"
    # Enough blocks that every token fits and the mapping spans >1 block.
    num_blocks = (num_tokens + block_size - 1) // block_size + 2
    num_slots = num_blocks * block_size
    rope_w = (
        head_size if rotary_dim == 0 else rotary_dim
    )  # rotary_dim==0 => full rotary
    max_pos = 4096

    qkv = torch.randn(
        (num_tokens, (num_heads_q + 2 * num_heads_kv) * head_size),
        dtype=dtype,
        device="cuda",
    )
    qw = torch.randn(head_size, dtype=dtype, device="cuda")
    kw = torch.randn(head_size, dtype=dtype, device="cuda")
    cos_sin = torch.randn((max_pos, rope_w), dtype=dtype, device="cuda")
    positions = torch.randint(
        0, max_pos, (num_tokens,), dtype=torch.int64, device="cuda"
    )
    # Same slots for both layouts; randperm spans blocks to exercise the stride.
    slot_mapping = torch.randperm(num_slots, device="cuda")[:num_tokens].to(torch.int64)
    per_tensor_k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
    per_tensor_v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")

    def run(blocks_first: bool):
        # blocks_first=True: new [num_blocks,2,...] (K=cache[:,0], stride 2x); False: original [2,num_blocks,...] (K=cache[0], contiguous).
        if blocks_first:
            kv = torch.zeros(
                (num_blocks, 2, block_size, num_heads_kv, head_size),
                dtype=cache_dtype,
                device="cuda",
            )
            k_cache, v_cache = kv[:, 0], kv[:, 1]
        else:
            kv = torch.zeros(
                (2, num_blocks, block_size, num_heads_kv, head_size),
                dtype=cache_dtype,
                device="cuda",
            )
            k_cache, v_cache = kv[0], kv[1]
        q_out = torch.empty(
            (num_tokens, num_heads_q, head_size), dtype=dtype, device="cuda"
        )
        k_out = torch.empty(
            (num_tokens, num_heads_kv, head_size), dtype=dtype, device="cuda"
        )
        v_out = torch.empty(
            (num_tokens, num_heads_kv, head_size), dtype=dtype, device="cuda"
        )
        aiter.fused_qk_norm_rope_cache_pts_quant_shuffle(
            qkv.clone(),
            qw,
            kw,
            cos_sin,
            positions,
            num_tokens,
            num_heads_q,
            num_heads_kv,
            num_heads_kv,
            head_size,
            is_neox_style,
            eps,
            q_out,
            k_cache,
            v_cache,
            slot_mapping,
            per_tensor_k_scale,
            per_tensor_v_scale,
            k_out,
            v_out,
            True,  # return_kv
            use_shuffle_layout,
            block_size,
            x,
            rotary_dim,
        )
        return q_out, k_cache, v_cache

    q_orig, k_orig, v_orig = run(blocks_first=False)  # [2, num_blocks, ...]
    q_new, k_new, v_new = run(blocks_first=True)  # [num_blocks, 2, ...]

    tag = (
        f"block_layout_parity qkv={dtype}, cache={cache_dtype}, tokens={num_tokens}, "
        f"Hq={num_heads_q}, Hkv={num_heads_kv}, D={head_size}, rotary_dim={rotary_dim}, "
        f"block_size={block_size}, blocks={num_blocks}, neox={is_neox_style}, "
        f"shuffle={use_shuffle_layout}"
    )
    # Only the block stride differs -> must match exactly; checkAllclose logs but doesn't raise, so assert on its ratio.
    for name, a, b in (
        ("q_out", q_orig, q_new),
        ("k_cache", k_orig, k_new),
        ("v_cache", v_orig, v_new),
    ):
        err = checkAllclose(
            a.float(), b.float(), rtol=0, atol=0, printLog=False, msg=f"{name} {tag}"
        )
        assert err == 0, f"{name} [2,B] vs [B,2] parity MISMATCH (err={err}): {tag}"
    print(f"[PASS] {tag}", flush=True)
    return {
        "qkv_dtype": str(dtype),
        "cache_dtype": str(cache_dtype),
        "num_tokens": num_tokens,
        "num_heads_q": num_heads_q,
        "num_heads_kv": num_heads_kv,
        "head_size": head_size,
        "rotary_dim": rotary_dim,
        "block_size": block_size,
        "num_blocks": num_blocks,
        "is_neox_style": "1" if is_neox_style else "0",
        "use_shuffle_layout": "1" if use_shuffle_layout else "0",
        "status": "PASS",
    }


parser = argparse.ArgumentParser(
    formatter_class=argparse.RawTextHelpFormatter,
    description="config input of test",
)
parser.add_argument(
    "-n",
    "--is_neox_styles",
    type=dtypes.str2bool,
    nargs="*",
    default=[False, True],
    help="""Whether to use the Neox-style or GPT-J-style rotary
            positional embeddings.
    e.g.: -n true   # for Neox-style
          or -n false # for GPT-J-style""",
)
parser.add_argument(
    "-t",
    "--token",
    type=int,
    nargs="*",
    default=[3, 127, 513, 778, 1024, 1257],
    help="""Number of tokens.
    e.g.: -t 513""",
)
parser.add_argument(
    "-hd",
    "--head",
    type=dtypes.str2tuple,
    nargs="*",
    default=[(32, 4), (64, 8), (4, 1)],
    help="""Number of heads.
    e.g.: -hd 32,4""",
)
parser.add_argument(
    "-hs",
    "--head_sizes",
    type=int,
    nargs="*",
    default=[64, 128, 256],
    help="""Head size.
    e.g.: -hs 64""",
)
parser.add_argument(
    "-m",
    "--max_positions",
    type=int,
    default=10000,
    help="""Max Positions.
    e.g.: -m 10000""",
)
parser.add_argument(
    "-b",
    "--num_blocks",
    type=int,
    default=1000,
    help="""Number of blocks.
    e.g.: -b 1000""",
)
parser.add_argument(
    "--batch",
    type=int,
    nargs="*",
    default=[1, 4],
    help="""Batch size. num_tokens is split across batch batches for block quant only.
    e.g.: --batch 4""",
)
parser.add_argument(
    "--decode_tokens_per_batch",
    type=int,
    default=1,
    help="""Tokens per batch in decode stage (1 or 4) for block quant only.
    e.g.: --decode_tokens_per_batch 4""",
)
parser.add_argument(
    "-p",
    "--page_size",
    type=int,
    default=16,
    help="""Page size (for per_head quant).
    e.g.: -p 16""",
)
parser.add_argument(
    "--block_page_size",
    type=int,
    default=64,
    help="""Page size for block quant (default 64, more stable than 16) for block quant only.
    e.g.: --block_page_size 64""",
)
parser.add_argument(
    "--block_sizes",
    type=int,
    nargs="*",
    default=[16, 32, 64],
    help="""Paged KV cache block sizes for the shuffle stride-aware parity sweep.
    e.g.: --block_sizes 16 32""",
)
parser.add_argument(
    "--parity_tokens",
    type=int,
    nargs="*",
    default=[3, 257, 1024],
    help="""Token counts for the shuffle stride-aware parity sweep.
    e.g.: --parity_tokens 3 257""",
)
parser.add_argument(
    "--rotary_modes",
    type=str,
    nargs="*",
    choices=["partial", "full"],
    default=["partial", "full"],
    help="""Rotary modes for the parity sweep: 'partial' (head_size-specific dim) and/or 'full' (rotary_dim=0).
    e.g.: --rotary_modes partial""",
)
parser.add_argument(
    "--qkv_dtypes",
    type=dtypes.str2Dtype,
    nargs="*",
    default=[torch.bfloat16, torch.float16],
    help="""QKV (activation) dtypes for the parity sweep; the cache dtype matches it, plus fp8.
    e.g.: --qkv_dtypes bf16""",
)
parser.add_argument(
    "--shuffle_layouts",
    type=dtypes.str2bool,
    nargs="*",
    default=[True, False],
    help="""Shuffle layouts for the parity sweep: True (x-packed shuffle) and/or
    False (contiguous); both honor the cache's per-block stride.
    e.g.: --shuffle_layouts false""",
)
parser.add_argument(
    "-d",
    "--dtype",
    type=dtypes.str2Dtype,
    default="bf16",
    help="""Data type.
    e.g.: -d bf16""",
)
parser.add_argument(
    "-kvd",
    "--kv_cache_dtypes",
    type=str,
    nargs="*",
    choices=["fp8_e4m3", "auto"],
    default=["fp8_e4m3", "auto"],
    help="""KV cache data type.
    e.g.: -kvd fp8_e4m3""",
)

parser.add_argument(
    "-q",
    "--quant_type",
    type=str,
    nargs="*",
    choices=["block", "per_head"],
    default=["block", "per_head"],
    help="""Quantization type.
    block: prefill + chunk (page_size-1 per batch, K=0.001) + decode (last slot)""",
)

if __name__ == "__main__":
    args = parser.parse_args()
    max_positions = args.max_positions
    df = []
    # rope
    block_df = []
    # partial rope: Qwen3.5-style has head_size=256 and rotary_dim=64
    # GLM 4.7 has head_size=128 and rotary_dim=64
    partial_rotary_configs = {256: 64, 128: 64, 64: 16}

    for is_neox_style in args.is_neox_styles:
        for num_token in args.token:
            for num_head, num_kv_head in args.head:
                for i, head_size in enumerate(args.head_sizes):
                    for kv_cache_dtype in args.kv_cache_dtypes:
                        if kv_cache_dtype == "fp8_e4m3":
                            cache_dtype = get_dtype_fp8()
                        else:
                            cache_dtype = args.dtype
                        for quant_type in args.quant_type:
                            if quant_type == "block":
                                for batch in args.batch:
                                    ret = test_qk_norm_rope_cache_block_quant(
                                        args.dtype,
                                        num_token,
                                        num_head,
                                        num_kv_head,
                                        num_kv_head,
                                        head_size,
                                        args.num_blocks,
                                        args.block_page_size,
                                        is_neox_style,
                                        1e-6,
                                        kv_cache_dtype,
                                        batch=batch,
                                        decode_tokens_per_batch=args.decode_tokens_per_batch,
                                    )
                                    block_df.append(ret)
                            else:
                                ret = test_qk_norm_rope_cache_quant(
                                    args.dtype,
                                    num_token,
                                    num_head,
                                    num_kv_head,
                                    num_kv_head,
                                    head_size,
                                    is_neox_style,
                                    1e-6,
                                    kv_cache_dtype,
                                    args.num_blocks,
                                    args.page_size,
                                    max_positions=args.max_positions,
                                )
                                df.append(ret)
                                partial_rotary_dim = partial_rotary_configs.get(
                                    head_size
                                )
                                if partial_rotary_dim is not None:
                                    assert partial_rotary_dim < head_size
                                    ret = test_qk_norm_rope_cache_quant(
                                        args.dtype,
                                        num_token,
                                        num_head,
                                        num_kv_head,
                                        num_kv_head,
                                        head_size,
                                        is_neox_style,
                                        1e-6,
                                        kv_cache_dtype,
                                        args.num_blocks,
                                        args.page_size,
                                        max_positions=args.max_positions,
                                        rotary_dim=partial_rotary_dim,
                                    )
                                    df.append(ret)
    df = pd.DataFrame(df)
    block_df = pd.DataFrame(block_df)
    if "per_head" in args.quant_type:
        for is_neox_style in args.is_neox_styles:
            for kv_cache_dtype in args.kv_cache_dtypes:
                test_qk_norm_rope_cache_quant_v_shuffle_layout(
                    args.dtype,
                    127,
                    32,
                    4,
                    4,
                    128,
                    is_neox_style,
                    1e-6,
                    kv_cache_dtype,
                    args.num_blocks,
                    args.page_size,
                    max_positions=args.max_positions,
                )
    df_md = df.to_markdown(index=False)
    block_df_md = block_df.to_markdown(index=False)
    aiter.logger.info("qk_norm_rope_cache_quant summary (markdown):\n%s", df_md)
    aiter.logger.info(
        "qk_norm_rope_cache_block_quant summary (markdown):\n%s", block_df_md
    )

    # Mixed prefill/decode test: 120 decode (1 tok) + 8 prefill (100 tok)
    # avg_tpb = ceil(920/128) = 8 < page_size=64, but max_tpb = 100 > page_size
    # Old code would wrongly pick decode fast path; with max_tokens_per_batch fix,
    # it correctly falls back to general path.
    if "block" in args.quant_type:
        for num_head, num_kv_head in args.head:
            for kv_cache_dtype in args.kv_cache_dtypes:
                test_mixed_prefill_decode_block_quant(
                    args.dtype,
                    num_head,
                    num_kv_head,
                    num_kv_head,
                    128,
                    args.num_blocks,
                    args.block_page_size,
                    True,
                    1e-6,
                    kv_cache_dtype,
                    num_decode_batches=120,
                    num_prefill_batches=8,
                    prefill_seq_len=100,
                )
    #
    dtype = torch.bfloat16
    batch_size = 2
    num_tokens1 = 3608
    num_heads_q = 24
    num_heads_k = 25
    df = []
    for head_size in args.head_sizes:
        for num_tokens0 in args.token:
            for is_neox_styles in args.is_neox_styles:
                ret = test_qk_norm_rope_2way(
                    dtype,
                    batch_size,
                    num_tokens0,
                    num_tokens1,
                    num_heads_q,
                    num_heads_k,
                    head_size,
                    not is_neox_styles,
                    eps=1e-6,
                )
                df.append(ret)
    df = pd.DataFrame(df)
    df_md = df.to_markdown(index=False)
    aiter.logger.info("qk_norm_rope_2way summary (markdown):\n%s", df_md)

    # 1way tests (Qwen-Image-2-style: single token stream norm + RoPE)
    df = []
    for head_size in args.head_sizes:
        for num_tokens in args.token:
            for is_neox_styles in args.is_neox_styles:
                ret = test_qk_norm_rope_1way(
                    dtype,
                    batch_size,
                    num_tokens,
                    num_heads_q,
                    num_heads_k,
                    head_size,
                    not is_neox_styles,
                    eps=1e-6,
                )
                df.append(ret)
    df = pd.DataFrame(df)
    df_md = df.to_markdown(index=False)
    aiter.logger.info("qk_norm_rope_1way summary (markdown):\n%s", df_md)

    # MiniMax TP1: q/k are normalized across the full flattened q/k vector.
    df = []
    minimax_heads_q = 8
    minimax_heads_k = 2
    for head_size in args.head_sizes:
        for num_tokens in args.token:
            rotary_dims = {head_size}
            partial_rotary_dim = partial_rotary_configs.get(head_size)
            if partial_rotary_dim is not None:
                rotary_dims.add(partial_rotary_dim)
            for rotary_dim in sorted(rotary_dims):
                for is_neox_style in args.is_neox_styles:
                    ret = test_minimax_qk_norm_rope(
                        args.dtype,
                        num_tokens,
                        minimax_heads_q,
                        minimax_heads_k,
                        head_size,
                        rotary_dim,
                        is_neox_style,
                        eps=1e-6,
                    )
                    df.append(ret)
    df = pd.DataFrame(df)
    df_md = df.to_markdown(index=False)
    aiter.logger.info("minimax_qk_norm_rope summary (markdown):\n%s", df_md)

    # partial rotary tests (Qwen3.5-style: head_size=256, rotary_dim=64)
    df = []
    partial_rotary_configs = {256: 64, 128: 32, 64: 16}

    for num_token in args.token:
        for num_head, num_kv_head in args.head:
            for head_size in args.head_sizes:
                rotary_dim = partial_rotary_configs[head_size]
                assert rotary_dim < head_size
                for is_neox_style in args.is_neox_styles:
                    ret = test_partial_rotary_pts_quant(
                        args.dtype,
                        num_token,
                        num_head,
                        num_kv_head,
                        num_kv_head,
                        head_size,
                        rotary_dim,
                        is_neox_style,
                        eps=1e-6,
                    )
                    df.append(ret)
    df = pd.DataFrame(df)
    df_md = df.to_markdown(index=False)
    aiter.logger.info("partial_rotary_pts_quant summary (markdown):\n%s", df_md)

    # Stride-aware shuffle parity: original [2, num_blocks, ...] vs new [num_blocks, 2, ...] paged layout must give identical KV cache.
    fp8 = get_dtype_fp8()
    df = []
    for is_neox_style in args.is_neox_styles:
        for qkv_dtype in args.qkv_dtypes:
            for cache_dtype in (qkv_dtype, fp8):
                for block_size in args.block_sizes:
                    for num_head, num_kv_head in args.head:
                        for head_size in args.head_sizes:
                            for rotary_mode in args.rotary_modes:
                                rotary_dim = (
                                    0
                                    if rotary_mode == "full"
                                    else partial_rotary_configs[head_size]
                                )
                                for num_tokens in args.parity_tokens:
                                    for use_shuffle_layout in args.shuffle_layouts:
                                        df.append(
                                            test_pts_quant_shuffle_block_layout_parity(
                                                qkv_dtype,
                                                num_tokens,
                                                num_head,
                                                num_kv_head,
                                                head_size,
                                                rotary_dim,
                                                is_neox_style,
                                                block_size=block_size,
                                                cache_dtype=cache_dtype,
                                                use_shuffle_layout=use_shuffle_layout,
                                            )
                                        )
    df = pd.DataFrame(df)
    aiter.logger.info(
        "pts_quant_shuffle_block_layout parity summary (markdown):\n%s",
        df.to_markdown(index=False),
    )