@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1be3628367

[Kernel][Perf] Add fused CUDA post-conv MTP decode kernel for Qwen3.5 GDN (#51674)

Signed-off-by: Jie Fang <jief@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Jie Fang committed Aug 14, 2026 at 11:44 UTC 1be36283678a9a94fc8fdaad6c95c2896d6b4015
12 files changed +1456 -7
CMakeLists.txt
+18
@@ -1118,6 +1118,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1118 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
1119 cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS
1120 "9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
1121 + cuda_archs_loose_intersection(FUSED_GDN_DECODE_ARCHS
1122 + "8.0;8.6;8.9;9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
1123 endif()
1124 if(FUSED_KDA_DECODE_ARCHS)
1125 set(FUSED_KDA_DECODE_SRC
@@ -1131,6 +1133,18 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1133 message(STATUS
1134 "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}")
1135 endif()
1136 + if(FUSED_GDN_DECODE_ARCHS)
1137 + set(FUSED_GDN_DECODE_SRC
1138 + "csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu")
1139 + set_gencode_flags_for_srcs(
1140 + SRCS "${FUSED_GDN_DECODE_SRC}"
1141 + CUDA_ARCHS "${FUSED_GDN_DECODE_ARCHS}")
1142 + set_property(SOURCE ${FUSED_GDN_DECODE_SRC} APPEND PROPERTY
1143 + COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>")
1144 + list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_GDN_DECODE_SRC}")
1145 + message(STATUS
1146 + "Building fused GDN decode for archs: ${FUSED_GDN_DECODE_ARCHS}")
1147 + endif()
1148
1149 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
1150 cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
@@ -1213,6 +1227,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1227 target_compile_definitions(_C_stable_libtorch PRIVATE
1228 VLLM_ENABLE_FUSED_KDA_DECODE=1)
1229 endif()
1230 + if(FUSED_GDN_DECODE_ARCHS)
1231 + target_compile_definitions(_C_stable_libtorch PRIVATE
1232 + VLLM_ENABLE_FUSED_GDN_DECODE=1)
1233 + endif()
1234 if(KIMI_K3_ATTN_RES_ARCHS)
1235 target_compile_definitions(_C_stable_libtorch PRIVATE
1236 VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu new
+557
@@ -0,0 +1,557 @@
1 +/*
2 + * SPDX-License-Identifier: Apache-2.0
3 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 + */
5 +
6 +#include <cstdint>
7 +#include <cuda_bf16.h>
8 +#include <cuda_fp16.h>
9 +#include <cuda_runtime.h>
10 +
11 +#include "../torch_utils.h"
12 +#include "../../cuda_compat.h"
13 +
14 +namespace {
15 +
16 +template <typename StateT>
17 +__device__ __forceinline__ void cp_async_16b(StateT* smem_ptr,
18 + const StateT* gmem_ptr) {
19 + const uint32_t smem_addr =
20 + static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
21 + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"
22 + :
23 + : "r"(smem_addr), "l"(gmem_ptr));
24 +}
25 +
26 +__device__ __forceinline__ void cp_async_commit() {
27 + asm volatile("cp.async.commit_group;\n" ::);
28 +}
29 +
30 +__device__ __forceinline__ void cp_async_wait_all() {
31 + asm volatile("cp.async.wait_all;\n" ::: "memory");
32 +}
33 +
34 +template <typename StateT, int ChunkV, int DimK, int Stages>
35 +__device__ __forceinline__ void copy_state_chunk(StateT* shared_state,
36 + const StateT* state, int chunk,
37 + int thread, int threads) {
38 + constexpr int kElementsPerCopy = 16 / sizeof(StateT);
39 + constexpr int kCopiesPerChunk = ChunkV * DimK / kElementsPerCopy;
40 + const int stage = chunk % Stages;
41 + for (int copy = thread; copy < kCopiesPerChunk; copy += threads) {
42 + const int element = copy * kElementsPerCopy;
43 + cp_async_16b(shared_state + stage * ChunkV * DimK + element,
44 + state + chunk * ChunkV * DimK + element);
45 + }
46 + cp_async_commit();
47 +}
48 +
49 +template <typename StateT>
50 +__device__ __forceinline__ float4 load_state4(const StateT* state);
51 +
52 +template <>
53 +__device__ __forceinline__ float4 load_state4<float>(const float* state) {
54 + return *reinterpret_cast<const float4*>(state);
55 +}
56 +
57 +template <>
58 +__device__ __forceinline__ float4
59 +load_state4<__nv_bfloat16>(const __nv_bfloat16* state) {
60 + const __nv_bfloat162 lo = *reinterpret_cast<const __nv_bfloat162*>(state);
61 + const __nv_bfloat162 hi = *reinterpret_cast<const __nv_bfloat162*>(state + 2);
62 + return make_float4(__bfloat162float(lo.x), __bfloat162float(lo.y),
63 + __bfloat162float(hi.x), __bfloat162float(hi.y));
64 +}
65 +
66 +template <typename StateT>
67 +__device__ __forceinline__ void store_state4(StateT* state, float4 value);
68 +
69 +template <>
70 +__device__ __forceinline__ void store_state4<float>(float* state,
71 + float4 value) {
72 + *reinterpret_cast<float4*>(state) = value;
73 +}
74 +
75 +template <>
76 +__device__ __forceinline__ void store_state4<__nv_bfloat16>(
77 + __nv_bfloat16* state, float4 value) {
78 + *reinterpret_cast<__nv_bfloat162*>(state) =
79 + __floats2bfloat162_rn(value.x, value.y);
80 + *reinterpret_cast<__nv_bfloat162*>(state + 2) =
81 + __floats2bfloat162_rn(value.z, value.w);
82 +}
83 +
84 +constexpr int kDimK = 128;
85 +constexpr int kDimV = 128;
86 +constexpr int kThreads = 256;
87 +constexpr int kWarps = kThreads / 32;
88 +constexpr int kChunkV = 32;
89 +constexpr int kNumChunks = kDimV / kChunkV;
90 +constexpr int kRowsPerWarp = kChunkV / kWarps;
91 +constexpr int kMaxMtpTokens = 8;
92 +constexpr int kDtBiasFloat32 = 0;
93 +constexpr int kDtBiasBFloat16 = 1;
94 +constexpr int kDtBiasFloat16 = 2;
95 +
96 +struct GdnDecodeStrides {
97 + int64_t mixed_row;
98 + int64_t a_row;
99 + int64_t b_row;
100 + int64_t gate_row;
101 + int64_t state_slot;
102 +};
103 +
104 +__device__ __forceinline__ float sigmoid_fast(float x) {
105 + return 1.0f / (1.0f + __expf(-x));
106 +}
107 +
108 +__device__ __forceinline__ float silu_fast(float x) {
109 + return x * sigmoid_fast(x);
110 +}
111 +
112 +__device__ __forceinline__ float softplus_fast(float x) {
113 + return x > 20.0f ? x : log1pf(__expf(x));
114 +}
115 +
116 +__device__ __forceinline__ float load_dt_bias(const void* dt_bias, int head,
117 + int dt_bias_type) {
118 + if (dt_bias_type == kDtBiasBFloat16) {
119 + return __bfloat162float(static_cast<const __nv_bfloat16*>(dt_bias)[head]);
120 + }
121 + if (dt_bias_type == kDtBiasFloat16) {
122 + return __half2float(static_cast<const __half*>(dt_bias)[head]);
123 + }
124 + return static_cast<const float*>(dt_bias)[head];
125 +}
126 +
127 +__device__ __forceinline__ float warp_reduce_sum(float value) {
128 +#pragma unroll
129 + for (int offset = 16; offset > 0; offset >>= 1) {
130 + value += __shfl_xor_sync(0xffffffffu, value, offset);
131 + }
132 + return value;
133 +}
134 +
135 +struct Sum2 {
136 + float x;
137 + float y;
138 +};
139 +
140 +__device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) {
141 +#pragma unroll
142 + for (int offset = 16; offset > 0; offset >>= 1) {
143 + x += __shfl_xor_sync(0xffffffffu, x, offset);
144 + y += __shfl_xor_sync(0xffffffffu, y, offset);
145 + }
146 + return {x, y};
147 +}
148 +
149 +template <typename StateT>
150 +__global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel(
151 + const __nv_bfloat16* __restrict__ mixed_qkv,
152 + const __nv_bfloat16* __restrict__ a, const __nv_bfloat16* __restrict__ b,
153 + const float* __restrict__ a_log, const void* __restrict__ dt_bias,
154 + const int* __restrict__ state_indices, const int* __restrict__ cu_seqlens,
155 + const int* __restrict__ num_accepted_tokens, StateT* __restrict__ state,
156 + const __nv_bfloat16* __restrict__ output_gate,
157 + const void* __restrict__ norm_weight, __nv_bfloat16* __restrict__ out,
158 + int H, int HV, int state_indices_width, int dt_bias_type,
159 + bool norm_weight_is_bf16, float scale, float norm_eps,
160 + GdnDecodeStrides strides) {
161 + const int request = blockIdx.x;
162 + const int value_head = blockIdx.y;
163 + const int tid = threadIdx.x;
164 + const int lane = tid & 31;
165 + const int warp = tid >> 5;
166 + const int bos = cu_seqlens[request];
167 + const int eos = cu_seqlens[request + 1];
168 + const int num_tokens = eos - bos;
169 + if (num_tokens <= 0) {
170 + return;
171 + }
172 +
173 + const int accepted = num_accepted_tokens[request];
174 + const int source_slot =
175 + accepted > 0 && accepted <= state_indices_width
176 + ? state_indices[request * state_indices_width + accepted - 1]
177 + : 0;
178 + if (source_slot <= 0 || num_tokens > kMaxMtpTokens) {
179 + for (int linear = tid; linear < num_tokens * kDimV; linear += kThreads) {
180 + const int token = bos + linear / kDimV;
181 + const int value = linear % kDimV;
182 + const int64_t out_offset =
183 + (static_cast<int64_t>(token) * HV + value_head) * kDimV + value;
184 + out[out_offset] = __float2bfloat16(0.0f);
185 + }
186 + return;
187 + }
188 +
189 + const int key_head = value_head / 8;
190 + __shared__ StateT shared_state[2][kChunkV][kDimK];
191 + __shared__ float shared_q[kMaxMtpTokens][kDimK];
192 + __shared__ float shared_k[kMaxMtpTokens][kDimK];
193 + __shared__ __nv_bfloat16 shared_v[kMaxMtpTokens][kDimV];
194 + __shared__ __nv_bfloat16 shared_out[kMaxMtpTokens][kDimV];
195 + __shared__ float shared_decay[kMaxMtpTokens];
196 + __shared__ float shared_beta[kMaxMtpTokens];
197 +
198 + StateT* source_state =
199 + state + static_cast<int64_t>(source_slot) * strides.state_slot +
200 + value_head * kDimV * kDimK;
201 + copy_state_chunk<StateT, kChunkV, kDimK, 2>(&shared_state[0][0][0],
202 + source_state, 0, tid, kThreads);
203 +
204 + if (warp < num_tokens) {
205 + const int t = warp;
206 + const int token = bos + t;
207 + const int64_t mixed_base = static_cast<int64_t>(token) * strides.mixed_row;
208 + float q_values[4];
209 + float k_values[4];
210 + float q_square = 0.0f;
211 + float k_square = 0.0f;
212 +#pragma unroll
213 + for (int i = 0; i < 4; ++i) {
214 + const int dim = lane + i * 32;
215 + q_values[i] =
216 + __bfloat162float(mixed_qkv[mixed_base + key_head * kDimK + dim]);
217 + k_values[i] = __bfloat162float(
218 + mixed_qkv[mixed_base + H * kDimK + key_head * kDimK + dim]);
219 + shared_v[t][dim] =
220 + mixed_qkv[mixed_base + 2 * H * kDimK + value_head * kDimV + dim];
221 + q_square += q_values[i] * q_values[i];
222 + k_square += k_values[i] * k_values[i];
223 + }
224 + const Sum2 qk_sums = warp_reduce_sum_pair(q_square, k_square);
225 + const float q_scale = __shfl_sync(
226 + 0xffffffffu, lane == 0 ? rsqrtf(qk_sums.x + 1.0e-6f) * scale : 0.0f, 0);
227 + const float k_scale = __shfl_sync(
228 + 0xffffffffu, lane == 0 ? rsqrtf(qk_sums.y + 1.0e-6f) : 0.0f, 0);
229 +#pragma unroll
230 + for (int i = 0; i < 4; ++i) {
231 + const int dim = lane + i * 32;
232 + shared_q[t][dim] = q_values[i] * q_scale;
233 + shared_k[t][dim] = k_values[i] * k_scale;
234 + }
235 + if (lane == 0) {
236 + const float a_value = __bfloat162float(
237 + a[static_cast<int64_t>(token) * strides.a_row + value_head]);
238 + const float b_value = __bfloat162float(
239 + b[static_cast<int64_t>(token) * strides.b_row + value_head]);
240 + const float g = -__expf(a_log[value_head]) *
241 + softplus_fast(a_value + load_dt_bias(dt_bias, value_head,
242 + dt_bias_type));
243 + shared_decay[t] = __expf(g);
244 + shared_beta[t] = sigmoid_fast(b_value);
245 + }
246 + }
247 + __syncthreads();
248 +
249 + const int k_base = lane * 4;
250 + int rows[kRowsPerWarp];
251 +#pragma unroll
252 + for (int row = 0; row < kRowsPerWarp; ++row) {
253 + rows[row] = warp + row * kWarps;
254 + }
255 +
256 +#pragma unroll
257 + for (int chunk = 0; chunk < kNumChunks; ++chunk) {
258 + cp_async_wait_all();
259 + __syncthreads();
260 + if (chunk + 1 < kNumChunks) {
261 + copy_state_chunk<StateT, kChunkV, kDimK, 2>(
262 + &shared_state[0][0][0], source_state, chunk + 1, tid, kThreads);
263 + }
264 +
265 + float h[kRowsPerWarp][4];
266 +#pragma unroll
267 + for (int row = 0; row < kRowsPerWarp; ++row) {
268 + const float4 state_value =
269 + load_state4(&shared_state[chunk & 1][rows[row]][k_base]);
270 + h[row][0] = state_value.x;
271 + h[row][1] = state_value.y;
272 + h[row][2] = state_value.z;
273 + h[row][3] = state_value.w;
274 + }
275 +
276 + for (int t = 0; t < num_tokens; ++t) {
277 + const float4 q4 = *reinterpret_cast<const float4*>(&shared_q[t][k_base]);
278 + const float4 k4 = *reinterpret_cast<const float4*>(&shared_k[t][k_base]);
279 + const float q_values[4] = {q4.x, q4.y, q4.z, q4.w};
280 + const float k_values[4] = {k4.x, k4.y, k4.z, k4.w};
281 +
282 + float dot_hk[kRowsPerWarp] = {0.0f, 0.0f, 0.0f, 0.0f};
283 +#pragma unroll
284 + for (int row = 0; row < kRowsPerWarp; ++row) {
285 +#pragma unroll
286 + for (int i = 0; i < 4; ++i) {
287 + h[row][i] *= shared_decay[t];
288 + dot_hk[row] += h[row][i] * k_values[i];
289 + }
290 + }
291 + const Sum2 dot_hk_01 = warp_reduce_sum_pair(dot_hk[0], dot_hk[1]);
292 + const Sum2 dot_hk_23 = warp_reduce_sum_pair(dot_hk[2], dot_hk[3]);
293 + const float reduced_hk[kRowsPerWarp] = {dot_hk_01.x, dot_hk_01.y,
294 + dot_hk_23.x, dot_hk_23.y};
295 +
296 + float dot_hq[kRowsPerWarp] = {0.0f, 0.0f, 0.0f, 0.0f};
297 +#pragma unroll
298 + for (int row = 0; row < kRowsPerWarp; ++row) {
299 + const int value = chunk * kChunkV + rows[row];
300 + const float delta =
301 + (__bfloat162float(shared_v[t][value]) - reduced_hk[row]) *
302 + shared_beta[t];
303 +#pragma unroll
304 + for (int i = 0; i < 4; ++i) {
305 + h[row][i] += k_values[i] * delta;
306 + dot_hq[row] += h[row][i] * q_values[i];
307 + }
308 + }
309 + const Sum2 dot_hq_01 = warp_reduce_sum_pair(dot_hq[0], dot_hq[1]);
310 + const Sum2 dot_hq_23 = warp_reduce_sum_pair(dot_hq[2], dot_hq[3]);
311 + if (lane == 0) {
312 + shared_out[t][chunk * kChunkV + rows[0]] =
313 + __float2bfloat16(dot_hq_01.x);
314 + shared_out[t][chunk * kChunkV + rows[1]] =
315 + __float2bfloat16(dot_hq_01.y);
316 + shared_out[t][chunk * kChunkV + rows[2]] =
317 + __float2bfloat16(dot_hq_23.x);
318 + shared_out[t][chunk * kChunkV + rows[3]] =
319 + __float2bfloat16(dot_hq_23.y);
320 + }
321 +
322 + const int destination_slot =
323 + state_indices[request * state_indices_width + t];
324 + if (destination_slot > 0) {
325 + StateT* destination_state =
326 + state +
327 + static_cast<int64_t>(destination_slot) * strides.state_slot +
328 + value_head * kDimV * kDimK;
329 +#pragma unroll
330 + for (int row = 0; row < kRowsPerWarp; ++row) {
331 + const int value = chunk * kChunkV + rows[row];
332 + store_state4(destination_state + value * kDimK + k_base,
333 + make_float4(h[row][0], h[row][1], h[row][2], h[row][3]));
334 + }
335 + }
336 + }
337 + }
338 + __syncthreads();
339 +
340 + if (warp < num_tokens) {
341 + const int t = warp;
342 + float output_values[4];
343 + float sum_square = 0.0f;
344 +#pragma unroll
345 + for (int i = 0; i < 4; ++i) {
346 + const int value = lane + i * 32;
347 + output_values[i] = __bfloat162float(shared_out[t][value]);
348 + sum_square += output_values[i] * output_values[i];
349 + }
350 + sum_square = warp_reduce_sum(sum_square);
351 + const float rstd =
352 + rsqrtf(sum_square / static_cast<float>(kDimV) + norm_eps);
353 + const int token = bos + t;
354 +#pragma unroll
355 + for (int i = 0; i < 4; ++i) {
356 + const int value = lane + i * 32;
357 + const float gate = silu_fast(__bfloat162float(
358 + output_gate[static_cast<int64_t>(token) * strides.gate_row +
359 + value_head * kDimV + value]));
360 + const float weight =
361 + norm_weight_is_bf16
362 + ? __bfloat162float(
363 + static_cast<const __nv_bfloat16*>(norm_weight)[value])
364 + : static_cast<const float*>(norm_weight)[value];
365 + const int64_t out_offset =
366 + (static_cast<int64_t>(token) * HV + value_head) * kDimV + value;
367 + out[out_offset] =
368 + __float2bfloat16(output_values[i] * rstd * weight * gate);
369 + }
370 + }
371 +}
372 +
373 +template <typename StateT>
374 +void launch_gdn_decode_post_conv_mtp(
375 + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a_log,
376 + torch::stable::Tensor const& dt_bias,
377 + torch::stable::Tensor const& state_indices,
378 + torch::stable::Tensor const& cu_seqlens,
379 + torch::stable::Tensor const& num_accepted_tokens,
380 + torch::stable::Tensor& state, torch::stable::Tensor const& norm_weight,
381 + torch::stable::Tensor& out, const __nv_bfloat16* a, const __nv_bfloat16* b,
382 + const __nv_bfloat16* output_gate, int num_key_heads, int num_value_heads,
383 + double scale, double norm_eps, GdnDecodeStrides strides) {
384 + using torch::headeronly::ScalarType;
385 +
386 + const auto dt_bias_scalar_type = dt_bias.scalar_type();
387 + const int dt_bias_type =
388 + dt_bias_scalar_type == ScalarType::Float
389 + ? kDtBiasFloat32
390 + : (dt_bias_scalar_type == ScalarType::BFloat16 ? kDtBiasBFloat16
391 + : kDtBiasFloat16);
392 + torch::stable::accelerator::DeviceGuard const device_guard(
393 + mixed_qkv.get_device_index());
394 + const cudaStream_t stream =
395 + get_current_cuda_stream(mixed_qkv.get_device_index());
396 + const int num_requests = static_cast<int>(state_indices.size(0));
397 + const dim3 grid(num_requests, num_value_heads);
398 + gdn_decode_post_conv_mtp_kernel<StateT><<<grid, kThreads, 0, stream>>>(
399 + static_cast<const __nv_bfloat16*>(mixed_qkv.data_ptr()), a, b,
400 + static_cast<const float*>(a_log.data_ptr()), dt_bias.data_ptr(),
401 + static_cast<const int*>(state_indices.data_ptr()),
402 + static_cast<const int*>(cu_seqlens.data_ptr()),
403 + static_cast<const int*>(num_accepted_tokens.data_ptr()),
404 + static_cast<StateT*>(state.data_ptr()), output_gate,
405 + norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()),
406 + num_key_heads, num_value_heads, static_cast<int>(state_indices.size(1)),
407 + dt_bias_type, norm_weight.scalar_type() == ScalarType::BFloat16,
408 + static_cast<float>(scale), static_cast<float>(norm_eps), strides);
409 + const cudaError_t error = cudaGetLastError();
410 + STD_TORCH_CHECK(error == cudaSuccess,
411 + "GDN decode MTP post-conv kernel launch failed: ",
412 + cudaGetErrorString(error));
413 +}
414 +
415 +} // namespace
416 +
417 +void fused_gdn_decode_post_conv_mtp(
418 + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a,
419 + torch::stable::Tensor const& b, torch::stable::Tensor const& a_log,
420 + torch::stable::Tensor const& dt_bias,
421 + torch::stable::Tensor const& state_indices,
422 + torch::stable::Tensor const& cu_seqlens,
423 + torch::stable::Tensor const& num_accepted_tokens,
424 + torch::stable::Tensor& state, torch::stable::Tensor const& output_gate,
425 + torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out,
426 + double scale, double norm_eps) {
427 + using torch::headeronly::ScalarType;
428 +
429 + STD_TORCH_CHECK(
430 + mixed_qkv.is_cuda() && mixed_qkv.scalar_type() == ScalarType::BFloat16,
431 + "mixed_qkv must be a CUDA bfloat16 tensor");
432 + STD_TORCH_CHECK(a.is_cuda() && a.scalar_type() == ScalarType::BFloat16,
433 + "a must be a CUDA bfloat16 tensor");
434 + STD_TORCH_CHECK(b.is_cuda() && b.scalar_type() == ScalarType::BFloat16,
435 + "b must be a CUDA bfloat16 tensor");
436 + STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float,
437 + "A_log must be a CUDA float32 tensor");
438 + const auto dt_bias_scalar_type = dt_bias.scalar_type();
439 + STD_TORCH_CHECK(
440 + dt_bias.is_cuda() && (dt_bias_scalar_type == ScalarType::Float ||
441 + dt_bias_scalar_type == ScalarType::BFloat16 ||
442 + dt_bias_scalar_type == ScalarType::Half),
443 + "dt_bias must be a CUDA float32, bfloat16, or float16 tensor");
444 + STD_TORCH_CHECK(
445 + state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int,
446 + "state_indices must be a CUDA int32 tensor");
447 + STD_TORCH_CHECK(
448 + cu_seqlens.is_cuda() && cu_seqlens.scalar_type() == ScalarType::Int,
449 + "cu_seqlens must be a CUDA int32 tensor");
450 + STD_TORCH_CHECK(num_accepted_tokens.is_cuda() &&
451 + num_accepted_tokens.scalar_type() == ScalarType::Int,
452 + "num_accepted_tokens must be a CUDA int32 tensor");
453 + const auto state_scalar_type = state.scalar_type();
454 + STD_TORCH_CHECK(
455 + state.is_cuda() && (state_scalar_type == ScalarType::Float ||
456 + state_scalar_type == ScalarType::BFloat16),
457 + "state must be a CUDA float32 or bfloat16 tensor");
458 + STD_TORCH_CHECK(output_gate.is_cuda() &&
459 + output_gate.scalar_type() == ScalarType::BFloat16,
460 + "output_gate must be a CUDA bfloat16 tensor");
461 + STD_TORCH_CHECK(norm_weight.is_cuda() &&
462 + (norm_weight.scalar_type() == ScalarType::Float ||
463 + norm_weight.scalar_type() == ScalarType::BFloat16),
464 + "norm_weight must be a CUDA float32 or bfloat16 tensor");
465 + STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16,
466 + "out must be a CUDA bfloat16 tensor");
467 +
468 + STD_TORCH_CHECK(mixed_qkv.dim() == 2,
469 + "mixed_qkv must have shape [L, 2 * H * 128 + HV * 128]");
470 + const int num_tokens = static_cast<int>(mixed_qkv.size(0));
471 + STD_TORCH_CHECK(num_tokens > 0,
472 + "GDN decode MTP fusion requires at least one token");
473 + STD_TORCH_CHECK(
474 + state.dim() == 4 && state.size(2) == kDimV && state.size(3) == kDimK,
475 + "state must have shape [slots, HV, 128, 128]");
476 + const int num_value_heads = static_cast<int>(state.size(1));
477 + const int64_t key_width =
478 + mixed_qkv.size(1) - static_cast<int64_t>(num_value_heads) * kDimV;
479 + STD_TORCH_CHECK(key_width > 0 && key_width % (2 * kDimK) == 0,
480 + "mixed_qkv width is inconsistent with state");
481 + const int num_key_heads = static_cast<int>(key_width / (2 * kDimK));
482 + STD_TORCH_CHECK(num_value_heads == 8 * num_key_heads,
483 + "GDN decode MTP fusion requires HV/H=8");
484 +
485 + STD_TORCH_CHECK(state_indices.dim() == 2 && state_indices.size(0) > 0 &&
486 + state_indices.size(1) > 0 &&
487 + state_indices.size(1) <= kMaxMtpTokens,
488 + "state_indices must have shape [N, S] with 1 <= S <= 8");
489 + const int num_requests = static_cast<int>(state_indices.size(0));
490 + STD_TORCH_CHECK(
491 + cu_seqlens.dim() == 1 && cu_seqlens.numel() == num_requests + 1,
492 + "cu_seqlens must have N + 1 elements");
493 + STD_TORCH_CHECK(num_accepted_tokens.dim() == 1 &&
494 + num_accepted_tokens.numel() == num_requests,
495 + "num_accepted_tokens must have N elements");
496 + STD_TORCH_CHECK(
497 + a.dim() == 2 && a.size(0) == num_tokens && a.size(1) == num_value_heads,
498 + "a must have shape [L, HV]");
499 + STD_TORCH_CHECK(
500 + b.dim() == 2 && b.size(0) == num_tokens && b.size(1) == num_value_heads,
501 + "b must have shape [L, HV]");
502 + STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_value_heads,
503 + "A_log must be contiguous with HV elements");
504 + STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == num_value_heads,
505 + "dt_bias must be contiguous with HV elements");
506 + STD_TORCH_CHECK(state_indices.is_contiguous(),
507 + "state_indices must be contiguous");
508 + STD_TORCH_CHECK(cu_seqlens.is_contiguous(), "cu_seqlens must be contiguous");
509 + STD_TORCH_CHECK(num_accepted_tokens.is_contiguous(),
510 + "num_accepted_tokens must be contiguous");
511 + STD_TORCH_CHECK(output_gate.dim() == 3 && output_gate.size(0) == num_tokens &&
512 + output_gate.size(1) == num_value_heads &&
513 + output_gate.size(2) == kDimV,
514 + "output_gate must have shape [L, HV, 128]");
515 + STD_TORCH_CHECK(norm_weight.is_contiguous() && norm_weight.numel() == kDimV,
516 + "norm_weight must be contiguous with 128 elements");
517 + STD_TORCH_CHECK(out.dim() == 3 && out.size(0) == num_tokens &&
518 + out.size(1) == num_value_heads && out.size(2) == kDimV,
519 + "out must have shape [L, HV, 128]");
520 + STD_TORCH_CHECK(mixed_qkv.stride(1) == 1,
521 + "mixed_qkv channels must be contiguous");
522 + STD_TORCH_CHECK(a.stride(1) == 1 && b.stride(1) == 1,
523 + "a and b heads must be contiguous");
524 + STD_TORCH_CHECK(state.stride(0) >= num_value_heads * kDimV * kDimK &&
525 + state.stride(1) == kDimV * kDimK &&
526 + state.stride(2) == kDimK && state.stride(3) == 1,
527 + "state must have contiguous [HV, 128, 128] slot contents");
528 + const int state_elements_per_copy =
529 + state_scalar_type == ScalarType::Float ? 4 : 8;
530 + STD_TORCH_CHECK(reinterpret_cast<uintptr_t>(state.data_ptr()) % 16 == 0 &&
531 + state.stride(0) % state_elements_per_copy == 0,
532 + "state slots must preserve 16-byte alignment");
533 + STD_TORCH_CHECK(output_gate.stride(2) == 1 && output_gate.stride(1) == kDimV,
534 + "output_gate head rows must be contiguous");
535 + STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
536 + STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative");
537 +
538 + const GdnDecodeStrides strides{mixed_qkv.stride(0), a.stride(0), b.stride(0),
539 + output_gate.stride(0), state.stride(0)};
540 + const auto* a_ptr = static_cast<const __nv_bfloat16*>(a.data_ptr());
541 + const auto* b_ptr = static_cast<const __nv_bfloat16*>(b.data_ptr());
542 + const auto* output_gate_ptr =
543 + static_cast<const __nv_bfloat16*>(output_gate.data_ptr());
544 + if (state_scalar_type == ScalarType::Float) {
545 + launch_gdn_decode_post_conv_mtp<float>(
546 + mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens,
547 + num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr,
548 + output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps,
549 + strides);
550 + } else {
551 + launch_gdn_decode_post_conv_mtp<__nv_bfloat16>(
552 + mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens,
553 + num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr,
554 + output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps,
555 + strides);
556 + }
557 +}
csrc/libtorch_stable/ops.h
+12
@@ -400,6 +400,18 @@ void fused_kda_decode(
400 torch::stable::Tensor& out, std::optional<double> lower_bound,
401 std::optional<torch::stable::Tensor> output_gate,
402 std::optional<torch::stable::Tensor> norm_weight, double norm_eps);
403 +
404 +void fused_gdn_decode_post_conv_mtp(
405 + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a,
406 + torch::stable::Tensor const& b, torch::stable::Tensor const& a_log,
407 + torch::stable::Tensor const& dt_bias,
408 + torch::stable::Tensor const& state_indices,
409 + torch::stable::Tensor const& cu_seqlens,
410 + torch::stable::Tensor const& num_accepted_tokens,
411 + torch::stable::Tensor& state, torch::stable::Tensor const& output_gate,
412 + torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out,
413 + double scale, double norm_eps);
414 +
415 #endif
416
417 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
csrc/libtorch_stable/torch_bindings.cpp
+14
@@ -534,6 +534,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
534 "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()");
535 #endif
536
537 +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE
538 + ops.def(
539 + "fused_gdn_decode_post_conv_mtp("
540 + "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, "
541 + "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, "
542 + "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, "
543 + "float scale, float norm_eps=1e-5) -> ()");
544 +#endif
545 +
546 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
547 ops.def(
548 "kimi_k3_attn_res("
@@ -797,6 +806,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
806 ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode));
807 #endif
808
809 +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE
810 + ops.impl("fused_gdn_decode_post_conv_mtp",
811 + TORCH_BOX(&fused_gdn_decode_post_conv_mtp));
812 +#endif
813 +
814 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
815 ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
816 #endif
tests/compile/test_config.py
+5
@@ -252,6 +252,11 @@ def test_splitting_ops_dynamic():
252 # populated when the engine decides to use piecewise compilation.
253 assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE
254 assert config.compilation_config.splitting_ops_contain_attention()
255 + splitting_ops = config.compilation_config.splitting_ops
256 + assert splitting_ops is not None
257 + assert {
258 + "vllm::qwen_gdn_attention_core_fused_norm_packed",
259 + } <= set(splitting_ops)
260
261 # When use_inductor_graph_partition=True
262 config = VllmConfig(
tests/kernels/mamba/test_gdn_fused_mtp.py new
+345
@@ -0,0 +1,345 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Model-path tests for fused Qwen3.5 GDN MTP decode."""
4 +
5 +from __future__ import annotations
6 +
7 +import types
8 +from unittest.mock import patch
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.platforms import current_platform
14 +
15 +if not (current_platform.is_cuda() and current_platform.has_device_capability(80)):
16 + pytest.skip(
17 + reason="Fused GDN MTP decode requires CUDA compute capability 8.0+.",
18 + allow_module_level=True,
19 + )
20 +
21 +from tests.v1.attention.utils import ( # noqa: E402
22 + BatchSpec,
23 + create_common_attn_metadata,
24 + create_vllm_config,
25 +)
26 +from vllm.config import SpeculativeConfig, set_current_vllm_config # noqa: E402
27 +from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: E402
28 +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa: E402
29 + ChunkGatedDeltaRule,
30 + QwenGatedDeltaNetAttention,
31 +)
32 +from vllm.model_executor.layers.mamba.mamba_utils import ( # noqa: E402
33 + MambaStateShapeCalculator,
34 +)
35 +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( # noqa: E402
36 + rmsnorm_fn,
37 +)
38 +from vllm.utils.torch_utils import _encode_layer_name # noqa: E402
39 +from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402
40 + GDNAttentionMetadataBuilder,
41 +)
42 +from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402
43 +
44 +NUM_SPEC = 3
45 +SPEC_TOKENS = NUM_SPEC + 1
46 +H = 1
47 +HV = 8
48 +K = 128
49 +V = 128
50 +CONV_KERNEL = 4
51 +CONV_DIM = 2 * H * K + HV * V
52 +BLOCK_SIZE = 16
53 +PREFIX = "model.layers.0.linear_attn"
54 +EPS = 1e-6
55 +
56 +
57 +class _TestGatedNorm:
58 + def __init__(self, weight: torch.Tensor) -> None:
59 + self.weight = weight
60 + self.bias = None
61 + self.eps = EPS
62 + self.group_size = None
63 + self.norm_before_gate = True
64 + self.activation = "silu"
65 +
66 + def __call__(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
67 + return rmsnorm_fn(
68 + x,
69 + self.weight,
70 + None,
71 + z=z,
72 + eps=self.eps,
73 + norm_before_gate=True,
74 + activation=self.activation,
75 + )
76 +
77 +
78 +def _make_vllm_config():
79 + config = create_vllm_config(
80 + model_name="Qwen/Qwen3.5-0.8B",
81 + block_size=BLOCK_SIZE,
82 + hf_config_override={"linear_key_head_dim": K},
83 + )
84 + config.additional_config = {"gdn_prefill_backend": "cutedsl"}
85 + config.cache_config.mamba_cache_mode = "none"
86 + config.speculative_config = SpeculativeConfig(
87 + method="ngram", num_speculative_tokens=NUM_SPEC
88 + )
89 + return config
90 +
91 +
92 +def _build_layer(
93 + vllm_config,
94 + conv_state: torch.Tensor,
95 + ssm_state: torch.Tensor,
96 + a_log: torch.Tensor,
97 + dt_bias: torch.Tensor,
98 + conv_weight: torch.Tensor,
99 + norm_weight: torch.Tensor,
100 +):
101 + layer = types.SimpleNamespace(
102 + prefix=PREFIX,
103 + enable_packed_recurrent_decode=False,
104 + disable_tp_for_ba_proj=False,
105 + tp_size=1,
106 + num_k_heads=H,
107 + num_v_heads=HV,
108 + head_k_dim=K,
109 + head_v_dim=V,
110 + key_dim=K,
111 + value_dim=HV * V,
112 + activation="silu",
113 + A_log=a_log,
114 + dt_bias=dt_bias,
115 + conv1d=types.SimpleNamespace(weight=conv_weight, bias=None),
116 + kv_cache=(conv_state, ssm_state),
117 + norm=_TestGatedNorm(norm_weight),
118 + layer_norm_epsilon=EPS,
119 + gdn_decode_kernel="cuda",
120 + )
121 + with set_current_vllm_config(vllm_config):
122 + layer.chunk_gated_delta_rule = ChunkGatedDeltaRule()
123 + for name in (
124 + "rearrange_mixed_qkv",
125 + "_forward_core",
126 + "_forward_core_decode_spec_post_conv_fused_norm",
127 + "_forward_core_decode_spec_fused_norm",
128 + "_can_use_fused_gdn_mtp_decode",
129 + "_rms_norm_gated_cuda",
130 + "_forward_core_fused_norm",
131 + "_forward_core_fused_norm_packed",
132 + "split_ba",
133 + ):
134 + setattr(
135 + layer,
136 + name,
137 + types.MethodType(getattr(QwenGatedDeltaNetAttention, name), layer),
138 + )
139 + return layer
140 +
141 +
142 +@torch.inference_mode()
143 +def test_fused_forward_uses_packed_entrypoint() -> None:
144 + """Fused mode keeps projected QKVZ and BA packed through the model op."""
145 + device = torch.device("cuda")
146 + num_tokens = 3
147 + hidden_states = torch.empty(num_tokens, 1, dtype=torch.bfloat16, device=device)
148 + mixed_qkvz = torch.randn(
149 + num_tokens, CONV_DIM + HV * V, dtype=torch.bfloat16, device=device
150 + )
151 + ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device)
152 + layer = types.SimpleNamespace(
153 + prefix=PREFIX,
154 + enable_fused_gdn_decode=True,
155 + norm=types.SimpleNamespace(
156 + weight=torch.empty(V, dtype=torch.bfloat16, device=device)
157 + ),
158 + num_v_heads=HV,
159 + tp_size=1,
160 + head_v_dim=V,
161 + in_proj_qkvz=lambda _: (mixed_qkvz, None),
162 + in_proj_ba=lambda _: (ba, None),
163 + out_proj=lambda x: (x, None),
164 + )
165 + layer.forward_cuda = types.MethodType(
166 + QwenGatedDeltaNetAttention.forward_cuda, layer
167 + )
168 +
169 + def packed_op(
170 + actual_qkvz: torch.Tensor,
171 + actual_ba: torch.Tensor,
172 + output: torch.Tensor,
173 + *,
174 + layer_name: str,
175 + ) -> None:
176 + assert actual_qkvz is mixed_qkvz
177 + assert actual_ba is ba
178 + assert layer_name == _encode_layer_name(PREFIX)
179 + output.fill_(1)
180 +
181 + with (
182 + patch.object(
183 + torch.ops.vllm,
184 + "qwen_gdn_attention_core_fused_norm_packed",
185 + side_effect=packed_op,
186 + ) as packed_mock,
187 + patch.object(torch.ops.vllm, "qwen_gdn_attention_core") as triton_mock,
188 + ):
189 + output = layer.forward_cuda(hidden_states)
190 +
191 + packed_mock.assert_called_once()
192 + triton_mock.assert_not_called()
193 + torch.testing.assert_close(output, torch.ones_like(output))
194 +
195 +
196 +@pytest.mark.parametrize(
197 + "seq_lens,query_lens,draft_tokens,expected_fused_calls",
198 + [
199 + pytest.param([128], [SPEC_TOKENS], [NUM_SPEC], 1, id="pure-mtp"),
200 + pytest.param(
201 + [128, 96],
202 + [SPEC_TOKENS, 64],
203 + [NUM_SPEC, -1],
204 + 0,
205 + id="mixed-mtp-falls-back",
206 + ),
207 + pytest.param([96], [64], [-1], 0, id="pure-prefill"),
208 + pytest.param([128], [1], [-1], 0, id="pure-decode"),
209 + ],
210 +)
211 +@torch.inference_mode()
212 +def test_fused_model_path_matches_reference(
213 + seq_lens: list[int],
214 + query_lens: list[int],
215 + draft_tokens: list[int],
216 + expected_fused_calls: int,
217 +) -> None:
218 + """Fused MTP and its mixed/prefill/decode fallbacks match the reference."""
219 + torch.manual_seed(1)
220 + device = torch.device("cuda")
221 + vllm_config = _make_vllm_config()
222 + builder = GDNAttentionMetadataBuilder(
223 + kv_cache_spec=MambaSpec(
224 + block_size=BLOCK_SIZE,
225 + shapes=((16, 64),),
226 + dtypes=(torch.float16,),
227 + num_speculative_blocks=NUM_SPEC,
228 + ),
229 + layer_names=[PREFIX],
230 + vllm_config=vllm_config,
231 + device=device,
232 + )
233 + batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
234 + common = create_common_attn_metadata(
235 + batch, BLOCK_SIZE, device, arange_block_indices=True
236 + )
237 + common.block_table_tensor.add_(1)
238 + with set_current_vllm_config(vllm_config):
239 + metadata = builder.build(
240 + common_prefix_len=0,
241 + common_attn_metadata=common,
242 + num_accepted_tokens=torch.ones(
243 + batch.batch_size, dtype=torch.int32, device=device
244 + ),
245 + num_decode_draft_tokens_cpu=torch.tensor(draft_tokens, dtype=torch.int32),
246 + )
247 +
248 + state_indices = [
249 + indices
250 + for indices in (
251 + metadata.spec_state_indices_tensor,
252 + metadata.non_spec_state_indices_tensor,
253 + )
254 + if indices is not None and indices.numel() > 0
255 + ]
256 + pool_size = max(int(indices.max().item()) for indices in state_indices) + 1
257 + conv_state_shape, temporal_state_shape = (
258 + MambaStateShapeCalculator.gated_delta_net_state_shape(
259 + 1, H, HV, K, V, CONV_KERNEL, NUM_SPEC
260 + )
261 + )
262 + conv_state_seed = 0.05 * torch.randn(
263 + pool_size, *conv_state_shape, dtype=torch.bfloat16, device=device
264 + )
265 + ssm_state_seed = 0.01 * torch.randn(
266 + pool_size, *temporal_state_shape, dtype=torch.float32, device=device
267 + )
268 + a_log = 0.1 * torch.randn(HV, dtype=torch.float32, device=device)
269 + dt_bias = 0.1 * torch.randn(HV, dtype=torch.float32, device=device)
270 + conv_weight = 0.1 * torch.randn(
271 + CONV_DIM, 1, CONV_KERNEL, dtype=torch.bfloat16, device=device
272 + )
273 + norm_weight = torch.randn(V, dtype=torch.float32, device=device)
274 + num_tokens = batch.compute_num_tokens()
275 + mixed_qkv = 0.1 * torch.randn(
276 + num_tokens, CONV_DIM, dtype=torch.bfloat16, device=device
277 + )
278 + b = 0.1 * torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device)
279 + a = 0.1 * torch.randn_like(b)
280 + output_gate = 0.1 * torch.randn(
281 + num_tokens, HV, V, dtype=torch.bfloat16, device=device
282 + )
283 + mixed_qkvz = torch.cat((mixed_qkv, output_gate.flatten(1)), dim=-1)
284 + ba = torch.cat((b, a), dim=-1)
285 + context = types.SimpleNamespace(attn_metadata={PREFIX: metadata})
286 +
287 + reference_layer = _build_layer(
288 + vllm_config,
289 + conv_state_seed.clone(),
290 + ssm_state_seed.clone(),
291 + a_log,
292 + dt_bias,
293 + conv_weight,
294 + norm_weight,
295 + )
296 + reference_out = torch.zeros_like(output_gate)
297 + with patch.object(
298 + qwen_gdn_linear_attn, "get_forward_context", return_value=context
299 + ):
300 + reference_layer._forward_core(
301 + mixed_qkv=mixed_qkv.clone(),
302 + b=b,
303 + a=a,
304 + core_attn_out=reference_out,
305 + )
306 + reference_out = reference_layer.norm(reference_out, output_gate)
307 +
308 + fused_layer = _build_layer(
309 + vllm_config,
310 + conv_state_seed.clone(),
311 + ssm_state_seed.clone(),
312 + a_log,
313 + dt_bias,
314 + conv_weight,
315 + norm_weight,
316 + )
317 + context.no_compile_layers = {PREFIX: fused_layer}
318 + fused_out = torch.zeros_like(output_gate)
319 + fused_op = qwen_gdn_linear_attn.ops.fused_gdn_decode_post_conv_mtp
320 + with (
321 + patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=context),
322 + patch.object(
323 + qwen_gdn_linear_attn.ops,
324 + "fused_gdn_decode_post_conv_mtp",
325 + wraps=fused_op,
326 + ) as fused_mock,
327 + ):
328 + torch.ops.vllm.qwen_gdn_attention_core_fused_norm_packed(
329 + mixed_qkvz.clone(),
330 + ba,
331 + fused_out,
332 + layer_name=_encode_layer_name(PREFIX),
333 + )
334 +
335 + assert fused_mock.call_count == expected_fused_calls
336 + torch.testing.assert_close(
337 + fused_layer.kv_cache[0], reference_layer.kv_cache[0], atol=0, rtol=0
338 + )
339 + torch.testing.assert_close(fused_out, reference_out, atol=3e-2, rtol=3e-2)
340 + torch.testing.assert_close(
341 + fused_layer.kv_cache[1],
342 + reference_layer.kv_cache[1],
343 + atol=3e-2,
344 + rtol=3e-2,
345 + )
tests/kernels/test_fused_gdn_post_conv.py
+139 -3
@@ -1,18 +1,25 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for fused_gdn_prefill_post_conv kernel.
3 +"""Tests for GDN post-convolution kernels.
4
5 -Verifies that the fused kernel matches the reference:
6 - split → rearrange → contiguous → l2norm → gating
5 +The prefill tests cover preparation and the MTP tests cover recurrent state
6 +updates plus output normalization and gating.
7 """
8
9 import pytest
10 import torch
11 import torch.nn.functional as F
12
13 +from vllm import _custom_ops as ops
14 +from vllm.third_party.flash_linear_attention.ops import (
15 + fused_sigmoid_gating_delta_rule_update,
16 +)
17 from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import (
18 fused_post_conv_prep,
19 )
20 +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
21 + rmsnorm_fn,
22 +)
23
24
25 def reference_post_conv(
@@ -207,3 +214,132 @@ def test_fused_post_conv_l0():
214 )
215 assert q.shape == (0, H, K)
216 assert g.shape == (0, HV)
217 +
218 +
219 +@pytest.mark.parametrize(
220 + "tp_size,query_lengths,state_dtype,norm_dtype",
221 + [
222 + pytest.param(16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"),
223 + pytest.param(4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"),
224 + pytest.param(16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"),
225 + pytest.param(4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"),
226 + pytest.param(16, (8,), torch.float32, torch.bfloat16, id="tp16-max"),
227 + pytest.param(4, (8,), torch.bfloat16, torch.float32, id="tp4-max"),
228 + ],
229 +)
230 +@torch.inference_mode()
231 +def test_fused_gdn_decode_post_conv_mtp_ratio8(
232 + tp_size: int,
233 + query_lengths: tuple[int, ...],
234 + state_dtype: torch.dtype,
235 + norm_dtype: torch.dtype,
236 +) -> None:
237 + if torch.cuda.get_device_capability() < (8, 0):
238 + pytest.skip("fused GDN decode MTP requires compute capability 8.0+")
239 + if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"):
240 + pytest.skip("fused GDN decode MTP op is not built")
241 +
242 + torch.manual_seed(0)
243 + device = "cuda"
244 + H = 16 // tp_size
245 + HV = 128 // tp_size
246 + K = V = 128
247 + num_reqs = len(query_lengths)
248 + state_width = max(query_lengths)
249 + num_tokens = sum(query_lengths)
250 + num_slots = num_reqs * state_width + 1
251 + scale = K**-0.5
252 + eps = 1e-6
253 +
254 + mixed_qkv = torch.randn(
255 + num_tokens,
256 + 2 * H * K + HV * V,
257 + dtype=torch.bfloat16,
258 + device=device,
259 + )
260 + query, key, value = torch.split(
261 + mixed_qkv,
262 + [H * K, H * K, HV * V],
263 + dim=-1,
264 + )
265 + query = query.view(1, num_tokens, H, K)
266 + key = key.view(1, num_tokens, H, K)
267 + value = value.view(1, num_tokens, HV, V)
268 + ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device)
269 + b, a = ba.chunk(2, dim=-1)
270 + assert not a.is_contiguous()
271 + assert not b.is_contiguous()
272 + A_log = 0.5 * torch.randn(HV, dtype=torch.float32, device=device)
273 + dt_bias = 0.1 * torch.randn(HV, dtype=torch.float32, device=device)
274 + output_gate = torch.randn(num_tokens, HV, V, dtype=torch.bfloat16, device=device)
275 + norm_weight = torch.randn(V, dtype=norm_dtype, device=device)
276 + state_ref = (
277 + 0.01 * torch.randn(num_slots, HV, V, K, dtype=torch.float32, device=device)
278 + ).to(state_dtype)
279 + state_actual = state_ref.clone()
280 + state_indices = torch.arange(1, num_slots, dtype=torch.int32, device=device).view(
281 + num_reqs, state_width
282 + )
283 + cu_seqlens = torch.tensor(
284 + [0, *torch.tensor(query_lengths).cumsum(0).tolist()],
285 + dtype=torch.int32,
286 + device=device,
287 + )
288 + num_accepted_tokens = torch.ones(num_reqs, dtype=torch.int32, device=device)
289 + if query_lengths[-1] == 0:
290 + state_indices[-1].zero_()
291 +
292 + for step, accepted_tokens in enumerate((1, min(2, state_width), state_width)):
293 + num_accepted_tokens.fill_(accepted_tokens)
294 + if query_lengths[-1] == 0:
295 + num_accepted_tokens[-1] = 1
296 + raw_ref, _ = fused_sigmoid_gating_delta_rule_update(
297 + A_log=A_log,
298 + a=a,
299 + b=b,
300 + dt_bias=dt_bias,
301 + q=query,
302 + k=key,
303 + v=value,
304 + initial_state=state_ref,
305 + inplace_final_state=True,
306 + cu_seqlens=cu_seqlens,
307 + ssm_state_indices=state_indices,
308 + num_accepted_tokens=num_accepted_tokens,
309 + scale=scale,
310 + use_qk_l2norm_in_kernel=True,
311 + )
312 + expected = rmsnorm_fn(
313 + raw_ref.squeeze(0),
314 + norm_weight,
315 + None,
316 + z=output_gate,
317 + eps=eps,
318 + norm_before_gate=True,
319 + activation="silu",
320 + )
321 + actual = ops.fused_gdn_decode_post_conv_mtp(
322 + mixed_qkv=mixed_qkv,
323 + a=a,
324 + b=b,
325 + A_log=A_log,
326 + dt_bias=dt_bias,
327 + state_indices=state_indices,
328 + cu_seqlens=cu_seqlens,
329 + num_accepted_tokens=num_accepted_tokens,
330 + state=state_actual,
331 + output_gate=output_gate,
332 + norm_weight=norm_weight,
333 + out=torch.empty_like(output_gate),
334 + scale=scale,
335 + norm_eps=eps,
336 + )
337 +
338 + output_error = (actual.float() - expected.float()).norm()
339 + output_relative_l2 = output_error / expected.float().norm().clamp_min(1e-20)
340 + assert output_relative_l2 < 5e-4, (
341 + f"MTP output relative L2 mismatch at step {step}: "
342 + f"{output_relative_l2.item():.6g}"
343 + )
344 +
345 + torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2)
tests/test_envs.py
+15
@@ -251,6 +251,21 @@ class TestEnvWithChoices:
251 env_func()
252
253
254 +def test_gdn_decode_kernel_env(monkeypatch: pytest.MonkeyPatch):
255 + env_func = environment_variables["VLLM_GDN_DECODE_KERNEL"]
256 + monkeypatch.delenv("VLLM_GDN_DECODE_KERNEL", raising=False)
257 + assert env_func() == "cuda"
258 +
259 + for value in ("cuda", "triton"):
260 + monkeypatch.setenv("VLLM_GDN_DECODE_KERNEL", value)
261 + assert env_func() == value
262 +
263 + for value in ("fused", "invalid"):
264 + monkeypatch.setenv("VLLM_GDN_DECODE_KERNEL", value)
265 + with pytest.raises(ValueError, match="VLLM_GDN_DECODE_KERNEL"):
266 + env_func()
267 +
268 +
269 class TestEnvListWithChoices:
270 """Test cases for env_list_with_choices function."""
271
vllm/_custom_ops.py
+37
@@ -2777,6 +2777,43 @@ def fused_kda_decode(
2777 return out
2778
2779
2780 +def fused_gdn_decode_post_conv_mtp(
2781 + mixed_qkv: torch.Tensor,
2782 + a: torch.Tensor,
2783 + b: torch.Tensor,
2784 + A_log: torch.Tensor,
2785 + dt_bias: torch.Tensor,
2786 + state_indices: torch.Tensor,
2787 + cu_seqlens: torch.Tensor,
2788 + num_accepted_tokens: torch.Tensor,
2789 + state: torch.Tensor,
2790 + output_gate: torch.Tensor,
2791 + norm_weight: torch.Tensor,
2792 + out: torch.Tensor | None = None,
2793 + scale: float = 128**-0.5,
2794 + norm_eps: float = 1e-5,
2795 +) -> torch.Tensor:
2796 + if out is None:
2797 + out = torch.empty_like(output_gate)
2798 + torch.ops._C.fused_gdn_decode_post_conv_mtp(
2799 + mixed_qkv,
2800 + a,
2801 + b,
2802 + A_log,
2803 + dt_bias,
2804 + state_indices,
2805 + cu_seqlens,
2806 + num_accepted_tokens,
2807 + state,
2808 + output_gate,
2809 + norm_weight,
2810 + out,
2811 + scale,
2812 + norm_eps,
2813 + )
2814 + return out
2815 +
2816 +
2817 def concat_and_cache_mla(
2818 kv_c: torch.Tensor,
2819 k_pe: torch.Tensor,
vllm/config/compilation.py
+1
@@ -769,6 +769,7 @@ class CompilationConfig:
769 "vllm::short_conv",
770 "vllm::linear_attention",
771 "vllm::qwen_gdn_attention_core",
772 + "vllm::qwen_gdn_attention_core_fused_norm_packed",
773 "vllm::gdn_attention_core_xpu",
774 "vllm::olmo_hybrid_gdn_full_forward",
775 "vllm::sparse_attn_indexer",
vllm/envs.py
+10
@@ -127,6 +127,7 @@ if TYPE_CHECKING:
127 VLLM_DISABLED_KERNELS: list[str] = []
128 VLLM_USE_HW_AGNOSTIC: bool = False
129 VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE: bool = True
130 + VLLM_GDN_DECODE_KERNEL: Literal["cuda", "triton"] = "cuda"
131 VLLM_DISABLE_PYNCCL: bool = False
132 VLLM_USE_OINK_OPS: bool = False
133 VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True
@@ -1209,6 +1210,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
1210 "VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE": lambda: bool(
1211 int(os.getenv("VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1"))
1212 ),
1213 + # Select the GDN MTP decode implementation. "cuda" uses the fused decode
1214 + # kernel where supported and falls back to "triton" otherwise; setting it
1215 + # explicitly to "cuda" raises when unsupported.
1216 + "VLLM_GDN_DECODE_KERNEL": env_with_choices(
1217 + "VLLM_GDN_DECODE_KERNEL",
1218 + "cuda",
1219 + ["cuda", "triton"],
1220 + case_sensitive=False,
1221 + ),
1222 # Disable pynccl (using torch.distributed instead)
1223 "VLLM_DISABLE_PYNCCL": lambda: (
1224 os.getenv("VLLM_DISABLE_PYNCCL", "False").lower() in ("true", "1")
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+303 -4
@@ -2,12 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Inference-only Qwen3-Next/Qwen3.5 model."""
4
5 +import os
6 from typing import Literal
7
8 import torch
9 from einops import rearrange
10 from torch import nn
11
12 +from vllm import _custom_ops as ops
13 from vllm import envs
14 from vllm._aiter_ops import rocm_aiter_ops
15 from vllm.config import (
@@ -84,6 +86,9 @@ if GDN_AITER_TRITON_AVAILABLE:
86
87 logger = init_logger(__name__)
88
89 +MAX_FUSED_GDN_MTP_TOKENS = 8
90 +FUSED_GDN_STATE_DTYPES = (torch.float32, torch.bfloat16)
91 +
92
93 def _resolve_gdn_prefill_backend(
94 vllm_config: VllmConfig,
@@ -484,12 +489,50 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
489 self.enable_packed_recurrent_decode = (
490 envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
491 )
492 + self.gdn_decode_kernel = envs.VLLM_GDN_DECODE_KERNEL.strip().lower()
493 + if self.gdn_decode_kernel == "cuda":
494 + reason = self._fused_gdn_decode_unsupported_reason(vllm_config)
495 + if reason is not None:
496 + if "VLLM_GDN_DECODE_KERNEL" in os.environ:
497 + raise ValueError(
498 + f"VLLM_GDN_DECODE_KERNEL=cuda is not supported: {reason}"
499 + )
500 + logger.info_once(
501 + "Falling back to the Triton GDN decode path: %s", reason
502 + )
503 + self.gdn_decode_kernel = "triton"
504 + self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda"
505 + logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel)
506
507 compilation_config = get_current_vllm_config().compilation_config
508 if prefix in compilation_config.static_forward_context:
509 raise ValueError(f"Duplicate layer name: {prefix}")
510 compilation_config.static_forward_context[prefix] = self
511
512 + def _fused_gdn_decode_unsupported_reason(
513 + self, vllm_config: VllmConfig
514 + ) -> str | None:
515 + conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()
516 + if (
517 + self.gqa_interleaved_layout
518 + or self.head_k_dim != 128
519 + or self.head_v_dim != 128
520 + or self.norm.activation != "silu"
521 + or vllm_config.model_config.dtype != torch.bfloat16
522 + or conv_state_dtype != torch.bfloat16
523 + or recurrent_state_dtype not in FUSED_GDN_STATE_DTYPES
524 + or not current_platform.has_device_capability(80)
525 + ):
526 + return (
527 + "the fused CUDA kernel requires a BF16 GDN model with "
528 + "K=V=128, SiLU gating, non-interleaved GQA layout, BF16 "
529 + "convolution cache, BF16 or FP32 recurrent state, and a "
530 + "GPU with compute capability 8.0+"
531 + )
532 + if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"):
533 + return "torch.ops._C.fused_gdn_decode_post_conv_mtp is not built"
534 + return None
535 +
536 def create_qkvz_proj(
537 self,
538 hidden_size: int,
@@ -846,6 +889,26 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
889 mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
890 ba, _ = self.in_proj_ba(hidden_states)
891
892 + use_fused_gdn_decode = (
893 + self.enable_fused_gdn_decode
894 + and hidden_states.dtype == torch.bfloat16
895 + and self.norm.weight.dtype in (torch.bfloat16, torch.float32)
896 + )
897 + if use_fused_gdn_decode:
898 + core_attn_out = torch.zeros(
899 + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
900 + dtype=hidden_states.dtype,
901 + device=hidden_states.device,
902 + )
903 + torch.ops.vllm.qwen_gdn_attention_core_fused_norm_packed(
904 + mixed_qkvz,
905 + ba,
906 + core_attn_out,
907 + layer_name=_encode_layer_name(self.prefix),
908 + )
909 + output, _ = self.out_proj(core_attn_out.flatten(-2))
910 + return output
911 +
912 if self.gqa_interleaved_layout:
913 # Qwen3-Next: unpack the interleaved GQA layout
914 query, key, value, z, b, a = self.fix_query_key_value_ordering(
@@ -862,8 +925,6 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
925 mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
926 z = z.reshape(z.size(0), -1, self.head_v_dim)
927 b, a = self.split_ba(ba)
865 - b = b.contiguous()
866 - a = a.contiguous()
928
929 # ============================================================
930 # Part 2: Core Attention (Custom Op)
@@ -878,8 +939,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
939
940 torch.ops.vllm.qwen_gdn_attention_core(
941 mixed_qkv,
881 - b,
882 - a,
942 + b.contiguous(),
943 + a.contiguous(),
944 core_attn_out,
945 layer_name=_encode_layer_name(self.prefix),
946 )
@@ -1622,6 +1683,211 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1683 )
1684 return
1685
1686 + def _forward_core_decode_spec_fused_norm(
1687 + self,
1688 + mixed_qkv: torch.Tensor,
1689 + b: torch.Tensor,
1690 + a: torch.Tensor,
1691 + output_gate: torch.Tensor,
1692 + core_attn_out: torch.Tensor,
1693 + attn_metadata: GDNAttentionMetadata,
1694 + ) -> None:
1695 + state_indices = attn_metadata.spec_state_indices_tensor
1696 + cu_seqlens = attn_metadata.spec_query_start_loc
1697 + num_accepted_tokens = attn_metadata.num_accepted_tokens
1698 + assert state_indices is not None
1699 + assert cu_seqlens is not None
1700 + assert num_accepted_tokens is not None
1701 +
1702 + num_requests = attn_metadata.num_spec_decodes
1703 + num_actual_tokens = attn_metadata.num_actual_tokens
1704 + conv_state = (
1705 + self.kv_cache[0]
1706 + if is_conv_state_dim_first()
1707 + else self.kv_cache[0].transpose(-1, -2)
1708 + )
1709 + conv_weights = self.conv1d.weight.view(
1710 + self.conv1d.weight.size(0), self.conv1d.weight.size(2)
1711 + )
1712 + mixed_qkv = causal_conv1d_update(
1713 + mixed_qkv[:num_actual_tokens],
1714 + conv_state,
1715 + conv_weights,
1716 + self.conv1d.bias,
1717 + self.activation,
1718 + conv_state_indices=state_indices[:num_requests, 0],
1719 + num_accepted_tokens=num_accepted_tokens[:num_requests],
1720 + query_start_loc=cu_seqlens[: num_requests + 1],
1721 + max_query_len=state_indices.size(1),
1722 + validate_data=False,
1723 + )
1724 + self._forward_core_decode_spec_post_conv_fused_norm(
1725 + mixed_qkv=mixed_qkv,
1726 + b=b[:num_actual_tokens],
1727 + a=a[:num_actual_tokens],
1728 + output_gate=output_gate[:num_actual_tokens],
1729 + core_attn_out=core_attn_out[:num_actual_tokens],
1730 + attn_metadata=attn_metadata,
1731 + )
1732 +
1733 + def _forward_core_decode_spec_post_conv_fused_norm(
1734 + self,
1735 + mixed_qkv: torch.Tensor,
1736 + b: torch.Tensor,
1737 + a: torch.Tensor,
1738 + output_gate: torch.Tensor,
1739 + core_attn_out: torch.Tensor,
1740 + attn_metadata: GDNAttentionMetadata,
1741 + ) -> None:
1742 + state_indices = attn_metadata.spec_state_indices_tensor
1743 + cu_seqlens = attn_metadata.spec_query_start_loc
1744 + num_accepted_tokens = attn_metadata.num_accepted_tokens
1745 + assert state_indices is not None
1746 + assert cu_seqlens is not None
1747 + assert num_accepted_tokens is not None
1748 +
1749 + num_requests = attn_metadata.num_spec_decodes
1750 + ops.fused_gdn_decode_post_conv_mtp(
1751 + mixed_qkv=mixed_qkv,
1752 + a=a,
1753 + b=b,
1754 + A_log=self.A_log,
1755 + dt_bias=self.dt_bias,
1756 + state_indices=state_indices[:num_requests],
1757 + cu_seqlens=cu_seqlens[: num_requests + 1],
1758 + num_accepted_tokens=num_accepted_tokens[:num_requests],
1759 + state=self.kv_cache[1],
1760 + output_gate=output_gate,
1761 + norm_weight=self.norm.weight,
1762 + out=core_attn_out,
1763 + scale=self.head_k_dim**-0.5,
1764 + norm_eps=self.layer_norm_epsilon,
1765 + )
1766 +
1767 + def _forward_core_fused_norm_packed(
1768 + self,
1769 + mixed_qkvz: torch.Tensor,
1770 + ba: torch.Tensor,
1771 + core_attn_out: torch.Tensor,
1772 + ) -> None:
1773 + forward_context = get_forward_context()
1774 + attn_metadata_raw = forward_context.attn_metadata
1775 + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
1776 + if attn_metadata_raw is None:
1777 + self._warmup_prefill_kernels(mixed_qkvz[:, :qkv_size], 0)
1778 + return
1779 +
1780 + assert isinstance(attn_metadata_raw, dict)
1781 + attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]
1782 + assert isinstance(attn_metadata, GDNAttentionMetadata)
1783 + mixed_qkv, output_gate_flat = mixed_qkvz.split(
1784 + [qkv_size, self.value_dim // self.tp_size], dim=-1
1785 + )
1786 + output_gate = output_gate_flat.reshape(
1787 + output_gate_flat.size(0), -1, self.head_v_dim
1788 + )
1789 + b, a = self.split_ba(ba)
1790 + self._forward_core_fused_norm(
1791 + mixed_qkv=mixed_qkv,
1792 + b=b,
1793 + a=a,
1794 + output_gate=output_gate,
1795 + core_attn_out=core_attn_out,
1796 + )
1797 +
1798 + def _can_use_fused_gdn_mtp_decode(
1799 + self, attn_metadata: GDNAttentionMetadata
1800 + ) -> bool:
1801 + state_indices = attn_metadata.spec_state_indices_tensor
1802 + return (
1803 + attn_metadata.spec_sequence_masks is not None
1804 + and attn_metadata.num_decodes == 0
1805 + and attn_metadata.num_spec_decodes > 0
1806 + and self.kv_cache[1].dtype in FUSED_GDN_STATE_DTYPES
1807 + and self.gdn_decode_kernel == "cuda"
1808 + and self.num_v_heads == 8 * self.num_k_heads
1809 + and state_indices is not None
1810 + and state_indices.size(1) <= MAX_FUSED_GDN_MTP_TOKENS
1811 + and hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp")
1812 + )
1813 +
1814 + def _rms_norm_gated_cuda(
1815 + self,
1816 + x: torch.Tensor,
1817 + output_gate: torch.Tensor,
1818 + out: torch.Tensor,
1819 + ) -> None:
1820 + from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
1821 + layer_norm_fwd,
1822 + )
1823 +
1824 + x_shape = x.shape
1825 + assert output_gate.shape == x_shape
1826 + assert out.shape == x_shape
1827 + x_2d = x.reshape(-1, x_shape[-1])
1828 + output_gate_2d = output_gate.reshape(-1, x_shape[-1])
1829 + out_2d = out.reshape(-1, x_shape[-1])
1830 + assert x_2d.stride(-1) == 1
1831 + assert output_gate_2d.stride(-1) == 1
1832 + assert out_2d.stride(-1) == 1
1833 + layer_norm_fwd(
1834 + x_2d,
1835 + self.norm.weight.contiguous(),
1836 + self.norm.bias,
1837 + self.norm.eps,
1838 + z=output_gate_2d,
1839 + out=out_2d,
1840 + group_size=(
1841 + x_shape[-1] if self.norm.group_size is None else self.norm.group_size
1842 + ),
1843 + norm_before_gate=self.norm.norm_before_gate,
1844 + is_rms_norm=True,
1845 + activation=self.norm.activation,
1846 + )
1847 +
1848 + def _forward_core_fused_norm(
1849 + self,
1850 + mixed_qkv: torch.Tensor,
1851 + b: torch.Tensor,
1852 + a: torch.Tensor,
1853 + output_gate: torch.Tensor,
1854 + core_attn_out: torch.Tensor,
1855 + ) -> None:
1856 + forward_context = get_forward_context()
1857 + attn_metadata_raw = forward_context.attn_metadata
1858 + if attn_metadata_raw is None:
1859 + self._warmup_prefill_kernels(mixed_qkv, 0)
1860 + return
1861 +
1862 + assert isinstance(attn_metadata_raw, dict)
1863 + attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]
1864 + assert isinstance(attn_metadata, GDNAttentionMetadata)
1865 + if (
1866 + self._can_use_fused_gdn_mtp_decode(attn_metadata)
1867 + and attn_metadata.num_prefills == 0
1868 + ):
1869 + self._forward_core_decode_spec_fused_norm(
1870 + mixed_qkv=mixed_qkv,
1871 + b=b,
1872 + a=a,
1873 + output_gate=output_gate,
1874 + core_attn_out=core_attn_out,
1875 + attn_metadata=attn_metadata,
1876 + )
1877 + return
1878 + self._forward_core(
1879 + mixed_qkv=mixed_qkv,
1880 + b=b.contiguous(),
1881 + a=a.contiguous(),
1882 + core_attn_out=core_attn_out,
1883 + )
1884 + num_actual_tokens = attn_metadata.num_actual_tokens
1885 + self._rms_norm_gated_cuda(
1886 + core_attn_out[:num_actual_tokens],
1887 + output_gate[:num_actual_tokens],
1888 + core_attn_out[:num_actual_tokens],
1889 + )
1890 +
1891
1892 def qwen_gdn_attention_core(
1893 qkv_or_qkvz: torch.Tensor,
@@ -1683,6 +1949,39 @@ direct_register_custom_op(
1949 )
1950
1951
1952 +def qwen_gdn_attention_core_fused_norm_packed(
1953 + mixed_qkvz: torch.Tensor,
1954 + ba: torch.Tensor,
1955 + core_attn_out: torch.Tensor,
1956 + layer_name: LayerNameType,
1957 +) -> None:
1958 + layer_name = _resolve_layer_name(layer_name)
1959 + forward_context: ForwardContext = get_forward_context()
1960 + self = forward_context.no_compile_layers[layer_name]
1961 + self._forward_core_fused_norm_packed(
1962 + mixed_qkvz=mixed_qkvz,
1963 + ba=ba,
1964 + core_attn_out=core_attn_out,
1965 + )
1966 +
1967 +
1968 +def gdn_attention_core_fused_norm_packed_fake(
1969 + mixed_qkvz: torch.Tensor,
1970 + ba: torch.Tensor,
1971 + core_attn_out: torch.Tensor,
1972 + layer_name: LayerNameType,
1973 +) -> None:
1974 + return
1975 +
1976 +
1977 +direct_register_custom_op(
1978 + op_name="qwen_gdn_attention_core_fused_norm_packed",
1979 + op_func=qwen_gdn_attention_core_fused_norm_packed,
1980 + mutates_args=["core_attn_out"],
1981 + fake_impl=gdn_attention_core_fused_norm_packed_fake,
1982 +)
1983 +
1984 +
1985 @triton.jit
1986 def fused_gdn_gating_kernel(
1987 g,