build/torch-rocm/fusions/fused_kv_cache.py

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

import torch
import triton
from typing import Tuple
from .._triton_kernels.fusions.fused_kv_cache import (
    _fused_qk_rope_cat_and_cache_mla_kernel as triton_fused_qk_rope_cat_and_cache_mla_kernel,
    _fused_qk_rope_reshape_and_cache_kernel,
    _fused_qk_rope_cosine_cache_llama_kernel,
)

try:
    from .._gluon_kernels.gfx1250.fusions.fused_kv_cache import (
        _fused_qk_rope_cat_and_cache_mla_kernel as gluon_fused_qk_rope_cat_and_cache_mla_kernel,
    )
except:  # noqa: E722
    gluon_fused_qk_rope_cat_and_cache_mla_kernel = None

from .._aiter_compat.torch_guard import torch_compile_guard
from ..utils.logger import AiterTritonLogger
from ..utils.types import e4m3_dtype
from ..utils._triton import arch_info

_LOGGER = AiterTritonLogger()

DEVICE_ARCH = arch_info.get_arch()


def fused_qk_rope_cat_and_cache_mla_fake_tensor(
    q_nope: torch.Tensor,
    q_pe: torch.Tensor,
    k_nope: torch.Tensor,
    k_pe: torch.Tensor,
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    pos: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    k_scale: torch.Tensor,
    is_neox: bool,
    num_decode_toks_for_zeros: int = 0,
    apply_scale: bool = True,
    q_out: torch.Tensor = None,
    decode_q_pe_out: torch.Tensor = None,
    k_pe_out: torch.Tensor = None,
    q_out_dtype: torch.dtype = None,
    shuffled_kv_cache: bool = False,
    upcast_operand: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    b, qh, d_nope = q_nope.shape
    _, _, d_pe = q_pe.shape
    bk, kh, dk_nope = k_nope.shape

    if q_out is None:
        q_out = torch.empty(
            (b, qh, d_nope + d_pe),
            dtype=q_out_dtype if q_out_dtype is not None else q_nope.dtype,
            device=q_nope.device,
        )

    if decode_q_pe_out is None:
        decode_q_pe_out = torch.empty(
            (num_decode_toks_for_zeros, qh, d_pe),
            dtype=q_nope.dtype,
            device=q_nope.device,
        )

    if k_pe_out is None:
        k_pe_out = torch.empty((bk, kh, d_pe), dtype=k_pe.dtype, device=k_pe.device)

    if num_decode_toks_for_zeros > 0:
        q_nope_zeros_out = torch.empty(
            (num_decode_toks_for_zeros, qh, dk_nope),
            dtype=q_nope.dtype,
            device=q_nope.device,
        )
    else:
        q_nope_zeros_out = torch.empty(
            (0, qh, dk_nope),
            dtype=q_nope.dtype,
            device=q_nope.device,
        )

    return q_out, decode_q_pe_out, k_pe_out, q_nope_zeros_out


@torch_compile_guard(gen_fake=fused_qk_rope_cat_and_cache_mla_fake_tensor)
def fused_qk_rope_cat_and_cache_mla(
    q_nope: torch.Tensor,
    q_pe: torch.Tensor,
    k_nope: torch.Tensor,
    k_pe: torch.Tensor,
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    pos: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    k_scale: torch.Tensor,
    is_neox: bool,
    num_decode_toks_for_zeros: int = 0,
    apply_scale: bool = True,
    q_out: torch.Tensor = None,
    decode_q_pe_out: torch.Tensor = None,
    k_pe_out: torch.Tensor = None,
    q_out_dtype: torch.dtype = None,
    shuffled_kv_cache: bool = False,
    upcast_operand: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Perform RoPE on q_pe and k_pe and concat q_nope with q_pe and k_nope with k_pe along the last dimension
    the concatenated k_nope and k_pe are copied to kv_cache inplace

    Key parameters:
    - q_nope: Matrix X with shape (B, QH, D1).
    - q_pe: Matrix W with shape (B, QH, D2).
    - k_nope: Matrix X with shape (B_slot, KH, D1).
    - k_pe: Matrix W with shape (B_slot, KH, D2).
    - kv_cache: Matrix W with shape (B_cache, KH, D1 + D2).
    - slot_mapping: Matrix W with shape (B_slot, ).

    B is the number of decode tokens, B_slot is the number of prefill + decode tokens, B_cache is the max number of tokens of kv_cache
    QH must be multiple of KH

    Returns:
    - q_out: The output matrix with shape (B, QH, D1+D2).
    - kv_cache: The output matrix with shape (B_max, KH, D1 + D2) (inplace).
    """
    _LOGGER.info(
        f"FUSED_QK_ROPE_CAT_AND_CACHE_MLA: q_nope={tuple(q_nope.shape)} q_pe={tuple(q_pe.shape)} k_nope={tuple(k_nope.shape)} k_pe={tuple(k_pe.shape)} "
        + f"pos={tuple(pos.shape)} cos={tuple(cos.shape)} sin={tuple(sin.shape)} kv_cache={tuple(kv_cache.shape)} slot_mapping={tuple(slot_mapping.shape)}"
    )

    b, qh, d_nope = q_nope.shape
    b2, qh2, d_pe = q_pe.shape
    bk, kh, dk_nope = k_nope.shape
    bk2, kh2, dk2 = k_pe.shape
    kv_cache_dtype = kv_cache.dtype
    assert kv_cache_dtype in [
        torch.bfloat16,
        e4m3_dtype,
        torch.uint8,
    ], "KV cache dtype must be BF16, FP8 or packed FP4"

    block_size = 1
    SCALE_K_WIDTH_NOPE = 4
    SCALE_K_WIDTH_ROPE = 4
    if kv_cache_dtype == torch.uint8:
        assert shuffled_kv_cache, "shuffle_kv_cache must be True for FP4 KV cache"
        b_cache, h_cache, block_size, d_cache = kv_cache.shape
        SCALE_K_LORA = d_nope // 16
        SCALE_K_ROPE = d_pe // 16
        SCALE_K_WIDTH_NOPE = (
            min(16, triton.next_power_of_2(SCALE_K_LORA))
            if SCALE_K_LORA >= 4
            else SCALE_K_LORA
        )
        SCALE_K_WIDTH_ROPE = (
            min(16, triton.next_power_of_2(SCALE_K_ROPE))
            if SCALE_K_ROPE >= 4
            else SCALE_K_ROPE
        )
    else:
        if shuffled_kv_cache:
            b_cache, h_cache, block_size, d_cache = kv_cache.shape
        else:
            b_cache, h_cache, d_cache = kv_cache.shape
    (b_slot,) = slot_mapping.shape

    # allow bk >= b to support prefill + decode mixed scenario
    assert (
        b == b2 and bk == bk2 and b_slot <= bk and b <= bk
    ), "Q batch dimensions should be identical (b == b2), K batch dimensions should be identical (bk == bk2), slot_mapping should not exceed K batch size (b_slot <= bk), and Q batch should not exceed K batch (b <= bk)"
    assert qh == qh2, "Q head should be identical"
    assert kh == kh2 == h_cache, "K head should be identical"
    assert d_pe == dk2, "D dimension of q_pe and k_pe should be identical"
    assert d_nope == dk_nope, "D dimension of q_nope and k_nope should be identical"
    if kv_cache.dtype == torch.uint8:
        assert (
            (d_nope + d_pe) // 2 + (d_nope + d_pe) // 16
        ) == d_cache, "The D dimension of kv_cache should be (d_nope + d_rope) // 2 + (d_nope + d_rope) // 16 for FP4 KV cache"
    else:
        assert (
            dk_nope + d_pe == d_cache
        ), "D dimension of k_nope and k_pe should be summed up to be the D dimension of kv_cache"
    assert qh % kh == 0, "Q heads must be multiple of H heads"
    d_freq = cos.shape[-1]
    assert (d_freq == d_pe // 2) or (
        d_freq == d_pe
    ), "cos/sin last dim should be the same or half of the qk last dim"
    assert (
        num_decode_toks_for_zeros >= 0
    ), "num_decode_toks_for_zeros must be non-negative to avoid invalid tensor creation"
    if isinstance(k_scale, torch.Tensor):
        assert k_scale.numel() == 1, "k_scale should be a single-element torch.Tensor"
    reuse_freqs_front_part = d_freq == d_pe // 2

    if q_out is None:
        q_out = torch.empty(
            (b, qh, d_nope + d_pe),
            dtype=q_out_dtype if q_out_dtype is not None else q_nope.dtype,
            device=q_nope.device,
        )
    else:
        b_q_out, qh_q_out, d_q_out = q_out.shape
        assert (
            b == b_q_out and qh == qh_q_out and d_nope + d_pe == d_q_out
        ), "q_out shape mismatch"

    if decode_q_pe_out is None:
        decode_q_pe_out = torch.empty(
            (num_decode_toks_for_zeros, qh, d_pe),
            dtype=q_nope.dtype,
            device=q_nope.device,
        )
    else:
        b_decode_q_pe_out, qh_decode_q_pe_out, d_decode_q_pe_out = decode_q_pe_out.shape
        assert (
            num_decode_toks_for_zeros == b_decode_q_pe_out
            and qh == qh_decode_q_pe_out
            and d_pe == d_decode_q_pe_out
        ), "decode_q_pe_out shape mismatch"

    if k_pe_out is None:
        k_pe_out = torch.empty((bk, kh, d_pe), dtype=k_pe.dtype, device=k_pe.device)
    else:
        b_k_pe_out, hk_k_pe_out, d_k_pe_out = k_pe_out.shape
        assert (
            bk == b_k_pe_out and kh == hk_k_pe_out and d_pe == d_k_pe_out
        ), "k_pe_out shape mismatch, expected (bk, kh, d_pe)"

    q_nope_zeros_out = torch.empty(
        (num_decode_toks_for_zeros, qh, d_nope),
        dtype=q_nope.dtype,
        device=q_nope.device,
    )

    if shuffled_kv_cache:
        kv_cache_stride_b = kv_cache.stride(0)
        kv_cache_stride_h = kv_cache.stride(1)
        kv_cache_stride_d = kv_cache.stride(3)
    else:
        kv_cache_stride_b = kv_cache.stride(0)
        kv_cache_stride_h = kv_cache.stride(1)
        kv_cache_stride_d = kv_cache.stride(2)

    assert (
        kv_cache_stride_d == 1
    ), "The stride of the last dimension of KV cache must be 1"

    n_pid = b * qh + (b_slot - b) * kh
    grid = (n_pid, 1, 1)
    if DEVICE_ARCH == "gfx1250":
        _kernel = gluon_fused_qk_rope_cat_and_cache_mla_kernel
    else:
        _kernel = triton_fused_qk_rope_cat_and_cache_mla_kernel

    _kernel[grid](
        q_nope,
        q_pe,
        k_nope,
        k_pe,
        pos,
        cos,
        sin,
        q_out,
        decode_q_pe_out,
        k_pe_out,
        q_nope_zeros_out,
        kv_cache,
        slot_mapping,
        b,
        b_slot,
        num_decode_toks_for_zeros,
        *q_nope.stride(),
        *q_pe.stride(),
        *k_nope.stride(),
        *k_pe.stride(),
        pos.stride(0),
        cos.stride(0),
        cos.stride(-1),
        *q_out.stride(),
        *decode_q_pe_out.stride(),
        *k_pe_out.stride(),
        *q_nope_zeros_out.stride(),
        kv_cache_stride_b,
        kv_cache_stride_h,
        kv_cache_stride_d,
        k_scale_ptr=k_scale,
        QH_PER_KH=qh // kh,
        QH=qh,
        KH=kh,
        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
        IS_NEOX=is_neox,
        BLOCK_D_nope=d_nope,
        BLOCK_D_pe=d_pe,
        BLOCK_D_HALF_pe=d_pe // 2,
        BLOCK_SIZE=block_size,
        SHUFFLED_KV_CACHE=shuffled_kv_cache,
        SCALE_K_WIDTH_NOPE=SCALE_K_WIDTH_NOPE,
        SCALE_K_WIDTH_ROPE=SCALE_K_WIDTH_ROPE,
        OUTPUT_Q_NOPE_ZEROS_AND_Q_PE=(num_decode_toks_for_zeros > 0),
        HAVE_K_SCALE=(k_scale is not None and apply_scale),
        UPCAST_OPERAND=upcast_operand,
        num_warps=1,
    )

    return q_out, decode_q_pe_out, k_pe_out, q_nope_zeros_out


def fused_qk_rope_reshape_and_cache(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    pos: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    k_scale: torch.Tensor,
    v_scale: torch.Tensor,
    is_neox: bool,
    flash_layout: bool,
    apply_scale: bool = True,
    offs: torch.Tensor = None,
    q_out: torch.Tensor = None,
    k_out: torch.Tensor = None,
    output_zeros: bool = True,
    zeros_out: torch.Tensor = None,
    upcast_operand: bool = False,
):
    """
    Perform RoPE on q and k and along the last dimension and copy k and v into key_cache and value_cache inplace

    Key parameters:
    - q: shape (T, QH, D).
    - k: shape (T, KH, D).
    - v: shape (T, KH, D).
    - if flash_layout:
    -     key_cache: shape (T_cache, block_size, KH, D).
    -     value_cache: shape (T_cache, block_size, KH, D).
    - else:
    -     key_cache: shape (T_cache, KH, D // x, block_size, x).
    -     value_cache: shape (T_cache, KH, D, block_size).
    - slot_mapping: shape (T_slot, ).

    T is the number of decode tokens, T_cache * block_size is the max number of tokens of kv_cache
    QH must be multiple of KH

    Returns:
    - q_out: same shape as input q.
    - k_out: same shape as input k.
    - key_cache: same shape as input key_cache (inplace).
    - value_cache: same shape as input value_cache (inplace).
    - zeros_out: same shape as input q.
    """
    _LOGGER.info(
        f"FUSED_QK_ROPE_RESHAPE_AND_CACHE: q={tuple(q.shape)} k={tuple(k.shape)} "
        + f"pos={tuple(pos.shape)} cos={tuple(cos.shape)} sin={tuple(sin.shape)} key_cache={tuple(key_cache.shape)} value_cache={tuple(value_cache.shape)} slot_mapping={tuple(slot_mapping.shape)}"
    )

    t, qh, d = q.shape
    tk, kh, dk = k.shape
    tv, vh, dv = v.shape
    kv_cache_dtype = key_cache.dtype
    assert kv_cache_dtype in [
        torch.bfloat16,
        e4m3_dtype,
        torch.uint8,
    ], "KV cache dtype must be BF16, FP8 or packed FP4"

    SCALE_K_WIDTH = 4
    x_cache = 8
    value_shuffle_layout = False
    if kv_cache_dtype == torch.uint8:
        # always shuffled
        t_cache, kh_cache, block_size, d_cache = key_cache.shape
        t_cache_v, vh_cache, block_size_v, d_cache_v = value_cache.shape
        assert block_size == block_size_v
        SCALE_K = dk // 16
        SCALE_K_WIDTH = (
            min(16, triton.next_power_of_2(SCALE_K)) if SCALE_K >= 4 else SCALE_K
        )
    else:
        if flash_layout:
            t_cache, block_size, kh_cache, dk_cache = key_cache.shape
            t_cache_v, block_size_v, vh_cache, dv_cache = value_cache.shape
            value_shuffle_layout = False
        else:
            t_cache, kh_cache, dkx_cache, block_size, x_cache = key_cache.shape
            if value_cache.ndim == 5:
                # value_cache shuffle: (num_blocks, num_kv_heads, block_size // x, head_size, x)
                t_cache_v, vh_cache, slot_chunk_v, dv_cache, x_v = value_cache.shape
                value_shuffle_layout = True
                block_size_v = slot_chunk_v * x_v
                assert block_size_v == block_size and x_v == x_cache, (
                    f"value_cache shuffle (T,KH,block_size//x,D,x) must match key: "
                    f"{block_size_v=} {block_size=} {x_v=} {x_cache=}"
                )
            else:
                t_cache_v, vh_cache, dv_cache, block_size_v = value_cache.shape
                value_shuffle_layout = False
    (t_slot,) = slot_mapping.shape

    assert (
        t == tk == tv and t <= t_slot
    ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more less that of q, k and v, {t=} {tk=} {tv=} {t_slot=}"
    assert (
        block_size == block_size_v
    ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}"
    assert (
        kh == vh == kh_cache == vh_cache
    ), "KV head should be identical for k, v, key_cache, and value_cache"
    assert (
        t_cache == t_cache_v
    ), "Number of tokens should be identical for key_cache, and value_cache"
    if kv_cache_dtype == torch.uint8:
        assert d == dk == dv, "D dimension should be identical for q, k, and v"
        assert (
            d_cache == d_cache_v
        ), "D dimension should be identical for key_cache and value_cache"
        assert (
            dk // 2 + dk // 16 == d_cache
        ), "D dimension of key_cache should be (dk // 2 + dk // 16) for FP4 KV cache"
    else:
        if flash_layout:
            assert (
                d == dk == dv == dk_cache == dv_cache
            ), "D dimension should be identical for q, k, and v"
        else:
            assert (
                d == dk == dv == dkx_cache * x_cache == dv_cache
            ), "D dimension should be identical for q, k, and v"
            assert x_cache == triton.next_power_of_2(
                x_cache
            ), "x_size should be power of 2"

    assert d == triton.next_power_of_2(d), "D dimension should be power of 2"
    assert block_size == triton.next_power_of_2(
        block_size
    ), "block_size should be power of 2"
    assert qh % kh == 0, "Q heads must be multiple of H heads"
    d_freq = cos.shape[-1]
    assert (d_freq == d // 2) or (
        d_freq == d
    ), "cos/sin last dim should be the same or half of the qk last dim"
    reuse_freqs_front_part = d_freq == d // 2

    if q_out is None:
        q_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device)

    if k_out is None:
        k_out = torch.empty((tk, kh, dk), dtype=k.dtype, device=q.device)

    if zeros_out is not None:
        tz, qhz, dz = zeros_out.shape
        assert (
            t == tz and qh == qhz and d == dz
        ), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}"
        output_zeros = True
    elif output_zeros:
        zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device)
    else:
        zeros_out = None

    if kv_cache_dtype == torch.uint8:
        t_cache, kh_cache, block_size, d_cache = key_cache.shape
        key_cache_stride_t = key_cache.stride(0)
        key_cache_stride_h = key_cache.stride(1)
        key_cache_stride_d = key_cache.stride(3)
        key_cache_stride_b = key_cache.stride(2)
        key_cache_stride_x = 0
        value_cache_stride_t = value_cache.stride(0)
        value_cache_stride_h = value_cache.stride(1)
        value_cache_stride_d = value_cache.stride(3)
        value_cache_stride_b = value_cache.stride(2)
        value_cache_stride_slot_chunk = 0
        value_cache_stride_x = 0
        assert (
            key_cache_stride_d == value_cache_stride_d == 1
        ), "The stride of the last dimension of key_cache and value_cache must be 1"
    elif value_shuffle_layout:
        key_cache_stride_t = key_cache.stride(0)
        key_cache_stride_h = key_cache.stride(1)
        key_cache_stride_d = key_cache.stride(2)
        key_cache_stride_b = key_cache.stride(3)
        key_cache_stride_x = key_cache.stride(4)
        value_cache_stride_t = value_cache.stride(0)
        value_cache_stride_h = value_cache.stride(1)
        value_cache_stride_d = value_cache.stride(3)
        value_cache_stride_b = 0
        value_cache_stride_slot_chunk = value_cache.stride(2)
        value_cache_stride_x = value_cache.stride(4)
    elif not flash_layout:
        key_cache_stride_t = key_cache.stride(0)
        key_cache_stride_h = key_cache.stride(1)
        key_cache_stride_d = key_cache.stride(2)
        key_cache_stride_b = key_cache.stride(3)
        key_cache_stride_x = key_cache.stride(4)
        value_cache_stride_t = value_cache.stride(0)
        value_cache_stride_h = value_cache.stride(1)
        value_cache_stride_d = value_cache.stride(2)
        value_cache_stride_b = value_cache.stride(3)
        value_cache_stride_slot_chunk = 0
        value_cache_stride_x = 0
    else:
        key_cache_stride_t = key_cache.stride(0)
        key_cache_stride_h = key_cache.stride(2)
        key_cache_stride_d = key_cache.stride(3)
        key_cache_stride_b = key_cache.stride(1)
        key_cache_stride_x = 0
        value_cache_stride_t = value_cache.stride(0)
        value_cache_stride_h = value_cache.stride(2)
        value_cache_stride_d = value_cache.stride(3)
        value_cache_stride_b = value_cache.stride(1)
        value_cache_stride_slot_chunk = 0
        value_cache_stride_x = 0

    n_pid = t * qh + (t_slot - t) * kh
    grid = (n_pid, 1, 1)
    _fused_qk_rope_reshape_and_cache_kernel[grid](
        q,
        k,
        v,
        pos,
        cos,
        sin,
        offs,
        key_cache,
        value_cache,
        slot_mapping,
        q_out,
        k_out,
        zeros_out,
        t,
        t_slot,
        *q.stride(),
        *k.stride(),
        *v.stride(),
        cos.stride(0),
        cos.stride(-1),
        *q_out.stride(),
        *k_out.stride(),
        key_cache_stride_t,
        key_cache_stride_h,
        key_cache_stride_d,
        key_cache_stride_b,
        key_cache_stride_x,
        value_cache_stride_t,
        value_cache_stride_h,
        value_cache_stride_d,
        value_cache_stride_b,
        value_cache_stride_slot_chunk,
        value_cache_stride_x,
        zeros_out.stride(0) if zeros_out is not None else 0,
        zeros_out.stride(1) if zeros_out is not None else 0,
        zeros_out.stride(2) if zeros_out is not None else 0,
        k_scale_ptr=k_scale,
        v_scale_ptr=v_scale,
        QH_PER_KH=qh // kh,
        QH=qh,
        KH=kh,
        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
        IS_NEOX=is_neox,
        BLOCK_D_pe=d,
        BLOCK_D_HALF_pe=d // 2,
        BLOCK_SIZE=block_size,
        X_SIZE=x_cache if not flash_layout else 0,
        SCALE_K_WIDTH=SCALE_K_WIDTH,
        FLASH_LAYOUT=flash_layout,
        VALUE_SHUFFLE_LAYOUT=value_shuffle_layout,
        HAVE_POS=(offs is not None),
        HAVE_K_SCALE=(k_scale is not None and apply_scale),
        HAVE_V_SCALE=(v_scale is not None and apply_scale),
        HAVE_ZEROS=output_zeros,
        UPCAST_OPERAND=upcast_operand,
        num_warps=1,
    )

    if zeros_out is not None:
        return q_out, k_out, key_cache, value_cache, zeros_out
    return q_out, k_out, key_cache, value_cache


def fused_qk_rope_cosine_cache_llama(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    pos: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    k_scale: torch.Tensor,
    v_scale: torch.Tensor,
    is_neox: bool,
    flash_layout: bool,
    apply_scale: bool = True,
    offs: torch.Tensor = None,
    q_out: torch.Tensor = None,
):
    """
    Perform RoPE on q and k and along the last dimension and copy k and v into key_cache and value_cache inplace

    Key parameters:
    - q: shape (T, QH, D).
    - k: shape (T, KH, D).
    - v: shape (T, KH, D).
    - if flash_layout:
    -     key_cache: shape (T_cache, block_size, KH, D).
    -     value_cache: shape (T_cache, block_size, KH, D).
    - else:
    -     key_cache: shape (T_cache, KH, D // x, block_size, x).
    -     value_cache: shape (T_cache, KH, D, block_size).
    - slot_mapping: shape (T_slot, ).

    T is the number of decode tokens, T_cache * block_size is the max number of tokens of kv_cache
    QH must be multiple of KH

    Returns:
    - q_out: same shape as input q.
    - key_cache: same shape as input key_cache (inplace).
    - value_cache: same shape as input value_cache (inplace).
    """
    _LOGGER.info(
        f"FUSED_QK_ROPE_COSINE_CACHE_LLAMA: q={tuple(q.shape)} k={tuple(k.shape)} "
        + f"pos={tuple(pos.shape)} cos={tuple(cos.shape)} sin={tuple(sin.shape)} key_cache={tuple(key_cache.shape)} value_cache={tuple(value_cache.shape)} slot_mapping={tuple(slot_mapping.shape)}"
    )

    t, qh, d = q.shape
    tk, kh, dk = k.shape
    tv, vh, dv = v.shape
    if flash_layout:
        t_cache, block_size, kh_cache, dk_cache = key_cache.shape
        t_cache_v, block_size_v, vh_cache, dv_cache = value_cache.shape
    else:
        t_cache, kh_cache, dkx_cache, block_size, x_cache = key_cache.shape
        t_cache_v, vh_cache, dv_cache, block_size_v = value_cache.shape
    (t_slot,) = slot_mapping.shape

    assert (
        t == tk == tv and t <= t_slot
    ), f"Number of tokens should be identical for q, k and v. The number of tokens of slot_mapping should be no less than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}"
    assert (
        block_size == block_size_v
    ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}"
    assert (
        kh == vh == kh_cache == vh_cache
    ), "KV head should be identical for k, v, key_cache, and value_cache"
    assert (
        t_cache == t_cache_v
    ), "Number of tokens should be identical for key_cache, and value_cache"
    if flash_layout:
        assert (
            d == dk == dv == dk_cache == dv_cache
        ), "D dimension should be identical for q, k, and v"
    else:
        assert (
            d == dk == dv == dkx_cache * x_cache == dv_cache
        ), "D dimension should be identical for q, k, and v"
        assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2"

    assert d == triton.next_power_of_2(d), "D dimension should be power of 2"
    assert qh % kh == 0, "Q heads must be multiple of H heads"
    d_freq = cos.shape[-1]
    assert (d_freq == d // 2) or (
        d_freq == d
    ), "cos/sin last dim should be the same or half of the qk last dim"
    reuse_freqs_front_part = d_freq == d // 2

    if q_out is None:
        q_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device)

    n_pid = t * qh + (t_slot - t) * kh
    grid = (n_pid, 1, 1)
    _fused_qk_rope_cosine_cache_llama_kernel[grid](
        q,
        k,
        v,
        pos,
        cos,
        sin,
        offs,
        key_cache,
        value_cache,
        slot_mapping,
        q_out,
        t,
        t_slot,
        *q.stride(),
        *k.stride(),
        *v.stride(),
        cos.stride(0),
        cos.stride(-1),
        *q_out.stride(),
        key_cache.stride(0) if not flash_layout else key_cache.stride(0),
        key_cache.stride(1) if not flash_layout else key_cache.stride(2),
        key_cache.stride(2) if not flash_layout else key_cache.stride(3),
        key_cache.stride(3) if not flash_layout else key_cache.stride(1),
        key_cache.stride(4) if not flash_layout else 0,
        value_cache.stride(0) if not flash_layout else value_cache.stride(0),
        value_cache.stride(1) if not flash_layout else value_cache.stride(2),
        value_cache.stride(2) if not flash_layout else value_cache.stride(3),
        value_cache.stride(3) if not flash_layout else value_cache.stride(1),
        k_scale_ptr=k_scale,
        v_scale_ptr=v_scale,
        QH_PER_KH=qh // kh,
        QH=qh,
        KH=kh,
        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
        IS_NEOX=is_neox,
        BLOCK_D_pe=d,
        BLOCK_D_HALF_pe=d // 2,
        BLOCK_SIZE=block_size,
        X_SIZE=x_cache if not flash_layout else 0,
        FLASH_LAYOUT=flash_layout,
        HAVE_POS=(offs is not None),
        HAVE_K_SCALE=(k_scale is not None and apply_scale),
        HAVE_V_SCALE=(v_scale is not None and apply_scale),
        num_warps=1,
    )
    return q_out, key_cache, value_cache