vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache

V4.1 state saving/compression and independently schedulable cache insertion.

Functions:

• fused_save_compress_norm –

  Pool each closed group into a normalized BF16 latent; save FP32 states.

• rope_quant_insert –

  Apply GPT-J RoPE and publish a latent to the compressed KV cache.

• save_ring_rows –

  Store FP32 [kv, score] rows to their ring slots, skipping slot -1.

_rope_quant_insert_mxfp8_kernel(latent, positions, cos_sin, cache, cache_slots, COS_STRIDE, CACHE_STRIDE, CACHE_BLOCK, COMPRESS_RATIO, SANITIZE_CACHE_NANS)

V4.1 record: RoPE first, then MXFP8-quantize all 512 dims.

The RoPE dims are quantized here too, so unlike the V4 kernel the rotation has to happen before the scales are picked.

Source code in vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py

@triton.jit
def _rope_quant_insert_mxfp8_kernel(
    latent,
    positions,
    cos_sin,
    cache,
    cache_slots,
    COS_STRIDE: tl.constexpr,
    CACHE_STRIDE: tl.constexpr,
    CACHE_BLOCK: tl.constexpr,
    COMPRESS_RATIO: tl.constexpr,
    SANITIZE_CACHE_NANS: tl.constexpr,
):
    """V4.1 record: RoPE first, then MXFP8-quantize all 512 dims.

    The RoPE dims are quantized here too, so unlike the V4 kernel the rotation
    has to happen before the scales are picked.
    """
    t = tl.program_id(0)
    slot = tl.load(cache_slots + t)
    if slot < 0:
        return
    position = tl.load(positions + t)
    if (position + 1) % COMPRESS_RATIO != 0:
        return
    d = tl.arange(0, 512)
    normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)

    # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
    even, odd = tl.split(tl.reshape(normed, (256, 2)))
    pair = tl.arange(0, 256) - 224
    cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
    c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
    s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
    rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
    if SANITIZE_CACHE_NANS:
        rotated = tl.where(rotated == rotated, rotated, 0.0)

    page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
    values = page + (slot % CACHE_BLOCK) * 512
    scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16

    quant = tl.reshape(rotated, (16, 32))
    amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4)
    exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0)))
    scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1))
    fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv)
    tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,)))

    max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
    encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded)
    tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8))

_rope_quant_insert_nvfp4_kernel(latent, positions, cos_sin, cache, cache_slots, COS_STRIDE, CACHE_STRIDE, CACHE_BLOCK, COMPRESS_RATIO, SANITIZE_CACHE_NANS)

V4.1 NVFP4 record: RoPE, then e2m1 with one e4m3 scale per 16 dims.

The scale is amax / 6 (6 is e2m1's largest magnitude) clamped to the e4m3 range, with no per-tensor scale on top.

Source code in vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py

@triton.jit
def _rope_quant_insert_nvfp4_kernel(
    latent,
    positions,
    cos_sin,
    cache,
    cache_slots,
    COS_STRIDE: tl.constexpr,
    CACHE_STRIDE: tl.constexpr,
    CACHE_BLOCK: tl.constexpr,
    COMPRESS_RATIO: tl.constexpr,
    SANITIZE_CACHE_NANS: tl.constexpr,
):
    """V4.1 NVFP4 record: RoPE, then e2m1 with one e4m3 scale per 16 dims.

    The scale is ``amax / 6`` (6 is e2m1's largest magnitude) clamped to the
    e4m3 range, with no per-tensor scale on top.
    """
    t = tl.program_id(0)
    slot = tl.load(cache_slots + t)
    if slot < 0:
        return
    position = tl.load(positions + t)
    if (position + 1) % COMPRESS_RATIO != 0:
        return
    d = tl.arange(0, 512)
    normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)

    # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
    even, odd = tl.split(tl.reshape(normed, (256, 2)))
    pair = tl.arange(0, 256) - 224
    cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
    c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
    s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
    rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
    if SANITIZE_CACHE_NANS:
        rotated = tl.where(rotated == rotated, rotated, 0.0)

    tiles = tl.reshape(rotated, (32, 16))
    amax = tl.max(tl.abs(tiles), 1)
    # 2**-9 is the smallest normal e4m3 magnitude; 448 the largest.
    scale = tl.clamp(amax * (1.0 / 6.0), 0.001953125, 448.0).to(tl.float8e4nv)
    # Round-to-nearest division: Triton's default div.full misplaces values
    # that land exactly on an e2m1 tie.
    scaled = tl.math.div_rn(tiles, tl.reshape(scale.to(tl.float32), (32, 1)))
    lo, hi = tl.split(tl.reshape(tl.reshape(scaled, (512,)), (256, 2)))

    page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
    tl.store(
        page + (slot % CACHE_BLOCK) * 256 + tl.arange(0, 256),
        _fp32x2_to_fp4x2(lo, hi),
    )
    scales = page + CACHE_BLOCK * 256 + (slot % CACHE_BLOCK) * 32
    tl.store(scales + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True))

