vllm.models.deepseek_v4.xpu.xpu_sparse_decode_fp8

XPU sparse decode for DeepSeek V4 with FP8 KV cache.

Strategy: dequantize FP8 UE8M0 KV cache pages to BF16 on the fly, then reuse the BF16 sparse MLA attention kernel (xpu_sparse_mla_bf16). This keeps the external KV cache layout identical to CUDA/ROCm.

Functions:

• dequant_gather_slots –

  Dequantize FP8 UE8M0 pages at scattered slot indices into bf16.

• xpu_sparse_decode_fp8 –

  XPU decode: dequant FP8 pages to BF16, then BF16 sparse MLA attention.

_dequant_gather_slots_kernel(out_ptr, cache_ptr, indices_ptr, cache_block_size, token_data_size, block_stride, fp8_dim, bf16_dim, scale_dim, quant_block, output_dim, n_quant_blocks)

Dequantize scattered FP8 slots into a flat BF16 workspace.

Grid: [total_slots] — one program per slot to gather.

Cache block layout (block_size tokens): [0, block_size576): Token data, each token 448 FP8 + 128 BF16 [block_size576, block_size576 + block_size8): Scales

Source code in vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py

@triton.jit
def _dequant_gather_slots_kernel(
    # Output workspace: [total_slots, OUTPUT_DIM] bf16
    out_ptr,
    # FP8 paged cache base pointer (uint8 flat)
    cache_ptr,
    # Global slot indices: [total_slots] int32
    indices_ptr,
    # Cache geometry
    cache_block_size: tl.constexpr,
    token_data_size: tl.constexpr,  # 576
    block_stride: tl.int64,  # k_cache.stride(0) — total uint8 per block
    fp8_dim: tl.constexpr,  # 448
    bf16_dim: tl.constexpr,  # 64
    scale_dim: tl.constexpr,  # 8
    quant_block: tl.constexpr,  # 64
    output_dim: tl.constexpr,  # 512
    n_quant_blocks: tl.constexpr,  # 7
):
    """Dequantize scattered FP8 slots into a flat BF16 workspace.

    Grid: [total_slots] — one program per slot to gather.

    Cache block layout (block_size tokens):
      [0, block_size*576): Token data, each token 448 FP8 + 128 BF16
      [block_size*576, block_size*576 + block_size*8): Scales
    """
    pid = tl.program_id(0)

    # Load global slot index
    slot_idx = tl.load(indices_ptr + pid).to(tl.int64)

    # Output pointer for this slot
    out_row_ptr = out_ptr + pid * output_dim

    # Handle invalid slots (index < 0): write zeros
    if slot_idx < 0:
        zero = tl.zeros([quant_block], dtype=tl.bfloat16)
        for i in tl.static_range(0, 512, 64):
            offsets = i + tl.arange(0, quant_block)
            mask = offsets < output_dim
            tl.store(out_row_ptr + offsets, zero, mask=mask)
        return

    # Compute block and position within block
    block_idx = slot_idx // cache_block_size
    pos_in_block = slot_idx % cache_block_size

    # Block base pointer
    block_base = cache_ptr + block_idx * block_stride

    # Token data: at offset pos_in_block * token_data_size within block
    token_data_ptr = block_base + pos_in_block * token_data_size

    # Scale: after all token data, at offset
    # cache_block_size * token_data_size + pos_in_block * scale_dim
    scale_region_offset = tl.cast(cache_block_size, tl.int64) * token_data_size
    token_scale_ptr = block_base + scale_region_offset + pos_in_block * scale_dim

    # ========== Dequantize FP8 portion (448 elements) ==========
    for qblock_idx in tl.static_range(n_quant_blocks):
        qblock_start = qblock_idx * quant_block
        offsets = qblock_start + tl.arange(0, quant_block)
        mask = offsets < fp8_dim

        # Load FP8 as uint8 and bitcast
        x_uint8 = tl.load(token_data_ptr + offsets, mask=mask, other=0)
        x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
        x_float = x_fp8.to(tl.float32)

        # Load UE8M0 scale: scale = 2^(stored_value - 127)
        encoded_scale = tl.load(token_scale_ptr + qblock_idx)
        exponent = encoded_scale.to(tl.float32) - 127.0
        scale = tl.exp2(exponent)

        # Dequantize and store as bf16
        x_dequant = x_float * scale
        tl.store(out_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask)

    # ========== Copy BF16 portion (64 elements) directly ==========
    bf16_src_ptr = (token_data_ptr + fp8_dim).to(tl.pointer_type(tl.bfloat16))
    bf16_out_ptr = (out_row_ptr + fp8_dim).to(tl.pointer_type(tl.bfloat16))

    for j in tl.static_range(bf16_dim // 16):
        chunk_offsets = j * 16 + tl.arange(0, 16)
        bf16_vals = tl.load(bf16_src_ptr + chunk_offsets)
        tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)

dequant_gather_slots(out, cache, indices, cache_block_size)

Dequantize FP8 UE8M0 pages at scattered slot indices into bf16.

Source code in vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py

