DeepInfer/deepinfer/neuron_ops/vendor/deep_gemm/attention.py

"""Sparse-indexer MQA logits; mirrors ``DeepGEMM-Ascend/csrc/apis/attention.hpp``.

All four logits kernels compute the same per-token score

    logits[q, t] = sum_h relu(dot(q[q, h], kv[t])) * weights[q, h]

in fp32, with ``q`` / ``kv`` stored as FP8 e4m3 or packed FP4 e2m1 with UE8M0
scales (one scale per 32 head-dim elements). Reference behaviour is taken from
``DeepGEMM-Ascend/tests/test_attention.py`` (``ref_mqa_logits``,
``ref_paged_mqa_logits``, and the sparse-vs-dense comparison in
``run_sparse_mqa_logits_case``).

Output conventions (bf16, row stride aligned to 512 B):

* dense contiguous: ``[num_q_tokens, max_seqlen_k]``; column ``j`` of row ``q`` is
  KV token ``cu_seq_len_k_start[q] + j`` (ks-relative), ``-inf`` outside ``[ks, ke)``.
* dense paged: ``[num_q_tokens, max_context_len]``; column ``j`` is context
  position ``j`` of that query's request, resolved through the block table.
* sparse (both variants): ``[num_q_tokens, num_max_sparse_blocks * sparse_block_kv]``;
  column ``b * B + o`` is KV token ``sparse_kv_block_indices[q, b] * B + ks % B + o``
  (``ks == 0`` for the paged variant, where the index is a context position).

Scale-factor layouts:

* Q scales are packed UE8M0 ``torch.int16`` of shape ``[..., head_dim / 64]``
  (two exponent bytes per element, low byte = even 32-element block).
* paged KV pages store ``[page_kv * data_bytes | page_kv * sf_bytes]``, with
  ``sf_bytes = head_dim / 32`` raw UE8M0 exponent bytes per token.

``get_*_metadata`` returns an opaque int32 tensor carrying a snapshot of the
scheduling inputs as python attributes, matching the upstream contract that
metadata must be regenerated when those inputs change (or, for the paged sparse
kernel, when the KV contents change).
"""

import torch

from ._quant import decode_data, exponents_to_scales, is_packed_fp4, unpack_ue8m0_exponents

HEAD_DIM = 128
SF_GRAN_K = 32
#: UE8M0 scale bytes stored per KV token by the sparse kernels (``kSparseSFBytes``).
SPARSE_SF_BYTES = HEAD_DIM // SF_GRAN_K

#: Rough budget for the temporary per-column buffers, in bytes. Correctness only:
#: it merely bounds the chunk size of the column loops.
_BUDGET_BYTES = 64 << 20

_NEG_INF = float("-inf")


def ceil_div(x: int, y: int) -> int:
    return (x + y - 1) // y


def _get_logits_stride(num_logits: int) -> int:
    """512-byte row alignment of the bf16 logits (``get_logits_stride``)."""
    align_elems = 512 // torch.bfloat16.itemsize
    return ceil_div(num_logits, align_elems) * align_elems


def _empty_logits(num_rows: int, num_logits: int, device: torch.device) -> torch.Tensor:
    stride = _get_logits_stride(num_logits)
    storage = torch.empty((num_rows, stride), dtype=torch.bfloat16, device=device)
    return storage.as_strided((num_rows, num_logits), (stride, 1))