fused_save_compress_norm(kv_score, positions, state_cache, slot_mapping, query_start_loc, token_to_req_indices, rms_norm_weight, rms_norm_eps, compress_ratio, latent_out, prev_rows=None, prev_row_indices=None)

Pool each closed group into a normalized BF16 latent; save FP32 states.

The latent feeds the main-cache insert and the indexer K path, which the attention layer schedules on separate streams.

Ratio 2 keeps one ring block per request holding the open group's rows: position p lives in row p % capacity and slot_mapping encodes block * capacity + p % capacity. The grid has one program per request followed by one per pair of packed tokens. A request program handles the group that the chunk's first token closes with its predecessor's ring row, then stores the chunk's last capacity rows to the ring; because the same program does both, ring reads and writes never race. A pair program handles the group that ends inside its pair, reading both rows from the raw input. Ratio 1 has no ring and one program per token; slot_mapping then only marks valid tokens.

Parameters:

• kv_score

  (Tensor) –

  FP32 [tokens, 512] for CR1, [tokens, 1024] for CR2.

• positions

  (Tensor) –

  Absolute positions of the packed request tokens.

• state_cache

  (Tensor | None) –

  Ring FP32 [blocks, capacity, 1024] KV/score states (CR2).

• slot_mapping

  (Tensor) –

  Ring slots (CR2) or main-cache slots (CR1).

• query_start_loc

  (Tensor | None) –

  [num_reqs + 1] token offsets of each request's chunk.

• token_to_req_indices

  (Tensor | None) –

  Request indices for the packed token rows.

• rms_norm_weight

  (Tensor) –

  BF16 [512] normalization weight.

• rms_norm_eps

  (float) –

  RMSNorm epsilon.

• compress_ratio

  (int) –

  Group size, either 1 or 2.

• latent_out

  (Tensor) –

  BF16 [tokens, 512], written only at valid group boundaries.

• prev_rows

  (Tensor | None, default: None ) –

  FP32 [rows, 1024] rows gathered from all PCP ranks (CR2). A chunk reads its predecessor here where prev_row_indices is set, and the ring is left to save_ring_rows.

• prev_row_indices

  (Tensor | None, default: None ) –

  [num_reqs] row in prev_rows holding the token before each chunk's first token, or -1 to read the ring.

Source code in vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py

