sgl-project--sglang/python/sglang/srt/layers/attention/utils.py

import torch
import triton
import triton.language as tl

from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.kernels.ops.attention.pad import (
    pad_sequence_with_mask as pad_sequence_with_mask,
)
from sglang.kernels.ops.attention.pad import (
    pad_sequence_with_mask_kernel as pad_sequence_with_mask_kernel,
)
from sglang.kernels.ops.attention.pad import (
    seqlens_expand_kernel as seqlens_expand_kernel,
)
from sglang.kernels.ops.attention.pad import (
    seqlens_expand_triton as seqlens_expand_triton,
)
from sglang.kernels.ops.kvcache.cache_ops import (
    concat_and_cast_mha_k_kernel as concat_and_cast_mha_k_kernel,
)
from sglang.kernels.ops.kvcache.cache_ops import (
    concat_and_cast_mha_k_triton as concat_and_cast_mha_k_triton,
)
from sglang.kernels.ops.kvcache.cache_ops import (
    launch_reshape_and_cache_flash as launch_reshape_and_cache_flash,
)
from sglang.kernels.ops.kvcache.cache_ops import (
    reshape_and_cache_flash as reshape_and_cache_flash,
)
from sglang.kernels.ops.kvcache.kv_indices import (
    create_flashinfer_kv_indices_triton as create_flashinfer_kv_indices_triton,
)
from sglang.kernels.ops.kvcache.kv_indices import (
    create_flashmla_kv_indices_triton as create_flashmla_kv_indices_triton,
)
from sglang.kernels.ops.kvcache.kv_indices import (
    get_num_kv_index_blocks_flashmla as get_num_kv_index_blocks_flashmla,
)
from sglang.kernels.ops.kvcache.kv_indices import (
    get_num_page_per_block_flashmla as get_num_page_per_block_flashmla,
)
from sglang.kernels.ops.kvcache.rope_cache import (
    fused_qk_rope_reshape_and_cache as fused_qk_rope_reshape_and_cache,
)
from sglang.srt.utils import is_cuda

_is_cuda = is_cuda()

if _is_cuda:
    from sglang.jit_kernel.concat_mla import concat_mla_absorb_q


# When num_kv_heads=1, we have tensors with degenerate strides,
# For example, as below, where we have stride[-3] == stride[-2]:
# - shape: [num_pages, 1, 64, 128]
# - stride: [8192, 128, 128, 1]
# This will cause TMA desc validation fail in flashinfer (trtllm-mha backend).
#
# See: https://github.com/flashinfer-ai/flashinfer/issues/2232
def canonicalize_stride(tensor: torch.Tensor) -> torch.Tensor:
    """
    Adjust degenerate strides for a tensor, make it canonical.
    """
    sizes = tensor.size()
    strides = tensor.stride()
    ndim = tensor.dim()

    need_fix = any(
        sizes[i] == 1 and strides[i] == strides[i + 1] for i in range(ndim - 1)
    )

    if not need_fix:
        return tensor

    # canonicalize the stride
    # Example:
    # - shape: [num_pages, 1, 64, 128]
    # - stride: [8192, 128, 128, 1] (wrong!)
    # Gives new stride: [8192, 8192, 128 ,1] (correct!)
    new_strides = [0] * ndim
    new_strides[-1] = 1
    for i in range(ndim - 2, -1, -1):
        new_strides[i] = new_strides[i + 1] * sizes[i + 1]

    return tensor.as_strided(sizes, new_strides)


