@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 4,725 lines 134 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from enum import IntEnum
5 from functools import cache
6 from typing import TYPE_CHECKING, Literal
7
8 import torch
9
10 import vllm.envs as envs
11 from vllm.logger import init_logger
12 from vllm.platforms import current_platform
13 from vllm.scalar_type import ScalarType
14 from vllm.utils.flashinfer import (
15 flashinfer_quant_nvfp4_8x4_sf_layout,
16 )
17 from vllm.utils.math_utils import cdiv
18
19 logger = init_logger(__name__)
20
21 current_platform.import_kernels()
22
23 if TYPE_CHECKING:
24
25 def register_fake(fn):
26 return lambda name: fn
27 else:
28 try:
29 from torch.library import register_fake
30 except ImportError:
31 from torch.library import impl_abstract as register_fake
32
33
34 # scaled_fp4_quant functional + out variant for torch.compile buffer management
35
36
37 def create_fp4_scale_tensor(
38 m: int,
39 n: int,
40 device: torch.device,
41 is_sf_swizzled_layout: bool,
42 ) -> torch.Tensor:
43 """Allocate the output scale tensor for scaled_fp4_quant.
44
45 When is_sf_swizzled_layout=True, we use rounded values to store the
46 swizzled scales. Due to the requirement of the Tensor Core, the minimum
47 tile is 128x4 for the scales. So, we first pad the scales to multiples
48 of 128 (rows) and 4 (cols). Then, the scales (in float8_e4m3fn) are
49 packed into an int32 for every 4 values. More:
50 https://docs.nvidia.com/cuda/parallel-thread-execution/
51 #tcgen05-mma-scale-factor-b-layout-4x
52 """
53 from vllm.utils.math_utils import round_up
54
55 block_size = 16
56 if is_sf_swizzled_layout:
57 rounded_m = round_up(m, 128)
58 scale_n = n // block_size
59 rounded_n = round_up(scale_n, 4)
60 # The NVFP4 quant kernel explicitly zeroes every padded scale entry,
61 # so no separate zero-initialization kernel is required here.
62 return torch.empty(
63 (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
64 )
65 else:
66 return torch.empty((m, n // block_size), device=device, dtype=torch.uint8)
67
68
69 def create_fp4_output_tensors(
70 m: int,
71 n: int,
72 device: torch.device,
73 is_sf_swizzled_layout: bool,
74 padded_n: int | None = None,
75 ) -> tuple[torch.Tensor, torch.Tensor]:
76 """Allocate both output tensors for scaled_fp4_quant:
77 (quantized_output, output_scale).
78
79 Must match the C++ scaled_fp4_quant_func allocation exactly when
80 ``padded_n`` is ``None``. When ``padded_n`` is provided, allocate a larger
81 packed-FP4 output/scale buffer so the quantization kernel can write
82 CUTLASS-compatible K padding directly
83 """
84 physical_n = padded_n if padded_n is not None else n
85 output = torch.empty((m, physical_n // 2), device=device, dtype=torch.uint8)
86 output_scale = create_fp4_scale_tensor(m, physical_n, device, is_sf_swizzled_layout)
87 return output, output_scale
88
89
90 if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"):
91
92 @register_fake("_C::scaled_fp4_quant")
93 def _scaled_fp4_quant_fake(
94 input: torch.Tensor,
95 input_scale: torch.Tensor,
96 is_sf_swizzled_layout: bool,
97 ) -> tuple[torch.Tensor, torch.Tensor]:
98 n = input.shape[-1]
99 m = input.numel() // n
100 return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout)
101
102
103 # page attention ops
104 def paged_attention_rocm(
105 out: torch.Tensor,
106 exp_sum: torch.Tensor,
107 max_logits: torch.Tensor,
108 tmp_out: torch.Tensor,
109 query: torch.Tensor,
110 key_cache: torch.Tensor,
111 value_cache: torch.Tensor,
112 num_kv_heads: int,
113 scale: float,
114 block_tables: torch.Tensor,
115 seq_lens: torch.Tensor,
116 query_start_loc: torch.Tensor | None,
117 block_size: int,
118 max_seq_len: int,
119 alibi_slopes: torch.Tensor | None,
120 kv_cache_dtype: str,
121 k_scale: torch.Tensor,
122 v_scale: torch.Tensor,
123 fp8_out_scale: torch.Tensor | None = None,
124 mfma_type: str = "fp8" if envs.VLLM_ROCM_FP8_MFMA_PAGE_ATTN else "f16",
125 ) -> None:
126 torch.ops._rocm_C.paged_attention(
127 out,
128 exp_sum,
129 max_logits,
130 tmp_out,
131 query,
132 key_cache,
133 value_cache,
134 num_kv_heads,
135 scale,
136 block_tables,
137 seq_lens,
138 query_start_loc,
139 block_size,
140 max_seq_len,
141 alibi_slopes,
142 kv_cache_dtype,
143 k_scale,
144 v_scale,
145 fp8_out_scale,
146 mfma_type,
147 )
148
149
150 def mla_decode_kvcache_cpu(
151 out: torch.Tensor,
152 query: torch.Tensor,
153 kv_cache: torch.Tensor,
154 scale: float,
155 block_tables: torch.Tensor,
156 seq_lens: torch.Tensor,
157 ) -> None:
158 torch.ops._C.mla_decode_kvcache(out, query, kv_cache, scale, block_tables, seq_lens)
159
160
161 def gather_mla_context_cache_cpu(
162 *,
163 src_cache: torch.Tensor,
164 dst: torch.Tensor,
165 block_table: torch.Tensor,
166 starts: torch.Tensor,
167 cu_seq_lens: torch.Tensor,
168 ) -> None:
169 page_size = src_cache.shape[1]
170 flat_cache = src_cache.view(-1, src_cache.shape[-1])
171 seq_lens = cu_seq_lens[1:] - cu_seq_lens[:-1]
172
173 out_start = 0
174 for req_idx, (start, seq_len) in enumerate(zip(starts.tolist(), seq_lens.tolist())):
175 if seq_len <= 0:
176 continue
177 positions = torch.arange(start, start + seq_len, device=block_table.device)
178 block_ids = block_table[req_idx, positions // page_size].to(torch.long)
179 offsets = positions % page_size
180 slots = block_ids * page_size + offsets
181 out_end = out_start + seq_len
182 dst[out_start:out_end].copy_(flat_cache[slots])
183 out_start = out_end
184
185
186 # merge attn states ops
187 def merge_attn_states(
188 output: torch.Tensor,
189 prefix_output: torch.Tensor,
190 prefix_lse: torch.Tensor,
191 suffix_output: torch.Tensor,
192 suffix_lse: torch.Tensor,
193 output_lse: torch.Tensor | None = None,
194 prefill_tokens_with_context: int | None = None,
195 output_scale: torch.Tensor | None = None,
196 ) -> None:
197 torch.ops._C.merge_attn_states(
198 output,
199 output_lse,
200 prefix_output,
201 prefix_lse,
202 suffix_output,
203 suffix_lse,
204 prefill_tokens_with_context,
205 output_scale,
206 )
207
208
209 # pos encoding ops
210 def rotary_embedding(
211 positions: torch.Tensor,
212 query: torch.Tensor,
213 key: torch.Tensor | None,
214 head_size: int,
215 cos_sin_cache: torch.Tensor,
216 is_neox: bool,
217 rope_dim_offset: int = 0,
218 inverse: bool = False,
219 ) -> None:
220 if rope_dim_offset == 0 and not inverse:
221 torch.ops._C.rotary_embedding(
222 positions, query, key, head_size, cos_sin_cache, is_neox
223 )
224 else:
225 torch.ops._C.rotary_embedding(
226 positions,
227 query,
228 key,
229 head_size,
230 cos_sin_cache,
231 is_neox,
232 rope_dim_offset,
233 inverse,
234 )
235
236
237 # layer norm ops
238 def rms_norm(
239 out: torch.Tensor,
240 input: torch.Tensor,
241 weight: torch.Tensor | None,
242 epsilon: float,
243 ) -> None:
244 torch.ops._C.rms_norm(out, input, weight, epsilon)
245
246
247 # Fused vocab-parallel embedding lookup
248 # (see csrc/.../vocab_parallel_embedding_kernels.cu).
249 def vocab_parallel_embedding(
250 input_ids: torch.Tensor,
251 weight: torch.Tensor,
252 org_vocab_start_index: int,
253 org_vocab_end_index: int,
254 num_org_vocab_padding: int,
255 added_vocab_start_index: int,
256 added_vocab_end_index: int,
257 ) -> torch.Tensor:
258 """Gather the embedding rows owned by this rank, zeros for the others.
259
260 Fuses the range mask, id shift, gather and output masking of the TP > 1
261 embedding path into a single kernel. Out-of-shard tokens produce zero rows,
262 so an all-reduce over the ranks yields the full embedding.
263
264 Args:
265 input_ids: ``[num_tokens]`` int32 or int64 token ids.
266 weight: ``[num_embeddings_per_partition, embedding_dim]`` shard.
267 org_vocab_start_index: First original vocab id owned by this rank.
268 org_vocab_end_index: One past the last original vocab id owned by
269 this rank.
270 num_org_vocab_padding: Padding rows between this rank's original and
271 added vocab shards.
272 added_vocab_start_index: First added vocab id owned by this rank.
273 added_vocab_end_index: One past the last added vocab id owned by
274 this rank.
275
276 Returns:
277 ``[num_tokens, embedding_dim]`` partial embeddings.
278
279 """
280 out = torch.empty(
281 input_ids.shape[0],
282 weight.shape[1],
283 dtype=weight.dtype,
284 device=weight.device,
285 )
286 torch.ops._C.vocab_parallel_embedding(
287 out,
288 input_ids,
289 weight,
290 org_vocab_start_index,
291 org_vocab_end_index,
292 num_org_vocab_padding,
293 added_vocab_start_index,
294 added_vocab_end_index,
295 )
296 return out
297
298
299 # LongCat n-gram embedding index kernel (see csrc/.../ngram_embedding_kernels.cu).
300 def ngram_compute_n_gram_ids(
301 ne_n: int,
302 ne_k: int,
303 ne_weights: torch.Tensor,
304 ne_mods: torch.Tensor,
305 exclusive_ne_embedder_size_sums: torch.Tensor,
306 exclusive_req_len_sums: torch.Tensor,
307 ne_token_table: torch.Tensor,
308 row_indices: torch.Tensor,
309 column_starts: torch.Tensor,
310 n_gram_ids: torch.Tensor,
311 ) -> None:
312 """Compute concatenated (offset) n-gram ids for a ragged prefill batch.
313
314 Writes ``n_gram_ids`` of shape ``[token_num, (ne_n-1)*ne_k]``.
315 """
316 torch.ops._C.ngram_compute_n_gram_ids(
317 ne_n,
318 ne_k,
319 ne_weights,
320 ne_mods,
321 exclusive_ne_embedder_size_sums,
322 exclusive_req_len_sums,
323 ne_token_table,
324 row_indices,
325 column_starts,
326 n_gram_ids,
327 )
328
329
330 def fused_add_rms_norm(
331 input: torch.Tensor,
332 residual: torch.Tensor,
333 weight: torch.Tensor | None,
334 epsilon: float,
335 ) -> None:
336 # Note: this func is batch invariant
337 torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
338
339
340 def fused_qk_norm_rope(
341 qkv: torch.Tensor,
342 num_heads_q: int,
343 num_heads_k: int,
344 num_heads_v: int,
345 head_dim: int,
346 eps: float,
347 q_weight: torch.Tensor,
348 k_weight: torch.Tensor,
349 cos_sin_cache: torch.Tensor,
350 is_neox: bool,
351 position_ids: torch.Tensor,
352 forced_token_heads_per_warp: int = -1,
353 ) -> None:
354 torch.ops._C.fused_qk_norm_rope(
355 qkv,
356 num_heads_q,
357 num_heads_k,
358 num_heads_v,
359 head_dim,
360 eps,
361 q_weight,
362 k_weight,
363 cos_sin_cache,
364 is_neox,
365 position_ids,
366 forced_token_heads_per_warp,
367 )
368
369
370 def apply_repetition_penalties_torch(
371 logits: torch.Tensor,
372 prompt_mask: torch.Tensor,
373 output_mask: torch.Tensor,
374 repetition_penalties: torch.Tensor,
375 ) -> None:
376 repetition_penalties = repetition_penalties.unsqueeze(dim=1).repeat(
377 1, logits.size(1)
378 )
379 # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
380 penalties = torch.where(prompt_mask | output_mask, repetition_penalties, 1.0)
381 # If logits are positive, divide by penalty, otherwise multiply by penalty.
382 scaling = torch.where(logits > 0, 1.0 / penalties, penalties)
383 logits *= scaling
384
385
386 def apply_repetition_penalties_cuda(
387 logits: torch.Tensor,
388 prompt_mask: torch.Tensor,
389 output_mask: torch.Tensor,
390 repetition_penalties: torch.Tensor,
391 ) -> None:
392 torch.ops._C.apply_repetition_penalties_(
393 logits, prompt_mask, output_mask, repetition_penalties
394 )
395
396
397 def apply_repetition_penalties(
398 logits: torch.Tensor,
399 prompt_mask: torch.Tensor,
400 output_mask: torch.Tensor,
401 repetition_penalties: torch.Tensor,
402 ) -> None:
403 """Apply repetition penalties to logits in-place.
404
405 Args:
406 logits: The logits tensor of shape [num_seqs, vocab_size].
407 prompt_mask: A boolean tensor indicating which tokens appear in the prompt.
408 output_mask: A boolean tensor indicating which tokens appear in the output.
409 repetition_penalties: The repetition penalties of shape (num_seqs, ).
410
411 """
412 if logits.is_cuda and logits.is_contiguous():
413 apply_repetition_penalties_cuda(
414 logits, prompt_mask, output_mask, repetition_penalties
415 )
416 else:
417 apply_repetition_penalties_torch(
418 logits, prompt_mask, output_mask, repetition_penalties
419 )
420
421
422 # fused quant layer norm ops
423 def rms_norm_dynamic_per_token_quant(
424 input: torch.Tensor,
425 weight: torch.Tensor,
426 epsilon: float,
427 quant_dtype: torch.dtype,
428 scale_ub: torch.Tensor | None = None,
429 residual: torch.Tensor | None = None,
430 ) -> tuple[torch.Tensor, torch.Tensor]:
431 output = torch.empty(input.shape, dtype=quant_dtype, device=input.device)
432 scales = torch.empty(
433 (input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
434 )
435
436 torch.ops._C.rms_norm_dynamic_per_token_quant(
437 output, input, weight, scales, epsilon, scale_ub, residual
438 )
439 return output, scales
440
441
442 # fused quant layer norm ops blocked
443 def rms_norm_per_block_quant(
444 input: torch.Tensor,
445 weight: torch.Tensor,
446 epsilon: float,
447 quant_dtype: torch.dtype,
448 group_size: list[int],
449 scale_ub: torch.Tensor | None = None,
450 residual: torch.Tensor | None = None,
451 is_scale_transposed: bool = False,
452 tma_alignment: int = 0,
453 ) -> tuple[torch.Tensor, torch.Tensor]:
454 assert len(group_size) == 2
455 output = torch.empty(input.shape, dtype=quant_dtype, device=input.device)
456 if is_scale_transposed:
457 if tma_alignment == 0:
458 scales = torch.empty(
459 (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
460 device=input.device,
461 dtype=torch.float32,
462 ).transpose(0, 1)
463 else:
464 m = input.shape[-2]
465 sf_k = input.shape[-1] // group_size[1]
466 tma_aligned_m = (m + tma_alignment - 1) // tma_alignment * tma_alignment
467 shape = input.shape[:-2] + (m, sf_k)
468 stride = (
469 (1, tma_aligned_m)
470 if input.dim() == 2
471 else (tma_aligned_m * sf_k, 1, tma_aligned_m)
472 )
473 scales = torch.empty_strided(
474 shape, stride, device=input.device, dtype=torch.float32
475 )
476 else:
477 scales = torch.empty(
478 (input.numel() // input.shape[-1], input.shape[-1] // group_size[1]),
479 device=input.device,
480 dtype=torch.float32,
481 )
482
483 assert tma_alignment in [0, 4], "Expected TMA alignment 0 or 4, but got " + str(
484 tma_alignment
485 )
486
487 torch.ops._C.rms_norm_per_block_quant(
488 output,
489 input,
490 weight,
491 scales,
492 epsilon,
493 scale_ub,
494 residual,
495 group_size[1],
496 is_scale_transposed,
497 )
498 return output, scales
499
500
501 # fused silu_and_mul + block quant
502 def silu_and_mul_per_block_quant(
503 input: torch.Tensor,
504 group_size: int, # Changed from list[int]
505 quant_dtype: torch.dtype,
506 scale_ub: torch.Tensor | None = None,
507 is_scale_transposed: bool = False,
508 ) -> tuple[torch.Tensor, torch.Tensor]:
509 assert input.ndim == 2, f"input must be 2D [batch, hidden*2], got {input.shape}"
510 assert input.shape[-1] % 2 == 0, (
511 f"input last dim must be even (gate||up layout), got {input.shape[-1]}"
512 )
513
514 # Output is half the width of input (after silu_and_mul)
515 num_tokens = input.shape[0]
516 hidden_size = input.shape[-1] // 2 # Divide by 2 because input is [gate || up]
517
518 # Allocate output tensor (FP8 or INT8)
519 output = torch.empty(
520 (num_tokens, hidden_size), device=input.device, dtype=quant_dtype
521 )
522
523 # Allocate scales tensor
524 num_groups = hidden_size // group_size # Directly use group_size
525 if is_scale_transposed:
526 scales = torch.empty(
527 (num_groups, num_tokens),
528 device=input.device,
529 dtype=torch.float32,
530 ).t()
531 else:
532 scales = torch.empty(
533 (num_tokens, num_groups),
534 device=input.device,
535 dtype=torch.float32,
536 )
537
538 # Call the C++ kernel
539 torch.ops._C.silu_and_mul_per_block_quant(
540 output,
541 input,
542 scales,
543 group_size, # Pass directly as int
544 scale_ub,
545 is_scale_transposed,
546 )
547
548 return output, scales
549
550
551 # quantization ops
552 # awq
553 def awq_dequantize(
554 qweight: torch.Tensor,
555 scales: torch.Tensor,
556 zeros: torch.Tensor,
557 split_k_iters: int,
558 thx: int,
559 thy: int,
560 ) -> torch.Tensor:
561 if envs.VLLM_USE_TRITON_AWQ:
562 from vllm.model_executor.layers.quantization.awq_triton import (
563 awq_dequantize_triton,
564 )
565
566 return awq_dequantize_triton(qweight, scales, zeros)
567 return torch.ops._C.awq_dequantize(qweight, scales, zeros, split_k_iters, thx, thy)
568
569
570 if hasattr(torch.ops._C, "awq_dequantize"):
571
572 @register_fake("_C::awq_dequantize")
573 def _awq_dequantize_fake(
574 qweight: torch.Tensor,
575 scales: torch.Tensor,
576 zeros: torch.Tensor,
577 split_k_iters: torch.SymInt,
578 thx: int,
579 thy: int,
580 ) -> torch.Tensor:
581 in_c = qweight.size(0)
582 qout_c = qweight.size(1)
583 out_c = qout_c * 8
584 return torch.empty((in_c, out_c), dtype=scales.dtype, device=scales.device)
585
586
587 def awq_gemm(
588 input: torch.Tensor,
589 qweight: torch.Tensor,
590 scales: torch.Tensor,
591 qzeros: torch.Tensor,
592 split_k_iters: int,
593 ) -> torch.Tensor:
594 if envs.VLLM_USE_TRITON_AWQ:
595 from vllm.model_executor.layers.quantization.awq_triton import awq_gemm_triton
596
597 return awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters)
598 return torch.ops._C.awq_gemm(input, qweight, scales, qzeros, split_k_iters)
599
600
601 if hasattr(torch.ops._C, "awq_gemm"):
602
603 @register_fake("_C::awq_gemm")
604 def _awq_gemm_fake(
605 input: torch.Tensor,
606 qweight: torch.Tensor,
607 scales: torch.Tensor,
608 qzeros: torch.Tensor,
609 split_k_iters: torch.SymInt,
610 ) -> torch.Tensor:
611 num_in_feats = input.size(0)
612 return torch.empty(
613 (split_k_iters, num_in_feats, qweight.size(1) * 8),
614 dtype=input.dtype,
615 device=input.device,
616 ).sum(0)
617
618
619 # gptq
620 def gptq_gemm(
621 a: torch.Tensor,
622 b_q_weight: torch.Tensor,
623 b_gptq_qzeros: torch.Tensor,
624 b_gptq_scales: torch.Tensor,
625 use_exllama: bool,
626 use_v2_format: bool,
627 bit: int,
628 ) -> torch.Tensor:
629 return torch.ops._C.gptq_gemm(
630 a,
631 b_q_weight,
632 b_gptq_qzeros,
633 b_gptq_scales,
634 use_exllama,
635 use_v2_format,
636 bit,
637 )
638
639
640 if hasattr(torch.ops._C, "gptq_gemm"):
641
642 @register_fake("_C::gptq_gemm")
643 def _gptq_gemm_fake(
644 a: torch.Tensor,
645 b_q_weight: torch.Tensor,
646 b_gptq_qzeros: torch.Tensor,
647 b_gptq_scales: torch.Tensor,
648 use_exllama: bool,
649 use_v2_format: bool,
650 bit: int,
651 ) -> torch.Tensor:
652 return torch.empty(
653 (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
654 )
655
656
657 def gptq_shuffle(q_weight: torch.Tensor, bit: int) -> None:
658 torch.ops._C.gptq_shuffle(q_weight, bit)
659
660
661 def gptq_gemm_rdna3(
662 a: torch.Tensor,
663 b_q_weight: torch.Tensor,
664 b_qzeros: torch.Tensor,
665 b_scales: torch.Tensor,
666 use_v2_format: bool,
667 ) -> torch.Tensor:
668 return torch.ops._rocm_C.gptq_gemm_rdna3(
669 a, b_q_weight, b_qzeros, b_scales, use_v2_format
670 )
671
672
673 if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3"):
674
675 @register_fake("_rocm_C::gptq_gemm_rdna3")
676 def _gptq_gemm_rdna3_fake(
677 a: torch.Tensor,
678 b_q_weight: torch.Tensor,
679 b_qzeros: torch.Tensor,
680 b_scales: torch.Tensor,
681 use_v2_format: bool,
682 ) -> torch.Tensor:
683 return torch.empty(
684 (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
685 )
686
687
688 if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3_wmma"):
689
690 @register_fake("_rocm_C::gptq_gemm_rdna3_wmma")
691 def _gptq_gemm_rdna3_wmma_fake(
692 a: torch.Tensor,
693 b_q_weight: torch.Tensor,
694 b_qzeros: torch.Tensor,
695 b_scales: torch.Tensor,
696 use_v2_format: bool,
697 ) -> torch.Tensor:
698 return torch.empty(
699 (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
700 )
701
702
703 def moe_gptq_gemm_rdna3(
704 a: torch.Tensor,
705 c: torch.Tensor,
706 b_q_weight: torch.Tensor,
707 b_scales: torch.Tensor,
708 b_qzeros: torch.Tensor,
709 topk_weights: torch.Tensor,
710 sorted_token_ids: torch.Tensor,
711 expert_ids: torch.Tensor,
712 num_tokens_post_padded: torch.Tensor,
713 top_k: int,
714 block_size_m: int,
715 mul_topk_weight: bool,
716 output_topk: int = 0,
717 ) -> None:
718 torch.ops._rocm_C.moe_gptq_gemm_rdna3(
719 a,
720 c,
721 b_q_weight,
722 b_scales,
723 b_qzeros,
724 topk_weights,
725 sorted_token_ids,
726 expert_ids,
727 num_tokens_post_padded,
728 top_k,
729 block_size_m,
730 mul_topk_weight,
731 output_topk,
732 )
733
734
735 # cutlass
736 def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool:
737 return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability)
738
739
740 def mxfp4_experts_quant_supported(cuda_device_capability: int) -> bool:
741 try:
742 return torch.ops._C.mxfp4_experts_quant_supported(cuda_device_capability)
743 except AttributeError:
744 # Return False on builds where the CUDA helper is not available.
745 return False
746
747
748 def cutlass_scaled_fp4_mm(
749 a: torch.Tensor,
750 b: torch.Tensor,
751 block_scale_a: torch.Tensor,
752 block_scale_b: torch.Tensor,
753 alpha: torch.Tensor,
754 out_dtype: torch.dtype,
755 ) -> torch.Tensor:
756 assert a.ndim == 2 and b.ndim == 2
757 m, n = a.shape[0], b.shape[0]
758 out = torch.empty((m, n), dtype=out_dtype, device=a.device)
759 torch.ops._C.cutlass_scaled_fp4_mm(out, a, b, block_scale_a, block_scale_b, alpha)
760 return out
761
762
763 def cutlass_scaled_mm_supports_fp8(cuda_device_capability: int) -> bool:
764 return torch.ops._C.cutlass_scaled_mm_supports_fp8(cuda_device_capability)
765
766
767 def cutlass_scaled_mm_supports_block_fp8(cuda_device_capability: int) -> bool:
768 return torch.ops._C.cutlass_scaled_mm_supports_block_fp8(cuda_device_capability)
769
770
771 def cutlass_scaled_mm(
772 a: torch.Tensor,
773 b: torch.Tensor,
774 scale_a: torch.Tensor,
775 scale_b: torch.Tensor,
776 out_dtype: torch.dtype,
777 bias: torch.Tensor | None = None,
778 ) -> torch.Tensor:
779 """`cutlass_scaled_mm` implements a fused version of
780 `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
781 where scale_a * a and scale_b * b are implemented using numpy-style
782 broadcasting.
783
784 In order to support blockwise scaling like found in DeepSeek V3 we also
785 support extended "group" broadcast rules. We extend the numpy-style
786 broadcasting rules with the following rule:
787 "if the extent of a dimension in the source shape is between 1 and
788 corresponding extent in the target shape we repeat each element along
789 that dimension src_shape[dim] // target_shape[dim] times consecutively"
790 example if we have:
791 a = [[1, 2], and target_shape = (2, 4)
792 [3, 4]]
793 then we would expand a to:
794 a = [[1, 1, 2, 2],
795 [3, 3, 4, 4]]
796 currently we only support the case:
797 scale_a.shape * [1, 128] == a.shape
798 scale_b.shape * [128, 128] == b.shape
799 """
800 assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
801 assert bias is None or bias.numel() == b.shape[1] and bias.dtype == out_dtype
802
803 # Massage the input to be 2D
804 target_shape = (*a.shape[:-1], b.shape[1])
805 a = a.view(-1, a.shape[-1])
806
807 cutlass_compatible_b = b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0
808 if current_platform.is_rocm() or not cutlass_compatible_b:
809 from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import ( # noqa
810 triton_scaled_mm,
811 )
812
813 out = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
814 else:
815 out = torch.empty((a.shape[0], b.shape[1]), dtype=out_dtype, device=a.device)
816 torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias)
817
818 return out.view(*target_shape)
819
820
821 def cutlass_scaled_mm_azp(
822 a: torch.Tensor,
823 b: torch.Tensor,
824 scale_a: torch.Tensor,
825 scale_b: torch.Tensor,
826 out_dtype: torch.dtype,
827 azp_adj: torch.Tensor,
828 azp: torch.Tensor | None = None,
829 bias: torch.Tensor | None = None,
830 ) -> torch.Tensor:
831 """Args:
832 azp_adj: In the per-tensor case, this should include the azp.
833 Always per-channel.
834 azp: Only set in the per-token case. Per-token if set.
835
836 """
837 assert b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0
838 assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
839 assert bias is None or bias.numel() == b.shape[1] and bias.dtype == out_dtype
840
841 # Massage the input to be 2D
842 target_shape = (*a.shape[:-1], b.shape[1])
843 a = a.view(-1, a.shape[-1])
844 assert azp is None or azp.numel() == a.shape[0]
845
846 out = torch.empty((a.shape[0], b.shape[1]), dtype=out_dtype, device=a.device)
847 torch.ops._C.cutlass_scaled_mm_azp(out, a, b, scale_a, scale_b, azp_adj, azp, bias)
848 return out.view(*target_shape)
849
850
851 def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
852 if cuda_device_capability < 90 or cuda_device_capability >= 110:
853 return False
854 try:
855 return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability)
856 except AttributeError:
857 # Return False on non-CUDA platforms where it is not available
858 return False
859
860
861 def get_cutlass_moe_mm_data(
862 topk_ids: torch.Tensor,
863 expert_offsets: torch.Tensor,
864 problem_sizes1: torch.Tensor,
865 problem_sizes2: torch.Tensor,
866 input_permutation: torch.Tensor,
867 output_permutation: torch.Tensor,
868 num_experts: int,
869 n: int,
870 k: int,
871 blockscale_offsets: torch.Tensor | None = None,
872 is_gated: bool = True,
873 ):
874 """Prepare data necessary to perform CUTLASS grouped matrix multiplications
875 used in CUTLASS-based fused MoE.
876
877 The function takes in topk_ids (token-expert mapping) and uses it to
878 compute:
879 - expert_offsets: Indices that mark at which token index each expert begins
880 its computation after the input is sorted with
881 input_permutation. The number of tokens computed with
882 expert E is expert_offsets[E + 1] - expert_offsets[E]
883 - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's
884 multiplication in two grouped MMs used in
885 the fused MoE operation.
886 - input_permutation: Permutation that must be used to shuffle the input
887 before executing the MMs.
888 - output_permutation: Permutation that must be used to shuffle the output
889 after executing the MMs.
890 - blockscale_offsets: Optional argument passed for fp4 moe. Indices that
891 mark at which block scale index each expert begins
892 its computation. The number of block scale rows
893 computed with expert E is blockscale_offsets[E + 1] -
894 blockscale_offsets[E]
895 - is_gated: Whether the activation is gated (gate + up). When True, the
896 first GEMM N dimension is 2*n; when False, it is n.
897 """
898 return torch.ops._C.get_cutlass_moe_mm_data(
899 topk_ids,
900 expert_offsets,
901 problem_sizes1,
902 problem_sizes2,
903 input_permutation,
904 output_permutation,
905 num_experts,
906 n,
907 k,
908 blockscale_offsets,
909 is_gated,
910 )
911
912
913 def get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
914 expert_first_token_offset: torch.Tensor,
915 problem_sizes1: torch.Tensor,
916 problem_sizes2: torch.Tensor,
917 n: int,
918 k: int,
919 swap_ab: bool,
920 ):
921 """Compute per-expert (M, N, K) problem sizes from expert_first_token_offset"""
922 return torch.ops._C.get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
923 expert_first_token_offset,
924 problem_sizes1,
925 problem_sizes2,
926 n,
927 k,
928 swap_ab,
929 )
930
931
932 def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor):
933 """Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor.
934 This is used in MoE to permute the input tensor before performing grouped matrix multiplications.
935 """
936 num_tokens_permuted = dst2src_map.shape[0]
937 output_tensor = torch.empty(
938 (num_tokens_permuted, input_tensor.shape[1]),
939 device=input_tensor.device,
940 dtype=input_tensor.dtype,
941 )
942 torch.ops._moe_C.shuffle_rows(input_tensor, dst2src_map, output_tensor)
943 return output_tensor
944
945
946 def get_cutlass_batched_moe_mm_data(
947 expert_offsets: torch.Tensor,
948 problem_sizes1: torch.Tensor,
949 problem_sizes2: torch.Tensor,
950 expert_num_tokens: torch.Tensor,
951 num_local_experts: int,
952 padded_m: int,
953 n: int,
954 k: int,
955 ):
956 """Prepare data necessary to perform CUTLASS grouped matrix multiplications
957 used in CUTLASS-based fused MoE.
958
959 The function takes in expert_num_tokens (token count per expert) and
960 non_zero_expert_idxs (consecutive indices of experts with non-zero token
961 counts) and uses them to compute:
962 - expert_offsets: Indices that mark at which token index each expert begins
963 its computation.
964 - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's
965 multiplication in two grouped MMs used in
966 the fused MoE operation.
967 """
968 return torch.ops._C.get_cutlass_batched_moe_mm_data(
969 expert_offsets,
970 problem_sizes1,
971 problem_sizes2,
972 expert_num_tokens,
973 num_local_experts,
974 padded_m,
975 n,
976 k,
977 )
978
979
980 def cutlass_moe_mm(
981 out_tensors: torch.Tensor,
982 a_tensors: torch.Tensor,
983 b_tensors: torch.Tensor,
984 a_scales: torch.Tensor,
985 b_scales: torch.Tensor,
986 expert_offsets: torch.Tensor,
987 problem_sizes: torch.Tensor,
988 a_strides: torch.Tensor,
989 b_strides: torch.Tensor,
990 c_strides: torch.Tensor,
991 per_act_token: bool,
992 per_out_ch: bool,
993 ):
994 """A single grouped matrix multiplication used in CUTLASS-based fused MoE.
995 The function executes fp8-quantized OUT = AB matrix multiplication.
996
997 - expert_offsets: Indices that mark at which token index each expert begins
998 its computation. The number of tokens computed with
999 expert E is expert_offsets[E + 1] - expert_offsets[E]
1000 - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
1001 MMs used in the fused MoE operation.
1002 - a/b/c_strides: The data strides passed to grouped matrix multiplication.
1003 """
1004 return torch.ops._C.cutlass_moe_mm(
1005 out_tensors,
1006 a_tensors,
1007 b_tensors,
1008 a_scales,
1009 b_scales,
1010 expert_offsets,
1011 problem_sizes,
1012 a_strides,
1013 b_strides,
1014 c_strides,
1015 per_act_token,
1016 per_out_ch,
1017 )
1018
1019
1020 def cutlass_fp4_moe_mm(
1021 out_tensors: torch.Tensor,
1022 a_tensors: torch.Tensor,
1023 b_tensors: torch.Tensor,
1024 a_scales: torch.Tensor,
1025 b_scales: torch.Tensor,
1026 alphas: torch.Tensor,
1027 problem_sizes: torch.Tensor,
1028 expert_offsets: torch.Tensor,
1029 sf_offsets: torch.Tensor,
1030 ):
1031 """An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs
1032 the gemms for each combination based on the specified problem sizes.
1033
1034 This is used as the MoE gemm during NVFP4 Quantized MoERunner forward.
1035 - a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized
1036 input and expert weights.
1037 - a_/b_scales: The blockscales in FP8-E4M3 precision
1038 - expert_offsets/sf_offsets: Indices that mark at which token index
1039 each expert begins its computation. The number of tokens
1040 computed with expert E is expert_offsets[E + 1] -
1041 expert_offsets[E] And the sf_size per expert is
1042 sf_offset[E+1] - sf_offset[E]
1043 - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
1044 MMs used in the fused MoE operation.
1045 """
1046 return torch.ops._C.cutlass_fp4_group_mm(
1047 out_tensors,
1048 a_tensors,
1049 b_tensors,
1050 a_scales,
1051 b_scales,
1052 alphas,
1053 problem_sizes,
1054 expert_offsets,
1055 sf_offsets,
1056 )
1057
1058
1059 def cutlass_mxfp4_moe_mm(
1060 out_tensors: torch.Tensor,
1061 a_tensors: torch.Tensor,
1062 b_tensors: torch.Tensor,
1063 a_scales: torch.Tensor,
1064 b_scales: torch.Tensor,
1065 problem_sizes: torch.Tensor,
1066 expert_offsets: torch.Tensor,
1067 sf_offsets: torch.Tensor,
1068 ):
1069 """An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
1070
1071 Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
1072 - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
1073 - a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout)
1074 - Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales).
1075 - expert_offsets/sf_offsets: expert boundary indices
1076 - problem_sizes: (num_experts, 3) with (M, N, K) per expert
1077 """
1078 return torch.ops._C.cutlass_mxfp4_group_mm(
1079 out_tensors,
1080 a_tensors,
1081 b_tensors,
1082 a_scales,
1083 b_scales,
1084 problem_sizes,
1085 expert_offsets,
1086 sf_offsets,
1087 )
1088
1089
1090 # gptq_marlin
1091 def gptq_marlin_repack(
1092 b_q_weight: torch.Tensor,
1093 size_k: int,
1094 size_n: int,
1095 num_bits: int,
1096 is_a_8bit: bool = False,
1097 ) -> torch.Tensor:
1098 return torch.ops._C.gptq_marlin_repack(
1099 b_q_weight, size_k, size_n, num_bits, is_a_8bit
1100 )
1101
1102
1103 if hasattr(torch.ops._C, "gptq_marlin_repack"):
1104
1105 @register_fake("_C::gptq_marlin_repack")
1106 def _gptq_marlin_repack_fake(
1107 b_q_weight: torch.Tensor,
1108 size_k: torch.SymInt,
1109 size_n: torch.SymInt,
1110 num_bits: int,
1111 is_a_8bit: bool = False,
1112 ) -> torch.Tensor:
1113 pack_factor = 32 // num_bits
1114 marlin_tile_size = 16
1115 return torch.empty(
1116 (size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor),
1117 dtype=b_q_weight.dtype,
1118 device=b_q_weight.device,
1119 )
1120
1121
1122 # awq_marlin
1123 def awq_marlin_repack(
1124 b_q_weight: torch.Tensor,
1125 size_k: int,
1126 size_n: int,
1127 num_bits: int,
1128 is_a_8bit: bool = False,
1129 ) -> torch.Tensor:
1130 return torch.ops._C.awq_marlin_repack(
1131 b_q_weight, size_k, size_n, num_bits, is_a_8bit
1132 )
1133
1134
1135 if hasattr(torch.ops._C, "awq_marlin_repack"):
1136
1137 @register_fake("_C::awq_marlin_repack")
1138 def _awq_marlin_repack_fake(
1139 b_q_weight: torch.Tensor,
1140 size_k: torch.SymInt,
1141 size_n: torch.SymInt,
1142 num_bits: int,
1143 is_a_8bit: bool = False,
1144 ) -> torch.Tensor:
1145 pack_factor = 32 // num_bits
1146 marlin_tile_size = 16
1147 return torch.empty(
1148 (size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor),
1149 dtype=b_q_weight.dtype,
1150 device=b_q_weight.device,
1151 )
1152
1153
1154 def gptq_marlin_moe_repack(
1155 b_q_weight: torch.Tensor,
1156 size_k: int,
1157 size_n: int,
1158 num_bits: int,
1159 is_a_8bit: bool = False,
1160 ) -> torch.Tensor:
1161 num_experts = b_q_weight.shape[0]
1162 assert size_k % 16 == 0
1163 output = torch.empty(
1164 (num_experts, size_k // 16, size_n * (num_bits // 2)),
1165 device=b_q_weight.device,
1166 dtype=b_q_weight.dtype,
1167 )
1168 for e in range(num_experts):
1169 output[e] = torch.ops._C.gptq_marlin_repack(
1170 b_q_weight[e], size_k, size_n, num_bits, is_a_8bit
1171 )
1172 return output
1173
1174
1175 def awq_marlin_moe_repack(
1176 b_q_weight: torch.Tensor,
1177 size_k: int,
1178 size_n: int,
1179 num_bits: int,
1180 is_a_8bit: bool = False,
1181 ) -> torch.Tensor:
1182 num_experts = b_q_weight.shape[0]
1183 assert size_k % 16 == 0
1184 output = torch.empty(
1185 (num_experts, size_k // 16, size_n * (num_bits // 2)),
1186 device=b_q_weight.device,
1187 dtype=b_q_weight.dtype,
1188 )
1189 for e in range(num_experts):
1190 output[e] = torch.ops._C.awq_marlin_repack(
1191 b_q_weight[e], size_k, size_n, num_bits, is_a_8bit
1192 )
1193 return output
1194
1195
1196 def marlin_int4_fp8_preprocess(
1197 qweight: torch.Tensor,
1198 qzeros_or_none: torch.Tensor | None = None,
1199 inplace: bool = False,
1200 ):
1201 return torch.ops._C.marlin_int4_fp8_preprocess(qweight, qzeros_or_none, inplace)
1202
1203
1204 def marlin_gemm(
1205 a: torch.Tensor,
1206 c: torch.Tensor | None,
1207 b_q_weight: torch.Tensor,
1208 b_bias: torch.Tensor | None,
1209 b_scales: torch.Tensor,
1210 a_scales: torch.Tensor | None,
1211 global_scale: torch.Tensor | None,
1212 b_zeros: torch.Tensor | None,
1213 workspace: torch.Tensor,
1214 b_q_type: ScalarType,
1215 size_m: int,
1216 size_n: int,
1217 size_k: int,
1218 use_atomic_add: bool = False,
1219 use_fp32_reduce: bool = False,
1220 is_zp_float: bool = False,
1221 ) -> torch.Tensor:
1222 return torch.ops._C.marlin_gemm(
1223 a,
1224 c,
1225 b_q_weight,
1226 b_bias,
1227 b_scales,
1228 a_scales,
1229 global_scale,
1230 b_zeros,
1231 workspace,
1232 b_q_type.id,
1233 size_m,
1234 size_n,
1235 size_k,
1236 use_atomic_add,
1237 use_fp32_reduce,
1238 is_zp_float,
1239 )
1240
1241
1242 if hasattr(torch.ops._C, "marlin_gemm"):
1243
1244 @register_fake("_C::marlin_gemm")
1245 def _marlin_gemm_fake(
1246 a: torch.Tensor,
1247 c: torch.Tensor | None,
1248 b_q_weight: torch.Tensor,
1249 b_bias: torch.Tensor | None,
1250 b_scales: torch.Tensor,
1251 a_scales: torch.Tensor | None,
1252 global_scale: torch.Tensor | None,
1253 b_zeros: torch.Tensor | None,
1254 workspace: torch.Tensor,
1255 b_q_type_id: int,
1256 size_m: torch.SymInt,
1257 size_n: torch.SymInt,
1258 size_k: torch.SymInt,
1259 use_atomic_add: bool = False,
1260 use_fp32_reduce: bool = False,
1261 is_zp_float: bool = False,
1262 ) -> torch.Tensor:
1263 dtype = a.dtype
1264 if dtype not in [torch.half, torch.bfloat16]:
1265 dtype = b_scales.dtype
1266 return torch.empty((size_m, size_n), device=a.device, dtype=dtype)
1267
1268
1269 # machete
1270 def machete_supported_schedules(
1271 a_type: torch.dtype,
1272 b_type: ScalarType,
1273 group_scales_type: torch.dtype | None,
1274 group_zeros_type: torch.dtype | None = None,
1275 channel_scales_type: torch.dtype | None = None,
1276 token_scales_type: torch.dtype | None = None,
1277 out_type: torch.dtype | None = None,
1278 ) -> list[str]:
1279 return torch.ops._C.machete_supported_schedules(
1280 a_type,
1281 b_type.id,
1282 group_scales_type,
1283 group_zeros_type,
1284 channel_scales_type,
1285 token_scales_type,
1286 out_type,
1287 )
1288
1289
1290 def machete_mm(
1291 a: torch.Tensor,
1292 # b_q Should be the tensor returned by machete_prepack_B
1293 b_q: torch.Tensor,
1294 b_type: ScalarType,
1295 out_type: torch.dtype | None = None,
1296 b_group_scales: torch.Tensor | None = None,
1297 b_group_zeros: torch.Tensor | None = None,
1298 b_group_size: int | None = None,
1299 b_channel_scales: torch.Tensor | None = None,
1300 a_token_scales: torch.Tensor | None = None,
1301 schedule: str | None = None,
1302 ) -> torch.Tensor:
1303 return torch.ops._C.machete_mm(
1304 a,
1305 b_q,
1306 b_type.id,
1307 out_type,
1308 b_group_scales,
1309 b_group_zeros,
1310 b_group_size,
1311 b_channel_scales,
1312 a_token_scales,
1313 schedule,
1314 )
1315
1316
1317 if hasattr(torch.ops._C, "machete_mm"):
1318
1319 @register_fake("_C::machete_mm")
1320 def machete_mm_fake(
1321 a: torch.Tensor,
1322 # b_q Should be the tensor returned by machete_prepack_B
1323 b_q: torch.Tensor,
1324 b_type: ScalarType,
1325 out_type: torch.dtype | None = None,
1326 b_group_scales: torch.Tensor | None = None,
1327 b_group_zeros: torch.Tensor | None = None,
1328 b_group_size: int | None = None,
1329 b_channel_scales: torch.Tensor | None = None,
1330 a_token_scales: torch.Tensor | None = None,
1331 schedule: str | None = None,
1332 ) -> torch.Tensor:
1333 m = a.size(0)
1334 n = b_q.size(1)
1335 return torch.empty((m, n), device=a.device, dtype=a.dtype)
1336
1337
1338 def machete_prepack_B(
1339 b_q_weight: torch.Tensor,
1340 a_type: torch.dtype,
1341 b_type: ScalarType,
1342 group_scales_type: torch.dtype | None,
1343 ) -> torch.Tensor:
1344 return torch.ops._C.machete_prepack_B(
1345 b_q_weight, a_type, b_type.id, group_scales_type
1346 )
1347
1348
1349 if hasattr(torch.ops._C, "machete_prepack_B"):
1350
1351 @register_fake("_C::machete_prepack_B")
1352 def machete_prepack_B_fake(
1353 b_q_weight: torch.Tensor,
1354 a_type: torch.dtype,
1355 b_type: ScalarType,
1356 group_scales_type: torch.dtype | None,
1357 ) -> torch.Tensor:
1358 return torch.empty_like(b_q_weight, memory_format=torch.contiguous_format)
1359
1360
1361 # CUTLASS W4A8
1362 def cutlass_w4a8_mm(
1363 a: torch.Tensor,
1364 # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b
1365 b_q: torch.Tensor,
1366 b_group_scales: torch.Tensor,
1367 b_group_size: int,
1368 b_channel_scales: torch.Tensor,
1369 a_token_scales: torch.Tensor,
1370 out_type: torch.dtype | None = None,
1371 maybe_schedule: str | None = None,
1372 ) -> torch.Tensor:
1373 return torch.ops._C.cutlass_w4a8_mm(
1374 a,
1375 b_q,
1376 b_group_scales,
1377 b_group_size,
1378 b_channel_scales,
1379 a_token_scales,
1380 out_type,
1381 maybe_schedule,
1382 )
1383
1384
1385 if hasattr(torch.ops._C, "cutlass_w4a8_mm"):
1386
1387 @register_fake("_C::cutlass_w4a8_mm")
1388 def cutlass_w4a8_mm_fake(
1389 a: torch.Tensor,
1390 # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b
1391 b_q: torch.Tensor,
1392 b_group_scales: torch.Tensor,
1393 b_group_size: int,
1394 b_channel_scales: torch.Tensor,
1395 a_token_scales: torch.Tensor,
1396 out_type: torch.dtype | None = None,
1397 maybe_schedule: str | None = None,
1398 ) -> torch.Tensor:
1399 m = a.size(0)
1400 n = b_q.size(1)
1401 out_dtype = out_type if out_type is not None else torch.bfloat16
1402 return torch.empty((m, n), device=a.device, dtype=out_dtype)
1403
1404
1405 def cutlass_pack_scale_fp8(scales: torch.Tensor) -> torch.Tensor:
1406 return torch.ops._C.cutlass_pack_scale_fp8(scales)
1407
1408
1409 if hasattr(torch.ops._C, "cutlass_pack_scale_fp8"):
1410
1411 @register_fake("_C::cutlass_pack_scale_fp8")
1412 def cutlass_pack_scale_fp8_fake(scales: torch.Tensor) -> torch.Tensor:
1413 return torch.empty_like(scales, memory_format=torch.contiguous_format)
1414
1415
1416 def cutlass_encode_and_reorder_int4b(b: torch.Tensor) -> torch.Tensor:
1417 return torch.ops._C.cutlass_encode_and_reorder_int4b(b)
1418
1419
1420 if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b"):
1421
1422 @register_fake("_C::cutlass_encode_and_reorder_int4b")
1423 def cutlass_encode_and_reorder_int4b_fake(b: torch.Tensor) -> torch.Tensor:
1424 return torch.empty_like(b, memory_format=torch.contiguous_format)
1425
1426
1427 def cutlass_w4a8_moe_mm(
1428 out_tensors: torch.Tensor,
1429 a_tensors: torch.Tensor,
1430 b_tensors: torch.Tensor,
1431 a_scales: torch.Tensor,
1432 b_scales: torch.Tensor,
1433 b_group_scales: torch.Tensor,
1434 b_group_size: int,
1435 expert_offsets: torch.Tensor,
1436 problem_sizes: torch.Tensor,
1437 a_strides: torch.Tensor,
1438 b_strides: torch.Tensor,
1439 c_strides: torch.Tensor,
1440 group_scale_strides: torch.Tensor,
1441 maybe_schedule: str | None = None,
1442 ) -> None:
1443 """Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the
1444 W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8)
1445 and both per-channel + per-token scaling in the epilogue.
1446
1447 Args:
1448 out_tensors:
1449 Output buffer for all experts (updated in-place).
1450 a_tensors:
1451 FP8 (E4M3FN) activations for all experts.
1452 b_tensors:
1453 INT4-packed weight matrix for all experts, packed to INT32
1454 a_scales:
1455 Per-token FP8 activation scales, applied in the epilogue.
1456 b_scales:
1457 Per-channel FP8 weight scales for each expert, applied in the epilogue.
1458 b_group_scales:
1459 FP8 scale values for group-wise INT4 weight blocks.
1460 b_group_size:
1461 Number of elements grouped under each entry of b_group_scales.
1462 expert_offsets:
1463 Cumulative token offsets
1464 problem_sizes:
1465 Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher.
1466 a_strides:
1467 Strides describing the memory layout of a_tensors.
1468 b_strides:
1469 Strides describing the memory layout of b_tensors.
1470 c_strides:
1471 Strides describing the memory layout of out_tensors.
1472 group_scale_strides:
1473 Strides describing the memory layout of b_group_scales.
1474 maybe_schedule:
1475 Optional override to choose a specific kernel or epilogue schedule.
1476
1477 Returns:
1478 out_tensors updated in-place with the dequantized INT4xFP8 grouped GEMM result.
1479
1480 """
1481 return torch.ops._C.cutlass_w4a8_moe_mm(
1482 out_tensors,
1483 a_tensors,
1484 b_tensors,
1485 a_scales,
1486 b_scales,
1487 b_group_scales,
1488 b_group_size,
1489 expert_offsets,
1490 problem_sizes,
1491 a_strides,
1492 b_strides,
1493 c_strides,
1494 group_scale_strides,
1495 maybe_schedule,
1496 )
1497
1498
1499 def cutlass_encode_and_reorder_int4b_grouped(
1500 b_tensors: torch.Tensor,
1501 ) -> tuple[torch.Tensor, torch.Tensor]:
1502 return torch.ops._C.cutlass_encode_and_reorder_int4b_grouped(b_tensors)
1503
1504
1505 if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b_grouped"):
1506
1507 @register_fake("_C::cutlass_encode_and_reorder_int4b_grouped")
1508 def cutlass_encode_and_reorder_int4b_grouped_fake(b: torch.Tensor) -> torch.Tensor:
1509 return torch.empty_like(b, memory_format=torch.contiguous_format)
1510
1511
1512 # fp4
1513 def scaled_fp4_quant(
1514 input: torch.Tensor,
1515 input_global_scale: torch.Tensor,
1516 is_sf_swizzled_layout: bool = True,
1517 backend: str = "none",
1518 padded_n: int | None = None,
1519 ) -> tuple[torch.Tensor, torch.Tensor]:
1520 """Quantize input tensor to FP4 and return quantized tensor and scale.
1521
1522 This function quantizes the last dimension of the given tensor `input`. For
1523 every 16 consecutive elements, a single dynamically computed scaling factor
1524 is shared. This scaling factor is quantized using the `input_global_scale`
1525 and is stored in a swizzled layout (see
1526 https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x).
1527
1528 Args:
1529 input: The input tensor to be quantized to FP4
1530 input_global_scale: A scalar scaling factor for the entire tensor.
1531 is_sf_swizzled_layout: Whether to store the scaling factors in the
1532 swizzled layout (default `True`).
1533 backend: Quantization kernel backend to dispatch to. For `"trtllm"`
1534 backends the 8x4 scale-factor layout is selected for small
1535 batches (m <= 32) instead of the 128x4 layout.
1536 padded_n: Optional padded K dimension. When provided, the quantized
1537 output and scale tensors are allocated for ``padded_n``
1538
1539 Returns:
1540 tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every
1541 two values are packed into a uint8 and float8_e4m3 scaling factors
1542 in the sizzled layout.
1543
1544 """
1545 assert not current_platform.is_rocm()
1546 assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}."
1547 other_dims = 1 if input.ndim == 1 else -1
1548 input = input.reshape(other_dims, input.shape[-1])
1549 m, n = input.shape
1550 block_size = 16
1551
1552 assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}."
1553 assert input.dtype in (torch.float16, torch.bfloat16), (
1554 f"input.dtype needs to be fp16 or bf16 but got {input.dtype}."
1555 )
1556 if padded_n is not None:
1557 assert padded_n >= n, f"padded_n must be >= n, got padded_n={padded_n}, n={n}."
1558 assert padded_n % block_size == 0, (
1559 f"padded_n has to be a multiple of {block_size}, but got {padded_n}."
1560 )
1561
1562 use_8x4_sf_layout = True if "trtllm" in backend and m <= 32 else False # noqa: SIM210
1563 if use_8x4_sf_layout and padded_n is not None and padded_n != n:
1564 # TODO: support this case
1565 raise ValueError("padded_n is not supported with TRTLLM 8x4 scale layout.")
1566 if use_8x4_sf_layout:
1567 output, output_scale = flashinfer_quant_nvfp4_8x4_sf_layout(
1568 input, input_global_scale
1569 )
1570 else:
1571 # Pre-allocate and call .out variant (same behavior as old in-place API)
1572 output, output_scale = create_fp4_output_tensors(
1573 m,
1574 n,
1575 input.device,
1576 is_sf_swizzled_layout,
1577 padded_n=padded_n,
1578 )
1579 torch.ops._C.scaled_fp4_quant.out(
1580 input,
1581 input_global_scale,
1582 is_sf_swizzled_layout,
1583 output=output,
1584 output_scale=output_scale,
1585 )
1586
1587 output_scale = output_scale.view(torch.float8_e4m3fn)
1588 return output, output_scale
1589
1590
1591 def scaled_fp4_experts_quant(
1592 input_tensor: torch.Tensor,
1593 input_global_scale: torch.Tensor,
1594 expert_offsets: torch.Tensor,
1595 blockscale_offsets: torch.Tensor,
1596 topk: int,
1597 ) -> tuple[torch.Tensor, torch.Tensor]:
1598 """Quantize input tensor to NVFP4 and return quantized tensor and scale, for
1599 packed MoE Inputs.
1600
1601 Args:
1602 input_tensor: The input tensor to be quantized to NVFP4
1603 input_global_scale: A scalar scaling factor for the entire tensor.
1604 expert_offsets: The expert offsets tensor
1605 blockscale_offsets: The blockscale offsets tensor
1606 topk: The number of experts each token is routed to
1607 Outputs:
1608 output: The quantized tensor in NVFP4
1609 output_scales: The blockscale tensor in FP8-E4M3
1610
1611 """
1612 assert not current_platform.is_rocm()
1613 assert input_tensor.ndim == 2, (
1614 f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
1615 )
1616
1617 # Control the maximum number of tokens per expert supported by the
1618 # NVFP4 MoE Expert Quantization. This is used to prevent the kernel
1619 # from running out of memory. This value can also be increased to support
1620 # larger models.
1621 MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
1622 m_numtopk, k = input_tensor.shape
1623
1624 assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
1625 f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
1626 f"{MAX_TOKENS_PER_EXPERT})"
1627 f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
1628 f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
1629 )
1630 scales_k = k // 16
1631 padded_k = (scales_k + (4 - 1)) // 4
1632
1633 # output is uint8 and packed fp4 values
1634 output = torch.empty(
1635 m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
1636 )
1637 output_scales = torch.empty(
1638 MAX_TOKENS_PER_EXPERT * topk,
1639 padded_k,
1640 dtype=torch.int32,
1641 device=input_tensor.device,
1642 )
1643 torch.ops._C.scaled_fp4_experts_quant(
1644 output,
1645 output_scales,
1646 input_tensor,
1647 input_global_scale,
1648 expert_offsets,
1649 blockscale_offsets,
1650 )
1651 output_scales = output_scales.view(torch.float8_e4m3fn)
1652 return output, output_scales
1653
1654
1655 def silu_and_mul_scaled_fp4_experts_quant(
1656 input_tensor: torch.Tensor,
1657 input_global_scale: torch.Tensor,
1658 expert_offsets: torch.Tensor,
1659 blockscale_offsets: torch.Tensor,
1660 topk: int,
1661 ) -> tuple[torch.Tensor, torch.Tensor]:
1662 """Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations.
1663
1664 Args:
1665 input_tensor: The input tensor with gate || up layout [m_topk, k*2]
1666 input_global_scale: A per-expert scaling factor [n_experts]
1667 expert_offsets: The expert offsets tensor [n_experts+1]
1668 blockscale_offsets: The blockscale offsets tensor [n_experts+1]
1669 topk: Number of top-k experts selected
1670 Outputs:
1671 output: The quantized tensor in NVFP4 [m_topk, k/2]
1672 output_scales: The blockscale tensor in FP8-E4M3
1673
1674 """
1675 assert not current_platform.is_rocm()
1676 assert input_tensor.ndim == 2, (
1677 f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
1678 )
1679
1680 # Control the maximum number of tokens per expert supported by the
1681 # NVFP4 MoE Expert Quantization. This is used to prevent the kernel
1682 # from running out of memory. This value can also be increased to support
1683 # larger models.
1684 MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
1685 m_numtopk, k_times_2 = input_tensor.shape
1686 assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
1687 k = k_times_2 // 2
1688
1689 assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
1690 f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
1691 f"{MAX_TOKENS_PER_EXPERT})"
1692 f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
1693 f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
1694 )
1695 scales_k = k // 16
1696 padded_k = (scales_k + (4 - 1)) // 4
1697
1698 # output is uint8 and packed fp4 values
1699 output = torch.empty(
1700 m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
1701 )
1702 output_scales = torch.empty(
1703 MAX_TOKENS_PER_EXPERT * topk,
1704 padded_k,
1705 dtype=torch.int32,
1706 device=input_tensor.device,
1707 )
1708 torch.ops._C.silu_and_mul_scaled_fp4_experts_quant(
1709 output,
1710 output_scales,
1711 input_tensor,
1712 input_global_scale,
1713 expert_offsets,
1714 blockscale_offsets,
1715 )
1716 output_scales = output_scales.view(torch.float8_e4m3fn)
1717 return output, output_scales
1718
1719
1720 def mxfp4_experts_quant(
1721 input_tensor: torch.Tensor,
1722 expert_offsets: torch.Tensor,
1723 blockscale_offsets: torch.Tensor,
1724 n_experts: int,
1725 topk: int,
1726 ) -> tuple[torch.Tensor, torch.Tensor]:
1727 """Quantize input tensor to MXFP4 for packed MoE inputs.
1728 Uses 32-element blocks with E8M0 (power-of-two) scale factors.
1729 MXFP4 has no global scale - only block-level E8M0 scale factors.
1730
1731 Args:
1732 input_tensor: [m_topk, k] BF16/FP16 activations
1733 expert_offsets: [n_experts+1] token boundaries per expert
1734 blockscale_offsets: [n_experts+1] SF row boundaries per expert
1735 n_experts: number of experts
1736 topk: number of top-k experts
1737 Returns:
1738 output: [m_topk, k//2] packed E2M1 values (uint8)
1739 output_scales: E8M0 blockscales in swizzled layout (uint8 view)
1740
1741 """
1742 assert not current_platform.is_rocm()
1743 assert input_tensor.ndim == 2
1744
1745 MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
1746 m_numtopk, k = input_tensor.shape
1747
1748 assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
1749 f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
1750 f"{MAX_TOKENS_PER_EXPERT})"
1751 f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use"
1752 f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
1753 )
1754 scales_k = k // 32
1755 padded_k = (scales_k + (4 - 1)) // 4
1756
1757 output = torch.empty(
1758 m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
1759 )
1760 output_scales = torch.empty(
1761 MAX_TOKENS_PER_EXPERT * topk,
1762 padded_k,
1763 dtype=torch.int32,
1764 device=input_tensor.device,
1765 )
1766 torch.ops._C.mxfp4_experts_quant(
1767 output,
1768 output_scales,
1769 input_tensor,
1770 expert_offsets,
1771 blockscale_offsets,
1772 n_experts,
1773 )
1774 # E8M0 SFs are stored as uint8
1775 output_scales = output_scales.view(torch.uint8)
1776 return output, output_scales
1777
1778
1779 def silu_and_mul_mxfp4_experts_quant(
1780 input_tensor: torch.Tensor,
1781 expert_offsets: torch.Tensor,
1782 blockscale_offsets: torch.Tensor,
1783 n_experts: int,
1784 topk: int,
1785 ) -> tuple[torch.Tensor, torch.Tensor]:
1786 """Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
1787 MXFP4 has no global scale - only block-level E8M0 scale factors.
1788 """
1789 assert not current_platform.is_rocm()
1790 assert input_tensor.ndim == 2
1791
1792 MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
1793 m_numtopk, k_times_2 = input_tensor.shape
1794 assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
1795 k = k_times_2 // 2
1796
1797 assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk
1798 scales_k = k // 32
1799 padded_k = (scales_k + (4 - 1)) // 4
1800
1801 output = torch.empty(
1802 m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
1803 )
1804 output_scales = torch.empty(
1805 MAX_TOKENS_PER_EXPERT * topk,
1806 padded_k,
1807 dtype=torch.int32,
1808 device=input_tensor.device,
1809 )
1810 torch.ops._C.silu_and_mul_mxfp4_experts_quant(
1811 output,
1812 output_scales,
1813 input_tensor,
1814 expert_offsets,
1815 blockscale_offsets,
1816 n_experts,
1817 )
1818 output_scales = output_scales.view(torch.uint8)
1819 return output, output_scales
1820
1821
1822 # fp8
1823 def scaled_fp8_quant(
1824 input: torch.Tensor,
1825 scale: torch.Tensor | None = None,
1826 num_token_padding: int | None = None,
1827 scale_ub: torch.Tensor | None = None,
1828 use_per_token_if_dynamic: bool = False,
1829 output: torch.Tensor | None = None,
1830 group_shape: tuple[int, int] | None = None,
1831 ) -> tuple[torch.Tensor, torch.Tensor]:
1832 """Quantize input tensor to FP8 and return quantized tensor and scale.
1833
1834 This function supports both static and dynamic quantization: If you
1835 provide the scale, it will use static scaling and if you omit it,
1836 the scale will be determined dynamically. The function also allows
1837 optional padding of the output tensors for downstream kernels that
1838 will benefit from padding.
1839
1840 Args:
1841 input: The input tensor to be quantized to FP8 (must be 2D: [M, N])
1842 scale: Optional scaling factor for the FP8 quantization. Supports:
1843 - 0D or [1]: per-tensor scaling
1844 - 1D: requires explicit group_shape to disambiguate per-channel
1845 vs per-token (use (-1, 1) for per-channel, (1, -1) for per-token)
1846 - 2D [M/group_m, N/group_n]: group scaling (e.g. [M, N/128] for
1847 DeepSeek-style (1,128) groups, or [M/128, N/128] for (128,128))
1848 scale_ub: Optional upper bound for scaling factor in dynamic
1849 per token case
1850 num_token_padding: If specified, pad the first dimension
1851 of the output to at least this value.
1852 use_per_token_if_dynamic: Whether to do per_tensor or per_token
1853 in the dynamic quantization case.
1854 output: Optional tensor to write the quantized result into. A new
1855 tensor is allocated when omitted.
1856 group_shape: Optional tuple (group_m, group_n) specifying the group
1857 shape for static quantization. Use -1 for "full extent" (e.g.,
1858 (-1, -1) for per-tensor, (-1, 1) for per-channel, etc.)
1859 Required for 1D scales; optional for 2D scales.
1860
1861 Returns:
1862 tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and
1863 scaling factor.
1864
1865 """
1866 # This code assumes batch_dim and num_tokens are flattened
1867 assert input.ndim == 2
1868 shape: tuple[int, int] | torch.Size = input.shape
1869 # For ROCm on MI300, the output fp8 dtype is torch.float_e3m3fnuz
1870 out_dtype: torch.dtype = current_platform.fp8_dtype()
1871 if num_token_padding:
1872 shape = (max(num_token_padding, input.shape[0]), shape[1])
1873 if output is None:
1874 output = torch.empty(shape, device=input.device, dtype=out_dtype)
1875 else:
1876 assert num_token_padding is None, "padding not supported if output passed in"
1877 assert output.dtype == out_dtype
1878
1879 if scale is None:
1880 if use_per_token_if_dynamic:
1881 scale = torch.empty((shape[0], 1), device=input.device, dtype=torch.float32)
1882 torch.ops._C.dynamic_per_token_scaled_fp8_quant(
1883 output, input, scale, scale_ub
1884 )
1885 else:
1886 scale = torch.empty(1, device=input.device, dtype=torch.float32)
1887 torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
1888 else:
1889 torch.ops._C.static_scaled_fp8_quant(output, input, scale, group_shape)
1890
1891 return output, scale
1892
1893
1894 # int8
1895 def scaled_int8_quant(
1896 input: torch.Tensor,
1897 scale: torch.Tensor | None = None,
1898 azp: torch.Tensor | None = None,
1899 symmetric: bool = True,
1900 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
1901 """Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
1902
1903 Args:
1904 input: The input tensor to be quantized to int8.
1905 scale: Optional scaling factor for the int8 quantization.
1906 When not provided, we invoke dynamic-per-token quantization.
1907 azp: Optional zero-point for the int8 quantization.
1908 Must be provided for asymmetric quantization if `scale` is provided.
1909 symmetric: Whether to use symmetric quantization (scale only, azp ignored).
1910
1911 Returns:
1912 Output int8 tensor, scales, and optionally azp.
1913
1914 """
1915 if current_platform.is_xpu():
1916 # XPU has no _C int8 quant op; use the torch.compile reference.
1917 if not symmetric:
1918 raise NotImplementedError(
1919 "asymmetric int8 activation quantization is unsupported on XPU"
1920 )
1921 if scale is not None:
1922 q = (input.to(torch.float32) / scale).round().clamp(-128, 127)
1923 return q.to(torch.int8), scale, None
1924
1925 from vllm._xpu_ops import xpu_ops
1926
1927 q, scales, _ = xpu_ops.dynamic_per_token_int8_quant_ref(
1928 input.contiguous(), True, 8
1929 )
1930 return q, scales.reshape(-1, 1).to(torch.float32), None
1931
1932 output = torch.empty_like(input, dtype=torch.int8)
1933 if scale is not None:
1934 # static-per-tensor quantization.
1935 assert symmetric == (azp is None), (
1936 "azp must only be provided for asymmetric quantization."
1937 )
1938 torch.ops._C.static_scaled_int8_quant(output, input, scale, azp)
1939 return output, scale, azp
1940
1941 # dynamic-per-token quantization.
1942 input_scales = torch.empty(
1943 (input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
1944 )
1945 input_azp = None if symmetric else torch.empty_like(input_scales, dtype=torch.int32)
1946 torch.ops._C.dynamic_scaled_int8_quant(
1947 output, input.contiguous(), input_scales, input_azp
1948 )
1949 return output, input_scales, input_azp
1950
1951
1952 # mamba
1953 def selective_scan_fwd(
1954 u: torch.Tensor,
1955 delta: torch.Tensor,
1956 A: torch.Tensor,
1957 B: torch.Tensor,
1958 C: torch.Tensor,
1959 D_: torch.Tensor | None,
1960 z_: torch.Tensor | None,
1961 delta_bias_: torch.Tensor | None,
1962 delta_softplus: bool,
1963 query_start_loc: torch.Tensor | None,
1964 cache_indices: torch.Tensor | None,
1965 has_initial_state: torch.Tensor | None,
1966 ssm_states: torch.Tensor,
1967 null_block_id: int,
1968 block_size: int = 1024,
1969 block_idx_first_scheduled_token: torch.Tensor | None = None,
1970 block_idx_last_scheduled_token: torch.Tensor | None = None,
1971 initial_state_idx: torch.Tensor | None = None,
1972 cu_chunk_seqlen: torch.Tensor | None = None,
1973 last_chunk_indices: torch.Tensor | None = None,
1974 ):
1975 torch.ops._C.selective_scan_fwd(
1976 u,
1977 delta,
1978 A,
1979 B,
1980 C,
1981 D_,
1982 z_,
1983 delta_bias_,
1984 delta_softplus,
1985 query_start_loc,
1986 cache_indices,
1987 has_initial_state,
1988 ssm_states,
1989 null_block_id,
1990 block_size,
1991 block_idx_first_scheduled_token,
1992 block_idx_last_scheduled_token,
1993 initial_state_idx,
1994 cu_chunk_seqlen,
1995 last_chunk_indices,
1996 )
1997
1998
1999 def causal_conv1d_update_cpu_vec(
2000 x: torch.Tensor,
2001 conv_state: torch.Tensor,
2002 weight: torch.Tensor,
2003 bias: torch.Tensor | None = None,
2004 activation: str | None = None,
2005 conv_state_indices: torch.Tensor | None = None,
2006 query_start_loc: torch.Tensor | None = None,
2007 pad_slot_id: int = 0,
2008 ) -> torch.Tensor:
2009 return torch.ops._C.causal_conv1d_update_cpu_vec(
2010 x,
2011 conv_state,
2012 weight,
2013 bias,
2014 activation,
2015 conv_state_indices,
2016 query_start_loc,
2017 pad_slot_id,
2018 )
2019
2020
2021 def selective_state_update_cpu(
2022 state: torch.Tensor,
2023 x: torch.Tensor,
2024 dt: torch.Tensor,
2025 A: torch.Tensor,
2026 B: torch.Tensor,
2027 C: torch.Tensor,
2028 D: torch.Tensor | None,
2029 z: torch.Tensor | None,
2030 dt_bias: torch.Tensor | None,
2031 dt_softplus: bool,
2032 state_batch_indices: torch.Tensor | None,
2033 dst_state_batch_indices: torch.Tensor | None,
2034 null_block_id: int,
2035 out: torch.Tensor,
2036 num_accepted_tokens: torch.Tensor | None,
2037 cu_seqlens: torch.Tensor | None,
2038 ):
2039 torch.ops._C.selective_state_update_cpu(
2040 state,
2041 x,
2042 dt,
2043 A,
2044 B,
2045 C,
2046 D,
2047 z,
2048 dt_bias,
2049 dt_softplus,
2050 state_batch_indices,
2051 dst_state_batch_indices,
2052 null_block_id,
2053 out,
2054 num_accepted_tokens,
2055 cu_seqlens,
2056 )
2057
2058
2059 def mamba_chunk_scan_fwd_cpu(
2060 out: torch.Tensor,
2061 final_states: torch.Tensor,
2062 x: torch.Tensor,
2063 dt: torch.Tensor,
2064 A: torch.Tensor,
2065 B: torch.Tensor,
2066 C: torch.Tensor,
2067 D: torch.Tensor | None,
2068 z: torch.Tensor | None,
2069 cu_seqlens: torch.Tensor,
2070 ) -> None:
2071 """Prefill SSM scan kernel. out and final_states are written in-place."""
2072 torch.ops._C.mamba_chunk_scan_fwd_cpu(
2073 out,
2074 final_states,
2075 x,
2076 dt,
2077 A,
2078 B,
2079 C,
2080 D,
2081 z,
2082 cu_seqlens,
2083 )
2084
2085
2086 # ROCm skinny gemms
2087 def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor:
2088 return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
2089
2090
2091 def wvSplitK(
2092 a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
2093 ) -> torch.Tensor:
2094 return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count)
2095
2096
2097 def wvSplitK_int4_g(
2098 weight: torch.Tensor,
2099 activation: torch.Tensor,
2100 scale: torch.Tensor,
2101 cu_count: int,
2102 group_size: int,
2103 zero_points: torch.Tensor | None = None,
2104 bias: torch.Tensor | None = None,
2105 ) -> torch.Tensor:
2106 # NOTE: the kernel is weight-major; `weight` is the packed int4 operand
2107 # (in_a, [out_features, K/2]) and `activation` is in_b ([num_tokens, K]).
2108 return torch.ops._rocm_C.wvSplitK_int4_g(
2109 weight, activation, scale, zero_points, bias, cu_count, group_size
2110 )
2111
2112
2113 if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"):
2114
2115 @register_fake("_rocm_C::wvSplitK_int4_g")
2116 def _wvSplitK_int4_g_fake(
2117 in_a: torch.Tensor, # packed int4 weight [out_features, K/2]
2118 in_b: torch.Tensor, # activation [num_tokens, K]
2119 in_scale: torch.Tensor,
2120 in_zero_points: torch.Tensor | None,
2121 in_bias: torch.Tensor | None,
2122 CuCount: int,
2123 group_size: int,
2124 ) -> torch.Tensor:
2125 # Kernel returns {in_b.size(0), in_a.size(0)} = [num_tokens, out_features].
2126 num_tokens = in_b.size(0)
2127 out_features = in_a.size(0)
2128 return torch.empty(
2129 (num_tokens, out_features), dtype=in_b.dtype, device=in_b.device
2130 )
2131
2132
2133 def wvSplitKrc(
2134 a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
2135 ) -> torch.Tensor:
2136 return torch.ops._rocm_C.wvSplitKrc(a, b, bias, cu_count)
2137
2138
2139 def wvSplitKQ(
2140 a: torch.Tensor,
2141 b: torch.Tensor,
2142 out_dtype: torch.dtype,
2143 scale_a: torch.Tensor,
2144 scale_b: torch.Tensor,
2145 cu_count: int,
2146 bias: torch.Tensor = None,
2147 ) -> torch.Tensor:
2148 out = torch.empty((b.shape[0], a.shape[0]), dtype=out_dtype, device=b.device)
2149 torch.ops._rocm_C.wvSplitKQ(a, b, bias, out, scale_a, scale_b, cu_count)
2150 return out
2151
2152
2153 # moe
2154 def moe_sum(
2155 input: torch.Tensor,
2156 output: torch.Tensor,
2157 topk_ids: torch.Tensor | None = None,
2158 expert_map: torch.Tensor | None = None,
2159 ):
2160 torch.ops._moe_C.moe_sum(input, output, topk_ids, expert_map)
2161
2162
2163 def moe_align_block_size(
2164 topk_ids: torch.Tensor,
2165 num_experts: int,
2166 block_size: int,
2167 sorted_token_ids: torch.Tensor,
2168 experts_ids: torch.Tensor,
2169 num_tokens_post_pad: torch.Tensor,
2170 expert_map: torch.Tensor | None = None,
2171 scatter_idx: torch.Tensor | None = None,
2172 ) -> None:
2173 if current_platform.is_xpu():
2174 if scatter_idx is not None:
2175 raise NotImplementedError("scatter_idx is not supported on XPU")
2176 torch.ops._moe_C.moe_align_block_size(
2177 topk_ids,
2178 num_experts,
2179 block_size,
2180 sorted_token_ids,
2181 experts_ids,
2182 num_tokens_post_pad,
2183 expert_map,
2184 )
2185 return
2186 torch.ops._moe_C.moe_align_block_size(
2187 topk_ids,
2188 num_experts,
2189 block_size,
2190 sorted_token_ids,
2191 experts_ids,
2192 num_tokens_post_pad,
2193 expert_map,
2194 scatter_idx,
2195 )
2196
2197
2198 def batched_moe_align_block_size(
2199 max_tokens_per_batch: int,
2200 block_size: int,
2201 expert_num_tokens: torch.Tensor,
2202 sorted_ids: torch.Tensor,
2203 expert_ids: torch.Tensor,
2204 num_tokens_post_pad: torch.Tensor,
2205 ) -> None:
2206 torch.ops._moe_C.batched_moe_align_block_size(
2207 max_tokens_per_batch,
2208 block_size,
2209 expert_num_tokens,
2210 sorted_ids,
2211 expert_ids,
2212 num_tokens_post_pad,
2213 )
2214
2215
2216 def moe_lora_align_block_size(
2217 topk_ids: torch.Tensor,
2218 token_lora_mapping: torch.Tensor,
2219 num_experts: int,
2220 block_size: int,
2221 max_loras: int,
2222 max_num_tokens_padded: int,
2223 max_num_m_blocks: int,
2224 sorted_token_ids: torch.Tensor,
2225 experts_ids: torch.Tensor,
2226 num_tokens_post_pad: torch.Tensor,
2227 adapter_enabled: torch.Tensor,
2228 lora_ids: torch.Tensor,
2229 expert_map: torch.Tensor | None = None,
2230 ) -> None:
2231 torch.ops._moe_C.moe_lora_align_block_size(
2232 topk_ids,
2233 token_lora_mapping,
2234 num_experts,
2235 block_size,
2236 max_loras,
2237 max_num_tokens_padded,
2238 max_num_m_blocks,
2239 sorted_token_ids,
2240 experts_ids,
2241 num_tokens_post_pad,
2242 adapter_enabled,
2243 lora_ids,
2244 expert_map,
2245 )
2246
2247
2248 def moe_wna16_gemm(
2249 input: torch.Tensor,
2250 output: torch.Tensor,
2251 b_qweight: torch.Tensor,
2252 b_scales: torch.Tensor,
2253 b_qzeros: torch.Tensor | None,
2254 topk_weights: torch.Tensor | None,
2255 sorted_token_ids: torch.Tensor,
2256 experts_ids: torch.Tensor,
2257 num_tokens_post_pad: torch.Tensor,
2258 top_k: int,
2259 BLOCK_SIZE_M: int,
2260 BLOCK_SIZE_N: int,
2261 BLOCK_SIZE_K: int,
2262 bit: int,
2263 ) -> torch.Tensor:
2264 if not current_platform.is_cuda():
2265 raise NotImplementedError(
2266 "The optimized moe_wna16_gemm kernel is only available on CUDA platforms"
2267 )
2268 torch.ops._moe_C.moe_wna16_gemm(
2269 input,
2270 output,
2271 b_qweight,
2272 b_scales,
2273 b_qzeros,
2274 topk_weights,
2275 sorted_token_ids,
2276 experts_ids,
2277 num_tokens_post_pad,
2278 top_k,
2279 BLOCK_SIZE_M,
2280 BLOCK_SIZE_N,
2281 BLOCK_SIZE_K,
2282 bit,
2283 )
2284
2285
2286 def fp32_router_gemm(
2287 hidden_states: torch.Tensor,
2288 router_weight: torch.Tensor,
2289 ) -> torch.Tensor:
2290 output = torch.empty(
2291 hidden_states.shape[0],
2292 router_weight.shape[0],
2293 device=hidden_states.device,
2294 dtype=torch.float32,
2295 )
2296 torch.ops._C.fp32_router_gemm(output, hidden_states, router_weight)
2297 return output
2298
2299
2300 def topk_softmax(
2301 topk_weights: torch.Tensor,
2302 topk_ids: torch.Tensor,
2303 token_expert_indices: torch.Tensor,
2304 gating_output: torch.Tensor,
2305 renormalize: bool = False,
2306 e_score_correction_bias: torch.Tensor | None = None,
2307 is_padding: torch.Tensor | None = None,
2308 ) -> None:
2309 torch.ops._moe_C.topk_softmax(
2310 topk_weights,
2311 topk_ids,
2312 token_expert_indices,
2313 gating_output,
2314 renormalize,
2315 e_score_correction_bias,
2316 is_padding,
2317 )
2318
2319
2320 def topk_sigmoid(
2321 topk_weights: torch.Tensor,
2322 topk_ids: torch.Tensor,
2323 token_expert_indices: torch.Tensor,
2324 gating_output: torch.Tensor,
2325 renormalize: bool = False,
2326 e_score_correction_bias: torch.Tensor | None = None,
2327 routed_scaling_factor: float = 1.0,
2328 is_padding: torch.Tensor | None = None,
2329 ) -> None:
2330 torch.ops._moe_C.topk_sigmoid(
2331 topk_weights,
2332 topk_ids,
2333 token_expert_indices,
2334 gating_output,
2335 renormalize,
2336 e_score_correction_bias,
2337 routed_scaling_factor,
2338 is_padding,
2339 )
2340
2341
2342 def topk_hash_softplus_sqrt(
2343 topk_weights: torch.Tensor,
2344 topk_indices: torch.Tensor,
2345 token_expert_indices: torch.Tensor,
2346 gating_output: torch.Tensor,
2347 renormalize: bool = False,
2348 routed_scaling_factor: float = 1.0,
2349 e_score_correction_bias: torch.Tensor | None = None,
2350 input_tokens: torch.Tensor | None = None,
2351 hash_indices_table: torch.Tensor | None = None,
2352 is_padding: torch.Tensor | None = None,
2353 bias_vl: torch.Tensor | None = None,
2354 image_sentinel_lo: int = 0,
2355 ) -> None:
2356 torch.ops._moe_C.topk_softplus_sqrt(
2357 topk_weights,
2358 topk_indices,
2359 token_expert_indices,
2360 gating_output,
2361 renormalize,
2362 routed_scaling_factor,
2363 e_score_correction_bias,
2364 input_tokens,
2365 hash_indices_table,
2366 is_padding,
2367 bias_vl,
2368 image_sentinel_lo,
2369 )
2370
2371
2372 def grouped_topk(
2373 scores: torch.Tensor,
2374 num_expert_group: int,
2375 topk_group: int,
2376 topk: int,
2377 renormalize: bool,
2378 routed_scaling_factor: float,
2379 bias: torch.Tensor,
2380 scoring_func: int = 0,
2381 ):
2382 """Perform grouped top-k routing for mixture of experts.
2383
2384 Args:
2385 scores: Raw inputs (logits if scoring_func=1, scores if scoring_func=0)
2386 num_expert_group: Number of expert groups
2387 topk_group: Number of groups to select
2388 topk: Number of experts to select per token
2389 renormalize: Whether to renormalize the output weights
2390 routed_scaling_factor: Scaling factor for routing weights
2391 bias: Bias tensor (e_score_correction_bias). Always fused in kernel.
2392 scoring_func: 0=none (no activation), 1=sigmoid
2393
2394 """
2395 if not (current_platform.is_cuda() or current_platform.is_xpu()):
2396 raise NotImplementedError(
2397 "The fused grouped_topk kernel is only available on CUDA and XPU platforms"
2398 )
2399 return torch.ops._moe_C.grouped_topk(
2400 scores,
2401 num_expert_group,
2402 topk_group,
2403 topk,
2404 renormalize,
2405 routed_scaling_factor,
2406 bias,
2407 scoring_func,
2408 )
2409
2410
2411 if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "grouped_topk"):
2412
2413 @register_fake("_moe_C::grouped_topk")
2414 def _grouped_topk_fake(
2415 scores: torch.Tensor,
2416 num_expert_group: int,
2417 topk_group: int,
2418 topk: int,
2419 renormalize: bool,
2420 routed_scaling_factor: float,
2421 bias: torch.Tensor,
2422 scoring_func: int = 0,
2423 ) -> tuple[torch.Tensor, torch.Tensor]:
2424 num_tokens = scores.size(0)
2425 return (
2426 scores.new_empty((num_tokens, topk), dtype=torch.float32),
2427 scores.new_empty((num_tokens, topk), dtype=torch.int32),
2428 )
2429
2430
2431 def moe_wna16_marlin_gemm(
2432 input: torch.Tensor,
2433 output: torch.Tensor | None,
2434 b_qweight: torch.Tensor,
2435 b_bias: torch.Tensor | None,
2436 b_scales: torch.Tensor,
2437 a_scales: torch.Tensor | None,
2438 global_scale: torch.Tensor | None,
2439 b_qzeros: torch.Tensor | None,
2440 workspace: torch.Tensor,
2441 sorted_token_ids: torch.Tensor,
2442 expert_ids: torch.Tensor,
2443 num_tokens_past_padded: torch.Tensor,
2444 topk_weights: torch.Tensor,
2445 moe_block_size: int,
2446 top_k: int,
2447 mul_topk_weights: bool,
2448 b_q_type: ScalarType,
2449 size_m: int,
2450 size_n: int,
2451 size_k: int,
2452 use_atomic_add: bool,
2453 use_fp32_reduce: bool,
2454 is_zp_float: bool,
2455 thread_k: int = -1,
2456 thread_n: int = -1,
2457 blocks_per_sm: int = -1,
2458 ) -> torch.Tensor:
2459 return torch.ops._moe_C.moe_wna16_marlin_gemm(
2460 input,
2461 output,
2462 b_qweight,
2463 b_bias,
2464 b_scales,
2465 a_scales,
2466 global_scale,
2467 b_qzeros,
2468 workspace,
2469 sorted_token_ids,
2470 expert_ids,
2471 num_tokens_past_padded,
2472 topk_weights,
2473 moe_block_size,
2474 top_k,
2475 mul_topk_weights,
2476 b_q_type.id,
2477 size_m,
2478 size_n,
2479 size_k,
2480 use_atomic_add,
2481 use_fp32_reduce,
2482 is_zp_float,
2483 thread_k,
2484 thread_n,
2485 blocks_per_sm,
2486 )
2487
2488
2489 if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_gemm"):
2490
2491 @register_fake("_moe_C::moe_wna16_marlin_gemm")
2492 def moe_wna16_marlin_gemm_fake(
2493 input: torch.Tensor,
2494 output: torch.Tensor | None,
2495 b_qweight: torch.Tensor,
2496 b_bias: torch.Tensor | None,
2497 b_scales: torch.Tensor,
2498 a_scales: torch.Tensor | None,
2499 global_scale: torch.Tensor | None,
2500 b_qzeros: torch.Tensor | None,
2501 workspace: torch.Tensor,
2502 sorted_token_ids: torch.Tensor,
2503 expert_ids: torch.Tensor,
2504 num_tokens_past_padded: torch.Tensor,
2505 topk_weights: torch.Tensor,
2506 moe_block_size: int,
2507 top_k: int,
2508 mul_topk_weights: bool,
2509 b_q_type: ScalarType,
2510 size_m: int,
2511 size_n: int,
2512 size_k: int,
2513 use_atomic_add: bool,
2514 use_fp32_reduce: bool,
2515 is_zp_float: bool,
2516 ):
2517 return torch.empty(
2518 (size_m * top_k, size_n), dtype=input.dtype, device=input.device
2519 )
2520
2521
2522 def reshape_and_cache(
2523 key: torch.Tensor,
2524 value: torch.Tensor,
2525 key_cache: torch.Tensor,
2526 value_cache: torch.Tensor,
2527 slot_mapping: torch.Tensor,
2528 kv_cache_dtype: str,
2529 k_scale: torch.Tensor,
2530 v_scale: torch.Tensor,
2531 ) -> None:
2532 torch.ops._C_cache_ops.reshape_and_cache(
2533 key,
2534 value,
2535 key_cache,
2536 value_cache,
2537 slot_mapping,
2538 kv_cache_dtype,
2539 k_scale,
2540 v_scale,
2541 )
2542
2543
2544 def reshape_and_cache_flash(
2545 key: torch.Tensor,
2546 value: torch.Tensor,
2547 key_cache: torch.Tensor,
2548 value_cache: torch.Tensor,
2549 slot_mapping: torch.Tensor,
2550 kv_cache_dtype: str,
2551 k_scale: torch.Tensor,
2552 v_scale: torch.Tensor,
2553 ) -> None:
2554 torch.ops._C_cache_ops.reshape_and_cache_flash(
2555 key,
2556 value,
2557 key_cache,
2558 value_cache,
2559 slot_mapping,
2560 kv_cache_dtype,
2561 k_scale,
2562 v_scale,
2563 )
2564
2565
2566 def fused_minimax_m3_qknorm_rope_kv_insert(
2567 qkv: torch.Tensor,
2568 q_norm_weight: torch.Tensor,
2569 k_norm_weight: torch.Tensor,
2570 cos_sin_cache: torch.Tensor,
2571 positions: torch.Tensor,
2572 num_heads: int,
2573 num_kv_heads: int,
2574 rotary_dim: int,
2575 eps: float,
2576 index_q_norm_weight: torch.Tensor | None = None,
2577 index_k_norm_weight: torch.Tensor | None = None,
2578 num_index_heads: int = 0,
2579 slot_mapping: torch.Tensor | None = None,
2580 index_slot_mapping: torch.Tensor | None = None,
2581 kv_cache: torch.Tensor | None = None,
2582 index_cache: torch.Tensor | None = None,
2583 block_size: int = 0,
2584 q_out: torch.Tensor | None = None,
2585 index_q_out: torch.Tensor | None = None,
2586 kv_cache_dtype: str = "auto",
2587 skip_index_branch: bool = False,
2588 q_fp8_out: torch.Tensor | None = None,
2589 q_fp8_scale: float = 1.0,
2590 kv_k_scale: torch.Tensor | None = None,
2591 kv_v_scale: torch.Tensor | None = None,
2592 ) -> None:
2593 """Fused MiniMax-M3 attention pre-processing (in-place).
2594
2595 Applies Gemma RMSNorm + partial NeoX RoPE to ``qkv`` in place. ``qkv`` is a
2596 single fused tensor:
2597
2598 - dense layer (``num_index_heads == 0``): ``[q | k | v]``;
2599 - sparse layer (``num_index_heads > 0``): ``[q | k | v | index_q |
2600 index_k]`` — the index branch is read straight out of ``qkv``.
2601
2602 When ``kv_cache`` is given (sparse serving), also scatter-inserts the
2603 normed/roped k & v into the paged KV cache by ``slot_mapping`` and the
2604 index key into ``index_cache`` by ``index_slot_mapping``. ``kv_cache_dtype``
2605 selects the cache storage/conversion path. If
2606 ``index_slot_mapping`` is omitted, ``slot_mapping`` is used for both caches.
2607
2608 If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N,
2609 niq*128]``) are given, the normed/roped q / index_q are written there
2610 instead of in place — folding the de-interleave into this kernel's store so
2611 callers skip a separate ``.contiguous()`` copy before the SM100 sparse
2612 attention's flat TMA descriptor.
2613
2614 If ``q_fp8_out`` is given, the same normalized q is also written in FP8
2615 E4M3 using ``q_fp8_scale`` as its dequantization scale. Without ``q_out``,
2616 q is then written only in FP8 and the q slice of ``qkv`` is left as is.
2617
2618 ``kv_cache_dtype="nvfp4"`` quantizes k/v into the packed HND
2619 ``[num_blocks, 2 * num_kv_heads, block_size, 72]`` uint8 cache using the
2620 device dequantization scales ``kv_k_scale``/``kv_v_scale``. Slot
2621 ``2 * head + side`` (K = 0, V = 1) holds that head's E2M1 data followed by
2622 its E4M3 block scales, so every head is one contiguous run of the page.
2623
2624 When ``skip_index_branch`` is true, sparse rows still keep their packed
2625 ``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v
2626 branches and main KV cache. This is used by MiniMax-M3 index-topk reuse
2627 layers that consume top-k block ids selected by an earlier sparse layer.
2628 """
2629 torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert(
2630 qkv,
2631 q_norm_weight,
2632 k_norm_weight,
2633 cos_sin_cache,
2634 positions,
2635 num_heads,
2636 num_kv_heads,
2637 rotary_dim,
2638 eps,
2639 index_q_norm_weight,
2640 index_k_norm_weight,
2641 num_index_heads,
2642 slot_mapping,
2643 index_slot_mapping,
2644 kv_cache,
2645 index_cache,
2646 block_size,
2647 q_out,
2648 index_q_out,
2649 kv_cache_dtype,
2650 skip_index_branch,
2651 q_fp8_out,
2652 q_fp8_scale,
2653 kv_k_scale,
2654 kv_v_scale,
2655 )
2656
2657
2658 def fused_kda_decode(
2659 x: torch.Tensor,
2660 weight: torch.Tensor,
2661 bias: torch.Tensor | None,
2662 conv_state: torch.Tensor,
2663 raw_g: torch.Tensor,
2664 raw_beta: torch.Tensor,
2665 A_log: torch.Tensor,
2666 dt_bias: torch.Tensor,
2667 state_indices: torch.Tensor,
2668 state: torch.Tensor,
2669 out: torch.Tensor | None = None,
2670 lower_bound: float | None = None,
2671 output_gate: torch.Tensor | None = None,
2672 norm_weight: torch.Tensor | None = None,
2673 norm_eps: float = 1e-5,
2674 ) -> torch.Tensor:
2675 if out is None:
2676 out = torch.empty(
2677 1,
2678 x.shape[0],
2679 raw_g.shape[2],
2680 raw_g.shape[3],
2681 dtype=x.dtype,
2682 device=x.device,
2683 )
2684 torch.ops._C.fused_kda_decode(
2685 x,
2686 weight,
2687 bias,
2688 conv_state,
2689 raw_g,
2690 raw_beta,
2691 A_log,
2692 dt_bias,
2693 state_indices,
2694 state,
2695 out,
2696 lower_bound,
2697 output_gate,
2698 norm_weight,
2699 norm_eps,
2700 )
2701 return out
2702
2703
2704 def fused_gdn_decode_post_conv_mtp(
2705 mixed_qkv: torch.Tensor,
2706 a: torch.Tensor,
2707 b: torch.Tensor,
2708 A_log: torch.Tensor,
2709 dt_bias: torch.Tensor,
2710 state_indices: torch.Tensor,
2711 cu_seqlens: torch.Tensor,
2712 num_accepted_tokens: torch.Tensor,
2713 state: torch.Tensor,
2714 output_gate: torch.Tensor,
2715 norm_weight: torch.Tensor,
2716 out: torch.Tensor | None = None,
2717 scale: float = 128**-0.5,
2718 norm_eps: float = 1e-5,
2719 output_gate_activation: str = "silu",
2720 ) -> torch.Tensor:
2721 if out is None:
2722 out = torch.empty_like(output_gate)
2723 torch.ops._C.fused_gdn_decode_post_conv_mtp(
2724 mixed_qkv,
2725 a,
2726 b,
2727 A_log,
2728 dt_bias,
2729 state_indices,
2730 cu_seqlens,
2731 num_accepted_tokens,
2732 state,
2733 output_gate,
2734 norm_weight,
2735 out,
2736 scale,
2737 norm_eps,
2738 output_gate_activation,
2739 )
2740 return out
2741
2742
2743 def concat_and_cache_mla(
2744 kv_c: torch.Tensor,
2745 k_pe: torch.Tensor,
2746 kv_cache: torch.Tensor,
2747 slot_mapping: torch.Tensor,
2748 kv_cache_dtype: str,
2749 scale: torch.Tensor,
2750 ) -> None:
2751 torch.ops._C_cache_ops.concat_and_cache_mla(
2752 kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale
2753 )
2754
2755
2756 def concat_and_cache_mla_grouped(
2757 kv_c: torch.Tensor,
2758 k_pe: torch.Tensor,
2759 kv_cache_ptrs: torch.Tensor,
2760 slot_mapping: torch.Tensor,
2761 block_size: int,
2762 block_stride: int,
2763 entry_stride: int,
2764 kv_scales: torch.Tensor | None = None,
2765 kv_cache_dtype: str = "auto",
2766 ) -> None:
2767 torch.ops._C_cache_ops.concat_and_cache_mla_grouped(
2768 kv_c,
2769 k_pe,
2770 kv_cache_ptrs,
2771 slot_mapping,
2772 block_size,
2773 block_stride,
2774 entry_stride,
2775 kv_scales,
2776 kv_cache_dtype,
2777 )
2778
2779
2780 def kimi_k3_attn_res(
2781 prefix: torch.Tensor,
2782 delta: torch.Tensor | None,
2783 blocks: torch.Tensor,
2784 norm_weight: torch.Tensor,
2785 qk_weight: torch.Tensor,
2786 output_norm_weight: torch.Tensor | None,
2787 num_blocks: int,
2788 block_write_idx: int,
2789 eps: float,
2790 output_norm_eps: float,
2791 ) -> torch.Tensor:
2792 output = torch.empty_like(prefix)
2793 torch.ops._C.kimi_k3_attn_res(
2794 prefix,
2795 delta,
2796 blocks,
2797 norm_weight,
2798 qk_weight,
2799 output_norm_weight,
2800 output,
2801 num_blocks,
2802 block_write_idx,
2803 eps,
2804 output_norm_eps,
2805 )
2806 return output
2807
2808
2809 def concat_and_cache_mla_rope_fused(
2810 positions: torch.Tensor,
2811 q_pe: torch.Tensor,
2812 k_pe: torch.Tensor,
2813 kv_c: torch.Tensor,
2814 cos_sin_cache: torch.Tensor,
2815 is_neox: bool,
2816 slot_mapping: torch.Tensor,
2817 kv_cache: torch.Tensor,
2818 kv_cache_dtype: str,
2819 kv_cache_scale: torch.Tensor,
2820 ) -> None:
2821 torch.ops._C_cache_ops.concat_and_cache_mla_rope_fused(
2822 positions,
2823 q_pe,
2824 k_pe,
2825 kv_c,
2826 cos_sin_cache,
2827 is_neox,
2828 slot_mapping,
2829 kv_cache,
2830 kv_cache_dtype,
2831 kv_cache_scale,
2832 )
2833
2834
2835 def swap_blocks(
2836 src: torch.Tensor,
2837 dst: torch.Tensor,
2838 block_size_in_bytes: int,
2839 block_mapping: torch.Tensor,
2840 ) -> None:
2841 """Copy specific blocks from one tensor to another.
2842
2843 This method assumes each of the two input tensors is composed of
2844 consecutive contiguous blocks, of size block_size_in_bytes.
2845 i.e. the memory layout for each tensor is:
2846 [block0] [block1] ... [block N]
2847
2848 block_mapping determines the subset of blocks to copy of the source tensor,
2849 and their matching destination block number on the destination tensor.
2850 block_mapping is expected to be a tensor of shape (num_blocks_to_copy, 2)
2851 where each block_mapping[i] represents a single copy operation, copying
2852 block #block_mapping[i][0] from the source tensor
2853 to block #block_mapping[i][1] on the destination tensor.
2854 block_mapping should have dtype int64.
2855
2856 The source and the destination tensors can be either on cpu or gpu,
2857 but not both on cpu.
2858 the block mapping tensor must on cpu.
2859 """
2860 torch.ops._C_cache_ops.swap_blocks(src, dst, block_size_in_bytes, block_mapping)
2861
2862
2863 def swap_blocks_batch(
2864 src_ptrs: torch.Tensor,
2865 dst_ptrs: torch.Tensor,
2866 sizes: torch.Tensor,
2867 is_src_access_order_any: bool = False,
2868 ) -> None:
2869 """Batch version of swap_blocks: submit all copies in a single driver call.
2870
2871 Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i]
2872 of sizes[i] bytes. All three tensors must be CPU tensors with the
2873 platform-appropriate pointer dtype: int64 on CUDA/ROCm (required by
2874 cache_kernels.cu) and uint64 on XPU (required by the XPU DMA engine).
2875 On CUDA 12.8+ this uses cuMemcpyBatchAsync for minimal submission
2876 overhead; on older CUDA it falls back to a loop of cudaMemcpyAsync.
2877
2878 is_src_access_order_any: if True, pass CU_MEMCPY_SRC_ACCESS_ORDER_ANY to
2879 cuMemcpyBatchAsync, letting the DMA engine prefetch source bytes
2880 out of stream order. Only safe when no GPU stream is concurrently
2881 writing to the source. Defaults to False (STREAM ordering), which
2882 is always safe.
2883 """
2884 if current_platform.is_xpu():
2885 torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes)
2886 else:
2887 torch.ops._C_cache_ops.swap_blocks_batch(
2888 src_ptrs, dst_ptrs, sizes, is_src_access_order_any
2889 )
2890
2891
2892 def convert_fp8(
2893 output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8"
2894 ) -> None:
2895 torch.ops._C_cache_ops.convert_fp8(output, input, scale, kv_dtype)
2896
2897
2898 def gather_and_maybe_dequant_cache(
2899 src_cache: torch.Tensor,
2900 dst: torch.Tensor,
2901 block_table: torch.Tensor,
2902 cu_seq_lens: torch.Tensor,
2903 token_to_seq: torch.Tensor,
2904 num_tokens: int,
2905 kv_cache_dtype: str,
2906 scale: torch.Tensor,
2907 seq_starts: torch.Tensor | None = None,
2908 ) -> None:
2909 torch.ops._C_cache_ops.gather_and_maybe_dequant_cache(
2910 src_cache,
2911 dst,
2912 block_table,
2913 cu_seq_lens,
2914 token_to_seq,
2915 num_tokens,
2916 kv_cache_dtype,
2917 scale,
2918 seq_starts,
2919 )
2920
2921
2922 def cp_gather_cache(
2923 src_cache: torch.Tensor,
2924 dst: torch.Tensor,
2925 block_table: torch.Tensor,
2926 cu_seq_lens: torch.Tensor,
2927 batch_size: int,
2928 seq_starts: torch.Tensor | None = None,
2929 ) -> None:
2930 torch.ops._C_cache_ops.cp_gather_cache(
2931 src_cache, dst, block_table, cu_seq_lens, batch_size, seq_starts
2932 )
2933
2934
2935 def cp_gather_and_upconvert_fp8_kv_cache(
2936 src_cache: torch.Tensor,
2937 dst: torch.Tensor,
2938 block_table: torch.Tensor,
2939 workspace_starts: torch.Tensor,
2940 batch_size: int,
2941 seq_starts: torch.Tensor | None = None,
2942 host_cache: torch.Tensor | None = None,
2943 host_row_ids: torch.Tensor | None = None,
2944 device_row_ids: torch.Tensor | None = None,
2945 ) -> None:
2946 """Gather and upconvert FP8 KV cache to BF16 workspace.
2947
2948 Args:
2949 src_cache: FP8 KV cache [num_blocks, block_size, 656]
2950 dst: BF16 output workspace [total_tokens, 576]
2951 block_table: Block indices [num_reqs, max_blocks]
2952 workspace_starts: Workspace start offsets [num_reqs]
2953 batch_size: Number of requests
2954 seq_starts: Optional source sequence offsets [num_reqs]
2955 host_cache: Optional pinned host rows used for non-resident entries
2956 host_row_ids: Host source row for each remapped cache row
2957 device_row_ids: Device source row, or -1, for each remapped cache row
2958
2959 """
2960 torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache(
2961 src_cache,
2962 dst,
2963 block_table,
2964 workspace_starts,
2965 batch_size,
2966 seq_starts,
2967 host_cache,
2968 host_row_ids,
2969 device_row_ids,
2970 )
2971
2972
2973 def cp_gather_and_upconvert_nvfp4_kv_cache(
2974 src_cache: torch.Tensor,
2975 dst: torch.Tensor,
2976 block_table: torch.Tensor,
2977 workspace_starts: torch.Tensor,
2978 batch_size: int,
2979 ) -> None:
2980 """Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace.
2981
2982 Args:
2983 src_cache: NVFP4 KV cache [num_blocks, block_size, 352] (uint8)
2984 dst: BF16 output workspace [total_tokens, 576]
2985 block_table: Block indices [num_reqs, max_blocks]
2986 workspace_starts: Workspace start offsets [num_reqs]
2987 batch_size: Number of requests
2988
2989 """
2990 torch.ops._C_cache_ops.cp_gather_and_upconvert_nvfp4_kv_cache(
2991 src_cache,
2992 dst,
2993 block_table,
2994 workspace_starts,
2995 batch_size,
2996 )
2997
2998
2999 def concat_mla_q(
3000 ql_nope: torch.Tensor,
3001 q_pe: torch.Tensor,
3002 q_out: torch.Tensor,
3003 ) -> None:
3004 """Concatenate query nope and rope for MLA/DSA attention.
3005
3006 Args:
3007 ql_nope: Query nope component [num_tokens, num_heads, nope_dim]
3008 q_pe: Query rope component [num_tokens, num_heads, rope_dim]
3009 q_out: Output tensor [num_tokens, num_heads, nope_dim + rope_dim]
3010
3011 """
3012 torch.ops._C_cache_ops.concat_mla_q(ql_nope, q_pe, q_out)
3013
3014
3015 def indexer_k_quant_and_cache(
3016 k: torch.Tensor,
3017 kv_cache: torch.Tensor,
3018 slot_mapping: torch.Tensor,
3019 quant_block_size: int,
3020 kv_cache_dtype: str,
3021 ) -> None:
3022 torch.ops._C_cache_ops.indexer_k_quant_and_cache(
3023 k, kv_cache, slot_mapping, quant_block_size, kv_cache_dtype
3024 )
3025
3026
3027 def top_k_per_row_prefill(
3028 logits: torch.Tensor,
3029 cu_seqlen_ks: torch.Tensor,
3030 cu_seqlen_ke: torch.Tensor,
3031 raw_topk_indices: torch.Tensor,
3032 num_rows: int,
3033 stride0: int,
3034 stride1: int,
3035 topk_tokens: int,
3036 ) -> None:
3037 torch.ops._C.top_k_per_row_prefill(
3038 logits,
3039 cu_seqlen_ks,
3040 cu_seqlen_ke,
3041 raw_topk_indices,
3042 num_rows,
3043 stride0,
3044 stride1,
3045 topk_tokens,
3046 )
3047
3048
3049 def top_k_per_row_decode(
3050 logits: torch.Tensor,
3051 next_n: int,
3052 seq_lens: torch.Tensor,
3053 raw_topk_indices: torch.Tensor,
3054 num_rows: int,
3055 stride0: int,
3056 stride1: int,
3057 topk_tokens: int,
3058 ) -> None:
3059 torch.ops._C.top_k_per_row_decode(
3060 logits,
3061 next_n,
3062 seq_lens,
3063 raw_topk_indices,
3064 num_rows,
3065 stride0,
3066 stride1,
3067 topk_tokens,
3068 )
3069
3070
3071 def cp_gather_indexer_k_quant_cache(
3072 kv_cache: torch.Tensor,
3073 dst_k: torch.Tensor,
3074 dst_scale: torch.Tensor,
3075 block_table: torch.Tensor,
3076 cu_seq_lens: torch.Tensor,
3077 ) -> None:
3078 torch.ops._C_cache_ops.cp_gather_indexer_k_quant_cache(
3079 kv_cache, dst_k, dst_scale, block_table, cu_seq_lens
3080 )
3081
3082
3083 def get_device_attribute(attribute: int, device: int) -> int:
3084 return torch.ops._C_cuda_utils.get_device_attribute(attribute, device)
3085
3086
3087 def get_max_shared_memory_per_block_device_attribute(device: int) -> int:
3088 # ruff: noqa: E501
3089 return torch.ops._C_cuda_utils.get_max_shared_memory_per_block_device_attribute(
3090 device
3091 )
3092
3093
3094 # custom ar
3095 def init_custom_ar(
3096 ipc_tensors: list[torch.Tensor],
3097 rank_data: torch.Tensor,
3098 rank: int,
3099 fully_connected: bool,
3100 ) -> int:
3101 return torch.ops._C_custom_ar.init_custom_ar(
3102 ipc_tensors, rank_data, rank, fully_connected
3103 )
3104
3105
3106 def all_reduce(
3107 fa: int,
3108 inp: torch.Tensor,
3109 out: torch.Tensor,
3110 reg_buffer: int,
3111 reg_buffer_sz_bytes: int,
3112 ) -> None:
3113 torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
3114
3115
3116 def custom_all_gather(
3117 fa: int,
3118 inp: torch.Tensor,
3119 out: torch.Tensor,
3120 reg_buffer: int,
3121 reg_buffer_sz_bytes: int,
3122 ) -> None:
3123 torch.ops._C_custom_ar.custom_all_gather(
3124 fa, inp, out, reg_buffer, reg_buffer_sz_bytes
3125 )
3126
3127
3128 def mnnvl_lamport_all_gather(
3129 fa: int,
3130 inp: torch.Tensor,
3131 out: torch.Tensor,
3132 local_buffer: int,
3133 multicast_buffer: int,
3134 epoch_buffer: int,
3135 stage_sz_bytes: int,
3136 ) -> None:
3137 torch.ops._C_custom_ar.mnnvl_lamport_all_gather(
3138 fa,
3139 inp,
3140 out,
3141 local_buffer,
3142 multicast_buffer,
3143 epoch_buffer,
3144 stage_sz_bytes,
3145 )
3146
3147
3148 def custom_reduce_scatter(
3149 fa: int,
3150 inp: torch.Tensor,
3151 out: torch.Tensor,
3152 reg_buffer: int,
3153 reg_buffer_sz_bytes: int,
3154 ) -> None:
3155 torch.ops._C_custom_ar.custom_reduce_scatter(
3156 fa, inp, out, reg_buffer, reg_buffer_sz_bytes
3157 )
3158
3159
3160 def mnnvl_lamport_reduce_scatter(
3161 fa: int,
3162 inp: torch.Tensor,
3163 out: torch.Tensor,
3164 local_buffer: int,
3165 epoch_buffer: int,
3166 stage_sz_bytes: int,
3167 ) -> None:
3168 torch.ops._C_custom_ar.mnnvl_lamport_reduce_scatter(
3169 fa, inp, out, local_buffer, epoch_buffer, stage_sz_bytes
3170 )
3171
3172
3173 def mnnvl_multimem_reduce_scatter(
3174 fa: int,
3175 inp: torch.Tensor,
3176 out: torch.Tensor,
3177 local_buffer: int,
3178 multicast_buffer: int,
3179 stage_sz_bytes: int,
3180 block_limit: int,
3181 ) -> None:
3182 torch.ops._C_custom_ar.mnnvl_multimem_reduce_scatter(
3183 fa,
3184 inp,
3185 out,
3186 local_buffer,
3187 multicast_buffer,
3188 stage_sz_bytes,
3189 block_limit,
3190 )
3191
3192
3193 def dispose(fa: int) -> None:
3194 torch.ops._C_custom_ar.dispose(fa)
3195
3196
3197 def meta_size() -> int:
3198 return torch.ops._C_custom_ar.meta_size()
3199
3200
3201 def register_buffer(fa: int, ipc_tensors: list[int]) -> None:
3202 return torch.ops._C_custom_ar.register_buffer(fa, ipc_tensors)
3203
3204
3205 def get_graph_buffer_ipc_meta(fa: int) -> tuple[list[int], list[int]]:
3206 return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa)
3207
3208
3209 def register_graph_buffers(
3210 fa: int, handles: list[list[int]], offsets: list[list[int]]
3211 ) -> None:
3212 torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets)
3213
3214
3215 def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]:
3216 return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size)
3217
3218
3219 def open_mem_handle(mem_handle: torch.Tensor):
3220 return torch.ops._C_custom_ar.open_mem_handle(mem_handle)
3221
3222
3223 def free_shared_buffer(ptr: int) -> None:
3224 torch.ops._C_custom_ar.free_shared_buffer(ptr)
3225
3226
3227 # quick all reduce
3228 def init_custom_qr(rank: int, world_size: int, qr_max_size: int | None = None) -> int:
3229 return torch.ops._C_custom_ar.init_custom_qr(rank, world_size, qr_max_size)
3230
3231
3232 def qr_destroy(fa: int) -> None:
3233 torch.ops._C_custom_ar.qr_destroy(fa)
3234
3235
3236 def qr_all_reduce(
3237 fa: int,
3238 inp: torch.Tensor,
3239 out: torch.Tensor,
3240 quant_level: int,
3241 cast_bf2half: bool = False,
3242 ) -> None:
3243 torch.ops._C_custom_ar.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half)
3244
3245
3246 def qr_get_handle(fa: int) -> torch.Tensor:
3247 return torch.ops._C_custom_ar.qr_get_handle(fa)
3248
3249
3250 def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None:
3251 return torch.ops._C_custom_ar.qr_open_handles(fa, handles)
3252
3253
3254 def qr_max_size() -> int:
3255 return torch.ops._C_custom_ar.qr_max_size()
3256
3257
3258 def sm100_cutlass_mla_decode(
3259 out: torch.Tensor,
3260 lse: torch.Tensor,
3261 q_nope: torch.Tensor,
3262 q_pe: torch.Tensor,
3263 kv_c_and_k_pe_cache: torch.Tensor,
3264 seq_lens: torch.Tensor,
3265 page_table: torch.Tensor,
3266 workspace: torch.Tensor,
3267 scale: float,
3268 num_kv_splits: int,
3269 ) -> torch.Tensor:
3270 torch.ops._C.sm100_cutlass_mla_decode(
3271 out,
3272 lse,
3273 q_nope,
3274 q_pe,
3275 kv_c_and_k_pe_cache,
3276 seq_lens,
3277 page_table,
3278 workspace,
3279 scale,
3280 num_kv_splits,
3281 )
3282 return out
3283
3284
3285 def sm100_cutlass_mla_get_workspace_size(
3286 max_seq_len: int, num_batches: int, sm_count: int, num_kv_splits: int
3287 ) -> int:
3288 return torch.ops._C.sm100_cutlass_mla_get_workspace_size(
3289 max_seq_len, num_batches, sm_count, num_kv_splits
3290 )
3291
3292
3293 def dsv3_fused_a_gemm(
3294 output: torch.Tensor,
3295 mat_a: torch.Tensor,
3296 mat_b: torch.Tensor,
3297 enable_pdl: bool = False,
3298 ) -> None:
3299 """Low-latency fused-A-style GEMM (SM 9.0+, BF16, 1-16 tokens).
3300
3301 Computes ``output = mat_a @ mat_b`` for the compiled Kimi K3 and
3302 DeepSeek V3 projection shapes. ``mat_a`` and ``output`` are row-major;
3303 ``mat_b`` is the column-major transposed weight. ``enable_pdl`` permits
3304 programmatic dependent launch for callers that have validated it.
3305
3306 Requires SM 9.0+.
3307 """
3308 torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl)
3309
3310
3311 def weight_packed_linear_cpu(
3312 x: torch.Tensor,
3313 packed_weight: torch.Tensor,
3314 n: int,
3315 bias: torch.Tensor | None,
3316 ) -> torch.Tensor:
3317 output = torch.empty((*x.shape[0:-1], n), dtype=x.dtype)
3318 torch.ops._C.weight_packed_linear(
3319 output.reshape(-1, n), x.reshape(-1, x.size(-1)), packed_weight, bias, True
3320 )
3321 return output
3322
3323
3324 class CPUQuantMethod(IntEnum):
3325 UNQUANT = 0
3326 INT8_W8A8 = 1
3327 FP8_W8A16 = 2
3328 INT4_W4A8 = 3
3329 MXFP4 = 4
3330 FP8_W8A8 = 5
3331
3332
3333 if hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
3334
3335 @register_fake("_C::dynamic_4bit_int_moe")
3336 def dynamic_4bit_int_moe_fake(
3337 x: torch.Tensor,
3338 topk_ids: torch.Tensor,
3339 topk_weights: torch.Tensor,
3340 w13_packed: torch.Tensor,
3341 w2_packed: torch.Tensor,
3342 hidden_size: int,
3343 intermediate_size: int,
3344 group_size: int,
3345 apply_router_weight_on_input: bool,
3346 activation_kind: int,
3347 ) -> torch.Tensor:
3348 return x.new_empty((x.size(0), hidden_size))
3349
3350
3351 def fused_experts_cpu(
3352 out: torch.Tensor,
3353 hidden_states: torch.Tensor,
3354 w1: torch.Tensor,
3355 w2: torch.Tensor,
3356 topk_weights: torch.Tensor,
3357 topk_ids: torch.Tensor,
3358 moe_comp_method: CPUQuantMethod,
3359 w1_scale: torch.Tensor | None,
3360 w2_scale: torch.Tensor | None,
3361 w1_zero: torch.Tensor | None,
3362 w2_zero: torch.Tensor | None,
3363 block_size: list[int] | None,
3364 w1_bias: torch.Tensor | None = None,
3365 w2_bias: torch.Tensor | None = None,
3366 alpha: float | None = None,
3367 limit: float | None = None,
3368 is_vnni: bool = True,
3369 a1_scale: torch.Tensor | None = None,
3370 ) -> None:
3371 torch.ops._C.fused_experts_cpu(
3372 out,
3373 hidden_states,
3374 w1,
3375 w2,
3376 topk_weights,
3377 topk_ids,
3378 moe_comp_method,
3379 w1_scale,
3380 w2_scale,
3381 w1_zero,
3382 w2_zero,
3383 a1_scale,
3384 block_size,
3385 w1_bias,
3386 w2_bias,
3387 alpha,
3388 limit,
3389 is_vnni,
3390 )
3391
3392
3393 if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
3394
3395 @register_fake("_C::int8_scaled_mm_with_quant")
3396 def int8_scaled_mm_with_quant_fake(
3397 mat1: torch.Tensor,
3398 mat2: torch.Tensor,
3399 scales2: torch.Tensor,
3400 bias: torch.Tensor | None,
3401 out_dtype: torch.dtype,
3402 is_vnni: bool,
3403 ) -> torch.Tensor:
3404 M = mat1.size(0)
3405 N = mat2.size(0)
3406 return torch.empty((M, N), dtype=out_dtype)
3407
3408
3409 class CPUQuantAlgo(IntEnum):
3410 AWQ = 0
3411 GPTQ = 1
3412
3413
3414 if hasattr(torch.ops._C, "convert_weight_packed_scale_zp"):
3415
3416 @register_fake("_C::convert_weight_packed_scale_zp")
3417 def convert_weight_packed_scale_zp_fake(
3418 qweight: torch.Tensor,
3419 qzeros: torch.Tensor,
3420 scales: torch.Tensor,
3421 quant_method_4bit: CPUQuantAlgo,
3422 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
3423 return (
3424 torch.empty_like(qweight),
3425 torch.empty_like(qzeros),
3426 torch.empty_like(scales),
3427 )
3428
3429
3430 def convert_weight_packed_scale_zp(
3431 qweight: torch.Tensor,
3432 qzeros: torch.Tensor,
3433 scales: torch.Tensor,
3434 quant_method_4bit: CPUQuantAlgo,
3435 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
3436 return torch.ops._C.convert_weight_packed_scale_zp(
3437 qweight,
3438 qzeros,
3439 scales,
3440 quant_method_4bit,
3441 )
3442
3443
3444 if hasattr(torch.ops._C, "int4_scaled_mm_cpu"):
3445
3446 @register_fake("_C::int4_scaled_mm_cpu")
3447 def int4_scaled_mm_cpu_fake(
3448 x: torch.Tensor,
3449 w: torch.Tensor,
3450 w_zeros: torch.Tensor,
3451 w_scales: torch.Tensor,
3452 bias: torch.Tensor | None,
3453 ) -> torch.Tensor:
3454 N = w_scales.size(0) * w_scales.size(-1)
3455 return torch.empty((x.size(0), N), dtype=x.dtype, device=x.device)
3456
3457
3458 def int4_scaled_mm_cpu(
3459 x: torch.Tensor,
3460 w: torch.Tensor,
3461 w_zeros: torch.Tensor,
3462 w_scales: torch.Tensor,
3463 bias: torch.Tensor | None,
3464 ) -> torch.Tensor:
3465 x_shape = x.shape
3466 x_2d = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x
3467
3468 out = torch.ops._C.int4_scaled_mm_cpu(
3469 x_2d,
3470 w,
3471 w_zeros,
3472 w_scales,
3473 bias,
3474 )
3475 out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out
3476 return out
3477
3478
3479 if hasattr(torch.ops._C, "fp8_scaled_mm_cpu"):
3480
3481 @register_fake("_C::fp8_scaled_mm_cpu")
3482 def fp8_scaled_mm_cpu_fake(
3483 mat1: torch.Tensor,
3484 mat2: torch.Tensor,
3485 scales2: torch.Tensor,
3486 block_size: list[int],
3487 bias: torch.Tensor | None,
3488 out_dtype: torch.dtype,
3489 is_vnni: bool,
3490 ) -> torch.Tensor:
3491 M = mat1.size(0)
3492 N = mat2.size(0)
3493 return torch.empty((M, N), dtype=out_dtype, device=mat1.device)
3494
3495
3496 _supports_cpu_fp8_w8a16 = bool(hasattr(torch.ops._C, "fp8_scaled_mm_cpu"))
3497
3498
3499 def fp8_scaled_mm_cpu(
3500 mat1: torch.Tensor,
3501 mat2: torch.Tensor,
3502 scales2: torch.Tensor,
3503 block_size: list[int],
3504 bias: torch.Tensor | None,
3505 out_dtype: torch.dtype,
3506 is_vnni: bool,
3507 ) -> torch.Tensor:
3508 return torch.ops._C.fp8_scaled_mm_cpu(
3509 mat1, mat2, scales2, block_size, bias, out_dtype, is_vnni
3510 )
3511
3512
3513 # FP8 W8A8 CPU kernels
3514 @cache
3515 def cpu_has_amx_fp8() -> bool:
3516 """Whether this CPU has native AMX-FP8 MMA support."""
3517 if not hasattr(torch.ops._C, "cpu_has_amx_fp8"):
3518 return False
3519 return bool(torch.ops._C.cpu_has_amx_fp8())
3520
3521
3522 if hasattr(torch.ops._C, "fp8_scaled_mm_with_quant"):
3523
3524 @register_fake("_C::fp8_scaled_mm_with_quant")
3525 def fp8_scaled_mm_with_quant_fake(
3526 act: torch.Tensor,
3527 act_scales: torch.Tensor | None,
3528 channelwise: bool,
3529 weight: torch.Tensor,
3530 weight_scales: torch.Tensor,
3531 bias: torch.Tensor | None,
3532 output_dtype: torch.dtype,
3533 ) -> torch.Tensor:
3534 M = act.reshape(-1, act.size(-1)).size(0)
3535 N = weight.size(0) * weight.size(-1)
3536 return torch.empty((M, N), dtype=output_dtype, device=act.device)
3537
3538
3539 def fp8_scaled_mm_with_quant(
3540 act: torch.Tensor,
3541 act_scales: torch.Tensor | None,
3542 channelwise: bool,
3543 weight: torch.Tensor,
3544 weight_scales: torch.Tensor,
3545 bias: torch.Tensor | None,
3546 output_dtype: torch.dtype,
3547 ) -> torch.Tensor:
3548 return torch.ops._C.fp8_scaled_mm_with_quant(
3549 act, act_scales, channelwise, weight, weight_scales, bias, output_dtype
3550 )
3551
3552
3553 def chunk_gated_delta_rule_cpu(
3554 query: torch.Tensor,
3555 key: torch.Tensor,
3556 value: torch.Tensor,
3557 g: torch.Tensor,
3558 beta: torch.Tensor,
3559 initial_state: torch.Tensor,
3560 output_final_state: bool,
3561 cu_seqlens: torch.Tensor,
3562 head_first: bool,
3563 use_qk_l2norm_in_kernel: bool,
3564 initial_state_indices: torch.Tensor,
3565 eps: float = 1e-5,
3566 ) -> tuple[torch.Tensor, torch.Tensor]:
3567 return torch.ops._C.chunk_gated_delta_rule_cpu(
3568 query,
3569 key,
3570 value,
3571 g,
3572 beta,
3573 initial_state,
3574 output_final_state,
3575 cu_seqlens,
3576 head_first,
3577 use_qk_l2norm_in_kernel,
3578 initial_state_indices,
3579 eps,
3580 )
3581
3582
3583 def fused_sigmoid_gating_delta_rule_update_cpu(
3584 A_log: torch.Tensor,
3585 dt_bias: torch.Tensor,
3586 q: torch.Tensor,
3587 k: torch.Tensor,
3588 v: torch.Tensor,
3589 a: torch.Tensor,
3590 b: torch.Tensor,
3591 initial_state_source: torch.Tensor,
3592 initial_state_indices: torch.Tensor,
3593 cu_seqlens: torch.Tensor,
3594 use_qk_l2norm_in_kernel: bool,
3595 softplus_beta: float = 1.0,
3596 softplus_threshold: float = 20.0,
3597 ) -> torch.Tensor:
3598 return torch.ops._C.fused_sigmoid_gating_delta_rule_update_cpu(
3599 A_log,
3600 dt_bias,
3601 q,
3602 k,
3603 v,
3604 a,
3605 b,
3606 initial_state_source,
3607 initial_state_indices,
3608 cu_seqlens,
3609 use_qk_l2norm_in_kernel,
3610 softplus_beta,
3611 softplus_threshold,
3612 )
3613
3614
3615 def fused_sigmoid_gating_delta_rule_update_spec_cpu(
3616 A_log: torch.Tensor,
3617 dt_bias: torch.Tensor,
3618 q: torch.Tensor,
3619 k: torch.Tensor,
3620 v: torch.Tensor,
3621 a: torch.Tensor,
3622 b: torch.Tensor,
3623 initial_state_source: torch.Tensor,
3624 spec_state_indices: torch.Tensor,
3625 num_accepted_tokens: torch.Tensor,
3626 cu_seqlens: torch.Tensor,
3627 use_qk_l2norm_in_kernel: bool,
3628 softplus_beta: float = 1.0,
3629 softplus_threshold: float = 20.0,
3630 ) -> torch.Tensor:
3631 return torch.ops._C.fused_sigmoid_gating_delta_rule_update_spec_cpu(
3632 A_log,
3633 dt_bias,
3634 q,
3635 k,
3636 v,
3637 a,
3638 b,
3639 initial_state_source,
3640 spec_state_indices,
3641 num_accepted_tokens,
3642 cu_seqlens,
3643 use_qk_l2norm_in_kernel,
3644 softplus_beta,
3645 softplus_threshold,
3646 )
3647
3648
3649 def fused_gdn_gating_cpu(
3650 A_log: torch.Tensor,
3651 a: torch.Tensor,
3652 b: torch.Tensor,
3653 dt_bias: torch.Tensor,
3654 ) -> tuple[torch.Tensor, torch.Tensor]:
3655 return torch.ops._C.fused_gdn_gating_cpu(
3656 A_log,
3657 a,
3658 b,
3659 dt_bias,
3660 )
3661
3662
3663 def causal_conv1d_weight_pack(
3664 weight: torch.Tensor,
3665 ) -> torch.Tensor:
3666 return torch.ops._C.causal_conv1d_weight_pack(
3667 weight,
3668 )
3669
3670
3671 def causal_conv1d_fwd_cpu(
3672 x: torch.Tensor,
3673 weight: torch.Tensor,
3674 bias: torch.Tensor | None,
3675 conv_states: torch.Tensor | None,
3676 query_start_loc: torch.Tensor | None,
3677 cache_indices: torch.Tensor | None,
3678 has_initial_state: torch.Tensor | None,
3679 silu_activation: bool,
3680 is_weight_packed: bool,
3681 ) -> torch.Tensor:
3682 return torch.ops._C.causal_conv1d_fwd_cpu(
3683 x,
3684 weight,
3685 bias,
3686 conv_states,
3687 query_start_loc,
3688 cache_indices,
3689 has_initial_state,
3690 silu_activation,
3691 -1,
3692 is_weight_packed,
3693 )
3694
3695
3696 def causal_conv1d_update_cpu(
3697 x: torch.Tensor,
3698 conv_states: torch.Tensor,
3699 weight: torch.Tensor,
3700 bias: torch.Tensor | None,
3701 silu_activation: bool,
3702 conv_state_indices: torch.Tensor | None,
3703 is_weight_packed: bool,
3704 num_accepted_tokens: torch.Tensor | None = None,
3705 ) -> torch.Tensor:
3706 return torch.ops._C.causal_conv1d_update_cpu(
3707 x,
3708 conv_states,
3709 weight,
3710 bias,
3711 silu_activation,
3712 num_accepted_tokens,
3713 conv_state_indices,
3714 -1,
3715 is_weight_packed,
3716 )
3717
3718
3719 class CPUDNNLGEMMHandler:
3720 def __init__(self) -> None:
3721 self.handler_tensor: torch.Tensor | None = None
3722 self.n = -1
3723 self.k = -1
3724 self.dtor = torch.ops._C.release_dnnl_matmul_handler
3725
3726 def __del__(self):
3727 if self.handler_tensor is not None:
3728 self.dtor(self.handler_tensor.item())
3729
3730
3731 _supports_onednn = bool(hasattr(torch.ops._C, "create_onednn_mm_handler"))
3732
3733
3734 def is_onednn_acl_supported():
3735 return torch.ops._C.is_onednn_acl_supported()
3736
3737
3738 def create_onednn_mm(
3739 weight: torch.Tensor, # [K, N]
3740 primitive_cache_size: int = 128,
3741 ) -> CPUDNNLGEMMHandler:
3742 handler = CPUDNNLGEMMHandler()
3743 handler.k, handler.n = weight.size()
3744 # store the handler pointer in a tensor it doesn't get inlined
3745 handler.handler_tensor = torch.tensor(
3746 torch.ops._C.create_onednn_mm_handler(weight, primitive_cache_size),
3747 dtype=torch.int64,
3748 )
3749 return handler
3750
3751
3752 def onednn_mm(
3753 dnnl_handler: CPUDNNLGEMMHandler,
3754 x: torch.Tensor,
3755 bias: torch.Tensor | None,
3756 ) -> torch.Tensor:
3757 output = torch.empty((*x.shape[0:-1], dnnl_handler.n), dtype=x.dtype)
3758 torch.ops._C.onednn_mm(
3759 output, x.reshape(-1, dnnl_handler.k), bias, dnnl_handler.handler_tensor
3760 )
3761
3762 return output
3763
3764
3765 def create_onednn_scaled_mm(
3766 weight: torch.Tensor, # [K, N]
3767 weight_scales: torch.Tensor,
3768 output_type: torch.dtype,
3769 dynamic_quant: bool,
3770 use_azp: bool,
3771 primitive_cache_size: int = 128,
3772 ) -> CPUDNNLGEMMHandler:
3773 handler = CPUDNNLGEMMHandler()
3774 handler.k, handler.n = weight.size()
3775 # store the handler pointer in a tensor so it doesn't get inlined
3776 handler.handler_tensor = torch.tensor(
3777 torch.ops._C.create_onednn_scaled_mm_handler(
3778 weight,
3779 weight_scales,
3780 output_type,
3781 dynamic_quant,
3782 use_azp,
3783 primitive_cache_size,
3784 ),
3785 dtype=torch.int64,
3786 )
3787 return handler
3788
3789
3790 def onednn_scaled_int8_quant(
3791 input: torch.Tensor,
3792 scale: torch.Tensor | None = None,
3793 azp: torch.Tensor | None = None,
3794 symmetric: bool = True,
3795 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
3796 """Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
3797
3798 Args:
3799 input: The input tensor to be quantized to int8.
3800 scale: Optional scaling factor for the int8 quantization.
3801 When not provided, we invoke dynamic-per-token quantization.
3802 azp: Optional zero-point for the int8 quantization.
3803 Must be provided for asymmetric quantization if `scale` is provided.
3804 symmetric: Whether to use symmetric quantization (scale only, azp ignored).
3805
3806 Returns:
3807 Output int8 tensor, scales, and optionally azp.
3808
3809 """
3810 output = torch.empty_like(input, dtype=torch.int8)
3811 token_num = input.numel() // input.shape[-1]
3812 input = input.view((token_num, input.shape[-1]))
3813 if scale is not None:
3814 # static-per-tensor quantization.
3815 assert symmetric == (azp is None), (
3816 "azp must only be provided for asymmetric quantization."
3817 )
3818 torch.ops._C.static_scaled_int8_quant(output, input, scale, azp)
3819 return output, scale, azp
3820
3821 # dynamic-per-token quantization.
3822 input_scales = torch.empty((token_num, 1), device=input.device, dtype=torch.float32)
3823 input_azp = None if symmetric else torch.empty_like(input_scales, dtype=torch.int32)
3824 torch.ops._C.dynamic_scaled_int8_quant(output, input, input_scales, input_azp)
3825 return output, input_scales, input_azp
3826
3827
3828 def onednn_scaled_mm(
3829 dnnl_handler: CPUDNNLGEMMHandler,
3830 x: torch.Tensor,
3831 output: torch.Tensor,
3832 input_scale: torch.Tensor | None,
3833 input_zp: torch.Tensor | None,
3834 input_zp_adj: torch.Tensor | None,
3835 bias: torch.Tensor | None,
3836 ) -> torch.Tensor:
3837 torch.ops._C.onednn_scaled_mm(
3838 output,
3839 x,
3840 input_scale,
3841 input_zp,
3842 input_zp_adj,
3843 bias,
3844 dnnl_handler.handler_tensor,
3845 )
3846
3847 return output
3848
3849
3850 def cpu_attn_get_scheduler_metadata(
3851 num_reqs: int,
3852 num_heads: int,
3853 num_kv_heads: int,
3854 head_dim: int,
3855 seq_lens: torch.Tensor,
3856 dtype: torch.dtype,
3857 query_start_loc: torch.Tensor,
3858 causal: bool,
3859 sliding_window_size: int,
3860 isa: str,
3861 enable_kv_split: bool,
3862 dynamic_causal: torch.Tensor | None = None,
3863 kv_cache_dtype: str = "auto",
3864 ) -> torch.Tensor:
3865 scheduler_metadata = torch.ops._C.get_scheduler_metadata(
3866 num_reqs,
3867 num_heads,
3868 num_kv_heads,
3869 head_dim,
3870 seq_lens,
3871 dtype,
3872 query_start_loc,
3873 causal,
3874 sliding_window_size,
3875 isa,
3876 enable_kv_split,
3877 dynamic_causal,
3878 kv_cache_dtype,
3879 )
3880 return scheduler_metadata
3881
3882
3883 def cpu_attn_reshape_and_cache(
3884 key: torch.Tensor,
3885 value: torch.Tensor,
3886 key_cache: torch.Tensor,
3887 value_cache: torch.Tensor,
3888 slot_mapping: torch.Tensor,
3889 isa: str,
3890 k_scale: float = 1.0,
3891 v_scale: float = 1.0,
3892 kv_cache_dtype: str = "auto",
3893 ) -> None:
3894 torch.ops._C.cpu_attn_reshape_and_cache(
3895 key,
3896 value,
3897 key_cache,
3898 value_cache,
3899 slot_mapping,
3900 isa,
3901 k_scale,
3902 v_scale,
3903 kv_cache_dtype,
3904 )
3905
3906
3907 def cpu_attention_with_kv_cache(
3908 query: torch.Tensor,
3909 key_cache: torch.Tensor,
3910 value_cache: torch.Tensor,
3911 output: torch.Tensor,
3912 query_start_loc: torch.Tensor,
3913 seq_lens: torch.Tensor,
3914 scale: float,
3915 causal: bool,
3916 alibi_slopes: torch.Tensor | None,
3917 sliding_window: int,
3918 block_table: torch.Tensor,
3919 softcap: float,
3920 scheduler_metadata: torch.Tensor,
3921 s_aux: torch.Tensor | None,
3922 dynamic_causal: torch.Tensor | None = None,
3923 k_scale: float = 1.0,
3924 v_scale: float = 1.0,
3925 kv_cache_dtype: str = "auto",
3926 ) -> None:
3927 torch.ops._C.cpu_attention_with_kv_cache(
3928 query,
3929 key_cache,
3930 value_cache,
3931 output,
3932 query_start_loc,
3933 seq_lens,
3934 scale,
3935 causal,
3936 alibi_slopes,
3937 sliding_window,
3938 block_table,
3939 softcap,
3940 scheduler_metadata,
3941 s_aux,
3942 dynamic_causal,
3943 k_scale,
3944 v_scale,
3945 kv_cache_dtype,
3946 )
3947
3948
3949 def cpu_mla_decode(
3950 query: torch.Tensor,
3951 k_buffer: torch.Tensor,
3952 v_buffer: torch.Tensor,
3953 output: torch.Tensor,
3954 key: torch.Tensor | None,
3955 value: torch.Tensor | None,
3956 loc: torch.Tensor | None,
3957 attn_logits: torch.Tensor,
3958 req_to_token: torch.Tensor,
3959 req_pool_indices: torch.Tensor,
3960 seq_lens: torch.Tensor,
3961 sm_scale: float,
3962 logit_cap: float,
3963 is_cross_attn: bool,
3964 sliding_window_size: int,
3965 encoder_lens: torch.Tensor | None,
3966 sinks: torch.Tensor | None,
3967 ) -> None:
3968 torch.ops._C.decode_attention_cpu(
3969 query,
3970 k_buffer,
3971 v_buffer,
3972 output,
3973 key,
3974 value,
3975 loc,
3976 attn_logits,
3977 req_to_token,
3978 req_pool_indices,
3979 seq_lens,
3980 sm_scale,
3981 logit_cap,
3982 is_cross_attn,
3983 sliding_window_size,
3984 encoder_lens,
3985 sinks,
3986 )
3987
3988
3989 def cpu_mla_extend(
3990 q_extend: torch.Tensor,
3991 k_extend: torch.Tensor | None,
3992 v_extend: torch.Tensor | None,
3993 o_extend: torch.Tensor,
3994 k_buffer: torch.Tensor,
3995 v_buffer: torch.Tensor,
3996 req_to_token: torch.Tensor,
3997 req_pool_indices: torch.Tensor,
3998 seq_lens: torch.Tensor,
3999 extend_seq_lens: torch.Tensor,
4000 extend_start_loc: torch.Tensor,
4001 max_len_extend: int,
4002 sm_scale: float,
4003 logit_cap: float,
4004 is_cross_attn: bool,
4005 sliding_window_size: int,
4006 encoder_lens: torch.Tensor | None,
4007 sinks: torch.Tensor | None,
4008 tree_mask: torch.Tensor | None = None,
4009 ) -> None:
4010 torch.ops._C.extend_attention_cpu(
4011 q_extend,
4012 k_extend,
4013 v_extend,
4014 o_extend,
4015 k_buffer,
4016 v_buffer,
4017 req_to_token,
4018 req_pool_indices,
4019 seq_lens,
4020 extend_seq_lens,
4021 extend_start_loc,
4022 max_len_extend,
4023 sm_scale,
4024 logit_cap,
4025 is_cross_attn,
4026 sliding_window_size,
4027 encoder_lens,
4028 sinks,
4029 tree_mask,
4030 )
4031
4032
4033 def bmm_cpu(
4034 out: torch.Tensor,
4035 mat1: torch.Tensor,
4036 mat2: torch.Tensor,
4037 is_vnni: bool,
4038 scale: torch.Tensor | None = None,
4039 ) -> None:
4040 torch.ops._C.bmm_cpu(out, mat1, mat2, is_vnni, scale)
4041
4042
4043 def amx_mla_concat_and_cache(
4044 kv_c_normed: torch.Tensor,
4045 k_pe: torch.Tensor,
4046 kv_cache: torch.Tensor,
4047 slot_mapping: torch.Tensor,
4048 ) -> None:
4049 torch.ops._C.concat_and_cache_mla_cpu(kv_c_normed, k_pe, kv_cache, slot_mapping)
4050
4051
4052 if hasattr(torch.ops._C, "hc_pre_fused_cpu"):
4053
4054 @register_fake("_C::hc_pre_fused_cpu")
4055 def hc_pre_fused_cpu_fake(
4056 x: torch.Tensor,
4057 hc_fn: torch.Tensor,
4058 hc_scale: torch.Tensor,
4059 hc_base: torch.Tensor,
4060 hc_mult: int,
4061 sinkhorn_iters: int,
4062 rms_eps: float,
4063 hc_eps: float,
4064 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
4065 t, _, d = x.shape
4066 y = x.new_empty((t, d))
4067 post = x.new_empty((t, hc_mult), dtype=torch.float32)
4068 comb = x.new_empty((t, hc_mult, hc_mult), dtype=torch.float32)
4069 return y, post, comb
4070
4071
4072 def hc_pre_fused_cpu(
4073 x: torch.Tensor,
4074 hc_fn: torch.Tensor,
4075 hc_scale: torch.Tensor,
4076 hc_base: torch.Tensor,
4077 hc_mult: int,
4078 sinkhorn_iters: int,
4079 rms_eps: float,
4080 hc_eps: float,
4081 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
4082 return torch.ops._C.hc_pre_fused_cpu(
4083 x, hc_fn, hc_scale, hc_base, hc_mult, sinkhorn_iters, rms_eps, hc_eps
4084 )
4085
4086
4087 if hasattr(torch.ops._C, "hc_post_fused_cpu"):
4088
4089 @register_fake("_C::hc_post_fused_cpu")
4090 def hc_post_fused_cpu_fake(
4091 x: torch.Tensor,
4092 residual: torch.Tensor,
4093 post: torch.Tensor,
4094 comb: torch.Tensor,
4095 ) -> torch.Tensor:
4096 t, d = x.shape
4097 hc = residual.size(1)
4098 return x.new_empty((t, hc, d))
4099
4100
4101 def hc_post_fused_cpu(
4102 x: torch.Tensor,
4103 residual: torch.Tensor,
4104 post: torch.Tensor,
4105 comb: torch.Tensor,
4106 ) -> torch.Tensor:
4107 return torch.ops._C.hc_post_fused_cpu(x, residual, post, comb)
4108
4109
4110 if hasattr(torch.ops._C, "hc_head_fused_cpu"):
4111
4112 @register_fake("_C::hc_head_fused_cpu")
4113 def hc_head_fused_cpu_fake(
4114 x: torch.Tensor,
4115 hc_fn: torch.Tensor,
4116 hc_scale: torch.Tensor,
4117 hc_base: torch.Tensor,
4118 hc_eps: float,
4119 norm_eps: float,
4120 ) -> torch.Tensor:
4121 t, _, d = x.shape
4122 return x.new_empty((t, d))
4123
4124
4125 def hc_head_fused_cpu(
4126 x: torch.Tensor,
4127 hc_fn: torch.Tensor,
4128 hc_scale: torch.Tensor,
4129 hc_base: torch.Tensor,
4130 hc_eps: float,
4131 norm_eps: float,
4132 ) -> torch.Tensor:
4133 return torch.ops._C.hc_head_fused_cpu(x, hc_fn, hc_scale, hc_base, hc_eps, norm_eps)
4134
4135
4136 if hasattr(torch.ops._C, "fused_qnorm_rope_kv_insert_cpu"):
4137
4138 @register_fake("_C::fused_qnorm_rope_kv_insert_cpu")
4139 def fused_qnorm_rope_kv_insert_cpu_fake(
4140 q: torch.Tensor,
4141 kv: torch.Tensor,
4142 positions: torch.Tensor,
4143 swa_kv_cache_2d: torch.Tensor,
4144 slot_mapping: torch.Tensor,
4145 cos_sin_cache: torch.Tensor,
4146 q_head_padded: int,
4147 eps: float,
4148 cache_block_size: int,
4149 ) -> torch.Tensor:
4150 num_tokens_full, _, head_dim = q.shape
4151 return q.new_empty((num_tokens_full, q_head_padded, head_dim))
4152
4153
4154 def fused_qnorm_rope_kv_insert_cpu(
4155 q: torch.Tensor,
4156 kv: torch.Tensor,
4157 positions: torch.Tensor,
4158 swa_kv_cache_2d: torch.Tensor,
4159 slot_mapping: torch.Tensor,
4160 cos_sin_cache: torch.Tensor,
4161 q_head_padded: int,
4162 eps: float,
4163 cache_block_size: int,
4164 ) -> torch.Tensor:
4165 return torch.ops._C.fused_qnorm_rope_kv_insert_cpu(
4166 q,
4167 kv,
4168 positions,
4169 swa_kv_cache_2d,
4170 slot_mapping,
4171 cos_sin_cache,
4172 q_head_padded,
4173 eps,
4174 cache_block_size,
4175 )
4176
4177
4178 def flash_mla_with_kvcache_cpu(
4179 out: torch.Tensor,
4180 q: torch.Tensor,
4181 window_cache_2d: torch.Tensor,
4182 window_slots: torch.Tensor,
4183 window_block_size: int,
4184 compressed_cache_2d: torch.Tensor,
4185 compressed_slots: torch.Tensor,
4186 compressed_block_size: int,
4187 attn_sink: torch.Tensor,
4188 scale: float,
4189 ) -> None:
4190 torch.ops._C.flash_mla_with_kvcache_cpu(
4191 out,
4192 q,
4193 window_cache_2d,
4194 window_slots,
4195 window_block_size,
4196 compressed_cache_2d,
4197 compressed_slots,
4198 compressed_block_size,
4199 attn_sink,
4200 scale,
4201 )
4202
4203
4204 def save_partial_states_cpu(
4205 kv: torch.Tensor,
4206 score: torch.Tensor,
4207 ape: torch.Tensor,
4208 positions: torch.Tensor,
4209 state_cache: torch.Tensor,
4210 slot_mapping: torch.Tensor,
4211 ) -> None:
4212 torch.ops._C.save_partial_states_cpu(
4213 kv, score, ape, positions, state_cache, slot_mapping
4214 )
4215
4216
4217 def compress_norm_rope_store_cpu(
4218 state_cache: torch.Tensor,
4219 gather_slots: torch.Tensor,
4220 positions: torch.Tensor,
4221 kv_slot_mapping: torch.Tensor,
4222 rms_norm_weight: torch.Tensor,
4223 rms_norm_eps: float,
4224 cos_sin_cache: torch.Tensor,
4225 kv_cache_2d: torch.Tensor,
4226 kv_cache_block_size: int,
4227 compress_ratio: int,
4228 ) -> None:
4229 torch.ops._C.compress_norm_rope_store_cpu(
4230 state_cache,
4231 gather_slots,
4232 positions,
4233 kv_slot_mapping,
4234 rms_norm_weight,
4235 rms_norm_eps,
4236 cos_sin_cache,
4237 kv_cache_2d,
4238 kv_cache_block_size,
4239 compress_ratio,
4240 )
4241
4242
4243 def compress_norm_rope_store_indexer_cpu(
4244 state_cache: torch.Tensor,
4245 gather_slots: torch.Tensor,
4246 positions: torch.Tensor,
4247 kv_slot_mapping: torch.Tensor,
4248 rms_norm_weight: torch.Tensor,
4249 rms_norm_eps: float,
4250 cos_sin_cache: torch.Tensor,
4251 kv_cache_2d: torch.Tensor,
4252 kv_cache_block_size: int,
4253 compress_ratio: int,
4254 ) -> None:
4255 torch.ops._C.compress_norm_rope_store_indexer_cpu(
4256 state_cache,
4257 gather_slots,
4258 positions,
4259 kv_slot_mapping,
4260 rms_norm_weight,
4261 rms_norm_eps,
4262 cos_sin_cache,
4263 kv_cache_2d,
4264 kv_cache_block_size,
4265 compress_ratio,
4266 )
4267
4268
4269 if hasattr(torch.ops._C, "fp8_paged_mqa_logits_cpu"):
4270
4271 @register_fake("_C::fp8_paged_mqa_logits_cpu")
4272 def fp8_paged_mqa_logits_cpu_fake(
4273 q_fp8: torch.Tensor,
4274 kvcache_fp8: torch.Tensor,
4275 weight: torch.Tensor,
4276 seq_lens: torch.Tensor,
4277 page_table: torch.Tensor,
4278 block_size: int,
4279 max_seq_len: int,
4280 ) -> torch.Tensor:
4281 batch_size = q_fp8.size(0)
4282 return q_fp8.new_empty((batch_size, max_seq_len), dtype=torch.float32)
4283
4284
4285 def fp8_paged_mqa_logits_cpu(
4286 q_fp8: torch.Tensor,
4287 kvcache_fp8: torch.Tensor,
4288 weight: torch.Tensor,
4289 seq_lens: torch.Tensor,
4290 page_table: torch.Tensor,
4291 block_size: int,
4292 max_seq_len: int,
4293 ) -> torch.Tensor:
4294 return torch.ops._C.fp8_paged_mqa_logits_cpu(
4295 q_fp8, kvcache_fp8, weight, seq_lens, page_table, block_size, max_seq_len
4296 )
4297
4298
4299 def topk_transform_512_cpu(
4300 scores: torch.Tensor,
4301 seq_lens: torch.Tensor,
4302 page_tables: torch.Tensor,
4303 out_page_indices: torch.Tensor,
4304 page_size: int,
4305 out_raw_indices: torch.Tensor | None = None,
4306 ) -> None:
4307 torch.ops._C.topk_transform_512_cpu(
4308 scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices
4309 )
4310
4311
4312 def fused_indexer_q_rope_quant_cpu(
4313 positions: torch.Tensor,
4314 index_q: torch.Tensor,
4315 index_q_cos_sin_cache: torch.Tensor,
4316 index_q_fp8: torch.Tensor,
4317 index_weights: torch.Tensor,
4318 index_weights_softmax_scale: float,
4319 index_weights_head_scale: float,
4320 index_weights_out: torch.Tensor,
4321 ) -> None:
4322 torch.ops._C.fused_indexer_q_rope_quant_cpu(
4323 positions,
4324 index_q,
4325 index_q_cos_sin_cache,
4326 index_q_fp8,
4327 index_weights,
4328 index_weights_softmax_scale,
4329 index_weights_head_scale,
4330 index_weights_out,
4331 )
4332
4333
4334 if hasattr(torch.ops._C, "biased_topk_cpu"):
4335
4336 @register_fake("_C::biased_topk_cpu")
4337 def biased_topk_cpu_fake(
4338 hidden_states: torch.Tensor,
4339 gating_output: torch.Tensor,
4340 correction_bias: torch.Tensor,
4341 topk: int,
4342 renormalize: bool,
4343 scoring_func: str,
4344 num_fused_shared_experts: int,
4345 routed_scaling_factor: float | None,
4346 apply_routed_scaling_factor_on_output: bool,
4347 ) -> tuple[torch.Tensor, torch.Tensor]:
4348 num_tokens = hidden_states.size(0)
4349 topk_weights = hidden_states.new_empty((num_tokens, topk), dtype=torch.float32)
4350 topk_ids = hidden_states.new_empty((num_tokens, topk), dtype=torch.int32)
4351 return topk_weights, topk_ids
4352
4353
4354 def biased_topk_cpu(
4355 hidden_states: torch.Tensor,
4356 gating_output: torch.Tensor,
4357 correction_bias: torch.Tensor,
4358 topk: int,
4359 renormalize: bool,
4360 scoring_func: str,
4361 num_fused_shared_experts: int,
4362 routed_scaling_factor: float | None,
4363 apply_routed_scaling_factor_on_output: bool,
4364 ) -> tuple[torch.Tensor, torch.Tensor]:
4365 return torch.ops._C.biased_topk_cpu(
4366 hidden_states,
4367 gating_output,
4368 correction_bias,
4369 topk,
4370 renormalize,
4371 scoring_func,
4372 num_fused_shared_experts,
4373 routed_scaling_factor,
4374 apply_routed_scaling_factor_on_output,
4375 )
4376
4377
4378 if hasattr(torch.ops._C, "hash_topk_cpu"):
4379
4380 @register_fake("_C::hash_topk_cpu")
4381 def hash_topk_cpu_fake(
4382 gating_output: torch.Tensor,
4383 tid2eid: torch.Tensor,
4384 topk: int,
4385 scoring_func: str,
4386 num_fused_shared_experts: int,
4387 num_experts: int,
4388 routed_scaling_factor: float,
4389 ) -> tuple[torch.Tensor, torch.Tensor]:
4390 num_tokens = gating_output.size(0)
4391 topk_weights = gating_output.new_empty((num_tokens, topk), dtype=torch.float32)
4392 topk_ids = gating_output.new_empty((num_tokens, topk), dtype=torch.int32)
4393 return topk_weights, topk_ids
4394
4395
4396 def hash_topk_cpu(
4397 gating_output: torch.Tensor,
4398 tid2eid: torch.Tensor,
4399 topk: int,
4400 scoring_func: str,
4401 num_fused_shared_experts: int,
4402 num_experts: int,
4403 routed_scaling_factor: float,
4404 ) -> tuple[torch.Tensor, torch.Tensor]:
4405 return torch.ops._C.hash_topk_cpu(
4406 gating_output,
4407 tid2eid,
4408 topk,
4409 scoring_func,
4410 num_fused_shared_experts,
4411 num_experts,
4412 routed_scaling_factor,
4413 )
4414
4415
4416 if hasattr(torch.ops._C, "inverse_gptj_rope_o_proj_cpu"):
4417
4418 @register_fake("_C::inverse_gptj_rope_o_proj_cpu")
4419 def inverse_gptj_rope_o_proj_cpu_fake(
4420 o: torch.Tensor,
4421 positions: torch.Tensor,
4422 cos_sin_cache: torch.Tensor,
4423 rope_dim: int,
4424 ) -> torch.Tensor:
4425 return o.new_empty(o.shape, dtype=torch.float32)
4426
4427
4428 def inverse_gptj_rope_o_proj_cpu(
4429 o: torch.Tensor,
4430 positions: torch.Tensor,
4431 cos_sin_cache: torch.Tensor,
4432 rope_dim: int,
4433 ) -> torch.Tensor:
4434 return torch.ops._C.inverse_gptj_rope_o_proj_cpu(
4435 o, positions, cos_sin_cache, rope_dim
4436 )
4437
4438
4439 def cpu_gemm_wna16(
4440 input: torch.Tensor,
4441 q_weight: torch.Tensor,
4442 scales: torch.Tensor,
4443 zeros: torch.Tensor | None,
4444 bias: torch.Tensor | None,
4445 pack_factor: int,
4446 isa_hint: str,
4447 ) -> torch.Tensor:
4448 # Match int4_scaled_mm_cpu: flatten >2-D activations to [M, K] for the
4449 # C++ kernel, then restore the original leading dims on the output.
4450 x_shape = input.shape
4451 x_2d = input.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else input
4452 out = torch.empty((x_2d.size(0), scales.size(1)), dtype=input.dtype)
4453 torch.ops._C.cpu_gemm_wna16(
4454 x_2d,
4455 q_weight,
4456 out,
4457 scales,
4458 zeros,
4459 bias,
4460 pack_factor,
4461 isa_hint,
4462 )
4463 out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out
4464 return out
4465
4466
4467 def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor:
4468 out = torch.empty_like(input)
4469 torch.ops._C.activation_lut_bf16(out, input, activation)
4470 return out
4471
4472
4473 def cpu_prepack_moe_weight(
4474 weight: torch.Tensor,
4475 isa: str,
4476 ) -> torch.Tensor:
4477 output = torch.empty_like(weight)
4478 torch.ops._C.prepack_moe_weight(weight, output, isa)
4479 return output
4480
4481
4482 def cpu_prepack_moe_weight_int8(
4483 weight: torch.Tensor,
4484 isa: str,
4485 ) -> torch.Tensor:
4486 output = torch.empty_like(weight)
4487 torch.ops._C.prepack_moe_weight_int8(weight, output, isa)
4488 return output
4489
4490
4491 def cpu_fused_moe(
4492 out: torch.Tensor,
4493 input: torch.Tensor,
4494 w13: torch.Tensor,
4495 w2: torch.Tensor,
4496 w13_bias: torch.Tensor | None,
4497 w2_bias: torch.Tensor | None,
4498 topk_weights: torch.Tensor,
4499 topk_ids: torch.Tensor,
4500 act: str,
4501 isa: str,
4502 skip_weighted: bool = False,
4503 ) -> None:
4504 torch.ops._C.cpu_fused_moe(
4505 out,
4506 input,
4507 w13,
4508 w2,
4509 w13_bias,
4510 w2_bias,
4511 topk_weights,
4512 topk_ids,
4513 skip_weighted,
4514 act,
4515 isa,
4516 )
4517
4518
4519 def cpu_fused_moe_int8(
4520 out: torch.Tensor,
4521 input: torch.Tensor,
4522 w13: torch.Tensor,
4523 w2: torch.Tensor,
4524 w13_scale: torch.Tensor,
4525 w2_scale: torch.Tensor,
4526 w13_bias: torch.Tensor | None,
4527 w2_bias: torch.Tensor | None,
4528 topk_weights: torch.Tensor,
4529 topk_ids: torch.Tensor,
4530 act: str,
4531 isa: str,
4532 skip_weighted: bool = False,
4533 ) -> None:
4534 torch.ops._C.cpu_fused_moe_int8(
4535 out,
4536 input,
4537 w13,
4538 w2,
4539 w13_scale,
4540 w2_scale,
4541 w13_bias,
4542 w2_bias,
4543 topk_weights,
4544 topk_ids,
4545 skip_weighted,
4546 act,
4547 isa,
4548 )
4549
4550
4551 if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
4552
4553 @register_fake("_qutlass_C::matmul_mxf4_bf16_tn")
4554 def _fake_matmul_mxf4_bf16_tn(
4555 a: torch.Tensor,
4556 b: torch.Tensor,
4557 a_sf: torch.Tensor,
4558 b_sf: torch.Tensor,
4559 alpha: torch.Tensor,
4560 ):
4561 return a.new_empty(*a.shape[:-1], b.shape[0], dtype=torch.bfloat16)
4562
4563
4564 def matmul_mxf4_bf16_tn(
4565 a: torch.Tensor,
4566 b: torch.Tensor,
4567 a_sf: torch.Tensor,
4568 b_sf: torch.Tensor,
4569 alpha: torch.Tensor,
4570 ) -> torch.Tensor:
4571 return torch.ops._qutlass_C.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha)
4572
4573
4574 if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxQuest"):
4575
4576 @register_fake("_qutlass_C::fusedQuantizeMxQuest")
4577 def _fake_fused_quantize_mx_quest(
4578 a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor
4579 ):
4580 return xh_e2m1, xh_e8m0
4581
4582
4583 if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxAbsMax"):
4584
4585 @register_fake("_qutlass_C::fusedQuantizeMxAbsMax")
4586 def _fake_fused_quantize_mx_absmax(
4587 a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor
4588 ):
4589 return xh_e2m1, xh_e8m0
4590
4591
4592 def fusedQuantizeMx(
4593 a: torch.Tensor, b: torch.Tensor, *, method: Literal["quest", "abs_max"] = "quest"
4594 ) -> tuple[torch.Tensor, torch.Tensor]:
4595 if a.dim() == 0:
4596 raise ValueError("`a` must have at least 1 dimension.")
4597 if a.size(-1) % 32 != 0:
4598 raise ValueError(f"last dim of `a` must be divisible by 32, got {a.size(-1)}.")
4599 if b.device != a.device:
4600 raise ValueError("`a` and `b` must be on the same device.")
4601
4602 xh_e2m1 = torch.empty(
4603 *a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device
4604 )
4605
4606 rows, cols = a.numel() // a.size(-1), a.size(-1) // 32
4607 n_row_blocks = cdiv(rows, 128)
4608 n_col_blocks = cdiv(cols, 4)
4609 padded_rows = n_row_blocks * 128
4610 padded_cols = n_col_blocks * 4
4611
4612 xh_e8m0 = torch.empty(
4613 padded_rows, padded_cols, dtype=torch.float8_e8m0fnu, device=a.device
4614 )
4615
4616 if not hasattr(torch.ops, "_qutlass_C"):
4617 raise RuntimeError(
4618 "The `_qutlass_C` extension is not loaded. "
4619 "Make sure your custom op library is imported before calling fusedQuantizeMx."
4620 )
4621
4622 if method == "quest":
4623 return torch.ops._qutlass_C.fusedQuantizeMxQuest(a, b, xh_e2m1, xh_e8m0)
4624 elif method == "abs_max":
4625 return torch.ops._qutlass_C.fusedQuantizeMxAbsMax(a, b, xh_e2m1, xh_e8m0)
4626 else:
4627 raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'")
4628
4629
4630 if hasattr(torch.ops._qutlass_C, "fusedQuantizeNvAbsMax"):
4631
4632 @register_fake("_qutlass_C::fusedQuantizeNvAbsMax")
4633 def _fake_fused_quantize_nv_absmax(
4634 a: torch.Tensor,
4635 b: torch.Tensor,
4636 xh_e2m1: torch.Tensor,
4637 xh_e4m3: torch.Tensor,
4638 global_scale: torch.Tensor,
4639 ):
4640 return xh_e2m1, xh_e4m3
4641
4642
4643 def fusedQuantizeNv(
4644 a: torch.Tensor, b: torch.Tensor, global_scale: torch.Tensor
4645 ) -> tuple[torch.Tensor, torch.Tensor]:
4646 xh_e2m1 = torch.empty(
4647 *a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device
4648 )
4649
4650 rows, cols = a.numel() // a.size(-1), a.size(-1) // 16
4651 n_row_blocks = cdiv(rows, 128)
4652 n_col_blocks = cdiv(cols, 4)
4653 padded_rows = n_row_blocks * 128
4654 padded_cols = n_col_blocks * 4
4655 xh_e4m3 = torch.empty(
4656 padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device
4657 )
4658
4659 safeFusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale)
4660 return xh_e2m1, xh_e4m3
4661
4662
4663 @torch.library.custom_op(
4664 "vllm::safeFusedQuantizeNv", mutates_args=("xh_e2m1", "xh_e4m3")
4665 )
4666 def safeFusedQuantizeNv(
4667 a: torch.Tensor,
4668 b: torch.Tensor,
4669 xh_e2m1: torch.Tensor,
4670 xh_e4m3: torch.Tensor,
4671 global_scale: torch.Tensor,
4672 ) -> None:
4673 """Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place
4674 rather than returning them, to prevent torch 2.12+ errors that outputs of custom
4675 operators may not alias any inputs to the custom operator.
4676 """
4677 torch.ops._qutlass_C.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale)
4678 return
4679
4680
4681 def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
4682 """Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832)
4683 kernels. Note that these kernels exploit the recursive properties of
4684 Sylvester Hadamards, and therefore do not require transform weight data
4685
4686 Note that sylvester hadamard transforms are also symmetric, which means that
4687 this function is also applies the (transpose <=> inverse) transform.
4688
4689 Args:
4690 x: value to be transformed inplace
4691 inplace: modify value in place
4692
4693 Returns:
4694 value after transformation
4695
4696 """
4697 return torch.ops._C.hadacore_transform(x, inplace)
4698
4699
4700 if hasattr(torch.ops._C, "hadacore_transform"):
4701
4702 @register_fake("_C::hadacore_transform")
4703 def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor:
4704 return torch.empty_like(x) if not inplace else x
4705
4706
4707 if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
4708
4709 @register_fake("_C::minimax_allreduce_rms_qk")
4710 def _minimax_allreduce_rms_qk_fake(
4711 qkv: torch.Tensor,
4712 norm_weight_q: torch.Tensor,
4713 norm_weight_k: torch.Tensor,
4714 workspace: torch.Tensor,
4715 q_size: int,
4716 kv_size: int,
4717 rank: int,
4718 nranks: int,
4719 eps: float,
4720 ) -> tuple[torch.Tensor, torch.Tensor]:
4721 token_num = qkv.shape[0]
4722 return (
4723 torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device),
4724 torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device),
4725 )