Skip to content

vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.dsa.attn_ds

fp8_ds_mla sparse core of a K3.1 DSA layer (FLASHMLA_SPARSE), as the stock decode: the BF16 query rows [B, 64 (padded heads), 576], the logical -> physical top-k and FlashMLA sparse decode over the packed fp8_ds_mla blocks; under DCP, FlashMLA over this rank's compacted slots.

Functions:

  • flashmla_core –

    As the stock decode: the packed blocks as 64-row pages of flat rows, the

  • flashmla_dcp_core –

    DCP: FlashMLA sparse decode of every rank's BF16 query heads (the gathered q

flashmla_core(impl, query, cache, topk, metadata)

As the stock decode: the packed blocks as 64-row pages of flat rows, the logical top-k -> physical, FlashMLA sparse decode -> [B, 64, 512].

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/dsa/attn_ds.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def flashmla_core(impl, query, cache, topk, metadata):
    """As the stock decode: the packed blocks as 64-row pages of flat rows, the
    logical top-k -> physical, FlashMLA sparse decode -> [B, 64, 512]."""
    from vllm.v1.attention.backends.mla.flashmla_sparse import _fp8_kv_pages
    from vllm.v1.attention.ops.flashmla import flash_mla_with_kvcache

    pages, stride_rows = _fp8_kv_pages(cache, metadata.block_size)
    impl.record_logical_topk_ready()
    physical = impl._convert_logical_to_physical_topk(
        topk,
        metadata,
        block_stride_rows=stride_rows,
        return_valid_counts=False,
    )
    kernel_metadata = metadata.fp8_extra_metadata
    latent, _ = flash_mla_with_kvcache(
        q=query.unsqueeze(0),
        k_cache=pages.view(torch.uint8).unsqueeze(-2),
        block_table=kernel_metadata.dummy_block_table,
        head_dim_v=512,
        cache_seqlens=kernel_metadata.cache_lens,
        tile_scheduler_metadata=kernel_metadata.scheduler_metadata,
        is_fp8_kvcache=True,
        indices=physical.unsqueeze(0),
        softmax_scale=impl.softmax_scale,
    )
    return latent[0]

flashmla_dcp_core(impl, query, pages, slots, valid, out, table, cache_lens)

DCP: FlashMLA sparse decode of every rank's BF16 query heads (the gathered q mailbox) over this rank's compacted slots (natural-log LSE; a fresh scheduler metadata per call follows the valid lengths) into out, its LSE and valid (DcpValue drops the rows this rank owns none of).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/dsa/attn_ds.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def flashmla_dcp_core(impl, query, pages, slots, valid, out, table, cache_lens):
    """DCP: FlashMLA sparse decode of every rank's BF16 query heads (the gathered q
    mailbox) over this rank's compacted ``slots`` (natural-log LSE; a fresh
    scheduler metadata per call follows the ``valid`` lengths) into ``out``, its
    LSE and ``valid`` (DcpValue drops the rows this rank owns none of)."""
    from vllm.v1.attention.ops.flashmla import (
        flash_mla_with_kvcache,
        get_mla_metadata,
    )

    rows, lse = flash_mla_with_kvcache(
        q=query.unsqueeze(1),
        k_cache=pages.view(torch.uint8).unsqueeze(-2),
        block_table=table,
        head_dim_v=512,
        cache_seqlens=cache_lens,
        tile_scheduler_metadata=get_mla_metadata()[0],
        is_fp8_kvcache=True,
        indices=slots.unsqueeze(1),
        topk_length=valid,
        softmax_scale=impl.softmax_scale,
        out=out,
    )
    return rows, lse, valid