def fused_save_compress_norm(
    kv_score: torch.Tensor,
    positions: torch.Tensor,
    state_cache: torch.Tensor | None,
    slot_mapping: torch.Tensor,
    query_start_loc: torch.Tensor | None,
    token_to_req_indices: torch.Tensor | None,
    rms_norm_weight: torch.Tensor,
    rms_norm_eps: float,
    compress_ratio: int,
    latent_out: torch.Tensor,
    prev_rows: torch.Tensor | None = None,
    prev_row_indices: torch.Tensor | None = None,
) -> None:
    """Pool each closed group into a normalized BF16 latent; save FP32 states.

    The latent feeds the main-cache insert and the indexer K path, which the
    attention layer schedules on separate streams.

    Ratio 2 keeps one ring block per request holding the open group's rows:
    position ``p`` lives in row ``p % capacity`` and ``slot_mapping`` encodes
    ``block * capacity + p % capacity``. The grid has one program per request
    followed by one per pair of packed tokens. A request program handles the
    group that the chunk's first token closes with its predecessor's ring row,
    then stores the chunk's last ``capacity`` rows to the ring; because the
    same program does both, ring reads and writes never race. A pair program
    handles the group that ends inside its pair, reading both rows from the
    raw input. Ratio 1 has no ring and one program per token; ``slot_mapping``
    then only marks valid tokens.

    Args:
        kv_score: FP32 [tokens, 512] for CR1, [tokens, 1024] for CR2.
        positions: Absolute positions of the packed request tokens.
        state_cache: Ring FP32 [blocks, capacity, 1024] KV/score states (CR2).
        slot_mapping: Ring slots (CR2) or main-cache slots (CR1).
        query_start_loc: [num_reqs + 1] token offsets of each request's chunk.
        token_to_req_indices: Request indices for the packed token rows.
        rms_norm_weight: BF16 [512] normalization weight.
        rms_norm_eps: RMSNorm epsilon.
        compress_ratio: Group size, either 1 or 2.
        latent_out: BF16 [tokens, 512], written only at valid group boundaries.
        prev_rows: FP32 [rows, 1024] rows gathered from all PCP ranks (CR2). A
            chunk reads its predecessor here where ``prev_row_indices`` is set,
            and the ring is left to ``save_ring_rows``.
        prev_row_indices: [num_reqs] row in ``prev_rows`` holding the token
            before each chunk's first token, or -1 to read the ring.

    """
    assert compress_ratio in (1, 2)
    assert kv_score.dtype == torch.float32
    assert kv_score.shape[1] == 512 * compress_ratio and kv_score.stride(1) == 1
    assert latent_out.shape == (kv_score.shape[0], 512)
    assert latent_out.is_contiguous() and latent_out.dtype == torch.bfloat16
    assert positions.is_contiguous() and slot_mapping.is_contiguous()
    # Rows stay 64-byte aligned, as the kernel's tl.multiple_of hint promises.
    assert kv_score.stride(0) % 16 == 0
    if compress_ratio == 2:
        assert state_cache is not None and query_start_loc is not None
        assert token_to_req_indices is not None
        assert query_start_loc.is_contiguous()
        assert token_to_req_indices.is_contiguous()
        assert state_cache.dtype == torch.float32
        assert state_cache.shape[2] == 1024 and state_cache.stride(2) == 1
        assert state_cache.stride(1) % 16 == 0
        state_stride, state_row_stride, state_block = (
            state_cache.stride(0),
            state_cache.stride(1),
            state_cache.shape[1],
        )
        num_reqs = query_start_loc.numel() - 1
        if prev_rows is not None:
            assert prev_row_indices is not None
            assert prev_row_indices.numel() >= num_reqs
            assert prev_rows.dtype == torch.float32 and prev_rows.shape[1] == 1024
            assert prev_rows.stride(1) == 1 and prev_rows.stride(0) % 16 == 0
    else:
        state_cache = query_start_loc = token_to_req_indices = None
        state_stride = state_row_stride = state_block = 1
        num_reqs = 0
    num_tokens = slot_mapping.numel()
    assert num_tokens <= min(kv_score.shape[0], positions.numel())
    if num_tokens == 0:
        return
    has_prev_rows = prev_rows is not None
    prev_row_stride = prev_rows.stride(0) if prev_rows is not None else 1
    grid = num_reqs + triton.cdiv(num_tokens, compress_ratio)
    _fused_save_compress_norm_kernel[(grid,)](
        kv_score,
        positions,
        state_cache,
        slot_mapping,
        query_start_loc,
        token_to_req_indices,
        rms_norm_weight,
        latent_out,
        prev_rows if has_prev_rows else kv_score,
        prev_row_indices if has_prev_rows else slot_mapping,
        num_tokens,
        num_reqs,
        RAW_STRIDE=kv_score.stride(0),
        PREV_ROW_STRIDE=prev_row_stride,
        STATE_STRIDE=state_stride,
        STATE_ROW_STRIDE=state_row_stride,
        STATE_BLOCK=state_block,
        COMPRESS_RATIO=compress_ratio,
        EPS=rms_norm_eps,
        JOIN_ROW_PTRS=_JOIN_ROW_PTRS,
        HAS_PREV_ROWS=has_prev_rows,
        num_warps=4,
        **({"launch_pdl": False} if current_platform.is_cuda() else {}),
    )

rope_quant_insert(latent, positions, cos_sin_cache, kv_cache, slot_mapping, compress_ratio, fp8_scale=None)

Apply GPT-J RoPE and publish a latent to the compressed KV cache.