def mla_quantize_and_rope_for_fp8(
    q_nope: torch.Tensor,
    q_rope: torch.Tensor,
    k_nope: torch.Tensor,
    k_rope: torch.Tensor,
    pos_ids: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    is_neox: bool,
    kv_lora_rank: int,
    qk_rope_head_dim: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    import flashinfer.rope

    """Quantize and apply RoPE for FP8 attention path.

        This function handles the FP8 quantization and RoPE application for MLA attention.
        It takes separate query/key nope and rope components, applies RoPE to the rope parts,
        quantizes all components to FP8, and merges the query components into a single tensor.

        Args:
            q_nope: Query no-position-encoding component [seq_len, num_heads, kv_lora_rank]
                - expected dtype: torch.bfloat16
            q_rope: Query RoPE component [seq_len, num_heads, qk_rope_head_dim]
                - expected dtype: torch.bfloat16
            k_nope: Key no-position-encoding component [seq_len, num_heads, kv_lora_rank]
                - expected dtype: torch.bfloat16
            k_rope: Key RoPE component [seq_len, num_heads, qk_rope_head_dim]
                - expected dtype: torch.bfloat16
            pos_ids: Position indices for each token
                - expected dtype: torch.int64 or torch.int32
            cos_sin_cache: Precomputed cosine/sine cache for RoPE
                - expected dtype: matches q_/k_ input dtype (torch.bfloat16)
            is_neox: Whether to use NeoX-style RoPE (interleaved) or GPT-style (half rotation)
            kv_lora_rank: Dimension of the no-position-encoding component
            qk_rope_head_dim: Dimension of the RoPE component

        Returns:
            tuple: (merged_q_out, k_nope_out, k_rope_out) quantized to FP8
                - merged_q_out: [seq_len, num_heads, kv_lora_rank + qk_rope_head_dim], dtype=torch.float8_e4m3fn
                - k_nope_out:   [seq_len, num_heads, kv_lora_rank], dtype=torch.float8_e4m3fn
                - k_rope_out:   [seq_len, num_heads, qk_rope_head_dim], dtype=torch.float8_e4m3fn
        """
    attn_dtype = torch.float8_e4m3fn
    q_len, num_heads = q_rope.shape[0], q_rope.shape[1]

    # Allocate output tensors with FP8 dtype
    # Query output will contain merged nope + rope components
    q_out = q_rope.new_empty(
        q_len,
        num_heads,
        kv_lora_rank + qk_rope_head_dim,
        dtype=attn_dtype,
    )

    # Key outputs maintain original shapes but with FP8 dtype
    k_rope_out = k_rope.new_empty(k_rope.shape, dtype=attn_dtype)
    k_nope_out = k_nope.new_empty(k_nope.shape, dtype=attn_dtype)

    # Apply RoPE and quantize all components in a single fused kernel call
    # This kernel handles:
    # 1. RoPE application to q_rope and k_rope using cos_sin_cache and positions
    # 2. Quantization of all components to FP8 format
    # 3. Output placement into pre-allocated tensors
    flashinfer.rope.mla_rope_quantize_fp8(
        q_rope=q_rope,
        k_rope=k_rope,
        q_nope=q_nope,
        k_nope=k_nope,
        cos_sin_cache=cos_sin_cache,
        pos_ids=pos_ids,
        is_neox=is_neox,
        quantize_dtype=attn_dtype,
        # Output tensor slicing: q_out contains [nope_part, rope_part]
        q_rope_out=q_out[..., kv_lora_rank:],  # RoPE part goes to end
        k_rope_out=k_rope_out,
        q_nope_out=q_out[..., :kv_lora_rank],  # Nope part goes to beginning
        k_nope_out=k_nope_out,
        # Quantization scales (set to 1.0 for no additional scaling)
        quant_scale_q=1.0,
        quant_scale_kv=1.0,
        enable_pdl=is_arch_support_pdl(),
    )

    return q_out, k_nope_out, k_rope_out


def concat_mla_absorb_q_general(q_nope, q_rope):
    if _is_cuda and q_nope.shape[-1] == 512 and q_rope.shape[-1] == 64:
        return concat_mla_absorb_q(q_nope, q_rope)
    else:
        return torch.cat([q_nope, q_rope], dim=-1)


@triton.jit
def reshape_and_cache_flash(
    key_ptr,
    value_ptr,
    key_cache_ptr,
    value_cache_ptr,
    slot_mapping_ptr,
    swa_slot_mapping_ptr,
    k_scale_ptr,
    v_scale_ptr,
    block_stride,
    key_stride,
    value_stride,
    num_heads,
    head_size,
    block_size,
    HEAD_BLOCK: tl.constexpr,
    BLOCK_D: tl.constexpr,
    HAS_SWA: tl.constexpr,
    USE_SCALE: tl.constexpr,
):
    """
    Triton kernel for reshaping per-token K/V tensors into paged KV cache layout.

    Source layout:
        key/value: [num_tokens, num_heads, head_size]

    Target cache layout:
        cache: [num_blocks, block_size, num_heads, head_size]

    Each Triton program instance handles:
        - one token (program_id(0))
        - one block of heads (program_id(1))

    Features:
        - optional SWA slot remapping
        - optional FP8 scale dequantization before cache write

    Args:
        key_ptr: Pointer to source key tensor.
        value_ptr: Pointer to source value tensor.
        key_cache_ptr: Pointer to destination key cache tensor.
        value_cache_ptr: Pointer to destination value cache tensor.
        slot_mapping_ptr: Maps token -> cache slot.
        swa_slot_mapping_ptr: Optional second-stage slot remap for SWA mode.
        k_scale_ptr: Optional key scaling factor pointer.
        v_scale_ptr: Optional value scaling factor pointer.
        block_stride: Stride between cache blocks.
        key_stride: Stride between source key tokens.
        value_stride: Stride between source value tokens.
        num_heads: Number of attention heads.
        head_size: Hidden dimension per head.
        block_size: Number of slots per cache block.
        HEAD_BLOCK: Number of heads processed per program.
        BLOCK_D: Vectorized dimension size (power-of-2 padded).
        HAS_SWA: Enable SWA remapping.
        USE_SCALE: Enable scale division before storing.
    """

    # ----------------------------------
    # program ids
    # pid0 = token
    # pid1 = head block
    # ----------------------------------
    token_idx = tl.program_id(0)
    head_block_idx = tl.program_id(1)

    # ----------------------------------
    # slot mapping
    # ----------------------------------
    slot_idx = tl.load(slot_mapping_ptr + token_idx)

    if HAS_SWA:
        slot_idx = tl.load(swa_slot_mapping_ptr + slot_idx)

    if slot_idx < 0:
        return

    block_idx = slot_idx // block_size
    block_offset = slot_idx % block_size

    # ----------------------------------
    # head range
    # ----------------------------------
    head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK)

    head_mask = head_idx < num_heads

    dim_idx = tl.arange(0, BLOCK_D)

    # shape = [HEAD_BLOCK, BLOCK_D]
    offs = head_idx[:, None] * head_size + dim_idx[None, :]

    mask = head_mask[:, None] & (dim_idx[None, :] < head_size)

    # ----------------------------------
    # source load
    # ----------------------------------
    src_key = token_idx * key_stride + offs
    src_value = token_idx * value_stride + offs

    k = tl.load(key_ptr + src_key, mask=mask)
    v = tl.load(value_ptr + src_value, mask=mask)

    # ----------------------------------
    # optional scale
    # ----------------------------------
    if USE_SCALE:
        k_scale = tl.load(k_scale_ptr)
        v_scale = tl.load(v_scale_ptr)

        k = k / k_scale
        v = v / v_scale

    # ----------------------------------
    # target layout
    # [block_idx, block_offset, head, dim]
    # ----------------------------------
    tgt = block_idx * block_stride + block_offset * num_heads * head_size + offs

    tl.store(key_cache_ptr + tgt, k, mask=mask)
    tl.store(value_cache_ptr + tgt, v, mask=mask)


