Skip to content

vllm.models.kimi_k3.amd.specialized_decode.runtime.layers

The layer kernels: each takes and returns what KimiDecoderLayer.forward does with AttnRes, (prefix_sum, block_residual, prefix_delta).

  • WholeLayer (KDA layers): K1 (AttnRes, in_proj, f_b, spec conv, the KDA recurrence, the gated norm) writes o_proj's input; K2 runs o_proj, its all-reduce, the MLP AttnRes and the latent MoE.
  • FfnLayer (MLA layers): vLLM's AttnRes and MLA up to o_proj's input, then K2.

Weights are views of the loaded parameters, taken at every launch.

Classes:

Functions:

  • compile_build –

    Compiles one launch's kernel (a K2Build or KdaPreBuild) without

WholeLayer

Bases: _Layer

Methods:

  • caches –

    The conv state as [slots, dim, state_len] (a view) and the recurrent

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
class WholeLayer(_Layer):
    def _k1_key(self, qlen: int, state_len: int, delta: bool):
        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda

        layer = self.layer
        attn = layer.self_attn
        return kda.KdaPreBuild(
            tokens=self.width,
            qlen=qlen,
            nblocks=layer.prev_valid_blocks,
            delta=delta,
            state_len=state_len,
            eps=layer.self_attention_res_norm.variance_epsilon,
            out_eps=layer.input_layernorm.variance_epsilon,
            onorm_eps=attn.o_norm.eps,
            lower_bound=attn.gate_lower_bound,
            write_idx=layer.block_write_idx if layer.is_block_write_layer else -1,
        )

    def builds(self) -> list:
        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda

        q = self.runner.num_spec + 1
        # the conv state keeps CONV_W - 1 + num_spec entries; a layer after
        # another always has a pending delta
        k1 = self._k1_key(q, kda.CONV_W - 1 + self.runner.num_spec, True)
        return [k1, *super().builds()]

    def caches(self) -> tuple[torch.Tensor, torch.Tensor]:
        """The conv state as [slots, dim, state_len] (a view) and the recurrent
        state."""
        conv_state, rstate = self.layer.self_attn.kv_cache
        if not is_conv_state_dim_first():
            conv_state = conv_state.transpose(-1, -2)
        return conv_state, rstate

    def launch(self, step: "Step", positions, prefix_sum, block_residual, prefix_delta):
        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda

        layer, runner = self.layer, self.runner
        attn = layer.self_attn
        m: Any = step.attn_metadata[attn.prefix]
        if self.first:
            runner.step_begin()
        conv_state, rstate = self.caches()
        s = prefix_sum.size(0)
        core = prefix_sum.new_empty(s, kda.PROJ)
        kda.kda_pre(
            self._k1_key(
                m.spec_state_indices_tensor.size(-1),
                conv_state.size(-1),
                prefix_delta is not None,
            ),
            prefix=prefix_sum,
            delta=prefix_delta,
            blocks=block_residual,
            ares_nw=layer.self_attention_res_norm.weight,
            ares_qk=layer.self_attention_res_proj.weight.view(kda.HIDDEN),
            in_nw=layer.input_layernorm.weight,
            w_in=attn.in_proj_qkvgfab.weight,
            w_fb=attn.f_b_proj.weight,
            conv_w=attn.conv1d.weight.view(kda.QKV, 4),
            conv_state=conv_state,
            a_log=attn.A_log,
            dt_bias=attn.dt_bias,
            on_w=attn.o_norm.weight,
            rstate=rstate,
            st_idx=m.spec_state_indices_tensor,
            num_acc=m.num_accepted_tokens,
            core_out=core,
            layer=layer.layer_idx,
            **runner.launch_args(),
        )
        if layer.is_block_write_layer:
            # the prefix restarts at the attention's output
            prefix_sum = torch.empty_like(prefix_sum)
            out = _k2(layer, runner, core, prefix_sum, block_residual, reset=True)
        else:
            out = _k2(layer, runner, core, prefix_sum, block_residual)
        return prefix_sum, block_residual, out

caches()

