vllm.models.deepseek_v41.amd.ops.cand_logits

The indexer logits of the DSpark candidate positions on gfx942.

Layers 24 to 36 take their top 512 only from the 2048 candidate blocks of 8 positions that layer 20 chose for each decode row. vLLM computes such a layer's logits for every position of the row with AITER's deepgemm_fp8_paged_mqa_logits. Then topk942's gatherCandidates copies the 16384 candidate logits of each row into a compact row, and topK512Candidates takes the top 512 of that row. At 128k context the dense call scores 8 times more positions than the top-k reads.

_cand_logits_kernel computes only the candidate logits and writes the compact row directly. Position i of row r holds the logit of column 8 * cand[r][i // 8] + i % 8 and that column as its id. It holds -inf and id -1 when the block entry is -1 or the column is at or past the row's end. This is what gatherCandidates writes.

Each logit is computed with the code of AITER's gfx942 kernel (_gluon_deepgemm_fp8_paged_mqa_logits_preshuffle in aiter/ops/triton/gluon/pa_mqa_logits.py). It uses the same MFMA layout. It multiplies the queries of all heads with the column's key starting from zero, multiplies by the key's scale, applies ReLU, multiplies by the head weights, and adds the heads up with the same reduction. The sum over the heads runs inside one wave, along the heads of one column, so its order does not depend on which position a column holds. Only the address of each column's key differs from the dense kernel. So every logit is bit-identical to the dense kernel's logit of that position (tests/kernels/models/amd/test_dsv41_gfx942_cand_logits.py).

Functions:

• candidate_logits –

  Writes the compact candidate rows described in the module docstring.

• split_cache –

  The FP8 keys and the FP32 key scales of an indexer cache of shape

candidate_logits(q_fp8, kv_cache, weights, seq_lens, next_n, block_table, candidates, cand_block, out_logits, out_ids, stages_per_program=0)

Writes the compact candidate rows described in the module docstring. stages_per_program overrides STAGES_PER_PROGRAM when it is not 0.

q_fp8 is [rows, heads, dim] FP8, with row r the draft position r % next_n of request r // next_n. kv_cache is the paged indexer cache, weights [rows, heads] FP32, block_table one row of pages for each request, seq_lens the lengths as the dense call gets them, candidates [rows, blocks] int32 with -1 padding, and out_logits and out_ids contiguous [rows, blocks * cand_block].

Source code in vllm/models/deepseek_v41/amd/ops/cand_logits.py

def candidate_logits(
    q_fp8: torch.Tensor,
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
    seq_lens: torch.Tensor,
    next_n: int,
    block_table: torch.Tensor,
    candidates: torch.Tensor,
    cand_block: int,
    out_logits: torch.Tensor,
    out_ids: torch.Tensor,
    stages_per_program: int = 0,
) -> None:
    """Writes the compact candidate rows described in the module docstring.
    stages_per_program overrides STAGES_PER_PROGRAM when it is not 0.

    q_fp8 is [rows, heads, dim] FP8, with row r the draft position r % next_n
    of request r // next_n. kv_cache is the paged indexer cache, weights
    [rows, heads] FP32, block_table one row of pages for each request,
    seq_lens the lengths as the dense call gets them, candidates
    [rows, blocks] int32 with -1 padding, and out_logits and out_ids
    contiguous [rows, blocks * cand_block]."""
    rows, heads, dim = q_fp8.shape
    page_size = kv_cache.shape[1]
    row_len = candidates.shape[1] * cand_block
    assert out_logits.shape == (rows, row_len) and out_logits.is_contiguous()
    assert out_ids.shape == (rows, row_len) and out_ids.is_contiguous()
    # The kernel reads keys in the cache's groups of 16, so a page must hold
    # whole groups. Each column finds its own page, so a candidate block may
    # cross a page boundary.
    assert page_size % 16 == 0, f"page size {page_size} is not a multiple of 16"
    keys, scales = split_cache(kv_cache)
    per_program = stages_per_program or STAGES_PER_PROGRAM
    stages = triton_cdiv(row_len, COLUMNS)
    splits = triton_cdiv(stages, per_program)
    _cand_logits_kernel[(rows * splits,)](
        q_fp8,
        q_fp8.stride(0),
        q_fp8.stride(1),
        keys,
        keys.stride(0),
        scales,
        scales.stride(0),
        seq_lens,
        next_n,
        block_table,
        block_table.stride(0),
        candidates,
        candidates.stride(0),
        weights,
        weights.stride(0),
        out_logits,
        out_ids,
        row_len,
        splits,
        SeqLensIs2D=seq_lens.dim() == 2,
        StagesPerProgram=per_program,
        ChunkQ=heads,
        HiddenDim=dim,
        KVBlockSize=page_size,
        CandBlock=cand_block,
        CDNA_VERSION=get_cdna_version(),
        num_warps=4,
    )

split_cache(kv_cache)

The FP8 keys and the FP32 key scales of an indexer cache of shape [pages, page size, 1, dim + 4], split the way AITER's deepgemm_fp8_paged_mqa_logits splits it: each page holds page size x dim FP8 values and then page size FP32 scales.

Source code in vllm/models/deepseek_v41/amd/ops/cand_logits.py

def split_cache(kv_cache: torch.Tensor):
    """The FP8 keys and the FP32 key scales of an indexer cache of shape
    [pages, page size, 1, dim + 4], split the way AITER's
    deepgemm_fp8_paged_mqa_logits splits it: each page holds page size x dim
    FP8 values and then page size FP32 scales."""
    from aiter import dtypes

    pages, page_size, _, index_dim = kv_cache.shape
    dim = index_dim - 4
    flat = kv_cache.view(-1, page_size * index_dim)
    keys = flat[..., : page_size * dim].view(dtypes.fp8)
    scales = flat[..., page_size * dim :].view(torch.float32)
    return keys, scales