def launch_reshape_and_cache_flash(
    key,
    value,
    key_cache,
    value_cache,
    slot_mapping,
    swa_slot_mapping=None,
    k_scale=None,
    v_scale=None,
):
    """
    Launch wrapper for reshape_and_cache_flash Triton kernel.

    This wrapper prepares launch configuration and dispatches the Triton kernel
    that writes token-major K/V tensors into paged KV cache layout.

    Args:
        key: Source key tensor [num_tokens, num_heads, head_size]
        value: Source value tensor [num_tokens, num_heads, head_size]
        key_cache: Destination key cache [num_blocks, block_size, num_heads, head_size]
        value_cache: Destination value cache [num_blocks, block_size, num_heads, head_size]
        slot_mapping: Token-to-cache slot mapping
        swa_slot_mapping: Optional SWA remapping table
        k_scale: Optional key scaling factor
        v_scale: Optional value scaling factor
    """

    num_tokens = key.shape[0]
    num_heads = key.shape[1]
    head_size = key.shape[2]

    HEAD_BLOCK = 4

    BLOCK_D = triton.next_power_of_2(head_size)

    grid = (
        num_tokens,
        triton.cdiv(num_heads, HEAD_BLOCK),
    )

    reshape_and_cache_flash[grid](
        key,
        value,
        key_cache,
        value_cache,
        slot_mapping,
        swa_slot_mapping,
        k_scale if k_scale is not None else key,
        v_scale if v_scale is not None else key,
        key_cache.stride(0),
        key.stride(0),
        value.stride(0),
        num_heads,
        head_size,
        key_cache.shape[1],
        HEAD_BLOCK=HEAD_BLOCK,
        BLOCK_D=BLOCK_D,
        HAS_SWA=(swa_slot_mapping is not None),
        USE_SCALE=(k_scale is not None),
    )


@triton.jit
def reshape_and_cache_shuffle_5d(
    key_ptr,
    value_ptr,
    key_cache_ptr,
    value_cache_ptr,
    slot_mapping_ptr,
    swa_slot_mapping_ptr,
    key_stride_token,
    value_stride_token,
    num_heads,
    head_size,
    block_size,
    X: tl.constexpr,
    HEAD_BLOCK: tl.constexpr,
    BLOCK_D: tl.constexpr,
    HAS_SWA: tl.constexpr,
):
    """Scatter per-token (num_tokens, num_heads, head_size) K/V into the
    SHUFFLE 5D "vectorized" KV cache layout used by aiter CK
    `mha_batch_prefill_func` and aiter `pa_decode_gluon`.

    K cache shape: (num_blocks, num_heads, head_size // X, block_size, X)
    V cache shape: (num_blocks, num_heads, block_size // X, head_size, X)
    where X = 16 // element_size (=8 for bf16/fp16, =16 for fp8).
    block_size must be divisible by X, and head_size must be divisible by X.

    Each program handles one token and a HEAD_BLOCK-wide slice of heads.
    """
    token_idx = tl.program_id(0)
    head_block_idx = tl.program_id(1)

    slot_idx = tl.load(slot_mapping_ptr + token_idx)
    if HAS_SWA:
        slot_idx = tl.load(swa_slot_mapping_ptr + slot_idx)
    if slot_idx < 0:
        return

    block_idx = slot_idx // block_size
    slot_in_page = slot_idx % block_size
    page_outer = slot_in_page // X
    page_inner = slot_in_page % X

    head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK)
    head_mask = head_idx < num_heads
    d = tl.arange(0, BLOCK_D)
    d_mask = d < head_size
    d_outer = d // X
    d_inner = d % X

    src_off = token_idx * key_stride_token + head_idx[:, None] * head_size + d[None, :]
    src_mask = head_mask[:, None] & d_mask[None, :]
    k = tl.load(key_ptr + src_off, mask=src_mask)
    src_off_v = (
        token_idx * value_stride_token + head_idx[:, None] * head_size + d[None, :]
    )
    v = tl.load(value_ptr + src_off_v, mask=src_mask)

    layer_stride = num_heads * head_size * block_size
    head_stride = head_size * block_size

    k_tgt = (
        block_idx * layer_stride
        + head_idx[:, None] * head_stride
        + d_outer[None, :] * block_size * X
        + slot_in_page * X
        + d_inner[None, :]
    )
    tl.store(key_cache_ptr + k_tgt, k, mask=src_mask)

    v_tgt = (
        block_idx * layer_stride
        + head_idx[:, None] * head_stride
        + page_outer * head_size * X
        + d[None, :] * X
        + page_inner
    )
    tl.store(value_cache_ptr + v_tgt, v, mask=src_mask)