The conv state as [slots, dim, state_len] (a view) and the recurrent state.

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py 
146
147
148
149
150
151
152
def caches(self) -> tuple[torch.Tensor, torch.Tensor]:
    """The conv state as [slots, dim, state_len] (a view) and the recurrent
    state."""
    conv_state, rstate = self.layer.self_attn.kv_cache
    if not is_conv_state_dim_first():
        conv_state = conv_state.transpose(-1, -2)
    return conv_state, rstate

_Layer

Methods:

  • builds –

    The builds this layer launches at its width, compiled before capture.

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py 
104
105
106
107
108
109
110
111
112
113
114
115
class _Layer:
    def __init__(self, layer, runner: MonoRunner, width: int):
        self.layer = layer
        self.runner = runner
        self.width = width
        # the step's first launch bumps the epoch (a graph captures it)
        self.first = layer.layer_idx == runner.first_layer

    def builds(self) -> list:
        """The builds this layer launches at its width, compiled before capture."""
        reset = getattr(self.layer, "is_block_write_layer", False)
        return [_k2_key(self.layer, self.runner, self.width, reset)]

builds()

The builds this layer launches at its width, compiled before capture.

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py 
112
113
114
115
def builds(self) -> list:
    """The builds this layer launches at its width, compiled before capture."""
    reset = getattr(self.layer, "is_block_write_layer", False)
    return [_k2_key(self.layer, self.runner, self.width, reset)]

_k2(layer, runner, core, prefix, blocks, reset=False)

K2 from core (o_proj's input): prefix updated in place (written, when reset), the MoE output returned.

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py 
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
def _k2(layer, runner: MonoRunner, core, prefix, blocks, reset=False):
    """K2 from ``core`` (o_proj's input): ``prefix`` updated in place (written,
    when ``reset``), the MoE output returned."""
    from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2
    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN
    from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe

    m = layer.mlp
    ex = getattr(m.experts, "routed_experts", m.experts)
    w_up = m.routed_expert_up_proj.weight
    if not getattr(m.routed_output_transform, "row_sharded", False):
        w_up = w_up.narrow(0, runner.rank * moe.UP_N, moe.UP_N)
    out = torch.empty_like(prefix)
    k2.k2_launch(
        _k2_key(layer, runner, core.size(0), reset),
        core=core,
        w_o=layer.self_attn.o_proj.weight,
        prefix=prefix,
        blocks=blocks,
        ares_nw=layer.mlp_res_norm.weight,
        ares_qk=layer.mlp_res_proj.weight.view(HIDDEN),
        in_nw=layer.post_attention_layernorm.weight,
        w_gate=m.gate.weight,
        bias=m.gate.e_score_correction_bias,
        w_ld=m.routed_expert_down_proj.weight,
        w_sgu=m.shared_experts.gate_up_proj.weight,
        w_sd=m.shared_experts.down_proj.weight,
        w13=ex.w13_weight,
        w13s=ex.w13_weight_scale,
        w2=ex.w2_weight,
        w2s=ex.w2_weight_scale,
        ln_w=m.routed_expert_norm.weight,
        w_up=w_up,
        out=out,
        layer=128 + layer.layer_idx,
        **runner.launch_args(),
        **runner.peer_args(),
    )
    return out

compile_build(key)

Compiles one launch's kernel (a K2Build or KdaPreBuild) without running it: FlyDSL's compile-only mode returns before the launch.

Source code in vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def compile_build(key) -> None:
    """Compiles one launch's kernel (a ``K2Build`` or ``KdaPreBuild``) without
    running it: FlyDSL's compile-only mode returns before the launch."""
    import os

    from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2
    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda

    mod = kda if isinstance(key, kda.KdaPreBuild) else k2
    launch = mod.build(key)
    # FlyDSL's compile-only switch has no FLYDSL_ prefix
    old = os.environ.get("COMPILE_ONLY")
    os.environ["COMPILE_ONLY"] = "1"
    try:
        launch(*[0] * len(mod.ABI.names), stream=torch.cuda.current_stream())
    finally:
        if old is None:
            del os.environ["COMPILE_ONLY"]
        else:
            os.environ["COMPILE_ONLY"] = old