def dequant_gather_slots(
    out: torch.Tensor,  # [total_slots, 512] bf16, pre-allocated
    cache: torch.Tensor,  # [num_blocks, block_size, head_bytes] uint8
    indices: torch.Tensor,  # [total_slots] int32, global slot IDs
    cache_block_size: int,  # block_size for this cache
) -> None:
    """Dequantize FP8 UE8M0 pages at scattered slot indices into bf16."""
    total_slots = indices.shape[0]
    if total_slots == 0:
        return

    block_stride = cache.stride(0)

    _dequant_gather_slots_kernel[(total_slots,)](
        out,
        cache,
        indices,
        cache_block_size=cache_block_size,
        token_data_size=TOKEN_DATA_SIZE,
        block_stride=block_stride,
        fp8_dim=TOKEN_FP8_DIM,
        bf16_dim=TOKEN_BF16_DIM,
        scale_dim=TOKEN_SCALE_DIM,
        quant_block=QUANT_BLOCK_SIZE,
        output_dim=OUTPUT_DIM,
        n_quant_blocks=7,
    )

xpu_sparse_decode_fp8(q, kv_cache, swa_kv_cache, swa_only, topk_indices, topk_lens, swa_indices, swa_lens, attn_sink, softmax_scale, head_dim, nope_head_dim, rope_head_dim, out)

XPU decode: dequant FP8 pages to BF16, then BF16 sparse MLA attention.

Keeps external FP8 KV cache layout identical to CUDA/ROCm. Performance is slower due to on-the-fly dequant, but correctness is guaranteed by reusing the validated BF16 attention kernel.

Source code in vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py

def xpu_sparse_decode_fp8(
    q: torch.Tensor,  # [num_tokens, num_heads, head_dim]
    kv_cache: torch.Tensor | None,  # [num_blocks, block_size, head_bytes] uint8
    swa_kv_cache: torch.Tensor,  # [num_blocks, swa_block_size, head_bytes] uint8
    swa_only: bool,
    topk_indices: torch.Tensor | None,  # [num_tokens, 1, topk] global slot IDs
    topk_lens: torch.Tensor | None,
    swa_indices: torch.Tensor,  # [num_tokens, 1, swa_k] global slot IDs
    swa_lens: torch.Tensor,
    attn_sink: torch.Tensor,
    softmax_scale: float,
    head_dim: int,
    nope_head_dim: int,
    rope_head_dim: int,
    out: torch.Tensor,  # [num_tokens, num_heads, head_dim]
) -> None:
    """XPU decode: dequant FP8 pages to BF16, then BF16 sparse MLA attention.

    Keeps external FP8 KV cache layout identical to CUDA/ROCm.
    Performance is slower due to on-the-fly dequant, but correctness is
    guaranteed by reusing the validated BF16 attention kernel.
    """
    num_tokens = q.shape[0]
    device = q.device

    if not swa_only and topk_indices is not None and kv_cache is not None:
        topk_idx_2d = (
            topk_indices.squeeze(1) if topk_indices.dim() == 3 else topk_indices
        )
        max_topk = topk_idx_2d.shape[1]
    else:
        topk_idx_2d = None
        max_topk = 0

    use_topk = topk_idx_2d is not None
    swa_idx_2d = swa_indices.squeeze(1) if swa_indices.dim() == 3 else swa_indices
    max_swa = swa_idx_2d.shape[1]
    K_total = max_topk + max_swa

    # Topk rows precede SWA rows in the flat workspace.
    n_topk_rows = num_tokens * max_topk
    workspace = torch.empty(
        (num_tokens * K_total, OUTPUT_DIM), dtype=torch.bfloat16, device=device
    )

    if topk_idx_2d is not None and kv_cache is not None and n_topk_rows > 0:
        dequant_gather_slots(
            workspace[:n_topk_rows],
            kv_cache,
            topk_idx_2d.reshape(-1).to(torch.int32),
            kv_cache.shape[1],
        )

    dequant_gather_slots(
        workspace[n_topk_rows:],
        swa_kv_cache,
        swa_idx_2d.reshape(-1).to(torch.int32),
        swa_kv_cache.shape[1],
    )

    # Use a shape-derived width so replay accepts changing lens without host sync.
    _BLOCK_N = 16
    width = max(_BLOCK_N, ((K_total + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N)

    pos = torch.arange(width, device=device, dtype=torch.int32).unsqueeze(0)
    tok = torch.arange(num_tokens, device=device, dtype=torch.int32).unsqueeze(1)

    slen = swa_lens.reshape(-1, 1).to(torch.int32).clamp(min=0, max=max_swa)
    if use_topk and topk_lens is not None:
        tlen = topk_lens.reshape(-1, 1).to(torch.int32).clamp(min=0, max=max_topk)
    else:
        tlen = torch.zeros_like(slen)

    in_topk = pos < tlen
    in_swa = (pos >= tlen) & (pos < tlen + slen)

    topk_ws = tok * max_topk + pos
    swa_ws = n_topk_rows + tok * max_swa + (pos - tlen)

    combined_indices = torch.where(
        in_topk,
        topk_ws,
        torch.where(in_swa, swa_ws, torch.full_like(topk_ws, -1)),
    )

    out_attn, _, _ = triton_bf16_mla_sparse_interface(
        q=q,
        kv=workspace.unsqueeze(1),
        indices=combined_indices.unsqueeze(1),
        sm_scale=softmax_scale,
        d_v=q.shape[-1],
        block_dpe=0,
    )
    out.copy_(out_attn)