def launch_reshape_and_cache_shuffle_5d(
    key: torch.Tensor,
    value: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    swa_slot_mapping=None,
):
    """Launcher for reshape_and_cache_shuffle_5d.

    Args:
        key/value: (num_tokens, num_heads, head_size) source tensors
        key_cache: (num_blocks, num_heads, head_size//X, block_size, X)
        value_cache: (num_blocks, num_heads, block_size//X, head_size, X)
        slot_mapping: per-token destination slot in [0, num_blocks*block_size)
    """
    num_tokens, num_heads, head_size = key.shape
    assert value.shape == key.shape, "K/V must share token-major shape"
    assert key_cache.dim() == 5 and value_cache.dim() == 5
    num_blocks, kc_H, kc_D_over_X, block_size, X = key_cache.shape
    assert kc_H == num_heads and kc_D_over_X * X == head_size
    vb_blocks, vc_H, vc_page_over_X, vc_D, vc_X = value_cache.shape
    assert (
        vc_H == num_heads
        and vc_page_over_X * X == block_size
        and vc_D == head_size
        and vc_X == X
    )
    assert block_size % X == 0 and head_size % X == 0

    HEAD_BLOCK = min(4, triton.next_power_of_2(num_heads))
    BLOCK_D = triton.next_power_of_2(head_size)
    grid = (num_tokens, triton.cdiv(num_heads, HEAD_BLOCK))

    reshape_and_cache_shuffle_5d[grid](
        key,
        value,
        key_cache,
        value_cache,
        slot_mapping,
        swa_slot_mapping if swa_slot_mapping is not None else slot_mapping,
        key.stride(0),
        value.stride(0),
        num_heads,
        head_size,
        block_size,
        X=X,
        HEAD_BLOCK=HEAD_BLOCK,
        BLOCK_D=BLOCK_D,
        HAS_SWA=(swa_slot_mapping is not None),
    )


@triton.jit
def gather_shuffle_5d_to_linear(
    key_cache_ptr,
    value_cache_ptr,
    key_out_ptr,  # (T, num_heads, head_size), store dtype
    value_out_ptr,  # (T, num_heads, head_size), store dtype
    slot_mapping_ptr,  # (T,) absolute pool slot id per token
    key_out_stride_token,
    value_out_stride_token,
    num_heads,
    head_size,
    block_size,
    X: tl.constexpr,
    HEAD_BLOCK: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    """Inverse of :func:`reshape_and_cache_shuffle_5d`.

    Gather one token's K/V from the SHUFFLE 5D paged cache into the
    canonical (T, H, D) layout that aiter's ``mha_batch_prefill_func``
    expects in LINEAR mode. Source addressing is identical to the
    writer kernel so any bit-exact round-trip is guaranteed.
    """
    token_idx = tl.program_id(0)
    head_block_idx = tl.program_id(1)

    slot_idx = tl.load(slot_mapping_ptr + token_idx)

    block_idx = slot_idx // block_size
    slot_in_page = slot_idx % block_size
    page_outer = slot_in_page // X
    page_inner = slot_in_page % X

    head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK)
    head_mask = head_idx < num_heads
    d = tl.arange(0, BLOCK_D)
    d_mask = d < head_size
    d_outer = d // X
    d_inner = d % X

    layer_stride = num_heads * head_size * block_size
    head_stride = head_size * block_size

    src_mask = head_mask[:, None] & d_mask[None, :]
    k_src = (
        block_idx * layer_stride
        + head_idx[:, None] * head_stride
        + d_outer[None, :] * block_size * X
        + slot_in_page * X
        + d_inner[None, :]
    )
    k = tl.load(key_cache_ptr + k_src, mask=src_mask)
    v_src = (
        block_idx * layer_stride
        + head_idx[:, None] * head_stride
        + page_outer * head_size * X
        + d[None, :] * X
        + page_inner
    )
    v = tl.load(value_cache_ptr + v_src, mask=src_mask)

    dst_k = (
        token_idx * key_out_stride_token + head_idx[:, None] * head_size + d[None, :]
    )
    tl.store(key_out_ptr + dst_k, k, mask=src_mask)
    dst_v = (
        token_idx * value_out_stride_token + head_idx[:, None] * head_size + d[None, :]
    )
    tl.store(value_out_ptr + dst_v, v, mask=src_mask)


