Skip to content

vllm.models.kimi_k3.amd.specialized_decode.runtime.runner

The specialized decode's per-rank state: one scratch buffer every launch shares (their mailbox tags differ), the device step epoch, the health words and the TP group's peer memory.

Mailbox tags carry the epoch the model bumps once a forward (step_begin, which a CUDA graph captures), so no buffer is zeroed between steps.

The peer buffer's first PEER_HEAD bytes hold the output rank's flags, a word a rank (alarm): a rank whose wait gives up sets its word there before it stores anything else. The kernels see the buffer past the head.

Classes:

  • MonoRunner –

    Collective: every TP rank builds it at the same point.

MonoRunner

Collective: every TP rank builds it at the same point.

Methods:

  • __init__ –

    dcp: M1's ranks, 1 or the TP group (default: the DCP group's).

  • step_begin –

    Once a forward, before its first mono launch (a graph captures it).

Attributes:

  • m1c (bool) –

    DCP1 at qlen 4 / 8 (M1b in K2): M1a writes a bf16 query, K2's launch

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py 
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
class MonoRunner:
    """Collective: every TP rank builds it at the same point."""

    def __init__(self, device: torch.device, dcp: int | None = None, m1: bool = True):
        """``dcp``: M1's ranks, 1 or the TP group (default: the DCP group's).
        ``m1``: the MLA layers run M1 (else vLLM's attention ops)."""
        from vllm.models.kimi_k3.amd.specialized_decode.mono import layer
        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import (
            kda,
            mla,
            mla_back,
        )

        self.scratch = torch.zeros(
            max(
                kda.scratch_bytes(
                    kda.KdaPreBuild(
                        tokens=MAX_TOKENS,
                        qlen=MAX_TOKENS,
                        nblocks=8,
                        delta=True,
                        state_len=3 + MAX_TOKENS,
                    )
                ),
                layer.scratch_bytes(layer.K2Build(tokens=MAX_TOKENS, nblocks=8)),
                mla.scratch_bytes(
                    mla.MlaPreBuild(tokens=MAX_TOKENS, nblocks=8, delta=True)
                ),
                # K2's tagged x rows (layer.OPT_TAG_X), past every layout
                layer.tag_end(),
            ),
            dtype=torch.uint8,
            device=device,
        )
        layer.TAG_READY = True
        self.epoch = torch.zeros(1, dtype=torch.int32, device=device)
        self.health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device=device)

        tp = get_tp_group()
        self.rank = tp.rank_in_group
        self.world = tp.world_size
        self.dcp = get_dcp_group().world_size if dcp is None else dcp
        # K2's regions, M1a's gathered query (QG), M1b's partials (CMB); M1
        # takes no peer memory at DCP1
        self.m1 = m1
        m1 = m1 and self.dcp > 1
        self.mla_peer_off = layer.peer_bytes(MAX_TOKENS)
        self.mla_post_off = self.mla_peer_off + (
            mla.peer_bytes(MAX_TOKENS, self.dcp) if m1 else 0
        )
        nbytes = (
            PEER_HEAD
            + self.mla_post_off
            + (mla_back.peer_bytes(MAX_TOKENS, self.dcp) if m1 else 0)
        )
        # K2's sentinel regions (layer.OPT_SENTS, all ones, both layer
        # parities) and RLOGP (layer.OPT_RLOG_PARTS), past every other region
        sent_lo = PEER_HEAD + layer.sent_offset()
        nbytes = max(nbytes, PEER_HEAD + layer.extra_peer_end())
        self.peers = PeerBuffer(nbytes, tp.cpu_group, self.rank, tp.world_size, device)
        self.peers.bytes.zero_()
        self.peers.bytes[sent_lo : sent_lo + layer.sent_bytes()].fill_(0xFF)
        torch.accelerator.synchronize()
        torch.distributed.barrier(group=tp.cpu_group)
        layer.SENT_READY = True
        self.peer_table = self.peers.addresses + PEER_HEAD
        self.alarm = int(self.peers.addresses[0].item()) + 4 * self.rank
        self.flags = self.peers.bytes[: 4 * tp.world_size].view(torch.int32)
        self.fault: dict = {}  # K2Build's fault_* fields (tests only)
        self.fuse_m1b = True  # an MLA layer's M1b runs in its K2's launch
        self.first_layer = -1  # the step's first kernel layer: bumps the epoch
        self.kda_prefix = ""  # a KDA layer's name: the step's metadata
        self.mla_prefix = ""  # an M1 layer's name: its decode metadata
        self.interleave = 1  # the DCP KV cache interleave
        self.num_spec = 0  # speculative tokens a request

    @property
    def m1c(self) -> bool:
        """DCP1 at qlen 4 / 8 (M1b in K2): M1a writes a bf16 query, K2's launch
        runs the decode attention (M1c) first, its M1b merging the KV splits."""
        return self.dcp == 1 and self.num_spec + 1 in M1C_QLENS and self.fuse_m1b

    def step_begin(self) -> None:
        """Once a forward, before its first mono launch (a graph captures it)."""
        self.epoch.add_(1)

    def launch_args(self) -> dict:
        return {
            "scratch": self.scratch,
            "epoch": self.epoch,
            "health": self.health,
            "alarm": self.alarm,
        }

    def peer_args(self) -> dict:
        return {"peers": self.peer_table, "rank": self.rank}

