A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 15:00 UTC
[Feat] DeepSeek V4 Rebased (#40860)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Signed-off-by: qizixi <zixi@inferact.ai> Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Yongye Zhu <yongye@inferact.ai> Co-authored-by: Simon Mo <simon@inferact.ai> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Roy Wang <yasong.wang@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: youkaichao <youkaichao@gmail.com> Co-authored-by: Zhewen Li <jerven.vllm@gmail.com> Co-authored-by: Zijing Liu <liuzijing2014@gmail.com> Co-authored-by: khluu <khluu000@gmail.com> Co-authored-by: qizixi <zixi@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai>
Yifan Qiao committed
Apr 26, 2026 at 18:31 UTC
4d51588e2381018348f1022dfa3a7698899805b7
150 files changed
+16313
-717
CMakeLists.txt
+5
-2
@@ -310,7 +310,9 @@ set(VLLM_EXT_SRC
310
"csrc/torch_bindings.cpp")
311
312
if(VLLM_GPU_LANG STREQUAL "CUDA")
313
- list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
313
+ list(APPEND VLLM_EXT_SRC
314
+ "csrc/minimax_reduce_rms_kernel.cu"
315
+ "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
316
317
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
318
@@ -1051,7 +1053,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1053
list(APPEND VLLM_MOE_EXT_SRC
1054
"csrc/moe/moe_wna16.cu"
1055
"csrc/moe/grouped_topk_kernels.cu"
1054
- "csrc/moe/router_gemm.cu")
1056
+ "csrc/moe/router_gemm.cu"
1057
+ "csrc/moe/topk_softplus_sqrt_kernels.cu")
1058
endif()
1059
1060
if(VLLM_GPU_LANG STREQUAL "CUDA")
cmake/external_projects/deepgemm.cmake
+6
-1
@@ -20,7 +20,7 @@ else()
20
FetchContent_Declare(
21
deepgemm
22
GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
23
- GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
23
+ GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
24
GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
25
GIT_PROGRESS TRUE
26
CONFIGURE_COMMAND ""
@@ -120,6 +120,11 @@ if(DEEPGEMM_ARCHS)
120
COMPONENT _deep_gemm_C
121
FILES_MATCHING PATTERN "*.py")
122
123
+ install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
124
+ DESTINATION vllm/third_party/deep_gemm/mega
125
+ COMPONENT _deep_gemm_C
126
+ FILES_MATCHING PATTERN "*.py")
127
+
128
# Generate envs.py (normally generated by DeepGEMM's setup.py build step)
129
file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
130
"# Pre-installed environment variables\npersistent_envs = dict()\n")
cmake/external_projects/flashmla.cmake
+1
-1
@@ -19,7 +19,7 @@ else()
19
FetchContent_Declare(
20
flashmla
21
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22
- GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
22
+ GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
23
GIT_PROGRESS TRUE
24
CONFIGURE_COMMAND ""
25
BUILD_COMMAND ""
csrc/cpu/pos_encoding.cpp
+6
-1
@@ -178,7 +178,12 @@ void rotary_embedding_gptj_impl(
178
179
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
180
std::optional<torch::Tensor> key, int64_t head_size,
181
- torch::Tensor& cos_sin_cache, bool is_neox) {
181
+ torch::Tensor& cos_sin_cache, bool is_neox,
182
+ int64_t rope_dim_offset, bool inverse) {
183
+ TORCH_CHECK(rope_dim_offset == 0,
184
+ "rope_dim_offset != 0 is not supported on CPU");
185
+ TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
186
+
187
int num_tokens = positions.numel();
188
int rot_dim = cos_sin_cache.size(1);
189
int num_heads = query.size(-1) / head_size;
csrc/cpu/torch_bindings.cpp
+2
-1
@@ -263,7 +263,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
263
ops.def(
264
"rotary_embedding(Tensor positions, Tensor! query,"
265
" Tensor!? key, int head_size,"
266
- " Tensor cos_sin_cache, bool is_neox) -> ()");
266
+ " Tensor cos_sin_cache, bool is_neox, int "
267
+ "rope_dim_offset=0, bool inverse=False) -> ()");
268
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
269
270
// Quantization
csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
new
+477
@@ -0,0 +1,477 @@
1
+/*
2
+ * SPDX-License-Identifier: Apache-2.0
3
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+ *
5
+ * Horizontally-fused DeepseekV4-MLA kernel:
6
+ * - Q side: per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM
7
+ * - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
8
+ * cache insert
9
+ *
10
+ * Structured after `applyMLARopeAndAssignQKVKernelGeneration` in
11
+ * TensorRT-LLM's mlaKernels.cu: one kernel, one grid, with head-slot
12
+ * dispatch choosing Q vs KV work per warp. The per-warp RMSNorm/RoPE
13
+ * skeleton is adapted from vllm-deepseek_v4's existing
14
+ * `fusedQKNormRopeKernel` (csrc/fused_qknorm_rope_kernel.cu).
15
+ *
16
+ * Assumptions (hard-coded for DeepseekV4 attention):
17
+ * HEAD_DIM = 512
18
+ * ROPE_DIM = 64 (RoPE applied to dims [NOPE_DIM, HEAD_DIM))
19
+ * NOPE_DIM = 448
20
+ * QUANT_BLOCK = 64 (UE8M0 FP8 quant block)
21
+ * FP8_MAX = 448.0f
22
+ * is_neox=false (GPT-J interleaved pairs)
23
+ * cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
24
+ * second along last dim; each half is rope_dim/2 = 32 values)
25
+ *
26
+ * Cache layout per paged-cache block (block_size tokens):
27
+ * [0, bs*576): token data, 448 fp8 + 128 bf16 each
28
+ * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token
29
+ */
30
+
31
+#include <cmath>
32
+#include <cuda_fp8.h>
33
+#include <cuda_runtime.h>
34
+#include <type_traits>
35
+
36
+#include <ATen/cuda/CUDAContext.h>
37
+#include <c10/cuda/CUDAGuard.h>
38
+#include <torch/cuda.h>
39
+
40
+#include "cuda_compat.h"
41
+#include "dispatch_utils.h"
42
+#include "type_convert.cuh"
43
+
44
+#ifndef FINAL_MASK
45
+ #define FINAL_MASK 0xffffffffu
46
+#endif
47
+
48
+namespace vllm {
49
+namespace deepseek_v4_fused_ops {
50
+
51
+namespace {
52
+inline int getSMVersion() {
53
+ auto* props = at::cuda::getCurrentDeviceProperties();
54
+ return props->major * 10 + props->minor;
55
+}
56
+} // namespace
57
+
58
+// ────────────────────────────────────────────────────────────────────────────
59
+// Constants
60
+// ────────────────────────────────────────────────────────────────────────────
61
+constexpr int kHeadDim = 512;
62
+constexpr int kRopeDim = 64;
63
+constexpr int kNopeDim = kHeadDim - kRopeDim; // 448
64
+constexpr int kQuantBlock = 64;
65
+constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7
66
+constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad)
67
+constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576
68
+constexpr float kFp8Max = 448.0f;
69
+
70
+// Per-warp layout: 32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
71
+constexpr int kNumLanes = 32;
72
+constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16
73
+
74
+// ────────────────────────────────────────────────────────────────────────────
75
+// Small inline helpers
76
+// ────────────────────────────────────────────────────────────────────────────
77
+__device__ __forceinline__ float warp4MaxAbs(float val) {
78
+ // Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
79
+ float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
80
+ val = fmaxf(val, peer);
81
+ peer = __shfl_xor_sync(FINAL_MASK, val, 2);
82
+ val = fmaxf(val, peer);
83
+ return val;
84
+}
85
+
86
+template <typename T>
87
+__device__ __forceinline__ float warpSum(float val) {
88
+#pragma unroll
89
+ for (int mask = 16; mask > 0; mask >>= 1) {
90
+ val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
91
+ }
92
+ return val;
93
+}
94
+
95
+// ────────────────────────────────────────────────────────────────────────────
96
+// Kernel
97
+// ────────────────────────────────────────────────────────────────────────────
98
+//
99
+// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
100
+// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
101
+// warp handles one (token, head_slot) pair. head_slot < num_heads_q →
102
+// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q → KV
103
+// branch (RoPE + UE8M0 quant + insert)
104
+//
105
+// With DP padding, q/kv/position_ids can have more rows than slot_mapping.
106
+// The Q branch covers all `num_tokens_full` rows (downstream attention uses
107
+// them). The KV branch only inserts the first `num_tokens_insert` tokens
108
+// (= slot_mapping length) into the paged cache.
109
+//
110
+template <typename scalar_t_in>
111
+__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
112
+ scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place
113
+ scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16
114
+ uint8_t* __restrict__ k_cache, // [num_blocks, block_stride]
115
+ int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64
116
+ int64_t const* __restrict__ position_ids, // [N] i64
117
+ float const* __restrict__ cos_sin_cache, // [max_pos, 64] fp32
118
+ float const eps,
119
+ int const num_tokens_full, // = q.size(0) = kv.size(0)
120
+ int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full
121
+ int const num_heads_q, // H
122
+ int const cache_block_size, // tokens per paged-cache block
123
+ int const kv_block_stride) { // bytes per paged-cache block
124
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
125
+ // BF16 _typeConvert specialization is unavailable on pre-Ampere. The
126
+ // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
127
+ // no-op stub for sm_70/sm_75 to keep multi-arch builds happy.
128
+ if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
129
+ return;
130
+ } else {
131
+#endif
132
+ using Converter = vllm::_typeConvert<scalar_t_in>;
133
+
134
+ int const warpsPerBlock = blockDim.x / 32;
135
+ int const warpId = threadIdx.x / 32;
136
+ int const laneId = threadIdx.x % 32;
137
+ int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
138
+
139
+ int const total_slots_per_token = num_heads_q + 1;
140
+ int const tokenIdx = globalWarpIdx / total_slots_per_token;
141
+ int const slotIdx = globalWarpIdx % total_slots_per_token;
142
+ if (tokenIdx >= num_tokens_full) return;
143
+
144
+ bool const isKV = (slotIdx == num_heads_q);
145
+ // KV branch: skip DP-padded tokens (no slot reserved for them).
146
+ if (isKV && tokenIdx >= num_tokens_insert) return;
147
+
148
+ // PDL: wait for predecessor kernel (upstream q/kv producer) to signal
149
+ // before touching any global memory. No-op when PDL is not enabled on
150
+ // the launch. The CUDA runtime wrapper emits the griddepcontrol.wait
151
+ // PTX with the required memory clobber internally.
152
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
153
+ cudaGridDependencySynchronize();
154
+#endif
155
+
156
+ // Dim range this lane owns within the 512-wide head.
157
+ int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
158
+
159
+ // ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
160
+ float elements[kElemsPerLane];
161
+ float sumOfSquares = 0.0f;
162
+
163
+ scalar_t_in const* src_ptr;
164
+ if (isKV) {
165
+ src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
166
+ } else {
167
+ int64_t const q_row_offset =
168
+ (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
169
+ dim_base;
170
+ src_ptr = q_inout + q_row_offset;
171
+ }
172
+
173
+ // Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector
174
+ // type and bitcast to scalar_t_in packed pairs for conversion.
175
+ uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
176
+ uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
177
+
178
+ {
179
+ typename Converter::packed_hip_type const* p0 =
180
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
181
+ typename Converter::packed_hip_type const* p1 =
182
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
183
+// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
184
+#pragma unroll
185
+ for (int i = 0; i < 4; i++) {
186
+ float2 f2 = Converter::convert(p0[i]);
187
+ elements[2 * i] = f2.x;
188
+ elements[2 * i + 1] = f2.y;
189
+ }
190
+#pragma unroll
191
+ for (int i = 0; i < 4; i++) {
192
+ float2 f2 = Converter::convert(p1[i]);
193
+ elements[8 + 2 * i] = f2.x;
194
+ elements[8 + 2 * i + 1] = f2.y;
195
+ }
196
+ }
197
+
198
+ // ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
199
+ // Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
200
+ // The downstream bf16 round only happens at the final store.
201
+ if (!isKV) {
202
+#pragma unroll
203
+ for (int i = 0; i < kElemsPerLane; i++) {
204
+ sumOfSquares += elements[i] * elements[i];
205
+ }
206
+ sumOfSquares = warpSum<float>(sumOfSquares);
207
+ float const rms_rcp =
208
+ rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
209
+#pragma unroll
210
+ for (int i = 0; i < kElemsPerLane; i++) {
211
+ elements[i] = elements[i] * rms_rcp;
212
+ }
213
+ }
214
+
215
+ // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
216
+ // All math in fp32. cos_sin_cache is loaded as fp32 (its native storage).
217
+ bool const is_rope_lane = dim_base >= kNopeDim;
218
+ if (is_rope_lane) {
219
+ int64_t const pos = position_ids[tokenIdx];
220
+ constexpr int kHalfRope = kRopeDim / 2; // 32
221
+ float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
222
+ float const* sin_ptr = cos_ptr + kHalfRope;
223
+
224
+ int const rope_local_base = dim_base - kNopeDim; // in [0, 64) step 16
225
+#pragma unroll
226
+ for (int p = 0; p < kElemsPerLane / 2; p++) {
227
+ int const pair_dim = rope_local_base + 2 * p;
228
+ int const half_idx = pair_dim / 2;
229
+ float const cos_v = VLLM_LDG(cos_ptr + half_idx);
230
+ float const sin_v = VLLM_LDG(sin_ptr + half_idx);
231
+ float const x_even = elements[2 * p];
232
+ float const x_odd = elements[2 * p + 1];
233
+ elements[2 * p] = x_even * cos_v - x_odd * sin_v;
234
+ elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
235
+ }
236
+ }
237
+
238
+ // ═══════════════════════════════════════════════════════════════════════
239
+ // Q branch: cast to bf16 and store back in place.
240
+ // ═══════════════════════════════════════════════════════════════════════
241
+ if (!isKV) {
242
+ uint4 out0, out1;
243
+ typename Converter::packed_hip_type* po0 =
244
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
245
+ typename Converter::packed_hip_type* po1 =
246
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
247
+#pragma unroll
248
+ for (int i = 0; i < 4; i++) {
249
+ po0[i] = Converter::convert(
250
+ make_float2(elements[2 * i], elements[2 * i + 1]));
251
+ }
252
+#pragma unroll
253
+ for (int i = 0; i < 4; i++) {
254
+ po1[i] = Converter::convert(
255
+ make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
256
+ }
257
+ scalar_t_in* dst =
258
+ q_inout +
259
+ (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
260
+ dim_base;
261
+ *reinterpret_cast<uint4*>(dst) = out0;
262
+ *reinterpret_cast<uint4*>(dst + 8) = out1;
263
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
264
+ cudaTriggerProgrammaticLaunchCompletion();
265
+#endif
266
+ return;
267
+ }
268
+
269
+ // ═══════════════════════════════════════════════════════════════════════
270
+ // KV branch.
271
+ // ═══════════════════════════════════════════════════════════════════════
272
+ int64_t const slot_id = slot_mapping[tokenIdx];
273
+ if (slot_id < 0) {
274
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
275
+ cudaTriggerProgrammaticLaunchCompletion();
276
+#endif
277
+ return;
278
+ }
279
+
280
+ int64_t const block_idx = slot_id / cache_block_size;
281
+ int64_t const pos_in_block = slot_id % cache_block_size;
282
+ uint8_t* block_base =
283
+ k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
284
+ uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
285
+ uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
286
+ uint8_t* token_scale_ptr =
287
+ block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
288
+ pos_in_block * kScaleBytesPerToken;
289
+
290
+ // Round K to bf16 first, matching the unfused reference path where K is
291
+ // materialized as bf16 before K quantization. absmax, clamp, and FP8
292
+ // quant below all run on these bf16-rounded values.
293
+#pragma unroll
294
+ for (int i = 0; i < kElemsPerLane; i++) {
295
+ elements[i] = Converter::convert(Converter::convert(elements[i]));
296
+ }
297
+
298
+ // Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
299
+ // shuffle requires full participation). RoPE lanes contribute garbage,
300
+ // but their values are gated out below via `!is_rope_lane`.
301
+ float local_absmax = 0.0f;
302
+#pragma unroll
303
+ for (int i = 0; i < kElemsPerLane; i++) {
304
+ local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
305
+ }
306
+ float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
307
+ float const exponent = ceilf(log2f(absmax / kFp8Max));
308
+ float const inv_scale = exp2f(-exponent);
309
+
310
+ if (!is_rope_lane) {
311
+ // ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
312
+ uint8_t out_bytes[kElemsPerLane];
313
+#pragma unroll
314
+ for (int i = 0; i < kElemsPerLane; i++) {
315
+ float scaled = elements[i] * inv_scale;
316
+ scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
317
+ __nv_fp8_storage_t s =
318
+ __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
319
+ out_bytes[i] = static_cast<uint8_t>(s);
320
+ }
321
+ // One 16-byte STG per lane.
322
+ *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
323
+ *reinterpret_cast<uint4 const*>(out_bytes);
324
+
325
+ // Lane (4k) of each 4-lane group writes the scale byte for block k<7.
326
+ if ((laneId & 3) == 0) {
327
+ int const q_block_idx = laneId >> 2; // 0..6 for NoPE lanes
328
+ float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
329
+ token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
330
+ }
331
+ // Lane 0 also writes the padding byte at index 7.
332
+ if (laneId == 0) {
333
+ token_scale_ptr[kNumQuantBlocks] = 0; // pad
334
+ }
335
+ } else {
336
+ // ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
337
+ uint4 out0, out1;
338
+ typename Converter::packed_hip_type* po0 =
339
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
340
+ typename Converter::packed_hip_type* po1 =
341
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
342
+#pragma unroll
343
+ for (int i = 0; i < 4; i++) {
344
+ po0[i] = Converter::convert(
345
+ make_float2(elements[2 * i], elements[2 * i + 1]));
346
+ }
347
+#pragma unroll
348
+ for (int i = 0; i < 4; i++) {
349
+ po1[i] = Converter::convert(
350
+ make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
351
+ }
352
+ int const rope_local_base = dim_base - kNopeDim; // in [0, 64)
353
+ scalar_t_in* bf16_dst =
354
+ reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
355
+ *reinterpret_cast<uint4*>(bf16_dst) = out0;
356
+ *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
357
+ }
358
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
359
+ cudaTriggerProgrammaticLaunchCompletion();
360
+#endif
361
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
362
+ }
363
+#endif
364
+}
365
+
366
+// ────────────────────────────────────────────────────────────────────────────
367
+// Launch wrapper
368
+// ────────────────────────────────────────────────────────────────────────────
369
+template <typename scalar_t_in>
370
+void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
371
+ scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
372
+ int64_t const* slot_mapping, int64_t const* position_ids,
373
+ float const* cos_sin_cache, float const eps, int const num_tokens_full,
374
+ int const num_tokens_insert, int const num_heads_q,
375
+ int const cache_block_size, int const kv_block_stride,
376
+ cudaStream_t stream) {
377
+ constexpr int kBlockSize = 256;
378
+ constexpr int kWarpsPerBlock = kBlockSize / 32;
379
+ int64_t const total_warps =
380
+ static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
381
+ int const grid =
382
+ static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
383
+
384
+ // PDL: enable programmatic stream serialization whenever the hardware
385
+ // supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable,
386
+ // so leave numAttrs = 0 and launch as a regular kernel.
387
+ static int const sm_version = getSMVersion();
388
+ // Host-side guard: the device kernel body is compiled as a no-op for
389
+ // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
390
+ // unavailable there. Refuse the launch loudly instead of silently
391
+ // skipping the work.
392
+ TORCH_CHECK(
393
+ sm_version >= 80,
394
+ "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
395
+ "(Ampere or newer); got sm_",
396
+ sm_version);
397
+ cudaLaunchConfig_t config;
398
+ config.gridDim = dim3(grid);
399
+ config.blockDim = dim3(kBlockSize);
400
+ config.dynamicSmemBytes = 0;
401
+ config.stream = stream;
402
+ cudaLaunchAttribute attrs[1];
403
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
404
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
405
+ config.attrs = attrs;
406
+ config.numAttrs = (sm_version >= 90) ? 1 : 0;
407
+
408
+ cudaLaunchKernelEx(
409
+ &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
410
+ q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
411
+ num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
412
+ kv_block_stride);
413
+}
414
+
415
+} // namespace deepseek_v4_fused_ops
416
+} // namespace vllm
417
+
418
+// ────────────────────────────────────────────────────────────────────────────
419
+// Torch op wrapper
420
+// ────────────────────────────────────────────────────────────────────────────
421
+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
422
+ torch::Tensor& q, // [N, H, 512] bf16, in place
423
+ torch::Tensor const& kv, // [N, 512] bf16 (read-only)
424
+ torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8
425
+ torch::Tensor const& slot_mapping, // [N] int64
426
+ torch::Tensor const& position_ids, // [N] int64
427
+ torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
428
+ double eps, int64_t cache_block_size) {
429
+ TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
430
+ TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
431
+ TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
432
+ TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
433
+ "slot_mapping must be int64 CUDA");
434
+ TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
435
+ "position_ids must be int64 CUDA");
436
+ TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
437
+ TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
438
+ TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
439
+ TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
440
+ TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
441
+ TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
442
+ "cos_sin_cache shape [max_pos, 64]");
443
+ TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
444
+ "cos_sin_cache must be float32");
445
+
446
+ // With DP padding, slot_mapping can be shorter than q/kv/positions.
447
+ // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
448
+ // KV quant+insert runs only on the first slot_mapping.size(0) rows.
449
+ int const num_tokens_full = static_cast<int>(q.size(0));
450
+ int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
451
+ TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
452
+ static_cast<int>(position_ids.size(0)) == num_tokens_full,
453
+ "q/kv/position_ids row counts must match");
454
+ TORCH_CHECK(num_tokens_insert <= num_tokens_full,
455
+ "slot_mapping must not exceed q row count");
456
+ int const num_heads_q = static_cast<int>(q.size(1));
457
+ int const cache_block_size_i = static_cast<int>(cache_block_size);
458
+ int const kv_block_stride = static_cast<int>(k_cache.stride(0));
459
+
460
+ at::cuda::OptionalCUDAGuard device_guard(device_of(q));
461
+ auto stream = at::cuda::getCurrentCUDAStream();
462
+
463
+ VLLM_DISPATCH_HALF_TYPES(
464
+ q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
465
+ using qkv_scalar_t = scalar_t;
466
+ vllm::deepseek_v4_fused_ops::
467
+ launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
468
+ reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
469
+ reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
470
+ reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
471
+ reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
472
+ reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
473
+ cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
474
+ num_tokens_full, num_tokens_insert, num_heads_q,
475
+ cache_block_size_i, kv_block_stride, stream);
476
+ });
477
+}
csrc/layernorm_kernels.cu
+15
-7
@@ -77,7 +77,8 @@ __global__ void rms_norm_kernel(
77
#pragma unroll
78
for (int j = 0; j < VEC_SIZE; j++) {
79
float x = static_cast<float>(src1.val[j]);
80
- dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
80
+ float w = static_cast<float>(src2.val[j]);
81
+ dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
82
}
83
v_out[i] = dst;
84
}
@@ -134,10 +135,17 @@ fused_add_rms_norm_kernel(
135
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
136
int id = blockIdx.x * vec_hidden_size + idx;
137
int64_t strided_id = blockIdx.x * vec_input_stride + idx;
137
- _f16Vec<scalar_t, width> temp = residual_v[id];
138
- temp *= s_variance;
139
- temp *= weight_v[idx];
140
- input_v[strided_id] = temp;
138
+ _f16Vec<scalar_t, width> res = residual_v[id];
139
+ _f16Vec<scalar_t, width> w = weight_v[idx];
140
+ _f16Vec<scalar_t, width> out;
141
+ using Converter = _typeConvert<scalar_t>;
142
+#pragma unroll
143
+ for (int j = 0; j < width; ++j) {
144
+ float x = Converter::convert(res.data[j]);
145
+ float wf = Converter::convert(w.data[j]);
146
+ out.data[j] = Converter::convert(x * s_variance * wf);
147
+ }
148
+ input_v[strided_id] = out;
149
}
150
}
151
@@ -174,8 +182,8 @@ fused_add_rms_norm_kernel(
182
183
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
184
float x = (float)residual[blockIdx.x * hidden_size + idx];
177
- input[blockIdx.x * input_stride + idx] =
178
- ((scalar_t)(x * s_variance)) * weight[idx];
185
+ float w = (float)weight[idx];
186
+ input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
187
}
188
}
189
csrc/layernorm_quant_kernels.cu
+28
-10
@@ -65,9 +65,16 @@ __global__ void rms_norm_static_fp8_quant_kernel(
65
#pragma unroll
66
for (int j = 0; j < VEC_SIZE; j++) {
67
float x = static_cast<float>(src1.val[j]);
68
- float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
68
+ float w = static_cast<float>(src2.val[j]);
69
+ // Round normalized result through scalar_t to match the precision of the
70
+ // unfused composite (rms_norm writes scalar_t, then
71
+ // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
72
+ // Without this round, the fused path is strictly more accurate and
73
+ // disagrees with the composite at exact E4M3 quantization tie boundaries.
74
+ scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
75
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
70
- scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
76
+ scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
77
+ scale_inv);
78
}
79
}
80
}
@@ -127,13 +134,21 @@ fused_add_rms_norm_static_fp8_quant_kernel(
134
135
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
136
int id = blockIdx.x * vec_hidden_size + idx;
130
- _f16Vec<scalar_t, width> temp = residual_v[id];
131
- temp *= s_variance;
132
- temp *= weight_v[idx];
137
+ _f16Vec<scalar_t, width> res = residual_v[id];
138
+ _f16Vec<scalar_t, width> w = weight_v[idx];
139
+ using Converter = _typeConvert<scalar_t>;
140
+ using HipT = typename Converter::hip_type;
141
#pragma unroll
142
for (int i = 0; i < width; ++i) {
135
- out[id * width + i] =
136
- scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
143
+ float x = Converter::convert(res.data[i]);
144
+ float wf = Converter::convert(w.data[i]);
145
+ // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
146
+ // to match the unfused composite path at FP8 boundaries. We use the
147
+ // backend's hip_type for the intermediate since c10::Half/BFloat16 has
148
+ // ambiguous conversions on CUDA and no implicit conversion on ROCm.
149
+ HipT out_norm_h = Converter::convert(x * s_variance * wf);
150
+ out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
151
+ Converter::convert(out_norm_h), scale_inv);
152
}
153
}
154
}
@@ -176,9 +191,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(
191
192
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
193
float x = (float)residual[blockIdx.x * hidden_size + idx];
179
- float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
180
- out[blockIdx.x * hidden_size + idx] =
181
- scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
194
+ float w = (float)weight[idx];
195
+ // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
196
+ // to match the unfused composite path at FP8 boundaries.
197
+ scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
198
+ out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
199
+ static_cast<float>(out_norm), scale_inv);
200
}
201
}
202
csrc/moe/moe_ops.h
+9
@@ -12,6 +12,15 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
12
torch::Tensor& gating_output, bool renormalize,
13
std::optional<torch::Tensor> bias);
14
15
+void topk_softplus_sqrt(torch::Tensor& topk_weights,
16
+ torch::Tensor& topk_indices,
17
+ torch::Tensor& token_expert_indices,
18
+ torch::Tensor& gating_output, bool renormalize,
19
+ double routed_scaling_factor,
20
+ const c10::optional<torch::Tensor>& correction_bias,
21
+ const c10::optional<torch::Tensor>& input_ids,
22
+ const c10::optional<torch::Tensor>& tid2eid);
23
+
24
void moe_sum(torch::Tensor& input, torch::Tensor& output);
25
26
void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
csrc/moe/topk_softplus_sqrt_kernels.cu
new
+715
@@ -0,0 +1,715 @@
1
+/*
2
+ * Adapted from
3
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
4
+ * Copyright (c) 2024, The vLLM team.
5
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
6
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
7
+ *
8
+ * Licensed under the Apache License, Version 2.0 (the "License");
9
+ * you may not use this file except in compliance with the License.
10
+ * You may obtain a copy of the License at
11
+ *
12
+ * http://www.apache.org/licenses/LICENSE-2.0
13
+ *
14
+ * Unless required by applicable law or agreed to in writing, software
15
+ * distributed under the License is distributed on an "AS IS" BASIS,
16
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
+ * See the License for the specific language governing permissions and
18
+ * limitations under the License.
19
+ */
20
+#include <type_traits>
21
+#include <torch/all.h>
22
+#include <ATen/cuda/CUDAContext.h>
23
+#include <c10/cuda/CUDAGuard.h>
24
+#include "../cuda_compat.h"
25
+#include "../cub_helpers.h"
26
+#ifndef USE_ROCM
27
+ #include <cuda_bf16.h>
28
+ #include <cuda_fp16.h>
29
+#else
30
+ #include <hip/hip_bf16.h>
31
+ #include <hip/hip_fp16.h>
32
+typedef __hip_bfloat16 __nv_bfloat16;
33
+typedef __hip_bfloat162 __nv_bfloat162;
34
+#endif
35
+
36
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
37
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
38
+
39
+namespace vllm {
40
+namespace moe {
41
+
42
+/// Aligned array type
43
+template <typename T,
44
+ /// Number of elements in the array
45
+ int N,
46
+ /// Alignment requirement in bytes
47
+ int Alignment = sizeof(T) * N>
48
+struct alignas(Alignment) AlignedArray {
49
+ T data[N];
50
+};
51
+
52
+template <typename T>
53
+__device__ __forceinline__ float toFloat(T value) {
54
+ if constexpr (std::is_same_v<T, float>) {
55
+ return value;
56
+ } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
57
+ return __bfloat162float(value);
58
+ } else if constexpr (std::is_same_v<T, __half>) {
59
+ return __half2float(value);
60
+ }
61
+}
62
+
63
+#define FINAL_MASK 0xffffffff
64
+template <typename T>
65
+__inline__ __device__ T warpReduceSum(T val) {
66
+#pragma unroll
67
+ for (int mask = 16; mask > 0; mask >>= 1)
68
+ val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
69
+ return val;
70
+}
71
+
72
+// ====================== TopK softplus_sqrt things
73
+// ===============================
74
+
75
+/*
76
+ A Top-K gating softplus_sqrt written to exploit when the number of experts in
77
+ the MoE layers are a small power of 2. This allows us to cleanly share the
78
+ rows among the threads in a single warp and eliminate communication between
79
+ warps (so no need to use shared mem).
80
+
81
+ It fuses the sigmoid, max and argmax into a single kernel.
82
+
83
+ Limitations:
84
+ 1) This implementation is optimized for when the number of experts is a small
85
+ power of 2. Additionally it also supports when number of experts is multiple
86
+ of 64 which is still faster than the computing sigmoid and topK separately
87
+ (only tested on CUDA yet). 2) This implementation assumes k is small, but will
88
+ work for any k.
89
+*/
90
+
91
+template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
92
+ int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
93
+ typename InputType = float>
94
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
95
+ void topkGatingSoftplusSqrt(
96
+ const InputType* input, const bool* finished, float* output,
97
+ const int num_rows, IndType* indices, int* source_rows, const int k,
98
+ const int start_expert, const int end_expert, const bool renormalize,
99
+ double routed_scaling_factor, const float* correction_bias,
100
+ const IndType* input_ids, const IndType* tid2eid) {
101
+ static_assert(std::is_same_v<InputType, float> ||
102
+ std::is_same_v<InputType, __nv_bfloat16> ||
103
+ std::is_same_v<InputType, __half>,
104
+ "InputType must be float, __nv_bfloat16, or __half");
105
+
106
+ // We begin by enforcing compile time assertions and setting up compile time
107
+ // constants.
108
+ static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
109
+ "BYTES_PER_LDG must be power of 2");
110
+ static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
111
+
112
+ // Number of bytes each thread pulls in per load
113
+ static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
114
+ static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
115
+ static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
116
+ static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
117
+
118
+ if constexpr (std::is_same_v<InputType, __nv_bfloat16> ||
119
+ std::is_same_v<InputType, __half>) {
120
+ static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
121
+ "ELTS_PER_LDG must be 1 or even for 16-bit conversion");
122
+ }
123
+
124
+ // Restrictions based on previous section.
125
+ static_assert(
126
+ VPT % ELTS_PER_LDG == 0,
127
+ "The elements per thread must be a multiple of the elements per ldg");
128
+ static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0,
129
+ "The threads per row must cleanly divide the threads per warp");
130
+ static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
131
+ "THREADS_PER_ROW must be power of 2");
132
+ static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM,
133
+ "THREADS_PER_ROW can be at most warp size");
134
+
135
+ // We have NUM_EXPERTS elements per row. We specialize for small #experts
136
+ static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
137
+ static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
138
+ static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
139
+
140
+ // Restrictions for previous section.
141
+ static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
142
+ "The elts per row must cleanly divide the total elt per warp");
143
+
144
+ // ===================== From this point, we finally start computing run-time
145
+ // variables. ========================
146
+
147
+ // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
148
+ // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
149
+ // rows. We start by computing the start row for each block.
150
+ const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
151
+
152
+ // Now, using the base row per thread block, we compute the base row per warp.
153
+ const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
154
+
155
+ // The threads in a warp are split into sub-groups that will work on a row.
156
+ // We compute row offset for each thread sub-group
157
+ const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
158
+ const int thread_row = warp_base_row + thread_row_in_warp;
159
+
160
+ // Threads with indices out of bounds should early exit here.
161
+ if (thread_row >= num_rows) {
162
+ return;
163
+ }
164
+ const bool row_is_active = finished ? !finished[thread_row] : true;
165
+
166
+ // We finally start setting up the read pointers for each thread. First, each
167
+ // thread jumps to the start of the row it will read.
168
+ const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
169
+
170
+ // Now, we compute the group each thread belong to in order to determine the
171
+ // first column to start loads.
172
+ const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
173
+ const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
174
+ const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
175
+
176
+ // Finally, we pull in the data from global mem
177
+ float row_chunk[VPT];
178
+
179
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
180
+ asm volatile("griddepcontrol.wait;");
181
+#endif
182
+
183
+ // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
184
+ // to float
185
+ if constexpr (std::is_same_v<InputType, float>) {
186
+ using VecType = AlignedArray<float, ELTS_PER_LDG>;
187
+ VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
188
+ const VecType* vec_thread_read_ptr =
189
+ reinterpret_cast<const VecType*>(thread_read_ptr);
190
+#pragma unroll
191
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
192
+ row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
193
+ }
194
+ } else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
195
+ if constexpr (ELTS_PER_LDG >= 2) {
196
+ using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
197
+ float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
198
+ const VecType* vec_thread_read_ptr =
199
+ reinterpret_cast<const VecType*>(thread_read_ptr);
200
+#pragma unroll
201
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
202
+ VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
203
+ int base_idx_f2 = ii * ELTS_PER_LDG / 2;
204
+#pragma unroll
205
+ for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
206
+ row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
207
+ *reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2));
208
+ }
209
+ }
210
+ } else { // ELTS_PER_LDG == 1
211
+#pragma unroll
212
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
213
+ const __nv_bfloat16* scalar_ptr =
214
+ thread_read_ptr + ii * THREADS_PER_ROW;
215
+ row_chunk[ii] = __bfloat162float(*scalar_ptr);
216
+ }
217
+ }
218
+ } else if constexpr (std::is_same_v<InputType, __half>) {
219
+ if constexpr (ELTS_PER_LDG >= 2) {
220
+ using VecType = AlignedArray<__half, ELTS_PER_LDG>;
221
+ float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
222
+ const VecType* vec_thread_read_ptr =
223
+ reinterpret_cast<const VecType*>(thread_read_ptr);
224
+#pragma unroll
225
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
226
+ VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
227
+ int base_idx_f2 = ii * ELTS_PER_LDG / 2;
228
+#pragma unroll
229
+ for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
230
+ row_chunk_f2[base_idx_f2 + jj] = __half22float2(
231
+ *reinterpret_cast<const __half2*>(vec.data + jj * 2));
232
+ }
233
+ }
234
+ } else { // ELTS_PER_LDG == 1
235
+#pragma unroll
236
+ for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
237
+ const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
238
+ row_chunk[ii] = __half2float(*scalar_ptr);
239
+ }
240
+ }
241
+ }
242
+ constexpr float threshold = 20.0f;
243
+ constexpr float beta = 1.0f;
244
+
245
+ // Hash MoE path: indices are predetermined from lookup table
246
+ if constexpr (USE_HASH) {
247
+ const IndType token_id = input_ids[thread_row];
248
+ const IndType* expert_indices_for_token = tid2eid + token_id * k;
249
+#pragma unroll
250
+ for (int ii = 0; ii < VPT; ++ii) {
251
+ float val = row_chunk[ii];
252
+ float val_b = val * beta;
253
+ val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
254
+ row_chunk[ii] = sqrtf(val);
255
+ }
256
+ float selected_sum = 0.f;
257
+#pragma unroll
258
+ for (int k_idx = 0; k_idx < k; ++k_idx) {
259
+ const int expert = expert_indices_for_token[k_idx];
260
+ const int idx = k * thread_row + k_idx;
261
+ for (int ii = 0; ii < VPT; ++ii) {
262
+ const int group_id = ii / ELTS_PER_LDG;
263
+ const int local_id = ii % ELTS_PER_LDG;
264
+ const int expert_idx = first_elt_read_by_thread +
265
+ group_id * THREADS_PER_ROW * ELTS_PER_LDG +
266
+ local_id;
267
+ if (expert == expert_idx) {
268
+ indices[idx] = expert;
269
+ selected_sum += row_chunk[ii];
270
+ break;
271
+ }
272
+ }
273
+ }
274
+ // Compute per-thread scale (using warp reduction when renormalizing).
275
+ if (renormalize) {
276
+ selected_sum = warpReduceSum(selected_sum);
277
+ }
278
+ float scale = static_cast<float>(routed_scaling_factor);
279
+ if (renormalize) {
280
+ const float denom = selected_sum > 0.f ? selected_sum : 1.f;
281
+ scale /= denom;
282
+ }
283
+
284
+#pragma unroll
285
+ for (int k_idx = 0; k_idx < k; ++k_idx) {
286
+ const int expert = expert_indices_for_token[k_idx];
287
+ const int idx = k * thread_row + k_idx;
288
+ for (int ii = 0; ii < VPT; ++ii) {
289
+ const int group_id = ii / ELTS_PER_LDG;
290
+ const int local_id = ii % ELTS_PER_LDG;
291
+ const int expert_idx = first_elt_read_by_thread +
292
+ group_id * THREADS_PER_ROW * ELTS_PER_LDG +
293
+ local_id;
294
+ if (expert == expert_idx) {
295
+ output[idx] = row_chunk[ii] * scale;
296
+ break;
297
+ }
298
+ }
299
+ }
300
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
301
+ asm volatile("griddepcontrol.launch_dependents;");
302
+#endif
303
+ return;
304
+ }
305
+
306
+#pragma unroll
307
+ for (int ii = 0; ii < VPT; ++ii) {
308
+ float val = row_chunk[ii];
309
+ float val_b = val * beta;
310
+ // Compute softplus: log(1 + exp(val)) with numerical stability
311
+ // When val > threshold, softplus(x) ≈ x to avoid exp overflow
312
+ val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
313
+ val = sqrtf(val);
314
+ if (correction_bias) {
315
+ const int group_id = ii / ELTS_PER_LDG;
316
+ const int local_id = ii % ELTS_PER_LDG;
317
+ const int expert_idx = first_elt_read_by_thread +
318
+ group_id * THREADS_PER_ROW * ELTS_PER_LDG +
319
+ local_id;
320
+ val = val + correction_bias[expert_idx];
321
+ }
322
+ row_chunk[ii] = val;
323
+ }
324
+
325
+ // Original TopK path: find top-k experts by score
326
+ // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find
327
+ // the topk elements in each row, along with the max index.
328
+ int start_col = first_elt_read_by_thread;
329
+ static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
330
+
331
+ float selected_sum = 0.f;
332
+ for (int k_idx = 0; k_idx < k; ++k_idx) {
333
+ // First, each thread does the local argmax
334
+ float max_val = row_chunk[0];
335
+ int expert = start_col;
336
+#pragma unroll
337
+ for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
338
+ ++ldg, col += COLS_PER_GROUP_LDG) {
339
+#pragma unroll
340
+ for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
341
+ float val = row_chunk[ldg * ELTS_PER_LDG + ii];
342
+
343
+ // No check on the experts here since columns with the smallest index
344
+ // are processed first and only updated if > (not >=)
345
+ if (val > max_val) {
346
+ max_val = val;
347
+ expert = col + ii;
348
+ }
349
+ }
350
+ }
351
+
352
+// Now, we perform the argmax reduce. We use the butterfly pattern so threads
353
+// reach consensus about the max. This will be useful for K > 1 so that the
354
+// threads can agree on "who" had the max value. That thread can then blank out
355
+// their max with -inf and the warp can run more iterations...
356
+#pragma unroll
357
+ for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
358
+ float other_max =
359
+ VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
360
+ int other_expert =
361
+ VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
362
+
363
+ // We want lower indices to "win" in every thread so we break ties this
364
+ // way
365
+ if (other_max > max_val ||
366
+ (other_max == max_val && other_expert < expert)) {
367
+ max_val = other_max;
368
+ expert = other_expert;
369
+ }
370
+ }
371
+
372
+ // Write the max for this k iteration to global memory.
373
+ if (thread_group_idx == 0) {
374
+ // Add a guard to ignore experts not included by this node
375
+ const bool node_uses_expert =
376
+ expert >= start_expert && expert < end_expert;
377
+ const bool should_process_row = row_is_active && node_uses_expert;
378
+
379
+ // The lead thread from each sub-group will write out the final results to
380
+ // global memory. (This will be a single) thread per row of the
381
+ // input/output matrices.
382
+ const int idx = k * thread_row + k_idx;
383
+ if (correction_bias != nullptr) {
384
+ max_val -= correction_bias[expert];
385
+ }
386
+ output[idx] = max_val;
387
+ indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
388
+ source_rows[idx] = k_idx * num_rows + thread_row;
389
+ if (renormalize) {
390
+ selected_sum += max_val;
391
+ }
392
+ }
393
+
394
+ // Finally, we clear the value in the thread with the current max if there
395
+ // is another iteration to run.
396
+ if (k_idx + 1 < k) {
397
+ const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
398
+ const int thread_to_clear_in_group =
399
+ (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
400
+
401
+ // Only the thread in the group which produced the max will reset the
402
+ // "winning" value to -inf.
403
+ if (thread_group_idx == thread_to_clear_in_group) {
404
+ const int offset_for_expert = expert % ELTS_PER_LDG;
405
+ // Safe to set to any negative value since row_chunk values must be
406
+ // between 0 and 1.
407
+ row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
408
+ -10000.f;
409
+ }
410
+ }
411
+ }
412
+
413
+ // Apply renormalization and routed scaling factor to final weights.
414
+ if (thread_group_idx == 0) {
415
+ float scale = static_cast<float>(routed_scaling_factor);
416
+ if (renormalize) {
417
+ const float denom = selected_sum > 0.f ? selected_sum : 1.f;
418
+ scale /= denom;
419
+ }
420
+ for (int k_idx = 0; k_idx < k; ++k_idx) {
421
+ const int idx = k * thread_row + k_idx;
422
+ output[idx] = output[idx] * scale;
423
+ }
424
+ }
425
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
426
+ asm volatile("griddepcontrol.launch_dependents;");
427
+#endif
428
+}
429
+
430
+namespace detail {
431
+// Constructs some constants needed to partition the work across threads at
432
+// compile time.
433
+template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM,
434
+ typename InputType>
435
+struct TopkConstants {
436
+ static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
437
+ static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 ||
438
+ EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0,
439
+ "");
440
+ static constexpr int VECs_PER_THREAD =
441
+ MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
442
+ static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
443
+ static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
444
+ static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
445
+};
446
+} // namespace detail
447
+
448
+#define DISPATCH_HASH(use_hash, USE_HASH, ...) \
449
+ if (use_hash) { \
450
+ const bool USE_HASH = true; \
451
+ static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \
452
+ __VA_ARGS__ \
453
+ } else { \
454
+ const bool USE_HASH = false; \
455
+ static_assert(USE_HASH == false, \
456
+ "USE_HASH must be compile-time constant"); \
457
+ __VA_ARGS__ \
458
+ }
459
+
460
+template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
461
+ int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
462
+void topkGatingSoftplusSqrtLauncherHelper(
463
+ const InputType* input, const bool* finished, float* output,
464
+ IndType* indices, int* source_row, const int num_rows, const int k,
465
+ const int start_expert, const int end_expert, const bool renormalize,
466
+ double routed_scaling_factor, const float* correction_bias,
467
+ const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
468
+ cudaStream_t stream) {
469
+ static constexpr int BYTES_PER_LDG =
470
+ MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
471
+ using Constants =
472
+ detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
473
+ static constexpr int VPT = Constants::VPT;
474
+ static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
475
+ const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
476
+ const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
477
+ dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
478
+ DISPATCH_HASH(use_hash, USE_HASH, {
479
+ auto* kernel =
480
+ &topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
481
+ WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
482
+#ifndef USE_ROCM
483
+ cudaLaunchConfig_t config = {};
484
+ config.gridDim = num_blocks;
485
+ config.blockDim = block_dim;
486
+ config.dynamicSmemBytes = 0;
487
+ config.stream = stream;
488
+ cudaLaunchAttribute attrs[1];
489
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
490
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
491
+ config.numAttrs = 1;
492
+ config.attrs = attrs;
493
+ cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
494
+ indices, source_row, k, start_expert, end_expert,
495
+ renormalize, routed_scaling_factor, correction_bias,
496
+ input_ids, tid2eid);
497
+#else
498
+ kernel<<<num_blocks, block_dim, 0, stream>>>(
499
+ input, finished, output, num_rows, indices, source_row, k, start_expert,
500
+ end_expert, renormalize, routed_scaling_factor, correction_bias,
501
+ input_ids, tid2eid);
502
+#endif
503
+ })
504
+}
505
+
506
+#ifndef USE_ROCM
507
+ #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
508
+ static_assert(WARP_SIZE == 32, \
509
+ "Unsupported warp size. Only 32 is supported for CUDA"); \
510
+ topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, \
511
+ MAX_BYTES>( \
512
+ gating_output, nullptr, topk_weights, topk_indices, \
513
+ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
514
+ routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
515
+ stream);
516
+#else
517
+ #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
518
+ if (WARP_SIZE == 64) { \
519
+ topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, \
520
+ MAX_BYTES>( \
521
+ gating_output, nullptr, topk_weights, topk_indices, \
522
+ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
523
+ routed_scaling_factor, correction_bias, use_hash, input_ids, \
524
+ tid2eid, stream); \
525
+ } else if (WARP_SIZE == 32) { \
526
+ topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
527
+ MAX_BYTES>( \
528
+ gating_output, nullptr, topk_weights, topk_indices, \
529
+ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
530
+ routed_scaling_factor, correction_bias, use_hash, input_ids, \
531
+ tid2eid, stream); \
532
+ } else { \
533
+ assert(false && \
534
+ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
535
+ }
536
+#endif
537
+
538
+template <typename IndType, typename InputType>
539
+void topkGatingSoftplusSqrtKernelLauncher(
540
+ const InputType* gating_output, float* topk_weights, IndType* topk_indices,
541
+ int* token_expert_indices, const int num_tokens, const int num_experts,
542
+ const int topk, const bool renormalize, double routed_scaling_factor,
543
+ const float* correction_bias, const bool use_hash, const IndType* input_ids,
544
+ const IndType* tid2eid, cudaStream_t stream) {
545
+ static constexpr int WARPS_PER_TB = 4;
546
+ static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
547
+#ifndef USE_ROCM
548
+ // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
549
+ // elements can be loaded by a warp
550
+ static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
551
+ (std::is_same_v<InputType, __nv_bfloat16> ||
552
+ std::is_same_v<InputType, __half>)
553
+ ? 4
554
+ : 8;
555
+#endif
556
+ switch (num_experts) {
557
+ case 1:
558
+ LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
559
+ break;
560
+ case 2:
561
+ LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
562
+ break;
563
+ case 4:
564
+ LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
565
+ break;
566
+ case 8:
567
+ LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
568
+ break;
569
+ case 16:
570
+ LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
571
+ break;
572
+ case 32:
573
+ LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
574
+ break;
575
+ case 64:
576
+ LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
577
+ break;
578
+ case 128:
579
+ LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
580
+ break;
581
+ case 256:
582
+ LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
583
+ break;
584
+ case 512:
585
+ LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
586
+ break;
587
+ // (CUDA only) support multiples of 64 when num_experts is not power of 2.
588
+ // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
589
+ // num_experts, alternatively we can test 4 bytes loading and enable it in
590
+ // future.
591
+#ifndef USE_ROCM
592
+ case 192:
593
+ LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
594
+ break;
595
+ case 320:
596
+ LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
597
+ break;
598
+ case 384:
599
+ LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
600
+ break;
601
+ case 448:
602
+ LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
603
+ break;
604
+ case 576:
605
+ LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
606
+ break;
607
+#endif
608
+ default: {
609
+ TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
610
+ }
611
+ }
612
+}
613
+
614
+} // namespace moe
615
+} // namespace vllm
616
+
617
+template <typename ComputeType>
618
+void dispatch_topk_softplus_sqrt_launch(
619
+ const ComputeType* gating_output, torch::Tensor& topk_weights,
620
+ torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
621
+ int num_tokens, int num_experts, int topk, bool renormalize,
622
+ double routed_scaling_factor,
623
+ const c10::optional<torch::Tensor>& correction_bias,
624
+ const c10::optional<torch::Tensor>& input_ids,
625
+ const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
626
+ const float* bias_ptr = nullptr;
627
+ if (correction_bias.has_value()) {
628
+ bias_ptr = correction_bias.value().data_ptr<float>();
629
+ }
630
+ bool use_hash = false;
631
+ if (tid2eid.has_value()) {
632
+ TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
633
+ use_hash = true;
634
+ }
635
+ if (topk_indices.scalar_type() == at::ScalarType::Int) {
636
+ const int* input_ids_ptr = nullptr;
637
+ const int* tid2eid_ptr = nullptr;
638
+ if (tid2eid.has_value()) {
639
+ input_ids_ptr = input_ids.value().data_ptr<int>();
640
+ tid2eid_ptr = tid2eid.value().data_ptr<int>();
641
+ }
642
+
643
+ vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
644
+ gating_output, topk_weights.data_ptr<float>(),
645
+ topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
646
+ num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
647
+ bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
648
+ } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
649
+ const uint32_t* input_ids_ptr = nullptr;
650
+ const uint32_t* tid2eid_ptr = nullptr;
651
+ if (tid2eid.has_value()) {
652
+ input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
653
+ tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
654
+ }
655
+ vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
656
+ gating_output, topk_weights.data_ptr<float>(),
657
+ topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
658
+ num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
659
+ bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
660
+ } else {
661
+ TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
662
+
663
+ const int64_t* input_ids_ptr = nullptr;
664
+ const int64_t* tid2eid_ptr = nullptr;
665
+ if (tid2eid.has_value()) {
666
+ input_ids_ptr = input_ids.value().data_ptr<int64_t>();
667
+ tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
668
+ }
669
+
670
+ vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
671
+ gating_output, topk_weights.data_ptr<float>(),
672
+ topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
673
+ num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
674
+ bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
675
+ }
676
+}
677
+
678
+void topk_softplus_sqrt(
679
+ torch::Tensor& topk_weights, // [num_tokens, topk]
680
+ torch::Tensor& topk_indices, // [num_tokens, topk]
681
+ torch::Tensor& token_expert_indices, // [num_tokens, topk]
682
+ torch::Tensor& gating_output, // [num_tokens, num_experts]
683
+ bool renormalize, double routed_scaling_factor,
684
+ const c10::optional<torch::Tensor>& correction_bias,
685
+ const c10::optional<torch::Tensor>& input_ids,
686
+ const c10::optional<torch::Tensor>& tid2eid) {
687
+ const int num_experts = gating_output.size(-1);
688
+ const auto num_tokens = gating_output.numel() / num_experts;
689
+ const int topk = topk_weights.size(-1);
690
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
691
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
692
+
693
+ if (gating_output.scalar_type() == at::ScalarType::Float) {
694
+ dispatch_topk_softplus_sqrt_launch<float>(
695
+ gating_output.data_ptr<float>(), topk_weights, topk_indices,
696
+ token_expert_indices, num_tokens, num_experts, topk, renormalize,
697
+ routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
698
+ } else if (gating_output.scalar_type() == at::ScalarType::Half) {
699
+ dispatch_topk_softplus_sqrt_launch<__half>(
700
+ reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
701
+ topk_weights, topk_indices, token_expert_indices, num_tokens,
702
+ num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
703
+ input_ids, tid2eid, stream);
704
+ } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
705
+ dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
706
+ reinterpret_cast<const __nv_bfloat16*>(
707
+ gating_output.data_ptr<at::BFloat16>()),
708
+ topk_weights, topk_indices, token_expert_indices, num_tokens,
709
+ num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
710
+ input_ids, tid2eid, stream);
711
+ } else {
712
+ TORCH_CHECK(false, "Unsupported gating_output data type: ",
713
+ gating_output.scalar_type());
714
+ }
715
+}
\ No newline at end of file
csrc/moe/torch_bindings.cpp
+8
@@ -16,6 +16,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
16
"bias) -> ()");
17
m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
18
19
+#ifndef USE_ROCM
20
+ m.def(
21
+ "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
22
+ "token_expert_indices, Tensor gating_output, bool renormalize, float "
23
+ "routed_scaling_factor, Tensor? "
24
+ "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
25
+ m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
26
+#endif
27
// Calculate the result of moe by summing up the partial results
28
// from all selected experts.
29
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
csrc/ops.h
+7
-1
@@ -100,6 +100,11 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
100
bool is_neox, torch::Tensor& position_ids,
101
int64_t forced_token_heads_per_warp);
102
103
+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
104
+ torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
105
+ torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
106
+ torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
107
+
108
void apply_repetition_penalties_(torch::Tensor& logits,
109
const torch::Tensor& prompt_mask,
110
const torch::Tensor& output_mask,
@@ -153,7 +158,8 @@ void silu_and_mul_per_block_quant(torch::Tensor& out,
158
159
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
160
std::optional<torch::Tensor> key, int64_t head_size,
156
- torch::Tensor& cos_sin_cache, bool is_neox);
161
+ torch::Tensor& cos_sin_cache, bool is_neox,
162
+ int64_t rope_dim_offset, bool inverse);
163
164
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
165
csrc/persistent_topk.cuh
+17
-16
@@ -18,7 +18,6 @@ namespace persistent {
18
// Constants
19
// ============================================================================
20
21
-constexpr int TopK = 2048;
21
constexpr int kThreadsPerBlock = 1024;
22
constexpr int RADIX = 256;
23
@@ -128,11 +127,12 @@ struct RadixRowState {
127
128
struct PersistentTopKParams {
129
const float* __restrict__ input; // [num_rows, stride]
131
- int32_t* __restrict__ output; // [num_rows, TopK]
130
+ int32_t* __restrict__ output; // [num_rows, top_k]
131
int32_t* __restrict__ lengths; // [num_rows]
132
RadixRowState* row_states; // large path: per-group state
133
uint32_t num_rows;
134
uint32_t stride;
135
+ uint32_t top_k; // actual k value for output stride
136
uint32_t chunk_size; // large path: elements per CTA
137
uint32_t ctas_per_group; // 1=medium, >1=large
138
uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit)
@@ -154,6 +154,7 @@ __device__ __forceinline__ uint32_t decode_bin(float x) {
154
return key >> 5;
155
}
156
157
+template <int TopK>
158
__device__ __noinline__ void histogram_2048_topk(
159
const float* __restrict__ logits, int32_t* __restrict__ output_indices,
160
int32_t seq_len) {
@@ -418,6 +419,7 @@ __device__ __noinline__ void histogram_2048_topk(
419
// by: DarkSharpness
420
// which at the same time is an optimized topk kernel copied from tilelang
421
// kernel
422
+template <int TopK>
423
__device__ __noinline__ void histogram_256_topk(
424
const float* __restrict__ logits, int* __restrict__ output_indices,
425
int logits_offset, int seq_len) {
@@ -649,7 +651,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
651
// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
652
// ============================================================================
653
652
-template <uint32_t VEC_SIZE>
654
+template <int TopK, uint32_t VEC_SIZE>
655
__device__ void radix_topk(const float* __restrict__ row_input,
656
int32_t* __restrict__ row_output, uint32_t seq_len,
657
uint32_t my_chunk_start, uint32_t chunk_size,
@@ -857,7 +859,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
859
// see filtered_topk.cuh)
860
// ============================================================================
861
860
-template <uint32_t VEC_SIZE = 1>
862
+template <int TopK = 2048, uint32_t VEC_SIZE = 1>
863
__global__ void __launch_bounds__(kThreadsPerBlock, 2)
864
persistent_topk_kernel(PersistentTopKParams params) {
865
const uint32_t tx = threadIdx.x;
@@ -915,7 +917,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
917
if (row_idx >= params.num_rows) break;
918
919
const uint32_t seq_len = params.lengths[row_idx];
918
- int32_t* row_output = params.output + row_idx * TopK;
920
+ int32_t* row_output = params.output + row_idx * params.top_k;
921
const float* row_input = params.input + row_idx * params.stride;
922
923
if (seq_len <= RADIX_THRESHOLD) {
@@ -927,19 +929,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
929
row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
930
}
931
} else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
930
- histogram_2048_topk(row_input, row_output, seq_len);
932
+ histogram_2048_topk<TopK>(row_input, row_output, seq_len);
933
} else {
932
- histogram_256_topk(row_input, row_output, 0, seq_len);
934
+ histogram_256_topk<TopK>(row_input, row_output, 0, seq_len);
935
}
936
}
937
continue;
938
}
939
940
const uint32_t my_chunk_start = cta_in_group * chunk_size;
939
- radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
940
- chunk_size, local_histogram, suffix_sum,
941
- shared_scalars, shared_ordered, state, cta_in_group,
942
- ctas_per_group, barrier_phase, iter, tx);
941
+ radix_topk<TopK, VEC_SIZE>(
942
+ row_input, row_output, seq_len, my_chunk_start, chunk_size,
943
+ local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
944
+ cta_in_group, ctas_per_group, barrier_phase, iter, tx);
945
}
946
}
947
@@ -1011,7 +1013,6 @@ struct FilteredTopKTraits<float> {
1013
}
1014
};
1015
1014
-constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
1016
constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
1017
constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
1018
16 * 1024; // 16K indices per buffer
@@ -1025,7 +1026,7 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
1026
* \tparam IdType Index type (int32_t)
1027
* \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
1028
*/
1028
-template <typename DType, typename IdType, int VEC_SIZE>
1029
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
1030
__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1031
FilteredTopKUnifiedKernel(const DType* __restrict__ input,
1032
IdType* __restrict__ output,
@@ -1059,7 +1060,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1060
alignas(128) __shared__ int s_counter;
1061
alignas(128) __shared__ int s_threshold_bin_id;
1062
alignas(128) __shared__ int s_num_input[2];
1062
- alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
1063
+ alignas(128) __shared__ int s_indices[MAX_K];
1064
1065
auto& s_histogram = s_histogram_buf[0];
1066
@@ -1280,7 +1281,7 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
1281
return static_cast<int>(g);
1282
}
1283
1283
-template <typename DType, typename IdType>
1284
+template <typename DType, typename IdType, uint32_t MAX_K = 2048>
1285
cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
1286
IdType* lengths, uint32_t num_rows,
1287
uint32_t top_k_val, uint32_t max_len,
@@ -1297,7 +1298,7 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
1298
1299
#define DISPATCH_VEC_SIZE(VS) \
1300
if (vec_size == VS) { \
1300
- auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>; \
1301
+ auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>; \
1302
FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \
1303
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \
1304
FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
csrc/pos_encoding_kernels.cu
+39
-33
@@ -9,28 +9,29 @@ namespace vllm {
9
10
template <typename scalar_t, bool IS_NEOX>
11
inline __device__ void apply_token_rotary_embedding(
12
- scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
13
- const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
12
+ scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
13
+ const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
14
+ const bool inverse) {
15
int x_index, y_index;
15
- scalar_t cos, sin;
16
+ float cos_f, sin_f;
17
if (IS_NEOX) {
17
- // GPT-NeoX style rotary embedding.
18
x_index = rot_offset;
19
y_index = embed_dim + rot_offset;
20
- cos = VLLM_LDG(cos_ptr + x_index);
21
- sin = VLLM_LDG(sin_ptr + x_index);
20
+ cos_f = VLLM_LDG(cos_ptr + x_index);
21
+ sin_f = VLLM_LDG(sin_ptr + x_index);
22
} else {
23
- // GPT-J style rotary embedding.
23
x_index = 2 * rot_offset;
24
y_index = 2 * rot_offset + 1;
26
- cos = VLLM_LDG(cos_ptr + x_index / 2);
27
- sin = VLLM_LDG(sin_ptr + x_index / 2);
25
+ cos_f = VLLM_LDG(cos_ptr + x_index / 2);
26
+ sin_f = VLLM_LDG(sin_ptr + x_index / 2);
27
}
29
-
30
- const scalar_t x = arr[x_index];
31
- const scalar_t y = arr[y_index];
32
- arr[x_index] = x * cos - y * sin;
33
- arr[y_index] = y * cos + x * sin;
28
+ if (inverse) {
29
+ sin_f = -sin_f;
30
+ }
31
+ const float x_f = static_cast<float>(arr[x_index]);
32
+ const float y_f = static_cast<float>(arr[y_index]);
33
+ arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f);
34
+ arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
35
}
36
37
template <typename scalar_t, bool IS_NEOX>
@@ -42,22 +43,23 @@ inline __device__ void apply_rotary_embedding(
43
// [batch_size, seq_len, num_kv_heads,
44
// head_size] or [num_tokens, num_kv_heads,
45
// head_size]
45
- const scalar_t* cache_ptr, const int head_size, const int num_heads,
46
+ const float* cache_ptr, const int head_size, const int num_heads,
47
const int num_kv_heads, const int rot_dim, const int token_idx,
48
const int64_t query_stride, const int64_t key_stride,
48
- const int64_t head_stride) {
49
+ const int64_t head_stride, const int64_t rope_dim_offset,
50
+ const bool inverse) {
51
const int embed_dim = rot_dim / 2;
50
- const scalar_t* cos_ptr = cache_ptr;
51
- const scalar_t* sin_ptr = cache_ptr + embed_dim;
52
+ const float* cos_ptr = cache_ptr;
53
+ const float* sin_ptr = cache_ptr + embed_dim;
54
55
const int nq = num_heads * embed_dim;
56
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
57
const int head_idx = i / embed_dim;
58
const int64_t token_head =
57
- token_idx * query_stride + head_idx * head_stride;
59
+ token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
60
const int rot_offset = i % embed_dim;
61
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
60
- query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
62
+ query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
63
}
64
65
if (key != nullptr) {
@@ -65,10 +67,10 @@ inline __device__ void apply_rotary_embedding(
67
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
68
const int head_idx = i / embed_dim;
69
const int64_t token_head =
68
- token_idx * key_stride + head_idx * head_stride;
70
+ token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
71
const int rot_offset = i % embed_dim;
72
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
71
- key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
73
+ key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
74
}
75
}
76
}
@@ -84,19 +86,18 @@ __global__ void rotary_embedding_kernel(
86
// [batch_size, seq_len, num_kv_heads,
87
// head_size] or [num_tokens, num_kv_heads,
88
// head_size]
87
- const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
88
- // 2]
89
+ const float* __restrict__ cos_sin_cache, // [max_position, rot_dim] fp32
90
const int rot_dim, const int64_t query_stride, const int64_t key_stride,
91
const int64_t head_stride, const int num_heads, const int num_kv_heads,
91
- const int head_size) {
92
- // Each thread block is responsible for one token.
92
+ const int head_size, const int64_t rope_dim_offset, const bool inverse) {
93
const int token_idx = blockIdx.x;
94
int64_t pos = positions[token_idx];
95
- const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
95
+ const float* cache_ptr = cos_sin_cache + pos * rot_dim;
96
97
apply_rotary_embedding<scalar_t, IS_NEOX>(
98
query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
99
- token_idx, query_stride, key_stride, head_stride);
99
+ token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
100
+ inverse);
101
}
102
103
} // namespace vllm
@@ -115,7 +116,7 @@ void rotary_embedding(
116
// [num_tokens, num_heads, head_size]
117
int64_t head_size,
118
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
118
- bool is_neox) {
119
+ bool is_neox, int64_t rope_dim_offset, bool inverse) {
120
// num_tokens = batch_size * seq_len
121
int64_t num_tokens = positions.numel();
122
int positions_ndim = positions.dim();
@@ -154,6 +155,8 @@ void rotary_embedding(
155
int seq_dim_idx = positions_ndim - 1;
156
int64_t query_stride = query.stride(seq_dim_idx);
157
int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
158
+
159
+ TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
160
// Determine head stride: for [*, heads, head_size] use stride of last dim;
161
// for flat [*, heads*head_size], heads blocks are contiguous of size
162
// head_size
@@ -165,20 +168,23 @@ void rotary_embedding(
168
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
169
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
170
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
171
+ auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
172
VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
173
if (is_neox) {
174
vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
175
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
176
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
173
- cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
174
- head_stride, num_heads, num_kv_heads, head_size);
177
+ cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
178
+ head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
179
+ inverse);
180
} else {
181
vllm::rotary_embedding_kernel<scalar_t, false>
182
<<<grid, block, 0, stream>>>(
183
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
184
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
180
- cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
181
- key_stride, head_stride, num_heads, num_kv_heads, head_size);
185
+ cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
186
+ head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
187
+ inverse);
188
}
189
});
190
}
csrc/sampler.cu
+7
-1
@@ -258,7 +258,13 @@ __device__ bool processHistogramStep(
258
auto processBins = [&](float logit, int idx) {
259
if (isPartialMatch<patternShift>(logit, logitPattern)) {
260
uint32_t binIdx = extractBinIdx<step>(logit);
261
- if (binIdx < thresholdBinIdx) {
261
+ // Only write elements with binIdx < thresholdBinIdx when:
262
+ // 1. This is step 0 and the threshold bin is small enough (no step 1)
263
+ // 2. This is step >= 1 (where pattern matching filters correctly)
264
+ // This prevents duplicates when step 0 and step 1 both run.
265
+ bool shouldWriteDirectly =
266
+ (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
267
+ if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
268
// The element is part of the top-k selection
269
int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
270
csrc/topk.cu
+59
-35
@@ -10,33 +10,17 @@
10
#include "persistent_topk.cuh"
11
#endif
12
13
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
14
- torch::Tensor& output, torch::Tensor& workspace, int64_t k,
15
- int64_t max_seq_len) {
13
+namespace {
14
+
15
#ifndef USE_ROCM
17
- TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
18
- TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
19
- TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
20
- TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
21
- TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
22
- TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
23
- TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
24
- TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
25
- "lengths must be 1D or 2D");
26
- TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
27
- TORCH_CHECK(output.dim() == 2, "output must be 2D");
16
+template <int TopK>
17
+void launch_persistent_topk(const torch::Tensor& logits,
18
+ const torch::Tensor& lengths, torch::Tensor& output,
19
+ torch::Tensor& workspace, int64_t max_seq_len) {
20
+ namespace P = vllm::persistent;
21
22
const int64_t num_rows = logits.size(0);
23
const int64_t stride = logits.size(1);
31
-
32
- TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
33
- TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
34
- "output size mismatch");
35
- namespace P = vllm::persistent;
36
-
37
- TORCH_CHECK(k == P::TopK, "k must be 2048");
38
- TORCH_CHECK(k <= stride, "k out of range");
39
-
24
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
25
26
static int num_sms = 0;
@@ -50,18 +34,17 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
34
}
35
36
if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
53
- cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
54
- logits.data_ptr<float>(), output.data_ptr<int32_t>(),
55
- lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
56
- static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
37
+ cudaError_t status =
38
+ vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
39
+ logits.data_ptr<float>(), output.data_ptr<int32_t>(),
40
+ lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
41
+ static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
42
TORCH_CHECK(status == cudaSuccess,
43
"FilteredTopK failed: ", cudaGetErrorString(status));
44
} else {
45
TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
46
TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
47
63
- // Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
64
- // large path. Empirically tuned.
48
int effective_max_smem;
49
if (num_rows <= 4) {
50
effective_max_smem =
@@ -101,7 +84,7 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
84
85
int occupancy = 1;
86
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
104
- &occupancy, P::persistent_topk_kernel<4>, P::kThreadsPerBlock,
87
+ &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
88
smem_size);
89
if (occupancy < 1) occupancy = 1;
90
@@ -121,15 +104,16 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
104
params.lengths = lengths.data_ptr<int32_t>();
105
params.num_rows = static_cast<uint32_t>(num_rows);
106
params.stride = static_cast<uint32_t>(stride);
107
+ params.top_k = static_cast<uint32_t>(TopK);
108
params.chunk_size = chunk_size;
109
params.row_states =
110
reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
111
params.ctas_per_group = ctas_per_group;
112
params.max_seq_len = static_cast<uint32_t>(max_seq_len);
113
130
- #define LAUNCH_PERSISTENT(VS) \
114
+ #define LAUNCH_PERSISTENT(TOPK_VAL, VS) \
115
do { \
132
- auto kernel = &P::persistent_topk_kernel<VS>; \
116
+ auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>; \
117
cudaError_t err = cudaFuncSetAttribute( \
118
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \
119
TORCH_CHECK(err == cudaSuccess, \
@@ -138,11 +122,11 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
122
} while (0)
123
124
if (vec_size == 4) {
141
- LAUNCH_PERSISTENT(4);
125
+ LAUNCH_PERSISTENT(TopK, 4);
126
} else if (vec_size == 2) {
143
- LAUNCH_PERSISTENT(2);
127
+ LAUNCH_PERSISTENT(TopK, 2);
128
} else {
145
- LAUNCH_PERSISTENT(1);
129
+ LAUNCH_PERSISTENT(TopK, 1);
130
}
131
#undef LAUNCH_PERSISTENT
132
}
@@ -150,6 +134,46 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
134
cudaError_t err = cudaGetLastError();
135
TORCH_CHECK(err == cudaSuccess,
136
"persistent_topk failed: ", cudaGetErrorString(err));
137
+}
138
+#endif
139
+
140
+} // anonymous namespace
141
+
142
+void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
143
+ torch::Tensor& output, torch::Tensor& workspace, int64_t k,
144
+ int64_t max_seq_len) {
145
+#ifndef USE_ROCM
146
+ TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
147
+ TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
148
+ TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
149
+ TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
150
+ TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
151
+ TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
152
+ TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
153
+ TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
154
+ "lengths must be 1D or 2D");
155
+ TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
156
+ TORCH_CHECK(output.dim() == 2, "output must be 2D");
157
+
158
+ const int64_t num_rows = logits.size(0);
159
+ const int64_t stride = logits.size(1);
160
+
161
+ TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
162
+ TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
163
+ "output size mismatch");
164
+ TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
165
+ "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
166
+
167
+ if (k == 512) {
168
+ launch_persistent_topk<512>(logits, lengths, output, workspace,
169
+ max_seq_len);
170
+ } else if (k == 1024) {
171
+ launch_persistent_topk<1024>(logits, lengths, output, workspace,
172
+ max_seq_len);
173
+ } else {
174
+ launch_persistent_topk<2048>(logits, lengths, output, workspace,
175
+ max_seq_len);
176
+ }
177
#else
178
TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
179
#endif
csrc/torch_bindings.cpp
+15
-1
@@ -177,6 +177,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
177
"int forced_token_heads_per_warp=-1) -> ()");
178
ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
179
180
+#ifndef USE_ROCM
181
+ // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
182
+ // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
183
+ // kernel launch.
184
+ ops.def(
185
+ "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
186
+ "Tensor! q, Tensor kv, Tensor! k_cache, "
187
+ "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
188
+ "float eps, int cache_block_size) -> ()");
189
+ ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
190
+ &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
191
+#endif
192
+
193
// Apply repetition penalties to logits in-place
194
ops.def(
195
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -240,7 +253,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
253
ops.def(
254
"rotary_embedding(Tensor positions, Tensor! query,"
255
" Tensor!? key, int head_size,"
243
- " Tensor cos_sin_cache, bool is_neox) -> ()");
256
+ " Tensor cos_sin_cache, bool is_neox, int "
257
+ "rope_dim_offset=0, bool inverse=False) -> ()");
258
ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
259
260
// Quantization ops
docs/design/attention_backends.md
+1
-1
@@ -213,7 +213,7 @@ configuration.
213
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
214
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
215
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
216
-| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
216
+| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
217
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
218
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
219
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
docs/models/supported_models.md
+4
-3
@@ -384,6 +384,7 @@ th {
384
| `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
385
| `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
386
| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
387
+| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | |
388
| `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
389
| `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
390
| `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
@@ -643,10 +644,10 @@ Some models are supported only via the [Transformers modeling backend](#transfor
644
!!! note
645
`Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
646
MobileNet-v5 vision backbone.
646
-
647
+
648
Performance is not yet fully optimized mainly due to:
648
-
649
- - Both audio and vision MM encoders use `transformers.AutoModel` implementation.
649
+
650
+ - Both audio and vision MM encoders use `transformers.AutoModel` implementation.
651
- There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
652
653
!!! note
requirements/cuda.txt
+2
@@ -11,6 +11,8 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor
11
# FlashInfer should be updated together with the Dockerfile
12
flashinfer-python==0.6.8.post1
13
flashinfer-cubin==0.6.8.post1
14
+apache-tvm-ffi==0.1.9
15
+tilelang==0.1.9
16
# Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
17
# breaking changes in 1.19.0
18
nvidia-cudnn-frontend>=1.13.0,<1.19.0
tests/compile/fusions_e2e/conftest.py
+5
@@ -116,6 +116,11 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
116
model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
117
model_kwargs["tensor_parallel_size"] = tp_size
118
119
+ # Cap warmup memory: tests use small max_model_len (1024) but the
120
+ # engine default max_num_batched_tokens is 16384. Warming up large
121
+ # models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM.
122
+ model_kwargs.setdefault("max_num_batched_tokens", 8192)
123
+
124
# Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in
125
# decompose_auto_functionalized when +rotary_embedding is forced into
126
# the compile graph. Disable qk_norm+rope fusion (which auto-enables
tests/kernels/attention/test_deepgemm_attention.py
+22
-12
@@ -9,8 +9,8 @@ from vllm.platforms import current_platform
9
from vllm.utils.deep_gemm import (
10
_ceil_to_ue8m0,
11
calc_diff,
12
- fp8_mqa_logits,
13
- fp8_paged_mqa_logits,
12
+ fp8_fp4_mqa_logits,
13
+ fp8_fp4_paged_mqa_logits,
14
get_num_sms,
15
get_paged_mqa_logits_metadata,
16
)
@@ -127,8 +127,8 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
127
128
q_fp8 = q.to(torch.float8_e4m3fn)
129
kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False)
130
- logits = fp8_mqa_logits(
131
- q_fp8, kv_fp8, weights, ks, ke, clean_logits=clean_logits
130
+ logits = fp8_fp4_mqa_logits(
131
+ (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits
132
)
133
134
ref_logits = _ref_fp8_mqa_logits(
@@ -150,7 +150,7 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
150
assert diff < 1e-3, f"{diff=}"
151
152
153
-def _ref_fp8_paged_mqa_logits(
153
+def _ref_fp8_fp4_paged_mqa_logits(
154
q: torch.Tensor,
155
kv_cache: torch.Tensor,
156
weights: torch.Tensor,
@@ -205,8 +205,10 @@ def _ref_fp8_paged_mqa_logits(
205
@pytest.mark.skipif(
206
not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
207
)
208
-@pytest.mark.parametrize("clean_logits", [True, False])
209
-def test_deepgemm_fp8_paged_mqa_logits(clean_logits: bool):
208
+def test_deepgemm_fp8_fp4_paged_mqa_logits():
209
+ # NOTE: clean_logits=True is incompatible with the 2D context_lens
210
+ # required by csrc/apis/attention.hpp; only the False path is exercised.
211
+ clean_logits = False
212
torch.manual_seed(0)
213
random.seed(0)
214
@@ -258,21 +260,29 @@ def test_deepgemm_fp8_paged_mqa_logits(clean_logits: bool):
260
q_fp8 = q.to(torch.float8_e4m3fn)
261
kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
262
263
+ # deep_gemm paged MQA logits requires 2D context_lens of
264
+ # shape (B, next_n) (csrc/apis/attention.hpp:332-335);
265
+ # see indexer.py:607-608. For each batch/next_n token, the
266
+ # effective context length is context_lens[b] - next_n + j + 1.
267
+ next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32)
268
+ context_lens_2d = (
269
+ context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange
270
+ ).contiguous()
271
schedule_metadata = get_paged_mqa_logits_metadata(
262
- context_lens, blocksize, get_num_sms()
272
+ context_lens_2d, blocksize, get_num_sms()
273
)
264
- logits = fp8_paged_mqa_logits(
265
- q_fp8,
274
+ logits = fp8_fp4_paged_mqa_logits(
275
+ (q_fp8, None),
276
kv_cache_fp8,
277
weights,
268
- context_lens,
278
+ context_lens_2d,
279
block_tables,
280
schedule_metadata,
281
max_model_len,
282
clean_logits=clean_logits,
283
)
284
275
- ref_logits = _ref_fp8_paged_mqa_logits(
285
+ ref_logits = _ref_fp8_fp4_paged_mqa_logits(
286
q,
287
kv_cache,
288
weights,
tests/kernels/core/test_fused_q_kv_rmsnorm.py
new
+81
@@ -0,0 +1,81 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Correctness + large-token-count launch tests for fused_q_kv_rmsnorm.
4
+
5
+Before the grid-dim fix the kernel used grid ``(2, num_tokens)``, which hit
6
+CUDA's 65535 grid-y cap for ``num_tokens >= 65536`` and failed with
7
+``Triton Error [CUDA]: invalid argument`` at every large chunked-prefill
8
+profile run. These tests pin the new grid layout.
9
+"""
10
+
11
+from __future__ import annotations
12
+
13
+import pytest
14
+import torch
15
+
16
+from vllm.platforms import current_platform
17
+from vllm.v1.attention.ops.deepseek_v4_ops import fused_q_kv_rmsnorm
18
+
19
+pytestmark = pytest.mark.skipif(
20
+ not current_platform.is_cuda_alike(),
21
+ reason="fused_q_kv_rmsnorm requires a CUDA/ROCm device",
22
+)
23
+
24
+
25
+def _ref_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
26
+ x_f32 = x.to(torch.float32)
27
+ variance = x_f32.pow(2).mean(dim=-1, keepdim=True)
28
+ y = x_f32 * torch.rsqrt(variance + eps) * w.to(torch.float32)
29
+ return y.to(x.dtype)
30
+
31
+
32
+@pytest.mark.parametrize("num_tokens", [1, 17, 1024, 8192])
33
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
34
+def test_fused_q_kv_rmsnorm_correctness(num_tokens: int, dtype: torch.dtype):
35
+ torch.manual_seed(0)
36
+ device = "cuda"
37
+ q_size, kv_size = 192, 576
38
+ qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
39
+ kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
40
+ qw = torch.randn(q_size, dtype=dtype, device=device)
41
+ kvw = torch.randn(kv_size, dtype=dtype, device=device)
42
+ eps = 1e-6
43
+
44
+ qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
45
+
46
+ qr_ref = _ref_rmsnorm(qr, qw, eps)
47
+ kv_ref = _ref_rmsnorm(kv, kvw, eps)
48
+
49
+ tol = dict(rtol=1e-2, atol=1e-2)
50
+ torch.testing.assert_close(qr_out, qr_ref, **tol)
51
+ torch.testing.assert_close(kv_out, kv_ref, **tol)
52
+
53
+
54
+@pytest.mark.parametrize("num_tokens", [65535, 65536, 131072])
55
+def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
56
+ """Regression guard: grid used to be (2, num_tokens), hitting CUDA's
57
+ 65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2)
58
+ lifts that bound to 2**31-1."""
59
+ device = "cuda"
60
+ dtype = torch.bfloat16
61
+ q_size, kv_size = 192, 576
62
+ qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
63
+ kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
64
+ qw = torch.randn(q_size, dtype=dtype, device=device)
65
+ kvw = torch.randn(kv_size, dtype=dtype, device=device)
66
+
67
+ qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-6)
68
+ # spot-check a couple of rows against the torch reference
69
+ for row in (0, num_tokens // 2, num_tokens - 1):
70
+ torch.testing.assert_close(
71
+ qr_out[row],
72
+ _ref_rmsnorm(qr[row : row + 1], qw, 1e-6)[0],
73
+ rtol=1e-2,
74
+ atol=1e-2,
75
+ )
76
+ torch.testing.assert_close(
77
+ kv_out[row],
78
+ _ref_rmsnorm(kv[row : row + 1], kvw, 1e-6)[0],
79
+ rtol=1e-2,
80
+ atol=1e-2,
81
+ )
tests/kernels/moe/test_deepgemm.py
+195
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""
4
-Unit-test DeepGEMM FP8 kernels (no DeepEP).
4
+Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP).
5
Compare DeepGEMM path against the Triton fallback inside vLLM's fused_experts.
6
"""
7
@@ -21,6 +21,8 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
21
maybe_make_prepare_finalize,
22
)
23
from vllm.model_executor.layers.fused_moe.config import (
24
+ FusedMoEQuantConfig,
25
+ FusedMoEQuantDesc,
26
fp8_w8a8_moe_quant_config,
27
)
28
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
@@ -204,3 +206,195 @@ def test_deepgemm_vs_triton(m, n, k, topk, num_experts, monkeypatch, workspace_i
206
f"DeepGEMM path was not executed during the test. "
207
f"Call counter: {call_counter['cnt']}"
208
)
209
+
210
+
211
+# ---------------------------------------------------------------------------
212
+# FP4 weight tests (DeepGEMM m_grouped_fp8_fp4_gemm_nt_contiguous)
213
+# ---------------------------------------------------------------------------
214
+
215
+
216
+def make_mxfp4_weights(
217
+ e: int,
218
+ n: int,
219
+ k: int,
220
+):
221
+ """
222
+ Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales,
223
+ plus BF16 reference weights for validation.
224
+
225
+ w1 shape: (E, 2N, K//2) uint8 — packed FP4
226
+ w2 shape: (E, K, N//2) uint8 — packed FP4
227
+ w1_s shape: (E, 2N, K//32) float32 — per-row block-32 scales
228
+ w2_s shape: (E, K, N//32) float32 — per-row block-32 scales
229
+ w1_bf16: (E, 2N, K) — original BF16 for reference
230
+ w2_bf16: (E, K, N) — original BF16 for reference
231
+ """
232
+ from deep_gemm.utils.math import per_token_cast_to_fp4
233
+
234
+ dtype = torch.bfloat16
235
+ gran_k = 32 # MXFP4 block size
236
+
237
+ # bf16 reference weights — scale by 1/sqrt(dim) for numerical stability
238
+ w1_bf16 = torch.randn(e, 2 * n, k, device="cuda", dtype=dtype) * (k**-0.5)
239
+ w2_bf16 = torch.randn(e, k, n, device="cuda", dtype=dtype) * (n**-0.5)
240
+
241
+ # Quantize per-expert to FP4
242
+ w1 = torch.empty(e, 2 * n, k // 2, device="cuda", dtype=torch.uint8)
243
+ w2 = torch.empty(e, k, n // 2, device="cuda", dtype=torch.uint8)
244
+ w1_s = torch.empty(
245
+ e, 2 * n, math.ceil(k / gran_k), device="cuda", dtype=torch.float32
246
+ )
247
+ w2_s = torch.empty(e, k, math.ceil(n / gran_k), device="cuda", dtype=torch.float32)
248
+
249
+ for i in range(e):
250
+ w1[i], w1_s[i] = per_token_cast_to_fp4(
251
+ w1_bf16[i].float(), use_ue8m0=True, gran_k=gran_k
252
+ )
253
+ w2[i], w2_s[i] = per_token_cast_to_fp4(
254
+ w2_bf16[i].float(), use_ue8m0=True, gran_k=gran_k
255
+ )
256
+
257
+ return w1, w2, w1_s, w2_s, w1_bf16, w2_bf16
258
+
259
+
260
+def _bf16_moe_reference(x, w1, w2, topk_weights, topk_ids):
261
+ """BF16 token-loop MoE reference for correctness testing."""
262
+ import torch.nn.functional as F
263
+
264
+ num_tokens, hidden_size = x.shape
265
+ intermediate = w1.shape[1] // 2
266
+ top_k = topk_ids.shape[1]
267
+
268
+ output = torch.zeros(num_tokens, hidden_size, dtype=torch.float32, device=x.device)
269
+ for t in range(num_tokens):
270
+ for kk in range(top_k):
271
+ e = topk_ids[t, kk].item()
272
+ w = topk_weights[t, kk].item()
273
+ fc1 = x[t : t + 1].float() @ w1[e].float().T
274
+ linear = fc1[:, :intermediate]
275
+ gate = fc1[:, intermediate:]
276
+ act = F.silu(gate) * linear
277
+ fc2 = act @ w2[e].float().T
278
+ output[t] += w * fc2[0]
279
+ return output.to(torch.bfloat16)
280
+
281
+
282
+def run_single_fp4_case(m, n, k, topk, num_experts):
283
+ """
284
+ Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert
285
+ DeepGEMM FP4 == BF16 reference within tolerance.
286
+ """
287
+ tokens_bf16 = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
288
+
289
+ # FP4 expert weight tensors + BF16 originals for reference
290
+ w1, w2, w1_s, w2_s, w1_bf16, w2_bf16 = make_mxfp4_weights(num_experts, n, k)
291
+
292
+ router_logits = torch.randn(m, num_experts, device="cuda", dtype=torch.float32)
293
+ topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1)
294
+ topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
295
+
296
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
297
+ GroupShape,
298
+ )
299
+ from vllm.platforms import current_platform
300
+
301
+ _fp8_dtype = current_platform.fp8_dtype()
302
+ _block_shape = GroupShape(128, 128)
303
+ quant_config = FusedMoEQuantConfig(
304
+ _a1=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
305
+ _a2=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
306
+ _w1=FusedMoEQuantDesc("mxfp4", None, w1_s, None, None, None),
307
+ _w2=FusedMoEQuantDesc("mxfp4", None, w2_s, None, None, None),
308
+ )
309
+ moe_config = make_dummy_moe_config()
310
+
311
+ from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
312
+ DeepGemmFP4Experts,
313
+ )
314
+
315
+ deep_gemm_fp4_experts = mk.FusedMoEKernel(
316
+ prepare_finalize=maybe_make_prepare_finalize(
317
+ moe=moe_config,
318
+ quant_config=quant_config,
319
+ allow_new_interface=True,
320
+ use_monolithic=False,
321
+ ),
322
+ fused_experts=DeepGemmFP4Experts(
323
+ moe_config=moe_config,
324
+ quant_config=quant_config,
325
+ ),
326
+ inplace=False,
327
+ )
328
+
329
+ # DeepGEMM FP4 path
330
+ out_deepgemm_fp4 = deep_gemm_fp4_experts.apply(
331
+ hidden_states=tokens_bf16,
332
+ w1=w1,
333
+ w2=w2,
334
+ topk_weights=topk_weights,
335
+ topk_ids=topk_ids,
336
+ global_num_experts=num_experts,
337
+ activation=MoEActivation.SILU,
338
+ apply_router_weight_on_input=False,
339
+ expert_map=None,
340
+ )
341
+
342
+ # BF16 reference using the same original weights
343
+ out_ref = _bf16_moe_reference(tokens_bf16, w1_bf16, w2_bf16, topk_weights, topk_ids)
344
+
345
+ # FP4 vs BF16 reference: quantization error from FP4 weights + FP8 activations
346
+ diff = calc_diff(out_deepgemm_fp4, out_ref)
347
+ assert diff < 0.05, f"FP4 diff exceeded 5%: {diff}"
348
+
349
+
350
+# DeepSeek V4 dims: H=4096, I=2048, so N=2*I=4096, K=H=4096.
351
+# FP4 quantization with block_k=32 needs large K for good accuracy.
352
+FP4_MNKs = [
353
+ (128, 4096, 4096), # DeepSeek V4 shape
354
+ (256, 2048, 2048), # Half-size variant
355
+]
356
+
357
+FP4_TOPKS = [2]
358
+FP4_NUM_EXPERTS = [8]
359
+
360
+
361
+@pytest.mark.parametrize(("m", "n", "k"), FP4_MNKs)
362
+@pytest.mark.parametrize("topk", FP4_TOPKS)
363
+@pytest.mark.parametrize("num_experts", FP4_NUM_EXPERTS)
364
+@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
365
+def test_deepgemm_fp4_vs_triton(
366
+ m, n, k, topk, num_experts, monkeypatch, workspace_init
367
+):
368
+ pytest.importorskip("deep_gemm.utils.math")
369
+ with monkeypatch.context() as mp:
370
+ mp.setenv("VLLM_USE_DEEP_GEMM", "1")
371
+
372
+ _DeepGemmFP4Experts = importlib.import_module(
373
+ "vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe"
374
+ ).DeepGemmFP4Experts
375
+
376
+ call_counter = {"cnt": 0}
377
+
378
+ orig_fn = _DeepGemmFP4Experts.apply
379
+
380
+ def _spy_apply(*args, **kwargs):
381
+ call_counter["cnt"] += 1
382
+ return orig_fn(*args, **kwargs)
383
+
384
+ monkeypatch.setattr(_DeepGemmFP4Experts, "apply", _spy_apply)
385
+ if topk > num_experts:
386
+ pytest.skip(f"topk={topk} > num_experts={num_experts}")
387
+
388
+ run_single_fp4_case(
389
+ m=m,
390
+ n=n,
391
+ k=k,
392
+ topk=topk,
393
+ num_experts=num_experts,
394
+ )
395
+
396
+ # ensure that the DeepGEMM FP4 path was indeed taken.
397
+ assert call_counter["cnt"] == 1, (
398
+ f"DeepGEMM FP4 path was not executed during the test. "
399
+ f"Call counter: {call_counter['cnt']}"
400
+ )
tests/kernels/moe/test_topk_softplus_sqrt.py
new
+186
@@ -0,0 +1,186 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import pytest
6
+import torch
7
+import torch.nn.functional as F
8
+
9
+from vllm.model_executor.layers.fused_moe.config import (
10
+ RoutingMethodType,
11
+ get_routing_method_type,
12
+)
13
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
14
+ fused_topk_bias,
15
+)
16
+from vllm.platforms import current_platform
17
+
18
+
19
+def _torch_topk_softplus_sqrt(
20
+ gating_output: torch.Tensor,
21
+ topk: int,
22
+ renormalize: bool,
23
+ routed_scaling_factor: float,
24
+ e_score_correction_bias: torch.Tensor | None = None,
25
+ input_ids: torch.Tensor | None = None,
26
+ hash_indices_table: torch.Tensor | None = None,
27
+):
28
+ scores = F.softplus(gating_output.float()).sqrt()
29
+ original_scores = scores
30
+ if e_score_correction_bias is not None:
31
+ scores_for_choice = scores + e_score_correction_bias.unsqueeze(0)
32
+ else:
33
+ scores_for_choice = scores
34
+
35
+ if hash_indices_table is not None:
36
+ assert input_ids is not None
37
+ topk_ids = hash_indices_table[input_ids.long()]
38
+ else:
39
+ topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=True)[1]
40
+
41
+ topk_weights = original_scores.gather(1, topk_ids.long())
42
+ if renormalize:
43
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
44
+ if routed_scaling_factor != 1.0:
45
+ topk_weights = topk_weights * routed_scaling_factor
46
+ return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
47
+
48
+
49
+def test_sqrtsoftplus_bias_uses_deepseek_v4_routing_method():
50
+ assert (
51
+ get_routing_method_type(
52
+ scoring_func="sqrtsoftplus",
53
+ top_k=8,
54
+ renormalize=True,
55
+ num_expert_group=None,
56
+ has_e_score_bias=True,
57
+ )
58
+ == RoutingMethodType.DeepseekV4
59
+ )
60
+ assert (
61
+ get_routing_method_type(
62
+ scoring_func="sqrtsoftplus",
63
+ top_k=8,
64
+ renormalize=False,
65
+ num_expert_group=None,
66
+ has_e_score_bias=True,
67
+ )
68
+ == RoutingMethodType.Unspecified
69
+ )
70
+
71
+
72
+@pytest.mark.skipif(
73
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
74
+)
75
+@pytest.mark.parametrize("num_tokens", [1, 33, 128])
76
+@pytest.mark.parametrize("hidden_size", [1024, 2048])
77
+@pytest.mark.parametrize("num_experts", [128, 256, 384, 512])
78
+@pytest.mark.parametrize("topk", [6, 8, 16])
79
+@pytest.mark.parametrize("renormalize", [True, False])
80
+@pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.5])
81
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
82
+def test_fused_topk_softplus_sqrt(
83
+ num_tokens: int,
84
+ hidden_size: int,
85
+ num_experts: int,
86
+ topk: int,
87
+ renormalize: bool,
88
+ routed_scaling_factor: float,
89
+ dtype: torch.dtype,
90
+):
91
+ torch.manual_seed(0)
92
+ hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
93
+ gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
94
+ e_score_correction_bias = torch.randn(
95
+ (num_experts,), dtype=torch.float32, device="cuda"
96
+ )
97
+
98
+ topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
99
+ gating_output=gating_output,
100
+ topk=topk,
101
+ renormalize=renormalize,
102
+ routed_scaling_factor=routed_scaling_factor,
103
+ e_score_correction_bias=e_score_correction_bias,
104
+ )
105
+
106
+ topk_weights, topk_ids = fused_topk_bias(
107
+ hidden_states=hidden_states,
108
+ gating_output=gating_output,
109
+ scoring_func="sqrtsoftplus",
110
+ e_score_correction_bias=e_score_correction_bias,
111
+ topk=topk,
112
+ renormalize=renormalize,
113
+ routed_scaling_factor=routed_scaling_factor,
114
+ )
115
+
116
+ # Different kernels may return the topk experts in different orders when
117
+ # scores tie; sort by expert id before comparing.
118
+ sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1)
119
+ sorted_ids, idx_ops = topk_ids.sort(dim=-1)
120
+ torch.testing.assert_close(sorted_ref_ids, sorted_ids, atol=0, rtol=0)
121
+
122
+ sorted_w_ref = topk_weights_ref.gather(1, idx_ref)
123
+ sorted_w = topk_weights.gather(1, idx_ops)
124
+ torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
125
+
126
+
127
+@pytest.mark.skipif(
128
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
129
+)
130
+@pytest.mark.parametrize("num_tokens", [1, 33, 128])
131
+@pytest.mark.parametrize("hidden_size", [1024, 2048])
132
+@pytest.mark.parametrize("num_experts", [256, 384, 512])
133
+@pytest.mark.parametrize("topk", [6, 8, 16])
134
+@pytest.mark.parametrize("renormalize", [True, False])
135
+@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5])
136
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
137
+def test_fused_topk_softplus_sqrt_hash(
138
+ num_tokens: int,
139
+ hidden_size: int,
140
+ num_experts: int,
141
+ topk: int,
142
+ renormalize: bool,
143
+ routed_scaling_factor: float,
144
+ dtype: torch.dtype,
145
+):
146
+ torch.manual_seed(0)
147
+ vocab_size = 1024
148
+ hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
149
+ gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
150
+ # Per-token fixed expert selection: for each vocab id pick `topk` distinct
151
+ # experts.
152
+ hash_indices_table = torch.stack(
153
+ [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
154
+ ).to(device="cuda", dtype=torch.int32)
155
+ input_ids = torch.randint(
156
+ 0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda"
157
+ )
158
+
159
+ topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
160
+ gating_output=gating_output,
161
+ topk=topk,
162
+ renormalize=renormalize,
163
+ routed_scaling_factor=routed_scaling_factor,
164
+ input_ids=input_ids,
165
+ hash_indices_table=hash_indices_table,
166
+ )
167
+
168
+ topk_weights, topk_ids = fused_topk_bias(
169
+ hidden_states=hidden_states,
170
+ gating_output=gating_output,
171
+ scoring_func="sqrtsoftplus",
172
+ e_score_correction_bias=None,
173
+ topk=topk,
174
+ renormalize=renormalize,
175
+ input_tokens=input_ids,
176
+ hash_indices_table=hash_indices_table,
177
+ routed_scaling_factor=routed_scaling_factor,
178
+ )
179
+
180
+ sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1)
181
+ sorted_ids, idx_ops = topk_ids.sort(dim=-1)
182
+ torch.testing.assert_close(sorted_ref_ids, sorted_ids, atol=0, rtol=0)
183
+
184
+ sorted_w_ref = topk_weights_ref.gather(1, idx_ref)
185
+ sorted_w = topk_weights.gather(1, idx_ops)
186
+ torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
tests/kernels/test_compressor_kv_cache.py
new
+311
@@ -0,0 +1,311 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
5
+
6
+Two paths tested:
7
+ A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
8
+ B) Indexer: head_dim=128 (all FP8), quant_block=128
9
+
10
+These serve as golden references for validating the future fused
11
+compressor+quant+cache kernel.
12
+"""
13
+
14
+import math
15
+
16
+import pytest
17
+import torch
18
+
19
+from vllm import _custom_ops as ops
20
+from vllm.v1.attention.ops.deepseek_v4_ops import (
21
+ dequantize_and_gather_k_cache,
22
+ quantize_and_insert_k_cache,
23
+)
24
+
25
+
26
+def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
27
+ """PyTorch reference for UE8M0 FP8 quantization (per-block, power-of-2 scale).
28
+
29
+ Returns (x_fp8, scales) where x_fp8 is float8_e4m3fn and scales are float32.
30
+ """
31
+ assert x.dim() == 1
32
+ n = x.numel()
33
+ n_blocks = math.ceil(n / block_size)
34
+ x_fp8 = torch.zeros(n, dtype=torch.float8_e4m3fn, device=x.device)
35
+ scales = torch.zeros(n_blocks, dtype=torch.float32, device=x.device)
36
+
37
+ for i in range(n_blocks):
38
+ start = i * block_size
39
+ end = min(start + block_size, n)
40
+ block = x[start:end].float()
41
+ amax = block.abs().max().clamp(min=1e-4)
42
+ raw_scale = amax / fp8_max
43
+ exponent = math.ceil(math.log2(raw_scale.item()))
44
+ scale = 2.0**exponent
45
+ scales[i] = scale
46
+ quantized = (block / scale).clamp(-fp8_max, fp8_max)
47
+ x_fp8[start:end] = quantized.to(torch.float8_e4m3fn)
48
+
49
+ return x_fp8, scales
50
+
51
+
52
+# ── Test A: DeepseekV4 Attention path ──────────────────────────────────────────────
53
+
54
+
55
+@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
56
+@pytest.mark.parametrize("block_size", [16, 64])
57
+def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size: int):
58
+ """compressed_kv → quantize_and_insert_k_cache → dequantize_and_gather_k_cache
59
+ → compare against original."""
60
+
61
+ HEAD_DIM = 512
62
+ NOPE_DIM = 448
63
+ HEAD_BYTES = 584 # 448 fp8 + 128 bf16 + 8 uint8 scale
64
+ FP8_MAX = 448.0
65
+ QUANT_BLOCK = 64
66
+
67
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
68
+ device = "cuda"
69
+
70
+ # Random compressed_kv (simulates compressor output)
71
+ compressed_kv = torch.randn(
72
+ num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device
73
+ )
74
+
75
+ # ── Quant + insert ──────────────────────────────────────────────────
76
+ k_cache = torch.zeros(
77
+ num_blocks, block_size, HEAD_BYTES, dtype=torch.uint8, device=device
78
+ )
79
+ k_cache_2d = k_cache.view(num_blocks, -1)
80
+
81
+ # Sequential slot mapping: token i → slot i
82
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
83
+
84
+ quantize_and_insert_k_cache(
85
+ compressed_kv, k_cache_2d, slot_mapping, block_size=block_size
86
+ )
87
+
88
+ # ── Gather + dequant ────────────────────────────────────────────────
89
+ num_reqs = 1
90
+ max_blocks_per_seq = num_blocks
91
+ out = torch.zeros(
92
+ num_reqs, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device
93
+ )
94
+ seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
95
+ # block_table: request 0 uses physical blocks 0, 1, ...
96
+ block_table = torch.arange(
97
+ max_blocks_per_seq, dtype=torch.int32, device=device
98
+ ).unsqueeze(0)
99
+
100
+ dequantize_and_gather_k_cache(
101
+ out, k_cache, seq_lens, None, block_table, block_size, offset=0
102
+ )
103
+
104
+ recovered = out[0, :num_tokens]
105
+
106
+ # ── NoPE portion (first 448): FP8 quantized, expect UE8M0 error ──
107
+ nope_orig = compressed_kv[:, :NOPE_DIM].float()
108
+ nope_recv = recovered[:, :NOPE_DIM].float()
109
+ nope_diff = (nope_recv - nope_orig).abs()
110
+
111
+ # Per-token check: FP8 e4m3 (3-bit mantissa) worst-case error is
112
+ # half-ULP at the largest representable value. At y ≈ 448 (max),
113
+ # ULP = 2^(8-3) = 32, so error ≤ 16 * scale.
114
+ for t in range(num_tokens):
115
+ _, scales = _ue8m0_reference(
116
+ compressed_kv[t, :NOPE_DIM].float(), QUANT_BLOCK, FP8_MAX
117
+ )
118
+ max_allowed = 16.0 * scales.max().item()
119
+ token_diff = nope_diff[t].max().item()
120
+ assert token_diff <= max_allowed, (
121
+ f"Token {t} nope diff {token_diff} exceeds max_allowed "
122
+ f"{max_allowed} (scale={scales.max().item()})"
123
+ )
124
+
125
+ # ── RoPE portion (last 64): stored as bf16, should be exact ─────
126
+ rope_diff = (recovered[:, NOPE_DIM:] - compressed_kv[:, NOPE_DIM:]).abs()
127
+ assert rope_diff.max().item() == 0.0, (
128
+ f"RoPE portion should be exact but got max diff {rope_diff.max().item()}"
129
+ )
130
+
131
+
132
+# ── Test B: Indexer path ────────────────────────────────────────────────────
133
+
134
+
135
+@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
136
+@pytest.mark.parametrize("block_size", [16, 64])
137
+def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int):
138
+ """k → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache
139
+ → manual dequant → compare against original."""
140
+
141
+ HEAD_DIM = 128
142
+ QUANT_BLOCK_SIZE = 128
143
+ # cache_stride = head_dim + (head_dim * 4 / quant_block_size) = 128 + 4 = 132
144
+ CACHE_STRIDE = HEAD_DIM + HEAD_DIM * 4 // QUANT_BLOCK_SIZE
145
+
146
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
147
+ device = "cuda"
148
+
149
+ # Random K (simulates compressor output for indexer)
150
+ k = torch.randn(num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device)
151
+
152
+ # ── Quant + insert ──────────────────────────────────────────────────
153
+ kv_cache = torch.zeros(
154
+ num_blocks, block_size, CACHE_STRIDE, dtype=torch.uint8, device=device
155
+ )
156
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
157
+
158
+ ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, "ue8m0")
159
+
160
+ # ── Gather ──────────────────────────────────────────────────────────
161
+ max_blocks_per_seq = num_blocks
162
+ block_table = torch.arange(
163
+ max_blocks_per_seq, dtype=torch.int32, device=device
164
+ ).unsqueeze(0)
165
+ cu_seq_lens = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
166
+
167
+ # dst_k: [total_seq_len, head_dim] as uint8 (raw FP8 bytes)
168
+ dst_k = torch.zeros(num_tokens, HEAD_DIM, dtype=torch.uint8, device=device)
169
+ # dst_scale: [total_seq_len, head_dim/quant_block*4] as uint8 (raw float32 bytes)
170
+ num_scale_bytes = HEAD_DIM * 4 // QUANT_BLOCK_SIZE # 4
171
+ dst_scale = torch.zeros(
172
+ num_tokens, num_scale_bytes, dtype=torch.uint8, device=device
173
+ )
174
+
175
+ ops.cp_gather_indexer_k_quant_cache(
176
+ kv_cache, dst_k, dst_scale, block_table, cu_seq_lens
177
+ )
178
+
179
+ # ── Manual dequant ──────────────────────────────────────────────────
180
+ k_fp8 = dst_k.view(torch.float8_e4m3fn).float() # [num_tokens, 128]
181
+ scale = dst_scale.view(torch.float32) # [num_tokens, 1]
182
+ k_recovered = k_fp8 * scale # [num_tokens, 128]
183
+
184
+ # ── Compare ─────────────────────────────────────────────────────────
185
+ diff = (k_recovered - k.float()).abs()
186
+ k_abs = k.float().abs()
187
+
188
+ for t in range(num_tokens):
189
+ amax = k_abs[t].max().clamp(min=1e-4).item()
190
+ # UE8M0: scale = 2^ceil(log2(amax / 448))
191
+ exponent = math.ceil(math.log2(amax / 448.0))
192
+ ue8m0_scale = 2.0**exponent
193
+ # FP8 e4m3 (3-bit mantissa): worst-case error = 16 * scale
194
+ max_allowed = 16.0 * ue8m0_scale
195
+ token_diff = diff[t].max().item()
196
+ assert token_diff <= max_allowed, (
197
+ f"Token {t} diff {token_diff} exceeds max_allowed "
198
+ f"{max_allowed} (scale={ue8m0_scale})"
199
+ )
200
+
201
+
202
+def test_indexer_gather_accepts_upper_bound_output():
203
+ """Gather only exact cu_seq_lens even when dst is over-allocated."""
204
+
205
+ head_dim = 128
206
+ quant_block_size = 128
207
+ cache_stride = head_dim + head_dim * 4 // quant_block_size
208
+ valid_tokens = 9
209
+ upper_bound_tokens = 13
210
+ block_size = 16
211
+ num_blocks = 2
212
+ sentinel = 123
213
+ device = "cuda"
214
+
215
+ k = torch.randn(valid_tokens, head_dim, dtype=torch.bfloat16, device=device)
216
+ kv_cache = torch.zeros(
217
+ num_blocks, block_size, cache_stride, dtype=torch.uint8, device=device
218
+ )
219
+ slot_mapping = torch.arange(valid_tokens, dtype=torch.int64, device=device)
220
+ ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, quant_block_size, "ue8m0")
221
+
222
+ block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
223
+ 0
224
+ )
225
+ cu_seq_lens = torch.tensor([0, valid_tokens], dtype=torch.int32, device=device)
226
+ dst_k = torch.full(
227
+ (upper_bound_tokens, head_dim), sentinel, dtype=torch.uint8, device=device
228
+ )
229
+ num_scale_bytes = head_dim * 4 // quant_block_size
230
+ dst_scale = torch.full(
231
+ (upper_bound_tokens, num_scale_bytes),
232
+ sentinel,
233
+ dtype=torch.uint8,
234
+ device=device,
235
+ )
236
+
237
+ ops.cp_gather_indexer_k_quant_cache(
238
+ kv_cache, dst_k, dst_scale, block_table, cu_seq_lens
239
+ )
240
+ torch.accelerator.synchronize()
241
+
242
+ k_recovered = dst_k[:valid_tokens].view(torch.float8_e4m3fn).float() * dst_scale[
243
+ :valid_tokens
244
+ ].view(torch.float32)
245
+ diff = (k_recovered - k.float()).abs()
246
+ max_allowed = (16.0 * dst_scale[:valid_tokens].view(torch.float32).max()).item()
247
+ assert diff.max().item() <= max_allowed
248
+ assert torch.all(dst_k[valid_tokens:] == sentinel)
249
+ assert torch.all(dst_scale[valid_tokens:] == sentinel)
250
+
251
+
252
+# ── Test C: DeepseekV4 attention with values at different magnitudes ───────────
253
+
254
+
255
+def test_deepseek_v4_quant_magnitude_range():
256
+ """Test that quantization handles a range of magnitudes correctly."""
257
+
258
+ HEAD_DIM = 512
259
+ NOPE_DIM = 448
260
+ HEAD_BYTES = 584
261
+ block_size = 16
262
+ num_tokens = 4
263
+ num_blocks = 2
264
+ device = "cuda"
265
+
266
+ # Create inputs with varying magnitudes: small, medium, large
267
+ compressed_kv = torch.zeros(
268
+ num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device
269
+ )
270
+ compressed_kv[0] = 0.001 # very small
271
+ compressed_kv[1] = 1.0 # unit scale
272
+ compressed_kv[2] = 100.0 # large
273
+ compressed_kv[3] = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
274
+
275
+ k_cache = torch.zeros(
276
+ num_blocks, block_size, HEAD_BYTES, dtype=torch.uint8, device=device
277
+ )
278
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
279
+
280
+ quantize_and_insert_k_cache(
281
+ compressed_kv, k_cache.view(num_blocks, -1), slot_mapping, block_size
282
+ )
283
+
284
+ out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device)
285
+ seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
286
+ block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
287
+ 0
288
+ )
289
+
290
+ dequantize_and_gather_k_cache(
291
+ out, k_cache, seq_lens, None, block_table, block_size, offset=0
292
+ )
293
+
294
+ recovered = out[0, :num_tokens]
295
+
296
+ # RoPE portion must be exact
297
+ rope_diff = (recovered[:, NOPE_DIM:] - compressed_kv[:, NOPE_DIM:]).abs().max()
298
+ assert rope_diff.item() == 0.0, f"RoPE diff {rope_diff.item()}"
299
+
300
+ # NoPE: relative error should be reasonable
301
+ for t in range(num_tokens):
302
+ orig = compressed_kv[t, :NOPE_DIM].float()
303
+ recv = recovered[t, :NOPE_DIM].float()
304
+ abs_diff = (recv - orig).abs().max().item()
305
+ magnitude = orig.abs().max().item()
306
+ if magnitude > 0.01:
307
+ rel_err = abs_diff / magnitude
308
+ assert rel_err < 0.15, (
309
+ f"Token {t}: rel_err={rel_err:.4f}, abs_diff={abs_diff:.6f}, "
310
+ f"magnitude={magnitude:.4f}"
311
+ )
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
new
+359
@@ -0,0 +1,359 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel:
5
+
6
+ fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
7
+ - Q side: per-head RMSNorm (no weight) + GPT-J RoPE on last 64 dims
8
+ - KV side: GPT-J RoPE on last 64 + UE8M0 FP8 quant + paged cache insert
9
+
10
+We compare against:
11
+ - PyTorch reference for RMSNorm + GPT-J RoPE on Q
12
+ - Existing Triton `quantize_and_insert_k_cache` + round-trip via
13
+ `dequantize_and_gather_k_cache` for KV
14
+
15
+The kernel is imported via
16
+`torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert`.
17
+"""
18
+
19
+import pytest
20
+import torch
21
+
22
+from vllm.v1.attention.ops.deepseek_v4_ops import (
23
+ dequantize_and_gather_k_cache,
24
+ quantize_and_insert_k_cache,
25
+)
26
+
27
+# ── Constants matching the kernel ────────────────────────────────────────────
28
+HEAD_DIM = 512
29
+ROPE_DIM = 64
30
+NOPE_DIM = HEAD_DIM - ROPE_DIM # 448
31
+QUANT_BLOCK = 64
32
+FP8_MAX = 448.0
33
+HEAD_BYTES = NOPE_DIM + ROPE_DIM * 2 + 8 # 448 + 128 + 8 = 584
34
+
35
+
36
+# ── PyTorch reference implementations ────────────────────────────────────────
37
+
38
+
39
+def make_cos_sin_cache(max_pos: int, rope_dim: int, dtype, device):
40
+ """Build a cos||sin cache matching DeepseekV4ScalingRotaryEmbedding layout.
41
+ cos_sin_cache[pos, :rope_dim/2] = cos(theta), [rope_dim/2:] = sin(theta).
42
+ """
43
+ base = 10000.0
44
+ inv_freq = 1.0 / (
45
+ base
46
+ ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)
47
+ )
48
+ t = torch.arange(max_pos, dtype=torch.float32, device=device)
49
+ freqs = torch.einsum("i,j -> ij", t, inv_freq) # [max_pos, rope_dim/2]
50
+ cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1) # [max_pos, rope_dim]
51
+ return cache.to(dtype)
52
+
53
+
54
+def apply_rope_gptj_last_k(
55
+ x: torch.Tensor, positions: torch.Tensor, cos_sin_cache: torch.Tensor
56
+) -> torch.Tensor:
57
+ """GPT-J-style (interleaved-pair) RoPE on the LAST rope_dim elements.
58
+
59
+ x: [..., head_dim] float32
60
+ positions: [num_tokens] int64 (positions[i] corresponds to x[i, ...])
61
+ cos_sin_cache: [max_pos, rope_dim] float (cos|sin layout)
62
+
63
+ Returns rotated x (same shape/dtype).
64
+ """
65
+ rope_dim = cos_sin_cache.shape[-1]
66
+ half = rope_dim // 2
67
+ head_dim = x.shape[-1]
68
+ nope_dim = head_dim - rope_dim
69
+
70
+ # Gather cos/sin for each token position: [num_tokens, rope_dim]
71
+ cs = cos_sin_cache[positions].to(torch.float32) # [N, rope_dim]
72
+ cos = cs[..., :half] # [N, half]
73
+ sin = cs[..., half:] # [N, half]
74
+
75
+ # Reshape leading dims so we can broadcast: x shape [..., head_dim].
76
+ # Bring token dim to front; assume x is [num_tokens, ..., head_dim].
77
+ # We rely on positions being per-token and all other dims sharing the same pos.
78
+ rope = x[..., nope_dim:].float() # [..., rope_dim]
79
+ # Make rope pairs: reshape last dim to [half, 2]
80
+ shape = rope.shape
81
+ rope = rope.reshape(*shape[:-1], half, 2)
82
+ even = rope[..., 0] # [..., half]
83
+ odd = rope[..., 1]
84
+
85
+ # Broadcast cos/sin over any heads dim in between. cos/sin are [N, half].
86
+ # Add singleton dims for intermediate axes.
87
+ for _ in range(rope.ndim - 3):
88
+ cos = cos.unsqueeze(1)
89
+ sin = sin.unsqueeze(1)
90
+
91
+ new_even = even * cos - odd * sin
92
+ new_odd = even * sin + odd * cos
93
+ rope_rotated = torch.stack((new_even, new_odd), dim=-1).reshape(shape)
94
+
95
+ out = x.clone().float()
96
+ out[..., nope_dim:] = rope_rotated
97
+ return out.to(x.dtype)
98
+
99
+
100
+def rmsnorm_no_weight(x: torch.Tensor, eps: float) -> torch.Tensor:
101
+ """RMSNorm with no learnable weight, matching
102
+ `RMSNorm(head_dim, has_weight=False)`."""
103
+ orig_dtype = x.dtype
104
+ xf = x.float()
105
+ variance = xf.pow(2).mean(dim=-1, keepdim=True)
106
+ return (xf * torch.rsqrt(variance + eps)).to(orig_dtype)
107
+
108
+
109
+# ── Dispatch to the CUDA op (skip test cleanly if it isn't built in) ─────────
110
+
111
+
112
+def _op_available() -> bool:
113
+ return hasattr(torch.ops._C, "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert")
114
+
115
+
116
+pytestmark = pytest.mark.skipif(
117
+ not torch.cuda.is_available() or not _op_available(),
118
+ reason="CUDA not available or fused DeepseekV4 op not built in",
119
+)
120
+
121
+
122
+def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs):
123
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
124
+ q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
125
+ )
126
+
127
+
128
+# ── Test 1: Q path numerical parity ──────────────────────────────────────────
129
+
130
+
131
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64])
132
+@pytest.mark.parametrize("n_heads", [8, 64])
133
+def test_q_path_matches_reference(num_tokens: int, n_heads: int):
134
+ torch.manual_seed(0)
135
+ device = "cuda"
136
+ dtype = torch.bfloat16
137
+ eps = 1e-6
138
+ max_pos = 4096
139
+
140
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
141
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
142
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
143
+
144
+ # Reference: RMSNorm (no weight) per head, then GPT-J RoPE on last 64.
145
+ q_ref = rmsnorm_no_weight(q, eps)
146
+ q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)
147
+
148
+ # Fused call with dummy KV tensors (KV branch will write slot_mapping=-1 → noop).
149
+ num_blocks = 2
150
+ bs = 16
151
+ kv = torch.zeros(num_tokens, HEAD_DIM, dtype=dtype, device=device)
152
+ k_cache = torch.zeros(
153
+ num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device
154
+ ).view(num_blocks, -1)
155
+ slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
156
+ q_fused = q.clone()
157
+ _call_fused(q_fused, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
158
+
159
+ torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
160
+
161
+
162
+# ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
163
+
164
+
165
+def _ue8m0_per_block_scales(kv_roped_nope_f32: torch.Tensor, qblock: int):
166
+ """Return per-token per-block max scale (used to bound FP8 error)."""
167
+ n_tok, nope = kv_roped_nope_f32.shape
168
+ n_blocks = nope // qblock
169
+ blocks = kv_roped_nope_f32.view(n_tok, n_blocks, qblock)
170
+ absmax = blocks.abs().amax(dim=-1).clamp(min=1e-4)
171
+ raw = absmax / FP8_MAX
172
+ exponent = torch.ceil(torch.log2(raw))
173
+ return torch.pow(2.0, exponent) # [n_tok, n_blocks]
174
+
175
+
176
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64])
177
+@pytest.mark.parametrize("block_size", [16, 64])
178
+def test_kv_path_matches_reference(num_tokens: int, block_size: int):
179
+ torch.manual_seed(1)
180
+ device = "cuda"
181
+ dtype = torch.bfloat16
182
+ eps = 1e-6
183
+ max_pos = 4096
184
+
185
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
186
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
187
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
188
+
189
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
190
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
191
+
192
+ # ── Reference path: RoPE on kv, then existing Triton quant+insert ──────
193
+ kv_ref = apply_rope_gptj_last_k(kv, positions, cos_sin_cache)
194
+ k_cache_ref = torch.zeros(
195
+ num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
196
+ )
197
+ quantize_and_insert_k_cache(
198
+ kv_ref, k_cache_ref, slot_mapping, block_size=block_size
199
+ )
200
+
201
+ # ── Fused path (dummy q, single head) ──────────────────────────────────
202
+ k_cache_fused = torch.zeros_like(k_cache_ref)
203
+ q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device)
204
+ _call_fused(
205
+ q_dummy,
206
+ kv,
207
+ k_cache_fused,
208
+ slot_mapping,
209
+ positions,
210
+ cos_sin_cache,
211
+ eps,
212
+ block_size,
213
+ )
214
+
215
+ # ── Round-trip compare via dequant+gather ──────────────────────────────
216
+ def _dequant(k_cache_2d):
217
+ num_reqs = 1
218
+ max_blocks = num_blocks
219
+ out = torch.zeros(
220
+ num_reqs, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device
221
+ )
222
+ seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
223
+ block_table = torch.arange(
224
+ max_blocks, dtype=torch.int32, device=device
225
+ ).unsqueeze(0)
226
+ # gather_lens arg is None (use seq_lens)
227
+ k_cache_3d = k_cache_2d.view(num_blocks, block_size, HEAD_BYTES)
228
+ dequantize_and_gather_k_cache(
229
+ out, k_cache_3d, seq_lens, None, block_table, block_size, offset=0
230
+ )
231
+ return out[0, :num_tokens]
232
+
233
+ recovered_ref = _dequant(k_cache_ref)
234
+ recovered_fused = _dequant(k_cache_fused)
235
+
236
+ # NoPE: per-block UE8M0 FP8 error bound (half-ULP at max = 16 * scale).
237
+ scales = _ue8m0_per_block_scales(kv_ref[:, :NOPE_DIM].float(), QUANT_BLOCK)
238
+ for t in range(num_tokens):
239
+ max_allowed = 16.0 * scales[t].max().item()
240
+ diff_ref = (
241
+ (recovered_ref[t, :NOPE_DIM] - kv_ref[t, :NOPE_DIM]).abs().max().item()
242
+ )
243
+ diff_fused = (
244
+ (recovered_fused[t, :NOPE_DIM] - kv_ref[t, :NOPE_DIM]).abs().max().item()
245
+ )
246
+ assert diff_ref <= max_allowed, (
247
+ f"ref NoPE token {t} diff {diff_ref} > {max_allowed}"
248
+ )
249
+ assert diff_fused <= max_allowed, (
250
+ f"fused NoPE token {t} diff {diff_fused} > {max_allowed}"
251
+ )
252
+
253
+ # RoPE region: bf16 stored exactly → zero diff.
254
+ rope_diff = (recovered_fused[:, NOPE_DIM:] - kv_ref[:, NOPE_DIM:]).abs().max()
255
+ assert rope_diff.item() == 0.0, f"RoPE portion not exact: {rope_diff.item()}"
256
+
257
+ # Exact byte equality of the two cache buffers — strong parity.
258
+ torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)
259
+
260
+
261
+# ── Test 2b: DP padding (slot_mapping shorter than q/kv) ─────────────────────
262
+
263
+
264
+@pytest.mark.parametrize("num_tokens", [4, 17])
265
+@pytest.mark.parametrize("pad", [1, 5])
266
+@pytest.mark.parametrize("block_size", [16, 64])
267
+def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
268
+ """slot_mapping.size(0) < q.size(0): the kernel must skip padded
269
+ tokens in the KV branch while still running Q-norm+RoPE on all rows."""
270
+ torch.manual_seed(3)
271
+ device = "cuda"
272
+ dtype = torch.bfloat16
273
+ eps = 1e-6
274
+ max_pos = 4096
275
+ total = num_tokens + pad
276
+
277
+ kv = torch.randn(total, HEAD_DIM, dtype=dtype, device=device)
278
+ positions = torch.arange(total, dtype=torch.int64, device=device)
279
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
280
+
281
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
282
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
283
+
284
+ # Reference: only the first num_tokens kv rows get inserted.
285
+ kv_ref = apply_rope_gptj_last_k(
286
+ kv[:num_tokens], positions[:num_tokens], cos_sin_cache
287
+ )
288
+ k_cache_ref = torch.zeros(
289
+ num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
290
+ )
291
+ quantize_and_insert_k_cache(
292
+ kv_ref, k_cache_ref, slot_mapping, block_size=block_size
293
+ )
294
+
295
+ # Fused: pass full-sized q/kv/positions, shorter slot_mapping.
296
+ q_dummy = torch.zeros(total, 1, HEAD_DIM, dtype=dtype, device=device)
297
+ k_cache_fused = torch.zeros_like(k_cache_ref)
298
+ _call_fused(
299
+ q_dummy,
300
+ kv,
301
+ k_cache_fused,
302
+ slot_mapping,
303
+ positions,
304
+ cos_sin_cache,
305
+ eps,
306
+ block_size,
307
+ )
308
+
309
+ torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)
310
+
311
+
312
+# ── Test 3: combined single-call Q + KV parity ───────────────────────────────
313
+
314
+
315
+@pytest.mark.parametrize("num_tokens", [1, 4, 17])
316
+@pytest.mark.parametrize("n_heads", [8, 64])
317
+@pytest.mark.parametrize("block_size", [16, 64])
318
+def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
319
+ torch.manual_seed(2)
320
+ device = "cuda"
321
+ dtype = torch.bfloat16
322
+ eps = 1e-6
323
+ max_pos = 4096
324
+
325
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
326
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
327
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
328
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
329
+
330
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
331
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
332
+
333
+ # Reference.
334
+ q_ref = rmsnorm_no_weight(q, eps)
335
+ q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)
336
+ kv_ref = apply_rope_gptj_last_k(kv, positions, cos_sin_cache)
337
+ k_cache_ref = torch.zeros(
338
+ num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
339
+ )
340
+ quantize_and_insert_k_cache(
341
+ kv_ref, k_cache_ref, slot_mapping, block_size=block_size
342
+ )
343
+
344
+ # Fused single call.
345
+ q_fused = q.clone()
346
+ k_cache_fused = torch.zeros_like(k_cache_ref)
347
+ _call_fused(
348
+ q_fused,
349
+ kv,
350
+ k_cache_fused,
351
+ slot_mapping,
352
+ positions,
353
+ cos_sin_cache,
354
+ eps,
355
+ block_size,
356
+ )
357
+
358
+ torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
359
+ torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)
tests/kernels/test_fused_indexer_q_rope_quant.py
new
+98
@@ -0,0 +1,98 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit test for fused_indexer_q_rope_quant.
4
+
5
+Compares the fused Triton kernel against the unfused reference flow used by
6
+the DeepseekV4 indexer in model_tracking:
7
+ q_rot = ops.rotary_embedding(positions, q, None, head_dim, cos_sin_cache,
8
+ is_neox_style=False,
9
+ rope_dim_offset=head_dim - rope_dim)
10
+ q_fp8, q_scale = per_token_group_quant_fp8(q_rot, head_dim, use_ue8m0=True)
11
+ weights_out = weights * q_scale * softmax_scale * head_scale
12
+
13
+Expects bit-exact equality on both q_fp8 and weights_out.
14
+"""
15
+
16
+import pytest
17
+import torch
18
+
19
+from vllm import _custom_ops as ops
20
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
21
+ per_token_group_quant_fp8,
22
+)
23
+from vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q import (
24
+ fused_indexer_q_rope_quant,
25
+)
26
+
27
+HEAD_DIM = 128
28
+ROPE_DIM = 64
29
+N_HEAD = 64
30
+MAX_POS = 4096
31
+
32
+
33
+def _reference(
34
+ positions: torch.Tensor,
35
+ q: torch.Tensor,
36
+ cos_sin_cache: torch.Tensor,
37
+ weights: torch.Tensor,
38
+ softmax_scale: float,
39
+ head_scale: float,
40
+) -> tuple[torch.Tensor, torch.Tensor]:
41
+ q_rot = q.clone()
42
+ ops.rotary_embedding(
43
+ positions,
44
+ q_rot,
45
+ None,
46
+ HEAD_DIM,
47
+ cos_sin_cache,
48
+ False, # is_neox_style=False → GPT-J interleaved
49
+ HEAD_DIM - ROPE_DIM, # rope_dim_offset → rotate the tail
50
+ False,
51
+ )
52
+ q_fp8, q_scale = per_token_group_quant_fp8(
53
+ q_rot.view(-1, HEAD_DIM).contiguous(),
54
+ HEAD_DIM,
55
+ use_ue8m0=True,
56
+ )
57
+ q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
58
+ q_scale = q_scale.view(-1, N_HEAD)
59
+
60
+ weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
61
+ return q_fp8, weights_out
62
+
63
+
64
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
65
+@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
66
+@torch.inference_mode()
67
+def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
68
+ device = "cuda"
69
+ torch.manual_seed(0)
70
+
71
+ q = torch.randn(num_tokens, N_HEAD, HEAD_DIM, dtype=torch.bfloat16, device=device)
72
+ positions = torch.randint(
73
+ 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device
74
+ )
75
+ cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=cache_dtype, device=device)
76
+ weights = torch.randn(num_tokens, N_HEAD, dtype=torch.bfloat16, device=device)
77
+ softmax_scale = HEAD_DIM**-0.5
78
+ head_scale = N_HEAD**-0.5
79
+
80
+ q_fp8_ref, weights_ref = _reference(
81
+ positions, q, cos_sin_cache, weights, softmax_scale, head_scale
82
+ )
83
+ q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
84
+ positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
85
+ )
86
+
87
+ # fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8.
88
+ ref_bits = q_fp8_ref.view(torch.int8)
89
+ fused_bits = q_fp8_fused.view(torch.int8)
90
+ assert torch.equal(ref_bits, fused_bits), (
91
+ f"q_fp8 mismatch: "
92
+ f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ"
93
+ )
94
+
95
+ assert torch.equal(weights_ref, weights_fused), (
96
+ f"weights mismatch: max abs diff "
97
+ f"{(weights_ref - weights_fused).abs().max().item()}"
98
+ )
tests/kernels/test_fused_inv_rope_fp8_quant.py
new
+908
@@ -0,0 +1,908 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel.
5
+
6
+Tests compare the fused kernel against a reference implementation built from
7
+the existing separate operations (inverse RoPE via rotate_neox + FP8 quant
8
+via per_token_group_quant_fp8).
9
+
10
+The reference faithfully reproduces the exact flow in deepseek_v4_attention.py:295-310:
11
+ 1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head)
12
+ 2. Reshape [T, H, head_dim] -> [T, G, D]
13
+ 3. Transpose+flatten to [G*T, D], quantize, reshape back
14
+ 4. Return o_fp8 and o_scale with strides (D, T*D, 1) and (S, T*S, 1)
15
+ (non-contiguous [T, G, ...] view backed by contiguous [G, T, ...] memory)
16
+
17
+Usage:
18
+ pytest tests/kernels/test_fused_inv_rope_fp8_quant.py -v
19
+"""
20
+
21
+import pytest
22
+import torch
23
+
24
+from vllm.v1.attention.ops.deepseek_v4_ops import fused_inv_rope_fp8_quant
25
+
26
+# -- Default dimensions matching DeepSeek V3/V4 --------------------------
27
+HEAD_DIM = 512
28
+NOPE_DIM = 448
29
+ROPE_DIM = 64
30
+QUANT_GROUP_SIZE = 128
31
+FP8_MAX = 448.0 # torch.finfo(torch.float8_e4m3fn).max
32
+FP8_DTYPE = torch.float8_e4m3fn
33
+EPS = 1e-10
34
+
35
+
36
+# =========================================================================
37
+# Helpers
38
+# =========================================================================
39
+
40
+
41
+def assert_dequant_close(
42
+ fp8_a: torch.Tensor,
43
+ scale_a: torch.Tensor,
44
+ fp8_b: torch.Tensor,
45
+ scale_b: torch.Tensor,
46
+ msg: str = "",
47
+):
48
+ """Compare two FP8-quantized tensors via their dequantized values.
49
+
50
+ Uses cosine-similarity-based diff (same as deep_gemm calc_diff).
51
+ Both fused and reference paths rotate in fp32 using an fp32
52
+ cos_sin_cache, so differences are only fp32 ordering ULPs that can
53
+ occasionally shift FP8 values at quantization boundaries.
54
+ """
55
+ S = scale_a.shape[-1]
56
+ shape = fp8_a.shape
57
+
58
+ dq_a = fp8_a.float() * scale_a.unsqueeze(-1).expand(
59
+ *shape[:-1], S, QUANT_GROUP_SIZE
60
+ ).reshape(shape)
61
+ dq_b = fp8_b.float() * scale_b.unsqueeze(-1).expand(
62
+ *shape[:-1], S, QUANT_GROUP_SIZE
63
+ ).reshape(shape)
64
+
65
+ # Cosine diff: 1 - cos_sim (0 = identical, higher = worse)
66
+ dq_a_flat = dq_a.flatten().float()
67
+ dq_b_flat = dq_b.flatten().float()
68
+ cos_sim = torch.nn.functional.cosine_similarity(
69
+ dq_a_flat.unsqueeze(0), dq_b_flat.unsqueeze(0)
70
+ ).item()
71
+ diff = 1.0 - cos_sim
72
+
73
+ assert diff < 1e-4, f"Dequant diff too large: {diff:.8f} (expected < 1e-4). {msg}"
74
+
75
+
76
+def rotate_gptj(x: torch.Tensor) -> torch.Tensor:
77
+ """GPT-J style rotation: interleaved pairs, negate-swap.
78
+
79
+ Matches vllm/model_executor/layers/rotary_embedding/common.py:23-27.
80
+ DeepseekV4 uses is_neox_style=False, so this is the correct rotation.
81
+ """
82
+ x1 = x[..., ::2]
83
+ x2 = x[..., 1::2]
84
+ x = torch.stack((-x2, x1), dim=-1)
85
+ return x.flatten(-2)
86
+
87
+
88
+def make_cos_sin_cache(
89
+ max_pos: int,
90
+ rope_dim: int = ROPE_DIM,
91
+ dtype: torch.dtype = torch.float32,
92
+ device: str = "cuda",
93
+) -> torch.Tensor:
94
+ """Create a synthetic cos_sin_cache matching the layout used by
95
+ DeepseekV4ScalingRotaryEmbedding._compute_cos_sin_cache.
96
+
97
+ Shape: [max_pos, rope_dim] where first half is cos, second half is sin.
98
+ The fused kernel requires fp32; callers can override dtype if passing
99
+ the cache into the bf16-only paths.
100
+ """
101
+ half = rope_dim // 2
102
+ # Use random but bounded frequencies so cos/sin are well-behaved
103
+ inv_freq = 1.0 / (
104
+ 10000.0 ** (torch.arange(0, half, device=device, dtype=torch.float32) / half)
105
+ )
106
+ t = torch.arange(max_pos, device=device, dtype=torch.float32)
107
+ freqs = torch.outer(t, inv_freq) # [max_pos, half]
108
+ cos = freqs.cos()
109
+ sin = freqs.sin()
110
+ cache = torch.cat((cos, sin), dim=-1) # [max_pos, rope_dim]
111
+ return cache.to(dtype)
112
+
113
+
114
+def reference_inv_rope(
115
+ o: torch.Tensor,
116
+ positions: torch.Tensor,
117
+ cos_sin_cache: torch.Tensor,
118
+ nope_dim: int = NOPE_DIM,
119
+ rope_dim: int = ROPE_DIM,
120
+) -> torch.Tensor:
121
+ """Apply inverse RoPE to the last rope_dim dimensions of each head.
122
+
123
+ Matches the GPT-J inverse rotation in pos_encoding_kernels.cu, which
124
+ promotes the cache to fp32 and performs the rotation in fp32. The
125
+ result is cast back to the input dtype.
126
+
127
+ Args:
128
+ o: [T, H, head_dim] bf16
129
+ positions: [T] int64
130
+ cos_sin_cache: [max_pos, rope_dim] fp32
131
+
132
+ Returns:
133
+ o with inverse RoPE applied on the rope portion (bf16).
134
+ """
135
+ assert cos_sin_cache.dtype == torch.float32
136
+ cos_sin = cos_sin_cache[positions] # [T, rope_dim] fp32
137
+ half = rope_dim // 2
138
+ cos = cos_sin[:, :half]
139
+ sin = cos_sin[:, half:]
140
+
141
+ # GPT-J style: repeat_interleave (not repeat) to match interleaved pairs
142
+ cos = cos.repeat_interleave(2, dim=-1).unsqueeze(1)
143
+ sin = sin.repeat_interleave(2, dim=-1).unsqueeze(1)
144
+ sin = -sin # inverse
145
+
146
+ o_pass = o[..., :nope_dim]
147
+ o_rot_f32 = o[..., nope_dim:].float()
148
+ o_rot_f32 = o_rot_f32 * cos + rotate_gptj(o_rot_f32) * sin
149
+ o_rot = o_rot_f32.to(o.dtype)
150
+
151
+ return torch.cat([o_pass, o_rot], dim=-1)
152
+
153
+
154
+def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
155
+ """Per-block UE8M0 FP8 quantization in pure float32.
156
+
157
+ Matches the Triton kernel logic exactly:
158
+ absmax -> 2^ceil(log2(absmax / fp8_max)) -> clamp(x / scale) -> fp8
159
+
160
+ Args:
161
+ x_f32: [..., quant_group_size] float32 — one or more 128-element blocks.
162
+
163
+ Returns:
164
+ x_fp8: same shape, float8_e4m3fn
165
+ scales: [...] float32, one scale per block
166
+ """
167
+ assert x_f32.shape[-1] == QUANT_GROUP_SIZE
168
+ absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS)
169
+ scale_raw = absmax * (1.0 / FP8_MAX)
170
+ scale = torch.exp2(torch.ceil(torch.log2(scale_raw)))
171
+ x_scaled = (x_f32 / scale).clamp(-FP8_MAX, FP8_MAX)
172
+ x_fp8 = x_scaled.to(FP8_DTYPE)
173
+ return x_fp8, scale.squeeze(-1)
174
+
175
+
176
+def reference_inv_rope_fp8_quant(
177
+ o: torch.Tensor,
178
+ positions: torch.Tensor,
179
+ cos_sin_cache: torch.Tensor,
180
+ n_groups: int,
181
+ heads_per_group: int,
182
+ nope_dim: int = NOPE_DIM,
183
+ rope_dim: int = ROPE_DIM,
184
+ quant_group_size: int = QUANT_GROUP_SIZE,
185
+) -> tuple[torch.Tensor, torch.Tensor]:
186
+ """Full reference: inverse RoPE in fp32 + UE8M0 FP8 quant in fp32.
187
+
188
+ Mimics the Triton kernel's precision path exactly:
189
+ Load bf16 -> cast to fp32 -> apply inverse RoPE with fp32 cos/sin ->
190
+ UE8M0 quant in fp32 -> write fp8 + scale
191
+
192
+ Returns:
193
+ o_fp8: [T, G, D] FP8 with strides (D, T*D, 1)
194
+ o_scale: [T, G, S] FP32 with strides (S, T*S, 1)
195
+ """
196
+ assert cos_sin_cache.dtype == torch.float32
197
+ T, _H, head_dim = o.shape
198
+ d = heads_per_group * head_dim
199
+ S = d // quant_group_size
200
+ half_rope = rope_dim // 2
201
+ chunks_per_head = head_dim // quant_group_size
202
+
203
+ # Reshape [T, H, head_dim] -> [T, G, heads_per_group, head_dim]
204
+ o_4d = o.view(T, n_groups, heads_per_group, head_dim)
205
+
206
+ # Lookup cos/sin directly in fp32
207
+ cos_sin = cos_sin_cache[positions] # [T, rope_dim] fp32
208
+ cos = cos_sin[:, :half_rope] # [T, half_rope] fp32
209
+ sin = cos_sin[:, half_rope:] # [T, half_rope] fp32
210
+
211
+ # Allocate outputs in [G, T, ...] contiguous layout
212
+ fp8_buf = torch.empty(n_groups, T, d, dtype=FP8_DTYPE, device=o.device)
213
+ scale_buf = torch.empty(n_groups, T, S, dtype=torch.float32, device=o.device)
214
+
215
+ # Process each quant block, matching the Triton kernel's per-program logic
216
+ for g in range(n_groups):
217
+ for qb in range(S):
218
+ head_in_group = qb // chunks_per_head
219
+ chunk_in_head = qb % chunks_per_head
220
+ offset = chunk_in_head * quant_group_size
221
+
222
+ # Load 128 bf16 elements and promote to fp32 for rotation+quant
223
+ block = o_4d[:, g, head_in_group, offset : offset + quant_group_size]
224
+ x = block.float()
225
+
226
+ # Apply inverse RoPE in fp32 if this is the last chunk
227
+ # GPT-J style: interleaved pairs (even=x, odd=y)
228
+ if chunk_in_head == chunks_per_head - 1:
229
+ rope_start = nope_dim % quant_group_size # 64
230
+ rope_region = x[:, rope_start:].clone()
231
+ x_vals = rope_region[:, ::2]
232
+ y_vals = rope_region[:, 1::2]
233
+ x_new = x_vals * cos + y_vals * sin
234
+ y_new = y_vals * cos - x_vals * sin
235
+ x = x.clone()
236
+ x[:, rope_start::2] = x_new
237
+ x[:, rope_start + 1 :: 2] = y_new
238
+
239
+ # UE8M0 quant in fp32
240
+ x_fp8, scale = _ref_ue8m0_quant_block(x)
241
+
242
+ # Write to [G, T, D] contiguous memory
243
+ fp8_buf[g, :, qb * quant_group_size : (qb + 1) * quant_group_size] = x_fp8
244
+ scale_buf[g, :, qb] = scale
245
+
246
+ # Return transposed views
247
+ return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1)
248
+
249
+
250
+# =========================================================================
251
+# Tests
252
+# =========================================================================
253
+
254
+
255
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
256
+@pytest.mark.parametrize(
257
+ "num_heads,n_groups",
258
+ [(64, 8), (32, 4), (128, 8)],
259
+ ids=["H64_G8", "H32_G4", "H128_G8"],
260
+)
261
+@pytest.mark.parametrize("seed", [0, 42])
262
+@torch.inference_mode()
263
+def test_correctness(num_tokens, num_heads, n_groups, seed):
264
+ """Compare fused kernel against reference for FP8 values and scales."""
265
+ torch.manual_seed(seed)
266
+
267
+ heads_per_group = num_heads // n_groups
268
+ max_pos = 4096
269
+ device = "cuda"
270
+
271
+ # Create inputs
272
+ o = torch.randn(
273
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
274
+ )
275
+ positions = torch.randint(
276
+ 0, max_pos, (num_tokens,), device=device, dtype=torch.long
277
+ )
278
+ cos_sin_cache = make_cos_sin_cache(
279
+ max_pos, ROPE_DIM, dtype=torch.float32, device=device
280
+ )
281
+
282
+ # Reference
283
+ ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
284
+ o.clone(),
285
+ positions,
286
+ cos_sin_cache,
287
+ n_groups,
288
+ heads_per_group,
289
+ )
290
+
291
+ # Fused kernel
292
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
293
+ o.clone(),
294
+ positions,
295
+ cos_sin_cache,
296
+ n_groups,
297
+ heads_per_group,
298
+ )
299
+
300
+ # Check shapes
301
+ d = heads_per_group * HEAD_DIM
302
+ S = d // QUANT_GROUP_SIZE
303
+ assert ref_fp8.shape == (num_tokens, n_groups, d)
304
+ assert fused_fp8.shape == (num_tokens, n_groups, d)
305
+ assert ref_scale.shape == (num_tokens, n_groups, S)
306
+ assert fused_scale.shape == (num_tokens, n_groups, S)
307
+
308
+ # Scales: exact match (both use identical UE8M0 algorithm)
309
+ # Scales may differ by one UE8M0 step (factor of 2) if fp32 rotation
310
+ # ordering shifts absmax across a power-of-2 boundary. Check ratio is
311
+ # close to 1.
312
+ scale_ratio = fused_scale / ref_scale.clamp(min=1e-30)
313
+ assert scale_ratio.max() <= 2.0 and scale_ratio.min() >= 0.5, (
314
+ f"Scale ratio out of [0.5, 2]: min={scale_ratio.min():.4f} "
315
+ f"max={scale_ratio.max():.4f}"
316
+ )
317
+
318
+ # Compare via dequant (Triton vs PyTorch fp32 may differ by ULPs)
319
+ assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale)
320
+
321
+
322
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
323
+@pytest.mark.parametrize(
324
+ "num_heads,n_groups",
325
+ [(64, 8), (128, 8)],
326
+ ids=["H64_G8", "H128_G8"],
327
+)
328
+@torch.inference_mode()
329
+def test_output_strides(num_tokens, num_heads, n_groups):
330
+ """Verify fused output layout:
331
+ - FP8: logical [T, G, D] backed by contiguous [G, T, D].
332
+ - Scale: MN-major TMA-aligned (column-major: T-stride=1).
333
+ """
334
+
335
+ heads_per_group = num_heads // n_groups
336
+ max_pos = 4096
337
+ device = "cuda"
338
+
339
+ o = torch.randn(
340
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
341
+ )
342
+ positions = torch.randint(
343
+ 0, max_pos, (num_tokens,), device=device, dtype=torch.long
344
+ )
345
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
346
+
347
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
348
+ o.clone(),
349
+ positions,
350
+ cos_sin_cache,
351
+ n_groups,
352
+ heads_per_group,
353
+ )
354
+
355
+ # FP8: logical [T, G, D] backed by [G, T, D] row-major
356
+ d = heads_per_group * HEAD_DIM
357
+ expected_fp8_stride = (d, num_tokens * d, 1)
358
+ assert fused_fp8.stride() == expected_fp8_stride, (
359
+ f"FP8 stride mismatch: got {fused_fp8.stride()}, expected {expected_fp8_stride}"
360
+ )
361
+
362
+ # Scale: MN-major TMA-aligned layout. After fp8_einsum permutes
363
+ # [T,G,S] -> [G,T,S], T-dim should have stride 1.
364
+ # Our output is [T,G,S] = transpose of [G,T,S].
365
+ # So fused_scale.permute(1,0,2) should have T-stride=1.
366
+ perm = fused_scale.permute(1, 0, 2) # [G, T, S]
367
+ assert perm.stride(1) == 1 or num_tokens == 1, (
368
+ f"Scale T-stride (after permute to [G,T,S]) should be 1, got {perm.stride(1)}"
369
+ )
370
+
371
+
372
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
373
+@torch.inference_mode()
374
+def test_per_group_contiguity(num_tokens):
375
+ """FP8 per-group slices must be contiguous. Scale per-group slices
376
+ are column-major (T-stride=1) — not row-major contiguous, which is
377
+ correct for TMA loads."""
378
+ num_heads, n_groups = 64, 8
379
+ heads_per_group = num_heads // n_groups
380
+ max_pos = 4096
381
+ device = "cuda"
382
+
383
+ o = torch.randn(
384
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
385
+ )
386
+ positions = torch.randint(
387
+ 0, max_pos, (num_tokens,), device=device, dtype=torch.long
388
+ )
389
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
390
+
391
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
392
+ o.clone(),
393
+ positions,
394
+ cos_sin_cache,
395
+ n_groups,
396
+ heads_per_group,
397
+ )
398
+
399
+ for g in range(n_groups):
400
+ fp8_slice = fused_fp8[:, g, :]
401
+ assert fp8_slice.is_contiguous(), (
402
+ f"o_fp8[:, {g}, :] is not contiguous: "
403
+ f"shape={list(fp8_slice.shape)}, stride={list(fp8_slice.stride())}"
404
+ )
405
+
406
+
407
+@torch.inference_mode()
408
+def test_scales_are_power_of_two():
409
+ """Verify all scales are exact powers of 2 (UE8M0 property)."""
410
+ num_tokens, num_heads, n_groups = 32, 64, 8
411
+ heads_per_group = num_heads // n_groups
412
+ max_pos = 4096
413
+ device = "cuda"
414
+
415
+ o = torch.randn(
416
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
417
+ )
418
+ positions = torch.randint(
419
+ 0, max_pos, (num_tokens,), device=device, dtype=torch.long
420
+ )
421
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
422
+
423
+ _, fused_scale = fused_inv_rope_fp8_quant(
424
+ o.clone(),
425
+ positions,
426
+ cos_sin_cache,
427
+ n_groups,
428
+ heads_per_group,
429
+ )
430
+
431
+ # log2 of a power-of-two is an exact integer
432
+ log2_scales = torch.log2(fused_scale)
433
+ residual = (log2_scales - log2_scales.round()).abs()
434
+ assert residual.max() < 1e-5, (
435
+ f"Not all scales are powers of 2: max log2 residual = {residual.max().item()}"
436
+ )
437
+
438
+
439
+@torch.inference_mode()
440
+def test_nope_dims_unchanged():
441
+ """Nope dimensions (first 448 per head) should only be quantized,
442
+ not rotated. Verify by dequantizing and comparing against
443
+ quantize-only reference (no RoPE)."""
444
+ num_tokens, num_heads, n_groups = 16, 64, 8
445
+ heads_per_group = num_heads // n_groups
446
+ max_pos = 4096
447
+ device = "cuda"
448
+ torch.manual_seed(0)
449
+
450
+ o = torch.randn(
451
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
452
+ )
453
+ positions = torch.randint(
454
+ 0, max_pos, (num_tokens,), device=device, dtype=torch.long
455
+ )
456
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
457
+
458
+ # Fused kernel result
459
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
460
+ o.clone(),
461
+ positions,
462
+ cos_sin_cache,
463
+ n_groups,
464
+ heads_per_group,
465
+ )
466
+
467
+ # Reference: quantize without RoPE (identity rotation)
468
+ # Create a zero-sin cache so RoPE is identity
469
+ zero_cache = torch.zeros_like(cos_sin_cache)
470
+ half = ROPE_DIM // 2
471
+ zero_cache[:, :half] = 1.0 # cos = 1
472
+ # sin = 0 (already zero)
473
+
474
+ norope_fp8, norope_scale = fused_inv_rope_fp8_quant(
475
+ o.clone(),
476
+ positions,
477
+ zero_cache,
478
+ n_groups,
479
+ heads_per_group,
480
+ )
481
+
482
+ # Extract nope quant blocks only (first 3 of every 4 blocks per head)
483
+ chunks_per_head = HEAD_DIM // QUANT_GROUP_SIZE # 4
484
+
485
+ for h in range(heads_per_group):
486
+ for c in range(chunks_per_head - 1): # skip last chunk (has rope)
487
+ qb = h * chunks_per_head + c
488
+ start = qb * QUANT_GROUP_SIZE
489
+ end = start + QUANT_GROUP_SIZE
490
+
491
+ fused_nope = fused_fp8[:, :, start:end].view(torch.uint8)
492
+ norope_nope = norope_fp8[:, :, start:end].view(torch.uint8)
493
+ assert torch.equal(fused_nope, norope_nope), (
494
+ f"Nope block (head={h}, chunk={c}) differs between "
495
+ f"fused and no-rope reference"
496
+ )
497
+
498
+ fused_s = fused_scale[:, :, qb]
499
+ norope_s = norope_scale[:, :, qb]
500
+ assert torch.equal(fused_s, norope_s), (
501
+ f"Nope scale (head={h}, chunk={c}) differs"
502
+ )
503
+
504
+
505
+@torch.inference_mode()
506
+def test_single_token():
507
+ """Edge case: single token."""
508
+ num_tokens, num_heads, n_groups = 1, 64, 8
509
+ heads_per_group = num_heads // n_groups
510
+ max_pos = 4096
511
+ device = "cuda"
512
+
513
+ o = torch.randn(
514
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
515
+ )
516
+ positions = torch.tensor([42], device=device, dtype=torch.long)
517
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
518
+
519
+ ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
520
+ o.clone(),
521
+ positions,
522
+ cos_sin_cache,
523
+ n_groups,
524
+ heads_per_group,
525
+ )
526
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
527
+ o.clone(),
528
+ positions,
529
+ cos_sin_cache,
530
+ n_groups,
531
+ heads_per_group,
532
+ )
533
+
534
+ assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale)
535
+
536
+
537
+@torch.inference_mode()
538
+def test_zero_positions():
539
+ """Edge case: all positions are 0."""
540
+ num_tokens, num_heads, n_groups = 16, 64, 8
541
+ heads_per_group = num_heads // n_groups
542
+ max_pos = 4096
543
+ device = "cuda"
544
+
545
+ o = torch.randn(
546
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
547
+ )
548
+ positions = torch.zeros(num_tokens, device=device, dtype=torch.long)
549
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
550
+
551
+ ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
552
+ o.clone(),
553
+ positions,
554
+ cos_sin_cache,
555
+ n_groups,
556
+ heads_per_group,
557
+ )
558
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
559
+ o.clone(),
560
+ positions,
561
+ cos_sin_cache,
562
+ n_groups,
563
+ heads_per_group,
564
+ )
565
+
566
+ assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale)
567
+
568
+
569
+@torch.inference_mode()
570
+def test_large_values():
571
+ """Edge case: values near FP8 saturation to test clamping."""
572
+ num_tokens, num_heads, n_groups = 8, 64, 8
573
+ heads_per_group = num_heads // n_groups
574
+ max_pos = 4096
575
+ device = "cuda"
576
+
577
+ # Create inputs with large values that will saturate FP8
578
+ o = torch.randn(
579
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
580
+ )
581
+ o = o * 1000.0 # scale up to force saturation
582
+ positions = torch.randint(
583
+ 0, max_pos, (num_tokens,), device=device, dtype=torch.long
584
+ )
585
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
586
+
587
+ ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
588
+ o.clone(),
589
+ positions,
590
+ cos_sin_cache,
591
+ n_groups,
592
+ heads_per_group,
593
+ )
594
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
595
+ o.clone(),
596
+ positions,
597
+ cos_sin_cache,
598
+ n_groups,
599
+ heads_per_group,
600
+ )
601
+
602
+ assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale)
603
+
604
+
605
+@torch.inference_mode()
606
+def test_dequant_numerical_accuracy():
607
+ """Verify dequantized values are close to the original (after inv RoPE)."""
608
+ num_tokens, num_heads, n_groups = 32, 64, 8
609
+ heads_per_group = num_heads // n_groups
610
+ max_pos = 4096
611
+ device = "cuda"
612
+ torch.manual_seed(0)
613
+
614
+ o = torch.randn(
615
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
616
+ )
617
+ positions = torch.randint(
618
+ 0, max_pos, (num_tokens,), device=device, dtype=torch.long
619
+ )
620
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
621
+
622
+ # Get the post-inv-RoPE values (ground truth before quantization)
623
+ o_after_rope = reference_inv_rope(o.clone(), positions, cos_sin_cache)
624
+ d = heads_per_group * HEAD_DIM
625
+ o_after_rope = o_after_rope.view(num_tokens, n_groups, d)
626
+
627
+ # Get fused quantized output
628
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
629
+ o.clone(),
630
+ positions,
631
+ cos_sin_cache,
632
+ n_groups,
633
+ heads_per_group,
634
+ )
635
+
636
+ # Dequantize: broadcast scale [T, G, S] to [T, G, D] via repeat
637
+ S = d // QUANT_GROUP_SIZE
638
+ scale_expanded = (
639
+ fused_scale.unsqueeze(-1)
640
+ .expand(num_tokens, n_groups, S, QUANT_GROUP_SIZE)
641
+ .reshape(num_tokens, n_groups, d)
642
+ )
643
+ dequant = fused_fp8.float() * scale_expanded
644
+
645
+ # Check relative error.
646
+ # FP8 e4m3 with UE8M0 (power-of-two scales that round UP) quantizes more
647
+ # coarsely than optimal scaling. Both paths rotate in fp32, so the bulk
648
+ # of the error comes from UE8M0 quantization itself (~10-12% typical).
649
+ o_gt = o_after_rope.transpose(0, 1).contiguous().transpose(0, 1)
650
+ dequant_contig = dequant.transpose(0, 1).contiguous().transpose(0, 1)
651
+
652
+ abs_err = (dequant_contig.float() - o_gt.float()).abs()
653
+ rel_err = abs_err / (o_gt.float().abs().clamp(min=1e-6))
654
+ mean_rel_err = rel_err.mean().item()
655
+
656
+ assert mean_rel_err < 0.15, (
657
+ f"Mean relative error too high: {mean_rel_err:.4f} (expected < 0.15)"
658
+ )
659
+
660
+
661
+def _unfused_inv_rope_fp8_quant(
662
+ o: torch.Tensor,
663
+ positions: torch.Tensor,
664
+ cos_sin_cache: torch.Tensor,
665
+ n_groups: int,
666
+ heads_per_group: int,
667
+ nope_dim: int = NOPE_DIM,
668
+ rope_dim: int = ROPE_DIM,
669
+) -> tuple[torch.Tensor, torch.Tensor]:
670
+ """Unfused path matching deepseek_v4_attention.py:295-310.
671
+
672
+ Uses the production CUDA RoPE kernel + per_token_group_quant_fp8.
673
+ """
674
+ from vllm import _custom_ops as ops
675
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
676
+ per_token_group_quant_fp8,
677
+ )
678
+
679
+ head_dim = o.shape[-1]
680
+ rope_dim_offset = head_dim - rope_dim
681
+
682
+ # Step 1: In-place CUDA RoPE (same as production)
683
+ ops.rotary_embedding(
684
+ positions,
685
+ o,
686
+ None,
687
+ head_dim,
688
+ cos_sin_cache,
689
+ False, # is_neox=False for DeepseekV4 (GPT-J style)
690
+ rope_dim_offset=rope_dim_offset,
691
+ inverse=True,
692
+ )
693
+
694
+ # Step 2: Reshape + quant + reshape (same as production)
695
+ T = o.shape[0]
696
+ d = heads_per_group * head_dim
697
+ o = o.view(T, n_groups, -1)
698
+ o_flat = o.transpose(0, 1).contiguous().reshape(-1, d)
699
+ o_fp8, o_scale = per_token_group_quant_fp8(
700
+ o_flat,
701
+ group_size=QUANT_GROUP_SIZE,
702
+ use_ue8m0=True,
703
+ )
704
+ o_fp8 = o_fp8.view(n_groups, T, d).transpose(0, 1)
705
+ o_scale = o_scale.view(n_groups, T, -1).transpose(0, 1)
706
+ return o_fp8, o_scale
707
+
708
+
709
+# =========================================================================
710
+# End-to-end test including fp8_einsum
711
+# =========================================================================
712
+
713
+
714
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128, 1024])
715
+@pytest.mark.parametrize(
716
+ "num_heads,n_groups",
717
+ [(64, 8)],
718
+ ids=["H64_G8"],
719
+)
720
+@torch.inference_mode()
721
+def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
722
+ """End-to-end: fused inv_rope+quant → fp8_einsum must match
723
+ unfused CUDA_rope+quant → fp8_einsum bitwise.
724
+
725
+ This catches stride/layout bugs that only manifest when the einsum
726
+ kernel actually consumes the quantized activations.
727
+ """
728
+ from deep_gemm.utils.math import ceil_div
729
+
730
+ from vllm.utils.deep_gemm import (
731
+ fp8_einsum,
732
+ per_block_cast_to_fp8,
733
+ transform_sf_into_required_layout,
734
+ )
735
+
736
+ heads_per_group = num_heads // n_groups
737
+ d = heads_per_group * HEAD_DIM
738
+ o_lora_rank = 1024
739
+ max_pos = 4096
740
+ device = "cuda"
741
+ torch.manual_seed(0)
742
+
743
+ o = torch.randn(
744
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
745
+ )
746
+ positions = torch.randint(
747
+ 0, max_pos, (num_tokens,), device=device, dtype=torch.long
748
+ )
749
+ cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
750
+
751
+ # -- Weight quantization (shared between both paths) --
752
+ w = torch.randn(n_groups, o_lora_rank, d, device=device, dtype=torch.bfloat16)
753
+ w_fp8 = torch.empty_like(w, dtype=torch.float8_e4m3fn)
754
+ w_scale = torch.empty(
755
+ n_groups,
756
+ ceil_div(o_lora_rank, 128),
757
+ ceil_div(d, 128),
758
+ device=device,
759
+ dtype=torch.float32,
760
+ )
761
+ for g in range(n_groups):
762
+ w_fp8[g], w_scale[g] = per_block_cast_to_fp8(w[g], use_ue8m0=True)
763
+
764
+ recipe = (1, 1, 128)
765
+ w_scale_t = transform_sf_into_required_layout(
766
+ sf=w_scale,
767
+ mn=o_lora_rank,
768
+ k=d,
769
+ recipe=(1, 128, 128),
770
+ num_groups=n_groups,
771
+ is_sfa=False,
772
+ )
773
+
774
+ # -- UNFUSED path --
775
+ ref_fp8, ref_scale = _unfused_inv_rope_fp8_quant(
776
+ o.clone(),
777
+ positions,
778
+ cos_sin_cache,
779
+ n_groups,
780
+ heads_per_group,
781
+ )
782
+ z_ref = torch.empty(
783
+ num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16
784
+ )
785
+ fp8_einsum(
786
+ "bhr,hdr->bhd", (ref_fp8, ref_scale), (w_fp8, w_scale_t), z_ref, recipe=recipe
787
+ )
788
+
789
+ # -- FUSED path --
790
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
791
+ o.clone(),
792
+ positions,
793
+ cos_sin_cache,
794
+ n_groups,
795
+ heads_per_group,
796
+ )
797
+ z_fused = torch.empty(
798
+ num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16
799
+ )
800
+ fp8_einsum(
801
+ "bhr,hdr->bhd",
802
+ (fused_fp8, fused_scale),
803
+ (w_fp8, w_scale_t),
804
+ z_fused,
805
+ recipe=recipe,
806
+ )
807
+
808
+ # -- Checks --
809
+ # Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
810
+ # come from fp32 ordering and UE8M0 boundary shifts only.
811
+ # Use relative diff (same metric as test_fp8_einsum.py).
812
+ from deep_gemm.testing import calc_diff
813
+
814
+ z_diff = calc_diff(z_fused, z_ref)
815
+ assert z_diff < 0.01, (
816
+ f"Einsum output diff too large: {z_diff:.6f} (expected < 0.01)"
817
+ )
818
+
819
+
820
+@pytest.mark.parametrize("num_tokens", [1, 32, 256])
821
+@torch.inference_mode()
822
+def test_with_real_deepseek_v4_rope(num_tokens, default_vllm_config):
823
+ """Test with real DeepseekV4ScalingRotaryEmbedding (GPT-J style,
824
+ mscale=0, YaRN scaling) matching the production config."""
825
+
826
+ num_heads = 64
827
+ n_groups = 8
828
+ heads_per_group = num_heads // n_groups
829
+ device = "cuda"
830
+ torch.manual_seed(0)
831
+
832
+ # Build YaRN-scaled cos_sin_cache matching real DeepSeek V3/V4 config
833
+ # (mscale=0 → mscale=1.0, so no magnitude scaling)
834
+ from vllm.model_executor.layers.rotary_embedding.common import (
835
+ yarn_find_correction_range,
836
+ yarn_linear_ramp_mask,
837
+ )
838
+
839
+ scaling_factor = 16
840
+ base = 10000.0
841
+ max_pos = 65536
842
+ beta_fast, beta_slow = 32, 1
843
+
844
+ pos_freqs = base ** (
845
+ torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) / ROPE_DIM
846
+ )
847
+ inv_freq_extra = 1.0 / pos_freqs
848
+ inv_freq_interp = 1.0 / (scaling_factor * pos_freqs)
849
+ low, high = yarn_find_correction_range(
850
+ beta_fast, beta_slow, ROPE_DIM, base, max_pos
851
+ )
852
+ mask = 1 - yarn_linear_ramp_mask(low, high, ROPE_DIM // 2, dtype=torch.float32).to(
853
+ device
854
+ )
855
+ inv_freq = inv_freq_interp * (1 - mask) + inv_freq_extra * mask
856
+ t = torch.arange(max_pos * scaling_factor, device=device, dtype=torch.float32)
857
+ freqs = torch.outer(t, inv_freq)
858
+ # mscale=0 → yarn_get_mscale returns 1.0
859
+ cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1) # fp32
860
+
861
+ o = torch.randn(
862
+ num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
863
+ )
864
+ positions = torch.randint(0, 4096, (num_tokens,), device=device, dtype=torch.long)
865
+
866
+ # UNFUSED: CUDA RoPE with is_neox=False (GPT-J)
867
+ from vllm import _custom_ops as ops
868
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
869
+ per_token_group_quant_fp8,
870
+ )
871
+
872
+ o_unfused = o.clone()
873
+ ops.rotary_embedding(
874
+ positions,
875
+ o_unfused,
876
+ None,
877
+ HEAD_DIM,
878
+ cos_sin_cache,
879
+ False, # is_neox=False (GPT-J style)
880
+ rope_dim_offset=NOPE_DIM,
881
+ inverse=True,
882
+ )
883
+ d = heads_per_group * HEAD_DIM
884
+ T = num_tokens
885
+ o_unfused = o_unfused.view(T, n_groups, d)
886
+ o_flat = o_unfused.transpose(0, 1).contiguous().reshape(-1, d)
887
+ ref_fp8, ref_scale = per_token_group_quant_fp8(
888
+ o_flat,
889
+ group_size=QUANT_GROUP_SIZE,
890
+ use_ue8m0=True,
891
+ )
892
+ ref_fp8 = ref_fp8.view(n_groups, T, d).transpose(0, 1)
893
+ ref_scale = ref_scale.view(n_groups, T, -1).transpose(0, 1)
894
+
895
+ # FUSED: use the real YaRN-scaled cos_sin_cache
896
+ fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
897
+ o.clone(),
898
+ positions,
899
+ cos_sin_cache,
900
+ n_groups,
901
+ heads_per_group,
902
+ )
903
+
904
+ # Scales must match exactly (same UE8M0 algorithm)
905
+ # Compare via dequant (Triton bf16 rotation may differ from CUDA by 1 ULP)
906
+ assert_dequant_close(
907
+ ref_fp8, ref_scale, fused_fp8, fused_scale, msg="Real DeepSeek V4 rope"
908
+ )
tests/kernels/test_top_k_per_row.py
+5
-10
@@ -718,7 +718,6 @@ def test_persistent_topk_stress() -> None:
718
pytest.param(
719
{
720
"seq_lens": [2000, 6000, 30000, 80000],
721
- "top_k": 2048,
721
"data_type": "random",
722
},
723
id="mixed_all_paths",
@@ -727,7 +726,6 @@ def test_persistent_topk_stress() -> None:
726
pytest.param(
727
{
728
"seq_lens": [2048, 4096, 8192, 16000],
730
- "top_k": 2048,
729
"data_type": "random",
730
},
731
id="all_decode_medium",
@@ -736,7 +734,6 @@ def test_persistent_topk_stress() -> None:
734
pytest.param(
735
{
736
"seq_lens": [70000, 100000, 163840],
739
- "top_k": 2048,
737
"data_type": "random",
738
},
739
id="all_large",
@@ -745,7 +742,6 @@ def test_persistent_topk_stress() -> None:
742
pytest.param(
743
{
744
"seq_lens": [32767, 32768, 32769, 32772],
748
- "top_k": 2048,
745
"data_type": "random",
746
},
747
id="large_threshold_boundary",
@@ -754,7 +750,6 @@ def test_persistent_topk_stress() -> None:
750
pytest.param(
751
{
752
"seq_lens": [5000],
757
- "top_k": 2048,
753
"data_type": "random",
754
},
755
id="single_row_medium",
@@ -772,15 +767,15 @@ def test_persistent_topk_stress() -> None:
767
pytest.param(
768
{
769
"seq_lens": [100, 2048, 10000, 80000],
775
- "top_k": 2048,
770
"data_type": "random",
771
},
772
id="trivial_medium_large_mix",
773
),
774
],
775
)
776
+@pytest.mark.parametrize("top_k", [512, 2048])
777
@torch.inference_mode()
783
-def test_persistent_topk(test_config: dict) -> None:
778
+def test_persistent_topk(test_config: dict, top_k: int) -> None:
779
"""
780
Tests specific to the persistent_topk kernel:
781
- Mixed medium/large rows in the same batch (dynamic per-row dispatch)
@@ -790,14 +785,15 @@ def test_persistent_topk(test_config: dict) -> None:
785
run_large_context_topk_test(
786
batch_size=len(test_config["seq_lens"]),
787
seq_lens=test_config["seq_lens"],
793
- top_k=test_config["top_k"],
788
+ top_k=top_k,
789
data_type=test_config.get("data_type", "random"),
790
)
791
792
793
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
794
+@pytest.mark.parametrize("top_k", [512, 2048])
795
@torch.inference_mode()
800
-def test_persistent_topk_padded_stride() -> None:
796
+def test_persistent_topk_padded_stride(top_k: int) -> None:
797
"""
798
Test persistent_topk with padded logits (large stride, small seq_len)
799
to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits
@@ -806,7 +802,6 @@ def test_persistent_topk_padded_stride() -> None:
802
set_random_seed(42)
803
torch.set_default_device("cuda:0")
804
809
- top_k = 2048
805
batch_size = 4
806
padded_stride = 163840 # DeepSeek-V3.2 max_model_len
807
actual_seq_lens = [3000, 5000, 8000, 12000]
tests/model_executor/test_routed_experts_capture.py
+3
-1
@@ -41,7 +41,9 @@ class DummyRouter(BaseRouter):
41
def routing_method_type(self) -> RoutingMethodType:
42
return RoutingMethodType.FUSED_TOPK
43
44
- def _compute_routing(self, hidden_states, router_logits, indices_type):
44
+ def _compute_routing(
45
+ self, hidden_states, router_logits, indices_type, *, input_ids=None
46
+ ):
47
topk_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
48
topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
49
return topk_weights, topk_ids
tests/models/registry.py
+9
@@ -260,6 +260,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
260
trust_remote_code=True,
261
),
262
"DeepseekV32ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3.2-Exp"),
263
+ "DeepseekV4ForCausalLM": _HfExamplesInfo(
264
+ "deepseek-ai/DeepSeek-V4-Flash", is_available_online=False
265
+ ),
266
"Ernie4_5ForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-0.3B-PT"),
267
"Ernie4_5_MoeForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT"),
268
"ExaoneForCausalLM": _HfExamplesInfo(
@@ -1482,6 +1485,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1485
speculative_model="luccafong/deepseek_mtp_draft_random",
1486
trust_remote_code=True,
1487
),
1488
+ "DeepSeekV4MTPModel": _HfExamplesInfo(
1489
+ "deepseek-ai/DeepSeek-V4-Flash",
1490
+ speculative_model="deepseek-ai/DeepSeek-V4-Flash",
1491
+ trust_remote_code=True,
1492
+ is_available_online=False,
1493
+ ),
1494
"ErnieMTPModel": _HfExamplesInfo(
1495
"baidu/ERNIE-4.5-21B-A3B-PT",
1496
trust_remote_code=True,
tests/models/test_deepseek_v4_mega_moe.py
new
+184
@@ -0,0 +1,184 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.model_executor.models.deepseek_v4 import (
10
+ DeepseekV4MegaMoEExperts,
11
+ _stage_deepseek_v4_mega_moe_inputs,
12
+ make_deepseek_v4_expert_params_mapping,
13
+)
14
+from vllm.platforms import current_platform
15
+
16
+pytestmark = pytest.mark.skipif(
17
+ not current_platform.is_cuda(),
18
+ reason="DeepSeek V4 MegaMoE requires CUDA",
19
+)
20
+
21
+
22
+def test_deepseek_v4_mega_moe_expert_mapping():
23
+ mapping = make_deepseek_v4_expert_params_mapping(2)
24
+
25
+ assert mapping == [
26
+ ("experts.w13_", "experts.0.w1.", 0, "w1"),
27
+ ("experts.w2_", "experts.0.w2.", 0, "w2"),
28
+ ("experts.w13_", "experts.0.w3.", 0, "w3"),
29
+ ("experts.w13_", "experts.1.w1.", 1, "w1"),
30
+ ("experts.w2_", "experts.1.w2.", 1, "w2"),
31
+ ("experts.w13_", "experts.1.w3.", 1, "w3"),
32
+ ]
33
+
34
+
35
+def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float():
36
+ raw = torch.tensor([0, 126, 127, 128], dtype=torch.uint8)
37
+
38
+ decoded = DeepseekV4MegaMoEExperts._ue8m0_uint8_to_float(raw)
39
+
40
+ assert torch.equal(decoded.view(torch.int32), raw.to(torch.int32) << 23)
41
+ assert decoded[0].item() == 0.0
42
+ assert decoded[1].item() == 0.5
43
+ assert decoded[2].item() == 1.0
44
+ assert decoded[3].item() == 2.0
45
+
46
+
47
+def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
48
+ vllm_config = SimpleNamespace(
49
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=4)
50
+ )
51
+ experts = DeepseekV4MegaMoEExperts(
52
+ vllm_config,
53
+ num_experts=4,
54
+ num_local_experts=2,
55
+ experts_start_idx=2,
56
+ top_k=2,
57
+ hidden_size=128,
58
+ intermediate_size=128,
59
+ )
60
+
61
+ nonlocal_weight = torch.ones(128, 64, dtype=torch.uint8)
62
+ assert (
63
+ experts.weight_loader(
64
+ experts.w13_weight,
65
+ nonlocal_weight,
66
+ "experts.w13_weight",
67
+ shard_id="w1",
68
+ expert_id=1,
69
+ return_success=True,
70
+ )
71
+ is False
72
+ )
73
+
74
+ w1 = torch.full((128, 64), 3, dtype=torch.uint8)
75
+ w3 = torch.full((128, 64), 7, dtype=torch.uint8)
76
+ w2 = torch.full((128, 64), 11, dtype=torch.uint8)
77
+
78
+ assert experts.weight_loader(
79
+ experts.w13_weight,
80
+ w1,
81
+ "experts.w13_weight",
82
+ shard_id="w1",
83
+ expert_id=2,
84
+ return_success=True,
85
+ )
86
+ assert experts.weight_loader(
87
+ experts.w13_weight,
88
+ w3,
89
+ "experts.w13_weight",
90
+ shard_id="w3",
91
+ expert_id=2,
92
+ return_success=True,
93
+ )
94
+ assert experts.weight_loader(
95
+ experts.w2_weight,
96
+ w2,
97
+ "experts.w2_weight",
98
+ shard_id="w2",
99
+ expert_id=2,
100
+ return_success=True,
101
+ )
102
+
103
+ assert torch.equal(experts.w13_weight[0, :128], w1)
104
+ assert torch.equal(experts.w13_weight[0, 128:], w3)
105
+ assert torch.equal(experts.w2_weight[0], w2)
106
+ assert torch.count_nonzero(experts.w13_weight[1]) == 0
107
+
108
+
109
+@pytest.mark.skipif(
110
+ not torch.cuda.is_available(),
111
+ reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
112
+)
113
+def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
114
+ from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
115
+
116
+ device = torch.device("cuda")
117
+ num_tokens = 7
118
+ hidden_size = 256
119
+ top_k = 8
120
+
121
+ generator = torch.Generator(device=device)
122
+ generator.manual_seed(0)
123
+ hidden_states = (
124
+ torch.randn(
125
+ num_tokens,
126
+ hidden_size,
127
+ device=device,
128
+ dtype=torch.float32,
129
+ generator=generator,
130
+ )
131
+ * 17.0
132
+ ).to(torch.bfloat16)
133
+ hidden_states[0, :32] = 0
134
+ hidden_states[1, 32:64] = 1.0e-6
135
+ hidden_states[2, 64:96] = -1.0e-6
136
+
137
+ topk_ids = torch.randint(
138
+ 0,
139
+ 256,
140
+ (num_tokens, top_k),
141
+ device=device,
142
+ dtype=torch.int32,
143
+ generator=generator,
144
+ )
145
+ topk_weights = torch.randn(
146
+ num_tokens,
147
+ top_k,
148
+ device=device,
149
+ dtype=torch.float32,
150
+ generator=generator,
151
+ )
152
+
153
+ ref_x, ref_x_sf = per_token_cast_to_fp8(
154
+ hidden_states,
155
+ use_ue8m0=True,
156
+ gran_k=32,
157
+ use_packed_ue8m0=True,
158
+ )
159
+ ref_topk_idx = topk_ids.to(torch.int64)
160
+ ref_topk_weights = topk_weights.clone()
161
+
162
+ fused_x = torch.empty_like(ref_x)
163
+ fused_x_sf = torch.empty_like(ref_x_sf)
164
+ fused_topk_idx = torch.empty_like(ref_topk_idx)
165
+ fused_topk_weights = torch.empty_like(ref_topk_weights)
166
+
167
+ _stage_deepseek_v4_mega_moe_inputs(
168
+ hidden_states,
169
+ topk_weights,
170
+ topk_ids,
171
+ fused_x,
172
+ fused_x_sf,
173
+ fused_topk_idx,
174
+ fused_topk_weights,
175
+ )
176
+ torch.accelerator.synchronize()
177
+
178
+ assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
179
+ assert torch.equal(fused_x_sf, ref_x_sf)
180
+ assert torch.equal(fused_topk_idx, ref_topk_idx)
181
+ assert torch.equal(
182
+ fused_topk_weights.view(torch.uint8),
183
+ ref_topk_weights.view(torch.uint8),
184
+ )
tests/reasoning/test_deepseekv3_reasoning_parser.py
+7
@@ -6,6 +6,7 @@ from transformers import AutoTokenizer
6
7
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
8
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
9
+from vllm.reasoning import ReasoningParserManager
10
from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
11
from vllm.reasoning.deepseek_v3_reasoning_parser import DeepSeekV3ReasoningParser
12
from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser
@@ -33,6 +34,12 @@ def test_parser_selection(tokenizer, thinking, expected_parser_type):
34
assert isinstance(parser._parser, expected_parser_type)
35
36
37
+def test_deepseek_v4_reasoning_parser_alias():
38
+ parser_cls = ReasoningParserManager.get_reasoning_parser("deepseek_v4")
39
+
40
+ assert parser_cls is DeepSeekV3ReasoningParser
41
+
42
+
43
def test_identity_reasoning_parser_basic(tokenizer):
44
parser = IdentityReasoningParser(tokenizer)
45
tests/tokenizers_/fixtures/deepseek_v4/test_input_1.json
new
+81
@@ -0,0 +1,81 @@
1
+{
2
+ "tools": [
3
+ {
4
+ "type": "function",
5
+ "function": {
6
+ "name": "get_weather",
7
+ "description": "Get the weather for a specific location",
8
+ "parameters": {
9
+ "type": "object",
10
+ "properties": {
11
+ "location": {
12
+ "type": "string",
13
+ "description": "The city name"
14
+ },
15
+ "unit": {
16
+ "type": "string",
17
+ "enum": ["celsius", "fahrenheit"],
18
+ "description": "Temperature unit"
19
+ }
20
+ },
21
+ "required": ["location"]
22
+ }
23
+ }
24
+ },
25
+ {
26
+ "type": "function",
27
+ "function": {
28
+ "name": "search",
29
+ "description": "Search the web for information",
30
+ "parameters": {
31
+ "type": "object",
32
+ "properties": {
33
+ "query": {
34
+ "type": "string",
35
+ "description": "Search query"
36
+ },
37
+ "num_results": {
38
+ "type": "integer",
39
+ "description": "Number of results to return"
40
+ }
41
+ },
42
+ "required": ["query"]
43
+ }
44
+ }
45
+ }
46
+ ],
47
+ "messages": [
48
+ {
49
+ "role": "system",
50
+ "content": "You are a helpful assistant."
51
+ },
52
+ {
53
+ "role": "user",
54
+ "content": "What's the weather in Beijing?"
55
+ },
56
+ {
57
+ "role": "assistant",
58
+ "reasoning": "The user wants to know the weather in Beijing. I should use the get_weather tool.",
59
+ "tool_calls": [
60
+ {
61
+ "id": "call_001",
62
+ "type": "function",
63
+ "function": {
64
+ "name": "get_weather",
65
+ "arguments": "{\"location\": \"Beijing\", \"unit\": \"celsius\"}"
66
+ }
67
+ }
68
+ ]
69
+ },
70
+ {
71
+ "role": "tool",
72
+ "tool_call_id": "call_001",
73
+ "content": "{\"temperature\": 22, \"condition\": \"sunny\", \"humidity\": 45}"
74
+ },
75
+ {
76
+ "role": "assistant",
77
+ "reasoning": "Got the weather data. Let me format a nice response.",
78
+ "content": "The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity."
79
+ }
80
+ ]
81
+}
tests/tokenizers_/fixtures/deepseek_v4/test_input_2.json
new
+24
@@ -0,0 +1,24 @@
1
+[
2
+ {
3
+ "role": "system",
4
+ "content": "You are a helpful assistant."
5
+ },
6
+ {
7
+ "role": "user",
8
+ "content": "Hello"
9
+ },
10
+ {
11
+ "role": "assistant",
12
+ "reasoning": "The user said hello, I should greet back.",
13
+ "content": "Hi there! How can I help you?"
14
+ },
15
+ {
16
+ "role": "user",
17
+ "content": "What is the capital of France?"
18
+ },
19
+ {
20
+ "role": "assistant",
21
+ "reasoning": "The user asks about the capital of France. It is Paris.",
22
+ "content": "The capital of France is Paris."
23
+ }
24
+]
\ No newline at end of file
tests/tokenizers_/fixtures/deepseek_v4/test_input_3.json
new
+159
@@ -0,0 +1,159 @@
1
+[
2
+ {
3
+ "role": "system",
4
+ "content": "该助手为DeepSeek，由深度求索公司创造。"
5
+ },
6
+ {
7
+ "role": "latest_reminder",
8
+ "content": "2026-02-21,星期六,广州,App,中文"
9
+ },
10
+ {
11
+ "role": "developer",
12
+ "content": "小柴胡冲剂和布洛芬能一起吃吗？\n\nCITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】",
13
+ "tools": [
14
+ {
15
+ "type": "function",
16
+ "function": {
17
+ "name": "search",
18
+ "description": "Web search. Split multiple queries with '||'.",
19
+ "parameters": {
20
+ "type": "object",
21
+ "properties": {
22
+ "queries": {
23
+ "type": "string",
24
+ "description": "query1||query2"
25
+ }
26
+ },
27
+ "required": [
28
+ "queries"
29
+ ],
30
+ "additionalProperties": false,
31
+ "$schema": "http://json-schema.org/draft-07/schema#"
32
+ }
33
+ }
34
+ },
35
+ {
36
+ "type": "function",
37
+ "function": {
38
+ "name": "open",
39
+ "description": "Batch open IDs (format 【{id}†...】) or URLs.",
40
+ "parameters": {
41
+ "type": "object",
42
+ "properties": {
43
+ "open_list": {
44
+ "type": "array",
45
+ "items": {
46
+ "type": "object",
47
+ "properties": {
48
+ "id": {
49
+ "description": "ID or URL",
50
+ "anyOf": [
51
+ {
52
+ "type": "integer"
53
+ },
54
+ {
55
+ "type": "string"
56
+ }
57
+ ],
58
+ "default": -1
59
+ },
60
+ "cursor": {
61
+ "type": "integer",
62
+ "description": "",
63
+ "default": -1
64
+ },
65
+ "loc": {
66
+ "type": "integer",
67
+ "description": "Start line",
68
+ "default": -1
69
+ },
70
+ "num_lines": {
71
+ "type": "integer",
72
+ "description": "",
73
+ "default": -1
74
+ },
75
+ "view_source": {
76
+ "type": "boolean",
77
+ "description": "",
78
+ "default": false
79
+ }
80
+ },
81
+ "additionalProperties": false
82
+ },
83
+ "description": ""
84
+ }
85
+ },
86
+ "required": [
87
+ "open_list"
88
+ ],
89
+ "additionalProperties": false,
90
+ "$schema": "http://json-schema.org/draft-07/schema#"
91
+ }
92
+ }
93
+ },
94
+ {
95
+ "type": "function",
96
+ "function": {
97
+ "name": "find",
98
+ "description": "Find exact text pattern in pages.",
99
+ "parameters": {
100
+ "type": "object",
101
+ "properties": {
102
+ "find_list": {
103
+ "type": "array",
104
+ "items": {
105
+ "type": "object",
106
+ "properties": {
107
+ "pattern": {
108
+ "type": "string",
109
+ "description": ""
110
+ },
111
+ "cursor": {
112
+ "type": "integer",
113
+ "description": "",
114
+ "default": -1
115
+ }
116
+ },
117
+ "required": [
118
+ "pattern"
119
+ ],
120
+ "additionalProperties": false
121
+ },
122
+ "description": ""
123
+ }
124
+ },
125
+ "required": [
126
+ "find_list"
127
+ ],
128
+ "additionalProperties": false,
129
+ "$schema": "http://json-schema.org/draft-07/schema#"
130
+ }
131
+ }
132
+ }
133
+ ]
134
+ },
135
+ {
136
+ "role": "assistant",
137
+ "content": "",
138
+ "reasoning": "用户想知道小柴胡冲剂和布洛芬能否一起服用。",
139
+ "tool_calls": [
140
+ {
141
+ "type": "function",
142
+ "function": {
143
+ "name": "search",
144
+ "arguments": "{\"queries\": \"小柴胡冲剂 布洛芬 相互作用 一起吃\"}"
145
+ }
146
+ }
147
+ ]
148
+ },
149
+ {
150
+ "role": "tool",
151
+ "content": "[0]"
152
+ },
153
+ {
154
+ "role": "assistant",
155
+ "content": "请及时就医。",
156
+ "reasoning": "现在开始组织回答。",
157
+ "tool_calls": []
158
+ }
159
+]
\ No newline at end of file
tests/tokenizers_/fixtures/deepseek_v4/test_input_4.json
new
+28
@@ -0,0 +1,28 @@
1
+[
2
+ {
3
+ "role": "system",
4
+ "content": "该助手为DeepSeek-V3，由深度求索公司创造。\n今天是2025年10月17日，星期五。"
5
+ },
6
+ {
7
+ "role": "latest_reminder",
8
+ "content": "2024-11-15,上海市,App,中文"
9
+ },
10
+ {
11
+ "role": "user",
12
+ "content": "热海大滚锅是世界著名温泉吗"
13
+ },
14
+ {
15
+ "role": "assistant",
16
+ "content": "关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。\n\n为了让你快速了解，我整理了一个简要的表格：\n\n| 维度 | 热海大滚锅的地位与特点 |\n| :--- | :--- |\n| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |\n| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |\n| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |\n| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |\n\n### 💡 游玩攻略与温馨提示\n\n如果你计划前往热海大滚锅，这里有一些实用信息供你参考：\n\n- **门票与开放时间**：\n - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。\n - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。\n\n- **特色体验**：\n - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。\n - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。\n\n- **注意事项**：\n - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。\n - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。\n\n希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。",
17
+ "mask": 1
18
+ },
19
+ {
20
+ "role": "user",
21
+ "content": "世界著名温泉有哪些",
22
+ "task": "action"
23
+ },
24
+ {
25
+ "role": "assistant",
26
+ "content": "Search"
27
+ }
28
+]
\ No newline at end of file
tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt
new
+36
@@ -0,0 +1,36 @@
1
+<｜begin▁of▁sentence｜>
2
+
3
+## Tools
4
+
5
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
6
+
7
+<｜DSML｜tool_calls>
8
+<｜DSML｜invoke name="$TOOL_NAME">
9
+<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
10
+...
11
+</｜DSML｜invoke>
12
+<｜DSML｜invoke name="$TOOL_NAME2">
13
+...
14
+</｜DSML｜invoke>
15
+</｜DSML｜tool_calls>
16
+
17
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
18
+
19
+If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
20
+
21
+Otherwise, output directly after </think> with tool calls or final response.
22
+
23
+### Available Tool Schemas
24
+
25
+{"name": "get_weather", "description": "Get the weather for a specific location", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city name"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "Temperature unit"}}, "required": ["location"]}}
26
+{"name": "search", "description": "Search the web for information", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}, "num_results": {"type": "integer", "description": "Number of results to return"}}, "required": ["query"]}}
27
+
28
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
29
+You are a helpful assistant.<｜User｜>What's the weather in Beijing?<｜Assistant｜><think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>
30
+
31
+<｜DSML｜tool_calls>
32
+<｜DSML｜invoke name="get_weather">
33
+<｜DSML｜parameter name="location" string="true">Beijing</｜DSML｜parameter>
34
+<｜DSML｜parameter name="unit" string="true">celsius</｜DSML｜parameter>
35
+</｜DSML｜invoke>
36
+</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>{"temperature": 22, "condition": "sunny", "humidity": 45}</tool_result><｜Assistant｜><think>Got the weather data. Let me format a nice response.</think>The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.<｜end▁of▁sentence｜>
\ No newline at end of file
tests/tokenizers_/fixtures/deepseek_v4/test_output_2.txt
new
+1
@@ -0,0 +1 @@
1
+<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>Hello<｜Assistant｜></think>Hi there! How can I help you?<｜end▁of▁sentence｜><｜User｜>What is the capital of France?<｜Assistant｜><think>The user asks about the capital of France. It is Paris.</think>The capital of France is Paris.<｜end▁of▁sentence｜>
\ No newline at end of file
tests/tokenizers_/fixtures/deepseek_v4/test_output_3.txt
new
+38
@@ -0,0 +1,38 @@
1
+<｜begin▁of▁sentence｜>该助手为DeepSeek，由深度求索公司创造。<｜latest_reminder｜>2026-02-21,星期六,广州,App,中文<｜User｜>小柴胡冲剂和布洛芬能一起吃吗？
2
+
3
+CITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】
4
+
5
+## Tools
6
+
7
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
8
+
9
+<｜DSML｜tool_calls>
10
+<｜DSML｜invoke name="$TOOL_NAME">
11
+<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
12
+...
13
+</｜DSML｜invoke>
14
+<｜DSML｜invoke name="$TOOL_NAME2">
15
+...
16
+</｜DSML｜invoke>
17
+</｜DSML｜tool_calls>
18
+
19
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
20
+
21
+If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
22
+
23
+Otherwise, output directly after </think> with tool calls or final response.
24
+
25
+### Available Tool Schemas
26
+
27
+{"name": "search", "description": "Web search. Split multiple queries with '||'.", "parameters": {"type": "object", "properties": {"queries": {"type": "string", "description": "query1||query2"}}, "required": ["queries"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
28
+{"name": "open", "description": "Batch open IDs (format 【{id}†...】) or URLs.", "parameters": {"type": "object", "properties": {"open_list": {"type": "array", "items": {"type": "object", "properties": {"id": {"description": "ID or URL", "anyOf": [{"type": "integer"}, {"type": "string"}], "default": -1}, "cursor": {"type": "integer", "description": "", "default": -1}, "loc": {"type": "integer", "description": "Start line", "default": -1}, "num_lines": {"type": "integer", "description": "", "default": -1}, "view_source": {"type": "boolean", "description": "", "default": false}}, "additionalProperties": false}, "description": ""}}, "required": ["open_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
29
+{"name": "find", "description": "Find exact text pattern in pages.", "parameters": {"type": "object", "properties": {"find_list": {"type": "array", "items": {"type": "object", "properties": {"pattern": {"type": "string", "description": ""}, "cursor": {"type": "integer", "description": "", "default": -1}}, "required": ["pattern"], "additionalProperties": false}, "description": ""}}, "required": ["find_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
30
+
31
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
32
+<｜Assistant｜><think>用户想知道小柴胡冲剂和布洛芬能否一起服用。</think>
33
+
34
+<｜DSML｜tool_calls>
35
+<｜DSML｜invoke name="search">
36
+<｜DSML｜parameter name="queries" string="true">小柴胡冲剂 布洛芬 相互作用 一起吃</｜DSML｜parameter>
37
+</｜DSML｜invoke>
38
+</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>[0]</tool_result><｜Assistant｜><think>现在开始组织回答。</think>请及时就医。<｜end▁of▁sentence｜>
\ No newline at end of file
tests/tokenizers_/fixtures/deepseek_v4/test_output_4.txt
new
+29
@@ -0,0 +1,29 @@
1
+<｜begin▁of▁sentence｜>该助手为DeepSeek-V3，由深度求索公司创造。
2
+今天是2025年10月17日，星期五。<｜latest_reminder｜>2024-11-15,上海市,App,中文<｜User｜>热海大滚锅是世界著名温泉吗<｜Assistant｜></think>关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。
3
+
4
+为了让你快速了解，我整理了一个简要的表格：
5
+
6
+| 维度 | 热海大滚锅的地位与特点 |
7
+| :--- | :--- |
8
+| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |
9
+| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |
10
+| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |
11
+| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |
12
+
13
+### 💡 游玩攻略与温馨提示
14
+
15
+如果你计划前往热海大滚锅，这里有一些实用信息供你参考：
16
+
17
+- **门票与开放时间**：
18
+ - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。
19
+ - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。
20
+
21
+- **特色体验**：
22
+ - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。
23
+ - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。
24
+
25
+- **注意事项**：
26
+ - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。
27
+ - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。
28
+
29
+希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。<｜end▁of▁sentence｜><｜User｜>世界著名温泉有哪些<｜Assistant｜></think><｜action｜>Search<｜end▁of▁sentence｜>
\ No newline at end of file
tests/tokenizers_/test_deepseek_v4.py
new
+224
@@ -0,0 +1,224 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import json
5
+from pathlib import Path
6
+from types import SimpleNamespace
7
+
8
+import pytest
9
+
10
+from vllm.entrypoints.chat_utils import parse_chat_messages
11
+from vllm.renderers.registry import RENDERER_REGISTRY
12
+from vllm.tokenizers.deepseek_v4 import get_deepseek_v4_tokenizer
13
+from vllm.tokenizers.registry import TokenizerRegistry
14
+
15
+FIXTURES_DIR = Path(__file__).parent / "fixtures" / "deepseek_v4"
16
+
17
+
18
+class FakeHfTokenizer:
19
+ vocab_size = 100
20
+
21
+ def get_added_vocab(self) -> dict[str, int]:
22
+ return {"</think>": 100}
23
+
24
+ def encode(
25
+ self,
26
+ text: str,
27
+ add_special_tokens: bool = False,
28
+ **kwargs,
29
+ ) -> list[int]:
30
+ self.last_encode = (text, add_special_tokens, kwargs)
31
+ return [len(text)]
32
+
33
+
34
+def _tokenizer():
35
+ return get_deepseek_v4_tokenizer(FakeHfTokenizer())
36
+
37
+
38
+def _model_config():
39
+ return SimpleNamespace(
40
+ multimodal_config=None,
41
+ allowed_local_media_path="",
42
+ allowed_media_domains=None,
43
+ )
44
+
45
+
46
+def _load_reference_case(case_id: int):
47
+ data = json.loads((FIXTURES_DIR / f"test_input_{case_id}.json").read_text())
48
+ if isinstance(data, dict):
49
+ return data["messages"], data.get("tools")
50
+ return data, None
51
+
52
+
53
+def _render_reference_case(case_id: int, **kwargs):
54
+ messages, tools = _load_reference_case(case_id)
55
+ conversation, _, _ = parse_chat_messages(
56
+ messages,
57
+ _model_config(),
58
+ content_format="string",
59
+ )
60
+ return _tokenizer().apply_chat_template(
61
+ conversation=conversation,
62
+ messages=messages,
63
+ tools=tools,
64
+ tokenize=False,
65
+ **kwargs,
66
+ )
67
+
68
+
69
+def test_deepseek_v4_tokenizer_registered():
70
+ assert TokenizerRegistry.load_tokenizer_cls("deepseek_v4").__name__ == (
71
+ "DeepseekV4Tokenizer"
72
+ )
73
+ assert RENDERER_REGISTRY.load_renderer_cls("deepseek_v4").__name__ == (
74
+ "DeepseekV4Renderer"
75
+ )
76
+
77
+
78
+def test_deepseek_v4_defaults_to_chat_mode():
79
+ prompt = _tokenizer().apply_chat_template(
80
+ [{"role": "user", "content": "Hello"}],
81
+ tokenize=False,
82
+ )
83
+
84
+ assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜></think>")
85
+
86
+
87
+@pytest.mark.parametrize("kwargs", [{"thinking": True}, {"enable_thinking": True}])
88
+def test_deepseek_v4_enables_thinking_with_compatible_kwargs(kwargs):
89
+ prompt = _tokenizer().apply_chat_template(
90
+ [{"role": "user", "content": "Hello"}],
91
+ tokenize=False,
92
+ **kwargs,
93
+ )
94
+
95
+ assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜><think>")
96
+
97
+
98
+def test_deepseek_v4_uses_v4_tool_prompt_from_request_tools():
99
+ tools = [
100
+ {
101
+ "type": "function",
102
+ "function": {
103
+ "name": "get_weather",
104
+ "description": "Get weather for a city",
105
+ "parameters": {
106
+ "type": "object",
107
+ "properties": {"city": {"type": "string"}},
108
+ "required": ["city"],
109
+ },
110
+ },
111
+ }
112
+ ]
113
+
114
+ prompt = _tokenizer().apply_chat_template(
115
+ [{"role": "user", "content": "Weather?"}],
116
+ tools=tools,
117
+ tokenize=False,
118
+ )
119
+
120
+ assert "## Tools" in prompt
121
+ assert "<｜DSML｜tool_calls>" in prompt
122
+ assert "</｜DSML｜tool_calls>" in prompt
123
+ assert "function_calls" not in prompt
124
+ assert '"name": "get_weather"' in prompt
125
+ assert prompt.endswith("<｜User｜>Weather?<｜Assistant｜></think>")
126
+
127
+
128
+def test_deepseek_v4_renders_parsed_history_tool_arguments():
129
+ messages = [
130
+ {"role": "user", "content": "List the repo"},
131
+ {
132
+ "role": "assistant",
133
+ "tool_calls": [
134
+ {
135
+ "id": "call_1",
136
+ "type": "function",
137
+ "function": {
138
+ "name": "str_replace_editor",
139
+ "arguments": '{"command": "view", "path": "/testbed"}',
140
+ },
141
+ }
142
+ ],
143
+ },
144
+ {
145
+ "role": "tool",
146
+ "tool_call_id": "call_1",
147
+ "content": "file list",
148
+ },
149
+ ]
150
+ tools = [
151
+ {
152
+ "type": "function",
153
+ "function": {
154
+ "name": "str_replace_editor",
155
+ "description": "Edit files",
156
+ "parameters": {
157
+ "type": "object",
158
+ "properties": {
159
+ "command": {"type": "string"},
160
+ "path": {"type": "string"},
161
+ },
162
+ "required": ["command", "path"],
163
+ },
164
+ },
165
+ }
166
+ ]
167
+ conversation, _, _ = parse_chat_messages(
168
+ messages,
169
+ _model_config(),
170
+ content_format="string",
171
+ )
172
+
173
+ prompt = _tokenizer().apply_chat_template(
174
+ conversation=conversation,
175
+ messages=messages,
176
+ tools=tools,
177
+ tokenize=False,
178
+ )
179
+
180
+ assert '<｜DSML｜parameter name="command" string="true">view' in prompt
181
+ assert '<｜DSML｜parameter name="path" string="true">/testbed' in prompt
182
+ assert 'parameter name="arguments"' not in prompt
183
+
184
+
185
+@pytest.mark.parametrize("reasoning_effort", ["none", "low", "medium", "high"])
186
+def test_deepseek_v4_accepts_openai_reasoning_effort_values(reasoning_effort):
187
+ prompt = _tokenizer().apply_chat_template(
188
+ [{"role": "user", "content": "Hello"}],
189
+ tokenize=False,
190
+ enable_thinking=True,
191
+ reasoning_effort=reasoning_effort,
192
+ )
193
+
194
+ assert prompt.endswith("<｜Assistant｜><think>")
195
+ assert "Reasoning Effort: Absolute maximum" not in prompt
196
+
197
+
198
+def test_deepseek_v4_preserves_reference_max_reasoning_effort():
199
+ prompt = _tokenizer().apply_chat_template(
200
+ [{"role": "user", "content": "Hello"}],
201
+ tokenize=False,
202
+ enable_thinking=True,
203
+ reasoning_effort="max",
204
+ )
205
+
206
+ assert prompt.startswith(
207
+ "<｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum"
208
+ )
209
+
210
+
211
+@pytest.mark.parametrize(
212
+ ("case_id", "kwargs"),
213
+ [
214
+ (1, {"thinking": True}),
215
+ (2, {"thinking": True}),
216
+ (3, {"thinking": True}),
217
+ (4, {}),
218
+ ],
219
+)
220
+def test_deepseek_v4_matches_reference_golden_fixtures(case_id, kwargs):
221
+ prompt = _render_reference_case(case_id, **kwargs)
222
+
223
+ expected = (FIXTURES_DIR / f"test_output_{case_id}.txt").read_text()
224
+ assert prompt == expected
tests/tool_parsers/test_deepseekv4_tool_parser.py
new
+123
@@ -0,0 +1,123 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Unit tests for DeepSeekV4ToolParser."""
5
+
6
+import json
7
+from unittest.mock import MagicMock
8
+
9
+from vllm.tool_parsers import ToolParserManager
10
+from vllm.tool_parsers.deepseekv4_tool_parser import DeepSeekV4ToolParser
11
+
12
+MOCK_TOKENIZER = MagicMock()
13
+MOCK_TOKENIZER.get_vocab.return_value = {}
14
+
15
+TC_START = "<｜DSML｜tool_calls>"
16
+TC_END = "</｜DSML｜tool_calls>"
17
+INV_START = '<｜DSML｜invoke name="'
18
+INV_END = "</｜DSML｜invoke>"
19
+PARAM_START = '<｜DSML｜parameter name="'
20
+PARAM_END = "</｜DSML｜parameter>"
21
+
22
+
23
+def make_parser(tools=None) -> DeepSeekV4ToolParser:
24
+ return DeepSeekV4ToolParser(MOCK_TOKENIZER, tools=tools)
25
+
26
+
27
+def make_request(tools=None) -> MagicMock:
28
+ req = MagicMock()
29
+ req.tools = tools
30
+ return req
31
+
32
+
33
+def build_tool_call(func_name: str, params: dict[str, str]) -> str:
34
+ param_strs = "".join(
35
+ f'{PARAM_START}{k}" string="true">{v}{PARAM_END}\n' for k, v in params.items()
36
+ )
37
+ return f'{TC_START}\n{INV_START}{func_name}">\n{param_strs}{INV_END}\n{TC_END}'
38
+
39
+
40
+def stream(parser: DeepSeekV4ToolParser, full_text: str, chunk_size: int = 7):
41
+ deltas = []
42
+ previous_text = ""
43
+ for start in range(0, len(full_text), chunk_size):
44
+ delta_text = full_text[start : start + chunk_size]
45
+ current_text = previous_text + delta_text
46
+ delta = parser.extract_tool_calls_streaming(
47
+ previous_text=previous_text,
48
+ current_text=current_text,
49
+ delta_text=delta_text,
50
+ previous_token_ids=[],
51
+ current_token_ids=[],
52
+ delta_token_ids=[1],
53
+ request=make_request(),
54
+ )
55
+ previous_text = current_text
56
+ if delta is not None:
57
+ deltas.append(delta)
58
+ return deltas
59
+
60
+
61
+def reconstruct_args(deltas, tool_index: int = 0) -> str:
62
+ fragments = []
63
+ for delta in deltas:
64
+ if delta.tool_calls:
65
+ for tool_call in delta.tool_calls:
66
+ if (
67
+ tool_call.index == tool_index
68
+ and tool_call.function
69
+ and tool_call.function.arguments
70
+ ):
71
+ fragments.append(tool_call.function.arguments)
72
+ return "".join(fragments)
73
+
74
+
75
+def test_registered():
76
+ assert ToolParserManager.get_tool_parser("deepseek_v4") is DeepSeekV4ToolParser
77
+
78
+
79
+def test_extract_tool_calls():
80
+ parser = make_parser()
81
+ model_output = "Let me check. " + build_tool_call(
82
+ "get_weather", {"location": "Beijing", "unit": "celsius"}
83
+ )
84
+
85
+ result = parser.extract_tool_calls(model_output, make_request())
86
+
87
+ assert result.tools_called
88
+ assert result.content == "Let me check. "
89
+ assert len(result.tool_calls) == 1
90
+ tool_call = result.tool_calls[0]
91
+ assert tool_call.function.name == "get_weather"
92
+ assert json.loads(tool_call.function.arguments) == {
93
+ "location": "Beijing",
94
+ "unit": "celsius",
95
+ }
96
+
97
+
98
+def test_function_calls_block_is_not_accepted():
99
+ parser = make_parser()
100
+ model_output = build_tool_call("search", {"query": "vllm"}).replace(
101
+ "tool_calls", "function_calls"
102
+ )
103
+
104
+ result = parser.extract_tool_calls(model_output, make_request())
105
+
106
+ assert not result.tools_called
107
+ assert result.content == model_output
108
+
109
+
110
+def test_streaming_extracts_complete_invokes():
111
+ parser = make_parser()
112
+ full_text = build_tool_call("search", {"query": "deepseek v4"})
113
+
114
+ deltas = stream(parser, full_text, chunk_size=5)
115
+
116
+ names = [
117
+ tool_call.function.name
118
+ for delta in deltas
119
+ if delta.tool_calls
120
+ for tool_call in delta.tool_calls
121
+ ]
122
+ assert names == ["search"]
123
+ assert json.loads(reconstruct_args(deltas)) == {"query": "deepseek v4"}
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
new
+92
@@ -0,0 +1,92 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from tests.v1.attention.utils import create_vllm_config
8
+from vllm.v1.attention.backend import CommonAttentionMetadata
9
+from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder
10
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
11
+
12
+
13
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
14
+def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_size():
15
+ """Regression test: DeepseekV4 compression path must compute slot_mapping from
16
+ compressed positions, not reuse the uncompressed common metadata mapping.
17
+ """
18
+ device = torch.device("cuda")
19
+
20
+ # storage_block_size = block_size // compress_ratio = 256 // 4 = 64
21
+ kv_cache_spec = MLAAttentionSpec(
22
+ block_size=256,
23
+ num_kv_heads=1,
24
+ head_size=128,
25
+ dtype=torch.bfloat16,
26
+ compress_ratio=4,
27
+ )
28
+ vllm_config = create_vllm_config(max_model_len=1024)
29
+ builder = DeepseekV32IndexerMetadataBuilder(
30
+ kv_cache_spec=kv_cache_spec,
31
+ layer_names=["dummy"],
32
+ vllm_config=vllm_config,
33
+ device=device,
34
+ )
35
+
36
+ # Construct a single request where:
37
+ # - num_computed = 240 (=> compressed_pos_start = 60)
38
+ # - query_len = 40 (=> num_groups = 10)
39
+ # => compressed positions are 60..69 which cross the storage block boundary at 64.
40
+ query_start_loc = torch.tensor([0, 40], dtype=torch.int32, device=device)
41
+ query_start_loc_cpu = query_start_loc.cpu()
42
+ seq_lens = torch.tensor([280], dtype=torch.int32, device=device) # 240 + 40
43
+
44
+ # Two blocks: compressed positions 0..63 map to block 5, 64..127 map to block 7.
45
+ block_table_tensor = torch.tensor([[5, 7]], dtype=torch.int32, device=device)
46
+
47
+ # Dummy uncompressed slot mapping (length == uncompressed num_actual_tokens).
48
+ slot_mapping = torch.full((40,), -123, dtype=torch.int64, device=device)
49
+
50
+ common = CommonAttentionMetadata(
51
+ query_start_loc=query_start_loc,
52
+ query_start_loc_cpu=query_start_loc_cpu,
53
+ seq_lens=seq_lens,
54
+ seq_lens_cpu_upper_bound=seq_lens.cpu(),
55
+ num_reqs=1,
56
+ num_actual_tokens=40,
57
+ max_query_len=40,
58
+ max_seq_len=280,
59
+ block_table_tensor=block_table_tensor,
60
+ slot_mapping=slot_mapping,
61
+ causal=True,
62
+ )
63
+
64
+ md = builder.build(common_prefix_len=0, common_attn_metadata=common)
65
+
66
+ # The compressed slot_mapping retains the original uncompressed size (40).
67
+ # Only every compress_ratio-th position gets a valid slot; the rest are -1.
68
+ assert md.slot_mapping.numel() == 40
69
+ valid_slots = md.slot_mapping[md.slot_mapping >= 0]
70
+ assert valid_slots.numel() == 10 # 40 tokens / compress_ratio 4
71
+
72
+ storage_bs = kv_cache_spec.storage_block_size # 64
73
+ # Compressed positions 60..63 land in block 5, positions 64..69 in block 7.
74
+ expected = torch.tensor(
75
+ [
76
+ 5 * storage_bs + 60,
77
+ 5 * storage_bs + 61,
78
+ 5 * storage_bs + 62,
79
+ 5 * storage_bs + 63,
80
+ ]
81
+ + [
82
+ 7 * storage_bs + 0,
83
+ 7 * storage_bs + 1,
84
+ 7 * storage_bs + 2,
85
+ 7 * storage_bs + 3,
86
+ 7 * storage_bs + 4,
87
+ 7 * storage_bs + 5,
88
+ ],
89
+ dtype=torch.int64,
90
+ device=device,
91
+ )
92
+ torch.testing.assert_close(valid_slots, expected)
tests/v1/core/test_kv_cache_utils.py
+3
-2
@@ -1855,10 +1855,11 @@ def test_generate_scheduler_kv_cache_config():
1855
1856
1857
def new_mla_spec(cache_dtype_str=None):
1858
+ # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576
1859
return MLAAttentionSpec(
1860
block_size=16,
1860
- num_kv_heads=16,
1861
- head_size=64,
1861
+ num_kv_heads=1,
1862
+ head_size=576,
1863
dtype=torch.float32,
1864
cache_dtype_str=cache_dtype_str,
1865
)
tests/v1/core/test_prefix_caching.py
+19
-20
@@ -557,19 +557,19 @@ def test_prefill_hybrid_model_eagle():
557
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
558
assert len(req1.block_hashes) == num_full_blocks
559
assert computed_blocks.get_block_ids() == (
560
- [1, 2, 3, 4],
561
- [0, 9, 10, 11],
562
- [0, 16, 17, 18],
560
+ [1, 2, 3, 4, 5],
561
+ [0, 0, 10, 11, 12],
562
+ [0, 0, 17, 18, 19],
563
)
564
- assert num_computed_tokens == 4 * block_size
564
+ assert num_computed_tokens == 5 * block_size
565
num_new_tokens = len(all_token_ids) - num_computed_tokens
566
blocks = manager.allocate_slots(
567
req1, num_new_tokens, num_computed_tokens, computed_blocks
568
)
569
assert blocks is not None and blocks.get_block_ids() == (
570
- [22, 23, 24],
571
- [25, 26, 27],
572
- [28, 29, 30],
570
+ [22, 23],
571
+ [24, 25],
572
+ [26, 27],
573
)
574
for block_per_group in computed_blocks.blocks:
575
for block in block_per_group:
@@ -591,7 +591,7 @@ def test_prefill_hybrid_model_eagle():
591
make_block_hash_with_group_id(block_hashes[0], 1),
592
make_block_hash_with_group_id(block_hashes[0], 2),
593
],
594
- 4,
594
+ 5,
595
)
596
597
# Evict the first block of full attention, makes total cache miss.
@@ -605,7 +605,7 @@ def test_prefill_hybrid_model_eagle():
605
0,
606
)
607
608
- # Evict the last block of all layers, reduces the hit length to 3.
608
+ # Evict the last block of all layers, reduces the hit length to 4.
609
_test_partial_request_hit(
610
manager,
611
block_size,
@@ -617,10 +617,10 @@ def test_prefill_hybrid_model_eagle():
617
make_block_hash_with_group_id(block_hashes[-1], 1),
618
make_block_hash_with_group_id(block_hashes[-1], 2),
619
],
620
- 3,
620
+ 4,
621
)
622
623
- # Evict the last block of full attention, reduces the hit length to 3.
623
+ # Evict the last block of full attention, reduces the hit length to 4.
624
_test_partial_request_hit(
625
manager,
626
block_size,
@@ -628,7 +628,7 @@ def test_prefill_hybrid_model_eagle():
628
"5",
629
all_token_ids,
630
[make_block_hash_with_group_id(block_hashes[-1], 0)],
631
- 3,
631
+ 4,
632
)
633
634
# Since the last block of full attention is dropped for eagle, evict
@@ -655,12 +655,11 @@ def test_prefill_hybrid_model_eagle():
655
3,
656
)
657
658
- # Evict different set of blocks for full attention and sliding window makes
659
- # total cache miss.
660
- # The cache hit length of full attention is 4 * block_size.
661
- # The cache hit length of sliding window is 3 * block_size.
662
- # Then it is cache miss as the two type of layers
663
- # have different hit length.
658
+ # Evict different set of blocks for full attention and sliding window.
659
+ # Full loses its last block so it drops to 4 full blocks after the eagle
660
+ # pop; SWA lost block 0 (outside the sliding window of the final hit),
661
+ # which is not required for the K+1 anchor at position 4. Coordinated
662
+ # single-drop aligns both groups at hit=4.
663
_test_partial_request_hit(
664
manager,
665
block_size,
@@ -672,7 +671,7 @@ def test_prefill_hybrid_model_eagle():
671
make_block_hash_with_group_id(block_hashes[0], 1),
672
make_block_hash_with_group_id(block_hashes[0], 2),
673
],
675
- 0,
674
+ 4,
675
)
676
677
@@ -893,7 +892,7 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):
892
# - 2 groups: 1 full + 1 other
893
_EAGLE_HYBRID_MODEL_TEST_CASES = [
894
# 2 groups: 1 full + 1 other
896
- pytest.param(["full", "sliding_window"], 2, id="2g-full+sw"),
895
+ pytest.param(["full", "sliding_window"], 3, id="2g-full+sw"),
896
]
897
898
tests/v1/core/test_scheduler.py
+2
@@ -1892,6 +1892,7 @@ def create_scheduler_with_priority(
1892
log_stats=True,
1893
structured_output_manager=StructuredOutputManager(vllm_config),
1894
block_size=block_size,
1895
+ hash_block_size=block_size,
1896
)
1897
1898
@@ -4008,6 +4009,7 @@ def _create_encoder_decoder_scheduler(
4009
vllm_config=vllm_config,
4010
kv_cache_config=kv_cache_config,
4011
block_size=block_size,
4012
+ hash_block_size=block_size,
4013
structured_output_manager=StructuredOutputManager(vllm_config),
4014
)
4015
tests/v1/kv_connector/unit/test_mooncake_connector.py
+27
-23
@@ -91,8 +91,10 @@ def test_basic_interface():
91
assert request_id in kv_connector_metadata.reqs_to_recv["my-engine-id"]
92
req_meta = kv_connector_metadata.reqs_to_recv["my-engine-id"][request_id]
93
94
+ # local_block_ids is list[list[int]] (per-group); flatten for comparison.
95
+ all_block_ids = [bid for group in req_meta.local_block_ids for bid in group]
96
for block_id, block in zip(
95
- req_meta.local_block_ids,
97
+ all_block_ids,
98
scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
99
request_id
100
],
@@ -228,15 +230,15 @@ def test_scheduler_request_finished():
230
231
# Case: Capped length (Successful prefill, need to send to decoder)
232
request.status = RequestStatus.FINISHED_LENGTH_CAPPED
231
- delay_free, _ = scheduler_connector.request_finished(request, block_ids=[10, 11])
233
+ delay_free, _ = scheduler_connector.request_finished(request, block_ids=([10, 11],))
234
assert delay_free is True
235
assert "id-1" in scheduler_connector._reqs_need_send
234
- assert scheduler_connector._reqs_need_send["id-1"][1] == [10, 11]
236
+ assert scheduler_connector._reqs_need_send["id-1"][1] == [[10, 11]]
237
238
# Case: Aborted (No need to transfer, free blocks immediately)
239
scheduler_connector._reqs_need_send.clear()
240
request.status = RequestStatus.FINISHED_ABORTED
239
- delay_free, _ = scheduler_connector.request_finished(request, block_ids=[12])
241
+ delay_free, _ = scheduler_connector.request_finished(request, block_ids=([12],))
242
assert delay_free is False
243
assert len(scheduler_connector._reqs_need_send) == 0
244
assert "id-1" in scheduler_connector._reqs_not_processed
@@ -334,7 +336,7 @@ async def test_kv_producer(monkeypatch):
336
send_meta = SendBlockMeta(
337
p_req_id="p-req-1",
338
transfer_id=transfer_id,
337
- local_block_ids=[10, 11],
339
+ local_block_ids=[[10, 11]],
340
ready=asyncio.Event(),
341
)
342
prefill_worker.reqs_need_send[transfer_id] = send_meta
@@ -346,7 +348,7 @@ async def test_kv_producer(monkeypatch):
348
remote_port=54321,
349
remote_tp_size=1,
350
remote_tp_rank=0,
349
- req_blocks={"d-req-1": (transfer_id, [20, 21])},
351
+ req_blocks={"d-req-1": (transfer_id, [[20, 21]])},
352
kv_caches_base_addr=[0x2000],
353
block_lens=[block_len],
354
)
@@ -389,7 +391,7 @@ async def test_kv_producer(monkeypatch):
391
prefill_worker.reqs_need_send[transfer_id] = send_meta
392
send_meta.sent = 0
393
send_meta.ready.set()
392
- xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20])
394
+ xfer_meta.req_blocks["d-req-1"] = (transfer_id, [[20]])
395
# Worker processes the consumer's request
396
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
397
# Verify transfer parameters are correct: 11 to 20
@@ -407,7 +409,7 @@ async def test_kv_producer(monkeypatch):
409
prefill_worker.reqs_need_send[transfer_id] = send_meta
410
send_meta.sent = 0
411
send_meta.ready.set()
410
- xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21, 22])
412
+ xfer_meta.req_blocks["d-req-1"] = (transfer_id, [[20, 21, 22]])
413
# Worker processes the consumer's request
414
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
415
# This should not be called because error.
@@ -424,7 +426,7 @@ async def test_kv_producer(monkeypatch):
426
prefill_worker.reqs_need_send[transfer_id] = send_meta
427
send_meta.sent = 0
428
send_meta.ready.clear()
427
- xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21])
429
+ xfer_meta.req_blocks["d-req-1"] = (transfer_id, [[20, 21]])
430
# Worker processes the consumer's request
431
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
432
# This should not be called because timeout.
@@ -443,7 +445,7 @@ async def test_kv_producer(monkeypatch):
445
prefill_worker.reqs_need_send[transfer_id] = send_meta
446
send_meta.sent = 0
447
send_meta.ready.set()
446
- xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21])
448
+ xfer_meta.req_blocks["d-req-1"] = (transfer_id, [[20, 21]])
449
# Worker processes the consumer's request
450
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
451
mock_send_blocks.assert_called_once()
@@ -481,7 +483,7 @@ async def test_kv_consumuer(monkeypatch):
483
"d-req-1": PullReqMeta(
484
d_req_id="d-req-1",
485
transfer_id="xfer-req-1",
484
- local_block_ids=[100, 101],
486
+ local_block_ids=[[100, 101]],
487
remote_engine_id="p-engine",
488
remote_bootstrap_addr="http://bootstrap:33333",
489
pull_tasks_count=1,
@@ -514,7 +516,7 @@ async def test_kv_consumuer(monkeypatch):
516
517
assert sent_meta.remote_hostname == "127.0.0.1"
518
assert sent_meta.remote_port == 54321
517
- assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [100, 101])
519
+ assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [[100, 101]])
520
521
# Verify internal state is updated correctly.
522
assert "d-req-1" in decode_worker.finished_recving_reqs
@@ -538,7 +540,7 @@ async def test_worker_get_finished_timeout(monkeypatch):
540
prefill_worker.reqs_need_send["tx-expired"] = SendBlockMeta(
541
p_req_id="p-req-expired",
542
transfer_id="tx-expired",
541
- local_block_ids=[1, 2],
543
+ local_block_ids=[[1, 2]],
544
ready=MagicMock(),
545
expire_time=time.perf_counter() - 100,
546
)
@@ -547,7 +549,7 @@ async def test_worker_get_finished_timeout(monkeypatch):
549
prefill_worker.reqs_need_send["tx-active"] = SendBlockMeta(
550
p_req_id="p-req-active",
551
transfer_id="tx-active",
550
- local_block_ids=[3, 4],
552
+ local_block_ids=[[3, 4]],
553
ready=MagicMock(),
554
expire_time=time.perf_counter() + 100,
555
)
@@ -703,7 +705,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
705
prefill_worker.sender_loop = asyncio.get_event_loop()
706
707
transfer_id = "xfer-hetero-1"
706
- local_block_ids = [10, 11]
708
+ local_block_ids = [[10, 11]]
709
send_meta = SendBlockMeta(
710
p_req_id="p-req-h1",
711
transfer_id=transfer_id,
@@ -720,9 +722,9 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
722
mock_socket.send_multipart = AsyncMock()
723
identity = b"consumer-hetero"
724
723
- # Assign different remote block IDs per D rank
725
+ # Assign different remote block IDs per D rank (nested per-group)
726
d_rank_remote_blocks = {
725
- rank: [20 + i * 10, 21 + i * 10] for i, rank in enumerate(target_d_ranks)
727
+ rank: [[20 + i * 10, 21 + i * 10]] for i, rank in enumerate(target_d_ranks)
728
}
729
730
with patch.object(
@@ -757,11 +759,15 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
759
dst_ptrs = call_args[2]
760
lengths = call_args[3]
761
762
+ # Flatten nested per-group block IDs for assertions
763
+ flat_local = [b for g in local_block_ids for b in g]
764
+ flat_remote = [b for g in remote_block_ids for b in g]
765
+
766
# Heterogeneous TP: blocks cannot be coalesced because
767
# local and remote block_lens differ
762
- assert len(src_ptrs) == len(local_block_ids)
763
- assert len(dst_ptrs) == len(local_block_ids)
764
- assert len(lengths) == len(local_block_ids)
768
+ assert len(src_ptrs) == len(flat_local)
769
+ assert len(dst_ptrs) == len(flat_local)
770
+ assert len(lengths) == len(flat_local)
771
772
# Compute expected offsets based on TP ratio
773
if d_tp_size <= P_TP_SIZE:
@@ -775,9 +781,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
781
expected_dst_off = 0
782
expected_xfer_len = remote_block_len
783
778
- for idx, (lblk, rblk) in enumerate(
779
- zip(local_block_ids, remote_block_ids)
780
- ):
784
+ for idx, (lblk, rblk) in enumerate(zip(flat_local, flat_remote)):
785
assert src_ptrs[idx] == (
786
0x1000 + lblk * local_block_len + expected_src_off
787
)
tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
new
+410
@@ -0,0 +1,410 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for MooncakeConnector HMA (Hybrid Memory Architecture) support.
4
+
5
+Covers sliding-window clipping, multi-group metadata shape, multi-group
6
+send trimming, and group-count invariant checking in _build_transfer_params.
7
+"""
8
+
9
+import asyncio
10
+from unittest.mock import patch
11
+
12
+import pytest
13
+
14
+from vllm.config import set_current_vllm_config
15
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
16
+ KVConnectorRole,
17
+ MooncakeConnector,
18
+ MooncakeConnectorMetadata,
19
+ MooncakeConnectorScheduler,
20
+ MooncakeXferMetadata,
21
+ SendBlockMeta,
22
+ TransferRegion,
23
+)
24
+
25
+from .test_mooncake_connector import FakeMooncakeWrapper, patch_worker_dependencies
26
+from .utils import create_request, create_vllm_config, make_kv_cache_config
27
+
28
+
29
+# ---------------------------------------------------------------------------
30
+# test_sw_sizes: blocks_per_sw computed from KVCacheConfig
31
+# ---------------------------------------------------------------------------
32
+@pytest.mark.cpu_test
33
+@pytest.mark.parametrize(
34
+ "swa_enabled,expected_blocks_per_sw",
35
+ [
36
+ # SWA enabled: FullAttentionSpec (0) + SlidingWindowSpec (2048/16=128+1)
37
+ (True, [0, 128 + 1]),
38
+ # SWA disabled: only FullAttentionSpec (0)
39
+ (False, [0]),
40
+ ],
41
+)
42
+def test_sw_sizes(swa_enabled, expected_blocks_per_sw):
43
+ """blocks_per_sw is correctly computed based on SWA enabled/disabled."""
44
+ block_size = 16
45
+ vllm_config = create_vllm_config(
46
+ kv_connector="MooncakeConnector",
47
+ kv_role="kv_both",
48
+ block_size=block_size,
49
+ )
50
+ # Override so HMA detection works
51
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
52
+ kv_cache_config = make_kv_cache_config(
53
+ block_size=block_size, swa_enabled=swa_enabled, sw_size=2048
54
+ )
55
+
56
+ scheduler = MooncakeConnectorScheduler(
57
+ vllm_config=vllm_config,
58
+ engine_id="test-engine",
59
+ kv_cache_config=kv_cache_config,
60
+ )
61
+ assert scheduler.blocks_per_sw == expected_blocks_per_sw
62
+
63
+
64
+# ---------------------------------------------------------------------------
65
+# test_is_hma_required: derived from kv_cache_config groups
66
+# ---------------------------------------------------------------------------
67
+@pytest.mark.cpu_test
68
+@pytest.mark.parametrize(
69
+ "swa_enabled,disable_hma,expected_is_hma",
70
+ [
71
+ (True, False, True), # SWA group present, HMA enabled
72
+ (True, True, False), # SWA group present, but HMA disabled
73
+ (False, False, False), # FA only, HMA not needed
74
+ ],
75
+)
76
+def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma):
77
+ """_is_hma_required is correctly derived from kv_cache_config."""
78
+ block_size = 16
79
+ vllm_config = create_vllm_config(
80
+ kv_connector="MooncakeConnector",
81
+ kv_role="kv_both",
82
+ block_size=block_size,
83
+ )
84
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hma
85
+ kv_cache_config = make_kv_cache_config(
86
+ block_size=block_size, swa_enabled=swa_enabled
87
+ )
88
+
89
+ scheduler = MooncakeConnectorScheduler(
90
+ vllm_config=vllm_config,
91
+ engine_id="test-engine",
92
+ kv_cache_config=kv_cache_config,
93
+ )
94
+ assert scheduler._is_hma_required is expected_is_hma
95
+
96
+
97
+# ---------------------------------------------------------------------------
98
+# test_get_sw_clipped_blocks: sliding-window clipping logic
99
+# ---------------------------------------------------------------------------
100
+@pytest.mark.cpu_test
101
+def test_get_sw_clipped_blocks():
102
+ """get_sw_clipped_blocks clips SWA group but keeps FA group intact."""
103
+ block_size = 16
104
+ vllm_config = create_vllm_config(
105
+ kv_connector="MooncakeConnector",
106
+ kv_role="kv_both",
107
+ block_size=block_size,
108
+ )
109
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
110
+ # SW=128 tokens → 128/16 = 8 blocks + 1 = 9 blocks_per_sw
111
+ kv_cache_config = make_kv_cache_config(
112
+ block_size=block_size, swa_enabled=True, sw_size=128
113
+ )
114
+
115
+ scheduler = MooncakeConnectorScheduler(
116
+ vllm_config=vllm_config,
117
+ engine_id="test-engine",
118
+ kv_cache_config=kv_cache_config,
119
+ )
120
+ assert scheduler.blocks_per_sw == [0, 9]
121
+
122
+ # FA group: 20 blocks, SW group: 20 blocks (exceeds window)
123
+ fa_blocks = list(range(20))
124
+ sw_blocks = list(range(100, 120))
125
+ block_ids = (fa_blocks, sw_blocks)
126
+
127
+ clipped = scheduler.get_sw_clipped_blocks(block_ids)
128
+
129
+ # FA: untouched (blocks_per_sw[0] = 0)
130
+ assert clipped[0] == fa_blocks
131
+ # SW: clipped to last 9 blocks
132
+ assert clipped[1] == sw_blocks[-9:]
133
+ assert len(clipped[1]) == 9
134
+
135
+
136
+@pytest.mark.cpu_test
137
+def test_get_sw_clipped_blocks_noop_no_hma():
138
+ """get_sw_clipped_blocks is a no-op when HMA is not required."""
139
+ block_size = 16
140
+ vllm_config = create_vllm_config(
141
+ kv_connector="MooncakeConnector",
142
+ kv_role="kv_both",
143
+ block_size=block_size,
144
+ )
145
+ # FA only → _is_hma_required = False
146
+ kv_cache_config = make_kv_cache_config(block_size=block_size, swa_enabled=False)
147
+
148
+ scheduler = MooncakeConnectorScheduler(
149
+ vllm_config=vllm_config,
150
+ engine_id="test-engine",
151
+ kv_cache_config=kv_cache_config,
152
+ )
153
+ assert scheduler._is_hma_required is False
154
+
155
+ block_ids = ([1, 2, 3],)
156
+ clipped = scheduler.get_sw_clipped_blocks(block_ids)
157
+ assert clipped == [[1, 2, 3]]
158
+
159
+
160
+# ---------------------------------------------------------------------------
161
+# test_metadata_hma_block_ids: MooncakeConnectorMetadata stores per-group IDs
162
+# ---------------------------------------------------------------------------
163
+@pytest.mark.cpu_test
164
+def test_metadata_hma_block_ids():
165
+ """MooncakeConnectorMetadata.add_new_req stores per-group block IDs."""
166
+ metadata = MooncakeConnectorMetadata()
167
+
168
+ # FA group: 6 blocks, SW group: 3 blocks (clipped)
169
+ fa_blocks = [0, 1, 2, 3, 4, 5]
170
+ sw_blocks = [10, 11, 12]
171
+
172
+ # Test recv path
173
+ metadata.add_new_req(
174
+ request_id="recv-req",
175
+ local_block_ids=[fa_blocks, sw_blocks],
176
+ kv_transfer_params={
177
+ "transfer_id": "recv-req",
178
+ "remote_engine_id": "remote-engine",
179
+ "remote_bootstrap_addr": "http://bootstrap:33333",
180
+ },
181
+ load_remote_cache=True,
182
+ )
183
+
184
+ assert "recv-req" in metadata.reqs_to_recv["remote-engine"]
185
+ req_meta = metadata.reqs_to_recv["remote-engine"]["recv-req"]
186
+ assert len(req_meta.local_block_ids) == 2
187
+ assert req_meta.local_block_ids[0] == fa_blocks
188
+ assert req_meta.local_block_ids[1] == sw_blocks
189
+
190
+ # Test send path
191
+ metadata.add_new_req(
192
+ request_id="send-req",
193
+ local_block_ids=[fa_blocks, sw_blocks],
194
+ kv_transfer_params={
195
+ "transfer_id": "send-req",
196
+ },
197
+ load_remote_cache=False,
198
+ )
199
+
200
+ assert "send-req" in metadata.reqs_to_send
201
+ transfer_id, stored_blocks = metadata.reqs_to_send["send-req"]
202
+ assert transfer_id == "send-req"
203
+ assert len(stored_blocks) == 2
204
+ assert stored_blocks[0] == fa_blocks
205
+ assert stored_blocks[1] == sw_blocks
206
+
207
+
208
+# ---------------------------------------------------------------------------
209
+# test_build_transfer_params_multi_group_trimming
210
+# ---------------------------------------------------------------------------
211
+@pytest.mark.asyncio
212
+@patch(
213
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
214
+ ".mooncake_connector.TransferEngine",
215
+ FakeMooncakeWrapper,
216
+)
217
+async def test_build_transfer_params_multi_group_trimming(monkeypatch):
218
+ """_build_transfer_params trims per-group blocks when local > remote."""
219
+
220
+ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
221
+ vllm_config = create_vllm_config(
222
+ kv_connector="MooncakeConnector", kv_role="kv_producer"
223
+ )
224
+
225
+ with set_current_vllm_config(vllm_config), patch_worker_dependencies():
226
+ connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
227
+ worker = connector.connector_worker
228
+
229
+ block_len = 4096
230
+ # Call _build_transfer_params directly (avoids send_kv_to_decode
231
+ # async event loop complexity).
232
+ transfer_id = "xfer-hma-trim"
233
+ send_meta = SendBlockMeta(
234
+ p_req_id="p-trim",
235
+ transfer_id=transfer_id,
236
+ # FA: 4 blocks, SW: 3 blocks (producer has more)
237
+ local_block_ids=[[10, 11, 12, 13], [20, 21, 22]],
238
+ ready=asyncio.Event(),
239
+ )
240
+
241
+ xfer_meta = MooncakeXferMetadata(
242
+ remote_hostname="consumer-host",
243
+ remote_port=54321,
244
+ remote_tp_size=1,
245
+ remote_tp_rank=0,
246
+ req_blocks={
247
+ "d-trim": (
248
+ transfer_id,
249
+ # FA: 2 blocks, SW: 2 blocks (consumer needs fewer)
250
+ [[30, 31], [40, 41]],
251
+ )
252
+ },
253
+ kv_caches_base_addr=[0x2000],
254
+ block_lens=[block_len],
255
+ )
256
+
257
+ local_regions = [
258
+ TransferRegion(
259
+ base_addr=0x1000, block_len=block_len, kv_block_len=block_len
260
+ ),
261
+ ]
262
+ remote_regions = [
263
+ TransferRegion(
264
+ base_addr=0x2000, block_len=block_len, kv_block_len=block_len
265
+ ),
266
+ ]
267
+
268
+ ready_reqs = [("d-trim", send_meta)]
269
+ (
270
+ src_ptrs,
271
+ dst_ptrs,
272
+ lengths,
273
+ err_reqs,
274
+ err_msg,
275
+ ) = await worker._build_transfer_params(
276
+ ready_reqs, xfer_meta, local_regions, remote_regions
277
+ )
278
+
279
+ # No errors
280
+ assert err_reqs == []
281
+ assert err_msg is None
282
+ # After trimming: FA [10..13] → last 2 → [12,13]; SW [20..22] → last 2 → [21,22]
283
+ # Flattened: [12,13,21,22] = 4 blocks → coalesced into some transfers
284
+ assert len(src_ptrs) > 0
285
+ assert len(dst_ptrs) == len(src_ptrs)
286
+ assert len(lengths) == len(src_ptrs)
287
+
288
+ worker.shutdown()
289
+
290
+
291
+# ---------------------------------------------------------------------------
292
+# test_build_transfer_params_group_count_mismatch
293
+# ---------------------------------------------------------------------------
294
+@pytest.mark.asyncio
295
+@patch(
296
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
297
+ ".mooncake_connector.TransferEngine",
298
+ FakeMooncakeWrapper,
299
+)
300
+async def test_build_transfer_params_group_count_mismatch(monkeypatch):
301
+ """_build_transfer_params reports an error when group counts differ."""
302
+
303
+ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
304
+ vllm_config = create_vllm_config(
305
+ kv_connector="MooncakeConnector", kv_role="kv_producer"
306
+ )
307
+
308
+ with set_current_vllm_config(vllm_config), patch_worker_dependencies():
309
+ connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
310
+ worker = connector.connector_worker
311
+
312
+ block_len = 4096
313
+ transfer_id = "xfer-mismatch"
314
+ send_meta = SendBlockMeta(
315
+ p_req_id="p-mismatch",
316
+ transfer_id=transfer_id,
317
+ # Producer has 2 groups
318
+ local_block_ids=[[10, 11], [20, 21]],
319
+ ready=asyncio.Event(),
320
+ )
321
+
322
+ # Consumer has only 1 group — group count mismatch
323
+ xfer_meta = MooncakeXferMetadata(
324
+ remote_hostname="consumer-host",
325
+ remote_port=54321,
326
+ remote_tp_size=1,
327
+ remote_tp_rank=0,
328
+ req_blocks={
329
+ "d-mismatch": (transfer_id, [[30, 31]]),
330
+ },
331
+ kv_caches_base_addr=[0x2000],
332
+ block_lens=[block_len],
333
+ )
334
+
335
+ local_regions = [
336
+ TransferRegion(
337
+ base_addr=0x1000, block_len=block_len, kv_block_len=block_len
338
+ ),
339
+ ]
340
+ remote_regions = [
341
+ TransferRegion(
342
+ base_addr=0x2000, block_len=block_len, kv_block_len=block_len
343
+ ),
344
+ ]
345
+
346
+ ready_reqs = [("d-mismatch", send_meta)]
347
+ (
348
+ src_ptrs,
349
+ dst_ptrs,
350
+ lengths,
351
+ err_reqs,
352
+ err_msg,
353
+ ) = await worker._build_transfer_params(
354
+ ready_reqs, xfer_meta, local_regions, remote_regions
355
+ )
356
+
357
+ # Mismatched req is reported via err_reqs/err_msg with no transfers built.
358
+ assert err_reqs == ["d-mismatch"]
359
+ assert err_msg == "KV group count mismatch"
360
+ assert src_ptrs == []
361
+ assert dst_ptrs == []
362
+ assert lengths == []
363
+
364
+ worker.shutdown()
365
+
366
+
367
+# ---------------------------------------------------------------------------
368
+# test_request_finished_with_hma_groups
369
+# ---------------------------------------------------------------------------
370
+@pytest.mark.cpu_test
371
+def test_request_finished_with_hma_groups():
372
+ """request_finished correctly handles per-group block_ids."""
373
+ block_size = 16
374
+ vllm_config = create_vllm_config(
375
+ kv_connector="MooncakeConnector",
376
+ kv_role="kv_producer",
377
+ block_size=block_size,
378
+ )
379
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
380
+ kv_cache_config = make_kv_cache_config(
381
+ block_size=block_size, swa_enabled=True, sw_size=128
382
+ )
383
+
384
+ scheduler = MooncakeConnectorScheduler(
385
+ vllm_config=vllm_config,
386
+ engine_id="test-engine",
387
+ kv_cache_config=kv_cache_config,
388
+ )
389
+
390
+ request = create_request(request_id=1, do_remote_decode=True)
391
+ request.kv_transfer_params["transfer_id"] = request.request_id
392
+
393
+ from vllm.v1.request import RequestStatus
394
+
395
+ request.status = RequestStatus.FINISHED_LENGTH_CAPPED
396
+
397
+ # 2 groups: FA with 10 blocks, SW with 20 blocks (will be clipped)
398
+ fa_blocks = list(range(10))
399
+ sw_blocks = list(range(100, 120))
400
+ block_ids = (fa_blocks, sw_blocks)
401
+
402
+ delay_free, _ = scheduler.request_finished(request, block_ids)
403
+ assert delay_free is True
404
+ assert request.request_id in scheduler._reqs_need_send
405
+
406
+ _, stored_blocks = scheduler._reqs_need_send[request.request_id]
407
+ # FA: untouched
408
+ assert stored_blocks[0] == fa_blocks
409
+ # SW: clipped to last 9 blocks (sw_size=128, block_size=16 → 8+1=9)
410
+ assert stored_blocks[1] == sw_blocks[-9:]
tests/v1/streaming_input/test_scheduler_streaming.py
+1
@@ -76,6 +76,7 @@ def create_scheduler() -> Scheduler:
76
log_stats=True,
77
structured_output_manager=StructuredOutputManager(vllm_config),
78
block_size=16,
79
+ hash_block_size=16,
80
)
81
82
tools/install_deepgemm.sh
+1
-1
@@ -7,7 +7,7 @@ set -e
7
# Default values
8
# Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake
9
DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
10
-DEEPGEMM_GIT_REF="477618cd51baffca09c4b0b87e97c03fe827ef03"
10
+DEEPGEMM_GIT_REF="891d57b4db1071624b5c8fa0d1e51cb317fa709f"
11
WHEEL_DIR=""
12
13
# Parse command line arguments
vllm/_custom_ops.py
+41
-3
@@ -404,10 +404,24 @@ def rotary_embedding(
404
head_size: int,
405
cos_sin_cache: torch.Tensor,
406
is_neox: bool,
407
+ rope_dim_offset: int = 0,
408
+ inverse: bool = False,
409
) -> None:
408
- torch.ops._C.rotary_embedding(
409
- positions, query, key, head_size, cos_sin_cache, is_neox
410
- )
410
+ if rope_dim_offset == 0 and not inverse:
411
+ torch.ops._C.rotary_embedding(
412
+ positions, query, key, head_size, cos_sin_cache, is_neox
413
+ )
414
+ else:
415
+ torch.ops._C.rotary_embedding(
416
+ positions,
417
+ query,
418
+ key,
419
+ head_size,
420
+ cos_sin_cache,
421
+ is_neox,
422
+ rope_dim_offset,
423
+ inverse,
424
+ )
425
426
427
# layer norm ops
@@ -2503,6 +2517,30 @@ def topk_sigmoid(
2517
)
2518
2519
2520
+def topk_hash_softplus_sqrt(
2521
+ topk_weights: torch.Tensor,
2522
+ topk_indices: torch.Tensor,
2523
+ token_expert_indices: torch.Tensor,
2524
+ gating_output: torch.Tensor,
2525
+ renormalize: bool = False,
2526
+ routed_scaling_factor: float = 1.0,
2527
+ e_score_correction_bias: torch.Tensor | None = None,
2528
+ input_tokens: torch.Tensor | None = None,
2529
+ hash_indices_table: torch.Tensor | None = None,
2530
+) -> None:
2531
+ torch.ops._moe_C.topk_softplus_sqrt(
2532
+ topk_weights,
2533
+ topk_indices,
2534
+ token_expert_indices,
2535
+ gating_output,
2536
+ renormalize,
2537
+ routed_scaling_factor,
2538
+ e_score_correction_bias,
2539
+ input_tokens,
2540
+ hash_indices_table,
2541
+ )
2542
+
2543
+
2544
def grouped_topk(
2545
scores: torch.Tensor,
2546
num_expert_group: int,
vllm/config/attention.py
+3
@@ -51,6 +51,9 @@ class AttentionConfig:
51
use_prefill_query_quantization: bool = False
52
"""If set, quantize query for attention in prefill."""
53
54
+ use_fp4_indexer_cache: bool = False
55
+ """If set, use fp4 indexer cache for dsv32 family model (not support yet)"""
56
+
57
def compute_hash(self) -> str:
58
"""
59
Provide a hash that uniquely identifies all the configs
vllm/config/cache.py
+14
@@ -51,6 +51,18 @@ class CacheConfig:
51
"""Whether block_size was explicitly provided. Derived automatically."""
52
user_specified_mamba_block_size: bool = field(default=False, init=False)
53
"""Whether mamba_block_size was explicitly provided. Derived automatically."""
54
+ hash_block_size: SkipValidation[int] | None = None # type: ignore
55
+ """Block size (in tokens) used for computing Request's block_hashes.
56
+
57
+ This can be set to a finer granularity than the physical KV cache block
58
+ sizes (e.g. 8) as long as every KV cache group's `block_size` is divisible
59
+ by it. This enables prefix-caching keys to be computed at the finest common
60
+ granularity and then merged for larger physical block sizes.
61
+
62
+ This config is not static default. If left unspecified, vLLM will choose a
63
+ default based on the resolved KV cache groups (typically the smallest KV
64
+ cache block size when there are multiple groups).
65
+ """
66
gpu_memory_utilization: float = Field(default=0.92, gt=0, le=1)
67
"""The fraction of GPU memory to be used for the model executor, which can
68
range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory
@@ -182,6 +194,8 @@ class CacheConfig:
194
"num_gpu_blocks_override",
195
"enable_prefix_caching",
196
"prefix_caching_hash_algo",
197
+ # Prefix-caching implementation detail (doesn't affect compiled graph).
198
+ "hash_block_size",
199
"mamba_page_size_padded",
200
"user_specified_block_size",
201
"user_specified_mamba_block_size",
vllm/config/compilation.py
+1
@@ -749,6 +749,7 @@ class CompilationConfig:
749
"vllm::kda_attention",
750
"vllm::sparse_attn_indexer",
751
"vllm::rocm_aiter_sparse_attn_indexer",
752
+ "vllm::deepseek_v4_attention",
753
]
754
755
def compute_hash(self) -> str:
vllm/config/kernel.py
+6
-4
@@ -50,7 +50,7 @@ class IrOpPriorityConfig:
50
name: {
51
provider: IrOp.registry[name].impls[provider].uuid() for provider in p
52
}
53
- for name, p in asdict(self).items()
53
+ for name, p in asdict(self).items() # type: ignore[call-overload]
54
}
55
56
return hash_factors(factors)
@@ -77,7 +77,7 @@ class IrOpPriorityConfig:
77
current_platform.import_ir_kernels()
78
79
with contextlib.ExitStack() as stack:
80
- for field in fields(self):
80
+ for field in fields(self): # type: ignore[arg-type]
81
op_priority = getattr(self, field.name)
82
assert op_priority is not None, (
83
f"IR op priority for {field.name} must be set"
@@ -98,7 +98,7 @@ class IrOpPriorityConfig:
98
A helper to create an IrOpPriorityConfig where fields not specified in kwargs
99
use the given default list.
100
"""
101
- for field in fields(cls):
101
+ for field in fields(cls): # type: ignore[arg-type]
102
if field.name not in kwargs:
103
kwargs[field.name] = list(default)
104
@@ -109,6 +109,7 @@ MoEBackend = Literal[
109
"auto",
110
"triton",
111
"deep_gemm",
112
+ "deep_gemm_mega_moe",
113
"cutlass",
114
"flashinfer_trtllm",
115
"flashinfer_cutlass",
@@ -136,8 +137,9 @@ class KernelConfig:
137
"""Backend for MoE expert computation kernels. Available options:
138
139
- "auto": Automatically select the best backend based on model and hardware
139
- - "triton": Use Triton-based fused MoE kernels
140
+ - "triton": Use Triton-based fused MoE kernels
141
- "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)
142
+ - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels
143
- "cutlass": Use vLLM CUTLASS kernels
144
- "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels
145
- "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
vllm/config/model.py
+5
-1
@@ -83,7 +83,7 @@ logger = init_logger(__name__)
83
RunnerOption = Literal["auto", RunnerType]
84
ConvertType = Literal["none", "embed", "classify"]
85
ConvertOption = Literal["auto", ConvertType]
86
-TokenizerMode = Literal["auto", "hf", "slow", "mistral", "deepseek_v32"]
86
+TokenizerMode = Literal["auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4"]
87
ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
88
LogprobsMode = Literal[
89
"raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs"
@@ -134,6 +134,7 @@ class ModelConfig:
134
- "slow" will always use the slow tokenizer.
135
- "mistral" will always use the tokenizer from `mistral_common`.
136
- "deepseek_v32" will always use the tokenizer from `deepseek_v32`.
137
+ - "deepseek_v4" will always use the tokenizer from `deepseek_v4`.
138
- "qwen_vl" will always use the tokenizer from `qwen_vl`.
139
- Other custom values can be supported via plugins."""
140
trust_remote_code: bool = False
@@ -565,6 +566,8 @@ class ModelConfig:
566
self.tokenizer_mode = "qwen_vl"
567
elif arch == "DeepseekV32ForCausalLM":
568
self.tokenizer_mode = "deepseek_v32"
569
+ elif arch == "DeepseekV4ForCausalLM":
570
+ self.tokenizer_mode = "deepseek_v4"
571
572
if self.tokenizer_mode != "auto":
573
logger.info(
@@ -952,6 +955,7 @@ class ModelConfig:
955
# imports during override detection (e.g., MXFP4 imports Triton)
956
"mxfp4",
957
"gpt_oss_mxfp4",
958
+ "deepseek_v4_fp8",
959
"cpu_awq",
960
"humming",
961
"gguf",
vllm/config/speculative.py
+11
-1
@@ -287,13 +287,23 @@ class SpeculativeConfig:
287
@staticmethod
288
def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
289
initial_architecture = hf_config.architectures[0]
290
- if hf_config.model_type in ("deepseek_v3", "deepseek_v32", "glm_moe_dsa"):
290
+ if hf_config.model_type in (
291
+ "deepseek_v3",
292
+ "deepseek_v32",
293
+ "glm_moe_dsa",
294
+ ):
295
hf_config.model_type = "deepseek_mtp"
296
if hf_config.model_type == "deepseek_mtp":
297
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
298
hf_config.update(
299
{"n_predict": n_predict, "architectures": ["DeepSeekMTPModel"]}
300
)
301
+ if hf_config.model_type == "deepseek_v4":
302
+ hf_config.model_type = "deepseek_mtp"
303
+ n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
304
+ hf_config.update(
305
+ {"n_predict": n_predict, "architectures": ["DeepSeekV4MTPModel"]}
306
+ )
307
if hf_config.model_type in ("pangu_ultra_moe"):
308
hf_config.model_type = "pangu_ultra_moe_mtp"
309
if hf_config.model_type == "pangu_ultra_moe_mtp":
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+143
-40
@@ -29,6 +29,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
29
KVConnectorBase_V1,
30
KVConnectorMetadata,
31
KVConnectorRole,
32
+ SupportsHMA,
33
)
34
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
35
MooncakeBootstrapServer,
@@ -43,10 +44,12 @@ from vllm.distributed.parallel_state import (
44
from vllm.forward_context import ForwardContext
45
from vllm.logger import init_logger
46
from vllm.platforms import current_platform
47
+from vllm.utils.math_utils import cdiv
48
from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
49
from vllm.v1.attention.backend import AttentionMetadata
50
from vllm.v1.attention.backends.utils import get_kv_cache_layout
51
from vllm.v1.core.sched.output import SchedulerOutput
52
+from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
53
from vllm.v1.request import RequestStatus
54
from vllm.v1.worker.utils import select_common_block_size
55
@@ -252,7 +255,7 @@ class MooncakeXferMetadata(
255
remote_port: int
256
remote_tp_size: int
257
remote_tp_rank: int
255
- req_blocks: dict[ReqId, tuple[TransferId, list[int]]]
258
+ req_blocks: dict[ReqId, tuple[TransferId, list[list[int]]]]
259
kv_caches_base_addr: list[int]
260
block_lens: list[int]
261
@@ -280,7 +283,7 @@ class MooncakeXferResponse(
283
class PullReqMeta:
284
d_req_id: ReqId
285
transfer_id: TransferId
283
- local_block_ids: list[int]
286
+ local_block_ids: list[list[int]]
287
remote_engine_id: EngineId
288
remote_bootstrap_addr: str
289
# Set expire time to avoid infinitely sending requests.
@@ -293,7 +296,7 @@ class PullReqMeta:
296
class SendBlockMeta:
297
p_req_id: ReqId
298
transfer_id: TransferId
296
- local_block_ids: list[int]
299
+ local_block_ids: list[list[int]]
300
ready: asyncio.Event
301
expire_time: float = float("inf")
302
need_send: int = 0
@@ -306,13 +309,13 @@ class MooncakeConnectorMetadata(KVConnectorMetadata):
309
# Use (engine_id, dp_rank) to group reqs with same dp.
310
# See comments in MooncakeBootstrapServer.
311
self.reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict)
309
- self.reqs_to_send: dict[ReqId, tuple[TransferId, list[int]]] = {}
312
+ self.reqs_to_send: dict[ReqId, tuple[TransferId, list[list[int]]]] = {}
313
self.reqs_not_processed: set[TransferId] = set()
314
315
def add_new_req(
316
self,
317
request_id: ReqId,
315
- local_block_ids: list[int],
318
+ local_block_ids: list[list[int]],
319
kv_transfer_params: dict[str, Any],
320
load_remote_cache: bool = True,
321
):
@@ -330,7 +333,7 @@ class MooncakeConnectorMetadata(KVConnectorMetadata):
333
self.reqs_to_send[request_id] = (transfer_id, local_block_ids)
334
335
333
-class MooncakeConnector(KVConnectorBase_V1):
336
+class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
337
def __init__(
338
self,
339
vllm_config: VllmConfig,
@@ -344,13 +347,18 @@ class MooncakeConnector(KVConnectorBase_V1):
347
self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
348
349
if role == KVConnectorRole.SCHEDULER:
350
+ assert kv_cache_config is not None, (
351
+ "kv_cache_config is required for SCHEDULER role"
352
+ )
353
self.connector_scheduler: MooncakeConnectorScheduler | None = (
348
- MooncakeConnectorScheduler(vllm_config, self.engine_id)
354
+ MooncakeConnectorScheduler(vllm_config, self.engine_id, kv_cache_config)
355
)
356
self.connector_worker: MooncakeConnectorWorker | None = None
357
elif role == KVConnectorRole.WORKER:
358
self.connector_scheduler = None
353
- self.connector_worker = MooncakeConnectorWorker(vllm_config, self.engine_id)
359
+ self.connector_worker = MooncakeConnectorWorker(
360
+ vllm_config, self.engine_id, kv_cache_config
361
+ )
362
363
@classmethod
364
def get_required_kvcache_layout(cls, vllm_config: VllmConfig):
@@ -401,6 +409,14 @@ class MooncakeConnector(KVConnectorBase_V1):
409
self,
410
request: "Request",
411
block_ids: list[int],
412
+ ) -> tuple[bool, dict[str, Any] | None]:
413
+ assert self.connector_scheduler is not None
414
+ return self.connector_scheduler.request_finished(request, (block_ids,))
415
+
416
+ def request_finished_all_groups(
417
+ self,
418
+ request: "Request",
419
+ block_ids: tuple[list[int], ...],
420
) -> tuple[bool, dict[str, Any] | None]:
421
assert self.connector_scheduler is not None
422
return self.connector_scheduler.request_finished(request, block_ids)
@@ -445,8 +461,14 @@ class MooncakeConnector(KVConnectorBase_V1):
461
class MooncakeConnectorScheduler:
462
"""Implementation of Scheduler side methods"""
463
448
- def __init__(self, vllm_config: VllmConfig, engine_id: str):
464
+ def __init__(
465
+ self,
466
+ vllm_config: VllmConfig,
467
+ engine_id: str,
468
+ kv_cache_config: "KVCacheConfig",
469
+ ):
470
self.vllm_config = vllm_config
471
+ self.block_size = vllm_config.cache_config.block_size
472
473
assert vllm_config.kv_transfer_config
474
self.is_kv_producer: bool = (
@@ -457,15 +479,49 @@ class MooncakeConnectorScheduler:
479
)
480
logger.info("Initializing Mooncake Transfer Engine Scheduler %s", engine_id)
481
482
+ self._is_hma_required = (
483
+ not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
484
+ and any(
485
+ not isinstance(g.kv_cache_spec, FullAttentionSpec)
486
+ for g in kv_cache_config.kv_cache_groups
487
+ )
488
+ )
489
+
490
# Requests that need to start recv/send.
491
# New requests are added by update_state_after_alloc in
492
# the scheduler. Used to make metadata passed to Worker.
463
- self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
464
- self._reqs_need_send: dict[ReqId, tuple[Request, list[int]]] = {}
493
+ self._reqs_need_recv: dict[ReqId, tuple[Request, list[list[int]]]] = {}
494
+ self._reqs_need_send: dict[ReqId, tuple[Request, list[list[int]]]] = {}
495
# Reqs to remove from processed set because they're not to send after
496
# remote prefill or aborted.
497
self._reqs_not_processed: set[TransferId] = set()
498
499
+ # Compute sliding window block counts per KV cache group.
500
+ sw_sizes_tokens: list[tuple[int, int]] = [
501
+ (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
502
+ if isinstance(g.kv_cache_spec, SlidingWindowSpec)
503
+ else (0, self.block_size)
504
+ for g in kv_cache_config.kv_cache_groups
505
+ ]
506
+ # cdiv(n_tokens, block_size) gives blocks/window; add 1 to
507
+ # conservatively account for boundary overlap.
508
+ self.blocks_per_sw = [
509
+ cdiv(n_tokens, block_size) + 1 if n_tokens else 0
510
+ for n_tokens, block_size in sw_sizes_tokens
511
+ ]
512
+
513
+ def get_sw_clipped_blocks(
514
+ self,
515
+ block_ids: tuple[list[int], ...] | list[list[int]],
516
+ ) -> list[list[int]]:
517
+ """Clip per-group block IDs to sliding window size."""
518
+ if len(block_ids) == 0 or not self._is_hma_required:
519
+ return list(block_ids)
520
+ return [
521
+ blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks
522
+ for i, blocks in enumerate(block_ids)
523
+ ]
524
+
525
def get_num_new_matched_tokens(
526
self, request: "Request", num_computed_tokens: int
527
) -> tuple[int, bool]:
@@ -530,9 +586,12 @@ class MooncakeConnectorScheduler:
586
# If remote_blocks and num_external_tokens = 0, we have
587
# a full prefix cache hit on the D worker. We need to call
588
# send_notif in _read_blocks to free the memory on the P.
533
- local_block_ids = (
534
- blocks.get_unhashed_block_ids() if num_external_tokens > 0 else []
589
+ unhashed_block_ids = (
590
+ blocks.get_unhashed_block_ids_all_groups()
591
+ if num_external_tokens > 0
592
+ else ()
593
)
594
+ local_block_ids = self.get_sw_clipped_blocks(unhashed_block_ids)
595
# Get unhashed blocks to pull from remote.
596
self._reqs_need_recv[request.request_id] = (request, local_block_ids)
597
else:
@@ -587,7 +646,7 @@ class MooncakeConnectorScheduler:
646
def request_finished(
647
self,
648
request: "Request",
590
- block_ids: list[int],
649
+ block_ids: tuple[list[int], ...],
650
) -> tuple[bool, dict[str, Any] | None]:
651
"""
652
Once a request is finished, determine whether request blocks
@@ -630,10 +689,13 @@ class MooncakeConnectorScheduler:
689
690
# TODO: check whether block_ids actually ever be 0. If not we could
691
# remove the conditional below
633
- delay_free_blocks = len(block_ids) > 0
692
+ delay_free_blocks = any(len(group) > 0 for group in block_ids)
693
694
if delay_free_blocks:
636
- self._reqs_need_send[request.request_id] = (request, block_ids)
695
+ self._reqs_need_send[request.request_id] = (
696
+ request,
697
+ self.get_sw_clipped_blocks(block_ids),
698
+ )
699
700
return delay_free_blocks, None
701
@@ -641,7 +703,12 @@ class MooncakeConnectorScheduler:
703
class MooncakeConnectorWorker:
704
"""Implementation of Worker side methods"""
705
644
- def __init__(self, vllm_config: VllmConfig, engine_id: str):
706
+ def __init__(
707
+ self,
708
+ vllm_config: VllmConfig,
709
+ engine_id: str,
710
+ kv_cache_config: "KVCacheConfig | None" = None,
711
+ ):
712
if TransferEngine is None:
713
logger.error("Mooncake is not available")
714
raise RuntimeError("Mooncake is not available")
@@ -752,6 +819,7 @@ class MooncakeConnectorWorker:
819
self.block_size = vllm_config.cache_config.block_size
820
self.model_config = vllm_config.model_config
821
self.cache_config = vllm_config.cache_config
822
+ self.kv_cache_config = kv_cache_config
823
self.use_mla = self.model_config.use_mla
824
self._sync_block_size_with_kernel()
825
@@ -1103,27 +1171,61 @@ class MooncakeConnectorWorker:
1171
remote_session = f"{agent_meta.remote_hostname}:{agent_meta.remote_port}"
1172
1173
for d_req_id, send_meta in ready_reqs:
1106
- _, remote_block_ids = agent_meta.req_blocks[d_req_id]
1107
- num_remote_blocks = len(remote_block_ids)
1108
- if num_remote_blocks == 0:
1174
+ _, remote_block_ids_per_group = agent_meta.req_blocks[d_req_id]
1175
+
1176
+ if not remote_block_ids_per_group or all(
1177
+ len(g) == 0 for g in remote_block_ids_per_group
1178
+ ):
1179
continue
1180
1111
- local_block_ids = send_meta.local_block_ids
1112
- # Partial prefix cache hit: just read uncomputed blocks.
1113
- num_local_blocks = len(local_block_ids)
1114
- if num_local_blocks < num_remote_blocks:
1181
+ # Per-group partial hit trimming, then flatten.
1182
+ # With HMA, groups share the same KV tensor but use different
1183
+ # block ranges. We trim and concatenate so the coalescer and
1184
+ # address math see one flat block list — same as non-HMA, but
1185
+ # now including blocks from every group.
1186
+ local_block_ids: list[int] = []
1187
+ remote_block_ids: list[int] = []
1188
+ has_block_error = False
1189
+ if len(send_meta.local_block_ids) != len(remote_block_ids_per_group):
1190
logger.error(
1116
- "req %s: local blocks(%d) less than remote blocks(%d)!",
1191
+ "req %s: KV group count mismatch: local=%d, remote=%d",
1192
d_req_id,
1118
- num_local_blocks,
1119
- num_remote_blocks,
1193
+ len(send_meta.local_block_ids),
1194
+ len(remote_block_ids_per_group),
1195
)
1196
+ err_reqs.append(d_req_id)
1197
+ if err_msg is None:
1198
+ err_msg = "KV group count mismatch"
1199
+ continue
1200
+ for local_group, remote_group in zip(
1201
+ send_meta.local_block_ids, remote_block_ids_per_group
1202
+ ):
1203
+ n_local = len(local_group)
1204
+ n_remote = len(remote_group)
1205
+ if n_local < n_remote:
1206
+ logger.error(
1207
+ "req %s: local blocks(%d) < remote blocks(%d) "
1208
+ "in a KV cache group",
1209
+ d_req_id,
1210
+ n_local,
1211
+ n_remote,
1212
+ )
1213
+ has_block_error = True
1214
+ break
1215
+ if n_local > n_remote:
1216
+ # Partial prefix cache hit: just read uncomputed blocks.
1217
+ local_group = local_group[-n_remote:]
1218
+ local_block_ids.extend(local_group)
1219
+ remote_block_ids.extend(remote_group)
1220
+
1221
+ if has_block_error:
1222
err_reqs.append(d_req_id)
1223
if err_msg is None:
1224
err_msg = "P num blocks less than D"
1225
continue
1125
- if num_local_blocks > num_remote_blocks:
1126
- local_block_ids = local_block_ids[-num_remote_blocks:]
1226
+
1227
+ if not local_block_ids:
1228
+ continue
1229
1230
# Group by indices
1231
group_local_block_ids, group_remote_block_ids = group_concurrent_contiguous(
@@ -1215,7 +1317,7 @@ class MooncakeConnectorWorker:
1317
logger.debug(
1318
"Sending kv_caches for request %s (%d blocks) to %s",
1319
d_req_id,
1218
- num_remote_blocks,
1320
+ len(local_block_ids),
1321
remote_session,
1322
)
1323
@@ -1273,23 +1375,24 @@ class MooncakeConnectorWorker:
1375
continue
1376
1377
seen_base_addresses.append(base_addr)
1276
- curr_tensor_size_bytes = cache.nbytes
1378
1379
if tensor_size_bytes is None:
1279
- tensor_size_bytes = curr_tensor_size_bytes
1380
+ tensor_size_bytes = cache.nbytes
1381
self.num_blocks = cache.shape[0]
1382
assert cache.shape[0] == self.num_blocks, (
1383
"All kv cache tensors must have the same number of blocks"
1384
)
1284
- assert curr_tensor_size_bytes % self.num_blocks == 0, (
1285
- "Mooncake expects each kv cache tensor size to be "
1286
- "divisible by the number of blocks."
1287
- )
1288
- self.block_len_per_layer.append(
1289
- curr_tensor_size_bytes // self.num_blocks
1290
- )
1385
+
1386
+ # Use stride-based block length so RDMA reaches the last
1387
+ # block's padding (e.g. DeepseekV4 MLA alignment). stride(0)
1388
+ # reflects the actual byte distance between consecutive
1389
+ # blocks in GPU memory, which matches or exceeds the
1390
+ # shape-based size.
1391
+ block_len = cache.stride(0) * cache.element_size()
1392
+
1393
+ self.block_len_per_layer.append(block_len)
1394
kv_data_ptrs.append(base_addr)
1292
- kv_data_lens.append(curr_tensor_size_bytes)
1395
+ kv_data_lens.append(self.num_blocks * block_len)
1396
1397
self.kv_caches_base_addr = seen_base_addresses
1398
self.seen_base_addresses = seen_base_addresses
vllm/entrypoints/chat_utils.py
+9
@@ -299,6 +299,9 @@ class CustomChatCompletionMessageParam(TypedDict, total=False):
299
tools: list[ChatCompletionFunctionToolParam] | None
300
"""The tools for developer role."""
301
302
+ task: str | None
303
+ """Model-specific task marker. Currently passed through for DeepSeek V4."""
304
+
305
306
ChatCompletionMessageParam: TypeAlias = (
307
OpenAIChatCompletionMessageParam
@@ -333,6 +336,9 @@ class ConversationMessage(TypedDict, total=False):
336
tools: list[ChatCompletionFunctionToolParam] | None
337
"""The tools for developer role."""
338
339
+ task: str | None
340
+ """Model-specific task marker. Currently passed through for DeepSeek V4."""
341
+
342
343
# Passed in by user
344
ChatTemplateContentFormatOption = Literal["auto", "string", "openai"]
@@ -1566,6 +1572,9 @@ def _parse_chat_message_content(
1572
if "name" in message and isinstance(message["name"], str):
1573
result_msg["name"] = message["name"]
1574
1575
+ if "task" in message and isinstance(message["task"], str):
1576
+ result_msg["task"] = message["task"]
1577
+
1578
if role == "developer":
1579
result_msg["tools"] = message.get("tools", None)
1580
return result
vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
+2
@@ -100,6 +100,8 @@ class DeepGemmFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
100
else params.weight_scale,
101
quant_block_shape=tuple(layer.weight_block_size),
102
use_e8m0=self.use_deep_gemm_e8m0,
103
+ is_bmm=getattr(layer, "is_bmm", False),
104
+ bmm_batch_size=getattr(layer, "bmm_batch_size", 0),
105
)
106
replace_parameter(layer, params.WEIGHT, dg_weight)
107
replace_parameter(layer, scale_attr, dg_weight_scale)
vllm/model_executor/layers/attention/mla_attention.py
+16
@@ -1422,6 +1422,20 @@ class MLADims:
1422
def get_mla_dims(model_config: ModelConfig) -> MLADims:
1423
hf_text_config = model_config.hf_text_config
1424
1425
+ # Check if this is a DeepseekV4 config (uses unified head_dim + rope_head_dim)
1426
+ if hasattr(hf_text_config, "compress_ratios"):
1427
+ # DeepseekV4 style config: unified head_dim with rope_head_dim
1428
+ head_dim = hf_text_config.head_dim
1429
+ rope_head_dim = hf_text_config.qk_rope_head_dim
1430
+ return MLADims(
1431
+ q_lora_rank=hf_text_config.q_lora_rank,
1432
+ kv_lora_rank=head_dim,
1433
+ qk_nope_head_dim=head_dim - rope_head_dim,
1434
+ qk_rope_head_dim=rope_head_dim,
1435
+ v_head_dim=head_dim,
1436
+ )
1437
+
1438
+ # DeepseekV2/V3 style config
1439
return MLADims(
1440
q_lora_rank=getattr(hf_text_config, "q_lora_rank", None),
1441
kv_lora_rank=hf_text_config.kv_lora_rank,
@@ -2191,6 +2205,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2205
qk_head_dim: int,
2206
v_head_dim: int,
2207
kv_b_proj: ColumnParallelLinear,
2208
+ # DSV3.2 MLA Specific Arguments
2209
indexer: object | None = None,
2210
q_pad_num_heads: int | None = None,
2211
) -> None:
@@ -2213,6 +2228,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2228
self.indexer = indexer
2229
self.q_pad_num_heads = q_pad_num_heads
2230
self.supports_quant_query_input = True
2231
+ self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
2232
2233
# Use flashinfer's optimized concat_mla_k kernel when available.
2234
# The kernel is optimized for DeepSeek V3 dimensions:
vllm/model_executor/layers/deepseek_compressor.py
new
+438
@@ -0,0 +1,438 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+from typing import Any, ClassVar, cast
6
+
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.config import VllmConfig, get_current_vllm_config
11
+from vllm.forward_context import get_forward_context
12
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
13
+from vllm.model_executor.layers.layernorm import RMSNorm
14
+from vllm.model_executor.layers.linear import (
15
+ MergedColumnParallelLinear,
16
+)
17
+from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32
18
+from vllm.platforms import current_platform
19
+from vllm.triton_utils import tl, triton
20
+from vllm.v1.attention.backend import (
21
+ AttentionBackend,
22
+ AttentionCGSupport,
23
+ AttentionMetadataBuilder,
24
+ CommonAttentionMetadata,
25
+ MultipleOf,
26
+)
27
+from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
28
+ _fused_kv_compress_norm_rope_insert_indexer_attn,
29
+ _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
30
+ _fused_kv_compress_norm_rope_insert_sparse_attn,
31
+)
32
+from vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q import (
33
+ MXFP4_BLOCK_SIZE,
34
+)
35
+from vllm.v1.kv_cache_interface import (
36
+ KVCacheSpec,
37
+ MLAAttentionSpec,
38
+ SlidingWindowMLASpec,
39
+)
40
+
41
+
42
+class CompressorBackend(AttentionBackend):
43
+ def __init__(self):
44
+ super().__init__()
45
+
46
+ @staticmethod
47
+ def get_name() -> str:
48
+ return "CompressorBackend"
49
+
50
+ @staticmethod
51
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
52
+ return [MultipleOf(1)]
53
+
54
+ @classmethod
55
+ def get_supported_head_sizes(cls) -> list[int]:
56
+ return [512, 1024]
57
+
58
+ @staticmethod
59
+ def get_builder_cls() -> type["CompressorMetadataBuilder"]:
60
+ return CompressorMetadataBuilder
61
+
62
+ @staticmethod
63
+ def get_kv_cache_shape(
64
+ num_blocks: int,
65
+ block_size: int,
66
+ num_kv_heads: int,
67
+ head_size: int,
68
+ cache_dtype_str: str = "auto",
69
+ ) -> tuple[int, ...]:
70
+ assert num_kv_heads == 1
71
+ return (num_blocks, block_size, head_size)
72
+
73
+ @staticmethod
74
+ def get_kv_cache_stride_order(
75
+ include_num_layers_dimension: bool = False,
76
+ ) -> tuple[int, ...]:
77
+ if include_num_layers_dimension:
78
+ return (0, 1, 2, 3)
79
+ return (0, 1, 2)
80
+
81
+
82
+@dataclass
83
+class CompressorMetadata:
84
+ block_table: torch.Tensor
85
+ slot_mapping: torch.Tensor
86
+ block_size: int
87
+
88
+ token_to_req_indices: torch.Tensor | None = None # [num_tokens]
89
+
90
+
91
+class CompressorMetadataBuilder(AttentionMetadataBuilder):
92
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
93
+
94
+ def __init__(self, *args, **kwargs):
95
+ super().__init__(*args, **kwargs)
96
+ assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec)
97
+ mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec)
98
+ self.block_size = mla_spec.block_size
99
+
100
+ self.token_to_req_indices = torch.zeros(
101
+ self.vllm_config.scheduler_config.max_num_batched_tokens,
102
+ dtype=torch.int32,
103
+ device=self.device,
104
+ )
105
+
106
+ def build(
107
+ self,
108
+ common_prefix_len: int,
109
+ common_attn_metadata: CommonAttentionMetadata,
110
+ fast_build: bool = False,
111
+ ) -> CompressorMetadata:
112
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
113
+ num_reqs = common_attn_metadata.num_reqs
114
+ query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
115
+ x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory()
116
+ token_to_req_indices = self.token_to_req_indices[: x.shape[0]]
117
+ token_to_req_indices.copy_(x, non_blocking=True)
118
+ return CompressorMetadata(
119
+ block_table=common_attn_metadata.block_table_tensor.clamp_(min=0),
120
+ slot_mapping=common_attn_metadata.slot_mapping,
121
+ block_size=self.block_size,
122
+ token_to_req_indices=token_to_req_indices,
123
+ )
124
+
125
+
126
+class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
127
+ def __init__(
128
+ self,
129
+ state_dim: int,
130
+ dtype: torch.dtype,
131
+ compress_ratio: int,
132
+ prefix: str,
133
+ ):
134
+ super().__init__()
135
+ self.state_dim = state_dim
136
+ self.dtype = dtype
137
+ self.prefix = prefix
138
+ self.kv_cache = torch.tensor([])
139
+ compilation_config = get_current_vllm_config().compilation_config
140
+ if prefix in compilation_config.static_forward_context:
141
+ raise ValueError(f"Duplicate layer name: {prefix}")
142
+ compilation_config.static_forward_context[prefix] = self
143
+
144
+ assert self.dtype == torch.float32
145
+ assert compress_ratio in [4, 128]
146
+ coff = 1 + (compress_ratio == 4)
147
+ self.sliding_window = coff * compress_ratio
148
+ # Block size is constrained by tensor sharing between compressor states
149
+ # and KV blocks. Since compressor states share the same physical tensor
150
+ # as KV blocks, they must use the same page size.
151
+ # The KV block shape [256//4, head_dim] = [64, 584] determines:
152
+ # - C4 compressor block shape [4, 2*512*2*4] -> block_size = 4
153
+ # - C128 compressor block shape [8, 512*2*4] -> block_size = 8
154
+ # TODO(yifan): make block size automatically determined and configurable.
155
+ if compress_ratio == 4:
156
+ self.block_size = 4
157
+ elif compress_ratio == 128:
158
+ self.block_size = 8
159
+ else:
160
+ raise ValueError(f"Invalid compress ratio: {compress_ratio}")
161
+
162
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
163
+ return SlidingWindowMLASpec( # only has one vector instead of K + V
164
+ block_size=self.block_size,
165
+ num_kv_heads=1,
166
+ head_size=self.state_dim,
167
+ dtype=self.dtype,
168
+ sliding_window=self.sliding_window,
169
+ alignment=576, # NOTE: FlashMLA requires 576B alignment
170
+ )
171
+
172
+ def forward(self): ...
173
+
174
+ def get_attn_backend(self) -> type[AttentionBackend]:
175
+ return CompressorBackend
176
+
177
+
178
+class DeepseekCompressor(nn.Module):
179
+ def __init__(
180
+ self,
181
+ vllm_config: VllmConfig,
182
+ compress_ratio: int,
183
+ hidden_size: int,
184
+ head_dim: int,
185
+ rotate: bool = False,
186
+ prefix: str = "",
187
+ k_cache_prefix="",
188
+ use_fp4_cache: bool = False,
189
+ ):
190
+ super().__init__()
191
+ self.compress_ratio = compress_ratio
192
+ self.hidden_size = hidden_size
193
+ self.head_dim = head_dim
194
+ self.rotate = rotate
195
+ self.prefix = prefix
196
+ self.k_cache_prefix = k_cache_prefix
197
+ self.use_fp4_cache = use_fp4_cache
198
+
199
+ config = vllm_config.model_config.hf_config
200
+ self.rope_head_dim = config.qk_rope_head_dim
201
+ self.nope_head_dim = self.head_dim - self.rope_head_dim
202
+ self.rms_norm_eps = config.rms_norm_eps
203
+ self.device = current_platform.device_type
204
+ self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
205
+ self.max_model_len = vllm_config.model_config.max_model_len
206
+
207
+ self.overlap = compress_ratio == 4
208
+ self.coff = 1 + self.overlap
209
+
210
+ state_dtype = torch.float32
211
+ self.ape = nn.Parameter(
212
+ torch.empty(
213
+ (compress_ratio, self.coff * self.head_dim),
214
+ dtype=state_dtype,
215
+ device=self.device,
216
+ ),
217
+ requires_grad=False,
218
+ )
219
+
220
+ self.fused_wkv_wgate = MergedColumnParallelLinear(
221
+ self.hidden_size,
222
+ [self.coff * self.head_dim, self.coff * self.head_dim],
223
+ bias=False,
224
+ return_bias=False,
225
+ quant_config=None,
226
+ disable_tp=True,
227
+ prefix=f"{prefix}.fused_wkv_wgate",
228
+ )
229
+ self.norm = RMSNorm(self.head_dim, self.rms_norm_eps)
230
+
231
+ self.state_cache = CompressorStateCache(
232
+ state_dim=2 * self.coff * self.head_dim, # kv_state + score_state
233
+ dtype=state_dtype,
234
+ compress_ratio=compress_ratio,
235
+ prefix=f"{prefix}.state_cache",
236
+ )
237
+
238
+ # Save reference to static_forward_context for forward-time KV cache lookup.
239
+ # get_current_vllm_config() is only available during __init__, not forward.
240
+ self._static_forward_context = (
241
+ vllm_config.compilation_config.static_forward_context
242
+ )
243
+
244
+ if self.head_dim == 512:
245
+ assert not use_fp4_cache, (
246
+ "MXFP4 cache is only supported for indexer (head=128)"
247
+ )
248
+ self._fused_kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
249
+ self._quant_block = 64
250
+ self._token_stride = self.nope_head_dim + self.rope_head_dim * 2
251
+ self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad
252
+ self._num_warps = 4
253
+ elif self.head_dim == 128:
254
+ if use_fp4_cache:
255
+ self._fused_kernel = (
256
+ _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
257
+ )
258
+ self._quant_block = MXFP4_BLOCK_SIZE
259
+ self._token_stride = self.head_dim // 2
260
+ self._scale_dim = self.head_dim // MXFP4_BLOCK_SIZE
261
+ else:
262
+ self._fused_kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
263
+ self._quant_block = 128
264
+ self._token_stride = self.head_dim
265
+ self._scale_dim = 4 # single float32 scale
266
+ self._num_warps = 1
267
+ else:
268
+ raise ValueError(
269
+ f"Unsupported head_dim for fused quant+cache: {self.head_dim}"
270
+ )
271
+
272
+ def forward(
273
+ self,
274
+ # [num_tokens, hidden_size]
275
+ x: torch.Tensor,
276
+ # [num_tokens]
277
+ positions: torch.Tensor,
278
+ rotary_emb,
279
+ ) -> None:
280
+ num_tokens, _ = x.shape
281
+ # bf16 weights/activations but fp32 output for numerical stability of
282
+ # the downstream compressor math.
283
+ kv_score = cublas_gemm_bf16_bf16_fp32(x, self.fused_wkv_wgate.weight)
284
+ # Each of shape [num_tokens, coff * self.head_dim]
285
+ # input bf16, output are fp32
286
+ kv, score = kv_score.split(
287
+ [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1
288
+ )
289
+
290
+ # Get the metadata and handle dummy profiling run.
291
+ attn_metadata = get_forward_context().attn_metadata
292
+ if not isinstance(attn_metadata, dict):
293
+ return
294
+
295
+ state_metadata = cast(
296
+ CompressorMetadata, attn_metadata[self.state_cache.prefix]
297
+ )
298
+ token_to_req_indices = state_metadata.token_to_req_indices
299
+ slot_mapping = state_metadata.slot_mapping
300
+ num_actual = slot_mapping.shape[0]
301
+ block_table = state_metadata.block_table
302
+ block_size = state_metadata.block_size
303
+
304
+ # [num_blocks, block_size, kv_dim+score_dim], where kv_dim == score_dim
305
+ state_cache = self.state_cache.kv_cache
306
+ # kv_state stored in first half, score_state stored in second half
307
+ state_width = state_cache.shape[-1] // 2
308
+
309
+ # Store the KV and score (with fused APE addition) in the state.
310
+ # NOTE: PDL is disabled — both this kernel and _fused_kernel below
311
+ # depend on preceding kernel outputs (kv/score from the cublas GEMM;
312
+ # state_cache from this kernel) but neither emits/waits on PDL grid
313
+ # dependency primitives, so launch_pdl=True caused a read-after-write
314
+ # race and non-deterministic output.
315
+ _save_partial_states_kernel[(num_actual,)](
316
+ kv,
317
+ kv.stride(0),
318
+ score,
319
+ score.stride(0),
320
+ self.ape,
321
+ self.ape.stride(0),
322
+ positions,
323
+ state_cache,
324
+ state_cache.stride(0),
325
+ state_cache.stride(1),
326
+ slot_mapping,
327
+ block_size,
328
+ HEAD_SIZE=kv.shape[-1],
329
+ TRITON_BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]),
330
+ STATE_WIDTH=state_width,
331
+ COMPRESS_RATIO=self.compress_ratio,
332
+ launch_pdl=False,
333
+ )
334
+
335
+ # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
336
+ # RoPE requirements (kernel applies forward GPT-J style rotation):
337
+ # - is_neox_style=False (interleaved pairs, NOT split-half)
338
+ # - cos_sin_cache layout: [max_pos, rope_head_dim] with first half cos,
339
+ # second half sin (per-pair, length rope_head_dim // 2 each)
340
+ # - applied to LAST rope_head_dim elements of head_dim
341
+ # - position used: (positions // compress_ratio) * compress_ratio
342
+ cos_sin_cache = rotary_emb.cos_sin_cache
343
+ k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
344
+ kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache
345
+
346
+ self._fused_kernel[(num_actual,)](
347
+ # state cache
348
+ state_cache,
349
+ state_cache.stride(0),
350
+ state_cache.stride(1),
351
+ # metadata
352
+ token_to_req_indices,
353
+ positions,
354
+ slot_mapping,
355
+ block_table,
356
+ block_table.stride(0),
357
+ block_size,
358
+ # RMSNorm
359
+ self.norm.weight,
360
+ self.rms_norm_eps,
361
+ # RoPE
362
+ cos_sin_cache,
363
+ cos_sin_cache.stride(0),
364
+ # KV cache
365
+ kv_cache,
366
+ k_cache_metadata.slot_mapping,
367
+ kv_cache.shape[1], # paged KV cache block size (tokens per block)
368
+ # constexprs
369
+ HEAD_SIZE=self.head_dim,
370
+ TRITON_BLOCK_SIZE=triton.next_power_of_2(self.head_dim),
371
+ STATE_WIDTH=state_width,
372
+ COMPRESS_RATIO=self.compress_ratio,
373
+ OVERLAP=self.overlap,
374
+ ROPE_HEAD_DIM=self.rope_head_dim,
375
+ FP8_MAX=448.0,
376
+ QUANT_BLOCK=self._quant_block,
377
+ TOKEN_STRIDE=self._token_stride,
378
+ SCALE_DIM=self._scale_dim,
379
+ KV_BLOCK_STRIDE=kv_cache.stride(0),
380
+ num_warps=self._num_warps,
381
+ launch_pdl=False,
382
+ )
383
+
384
+
385
+@triton.jit
386
+def _save_partial_states_kernel(
387
+ kv_ptr,
388
+ kv_stride,
389
+ score_ptr,
390
+ score_stride,
391
+ ape_ptr,
392
+ ape_stride,
393
+ positions_ptr,
394
+ state_cache_ptr,
395
+ state_cache_stride0,
396
+ state_cache_stride1,
397
+ slot_mapping_ptr,
398
+ block_size,
399
+ HEAD_SIZE: tl.constexpr,
400
+ TRITON_BLOCK_SIZE: tl.constexpr,
401
+ # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
402
+ STATE_WIDTH: tl.constexpr,
403
+ COMPRESS_RATIO: tl.constexpr,
404
+):
405
+ token_idx = tl.program_id(0)
406
+ slot_id = tl.load(slot_mapping_ptr + token_idx)
407
+
408
+ # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
409
+ # by vLLM). During CUDA graph replay the batch may contain padding
410
+ # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
411
+ # illegal memory access.
412
+ if slot_id < 0:
413
+ return
414
+
415
+ block_idx = slot_id // block_size
416
+ pos_in_block = slot_id % block_size
417
+ base_ptr = (
418
+ state_cache_ptr
419
+ + block_idx * state_cache_stride0
420
+ + pos_in_block * state_cache_stride1
421
+ )
422
+
423
+ block = tl.arange(0, TRITON_BLOCK_SIZE)
424
+ mask = block < HEAD_SIZE
425
+
426
+ kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
427
+ tl.store(base_ptr + block, kv, mask=mask)
428
+
429
+ # Fused: score += ape[position % compress_ratio]
430
+ position = tl.load(positions_ptr + token_idx)
431
+ ape_row = position % COMPRESS_RATIO
432
+ ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
433
+ score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
434
+ tl.store(
435
+ base_ptr + STATE_WIDTH + block,
436
+ score + ape,
437
+ mask=mask,
438
+ )
vllm/model_executor/layers/deepseek_v4_attention.py
new
+1076
@@ -0,0 +1,1076 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+DeepseekV4 MLA Attention Layer
5
+"""
6
+
7
+from dataclasses import dataclass
8
+from typing import TYPE_CHECKING, cast
9
+
10
+import torch
11
+import torch.nn as nn
12
+import torch.nn.functional as F
13
+from transformers import DeepseekV2Config, DeepseekV3Config
14
+
15
+from vllm.model_executor.layers.linear import (
16
+ ReplicatedLinear,
17
+)
18
+from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
19
+from vllm.utils.deep_gemm import fp8_einsum
20
+from vllm.utils.torch_utils import direct_register_custom_op
21
+from vllm.v1.attention.ops.deepseek_v4_ops import (
22
+ combine_topk_swa_indices,
23
+ compute_global_topk_indices_and_lens,
24
+ dequantize_and_gather_k_cache,
25
+ fused_indexer_q_rope_quant,
26
+ fused_inv_rope_fp8_quant,
27
+ fused_q_kv_rmsnorm,
28
+)
29
+
30
+if TYPE_CHECKING:
31
+ from vllm.v1.attention.backends.mla.sparse_swa import (
32
+ DeepseekSparseSWAMetadata,
33
+ )
34
+
35
+from vllm.config import (
36
+ CacheConfig,
37
+ VllmConfig,
38
+ get_current_vllm_config,
39
+)
40
+from vllm.distributed import get_tensor_model_parallel_world_size
41
+from vllm.forward_context import ForwardContext, get_forward_context
42
+from vllm.logger import init_logger
43
+from vllm.model_executor.custom_op import PluggableLayer
44
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
45
+from vllm.model_executor.layers.deepseek_compressor import DeepseekCompressor
46
+from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
47
+from vllm.model_executor.layers.quantization import QuantizationConfig
48
+from vllm.model_executor.layers.quantization.input_quant_fp8 import (
49
+ QuantFP8,
50
+)
51
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
52
+ GroupShape,
53
+)
54
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
55
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
56
+from vllm.v1.attention.backends.mla.flashmla_sparse import (
57
+ DeepseekV4FlashMLASparseBackend,
58
+ FlashMLASparseBackend,
59
+ FlashMLASparseMetadata,
60
+)
61
+from vllm.v1.attention.backends.mla.indexer import (
62
+ DeepseekV4IndexerBackend,
63
+ get_max_prefill_buffer_size,
64
+)
65
+from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
66
+from vllm.v1.attention.ops.flashmla import (
67
+ flash_mla_sparse_fwd,
68
+ flash_mla_with_kvcache,
69
+)
70
+from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
71
+from vllm.v1.worker.workspace import current_workspace_manager
72
+
73
+logger = init_logger(__name__)
74
+
75
+# Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
76
+# workspace allocated at _forward_prefill (and the matching profile-time
77
+# reservation in attention_impl's dummy-run branch).
78
+PREFILL_CHUNK_SIZE = 4
79
+
80
+
81
+@dataclass
82
+class DeepseekV4MLAModules:
83
+ """Modules used in DeepseekV4 MLA."""
84
+
85
+ vllm_config: VllmConfig
86
+ fused_wqa_wkv: torch.nn.Module
87
+ q_norm: torch.nn.Module
88
+ wq_b: torch.nn.Module
89
+ kv_norm: torch.nn.Module
90
+ wo_a: torch.nn.Module
91
+ wo_b: torch.nn.Module
92
+ attn_sink: torch.nn.Module
93
+ rotary_emb: torch.nn.Module
94
+ indexer: torch.nn.Module | None
95
+ indexer_rotary_emb: torch.nn.Module
96
+ topk_indices_buffer: torch.Tensor | None
97
+ aux_stream: torch.cuda.Stream | None = None
98
+
99
+
100
+# --8<-- [start:multi_head_latent_attention]
101
+@PluggableLayer.register("deepseek_v4_multi_head_latent_attention")
102
+class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
103
+ """Pluggable MLA layer which allows OOT backends to add
104
+ custom implementations of the outer MLA layer (including rope & o_proj).
105
+ Note that currently oot platforms can still use CustomOp.register_oot to
106
+ replace MLA layer entirely, although we use PluggableLayer to register
107
+ this layer now.
108
+
109
+ This class takes positions and hidden_states as input.
110
+ The input tensors can either contain prefill tokens or decode tokens.
111
+ The class does the following:
112
+
113
+ 1. MLA Preprocess.
114
+ 2. Perform multi-head attention to prefill tokens and
115
+ multi-query attention to decode tokens separately.
116
+ 3. Return the output tensor.
117
+ """
118
+
119
+ # --8<-- [end:multi_head_latent_attention]
120
+
121
+ def __init__(
122
+ self,
123
+ hidden_size: int,
124
+ num_heads: int,
125
+ head_dim: int,
126
+ scale: float,
127
+ qk_nope_head_dim: int,
128
+ qk_rope_head_dim: int,
129
+ v_head_dim: int,
130
+ q_lora_rank: int | None,
131
+ kv_lora_rank: int,
132
+ o_lora_rank: int | None,
133
+ mla_modules: DeepseekV4MLAModules,
134
+ window_size: int,
135
+ compress_ratio: int | None,
136
+ cache_config: CacheConfig | None = None,
137
+ quant_config: QuantizationConfig | None = None,
138
+ prefix: str = "",
139
+ ) -> None:
140
+ super().__init__()
141
+ self.hidden_size = hidden_size
142
+ self.n_local_heads = num_heads
143
+ self.head_dim = head_dim
144
+ self.scale = scale
145
+
146
+ # FlashMLA sparse kernel only supports 64 or 128 heads; pad up to the
147
+ # next supported size. Must match DeepseekV4MLAAttention.padded_heads.
148
+ if num_heads <= 64:
149
+ self.padded_heads = 64
150
+ elif num_heads <= 128:
151
+ self.padded_heads = 128
152
+ else:
153
+ raise ValueError(
154
+ f"DeepseekV4 attention does not support {num_heads} heads "
155
+ "(must be <= 128)."
156
+ )
157
+
158
+ self.q_lora_rank = q_lora_rank
159
+ self.kv_lora_rank = kv_lora_rank
160
+ self.window_size = window_size
161
+ self.compress_ratio = compress_ratio if compress_ratio is not None else 1
162
+ self.prefix = prefix
163
+
164
+ # Extract config from vllm_config
165
+ config = mla_modules.vllm_config.model_config.hf_config
166
+ tp_size = get_tensor_model_parallel_world_size()
167
+
168
+ # DeepseekV4-specific attributes (num_heads is already TP-adjusted)
169
+ self.eps = config.rms_norm_eps
170
+ self.rope_head_dim = config.qk_rope_head_dim
171
+ self.nope_head_dim = head_dim - self.rope_head_dim
172
+ self.n_local_groups = config.o_groups // tp_size
173
+ self.o_lora_rank = config.o_lora_rank
174
+
175
+ # Store projection modules
176
+ self.fused_wqa_wkv = mla_modules.fused_wqa_wkv
177
+ self.q_norm = mla_modules.q_norm
178
+ self.wq_b = mla_modules.wq_b
179
+
180
+ self.kv_norm = mla_modules.kv_norm
181
+ self.wo_a = mla_modules.wo_a
182
+
183
+ self._wo_a_act_quant = QuantFP8(
184
+ static=False,
185
+ group_shape=GroupShape(1, 128),
186
+ use_ue8m0=True,
187
+ )
188
+ # Bypass packed-for-deepgemm path — we need FP32 scales (not packed
189
+ # INT32) so fp8_einsum can handle layout transform internally.
190
+ self._wo_a_act_quant.use_deep_gemm_supported = False
191
+ self.wo_b = mla_modules.wo_b
192
+
193
+ # Pick fp8_einsum recipe based on GPU arch:
194
+ # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128
195
+ # SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1
196
+ from vllm.platforms import current_platform
197
+
198
+ cap = current_platform.get_device_capability()
199
+ assert cap is not None, "DeepseekV4 attention requires a CUDA device"
200
+ self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128)
201
+ self._tma_aligned_scales = cap.major >= 10
202
+
203
+ self.rotary_emb = mla_modules.rotary_emb
204
+ self.indexer_rotary_emb = mla_modules.indexer_rotary_emb
205
+ self.topk_indices_buffer = mla_modules.topk_indices_buffer
206
+
207
+ self.indexer = mla_modules.indexer
208
+
209
+ # Per-head RMS normalization for Q (no learnable weights)
210
+ self.q_head_norm = RMSNorm(head_dim, eps=self.eps, has_weight=False)
211
+
212
+ # TODO(yifan): currently hardcoded for FP8 sparse, make it more generic
213
+ head_bytes = (
214
+ self.nope_head_dim # 448 fp8 NoPE
215
+ + self.rope_head_dim * 2 # 64 bf16 RoPE
216
+ + self.nope_head_dim // 64 # 7B scale factors
217
+ + 1 # 1B pad
218
+ )
219
+
220
+ self.aux_stream = mla_modules.aux_stream
221
+ self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
222
+
223
+ assert cache_config is not None, "DeepseekV4 attention requires cache_config"
224
+ self.swa_cache_layer = DeepseekV4SWACache(
225
+ head_dim=self.head_dim,
226
+ window_size=self.window_size,
227
+ dtype=torch.uint8,
228
+ prefix=f"{prefix}.swa_cache",
229
+ cache_config=cache_config,
230
+ )
231
+
232
+ self.mla_attn = DeepseekV4MLAAttention(
233
+ num_heads=self.n_local_heads,
234
+ head_dim=self.head_dim,
235
+ scale=self.scale,
236
+ qk_nope_head_dim=self.nope_head_dim,
237
+ qk_rope_head_dim=self.rope_head_dim,
238
+ q_lora_rank=self.q_lora_rank,
239
+ kv_lora_rank=self.kv_lora_rank,
240
+ compress_ratio=self.compress_ratio,
241
+ window_size=self.window_size,
242
+ head_bytes=head_bytes,
243
+ swa_cache_layer=self.swa_cache_layer,
244
+ attn_sink=mla_modules.attn_sink, # already padded with -inf
245
+ cache_config=cache_config,
246
+ quant_config=quant_config,
247
+ prefix=prefix,
248
+ indexer=self.indexer,
249
+ topk_indices_buffer=self.topk_indices_buffer,
250
+ )
251
+ # Register this layer in the compilation config's static forward context
252
+ # This allows the custom op to retrieve the layer during execution
253
+ compilation_config = mla_modules.vllm_config.compilation_config
254
+ # HACK
255
+ self.layer_name = prefix + ".deepseek_v4_multi_head_latent_attention"
256
+ if self.layer_name in compilation_config.static_forward_context:
257
+ raise ValueError(f"Duplicate layer name: {self.layer_name}")
258
+ compilation_config.static_forward_context[self.layer_name] = self
259
+
260
+ # Create the compressor for layers with compress_ratio > 1; after
261
+ # creating the DeepseekV4MLAAttention layer to get its cache.
262
+ self.compressor = None
263
+ if self.compress_ratio > 1:
264
+ self.compressor = DeepseekCompressor(
265
+ vllm_config=mla_modules.vllm_config,
266
+ compress_ratio=self.compress_ratio,
267
+ hidden_size=self.hidden_size,
268
+ head_dim=self.head_dim,
269
+ rotate=True,
270
+ prefix=f"{prefix}.compressor",
271
+ k_cache_prefix=self.mla_attn.prefix,
272
+ )
273
+
274
+ def forward(
275
+ self,
276
+ positions: torch.Tensor,
277
+ hidden_states: torch.Tensor,
278
+ llama_4_scaling: torch.Tensor | None = None,
279
+ ) -> torch.Tensor:
280
+ qr_kv, _ = self.fused_wqa_wkv(hidden_states)
281
+ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
282
+
283
+ # Pre-allocate attention output with FlashMLA-padded head count.
284
+ # The op writes into `o_padded`; we slice to n_local_heads after.
285
+ num_tokens = hidden_states.shape[0]
286
+ o_padded = torch.empty(
287
+ (num_tokens, self.padded_heads, self.head_dim),
288
+ dtype=hidden_states.dtype,
289
+ device=hidden_states.device,
290
+ )
291
+
292
+ # Attention (inside custom op for torch.compile boundary)
293
+ torch.ops.vllm.deepseek_v4_attention(
294
+ hidden_states,
295
+ qr,
296
+ kv,
297
+ positions,
298
+ o_padded,
299
+ self.layer_name,
300
+ )
301
+ o = o_padded[:, : self.n_local_heads, :]
302
+
303
+ # O projection: inverse RoPE + FP8 quant + einsum + wo_b
304
+ o_fp8, o_scale = fused_inv_rope_fp8_quant(
305
+ o,
306
+ positions,
307
+ self.rotary_emb.cos_sin_cache,
308
+ n_groups=self.n_local_groups,
309
+ heads_per_group=self.n_local_heads // self.n_local_groups,
310
+ nope_dim=self.nope_head_dim,
311
+ rope_dim=self.rope_head_dim,
312
+ tma_aligned_scales=self._tma_aligned_scales,
313
+ )
314
+
315
+ wo_a_fp8 = self.wo_a.weight
316
+ wo_a_scale = self.wo_a.weight_scale_inv
317
+
318
+ z = torch.empty(
319
+ (num_tokens, self.n_local_groups, self.o_lora_rank),
320
+ device=o.device,
321
+ dtype=torch.bfloat16,
322
+ )
323
+ torch.ops.vllm.deepseek_v4_fp8_einsum(
324
+ o_fp8,
325
+ o_scale,
326
+ wo_a_fp8,
327
+ wo_a_scale,
328
+ z,
329
+ "bhr,hdr->bhd",
330
+ list(self._einsum_recipe),
331
+ )
332
+
333
+ return self.wo_b(z.flatten(1))
334
+
335
+ def attention_impl(
336
+ self,
337
+ hidden_states: torch.Tensor,
338
+ qr: torch.Tensor,
339
+ kv: torch.Tensor,
340
+ positions: torch.Tensor,
341
+ out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place
342
+ ) -> None:
343
+ forward_context = get_forward_context()
344
+ attn_metadata = forward_context.attn_metadata
345
+
346
+ qr, kv = fused_q_kv_rmsnorm(
347
+ qr,
348
+ kv,
349
+ self.q_norm.weight.data,
350
+ self.kv_norm.weight.data,
351
+ self.eps,
352
+ )
353
+ q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
354
+
355
+ # Overlap kv_insert with whichever of indexer/compressor is present.
356
+ # Indexer implies compressor; when both exist, compressor rides on the
357
+ # aux stream alongside kv_insert so the heavy indexer owns default.
358
+ if self.indexer is not None:
359
+ indexer = self.indexer
360
+ # Local ref so the closure keeps a non-None type for mypy.
361
+ assert self.compressor is not None
362
+ compressor = self.compressor
363
+
364
+ def kv_insert_and_compress() -> None:
365
+ self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
366
+ compressor(hidden_states, positions, self.rotary_emb)
367
+
368
+ maybe_execute_in_parallel(
369
+ lambda: indexer(hidden_states, qr, positions, self.indexer_rotary_emb),
370
+ kv_insert_and_compress,
371
+ self.ln_events[0],
372
+ self.ln_events[1],
373
+ self.aux_stream,
374
+ )
375
+ elif self.compressor is not None:
376
+ # Compressor on default, kv_insert on aux.
377
+ compressor = self.compressor
378
+ maybe_execute_in_parallel(
379
+ lambda: compressor(hidden_states, positions, self.rotary_emb),
380
+ lambda: self._fused_qnorm_rope_kv_insert(
381
+ q, kv, positions, attn_metadata
382
+ ),
383
+ self.ln_events[0],
384
+ self.ln_events[1],
385
+ self.aux_stream,
386
+ )
387
+ else:
388
+ # SWA-only layer: no compressor, no overlap.
389
+ self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
390
+
391
+ # Handle dummy run (no metadata).
392
+ if not isinstance(attn_metadata, dict):
393
+ # Reserve _forward_prefill's bf16-gather workspace; the dummy
394
+ # run returns before mla_attn runs, so without this the shared
395
+ # workspace locks below the real prefill size.
396
+ sub = self.mla_attn
397
+ swa_only = sub.compress_ratio <= 1
398
+ N = (
399
+ 0
400
+ if swa_only
401
+ else (sub.max_model_len + sub.compress_ratio - 1) // sub.compress_ratio
402
+ )
403
+ M = N + sub.window_size + sub.max_num_batched_tokens
404
+ current_workspace_manager().get_simultaneous(
405
+ ((PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
406
+ )
407
+ out.zero_()
408
+ return
409
+
410
+ # Pad q to FlashMLA-required head count (64 or 128)
411
+ if self.n_local_heads < self.padded_heads:
412
+ pad_size = self.padded_heads - self.n_local_heads
413
+ q = F.pad(q, (0, 0, 0, pad_size), value=0.0)
414
+
415
+ # MLA attention writes into the pre-allocated `out` buffer
416
+ # ([num_tokens, padded_heads, head_dim]).
417
+ self.mla_attn(q, kv, positions, output=out)
418
+
419
+ def _fused_qnorm_rope_kv_insert(
420
+ self,
421
+ q: torch.Tensor,
422
+ kv: torch.Tensor,
423
+ positions: torch.Tensor,
424
+ attn_metadata: (
425
+ dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
426
+ ),
427
+ ) -> None:
428
+ if not isinstance(attn_metadata, dict):
429
+ return
430
+
431
+ swa_metadata = cast(
432
+ "DeepseekSparseSWAMetadata | None",
433
+ attn_metadata.get(self.swa_cache_layer.prefix),
434
+ )
435
+ assert swa_metadata is not None
436
+
437
+ swa_kv_cache = self.swa_cache_layer.kv_cache
438
+ swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
439
+
440
+ # Horizontally fused:
441
+ # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE
442
+ # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert
443
+ # kv is unchanged; mla_attn reads kv solely via swa_kv_cache.
444
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
445
+ q,
446
+ kv,
447
+ swa_kv_cache_2d,
448
+ swa_metadata.slot_mapping,
449
+ positions.to(torch.int64),
450
+ self.rotary_emb.cos_sin_cache,
451
+ self.eps,
452
+ swa_metadata.block_size,
453
+ )
454
+
455
+
456
+def deepseek_v4_attention(
457
+ hidden_states: torch.Tensor,
458
+ qr: torch.Tensor,
459
+ kv: torch.Tensor,
460
+ positions: torch.Tensor,
461
+ out: torch.Tensor,
462
+ layer_name: str,
463
+) -> None:
464
+ forward_context: ForwardContext = get_forward_context()
465
+ self = forward_context.no_compile_layers[layer_name]
466
+ self.attention_impl(hidden_states, qr, kv, positions, out)
467
+
468
+
469
+def deepseek_v4_attention_fake(
470
+ hidden_states: torch.Tensor,
471
+ qr: torch.Tensor,
472
+ kv: torch.Tensor,
473
+ positions: torch.Tensor,
474
+ out: torch.Tensor,
475
+ layer_name: str,
476
+) -> None:
477
+ return None
478
+
479
+
480
+direct_register_custom_op(
481
+ op_name="deepseek_v4_attention",
482
+ op_func=deepseek_v4_attention,
483
+ mutates_args=["out"],
484
+ fake_impl=deepseek_v4_attention_fake,
485
+)
486
+
487
+
488
+def deepseek_v4_fp8_einsum(
489
+ a: torch.Tensor,
490
+ a_scale: torch.Tensor,
491
+ b: torch.Tensor,
492
+ b_scale: torch.Tensor,
493
+ out: torch.Tensor,
494
+ equation: str,
495
+ recipe: list[int],
496
+) -> None:
497
+ fp8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe))
498
+
499
+
500
+def deepseek_v4_fp8_einsum_fake(
501
+ a: torch.Tensor,
502
+ a_scale: torch.Tensor,
503
+ b: torch.Tensor,
504
+ b_scale: torch.Tensor,
505
+ out: torch.Tensor,
506
+ equation: str,
507
+ recipe: list[int],
508
+) -> None:
509
+ return None
510
+
511
+
512
+direct_register_custom_op(
513
+ op_name="deepseek_v4_fp8_einsum",
514
+ op_func=deepseek_v4_fp8_einsum,
515
+ mutates_args=["out"],
516
+ fake_impl=deepseek_v4_fp8_einsum_fake,
517
+)
518
+
519
+
520
+class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
521
+ # FlashMLA FP8 sparse only supports 64 or 128 heads
522
+ SUPPORTED_HEAD_COUNTS = (64, 128)
523
+
524
+ def __init__(
525
+ self,
526
+ num_heads: int,
527
+ head_dim: int,
528
+ scale: float,
529
+ qk_nope_head_dim: int,
530
+ qk_rope_head_dim: int,
531
+ q_lora_rank: int | None,
532
+ kv_lora_rank: int,
533
+ compress_ratio: int,
534
+ window_size: int,
535
+ head_bytes: int,
536
+ swa_cache_layer: DeepseekV4SWACache,
537
+ attn_sink: torch.Tensor,
538
+ cache_config: CacheConfig | None = None,
539
+ quant_config: QuantizationConfig | None = None,
540
+ prefix: str = "",
541
+ # Sparse MLA Args
542
+ indexer: object | None = None,
543
+ topk_indices_buffer: torch.Tensor | None = None,
544
+ aux_stream: torch.cuda.Stream | None = None,
545
+ **extra_impl_args,
546
+ ) -> None:
547
+ super().__init__()
548
+ self.num_heads = num_heads
549
+ self.num_kv_heads = 1
550
+ self.head_dim = head_dim
551
+ self.scale = scale
552
+ self.window_size = window_size
553
+ self.head_bytes = head_bytes
554
+ self.compress_ratio = compress_ratio
555
+ self.q_lora_rank = q_lora_rank
556
+ self.kv_lora_rank = kv_lora_rank
557
+ self.nope_head_dim = qk_nope_head_dim
558
+ self.rope_head_dim = qk_rope_head_dim
559
+ self.indexer = indexer
560
+ self.topk_indices_buffer = topk_indices_buffer
561
+
562
+ self.prefix = prefix # Alias for compatibility with compressor
563
+
564
+ self.aux_stream = aux_stream
565
+ self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
566
+
567
+ # Determine padded head count for FlashMLA
568
+ if num_heads not in self.SUPPORTED_HEAD_COUNTS:
569
+ if num_heads < 64:
570
+ self.padded_heads = 64
571
+ elif num_heads < 128:
572
+ self.padded_heads = 128
573
+ else:
574
+ raise ValueError(
575
+ f"DeepseekV4MLAAttention does not support {num_heads} heads. "
576
+ f"Supported: <= 128 (will be padded to 64 or 128)"
577
+ )
578
+ else:
579
+ self.padded_heads = num_heads
580
+
581
+ # Store attention sink
582
+ assert attn_sink is not None
583
+ self.attn_sink: torch.Tensor = attn_sink
584
+ # Store SWA cache
585
+ assert swa_cache_layer is not None
586
+ self.swa_cache_layer: DeepseekV4SWACache = swa_cache_layer
587
+
588
+ # Get vllm config for cache setup
589
+ vllm_config = get_current_vllm_config()
590
+ self.max_num_batched_tokens = (
591
+ vllm_config.scheduler_config.max_num_batched_tokens
592
+ )
593
+ self.max_model_len = vllm_config.model_config.max_model_len
594
+ # DeepseekV4 only supports fp8 kv-cache format for now
595
+ kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8"
596
+
597
+ assert kv_cache_dtype.startswith("fp8"), (
598
+ f"DeepseekV4 only supports fp8 kv-cache format for now, "
599
+ f"got {kv_cache_dtype}"
600
+ )
601
+ assert issubclass(self.get_attn_backend(), FlashMLASparseBackend), (
602
+ "Only FlashMLA Sparse Attention backend is supported for DeepseekV4 for now"
603
+ )
604
+ # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format
605
+ # Automatically convert fp8 kv-cache format to "fp8_ds_mla"
606
+ if (
607
+ issubclass(self.get_attn_backend(), FlashMLASparseBackend)
608
+ and kv_cache_dtype.startswith("fp8")
609
+ and kv_cache_dtype != "fp8_ds_mla"
610
+ ):
611
+ assert cache_config is not None
612
+ cache_config.cache_dtype = "fp8_ds_mla"
613
+ kv_cache_dtype = "fp8_ds_mla"
614
+ logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
615
+
616
+ self.kv_cache_dtype = kv_cache_dtype
617
+
618
+ # Register with compilation context for metadata lookup
619
+ compilation_config = vllm_config.compilation_config
620
+ if prefix and prefix in compilation_config.static_forward_context:
621
+ raise ValueError(f"Duplicate layer name: {prefix}")
622
+ if prefix:
623
+ compilation_config.static_forward_context[prefix] = self
624
+
625
+ self.kv_cache = torch.tensor([])
626
+
627
+ def get_attn_backend(self) -> type[AttentionBackend]:
628
+ return DeepseekV4FlashMLASparseBackend
629
+
630
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
631
+ if (
632
+ self.compress_ratio <= 1
633
+ ): # SWA part. Allocated separately as DeepseekV4SWACache.
634
+ return None
635
+ return MLAAttentionSpec(
636
+ block_size=vllm_config.cache_config.block_size,
637
+ num_kv_heads=1,
638
+ head_size=self.head_dim,
639
+ dtype=torch.uint8,
640
+ compress_ratio=self.compress_ratio,
641
+ cache_dtype_str=self.kv_cache_dtype,
642
+ alignment=576, # NOTE: FlashMLA requires 576B alignment
643
+ model_version="deepseek_v4",
644
+ )
645
+
646
+ def forward(
647
+ self,
648
+ q: torch.Tensor,
649
+ kv: torch.Tensor,
650
+ positions: torch.Tensor,
651
+ output: torch.Tensor,
652
+ ) -> None:
653
+ assert output.shape == q.shape, (
654
+ f"output buffer shape {output.shape} must match q shape {q.shape}"
655
+ )
656
+ assert output.dtype == q.dtype, (
657
+ f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
658
+ )
659
+
660
+ # Get SWA and indexer metadata from forward context
661
+ forward_context = get_forward_context()
662
+ attn_metadata = forward_context.attn_metadata
663
+ assert isinstance(attn_metadata, dict)
664
+ flashmla_metadata = cast(
665
+ FlashMLASparseMetadata | None, attn_metadata.get(self.prefix)
666
+ )
667
+ swa_metadata = cast(
668
+ "DeepseekSparseSWAMetadata | None",
669
+ attn_metadata.get(self.swa_cache_layer.prefix),
670
+ )
671
+ assert swa_metadata is not None
672
+
673
+ swa_only = self.compress_ratio <= 1
674
+ # SWA-only layers (compress_ratio <= 1) don't have their own KV cache
675
+ # allocation, so self.kv_cache may be empty after profiling cleanup.
676
+ self_kv_cache = self.kv_cache if not swa_only else None
677
+ swa_kv_cache = self.swa_cache_layer.kv_cache
678
+
679
+ # Split prefill and decode
680
+ num_decodes = swa_metadata.num_decodes
681
+ num_prefills = swa_metadata.num_prefills
682
+ num_decode_tokens = swa_metadata.num_decode_tokens
683
+
684
+ if num_prefills > 0:
685
+ self._forward_prefill(
686
+ q=q[num_decode_tokens:],
687
+ positions=positions[num_decode_tokens:],
688
+ compressed_k_cache=self_kv_cache,
689
+ swa_k_cache=swa_kv_cache,
690
+ output=output[num_decode_tokens:],
691
+ attn_metadata=flashmla_metadata,
692
+ swa_metadata=swa_metadata,
693
+ )
694
+ if num_decodes > 0:
695
+ self._forward_decode(
696
+ q=q[:num_decode_tokens],
697
+ kv_cache=self_kv_cache,
698
+ swa_metadata=swa_metadata,
699
+ attn_metadata=flashmla_metadata,
700
+ swa_only=swa_only,
701
+ output=output[:num_decode_tokens],
702
+ )
703
+
704
+ def _forward_decode(
705
+ self,
706
+ q: torch.Tensor,
707
+ kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1
708
+ swa_metadata: "DeepseekSparseSWAMetadata",
709
+ attn_metadata: FlashMLASparseMetadata | None,
710
+ swa_only: bool,
711
+ output: torch.Tensor,
712
+ ) -> None:
713
+ num_decodes = swa_metadata.num_decodes
714
+ num_decode_tokens = swa_metadata.num_decode_tokens
715
+
716
+ topk_indices = None
717
+ topk_lens = None
718
+ if not swa_only:
719
+ assert attn_metadata is not None
720
+ assert swa_metadata.is_valid_token is not None
721
+ block_size = attn_metadata.block_size // self.compress_ratio
722
+ is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
723
+ if self.compress_ratio == 4:
724
+ # C4A: local indices differ per layer (filled by Indexer).
725
+ assert self.topk_indices_buffer is not None
726
+ global_indices, topk_lens = compute_global_topk_indices_and_lens(
727
+ self.topk_indices_buffer[:num_decode_tokens],
728
+ swa_metadata.token_to_req_indices,
729
+ attn_metadata.block_table[:num_decodes],
730
+ block_size,
731
+ is_valid,
732
+ )
733
+ topk_indices = global_indices.view(num_decode_tokens, 1, -1)
734
+ else:
735
+ # C128A: pre-computed during metadata build.
736
+ topk_indices = attn_metadata.c128a_global_decode_topk_indices
737
+ topk_lens = attn_metadata.c128a_decode_topk_lens
738
+
739
+ swa_indices = swa_metadata.decode_swa_indices
740
+ swa_lens = swa_metadata.decode_swa_lens
741
+
742
+ # We treat queries in the same seq as different queries
743
+ # and later we only attend by generated indices.
744
+ # q arrives pre-padded to self.padded_heads by the outer wrapper.
745
+ q = q.unsqueeze(1)
746
+
747
+ # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes)
748
+ # Use unsqueeze to preserve strides (handles padded blocks correctly)
749
+ swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2)
750
+ # Reshape KV cache to (num_blocks, block_size, 1, head_bytes)
751
+ if kv_cache is not None:
752
+ kv_cache = kv_cache.unsqueeze(-2)
753
+
754
+ # One FlashMLASchedMeta per layer type, shared across all same-type
755
+ # layers within this decode step. The first forward call per type
756
+ # triggers the in-kernel planner (allocating tile_scheduler_metadata
757
+ # and num_splits via PyTorch's graph-aware allocator so CUDA graph
758
+ # capture reuses the same addresses on replay); subsequent same-type
759
+ # layers see have_initialized=True and skip the planner.
760
+ if self.compress_ratio <= 1:
761
+ tile_metadata = swa_metadata.tile_sched_swaonly
762
+ elif self.compress_ratio == 4:
763
+ tile_metadata = swa_metadata.tile_sched_c4a
764
+ elif self.compress_ratio == 128:
765
+ tile_metadata = swa_metadata.tile_sched_c128a
766
+ else:
767
+ raise ValueError(
768
+ f"Unsupported compress_ratio={self.compress_ratio}; "
769
+ "expected 1, 4, or 128."
770
+ )
771
+ assert tile_metadata is not None, (
772
+ "swa_metadata missing tile_sched entry for "
773
+ f"compress_ratio={self.compress_ratio}; "
774
+ "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not "
775
+ "allocate one for this layer type."
776
+ )
777
+
778
+ out, _ = flash_mla_with_kvcache(
779
+ q=q,
780
+ k_cache=swa_cache,
781
+ block_table=None,
782
+ head_dim_v=512,
783
+ tile_scheduler_metadata=tile_metadata,
784
+ cache_seqlens=None,
785
+ is_fp8_kvcache=True,
786
+ indices=swa_indices,
787
+ topk_length=swa_lens,
788
+ softmax_scale=self.scale,
789
+ attn_sink=self.attn_sink,
790
+ extra_k_cache=kv_cache if not swa_only else None,
791
+ extra_indices_in_kvcache=topk_indices,
792
+ extra_topk_length=topk_lens,
793
+ out=output.unsqueeze(1),
794
+ )
795
+
796
+ def _forward_prefill(
797
+ self,
798
+ q: torch.Tensor,
799
+ positions: torch.Tensor,
800
+ compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1
801
+ swa_k_cache: torch.Tensor,
802
+ output: torch.Tensor,
803
+ attn_metadata: FlashMLASparseMetadata | None,
804
+ swa_metadata: "DeepseekSparseSWAMetadata",
805
+ ) -> None:
806
+ swa_only = attn_metadata is None
807
+
808
+ num_prefills = swa_metadata.num_prefills
809
+ num_prefill_tokens = swa_metadata.num_prefill_tokens
810
+ num_decodes = swa_metadata.num_decodes
811
+ num_decode_tokens = swa_metadata.num_decode_tokens
812
+
813
+ # Use pre-computed prefill metadata.
814
+ seq_lens = swa_metadata.prefill_seq_lens
815
+ gather_lens = swa_metadata.prefill_gather_lens
816
+ assert seq_lens is not None
817
+ assert gather_lens is not None
818
+
819
+ # Derive prefill-local token offsets from the full query_start_loc_cpu.
820
+ query_start_loc_cpu = swa_metadata.query_start_loc_cpu
821
+ query_start_loc = swa_metadata.query_start_loc
822
+ assert query_start_loc_cpu is not None
823
+ assert query_start_loc is not None
824
+ prefill_token_base = query_start_loc_cpu[num_decodes]
825
+
826
+ if not swa_only:
827
+ if self.compress_ratio == 4:
828
+ assert self.topk_indices_buffer is not None
829
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:]
830
+ topk_indices = topk_indices[:num_prefill_tokens]
831
+ else:
832
+ # C128A: pre-computed during metadata build.
833
+ assert attn_metadata is not None
834
+ topk_indices = attn_metadata.c128a_prefill_topk_indices
835
+ top_k = topk_indices.shape[-1]
836
+ # Compressed region must fit the full compressed pool (seq_len //
837
+ # compress_ratio), not just top_k. top_k bounds how many indices
838
+ # the indexer selects, not the pool size it indexes into.
839
+ N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
840
+ else:
841
+ # NOTE(woosuk): topk_indices will not be used for SWA-only layers.
842
+ assert self.topk_indices_buffer is not None
843
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:]
844
+ top_k = 0
845
+ N = 0
846
+
847
+ M = N + self.window_size + self.max_num_batched_tokens
848
+ num_chunks = (num_prefills + PREFILL_CHUNK_SIZE - 1) // PREFILL_CHUNK_SIZE
849
+
850
+ workspace_manager = current_workspace_manager()
851
+ kv = workspace_manager.get_simultaneous(
852
+ ((PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
853
+ )[0]
854
+ for chunk_idx in range(num_chunks):
855
+ chunk_start = chunk_idx * PREFILL_CHUNK_SIZE
856
+ chunk_end = min(chunk_start + PREFILL_CHUNK_SIZE, num_prefills)
857
+ chunk_size = chunk_end - chunk_start
858
+ if not swa_only:
859
+ # Gather compressed KV
860
+ assert attn_metadata is not None
861
+ block_table = attn_metadata.block_table[num_decodes:]
862
+ dequantize_and_gather_k_cache(
863
+ kv[:chunk_size],
864
+ compressed_k_cache,
865
+ seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
866
+ gather_lens=None,
867
+ block_table=block_table[chunk_start:chunk_end],
868
+ block_size=attn_metadata.block_size // self.compress_ratio,
869
+ offset=0,
870
+ )
871
+
872
+ # Gather SWA KV
873
+ swa_block_table = swa_metadata.block_table[num_decodes:]
874
+ dequantize_and_gather_k_cache(
875
+ kv[:chunk_size],
876
+ swa_k_cache,
877
+ seq_lens=seq_lens[chunk_start:chunk_end],
878
+ gather_lens=gather_lens[chunk_start:chunk_end],
879
+ block_table=swa_block_table[chunk_start:chunk_end],
880
+ block_size=swa_metadata.block_size,
881
+ offset=N,
882
+ )
883
+
884
+ # Combine the topk indices and SWA indices for gathered KV cache
885
+ query_start = (
886
+ query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
887
+ )
888
+ query_end = (
889
+ query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
890
+ )
891
+
892
+ combined_indices, combined_lens = combine_topk_swa_indices(
893
+ topk_indices[query_start:query_end],
894
+ query_start_loc[
895
+ num_decodes + chunk_start : num_decodes + chunk_end + 1
896
+ ],
897
+ seq_lens[chunk_start:chunk_end],
898
+ gather_lens[chunk_start:chunk_end],
899
+ self.window_size,
900
+ self.compress_ratio,
901
+ top_k,
902
+ M,
903
+ N,
904
+ )
905
+
906
+ output_chunk, _, _ = flash_mla_sparse_fwd(
907
+ q=q[query_start:query_end],
908
+ kv=kv.view(-1, 1, q.shape[-1]),
909
+ indices=combined_indices.unsqueeze(1),
910
+ sm_scale=self.scale,
911
+ attn_sink=self.attn_sink,
912
+ topk_length=combined_lens,
913
+ out=output[query_start:query_end],
914
+ )
915
+
916
+
917
+class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
918
+ def __init__(
919
+ self,
920
+ head_dim: int,
921
+ dtype: torch.dtype,
922
+ prefix: str,
923
+ cache_config: CacheConfig,
924
+ compress_ratio: int = 1,
925
+ ):
926
+ super().__init__()
927
+ self.kv_cache = torch.tensor([])
928
+ self.head_dim = head_dim
929
+ self.prefix = prefix
930
+ self.cache_config = cache_config
931
+ self.dtype = dtype
932
+ self.compress_ratio = compress_ratio
933
+ compilation_config = get_current_vllm_config().compilation_config
934
+ if prefix in compilation_config.static_forward_context:
935
+ raise ValueError(f"Duplicate layer name: {prefix}")
936
+ compilation_config.static_forward_context[prefix] = self
937
+
938
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
939
+ # head_dim already carries the fp8 scale padding
940
+ # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout.
941
+ return MLAAttentionSpec(
942
+ block_size=self.cache_config.block_size,
943
+ num_kv_heads=1,
944
+ head_size=self.head_dim,
945
+ dtype=self.dtype,
946
+ compress_ratio=self.compress_ratio,
947
+ # DeepseekV4 aligns indexer pages to FlashMLA's 576B so they can pack with
948
+ # the indexer's compressor state cache. V3.2 keeps the legacy layout.
949
+ alignment=576,
950
+ )
951
+
952
+ def forward(self): ...
953
+
954
+ def get_attn_backend(self) -> type[AttentionBackend]:
955
+ return DeepseekV4IndexerBackend
956
+
957
+
958
+class DeepseekV4Indexer(nn.Module):
959
+ def __init__(
960
+ self,
961
+ vllm_config: VllmConfig,
962
+ config: DeepseekV2Config | DeepseekV3Config,
963
+ hidden_size: int,
964
+ q_lora_rank: int,
965
+ quant_config: QuantizationConfig | None,
966
+ cache_config: CacheConfig | None,
967
+ topk_indices_buffer: torch.Tensor | None,
968
+ compress_ratio: int = 1,
969
+ prefix: str = "",
970
+ ):
971
+ super().__init__()
972
+ self.vllm_config = vllm_config
973
+ self.config = config
974
+ self.quant_config = quant_config
975
+ # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
976
+ self.topk_tokens = config.index_topk
977
+ self.n_head = config.index_n_heads # 64
978
+ self.head_dim = config.index_head_dim # 128
979
+ self.rope_dim = config.qk_rope_head_dim # 64
980
+ self.q_lora_rank = q_lora_rank # 1536
981
+ self.compress_ratio = compress_ratio
982
+ self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache
983
+ logger.info_once(
984
+ "Using %s indexer cache for Lighening Indexer.",
985
+ "MXFP4" if self.use_fp4_kv else "FP8",
986
+ )
987
+
988
+ # no tensor parallel, just replicated
989
+ self.wq_b = ReplicatedLinear(
990
+ self.q_lora_rank,
991
+ self.head_dim * self.n_head,
992
+ bias=False,
993
+ quant_config=quant_config,
994
+ prefix=f"{prefix}.wq_b",
995
+ )
996
+ self.weights_proj = ReplicatedLinear(
997
+ hidden_size,
998
+ self.n_head,
999
+ bias=False,
1000
+ quant_config=None,
1001
+ prefix=f"{prefix}.weights_proj",
1002
+ )
1003
+ self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
1004
+ self.softmax_scale = self.head_dim**-0.5
1005
+
1006
+ self.scale_fmt = "ue8m0"
1007
+ self.quant_block_size = 128 # TODO: get from config
1008
+ self.topk_indices_buffer = topk_indices_buffer
1009
+
1010
+ self.max_model_len = (
1011
+ vllm_config.model_config.max_model_len // self.compress_ratio
1012
+ )
1013
+ self.prefix = prefix
1014
+
1015
+ self.max_total_seq_len = (
1016
+ get_max_prefill_buffer_size(vllm_config) // self.compress_ratio
1017
+ )
1018
+
1019
+ assert cache_config is not None, "Deepseek V4 indexer requires cache_config"
1020
+ # NOTE(yifan): FP8 indxer cache use the same layout as V3.2:
1021
+ # head_dim bytes = 128 fp8 + 4 fp32 scale = 132.
1022
+ # For FP4 indexer cache, we still allocate the same amount of memory as FP8,
1023
+ # but only use the first half of the memory.
1024
+ k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4
1025
+ self.k_cache = DeepseekV4IndexerCache(
1026
+ head_dim=k_cache_head_dim,
1027
+ dtype=torch.uint8,
1028
+ prefix=f"{prefix}.k_cache",
1029
+ cache_config=cache_config,
1030
+ compress_ratio=self.compress_ratio,
1031
+ )
1032
+ self.compressor = DeepseekCompressor(
1033
+ vllm_config=vllm_config,
1034
+ compress_ratio=self.compress_ratio,
1035
+ hidden_size=hidden_size,
1036
+ head_dim=self.head_dim,
1037
+ rotate=True,
1038
+ prefix=f"{prefix}.compressor",
1039
+ k_cache_prefix=self.k_cache.prefix,
1040
+ use_fp4_cache=self.use_fp4_kv,
1041
+ )
1042
+
1043
+ self.indexer_op = SparseAttnIndexer(
1044
+ self.k_cache,
1045
+ self.quant_block_size,
1046
+ self.scale_fmt,
1047
+ self.topk_tokens,
1048
+ self.head_dim,
1049
+ self.max_model_len,
1050
+ self.max_total_seq_len,
1051
+ self.topk_indices_buffer,
1052
+ skip_k_cache_insert=True,
1053
+ use_fp4_cache=self.use_fp4_kv,
1054
+ )
1055
+
1056
+ def forward(
1057
+ self,
1058
+ hidden_states: torch.Tensor,
1059
+ qr: torch.Tensor,
1060
+ positions: torch.Tensor,
1061
+ rotary_emb: nn.Module,
1062
+ ) -> torch.Tensor:
1063
+ q, _ = self.wq_b(qr)
1064
+ q = q.view(-1, self.n_head, self.head_dim)
1065
+ k = self.compressor(hidden_states, positions, rotary_emb)
1066
+ weights, _ = self.weights_proj(hidden_states)
1067
+ q_quant, weights = fused_indexer_q_rope_quant(
1068
+ positions,
1069
+ q,
1070
+ rotary_emb.cos_sin_cache,
1071
+ weights,
1072
+ self.softmax_scale,
1073
+ self.n_head**-0.5,
1074
+ use_fp4=self.use_fp4_kv,
1075
+ )
1076
+ return self.indexer_op(hidden_states, q_quant, k, weights)
vllm/model_executor/layers/fused_moe/config.py
+45
-1
@@ -117,8 +117,10 @@ class RoutingMethodType(IntEnum):
117
Custom = (6,)
118
# Simulated
119
Simulated = (7,)
120
+ # Deepseek V4 -> sqrtsoftplus + Bias + Normalize
121
+ DeepseekV4 = (8,)
122
# Unspecified
121
- Unspecified = 8.0
123
+ Unspecified = 9.0
124
125
126
def get_routing_method_type(
@@ -128,6 +130,14 @@ def get_routing_method_type(
130
num_expert_group: int | None,
131
has_e_score_bias: bool,
132
) -> RoutingMethodType:
133
+ if scoring_func == "sqrtsoftplus":
134
+ # DeepSeek V4 uses sqrtsoftplus routing with optional routing bias
135
+ # and top-k renormalization.
136
+ if renormalize:
137
+ return RoutingMethodType.DeepseekV4
138
+ else:
139
+ return RoutingMethodType.Unspecified
140
+
141
if has_e_score_bias:
142
if (num_expert_group or 0) > 0 and scoring_func == "sigmoid":
143
return RoutingMethodType.DeepSeekV3
@@ -230,6 +240,13 @@ class FusedMoEQuantConfig:
240
_w2: FusedMoEQuantDesc
241
is_nvfp4_scale_swizzled: bool = True
242
243
+ # MXFP4-specific TRTLLM parameters for SwiGLU activation clamping.
244
+ # These correspond to gemm1_alpha, gemm1_beta, gemm1_clamp_limit
245
+ # in TrtLlmMxfp4ExpertsBase.
246
+ gemm1_alpha: float | None = None
247
+ gemm1_beta: float | None = None
248
+ gemm1_clamp_limit: float | None = None
249
+
250
def __post_init__(self):
251
assert not self.per_act_token_quant or self.block_shape is None, (
252
"illegal quantization"
@@ -477,6 +494,9 @@ class FusedMoEQuantConfig:
494
w2_zp: torch.Tensor | None = None,
495
weight_dtype: torch.dtype | str | None = None,
496
is_nvfp4_scale_swizzled: bool = True,
497
+ gemm1_alpha: float | None = None,
498
+ gemm1_beta: float | None = None,
499
+ gemm1_clamp_limit: float | None = None,
500
) -> "FusedMoEQuantConfig":
501
"""
502
General builder function for a FusedMoEQuantConfig.
@@ -507,6 +527,9 @@ class FusedMoEQuantConfig:
527
- w1_zp: Optional w1 zero points for int4/int8 quantization.
528
- w2_zp: Optional w2 zero points for int4/int8 quantization.
529
- is_nvfp4_scale_swizzled: Whether to swizzle the nvfp4 scale swizzling.
530
+ - gemm1_alpha: Optional MXFP4 TRTLLM SwiGLU alpha parameter.
531
+ - gemm1_beta: Optional MXFP4 TRTLLM SwiGLU beta parameter.
532
+ - gemm1_clamp_limit: Optional MXFP4 TRTLLM SwiGLU clamp limit.
533
"""
534
assert not isinstance(quant_dtype, str) or quant_dtype in {
535
"nvfp4",
@@ -540,6 +563,9 @@ class FusedMoEQuantConfig:
563
weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias
564
),
565
is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled,
566
+ gemm1_alpha=gemm1_alpha,
567
+ gemm1_beta=gemm1_beta,
568
+ gemm1_clamp_limit=gemm1_clamp_limit,
569
)
570
assert quant_config.per_act_token_quant == per_act_token_quant
571
assert quant_config.per_out_ch_quant == per_out_ch_quant
@@ -650,6 +676,9 @@ def mxfp4_w4a16_moe_quant_config(
676
w2_scale: Union[torch.Tensor, "PrecisionConfig"],
677
w1_bias: torch.Tensor | None = None,
678
w2_bias: torch.Tensor | None = None,
679
+ gemm1_alpha: float | None = None,
680
+ gemm1_beta: float | None = None,
681
+ gemm1_clamp_limit: float | None = None,
682
) -> FusedMoEQuantConfig:
683
"""
684
Construct a quant config for unquantized activations and mxfp4 weights.
@@ -659,6 +688,9 @@ def mxfp4_w4a16_moe_quant_config(
688
_a2=FusedMoEQuantDesc(),
689
_w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
690
_w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
691
+ gemm1_alpha=gemm1_alpha,
692
+ gemm1_beta=gemm1_beta,
693
+ gemm1_clamp_limit=gemm1_clamp_limit,
694
)
695
696
@@ -670,6 +702,9 @@ def mxfp4_mxfp8_moe_quant_config(
702
w1_bias: torch.Tensor | None = None,
703
w2_bias: torch.Tensor | None = None,
704
block_shape: list[int] | None = None,
705
+ gemm1_alpha: float | None = None,
706
+ gemm1_beta: float | None = None,
707
+ gemm1_clamp_limit: float | None = None,
708
) -> FusedMoEQuantConfig:
709
"""
710
Construct a quant config for mxfp4 activations and mxfp4 weights.
@@ -679,6 +714,9 @@ def mxfp4_mxfp8_moe_quant_config(
714
_a2=FusedMoEQuantDesc("mxfp8"),
715
_w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
716
_w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
717
+ gemm1_alpha=gemm1_alpha,
718
+ gemm1_beta=gemm1_beta,
719
+ gemm1_clamp_limit=gemm1_clamp_limit,
720
)
721
722
@@ -712,6 +750,9 @@ def ocp_mx_moe_quant_config(
750
w1_bias: torch.Tensor | None = None,
751
w2_bias: torch.Tensor | None = None,
752
block_shape: list[int] | None = None,
753
+ gemm1_alpha: float | None = None,
754
+ gemm1_beta: float | None = None,
755
+ gemm1_clamp_limit: float | None = None,
756
) -> FusedMoEQuantConfig:
757
"""
758
Construct a quant config for mxfp4 activations and mxfp4 weights.
@@ -729,6 +770,9 @@ def ocp_mx_moe_quant_config(
770
per_act_token_quant=False,
771
per_out_ch_quant=False,
772
block_shape=block_shape,
773
+ gemm1_alpha=gemm1_alpha,
774
+ gemm1_beta=gemm1_beta,
775
+ gemm1_clamp_limit=gemm1_clamp_limit,
776
)
777
778
vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
+234
-1
@@ -25,15 +25,20 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
25
per_token_group_quant_fp8_packed_for_deepgemm,
26
silu_mul_per_token_group_quant_fp8_colmajor,
27
)
28
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
29
+ silu_mul_quant_fp8_packed_triton as fused_silu_mul_fp8_quant_packed,
30
+)
31
from vllm.model_executor.layers.quantization.utils.quant_utils import (
32
QuantKey,
33
kFp8Dynamic128Sym,
34
kFp8Static128BlockSym,
35
+ kMxfp4Static,
36
)
37
from vllm.utils.deep_gemm import (
38
DeepGemmQuantScaleFMT,
39
get_mk_alignment_for_contiguous_layout,
40
is_deep_gemm_supported,
41
+ m_grouped_fp8_fp4_gemm_nt_contiguous,
42
m_grouped_fp8_gemm_nt_contiguous,
43
)
44
from vllm.utils.import_utils import has_deep_gemm
@@ -197,8 +202,14 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
202
M_sum, N = input.size()
203
activation_out_dim = self.adjust_N_for_activation(N, activation)
204
200
- # 1. DeepGemm UE8M0: use packed per-token-group quant
205
+ # 1. DeepGemm UE8M0: fused SiLU+mul+clamp+quant+pack
206
if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
207
+ if activation == MoEActivation.SILU:
208
+ return fused_silu_mul_fp8_quant_packed(
209
+ input=input,
210
+ output_q=output,
211
+ group_size=block_k,
212
+ )
213
act_out = torch.empty(
214
(M_sum, activation_out_dim), dtype=input.dtype, device=input.device
215
)
@@ -312,3 +323,225 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
323
expert_map=expert_map,
324
output=output,
325
)
326
+
327
+
328
+class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
329
+ """DeepGemm-based fused MoE expert implementation for FP4 weights.
330
+
331
+ Uses m_grouped_fp8_fp4_gemm_nt_contiguous with FP8 activations and
332
+ MXFP4 (FP4 E2M1 packed as uint8) weights. Requires SM100+ (Blackwell).
333
+ """
334
+
335
+ # FP8 activation block size (hardcoded since mxfp4_w4a8 quant config
336
+ # does not set a block_shape on the activation descriptor).
337
+ _ACT_BLOCK_K = 128
338
+ # FP4 weight block size
339
+ _WEIGHT_BLOCK_K = 32
340
+
341
+ def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig):
342
+ super().__init__(moe_config=moe_config, quant_config=quant_config)
343
+ assert quant_config.weight_quant_dtype == "mxfp4"
344
+ assert not quant_config.per_act_token_quant
345
+ assert not quant_config.per_out_ch_quant
346
+
347
+ self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
348
+
349
+ @staticmethod
350
+ def activation_format() -> mk.FusedMoEActivationFormat:
351
+ return mk.FusedMoEActivationFormat.Standard
352
+
353
+ @staticmethod
354
+ def _supports_current_device() -> bool:
355
+ from vllm.platforms import current_platform
356
+
357
+ return (
358
+ is_deep_gemm_supported()
359
+ and current_platform.is_device_capability_family(100)
360
+ )
361
+
362
+ @staticmethod
363
+ def _supports_no_act_and_mul() -> bool:
364
+ return False
365
+
366
+ @staticmethod
367
+ def _supports_quant_scheme(
368
+ weight_key: QuantKey | None,
369
+ activation_key: QuantKey | None,
370
+ ) -> bool:
371
+ SUPPORTED_W_A = [
372
+ (kMxfp4Static, kFp8Dynamic128Sym),
373
+ ]
374
+ return (weight_key, activation_key) in SUPPORTED_W_A
375
+
376
+ @staticmethod
377
+ def _supports_activation(activation: MoEActivation) -> bool:
378
+ return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP]
379
+
380
+ @staticmethod
381
+ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
382
+ return not (
383
+ moe_parallel_config.use_fi_nvl_two_sided_kernels
384
+ or moe_parallel_config.use_fi_nvl_one_sided_kernels
385
+ )
386
+
387
+ def supports_expert_map(self) -> bool:
388
+ return True
389
+
390
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
391
+ return TopKWeightAndReduceNoOP()
392
+
393
+ def workspace_shapes(
394
+ self,
395
+ M: int,
396
+ N: int,
397
+ K: int,
398
+ topk: int,
399
+ global_num_experts: int,
400
+ local_num_experts: int,
401
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
402
+ activation: MoEActivation,
403
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
404
+ block_m = get_mk_alignment_for_contiguous_layout()[0]
405
+ M_sum = compute_aligned_M(
406
+ M, topk, local_num_experts, block_m, expert_tokens_meta
407
+ )
408
+ assert M_sum % block_m == 0
409
+
410
+ activation_out_dim = self.adjust_N_for_activation(N, activation)
411
+ workspace1 = (M_sum, max(activation_out_dim, K))
412
+ workspace2 = (M_sum, max(N, K))
413
+ output = (M, K)
414
+ return (workspace1, workspace2, output)
415
+
416
+ def _act_mul_quant(
417
+ self, input: torch.Tensor, output: torch.Tensor, activation: MoEActivation
418
+ ) -> tuple[torch.Tensor, torch.Tensor]:
419
+ block_k = self._ACT_BLOCK_K
420
+ scale_fmt = DeepGemmQuantScaleFMT.from_oracle()
421
+
422
+ M_sum, N = input.size()
423
+ activation_out_dim = self.adjust_N_for_activation(N, activation)
424
+
425
+ if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
426
+ assert activation == MoEActivation.SILU
427
+ return fused_silu_mul_fp8_quant_packed(
428
+ input=input,
429
+ output_q=output,
430
+ group_size=block_k,
431
+ clamp_limit=self.gemm1_clamp_limit,
432
+ )
433
+
434
+ if activation == MoEActivation.SILU:
435
+ use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
436
+ return silu_mul_per_token_group_quant_fp8_colmajor(
437
+ input=input,
438
+ output=output,
439
+ use_ue8m0=use_ue8m0,
440
+ )
441
+
442
+ act_out = torch.empty(
443
+ (M_sum, activation_out_dim), dtype=input.dtype, device=input.device
444
+ )
445
+ self.activation(activation, act_out, input)
446
+ return per_token_group_quant_fp8(
447
+ act_out, block_k, column_major_scales=True, out_q=output
448
+ )
449
+
450
+ def apply(
451
+ self,
452
+ output: torch.Tensor,
453
+ hidden_states: torch.Tensor,
454
+ w1: torch.Tensor,
455
+ w2: torch.Tensor,
456
+ topk_weights: torch.Tensor,
457
+ topk_ids: torch.Tensor,
458
+ activation: MoEActivation,
459
+ global_num_experts: int,
460
+ expert_map: torch.Tensor | None,
461
+ a1q_scale: torch.Tensor | None,
462
+ a2_scale: torch.Tensor | None,
463
+ workspace13: torch.Tensor,
464
+ workspace2: torch.Tensor,
465
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
466
+ apply_router_weight_on_input: bool,
467
+ ):
468
+ assert a1q_scale is not None
469
+ assert a2_scale is None
470
+ assert self.w1_scale is not None
471
+ assert self.w2_scale is not None
472
+
473
+ a1q = hidden_states
474
+ _, N, _ = w1.size()
475
+ # K comes from activations (full hidden dim), not from w1 which is
476
+ # packed FP4 (E, N, K//2).
477
+ K = a1q.size(1)
478
+
479
+ local_num_experts = w1.size(0)
480
+ if global_num_experts == -1:
481
+ global_num_experts = local_num_experts
482
+
483
+ M_sum = compute_aligned_M(
484
+ M=topk_ids.size(0),
485
+ num_topk=topk_ids.size(1),
486
+ local_num_experts=local_num_experts,
487
+ alignment=get_mk_alignment_for_contiguous_layout()[0],
488
+ expert_tokens_meta=expert_tokens_meta,
489
+ )
490
+
491
+ a1q_perm = _resize_cache(
492
+ workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K)
493
+ )
494
+ a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
495
+ aq=a1q,
496
+ aq_scale=a1q_scale,
497
+ topk_ids=topk_ids,
498
+ local_num_experts=local_num_experts,
499
+ expert_map=expert_map,
500
+ expert_tokens_meta=expert_tokens_meta,
501
+ aq_out=a1q_perm,
502
+ )
503
+ assert a1q.size(0) == M_sum
504
+
505
+ # FC1: FP8 activations x FP4 weights
506
+ # DeepGEMM 2.4.2 requires FP4-packed weights as int8 (kPackedFP4).
507
+ mm1_out = _resize_cache(workspace2, (M_sum, N))
508
+ m_grouped_fp8_fp4_gemm_nt_contiguous(
509
+ (a1q, a1q_scale),
510
+ (w1.view(torch.int8), self.w1_scale),
511
+ mm1_out,
512
+ expert_ids,
513
+ recipe_a=(1, self._ACT_BLOCK_K),
514
+ recipe_b=(1, self._WEIGHT_BLOCK_K),
515
+ )
516
+
517
+ # SwiGLU activation + FP8 requant
518
+ activation_out_dim = self.adjust_N_for_activation(N, activation)
519
+ quant_out = _resize_cache(
520
+ workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
521
+ )
522
+ a2q, a2q_scale = self._act_mul_quant(
523
+ input=mm1_out.view(-1, N), output=quant_out, activation=activation
524
+ )
525
+
526
+ # FC2: FP8 activations x FP4 weights
527
+ mm2_out = _resize_cache(workspace2, (M_sum, K))
528
+ m_grouped_fp8_fp4_gemm_nt_contiguous(
529
+ (a2q, a2q_scale),
530
+ (w2.view(torch.int8), self.w2_scale),
531
+ mm2_out,
532
+ expert_ids,
533
+ recipe_a=(1, self._ACT_BLOCK_K),
534
+ recipe_b=(1, self._WEIGHT_BLOCK_K),
535
+ )
536
+
537
+ if apply_router_weight_on_input:
538
+ topk_weights = torch.ones_like(topk_weights)
539
+
540
+ deepgemm_unpermute_and_reduce(
541
+ a=mm2_out,
542
+ topk_ids=topk_ids,
543
+ topk_weights=topk_weights,
544
+ inv_perm=inv_perm,
545
+ expert_map=expert_map,
546
+ output=output,
547
+ )
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+193
-2
@@ -28,8 +28,162 @@ from vllm.platforms import current_platform
28
from vllm.triton_utils import tl, triton
29
from vllm.utils.import_utils import has_triton_kernels
30
31
+from ..utils import swiglu_limit_func
32
+
33
logger = init_logger(__name__)
34
35
+
36
+def _patch_make_bitmatrix_metadata() -> None:
37
+ """Monkey-patch make_bitmatrix_metadata to support non-power-of-2 top_k.
38
+
39
+ triton's tl.arange requires a power-of-2 range. The original kernel
40
+ computes BLOCK_SIZE = BLOCK_PER_TOK * TOKS_PER_ROW (= 32 * top_k). For
41
+ DeepSeek-V4 with top_k=6 this gives 192, which is not a power of 2 and
42
+ causes a compile error at the first forward pass.
43
+
44
+ Fix: define a drop-in replacement kernel that accepts an extra constexpr
45
+ BLOCK_SIZE_PADDED (next power of 2 >= BLOCK_SIZE) and uses it for the
46
+ tl.arange call while keeping the actual BLOCK_SIZE as the stride between
47
+ thread-blocks so that all flat indices into NonzeroIndx stay correct.
48
+ Elements beyond BLOCK_SIZE are masked out (col_indx = 0xffff) and ignored.
49
+
50
+ This function is called once at module load time and patches the function
51
+ inside the triton_kernels tensor module so that SparseMatrix.__post_init__
52
+ picks up the fixed version transparently.
53
+ """
54
+ import torch
55
+ import triton
56
+ import triton.language as tl
57
+
58
+ try:
59
+ from vllm.third_party.triton_kernels.tensor_details import (
60
+ bitmatrix as _bm,
61
+ )
62
+ from vllm.third_party.triton_kernels.tensor_details.bitmatrix import (
63
+ BitmatrixMetadata,
64
+ _keyed_add,
65
+ cdiv,
66
+ )
67
+ from vllm.third_party.triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501
68
+ sum_bitmatrix_rows,
69
+ )
70
+ except ImportError:
71
+ return
72
+
73
+ @triton.jit
74
+ def _stage2_pow2(
75
+ ColSortedIndx,
76
+ RowSortedIndx,
77
+ NonzeroIndx,
78
+ n_tokens,
79
+ ColPartialSum,
80
+ stride_pm,
81
+ stride_pn,
82
+ ColOffs,
83
+ TOKS_PER_ROW: tl.constexpr,
84
+ BLOCK_PER_TOK: tl.constexpr,
85
+ BLOCK_SIZE_PADDED: tl.constexpr,
86
+ ):
87
+ # Actual number of elements per block (may not be a power of 2).
88
+ BLOCK_SIZE: tl.constexpr = BLOCK_PER_TOK * TOKS_PER_ROW
89
+ tl.static_assert(BLOCK_SIZE_PADDED <= 32768)
90
+ if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr():
91
+ n_tokens = tl.load(n_tokens)
92
+ nonzero_indx_size = n_tokens * TOKS_PER_ROW
93
+ pid_m = tl.program_id(0)
94
+ # Use BLOCK_SIZE_PADDED (a power of 2) for tl.arange, but stride by
95
+ # the actual BLOCK_SIZE so flat positions in NonzeroIndx are correct.
96
+ # Elements with offs_local >= BLOCK_SIZE have offs_global beyond the
97
+ # valid range, get col_indx = 0xffff, and are filtered by the mask
98
+ # below without producing any output.
99
+ offs_local = tl.arange(0, BLOCK_SIZE_PADDED)
100
+ offs_global = pid_m * BLOCK_SIZE + offs_local
101
+ mask = offs_global < nonzero_indx_size
102
+ col_indx = tl.load(NonzeroIndx + offs_global, mask=mask, other=-1).to(tl.uint32)
103
+ kv_pairs = ((col_indx << 16) | offs_local).to(tl.uint32)
104
+ kv_pairs = tl.sort(kv_pairs, 0)
105
+ col_indx = kv_pairs >> 16
106
+ offs_global = pid_m * BLOCK_SIZE + (kv_pairs & 0xFFFF)
107
+ mask = col_indx != 0xFFFF
108
+ x = kv_pairs & 0xFFFF0000 | 0x00000001
109
+ cols_and_inclusive_run_lengths = tl.associative_scan(x, 0, _keyed_add)
110
+ exclusive_run_lengths = (cols_and_inclusive_run_lengths - 1) & 0xFFFF
111
+ row_sorted_indx = tl.load(
112
+ ColPartialSum + pid_m * stride_pm + col_indx * stride_pn, mask=mask
113
+ )
114
+ row_sorted_indx += tl.load(ColOffs + col_indx, mask=mask)
115
+ row_sorted_indx += exclusive_run_lengths
116
+ tl.store(RowSortedIndx + offs_global, row_sorted_indx, mask=mask)
117
+ tl.store(ColSortedIndx + row_sorted_indx, offs_global, mask=mask)
118
+
119
+ def _make_bitmatrix_metadata_pow2_safe(nonzero_indx, bitmatrix):
120
+ assert nonzero_indx.ndim == 2
121
+ PARTIAL_BLOCK_M = 32
122
+ col_sum, col_partial_sum = sum_bitmatrix_rows(
123
+ bitmatrix, partials_block_size=PARTIAL_BLOCK_M
124
+ )
125
+ device = bitmatrix.device
126
+ n_indx = nonzero_indx.numel()
127
+ n_cols = bitmatrix.shape[1]
128
+ col_offs = torch.empty(n_cols, dtype=torch.int32, device=device)
129
+ combined_indx = torch.empty(n_indx * 2, dtype=torch.int32, device=device)
130
+ col_sorted_indx = combined_indx[:n_indx]
131
+ row_sorted_indx = combined_indx[n_indx:]
132
+ MEMSET_BLOCK = 1024
133
+ memset_grid = (cdiv(n_indx * 2, MEMSET_BLOCK) + n_cols + 1,)
134
+ _bm._bitmatrix_metadata_compute_stage1[memset_grid](
135
+ combined_indx,
136
+ n_indx * 2,
137
+ -1,
138
+ MEMSET_BLOCK,
139
+ col_sum,
140
+ col_offs,
141
+ col_sum.shape[0],
142
+ col_partial_sum,
143
+ col_partial_sum.shape[0],
144
+ col_partial_sum.stride(0),
145
+ col_partial_sum.stride(1),
146
+ BLOCK_M=512,
147
+ BLOCK_N=512,
148
+ )
149
+ toks_per_row = nonzero_indx.shape[-1]
150
+ block_size = PARTIAL_BLOCK_M * toks_per_row
151
+ # Next power of 2 >= block_size (required by tl.arange).
152
+ block_size_padded = 1 << (max(block_size, 1) - 1).bit_length()
153
+ compute_grid = (cdiv(bitmatrix.shape_max[0], PARTIAL_BLOCK_M),)
154
+ _stage2_pow2[compute_grid](
155
+ col_sorted_indx,
156
+ row_sorted_indx,
157
+ nonzero_indx,
158
+ bitmatrix.shape[0],
159
+ col_partial_sum,
160
+ col_partial_sum.stride(0),
161
+ col_partial_sum.stride(1),
162
+ col_offs,
163
+ TOKS_PER_ROW=toks_per_row,
164
+ BLOCK_PER_TOK=PARTIAL_BLOCK_M,
165
+ BLOCK_SIZE_PADDED=block_size_padded,
166
+ )
167
+ return BitmatrixMetadata(
168
+ col_sum=col_sum,
169
+ col_sorted_indx=col_sorted_indx,
170
+ row_sorted_indx=row_sorted_indx,
171
+ )
172
+
173
+ # The most reliable patch point: SparseMatrix.__post_init__ looks up
174
+ # make_bitmatrix_metadata via its own __globals__ dict (the tensor.py
175
+ # module dict). Patching through __globals__ works regardless of how
176
+ # sys.modules maps "triton_kernels.tensor" vs
177
+ # "vllm.third_party.triton_kernels.tensor".
178
+ from triton_kernels.tensor import SparseMatrix as _SparseMatrix
179
+
180
+ _SparseMatrix.__post_init__.__globals__["make_bitmatrix_metadata"] = (
181
+ _make_bitmatrix_metadata_pow2_safe
182
+ )
183
+ # Also patch the bitmatrix module itself in case it is imported directly.
184
+ _bm.make_bitmatrix_metadata = _make_bitmatrix_metadata_pow2_safe
185
+
186
+
187
use_legacy_triton_kernels = False
188
189
if has_triton_kernels():
@@ -59,6 +213,8 @@ if has_triton_kernels():
213
use_legacy_triton_kernels = True
214
else:
215
raise
216
+ if not use_legacy_triton_kernels:
217
+ _patch_make_bitmatrix_metadata()
218
except (AttributeError, ImportError) as e:
219
logger.error(
220
"Failed to import Triton kernels. Please make sure your triton "
@@ -497,6 +653,8 @@ class BaseOAITritonExperts(mk.FusedMoEExpertsModular):
653
return False
654
# (9,0) <= cap < (11,0) covers CUDA SM90 (Hopper), SM100+ (Blackwell)
655
# and ROCm gfx942/gfx950 (which map to 9.4/9.5).
656
+ if not has_triton_kernels():
657
+ return False
658
return (9, 0) <= (cap.major, cap.minor) < (11, 0)
659
660
@staticmethod
@@ -698,6 +856,37 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
856
def moe_sum(self, input: torch.Tensor, output: torch.Tensor):
857
ops.moe_sum(input, output)
858
859
+ def activation(
860
+ self,
861
+ activation: MoEActivation,
862
+ output: torch.Tensor,
863
+ input: torch.Tensor,
864
+ ) -> None:
865
+ quant_config = self.quant_config or FUSED_MOE_UNQUANTIZED_CONFIG
866
+ if activation == MoEActivation.SWIGLUOAI:
867
+ alpha = (
868
+ quant_config.gemm1_alpha
869
+ if quant_config.gemm1_alpha is not None
870
+ else 1.702
871
+ )
872
+ limit = (
873
+ quant_config.gemm1_clamp_limit
874
+ if quant_config.gemm1_clamp_limit is not None
875
+ else 7.0
876
+ )
877
+ torch.ops._C.swigluoai_and_mul(output, input, alpha, limit)
878
+ elif (
879
+ activation == MoEActivation.SILU
880
+ and quant_config.gemm1_clamp_limit is not None
881
+ ):
882
+ swiglu_limit_func(
883
+ output,
884
+ input,
885
+ quant_config.gemm1_clamp_limit,
886
+ )
887
+ else:
888
+ super().activation(activation, output, input)
889
+
890
def apply(
891
self,
892
output: torch.Tensor,
@@ -812,9 +1001,9 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
1001
act_input,
1002
)
1003
815
- # matmul_ogs grouped reduction fuse sum across multiple experts:
1004
+ # matmul_ogs grouped reduction fuses sum across multiple experts:
1005
# y[dst_indx // n_expts_act, :] += x
817
- # Need to set n_expts_act to 1 to unfuse moe_sum
1006
+ # Set n_expts_act to 1 to unfuse the sum so we can do it manually via moe_sum.
1007
routing_data.n_expts_act = 1
1008
1009
matmul_ogs(
@@ -878,6 +1067,8 @@ class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
1067
return False
1068
# (9,0) <= cap < (11,0) covers CUDA SM90 (Hopper), SM100+ (Blackwell)
1069
# and ROCm gfx942/gfx950 (which map to 9.4/9.5).
1070
+ if not has_triton_kernels():
1071
+ return False
1072
return (9, 0) <= (cap.major, cap.minor) < (11, 0)
1073
1074
@staticmethod
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+84
-60
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.config import (
14
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15
TopKWeightAndReduceNoOP,
16
)
17
+from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
18
from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
QuantKey,
20
kMxfp4Static,
@@ -32,10 +33,8 @@ class TrtLlmMxfp4ExpertsBase:
33
self,
34
moe_config: FusedMoEConfig,
35
quant_config: FusedMoEQuantConfig,
36
+ **kwargs,
37
):
36
- # NOTE: FusedMoEExperts.__init__ is called by the concrete subclass
37
- # (Monolithic/Modular) via MRO, not here, to avoid mypy issues with
38
- # multiple inheritance. This matches the NvFP4 expert pattern.
38
self.moe_config = moe_config
39
self.quant_config = quant_config
40
@@ -48,23 +47,34 @@ class TrtLlmMxfp4ExpertsBase:
47
self.local_num_experts = moe_config.num_local_experts
48
self.ep_rank = moe_config.moe_parallel_config.ep_rank
49
51
- # MXFP4-specific TRTLLM parameters
50
+ # MXFP4-specific TRTLLM parameters from quant_config
51
device = torch.accelerator.current_device_index()
53
- self.gemm1_alpha = torch.tensor(
54
- [1.702] * self.local_num_experts,
55
- dtype=torch.float32,
56
- device=device,
57
- )
58
- self.gemm1_beta = torch.tensor(
59
- [1.0] * self.local_num_experts,
60
- dtype=torch.float32,
61
- device=device,
62
- )
63
- self.gemm1_clamp_limit = torch.tensor(
64
- [7.0] * self.local_num_experts,
65
- dtype=torch.float32,
66
- device=device,
67
- )
52
+ if quant_config.gemm1_alpha is not None:
53
+ self.gemm1_alpha = torch.tensor(
54
+ [quant_config.gemm1_alpha] * self.local_num_experts,
55
+ dtype=torch.float32,
56
+ device=device,
57
+ )
58
+ else:
59
+ self.gemm1_alpha = None
60
+
61
+ if quant_config.gemm1_beta is not None:
62
+ self.gemm1_beta = torch.tensor(
63
+ [quant_config.gemm1_beta] * self.local_num_experts,
64
+ dtype=torch.float32,
65
+ device=device,
66
+ )
67
+ else:
68
+ self.gemm1_beta = None
69
+
70
+ if quant_config.gemm1_clamp_limit is not None:
71
+ self.gemm1_clamp_limit = torch.tensor(
72
+ [quant_config.gemm1_clamp_limit] * self.local_num_experts,
73
+ dtype=torch.float32,
74
+ device=device,
75
+ )
76
+ else:
77
+ self.gemm1_clamp_limit = None
78
79
from vllm.config import get_current_vllm_config
80
@@ -97,7 +107,7 @@ class TrtLlmMxfp4ExpertsBase:
107
108
@staticmethod
109
def _supports_activation(activation: MoEActivation) -> bool:
100
- return activation == MoEActivation.SWIGLUOAI
110
+ return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU)
111
112
@staticmethod
113
def activation_format() -> mk.FusedMoEActivationFormat:
@@ -190,36 +200,41 @@ class TrtLlmMxfp4ExpertsMonolithic(
200
201
output = torch.empty_like(hidden_states)
202
193
- return trtllm_fp4_block_scale_moe(
194
- routing_logits=router_logits.to(torch.bfloat16),
195
- routing_bias=None,
196
- hidden_states=x_quant,
197
- hidden_states_scale=x_scale,
198
- gemm1_weights=w1,
199
- gemm1_weights_scale=self.w1_scale,
200
- gemm1_bias=self.w1_bias,
201
- gemm1_alpha=self.gemm1_alpha,
202
- gemm1_beta=self.gemm1_beta,
203
- gemm1_clamp_limit=self.gemm1_clamp_limit,
204
- gemm2_weights=w2,
205
- gemm2_weights_scale=self.w2_scale,
206
- gemm2_bias=self.w2_bias,
207
- output1_scale_scalar=None,
208
- output1_scale_gate_scalar=None,
209
- output2_scale_scalar=None,
210
- num_experts=global_num_experts,
211
- top_k=self.topk,
212
- n_group=None,
213
- topk_group=None,
214
- intermediate_size=self.intermediate_size_per_partition,
215
- local_expert_offset=self.ep_rank * self.local_num_experts,
216
- local_num_experts=self.local_num_experts,
217
- routed_scaling_factor=None,
218
- routing_method_type=self.routing_method_type,
219
- do_finalize=True,
220
- tune_max_num_tokens=max(self.max_capture_size, 1),
221
- output=output,
222
- )[0]
203
+ from vllm.utils.flashinfer import _is_fi_autotuning, autotune
204
+
205
+ with autotune(_is_fi_autotuning):
206
+ trtllm_fp4_block_scale_moe(
207
+ routing_logits=router_logits.to(torch.bfloat16),
208
+ routing_bias=None,
209
+ hidden_states=x_quant,
210
+ hidden_states_scale=x_scale,
211
+ gemm1_weights=w1,
212
+ gemm1_weights_scale=self.w1_scale,
213
+ gemm1_bias=self.w1_bias,
214
+ gemm1_alpha=self.gemm1_alpha,
215
+ gemm1_beta=self.gemm1_beta,
216
+ gemm1_clamp_limit=self.gemm1_clamp_limit,
217
+ gemm2_weights=w2,
218
+ gemm2_weights_scale=self.w2_scale,
219
+ gemm2_bias=self.w2_bias,
220
+ output1_scale_scalar=None,
221
+ output1_scale_gate_scalar=None,
222
+ output2_scale_scalar=None,
223
+ num_experts=global_num_experts,
224
+ top_k=self.topk,
225
+ n_group=None,
226
+ topk_group=None,
227
+ intermediate_size=self.intermediate_size_per_partition,
228
+ local_expert_offset=self.ep_rank * self.local_num_experts,
229
+ local_num_experts=self.local_num_experts,
230
+ routed_scaling_factor=None,
231
+ routing_method_type=self.routing_method_type,
232
+ do_finalize=True,
233
+ tune_max_num_tokens=max(self.max_capture_size, 1),
234
+ output=output,
235
+ )
236
+
237
+ return output
238
239
240
class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModular):
@@ -239,6 +254,16 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
254
) -> bool:
255
return True
256
257
+ @staticmethod
258
+ def _supports_routing_method(
259
+ routing_method: RoutingMethodType,
260
+ weight_key: QuantKey | None,
261
+ activation_key: QuantKey | None,
262
+ ) -> bool:
263
+ # Modular kernel handles only the expert computation;
264
+ # routing is done externally, so accept any routing method.
265
+ return True
266
+
267
def supports_expert_map(self) -> bool:
268
return True
269
@@ -282,7 +307,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
307
):
308
topk = topk_ids.size(-1)
309
local_num_experts = w1.size(0)
285
- intermediate_size = w2.size(1)
310
+ intermediate_size = self.intermediate_size_per_partition
311
local_expert_offset = self.moe_config.ep_rank * local_num_experts
312
313
# Handle input quantization
@@ -302,9 +327,8 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
327
x_quant = hidden_states
328
x_scale = None
329
305
- packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(
306
- torch.bfloat16
307
- ).view(torch.int16)
330
+ # Pack topk ids and weights into format expected by the kernel.
331
+ packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
332
333
assert self.w1_scale is not None
334
assert self.w2_scale is not None
@@ -333,7 +357,10 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
357
"local_expert_offset": local_expert_offset,
358
"local_num_experts": local_num_experts,
359
"routed_scaling_factor": None,
336
- "routing_method_type": self.routing_method_type,
360
+ # Modular kernel receives pre-routed tokens, so routing
361
+ # is already done. Use Renormalize as a safe default that
362
+ # the TRTLLM C++ kernel supports.
363
+ "routing_method_type": RoutingMethodType.Renormalize,
364
"do_finalize": True,
365
"output": output,
366
"tune_max_num_tokens": max(self.max_capture_size, 1),
@@ -341,12 +368,9 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
368
369
from flashinfer import trtllm_fp4_block_scale_routed_moe
370
344
- from vllm.utils.flashinfer import autotune
371
+ from vllm.utils.flashinfer import _is_fi_autotuning, autotune
372
346
- with autotune(False):
347
- # Enable autotune when,
348
- # https://github.com/flashinfer-ai/flashinfer/issues/2023 is
349
- # resolved.
373
+ with autotune(_is_fi_autotuning):
374
trtllm_fp4_block_scale_routed_moe(**kwargs)
375
376
return output
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+19
-5
@@ -50,6 +50,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
50
from vllm.platforms import current_platform
51
from vllm.scalar_type import ScalarType, scalar_types
52
53
+from .utils import swiglu_limit_func
54
+
55
56
def _fused_marlin_moe(
57
hidden_states: torch.Tensor,
@@ -88,6 +90,7 @@ def _fused_marlin_moe(
90
output: torch.Tensor | None = None,
91
input_dtype: torch.dtype | None = None,
92
is_k_full: bool = True,
93
+ clamp_limit: float | None = None,
94
) -> torch.Tensor:
95
assert hidden_states.ndim == 2
96
M, K = hidden_states.size()
@@ -155,11 +158,18 @@ def _fused_marlin_moe(
158
use_fp32_reduce=True,
159
is_zp_float=False,
160
)
158
- activation_func(
159
- activation,
160
- intermediate_cache2,
161
- intermediate_cache1.view(-1, w13_num_shards * N),
162
- )
161
+ if clamp_limit is not None and activation == MoEActivation.SILU:
162
+ swiglu_limit_func(
163
+ intermediate_cache2,
164
+ intermediate_cache1.view(-1, w13_num_shards * N),
165
+ clamp_limit,
166
+ )
167
+ else:
168
+ activation_func(
169
+ activation,
170
+ intermediate_cache2,
171
+ intermediate_cache1.view(-1, w13_num_shards * N),
172
+ )
173
174
if output is None:
175
output = intermediate_cache3
@@ -247,6 +257,7 @@ def fused_marlin_moe(
257
output: torch.Tensor | None = None,
258
input_dtype: torch.dtype | None = None,
259
inplace: bool = False,
260
+ clamp_limit: float | None = None,
261
) -> torch.Tensor:
262
"""
263
This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -363,6 +374,7 @@ def fused_marlin_moe(
374
output=None,
375
input_dtype=input_dtype,
376
is_k_full=is_k_full,
377
+ clamp_limit=clamp_limit,
378
).view(-1, topk, K)
379
380
if output is None:
@@ -557,6 +569,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
569
self.w2_g_idx_sort_indices = w2_g_idx_sort_indices
570
self.is_k_full = is_k_full
571
self.input_dtype = get_marlin_input_dtype()
572
+ self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
573
574
super().__init__(
575
moe_config=moe_config,
@@ -850,6 +863,7 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
863
sort_indices2=self.w2_g_idx_sort_indices,
864
is_k_full=self.is_k_full,
865
input_dtype=self.input_dtype,
866
+ clamp_limit=self.gemm1_clamp_limit,
867
)
868
869
def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+1
@@ -169,5 +169,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
169
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
170
x: torch.Tensor,
171
router_logits: torch.Tensor,
172
+ input_ids: torch.Tensor | None = None,
173
) -> torch.Tensor:
174
raise NotImplementedError
vllm/model_executor/layers/fused_moe/layer.py
+7
@@ -268,6 +268,7 @@ class FusedMoE(PluggableLayer):
268
custom_routing_function: Callable | None = None,
269
scoring_func: str = "softmax",
270
routed_scaling_factor: float = 1.0,
271
+ swiglu_limit: float | None = None,
272
e_score_correction_bias: torch.Tensor | None = None,
273
apply_router_weight_on_input: bool = False,
274
activation: str = "silu",
@@ -285,6 +286,7 @@ class FusedMoE(PluggableLayer):
286
routed_output_transform: torch.nn.Module | None = None,
287
apply_routed_scale_to_output: bool = False,
288
zero_expert_type: str | None = None,
289
+ hash_indices_table: torch.Tensor | None = None,
290
):
291
super().__init__()
292
@@ -294,6 +296,7 @@ class FusedMoE(PluggableLayer):
296
297
vllm_config = get_current_vllm_config()
298
self.vllm_config = vllm_config
299
+ self.swiglu_limit = swiglu_limit
300
301
# FIXME (varun): We should have a better way of inferring the activation
302
# datatype. This works for now as the tensor datatype entering the MoE
@@ -455,6 +458,7 @@ class FusedMoE(PluggableLayer):
458
self.e_score_correction_bias = e_score_correction_bias
459
# TODO(bnell): end attributes
460
461
+ self.hash_indices_table = hash_indices_table
462
self.apply_router_weight_on_input = apply_router_weight_on_input
463
self.activation = MoEActivation.from_str(activation)
464
@@ -479,6 +483,7 @@ class FusedMoE(PluggableLayer):
483
indices_type_getter=lambda: self.quant_method.topk_indices_dtype,
484
zero_expert_type=zero_expert_type,
485
num_logical_experts=self.logical_num_experts,
486
+ hash_indices_table=self.hash_indices_table,
487
)
488
self.routing_method_type: RoutingMethodType = self.router.routing_method_type
489
@@ -1541,10 +1546,12 @@ class FusedMoE(PluggableLayer):
1546
self,
1547
hidden_states: torch.Tensor,
1548
router_logits: torch.Tensor,
1549
+ input_ids: torch.Tensor | None = None,
1550
) -> torch.Tensor:
1551
return self.runner.forward(
1552
hidden_states,
1553
router_logits,
1554
+ input_ids,
1555
)
1556
1557
@property
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+409
-10
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
17
)
18
from vllm.model_executor.layers.fused_moe.config import (
19
FusedMoEQuantConfig,
20
+ FusedMoEQuantDesc,
21
mxfp4_mxfp8_moe_quant_config,
22
mxfp4_w4a16_moe_quant_config,
23
ocp_mx_moe_quant_config,
@@ -24,6 +25,7 @@ from vllm.model_executor.layers.fused_moe.config import (
25
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
26
from vllm.model_executor.layers.quantization.utils.quant_utils import (
27
QuantKey,
28
+ kFp8Dynamic128Sym,
29
kMxfp4Static,
30
kMxfp8Dynamic,
31
)
@@ -46,6 +48,8 @@ if has_triton_kernels():
48
49
class Mxfp4MoeBackend(Enum):
50
NONE = "None"
51
+ # DeepGEMM FP8xFP4 backend (SM100+)
52
+ DEEPGEMM_MXFP4 = "DEEPGEMM_MXFP4"
53
# FlashInfer TRTLLM backends
54
FLASHINFER_TRTLLM_MXFP4_MXFP8 = "FLASHINFER_TRTLLM_MXFP4_MXFP8"
55
FLASHINFER_TRTLLM_MXFP4_BF16 = "FLASHINFER_TRTLLM_MXFP4_BF16"
@@ -81,7 +85,14 @@ TRITON_BACKENDS = (
85
def backend_to_kernel_cls(
86
backend: Mxfp4MoeBackend,
87
) -> list[type[mk.FusedMoEExperts]]:
84
- if backend in (
88
+ if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
89
+ from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
90
+ DeepGemmFP4Experts,
91
+ )
92
+
93
+ return [DeepGemmFP4Experts]
94
+
95
+ elif backend in (
96
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
97
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
98
):
@@ -159,11 +170,13 @@ def backend_to_kernel_cls(
170
def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend:
171
"""Map user's moe_backend string to Mxfp4MoeBackend."""
172
mapping: dict[str, Mxfp4MoeBackend] = {
173
+ "deep_gemm": Mxfp4MoeBackend.DEEPGEMM_MXFP4,
174
"flashinfer_trtllm": Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
175
"flashinfer_trtllm_afp8": Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
176
"flashinfer_cutlass": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
177
"flashinfer_cutlass_afp8": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
178
"triton": Mxfp4MoeBackend.TRITON,
179
+ "triton_unfused": Mxfp4MoeBackend.TRITON_UNFUSED,
180
"marlin": Mxfp4MoeBackend.MARLIN,
181
"aiter": Mxfp4MoeBackend.AITER,
182
"xpu": Mxfp4MoeBackend.XPU,
@@ -177,7 +190,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend:
190
)
191
192
180
-def _get_priority_backends() -> list[Mxfp4MoeBackend]:
193
+def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:
194
"""
195
Get available backends in priority order based on platform and config.
196
Only includes BF16 backends. MXFP8 backends are selected via env vars.
@@ -187,7 +200,9 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
200
Mxfp4MoeBackend.AITER,
201
Mxfp4MoeBackend.TRITON,
202
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
190
- Mxfp4MoeBackend.TRITON_UNFUSED,
203
+ # TRITON_UNFUSED has bug with MTP support
204
+ # TODO re-enable after kernel is fixed
205
+ # TRITON_UNFUSED
206
Mxfp4MoeBackend.MARLIN,
207
Mxfp4MoeBackend.BATCHED_MARLIN,
208
Mxfp4MoeBackend.XPU,
@@ -196,8 +211,28 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
211
return _AVAILABLE_BACKENDS
212
213
214
+def _get_priority_backends() -> list[Mxfp4MoeBackend]:
215
+ """
216
+ Get available backends in priority order. SM100+ prefers DeepGEMM FP4 /
217
+ TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the
218
+ backend-level ``is_supported_config`` check filters by device capability).
219
+ """
220
+ _AVAILABLE_BACKENDS = [
221
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
222
+ Mxfp4MoeBackend.DEEPGEMM_MXFP4,
223
+ # TRITON_UNFUSED has bug with MTP support
224
+ # TODO re-enable after kernel is fixed
225
+ # TRITON_UNFUSED
226
+ Mxfp4MoeBackend.MARLIN,
227
+ Mxfp4MoeBackend.BATCHED_MARLIN,
228
+ ]
229
+ return _AVAILABLE_BACKENDS
230
+
231
+
232
def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:
200
- """Map backend to its activation key (MXFP8 or None for BF16)."""
233
+ """Map backend to its activation key (FP8, MXFP8, or None for BF16)."""
234
+ if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
235
+ return kFp8Dynamic128Sym
236
if backend in (
237
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
238
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
@@ -290,7 +325,7 @@ def select_gpt_oss_mxfp4_moe_backend(
325
)
326
327
# Select kernels in order of backend.
293
- AVAILABLE_BACKENDS = _get_priority_backends()
328
+ AVAILABLE_BACKENDS = _get_priority_backends_for_gpt_oss()
329
330
# Handle explicit FlashInfer MXFP4 BF16 configuration.
331
if envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16"):
@@ -387,11 +422,95 @@ def select_gpt_oss_mxfp4_moe_backend(
422
return Mxfp4MoeBackend.NONE, None
423
424
425
+def select_mxfp4_moe_backend(
426
+ config: FusedMoEConfig,
427
+) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:
428
+ """
429
+ Select the MXFP4 MoE backend with MXFP8 activation as top priority.
430
+ Falls back through BF16 and other backends.
431
+ """
432
+ activation_format = (
433
+ mk.FusedMoEActivationFormat.BatchedExperts
434
+ if config.moe_parallel_config.use_batched_activation_format
435
+ else mk.FusedMoEActivationFormat.Standard
436
+ )
437
+
438
+ def _make_log_backend(backend: Mxfp4MoeBackend):
439
+ return f"Using '{backend.value}' Mxfp4 MoE backend."
440
+
441
+ def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str:
442
+ if reason:
443
+ return (
444
+ f"Mxfp4 MoE backend '{backend.value}' does not support the "
445
+ f"deployment configuration since {reason}."
446
+ )
447
+ return (
448
+ f"Mxfp4 MoE backend '{backend.value}' does not support the "
449
+ "deployment configuration."
450
+ )
451
+
452
+ def _return_or_raise(
453
+ backend: Mxfp4MoeBackend,
454
+ config: FusedMoEConfig,
455
+ weight_key: QuantKey | None,
456
+ activation_key: QuantKey | None,
457
+ activation_format: mk.FusedMoEActivationFormat,
458
+ ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts]]:
459
+ reason: str | None = None
460
+ for k_cls in backend_to_kernel_cls(backend):
461
+ supported, reason = k_cls.is_supported_config(
462
+ k_cls, config, weight_key, activation_key, activation_format
463
+ )
464
+ if supported:
465
+ logger.info_once(_make_log_backend(backend), scope="local")
466
+ return backend, k_cls
467
+ raise ValueError(_make_log_unsupported(backend, reason))
468
+
469
+ # Honor explicit moe_backend (e.g. "marlin", "triton_unfused") before
470
+ # falling back to the auto priority list.
471
+ runner_backend = config.moe_backend
472
+ if runner_backend != "auto":
473
+ requested_backend = map_mxfp4_backend(runner_backend)
474
+ if (
475
+ activation_format == mk.FusedMoEActivationFormat.BatchedExperts
476
+ and requested_backend == Mxfp4MoeBackend.MARLIN
477
+ ):
478
+ requested_backend = Mxfp4MoeBackend.BATCHED_MARLIN
479
+ return _return_or_raise(
480
+ requested_backend,
481
+ config,
482
+ kMxfp4Static,
483
+ _backend_activation_key(requested_backend),
484
+ activation_format,
485
+ )
486
+
487
+ # Iterate priority backends: TRTLLM MXFP8, then Triton.
488
+ for backend in _get_priority_backends():
489
+ activation_key = _backend_activation_key(backend)
490
+ for k_cls in backend_to_kernel_cls(backend):
491
+ supported, reason = k_cls.is_supported_config(
492
+ k_cls, config, kMxfp4Static, activation_key, activation_format
493
+ )
494
+ if supported:
495
+ logger.info_once(_make_log_backend(backend), scope="local")
496
+ return backend, k_cls
497
+ else:
498
+ logger.debug_once(_make_log_unsupported(backend, reason), scope="local")
499
+
500
+ raise NotImplementedError(
501
+ "No MXFP4 MoE backend supports the deployment configuration."
502
+ )
503
+
504
+
505
def mxfp4_round_up_hidden_size_and_intermediate_size(
506
backend: Mxfp4MoeBackend, hidden_size: int, intermediate_size: int
507
) -> tuple[int, int]:
508
"""Round up hidden_size and intermediate_size based on backend requirements."""
394
- if backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):
509
+ if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
510
+ # DeepGEMM requires M/N/K alignment
511
+ intermediate_size = round_up(intermediate_size, 128)
512
+ hidden_size = round_up(hidden_size, 128)
513
+ elif backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):
514
intermediate_size = round_up(intermediate_size, 128)
515
if current_platform.is_xpu():
516
hidden_size = round_up(hidden_size, 128)
@@ -434,6 +553,20 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
553
]:
554
"""Convert loaded weights into backend-specific kernel format."""
555
556
+ if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
557
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
558
+ _upcast_e8m0_to_fp32,
559
+ )
560
+
561
+ return (
562
+ w13_weight.data,
563
+ w2_weight.data,
564
+ _upcast_e8m0_to_fp32(w13_weight_scale.data),
565
+ _upcast_e8m0_to_fp32(w2_weight_scale.data),
566
+ w13_bias,
567
+ w2_bias,
568
+ )
569
+
570
num_experts = w13_weight.shape[0]
571
intermediate_size = w13_weight.shape[1] // 2
572
hidden_size = w13_weight.shape[2] * 2
@@ -738,9 +871,10 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
871
elif mxfp4_backend in TRITON_BACKENDS:
872
from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
873
741
- assert w13_bias is not None and w2_bias is not None
742
- w13_bias = w13_bias.to(torch.float32)
743
- w2_bias = w2_bias.to(torch.float32)
874
+ if w13_bias is not None:
875
+ w13_bias = w13_bias.to(torch.float32)
876
+ if w2_bias is not None:
877
+ w2_bias = w2_bias.to(torch.float32)
878
879
w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
880
w13_weight,
@@ -797,15 +931,271 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
931
)
932
933
934
+def convert_weight_to_mxfp4_moe_kernel_format(
935
+ mxfp4_backend: Mxfp4MoeBackend,
936
+ layer: torch.nn.Module,
937
+ w13_weight: torch.Tensor,
938
+ w2_weight: torch.Tensor,
939
+ w13_weight_scale: torch.Tensor,
940
+ w2_weight_scale: torch.Tensor,
941
+ w13_bias: torch.Tensor | None = None,
942
+ w2_bias: torch.Tensor | None = None,
943
+ _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None,
944
+) -> tuple[
945
+ torch.Tensor,
946
+ torch.Tensor,
947
+ Union[torch.Tensor, "PrecisionConfig"],
948
+ Union[torch.Tensor, "PrecisionConfig"],
949
+ torch.Tensor | None,
950
+ torch.Tensor | None,
951
+]:
952
+ """Convert loaded weights into backend-specific kernel format.
953
+
954
+ Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends.
955
+ """
956
+
957
+ if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
958
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
959
+ _upcast_e8m0_to_fp32,
960
+ )
961
+
962
+ # Weights stay as uint8 packed FP4 — no layout change needed.
963
+ # Convert E8M0 uint8 scales to float32.
964
+ return (
965
+ w13_weight.data,
966
+ w2_weight.data,
967
+ _upcast_e8m0_to_fp32(w13_weight_scale.data),
968
+ _upcast_e8m0_to_fp32(w2_weight_scale.data),
969
+ w13_bias,
970
+ w2_bias,
971
+ )
972
+
973
+ if mxfp4_backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):
974
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
975
+ prepare_moe_mxfp4_layer_for_marlin,
976
+ )
977
+
978
+ return prepare_moe_mxfp4_layer_for_marlin(
979
+ layer,
980
+ w13_weight,
981
+ w2_weight,
982
+ w13_weight_scale,
983
+ w2_weight_scale,
984
+ w13_bias,
985
+ w2_bias,
986
+ )
987
+
988
+ num_experts = w13_weight.shape[0]
989
+ intermediate_size = w13_weight.shape[1] // 2
990
+ hidden_size = w13_weight.shape[2] * 2
991
+
992
+ sf_block_size = 32 # mxfp4 block size
993
+
994
+ if mxfp4_backend in TRTLLM_BACKENDS:
995
+ assert _cache_permute_indices is not None
996
+ from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
997
+ from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache
998
+
999
+ w13_weight = w13_weight.data
1000
+ w2_weight = w2_weight.data
1001
+ w13_weight_scale = w13_weight_scale.data
1002
+ w2_weight_scale = w2_weight_scale.data
1003
+ if w13_bias is not None:
1004
+ w13_bias = w13_bias.data.to(torch.float32)
1005
+ if w2_bias is not None:
1006
+ w2_bias = w2_bias.data.to(torch.float32)
1007
+
1008
+ # Swap w1/w3 and interleave to match TRTLLM SwiGLU convention.
1009
+ # Standard loading gives contiguous [w1/gate, w3/up].
1010
+ # TRTLLM kernel expects interleaved [w3_0, w1_0, w3_1, w1_1, ...].
1011
+ w1_weight = w13_weight[:, :intermediate_size, :]
1012
+ w3_weight = w13_weight[:, intermediate_size:, :]
1013
+ w13_weight = torch.stack([w3_weight, w1_weight], dim=2).reshape(
1014
+ w13_weight.shape
1015
+ )
1016
+
1017
+ w1_scale = w13_weight_scale[:, :intermediate_size, :]
1018
+ w3_scale = w13_weight_scale[:, intermediate_size:, :]
1019
+ w13_weight_scale = torch.stack([w3_scale, w1_scale], dim=2).reshape(
1020
+ w13_weight_scale.shape
1021
+ )
1022
+
1023
+ if w13_bias is not None:
1024
+ b1 = w13_bias[:, :intermediate_size]
1025
+ b3 = w13_bias[:, intermediate_size:]
1026
+ w13_bias = torch.stack([b3, b1], dim=2).reshape(w13_bias.shape)
1027
+
1028
+ # Shuffle weights and scaling factors for transposed mma output.
1029
+ # Permute indices depend only on shape (cached by torch.Size),
1030
+ # so compute once and apply to all experts via batched indexing.
1031
+ epilogue_tile_m = 128
1032
+
1033
+ # w13 weight permute
1034
+ w13_perm = get_w2_permute_indices_with_cache(
1035
+ _cache_permute_indices,
1036
+ w13_weight[0].view(torch.uint8),
1037
+ epilogue_tile_m,
1038
+ ).to(w13_weight.device)
1039
+ w13_weight = w13_weight.view(torch.uint8)[:, w13_perm].contiguous()
1040
+
1041
+ # w13 scale permute + interleave
1042
+ w13_sf_perm = get_w2_permute_indices_with_cache(
1043
+ _cache_permute_indices,
1044
+ w13_weight_scale[0].view(torch.uint8),
1045
+ epilogue_tile_m,
1046
+ num_elts_per_sf=16,
1047
+ ).to(w13_weight_scale.device)
1048
+ w13_s = w13_weight_scale.view(torch.uint8)[:, w13_sf_perm].contiguous()
1049
+ E, N_s, K_s = w13_s.shape
1050
+ w13_weight_scale = (
1051
+ nvfp4_block_scale_interleave(w13_s.reshape(E * N_s, K_s))
1052
+ .reshape(num_experts, 2 * intermediate_size, hidden_size // sf_block_size)
1053
+ .view(torch.float8_e4m3fn)
1054
+ )
1055
+
1056
+ # w2 weight permute
1057
+ w2_perm = get_w2_permute_indices_with_cache(
1058
+ _cache_permute_indices,
1059
+ w2_weight[0].view(torch.uint8),
1060
+ epilogue_tile_m,
1061
+ ).to(w2_weight.device)
1062
+ w2_weight = w2_weight.view(torch.uint8)[:, w2_perm].contiguous()
1063
+
1064
+ # w2 scale permute + interleave
1065
+ w2_sf_perm = get_w2_permute_indices_with_cache(
1066
+ _cache_permute_indices,
1067
+ w2_weight_scale[0].view(torch.uint8),
1068
+ epilogue_tile_m,
1069
+ num_elts_per_sf=16,
1070
+ ).to(w2_weight_scale.device)
1071
+ w2_s = w2_weight_scale.view(torch.uint8)[:, w2_sf_perm].contiguous()
1072
+ E2, N2_s, K2_s = w2_s.shape
1073
+ w2_weight_scale = (
1074
+ nvfp4_block_scale_interleave(w2_s.reshape(E2 * N2_s, K2_s))
1075
+ .reshape(num_experts, hidden_size, intermediate_size // sf_block_size)
1076
+ .view(torch.float8_e4m3fn)
1077
+ )
1078
+
1079
+ # w13 bias permute
1080
+ if w13_bias is not None:
1081
+ w13_b_perm = get_w2_permute_indices_with_cache(
1082
+ _cache_permute_indices,
1083
+ w13_bias[0].reshape(-1, 1),
1084
+ epilogue_tile_m,
1085
+ ).to(w13_bias.device)
1086
+ w13_bias = w13_bias.reshape(num_experts, -1, 1)[:, w13_b_perm].reshape(
1087
+ num_experts, -1
1088
+ )
1089
+
1090
+ # w2 bias permute
1091
+ if w2_bias is not None:
1092
+ w2_b_perm = get_w2_permute_indices_with_cache(
1093
+ _cache_permute_indices,
1094
+ w2_bias[0].reshape(-1, 1),
1095
+ epilogue_tile_m,
1096
+ ).to(w2_bias.device)
1097
+ w2_bias = w2_bias.reshape(num_experts, -1, 1)[:, w2_b_perm].reshape(
1098
+ num_experts, -1
1099
+ )
1100
+
1101
+ return (
1102
+ w13_weight,
1103
+ w2_weight,
1104
+ w13_weight_scale,
1105
+ w2_weight_scale,
1106
+ w13_bias,
1107
+ w2_bias,
1108
+ )
1109
+
1110
+ elif mxfp4_backend in TRITON_BACKENDS:
1111
+ from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1112
+
1113
+ if mxfp4_backend == Mxfp4MoeBackend.TRITON:
1114
+
1115
+ def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
1116
+ shape = w.shape
1117
+ n = shape[-1]
1118
+ first = w[..., : n // 2]
1119
+ second = w[..., n // 2 :]
1120
+ stacked = torch.stack((first, second), dim=-1)
1121
+ return stacked.reshape(shape)
1122
+
1123
+ w13_weight = shuffle_weight(w13_weight)
1124
+ w13_weight_scale = shuffle_weight(w13_weight_scale)
1125
+
1126
+ if w13_bias is not None:
1127
+ w13_bias = shuffle_weight(w13_bias.to(torch.float32))
1128
+ else:
1129
+ if w13_bias is not None:
1130
+ w13_bias = w13_bias.to(torch.float32)
1131
+
1132
+ if w2_bias is not None:
1133
+ w2_bias = w2_bias.to(torch.float32)
1134
+
1135
+ w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
1136
+ w13_weight,
1137
+ w13_weight_scale,
1138
+ )
1139
+ w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
1140
+ w2_weight,
1141
+ w2_weight_scale,
1142
+ )
1143
+
1144
+ w13_precision_config = PrecisionConfig(
1145
+ weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
1146
+ )
1147
+ w2_precision_config = PrecisionConfig(
1148
+ weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
1149
+ )
1150
+
1151
+ del layer.w13_weight
1152
+ del layer.w2_weight
1153
+
1154
+ return (
1155
+ w13_weight,
1156
+ w2_weight,
1157
+ w13_precision_config,
1158
+ w2_precision_config,
1159
+ w13_bias,
1160
+ w2_bias,
1161
+ )
1162
+ else:
1163
+ raise ValueError(
1164
+ f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
1165
+ f"Expected TRTLLM or Triton backend."
1166
+ )
1167
+
1168
+
1169
def make_mxfp4_moe_quant_config(
1170
mxfp4_backend: Mxfp4MoeBackend,
1171
w1_scale: Union[torch.Tensor, "PrecisionConfig"],
1172
w2_scale: Union[torch.Tensor, "PrecisionConfig"],
1173
+ gemm1_alpha: float | None = None,
1174
+ gemm1_beta: float | None = None,
1175
+ swiglu_limit: float | None = None,
1176
w1_bias: torch.Tensor | None = None,
1177
w2_bias: torch.Tensor | None = None,
1178
) -> FusedMoEQuantConfig | None:
1179
"""Create a FusedMoEQuantConfig for the given MXFP4 backend."""
808
- if mxfp4_backend in (
1180
+ if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
1181
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
1182
+ GroupShape,
1183
+ )
1184
+
1185
+ # DeepGEMM FP4 uses FP8 per-token-group activation quantization
1186
+ # with block 128, matching the FP8 DeepGEMM path.
1187
+ _fp8_dtype = current_platform.fp8_dtype()
1188
+ _block_shape = GroupShape(128, 128)
1189
+ return FusedMoEQuantConfig(
1190
+ _a1=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
1191
+ _a2=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
1192
+ _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
1193
+ _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
1194
+ gemm1_alpha=gemm1_alpha,
1195
+ gemm1_beta=gemm1_beta,
1196
+ gemm1_clamp_limit=swiglu_limit,
1197
+ )
1198
+ elif mxfp4_backend in (
1199
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
1200
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
1201
):
@@ -814,6 +1204,9 @@ def make_mxfp4_moe_quant_config(
1204
w2_bias=w2_bias,
1205
w1_scale=w1_scale,
1206
w2_scale=w2_scale,
1207
+ gemm1_alpha=gemm1_alpha,
1208
+ gemm1_beta=gemm1_beta,
1209
+ gemm1_clamp_limit=swiglu_limit,
1210
)
1211
elif mxfp4_backend in (
1212
Mxfp4MoeBackend.MARLIN,
@@ -829,6 +1222,9 @@ def make_mxfp4_moe_quant_config(
1222
w2_bias=w2_bias,
1223
w1_scale=w1_scale,
1224
w2_scale=w2_scale,
1225
+ gemm1_alpha=gemm1_alpha,
1226
+ gemm1_beta=gemm1_beta,
1227
+ gemm1_clamp_limit=swiglu_limit,
1228
)
1229
else:
1230
return ocp_mx_moe_quant_config(
@@ -837,6 +1233,9 @@ def make_mxfp4_moe_quant_config(
1233
w2_bias=w2_bias,
1234
w1_scale=w1_scale,
1235
w2_scale=w2_scale,
1236
+ gemm1_alpha=gemm1_alpha,
1237
+ gemm1_beta=gemm1_beta,
1238
+ gemm1_clamp_limit=swiglu_limit,
1239
)
1240
1241
vllm/model_executor/layers/fused_moe/router/base_router.py
+5
-1
@@ -228,6 +228,8 @@ class BaseRouter(FusedMoERouter):
228
hidden_states: torch.Tensor,
229
router_logits: torch.Tensor,
230
indices_type: torch.dtype | None,
231
+ *,
232
+ input_ids: torch.Tensor | None = None,
233
) -> tuple[torch.Tensor, torch.Tensor]:
234
"""
235
Compute the actual routing logic.
@@ -249,6 +251,8 @@ class BaseRouter(FusedMoERouter):
251
self,
252
hidden_states: torch.Tensor,
253
router_logits: torch.Tensor,
254
+ *,
255
+ input_ids: torch.Tensor | None = None,
256
) -> tuple[torch.Tensor, torch.Tensor]:
257
"""
258
Route the input hidden states to the top-k experts based on the
@@ -278,7 +282,7 @@ class BaseRouter(FusedMoERouter):
282
283
# Step 3: Compute routing (delegated to subclass)
284
topk_weights, topk_ids = self._compute_routing(
281
- hidden_states, router_logits, indices_type
285
+ hidden_states, router_logits, indices_type, input_ids=input_ids
286
)
287
288
# Capture logical ids before EPLB mapping.
vllm/model_executor/layers/fused_moe/router/custom_routing_router.py
+2
@@ -46,6 +46,8 @@ class CustomRoutingRouter(BaseRouter):
46
hidden_states: torch.Tensor,
47
router_logits: torch.Tensor,
48
indices_type: torch.dtype | None,
49
+ *,
50
+ input_ids: torch.Tensor | None = None,
51
) -> tuple[torch.Tensor, torch.Tensor]:
52
"""Compute routing using the custom routing function."""
53
topk_weights, topk_ids = self.custom_routing_function(
vllm/model_executor/layers/fused_moe/router/fused_moe_router.py
+2
@@ -31,6 +31,8 @@ class FusedMoERouter(ABC):
31
self,
32
hidden_states: torch.Tensor,
33
router_logits: torch.Tensor,
34
+ *,
35
+ input_ids: torch.Tensor | None = None,
36
) -> tuple[torch.Tensor, torch.Tensor]:
37
"""
38
Route the input hidden states to the top-k experts based on the
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+84
-16
@@ -4,6 +4,7 @@ import functools
4
from collections.abc import Callable
5
6
import torch
7
+import torch.nn.functional as F
8
9
import vllm._custom_ops as ops
10
import vllm.envs as envs
@@ -56,6 +57,32 @@ def vllm_topk_sigmoid(
57
return topk_weights, topk_indices
58
59
60
+def vllm_topk_softplus_sqrt(
61
+ topk_weights: torch.Tensor,
62
+ topk_indices: torch.Tensor,
63
+ token_expert_indices: torch.Tensor,
64
+ gating_output: torch.Tensor,
65
+ renormalize: bool = False,
66
+ e_score_correction_bias: torch.Tensor | None = None,
67
+ input_tokens: torch.Tensor | None = None,
68
+ hash_indices_table: torch.Tensor | None = None,
69
+ routed_scaling_factor: float = 1.0,
70
+) -> tuple[torch.Tensor, ...]:
71
+ ops.topk_hash_softplus_sqrt(
72
+ topk_weights,
73
+ topk_indices,
74
+ token_expert_indices,
75
+ gating_output,
76
+ renormalize,
77
+ routed_scaling_factor,
78
+ e_score_correction_bias,
79
+ input_tokens,
80
+ hash_indices_table,
81
+ )
82
+
83
+ return topk_weights, topk_indices
84
+
85
+
86
@functools.lru_cache(maxsize=8)
87
def _aiter_get_num_expert_group(num_experts: int) -> int:
88
_AITER_MAX_EXPERTS_PER_GROUP = 32
@@ -72,11 +99,14 @@ def _aiter_get_num_expert_group(num_experts: int) -> int:
99
def fused_topk_bias(
100
hidden_states: torch.Tensor,
101
gating_output: torch.Tensor,
102
+ scoring_func: str,
103
e_score_correction_bias: torch.Tensor,
104
topk: int,
105
renormalize: bool,
78
- scoring_func: str = "softmax",
106
indices_type: torch.dtype | None = None,
107
+ input_tokens: torch.Tensor | None = None,
108
+ hash_indices_table: torch.Tensor | None = None,
109
+ routed_scaling_factor: float = 1.0,
110
):
111
if not rocm_aiter_ops.is_fused_moe_enabled():
112
assert hidden_states.size(0) == gating_output.size(0), (
@@ -107,6 +137,8 @@ def fused_topk_bias(
137
renormalize,
138
e_score_correction_bias,
139
)
140
+ if routed_scaling_factor != 1.0:
141
+ topk_weights *= routed_scaling_factor
142
return topk_weights, topk_ids
143
elif scoring_func == "sigmoid":
144
topk_weights, topk_ids = vllm_topk_sigmoid(
@@ -117,9 +149,24 @@ def fused_topk_bias(
149
renormalize,
150
e_score_correction_bias,
151
)
152
+ if routed_scaling_factor != 1.0:
153
+ topk_weights *= routed_scaling_factor
154
return topk_weights, topk_ids
155
+ elif scoring_func == "sqrtsoftplus":
156
+ return vllm_topk_softplus_sqrt(
157
+ topk_weights,
158
+ topk_ids,
159
+ token_expert_indices,
160
+ gating_output,
161
+ renormalize,
162
+ e_score_correction_bias,
163
+ input_tokens,
164
+ hash_indices_table,
165
+ routed_scaling_factor,
166
+ )
167
else:
168
raise ValueError(f"Unsupported scoring function: {scoring_func}")
169
+
170
elif rocm_aiter_ops.is_fused_moe_enabled() and scoring_func == "sigmoid":
171
M = hidden_states.size(0)
172
num_experts = gating_output.shape[-1]
@@ -143,6 +190,8 @@ def fused_topk_bias(
190
topk_group=num_expert_group,
191
need_renorm=renormalize,
192
)
193
+ if routed_scaling_factor != 1.0:
194
+ topk_weights *= routed_scaling_factor
195
return topk_weights, topk_ids
196
197
n_routed_experts = gating_output.shape[-1]
@@ -150,20 +199,31 @@ def fused_topk_bias(
199
scores = gating_output.softmax(dim=-1)
200
elif scoring_func == "sigmoid":
201
scores = gating_output.sigmoid()
202
+ elif scoring_func == "sqrtsoftplus":
203
+ scores = F.softplus(gating_output).sqrt()
204
else:
205
raise ValueError(f"Unsupported scoring function: {scoring_func}")
155
-
156
- scores_for_choice = scores.view(
157
- -1, n_routed_experts
158
- ) + e_score_correction_bias.unsqueeze(0)
159
-
206
+ if e_score_correction_bias is not None:
207
+ scores_for_choice = scores.view(
208
+ -1, n_routed_experts
209
+ ) + e_score_correction_bias.unsqueeze(0)
210
+ else:
211
+ scores_for_choice = scores.view(-1, n_routed_experts)
212
# For batch invariance, use sorted=True to ensure deterministic expert selection
161
- use_sorted = envs.VLLM_BATCH_INVARIANT
162
- topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[1]
213
+ if hash_indices_table is not None:
214
+ topk_indices = hash_indices_table[input_tokens]
215
+ else:
216
+ use_sorted = envs.VLLM_BATCH_INVARIANT
217
+ topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[
218
+ 1
219
+ ]
220
topk_weights = scores.gather(1, topk_indices)
221
if renormalize:
222
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
166
- return topk_weights.to(torch.float32), topk_indices.to(
223
+ topk_weights = topk_weights.to(torch.float32)
224
+ if routed_scaling_factor != 1.0:
225
+ topk_weights *= routed_scaling_factor
226
+ return topk_weights, topk_indices.to(
227
torch.int32 if indices_type is None else indices_type
228
)
229
@@ -176,12 +236,14 @@ class FusedTopKBiasRouter(BaseRouter):
236
top_k: int,
237
global_num_experts: int,
238
eplb_state: EplbLayerState,
179
- e_score_correction_bias: torch.Tensor,
180
- scoring_func: str,
239
+ e_score_correction_bias: torch.Tensor | None = None,
240
renormalize: bool = True,
241
routed_scaling_factor: float = 1.0,
242
enable_eplb: bool = False,
243
indices_type_getter: Callable[[], torch.dtype | None] | None = None,
244
+ *,
245
+ scoring_func: str = "sigmoid",
246
+ hash_indices_table: torch.Tensor | None = None,
247
):
248
super().__init__(
249
top_k=top_k,
@@ -194,6 +256,8 @@ class FusedTopKBiasRouter(BaseRouter):
256
self.renormalize = renormalize
257
self.scoring_func = scoring_func
258
self.routed_scaling_factor = routed_scaling_factor
259
+ self.scoring_func = scoring_func
260
+ self._hash_indices_table = hash_indices_table
261
262
@property
263
def routing_method_type(self) -> RoutingMethodType:
@@ -210,19 +274,23 @@ class FusedTopKBiasRouter(BaseRouter):
274
hidden_states: torch.Tensor,
275
router_logits: torch.Tensor,
276
indices_type: torch.dtype | None,
277
+ *,
278
+ input_ids: torch.Tensor | None = None,
279
) -> tuple[torch.Tensor, torch.Tensor]:
280
"""Compute routing using fused top-k with bias."""
281
topk_weights, topk_ids = fused_topk_bias(
282
hidden_states=hidden_states,
283
gating_output=router_logits,
218
- e_score_correction_bias=self.e_score_correction_bias.data,
284
+ scoring_func=self.scoring_func,
285
+ e_score_correction_bias=self.e_score_correction_bias.data
286
+ if self.e_score_correction_bias is not None
287
+ else None,
288
topk=self.top_k,
289
renormalize=self.renormalize,
221
- scoring_func=self.scoring_func,
290
indices_type=indices_type,
291
+ input_tokens=input_ids,
292
+ hash_indices_table=self._hash_indices_table,
293
+ routed_scaling_factor=self.routed_scaling_factor,
294
)
295
225
- if self.routed_scaling_factor != 1.0:
226
- topk_weights *= self.routed_scaling_factor
227
-
296
return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/fused_topk_router.py
+2
@@ -151,6 +151,8 @@ class FusedTopKRouter(BaseRouter):
151
hidden_states: torch.Tensor,
152
router_logits: torch.Tensor,
153
indices_type: torch.dtype | None,
154
+ *,
155
+ input_ids: torch.Tensor | None = None,
156
) -> tuple[torch.Tensor, torch.Tensor]:
157
"""Compute routing using standard fused top-k."""
158
topk_weights, topk_ids, token_expert_indices = fused_topk(
vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
+3
@@ -292,6 +292,8 @@ class GroupedTopKRouter(BaseRouter):
292
hidden_states: torch.Tensor,
293
router_logits: torch.Tensor,
294
indices_type: torch.dtype | None,
295
+ *,
296
+ input_ids: torch.Tensor | None = None,
297
) -> tuple[torch.Tensor, torch.Tensor]:
298
"""Compute routing using grouped top-k."""
299
@@ -308,6 +310,7 @@ class GroupedTopKRouter(BaseRouter):
310
topk_weights, topk_ids = fused_topk_bias(
311
hidden_states=hidden_states,
312
gating_output=router_logits,
313
+ scoring_func=self.scoring_func,
314
e_score_correction_bias=self.e_score_correction_bias.data,
315
topk=self.top_k,
316
renormalize=self.renormalize,
vllm/model_executor/layers/fused_moe/router/router_factory.py
+9
-2
@@ -55,6 +55,7 @@ def create_fused_moe_router(
55
# zero expert parameters
56
zero_expert_type: str | None = None,
57
num_logical_experts: int | None = None,
58
+ hash_indices_table: torch.Tensor | None = None,
59
) -> FusedMoERouter:
60
"""
61
Factory function to create the appropriate FusedMoERouter subclass based on
@@ -99,6 +100,9 @@ def create_fused_moe_router(
100
num_logical_experts: Number of real (non-zero) experts. Required when
101
zero_expert_type is not None.
102
103
+ Hash Indices Table:
104
+ Used to map input_ids to experts, need for Deepseek V4
105
+
106
Returns:
107
An instance of the appropriate FusedMoERouter subclass
108
"""
@@ -179,17 +183,20 @@ def create_fused_moe_router(
183
indices_type_getter=indices_type_getter,
184
)
185
182
- if e_score_correction_bias is not None:
186
+ assert scoring_func in ["sigmoid", "softmax", "sqrtsoftplus"]
187
+
188
+ if e_score_correction_bias is not None or hash_indices_table is not None:
189
return FusedTopKBiasRouter(
190
top_k=top_k,
191
global_num_experts=global_num_experts,
192
eplb_state=eplb_state,
193
e_score_correction_bias=e_score_correction_bias,
188
- scoring_func=scoring_func,
194
renormalize=renormalize,
195
routed_scaling_factor=routed_scaling_factor,
196
enable_eplb=enable_eplb,
197
indices_type_getter=indices_type_getter,
198
+ scoring_func=scoring_func,
199
+ hash_indices_table=hash_indices_table,
200
)
201
202
return FusedTopKRouter(
vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py
+2
@@ -334,6 +334,8 @@ class RoutingSimulatorRouter(BaseRouter):
334
hidden_states: torch.Tensor,
335
router_logits: torch.Tensor,
336
indices_type: torch.dtype | None,
337
+ *,
338
+ input_ids: torch.Tensor | None = None,
339
) -> tuple[torch.Tensor, torch.Tensor]:
340
"""Use routing simulator to compute routing."""
341
routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
vllm/model_executor/layers/fused_moe/router/zero_expert_router.py
+2
@@ -72,6 +72,8 @@ class ZeroExpertRouter(BaseRouter):
72
hidden_states: torch.Tensor,
73
router_logits: torch.Tensor,
74
indices_type: torch.dtype | None,
75
+ *,
76
+ input_ids: torch.Tensor | None = None,
77
) -> tuple[torch.Tensor, torch.Tensor]:
78
"""Compute routing with full bias, compute zero expert output,
79
mask zero expert IDs."""
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+13
@@ -91,6 +91,7 @@ def _moe_forward(
91
hidden_states: torch.Tensor,
92
router_logits: torch.Tensor,
93
shared_experts_input: torch.Tensor | None,
94
+ input_ids: torch.Tensor | None,
95
layer_name: _layer_name_type,
96
) -> torch.Tensor:
97
layer = get_layer_from_name(_resolve_layer_name(layer_name))
@@ -99,6 +100,7 @@ def _moe_forward(
100
hidden_states,
101
router_logits,
102
shared_experts_input,
103
+ input_ids,
104
)
105
106
@@ -106,6 +108,7 @@ def _moe_forward_fake(
108
hidden_states: torch.Tensor,
109
router_logits: torch.Tensor,
110
shared_experts_input: torch.Tensor | None,
111
+ input_ids: torch.Tensor | None,
112
layer_name: _layer_name_type,
113
) -> torch.Tensor:
114
return torch.empty_like(hidden_states)
@@ -115,6 +118,7 @@ def _moe_forward_shared(
118
hidden_states: torch.Tensor,
119
router_logits: torch.Tensor,
120
shared_experts_input: torch.Tensor | None,
121
+ input_ids: torch.Tensor | None,
122
layer_name: _layer_name_type,
123
) -> tuple[torch.Tensor, torch.Tensor]:
124
layer = get_layer_from_name(_resolve_layer_name(layer_name))
@@ -123,6 +127,7 @@ def _moe_forward_shared(
127
hidden_states,
128
router_logits,
129
shared_experts_input,
130
+ input_ids,
131
)
132
133
@@ -130,6 +135,7 @@ def _moe_forward_shared_fake(
135
hidden_states: torch.Tensor,
136
router_logits: torch.Tensor,
137
shared_experts_input: torch.Tensor | None,
138
+ input_ids: torch.Tensor | None,
139
layer_name: _layer_name_type,
140
) -> tuple[torch.Tensor, torch.Tensor]:
141
# Output shapes:
@@ -433,6 +439,7 @@ class MoERunner(MoERunnerInterface):
439
hidden_states: torch.Tensor,
440
router_logits: torch.Tensor,
441
shared_experts_input: torch.Tensor | None,
442
+ input_ids: torch.Tensor | None = None,
443
) -> tuple[torch.Tensor | None, torch.Tensor]:
444
"""Run expert routing and the fused MoE kernel via the quant method.
445
@@ -449,11 +456,13 @@ class MoERunner(MoERunnerInterface):
456
layer=layer,
457
x=hidden_states,
458
router_logits=router_logits,
459
+ input_ids=input_ids,
460
)
461
else:
462
topk_weights, topk_ids = self.router.select_experts(
463
hidden_states=hidden_states,
464
router_logits=router_logits,
465
+ input_ids=input_ids,
466
)
467
468
# Passing shared_experts_input in case SharedExpertsOrder is
@@ -523,6 +532,7 @@ class MoERunner(MoERunnerInterface):
532
self,
533
hidden_states: torch.Tensor,
534
router_logits: torch.Tensor,
535
+ input_ids: torch.Tensor | None = None,
536
) -> torch.Tensor:
537
"""Invoke the fused moe layer.
538
@@ -565,6 +575,7 @@ class MoERunner(MoERunnerInterface):
575
hidden_states,
576
router_logits,
577
shared_experts_input,
578
+ input_ids,
579
self._encode_layer_name(),
580
)
581
@@ -672,6 +683,7 @@ class MoERunner(MoERunnerInterface):
683
hidden_states: torch.Tensor,
684
router_logits: torch.Tensor,
685
shared_experts_input: torch.Tensor | None,
686
+ input_ids: torch.Tensor | None = None,
687
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
688
"""Entry point called by the custom op to run the MoE computation.
689
@@ -712,6 +724,7 @@ class MoERunner(MoERunnerInterface):
724
hidden_states=hidden_states,
725
router_logits=router_logits,
726
shared_experts_input=shared_experts_input,
727
+ input_ids=input_ids,
728
)
729
730
return self._maybe_combine(
vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py
+1
@@ -26,6 +26,7 @@ class MoERunnerInterface(ABC):
26
self,
27
hidden_states: torch.Tensor,
28
router_logits: torch.Tensor,
29
+ input_ids: torch.Tensor | None = None,
30
) -> torch.Tensor:
31
raise NotImplementedError
32
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+1
@@ -309,6 +309,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
309
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
310
x: torch.Tensor,
311
router_logits: torch.Tensor,
312
+ input_ids: torch.Tensor | None = None,
313
) -> torch.Tensor:
314
assert self.is_monolithic
315
if self.unquantized_backend == UnquantizedMoeBackend.CPU:
vllm/model_executor/layers/fused_moe/utils.py
+18
@@ -4,6 +4,7 @@ import functools
4
from math import prod
5
6
import torch
7
+import torch.nn.functional as F
8
9
from vllm import _custom_ops as ops
10
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
@@ -384,3 +385,20 @@ def trtllm_moe_pack_topk_ids_weights(
385
return (topk_ids.to(torch.int32) << 16) | topk_weights.to(torch.bfloat16).view(
386
torch.int16
387
)
388
+
389
+
390
+@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
391
+def swiglu_limit_func(
392
+ output: torch.Tensor,
393
+ input: torch.Tensor, # first half is gate, second half is up
394
+ swiglu_limit: float = 0.0,
395
+) -> None:
396
+ d = input.shape[1] // 2
397
+ gate = input[:, :d]
398
+ up = input[:, d:]
399
+
400
+ if swiglu_limit > 0:
401
+ gate = torch.clamp(gate, max=swiglu_limit)
402
+ up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit)
403
+
404
+ output.copy_(F.silu(gate) * up)
vllm/model_executor/layers/mhc.py
new
+450
@@ -0,0 +1,450 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
+from functools import cache
5
+from typing import TYPE_CHECKING
6
+
7
+import torch
8
+
9
+from vllm.platforms import current_platform
10
+from vllm.utils.import_utils import has_tilelang
11
+from vllm.utils.math_utils import cdiv
12
+from vllm.utils.torch_utils import direct_register_custom_op
13
+
14
+# tilelang is only available on CUDA platforms
15
+if TYPE_CHECKING or current_platform.is_cuda_alike():
16
+ if not has_tilelang():
17
+ raise ImportError(
18
+ "tilelang is required for mhc but is not installed. Install it with "
19
+ "`pip install tilelang`."
20
+ )
21
+ import tilelang
22
+ import tilelang.language as T
23
+else:
24
+ tilelang = None # type: ignore[assignment]
25
+ T = None # type: ignore[assignment]
26
+
27
+
28
+@cache
29
+def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
30
+ device_props = torch.cuda.get_device_properties(0)
31
+ n_sms = device_props.multi_processor_count
32
+ split_k = n_sms // grid_size
33
+ if k is not None:
34
+ # avoid split_k for small k
35
+ num_block_k = cdiv(k, block_k)
36
+ split_k = min(split_k, num_block_k // 4)
37
+ split_k = max(split_k, 1)
38
+ return split_k
39
+
40
+
41
+@tilelang.jit(
42
+ pass_configs={
43
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
44
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
45
+ tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
46
+ },
47
+)
48
+def mhc_pre_big_fuse_tilelang(
49
+ gemm_out_mul,
50
+ gemm_out_sqrsum,
51
+ hc_scale,
52
+ hc_base,
53
+ residual,
54
+ post_mix,
55
+ comb_mix,
56
+ layer_input,
57
+ hidden_size: int,
58
+ rms_eps: float,
59
+ hc_pre_eps: float,
60
+ hc_sinkhorn_eps: float,
61
+ hc_post_mult_value: float,
62
+ sinkhorn_repeat: int,
63
+ n_splits: int = 16,
64
+ hc_mult: int = 4,
65
+):
66
+ """Deeply fused kernels, everything other than gemm & sqrsum in mHC pre block."""
67
+ num_tokens = T.dynamic("num_tokens")
68
+ hc_mult3 = hc_mult * (2 + hc_mult)
69
+ hidden_block = math.gcd(512, hidden_size)
70
+
71
+ gemm_out_mul: T.Tensor[[n_splits, num_tokens, hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
72
+ gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type]
73
+ hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type]
74
+ hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
75
+ residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
76
+ # outputs
77
+ post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
78
+ comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type]
79
+ layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
80
+
81
+ with T.Kernel(num_tokens, threads=96) as i:
82
+ T.pdl_sync()
83
+ ##################################################################
84
+ # _pre_norm_fn_fwd_norm
85
+ rms = T.alloc_fragment(1, T.float32)
86
+ mixes = T.alloc_fragment(hc_mult3, T.float32)
87
+ T.clear(mixes)
88
+ rms[0] = 0
89
+ for i_split in T.serial(n_splits):
90
+ rms[0] += gemm_out_sqrsum[i_split, i]
91
+ rms[0] = T.rsqrt(rms[0] / (hc_mult * hidden_size) + rms_eps)
92
+ for j in T.Parallel(hc_mult3):
93
+ mixes[j] = 0
94
+ for i_split in T.serial(n_splits):
95
+ mixes[j] += gemm_out_mul[i_split, i, j]
96
+ mixes[j] *= rms[0]
97
+ mixes_shared = T.alloc_shared(hc_mult3, T.float32)
98
+ T.copy(mixes, mixes_shared)
99
+
100
+ if T.get_thread_binding() < 32:
101
+ ##################################################################
102
+ # _pre_split_mixes_fwd (post & comb)
103
+ cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
104
+ for j in T.Parallel(hc_mult):
105
+ post_mix[i, j] = (
106
+ T.sigmoid(
107
+ mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
108
+ )
109
+ * hc_post_mult_value
110
+ )
111
+ for j, k in T.Parallel(hc_mult, hc_mult):
112
+ cm[j, k] = (
113
+ mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
114
+ + hc_base[j * hc_mult + k + hc_mult * 2]
115
+ )
116
+
117
+ ##################################################################
118
+ # _sinkhorn_fwd
119
+ row_sum = T.alloc_fragment(hc_mult, T.float32)
120
+ col_sum = T.alloc_fragment(hc_mult, T.float32)
121
+
122
+ # comb = comb.softmax(-1) + eps
123
+ row_max = T.alloc_fragment(hc_mult, T.float32)
124
+ T.reduce_max(cm, row_max, dim=1)
125
+ for j, k in T.Parallel(hc_mult, hc_mult):
126
+ cm[j, k] = T.exp(cm[j, k] - row_max[j])
127
+ T.reduce_sum(cm, row_sum, dim=1)
128
+ for j, k in T.Parallel(hc_mult, hc_mult):
129
+ cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps
130
+
131
+ # comb = comb / (comb.sum(-2) + eps)
132
+ T.reduce_sum(cm, col_sum, dim=0)
133
+ for j, k in T.Parallel(hc_mult, hc_mult):
134
+ cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
135
+
136
+ for _ in T.serial(sinkhorn_repeat - 1):
137
+ # comb = comb / (comb.sum(-1) + eps)
138
+ T.reduce_sum(cm, row_sum, dim=1)
139
+ for j, k in T.Parallel(hc_mult, hc_mult):
140
+ cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)
141
+
142
+ # comb = comb / (comb.sum(-2) + eps)
143
+ T.reduce_sum(cm, col_sum, dim=0)
144
+ for j, k in T.Parallel(hc_mult, hc_mult):
145
+ cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
146
+
147
+ # save comb_mix to global memory
148
+ for j, k in T.Parallel(hc_mult, hc_mult):
149
+ comb_mix[i, j * hc_mult + k] = cm[j, k]
150
+ else:
151
+ ##################################################################
152
+ # _pre_split_mixes_fwd (pre)
153
+ pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
154
+ for j in T.Parallel(hc_mult):
155
+ pre_mix_shared[j] = (
156
+ T.sigmoid(
157
+ mixes_shared[j] * hc_scale[0] + hc_base[j],
158
+ )
159
+ + hc_pre_eps
160
+ )
161
+ ###################################################################
162
+ # _pre_apply_mix_fwd
163
+ for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
164
+ xs = T.alloc_shared((hc_mult, hidden_block), T.float32)
165
+ xl = T.alloc_fragment((hc_mult, hidden_block), T.float32)
166
+ T.copy(residual[i, 0, i0_h * hidden_block], xs)
167
+ T.copy(xs, xl)
168
+
169
+ ol = T.alloc_fragment(hidden_block, T.float32)
170
+ T.clear(ol)
171
+
172
+ for i_hc in T.serial(hc_mult):
173
+ pre = pre_mix_shared[i_hc]
174
+ for i1_h in T.Parallel(hidden_block):
175
+ ol[i1_h] += pre * xl[i_hc, i1_h]
176
+
177
+ T.copy(ol, layer_input[i, i0_h * hidden_block])
178
+ T.pdl_trigger()
179
+
180
+
181
+def mhc_pre(
182
+ residual: torch.Tensor,
183
+ fn: torch.Tensor,
184
+ hc_scale: torch.Tensor,
185
+ hc_base: torch.Tensor,
186
+ rms_eps: float,
187
+ hc_pre_eps: float,
188
+ hc_sinkhorn_eps: float,
189
+ hc_post_mult_value: float,
190
+ sinkhorn_repeat: int,
191
+ n_splits: int = 1,
192
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
193
+ """
194
+ Forward pass for mHC pre block.
195
+
196
+ Args:
197
+ residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
198
+ fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
199
+ hc_scale: shape (3,), dtype torch.float32
200
+ hc_base: shape (hc_mult3,), dtype torch.float32
201
+ rms_eps: RMS normalization epsilon
202
+ hc_pre_eps: pre-mix epsilon
203
+ hc_sinkhorn_eps: sinkhorn epsilon
204
+ hc_post_mult_value: post-mix multiplier value
205
+ sinkhorn_repeat: number of sinkhorn iterations
206
+ n_splits: split-k factor;
207
+
208
+ Returns:
209
+ post_mix: shape (..., hc_mult), dtype torch.float32
210
+ comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
211
+ layer_input: shape (..., hidden_size), dtype torch.bfloat16
212
+ """
213
+
214
+ # Validate shapes
215
+ assert residual.dtype == torch.bfloat16
216
+ assert fn.dtype == torch.float32
217
+ assert hc_scale.dtype == torch.float32
218
+ assert hc_base.dtype == torch.float32
219
+
220
+ hc_mult = residual.shape[-2]
221
+ hidden_size = residual.shape[-1]
222
+ hc_mult2 = hc_mult * hc_mult
223
+ hc_mult3 = hc_mult * 2 + hc_mult2
224
+
225
+ hc_hidden_size = hc_mult * hidden_size
226
+ assert fn.shape[0] == hc_mult3
227
+ assert fn.shape[1] == hc_hidden_size
228
+ assert hc_scale.shape == (3,)
229
+ assert hc_base.shape == (hc_mult3,)
230
+
231
+ outer_shape = residual.shape[:-2]
232
+
233
+ residual_flat = residual.view(-1, hc_mult, hidden_size)
234
+ num_tokens = residual_flat.shape[0]
235
+ fn_flat = fn
236
+
237
+ # these number are from deepgemm kernel impl
238
+ block_k = 64
239
+ block_m = 64
240
+ n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
241
+
242
+ post_mix = torch.empty(
243
+ num_tokens,
244
+ hc_mult,
245
+ dtype=torch.float32,
246
+ device=residual.device,
247
+ )
248
+ comb_mix = torch.empty(
249
+ num_tokens,
250
+ hc_mult2,
251
+ dtype=torch.float32,
252
+ device=residual.device,
253
+ )
254
+ layer_input = torch.empty(
255
+ num_tokens,
256
+ hidden_size,
257
+ dtype=torch.bfloat16,
258
+ device=residual.device,
259
+ )
260
+
261
+ gemm_out_mul = torch.empty(
262
+ n_splits,
263
+ num_tokens,
264
+ hc_mult3,
265
+ dtype=torch.float32,
266
+ device=residual.device,
267
+ )
268
+ gemm_out_sqrsum = torch.empty(
269
+ n_splits,
270
+ num_tokens,
271
+ dtype=torch.float32,
272
+ device=residual.device,
273
+ )
274
+
275
+ from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
276
+
277
+ tf32_hc_prenorm_gemm(
278
+ residual_flat.view(num_tokens, hc_mult * hidden_size),
279
+ fn_flat,
280
+ gemm_out_mul,
281
+ gemm_out_sqrsum,
282
+ n_splits,
283
+ )
284
+
285
+ mhc_pre_big_fuse_tilelang(
286
+ gemm_out_mul,
287
+ gemm_out_sqrsum,
288
+ hc_scale,
289
+ hc_base,
290
+ residual_flat,
291
+ post_mix,
292
+ comb_mix,
293
+ layer_input,
294
+ hidden_size,
295
+ rms_eps,
296
+ hc_pre_eps,
297
+ hc_sinkhorn_eps,
298
+ hc_post_mult_value,
299
+ sinkhorn_repeat,
300
+ n_splits,
301
+ hc_mult,
302
+ )
303
+
304
+ post_mix = post_mix.view(*outer_shape, hc_mult, 1)
305
+ comb_mix = comb_mix.view(*outer_shape, hc_mult, hc_mult)
306
+ layer_input = layer_input.view(*outer_shape, hidden_size)
307
+
308
+ return post_mix, comb_mix, layer_input
309
+
310
+
311
+def _mhc_pre_fake(
312
+ residual: torch.Tensor,
313
+ fn: torch.Tensor,
314
+ hc_scale: torch.Tensor,
315
+ hc_base: torch.Tensor,
316
+ rms_eps: float,
317
+ hc_pre_eps: float,
318
+ hc_sinkhorn_eps: float,
319
+ hc_post_mult_value: float,
320
+ sinkhorn_repeat: int,
321
+ n_splits: int = 1,
322
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
323
+ hc_mult = residual.shape[-2]
324
+ hidden_size = residual.shape[-1]
325
+ outer_shape = residual.shape[:-2]
326
+
327
+ # Create empty tensors with correct shapes for meta device / shape inference
328
+ post_mix = torch.empty(
329
+ *outer_shape,
330
+ hc_mult,
331
+ 1,
332
+ dtype=torch.float32,
333
+ device=residual.device,
334
+ )
335
+ comb_mix = torch.empty(
336
+ *outer_shape,
337
+ hc_mult,
338
+ hc_mult,
339
+ dtype=torch.float32,
340
+ device=residual.device,
341
+ )
342
+ layer_input = torch.empty(
343
+ *outer_shape,
344
+ hidden_size,
345
+ dtype=torch.bfloat16,
346
+ device=residual.device,
347
+ )
348
+
349
+ return post_mix, comb_mix, layer_input
350
+
351
+
352
+@tilelang.jit(
353
+ pass_configs={
354
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
355
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
356
+ tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
357
+ },
358
+)
359
+def mhc_post_tilelang(
360
+ a,
361
+ b,
362
+ c,
363
+ d,
364
+ x,
365
+ hc: int,
366
+ hidden: int,
367
+ n_thr: int = 128,
368
+ h_blk: int = 1024,
369
+) -> tilelang.JITKernel:
370
+ # rename for shorter code
371
+ n = T.dynamic("num_tokens")
372
+ h = hidden
373
+
374
+ h_blk = math.gcd(hidden, h_blk)
375
+ a: T.Tensor((n, hc, hc), T.float32) # type: ignore[no-redef, valid-type]
376
+ b: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
377
+ c: T.Tensor((n, hc), T.float32) # type: ignore[no-redef, valid-type]
378
+ d: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type]
379
+ x: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
380
+ with T.Kernel(n, threads=n_thr) as i_n:
381
+ x_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
382
+ b_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
383
+ d_shared = T.alloc_shared(h_blk, T.bfloat16)
384
+
385
+ x_local = T.alloc_fragment((hc, h_blk), T.float32)
386
+ b_local = T.alloc_fragment((hc, h_blk), T.float32)
387
+ d_local = T.alloc_fragment(h_blk, T.float32)
388
+
389
+ a_local = T.alloc_fragment((hc, hc), T.float32)
390
+ c_local = T.alloc_fragment(hc, T.float32)
391
+ T.pdl_sync()
392
+ T.copy(a[i_n, 0, 0], a_local)
393
+ T.copy(c[i_n, 0], c_local)
394
+
395
+ for i0_h in T.Pipelined(T.ceildiv(h, h_blk), num_stages=2):
396
+ T.copy(b[i_n, 0, i0_h * h_blk], b_shared)
397
+ T.copy(d[i_n, i0_h * h_blk], d_shared)
398
+
399
+ T.copy(b_shared, b_local)
400
+ T.copy(d_shared, d_local)
401
+ for i_hco, i1_h in T.Parallel(hc, h_blk):
402
+ x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h]
403
+ for i_hci in T.serial(hc):
404
+ x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]
405
+ T.copy(x_local, x_shared)
406
+
407
+ T.copy(x_shared, x[i_n, 0, i0_h * h_blk])
408
+ T.pdl_trigger()
409
+
410
+
411
+def mhc_post(
412
+ x: torch.Tensor,
413
+ residual: torch.Tensor,
414
+ post_layer_mix: torch.Tensor,
415
+ comb_res_mix: torch.Tensor,
416
+) -> torch.Tensor:
417
+ out = torch.empty_like(residual)
418
+ mhc_post_tilelang(
419
+ comb_res_mix,
420
+ residual,
421
+ post_layer_mix.squeeze(-1),
422
+ x,
423
+ out,
424
+ residual.shape[-2],
425
+ residual.shape[-1],
426
+ )
427
+ return out
428
+
429
+
430
+def _mhc_post_fake(
431
+ x: torch.Tensor,
432
+ residual: torch.Tensor,
433
+ post_layer_mix: torch.Tensor,
434
+ comb_res_mix: torch.Tensor,
435
+) -> torch.Tensor:
436
+ return torch.empty_like(residual)
437
+
438
+
439
+direct_register_custom_op(
440
+ op_name="mhc_pre",
441
+ op_func=mhc_pre,
442
+ mutates_args=[],
443
+ fake_impl=_mhc_pre_fake,
444
+)
445
+direct_register_custom_op(
446
+ op_name="mhc_post",
447
+ op_func=mhc_post,
448
+ mutates_args=[],
449
+ fake_impl=_mhc_post_fake,
450
+)
vllm/model_executor/layers/quantization/__init__.py
+3
@@ -32,6 +32,7 @@ QuantizationMethods = Literal[
32
"inc",
33
"mxfp4",
34
"gpt_oss_mxfp4",
35
+ "deepseek_v4_fp8",
36
"cpu_awq",
37
"online",
38
# Below are values of the OnlineQuantScheme enum, specified as strings to
@@ -112,6 +113,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
113
# lazy import to avoid triggering `torch.compile` too early
114
from vllm.config.quantization import OnlineQuantScheme
115
from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
116
+ from vllm.model_executor.models.deepseek_v4 import DeepseekV4FP8Config
117
118
from .awq import AWQConfig
119
from .awq_marlin import AWQMarlinConfig
@@ -163,6 +165,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
165
"inc": INCConfig,
166
"mxfp4": Mxfp4Config,
167
"gpt_oss_mxfp4": GptOssMxfp4Config,
168
+ "deepseek_v4_fp8": DeepseekV4FP8Config,
169
"cpu_awq": CPUAWQConfig,
170
"humming": HummingConfig,
171
"online": OnlineQuantizationConfig,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+1
@@ -265,6 +265,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
265
layer: FusedMoE,
266
x: torch.Tensor,
267
router_logits: torch.Tensor,
268
+ input_ids: torch.Tensor | None = None,
269
) -> torch.Tensor:
270
assert self.is_monolithic
271
assert self.moe_kernel is not None
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py
+1
@@ -305,6 +305,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
305
layer: FusedMoE,
306
x: torch.Tensor,
307
router_logits: torch.Tensor,
308
+ input_ids: torch.Tensor | None = None,
309
) -> torch.Tensor:
310
assert not layer.enable_eplb, "EPLB not supported for W4A8-int MoE yet."
311
assert layer.activation in (
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
+1
@@ -367,6 +367,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
367
layer: FusedMoE,
368
x: torch.Tensor,
369
router_logits: torch.Tensor,
370
+ input_ids: torch.Tensor | None = None,
371
) -> torch.Tensor:
372
assert self.moe_kernel is not None
373
return self.moe_kernel.apply_monolithic(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py
+1
@@ -168,6 +168,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
168
layer: FusedMoE,
169
x: torch.Tensor,
170
router_logits: torch.Tensor,
171
+ input_ids: torch.Tensor | None = None,
172
) -> torch.Tensor:
173
assert self.moe_kernel is not None
174
return self.moe_kernel.apply_monolithic(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+1
@@ -517,6 +517,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
517
layer: FusedMoE,
518
x: torch.Tensor,
519
router_logits: torch.Tensor,
520
+ input_ids: torch.Tensor | None = None,
521
) -> torch.Tensor:
522
assert self.kernel_backend == "Flashinfer"
523
return flashinfer_trtllm_mxint4_moe(
vllm/model_executor/layers/quantization/fp8.py
+3
@@ -269,6 +269,7 @@ class Fp8LinearMethod(LinearMethodBase):
269
270
def __init__(self, quant_config: Fp8Config):
271
self.quant_config = quant_config
272
+ self.is_scale_e8m0 = getattr(quant_config, "is_scale_e8m0", False)
273
self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
274
self.out_dtype = torch.get_default_dtype()
275
self.input_dtype = get_current_vllm_config().model_config.dtype
@@ -362,6 +363,7 @@ class Fp8LinearMethod(LinearMethodBase):
363
input_size_per_partition,
364
self.weight_block_size,
365
weight_loader,
366
+ scale_dtype=(torch.float8_e8m0fnu if self.is_scale_e8m0 else None),
367
)
368
# The weight_scale_inv name is intentional for deepseekv3
369
layer.register_parameter("weight_scale_inv", scale)
@@ -866,6 +868,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
868
layer: FusedMoE,
869
x: torch.Tensor,
870
router_logits: torch.Tensor,
871
+ input_ids: torch.Tensor | None = None,
872
) -> torch.Tensor:
873
assert self.is_monolithic
874
assert self.moe_kernel is not None
vllm/model_executor/layers/quantization/modelopt.py
+3
@@ -950,6 +950,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
950
layer: FusedMoE,
951
x: torch.Tensor,
952
router_logits: torch.Tensor,
953
+ input_ids: torch.Tensor | None = None,
954
) -> torch.Tensor:
955
assert self.is_monolithic
956
assert self.moe_kernel is not None
@@ -1442,6 +1443,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1443
layer: FusedMoE,
1444
x: torch.Tensor,
1445
router_logits: torch.Tensor,
1446
+ input_ids: torch.Tensor | None = None,
1447
) -> torch.Tensor:
1448
assert self.is_monolithic
1449
assert self.moe_kernel is not None
@@ -1920,6 +1922,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1922
layer: FusedMoE,
1923
x: torch.Tensor,
1924
router_logits: torch.Tensor,
1925
+ input_ids: torch.Tensor | None = None,
1926
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1927
from flashinfer.fused_moe.core import (
1928
ActivationType,
vllm/model_executor/layers/quantization/mxfp4.py
+333
@@ -20,10 +20,12 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
20
TRITON_BACKENDS,
21
Mxfp4MoeBackend,
22
convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
23
+ convert_weight_to_mxfp4_moe_kernel_format,
24
make_mxfp4_moe_kernel,
25
make_mxfp4_moe_quant_config,
26
mxfp4_round_up_hidden_size_and_intermediate_size,
27
select_gpt_oss_mxfp4_moe_backend,
28
+ select_mxfp4_moe_backend,
29
)
30
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
31
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -217,6 +219,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
219
)
220
layer.register_parameter("w13_weight_scale", w13_weight_scale)
221
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
222
+ w13_weight_scale.quant_method = "block"
223
224
# down_proj (row parallel)
225
w2_weight = torch.nn.Parameter(
@@ -242,6 +245,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
245
)
246
layer.register_parameter("w2_weight_scale", w2_weight_scale)
247
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
248
+ w2_weight_scale.quant_method = "block"
249
250
if self.moe.has_bias:
251
w13_bias = torch.nn.Parameter(
@@ -397,6 +401,9 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
401
w2_scale=w2_scale,
402
w1_bias=w1_bias,
403
w2_bias=w2_bias,
404
+ gemm1_alpha=1.702,
405
+ gemm1_beta=1.0,
406
+ swiglu_limit=7.0,
407
)
408
409
def select_gemm_impl(
@@ -437,6 +444,332 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
444
layer: FusedMoE,
445
x: torch.Tensor,
446
router_logits: torch.Tensor,
447
+ input_ids: torch.Tensor | None = None,
448
+ ) -> torch.Tensor:
449
+ assert self.is_monolithic
450
+ assert self.moe_kernel is not None
451
+ return self.moe_kernel.apply_monolithic(
452
+ hidden_states=x,
453
+ w1=layer.w13_weight,
454
+ w2=layer.w2_weight,
455
+ router_logits=router_logits,
456
+ activation=layer.activation,
457
+ global_num_experts=layer.global_num_experts,
458
+ expert_map=layer.expert_map,
459
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
460
+ )
461
+
462
+
463
+class Mxfp4MoEMethod(FusedMoEMethodBase):
464
+ """MXFP4 MoE quantization method."""
465
+
466
+ def __init__(self, moe: FusedMoEConfig):
467
+ super().__init__(moe)
468
+ self.weight_dtype = "mxfp4"
469
+ self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
470
+
471
+ self.max_capture_size = (
472
+ get_current_vllm_config().compilation_config.max_cudagraph_capture_size
473
+ )
474
+
475
+ self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
476
+ self.moe_kernel: mk.FusedMoEKernel | None = None
477
+
478
+ # Used for triton kernel precision configs
479
+ self.w13_precision_config = None
480
+ self.w2_precision_config = None
481
+
482
+ @property
483
+ def skip_forward_padding(self) -> bool:
484
+ # SM100_FI_MXFP4_MXFP8_TRTLLM supports padding with mxfp8 quant
485
+ # so can skip the padding in the forward before applying the moe method
486
+ return self.mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8
487
+
488
+ def maybe_roundup_sizes(
489
+ self,
490
+ hidden_size: int,
491
+ intermediate_size_per_partition: int,
492
+ act_dtype: torch.dtype,
493
+ moe_parallel_config: FusedMoEParallelConfig,
494
+ ) -> tuple[int, int]:
495
+ hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
496
+ hidden_size=hidden_size,
497
+ intermediate_size_per_partition=intermediate_size_per_partition,
498
+ act_dtype=act_dtype,
499
+ moe_parallel_config=moe_parallel_config,
500
+ )
501
+ return mxfp4_round_up_hidden_size_and_intermediate_size(
502
+ self.mxfp4_backend, hidden_size, intermediate_size_per_partition
503
+ )
504
+
505
+ def create_weights(
506
+ self,
507
+ layer: torch.nn.Module,
508
+ num_experts: int,
509
+ hidden_size: int,
510
+ intermediate_size_per_partition: int,
511
+ params_dtype: torch.dtype,
512
+ **extra_weight_attrs,
513
+ ):
514
+ self.num_experts = num_experts
515
+ weight_dtype = torch.uint8
516
+ scale_dtype = torch.uint8
517
+ mxfp4_block = 32
518
+
519
+ layer.params_dtype = params_dtype
520
+ layer.num_experts = num_experts
521
+ self.intermediate_size = intermediate_size_per_partition
522
+ self.hidden_size = hidden_size
523
+
524
+ # Fused gate_up_proj (column parallel)
525
+ w13_weight = torch.nn.Parameter(
526
+ torch.zeros(
527
+ num_experts,
528
+ 2 * intermediate_size_per_partition,
529
+ hidden_size // 2,
530
+ dtype=weight_dtype,
531
+ ),
532
+ requires_grad=False,
533
+ )
534
+ layer.register_parameter("w13_weight", w13_weight)
535
+ set_weight_attrs(w13_weight, extra_weight_attrs)
536
+
537
+ w13_weight_scale = torch.nn.Parameter(
538
+ torch.zeros(
539
+ num_experts,
540
+ 2 * intermediate_size_per_partition,
541
+ hidden_size // mxfp4_block,
542
+ dtype=scale_dtype,
543
+ ),
544
+ requires_grad=False,
545
+ )
546
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
547
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
548
+ w13_weight_scale.quant_method = "block"
549
+
550
+ # down_proj (row parallel)
551
+ w2_weight = torch.nn.Parameter(
552
+ torch.zeros(
553
+ num_experts,
554
+ hidden_size,
555
+ intermediate_size_per_partition // 2,
556
+ dtype=weight_dtype,
557
+ ),
558
+ requires_grad=False,
559
+ )
560
+ layer.register_parameter("w2_weight", w2_weight)
561
+ set_weight_attrs(w2_weight, extra_weight_attrs)
562
+
563
+ w2_weight_scale = torch.nn.Parameter(
564
+ torch.zeros(
565
+ num_experts,
566
+ hidden_size,
567
+ intermediate_size_per_partition // mxfp4_block,
568
+ dtype=scale_dtype,
569
+ ),
570
+ requires_grad=False,
571
+ )
572
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
573
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
574
+ w2_weight_scale.quant_method = "block"
575
+
576
+ if self.moe.has_bias:
577
+ w13_bias = torch.nn.Parameter(
578
+ torch.zeros(
579
+ num_experts,
580
+ 2 * intermediate_size_per_partition,
581
+ dtype=torch.bfloat16,
582
+ ),
583
+ requires_grad=False,
584
+ )
585
+ layer.register_parameter("w13_bias", w13_bias)
586
+ set_weight_attrs(w13_bias, extra_weight_attrs)
587
+
588
+ w2_bias = torch.nn.Parameter(
589
+ torch.zeros(
590
+ num_experts,
591
+ hidden_size,
592
+ dtype=torch.bfloat16,
593
+ ),
594
+ requires_grad=False,
595
+ )
596
+ layer.register_parameter("w2_bias", w2_bias)
597
+ set_weight_attrs(w2_bias, extra_weight_attrs)
598
+
599
+ def _setup_kernel(
600
+ self,
601
+ layer: FusedMoE,
602
+ w13: torch.Tensor,
603
+ w2: torch.Tensor,
604
+ w13_scale: torch.Tensor,
605
+ w2_scale: torch.Tensor,
606
+ w13_bias: torch.Tensor | None = None,
607
+ w2_bias: torch.Tensor | None = None,
608
+ ) -> None:
609
+ num_experts = self.num_experts
610
+ intermediate_size = self.intermediate_size
611
+ hidden_size = self.hidden_size
612
+ sf_block_size = 32
613
+
614
+ # Shape assertions
615
+ assert (
616
+ w13.dim() == 3
617
+ and w13.shape[0] == num_experts
618
+ and w13.shape[1] == intermediate_size * 2
619
+ and w13.shape[2] == hidden_size // 2
620
+ )
621
+ assert (
622
+ w13_scale.dim() == 3
623
+ and w13_scale.shape[0] == num_experts
624
+ and w13_scale.shape[1] == intermediate_size * 2
625
+ and w13_scale.shape[2] == hidden_size // sf_block_size
626
+ )
627
+ assert (
628
+ w2.dim() == 3
629
+ and w2.shape[0] == num_experts
630
+ and w2.shape[1] == hidden_size
631
+ and w2.shape[2] == intermediate_size // 2
632
+ )
633
+ assert (
634
+ w2_scale.dim() == 3
635
+ and w2_scale.shape[1] == hidden_size
636
+ and w2_scale.shape[2] == intermediate_size // sf_block_size
637
+ )
638
+ if w13_bias is not None:
639
+ assert (
640
+ w13_bias.dim() == 2
641
+ and w13_bias.shape[0] == num_experts
642
+ and w13_bias.shape[1] == intermediate_size * 2
643
+ )
644
+ if w2_bias is not None:
645
+ assert (
646
+ w2_bias.dim() == 2
647
+ and w2_bias.shape[0] == num_experts
648
+ and w2_bias.shape[1] == hidden_size
649
+ )
650
+
651
+ # Convert weights to kernel format
652
+ w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (
653
+ convert_weight_to_mxfp4_moe_kernel_format(
654
+ mxfp4_backend=self.mxfp4_backend,
655
+ layer=layer,
656
+ w13_weight=w13,
657
+ w2_weight=w2,
658
+ w13_weight_scale=w13_scale,
659
+ w2_weight_scale=w2_scale,
660
+ w13_bias=w13_bias,
661
+ w2_bias=w2_bias,
662
+ _cache_permute_indices=self._cache_permute_indices,
663
+ )
664
+ )
665
+
666
+ # For TRITON backends, weights are wrapped tensors from triton_kernels
667
+ # that don't support .detach(). Manually assign parameters.
668
+ if self.mxfp4_backend not in TRITON_BACKENDS:
669
+ replace_parameter(layer, "w13_weight", w13)
670
+ replace_parameter(layer, "w2_weight", w2)
671
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
672
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
673
+ else:
674
+ layer.w13_weight = w13
675
+ layer.w2_weight = w2
676
+ self.w13_precision_config = w13_scale
677
+ self.w2_precision_config = w2_scale
678
+
679
+ if w13_bias is not None and w2_bias is not None:
680
+ replace_parameter(layer, "w13_bias", w13_bias)
681
+ replace_parameter(layer, "w2_bias", w2_bias)
682
+
683
+ # Build quant config
684
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
685
+
686
+ # Build kernel (modular or monolithic)
687
+ if self.moe_quant_config is not None and self.experts_cls is not None:
688
+ self.moe_kernel = make_mxfp4_moe_kernel(
689
+ moe_quant_config=self.moe_quant_config,
690
+ moe_config=self.moe,
691
+ mxfp4_backend=self.mxfp4_backend,
692
+ experts_cls=self.experts_cls,
693
+ routing_tables=layer._maybe_init_expert_routing_tables(),
694
+ shared_experts=layer.shared_experts,
695
+ )
696
+
697
+ def process_weights_after_loading(self, layer):
698
+ w13 = layer.w13_weight
699
+ w2 = layer.w2_weight
700
+ w13_scale = layer.w13_weight_scale
701
+ w2_scale = layer.w2_weight_scale
702
+ w13_bias = getattr(layer, "w13_bias", None)
703
+ w2_bias = getattr(layer, "w2_bias", None)
704
+
705
+ if self.mxfp4_backend == Mxfp4MoeBackend.NONE:
706
+ return
707
+
708
+ self._setup_kernel(layer, w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
709
+
710
+ def get_fused_moe_quant_config(
711
+ self, layer: torch.nn.Module
712
+ ) -> FusedMoEQuantConfig | None:
713
+ w1_scale = layer.w13_weight_scale
714
+ w2_scale = layer.w2_weight_scale
715
+ w1_bias = getattr(layer, "w13_bias", None)
716
+ w2_bias = getattr(layer, "w2_bias", None)
717
+ swiglu_limit = getattr(layer, "swiglu_limit", None)
718
+
719
+ if self.mxfp4_backend in TRITON_BACKENDS:
720
+ assert self.w13_precision_config is not None
721
+ assert self.w2_precision_config is not None
722
+ w1_scale = self.w13_precision_config
723
+ w2_scale = self.w2_precision_config
724
+
725
+ return make_mxfp4_moe_quant_config(
726
+ mxfp4_backend=self.mxfp4_backend,
727
+ w1_scale=w1_scale,
728
+ w2_scale=w2_scale,
729
+ w1_bias=w1_bias,
730
+ w2_bias=w2_bias,
731
+ swiglu_limit=swiglu_limit,
732
+ )
733
+
734
+ def select_gemm_impl(
735
+ self,
736
+ prepare_finalize: mk.FusedMoEPrepareAndFinalize,
737
+ layer: torch.nn.Module,
738
+ ) -> mk.FusedMoEExpertsModular:
739
+ raise ValueError(
740
+ f"{self.__class__.__name__} uses the new modular kernel "
741
+ "initialization logic. This function should not be called."
742
+ )
743
+
744
+ def apply(
745
+ self,
746
+ layer: FusedMoE,
747
+ x: torch.Tensor,
748
+ topk_weights: torch.Tensor,
749
+ topk_ids: torch.Tensor,
750
+ shared_experts_input: torch.Tensor | None,
751
+ ) -> torch.Tensor:
752
+ assert not self.is_monolithic
753
+ assert self.moe_kernel is not None
754
+ return self.moe_kernel.apply(
755
+ hidden_states=x,
756
+ w1=layer.w13_weight,
757
+ w2=layer.w2_weight,
758
+ topk_weights=topk_weights,
759
+ topk_ids=topk_ids,
760
+ activation=layer.activation,
761
+ global_num_experts=layer.global_num_experts,
762
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
763
+ expert_map=layer.expert_map,
764
+ shared_experts_input=shared_experts_input,
765
+ )
766
+
767
+ def apply_monolithic(
768
+ self,
769
+ layer: FusedMoE,
770
+ x: torch.Tensor,
771
+ router_logits: torch.Tensor,
772
+ input_ids: torch.Tensor | None = None,
773
) -> torch.Tensor:
774
assert self.is_monolithic
775
assert self.moe_kernel is not None
vllm/model_executor/layers/quantization/online/moe_base.py
+1
@@ -130,6 +130,7 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
130
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
131
x: torch.Tensor,
132
router_logits: torch.Tensor,
133
+ input_ids: torch.Tensor | None = None,
134
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
135
assert self.is_monolithic
136
assert self.moe_kernel is not None
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1
@@ -1457,6 +1457,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1457
layer: FusedMoE,
1458
x: torch.Tensor,
1459
router_logits: torch.Tensor,
1460
+ input_ids: torch.Tensor | None = None,
1461
) -> torch.Tensor:
1462
assert self.is_monolithic
1463
assert self.moe_kernel is not None
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+201
-10
@@ -149,6 +149,148 @@ def _per_token_group_quant_fp8(
149
tl.store(y_s_ptr, y_s)
150
151
152
+@triton.jit
153
+def _silu_mul_quant_fp8_packed_kernel(
154
+ input_ptr,
155
+ output_q_ptr,
156
+ output_scale_ptr,
157
+ M,
158
+ input_stride_m,
159
+ output_q_stride_m,
160
+ output_scale_stride_k,
161
+ clamp_limit,
162
+ N: tl.constexpr,
163
+ NUM_GROUPS: tl.constexpr,
164
+ fp8_min: tl.constexpr,
165
+ fp8_max: tl.constexpr,
166
+ GROUP_SIZE: tl.constexpr,
167
+ BLOCK_M: tl.constexpr,
168
+ HAS_CLAMP: tl.constexpr,
169
+):
170
+ N_2: tl.constexpr = N // 2
171
+
172
+ pid_pack = tl.program_id(0)
173
+ pid_m = tl.program_id(1)
174
+ m_offset = pid_m * BLOCK_M
175
+
176
+ if m_offset >= M:
177
+ return
178
+
179
+ offs_m = tl.arange(0, BLOCK_M)
180
+ offs_n = tl.arange(0, GROUP_SIZE)
181
+ row_mask = (m_offset + offs_m) < M
182
+
183
+ base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m
184
+ base_out_offset = (m_offset + offs_m[:, None]) * output_q_stride_m
185
+
186
+ packed_scale = tl.zeros((BLOCK_M,), dtype=tl.int32)
187
+
188
+ for pack_idx in tl.static_range(4):
189
+ group_id = pid_pack * 4 + pack_idx
190
+
191
+ if group_id < NUM_GROUPS:
192
+ n_offset = group_id * GROUP_SIZE
193
+
194
+ act_ptrs = input_ptr + base_row_offset + n_offset + offs_n[None, :]
195
+ act_in = tl.load(act_ptrs, mask=row_mask[:, None], other=0.0)
196
+
197
+ mul_ptrs = act_ptrs + N_2
198
+ mul_in = tl.load(mul_ptrs, mask=row_mask[:, None], other=0.0)
199
+
200
+ act_f32 = act_in.to(tl.float32)
201
+ mul_f32 = mul_in.to(tl.float32)
202
+
203
+ if HAS_CLAMP:
204
+ act_f32 = tl.minimum(act_f32, clamp_limit)
205
+ mul_f32 = tl.clamp(mul_f32, -clamp_limit, clamp_limit)
206
+
207
+ y = (act_f32 / (1.0 + tl.exp(-act_f32))) * mul_f32
208
+ # Round through bf16 to match unfused precision path
209
+ y = y.to(tl.bfloat16).to(tl.float32)
210
+
211
+ absmax = tl.max(tl.abs(y), axis=1)
212
+
213
+ scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
214
+ exponent = tl.ceil(tl.log2(scale_raw))
215
+ scale = tl.math.exp2(exponent)
216
+
217
+ y_q = tl.clamp(y / scale[:, None], fp8_min, fp8_max)
218
+
219
+ out_q_ptrs = output_q_ptr + base_out_offset + n_offset + offs_n[None, :]
220
+ tl.store(
221
+ out_q_ptrs,
222
+ y_q.to(output_q_ptr.dtype.element_ty),
223
+ mask=row_mask[:, None],
224
+ )
225
+
226
+ exponent_biased = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
227
+ packed_scale = packed_scale | (exponent_biased << (pack_idx * 8))
228
+
229
+ scale_ptrs = output_scale_ptr + pid_pack * output_scale_stride_k + m_offset + offs_m
230
+ tl.store(scale_ptrs, packed_scale, mask=row_mask)
231
+
232
+
233
+def silu_mul_quant_fp8_packed_triton(
234
+ input: torch.Tensor,
235
+ group_size: int = 128,
236
+ output_q: torch.Tensor | None = None,
237
+ clamp_limit: float | None = None,
238
+) -> tuple[torch.Tensor, torch.Tensor]:
239
+ assert input.dim() == 2
240
+ assert input.is_contiguous()
241
+
242
+ M, N = input.shape
243
+ N_2 = N // 2
244
+
245
+ assert N_2 % group_size == 0
246
+
247
+ fp8_dtype = torch.float8_e4m3fn
248
+ finfo = torch.finfo(fp8_dtype)
249
+ fp8_min, fp8_max = finfo.min, finfo.max
250
+
251
+ num_groups_per_row = N_2 // group_size
252
+ num_packed_groups = (num_groups_per_row + 3) // 4
253
+ tma_aligned_M = ((M + 3) // 4) * 4
254
+
255
+ if output_q is None:
256
+ output_q = torch.empty((M, N_2), dtype=fp8_dtype, device=input.device)
257
+
258
+ output_scale_packed = torch.zeros(
259
+ (num_packed_groups, tma_aligned_M),
260
+ dtype=torch.int32,
261
+ device=input.device,
262
+ ).T[:M, :]
263
+
264
+ BLOCK_M = 8
265
+ grid = (num_packed_groups, (M + BLOCK_M - 1) // BLOCK_M)
266
+
267
+ num_warps = max(4, group_size // 32)
268
+ num_stages = 2
269
+
270
+ has_clamp = clamp_limit is not None
271
+ _silu_mul_quant_fp8_packed_kernel[grid](
272
+ input,
273
+ output_q,
274
+ output_scale_packed,
275
+ M,
276
+ input.stride(0),
277
+ output_q.stride(0),
278
+ output_scale_packed.stride(1),
279
+ clamp_limit if has_clamp else 0.0,
280
+ N=N,
281
+ NUM_GROUPS=num_groups_per_row,
282
+ fp8_min=fp8_min,
283
+ fp8_max=fp8_max,
284
+ GROUP_SIZE=group_size,
285
+ BLOCK_M=BLOCK_M,
286
+ HAS_CLAMP=has_clamp,
287
+ num_warps=num_warps,
288
+ num_stages=num_stages,
289
+ )
290
+
291
+ return output_q, output_scale_packed
292
+
293
+
294
@triton.jit
295
def _silu_mul_per_token_group_quant_fp8_colmajor(
296
y_ptr, # [M, N]
@@ -823,19 +965,65 @@ def requant_weight_ue8m0_inplace(
965
s_old.copy_(s_requant)
966
967
968
+def _upcast_e8m0_to_fp32(scale: torch.Tensor) -> torch.Tensor:
969
+ """Upcast E8M0 (exponent-only) scale to float32.
970
+
971
+ E8M0 stores only the 8-bit biased exponent (bias=127). To convert
972
+ to float32 we place those 8 bits into the exponent field of an
973
+ IEEE-754 float32 (bits 23-30) with sign=0 and mantissa=0.
974
+ """
975
+ exp_bits = scale.view(torch.uint8).to(torch.int32)
976
+ fp32_bits = exp_bits << 23
977
+ return fp32_bits.view(torch.float32)
978
+
979
+
980
def deepgemm_post_process_fp8_weight_block(
827
- wq: torch.Tensor, ws: torch.Tensor, quant_block_shape: tuple[int], use_e8m0: bool
981
+ wq: torch.Tensor,
982
+ ws: torch.Tensor,
983
+ quant_block_shape: tuple[int, ...],
984
+ use_e8m0: bool,
985
+ is_bmm: bool = False,
986
+ bmm_batch_size: int = 0,
987
) -> tuple[torch.Tensor, torch.Tensor]:
988
assert wq.dtype == torch.float8_e4m3fn, (
989
"Expected quantized tensor dtype "
990
f"to be torch.float8_e4m3fn, got {wq.dtype} instead."
991
)
833
- assert ws.dtype == torch.float32, (
834
- f"Expected tensor scales dtype to be torch.float32, got {ws.dtype} instead"
835
- )
992
837
- if use_e8m0:
838
- requant_weight_ue8m0_inplace(wq, ws, block_size=quant_block_shape)
993
+ if ws.dtype == torch.float8_e8m0fnu:
994
+ # Scales already in E8M0 from checkpoint — upcast to fp32
995
+ # and skip requantization (weights already have power-of-two scales).
996
+ ws = _upcast_e8m0_to_fp32(ws)
997
+ else:
998
+ assert ws.dtype == torch.float32, (
999
+ f"Expected tensor scales dtype to be torch.float32 or "
1000
+ f"torch.float8_e8m0fnu, got {ws.dtype} instead"
1001
+ )
1002
+ if use_e8m0:
1003
+ requant_weight_ue8m0_inplace(wq, ws, block_size=quant_block_shape)
1004
+
1005
+ if is_bmm:
1006
+ # Reshape 2D weight/scale to 3D for grouped BMM (einsum):
1007
+ # wq: (g*r, d) -> (g, r, d)
1008
+ # ws: (g*r/128, d/128) -> (g, r/128, d/128)
1009
+ g = bmm_batch_size
1010
+ assert wq.ndim == 2 and ws.ndim == 2
1011
+ d = wq.size(1)
1012
+ r = wq.size(0) // g
1013
+ wq = wq.view(g, r, d)
1014
+ ws = ws.view(g, r // quant_block_shape[0], d // quant_block_shape[1])
1015
+ # Pre-transform scale with recipe=(1, 128, 128) to broadcast + pack
1016
+ # into TMA-aligned UE8M0 (INT32) layout. At runtime fp8_einsum uses
1017
+ # recipe=(1, 1, 128) which sees INT dtype and skips re-transform.
1018
+ dg_ws = transform_sf_into_required_layout(
1019
+ sf=ws,
1020
+ mn=r,
1021
+ k=d,
1022
+ recipe=(1, quant_block_shape[0], quant_block_shape[1]),
1023
+ num_groups=g,
1024
+ is_sfa=False,
1025
+ )
1026
+ return wq, dg_ws
1027
1028
original_ndim = wq.ndim
1029
if wq.ndim == 2:
@@ -984,11 +1172,13 @@ def create_fp8_scale_parameter(
1172
input_size_per_partition: int,
1173
block_size: list[int] | None,
1174
weight_loader: Callable | None,
1175
+ scale_dtype: torch.dtype | None = None,
1176
) -> torch.nn.Parameter:
1177
"""Create scale parameter based on quantization strategy."""
1178
+ dtype = scale_dtype if scale_dtype is not None else torch.float32
1179
if parameter_type == ChannelQuantScaleParameter:
1180
scale = parameter_type(
991
- data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
1181
+ data=torch.empty((sum(output_partition_sizes), 1), dtype=dtype),
1182
output_dim=0,
1183
weight_loader=weight_loader,
1184
)
@@ -1000,7 +1190,7 @@ def create_fp8_scale_parameter(
1190
data=torch.empty(
1191
(output_size_per_partition + block_n - 1) // block_n,
1192
(input_size_per_partition + block_k - 1) // block_k,
1003
- dtype=torch.float32,
1193
+ dtype=dtype,
1194
),
1195
input_dim=1,
1196
output_dim=0,
@@ -1008,13 +1198,14 @@ def create_fp8_scale_parameter(
1198
)
1199
elif parameter_type == PerTensorScaleParameter:
1200
scale = parameter_type(
1011
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1201
+ data=torch.empty(len(output_partition_sizes), dtype=dtype),
1202
weight_loader=weight_loader,
1203
)
1204
else:
1205
raise ValueError(f"Unknown parameter type: {parameter_type}")
1206
1017
- scale[:] = torch.finfo(torch.float32).min
1207
+ if dtype == torch.float32:
1208
+ scale[:] = torch.finfo(torch.float32).min
1209
set_weight_attrs(scale, {"scale_type": "weight_scale"})
1210
return scale
1211
vllm/model_executor/layers/rotary_embedding/__init__.py
+16
-7
@@ -7,7 +7,10 @@ from typing import Any
7
import torch
8
9
from .base import RotaryEmbedding
10
-from .deepseek_scaling_rope import DeepseekScalingRotaryEmbedding
10
+from .deepseek_scaling_rope import (
11
+ DeepseekScalingRotaryEmbedding,
12
+ DeepseekV4ScalingRotaryEmbedding,
13
+)
14
from .dual_chunk_rope import DualChunkRotaryEmbedding
15
from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding
16
from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding
@@ -60,11 +63,13 @@ def get_rope(
63
rope_parameters = rope_parameters or {}
64
base = rope_parameters.get("rope_theta", 10000)
65
scaling_type = rope_parameters.get("rope_type", "default")
63
- partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0)
64
-
65
- if partial_rotary_factor <= 0.0 or partial_rotary_factor > 1.0:
66
- raise ValueError(f"{partial_rotary_factor=} must be between 0.0 and 1.0")
67
- rotary_dim = int(head_size * partial_rotary_factor)
66
+ if rotary_dim := rope_parameters.get("rope_dim", None):
67
+ pass
68
+ else:
69
+ partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0)
70
+ if partial_rotary_factor <= 0.0 or partial_rotary_factor > 1.0:
71
+ raise ValueError(f"{partial_rotary_factor=} must be between 0.0 and 1.0")
72
+ rotary_dim = int(head_size * partial_rotary_factor)
73
74
key = (
75
head_size,
@@ -289,7 +294,11 @@ def get_rope(
294
"mscale_all_dim",
295
)
296
}
292
- rotary_emb = DeepseekScalingRotaryEmbedding(
297
+ if rope_parameters.get("is_deepseek_v4", False):
298
+ cls = DeepseekV4ScalingRotaryEmbedding
299
+ else:
300
+ cls = DeepseekScalingRotaryEmbedding
301
+ rotary_emb = cls(
302
head_size,
303
rotary_dim,
304
original_max_position,
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+117
-4
@@ -132,10 +132,8 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
132
]
133
cos, sin = cos_sin.chunk(2, dim=-1)
134
if self.is_neox_style:
135
- # NOTE(woosuk): Here we assume that the positions tensor has the
136
- # shape [batch_size, seq_len].
137
- cos = cos.repeat(1, 1, 2).unsqueeze(-2)
138
- sin = sin.repeat(1, 1, 2).unsqueeze(-2)
135
+ cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2)
136
+ sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2)
137
else:
138
cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
139
sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
@@ -197,3 +195,118 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
195
return query, key
196
else:
197
return self.forward_native(positions, query, key, offsets)
198
+
199
+
200
+class DeepseekV4ScalingRotaryEmbedding(DeepseekScalingRotaryEmbedding):
201
+ """RotaryEmbedding extended with YaRN method.
202
+
203
+ Credits to Peng et al. github.com/jquesnelle/yarn
204
+
205
+ Compared to DeepseekScalingRotaryEmbedding:
206
+ - Applies RoPE to the last rotary_dim
207
+ - The forward method requires an inverse parameter to indicate
208
+ whether to negate the sin
209
+ - Supports applying RoPE to query only (without key)
210
+ - cos_sin_cache stored as fp32 for higher precision RoPE
211
+ """
212
+
213
+ def __init__(self, *args, **kwargs):
214
+ super().__init__(*args, **kwargs)
215
+ cache_fp32 = self._compute_cos_sin_cache()
216
+ self.register_buffer("cos_sin_cache", cache_fp32, persistent=False)
217
+
218
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
219
+ inv_freq = self._compute_inv_freq(self.scaling_factor)
220
+ t = torch.arange(
221
+ self.max_position_embeddings * self.scaling_factor,
222
+ device=current_platform.device_type,
223
+ dtype=torch.float32,
224
+ )
225
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
226
+ cos = freqs.cos() * self.mscale
227
+ sin = freqs.sin() * self.mscale
228
+ cache = torch.cat((cos, sin), dim=-1)
229
+ return cache
230
+
231
+ def forward_native(
232
+ self,
233
+ positions: torch.Tensor,
234
+ query: torch.Tensor,
235
+ key: torch.Tensor | None = None,
236
+ offsets: torch.Tensor | None = None,
237
+ inverse: bool = False,
238
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
239
+ """PyTorch-native implementation equivalent to forward()."""
240
+
241
+ head_size = query.size(-1)
242
+ query_rot = query[..., -self.rotary_dim :]
243
+ key_rot = key[..., -self.rotary_dim :] if key is not None else None
244
+
245
+ if self.rotary_dim < head_size:
246
+ query_pass = query[..., : -self.rotary_dim]
247
+ key_pass = key[..., : -self.rotary_dim] if key is not None else None
248
+
249
+ cos_sin = self.cos_sin_cache[
250
+ torch.add(positions, offsets) if offsets is not None else positions
251
+ ]
252
+ cos, sin = cos_sin.chunk(2, dim=-1)
253
+ if self.is_neox_style:
254
+ cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2)
255
+ sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2)
256
+ else:
257
+ cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
258
+ sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
259
+ if inverse:
260
+ sin = -sin
261
+ rotate_fn = rotate_neox if self.is_neox_style else rotate_gptj
262
+ orig_dtype = query.dtype
263
+ query_rot = (query_rot * cos + rotate_fn(query_rot) * sin).to(orig_dtype)
264
+ if key_rot is not None:
265
+ key_rot = (key_rot * cos + rotate_fn(key_rot) * sin).to(orig_dtype)
266
+
267
+ if self.rotary_dim < head_size:
268
+ query = torch.cat((query_pass, query_rot), dim=-1)
269
+ key = torch.cat((key_pass, key_rot), dim=-1) if key is not None else None
270
+ else:
271
+ query = query_rot
272
+ key = key_rot
273
+
274
+ return query, key
275
+
276
+ def forward_hip(
277
+ self,
278
+ positions: torch.Tensor,
279
+ query: torch.Tensor,
280
+ key: torch.Tensor | None = None,
281
+ offsets: torch.Tensor | None = None,
282
+ inverse: bool = False,
283
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
284
+ return self.forward_native(positions, query, key, offsets)
285
+
286
+ def forward_cuda(
287
+ self,
288
+ positions: torch.Tensor,
289
+ query: torch.Tensor,
290
+ key: torch.Tensor | None = None,
291
+ offsets: torch.Tensor | None = None,
292
+ inverse: bool = False,
293
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
294
+ from vllm import _custom_ops as ops
295
+
296
+ # The indexer and attention have different head_dim,
297
+ # we obtain the corresponding head_dim via the query.
298
+ head_size = query.size(-1)
299
+ rope_dim_offset = head_size - self.rotary_dim
300
+ # ops.rotary_embedding() is an in-place operation
301
+ # that updates the query and key tensors.
302
+ ops.rotary_embedding(
303
+ torch.add(positions, offsets) if offsets is not None else positions,
304
+ query,
305
+ key,
306
+ head_size,
307
+ self.cos_sin_cache,
308
+ self.is_neox_style,
309
+ rope_dim_offset=rope_dim_offset,
310
+ inverse=inverse,
311
+ )
312
+ return query, key
vllm/model_executor/layers/sparse_attn_indexer.py
+185
-50
@@ -10,7 +10,11 @@ from vllm.forward_context import get_forward_context
10
from vllm.logger import init_logger
11
from vllm.model_executor.custom_op import CustomOp
12
from vllm.platforms import current_platform
13
-from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits, has_deep_gemm
13
+from vllm.utils.deep_gemm import (
14
+ fp8_fp4_mqa_logits,
15
+ fp8_fp4_paged_mqa_logits,
16
+ has_deep_gemm,
17
+)
18
from vllm.utils.torch_utils import (
19
LayerNameType,
20
_encode_layer_name,
@@ -32,12 +36,57 @@ logger = init_logger(__name__)
36
37
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
38
39
+# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
40
+MXFP4_BLOCK_SIZE = 32
41
+
42
+
43
+def _gather_workspace_shapes(
44
+ total_seq_lens: int,
45
+ head_dim: int,
46
+ fp8_dtype: torch.dtype,
47
+ use_fp4_cache: bool,
48
+) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]:
49
+ """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for
50
+ the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32
51
+ scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 +
52
+ (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales."""
53
+ if use_fp4_cache:
54
+ return (
55
+ ((total_seq_lens, head_dim // 2), torch.uint8),
56
+ ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8),
57
+ )
58
+ return (
59
+ ((total_seq_lens, head_dim), fp8_dtype),
60
+ ((total_seq_lens, 4), torch.uint8),
61
+ )
62
+
63
+
64
+def kv_cache_as_quant_view(
65
+ kv_cache: torch.Tensor,
66
+ head_dim: int,
67
+ use_fp4_cache: bool,
68
+) -> torch.Tensor:
69
+ """4D ``[num_blocks, block_size, 1, head_width]`` view expected by
70
+ DeepGEMM, from the 3D indexer kv-cache allocation."""
71
+ if use_fp4_cache:
72
+ assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8
73
+ num_blocks, block_size, _ = kv_cache.shape
74
+ page_bytes = int(kv_cache.stride(0))
75
+ fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
76
+ return torch.as_strided(
77
+ kv_cache,
78
+ size=(num_blocks, block_size, 1, fp4_bytes),
79
+ stride=(page_bytes, fp4_bytes, fp4_bytes, 1),
80
+ )
81
+ return kv_cache.unsqueeze(-2)
82
+
83
84
def sparse_attn_indexer(
85
hidden_states: torch.Tensor,
86
k_cache_prefix: LayerNameType,
87
kv_cache: torch.Tensor,
40
- q_fp8: torch.Tensor,
88
+ q_quant: torch.Tensor,
89
+ q_scale: torch.Tensor | None,
90
k: torch.Tensor,
91
weights: torch.Tensor,
92
quant_block_size: int,
@@ -47,6 +96,8 @@ def sparse_attn_indexer(
96
max_model_len: int,
97
total_seq_lens: int,
98
topk_indices_buffer: torch.Tensor,
99
+ skip_k_cache_insert: bool,
100
+ use_fp4_cache: bool = False,
101
) -> torch.Tensor:
102
# careful! this will be None in dummy run
103
attn_metadata = get_forward_context().attn_metadata
@@ -56,9 +107,12 @@ def sparse_attn_indexer(
107
# assert isinstance(attn_metadata, dict)
108
if not isinstance(attn_metadata, dict):
109
# Reserve workspace for indexer during profiling run
110
+ values_spec, scales_spec = _gather_workspace_shapes(
111
+ total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
112
+ )
113
current_workspace_manager().get_simultaneous(
60
- ((total_seq_lens, head_dim), torch.float8_e4m3fn),
61
- ((total_seq_lens, 4), torch.uint8),
114
+ values_spec,
115
+ scales_spec,
116
((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
117
)
118
@@ -73,7 +127,8 @@ def sparse_attn_indexer(
127
hidden_states,
128
k_cache_prefix,
129
kv_cache,
76
- q_fp8,
130
+ q_quant,
131
+ q_scale,
132
k,
133
weights,
134
quant_block_size,
@@ -83,6 +138,8 @@ def sparse_attn_indexer(
138
max_model_len,
139
total_seq_lens,
140
topk_indices_buffer,
141
+ skip_k_cache_insert,
142
+ use_fp4_cache,
143
)
144
attn_metadata_narrowed = attn_metadata[k_cache_prefix]
145
assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
@@ -91,49 +148,81 @@ def sparse_attn_indexer(
148
has_prefill = attn_metadata_narrowed.num_prefills > 0
149
num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
150
151
+ # q_scale is required iff the FP4 cache path is enabled; the FP8 path
152
+ # folds the Q scale into `weights` inside fused_indexer_q_rope_quant.
153
+ if use_fp4_cache:
154
+ assert q_scale is not None, "use_fp4_cache=True requires q_scale"
155
+ else:
156
+ assert q_scale is None, "q_scale must be None when use_fp4_cache=False"
157
+
158
# During speculative decoding, k may be padded to the CUDA graph batch
159
# size while slot_mapping only covers actual tokens. Truncate k to avoid
160
# out-of-bounds reads in the kernel.
161
num_tokens = slot_mapping.shape[0]
98
- k = k[:num_tokens]
99
-
100
- # scale_fmt can be None, but the function expects str
101
- assert scale_fmt is not None
102
- ops.indexer_k_quant_and_cache(
103
- k,
104
- kv_cache,
105
- slot_mapping,
106
- quant_block_size,
107
- scale_fmt,
108
- )
162
+ if k is not None:
163
+ k = k[:num_tokens]
164
+
165
+ if not skip_k_cache_insert:
166
+ # scale_fmt can be None, but the function expects str
167
+ assert scale_fmt is not None
168
+ assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
169
+ ops.indexer_k_quant_and_cache(
170
+ k,
171
+ kv_cache,
172
+ slot_mapping,
173
+ quant_block_size,
174
+ scale_fmt,
175
+ )
176
177
topk_indices_buffer[: hidden_states.shape[0]] = -1
178
if has_prefill:
179
prefill_metadata = attn_metadata_narrowed.prefill
180
assert prefill_metadata is not None
181
115
- # Get the full shared workspace buffers once (will allocate on first use)
182
+ # Get the full shared workspace buffers once (will allocate on first use).
183
+ # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and
184
+ # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0
185
+ # scales) based on use_fp4_cache.
186
workspace_manager = current_workspace_manager()
117
- k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
118
- ((total_seq_lens, head_dim), fp8_dtype),
119
- ((total_seq_lens, 4), torch.uint8),
187
+ values_spec, scales_spec = _gather_workspace_shapes(
188
+ total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
189
+ )
190
+ k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
191
+ values_spec,
192
+ scales_spec,
193
)
194
for chunk in prefill_metadata.chunks:
122
- k_fp8 = k_fp8_full[: chunk.total_seq_lens]
195
+ k_quant = k_quant_full[: chunk.total_seq_lens]
196
k_scale = k_scale_full[: chunk.total_seq_lens]
197
198
if not chunk.skip_kv_gather:
199
ops.cp_gather_indexer_k_quant_cache(
200
kv_cache,
128
- k_fp8,
201
+ k_quant,
202
k_scale,
203
chunk.block_table,
204
chunk.cu_seq_lens,
205
)
206
134
- logits = fp8_mqa_logits(
135
- q_fp8[chunk.token_start : chunk.token_end],
136
- (k_fp8, k_scale.view(torch.float32).flatten()),
207
+ q_slice = q_quant[chunk.token_start : chunk.token_end]
208
+ q_scale_slice = (
209
+ q_scale[chunk.token_start : chunk.token_end]
210
+ if q_scale is not None
211
+ else None
212
+ )
213
+ # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
214
+ # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
215
+ if use_fp4_cache:
216
+ q_slice_cast = q_slice.view(torch.int8)
217
+ k_quant_cast = k_quant.view(torch.int8)
218
+ k_scale_cast = k_scale.view(torch.int32).squeeze(-1)
219
+ else:
220
+ q_slice_cast = q_slice
221
+ k_quant_cast = k_quant
222
+ k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
223
+ logits = fp8_fp4_mqa_logits(
224
+ (q_slice_cast, q_scale_slice),
225
+ (k_quant_cast, k_scale_cast),
226
weights[chunk.token_start : chunk.token_end],
227
chunk.cu_seqlen_ks,
228
chunk.cu_seqlen_ke,
@@ -171,32 +260,55 @@ def sparse_attn_indexer(
260
if has_decode:
261
decode_metadata = attn_metadata_narrowed.decode
262
assert decode_metadata is not None
174
- # kv_cache shape [
175
- # kv_cache size requirement [num_block, block_size, n_head, head_dim],
176
- # we only have [num_block, block_size, head_dim],
177
- kv_cache = kv_cache.unsqueeze(-2)
263
+ kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
264
decode_lens = decode_metadata.decode_lens
265
if decode_metadata.requires_padding:
266
# pad in edge case where we have short chunked prefill length <
267
# decode_threshold since we unstrictly split
268
# prefill and decode by decode_threshold
183
- # (currently set to 1 + speculative tokens)
184
- padded_q_fp8_decode_tokens = pack_seq_triton(
185
- q_fp8[:num_decode_tokens], decode_lens
186
- )
269
+ # (currently set to 1 + speculative tokens).
270
+ # FP8 Q is float8_e4m3fn (pack_seq_triton's fp32 pad path is OK —
271
+ # downstream context_lens masks stale slots). MXFP4 Q is two
272
+ # uint8 tensors (values + ue8m0 scales) — use the dedicated uint8
273
+ # packer with pad_byte=0 so padded slots dequantize to 0 and
274
+ # can't produce NaN/Inf in the logits kernel.
275
+ if q_scale is not None:
276
+ padded_q_quant_decode_tokens = pack_seq_triton(
277
+ q_quant[:num_decode_tokens], decode_lens, pad_value=0
278
+ )
279
+ padded_q_scale = pack_seq_triton(
280
+ q_scale[:num_decode_tokens], decode_lens, pad_value=0
281
+ )
282
+ else:
283
+ padded_q_quant_decode_tokens = pack_seq_triton(
284
+ q_quant[:num_decode_tokens], decode_lens
285
+ )
286
+ padded_q_scale = None
287
else:
188
- padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
189
- decode_lens.shape[0], -1, *q_fp8.shape[1:]
288
+ padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape(
289
+ decode_lens.shape[0], -1, *q_quant.shape[1:]
290
)
291
+ if q_scale is not None:
292
+ padded_q_scale = q_scale[:num_decode_tokens].reshape(
293
+ decode_lens.shape[0], -1, *q_scale.shape[1:]
294
+ )
295
+ else:
296
+ padded_q_scale = None
297
# TODO: move and optimize below logic with triton kernels
192
- batch_size = padded_q_fp8_decode_tokens.shape[0]
193
- next_n = padded_q_fp8_decode_tokens.shape[1]
298
+ batch_size = padded_q_quant_decode_tokens.shape[0]
299
+ next_n = padded_q_quant_decode_tokens.shape[1]
300
num_padded_tokens = batch_size * next_n
301
seq_lens = decode_metadata.seq_lens[:batch_size]
196
- # seq_lens is (B, next_n) for native spec decode, (B,) otherwise.
197
- # fp8_paged_mqa_logits and all topk kernels accept both shapes.
198
- logits = fp8_paged_mqa_logits(
199
- padded_q_fp8_decode_tokens,
302
+ # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1)
303
+ # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens;
304
+ # the downstream topk kernels accept both 1D and 2D.
305
+ padded_q_quant_cast = (
306
+ padded_q_quant_decode_tokens.view(torch.int8)
307
+ if use_fp4_cache
308
+ else padded_q_quant_decode_tokens
309
+ )
310
+ logits = fp8_fp4_paged_mqa_logits(
311
+ (padded_q_quant_cast, padded_q_scale),
312
kv_cache,
313
weights[:num_padded_tokens],
314
seq_lens,
@@ -208,7 +320,7 @@ def sparse_attn_indexer(
320
num_rows = logits.shape[0]
321
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
322
211
- if current_platform.is_cuda():
323
+ if current_platform.is_cuda() and topk_tokens in (512, 1024, 2048):
324
workspace_manager = current_workspace_manager()
325
(topk_workspace,) = workspace_manager.get_simultaneous(
326
((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
@@ -263,7 +375,8 @@ def sparse_attn_indexer_fake(
375
hidden_states: torch.Tensor,
376
k_cache_prefix: LayerNameType,
377
kv_cache: torch.Tensor,
266
- q_fp8: torch.Tensor,
378
+ q_quant: torch.Tensor,
379
+ q_scale: torch.Tensor | None,
380
k: torch.Tensor,
381
weights: torch.Tensor,
382
quant_block_size: int,
@@ -273,6 +386,8 @@ def sparse_attn_indexer_fake(
386
max_model_len: int,
387
total_seq_lens: int,
388
topk_indices_buffer: torch.Tensor | None,
389
+ skip_k_cache_insert: bool,
390
+ use_fp4_cache: bool = False,
391
) -> torch.Tensor:
392
return topk_indices_buffer
393
@@ -309,6 +424,8 @@ class SparseAttnIndexer(CustomOp):
424
max_model_len: int,
425
max_total_seq_len: int,
426
topk_indices_buffer: torch.Tensor,
427
+ skip_k_cache_insert: bool = False,
428
+ use_fp4_cache: bool = False,
429
):
430
super().__init__()
431
self.k_cache = k_cache
@@ -319,6 +436,8 @@ class SparseAttnIndexer(CustomOp):
436
self.max_model_len = max_model_len
437
self.max_total_seq_len = max_total_seq_len
438
self.topk_indices_buffer = topk_indices_buffer
439
+ self.skip_k_cache_insert = skip_k_cache_insert
440
+ self.use_fp4_cache = use_fp4_cache
441
if current_platform.is_cuda() and not has_deep_gemm():
442
raise RuntimeError(
443
"Sparse Attention Indexer CUDA op requires DeepGEMM to be installed."
@@ -327,14 +446,14 @@ class SparseAttnIndexer(CustomOp):
446
def forward_native(
447
self,
448
hidden_states: torch.Tensor,
330
- q_fp8: torch.Tensor,
449
+ q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
450
k: torch.Tensor,
451
weights: torch.Tensor,
452
):
453
if current_platform.is_cuda() or current_platform.is_xpu():
335
- return self.forward_cuda(hidden_states, q_fp8, k, weights)
454
+ return self.forward_cuda(hidden_states, q_quant, k, weights)
455
elif current_platform.is_rocm():
337
- return self.forward_hip(hidden_states, q_fp8, k, weights)
456
+ return self.forward_hip(hidden_states, q_quant, k, weights)
457
else:
458
raise NotImplementedError(
459
"SparseAttnIndexer native forward is only implemented for "
@@ -344,15 +463,22 @@ class SparseAttnIndexer(CustomOp):
463
def forward_cuda(
464
self,
465
hidden_states: torch.Tensor,
347
- q_fp8: torch.Tensor,
466
+ q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
467
k: torch.Tensor,
468
weights: torch.Tensor,
469
):
470
+ # FP8 path: single tensor (per-token scale is folded into `weights`).
471
+ # FP4 path: (values, scales) tuple with scales required by the kernel.
472
+ if isinstance(q_quant, tuple):
473
+ q_values, q_scale = q_quant
474
+ else:
475
+ q_values, q_scale = q_quant, None
476
return torch.ops.vllm.sparse_attn_indexer(
477
hidden_states,
478
_encode_layer_name(self.k_cache.prefix),
479
self.k_cache.kv_cache,
355
- q_fp8,
480
+ q_values,
481
+ q_scale,
482
k,
483
weights,
484
self.quant_block_size,
@@ -362,21 +488,30 @@ class SparseAttnIndexer(CustomOp):
488
self.max_model_len,
489
self.max_total_seq_len,
490
self.topk_indices_buffer,
491
+ self.skip_k_cache_insert,
492
+ self.use_fp4_cache,
493
)
494
495
def forward_hip(
496
self,
497
hidden_states: torch.Tensor,
370
- q_fp8: torch.Tensor,
498
+ q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
499
k: torch.Tensor,
500
weights: torch.Tensor,
501
):
502
+ assert not self.skip_k_cache_insert, (
503
+ "AMD platform doesn't support skip cache insert yet"
504
+ )
505
+ assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
506
+ assert isinstance(q_quant, torch.Tensor), (
507
+ "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
508
+ )
509
if rocm_aiter_ops.is_enabled():
510
return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
511
hidden_states,
512
_encode_layer_name(self.k_cache.prefix),
513
self.k_cache.kv_cache,
379
- q_fp8,
514
+ q_quant,
515
k,
516
weights,
517
self.quant_block_size,
vllm/model_executor/layers/utils.py
+7
@@ -299,6 +299,13 @@ def cpu_unquantized_gemm(
299
return layer.cpu_linear(x, weight, bias)
300
301
302
+def cublas_gemm_bf16_bf16_fp32(
303
+ x: torch.Tensor,
304
+ weight: torch.Tensor,
305
+):
306
+ return ops.router_gemm_bf16_fp32(x, weight)
307
+
308
+
309
def dispatch_unquantized_gemm() -> Callable[..., torch.Tensor]:
310
if current_platform.is_rocm():
311
return rocm_unquantized_gemm
vllm/model_executor/models/config.py
+26
@@ -107,6 +107,31 @@ class Gemma4Config(VerifyAndUpdateConfig):
107
)
108
109
110
+class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig):
111
+ @staticmethod
112
+ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
113
+ quant_config = getattr(model_config.hf_config, "quantization_config", None)
114
+ if quant_config is not None and quant_config.get("quant_method") == "fp8":
115
+ model_type = getattr(model_config.hf_config, "model_type", None)
116
+ if model_type == "deepseek_v4":
117
+ model_config.hf_config.quantization_config["quant_method"] = (
118
+ "deepseek_v4_fp8"
119
+ )
120
+
121
+ hf_text_quant_config = getattr(
122
+ model_config.hf_text_config, "quantization_config", None
123
+ )
124
+ if (
125
+ hf_text_quant_config is not None
126
+ and hf_text_quant_config.get("quant_method") == "fp8"
127
+ ):
128
+ model_type = getattr(model_config.hf_text_config, "model_type", None)
129
+ if model_type == "deepseek_v4":
130
+ model_config.hf_text_config.quantization_config["quant_method"] = (
131
+ "deepseek_v4_fp8"
132
+ )
133
+
134
+
135
class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
136
@staticmethod
137
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -635,6 +660,7 @@ class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
660
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
661
"ColBERTJinaRobertaModel": JinaRobertaModelConfig,
662
"ColQwen3_5": Qwen3_5ForConditionalGenerationConfig,
663
+ "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,
664
"DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
665
"Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501
666
"FalconMambaForCausalLM": MambaModelConfig,
vllm/model_executor/models/deepseek_v4.py
new
+1437
@@ -0,0 +1,1437 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import typing
4
+from collections.abc import Callable, Iterable
5
+from itertools import islice
6
+
7
+import regex as re
8
+import torch
9
+import torch.nn as nn
10
+import torch.nn.functional as F
11
+
12
+from vllm.compilation.decorators import support_torch_compile
13
+from vllm.config import VllmConfig
14
+from vllm.distributed import (
15
+ get_ep_group,
16
+ get_tensor_model_parallel_rank,
17
+ get_tensor_model_parallel_world_size,
18
+)
19
+from vllm.forward_context import get_forward_context
20
+from vllm.model_executor.layers.deepseek_v4_attention import (
21
+ DeepseekV4Indexer,
22
+ DeepseekV4MLAModules,
23
+ DeepseekV4MultiHeadLatentAttentionWrapper,
24
+)
25
+from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
26
+from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
27
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
28
+ fused_topk_bias,
29
+)
30
+from vllm.model_executor.layers.layernorm import RMSNorm
31
+from vllm.model_executor.layers.linear import (
32
+ ColumnParallelLinear,
33
+ MergedColumnParallelLinear,
34
+ RowParallelLinear,
35
+)
36
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
+from vllm.model_executor.layers.quantization import QuantizationMethods
38
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
39
+from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
40
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
41
+ is_layer_skipped,
42
+)
43
+from vllm.model_executor.layers.rotary_embedding import get_rope
44
+from vllm.model_executor.layers.vocab_parallel_embedding import (
45
+ ParallelLMHead,
46
+ VocabParallelEmbedding,
47
+)
48
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
49
+from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLP
50
+from vllm.model_executor.utils import set_weight_attrs
51
+from vllm.platforms import current_platform
52
+from vllm.sequence import IntermediateTensors
53
+from vllm.triton_utils import tl, triton
54
+from vllm.utils.multi_stream_utils import AuxStreamType
55
+from vllm.utils.torch_utils import direct_register_custom_op
56
+
57
+from .utils import (
58
+ AutoWeightsLoader,
59
+ WeightsMapper,
60
+ extract_layer_index,
61
+ make_layers,
62
+ maybe_prefix,
63
+)
64
+
65
+
66
+class DeepseekV4FP8Config(Fp8Config):
67
+ """FP8 config that routes MoE layers to MXFP4 quantization.
68
+
69
+ DeepSeek V4 checkpoints use FP8 for linear/attention layers but
70
+ MXFP4 for MoE expert weights. This config inherits standard FP8
71
+ behavior and overrides only the MoE dispatch.
72
+ """
73
+
74
+ def __init__(self, *args, **kwargs):
75
+ super().__init__(*args, **kwargs)
76
+ self.is_scale_e8m0: bool = True
77
+
78
+ @classmethod
79
+ def get_name(cls) -> QuantizationMethods:
80
+ return "deepseek_v4_fp8"
81
+
82
+ @classmethod
83
+ def override_quantization_method(
84
+ cls, hf_quant_cfg, user_quant, hf_config=None
85
+ ) -> QuantizationMethods | None:
86
+ if not (
87
+ isinstance(hf_quant_cfg, dict)
88
+ and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
89
+ ):
90
+ return None
91
+ model_type = getattr(hf_config, "model_type", None)
92
+ if model_type == "deepseek_v4" or user_quant == "deepseek_v4_fp8":
93
+ return "deepseek_v4_fp8"
94
+ return None
95
+
96
+ def get_quant_method(self, layer, prefix):
97
+ if isinstance(layer, FusedMoE):
98
+ if is_layer_skipped(
99
+ prefix=prefix,
100
+ ignored_layers=self.ignored_layers,
101
+ fused_mapping=self.packed_modules_mapping,
102
+ ):
103
+ return UnquantizedFusedMoEMethod(layer.moe_config)
104
+ return Mxfp4MoEMethod(layer.moe_config)
105
+ return super().get_quant_method(layer, prefix)
106
+
107
+ def is_mxfp4_quant(self, prefix, layer):
108
+ return isinstance(layer, FusedMoE)
109
+
110
+
111
+@triton.jit
112
+def _deepseek_v4_stage_mega_moe_inputs_kernel(
113
+ hidden_states,
114
+ x_fp8,
115
+ x_sf,
116
+ topk_ids,
117
+ topk_weights,
118
+ topk_idx_out,
119
+ topk_weights_out,
120
+ hidden_stride_m: tl.constexpr,
121
+ hidden_stride_k: tl.constexpr,
122
+ x_stride_m: tl.constexpr,
123
+ x_stride_k: tl.constexpr,
124
+ x_sf_stride_m: tl.constexpr,
125
+ x_sf_stride_k: tl.constexpr,
126
+ topk_ids_stride_m: tl.constexpr,
127
+ topk_ids_stride_k: tl.constexpr,
128
+ topk_weights_stride_m: tl.constexpr,
129
+ topk_weights_stride_k: tl.constexpr,
130
+ topk_idx_stride_m: tl.constexpr,
131
+ topk_idx_stride_k: tl.constexpr,
132
+ topk_weights_out_stride_m: tl.constexpr,
133
+ topk_weights_out_stride_k: tl.constexpr,
134
+ hidden_size: tl.constexpr,
135
+ top_k: tl.constexpr,
136
+ BLOCK_K: tl.constexpr,
137
+ GROUP_K: tl.constexpr,
138
+ BLOCK_TOPK: tl.constexpr,
139
+) -> None:
140
+ token_id = tl.program_id(0)
141
+ k_block_id = tl.program_id(1)
142
+
143
+ k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
144
+ k_mask = k_offsets < hidden_size
145
+ hidden = tl.load(
146
+ hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
147
+ mask=k_mask,
148
+ other=0.0,
149
+ ).to(tl.float32)
150
+
151
+ num_groups: tl.constexpr = BLOCK_K // GROUP_K
152
+ hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
153
+ amax = tl.max(hidden_groups, axis=1)
154
+ amax = tl.maximum(amax, 1.0e-4)
155
+
156
+ scale = amax / 448.0
157
+ scale_bits = scale.to(tl.uint32, bitcast=True)
158
+ scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
159
+ tl.uint32
160
+ )
161
+ scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
162
+ rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
163
+
164
+ hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
165
+ scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
166
+ scaled = tl.reshape(scaled, [BLOCK_K])
167
+ fp8 = scaled.to(tl.float8e4nv)
168
+ tl.store(
169
+ x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
170
+ fp8,
171
+ mask=k_mask,
172
+ )
173
+
174
+ scale_offsets = tl.arange(0, num_groups)
175
+ packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
176
+ tl.store(
177
+ x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
178
+ packed_scale,
179
+ )
180
+
181
+ if k_block_id == 0:
182
+ topk_offsets = tl.arange(0, BLOCK_TOPK)
183
+ topk_mask = topk_offsets < top_k
184
+
185
+ ids = tl.load(
186
+ topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
187
+ mask=topk_mask,
188
+ other=0,
189
+ ).to(tl.int64)
190
+ tl.store(
191
+ topk_idx_out
192
+ + token_id * topk_idx_stride_m
193
+ + topk_offsets * topk_idx_stride_k,
194
+ ids,
195
+ mask=topk_mask,
196
+ )
197
+
198
+ weights = tl.load(
199
+ topk_weights
200
+ + token_id * topk_weights_stride_m
201
+ + topk_offsets * topk_weights_stride_k,
202
+ mask=topk_mask,
203
+ other=0.0,
204
+ )
205
+ tl.store(
206
+ topk_weights_out
207
+ + token_id * topk_weights_out_stride_m
208
+ + topk_offsets * topk_weights_out_stride_k,
209
+ weights,
210
+ mask=topk_mask,
211
+ )
212
+
213
+
214
+def _stage_deepseek_v4_mega_moe_inputs(
215
+ hidden_states: torch.Tensor,
216
+ topk_weights: torch.Tensor,
217
+ topk_ids: torch.Tensor,
218
+ x_fp8: torch.Tensor,
219
+ x_sf: torch.Tensor,
220
+ topk_idx_out: torch.Tensor,
221
+ topk_weights_out: torch.Tensor,
222
+) -> None:
223
+ num_tokens, hidden_size = hidden_states.shape
224
+ if num_tokens == 0:
225
+ return
226
+ if hidden_size % 128 != 0:
227
+ raise ValueError(
228
+ "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
229
+ "a multiple of 128."
230
+ )
231
+ top_k = topk_ids.shape[1]
232
+ if topk_weights.shape != topk_ids.shape:
233
+ raise ValueError(
234
+ "DeepSeek V4 MegaMoE input staging requires topk_weights and "
235
+ "topk_ids to have the same shape."
236
+ )
237
+
238
+ block_k = 128
239
+ grid = (num_tokens, triton.cdiv(hidden_size, block_k))
240
+ block_topk = triton.next_power_of_2(top_k)
241
+ _deepseek_v4_stage_mega_moe_inputs_kernel[grid](
242
+ hidden_states,
243
+ x_fp8,
244
+ x_sf,
245
+ topk_ids,
246
+ topk_weights,
247
+ topk_idx_out,
248
+ topk_weights_out,
249
+ hidden_states.stride(0),
250
+ hidden_states.stride(1),
251
+ x_fp8.stride(0),
252
+ x_fp8.stride(1),
253
+ x_sf.stride(0),
254
+ x_sf.stride(1),
255
+ topk_ids.stride(0),
256
+ topk_ids.stride(1),
257
+ topk_weights.stride(0),
258
+ topk_weights.stride(1),
259
+ topk_idx_out.stride(0),
260
+ topk_idx_out.stride(1),
261
+ topk_weights_out.stride(0),
262
+ topk_weights_out.stride(1),
263
+ hidden_size,
264
+ top_k,
265
+ BLOCK_K=block_k,
266
+ GROUP_K=32,
267
+ BLOCK_TOPK=block_topk,
268
+ num_warps=4,
269
+ )
270
+
271
+
272
+def make_deepseek_v4_expert_params_mapping(
273
+ num_experts: int,
274
+) -> list[tuple[str, str, int, str]]:
275
+ return [
276
+ (
277
+ "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_",
278
+ f"experts.{expert_id}.{weight_name}.",
279
+ expert_id,
280
+ shard_id,
281
+ )
282
+ for expert_id in range(num_experts)
283
+ for shard_id, weight_name in [
284
+ ("w1", "w1"),
285
+ ("w2", "w2"),
286
+ ("w3", "w3"),
287
+ ]
288
+ ]
289
+
290
+
291
+class DeepseekV4MegaMoEExperts(nn.Module):
292
+ _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
293
+
294
+ def __init__(
295
+ self,
296
+ vllm_config: VllmConfig,
297
+ *,
298
+ num_experts: int,
299
+ num_local_experts: int,
300
+ experts_start_idx: int,
301
+ top_k: int,
302
+ hidden_size: int,
303
+ intermediate_size: int,
304
+ prefix: str = "",
305
+ ):
306
+ super().__init__()
307
+ self.prefix = prefix
308
+ self.num_experts = num_experts
309
+ self.num_local_experts = num_local_experts
310
+ self.experts_start_idx = experts_start_idx
311
+ self.experts_end_idx = experts_start_idx + num_local_experts
312
+ self.top_k = top_k
313
+ self.hidden_size = hidden_size
314
+ self.intermediate_size = intermediate_size
315
+ self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
316
+
317
+ weight_attrs = {"weight_loader": self.weight_loader}
318
+ self.w13_weight = nn.Parameter(
319
+ torch.zeros(
320
+ num_local_experts,
321
+ 2 * intermediate_size,
322
+ hidden_size // 2,
323
+ dtype=torch.uint8,
324
+ ),
325
+ requires_grad=False,
326
+ )
327
+ set_weight_attrs(self.w13_weight, weight_attrs)
328
+
329
+ self.w13_weight_scale = nn.Parameter(
330
+ torch.zeros(
331
+ num_local_experts,
332
+ 2 * intermediate_size,
333
+ hidden_size // 32,
334
+ dtype=torch.uint8,
335
+ ),
336
+ requires_grad=False,
337
+ )
338
+ set_weight_attrs(self.w13_weight_scale, weight_attrs)
339
+ self.w13_weight_scale.quant_method = "block"
340
+
341
+ self.w2_weight = nn.Parameter(
342
+ torch.zeros(
343
+ num_local_experts,
344
+ hidden_size,
345
+ intermediate_size // 2,
346
+ dtype=torch.uint8,
347
+ ),
348
+ requires_grad=False,
349
+ )
350
+ set_weight_attrs(self.w2_weight, weight_attrs)
351
+
352
+ self.w2_weight_scale = nn.Parameter(
353
+ torch.zeros(
354
+ num_local_experts,
355
+ hidden_size,
356
+ intermediate_size // 32,
357
+ dtype=torch.uint8,
358
+ ),
359
+ requires_grad=False,
360
+ )
361
+ set_weight_attrs(self.w2_weight_scale, weight_attrs)
362
+ self.w2_weight_scale.quant_method = "block"
363
+
364
+ self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
365
+ self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
366
+
367
+ # Register in the static forward context so the custom-op wrapper
368
+ # can look up this module by name from within a torch.compile graph.
369
+ compilation_config = vllm_config.compilation_config
370
+ if prefix in compilation_config.static_forward_context:
371
+ raise ValueError(f"Duplicate layer name: {prefix}")
372
+ compilation_config.static_forward_context[prefix] = self
373
+
374
+ def _map_global_expert_id(self, expert_id: int) -> int:
375
+ if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:
376
+ return -1
377
+ return expert_id - self.experts_start_idx
378
+
379
+ def weight_loader(
380
+ self,
381
+ param: nn.Parameter,
382
+ loaded_weight: torch.Tensor,
383
+ weight_name: str,
384
+ shard_id: str,
385
+ expert_id: int,
386
+ return_success: bool = False,
387
+ ) -> bool | None:
388
+ local_expert_id = self._map_global_expert_id(expert_id)
389
+ if local_expert_id == -1:
390
+ return False if return_success else None
391
+
392
+ expert_data = param.data[local_expert_id]
393
+ if shard_id in ("w1", "w3"):
394
+ if "w13_" not in weight_name:
395
+ return False if return_success else None
396
+ shard_offset = 0 if shard_id == "w1" else self.intermediate_size
397
+ expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size)
398
+ elif shard_id == "w2":
399
+ if "w2_" not in weight_name:
400
+ return False if return_success else None
401
+ else:
402
+ raise ValueError(f"Unsupported expert shard id: {shard_id}")
403
+
404
+ if expert_data.shape != loaded_weight.shape:
405
+ raise ValueError(
406
+ f"DeepSeek V4 MegaMoE expert weight shape mismatch for "
407
+ f"{weight_name}: parameter shard {tuple(expert_data.shape)} "
408
+ f"vs checkpoint {tuple(loaded_weight.shape)}"
409
+ )
410
+ expert_data.copy_(loaded_weight)
411
+ return True if return_success else None
412
+
413
+ @staticmethod
414
+ def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor:
415
+ return (sf.to(torch.int32) << 23).view(torch.float32)
416
+
417
+ def _check_runtime_supported(self) -> None:
418
+ if not torch.cuda.is_available():
419
+ raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.")
420
+ device = self.w13_weight.device
421
+ if device.type != "cuda":
422
+ raise NotImplementedError(
423
+ "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA."
424
+ )
425
+ if torch.cuda.get_device_capability(device)[0] != 10:
426
+ raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.")
427
+ if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0:
428
+ raise ValueError(
429
+ "DeepGEMM MegaMoE requires hidden and intermediate sizes "
430
+ "to be multiples of 128."
431
+ )
432
+
433
+ def finalize_weights(self) -> None:
434
+ if self._transformed_l1_weights is not None:
435
+ return
436
+
437
+ self._check_runtime_supported()
438
+ import vllm.third_party.deep_gemm as deep_gemm
439
+
440
+ w13_scale = deep_gemm.transform_sf_into_required_layout(
441
+ self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
442
+ 2 * self.intermediate_size,
443
+ self.hidden_size,
444
+ (1, 32),
445
+ self.num_local_experts,
446
+ )
447
+ w2_scale = deep_gemm.transform_sf_into_required_layout(
448
+ self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
449
+ self.hidden_size,
450
+ self.intermediate_size,
451
+ (1, 32),
452
+ self.num_local_experts,
453
+ )
454
+ self._transformed_l1_weights, self._transformed_l2_weights = (
455
+ deep_gemm.transform_weights_for_mega_moe(
456
+ (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
457
+ (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
458
+ )
459
+ )
460
+ # Drop the original loader-side parameters: the MegaMoE kernels only
461
+ # consume the transformed views above. transform_weights_for_mega_moe
462
+ # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
463
+ # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
464
+ # aliases the original storage, and _transformed_l2_weights still holds
465
+ # it, so the storage stays live after we drop the Parameter.
466
+ self.w13_weight = None
467
+ self.w13_weight_scale = None
468
+ self.w2_weight = None
469
+ self.w2_weight_scale = None
470
+
471
+ def get_symm_buffer(self):
472
+ import vllm.third_party.deep_gemm as deep_gemm
473
+
474
+ group = get_ep_group().device_group
475
+ device = torch.accelerator.current_device_index()
476
+ key = (
477
+ id(group),
478
+ device,
479
+ self.num_experts,
480
+ self.max_num_tokens,
481
+ self.top_k,
482
+ self.hidden_size,
483
+ self.intermediate_size,
484
+ )
485
+ symm_buffer = self._symm_buffer_cache.get(key)
486
+ if symm_buffer is None:
487
+ symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
488
+ group,
489
+ self.num_experts,
490
+ self.max_num_tokens,
491
+ self.top_k,
492
+ self.hidden_size,
493
+ self.intermediate_size,
494
+ )
495
+ self._symm_buffer_cache[key] = symm_buffer
496
+ return symm_buffer
497
+
498
+ def forward(
499
+ self,
500
+ hidden_states: torch.Tensor,
501
+ topk_weights: torch.Tensor,
502
+ topk_ids: torch.Tensor,
503
+ *,
504
+ activation_clamp: float | None,
505
+ fast_math: bool = True,
506
+ ) -> torch.Tensor:
507
+ if hidden_states.shape[0] > self.max_num_tokens:
508
+ raise ValueError(
509
+ f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
510
+ f"but the symmetric buffer was sized for {self.max_num_tokens}."
511
+ )
512
+ y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
513
+ torch.ops.vllm.deepseek_v4_mega_moe_experts(
514
+ hidden_states,
515
+ topk_weights,
516
+ topk_ids,
517
+ y,
518
+ self.prefix,
519
+ activation_clamp,
520
+ fast_math,
521
+ )
522
+ return y
523
+
524
+ def _run_mega_moe(
525
+ self,
526
+ hidden_states: torch.Tensor,
527
+ topk_weights: torch.Tensor,
528
+ topk_ids: torch.Tensor,
529
+ y: torch.Tensor,
530
+ activation_clamp: float | None,
531
+ fast_math: bool,
532
+ ) -> None:
533
+ import vllm.third_party.deep_gemm as deep_gemm
534
+
535
+ symm_buffer = self.get_symm_buffer()
536
+ num_tokens = hidden_states.shape[0]
537
+ _stage_deepseek_v4_mega_moe_inputs(
538
+ hidden_states,
539
+ topk_weights,
540
+ topk_ids,
541
+ symm_buffer.x[:num_tokens],
542
+ symm_buffer.x_sf[:num_tokens],
543
+ symm_buffer.topk_idx[:num_tokens],
544
+ symm_buffer.topk_weights[:num_tokens],
545
+ )
546
+
547
+ # This method must have been already called during the weight loading phase.
548
+ # We call it again here to cover the dummy weight loading case.
549
+ self.finalize_weights()
550
+
551
+ assert self._transformed_l1_weights is not None
552
+ assert self._transformed_l2_weights is not None
553
+ deep_gemm.fp8_fp4_mega_moe(
554
+ y,
555
+ self._transformed_l1_weights,
556
+ self._transformed_l2_weights,
557
+ symm_buffer,
558
+ activation_clamp=activation_clamp,
559
+ fast_math=fast_math,
560
+ )
561
+
562
+
563
+DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
564
+
565
+
566
+def _deepseek_v4_mega_moe_experts_op(
567
+ hidden_states: torch.Tensor,
568
+ topk_weights: torch.Tensor,
569
+ topk_ids: torch.Tensor,
570
+ out: torch.Tensor,
571
+ layer_name: str,
572
+ activation_clamp: float | None,
573
+ fast_math: bool,
574
+) -> None:
575
+ self = get_forward_context().no_compile_layers[layer_name]
576
+ self._run_mega_moe(
577
+ hidden_states,
578
+ topk_weights,
579
+ topk_ids,
580
+ out,
581
+ activation_clamp,
582
+ fast_math,
583
+ )
584
+
585
+
586
+def _deepseek_v4_mega_moe_experts_op_fake(
587
+ hidden_states: torch.Tensor,
588
+ topk_weights: torch.Tensor,
589
+ topk_ids: torch.Tensor,
590
+ out: torch.Tensor,
591
+ layer_name: str,
592
+ activation_clamp: float | None,
593
+ fast_math: bool,
594
+) -> None:
595
+ return None
596
+
597
+
598
+direct_register_custom_op(
599
+ op_name="deepseek_v4_mega_moe_experts",
600
+ op_func=_deepseek_v4_mega_moe_experts_op,
601
+ mutates_args=["out"],
602
+ fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
603
+)
604
+
605
+
606
+class DeepseekV4MoE(nn.Module):
607
+ def __init__(
608
+ self,
609
+ vllm_config: VllmConfig,
610
+ prefix: str = "",
611
+ ):
612
+ super().__init__()
613
+
614
+ self.tp_size = get_tensor_model_parallel_world_size()
615
+ config = vllm_config.model_config.hf_config
616
+ quant_config = vllm_config.quant_config
617
+ self.prefix = prefix
618
+ if vllm_config.parallel_config.enable_expert_parallel:
619
+ self.use_mega_moe = (
620
+ vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
621
+ )
622
+ else:
623
+ self.use_mega_moe = False
624
+
625
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
626
+ self.hidden_size = config.hidden_size
627
+
628
+ self.n_routed_experts = config.n_routed_experts
629
+ self.n_activated_experts = config.num_experts_per_tok
630
+ self.moe_intermediate_size = config.moe_intermediate_size
631
+ self.swiglu_limit = config.swiglu_limit
632
+ self.renormalize = config.norm_topk_prob
633
+ self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
634
+ if self.use_mega_moe and self.scoring_func != "sqrtsoftplus":
635
+ raise NotImplementedError(
636
+ "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
637
+ )
638
+
639
+ self.gate = GateLinear(
640
+ config.hidden_size,
641
+ config.n_routed_experts,
642
+ out_dtype=torch.float32,
643
+ bias=False,
644
+ prefix=f"{prefix}.gate",
645
+ )
646
+ self.gate.e_score_correction_bias = None
647
+ self.gate.tid2eid = None
648
+ is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
649
+ self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
650
+
651
+ if is_hash_moe:
652
+ # hash MoE doesn't use e_score_correction_bias
653
+ # Use randint instead of empty to avoid garbage values causing
654
+ # invalid memory access in dummy mode (--load-format="dummy")
655
+ self.gate.tid2eid = nn.Parameter(
656
+ torch.randint(
657
+ 0,
658
+ config.n_routed_experts,
659
+ (config.vocab_size, config.num_experts_per_tok),
660
+ dtype=self.hash_indices_dtype,
661
+ ),
662
+ requires_grad=False,
663
+ )
664
+ elif getattr(config, "topk_method", None) == "noaux_tc":
665
+ self.gate.e_score_correction_bias = nn.Parameter(
666
+ torch.empty(config.n_routed_experts, dtype=torch.float32),
667
+ requires_grad=False,
668
+ )
669
+
670
+ if config.n_shared_experts is None:
671
+ self.shared_experts = None
672
+ else:
673
+ intermediate_size = config.moe_intermediate_size * config.n_shared_experts
674
+
675
+ self.shared_experts = DeepseekV2MLP(
676
+ hidden_size=config.hidden_size,
677
+ intermediate_size=intermediate_size,
678
+ hidden_act=config.hidden_act,
679
+ quant_config=quant_config,
680
+ reduce_results=self.use_mega_moe,
681
+ prefix=f"{prefix}.shared_experts",
682
+ )
683
+
684
+ if self.use_mega_moe:
685
+ self._init_mega_moe_experts(vllm_config, config, prefix)
686
+ else:
687
+ self._init_fused_moe_experts(config, quant_config, prefix)
688
+
689
+ def _init_mega_moe_experts(
690
+ self,
691
+ vllm_config: VllmConfig,
692
+ config,
693
+ prefix: str,
694
+ ) -> None:
695
+ self.ep_group = get_ep_group()
696
+ self.ep_size = self.ep_group.world_size
697
+ self.ep_rank = self.ep_group.rank_in_group
698
+ assert config.n_routed_experts % self.ep_size == 0
699
+
700
+ self.n_local_experts = config.n_routed_experts // self.ep_size
701
+ self.experts_start_idx = self.ep_rank * self.n_local_experts
702
+ self.experts_end_idx = self.experts_start_idx + self.n_local_experts
703
+
704
+ self.experts = DeepseekV4MegaMoEExperts(
705
+ vllm_config,
706
+ num_experts=config.n_routed_experts,
707
+ num_local_experts=self.n_local_experts,
708
+ experts_start_idx=self.experts_start_idx,
709
+ top_k=config.num_experts_per_tok,
710
+ hidden_size=config.hidden_size,
711
+ intermediate_size=config.moe_intermediate_size,
712
+ prefix=f"{prefix}.experts",
713
+ )
714
+
715
+ def _init_fused_moe_experts(
716
+ self,
717
+ config,
718
+ quant_config,
719
+ prefix: str,
720
+ ) -> None:
721
+ self.tp_rank = get_tensor_model_parallel_rank()
722
+ assert config.n_routed_experts % self.tp_size == 0
723
+
724
+ self.n_local_experts = config.n_routed_experts // self.tp_size
725
+ self.experts_start_idx = self.tp_rank * self.n_local_experts
726
+ self.experts_end_idx = self.experts_start_idx + self.n_local_experts
727
+
728
+ self.experts = FusedMoE(
729
+ shared_experts=self.shared_experts,
730
+ gate=self.gate,
731
+ num_experts=config.n_routed_experts,
732
+ top_k=config.num_experts_per_tok,
733
+ hidden_size=config.hidden_size,
734
+ intermediate_size=config.moe_intermediate_size,
735
+ renormalize=config.norm_topk_prob,
736
+ quant_config=quant_config,
737
+ prefix=f"{prefix}.experts",
738
+ scoring_func=self.scoring_func,
739
+ routed_scaling_factor=self.routed_scaling_factor,
740
+ e_score_correction_bias=self.gate.e_score_correction_bias,
741
+ hash_indices_table=self.gate.tid2eid,
742
+ swiglu_limit=self.swiglu_limit,
743
+ router_logits_dtype=torch.float32,
744
+ )
745
+
746
+ def forward(
747
+ self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
748
+ ) -> torch.Tensor:
749
+ if self.gate.tid2eid is not None:
750
+ if input_ids is None:
751
+ raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
752
+ input_ids = input_ids.to(dtype=self.hash_indices_dtype)
753
+ if not self.use_mega_moe:
754
+ return self._forward_fused_moe(hidden_states, input_ids)
755
+
756
+ org_shape = hidden_states.shape
757
+ router_logits, _ = self.gate(hidden_states)
758
+ topk_weights, topk_ids = fused_topk_bias(
759
+ hidden_states=hidden_states,
760
+ gating_output=router_logits,
761
+ scoring_func=self.scoring_func,
762
+ e_score_correction_bias=self.gate.e_score_correction_bias.data
763
+ if self.gate.e_score_correction_bias is not None
764
+ else None,
765
+ topk=self.n_activated_experts,
766
+ renormalize=self.renormalize,
767
+ indices_type=self.hash_indices_dtype,
768
+ input_tokens=input_ids,
769
+ hash_indices_table=self.gate.tid2eid,
770
+ routed_scaling_factor=self.routed_scaling_factor,
771
+ )
772
+ activation_clamp = (
773
+ float(self.swiglu_limit) if self.swiglu_limit is not None else None
774
+ )
775
+ final_hidden_states = self.experts(
776
+ hidden_states,
777
+ topk_weights,
778
+ topk_ids,
779
+ activation_clamp=activation_clamp,
780
+ )
781
+
782
+ if self.shared_experts is not None:
783
+ shared_output = self.shared_experts(hidden_states)
784
+ final_hidden_states += shared_output
785
+
786
+ return final_hidden_states.view(org_shape)
787
+
788
+ def _forward_fused_moe(
789
+ self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
790
+ ) -> torch.Tensor:
791
+ org_shape = hidden_states.shape
792
+ if self.experts.is_internal_router:
793
+ # In this case, the gate/router runs inside the FusedMoE class
794
+ final_hidden_states = self.experts(
795
+ hidden_states=hidden_states,
796
+ router_logits=hidden_states,
797
+ input_ids=input_ids,
798
+ )
799
+ else:
800
+ router_logits, _ = self.gate(hidden_states)
801
+ final_hidden_states = self.experts(
802
+ hidden_states=hidden_states,
803
+ router_logits=router_logits,
804
+ input_ids=input_ids,
805
+ )
806
+
807
+ return final_hidden_states.view(org_shape)
808
+
809
+ def finalize_mega_moe_weights(self) -> None:
810
+ if self.use_mega_moe:
811
+ self.experts.finalize_weights()
812
+
813
+
814
+class DeepseekV4Attention(nn.Module):
815
+ def __init__(
816
+ self,
817
+ vllm_config: VllmConfig,
818
+ prefix: str,
819
+ topk_indices_buffer: torch.Tensor | None = None,
820
+ aux_stream: torch.cuda.Stream | None = None,
821
+ ):
822
+ super().__init__()
823
+ config = vllm_config.model_config.hf_config
824
+ quant_config = vllm_config.quant_config
825
+ layer_id = extract_layer_index(prefix)
826
+
827
+ self.layer_id = layer_id
828
+ self.hidden_size = config.hidden_size
829
+ self.n_heads = config.num_attention_heads
830
+ tp_size = get_tensor_model_parallel_world_size()
831
+ assert self.n_heads % tp_size == 0
832
+
833
+ self.n_local_heads = self.n_heads // tp_size
834
+ self.q_lora_rank = config.q_lora_rank
835
+ self.o_lora_rank = config.o_lora_rank
836
+ self.head_dim = config.head_dim
837
+ self.rope_head_dim = config.qk_rope_head_dim
838
+ self.nope_head_dim = self.head_dim - self.rope_head_dim
839
+ self.n_groups = config.o_groups
840
+ self.n_local_groups = self.n_groups // tp_size
841
+ self.window_size = config.sliding_window
842
+ # NOTE(zyongye) Compress ratio can't be 0
843
+ # we do this for because MTP layer is not included
844
+ # in the compress ratio list
845
+ if layer_id < config.num_hidden_layers:
846
+ self.compress_ratio = max(1, config.compress_ratios[layer_id])
847
+ else:
848
+ self.compress_ratio = 1
849
+ self.eps = config.rms_norm_eps
850
+ self.max_position_embeddings = config.max_position_embeddings
851
+
852
+ # Padded to min 64 heads for FlashMLA, initialized to -inf
853
+ # (no sink effect). Weight loading fills the first n_local_heads slots.
854
+ padded_heads = max(self.n_local_heads, 64)
855
+ self.attn_sink = nn.Parameter(
856
+ torch.full((padded_heads,), -float("inf"), dtype=torch.float32),
857
+ requires_grad=False,
858
+ )
859
+
860
+ self.fused_wqa_wkv = MergedColumnParallelLinear(
861
+ self.hidden_size,
862
+ [self.q_lora_rank, self.head_dim],
863
+ bias=False,
864
+ quant_config=quant_config,
865
+ prefix=f"{prefix}.fused_wqa_wkv",
866
+ disable_tp=True, # fused ReplicatedLinear
867
+ )
868
+ self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
869
+ self.wq_b = ColumnParallelLinear(
870
+ self.q_lora_rank,
871
+ self.n_heads * self.head_dim,
872
+ bias=False,
873
+ quant_config=quant_config,
874
+ return_bias=False,
875
+ prefix=f"{prefix}.wq_b",
876
+ )
877
+
878
+ self.kv_norm = RMSNorm(self.head_dim, self.eps)
879
+ self.wo_a = ColumnParallelLinear(
880
+ self.n_heads * self.head_dim // self.n_groups,
881
+ self.n_groups * self.o_lora_rank,
882
+ bias=False,
883
+ quant_config=quant_config,
884
+ return_bias=False,
885
+ prefix=f"{prefix}.wo_a",
886
+ )
887
+ self.wo_a.is_bmm = True
888
+ self.wo_a.bmm_batch_size = self.n_local_groups
889
+ self.wo_b = RowParallelLinear(
890
+ self.n_groups * self.o_lora_rank,
891
+ self.hidden_size,
892
+ bias=False,
893
+ quant_config=quant_config,
894
+ return_bias=False,
895
+ prefix=f"{prefix}.wo_b",
896
+ )
897
+ self.softmax_scale = self.head_dim**-0.5
898
+ self.scale_fmt = config.quantization_config["scale_fmt"]
899
+
900
+ self.rope_parameters = config.rope_scaling
901
+
902
+ # Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it)
903
+ rope_parameters = config.rope_parameters
904
+ rope_parameters["rope_theta"] = (
905
+ config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta
906
+ )
907
+ if config.rope_parameters["rope_type"] != "default":
908
+ config.rope_parameters["rope_type"] = (
909
+ "deepseek_yarn"
910
+ if config.rope_parameters.get("apply_yarn_scaling", True)
911
+ else "deepseek_llama_scaling"
912
+ )
913
+ rope_parameters["mscale"] = 0 # Disable mscale
914
+ rope_parameters["mscale_all_dim"] = 0 # Disable mscale
915
+ rope_parameters["is_deepseek_v4"] = True
916
+ rope_parameters["rope_dim"] = self.rope_head_dim
917
+ self.rotary_emb = get_rope(
918
+ self.head_dim,
919
+ max_position=self.max_position_embeddings,
920
+ rope_parameters=rope_parameters,
921
+ is_neox_style=False,
922
+ dtype=config.torch_dtype,
923
+ )
924
+
925
+ self.indexer = None
926
+ if self.compress_ratio == 4:
927
+ # Only C4A uses sparse attention and hence has indexer.
928
+ self.indexer = DeepseekV4Indexer(
929
+ vllm_config,
930
+ config=config,
931
+ hidden_size=self.hidden_size,
932
+ q_lora_rank=self.q_lora_rank,
933
+ quant_config=quant_config,
934
+ cache_config=vllm_config.cache_config,
935
+ topk_indices_buffer=topk_indices_buffer,
936
+ compress_ratio=self.compress_ratio,
937
+ prefix=f"{prefix}.indexer",
938
+ )
939
+
940
+ mla_modules = DeepseekV4MLAModules(
941
+ vllm_config=vllm_config,
942
+ fused_wqa_wkv=self.fused_wqa_wkv,
943
+ q_norm=self.q_norm,
944
+ wq_b=self.wq_b,
945
+ kv_norm=self.kv_norm,
946
+ wo_a=self.wo_a,
947
+ wo_b=self.wo_b,
948
+ attn_sink=self.attn_sink,
949
+ rotary_emb=self.rotary_emb,
950
+ indexer=self.indexer,
951
+ indexer_rotary_emb=self.rotary_emb,
952
+ topk_indices_buffer=topk_indices_buffer,
953
+ aux_stream=aux_stream,
954
+ )
955
+ self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper(
956
+ hidden_size=self.hidden_size,
957
+ num_heads=self.n_local_heads,
958
+ head_dim=self.head_dim,
959
+ scale=self.softmax_scale,
960
+ qk_nope_head_dim=self.nope_head_dim,
961
+ qk_rope_head_dim=self.rope_head_dim,
962
+ v_head_dim=self.head_dim,
963
+ q_lora_rank=self.q_lora_rank,
964
+ kv_lora_rank=self.head_dim,
965
+ o_lora_rank=self.o_lora_rank,
966
+ mla_modules=mla_modules,
967
+ window_size=self.window_size,
968
+ compress_ratio=self.compress_ratio,
969
+ cache_config=vllm_config.cache_config,
970
+ quant_config=quant_config,
971
+ prefix=prefix,
972
+ )
973
+
974
+ def forward(
975
+ self,
976
+ positions: torch.Tensor,
977
+ hidden_states: torch.Tensor,
978
+ llama_4_scaling: torch.Tensor | None,
979
+ ):
980
+ return self.mla_attn(positions, hidden_states, llama_4_scaling)
981
+
982
+
983
+class DeepseekV4DecoderLayer(nn.Module):
984
+ def __init__(
985
+ self,
986
+ vllm_config,
987
+ prefix,
988
+ topk_indices_buffer: torch.Tensor | None = None,
989
+ aux_stream_dict: dict[AuxStreamType, torch.cuda.Stream] | None = None,
990
+ ):
991
+ super().__init__()
992
+ config = vllm_config.model_config.hf_config
993
+ self.hidden_size = config.hidden_size
994
+
995
+ self.rms_norm_eps = config.rms_norm_eps
996
+ self.attn = DeepseekV4Attention(
997
+ vllm_config,
998
+ prefix=f"{prefix}.attn",
999
+ topk_indices_buffer=topk_indices_buffer,
1000
+ aux_stream=aux_stream_dict.get(AuxStreamType.Attention)
1001
+ if aux_stream_dict is not None
1002
+ else None,
1003
+ )
1004
+ self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
1005
+
1006
+ self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1007
+ self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1008
+ self.hc_mult = config.hc_mult
1009
+ self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
1010
+ self.hc_eps = config.hc_eps
1011
+ self.hc_post_alpha = 2.0
1012
+ mix_hc = (2 + self.hc_mult) * self.hc_mult
1013
+ hc_dim = self.hc_mult * self.hidden_size
1014
+ self.hc_attn_fn = nn.Parameter(
1015
+ torch.empty(
1016
+ (mix_hc, hc_dim),
1017
+ dtype=torch.float32,
1018
+ ),
1019
+ requires_grad=False,
1020
+ )
1021
+ self.hc_ffn_fn = nn.Parameter(
1022
+ torch.empty(
1023
+ (mix_hc, hc_dim),
1024
+ dtype=torch.float32,
1025
+ ),
1026
+ requires_grad=False,
1027
+ )
1028
+ self.hc_attn_base = nn.Parameter(
1029
+ torch.empty(
1030
+ mix_hc,
1031
+ dtype=torch.float32,
1032
+ ),
1033
+ requires_grad=False,
1034
+ )
1035
+ self.hc_ffn_base = nn.Parameter(
1036
+ torch.empty(
1037
+ mix_hc,
1038
+ dtype=torch.float32,
1039
+ ),
1040
+ requires_grad=False,
1041
+ )
1042
+ self.hc_attn_scale = nn.Parameter(
1043
+ torch.empty(
1044
+ 3,
1045
+ dtype=torch.float32,
1046
+ ),
1047
+ requires_grad=False,
1048
+ )
1049
+ self.hc_ffn_scale = nn.Parameter(
1050
+ torch.empty(
1051
+ 3,
1052
+ dtype=torch.float32,
1053
+ ),
1054
+ requires_grad=False,
1055
+ )
1056
+
1057
+ def hc_pre(
1058
+ self,
1059
+ x: torch.Tensor,
1060
+ hc_fn: torch.Tensor,
1061
+ hc_scale: torch.Tensor,
1062
+ hc_base: torch.Tensor,
1063
+ ):
1064
+ # Lazy import to avoid top-level tilelang dependency.
1065
+ # Registers both torch.ops.vllm.mhc_pre and mhc_post,
1066
+ # so hc_post() doesn't need its own import.
1067
+ import vllm.model_executor.layers.mhc # noqa: F401
1068
+
1069
+ post_mix, res_mix, layer_input = torch.ops.vllm.mhc_pre(
1070
+ residual=x,
1071
+ fn=hc_fn,
1072
+ hc_scale=hc_scale,
1073
+ hc_base=hc_base,
1074
+ rms_eps=self.rms_norm_eps,
1075
+ hc_pre_eps=self.hc_eps,
1076
+ hc_sinkhorn_eps=self.hc_eps,
1077
+ hc_post_mult_value=self.hc_post_alpha,
1078
+ sinkhorn_repeat=self.hc_sinkhorn_iters,
1079
+ )
1080
+ return layer_input, post_mix, res_mix
1081
+
1082
+ def hc_post(
1083
+ self,
1084
+ x: torch.Tensor,
1085
+ residual: torch.Tensor,
1086
+ post: torch.Tensor,
1087
+ comb: torch.Tensor,
1088
+ ):
1089
+ return torch.ops.vllm.mhc_post(x, residual, post, comb)
1090
+
1091
+ def forward(
1092
+ self,
1093
+ x: torch.Tensor,
1094
+ positions: torch.Tensor,
1095
+ input_ids: torch.Tensor | None,
1096
+ ) -> torch.Tensor:
1097
+ residual = x
1098
+ x, post, comb = self.hc_pre(
1099
+ x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1100
+ )
1101
+ x = self.attn_norm(x)
1102
+ x = self.attn(positions, x, None)
1103
+ x = self.hc_post(x, residual, post, comb)
1104
+
1105
+ residual = x
1106
+ x, post, comb = self.hc_pre(
1107
+ x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
1108
+ )
1109
+ x = self.ffn_norm(x)
1110
+ x = self.ffn(x, input_ids)
1111
+ x = self.hc_post(x, residual, post, comb)
1112
+ return x
1113
+
1114
+
1115
+@support_torch_compile
1116
+class DeepseekV4Model(nn.Module):
1117
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1118
+ super().__init__()
1119
+
1120
+ config = vllm_config.model_config.hf_config
1121
+ quant_config = vllm_config.quant_config
1122
+ self.config = config
1123
+
1124
+ self.vocab_size = config.vocab_size
1125
+ self.hc_eps = config.hc_eps
1126
+ self.hc_mult = config.hc_mult
1127
+ self.hc_dim = self.hc_mult * config.hidden_size
1128
+ self.rms_norm_eps = config.rms_norm_eps
1129
+
1130
+ aux_stream_list = [torch.cuda.Stream() for _ in range(1)]
1131
+ self.aux_stream_dict = {
1132
+ AuxStreamType.Attention: aux_stream_list[0],
1133
+ }
1134
+
1135
+ self.device = current_platform.device_type
1136
+ # Reserved topk indices buffer for all Indexer layers to reuse.
1137
+ self.topk_indices_buffer = torch.empty(
1138
+ vllm_config.scheduler_config.max_num_batched_tokens,
1139
+ config.index_topk,
1140
+ dtype=torch.int32,
1141
+ device=self.device,
1142
+ )
1143
+
1144
+ self.embed_tokens = VocabParallelEmbedding(
1145
+ config.vocab_size,
1146
+ config.hidden_size,
1147
+ quant_config=quant_config,
1148
+ prefix=f"{prefix}.embed_tokens",
1149
+ )
1150
+
1151
+ self.start_layer, self.end_layer, self.layers = make_layers(
1152
+ config.num_hidden_layers,
1153
+ lambda prefix: DeepseekV4DecoderLayer(
1154
+ vllm_config,
1155
+ prefix=prefix,
1156
+ topk_indices_buffer=self.topk_indices_buffer,
1157
+ aux_stream_dict=self.aux_stream_dict,
1158
+ ),
1159
+ prefix=f"{prefix}.layers",
1160
+ )
1161
+
1162
+ self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
1163
+
1164
+ self.hc_head_fn = nn.Parameter(
1165
+ torch.empty(
1166
+ self.hc_mult,
1167
+ self.hc_dim,
1168
+ dtype=torch.float32,
1169
+ ),
1170
+ requires_grad=False,
1171
+ )
1172
+ self.hc_head_base = nn.Parameter(
1173
+ torch.empty(
1174
+ self.hc_mult,
1175
+ dtype=torch.float32,
1176
+ ),
1177
+ requires_grad=False,
1178
+ )
1179
+ self.hc_head_scale = nn.Parameter(
1180
+ torch.empty(1, dtype=torch.float32),
1181
+ requires_grad=False,
1182
+ )
1183
+
1184
+ # Pre-hc_head residual stream buffer for the MTP draft. Stable
1185
+ # address (outside the cudagraph pool) so the copy_ in forward()
1186
+ # refreshes it correctly across captured shapes.
1187
+ self._mtp_hidden_buffer = torch.empty(
1188
+ vllm_config.scheduler_config.max_num_batched_tokens,
1189
+ self.hc_dim,
1190
+ dtype=vllm_config.model_config.dtype,
1191
+ device=self.device,
1192
+ )
1193
+
1194
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1195
+ return self.embed_tokens(input_ids)
1196
+
1197
+ def forward(
1198
+ self,
1199
+ input_ids: torch.Tensor,
1200
+ positions: torch.Tensor,
1201
+ intermediate_tensors: IntermediateTensors | None,
1202
+ inputs_embeds: torch.Tensor | None = None,
1203
+ ) -> torch.Tensor | IntermediateTensors:
1204
+ hidden_states = self.embed_input_ids(input_ids)
1205
+ hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
1206
+
1207
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
1208
+ hidden_states = layer(
1209
+ hidden_states,
1210
+ positions,
1211
+ input_ids,
1212
+ )
1213
+
1214
+ # Stash pre-hc_head residual for the MTP draft (captured copy_).
1215
+ num_tokens = hidden_states.shape[0]
1216
+ self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1217
+
1218
+ hidden_states = hc_head(
1219
+ hidden_states,
1220
+ self.hc_head_fn,
1221
+ self.hc_head_scale,
1222
+ self.hc_head_base,
1223
+ self.rms_norm_eps,
1224
+ self.hc_eps,
1225
+ )
1226
+ hidden_states = self.norm(hidden_states)
1227
+ return hidden_states
1228
+
1229
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1230
+ stacked_params_mapping = [
1231
+ # (param_name, shard_name, shard_id)
1232
+ ("gate_up_proj", "w1", 0),
1233
+ ("gate_up_proj", "w3", 1),
1234
+ ("attn.fused_wqa_wkv", "attn.wq_a", 0),
1235
+ ("attn.fused_wqa_wkv", "attn.wkv", 1),
1236
+ ("compressor.fused_wkv_wgate", "compressor.wkv", 0),
1237
+ ("compressor.fused_wkv_wgate", "compressor.wgate", 1),
1238
+ ]
1239
+ params_dict = dict(self.named_parameters())
1240
+ loaded_params: set[str] = set()
1241
+
1242
+ # TP for attention
1243
+ tp_size = get_tensor_model_parallel_world_size()
1244
+ tp_rank = get_tensor_model_parallel_rank()
1245
+ n_head = self.config.num_attention_heads
1246
+ n_local_head = n_head // tp_size
1247
+ head_rank_start = n_local_head * tp_rank
1248
+ head_rank_end = n_local_head * (tp_rank + 1)
1249
+
1250
+ # Pre-compute expert mapping ONCE.
1251
+ expert_mapping = self.get_expert_mapping()
1252
+
1253
+ for name, loaded_weight in weights:
1254
+ for param_name, weight_name, shard_id in stacked_params_mapping:
1255
+ # Skip non-stacked layers and experts (experts handled below).
1256
+ if ".experts." in name:
1257
+ continue
1258
+ if weight_name not in name:
1259
+ continue
1260
+ name = name.replace(weight_name, param_name)
1261
+
1262
+ param = params_dict[name]
1263
+ weight_loader = param.weight_loader
1264
+ weight_loader(param, loaded_weight, shard_id)
1265
+ loaded_params.add(name)
1266
+ break
1267
+ else:
1268
+ if ".experts." in name:
1269
+ # E8M0 scales are stored as float8_e8m0fnu in
1270
+ # checkpoints but the MoE param is uint8. copy_()
1271
+ # would do a numeric conversion (e.g. 2^-7 → 0),
1272
+ # destroying the raw exponent bytes.
1273
+ if (
1274
+ "weight_scale" in name
1275
+ and loaded_weight.dtype == torch.float8_e8m0fnu
1276
+ ):
1277
+ loaded_weight = loaded_weight.view(torch.uint8)
1278
+ for mapping in expert_mapping:
1279
+ param_name, weight_name, expert_id, shard_id = mapping
1280
+ if weight_name not in name:
1281
+ continue
1282
+ name_mapped = name.replace(weight_name, param_name)
1283
+ param = params_dict[name_mapped]
1284
+ # We should ask the weight loader to return success or not
1285
+ # here since otherwise we may skip experts with other
1286
+ # available replicas.
1287
+ weight_loader = typing.cast(
1288
+ Callable[..., bool], param.weight_loader
1289
+ )
1290
+ success = weight_loader(
1291
+ param,
1292
+ loaded_weight,
1293
+ name_mapped,
1294
+ shard_id=shard_id,
1295
+ expert_id=expert_id,
1296
+ return_success=True,
1297
+ )
1298
+ if success:
1299
+ name = name_mapped
1300
+ break
1301
+ loaded_params.add(name_mapped)
1302
+ continue
1303
+ elif "attn_sink" in name:
1304
+ narrow_weight = loaded_weight[head_rank_start:head_rank_end]
1305
+ n = narrow_weight.shape[0]
1306
+ params_dict[name][:n].copy_(narrow_weight)
1307
+ loaded_params.add(name)
1308
+ continue
1309
+ else:
1310
+ param = params_dict[name]
1311
+ weight_loader = getattr(
1312
+ param, "weight_loader", default_weight_loader
1313
+ )
1314
+ weight_loader(param, loaded_weight)
1315
+ loaded_params.add(name)
1316
+ continue
1317
+
1318
+ return loaded_params
1319
+
1320
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1321
+ first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
1322
+ if first_layer.ffn.use_mega_moe:
1323
+ return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
1324
+ # Params for weights, fp8 weight scales, fp8 activation scales
1325
+ # (param_name, weight_name, expert_id, shard_id)
1326
+ return FusedMoE.make_expert_params_mapping(
1327
+ self,
1328
+ ckpt_gate_proj_name="w1",
1329
+ ckpt_down_proj_name="w2",
1330
+ ckpt_up_proj_name="w3",
1331
+ num_experts=self.config.n_routed_experts,
1332
+ )
1333
+
1334
+ def finalize_mega_moe_weights(self) -> None:
1335
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
1336
+ layer.ffn.finalize_mega_moe_weights()
1337
+
1338
+
1339
+@torch.compile(backend=current_platform.simple_compile_backend)
1340
+def hc_head(
1341
+ hidden_states: torch.Tensor,
1342
+ hc_fn: torch.Tensor,
1343
+ hc_scale: torch.Tensor,
1344
+ hc_base: torch.Tensor,
1345
+ rms_norm_eps: float,
1346
+ hc_eps: float,
1347
+) -> torch.Tensor:
1348
+ x = hidden_states
1349
+ shape, dtype = x.size(), x.dtype
1350
+ x = x.flatten(1).float()
1351
+ rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + rms_norm_eps)
1352
+ mixes = F.linear(x, hc_fn) * rsqrt
1353
+ pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
1354
+ y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1)
1355
+ return y.to(dtype)
1356
+
1357
+
1358
+class DeepseekV4ForCausalLM(nn.Module):
1359
+ model_cls = DeepseekV4Model
1360
+
1361
+ hf_to_vllm_mapper = WeightsMapper(
1362
+ orig_to_new_prefix={
1363
+ "layers.": "model.layers.",
1364
+ "embed.": "model.embed.",
1365
+ "norm.": "model.norm.",
1366
+ "hc_head": "model.hc_head",
1367
+ "mtp.": "model.mtp.",
1368
+ },
1369
+ orig_to_new_regex={
1370
+ # Routed MoE expert scales: experts.N.wX.scale -> .weight_scale
1371
+ re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
1372
+ # Everything else (FP8 linear + shared experts): .scale -> .weight_scale_inv
1373
+ re.compile(r"\.scale$"): ".weight_scale_inv",
1374
+ },
1375
+ orig_to_new_suffix={
1376
+ "head.weight": "lm_head.weight",
1377
+ "embed.weight": "embed_tokens.weight",
1378
+ ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
1379
+ },
1380
+ orig_to_new_substr={
1381
+ ".attn.compressor.": ".attn.mla_attn.compressor.",
1382
+ ".shared_experts.w2": ".shared_experts.down_proj",
1383
+ },
1384
+ )
1385
+
1386
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1387
+ super().__init__()
1388
+
1389
+ config = vllm_config.model_config.hf_config
1390
+ self.config = config
1391
+
1392
+ self.model = self.model_cls(
1393
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
1394
+ )
1395
+ self.lm_head = ParallelLMHead(
1396
+ config.vocab_size,
1397
+ config.hidden_size,
1398
+ prefix=maybe_prefix(prefix, "lm_head"),
1399
+ )
1400
+ self.logits_processor = LogitsProcessor(config.vocab_size)
1401
+
1402
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1403
+ return self.model.embed_input_ids(input_ids)
1404
+
1405
+ def compute_logits(
1406
+ self,
1407
+ hidden_states: torch.Tensor,
1408
+ ) -> torch.Tensor | None:
1409
+ logits = self.logits_processor(self.lm_head, hidden_states)
1410
+ return logits
1411
+
1412
+ def forward(
1413
+ self,
1414
+ input_ids: torch.Tensor,
1415
+ positions: torch.Tensor,
1416
+ intermediate_tensors: IntermediateTensors | None = None,
1417
+ inputs_embeds: torch.Tensor | None = None,
1418
+ ) -> torch.Tensor | IntermediateTensors:
1419
+ hidden_states = self.model(
1420
+ input_ids, positions, intermediate_tensors, inputs_embeds
1421
+ )
1422
+ return hidden_states
1423
+
1424
+ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
1425
+ """Pre-hc_head residual stream buffer (max_num_batched_tokens,
1426
+ hc_mult * hidden_size) for the MTP draft model. Populated by
1427
+ forward(); valid after each target step."""
1428
+ return getattr(self.model, "_mtp_hidden_buffer", None)
1429
+
1430
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1431
+ loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
1432
+ loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1433
+ self.model.finalize_mega_moe_weights()
1434
+ return loaded_params
1435
+
1436
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1437
+ return self.model.get_expert_mapping()
vllm/model_executor/models/deepseek_v4_mtp.py
new
+483
@@ -0,0 +1,483 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""MTP draft model for DeepSeek V4 (internal codename: DeepseekV4).
4
+
5
+Split from ``deepseek_mtp.py`` because the V4 architecture introduces several
6
+pieces that have no analogue in V3/V32:
7
+ * separate ``e_proj`` / ``h_proj`` with fp8 linear quantization (instead of
8
+ the fused ``eh_proj``);
9
+ * ``hc_head`` hypercompressed vocab projection applied in ``compute_logits``;
10
+ * ``DeepseekV4DecoderLayer`` with its own aux-stream management;
11
+ * V4-specific checkpoint weight-name remapping in ``load_weights``.
12
+"""
13
+
14
+import typing
15
+from collections.abc import Callable, Iterable
16
+
17
+import regex as re
18
+import torch
19
+import torch.nn as nn
20
+
21
+from vllm.compilation.decorators import support_torch_compile
22
+from vllm.config import VllmConfig
23
+from vllm.distributed import (
24
+ get_tensor_model_parallel_rank,
25
+ get_tensor_model_parallel_world_size,
26
+)
27
+from vllm.logger import init_logger
28
+from vllm.model_executor.layers.fused_moe import FusedMoE
29
+from vllm.model_executor.layers.layernorm import RMSNorm
30
+from vllm.model_executor.layers.linear import ReplicatedLinear
31
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
32
+from vllm.model_executor.layers.vocab_parallel_embedding import (
33
+ VocabParallelEmbedding,
34
+)
35
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36
+from vllm.platforms import current_platform
37
+from vllm.sequence import IntermediateTensors
38
+from vllm.utils.multi_stream_utils import AuxStreamType
39
+
40
+from .deepseek_mtp import SharedHead
41
+from .deepseek_v2 import get_spec_layer_idx_from_weight_name
42
+from .deepseek_v4 import (
43
+ DeepseekV4DecoderLayer,
44
+ hc_head,
45
+ make_deepseek_v4_expert_params_mapping,
46
+)
47
+from .utils import maybe_prefix
48
+
49
+logger = init_logger(__name__)
50
+
51
+# MoE expert scales are fused into per-layer w13/w2 tensors; other FP8 linear
52
+# scales use `.weight_scale_inv`. Mirrors the regex in
53
+# DeepseekV4ForCausalLM.hf_to_vllm_mapper.
54
+_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
55
+
56
+
57
+class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
58
+ def __init__(
59
+ self,
60
+ vllm_config: VllmConfig,
61
+ topk_indices_buffer: torch.Tensor,
62
+ prefix: str,
63
+ ) -> None:
64
+ super().__init__()
65
+
66
+ config = vllm_config.speculative_config.draft_model_config.hf_config
67
+ self.config = config
68
+ quant_config = vllm_config.quant_config
69
+ self.rms_norm_eps = config.rms_norm_eps
70
+
71
+ self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
72
+ self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
73
+
74
+ # V4 keeps e_ and h_ proj separate (with fp8 linear quant) rather than
75
+ # fusing them the way V3 does with eh_proj.
76
+ self.e_proj = ReplicatedLinear(
77
+ config.hidden_size,
78
+ config.hidden_size,
79
+ bias=False,
80
+ return_bias=False,
81
+ quant_config=quant_config,
82
+ )
83
+ self.h_proj = ReplicatedLinear(
84
+ config.hidden_size,
85
+ config.hidden_size,
86
+ bias=False,
87
+ return_bias=False,
88
+ quant_config=quant_config,
89
+ )
90
+
91
+ self.hc_eps = config.hc_eps
92
+ self.hc_mult = config.hc_mult
93
+ self.hc_dim = self.hc_mult * config.hidden_size
94
+ self.hc_head_fn = nn.Parameter(
95
+ torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32),
96
+ requires_grad=False,
97
+ )
98
+ self.hc_head_base = nn.Parameter(
99
+ torch.empty(self.hc_mult, dtype=torch.float32),
100
+ requires_grad=False,
101
+ )
102
+ self.hc_head_scale = nn.Parameter(
103
+ torch.empty(1, dtype=torch.float32),
104
+ requires_grad=False,
105
+ )
106
+
107
+ self.shared_head = SharedHead(
108
+ config=config, prefix=prefix, quant_config=quant_config
109
+ )
110
+ self.aux_stream_dict = {
111
+ AuxStreamType.Attention: torch.cuda.Stream(),
112
+ }
113
+ self.mtp_block = DeepseekV4DecoderLayer(
114
+ vllm_config,
115
+ prefix,
116
+ topk_indices_buffer=topk_indices_buffer,
117
+ aux_stream_dict=self.aux_stream_dict,
118
+ )
119
+
120
+ def forward(
121
+ self,
122
+ input_ids: torch.Tensor,
123
+ positions: torch.Tensor,
124
+ previous_hidden_states: torch.Tensor,
125
+ inputs_embeds: torch.Tensor | None = None,
126
+ spec_step_index: int = 0,
127
+ ) -> torch.Tensor:
128
+ assert inputs_embeds is not None
129
+ # masking inputs at position 0, as not needed by MTP
130
+ inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
131
+ inputs_embeds = self.enorm(inputs_embeds)
132
+
133
+ # Target stashes pre-hc_head residual as flat (T, hc_mult * D);
134
+ # reshape to (T, hc_mult, D) — the training-time layout.
135
+ previous_hidden_states = previous_hidden_states.view(
136
+ -1, self.hc_mult, self.config.hidden_size
137
+ )
138
+ previous_hidden_states = self.hnorm(previous_hidden_states)
139
+ hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
140
+ inputs_embeds
141
+ ).unsqueeze(-2)
142
+ hidden_states = self.mtp_block(
143
+ positions=positions, x=hidden_states, input_ids=None
144
+ )
145
+ # Return the flat pre-hc_head residual so it can be re-fed as the
146
+ # next spec step's `previous_hidden_states` when
147
+ # num_speculative_tokens > 1. hc_head is deferred to compute_logits.
148
+ return hidden_states.flatten(1)
149
+
150
+
151
+class DeepSeekV4MultiTokenPredictor(nn.Module):
152
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
153
+ super().__init__()
154
+ config = vllm_config.model_config.hf_config
155
+ self.mtp_start_layer_idx = config.num_hidden_layers
156
+ self.num_mtp_layers = config.num_nextn_predict_layers
157
+ self.device = current_platform.device_type
158
+
159
+ topk_tokens = config.index_topk
160
+ self.topk_indices_buffer = torch.empty(
161
+ vllm_config.scheduler_config.max_num_batched_tokens,
162
+ topk_tokens,
163
+ dtype=torch.int32,
164
+ device=self.device,
165
+ )
166
+
167
+ # to map the exact layer index from weights
168
+ self.layers = torch.nn.ModuleDict(
169
+ {
170
+ str(idx): DeepSeekV4MultiTokenPredictorLayer(
171
+ vllm_config,
172
+ self.topk_indices_buffer,
173
+ f"{prefix}.layers.{idx}",
174
+ )
175
+ for idx in range(
176
+ self.mtp_start_layer_idx,
177
+ self.mtp_start_layer_idx + self.num_mtp_layers,
178
+ )
179
+ }
180
+ )
181
+ self.embed_tokens = VocabParallelEmbedding(
182
+ config.vocab_size,
183
+ config.hidden_size,
184
+ prefix=maybe_prefix(prefix, "embed_tokens"),
185
+ )
186
+ self.logits_processor = LogitsProcessor(config.vocab_size)
187
+
188
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
189
+ return self.embed_tokens(input_ids)
190
+
191
+ def forward(
192
+ self,
193
+ input_ids: torch.Tensor,
194
+ positions: torch.Tensor,
195
+ previous_hidden_states: torch.Tensor,
196
+ inputs_embeds: torch.Tensor | None = None,
197
+ spec_step_idx: int = 0,
198
+ ) -> torch.Tensor:
199
+ if inputs_embeds is None:
200
+ inputs_embeds = self.embed_tokens(input_ids)
201
+ current_step_idx = spec_step_idx % self.num_mtp_layers
202
+ return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
203
+ input_ids,
204
+ positions,
205
+ previous_hidden_states,
206
+ inputs_embeds,
207
+ current_step_idx,
208
+ )
209
+
210
+ def compute_logits(
211
+ self,
212
+ hidden_states: torch.Tensor,
213
+ spec_step_idx: int = 0,
214
+ ) -> torch.Tensor:
215
+ current_step_idx = spec_step_idx % self.num_mtp_layers
216
+ mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
217
+ # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply
218
+ # hc_head here so logits are computed from the dense hidden state.
219
+ hidden_states = hidden_states.view(
220
+ -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size
221
+ )
222
+ hidden_states = hc_head(
223
+ hidden_states,
224
+ mtp_layer.hc_head_fn,
225
+ mtp_layer.hc_head_scale,
226
+ mtp_layer.hc_head_base,
227
+ mtp_layer.rms_norm_eps,
228
+ mtp_layer.hc_eps,
229
+ )
230
+ logits = self.logits_processor(
231
+ mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
232
+ )
233
+ return logits
234
+
235
+
236
+@support_torch_compile
237
+class DeepSeekV4MTP(nn.Module):
238
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
239
+ super().__init__()
240
+ self.config = vllm_config.model_config.hf_config
241
+ self.quant_config = vllm_config.quant_config
242
+ self.model = DeepSeekV4MultiTokenPredictor(
243
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
244
+ )
245
+
246
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
247
+ return self.model.embed_input_ids(input_ids)
248
+
249
+ def forward(
250
+ self,
251
+ input_ids: torch.Tensor | None,
252
+ positions: torch.Tensor,
253
+ hidden_states: torch.Tensor,
254
+ intermediate_tensors: IntermediateTensors | None = None,
255
+ inputs_embeds: torch.Tensor | None = None,
256
+ spec_step_idx: int = 0,
257
+ ) -> torch.Tensor:
258
+ hidden_states = self.model(
259
+ input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
260
+ )
261
+ return hidden_states
262
+
263
+ def compute_logits(
264
+ self,
265
+ hidden_states: torch.Tensor,
266
+ spec_step_idx: int = 0,
267
+ ) -> torch.Tensor | None:
268
+ return self.model.compute_logits(hidden_states, spec_step_idx)
269
+
270
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
271
+ # Weight name remapping for checkpoint compatibility.
272
+ # Maps checkpoint weight paths to model parameter paths.
273
+ WEIGHT_NAME_REMAPPING: dict[str, str] = {
274
+ ".emb.tok_emb.weight": ".embed_tokens.weight",
275
+ ".head.weight": ".shared_head.head.weight",
276
+ ".norm.weight": ".shared_head.norm.weight",
277
+ }
278
+
279
+ def _remap_weight_name(name: str) -> str:
280
+ """Remap checkpoint weight names to model parameter names."""
281
+ for old_pattern, new_pattern in WEIGHT_NAME_REMAPPING.items():
282
+ if old_pattern in name:
283
+ name = name.replace(old_pattern, new_pattern)
284
+ return name
285
+
286
+ def _find_mtp_layer_idx(name: str) -> int:
287
+ subnames = name.split(".")
288
+ for subname in subnames:
289
+ try:
290
+ # we return the first encountered integer
291
+ return int(subname)
292
+ except ValueError:
293
+ continue
294
+ return 0
295
+
296
+ stacked_params_mapping = [
297
+ # (param_name, shard_name, shard_id)
298
+ ("gate_up_proj", "w1", 0),
299
+ ("gate_up_proj", "w3", 1),
300
+ ("attn.fused_wqa_wkv", "attn.wq_a", 0),
301
+ ("attn.fused_wqa_wkv", "attn.wkv", 1),
302
+ ]
303
+ params_dict = dict(self.named_parameters())
304
+ loaded_params: set[str] = set()
305
+
306
+ # TP for attention
307
+ tp_size = get_tensor_model_parallel_world_size()
308
+ tp_rank = get_tensor_model_parallel_rank()
309
+ n_head = self.config.num_attention_heads
310
+ n_local_head = n_head // tp_size
311
+ head_rank_start = n_local_head * tp_rank
312
+ head_rank_end = n_local_head * (tp_rank + 1)
313
+
314
+ # Pre-compute expert mapping ONCE.
315
+ first_layer = next(iter(self.model.layers.values()))
316
+ if first_layer.mtp_block.ffn.use_mega_moe:
317
+ expert_mapping = make_deepseek_v4_expert_params_mapping(
318
+ self.config.n_routed_experts
319
+ )
320
+ else:
321
+ expert_mapping = FusedMoE.make_expert_params_mapping(
322
+ self,
323
+ ckpt_gate_proj_name="w1",
324
+ ckpt_down_proj_name="w2",
325
+ ckpt_up_proj_name="w3",
326
+ num_experts=self.config.n_routed_experts,
327
+ )
328
+
329
+ for name, loaded_weight in weights:
330
+ mtp_layer_idx = _find_mtp_layer_idx(name)
331
+ # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to
332
+ # `model.layers.{num_hidden_layers + i}.*` so that
333
+ # get_spec_layer_idx_from_weight_name can identify them.
334
+ name = name.replace(
335
+ f"mtp.{mtp_layer_idx}.",
336
+ f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.",
337
+ )
338
+
339
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
340
+ if spec_layer is None:
341
+ continue
342
+
343
+ name = _remap_weight_name(name)
344
+ name = self._rewrite_spec_layer_name(spec_layer, name)
345
+
346
+ if spec_layer != self.model.mtp_start_layer_idx and ".layers" not in name:
347
+ continue
348
+ if name.endswith(".scale"):
349
+ suffix = (
350
+ ".weight_scale"
351
+ if _EXPERT_SCALE_RE.search(name)
352
+ else ".weight_scale_inv"
353
+ )
354
+ name = name.removesuffix(".scale") + suffix
355
+ for param_name, weight_name, shard_id in stacked_params_mapping:
356
+ # Skip non-stacked layers and experts (experts handled below).
357
+ if ".experts." in name:
358
+ continue
359
+ if weight_name not in name:
360
+ continue
361
+ name = name.replace(weight_name, param_name)
362
+
363
+ param = params_dict[name]
364
+ weight_loader = param.weight_loader
365
+ weight_loader(param, loaded_weight, shard_id)
366
+ loaded_params.add(name)
367
+ break
368
+ else:
369
+ if ".experts." in name:
370
+ # Reinterpret E8M0 scales as uint8 to preserve raw
371
+ # exponent bytes; numeric copy_() would zero them.
372
+ # Mirrors the main DeepseekV4 loader.
373
+ if (
374
+ "weight_scale" in name
375
+ and loaded_weight.dtype == torch.float8_e8m0fnu
376
+ ):
377
+ loaded_weight = loaded_weight.view(torch.uint8)
378
+ for mapping in expert_mapping:
379
+ param_name, weight_name, expert_id, shard_id = mapping
380
+ if weight_name not in name:
381
+ continue
382
+ name_mapped = name.replace(weight_name, param_name)
383
+ param = params_dict[name_mapped]
384
+ # We should ask the weight loader to return success or not
385
+ # here since otherwise we may skip experts with other
386
+ # available replicas.
387
+ weight_loader = typing.cast(
388
+ Callable[..., bool], param.weight_loader
389
+ )
390
+ success = weight_loader(
391
+ param,
392
+ loaded_weight,
393
+ name_mapped,
394
+ shard_id=shard_id,
395
+ expert_id=expert_id,
396
+ return_success=True,
397
+ )
398
+ if success:
399
+ name = name_mapped
400
+ loaded_params.add(name_mapped)
401
+ break
402
+ continue
403
+ elif "attn_sink" in name:
404
+ narrow_weight = loaded_weight[head_rank_start:head_rank_end]
405
+ n = narrow_weight.shape[0]
406
+ params_dict[name][:n].copy_(narrow_weight)
407
+ loaded_params.add(name)
408
+ continue
409
+ else:
410
+ if ".shared_experts.w2" in name:
411
+ name = name.replace(
412
+ ".shared_experts.w2", ".shared_experts.down_proj"
413
+ )
414
+ if name.endswith(".ffn.gate.bias"):
415
+ name = name.replace(".bias", ".e_score_correction_bias")
416
+ param = params_dict[name]
417
+ weight_loader = getattr(
418
+ param, "weight_loader", default_weight_loader
419
+ )
420
+ weight_loader(param, loaded_weight)
421
+ loaded_params.add(name)
422
+ continue
423
+
424
+ loaded_layers: set[int] = set()
425
+ for param_name in loaded_params:
426
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
427
+ if spec_layer is not None:
428
+ loaded_layers.add(spec_layer)
429
+ for layer_idx in range(
430
+ self.model.mtp_start_layer_idx,
431
+ self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
432
+ ):
433
+ if layer_idx not in loaded_layers:
434
+ raise ValueError(
435
+ f"MTP speculative decoding layer {layer_idx} weights "
436
+ f"missing from checkpoint. The checkpoint may have "
437
+ f"been quantized without including the MTP layers. "
438
+ f"Use a checkpoint that includes MTP layer weights, "
439
+ f"or disable speculative decoding."
440
+ )
441
+ self.finalize_mega_moe_weights()
442
+ logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
443
+ return loaded_params
444
+
445
+ def finalize_mega_moe_weights(self) -> None:
446
+ for layer in self.model.layers.values():
447
+ layer.mtp_block.ffn.finalize_mega_moe_weights()
448
+
449
+ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
450
+ """
451
+ Rewrite the weight name to match the format of the original model.
452
+ Add .mtp_block for modules in transformer layer block for spec layer
453
+ and rename shared layer weights to be top level.
454
+ """
455
+ spec_layer_weight_names = [
456
+ "embed_tokens",
457
+ "enorm",
458
+ "hnorm",
459
+ "h_proj",
460
+ "e_proj",
461
+ "shared_head",
462
+ "hc_head_fn",
463
+ "hc_head_base",
464
+ "hc_head_scale",
465
+ ]
466
+ shared_weight_names = ["embed_tokens"]
467
+ spec_layer_weight = False
468
+ shared_weight = False
469
+ for weight_name in spec_layer_weight_names:
470
+ if weight_name in name:
471
+ spec_layer_weight = True
472
+ if weight_name in shared_weight_names:
473
+ shared_weight = True
474
+ break
475
+ if not spec_layer_weight:
476
+ # treat rest weights as weights for transformer layer block
477
+ name = name.replace(
478
+ f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
479
+ )
480
+ elif shared_weight:
481
+ # treat shared weights as top level weights
482
+ name = name.replace(f"model.layers.{spec_layer}.", "model.")
483
+ return name
vllm/model_executor/models/registry.py
+2
@@ -96,6 +96,7 @@ _TEXT_GENERATION_MODELS = {
96
"DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
97
"DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
98
"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
99
+ "DeepseekV4ForCausalLM": ("deepseek_v4", "DeepseekV4ForCausalLM"),
100
"Dots1ForCausalLM": ("dots1", "Dots1ForCausalLM"),
101
"Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"),
102
"Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
@@ -586,6 +587,7 @@ _SPECULATIVE_DECODING_MODELS = {
587
"Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
588
"EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
589
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
590
+ "DeepSeekV4MTPModel": ("deepseek_v4_mtp", "DeepSeekV4MTP"),
591
"ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
592
"ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
593
"Exaone4_5_MTP": ("exaone4_5_mtp", "Exaone4_5_MTP"),
vllm/reasoning/__init__.py
+4
@@ -28,6 +28,10 @@ _REASONING_PARSERS_TO_REGISTER = {
28
"deepseek_v3_reasoning_parser",
29
"DeepSeekV3ReasoningParser",
30
),
31
+ "deepseek_v4": (
32
+ "deepseek_v3_reasoning_parser",
33
+ "DeepSeekV3ReasoningParser",
34
+ ),
35
"ernie45": (
36
"ernie45_reasoning_parser",
37
"Ernie45ReasoningParser",
vllm/renderers/deepseek_v4.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.config import VllmConfig
5
+from vllm.entrypoints.chat_utils import (
6
+ ChatCompletionMessageParam,
7
+ ConversationMessage,
8
+ parse_chat_messages,
9
+ parse_chat_messages_async,
10
+)
11
+from vllm.logger import init_logger
12
+from vllm.tokenizers.deepseek_v4 import DeepseekV4Tokenizer
13
+from vllm.utils.async_utils import make_async
14
+
15
+from .base import BaseRenderer
16
+from .inputs import DictPrompt
17
+from .inputs.preprocess import parse_dec_only_prompt
18
+from .params import ChatParams
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+class DeepseekV4Renderer(BaseRenderer[DeepseekV4Tokenizer]):
24
+ def __init__(
25
+ self,
26
+ config: VllmConfig,
27
+ tokenizer: DeepseekV4Tokenizer | None,
28
+ ) -> None:
29
+ super().__init__(config, tokenizer)
30
+
31
+ self._apply_chat_template_async = make_async(
32
+ self._apply_chat_template, executor=self._executor
33
+ )
34
+
35
+ def _apply_chat_template(self, *args, **kwargs):
36
+ return self.get_tokenizer().apply_chat_template(*args, **kwargs)
37
+
38
+ def render_messages(
39
+ self,
40
+ messages: list[ChatCompletionMessageParam],
41
+ params: ChatParams,
42
+ ) -> tuple[list[ConversationMessage], DictPrompt]:
43
+ conversation, mm_data, mm_uuids = parse_chat_messages(
44
+ messages,
45
+ self.model_config,
46
+ content_format="string",
47
+ media_io_kwargs=params.media_io_kwargs,
48
+ mm_processor_kwargs=params.mm_processor_kwargs,
49
+ )
50
+
51
+ prompt_raw = self._apply_chat_template(
52
+ conversation=conversation,
53
+ messages=messages,
54
+ **params.get_apply_chat_template_kwargs(),
55
+ )
56
+
57
+ prompt = parse_dec_only_prompt(prompt_raw)
58
+ if mm_data is not None:
59
+ prompt["multi_modal_data"] = mm_data
60
+ if mm_uuids is not None:
61
+ prompt["multi_modal_uuids"] = mm_uuids
62
+
63
+ return conversation, prompt
64
+
65
+ async def render_messages_async(
66
+ self,
67
+ messages: list[ChatCompletionMessageParam],
68
+ params: ChatParams,
69
+ ) -> tuple[list[ConversationMessage], DictPrompt]:
70
+ conversation, mm_data, mm_uuids = await parse_chat_messages_async(
71
+ messages,
72
+ self.model_config,
73
+ content_format="string",
74
+ media_io_kwargs=params.media_io_kwargs,
75
+ mm_processor_kwargs=params.mm_processor_kwargs,
76
+ )
77
+
78
+ prompt_raw = await self._apply_chat_template_async(
79
+ conversation=conversation,
80
+ messages=messages,
81
+ **params.get_apply_chat_template_kwargs(),
82
+ )
83
+
84
+ prompt = parse_dec_only_prompt(prompt_raw)
85
+ if mm_data is not None:
86
+ prompt["multi_modal_data"] = mm_data
87
+ if mm_uuids is not None:
88
+ prompt["multi_modal_uuids"] = mm_uuids
89
+
90
+ return conversation, prompt
vllm/renderers/registry.py
+1
@@ -21,6 +21,7 @@ logger = init_logger(__name__)
21
22
_VLLM_RENDERERS = {
23
"deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"),
24
+ "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"),
25
"hf": ("hf", "HfRenderer"),
26
"grok2": ("grok2", "Grok2Renderer"),
27
"kimi_audio": ("hf", "HfRenderer"),
vllm/tokenizers/deepseek_v4.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copy
4
+from typing import Any
5
+
6
+from transformers import PreTrainedTokenizerFast
7
+
8
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
9
+
10
+from .deepseek_v4_encoding import encode_messages
11
+from .hf import HfTokenizer, get_cached_tokenizer
12
+from .protocol import TokenizerLike
13
+
14
+
15
+def get_deepseek_v4_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
16
+ """
17
+ Wraps a tokenizer to use the custom DeepSeek V4 chat template encoding.
18
+ """
19
+ dsv4_tokenizer = copy.copy(tokenizer)
20
+
21
+ added_vocab = tokenizer.get_added_vocab()
22
+ added_vocab_size = len(added_vocab)
23
+ tokenizer_vocab_size = tokenizer.vocab_size
24
+
25
+ class _DeepseekV4Tokenizer(tokenizer.__class__): # type: ignore
26
+ def apply_chat_template(
27
+ self,
28
+ messages: list["ChatCompletionMessageParam"],
29
+ tools: list[dict[str, Any]] | None = None,
30
+ **kwargs,
31
+ ) -> str | list[int]:
32
+ thinking = kwargs.get("thinking", False)
33
+ enable_thinking = kwargs.get("enable_thinking", False)
34
+ thinking = thinking or enable_thinking
35
+ thinking_mode = "thinking" if thinking else "chat"
36
+
37
+ conversation = kwargs.get("conversation", messages)
38
+ messages = conversation.copy()
39
+ if tools is not None and len(tools) > 0:
40
+ messages.insert(0, {"role": "system"})
41
+ messages[0]["tools"] = tools # type: ignore[typeddict-unknown-key]
42
+
43
+ # The V4 reference currently accepts only "max", "high", or None.
44
+ reasoning_effort = kwargs.get("reasoning_effort")
45
+ if reasoning_effort not in ("max", "high"):
46
+ reasoning_effort = None
47
+
48
+ encode_config = dict(
49
+ thinking_mode=thinking_mode,
50
+ drop_thinking=kwargs.get("drop_thinking", True),
51
+ reasoning_effort=reasoning_effort,
52
+ )
53
+
54
+ prompt_str = encode_messages(messages, **encode_config) # type: ignore
55
+
56
+ if kwargs.get("tokenize", True):
57
+ tokenizer_kwargs = {
58
+ k: kwargs[k] for k in ("truncation", "max_length") if k in kwargs
59
+ }
60
+ return self.encode(
61
+ prompt_str,
62
+ add_special_tokens=False,
63
+ **tokenizer_kwargs,
64
+ )
65
+
66
+ return prompt_str
67
+
68
+ def num_special_tokens_to_add(self) -> int:
69
+ return len(self.encode(""))
70
+
71
+ def __len__(self) -> int:
72
+ return tokenizer_vocab_size + added_vocab_size
73
+
74
+ def get_added_vocab(self) -> dict[str, int]:
75
+ return added_vocab.copy()
76
+
77
+ def __reduce__(self):
78
+ return get_deepseek_v4_tokenizer, (tokenizer,)
79
+
80
+ _DeepseekV4Tokenizer.__name__ = f"DSV4{tokenizer.__class__.__name__}"
81
+
82
+ dsv4_tokenizer.__class__ = _DeepseekV4Tokenizer
83
+ return dsv4_tokenizer
84
+
85
+
86
+class DeepseekV4Tokenizer(TokenizerLike):
87
+ @classmethod
88
+ def from_pretrained(cls, *args, **kwargs) -> HfTokenizer:
89
+ tokenizer = PreTrainedTokenizerFast.from_pretrained(*args, **kwargs)
90
+ return get_cached_tokenizer(get_deepseek_v4_tokenizer(tokenizer))
vllm/tokenizers/deepseek_v4_encoding.py
new
+757
@@ -0,0 +1,757 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa
4
+# fmt: off
5
+
6
+"""
7
+DeepSeek-V4 Encoding
8
+
9
+A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages
10
+with tool calling, thinking mode, and quick instruction task support.
11
+"""
12
+
13
+from typing import Any, Dict, List, Union, Optional, Tuple
14
+import copy
15
+import json
16
+
17
+import regex as re
18
+
19
+# ============================================================
20
+# Special Tokens
21
+# ============================================================
22
+
23
+bos_token: str = "<｜begin▁of▁sentence｜>"
24
+eos_token: str = "<｜end▁of▁sentence｜>"
25
+thinking_start_token: str = "<think>"
26
+thinking_end_token: str = "</think>"
27
+dsml_token: str = "｜DSML｜"
28
+
29
+USER_SP_TOKEN = "<｜User｜>"
30
+ASSISTANT_SP_TOKEN = "<｜Assistant｜>"
31
+LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>"
32
+
33
+# Task special tokens for internal classification tasks
34
+DS_TASK_SP_TOKENS = {
35
+ "action": "<｜action｜>",
36
+ "query": "<｜query｜>",
37
+ "authority": "<｜authority｜>",
38
+ "domain": "<｜domain｜>",
39
+ "title": "<｜title｜>",
40
+ "read_url": "<｜read_url｜>",
41
+}
42
+VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())
43
+
44
+# ============================================================
45
+# Templates
46
+# ============================================================
47
+
48
+system_msg_template: str = "{content}"
49
+user_msg_template: str = "{content}"
50
+latest_reminder_msg_template: str = "{content}"
51
+assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token
52
+assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"
53
+thinking_template: str = "{reasoning}"
54
+
55
+response_format_template: str = (
56
+ "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
57
+)
58
+tool_call_template: str = (
59
+ "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n</{dsml_token}invoke>"
60
+)
61
+tool_calls_template = (
62
+ "<{dsml_token}{tc_block_name}>\n{tool_calls}\n</{dsml_token}{tc_block_name}>"
63
+)
64
+tool_calls_block_name: str = "tool_calls"
65
+
66
+tool_output_template: str = (
67
+ "<tool_result>{content}</tool_result>"
68
+)
69
+
70
+REASONING_EFFORT_MAX = (
71
+ "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
72
+ "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n"
73
+ "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"
74
+)
75
+
76
+TOOLS_TEMPLATE = """## Tools
77
+
78
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
79
+
80
+<{dsml_token}tool_calls>
81
+<{dsml_token}invoke name="$TOOL_NAME">
82
+<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
83
+...
84
+</{dsml_token}invoke>
85
+<{dsml_token}invoke name="$TOOL_NAME2">
86
+...
87
+</{dsml_token}invoke>
88
+</{dsml_token}tool_calls>
89
+
90
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
91
+
92
+If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
93
+
94
+Otherwise, output directly after {thinking_end_token} with tool calls or final response.
95
+
96
+### Available Tool Schemas
97
+
98
+{tool_schemas}
99
+
100
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
101
+"""
102
+
103
+# ============================================================
104
+# Utility Functions
105
+# ============================================================
106
+
107
+def to_json(value: Any) -> str:
108
+ """Serialize a value to JSON string."""
109
+ try:
110
+ return json.dumps(value, ensure_ascii=False)
111
+ except Exception:
112
+ return json.dumps(value, ensure_ascii=True)
113
+
114
+
115
+def tools_from_openai_format(tools):
116
+ """Extract function definitions from OpenAI-format tool list."""
117
+ return [tool["function"] for tool in tools]
118
+
119
+
120
+def tool_calls_from_openai_format(tool_calls):
121
+ """Convert OpenAI-format tool calls to internal format."""
122
+ return [
123
+ {
124
+ "name": tool_call["function"]["name"],
125
+ "arguments": tool_call["function"]["arguments"],
126
+ }
127
+ for tool_call in tool_calls
128
+ ]
129
+
130
+
131
+def tool_calls_to_openai_format(tool_calls):
132
+ """Convert internal tool calls to OpenAI format."""
133
+ return [
134
+ {
135
+ "type": "function",
136
+ "function": {
137
+ "name": tool_call["name"],
138
+ "arguments": tool_call["arguments"],
139
+ }
140
+ }
141
+ for tool_call in tool_calls
142
+ ]
143
+
144
+
145
+def encode_arguments_to_dsml(tool_call: Dict[str, Any]) -> str:
146
+ """
147
+ Encode tool call arguments into DSML parameter format.
148
+
149
+ Args:
150
+ tool_call: Dict with "name" and "arguments" keys.
151
+
152
+ Returns:
153
+ DSML-formatted parameter string.
154
+ """
155
+ p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>'
156
+ P_dsml_strs = []
157
+
158
+ if isinstance(tool_call["arguments"], str):
159
+ arguments = json.loads(tool_call["arguments"])
160
+ else:
161
+ arguments = tool_call["arguments"]
162
+
163
+ for k, v in arguments.items():
164
+ p_dsml_str = p_dsml_template.format(
165
+ dsml_token=dsml_token,
166
+ key=k,
167
+ is_str="true" if isinstance(v, str) else "false",
168
+ value=v if isinstance(v, str) else to_json(v),
169
+ )
170
+ P_dsml_strs.append(p_dsml_str)
171
+
172
+ return "\n".join(P_dsml_strs)
173
+
174
+
175
+def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
176
+ """
177
+ Decode DSML parameters back to a tool call dict.
178
+
179
+ Args:
180
+ tool_name: Name of the tool.
181
+ tool_args: Dict mapping param_name -> (value, is_string_flag).
182
+
183
+ Returns:
184
+ Dict with "name" and "arguments" (JSON string) keys.
185
+ """
186
+ def _decode_value(key: str, value: str, string: str):
187
+ if string == "true":
188
+ value = to_json(value)
189
+ return f"{to_json(key)}: {value}"
190
+
191
+ tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"
192
+ return dict(name=tool_name, arguments=tool_args_json)
193
+
194
+
195
+def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
196
+ """
197
+ Render tool schemas into the system prompt format.
198
+
199
+ Args:
200
+ tools: List of tool schema dicts (each with name, description, parameters).
201
+
202
+ Returns:
203
+ Formatted tools section string.
204
+ """
205
+ tools_json = [to_json(t) for t in tools]
206
+
207
+ return TOOLS_TEMPLATE.format(
208
+ tool_schemas="\n".join(tools_json),
209
+ dsml_token=dsml_token,
210
+ thinking_start_token=thinking_start_token,
211
+ thinking_end_token=thinking_end_token,
212
+ )
213
+
214
+
215
+def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
216
+ """Find the index of the last user/developer message."""
217
+ last_user_index = -1
218
+ for idx in range(len(messages) - 1, -1, -1):
219
+ if messages[idx].get("role") in ["user", "developer"]:
220
+ last_user_index = idx
221
+ break
222
+ return last_user_index
223
+
224
+
225
+# ============================================================
226
+# Message Rendering
227
+# ============================================================
228
+
229
+def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str:
230
+ """
231
+ Render a single message at the given index into its encoded string form.
232
+
233
+ This is the core function that converts each message in the conversation
234
+ into the DeepSeek-V4 format.
235
+
236
+ Args:
237
+ index: Index of the message to render.
238
+ messages: Full list of messages in the conversation.
239
+ thinking_mode: Either "chat" or "thinking".
240
+ drop_thinking: Whether to drop reasoning content from earlier turns.
241
+ reasoning_effort: Optional reasoning effort level ("max", "high", or None).
242
+
243
+ Returns:
244
+ Encoded string for this message.
245
+ """
246
+ assert 0 <= index < len(messages)
247
+ assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"
248
+
249
+ prompt = ""
250
+ msg = messages[index]
251
+ last_user_idx = find_last_user_index(messages)
252
+
253
+ role = msg.get("role")
254
+ content = msg.get("content")
255
+ tools = msg.get("tools")
256
+ response_format = msg.get("response_format")
257
+ tool_calls = msg.get("tool_calls")
258
+ reasoning = msg.get("reasoning")
259
+ wo_eos = msg.get("wo_eos", False)
260
+
261
+ if tools:
262
+ tools = tools_from_openai_format(tools)
263
+ if tool_calls:
264
+ tool_calls = tool_calls_from_openai_format(tool_calls)
265
+
266
+ # Reasoning effort prefix (only at index 0 in thinking mode with max effort)
267
+ assert reasoning_effort in ['max', None, 'high'], f"Invalid reasoning effort: {reasoning_effort}"
268
+ if index == 0 and thinking_mode == "thinking" and reasoning_effort == 'max':
269
+ prompt += REASONING_EFFORT_MAX
270
+
271
+ if role == "system":
272
+ prompt += system_msg_template.format(content=content or "")
273
+ if tools:
274
+ prompt += "\n\n" + render_tools(tools)
275
+ if response_format:
276
+ prompt += "\n\n" + response_format_template.format(schema=to_json(response_format))
277
+
278
+ elif role == "developer":
279
+ assert content, f"Invalid message for role `{role}`: {msg}"
280
+
281
+ content_developer = USER_SP_TOKEN
282
+ content_developer += content
283
+
284
+ if tools:
285
+ content_developer += "\n\n" + render_tools(tools)
286
+ if response_format:
287
+ content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format))
288
+
289
+ prompt += user_msg_template.format(content=content_developer)
290
+
291
+ elif role == "user":
292
+ prompt += USER_SP_TOKEN
293
+
294
+ # Handle content blocks (tool results mixed with text)
295
+ content_blocks = msg.get("content_blocks")
296
+ if content_blocks:
297
+ parts = []
298
+ for block in content_blocks:
299
+ block_type = block.get("type")
300
+ if block_type == "text":
301
+ parts.append(block.get("text", ""))
302
+ elif block_type == "tool_result":
303
+ tool_content = block.get("content", "")
304
+ if isinstance(tool_content, list):
305
+ text_parts = []
306
+ for b in tool_content:
307
+ if b.get("type") == "text":
308
+ text_parts.append(b.get("text", ""))
309
+ else:
310
+ text_parts.append(f"[Unsupported {b.get('type')}]")
311
+ tool_content = "\n\n".join(text_parts)
312
+ parts.append(tool_output_template.format(content=tool_content))
313
+ else:
314
+ parts.append(f"[Unsupported {block_type}]")
315
+ prompt += "\n\n".join(parts)
316
+ else:
317
+ prompt += content or ""
318
+
319
+ elif role == "latest_reminder":
320
+ prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
321
+
322
+ elif role == "tool":
323
+ raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()")
324
+
325
+ elif role == "assistant":
326
+ thinking_part = ""
327
+ tc_content = ""
328
+
329
+ if tool_calls:
330
+ tc_list = [
331
+ tool_call_template.format(
332
+ dsml_token=dsml_token,
333
+ name=tc.get("name"),
334
+ arguments=encode_arguments_to_dsml(tc)
335
+ )
336
+ for tc in tool_calls
337
+ ]
338
+ tc_content += '\n\n' + tool_calls_template.format(
339
+ dsml_token=dsml_token,
340
+ tool_calls="\n".join(tc_list),
341
+ tc_block_name=tool_calls_block_name,
342
+ )
343
+
344
+ summary_content = content or ""
345
+ reasoning = reasoning or ""
346
+
347
+ # Check if previous message has a task - if so, this is a task output (no thinking)
348
+ prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
349
+
350
+ if thinking_mode == "thinking" and not prev_has_task:
351
+ if not drop_thinking or index > last_user_idx:
352
+ thinking_part = thinking_template.format(reasoning=reasoning) + thinking_end_token
353
+ else:
354
+ thinking_part = ""
355
+
356
+ if wo_eos:
357
+ prompt += assistant_msg_wo_eos_template.format(
358
+ reasoning=thinking_part,
359
+ content=summary_content,
360
+ tool_calls=tc_content,
361
+ )
362
+ else:
363
+ prompt += assistant_msg_template.format(
364
+ reasoning=thinking_part,
365
+ content=summary_content,
366
+ tool_calls=tc_content,
367
+ )
368
+ else:
369
+ raise NotImplementedError(f"Unknown role: {role}")
370
+
371
+ # Append transition tokens based on what follows
372
+ if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:
373
+ return prompt
374
+
375
+ task = messages[index].get("task")
376
+ if task is not None:
377
+ # Task special token for internal classification tasks
378
+ assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
379
+ task_sp_token = DS_TASK_SP_TOKENS[task]
380
+
381
+ if task != "action":
382
+ # Non-action tasks: append task sp token directly after the message
383
+ prompt += task_sp_token
384
+ else:
385
+ # Action task: append Assistant + thinking token + action sp token
386
+ prompt += ASSISTANT_SP_TOKEN
387
+ prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token
388
+ prompt += task_sp_token
389
+
390
+ elif messages[index].get("role") in ["user", "developer"]:
391
+ # Normal generation: append Assistant + thinking token
392
+ prompt += ASSISTANT_SP_TOKEN
393
+ if not drop_thinking and thinking_mode == "thinking":
394
+ prompt += thinking_start_token
395
+ elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:
396
+ prompt += thinking_start_token
397
+ else:
398
+ prompt += thinking_end_token
399
+
400
+ return prompt
401
+
402
+
403
+# ============================================================
404
+# Preprocessing
405
+# ============================================================
406
+
407
+def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
408
+ """
409
+ Merge tool messages into the preceding user message using content_blocks format.
410
+
411
+ DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
412
+ are encoded as <tool_result> blocks within user messages.
413
+
414
+ This function converts a standard OpenAI-format conversation (with separate
415
+ "tool" role messages) into V4 format where tool results are merged into
416
+ user messages.
417
+
418
+ Args:
419
+ messages: List of message dicts in OpenAI format.
420
+
421
+ Returns:
422
+ Processed message list with tool messages merged into user messages.
423
+ """
424
+ merged: List[Dict[str, Any]] = []
425
+
426
+ for msg in messages:
427
+ msg = copy.deepcopy(msg)
428
+ role = msg.get("role")
429
+
430
+ if role == "tool":
431
+ # Convert tool message to a user message with tool_result block
432
+ tool_block = {
433
+ "type": "tool_result",
434
+ "tool_use_id": msg.get("tool_call_id", ""),
435
+ "content": msg.get("content", ""),
436
+ }
437
+ # Merge into previous message if it's already a user (merged tool)
438
+ if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:
439
+ merged[-1]["content_blocks"].append(tool_block)
440
+ else:
441
+ merged.append({
442
+ "role": "user",
443
+ "content_blocks": [tool_block],
444
+ })
445
+ elif role == "user":
446
+ text_block = {"type": "text", "text": msg.get("content", "")}
447
+ if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None:
448
+ merged[-1]["content_blocks"].append(text_block)
449
+ else:
450
+ new_msg = {
451
+ "role": "user",
452
+ "content": msg.get("content", ""),
453
+ "content_blocks": [text_block],
454
+ }
455
+ # Preserve extra fields (task, wo_eos, mask, etc.)
456
+ for key in ("task", "wo_eos", "mask"):
457
+ if key in msg:
458
+ new_msg[key] = msg[key]
459
+ merged.append(new_msg)
460
+ else:
461
+ merged.append(msg)
462
+
463
+ return merged
464
+
465
+
466
+def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
467
+ """
468
+ Sort tool_result blocks within user messages by the order of tool_calls
469
+ in the preceding assistant message.
470
+
471
+ Args:
472
+ messages: Preprocessed message list (after merge_tool_messages).
473
+
474
+ Returns:
475
+ Message list with sorted tool result blocks.
476
+ """
477
+ last_tool_call_order: Dict[str, int] = {}
478
+
479
+ for msg in messages:
480
+ role = msg.get("role")
481
+ if role == "assistant" and msg.get("tool_calls"):
482
+ last_tool_call_order = {}
483
+ for idx, tc in enumerate(msg["tool_calls"]):
484
+ tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
485
+ if tc_id:
486
+ last_tool_call_order[tc_id] = idx
487
+
488
+ elif role == "user" and msg.get("content_blocks"):
489
+ tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]
490
+ if len(tool_blocks) > 1 and last_tool_call_order:
491
+ sorted_blocks = sorted(
492
+ tool_blocks,
493
+ key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0)
494
+ )
495
+ sorted_idx = 0
496
+ new_blocks = []
497
+ for block in msg["content_blocks"]:
498
+ if block.get("type") == "tool_result":
499
+ new_blocks.append(sorted_blocks[sorted_idx])
500
+ sorted_idx += 1
501
+ else:
502
+ new_blocks.append(block)
503
+ msg["content_blocks"] = new_blocks
504
+
505
+ return messages
506
+
507
+
508
+# ============================================================
509
+# Main Encoding Function
510
+# ============================================================
511
+
512
+def encode_messages(
513
+ messages: List[Dict[str, Any]],
514
+ thinking_mode: str,
515
+ context: Optional[List[Dict[str, Any]]] = None,
516
+ drop_thinking: bool = True,
517
+ add_default_bos_token: bool = True,
518
+ reasoning_effort: Optional[str] = None,
519
+) -> str:
520
+ """
521
+ Encode a list of messages into the DeepSeek-V4 prompt format.
522
+
523
+ This is the main entry point for encoding conversations. It handles:
524
+ - BOS token insertion
525
+ - Thinking mode with optional reasoning content dropping
526
+ - Tool message merging into user messages
527
+ - Multi-turn conversation context
528
+
529
+ Args:
530
+ messages: List of message dicts to encode.
531
+ thinking_mode: Either "chat" or "thinking".
532
+ context: Optional preceding context messages (already encoded prefix).
533
+ drop_thinking: If True, drop reasoning from earlier assistant turns
534
+ (only keep reasoning for messages after the last user message).
535
+ add_default_bos_token: Whether to prepend BOS token at conversation start.
536
+ reasoning_effort: Optional reasoning effort level ("max", "high", or None).
537
+
538
+ Returns:
539
+ The encoded prompt string.
540
+ """
541
+ context = context if context else []
542
+
543
+ # Preprocess: merge tool messages and sort tool results
544
+ messages = merge_tool_messages(messages)
545
+ messages = sort_tool_results_by_call_order(context + messages)[len(context):]
546
+ if context:
547
+ context = merge_tool_messages(context)
548
+ context = sort_tool_results_by_call_order(context)
549
+
550
+ full_messages = context + messages
551
+
552
+ prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
553
+
554
+ # Resolve drop_thinking: if any message has tools defined, don't drop thinking
555
+ effective_drop_thinking = drop_thinking
556
+ if any(m.get("tools") for m in full_messages):
557
+ effective_drop_thinking = False
558
+
559
+ if thinking_mode == "thinking" and effective_drop_thinking:
560
+ full_messages = _drop_thinking_messages(full_messages)
561
+ # After dropping, recalculate how many messages to render
562
+ # (context may have shrunk too)
563
+ num_to_render = len(full_messages) - len(_drop_thinking_messages(context))
564
+ context_len = len(full_messages) - num_to_render
565
+ else:
566
+ num_to_render = len(messages)
567
+ context_len = len(context)
568
+
569
+ for idx in range(num_to_render):
570
+ prompt += render_message(
571
+ idx + context_len,
572
+ full_messages,
573
+ thinking_mode=thinking_mode,
574
+ drop_thinking=effective_drop_thinking,
575
+ reasoning_effort=reasoning_effort,
576
+ )
577
+
578
+ return prompt
579
+
580
+
581
+def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
582
+ """
583
+ Drop reasoning and non-essential messages before the last user message.
584
+
585
+ Behavior:
586
+ - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
587
+ - Messages at or after the last user index are always kept.
588
+ - Assistant messages before the last user get reasoning removed.
589
+ - Developer messages before the last user are dropped entirely.
590
+ """
591
+ last_user_idx = find_last_user_index(messages)
592
+ result = []
593
+ keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
594
+
595
+ for idx, msg in enumerate(messages):
596
+ role = msg.get("role")
597
+ if role in keep_roles or idx >= last_user_idx:
598
+ result.append(msg)
599
+ elif role == "assistant":
600
+ msg = copy.copy(msg)
601
+ msg.pop("reasoning", None)
602
+ result.append(msg)
603
+ # developer and other roles before last_user_idx are dropped
604
+
605
+ return result
606
+
607
+
608
+# ============================================================
609
+# Parsing (Decoding model output)
610
+# ============================================================
611
+
612
+def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
613
+ """
614
+ Read text from index until one of the stop strings is found.
615
+
616
+ Returns:
617
+ Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
618
+ """
619
+ min_pos = len(text)
620
+ matched_stop = None
621
+
622
+ for s in stop:
623
+ pos = text.find(s, index)
624
+ if pos != -1 and pos < min_pos:
625
+ min_pos = pos
626
+ matched_stop = s
627
+
628
+ if matched_stop:
629
+ content = text[index:min_pos]
630
+ return min_pos + len(matched_stop), content, matched_stop
631
+ else:
632
+ content = text[index:]
633
+ return len(text), content, None
634
+
635
+
636
+def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
637
+ """
638
+ Parse DSML tool calls from text starting at the given index.
639
+
640
+ Args:
641
+ index: Starting position in text.
642
+ text: The full text to parse.
643
+
644
+ Returns:
645
+ Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
646
+ Each tool call dict has "name" and "arguments" keys.
647
+ """
648
+ tool_calls: List[Dict[str, Any]] = []
649
+ stop_token = None
650
+ tool_calls_end_token = f"</{dsml_token}{tool_calls_block_name}>"
651
+
652
+ while index < len(text):
653
+ index, content_before, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token])
654
+ if content_before != ">\n":
655
+ raise ValueError(f"Tool call format error: expected '>\\n' but got '{content_before}'")
656
+
657
+ if stop_token == tool_calls_end_token:
658
+ break
659
+
660
+ if stop_token is None:
661
+ raise ValueError("Missing special token in tool calls")
662
+
663
+ index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])
664
+
665
+ p_tool_name = re.findall(r'^\s*name="(.*?)">\n$', tool_name_content, flags=re.DOTALL)
666
+ if len(p_tool_name) != 1:
667
+ raise ValueError(f"Tool name format error: '{tool_name_content}'")
668
+ tool_name = p_tool_name[0]
669
+
670
+ tool_args: Dict[str, Tuple[str, str]] = {}
671
+ while stop_token == f"<{dsml_token}parameter":
672
+ index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"])
673
+
674
+ param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL)
675
+ if len(param_kv) != 1:
676
+ raise ValueError(f"Parameter format error: '{param_content}'")
677
+ param_name, string, param_value = param_kv[0]
678
+
679
+ if param_name in tool_args:
680
+ raise ValueError(f"Duplicate parameter name: '{param_name}'")
681
+ tool_args[param_name] = (param_value, string)
682
+
683
+ index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])
684
+ if content != ">\n":
685
+ raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'")
686
+
687
+ tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
688
+ tool_calls.append(tool_call)
689
+
690
+ return index, stop_token, tool_calls
691
+
692
+
693
+def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
694
+ """
695
+ Parse a model completion text into a structured assistant message.
696
+
697
+ This function takes the raw text output from the model (a single assistant turn)
698
+ and extracts:
699
+ - reasoning (thinking block)
700
+ - content (summary/response)
701
+ - tool_calls (if any)
702
+
703
+ NOTE: This function is designed to parse only correctly formatted strings and
704
+ will raise ValueError for malformed output.
705
+
706
+ Args:
707
+ text: The raw completion text (including EOS token).
708
+ thinking_mode: Either "chat" or "thinking".
709
+
710
+ Returns:
711
+ Dict with keys: "role", "content", "reasoning", "tool_calls".
712
+ tool_calls are in OpenAI format.
713
+ """
714
+ summary_content, reasoning = "", ""
715
+ tool_calls: List[Dict[str, str]] = []
716
+ index, stop_token = 0, None
717
+ tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"
718
+
719
+ is_thinking = thinking_mode == "thinking"
720
+ is_tool_calling = False
721
+
722
+ if is_thinking:
723
+ index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token])
724
+ reasoning = content_delta
725
+ if stop_token != thinking_end_token:
726
+ raise ValueError("Invalid thinking format: missing </think>")
727
+
728
+ index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token])
729
+ summary_content = content_delta
730
+ if stop_token == tool_calls_start_token:
731
+ is_tool_calling = True
732
+ else:
733
+ if stop_token != eos_token:
734
+ raise ValueError("Invalid format: missing EOS token")
735
+
736
+ if is_tool_calling:
737
+ index, stop_token, tool_calls = parse_tool_calls(index, text)
738
+
739
+ index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
740
+ if tool_ends_text:
741
+ raise ValueError("Unexpected content after tool calls")
742
+
743
+ if len(text) != index or stop_token not in [eos_token, None]:
744
+ raise ValueError("Unexpected content at end")
745
+
746
+ for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]:
747
+ if sp_token in summary_content or sp_token in reasoning:
748
+ raise ValueError(f"Unexpected special token '{sp_token}' in content")
749
+
750
+ return {
751
+ "role": "assistant",
752
+ "content": summary_content,
753
+ "reasoning": reasoning,
754
+ "tool_calls": tool_calls_to_openai_format(tool_calls)
755
+ }
756
+
757
+# fmt: on
vllm/tokenizers/registry.py
+1
@@ -42,6 +42,7 @@ _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl"}
42
43
_VLLM_TOKENIZERS = {
44
"deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),
45
+ "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"),
46
"grok2": ("grok2", "Grok2Tokenizer"),
47
"hf": ("hf", "CachedHfTokenizer"),
48
"kimi_audio": ("kimi_audio", "KimiAudioTokenizer"),
vllm/tool_parsers/__init__.py
+4
@@ -34,6 +34,10 @@ _TOOL_PARSERS_TO_REGISTER = {
34
"deepseekv32_tool_parser",
35
"DeepSeekV32ToolParser",
36
),
37
+ "deepseek_v4": (
38
+ "deepseekv4_tool_parser",
39
+ "DeepSeekV4ToolParser",
40
+ ),
41
"ernie45": (
42
"ernie45_tool_parser",
43
"Ernie45ToolParser",
vllm/tool_parsers/deepseekv32_tool_parser.py
+8
-5
@@ -46,21 +46,24 @@ class DeepSeekV32ToolParser(ToolParser):
46
</｜DSML｜function_calls>
47
"""
48
49
+ tool_call_start_token: str = "<｜DSML｜function_calls>"
50
+ tool_call_end_token: str = "</｜DSML｜function_calls>"
51
+
52
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
53
super().__init__(tokenizer, tools)
54
55
self.prev_tool_call_arr: list[dict] = []
56
54
- # Sentinel token
55
- self.tool_call_start_token: str = "<｜DSML｜function_calls>"
56
-
57
# Streaming state
58
self.current_tool_index: int = 0
59
self._sent_content_idx: int = 0
60
61
# Regex patterns for complete parsing
62
self.tool_call_complete_regex = re.compile(
63
- r"<｜DSML｜function_calls>(.*?)</｜DSML｜function_calls>", re.DOTALL
63
+ re.escape(self.tool_call_start_token)
64
+ + r"(.*?)"
65
+ + re.escape(self.tool_call_end_token),
66
+ re.DOTALL,
67
)
68
self.invoke_complete_regex = re.compile(
69
r'<｜DSML｜invoke\s+name="([^"]+)"\s*>(.*?)</｜DSML｜invoke>', re.DOTALL
@@ -86,7 +89,7 @@ class DeepSeekV32ToolParser(ToolParser):
89
request = super().adjust_request(request)
90
if request.tools and request.tool_choice != "none":
91
# Ensure tool call tokens
89
- # (<｜DSML｜function_calls>, </｜DSML｜function_calls>)
92
+ # (e.g. <｜DSML｜function_calls>, </｜DSML｜function_calls>)
93
# are not skippedduring decoding.
94
# Even though they are not marked as special tokens,
95
# setting skip_special_tokens=False ensures proper handling in
vllm/tool_parsers/deepseekv4_tool_parser.py
new
+16
@@ -0,0 +1,16 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser
5
+
6
+
7
+class DeepSeekV4ToolParser(DeepSeekV32ToolParser):
8
+ """
9
+ DeepSeek V4 DSML tool parser.
10
+
11
+ V4 keeps the V3.2 DSML invoke/parameter grammar, but wraps tool calls in
12
+ ``<｜DSML｜tool_calls>`` instead of ``<｜DSML｜function_calls>``.
13
+ """
14
+
15
+ tool_call_start_token: str = "<｜DSML｜tool_calls>"
16
+ tool_call_end_token: str = "</｜DSML｜tool_calls>"
vllm/transformers_utils/config.py
+1
@@ -89,6 +89,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
89
qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig",
90
deepseek_vl_v2="DeepseekVLV2Config",
91
deepseek_v32="DeepseekV3Config",
92
+ deepseek_v4="DeepseekV4Config",
93
flex_olmo="FlexOlmoConfig",
94
fireredlid="FireRedLIDConfig",
95
funaudiochat="FunAudioChatConfig",
vllm/transformers_utils/configs/__init__.py
+2
@@ -26,6 +26,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
26
"OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
27
"Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3",
28
"DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
29
+ "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4",
30
"DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
31
"EAGLEConfig": "vllm.transformers_utils.configs.eagle",
32
"FireRedLIDConfig": "vllm.transformers_utils.configs.fireredlid",
@@ -88,6 +89,7 @@ __all__ = [
89
"Qwen3VLNemotronEmbedConfig",
90
"DeepseekVLV2Config",
91
"DeepseekV3Config",
92
+ "DeepseekV4Config",
93
"DotsOCRConfig",
94
"EAGLEConfig",
95
"FlexOlmoConfig",
vllm/transformers_utils/configs/deepseek_v4.py
new
+23
@@ -0,0 +1,23 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Any
4
+
5
+from transformers import PretrainedConfig
6
+
7
+
8
+class DeepseekV4Config(PretrainedConfig):
9
+ model_type = "deepseek_v4"
10
+
11
+ def __init__(
12
+ self,
13
+ max_position_embeddings: int = 1048576,
14
+ rope_scaling: dict[str, Any] | None = None,
15
+ rope_parameters: dict[str, Any] | None = None,
16
+ rope_theta: float = 10000.0,
17
+ **kwargs,
18
+ ):
19
+ self.max_position_embeddings = max_position_embeddings
20
+ self.rope_scaling = rope_scaling
21
+ self.rope_theta = rope_theta
22
+ self.rope_parameters = rope_scaling or rope_parameters
23
+ super().__init__(**kwargs)
vllm/transformers_utils/model_arch_config_convertor.py
+9
-1
@@ -47,6 +47,9 @@ class ModelArchConfigConvertorBase:
47
48
def get_head_size(self) -> int:
49
if self.is_deepseek_mla():
50
+ # special case for deepseek_v4
51
+ if hasattr(self.hf_text_config, "compress_ratios"):
52
+ return self.hf_text_config.head_dim
53
qk_rope_head_dim = getattr(self.hf_text_config, "qk_rope_head_dim", 0)
54
if not envs.VLLM_MLA_DISABLE:
55
return self.hf_text_config.kv_lora_rank + qk_rope_head_dim
@@ -222,6 +225,7 @@ class ModelArchConfigConvertorBase:
225
"deepseek_v2",
226
"deepseek_v3",
227
"deepseek_v32",
228
+ "deepseek_v4",
229
"deepseek_mtp",
230
"glm_moe_dsa",
231
"glm4_moe_lite",
@@ -233,7 +237,11 @@ class ModelArchConfigConvertorBase:
237
"pangu_ultra_moe_mtp",
238
"bailing_hybrid",
239
):
236
- return getattr(self.hf_text_config, "kv_lora_rank", None) is not None
240
+ # check is deepseek_v4 model
241
+ if hasattr(self.hf_text_config, "compress_ratios"):
242
+ return getattr(self.hf_text_config, "head_dim", None) is not None
243
+ else:
244
+ return getattr(self.hf_text_config, "kv_lora_rank", None) is not None
245
elif self.hf_text_config.model_type == "eagle":
246
# if the model is an EAGLE module, check for the
247
# underlying architecture
vllm/utils/deep_gemm.py
+125
-40
@@ -125,12 +125,16 @@ def _missing(*_: Any, **__: Any) -> NoReturn:
125
)
126
127
128
+_cublaslt_gemm_nt_impl: Callable[..., Any] | None = None
129
_fp8_gemm_nt_impl: Callable[..., Any] | None = None
130
+_fp8_einsum_impl: Callable[..., Any] | None = None
131
_grouped_impl: Callable[..., Any] | None = None
132
_grouped_masked_impl: Callable[..., Any] | None = None
131
-_fp8_mqa_logits_impl: Callable[..., Any] | None = None
132
-_fp8_paged_mqa_logits_impl: Callable[..., Any] | None = None
133
+_grouped_fp4_impl: Callable[..., Any] | None = None
134
+_fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None
135
+_fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None
136
_get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
137
+_tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
138
_get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
139
_get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
140
_transform_sf_into_required_layout_impl: Callable[..., Any] | None = None
@@ -173,20 +177,27 @@ def _import_deep_gemm():
177
178
def _lazy_init() -> None:
179
"""Import deep_gemm and resolve symbols on first use."""
176
- global _fp8_gemm_nt_impl, _grouped_impl, _grouped_masked_impl
177
- global _fp8_mqa_logits_impl, _fp8_paged_mqa_logits_impl
180
+ global _cublaslt_gemm_nt_impl
181
+ global _fp8_gemm_nt_impl, _fp8_einsum_impl
182
+ global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl
183
+ global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl
184
global _get_paged_mqa_logits_metadata_impl
185
+ global _tf32_hc_prenorm_gemm_impl
186
global _get_mn_major_tma_aligned_tensor_impl
187
global _get_mk_alignment_for_contiguous_layout_impl
188
global _transform_sf_into_required_layout_impl
189
# fast path
190
if (
184
- _fp8_gemm_nt_impl is not None
191
+ _cublaslt_gemm_nt_impl is not None
192
+ or _fp8_gemm_nt_impl is not None
193
+ or _fp8_einsum_impl is not None
194
or _grouped_impl is not None
195
or _grouped_masked_impl is not None
187
- or _fp8_mqa_logits_impl is not None
188
- or _fp8_paged_mqa_logits_impl is not None
196
+ or _grouped_fp4_impl is not None
197
+ or _fp8_fp4_mqa_logits_impl is not None
198
+ or _fp8_fp4_paged_mqa_logits_impl is not None
199
or _get_paged_mqa_logits_metadata_impl is not None
200
+ or _tf32_hc_prenorm_gemm_impl is not None
201
or _get_mk_alignment_for_contiguous_layout_impl is not None
202
or _transform_sf_into_required_layout_impl is not None
203
):
@@ -206,14 +217,20 @@ def _lazy_init() -> None:
217
if _dg is None:
218
return
219
220
+ _cublaslt_gemm_nt_impl = getattr(_dg, "cublaslt_gemm_nt", None)
221
_fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
222
+ _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None)
223
_grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None)
224
_grouped_masked_impl = getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None)
212
- _fp8_mqa_logits_impl = getattr(_dg, "fp8_mqa_logits", None)
213
- _fp8_paged_mqa_logits_impl = getattr(_dg, "fp8_paged_mqa_logits", None)
225
+ _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None)
226
+ # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that
227
+ # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`.
228
+ _fp8_fp4_mqa_logits_impl = getattr(_dg, "fp8_fp4_mqa_logits", None)
229
+ _fp8_fp4_paged_mqa_logits_impl = getattr(_dg, "fp8_fp4_paged_mqa_logits", None)
230
_get_paged_mqa_logits_metadata_impl = getattr(
231
_dg, "get_paged_mqa_logits_metadata", None
232
)
233
+ _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None)
234
_get_mn_major_tma_aligned_tensor_impl = getattr(
235
_dg, "get_mn_major_tma_aligned_tensor", None
236
)
@@ -259,6 +276,13 @@ def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
276
return _get_mn_major_tma_aligned_tensor_impl(x)
277
278
279
+def cublaslt_gemm_nt(*args, **kwargs):
280
+ _lazy_init()
281
+ if _cublaslt_gemm_nt_impl is None:
282
+ return _missing(*args, **kwargs)
283
+ return _cublaslt_gemm_nt_impl(*args, **kwargs)
284
+
285
+
286
def fp8_gemm_nt(*args, **kwargs):
287
_lazy_init()
288
if _fp8_gemm_nt_impl is None:
@@ -271,6 +295,13 @@ def fp8_gemm_nt(*args, **kwargs):
295
return _fp8_gemm_nt_impl(*args, disable_ue8m0_cast=not use_ue8m0, **kwargs)
296
297
298
+def fp8_einsum(*args, **kwargs):
299
+ _lazy_init()
300
+ if _fp8_einsum_impl is None:
301
+ return _missing(*args, **kwargs)
302
+ return _fp8_einsum_impl(*args, **kwargs)
303
+
304
+
305
def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs):
306
_lazy_init()
307
if _grouped_impl is None:
@@ -280,6 +311,15 @@ def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs):
311
)
312
313
314
+def m_grouped_fp8_fp4_gemm_nt_contiguous(*args, **kwargs):
315
+ _lazy_init()
316
+ if _grouped_fp4_impl is None:
317
+ return _missing(*args, **kwargs)
318
+ return _grouped_fp4_impl(
319
+ *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
320
+ )
321
+
322
+
323
def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):
324
_lazy_init()
325
if _grouped_masked_impl is None:
@@ -298,37 +338,48 @@ def transform_sf_into_required_layout(*args, **kwargs):
338
)
339
340
301
-def fp8_mqa_logits(
302
- q: torch.Tensor,
341
+def fp8_fp4_mqa_logits(
342
+ q: tuple[torch.Tensor, torch.Tensor | None],
343
kv: tuple[torch.Tensor, torch.Tensor],
344
weights: torch.Tensor,
345
cu_seqlen_ks: torch.Tensor,
346
cu_seqlen_ke: torch.Tensor,
347
clean_logits: bool,
348
) -> torch.Tensor:
309
- """Compute FP8 MQA logits for a single sequence without KV paging.
349
+ """Compute MQA logits for a single sequence without KV paging.
350
+
351
+ Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes
352
+ ``q = (values, scales_or_None)`` where ``scales`` is None for FP8 Q
353
+ (per-token scale is folded into ``weights``) and a packed block-scale
354
+ tensor for MXFP4 Q.
355
356
Args:
312
- q: Query tensor of shape [M, H, D]. Casted to
313
- `torch.float8_e4m3fn` by caller.
314
- kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
315
- dtype `torch.float8_e4m3fn` and `k_scales` has shape [N])
316
- with dtype `torch.float32`.
357
+ q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is [M, H, D]
358
+ float8_e4m3fn and q_scale is None (per-token scale is folded
359
+ into ``weights``). FP4 path: q_values is packed uint8 and
360
+ q_scale is the companion block-scale tensor.
361
+ kv: Tuple `(k_packed, k_scales)` — FP8 layout is [N, D]
362
+ float8_e4m3fn plus fp32 scales [N]; FP4 layout is packed uint8.
363
weights: weights of shape [M, H], dtype `torch.float32`.
318
- cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
319
- shape [M], dtype int32.
320
- cu_seqlen_ke: End indices (exclusive) for valid K per query position,
321
- shape [M], dtype int32.
364
+ cu_seqlen_ks: Start indices (inclusive) for valid K per query
365
+ position, shape [M], dtype int32.
366
+ cu_seqlen_ke: End indices (exclusive) for valid K per query
367
+ position, shape [M], dtype int32.
368
clean_logits: Whether to clean the unfilled logits into `-inf`.
369
370
Returns:
371
Logits tensor of shape [M, N], dtype `torch.float32`.
372
"""
373
_lazy_init()
328
- if _fp8_mqa_logits_impl is None:
374
+ if _fp8_fp4_mqa_logits_impl is None:
375
return _missing()
330
- return _fp8_mqa_logits_impl(
331
- q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=clean_logits
376
+ return _fp8_fp4_mqa_logits_impl(
377
+ q,
378
+ kv,
379
+ weights,
380
+ cu_seqlen_ks,
381
+ cu_seqlen_ke,
382
+ clean_logits=clean_logits,
383
)
384
385
@@ -344,7 +395,7 @@ def get_paged_mqa_logits_metadata(
395
num_sms: Number of SMs available. 132 for Hopper
396
397
Returns:
347
- Backend-specific tensor consumed by `fp8_paged_mqa_logits` to
398
+ Backend-specific tensor consumed by `fp8_fp4_paged_mqa_logits` to
399
schedule work across SMs.
400
"""
401
_lazy_init()
@@ -353,9 +404,9 @@ def get_paged_mqa_logits_metadata(
404
return _get_paged_mqa_logits_metadata_impl(context_lens, block_size, num_sms)
405
406
356
-def fp8_paged_mqa_logits(
357
- q_fp8: torch.Tensor,
358
- kv_cache_fp8: torch.Tensor,
407
+def fp8_fp4_paged_mqa_logits(
408
+ q: tuple[torch.Tensor, torch.Tensor | None],
409
+ kv_cache: torch.Tensor,
410
weights: torch.Tensor,
411
context_lens: torch.Tensor,
412
block_tables: torch.Tensor,
@@ -363,14 +414,20 @@ def fp8_paged_mqa_logits(
414
max_model_len: int,
415
clean_logits: bool,
416
) -> torch.Tensor:
366
- """Compute FP8 MQA logits using paged KV-cache.
417
+ """Compute MQA logits using a paged KV-cache.
418
+
419
+ Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes
420
+ ``q = (values, scales_or_None)``; pass ``(q_tensor, None)`` for the FP8
421
+ path and ``(q_values, q_scale)`` for MXFP4.
422
423
Args:
369
- q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to
370
- `torch.float8_e4m3fn` by caller.
371
- kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape
372
- [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last
373
- 4 bytes per (block,pos) store the `float` dequant scale.
424
+ q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is
425
+ [B, next_n, H, D] float8_e4m3fn and q_scale is None. FP4 path:
426
+ q_values is packed uint8 and q_scale is the companion
427
+ block-scale tensor.
428
+ kv_cache: Paged KV-cache. FP8 layout is [num_blocks, block_size, 1,
429
+ D+4], dtype `torch.uint8`, with the last 4 bytes per (block, pos)
430
+ storing the float dequant scale.
431
weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
432
context_lens: Tensor of shape [B], dtype int32; effective context length
433
for each batch element.
@@ -386,11 +443,11 @@ def fp8_paged_mqa_logits(
443
`torch.float32`.
444
"""
445
_lazy_init()
389
- if _fp8_paged_mqa_logits_impl is None:
446
+ if _fp8_fp4_paged_mqa_logits_impl is None:
447
return _missing()
391
- return _fp8_paged_mqa_logits_impl(
392
- q_fp8,
393
- kv_cache_fp8,
448
+ return _fp8_fp4_paged_mqa_logits_impl(
449
+ q,
450
+ kv_cache,
451
weights,
452
context_lens,
453
block_tables,
@@ -400,6 +457,32 @@ def fp8_paged_mqa_logits(
457
)
458
459
460
+def tf32_hc_prenorm_gemm(
461
+ x: torch.Tensor,
462
+ fn: torch.Tensor,
463
+ out: torch.Tensor,
464
+ sqrsum: torch.Tensor,
465
+ num_split: int,
466
+) -> torch.Tensor:
467
+ """
468
+ Perform the following computation:
469
+ out = x.float() @ fn.T
470
+ sqrsum = x.float().square().sum(-1)
471
+
472
+ See the caller function for shape requirement
473
+ """
474
+ _lazy_init()
475
+ if _tf32_hc_prenorm_gemm_impl is None:
476
+ return _missing()
477
+ return _tf32_hc_prenorm_gemm_impl(
478
+ x,
479
+ fn,
480
+ out,
481
+ sqrsum,
482
+ num_split,
483
+ )
484
+
485
+
486
def _ceil_to_ue8m0(x: torch.Tensor):
487
return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
488
@@ -482,10 +565,12 @@ __all__ = [
565
"calc_diff",
566
"DeepGemmQuantScaleFMT",
567
"fp8_gemm_nt",
568
+ "fp8_einsum",
569
"m_grouped_fp8_gemm_nt_contiguous",
570
+ "m_grouped_fp8_fp4_gemm_nt_contiguous",
571
"fp8_m_grouped_gemm_nt_masked",
487
- "fp8_mqa_logits",
488
- "fp8_paged_mqa_logits",
572
+ "fp8_fp4_mqa_logits",
573
+ "fp8_fp4_paged_mqa_logits",
574
"get_paged_mqa_logits_metadata",
575
"per_block_cast_to_fp8",
576
"is_deep_gemm_e8m0_used",
vllm/utils/multi_stream_utils.py
+10
@@ -2,11 +2,21 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable
5
+from enum import Enum
6
from typing import Any
7
8
import torch
9
10
11
+class AuxStreamType(Enum):
12
+ Attention = 1
13
+
14
+
15
+class EventType(Enum):
16
+ Main = 0
17
+ Attention = 1
18
+
19
+
20
def maybe_execute_in_parallel(
21
fn0: Callable[[], Any],
22
fn1: Callable[[], Any],
vllm/v1/attention/backend.py
+5
@@ -392,6 +392,11 @@ class CommonAttentionMetadata:
392
dcp_local_seq_lens_cpu: torch.Tensor | None = None
393
"""Sequence lengths of the local rank in decode context parallelism world"""
394
395
+ positions: torch.Tensor | None = None
396
+ """(num_actual_tokens,) token positions. Optional; set when the caller
397
+ has positions available so that builders can pre-compute position-dependent
398
+ metadata (e.g. C128A topk indices for DeepSeek V4)."""
399
+
400
is_prefilling: torch.Tensor | None = None
401
"""(batch_size,) bool tensor: True if request is still in prefill phase
402
(num_computed_tokens < num_prompt_tokens). Used by some backends to
vllm/v1/attention/backends/mla/compressor_utils.py
new
+86
@@ -0,0 +1,86 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.triton_utils import tl, triton
6
+
7
+
8
+@triton.jit
9
+def _compressed_slot_mapping_kernel(
10
+ # [num_tokens]
11
+ slot_mapping_ptr,
12
+ # [num_reqs + 1]
13
+ query_start_loc_ptr,
14
+ # [num_reqs]
15
+ seq_lens_ptr,
16
+ # [num_reqs, max_num_blocks]
17
+ block_table_ptr,
18
+ block_table_stride,
19
+ block_size,
20
+ COMPRESS_RATIO: tl.constexpr,
21
+ PAD_ID: tl.constexpr,
22
+ TRITON_BLOCK_SIZE: tl.constexpr,
23
+):
24
+ batch_idx = tl.program_id(0)
25
+
26
+ query_start = tl.load(query_start_loc_ptr + batch_idx)
27
+ query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
28
+ query_len = query_end - query_start
29
+
30
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
31
+ start_pos = seq_len - query_len
32
+
33
+ for i in range(0, query_len, TRITON_BLOCK_SIZE):
34
+ offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
35
+ mask = offset < query_len
36
+
37
+ pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
38
+ is_valid = (pos + 1) % COMPRESS_RATIO == 0
39
+ pos_after_compress = pos // COMPRESS_RATIO
40
+
41
+ block_ids = pos_after_compress // block_size
42
+ block_numbers = tl.load(
43
+ block_table_ptr + batch_idx * block_table_stride + block_ids,
44
+ mask=mask & is_valid,
45
+ )
46
+ slot_ids = block_numbers * block_size + pos_after_compress % block_size
47
+
48
+ # NOTE
49
+ slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
50
+ tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)
51
+
52
+
53
+def get_compressed_slot_mapping(
54
+ num_tokens: int,
55
+ query_start_loc: torch.Tensor,
56
+ seq_lens: torch.Tensor,
57
+ block_table: torch.Tensor,
58
+ block_size: int,
59
+ compress_ratio: int,
60
+ out: torch.Tensor | None = None,
61
+) -> torch.Tensor:
62
+ if out is not None:
63
+ # Guard: for padded / invalid sequences.
64
+ # Negative positions produce bogus block indices that lead to illegal memory
65
+ # accesses inside the block_table load.
66
+ # NOTE: Fill -1 to the whole tensor, not just the first `num_tokens`.
67
+ out.fill_(-1)
68
+ slot_mapping = out[:num_tokens]
69
+ else:
70
+ slot_mapping = torch.full(
71
+ (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device
72
+ )
73
+
74
+ num_reqs = block_table.shape[0]
75
+ _compressed_slot_mapping_kernel[(num_reqs,)](
76
+ slot_mapping,
77
+ query_start_loc,
78
+ seq_lens,
79
+ block_table,
80
+ block_table.stride(0),
81
+ block_size,
82
+ compress_ratio,
83
+ PAD_ID=-1,
84
+ TRITON_BLOCK_SIZE=1024,
85
+ )
86
+ return slot_mapping
vllm/v1/attention/backends/mla/flashmla_sparse.py
+328
-23
@@ -15,6 +15,8 @@ from vllm.model_executor.layers.attention.mla_attention import (
15
)
16
from vllm.platforms import current_platform
17
from vllm.platforms.interface import DeviceCapability
18
+from vllm.triton_utils import tl, triton
19
+from vllm.utils.math_utils import cdiv
20
from vllm.utils.platform_utils import num_compute_units
21
from vllm.utils.torch_utils import is_quantized_kv_cache
22
from vllm.v1.attention.backend import (
@@ -27,6 +29,7 @@ from vllm.v1.attention.backend import (
29
MultipleOf,
30
SparseMLAAttentionImpl,
31
)
32
+from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
33
from vllm.v1.attention.backends.mla.sparse_utils import (
34
triton_convert_req_index_to_global_index,
35
)
@@ -65,8 +68,8 @@ MIN_HEADS_FOR_BF16_PREFILL = 32
68
"""
69
NOTE: FlashMLA Sparse uses an fp8 cache with the following format
70
68
-In the "FP8 with scale" format, each token's KV cache is 656 Bytes,
69
-structured as:
71
+For DeepSeek V3.2, in the "FP8 with scale" format, each token's KV cache is 656
72
+Bytes, structured as:
73
- **First 512 bytes:** The "quantized NoPE" part, containing 512
74
`float8_e4m3` values.
75
- **Next 16 bytes:** Scale factors, containing 4 `float32` values.
@@ -74,6 +77,16 @@ structured as:
77
the second for the next 128, and so on.
78
- **Last 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This
79
part is not quantized for accuracy.
80
+
81
+For DeepSeek V4, in the "FP8 with scale" format, each token's KV cache is 584
82
+Bytes, structured as:
83
+- **First 448 bytes:** The "quantized NoPE" part, containing 448
84
+ `float8_e4m3` values.
85
+- **Next 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This
86
+ part is not quantized for accuracy.
87
+- **Last 8 bytes:** Scale factors, containing 7 `ue8m0` values + 1B pad.
88
+ The first `ue8m0` is the scale for the first 64 `float8_e4m3` values,
89
+ the second for the next 64, and so on.
90
"""
91
92
@@ -104,7 +117,8 @@ class FlashMLASparseBackend(AttentionBackend):
117
118
@classmethod
119
def get_supported_head_sizes(cls) -> list[int]:
107
- return [576]
120
+ # V3.2: 576 (512 NoPE + 64 RoPE); DeepseekV4: 512 (448 NoPE + 64 RoPE)
121
+ return [512, 576]
122
123
@classmethod
124
def is_mla(cls) -> bool:
@@ -127,13 +141,37 @@ class FlashMLASparseBackend(AttentionBackend):
141
cache_dtype_str: str = "auto",
142
) -> tuple[int, ...]:
143
if cache_dtype_str == "fp8_ds_mla":
130
- # custom storage format is 656 bytes
131
- # see FlashMLA readme.md for details
144
+ # V3.2 main MLA: 656-byte custom storage format. See module docstring.
145
return (num_blocks, block_size, 656)
146
else:
147
return (num_blocks, block_size, head_size)
148
149
150
+class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
151
+ @staticmethod
152
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
153
+ return [256]
154
+
155
+ @staticmethod
156
+ def get_name() -> str:
157
+ return "V4_FLASHMLA_SPARSE"
158
+
159
+ @staticmethod
160
+ def get_kv_cache_shape(
161
+ num_blocks: int,
162
+ block_size: int,
163
+ num_kv_heads: int,
164
+ head_size: int,
165
+ cache_dtype_str: str = "auto",
166
+ ) -> tuple[int, ...]:
167
+ if cache_dtype_str == "fp8_ds_mla":
168
+ # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
169
+ # head_size passed in is the semantic head_dim (512).
170
+ return (num_blocks, block_size, 584)
171
+ else:
172
+ return (num_blocks, block_size, head_size)
173
+
174
+
175
@dataclass
176
class FlashMLASparseMetadata(AttentionMetadata):
177
num_reqs: int
@@ -159,6 +197,7 @@ class FlashMLASparseMetadata(AttentionMetadata):
197
class FP8SeparatePrefillDecode:
198
@dataclass
199
class Decode:
200
+ seq_lens: torch.Tensor
201
kernel_metadata: "FlashMLASparseMetadata.FP8KernelMetadata"
202
decode_query_len: int # needed for reshape in spec decode
203
@@ -206,6 +245,13 @@ class FlashMLASparseMetadata(AttentionMetadata):
245
fp8_extra_metadata: FP8SeparatePrefillDecode | FP8KernelMetadata | None = None
246
fp8_use_mixed_batch: bool = False
247
248
+ # Pre-computed C128A metadata (DeepseekV4 only, compress_ratio == 128).
249
+ # Decode: global slot ids + valid-entry counts (fused from positions).
250
+ c128a_global_decode_topk_indices: torch.Tensor | None = None
251
+ c128a_decode_topk_lens: torch.Tensor | None = None
252
+ # Prefill: local topk indices (used by combine_topk_swa_indices).
253
+ c128a_prefill_topk_indices: torch.Tensor | None = None
254
+
255
256
def get_prefill_workspace_size(max_model_len: int):
257
# NOTE(Lucas): 5 is a magic number for controlling the prefill buffer size.
@@ -235,8 +281,9 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
281
parallel_config = vllm_config.parallel_config
282
self.device = device
283
238
- # Treat requests with query length <= 1 as decodes to match the
239
- # DeepGEMM indexer constraint (fp8_paged_mqa_logits only supports next_n <= 2)
284
+ # Classify single-token queries (plus num_speculative_tokens via
285
+ # supports_spec_as_decode=True) as decodes; longer queries go to
286
+ # prefill.
287
self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
288
289
sm_count = num_compute_units(device.index)
@@ -300,6 +347,68 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
347
device=device,
348
)
349
350
+ # DeepseekV4: has compress_ratios in hf_config.
351
+ hf_config = vllm_config.model_config.hf_config
352
+ self.is_deepseek_v4 = (
353
+ hasattr(hf_config, "compress_ratios") and len(hf_config.compress_ratios) > 0
354
+ )
355
+ self.compress_ratio = 1
356
+ if self.is_deepseek_v4:
357
+ assert hasattr(self.kv_cache_spec, "compress_ratio")
358
+ self.compress_ratio = self.kv_cache_spec.compress_ratio
359
+ # Pre-allocate compressed slot mapping buffer for CUDA graph
360
+ # address stability when compress_ratio > 1.
361
+ if self.compress_ratio > 1:
362
+ max_num_batched_tokens = (
363
+ vllm_config.scheduler_config.max_num_batched_tokens
364
+ )
365
+ self.compressed_slot_mapping_buffer = torch.empty(
366
+ max_num_batched_tokens,
367
+ dtype=torch.int64,
368
+ device=self.device,
369
+ )
370
+
371
+ # Pre-allocate C128A topk buffers for CUDA graph address stability.
372
+ if self.compress_ratio == 128:
373
+ max_num_batched_tokens = (
374
+ vllm_config.scheduler_config.max_num_batched_tokens
375
+ )
376
+ # Pad to B_TOPK alignment (128 covers both h_q=64 B_TOPK=64 and
377
+ # h_q=128 B_TOPK=128). FlashMLA decode asserts extra_topk % B_TOPK
378
+ # == 0; unaligned widths (e.g. 17 = ceil(2136/128)) crash the
379
+ # sm100 head64 kernel. Padded slots stay -1 and decode_lens caps
380
+ # them via topk_length, so the pad is a no-op at kernel level.
381
+ # Mirrors _SPARSE_PREFILL_TOPK_ALIGNMENT in cache_utils.py.
382
+ _C128A_TOPK_ALIGNMENT = 128
383
+ c128a_max_compressed = cdiv(
384
+ self.model_config.max_model_len, self.compress_ratio
385
+ )
386
+ c128a_max_compressed = (
387
+ cdiv(c128a_max_compressed, _C128A_TOPK_ALIGNMENT)
388
+ * _C128A_TOPK_ALIGNMENT
389
+ )
390
+ # Stored so _build_c128a_metadata passes it as the kernel's
391
+ # max_compressed_tokens, matching the buffer stride. Otherwise
392
+ # the kernel's default 8192 iterates past row width and spills
393
+ # writes into adjacent rows (present in both decode and prefill
394
+ # branches of _build_c128a_topk_metadata_kernel).
395
+ self.c128a_max_compressed = c128a_max_compressed
396
+ self.c128a_global_decode_buffer = torch.empty(
397
+ (max_num_batched_tokens, c128a_max_compressed),
398
+ dtype=torch.int32,
399
+ device=self.device,
400
+ )
401
+ self.c128a_decode_lens_buffer = torch.empty(
402
+ max_num_batched_tokens,
403
+ dtype=torch.int32,
404
+ device=self.device,
405
+ )
406
+ self.c128a_prefill_buffer = torch.empty(
407
+ (max_num_batched_tokens, c128a_max_compressed),
408
+ dtype=torch.int32,
409
+ device=self.device,
410
+ )
411
+
412
def _build_fp8_mixed_decode_prefill(
413
self,
414
common_attn_metadata: CommonAttentionMetadata,
@@ -460,15 +569,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
569
decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item()
570
571
# Use padded head count since that's what the kernel will see
463
- padded_heads = self.fp8_decode_padded_heads
464
- scheduler_metadata, _ = get_mla_metadata(
465
- cache_seqlens=self.topk_tokens_tensor[:num_decodes],
466
- num_q_tokens_per_head_k=decode_query_len * padded_heads,
467
- topk=self.topk_tokens,
468
- num_heads_q=padded_heads,
469
- num_heads_k=1,
470
- is_fp8_kvcache=True,
471
- )
572
+ scheduler_metadata, _ = get_mla_metadata()
573
574
kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata(
575
scheduler_metadata=scheduler_metadata,
@@ -476,6 +577,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
577
cache_lens=self.max_model_len_tensor[:num_decodes],
578
)
579
fp8_metadata.decode = FP8Meta.Decode(
580
+ seq_lens=common_attn_metadata.seq_lens[:num_decodes],
581
kernel_metadata=kernel_meta,
582
decode_query_len=decode_query_len,
583
)
@@ -502,35 +604,109 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
604
)
605
req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
606
607
+ slot_mapping = cm.slot_mapping
608
+ if self.compress_ratio > 1:
609
+ slot_mapping = get_compressed_slot_mapping(
610
+ common_attn_metadata.num_actual_tokens,
611
+ common_attn_metadata.query_start_loc,
612
+ common_attn_metadata.seq_lens,
613
+ common_attn_metadata.block_table_tensor.clamp(min=0),
614
+ int(self.kv_cache_spec.storage_block_size),
615
+ self.compress_ratio,
616
+ out=self.compressed_slot_mapping_buffer,
617
+ )
618
+
619
fp8_extra_metadata: (
620
FlashMLASparseMetadata.FP8SeparatePrefillDecode
621
| FlashMLASparseMetadata.FP8KernelMetadata
622
| None
623
) = None
510
- fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL
511
- if self.use_fp8_kv_cache:
624
+ fp8_use_mixed_batch = (
625
+ self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.is_deepseek_v4
626
+ )
627
+ # DeepseekV4 has its own attention impl (DeepseekV4MLAAttention) that does not
628
+ # consume fp8_extra_metadata. Skipping the build here avoids a
629
+ # forced D2H sync on seq_lens that would otherwise fire on every
630
+ # prefill-bearing step, lifting GPU utilization on long-prefill
631
+ # workloads (e.g. LongBench) from ~83% to ~100%.
632
+ if self.use_fp8_kv_cache and not self.is_deepseek_v4:
633
if fp8_use_mixed_batch:
634
fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(cm)
635
else:
636
fp8_extra_metadata = self._build_fp8_separate_prefill_decode(cm)
637
638
+ # Pre-compute C128A topk indices for DeepseekV4.
639
+ c128a_fields = {}
640
+ if self.is_deepseek_v4 and self.compress_ratio == 128:
641
+ c128a_fields = self._build_c128a_metadata(cm, req_id_per_token)
642
+
643
metadata = FlashMLASparseMetadata(
644
num_reqs=cm.num_reqs,
645
max_query_len=cm.max_query_len,
646
max_seq_len=cm.max_seq_len,
647
num_actual_tokens=cm.num_actual_tokens,
648
query_start_loc=cm.query_start_loc,
523
- slot_mapping=cm.slot_mapping,
649
+ slot_mapping=slot_mapping,
650
block_table=cm.block_table_tensor,
651
req_id_per_token=req_id_per_token,
652
block_size=self.kv_cache_spec.block_size,
653
topk_tokens=self.topk_tokens,
654
fp8_extra_metadata=fp8_extra_metadata,
655
fp8_use_mixed_batch=fp8_use_mixed_batch,
656
+ **c128a_fields,
657
)
658
659
return metadata
660
661
+ def _build_c128a_metadata(
662
+ self,
663
+ cm: CommonAttentionMetadata,
664
+ req_id_per_token: torch.Tensor,
665
+ ) -> dict[str, torch.Tensor | None]:
666
+ """Pre-compute C128A topk indices for DeepseekV4 (compress_ratio >= 128)."""
667
+ # Must match SWA's decode split (no `require_uniform=True`) so
668
+ # `c128a_global_decode_topk_indices.shape[0]` lines up with q in
669
+ # `_forward_decode`. The per-token C128A kernel handles non-uniform
670
+ # query lengths.
671
+ (num_decodes, _, num_decode_tokens, num_prefill_tokens) = (
672
+ split_decodes_and_prefills(
673
+ cm,
674
+ decode_threshold=self.reorder_batch_threshold or 1,
675
+ )
676
+ )
677
+
678
+ num_total = num_decode_tokens + num_prefill_tokens
679
+ if num_total == 0:
680
+ return {}
681
+
682
+ assert cm.positions is not None, (
683
+ "positions is required for C128A metadata build"
684
+ )
685
+ block_size = self.kv_cache_spec.block_size // self.compress_ratio
686
+ global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(
687
+ cm.positions[:num_total],
688
+ self.compress_ratio,
689
+ num_decode_tokens,
690
+ req_id_per_token,
691
+ cm.block_table_tensor[:num_decodes],
692
+ block_size,
693
+ cm.slot_mapping,
694
+ self.c128a_global_decode_buffer,
695
+ self.c128a_decode_lens_buffer,
696
+ self.c128a_prefill_buffer,
697
+ max_compressed_tokens=self.c128a_max_compressed,
698
+ )
699
+
700
+ result: dict[str, torch.Tensor | None] = {}
701
+ if num_decode_tokens > 0:
702
+ result["c128a_global_decode_topk_indices"] = global_decode.view(
703
+ num_decode_tokens, 1, -1
704
+ )
705
+ result["c128a_decode_topk_lens"] = decode_lens
706
+ if num_prefill_tokens > 0:
707
+ result["c128a_prefill_topk_indices"] = prefill_local
708
+ return result
709
+
710
711
class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
712
@staticmethod
@@ -552,7 +728,7 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
728
attn_type: str,
729
kv_sharing_target_layer_name: str | None,
730
# MLA Specific Arguments
555
- topk_indice_buffer: torch.Tensor | None = None,
731
+ topk_indices_buffer: torch.Tensor | None = None,
732
indexer: "Indexer | None" = None,
733
**mla_args,
734
) -> None:
@@ -615,7 +791,11 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
791
NUM_TOPK_TOKENS=topk_indices.shape[1],
792
)
793
618
- return self._bf16_flash_mla_kernel(q, kv_c_and_k_pe_cache, topk_indices)
794
+ return self._bf16_flash_mla_kernel(
795
+ q,
796
+ kv_c_and_k_pe_cache,
797
+ topk_indices,
798
+ )
799
800
def _forward_fp8_kv_separate_prefill_decode(
801
self,
@@ -656,7 +836,10 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
836
fp8_metadata = attn_metadata.fp8_extra_metadata
837
assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode)
838
659
- def _fp8_decode(q: torch.Tensor, topk_indices: torch.Tensor) -> torch.Tensor:
839
+ def _fp8_decode(
840
+ q: torch.Tensor,
841
+ topk_indices: torch.Tensor,
842
+ ) -> torch.Tensor:
843
# Reshape q: (num_decode_tokens, num_heads, head_dim)
844
# -> (num_decodes, seq_len, num_heads, head_dim)
845
q = reshape_query_for_spec_decode(q, num_decodes)
@@ -692,7 +875,8 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
875
876
if num_decode_tokens > 0:
877
attn_out[:num_decode_tokens] = _fp8_decode(
695
- q[:num_decode_tokens], topk_indices[:num_decode_tokens]
878
+ q[:num_decode_tokens],
879
+ topk_indices[:num_decode_tokens],
880
)
881
882
assert fp8_metadata.prefill is not None
@@ -823,6 +1007,7 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
1007
output = flash_mla_sparse_fwd(
1008
q, kv_c_and_k_pe_cache, topk_indices, self.softmax_scale
1009
)[0]
1010
+
1011
output = output[:, : self.num_heads, :]
1012
return output
1013
@@ -864,3 +1049,123 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
1049
)
1050
1051
return attn_out, None
1052
+
1053
+
1054
+def build_c128a_topk_metadata(
1055
+ positions: torch.Tensor,
1056
+ compress_ratio: int,
1057
+ num_decode_tokens: int,
1058
+ token_to_req_indices: torch.Tensor,
1059
+ block_table: torch.Tensor,
1060
+ block_size: int,
1061
+ slot_mapping: torch.Tensor,
1062
+ global_decode_buffer: torch.Tensor,
1063
+ decode_lens_buffer: torch.Tensor,
1064
+ prefill_buffer: torch.Tensor,
1065
+ max_compressed_tokens: int = 8192,
1066
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1067
+ """Single kernel for all C128A tokens (decode + prefill).
1068
+
1069
+ Decode tokens: position → block_table lookup → global slot ids + topk_lens.
1070
+ Prefill tokens: position → local indices [0, ..., n-1, -1, ...].
1071
+
1072
+ Writes into pre-allocated buffers for CUDA graph address stability.
1073
+ Returns slices of the buffers.
1074
+ """
1075
+ num_tokens = positions.shape[0]
1076
+ num_prefill_tokens = num_tokens - num_decode_tokens
1077
+
1078
+ global_decode = global_decode_buffer[:num_decode_tokens]
1079
+ decode_lens = decode_lens_buffer[:num_decode_tokens]
1080
+ prefill_local = prefill_buffer[:num_prefill_tokens]
1081
+
1082
+ if num_tokens == 0:
1083
+ return global_decode, decode_lens, prefill_local
1084
+
1085
+ _build_c128a_topk_metadata_kernel[(num_tokens,)](
1086
+ global_decode_buffer,
1087
+ global_decode_buffer.stride(0),
1088
+ decode_lens_buffer,
1089
+ prefill_buffer,
1090
+ prefill_buffer.stride(0),
1091
+ positions,
1092
+ compress_ratio,
1093
+ max_compressed_tokens,
1094
+ num_decode_tokens,
1095
+ token_to_req_indices,
1096
+ block_table,
1097
+ block_table.stride(0),
1098
+ block_size,
1099
+ slot_mapping,
1100
+ BLOCK_SIZE=1024,
1101
+ )
1102
+ return global_decode, decode_lens, prefill_local
1103
+
1104
+
1105
+@triton.jit
1106
+def _build_c128a_topk_metadata_kernel(
1107
+ # Decode outputs
1108
+ global_decode_ptr,
1109
+ global_decode_stride,
1110
+ decode_lens_ptr,
1111
+ # Prefill output
1112
+ prefill_local_ptr,
1113
+ prefill_local_stride,
1114
+ # Inputs
1115
+ positions_ptr,
1116
+ compress_ratio,
1117
+ max_compressed_tokens,
1118
+ num_decode_tokens,
1119
+ token_to_req_indices_ptr,
1120
+ block_table_ptr,
1121
+ block_table_stride,
1122
+ block_size,
1123
+ slot_mapping_ptr,
1124
+ BLOCK_SIZE: tl.constexpr,
1125
+):
1126
+ token_idx = tl.program_id(0)
1127
+ position = tl.load(positions_ptr + token_idx)
1128
+ num_compressed = (position + 1) // compress_ratio
1129
+ num_compressed = tl.minimum(num_compressed, max_compressed_tokens)
1130
+ is_decode = token_idx < num_decode_tokens
1131
+
1132
+ if is_decode:
1133
+ # --- Decode: block-table lookup → global slot ids + count ---
1134
+ is_valid_token = tl.load(slot_mapping_ptr + token_idx) >= 0
1135
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
1136
+ count = tl.zeros((), dtype=tl.int32)
1137
+ for i in range(0, max_compressed_tokens, BLOCK_SIZE):
1138
+ offset = i + tl.arange(0, BLOCK_SIZE)
1139
+ mask = offset < max_compressed_tokens
1140
+ is_valid = offset < num_compressed
1141
+
1142
+ block_indices = offset // block_size
1143
+ block_numbers = tl.load(
1144
+ block_table_ptr + req_idx * block_table_stride + block_indices,
1145
+ mask=mask & is_valid,
1146
+ )
1147
+ block_offsets = offset % block_size
1148
+ slot_ids = block_numbers * block_size + block_offsets
1149
+ slot_ids = tl.where(is_valid, slot_ids, -1)
1150
+ tl.store(
1151
+ global_decode_ptr + token_idx * global_decode_stride + offset,
1152
+ slot_ids,
1153
+ mask=mask,
1154
+ )
1155
+ count += tl.sum(is_valid.to(tl.int32), axis=0)
1156
+
1157
+ tl.store(
1158
+ decode_lens_ptr + token_idx,
1159
+ tl.where(is_valid_token, count, 0),
1160
+ )
1161
+ else:
1162
+ # --- Prefill: write local indices ---
1163
+ pfx_idx = token_idx - num_decode_tokens
1164
+ for i in range(0, max_compressed_tokens, BLOCK_SIZE):
1165
+ offset = i + tl.arange(0, BLOCK_SIZE)
1166
+ mask = offset < max_compressed_tokens
1167
+ tl.store(
1168
+ prefill_local_ptr + pfx_idx * prefill_local_stride + offset,
1169
+ tl.where(offset < num_compressed, offset, -1),
1170
+ mask=mask,
1171
+ )
vllm/v1/attention/backends/mla/indexer.py
+272
-148
@@ -22,10 +22,11 @@ from vllm.v1.attention.backend import (
22
CommonAttentionMetadata,
23
MultipleOf,
24
)
25
+from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
26
from vllm.v1.attention.backends.utils import (
27
split_decodes_and_prefills,
28
)
28
-from vllm.v1.kv_cache_interface import AttentionSpec
29
+from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
30
from vllm.v1.worker.cp_utils import get_total_cp_world_size
31
32
logger = init_logger(__name__)
@@ -154,6 +155,16 @@ class DeepseekV32IndexerBackend(AttentionBackend):
155
return (0, 1, 2)
156
157
158
+class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
159
+ @staticmethod
160
+ def get_name() -> str:
161
+ return "DEEPSEEK_V4_INDEXER"
162
+
163
+ @staticmethod
164
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
165
+ return [256]
166
+
167
+
168
@dataclass
169
class DeepseekV32IndexerPrefillChunkMetadata:
170
block_table: torch.Tensor
@@ -179,7 +190,7 @@ class DeepSeekV32IndexerDecodeMetadata:
190
# seq_lens: per-token effective context lengths.
191
# - flatten path / plain decode: 1D (batch_size,)
192
# - native MTP path: 2D (B, next_n) where [b,j] = L_b - next_n + j + 1
182
- # Both fp8_paged_mqa_logits and the topk kernels accept both shapes.
193
+ # Both fp8_fp4_paged_mqa_logits and the topk kernels accept both shapes.
194
seq_lens: torch.Tensor
195
decode_lens: torch.Tensor
196
requires_padding: bool
@@ -191,16 +202,8 @@ class DeepseekV32IndexerMetadata:
202
# FIXME (zyongye)
203
# hacky way to access the data now, need to be in chunked meta
204
seq_lens: torch.Tensor
194
-
195
- num_reqs: int
196
- max_query_len: int
205
max_seq_len: int
198
-
199
- num_actual_tokens: int # Number of tokens excluding padding.
200
- query_start_loc: torch.Tensor
206
slot_mapping: torch.Tensor
202
- # The dimension of the attention heads
203
- head_dim: int
207
208
# New for MLA (compared to FlashAttention)
209
# For handling prefill decode split
@@ -213,71 +216,6 @@ class DeepseekV32IndexerMetadata:
216
prefill: DeepseekV32IndexerPrefillMetadata | None = None
217
218
216
-# TODO (zyongye) optimize this, this is now vibe coded
217
-def kv_spans_from_batches(
218
- start_seq_loc: torch.Tensor, seq_len_per_batch: torch.Tensor, device: torch.device
219
-) -> tuple[torch.Tensor, torch.Tensor]:
220
- """
221
- Args:
222
- start_seq_loc: 1D long tensor [B+1], cumulative counts of
223
- selected tokens per batch.
224
- Example: [0, 2, 4, 7] ->
225
- batch sizes (selected) [2, 2, 3], N=7 tokens total.
226
- seq_len_per_batch: 1D long tensor [B],
227
- full sequence length (KV length) of each batch.
228
- Example: [5, 9, 4].
229
-
230
- Returns:
231
- start_tensor: 1D long tensor [N], start offset in the
232
- concatenated KV cache for each token's batch.
233
- end_location: 1D long tensor [N],
234
- **exclusive** end = start + token's local position.
235
- (So the attended KV slice is kv[start:end].)
236
-
237
- Assumes each batch contributes its full `seq_len_per_batch[i]`
238
- keys to the KV cache, andthe selected tokens within a batch
239
- are the **last** `counts[i]` positions of that sequence.
240
- """
241
- q = start_seq_loc.to(dtype=torch.long)
242
- L = seq_len_per_batch.to(dtype=torch.long)
243
- assert q.dim() == 1 and L.dim() == 1
244
- assert q.numel() == L.numel() + 1, "start_seq_loc must have length B+1"
245
-
246
- # Selected tokens per batch and totals
247
- counts = q[1:] - q[:-1] # [B]
248
- N = int(q[-1].item()) # total selected tokens
249
- B = L.numel()
250
-
251
- if N == 0:
252
- return (
253
- torch.empty(0, dtype=torch.long, device=device),
254
- torch.empty(0, dtype=torch.long, device=device),
255
- )
256
-
257
- # KV start offsets per batch in the concatenated KV cache
258
- kv_starts_per_batch = torch.cumsum(L, dim=0) - L # [B]
259
-
260
- # For each selected token, which batch does it belong to?
261
- batch_id = torch.repeat_interleave(torch.arange(B), counts) # [N]
262
-
263
- # Map batch KV start to each token
264
- start_tensor = kv_starts_per_batch[batch_id] # [N]
265
-
266
- # End-align local positions inside each batch:
267
- # local_pos = L[b] - counts[b] + (1..counts[b]) for each batch b
268
- L_expand = torch.repeat_interleave(L, counts) # [N]
269
- m_expand = torch.repeat_interleave(counts, counts) # [N]
270
- # position within the selected block: 1..counts[b]
271
- pos_within = (
272
- torch.arange(N, dtype=torch.long) - torch.repeat_interleave(q[:-1], counts) + 1
273
- )
274
-
275
- local_pos = L_expand - m_expand + pos_within # [N], 1-based
276
- end_location = start_tensor + local_pos # exclusive end
277
-
278
- return start_tensor.int().to(device), end_location.int().to(device)
279
-
280
-
219
def get_max_prefill_buffer_size(vllm_config: VllmConfig):
220
max_model_len = vllm_config.model_config.max_model_len
221
# NOTE(Chen): 40 is a magic number for controlling the prefill buffer size.
@@ -293,7 +231,7 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
231
232
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
233
reorder_batch_threshold: int = 1
296
- natively_supported_next_n: list[int] = [1, 2]
234
+ natively_supported_next_n_fp4: list[int] = [1, 2]
235
# TODO (matt): integrate kernel with next_n = 4 support
236
237
@classmethod
@@ -314,9 +252,30 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
252
if self.vllm_config.speculative_config
253
else 0
254
)
255
+ self.use_fp4_indexer_cache = (
256
+ self.vllm_config.attention_config.use_fp4_indexer_cache
257
+ )
258
+
259
+ assert (
260
+ current_platform.is_device_capability_family(100)
261
+ or not self.use_fp4_indexer_cache
262
+ ), (
263
+ "use_fp4_indexer_cache requires Blackwell datacenter GPUs "
264
+ "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
265
+ "earlier architectures are not supported."
266
+ )
267
+
268
next_n = self.num_speculative_tokens + 1
269
self.reorder_batch_threshold += self.num_speculative_tokens
319
- self.use_flattening = next_n not in self.natively_supported_next_n
270
+ # NOTE(zyongye) fp4 indexer cache only natively supports next_n in
271
+ # natively_supported_next_n_fp4; for other next_n values we fall back
272
+ # to the flattening path. Outside the SM100 datacenter family the FP8
273
+ # paged MQA logits kernel has the same [1, 2] constraint (deepgemm
274
+ # smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there too.
275
+ self.use_flattening = (
276
+ self.use_fp4_indexer_cache
277
+ or not current_platform.is_device_capability_family(100)
278
+ ) and next_n not in self.natively_supported_next_n_fp4
279
280
sm_count = num_compute_units(self.device.index)
281
self.num_sms = sm_count
@@ -331,7 +290,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
290
)
291
if not self.use_flattening and next_n > 1:
292
# Native MTP: 2D buffer for per-token seq_lens.
334
- # Flattening path is never used, so no expanded_seq_lens_buffer.
293
self.decode_seq_lens_buffer = torch.zeros(
294
(scheduler_config.max_num_seqs, next_n),
295
dtype=torch.int32,
@@ -367,53 +325,27 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
325
(self.num_sms + 1, 2), dtype=torch.int32, device=self.device
326
)
327
370
- def build_one_prefill_chunk(
371
- self,
372
- req_slice: slice,
373
- query_slice: slice,
374
- query_start_loc_cpu,
375
- seq_lens_cpu,
376
- block_table,
377
- skip_kv_gather: bool = False,
378
- ) -> DeepseekV32IndexerPrefillChunkMetadata:
379
- prefill_query_start_loc = (
380
- query_start_loc_cpu[req_slice.start : req_slice.stop + 1]
381
- - query_start_loc_cpu[req_slice.start]
382
- )
383
- cu_seqlen_ks, cu_seqlen_ke = kv_spans_from_batches(
384
- prefill_query_start_loc, seq_lens_cpu[req_slice], self.device
385
- )
386
- token_start = query_start_loc_cpu[req_slice.start].item()
387
- total_seq_lens = seq_lens_cpu[req_slice].sum()
388
- num_reqs = req_slice.stop - req_slice.start
389
- seq_idx = torch.arange(0, num_reqs, dtype=torch.int32)
390
- token_to_seq = torch.repeat_interleave(seq_idx, seq_lens_cpu[req_slice]).to(
391
- self.device
392
- )
393
- assert total_seq_lens <= self.max_prefill_buffer_size
394
- cu_seq_lens = (
395
- torch.cat(
396
- [
397
- torch.zeros(1, dtype=torch.int32),
398
- seq_lens_cpu[req_slice].cumsum(dim=0),
399
- ]
328
+ # KV compression. Default to 1 for no compression.
329
+ self.compress_ratio = 1
330
+ # Get compress_ratio for DeepseekV4 support
331
+ if isinstance(self.kv_cache_spec, MLAAttentionSpec):
332
+ self.compress_ratio = self.kv_cache_spec.compress_ratio
333
+
334
+ # Pre-allocate buffers for CUDA graph compatibility when
335
+ if self.compress_ratio > 1:
336
+ # compress_ratio > 1 (DeepseekV4)
337
+ # Compressed slot mapping output buffer
338
+ self.compressed_slot_mapping_buffer = torch.zeros(
339
+ (scheduler_config.max_num_batched_tokens,),
340
+ dtype=torch.int64,
341
+ device=self.device,
342
+ )
343
+ # Buffer for compressed seq_lens in decode path
344
+ self.expanded_seq_lens_buffer = torch.zeros(
345
+ (scheduler_config.max_num_batched_tokens,),
346
+ dtype=torch.int32,
347
+ device=self.device,
348
)
401
- .to(torch.int32)
402
- .to(self.device)
403
- )
404
-
405
- return DeepseekV32IndexerPrefillChunkMetadata(
406
- cu_seqlen_ks=cu_seqlen_ks[query_slice],
407
- cu_seqlen_ke=cu_seqlen_ke[query_slice],
408
- cu_seq_lens=cu_seq_lens,
409
- token_to_seq=token_to_seq,
410
- total_seq_lens=total_seq_lens,
411
- block_table=block_table[req_slice],
412
- token_start=token_start + query_slice.start,
413
- token_end=token_start + query_slice.stop,
414
- num_reqs=num_reqs,
415
- skip_kv_gather=skip_kv_gather,
416
- )
349
350
def _prepare_decode_tensors(
351
self,
@@ -520,11 +452,15 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
452
requires_padding = min_decode_len != max_decode_len
453
if use_native and next_n > 1:
454
assert self.decode_seq_lens_buffer.dim() == 2
523
- # (B, next_n): token j attends to L - next_n + j + 1 KV tokens
524
- self.decode_seq_lens_buffer[:num_decodes] = (
525
- seq_lens.unsqueeze(1) - next_n + 1 + self.offsets_buffer
455
+ # (B, max_decode_len): token j attends to
456
+ # L - max_decode_len + j + 1 KV tokens.
457
+ self.decode_seq_lens_buffer[:num_decodes, :max_decode_len] = (
458
+ seq_lens.unsqueeze(1)
459
+ - max_decode_len
460
+ + 1
461
+ + self.offsets_buffer[:max_decode_len]
462
)
527
- seq_lens = self.decode_seq_lens_buffer[:num_decodes]
463
+ seq_lens = self.decode_seq_lens_buffer[:num_decodes, :max_decode_len]
464
return seq_lens, block_table, decode_lens, num_decodes, requires_padding
465
466
def build(
@@ -535,8 +471,12 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
471
) -> DeepseekV32IndexerMetadata:
472
num_reqs = common_attn_metadata.num_reqs
473
num_tokens = common_attn_metadata.num_actual_tokens
538
-
474
+ query_start_loc = common_attn_metadata.query_start_loc
475
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
476
+ seq_lens = common_attn_metadata.seq_lens
477
+ slot_mapping = common_attn_metadata.slot_mapping
478
+ block_table = common_attn_metadata.block_table_tensor
479
+
480
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
481
split_decodes_and_prefills(
482
common_attn_metadata,
@@ -548,8 +488,32 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
488
assert num_decodes + num_prefills == num_reqs
489
assert num_decode_tokens + num_prefill_tokens == num_tokens
490
491
+ compressed_slot_mapping = slot_mapping
492
+ compressed_seq_lens = seq_lens
493
+ if self.compress_ratio > 1:
494
+ compressed_slot_mapping = get_compressed_slot_mapping(
495
+ num_tokens,
496
+ query_start_loc,
497
+ seq_lens,
498
+ block_table,
499
+ self.kv_cache_spec.storage_block_size,
500
+ self.compress_ratio,
501
+ out=self.compressed_slot_mapping_buffer,
502
+ )
503
+ compressed_seq_lens = seq_lens // self.compress_ratio
504
+
505
prefill_metadata = None
506
if num_prefills > 0:
507
+ # This CPU value is an upper bound for async-spec extend rows. It
508
+ # is safe for chunking/allocation because CUDA metadata below is
509
+ # built from exact device seq_lens and gather ignores the tail.
510
+ assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
511
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
512
+ compressed_seq_lens_cpu = (
513
+ seq_lens_cpu // self.compress_ratio
514
+ if self.compress_ratio > 1
515
+ else seq_lens_cpu
516
+ )
517
prefill_query_lens_cpu = torch.diff(
518
query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1]
519
)
@@ -559,26 +523,32 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
523
assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
524
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
525
chunk_specs = split_indexer_prefill_chunks(
562
- seq_lens_cpu[num_decodes:],
526
+ compressed_seq_lens_cpu[num_decodes:],
527
prefill_query_lens_cpu,
528
self.max_prefill_buffer_size,
529
max_logits_bytes,
530
request_offset=num_decodes,
531
)
568
- chunks = [
569
- self.build_one_prefill_chunk(
570
- req_slice,
571
- query_slice,
532
+
533
+ chunks = []
534
+ for req_slice, query_slice in chunk_specs:
535
+ metadata = build_prefill_chunk_metadata(
536
+ req_slice.start,
537
+ req_slice.stop,
538
+ query_start_loc,
539
query_start_loc_cpu,
573
- seq_lens_cpu,
540
+ seq_lens,
541
+ compressed_seq_lens,
542
+ compressed_seq_lens_cpu,
543
common_attn_metadata.block_table_tensor,
544
+ self.compress_ratio,
545
+ query_slice=query_slice,
546
skip_kv_gather=query_slice.start > 0,
547
)
577
- for req_slice, query_slice in chunk_specs
578
- ]
579
- prefill_metadata = DeepseekV32IndexerPrefillMetadata(
580
- chunks=chunks,
581
- )
548
+ # Skip when total_seq_lens is 0 (i.e., no compressed token).
549
+ if metadata is not None:
550
+ chunks.append(metadata)
551
+ prefill_metadata = DeepseekV32IndexerPrefillMetadata(chunks)
552
553
decode_metadata = None
554
if num_decodes > 0:
@@ -596,7 +566,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
566
567
max_decode_len = int(decode_lens_cpu.max().item())
568
next_n = 1 + self.num_speculative_tokens
599
- use_native = not self.use_flattening and max_decode_len == next_n
569
+ use_native = not self.use_flattening and max_decode_len <= next_n
570
571
seq_lens, block_table, decode_lens, batch_size, requires_padding = (
572
self._prepare_decode_tensors(
@@ -613,11 +583,35 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
583
)
584
)
585
586
+ # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
587
+ # compressed tokens. Convert uncompressed seq_lens to compressed.
588
+ if self.compress_ratio > 1:
589
+ # True iff seq_lens aliases decode_seq_lens_buffer (flatten or
590
+ # native wrote it); False iff it aliases common_attn_metadata.
591
+ seq_lens_is_local_view = (use_native and next_n > 1) or (
592
+ not use_native and max_decode_len > 1
593
+ )
594
+ if seq_lens_is_local_view:
595
+ seq_lens //= self.compress_ratio
596
+ else:
597
+ # Copy to avoid mutating shared state; keeps CG address stable.
598
+ self.expanded_seq_lens_buffer[:num_decodes] = (
599
+ seq_lens // self.compress_ratio
600
+ )
601
+ self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0
602
+ seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]
603
+
604
+ # Non-MTP: deep_gemm paged MQA logits requires 2D context_lens
605
+ # (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream
606
+ # kernels see the same (B, next_n) layout as the MTP path.
607
+ if seq_lens.dim() == 1:
608
+ seq_lens = seq_lens.unsqueeze(-1)
609
+
610
# DeepGEMM is required for the paged MQA logits on CUDA devices
611
if current_platform.is_cuda() and has_deep_gemm():
612
self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
613
seq_lens,
620
- self.kv_cache_spec.block_size,
614
+ self.kv_cache_spec.storage_block_size,
615
self.num_sms,
616
)
617
@@ -631,13 +625,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
625
626
attn_metadata = DeepseekV32IndexerMetadata(
627
seq_lens=common_attn_metadata.seq_lens,
634
- num_reqs=common_attn_metadata.num_reqs,
635
- max_query_len=common_attn_metadata.max_query_len,
628
max_seq_len=common_attn_metadata.max_seq_len,
637
- num_actual_tokens=common_attn_metadata.num_actual_tokens,
638
- query_start_loc=common_attn_metadata.query_start_loc,
639
- slot_mapping=common_attn_metadata.slot_mapping,
640
- head_dim=128,
629
+ slot_mapping=compressed_slot_mapping,
630
num_decodes=num_decodes,
631
num_decode_tokens=num_decode_tokens,
632
num_prefills=num_prefills,
@@ -647,3 +636,138 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
636
)
637
638
return attn_metadata
639
+
640
+
641
+def build_prefill_chunk_metadata(
642
+ start_idx: int,
643
+ end_idx: int,
644
+ query_start_loc: torch.Tensor,
645
+ query_start_loc_cpu: torch.Tensor,
646
+ uncompressed_seq_lens: torch.Tensor,
647
+ compressed_seq_lens: torch.Tensor,
648
+ compressed_seq_lens_cpu: torch.Tensor,
649
+ block_table: torch.Tensor,
650
+ compress_ratio: int,
651
+ query_slice: slice | None = None,
652
+ skip_kv_gather: bool = False,
653
+) -> DeepseekV32IndexerPrefillChunkMetadata | None:
654
+ total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
655
+ if total_seq_lens == 0:
656
+ return None
657
+
658
+ num_reqs = end_idx - start_idx
659
+ device = block_table.device
660
+ token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device)
661
+
662
+ cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
663
+ # Assigning to slice avoids cpu sync.
664
+ cu_seq_lens[:1] = 0
665
+ torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
666
+
667
+ query_start_loc = (
668
+ query_start_loc[start_idx : end_idx + 1] - query_start_loc[start_idx]
669
+ )
670
+
671
+ total_query_len = int(
672
+ (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
673
+ )
674
+ if query_slice is not None:
675
+ qs_start = query_slice.start
676
+ qs_stop = query_slice.stop
677
+ else:
678
+ qs_start = 0
679
+ qs_stop = total_query_len
680
+ output_query_len = qs_stop - qs_start
681
+
682
+ cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
683
+ cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
684
+
685
+ _build_prefill_chunk_metadata_kernel[(num_reqs,)](
686
+ query_start_loc,
687
+ uncompressed_seq_lens[start_idx:end_idx],
688
+ cu_seq_lens,
689
+ token_to_seq,
690
+ cu_seq_len_ks,
691
+ cu_seq_len_ke,
692
+ qs_start,
693
+ qs_stop,
694
+ BLOCK_SIZE=1024,
695
+ COMPRESS_RATIO=compress_ratio,
696
+ )
697
+
698
+ token_start = query_start_loc_cpu[start_idx].item()
699
+ if query_slice is not None:
700
+ token_end = token_start + qs_stop
701
+ token_start = token_start + qs_start
702
+ skip_kv_gather = skip_kv_gather or qs_start > 0
703
+ else:
704
+ token_end = query_start_loc_cpu[end_idx].item()
705
+
706
+ return DeepseekV32IndexerPrefillChunkMetadata(
707
+ cu_seqlen_ks=cu_seq_len_ks,
708
+ cu_seqlen_ke=cu_seq_len_ke,
709
+ cu_seq_lens=cu_seq_lens,
710
+ token_to_seq=token_to_seq,
711
+ total_seq_lens=total_seq_lens,
712
+ block_table=block_table[start_idx:end_idx],
713
+ token_start=token_start,
714
+ token_end=token_end,
715
+ num_reqs=num_reqs,
716
+ skip_kv_gather=skip_kv_gather,
717
+ )
718
+
719
+
720
+@triton.jit
721
+def _build_prefill_chunk_metadata_kernel(
722
+ # Inputs
723
+ query_start_loc_ptr,
724
+ uncompressed_seq_lens_ptr,
725
+ cu_compressed_seq_lens_ptr,
726
+ # Outputs
727
+ token_to_seq_ptr,
728
+ cu_compressed_seq_len_ks_ptr,
729
+ cu_compressed_seq_len_ke_ptr,
730
+ query_slice_start,
731
+ query_slice_stop,
732
+ BLOCK_SIZE: tl.constexpr,
733
+ COMPRESS_RATIO: tl.constexpr,
734
+):
735
+ batch_idx = tl.program_id(0)
736
+
737
+ query_start = tl.load(query_start_loc_ptr + batch_idx)
738
+ query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
739
+ query_len = query_end - query_start
740
+
741
+ seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
742
+ seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
743
+ compressed_seq_len = seq_end - seq_start
744
+
745
+ uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
746
+ start_pos = uncompressed_seq_len - query_len
747
+
748
+ for i in range(0, query_len, BLOCK_SIZE):
749
+ offset = i + tl.arange(0, BLOCK_SIZE)
750
+ abs_pos = query_start + offset
751
+ mask = (
752
+ (offset < query_len)
753
+ & (abs_pos >= query_slice_start)
754
+ & (abs_pos < query_slice_stop)
755
+ )
756
+ out_pos = abs_pos - query_slice_start
757
+
758
+ # Compute cu_seq_len_ks
759
+ tl.store(cu_compressed_seq_len_ks_ptr + out_pos, seq_start, mask=mask)
760
+
761
+ # Compute cu_seq_len_ke
762
+ seq_len_per_token = (start_pos + 1 + offset) // COMPRESS_RATIO
763
+ tl.store(
764
+ cu_compressed_seq_len_ke_ptr + out_pos,
765
+ seq_start + seq_len_per_token,
766
+ mask=mask,
767
+ )
768
+
769
+ # Compute token_to_seq
770
+ for i in range(0, compressed_seq_len, BLOCK_SIZE):
771
+ offset = i + tl.arange(0, BLOCK_SIZE)
772
+ mask = offset < compressed_seq_len
773
+ tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)
vllm/v1/attention/backends/mla/sparse_swa.py
new
+492
@@ -0,0 +1,492 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+from typing import ClassVar, cast
5
+
6
+import torch
7
+
8
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
9
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10
+from vllm.triton_utils import tl, triton
11
+from vllm.v1.attention.backend import (
12
+ AttentionBackend,
13
+ AttentionCGSupport,
14
+ AttentionMetadataBuilder,
15
+ CommonAttentionMetadata,
16
+ MultipleOf,
17
+)
18
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
19
+from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata
20
+from vllm.v1.kv_cache_interface import (
21
+ KVCacheSpec,
22
+ MLAAttentionSpec,
23
+ SlidingWindowMLASpec,
24
+)
25
+
26
+# DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct
27
+# (topk, extra_topk, extra_page_block_size) config, so they cannot share a
28
+# FlashMLA tile-scheduler plan. Within a type, all ~60 DeepseekV4 layers share one
29
+# plan per step because b / s_q / h_q / page_block_sizes / topks are identical.
30
+_LAYER_TYPE_SWAONLY = "swaonly"
31
+_LAYER_TYPE_C4A = "c4a"
32
+_LAYER_TYPE_C128A = "c128a"
33
+
34
+
35
+def _layer_type_for(compress_ratio: int) -> str:
36
+ if compress_ratio <= 1:
37
+ return _LAYER_TYPE_SWAONLY
38
+ if compress_ratio == 4:
39
+ return _LAYER_TYPE_C4A
40
+ if compress_ratio == 128:
41
+ return _LAYER_TYPE_C128A
42
+ raise ValueError(
43
+ f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; "
44
+ "expected 1, 4, or 128."
45
+ )
46
+
47
+
48
+class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
49
+ def __init__(
50
+ self,
51
+ head_dim: int,
52
+ window_size: int,
53
+ dtype: torch.dtype,
54
+ prefix: str,
55
+ cache_config: CacheConfig,
56
+ ):
57
+ super().__init__()
58
+ self.kv_cache = torch.tensor([])
59
+ self.head_dim = head_dim
60
+ self.window_size = window_size
61
+ self.prefix = prefix
62
+ self.cache_config = cache_config
63
+ self.dtype = dtype
64
+ compilation_config = get_current_vllm_config().compilation_config
65
+ if prefix in compilation_config.static_forward_context:
66
+ raise ValueError(f"Duplicate layer name: {prefix}")
67
+ compilation_config.static_forward_context[prefix] = self
68
+
69
+ # Block size is constrained by tensor sharing between SWA and C4A KV blocks.
70
+ # Since both block types share the same physical tensor, they must use the
71
+ # same page size. The C4A KV block shape [256//4, head_dim] = [64, head_dim]
72
+ # determines the SWA block size of 64 tokens per block.
73
+ # TODO(yifan): make SWA block size automatically determined and configurable.
74
+ self.block_size = 64
75
+ assert self.dtype == torch.uint8
76
+
77
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
78
+ return SlidingWindowMLASpec(
79
+ block_size=self.block_size,
80
+ num_kv_heads=1,
81
+ head_size=self.head_dim,
82
+ dtype=self.dtype,
83
+ sliding_window=self.window_size,
84
+ cache_dtype_str=self.cache_config.cache_dtype,
85
+ alignment=576, # NOTE: FlashMLA requires 576B alignment
86
+ model_version="deepseek_v4",
87
+ )
88
+
89
+ def forward(self): ...
90
+
91
+ def get_attn_backend(self) -> type[AttentionBackend]:
92
+ return DeepseekSparseSWABackend
93
+
94
+
95
+class DeepseekSparseSWABackend(AttentionBackend):
96
+ @staticmethod
97
+ def get_name() -> str:
98
+ return "DEEPSEEK_SPARSE_SWA"
99
+
100
+ @staticmethod
101
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
102
+ return [MultipleOf(64)]
103
+
104
+ @classmethod
105
+ def get_preferred_block_size(cls, default_block_size: int) -> int:
106
+ return 256
107
+
108
+ @classmethod
109
+ def get_supported_head_sizes(cls) -> list[int]:
110
+ return [512]
111
+
112
+ @staticmethod
113
+ def get_builder_cls() -> type["DeepseekSparseSWAMetadataBuilder"]:
114
+ return DeepseekSparseSWAMetadataBuilder
115
+
116
+ @staticmethod
117
+ def get_kv_cache_shape(
118
+ num_blocks: int,
119
+ block_size: int,
120
+ num_kv_heads: int,
121
+ head_size: int,
122
+ cache_dtype_str: str = "auto",
123
+ ) -> tuple[int, ...]:
124
+ assert num_kv_heads == 1
125
+ if cache_dtype_str == "fp8_ds_mla":
126
+ # DeepseekV4 SWA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
127
+ # head_size passed in is the semantic head_dim (512).
128
+ return (num_blocks, block_size, 584)
129
+ else:
130
+ return (num_blocks, block_size, head_size)
131
+
132
+ @staticmethod
133
+ def get_kv_cache_stride_order(
134
+ include_num_layers_dimension: bool = False,
135
+ ) -> tuple[int, ...]:
136
+ if include_num_layers_dimension:
137
+ return (0, 1, 2, 3)
138
+ return (0, 1, 2)
139
+
140
+
141
+@dataclass
142
+class DeepseekSparseSWAMetadata:
143
+ block_table: torch.Tensor
144
+ slot_mapping: torch.Tensor
145
+ block_size: int
146
+ seq_lens: torch.Tensor | None = None # [num_seqs]
147
+ query_start_loc: torch.Tensor | None = None # [num_seqs + 1]
148
+ query_start_loc_cpu: torch.Tensor | None = None # [num_seqs + 1]
149
+
150
+ is_valid_token: torch.Tensor | None = None # [num_tokens]
151
+ token_to_req_indices: torch.Tensor | None = None # [num_tokens]
152
+ decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, window_size]
153
+ decode_swa_lens: torch.Tensor | None = None # [num_decode_tokens]
154
+
155
+ # Number of decode/prefill requests/tokens (batch is reordered: decodes first)
156
+ num_decodes: int = 0
157
+ num_prefills: int = 0
158
+ num_decode_tokens: int = 0
159
+ num_prefill_tokens: int = 0
160
+
161
+ # Pre-computed prefill metadata shared across all DeepseekV4 attention layers.
162
+ prefill_seq_lens: torch.Tensor | None = None
163
+ prefill_gather_lens: torch.Tensor | None = None
164
+
165
+ # Per-layer-type FlashMLA tile-scheduler metadata. One FlashMLASchedMeta
166
+ # per present DeepseekV4 layer type, shared across all ~60 layers of that type
167
+ # within a decode step. The first forward call of a given type triggers
168
+ # the in-kernel planner (which also allocates tile_scheduler_metadata and
169
+ # num_splits via PyTorch's graph-aware allocator); subsequent same-type
170
+ # calls skip planning and reuse the plan. Fresh instance per build(), so
171
+ # have_initialized is always False at the start of a step and the plan
172
+ # is re-derived from current seq_lens / topk_length on replay.
173
+ # None for layer types the model does not use (or when num_decode_tokens
174
+ # is zero).
175
+ tile_sched_swaonly: "FlashMLASchedMeta | None" = None
176
+ tile_sched_c4a: "FlashMLASchedMeta | None" = None
177
+ tile_sched_c128a: "FlashMLASchedMeta | None" = None
178
+
179
+
180
+class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
181
+ """Builds metadata for DeepseekV4 SWA cache.
182
+
183
+ Similar to the indexer, this handles mixed batches by:
184
+ 1. Using split_decodes_and_prefills() to determine the boundary
185
+ 2. Building separate metadata for decode and prefill portions
186
+
187
+ Supports:
188
+ - Mixed decode/prefill batches
189
+ - MTP (Multi-Token Prediction) where decode has query_len > 1
190
+ - Chunked prefill (aligns with the indexer's chunking)
191
+ """
192
+
193
+ # Base threshold: query_len <= 1 is decode
194
+ reorder_batch_threshold: int = 1
195
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
196
+
197
+ def __init__(self, *args, **kwargs):
198
+ super().__init__(*args, **kwargs)
199
+ assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec)
200
+ mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec)
201
+ self.head_size = mla_spec.head_size # Already considered quantization.
202
+ self.compress_ratio = mla_spec.compress_ratio
203
+ self.block_size = mla_spec.block_size
204
+
205
+ # Handle MTP: adjust decode_threshold like the indexer does
206
+ self.num_speculative_tokens = (
207
+ self.vllm_config.speculative_config.num_speculative_tokens
208
+ if self.vllm_config.speculative_config
209
+ else 0
210
+ )
211
+ # With MTP, decode can have query_len up to 1 + num_speculative_tokens.
212
+ # Must match the threshold used by the indexer and flashmla_sparse so
213
+ # that all backends agree on the decode/prefill split.
214
+ self.decode_threshold = (
215
+ self.reorder_batch_threshold + self.num_speculative_tokens
216
+ )
217
+
218
+ hf_config = self.vllm_config.model_config.hf_config
219
+ assert hasattr(hf_config, "sliding_window")
220
+ self.window_size = hf_config.sliding_window
221
+
222
+ # Detect which DeepseekV4 layer types this model uses so we only build a
223
+ # FlashMLA tile-scheduler plan for types that will actually be called.
224
+ # Models without compress_ratios (pure SWA) fall back to swaonly.
225
+ compress_ratios = getattr(hf_config, "compress_ratios", None) or [1]
226
+ self._layer_types: set[str] = set()
227
+ for ratio in compress_ratios:
228
+ self._layer_types.add(_layer_type_for(int(ratio)))
229
+
230
+ max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
231
+ self.token_to_req_indices = torch.zeros(
232
+ max_tokens,
233
+ dtype=torch.int32,
234
+ device=self.device,
235
+ )
236
+ self.decode_swa_indices = torch.zeros(
237
+ max_tokens,
238
+ 1,
239
+ self.window_size,
240
+ dtype=torch.int32,
241
+ device=self.device,
242
+ )
243
+ self.decode_swa_lens = torch.zeros(
244
+ max_tokens,
245
+ dtype=torch.int32,
246
+ device=self.device,
247
+ )
248
+ self.is_valid_token = torch.zeros(
249
+ max_tokens,
250
+ dtype=torch.bool,
251
+ device=self.device,
252
+ )
253
+
254
+ def build(
255
+ self,
256
+ common_prefix_len: int,
257
+ common_attn_metadata: CommonAttentionMetadata,
258
+ fast_build: bool = False,
259
+ ) -> DeepseekSparseSWAMetadata:
260
+ """Build SWA metadata for mixed decode/prefill batches.
261
+
262
+ The batch is assumed to be reordered with decodes first (by vLLM scheduler).
263
+ We use split_decodes_and_prefills() to find the boundary, then build
264
+ separate window_topk_idxs for each portion.
265
+
266
+ For prefill, we use chunked prefill to align with the indexer's chunking.
267
+ """
268
+ num_reqs = common_attn_metadata.num_reqs
269
+ seq_lens = common_attn_metadata.seq_lens
270
+ query_start_loc = common_attn_metadata.query_start_loc
271
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
272
+ block_table = common_attn_metadata.block_table_tensor
273
+ slot_mapping = common_attn_metadata.slot_mapping
274
+
275
+ # Split into decode and prefill portions using configurable threshold
276
+ (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = (
277
+ split_decodes_and_prefills(
278
+ common_attn_metadata, decode_threshold=self.decode_threshold
279
+ )
280
+ )
281
+
282
+ # NOTE: Ensure all metadata tensors maintain fixed memory addresses
283
+ # for CUDA graph compatibility.
284
+ query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
285
+ x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory()
286
+ token_to_req_indices = self.token_to_req_indices[: x.shape[0]]
287
+ token_to_req_indices.copy_(x, non_blocking=True)
288
+
289
+ is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
290
+ is_valid_token.copy_(slot_mapping >= 0)
291
+
292
+ if num_decode_tokens > 0:
293
+ self.decode_swa_lens[num_decode_tokens:] = 0
294
+ _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
295
+ self.decode_swa_indices,
296
+ self.decode_swa_indices.stride(0),
297
+ self.decode_swa_lens,
298
+ self.window_size,
299
+ query_start_loc,
300
+ seq_lens,
301
+ token_to_req_indices,
302
+ is_valid_token,
303
+ block_table,
304
+ block_table.stride(0),
305
+ self.block_size,
306
+ TRITON_BLOCK_SIZE=1024,
307
+ )
308
+
309
+ # Pre-compute DeepseekV4 prefill metadata shared across all attention layers.
310
+ deepseek_v4_fields = self._build_deepseek_v4_metadata(
311
+ num_decodes,
312
+ num_prefills,
313
+ seq_lens,
314
+ query_start_loc,
315
+ )
316
+
317
+ # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta
318
+ # per present DeepseekV4 layer type; the first flash_mla_with_kvcache call of
319
+ # each type triggers the planner and all same-type layers reuse the
320
+ # resulting plan for the rest of the step.
321
+ tile_sched = self.build_tile_scheduler(num_decode_tokens)
322
+
323
+ return DeepseekSparseSWAMetadata(
324
+ seq_lens=seq_lens,
325
+ query_start_loc=query_start_loc,
326
+ query_start_loc_cpu=query_start_loc_cpu,
327
+ block_table=block_table,
328
+ slot_mapping=slot_mapping,
329
+ is_valid_token=is_valid_token,
330
+ token_to_req_indices=token_to_req_indices,
331
+ decode_swa_indices=self.decode_swa_indices[:num_decode_tokens],
332
+ decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
333
+ block_size=self.block_size,
334
+ num_decodes=num_decodes,
335
+ num_prefills=num_prefills,
336
+ num_decode_tokens=num_decode_tokens,
337
+ num_prefill_tokens=num_prefill_tokens,
338
+ tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY],
339
+ tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A],
340
+ tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A],
341
+ **deepseek_v4_fields,
342
+ )
343
+
344
+ def build_tile_scheduler(
345
+ self, num_decode_tokens: int
346
+ ) -> dict[str, FlashMLASchedMeta | None]:
347
+ """Allocate one empty ``FlashMLASchedMeta`` per present DeepseekV4 layer type.
348
+
349
+ Returned instances have ``tile_scheduler_metadata`` / ``num_splits``
350
+ set to ``None``; the FlashMLA C++ decode path will allocate them and
351
+ run the tile-scheduler planner on the first ``flash_mla_with_kvcache``
352
+ call of each type. Subsequent same-type calls reuse the plan because
353
+ the tensors (and ``have_initialized``) are populated on the struct.
354
+
355
+ Returns all-``None`` when there are no decode tokens this step, so
356
+ ``_forward_decode`` sees a clean sentinel.
357
+ """
358
+ out: dict[str, FlashMLASchedMeta | None] = {
359
+ _LAYER_TYPE_SWAONLY: None,
360
+ _LAYER_TYPE_C4A: None,
361
+ _LAYER_TYPE_C128A: None,
362
+ }
363
+ if num_decode_tokens == 0:
364
+ return out
365
+ for layer_type in self._layer_types:
366
+ # get_mla_metadata() is the official FlashMLA entry point that
367
+ # returns a fresh empty FlashMLASchedMeta; using it keeps this
368
+ # call site aligned with the rest of the vLLM FlashMLA backends
369
+ # that already go through the same stub.
370
+ out[layer_type] = get_mla_metadata()[0]
371
+ return out
372
+
373
+ def _build_deepseek_v4_metadata(
374
+ self,
375
+ num_decodes: int,
376
+ num_prefills: int,
377
+ seq_lens: torch.Tensor,
378
+ query_start_loc: torch.Tensor,
379
+ ) -> dict[str, torch.Tensor | None]:
380
+ """Pre-compute DeepseekV4 prefill metadata during the metadata build phase.
381
+
382
+ Returns a dict of keyword arguments to pass to the
383
+ DeepseekSparseSWAMetadata constructor.
384
+
385
+ Note: C128A topk indices are computed by the FlashMLASparse builder
386
+ (which owns the C128A block_table), not here.
387
+ """
388
+ result: dict[str, torch.Tensor | None] = {}
389
+
390
+ # --- Prefill query metadata (single Triton kernel + CPU slicing) ---
391
+ if num_prefills > 0:
392
+ pfx_gather_lens = torch.empty(
393
+ num_prefills, dtype=torch.int32, device=seq_lens.device
394
+ )
395
+ _compute_prefill_metadata_kernel[(1,)](
396
+ pfx_gather_lens,
397
+ seq_lens,
398
+ query_start_loc,
399
+ num_prefills,
400
+ num_decodes,
401
+ self.window_size,
402
+ BLOCK_SIZE=triton.next_power_of_2(num_prefills),
403
+ )
404
+
405
+ result["prefill_seq_lens"] = seq_lens[num_decodes:]
406
+ result["prefill_gather_lens"] = pfx_gather_lens
407
+
408
+ return result
409
+
410
+
411
+@triton.jit
412
+def _compute_prefill_metadata_kernel(
413
+ # Outputs
414
+ prefill_gather_lens_ptr,
415
+ # Inputs
416
+ seq_lens_ptr,
417
+ query_start_loc_ptr,
418
+ num_prefills,
419
+ num_decodes,
420
+ window_size,
421
+ BLOCK_SIZE: tl.constexpr,
422
+):
423
+ """Compute prefill gather_lens in a single pass."""
424
+ offset = tl.arange(0, BLOCK_SIZE)
425
+ mask = offset < num_prefills
426
+
427
+ seq_len = tl.load(seq_lens_ptr + num_decodes + offset, mask=mask)
428
+ qsl_start = tl.load(query_start_loc_ptr + num_decodes + offset, mask=mask)
429
+ qsl_end = tl.load(query_start_loc_ptr + num_decodes + offset + 1, mask=mask)
430
+
431
+ query_len = qsl_end - qsl_start
432
+ prefix_len = seq_len - query_len
433
+ gather_len = query_len + tl.minimum(prefix_len, window_size - 1)
434
+
435
+ tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
436
+
437
+
438
+@triton.jit
439
+def _compute_swa_indices_and_lens_kernel(
440
+ swa_indices_ptr,
441
+ swa_indices_stride,
442
+ swa_lens_ptr,
443
+ window_size,
444
+ query_start_loc_ptr,
445
+ seq_lens_ptr,
446
+ token_to_req_indices_ptr,
447
+ is_valid_token_ptr,
448
+ block_table_ptr,
449
+ block_table_stride,
450
+ block_size,
451
+ TRITON_BLOCK_SIZE: tl.constexpr,
452
+):
453
+ token_idx = tl.program_id(0)
454
+ is_valid = tl.load(is_valid_token_ptr + token_idx)
455
+ if not is_valid:
456
+ tl.store(swa_lens_ptr + token_idx, 0)
457
+ return
458
+
459
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
460
+
461
+ query_start = tl.load(query_start_loc_ptr + req_idx)
462
+ query_end = tl.load(query_start_loc_ptr + req_idx + 1)
463
+ query_len = query_end - query_start
464
+
465
+ seq_len = tl.load(seq_lens_ptr + req_idx)
466
+ prefix_len = seq_len - query_len
467
+
468
+ pos = prefix_len + token_idx - query_start
469
+ start_pos = tl.maximum(pos - window_size + 1, 0)
470
+ end_pos = pos + 1
471
+
472
+ swa_len = end_pos - start_pos
473
+ tl.store(swa_lens_ptr + token_idx, swa_len)
474
+
475
+ for i in range(0, window_size, TRITON_BLOCK_SIZE):
476
+ offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
477
+
478
+ pos_offset = start_pos + offset
479
+ block_indices = pos_offset // block_size
480
+ block_numbers = tl.load(
481
+ block_table_ptr + req_idx * block_table_stride + block_indices,
482
+ mask=pos_offset < end_pos,
483
+ )
484
+ block_offsets = pos_offset % block_size
485
+ slot_ids = block_numbers * block_size + block_offsets
486
+
487
+ slot_ids = tl.where(offset < swa_len, slot_ids, -1)
488
+ tl.store(
489
+ swa_indices_ptr + token_idx * swa_indices_stride + offset,
490
+ slot_ids,
491
+ mask=offset < window_size,
492
+ )
vllm/v1/attention/backends/utils.py
+1
-1
@@ -356,7 +356,7 @@ def make_local_attention_virtual_batches(
356
block_table_tensor=block_table_local,
357
slot_mapping=common_attn_metadata.slot_mapping,
358
causal=True,
359
- seq_lens_cpu_upper_bound=seq_lens_cpu,
359
+ seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound,
360
_seq_lens_cpu=seq_lens_cpu,
361
_num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local),
362
), make_block_table
vllm/v1/attention/ops/common.py
+34
-17
@@ -265,6 +265,7 @@ def _pack_seq_kernel(
265
D: tl.constexpr,
266
Lmax: tl.constexpr,
267
PAD_VALUE: tl.constexpr,
268
+ PAD_IS_UINT8: tl.constexpr,
269
BLOCK_T: tl.constexpr, # timesteps per program
270
BLOCK_D: tl.constexpr, # features per program
271
):
@@ -294,9 +295,15 @@ def _pack_seq_kernel(
295
# out_ptr: row-major [B, Lmax, D]
296
out_row_ptr = out_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :]
297
297
- # Initialize with PAD (cast will occur as needed based on out_ptr dtype)
298
+ # Initialize with PAD. PAD_IS_UINT8 selects the pad tensor's dtype so
299
+ # integer-typed outputs (e.g. MXFP4 packed nibbles, ue8m0 scale bytes)
300
+ # get an exact-byte pad rather than going through an fp32→uint8 cast
301
+ # that's implementation-defined outside of value 0.
302
d_mask = off_d[None, :] < D
299
- pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.float32)
303
+ if PAD_IS_UINT8:
304
+ pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.uint8)
305
+ else:
306
+ pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.float32)
307
tl.store(out_row_ptr, pad_vals, mask=t_mask[:, None] & d_mask)
308
309
# Load & write only where within seq_len
@@ -307,23 +314,36 @@ def _pack_seq_kernel(
314
def pack_seq_triton(
315
x: torch.Tensor,
316
lengths: torch.Tensor,
310
- pad_value: float = -float("inf"),
317
+ pad_value: float | int = -float("inf"),
318
block_t: int = 64,
319
block_d: int = 64,
320
) -> torch.Tensor:
314
- """
315
- Pack sequences of different lengths into a batched tensor.
321
+ """Pack sequences of different lengths into a batched tensor.
322
+
323
+ Supports float dtypes (any, via fp32 pad) and ``torch.uint8`` (exact-byte
324
+ pad — e.g. MXFP4 packed nibbles or ue8m0 scale bytes). For uint8 inputs
325
+ ``pad_value`` must be an integer in ``[0, 255]``.
326
327
Args:
318
- x: [N, ...] - input tensor where N is total number of tokens
319
- lengths: [B] - sequence lengths for each batch
320
- pad_value: value to use for padding
321
- block_t: block size for time dimension
322
- block_d: block size for feature dimension
328
+ x: [N, ...] — input tensor where N is total number of tokens.
329
+ lengths: [B] — sequence lengths for each batch.
330
+ pad_value: value to use for padding. Defaults to ``-inf`` which is
331
+ only sensible for float dtypes; pass ``0`` (or any byte) for
332
+ uint8 inputs.
333
+ block_t: block size for time dimension.
334
+ block_d: block size for feature dimension.
335
336
Returns:
325
- packed: [B, Lmax, ...] - packed tensor
337
+ packed: [B, Lmax, ...] — packed tensor.
338
"""
339
+ is_uint8 = x.dtype == torch.uint8
340
+ if is_uint8:
341
+ assert isinstance(pad_value, int) and 0 <= pad_value <= 255, (
342
+ f"uint8 pack requires an integer pad in [0, 255], got {pad_value!r}"
343
+ )
344
+ pad_constexpr: int | float = int(pad_value)
345
+ else:
346
+ pad_constexpr = float(pad_value)
347
348
# Handle multi-dimensional input by reshaping to (N, -1)
349
original_shape = x.shape
@@ -338,8 +358,6 @@ def pack_seq_triton(
358
B = lengths.numel()
359
Lmax = int(lengths.max().item())
360
341
- # Starts are computed inside the kernel from lengths
342
-
361
out = torch.empty((B, Lmax, D), device=x.device, dtype=x.dtype)
362
363
grid = (B, triton.cdiv(Lmax, block_t), triton.cdiv(D, block_d))
@@ -350,17 +368,16 @@ def pack_seq_triton(
368
N,
369
D,
370
Lmax,
353
- PAD_VALUE=float(pad_value),
371
+ PAD_VALUE=pad_constexpr,
372
+ PAD_IS_UINT8=is_uint8,
373
BLOCK_T=block_t,
374
BLOCK_D=block_d,
375
num_warps=4,
376
num_stages=2,
377
)
378
360
- # Reshape output back to original dimensions (except first dimension)
379
if len(original_shape) > 2:
362
- output_shape = (B, Lmax) + original_shape[1:]
363
- out = out.reshape(output_shape)
380
+ out = out.reshape((B, Lmax) + original_shape[1:])
381
382
return out
383
vllm/v1/attention/ops/deepseek_v4_ops/__init__.py
new
+23
@@ -0,0 +1,23 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from .cache_utils import (
5
+ combine_topk_swa_indices,
6
+ compute_global_topk_indices_and_lens,
7
+ dequantize_and_gather_k_cache,
8
+ quantize_and_insert_k_cache,
9
+)
10
+from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
11
+from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
12
+from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
13
+
14
+__all__ = [
15
+ "MXFP4_BLOCK_SIZE",
16
+ "combine_topk_swa_indices",
17
+ "compute_global_topk_indices_and_lens",
18
+ "dequantize_and_gather_k_cache",
19
+ "fused_indexer_q_rope_quant",
20
+ "fused_inv_rope_fp8_quant",
21
+ "fused_q_kv_rmsnorm",
22
+ "quantize_and_insert_k_cache",
23
+]
vllm/v1/attention/ops/deepseek_v4_ops/cache_utils.py
new
+563
@@ -0,0 +1,563 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index
5
+preparation.
6
+
7
+- quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into
8
+ the paged cache.
9
+- dequantize_and_gather_k_cache: gather and dequantize FP8 K from the paged
10
+ cache for sparse/SWA prefill.
11
+- compute_global_topk_indices_and_lens: map local topk indices to global KV
12
+ cache slots and count valid entries.
13
+- combine_topk_swa_indices: concatenate topk compressed indices with SWA
14
+ window indices for sparse prefill.
15
+"""
16
+
17
+import torch
18
+
19
+from vllm.triton_utils import tl, triton
20
+
21
+
22
+@triton.jit
23
+def quantize_and_insert_k_kernel(
24
+ # Input tensors
25
+ k_ptr, # [num_tokens, 512] bf16
26
+ slot_mapping_ptr, # [num_tokens] int64
27
+ # Output tensor
28
+ k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view)
29
+ # Dimensions
30
+ num_tokens,
31
+ input_dim: tl.constexpr, # 512
32
+ fp8_dim: tl.constexpr, # 448
33
+ bf16_dim: tl.constexpr, # 64
34
+ scale_dim: tl.constexpr, # 8
35
+ quant_block: tl.constexpr, # 64 (quantization block size)
36
+ cache_block_size: tl.constexpr, # 64 (paged cache block size)
37
+ token_data_size: tl.constexpr, # 576 bytes per token data
38
+ block_stride: tl.constexpr, # total bytes per block (padded)
39
+ fp8_max: tl.constexpr,
40
+ n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding)
41
+):
42
+ """
43
+ Quantize K tensor and insert into paged K cache.
44
+
45
+ K Cache block layout (block_size=64 tokens):
46
+ - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16
47
+ - [64*576, 64*576 + 64*8): Scales, each token has 8 uint8 scales
48
+ - [64*576 + 64*8, block_stride): Padding
49
+
50
+ One program per token.
51
+ """
52
+ pid = tl.program_id(0)
53
+
54
+ if pid >= num_tokens:
55
+ return
56
+
57
+ # Get slot mapping
58
+ slot_idx = tl.load(slot_mapping_ptr + pid)
59
+ if slot_idx == -1:
60
+ return
61
+
62
+ block_idx = slot_idx // cache_block_size
63
+ pos_in_block = slot_idx % cache_block_size
64
+
65
+ # Input pointer for this token
66
+ input_row_ptr = k_ptr + pid * input_dim
67
+
68
+ # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks
69
+ # (e.g. >= 57K at block_stride ~37K). Matches gather path below.
70
+ cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride
71
+
72
+ # Token data pointer: token data is stored contiguously at start of block
73
+ # Each token's data is at offset pos_in_block * token_data_size
74
+ token_data_ptr = cache_block_ptr + pos_in_block * token_data_size
75
+
76
+ # Scale pointer: scales are stored after ALL token data in the block
77
+ # Scale for this token is at offset (64 * 576) + pos_in_block * 8
78
+ token_scale_ptr = (
79
+ cache_block_ptr + cache_block_size * token_data_size + pos_in_block * scale_dim
80
+ )
81
+
82
+ # Token data layout: [0:448] fp8, [448:576] bf16
83
+ token_fp8_ptr = token_data_ptr
84
+ token_bf16_ptr = token_data_ptr + fp8_dim
85
+
86
+ # ========== Quantize and store FP8 portion (first 448 elements) ==========
87
+ # Using UE8M0 quantization strategy (scale is power of 2, stored as uint8 exponent)
88
+ for qblock_idx in tl.static_range(n_quant_blocks):
89
+ qblock_start = qblock_idx * quant_block
90
+
91
+ if qblock_start < fp8_dim:
92
+ offsets = qblock_start + tl.arange(0, quant_block)
93
+ mask = offsets < fp8_dim
94
+
95
+ # Load bf16 input
96
+ x = tl.load(input_row_ptr + offsets, mask=mask, other=0.0)
97
+
98
+ # Compute absmax scale (same as CUDA kernel)
99
+ abs_x = tl.abs(x)
100
+ block_max = tl.max(abs_x, axis=0)
101
+ block_max = tl.maximum(block_max, 1e-4) # Match CUDA: fmaxf(amax, 1e-4)
102
+
103
+ # UE8M0: Round scale UP to next power of 2
104
+ # scale = 2^ceil(log2(block_max / fp8_max))
105
+ raw_scale = block_max / fp8_max
106
+ log_scale = tl.log2(raw_scale)
107
+ exponent = tl.ceil(log_scale) # Round UP to next integer exponent
108
+ scale = tl.exp2(exponent) # scale = 2^exponent (power of 2)
109
+
110
+ # Quantize to fp8: fp8_value = bf16_value / scale
111
+ x_scaled = x / scale
112
+ x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max)
113
+
114
+ # Convert to fp8, then bitcast to uint8 for storage
115
+ x_fp8 = x_clamped.to(tl.float8e4nv)
116
+ x_uint8 = x_fp8.to(tl.uint8, bitcast=True)
117
+
118
+ # Store as uint8 (1 byte each)
119
+ tl.store(token_fp8_ptr + offsets, x_uint8, mask=mask)
120
+
121
+ # UE8M0 scale encoding: stored_value = exponent + 127 (bias)
122
+ # During dequant: scale = 2^(stored_value - 127)
123
+ encoded_scale = exponent + 127.0
124
+ encoded_scale = tl.maximum(tl.minimum(encoded_scale, 255.0), 0.0)
125
+ tl.store(token_scale_ptr + qblock_idx, encoded_scale.to(tl.uint8))
126
+
127
+ # Padding scale at index 7
128
+ tl.store(token_scale_ptr + 7, tl.zeros((), dtype=tl.uint8))
129
+
130
+ # ========== Store BF16 portion (last 64 elements, no quantization) ==========
131
+ bf16_input_offset = fp8_dim
132
+
133
+ # Process bf16 in chunks of 16
134
+ bf16_out_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
135
+ for i in tl.static_range(bf16_dim // 16):
136
+ chunk_offsets = i * 16 + tl.arange(0, 16)
137
+ bf16_vals = tl.load(input_row_ptr + bf16_input_offset + chunk_offsets)
138
+ tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)
139
+
140
+
141
+def quantize_and_insert_k_cache(
142
+ k: torch.Tensor, # [num_tokens, 512] bf16
143
+ k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8
144
+ slot_mapping: torch.Tensor, # [num_tokens] int64
145
+ block_size: int = 64,
146
+ is_ue8m0: bool = True,
147
+):
148
+ """
149
+ Quantize K tensor and insert into paged K cache.
150
+
151
+ K Cache block layout (block_size=64 tokens):
152
+ - First 64 * 576 = 36864 bytes: Token data
153
+ - Each token: 448 bytes (fp8) + 128 bytes (bf16)
154
+ - Next 64 * 8 = 512 bytes: Scales
155
+ - Each token: 8 bytes (uint8 scales, 7 real + 1 padding)
156
+ - Padded to multiple of 576
157
+ """
158
+ assert k.dim() == 2 and k.shape[1] == 512, (
159
+ f"K must be [num_tokens, 512], got {k.shape}"
160
+ )
161
+ assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}"
162
+ assert is_ue8m0, "Only support ue8m0 quantization."
163
+
164
+ # NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to
165
+ # padding. Always use slot_mapping.shape[0] as the token count.
166
+ num_tokens = slot_mapping.shape[0]
167
+ block_stride = k_cache.stride(0) # bytes per block
168
+
169
+ TOKEN_FP8_DIM = 448
170
+ TOKEN_BF16_DIM = 64
171
+ TOKEN_SCALE_DIM = 8
172
+ QUANT_BLOCK_SIZE = 64
173
+ FP8_MAX = 448.0
174
+ TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
175
+
176
+ grid = (num_tokens,)
177
+
178
+ quantize_and_insert_k_kernel[grid](
179
+ k,
180
+ slot_mapping,
181
+ k_cache,
182
+ num_tokens,
183
+ input_dim=512,
184
+ fp8_dim=TOKEN_FP8_DIM,
185
+ bf16_dim=TOKEN_BF16_DIM,
186
+ scale_dim=TOKEN_SCALE_DIM,
187
+ quant_block=QUANT_BLOCK_SIZE,
188
+ cache_block_size=block_size,
189
+ token_data_size=TOKEN_DATA_SIZE,
190
+ block_stride=block_stride,
191
+ fp8_max=FP8_MAX,
192
+ n_quant_blocks=8,
193
+ )
194
+
195
+
196
+@triton.jit
197
+def _dequantize_and_gather_k_kernel(
198
+ out_ptr,
199
+ out_stride0,
200
+ out_stride1,
201
+ k_cache_ptr,
202
+ seq_lens_ptr,
203
+ block_table_ptr,
204
+ offset,
205
+ gather_lens_ptr,
206
+ # Constants
207
+ max_blocks_per_seq: tl.constexpr,
208
+ fp8_dim: tl.constexpr, # 448
209
+ bf16_dim: tl.constexpr, # 64
210
+ scale_dim: tl.constexpr, # 8
211
+ quant_block: tl.constexpr, # 64 (quantization block size)
212
+ cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size)
213
+ token_data_size: tl.constexpr, # 576 bytes per token data
214
+ block_stride: tl.constexpr, # total bytes per block (padded) int32
215
+ output_dim: tl.constexpr, # 512
216
+ fp8_max: tl.constexpr,
217
+ n_quant_blocks: tl.constexpr, # 7 real blocks
218
+):
219
+ batch_idx = tl.program_id(0)
220
+ worker_id = tl.program_id(1)
221
+ num_workers = tl.num_programs(1)
222
+
223
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
224
+ if gather_lens_ptr is not None: # noqa: SIM108
225
+ gather_len = tl.load(gather_lens_ptr + batch_idx)
226
+ else:
227
+ # Gather all tokens
228
+ gather_len = seq_len
229
+ start_pos = seq_len - gather_len
230
+
231
+ for i in range(worker_id, gather_len, num_workers):
232
+ # Calculate the actual token index in the sequence
233
+ pos = start_pos + i
234
+
235
+ # Calculate which block and position within block
236
+ block_in_seq = pos // cache_block_size
237
+ pos_in_block = pos % cache_block_size
238
+
239
+ # Get physical block index from block table
240
+ block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
241
+ physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32
242
+
243
+ # int64: physical_block_idx * block_stride can exceed 2^31 with many
244
+ # KV-cache blocks (e.g. >= 57K at block_stride ~37K).
245
+ cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
246
+
247
+ # Token data pointer
248
+ token_data_ptr = cache_block_ptr + pos_in_block * token_data_size
249
+
250
+ # Scale pointer: after all token data
251
+ token_scale_ptr = (
252
+ cache_block_ptr
253
+ + cache_block_size * token_data_size
254
+ + pos_in_block * scale_dim
255
+ )
256
+
257
+ # Token data layout: [0:448] fp8, [448:576] bf16
258
+ token_fp8_ptr = token_data_ptr
259
+ token_bf16_ptr = token_data_ptr + fp8_dim
260
+
261
+ # Output pointer for this token (flattened)
262
+ output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
263
+
264
+ # ========== Dequantize FP8 portion using UE8M0 ==========
265
+ for qblock_idx in tl.static_range(n_quant_blocks):
266
+ qblock_start = qblock_idx * quant_block
267
+
268
+ if qblock_start < fp8_dim:
269
+ offsets = qblock_start + tl.arange(0, quant_block)
270
+ mask = offsets < fp8_dim
271
+
272
+ # Load quantized fp8 values (stored as uint8)
273
+ x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0)
274
+
275
+ # Bitcast uint8 back to fp8
276
+ x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
277
+
278
+ # Convert fp8 to float32 for computation
279
+ x_float = x_fp8.to(tl.float32)
280
+
281
+ # Load and decode UE8M0 scale
282
+ # UE8M0: scale = 2^(stored_value - 127)
283
+ encoded_scale = tl.load(token_scale_ptr + qblock_idx)
284
+ exponent = encoded_scale.to(tl.float32) - 127.0
285
+ scale = tl.exp2(exponent)
286
+
287
+ # Dequantize: bf16_value = fp8_value * scale
288
+ x_dequant = x_float * scale
289
+
290
+ # Store as bf16
291
+ tl.store(output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask)
292
+
293
+ # ========== Copy BF16 portion directly ==========
294
+ bf16_output_offset = fp8_dim # After 448 elements in output
295
+
296
+ # Read bf16 from cache
297
+ bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
298
+
299
+ # Process in chunks of 16
300
+ for j in tl.static_range(bf16_dim // 16):
301
+ chunk_offsets = j * 16 + tl.arange(0, 16)
302
+ bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets)
303
+ tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals)
304
+
305
+
306
+def dequantize_and_gather_k_cache(
307
+ # [num_reqs, max_num_tokens, head_size]
308
+ out: torch.Tensor,
309
+ # [num_blocks, block_size, head_bytes]
310
+ k_cache: torch.Tensor,
311
+ # [num_reqs]
312
+ seq_lens: torch.Tensor,
313
+ # [num_reqs]
314
+ gather_lens: torch.Tensor | None,
315
+ # [num_reqs, max_blocks_per_seq]
316
+ block_table: torch.Tensor,
317
+ block_size: int,
318
+ offset: int,
319
+) -> None:
320
+ TOKEN_FP8_DIM = 448
321
+ TOKEN_BF16_DIM = 64
322
+ TOKEN_SCALE_DIM = 8
323
+ QUANT_BLOCK_SIZE = 64
324
+ FP8_MAX = 448.0
325
+ TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
326
+
327
+ num_reqs = seq_lens.shape[0]
328
+ NUM_WORKERS = 128
329
+ _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)](
330
+ out,
331
+ out.stride(0),
332
+ out.stride(1),
333
+ k_cache,
334
+ seq_lens,
335
+ block_table,
336
+ offset,
337
+ gather_lens,
338
+ max_blocks_per_seq=block_table.shape[-1],
339
+ fp8_dim=TOKEN_FP8_DIM,
340
+ bf16_dim=TOKEN_BF16_DIM,
341
+ scale_dim=TOKEN_SCALE_DIM,
342
+ quant_block=QUANT_BLOCK_SIZE,
343
+ cache_block_size=block_size,
344
+ token_data_size=TOKEN_DATA_SIZE,
345
+ block_stride=k_cache.stride(0),
346
+ output_dim=512,
347
+ fp8_max=FP8_MAX,
348
+ n_quant_blocks=7,
349
+ )
350
+
351
+
352
+def compute_global_topk_indices_and_lens(
353
+ topk_indices: torch.Tensor,
354
+ token_to_req_indices: torch.Tensor,
355
+ block_table: torch.Tensor,
356
+ block_size: int,
357
+ is_valid_token: torch.Tensor,
358
+) -> tuple[torch.Tensor, torch.Tensor]:
359
+ """Map local topk indices to global KV cache slots and count valid entries.
360
+
361
+ Fuses three operations into a single kernel:
362
+ 1. Block-table lookup (local index → global slot id)
363
+ 2. Valid-entry counting (topk_lens per token)
364
+ 3. Masking padding tokens to length 0
365
+ """
366
+ num_tokens = topk_indices.shape[0]
367
+ global_topk_indices = torch.empty_like(topk_indices)
368
+ topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
369
+ _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
370
+ global_topk_indices,
371
+ global_topk_indices.stride(0),
372
+ topk_lens,
373
+ topk_indices,
374
+ topk_indices.stride(0),
375
+ topk_indices.shape[-1],
376
+ token_to_req_indices,
377
+ block_table,
378
+ block_table.stride(0),
379
+ block_size,
380
+ is_valid_token,
381
+ TRITON_BLOCK_SIZE=1024,
382
+ )
383
+ return global_topk_indices, topk_lens
384
+
385
+
386
+@triton.jit
387
+def _compute_global_topk_indices_and_lens_kernel(
388
+ global_topk_indices_ptr,
389
+ global_topk_indices_stride,
390
+ topk_lens_ptr,
391
+ topk_indices_ptr,
392
+ topk_indices_stride,
393
+ topk,
394
+ token_to_req_indices_ptr,
395
+ block_table_ptr,
396
+ block_table_stride,
397
+ block_size,
398
+ is_valid_token_ptr,
399
+ TRITON_BLOCK_SIZE: tl.constexpr,
400
+):
401
+ token_idx = tl.program_id(0)
402
+ is_valid_token = tl.load(is_valid_token_ptr + token_idx)
403
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
404
+
405
+ count = tl.zeros((), dtype=tl.int32)
406
+ for i in range(0, topk, TRITON_BLOCK_SIZE):
407
+ offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
408
+ mask = offset < topk
409
+
410
+ local_idx = tl.load(
411
+ topk_indices_ptr + token_idx * topk_indices_stride + offset,
412
+ mask=mask,
413
+ other=-1,
414
+ )
415
+ is_valid = local_idx >= 0
416
+
417
+ block_indices = local_idx // block_size
418
+ block_numbers = tl.load(
419
+ block_table_ptr + req_idx * block_table_stride + block_indices,
420
+ mask=mask & is_valid,
421
+ )
422
+ block_offsets = local_idx % block_size
423
+
424
+ slot_ids = block_numbers * block_size + block_offsets
425
+ slot_ids = tl.where(is_valid, slot_ids, -1)
426
+ tl.store(
427
+ global_topk_indices_ptr + token_idx * global_topk_indices_stride + offset,
428
+ slot_ids,
429
+ mask=mask,
430
+ )
431
+ count += tl.sum(is_valid.to(tl.int32), axis=0)
432
+
433
+ # Zero out length for padding tokens.
434
+ tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0))
435
+
436
+
437
+# FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see
438
+# flashmla/csrc/sm100/prefill/sparse/fwd/head{64,128}/phase1.cuh). B_TOPK is
439
+# 64 for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both.
440
+# The extra slots stay as -1 sentinels and `combined_lens` caps the valid
441
+# range via `topk_length`, so padding is a no-op at kernel level.
442
+_SPARSE_PREFILL_TOPK_ALIGNMENT = 128
443
+
444
+
445
+def combine_topk_swa_indices(
446
+ topk_indices: torch.Tensor,
447
+ query_start_loc: torch.Tensor,
448
+ seq_lens: torch.Tensor,
449
+ gather_lens: torch.Tensor,
450
+ window_size: int,
451
+ compress_ratio: int,
452
+ topk: int,
453
+ M: int,
454
+ N: int,
455
+) -> tuple[torch.Tensor, torch.Tensor]:
456
+ num_tokens = topk_indices.shape[0]
457
+ num_reqs = seq_lens.shape[0]
458
+ combined_topk = (
459
+ (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
460
+ // _SPARSE_PREFILL_TOPK_ALIGNMENT
461
+ * _SPARSE_PREFILL_TOPK_ALIGNMENT
462
+ )
463
+ combined_indices = torch.full(
464
+ (num_tokens, combined_topk),
465
+ fill_value=-1,
466
+ dtype=torch.int32,
467
+ device=topk_indices.device,
468
+ )
469
+ combined_lens = torch.empty(
470
+ num_tokens, dtype=torch.int32, device=topk_indices.device
471
+ )
472
+
473
+ NUM_WORKERS = 128
474
+ _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)](
475
+ combined_indices,
476
+ combined_indices.stride(0),
477
+ combined_lens,
478
+ topk_indices,
479
+ topk_indices.stride(0),
480
+ query_start_loc,
481
+ seq_lens,
482
+ gather_lens,
483
+ M,
484
+ N,
485
+ TOP_K=topk,
486
+ COMPRESS_RATIO=compress_ratio,
487
+ WINDOW_SIZE=window_size,
488
+ PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]),
489
+ )
490
+ return combined_indices, combined_lens
491
+
492
+
493
+@triton.jit
494
+def _combine_topk_swa_indices_kernel(
495
+ combined_indices_ptr,
496
+ combined_indices_stride,
497
+ combined_lens_ptr,
498
+ topk_indices_ptr,
499
+ topk_indices_stride,
500
+ query_start_loc_ptr,
501
+ seq_lens_ptr,
502
+ gather_lens_ptr,
503
+ M,
504
+ N,
505
+ TOP_K: tl.constexpr,
506
+ COMPRESS_RATIO: tl.constexpr,
507
+ WINDOW_SIZE: tl.constexpr,
508
+ PADDED_TOP_K: tl.constexpr,
509
+):
510
+ batch_idx = tl.program_id(0)
511
+ worker_id = tl.program_id(1)
512
+ num_workers = tl.num_programs(1)
513
+
514
+ # query_start_loc is a global tensor; rebase to chunk-local offsets
515
+ # by subtracting the chunk's starting value.
516
+ base = tl.load(query_start_loc_ptr)
517
+ query_start = tl.load(query_start_loc_ptr + batch_idx) - base
518
+ query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
519
+ query_len = query_end - query_start
520
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
521
+ gather_len = tl.load(gather_lens_ptr + batch_idx)
522
+ start_pos = seq_len - query_len
523
+ # The SWA portion of the gathered buffer starts from position
524
+ # (seq_len - gather_len), not position 0. We need this offset
525
+ # to correctly index into the gathered buffer.
526
+ gather_start = seq_len - gather_len
527
+
528
+ for token_idx in range(query_start + worker_id, query_end, num_workers):
529
+ # topk_len is fully determined by the query token's absolute position:
530
+ # both the C4A indexer and the C128A metadata builder emit
531
+ # min((pos + 1) // compress_ratio, topk_tokens) valid entries.
532
+ # Caller passes TOP_K=0 for SWA-only layers to zero this out.
533
+ token_idx_in_query = token_idx - query_start
534
+ pos = start_pos + token_idx_in_query
535
+ topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
536
+ swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
537
+
538
+ offset = tl.arange(0, PADDED_TOP_K)
539
+ mask = offset < topk_len
540
+ topk_indices = tl.load(
541
+ topk_indices_ptr + token_idx * topk_indices_stride + offset,
542
+ mask=mask,
543
+ )
544
+ tl.store(
545
+ combined_indices_ptr + token_idx * combined_indices_stride + offset,
546
+ topk_indices + M * batch_idx,
547
+ mask=mask,
548
+ )
549
+ offset = tl.arange(0, WINDOW_SIZE)
550
+ # Index into gathered buffer: N + (position - gather_start)
551
+ # For positions [pos - swa_len + 1, pos], the buffer indices are:
552
+ # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start]
553
+ tl.store(
554
+ combined_indices_ptr
555
+ + token_idx * combined_indices_stride
556
+ + topk_len
557
+ + offset,
558
+ M * batch_idx + N + offset + pos - swa_len + 1 - gather_start,
559
+ mask=offset < swa_len,
560
+ )
561
+
562
+ combined_len = topk_len + swa_len
563
+ tl.store(combined_lens_ptr + token_idx, combined_len)
vllm/v1/attention/ops/deepseek_v4_ops/fused_compress_quant_cache.py
new
+584
@@ -0,0 +1,584 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Fused compressor + FP8/MXFP4 UE8M0 quantization + KV cache insert kernels.
5
+
6
+Three specialized kernels:
7
+ - _fused_kv_compress_norm_rope_insert_sparse_attn:
8
+ head=512, nope=448 FP8 + rope=64 bf16
9
+ - _fused_kv_compress_norm_rope_insert_indexer_attn:
10
+ head=128, all FP8, 1 block/token
11
+ - _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn:
12
+ head=128, MXFP4 (block=32), 4 ue8m0 bytes
13
+
14
+RoPE is register-based via tl.reshape -> tl.split -> tl.interleave (or the
15
+even/odd halves are consumed directly for MXFP4, no interleave needed).
16
+FP8 UE8M0 quant uses tl.reshape to tile [N_QUANT_BLOCKS, QUANT_BLOCK] for
17
+per-block absmax entirely in registers. MXFP4 does the same tiling on the
18
+even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
19
+and N_QUANT_BLOCKS ue8m0 bytes.
20
+"""
21
+
22
+from vllm.triton_utils import tl, triton
23
+
24
+from .fused_indexer_q import _e2m1_nibble
25
+
26
+
27
+# =============================================================================
28
+# DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16)
29
+# =============================================================================
30
+@triton.jit
31
+def _fused_kv_compress_norm_rope_insert_sparse_attn(
32
+ # ── state cache (compressor internal state) ──
33
+ state_cache_ptr,
34
+ state_cache_stride0,
35
+ state_cache_stride1,
36
+ # ── metadata ──
37
+ token_to_req_indices_ptr,
38
+ positions_ptr,
39
+ slot_mapping_ptr,
40
+ block_table_ptr,
41
+ block_table_stride,
42
+ block_size,
43
+ # ── RMSNorm ──
44
+ rms_norm_weight_ptr,
45
+ rms_norm_eps,
46
+ # ── RoPE ──
47
+ cos_sin_cache_ptr,
48
+ cos_sin_stride,
49
+ # ── KV cache output ──
50
+ k_cache_ptr,
51
+ kv_slot_mapping_ptr,
52
+ kv_cache_block_size,
53
+ # ── constexprs ──
54
+ HEAD_SIZE: tl.constexpr,
55
+ TRITON_BLOCK_SIZE: tl.constexpr,
56
+ STATE_WIDTH: tl.constexpr,
57
+ COMPRESS_RATIO: tl.constexpr,
58
+ OVERLAP: tl.constexpr,
59
+ ROPE_HEAD_DIM: tl.constexpr,
60
+ FP8_MAX: tl.constexpr, # 448.0
61
+ QUANT_BLOCK: tl.constexpr, # 64 for DeepseekV4
62
+ TOKEN_STRIDE: tl.constexpr, # 576 for DeepseekV4
63
+ SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad)
64
+ KV_BLOCK_STRIDE: tl.constexpr,
65
+):
66
+ """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope).
67
+
68
+ One program per token; early-exits for non-boundary positions.
69
+
70
+ Cache block layout (``block_size`` tokens):
71
+ [0, bs*576): token data (448 fp8 + 128 bf16 each)
72
+ [bs*576, +bs*8): uint8 UE8M0 scales (7 real + 1 pad each)
73
+ """
74
+ token_idx = tl.program_id(0)
75
+
76
+ slot_id = tl.load(slot_mapping_ptr + token_idx)
77
+ if slot_id < 0:
78
+ return
79
+
80
+ position = tl.load(positions_ptr + token_idx)
81
+ if (position + 1) % COMPRESS_RATIO != 0:
82
+ return
83
+
84
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
85
+
86
+ # ── Gather state cache entries ────────────────────────────────────
87
+ start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1
88
+ tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO)
89
+ pos = start + tokens
90
+ mask_pos = pos >= 0
91
+
92
+ block_indices = pos // block_size
93
+ block_numbers = tl.load(
94
+ block_table_ptr + req_idx * block_table_stride + block_indices,
95
+ mask=mask_pos,
96
+ other=0,
97
+ )
98
+ block_offsets = pos % block_size
99
+ head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE
100
+
101
+ block = tl.arange(0, TRITON_BLOCK_SIZE)
102
+ mask = block < HEAD_SIZE
103
+ block_numbers_i64 = block_numbers.to(tl.int64)
104
+
105
+ # Precomputed row base shared by score and kv loads
106
+ row_base = (
107
+ state_cache_ptr
108
+ + block_numbers_i64 * state_cache_stride0
109
+ + block_offsets * state_cache_stride1
110
+ + head_offset
111
+ )
112
+
113
+ combined_mask = mask_pos[:, None] & mask[None, :]
114
+
115
+ # ── Softmax + weighted sum ───────────────────────────────────────
116
+ score = tl.load(
117
+ row_base[:, None] + STATE_WIDTH + block[None, :],
118
+ mask=combined_mask,
119
+ other=float("-inf"),
120
+ )
121
+ score = tl.softmax(score, dim=0)
122
+
123
+ kv = tl.load(
124
+ row_base[:, None] + block[None, :],
125
+ mask=combined_mask,
126
+ other=0.0,
127
+ )
128
+
129
+ compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32
130
+
131
+ # ── RMSNorm (fp32 throughout) ──────────────────────────────────────
132
+ rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0)
133
+ variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE
134
+ rrms = tl.rsqrt(variance + rms_norm_eps)
135
+ normed = compressed_kv * rrms * rms_w
136
+
137
+ # ── KV cache pointers ────────────────────────────────────────────
138
+ kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
139
+ if kv_slot_idx < 0:
140
+ return
141
+ kv_block_idx = kv_slot_idx // kv_cache_block_size
142
+ kv_pos_in_block = kv_slot_idx % kv_cache_block_size
143
+
144
+ cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
145
+ fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
146
+ scale_ptr = (
147
+ cache_block_ptr
148
+ + kv_cache_block_size * TOKEN_STRIDE
149
+ + kv_pos_in_block * SCALE_DIM
150
+ )
151
+
152
+ NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM # 448
153
+ HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 # 32
154
+
155
+ # FP8 UE8M0 quant: cast fp32 → bf16 → fp32 before quant to match reference.
156
+ N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK
157
+ N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK # 7
158
+ INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX
159
+
160
+ quant_input = normed.to(tl.bfloat16).to(tl.float32)
161
+ quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK))
162
+ abs_2d = tl.abs(quant_2d)
163
+ block_absmax = tl.max(abs_2d, axis=1) # [N_QUANT_BLOCKS] fp32
164
+ block_absmax = tl.maximum(block_absmax, 1e-4)
165
+
166
+ raw_scales = block_absmax * INV_FP8_MAX
167
+ exponents = tl.ceil(tl.log2(raw_scales))
168
+ inv_scales = tl.exp2(-exponents)
169
+ inv_scales_col = tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1))
170
+ x_scaled = quant_2d * inv_scales_col
171
+ x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX)
172
+ x_fp8 = x_clamped.to(tl.float8e4nv)
173
+ x_uint8 = x_fp8.to(tl.uint8, bitcast=True)
174
+ x_uint8_flat = tl.reshape(x_uint8, (TRITON_BLOCK_SIZE,))
175
+
176
+ nope_mask = block < NOPE_HEAD_DIM
177
+ tl.store(fp8_ptr + block, x_uint8_flat, mask=nope_mask)
178
+
179
+ scale_idx = tl.arange(0, N_QUANT_BLOCKS)
180
+ encoded = exponents + 127.0
181
+ encoded = tl.maximum(tl.minimum(encoded, 255.0), 0.0)
182
+ tl.store(
183
+ scale_ptr + scale_idx,
184
+ encoded.to(tl.uint8),
185
+ mask=scale_idx < N_NOPE_BLOCKS,
186
+ )
187
+ tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8))
188
+
189
+ # Register-based GPT-J RoPE in fp32.
190
+ NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2
191
+ NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
192
+
193
+ pair_2d = tl.reshape(normed, (NUM_PAIRS, 2))
194
+ even, odd = tl.split(pair_2d) # each [NUM_PAIRS] fp32
195
+
196
+ pair_idx = tl.arange(0, NUM_PAIRS)
197
+ rope_pair_local = pair_idx - NOPE_PAIRS
198
+ is_rope_pair = rope_pair_local >= 0
199
+ cs_idx = tl.maximum(rope_pair_local, 0)
200
+
201
+ compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
202
+ cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
203
+ cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
204
+ sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
205
+
206
+ new_even = even * cos_v - odd * sin_v
207
+ new_odd = odd * cos_v + even * sin_v
208
+ result = tl.interleave(new_even, new_odd) # [TRITON_BLOCK_SIZE] fp32
209
+
210
+ # Store rotated rope portion as bf16 into the cache's bf16 area.
211
+ bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16))
212
+ rope_local = block - NOPE_HEAD_DIM
213
+ is_rope = (block >= NOPE_HEAD_DIM) & mask
214
+ tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope)
215
+
216
+
217
+# =============================================================================
218
+# Indexer path (head=128, all FP8, single quant block)
219
+# =============================================================================
220
+@triton.jit
221
+def _fused_kv_compress_norm_rope_insert_indexer_attn(
222
+ # ── state cache (compressor internal state) ──
223
+ state_cache_ptr,
224
+ state_cache_stride0,
225
+ state_cache_stride1,
226
+ # ── metadata ──
227
+ token_to_req_indices_ptr,
228
+ positions_ptr,
229
+ slot_mapping_ptr,
230
+ block_table_ptr,
231
+ block_table_stride,
232
+ block_size,
233
+ # ── RMSNorm ──
234
+ rms_norm_weight_ptr,
235
+ rms_norm_eps,
236
+ # ── RoPE ──
237
+ cos_sin_cache_ptr,
238
+ cos_sin_stride,
239
+ # ── KV cache output ──
240
+ k_cache_ptr,
241
+ kv_slot_mapping_ptr,
242
+ kv_cache_block_size,
243
+ # ── constexprs ──
244
+ HEAD_SIZE: tl.constexpr,
245
+ TRITON_BLOCK_SIZE: tl.constexpr,
246
+ STATE_WIDTH: tl.constexpr,
247
+ COMPRESS_RATIO: tl.constexpr,
248
+ OVERLAP: tl.constexpr,
249
+ ROPE_HEAD_DIM: tl.constexpr,
250
+ FP8_MAX: tl.constexpr, # 448.0
251
+ QUANT_BLOCK: tl.constexpr, # 128 for indexer
252
+ TOKEN_STRIDE: tl.constexpr, # 128 for indexer
253
+ SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32)
254
+ KV_BLOCK_STRIDE: tl.constexpr,
255
+):
256
+ """Fused compress → RMSNorm → RoPE → FP8 quant → store.
257
+
258
+ One program per token; early-exits for non-boundary positions.
259
+
260
+ Cache block layout:
261
+ [0, bs*128): FP8 data (128 bytes/token)
262
+ [bs*128, +bs*4): float32 scales (4 bytes/token)
263
+
264
+ For head_dim=128 we have exactly one quant block, so we skip the
265
+ [N_QUANT_BLOCKS, QUANT_BLOCK] reshape entirely and use a flat
266
+ ``tl.max`` reduction.
267
+ """
268
+ token_idx = tl.program_id(0)
269
+
270
+ slot_id = tl.load(slot_mapping_ptr + token_idx)
271
+ if slot_id < 0:
272
+ return
273
+
274
+ position = tl.load(positions_ptr + token_idx)
275
+ if (position + 1) % COMPRESS_RATIO != 0:
276
+ return
277
+
278
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
279
+
280
+ # ── Gather state cache entries ────────────────────────────────────
281
+ start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1
282
+ tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO)
283
+ pos = start + tokens
284
+ mask_pos = pos >= 0
285
+
286
+ block_indices = pos // block_size
287
+ block_numbers = tl.load(
288
+ block_table_ptr + req_idx * block_table_stride + block_indices,
289
+ mask=mask_pos,
290
+ other=0,
291
+ )
292
+ block_offsets = pos % block_size
293
+ head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE
294
+
295
+ block = tl.arange(0, TRITON_BLOCK_SIZE)
296
+ mask = block < HEAD_SIZE
297
+ block_numbers_i64 = block_numbers.to(tl.int64)
298
+
299
+ row_base = (
300
+ state_cache_ptr
301
+ + block_numbers_i64 * state_cache_stride0
302
+ + block_offsets * state_cache_stride1
303
+ + head_offset
304
+ )
305
+
306
+ combined_mask = mask_pos[:, None] & mask[None, :]
307
+
308
+ score = tl.load(
309
+ row_base[:, None] + STATE_WIDTH + block[None, :],
310
+ mask=combined_mask,
311
+ other=float("-inf"),
312
+ )
313
+ score = tl.softmax(score, dim=0)
314
+
315
+ kv = tl.load(
316
+ row_base[:, None] + block[None, :],
317
+ mask=combined_mask,
318
+ other=0.0,
319
+ )
320
+
321
+ compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32
322
+
323
+ # ── RMSNorm (fp32 throughout) ──────────────────────────────────────
324
+ rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0)
325
+ variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE
326
+ rrms = tl.rsqrt(variance + rms_norm_eps)
327
+ normed = compressed_kv * rrms * rms_w
328
+
329
+ # ── KV cache pointers ────────────────────────────────────────────
330
+ kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
331
+ if kv_slot_idx < 0:
332
+ return
333
+ kv_block_idx = kv_slot_idx // kv_cache_block_size
334
+ kv_pos_in_block = kv_slot_idx % kv_cache_block_size
335
+
336
+ cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
337
+ fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
338
+ scale_ptr = (
339
+ cache_block_ptr
340
+ + kv_cache_block_size * TOKEN_STRIDE
341
+ + kv_pos_in_block * SCALE_DIM
342
+ )
343
+
344
+ NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM
345
+ HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2
346
+
347
+ # ── Register-based GPT-J forward RoPE in fp32 ─────────────────────
348
+ NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2
349
+ NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
350
+
351
+ normed_2d = tl.reshape(normed, (NUM_PAIRS, 2))
352
+ even, odd = tl.split(normed_2d) # each [NUM_PAIRS] fp32
353
+
354
+ pair_idx = tl.arange(0, NUM_PAIRS)
355
+ rope_pair_local = pair_idx - NOPE_PAIRS
356
+ is_rope_pair = rope_pair_local >= 0
357
+ cs_idx = tl.maximum(rope_pair_local, 0)
358
+
359
+ compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
360
+ cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
361
+ cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
362
+ sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
363
+
364
+ new_even = even * cos_v - odd * sin_v
365
+ new_odd = odd * cos_v + even * sin_v
366
+ result = tl.interleave(new_even, new_odd) # fp32
367
+
368
+ # ── FP8 UE8M0 quant: single block, flat reduction ────────────────
369
+ tl.static_assert(
370
+ TRITON_BLOCK_SIZE == QUANT_BLOCK,
371
+ "Indexer expects one quant block (QUANT_BLOCK == TRITON_BLOCK_SIZE)",
372
+ )
373
+ INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX
374
+
375
+ result_bf16 = result.to(tl.bfloat16).to(tl.float32)
376
+ absmax = tl.max(tl.abs(result_bf16), axis=0) # scalar
377
+ absmax = tl.maximum(absmax, 1e-4)
378
+ raw_scale = absmax * INV_FP8_MAX
379
+ exponent = tl.ceil(tl.log2(raw_scale))
380
+ inv_scale = tl.exp2(-exponent)
381
+
382
+ x_scaled = result_bf16 * inv_scale
383
+ x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX)
384
+ x_fp8 = x_clamped.to(tl.float8e4nv)
385
+ x_uint8 = x_fp8.to(tl.uint8, bitcast=True)
386
+
387
+ tl.store(fp8_ptr + block, x_uint8, mask=mask)
388
+
389
+ # Single float32 scale
390
+ scale_val = tl.exp2(exponent)
391
+ tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val)
392
+
393
+
394
+# =============================================================================
395
+# Indexer path (head=128, MXFP4: 2 nibbles/byte + ue8m0 per 32-elem block)
396
+# =============================================================================
397
+@triton.jit
398
+def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn(
399
+ # ── state cache (compressor internal state) ──
400
+ state_cache_ptr,
401
+ state_cache_stride0,
402
+ state_cache_stride1,
403
+ # ── metadata ──
404
+ token_to_req_indices_ptr,
405
+ positions_ptr,
406
+ slot_mapping_ptr,
407
+ block_table_ptr,
408
+ block_table_stride,
409
+ block_size,
410
+ # ── RMSNorm ──
411
+ rms_norm_weight_ptr,
412
+ rms_norm_eps,
413
+ # ── RoPE ──
414
+ cos_sin_cache_ptr,
415
+ cos_sin_stride,
416
+ # ── KV cache output ──
417
+ k_cache_ptr,
418
+ kv_slot_mapping_ptr,
419
+ kv_cache_block_size,
420
+ # ── constexprs ──
421
+ HEAD_SIZE: tl.constexpr,
422
+ TRITON_BLOCK_SIZE: tl.constexpr,
423
+ STATE_WIDTH: tl.constexpr,
424
+ COMPRESS_RATIO: tl.constexpr,
425
+ OVERLAP: tl.constexpr,
426
+ ROPE_HEAD_DIM: tl.constexpr,
427
+ FP8_MAX: tl.constexpr, # unused for MXFP4 (kept for signature parity)
428
+ QUANT_BLOCK: tl.constexpr, # 32 for MXFP4
429
+ TOKEN_STRIDE: tl.constexpr, # HEAD_SIZE // 2 = 64 packed bytes/token
430
+ SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token
431
+ KV_BLOCK_STRIDE: tl.constexpr,
432
+):
433
+ """Fused compress → RMSNorm → RoPE → MXFP4 quant → store.
434
+
435
+ One program per token; early-exits for non-boundary positions.
436
+
437
+ Cache block layout (``block_size`` tokens per cache block):
438
+ [0, bs*TOKEN_STRIDE): packed MXFP4 nibbles (2 values/byte)
439
+ [bs*TOKEN_STRIDE, +bs*SCALE_DIM): ue8m0 scale bytes (one per 32-elem block)
440
+
441
+ MXFP4 format:
442
+ - E2M1 4-bit values packed two per byte (low nibble first, then high).
443
+ - Per-32-element block scale = 2^ceil(log2(amax / 6.0)), stored ue8m0
444
+ (byte = exponent + 127).
445
+ - Max representable magnitude = 6.0.
446
+ """
447
+ token_idx = tl.program_id(0)
448
+
449
+ slot_id = tl.load(slot_mapping_ptr + token_idx)
450
+ if slot_id < 0:
451
+ return
452
+
453
+ position = tl.load(positions_ptr + token_idx)
454
+ if (position + 1) % COMPRESS_RATIO != 0:
455
+ return
456
+
457
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
458
+
459
+ # ── Gather state cache entries ────────────────────────────────────
460
+ start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1
461
+ tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO)
462
+ pos = start + tokens
463
+ mask_pos = pos >= 0
464
+
465
+ block_indices = pos // block_size
466
+ block_numbers = tl.load(
467
+ block_table_ptr + req_idx * block_table_stride + block_indices,
468
+ mask=mask_pos,
469
+ other=0,
470
+ )
471
+ block_offsets = pos % block_size
472
+ head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE
473
+
474
+ block = tl.arange(0, TRITON_BLOCK_SIZE)
475
+ mask = block < HEAD_SIZE
476
+ block_numbers_i64 = block_numbers.to(tl.int64)
477
+
478
+ row_base = (
479
+ state_cache_ptr
480
+ + block_numbers_i64 * state_cache_stride0
481
+ + block_offsets * state_cache_stride1
482
+ + head_offset
483
+ )
484
+
485
+ combined_mask = mask_pos[:, None] & mask[None, :]
486
+
487
+ score = tl.load(
488
+ row_base[:, None] + STATE_WIDTH + block[None, :],
489
+ mask=combined_mask,
490
+ other=float("-inf"),
491
+ )
492
+ score = tl.softmax(score, dim=0)
493
+
494
+ kv = tl.load(
495
+ row_base[:, None] + block[None, :],
496
+ mask=combined_mask,
497
+ other=0.0,
498
+ )
499
+
500
+ compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32
501
+
502
+ # ── RMSNorm (fp32 throughout) ──────────────────────────────────────
503
+ rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0)
504
+ variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE
505
+ rrms = tl.rsqrt(variance + rms_norm_eps)
506
+ normed = compressed_kv * rrms * rms_w
507
+
508
+ # ── KV cache pointers (segregated: values first, then scales) ────
509
+ kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
510
+ if kv_slot_idx < 0:
511
+ return
512
+ kv_block_idx = kv_slot_idx // kv_cache_block_size
513
+ kv_pos_in_block = kv_slot_idx % kv_cache_block_size
514
+
515
+ cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
516
+ val_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
517
+ scale_ptr = (
518
+ cache_block_ptr
519
+ + kv_cache_block_size * TOKEN_STRIDE
520
+ + kv_pos_in_block * SCALE_DIM
521
+ )
522
+
523
+ NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM
524
+ HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2
525
+
526
+ # ── Register-based GPT-J forward RoPE in fp32 ─────────────────────
527
+ # We keep the even/odd halves (no tl.interleave afterwards) because the
528
+ # MXFP4 per-block absmax / pack naturally operates on (even, odd) pairs.
529
+ NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2
530
+ NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
531
+
532
+ normed_2d = tl.reshape(normed, (NUM_PAIRS, 2))
533
+ even, odd = tl.split(normed_2d) # each [NUM_PAIRS] fp32
534
+
535
+ pair_idx = tl.arange(0, NUM_PAIRS)
536
+ rope_pair_local = pair_idx - NOPE_PAIRS
537
+ is_rope_pair = rope_pair_local >= 0
538
+ cs_idx = tl.maximum(rope_pair_local, 0)
539
+
540
+ compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
541
+ cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
542
+ cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
543
+ sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
544
+
545
+ new_even = even * cos_v - odd * sin_v
546
+ new_odd = odd * cos_v + even * sin_v
547
+
548
+ # bf16 roundtrip for parity with reference / Q-side kernel numerics.
549
+ new_even = new_even.to(tl.bfloat16).to(tl.float32)
550
+ new_odd = new_odd.to(tl.bfloat16).to(tl.float32)
551
+
552
+ # ── MXFP4 quant: tile even/odd halves into (N_BLOCKS, HALF_BLOCK) ──
553
+ # Each MXFP4 block of QUANT_BLOCK elements = HALF_BLOCK consecutive pairs,
554
+ # so (N_BLOCKS, HALF_BLOCK) rows of even/odd each land exactly one block.
555
+ N_QUANT_BLOCKS: tl.constexpr = HEAD_SIZE // QUANT_BLOCK
556
+ HALF_BLOCK: tl.constexpr = QUANT_BLOCK // 2
557
+ tl.static_assert(TRITON_BLOCK_SIZE == HEAD_SIZE)
558
+ tl.static_assert(HEAD_SIZE % QUANT_BLOCK == 0)
559
+ tl.static_assert(TOKEN_STRIDE == HEAD_SIZE // 2)
560
+ tl.static_assert(SCALE_DIM == N_QUANT_BLOCKS)
561
+
562
+ even_2d = tl.reshape(new_even, (N_QUANT_BLOCKS, HALF_BLOCK))
563
+ odd_2d = tl.reshape(new_odd, (N_QUANT_BLOCKS, HALF_BLOCK))
564
+
565
+ amax = tl.maximum(
566
+ tl.max(tl.abs(even_2d), axis=1),
567
+ tl.max(tl.abs(odd_2d), axis=1),
568
+ )
569
+ amax = tl.maximum(amax, 1e-4)
570
+
571
+ # ue8m0 block scale: 2^ceil(log2(amax / 6.0)), stored as (exp + 127) byte.
572
+ log2_ratio = tl.ceil(tl.log2(amax / 6.0))
573
+ log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
574
+ inv_scale = tl.exp2(-log2_ratio)
575
+ ue8m0 = (log2_ratio + 127.0).to(tl.uint8) # [N_QUANT_BLOCKS]
576
+
577
+ inv_scale_col = tl.reshape(inv_scale, (N_QUANT_BLOCKS, 1))
578
+ lo_nib = _e2m1_nibble(even_2d * inv_scale_col) # (N_BLOCKS, HALF_BLOCK) uint8
579
+ hi_nib = _e2m1_nibble(odd_2d * inv_scale_col)
580
+ packed = lo_nib | (hi_nib << 4)
581
+ packed_flat = tl.reshape(packed, (TOKEN_STRIDE,))
582
+
583
+ tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat)
584
+ tl.store(scale_ptr + tl.arange(0, SCALE_DIM), ue8m0)
vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q.py
new
+415
@@ -0,0 +1,415 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.triton_utils import tl, triton
6
+
7
+# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale.
8
+MXFP4_BLOCK_SIZE = 32
9
+
10
+
11
+@triton.jit
12
+def _get_cos_sin(
13
+ cos_sin_cache_ptr,
14
+ cos_sin_cache_stride,
15
+ pos,
16
+ HALF_ROT_DIM: tl.constexpr,
17
+):
18
+ block = tl.arange(0, HALF_ROT_DIM)
19
+ cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block)
20
+ cos = cos.to(tl.float32)
21
+ sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM)
22
+ sin = sin.to(tl.float32)
23
+ return cos, sin
24
+
25
+
26
+@triton.jit
27
+def _e2m1_nibble(x):
28
+ """Quantize fp32 x (already scale-divided) to E2M1 4-bit nibble in uint8.
29
+ Matches torch.bucketize with boundaries
30
+ [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0] and right=False (each boundary
31
+ belongs to the lower bucket), plus sign bit."""
32
+ abs_x = tl.minimum(tl.abs(x), 6.0)
33
+ code = tl.where(
34
+ abs_x <= 0.25,
35
+ 0.0,
36
+ tl.where(
37
+ abs_x <= 0.75,
38
+ 1.0,
39
+ tl.where(
40
+ abs_x <= 1.25,
41
+ 2.0,
42
+ tl.where(
43
+ abs_x <= 1.75,
44
+ 3.0,
45
+ tl.where(
46
+ abs_x <= 2.5,
47
+ 4.0,
48
+ tl.where(abs_x <= 3.5, 5.0, tl.where(abs_x <= 5.0, 6.0, 7.0)),
49
+ ),
50
+ ),
51
+ ),
52
+ ),
53
+ )
54
+ code_u8 = code.to(tl.uint8)
55
+ sign = ((x < 0) & (code_u8 != 0)).to(tl.uint8)
56
+ return code_u8 | (sign << 3)
57
+
58
+
59
+@triton.jit
60
+def _quantize_mxfp4_pair(x_lo, x_hi):
61
+ """Quantize a block of MXFP4_BLOCK_SIZE fp32 values given as two
62
+ interleaved halves (x_lo = values at even positions in the block,
63
+ x_hi = values at odd positions). Returns:
64
+ - packed : uint8[BLOCK/2] (low nibble = quant(x_lo), high = quant(x_hi))
65
+ - ue8m0 : scalar uint8 (block scale = 2^(ue8m0 - 127))
66
+ """
67
+ amax = tl.maximum(tl.max(tl.abs(x_lo)), tl.max(tl.abs(x_hi)))
68
+ amax = tl.maximum(amax, 1e-4)
69
+ # ue8m0 block scale: 2^ceil(log2(amax/6.0)).
70
+ log2_ratio = tl.math.ceil(tl.math.log2(amax / 6.0))
71
+ log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
72
+ scale = tl.math.exp2(log2_ratio)
73
+ ue8m0 = (log2_ratio + 127.0).to(tl.uint8)
74
+
75
+ inv_scale = 1.0 / scale
76
+ lo_nib = _e2m1_nibble(x_lo * inv_scale)
77
+ hi_nib = _e2m1_nibble(x_hi * inv_scale)
78
+ packed = lo_nib | (hi_nib << 4)
79
+ return packed, ue8m0
80
+
81
+
82
+@triton.jit
83
+def _fused_indexer_q_rope_quant_kernel(
84
+ pos_ptr,
85
+ # Index Q RoPE
86
+ index_q_ptr,
87
+ index_q_stride0,
88
+ index_q_stride1,
89
+ index_q_cos_sin_ptr,
90
+ index_q_cos_sin_stride,
91
+ INDEX_Q_HALF_ROT_DIM: tl.constexpr,
92
+ # Index Q Quantize
93
+ index_q_fp8_ptr,
94
+ index_q_fp8_stride0,
95
+ index_q_fp8_stride1,
96
+ INDEX_Q_HEAD_DIM: tl.constexpr,
97
+ # Index weights
98
+ index_weights_ptr,
99
+ index_weights_stride,
100
+ index_weights_softmax_scale,
101
+ index_weights_head_scale,
102
+ index_weights_out_ptr,
103
+ index_weights_out_stride,
104
+):
105
+ # Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding
106
+ # + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the
107
+ # LAST rope_dim dims of each head; the leading [0, NOPE_DIM) is passed
108
+ # through unchanged.
109
+ INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM
110
+ INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM
111
+ tl.static_assert(INDEX_Q_NOPE_DIM >= 0)
112
+
113
+ tok_idx = tl.program_id(0)
114
+ head_idx = tl.program_id(1)
115
+
116
+ pos = tl.load(pos_ptr + tok_idx)
117
+ cos, sin = _get_cos_sin(
118
+ index_q_cos_sin_ptr,
119
+ index_q_cos_sin_stride,
120
+ pos,
121
+ INDEX_Q_HALF_ROT_DIM,
122
+ )
123
+ half_offset = tl.arange(0, INDEX_Q_HALF_ROT_DIM)
124
+ base_ptr = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
125
+
126
+ # Interleaved (GPT-J) RoPE on dims [NOPE_DIM, HEAD_DIM):
127
+ # even = q[NOPE_DIM + 2*i], odd = q[NOPE_DIM + 2*i + 1]
128
+ rot_base = base_ptr + INDEX_Q_NOPE_DIM
129
+ x_even = tl.load(rot_base + half_offset * 2).to(tl.float32)
130
+ x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32)
131
+ r_even = x_even * cos - x_odd * sin
132
+ r_odd = x_odd * cos + x_even * sin
133
+
134
+ # Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax.
135
+ # Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py).
136
+ r_even = r_even.to(tl.bfloat16).to(tl.float32)
137
+ r_odd = r_odd.to(tl.bfloat16).to(tl.float32)
138
+
139
+ amax = tl.maximum(tl.max(tl.abs(r_even)), tl.max(tl.abs(r_odd)))
140
+ if INDEX_Q_NOPE_DIM > 0:
141
+ nope_offset = tl.arange(0, INDEX_Q_NOPE_DIM)
142
+ x_nope = tl.load(base_ptr + nope_offset).to(tl.float32)
143
+ amax = tl.maximum(amax, tl.max(tl.abs(x_nope)))
144
+ index_q_scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
145
+ index_q_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(index_q_scale)))
146
+
147
+ # Store quantized values to index_q_fp8
148
+ fp8_base_ptr = (
149
+ index_q_fp8_ptr + tok_idx * index_q_fp8_stride0 + head_idx * index_q_fp8_stride1
150
+ )
151
+ if INDEX_Q_NOPE_DIM > 0:
152
+ tl.store(
153
+ fp8_base_ptr + nope_offset,
154
+ tl.div_rn(x_nope, index_q_scale).to(tl.float8e4nv),
155
+ )
156
+ fp8_rot_base = fp8_base_ptr + INDEX_Q_NOPE_DIM
157
+ tl.store(
158
+ fp8_rot_base + half_offset * 2,
159
+ tl.div_rn(r_even, index_q_scale).to(tl.float8e4nv),
160
+ )
161
+ tl.store(
162
+ fp8_rot_base + half_offset * 2 + 1,
163
+ tl.div_rn(r_odd, index_q_scale).to(tl.float8e4nv),
164
+ )
165
+
166
+ # FP8 weight-fold contract:
167
+ # index_weights_out = index_weights * q_scale * softmax_scale * head_scale
168
+ # The per-token-per-head q_scale (fp32) IS folded into the output weights
169
+ # here because FP8 Q is stored WITHOUT a companion scale tensor — the
170
+ # downstream fp8_fp4_mqa_logits/fp8_fp4_paged_mqa_logits kernels use `weights` to
171
+ # apply per-token Q scale inline. See the MXFP4 kernel below for the
172
+ # contrasting convention (scales live with the Q values, weights are NOT
173
+ # q-scaled).
174
+ index_weights = tl.load(
175
+ index_weights_ptr + tok_idx * index_weights_stride + head_idx
176
+ )
177
+ index_weights = index_weights.to(tl.float32)
178
+ index_weights *= index_q_scale
179
+ index_weights *= index_weights_softmax_scale
180
+ index_weights *= index_weights_head_scale
181
+ tl.store(
182
+ index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
183
+ index_weights,
184
+ )
185
+
186
+
187
+@triton.jit
188
+def _fused_indexer_q_rope_mxfp4_kernel(
189
+ pos_ptr,
190
+ # Index Q RoPE input (fp/bf16)
191
+ index_q_ptr,
192
+ index_q_stride0,
193
+ index_q_stride1,
194
+ index_q_cos_sin_ptr,
195
+ index_q_cos_sin_stride,
196
+ INDEX_Q_HALF_ROT_DIM: tl.constexpr,
197
+ # MXFP4 Q outputs
198
+ index_q_mxfp4_ptr, # uint8, (T, H, HEAD_DIM // 2)
199
+ index_q_mxfp4_stride0,
200
+ index_q_mxfp4_stride1,
201
+ index_q_scale_ptr, # uint8 ue8m0, (T, H, HEAD_DIM // BLOCK)
202
+ index_q_scale_stride0,
203
+ index_q_scale_stride1,
204
+ INDEX_Q_HEAD_DIM: tl.constexpr,
205
+ MXFP4_BLOCK: tl.constexpr,
206
+ # Weights (NO per-token q_scale fold for MXFP4; per-block scales stay
207
+ # with the Q values in the output scale tensor).
208
+ index_weights_ptr,
209
+ index_weights_stride,
210
+ index_weights_softmax_scale,
211
+ index_weights_head_scale,
212
+ index_weights_out_ptr,
213
+ index_weights_out_stride,
214
+):
215
+ INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM
216
+ INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM
217
+ NUM_NOPE_BLOCKS: tl.constexpr = INDEX_Q_NOPE_DIM // MXFP4_BLOCK
218
+ NUM_ROPE_BLOCKS: tl.constexpr = INDEX_Q_ROT_DIM // MXFP4_BLOCK
219
+ HALF_BLOCK: tl.constexpr = MXFP4_BLOCK // 2
220
+ tl.static_assert(INDEX_Q_NOPE_DIM >= 0)
221
+ tl.static_assert(INDEX_Q_NOPE_DIM % MXFP4_BLOCK == 0)
222
+ tl.static_assert(INDEX_Q_ROT_DIM % MXFP4_BLOCK == 0)
223
+ tl.static_assert(MXFP4_BLOCK % 2 == 0)
224
+
225
+ tok_idx = tl.program_id(0)
226
+ head_idx = tl.program_id(1)
227
+
228
+ pos = tl.load(pos_ptr + tok_idx)
229
+
230
+ q_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
231
+ out_base = (
232
+ index_q_mxfp4_ptr
233
+ + tok_idx * index_q_mxfp4_stride0
234
+ + head_idx * index_q_mxfp4_stride1
235
+ )
236
+ scale_base = (
237
+ index_q_scale_ptr
238
+ + tok_idx * index_q_scale_stride0
239
+ + head_idx * index_q_scale_stride1
240
+ )
241
+
242
+ half_off = tl.arange(0, HALF_BLOCK)
243
+
244
+ # ---- NoPE blocks: direct load, pair as (even-index, odd-index) values ----
245
+ for b in tl.static_range(NUM_NOPE_BLOCKS):
246
+ base = b * MXFP4_BLOCK
247
+ x_lo = tl.load(q_base + base + half_off * 2).to(tl.float32)
248
+ x_hi = tl.load(q_base + base + half_off * 2 + 1).to(tl.float32)
249
+ packed, ue8m0 = _quantize_mxfp4_pair(x_lo, x_hi)
250
+ tl.store(out_base + base // 2 + half_off, packed)
251
+ tl.store(scale_base + b, ue8m0)
252
+
253
+ # ---- RoPE blocks: apply GPT-J interleaved RoPE to the block's 16 pairs,
254
+ # then quantize. Each block covers HALF_BLOCK (=16) cos/sin pairs. ----
255
+ rot_q_base = q_base + INDEX_Q_NOPE_DIM
256
+ for b in tl.static_range(NUM_ROPE_BLOCKS):
257
+ pair_off = b * HALF_BLOCK + half_off # indices in [0, HALF_ROT_DIM)
258
+ cos_b = tl.load(
259
+ index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + pair_off
260
+ ).to(tl.float32)
261
+ sin_b = tl.load(
262
+ index_q_cos_sin_ptr
263
+ + pos * index_q_cos_sin_stride
264
+ + pair_off
265
+ + INDEX_Q_HALF_ROT_DIM
266
+ ).to(tl.float32)
267
+ x_even = tl.load(rot_q_base + pair_off * 2).to(tl.float32)
268
+ x_odd = tl.load(rot_q_base + pair_off * 2 + 1).to(tl.float32)
269
+ r_even = x_even * cos_b - x_odd * sin_b
270
+ r_odd = x_odd * cos_b + x_even * sin_b
271
+ # bf16 roundtrip for parity with the FP8 kernel / reference numerics.
272
+ r_even = r_even.to(tl.bfloat16).to(tl.float32)
273
+ r_odd = r_odd.to(tl.bfloat16).to(tl.float32)
274
+ packed, ue8m0 = _quantize_mxfp4_pair(r_even, r_odd)
275
+ rope_byte_off = (INDEX_Q_NOPE_DIM + b * MXFP4_BLOCK) // 2
276
+ tl.store(out_base + rope_byte_off + half_off, packed)
277
+ tl.store(scale_base + NUM_NOPE_BLOCKS + b, ue8m0)
278
+
279
+ # MXFP4 weight-fold contract:
280
+ # index_weights_out = index_weights * softmax_scale * head_scale
281
+ # NOTE: q_scale is NOT folded here (contrast with the FP8 kernel above).
282
+ # MXFP4 Q emits a separate ue8m0 scale tensor of shape
283
+ # (T, H, HEAD_DIM // MXFP4_BLOCK) alongside the packed values, so each
284
+ # per-block scale is applied by the downstream MXFP4 logits kernel when
285
+ # dequantizing Q — there is no per-token scalar to fold into `weights`.
286
+ index_weights = tl.load(
287
+ index_weights_ptr + tok_idx * index_weights_stride + head_idx
288
+ ).to(tl.float32)
289
+ index_weights *= index_weights_softmax_scale
290
+ index_weights *= index_weights_head_scale
291
+ tl.store(
292
+ index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
293
+ index_weights,
294
+ )
295
+
296
+
297
+def fused_indexer_q_rope_quant(
298
+ positions: torch.Tensor,
299
+ index_q: torch.Tensor,
300
+ index_q_cos_sin_cache: torch.Tensor,
301
+ # Index weights
302
+ index_weights: torch.Tensor,
303
+ index_weights_softmax_scale: float,
304
+ index_weights_head_scale: float,
305
+ use_fp4: bool = False,
306
+) -> tuple[
307
+ torch.Tensor | tuple[torch.Tensor, torch.Tensor],
308
+ torch.Tensor,
309
+]:
310
+ """Fused RoPE + quantize Q for the sparse indexer.
311
+
312
+ Weight-fold semantics (important — the two paths differ):
313
+
314
+ FP8 path (use_fp4=False, default):
315
+ q_fp8 : (T, H, HEAD_DIM) float8_e4m3fn, per-token-per-head
316
+ scalar scale (NOT stored — folded into weights below)
317
+ weights_out = weights * q_scale * softmax_scale * head_scale
318
+ Rationale: a single per-token q_scale is a scalar the downstream FP8
319
+ logits kernel would otherwise multiply in. Folding it into `weights`
320
+ avoids emitting a separate tensor and is free for the logits kernel.
321
+
322
+ MXFP4 path (use_fp4=True):
323
+ q_packed : (T, H, HEAD_DIM // 2) uint8 (2 E2M1 nibbles per byte)
324
+ q_scale : (T, H, HEAD_DIM // MXFP4_BLOCK_SIZE) uint8 ue8m0 bytes
325
+ weights_out = weights * softmax_scale * head_scale
326
+ Rationale: MXFP4 has PER-BLOCK (32-element) scales that live with
327
+ the Q values — they cannot be folded into a per-token weight
328
+ scalar, so `weights` carries only the softmax and head scales.
329
+
330
+ Returns (q_quant, weights_out) where q_quant is either a Tensor (FP8) or
331
+ a (values, scales) tuple (MXFP4). This matches the union type accepted
332
+ by `SparseAttnIndexer.forward_*`.
333
+ """
334
+ assert positions.ndim == 1
335
+ assert index_q.ndim == 3
336
+ assert index_q_cos_sin_cache.ndim == 2
337
+
338
+ num_tokens = positions.shape[0]
339
+ num_index_q_heads = index_q.shape[1]
340
+ index_q_head_dim = index_q.shape[2]
341
+
342
+ index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
343
+
344
+ if use_fp4:
345
+ assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
346
+ f"head_dim={index_q_head_dim} must be a multiple of MXFP4 block "
347
+ f"size {MXFP4_BLOCK_SIZE}"
348
+ )
349
+ num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE
350
+ index_q_packed = torch.empty(
351
+ (num_tokens, num_index_q_heads, index_q_head_dim // 2),
352
+ dtype=torch.uint8,
353
+ device=index_q.device,
354
+ )
355
+ index_q_scale = torch.empty(
356
+ (num_tokens, num_index_q_heads, num_scale_blocks),
357
+ dtype=torch.uint8,
358
+ device=index_q.device,
359
+ )
360
+ _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)](
361
+ positions,
362
+ index_q,
363
+ index_q.stride(0),
364
+ index_q.stride(1),
365
+ index_q_cos_sin_cache,
366
+ index_q_cos_sin_cache.stride(0),
367
+ index_q_cos_sin_cache.shape[-1] // 2,
368
+ index_q_packed,
369
+ index_q_packed.stride(0),
370
+ index_q_packed.stride(1),
371
+ index_q_scale,
372
+ index_q_scale.stride(0),
373
+ index_q_scale.stride(1),
374
+ index_q_head_dim,
375
+ MXFP4_BLOCK_SIZE,
376
+ index_weights,
377
+ index_weights.stride(0),
378
+ index_weights_softmax_scale,
379
+ index_weights_head_scale,
380
+ index_weights_out,
381
+ index_weights_out.stride(0),
382
+ num_warps=1, # TODO: Tune this
383
+ )
384
+ # Values stay uint8 (2 E2M1 nibbles per byte). Scales are 4 ue8m0
385
+ # bytes per (token, head) reinterpreted as one int32, then squeezed
386
+ # from (T, H, 1) to (T, H) to match DeepGEMM's expected q_sf rank
387
+ # (prefill wants 2-D (seq_len, num_heads); decode reshapes this to
388
+ # 3-D (batch, next_n, num_heads)).
389
+ return (
390
+ index_q_packed,
391
+ index_q_scale.view(torch.int32).squeeze(-1),
392
+ ), index_weights_out
393
+
394
+ index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
395
+ _fused_indexer_q_rope_quant_kernel[(num_tokens, num_index_q_heads)](
396
+ positions,
397
+ index_q,
398
+ index_q.stride(0),
399
+ index_q.stride(1),
400
+ index_q_cos_sin_cache,
401
+ index_q_cos_sin_cache.stride(0),
402
+ index_q_cos_sin_cache.shape[-1] // 2,
403
+ index_q_fp8,
404
+ index_q_fp8.stride(0),
405
+ index_q_fp8.stride(1),
406
+ index_q_head_dim,
407
+ index_weights,
408
+ index_weights.stride(0),
409
+ index_weights_softmax_scale,
410
+ index_weights_head_scale,
411
+ index_weights_out,
412
+ index_weights_out.stride(0),
413
+ num_warps=1, # TODO: Tune this
414
+ )
415
+ return index_q_fp8, index_weights_out
vllm/v1/attention/ops/deepseek_v4_ops/fused_inv_rope_fp8_quant.py
new
+242
@@ -0,0 +1,242 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Fused inverse RoPE + block-scaled FP8 quantization kernel for DeepseekV4 attention.
5
+
6
+Output scale format is pre-transformed (MN-major TMA-aligned; FP32 on SM90,
7
+INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layout.
8
+"""
9
+
10
+import torch
11
+
12
+from vllm.triton_utils import tl, triton
13
+
14
+
15
+@triton.jit
16
+def _fused_inv_rope_fp8_quant_per_head(
17
+ o_ptr,
18
+ positions_ptr,
19
+ cos_sin_cache_ptr,
20
+ fp8_ptr,
21
+ scale_ptr,
22
+ num_tokens,
23
+ heads_per_group: tl.constexpr,
24
+ o_stride_token,
25
+ o_stride_head,
26
+ cache_stride_pos,
27
+ fp8_stride_group,
28
+ fp8_stride_token,
29
+ scale_stride_group,
30
+ scale_stride_k,
31
+ fp8_max: tl.constexpr,
32
+ eps: tl.constexpr,
33
+ QUANT_GROUP_SIZE: tl.constexpr,
34
+ CHUNKS_PER_HEAD: tl.constexpr,
35
+ ROPE_START: tl.constexpr,
36
+ HALF_ROPE: tl.constexpr,
37
+ TMA_ALIGNED_SCALES: tl.constexpr,
38
+):
39
+ # int64: stride multiply overflows int32 past num_tokens=32768 (IMA).
40
+ pid_token = tl.program_id(0).to(tl.int64)
41
+ pid_gh = tl.program_id(1).to(tl.int64)
42
+
43
+ g = pid_gh // heads_per_group
44
+ head_in_group = pid_gh % heads_per_group
45
+ global_head = pid_gh
46
+ qb_start = head_in_group * CHUNKS_PER_HEAD
47
+
48
+ # Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant.
49
+ if pid_token >= num_tokens:
50
+ if TMA_ALIGNED_SCALES:
51
+ scale_addr = (
52
+ scale_ptr
53
+ + g * scale_stride_group
54
+ + pid_token
55
+ + head_in_group * scale_stride_k
56
+ )
57
+ tl.store(scale_addr, tl.zeros((), dtype=tl.int32))
58
+ else:
59
+ block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
60
+ qb_indices = qb_start + block_offsets
61
+ scale_addrs = (
62
+ scale_ptr
63
+ + g * scale_stride_group
64
+ + pid_token
65
+ + qb_indices * scale_stride_k
66
+ )
67
+ tl.store(scale_addrs, tl.zeros((CHUNKS_PER_HEAD,), dtype=tl.float32))
68
+ return
69
+
70
+ input_base = o_ptr + pid_token * o_stride_token + global_head * o_stride_head
71
+
72
+ HEAD_DIM: tl.constexpr = CHUNKS_PER_HEAD * QUANT_GROUP_SIZE
73
+ offsets = tl.arange(0, HEAD_DIM)
74
+ x = tl.load(input_base + offsets).to(tl.float32)
75
+
76
+ rope_abs_start: tl.constexpr = (CHUNKS_PER_HEAD - 1) * QUANT_GROUP_SIZE + ROPE_START
77
+ pos = tl.load(positions_ptr + pid_token)
78
+ cache_base = cos_sin_cache_ptr + pos * cache_stride_pos
79
+ is_rope = offsets >= rope_abs_start
80
+ rope_local = offsets - rope_abs_start
81
+
82
+ x_partner = tl.load(input_base + (offsets ^ 1), mask=is_rope, other=0.0).to(
83
+ tl.float32
84
+ )
85
+ cs_idx = tl.maximum(rope_local >> 1, 0)
86
+ cos_v = tl.load(cache_base + cs_idx, mask=is_rope, other=1.0)
87
+ sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0)
88
+ x_add = x * cos_v + x_partner * sin_v
89
+ x_sub = x * cos_v - x_partner * sin_v
90
+ is_even = (rope_local & 1) == 0
91
+ rotated = tl.where(is_even, x_add, x_sub)
92
+ x = tl.where(is_rope, rotated, x)
93
+
94
+ x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE))
95
+ block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps)
96
+ scale_raw = block_absmax * (1.0 / fp8_max)
97
+ scales = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
98
+
99
+ scales_exp = tl.reshape(
100
+ tl.broadcast_to(
101
+ tl.reshape(scales, (CHUNKS_PER_HEAD, 1)),
102
+ (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE),
103
+ ),
104
+ (HEAD_DIM,),
105
+ )
106
+ x_quant = tl.clamp(x / scales_exp, -fp8_max, fp8_max).to(tl.float8e4nv)
107
+
108
+ fp8_base = (
109
+ fp8_ptr
110
+ + g * fp8_stride_group
111
+ + pid_token * fp8_stride_token
112
+ + qb_start * QUANT_GROUP_SIZE
113
+ )
114
+ tl.store(fp8_base + offsets, x_quant)
115
+
116
+ block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
117
+ qb_indices = qb_start + block_offsets
118
+ if TMA_ALIGNED_SCALES:
119
+ scale_bits = scales.to(tl.int32, bitcast=True)
120
+ ue8m0_bytes = (scale_bits >> 23) & 0xFF
121
+ packed_val = tl.sum(ue8m0_bytes << (block_offsets * 8))
122
+ scale_addr = (
123
+ scale_ptr
124
+ + g * scale_stride_group
125
+ + pid_token
126
+ + head_in_group * scale_stride_k
127
+ )
128
+ tl.store(scale_addr, packed_val)
129
+ else:
130
+ scale_addrs = (
131
+ scale_ptr + g * scale_stride_group + pid_token + qb_indices * scale_stride_k
132
+ )
133
+ tl.store(scale_addrs, scales)
134
+
135
+
136
+def fused_inv_rope_fp8_quant(
137
+ o: torch.Tensor,
138
+ positions: torch.Tensor,
139
+ cos_sin_cache: torch.Tensor,
140
+ n_groups: int,
141
+ heads_per_group: int,
142
+ nope_dim: int = 448,
143
+ rope_dim: int = 64,
144
+ quant_group_size: int = 128,
145
+ tma_aligned_scales: bool = False,
146
+) -> tuple[torch.Tensor, torch.Tensor]:
147
+ """Fused inverse RoPE + block-scaled FP8 quantization.
148
+
149
+ Args:
150
+ o: Attention output [num_tokens, num_heads, head_dim] bf16.
151
+ positions: Token positions [num_tokens] int64.
152
+ cos_sin_cache: Precomputed [max_pos, rope_dim] with cos||sin.
153
+ n_groups: Number of output groups.
154
+ heads_per_group: Heads per group.
155
+ nope_dim: Non-RoPE dimensions per head (default 448).
156
+ rope_dim: RoPE dimensions per head (default 64).
157
+ quant_group_size: FP8 quantization block size (default 128).
158
+ tma_aligned_scales: Output INT32 packed UE8M0 for SM100 (True)
159
+ or FP32 for SM90 (False).
160
+
161
+ Returns:
162
+ o_fp8: [T, G, D] float8_e4m3fn, strides (D, T*D, 1).
163
+ o_scale: Pre-transformed scale tensor for fp8_einsum.
164
+ """
165
+ from vllm.utils.deep_gemm import get_tma_aligned_size
166
+
167
+ num_tokens, num_heads, head_dim = o.shape
168
+ assert num_heads == n_groups * heads_per_group
169
+ assert head_dim == nope_dim + rope_dim
170
+ assert head_dim % quant_group_size == 0
171
+ assert nope_dim % quant_group_size == (quant_group_size - rope_dim)
172
+ assert rope_dim % 2 == 0
173
+ assert cos_sin_cache.shape[-1] == rope_dim
174
+ assert cos_sin_cache.dtype == torch.float32
175
+
176
+ d = heads_per_group * head_dim
177
+ num_scale_blocks = d // quant_group_size
178
+ chunks_per_head = head_dim // quant_group_size
179
+
180
+ fp8_dtype = torch.float8_e4m3fn
181
+ fp8_max = torch.finfo(fp8_dtype).max
182
+
183
+ fp8_buf = torch.empty(
184
+ (n_groups, num_tokens, d),
185
+ dtype=fp8_dtype,
186
+ device=o.device,
187
+ )
188
+
189
+ tma_aligned_T = get_tma_aligned_size(num_tokens, 4)
190
+ if tma_aligned_scales:
191
+ packed_sf_k = (num_scale_blocks + 3) // 4
192
+ scale_buf = torch.empty(
193
+ n_groups * packed_sf_k * tma_aligned_T,
194
+ dtype=torch.int32,
195
+ device=o.device,
196
+ ).as_strided(
197
+ (n_groups, num_tokens, packed_sf_k),
198
+ (packed_sf_k * tma_aligned_T, 1, tma_aligned_T),
199
+ )
200
+ else:
201
+ scale_buf = torch.empty(
202
+ n_groups * num_scale_blocks * tma_aligned_T,
203
+ dtype=torch.float32,
204
+ device=o.device,
205
+ ).as_strided(
206
+ (n_groups, num_tokens, num_scale_blocks),
207
+ (num_scale_blocks * tma_aligned_T, 1, tma_aligned_T),
208
+ )
209
+
210
+ common_args = dict(
211
+ heads_per_group=heads_per_group,
212
+ o_stride_token=o.stride(0),
213
+ o_stride_head=o.stride(1),
214
+ cache_stride_pos=cos_sin_cache.stride(0),
215
+ fp8_stride_group=fp8_buf.stride(0),
216
+ fp8_stride_token=fp8_buf.stride(1),
217
+ scale_stride_group=scale_buf.stride(0),
218
+ scale_stride_k=scale_buf.stride(2),
219
+ fp8_max=fp8_max,
220
+ eps=1e-10,
221
+ QUANT_GROUP_SIZE=quant_group_size,
222
+ CHUNKS_PER_HEAD=chunks_per_head,
223
+ ROPE_START=nope_dim % quant_group_size,
224
+ HALF_ROPE=rope_dim // 2,
225
+ TMA_ALIGNED_SCALES=tma_aligned_scales,
226
+ num_stages=1,
227
+ launch_pdl=False,
228
+ )
229
+
230
+ grid = (tma_aligned_T, n_groups * heads_per_group)
231
+ _fused_inv_rope_fp8_quant_per_head[grid](
232
+ o,
233
+ positions,
234
+ cos_sin_cache,
235
+ fp8_buf,
236
+ scale_buf,
237
+ num_tokens,
238
+ **common_args,
239
+ num_warps=1,
240
+ )
241
+
242
+ return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1)
vllm/v1/attention/ops/deepseek_v4_ops/fused_qk_rmsnorm.py
new
+96
@@ -0,0 +1,96 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.triton_utils import tl, triton
6
+
7
+
8
+@triton.jit
9
+def _fused_q_kv_rmsnorm_kernel(
10
+ q_ptr,
11
+ q_out_ptr,
12
+ q_weight_ptr,
13
+ q_in_stride,
14
+ q_out_stride,
15
+ kv_ptr,
16
+ kv_out_ptr,
17
+ kv_weight_ptr,
18
+ kv_in_stride,
19
+ kv_out_stride,
20
+ eps,
21
+ Q_SIZE: tl.constexpr,
22
+ KV_SIZE: tl.constexpr,
23
+ BLOCK_SIZE: tl.constexpr,
24
+):
25
+ # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y.
26
+ # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes
27
+ # the launch at max-num-batched-tokens >= 65536 with "invalid argument".
28
+ # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32
29
+ # past num_tokens ~87K under large chunked prefill.
30
+ token_idx = tl.program_id(0).to(tl.int64)
31
+ pid_task = tl.program_id(1)
32
+
33
+ if pid_task == 0:
34
+ SIZE = Q_SIZE
35
+ row_in = q_ptr + token_idx * q_in_stride
36
+ weight_ptr = q_weight_ptr
37
+ row_out = q_out_ptr + token_idx * q_out_stride
38
+ else:
39
+ SIZE = KV_SIZE
40
+ row_in = kv_ptr + token_idx * kv_in_stride
41
+ weight_ptr = kv_weight_ptr
42
+ row_out = kv_out_ptr + token_idx * kv_out_stride
43
+
44
+ # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
45
+ # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
46
+ # keep x, rrms, and w all in fp32 and perform a single cast at store.
47
+ block = tl.arange(0, BLOCK_SIZE)
48
+ mask = block < SIZE
49
+ x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
50
+ variance = tl.sum(x * x, axis=0) / SIZE
51
+ rrms = tl.rsqrt(variance + eps)
52
+ w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
53
+ y = x * rrms * w
54
+ tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
55
+
56
+
57
+def fused_q_kv_rmsnorm(
58
+ qr: torch.Tensor,
59
+ kv: torch.Tensor,
60
+ q_weight: torch.Tensor,
61
+ kv_weight: torch.Tensor,
62
+ eps: float,
63
+) -> tuple[torch.Tensor, torch.Tensor]:
64
+ assert qr.ndim == 2 and kv.ndim == 2
65
+ assert qr.shape[0] == kv.shape[0], (
66
+ f"token dim mismatch: qr={qr.shape}, kv={kv.shape}"
67
+ )
68
+ assert qr.stride(-1) == 1 and kv.stride(-1) == 1
69
+ assert q_weight.is_contiguous() and kv_weight.is_contiguous()
70
+
71
+ q_size = qr.shape[1]
72
+ kv_size = kv.shape[1]
73
+ num_tokens = qr.shape[0]
74
+ qr_out = torch.empty_like(qr)
75
+ kv_out = torch.empty_like(kv)
76
+ if num_tokens == 0:
77
+ return qr_out, kv_out
78
+
79
+ block_size = triton.next_power_of_2(max(q_size, kv_size))
80
+ _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)](
81
+ qr,
82
+ qr_out,
83
+ q_weight,
84
+ qr.stride(0),
85
+ qr_out.stride(0),
86
+ kv,
87
+ kv_out,
88
+ kv_weight,
89
+ kv.stride(0),
90
+ kv_out.stride(0),
91
+ eps,
92
+ Q_SIZE=q_size,
93
+ KV_SIZE=kv_size,
94
+ BLOCK_SIZE=block_size,
95
+ )
96
+ return qr_out, kv_out
vllm/v1/core/kv_cache_coordinator.py
+60
-33
@@ -54,8 +54,14 @@ class KVCacheCoordinator(ABC):
54
metrics_collector,
55
)
56
57
- # Needs special handling for find_longest_cache_hit if eagle is enabled
58
- self.use_eagle = use_eagle
57
+ # KV cache group indices that get the EAGLE last-block drop.
58
+ self.eagle_group_ids: set[int] = {
59
+ i for i, g in enumerate(kv_cache_config.kv_cache_groups) if g.is_eagle_group
60
+ }
61
+ # Conservatively fall back to flag all groups when no group is flagged.
62
+ if use_eagle and not self.eagle_group_ids:
63
+ self.eagle_group_ids = set(range(len(kv_cache_config.kv_cache_groups)))
64
+
65
self.single_type_managers = tuple(
66
get_manager_for_kv_cache_spec(
67
kv_cache_spec=kv_cache_group.kv_cache_spec,
@@ -357,7 +363,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
363
kv_cache_group_ids=[0],
364
block_pool=self.block_pool,
365
kv_cache_spec=self.kv_cache_spec,
360
- use_eagle=self.use_eagle,
366
+ use_eagle=0 in self.eagle_group_ids,
367
alignment_tokens=self.block_size,
368
dcp_world_size=self.dcp_world_size,
369
pcp_world_size=self.pcp_world_size,
@@ -450,6 +456,14 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
456
block_sizes = [spec.block_size for spec, _, _ in attention_groups]
457
self.lcm_block_size = lcm(*block_sizes)
458
459
+ # Attention-group indices (into ``self.attention_groups``) that
460
+ # contain at least one EAGLE/MTP KV cache group.
461
+ self.eagle_attn_group_indices: set[int] = {
462
+ i
463
+ for i, (_, group_ids, _) in enumerate(self.attention_groups)
464
+ if any(gid in self.eagle_group_ids for gid in group_ids)
465
+ }
466
+
467
def find_longest_cache_hit(
468
self,
469
block_hashes: list[BlockHash],
@@ -485,49 +499,62 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
499
hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups
500
501
# Simple hybrid (1 full attn + 1 other): one iteration suffices.
488
- # Full attn is always first if it exists. This avoids EAGLE drops
489
- # being applied multiple times to non-full-attn groups.
490
- # FIXME (yifan): However, for complex hybrid models with multiple attn
491
- # groups, we still have the EAGLE spiral block dropping problem. See
492
- # discussion in issue https://github.com/vllm-project/vllm/issues/32802.
502
+ # Full attn is always first if it exists.
503
is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(
504
self.attention_groups[0][0], FullAttentionSpec
505
)
506
507
+ # Attention-group indices whose EAGLE drop is verified at the current
508
+ # ``curr_hit_length``. Each eagle group applies the drop at most once
509
+ # per candidate length (see issue #32802).
510
+ eagle_verified: set[int] = set()
511
+
512
while True:
513
curr_hit_length = hit_length
514
500
- for spec, group_ids, manager_cls in self.attention_groups:
501
- is_full_attn = isinstance(spec, FullAttentionSpec)
502
-
503
- # Full attention: reuse cached blocks (downward-closed property)
515
+ for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups):
516
cached_blocks = hit_blocks_by_group[group_ids[0]]
505
- if is_full_attn and cached_blocks is not None:
506
- # For full attention, we only need to compute the cache hit
507
- # length once. Starting from the second iteration, if the
508
- # curr_hit_length is reduced by other groups, we can simply
509
- # keep the first (curr_hit_length // block_size) blocks from
510
- # the last iteration.
511
- num_blocks = curr_hit_length // spec.block_size
512
- curr_hit_length = num_blocks * spec.block_size
513
- else:
514
- hit_blocks = manager_cls.find_longest_cache_hit(
515
- block_hashes=_get_block_hashes(spec),
516
- max_length=curr_hit_length,
517
- kv_cache_group_ids=group_ids,
518
- block_pool=self.block_pool,
519
- kv_cache_spec=spec,
520
- use_eagle=self.use_eagle,
521
- alignment_tokens=self.lcm_block_size,
517
+ if isinstance(spec, FullAttentionSpec) and cached_blocks is not None:
518
+ # Full attention is downward-closed: we only need to look
519
+ # up cached blocks once; on subsequent iterations just trim
520
+ # to the (reduced) current hit length.
521
+ curr_hit_length = (
522
+ curr_hit_length // spec.block_size * spec.block_size
523
)
523
- curr_hit_length = len(hit_blocks[0]) * spec.block_size
524
- for group_id, blocks in zip(group_ids, hit_blocks):
525
- hit_blocks_by_group[group_id] = blocks
524
+ continue
525
+
526
+ use_eagle = (
527
+ idx in self.eagle_attn_group_indices and idx not in eagle_verified
528
+ )
529
+
530
+ _max_length = curr_hit_length
531
+ if use_eagle:
532
+ # Eagle needs to match one more block and then pop the last.
533
+ _max_length = min(
534
+ curr_hit_length + spec.block_size, max_cache_hit_length
535
+ )
536
+ hit_blocks = manager_cls.find_longest_cache_hit(
537
+ block_hashes=_get_block_hashes(spec),
538
+ max_length=_max_length,
539
+ kv_cache_group_ids=group_ids,
540
+ block_pool=self.block_pool,
541
+ kv_cache_spec=spec,
542
+ use_eagle=use_eagle,
543
+ alignment_tokens=self.lcm_block_size,
544
+ )
545
+ _new_hit_length = len(hit_blocks[0]) * spec.block_size
546
+ if use_eagle:
547
+ eagle_verified.add(idx)
548
+ elif _new_hit_length < curr_hit_length:
549
+ # length shrunk; invalidate previous eagle verifications
550
+ eagle_verified.clear()
551
+ curr_hit_length = _new_hit_length
552
+ for group_id, blocks in zip(group_ids, hit_blocks):
553
+ hit_blocks_by_group[group_id] = blocks
554
555
if curr_hit_length >= hit_length:
556
break
557
hit_length = curr_hit_length
530
- # Simple hybrid: exit after one iteration
558
if is_simple_hybrid:
559
break
560
vllm/v1/core/kv_cache_utils.py
+412
-10
@@ -4,18 +4,19 @@
4
5
import copy
6
import hashlib
7
+import math
8
import os
9
from collections import defaultdict
10
from collections.abc import Callable, Iterable, Iterator, Sequence
11
from dataclasses import dataclass, replace
12
from functools import partial
12
-from typing import Any, NewType, TypeAlias, overload
13
+from typing import Any, NewType, TypeAlias, cast, overload
14
15
from vllm import envs
16
from vllm.config import VllmConfig
17
from vllm.logger import init_logger
18
from vllm.utils.hashing import sha256_cbor, xxhash_cbor
18
-from vllm.utils.math_utils import cdiv
19
+from vllm.utils.math_utils import cdiv, round_up
20
from vllm.utils.mem_utils import format_gib
21
from vllm.v1.kv_cache_interface import (
22
ChunkedLocalAttentionSpec,
@@ -24,6 +25,9 @@ from vllm.v1.kv_cache_interface import (
25
KVCacheGroupSpec,
26
KVCacheSpec,
27
KVCacheTensor,
28
+ MambaSpec,
29
+ MLAAttentionSpec,
30
+ SlidingWindowMLASpec,
31
SlidingWindowSpec,
32
UniformTypeKVCacheSpecs,
33
)
@@ -562,6 +566,72 @@ def hash_block_tokens(
566
)
567
568
569
+def resolve_kv_cache_block_sizes(
570
+ kv_cache_config: KVCacheConfig,
571
+ vllm_config: VllmConfig,
572
+) -> tuple[int, int]:
573
+ """Resolve (scheduler_block_size, hash_block_size).
574
+
575
+ - ``scheduler_block_size`` is the token-alignment invariant used by the
576
+ scheduler (e.g. for ``num_computed_tokens`` rounding). Single group:
577
+ ``cache_config.block_size * dcp * pcp``. Multiple groups: LCM of every
578
+ group's block size — context parallelism is not supported here.
579
+ - ``hash_block_size`` is the granularity at which ``Request.block_hashes``
580
+ is computed. Single group: equals scheduler block size. Multiple groups:
581
+ ``cache_config.hash_block_size`` override if set, else the GCD of group
582
+ block sizes; every group's block size must be divisible by it. Returns
583
+ the scheduler block size (i.e. disables finer hashing) if block hashing
584
+ is inactive or a mamba group's block size diverges from the cache
585
+ block size (mamba_cache_mode != "align").
586
+ """
587
+ cache_config = vllm_config.cache_config
588
+ dcp = vllm_config.parallel_config.decode_context_parallel_size
589
+ pcp = vllm_config.parallel_config.prefill_context_parallel_size
590
+ groups = kv_cache_config.kv_cache_groups
591
+
592
+ if len(groups) <= 1: # Single group: block_size * dcp * pcp
593
+ bs = cache_config.block_size * dcp * pcp
594
+ return bs, bs
595
+
596
+ if dcp != 1 or pcp != 1:
597
+ raise ValueError(
598
+ "Hybrid KV cache groups with multiple block sizes do not "
599
+ "support context parallelism (dcp_world_size/pcp_world_size > 1)."
600
+ )
601
+
602
+ group_block_sizes = [g.kv_cache_spec.block_size for g in groups]
603
+ scheduler_block_size = math.lcm(*group_block_sizes)
604
+
605
+ # Block hashes are only consumed by prefix caching and KV connectors
606
+ # (P/D, offloading); when neither is active, keep hash_block_size equal
607
+ # to the scheduler block size.
608
+ connector_enabled = vllm_config.kv_transfer_config is not None
609
+ if not (cache_config.enable_prefix_caching or connector_enabled):
610
+ return scheduler_block_size, scheduler_block_size
611
+
612
+ # Mamba groups with block_size != cache_config.block_size
613
+ # (mamba_cache_mode != "align") break divisibility; back off to the
614
+ # scheduler block size.
615
+ if any(
616
+ isinstance(g.kv_cache_spec, MambaSpec)
617
+ and g.kv_cache_spec.block_size != cache_config.block_size
618
+ for g in groups
619
+ ):
620
+ return scheduler_block_size, scheduler_block_size
621
+
622
+ requested = cache_config.hash_block_size
623
+ hash_block_size = (
624
+ requested if requested is not None else math.gcd(*group_block_sizes)
625
+ )
626
+ if any(bs % hash_block_size != 0 for bs in group_block_sizes):
627
+ raise ValueError(
628
+ f"Invalid hash_block_size={hash_block_size}; all KV cache group "
629
+ f"block sizes must be divisible by hash_block_size. "
630
+ f"Got group block sizes={group_block_sizes}."
631
+ )
632
+ return scheduler_block_size, hash_block_size
633
+
634
+
635
def get_request_block_hasher(
636
block_size: int,
637
caching_hash_fn: Callable[[Any], bytes],
@@ -1089,6 +1159,63 @@ def _get_kv_cache_groups_uniform_page_size(
1159
return create_kv_cache_group_specs(kv_cache_spec, grouped_layers)
1160
1161
1162
+def _get_kv_cache_config_deepseek_v4(
1163
+ vllm_config: VllmConfig,
1164
+ kv_cache_groups: list[KVCacheGroupSpec],
1165
+ available_memory: int,
1166
+) -> tuple[int, list[KVCacheTensor]]:
1167
+ """DeepseekV4 KV cache tensor layout planning.
1168
+
1169
+ Precondition: kv_cache_groups[0] is the full-MLA group; its page sizes
1170
+ define the canonical bucket set. Non-full-MLA groups must have been
1171
+ page_size-padded upstream (see _get_kv_cache_groups_uniform_groups) so
1172
+ every layer's page_size matches one of the full-MLA bucket sizes.
1173
+
1174
+ For each group, bucket its layers by page_size_bytes and place each
1175
+ layer at tuple_idx = position-within-bucket. Emit one KVCacheTensor
1176
+ per (tuple_idx, bucket) whose shared_by is the union of per-group
1177
+ layers at that slot.
1178
+ """
1179
+ full_mla_spec = kv_cache_groups[0].kv_cache_spec
1180
+ assert isinstance(full_mla_spec, UniformTypeKVCacheSpecs)
1181
+ page_sizes = sorted(full_mla_spec.get_page_sizes())
1182
+ layer_tuple_page_bytes = sum(page_sizes)
1183
+
1184
+ # Pre-bucket each group's layers by page_size (registration order within
1185
+ # bucket). bucketed[g_idx][page_size] = [layer_name, ...].
1186
+ bucketed: list[dict[int, list[str]]] = []
1187
+ for group in kv_cache_groups:
1188
+ assert isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1189
+ specs = group.kv_cache_spec.kv_cache_specs
1190
+ b: dict[int, list[str]] = defaultdict(list)
1191
+ for name in group.layer_names:
1192
+ b[specs[name].page_size_bytes].append(name)
1193
+ bucketed.append(b)
1194
+
1195
+ # num_layer_tuples = longest bucket list across all groups. For the
1196
+ # full-MLA group this equals the count of layers in the largest
1197
+ # per-page-size bucket (= get_num_layer_tuples()); for SWA sub-groups
1198
+ # this equals the sub-group size (each has a single page_size).
1199
+ num_layer_tuples = max(len(layers) for b in bucketed for layers in b.values())
1200
+
1201
+ num_blocks = available_memory // (layer_tuple_page_bytes * num_layer_tuples)
1202
+ num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1203
+
1204
+ kv_cache_tensors: list[KVCacheTensor] = []
1205
+ for tuple_idx in range(num_layer_tuples):
1206
+ for ps in page_sizes:
1207
+ shared_by: list[str] = []
1208
+ for b in bucketed:
1209
+ bucket = b.get(ps)
1210
+ if bucket is not None and tuple_idx < len(bucket):
1211
+ shared_by.append(bucket[tuple_idx])
1212
+ kv_cache_tensors.append(
1213
+ KVCacheTensor(size=ps * num_blocks, shared_by=shared_by)
1214
+ )
1215
+
1216
+ return num_blocks, kv_cache_tensors
1217
+
1218
+
1219
def get_kv_cache_config_from_groups(
1220
vllm_config: VllmConfig,
1221
kv_cache_groups: list[KVCacheGroupSpec],
@@ -1120,7 +1247,7 @@ def get_kv_cache_config_from_groups(
1247
kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
1248
):
1249
# Special case: all layers have the same type of KV cache but with
1123
- # different hidden size. Allocate different amount of memory for each
1250
+ # different hidden sizes. Allocate different amount of memory for each
1251
# layer based on its hidden size.
1252
num_blocks = (
1253
available_memory // kv_cache_groups[0].kv_cache_spec.page_size_bytes
@@ -1136,6 +1263,15 @@ def get_kv_cache_config_from_groups(
1263
)
1264
for layer_name in kv_cache_groups[0].layer_names
1265
]
1266
+ elif all(
1267
+ isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1268
+ for group in kv_cache_groups
1269
+ ):
1270
+ # DeepseekV4: UniformTypeKVCacheSpecs but multiple groups.
1271
+ # Delegate to the DeepseekV4-specific allocator.
1272
+ num_blocks, kv_cache_tensors = _get_kv_cache_config_deepseek_v4(
1273
+ vllm_config, kv_cache_groups, available_memory
1274
+ )
1275
else:
1276
# General case:
1277
# We will have group_size memory pools, each is shared by one layer from
@@ -1206,9 +1342,41 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1342
has_chunked_local_attention = any(
1343
isinstance(spec, ChunkedLocalAttentionSpec) for spec in kv_cache_spec.values()
1344
)
1345
+ has_swa_mla = any(
1346
+ isinstance(spec, SlidingWindowMLASpec) for spec in kv_cache_spec.values()
1347
+ )
1348
+
1349
+ uniform_block_size: int | None = None
1350
+ if has_swa_mla:
1351
+ # For DeepseekV4, block sizes can be different for different KV cache groups.
1352
+ # E.g., Full MLA: 256; SWA MLA: 64; C4 partial states: 4, C128 states: 8.
1353
+ assert has_full_attention
1354
+ any_full_spec = next(
1355
+ iter(
1356
+ spec
1357
+ for spec in kv_cache_spec.values()
1358
+ if isinstance(spec, FullAttentionSpec)
1359
+ )
1360
+ )
1361
+ uniform_block_size = any_full_spec.block_size
1362
+
1363
if has_full_attention and (has_sliding_window or has_chunked_local_attention):
1364
for layer_name, spec in kv_cache_spec.items():
1211
- if isinstance(spec, SlidingWindowSpec):
1365
+ if isinstance(spec, SlidingWindowMLASpec):
1366
+ kv_cache_spec[layer_name] = MLAAttentionSpec(
1367
+ block_size=uniform_block_size
1368
+ if uniform_block_size is not None
1369
+ else spec.block_size,
1370
+ num_kv_heads=spec.num_kv_heads,
1371
+ head_size=spec.head_size,
1372
+ dtype=spec.dtype,
1373
+ page_size_padded=spec.page_size_padded,
1374
+ cache_dtype_str=spec.cache_dtype_str,
1375
+ alignment=spec.alignment,
1376
+ compress_ratio=spec.compress_ratio,
1377
+ model_version=spec.model_version,
1378
+ )
1379
+ elif isinstance(spec, SlidingWindowSpec):
1380
kv_cache_spec[layer_name] = FullAttentionSpec(
1381
block_size=spec.block_size,
1382
num_kv_heads=spec.num_kv_heads,
@@ -1237,6 +1405,204 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1405
)
1406
1407
1408
+def group_and_unify_kv_cache_specs(
1409
+ kv_cache_spec: dict[str, KVCacheSpec],
1410
+) -> list[UniformTypeKVCacheSpecs] | None:
1411
+ """
1412
+ Group the KV cache specs and unify each group into one UniformTypeKVCacheSpecs.
1413
+ Currently, this is only used for DeepseekV4.
1414
+ """
1415
+ if not any(
1416
+ isinstance(spec, SlidingWindowMLASpec) for spec in kv_cache_spec.values()
1417
+ ):
1418
+ return None
1419
+
1420
+ mla_specs: dict[str, KVCacheSpec] = {}
1421
+ grouped_swa_mla_specs: dict[tuple[int, int], dict[str, KVCacheSpec]] = defaultdict(
1422
+ dict
1423
+ )
1424
+ # NOTE: Here we group SWA layers by (block_size, sliding_window), which separates
1425
+ # SWA layers, C4I+C4A layers, and C128A layers into three different groups. It can
1426
+ # be fragile with only block_size and sliding_window as keys, but fine for now.
1427
+ for name, spec in kv_cache_spec.items():
1428
+ if isinstance(spec, SlidingWindowMLASpec):
1429
+ grouped_swa_mla_specs[(spec.block_size, spec.sliding_window)][name] = spec
1430
+ elif isinstance(spec, MLAAttentionSpec):
1431
+ mla_specs[name] = spec
1432
+
1433
+ assert len(mla_specs) > 0
1434
+ mla_uniform_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs)
1435
+ assert mla_uniform_spec is not None
1436
+
1437
+ swa_uniform_specs: list[UniformTypeKVCacheSpecs] = []
1438
+ for spec_dict in grouped_swa_mla_specs.values():
1439
+ uniform_spec = UniformTypeKVCacheSpecs.from_specs(spec_dict)
1440
+ assert uniform_spec is not None
1441
+ swa_uniform_specs.append(uniform_spec)
1442
+
1443
+ return [mla_uniform_spec, *swa_uniform_specs]
1444
+
1445
+
1446
+def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int:
1447
+ """Pick a chunk size that minimizes total upward padding.
1448
+
1449
+ Each x is rounded up to a multiple of d:
1450
+
1451
+ x -> ceil(x / d) * d
1452
+
1453
+ Total padding is:
1454
+
1455
+ pad(d) = sum_i (ceil(x_i / d) * d - x_i)
1456
+
1457
+ We brute-force d in [lower_bound, max(values)] (fine for small lists / small
1458
+ maxima) and return the d with minimum padding. Ties prefer larger d.
1459
+ """
1460
+ if not values:
1461
+ raise ValueError("values must be non-empty")
1462
+ if any(x <= 0 for x in values):
1463
+ raise ValueError(f"values must be positive, got: {list(values)!r}")
1464
+
1465
+ min_d = max(1, lower_bound if lower_bound is not None else 1)
1466
+ max_d = max(values)
1467
+ if min_d > max_d:
1468
+ return min_d
1469
+
1470
+ best_d = min_d
1471
+ best_pad: int | None = None
1472
+ for d in range(min_d, max_d + 1):
1473
+ pad = sum((d - (x % d)) % d for x in values)
1474
+ if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d):
1475
+ best_pad = pad
1476
+ best_d = d
1477
+
1478
+ return best_d
1479
+
1480
+
1481
+def _get_kv_cache_groups_uniform_groups(
1482
+ grouped_specs: list[UniformTypeKVCacheSpecs],
1483
+) -> list[KVCacheGroupSpec]:
1484
+ """
1485
+ Generate the KV cache groups from the grouped specs.
1486
+ """
1487
+ assert len(grouped_specs) > 0 and all(
1488
+ isinstance(spec, UniformTypeKVCacheSpecs) for spec in grouped_specs
1489
+ )
1490
+ # For now, we restrict the first grouped_spec to be UniformTypeKVCacheSpecs
1491
+ # containing only MLAAttentionSpec.
1492
+ full_mla_spec = grouped_specs[0]
1493
+ assert all(
1494
+ isinstance(spec, MLAAttentionSpec)
1495
+ for spec in full_mla_spec.kv_cache_specs.values()
1496
+ )
1497
+ full_mla_group = KVCacheGroupSpec(
1498
+ layer_names=list(full_mla_spec.kv_cache_specs.keys()),
1499
+ kv_cache_spec=full_mla_spec,
1500
+ )
1501
+
1502
+ # We define a layer tuple as a group of layers with different page sizes, and
1503
+ # one UniformTypeKVCacheSpecs contains a list of layer tuples.
1504
+ # For example, if we have 11 C4 layers and 10 C128 layers, we can define a layer
1505
+ # tuple as [C4I, C4A, C128], and the full_mla_group will contain "11" layer tuples.
1506
+ # The other uniform KV cache specs will be similarly partitioned into layer tuples.
1507
+ # Say we have 21 SWA layers, all with the same page size, then we will have "21"
1508
+ # layer tuples.
1509
+ num_layer_tuples_per_group: list[int] = [
1510
+ g_spec.get_num_layer_tuples() for g_spec in grouped_specs
1511
+ ]
1512
+ # Choose `num_layer_tuples` to minimize total padding across groups.
1513
+ num_layer_tuples = _approximate_gcd(
1514
+ num_layer_tuples_per_group, lower_bound=num_layer_tuples_per_group[0]
1515
+ )
1516
+ # Round up to the nearest multiple of `num_layer_tuples` (i.e., padding)
1517
+ num_layer_tuples_per_group = [
1518
+ round_up(x, num_layer_tuples) for x in num_layer_tuples_per_group
1519
+ ]
1520
+
1521
+ swa_mla_specs = grouped_specs[1:]
1522
+ assert all(
1523
+ isinstance(spec, SlidingWindowMLASpec)
1524
+ for group in swa_mla_specs
1525
+ for spec in group.kv_cache_specs.values()
1526
+ )
1527
+
1528
+ # Split each SWA UniformKV group into smaller groups to align their #(layer tuples)
1529
+ # Possibly padding layer tuples for this.
1530
+ # Additionally, we also pad KV blocks in each SWA layer, to align the page size
1531
+ # with the corresponding layer in the full-MLA group.
1532
+ all_page_sizes = full_mla_spec.get_page_sizes()
1533
+ swa_mla_groups = []
1534
+ for sm_spec in swa_mla_specs:
1535
+ sm_page_sizes = sm_spec.get_page_sizes()
1536
+ layers_per_size: dict[int, list[str]] = defaultdict(list)
1537
+ assert max(sm_page_sizes) <= max(all_page_sizes)
1538
+
1539
+ # Unify page size by padding layers' page_size to the nearest larger page_size.
1540
+ # Compute candidate (nearest larger page_size) for each unique page size.
1541
+ size_to_candidate: dict[int, int] = {}
1542
+ for ps in sm_page_sizes:
1543
+ size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps)
1544
+ # Pad and collect layer names per page size.
1545
+ for layer_name, layer_spec in sm_spec.kv_cache_specs.items():
1546
+ current_size = layer_spec.page_size_bytes
1547
+ candidate = size_to_candidate[current_size]
1548
+ if current_size < candidate:
1549
+ object.__setattr__(layer_spec, "page_size_padded", candidate)
1550
+ layers_per_size[candidate].append(layer_name)
1551
+ # NOTE(yifan): for now, inside a UniformKV group, each page_size should
1552
+ # have the same number of layers. This also means we don't need to pad layers
1553
+ # inside a partial-full layer tuple.
1554
+ assert len(set(len(layers) for layers in layers_per_size.values())) == 1
1555
+ num_layers_per_size = len(next(iter(layers_per_size.values())))
1556
+
1557
+ # Split layers inside each UniformKV group for aligned #(layers).
1558
+ # See `_get_kv_cache_groups_uniform_page_size` for more details.
1559
+ num_tuple_groups = cdiv(num_layers_per_size, num_layer_tuples)
1560
+ layer_tuples = list(zip(*layers_per_size.values()))
1561
+ for i in range(num_tuple_groups):
1562
+ group_layer_tuples = layer_tuples[i::num_tuple_groups]
1563
+ # Flatten tuples and build dict for from_specs
1564
+ group_layer_names = [
1565
+ name for layer_tuple in group_layer_tuples for name in layer_tuple
1566
+ ]
1567
+ group_layer_specs = {
1568
+ name: sm_spec.kv_cache_specs[name] for name in group_layer_names
1569
+ }
1570
+ sub_sm_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs)
1571
+ assert sub_sm_spec is not None
1572
+ swa_mla_groups.append(
1573
+ KVCacheGroupSpec(
1574
+ layer_names=group_layer_names,
1575
+ kv_cache_spec=sub_sm_spec,
1576
+ )
1577
+ )
1578
+
1579
+ return [full_mla_group, *swa_mla_groups]
1580
+
1581
+
1582
+def _annotate_eagle_groups_deepseek_v4(
1583
+ vllm_config: VllmConfig,
1584
+ kv_cache_spec: dict[str, KVCacheSpec],
1585
+ kv_cache_groups: list[KVCacheGroupSpec],
1586
+) -> None:
1587
+ spec_config = vllm_config.speculative_config
1588
+ if spec_config is None or not spec_config.use_eagle():
1589
+ return
1590
+ # Detection uses the merged MLA spec's model_version.
1591
+ if not any(
1592
+ getattr(spec, "model_version", None) == "deepseek_v4"
1593
+ for spec in kv_cache_spec.values()
1594
+ ):
1595
+ return
1596
+ # DeepseekV4's MTP attention layer is always the last layer, and we flag whichever
1597
+ # group contains it.
1598
+ # FIXME(yifan): avoid/generalize this hacky check.
1599
+ last_layer = next(reversed(kv_cache_spec))
1600
+ for group in kv_cache_groups:
1601
+ if last_layer in group.layer_names:
1602
+ group.is_eagle_group = True
1603
+ break
1604
+
1605
+
1606
def get_kv_cache_groups(
1607
vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec]
1608
) -> list[KVCacheGroupSpec]:
@@ -1268,6 +1634,14 @@ def get_kv_cache_groups(
1634
# full attention, or all layers are sliding window attention with the
1635
# same window size). Put all layers into one group.
1636
return _get_kv_cache_groups_uniform_type(uniform_spec)
1637
+ elif grouped_specs := group_and_unify_kv_cache_specs(kv_cache_spec):
1638
+ # DeepseekV4 case: All layers need the same number of token slots,
1639
+ # yet some layers are full attention while others are sliding window
1640
+ # attention in different sizes. Need to group layers into multiple
1641
+ # UniformTypeKVCacheSpecs.
1642
+ kv_cache_groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
1643
+ _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups)
1644
+ return kv_cache_groups
1645
1646
# As KVCacheManager can only allocate memory of one size, we need to unify
1647
# the page size of the layers. For cases cannot be unified, this function
@@ -1360,15 +1734,40 @@ def _max_memory_usage_bytes_from_groups(
1734
if not kv_cache_groups:
1735
return 0
1736
1363
- # UniformTypeKVCacheSpecs special case (single group, per-layer specs)
1737
if len(kv_cache_groups) == 1 and isinstance(
1738
kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
1739
):
1740
+ # UniformTypeKVCacheSpecs special case (single group, per-layer specs)
1741
per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs
1742
return sum(
1743
spec.max_memory_usage_bytes(vllm_config)
1744
for spec in per_layer_specs.values()
1745
)
1746
+ elif all(
1747
+ isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1748
+ for group in kv_cache_groups
1749
+ ):
1750
+ # Special case (only DeepseekV4 for now): all groups are
1751
+ # UniformTypeKVCacheSpecs.
1752
+ # They must already be page_size aligned and share a common padded
1753
+ # layer-tuple layout. Even groups with fewer actual tuples still reserve
1754
+ # the global number of tuple slots in the shared tensor layout.
1755
+ full_mla_spec = cast(UniformTypeKVCacheSpecs, kv_cache_groups[0].kv_cache_spec)
1756
+ layer_tuple_bytes = sum(full_mla_spec.get_page_sizes())
1757
+ num_layer_tuples = max(
1758
+ cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).get_num_layer_tuples()
1759
+ for group in kv_cache_groups
1760
+ )
1761
+
1762
+ total_max_mem_usage_bytes = 0
1763
+ for group in kv_cache_groups:
1764
+ group_spec = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec)
1765
+ g_max_mem_usage_pages = group_spec.max_memory_usage_pages(vllm_config)
1766
+ g_max_mem_usage_page_bytes = (
1767
+ num_layer_tuples * g_max_mem_usage_pages * layer_tuple_bytes
1768
+ )
1769
+ total_max_mem_usage_bytes += g_max_mem_usage_page_bytes
1770
+ return total_max_mem_usage_bytes
1771
1772
# General case: group_size pools, each shared by one layer per group
1773
# Memory = group_size * page_size * blocks_for_max_len
@@ -1515,7 +1914,13 @@ def _project_kv_cache_groups_to_worker(
1914
for layer_name in worker_layer_names
1915
},
1916
)
1518
- projected_groups.append(KVCacheGroupSpec(worker_layer_names, group_spec))
1917
+ projected_groups.append(
1918
+ KVCacheGroupSpec(
1919
+ worker_layer_names,
1920
+ group_spec,
1921
+ is_eagle_group=group.is_eagle_group and bool(worker_layer_names),
1922
+ )
1923
+ )
1924
return projected_groups
1925
1926
@@ -1698,10 +2103,7 @@ class BlockHashListWithBlockSize:
2103
def _get_value_at(self, idx: int) -> BlockHash:
2104
base = idx * self.scale_factor
2105
end = base + self.scale_factor
1701
- merged_hash: bytes = self.block_hashes[base]
1702
- for i in range(base + 1, end):
1703
- merged_hash += self.block_hashes[i]
1704
- return BlockHash(merged_hash)
2106
+ return BlockHash(b"".join(self.block_hashes[base:end]))
2107
2108
2109
BlockHashList = list[BlockHash] | BlockHashListWithBlockSize
vllm/v1/core/sched/interface.py
+1
@@ -41,6 +41,7 @@ class SchedulerInterface(ABC):
41
kv_cache_config: "KVCacheConfig",
42
structured_output_manager: "StructuredOutputManager",
43
block_size: int,
44
+ hash_block_size: int,
45
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
46
include_finished_set: bool = False,
47
log_stats: bool = False,
vllm/v1/core/sched/scheduler.py
+5
-2
@@ -71,6 +71,7 @@ class Scheduler(SchedulerInterface):
71
kv_cache_config: KVCacheConfig,
72
structured_output_manager: StructuredOutputManager,
73
block_size: int,
74
+ hash_block_size: int | None = None,
75
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
76
include_finished_set: bool = False,
77
log_stats: bool = False,
@@ -222,6 +223,8 @@ class Scheduler(SchedulerInterface):
223
self.num_lookahead_tokens = self.num_spec_tokens
224
225
# Create the KV cache manager.
226
+ if hash_block_size is None:
227
+ hash_block_size = block_size
228
self.kv_cache_manager = KVCacheManager(
229
kv_cache_config=kv_cache_config,
230
max_model_len=self.max_model_len,
@@ -231,7 +234,7 @@ class Scheduler(SchedulerInterface):
234
enable_kv_cache_events=self.enable_kv_cache_events,
235
dcp_world_size=self.dcp_world_size,
236
pcp_world_size=self.pcp_world_size,
234
- hash_block_size=self.block_size,
237
+ hash_block_size=hash_block_size,
238
metrics_collector=self.kv_metrics_collector,
239
)
240
# Bind GPU block pool to the KV connector. This must happen after
@@ -2018,7 +2021,7 @@ class Scheduler(SchedulerInterface):
2021
# the connector.
2022
self.kv_cache_manager.remove_skipped_blocks(
2023
request_id=request.request_id,
2021
- total_computed_tokens=request.num_tokens,
2024
+ total_computed_tokens=request.num_computed_tokens,
2025
)
2026
2027
block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
vllm/v1/core/single_type_kv_cache_manager.py
+6
-6
@@ -20,6 +20,7 @@ from vllm.v1.kv_cache_interface import (
20
MambaSpec,
21
MLAAttentionSpec,
22
SinkFullAttentionSpec,
23
+ SlidingWindowMLASpec,
24
SlidingWindowSpec,
25
TQFullAttentionSpec,
26
)
@@ -534,12 +535,10 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
535
):
536
# Skip prefix matching check if the block is not aligned with
537
# `alignment_tokens`.
537
- if (
538
- num_contiguous_blocks == 0
539
- and block_size != alignment_tokens # Faster for common case.
540
- and (i + 1) * block_size % alignment_tokens != 0
541
- ):
542
- continue
538
+ if num_contiguous_blocks == 0 and block_size != alignment_tokens:
539
+ post_pop_blocks = i if use_eagle else i + 1
540
+ if (post_pop_blocks * block_size) % alignment_tokens != 0:
541
+ continue
542
# Add the cached block to the computed blocks.
543
for computed, cached in zip(computed_blocks, cached_block):
544
computed[i] = cached
@@ -1118,6 +1117,7 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
1117
TQFullAttentionSpec: FullAttentionManager,
1118
MLAAttentionSpec: FullAttentionManager,
1119
SlidingWindowSpec: SlidingWindowManager,
1120
+ SlidingWindowMLASpec: SlidingWindowManager,
1121
ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,
1122
MambaSpec: MambaManager,
1123
CrossAttentionSpec: CrossAttentionManager,
vllm/v1/engine/core.py
+5
-5
@@ -45,6 +45,7 @@ from vllm.v1.core.kv_cache_utils import (
45
get_kv_cache_configs,
46
get_request_block_hasher,
47
init_none_hash,
48
+ resolve_kv_cache_block_sizes,
49
)
50
from vllm.v1.core.sched.interface import PauseState, SchedulerInterface
51
from vllm.v1.core.sched.output import SchedulerOutput
@@ -137,10 +138,8 @@ class EngineCore:
138
logger.warning("Disabling chunked prefill for model without KVCache")
139
vllm_config.scheduler_config.enable_chunked_prefill = False
140
140
- scheduler_block_size = (
141
- vllm_config.cache_config.block_size
142
- * vllm_config.parallel_config.decode_context_parallel_size
143
- * vllm_config.parallel_config.prefill_context_parallel_size
141
+ scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
142
+ kv_cache_config, vllm_config
143
)
144
145
self.scheduler: SchedulerInterface = Scheduler(
@@ -150,6 +149,7 @@ class EngineCore:
149
include_finished_set=include_finished_set,
150
log_stats=self.log_stats,
151
block_size=scheduler_block_size,
152
+ hash_block_size=hash_block_size,
153
)
154
self.use_spec_decode = vllm_config.speculative_config is not None
155
if self.scheduler.connector is not None: # type: ignore
@@ -207,7 +207,7 @@ class EngineCore:
207
init_none_hash(caching_hash_fn)
208
209
self.request_block_hasher = get_request_block_hasher(
210
- scheduler_block_size, caching_hash_fn
210
+ hash_block_size, caching_hash_fn
211
)
212
213
self.step_fn = (
vllm/v1/kv_cache_interface.py
+138
-8
@@ -4,6 +4,7 @@
4
from __future__ import annotations
5
6
import copy
7
+from collections import Counter
8
from dataclasses import dataclass, fields, replace
9
from enum import IntEnum
10
from math import prod
@@ -13,11 +14,11 @@ import torch
14
from typing_extensions import Self
15
16
from vllm.logger import init_logger
17
+from vllm.utils.math_utils import cdiv, round_up
18
+from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
19
20
if TYPE_CHECKING:
21
from vllm.config import VllmConfig
19
-from vllm.utils.math_utils import cdiv
20
-from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
22
23
logger = init_logger(__name__)
24
@@ -95,6 +96,10 @@ class KVCacheSpec:
96
"""
97
raise NotImplementedError
98
99
+ @property
100
+ def storage_block_size(self) -> int:
101
+ return self.block_size
102
+
103
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
104
"""
105
The maximum possible memory usage of this KV cache in bytes.
@@ -269,6 +274,15 @@ class FullAttentionSpec(AttentionSpec):
274
)
275
276
277
+def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec):
278
+ if spec.alignment is None:
279
+ return
280
+ actual_page_size = spec.real_page_size_bytes
281
+ padded_page_size = round_up(actual_page_size, spec.alignment)
282
+ if padded_page_size != actual_page_size:
283
+ object.__setattr__(spec, "page_size_padded", padded_page_size)
284
+
285
+
286
@dataclass(frozen=True, kw_only=True)
287
class TQFullAttentionSpec(FullAttentionSpec):
288
"""FullAttentionSpec with TQ-aware page size.
@@ -299,15 +313,31 @@ class TQFullAttentionSpec(FullAttentionSpec):
313
class MLAAttentionSpec(FullAttentionSpec):
314
# TODO(Lucas/Chen): less hacky way to do this
315
cache_dtype_str: str | None = None
316
+ # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults.
317
+ alignment: int | None = None # Default to None for no padding.
318
+ compress_ratio: int = 1 # Default to 1 for no compression.
319
+ model_version: str | None = None
320
+
321
+ def __post_init__(self):
322
+ super().__post_init__()
323
+ _apply_alignment_padding(self)
324
+
325
+ @property
326
+ def storage_block_size(self) -> int:
327
+ return self.block_size // self.compress_ratio
328
329
@property
330
def real_page_size_bytes(self) -> int:
331
if self.cache_dtype_str == "fp8_ds_mla":
306
- # See `vllm/v1/attention/backends/mla/flashmla_sparse.py`
307
- # for details.
332
+ if self.model_version == "deepseek_v4":
333
+ # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token.
334
+ # head_size stays semantic (512); bytes are determined here.
335
+ return self.storage_block_size * 584
336
+ # V3.2 main MLA: 656-byte custom layout (kv_lora_rank=512 +
337
+ # qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py.
338
return self.block_size * 656
339
return (
310
- self.block_size
340
+ self.storage_block_size
341
* self.num_kv_heads
342
* self.head_size
343
* get_dtype_size(self.dtype)
@@ -319,9 +349,15 @@ class MLAAttentionSpec(FullAttentionSpec):
349
"All attention layers in the same KV cache group must be MLAAttentionSpec."
350
)
351
cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
322
- assert len(cache_dtype_str_set) == 1, (
352
+ compress_ratio_set = set(spec.compress_ratio for spec in specs)
353
+ model_version_set = set(spec.model_version for spec in specs)
354
+ assert (
355
+ len(cache_dtype_str_set) == 1
356
+ and len(compress_ratio_set) == 1
357
+ and len(model_version_set) == 1
358
+ ), (
359
"All attention layers in the same KV cache group must use the same "
324
- "quantization method."
360
+ "quantization method, compress ratio, and model version."
361
)
362
return cls(
363
block_size=specs[0].block_size,
@@ -331,6 +367,8 @@ class MLAAttentionSpec(FullAttentionSpec):
367
kv_quant_mode=specs[0].kv_quant_mode,
368
page_size_padded=specs[0].page_size_padded,
369
cache_dtype_str=cache_dtype_str_set.pop(),
370
+ compress_ratio=compress_ratio_set.pop(),
371
+ model_version=model_version_set.pop(),
372
)
373
374
@@ -393,6 +431,71 @@ class SlidingWindowSpec(AttentionSpec):
431
return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes
432
433
434
+@dataclass(frozen=True, kw_only=True)
435
+class SlidingWindowMLASpec(SlidingWindowSpec):
436
+ """Sliding window attention with MLA cache format."""
437
+
438
+ cache_dtype_str: str | None = None
439
+ # DeepseekV4-only: see MLAAttentionSpec.model_version.
440
+ alignment: int | None = None # Default to None for no padding.
441
+ compress_ratio: int = 1
442
+ model_version: str | None = None
443
+
444
+ def __post_init__(self):
445
+ _apply_alignment_padding(self)
446
+
447
+ @property
448
+ def storage_block_size(self) -> int:
449
+ return self.block_size // self.compress_ratio
450
+
451
+ @property
452
+ def real_page_size_bytes(self) -> int:
453
+ if self.model_version == "deepseek_v4":
454
+ # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token.
455
+ return self.storage_block_size * 584
456
+ assert self.model_version is None, (
457
+ f"Unsupported model version: {self.model_version}"
458
+ )
459
+ return (
460
+ self.storage_block_size
461
+ * self.num_kv_heads
462
+ * self.head_size
463
+ * get_dtype_size(self.dtype)
464
+ )
465
+
466
+ @classmethod
467
+ def merge(cls, specs: list[Self]) -> Self:
468
+ assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), (
469
+ "All attention layers in the same KV cache group must be "
470
+ "SlidingWindowMLASpec."
471
+ )
472
+ cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
473
+ compress_ratio_set = set(spec.compress_ratio for spec in specs)
474
+ model_version_set = set(spec.model_version for spec in specs)
475
+ sliding_window_set = set(spec.sliding_window for spec in specs)
476
+ assert (
477
+ len(cache_dtype_str_set) == 1
478
+ and len(compress_ratio_set) == 1
479
+ and len(model_version_set) == 1
480
+ and len(sliding_window_set) == 1
481
+ ), (
482
+ "All attention layers in the same KV cache group must use the same "
483
+ "quantization method, compress ratio, model version and sliding "
484
+ "window size."
485
+ )
486
+ return cls(
487
+ block_size=specs[0].block_size,
488
+ num_kv_heads=specs[0].num_kv_heads,
489
+ head_size=specs[0].head_size,
490
+ dtype=specs[0].dtype,
491
+ page_size_padded=specs[0].page_size_padded,
492
+ sliding_window=sliding_window_set.pop(),
493
+ cache_dtype_str=cache_dtype_str_set.pop(),
494
+ compress_ratio=compress_ratio_set.pop(),
495
+ model_version=model_version_set.pop(),
496
+ )
497
+
498
+
499
@dataclass(frozen=True)
500
class MambaSpec(KVCacheSpec):
501
shapes: tuple[tuple[int, ...], ...]
@@ -527,7 +630,17 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
630
# Different block sizes, not uniform.
631
return False
632
one_spec = next(iter(kv_cache_specs.values()))
530
- if isinstance(one_spec, FullAttentionSpec):
633
+ # NOTE: Check subclasses before parent classes since isinstance()
634
+ # returns True for subclasses.
635
+ if isinstance(one_spec, SlidingWindowMLASpec):
636
+ # SlidingWindowMLASpec is uniform if all specs are SlidingWindowMLASpec
637
+ # with the same sliding_window size.
638
+ return all(
639
+ isinstance(spec, SlidingWindowMLASpec)
640
+ and spec.sliding_window == one_spec.sliding_window
641
+ for spec in kv_cache_specs.values()
642
+ )
643
+ elif isinstance(one_spec, FullAttentionSpec):
644
return all(
645
isinstance(spec, FullAttentionSpec) for spec in kv_cache_specs.values()
646
)
@@ -571,6 +684,21 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
684
else:
685
return None
686
687
+ # NOTE: below util functions are only used by DeepseekV4 for now.
688
+ def get_page_sizes(self) -> list[int]:
689
+ return list(set(spec.page_size_bytes for spec in self.kv_cache_specs.values()))
690
+
691
+ def get_num_layer_tuples(self) -> int:
692
+ return Counter(
693
+ spec.page_size_bytes for spec in self.kv_cache_specs.values()
694
+ ).most_common(1)[0][1]
695
+
696
+ def max_memory_usage_pages(self, vllm_config: VllmConfig) -> int:
697
+ return max(
698
+ cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
699
+ for spec in self.kv_cache_specs.values()
700
+ )
701
+
702
703
@dataclass
704
class KVCacheTensor:
@@ -593,6 +721,8 @@ class KVCacheGroupSpec:
721
layer_names: list[str]
722
# The KV cache spec of this manager layer
723
kv_cache_spec: KVCacheSpec
724
+ # Whether this group contains EAGLE/MTP draft attention layers.
725
+ is_eagle_group: bool = False
726
727
728
@dataclass
vllm/v1/spec_decode/llm_base_proposer.py
+26
-3
@@ -84,6 +84,15 @@ class SpecDecodeBaseProposer:
84
self.hidden_size = self.draft_model_config.get_hidden_size()
85
self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()
86
87
+ # DeepSeek V4 MTP consumes the target's pre-hc_head residual stream,
88
+ # shape (T, hc_mult * hidden_size). Expand the hidden_states buffer
89
+ # so target_hidden_states fits; detect DeepseekV4 via draft hf_config.
90
+ draft_hf_config = self.draft_model_config.hf_config
91
+ if hasattr(draft_hf_config, "compress_ratios") and hasattr(
92
+ draft_hf_config, "hc_mult"
93
+ ):
94
+ self.hidden_size = self.hidden_size * draft_hf_config.hc_mult
95
+
96
# Unifying eagle, draft model, and parallel drafting support.
97
# DFlash always uses parallel drafting (all tokens in one pass),
98
# but has an additional slot for the next_token_id (does not shift like EAGLE)
@@ -1308,9 +1317,12 @@ class SpecDecodeBaseProposer:
1317
self.vllm_config,
1318
AttentionLayerBase, # type: ignore[type-abstract]
1319
)
1311
- self._draft_attn_layer_names = (
1312
- set(all_attn_layers.keys()) - target_attn_layer_names
1313
- )
1320
+ # Filter to only layers that have KV cache specs.
1321
+ self._draft_attn_layer_names = {
1322
+ name
1323
+ for name in (set(all_attn_layers.keys()) - target_attn_layer_names)
1324
+ if all_attn_layers[name].get_kv_cache_spec(self.vllm_config) is not None
1325
+ }
1326
1327
if self.supports_mm_inputs:
1328
# Even if the target model is multimodal, we can also use
@@ -1514,6 +1526,17 @@ class SpecDecodeBaseProposer:
1526
"Shared target model lm_head with MTP shared_head.head."
1527
)
1528
1529
+ if hasattr(target_language_model.model, "topk_indices_buffer"):
1530
+ if hasattr(self.model.model, "topk_indices_buffer"):
1531
+ del self.model.model.topk_indices_buffer
1532
+ self.model.model.topk_indices_buffer = (
1533
+ target_language_model.model.topk_indices_buffer
1534
+ )
1535
+ logger.info(
1536
+ "Detected MTP model with topk_indices_buffer. "
1537
+ "Sharing target model topk_indices_buffer with the draft model."
1538
+ )
1539
+
1540
if self.use_local_argmax_reduction:
1541
if not hasattr(self.model, "get_top_tokens"):
1542
raise ValueError(
vllm/v1/worker/gpu/attn_utils.py
+26
-3
@@ -9,6 +9,7 @@ import torch
9
10
from vllm.config import VllmConfig, get_layers_from_vllm_config
11
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
12
+from vllm.utils.torch_utils import get_dtype_size
13
from vllm.v1.attention.backend import (
14
AttentionBackend,
15
AttentionCGSupport,
@@ -162,7 +163,7 @@ def _reshape_kv_cache(
163
attn_backend = attn_backends[layer_name]
164
kv_cache_shape = attn_backend.get_kv_cache_shape(
165
num_blocks,
165
- kv_cache_spec.block_size,
166
+ kv_cache_spec.storage_block_size,
167
kv_cache_spec.num_kv_heads,
168
kv_cache_spec.head_size,
169
cache_dtype,
@@ -183,8 +184,28 @@ def _reshape_kv_cache(
184
185
dtype = kv_cache_spec.dtype
186
raw_tensor = raw_tensor.view(dtype)
186
- raw_tensor = raw_tensor.view(kv_cache_shape)
187
- kv_caches[layer_name] = raw_tensor.permute(*inv_order)
187
+ if kv_cache_spec.page_size_padded is not None:
188
+ # Use strided view to handle page_size_bytes that
189
+ # include padding. This follows the same pattern as
190
+ # MambaSpec handling in gpu_model_runner.py.
191
+ # NOTE: This assumes kv_cache_shape[0] == num_blocks
192
+ # (i.e. the first physical dimension is the block
193
+ # index), which holds for MLA backends but NOT for
194
+ # standard attention backends whose shape starts with
195
+ # a K/V dimension of size 2.
196
+ dtype_size = get_dtype_size(dtype)
197
+ page_stride = kv_cache_spec.page_size_bytes // dtype_size
198
+ strides = list(torch.empty(kv_cache_shape).stride())
199
+ strides[inv_order[0]] = page_stride
200
+ kv_cache = torch.as_strided(
201
+ raw_tensor,
202
+ size=kv_cache_shape,
203
+ stride=tuple(strides),
204
+ )
205
+ else:
206
+ # No padding — safe to use a contiguous view.
207
+ kv_cache = raw_tensor.view(kv_cache_shape)
208
+ kv_caches[layer_name] = kv_cache.permute(*inv_order)
209
return kv_caches
210
211
@@ -230,6 +251,7 @@ def build_attn_metadata(
251
seq_lens_cpu_upper_bound: torch.Tensor | None = None,
252
dcp_local_seq_lens: torch.Tensor | None = None,
253
encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]] | None = None,
254
+ positions: torch.Tensor | None = None,
255
) -> dict[str, Any]:
256
seq_lens = seq_lens[:num_reqs]
257
if dcp_local_seq_lens is not None:
@@ -256,6 +278,7 @@ def build_attn_metadata(
278
slot_mapping=slot_mapping,
279
causal=True,
280
dcp_local_seq_lens=dcp_local_seq_lens,
281
+ positions=positions,
282
)
283
if encoder_seq_lens and i in encoder_seq_lens:
284
encoder_seq_lens_gpu, encoder_seq_lens_cpu = encoder_seq_lens[i]
vllm/v1/worker/gpu/model_runner.py
+19
-3
@@ -486,11 +486,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
486
device=self.device,
487
),
488
)
489
+
490
+ # Let the target override the hidden state fed to the drafter
491
+ # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
492
+ # target returns a persistent buffer sized at max_num_batched_tokens;
493
+ # slice to the active token count that propose() expects.
494
+ spec_hidden_states = hidden_states
495
+ if hasattr(self.model, "get_mtp_target_hidden_states"):
496
+ pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
497
+ spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]
498
self.speculator.propose(
499
input_batch=input_batch,
500
attn_metadata=attn_metadata,
501
slot_mappings=slot_mappings_by_layer,
493
- last_hidden_states=hidden_states,
502
+ last_hidden_states=spec_hidden_states,
503
aux_hidden_states=aux_hidden_states,
504
num_sampled=torch.ones(
505
input_batch.num_reqs, dtype=torch.int32, device=self.device
@@ -808,7 +817,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
817
out=seq_lens_cpu_upper_bound_np[:num_reqs],
818
)
819
seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)
811
-
820
return InputBatch(
821
req_ids=req_ids,
822
num_reqs=num_reqs,
@@ -1233,11 +1241,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1241
1242
if self.speculator is not None:
1243
assert self.sampler is not None
1244
+ # Let the target override the hidden state fed to the drafter
1245
+ # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
1246
+ # target returns a persistent buffer sized at max_num_batched_tokens;
1247
+ # slice to the active token count that propose() expects.
1248
+ spec_hidden_states = hidden_states
1249
+ if hasattr(self.model, "get_mtp_target_hidden_states"):
1250
+ pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
1251
+ spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]
1252
draft_tokens = self.speculator.propose(
1253
input_batch,
1254
attn_metadata,
1255
slot_mappings_by_layer,
1240
- hidden_states,
1256
+ spec_hidden_states,
1257
aux_hidden_states,
1258
num_sampled,
1259
num_rejected,
vllm/v1/worker/gpu/model_states/default.py
+1
@@ -193,5 +193,6 @@ class DefaultModelState(ModelState):
193
kv_cache_config=kv_cache_config,
194
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
195
dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
196
+ positions=input_batch.positions,
197
)
198
return attn_metadata
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+5
@@ -53,6 +53,11 @@ class EagleSpeculator:
53
# the draft model's hidden size can be different from the target model's
54
# hidden size (e.g., Llama 3.3 70B).
55
self.hidden_size = self.draft_model_config.get_hidden_size()
56
+ # Widen for HC-multiplexed residuals (e.g. DeepSeek V4 feeds the MTP
57
+ # draft the target's pre-hc_head (T, hc_mult * hidden_size) residual).
58
+ # Non-HC models default to hc_mult=1 and are unaffected.
59
+ hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1)
60
+ self.hidden_size = self.hidden_size * hc_mult
61
self.vocab_size = self.draft_model_config.get_vocab_size()
62
self.dtype = vllm_config.model_config.dtype
63
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+15
@@ -49,4 +49,19 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
49
del eagle_model.lm_head
50
eagle_model.lm_head = target_model.lm_head
51
52
+ # MTP models call compute_logits via shared_head.head (a
53
+ # ParallelLMHead inside each MTP layer), not self.model.lm_head.
54
+ # If the checkpoint omits a copy of the lm_head weights at the
55
+ # MTP layer path, shared_head.head stays uninitialised and
56
+ # produces zero/NaN logits. Share it explicitly from the target.
57
+ inner = getattr(eagle_model, "model", None)
58
+ layers = getattr(inner, "layers", None) if inner is not None else None
59
+ if layers is not None:
60
+ items = layers.values() if isinstance(layers, nn.ModuleDict) else layers
61
+ for layer in items:
62
+ sh = getattr(layer, "shared_head", None)
63
+ if sh is not None and hasattr(sh, "head"):
64
+ del sh.head
65
+ sh.head = target_model.lm_head
66
+
67
return eagle_model
vllm/v1/worker/gpu/states.py
+1
@@ -104,6 +104,7 @@ class RequestState:
104
self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
105
self.num_computed_tokens_np[req_idx] = num_computed_tokens
106
self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)
107
+ self.num_computed_tokens_np[req_idx] = num_computed_tokens
108
109
if num_computed_tokens > 0 and num_computed_tokens <= prefill_len:
110
# For PD disagg or resumed requests: set last_sampled to the last
vllm/v1/worker/gpu_model_runner.py
+43
-7
@@ -2184,6 +2184,7 @@ class GPUModelRunner(
2184
slot_mapping=slot_mapping_gid_0,
2185
causal=True,
2186
is_prefilling=is_prefilling,
2187
+ positions=self.positions[:num_tokens_padded],
2188
)
2189
2190
if self.dcp_world_size > 1:
@@ -4671,6 +4672,16 @@ class GPUModelRunner(
4672
next_token_ids, valid_sampled_tokens_count
4673
)
4674
4675
+ # Let the target override the hidden state fed to the drafter
4676
+ # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). Safe to
4677
+ # rebind here: hidden_states was already consumed for sampling
4678
+ # above and is not used again in this branch.
4679
+ alt = getattr(
4680
+ self.get_model(), "get_mtp_target_hidden_states", lambda: None
4681
+ )()
4682
+ if alt is not None:
4683
+ hidden_states = alt
4684
+
4685
num_rejected_tokens_gpu = None
4686
if spec_decode_metadata is None:
4687
token_indices_to_sample = None
@@ -6587,9 +6598,15 @@ class GPUModelRunner(
6598
)
6599
kernel_num_blocks = num_blocks * num_blocks_per_kv_block
6600
6601
+ # For MLA with compression, storage_block_size != block_size
6602
+ if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
6603
+ shape_block_size = kv_cache_spec.storage_block_size
6604
+ else:
6605
+ shape_block_size = kernel_block_size
6606
+
6607
kv_cache_shape = attn_backend.get_kv_cache_shape(
6608
kernel_num_blocks,
6592
- kernel_block_size,
6609
+ shape_block_size,
6610
kv_cache_spec.num_kv_heads,
6611
kv_cache_spec.head_size,
6612
cache_dtype_str=self.cache_config.cache_dtype,
@@ -6613,12 +6630,31 @@ class GPUModelRunner(
6630
kv_cache_stride_order.index(i)
6631
for i in range(len(kv_cache_stride_order))
6632
]
6616
- kv_caches[layer_name] = (
6617
- kv_cache_raw_tensors[layer_name]
6618
- .view(dtype)
6619
- .view(kv_cache_shape)
6620
- .permute(*inv_order)
6621
- )
6633
+
6634
+ raw_tensor = kv_cache_raw_tensors[layer_name].view(dtype)
6635
+ if kv_cache_spec.page_size_padded is not None:
6636
+ # Use strided view to handle page_size_bytes that
6637
+ # include padding. This follows
6638
+ # the same pattern as MambaSpec handling below.
6639
+ # NOTE: This assumes kv_cache_shape[0] == num_blocks
6640
+ # (i.e. the first physical dimension is the block
6641
+ # index), which holds for MLA backends but NOT for
6642
+ # standard attention backends whose shape starts with
6643
+ # a K/V dimension of size 2.
6644
+ dtype_size = get_dtype_size(dtype)
6645
+ page_stride = kv_cache_spec.page_size_bytes // dtype_size
6646
+ strides = list(torch.empty(kv_cache_shape).stride())
6647
+ strides[inv_order[0]] = page_stride
6648
+ kv_cache = torch.as_strided(
6649
+ raw_tensor,
6650
+ size=kv_cache_shape,
6651
+ stride=tuple(strides),
6652
+ )
6653
+ else:
6654
+ # No padding — safe to use a contiguous view.
6655
+ kv_cache = raw_tensor.view(kv_cache_shape)
6656
+ kv_caches[layer_name] = kv_cache.permute(*inv_order)
6657
+
6658
elif isinstance(kv_cache_spec, MambaSpec):
6659
has_mamba = True
6660
raw_tensor = kv_cache_raw_tensors[layer_name]