vllm.v1.attention.ops.mxfp4_mla.store

Triton store kernel for the mxfp4_mla KV cache.

concat_and_cache_mla takes a single per-tensor scale, so it cannot write MXFP4's per-group scales. A row is packed data then scales, row_bytes(d) == d // 2 + d // 32, addressed as slot * row_bytes.

Functions:

• store_mxfp4_mla –

  Quantize the MLA latent to MXFP4 and write it into the paged cache.

_ceil_pow2_exponent(v)

Smallest integer e with 2**e >= v, for positive fp32 v.

Done on the bit pattern rather than with log2: fp32 log2 is not exactly representable at powers of two, and rounding one of those up would waste a whole exponent step and halve the usable E2M1 range for that group.

For v = m * 2**e with m in [1, 2), the biased exponent field is e + 127 and the mantissa field is zero exactly when v is a power of two. So ceil(log2(v)) is e when the mantissa is zero and e + 1 otherwise.

Source code in vllm/v1/attention/ops/mxfp4_mla/store.py

@triton.jit
def _ceil_pow2_exponent(v):
    """Smallest integer `e` with `2**e >= v`, for positive fp32 `v`.

    Done on the bit pattern rather than with `log2`: fp32 `log2` is not
    exactly representable at powers of two, and rounding one of those up would
    waste a whole exponent step and halve the usable E2M1 range for that group.

    For `v = m * 2**e` with `m` in `[1, 2)`, the biased exponent field is
    `e + 127` and the mantissa field is zero exactly when `v` is a power of
    two. So `ceil(log2(v))` is `e` when the mantissa is zero and `e + 1`
    otherwise.
    """
    bits = v.to(tl.int32, bitcast=True)
    exponent = ((bits >> 23) & 0xFF) - 127
    mantissa = bits & 0x7FFFFF
    return tl.where(mantissa == 0, exponent, exponent + 1)

_quantize_e2m1(normalized)

fp32 values pre-divided by their group scale -> 4-bit E2M1 codes.

Source code in vllm/v1/attention/ops/mxfp4_mla/store.py

@triton.jit
def _quantize_e2m1(normalized):
    """fp32 values pre-divided by their group scale -> 4-bit E2M1 codes."""
    a = tl.abs(normalized)
    index = (
        (a > _M0).to(tl.int32)
        + (a >= _M1).to(tl.int32)
        + (a > _M2).to(tl.int32)
        + (a >= _M3).to(tl.int32)
        + (a > _M4).to(tl.int32)
        + (a >= _M5).to(tl.int32)
        + (a > _M6).to(tl.int32)
    )
    # The sign bit is kept even when the magnitude rounds to zero, so a small
    # negative encodes as -0.0 (code 8). This matches AITER's reference exactly.
    sign = (normalized < 0).to(tl.int32) << 3
    return (index | sign).to(tl.uint8)

store_mxfp4_mla(kv_c_normed, slot_mapping, cache, group=GROUP_SIZE)

Quantize the MLA latent to MXFP4 and write it into the paged cache.

Parameters:

• kv_c_normed

  (Tensor) –

  [num_tokens, latent], the post-RMSNorm latent.

• slot_mapping

  (Tensor) –

  [num_tokens] global slot indices; negative entries (PAD_SLOT_ID) are skipped.

• cache

  (Tensor) –

  uint8, either the paged [num_blocks, block_size, row_bytes(latent)] layer view, whose block stride may exceed block_size rows when other layers' pages sit between blocks, or any contiguous tensor viewable as [num_slots, row_bytes(latent)].

• group

  (int, default: GROUP_SIZE ) –

  number of latent elements sharing one E8M0 scale.

Source code in vllm/v1/attention/ops/mxfp4_mla/store.py

def store_mxfp4_mla(
    kv_c_normed: torch.Tensor,
    slot_mapping: torch.Tensor,
    cache: torch.Tensor,
    group: int = GROUP_SIZE,
) -> None:
    """Quantize the MLA latent to MXFP4 and write it into the paged cache.

    Args:
        kv_c_normed: `[num_tokens, latent]`, the post-RMSNorm latent.
        slot_mapping: `[num_tokens]` global slot indices; negative entries
            (`PAD_SLOT_ID`) are skipped.
        cache: `uint8`, either the paged `[num_blocks, block_size,
            row_bytes(latent)]` layer view, whose block stride may exceed
            `block_size` rows when other layers' pages sit between blocks, or
            any contiguous tensor viewable as `[num_slots, row_bytes(latent)]`.
        group: number of latent elements sharing one E8M0 scale.

    """
    if kv_c_normed.ndim != 2:
        raise ValueError(
            f"expected kv_c_normed=[tokens, latent], got {kv_c_normed.shape}"
        )
    if cache.dtype != torch.uint8:
        raise ValueError(f"expected a uint8 cache, got {cache.dtype}")

    num_tokens, latent = kv_c_normed.shape
    pitch = row_bytes(latent, group)
    if cache.ndim == 3:
        num_blocks, block_size, width = cache.shape
        if (
            width != pitch
            or cache.stride()[1:] != (pitch, 1)
            or cache.stride(0) % pitch
        ):
            raise ValueError(
                f"paged cache {tuple(cache.shape)} with strides {cache.stride()} "
                f"is not a whole number of {pitch}-byte rows per block"
            )
        block_stride_rows = cache.stride(0) // pitch
    else:
        cache = cache.view(-1)
        if cache.numel() % pitch:
            raise ValueError(
                f"cache of {cache.numel()} bytes is not a multiple of {pitch}"
            )
        num_blocks, block_size = 1, cache.numel() // pitch
        block_stride_rows = block_size
    num_slots = num_blocks * block_size

    if num_tokens == 0:
        return

    _mxfp4_mla_store_kernel[(num_tokens,)](
        kv_c_normed,
        slot_mapping,
        cache,
        kv_c_normed.stride(0),
        num_slots,
        block_size,
        block_stride_rows,
        LATENT=latent,
        GROUP=group,
        NUM_GROUPS=latent // group,
        HALF_GROUP=group // 2,
        ROW_BYTES=pitch,
        SCALE_OFFSET=latent // 2,
        num_warps=4,
    )