def _chunk_columns(num_q: int, num_heads: int, gather_keys: bool, multiple: int = 1) -> int:
    per_column = max(1, num_q * 4 * (num_heads + (HEAD_DIM if gather_keys else 0)))
    chunk = max(1, min(4096, _BUDGET_BYTES // per_column))
    return max(multiple, chunk // multiple * multiple)


def _take_rows(t: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
    """``t[index]``; 1-byte float rows are gathered as ``uint8`` bytes (bit-identical, and advanced indexing on
    float8 is wrong on the Neuron device, torch-neuronx 2.14)."""
    if t.dtype.is_floating_point and t.dtype.itemsize == 1:
        return t.view(torch.uint8)[index].view(t.dtype)
    return t[index]


def _decode_rows(data: torch.Tensor, sf: torch.Tensor) -> torch.Tensor:
    """Dequantize MQA rows ``[..., head_dim_storage]`` with UE8M0 scales."""
    values = decode_data(data)
    scales = exponents_to_scales(unpack_ue8m0_exponents(sf)).repeat_interleave(SF_GRAN_K, dim=-1)
    assert scales.shape[-1] >= values.shape[-1], f"{scales.shape=}, {values.shape=}"
    return values * scales[..., : values.shape[-1]]


def _decode_q(q_data: torch.Tensor, q_sf: torch.Tensor | None, num_q: int, num_heads: int) -> torch.Tensor:
    assert q_sf is not None, "MQA logits requires Q scale factors"
    data = q_data.reshape(num_q, num_heads, -1)
    sf = q_sf.reshape(num_q, num_heads, -1)
    return _decode_rows(data, sf)


def _weighted_relu(scores: torch.Tensor, weights_f32: torch.Tensor) -> torch.Tensor:
    return (scores.relu() * weights_f32[:, :, None]).sum(dim=1)


def _scores_per_query_keys(q_f32: torch.Tensor, keys_f32: torch.Tensor, weights_f32: torch.Tensor) -> torch.Tensor:
    """``[Q, H, D] x [Q, T, D] -> [Q, T]``."""
    return _weighted_relu(torch.einsum("qhd,qtd->qht", q_f32, keys_f32), weights_f32)


def _scores_shared_keys(q_f32: torch.Tensor, keys_f32: torch.Tensor, weights_f32: torch.Tensor) -> torch.Tensor:
    """``[Q, H, D] x [T, D] -> [Q, T]``."""
    return _weighted_relu(torch.einsum("qhd,td->qht", q_f32, keys_f32), weights_f32)


def _make_metadata(device: torch.device, **attrs) -> torch.Tensor:
    metadata = torch.zeros((1,), dtype=torch.int32, device=device)
    for key, value in attrs.items():
        setattr(metadata, key, value)
    return metadata


# --------------------------------------------------------------------------- #
# Dense, token-major KV
# --------------------------------------------------------------------------- #


def fp8_fp4_mqa_logits(
    q: tuple[torch.Tensor, torch.Tensor | None],
    kv: tuple[torch.Tensor, torch.Tensor],
    weights: torch.Tensor,
    cu_seq_len_k_start: torch.Tensor,
    cu_seq_len_k_end: torch.Tensor,
    max_seqlen_k: int,
) -> torch.Tensor:
    q_data, q_sf = q
    kv_data, kv_sf = kv
    assert q_data.dim() == 3 and kv_data.dim() == 2, f"{q_data.shape=}, {kv_data.shape=}"
    assert weights.dtype == torch.bfloat16, f"{weights.dtype=}"
    num_q_tokens, num_heads = q_data.shape[0], q_data.shape[1]
    num_kv_tokens = kv_data.shape[0]
    assert weights.shape == (num_q_tokens, num_heads), f"{weights.shape=}"
    assert cu_seq_len_k_start.dtype == cu_seq_len_k_end.dtype == torch.int32
    assert max_seqlen_k > 0

    device = q_data.device
    q_f32 = _decode_q(q_data, q_sf, num_q_tokens, num_heads)
    weights_f32 = weights.float()
    ks = cu_seq_len_k_start.to(torch.int64)
    ke = cu_seq_len_k_end.to(torch.int64).clamp(max=num_kv_tokens)

    # The extra sink column absorbs every out-of-range scatter target.
    scratch = torch.full((num_q_tokens, max_seqlen_k + 1), _NEG_INF, dtype=torch.float32, device=device)
    chunk = _chunk_columns(num_q_tokens, num_heads, gather_keys=False)
    for t0 in range(0, num_kv_tokens, chunk):
        t1 = min(t0 + chunk, num_kv_tokens)
        keys = _decode_rows(kv_data[t0:t1], kv_sf[t0:t1])
        values = _scores_shared_keys(q_f32, keys, weights_f32)
        positions = torch.arange(t0, t1, device=device, dtype=torch.int64)[None, :]
        columns = positions - ks[:, None]
        valid = (positions >= ks[:, None]) & (positions < ke[:, None]) & (columns < max_seqlen_k)
        columns = torch.where(valid, columns, torch.full_like(columns, max_seqlen_k))
        scratch.scatter_(1, columns, torch.where(valid, values, torch.full_like(values, _NEG_INF)))

    logits = _empty_logits(num_q_tokens, max_seqlen_k, device)
    logits.copy_(scratch[:, :max_seqlen_k].to(torch.bfloat16))
    return logits


# --------------------------------------------------------------------------- #
# Paged KV helpers
# --------------------------------------------------------------------------- #


def _paged_page_views(kv_cache: torch.Tensor, is_fp4: bool) -> tuple[torch.Tensor, torch.Tensor]:
    """Split paged MQA storage into per-token data / SF byte views.

    ``kv_cache`` is ``[num_pages, page_kv, 1, token_bytes]`` uint8, but each page is
    physically ``[page_kv * data_bytes | page_kv * sf_bytes]``
    (``tests/test_attention.py::make_random_paged_mqa_kv``); the 4D shape only
    describes the page capacity, not a per-token interleaving.
    """
    assert kv_cache.dim() == 4 and kv_cache.dtype == torch.uint8, f"{kv_cache.shape=}, {kv_cache.dtype=}"
    num_pages, page_kv = kv_cache.shape[0], kv_cache.shape[1]
    token_bytes = kv_cache.shape[3]
    data_bytes = HEAD_DIM // 2 if is_fp4 else HEAD_DIM
    assert kv_cache.shape[2] == 1 and token_bytes == data_bytes + SPARSE_SF_BYTES, f"{kv_cache.shape=}"
    assert kv_cache.stride(3) == 1 and kv_cache.stride(1) == token_bytes, f"{kv_cache.stride()=}"
    pages = kv_cache.as_strided((num_pages, page_kv * token_bytes), (kv_cache.stride(0), 1))
    data = pages[:, : page_kv * data_bytes].reshape(num_pages, page_kv, data_bytes)
    sf = pages[:, page_kv * data_bytes :].reshape(num_pages, page_kv, SPARSE_SF_BYTES)
    return data, sf


def _gather_paged_bytes(
    kv_cache: torch.Tensor,
    page_ids: torch.Tensor,
    offsets: torch.Tensor,
    is_fp4: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
    data_view, sf_view = _paged_page_views(kv_cache, is_fp4)
    return data_view[page_ids, offsets], sf_view[page_ids, offsets]


def _decode_paged_bytes(data_bytes: torch.Tensor, sf_bytes: torch.Tensor, is_fp4: bool) -> torch.Tensor:
    data_bytes = data_bytes.contiguous()
    data = data_bytes.view(torch.int8) if is_fp4 else data_bytes.view(torch.float8_e4m3fn)
    return _decode_rows(data, sf_bytes)


# --------------------------------------------------------------------------- #
# Dense, paged KV
# --------------------------------------------------------------------------- #


def get_paged_mqa_logits_metadata(context_lens: torch.Tensor, num_heads: int, indices: torch.Tensor) -> torch.Tensor:
    """Scheduling metadata for the dense paged kernel (scheduling only, no data).

    The shape matches upstream (``[num_sms + 1, 2]`` int32); the contents are unused
    here because this implementation does not split the problem across cores.
    """
    from .runtime import get_num_sms

    assert context_lens.dim() == 2 and context_lens.shape[1] == 1, f"{context_lens.shape=}"
    assert context_lens.dtype == torch.int32 and indices.dtype == torch.int32
    assert indices.dim() == 1 and indices.shape[0] == context_lens.shape[0], f"{indices.shape=}"
    return torch.zeros((get_num_sms() + 1, 2), dtype=torch.int32, device=context_lens.device)


def fp8_fp4_paged_mqa_logits(
    q: tuple[torch.Tensor, torch.Tensor | None],
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
    context_lens: torch.Tensor,
    block_table: torch.Tensor,
    schedule_meta: torch.Tensor,
    max_context_len: int,
    indices: torch.Tensor,
    use_ascend_kv_layout: bool = False,
) -> torch.Tensor:
    if use_ascend_kv_layout:
        raise NotImplementedError("use_ascend_kv_layout=True is not supported by the Neuron deep_gemm port")
    q_data, q_sf = q
    assert q_data.dim() == 4 and q_data.shape[1] == 1, f"{q_data.shape=}"
    num_q_tokens, num_heads = q_data.shape[0], q_data.shape[2]
    assert weights.shape == (num_q_tokens, num_heads) and weights.dtype == torch.bfloat16, f"{weights.shape=}"
    assert context_lens.shape == (num_q_tokens, 1) and context_lens.dtype == torch.int32
    assert block_table.dim() == 2 and block_table.shape[0] == num_q_tokens, f"{block_table.shape=}"
    is_fp4 = is_packed_fp4(q_data.dtype)
    page_kv = kv_cache.shape[1]
    device = q_data.device

    q_f32 = _decode_q(q_data, q_sf, num_q_tokens, num_heads)
    weights_f32 = weights.float()
    table = block_table.to(torch.int64)
    lens = context_lens.reshape(num_q_tokens).to(torch.int64)
    max_positions = table.shape[1] * page_kv

    logits = _empty_logits(num_q_tokens, max_context_len, device)
    chunk = _chunk_columns(num_q_tokens, num_heads, gather_keys=True)
    for c0 in range(0, max_context_len, chunk):
        c1 = min(c0 + chunk, max_context_len)
        positions = torch.arange(c0, c1, device=device, dtype=torch.int64)[None, :]
        valid = (positions < lens[:, None]) & (positions < max_positions)
        safe = positions.clamp(max=max(0, max_positions - 1)).expand(num_q_tokens, -1)
        # Block-table slots past the context may hold stale/garbage page ids: read page 0 there
        # (masked to -inf below) instead of dereferencing them.
        page_ids = torch.where(valid, table.gather(1, safe // page_kv), 0)
        keys = _decode_paged_bytes(*_gather_paged_bytes(kv_cache, page_ids, safe % page_kv, is_fp4), is_fp4)
        values = _scores_per_query_keys(q_f32, keys, weights_f32)
        logits[:, c0:c1] = torch.where(valid, values, torch.full_like(values, _NEG_INF)).to(torch.bfloat16)
    return logits


# --------------------------------------------------------------------------- #
# Sparse logits (shared helpers)
# --------------------------------------------------------------------------- #


def _sparse_token_indices(
    blocks: torch.Tensor,
    sparse_block_kv: int,
    block_offset: torch.Tensor,
    c0: int,
    c1: int,
) -> torch.Tensor:
    """KV token index of output columns ``[c0, c1)``: ``block * B + ks % B + o``."""
    device = blocks.device
    columns = torch.arange(c0, c1, device=device, dtype=torch.int64)
    block_ids = blocks.to(torch.int64).gather(
        1, (columns // sparse_block_kv)[None, :].expand(blocks.shape[0], -1)
    )
    return block_ids * sparse_block_kv + block_offset[:, None] + (columns % sparse_block_kv)[None, :]


def _num_selected_blocks(lengths: torch.Tensor, sparse_block_kv: int, capacity: int) -> torch.Tensor:
    """``min(capacity, ceil(max(0, length) / B))`` - the prefix the kernel consumes."""
    blocks = (lengths.clamp(min=0) + sparse_block_kv - 1) // sparse_block_kv
    return blocks.clamp(max=capacity)


def _check_sparse_indices(sparse_kv_block_indices: torch.Tensor, qk_dtype: torch.dtype, sparse_block_kv: int) -> None:
    assert sparse_kv_block_indices.dim() == 2 and sparse_kv_block_indices.dtype == torch.int32
    assert sparse_kv_block_indices.shape[0] > 0 and 0 < sparse_kv_block_indices.shape[1] <= 4096
    assert qk_dtype == torch.float8_e4m3fn or is_packed_fp4(qk_dtype), f"{qk_dtype=}"
    assert sparse_block_kv in (8, 16), f"{sparse_block_kv=}"


# --------------------------------------------------------------------------- #
# Sparse, token-major KV
# --------------------------------------------------------------------------- #


def get_sparse_mqa_logits_metadata(
    cu_seq_len_k_start: torch.Tensor,
    cu_seq_len_k_end: torch.Tensor,
    num_kv_tokens: int,
    sparse_kv_block_indices: torch.Tensor,
    qk_dtype: torch.dtype,
    sparse_block_kv: int,
) -> torch.Tensor:
    _check_sparse_indices(sparse_kv_block_indices, qk_dtype, sparse_block_kv)
    num_q_tokens = sparse_kv_block_indices.shape[0]
    for bounds in (cu_seq_len_k_start, cu_seq_len_k_end):
        assert bounds.dim() == 1 and bounds.numel() == num_q_tokens and bounds.dtype == torch.int32
    assert num_kv_tokens > 0
    return _make_metadata(
        sparse_kv_block_indices.device,
        _ks=cu_seq_len_k_start.clone(),
        _ke=cu_seq_len_k_end.clone(),
        _num_kv_tokens=int(num_kv_tokens),
        _blocks=sparse_kv_block_indices.clone(),
        _qk_dtype=qk_dtype,
        _sparse_block_kv=int(sparse_block_kv),
    )


def fp8_fp4_sparse_mqa_logits(
    q: tuple[torch.Tensor, torch.Tensor | None],
    kv: tuple[torch.Tensor, torch.Tensor],
    weights: torch.Tensor,
    metadata: torch.Tensor,
    num_max_sparse_blocks: int,
    sparse_block_kv: int,
) -> torch.Tensor:
    q_data, q_sf = q
    kv_data, kv_sf = kv
    assert q_data.dim() == 3, f"{q_data.shape=}"
    num_q_tokens, num_heads = q_data.shape[0], q_data.shape[1]
    assert weights.shape == (num_q_tokens, num_heads) and weights.dtype == torch.bfloat16
    blocks = getattr(metadata, "_blocks", None)
    assert blocks is not None, "metadata must come from get_sparse_mqa_logits_metadata"
    assert blocks.shape == (num_q_tokens, num_max_sparse_blocks), f"{blocks.shape=}"
    assert int(metadata._sparse_block_kv) == sparse_block_kv

    device = q_data.device
    num_kv_tokens = min(int(metadata._num_kv_tokens), kv_data.shape[0])
    ks = metadata._ks.to(torch.int64)
    ke = metadata._ke.to(torch.int64).clamp(max=num_kv_tokens)
    q_f32 = _decode_q(q_data, q_sf, num_q_tokens, num_heads)
    weights_f32 = weights.float()
    num_blocks = _num_selected_blocks(ke - ks, sparse_block_kv, num_max_sparse_blocks)
    num_logits = num_max_sparse_blocks * sparse_block_kv
    logits = _empty_logits(num_q_tokens, num_logits, device)

    chunk = _chunk_columns(num_q_tokens, num_heads, gather_keys=True, multiple=sparse_block_kv)
    for c0 in range(0, num_logits, chunk):
        c1 = min(c0 + chunk, num_logits)
        tokens = _sparse_token_indices(blocks, sparse_block_kv, ks % sparse_block_kv, c0, c1)
        columns = torch.arange(c0, c1, device=device, dtype=torch.int64)[None, :]
        valid = (columns < (num_blocks * sparse_block_kv)[:, None]) & (tokens >= ks[:, None]) & (tokens < ke[:, None])
        safe = tokens.clamp(0, max(0, num_kv_tokens - 1))
        flat = safe.reshape(-1)
        keys = _decode_rows(_take_rows(kv_data, flat), kv_sf[flat]).reshape(safe.shape[0], safe.shape[1], HEAD_DIM)
        values = _scores_per_query_keys(q_f32, keys, weights_f32)
        logits[:, c0:c1] = torch.where(valid, values, torch.full_like(values, _NEG_INF)).to(torch.bfloat16)
    return logits


# --------------------------------------------------------------------------- #
# Sparse, paged KV
# --------------------------------------------------------------------------- #


def get_paged_sparse_mqa_logits_metadata(
    context_lens: torch.Tensor,
    block_table: torch.Tensor,
    indices: torch.Tensor,
    page_kv: int,
    sparse_kv_block_indices: torch.Tensor,
    qk_dtype: torch.dtype,
    sparse_block_kv: int,
) -> torch.Tensor:
    _check_sparse_indices(sparse_kv_block_indices, qk_dtype, sparse_block_kv)
    assert page_kv in (64, 128), f"{page_kv=}"
    num_q_tokens, capacity = sparse_kv_block_indices.shape
    assert capacity % (page_kv // sparse_block_kv) == 0, f"{capacity=}, {page_kv=}, {sparse_block_kv=}"
    assert context_lens.dim() == 1 and context_lens.numel() == num_q_tokens and context_lens.dtype == torch.int32
    assert indices.dim() == 1 and indices.numel() == num_q_tokens and indices.dtype == torch.int32
    assert block_table.dim() == 2 and block_table.shape[0] == num_q_tokens and block_table.dtype == torch.int32
    return _make_metadata(
        sparse_kv_block_indices.device,
        _context_lens=context_lens.clone(),
        _block_table=block_table.clone(),
        _indices=indices.clone(),
        _page_kv=int(page_kv),
        _blocks=sparse_kv_block_indices.clone(),
        _qk_dtype=qk_dtype,
        _sparse_block_kv=int(sparse_block_kv),
        _kv_snapshot=None,
    )


def fp8_fp4_paged_sparse_mqa_logits(
    q: tuple[torch.Tensor, torch.Tensor | None],
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
    metadata: torch.Tensor,
    num_max_sparse_blocks: int,
    sparse_block_kv: int,
) -> torch.Tensor:
    """Logits over the selected KV blocks of a paged cache.

    The first call snapshots the selected KV bytes onto ``metadata``; later calls
    reuse that snapshot with the current Q and weights, which is upstream's
    FIRST/REUSE protocol.
    """
    q_data, q_sf = q
    assert q_data.dim() == 4 and q_data.shape[1] == 1, f"{q_data.shape=}"
    num_q_tokens, num_heads = q_data.shape[0], q_data.shape[2]
    assert weights.shape == (num_q_tokens, num_heads) and weights.dtype == torch.bfloat16
    blocks = getattr(metadata, "_blocks", None)
    assert blocks is not None, "metadata must come from get_paged_sparse_mqa_logits_metadata"
    assert blocks.shape == (num_q_tokens, num_max_sparse_blocks), f"{blocks.shape=}"
    assert int(metadata._sparse_block_kv) == sparse_block_kv
    page_kv = int(metadata._page_kv)
    assert kv_cache.shape[1] == page_kv, f"{kv_cache.shape=}, {page_kv=}"

    device = q_data.device
    is_fp4 = is_packed_fp4(q_data.dtype)
    q_f32 = _decode_q(q_data, q_sf, num_q_tokens, num_heads)
    weights_f32 = weights.float()
    lens = metadata._context_lens.to(torch.int64)
    table = metadata._block_table.to(torch.int64)
    num_blocks = _num_selected_blocks(lens, sparse_block_kv, num_max_sparse_blocks)
    num_logits = num_max_sparse_blocks * sparse_block_kv
    max_positions = table.shape[1] * page_kv

    # Paged sparse blocks are context positions, so the block offset is always 0.
    tokens = _sparse_token_indices(blocks, sparse_block_kv, torch.zeros_like(lens), 0, num_logits)
    columns = torch.arange(num_logits, device=device, dtype=torch.int64)[None, :]
    valid = (columns < (num_blocks * sparse_block_kv)[:, None]) & (tokens < lens[:, None]) & (tokens < max_positions)

    if metadata._kv_snapshot is None:
        safe = tokens.clamp(0, max(0, max_positions - 1))
        page_ids = torch.where(valid, table.gather(1, safe // page_kv), 0)  # see the dense paged kernel
        metadata._kv_snapshot = _gather_paged_bytes(kv_cache, page_ids, safe % page_kv, is_fp4)
    data_bytes, sf_bytes = metadata._kv_snapshot

    logits = _empty_logits(num_q_tokens, num_logits, device)
    chunk = _chunk_columns(num_q_tokens, num_heads, gather_keys=True, multiple=sparse_block_kv)
    for c0 in range(0, num_logits, chunk):
        c1 = min(c0 + chunk, num_logits)
        keys = _decode_paged_bytes(data_bytes[:, c0:c1], sf_bytes[:, c0:c1], is_fp4)
        values = _scores_per_query_keys(q_f32, keys, weights_f32)
        logits[:, c0:c1] = torch.where(
            valid[:, c0:c1], values, torch.full_like(values, _NEG_INF)
        ).to(torch.bfloat16)
    return logits