inference/kernel.py

"""Step-5 Triton kernels for DSA attention, FP8 experts and fused decoder operations.

Requires the Triton 3.x runtime specified in requirements.txt. Importing does not
initialize CUDA; executing kernels requires CUDA tensors.
"""

# FP8 quantizer/GEMM portions: Copyright (c) 2023 DeepSeek, MIT License.

from __future__ import annotations

import math
from dataclasses import dataclass

import torch
import triton
import triton.language as tl
import triton.language.extra.libdevice as _libdevice


# DSA indexer: normalize, compress regions, score and select. Summaries use
# logical tensor layout shared by the producers and consumers below.


# Step-5 compresses each region of 8 consecutive tokens into one summary vector.
REGION_BLOCK_SIZE = 8
PROXY_DIM = 256


def shared_table_stride(block_table: torch.Tensor) -> int:
    """Row stride to pass for a block table, zero when the table is shared.

    A single-row table means "every request uses this mapping". Expressing that as a zero
    row stride rather than a branch inside the kernel keeps the two cases on one code path,
    and -- more importantly -- makes the shared case safe: the alternative, indexing a
    one-row table by a nonzero request id, reads past the tensor and returns plausible
    garbage instead of failing.
    """
    if block_table.ndim != 2:
        raise ValueError(
            f"block_table must be [num_reqs, pages], got {tuple(block_table.shape)}"
        )
    return 0 if block_table.shape[0] == 1 else block_table.stride(0)