def launch_gather_shuffle_5d_to_linear(
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
):
    """Inverse of :func:`launch_reshape_and_cache_shuffle_5d`.

    Returns ``(key_out, value_out)`` each shaped
    ``(T, num_heads, head_size)`` in ``key_cache.dtype`` /
    ``value_cache.dtype``. The caller is responsible for passing the
    right per-tensor descales downstream when ``store_dtype`` is fp8.

    Args:
        key_cache:   (num_blocks, num_heads, head_size // X, block_size, X)
        value_cache: (num_blocks, num_heads, block_size // X, head_size, X)
        slot_mapping: (T,) per-token absolute slot id in
            ``[0, num_blocks * block_size)``
    """
    assert key_cache.dim() == 5 and value_cache.dim() == 5
    num_blocks, num_heads, kc_D_over_X, block_size, X = key_cache.shape
    vc_blocks, vc_H, vc_page_over_X, vc_D, vc_X = value_cache.shape
    assert vc_blocks == num_blocks and vc_H == num_heads
    assert vc_page_over_X * X == block_size and vc_X == X
    head_size = kc_D_over_X * X
    assert vc_D == head_size

    num_tokens = slot_mapping.numel()
    key_out = torch.empty(
        (num_tokens, num_heads, head_size),
        dtype=key_cache.dtype,
        device=key_cache.device,
    )
    value_out = torch.empty(
        (num_tokens, num_heads, head_size),
        dtype=value_cache.dtype,
        device=value_cache.device,
    )

    HEAD_BLOCK = min(4, triton.next_power_of_2(num_heads))
    BLOCK_D = triton.next_power_of_2(head_size)
    grid = (num_tokens, triton.cdiv(num_heads, HEAD_BLOCK))

    gather_shuffle_5d_to_linear[grid](
        key_cache,
        value_cache,
        key_out,
        value_out,
        slot_mapping,
        key_out.stride(0),
        value_out.stride(0),
        num_heads,
        head_size,
        block_size,
        X=X,
        HEAD_BLOCK=HEAD_BLOCK,
        BLOCK_D=BLOCK_D,
    )
    return key_out, value_out


@triton.jit
def _get_gptj_rotated_x(
    x,
    x_rotated_mask,
    BLOCK_D: tl.constexpr,
    BLOCK_D_HALF: tl.constexpr,
):
    # GPT-J rotary layout:
    # Pair adjacent dimensions and apply:
    # [x0, x1, x2, x3] -> [-x1, x0, -x3, x2]

    # Apply sign inversion on odd positions.
    x_rotated = tl.where(x_rotated_mask, x, -x)
    # Reshape into (D/2, 2) pairs.
    x_rotated = tl.reshape(x_rotated, (BLOCK_D_HALF, 2))
    # Swap each pair.
    x_rotated = tl.flip(x_rotated, 1)
    # Flatten back to original shape.
    x_rotated = tl.reshape(x_rotated, (BLOCK_D,))
    return x_rotated


@triton.jit
def _get_neox_rotated_x(
    x,
    x_rotated_mask,
    BLOCK_D: tl.constexpr,
    BLOCK_D_HALF: tl.constexpr,
):
    # GPT-NeoX rotary layout:
    # Split head dimension into two halves:
    # [x0, x1, x2, x3] -> [-x2, -x3, x0, x1]

    # Keep first half positive, second half negative.
    x_rotated = tl.where(x_rotated_mask, x, -x)
    # Reshape into (2, D/2).
    x_rotated = tl.reshape(x_rotated, (2, BLOCK_D_HALF))
    # Reverse each half.
    x_rotated = tl.flip(x_rotated, 1)
    # Flatten and reverse full vector.
    x_rotated = tl.reshape(x_rotated, (BLOCK_D,))
    x_rotated = tl.flip(x_rotated, 0)
    return x_rotated


@triton.jit
def _unit_rope(
    x_ptrs,
    cos,
    sin,
    d_pe_offs,
    IS_NEOX: tl.constexpr,
    BLOCK_D_pe: tl.constexpr,
    BLOCK_D_HALF_pe: tl.constexpr,
):
    # Load one full attention head vector.
    x_pe = tl.load(x_ptrs)

    # Stage 1: Build rotated vector according to rotary layout.
    if IS_NEOX:
        x_rotated_mask = d_pe_offs < BLOCK_D_HALF_pe
        x_pe_rotated = _get_neox_rotated_x(
            x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe
        )
    else:
        x_rotated_mask = d_pe_offs % 2 == 0
        x_pe_rotated = _get_gptj_rotated_x(
            x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe
        )

    # Stage 2: Apply RoPE transform:
    # x' = x*cos + rotate(x)*sin
    x_pe = x_pe * cos + x_pe_rotated * sin

    return x_pe


@triton.jit
def _load_cos_sin(
    cos_sin_ptr,
    pos,
    d_cos_offs,
    stride_t,
    stride_d,
    freq_dim,
):
    base = pos * stride_t
    cos = tl.load(cos_sin_ptr + base + d_cos_offs * stride_d)
    sin = tl.load(cos_sin_ptr + base + (d_cos_offs + freq_dim) * stride_d)
    return cos, sin


