@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,175 lines 43.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Callable
4
5 import torch
6
7 from vllm.platforms import current_platform
8 from vllm.triton_utils import tl, triton
9 from vllm.utils.torch_utils import is_quantized_kv_cache
10
11 # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the
12 # has_indexer=False path so the indexer args don't allocate every call.
13 _DUMMY_CACHE: dict[tuple, torch.Tensor] = {}
14
15 # Tile shape of the indexer-K cache's shuffled layout, used when the cache reports
16 # uses_shuffled_layout. The block tile is a token count; the head tile is a byte
17 # count, which the kernel converts to cache elements.
18 _INDEXER_CACHE_BLOCK_TILE = 16
19 _INDEXER_CACHE_HEAD_TILE_BYTES = 16
20
21 # FNUZ (e4m3fnuz) on gfx942, OCP (e4m3fn) elsewhere. 224.0 keeps the ue8m0
22 # power-of-two scale inside the FNUZ range after its ceil.
23 _FP8_DTYPE = current_platform.fp8_dtype()
24 _USE_FNUZ = torch.float8_e4m3fnuz == _FP8_DTYPE
25 _FP8_MAX = 224.0 if _USE_FNUZ else 448.0
26
27
28 def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
29 key = (shape, dtype, device)
30 t = _DUMMY_CACHE.get(key)
31 if t is None:
32 t = torch.empty(shape, dtype=dtype, device=device)
33 _DUMMY_CACHE[key] = t
34 return t
35
36
37 @triton.jit
38 def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
39 x = x.to(tl.float32)
40 mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
41 rrms = tl.rsqrt(mean_sq + eps)
42 w = w.to(tl.float32)
43 return (x * rrms) * w
44
45
46 @triton.jit
47 def _get_cos_sin(
48 cos_sin_cache_ptr,
49 cos_sin_cache_stride,
50 pos,
51 HALF_ROT_DIM: tl.constexpr,
52 ):
53 block = tl.arange(0, HALF_ROT_DIM)
54 cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block)
55 cos = cos.to(tl.float32)
56 sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM)
57 sin = sin.to(tl.float32)
58 return cos, sin
59
60
61 @triton.jit
62 def _fp8_ue8m0_quantize(vals, FP8_MAX: tl.constexpr, USE_FNUZ: tl.constexpr):
63 """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
64
65 Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
66 """
67 vals = vals.to(tl.float32)
68 amax = tl.max(tl.abs(vals))
69 scale = tl.div_rn(tl.maximum(amax, 1e-4), FP8_MAX)
70 scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
71 fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
72 fp8_vals = tl.div_rn(vals, scale).to(fp8_dtype)
73 return fp8_vals, scale
74
75
76 @triton.jit
77 def _fp8_quant_and_cache_write(
78 vals,
79 mask,
80 slot_idx,
81 kv_cache_ptr,
82 kv_cache_scale_ptr,
83 cache_block_size,
84 cache_block_stride,
85 offsets,
86 HEAD_DIM: tl.constexpr,
87 SHUFFLE: tl.constexpr,
88 BLOCK_TILE: tl.constexpr,
89 HEAD_TILE: tl.constexpr,
90 FP8_MAX: tl.constexpr,
91 USE_FNUZ: tl.constexpr,
92 ):
93 k_fp8, scale = _fp8_ue8m0_quantize(vals, FP8_MAX, USE_FNUZ)
94
95 block_idx = slot_idx // cache_block_size
96 block_offset = slot_idx % cache_block_size
97 block_start = block_idx * cache_block_stride
98
99 # Shuffled layout: [blk/BLOCK_TILE, HEAD_DIM/HEAD_TILE, BLOCK_TILE, HEAD_TILE],
100 # so one contiguous run holds BLOCK_TILE tokens x HEAD_TILE bytes and the
101 # reader's coalesced loads land on the bytes we wrote.
102 if SHUFFLE:
103 value_off = (
104 block_offset // BLOCK_TILE * BLOCK_TILE * HEAD_DIM
105 + block_offset % BLOCK_TILE * HEAD_TILE
106 + offsets // HEAD_TILE * BLOCK_TILE * HEAD_TILE
107 + offsets % HEAD_TILE
108 )
109 else:
110 value_off = block_offset * HEAD_DIM + offsets
111
112 tl.store(kv_cache_ptr + block_start + value_off, k_fp8, mask=mask)
113 scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
114 tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
115
116
117 @triton.jit
118 def _fused_norm_rope_kernel(
119 pos_ptr,
120 # Q RMS norm
121 q_c_ptr,
122 q_c_stride,
123 q_rms_norm_w_ptr,
124 q_rms_eps,
125 q_c_out_ptr,
126 q_c_out_stride,
127 Q_DIM: tl.constexpr,
128 Q_BLOCK_SIZE: tl.constexpr,
129 # KV RMS norm
130 kv_ptr,
131 kv_stride,
132 kv_rms_norm_w_ptr,
133 kv_rms_eps,
134 kv_out_ptr,
135 kv_out_stride,
136 KV_DIM: tl.constexpr,
137 # KV RoPE
138 kpe_ptr,
139 kpe_stride,
140 kpe_rope_cos_sin_cache_ptr,
141 kpe_rope_cos_sin_cache_stride,
142 kpe_out_ptr,
143 kpe_out_stride,
144 KPE_HALF_ROT_DIM: tl.constexpr,
145 # Index K layer norm
146 index_k_ptr,
147 index_k_stride,
148 index_k_layer_norm_w_ptr,
149 index_k_layer_norm_bias_ptr,
150 index_k_layer_norm_eps,
151 INDEX_K_DIM: tl.constexpr,
152 INDEX_K_BLOCK_SIZE: tl.constexpr,
153 # Index K RoPE
154 index_k_rope_cos_sin_cache_ptr,
155 index_k_rope_cos_sin_cache_stride,
156 index_k_out_ptr,
157 index_k_out_stride,
158 INDEX_K_HALF_ROT_DIM: tl.constexpr,
159 # Cache params
160 slot_mapping_ptr,
161 indexer_slot_mapping_ptr,
162 # Index K FP8 cache
163 indexer_cache_ptr,
164 indexer_cache_scale_ptr,
165 indexer_cache_block_size,
166 indexer_cache_stride,
167 INDEXER_CACHE_SHUFFLE: tl.constexpr,
168 INDEXER_CACHE_BLOCK_TILE: tl.constexpr,
169 INDEXER_CACHE_HEAD_TILE: tl.constexpr,
170 # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
171 mla_cache_ptr,
172 mla_cache_block_stride,
173 mla_cache_entry_stride,
174 MLA_CACHE_BLOCK_SIZE: tl.constexpr,
175 MLA_CACHE_FP8: tl.constexpr,
176 mla_cache_scale_ptr,
177 # fp8_ds_mla cache views (block-scaled fp8 NoPE + unquantized bf16 RoPE).
178 # mla_cache_ptr is the fp8 (1-byte) view, so the block/entry strides above
179 # are byte offsets; these two share the same buffer as fp32 / bf16 views.
180 mla_cache_ds_scale_ptr,
181 mla_cache_ds_rope_ptr,
182 MLA_CACHE_DS_MLA: tl.constexpr,
183 MLA_CACHE_NVFP4: tl.constexpr,
184 MLA_NUM_TILES: tl.constexpr,
185 MLA_TILE_DIM: tl.constexpr,
186 # Top k indices
187 topk_indices_ptr,
188 topk_indices_stride,
189 TOPK: tl.constexpr,
190 TOPK_BLOCK_SIZE: tl.constexpr,
191 HAS_INDEXER: tl.constexpr,
192 INDEX_ROPE_INTERLEAVE: tl.constexpr,
193 USE_PDL: tl.constexpr,
194 FP8_MAX: tl.constexpr,
195 USE_FNUZ: tl.constexpr,
196 ):
197 tok_idx = tl.program_id(0).to(tl.int64)
198 pid = tl.program_id(1)
199 if USE_PDL:
200 tl.extra.cuda.gdc_wait()
201 tl.extra.cuda.gdc_launch_dependents()
202 fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
203 if pid == 3:
204 if not HAS_INDEXER:
205 # Shared layer: reuse the previous indexer layer's top-k; do not
206 # clear the buffer.
207 return
208 # Fill top k indices buffer with -1
209 for i in range(0, TOPK, TOPK_BLOCK_SIZE):
210 offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
211 mask = offset < TOPK
212 tl.store(
213 topk_indices_ptr + tok_idx * topk_indices_stride + offset,
214 -1,
215 mask=mask,
216 )
217 return
218
219 if pid == 2:
220 # Q RMS norm. Runs for every row: under DCP a negative slot only
221 # means another rank owns this token's KV slot, but the query is
222 # still needed on every rank (queries are not sharded), so the
223 # slot-based skip below must not gate it. Padding rows do harmless
224 # row-local extra work (no position load, no cache write).
225 q_block = tl.arange(0, Q_BLOCK_SIZE)
226 q_mask = q_block < Q_DIM
227 q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
228 q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
229 q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
230 tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
231 return
232
233 if slot_mapping_ptr is None:
234 if kv_out_ptr is None and kpe_out_ptr is None and index_k_out_ptr is None:
235 return
236 elif tl.load(slot_mapping_ptr + tok_idx) < 0 and (
237 pid != 1 or (kv_out_ptr is None and kpe_out_ptr is None)
238 ):
239 # Padding, or (under DCP) a token whose KV slot another rank owns.
240 # Dense prefill still consumes the materialized normalized/rotated K
241 # rows on every rank, so pid 1 must produce them even when this rank
242 # must not write the owner-local cache slot.
243 return
244
245 if pid == 1:
246 # KV RMS Norm + KV RoPE + MLA concat_and_cache.
247 # Merged so the normed kv_c and RoPE'd k_pe can be written
248 # to the MLA KV cache directly without a separate kernel.
249
250 # KV RMS Norm (result stays in registers for MLA cache write)
251 kv_block = tl.arange(0, KV_DIM)
252 kv_c = tl.load(kv_ptr + tok_idx * kv_stride + kv_block)
253 kv_c_rms_w = tl.load(kv_rms_norm_w_ptr + kv_block)
254 kv_c = _rms_norm(kv_c, kv_c_rms_w, kv_rms_eps, KV_DIM)
255
256 # KV RoPE (interleaved) on k_pe — in registers only.
257 # k_pe is not needed after the cache write (MLA decode reads
258 # from kv_cache), so we skip writing back to kpe_ptr.
259 pos = tl.load(pos_ptr + tok_idx)
260 cos, sin = _get_cos_sin(
261 kpe_rope_cos_sin_cache_ptr,
262 kpe_rope_cos_sin_cache_stride,
263 pos,
264 KPE_HALF_ROT_DIM,
265 )
266 dim_off = tl.arange(0, KPE_HALF_ROT_DIM)
267 kpe_base = kpe_ptr + tok_idx * kpe_stride
268 x1 = tl.load(kpe_base + dim_off * 2).to(tl.float32)
269 x2 = tl.load(kpe_base + dim_off * 2 + 1).to(tl.float32)
270 r1 = x1 * cos - x2 * sin
271 r2 = x2 * cos + x1 * sin
272
273 # PCP materializes K for the cross-rank gather, then inserts it into cache.
274 if kv_out_ptr is not None:
275 tl.store(kv_out_ptr + tok_idx * kv_out_stride + kv_block, kv_c)
276 if kpe_out_ptr is not None:
277 tl.store(
278 kpe_out_ptr + tok_idx * kpe_out_stride + dim_off * 2,
279 r1.to(kpe_out_ptr.dtype.element_ty),
280 )
281 tl.store(
282 kpe_out_ptr + tok_idx * kpe_out_stride + dim_off * 2 + 1,
283 r2.to(kpe_out_ptr.dtype.element_ty),
284 )
285
286 if slot_mapping_ptr is not None:
287 # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.
288 if mla_cache_entry_stride == 0:
289 return
290
291 slot_idx = tl.load(slot_mapping_ptr + tok_idx)
292 if slot_idx < 0:
293 return
294 mla_block_size = MLA_CACHE_BLOCK_SIZE
295 mla_block_idx = slot_idx // mla_block_size
296 mla_block_off = slot_idx % mla_block_size
297
298 if MLA_CACHE_NVFP4:
299 # nvfp4_ds_mla layout (352 B/token, KV_DIM == 512):
300 # [0, 256) 512 e2m1 NoPE, packed 2/byte (low nibble = even)
301 # [256, 320) 64 e4m3 RoPE, unscaled
302 # [320, 352) 32 e4m3 NoPE scales, byte-permuted
303 # Keep in lockstep with concat_and_cache_nvfp4_ds_mla (CUDA) and
304 # nvfp4_sf_byte() in FlashMLA: byte(s) = 8*(s&3) + (s>>2).
305 byte_base = (
306 mla_block_idx * mla_cache_block_stride
307 + mla_block_off * mla_cache_entry_stride
308 )
309 kv_2d = tl.reshape(kv_c.to(tl.float32), (MLA_NUM_TILES, MLA_TILE_DIM))
310 amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)
311 # sf = e4m3(max(amax/6, 2^-9)), round-to-nearest like
312 # cvt_warp_fp16_to_fp4.
313 sf8 = tl.maximum(amax * (1.0 / 6.0), 0.001953125).to(tl.float8e4nv)
314 q = kv_2d * (1.0 / sf8.to(tl.float32))
315 qp = tl.reshape(q, (KV_DIM // 2, 2))
316 sel = tl.arange(0, 2)[None, :]
317 even = tl.sum(tl.where(sel == 0, qp, 0.0), axis=1)
318 odd = tl.sum(tl.where(sel == 1, qp, 0.0), axis=1)
319 # Same instruction the CUDA writer uses, so the two are bit-exact
320 # by construction: dest byte = {high: odd, low: even}.
321 packed = tl.inline_asm_elementwise(
322 "{ .reg .b8 t; .reg .b32 w;"
323 " cvt.rn.satfinite.e2m1x2.f32 t, $2, $1;"
324 " mov.b32 w, {t, t, t, t}; and.b32 $0, w, 255; }",
325 "=r,f,f",
326 [even, odd],
327 dtype=tl.int32,
328 is_pure=True,
329 pack=1,
330 ).to(tl.uint8)
331 tl.store(
332 mla_cache_ptr + byte_base + tl.arange(0, KV_DIM // 2),
333 packed.to(tl.float8e4nv, bitcast=True),
334 )
335 # RoPE: unscaled e4m3, interleaved like the fp8_ds_mla path
336 rope_base = mla_cache_ptr + byte_base + KV_DIM // 2
337 tl.store(rope_base + dim_off * 2, r1.to(tl.float8e4nv))
338 tl.store(rope_base + dim_off * 2 + 1, r2.to(tl.float8e4nv))
339 # scales, permuted: byte(s) = 8*(s&3) + (s>>2)
340 sf_tile = tl.arange(0, MLA_NUM_TILES)
341 sf_perm = 8 * (sf_tile % 4) + (sf_tile // 4)
342 tl.store(
343 mla_cache_ptr + byte_base + (KV_DIM // 2 + 64) + sf_perm,
344 tl.reshape(sf8, (MLA_NUM_TILES,)),
345 )
346 return
347
348 if MLA_CACHE_DS_MLA:
349 # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element
350 # tile of the NoPE is dynamically quantized to fp8 with its own
351 # power-of-two scale stored as float32; the RoPE tail is bf16.
352 # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values
353 # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales
354 # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE
355 # mla_cache_block_stride / mla_cache_entry_stride are byte strides
356 # (mla_cache_ptr is the 1-byte fp8 view of the uint8 cache).
357 byte_base = (
358 mla_block_idx * mla_cache_block_stride
359 + mla_block_off * mla_cache_entry_stride
360 )
361 kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM))
362 tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)
363 # scale = amax / FP8_MAX (fp8 e4m3 max), matching the reference
364 # concat_and_cache_ds_mla kernel; floored to FLT_MIN.
365 tile_scale = tl.maximum(tile_amax * (1.0 / FP8_MAX), 1e-4)
366 tile_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(tile_scale)))
367 kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(fp8_dtype), (KV_DIM,))
368 tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8)
369 tile_off = tl.arange(0, MLA_NUM_TILES)
370 tl.store(
371 mla_cache_ds_scale_ptr + byte_base // 4 + KV_DIM // 4 + tile_off,
372 tl.reshape(tile_scale, (MLA_NUM_TILES,)),
373 )
374 rope_dst = mla_cache_ds_rope_ptr + byte_base // 2 + (KV_DIM // 2 + 8)
375 tl.store(rope_dst + dim_off * 2, r1.to(tl.bfloat16))
376 tl.store(rope_dst + dim_off * 2 + 1, r2.to(tl.bfloat16))
377 return
378
379 dst = (
380 mla_cache_ptr
381 + mla_block_idx * mla_cache_block_stride
382 + mla_block_off * mla_cache_entry_stride
383 )
384 # kv_c_normed (KV_DIM elements)
385 if MLA_CACHE_FP8:
386 scale = tl.load(mla_cache_scale_ptr)
387 kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(fp8_dtype)
388 tl.store(dst + kv_block, kv_c_fp8)
389 else:
390 tl.store(dst + kv_block, kv_c)
391 # k_pe_roped (from registers, interleaved layout)
392 if MLA_CACHE_FP8:
393 tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(fp8_dtype))
394 tl.store(
395 dst + KV_DIM + dim_off * 2 + 1,
396 (r2 / scale).to(fp8_dtype),
397 )
398 else:
399 tl.store(dst + KV_DIM + dim_off * 2, r1)
400 tl.store(dst + KV_DIM + dim_off * 2 + 1, r2)
401 elif pid == 0:
402 if not HAS_INDEXER:
403 # Shared layer: no indexer K to process.
404 return
405 # Fused: Index K LayerNorm + RoPE + FP8 quant + cache write.
406 # Eliminates the separate indexer_k_quant_and_cache kernel launch.
407
408 index_k_block = tl.arange(0, INDEX_K_BLOCK_SIZE)
409 index_k_mask = index_k_block < INDEX_K_DIM
410 index_k = tl.load(
411 index_k_ptr + tok_idx * index_k_stride + index_k_block,
412 mask=index_k_mask,
413 other=0.0,
414 ).to(tl.float32)
415 index_k_w = tl.load(
416 index_k_layer_norm_w_ptr + index_k_block, mask=index_k_mask
417 ).to(tl.float32)
418 index_k_b = tl.load(
419 index_k_layer_norm_bias_ptr + index_k_block, mask=index_k_mask
420 ).to(tl.float32)
421
422 # 1. LayerNorm. Keep (mean, rstd) so the RoPE rotation partner can be
423 # re-normalized in registers below, avoiding a global scratch buffer.
424 mean = tl.sum(index_k, axis=0) / INDEX_K_DIM
425 diff = tl.where(index_k_mask, index_k - mean, 0.0)
426 var = tl.sum(diff * diff, axis=0) / INDEX_K_DIM
427 rstd = tl.rsqrt(var + index_k_layer_norm_eps)
428 normed = (index_k - mean) * rstd * index_k_w + index_k_b
429
430 # 2. RoPE on the rotation region. Supports both interleaved (adjacent
431 # pairs, e.g. GLM-5.2) and NeoX (split-half, e.g. DeepSeek-V3.2). The
432 # rotation partner is gathered from the read-only inputs and
433 # re-normalized with the same (mean, rstd) — no scratch, no atomics.
434 pos = tl.load(pos_ptr + tok_idx)
435 in_rope = index_k_block < 2 * INDEX_K_HALF_ROT_DIM
436 if INDEX_ROPE_INTERLEAVE:
437 # pair i = block // 2; partner = block ^ 1; even -> -sin, odd -> +sin.
438 cos_idx = index_k_block // 2
439 partner_offs = tl.where(in_rope, index_k_block ^ 1, index_k_block)
440 sign = tl.where(index_k_block % 2 == 0, -1.0, 1.0)
441 else:
442 # NeoX: pair across halves; partner = block ^ HALF.
443 cos_idx = index_k_block % INDEX_K_HALF_ROT_DIM
444 partner_offs = tl.where(
445 in_rope, index_k_block ^ INDEX_K_HALF_ROT_DIM, index_k_block
446 )
447 sign = tl.where(index_k_block < INDEX_K_HALF_ROT_DIM, -1.0, 1.0)
448 cos_full = tl.load(
449 index_k_rope_cos_sin_cache_ptr
450 + pos * index_k_rope_cos_sin_cache_stride
451 + cos_idx,
452 mask=in_rope,
453 other=1.0,
454 ).to(tl.float32)
455 sin_full = tl.load(
456 index_k_rope_cos_sin_cache_ptr
457 + pos * index_k_rope_cos_sin_cache_stride
458 + INDEX_K_HALF_ROT_DIM
459 + cos_idx,
460 mask=in_rope,
461 other=0.0,
462 ).to(tl.float32)
463 # normed[partner_offs] == (raw_partner - mean) * rstd * w_partner +
464 # b_partner: gather the raw partner and its norm affine (read-only
465 # loads), then apply the same per-token mean/rstd.
466 raw_partner = tl.load(
467 index_k_ptr + tok_idx * index_k_stride + partner_offs,
468 mask=index_k_mask,
469 other=0.0,
470 ).to(tl.float32)
471 w_partner = tl.load(
472 index_k_layer_norm_w_ptr + partner_offs, mask=index_k_mask
473 ).to(tl.float32)
474 b_partner = tl.load(
475 index_k_layer_norm_bias_ptr + partner_offs, mask=index_k_mask
476 ).to(tl.float32)
477 normed_partner = (raw_partner - mean) * rstd * w_partner + b_partner
478 roped = normed * cos_full + sign * normed_partner * sin_full
479 result = tl.where(in_rope, roped, normed)
480
481 if index_k_out_ptr is not None:
482 tl.store(
483 index_k_out_ptr + tok_idx * index_k_out_stride + index_k_block,
484 result.to(index_k_out_ptr.dtype.element_ty),
485 mask=index_k_mask,
486 )
487
488 # PCP inserts index K after gathering; other paths write it directly.
489 if indexer_cache_ptr is not None and indexer_slot_mapping_ptr is not None:
490 slot_idx = tl.load(indexer_slot_mapping_ptr + tok_idx)
491 _fp8_quant_and_cache_write(
492 result,
493 index_k_mask,
494 slot_idx,
495 indexer_cache_ptr,
496 indexer_cache_scale_ptr,
497 indexer_cache_block_size,
498 indexer_cache_stride,
499 index_k_block,
500 INDEX_K_DIM,
501 INDEXER_CACHE_SHUFFLE,
502 INDEXER_CACHE_BLOCK_TILE,
503 INDEXER_CACHE_HEAD_TILE,
504 FP8_MAX,
505 USE_FNUZ,
506 )
507
508
509 def fused_norm_rope(
510 positions: torch.Tensor,
511 q_c: torch.Tensor,
512 q_rms_norm_w: torch.Tensor,
513 q_rms_eps: float,
514 kv_c: torch.Tensor,
515 kv_rms_norm_w: torch.Tensor,
516 kv_rms_eps: float,
517 k_pe: torch.Tensor,
518 k_rope_cos_sin_cache: torch.Tensor,
519 index_k: torch.Tensor | None,
520 index_k_layer_norm_w: torch.Tensor | None,
521 index_k_layer_norm_bias: torch.Tensor | None,
522 index_k_layer_norm_eps: float,
523 index_k_rope_cos_sin_cache: torch.Tensor | None,
524 topk_indices_buffer: torch.Tensor,
525 # Cache params for fused writes
526 slot_mapping: torch.Tensor | None = None,
527 indexer_slot_mapping: torch.Tensor | None = None,
528 indexer_k_cache: torch.Tensor | None = None,
529 indexer_cache_shuffled: bool = False,
530 mla_kv_cache: torch.Tensor | None = None,
531 mla_kv_cache_dtype: str = "auto",
532 mla_k_scale: torch.Tensor | None = None,
533 has_indexer: bool = True,
534 index_rope_interleave: bool = False,
535 q_c_out: torch.Tensor | None = None,
536 kv_c_out: torch.Tensor | None = None,
537 k_pe_out: torch.Tensor | None = None,
538 index_k_out: torch.Tensor | None = None,
539 ) -> torch.Tensor:
540 assert positions.ndim == 1
541 assert q_c.ndim == 2
542 assert kv_c.ndim == 2
543 assert k_pe.ndim == 2
544 assert topk_indices_buffer.ndim == 2
545
546 num_tokens = positions.shape[0]
547 q_dim = q_c.shape[-1]
548 kv_dim = kv_c.shape[-1]
549 device = positions.device
550
551 # Shared (no-indexer) layers: substitute cached 1-element dummies so the
552 # kernel launches cleanly; pid 0/3 (indexer + topk fill) skipped by
553 # HAS_INDEXER and never dereference them.
554 if not has_indexer:
555 indexer_k_cache = None
556 index_k = _dummy((1, 1), q_c.dtype, device)
557 index_k_layer_norm_w = _dummy((1,), torch.float32, device)
558 index_k_layer_norm_bias = _dummy((1,), torch.float32, device)
559 index_k_rope_cos_sin_cache = k_rope_cos_sin_cache
560 assert index_k is not None
561 assert index_k_rope_cos_sin_cache is not None
562 index_k_dim = index_k.shape[-1]
563 topk = topk_indices_buffer.shape[-1]
564 if indexer_k_cache is not None or mla_kv_cache is not None:
565 assert slot_mapping is not None
566 else:
567 slot_mapping = None
568
569 # --- Indexer K cache setup ---
570 if indexer_k_cache is not None:
571 assert slot_mapping is not None
572 if indexer_slot_mapping is None:
573 indexer_slot_mapping = slot_mapping
574 idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
575 idx_cache_block_size = indexer_k_cache.shape[1]
576 idx_cache_stride = indexer_k_cache.shape[2]
577 # The caller's cache reports whether its reader expects the shuffled
578 # layout; see DeepseekV32IndexerCache.uses_shuffled_layout.
579 idx_cache_shuffle = indexer_cache_shuffled
580 if idx_cache_shuffle:
581 assert idx_cache_block_size % _INDEXER_CACHE_BLOCK_TILE == 0, (
582 f"indexer K cache block size {idx_cache_block_size} must be a "
583 f"multiple of {_INDEXER_CACHE_BLOCK_TILE} for the shuffled layout"
584 )
585 idx_cache_stride = indexer_k_cache.stride(0)
586 if indexer_k_cache.dtype == torch.uint8:
587 indexer_k_cache = indexer_k_cache.view(_FP8_DTYPE)
588 # The head tile is a byte count; the kernel indexes the cache in elements.
589 idx_cache_head_tile = (
590 _INDEXER_CACHE_HEAD_TILE_BYTES // indexer_k_cache.element_size()
591 )
592 else:
593 idx_cache_scale_view = None
594 idx_cache_block_size = 1
595 idx_cache_stride = 0
596 idx_cache_shuffle = False
597 idx_cache_head_tile = _INDEXER_CACHE_HEAD_TILE_BYTES
598
599 # --- MLA KV cache setup ---
600 mla_cache_nvfp4 = mla_kv_cache_dtype == "nvfp4_ds_mla"
601 mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" or mla_cache_nvfp4
602 mla_cache_fp8 = is_quantized_kv_cache(mla_kv_cache_dtype) and not mla_cache_ds_mla
603 mla_num_tiles = 1
604 mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device)
605 mla_ds_rope_view = torch.empty(0, dtype=torch.bfloat16, device=device)
606 if mla_kv_cache is not None:
607 mla_block_size = mla_kv_cache.shape[1]
608 if mla_cache_ds_mla:
609 # Custom byte-addressed layout: 656 B/token for fp8_ds_mla, 352 B
610 # for nvfp4_ds_mla. mla_cache_ptr is the 1-byte fp8 view, so
611 # block/entry strides are byte offsets and the fp32/bf16 views
612 # (fp8_ds_mla only) share the same buffer.
613 assert kv_dim == 512, "ds_mla layouts require kv_lora_rank == 512"
614 mla_num_tiles = kv_dim // 16 if mla_cache_nvfp4 else kv_dim // 128
615 u8_cache = mla_kv_cache.view(torch.uint8)
616 mla_block_stride = u8_cache.stride(0)
617 mla_entry_stride = u8_cache.stride(1)
618 mla_ds_scale_view = u8_cache.view(torch.float32)
619 mla_ds_rope_view = u8_cache.view(torch.bfloat16)
620 mla_kv_cache = u8_cache.view(_FP8_DTYPE)
621 else:
622 mla_block_stride = mla_kv_cache.stride(0)
623 mla_entry_stride = mla_kv_cache.stride(1)
624 if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
625 mla_kv_cache = mla_kv_cache.view(_FP8_DTYPE)
626 if mla_k_scale is None:
627 mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
628 else:
629 # Dummy cache values; a zero entry stride disables the cache write.
630 mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)
631 mla_block_stride = 0
632 mla_entry_stride = 0
633 mla_block_size = 1
634 mla_k_scale = _dummy((1,), torch.float32, device)
635
636 if q_c_out is None:
637 q_c_out = torch.empty_like(q_c)
638 kv_c_out_stride = 0
639 k_pe_out_stride = 0
640 index_k_out_stride = 0
641 if kv_c_out is not None:
642 assert kv_c_out.shape == kv_c.shape
643 kv_c_out_stride = kv_c_out.stride(0)
644 if k_pe_out is not None:
645 assert k_pe_out.shape == k_pe.shape
646 k_pe_out_stride = k_pe_out.stride(0)
647 if index_k_out is not None:
648 assert index_k_out.shape == index_k.shape
649 index_k_out_stride = index_k_out.stride(0)
650 use_pdl = current_platform.is_arch_support_pdl()
651 pdl_kwargs = (
652 {"USE_PDL": use_pdl, "launch_pdl": True} if use_pdl else {"USE_PDL": use_pdl}
653 )
654 _fused_norm_rope_kernel[(num_tokens, 4)](
655 positions,
656 # Q RMS norm
657 q_c,
658 q_c.stride(0),
659 q_rms_norm_w,
660 q_rms_eps,
661 q_c_out,
662 q_c_out.stride(0),
663 q_dim,
664 triton.next_power_of_2(q_dim),
665 # KV RMS norm
666 kv_c,
667 kv_c.stride(0),
668 kv_rms_norm_w,
669 kv_rms_eps,
670 kv_c_out,
671 kv_c_out_stride,
672 kv_dim,
673 # KV RoPE
674 k_pe,
675 k_pe.stride(0),
676 k_rope_cos_sin_cache,
677 k_rope_cos_sin_cache.stride(0),
678 k_pe_out,
679 k_pe_out_stride,
680 k_rope_cos_sin_cache.shape[-1] // 2,
681 # Index K layer norm + RoPE + FP8 quant
682 index_k,
683 index_k.stride(0),
684 index_k_layer_norm_w,
685 index_k_layer_norm_bias,
686 index_k_layer_norm_eps,
687 index_k_dim,
688 triton.next_power_of_2(index_k_dim),
689 index_k_rope_cos_sin_cache,
690 index_k_rope_cos_sin_cache.stride(0),
691 index_k_out,
692 index_k_out_stride,
693 index_k_rope_cos_sin_cache.shape[-1] // 2,
694 # Cache params
695 slot_mapping,
696 indexer_slot_mapping,
697 indexer_k_cache,
698 idx_cache_scale_view,
699 idx_cache_block_size,
700 idx_cache_stride,
701 idx_cache_shuffle,
702 _INDEXER_CACHE_BLOCK_TILE,
703 idx_cache_head_tile,
704 # MLA KV cache (uses same slot_mapping)
705 mla_kv_cache,
706 mla_block_stride,
707 mla_entry_stride,
708 mla_block_size,
709 mla_cache_fp8,
710 mla_k_scale,
711 mla_ds_scale_view,
712 mla_ds_rope_view,
713 mla_cache_ds_mla and not mla_cache_nvfp4,
714 mla_cache_nvfp4,
715 mla_num_tiles,
716 kv_dim // mla_num_tiles if mla_cache_ds_mla else 1,
717 # Top k indices buffer
718 topk_indices_buffer,
719 topk_indices_buffer.stride(0),
720 topk,
721 TOPK_BLOCK_SIZE=1024,
722 HAS_INDEXER=has_indexer,
723 INDEX_ROPE_INTERLEAVE=index_rope_interleave,
724 FP8_MAX=_FP8_MAX,
725 USE_FNUZ=_USE_FNUZ,
726 **pdl_kwargs,
727 )
728 return q_c_out
729
730
731 @triton.jit
732 def _fused_q_kernel(
733 pos_ptr,
734 # MQA query PE: RoPE + FP8 pack into output tail
735 q_pe_ptr,
736 q_pe_stride0,
737 q_pe_stride1,
738 NUM_Q_HEADS: tl.constexpr,
739 q_pe_cos_sin_ptr,
740 q_pe_cos_sin_stride,
741 Q_PE_HALF_ROT_DIM: tl.constexpr,
742 # Index Q RoPE
743 index_q_ptr,
744 index_q_stride0,
745 index_q_stride1,
746 NUM_INDEX_Q_HEADS: tl.constexpr,
747 index_q_cos_sin_ptr,
748 index_q_cos_sin_stride,
749 INDEX_Q_HALF_ROT_DIM: tl.constexpr,
750 # Index Q Quantize
751 index_q_fp8_ptr,
752 index_q_fp8_stride0,
753 index_q_fp8_stride1,
754 INDEX_Q_HEAD_DIM: tl.constexpr,
755 # MQA query pack: quantize ql_nope and RoPE+quantize q_pe into mqa_q_fp8
756 ql_nope_ptr,
757 ql_nope_stride0,
758 ql_nope_stride1,
759 mqa_q_fp8_ptr,
760 mqa_q_fp8_stride0,
761 mqa_q_fp8_stride1,
762 q_scale_ptr,
763 QL_NOPE_DIM: tl.constexpr,
764 QL_NOPE_BLOCK: tl.constexpr,
765 # bf16 MQA query RoPE output (when QUANTIZE_MQA is False); the NoPE part is
766 # consumed directly from ql_nope, so only the RoPE'd q_pe is written here.
767 q_pe_out_ptr,
768 q_pe_out_stride0,
769 q_pe_out_stride1,
770 # Index weights
771 index_weights_ptr,
772 index_weights_stride,
773 index_weights_softmax_scale,
774 index_weights_head_scale,
775 index_weights_out_ptr,
776 index_weights_out_stride,
777 HAS_INDEXER: tl.constexpr,
778 INDEX_ROPE_INTERLEAVE: tl.constexpr,
779 QUANTIZE_MQA: tl.constexpr,
780 USE_PDL: tl.constexpr,
781 FP8_MAX: tl.constexpr,
782 USE_FNUZ: tl.constexpr,
783 ):
784 tok_idx = tl.program_id(0).to(tl.int64)
785 pid = tl.program_id(1)
786 head_idx = tl.program_id(2)
787 if USE_PDL:
788 tl.extra.cuda.gdc_wait()
789 tl.extra.cuda.gdc_launch_dependents()
790 fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
791
792 if pid == 2:
793 # ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16
794 # query path ql_nope is consumed as-is (no pack), so skip entirely.
795 if not QUANTIZE_MQA:
796 return
797 if 2 * head_idx >= NUM_Q_HEADS:
798 return
799
800 scale = tl.load(q_scale_ptr)
801 for local_head in range(2):
802 q_head_idx = head_idx * 2 + local_head
803 if q_head_idx < NUM_Q_HEADS:
804 ql_nope_off = tl.arange(0, QL_NOPE_BLOCK)
805 ql_nope_mask = ql_nope_off < QL_NOPE_DIM
806 ql_nope = tl.load(
807 ql_nope_ptr
808 + tok_idx * ql_nope_stride0
809 + q_head_idx * ql_nope_stride1
810 + ql_nope_off,
811 mask=ql_nope_mask,
812 ).to(tl.float32)
813 ql_nope_fp8 = (ql_nope / scale).to(fp8_dtype)
814 tl.store(
815 mqa_q_fp8_ptr
816 + tok_idx * mqa_q_fp8_stride0
817 + q_head_idx * mqa_q_fp8_stride1
818 + ql_nope_off,
819 ql_nope_fp8,
820 mask=ql_nope_mask,
821 )
822 return
823 elif pid == 0:
824 # q_pe RoPE + quantize + pack into the tail of mqa_q_fp8.
825 if 2 * head_idx >= NUM_Q_HEADS:
826 return
827
828 pos = tl.load(pos_ptr + tok_idx)
829 cos, sin = _get_cos_sin(
830 q_pe_cos_sin_ptr,
831 q_pe_cos_sin_stride,
832 pos,
833 Q_PE_HALF_ROT_DIM,
834 )
835
836 scale = tl.load(q_scale_ptr)
837 for local_head in range(2):
838 q_head_idx = head_idx * 2 + local_head
839 if q_head_idx < NUM_Q_HEADS:
840 rot_off = tl.arange(0, Q_PE_HALF_ROT_DIM)
841 x1 = tl.load(
842 q_pe_ptr
843 + tok_idx * q_pe_stride0
844 + q_head_idx * q_pe_stride1
845 + rot_off * 2,
846 ).to(tl.float32)
847 x2 = tl.load(
848 q_pe_ptr
849 + tok_idx * q_pe_stride0
850 + q_head_idx * q_pe_stride1
851 + rot_off * 2
852 + 1
853 ).to(tl.float32)
854 r1 = x1 * cos - x2 * sin
855 r2 = x2 * cos + x1 * sin
856 if QUANTIZE_MQA:
857 tl.store(
858 mqa_q_fp8_ptr
859 + tok_idx * mqa_q_fp8_stride0
860 + q_head_idx * mqa_q_fp8_stride1
861 + QL_NOPE_DIM
862 + rot_off * 2,
863 (r1 / scale).to(fp8_dtype),
864 )
865 tl.store(
866 mqa_q_fp8_ptr
867 + tok_idx * mqa_q_fp8_stride0
868 + q_head_idx * mqa_q_fp8_stride1
869 + QL_NOPE_DIM
870 + rot_off * 2
871 + 1,
872 (r2 / scale).to(fp8_dtype),
873 )
874 else:
875 # bf16 query: write the RoPE'd q_pe unquantized.
876 out_ty = q_pe_out_ptr.dtype.element_ty
877 q_pe_dst = (
878 q_pe_out_ptr
879 + tok_idx * q_pe_out_stride0
880 + q_head_idx * q_pe_out_stride1
881 )
882 tl.store(q_pe_dst + rot_off * 2, r1.to(out_ty))
883 tl.store(q_pe_dst + rot_off * 2 + 1, r2.to(out_ty))
884 return
885 elif pid == 1:
886 # Index Q RoPE + fp8 quant, all in registers. The roped bf16 index_q is
887 # never consumed (only the fp8 below is), so we avoid an in-place
888 # store-then-reload round-trip.
889 if not HAS_INDEXER:
890 return
891 if head_idx >= NUM_INDEX_Q_HEADS:
892 return
893
894 pos = tl.load(pos_ptr + tok_idx)
895 index_q_block = tl.arange(0, INDEX_Q_HEAD_DIM)
896 iq_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
897 index_q = tl.load(iq_base + index_q_block).to(tl.float32)
898
899 # RoPE in registers (interleaved for GLM-5.2, NeoX for DeepSeek-V3.2),
900 # gathering the rotation partner from the read-only input.
901 in_rope = index_q_block < 2 * INDEX_Q_HALF_ROT_DIM
902 if INDEX_ROPE_INTERLEAVE:
903 cos_idx = index_q_block // 2
904 partner_offs = tl.where(in_rope, index_q_block ^ 1, index_q_block)
905 sign = tl.where(index_q_block % 2 == 0, -1.0, 1.0)
906 else:
907 cos_idx = index_q_block % INDEX_Q_HALF_ROT_DIM
908 partner_offs = tl.where(
909 in_rope, index_q_block ^ INDEX_Q_HALF_ROT_DIM, index_q_block
910 )
911 sign = tl.where(index_q_block < INDEX_Q_HALF_ROT_DIM, -1.0, 1.0)
912 cos_full = tl.load(
913 index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + cos_idx,
914 mask=in_rope,
915 other=1.0,
916 ).to(tl.float32)
917 sin_full = tl.load(
918 index_q_cos_sin_ptr
919 + pos * index_q_cos_sin_stride
920 + INDEX_Q_HALF_ROT_DIM
921 + cos_idx,
922 mask=in_rope,
923 other=0.0,
924 ).to(tl.float32)
925 partner = tl.load(iq_base + partner_offs).to(tl.float32)
926 roped = index_q * cos_full + sign * partner * sin_full
927 index_q = tl.where(in_rope, roped, index_q)
928
929 # Index Q Quantize (from registers)
930 index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q, FP8_MAX, USE_FNUZ)
931 tl.store(
932 index_q_fp8_ptr
933 + tok_idx * index_q_fp8_stride0
934 + head_idx * index_q_fp8_stride1
935 + index_q_block,
936 index_q_fp8,
937 )
938
939 # Index weights update
940 index_weights = tl.load(
941 index_weights_ptr + tok_idx * index_weights_stride + head_idx
942 )
943 index_weights = index_weights.to(tl.float32)
944 index_weights *= index_q_scale
945 index_weights *= index_weights_softmax_scale
946 index_weights *= index_weights_head_scale
947 tl.store(
948 index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
949 index_weights,
950 )
951
952
953 def fused_q(
954 positions: torch.Tensor,
955 q_pe: torch.Tensor,
956 q_pe_cos_sin_cache: torch.Tensor,
957 index_q: torch.Tensor | None,
958 index_q_cos_sin_cache: torch.Tensor | None,
959 ql_nope: torch.Tensor,
960 q_scale: torch.Tensor,
961 # Index weights
962 index_weights: torch.Tensor | None,
963 index_weights_softmax_scale: float,
964 index_weights_head_scale: float,
965 has_indexer: bool = True,
966 index_rope_interleave: bool = False,
967 quantize_mqa: bool = True,
968 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
969 """Fuse the MQA-query and indexer-query RoPE/quantization.
970
971 Returns ``(index_q_fp8, index_weights_out, mqa_q)``. When ``quantize_mqa``
972 is True (FlashInfer sparse, fp8 query) ``mqa_q`` is a single fp8 tensor
973 packing ``[ql_nope; q_pe]``. When False (FlashMLA sparse, bf16 query) it is
974 the RoPE'd ``q_pe`` in bf16; the caller pairs it with ``ql_nope`` as the
975 ``(ql_nope, q_pe)`` tuple the backend expects.
976 """
977 assert positions.ndim == 1
978 assert positions.dtype == torch.int64
979 assert q_pe.ndim == 3
980 assert q_pe_cos_sin_cache.ndim == 2
981 assert ql_nope.ndim == 3
982 assert ql_nope.shape[:2] == q_pe.shape[:2]
983 assert q_scale.dtype == torch.float32 and q_scale.numel() == 1
984 num_tokens = positions.shape[0]
985 num_q_heads = q_pe.shape[1]
986 # Grid's 3rd dim must cover the MQA-pack heads (pid 0/2 iterate 2 heads
987 # each) and, when present, the indexer heads (pid 1).
988 mqa_grid_heads = (num_q_heads + 1) // 2
989 if not has_indexer:
990 # Shared layer: cached 1-element dummies; pid 1 skipped by HAS_INDEXER
991 # and never dereferences them.
992 index_q = _dummy((1, 1, 1), q_pe.dtype, q_pe.device)
993 index_q_cos_sin_cache = q_pe_cos_sin_cache
994 index_weights = _dummy((1, 1), torch.float32, q_pe.device)
995 assert index_q is not None and index_q.ndim == 3
996 assert index_q_cos_sin_cache is not None
997 assert index_weights is not None
998 num_index_q_heads = index_q.shape[1]
999 index_q_head_dim = index_q.shape[2]
1000 # fused_q is shared with the ROCm path, and the CuTeDSL module imports
1001 # cutlass at module scope, so only reach for it on CUDA.
1002 cutedsl_kernel: Callable[..., None] | None = None
1003 if current_platform.is_cuda():
1004 from vllm.models.deepseek_v32.nvidia.ops.fused_q_cutedsl import (
1005 fused_q_cutedsl,
1006 is_fused_q_cutedsl_supported,
1007 )
1008
1009 if is_fused_q_cutedsl_supported(
1010 q_pe,
1011 index_q,
1012 ql_nope,
1013 has_indexer=has_indexer,
1014 quantize_mqa=quantize_mqa,
1015 ):
1016 cutedsl_kernel = fused_q_cutedsl
1017 grid_heads = max(mqa_grid_heads, num_index_q_heads)
1018 if quantize_mqa:
1019 # fp8 path: pack [ql_nope; q_pe] into a single fp8 tensor.
1020 mqa_q_fp8 = torch.empty(
1021 q_pe.shape[0],
1022 q_pe.shape[1],
1023 ql_nope.shape[2] + q_pe.shape[2],
1024 dtype=_FP8_DTYPE,
1025 device=q_pe.device,
1026 )
1027 # Placeholder; pid 0 packs q_pe into mqa_q_fp8 instead.
1028 q_pe_out = mqa_q_fp8
1029 mqa_q = mqa_q_fp8
1030 else:
1031 # bf16 path: only the RoPE'd q_pe is produced; ql_nope used directly.
1032 q_pe_out = torch.empty_like(q_pe)
1033 mqa_q_fp8 = q_pe_out # unused placeholder for the fp8 pack pointer
1034 mqa_q = q_pe_out
1035
1036 index_q_fp8 = torch.empty_like(index_q, dtype=_FP8_DTYPE)
1037 index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
1038 if cutedsl_kernel is not None:
1039 cutedsl_kernel(
1040 positions,
1041 q_pe,
1042 q_pe_cos_sin_cache,
1043 ql_nope,
1044 q_scale,
1045 mqa_q,
1046 index_q,
1047 index_q_cos_sin_cache,
1048 index_weights,
1049 index_weights_softmax_scale,
1050 index_weights_head_scale,
1051 index_q_fp8,
1052 index_weights_out,
1053 has_indexer=has_indexer,
1054 index_rope_interleave=index_rope_interleave,
1055 quantize_mqa=quantize_mqa,
1056 )
1057 return index_q_fp8, index_weights_out, mqa_q
1058
1059 use_pdl = current_platform.is_arch_support_pdl()
1060 pdl_kwargs = (
1061 {"USE_PDL": use_pdl, "launch_pdl": True} if use_pdl else {"USE_PDL": use_pdl}
1062 )
1063 _fused_q_kernel[(num_tokens, 3, grid_heads)](
1064 positions,
1065 q_pe,
1066 q_pe.stride(0),
1067 q_pe.stride(1),
1068 num_q_heads,
1069 q_pe_cos_sin_cache,
1070 q_pe_cos_sin_cache.stride(0),
1071 q_pe_cos_sin_cache.shape[-1] // 2,
1072 index_q,
1073 index_q.stride(0),
1074 index_q.stride(1),
1075 num_index_q_heads,
1076 index_q_cos_sin_cache,
1077 index_q_cos_sin_cache.stride(0),
1078 index_q_cos_sin_cache.shape[-1] // 2,
1079 index_q_fp8,
1080 index_q_fp8.stride(0),
1081 index_q_fp8.stride(1),
1082 index_q_head_dim,
1083 ql_nope,
1084 ql_nope.stride(0),
1085 ql_nope.stride(1),
1086 mqa_q_fp8,
1087 mqa_q_fp8.stride(0),
1088 mqa_q_fp8.stride(1),
1089 q_scale,
1090 ql_nope.shape[2],
1091 triton.next_power_of_2(ql_nope.shape[2]),
1092 q_pe_out,
1093 q_pe_out.stride(0),
1094 q_pe_out.stride(1),
1095 index_weights,
1096 index_weights.stride(0),
1097 index_weights_softmax_scale,
1098 index_weights_head_scale,
1099 index_weights_out,
1100 index_weights_out.stride(0),
1101 HAS_INDEXER=has_indexer,
1102 INDEX_ROPE_INTERLEAVE=index_rope_interleave,
1103 QUANTIZE_MQA=quantize_mqa,
1104 FP8_MAX=_FP8_MAX,
1105 USE_FNUZ=_USE_FNUZ,
1106 **pdl_kwargs,
1107 # num_warps=1 is optimal here: each program is a single 128-element
1108 # rope+quant, so the kernel is program-count/occupancy bound, not
1109 # per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).
1110 num_warps=1,
1111 )
1112 return index_q_fp8, index_weights_out, mqa_q
1113
1114
1115 @triton.jit
1116 def _fused_eh_norm_kernel(
1117 pos_ptr,
1118 embeds_ptr,
1119 embeds_stride,
1120 prev_ptr,
1121 prev_stride,
1122 enorm_w_ptr,
1123 hnorm_w_ptr,
1124 eps,
1125 out_ptr,
1126 out_stride,
1127 H: tl.constexpr,
1128 BLOCK: tl.constexpr,
1129 ):
1130 """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm
1131 and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out``
1132 ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat."""
1133 tok = tl.program_id(0)
1134 off = tl.arange(0, BLOCK)
1135 mask = off < H
1136
1137 pos = tl.load(pos_ptr + tok)
1138 e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0)
1139 e = tl.where(pos == 0, 0.0, e.to(tl.float32))
1140 ew = tl.load(enorm_w_ptr + off, mask=mask)
1141 e_normed = _rms_norm(e, ew, eps, H)
1142 tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask)
1143
1144 p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0)
1145 hw = tl.load(hnorm_w_ptr + off, mask=mask)
1146 p_normed = _rms_norm(p, hw, eps, H)
1147 tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask)
1148
1149
1150 def fused_eh_norm(
1151 positions: torch.Tensor,
1152 inputs_embeds: torch.Tensor,
1153 previous_hidden: torch.Tensor,
1154 enorm_w: torch.Tensor,
1155 hnorm_w: torch.Tensor,
1156 eps: float,
1157 ) -> torch.Tensor:
1158 """Returns cat([enorm(masked embeds), hnorm(prev_hidden)]) -> [N, 2H]."""
1159 n, h = inputs_embeds.shape
1160 out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device)
1161 _fused_eh_norm_kernel[(n,)](
1162 positions,
1163 inputs_embeds,
1164 inputs_embeds.stride(0),
1165 previous_hidden,
1166 previous_hidden.stride(0),
1167 enorm_w,
1168 hnorm_w,
1169 eps,
1170 out,
1171 out.stride(0),
1172 h,
1173 triton.next_power_of_2(h),
1174 )
1175 return out