@triton.jit
def _affine_normalize(
    raw,
    mean,
    rstd,
    weight_ptr,
    bias_ptr,
    offsets,
    mask,
    weight_bias,
    HAS_BIAS: tl.constexpr,
):
    """Apply ``(x - mean) * rstd * (w + weight_bias) [+ b]`` to one dim slice.

    A device function rather than a closure because Triton cannot compile nested
    ``def`` inside a kernel, and the normalization has to be applied to three different
    dim slices of the same head: the tail, and each half of the rotated span.
    """
    scale = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
    result = (raw - mean) * rstd * (scale + weight_bias)
    if HAS_BIAS:
        result += tl.load(bias_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
    return result


@triton.jit
def _indexer_norm_rope_kernel(
    x_ptr,
    out_ptr,
    weight_ptr,
    bias_ptr,
    cos_ptr,
    sin_ptr,
    positions_ptr,
    stride_x_token,
    stride_out_token,
    stride_cos,
    eps,
    weight_bias,
    head_dim: tl.constexpr,
    rotary_dim: tl.constexpr,
    SUBTRACT_MEAN: tl.constexpr,
    HAS_BIAS: tl.constexpr,
    BLOCK_D: tl.constexpr,
    BLOCK_R: tl.constexpr,
) -> None:
    """Normalize one head of one token, then rotate its leading ``2 * rotary_dim`` dims.

    ``SUBTRACT_MEAN`` selects LayerNorm over RMSNorm. The indexer needs both: q is
    RMSNorm with a zero-centered weight, k is LayerNorm with a weight and a bias. They
    differ by one reduction and are trivial to confuse, which would shift every key
    without changing any shape.

    The rotated span is written from a separate load rather than by patching the
    normalized vector in place, because Triton cannot index a value tensor by another
    tensor -- and overlapping stores from one program would depend on store ordering the
    compiler does not promise.
    """
    token = tl.program_id(0)
    head = tl.program_id(1)
    base = token * stride_x_token + head * head_dim
    out_base = token * stride_out_token + head * head_dim

    dims = tl.arange(0, BLOCK_D)
    dim_ok = dims < head_dim
    values = tl.load(x_ptr + base + dims, mask=dim_ok, other=0.0).to(tl.float32)

    sum_sq = tl.sum(values * values)
    if SUBTRACT_MEAN:
        # Match the deployed LayerNorm reduction and fast-rsqrt path.  Computing
        # ``mean((x - mean) ** 2)`` is algebraically equivalent, but changes the
        # FP32 reduction/rounding before every sparse-indexer key.
        mean = tl.sum(values) / head_dim
        variance = sum_sq / head_dim - mean * mean
    else:
        mean = 0.0
        variance = sum_sq / head_dim
    rstd = tl.rsqrt(variance + eps)

    # Dims past the rotated span keep the plain normalized value.
    tl.store(
        out_ptr + out_base + dims,
        _affine_normalize(
            values,
            mean,
            rstd,
            weight_ptr,
            bias_ptr,
            dims,
            dim_ok,
            weight_bias,
            HAS_BIAS=HAS_BIAS,
        ).to(out_ptr.dtype.element_ty),
        mask=dim_ok & (dims >= 2 * rotary_dim),
    )

    half = tl.arange(0, BLOCK_R)
    half_ok = half < rotary_dim
    position = tl.load(positions_ptr + token).to(tl.int32)
    cos = tl.load(cos_ptr + position * stride_cos + half, mask=half_ok, other=0.0).to(tl.float32)
    sin = tl.load(sin_ptr + position * stride_cos + half, mask=half_ok, other=0.0).to(tl.float32)

    real = _affine_normalize(
        tl.load(x_ptr + base + half, mask=half_ok, other=0.0).to(tl.float32),
        mean,
        rstd,
        weight_ptr,
        bias_ptr,
        half,
        half_ok,
        weight_bias,
        HAS_BIAS=HAS_BIAS,
    ).to(out_ptr.dtype.element_ty).to(tl.float32)
    imaginary = _affine_normalize(
        tl.load(x_ptr + base + rotary_dim + half, mask=half_ok, other=0.0).to(tl.float32),
        mean,
        rstd,
        weight_ptr,
        bias_ptr,
        rotary_dim + half,
        half_ok,
        weight_bias,
        HAS_BIAS=HAS_BIAS,
    ).to(out_ptr.dtype.element_ty).to(tl.float32)
    tl.store(
        out_ptr + out_base + half,
        (real * cos - imaginary * sin).to(out_ptr.dtype.element_ty),
        mask=half_ok,
    )
    tl.store(
        out_ptr + out_base + rotary_dim + half,
        (real * sin + imaginary * cos).to(out_ptr.dtype.element_ty),
        mask=half_ok,
    )


def _norm_rope(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None,
    cos: torch.Tensor,
    sin: torch.Tensor,
    positions: torch.Tensor,
    *,
    num_heads: int,
    head_dim: int,
    rotary_dim: int,
    eps: float,
    weight_bias: float,
    subtract_mean: bool,
) -> torch.Tensor:
    out = torch.empty_like(x)
    _indexer_norm_rope_kernel[(x.shape[0], num_heads)](
        x,
        out,
        weight,
        bias if bias is not None else weight,
        cos,
        sin,
        positions,
        x.stride(0),
        out.stride(0),
        cos.stride(0),
        eps,
        weight_bias,
        head_dim=head_dim,
        rotary_dim=rotary_dim,
        SUBTRACT_MEAN=subtract_mean,
        HAS_BIAS=bias is not None,
        BLOCK_D=triton.next_power_of_2(head_dim),
        BLOCK_R=triton.next_power_of_2(rotary_dim),
    )
    return out


def indexer_norm_rope(
    index_q: torch.Tensor,
    index_k: torch.Tensor,
    index_z: torch.Tensor,
    q_norm_weight: torch.Tensor,
    k_norm_weight: torch.Tensor,
    k_norm_bias: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    positions: torch.Tensor,
    *,
    head_dim: int,
    num_q_heads: int,
    num_k_heads: int,
    rotary_dim: int,
    eps: float = 1e-6,
    q_norm_weight_bias: float = 1.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Prepare the indexer's q/k/z from their fused projection.

    Args:
        index_q: ``[tokens, num_q_heads * head_dim]``.
        index_k: ``[tokens, num_k_heads * head_dim]``.
        index_z: ``[tokens, num_k_heads * head_dim]`` -- returned untouched.
        cos, sin: ``[seq, rotary_dim]``, half the width of the rotated span.
        q_norm_weight_bias: ``1.0`` for a zero-centered RMSNorm weight (``x * (1 + w)``),
            ``0.0`` otherwise. This follows the checkpoint's norm class, so it must be
            threaded through from config rather than hardcoded.
    """
    return (
        _norm_rope(
            index_q,
            q_norm_weight,
            None,
            cos,
            sin,
            positions,
            num_heads=num_q_heads,
            head_dim=head_dim,
            rotary_dim=rotary_dim,
            eps=eps,
            weight_bias=q_norm_weight_bias,
            subtract_mean=False,
        ),
        _norm_rope(
            index_k,
            k_norm_weight,
            k_norm_bias,
            cos,
            sin,
            positions,
            num_heads=num_k_heads,
            head_dim=head_dim,
            rotary_dim=rotary_dim,
            eps=eps,
            weight_bias=0.0,
            subtract_mean=True,
        ),
        index_z.contiguous(),
    )


@triton.jit
def _csa_compress_regions_kernel(
    index_k_ptr,
    index_z_ptr,
    token_start_ptr,
    token_count_ptr,
    summary_ptr,
    mean_ptr,
    stride_token,
    stride_head,
    stride_summary_region,
    stride_summary_head,
    num_heads: tl.constexpr,
    proxy_dim: tl.constexpr,
    BLOCK_T: tl.constexpr,
    BLOCK_D: tl.constexpr,
) -> None:
    """One program per (region, head): compress a region's tokens into one vector.

    The weight is a softmax over the token axis of ``index_z`` taken **per dimension**,
    so this is an element-wise weighted mean, not a convex combination of whole key
    vectors:

        s_d = sum_t exp(z_td - m) * k_td / sum_t exp(z_td - m)

    ``m`` is one scalar maximum over both axes. It keeps ``exp`` in range;
    sharing the shift across dimensions also fixes the floating-point rounding
    used by Step-5's region-summary quantization.
    """
    region = tl.program_id(0)
    head = tl.program_id(1)

    start = tl.load(token_start_ptr + region).to(tl.int64)
    count = tl.load(token_count_ptr + region).to(tl.int64)

    tokens = tl.arange(0, BLOCK_T)
    dims = tl.arange(0, BLOCK_D)
    token_ok = tokens < count
    dim_ok = dims < proxy_dim
    mask = token_ok[:, None] & dim_ok[None, :]

    offsets = (
        (start + tokens.to(tl.int64))[:, None] * stride_token
        + head * stride_head
        + dims[None, :]
    )
    logits = tl.load(index_z_ptr + offsets, mask=mask, other=float("-inf")).to(tl.float32)
    values = tl.load(index_k_ptr + offsets, mask=mask, other=0.0).to(tl.float32)

    shift = tl.max(tl.where(mask, logits, float("-inf")))
    weights = tl.where(mask, tl.exp(logits - shift), 0.0)
    denominator = tl.sum(weights, axis=0)
    numerator = tl.sum(weights * values, axis=0)
    summary = tl.where(
        denominator > 0.0, numerator / tl.maximum(denominator, 1e-20), 0.0
    )

    out = region * stride_summary_region + head * stride_summary_head + dims
    tl.store(summary_ptr + out, summary, mask=dim_ok)
    # Step-5 rounds through BF16 before E4M3. BF16 has 7 stored fraction bits
    # and E4M3 has 3; the intermediate rounding can move a value onto an E4M3
    # midpoint. Direct FP32-to-E4M3 conversion can therefore change the bytes
    # checked by the quantization gate.
    tl.store(mean_ptr + out, summary.to(tl.bfloat16).to(tl.float8e4nv), mask=dim_ok)


def csa_region_layout(
    seq_lens: torch.Tensor, *, regions_per_seq: int, region_size: int = REGION_BLOCK_SIZE
) -> tuple[torch.Tensor, torch.Tensor]:
    """Token span of every region, for sequences packed back to back.

    ``regions_per_seq`` is the stride of the region id space, not the number of regions a
    sequence actually fills -- a sequence's regions have to land at a fixed offset so the
    summary store can be indexed without a second lookup. Regions past a sequence's end
    get a count of zero and are skipped rather than written, which keeps a stale summary
    from a previous request out of the selection.
    """
    device = seq_lens.device
    lengths = seq_lens.to(torch.int64)
    token_starts = torch.cumsum(lengths, dim=0) - lengths

    regions = torch.arange(regions_per_seq, device=device, dtype=torch.int64)
    offset = regions.view(1, -1) * region_size
    count = (lengths.view(-1, 1) - offset).clamp(0, region_size)
    start = token_starts.view(-1, 1) + offset
    return start.reshape(-1).to(torch.int32), count.reshape(-1).to(torch.int32)


def csa_compress_regions(
    index_k: torch.Tensor,
    index_z: torch.Tensor,
    token_start: torch.Tensor,
    token_count: torch.Tensor,
    *,
    region_size: int = REGION_BLOCK_SIZE,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compress packed tokens into per-region summaries.

    Args:
        index_k: ``[tokens, heads, proxy_dim]`` proxy keys.
        index_z: ``[tokens, heads, proxy_dim]`` per-dimension softmax logits.
        token_start: ``[regions]`` int32, first token of each region.
        token_count: ``[regions]`` int32, tokens present in each region; 0 = skip.

    Returns:
        ``(summary_fp32, summary_fp8)``, both ``[regions, heads, proxy_dim]``. The fp8
        tensor is what selection and sparse attention consume; the fp32 one exists
        because it is the only place the unrounded value is observable, which is what
        lets a parity failure be attributed to the reduction rather than the rounding.
    """
    if index_k.shape != index_z.shape:
        raise ValueError(
            f"index_k and index_z must match, got {tuple(index_k.shape)} and "
            f"{tuple(index_z.shape)}"
        )
    if index_k.ndim != 3:
        raise ValueError(f"index_k must be [tokens, heads, proxy_dim], got {tuple(index_k.shape)}")
    if index_k.stride(2) != 1:
        raise ValueError("index_k must be contiguous in the proxy dimension")

    regions = int(token_start.numel())
    _, heads, proxy_dim = index_k.shape
    summary = torch.zeros((regions, heads, proxy_dim), device=index_k.device, dtype=torch.float32)
    mean = torch.zeros((regions, heads, proxy_dim), device=index_k.device, dtype=torch.float8_e4m3fn)

    _csa_compress_regions_kernel[(regions, heads)](
        index_k,
        index_z,
        token_start,
        token_count,
        summary,
        mean,
        index_k.stride(0),
        index_k.stride(1),
        summary.stride(0),
        summary.stride(1),
        num_heads=heads,
        proxy_dim=proxy_dim,
        BLOCK_T=triton.next_power_of_2(region_size),
        BLOCK_D=triton.next_power_of_2(proxy_dim),
    )
    return summary, mean


@triton.jit
def _indexer_logits_kernel(
    index_q_ptr,
    index_k_ptr,
    weights_ptr,
    out_ptr,
    seq_q,
    seq_k,
    stride_q_token,
    stride_q_group,
    stride_q_head,
    stride_k_token,
    stride_w_token,
    stride_w_group,
    stride_out_row,
    heads_per_group: tl.constexpr,
    proxy_dim: tl.constexpr,
    BLOCK_Q: tl.constexpr,
    BLOCK_K: tl.constexpr,
) -> None:
    """Weighted-ReLU indexer scores for one (query tile, key tile, provider group).

        score[q, r] = sum_h relu(dot(q_qh, k_r)) * w_qh

    The ReLU is per head and *inside* the sum, so the heads cannot be folded into a
    single matmul -- hence the loop. The per-head weights are signed, so applying
    ReLU to each dot does **not** make the final weighted sum non-negative. Exact
    zeros are still common because a score is zero whenever every head dot is
    negative; for Step-5's 4 heads per group that event is about 6.25%. Those zeros
    are what the top-k tie-break then has to resolve, so they are load-bearing,
    not noise.
    """
    q_block = tl.program_id(0)
    k_block = tl.program_id(1)
    group = tl.program_id(2)

    queries = q_block * BLOCK_Q + tl.arange(0, BLOCK_Q)
    keys = k_block * BLOCK_K + tl.arange(0, BLOCK_K)
    dims = tl.arange(0, proxy_dim)
    query_ok = queries < seq_q
    key_ok = keys < seq_k

    key_tile = tl.load(
        index_k_ptr + dims[:, None] + keys[None, :] * stride_k_token,
        mask=key_ok[None, :],
        other=0.0,
    )

    accumulator = tl.zeros((BLOCK_Q, BLOCK_K), dtype=tl.float32)
    for head in tl.static_range(heads_per_group):
        query_tile = tl.load(
            index_q_ptr
            + queries[:, None] * stride_q_token
            + group * stride_q_group
            + head * stride_q_head
            + dims[None, :],
            mask=query_ok[:, None],
            other=0.0,
        )
        head_weight = tl.load(
            weights_ptr + queries * stride_w_token + group * stride_w_group + head,
            mask=query_ok,
            other=0.0,
        ).to(tl.float32)
        scores = tl.dot(query_tile, key_tile, out_dtype=tl.float32)
        accumulator += tl.maximum(scores, 0.0) * head_weight[:, None]

    rows = group * seq_q + queries
    tl.store(
        out_ptr + rows[:, None] * stride_out_row + keys[None, :],
        accumulator,
        mask=query_ok[:, None] & key_ok[None, :],
    )


def round_activations_e4m3(tensor: torch.Tensor) -> torch.Tensor:
    """Apply the indexer's e4m3 activation rounding, then widen back to bfloat16.

    The deployed kernel quantizes indexer activations to e4m3 before the score matmul,
    even though the checkpoint stores these weights in bf16. Skipping the rounding
    changes which regions get selected, so it is part of the semantics and not an
    optimization we can decline.

    Widening back to bf16 instead of feeding e4m3 into the matmul is deliberate. It is
    lossless -- e4m3 keeps 3 mantissa bits and bf16 keeps 8 -- and it keeps the kernel
    off Triton's fp8 MMA path, which is where portability across Triton versions and
    non-NVIDIA backends breaks down. The cost is peak throughput on Hopper, not
    accuracy.
    """
    return tensor.to(torch.float8_e4m3fn).to(torch.bfloat16)


def indexer_logits(
    index_q: torch.Tensor,
    weights: torch.Tensor,
    index_k: torch.Tensor,
    *,
    block_q: int = 64,
    block_k: int = 128,
) -> torch.Tensor:
    """Weighted-ReLU indexer scores over proxy keys.

    Args:
        index_q: ``[seq_q, groups, heads_per_group, proxy_dim]``.
        weights: ``[seq_q, groups, heads_per_group]``, already carrying the
            ``heads_per_group ** -0.5`` prescale.
        index_k: ``[seq_k, 1, proxy_dim]`` -- one shared key head (the indexer is MQA
            even though the main attention is GQA).

    Returns:
        ``[groups * seq_q, seq_k]`` float32, group-major then query. That row order is
        what the region selector consumes, so it is a contract rather than a convenience.
    """
    if index_k.shape[1] != 1:
        raise ValueError(f"indexer keys are MQA; expected one head, got {index_k.shape[1]}")
    seq_q, groups, heads_per_group, proxy_dim = index_q.shape
    seq_k = index_k.shape[0]
    if index_k.shape[2] != proxy_dim:
        raise ValueError(f"proxy_dim mismatch: q={proxy_dim} k={index_k.shape[2]}")

    quantized_q = round_activations_e4m3(index_q)
    quantized_k = round_activations_e4m3(index_k)
    out = torch.empty((groups * seq_q, seq_k), device=index_q.device, dtype=torch.float32)

    _indexer_logits_kernel[
        (triton.cdiv(seq_q, block_q), triton.cdiv(seq_k, block_k), groups)
    ](
        quantized_q,
        quantized_k,
        weights,
        out,
        seq_q,
        seq_k,
        quantized_q.stride(0),
        quantized_q.stride(1),
        quantized_q.stride(2),
        quantized_k.stride(0),
        weights.stride(0),
        weights.stride(1),
        out.stride(0),
        heads_per_group=heads_per_group,
        proxy_dim=proxy_dim,
        BLOCK_Q=block_q,
        BLOCK_K=block_k,
    )
    return out


# Prefill selection packs one entry as ``phys_region | (valid_tokens << 24)``.
# These are ``tl.constexpr`` because Triton kernels cannot read plain module globals.
PREFILL_REGION_VALID_SHIFT = tl.constexpr(24)
RADIX_BITS = tl.constexpr(8)
RADIX_BUCKETS = tl.constexpr(256)
RADIX_PASSES = tl.constexpr(4)


@triton.jit
def _ordered_float_key(scores):
    """Map fp32 values to monotonically ordered unsigned-32 keys held in int64.

    Negative IEEE-754 values reverse their bit ordering, while non-negative values
    become ordered after setting the sign bit. Keeping the resulting 32-bit pattern
    in a non-negative int64 lets ordinary Triton integer comparisons implement the
    required unsigned ordering. Branching on ``scores < 0`` deliberately maps
    ``-0.0`` and ``+0.0`` to the same key, matching numeric equality.

    Raw ordered-float keys place NaNs above ``+inf``. Production
    indexer scores are expected to be finite; this helper does not add a costly
    full-tensor finiteness scan to the hot path.
    """
    bits = scores.to(tl.int32, bitcast=True).to(tl.int64) & 0xFFFFFFFF
    return tl.where(
        scores < 0.0,
        (~bits) & 0xFFFFFFFF,
        bits | 0x80000000,
    )


@triton.jit
def _radix_topk_threshold(
    logits_ptr,
    row,
    stride_logits_row,
    visible,
    budget,
    seq_regions,
    BLOCK_R: tl.constexpr,
):
    """Find the exact score threshold separating the top ``budget`` entries of one row.

    Returns ``(threshold, remaining)``: the fp32 bit pattern of the ``budget``-th largest
    score, and how many entries *at* that exact key still belong in the selection. A
    threshold plus a tie count is what lets the caller take a **prefix** of a tied group
    rather than all or none of it, without ever materializing a sort.

    Scores may be signed. :func:`_ordered_float_key` converts their IEEE-754 bit
    patterns into unsigned-monotone keys before the radix passes.
    """
    threshold = budget.to(tl.int64) * 0
    remaining = budget
    bucket_ids = tl.arange(0, RADIX_BUCKETS)

    # Narrow the threshold one byte at a time, most significant first.
    for radix_pass in tl.static_range(RADIX_PASSES):
        shift = 32 - RADIX_BITS * (radix_pass + 1)
        histogram = tl.zeros((RADIX_BUCKETS,), dtype=tl.int32)
        for base in tl.range(0, seq_regions, BLOCK_R):
            regions = base + tl.arange(0, BLOCK_R)
            live = regions < visible
            scores = tl.load(
                logits_ptr + row * stride_logits_row + regions, mask=live, other=0.0
            )
            keys = _ordered_float_key(scores)
            if radix_pass == 0:
                # No high bits are fixed yet. Testing them would need a shift by 32,
                # which is undefined.
                matches = live
            else:
                high = shift + RADIX_BITS
                matches = live & ((keys >> high) == (threshold >> high))
            # Park non-matching lanes in bucket 0 and subtract them back out. Relying on
            # tl.histogram to drop an out-of-range sentinel would corrupt the bucket
            # totals if it clamps instead -- and the totals are what pick the cut, so the
            # failure would look like a wrong threshold rather than a bad histogram.
            buckets = tl.where(
                matches,
                ((keys >> shift) & (RADIX_BUCKETS - 1)).to(tl.int32),
                0,
            )
            counts = tl.histogram(buckets, RADIX_BUCKETS)
            parked = tl.sum((~matches).to(tl.int32))
            histogram += counts - tl.where(bucket_ids == 0, parked, 0)

        total = tl.sum(histogram)
        above = total - tl.cumsum(histogram, axis=0)
        # ``above`` is non-increasing, so the buckets satisfying above < remaining form a
        # suffix; the smallest of them is the one holding the cut.
        chosen = tl.min(tl.where(above < remaining, bucket_ids, RADIX_BUCKETS))
        remaining -= tl.sum(tl.where(bucket_ids > chosen, histogram, 0))
        # Cast before shifting the MSB byte. An int32 bucket >= 128 would otherwise
        # become negative first and sign-extend into the int64 threshold.
        threshold |= chosen.to(tl.int64) << shift

    return threshold, remaining


@triton.jit
def _region_topk_pack_kernel(
    logits_ptr,
    lengths_ptr,
    q_positions_ptr,
    block_table_ptr,
    request_ptr,
    out_ptr,
    seq_regions,
    stride_logits_row,
    stride_table_row,
    stride_out_row,
    topk: tl.constexpr,
    region_size: tl.constexpr,
    regions_per_page: tl.constexpr,
    BLOCK_R: tl.constexpr,
    BLOCK_TOPK: tl.constexpr,
) -> None:
    """Select the top-``topk`` visible regions of one row and pack them.

    Selection is ``torch.topk`` with ties broken by ascending region id: order the
    visible scores by ``(-score, id)`` and keep the first ``min(topk, visible)``. Ties at
    the cut take a *prefix* of the tied group, they do not drop it -- and that case is
    common rather than exotic, because post-ReLU indexer scores are exactly zero about
    6.25% of the time, so the group sitting at the cut is usually the zero group.
    """
    row = tl.program_id(0)
    visible = tl.load(lengths_ptr + row).to(tl.int32)
    visible = tl.minimum(tl.maximum(visible, 0), seq_regions)
    q_position = tl.load(q_positions_ptr + row).to(tl.int32)
    request = tl.load(request_ptr + row).to(tl.int32)

    slots = tl.arange(0, BLOCK_TOPK)
    tl.store(out_ptr + row * stride_out_row + slots, -1, mask=slots < topk)
    if visible <= 0:
        return

    budget = tl.minimum(topk, visible)
    threshold, remaining = _radix_topk_threshold(
        logits_ptr, row, stride_logits_row, visible, budget, seq_regions, BLOCK_R=BLOCK_R
    )

    # Emit in ascending region order. The selection contract specifies a set;
    # ascending order also matches the standalone metadata converter.
    cursor = 0
    ties_taken = 0
    for base in tl.range(0, seq_regions, BLOCK_R):
        regions = base + tl.arange(0, BLOCK_R)
        live = regions < visible
        scores = tl.load(
            logits_ptr + row * stride_logits_row + regions, mask=live, other=0.0
        )
        keys = _ordered_float_key(scores)
        greater = live & (keys > threshold)
        tied = live & (keys == threshold)
        tied_rank = ties_taken + tl.cumsum(tied.to(tl.int32), axis=0) - 1
        selected = greater | (tied & (tied_rank < remaining))

        positions = cursor + tl.cumsum(selected.to(tl.int32), axis=0) - 1
        safe_regions = tl.where(selected, regions, 0)
        pages = tl.load(
            block_table_ptr
            + request * stride_table_row
            + safe_regions // regions_per_page,
            mask=selected,
            other=0,
        )
        physical = pages * regions_per_page + safe_regions % regions_per_page
        valid_tokens = tl.minimum(
            tl.maximum(q_position + 1 - safe_regions * region_size, 0), region_size
        )
        packed = physical | (valid_tokens << PREFILL_REGION_VALID_SHIFT)
        # An unmapped page is a legal block-table state, not an error. It has to become a
        # ``-1`` slot rather than an address computed from a negative page, which would
        # read from before the cache.
        packed = tl.where(pages >= 0, packed, -1)
        tl.store(
            out_ptr + row * stride_out_row + positions,
            packed,
            mask=selected & (positions < topk),
        )
        cursor += tl.sum(selected.to(tl.int32))
        ties_taken += tl.sum(tied.to(tl.int32))


@triton.jit
def _region_topk_ids_kernel(
    logits_ptr,
    lengths_ptr,
    out_ptr,
    seq_regions,
    stride_logits_row,
    stride_out_row,
    topk: tl.constexpr,
    BLOCK_R: tl.constexpr,
    BLOCK_TOPK: tl.constexpr,
) -> None:
    """Select the top-``topk`` visible regions and emit **logical ids**, ascending.

    The same selection as :func:`_region_topk_pack_kernel` with the packing left off.
    Prefill and decode pack the identical selection into two different words, and the
    decode word needs the logical region id, which the prefill word does not carry -- so
    decode metadata cannot be derived from prefill metadata, only from the raw ids.
    """
    row = tl.program_id(0)
    visible = tl.load(lengths_ptr + row).to(tl.int32)
    visible = tl.minimum(tl.maximum(visible, 0), seq_regions)

    slots = tl.arange(0, BLOCK_TOPK)
    tl.store(out_ptr + row * stride_out_row + slots, -1, mask=slots < topk)
    if visible <= 0:
        return

    budget = tl.minimum(topk, visible)
    threshold, remaining = _radix_topk_threshold(
        logits_ptr, row, stride_logits_row, visible, budget, seq_regions, BLOCK_R=BLOCK_R
    )

    cursor = 0
    ties_taken = 0
    for base in tl.range(0, seq_regions, BLOCK_R):
        regions = base + tl.arange(0, BLOCK_R)
        live = regions < visible
        scores = tl.load(
            logits_ptr + row * stride_logits_row + regions, mask=live, other=0.0
        )
        keys = _ordered_float_key(scores)
        greater = live & (keys > threshold)
        tied = live & (keys == threshold)
        tied_rank = ties_taken + tl.cumsum(tied.to(tl.int32), axis=0) - 1
        selected = greater | (tied & (tied_rank < remaining))

        positions = cursor + tl.cumsum(selected.to(tl.int32), axis=0) - 1
        tl.store(
            out_ptr + row * stride_out_row + positions,
            regions,
            mask=selected & (positions < topk),
        )
        cursor += tl.sum(selected.to(tl.int32))
        ties_taken += tl.sum(tied.to(tl.int32))


def region_topk_ids(
    logits: torch.Tensor,
    lengths: torch.Tensor,
    *,
    topk: int,
    block_r: int = 1024,
) -> torch.Tensor:
    """Top-k region selection, returning logical region ids in ascending order.

    Args:
        logits: ``[rows, seq_regions]`` float32 indexer scores, which may be signed.
        lengths: ``[rows]`` int32, candidate region count per row. For DSA this is the
            **history** length, ``q_position // region_size``, which excludes the region
            containing the query itself -- that one is appended by the metadata builder
            rather than made to compete.

    Returns:
        ``[rows, topk]`` int32 logical region ids, ``-1`` padding the tail.
    """
    if logits.dtype != torch.float32:
        raise ValueError(f"selector expects float32 scores, got {logits.dtype}")

    rows, seq_regions = logits.shape
    out = torch.empty((rows, int(topk)), device=logits.device, dtype=torch.int32)
    _region_topk_ids_kernel[(rows,)](
        logits,
        lengths,
        out,
        seq_regions,
        logits.stride(0),
        out.stride(0),
        topk=int(topk),
        BLOCK_R=min(block_r, triton.next_power_of_2(seq_regions)),
        BLOCK_TOPK=triton.next_power_of_2(int(topk)),
    )
    return out


def region_topk_pack(
    logits: torch.Tensor,
    lengths: torch.Tensor,
    q_positions: torch.Tensor,
    block_table: torch.Tensor,
    request_indices: torch.Tensor | None = None,
    *,
    topk: int,
    region_size: int = REGION_BLOCK_SIZE,
    regions_per_page: int,
    block_r: int = 1024,
) -> torch.Tensor:
    """Top-k region selection fused with prefill metadata packing.

    Args:
        logits: ``[rows, seq_regions]`` float32 indexer scores, which may be signed, where
            ``rows = groups * seq_q`` in group-major order.
        lengths: ``[rows]`` int32, causally visible region count per row.
        q_positions: ``[rows]`` int32, absolute token position of each row's query.
        block_table: ``[num_reqs, pages]`` int32, logical page -> physical page, one row
            per request. ``-1`` marks an unmapped page. A table with a single row is
            treated as shared by every request, regardless of ``request_indices``.
        request_indices: ``[rows]`` int32, which block-table row each score row belongs
            to. Defaults to all-zero.

    Returns:
        ``[rows, topk]`` int32, each valid slot holding
        ``phys_region | (valid_tokens << 24)``, with ``-1`` padding the tail. A row fills
        exactly ``min(topk, visible)`` slots.
    """
    if logits.dtype != torch.float32:
        raise ValueError(f"selector expects float32 scores, got {logits.dtype}")
    if block_table.ndim != 2:
        raise ValueError(
            f"block_table must be [num_reqs, pages], got {tuple(block_table.shape)}"
        )

    rows, seq_regions = logits.shape
    if request_indices is None:
        request_indices = torch.zeros((rows,), device=logits.device, dtype=torch.int32)
    out = torch.empty((rows, int(topk)), device=logits.device, dtype=torch.int32)
    _region_topk_pack_kernel[(rows,)](
        logits,
        lengths,
        q_positions,
        block_table,
        request_indices,
        out,
        seq_regions,
        logits.stride(0),
        shared_table_stride(block_table),
        out.stride(0),
        topk=int(topk),
        region_size=int(region_size),
        regions_per_page=int(regions_per_page),
        BLOCK_R=min(block_r, triton.next_power_of_2(seq_regions)),
        BLOCK_TOPK=triton.next_power_of_2(int(topk)),
    )
    return out


# Sparse metadata. The query's current region is appended after selecting
# past regions. Prefill packs phys | valid_tokens << 24; decode packs
# start_token | phys << 32 and uses region_counts to exclude padding.


DECODE_PHYS_SHIFT = tl.constexpr(32)
# Larger than any real packed word, so invalid slots sort to the tail.
DECODE_SORT_SENTINEL = tl.constexpr(1 << 62)


@triton.jit
def _append_current_region_kernel(
    history_ptr,
    q_positions_ptr,
    block_table_ptr,
    request_ptr,
    out_ptr,
    counts_ptr,
    stride_history_row,
    stride_table_row,
    stride_out_row,
    topk: tl.constexpr,
    region_size: tl.constexpr,
    regions_per_page: tl.constexpr,
    BLOCK_K: tl.constexpr,
) -> None:
    """Widen prefill history metadata to ``topk + 1`` slots and append the current region.

    The history is copied verbatim -- including its ``-1`` slots, which stay *in place*
    rather than being compacted, because prefill's ``region_counts`` is positional
    (``min(current_region, topk) + 1``) and not a count of live entries. The appended slot
    lands at exactly the index the selector stopped filling.

    The two stores are kept **disjoint** (``slots != history_count``) rather than letting
    the second overwrite the first. Writing the whole row and then patching one slot is a
    cross-thread race: the row store is a vector store spread over lanes, the patch is a
    scalar store from one lane, and nothing orders them. It loses roughly one row in a
    thousand, so correctness requires disjoint writes even when small tests pass.
    """
    row = tl.program_id(0)
    q_position = tl.load(q_positions_ptr + row).to(tl.int32)
    request = tl.load(request_ptr + row).to(tl.int32)
    current_region = tl.maximum(q_position // region_size, 0)
    history_count = tl.minimum(current_region, topk)

    slots = tl.arange(0, BLOCK_K)
    history = tl.load(
        history_ptr + row * stride_history_row + slots, mask=slots < topk, other=-1
    )
    tl.store(
        out_ptr + row * stride_out_row + slots,
        history,
        mask=(slots < topk + 1) & (slots != history_count),
    )

    page = tl.load(
        block_table_ptr + request * stride_table_row + current_region // regions_per_page
    )
    physical = page * regions_per_page + current_region % regions_per_page
    valid_tokens = tl.minimum(
        tl.maximum(q_position + 1 - current_region * region_size, 0), region_size
    )
    packed = physical | (valid_tokens << PREFILL_REGION_VALID_SHIFT)
    tl.store(
        out_ptr + row * stride_out_row + history_count,
        tl.where(page >= 0, packed, -1),
    )
    tl.store(counts_ptr + row, history_count + 1)


@triton.jit
def _decode_pack_kernel(
    ids_ptr,
    seqlens_ptr,
    block_table_ptr,
    request_ptr,
    out_ptr,
    counts_ptr,
    stride_ids_row,
    stride_table_row,
    stride_out_row,
    topk: tl.constexpr,
    region_size: tl.constexpr,
    regions_per_page: tl.constexpr,
    BLOCK_K: tl.constexpr,
) -> None:
    """Pack selected logical regions into the int64 decode word.

    One row fits in one tile, so the compaction is a sort rather than a scatter: invalid
    slots are mapped to a sentinel above every real word, and sorting by the packed value
    both drops them to the tail and produces physical-major order. The current region
    is written at the surviving count after sorting, so it stays last.
    """
    row = tl.program_id(0)
    request = tl.load(request_ptr + row).to(tl.int32)
    kv_seqlen = tl.load(seqlens_ptr + request).to(tl.int32)
    valid_regions = (kv_seqlen + region_size - 1) // region_size

    slots = tl.arange(0, BLOCK_K)
    ids = tl.load(ids_ptr + row * stride_ids_row + slots, mask=slots < topk, other=-1)

    live = (slots < topk) & (ids >= 0) & (ids < valid_regions)
    safe_ids = tl.where(live, ids, 0)
    pages = tl.load(
        block_table_ptr + request * stride_table_row + safe_ids // regions_per_page,
        mask=live,
        other=-1,
    )
    live = live & (pages >= 0)

    physical = pages * regions_per_page + safe_ids % regions_per_page
    packed = (safe_ids * region_size).to(tl.int64) | (
        physical.to(tl.int64) << DECODE_PHYS_SHIFT
    )
    packed = tl.where(live, packed, DECODE_SORT_SENTINEL)
    packed = tl.sort(packed, descending=False)

    count = tl.sum(live.to(tl.int32))
    tl.store(
        out_ptr + row * stride_out_row + slots,
        tl.where(slots < count, packed, 0),
        # Slot ``count`` is left to the scalar store below. Overlapping the two would be a
        # cross-thread race between a vector store and a scalar one -- see
        # :func:`_append_current_region_kernel`.
        mask=(slots < topk + 1) & (slots != count),
    )

    current_region = tl.maximum((kv_seqlen - 1) // region_size, 0)
    current_page = tl.load(
        block_table_ptr + request * stride_table_row + current_region // regions_per_page
    )
    current_live = (current_region < valid_regions) & (current_page >= 0)
    current_physical = (
        current_page * regions_per_page + current_region % regions_per_page
    )
    current_packed = (current_region * region_size).to(tl.int64) | (
        current_physical.to(tl.int64) << DECODE_PHYS_SHIFT
    )
    # Written unconditionally so the slot is defined even when the current region is
    # unmapped, in which case it becomes ordinary zero padding.
    tl.store(
        out_ptr + row * stride_out_row + count,
        tl.where(current_live, current_packed, 0),
    )
    tl.store(counts_ptr + row, count + current_live.to(tl.int32))


def prefill_sparse_meta(
    logits: torch.Tensor,
    q_positions: torch.Tensor,
    block_table: torch.Tensor,
    request_indices: torch.Tensor | None = None,
    *,
    topk: int,
    region_size: int = REGION_BLOCK_SIZE,
    regions_per_page: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Select regions and build prefill sparse metadata.

    Args:
        logits: ``[rows, seq_regions]`` float32 indexer scores, which may be signed,
            group-major (``rows = num_kv_groups * total_q``). The scores are weighted
            sums of per-head ReLU terms and unconstrained signed indexer weights.
        q_positions: ``[rows]`` int32, absolute token position of each row's query.
        block_table: ``[num_reqs, pages]`` int32, logical page -> physical page, one row
            per request. ``-1`` marks an unmapped page, which yields a ``-1`` slot rather
            than an error.
        request_indices: ``[rows]`` int32, which block-table row each score row uses.
            Defaults to all-zero, i.e. every row shares one table.

    Returns:
        ``(packed [rows, topk + 1] int32, region_counts [rows] int32)``, ready for
        :func:`inference.sparse_attn.sparse_attention_prefill`.

    The history length is derived here as ``q_position // region_size`` rather than taken
    from the caller, because getting it wrong by one is the difference between the current
    region being force-included and it competing for a slot.
    """
    rows = int(logits.shape[0])
    if request_indices is None:
        request_indices = torch.zeros((rows,), device=logits.device, dtype=torch.int32)
    history_lengths = torch.div(
        q_positions.to(torch.int32), region_size, rounding_mode="floor"
    ).contiguous()

    history = region_topk_pack(
        logits,
        history_lengths,
        q_positions,
        block_table,
        request_indices,
        topk=topk,
        region_size=region_size,
        regions_per_page=regions_per_page,
    )

    packed = torch.empty((rows, int(topk) + 1), device=logits.device, dtype=torch.int32)
    counts = torch.empty((rows,), device=logits.device, dtype=torch.int32)
    _append_current_region_kernel[(rows,)](
        history,
        q_positions,
        block_table,
        request_indices,
        packed,
        counts,
        history.stride(0),
        shared_table_stride(block_table),
        packed.stride(0),
        topk=int(topk),
        region_size=int(region_size),
        regions_per_page=int(regions_per_page),
        BLOCK_K=triton.next_power_of_2(int(topk) + 1),
    )
    return packed, counts


def decode_sparse_meta(
    logits: torch.Tensor,
    kv_seqlens: torch.Tensor,
    block_table: torch.Tensor,
    request_indices: torch.Tensor | None = None,
    *,
    topk: int,
    region_size: int = REGION_BLOCK_SIZE,
    regions_per_page: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Select regions and build decode sparse metadata.

    Args:
        logits: ``[num_kv_groups * num_reqs, seq_regions]`` float32 indexer scores,
            which may be signed, group-major. The selector orders them numerically rather
            than by the magnitude encoded in their IEEE-754 bit pattern.
        kv_seqlens: ``[num_reqs]`` int32, context length *including* the current token.
        block_table: ``[num_reqs, pages]`` int32, logical page -> physical page.
        request_indices: ``[rows]`` int32. Defaults to the group-major layout implied by
            the row count, ``arange(num_reqs)`` tiled once per KV group.

    Returns:
        ``(packed [rows, topk + 1] int64, region_counts [rows] int32)``, ready for
        :func:`inference.sparse_attn.sparse_attention_decode`.

    ``region_counts`` here is a count of *surviving* entries, so it varies with how much of
    the block table is mapped. The tail beyond it is zero, and a zero word is not
    self-identifying -- the count is the only bound.
    """
    rows, _ = logits.shape
    num_reqs = int(kv_seqlens.shape[0])
    if rows % num_reqs:
        raise ValueError(f"{rows} score rows do not divide into {num_reqs} requests")

    if request_indices is None:
        request_indices = (
            torch.arange(num_reqs, device=logits.device, dtype=torch.int32)
            .repeat(rows // num_reqs)
            .contiguous()
        )
    # Decode's query position is the last token, so the history is everything strictly
    # before the current region.
    history_lengths = torch.div(
        (kv_seqlens.to(torch.int32) - 1).clamp_min(0), region_size, rounding_mode="floor"
    )[request_indices.to(torch.long)].contiguous()

    ids = region_topk_ids(logits, history_lengths, topk=topk)

    packed = torch.empty((rows, int(topk) + 1), device=logits.device, dtype=torch.int64)
    counts = torch.empty((rows,), device=logits.device, dtype=torch.int32)
    _decode_pack_kernel[(rows,)](
        ids,
        kv_seqlens,
        block_table,
        request_indices,
        packed,
        counts,
        ids.stride(0),
        shared_table_stride(block_table),
        packed.stride(0),
        topk=int(topk),
        region_size=int(region_size),
        regions_per_page=int(regions_per_page),
        BLOCK_K=triton.next_power_of_2(int(topk) + 1),
    )
    return packed, counts


# Sparse attention over each query's selected regions. Causality is encoded
# in metadata; softmax covers only gathered tokens.


REGION_VALID_SHIFT = tl.constexpr(24)
REGION_ID_MASK = tl.constexpr((1 << 24) - 1)
DECODE_LOW_MASK = tl.constexpr((1 << 32) - 1)
NEGATIVE_INFINITY = tl.constexpr(float("-inf"))


@triton.jit
def _sparse_attn_prefill_kernel(
    q_ptr,
    k_ptr,
    v_ptr,
    packed_ptr,
    counts_ptr,
    out_ptr,
    lse_ptr,
    stride_q_token,
    stride_q_head,
    stride_k_token,
    stride_k_group,
    stride_v_token,
    stride_v_group,
    stride_packed_row,
    stride_out_token,
    stride_out_head,
    stride_lse_token,
    softmax_scale,
    total_q,
    topk: tl.constexpr,
    heads_per_group: tl.constexpr,
    head_dim: tl.constexpr,
    region_size: tl.constexpr,
    BLOCK_H: tl.constexpr,
    BLOCK_D: tl.constexpr,
    BLOCK_R: tl.constexpr,
) -> None:
    """Attention for one query against one KV group's selected regions.

    Online softmax over region tiles. The accumulator is per (head, dim) and the running
    max is per head, exactly as in dense flash attention -- the only difference is that
    the K/V rows come from a gather driven by the packed region list instead of a
    contiguous range.
    """
    query = tl.program_id(0)
    group = tl.program_id(1)
    row = group * total_q + query
    count = tl.load(counts_ptr + row).to(tl.int32)

    heads = tl.arange(0, BLOCK_H)
    dims = tl.arange(0, BLOCK_D)
    head_ok = heads < heads_per_group
    dim_ok = dims < head_dim
    tile_ok = head_ok[:, None] & dim_ok[None, :]

    queries = tl.load(
        q_ptr
        + query * stride_q_token
        + (group * heads_per_group + heads)[:, None] * stride_q_head
        + dims[None, :],
        mask=tile_ok,
        other=0.0,
    )

    running_max = tl.full((BLOCK_H,), NEGATIVE_INFINITY, dtype=tl.float32)
    running_sum = tl.zeros((BLOCK_H,), dtype=tl.float32)
    accumulator = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)

    offsets_in_region = tl.arange(0, region_size)
    for base in tl.range(0, topk, BLOCK_R):
        slots = base + tl.arange(0, BLOCK_R)
        slot_ok = slots < count
        meta = tl.load(
            packed_ptr + row * stride_packed_row + slots, mask=slot_ok, other=-1
        )
        physical = meta & REGION_ID_MASK
        valid = meta >> REGION_VALID_SHIFT
        # A slot is live only if it was within the count, not a -1 sentinel, and actually
        # exposes tokens. The third condition is not redundant: a region whose first token
        # is past the query position packs a valid id with zero visible tokens.
        live = slot_ok & (meta >= 0) & (valid > 0)

        tokens = tl.reshape(
            physical[:, None] * region_size + offsets_in_region[None, :],
            (BLOCK_R * region_size,),
        )
        token_ok = tl.reshape(
            live[:, None] & (offsets_in_region[None, :] < valid[:, None]),
            (BLOCK_R * region_size,),
        )
        safe_tokens = tl.where(token_ok, tokens, 0)

        keys = tl.load(
            k_ptr
            + safe_tokens[:, None] * stride_k_token
            + group * stride_k_group
            + dims[None, :],
            mask=token_ok[:, None] & dim_ok[None, :],
            other=0.0,
        )
        values = tl.load(
            v_ptr
            + safe_tokens[:, None] * stride_v_token
            + group * stride_v_group
            + dims[None, :],
            mask=token_ok[:, None] & dim_ok[None, :],
            other=0.0,
        )

        scores = tl.dot(queries, tl.trans(keys), out_dtype=tl.float32) * softmax_scale
        scores = tl.where(token_ok[None, :] & head_ok[:, None], scores, NEGATIVE_INFINITY)

        tile_max = tl.maximum(running_max, tl.max(scores, axis=1))
        # A tile can be entirely masked -- every slot exhausted, or a row with no
        # selection at all. Then both maxima are -inf and ``m - m`` is NaN, which would
        # poison the accumulator for the rest of the loop. Substituting a finite shift
        # keeps the arithmetic well defined and leaves the running state untouched.
        empty = tile_max == NEGATIVE_INFINITY
        shift = tl.where(empty, 0.0, tile_max)
        rescale = tl.where(empty, 1.0, tl.exp(running_max - shift))

        probabilities = tl.exp(scores - shift[:, None])
        running_sum = running_sum * rescale + tl.sum(probabilities, axis=1)
        accumulator = accumulator * rescale[:, None] + tl.dot(
            probabilities.to(values.dtype), values, out_dtype=tl.float32
        )
        running_max = tile_max

    selected = running_sum > 0
    output = tl.where(selected[:, None], accumulator / running_sum[:, None], 0.0)
    tl.store(
        out_ptr
        + query * stride_out_token
        + (group * heads_per_group + heads)[:, None] * stride_out_head
        + dims[None, :],
        output.to(out_ptr.dtype.element_ty),
        mask=tile_ok,
    )
    tl.store(
        lse_ptr + query * stride_lse_token + group * heads_per_group + heads,
        tl.where(selected, running_max + tl.log(running_sum), NEGATIVE_INFINITY),
        mask=head_ok,
    )


def sparse_attention_prefill(
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    packed_regions: torch.Tensor,
    region_counts: torch.Tensor,
    *,
    num_kv_groups: int,
    region_size: int = 8,
    softmax_scale: float | None = None,
    block_regions: int = 16,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Sparse attention over per-query selected regions.

    Args:
        query: ``[total_q, num_heads, head_dim]``.
        key_cache, value_cache: ``[tokens, num_kv_groups, head_dim]``, i.e. the paged
            cache already flattened across pages. Flattening is safe because a region's
            tokens are contiguous in the physical token space by construction.
        packed_regions: ``[num_kv_groups * total_q, topk]`` int32 from the selector,
            each slot ``phys_region | (valid_tokens << 24)``, ``-1`` padding the tail.
        region_counts: ``[num_kv_groups * total_q]`` int32, how many slots to read. This
            truncates *positionally* -- it does not re-select by score.

    Returns:
        ``(out, lse)`` shaped ``[total_q, num_heads, head_dim]`` and
        ``[total_q, num_heads]``. Rows with no selected region get a zero output and
        ``lse = -inf``.

    The caller must guarantee the selected regions are distinct. Our selector does (the
    stage 3 gate checks it), but a union-with-sliding-window path would not, and duplicate
    regions would silently double-weight those tokens in the softmax rather than fail.
    """
    total_q, num_heads, head_dim = query.shape
    if num_heads % num_kv_groups:
        raise ValueError(f"{num_heads} heads do not divide into {num_kv_groups} groups")
    if packed_regions.shape[0] != num_kv_groups * total_q:
        raise ValueError(
            f"expected {num_kv_groups * total_q} metadata rows, got {packed_regions.shape[0]}"
        )
    if key_cache.shape != value_cache.shape:
        raise ValueError(
            f"key/value cache shape mismatch: {tuple(key_cache.shape)} vs {tuple(value_cache.shape)}"
        )

    heads_per_group = num_heads // num_kv_groups
    topk = int(packed_regions.shape[1])
    scale = softmax_scale if softmax_scale is not None else head_dim**-0.5

    out = torch.empty_like(query)
    lse = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32)
    _sparse_attn_prefill_kernel[(total_q, num_kv_groups)](
        query,
        key_cache,
        value_cache,
        packed_regions,
        region_counts,
        out,
        lse,
        query.stride(0),
        query.stride(1),
        key_cache.stride(0),
        key_cache.stride(1),
        value_cache.stride(0),
        value_cache.stride(1),
        packed_regions.stride(0),
        out.stride(0),
        out.stride(1),
        lse.stride(0),
        scale,
        total_q,
        topk=topk,
        heads_per_group=heads_per_group,
        head_dim=head_dim,
        region_size=region_size,
        BLOCK_H=max(16, triton.next_power_of_2(heads_per_group)),
        BLOCK_D=triton.next_power_of_2(head_dim),
        BLOCK_R=min(block_regions, triton.next_power_of_2(topk)),
    )
    return out, lse


@triton.jit
def _sparse_attn_decode_kernel(
    q_ptr,
    k_ptr,
    v_ptr,
    packed_ptr,
    counts_ptr,
    seqlens_ptr,
    out_ptr,
    lse_ptr,
    stride_q_token,
    stride_q_head,
    stride_k_token,
    stride_k_group,
    stride_v_token,
    stride_v_group,
    stride_packed_row,
    stride_out_split,
    stride_out_token,
    stride_out_head,
    stride_lse_split,
    stride_lse_token,
    softmax_scale,
    num_reqs,
    topk: tl.constexpr,
    heads_per_group: tl.constexpr,
    head_dim: tl.constexpr,
    region_size: tl.constexpr,
    slots_per_split: tl.constexpr,
    BLOCK_H: tl.constexpr,
    BLOCK_D: tl.constexpr,
    BLOCK_R: tl.constexpr,
) -> None:
    """One decode query against one KV group's selected regions, over one slot window.

    Structurally the same online softmax as prefill, but the metadata contract differs in
    the one way that matters: decode packs ``start_token | (phys_region << 32)`` in int64
    and does **not** carry ``valid_tokens``. The visible length is recomputed here as
    ``clamp(kv_seqlen - start_token, 0, region_size)``, which is why the logical start
    token has to be in the word at all -- the physical region alone cannot say how much of
    the last region exists yet.

    That also disposes of both padding sentinels without a special case. ``-1`` and
    ``0x7FFF_FFFF_FFFF_FFFF`` both leave ``0xFFFF_FFFF`` in the low word, so
    ``kv_seqlen - start_token`` is hugely negative and the slot clamps to zero tokens.
    """
    request = tl.program_id(0)
    group = tl.program_id(1)
    split = tl.program_id(2)
    row = group * num_reqs + request
    count = tl.load(counts_ptr + row).to(tl.int32)
    kv_seqlen = tl.load(seqlens_ptr + request).to(tl.int64)

    heads = tl.arange(0, BLOCK_H)
    dims = tl.arange(0, BLOCK_D)
    head_ok = heads < heads_per_group
    dim_ok = dims < head_dim
    tile_ok = head_ok[:, None] & dim_ok[None, :]

    queries = tl.load(
        q_ptr
        + request * stride_q_token
        + (group * heads_per_group + heads)[:, None] * stride_q_head
        + dims[None, :],
        mask=tile_ok,
        other=0.0,
    )

    running_max = tl.full((BLOCK_H,), NEGATIVE_INFINITY, dtype=tl.float32)
    running_sum = tl.zeros((BLOCK_H,), dtype=tl.float32)
    accumulator = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)

    offsets_in_region = tl.arange(0, region_size)
    window_begin = split * slots_per_split
    for base in tl.range(window_begin, window_begin + slots_per_split, BLOCK_R):
        slots = base + tl.arange(0, BLOCK_R)
        slot_ok = (slots < count) & (slots < topk)
        meta = tl.load(
            packed_ptr + row * stride_packed_row + slots, mask=slot_ok, other=-1
        )
        physical = (meta >> DECODE_PHYS_SHIFT) & DECODE_LOW_MASK
        start_token = meta & DECODE_LOW_MASK
        valid = tl.minimum(tl.maximum(kv_seqlen - start_token, 0), region_size)
        live = slot_ok & (valid > 0)

        # ``physical`` is a sentinel-sized integer on dead slots, so it must be
        # neutralised before it becomes an address -- masking the load is not enough,
        # since the pointer arithmetic itself would overflow.
        safe_physical = tl.where(live, physical, 0)
        tokens = tl.reshape(
            safe_physical[:, None] * region_size + offsets_in_region[None, :],
            (BLOCK_R * region_size,),
        )
        token_ok = tl.reshape(
            live[:, None] & (offsets_in_region[None, :] < valid[:, None]),
            (BLOCK_R * region_size,),
        )

        keys = tl.load(
            k_ptr + tokens[:, None] * stride_k_token + group * stride_k_group + dims[None, :],
            mask=token_ok[:, None] & dim_ok[None, :],
            other=0.0,
        )
        values = tl.load(
            v_ptr + tokens[:, None] * stride_v_token + group * stride_v_group + dims[None, :],
            mask=token_ok[:, None] & dim_ok[None, :],
            other=0.0,
        )

        scores = tl.dot(queries, tl.trans(keys), out_dtype=tl.float32) * softmax_scale
        scores = tl.where(token_ok[None, :] & head_ok[:, None], scores, NEGATIVE_INFINITY)

        tile_max = tl.maximum(running_max, tl.max(scores, axis=1))
        empty = tile_max == NEGATIVE_INFINITY
        shift = tl.where(empty, 0.0, tile_max)
        rescale = tl.where(empty, 1.0, tl.exp(running_max - shift))

        probabilities = tl.exp(scores - shift[:, None])
        running_sum = running_sum * rescale + tl.sum(probabilities, axis=1)
        accumulator = accumulator * rescale[:, None] + tl.dot(
            probabilities.to(values.dtype), values, out_dtype=tl.float32
        )
        running_max = tile_max

    selected = running_sum > 0
    output = tl.where(selected[:, None], accumulator / running_sum[:, None], 0.0)
    tl.store(
        out_ptr
        + split * stride_out_split
        + request * stride_out_token
        + (group * heads_per_group + heads)[:, None] * stride_out_head
        + dims[None, :],
        output.to(out_ptr.dtype.element_ty),
        mask=tile_ok,
    )
    tl.store(
        lse_ptr
        + split * stride_lse_split
        + request * stride_lse_token
        + group * heads_per_group
        + heads,
        tl.where(selected, running_max + tl.log(running_sum), NEGATIVE_INFINITY),
        mask=head_ok,
    )


@triton.jit
def _merge_split_states_kernel(
    partial_out_ptr,
    partial_lse_ptr,
    out_ptr,
    lse_ptr,
    stride_po_split,
    stride_po_token,
    stride_po_head,
    stride_pl_split,
    stride_pl_token,
    stride_out_token,
    stride_out_head,
    stride_lse_token,
    num_splits,
    head_dim: tl.constexpr,
    BLOCK_S: tl.constexpr,
    BLOCK_D: tl.constexpr,
) -> None:
    """Combine per-split attention states into one, weighting by the partial lse.

    Each split produced an output already divided by its own softmax denominator, so the
    combination is a softmax over the partial log-sum-exps rather than a plain sum:
    ``out = sum_s softmax(lse)_s * out_s`` and ``lse = logsumexp_s(lse_s)``. A split that
    covered no live slot reports ``lse = -inf`` and drops out with weight zero; if *every*
    split does, the row selected nothing and must come back as a zero output with
    ``lse = -inf`` rather than ``0/0``.
    """
    token = tl.program_id(0)
    head = tl.program_id(1)

    splits = tl.arange(0, BLOCK_S)
    dims = tl.arange(0, BLOCK_D)
    split_ok = splits < num_splits
    dim_ok = dims < head_dim

    partial_lse = tl.load(
        partial_lse_ptr + splits * stride_pl_split + token * stride_pl_token + head,
        mask=split_ok,
        other=NEGATIVE_INFINITY,
    )
    peak = tl.max(partial_lse)
    empty = peak == NEGATIVE_INFINITY
    shift = tl.where(empty, 0.0, peak)
    weights = tl.where(split_ok, tl.exp(partial_lse - shift), 0.0)
    total = tl.sum(weights)

    states = tl.load(
        partial_out_ptr
        + splits[:, None] * stride_po_split
        + token * stride_po_token
        + head * stride_po_head
        + dims[None, :],
        mask=split_ok[:, None] & dim_ok[None, :],
        other=0.0,
    ).to(tl.float32)
    merged = tl.sum(states * weights[:, None], axis=0) / tl.maximum(total, 1.0e-20)

    tl.store(
        out_ptr + token * stride_out_token + head * stride_out_head + dims,
        merged.to(out_ptr.dtype.element_ty),
        mask=dim_ok,
    )
    tl.store(
        lse_ptr + token * stride_lse_token + head,
        tl.where(empty, NEGATIVE_INFINITY, shift + tl.log(total)),
    )


def merge_split_states(
    partial_out: torch.Tensor, partial_lse: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Reduce ``[splits, tokens, heads, dim]`` partial states to ``[tokens, heads, dim]``.

    Each request uses the same split count. Partial outputs are combined using
    their natural-log softmax normalization values.
    """
    num_splits, tokens, num_heads, head_dim = partial_out.shape
    out = torch.empty((tokens, num_heads, head_dim), device=partial_out.device, dtype=partial_out.dtype)
    lse = torch.empty((tokens, num_heads), device=partial_out.device, dtype=torch.float32)
    _merge_split_states_kernel[(tokens, num_heads)](
        partial_out,
        partial_lse,
        out,
        lse,
        partial_out.stride(0),
        partial_out.stride(1),
        partial_out.stride(2),
        partial_lse.stride(0),
        partial_lse.stride(1),
        out.stride(0),
        out.stride(1),
        lse.stride(0),
        num_splits,
        head_dim=head_dim,
        BLOCK_S=triton.next_power_of_2(num_splits),
        BLOCK_D=triton.next_power_of_2(head_dim),
    )
    return out, lse


def sparse_attention_decode(
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    packed_regions: torch.Tensor,
    region_counts: torch.Tensor,
    kv_seqlens: torch.Tensor,
    *,
    num_kv_groups: int,
    region_size: int = 8,
    softmax_scale: float | None = None,
    num_splits: int = 1,
    block_regions: int = 16,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Sparse attention for one query token per request.

    Args:
        query: ``[num_reqs, num_heads, head_dim]``.
        key_cache, value_cache: ``[tokens, num_kv_groups, head_dim]``, pages flattened.
        packed_regions: ``[num_kv_groups * num_reqs, topk]`` **int64**, each slot
            ``start_token | (phys_region << 32)``. Note this is a different packing from
            prefill's int32 ``phys_region | (valid_tokens << 24)``: decode carries the
            logical start token and derives the visible length from ``kv_seqlens``.
        region_counts: ``[num_kv_groups * num_reqs]`` int32, slots to read per row.
        kv_seqlens: ``[num_reqs]`` int32, context length including the current token.
        num_splits: split the slot window this many ways and return partial states. With
            ``num_splits == 1`` the returned tensors are already the final result;
            otherwise pass them to :func:`merge_split_states`.

    Returns:
        ``num_splits == 1``: ``(out [num_reqs, num_heads, head_dim], lse [num_reqs, num_heads])``.
        Otherwise the same with a leading ``num_splits`` axis.

    As in prefill, the selected regions must be distinct. Sparse decode is
    stricter about this than the prefill one: it concatenates the gathered spans with no
    dedup at all, so a duplicated region would be counted twice in the softmax.
    """
    num_reqs, num_heads, head_dim = query.shape
    if num_heads % num_kv_groups:
        raise ValueError(f"{num_heads} heads do not divide into {num_kv_groups} groups")
    if packed_regions.shape[0] != num_kv_groups * num_reqs:
        raise ValueError(
            f"expected {num_kv_groups * num_reqs} metadata rows, got {packed_regions.shape[0]}"
        )
    if packed_regions.dtype != torch.int64:
        raise ValueError(f"decode metadata must be int64, got {packed_regions.dtype}")
    if kv_seqlens.shape[0] != num_reqs:
        raise ValueError(f"expected {num_reqs} kv_seqlens, got {kv_seqlens.shape[0]}")

    heads_per_group = num_heads // num_kv_groups
    topk = int(packed_regions.shape[1])
    scale = softmax_scale if softmax_scale is not None else head_dim**-0.5
    # Every split walks the same number of slots so the window bound is a constexpr; the
    # tail split simply finds every slot masked out.
    slots_per_split = -(-topk // num_splits)

    out = torch.empty((num_splits, num_reqs, num_heads, head_dim), device=query.device, dtype=query.dtype)
    lse = torch.empty((num_splits, num_reqs, num_heads), device=query.device, dtype=torch.float32)
    _sparse_attn_decode_kernel[(num_reqs, num_kv_groups, num_splits)](
        query,
        key_cache,
        value_cache,
        packed_regions,
        region_counts,
        kv_seqlens,
        out,
        lse,
        query.stride(0),
        query.stride(1),
        key_cache.stride(0),
        key_cache.stride(1),
        value_cache.stride(0),
        value_cache.stride(1),
        packed_regions.stride(0),
        out.stride(0),
        out.stride(1),
        out.stride(2),
        lse.stride(0),
        lse.stride(1),
        scale,
        num_reqs,
        topk=topk,
        heads_per_group=heads_per_group,
        head_dim=head_dim,
        region_size=region_size,
        slots_per_split=slots_per_split,
        BLOCK_H=max(16, triton.next_power_of_2(heads_per_group)),
        BLOCK_D=triton.next_power_of_2(head_dim),
        BLOCK_R=min(block_regions, triton.next_power_of_2(slots_per_split)),
    )
    if num_splits == 1:
        return out[0], lse[0]
    return out, lse


# DSA state: main K/V, summaries of completed regions and pending region tokens.
# Block tables map request-local pages to physical cache storage.


def build_rope_cache(
    *,
    rotary_span: int,
    theta: float,
    max_position: int,
    device: torch.device,
    dtype: torch.dtype = torch.float32,
    scaling: dict | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Cos/sin tables for a NEOX-style partial rotation.

    Args:
        rotary_span: how many of the head's dimensions rotate. The returned tables are
            ``[max_position, rotary_span // 2]`` -- *half* the span, because each frequency
            drives the pair ``(d, d + rotary_span // 2)``. The kernels take that halved
            width as their ``rotary_dim``, which is easy to misread as the span itself.
        scaling: the ``rope_scaling`` block, or ``None``. Only ``llama3`` is implemented,
            because it is the only type step5's config asks for; an unrecognised type
            raises rather than silently falling back to unscaled, which would look like a
            long-context quality regression rather than a bug.
    """
    if rotary_span % 2:
        raise ValueError(f"rotary span must be even, got {rotary_span}")
    half = rotary_span // 2
    # Compute frequencies and trigonometry in FP32, then round to the requested
    # cache dtype. This precision boundary determines the BF16 RoPE table.
    inv_freq = 1.0 / (
        theta ** (torch.arange(0, rotary_span, 2, device=device, dtype=torch.float32) / rotary_span)
    )
    if scaling:
        inv_freq = _apply_llama3_scaling(inv_freq, scaling)

    positions = torch.arange(max_position, device=device, dtype=torch.float32)
    freqs = torch.outer(positions, inv_freq)
    assert freqs.shape[1] == half
    return freqs.cos().to(dtype).contiguous(), freqs.sin().to(dtype).contiguous()


def _apply_llama3_scaling(inv_freq: torch.Tensor, scaling: dict) -> torch.Tensor:
    rope_type = str(scaling.get("rope_type", scaling.get("type", ""))).lower()
    if rope_type != "llama3":
        raise NotImplementedError(f"unsupported rope_scaling type {rope_type!r}")

    factor = float(scaling["factor"])
    low_freq_factor = float(scaling["low_freq_factor"])
    high_freq_factor = float(scaling["high_freq_factor"])
    original_max = float(scaling["original_max_position_embeddings"])

    low_wavelen = original_max / low_freq_factor
    high_wavelen = original_max / high_freq_factor
    wavelen = 2 * math.pi / inv_freq

    # Long wavelengths are stretched by the full factor, short ones are left alone, and the
    # band between is interpolated. Applying the stretch everywhere is the classic mistake:
    # it destroys the high-frequency positional detail the short-range heads rely on.
    scaled = torch.where(wavelen > low_wavelen, inv_freq / factor, inv_freq)
    smooth = (original_max / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
    smoothed = (1 - smooth) * scaled / factor + smooth * scaled
    in_band = (wavelen >= high_wavelen) & (wavelen <= low_wavelen)
    return torch.where(in_band, smoothed, scaled)


@dataclass(frozen=True)
class DSAGeometry:
    """The region/page arithmetic, in one place.

    ``regions_per_page`` is derived rather than configured: it is the ratio the metadata
    kernels use to split a logical region id into a page lookup and an offset, so a
    configured value that disagreed with ``page_size // region_size`` would address the
    wrong tokens.
    """

    proxy_dim: int
    topk: int
    region_size: int = REGION_BLOCK_SIZE
    page_size: int = 16

    @property
    def regions_per_page(self) -> int:
        if self.page_size % self.region_size:
            raise ValueError(
                f"page size {self.page_size} must be a multiple of region size {self.region_size}"
            )
        return self.page_size // self.region_size


class DSALayerCache:
    """Per-layer DSA state for one batch of in-flight sequences.

    Allocated once for a whole generation, so ``max_tokens`` is a capacity rather than a
    length; ``lengths`` tracks how much of it is live. Every tensor is addressed by the
    conventions the module docstring lists, and mixing them up is the failure this class
    exists to prevent. ``num_kv_groups`` is the *local* count: it is one for the production
    TP=8 layout, with the corresponding global KV head selected while pre-sharding.
    """

    def __init__(
        self,
        *,
        batch: int,
        max_tokens: int,
        num_kv_groups: int,
        head_dim: int,
        geometry: DSAGeometry,
        device: torch.device,
        dtype: torch.dtype,
    ) -> None:
        pages = -(-max_tokens // geometry.page_size)
        self.geometry = geometry
        self.batch = batch
        # Rounded up to a whole page so a request's token span never straddles the boundary
        # between two requests' page ranges.
        self.max_tokens = pages * geometry.page_size
        self.max_regions = self.max_tokens // geometry.region_size

        self.key = torch.zeros(
            (batch * self.max_tokens, num_kv_groups, head_dim), device=device, dtype=dtype
        )
        self.value = torch.zeros_like(self.key)
        self.summary = torch.zeros(
            (batch * self.max_regions, 1, geometry.proxy_dim),
            device=device,
            dtype=torch.float8_e4m3fn,
        )
        self.pending_key = torch.zeros(
            (batch, geometry.region_size, 1, geometry.proxy_dim), device=device, dtype=torch.bfloat16
        )
        self.pending_z = torch.zeros_like(self.pending_key)
        self.lengths = torch.zeros((batch,), device=device, dtype=torch.int32)

        # Request b owns pages [b * pages, (b + 1) * pages), so the mapping is an offset
        # identity. Materialising it as a real table rather than special-casing the
        # contiguous layout keeps the one code path the metadata gates actually cover.
        self.block_table = (
            torch.arange(pages, device=device, dtype=torch.int32).view(1, -1)
            + torch.arange(batch, device=device, dtype=torch.int32).view(-1, 1) * pages
        ).contiguous()

    def token_slice(self, request: int, start: int, count: int) -> slice:
        base = request * self.max_tokens + start
        return slice(base, base + count)

    def write_kv(self, request: int, start: int, key: torch.Tensor, value: torch.Tensor) -> None:
        """Store ``key``/``value`` for one request's tokens at logical position ``start``."""
        span = self.token_slice(request, start, key.shape[0])
        self.key[span] = key
        self.value[span] = value


class Step5SparseIndexer:
    """The indexer's projections and scoring, borrowing its parameters from the attention.

    Not an ``nn.Module``: the checkpoint stores the indexer weights directly under
    ``self_attn`` (``self_attn.sparse_indexer_q.weight``), so wrapping them in a submodule
    would rename every key and require a conversion step to load.
    """

    def __init__(self, owner, *, geometry: DSAGeometry, num_kv_groups: int) -> None:
        self.owner = owner
        self.geometry = geometry
        # Provider groups and attention KV groups are the same partition.  Under TP they
        # are both local counts; for TP=8, one provider group (four indexer heads) drives
        # one local KV group (eight attention Q heads).
        self.num_kv_groups = num_kv_groups
        self.num_heads = int(owner.sparse_indexer_w.weight.shape[0])
        q_out = int(owner.sparse_indexer_q.weight.shape[0])
        if q_out != self.num_heads * geometry.proxy_dim:
            raise ValueError(
                "local sparse_indexer_q/w shapes disagree: "
                f"q rows={q_out}, w heads={self.num_heads}, proxy_dim={geometry.proxy_dim}. "
                "This commonly means stale TP shards from the old replicated-provider layout."
            )
        if self.num_heads % num_kv_groups:
            raise ValueError(
                f"{self.num_heads} indexer heads do not divide into {num_kv_groups} groups"
            )
        self.heads_per_group = self.num_heads // num_kv_groups
        # The indexer's groups *are* the attention's KV groups, so this ratio is fixed under
        # every tensor-parallel degree and the prescale below is a TP invariant.
        self.weight_prescale = float(self.heads_per_group) ** -0.5

    def project(
        self, hidden_states: torch.Tensor, positions: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Project and normalise the indexer's q/k/z, and scale its per-head weights.

        Returns:
            ``(index_q [tokens, groups, heads_per_group, proxy], index_k [tokens, 1, proxy],
            index_z [tokens, 1, proxy], weights [tokens, groups, heads_per_group] fp32)``.
        """

        owner = self.owner
        proxy_dim = self.geometry.proxy_dim
        tokens = hidden_states.shape[0]

        index_q = owner.sparse_indexer_q(hidden_states)
        index_k = owner.sparse_indexer_k(hidden_states)
        index_z = owner.sparse_indexer_z(hidden_states)

        index_q, index_k, index_z = indexer_norm_rope(
            index_q.contiguous(),
            index_k.contiguous(),
            index_z.contiguous(),
            # Keep registered norm parameters in FP32 for checkpoint loading;
            # round them to the activation dtype before the fused operation.
            owner.sparse_indexer_q_norm.weight.to(dtype=index_q.dtype),
            owner.sparse_indexer_k_norm.weight.to(dtype=index_k.dtype),
            owner.sparse_indexer_k_norm.bias.to(dtype=index_k.dtype),
            owner.sparse_indexer_rope_cos,
            owner.sparse_indexer_rope_sin,
            positions,
            head_dim=proxy_dim,
            num_q_heads=self.num_heads,
            num_k_heads=1,
            rotary_dim=owner.sparse_indexer_rope_cos.shape[1],
            eps=owner.sparse_indexer_q_norm.variance_epsilon,
            q_norm_weight_bias=1.0,
        )

        # Store the checkpoint weight in FP32, round it to the activation dtype
        # for the GEMM, then widen the result before applying the prescale.
        weights = torch.nn.functional.linear(
            hidden_states,
            owner.sparse_indexer_w.weight.to(dtype=hidden_states.dtype),
        ).float()
        weights = weights.view(tokens, self.num_kv_groups, self.heads_per_group)
        weights = weights * self.weight_prescale

        return (
            index_q.view(tokens, self.num_kv_groups, self.heads_per_group, proxy_dim),
            index_k.view(tokens, 1, proxy_dim),
            index_z.view(tokens, 1, proxy_dim),
            weights.contiguous(),
        )


def update_summaries_prefill(
    cache: DSALayerCache,
    index_k: torch.Tensor,
    index_z: torch.Tensor,
    seq_lens: torch.Tensor,
) -> None:
    """Compress a packed prefill into per-region summaries, buffering the ragged tail.

    Only regions that are *complete* get a summary. The tail tokens are held instead, and
    the region they belong to is summarised later, once it fills -- see the module
    docstring for why a partial summary must never be stored.
    """
    geometry = cache.geometry
    region_size = geometry.region_size
    starts, counts = csa_region_layout(
        seq_lens, regions_per_seq=cache.max_regions, region_size=region_size
    )
    # A count below ``region_size`` marks the ragged tail; zero makes the kernel skip it.
    counts = torch.where(counts == region_size, counts, torch.zeros_like(counts))

    _, summary_fp8 = csa_compress_regions(
        index_k, index_z, starts, counts, region_size=region_size
    )
    live = counts > 0
    cache.summary[live] = summary_fp8[live]

    token_starts = torch.cumsum(seq_lens.to(torch.int64), dim=0) - seq_lens.to(torch.int64)
    for request in range(cache.batch):
        length = int(seq_lens[request])
        tail = length % region_size
        cache.pending_key[request, :tail] = index_k[
            int(token_starts[request]) + length - tail : int(token_starts[request]) + length
        ].to(cache.pending_key.dtype)
        cache.pending_z[request, :tail] = index_z[
            int(token_starts[request]) + length - tail : int(token_starts[request]) + length
        ].to(cache.pending_z.dtype)


def update_summaries_decode(
    cache: DSALayerCache, index_k: torch.Tensor, index_z: torch.Tensor
) -> None:
    """Append one token per request and summarise any region that just filled.

    ``cache.lengths`` must already have been advanced past the new token, so the pending
    slot is derived from it rather than tracked separately -- one counter that can be wrong
    instead of two that can disagree.
    """
    geometry = cache.geometry
    region_size = geometry.region_size
    lengths = cache.lengths.tolist()

    for request in range(cache.batch):
        length = int(lengths[request])
        slot = (length - 1) % region_size
        cache.pending_key[request, slot] = index_k[request].to(cache.pending_key.dtype)
        cache.pending_z[request, slot] = index_z[request].to(cache.pending_z.dtype)
        if slot != region_size - 1:
            continue

        region = (length - 1) // region_size
        starts = torch.zeros((1,), device=index_k.device, dtype=torch.int32)
        counts = torch.full((1,), region_size, device=index_k.device, dtype=torch.int32)
        _, summary_fp8 = csa_compress_regions(
            cache.pending_key[request].contiguous(),
            cache.pending_z[request].contiguous(),
            starts,
            counts,
            region_size=region_size,
        )
        cache.summary[request * cache.max_regions + region] = summary_fp8[0]


def score_regions(
    indexer: Step5SparseIndexer,
    cache: DSALayerCache,
    index_q: torch.Tensor,
    weights: torch.Tensor,
    seq_lens: torch.Tensor,
    past_lens: torch.Tensor,
) -> torch.Tensor:
    """Indexer scores for every query against its own request's region summaries.

    Returns ``[groups * total_q, width]`` float32 in the group-major row order the selector
    requires, where ``width`` is the largest number of regions any row can see. Rows whose own
    history is shorter are left at zero past their bound; that is safe only because the
    selector's candidate range is bounded per row by its own history length and never reads
    further.

    ``width`` is computed rather than taken as ``cache.max_regions``, and the difference is not
    cosmetic. ``max_regions`` is the cache's *capacity*: at a 512k context it is 65536, so a
    4k-token prefill would allocate ``4 groups x 4096 queries x 65536`` fp32 -- 4.3 GB, for a
    tensor whose live corner is a thousandth of that. Sizing to the live width makes the
    allocation track the context rather than the reservation.

    The loop is over requests, not a batched GEMM, because each request scores against a
    different key set. At decode that is one tiny matmul per request; at prefill it is dwarfed
    by the attention it feeds.
    """
    groups = indexer.num_kv_groups
    total_q = int(index_q.shape[0])
    region_size = cache.geometry.region_size

    # A row can see every region completed before its query, so the widest row is the one with
    # the largest final position.
    last_positions = [
        int(past_lens[request]) + int(seq_lens[request]) - 1
        for request in range(cache.batch)
        if int(seq_lens[request]) > 0
    ]
    width = max((position // region_size for position in last_positions), default=0)
    logits = torch.zeros(
        (groups * total_q, max(width, 1)), device=index_q.device, dtype=torch.float32
    )

    offset = 0
    for request in range(cache.batch):
        length = int(seq_lens[request])
        if length == 0:
            continue
        regions = (int(past_lens[request]) + length - 1) // region_size
        if regions == 0:
            offset += length
            continue

        base = request * cache.max_regions
        block = indexer_logits(
            index_q[offset : offset + length].contiguous(),
            weights[offset : offset + length].contiguous(),
            cache.summary[base : base + regions].to(index_q.dtype),
        )
        for group in range(groups):
            rows = slice(group * total_q + offset, group * total_q + offset + length)
            logits[rows, :regions] = block[group * length : (group + 1) * length]
        offset += length

    return logits


def prefill_metadata(
    indexer: Step5SparseIndexer,
    cache: DSALayerCache,
    index_q: torch.Tensor,
    weights: torch.Tensor,
    seq_lens: torch.Tensor,
    past_lens: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Score, select, and pack prefill metadata for a packed batch."""
    groups = indexer.num_kv_groups
    device = index_q.device

    logits = score_regions(indexer, cache, index_q, weights, seq_lens, past_lens)
    positions, requests = _packed_positions(seq_lens, past_lens, device)
    return prefill_sparse_meta(
        logits,
        positions.repeat(groups).contiguous(),
        cache.block_table,
        requests.repeat(groups).contiguous(),
        topk=indexer.geometry.topk,
        region_size=indexer.geometry.region_size,
        regions_per_page=indexer.geometry.regions_per_page,
    )


def decode_metadata(
    indexer: Step5SparseIndexer,
    cache: DSALayerCache,
    index_q: torch.Tensor,
    weights: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Score, select, and pack decode metadata for one token per request."""
    ones = torch.ones((cache.batch,), device=index_q.device, dtype=torch.int32)
    logits = score_regions(
        indexer, cache, index_q, weights, ones, (cache.lengths - 1).clamp_min(0)
    )
    return decode_sparse_meta(
        logits,
        cache.lengths,
        cache.block_table,
        topk=indexer.geometry.topk,
        region_size=indexer.geometry.region_size,
        regions_per_page=indexer.geometry.regions_per_page,
    )


def _packed_positions(
    seq_lens: torch.Tensor, past_lens: torch.Tensor, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor]:
    """Absolute position and owning request of every token in a packed batch."""
    positions = []
    requests = []
    for request in range(seq_lens.shape[0]):
        length = int(seq_lens[request])
        past = int(past_lens[request])
        positions.append(torch.arange(past, past + length, device=device, dtype=torch.int32))
        requests.append(torch.full((length,), request, device=device, dtype=torch.int32))
    return torch.cat(positions), torch.cat(requests)


# FP8 expert GEMM. Activations use 128-wide K blocks; weights use 128x128
# blocks. Both scales multiply the FP8 values when reconstructing real values.
# A missing weight scale selects the BF16 path.


# fp8_e4m3fn has 4 exponent bits and 3 mantissa bits, with range [-448, 448].
# The activation quantizer maps each block's range into that envelope by
# dividing by ``max(amax, 1e-10) / 448``. The 1e-10 floor preserves
# low-amplitude blocks in Step-5's FP8 activation path.
FP8_E4M3_MAX: float = 448.0
FP8_AMAX_FLOOR: float = 1e-10

# K-block size shared by the activation quantizer and the weight scale grid.
# The kernel asserts BLOCK_K == GROUP_K so one K-iteration consumes exactly
# one act scale entry and one weight scale entry per row.
BLOCK_K: int = 128
# Weight N-block. The kernel asserts BLOCK_N == GROUP_N so one output tile
# sits inside exactly one weight N-block, keeping the weight scale a scalar
# per (tile, K-iter) instead of a vector.
BLOCK_N: int = 128


@triton.jit
def _act_quant_kernel(
    x_ptr,
    y_ptr,
    scale_ptr,
    M,
    K,
    stride_xm,
    stride_xk,
    stride_ym,
    stride_yk,
    stride_sm,
    stride_sk,
    BLOCK_M: tl.constexpr,
    BLOCK_K: tl.constexpr,
) -> None:
    """Per-row 1-D block quantization to fp8_e4m3fn along K.

    One program handles a ``[BLOCK_M, BLOCK_K]`` tile -- exactly one quantization
    group wide. The amax is computed over the *live* elements of each row, so
    the kernel handles a final ragged K-block where the tile extends past K:
    the scale covers the live elements only, and the OOB quantized slots are
    not written (the GEMM's masked load will read them as zero, so the round
    trip stays numerically faithful).
    """
    pid_m = tl.program_id(0)
    pid_k = tl.program_id(1)

    rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    cols = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
    row_mask = rows < M
    col_mask = cols < K
    tile_mask = row_mask[:, None] & col_mask[None, :]

    x = tl.load(
        x_ptr + rows[:, None] * stride_xm + cols[None, :] * stride_xk,
        mask=tile_mask,
        other=0.0,
    ).to(tl.float32)

    # OOB elements are zero via ``other=0.0``, so they do not inflate the amax.
    amax = tl.max(tl.abs(x), axis=1)
    amax = tl.maximum(amax, 1e-10)
    scale = amax * (1.0 / 448.0)
    inv_scale = 1.0 / scale

    # Clamp to [-448, 448] before the round-to-fp8: a value just above 448 would
    # otherwise wrap to NaN instead of saturating.
    q = tl.clamp(x * inv_scale[:, None], -448.0, 448.0)

    tl.store(
        y_ptr + rows[:, None] * stride_ym + cols[None, :] * stride_yk,
        q.to(y_ptr.dtype.element_ty),
        mask=tile_mask,
    )
    tl.store(
        scale_ptr + rows * stride_sm + pid_k * stride_sk,
        scale,
        mask=row_mask,
    )


def act_quant(
    x: torch.Tensor, block_size: int = 128
) -> tuple[torch.Tensor, torch.Tensor]:
    """Block-wise FP8 quantization along the last dim of ``x``.

    Args:
        x: input tensor of shape ``[..., K]``. The last dim is the reduction
            dim K. Any leading dims are flattened; the result is reshaped back.
        block_size: K-block size, 128 by default. Must be a power of two so
            Triton accepts it as a constexpr bound on ``tl.arange``.

    Returns:
        ``(y, scale)`` where ``y`` has ``x``'s shape and dtype
        ``torch.float8_e4m3fn``, and ``scale`` has shape ``[..., ceil(K/block_size)]``
        and dtype ``torch.float32``. The recovery formula is
        ``real_x = y.float() * scale.unsqueeze(-1)``.

    K need not be divisible by block_size. For a partial final block, amax
    includes only live elements.
    """
    if block_size & (block_size - 1):
        raise ValueError(f"block_size must be a power of two, got {block_size}")
    leading = x.shape[:-1]
    K = x.shape[-1]
    x2 = x.contiguous()
    flat = x2.view(-1, K)
    M = flat.shape[0]
    y = torch.empty_like(flat, dtype=torch.float8_e4m3fn)
    scale = torch.empty(
        M, (K + block_size - 1) // block_size, device=x.device, dtype=torch.float32
    )
    BLOCK_M = 16 if M < 32 else 32
    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(K, block_size))
    _act_quant_kernel[grid](
        flat,
        y,
        scale,
        M,
        K,
        flat.stride(0),
        flat.stride(1),
        y.stride(0),
        y.stride(1),
        scale.stride(0),
        scale.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_K=block_size,
    )
    return y.view(*leading, K), scale.view(*leading, -1)


@triton.jit
def _fp8_gemm_kernel(
    A_ptr,
    B_ptr,
    C_ptr,
    A_scale_ptr,
    W_scale_ptr,
    M,
    N,
    K,
    stride_am,
    stride_ak,
    stride_bn,
    stride_bk,
    stride_cm,
    stride_cn,
    stride_asm,
    stride_ask,
    stride_wsn,
    stride_wsk,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_N: tl.constexpr,
) -> None:
    """fp8 GEMM with per-128 block scaling on both operands.

    One program computes a ``[BLOCK_M, BLOCK_N]`` output tile. ``BLOCK_N`` must
    equal ``GROUP_N`` (128) so the tile sits inside exactly one weight N-block,
    and ``BLOCK_K`` must equal the activation block size (128) so one K-iter
    consumes exactly one act scale entry and one weight scale entry per row.

    Loads B transposed (as ``[BLOCK_K, BLOCK_N]``) so ``tl.dot`` sees ``A @ B``
    in the natural orientation. The strides are the *original* N/K strides of
    B; the transpose is by index, not by stride.

    The two scales are folded into the fp32 accumulator after the fp8 dot, not
    into the fp8 operands. Folding into the operands would re-quantize and
    lose the precision fp8 GEMM is supposed to keep; the post-dot multiply is
    the whole reason the accumulator is fp32.
    """
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    m_mask = offs_m < M
    n_mask = offs_n < N

    a_ptrs = A_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = B_ptr + offs_n[None, :] * stride_bn + offs_k[:, None] * stride_bk

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # Weight scale N index: tile (pid_n) sits in weight N-block ``pid_n`` when
    # ``BLOCK_N == GROUP_N``. Tiles never span weight N-block boundaries.
    n_block = pid_n * BLOCK_N // GROUP_N

    k_iters = (K + BLOCK_K - 1) // BLOCK_K
    for k_idx in tl.range(0, k_iters):
        k_offset = k_idx * BLOCK_K
        k_mask = (k_offset + offs_k) < K

        a = tl.load(
            a_ptrs + k_offset * stride_ak,
            mask=m_mask[:, None] & k_mask[None, :],
            other=0.0,
        )
        b = tl.load(
            b_ptrs + k_offset * stride_bk,
            mask=n_mask[None, :] & k_mask[:, None],
            other=0.0,
        )
        # fp8 x fp8 -> fp32 on Hopper; this is the WGMMA path that makes the
        # scheme worth the precision cost. The accumulator is fp32 so the post-
        # dot scale multiply keeps full precision.
        partial = tl.dot(a, b, out_dtype=tl.float32)

        a_s = tl.load(
            A_scale_ptr + offs_m * stride_asm + k_idx * stride_ask,
            mask=m_mask,
            other=0.0,
        )
        # Scalar weight scale for this (n_block, k_idx) -- one fp32 per tile.
        w_s = tl.load(W_scale_ptr + n_block * stride_wsn + k_idx * stride_wsk)

        acc += partial * a_s[:, None] * w_s

    c_ptrs = C_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(
        c_ptrs,
        acc.to(C_ptr.dtype.element_ty),
        mask=m_mask[:, None] & n_mask[None, :],
    )


def fp8_gemm(
    x_fp8: torch.Tensor,
    x_scale: torch.Tensor,
    w_fp8: torch.Tensor,
    w_scale_inv: torch.Tensor,
) -> torch.Tensor:
    """Block-scaled fp8 matmul ``y = x @ w.T``.

    Args:
        x_fp8: ``[M, K]`` ``torch.float8_e4m3fn`` activation, already quantized
            by :func:`act_quant`.
        x_scale: ``[M, ceil(K/128)]`` ``torch.float32`` activation scales.
        w_fp8: ``[N, K]`` ``torch.float8_e4m3fn`` weight (one expert's slice).
        w_scale_inv: ``[ceil(N/128), ceil(K/128)]`` ``torch.float32`` weight
            scales stored in the *inverse* convention
            (``real_weight = w_fp8 * w_scale_inv[block]``).

    Returns:
        ``y`` of shape ``[M, N]`` in ``torch.bfloat16``.
    """
    if x_fp8.dtype != torch.float8_e4m3fn or w_fp8.dtype != torch.float8_e4m3fn:
        raise TypeError("operands must be float8_e4m3fn")
    if x_scale.dtype != torch.float32 or w_scale_inv.dtype != torch.float32:
        raise TypeError("scales must be float32")
    if x_fp8.dim() != 2 or w_fp8.dim() != 2:
        raise ValueError("operands must be 2-D; group by expert at the call site")

    M, K = x_fp8.shape
    N, K_w = w_fp8.shape
    if K != K_w:
        raise ValueError(f"K mismatch: x has {K}, weight has {K_w}")
    k_blocks = (K + BLOCK_K - 1) // BLOCK_K
    if x_scale.shape != (M, k_blocks):
        raise ValueError(
            f"x_scale shape {tuple(x_scale.shape)} != ({M}, {k_blocks})"
        )
    if w_scale_inv.shape != ((N + BLOCK_N - 1) // BLOCK_N, k_blocks):
        raise ValueError(
            f"w_scale_inv shape {tuple(w_scale_inv.shape)} != "
            f"({(N + BLOCK_N - 1) // BLOCK_N}, {k_blocks})"
        )

    x_fp8 = x_fp8.contiguous()
    w_fp8 = w_fp8.contiguous()
    x_scale = x_scale.contiguous()
    w_scale_inv = w_scale_inv.contiguous()

    y = torch.empty((M, N), device=x_fp8.device, dtype=torch.bfloat16)
    BLOCK_M = 64 if M >= 64 else 32
    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    _fp8_gemm_kernel[grid](
        x_fp8,
        w_fp8,
        y,
        x_scale,
        w_scale_inv,
        M,
        N,
        K,
        x_fp8.stride(0),
        x_fp8.stride(1),
        w_fp8.stride(0),
        w_fp8.stride(1),
        y.stride(0),
        y.stride(1),
        x_scale.stride(0),
        x_scale.stride(1),
        w_scale_inv.stride(0),
        w_scale_inv.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        GROUP_N=BLOCK_N,
    )
    return y


def linear_fp8_or_bf16(
    x: torch.Tensor,
    weight: torch.Tensor,
    weight_scale_inv: torch.Tensor | None,
) -> torch.Tensor:
    """``y = x @ weight.T`` with fp8 path when a scale is present, BF16 fallback.

    The dispatcher is the only place that decides fp8 vs BF16: layers 88/89/90
    ship BF16 experts with no scale, and the criterion is *presence of the
    scale tensor*, never the layer index. A layer can graduate off the
    not-convert list without code changes here.

    Args:
        x: ``[..., K]`` activation. Cast to BF16 internally on the fallback path
            (the fp8 path quantizes through BF16 anyway).
        weight: ``[N, K]`` weight; ``torch.float8_e4m3fn`` when
            ``weight_scale_inv`` is present, otherwise BF16.
        weight_scale_inv: ``[ceil(N/128), ceil(K/128)]`` fp32 inverse scales, or
            ``None`` to select the BF16 fallback path.

    Returns:
        ``[..., N]`` BF16 activation.
    """
    if weight_scale_inv is None:
        if weight.dtype != torch.bfloat16:
            weight = weight.to(torch.bfloat16)
        x_bf = x.to(torch.bfloat16) if x.dtype != torch.bfloat16 else x
        return torch.nn.functional.linear(x_bf, weight)

    leading = x.shape[:-1]
    K = x.shape[-1]
    x_bf = x.to(torch.bfloat16) if x.dtype != torch.bfloat16 else x
    x_flat = x_bf.reshape(-1, K).contiguous()
    x_fp8, x_s = act_quant(x_flat)
    y = fp8_gemm(x_fp8, x_s, weight, weight_scale_inv)
    return y.view(*leading, -1)


# MoE primitives. SwiGLU rounds after FP32 clamps and multiplication;
# weighted gather accumulates expert contributions in top-k slot order.


_BLOCK_SIZE = 1024


@triton.jit
def _clamped_swiglu_kernel(
    gate_ptr,
    up_ptr,
    output_ptr,
    num_elements,
    LIMIT: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    """FP32 SWIGLUstep over two independent, contiguous input tensors."""
    block = tl.program_id(0).to(tl.int64)
    offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < num_elements

    gate = tl.load(gate_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
    up = tl.load(up_ptr + offsets, mask=mask, other=0.0).to(tl.float32)

    gate_silu = gate * tl.sigmoid(gate)
    gate_clamped = tl.minimum(gate_silu, LIMIT)
    up_clamped = tl.minimum(tl.maximum(up, -LIMIT), LIMIT)
    output = gate_clamped * up_clamped

    tl.store(
        output_ptr + offsets,
        output.to(output_ptr.dtype.element_ty),
        mask=mask,
    )

@triton.jit
def _weighted_topk_gather_kernel(
    contributions_ptr,
    weights_ptr,
    output_ptr,
    hidden_size,
    TOP_K: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    """Sum one ``[top_k, hidden]`` slice in logical slot order."""
    token = tl.program_id(0).to(tl.int64)
    hidden_block = tl.program_id(1).to(tl.int64)
    hidden_offsets = hidden_block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    hidden_mask = hidden_offsets < hidden_size

    accumulator = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
    # TOP_K is constexpr, just as ``topk_num`` is in Optimus ep_gather.
    # Keeping the update in this loop produces the required slot-ordered
    # FP32 accumulation instead of a tree reduction over the K dimension.
    for slot in range(0, TOP_K):
        contribution_offset = (token * TOP_K + slot) * hidden_size + hidden_offsets
        contribution = tl.load(
            contributions_ptr + contribution_offset,
            mask=hidden_mask,
            other=0.0,
        ).to(tl.float32)
        weight = tl.load(weights_ptr + token * TOP_K + slot).to(tl.float32)
        accumulator += contribution * weight

    output_offset = token * hidden_size + hidden_offsets
    tl.store(
        output_ptr + output_offset,
        accumulator.to(output_ptr.dtype.element_ty),
        mask=hidden_mask,
    )


def _require_cuda(tensors: tuple[torch.Tensor, ...]) -> None:
    if any(tensor.device.type != "cuda" for tensor in tensors):
        devices = ", ".join(str(tensor.device) for tensor in tensors)
        raise RuntimeError(
            f"Step-5 MoE Triton kernels require CUDA tensors, got [{devices}]"
        )


def clamped_swiglu(
    gate: torch.Tensor,
    up: torch.Tensor,
    limit: float = 7.0,
) -> torch.Tensor:
    """Apply Step-5's clamped SwiGLU to independent BF16 projections.

    Computes

    ``silu(gate.float()).clamp(max=limit) * up.float().clamp(-limit, limit)``

    in FP32 and casts the product to the input dtype once.  ``gate`` and ``up``
    must have the same shape and be contiguous BF16 CUDA tensors.
    """
    if gate.shape != up.shape:
        raise ValueError(
            f"gate/up shape mismatch: {tuple(gate.shape)} != {tuple(up.shape)}"
        )
    if gate.dtype != torch.bfloat16 or up.dtype != torch.bfloat16:
        raise TypeError(
            f"gate and up must both be torch.bfloat16, got {gate.dtype} and {up.dtype}"
        )
    if gate.device != up.device:
        raise ValueError(f"gate/up device mismatch: {gate.device} != {up.device}")
    if not gate.is_contiguous() or not up.is_contiguous():
        raise ValueError("gate and up must both be contiguous")

    try:
        limit_value = float(limit)
    except (TypeError, ValueError) as exc:
        raise TypeError(f"limit must be a real scalar, got {limit!r}") from exc
    if not math.isfinite(limit_value) or limit_value < 0.0:
        raise ValueError(f"limit must be finite and non-negative, got {limit_value}")

    _require_cuda((gate, up))
    output = torch.empty_like(gate)
    if gate.numel() == 0:
        return output

    grid = (triton.cdiv(gate.numel(), _BLOCK_SIZE),)
    _clamped_swiglu_kernel[grid](
        gate,
        up,
        output,
        gate.numel(),
        LIMIT=limit_value,
        BLOCK_SIZE=_BLOCK_SIZE,
        num_warps=4,
    )
    return output


def weighted_topk_gather(
    contributions: torch.Tensor,
    weights: torch.Tensor,
) -> torch.Tensor:
    """Weight and sum top-k contributions in deterministic slot order.

    Args:
        contributions: contiguous BF16 tensor of shape ``[T, K, H]``.
        weights: contiguous FP32 tensor of shape ``[T, K]``.

    Returns:
        Contiguous BF16 tensor of shape ``[T, H]``.  For every token and hidden
        element the kernel starts from FP32 zero and performs
        ``acc += contributions[:, slot].float() * weights[:, slot]`` for slots
        ``0, 1, ..., K - 1``, then casts the accumulator to BF16 once.
    """
    if contributions.ndim != 3:
        raise ValueError(
            f"contributions must have shape [T, K, H], got {tuple(contributions.shape)}"
        )
    if weights.ndim != 2:
        raise ValueError(f"weights must have shape [T, K], got {tuple(weights.shape)}")

    tokens, top_k, hidden_size = contributions.shape
    if weights.shape != (tokens, top_k):
        raise ValueError(f"weights shape {tuple(weights.shape)} != ({tokens}, {top_k})")
    if top_k <= 0:
        raise ValueError(f"top-k dimension must be positive, got {top_k}")
    if hidden_size <= 0:
        raise ValueError(f"hidden dimension must be positive, got {hidden_size}")
    if contributions.dtype != torch.bfloat16:
        raise TypeError(
            f"contributions must be torch.bfloat16, got {contributions.dtype}"
        )
    if weights.dtype != torch.float32:
        raise TypeError(f"weights must be torch.float32, got {weights.dtype}")
    if contributions.device != weights.device:
        raise ValueError(
            "contributions/weights device mismatch: "
            f"{contributions.device} != {weights.device}"
        )
    if not contributions.is_contiguous() or not weights.is_contiguous():
        raise ValueError("contributions and weights must both be contiguous")

    _require_cuda((contributions, weights))
    output = torch.empty(
        (tokens, hidden_size),
        dtype=torch.bfloat16,
        device=contributions.device,
    )
    if tokens == 0:
        return output

    grid = (tokens, triton.cdiv(hidden_size, _BLOCK_SIZE))
    _weighted_topk_gather_kernel[grid](
        contributions,
        weights,
        output,
        hidden_size,
        TOP_K=top_k,
        BLOCK_SIZE=_BLOCK_SIZE,
        num_warps=4,
    )
    return output


# Fused Q/K normalization and NeoX RoPE. The BF16 materialization between
# normalization and rotary arithmetic, and the reduction order, are deliberate.


HEAD_DIM = 192
_SUPPORTED_ROTARY_PAIRS = frozenset((32, 96))


@triton.jit
def _lane_square_sum_192(
    qkv_ptr,
    head_base,
    LANE: tl.constexpr,
):
    """Step-5's 32-pair layout: three four-value vectors for one of 16 lanes."""
    acc = 0.0
    # The order is part of the numerical contract:
    #   block 0 dims [0, 64), block 1 [64, 128), block 2 [128, 192),
    # with four adjacent values owned by each logical lane in every block.
    for block in range(3):
        for element in range(4):
            dim = block * 64 + LANE * 4 + element
            value = tl.load(qkv_ptr + head_base + dim).to(tl.float32)
            acc = _libdevice.fma_rn(value, value, acc)
    return acc

@triton.jit
def _wide_lane_square_sum_192(
    qkv_ptr,
    head_base,
    LANE: tl.constexpr,
):
    """32-lane normalization layout for the 96-pair Step-5 path.

    Each lane covers eight adjacent values, padding the 192-value row to 256
    values. Lanes 24..31 contribute eight explicit zeros apiece.
    """
    acc = 0.0
    for element in range(8):
        dim = LANE * 8 + element
        value = tl.load(
            qkv_ptr + head_base + dim,
            mask=dim < 192,
            other=0.0,
        ).to(tl.float32)
        acc = _libdevice.fma_rn(value, value, acc)
    return acc

@triton.jit
def _qknorm_rope_192_kernel(
    qkv_ptr,
    q_out_ptr,
    k_out_ptr,
    q_weight_ptr,
    k_weight_ptr,
    cos_ptr,
    sin_ptr,
    positions_ptr,
    stride_qkv_token,
    stride_q_out_token,
    stride_k_out_token,
    stride_cos_token,
    stride_sin_token,
    NUM_Q_HEADS: tl.constexpr,
    NUM_KV_HEADS: tl.constexpr,
    ROTARY_PAIRS: tl.constexpr,
    EPS: tl.constexpr,
    NORM_WEIGHT_BIAS: tl.constexpr,
    BLOCK_D: tl.constexpr,
    BLOCK_R: tl.constexpr,
):
    """One program per token and Q/K head.

    The physical Triton layout is deliberately not used to define the
    reduction.  Explicit logical lane sums followed by the exact ascending
    XOR tree make the result independent of Triton's evolving warp/layout
    lowering.  Step-5's 32-pair path uses 16 lanes; the 96-pair path uses
    32 so all rotary pairs are covered.
    """
    token = tl.program_id(0)
    packed_head = tl.program_id(1)
    is_q = packed_head < NUM_Q_HEADS
    head_base = token * stride_qkv_token + packed_head * 192

    if ROTARY_PAIRS == 96:
        # Use 32 lanes for the 96-pair path. Normalization covers eight
        # contiguous values per lane, padded to 256. Preserve the local
        # FMA chains and XOR(1,2,4,8,16) reduction order.
        p0 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=0)
        p1 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=1)
        p2 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=2)
        p3 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=3)
        p4 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=4)
        p5 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=5)
        p6 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=6)
        p7 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=7)
        p8 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=8)
        p9 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=9)
        p10 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=10)
        p11 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=11)
        p12 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=12)
        p13 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=13)
        p14 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=14)
        p15 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=15)
        p16 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=16)
        p17 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=17)
        p18 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=18)
        p19 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=19)
        p20 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=20)
        p21 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=21)
        p22 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=22)
        p23 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=23)
        p24 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=24)
        p25 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=25)
        p26 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=26)
        p27 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=27)
        p28 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=28)
        p29 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=29)
        p30 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=30)
        p31 = _wide_lane_square_sum_192(qkv_ptr, head_base, LANE=31)

        s01 = _libdevice.add_rn(p0, p1)
        s23 = _libdevice.add_rn(p2, p3)
        s45 = _libdevice.add_rn(p4, p5)
        s67 = _libdevice.add_rn(p6, p7)
        s89 = _libdevice.add_rn(p8, p9)
        s1011 = _libdevice.add_rn(p10, p11)
        s1213 = _libdevice.add_rn(p12, p13)
        s1415 = _libdevice.add_rn(p14, p15)
        s1617 = _libdevice.add_rn(p16, p17)
        s1819 = _libdevice.add_rn(p18, p19)
        s2021 = _libdevice.add_rn(p20, p21)
        s2223 = _libdevice.add_rn(p22, p23)
        s2425 = _libdevice.add_rn(p24, p25)
        s2627 = _libdevice.add_rn(p26, p27)
        s2829 = _libdevice.add_rn(p28, p29)
        s3031 = _libdevice.add_rn(p30, p31)
        s03 = _libdevice.add_rn(s01, s23)
        s47 = _libdevice.add_rn(s45, s67)
        s811 = _libdevice.add_rn(s89, s1011)
        s1215 = _libdevice.add_rn(s1213, s1415)
        s1619 = _libdevice.add_rn(s1617, s1819)
        s2023 = _libdevice.add_rn(s2021, s2223)
        s2427 = _libdevice.add_rn(s2425, s2627)
        s2831 = _libdevice.add_rn(s2829, s3031)
        s07 = _libdevice.add_rn(s03, s47)
        s815 = _libdevice.add_rn(s811, s1215)
        s1623 = _libdevice.add_rn(s1619, s2023)
        s2431 = _libdevice.add_rn(s2427, s2831)
        s015 = _libdevice.add_rn(s07, s815)
        s1631 = _libdevice.add_rn(s1623, s2431)
        sum_squares = _libdevice.add_rn(s015, s1631)
    else:
        # 32-pair full-attention layers use 16 lanes.  Each pN is the
        # sequential 12-value fma.rn chain of logical lane N.
        p0 = _lane_square_sum_192(qkv_ptr, head_base, LANE=0)
        p1 = _lane_square_sum_192(qkv_ptr, head_base, LANE=1)
        p2 = _lane_square_sum_192(qkv_ptr, head_base, LANE=2)
        p3 = _lane_square_sum_192(qkv_ptr, head_base, LANE=3)
        p4 = _lane_square_sum_192(qkv_ptr, head_base, LANE=4)
        p5 = _lane_square_sum_192(qkv_ptr, head_base, LANE=5)
        p6 = _lane_square_sum_192(qkv_ptr, head_base, LANE=6)
        p7 = _lane_square_sum_192(qkv_ptr, head_base, LANE=7)
        p8 = _lane_square_sum_192(qkv_ptr, head_base, LANE=8)
        p9 = _lane_square_sum_192(qkv_ptr, head_base, LANE=9)
        p10 = _lane_square_sum_192(qkv_ptr, head_base, LANE=10)
        p11 = _lane_square_sum_192(qkv_ptr, head_base, LANE=11)
        p12 = _lane_square_sum_192(qkv_ptr, head_base, LANE=12)
        p13 = _lane_square_sum_192(qkv_ptr, head_base, LANE=13)
        p14 = _lane_square_sum_192(qkv_ptr, head_base, LANE=14)
        p15 = _lane_square_sum_192(qkv_ptr, head_base, LANE=15)

        # Lane 0 after ascending shfl_xor offsets 1, 2, 4, 8.
        s01 = _libdevice.add_rn(p0, p1)
        s23 = _libdevice.add_rn(p2, p3)
        s45 = _libdevice.add_rn(p4, p5)
        s67 = _libdevice.add_rn(p6, p7)
        s89 = _libdevice.add_rn(p8, p9)
        s1011 = _libdevice.add_rn(p10, p11)
        s1213 = _libdevice.add_rn(p12, p13)
        s1415 = _libdevice.add_rn(p14, p15)
        s03 = _libdevice.add_rn(s01, s23)
        s47 = _libdevice.add_rn(s45, s67)
        s811 = _libdevice.add_rn(s89, s1011)
        s1215 = _libdevice.add_rn(s1213, s1415)
        s07 = _libdevice.add_rn(s03, s47)
        s815 = _libdevice.add_rn(s811, s1215)
        sum_squares = _libdevice.add_rn(s07, s815)

    mean_square = _libdevice.div_rn(sum_squares, 192.0)
    inverse_rms = _libdevice.rsqrt(_libdevice.add_rn(mean_square, EPS))

    dims = tl.arange(0, BLOCK_D)
    dim_mask = dims < 192
    values = tl.load(
        qkv_ptr + head_base + dims,
        mask=dim_mask,
        other=0.0,
    ).to(tl.float32)
    q_weight = tl.load(
        q_weight_ptr + dims,
        mask=dim_mask,
        other=0.0,
    ).to(tl.float32)
    k_weight = tl.load(
        k_weight_ptr + dims,
        mask=dim_mask,
        other=0.0,
    ).to(tl.float32)
    weight = tl.where(is_q, q_weight, k_weight)
    weight = _libdevice.add_rn(weight, NORM_WEIGHT_BIAS)
    x_hat = _libdevice.mul_rn(values, inverse_rms)
    normalized = _libdevice.mul_rn(x_hat, weight).to(tl.bfloat16)

    q_base = token * stride_q_out_token + packed_head * 192
    k_base = token * stride_k_out_token + (packed_head - NUM_Q_HEADS) * 192
    tl.store(
        q_out_ptr + q_base + dims,
        normalized,
        mask=dim_mask & is_q,
    )
    tl.store(
        k_out_ptr + k_base + dims,
        normalized,
        mask=dim_mask & ~is_q,
    )

    # Store and reload to round normalized values to the activation dtype
    # before RoPE. The barrier makes the stores visible to the reloads.
    tl.debug_barrier()

    pairs = tl.arange(0, BLOCK_R)
    pair_mask = pairs < ROTARY_PAIRS
    dim0 = pairs
    dim1 = pairs + ROTARY_PAIRS
    q_value0 = tl.load(
        q_out_ptr + q_base + dim0,
        mask=pair_mask & is_q,
        other=0.0,
    ).to(tl.float32)
    q_value1 = tl.load(
        q_out_ptr + q_base + dim1,
        mask=pair_mask & is_q,
        other=0.0,
    ).to(tl.float32)
    k_value0 = tl.load(
        k_out_ptr + k_base + dim0,
        mask=pair_mask & ~is_q,
        other=0.0,
    ).to(tl.float32)
    k_value1 = tl.load(
        k_out_ptr + k_base + dim1,
        mask=pair_mask & ~is_q,
        other=0.0,
    ).to(tl.float32)
    value0 = tl.where(is_q, q_value0, k_value0)
    value1 = tl.where(is_q, q_value1, k_value1)

    position = tl.load(positions_ptr + token)
    cos_value = tl.load(
        cos_ptr + position * stride_cos_token + pairs,
        mask=pair_mask,
        other=0.0,
    ).to(tl.float32)
    sin_value = tl.load(
        sin_ptr + position * stride_sin_token + pairs,
        mask=pair_mask,
        other=0.0,
    ).to(tl.float32)

    # Step-5 RoPE rounding contract:
    #   rot0 = mul(value0, cos) - mul(value1, sin)
    #   rot1 = fma(value0, sin, mul(value1, cos))
    rotated0 = _libdevice.sub_rn(
        _libdevice.mul_rn(value0, cos_value),
        _libdevice.mul_rn(value1, sin_value),
    ).to(tl.bfloat16)
    rotated1 = _libdevice.fma_rn(
        value0,
        sin_value,
        _libdevice.mul_rn(value1, cos_value),
    ).to(tl.bfloat16)

    tl.store(
        q_out_ptr + q_base + dim0,
        rotated0,
        mask=pair_mask & is_q,
    )
    tl.store(
        q_out_ptr + q_base + dim1,
        rotated1,
        mask=pair_mask & is_q,
    )
    tl.store(
        k_out_ptr + k_base + dim0,
        rotated0,
        mask=pair_mask & ~is_q,
    )
    tl.store(
        k_out_ptr + k_base + dim1,
        rotated1,
        mask=pair_mask & ~is_q,
    )

