@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dd72658e7d

[Perf][GLM-5.2] Blackwell decode optimizations (#48597)

Signed-off-by: Peiyuan Zhou <peiyuanzhou1994@gmail.com> Signed-off-by: zhou <zhou@zhoudeMacBook-Pro.local> Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: zhou <zhou@zhoudeMacBook-Pro.local> Co-authored-by: stefankoncarevic <Stefan.Koncarevic@amd.com>

xiaozhoupy committed Jul 24, 2026 at 12:36 UTC dd72658e7db0c6a674f473f6f9f0f5c2ebe7e523
29 files changed +2139 -138
CMakeLists.txt
+17
@@ -727,6 +727,23 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
727 "(requires SM90+ and CUDA >= 12.0).")
728 endif()
729
730 + # BF16 skinny GEMM (M<=32; weight-bandwidth-bound decode shapes).
731 + # Requires SM90+.
732 + cuda_archs_sm90plus(BF16_SKINNY_GEMM_ARCHS "${CUDA_ARCHS}")
733 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND BF16_SKINNY_GEMM_ARCHS)
734 + set(BF16_SKINNY_GEMM_SRCS
735 + "csrc/libtorch_stable/bf16_skinny_gemm_entry.cu"
736 + "csrc/libtorch_stable/bf16_skinny_gemm.cu")
737 + set_gencode_flags_for_srcs(
738 + SRCS "${BF16_SKINNY_GEMM_SRCS}"
739 + CUDA_ARCHS "${BF16_SKINNY_GEMM_ARCHS}")
740 + list(APPEND VLLM_STABLE_EXT_SRC "${BF16_SKINNY_GEMM_SRCS}")
741 + message(STATUS "Building bf16_skinny_gemm for archs: ${BF16_SKINNY_GEMM_ARCHS}")
742 + else()
743 + message(STATUS "Not building bf16_skinny_gemm as no compatible archs found "
744 + "(requires SM90+ and CUDA >= 12.0).")
745 + endif()
746 +
747 # Only build AllSpark kernels if we are building for at least some compatible archs.
748 cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
749 if (ALLSPARK_ARCHS)
csrc/libtorch_stable/bf16_skinny_gemm.cu new
+262
@@ -0,0 +1,262 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +//
4 +// Skinny GEMM: activation(bf16) x weight(bf16)^T -> bf16, for decode-time
5 +// M <= 32 with a large reduction dim. Replaces cuBLAS splitK (GEMM +
6 +// splitKreduce) with a single block-per-output-column kernel; these shapes
7 +// are weight-bandwidth-bound, so one coalesced pass over the weight at
8 +// fp32 accumulation is optimal. Adapted from fp32_router_gemm.cu.
9 +//
10 +// First user: the DeepSeek-V32/GLM-5.2 MTP eh_proj (K=2*hidden=12288,
11 +// N=hidden/TP), whose cuBLAS splitK pick costs ~34.6us vs the ~19us
12 +// bandwidth floor per replicated read (and ~4us once column-parallel).
13 +
14 +#include <cuda_bf16.h>
15 +#include <cuda_runtime.h>
16 +
17 +// ---------------------------------------------------------------------------
18 +// Load helpers (8 x bf16 = one uint4 load, converted to fp32)
19 +// ---------------------------------------------------------------------------
20 +
21 +namespace skinny {
22 +
23 +constexpr int VPT = 8; // bf16 values per thread per load
24 +
25 +__device__ __forceinline__ void load_bf16x8(__nv_bfloat16 const* ptr,
26 + float* dst) {
27 + uint4 v = *reinterpret_cast<uint4 const*>(ptr);
28 + __nv_bfloat16 const* p = reinterpret_cast<__nv_bfloat16 const*>(&v);
29 +#pragma unroll
30 + for (int i = 0; i < VPT; i++) dst[i] = __bfloat162float(p[i]);
31 +}
32 +
33 +// Streaming variant for the weight: each row is read exactly once across the
34 +// whole grid, so bypass L2 residency (evict-first). Measured -1.4us at M=1.
35 +__device__ __forceinline__ void load_bf16x8_cs(__nv_bfloat16 const* ptr,
36 + float* dst) {
37 + uint4 v = __ldcs(reinterpret_cast<uint4 const*>(ptr));
38 + __nv_bfloat16 const* p = reinterpret_cast<__nv_bfloat16 const*>(&v);
39 +#pragma unroll
40 + for (int i = 0; i < VPT; i++) dst[i] = __bfloat162float(p[i]);
41 +}
42 +
43 +// ---------------------------------------------------------------------------
44 +// Kernel: each block computes kNPB output columns for all kNumTokens rows.
45 +// grid = kN / kNPB, block = kBlockSize threads. K is reduced VPT elements
46 +// per thread per iteration; fp32 accumulation, warp butterfly + smem
47 +// finalize, bf16 store.
48 +// ---------------------------------------------------------------------------
49 +
50 +template <int kBlockSize, int kNumTokens, int kNPB, int kPF, int kN, int kK>
51 +__global__ __launch_bounds__(kBlockSize, 1) void bf16_skinny_gemm_kernel(
52 + __nv_bfloat16* out, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b,
53 + int64_t out_stride) {
54 + constexpr int k_elems_per_iter = VPT * kBlockSize;
55 + constexpr int k_iterations = kK / k_elems_per_iter;
56 + static_assert(kK % k_elems_per_iter == 0);
57 + constexpr int kWarpSize = 32;
58 + constexpr int kNumWarps = kBlockSize / kWarpSize;
59 +
60 + int const n_base = blockIdx.x * kNPB;
61 + int const tid = threadIdx.x;
62 + int const warpId = tid / kWarpSize;
63 + int const laneId = tid % kWarpSize;
64 +
65 + float acc[kNumTokens][kNPB] = {};
66 + __shared__ float sm_reduction[kNumTokens][kNPB][kNumWarps];
67 +
68 + // Register prefetch (kPF > 0): W does not depend on the predecessor, so
69 + // the first kPF iterations' weight chunks are loaded raw BEFORE the
70 + // dependency sync; with a PDL-releasing producer (fused_eh_norm fires
71 + // gdc_launch_dependents early) these DRAM round trips overlap the norm.
72 + // Pair-measured on B300 (norm+gemm in one graph, full 6144x12288):
73 + // M=1 pf2 28.43us vs pf0 29.00us; deeper prefetch or M >= 2 regresses
74 + // (register pressure), hence the per-M selection in the launcher.
75 + uint4 w_pre[kPF > 0 ? kPF : 1][kNPB];
76 +#pragma unroll
77 + for (int pf = 0; pf < kPF; pf++) {
78 +#pragma unroll
79 + for (int n = 0; n < kNPB; n++) {
80 + w_pre[pf][n] =
81 + *reinterpret_cast<uint4 const*>(mat_b + (size_t)(n_base + n) * kK +
82 + pf * k_elems_per_iter + tid * VPT);
83 + }
84 + }
85 +
86 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
87 + cudaGridDependencySynchronize();
88 +#endif
89 +
90 +#pragma unroll
91 + for (int ki = 0; ki < k_iterations; ki++) {
92 + int const k_base = ki * k_elems_per_iter + tid * VPT;
93 +
94 + float b_float[kNPB][VPT];
95 + if (ki < kPF) {
96 +#pragma unroll
97 + for (int n = 0; n < kNPB; n++) {
98 + __nv_bfloat16 const* p =
99 + reinterpret_cast<__nv_bfloat16 const*>(&w_pre[ki][n]);
100 +#pragma unroll
101 + for (int v = 0; v < VPT; v++) b_float[n][v] = __bfloat162float(p[v]);
102 + }
103 + } else {
104 +#pragma unroll
105 + for (int n = 0; n < kNPB; n++) {
106 + load_bf16x8_cs(mat_b + (size_t)(n_base + n) * kK + k_base, b_float[n]);
107 + }
108 + }
109 +
110 +#pragma unroll
111 + for (int m = 0; m < kNumTokens; m++) {
112 + float a_float[VPT];
113 + load_bf16x8(mat_a + (size_t)m * kK + k_base, a_float);
114 +#pragma unroll
115 + for (int n = 0; n < kNPB; n++) {
116 +#pragma unroll
117 + for (int k = 0; k < VPT; k++) {
118 + acc[m][n] += a_float[k] * b_float[n][k];
119 + }
120 + }
121 + }
122 + }
123 +
124 + // Warp-level butterfly reduction
125 +#pragma unroll
126 + for (int m = 0; m < kNumTokens; m++) {
127 +#pragma unroll
128 + for (int n = 0; n < kNPB; n++) {
129 + float sum = acc[m][n];
130 + sum += __shfl_xor_sync(0xffffffff, sum, 16);
131 + sum += __shfl_xor_sync(0xffffffff, sum, 8);
132 + sum += __shfl_xor_sync(0xffffffff, sum, 4);
133 + sum += __shfl_xor_sync(0xffffffff, sum, 2);
134 + sum += __shfl_xor_sync(0xffffffff, sum, 1);
135 + if (laneId == 0) sm_reduction[m][n][warpId] = sum;
136 + }
137 + }
138 +
139 + __syncthreads();
140 +
141 + // Parallel finalize: one thread per (m, n) output.
142 + for (int idx = tid; idx < kNumTokens * kNPB; idx += kBlockSize) {
143 + int const m = idx / kNPB;
144 + int const n = idx % kNPB;
145 + float final_sum = 0.0f;
146 +#pragma unroll
147 + for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][n][w];
148 + out[(size_t)m * out_stride + n_base + n] = __float2bfloat16(final_sum);
149 + }
150 +
151 + // Trigger after our stores: harmless hardening, not a guarantee — the
152 + // trigger only permits dependent-launch scheduling and carries no memory
153 + // visibility semantics, so a PDL consumer must still gridsync before
154 + // reading our output. Every current consumer is a plain launch (full
155 + // stream order); firing late just avoids an unnecessarily early launch
156 + // window and matches fp32_router_gemm.cu.
157 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
158 + cudaTriggerProgrammaticLaunchCompletion();
159 +#endif
160 +}
161 +
162 +} // namespace skinny
163 +
164 +// ---------------------------------------------------------------------------
165 +// Launcher
166 +// ---------------------------------------------------------------------------
167 +
168 +template <int kBlockSize, int kNPB, int kNumTokens, int kN, int kK>
169 +void invokeBf16SkinnyGemm(__nv_bfloat16* output, __nv_bfloat16 const* mat_a,
170 + __nv_bfloat16 const* mat_b, int64_t out_stride,
171 + cudaStream_t stream) {
172 + static_assert(kN % kNPB == 0);
173 + // Weight prefetch depth: only M=1 measured a win (see kernel comment).
174 + constexpr int kPF = (kNumTokens == 1) ? 2 : 0;
175 + cudaLaunchConfig_t config;
176 + config.gridDim = kN / kNPB;
177 + config.blockDim = kBlockSize;
178 + config.dynamicSmemBytes = 0;
179 + config.stream = stream;
180 + cudaLaunchAttribute attrs[1];
181 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
182 + attrs[0].val.programmaticStreamSerializationAllowed = 1;
183 + config.numAttrs = 1;
184 + config.attrs = attrs;
185 + cudaLaunchKernelEx(&config,
186 + skinny::bf16_skinny_gemm_kernel<kBlockSize, kNumTokens,
187 + kNPB, kPF, kN, kK>,
188 + output, mat_a, mat_b, out_stride);
189 +}
190 +
191 +// ---------------------------------------------------------------------------
192 +// Explicit instantiations. M = 1..32; (N, K) pairs:
193 +// (768, 12288) eh_proj shard, TP8
194 +// (1536, 12288) eh_proj shard, TP4
195 +// (6144, 12288) eh_proj unsharded
196 +// kNPB (B300 sweep, M=1): 6144 -> 2 (3072 blocks, 23.2us vs 25.3 at kNPB=8;
197 +// narrow blocks minimize wave quantization); shards 768/1536 keep 4.
198 +// ---------------------------------------------------------------------------
199 +
200 +#define INSTANTIATE(M, NPB, N, K) \
201 + template void invokeBf16SkinnyGemm<128, NPB, M, N, K>( \
202 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int64_t, \
203 + cudaStream_t);
204 +
205 +#define INSTANTIATE_ALL_M(NPB, N, K) \
206 + INSTANTIATE(1, NPB, N, K) \
207 + INSTANTIATE(2, NPB, N, K) \
208 + INSTANTIATE(3, NPB, N, K) \
209 + INSTANTIATE(4, NPB, N, K) \
210 + INSTANTIATE(5, NPB, N, K) \
211 + INSTANTIATE(6, NPB, N, K) \
212 + INSTANTIATE(7, NPB, N, K) \
213 + INSTANTIATE(8, NPB, N, K) \
214 + INSTANTIATE(9, NPB, N, K) \
215 + INSTANTIATE(10, NPB, N, K) \
216 + INSTANTIATE(11, NPB, N, K) \
217 + INSTANTIATE(12, NPB, N, K) \
218 + INSTANTIATE(13, NPB, N, K) \
219 + INSTANTIATE(14, NPB, N, K) \
220 + INSTANTIATE(15, NPB, N, K) \
221 + INSTANTIATE(16, NPB, N, K) \
222 + INSTANTIATE(17, NPB, N, K) \
223 + INSTANTIATE(18, NPB, N, K) \
224 + INSTANTIATE(19, NPB, N, K) \
225 + INSTANTIATE(20, NPB, N, K) \
226 + INSTANTIATE(21, NPB, N, K) \
227 + INSTANTIATE(22, NPB, N, K) \
228 + INSTANTIATE(23, NPB, N, K) \
229 + INSTANTIATE(24, NPB, N, K) \
230 + INSTANTIATE(25, NPB, N, K) \
231 + INSTANTIATE(26, NPB, N, K) \
232 + INSTANTIATE(27, NPB, N, K) \
233 + INSTANTIATE(28, NPB, N, K) \
234 + INSTANTIATE(29, NPB, N, K) \
235 + INSTANTIATE(30, NPB, N, K) \
236 + INSTANTIATE(31, NPB, N, K) \
237 + INSTANTIATE(32, NPB, N, K)
238 +
239 +INSTANTIATE_ALL_M(4, 768, 12288)
240 +INSTANTIATE_ALL_M(4, 1536, 12288)
241 +INSTANTIATE_ALL_M(2, 6144, 12288)
242 +// LL-mode (M<=8 wiring guard) backbone shapes, B300 sweep vs cuBLAS:
243 +// q_b_proj (2048, 2048): 1.67x/1.29x/1.15x at M=4/6/8 (NPB=4 within
244 +// 0.1us of per-M best)
245 +// shared-expert gate_up (512, 6144): 1.95x/1.58x/1.40x at M=4/6/8
246 +// cuBLAS keeps qkv_a (2624,6144) and o_proj (6144,2048) — already at
247 +// 3.4-3.8 TB/s there; the GEMV loses on activation re-reads.
248 +INSTANTIATE_ALL_M(4, 2048, 2048)
249 +INSTANTIATE_ALL_M(4, 512, 6144)
250 +// fused_qkv_a (2624, 6144), 32MB: skinny wins ONLY at M<=2 (B300: M=1
251 +// 6.99us vs cuBLAS 9.21 = 1.32x, M=2 1.24x; M>=4 cuBLAS holds at 3.5TB/s
252 +// and every alternative loses — cublasLt top-8 3.6TB/s wall, DeepGEMM
253 +// 0.71x, wmma+cp.async custom 0.30x pending a TMA rewrite).
254 +INSTANTIATE_ALL_M(4, 2624, 6144)
255 +// DSv3.2 (TP8) siblings of the GLM shapes above, same dual-chip matrix:
256 +// fused_qkv_a (2112, 7168), 30MB: wins M<=2 (M=1 1.30-1.34x)
257 +// MTP eh_proj (7168, 14336), 205MB: wins M<=2 (M=1 1.12-1.16x)
258 +INSTANTIATE_ALL_M(4, 2112, 7168)
259 +INSTANTIATE_ALL_M(2, 7168, 14336)
260 +
261 +#undef INSTANTIATE_ALL_M
262 +#undef INSTANTIATE
csrc/libtorch_stable/bf16_skinny_gemm_entry.cu new
+170
@@ -0,0 +1,170 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +#include <torch/csrc/stable/library.h>
5 +#include <torch/csrc/stable/tensor.h>
6 +#include <torch/headeronly/core/ScalarType.h>
7 +
8 +#include "core/registration.h"
9 +#include "libtorch_stable/torch_utils.h"
10 +
11 +#include <cuda_bf16.h>
12 +#include <cuda_runtime.h>
13 +
14 +#include <stdexcept>
15 +
16 +namespace {
17 +
18 +inline int getSMVersion() {
19 + auto* props = get_device_prop();
20 + return props->major * 10 + props->minor;
21 +}
22 +
23 +} // namespace
24 +
25 +static constexpr int SKINNY_MAX_TOKENS = 32;
26 +
27 +// Supported (N, K) pairs (must match the instantiations in
28 +// bf16_skinny_gemm.cu): eh_proj shard TP8 / TP4 / unsharded.
29 +static inline bool bf16_skinny_gemm_supported(int n, int k) {
30 + if (k == 12288 && (n == 768 || n == 1536 || n == 6144)) return true;
31 + // LL-mode backbone shapes (wire callers with an M <= 8 guard; the GEMV
32 + // family loses to cuBLAS at larger M).
33 + if (k == 2048 && n == 2048) return true; // q_b_proj (TP8)
34 + if (k == 6144 && n == 2624) return true; // fused_qkv_a (wire M <= 2 only)
35 + if (k == 7168 && n == 2112) return true; // DSv3.2 fused_qkv_a (M <= 2)
36 + if (k == 14336 && n == 7168) return true; // DSv3.2 eh_proj (M <= 2)
37 + if (k == 6144 && n == 512) return true; // shared-expert gate_up (TP8)
38 + return false;
39 +}
40 +
41 +// Forward declarations - template params must match bf16_skinny_gemm.cu
42 +template <int kBlockSize, int kNPB, int kNumTokens, int kN, int kK>
43 +void invokeBf16SkinnyGemm(__nv_bfloat16* output, __nv_bfloat16 const* mat_a,
44 + __nv_bfloat16 const* mat_b, int64_t out_stride,
45 + cudaStream_t stream);
46 +
47 +template <int kNPB, int kN, int kK, int kBegin, int kEnd>
48 +struct SkinnyLoopUnroller {
49 + static void unroll(int num_tokens, __nv_bfloat16* output,
50 + __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b,
51 + int64_t out_stride, cudaStream_t stream) {
52 + if (num_tokens == kBegin) {
53 + invokeBf16SkinnyGemm<128, kNPB, kBegin, kN, kK>(output, mat_a, mat_b,
54 + out_stride, stream);
55 + } else {
56 + SkinnyLoopUnroller<kNPB, kN, kK, kBegin + 1, kEnd>::unroll(
57 + num_tokens, output, mat_a, mat_b, out_stride, stream);
58 + }
59 + }
60 +};
61 +
62 +template <int kNPB, int kN, int kK, int kEnd>
63 +struct SkinnyLoopUnroller<kNPB, kN, kK, kEnd, kEnd> {
64 + static void unroll(int num_tokens, __nv_bfloat16* output,
65 + __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b,
66 + int64_t out_stride, cudaStream_t stream) {
67 + if (num_tokens == kEnd) {
68 + invokeBf16SkinnyGemm<128, kNPB, kEnd, kN, kK>(output, mat_a, mat_b,
69 + out_stride, stream);
70 + } else {
71 + throw std::invalid_argument(
72 + "bf16_skinny_gemm: num_tokens must be in [1, 32]");
73 + }
74 + }
75 +};
76 +
77 +static void dispatchBf16SkinnyGemm(int n, int k, int num_tokens,
78 + __nv_bfloat16* output,
79 + __nv_bfloat16 const* mat_a,
80 + __nv_bfloat16 const* mat_b,
81 + int64_t out_stride, cudaStream_t stream) {
82 + if (n == 768 && k == 12288) {
83 + SkinnyLoopUnroller<4, 768, 12288, 1, SKINNY_MAX_TOKENS>::unroll(
84 + num_tokens, output, mat_a, mat_b, out_stride, stream);
85 + } else if (n == 1536 && k == 12288) {
86 + SkinnyLoopUnroller<4, 1536, 12288, 1, SKINNY_MAX_TOKENS>::unroll(
87 + num_tokens, output, mat_a, mat_b, out_stride, stream);
88 + } else if (n == 6144 && k == 12288) {
89 + SkinnyLoopUnroller<2, 6144, 12288, 1, SKINNY_MAX_TOKENS>::unroll(
90 + num_tokens, output, mat_a, mat_b, out_stride, stream);
91 + } else if (n == 2048 && k == 2048) {
92 + SkinnyLoopUnroller<4, 2048, 2048, 1, SKINNY_MAX_TOKENS>::unroll(
93 + num_tokens, output, mat_a, mat_b, out_stride, stream);
94 + } else if (n == 2624 && k == 6144) {
95 + SkinnyLoopUnroller<4, 2624, 6144, 1, SKINNY_MAX_TOKENS>::unroll(
96 + num_tokens, output, mat_a, mat_b, out_stride, stream);
97 + } else if (n == 2112 && k == 7168) {
98 + SkinnyLoopUnroller<4, 2112, 7168, 1, SKINNY_MAX_TOKENS>::unroll(
99 + num_tokens, output, mat_a, mat_b, out_stride, stream);
100 + } else if (n == 7168 && k == 14336) {
101 + SkinnyLoopUnroller<2, 7168, 14336, 1, SKINNY_MAX_TOKENS>::unroll(
102 + num_tokens, output, mat_a, mat_b, out_stride, stream);
103 + } else if (n == 512 && k == 6144) {
104 + SkinnyLoopUnroller<4, 512, 6144, 1, SKINNY_MAX_TOKENS>::unroll(
105 + num_tokens, output, mat_a, mat_b, out_stride, stream);
106 + } else {
107 + throw std::invalid_argument("bf16_skinny_gemm: unsupported (N, K) pair");
108 + }
109 +}
110 +
111 +void bf16_skinny_gemm(
112 + torch::stable::Tensor& output, // [num_tokens, N] bf16
113 + torch::stable::Tensor const& mat_a, // [num_tokens, K] bf16
114 + torch::stable::Tensor const& mat_b // [N, K] bf16
115 +) {
116 + STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
117 + STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(),
118 + "bf16_skinny_gemm: all tensors must be CUDA tensors");
119 + STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() &&
120 + output.get_device_index() == mat_b.get_device_index(),
121 + "bf16_skinny_gemm: all tensors must be on the same device");
122 + STD_TORCH_CHECK(mat_a.is_contiguous() && mat_b.is_contiguous(),
123 + "bf16_skinny_gemm: inputs must be contiguous");
124 + // Output may be a column-slice view of a wider padded buffer: unit column
125 + // stride, row stride >= N (rows must not overlap).
126 + STD_TORCH_CHECK(output.stride(1) == 1,
127 + "bf16_skinny_gemm: output columns must be contiguous");
128 +
129 + const int num_tokens = mat_a.size(0);
130 + const int n = mat_b.size(0);
131 + const int k = mat_a.size(1);
132 +
133 + STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == n,
134 + "bf16_skinny_gemm: output must be [num_tokens, N]");
135 + STD_TORCH_CHECK(mat_b.size(1) == k,
136 + "bf16_skinny_gemm: mat_a and mat_b must share K");
137 + STD_TORCH_CHECK(bf16_skinny_gemm_supported(n, k),
138 + "bf16_skinny_gemm: unsupported (N, K) pair");
139 + const int64_t out_stride = output.stride(0);
140 + STD_TORCH_CHECK(num_tokens <= 1 || out_stride >= n,
141 + "bf16_skinny_gemm: output rows overlap");
142 + STD_TORCH_CHECK(num_tokens >= 0 && num_tokens <= SKINNY_MAX_TOKENS,
143 + "bf16_skinny_gemm: num_tokens must be in [0, 32]");
144 + STD_TORCH_CHECK(
145 + mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
146 + mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
147 + output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
148 + "bf16_skinny_gemm: all tensors must be bfloat16");
149 +
150 + // Empty batch (e.g. an empty rank at a DP/PP boundary): nothing to compute.
151 + if (num_tokens == 0) {
152 + return;
153 + }
154 +
155 + const torch::stable::accelerator::DeviceGuard device_guard(
156 + mat_a.get_device_index());
157 + STD_TORCH_CHECK(getSMVersion() >= 90, "bf16_skinny_gemm: requires SM90+");
158 +
159 + auto stream = get_current_cuda_stream(mat_a.get_device_index());
160 + dispatchBf16SkinnyGemm(
161 + n, k, num_tokens,
162 + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
163 + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
164 + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), out_stride,
165 + stream);
166 +}
167 +
168 +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
169 + m.impl("bf16_skinny_gemm", TORCH_BOX(&bf16_skinny_gemm));
170 +}
csrc/libtorch_stable/dsv3_fused_a_gemm.cu
+112 -36
@@ -391,7 +391,7 @@ struct MmaComputer {
391 static constexpr int n_iter_cnt =
392 (tile_n + 7) /
393 8; // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
394 - static_assert(m_iter_cnt == 1);
394 + static_assert(m_iter_cnt == 1 || m_iter_cnt == 2);
395 static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
396
397 __device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
@@ -416,13 +416,18 @@ struct MmaComputer {
416 public:
417 __device__ void prepare() {
418 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
419 + // Fragment addressing is per 16-row ldmatrix tile; m_iter selects the
420 + // 16-row half within tile_m.
421 #pragma unroll
420 - for (int i = 0; i < k_phase_cnt; i++) {
421 - int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
422 - int m_idx = linear_idx % tile_m;
423 - int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
424 - k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
425 - a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
422 + for (int m = 0; m < m_iter_cnt; m++) {
423 + #pragma unroll
424 + for (int i = 0; i < k_phase_cnt; i++) {
425 + int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
426 + int m_idx = linear_idx % 16 + m * 16;
427 + int k_idx = linear_idx / 16 + warp_k_offset_in_tile_k;
428 + k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
429 + a_smem_offsets[m][i] = m_idx * tile_k + k_idx;
430 + }
431 }
432 #pragma unroll
433 for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
@@ -446,11 +451,14 @@ struct MmaComputer {
451 wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
452
453 #pragma unroll
449 - for (int i = 0; i < k_phase_cnt; i++) {
450 - int smem_offset = a_smem_offsets[0][i];
451 - bf16_t* smem_ptr_this_iter =
452 - smem_a + stage_idx * tile_m * tile_k + smem_offset;
453 - ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
454 + for (int m = 0; m < m_iter_cnt; m++) {
455 + #pragma unroll
456 + for (int i = 0; i < k_phase_cnt; i++) {
457 + int smem_offset = a_smem_offsets[m][i];
458 + bf16_t* smem_ptr_this_iter =
459 + smem_a + stage_idx * tile_m * tile_k + smem_offset;
460 + ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[m][i]));
461 + }
462 }
463
464 #pragma unroll
@@ -469,9 +477,12 @@ struct MmaComputer {
477 for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
478 #pragma unroll
479 for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
472 - hmma_16_8_16_f32acc_bf16ab(
473 - acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
474 - b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
480 + #pragma unroll
481 + for (int m = 0; m < m_iter_cnt; m++) {
482 + hmma_16_8_16_f32acc_bf16ab(
483 + acc_reg[m][n_iter_idx], a_reg[m][k_iter_idx],
484 + b_reg[n_iter_idx][k_iter_idx], acc_reg[m][n_iter_idx]);
485 + }
486 }
487 }
488 ::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
@@ -486,14 +497,14 @@ struct MmaComputer {
497 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
498 asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
499 // reorganize the acc_reg
489 - constexpr int thread_m = 2;
500 + constexpr int thread_m = 2 * m_iter_cnt;
501 constexpr int thread_n = 2 * n_iter_cnt;
502 constexpr int cta_mma_n = n_iter_cnt * 8;
503 float acc_reg_reorg[thread_m][thread_n];
504
505 for (int i = 0; i < thread_m; i++) {
506 for (int j = 0; j < thread_n; j++) {
496 - acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
507 + acc_reg_reorg[i][j] = acc_reg[i / 2][j / 2][(j % 2) + (i % 2) * 2];
508 }
509 }
510
@@ -514,7 +525,8 @@ struct MmaComputer {
525 for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
526 #pragma unroll
527 for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
517 - int m_idx = (lane_idx / 4) + m_idx_thread * 8;
528 + int m_idx =
529 + (lane_idx / 4) + (m_idx_thread % 2) * 8 + (m_idx_thread / 2) * 16;
530 int n_idx =
531 ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
532 smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
@@ -587,7 +599,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
599 static_assert(
600 tile_k == 128 || tile_k == 256 || tile_k == 512 ||
601 tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK.
590 - static_assert(tile_m == 16);
602 + static_assert(tile_m == 16 || tile_m == 32);
603 constexpr int g2s_vec_bytes = 16;
604 constexpr int a_elem_bytes = 2;
605 constexpr int b_elem_bytes = 2;
@@ -647,7 +659,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
659 #endif
660 }
661
650 -template <typename T, int kHdIn, int kHdOut, int kTileN>
662 +template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileM = 16>
663 void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
664 cudaStream_t const stream) {
665 constexpr int gemm_m = kHdOut; // 2112
@@ -655,7 +667,7 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
667 constexpr int gemm_k = kHdIn; // 7168
668 constexpr int batch_size = 1;
669 std::swap(mat_a, mat_b);
658 - constexpr int tile_m = 16;
670 + constexpr int tile_m = kTileM;
671 constexpr int tile_n = kTileN; // 8 or 16
672 constexpr int tile_k = std::max(256, 1024 / tile_n); // 256
673 constexpr int max_stage_cnt =
@@ -702,6 +714,44 @@ template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
714 __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
715 cudaStream_t);
716
717 +// GLM-5.2 fused_qkv_a (K=6144 -> N=2624). tile_m=32 so the grid is
718 +// 2624/32 = 82 CTAs, a single wave on B300 (148 SMs); tile_m=16's 164 CTAs
719 +// straddle two waves and drop to 2.9 TB/s vs 3.9-4.0 here. Beats cuBLAS
720 +// at every decode M: 1.11-1.15x for M<=8 (tile_n=8), 1.12x at M=16
721 +// (tile_n=16). The M<=2 dispatch still belongs to bf16_skinny_gemm
722 +// (4.5 TB/s).
723 +template void invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 8, 32>(
724 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
725 + cudaStream_t);
726 +
727 +template void invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 16, 32>(
728 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
729 + cudaStream_t);
730 +
731 +// GLM-5.2 q_b_proj TP8 (K=2048 -> N=2048). tile_m=16 keeps 128 CTAs (already
732 +// a single wave) and measures ahead of tile_m=32 here. Beats cuBLAS
733 +// 1.79-1.87x at M=3..8 (tile_n=8) and 1.19-1.30x at M=9..16 (tile_n=16);
734 +// M<=2 belongs to bf16_skinny_gemm, M>=20 to cuBLAS.
735 +template void invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 8>(
736 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
737 + cudaStream_t);
738 +
739 +template void invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 16>(
740 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
741 + cudaStream_t);
742 +
743 +// DSv3.2 q_b_proj TP8 (K=1536 -> N=3072). tile_m=32 keeps 96 CTAs (single
744 +// wave; tile_m=16's 192 straddle two). Beats cuBLAS 1.6-1.8x at M=1..8 and
745 +// 1.05-1.18x at M=9..16 on B300/B200; the whole 1..16 range dispatches here
746 +// (no skinny tier: K=1536 does not fit the GEMV's 128x8 K-step).
747 +template void invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 8, 32>(
748 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
749 + cudaStream_t);
750 +
751 +template void invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 16, 32>(
752 + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
753 + cudaStream_t);
754 +
755 void dsv3_fused_a_gemm(torch::stable::Tensor& output,
756 torch::stable::Tensor const& mat_a,
757 torch::stable::Tensor const& mat_b) {
@@ -710,12 +760,15 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
760 int const hd_in = mat_a.size(1);
761 int const hd_out = mat_b.size(1);
762
713 - constexpr int kHdIn = 7168;
714 - constexpr int kHdOut = 2112;
763 + bool const is_dsv3 = hd_in == 7168 && hd_out == 2112;
764 + bool const is_glm = hd_in == 6144 && hd_out == 2624;
765 + bool const is_glm_qb = hd_in == 2048 && hd_out == 2048;
766 + bool const is_ds_qb = hd_in == 1536 && hd_out == 3072;
767 STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
768 "required 1 <= mat_a.shape[0] <= 16");
717 - STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
718 - STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
769 + STD_TORCH_CHECK(is_dsv3 || is_glm || is_glm_qb || is_ds_qb,
770 + "supported (hd_in, hd_out): (7168, 2112), (6144, 2624), "
771 + "(2048, 2048), (1536, 3072)");
772 STD_TORCH_CHECK(output.size(0) == num_tokens,
773 "required output.shape[0] == mat_a.shape[0]");
774 STD_TORCH_CHECK(output.size(1) == hd_out,
@@ -738,18 +791,41 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
791 STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
792
793 auto stream = get_current_cuda_stream(mat_a.get_device_index());
741 - if (num_tokens <= 8) {
742 - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
743 - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
744 - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
745 - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
746 - stream);
794 + auto* out = reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
795 + auto* a = reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
796 + auto* b = reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
797 + if (is_dsv3) {
798 + if (num_tokens <= 8) {
799 + invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(out, a, b, num_tokens,
800 + stream);
801 + } else {
802 + invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(out, a, b, num_tokens,
803 + stream);
804 + }
805 + } else if (is_glm) {
806 + if (num_tokens <= 8) {
807 + invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 8, 32>(out, a, b, num_tokens,
808 + stream);
809 + } else {
810 + invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 16, 32>(out, a, b, num_tokens,
811 + stream);
812 + }
813 + } else if (is_glm_qb) {
814 + if (num_tokens <= 8) {
815 + invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 8>(out, a, b, num_tokens,
816 + stream);
817 + } else {
818 + invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 16>(out, a, b, num_tokens,
819 + stream);
820 + }
821 } else {
748 - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
749 - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
750 - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
751 - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
752 - stream);
822 + if (num_tokens <= 8) {
823 + invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 8, 32>(out, a, b, num_tokens,
824 + stream);
825 + } else {
826 + invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 16, 32>(out, a, b, num_tokens,
827 + stream);
828 + }
829 }
830 }
831
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
+42 -10
@@ -49,6 +49,11 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
49 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
50 "Vec size is not matched.");
51
52 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
53 + cudaGridDependencySynchronize();
54 + cudaTriggerProgrammaticLaunchCompletion();
55 +#endif
56 +
57 // Precompute SF layout parameter (constant for entire kernel).
58 int32_t const numKTiles = (outputCols + 63) / 64;
59
@@ -123,6 +128,11 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
128 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
129 "Vec size is not matched.");
130
131 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
132 + cudaGridDependencySynchronize();
133 + cudaTriggerProgrammaticLaunchCompletion();
134 +#endif
135 +
136 int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
137 int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
138
@@ -201,6 +211,8 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
211 const torch::stable::accelerator::DeviceGuard device_guard(
212 input.get_device_index());
213 auto stream = get_current_cuda_stream(input.get_device_index());
214 + auto* device_props = get_device_prop();
215 + int const sm_version = device_props->major * 10 + device_props->minor;
216
217 int output_sf_n_unpadded = int(output_n / CVT_FP4_SF_VEC_SIZE);
218
@@ -224,10 +236,21 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
236 input.scalar_type(), "nvfp4_quant_kernel", [&] {
237 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
238 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
227 - vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
228 - m, n, output_n, num_padded_cols, input_ptr, input_sf_ptr,
229 - reinterpret_cast<uint32_t*>(output_ptr),
230 - reinterpret_cast<uint32_t*>(sf_out));
239 + cudaLaunchConfig_t config = {};
240 + config.gridDim = grid;
241 + config.blockDim = block;
242 + config.dynamicSmemBytes = 0;
243 + config.stream = stream;
244 + cudaLaunchAttribute attrs[1];
245 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
246 + attrs[0].val.programmaticStreamSerializationAllowed = 1;
247 + config.numAttrs = (sm_version >= 90) ? 1 : 0;
248 + config.attrs = attrs;
249 + cudaLaunchKernelEx(&config, vllm::cvt_fp16_to_fp4<cuda_type, false>,
250 + m, n, output_n, num_padded_cols, input_ptr,
251 + input_sf_ptr,
252 + reinterpret_cast<uint32_t*>(output_ptr),
253 + reinterpret_cast<uint32_t*>(sf_out));
254 });
255 } else {
256 int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD;
@@ -240,12 +263,21 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
263 input.scalar_type(), "nvfp4_quant_kernel", [&] {
264 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
265 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
243 - vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
244 - <<<grid, block, 0, stream>>>(
245 - m, n, output_n, output_sf_n_unpadded, num_packed_cols,
246 - input_ptr, input_sf_ptr,
247 - reinterpret_cast<uint32_t*>(output_ptr),
248 - reinterpret_cast<uint32_t*>(sf_out));
266 + cudaLaunchConfig_t config = {};
267 + config.gridDim = grid;
268 + config.blockDim = block;
269 + config.dynamicSmemBytes = 0;
270 + config.stream = stream;
271 + cudaLaunchAttribute attrs[1];
272 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
273 + attrs[0].val.programmaticStreamSerializationAllowed = 1;
274 + config.numAttrs = (sm_version >= 90) ? 1 : 0;
275 + config.attrs = attrs;
276 + cudaLaunchKernelEx(
277 + &config, vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>, m, n,
278 + output_n, output_sf_n_unpadded, num_packed_cols, input_ptr,
279 + input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr),
280 + reinterpret_cast<uint32_t*>(sf_out));
281 });
282 }
283 }
csrc/libtorch_stable/torch_bindings.cpp
+4
@@ -330,6 +330,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
330 // conditionally compiled so impl registration is in source file
331 ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
332
333 + // BF16 skinny GEMM (M<=32, weight-BW-bound shapes, e.g. MTP eh_proj).
334 + // conditionally compiled so impl registration is in source file
335 + ops.def("bf16_skinny_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
336 +
337 // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
338 ops.def(
339 "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
recipes/glm5.2-ll-b300-tp8-mtp5.md new
+41
@@ -0,0 +1,41 @@
1 +# GLM-5.2-NVFP4 — Low-Latency (LL) serving recipe · 8×B300 · TP8 · MTP=5
2 +
3 +Single-node 8×B300 (sm_100/103), pure tensor parallel, NVFP4 weights, fp8 KV cache,
4 +MTP=5 speculative decode, model runner v2.
5 +
6 +Validated on `glm5.2-LL` @ `8d407aee1`.
7 +
8 +## Versions
9 +
10 +| Component | Version |
11 +| --- | --- |
12 +| vLLM | `glm5.2-LL` @ `8d407aee1` |
13 +| flashinfer-python | `0.6.15` |
14 +| Model | `nvidia/GLM-5.2-NVFP4` (`modelopt_fp4`) |
15 +
16 +## Build (once)
17 +
18 +```bash
19 +pip install 'flashinfer-python==0.6.15'
20 +
21 +CUDA_HOME=/usr/local/cuda-13.0 TORCH_CUDA_ARCH_LIST="10.0 10.3" MAX_JOBS=96 \
22 +VLLM_USE_PRECOMPILED=0 VLLM_USE_PRECOMPILED_RUST=0 \
23 +pip install -e . --no-deps --no-build-isolation
24 +```
25 +
26 +## Server
27 +
28 +```bash
29 +export VLLM_USE_V2_MODEL_RUNNER=1 # model runner v2
30 +export VLLM_DEEP_GEMM_WARMUP=skip
31 +
32 +vllm serve nvidia/GLM-5.2-NVFP4 \
33 + --trust-remote-code \
34 + --tensor-parallel-size 8 \
35 + --quantization modelopt_fp4 --kv-cache-dtype fp8_e4m3 \
36 + --max-model-len 32768 --max-num-batched-tokens 16384 --max-num-seqs 256 \
37 + --no-enable-prefix-caching --gpu-memory-utilization 0.85 \
38 + --speculative-config '{"method":"mtp","num_speculative_tokens":5}' \
39 + --kernel-config '{"ir_op_priority":{"rms_norm":["vllm_c","native"],"fused_add_rms_norm":["vllm_c","native"]}}' \
40 + --host 0.0.0.0 --port 8000
41 +```
tests/kernels/test_bf16_skinny_gemm.py new
+70
@@ -0,0 +1,70 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Numerical correctness for the bf16 skinny GEMM (decode-M GEMV) kernel.
4 +
5 +Covers every supported (N, K) shape across M = 0/1/2/8/16/32 (M=32 exercises
6 +the tile_m=32 MMA path), plus strided (column-slice) output. M=0 must return
7 +an empty result without launching the kernel (empty ranks at DP/PP
8 +boundaries)."""
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm import _custom_ops as ops
14 +from vllm.platforms import current_platform
15 +
16 +# (N, K) pairs must match bf16_skinny_gemm_supported() in the entry .cu.
17 +SHAPES = [
18 + (768, 12288),
19 + (1536, 12288),
20 + (6144, 12288),
21 + (2048, 2048),
22 + (2624, 6144),
23 + (2112, 7168),
24 + (7168, 14336),
25 + (512, 6144),
26 +]
27 +MS = [0, 1, 2, 8, 16, 32]
28 +
29 +pytestmark = pytest.mark.skipif(
30 + not current_platform.is_cuda() or not current_platform.has_device_capability(90),
31 + reason="bf16_skinny_gemm requires CUDA SM90+",
32 +)
33 +
34 +
35 +def _rel_err(out: torch.Tensor, ref: torch.Tensor) -> float:
36 + return (out.float() - ref).norm().item() / ref.norm().clamp_min(1e-6).item()
37 +
38 +
39 +@pytest.mark.parametrize("n,k", SHAPES)
40 +@pytest.mark.parametrize("m", MS)
41 +def test_bf16_skinny_gemm_matches_reference(n: int, k: int, m: int):
42 + torch.manual_seed(0)
43 + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
44 + w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
45 +
46 + out = ops.bf16_skinny_gemm(x, w)
47 + assert out.shape == (m, n)
48 + assert out.dtype == torch.bfloat16
49 +
50 + if m == 0: # empty batch: nothing to check beyond the empty shape
51 + return
52 + ref = x.float() @ w.float().t()
53 + rel = _rel_err(out, ref)
54 + assert rel < 2e-2, f"rel err {rel:.4f} for (m,n,k)=({m},{n},{k})"
55 +
56 +
57 +@pytest.mark.parametrize("n,k", [(2048, 2048), (512, 6144)])
58 +def test_bf16_skinny_gemm_strided_output(n: int, k: int):
59 + """Output is a column-slice of a wider padded buffer (row stride > N)."""
60 + torch.manual_seed(0)
61 + m = 8
62 + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
63 + w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
64 +
65 + wide = torch.empty(m, n + 64, dtype=torch.bfloat16, device="cuda")
66 + view = wide[:, :n]
67 + torch.ops._C.bf16_skinny_gemm(view, x, w)
68 +
69 + ref = x.float() @ w.float().t()
70 + assert _rel_err(view, ref) < 2e-2
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+77
@@ -49,6 +49,32 @@ pytestmark = pytest.mark.skipif(
49 )
50
51
52 +def test_platform_capability_queries_are_constant_during_compile(monkeypatch):
53 + monkeypatch.setattr(
54 + K.current_platform, "has_device_capability", lambda capability: True
55 + )
56 + monkeypatch.setattr(K, "has_cutedsl", lambda: True)
57 + monkeypatch.setattr(K.current_platform, "is_arch_support_pdl", lambda: True)
58 +
59 + def capability_branches(x: torch.Tensor) -> torch.Tensor:
60 + if K._can_use_fused_q_cutedsl() and K._is_arch_support_pdl():
61 + return x + 1
62 + return x - 1
63 +
64 + compiled = torch.compile(capability_branches, backend="eager", fullgraph=True)
65 + x = torch.zeros(1, device="cuda")
66 + torch.testing.assert_close(compiled(x), torch.ones_like(x))
67 +
68 +
69 +def test_pdl_is_disabled_before_blackwell(monkeypatch):
70 + monkeypatch.setattr(
71 + K.current_platform, "has_device_capability", lambda capability: False
72 + )
73 + monkeypatch.setattr(K.current_platform, "is_arch_support_pdl", lambda: True)
74 +
75 + assert not K._is_arch_support_pdl()
76 +
77 +
78 # ── reference helpers ────────────────────────────────────────────────────────
79
80
@@ -292,6 +318,57 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
318 assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
319
320
321 +def test_fused_norm_rope_profile_without_cache_compiles():
322 + """The cache-free profiling path must still compile and produce Q."""
323 + torch.manual_seed(2)
324 + dev = "cuda"
325 + num_tokens = 4
326 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
327 + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
328 + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
329 + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
330 + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
331 + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
332 + index_k = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
333 + index_w = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
334 + index_b = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
335 + cos_sin = make_cos_sin(64, ROPE_DIM, dev)
336 + topk = torch.empty(num_tokens, 2048, device=dev, dtype=torch.int32)
337 +
338 + def profile_run(
339 + q_c: torch.Tensor,
340 + kv_c: torch.Tensor,
341 + k_pe: torch.Tensor,
342 + index_k: torch.Tensor,
343 + ) -> torch.Tensor:
344 + return K.fused_norm_rope(
345 + pos,
346 + q_c,
347 + qw,
348 + EPS,
349 + kv_c,
350 + kvw,
351 + EPS,
352 + k_pe,
353 + cos_sin,
354 + index_k,
355 + index_w,
356 + index_b,
357 + EPS,
358 + cos_sin,
359 + topk,
360 + slot_mapping=None,
361 + indexer_k_cache=None,
362 + mla_kv_cache=None,
363 + has_indexer=True,
364 + index_rope_interleave=False,
365 + )
366 +
367 + compiled = torch.compile(profile_run, fullgraph=True)
368 + q_out = compiled(q_c, kv_c, k_pe, index_k)
369 + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (profiling)")
370 +
371 +
372 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
373 def test_fused_norm_rope_ds_mla(num_tokens: int):
374 """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).
tests/models/registry.py
+4
@@ -1604,6 +1604,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1604 speculative_model="luccafong/deepseek_mtp_draft_random",
1605 trust_remote_code=True,
1606 ),
1607 + "DeepseekV32MTPModel": _HfExamplesInfo(
1608 + "deepseek-ai/DeepSeek-V3.2-Exp",
1609 + speculative_model="deepseek-ai/DeepSeek-V3.2-Exp",
1610 + ),
1611 "DeepSeekV4MTPModel": _HfExamplesInfo(
1612 "deepseek-ai/DeepSeek-V4-Flash",
1613 speculative_model="deepseek-ai/DeepSeek-V4-Flash",
vllm/_custom_ops.py
+25
@@ -2371,6 +2371,20 @@ def fp32_router_gemm(
2371 return output
2372
2373
2374 +def bf16_skinny_gemm(
2375 + x: torch.Tensor,
2376 + weight: torch.Tensor,
2377 +) -> torch.Tensor:
2378 + """Skinny bf16 GEMM (M<=32): x [M, K] @ weight [N, K]^T -> [M, N].
2379 +
2380 + Single block-per-column kernel with fp32 accumulation; replaces cuBLAS
2381 + splitK for weight-bandwidth-bound decode shapes (e.g. MTP eh_proj).
2382 + """
2383 + output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device)
2384 + torch.ops._C.bf16_skinny_gemm(output, x, weight)
2385 + return output
2386 +
2387 +
2388 if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):
2389
2390 @register_fake("_C::fp32_router_gemm")
@@ -2382,6 +2396,17 @@ if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):
2396 return
2397
2398
2399 +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "bf16_skinny_gemm"):
2400 +
2401 + @register_fake("_C::bf16_skinny_gemm")
2402 + def bf16_skinny_gemm_fake(
2403 + output: torch.Tensor,
2404 + mat_a: torch.Tensor,
2405 + mat_b: torch.Tensor,
2406 + ) -> None:
2407 + return
2408 +
2409 +
2410 def topk_softmax(
2411 topk_weights: torch.Tensor,
2412 topk_ids: torch.Tensor,
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+111 -35
@@ -202,11 +202,7 @@ if flashinfer_comm is not None:
202 element_size = allreduce_in.element_size()
203 current_tensor_size = num_tokens * hidden_size * element_size
204 max_tensor_size = max_token_num * hidden_size * element_size
205 - assert current_tensor_size <= max_tensor_size, (
206 - f"Current tensor size {current_tensor_size} is larger than "
207 - f"max token num {max_token_num} * hidden size {hidden_size} * "
208 - f"element size {element_size}"
209 - )
205 + oversize = current_tensor_size > max_tensor_size
206 curr_device = current_platform.get_device_capability()
207 device_capability = curr_device.to_int() if curr_device is not None else None
208
@@ -219,6 +215,47 @@ if flashinfer_comm is not None:
215 get_workspace_fn = (
216 get_fi_ar_quant_workspace if is_quant_pattern else get_fi_ar_workspace
217 )
218 +
219 + orig_norm_out = norm_out
220 +
221 + def unfused_fallback() -> None:
222 + # Fused-path equivalent honoring the op's output-buffer contract:
223 + # with norm_out given, normed -> norm_out and the new residual ->
224 + # the allreduce_in buffer; otherwise normed -> allreduce_in and
225 + # the new residual -> the residual buffer. Runs fine during CUDA
226 + # graph capture (the branch is resolved per captured shape).
227 + #
228 + # weight_bias mirrors the fused kernel's gamma offset (Gemma
229 + # patterns pass 1.0 for rms_norm(x, weight + 1)); fold it into
230 + # the gamma in fp32 like the reference before the plain rms_norm.
231 + gamma = (
232 + rms_gamma
233 + if weight_bias == 0.0
234 + else (rms_gamma.float() + weight_bias).to(rms_gamma.dtype)
235 + )
236 + reduced = tensor_model_parallel_all_reduce(allreduce_in)
237 + new_residual = reduced + residual
238 + if orig_norm_out is None:
239 + torch.ops._C.rms_norm(allreduce_in, new_residual, gamma, rms_eps)
240 + residual.copy_(new_residual)
241 + else:
242 + torch.ops._C.rms_norm(orig_norm_out, new_residual, gamma, rms_eps)
243 + allreduce_in.copy_(new_residual)
244 +
245 + if oversize:
246 + assert not is_quant_pattern, (
247 + "flashinfer fused allreduce: input exceeds the workspace budget "
248 + "and the unfused fallback does not cover quant patterns"
249 + )
250 + logger.warning_once(
251 + "fused allreduce+rmsnorm: %d tokens exceeds the workspace "
252 + "budget (max_token_num=%d); using the unfused fallback for "
253 + "this shape.",
254 + num_tokens,
255 + max_token_num,
256 + )
257 + unfused_fallback()
258 + return
259 workspace = get_workspace_fn(
260 world_size=world_size,
261 rank=get_tensor_model_parallel_rank(),
@@ -252,36 +289,75 @@ if flashinfer_comm is not None:
289 if workspace.backend in ("trtllm", "mnnvl"):
290 layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4
291
255 - flashinfer_comm.allreduce_fusion(
256 - input=allreduce_in,
257 - workspace=workspace,
258 - pattern=pattern_code,
259 - launch_with_pdl=launch_with_pdl,
260 - output=None,
261 - residual_out=residual_out,
262 - norm_out=norm_out,
263 - quant_out=quant_out,
264 - scale_out=scale_out,
265 - residual_in=residual,
266 - rms_gamma=rms_gamma,
267 - rms_eps=rms_eps,
268 - scale_factor=scale_factor,
269 - layout_code=layout_code,
270 - use_oneshot=use_oneshot,
271 - fp32_acc=fp32_acc,
272 - weight_bias=weight_bias,
273 - # The one-shot Lamport all-reduce signals PDL completion before its
274 - # output buffer is committed when trigger_completion_at_end is
275 - # False, so the next PDL-launched kernel can read the uninitialized
276 - # Lamport buffer and produce NaN. This only fires for
277 - # num_tokens <= PDL_ADVANCE_LAUNCH_TOKENS (the batch=1 / spec-decode
278 - # shapes, where the one-shot path is always selected). Complete at
279 - # the end for the one-shot path; the two-shot path is synchronized
280 - # and keeps the early completion. Related one-shot instability in
281 - # the same kernel: flashinfer-ai/flashinfer#1223.
282 - trigger_completion_at_end=(use_oneshot is True)
283 - or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
284 - )
292 + # vllm's one-shot size table is tuned for the trtllm backend. For
293 + # mnnvl, flashinfer caps the workspace's one-shot region at its own
294 + # MNNVL_ONE_SHOT_THRESHOLD (~4MB, e.g. 44 tokens at hidden 6144 bf16
295 + # TP8); forcing use_oneshot=True past that raises ValueError instead
296 + # of using two-shot, so every decode capture between the one-shot cap
297 + # and the fusion threshold silently lost fusion. Let flashinfer's
298 + # AUTO heuristic pick the strategy (it matches its own workspace
299 + # sizing), and keep trigger_completion_at_end=True since AUTO may
300 + # select one-shot (see the Lamport NaN note below).
301 + mnnvl_auto = workspace.backend == "mnnvl"
302 + if mnnvl_auto:
303 + use_oneshot = None
304 +
305 + try:
306 + flashinfer_comm.allreduce_fusion(
307 + input=allreduce_in,
308 + workspace=workspace,
309 + pattern=pattern_code,
310 + launch_with_pdl=launch_with_pdl,
311 + output=None,
312 + residual_out=residual_out,
313 + norm_out=norm_out,
314 + quant_out=quant_out,
315 + scale_out=scale_out,
316 + residual_in=residual,
317 + rms_gamma=rms_gamma,
318 + rms_eps=rms_eps,
319 + scale_factor=scale_factor,
320 + layout_code=layout_code,
321 + use_oneshot=use_oneshot,
322 + fp32_acc=fp32_acc,
323 + weight_bias=weight_bias,
324 + # The one-shot Lamport all-reduce signals PDL completion before
325 + # its output buffer is committed when trigger_completion_at_end
326 + # is False, so the next PDL-launched kernel can read the
327 + # uninitialized Lamport buffer and produce NaN. This only fires
328 + # for num_tokens <= PDL_ADVANCE_LAUNCH_TOKENS (the batch=1 /
329 + # spec-decode shapes, where the one-shot path is always
330 + # selected). Complete at the end for the one-shot path; the
331 + # two-shot path is synchronized and keeps the early completion.
332 + # Related one-shot instability in the same kernel:
333 + # flashinfer-ai/flashinfer#1223.
334 + trigger_completion_at_end=True
335 + if mnnvl_auto
336 + else (use_oneshot or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS),
337 + )
338 + except ValueError as e:
339 + # The globally cached workspace can be smaller than this call's
340 + # budget (it is created once by whichever caller runs first, and
341 + # the mnnvl workspace's is_buffer_size_sufficient() checks trtllm
342 + # metadata, so vllm-side accounting cannot see the real capacity).
343 + # flashinfer raises before launching anything, so falling back
344 + # mid-capture is safe. Observed with spec-decode draft graphs
345 + # whose token count (max_num_reqs * (1 + num_spec)) exceeds the
346 + # workspace sized for the fusion threshold.
347 + assert not is_quant_pattern, (
348 + "flashinfer fused allreduce failed and the unfused fallback "
349 + f"does not cover quant patterns: {e}"
350 + )
351 + logger.warning_once(
352 + "flashinfer fused allreduce+rmsnorm unavailable for this "
353 + "shape (num_tokens=%d, max_token_num=%d, world_size=%d): %s. "
354 + "Using the unfused fallback.",
355 + num_tokens,
356 + max_token_num,
357 + world_size,
358 + e,
359 + )
360 + unfused_fallback()
361
362 def call_trtllm_fused_allreduce_norm_fake(
363 allreduce_in: torch.Tensor,
vllm/config/speculative.py
+7 -1
@@ -330,6 +330,7 @@ class SpeculativeConfig:
330 @staticmethod
331 def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
332 initial_architecture = hf_config.architectures[0]
333 + use_sparse_mtp = hf_config.model_type == "glm_moe_dsa"
334 if hf_config.model_type in (
335 "deepseek_v3",
336 "deepseek_v32",
@@ -339,7 +340,12 @@ class SpeculativeConfig:
340 if hf_config.model_type == "deepseek_mtp":
341 n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
342 hf_config.update(
342 - {"n_predict": n_predict, "architectures": ["DeepSeekMTPModel"]}
343 + {
344 + "n_predict": n_predict,
345 + "architectures": [
346 + "DeepseekV32MTPModel" if use_sparse_mtp else "DeepSeekMTPModel"
347 + ],
348 + }
349 )
350 if hf_config.model_type == "deepseek_v4":
351 hf_config.model_type = "deepseek_mtp"
vllm/cute_utils/__init__.py
+5
@@ -33,6 +33,11 @@ EVICT_NORMAL = Int64(0x1000000000000000)
33 EVICT_FIRST = Int64(0x12F0000000000000)
34 EVICT_LAST = Int64(0x14F0000000000000)
35
36 +TORCH_TO_CUTE_DTYPE = {
37 + torch.bfloat16: BFloat16,
38 + torch.float32: Float32,
39 +}
40 +
41
42 @dsl_user_op
43 def recast_val(x, dtype, *, loc=None, ip=None):
vllm/cute_utils/cvt.py
+22 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from cutlass import Constexpr, Float32, Uint32, cute
3 +from cutlass import BFloat16, Constexpr, Float32, Uint16, Uint32, cute
4 from cutlass._mlir import ir
5 from cutlass._mlir.dialects import llvm, vector
6 from cutlass.cutlass_dsl import T, dsl_user_op
@@ -38,6 +38,11 @@ def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]:
38 )
39
40 elif isinstance(data, (cute.Tensor, cute.TensorSSA)):
41 + # recast to 32-bit registers if needed
42 + if data.element_type == BFloat16:
43 + data = cute.recast_tensor(data, Uint32)
44 + assert data.element_type == Uint32
45 +
46 # NOTE: the output is always 1D
47 size = cute.size(data.shape)
48 out = cute.make_rmem_tensor(size * 2, Float32)
@@ -84,6 +89,22 @@ def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
89 return cute.TensorSSA(vec, 2, Uint32)
90
91
92 +@dsl_user_op
93 +def fp32x2_to_fp8x2(a0: Float32, a1: Float32, *, loc=None, ip=None) -> Uint16:
94 + out = llvm.inline_asm(
95 + T.i16(),
96 + [
97 + a0.ir_value(loc=loc, ip=ip),
98 + a1.ir_value(loc=loc, ip=ip),
99 + ],
100 + "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;",
101 + "=h,f,f",
102 + has_side_effects=False,
103 + is_align_stack=False,
104 + )
105 + return Uint16(out)
106 +
107 +
108 @dsl_user_op
109 def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
110 out = llvm.inline_asm(
vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
+15 -4
@@ -43,7 +43,7 @@ try:
43 if flashinfer_comm is not None
44 else None
45 )
46 -except ImportError:
46 +except (ImportError, AttributeError):
47 flashinfer_trtllm_fused_allreduce_norm = None # type: ignore[assignment]
48 get_fi_ar_workspace = None # type: ignore[assignment]
49 _AR_RESIDUAL_RMS_NORM = None
@@ -52,12 +52,23 @@ except ImportError:
52 _FI_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16)
53
54
55 +@torch.compiler.assume_constant_result
56 +def _fi_ar_max_size_mb() -> dict[int, float]:
57 + """Flashinfer all-reduce fusion size table for the current device.
58 +
59 + Device capability is constant; marking the result constant keeps this out
60 + of the traced graph. Otherwise the per-forward call reaches
61 + ``current_platform.get_device_capability()`` inside ``torch.compile`` and
62 + graph-breaks ("can't handle functions not implemented in python")."""
63 + from vllm.config.compilation import PassConfig
64 +
65 + return PassConfig.default_fi_allreduce_fusion_max_size_mb()
66 +
67 +
68 def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None:
69 """Workspace token budget for flashinfer fused all-reduce, or None if the
70 current world size / device is unsupported. Mirrors ``FlashInferAllReduce``."""
58 - from vllm.config.compilation import PassConfig
59 -
60 - max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
71 + max_size_mb = _fi_ar_max_size_mb().get(tp_size)
72 if not max_size_mb:
73 return None
74 element_size = torch.tensor([], dtype=dtype).element_size()
vllm/model_executor/models/registry.py
+2 -1
@@ -113,7 +113,7 @@ _TEXT_GENERATION_MODELS = {
113 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
114 "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
115 "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
116 - "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
116 + "GlmMoeDsaForCausalLM": ("vllm.models.deepseek_v32", "GlmMoeDsaForCausalLM"),
117 "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
118 "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
119 "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
@@ -628,6 +628,7 @@ _SPECULATIVE_DECODING_MODELS = {
628 "Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
629 "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
630 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
631 + "DeepseekV32MTPModel": ("vllm.models.deepseek_v32", "DeepseekV32MTP"),
632 "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
633 "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
634 "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"),
vllm/models/deepseek_v32/__init__.py
+18 -4
@@ -6,17 +6,31 @@ DeepSeek V3.2 introduced the DeepSeek Sparse Attention (DSA) architecture:
6 MLA + a "lightning indexer" that selects the top-k tokens for a sparse MLA
7 attend. The same model code serves any DSA checkpoint, including GLM-5.2
8 (``glm_moe_dsa``), which reuses this architecture.
9 +
10 +The optimized kernels under ``nvidia/`` target the Blackwell (SM100) family.
11 +Every other platform — ROCm, XPU, pre-SM100 CUDA (e.g. H100), CPU — falls back
12 +to the generic ``deepseek_v2`` implementation, which already handles the DSA
13 +(index_topk) architecture and is ``torch.compile``-friendly there. This matches
14 +main's behavior on those platforms (no hard failure).
15 """
16
17 from vllm.platforms import current_platform
18
13 -if current_platform.is_rocm() or current_platform.is_xpu():
14 - raise NotImplementedError("deepseek_v32 currently supports NVIDIA SM100 only.")
19 +if current_platform.is_cuda() and current_platform.is_device_capability_family(100):
20 + from .nvidia.model import DeepseekV32ForCausalLM
21 + from .nvidia.mtp import DeepseekV32MTP
22
16 -from .nvidia.model import DeepseekV32ForCausalLM
17 -from .nvidia.mtp import DeepseekV32MTP
23 + # GLM-5.2 (glm_moe_dsa) reuses the same optimized DSA module on SM100.
24 + GlmMoeDsaForCausalLM = DeepseekV32ForCausalLM
25 +else:
26 + from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepseekV32MTP
27 + from vllm.model_executor.models.deepseek_v2 import (
28 + DeepseekV3ForCausalLM as DeepseekV32ForCausalLM,
29 + )
30 + from vllm.model_executor.models.deepseek_v2 import GlmMoeDsaForCausalLM
31
32 __all__ = [
33 "DeepseekV32ForCausalLM",
34 "DeepseekV32MTP",
35 + "GlmMoeDsaForCausalLM",
36 ]
vllm/models/deepseek_v32/nvidia/attention.py
+72 -12
@@ -4,6 +4,7 @@ import torch
4 import torch.nn as nn
5 from transformers import DeepseekV2Config, DeepseekV3Config
6
7 +import vllm._custom_ops as ops
8 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
9 from vllm.config import CacheConfig, VllmConfig
10 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -21,17 +22,15 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
22 per_token_group_quant_fp8,
23 )
24 from vllm.model_executor.layers.rotary_embedding import get_rope
24 -from vllm.model_executor.layers.sparse_attn_indexer import (
25 - SparseAttnIndexer,
26 - sparse_attn_indexer,
27 -)
25 +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
26 from vllm.model_executor.models.deepseek_v2 import (
27 DeepSeekV2FusedQkvAProjLinear,
28 DeepseekV32IndexerCache,
29 yarn_get_mscale,
30 )
31 from vllm.model_executor.models.utils import extract_layer_index
34 -from vllm.utils.torch_utils import is_quantized_kv_cache
32 +from vllm.platforms import current_platform
33 +from vllm.utils.torch_utils import _encode_layer_name, is_quantized_kv_cache
34
35 from .kernels import fused_norm_rope, fused_q
36
@@ -171,6 +170,11 @@ class DeepseekV32Attention(MLAAttention):
170 ) -> None:
171 quant_config = vllm_config.quant_config
172 cache_config = vllm_config.cache_config
173 + if cache_config is not None and cache_config.cache_dtype == "auto":
174 + # This implementation requires an FP8 sparse cache. Start with the
175 + # generic FP8 format; MLAAttention will canonicalize it for the
176 + # selected sparse backend.
177 + cache_config.cache_dtype = "fp8"
178
179 hidden_size = config.hidden_size
180 qk_nope_head_dim = config.qk_nope_head_dim
@@ -331,6 +335,37 @@ class DeepseekV32Attention(MLAAttention):
335 rope_parameters=config.rope_parameters,
336 is_neox_style=False,
337 )
338 + # Decode-M GEMM dispatch for the bf16 A/B projections, measured on
339 + # B300 and B200 (graph replay + rotating weights vs fp64):
340 + # bf16_skinny_gemm for M <= skinny_max, dsv3_fused_a_gemm up to
341 + # M = 16, cuBLAS above. Blackwell-only: the boundaries and the
342 + # fused_a tile choices are tied to the 148/152-SM single-wave
343 + # geometry. DSv3.2 q_b has no skinny tier (K=1536 does not fit the
344 + # GEMV's 128x8 K-step; fused_a still wins 1.6-1.8x there).
345 + # Quantized linear methods may not register .weight until
346 + # process_weights_after_loading (e.g. compressed-tensors NVFP4
347 + # registers weight_packed at init) — probe with getattr.
348 + # (N, K) -> skinny_max
349 + qkv_a_dispatch = {(2624, 6144): 2, (2112, 7168): 2}
350 + q_b_dispatch = {(2048, 2048): 2, (3072, 1536): 0}
351 + ok = (
352 + current_platform.is_device_capability_family(100)
353 + and hasattr(torch.ops._C, "bf16_skinny_gemm")
354 + and hasattr(torch.ops._C, "dsv3_fused_a_gemm")
355 + )
356 + qkv_a_weight = getattr(self.fused_qkv_a_proj, "weight", None)
357 + q_b_weight = getattr(self.q_b_proj, "weight", None)
358 + self._qkv_a_skinny_max = (
359 + qkv_a_dispatch.get(tuple(qkv_a_weight.shape))
360 + if ok and qkv_a_weight is not None and qkv_a_weight.dtype == torch.bfloat16
361 + else None
362 + )
363 + self._q_b_skinny_max = (
364 + q_b_dispatch.get(tuple(q_b_weight.shape))
365 + if ok and q_b_weight is not None and q_b_weight.dtype == torch.bfloat16
366 + else None
367 + )
368 +
369 # Lightning indexer uses its own RoPE; interleave maps to non-NeoX.
370 self.indexer_rope_emb = get_rope(
371 qk_rope_head_dim,
@@ -339,13 +374,32 @@ class DeepseekV32Attention(MLAAttention):
374 is_neox_style=not getattr(config, "indexer_rope_interleave", False),
375 )
376
377 + def _decode_m_gemm(
378 + self, x: torch.Tensor, weight: torch.Tensor, skinny_max: int
379 + ) -> torch.Tensor:
380 + """x @ weight.T for decode M <= 16."""
381 + if x.shape[0] == 0:
382 + return torch.empty((0, weight.shape[0]), dtype=x.dtype, device=x.device)
383 + if x.shape[0] <= skinny_max:
384 + return ops.bf16_skinny_gemm(x, weight)
385 + out = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device)
386 + ops.dsv3_fused_a_gemm(out, x, weight.t())
387 + return out
388 +
389 def forward( # type: ignore[override]
390 self,
391 positions: torch.Tensor,
392 hidden_states: torch.Tensor,
393 ) -> torch.Tensor:
394 # Captured: A-projections (+ indexer A-GEMM on indexer layers).
348 - qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
395 + if self._qkv_a_skinny_max is not None and hidden_states.shape[0] <= 16:
396 + qkv_lora = self._decode_m_gemm(
397 + hidden_states,
398 + self.fused_qkv_a_proj.weight,
399 + self._qkv_a_skinny_max,
400 + )
401 + else:
402 + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
403 q_c, kv_c, k_pe = qkv_lora.split(
404 [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
405 )
@@ -398,7 +452,9 @@ class DeepseekV32Attention(MLAAttention):
452 assert isinstance(slot_mapping, dict)
453 mla_slot = slot_mapping.get(self.layer_name)
454
401 - if self.indexer is not None:
455 + # skip_topk (index_share_for_mtp_iteration reuse mode): don't run the
456 + # indexer this step; keep the top-k already in the shared buffer.
457 + if self.indexer is not None and not self.skip_topk:
458 has_indexer = True
459 indexer_k_norm_w = self.indexer.k_norm.weight
460 indexer_k_norm_bias = self.indexer.k_norm.bias
@@ -451,12 +507,16 @@ class DeepseekV32Attention(MLAAttention):
507 index_rope_interleave=self._index_rope_interleave,
508 )
509
454 - q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
510 + if self._q_b_skinny_max is not None and q_c.shape[0] <= 16:
511 + q = self._decode_m_gemm(q_c, self.q_b_proj.weight, self._q_b_skinny_max)
512 + else:
513 + q = self.q_b_proj(q_c)[0]
514 + q = q.view(-1, self.num_local_heads, self.qk_head_dim)
515 q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
516 q_nope = q_nope.transpose(0, 1)
517 ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1)
518
459 - if self.indexer is not None:
519 + if self.indexer is not None and not self.skip_topk:
520 index_q = self.indexer.wq_b(q_c)[0]
521 index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim)
522 else:
@@ -478,10 +538,10 @@ class DeepseekV32Attention(MLAAttention):
538 quantize_mqa=self._fp8_query,
539 )
540
481 - if self.indexer is not None:
482 - sparse_attn_indexer(
541 + if self.indexer is not None and not self.skip_topk:
542 + torch.ops.vllm.sparse_attn_indexer(
543 q_c,
484 - self.indexer.k_cache.prefix,
544 + _encode_layer_name(self.indexer.k_cache.prefix),
545 self.indexer.k_cache.kv_cache,
546 index_q_fp8,
547 None, # q_scale folded into weights on the fp8 path
vllm/models/deepseek_v32/nvidia/kernels.py
+284 -16
@@ -2,13 +2,80 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 +from vllm.platforms import current_platform
6 from vllm.triton_utils import tl, triton
7 +from vllm.utils.import_utils import has_cutedsl
8 +from vllm.utils.torch_utils import direct_register_custom_op
9
10 # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the
11 # has_indexer=False path so the indexer args don't allocate every call.
12 _DUMMY_CACHE: dict[tuple, torch.Tensor] = {}
13
14
15 +@torch.compiler.assume_constant_result
16 +def _can_use_fused_q_cutedsl() -> bool:
17 + return current_platform.has_device_capability(100) and has_cutedsl()
18 +
19 +
20 +@torch.compiler.assume_constant_result
21 +def _is_arch_support_pdl() -> bool:
22 + return (
23 + current_platform.has_device_capability(100)
24 + and current_platform.is_arch_support_pdl()
25 + )
26 +
27 +
28 +def _fused_q_cutedsl_impl(
29 + positions: torch.Tensor,
30 + q_pe: torch.Tensor,
31 + rope_cache: torch.Tensor,
32 + ql_nope: torch.Tensor,
33 + q_scale: torch.Tensor,
34 + mqa_output: torch.Tensor,
35 + idx_q: torch.Tensor,
36 + idx_rope_cache: torch.Tensor,
37 + idx_weights: torch.Tensor,
38 + idx_weights_softmax_scale: float,
39 + idx_weights_head_scale: float,
40 + idx_q_fp8: torch.Tensor,
41 + idx_weights_out: torch.Tensor,
42 + has_indexer: bool,
43 + index_rope_interleave: bool,
44 +) -> None:
45 + from .ops.fused_q_cutedsl import fused_q_cutedsl
46 +
47 + fused_q_cutedsl(
48 + positions,
49 + q_pe,
50 + rope_cache,
51 + ql_nope,
52 + q_scale,
53 + mqa_output,
54 + idx_q,
55 + idx_rope_cache,
56 + idx_weights,
57 + idx_weights_softmax_scale,
58 + idx_weights_head_scale,
59 + idx_q_fp8,
60 + idx_weights_out,
61 + has_indexer=has_indexer,
62 + index_rope_interleave=index_rope_interleave,
63 + )
64 +
65 +
66 +def _fused_q_cutedsl_fake(*args, **kwargs) -> None:
67 + pass
68 +
69 +
70 +direct_register_custom_op(
71 + op_name="fused_q_cutedsl",
72 + op_func=_fused_q_cutedsl_impl,
73 + mutates_args=["mqa_output", "idx_q_fp8", "idx_weights_out"],
74 + fake_impl=_fused_q_cutedsl_fake,
75 + dispatch_key="CUDA",
76 +)
77 +
78 +
79 def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
80 key = (shape, dtype, device)
81 t = _DUMMY_CACHE.get(key)
@@ -147,9 +214,14 @@ def _fused_norm_rope_kernel(
214 TOPK_BLOCK_SIZE: tl.constexpr,
215 HAS_INDEXER: tl.constexpr,
216 INDEX_ROPE_INTERLEAVE: tl.constexpr,
217 + USE_PDL: tl.constexpr,
218 ):
219 pid = tl.program_id(0)
220 tok_idx = tl.program_id(1)
221 + if USE_PDL:
222 + tl.extra.cuda.gdc_wait()
223 + tl.extra.cuda.gdc_launch_dependents()
224 +
225 if pid == 3:
226 if not HAS_INDEXER:
227 # Shared layer: reuse the previous indexer layer's top-k; do not
@@ -166,6 +238,17 @@ def _fused_norm_rope_kernel(
238 )
239 return
240
241 + if pid == 2:
242 + # Q RMS norm does not depend on cache availability. In particular,
243 + # profiling uses a negative dummy slot mapping but still consumes Q.
244 + q_block = tl.arange(0, Q_BLOCK_SIZE)
245 + q_mask = q_block < Q_DIM
246 + q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
247 + q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
248 + q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
249 + tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
250 + return
251 +
252 if slot_mapping_ptr is None:
253 # Memory profiling run.
254 return
@@ -174,15 +257,7 @@ def _fused_norm_rope_kernel(
257 # Padding
258 return
259
177 - if pid == 2:
178 - # Q RMS norm
179 - q_block = tl.arange(0, Q_BLOCK_SIZE)
180 - q_mask = q_block < Q_DIM
181 - q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
182 - q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
183 - q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
184 - tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
185 - elif pid == 1:
260 + if pid == 1:
261 # KV RMS Norm + KV RoPE + MLA concat_and_cache.
262 # Merged so the normed kv_c and RoPE'd k_pe can be written
263 # to the MLA KV cache directly without a separate kernel.
@@ -211,9 +286,6 @@ def _fused_norm_rope_kernel(
286 r2 = x2 * cos + x1 * sin
287
288 # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.
214 - if mla_cache_entry_stride == 0:
215 - return
216 -
289 mla_block_size = mla_cache_block_stride // mla_cache_entry_stride
290 mla_block_idx = slot_idx // mla_block_size
291 mla_block_off = slot_idx % mla_block_size
@@ -363,7 +435,7 @@ def _fused_norm_rope_kernel(
435 )
436
437
366 -def fused_norm_rope(
438 +def _fused_norm_rope_impl(
439 positions: torch.Tensor,
440 q_c: torch.Tensor,
441 q_rms_norm_w: torch.Tensor,
@@ -465,12 +537,17 @@ def fused_norm_rope(
537 # Dummy values — pid 2 will skip the MLA cache write because
538 # slot_mapping is all -1.
539 mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)
468 - mla_block_stride = 0
469 - mla_entry_stride = 0
540 + # Torch's Triton HOP mutation analysis compiles every program-id
541 + # branch even though the dummy slot mapping returns before cache
542 + # access. Keep the unused strides nonzero so that analysis can lower
543 + # the block-size calculation safely.
544 + mla_block_stride = 1
545 + mla_entry_stride = 1
546 mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
547
548 if q_c_out is None:
549 q_c_out = torch.empty_like(q_c)
550 + use_pdl = _is_arch_support_pdl()
551 _fused_norm_rope_kernel[(4, num_tokens)](
552 positions,
553 # Q RMS norm
@@ -529,6 +606,157 @@ def fused_norm_rope(
606 TOPK_BLOCK_SIZE=1024,
607 HAS_INDEXER=has_indexer,
608 INDEX_ROPE_INTERLEAVE=index_rope_interleave,
609 + USE_PDL=use_pdl,
610 + launch_pdl=use_pdl,
611 + )
612 + return q_c_out
613 +
614 +
615 +def _fused_norm_rope_op(
616 + positions: torch.Tensor,
617 + q_c: torch.Tensor,
618 + q_rms_norm_w: torch.Tensor,
619 + q_rms_eps: float,
620 + kv_c: torch.Tensor,
621 + kv_rms_norm_w: torch.Tensor,
622 + kv_rms_eps: float,
623 + k_pe: torch.Tensor,
624 + k_rope_cos_sin_cache: torch.Tensor,
625 + index_k: torch.Tensor | None,
626 + index_k_layer_norm_w: torch.Tensor | None,
627 + index_k_layer_norm_bias: torch.Tensor | None,
628 + index_k_layer_norm_eps: float,
629 + index_k_rope_cos_sin_cache: torch.Tensor | None,
630 + topk_indices_buffer: torch.Tensor,
631 + slot_mapping: torch.Tensor | None,
632 + indexer_k_cache: torch.Tensor | None,
633 + mla_kv_cache: torch.Tensor | None,
634 + mla_kv_cache_dtype: str,
635 + mla_k_scale: torch.Tensor | None,
636 + has_indexer: bool,
637 + index_rope_interleave: bool,
638 + q_c_out: torch.Tensor,
639 +) -> None:
640 + _fused_norm_rope_impl(
641 + positions,
642 + q_c,
643 + q_rms_norm_w,
644 + q_rms_eps,
645 + kv_c,
646 + kv_rms_norm_w,
647 + kv_rms_eps,
648 + k_pe,
649 + k_rope_cos_sin_cache,
650 + index_k,
651 + index_k_layer_norm_w,
652 + index_k_layer_norm_bias,
653 + index_k_layer_norm_eps,
654 + index_k_rope_cos_sin_cache,
655 + topk_indices_buffer,
656 + slot_mapping=slot_mapping,
657 + indexer_k_cache=indexer_k_cache,
658 + mla_kv_cache=mla_kv_cache,
659 + mla_kv_cache_dtype=mla_kv_cache_dtype,
660 + mla_k_scale=mla_k_scale,
661 + has_indexer=has_indexer,
662 + index_rope_interleave=index_rope_interleave,
663 + q_c_out=q_c_out,
664 + )
665 +
666 +
667 +def _fused_norm_rope_fake(
668 + positions: torch.Tensor,
669 + q_c: torch.Tensor,
670 + q_rms_norm_w: torch.Tensor,
671 + q_rms_eps: float,
672 + kv_c: torch.Tensor,
673 + kv_rms_norm_w: torch.Tensor,
674 + kv_rms_eps: float,
675 + k_pe: torch.Tensor,
676 + k_rope_cos_sin_cache: torch.Tensor,
677 + index_k: torch.Tensor | None,
678 + index_k_layer_norm_w: torch.Tensor | None,
679 + index_k_layer_norm_bias: torch.Tensor | None,
680 + index_k_layer_norm_eps: float,
681 + index_k_rope_cos_sin_cache: torch.Tensor | None,
682 + topk_indices_buffer: torch.Tensor,
683 + slot_mapping: torch.Tensor | None,
684 + indexer_k_cache: torch.Tensor | None,
685 + mla_kv_cache: torch.Tensor | None,
686 + mla_kv_cache_dtype: str,
687 + mla_k_scale: torch.Tensor | None,
688 + has_indexer: bool,
689 + index_rope_interleave: bool,
690 + q_c_out: torch.Tensor,
691 +) -> None:
692 + return
693 +
694 +
695 +direct_register_custom_op(
696 + op_name="fused_norm_rope_deepseek_v32",
697 + op_func=_fused_norm_rope_op,
698 + mutates_args=[
699 + "topk_indices_buffer",
700 + "indexer_k_cache",
701 + "mla_kv_cache",
702 + "q_c_out",
703 + ],
704 + fake_impl=_fused_norm_rope_fake,
705 + dispatch_key="CUDA",
706 +)
707 +
708 +
709 +def fused_norm_rope(
710 + positions: torch.Tensor,
711 + q_c: torch.Tensor,
712 + q_rms_norm_w: torch.Tensor,
713 + q_rms_eps: float,
714 + kv_c: torch.Tensor,
715 + kv_rms_norm_w: torch.Tensor,
716 + kv_rms_eps: float,
717 + k_pe: torch.Tensor,
718 + k_rope_cos_sin_cache: torch.Tensor,
719 + index_k: torch.Tensor | None,
720 + index_k_layer_norm_w: torch.Tensor | None,
721 + index_k_layer_norm_bias: torch.Tensor | None,
722 + index_k_layer_norm_eps: float,
723 + index_k_rope_cos_sin_cache: torch.Tensor | None,
724 + topk_indices_buffer: torch.Tensor,
725 + slot_mapping: torch.Tensor | None = None,
726 + indexer_k_cache: torch.Tensor | None = None,
727 + mla_kv_cache: torch.Tensor | None = None,
728 + mla_kv_cache_dtype: str = "auto",
729 + mla_k_scale: torch.Tensor | None = None,
730 + has_indexer: bool = True,
731 + index_rope_interleave: bool = False,
732 + q_c_out: torch.Tensor | None = None,
733 +) -> torch.Tensor:
734 + if q_c_out is None:
735 + q_c_out = torch.empty_like(q_c)
736 + torch.ops.vllm.fused_norm_rope_deepseek_v32(
737 + positions,
738 + q_c,
739 + q_rms_norm_w,
740 + q_rms_eps,
741 + kv_c,
742 + kv_rms_norm_w,
743 + kv_rms_eps,
744 + k_pe,
745 + k_rope_cos_sin_cache,
746 + index_k,
747 + index_k_layer_norm_w,
748 + index_k_layer_norm_bias,
749 + index_k_layer_norm_eps,
750 + index_k_rope_cos_sin_cache,
751 + topk_indices_buffer,
752 + slot_mapping,
753 + indexer_k_cache,
754 + mla_kv_cache,
755 + mla_kv_cache_dtype,
756 + mla_k_scale,
757 + has_indexer,
758 + index_rope_interleave,
759 + q_c_out,
760 )
761 return q_c_out
762
@@ -582,10 +810,14 @@ def _fused_q_kernel(
810 HAS_INDEXER: tl.constexpr,
811 INDEX_ROPE_INTERLEAVE: tl.constexpr,
812 QUANTIZE_MQA: tl.constexpr,
813 + USE_PDL: tl.constexpr,
814 ):
815 pid = tl.program_id(0)
816 tok_idx = tl.program_id(1)
817 head_idx = tl.program_id(2)
818 + if USE_PDL:
819 + tl.extra.cuda.gdc_wait()
820 + tl.extra.cuda.gdc_launch_dependents()
821
822 if pid == 2:
823 # ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16
@@ -773,10 +1005,12 @@ def fused_q(
1005 ``(ql_nope, q_pe)`` tuple the backend expects.
1006 """
1007 assert positions.ndim == 1
1008 + assert positions.dtype == torch.int64
1009 assert q_pe.ndim == 3
1010 assert q_pe_cos_sin_cache.ndim == 2
1011 assert ql_nope.ndim == 3
1012 assert ql_nope.shape[:2] == q_pe.shape[:2]
1013 + assert q_scale.dtype == torch.float32 and q_scale.numel() == 1
1014
1015 num_tokens = positions.shape[0]
1016 num_q_heads = q_pe.shape[1]
@@ -794,7 +1028,18 @@ def fused_q(
1028 assert index_weights is not None
1029 num_index_q_heads = index_q.shape[1]
1030 index_q_head_dim = index_q.shape[2]
797 - grid_heads = max(mqa_grid_heads, num_index_q_heads)
1031 + use_cutedsl = False
1032 + if _can_use_fused_q_cutedsl():
1033 + from .ops.fused_q_cutedsl import is_fused_q_cutedsl_supported
1034 +
1035 + use_cutedsl = is_fused_q_cutedsl_supported(
1036 + q_pe,
1037 + index_q,
1038 + ql_nope,
1039 + has_indexer=has_indexer,
1040 + quantize_mqa=quantize_mqa,
1041 + )
1042 + grid_heads = max(mqa_grid_heads, num_index_q_heads if has_indexer else 1)
1043 if quantize_mqa:
1044 # fp8 path: pack [ql_nope; q_pe] into a single fp8 tensor.
1045 mqa_q_fp8 = torch.empty(
@@ -815,6 +1060,27 @@ def fused_q(
1060
1061 index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
1062 index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
1063 + if use_cutedsl:
1064 + torch.ops.vllm.fused_q_cutedsl(
1065 + positions,
1066 + q_pe,
1067 + q_pe_cos_sin_cache,
1068 + ql_nope,
1069 + q_scale,
1070 + mqa_q,
1071 + index_q,
1072 + index_q_cos_sin_cache,
1073 + index_weights,
1074 + index_weights_softmax_scale,
1075 + index_weights_head_scale,
1076 + index_q_fp8,
1077 + index_weights_out,
1078 + has_indexer=has_indexer,
1079 + index_rope_interleave=index_rope_interleave,
1080 + )
1081 + return index_q_fp8, index_weights_out, mqa_q
1082 +
1083 + use_pdl = _is_arch_support_pdl()
1084 _fused_q_kernel[(3, num_tokens, grid_heads)](
1085 positions,
1086 q_pe,
@@ -856,6 +1122,8 @@ def fused_q(
1122 HAS_INDEXER=has_indexer,
1123 INDEX_ROPE_INTERLEAVE=index_rope_interleave,
1124 QUANTIZE_MQA=quantize_mqa,
1125 + USE_PDL=use_pdl,
1126 + launch_pdl=use_pdl,
1127 # num_warps=1 is optimal here: each program is a single 128-element
1128 # rope+quant, so the kernel is program-count/occupancy bound, not
1129 # per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).
vllm/models/deepseek_v32/nvidia/model.py
+4
@@ -6,6 +6,7 @@ from itertools import islice
6
7 import torch
8
9 +from vllm.compilation.decorators import support_torch_compile
10 from vllm.config import VllmConfig
11 from vllm.distributed import (
12 get_pp_group,
@@ -39,6 +40,7 @@ from vllm.model_executor.models.utils import (
40 sequence_parallel_chunk,
41 )
42 from vllm.sequence import IntermediateTensors
43 +from vllm.v1.attention.backends.mla.sparse_utils import register_phys_shadow
44
45 from .attention import DeepseekV32Attention
46 from .fused_ops import fused_allreduce_rms_norm
@@ -177,6 +179,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
179 return hidden_states, residual
180
181
182 +@support_torch_compile
183 class DeepseekV32Model(torch.nn.Module):
184 fall_back_to_pt_during_load = False
185
@@ -200,6 +203,7 @@ class DeepseekV32Model(torch.nn.Module):
203 dtype=torch.int32,
204 device=self.device,
205 )
206 + register_phys_shadow(topk_indices_buffer)
207
208 if get_pp_group().is_first_rank:
209 self.embed_tokens = VocabParallelEmbedding(
vllm/models/deepseek_v32/nvidia/mtp.py
+76 -6
@@ -6,9 +6,9 @@ from collections.abc import Callable, Iterable
6 import torch
7 import torch.nn as nn
8
9 +import vllm._custom_ops as ops
10 from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.config import VllmConfig
11 -from vllm.distributed import tensor_model_parallel_all_reduce
12 from vllm.model_executor.layers.fused_moe import (
13 fused_moe_make_expert_params_mapping,
14 )
@@ -40,7 +40,12 @@ from vllm.model_executor.models.utils import (
40 )
41 from vllm.platforms import current_platform
42 from vllm.sequence import IntermediateTensors
43 +from vllm.v1.attention.backends.mla.sparse_utils import (
44 + phys_shadow,
45 + register_phys_shadow,
46 +)
47
48 +from .fused_ops import fused_allreduce_rms_norm
49 from .kernels import fused_eh_norm
50 from .model import DeepseekV32DecoderLayer
51
@@ -56,6 +61,19 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
61 self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
62 self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
63 self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
64 + # bf16 skinny GEMM for the eh_proj, B300+B200 measured: GLM-5.2
65 + # (6144, 12288) wins at M <= 3, DSv3.2 (7168, 14336) at M <= 2;
66 + # cuBLAS holds above (the 151/205MB weights stream at ~5.7TB/s).
67 + eh_dispatch = {(6144, 12288): 3, (7168, 14336): 2}
68 + eh_weight = getattr(self.eh_proj, "weight", None)
69 + self._eh_skinny_max = (
70 + eh_dispatch.get(tuple(eh_weight.shape), 0)
71 + if current_platform.is_device_capability_family(100)
72 + and hasattr(torch.ops._C, "bf16_skinny_gemm")
73 + and eh_weight is not None
74 + and eh_weight.dtype == torch.bfloat16
75 + else 0
76 + )
77
78 topk_indices_buffer = torch.empty(
79 vllm_config.scheduler_config.max_num_batched_tokens,
@@ -63,6 +81,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
81 dtype=torch.int32,
82 device=current_platform.device_type,
83 )
84 + register_phys_shadow(topk_indices_buffer)
85 self.shared_head = SharedHead(
86 config=config, prefix=prefix, quant_config=quant_config
87 )
@@ -91,7 +110,13 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
110 self.hnorm.weight,
111 self.enorm.variance_epsilon,
112 )
94 - hidden_states = self.eh_proj(eh_input)
113 + if (
114 + eh_input.shape[0] <= self._eh_skinny_max
115 + and eh_input.dtype == torch.bfloat16
116 + ):
117 + hidden_states = ops.bf16_skinny_gemm(eh_input, self.eh_proj.weight)
118 + else:
119 + hidden_states = self.eh_proj(eh_input)
120 hidden_states, residual = self.mtp_block(
121 positions=positions, hidden_states=hidden_states, residual=None
122 )
@@ -101,9 +126,6 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
126 positions.shape[0],
127 is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe,
128 )
104 - if not self.mtp_block.use_sequence_parallel_moe:
105 - # Without sequence parallelism, the MoE output is left un-reduced.
106 - hidden_states = tensor_model_parallel_all_reduce(hidden_states)
129 # Recycle the POST-final-norm hidden into the next draft step. The
130 # residual-add is fused into the final RMSNorm so it is computed
131 # exactly once, and the result is returned for both tuple positions:
@@ -114,7 +136,14 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
136 # is understood by both the V2 speculator (isinstance-tuple check) and
137 # the legacy proposer (model_returns_tuple is True for the
138 # DeepSeekMTPModel architecture).
117 - hidden_states, _ = self.shared_head.norm(hidden_states, residual)
139 + if self.mtp_block.use_sequence_parallel_moe:
140 + hidden_states, _ = self.shared_head.norm(hidden_states, residual)
141 + else:
142 + # The MoE output is left un-reduced; fuse its all-reduce into the
143 + # final norm, as the main model does at layer boundaries.
144 + hidden_states, _ = fused_allreduce_rms_norm(
145 + hidden_states, residual, self.shared_head.norm
146 + )
147 return hidden_states, hidden_states
148
149
@@ -157,6 +186,15 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
186 if self_attn is not None and hasattr(self_attn, "topk_indices_buffer"):
187 topk_indices_buffer = self_attn.topk_indices_buffer
188 topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids]
189 + # The physical-index shadow mirrors this buffer row-wise
190 + # (shadow[j] == convert(logical[j])); re-arrange it the same
191 + # way or skip_topk draft steps read the step-0 multi-token
192 + # layout's rows for the wrong tokens/requests.
193 + shadow = phys_shadow(topk_indices_buffer)
194 + if shadow is not None:
195 + phys, seq = shadow
196 + phys[:num_slots] = phys[slot_ids]
197 + seq[:num_slots] = seq[slot_ids]
198
199 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
200 return self.embed_tokens(input_ids)
@@ -192,6 +230,19 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
230 # second RMSNorm.
231 return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
232
233 + def get_top_tokens(
234 + self,
235 + hidden_states: torch.Tensor,
236 + spec_step_idx: int = 0,
237 + ) -> torch.Tensor:
238 + # Greedy draft token ids via vocab-parallel local argmax: skips the
239 + # full-vocab logits all-gather (use_local_argmax_reduction).
240 + current_step_idx = spec_step_idx % self.num_mtp_layers
241 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
242 + return self.logits_processor.get_top_tokens(
243 + mtp_layer.shared_head.head, hidden_states
244 + )
245 +
246
247 class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
248 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -240,6 +291,25 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
291 ) -> torch.Tensor | None:
292 return self.model.compute_logits(hidden_states, spec_step_idx)
293
294 + def get_top_tokens(
295 + self,
296 + hidden_states: torch.Tensor,
297 + spec_step_idx: int | None = None,
298 + ) -> torch.Tensor:
299 + # The V2 speculator calls this without spec_step_idx; that is only
300 + # correct while a single MTP layer is cycled for every draft step
301 + # (num_nextn_predict_layers == 1, as in GLM-5.2/DSV3.2). Guard the
302 + # assumption so a future multi-layer MTP fails loudly instead of
303 + # silently using step 0's head for every step.
304 + if spec_step_idx is None:
305 + assert self.model.num_mtp_layers == 1, (
306 + "get_top_tokens called without spec_step_idx on a "
307 + "multi-layer MTP; thread the draft step index through "
308 + "the speculator."
309 + )
310 + spec_step_idx = 0
311 + return self.model.get_top_tokens(hidden_states, spec_step_idx)
312 +
313 def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
314 spec_layer_weight_names = [
315 "embed_tokens",
vllm/models/deepseek_v32/nvidia/ops/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py new
+532
@@ -0,0 +1,532 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from functools import cache
4 +
5 +import cutlass
6 +import cutlass.cute as cute
7 +import torch
8 +from cuda.bindings.driver import CUstream
9 +from cutlass import BFloat16, Float8E4M3FN, Float32, Int64, Uint8, Uint16, Uint32
10 +
11 +from vllm.cute_utils import TORCH_TO_CUTE_DTYPE, cvt
12 +
13 +
14 +def _make_fake_tensor(dtype, shape, divisibility):
15 + stride = tuple(
16 + 1 if i == len(shape) - 1 else cute.sym_int64(divisibility=divisibility)
17 + for i in range(len(shape))
18 + )
19 + return cute.runtime.make_fake_tensor(
20 + dtype,
21 + shape,
22 + stride,
23 + assumed_align=divisibility * dtype.width // 8,
24 + )
25 +
26 +
27 +def is_fused_q_cutedsl_supported(
28 + q_pe: torch.Tensor,
29 + index_q: torch.Tensor | None,
30 + ql_nope: torch.Tensor,
31 + *,
32 + has_indexer: bool,
33 + quantize_mqa: bool,
34 +) -> bool:
35 + if not (
36 + quantize_mqa
37 + and q_pe.dtype == ql_nope.dtype == torch.bfloat16
38 + and q_pe.shape[-1] == 64
39 + and ql_nope.shape[-1] == 512
40 + ):
41 + return False
42 + return not has_indexer or (
43 + index_q is not None
44 + and index_q.dtype == torch.bfloat16
45 + and index_q.shape[1] % 16 == 0
46 + and index_q.shape[-1] == 128
47 + )
48 +
49 +
50 +def fused_q_cutedsl(
51 + positions: torch.Tensor,
52 + q_pe: torch.Tensor,
53 + rope_cache: torch.Tensor,
54 + ql_nope: torch.Tensor,
55 + q_scale: torch.Tensor,
56 + mqa_output: torch.Tensor,
57 + idx_q: torch.Tensor,
58 + idx_rope_cache: torch.Tensor,
59 + idx_weights: torch.Tensor,
60 + idx_weights_softmax_scale: float,
61 + idx_weights_head_scale: float,
62 + idx_q_fp8: torch.Tensor,
63 + idx_weights_out: torch.Tensor,
64 + has_indexer: bool = True,
65 + index_rope_interleave: bool = True,
66 +) -> None:
67 + _, num_heads, rope_dim = q_pe.shape
68 + _, _, nope_dim = ql_nope.shape
69 + _, num_idx_heads, idx_dim = idx_q.shape
70 +
71 + if has_indexer:
72 + idx_rope_type = TORCH_TO_CUTE_DTYPE[idx_rope_cache.dtype]
73 + idx_weights_type = TORCH_TO_CUTE_DTYPE[idx_weights.dtype]
74 + else:
75 + idx_dim = num_idx_heads = 0
76 + idx_q = idx_rope_cache = idx_q_fp8 = None
77 + idx_weights = idx_weights_out = None
78 + idx_rope_type = idx_weights_type = None
79 +
80 + rope_type = TORCH_TO_CUTE_DTYPE[rope_cache.dtype]
81 + compiled = FusedQKernel.compile(
82 + rope_dim,
83 + nope_dim,
84 + num_heads,
85 + rope_type,
86 + idx_dim,
87 + num_idx_heads,
88 + idx_rope_type,
89 + idx_weights_type,
90 + index_rope_interleave,
91 + )
92 + compiled(
93 + positions,
94 + q_pe,
95 + rope_cache,
96 + ql_nope,
97 + q_scale.view(1),
98 + mqa_output,
99 + idx_q,
100 + idx_rope_cache,
101 + idx_weights,
102 + idx_q_fp8,
103 + idx_weights_out,
104 + float(idx_weights_softmax_scale * idx_weights_head_scale),
105 + )
106 +
107 +
108 +class FusedQKernel:
109 + def __init__(
110 + self,
111 + rope_dim: int,
112 + nope_dim: int,
113 + num_heads: int,
114 + idx_dim: int,
115 + num_idx_heads: int,
116 + index_rope_interleave: bool,
117 + ) -> None:
118 + assert rope_dim == 64
119 + assert nope_dim == 512
120 + assert idx_dim in (128, 0)
121 +
122 + self.rope_dim = rope_dim
123 + self.nope_dim = nope_dim
124 + self.num_heads = num_heads
125 + self.idx_dim = idx_dim
126 + self.num_idx_heads = num_idx_heads
127 + self.index_rope_interleave = index_rope_interleave
128 +
129 + # mqa: rope_dim=64, nope_dim=512, num_heads=64/TP
130 + # indexer: rope_dim=64, nope_dim=64, num_heads=32
131 +
132 + self.num_warps = 4
133 + assert num_heads % self.num_warps == 0
134 + assert num_idx_heads % (4 * self.num_warps) == 0
135 + self.num_ctas_per_tok = num_heads // self.num_warps
136 + self.num_ctas_per_idx_tok = num_idx_heads // (4 * self.num_warps)
137 +
138 + @cute.jit
139 + def __call__(
140 + self,
141 + positions: cute.Tensor,
142 + q_pe: cute.Tensor,
143 + rope_cache: cute.Tensor,
144 + ql_nope: cute.Tensor,
145 + q_scale: cute.Tensor,
146 + mqa_output: cute.Tensor,
147 + idx_q: cute.Tensor,
148 + idx_rope_cache: cute.Tensor,
149 + idx_weights: cute.Tensor,
150 + idx_q_fp8: cute.Tensor,
151 + idx_weights_out: cute.Tensor,
152 + weight_scale: Float32,
153 + stream: CUstream,
154 + ):
155 + num_tokens = positions.shape[0]
156 + if cutlass.const_expr(self.idx_dim == 0):
157 + grid = (num_tokens, self.num_ctas_per_tok, 1)
158 + else:
159 + num_mqa_ctas = num_tokens * self.num_ctas_per_tok
160 + num_idx_ctas = num_tokens * self.num_ctas_per_idx_tok
161 + grid = (num_mqa_ctas + num_idx_ctas, 1, 1)
162 +
163 + self.kernel(
164 + positions,
165 + q_pe,
166 + rope_cache,
167 + ql_nope,
168 + q_scale,
169 + mqa_output,
170 + idx_q,
171 + idx_rope_cache,
172 + idx_weights,
173 + idx_q_fp8,
174 + idx_weights_out,
175 + weight_scale,
176 + ).launch(
177 + grid=grid,
178 + block=(self.num_warps * 32, 1, 1),
179 + stream=stream,
180 + use_pdl=True,
181 + )
182 +
183 + @cute.kernel
184 + def kernel(
185 + self,
186 + positions: cute.Tensor,
187 + q_pe: cute.Tensor,
188 + q_pe_rope_cache: cute.Tensor,
189 + ql_nope: cute.Tensor,
190 + q_scale: cute.Tensor,
191 + mqa_output: cute.Tensor,
192 + idx_q: cute.Tensor,
193 + idx_q_rope_cache: cute.Tensor,
194 + idx_weights: cute.Tensor,
195 + idx_q_fp8: cute.Tensor,
196 + idx_weights_out: cute.Tensor,
197 + weight_scale: Float32,
198 + ):
199 + if cutlass.const_expr(self.idx_dim == 0):
200 + token_id, group_id, _ = cute.arch.block_idx()
201 + self.mqa(
202 + positions,
203 + q_pe,
204 + q_pe_rope_cache,
205 + ql_nope,
206 + q_scale,
207 + mqa_output,
208 + token_id,
209 + group_id,
210 + )
211 + else:
212 + # CTA-specialization
213 + bid, _, _ = cute.arch.block_idx()
214 + num_mqa_ctas = positions.shape[0] * self.num_ctas_per_tok
215 + if bid < num_mqa_ctas:
216 + self.mqa(
217 + positions,
218 + q_pe,
219 + q_pe_rope_cache,
220 + ql_nope,
221 + q_scale,
222 + mqa_output,
223 + bid // self.num_ctas_per_tok,
224 + bid % self.num_ctas_per_tok,
225 + )
226 + else:
227 + bid -= num_mqa_ctas
228 + self.indexer(
229 + positions,
230 + idx_q,
231 + idx_q_rope_cache,
232 + idx_weights,
233 + idx_q_fp8,
234 + idx_weights_out,
235 + weight_scale,
236 + bid // self.num_ctas_per_idx_tok,
237 + bid % self.num_ctas_per_idx_tok,
238 + )
239 +
240 + @cute.jit
241 + def mqa(
242 + self,
243 + positions: cute.Tensor,
244 + q_pe: cute.Tensor,
245 + q_pe_rope_cache: cute.Tensor,
246 + ql_nope: cute.Tensor,
247 + q_scale: cute.Tensor,
248 + mqa_output: cute.Tensor,
249 + token_id,
250 + group_id,
251 + ):
252 + tid, _, _ = cute.arch.thread_idx()
253 + warp_id = cute.arch.make_warp_uniform(tid // 32)
254 + lane_id = tid % 32
255 + head_id = group_id * self.num_warps + warp_id
256 +
257 + cute.arch.griddepcontrol_wait()
258 +
259 + pos = positions[token_id]
260 + inv_scale = 1.0 / q_scale[0]
261 +
262 + cp_op = cute.nvgpu.CopyUniversalOp()
263 + cp_32B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256)
264 + cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
265 + cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32)
266 + cp_2B = cute.make_copy_atom(cp_op, Uint8, num_bits_per_copy=16)
267 +
268 + ##### issue all loads asap #####
269 + rQ_nope_bf16 = cute.make_rmem_tensor(16, BFloat16)
270 + rQ_rope_bf16 = cute.make_rmem_tensor(2, BFloat16)
271 +
272 + src_ql_nope = cute.local_tile(
273 + ql_nope[token_id, head_id, None], (16,), (lane_id,)
274 + )
275 + src_q_rope = cute.local_tile(q_pe[token_id, head_id, None], (2,), (lane_id,))
276 + cute.copy(cp_32B, src_ql_nope, rQ_nope_bf16)
277 + cute.copy(cp_4B, src_q_rope, rQ_rope_bf16)
278 +
279 + rCos_raw = q_pe_rope_cache[pos, 0 + lane_id]
280 + rSin_raw = q_pe_rope_cache[pos, 32 + lane_id]
281 +
282 + ##### process NoPE #####
283 + rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16).load() * inv_scale
284 + rQ_nope_f8 = cute.make_rmem_tensor(16, Float8E4M3FN)
285 + rQ_nope_f8.store(rQ_nope_f32.to(Float8E4M3FN))
286 + dst_Q_nope = cute.local_tile(
287 + mqa_output[token_id, head_id, None], (16,), (lane_id,)
288 + )
289 + cute.copy(cp_16B, rQ_nope_f8, dst_Q_nope)
290 +
291 + ##### process RoPE ######
292 + rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16)
293 + rCos = rCos_raw.to(Float32)
294 + rSin = rSin_raw.to(Float32)
295 + r0 = (rQ_rope_f32[0] * rCos - rQ_rope_f32[1] * rSin) * inv_scale
296 + r1 = (rQ_rope_f32[1] * rCos + rQ_rope_f32[0] * rSin) * inv_scale
297 +
298 + cute.arch.griddepcontrol_launch_dependents()
299 +
300 + # TensorSSA fp32->fp8 cvt has a bug. rely on direct PTX
301 + rQ_rope_f8 = cute.make_rmem_tensor(2, Float8E4M3FN)
302 + cute.recast_tensor(rQ_rope_f8, Uint16)[0] = cvt.fp32x2_to_fp8x2(r0, r1)
303 + dst_Q_rope = cute.local_tile(
304 + mqa_output[token_id, head_id, None], (2,), (256 + lane_id,)
305 + )
306 + cute.copy(cp_2B, rQ_rope_f8, dst_Q_rope)
307 +
308 + @cute.jit
309 + def indexer(
310 + self,
311 + positions: cute.Tensor,
312 + idx_q: cute.Tensor,
313 + idx_q_rope_cache: cute.Tensor,
314 + idx_weights: cute.Tensor,
315 + idx_q_fp8: cute.Tensor,
316 + idx_weights_out: cute.Tensor,
317 + weight_scale: Float32,
318 + token_id,
319 + group_id,
320 + ):
321 + tid, _, _ = cute.arch.thread_idx()
322 + subwarp_id = tid // 8
323 + sublane_id = tid % 8
324 + head_id = group_id * (4 * self.num_warps) + subwarp_id
325 + cute.arch.griddepcontrol_wait()
326 +
327 + pos = positions[token_id]
328 +
329 + cp_op = cute.nvgpu.CopyUniversalOp()
330 + cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
331 + cp_8B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
332 + cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32)
333 +
334 + ##### issue all loads first #####
335 + rQ_rope_bf16 = cute.make_rmem_tensor(8, BFloat16)
336 + if cutlass.const_expr(self.index_rope_interleave):
337 + src_idx_q_rope = cute.local_tile(
338 + idx_q[token_id, head_id, None], (8,), (sublane_id,)
339 + )
340 + cute.copy(cp_16B, src_idx_q_rope, rQ_rope_bf16)
341 + else:
342 + src_idx_q_rope = cute.zipped_divide(
343 + idx_q[token_id, head_id, None], (4,)
344 + ) # (4,32)
345 + cute.copy(
346 + cp_8B,
347 + src_idx_q_rope[None, 0 + sublane_id],
348 + cute.local_tile(rQ_rope_bf16, (4,), (0,)),
349 + )
350 + cute.copy(
351 + cp_8B,
352 + src_idx_q_rope[None, 8 + sublane_id],
353 + cute.local_tile(rQ_rope_bf16, (4,), (1,)),
354 + )
355 +
356 + rQ_nope_bf16 = cute.make_rmem_tensor(8, BFloat16)
357 + src_idx_q_nope = cute.local_tile(
358 + idx_q[token_id, head_id, None], (8,), (8 + sublane_id,)
359 + )
360 + cute.copy(cp_16B, src_idx_q_nope, rQ_nope_bf16)
361 +
362 + rCos_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type)
363 + rSin_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type)
364 + rope_cache_view = cute.zipped_divide(idx_q_rope_cache[pos, None], (4,))
365 +
366 + if cutlass.const_expr(idx_q_rope_cache.element_type == Float32):
367 + cute.copy(cp_16B, rope_cache_view[None, sublane_id], rCos_raw)
368 + cute.copy(cp_16B, rope_cache_view[None, 8 + sublane_id], rSin_raw)
369 + elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16):
370 + cute.copy(cp_8B, rope_cache_view[None, sublane_id], rCos_raw)
371 + cute.copy(cp_8B, rope_cache_view[None, 8 + sublane_id], rSin_raw)
372 +
373 + # unpack to FP32
374 + rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16)
375 + rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16)
376 + if cutlass.const_expr(idx_q_rope_cache.element_type == Float32):
377 + rCos = rCos_raw
378 + rSin = rSin_raw
379 + elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16):
380 + rCos = cvt.bf16x2_to_fp32x2(rCos_raw)
381 + rSin = cvt.bf16x2_to_fp32x2(rSin_raw)
382 +
383 + # apply rope
384 + for i in cutlass.range_constexpr(4):
385 + if cutlass.const_expr(self.index_rope_interleave):
386 + r0 = rQ_rope_f32[i * 2 + 0] * rCos[i] - rQ_rope_f32[i * 2 + 1] * rSin[i]
387 + r1 = rQ_rope_f32[i * 2 + 1] * rCos[i] + rQ_rope_f32[i * 2 + 0] * rSin[i]
388 + rQ_rope_f32[i * 2 + 0] = r0
389 + rQ_rope_f32[i * 2 + 1] = r1
390 + else:
391 + r0 = rQ_rope_f32[0 + i] * rCos[i] - rQ_rope_f32[4 + i] * rSin[i]
392 + r1 = rQ_rope_f32[4 + i] * rCos[i] + rQ_rope_f32[0 + i] * rSin[i]
393 + rQ_rope_f32[0 + i] = r0
394 + rQ_rope_f32[4 + i] = r1
395 +
396 + # amax
397 + amax = Float32(1e-4)
398 + for i in cutlass.range_constexpr(8):
399 + amax = cute.arch.fmax(amax, cute.math.absf(rQ_rope_f32[i]))
400 + amax = cute.arch.fmax(amax, cute.math.absf(rQ_nope_f32[i]))
401 +
402 + # warp reduction among 8 lanes
403 + for i in cutlass.range_constexpr(3):
404 + other = cute.arch.shuffle_sync_bfly(amax, 1 << i)
405 + amax = cute.arch.fmax(amax, other)
406 +
407 + # compute scale from amax
408 + # exp2(ceil(log2(scale))) via bit manipulation
409 + scale = amax * (1.0 / 448.0)
410 + bits = scale.bitcast(Uint32)
411 + exp_bits = (bits + Uint32(0x007FFFFF)) & Uint32(0x7F800000)
412 + scale = exp_bits.bitcast(Float32)
413 + inv_scale = (Uint32(0x7F000000) - exp_bits).bitcast(Float32)
414 +
415 + for i in cutlass.range_constexpr(8):
416 + rQ_nope_f32[i] *= inv_scale
417 + rQ_rope_f32[i] *= inv_scale
418 +
419 + cute.arch.griddepcontrol_launch_dependents()
420 +
421 + # quantize and store
422 + rQ_nope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN)
423 + rQ_nope_f8.store(rQ_nope_f32.load().to(Float8E4M3FN))
424 + dst_idx_q_nope = cute.local_tile(
425 + idx_q_fp8[token_id, head_id, None], (8,), (8 + sublane_id,)
426 + )
427 + cute.copy(cp_8B, rQ_nope_f8, dst_idx_q_nope)
428 +
429 + rQ_rope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN)
430 + rQ_rope_f8.store(rQ_rope_f32.load().to(Float8E4M3FN))
431 + if cutlass.const_expr(self.index_rope_interleave):
432 + dst_idx_q_rope = cute.local_tile(
433 + idx_q_fp8[token_id, head_id, None], (8,), (sublane_id,)
434 + )
435 + cute.copy(cp_8B, rQ_rope_f8, dst_idx_q_rope)
436 + else:
437 + dst_idx_q_rope = cute.zipped_divide(
438 + idx_q_fp8[token_id, head_id, None], (4,)
439 + ) # (4,32)
440 + cute.copy(
441 + cp_4B,
442 + cute.local_tile(rQ_rope_f8, (4,), (0,)),
443 + dst_idx_q_rope[None, 0 + sublane_id],
444 + )
445 + cute.copy(
446 + cp_4B,
447 + cute.local_tile(rQ_rope_f8, (4,), (1,)),
448 + dst_idx_q_rope[None, 8 + sublane_id],
449 + )
450 +
451 + # scale indexer weights
452 + if sublane_id == 0:
453 + w = idx_weights[token_id, head_id].to(Float32)
454 + idx_weights_out[token_id, head_id] = w * scale * weight_scale
455 +
456 + @cache
457 + @staticmethod
458 + def compile(
459 + rope_dim: int,
460 + nope_dim: int,
461 + num_heads: int,
462 + rope_type: type[cutlass.Numeric],
463 + idx_dim: int,
464 + num_idx_heads: int,
465 + idx_rope_type: type[cutlass.Numeric] | None,
466 + idx_weights_type: type[cutlass.Numeric] | None,
467 + index_rope_interleave: bool,
468 + ):
469 + num_tokens = cute.sym_int()
470 + max_pos = cute.sym_int()
471 +
472 + positions = _make_fake_tensor(Int64, (num_tokens,), divisibility=1)
473 + q_pe = _make_fake_tensor(
474 + BFloat16, (num_tokens, num_heads, rope_dim), divisibility=16
475 + )
476 + rope_cache = _make_fake_tensor(rope_type, (max_pos, rope_dim), divisibility=8)
477 + ql_nope = _make_fake_tensor(
478 + BFloat16, (num_tokens, num_heads, nope_dim), divisibility=16
479 + )
480 + q_scale = _make_fake_tensor(Float32, (1,), divisibility=4)
481 + mqa_output = _make_fake_tensor(
482 + Float8E4M3FN,
483 + (num_tokens, num_heads, nope_dim + rope_dim),
484 + divisibility=16,
485 + )
486 +
487 + if idx_rope_type is not None:
488 + index_q = _make_fake_tensor(
489 + BFloat16, (num_tokens, num_idx_heads, idx_dim), divisibility=16
490 + )
491 + index_rope_cache = _make_fake_tensor(
492 + idx_rope_type, (max_pos, rope_dim), divisibility=8
493 + )
494 + index_weights = _make_fake_tensor(
495 + idx_weights_type, (num_tokens, num_idx_heads), divisibility=8
496 + )
497 + index_q_fp8 = _make_fake_tensor(
498 + Float8E4M3FN, (num_tokens, num_idx_heads, idx_dim), divisibility=16
499 + )
500 + index_weights_out = _make_fake_tensor(
501 + Float32, (num_tokens, num_idx_heads), divisibility=4
502 + )
503 + else:
504 + index_q = index_rope_cache = index_q_fp8 = None
505 + index_weights = index_weights_out = None
506 +
507 + kernel = FusedQKernel(
508 + rope_dim,
509 + nope_dim,
510 + num_heads,
511 + idx_dim,
512 + num_idx_heads,
513 + index_rope_interleave,
514 + )
515 + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
516 + return cute.compile(
517 + kernel,
518 + positions,
519 + q_pe,
520 + rope_cache,
521 + ql_nope,
522 + q_scale,
523 + mqa_output,
524 + index_q,
525 + index_rope_cache,
526 + index_weights,
527 + index_q_fp8,
528 + index_weights_out,
529 + Float32(0.0),
530 + stream,
531 + options="--enable-tvm-ffi",
532 + )
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+31
@@ -28,6 +28,7 @@ from vllm.v1.attention.backend import (
28 MultipleOf,
29 )
30 from vllm.v1.attention.backends.mla.sparse_utils import (
31 + phys_shadow,
32 triton_convert_req_index_to_global_index,
33 triton_filter_and_convert_dcp_index,
34 )
@@ -404,6 +405,36 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
405 NUM_TOPK_TOKENS=topk_indices.shape[1],
406 return_valid_counts=True,
407 )
408 + elif (shadow := phys_shadow(self.topk_indices_buffer)) is not None:
409 + # skip_topk layers reuse the previous fresh layer's PHYSICAL
410 + # indices: the logical top-k indices are shared (one buffer for
411 + # the whole forward) and block_table/req_id are per-step constant,
412 + # so re-converting them per layer produces identical output. This
413 + # covers both the backbone (index_topk_freq > 1) and MTP draft
414 + # steps 1+ (index_share_for_mtp_iteration sets skip_topk; the
415 + # draft compact re-arranges the shadow together with the logical
416 + # buffer). The fresh/skip split is fixed per captured graph, so
417 + # the branch is cudagraph-safe.
418 + phys_buf, seq_buf = shadow
419 + wrote_fresh = getattr(layer, "indexer", None) is not None and not getattr(
420 + layer, "skip_topk", False
421 + )
422 + if wrote_fresh:
423 + topk_indices_physical, seq_lens = (
424 + triton_convert_req_index_to_global_index(
425 + attn_metadata.req_id_per_token[:num_actual_toks],
426 + attn_metadata.block_table,
427 + topk_indices,
428 + BLOCK_SIZE=attn_metadata.block_size,
429 + NUM_TOPK_TOKENS=topk_indices.shape[1],
430 + return_valid_counts=True,
431 + out=phys_buf[:num_actual_toks],
432 + valid_counts_out=seq_buf[:num_actual_toks],
433 + )
434 + )
435 + else:
436 + topk_indices_physical = phys_buf[:num_actual_toks]
437 + seq_lens = seq_buf[:num_actual_toks]
438 else:
439 topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(
440 attn_metadata.req_id_per_token[:num_actual_toks],
vllm/v1/attention/backends/mla/sparse_utils.py
+56 -4
@@ -128,6 +128,8 @@ def triton_convert_req_index_to_global_index(
128 prefill_workspace_request_ids: torch.Tensor | None = None,
129 prefill_workspace_starts: torch.Tensor | None = None,
130 return_valid_counts: bool = False,
131 + out: torch.Tensor | None = None,
132 + valid_counts_out: torch.Tensor | None = None,
133 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
134 """
135 out[token_id, indice_id] =
@@ -178,14 +180,20 @@ def triton_convert_req_index_to_global_index(
180 req_id_c = req_id.contiguous()
181 block_table_c = block_table.contiguous()
182 token_indices_c = token_indices.contiguous()
181 - out = torch.empty_like(token_indices_c)
183 + # `out`/`valid_counts_out` allow writing into a stable pre-allocated buffer
184 + # (used by the shared-physical-index cache to stay cudagraph-safe).
185 + out = torch.empty_like(token_indices_c) if out is None else out
186
187 # Allocate valid count buffer if needed (must be zero-initialized for atomics)
188 valid_counts: torch.Tensor | None = None
189 if return_valid_counts:
186 - valid_counts = torch.zeros(
187 - num_tokens, dtype=torch.int32, device=token_indices.device
188 - )
190 + if valid_counts_out is None:
191 + valid_counts = torch.zeros(
192 + num_tokens, dtype=torch.int32, device=token_indices.device
193 + )
194 + else:
195 + valid_counts = valid_counts_out
196 + valid_counts.zero_()
197
198 # Strides in elements
199 bt_stride0, bt_stride1 = block_table_c.stride()
@@ -341,3 +349,47 @@ def triton_filter_and_convert_dcp_index(
349 assert valid_counts is not None
350 return out, valid_counts
351 return out
352 +
353 +
354 +# --- Physical-index shadow (DSA index_topk_freq > 1) -------------------------
355 +# GLM-5.2 uses index_topk_freq=4: only ~1/4 of layers write a fresh top-k into
356 +# the single shared topk_indices_buffer; the rest reuse it. Since block_table
357 +# and req_id_per_token are constant within a decode step, the physical indices
358 +# (block_table lookup of the logical top-k) are identical across a freq-group.
359 +# Fresh layers convert once into a STABLE shadow of the logical buffer and
360 +# skip layers read it -- eliminating ~3/4 of the per-layer convert kernels.
361 +#
362 +# Invariant: shadow[j] == convert(logical[j]) row-wise. Whoever re-arranges
363 +# the logical buffer (e.g. the MTP draft compact between the multi-token
364 +# step-0 layout and the single-token steps-1+ layout) must apply the same
365 +# gather to the shadow.
366 +#
367 +# Shadows are registered ONLY at buffer creation time (model init, never
368 +# inside cudagraph capture) via register_phys_shadow; every other access is
369 +# the read-only phys_shadow lookup, so addresses are stable across graph
370 +# replays and no allocation can happen during capture.
371 +import weakref as _weakref # noqa: E402
372 +
373 +_PHYS_SHADOWS: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
374 +
375 +
376 +def register_phys_shadow(topk_buf: torch.Tensor) -> None:
377 + """Allocate the (physical-index, valid-count) shadow for a logical top-k
378 + buffer. Call once where the buffer is created. The finalizer drops the
379 + entry with the owning buffer, so a recycled id can never alias a stale
380 + shadow."""
381 + key = id(topk_buf)
382 + if key not in _PHYS_SHADOWS:
383 + _PHYS_SHADOWS[key] = (
384 + torch.empty_like(topk_buf),
385 + torch.empty(topk_buf.shape[0], dtype=torch.int32, device=topk_buf.device),
386 + )
387 + _weakref.finalize(topk_buf, _PHYS_SHADOWS.pop, key, None)
388 +
389 +
390 +def phys_shadow(
391 + topk_buf: torch.Tensor,
392 +) -> tuple[torch.Tensor, torch.Tensor] | None:
393 + """The registered shadow for this buffer, or None if it was never
394 + registered. Never allocates; capture-safe."""
395 + return _PHYS_SHADOWS.get(id(topk_buf))
vllm/v1/spec_decode/llm_base_proposer.py
+3 -2
@@ -1009,8 +1009,9 @@ class SpecDecodeBaseProposer:
1009 # DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final-
1010 # norm hidden, so its forward returns (logit_hidden,
1011 # recycle_hidden). Other MTP families return a single tensor.
1012 - return "DeepSeekMTPModel" in (
1013 - self.draft_model_config.hf_config.architectures or []
1012 + return bool(
1013 + {"DeepSeekMTPModel", "DeepseekV32MTPModel"}
1014 + & set(self.draft_model_config.hf_config.architectures or [])
1015 )
1016 return self.method not in ("mtp", "draft_model", "dflash")
1017
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+35 -5
@@ -48,6 +48,18 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
48 self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None
49 self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None
50
51 + # Lifecycle hooks for model-specific optimizations. Subclasses override
52 + # the ones they need. These fire in both `capture` and `propose` so that
53 + # any state they toggle (e.g. attention flags baked into a CUDA graph) is
54 + # identical at capture time and replay time.
55 + def on_prefill_begin(self, num_reqs: int, num_tokens: int) -> None: ...
56 +
57 + def on_prefill_end(self, num_reqs: int, num_tokens: int) -> None: ...
58 +
59 + def on_multi_step_decode_begin(self, num_reqs: int) -> None: ...
60 +
61 + def on_multi_step_decode_end(self, num_reqs: int) -> None: ...
62 +
63 @property
64 def advance_draft_positions(self) -> bool:
65 """
@@ -98,6 +110,9 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
110 assert self.prefill_cudagraph_manager is not None
111 if self.prefill_cudagraph_manager.use_breakable_cg:
112 self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
113 +
114 + self.on_prefill_begin(self.max_num_reqs, self.max_num_tokens)
115 +
116 self.prefill_cudagraph_manager.capture(
117 self._prefill,
118 self.model_state,
@@ -108,9 +123,13 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
123 progress_bar_desc="Capturing prefill CUDA graphs",
124 )
125
126 + self.on_prefill_end(self.max_num_reqs, self.max_num_tokens)
127 +
128 if self.num_speculative_steps == 1:
129 return
130
131 + self.on_multi_step_decode_begin(self.max_num_reqs)
132 +
133 # Capture the decode draft generation routine (model forward +
134 # sample + update_draft_inputs) for a single
135 # step.
@@ -125,6 +144,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
144 progress_bar_desc="Capturing decode CUDA graphs",
145 )
146
147 + self.on_multi_step_decode_end(self.max_num_reqs)
148 +
149 @torch.inference_mode()
150 def propose(
151 self,
@@ -153,7 +174,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
174 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
175 is_profile: bool = False,
176 ) -> torch.Tensor:
156 - num_tokens = input_batch.num_tokens_after_padding
177 + num_tokens = input_batch.num_tokens
178 + num_tokens_padded = input_batch.num_tokens_after_padding
179 num_reqs = input_batch.num_reqs
180 max_query_len = input_batch.num_scheduled_tokens.max()
181 max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
@@ -174,7 +196,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
196 )
197 else:
198 hidden_states = last_hidden_states
177 - self.hidden_states[:num_tokens].copy_(hidden_states)
199 + self.hidden_states[:num_tokens_padded].copy_(hidden_states)
200
201 self._copy_request_inputs(
202 num_reqs,
@@ -202,20 +224,22 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
224 num_reqs,
225 # Use the actual number of tokens without padding added by
226 # the target model during FULL cudagraph.
205 - input_batch.num_tokens,
227 + num_tokens,
228 max_query_len,
229 )
230 prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
231 self.prefill_cudagraph_manager,
232 num_reqs,
211 - num_tokens,
233 + num_tokens_padded,
234 uniform_token_count,
235 dp_size=self.dp_size,
236 dp_rank=self.dp_rank,
237 need_eager=is_profile,
238 )
239
218 - self._prepare_eplb_forward(input_batch.num_tokens)
240 + self._prepare_eplb_forward(num_tokens)
241 +
242 + self.on_prefill_begin(num_reqs, num_tokens)
243
244 if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
245 # Replay the full graph for draft prefill.
@@ -235,6 +259,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
259 mm_inputs=mm_inputs,
260 )
261
262 + self.on_prefill_end(num_reqs, num_tokens)
263 +
264 if self.num_speculative_steps == 1:
265 # Early exit.
266 return self.draft_tokens[:num_reqs, :1]
@@ -262,6 +288,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
288 need_eager=is_profile,
289 )
290
291 + self.on_multi_step_decode_begin(num_reqs)
292 +
293 # Generate the remaining num_speculative_steps - 1 draft tokens.
294 self._multi_step_decode(
295 num_reqs,
@@ -271,6 +299,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
299 input_batch.seq_lens_cpu_upper_bound,
300 )
301
302 + self.on_multi_step_decode_end(num_reqs)
303 +
304 return self.draft_tokens[:num_reqs]
305
306 @torch.inference_mode()
vllm/v1/worker/gpu/spec_decode/mtp/speculator.py
+40 -1
@@ -1,8 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import torch
5 import torch.nn as nn
6
7 +from vllm.config import VllmConfig
8 from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
9 AutoRegressiveSpeculator,
10 )
@@ -10,9 +12,46 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
12
13
14 class MTPSpeculator(AutoRegressiveSpeculator):
15 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
16 + super().__init__(vllm_config, device)
17 +
18 + spec_config = vllm_config.speculative_config
19 + draft_hf_config = (
20 + spec_config.draft_model_config.hf_config
21 + if spec_config is not None
22 + else None
23 + )
24 + # Detect index_share_for_mtp_iteration. When True, the proposer
25 + # toggles skip_topk so step 0 computes MTP's own indices and
26 + # steps 1+ reuse them.
27 + self.share_mtp_topk_indices = getattr(
28 + draft_hf_config, "index_share_for_mtp_iteration", False
29 + )
30 +
31 def load_draft_model(
32 self,
33 target_model: nn.Module,
34 target_attn_layer_names: set[str],
35 ) -> nn.Module:
18 - return load_eagle_model(target_model, self.vllm_config)
36 + draft_model = load_eagle_model(target_model, self.vllm_config)
37 + self.share_mtp_topk_indices = self.share_mtp_topk_indices and hasattr(
38 + draft_model.model, "set_skip_topk"
39 + )
40 + return draft_model
41 +
42 + def on_prefill_end(self, num_reqs: int, num_tokens: int) -> None:
43 + # Step 0 (prefill) wrote topk indices for every query token in the
44 + # multi-token batch. Compact them down to each request's last token so
45 + # steps 1+ can reuse them from the shared buffer.
46 + if self.share_mtp_topk_indices and self.num_speculative_steps > 1:
47 + self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs])
48 +
49 + def on_multi_step_decode_begin(self, num_reqs: int) -> None:
50 + # Switch to reuse mode so draft steps 1+ skip the indexer op and read
51 + # the indices that step 0 wrote into the shared buffer.
52 + if self.share_mtp_topk_indices:
53 + self.model.model.set_skip_topk(True)
54 +
55 + def on_multi_step_decode_end(self, num_reqs: int) -> None:
56 + if self.share_mtp_topk_indices:
57 + self.model.model.set_skip_topk(False)