@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 213f681f81

Revert "[Perf][GLM-5.2] Blackwell decode optimizations" (#49768)

Woosuk Kwon committed Jul 24, 2026 at 16:04 UTC 213f681f8117b9026ca8189583d2855d70104401
29 files changed +138 -2139
CMakeLists.txt
-17
@@ -732,23 +732,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
732 "(requires SM90+ and CUDA >= 12.0).")
733 endif()
734
735 - # BF16 skinny GEMM (M<=32; weight-bandwidth-bound decode shapes).
736 - # Requires SM90+.
737 - cuda_archs_sm90plus(BF16_SKINNY_GEMM_ARCHS "${CUDA_ARCHS}")
738 - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND BF16_SKINNY_GEMM_ARCHS)
739 - set(BF16_SKINNY_GEMM_SRCS
740 - "csrc/libtorch_stable/bf16_skinny_gemm_entry.cu"
741 - "csrc/libtorch_stable/bf16_skinny_gemm.cu")
742 - set_gencode_flags_for_srcs(
743 - SRCS "${BF16_SKINNY_GEMM_SRCS}"
744 - CUDA_ARCHS "${BF16_SKINNY_GEMM_ARCHS}")
745 - list(APPEND VLLM_STABLE_EXT_SRC "${BF16_SKINNY_GEMM_SRCS}")
746 - message(STATUS "Building bf16_skinny_gemm for archs: ${BF16_SKINNY_GEMM_ARCHS}")
747 - else()
748 - message(STATUS "Not building bf16_skinny_gemm as no compatible archs found "
749 - "(requires SM90+ and CUDA >= 12.0).")
750 - endif()
751 -
735 # Only build AllSpark kernels if we are building for at least some compatible archs.
736 cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
737 if (ALLSPARK_ARCHS)
csrc/libtorch_stable/bf16_skinny_gemm.cu deleted
-262
@@ -1,262 +0,0 @@
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 deleted
-170
@@ -1,170 +0,0 @@
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
+36 -112
@@ -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 || m_iter_cnt == 2);
394 + static_assert(m_iter_cnt == 1);
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,18 +416,13 @@ 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.
419 #pragma unroll
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 - }
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;
426 }
427 #pragma unroll
428 for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
@@ -451,14 +446,11 @@ struct MmaComputer {
446 wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
447
448 #pragma unroll
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 - }
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 }
455
456 #pragma unroll
@@ -477,12 +469,9 @@ struct MmaComputer {
469 for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
470 #pragma unroll
471 for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; 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 - }
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]);
475 }
476 }
477 ::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
@@ -497,14 +486,14 @@ struct MmaComputer {
486 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
487 asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
488 // reorganize the acc_reg
500 - constexpr int thread_m = 2 * m_iter_cnt;
489 + constexpr int thread_m = 2;
490 constexpr int thread_n = 2 * n_iter_cnt;
491 constexpr int cta_mma_n = n_iter_cnt * 8;
492 float acc_reg_reorg[thread_m][thread_n];
493
494 for (int i = 0; i < thread_m; i++) {
495 for (int j = 0; j < thread_n; j++) {
507 - acc_reg_reorg[i][j] = acc_reg[i / 2][j / 2][(j % 2) + (i % 2) * 2];
496 + acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
497 }
498 }
499
@@ -525,8 +514,7 @@ struct MmaComputer {
514 for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
515 #pragma unroll
516 for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
528 - int m_idx =
529 - (lane_idx / 4) + (m_idx_thread % 2) * 8 + (m_idx_thread / 2) * 16;
517 + int m_idx = (lane_idx / 4) + m_idx_thread * 8;
518 int n_idx =
519 ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
520 smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
@@ -599,7 +587,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
587 static_assert(
588 tile_k == 128 || tile_k == 256 || tile_k == 512 ||
589 tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK.
602 - static_assert(tile_m == 16 || tile_m == 32);
590 + static_assert(tile_m == 16);
591 constexpr int g2s_vec_bytes = 16;
592 constexpr int a_elem_bytes = 2;
593 constexpr int b_elem_bytes = 2;
@@ -659,7 +647,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
647 #endif
648 }
649
662 -template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileM = 16>
650 +template <typename T, int kHdIn, int kHdOut, int kTileN>
651 void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
652 cudaStream_t const stream) {
653 constexpr int gemm_m = kHdOut; // 2112
@@ -667,7 +655,7 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
655 constexpr int gemm_k = kHdIn; // 7168
656 constexpr int batch_size = 1;
657 std::swap(mat_a, mat_b);
670 - constexpr int tile_m = kTileM;
658 + constexpr int tile_m = 16;
659 constexpr int tile_n = kTileN; // 8 or 16
660 constexpr int tile_k = std::max(256, 1024 / tile_n); // 256
661 constexpr int max_stage_cnt =
@@ -714,44 +702,6 @@ template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
702 __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
703 cudaStream_t);
704
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 -
705 void dsv3_fused_a_gemm(torch::stable::Tensor& output,
706 torch::stable::Tensor const& mat_a,
707 torch::stable::Tensor const& mat_b) {
@@ -760,15 +710,12 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
710 int const hd_in = mat_a.size(1);
711 int const hd_out = mat_b.size(1);
712
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;
713 + constexpr int kHdIn = 7168;
714 + constexpr int kHdOut = 2112;
715 STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
716 "required 1 <= mat_a.shape[0] <= 16");
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)");
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");
719 STD_TORCH_CHECK(output.size(0) == num_tokens,
720 "required output.shape[0] == mat_a.shape[0]");
721 STD_TORCH_CHECK(output.size(1) == hd_out,
@@ -791,41 +738,18 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
738 STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
739
740 auto stream = get_current_cuda_stream(mat_a.get_device_index());
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 - }
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);
747 } else {
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 - }
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);
753 }
754 }
755
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
+10 -42
@@ -49,11 +49,6 @@ __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 -
52 // Precompute SF layout parameter (constant for entire kernel).
53 int32_t const numKTiles = (outputCols + 63) / 64;
54
@@ -128,11 +123,6 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
123 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
124 "Vec size is not matched.");
125
131 -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
132 - cudaGridDependencySynchronize();
133 - cudaTriggerProgrammaticLaunchCompletion();
134 -#endif
135 -
126 int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
127 int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
128
@@ -211,8 +201,6 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
201 const torch::stable::accelerator::DeviceGuard device_guard(
202 input.get_device_index());
203 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;
204
205 int output_sf_n_unpadded = int(output_n / CVT_FP4_SF_VEC_SIZE);
206
@@ -236,21 +224,10 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
224 input.scalar_type(), "nvfp4_quant_kernel", [&] {
225 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
226 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
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));
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));
231 });
232 } else {
233 int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD;
@@ -263,21 +240,12 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
240 input.scalar_type(), "nvfp4_quant_kernel", [&] {
241 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
242 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
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));
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));
249 });
250 }
251 }
csrc/libtorch_stable/torch_bindings.cpp
-4
@@ -330,10 +330,6 @@ 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 -
333 // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
334 ops.def(
335 "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
recipes/glm5.2-ll-b300-tp8-mtp5.md deleted
-41
@@ -1,41 +0,0 @@
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 deleted
-70
@@ -1,70 +0,0 @@
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,32 +49,6 @@ 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 -
52 # ── reference helpers ────────────────────────────────────────────────────────
53
54
@@ -318,57 +292,6 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
292 assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
293
294
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 -
295 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
296 def test_fused_norm_rope_ds_mla(num_tokens: int):
297 """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).
tests/models/registry.py
-4
@@ -1600,10 +1600,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1600 speculative_model="luccafong/deepseek_mtp_draft_random",
1601 trust_remote_code=True,
1602 ),
1603 - "DeepseekV32MTPModel": _HfExamplesInfo(
1604 - "deepseek-ai/DeepSeek-V3.2-Exp",
1605 - speculative_model="deepseek-ai/DeepSeek-V3.2-Exp",
1606 - ),
1603 "DeepSeekV4MTPModel": _HfExamplesInfo(
1604 "deepseek-ai/DeepSeek-V4-Flash",
1605 speculative_model="deepseek-ai/DeepSeek-V4-Flash",
vllm/_custom_ops.py
-25
@@ -2371,20 +2371,6 @@ 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 -
2374 if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):
2375
2376 @register_fake("_C::fp32_router_gemm")
@@ -2396,17 +2382,6 @@ if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):
2382 return
2383
2384
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 -
2385 def topk_softmax(
2386 topk_weights: torch.Tensor,
2387 topk_ids: torch.Tensor,
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+35 -111
@@ -202,7 +202,11 @@ 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 - oversize = current_tensor_size > max_tensor_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 + )
210 curr_device = current_platform.get_device_capability()
211 device_capability = curr_device.to_int() if curr_device is not None else None
212
@@ -215,47 +219,6 @@ if flashinfer_comm is not None:
219 get_workspace_fn = (
220 get_fi_ar_quant_workspace if is_quant_pattern else get_fi_ar_workspace
221 )
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
222 workspace = get_workspace_fn(
223 world_size=world_size,
224 rank=get_tensor_model_parallel_rank(),
@@ -289,75 +252,36 @@ if flashinfer_comm is not None:
252 if workspace.backend in ("trtllm", "mnnvl"):
253 layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4
254
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()
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 + )
285
286 def call_trtllm_fused_allreduce_norm_fake(
287 allreduce_in: torch.Tensor,
vllm/config/speculative.py
+1 -7
@@ -330,7 +330,6 @@ 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"
333 if hf_config.model_type in (
334 "deepseek_v3",
335 "deepseek_v32",
@@ -340,12 +339,7 @@ class SpeculativeConfig:
339 if hf_config.model_type == "deepseek_mtp":
340 n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
341 hf_config.update(
343 - {
344 - "n_predict": n_predict,
345 - "architectures": [
346 - "DeepseekV32MTPModel" if use_sparse_mtp else "DeepSeekMTPModel"
347 - ],
348 - }
342 + {"n_predict": n_predict, "architectures": ["DeepSeekMTPModel"]}
343 )
344 if hf_config.model_type == "deepseek_v4":
345 hf_config.model_type = "deepseek_mtp"
vllm/cute_utils/__init__.py
-5
@@ -33,11 +33,6 @@ 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 -
36
37 @dsl_user_op
38 def recast_val(x, dtype, *, loc=None, ip=None):
vllm/cute_utils/cvt.py
+1 -22
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from cutlass import BFloat16, Constexpr, Float32, Uint16, Uint32, cute
3 +from cutlass import Constexpr, Float32, 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,11 +38,6 @@ 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 -
41 # NOTE: the output is always 1D
42 size = cute.size(data.shape)
43 out = cute.make_rmem_tensor(size * 2, Float32)
@@ -89,22 +84,6 @@ def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
84 return cute.TensorSSA(vec, 2, Uint32)
85
86
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 -
87 @dsl_user_op
88 def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
89 out = llvm.inline_asm(
vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
+4 -15
@@ -43,7 +43,7 @@ try:
43 if flashinfer_comm is not None
44 else None
45 )
46 -except (ImportError, AttributeError):
46 +except ImportError:
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,23 +52,12 @@ except (ImportError, AttributeError):
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 -
55 def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None:
56 """Workspace token budget for flashinfer fused all-reduce, or None if the
57 current world size / device is unsupported. Mirrors ``FlashInferAllReduce``."""
71 - max_size_mb = _fi_ar_max_size_mb().get(tp_size)
58 + from vllm.config.compilation import PassConfig
59 +
60 + max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
61 if not max_size_mb:
62 return None
63 element_size = torch.tensor([], dtype=dtype).element_size()
vllm/model_executor/models/registry.py
+1 -2
@@ -114,7 +114,7 @@ _TEXT_GENERATION_MODELS = {
114 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
115 "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
116 "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
117 - "GlmMoeDsaForCausalLM": ("vllm.models.deepseek_v32", "GlmMoeDsaForCausalLM"),
117 + "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
118 "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
119 "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
120 "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
@@ -628,7 +628,6 @@ _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"),
631 "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
632 "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
633 "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"),
vllm/models/deepseek_v32/__init__.py
+4 -18
@@ -6,31 +6,17 @@ 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).
9 """
10
11 from vllm.platforms import current_platform
12
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
13 +if current_platform.is_rocm() or current_platform.is_xpu():
14 + raise NotImplementedError("deepseek_v32 currently supports NVIDIA SM100 only.")
15
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
16 +from .nvidia.model import DeepseekV32ForCausalLM
17 +from .nvidia.mtp import DeepseekV32MTP
18
19 __all__ = [
20 "DeepseekV32ForCausalLM",
21 "DeepseekV32MTP",
35 - "GlmMoeDsaForCausalLM",
22 ]
vllm/models/deepseek_v32/nvidia/attention.py
+12 -72
@@ -4,7 +4,6 @@ import torch
4 import torch.nn as nn
5 from transformers import DeepseekV2Config, DeepseekV3Config
6
7 -import vllm._custom_ops as ops
7 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
8 from vllm.config import CacheConfig, VllmConfig
9 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -22,15 +21,17 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
21 per_token_group_quant_fp8,
22 )
23 from vllm.model_executor.layers.rotary_embedding import get_rope
25 -from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
24 +from vllm.model_executor.layers.sparse_attn_indexer import (
25 + SparseAttnIndexer,
26 + sparse_attn_indexer,
27 +)
28 from vllm.model_executor.models.deepseek_v2 import (
29 DeepSeekV2FusedQkvAProjLinear,
30 DeepseekV32IndexerCache,
31 yarn_get_mscale,
32 )
33 from vllm.model_executor.models.utils import extract_layer_index
32 -from vllm.platforms import current_platform
33 -from vllm.utils.torch_utils import _encode_layer_name, is_quantized_kv_cache
34 +from vllm.utils.torch_utils import is_quantized_kv_cache
35
36 from .kernels import fused_norm_rope, fused_q
37
@@ -170,11 +171,6 @@ class DeepseekV32Attention(MLAAttention):
171 ) -> None:
172 quant_config = vllm_config.quant_config
173 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"
174
175 hidden_size = config.hidden_size
176 qk_nope_head_dim = config.qk_nope_head_dim
@@ -335,37 +331,6 @@ class DeepseekV32Attention(MLAAttention):
331 rope_parameters=config.rope_parameters,
332 is_neox_style=False,
333 )
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 -
334 # Lightning indexer uses its own RoPE; interleave maps to non-NeoX.
335 self.indexer_rope_emb = get_rope(
336 qk_rope_head_dim,
@@ -374,32 +339,13 @@ class DeepseekV32Attention(MLAAttention):
339 is_neox_style=not getattr(config, "indexer_rope_interleave", False),
340 )
341
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 -
342 def forward( # type: ignore[override]
343 self,
344 positions: torch.Tensor,
345 hidden_states: torch.Tensor,
346 ) -> torch.Tensor:
347 # Captured: A-projections (+ indexer A-GEMM on indexer layers).
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]
348 + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
349 q_c, kv_c, k_pe = qkv_lora.split(
350 [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
351 )
@@ -452,9 +398,7 @@ class DeepseekV32Attention(MLAAttention):
398 assert isinstance(slot_mapping, dict)
399 mla_slot = slot_mapping.get(self.layer_name)
400
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:
401 + if self.indexer is not None:
402 has_indexer = True
403 indexer_k_norm_w = self.indexer.k_norm.weight
404 indexer_k_norm_bias = self.indexer.k_norm.bias
@@ -507,16 +451,12 @@ class DeepseekV32Attention(MLAAttention):
451 index_rope_interleave=self._index_rope_interleave,
452 )
453
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)
454 + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
455 q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
456 q_nope = q_nope.transpose(0, 1)
457 ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1)
458
519 - if self.indexer is not None and not self.skip_topk:
459 + if self.indexer is not None:
460 index_q = self.indexer.wq_b(q_c)[0]
461 index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim)
462 else:
@@ -538,10 +478,10 @@ class DeepseekV32Attention(MLAAttention):
478 quantize_mqa=self._fp8_query,
479 )
480
541 - if self.indexer is not None and not self.skip_topk:
542 - torch.ops.vllm.sparse_attn_indexer(
481 + if self.indexer is not None:
482 + sparse_attn_indexer(
483 q_c,
544 - _encode_layer_name(self.indexer.k_cache.prefix),
484 + self.indexer.k_cache.prefix,
485 self.indexer.k_cache.kv_cache,
486 index_q_fp8,
487 None, # q_scale folded into weights on the fp8 path
vllm/models/deepseek_v32/nvidia/kernels.py
+16 -284
@@ -2,80 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 -from vllm.platforms import current_platform
5 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
6
7 # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the
8 # has_indexer=False path so the indexer args don't allocate every call.
9 _DUMMY_CACHE: dict[tuple, torch.Tensor] = {}
10
11
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 -
12 def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
13 key = (shape, dtype, device)
14 t = _DUMMY_CACHE.get(key)
@@ -214,14 +147,9 @@ def _fused_norm_rope_kernel(
147 TOPK_BLOCK_SIZE: tl.constexpr,
148 HAS_INDEXER: tl.constexpr,
149 INDEX_ROPE_INTERLEAVE: tl.constexpr,
217 - USE_PDL: tl.constexpr,
150 ):
151 pid = tl.program_id(0)
152 tok_idx = tl.program_id(1)
221 - if USE_PDL:
222 - tl.extra.cuda.gdc_wait()
223 - tl.extra.cuda.gdc_launch_dependents()
224 -
153 if pid == 3:
154 if not HAS_INDEXER:
155 # Shared layer: reuse the previous indexer layer's top-k; do not
@@ -238,17 +166,6 @@ def _fused_norm_rope_kernel(
166 )
167 return
168
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 -
169 if slot_mapping_ptr is None:
170 # Memory profiling run.
171 return
@@ -257,7 +174,15 @@ def _fused_norm_rope_kernel(
174 # Padding
175 return
176
260 - if pid == 1:
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:
186 # KV RMS Norm + KV RoPE + MLA concat_and_cache.
187 # Merged so the normed kv_c and RoPE'd k_pe can be written
188 # to the MLA KV cache directly without a separate kernel.
@@ -286,6 +211,9 @@ def _fused_norm_rope_kernel(
211 r2 = x2 * cos + x1 * sin
212
213 # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.
214 + if mla_cache_entry_stride == 0:
215 + return
216 +
217 mla_block_size = mla_cache_block_stride // mla_cache_entry_stride
218 mla_block_idx = slot_idx // mla_block_size
219 mla_block_off = slot_idx % mla_block_size
@@ -435,7 +363,7 @@ def _fused_norm_rope_kernel(
363 )
364
365
438 -def _fused_norm_rope_impl(
366 +def fused_norm_rope(
367 positions: torch.Tensor,
368 q_c: torch.Tensor,
369 q_rms_norm_w: torch.Tensor,
@@ -537,17 +465,12 @@ def _fused_norm_rope_impl(
465 # Dummy values — pid 2 will skip the MLA cache write because
466 # slot_mapping is all -1.
467 mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)
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
468 + mla_block_stride = 0
469 + mla_entry_stride = 0
470 mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
471
472 if q_c_out is None:
473 q_c_out = torch.empty_like(q_c)
550 - use_pdl = _is_arch_support_pdl()
474 _fused_norm_rope_kernel[(4, num_tokens)](
475 positions,
476 # Q RMS norm
@@ -606,157 +529,6 @@ def _fused_norm_rope_impl(
529 TOPK_BLOCK_SIZE=1024,
530 HAS_INDEXER=has_indexer,
531 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,
532 )
533 return q_c_out
534
@@ -810,14 +582,10 @@ def _fused_q_kernel(
582 HAS_INDEXER: tl.constexpr,
583 INDEX_ROPE_INTERLEAVE: tl.constexpr,
584 QUANTIZE_MQA: tl.constexpr,
813 - USE_PDL: tl.constexpr,
585 ):
586 pid = tl.program_id(0)
587 tok_idx = tl.program_id(1)
588 head_idx = tl.program_id(2)
818 - if USE_PDL:
819 - tl.extra.cuda.gdc_wait()
820 - tl.extra.cuda.gdc_launch_dependents()
589
590 if pid == 2:
591 # ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16
@@ -1005,12 +773,10 @@ def fused_q(
773 ``(ql_nope, q_pe)`` tuple the backend expects.
774 """
775 assert positions.ndim == 1
1008 - assert positions.dtype == torch.int64
776 assert q_pe.ndim == 3
777 assert q_pe_cos_sin_cache.ndim == 2
778 assert ql_nope.ndim == 3
779 assert ql_nope.shape[:2] == q_pe.shape[:2]
1013 - assert q_scale.dtype == torch.float32 and q_scale.numel() == 1
780
781 num_tokens = positions.shape[0]
782 num_q_heads = q_pe.shape[1]
@@ -1028,18 +794,7 @@ def fused_q(
794 assert index_weights is not None
795 num_index_q_heads = index_q.shape[1]
796 index_q_head_dim = index_q.shape[2]
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)
797 + grid_heads = max(mqa_grid_heads, num_index_q_heads)
798 if quantize_mqa:
799 # fp8 path: pack [ql_nope; q_pe] into a single fp8 tensor.
800 mqa_q_fp8 = torch.empty(
@@ -1060,27 +815,6 @@ def fused_q(
815
816 index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
817 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()
818 _fused_q_kernel[(3, num_tokens, grid_heads)](
819 positions,
820 q_pe,
@@ -1122,8 +856,6 @@ def fused_q(
856 HAS_INDEXER=has_indexer,
857 INDEX_ROPE_INTERLEAVE=index_rope_interleave,
858 QUANTIZE_MQA=quantize_mqa,
1125 - USE_PDL=use_pdl,
1126 - launch_pdl=use_pdl,
859 # num_warps=1 is optimal here: each program is a single 128-element
860 # rope+quant, so the kernel is program-count/occupancy bound, not
861 # per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).
vllm/models/deepseek_v32/nvidia/model.py
-4
@@ -6,7 +6,6 @@ from itertools import islice
6
7 import torch
8
9 -from vllm.compilation.decorators import support_torch_compile
9 from vllm.config import VllmConfig
10 from vllm.distributed import (
11 get_pp_group,
@@ -40,7 +39,6 @@ from vllm.model_executor.models.utils import (
39 sequence_parallel_chunk,
40 )
41 from vllm.sequence import IntermediateTensors
43 -from vllm.v1.attention.backends.mla.sparse_utils import register_phys_shadow
42
43 from .attention import DeepseekV32Attention
44 from .fused_ops import fused_allreduce_rms_norm
@@ -179,7 +177,6 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
177 return hidden_states, residual
178
179
182 -@support_torch_compile
180 class DeepseekV32Model(torch.nn.Module):
181 fall_back_to_pt_during_load = False
182
@@ -203,7 +200,6 @@ class DeepseekV32Model(torch.nn.Module):
200 dtype=torch.int32,
201 device=self.device,
202 )
206 - register_phys_shadow(topk_indices_buffer)
203
204 if get_pp_group().is_first_rank:
205 self.embed_tokens = VocabParallelEmbedding(
vllm/models/deepseek_v32/nvidia/mtp.py
+6 -76
@@ -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
9 from vllm._aiter_ops import rocm_aiter_ops
10 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,12 +40,7 @@ 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 -)
43
48 -from .fused_ops import fused_allreduce_rms_norm
44 from .kernels import fused_eh_norm
45 from .model import DeepseekV32DecoderLayer
46
@@ -61,19 +56,6 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
56 self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
57 self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
58 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 - )
59
60 topk_indices_buffer = torch.empty(
61 vllm_config.scheduler_config.max_num_batched_tokens,
@@ -81,7 +63,6 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
63 dtype=torch.int32,
64 device=current_platform.device_type,
65 )
84 - register_phys_shadow(topk_indices_buffer)
66 self.shared_head = SharedHead(
67 config=config, prefix=prefix, quant_config=quant_config
68 )
@@ -110,13 +91,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
91 self.hnorm.weight,
92 self.enorm.variance_epsilon,
93 )
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)
94 + hidden_states = self.eh_proj(eh_input)
95 hidden_states, residual = self.mtp_block(
96 positions=positions, hidden_states=hidden_states, residual=None
97 )
@@ -126,6 +101,9 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
101 positions.shape[0],
102 is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe,
103 )
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)
107 # Recycle the POST-final-norm hidden into the next draft step. The
108 # residual-add is fused into the final RMSNorm so it is computed
109 # exactly once, and the result is returned for both tuple positions:
@@ -136,14 +114,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
114 # is understood by both the V2 speculator (isinstance-tuple check) and
115 # the legacy proposer (model_returns_tuple is True for the
116 # DeepSeekMTPModel architecture).
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 - )
117 + hidden_states, _ = self.shared_head.norm(hidden_states, residual)
118 return hidden_states, hidden_states
119
120
@@ -186,15 +157,6 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
157 if self_attn is not None and hasattr(self_attn, "topk_indices_buffer"):
158 topk_indices_buffer = self_attn.topk_indices_buffer
159 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]
160
161 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
162 return self.embed_tokens(input_ids)
@@ -230,19 +192,6 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
192 # second RMSNorm.
193 return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
194
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 -
195
196 class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
197 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -291,25 +240,6 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
240 ) -> torch.Tensor | None:
241 return self.model.compute_logits(hidden_states, spec_step_idx)
242
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 -
243 def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
244 spec_layer_weight_names = [
245 "embed_tokens",
vllm/models/deepseek_v32/nvidia/ops/__init__.py deleted
-2
@@ -1,2 +0,0 @@
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 deleted
-532
@@ -1,532 +0,0 @@
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,7 +28,6 @@ from vllm.v1.attention.backend import (
28 MultipleOf,
29 )
30 from vllm.v1.attention.backends.mla.sparse_utils import (
31 - phys_shadow,
31 triton_convert_req_index_to_global_index,
32 triton_filter_and_convert_dcp_index,
33 )
@@ -405,36 +404,6 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
404 NUM_TOPK_TOKENS=topk_indices.shape[1],
405 return_valid_counts=True,
406 )
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]
407 else:
408 topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(
409 attn_metadata.req_id_per_token[:num_actual_toks],
vllm/v1/attention/backends/mla/sparse_utils.py
+4 -56
@@ -128,8 +128,6 @@ 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,
131 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
132 """
133 out[token_id, indice_id] =
@@ -180,20 +178,14 @@ def triton_convert_req_index_to_global_index(
178 req_id_c = req_id.contiguous()
179 block_table_c = block_table.contiguous()
180 token_indices_c = token_indices.contiguous()
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
181 + out = torch.empty_like(token_indices_c)
182
183 # Allocate valid count buffer if needed (must be zero-initialized for atomics)
184 valid_counts: torch.Tensor | None = None
185 if return_valid_counts:
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_()
186 + valid_counts = torch.zeros(
187 + num_tokens, dtype=torch.int32, device=token_indices.device
188 + )
189
190 # Strides in elements
191 bt_stride0, bt_stride1 = block_table_c.stride()
@@ -349,47 +341,3 @@ def triton_filter_and_convert_dcp_index(
341 assert valid_counts is not None
342 return out, valid_counts
343 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
+2 -3
@@ -1009,9 +1009,8 @@ 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 bool(
1013 - {"DeepSeekMTPModel", "DeepseekV32MTPModel"}
1014 - & set(self.draft_model_config.hf_config.architectures or [])
1012 + return "DeepSeekMTPModel" in (
1013 + self.draft_model_config.hf_config.architectures or []
1014 )
1015 return self.method not in ("mtp", "draft_model", "dflash")
1016
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+5 -35
@@ -48,18 +48,6 @@ 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 -
51 @property
52 def advance_draft_positions(self) -> bool:
53 """
@@ -110,9 +98,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
98 assert self.prefill_cudagraph_manager is not None
99 if self.prefill_cudagraph_manager.use_breakable_cg:
100 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 -
101 self.prefill_cudagraph_manager.capture(
102 self._prefill,
103 self.model_state,
@@ -123,13 +108,9 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
108 progress_bar_desc="Capturing prefill CUDA graphs",
109 )
110
126 - self.on_prefill_end(self.max_num_reqs, self.max_num_tokens)
127 -
111 if self.num_speculative_steps == 1:
112 return
113
131 - self.on_multi_step_decode_begin(self.max_num_reqs)
132 -
114 # Capture the decode draft generation routine (model forward +
115 # sample + update_draft_inputs) for a single
116 # step.
@@ -144,8 +125,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
125 progress_bar_desc="Capturing decode CUDA graphs",
126 )
127
147 - self.on_multi_step_decode_end(self.max_num_reqs)
148 -
128 @torch.inference_mode()
129 def propose(
130 self,
@@ -174,8 +153,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
153 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
154 is_profile: bool = False,
155 ) -> torch.Tensor:
177 - num_tokens = input_batch.num_tokens
178 - num_tokens_padded = input_batch.num_tokens_after_padding
156 + num_tokens = input_batch.num_tokens_after_padding
157 num_reqs = input_batch.num_reqs
158 max_query_len = input_batch.num_scheduled_tokens.max()
159 max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
@@ -196,7 +174,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
174 )
175 else:
176 hidden_states = last_hidden_states
199 - self.hidden_states[:num_tokens_padded].copy_(hidden_states)
177 + self.hidden_states[:num_tokens].copy_(hidden_states)
178
179 self._copy_request_inputs(
180 num_reqs,
@@ -224,22 +202,20 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
202 num_reqs,
203 # Use the actual number of tokens without padding added by
204 # the target model during FULL cudagraph.
227 - num_tokens,
205 + input_batch.num_tokens,
206 max_query_len,
207 )
208 prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
209 self.prefill_cudagraph_manager,
210 num_reqs,
233 - num_tokens_padded,
211 + num_tokens,
212 uniform_token_count,
213 dp_size=self.dp_size,
214 dp_rank=self.dp_rank,
215 need_eager=is_profile,
216 )
217
240 - self._prepare_eplb_forward(num_tokens)
241 -
242 - self.on_prefill_begin(num_reqs, num_tokens)
218 + self._prepare_eplb_forward(input_batch.num_tokens)
219
220 if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
221 # Replay the full graph for draft prefill.
@@ -259,8 +235,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
235 mm_inputs=mm_inputs,
236 )
237
262 - self.on_prefill_end(num_reqs, num_tokens)
263 -
238 if self.num_speculative_steps == 1:
239 # Early exit.
240 return self.draft_tokens[:num_reqs, :1]
@@ -288,8 +262,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
262 need_eager=is_profile,
263 )
264
291 - self.on_multi_step_decode_begin(num_reqs)
292 -
265 # Generate the remaining num_speculative_steps - 1 draft tokens.
266 self._multi_step_decode(
267 num_reqs,
@@ -299,8 +271,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
271 input_batch.seq_lens_cpu_upper_bound,
272 )
273
302 - self.on_multi_step_decode_end(num_reqs)
303 -
274 return self.draft_tokens[:num_reqs]
275
276 @torch.inference_mode()
vllm/v1/worker/gpu/spec_decode/mtp/speculator.py
+1 -40
@@ -1,10 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import torch
4 import torch.nn as nn
5
7 -from vllm.config import VllmConfig
6 from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
7 AutoRegressiveSpeculator,
8 )
@@ -12,46 +10,9 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
10
11
12 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 -
13 def load_draft_model(
14 self,
15 target_model: nn.Module,
16 target_attn_layer_names: set[str],
17 ) -> nn.Module:
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)
18 + return load_eagle_model(target_model, self.vllm_config)