@triton.jit
def _copy_v_192_kernel(
    qkv_ptr,
    v_out_ptr,
    stride_qkv_token,
    stride_v_out_token,
    NUM_Q_HEADS: tl.constexpr,
    NUM_KV_HEADS: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    token = tl.program_id(0)
    head = tl.program_id(1)
    dims = tl.arange(0, BLOCK_D)
    mask = dims < 192
    input_base = (
        token * stride_qkv_token + (NUM_Q_HEADS + NUM_KV_HEADS + head) * 192
    )
    output_base = token * stride_v_out_token + head * 192
    # BF16 load/store only: V must retain its original bits.
    value = tl.load(qkv_ptr + input_base + dims, mask=mask)
    tl.store(v_out_ptr + output_base + dims, value, mask=mask)


def _validate_qknorm_rope_inputs(
    qkv: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    positions: torch.Tensor,
    *,
    head_dim: int,
    num_q_heads: int,
    num_kv_heads: int,
    rotary_pairs: int,
    eps: float,
    norm_weight_bias: float,
) -> None:
    tensors = {
        "qkv": qkv,
        "q_weight": q_weight,
        "k_weight": k_weight,
        "cos": cos,
        "sin": sin,
        "positions": positions,
    }
    for name, tensor in tensors.items():
        if not isinstance(tensor, torch.Tensor):
            raise TypeError(f"{name} must be a torch.Tensor")

    if head_dim != HEAD_DIM:
        raise ValueError(f"head_dim must be {HEAD_DIM}, got {head_dim}")
    if num_q_heads <= 0 or num_kv_heads <= 0:
        raise ValueError("num_q_heads and num_kv_heads must be positive")
    if rotary_pairs not in _SUPPORTED_ROTARY_PAIRS:
        raise ValueError(
            "rotary_pairs must be 32 (full-attention layers) or 96 "
            f"(sliding-attention layers), got {rotary_pairs}"
        )
    if not math.isfinite(eps) or eps <= 0.0:
        raise ValueError(f"eps must be finite and positive, got {eps}")
    if not math.isfinite(norm_weight_bias):
        raise ValueError(f"norm_weight_bias must be finite, got {norm_weight_bias}")

    packed_width = (num_q_heads + 2 * num_kv_heads) * head_dim
    if qkv.ndim != 2 or tuple(qkv.shape) != (qkv.shape[0], packed_width):
        raise ValueError(
            f"qkv must have shape [tokens, {packed_width}], got {tuple(qkv.shape)}"
        )
    if qkv.stride(1) != 1:
        raise ValueError("qkv must be contiguous in its last dimension")
    if qkv.dtype != torch.bfloat16:
        raise TypeError(f"qkv must have dtype torch.bfloat16, got {qkv.dtype}")

    for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)):
        if weight.ndim != 1 or weight.numel() != head_dim:
            raise ValueError(
                f"{name} must have shape [{head_dim}], got {tuple(weight.shape)}"
            )
        if not weight.is_contiguous():
            raise ValueError(f"{name} must be contiguous")
        if weight.dtype != torch.float32:
            raise TypeError(f"{name} must have dtype torch.float32, got {weight.dtype}")

    if cos.ndim != 2 or sin.ndim != 2:
        raise ValueError(
            f"cos and sin must be 2D, got {tuple(cos.shape)} and {tuple(sin.shape)}"
        )
    if cos.shape != sin.shape:
        raise ValueError(
            f"cos and sin must have equal shapes, got {tuple(cos.shape)} "
            f"and {tuple(sin.shape)}"
        )
    if cos.shape[1] < rotary_pairs:
        raise ValueError(
            f"cos/sin width must be at least {rotary_pairs}, got {cos.shape[1]}"
        )
    if cos.stride(1) != 1 or sin.stride(1) != 1:
        raise ValueError("cos and sin must be contiguous in their last dimension")
    if cos.dtype != torch.bfloat16 or sin.dtype != torch.bfloat16:
        raise TypeError(
            "cos and sin must have dtype torch.bfloat16, got "
            f"{cos.dtype} and {sin.dtype}"
        )

    if positions.ndim != 1 or positions.numel() != qkv.shape[0]:
        raise ValueError(
            "positions must have one entry per qkv token, got "
            f"shape {tuple(positions.shape)} for {qkv.shape[0]} tokens"
        )
    if not positions.is_contiguous():
        raise ValueError("positions must be contiguous")
    if positions.dtype not in (torch.int32, torch.int64):
        raise TypeError(
            "positions must have dtype torch.int32 or torch.int64, got "
            f"{positions.dtype}"
        )

    # Device checks intentionally come last.  Shape and dtype are properties of
    # the public contract too, and validating them first lets CPU-only tooling
    # diagnose malformed model inputs without needing a CUDA allocation.
    for name, tensor in tensors.items():
        if not tensor.is_cuda:
            raise ValueError(f"{name} must be a CUDA tensor")
    devices = {tensor.device for tensor in tensors.values()}
    if len(devices) != 1:
        raise ValueError("all inputs must be on the same CUDA device")


