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, 12:30 UTC
[Bugfix] fix inductor error for dpsk v4 (#41135)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Apr 29, 2026 at 12:18 UTC
2ae73c758ceed55ad2f70a69b47c8a994fce5662
1 file changed
+106
-36
vllm/v1/attention/ops/deepseek_v4_ops/fused_inv_rope_fp8_quant.py
+106
-36
@@ -10,6 +10,7 @@ INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layo
10
import torch
11
12
from vllm.triton_utils import tl, triton
13
+from vllm.utils.torch_utils import direct_register_custom_op
14
15
16
@triton.jit
@@ -180,34 +181,74 @@ def fused_inv_rope_fp8_quant(
181
fp8_dtype = torch.float8_e4m3fn
182
fp8_max = torch.finfo(fp8_dtype).max
183
183
- fp8_buf = torch.empty(
184
- (n_groups, num_tokens, d),
185
- dtype=fp8_dtype,
186
- device=o.device,
187
- )
188
-
184
tma_aligned_T = get_tma_aligned_size(num_tokens, 4)
185
if tma_aligned_scales:
186
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
- )
187
+ scale_inner = packed_sf_k
188
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
- )
189
+ scale_inner = num_scale_blocks
190
+
191
+ # Run kernel through a custom op so inductor sees an opaque boundary.
192
+ # It's a pytorch bug, see https://github.com/vllm-project/vllm/issues/41106
193
+ fp8_buf, scale_buf = torch.ops.vllm.fused_inv_rope_fp8_quant_kernel(
194
+ o,
195
+ positions,
196
+ cos_sin_cache,
197
+ heads_per_group,
198
+ quant_group_size,
199
+ chunks_per_head,
200
+ nope_dim % quant_group_size,
201
+ rope_dim // 2,
202
+ tma_aligned_scales,
203
+ fp8_max,
204
+ tma_aligned_T,
205
+ num_tokens,
206
+ n_groups,
207
+ d,
208
+ scale_inner,
209
+ )
210
+ return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1)
211
+
212
210
- common_args = dict(
213
+def _fused_inv_rope_fp8_quant_kernel_impl(
214
+ o: torch.Tensor,
215
+ positions: torch.Tensor,
216
+ cos_sin_cache: torch.Tensor,
217
+ heads_per_group: int,
218
+ quant_group_size: int,
219
+ chunks_per_head: int,
220
+ rope_start: int,
221
+ half_rope: int,
222
+ tma_aligned_scales: bool,
223
+ fp8_max: float,
224
+ tma_aligned_T: int,
225
+ num_tokens: int,
226
+ n_groups: int,
227
+ d: int,
228
+ scale_inner: int,
229
+) -> tuple[torch.Tensor, torch.Tensor]:
230
+ fp8_buf = torch.empty(
231
+ (n_groups, num_tokens, d),
232
+ dtype=torch.float8_e4m3fn,
233
+ device=o.device,
234
+ )
235
+ scale_dtype = torch.int32 if tma_aligned_scales else torch.float32
236
+ scale_buf = torch.empty(
237
+ n_groups * scale_inner * tma_aligned_T,
238
+ dtype=scale_dtype,
239
+ device=o.device,
240
+ ).as_strided(
241
+ (n_groups, num_tokens, scale_inner),
242
+ (scale_inner * tma_aligned_T, 1, tma_aligned_T),
243
+ )
244
+ grid = (tma_aligned_T, n_groups * heads_per_group)
245
+ _fused_inv_rope_fp8_quant_per_head[grid](
246
+ o,
247
+ positions,
248
+ cos_sin_cache,
249
+ fp8_buf,
250
+ scale_buf,
251
+ num_tokens,
252
heads_per_group=heads_per_group,
253
o_stride_token=o.stride(0),
254
o_stride_head=o.stride(1),
@@ -220,23 +261,52 @@ def fused_inv_rope_fp8_quant(
261
eps=1e-10,
262
QUANT_GROUP_SIZE=quant_group_size,
263
CHUNKS_PER_HEAD=chunks_per_head,
223
- ROPE_START=nope_dim % quant_group_size,
224
- HALF_ROPE=rope_dim // 2,
264
+ ROPE_START=rope_start,
265
+ HALF_ROPE=half_rope,
266
TMA_ALIGNED_SCALES=tma_aligned_scales,
267
num_stages=1,
268
launch_pdl=False,
269
+ num_warps=1,
270
)
271
+ return fp8_buf, scale_buf
272
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,
273
+
274
+def _fused_inv_rope_fp8_quant_kernel_fake(
275
+ o: torch.Tensor,
276
+ positions: torch.Tensor,
277
+ cos_sin_cache: torch.Tensor,
278
+ heads_per_group: int,
279
+ quant_group_size: int,
280
+ chunks_per_head: int,
281
+ rope_start: int,
282
+ half_rope: int,
283
+ tma_aligned_scales: bool,
284
+ fp8_max: float,
285
+ tma_aligned_T: int,
286
+ num_tokens: int,
287
+ n_groups: int,
288
+ d: int,
289
+ scale_inner: int,
290
+) -> tuple[torch.Tensor, torch.Tensor]:
291
+ fp8_buf = torch.empty(
292
+ (n_groups, num_tokens, d),
293
+ dtype=torch.float8_e4m3fn,
294
+ device=o.device,
295
)
296
+ scale_dtype = torch.int32 if tma_aligned_scales else torch.float32
297
+ scale_buf = torch.empty(
298
+ n_groups * scale_inner * tma_aligned_T,
299
+ dtype=scale_dtype,
300
+ device=o.device,
301
+ ).as_strided(
302
+ (n_groups, num_tokens, scale_inner),
303
+ (scale_inner * tma_aligned_T, 1, tma_aligned_T),
304
+ )
305
+ return fp8_buf, scale_buf
306
242
- return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1)
307
+
308
+direct_register_custom_op(
309
+ op_name="fused_inv_rope_fp8_quant_kernel",
310
+ op_func=_fused_inv_rope_fp8_quant_kernel_impl,
311
+ fake_impl=_fused_inv_rope_fp8_quant_kernel_fake,
312
+)