@triton.jit
def _fused_qk_rope_reshape_and_cache_kernel(
    q_ptr,
    k_ptr,
    v_ptr,
    pos_ptr,
    cos_sin_ptr,
    offs_ptr,
    key_cache_ptr,
    value_cache_ptr,
    slot_mapping_ptr,
    swa_slot_mapping_ptr,
    q_out_ptr,
    k_out_ptr,
    zeros_out_ptr,
    T,
    T_slot,
    q_stride_t,
    q_stride_h,
    q_stride_d,
    k_stride_t,
    k_stride_h,
    k_stride_d,
    v_stride_t,
    v_stride_h,
    v_stride_d,
    cos_sin_stride_t,
    cos_sin_stride_d,
    q_out_stride_t,
    q_out_stride_h,
    q_out_stride_d,
    k_out_stride_t,
    k_out_stride_h,
    k_out_stride_d,
    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_t,
    zeros_out_stride_h,
    zeros_out_stride_d,
    k_scale_ptr,
    v_scale_ptr,
    QH_PER_KH: tl.constexpr,
    QH: tl.constexpr,
    KH: tl.constexpr,
    REUSE_FREQS_FRONT_PART: tl.constexpr,
    IS_NEOX: tl.constexpr,
    BLOCK_D_pe: tl.constexpr,
    BLOCK_D_HALF_pe: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    X_SIZE: tl.constexpr,
    FLASH_LAYOUT: tl.constexpr,
    VALUE_SHUFFLE_LAYOUT: tl.constexpr = False,
    HAVE_POS: tl.constexpr = False,
    HAVE_K_SCALE: tl.constexpr = False,
    HAVE_V_SCALE: tl.constexpr = False,
    HAVE_ZEROS: tl.constexpr = False,
    HAS_SWA: tl.constexpr = False,
):
    # ============================================================
    # Stage 0: Static stride assumptions for Triton compiler
    #
    # These assumptions help Triton optimize pointer arithmetic and
    # simplify generated address calculations.
    # ============================================================

    tl.assume(q_stride_t >= 0)
    tl.assume(q_stride_h >= 0)
    tl.assume(q_stride_d >= 0)
    tl.assume(k_stride_t >= 0)
    tl.assume(k_stride_h >= 0)
    tl.assume(k_stride_d >= 0)
    tl.assume(v_stride_t >= 0)
    tl.assume(v_stride_h >= 0)
    tl.assume(v_stride_d >= 0)
    tl.assume(cos_sin_stride_t >= 0)
    tl.assume(cos_sin_stride_d >= 0)
    tl.assume(q_out_stride_t >= 0)
    tl.assume(q_out_stride_h >= 0)
    tl.assume(q_out_stride_d >= 0)
    tl.assume(k_out_stride_t >= 0)
    tl.assume(k_out_stride_h >= 0)
    tl.assume(k_out_stride_d >= 0)
    tl.assume(key_cache_stride_t >= 0)
    tl.assume(key_cache_stride_h >= 0)
    tl.assume(key_cache_stride_d >= 0)
    tl.assume(key_cache_stride_b >= 0)
    tl.assume(key_cache_stride_x >= 0)
    tl.assume(value_cache_stride_t >= 0)
    tl.assume(value_cache_stride_h >= 0)
    tl.assume(value_cache_stride_d >= 0)
    tl.assume(value_cache_stride_b >= 0)
    tl.assume(value_cache_stride_slot_chunk >= 0)
    tl.assume(value_cache_stride_x >= 0)
    tl.assume(zeros_out_stride_t >= 0)
    tl.assume(zeros_out_stride_h >= 0)
    tl.assume(zeros_out_stride_d >= 0)

    # ============================================================
    # Stage 1: Program instance mapping
    #
    # Each program handles:
    #   - one (token, q_head) for Q path
    #   - selected KV ownership for cache write path
    #
    # pid layout:
    #   [0, T*QH)            -> decode Q path
    #   [T*QH, extra KV)     -> KV-only path
    # ============================================================

    pid = tl.program_id(0)
    tl.assume(pid >= 0)

    d_pe_offs = tl.arange(0, BLOCK_D_pe).to(tl.int64)

    # ============================================================
    # Stage 2: Main decode path (Q always active)
    # ============================================================

    if pid < T * QH:
        pid_t = pid // QH
        pid_hq = pid % QH

        # --------------------------------------------------------
        # Stage 2.1: Compute rotary frequency offsets
        #
        # RoPE frequencies may be stored as:
        #   D/2 frequencies (shared front-half)
        #   D frequencies (full explicit)
        # --------------------------------------------------------

        if REUSE_FREQS_FRONT_PART:
            if IS_NEOX:
                d_cos_offs = d_pe_offs
                d_cos_offs = tl.where(
                    (d_cos_offs >= BLOCK_D_HALF_pe) & (d_cos_offs < BLOCK_D_pe),
                    d_cos_offs - BLOCK_D_HALF_pe,
                    d_cos_offs,
                ).to(d_cos_offs.dtype)
                # d_cos_mask = d_cos_offs < BLOCK_D_pe
            else:
                d_cos_offs = d_pe_offs // 2
                # d_cos_mask = d_cos_offs < BLOCK_D_HALF_pe
        else:
            d_cos_offs = d_pe_offs
            # d_cos_mask = d_cos_offs < BLOCK_D_pe

        # --------------------------------------------------------
        # Stage 2.2: Load token position and optional offset
        #
        # offs_ptr is used by chunked prefill / sliding-window decode.
        # --------------------------------------------------------
        pos = tl.load(pos_ptr + pid_t)
        if HAVE_POS:
            offset = tl.load(offs_ptr + pid_t)
            pos = pos + offset

        # --------------------------------------------------------
        # Stage 2.3: Load cosine / sine table
        # --------------------------------------------------------
        # cos_offs = pos * cos_stride_t + d_cos_offs * cos_stride_d
        # cos = tl.load(cos_ptr + cos_offs)
        # sin = tl.load(sin_ptr + cos_offs)

        freq_dim = BLOCK_D_HALF_pe if REUSE_FREQS_FRONT_PART else BLOCK_D_pe

        cos, sin = _load_cos_sin(
            cos_sin_ptr,
            pos,
            d_cos_offs,
            cos_sin_stride_t,
            cos_sin_stride_d,
            freq_dim,
        )

        # --------------------------------------------------------
        # Stage 2.4: Apply RoPE to Q
        # --------------------------------------------------------
        q_ptrs = (
            q_ptr + pid_t * q_stride_t + pid_hq * q_stride_h + d_pe_offs * q_stride_d
        )
        q_pe = _unit_rope(
            q_ptrs,
            cos,
            sin,
            d_pe_offs,
            IS_NEOX,
            BLOCK_D_pe,
            BLOCK_D_HALF_pe,
        )

        # Store rotated Q output.
        q_out_ptrs = (
            q_out_ptr
            + pid_t * q_out_stride_t
            + pid_hq * q_out_stride_h
            + d_pe_offs * q_out_stride_d
        )
        tl.store(q_out_ptrs, q_pe.to(q_out_ptr.dtype.element_ty))

        if HAVE_ZEROS:
            z = tl.zeros((BLOCK_D_pe,), dtype=zeros_out_ptr.dtype.element_ty)
            zeros_out_ptrs = (
                zeros_out_ptr
                + pid_t * zeros_out_stride_t
                + pid_hq * zeros_out_stride_h
                + d_pe_offs * zeros_out_stride_d
            )
            tl.store(zeros_out_ptrs, z)

        # ========================================================
        # Stage 3: KV ownership path
        #
        # Only one Q group leader writes KV:
        #   pid_hq % QH_PER_KH == 0
        #
        # This prevents duplicated KV cache writes.
        # ========================================================

        if pid_hq % QH_PER_KH == 0:
            # ----------------------------------------------------
            # Stage 3.1: Resolve cache slot
            # ----------------------------------------------------
            pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64)
            if HAS_SWA:
                pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot)

            # ------------------------------------------------
            # Stage 3.2: Apply RoPE to K
            # ------------------------------------------------
            if pid_slot >= 0:
                pid_t_slot = pid_slot // BLOCK_SIZE
                pid_b = pid_slot % BLOCK_SIZE
                pid_hk = pid_hq // QH_PER_KH
                if HAVE_K_SCALE:
                    k_scale = tl.load(k_scale_ptr)
                else:
                    k_scale = 1
                k_ptrs = (
                    k_ptr
                    + pid_t * k_stride_t
                    + pid_hk * k_stride_h
                    + d_pe_offs * k_stride_d
                )
                k_pe = _unit_rope(
                    k_ptrs,
                    cos,
                    sin,
                    d_pe_offs,
                    IS_NEOX,
                    BLOCK_D_pe,
                    BLOCK_D_HALF_pe,
                )

                k_out_ptrs = (
                    k_out_ptr
                    + pid_t * k_out_stride_t
                    + pid_hk * k_out_stride_h
                    + d_pe_offs * k_out_stride_d
                )
                tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty))

                # ------------------------------------------------
                # Stage 3.3: Optional fp8 scaling before cache
                # ------------------------------------------------

                k_scale_rcprl = 1 / k_scale
                k_pe = k_pe * k_scale_rcprl

                # ------------------------------------------------
                # Stage 3.4: Write K cache
                #
                # Two layouts supported:
                #   FLASH_LAYOUT
                #   paged KV layout
                # ------------------------------------------------

                if FLASH_LAYOUT:
                    k_out_ptrs = (
                        key_cache_ptr
                        + pid_t_slot * key_cache_stride_t
                        + pid_b * key_cache_stride_b
                        + pid_hk * key_cache_stride_h
                        + d_pe_offs * key_cache_stride_d
                    )
                else:
                    k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE))
                    dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64)
                    x_offs = tl.arange(0, X_SIZE).to(tl.int64)
                    k_out_ptrs = (
                        key_cache_ptr
                        + pid_t_slot * key_cache_stride_t
                        + pid_hk * key_cache_stride_h
                        + dx_offs[:, None] * key_cache_stride_d
                        + pid_b * key_cache_stride_b
                        + x_offs[None, :] * key_cache_stride_x
                    )

                tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty))

                # ------------------------------------------------
                # Stage 3.5: Write V cache
                #
                # Supports:
                #   normal layout
                #   shuffle layout
                # ------------------------------------------------

                v_ptrs = (
                    v_ptr
                    + pid_t * v_stride_t
                    + pid_hk * v_stride_h
                    + d_pe_offs * v_stride_d
                )
                if HAVE_V_SCALE:
                    v_scale = tl.load(v_scale_ptr)
                else:
                    v_scale = 1
                v_scale_rcprl = 1 / v_scale
                v = tl.load(v_ptrs) * v_scale_rcprl
                if VALUE_SHUFFLE_LAYOUT:
                    slot_chunk = pid_b // X_SIZE
                    x_off = pid_b % X_SIZE
                    v_out_ptrs = (
                        value_cache_ptr
                        + pid_t_slot * value_cache_stride_t
                        + pid_hk * value_cache_stride_h
                        + slot_chunk * value_cache_stride_slot_chunk
                        + d_pe_offs.to(tl.int64) * value_cache_stride_d
                        + x_off * value_cache_stride_x
                    )
                else:
                    v_out_ptrs = (
                        value_cache_ptr
                        + pid_t_slot * value_cache_stride_t
                        + pid_hk * value_cache_stride_h
                        + d_pe_offs.to(tl.int64) * value_cache_stride_d
                        + pid_b * value_cache_stride_b
                    )
                tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty))
    # ============================================================
    # Stage 4: Extra KV-only path
    #
    # Handles tokens that only require cache update:
    #   T_slot > T
    #
    # No Q / no RoPE on Q branch.
    # ============================================================
    else:
        pid = pid - T * QH + T * KH
        if pid < T_slot * KH:
            pid_t = pid // KH
            pid_hk = pid % KH
            pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64)
            if HAS_SWA:
                pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot)

            if pid_slot >= 0:
                pid_t_slot = pid_slot // BLOCK_SIZE
                pid_b = pid_slot % BLOCK_SIZE
                if HAVE_K_SCALE:
                    k_scale = tl.load(k_scale_ptr)
                else:
                    k_scale = 1
                k_ptrs = (
                    k_ptr
                    + pid_t * k_stride_t
                    + pid_hk * k_stride_h
                    + d_pe_offs * k_stride_d
                )

                k_pe = tl.load(k_ptrs)

                k_out_ptrs = (
                    k_out_ptr
                    + pid_t * k_out_stride_t
                    + pid_hk * k_out_stride_h
                    + d_pe_offs * k_out_stride_d
                )
                tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty))

                k_scale_rcprl = 1 / k_scale
                k_pe = k_pe * k_scale_rcprl

                if FLASH_LAYOUT:
                    k_out_ptrs = (
                        key_cache_ptr
                        + pid_t_slot * key_cache_stride_t
                        + d_pe_offs * key_cache_stride_d
                        + pid_b * key_cache_stride_b
                        + pid_hk * key_cache_stride_h
                    )
                else:
                    k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE))
                    dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64)
                    x_offs = tl.arange(0, X_SIZE).to(tl.int64)
                    k_out_ptrs = (
                        key_cache_ptr
                        + pid_t_slot * key_cache_stride_t
                        + pid_hk * key_cache_stride_h
                        + dx_offs[:, None] * key_cache_stride_d
                        + pid_b * key_cache_stride_b
                        + x_offs[None, :] * key_cache_stride_x
                    )
                tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty))

                v_ptrs = (
                    v_ptr
                    + pid_t * v_stride_t
                    + pid_hk * v_stride_h
                    + d_pe_offs * v_stride_d
                )
                if HAVE_V_SCALE:
                    v_scale = tl.load(v_scale_ptr)
                else:
                    v_scale = 1
                v_scale_rcprl = 1 / v_scale
                v = tl.load(v_ptrs) * v_scale_rcprl
                if VALUE_SHUFFLE_LAYOUT:
                    slot_chunk = pid_b // X_SIZE
                    x_off = pid_b % X_SIZE
                    v_out_ptrs = (
                        value_cache_ptr
                        + pid_t_slot * value_cache_stride_t
                        + pid_hk * value_cache_stride_h
                        + slot_chunk * value_cache_stride_slot_chunk
                        + d_pe_offs * value_cache_stride_d
                        + x_off * value_cache_stride_x
                    )
                else:
                    v_out_ptrs = (
                        value_cache_ptr
                        + pid_t_slot * value_cache_stride_t
                        + pid_hk * value_cache_stride_h
                        + d_pe_offs * value_cache_stride_d
                        + pid_b * value_cache_stride_b
                    )
                tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty))


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_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,
    swa_slot_mapping=None,
):
    """
    Perform RoPE on q and k and along the last dimension and copy k and v in to key_cache and value_cache inplace

    Key parameters:
    - q: shape (T, QH, D).
    - k: shape (T_slot, KH, D).
    - v: shape (T_slot, 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_cahce * 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.
    """

    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
        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_slot <= tk
    ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more 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 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_sin.shape[-1] // 2
    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

    n_pid = t * qh + (t_slot - t) * kh if t_slot >= t else t * qh
    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,
        swa_slot_mapping,
        q_out,
        k_out,
        zeros_out,
        t,
        t_slot,
        *q.stride(),
        *k.stride(),
        *v.stride(),
        cos_sin.stride(0),
        cos_sin.stride(-1),
        *q_out.stride(),
        *k_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(3)
            if (not flash_layout and value_shuffle_layout)
            else (value_cache.stride(2) if not flash_layout else value_cache.stride(3))
        ),
        (
            0
            if (not flash_layout and value_shuffle_layout)
            else (value_cache.stride(3) if not flash_layout else value_cache.stride(1))
        ),
        value_cache.stride(2) if (not flash_layout and value_shuffle_layout) else 0,
        value_cache.stride(4) if (not flash_layout and value_shuffle_layout) else 0,
        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,
        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,
        HAS_SWA=(swa_slot_mapping is not None),
        num_warps=1,
    )

    if zeros_out is not None:
        return q_out.view(-1, qh * d), k_out, key_cache, value_cache, zeros_out
    return q_out.view(-1, qh * d), k_out, key_cache, value_cache


def assert_buffer_fits(used: int, capacity: int, what: str, **context) -> None:
    """Safety guard: a preallocated cuda-graph buffer must hold the runtime write.

    The kv_indices / page_table scatter kernels bound writes only per-row, not
    against the destination buffer, so an undersized buffer silently overflows
    into the adjacent row. Fail fast on the host-known extent instead. All args
    are host ints, so this is always-on (no device sync, unlike async probes).
    """
    assert used <= capacity, f"{what}: used {used} > capacity {capacity}" + (
        f" ({', '.join(f'{k}={v}' for k, v in context.items())})" if context else ""
    )