vllm.v1.attention.ops.mxfp4_mla_store

Old import path, kept for PRs stacked on this one that still import it.

Use vllm.v1.attention.ops.mxfp4_mla.store.

Functions:

• store_mxfp4_mla –

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

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,
    )