@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7f9173dfa2

[ROCm][Perf] Kimi-K3 Fused kernel for KDA decode (#50654)

Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

kliuae committed Aug 12, 2026 at 11:39 UTC 7f9173dfa24d70846361d867dfe32f7440875460
6 files changed +1592 -9
CMakeLists.txt
+22
@@ -1164,6 +1164,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1164 # if CUDA endif
1165 endif()
1166
1167 + # HIP counterpart of the FUSED_KDA_DECODE block above: same `fused_kda_decode`
1168 + # op, same VLLM_ENABLE_FUSED_KDA_DECODE guard in torch_bindings.cpp. Only
1169 + # measured on gfx950, so it is only built when gfx950 is in the arch list.
1170 + # A multi-arch build that includes gfx950 still emits this source for every
1171 + # arch in that list. The runtime gate in kda_decode.py keeps the kernel off
1172 + # non-gfx950 devices.
1173 + if(VLLM_GPU_LANG STREQUAL "HIP")
1174 + set(FUSED_KDA_DECODE_HIP_ARCHS ${VLLM_GPU_ARCHES})
1175 + list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx950")
1176 + if(FUSED_KDA_DECODE_HIP_ARCHS)
1177 + set(FUSED_KDA_DECODE_HIP_SRC
1178 + "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu")
1179 + list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_HIP_SRC}")
1180 + message(STATUS
1181 + "Building fused KDA decode for HIP archs: ${FUSED_KDA_DECODE_HIP_ARCHS}")
1182 + endif()
1183 + endif()
1184 +
1185 message(STATUS "Enabling C_stable extension.")
1186 list(REMOVE_DUPLICATES VLLM_STABLE_EXT_SRC)
1187 define_extension_target(
@@ -1204,6 +1222,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1222 CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
1223 elseif(VLLM_GPU_LANG STREQUAL "HIP")
1224 target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
1225 + if(FUSED_KDA_DECODE_HIP_SRC)
1226 + target_compile_definitions(_C_stable_libtorch PRIVATE
1227 + VLLM_ENABLE_FUSED_KDA_DECODE=1)
1228 + endif()
1229 endif()
1230
1231 # On ROCm, _C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in
benchmarks/kernels/benchmark_kimi_k3_kda_decode.py new
+291
@@ -0,0 +1,291 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Microbenchmark for the fused Kimi-K3 KDA decode kernel on ROCm.
4 +
5 +Compares the single fused launch against the three-kernel Triton chain it
6 +replaces (packed causal conv1d update -> recurrent delta-rule decode -> gated
7 +output RMSNorm), at Kimi-K3 KDA shapes: 96 heads x 128, conv width 4,
8 +gate_lower_bound -5.0. Per-step figures scale the per-layer number by the 69
9 +KDA layers in the model.
10 +
11 +The recurrent state dominates the traffic (num_heads x 128 x 128 fp32 read and
12 +written per token), so the reported bandwidth is the useful metric: the fusion
13 +removes launches and the intermediate QKV / core-output round trips, not the
14 +state traffic itself.
15 +
16 +Example:
17 + python benchmarks/kernels/benchmark_kimi_k3_kda_decode.py \
18 + --tokens 1 8 32 64 128 --heads 12
19 +"""
20 +
21 +import argparse
22 +import functools
23 +
24 +import torch
25 +
26 +from vllm.triton_utils import triton
27 +
28 +HEAD_DIM = 128
29 +CONV_WIDTH = 4
30 +GATE_LOWER_BOUND = -5.0
31 +NORM_EPS = 1e-5
32 +NUM_KDA_LAYERS = 69
33 +DTYPE = torch.bfloat16
34 +
35 +
36 +def _bench(fn) -> float:
37 + return triton.testing.do_bench(fn, warmup=50, rep=300, return_mode="median")
38 +
39 +
40 +def _bench_graph_layers(calls: list) -> float:
41 + """Per-layer milliseconds for a graph holding one call per KDA layer."""
42 + stream = torch.cuda.Stream()
43 + stream.wait_stream(torch.cuda.current_stream())
44 + with torch.cuda.stream(stream):
45 + for call in calls[:3]:
46 + call()
47 + torch.cuda.current_stream().wait_stream(stream)
48 + graph = torch.cuda.CUDAGraph()
49 + with torch.cuda.graph(graph):
50 + for call in calls:
51 + call()
52 + total = triton.testing.do_bench(
53 + graph.replay, warmup=50, rep=300, return_mode="median"
54 + )
55 + return total / len(calls)
56 +
57 +
58 +def _bench_graph(fn, repeats: int = NUM_KDA_LAYERS) -> float:
59 + """Per-call milliseconds under CUDA-graph replay, as decode actually runs.
60 +
61 + Eager timings credit a fusion for every Python dispatch it removes; inside a
62 + captured graph those are gone, so this is the number that decides whether
63 + the kernel is worth it in the server. The graph holds `repeats` calls (one
64 + per KDA layer) because a one-call graph is swamped by the ~19 us HIP
65 + graph-launch overhead, which a real 93-layer graph amortises away.
66 + """
67 + stream = torch.cuda.Stream()
68 + stream.wait_stream(torch.cuda.current_stream())
69 + with torch.cuda.stream(stream):
70 + for _ in range(3):
71 + fn()
72 + torch.cuda.current_stream().wait_stream(stream)
73 + graph = torch.cuda.CUDAGraph()
74 + with torch.cuda.graph(graph):
75 + for _ in range(repeats):
76 + fn()
77 + total = triton.testing.do_bench(
78 + graph.replay, warmup=50, rep=300, return_mode="median"
79 + )
80 + return total / repeats
81 +
82 +
83 +class Inputs:
84 + def __init__(self, num_tokens: int, num_heads: int) -> None:
85 + torch.manual_seed(0)
86 + device = "cuda"
87 + dim = num_heads * HEAD_DIM
88 + num_slots = num_tokens + 8
89 + self.num_heads = num_heads
90 + self.num_tokens = num_tokens
91 + self.mixed_qkv = torch.randn(num_tokens, 3 * dim, device=device, dtype=DTYPE)
92 + self.conv_weights = torch.randn(
93 + 3 * dim, CONV_WIDTH, device=device, dtype=torch.float32
94 + )
95 + self.decode_conv1d_weight = torch.stack(
96 + [
97 + self.conv_weights[i * dim : (i + 1) * dim].transpose(0, 1).contiguous()
98 + for i in range(3)
99 + ]
100 + )
101 + self.conv_state = torch.randn(
102 + num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE
103 + )
104 + self.recurrent_state = torch.randn(
105 + num_slots, num_heads, HEAD_DIM, HEAD_DIM, device=device, dtype=torch.float32
106 + )
107 + self.g1 = torch.randn(
108 + 1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE
109 + )
110 + self.g2 = torch.randn(
111 + num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE
112 + )
113 + self.beta = torch.randn(1, num_tokens, num_heads, device=device, dtype=DTYPE)
114 + self.A_log = torch.randn(num_heads, device=device, dtype=torch.float32)
115 + self.dt_bias = torch.randn(dim, device=device, dtype=torch.float32)
116 + self.norm_weight_bf16 = torch.ones(HEAD_DIM, device=device, dtype=DTYPE)
117 + self.decode_norm_weight = self.norm_weight_bf16.float()
118 + # Slots start at 1: slot 0 is NULL_BLOCK_ID, which the fused kernel
119 + # treats as a padded row and skips, so timing it measures nothing.
120 + self.state_indices = torch.arange(
121 + 1, num_tokens + 1, device=device, dtype=torch.int32
122 + )
123 + self.conv_state_t = self.conv_state.transpose(-1, -2)
124 + self.out = torch.empty(
125 + 1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE
126 + )
127 + self.conv_out = torch.empty_like(self.mixed_qkv)
128 +
129 + def state_bytes(self) -> int:
130 + """Recurrent-state traffic, identical for both implementations."""
131 + return self.num_tokens * self.num_heads * HEAD_DIM * HEAD_DIM * 4 * 2
132 +
133 +
134 +def _gated_rmsnorm(
135 + x: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float
136 +) -> torch.Tensor:
137 + x_float = x.float()
138 + variance = x_float.pow(2).mean(dim=-1, keepdim=True)
139 + normed = x_float * torch.rsqrt(variance + eps) * weight.float()
140 + return (normed * torch.sigmoid(gate.float())).to(x.dtype)
141 +
142 +
143 +def _triton_gated_norm(inp: Inputs, core_attn_out: torch.Tensor) -> torch.Tensor:
144 + from vllm.third_party.flash_linear_attention.ops.kda import rms_norm_gated
145 +
146 + return rms_norm_gated(
147 + core_attn_out,
148 + inp.g2,
149 + inp.norm_weight_bf16,
150 + None,
151 + activation="sigmoid",
152 + eps=NORM_EPS,
153 + )
154 +
155 +
156 +def triton_chain(inp: Inputs, fused_norm: bool = False) -> None:
157 + from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
158 + causal_conv1d_update,
159 + )
160 + from vllm.models.kimi_k3.amd.ops.third_party.kda import (
161 + fused_recurrent_kda_packed_decode,
162 + )
163 +
164 + causal_conv1d_update(
165 + inp.mixed_qkv,
166 + inp.conv_state_t,
167 + inp.conv_weights,
168 + None,
169 + activation="silu",
170 + conv_state_indices=inp.state_indices,
171 + validate_data=False,
172 + out=inp.conv_out,
173 + )
174 + core_attn_out, _ = fused_recurrent_kda_packed_decode(
175 + mixed_qkv=inp.conv_out,
176 + raw_g=inp.g1,
177 + raw_beta=inp.beta,
178 + A_log=inp.A_log,
179 + dt_bias=inp.dt_bias,
180 + lower_bound=GATE_LOWER_BOUND,
181 + initial_state=inp.recurrent_state,
182 + state_indices=inp.state_indices,
183 + )
184 + if fused_norm:
185 + inp.out.copy_(_triton_gated_norm(inp, core_attn_out))
186 + else:
187 + inp.out.copy_(
188 + _gated_rmsnorm(core_attn_out, inp.g2, inp.norm_weight_bf16, NORM_EPS)
189 + )
190 +
191 +
192 +def fused(inp: Inputs) -> None:
193 + from vllm import _custom_ops as ops
194 +
195 + ops.fused_kda_decode(
196 + x=inp.mixed_qkv,
197 + weight=inp.decode_conv1d_weight,
198 + bias=None,
199 + conv_state=inp.conv_state_t,
200 + raw_g=inp.g1,
201 + raw_beta=inp.beta,
202 + A_log=inp.A_log,
203 + dt_bias=inp.dt_bias,
204 + state_indices=inp.state_indices,
205 + state=inp.recurrent_state,
206 + out=inp.out,
207 + lower_bound=GATE_LOWER_BOUND,
208 + output_gate=inp.g2,
209 + norm_weight=inp.decode_norm_weight,
210 + norm_eps=NORM_EPS,
211 + )
212 +
213 +
214 +def main() -> None:
215 + parser = argparse.ArgumentParser(description=__doc__)
216 + parser.add_argument("--tokens", type=int, nargs="+", default=[1, 8, 32, 64, 128])
217 + parser.add_argument(
218 + "--graph",
219 + action="store_true",
220 + help="time under CUDA-graph replay instead of eager dispatch",
221 + )
222 + parser.add_argument(
223 + "--layers",
224 + type=int,
225 + default=1,
226 + help=(
227 + "with --graph, give each of N layers its own state buffers so the "
228 + "recurrent state cannot stay resident in Infinity Cache. "
229 + f"Use {NUM_KDA_LAYERS} for a production-shaped step (needs "
230 + "~0.9 GB per layer at 128 tokens / 12 heads)."
231 + ),
232 + )
233 + parser.add_argument(
234 + "--heads",
235 + type=int,
236 + nargs="+",
237 + default=[12],
238 + help="KDA heads per rank (96 total / TP size)",
239 + )
240 + args = parser.parse_args()
241 +
242 + if not hasattr(torch.ops._C, "fused_kda_decode"):
243 + raise SystemExit("vLLM was built without the fused KDA decode kernel")
244 +
245 + props = torch.cuda.get_device_properties(0)
246 + bench = _bench_graph if args.graph else _bench
247 + mode = "cuda-graph replay" if args.graph else "eager dispatch"
248 + print(f"device: {props.name} ({props.gcnArchName}) timing: {mode}")
249 + print(
250 + f"{'heads':>6} {'tokens':>7} {'eager-norm':>11} {'triton-norm':>12} "
251 + f"{'fused us':>9} {'speedup':>8} {'state TB/s':>11} {'saved ms/step':>14}"
252 + )
253 + for num_heads in args.heads:
254 + for num_tokens in args.tokens:
255 + if args.graph and args.layers > 1:
256 + layers = [Inputs(num_tokens, num_heads) for _ in range(args.layers)]
257 + eager_ms = _bench_graph_layers(
258 + [functools.partial(triton_chain, i, False) for i in layers]
259 + )
260 + triton_ms = _bench_graph_layers(
261 + [functools.partial(triton_chain, i, True) for i in layers]
262 + )
263 + fused_ms = _bench_graph_layers(
264 + [functools.partial(fused, i) for i in layers]
265 + )
266 + inp = layers[0]
267 + del layers
268 + torch.accelerator.empty_cache()
269 + else:
270 + inp = Inputs(num_tokens, num_heads)
271 + eager_ms = bench(functools.partial(triton_chain, inp, False))
272 + triton_ms = bench(functools.partial(triton_chain, inp, True))
273 + fused_ms = bench(functools.partial(fused, inp))
274 + bandwidth = inp.state_bytes() * 1e3 / fused_ms / 1e12
275 + saved = (triton_ms - fused_ms) * NUM_KDA_LAYERS
276 + print(
277 + f"{num_heads:>6} {num_tokens:>7} {eager_ms * 1e3:>11.2f} "
278 + f"{triton_ms * 1e3:>12.2f} {fused_ms * 1e3:>9.2f} "
279 + f"{triton_ms / fused_ms:>7.2f}x {bandwidth:>10.2f} {saved:>14.3f}"
280 + )
281 + print(
282 + f"\n'eager-norm' is the chain as it runs today (FusedRMSNormGated falls "
283 + f"back to ~10 eager ops when custom_ops are off);\n'triton-norm' uses the "
284 + f"Triton rms_norm_gated kernel, which is the honest baseline for this "
285 + f"fusion.\nspeedup and saved ms/step are against 'triton-norm', over "
286 + f"{NUM_KDA_LAYERS} KDA layers per forward pass."
287 + )
288 +
289 +
290 +if __name__ == "__main__":
291 + main()
csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu new
+800
@@ -0,0 +1,800 @@
1 +/*
2 + * SPDX-License-Identifier: Apache-2.0
3 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 + *
5 + * ROCm port of the fused Kimi-K3 KDA decode kernel.
6 + *
7 + * Fuses, per (sequence, value head), the causal conv1d update over Q/K/V, the
8 + * gated delta-rule recurrence, and the gated output RMSNorm -- the work the
9 + * ROCm path otherwise spends three Triton launches and two copies on, in each
10 + * of the 69 KDA layers.
11 + *
12 + * The recurrent state is [128, 128] fp32 per (sequence, head), read and written
13 + * once per token, so the kernel is bound by that 128 KiB of HBM traffic and the
14 + * shape is chosen to reach streaming bandwidth rather than to minimize FLOPs.
15 + */
16 +
17 +#include <cstdint>
18 +#include <array>
19 +#include <optional>
20 +
21 +#include <hip/hip_runtime.h>
22 +#include <hip/hip_bf16.h>
23 +
24 +#include <torch/headeronly/core/ScalarType.h>
25 +
26 +#include "../torch_utils.h"
27 +
28 +namespace {
29 +
30 +using bf16_t = __hip_bfloat16;
31 +// HIP's float4 is a class type, which the non-temporal builtins reject.
32 +using f32x4 = float __attribute__((ext_vector_type(4)));
33 +
34 +constexpr int kDimK = 128;
35 +constexpr int kDimV = 128;
36 +constexpr int kKernelWidth = 4;
37 +constexpr int kConvStateWidth = kKernelWidth - 1;
38 +constexpr int kThreads = 256;
39 +// One state row is split across kLanes lanes, each holding kElems = 8 floats in
40 +// two float4s. 16 lanes is the widest group a pure-DPP reduction can span
41 +// (row_mirror is xor 15), and it keeps 8 loads per lane in flight.
42 +constexpr int kLanes = 16;
43 +constexpr int kElems = kDimK / kLanes;
44 +constexpr int kVec = kElems / 4;
45 +constexpr int kRows = kThreads / kLanes;
46 +constexpr int kIters = kDimV / kRows;
47 +constexpr int kGroups = kThreads / 16;
48 +// s_reduce layout: [3 * kGroups) cross-group partials, then the three results,
49 +// then the epilogue slot, then the per-row-group sums of squares.
50 +constexpr int kSumsqSlot = 3 * kGroups + 3;
51 +constexpr int kPartialBase = 3 * kGroups + 4;
52 +
53 +struct KdaDecodeStrides {
54 + int64_t x_row;
55 + int64_t beta_row;
56 + int64_t onorm_row;
57 + int64_t conv_slot;
58 + int64_t state_slot;
59 +};
60 +
61 +__device__ __forceinline__ float bf16_load(const bf16_t* ptr, int64_t idx) {
62 + return __bfloat162float(ptr[idx]);
63 +}
64 +
65 +__device__ __forceinline__ float bf16_load(const float* ptr, int64_t idx) {
66 + return ptr[idx];
67 +}
68 +
69 +template <int kChannels>
70 +__device__ __forceinline__ float conv_weight_load(const float* ptr, int channel,
71 + int width) {
72 + return ptr[width * kChannels + channel];
73 +}
74 +
75 +__device__ __forceinline__ bf16_t bf16_store(float value) {
76 + return __float2bfloat16(value);
77 +}
78 +
79 +__device__ __forceinline__ float sigmoid_fast(float x) {
80 + return 1.0f / (1.0f + __expf(-x));
81 +}
82 +
83 +__device__ __forceinline__ float silu_fast(float x) {
84 + return x * sigmoid_fast(x);
85 +}
86 +
87 +__device__ __forceinline__ float softplus_fast(float x) {
88 + return x > 20.0f ? x : log1pf(__expf(x));
89 +}
90 +
91 +__device__ __forceinline__ f32x4 load_state(const float* p) {
92 + return __builtin_nontemporal_load(reinterpret_cast<const f32x4*>(p));
93 +}
94 +
95 +__device__ __forceinline__ void store_state(float* p, f32x4 v) {
96 + __builtin_nontemporal_store(v, reinterpret_cast<f32x4*>(p));
97 +}
98 +
99 +// DPP row permutations within a wavefront: quad_perm[1,0,3,2] is xor 1,
100 +// quad_perm[2,3,0,1] is xor 2, row_half_mirror is xor 7, row_mirror is xor 15.
101 +template <int kCtrl>
102 +__device__ __forceinline__ float dpp_add(float x) {
103 + const int shuffled = __builtin_amdgcn_update_dpp(
104 + 0, __builtin_bit_cast(int, x), kCtrl, 0xf, 0xf, /*bound_ctrl=*/true);
105 + return x + __builtin_bit_cast(float, shuffled);
106 +}
107 +
108 +// Sum across each aligned group of 16 lanes; every lane ends with the total.
109 +__device__ __forceinline__ float row_reduce(float x) {
110 + x = dpp_add<0xb1>(x);
111 + x = dpp_add<0x4e>(x);
112 + x = dpp_add<0x141>(x);
113 + x = dpp_add<0x140>(x);
114 + return x;
115 +}
116 +
117 +struct Sum3 {
118 + float x, y, z;
119 +};
120 +
121 +// DPP within each 16-lane group, one LDS round trip across the kGroups groups.
122 +__device__ __forceinline__ Sum3 block_reduce_sum3(float x, float y, float z,
123 + float* scratch) {
124 + const int lane = threadIdx.x & 15;
125 + const int grp = threadIdx.x / 16;
126 + x = row_reduce(x);
127 + y = row_reduce(y);
128 + z = row_reduce(z);
129 + if (lane == 0) {
130 + scratch[grp] = x;
131 + scratch[kGroups + grp] = y;
132 + scratch[2 * kGroups + grp] = z;
133 + }
134 + __syncthreads();
135 + if (threadIdx.x < kGroups) {
136 + const float bx = row_reduce(scratch[lane]);
137 + const float by = row_reduce(scratch[kGroups + lane]);
138 + const float bz = row_reduce(scratch[2 * kGroups + lane]);
139 + if (lane == 0) {
140 + scratch[3 * kGroups] = bx;
141 + scratch[3 * kGroups + 1] = by;
142 + scratch[3 * kGroups + 2] = bz;
143 + }
144 + }
145 + __syncthreads();
146 + return {scratch[3 * kGroups], scratch[3 * kGroups + 1],
147 + scratch[3 * kGroups + 2]};
148 +}
149 +
150 +template <bool kApplyOnorm, int kFixedHeads, bool kUpdateConvState,
151 + bool kUseLowerBound, bool kApplyBetaSigmoid>
152 +__global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
153 + const bf16_t* __restrict__ x_q, const bf16_t* __restrict__ x_k,
154 + const bf16_t* __restrict__ x_v, const float* __restrict__ w_q_t,
155 + const float* __restrict__ w_k_t, const float* __restrict__ w_v_t,
156 + const float* __restrict__ bias_q, const float* __restrict__ bias_k,
157 + const float* __restrict__ bias_v, bf16_t* __restrict__ cs_q,
158 + bf16_t* __restrict__ cs_k, bf16_t* __restrict__ cs_v,
159 + const float* __restrict__ a_log, const bf16_t* __restrict__ g,
160 + const float* __restrict__ dt_bias, const bf16_t* __restrict__ beta,
161 + const bf16_t* __restrict__ onorm_g, const float* __restrict__ onorm_weight,
162 + const int* __restrict__ ssm_state_indices, float* __restrict__ state,
163 + bf16_t* __restrict__ out, float lower_bound, float scale, float onorm_eps,
164 + KdaDecodeStrides strides) {
165 + const int tid = threadIdx.x;
166 +
167 + // Heads are the fast grid axis so that adjacent workgroups walk adjacent
168 + // 64 KiB state slices of the same slot.
169 + const int i_hv = blockIdx.x;
170 + const int i_n = blockIdx.y;
171 + const int i_h = i_hv;
172 + const int bos = i_n;
173 + const int slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n];
174 +
175 + constexpr int kLocalDim = kFixedHeads * kDimK;
176 + // gdn_attn.py fills the tail of a CUDA-graph decode batch with
177 + // NULL_BLOCK_ID (0). Both kernels this one replaces skip it: the Triton
178 + // recurrent kernel zeroes the output and returns, and causal_conv1d_update
179 + // leaves the conv state alone. Without this, every padded row would read and
180 + // rewrite slot 0's state for nothing. `slot` is block-uniform, so returning
181 + // here is safe with respect to the __syncthreads() calls below.
182 + if (slot <= 0) {
183 + if (tid < kDimV) {
184 + out[static_cast<int64_t>(i_n) * kLocalDim + i_hv * kDimV + tid] =
185 + bf16_store(0.0f);
186 + }
187 + return;
188 + }
189 + constexpr int kPackedDim = 3 * kLocalDim;
190 + const int hk_off = i_h * kDimK;
191 + const int hv_off = i_hv * kDimV;
192 + float* const state_for_slot = state + slot * strides.state_slot;
193 + const int64_t conv_slot_offset = slot * strides.conv_slot;
194 + bf16_t* const cs_q_for_slot = cs_q + conv_slot_offset;
195 + bf16_t* const cs_k_for_slot = cs_k + conv_slot_offset;
196 + bf16_t* const cs_v_for_slot = cs_v + conv_slot_offset;
197 +
198 + __shared__ float s_q[kDimK];
199 + __shared__ float s_k[kDimK];
200 + __shared__ float s_decay[kDimK];
201 + __shared__ float s_v[kDimV];
202 + __shared__ float s_o[kDimV];
203 + __shared__ float s_reduce[kPartialBase + kRows];
204 + float pre_onorm_gate = 0.0f;
205 + float pre_onorm_weight = 0.0f;
206 +
207 + const int rlane = tid & (kLanes - 1);
208 + const int rgroup = tid / kLanes;
209 + float* const row_base =
210 + state_for_slot + static_cast<int64_t>(hv_off) * kDimK + 4 * rlane;
211 +
212 + // Nothing in the prologue feeds these addresses, so issue the whole state
213 + // slice now and let it fly under the conv1d, the RMS reduction and the two
214 + // barriers that follow.
215 + f32x4 h_raw[kIters][kVec];
216 +#pragma unroll
217 + for (int t = 0; t < kIters; ++t) {
218 + const float* src =
219 + row_base + static_cast<int64_t>(t * kRows + rgroup) * kDimK;
220 +#pragma unroll
221 + for (int j = 0; j < kVec; ++j) {
222 + h_raw[t][j] = load_state(src + j * 4 * kLanes);
223 + }
224 + }
225 +
226 + // Q and K stay in registers: the RMS reduction below only ever reads each
227 + // thread's own element, so it needs no barrier ahead of it.
228 + float qv = 0.0f;
229 + float kv = 0.0f;
230 + if (tid < kDimK) {
231 + const int k = tid;
232 + const int hk = hk_off + k;
233 + // A_log is uniform across the block; the scalar unit hoists this load.
234 + const float exp_a = __expf(a_log[i_h]);
235 +
236 + float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk);
237 + float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk);
238 + if constexpr (kUpdateConvState) {
239 + bf16_t q_shift0 = __float2bfloat16(0.0f);
240 + bf16_t q_shift1 = __float2bfloat16(0.0f);
241 + bf16_t k_shift0 = __float2bfloat16(0.0f);
242 + bf16_t k_shift1 = __float2bfloat16(0.0f);
243 +#pragma unroll
244 + for (int w = 0; w < kConvStateWidth; ++w) {
245 + const bf16_t q_state = cs_q_for_slot[hk + w * kPackedDim];
246 + const bf16_t k_state = cs_k_for_slot[hk + w * kPackedDim];
247 + q_acc += __bfloat162float(q_state) *
248 + conv_weight_load<kLocalDim>(w_q_t, hk, w);
249 + k_acc += __bfloat162float(k_state) *
250 + conv_weight_load<kLocalDim>(w_k_t, hk, w);
251 + if (w == 1) {
252 + q_shift0 = q_state;
253 + k_shift0 = k_state;
254 + } else if (w == 2) {
255 + q_shift1 = q_state;
256 + k_shift1 = k_state;
257 + }
258 + }
259 + const int64_t xq_idx = bos * strides.x_row + i_h * kDimK + k;
260 + const bf16_t q_new = x_q[xq_idx];
261 + const bf16_t k_new = x_k[xq_idx];
262 + q_acc += __bfloat162float(q_new) *
263 + conv_weight_load<kLocalDim>(w_q_t, hk, kKernelWidth - 1);
264 + k_acc += __bfloat162float(k_new) *
265 + conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1);
266 +
267 + cs_q_for_slot[hk] = q_shift0;
268 + cs_q_for_slot[hk + kPackedDim] = q_shift1;
269 + cs_q_for_slot[hk + 2 * kPackedDim] = q_new;
270 + cs_k_for_slot[hk] = k_shift0;
271 + cs_k_for_slot[hk + kPackedDim] = k_shift1;
272 + cs_k_for_slot[hk + 2 * kPackedDim] = k_new;
273 + } else {
274 +#pragma unroll
275 + for (int w = 0; w < kConvStateWidth; ++w) {
276 + const int cs_idx = hk + w * kPackedDim;
277 + q_acc += bf16_load(cs_q_for_slot, cs_idx) *
278 + conv_weight_load<kLocalDim>(w_q_t, hk, w);
279 + k_acc += bf16_load(cs_k_for_slot, cs_idx) *
280 + conv_weight_load<kLocalDim>(w_k_t, hk, w);
281 + }
282 + q_acc += bf16_load(x_q, bos * strides.x_row + i_h * kDimK + k) *
283 + conv_weight_load<kLocalDim>(w_q_t, hk, kKernelWidth - 1);
284 + k_acc += bf16_load(x_k, bos * strides.x_row + i_h * kDimK + k) *
285 + conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1);
286 + }
287 +
288 + qv = silu_fast(q_acc);
289 + kv = silu_fast(k_acc);
290 +
291 + const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k;
292 + const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk];
293 + if constexpr (kUseLowerBound) {
294 + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw));
295 + } else {
296 + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw));
297 + }
298 + }
299 +
300 + if (tid >= kThreads - kDimV) {
301 + const int v = tid - (kThreads - kDimV);
302 + const int hvv = hv_off + v;
303 +
304 + float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv);
305 + if constexpr (kUpdateConvState) {
306 + bf16_t v_shift0 = __float2bfloat16(0.0f);
307 + bf16_t v_shift1 = __float2bfloat16(0.0f);
308 +#pragma unroll
309 + for (int w = 0; w < kConvStateWidth; ++w) {
310 + const bf16_t v_state = cs_v_for_slot[hvv + w * kPackedDim];
311 + v_acc += __bfloat162float(v_state) *
312 + conv_weight_load<kLocalDim>(w_v_t, hvv, w);
313 + if (w == 1) {
314 + v_shift0 = v_state;
315 + } else if (w == 2) {
316 + v_shift1 = v_state;
317 + }
318 + }
319 + const int64_t xv_idx = bos * strides.x_row + i_hv * kDimV + v;
320 + const bf16_t v_new = x_v[xv_idx];
321 + v_acc += __bfloat162float(v_new) *
322 + conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1);
323 + cs_v_for_slot[hvv] = v_shift0;
324 + cs_v_for_slot[hvv + kPackedDim] = v_shift1;
325 + cs_v_for_slot[hvv + 2 * kPackedDim] = v_new;
326 + } else {
327 +#pragma unroll
328 + for (int w = 0; w < kConvStateWidth; ++w) {
329 + const int cs_idx = hvv + w * kPackedDim;
330 + v_acc += bf16_load(cs_v_for_slot, cs_idx) *
331 + conv_weight_load<kLocalDim>(w_v_t, hvv, w);
332 + }
333 + v_acc += bf16_load(x_v, bos * strides.x_row + i_hv * kDimV + v) *
334 + conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1);
335 + }
336 + s_v[v] = silu_fast(v_acc);
337 +
338 + if constexpr (kApplyOnorm) {
339 + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v;
340 + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx));
341 + pre_onorm_weight = onorm_weight[v];
342 + }
343 + }
344 +
345 + // Block-uniform address, so this is a scalar load.
346 + const float beta_raw = bf16_load(beta, bos * strides.beta_row + i_hv);
347 + const float beta_v = kApplyBetaSigmoid ? sigmoid_fast(beta_raw) : beta_raw;
348 +
349 + // One reduction yields both RMS norms and the raw q.k that folds the rank-1
350 + // state update into the output below.
351 + const Sum3 sums = block_reduce_sum3(qv * qv, kv * kv, qv * kv, s_reduce);
352 + const float rq = rsqrtf(sums.x + 1.0e-6f) * scale;
353 + const float rk = rsqrtf(sums.y + 1.0e-6f);
354 + const float kq = sums.z * rq * rk;
355 + if (tid < kDimK) {
356 + s_q[tid] = qv * rq;
357 + s_k[tid] = kv * rk;
358 + }
359 + __syncthreads();
360 +
361 + float r_q[kElems], r_k[kElems], r_decay[kElems];
362 +#pragma unroll
363 + for (int j = 0; j < kVec; ++j) {
364 + const int kb = 4 * rlane + j * 4 * kLanes;
365 + const float4 q4 = *reinterpret_cast<const float4*>(s_q + kb);
366 + const float4 k4 = *reinterpret_cast<const float4*>(s_k + kb);
367 + const float4 d4 = *reinterpret_cast<const float4*>(s_decay + kb);
368 + r_q[4 * j + 0] = q4.x;
369 + r_q[4 * j + 1] = q4.y;
370 + r_q[4 * j + 2] = q4.z;
371 + r_q[4 * j + 3] = q4.w;
372 + r_k[4 * j + 0] = k4.x;
373 + r_k[4 * j + 1] = k4.y;
374 + r_k[4 * j + 2] = k4.z;
375 + r_k[4 * j + 3] = k4.w;
376 + r_decay[4 * j + 0] = d4.x;
377 + r_decay[4 * j + 1] = d4.y;
378 + r_decay[4 * j + 2] = d4.z;
379 + r_decay[4 * j + 3] = d4.w;
380 + }
381 +
382 + // Per row v, with h' = h * decay:
383 + // u = h'.k vn = (v - u) * beta
384 + // h'' = h' + k * vn o = h''.q = h'.q + vn * (k.q)
385 + // Folding the rank-1 update into the output lets h'.k and h'.q share one
386 + // reduction and keeps the state store off the second one's critical path.
387 + float sumsq_local = 0.0f;
388 +#pragma unroll
389 + for (int t = 0; t < kIters; ++t) {
390 + float h[kElems];
391 +#pragma unroll
392 + for (int j = 0; j < kVec; ++j) {
393 +#pragma unroll
394 + for (int e = 0; e < 4; ++e) {
395 + h[4 * j + e] = h_raw[t][j][e] * r_decay[4 * j + e];
396 + }
397 + }
398 + // Four partial sums per dot product: at decode the grid rarely fills the
399 + // wave slots, so the serial FMA chain is exposed latency, not throughput.
400 + float part_k[4] = {0.0f, 0.0f, 0.0f, 0.0f};
401 + float part_q[4] = {0.0f, 0.0f, 0.0f, 0.0f};
402 +#pragma unroll
403 + for (int e = 0; e < kElems; ++e) {
404 + part_k[e & 3] += h[e] * r_k[e];
405 + part_q[e & 3] += h[e] * r_q[e];
406 + }
407 + const float dot_hk =
408 + row_reduce((part_k[0] + part_k[1]) + (part_k[2] + part_k[3]));
409 + const float dot_hq =
410 + row_reduce((part_q[0] + part_q[1]) + (part_q[2] + part_q[3]));
411 +
412 + const int v = t * kRows + rgroup;
413 + const float v_new = (s_v[v] - dot_hk) * beta_v;
414 + float* dst = row_base + static_cast<int64_t>(v) * kDimK;
415 +#pragma unroll
416 + for (int j = 0; j < kVec; ++j) {
417 + f32x4 updated;
418 +#pragma unroll
419 + for (int e = 0; e < 4; ++e) {
420 + updated[e] = h[4 * j + e] + r_k[4 * j + e] * v_new;
421 + }
422 + store_state(dst + j * 4 * kLanes, updated);
423 + }
424 + if (rlane == 0) {
425 + const float o = dot_hq + v_new * kq;
426 + s_o[v] = o;
427 + sumsq_local += o * o;
428 + }
429 + }
430 +
431 + if constexpr (kApplyOnorm) {
432 + // Each row-group leader already holds the squares of every output it
433 + // produced, so only the kRows group totals cross LDS.
434 + if (rlane == 0) s_reduce[kPartialBase + rgroup] = sumsq_local;
435 + __syncthreads();
436 + if (tid < 16) {
437 + float total = 0.0f;
438 +#pragma unroll
439 + for (int i = 0; i < kRows / 16; ++i) {
440 + total += s_reduce[kPartialBase + tid + i * 16];
441 + }
442 + total = row_reduce(total);
443 + if (tid == 0) s_reduce[kSumsqSlot] = total;
444 + }
445 + __syncthreads();
446 + const float rstd =
447 + rsqrtf(s_reduce[kSumsqSlot] / static_cast<float>(kDimV) + onorm_eps);
448 + if (tid >= kThreads - kDimV) {
449 + const int v = tid - (kThreads - kDimV);
450 + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + v;
451 + out[out_idx] =
452 + bf16_store(s_o[v] * rstd * pre_onorm_weight * pre_onorm_gate);
453 + }
454 + } else {
455 + __syncthreads();
456 + if (tid < kDimV) {
457 + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid;
458 + out[out_idx] = bf16_store(s_o[tid]);
459 + }
460 + }
461 +}
462 +
463 +template <int kHeads, bool kApplyOnorm, bool kUpdateConvState,
464 + bool kUseLowerBound, bool kApplyBetaSigmoid>
465 +void launch_kda_decode_raw(
466 + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t,
467 + const void* w_k_t, const void* w_v_t, const void* bias_q,
468 + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v,
469 + const float* a_log, const void* g, const float* dt_bias, const void* beta,
470 + const void* onorm_g, const float* onorm_weight,
471 + const int* ssm_state_indices, float* state, void* out, int B,
472 + float lower_bound, float scale, float onorm_eps, KdaDecodeStrides strides,
473 + hipStream_t stream) {
474 + kda_decode_fusion_kernel<kApplyOnorm, kHeads, kUpdateConvState,
475 + kUseLowerBound, kApplyBetaSigmoid>
476 + <<<dim3(kHeads, B), dim3(kThreads), 0, stream>>>(
477 + reinterpret_cast<const bf16_t*>(x_q),
478 + reinterpret_cast<const bf16_t*>(x_k),
479 + reinterpret_cast<const bf16_t*>(x_v),
480 + reinterpret_cast<const float*>(w_q_t),
481 + reinterpret_cast<const float*>(w_k_t),
482 + reinterpret_cast<const float*>(w_v_t),
483 + reinterpret_cast<const float*>(bias_q),
484 + reinterpret_cast<const float*>(bias_k),
485 + reinterpret_cast<const float*>(bias_v),
486 + reinterpret_cast<bf16_t*>(cs_q), reinterpret_cast<bf16_t*>(cs_k),
487 + reinterpret_cast<bf16_t*>(cs_v), a_log,
488 + reinterpret_cast<const bf16_t*>(g), dt_bias,
489 + reinterpret_cast<const bf16_t*>(beta),
490 + reinterpret_cast<const bf16_t*>(onorm_g), onorm_weight,
491 + ssm_state_indices, state, reinterpret_cast<bf16_t*>(out), lower_bound,
492 + scale, onorm_eps, strides);
493 +}
494 +
495 +struct KdaDecodeLaunchParams {
496 + const void* x_q;
497 + const void* x_k;
498 + const void* x_v;
499 + const void* w_q_t;
500 + const void* w_k_t;
501 + const void* w_v_t;
502 + const void* bias_q;
503 + const void* bias_k;
504 + const void* bias_v;
505 + void* cs_q;
506 + void* cs_k;
507 + void* cs_v;
508 + const float* a_log;
509 + const void* g;
510 + const float* dt_bias;
511 + const void* beta;
512 + const void* onorm_g;
513 + const float* onorm_weight;
514 + const int* ssm_state_indices;
515 + float* state;
516 + void* out;
517 + int B;
518 + int H;
519 + bool update_conv_cache;
520 + float lower_bound;
521 + float scale;
522 + float onorm_eps;
523 + KdaDecodeStrides strides;
524 + hipStream_t stream;
525 +};
526 +
527 +template <int kHeads, bool kApplyOnorm, bool kUseLowerBound,
528 + bool kApplyBetaSigmoid>
529 +void dispatch_kda_decode_conv(const KdaDecodeLaunchParams& p) {
530 +#define LAUNCH_KDA_DECODE(UPDATE_CONV) \
531 + launch_kda_decode_raw<kHeads, kApplyOnorm, UPDATE_CONV, kUseLowerBound, \
532 + kApplyBetaSigmoid>( \
533 + p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k, \
534 + p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, \
535 + p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.state, p.out, p.B, \
536 + p.lower_bound, p.scale, p.onorm_eps, p.strides, p.stream)
537 + if (p.update_conv_cache) {
538 + LAUNCH_KDA_DECODE(true);
539 + } else {
540 + LAUNCH_KDA_DECODE(false);
541 + }
542 +#undef LAUNCH_KDA_DECODE
543 +}
544 +
545 +template <bool kApplyOnorm, bool kUseLowerBound, bool kApplyBetaSigmoid>
546 +void dispatch_kda_decode_heads(const KdaDecodeLaunchParams& p) {
547 + switch (p.H) {
548 + case 12:
549 + dispatch_kda_decode_conv<12, kApplyOnorm, kUseLowerBound,
550 + kApplyBetaSigmoid>(p);
551 + break;
552 + case 24:
553 + dispatch_kda_decode_conv<24, kApplyOnorm, kUseLowerBound,
554 + kApplyBetaSigmoid>(p);
555 + break;
556 + case 48:
557 + dispatch_kda_decode_conv<48, kApplyOnorm, kUseLowerBound,
558 + kApplyBetaSigmoid>(p);
559 + break;
560 + case 96:
561 + dispatch_kda_decode_conv<96, kApplyOnorm, kUseLowerBound,
562 + kApplyBetaSigmoid>(p);
563 + break;
564 + default:
565 + STD_TORCH_CHECK(false, "Unsupported number of heads: ", p.H);
566 + }
567 +}
568 +
569 +template <bool kApplyOnorm, bool kUseLowerBound>
570 +void dispatch_kda_decode_beta(const KdaDecodeLaunchParams& p,
571 + bool apply_beta_sigmoid) {
572 + if (apply_beta_sigmoid) {
573 + dispatch_kda_decode_heads<kApplyOnorm, kUseLowerBound, true>(p);
574 + } else {
575 + dispatch_kda_decode_heads<kApplyOnorm, kUseLowerBound, false>(p);
576 + }
577 +}
578 +
579 +template <bool kApplyOnorm>
580 +void dispatch_kda_decode_decay(const KdaDecodeLaunchParams& p,
581 + bool use_lower_bound, bool apply_beta_sigmoid) {
582 + if (use_lower_bound) {
583 + dispatch_kda_decode_beta<kApplyOnorm, true>(p, apply_beta_sigmoid);
584 + } else {
585 + dispatch_kda_decode_beta<kApplyOnorm, false>(p, apply_beta_sigmoid);
586 + }
587 +}
588 +
589 +void dispatch_kda_decode_features(const KdaDecodeLaunchParams& p,
590 + bool apply_onorm, bool use_lower_bound,
591 + bool apply_beta_sigmoid) {
592 + if (apply_onorm) {
593 + dispatch_kda_decode_decay<true>(p, use_lower_bound, apply_beta_sigmoid);
594 + } else {
595 + dispatch_kda_decode_decay<false>(p, use_lower_bound, apply_beta_sigmoid);
596 + }
597 +}
598 +
599 +} // namespace
600 +
601 +void fused_kda_decode(
602 + torch::stable::Tensor const& x, torch::stable::Tensor const& weight,
603 + std::optional<torch::stable::Tensor> bias,
604 + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g,
605 + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log,
606 + torch::stable::Tensor const& dt_bias,
607 + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state,
608 + torch::stable::Tensor& out, std::optional<double> lower_bound,
609 + std::optional<torch::stable::Tensor> output_gate,
610 + std::optional<torch::stable::Tensor> norm_weight, double norm_eps) {
611 + using torch::headeronly::ScalarType;
612 + constexpr int kHeadDim = 128;
613 + constexpr int kConvWidth = 4;
614 +
615 + STD_TORCH_CHECK(x.is_cuda() && x.scalar_type() == ScalarType::BFloat16,
616 + "x must be a GPU bfloat16 tensor");
617 + STD_TORCH_CHECK(weight.is_cuda() && weight.scalar_type() == ScalarType::Float,
618 + "weight must be a GPU float32 tensor");
619 + STD_TORCH_CHECK(
620 + conv_state.is_cuda() && conv_state.scalar_type() == ScalarType::BFloat16,
621 + "conv_state must be a GPU bfloat16 tensor");
622 + STD_TORCH_CHECK(
623 + raw_g.is_cuda() && raw_g.scalar_type() == ScalarType::BFloat16,
624 + "raw_g must be a GPU bfloat16 tensor");
625 + STD_TORCH_CHECK(
626 + raw_beta.is_cuda() && raw_beta.scalar_type() == ScalarType::BFloat16,
627 + "raw_beta must be a GPU bfloat16 tensor");
628 + STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float,
629 + "A_log must be a GPU float32 tensor");
630 + STD_TORCH_CHECK(
631 + dt_bias.is_cuda() && dt_bias.scalar_type() == ScalarType::Float,
632 + "dt_bias must be a GPU float32 tensor");
633 + STD_TORCH_CHECK(state.is_cuda() && state.scalar_type() == ScalarType::Float,
634 + "state must be a GPU float32 tensor");
635 + STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16,
636 + "out must be a GPU bfloat16 tensor");
637 + STD_TORCH_CHECK(
638 + state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int,
639 + "state_indices must be a GPU int32 tensor");
640 +
641 + STD_TORCH_CHECK(x.dim() == 2, "x must have shape [B, 3 * H * 128]");
642 + int const batch_size = static_cast<int>(x.size(0));
643 + int64_t const qkv_width = x.size(1);
644 + STD_TORCH_CHECK(qkv_width % (3 * kHeadDim) == 0,
645 + "x must have shape [B, 3 * H * 128]");
646 + int64_t const num_heads = qkv_width / (3 * kHeadDim);
647 + STD_TORCH_CHECK(
648 + num_heads == 12 || num_heads == 24 || num_heads == 48 || num_heads == 96,
649 + "H must be 12, 24, 48, or 96, got ", num_heads);
650 + STD_TORCH_CHECK(batch_size > 0,
651 + "KDA decode fusion requires at least one row");
652 + // Rows are the y grid axis so that heads, which are contiguous in the state,
653 + // vary fastest across workgroups.
654 + STD_TORCH_CHECK(batch_size <= 65535, "KDA decode fusion supports at most ",
655 + 65535, " rows, got ", batch_size);
656 + int const dim = num_heads * kHeadDim;
657 +
658 + STD_TORCH_CHECK(weight.dim() == 3 && weight.is_contiguous() &&
659 + weight.size(0) == 3 && weight.size(1) == kConvWidth &&
660 + weight.size(2) == dim,
661 + "weight must have shape [3, 4, H * 128]");
662 + STD_TORCH_CHECK(conv_state.dim() == 3 && conv_state.size(1) == 3 * dim &&
663 + conv_state.size(2) == kConvWidth - 1,
664 + "conv_state must have shape [slots, 3 * H * 128, 3]");
665 + STD_TORCH_CHECK(raw_g.dim() == 4 && raw_g.size(0) == 1 &&
666 + raw_g.size(1) == batch_size &&
667 + raw_g.size(2) == num_heads && raw_g.size(3) == kHeadDim,
668 + "raw_g must have shape [1, B, H, 128]");
669 + STD_TORCH_CHECK(raw_beta.dim() == 3 && raw_beta.size(0) == 1 &&
670 + raw_beta.size(1) == batch_size &&
671 + raw_beta.size(2) == num_heads,
672 + "raw_beta must have shape [1, B, H]");
673 + STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_heads,
674 + "A_log must be contiguous with H elements");
675 + STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == dim,
676 + "dt_bias must be contiguous with H * 128 elements");
677 + STD_TORCH_CHECK(
678 + state_indices.is_contiguous() && state_indices.numel() == batch_size,
679 + "state_indices must be contiguous with B elements");
680 + STD_TORCH_CHECK(state.dim() == 4 && state.size(1) == num_heads &&
681 + state.size(2) == kHeadDim && state.size(3) == kHeadDim,
682 + "state must have shape [slots, H, 128, 128]");
683 + STD_TORCH_CHECK(out.dim() == 4 && out.size(0) == 1 &&
684 + out.size(1) == batch_size && out.size(2) == num_heads &&
685 + out.size(3) == kHeadDim,
686 + "out must have shape [1, B, H, 128]");
687 + STD_TORCH_CHECK(x.stride(1) == 1,
688 + "x must be contiguous in its channel dimension");
689 + STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
690 + conv_state.stride(1) == 1 &&
691 + conv_state.stride(2) == 3 * dim,
692 + "conv_state must use the SD cache layout");
693 + STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim &&
694 + state.stride(1) == kHeadDim * kHeadDim &&
695 + state.stride(2) == kHeadDim && state.stride(3) == 1,
696 + "state must have contiguous [H, 128, 128] slot contents");
697 + STD_TORCH_CHECK(raw_g.is_contiguous(), "raw_g must be contiguous");
698 + STD_TORCH_CHECK(raw_beta.stride(2) == 1,
699 + "raw_beta must be contiguous in its head dimension");
700 + STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
701 +
702 + bool const apply_onorm = output_gate.has_value();
703 + STD_TORCH_CHECK(apply_onorm == norm_weight.has_value(),
704 + "output_gate and norm_weight must be provided together");
705 + void const* output_gate_ptr = nullptr;
706 + float const* norm_weight_ptr = nullptr;
707 + int64_t output_gate_row_stride = 0;
708 + if (apply_onorm) {
709 + STD_TORCH_CHECK(output_gate->is_cuda() &&
710 + output_gate->scalar_type() == ScalarType::BFloat16,
711 + "output_gate must be a GPU bfloat16 tensor");
712 + bool const gate_is_3d =
713 + output_gate->dim() == 3 && output_gate->size(0) == batch_size &&
714 + output_gate->size(1) == num_heads && output_gate->size(2) == kHeadDim;
715 + bool const gate_is_4d =
716 + output_gate->dim() == 4 && output_gate->size(0) == 1 &&
717 + output_gate->size(1) == batch_size &&
718 + output_gate->size(2) == num_heads && output_gate->size(3) == kHeadDim;
719 + STD_TORCH_CHECK(gate_is_3d || gate_is_4d,
720 + "output_gate must have shape [B, H, 128] or "
721 + "[1, B, H, 128]");
722 + int const row_dim = gate_is_3d ? 0 : 1;
723 + STD_TORCH_CHECK(output_gate->stride(output_gate->dim() - 1) == 1,
724 + "output_gate must be contiguous in its last dimension");
725 + STD_TORCH_CHECK(output_gate->stride(row_dim + 1) == kHeadDim,
726 + "output_gate must have contiguous head rows");
727 + STD_TORCH_CHECK(norm_weight->is_cuda() &&
728 + norm_weight->scalar_type() == ScalarType::Float,
729 + "norm_weight must be a GPU float32 tensor");
730 + STD_TORCH_CHECK(
731 + norm_weight->is_contiguous() && norm_weight->numel() == kHeadDim,
732 + "norm_weight must be contiguous with 128 elements");
733 + STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative");
734 + output_gate_ptr = output_gate->data_ptr();
735 + norm_weight_ptr = static_cast<float const*>(norm_weight->data_ptr());
736 + output_gate_row_stride = output_gate->stride(row_dim);
737 + }
738 +
739 + void const* bias_ptr = nullptr;
740 + if (bias.has_value()) {
741 + STD_TORCH_CHECK(bias->is_cuda() && bias->scalar_type() == ScalarType::Float,
742 + "bias must be a GPU float32 tensor");
743 + STD_TORCH_CHECK(bias->is_contiguous() && bias->numel() == 3 * dim,
744 + "bias must be contiguous with 3 * H * 128 elements");
745 + bias_ptr = bias->data_ptr();
746 + }
747 +
748 + auto const* x_ptr = static_cast<char const*>(x.data_ptr());
749 + auto const* weight_ptr = static_cast<char const*>(weight.data_ptr());
750 + auto* conv_ptr = static_cast<char*>(conv_state.data_ptr());
751 + auto const* bias_bytes = static_cast<char const*>(bias_ptr);
752 + int64_t const segment_bytes = dim * sizeof(bf16_t);
753 + int64_t const weight_segment_bytes =
754 + dim * kConvWidth * static_cast<int64_t>(sizeof(float));
755 + int64_t const conv_segment_bytes =
756 + dim * conv_state.stride(1) * sizeof(bf16_t);
757 + int64_t const bias_segment_bytes = dim * sizeof(float);
758 + bool const use_lower_bound = lower_bound.has_value();
759 + float const lower_bound_value =
760 + use_lower_bound ? static_cast<float>(*lower_bound) : 0.0f;
761 +
762 + torch::stable::accelerator::DeviceGuard const device_guard(
763 + x.get_device_index());
764 + KdaDecodeLaunchParams const params{
765 + x_ptr,
766 + x_ptr + segment_bytes,
767 + x_ptr + 2 * segment_bytes,
768 + weight_ptr,
769 + weight_ptr + weight_segment_bytes,
770 + weight_ptr + 2 * weight_segment_bytes,
771 + bias_bytes,
772 + bias_bytes == nullptr ? nullptr : bias_bytes + bias_segment_bytes,
773 + bias_bytes == nullptr ? nullptr : bias_bytes + 2 * bias_segment_bytes,
774 + conv_ptr,
775 + conv_ptr + conv_segment_bytes,
776 + conv_ptr + 2 * conv_segment_bytes,
777 + static_cast<float const*>(a_log.data_ptr()),
778 + raw_g.data_ptr(),
779 + static_cast<float const*>(dt_bias.data_ptr()),
780 + raw_beta.data_ptr(),
781 + output_gate_ptr,
782 + norm_weight_ptr,
783 + static_cast<int const*>(state_indices.data_ptr()),
784 + static_cast<float*>(state.data_ptr()),
785 + out.data_ptr(),
786 + batch_size,
787 + static_cast<int>(num_heads),
788 + true,
789 + lower_bound_value,
790 + 0.08838834764831845f,
791 + static_cast<float>(norm_eps),
792 + KdaDecodeStrides{x.stride(0), raw_beta.stride(1), output_gate_row_stride,
793 + conv_state.stride(0), state.stride(0)},
794 + get_current_cuda_stream(x.get_device_index())};
795 + dispatch_kda_decode_features(params, apply_onorm, use_lower_bound, true);
796 + hipError_t const error = hipGetLastError();
797 + STD_TORCH_CHECK(
798 + error == hipSuccess,
799 + "Kimi K3 KDA decode kernel launch failed: ", hipGetErrorString(error));
800 +}
tests/models/kimi_k3/test_amd_kda_decode.py new
+291
@@ -0,0 +1,291 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""The ROCm fused KDA decode kernel must match the Triton chain it replaces.
4 +
5 +The fused kernel folds the packed causal conv1d update, the gated delta-rule
6 +recurrence and the gated output RMSNorm into one launch, and updates both the
7 +conv state and the recurrent state in place. Every one of those outputs is
8 +compared against the three-kernel path the AMD layer falls back to.
9 +"""
10 +
11 +import pytest
12 +import torch
13 +
14 +from vllm.platforms import current_platform
15 +
16 +
17 +def _on_gfx950() -> bool:
18 + if not current_platform.is_rocm():
19 + return False
20 + from vllm.platforms.rocm import on_gfx950
21 +
22 + return on_gfx950()
23 +
24 +
25 +pytestmark = pytest.mark.skipif(
26 + not _on_gfx950(),
27 + reason="The fused KDA decode kernel is only built for gfx950",
28 +)
29 +
30 +# Kimi-K3 KDA: 96 heads x 128, conv width 4, gate_lower_bound -5.0.
31 +HEAD_DIM = 128
32 +CONV_WIDTH = 4
33 +GATE_LOWER_BOUND = -5.0
34 +NORM_EPS = 1e-5
35 +DTYPE = torch.bfloat16
36 +
37 +
38 +def _requires_kernel() -> None:
39 + if not hasattr(torch.ops._C, "fused_kda_decode"):
40 + pytest.skip("vLLM was built without the fused KDA decode kernel")
41 +
42 +
43 +class KdaDecodeInputs:
44 + """One decode step of a Kimi-K3 KDA layer, in the layer's own layout."""
45 +
46 + def __init__(
47 + self,
48 + num_tokens: int,
49 + num_heads: int,
50 + num_slots: int,
51 + seed: int = 0,
52 + ) -> None:
53 + torch.manual_seed(seed)
54 + device = "cuda"
55 + dim = num_heads * HEAD_DIM
56 + self.num_heads = num_heads
57 + self.dim = dim
58 +
59 + self.mixed_qkv = (
60 + torch.randn(num_tokens, 3 * dim, device=device, dtype=DTYPE) * 0.5
61 + )
62 + # conv1d weight as the layer holds it: [3 * dim, width] fp32.
63 + self.conv_weights = (
64 + torch.randn(3 * dim, CONV_WIDTH, device=device, dtype=torch.float32) * 0.3
65 + )
66 + # Width-major fp32 copy the fused kernel indexes: [3, width, dim].
67 + self.decode_conv1d_weight = torch.stack(
68 + [
69 + self.conv_weights[i * dim : (i + 1) * dim].transpose(0, 1).contiguous()
70 + for i in range(3)
71 + ]
72 + )
73 + # SD cache layout: [slots, width - 1, 3 * dim]; the layer transposes it.
74 + self.conv_state = (
75 + torch.randn(num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE)
76 + * 0.5
77 + )
78 + self.recurrent_state = (
79 + torch.randn(
80 + num_slots,
81 + num_heads,
82 + HEAD_DIM,
83 + HEAD_DIM,
84 + device=device,
85 + dtype=torch.float32,
86 + )
87 + * 0.1
88 + )
89 + self.g1 = (
90 + torch.randn(1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE)
91 + * 0.5
92 + )
93 + self.g2 = (
94 + torch.randn(num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE)
95 + * 0.5
96 + )
97 + self.beta = torch.randn(1, num_tokens, num_heads, device=device, dtype=DTYPE)
98 + self.A_log = torch.randn(num_heads, device=device, dtype=torch.float32) * 0.5
99 + self.dt_bias = torch.randn(dim, device=device, dtype=torch.float32) * 0.1
100 + self.norm_weight_bf16 = 1 + 0.1 * torch.randn(
101 + HEAD_DIM, device=device, dtype=DTYPE
102 + )
103 + self.decode_norm_weight = self.norm_weight_bf16.float()
104 + # Distinct, shuffled slots: the kernel must honour the indirection.
105 + # Slot 0 is NULL_BLOCK_ID and never allocated for live state, and the
106 + # Triton reference short-circuits it, so draw from [1, num_slots).
107 + assert num_slots > num_tokens, "need a spare slot to detect stray writes"
108 + self.state_indices = (
109 + torch.randperm(num_slots - 1, device=device)[:num_tokens] + 1
110 + ).to(torch.int32)
111 +
112 + def conv_state_view(self, state: torch.Tensor) -> torch.Tensor:
113 + return state.transpose(-1, -2)
114 +
115 +
116 +def _gated_rmsnorm(
117 + x: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float
118 +) -> torch.Tensor:
119 + x_float = x.float()
120 + variance = x_float.pow(2).mean(dim=-1, keepdim=True)
121 + normed = x_float * torch.rsqrt(variance + eps) * weight.float()
122 + return (normed * torch.sigmoid(gate.float())).to(x.dtype)
123 +
124 +
125 +def _run_triton_chain(
126 + inp: KdaDecodeInputs,
127 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
128 + """conv1d update -> recurrent decode -> gated norm, as the layer runs it."""
129 + from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
130 + causal_conv1d_update,
131 + )
132 + from vllm.models.kimi_k3.amd.ops.third_party.kda import (
133 + fused_recurrent_kda_packed_decode,
134 + )
135 +
136 + conv_state = inp.conv_state.clone()
137 + recurrent_state = inp.recurrent_state.clone()
138 + conv_out = torch.empty_like(inp.mixed_qkv)
139 + causal_conv1d_update(
140 + inp.mixed_qkv,
141 + inp.conv_state_view(conv_state),
142 + inp.conv_weights,
143 + None,
144 + activation="silu",
145 + conv_state_indices=inp.state_indices,
146 + validate_data=True,
147 + out=conv_out,
148 + )
149 + core_attn_out, _ = fused_recurrent_kda_packed_decode(
150 + mixed_qkv=conv_out,
151 + raw_g=inp.g1,
152 + raw_beta=inp.beta,
153 + A_log=inp.A_log,
154 + dt_bias=inp.dt_bias,
155 + lower_bound=GATE_LOWER_BOUND,
156 + initial_state=recurrent_state,
157 + state_indices=inp.state_indices,
158 + )
159 + out = _gated_rmsnorm(core_attn_out, inp.g2, inp.norm_weight_bf16, NORM_EPS)
160 + return out, conv_state, recurrent_state
161 +
162 +
163 +def _run_fused(
164 + inp: KdaDecodeInputs,
165 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
166 + from vllm import _custom_ops as ops
167 +
168 + conv_state = inp.conv_state.clone()
169 + recurrent_state = inp.recurrent_state.clone()
170 + num_tokens = inp.mixed_qkv.shape[0]
171 + out = torch.empty(
172 + 1,
173 + num_tokens,
174 + inp.num_heads,
175 + HEAD_DIM,
176 + device=inp.mixed_qkv.device,
177 + dtype=DTYPE,
178 + )
179 + ops.fused_kda_decode(
180 + x=inp.mixed_qkv,
181 + weight=inp.decode_conv1d_weight,
182 + bias=None,
183 + conv_state=inp.conv_state_view(conv_state),
184 + raw_g=inp.g1,
185 + raw_beta=inp.beta,
186 + A_log=inp.A_log,
187 + dt_bias=inp.dt_bias,
188 + state_indices=inp.state_indices,
189 + state=recurrent_state,
190 + out=out,
191 + lower_bound=GATE_LOWER_BOUND,
192 + output_gate=inp.g2,
193 + norm_weight=inp.decode_norm_weight,
194 + norm_eps=NORM_EPS,
195 + )
196 + return out, conv_state, recurrent_state
197 +
198 +
199 +@pytest.mark.parametrize("num_heads", [12, 24, 96])
200 +@pytest.mark.parametrize("num_tokens", [1, 7, 128])
201 +@torch.inference_mode()
202 +def test_fused_kda_decode_matches_triton_chain(num_heads: int, num_tokens: int) -> None:
203 + _requires_kernel()
204 + inp = KdaDecodeInputs(num_tokens, num_heads, num_slots=max(num_tokens, 4) + 3)
205 +
206 + expected_out, expected_conv, expected_state = _run_triton_chain(inp)
207 + actual_out, actual_conv, actual_state = _run_fused(inp)
208 +
209 + # The fused kernel keeps the recurrent output in fp32 through the norm,
210 + # while the Triton chain rounds it to BF16 in between.
211 + torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2)
212 + # The conv state is a pure shift-and-append: it must be bit-exact.
213 + torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0)
214 + torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3)
215 +
216 +
217 +@torch.inference_mode()
218 +def test_fused_kda_decode_leaves_untouched_slots_alone() -> None:
219 + """Only the slots named by state_indices may be written."""
220 + _requires_kernel()
221 + num_tokens, num_heads, num_slots = 4, 12, 9
222 + inp = KdaDecodeInputs(num_tokens, num_heads, num_slots, seed=3)
223 +
224 + _, actual_conv, actual_state = _run_fused(inp)
225 +
226 + touched = set(inp.state_indices.tolist())
227 + untouched = [slot for slot in range(num_slots) if slot not in touched]
228 + assert untouched, "test needs at least one unused slot"
229 + torch.testing.assert_close(
230 + actual_conv[untouched], inp.conv_state[untouched], atol=0, rtol=0
231 + )
232 + torch.testing.assert_close(
233 + actual_state[untouched], inp.recurrent_state[untouched], atol=0, rtol=0
234 + )
235 +
236 +
237 +@torch.inference_mode()
238 +def test_fused_kda_decode_without_output_norm() -> None:
239 + """Omitting the gate/norm pair returns the raw recurrent output."""
240 + _requires_kernel()
241 + from vllm import _custom_ops as ops
242 +
243 + inp = KdaDecodeInputs(5, 12, num_slots=8, seed=7)
244 + _, _, expected_state = _run_triton_chain(inp)
245 +
246 + conv_state = inp.conv_state.clone()
247 + recurrent_state = inp.recurrent_state.clone()
248 + out = torch.empty(1, 5, inp.num_heads, HEAD_DIM, device="cuda", dtype=DTYPE)
249 + ops.fused_kda_decode(
250 + x=inp.mixed_qkv,
251 + weight=inp.decode_conv1d_weight,
252 + bias=None,
253 + conv_state=inp.conv_state_view(conv_state),
254 + raw_g=inp.g1,
255 + raw_beta=inp.beta,
256 + A_log=inp.A_log,
257 + dt_bias=inp.dt_bias,
258 + state_indices=inp.state_indices,
259 + state=recurrent_state,
260 + out=out,
261 + lower_bound=GATE_LOWER_BOUND,
262 + )
263 +
264 + torch.testing.assert_close(recurrent_state, expected_state, atol=2e-3, rtol=2e-3)
265 + assert torch.isfinite(out).all()
266 +
267 +
268 +@torch.inference_mode()
269 +def test_fused_kda_decode_skips_null_block_padding() -> None:
270 + """A CUDA-graph decode batch is padded with NULL_BLOCK_ID (0).
271 +
272 + gdn_attn.py fills the tail of ``non_spec_state_indices_tensor`` with 0, and
273 + the Triton chain zeroes those rows' output while leaving slot 0 alone. The
274 + fused kernel must do the same rather than read-modify-write slot 0's state
275 + once per padded row.
276 + """
277 + _requires_kernel()
278 + num_real, num_padded, num_heads = 3, 5, 12
279 + inp = KdaDecodeInputs(num_real + num_padded, num_heads, num_slots=12, seed=11)
280 + inp.state_indices[num_real:] = 0
281 +
282 + expected_out, expected_conv, expected_state = _run_triton_chain(inp)
283 + actual_out, actual_conv, actual_state = _run_fused(inp)
284 +
285 + torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2)
286 + assert not actual_out[0, num_real:].any(), "padded rows must produce zeros"
287 + # Slot 0 must be untouched by both paths.
288 + torch.testing.assert_close(actual_conv[0], inp.conv_state[0], atol=0, rtol=0)
289 + torch.testing.assert_close(actual_state[0], inp.recurrent_state[0], atol=0, rtol=0)
290 + torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0)
291 + torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3)
vllm/models/kimi_k3/amd/kda.py
+89 -9
@@ -5,10 +5,12 @@ import torch
5 from einops import rearrange
6 from torch import nn
7
8 +from vllm import _custom_ops as ops
9 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
10 from vllm.config import VllmConfig
11 from vllm.distributed import divide
12 from vllm.forward_context import get_forward_context
13 +from vllm.logger import init_logger
14 from vllm.model_executor.layers.linear import (
15 ColumnParallelLinear,
16 RowParallelLinear,
@@ -38,6 +40,11 @@ from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
40 )
41 from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
42 from vllm.model_executor.utils import set_weight_attrs
43 +from vllm.models.kimi_k3.amd.ops.kda_decode import (
44 + is_fused_kda_decode_supported,
45 + make_decode_conv1d_weight_loader,
46 + make_decode_norm_weight_loader,
47 +)
48 from vllm.models.kimi_k3.amd.ops.third_party.kda import (
49 chunk_kda_with_fused_gate,
50 fused_recurrent_kda,
@@ -47,6 +54,8 @@ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
54 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
55 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
56
57 +logger = init_logger(__name__)
58 +
59
60 class KimiK3DeltaAttention(GatedDeltaNetAttention):
61 def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
@@ -137,16 +146,44 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
146 )
147 self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
148 delattr(self.conv1d.weight, "weight_loader")
140 - set_weight_attrs(
141 - self.conv1d.weight,
142 - {
143 - "weight_loader": _make_fused_conv1d_weight_loader(
144 - [self.projection_size] * 3,
145 - self.tp_size,
146 - self.tp_rank,
147 - )
148 - },
149 + # ROCm can fuse the whole decode step (conv + recurrence + gated norm)
150 + # into one kernel, which wants a width-major fp32 conv weight staged at
151 + # load time. Everything else keeps the [channel, width] layout.
152 + conv_state_dtype, _ = self.get_state_dtype()
153 + decode_conv1d_weight = None
154 + if is_fused_kda_decode_supported(
155 + self.local_num_heads,
156 + self.head_dim,
157 + self.conv_size,
158 + self.num_spec,
159 + vllm_config.model_config.dtype,
160 + conv_state_dtype,
161 + ):
162 + logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.")
163 + decode_conv1d_weight = torch.empty(
164 + 3,
165 + self.conv_size,
166 + self.local_projection_size,
167 + dtype=self.conv1d.weight.dtype,
168 + device=self.conv1d.weight.device,
169 + )
170 + self.register_buffer(
171 + "decode_conv1d_weight", decode_conv1d_weight, persistent=False
172 )
173 + if decode_conv1d_weight is None:
174 + conv1d_weight_loader = _make_fused_conv1d_weight_loader(
175 + [self.projection_size] * 3,
176 + self.tp_size,
177 + self.tp_rank,
178 + )
179 + else:
180 + conv1d_weight_loader = make_decode_conv1d_weight_loader(
181 + [self.projection_size] * 3,
182 + self.tp_size,
183 + self.tp_rank,
184 + decode_conv1d_weight,
185 + )
186 + set_weight_attrs(self.conv1d.weight, {"weight_loader": conv1d_weight_loader})
187
188 self.A_log = nn.Parameter(
189 torch.empty(self.local_num_heads, dtype=torch.float32)
@@ -175,6 +212,22 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
212 )
213
214 self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
215 + decode_norm_weight = None
216 + if decode_conv1d_weight is not None:
217 + # Upcast once at load time; a BF16 norm weight slows the fused
218 + # decode kernel's epilogue.
219 + decode_norm_weight = torch.empty(
220 + self.head_dim,
221 + dtype=torch.float32,
222 + device=self.o_norm.weight.device,
223 + )
224 + if hasattr(self.o_norm.weight, "weight_loader"):
225 + delattr(self.o_norm.weight, "weight_loader")
226 + set_weight_attrs(
227 + self.o_norm.weight,
228 + {"weight_loader": make_decode_norm_weight_loader(decode_norm_weight)},
229 + )
230 + self.register_buffer("decode_norm_weight", decode_norm_weight, persistent=False)
231 self.o_proj = RowParallelLinear(
232 self.projection_size,
233 self.hidden_size,
@@ -280,6 +333,33 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
333 if not is_conv_state_dim_first():
334 conv_state = conv_state.transpose(-1, -2)
335
336 + if (
337 + self.decode_conv1d_weight is not None
338 + and self.decode_norm_weight is not None
339 + and spec_sequence_masks is None
340 + and m.num_prefills == 0
341 + and m.num_decodes > 0
342 + ):
343 + assert non_spec_state_indices_tensor is not None
344 + ops.fused_kda_decode(
345 + x=mixed_qkv,
346 + weight=self.decode_conv1d_weight,
347 + bias=self.conv1d.bias,
348 + conv_state=conv_state,
349 + raw_g=g1,
350 + raw_beta=beta,
351 + A_log=self.A_log,
352 + dt_bias=self.dt_bias,
353 + state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
354 + state=recurrent_state,
355 + out=core_attn_out[:, :num_actual_tokens],
356 + lower_bound=self.gate_lower_bound,
357 + output_gate=g2[:num_actual_tokens],
358 + norm_weight=self.decode_norm_weight,
359 + norm_eps=self.o_norm.eps,
360 + )
361 + return
362 +
363 conv_weights = self.conv1d.weight.view(
364 self.conv1d.weight.size(0), self.conv1d.weight.size(2)
365 )
vllm/models/kimi_k3/amd/ops/kda_decode.py new
+99
@@ -0,0 +1,99 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ROCm entry points for the fused Kimi-K3 KDA decode kernel.
4 +
5 +The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu``
6 +replaces, for a pure non-speculative decode batch, the three Triton launches
7 +and two copies the AMD KDA layer otherwise runs per layer: the packed causal
8 +conv1d update, the recurrent delta-rule step, and the gated output RMSNorm.
9 +
10 +The kernel wants a width-major conv weight and an fp32 norm weight, so both are
11 +staged once at load time by the weight loaders below.
12 +"""
13 +
14 +from collections.abc import Callable
15 +
16 +import torch
17 +
18 +from vllm.logger import init_logger
19 +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
20 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21 +
22 +logger = init_logger(__name__)
23 +
24 +# Head counts the kernel is instantiated for (Kimi-K3 has 96 KDA heads, so this
25 +# covers TP1/2/4/8).
26 +SUPPORTED_NUM_HEADS = (12, 24, 48, 96)
27 +
28 +
29 +def is_fused_kda_decode_supported(
30 + num_heads: int,
31 + head_dim: int,
32 + conv_width: int,
33 + num_spec: int,
34 + input_dtype: torch.dtype,
35 + conv_state_dtype: torch.dtype,
36 +) -> bool:
37 + """Whether the fused decode kernel can serve this layer on this device."""
38 + from vllm.platforms.rocm import on_gfx950
39 +
40 + if (
41 + num_heads not in SUPPORTED_NUM_HEADS
42 + or head_dim != 128
43 + or conv_width != 4
44 + or num_spec != 0
45 + or input_dtype != torch.bfloat16
46 + or conv_state_dtype != torch.bfloat16
47 + or is_conv_state_dim_first()
48 + or not hasattr(torch.ops._C, "fused_kda_decode")
49 + ):
50 + return False
51 + # TODO: Verify on other archs; only measured on gfx950 for now
52 + return on_gfx950()
53 +
54 +
55 +def make_decode_conv1d_weight_loader(
56 + dims: list[int],
57 + tp_size: int,
58 + tp_rank: int,
59 + decode_conv1d_weight: torch.Tensor | None,
60 +) -> Callable[..., None]:
61 + """Load the packed conv1d weight, mirroring a width-major fp32 copy.
62 +
63 + The fused kernel indexes the weight as ``[qkv, width, channel]`` so the
64 + channel dimension is the contiguous one; the Triton prefill and fallback
65 + decode kernels keep the ``[channel, width]`` layout.
66 + """
67 + sharded_dims = [dim // tp_size for dim in dims]
68 +
69 + def weight_loader(
70 + param: torch.Tensor,
71 + loaded_weight: torch.Tensor,
72 + loaded_shard_id: int,
73 + ) -> None:
74 + if loaded_weight.dim() == 2:
75 + loaded_weight = loaded_weight.unsqueeze(1)
76 + shard_size = sharded_dims[loaded_shard_id]
77 + source_start = tp_rank * shard_size
78 + target_start = sum(sharded_dims[:loaded_shard_id])
79 + loaded_shard = loaded_weight[source_start : source_start + shard_size]
80 + param.data[target_start : target_start + shard_size].copy_(loaded_shard)
81 + if decode_conv1d_weight is not None and not param.is_meta:
82 + decode_conv1d_weight[loaded_shard_id].copy_(
83 + loaded_shard.squeeze(1).transpose(0, 1)
84 + )
85 +
86 + return weight_loader
87 +
88 +
89 +def make_decode_norm_weight_loader(
90 + decode_norm_weight: torch.Tensor,
91 +) -> Callable[..., None]:
92 + """Load the gated-norm weight, mirroring an fp32 copy for the kernel."""
93 +
94 + def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
95 + default_weight_loader(param, loaded_weight)
96 + if not param.is_meta:
97 + decode_norm_weight.copy_(param.data)
98 +
99 + return weight_loader