@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 731 lines 27.9 KB
Raw
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
6
7 import torch
8
9 from vllm.model_executor.warmup.jit_warmup import kernel_launcher
10 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
11 LaunchSpec,
12 TritonWarmupTensor,
13 VllmTritonJitKernel,
14 )
15 from vllm.platforms import current_platform
16 from vllm.triton_utils import tl, triton
17 from vllm.utils.import_utils import has_cutedsl
18
19 # MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale.
20 MXFP4_BLOCK_SIZE = 32
21
22
23 @triton.jit
24 def _get_cos_sin(
25 cos_sin_cache_ptr,
26 cos_sin_cache_stride,
27 pos,
28 HALF_ROT_DIM: tl.constexpr,
29 ):
30 block = tl.arange(0, HALF_ROT_DIM)
31 cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block)
32 cos = cos.to(tl.float32)
33 sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM)
34 sin = sin.to(tl.float32)
35 return cos, sin
36
37
38 @triton.jit
39 def _fp32x2_to_fp4x2(x_lo, x_hi):
40 # NOTE: $1 is high nibble, $2 is low nibble
41 return tl.inline_asm_elementwise(
42 """
43 {
44 .reg .b8 tmp;
45 cvt.rn.satfinite.e2m1x2.f32 tmp, $1, $2;
46 cvt.u32.u8 $0, tmp;
47 }
48 """,
49 constraints="=r,f,f",
50 args=[x_hi, x_lo],
51 dtype=tl.uint32,
52 is_pure=True,
53 pack=1,
54 ).to(tl.uint8)
55
56
57 @triton.jit
58 def _quantize_mxfp4_pair(x_lo, x_hi):
59 """Quantize a block of MXFP4_BLOCK_SIZE fp32 values given as two
60 interleaved halves (x_lo = values at even positions in the block,
61 x_hi = values at odd positions). Returns:
62 - packed : uint8[BLOCK/2] (low nibble = quant(x_lo), high = quant(x_hi))
63 - ue8m0 : scalar uint8 (block scale = 2^(ue8m0 - 127))
64 """
65 amax = tl.maximum(tl.max(tl.abs(x_lo)), tl.max(tl.abs(x_hi)))
66 # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
67 amax = tl.maximum(amax, 6.0 * (2**-126))
68 # ue8m0 block scale: 2^ceil(log2(amax/6.0)).
69 log2_ratio = tl.math.ceil(tl.math.log2(amax * (1.0 / 6.0)))
70 log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
71 scale = tl.math.exp2(log2_ratio)
72 ue8m0 = (log2_ratio + 127.0).to(tl.uint8)
73
74 inv_scale = 1.0 / scale
75 packed = _fp32x2_to_fp4x2(x_lo * inv_scale, x_hi * inv_scale)
76 return packed, ue8m0
77
78
79 class FusedIndexerQRopeQuantTritonKernel(
80 VllmTritonJitKernel["FusedIndexerQRopeQuantTritonKernel.CompileKey"]
81 ):
82 @dataclass(frozen=True)
83 class CompileKey:
84 dtype: torch.dtype
85 num_heads: int
86 index_q_half_rot_dim: int
87 index_q_head_dim: int
88 fp8_max: float
89 use_fnuz: bool
90 use_explicit_fma: bool
91
92 @staticmethod
93 @triton.jit
94 def kernel(
95 pos_ptr,
96 # Index Q RoPE
97 index_q_ptr,
98 index_q_stride0,
99 index_q_stride1,
100 index_q_cos_sin_ptr,
101 index_q_cos_sin_stride,
102 INDEX_Q_HALF_ROT_DIM: tl.constexpr,
103 # Index Q Quantize
104 index_q_fp8_ptr,
105 index_q_fp8_stride0,
106 index_q_fp8_stride1,
107 INDEX_Q_HEAD_DIM: tl.constexpr,
108 # Index weights
109 index_weights_ptr,
110 index_weights_stride,
111 index_weights_softmax_scale,
112 index_weights_head_scale,
113 index_weights_out_ptr,
114 index_weights_out_stride,
115 FP8_MAX: tl.constexpr = 448.0,
116 USE_FNUZ: tl.constexpr = False,
117 USE_EXPLICIT_FMA: tl.constexpr = False,
118 ):
119 # Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding
120 # + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the
121 # LAST rope_dim dims of each head; the leading [0, NOPE_DIM) is passed
122 # through unchanged.
123 INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM
124 INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM
125 tl.static_assert(INDEX_Q_NOPE_DIM >= 0)
126
127 tok_idx = tl.program_id(0)
128 head_idx = tl.program_id(1)
129
130 pos = tl.load(pos_ptr + tok_idx)
131 cos, sin = _get_cos_sin(
132 index_q_cos_sin_ptr,
133 index_q_cos_sin_stride,
134 pos,
135 INDEX_Q_HALF_ROT_DIM,
136 )
137 half_offset = tl.arange(0, INDEX_Q_HALF_ROT_DIM)
138 base_ptr = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
139
140 # Interleaved (GPT-J) RoPE on dims [NOPE_DIM, HEAD_DIM):
141 # even = q[NOPE_DIM + 2*i], odd = q[NOPE_DIM + 2*i + 1]
142 rot_base = base_ptr + INDEX_Q_NOPE_DIM
143 x_even = tl.load(rot_base + half_offset * 2).to(tl.float32)
144 x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32)
145 if USE_EXPLICIT_FMA:
146 # Match HIP rotary_embedding contraction before bf16 materialization.
147 r_even = tl.fma(x_even, cos, -(x_odd * sin))
148 r_odd = tl.fma(x_odd, cos, x_even * sin)
149 else:
150 r_even = x_even * cos - x_odd * sin
151 r_odd = x_odd * cos + x_even * sin
152
153 # Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax.
154 # Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py).
155 r_even = r_even.to(tl.bfloat16).to(tl.float32)
156 r_odd = r_odd.to(tl.bfloat16).to(tl.float32)
157
158 amax = tl.maximum(tl.max(tl.abs(r_even)), tl.max(tl.abs(r_odd)))
159 if INDEX_Q_NOPE_DIM > 0:
160 nope_offset = tl.arange(0, INDEX_Q_NOPE_DIM)
161 x_nope = tl.load(base_ptr + nope_offset).to(tl.float32)
162 amax = tl.maximum(amax, tl.max(tl.abs(x_nope)))
163 index_q_scale = tl.div_rn(tl.maximum(amax, 1e-4), FP8_MAX)
164 index_q_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(index_q_scale)))
165
166 # Store quantized values to index_q_fp8. FNUZ (e4m3fnuz) on gfx942, OCP
167 # (e4m3fn) elsewhere -- matches the K cache.
168 fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
169 fp8_base_ptr = (
170 index_q_fp8_ptr
171 + tok_idx * index_q_fp8_stride0
172 + head_idx * index_q_fp8_stride1
173 )
174 if INDEX_Q_NOPE_DIM > 0:
175 tl.store(
176 fp8_base_ptr + nope_offset,
177 tl.div_rn(x_nope, index_q_scale).to(fp8_dtype),
178 )
179 fp8_rot_base = fp8_base_ptr + INDEX_Q_NOPE_DIM
180 tl.store(
181 fp8_rot_base + half_offset * 2,
182 tl.div_rn(r_even, index_q_scale).to(fp8_dtype),
183 )
184 tl.store(
185 fp8_rot_base + half_offset * 2 + 1,
186 tl.div_rn(r_odd, index_q_scale).to(fp8_dtype),
187 )
188
189 # FP8 weight-fold contract:
190 # index_weights_out = index_weights * q_scale * softmax_scale * head_scale
191 # The per-token-per-head q_scale (fp32) IS folded into the output weights
192 # here because FP8 Q is stored WITHOUT a companion scale tensor — the
193 # downstream fp8_fp4_mqa_logits/fp8_fp4_paged_mqa_logits kernels use
194 # `weights` to apply per-token Q scale inline. See the MXFP4 kernel below for
195 # contrasting convention (scales live with the Q values, weights are NOT
196 # q-scaled).
197 index_weights = tl.load(
198 index_weights_ptr + tok_idx * index_weights_stride + head_idx
199 )
200 index_weights = index_weights.to(tl.float32)
201 index_weights *= index_q_scale
202 index_weights *= index_weights_softmax_scale
203 index_weights *= index_weights_head_scale
204 tl.store(
205 index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
206 index_weights,
207 )
208
209 def dispatch( # type: ignore[override]
210 self,
211 *,
212 dtype: torch.dtype,
213 num_heads: int,
214 head_dim: int,
215 rope_dim: int,
216 use_fnuz: bool,
217 ) -> CompileKey:
218 return self.CompileKey(
219 dtype=dtype,
220 num_heads=num_heads,
221 index_q_half_rot_dim=rope_dim // 2,
222 index_q_head_dim=head_dim,
223 fp8_max=224.0 if use_fnuz else 448.0,
224 use_fnuz=use_fnuz,
225 use_explicit_fma=current_platform.is_rocm(),
226 )
227
228 def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
229 hf_config = vllm_config.model_config.hf_config
230 num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0)
231 head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
232 rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
233 if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
234 return []
235
236 return self._trace_dispatch(self.dispatch)(
237 dtype=vllm_config.model_config.dtype,
238 num_heads=num_heads,
239 head_dim=head_dim,
240 rope_dim=rope_dim,
241 use_fnuz=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,
242 )
243
244 def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
245 q_stride0 = compile_key.num_heads * compile_key.index_q_head_dim
246 return dict(
247 positions=TritonWarmupTensor(torch.int64),
248 index_q=TritonWarmupTensor(
249 compile_key.dtype,
250 shape=(1, compile_key.num_heads, compile_key.index_q_head_dim),
251 strides=(q_stride0, compile_key.index_q_head_dim, 1),
252 ),
253 index_q_cos_sin_cache=TritonWarmupTensor(
254 torch.float32,
255 shape=(1, compile_key.index_q_half_rot_dim * 2),
256 ),
257 index_weights=TritonWarmupTensor(
258 compile_key.dtype,
259 shape=(1, compile_key.num_heads),
260 ),
261 index_weights_softmax_scale=1.0,
262 index_weights_head_scale=1.0,
263 index_q_fp8=TritonWarmupTensor(
264 current_platform.fp8_dtype(),
265 shape=(1, compile_key.num_heads, compile_key.index_q_head_dim),
266 strides=(q_stride0, compile_key.index_q_head_dim, 1),
267 ),
268 index_weights_out=TritonWarmupTensor(
269 torch.float32,
270 shape=(1, compile_key.num_heads),
271 ),
272 fp8_max=compile_key.fp8_max,
273 use_fnuz=compile_key.use_fnuz,
274 )
275
276 @kernel_launcher
277 def __call__(
278 self,
279 positions: torch.Tensor,
280 index_q: torch.Tensor,
281 index_q_cos_sin_cache: torch.Tensor,
282 index_weights: torch.Tensor,
283 index_weights_softmax_scale: float,
284 index_weights_head_scale: float,
285 index_q_fp8: torch.Tensor,
286 index_weights_out: torch.Tensor,
287 *,
288 fp8_max: float,
289 use_fnuz: bool,
290 ) -> LaunchSpec:
291 num_tokens = positions.shape[0]
292 num_index_q_heads = index_q.shape[1]
293 return (num_tokens, num_index_q_heads), dict(
294 pos_ptr=positions,
295 index_q_stride0=index_q.stride(0),
296 index_q_stride1=index_q.stride(1),
297 index_q_cos_sin_ptr=index_q_cos_sin_cache,
298 index_q_cos_sin_stride=index_q_cos_sin_cache.stride(0),
299 INDEX_Q_HALF_ROT_DIM=index_q_cos_sin_cache.shape[-1] // 2,
300 index_q_fp8_stride0=index_q_fp8.stride(0),
301 index_q_fp8_stride1=index_q_fp8.stride(1),
302 INDEX_Q_HEAD_DIM=index_q.shape[2],
303 index_weights_stride=index_weights.stride(0),
304 index_weights_out_stride=index_weights_out.stride(0),
305 FP8_MAX=fp8_max,
306 USE_FNUZ=use_fnuz,
307 USE_EXPLICIT_FMA=current_platform.is_rocm(),
308 num_warps=1,
309 )
310
311
312 def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]:
313 """Weights dtypes the model's indexer layers ask for: fp32 for the dense
314 scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer
315 (`SparseMQAIndexer`) is enabled."""
316 if vllm_config.attention_config.indexer_sparse_logits:
317 return (torch.float32, torch.bfloat16)
318 return (torch.float32,)
319
320
321 class FusedIndexerQRopeMxFp4TritonKernel(
322 VllmTritonJitKernel["FusedIndexerQRopeMxFp4TritonKernel.CompileKey"]
323 ):
324 @dataclass(frozen=True)
325 class CompileKey:
326 dtype: torch.dtype
327 weights_out_dtype: torch.dtype
328 num_heads: int
329 index_q_half_rot_dim: int
330 index_q_head_dim: int
331 mxfp4_block: int
332
333 @staticmethod
334 @triton.jit
335 def kernel(
336 pos_ptr,
337 # Index Q RoPE input (fp/bf16)
338 index_q_ptr,
339 index_q_stride0,
340 index_q_stride1,
341 index_q_cos_sin_ptr,
342 index_q_cos_sin_stride,
343 INDEX_Q_HALF_ROT_DIM: tl.constexpr,
344 # MXFP4 Q outputs
345 index_q_mxfp4_ptr, # uint8, (T, H, HEAD_DIM // 2)
346 index_q_mxfp4_stride0,
347 index_q_mxfp4_stride1,
348 index_q_scale_ptr, # uint8 ue8m0, (T, H, HEAD_DIM // BLOCK)
349 index_q_scale_stride0,
350 index_q_scale_stride1,
351 INDEX_Q_HEAD_DIM: tl.constexpr,
352 MXFP4_BLOCK: tl.constexpr,
353 # Weights (NO per-token q_scale fold for MXFP4; per-block scales stay
354 # with the Q values in the output scale tensor).
355 index_weights_ptr,
356 index_weights_stride,
357 index_weights_softmax_scale,
358 index_weights_head_scale,
359 index_weights_out_ptr,
360 index_weights_out_stride,
361 ):
362 INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM
363 INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM
364 NUM_NOPE_BLOCKS: tl.constexpr = INDEX_Q_NOPE_DIM // MXFP4_BLOCK
365 NUM_ROPE_BLOCKS: tl.constexpr = INDEX_Q_ROT_DIM // MXFP4_BLOCK
366 HALF_BLOCK: tl.constexpr = MXFP4_BLOCK // 2
367 tl.static_assert(INDEX_Q_NOPE_DIM >= 0)
368 tl.static_assert(INDEX_Q_NOPE_DIM % MXFP4_BLOCK == 0)
369 tl.static_assert(INDEX_Q_ROT_DIM % MXFP4_BLOCK == 0)
370 tl.static_assert(MXFP4_BLOCK % 2 == 0)
371
372 tok_idx = tl.program_id(0)
373 head_idx = tl.program_id(1)
374
375 pos = tl.load(pos_ptr + tok_idx)
376
377 q_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
378 out_base = (
379 index_q_mxfp4_ptr
380 + tok_idx * index_q_mxfp4_stride0
381 + head_idx * index_q_mxfp4_stride1
382 )
383 scale_base = (
384 index_q_scale_ptr
385 + tok_idx * index_q_scale_stride0
386 + head_idx * index_q_scale_stride1
387 )
388
389 half_off = tl.arange(0, HALF_BLOCK)
390
391 # ---- NoPE blocks: direct load, pair as (even-index, odd-index) values ----
392 for b in tl.static_range(NUM_NOPE_BLOCKS):
393 base = b * MXFP4_BLOCK
394 x_lo = tl.load(q_base + base + half_off * 2).to(tl.float32)
395 x_hi = tl.load(q_base + base + half_off * 2 + 1).to(tl.float32)
396 packed, ue8m0 = _quantize_mxfp4_pair(x_lo, x_hi)
397 tl.store(out_base + base // 2 + half_off, packed)
398 tl.store(scale_base + b, ue8m0)
399
400 # ---- RoPE blocks: apply GPT-J interleaved RoPE to the block's 16 pairs,
401 # then quantize. Each block covers HALF_BLOCK (=16) cos/sin pairs. ----
402 rot_q_base = q_base + INDEX_Q_NOPE_DIM
403 for b in tl.static_range(NUM_ROPE_BLOCKS):
404 pair_off = b * HALF_BLOCK + half_off # indices in [0, HALF_ROT_DIM)
405 cos_b = tl.load(
406 index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + pair_off
407 ).to(tl.float32)
408 sin_b = tl.load(
409 index_q_cos_sin_ptr
410 + pos * index_q_cos_sin_stride
411 + pair_off
412 + INDEX_Q_HALF_ROT_DIM
413 ).to(tl.float32)
414 x_even = tl.load(rot_q_base + pair_off * 2).to(tl.float32)
415 x_odd = tl.load(rot_q_base + pair_off * 2 + 1).to(tl.float32)
416 r_even = x_even * cos_b - x_odd * sin_b
417 r_odd = x_odd * cos_b + x_even * sin_b
418 # bf16 roundtrip for parity with the FP8 kernel / reference numerics.
419 r_even = r_even.to(tl.bfloat16).to(tl.float32)
420 r_odd = r_odd.to(tl.bfloat16).to(tl.float32)
421 packed, ue8m0 = _quantize_mxfp4_pair(r_even, r_odd)
422 rope_byte_off = (INDEX_Q_NOPE_DIM + b * MXFP4_BLOCK) // 2
423 tl.store(out_base + rope_byte_off + half_off, packed)
424 tl.store(scale_base + NUM_NOPE_BLOCKS + b, ue8m0)
425
426 # MXFP4 weight-fold contract:
427 # index_weights_out = index_weights * softmax_scale * head_scale
428 # NOTE: q_scale is NOT folded here (contrast with the FP8 kernel above).
429 # MXFP4 Q emits a separate ue8m0 scale tensor of shape
430 # (T, H, HEAD_DIM // MXFP4_BLOCK) alongside the packed values, so each
431 # per-block scale is applied by the downstream MXFP4 logits kernel when
432 # dequantizing Q — there is no per-token scalar to fold into `weights`.
433 index_weights = tl.load(
434 index_weights_ptr + tok_idx * index_weights_stride + head_idx
435 ).to(tl.float32)
436 index_weights *= index_weights_softmax_scale
437 index_weights *= index_weights_head_scale
438 tl.store(
439 index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
440 index_weights,
441 )
442
443 def dispatch( # type: ignore[override]
444 self,
445 *,
446 dtype: torch.dtype,
447 weights_out_dtype: torch.dtype,
448 num_heads: int,
449 head_dim: int,
450 rope_dim: int,
451 ) -> CompileKey:
452 return self.CompileKey(
453 dtype=dtype,
454 weights_out_dtype=weights_out_dtype,
455 num_heads=num_heads,
456 index_q_half_rot_dim=rope_dim // 2,
457 index_q_head_dim=head_dim,
458 mxfp4_block=MXFP4_BLOCK_SIZE,
459 )
460
461 def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
462 hf_config = vllm_config.model_config.hf_config
463 num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0)
464 head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
465 rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
466 if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
467 return []
468
469 return self._trace_dispatch(self.dispatch)(
470 dtype=vllm_config.model_config.dtype,
471 weights_out_dtype=_indexer_weights_out_dtypes(vllm_config),
472 num_heads=num_heads,
473 head_dim=head_dim,
474 rope_dim=rope_dim,
475 )
476
477 def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
478 q_stride0 = compile_key.num_heads * compile_key.index_q_head_dim
479 packed_head_dim = compile_key.index_q_head_dim // 2
480 scale_head_dim = compile_key.index_q_head_dim // compile_key.mxfp4_block
481 return dict(
482 positions=TritonWarmupTensor(torch.int64),
483 index_q=TritonWarmupTensor(
484 compile_key.dtype,
485 shape=(1, compile_key.num_heads, compile_key.index_q_head_dim),
486 strides=(q_stride0, compile_key.index_q_head_dim, 1),
487 ),
488 index_q_cos_sin_cache=TritonWarmupTensor(
489 torch.float32,
490 shape=(1, compile_key.index_q_half_rot_dim * 2),
491 ),
492 index_weights=TritonWarmupTensor(
493 compile_key.dtype,
494 shape=(1, compile_key.num_heads),
495 ),
496 index_weights_softmax_scale=1.0,
497 index_weights_head_scale=1.0,
498 index_q_packed=TritonWarmupTensor(
499 torch.uint8,
500 shape=(1, compile_key.num_heads, packed_head_dim),
501 ),
502 index_q_scale=TritonWarmupTensor(
503 torch.uint8,
504 shape=(1, compile_key.num_heads, scale_head_dim),
505 ),
506 index_weights_out=TritonWarmupTensor(
507 compile_key.weights_out_dtype,
508 shape=(1, compile_key.num_heads),
509 ),
510 )
511
512 @kernel_launcher
513 def __call__(
514 self,
515 positions: torch.Tensor,
516 index_q: torch.Tensor,
517 index_q_cos_sin_cache: torch.Tensor,
518 index_weights: torch.Tensor,
519 index_weights_softmax_scale: float,
520 index_weights_head_scale: float,
521 index_q_packed: torch.Tensor,
522 index_q_scale: torch.Tensor,
523 index_weights_out: torch.Tensor,
524 ) -> LaunchSpec:
525 num_tokens = positions.shape[0]
526 num_index_q_heads = index_q.shape[1]
527 return (num_tokens, num_index_q_heads), dict(
528 pos_ptr=positions,
529 index_q_stride0=index_q.stride(0),
530 index_q_stride1=index_q.stride(1),
531 index_q_cos_sin_ptr=index_q_cos_sin_cache,
532 index_q_cos_sin_stride=index_q_cos_sin_cache.stride(0),
533 INDEX_Q_HALF_ROT_DIM=index_q_cos_sin_cache.shape[-1] // 2,
534 index_q_mxfp4_ptr=index_q_packed,
535 index_q_mxfp4_stride0=index_q_packed.stride(0),
536 index_q_mxfp4_stride1=index_q_packed.stride(1),
537 index_q_scale_stride0=index_q_scale.stride(0),
538 index_q_scale_stride1=index_q_scale.stride(1),
539 INDEX_Q_HEAD_DIM=index_q.shape[2],
540 MXFP4_BLOCK=MXFP4_BLOCK_SIZE,
541 index_weights_stride=index_weights.stride(0),
542 index_weights_out_stride=index_weights_out.stride(0),
543 num_warps=1,
544 )
545
546
547 def fused_indexer_q_rope_quant(
548 positions: torch.Tensor,
549 index_q: torch.Tensor,
550 index_q_cos_sin_cache: torch.Tensor,
551 # Index weights
552 index_weights: torch.Tensor,
553 index_weights_softmax_scale: float,
554 index_weights_head_scale: float,
555 use_fp4: bool = False,
556 weights_out_dtype: torch.dtype = torch.float32,
557 ) -> tuple[
558 torch.Tensor | tuple[torch.Tensor, torch.Tensor],
559 torch.Tensor,
560 ]:
561 """Fused RoPE + quantize Q for the sparse indexer.
562
563 ``weights_out_dtype`` is the dtype the downstream scoring kernel takes:
564 fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse
565 MQA-logits kernels (CUDA MXFP4 path only).
566
567 Weight-fold semantics (important — the two paths differ):
568
569 FP8 path (use_fp4=False, default):
570 q_fp8 : (T, H, HEAD_DIM) platform fp8 (e4m3fnuz on gfx942,
571 e4m3fn elsewhere); per-token-per-head scalar scale
572 (NOT stored — folded into weights below)
573 weights_out = weights * q_scale * softmax_scale * head_scale
574 Rationale: a single per-token q_scale is a scalar the downstream FP8
575 logits kernel would otherwise multiply in. Folding it into `weights`
576 avoids emitting a separate tensor and is free for the logits kernel.
577
578 MXFP4 path (use_fp4=True):
579 q_packed : (T, H, HEAD_DIM // 2) uint8 (2 E2M1 nibbles per byte)
580 q_scale : (T, H, HEAD_DIM // MXFP4_BLOCK_SIZE) uint8 ue8m0 bytes
581 weights_out = weights * softmax_scale * head_scale
582 Rationale: MXFP4 has PER-BLOCK (32-element) scales that live with
583 the Q values — they cannot be folded into a per-token weight
584 scalar, so `weights` carries only the softmax and head scales.
585
586 Returns (q_quant, weights_out) where q_quant is either a Tensor (FP8) or
587 a (values, scales) tuple (MXFP4). This matches the union type accepted
588 by `SparseAttnIndexer.forward_*`.
589 """
590 assert positions.ndim == 1
591 assert index_q.ndim == 3
592 assert index_q_cos_sin_cache.ndim == 2
593
594 num_tokens = positions.shape[0]
595 num_index_q_heads = index_q.shape[1]
596 index_q_head_dim = index_q.shape[2]
597
598 index_weights_out = torch.empty_like(index_weights, dtype=weights_out_dtype)
599
600 if use_fp4:
601 assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
602 f"head_dim={index_q_head_dim} must be a multiple of MXFP4 block "
603 f"size {MXFP4_BLOCK_SIZE}"
604 )
605 num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE
606 index_q_packed = torch.empty(
607 (num_tokens, num_index_q_heads, index_q_head_dim // 2),
608 dtype=torch.uint8,
609 device=index_q.device,
610 )
611 index_q_scale = torch.empty(
612 (num_tokens, num_index_q_heads, num_scale_blocks),
613 dtype=torch.uint8,
614 device=index_q.device,
615 )
616 if has_cutedsl():
617 # lazily import, otherwise some tests fail due to CUDA driver init failure.
618 from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
619 _INDEXER_Q_MXFP4_KERNEL,
620 )
621
622 _INDEXER_Q_MXFP4_KERNEL(
623 positions=positions,
624 q=index_q,
625 cos_sin_cache=index_q_cos_sin_cache,
626 weights=index_weights,
627 weights_softmax_scale=index_weights_softmax_scale,
628 weights_head_scale=index_weights_head_scale,
629 q_packed=index_q_packed,
630 q_scale=index_q_scale,
631 weights_out=index_weights_out,
632 )
633 elif current_platform.is_xpu():
634 assert weights_out_dtype == torch.float32, weights_out_dtype
635 torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4(
636 index_q,
637 positions,
638 index_q_cos_sin_cache,
639 index_weights,
640 index_weights_softmax_scale,
641 index_weights_head_scale,
642 index_q_packed,
643 index_q_scale,
644 index_weights_out,
645 )
646 else:
647 _FUSED_INDEXER_Q_ROPE_MXFP4_TRITON_KERNEL(
648 positions,
649 index_q,
650 index_q_cos_sin_cache,
651 index_weights,
652 index_weights_softmax_scale,
653 index_weights_head_scale,
654 index_q_packed,
655 index_q_scale,
656 index_weights_out,
657 )
658
659 # Values stay uint8 (2 E2M1 nibbles per byte). Scales are 4 ue8m0
660 # bytes per (token, head) reinterpreted as one int32, then squeezed
661 # from (T, H, 1) to (T, H) to match DeepGEMM's expected q_sf rank
662 # (prefill wants 2-D (seq_len, num_heads); decode reshapes this to
663 # 3-D (batch, next_n, num_heads)).
664 return (
665 index_q_packed,
666 index_q_scale.view(torch.int32).squeeze(-1),
667 ), index_weights_out
668
669 assert weights_out_dtype == torch.float32, weights_out_dtype
670 fp8_dtype = current_platform.fp8_dtype()
671 use_fnuz = fp8_dtype == torch.float8_e4m3fnuz
672 fp8_max = 224.0 if use_fnuz else 448.0
673 index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
674 if has_cutedsl():
675 # lazily import, otherwise some tests fail due to CUDA driver init failure.
676 from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
677 _INDEXER_Q_FP8_KERNEL,
678 )
679
680 _INDEXER_Q_FP8_KERNEL(
681 positions=positions,
682 q=index_q,
683 cos_sin_cache=index_q_cos_sin_cache,
684 weights=index_weights,
685 weights_softmax_scale=index_weights_softmax_scale,
686 weights_head_scale=index_weights_head_scale,
687 q_fp8=index_q_fp8.view(torch.uint8),
688 weights_out=index_weights_out,
689 )
690 elif current_platform.is_xpu():
691 torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8(
692 index_q,
693 positions,
694 index_q_cos_sin_cache,
695 index_weights,
696 index_weights_softmax_scale,
697 index_weights_head_scale,
698 index_q_fp8,
699 index_weights_out,
700 )
701 elif current_platform.is_cpu():
702 from vllm._custom_ops import fused_indexer_q_rope_quant_cpu
703
704 fused_indexer_q_rope_quant_cpu(
705 positions.contiguous(),
706 index_q.to(torch.float32),
707 index_q_cos_sin_cache.to(torch.float32).contiguous(),
708 index_q_fp8,
709 index_weights.to(torch.float32),
710 index_weights_softmax_scale,
711 index_weights_head_scale,
712 index_weights_out,
713 )
714 else:
715 _FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL(
716 positions,
717 index_q,
718 index_q_cos_sin_cache,
719 index_weights,
720 index_weights_softmax_scale,
721 index_weights_head_scale,
722 index_q_fp8,
723 index_weights_out,
724 fp8_max=fp8_max,
725 use_fnuz=use_fnuz,
726 )
727 return index_q_fp8, index_weights_out
728
729
730 _FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL = FusedIndexerQRopeQuantTritonKernel()
731 _FUSED_INDEXER_Q_ROPE_MXFP4_TRITON_KERNEL = FusedIndexerQRopeMxFp4TritonKernel()