def fused_qknorm_rope(
    qkv: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    positions: torch.Tensor,
    *,
    head_dim: int,
    num_q_heads: int,
    num_kv_heads: int,
    rotary_pairs: int,
    eps: float = 1e-5,
    norm_weight_bias: float = 1.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Normalize and rotate packed TP-local Q/K heads, and copy V exactly.

    Args:
        qkv: BF16 ``[tokens, (num_q_heads + 2 * num_kv_heads) * 192]``.
            A wider physical row stride is allowed, but the last dimension must
            be contiguous.
        q_weight, k_weight: FP32 zero-centered RMSNorm weights of shape
            ``[192]``.  The effective scale is ``weight + norm_weight_bias``;
            Step-5 uses the default bias of ``1.0``.
        cos, sin: BF16 NeoX tables ``[max_position, >= rotary_pairs]``.
        positions: contiguous int32 or int64 position for every token.
        rotary_pairs: ``96`` on sliding layers and ``32`` on full-attention
            layers.  The leading ``2 * rotary_pairs`` dimensions are rotated.

    Returns:
        Flat contiguous BF16 tensors ``q``, ``k``, ``v`` with shapes
        ``[tokens, num_q_heads * 192]`` and
        ``[tokens, num_kv_heads * 192]`` (for both K and V).

    The 96-pair sliding path uses 32 lanes and a padded 256-value normalization
    tile, rotating all 96 pairs. The 32-pair full-attention path uses a 16-lane
    blocked reduction layout.
    """
    _validate_qknorm_rope_inputs(
        qkv,
        q_weight,
        k_weight,
        cos,
        sin,
        positions,
        head_dim=head_dim,
        num_q_heads=num_q_heads,
        num_kv_heads=num_kv_heads,
        rotary_pairs=rotary_pairs,
        eps=eps,
        norm_weight_bias=norm_weight_bias,
    )

    tokens = qkv.shape[0]
    q_out = torch.empty(
        (tokens, num_q_heads * HEAD_DIM),
        device=qkv.device,
        dtype=qkv.dtype,
    )
    k_out = torch.empty(
        (tokens, num_kv_heads * HEAD_DIM),
        device=qkv.device,
        dtype=qkv.dtype,
    )
    v_out = torch.empty_like(k_out)
    if tokens == 0:
        return q_out, k_out, v_out

    _qknorm_rope_192_kernel[(tokens, num_q_heads + num_kv_heads)](
        qkv,
        q_out,
        k_out,
        q_weight,
        k_weight,
        cos,
        sin,
        positions,
        qkv.stride(0),
        q_out.stride(0),
        k_out.stride(0),
        cos.stride(0),
        sin.stride(0),
        NUM_Q_HEADS=num_q_heads,
        NUM_KV_HEADS=num_kv_heads,
        ROTARY_PAIRS=rotary_pairs,
        EPS=float(eps),
        NORM_WEIGHT_BIAS=float(norm_weight_bias),
        BLOCK_D=256,
        BLOCK_R=128,
        num_warps=4,
    )
    _copy_v_192_kernel[(tokens, num_kv_heads)](
        qkv,
        v_out,
        qkv.stride(0),
        v_out.stride(0),
        NUM_Q_HEADS=num_q_heads,
        NUM_KV_HEADS=num_kv_heads,
        BLOCK_D=256,
        num_warps=4,
    )
    return q_out, k_out, v_out


__all__ = [
    "REGION_BLOCK_SIZE",
    "PROXY_DIM",
    "shared_table_stride",
    "indexer_norm_rope",
    "csa_region_layout",
    "csa_compress_regions",
    "round_activations_e4m3",
    "indexer_logits",
    "region_topk_ids",
    "region_topk_pack",
    "prefill_sparse_meta",
    "decode_sparse_meta",
    "sparse_attention_prefill",
    "merge_split_states",
    "sparse_attention_decode",
    "build_rope_cache",
    "DSAGeometry",
    "DSALayerCache",
    "Step5SparseIndexer",
    "update_summaries_prefill",
    "update_summaries_decode",
    "score_regions",
    "prefill_metadata",
    "decode_metadata",
    "act_quant",
    "fp8_gemm",
    "linear_fp8_or_bf16",
    "clamped_swiglu",
    "weighted_topk_gather",
    "HEAD_DIM",
    "fused_qknorm_rope",
]