vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache ¶

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,
) -> 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.
    """
    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
    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
    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,
        num_tokens,
        num_reqs,
        RAW_STRIDE=kv_score.stride(0),
        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,
        num_warps=4,
        **({"launch_pdl": False} if current_platform.is_cuda() else {}),
    )

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 the fp8_ds_mla paged layout (576 value bytes and eight
    segregated UE8M0 scale bytes per token, including one zero padding scale);
    ``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:
        assert kv_cache.shape[-1] == 584
        _rope_quant_insert_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,
    )