m1c property

DCP1 at qlen 4 / 8 (M1b in K2): M1a writes a bf16 query, K2's launch runs the decode attention (M1c) first, its M1b merging the KV splits.

__init__(device, dcp=None, m1=True)

dcp: M1's ranks, 1 or the TP group (default: the DCP group's). m1: the MLA layers run M1 (else vLLM's attention ops).

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def __init__(self, device: torch.device, dcp: int | None = None, m1: bool = True):
    """``dcp``: M1's ranks, 1 or the TP group (default: the DCP group's).
    ``m1``: the MLA layers run M1 (else vLLM's attention ops)."""
    from vllm.models.kimi_k3.amd.specialized_decode.mono import layer
    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import (
        kda,
        mla,
        mla_back,
    )

    self.scratch = torch.zeros(
        max(
            kda.scratch_bytes(
                kda.KdaPreBuild(
                    tokens=MAX_TOKENS,
                    qlen=MAX_TOKENS,
                    nblocks=8,
                    delta=True,
                    state_len=3 + MAX_TOKENS,
                )
            ),
            layer.scratch_bytes(layer.K2Build(tokens=MAX_TOKENS, nblocks=8)),
            mla.scratch_bytes(
                mla.MlaPreBuild(tokens=MAX_TOKENS, nblocks=8, delta=True)
            ),
            # K2's tagged x rows (layer.OPT_TAG_X), past every layout
            layer.tag_end(),
        ),
        dtype=torch.uint8,
        device=device,
    )
    layer.TAG_READY = True
    self.epoch = torch.zeros(1, dtype=torch.int32, device=device)
    self.health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device=device)

    tp = get_tp_group()
    self.rank = tp.rank_in_group
    self.world = tp.world_size
    self.dcp = get_dcp_group().world_size if dcp is None else dcp
    # K2's regions, M1a's gathered query (QG), M1b's partials (CMB); M1
    # takes no peer memory at DCP1
    self.m1 = m1
    m1 = m1 and self.dcp > 1
    self.mla_peer_off = layer.peer_bytes(MAX_TOKENS)
    self.mla_post_off = self.mla_peer_off + (
        mla.peer_bytes(MAX_TOKENS, self.dcp) if m1 else 0
    )
    nbytes = (
        PEER_HEAD
        + self.mla_post_off
        + (mla_back.peer_bytes(MAX_TOKENS, self.dcp) if m1 else 0)
    )
    # K2's sentinel regions (layer.OPT_SENTS, all ones, both layer
    # parities) and RLOGP (layer.OPT_RLOG_PARTS), past every other region
    sent_lo = PEER_HEAD + layer.sent_offset()
    nbytes = max(nbytes, PEER_HEAD + layer.extra_peer_end())
    self.peers = PeerBuffer(nbytes, tp.cpu_group, self.rank, tp.world_size, device)
    self.peers.bytes.zero_()
    self.peers.bytes[sent_lo : sent_lo + layer.sent_bytes()].fill_(0xFF)
    torch.accelerator.synchronize()
    torch.distributed.barrier(group=tp.cpu_group)
    layer.SENT_READY = True
    self.peer_table = self.peers.addresses + PEER_HEAD
    self.alarm = int(self.peers.addresses[0].item()) + 4 * self.rank
    self.flags = self.peers.bytes[: 4 * tp.world_size].view(torch.int32)
    self.fault: dict = {}  # K2Build's fault_* fields (tests only)
    self.fuse_m1b = True  # an MLA layer's M1b runs in its K2's launch
    self.first_layer = -1  # the step's first kernel layer: bumps the epoch
    self.kda_prefix = ""  # a KDA layer's name: the step's metadata
    self.mla_prefix = ""  # an M1 layer's name: its decode metadata
    self.interleave = 1  # the DCP KV cache interleave
    self.num_spec = 0  # speculative tokens a request

step_begin()

Once a forward, before its first mono launch (a graph captures it).

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py 
106
107
108
def step_begin(self) -> None:
    """Once a forward, before its first mono launch (a graph captures it)."""
    self.epoch.add_(1)