The BF16 latent supplies both NoPE quantization and RoPE input. It is read only for valid slots at group boundaries. The cache dtype selects the layout: uint8 is a paged FlashMLA layout, whose record the per-token byte width names -- 584 B for V4 (576 value bytes and eight segregated UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of 32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32 e4m3 scales of 16 dims each), which only the compressed cache uses. bfloat16 and float8_e4m3fn are the plain [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the per-tensor fp8_scale.

Source code in vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py

def rope_quant_insert(
    latent: torch.Tensor,
    positions: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    compress_ratio: int,
    fp8_scale: torch.Tensor | None = None,
) -> None:
    """Apply GPT-J RoPE and publish a latent to the compressed KV cache.

    The BF16 latent supplies both NoPE quantization and RoPE input. It is read
    only for valid slots at group boundaries. The cache dtype selects the
    layout: ``uint8`` is a paged FlashMLA layout, whose record the per-token
    byte width names -- 584 B for V4 (576 value bytes and eight segregated
    UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1
    (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of
    32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32
    e4m3 scales of 16 dims each), which only the compressed cache uses.
    ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows
    read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``.
    """
    assert compress_ratio in (1, 2)
    assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16
    assert latent.is_contiguous()
    num_tokens = slot_mapping.numel()
    assert num_tokens <= min(latent.shape[0], positions.numel())
    if num_tokens == 0:
        return
    launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
    if kv_cache.dtype == torch.uint8:
        kernel = {
            584: _rope_quant_insert_kernel,
            528: _rope_quant_insert_mxfp8_kernel,
            288: _rope_quant_insert_nvfp4_kernel,
        }.get(kv_cache.shape[-1])
        assert kernel is not None, (
            f"unsupported paged KV record width {kv_cache.shape[-1]}"
        )
        kernel[(num_tokens,)](
            latent,
            positions,
            cos_sin_cache,
            kv_cache,
            slot_mapping,
            COS_STRIDE=cos_sin_cache.stride(0),
            CACHE_STRIDE=kv_cache.stride(0),
            CACHE_BLOCK=kv_cache.shape[1],
            COMPRESS_RATIO=compress_ratio,
            SANITIZE_CACHE_NANS=_ON_GFX950,
            num_warps=4,
            **launch_kwargs,
        )
        return

    assert kv_cache.dtype in (torch.bfloat16, torch.float8_e4m3fn)
    assert kv_cache.shape[-1] == 512 and kv_cache.stride(-1) == 1
    store_fp8 = kv_cache.dtype == torch.float8_e4m3fn
    if store_fp8:
        assert fp8_scale is not None and fp8_scale.numel() == 1
        assert fp8_scale.dtype == torch.float32
    _rope_plain_insert_kernel[(num_tokens,)](
        latent,
        positions,
        cos_sin_cache,
        kv_cache,
        slot_mapping,
        fp8_scale if store_fp8 else None,
        COS_STRIDE=cos_sin_cache.stride(0),
        CACHE_STRIDE=kv_cache.stride(0),
        ROW_STRIDE=kv_cache.stride(1),
        CACHE_BLOCK=kv_cache.shape[1],
        COMPRESS_RATIO=compress_ratio,
        STORE_FP8=store_fp8,
        num_warps=4,
        **launch_kwargs,
    )

save_ring_rows(rows, slot_mapping, state_cache)

Store FP32 [kv, score] rows to their ring slots, skipping slot -1.

Source code in vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py

def save_ring_rows(
    rows: torch.Tensor, slot_mapping: torch.Tensor, state_cache: torch.Tensor
) -> None:
    """Store FP32 [kv, score] rows to their ring slots, skipping slot -1."""
    assert rows.dtype == torch.float32 and rows.shape[1] == 1024
    assert rows.stride(1) == 1 and state_cache.stride(2) == 1
    assert slot_mapping.numel() == rows.shape[0]
    if rows.shape[0] == 0:
        return
    _save_ring_rows_kernel[(rows.shape[0],)](
        rows,
        slot_mapping,
        state_cache,
        ROW_STRIDE=rows.stride(0),
        STATE_STRIDE=state_cache.stride(0),
        STATE_ROW_STRIDE=state_cache.stride(1),
        STATE_BLOCK=state_cache.shape[1],
        num_warps=4,
    )