iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity
Files
382970ee6ca490aeaaaf4e32c53695b581ff61ba
vLLM/vllm/_custom_ops.py
T
liuzhenweiandKunshang Ji a15dbbd63d [XPU][CI] fallback to 7-args of moe_align_block_size (#57855) 
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-09-21 13:15:29 +08:00

4794 lines
135 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from enum import IntEnum
from functools import cache
from typing import TYPE_CHECKING, Literal
import torch
import vllm.envs as envs
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.scalar_type import ScalarType
from vllm.utils.flashinfer import (
flashinfer_quant_nvfp4_8x4_sf_layout,
)
from vllm.utils.math_utils import cdiv
logger = init_logger(__name__)
current_platform.import_kernels()
if TYPE_CHECKING:
def register_fake(fn):
return lambda name: fn
else:
try:
from torch.library import register_fake
except ImportError:
from torch.library import impl_abstract as register_fake
# scaled_fp4_quant functional + out variant for torch.compile buffer management
def create_fp4_scale_tensor(
m: int,
n: int,
device: torch.device,
is_sf_swizzled_layout: bool,
) -> torch.Tensor:
"""Allocate the output scale tensor for scaled_fp4_quant.
When is_sf_swizzled_layout=True, we use rounded values to store the
swizzled scales. Due to the requirement of the Tensor Core, the minimum
tile is 128x4 for the scales. So, we first pad the scales to multiples
of 128 (rows) and 4 (cols). Then, the scales (in float8_e4m3fn) are
packed into an int32 for every 4 values. More:
https://docs.nvidia.com/cuda/parallel-thread-execution/
#tcgen05-mma-scale-factor-b-layout-4x
"""
from vllm.utils.math_utils import round_up
block_size = 16
if is_sf_swizzled_layout:
rounded_m = round_up(m, 128)
scale_n = n // block_size
rounded_n = round_up(scale_n, 4)
# The NVFP4 quant kernel explicitly zeroes every padded scale entry,
# so no separate zero-initialization kernel is required here.
return torch.empty(
(rounded_m, rounded_n // 4), device=device, dtype=torch.int32
)
else:
return torch.empty((m, n // block_size), device=device, dtype=torch.uint8)
def create_fp4_output_tensors(
m: int,
n: int,
device: torch.device,
is_sf_swizzled_layout: bool,
padded_n: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Allocate both output tensors for scaled_fp4_quant:
(quantized_output, output_scale).
Must match the C++ scaled_fp4_quant_func allocation exactly when
``padded_n`` is ``None``. When ``padded_n`` is provided, allocate a larger
packed-FP4 output/scale buffer so the quantization kernel can write
CUTLASS-compatible K padding directly
"""
physical_n = padded_n if padded_n is not None else n
output = torch.empty((m, physical_n // 2), device=device, dtype=torch.uint8)
output_scale = create_fp4_scale_tensor(m, physical_n, device, is_sf_swizzled_layout)
return output, output_scale
if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"):
@register_fake("_C::scaled_fp4_quant")
def _scaled_fp4_quant_fake(
input: torch.Tensor,
input_scale: torch.Tensor,
is_sf_swizzled_layout: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
n = input.shape[-1]
m = input.numel() // n
return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout)
# page attention ops
def paged_attention_rocm(
out: torch.Tensor,
exp_sum: torch.Tensor,
max_logits: torch.Tensor,
tmp_out: torch.Tensor,
query: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
num_kv_heads: int,
scale: float,
block_tables: torch.Tensor,
seq_lens: torch.Tensor,
query_start_loc: torch.Tensor | None,
block_size: int,
max_seq_len: int,
alibi_slopes: torch.Tensor | None,
kv_cache_dtype: str,
k_scale: torch.Tensor,
v_scale: torch.Tensor,
fp8_out_scale: torch.Tensor | None = None,
mfma_type: str = "fp8" if envs.VLLM_ROCM_FP8_MFMA_PAGE_ATTN else "f16",
) -> None:
torch.ops._rocm_C.paged_attention(
out,
exp_sum,
max_logits,
tmp_out,
query,
key_cache,
value_cache,
num_kv_heads,
scale,
block_tables,
seq_lens,
query_start_loc,
block_size,
max_seq_len,
alibi_slopes,
kv_cache_dtype,
k_scale,
v_scale,
fp8_out_scale,
mfma_type,
)
def mla_decode_kvcache_cpu(
out: torch.Tensor,
query: torch.Tensor,
kv_cache: torch.Tensor,
scale: float,
block_tables: torch.Tensor,
seq_lens: torch.Tensor,
) -> None:
torch.ops._C.mla_decode_kvcache(out, query, kv_cache, scale, block_tables, seq_lens)
def gather_mla_context_cache_cpu(
*,
src_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
starts: torch.Tensor,
cu_seq_lens: torch.Tensor,
) -> None:
page_size = src_cache.shape[1]
flat_cache = src_cache.view(-1, src_cache.shape[-1])
seq_lens = cu_seq_lens[1:] - cu_seq_lens[:-1]
out_start = 0
for req_idx, (start, seq_len) in enumerate(zip(starts.tolist(), seq_lens.tolist())):
if seq_len <= 0:
continue
positions = torch.arange(start, start + seq_len, device=block_table.device)
block_ids = block_table[req_idx, positions // page_size].to(torch.long)
offsets = positions % page_size
slots = block_ids * page_size + offsets
out_end = out_start + seq_len
dst[out_start:out_end].copy_(flat_cache[slots])
out_start = out_end
# merge attn states ops
def merge_attn_states(
output: torch.Tensor,
prefix_output: torch.Tensor,
prefix_lse: torch.Tensor,
suffix_output: torch.Tensor,
suffix_lse: torch.Tensor,
output_lse: torch.Tensor | None = None,
prefill_tokens_with_context: int | None = None,
output_scale: torch.Tensor | None = None,
) -> None:
torch.ops._C.merge_attn_states(
output,
output_lse,
prefix_output,
prefix_lse,
suffix_output,
suffix_lse,
prefill_tokens_with_context,
output_scale,
)
# pos encoding ops
def rotary_embedding(
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor | None,
head_size: int,
cos_sin_cache: torch.Tensor,
is_neox: bool,
rope_dim_offset: int = 0,
inverse: bool = False,
) -> None:
if rope_dim_offset == 0 and not inverse:
torch.ops._C.rotary_embedding(
positions, query, key, head_size, cos_sin_cache, is_neox
)
else:
torch.ops._C.rotary_embedding(
positions,
query,
key,
head_size,
cos_sin_cache,
is_neox,
rope_dim_offset,
inverse,
)
# layer norm ops
def rms_norm(
out: torch.Tensor,
input: torch.Tensor,
weight: torch.Tensor | None,
epsilon: float,
) -> None:
torch.ops._C.rms_norm(out, input, weight, epsilon)
# Fused vocab-parallel embedding lookup
# (see csrc/.../vocab_parallel_embedding_kernels.cu).
def vocab_parallel_embedding(
input_ids: torch.Tensor,
weight: torch.Tensor,
org_vocab_start_index: int,
org_vocab_end_index: int,
num_org_vocab_padding: int,
added_vocab_start_index: int,
added_vocab_end_index: int,
) -> torch.Tensor:
"""Gather the embedding rows owned by this rank, zeros for the others.
Fuses the range mask, id shift, gather and output masking of the TP > 1
embedding path into a single kernel. Out-of-shard tokens produce zero rows,
so an all-reduce over the ranks yields the full embedding.
Args:
input_ids: ``[num_tokens]`` int32 or int64 token ids.
weight: ``[num_embeddings_per_partition, embedding_dim]`` shard.
org_vocab_start_index: First original vocab id owned by this rank.
org_vocab_end_index: One past the last original vocab id owned by
this rank.
num_org_vocab_padding: Padding rows between this rank's original and
added vocab shards.
added_vocab_start_index: First added vocab id owned by this rank.
added_vocab_end_index: One past the last added vocab id owned by
this rank.
Returns:
``[num_tokens, embedding_dim]`` partial embeddings.
"""
out = torch.empty(
input_ids.shape[0],
weight.shape[1],
dtype=weight.dtype,
device=weight.device,
)
torch.ops._C.vocab_parallel_embedding(
out,
input_ids,
weight,
org_vocab_start_index,
org_vocab_end_index,
num_org_vocab_padding,
added_vocab_start_index,
added_vocab_end_index,
)
return out
# LongCat n-gram embedding index kernel (see csrc/.../ngram_embedding_kernels.cu).
def ngram_compute_n_gram_ids(
ne_n: int,
ne_k: int,
ne_weights: torch.Tensor,
ne_mods: torch.Tensor,
exclusive_ne_embedder_size_sums: torch.Tensor,
exclusive_req_len_sums: torch.Tensor,
ne_token_table: torch.Tensor,
row_indices: torch.Tensor,
column_starts: torch.Tensor,
n_gram_ids: torch.Tensor,
) -> None:
"""Compute concatenated (offset) n-gram ids for a ragged prefill batch.
Writes ``n_gram_ids`` of shape ``[token_num, (ne_n-1)*ne_k]``.
"""
torch.ops._C.ngram_compute_n_gram_ids(
ne_n,
ne_k,
ne_weights,
ne_mods,
exclusive_ne_embedder_size_sums,
exclusive_req_len_sums,
ne_token_table,
row_indices,
column_starts,
n_gram_ids,
)
def fused_add_rms_norm(
input: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor | None,
epsilon: float,
) -> None:
# Note: this func is batch invariant
torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
def fused_qk_norm_rope(
qkv: torch.Tensor,
num_heads_q: int,
num_heads_k: int,
num_heads_v: int,
head_dim: int,
eps: float,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
cos_sin_cache: torch.Tensor,
is_neox: bool,
position_ids: torch.Tensor,
forced_token_heads_per_warp: int = -1,
) -> None:
torch.ops._C.fused_qk_norm_rope(
qkv,
num_heads_q,
num_heads_k,
num_heads_v,
head_dim,
eps,
q_weight,
k_weight,
cos_sin_cache,
is_neox,
position_ids,
forced_token_heads_per_warp,
)
def apply_repetition_penalties_torch(
logits: torch.Tensor,
prompt_mask: torch.Tensor,
output_mask: torch.Tensor,
repetition_penalties: torch.Tensor,
) -> None:
repetition_penalties = repetition_penalties.unsqueeze(dim=1).repeat(
1, logits.size(1)
)
# If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
penalties = torch.where(prompt_mask | output_mask, repetition_penalties, 1.0)
# If logits are positive, divide by penalty, otherwise multiply by penalty.
scaling = torch.where(logits > 0, 1.0 / penalties, penalties)
logits *= scaling
def apply_repetition_penalties_cuda(
logits: torch.Tensor,
prompt_mask: torch.Tensor,
output_mask: torch.Tensor,
repetition_penalties: torch.Tensor,
) -> None:
torch.ops._C.apply_repetition_penalties_(
logits, prompt_mask, output_mask, repetition_penalties
)
def apply_repetition_penalties(
logits: torch.Tensor,
prompt_mask: torch.Tensor,
output_mask: torch.Tensor,
repetition_penalties: torch.Tensor,
) -> None:
"""Apply repetition penalties to logits in-place.
Args:
logits: The logits tensor of shape [num_seqs, vocab_size].
prompt_mask: A boolean tensor indicating which tokens appear in the prompt.
output_mask: A boolean tensor indicating which tokens appear in the output.
repetition_penalties: The repetition penalties of shape (num_seqs, ).
"""
if logits.is_cuda and logits.is_contiguous():
apply_repetition_penalties_cuda(
logits, prompt_mask, output_mask, repetition_penalties
)
else:
apply_repetition_penalties_torch(
logits, prompt_mask, output_mask, repetition_penalties
)
# fused quant layer norm ops
def rms_norm_dynamic_per_token_quant(
input: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
quant_dtype: torch.dtype,
scale_ub: torch.Tensor | None = None,
residual: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
output = torch.empty(input.shape, dtype=quant_dtype, device=input.device)
scales = torch.empty(
(input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
)
torch.ops._C.rms_norm_dynamic_per_token_quant(
output, input, weight, scales, epsilon, scale_ub, residual
)
return output, scales
# fused quant layer norm ops blocked
def rms_norm_per_block_quant(
input: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
quant_dtype: torch.dtype,
group_size: list[int],
scale_ub: torch.Tensor | None = None,
residual: torch.Tensor | None = None,
is_scale_transposed: bool = False,
tma_alignment: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
assert len(group_size) == 2
output = torch.empty(input.shape, dtype=quant_dtype, device=input.device)
if is_scale_transposed:
if tma_alignment == 0:
scales = torch.empty(
(input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
device=input.device,
dtype=torch.float32,
).transpose(0, 1)
else:
m = input.shape[-2]
sf_k = input.shape[-1] // group_size[1]
tma_aligned_m = (m + tma_alignment - 1) // tma_alignment * tma_alignment
shape = input.shape[:-2] + (m, sf_k)
stride = (
(1, tma_aligned_m)
if input.dim() == 2
else (tma_aligned_m * sf_k, 1, tma_aligned_m)
)
scales = torch.empty_strided(
shape, stride, device=input.device, dtype=torch.float32
)
else:
scales = torch.empty(
(input.numel() // input.shape[-1], input.shape[-1] // group_size[1]),
device=input.device,
dtype=torch.float32,
)
assert tma_alignment in [0, 4], "Expected TMA alignment 0 or 4, but got " + str(
tma_alignment
)
torch.ops._C.rms_norm_per_block_quant(
output,
input,
weight,
scales,
epsilon,
scale_ub,
residual,
group_size[1],
is_scale_transposed,
)
return output, scales
# fused silu_and_mul + block quant
def silu_and_mul_per_block_quant(
input: torch.Tensor,
group_size: int, # Changed from list[int]
quant_dtype: torch.dtype,
scale_ub: torch.Tensor | None = None,
is_scale_transposed: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
assert input.ndim == 2, f"input must be 2D [batch, hidden*2], got {input.shape}"
assert input.shape[-1] % 2 == 0, (
f"input last dim must be even (gate||up layout), got {input.shape[-1]}"
)
# Output is half the width of input (after silu_and_mul)
num_tokens = input.shape[0]
hidden_size = input.shape[-1] // 2 # Divide by 2 because input is [gate || up]
# Allocate output tensor (FP8 or INT8)
output = torch.empty(
(num_tokens, hidden_size), device=input.device, dtype=quant_dtype
)
# Allocate scales tensor
num_groups = hidden_size // group_size # Directly use group_size
if is_scale_transposed:
scales = torch.empty(
(num_groups, num_tokens),
device=input.device,
dtype=torch.float32,
).t()
else:
scales = torch.empty(
(num_tokens, num_groups),
device=input.device,
dtype=torch.float32,
)
# Call the C++ kernel
torch.ops._C.silu_and_mul_per_block_quant(
output,
input,
scales,
group_size, # Pass directly as int
scale_ub,
is_scale_transposed,
)
return output, scales
# quantization ops
# awq
def awq_dequantize(
qweight: torch.Tensor,
scales: torch.Tensor,
zeros: torch.Tensor,
split_k_iters: int,
thx: int,
thy: int,
) -> torch.Tensor:
if envs.VLLM_USE_TRITON_AWQ:
from vllm.model_executor.layers.quantization.awq_triton import (
awq_dequantize_triton,
)
return awq_dequantize_triton(qweight, scales, zeros)
return torch.ops._C.awq_dequantize(qweight, scales, zeros, split_k_iters, thx, thy)
if hasattr(torch.ops._C, "awq_dequantize"):
@register_fake("_C::awq_dequantize")
def _awq_dequantize_fake(
qweight: torch.Tensor,
scales: torch.Tensor,
zeros: torch.Tensor,
split_k_iters: torch.SymInt,
thx: int,
thy: int,
) -> torch.Tensor:
in_c = qweight.size(0)
qout_c = qweight.size(1)
out_c = qout_c * 8
return torch.empty((in_c, out_c), dtype=scales.dtype, device=scales.device)
def awq_gemm(
input: torch.Tensor,
qweight: torch.Tensor,
scales: torch.Tensor,
qzeros: torch.Tensor,
split_k_iters: int,
) -> torch.Tensor:
if envs.VLLM_USE_TRITON_AWQ:
from vllm.model_executor.layers.quantization.awq_triton import awq_gemm_triton
return awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters)
return torch.ops._C.awq_gemm(input, qweight, scales, qzeros, split_k_iters)
if hasattr(torch.ops._C, "awq_gemm"):
@register_fake("_C::awq_gemm")
def _awq_gemm_fake(
input: torch.Tensor,
qweight: torch.Tensor,
scales: torch.Tensor,
qzeros: torch.Tensor,
split_k_iters: torch.SymInt,
) -> torch.Tensor:
num_in_feats = input.size(0)
return torch.empty(
(split_k_iters, num_in_feats, qweight.size(1) * 8),
dtype=input.dtype,
device=input.device,
).sum(0)
# gptq
def gptq_gemm(
a: torch.Tensor,
b_q_weight: torch.Tensor,
b_gptq_qzeros: torch.Tensor,
b_gptq_scales: torch.Tensor,
use_exllama: bool,
use_v2_format: bool,
bit: int,
) -> torch.Tensor:
return torch.ops._C.gptq_gemm(
a,
b_q_weight,
b_gptq_qzeros,
b_gptq_scales,
use_exllama,
use_v2_format,
bit,
)
if hasattr(torch.ops._C, "gptq_gemm"):
@register_fake("_C::gptq_gemm")
def _gptq_gemm_fake(
a: torch.Tensor,
b_q_weight: torch.Tensor,
b_gptq_qzeros: torch.Tensor,
b_gptq_scales: torch.Tensor,
use_exllama: bool,
use_v2_format: bool,
bit: int,
) -> torch.Tensor:
return torch.empty(
(a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
)
def gptq_shuffle(q_weight: torch.Tensor, bit: int) -> None:
torch.ops._C.gptq_shuffle(q_weight, bit)
def gptq_gemm_rdna3(
a: torch.Tensor,
b_q_weight: torch.Tensor,
b_qzeros: torch.Tensor,
b_scales: torch.Tensor,
use_v2_format: bool,
) -> torch.Tensor:
return torch.ops._rocm_C.gptq_gemm_rdna3(
a, b_q_weight, b_qzeros, b_scales, use_v2_format
)
if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3"):
@register_fake("_rocm_C::gptq_gemm_rdna3")
def _gptq_gemm_rdna3_fake(
a: torch.Tensor,
b_q_weight: torch.Tensor,
b_qzeros: torch.Tensor,
b_scales: torch.Tensor,
use_v2_format: bool,
) -> torch.Tensor:
return torch.empty(
(a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
)
if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3_wmma"):
@register_fake("_rocm_C::gptq_gemm_rdna3_wmma")
def _gptq_gemm_rdna3_wmma_fake(
a: torch.Tensor,
b_q_weight: torch.Tensor,
b_qzeros: torch.Tensor,
b_scales: torch.Tensor,
use_v2_format: bool,
) -> torch.Tensor:
return torch.empty(
(a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device
)
def moe_gptq_gemm_rdna3(
a: torch.Tensor,
c: torch.Tensor,
b_q_weight: torch.Tensor,
b_scales: torch.Tensor,
b_qzeros: torch.Tensor,
topk_weights: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor,
top_k: int,
block_size_m: int,
mul_topk_weight: bool,
output_topk: int = 0,
) -> None:
torch.ops._rocm_C.moe_gptq_gemm_rdna3(
a,
c,
b_q_weight,
b_scales,
b_qzeros,
topk_weights,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
top_k,
block_size_m,
mul_topk_weight,
output_topk,
)
if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
@register_fake("_C::allspark_w8a16_gemm")
def _allspark_w8a16_gemm_fake(
a: torch.Tensor,
b_qweight: torch.Tensor,
b_scales: torch.Tensor,
b_qzeros: torch.Tensor | None,
n: torch.SymInt,
group_size: torch.SymInt,
sm_count: torch.SymInt,
sm_version: torch.SymInt,
CUBLAS_M_THRESHOLD: torch.SymInt,
has_zp: bool,
n32k16_reorder: bool,
) -> torch.Tensor:
m = a.size(0)
return torch.empty((m, n), device=a.device, dtype=a.dtype)
# cutlass
def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool:
return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability)
def mxfp4_experts_quant_supported(cuda_device_capability: int) -> bool:
try:
return torch.ops._C.mxfp4_experts_quant_supported(cuda_device_capability)
except AttributeError:
# Return False on builds where the CUDA helper is not available.
return False
def cutlass_scaled_fp4_mm(
a: torch.Tensor,
b: torch.Tensor,
block_scale_a: torch.Tensor,
block_scale_b: torch.Tensor,
alpha: torch.Tensor,
out_dtype: torch.dtype,
) -> torch.Tensor:
assert a.ndim == 2 and b.ndim == 2
m, n = a.shape[0], b.shape[0]
out = torch.empty((m, n), dtype=out_dtype, device=a.device)
torch.ops._C.cutlass_scaled_fp4_mm(out, a, b, block_scale_a, block_scale_b, alpha)
return out
def cutlass_scaled_mm_supports_fp8(cuda_device_capability: int) -> bool:
return torch.ops._C.cutlass_scaled_mm_supports_fp8(cuda_device_capability)
def cutlass_scaled_mm_supports_block_fp8(cuda_device_capability: int) -> bool:
return torch.ops._C.cutlass_scaled_mm_supports_block_fp8(cuda_device_capability)
def cutlass_scaled_mm(
a: torch.Tensor,
b: torch.Tensor,
scale_a: torch.Tensor,
scale_b: torch.Tensor,
out_dtype: torch.dtype,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
"""`cutlass_scaled_mm` implements a fused version of
`output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
where scale_a * a and scale_b * b are implemented using numpy-style
broadcasting.
In order to support blockwise scaling like found in DeepSeek V3 we also
support extended "group" broadcast rules. We extend the numpy-style
broadcasting rules with the following rule:
"if the extent of a dimension in the source shape is between 1 and
corresponding extent in the target shape we repeat each element along
that dimension src_shape[dim] // target_shape[dim] times consecutively"
example if we have:
a = [[1, 2], and target_shape = (2, 4)
[3, 4]]
then we would expand a to:
a = [[1, 1, 2, 2],
[3, 3, 4, 4]]
currently we only support the case:
scale_a.shape * [1, 128] == a.shape
scale_b.shape * [128, 128] == b.shape
"""
assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
assert bias is None or bias.numel() == b.shape[1] and bias.dtype == out_dtype
# Massage the input to be 2D
target_shape = (*a.shape[:-1], b.shape[1])
a = a.view(-1, a.shape[-1])
cutlass_compatible_b = b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0
if current_platform.is_rocm() or not cutlass_compatible_b:
from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import ( # noqa
triton_scaled_mm,
)
out = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
else:
out = torch.empty((a.shape[0], b.shape[1]), dtype=out_dtype, device=a.device)
torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias)
return out.view(*target_shape)
def cutlass_scaled_mm_azp(
a: torch.Tensor,
b: torch.Tensor,
scale_a: torch.Tensor,
scale_b: torch.Tensor,
out_dtype: torch.dtype,
azp_adj: torch.Tensor,
azp: torch.Tensor | None = None,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
"""Args:
azp_adj: In the per-tensor case, this should include the azp.
Always per-channel.
azp: Only set in the per-token case. Per-token if set.
"""
assert b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0
assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
assert bias is None or bias.numel() == b.shape[1] and bias.dtype == out_dtype
# Massage the input to be 2D
target_shape = (*a.shape[:-1], b.shape[1])
a = a.view(-1, a.shape[-1])
assert azp is None or azp.numel() == a.shape[0]
out = torch.empty((a.shape[0], b.shape[1]), dtype=out_dtype, device=a.device)
torch.ops._C.cutlass_scaled_mm_azp(out, a, b, scale_a, scale_b, azp_adj, azp, bias)
return out.view(*target_shape)
def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
if cuda_device_capability < 90 or cuda_device_capability >= 110:
return False
try:
return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability)
except AttributeError:
# Return False on non-CUDA platforms where it is not available
return False
def get_cutlass_moe_mm_data(
topk_ids: torch.Tensor,
expert_offsets: torch.Tensor,
problem_sizes1: torch.Tensor,
problem_sizes2: torch.Tensor,
input_permutation: torch.Tensor,
output_permutation: torch.Tensor,
num_experts: int,
n: int,
k: int,
blockscale_offsets: torch.Tensor | None = None,
is_gated: bool = True,
):
"""Prepare data necessary to perform CUTLASS grouped matrix multiplications
used in CUTLASS-based fused MoE.
The function takes in topk_ids (token-expert mapping) and uses it to
compute:
- expert_offsets: Indices that mark at which token index each expert begins
its computation after the input is sorted with
input_permutation. The number of tokens computed with
expert E is expert_offsets[E + 1] - expert_offsets[E]
- problem_sizes1, problem_sizes2: MxNxK sizes of each expert's
multiplication in two grouped MMs used in
the fused MoE operation.
- input_permutation: Permutation that must be used to shuffle the input
before executing the MMs.
- output_permutation: Permutation that must be used to shuffle the output
after executing the MMs.
- blockscale_offsets: Optional argument passed for fp4 moe. Indices that
mark at which block scale index each expert begins
its computation. The number of block scale rows
computed with expert E is blockscale_offsets[E + 1] -
blockscale_offsets[E]
- is_gated: Whether the activation is gated (gate + up). When True, the
first GEMM N dimension is 2*n; when False, it is n.
"""
return torch.ops._C.get_cutlass_moe_mm_data(
topk_ids,
expert_offsets,
problem_sizes1,
problem_sizes2,
input_permutation,
output_permutation,
num_experts,
n,
k,
blockscale_offsets,
is_gated,
)
def get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
expert_first_token_offset: torch.Tensor,
problem_sizes1: torch.Tensor,
problem_sizes2: torch.Tensor,
n: int,
k: int,
swap_ab: bool,
):
"""Compute per-expert (M, N, K) problem sizes from expert_first_token_offset"""
return torch.ops._C.get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
expert_first_token_offset,
problem_sizes1,
problem_sizes2,
n,
k,
swap_ab,
)
def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor):
"""Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor.
This is used in MoE to permute the input tensor before performing grouped matrix multiplications.
"""
num_tokens_permuted = dst2src_map.shape[0]
output_tensor = torch.empty(
(num_tokens_permuted, input_tensor.shape[1]),
device=input_tensor.device,
dtype=input_tensor.dtype,
)
torch.ops._moe_C.shuffle_rows(input_tensor, dst2src_map, output_tensor)
return output_tensor
def get_cutlass_batched_moe_mm_data(
expert_offsets: torch.Tensor,
problem_sizes1: torch.Tensor,
problem_sizes2: torch.Tensor,
expert_num_tokens: torch.Tensor,
num_local_experts: int,
padded_m: int,
n: int,
k: int,
):
"""Prepare data necessary to perform CUTLASS grouped matrix multiplications
used in CUTLASS-based fused MoE.
The function takes in expert_num_tokens (token count per expert) and
non_zero_expert_idxs (consecutive indices of experts with non-zero token
counts) and uses them to compute:
- expert_offsets: Indices that mark at which token index each expert begins
its computation.
- problem_sizes1, problem_sizes2: MxNxK sizes of each expert's
multiplication in two grouped MMs used in
the fused MoE operation.
"""
return torch.ops._C.get_cutlass_batched_moe_mm_data(
expert_offsets,
problem_sizes1,
problem_sizes2,
expert_num_tokens,
num_local_experts,
padded_m,
n,
k,
)
def cutlass_moe_mm(
out_tensors: torch.Tensor,
a_tensors: torch.Tensor,
b_tensors: torch.Tensor,
a_scales: torch.Tensor,
b_scales: torch.Tensor,
expert_offsets: torch.Tensor,
problem_sizes: torch.Tensor,
a_strides: torch.Tensor,
b_strides: torch.Tensor,
c_strides: torch.Tensor,
per_act_token: bool,
per_out_ch: bool,
):
"""A single grouped matrix multiplication used in CUTLASS-based fused MoE.
The function executes fp8-quantized OUT = AB matrix multiplication.
- expert_offsets: Indices that mark at which token index each expert begins
its computation. The number of tokens computed with
expert E is expert_offsets[E + 1] - expert_offsets[E]
- problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
MMs used in the fused MoE operation.
- a/b/c_strides: The data strides passed to grouped matrix multiplication.
"""
return torch.ops._C.cutlass_moe_mm(
out_tensors,
a_tensors,
b_tensors,
a_scales,
b_scales,
expert_offsets,
problem_sizes,
a_strides,
b_strides,
c_strides,
per_act_token,
per_out_ch,
)
def cutlass_fp4_moe_mm(
out_tensors: torch.Tensor,
a_tensors: torch.Tensor,
b_tensors: torch.Tensor,
a_scales: torch.Tensor,
b_scales: torch.Tensor,
alphas: torch.Tensor,
problem_sizes: torch.Tensor,
expert_offsets: torch.Tensor,
sf_offsets: torch.Tensor,
):
"""An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs
the gemms for each combination based on the specified problem sizes.
This is used as the MoE gemm during NVFP4 Quantized MoERunner forward.
- a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized
input and expert weights.
- a_/b_scales: The blockscales in FP8-E4M3 precision
- expert_offsets/sf_offsets: Indices that mark at which token index
each expert begins its computation. The number of tokens
computed with expert E is expert_offsets[E + 1] -
expert_offsets[E] And the sf_size per expert is
sf_offset[E+1] - sf_offset[E]
- problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
MMs used in the fused MoE operation.
"""
return torch.ops._C.cutlass_fp4_group_mm(
out_tensors,
a_tensors,
b_tensors,
a_scales,
b_scales,
alphas,
problem_sizes,
expert_offsets,
sf_offsets,
)
def cutlass_mxfp4_moe_mm(
out_tensors: torch.Tensor,
a_tensors: torch.Tensor,
b_tensors: torch.Tensor,
a_scales: torch.Tensor,
b_scales: torch.Tensor,
problem_sizes: torch.Tensor,
expert_offsets: torch.Tensor,
sf_offsets: torch.Tensor,
):
"""An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
- a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
- a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout)
- Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales).
- expert_offsets/sf_offsets: expert boundary indices
- problem_sizes: (num_experts, 3) with (M, N, K) per expert
"""
return torch.ops._C.cutlass_mxfp4_group_mm(
out_tensors,
a_tensors,
b_tensors,
a_scales,
b_scales,
problem_sizes,
expert_offsets,
sf_offsets,
)
# gptq_marlin
def gptq_marlin_repack(
b_q_weight: torch.Tensor,
size_k: int,
size_n: int,
num_bits: int,
is_a_8bit: bool = False,
) -> torch.Tensor:
return torch.ops._C.gptq_marlin_repack(
b_q_weight, size_k, size_n, num_bits, is_a_8bit
)
if hasattr(torch.ops._C, "gptq_marlin_repack"):
@register_fake("_C::gptq_marlin_repack")
def _gptq_marlin_repack_fake(
b_q_weight: torch.Tensor,
size_k: torch.SymInt,
size_n: torch.SymInt,
num_bits: int,
is_a_8bit: bool = False,
) -> torch.Tensor:
pack_factor = 32 // num_bits
marlin_tile_size = 16
return torch.empty(
(size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor),
dtype=b_q_weight.dtype,
device=b_q_weight.device,
)
# awq_marlin
def awq_marlin_repack(
b_q_weight: torch.Tensor,
size_k: int,
size_n: int,
num_bits: int,
is_a_8bit: bool = False,
) -> torch.Tensor:
return torch.ops._C.awq_marlin_repack(
b_q_weight, size_k, size_n, num_bits, is_a_8bit
)
if hasattr(torch.ops._C, "awq_marlin_repack"):
@register_fake("_C::awq_marlin_repack")
def _awq_marlin_repack_fake(
b_q_weight: torch.Tensor,
size_k: torch.SymInt,
size_n: torch.SymInt,
num_bits: int,
is_a_8bit: bool = False,
) -> torch.Tensor:
pack_factor = 32 // num_bits
marlin_tile_size = 16
return torch.empty(
(size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor),
dtype=b_q_weight.dtype,
device=b_q_weight.device,
)
def gptq_marlin_moe_repack(
b_q_weight: torch.Tensor,
size_k: int,
size_n: int,
num_bits: int,
is_a_8bit: bool = False,
) -> torch.Tensor:
num_experts = b_q_weight.shape[0]
assert size_k % 16 == 0
output = torch.empty(
(num_experts, size_k // 16, size_n * (num_bits // 2)),
device=b_q_weight.device,
dtype=b_q_weight.dtype,
)
for e in range(num_experts):
output[e] = torch.ops._C.gptq_marlin_repack(
b_q_weight[e], size_k, size_n, num_bits, is_a_8bit
)
return output
def awq_marlin_moe_repack(
b_q_weight: torch.Tensor,
size_k: int,
size_n: int,
num_bits: int,
is_a_8bit: bool = False,
) -> torch.Tensor:
num_experts = b_q_weight.shape[0]
assert size_k % 16 == 0
output = torch.empty(
(num_experts, size_k // 16, size_n * (num_bits // 2)),
device=b_q_weight.device,
dtype=b_q_weight.dtype,
)
for e in range(num_experts):
output[e] = torch.ops._C.awq_marlin_repack(
b_q_weight[e], size_k, size_n, num_bits, is_a_8bit
)
return output
def marlin_int4_fp8_preprocess(
qweight: torch.Tensor,
qzeros_or_none: torch.Tensor | None = None,
inplace: bool = False,
):
return torch.ops._C.marlin_int4_fp8_preprocess(qweight, qzeros_or_none, inplace)
def marlin_gemm(
a: torch.Tensor,
c: torch.Tensor | None,
b_q_weight: torch.Tensor,
b_bias: torch.Tensor | None,
b_scales: torch.Tensor,
a_scales: torch.Tensor | None,
global_scale: torch.Tensor | None,
b_zeros: torch.Tensor | None,
workspace: torch.Tensor,
b_q_type: ScalarType,
size_m: int,
size_n: int,
size_k: int,
use_atomic_add: bool = False,
use_fp32_reduce: bool = False,
is_zp_float: bool = False,
) -> torch.Tensor:
return torch.ops._C.marlin_gemm(
a,
c,
b_q_weight,
b_bias,
b_scales,
a_scales,
global_scale,
b_zeros,
workspace,
b_q_type.id,
size_m,
size_n,
size_k,
use_atomic_add,
use_fp32_reduce,
is_zp_float,
)
if hasattr(torch.ops._C, "marlin_gemm"):
@register_fake("_C::marlin_gemm")
def _marlin_gemm_fake(
a: torch.Tensor,
c: torch.Tensor | None,
b_q_weight: torch.Tensor,
b_bias: torch.Tensor | None,
b_scales: torch.Tensor,
a_scales: torch.Tensor | None,
global_scale: torch.Tensor | None,
b_zeros: torch.Tensor | None,
workspace: torch.Tensor,
b_q_type_id: int,
size_m: torch.SymInt,
size_n: torch.SymInt,
size_k: torch.SymInt,
use_atomic_add: bool = False,
use_fp32_reduce: bool = False,
is_zp_float: bool = False,
) -> torch.Tensor:
dtype = a.dtype
if dtype not in [torch.half, torch.bfloat16]:
dtype = b_scales.dtype
return torch.empty((size_m, size_n), device=a.device, dtype=dtype)
# machete
def machete_supported_schedules(
a_type: torch.dtype,
b_type: ScalarType,
group_scales_type: torch.dtype | None,
group_zeros_type: torch.dtype | None = None,
channel_scales_type: torch.dtype | None = None,
token_scales_type: torch.dtype | None = None,
out_type: torch.dtype | None = None,
) -> list[str]:
return torch.ops._C.machete_supported_schedules(
a_type,
b_type.id,
group_scales_type,
group_zeros_type,
channel_scales_type,
token_scales_type,
out_type,
)
def machete_mm(
a: torch.Tensor,
# b_q Should be the tensor returned by machete_prepack_B
b_q: torch.Tensor,
b_type: ScalarType,
out_type: torch.dtype | None = None,
b_group_scales: torch.Tensor | None = None,
b_group_zeros: torch.Tensor | None = None,
b_group_size: int | None = None,
b_channel_scales: torch.Tensor | None = None,
a_token_scales: torch.Tensor | None = None,
schedule: str | None = None,
) -> torch.Tensor:
return torch.ops._C.machete_mm(
a,
b_q,
b_type.id,
out_type,
b_group_scales,
b_group_zeros,
b_group_size,
b_channel_scales,
a_token_scales,
schedule,
)
if hasattr(torch.ops._C, "machete_mm"):
@register_fake("_C::machete_mm")
def machete_mm_fake(
a: torch.Tensor,
# b_q Should be the tensor returned by machete_prepack_B
b_q: torch.Tensor,
b_type: ScalarType,
out_type: torch.dtype | None = None,
b_group_scales: torch.Tensor | None = None,
b_group_zeros: torch.Tensor | None = None,
b_group_size: int | None = None,
b_channel_scales: torch.Tensor | None = None,
a_token_scales: torch.Tensor | None = None,
schedule: str | None = None,
) -> torch.Tensor:
m = a.size(0)
n = b_q.size(1)
return torch.empty((m, n), device=a.device, dtype=a.dtype)
def machete_prepack_B(
b_q_weight: torch.Tensor,
a_type: torch.dtype,
b_type: ScalarType,
group_scales_type: torch.dtype | None,
) -> torch.Tensor:
return torch.ops._C.machete_prepack_B(
b_q_weight, a_type, b_type.id, group_scales_type
)
if hasattr(torch.ops._C, "machete_prepack_B"):
@register_fake("_C::machete_prepack_B")
def machete_prepack_B_fake(
b_q_weight: torch.Tensor,
a_type: torch.dtype,
b_type: ScalarType,
group_scales_type: torch.dtype | None,
) -> torch.Tensor:
return torch.empty_like(b_q_weight, memory_format=torch.contiguous_format)
# CUTLASS W4A8
def cutlass_w4a8_mm(
a: torch.Tensor,
# b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b
b_q: torch.Tensor,
b_group_scales: torch.Tensor,
b_group_size: int,
b_channel_scales: torch.Tensor,
a_token_scales: torch.Tensor,
out_type: torch.dtype | None = None,
maybe_schedule: str | None = None,
) -> torch.Tensor:
return torch.ops._C.cutlass_w4a8_mm(
a,
b_q,
b_group_scales,
b_group_size,
b_channel_scales,
a_token_scales,
out_type,
maybe_schedule,
)
if hasattr(torch.ops._C, "cutlass_w4a8_mm"):
@register_fake("_C::cutlass_w4a8_mm")
def cutlass_w4a8_mm_fake(
a: torch.Tensor,
# b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b
b_q: torch.Tensor,
b_group_scales: torch.Tensor,
b_group_size: int,
b_channel_scales: torch.Tensor,
a_token_scales: torch.Tensor,
out_type: torch.dtype | None = None,
maybe_schedule: str | None = None,
) -> torch.Tensor:
m = a.size(0)
n = b_q.size(1)
out_dtype = out_type if out_type is not None else torch.bfloat16
return torch.empty((m, n), device=a.device, dtype=out_dtype)
def cutlass_pack_scale_fp8(scales: torch.Tensor) -> torch.Tensor:
return torch.ops._C.cutlass_pack_scale_fp8(scales)
if hasattr(torch.ops._C, "cutlass_pack_scale_fp8"):
@register_fake("_C::cutlass_pack_scale_fp8")
def cutlass_pack_scale_fp8_fake(scales: torch.Tensor) -> torch.Tensor:
return torch.empty_like(scales, memory_format=torch.contiguous_format)
def cutlass_encode_and_reorder_int4b(b: torch.Tensor) -> torch.Tensor:
return torch.ops._C.cutlass_encode_and_reorder_int4b(b)
if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b"):
@register_fake("_C::cutlass_encode_and_reorder_int4b")
def cutlass_encode_and_reorder_int4b_fake(b: torch.Tensor) -> torch.Tensor:
return torch.empty_like(b, memory_format=torch.contiguous_format)
def cutlass_w4a8_moe_mm(
out_tensors: torch.Tensor,
a_tensors: torch.Tensor,
b_tensors: torch.Tensor,
a_scales: torch.Tensor,
b_scales: torch.Tensor,
b_group_scales: torch.Tensor,
b_group_size: int,
expert_offsets: torch.Tensor,
problem_sizes: torch.Tensor,
a_strides: torch.Tensor,
b_strides: torch.Tensor,
c_strides: torch.Tensor,
group_scale_strides: torch.Tensor,
maybe_schedule: str | None = None,
) -> None:
"""Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the
W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8)
and both per-channel + per-token scaling in the epilogue.
Args:
out_tensors:
Output buffer for all experts (updated in-place).
a_tensors:
FP8 (E4M3FN) activations for all experts.
b_tensors:
INT4-packed weight matrix for all experts, packed to INT32
a_scales:
Per-token FP8 activation scales, applied in the epilogue.
b_scales:
Per-channel FP8 weight scales for each expert, applied in the epilogue.
b_group_scales:
FP8 scale values for group-wise INT4 weight blocks.
b_group_size:
Number of elements grouped under each entry of b_group_scales.
expert_offsets:
Cumulative token offsets
problem_sizes:
Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher.
a_strides:
Strides describing the memory layout of a_tensors.
b_strides:
Strides describing the memory layout of b_tensors.
c_strides:
Strides describing the memory layout of out_tensors.
group_scale_strides:
Strides describing the memory layout of b_group_scales.
maybe_schedule:
Optional override to choose a specific kernel or epilogue schedule.
Returns:
out_tensors updated in-place with the dequantized INT4xFP8 grouped GEMM result.
"""
return torch.ops._C.cutlass_w4a8_moe_mm(
out_tensors,
a_tensors,
b_tensors,
a_scales,
b_scales,
b_group_scales,
b_group_size,
expert_offsets,
problem_sizes,
a_strides,
b_strides,
c_strides,
group_scale_strides,
maybe_schedule,
)
def cutlass_encode_and_reorder_int4b_grouped(
b_tensors: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.ops._C.cutlass_encode_and_reorder_int4b_grouped(b_tensors)
if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b_grouped"):
@register_fake("_C::cutlass_encode_and_reorder_int4b_grouped")
def cutlass_encode_and_reorder_int4b_grouped_fake(b: torch.Tensor) -> torch.Tensor:
return torch.empty_like(b, memory_format=torch.contiguous_format)
# fp4
def scaled_fp4_quant(
input: torch.Tensor,
input_global_scale: torch.Tensor,
is_sf_swizzled_layout: bool = True,
backend: str = "none",
padded_n: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize input tensor to FP4 and return quantized tensor and scale.
This function quantizes the last dimension of the given tensor `input`. For
every 16 consecutive elements, a single dynamically computed scaling factor
is shared. This scaling factor is quantized using the `input_global_scale`
and is stored in a swizzled layout (see
https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x).
Args:
input: The input tensor to be quantized to FP4
input_global_scale: A scalar scaling factor for the entire tensor.
is_sf_swizzled_layout: Whether to store the scaling factors in the
swizzled layout (default `True`).
backend: Quantization kernel backend to dispatch to. For `"trtllm"`
backends the 8x4 scale-factor layout is selected for small
batches (m <= 32) instead of the 128x4 layout.
padded_n: Optional padded K dimension. When provided, the quantized
output and scale tensors are allocated for ``padded_n``
Returns:
tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every
two values are packed into a uint8 and float8_e4m3 scaling factors
in the sizzled layout.
"""
assert not current_platform.is_rocm()
assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}."
other_dims = 1 if input.ndim == 1 else -1
input = input.reshape(other_dims, input.shape[-1])
m, n = input.shape
block_size = 16
assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}."
assert input.dtype in (torch.float16, torch.bfloat16), (
f"input.dtype needs to be fp16 or bf16 but got {input.dtype}."
)
if padded_n is not None:
assert padded_n >= n, f"padded_n must be >= n, got padded_n={padded_n}, n={n}."
assert padded_n % block_size == 0, (
f"padded_n has to be a multiple of {block_size}, but got {padded_n}."
)
use_8x4_sf_layout = True if "trtllm" in backend and m <= 32 else False # noqa: SIM210
if use_8x4_sf_layout and padded_n is not None and padded_n != n:
# TODO: support this case
raise ValueError("padded_n is not supported with TRTLLM 8x4 scale layout.")
if use_8x4_sf_layout:
output, output_scale = flashinfer_quant_nvfp4_8x4_sf_layout(
input, input_global_scale
)
else:
# Pre-allocate and call .out variant (same behavior as old in-place API)
output, output_scale = create_fp4_output_tensors(
m,
n,
input.device,
is_sf_swizzled_layout,
padded_n=padded_n,
)
torch.ops._C.scaled_fp4_quant.out(
input,
input_global_scale,
is_sf_swizzled_layout,
output=output,
output_scale=output_scale,
)
output_scale = output_scale.view(torch.float8_e4m3fn)
return output, output_scale
def scaled_fp4_experts_quant(
input_tensor: torch.Tensor,
input_global_scale: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
topk: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize input tensor to NVFP4 and return quantized tensor and scale, for
packed MoE Inputs.
Args:
input_tensor: The input tensor to be quantized to NVFP4
input_global_scale: A scalar scaling factor for the entire tensor.
expert_offsets: The expert offsets tensor
blockscale_offsets: The blockscale offsets tensor
topk: The number of experts each token is routed to
Outputs:
output: The quantized tensor in NVFP4
output_scales: The blockscale tensor in FP8-E4M3
"""
assert not current_platform.is_rocm()
assert input_tensor.ndim == 2, (
f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
)
# Control the maximum number of tokens per expert supported by the
# NVFP4 MoE Expert Quantization. This is used to prevent the kernel
# from running out of memory. This value can also be increased to support
# larger models.
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
m_numtopk, k = input_tensor.shape
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
f"{MAX_TOKENS_PER_EXPERT})"
f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
)
scales_k = k // 16
padded_k = (scales_k + (4 - 1)) // 4
# output is uint8 and packed fp4 values
output = torch.empty(
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
)
output_scales = torch.empty(
MAX_TOKENS_PER_EXPERT * topk,
padded_k,
dtype=torch.int32,
device=input_tensor.device,
)
torch.ops._C.scaled_fp4_experts_quant(
output,
output_scales,
input_tensor,
input_global_scale,
expert_offsets,
blockscale_offsets,
)
output_scales = output_scales.view(torch.float8_e4m3fn)
return output, output_scales
def silu_and_mul_scaled_fp4_experts_quant(
input_tensor: torch.Tensor,
input_global_scale: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
topk: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations.
Args:
input_tensor: The input tensor with gate || up layout [m_topk, k*2]
input_global_scale: A per-expert scaling factor [n_experts]
expert_offsets: The expert offsets tensor [n_experts+1]
blockscale_offsets: The blockscale offsets tensor [n_experts+1]
topk: Number of top-k experts selected
Outputs:
output: The quantized tensor in NVFP4 [m_topk, k/2]
output_scales: The blockscale tensor in FP8-E4M3
"""
assert not current_platform.is_rocm()
assert input_tensor.ndim == 2, (
f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
)
# Control the maximum number of tokens per expert supported by the
# NVFP4 MoE Expert Quantization. This is used to prevent the kernel
# from running out of memory. This value can also be increased to support
# larger models.
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
m_numtopk, k_times_2 = input_tensor.shape
assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
k = k_times_2 // 2
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
f"{MAX_TOKENS_PER_EXPERT})"
f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
)
scales_k = k // 16
padded_k = (scales_k + (4 - 1)) // 4
# output is uint8 and packed fp4 values
output = torch.empty(
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
)
output_scales = torch.empty(
MAX_TOKENS_PER_EXPERT * topk,
padded_k,
dtype=torch.int32,
device=input_tensor.device,
)
torch.ops._C.silu_and_mul_scaled_fp4_experts_quant(
output,
output_scales,
input_tensor,
input_global_scale,
expert_offsets,
blockscale_offsets,
)
output_scales = output_scales.view(torch.float8_e4m3fn)
return output, output_scales
def mxfp4_experts_quant(
input_tensor: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
n_experts: int,
topk: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize input tensor to MXFP4 for packed MoE inputs.
Uses 32-element blocks with E8M0 (power-of-two) scale factors.
MXFP4 has no global scale - only block-level E8M0 scale factors.
Args:
input_tensor: [m_topk, k] BF16/FP16 activations
expert_offsets: [n_experts+1] token boundaries per expert
blockscale_offsets: [n_experts+1] SF row boundaries per expert
n_experts: number of experts
topk: number of top-k experts
Returns:
output: [m_topk, k//2] packed E2M1 values (uint8)
output_scales: E8M0 blockscales in swizzled layout (uint8 view)
"""
assert not current_platform.is_rocm()
assert input_tensor.ndim == 2
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
m_numtopk, k = input_tensor.shape
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
f"{MAX_TOKENS_PER_EXPERT})"
f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use"
f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
)
scales_k = k // 32
padded_k = (scales_k + (4 - 1)) // 4
output = torch.empty(
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
)
output_scales = torch.empty(
MAX_TOKENS_PER_EXPERT * topk,
padded_k,
dtype=torch.int32,
device=input_tensor.device,
)
torch.ops._C.mxfp4_experts_quant(
output,
output_scales,
input_tensor,
expert_offsets,
blockscale_offsets,
n_experts,
)
# E8M0 SFs are stored as uint8
output_scales = output_scales.view(torch.uint8)
return output, output_scales
def silu_and_mul_mxfp4_experts_quant(
input_tensor: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
n_experts: int,
topk: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
MXFP4 has no global scale - only block-level E8M0 scale factors.
"""
assert not current_platform.is_rocm()
assert input_tensor.ndim == 2
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
m_numtopk, k_times_2 = input_tensor.shape
assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
k = k_times_2 // 2
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk
scales_k = k // 32
padded_k = (scales_k + (4 - 1)) // 4
output = torch.empty(
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
)
output_scales = torch.empty(
MAX_TOKENS_PER_EXPERT * topk,
padded_k,
dtype=torch.int32,
device=input_tensor.device,
)
torch.ops._C.silu_and_mul_mxfp4_experts_quant(
output,
output_scales,
input_tensor,
expert_offsets,
blockscale_offsets,
n_experts,
)
output_scales = output_scales.view(torch.uint8)
return output, output_scales
# fp8
def scaled_fp8_quant(
input: torch.Tensor,
scale: torch.Tensor | None = None,
num_token_padding: int | None = None,
scale_ub: torch.Tensor | None = None,
use_per_token_if_dynamic: bool = False,
output: torch.Tensor | None = None,
group_shape: tuple[int, int] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize input tensor to FP8 and return quantized tensor and scale.
This function supports both static and dynamic quantization: If you
provide the scale, it will use static scaling and if you omit it,
the scale will be determined dynamically. The function also allows
optional padding of the output tensors for downstream kernels that
will benefit from padding.
Args:
input: The input tensor to be quantized to FP8 (must be 2D: [M, N])
scale: Optional scaling factor for the FP8 quantization. Supports:
- 0D or [1]: per-tensor scaling
- 1D: requires explicit group_shape to disambiguate per-channel
vs per-token (use (-1, 1) for per-channel, (1, -1) for per-token)
- 2D [M/group_m, N/group_n]: group scaling (e.g. [M, N/128] for
DeepSeek-style (1,128) groups, or [M/128, N/128] for (128,128))
scale_ub: Optional upper bound for scaling factor in dynamic
per token case
num_token_padding: If specified, pad the first dimension
of the output to at least this value.
use_per_token_if_dynamic: Whether to do per_tensor or per_token
in the dynamic quantization case.
output: Optional tensor to write the quantized result into. A new
tensor is allocated when omitted.
group_shape: Optional tuple (group_m, group_n) specifying the group
shape for static quantization. Use -1 for "full extent" (e.g.,
(-1, -1) for per-tensor, (-1, 1) for per-channel, etc.)
Required for 1D scales; optional for 2D scales.
Returns:
tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and
scaling factor.
"""
# This code assumes batch_dim and num_tokens are flattened
assert input.ndim == 2
shape: tuple[int, int] | torch.Size = input.shape
# For ROCm on MI300, the output fp8 dtype is torch.float_e3m3fnuz
out_dtype: torch.dtype = current_platform.fp8_dtype()
if num_token_padding:
shape = (max(num_token_padding, input.shape[0]), shape[1])
if output is None:
output = torch.empty(shape, device=input.device, dtype=out_dtype)
else:
assert num_token_padding is None, "padding not supported if output passed in"
assert output.dtype == out_dtype
if scale is None:
if use_per_token_if_dynamic:
scale = torch.empty((shape[0], 1), device=input.device, dtype=torch.float32)
torch.ops._C.dynamic_per_token_scaled_fp8_quant(
output, input, scale, scale_ub
)
else:
scale = torch.empty(1, device=input.device, dtype=torch.float32)
torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
else:
torch.ops._C.static_scaled_fp8_quant(output, input, scale, group_shape)
return output, scale
# gptq allspark
def allspark_repack_weight(
qweight: torch.Tensor,
scale: torch.Tensor,
zero_point: torch.Tensor | None = None,
has_zp: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format
for Ampere W8A16 Fused Gemm kernel
Args:
qweight: uint8 weight tensor, original k x n format.
scale: fp16/bf16 weight scale tensor, 1 x n format.
zero_point: fp16/bf16 weight zero_point tensor, 1 x n format.
Must be provided for asymmetric quantization.
has_zp: if use symmetric quantization, has_zp = False.
if use asymmetric quantization, has_zp = True.
Returns:
tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] :
rearranged weight, scale, and optionally zero_point.
"""
K = qweight.shape[0]
N = qweight.shape[1]
N_32align = (N + 32 - 1) // 32 * 32
qweight_reorder = torch.empty(
(N_32align, K), device=qweight.device, dtype=qweight.dtype
)
scale_reorder = torch.empty((1, N_32align), device=scale.device, dtype=scale.dtype)
zero_point_reorder = None
if has_zp:
assert zero_point is not None, (
"zero_point must be provided for asymmetric quantization."
)
zero_point_reorder = torch.empty(
(1, N_32align), device=zero_point.device, dtype=zero_point.dtype
)
torch.ops._C.rearrange_kn_weight_as_n32k16_order(
qweight,
scale,
zero_point,
has_zp,
qweight_reorder,
scale_reorder,
zero_point_reorder,
K,
N,
N_32align,
)
return qweight_reorder, scale_reorder, zero_point_reorder
def allspark_w8a16_gemm(
a: torch.Tensor,
b_qweight: torch.Tensor,
b_scales: torch.Tensor,
b_qzeros: torch.Tensor | None,
n: int,
group_size: int,
sm_count: int,
sm_version: int,
CUBLAS_M_THRESHOLD: int,
has_zp: bool,
n32k16_reorder: bool,
) -> torch.Tensor:
return torch.ops._C.allspark_w8a16_gemm(
a,
b_qweight,
b_scales,
b_qzeros,
n,
group_size,
sm_count,
sm_version,
CUBLAS_M_THRESHOLD,
has_zp,
n32k16_reorder,
)
# int8
def scaled_int8_quant(
input: torch.Tensor,
scale: torch.Tensor | None = None,
azp: torch.Tensor | None = None,
symmetric: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
"""Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
Args:
input: The input tensor to be quantized to int8.
scale: Optional scaling factor for the int8 quantization.
When not provided, we invoke dynamic-per-token quantization.
azp: Optional zero-point for the int8 quantization.
Must be provided for asymmetric quantization if `scale` is provided.
symmetric: Whether to use symmetric quantization (scale only, azp ignored).
Returns:
Output int8 tensor, scales, and optionally azp.
"""
if current_platform.is_xpu():
# XPU has no _C int8 quant op; use the torch.compile reference.
if not symmetric:
raise NotImplementedError(
"asymmetric int8 activation quantization is unsupported on XPU"
)
if scale is not None:
q = (input.to(torch.float32) / scale).round().clamp(-128, 127)
return q.to(torch.int8), scale, None
from vllm._xpu_ops import xpu_ops
q, scales, _ = xpu_ops.dynamic_per_token_int8_quant_ref(
input.contiguous(), True, 8
)
return q, scales.reshape(-1, 1).to(torch.float32), None
output = torch.empty_like(input, dtype=torch.int8)
if scale is not None:
# static-per-tensor quantization.
assert symmetric == (azp is None), (
"azp must only be provided for asymmetric quantization."
)
torch.ops._C.static_scaled_int8_quant(output, input, scale, azp)
return output, scale, azp
# dynamic-per-token quantization.
input_scales = torch.empty(
(input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
)
input_azp = None if symmetric else torch.empty_like(input_scales, dtype=torch.int32)
torch.ops._C.dynamic_scaled_int8_quant(
output, input.contiguous(), input_scales, input_azp
)
return output, input_scales, input_azp
# mamba
def selective_scan_fwd(
u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D_: torch.Tensor | None,
z_: torch.Tensor | None,
delta_bias_: torch.Tensor | None,
delta_softplus: bool,
query_start_loc: torch.Tensor | None,
cache_indices: torch.Tensor | None,
has_initial_state: torch.Tensor | None,
ssm_states: torch.Tensor,
null_block_id: int,
block_size: int = 1024,
block_idx_first_scheduled_token: torch.Tensor | None = None,
block_idx_last_scheduled_token: torch.Tensor | None = None,
initial_state_idx: torch.Tensor | None = None,
cu_chunk_seqlen: torch.Tensor | None = None,
last_chunk_indices: torch.Tensor | None = None,
):
torch.ops._C.selective_scan_fwd(
u,
delta,
A,
B,
C,
D_,
z_,
delta_bias_,
delta_softplus,
query_start_loc,
cache_indices,
has_initial_state,
ssm_states,
null_block_id,
block_size,
block_idx_first_scheduled_token,
block_idx_last_scheduled_token,
initial_state_idx,
cu_chunk_seqlen,
last_chunk_indices,
)
def causal_conv1d_update_cpu_vec(
x: torch.Tensor,
conv_state: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None = None,
activation: str | None = None,
conv_state_indices: torch.Tensor | None = None,
query_start_loc: torch.Tensor | None = None,
pad_slot_id: int = 0,
) -> torch.Tensor:
return torch.ops._C.causal_conv1d_update_cpu_vec(
x,
conv_state,
weight,
bias,
activation,
conv_state_indices,
query_start_loc,
pad_slot_id,
)
def selective_state_update_cpu(
state: torch.Tensor,
x: torch.Tensor,
dt: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None,
z: torch.Tensor | None,
dt_bias: torch.Tensor | None,
dt_softplus: bool,
state_batch_indices: torch.Tensor | None,
dst_state_batch_indices: torch.Tensor | None,
null_block_id: int,
out: torch.Tensor,
num_accepted_tokens: torch.Tensor | None,
cu_seqlens: torch.Tensor | None,
):
torch.ops._C.selective_state_update_cpu(
state,
x,
dt,
A,
B,
C,
D,
z,
dt_bias,
dt_softplus,
state_batch_indices,
dst_state_batch_indices,
null_block_id,
out,
num_accepted_tokens,
cu_seqlens,
)
def mamba_chunk_scan_fwd_cpu(
out: torch.Tensor,
final_states: torch.Tensor,
x: torch.Tensor,
dt: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None,
z: torch.Tensor | None,
cu_seqlens: torch.Tensor,
) -> None:
"""Prefill SSM scan kernel. out and final_states are written in-place."""
torch.ops._C.mamba_chunk_scan_fwd_cpu(
out,
final_states,
x,
dt,
A,
B,
C,
D,
z,
cu_seqlens,
)
# ROCm skinny gemms
def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor:
return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
def wvSplitK(
a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
) -> torch.Tensor:
return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count)
def wvSplitK_int4_g(
weight: torch.Tensor,
activation: torch.Tensor,
scale: torch.Tensor,
cu_count: int,
group_size: int,
zero_points: torch.Tensor | None = None,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
# NOTE: the kernel is weight-major; `weight` is the packed int4 operand
# (in_a, [out_features, K/2]) and `activation` is in_b ([num_tokens, K]).
return torch.ops._rocm_C.wvSplitK_int4_g(
weight, activation, scale, zero_points, bias, cu_count, group_size
)
if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"):
@register_fake("_rocm_C::wvSplitK_int4_g")
def _wvSplitK_int4_g_fake(
in_a: torch.Tensor, # packed int4 weight [out_features, K/2]
in_b: torch.Tensor, # activation [num_tokens, K]
in_scale: torch.Tensor,
in_zero_points: torch.Tensor | None,
in_bias: torch.Tensor | None,
CuCount: int,
group_size: int,
) -> torch.Tensor:
# Kernel returns {in_b.size(0), in_a.size(0)} = [num_tokens, out_features].
num_tokens = in_b.size(0)
out_features = in_a.size(0)
return torch.empty(
(num_tokens, out_features), dtype=in_b.dtype, device=in_b.device
)
def wvSplitKrc(
a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
) -> torch.Tensor:
return torch.ops._rocm_C.wvSplitKrc(a, b, bias, cu_count)
def wvSplitKQ(
a: torch.Tensor,
b: torch.Tensor,
out_dtype: torch.dtype,
scale_a: torch.Tensor,
scale_b: torch.Tensor,
cu_count: int,
bias: torch.Tensor = None,
) -> torch.Tensor:
out = torch.empty((b.shape[0], a.shape[0]), dtype=out_dtype, device=b.device)
torch.ops._rocm_C.wvSplitKQ(a, b, bias, out, scale_a, scale_b, cu_count)
return out
# moe
def moe_sum(
input: torch.Tensor,
output: torch.Tensor,
topk_ids: torch.Tensor | None = None,
expert_map: torch.Tensor | None = None,
):
torch.ops._moe_C.moe_sum(input, output, topk_ids, expert_map)
def moe_align_block_size(
topk_ids: torch.Tensor,
num_experts: int,
block_size: int,
sorted_token_ids: torch.Tensor,
experts_ids: torch.Tensor,
num_tokens_post_pad: torch.Tensor,
expert_map: torch.Tensor | None = None,
scatter_idx: torch.Tensor | None = None,
) -> None:
if current_platform.is_xpu():
if scatter_idx is not None:
raise NotImplementedError("scatter_idx is not supported on XPU")
torch.ops._moe_C.moe_align_block_size(
topk_ids,
num_experts,
block_size,
sorted_token_ids,
experts_ids,
num_tokens_post_pad,
expert_map,
)
return
torch.ops._moe_C.moe_align_block_size(
topk_ids,
num_experts,
block_size,
sorted_token_ids,
experts_ids,
num_tokens_post_pad,
expert_map,
scatter_idx,
)
def batched_moe_align_block_size(
max_tokens_per_batch: int,
block_size: int,
expert_num_tokens: torch.Tensor,
sorted_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_post_pad: torch.Tensor,
) -> None:
torch.ops._moe_C.batched_moe_align_block_size(
max_tokens_per_batch,
block_size,
expert_num_tokens,
sorted_ids,
expert_ids,
num_tokens_post_pad,
)
def moe_lora_align_block_size(
topk_ids: torch.Tensor,
token_lora_mapping: torch.Tensor,
num_experts: int,
block_size: int,
max_loras: int,
max_num_tokens_padded: int,
max_num_m_blocks: int,
sorted_token_ids: torch.Tensor,
experts_ids: torch.Tensor,
num_tokens_post_pad: torch.Tensor,
adapter_enabled: torch.Tensor,
lora_ids: torch.Tensor,
expert_map: torch.Tensor | None = None,
) -> None:
torch.ops._moe_C.moe_lora_align_block_size(
topk_ids,
token_lora_mapping,
num_experts,
block_size,
max_loras,
max_num_tokens_padded,
max_num_m_blocks,
sorted_token_ids,
experts_ids,
num_tokens_post_pad,
adapter_enabled,
lora_ids,
expert_map,
)
def moe_wna16_gemm(
input: torch.Tensor,
output: torch.Tensor,
b_qweight: torch.Tensor,
b_scales: torch.Tensor,
b_qzeros: torch.Tensor | None,
topk_weights: torch.Tensor | None,
sorted_token_ids: torch.Tensor,
experts_ids: torch.Tensor,
num_tokens_post_pad: torch.Tensor,
top_k: int,
BLOCK_SIZE_M: int,
BLOCK_SIZE_N: int,
BLOCK_SIZE_K: int,
bit: int,
) -> torch.Tensor:
if not current_platform.is_cuda():
raise NotImplementedError(
"The optimized moe_wna16_gemm kernel is only available on CUDA platforms"
)
torch.ops._moe_C.moe_wna16_gemm(
input,
output,
b_qweight,
b_scales,
b_qzeros,
topk_weights,
sorted_token_ids,
experts_ids,
num_tokens_post_pad,
top_k,
BLOCK_SIZE_M,
BLOCK_SIZE_N,
BLOCK_SIZE_K,
bit,
)
def fp32_router_gemm(
hidden_states: torch.Tensor,
router_weight: torch.Tensor,
) -> torch.Tensor:
output = torch.empty(
hidden_states.shape[0],
router_weight.shape[0],
device=hidden_states.device,
dtype=torch.float32,
)
torch.ops._C.fp32_router_gemm(output, hidden_states, router_weight)
return output
def topk_softmax(
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
token_expert_indices: torch.Tensor,
gating_output: torch.Tensor,
renormalize: bool = False,
e_score_correction_bias: torch.Tensor | None = None,
is_padding: torch.Tensor | None = None,
) -> None:
torch.ops._moe_C.topk_softmax(
topk_weights,
topk_ids,
token_expert_indices,
gating_output,
renormalize,
e_score_correction_bias,
is_padding,
)
def topk_sigmoid(
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
token_expert_indices: torch.Tensor,
gating_output: torch.Tensor,
renormalize: bool = False,
e_score_correction_bias: torch.Tensor | None = None,
routed_scaling_factor: float = 1.0,
is_padding: torch.Tensor | None = None,
) -> None:
torch.ops._moe_C.topk_sigmoid(
topk_weights,
topk_ids,
token_expert_indices,
gating_output,
renormalize,
e_score_correction_bias,
routed_scaling_factor,
is_padding,
)
def topk_hash_softplus_sqrt(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
token_expert_indices: torch.Tensor,
gating_output: torch.Tensor,
renormalize: bool = False,
routed_scaling_factor: float = 1.0,
e_score_correction_bias: torch.Tensor | None = None,
input_tokens: torch.Tensor | None = None,
hash_indices_table: torch.Tensor | None = None,
is_padding: torch.Tensor | None = None,
bias_vl: torch.Tensor | None = None,
image_sentinel_lo: int = 0,
) -> None:
torch.ops._moe_C.topk_softplus_sqrt(
topk_weights,
topk_indices,
token_expert_indices,
gating_output,
renormalize,
routed_scaling_factor,
e_score_correction_bias,
input_tokens,
hash_indices_table,
is_padding,
bias_vl,
image_sentinel_lo,
)
def grouped_topk(
scores: torch.Tensor,
num_expert_group: int,
topk_group: int,
topk: int,
renormalize: bool,
routed_scaling_factor: float,
bias: torch.Tensor,
scoring_func: int = 0,
):
"""Perform grouped top-k routing for mixture of experts.
Args:
scores: Raw inputs (logits if scoring_func=1, scores if scoring_func=0)
num_expert_group: Number of expert groups
topk_group: Number of groups to select
topk: Number of experts to select per token
renormalize: Whether to renormalize the output weights
routed_scaling_factor: Scaling factor for routing weights
bias: Bias tensor (e_score_correction_bias). Always fused in kernel.
scoring_func: 0=none (no activation), 1=sigmoid
"""
if not (current_platform.is_cuda() or current_platform.is_xpu()):
raise NotImplementedError(
"The fused grouped_topk kernel is only available on CUDA and XPU platforms"
)
return torch.ops._moe_C.grouped_topk(
scores,
num_expert_group,
topk_group,
topk,
renormalize,
routed_scaling_factor,
bias,
scoring_func,
)
if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "grouped_topk"):
@register_fake("_moe_C::grouped_topk")
def _grouped_topk_fake(
scores: torch.Tensor,
num_expert_group: int,
topk_group: int,
topk: int,
renormalize: bool,
routed_scaling_factor: float,
bias: torch.Tensor,
scoring_func: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
num_tokens = scores.size(0)
return (
scores.new_empty((num_tokens, topk), dtype=torch.float32),
scores.new_empty((num_tokens, topk), dtype=torch.int32),
)
def moe_wna16_marlin_gemm(
input: torch.Tensor,
output: torch.Tensor | None,
b_qweight: torch.Tensor,
b_bias: torch.Tensor | None,
b_scales: torch.Tensor,
a_scales: torch.Tensor | None,
global_scale: torch.Tensor | None,
b_qzeros: torch.Tensor | None,
workspace: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_past_padded: torch.Tensor,
topk_weights: torch.Tensor,
moe_block_size: int,
top_k: int,
mul_topk_weights: bool,
b_q_type: ScalarType,
size_m: int,
size_n: int,
size_k: int,
use_atomic_add: bool,
use_fp32_reduce: bool,
is_zp_float: bool,
thread_k: int = -1,
thread_n: int = -1,
blocks_per_sm: int = -1,
) -> torch.Tensor:
return torch.ops._moe_C.moe_wna16_marlin_gemm(
input,
output,
b_qweight,
b_bias,
b_scales,
a_scales,
global_scale,
b_qzeros,
workspace,
sorted_token_ids,
expert_ids,
num_tokens_past_padded,
topk_weights,
moe_block_size,
top_k,
mul_topk_weights,
b_q_type.id,
size_m,
size_n,
size_k,
use_atomic_add,
use_fp32_reduce,
is_zp_float,
thread_k,
thread_n,
blocks_per_sm,
)
if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_gemm"):
@register_fake("_moe_C::moe_wna16_marlin_gemm")
def moe_wna16_marlin_gemm_fake(
input: torch.Tensor,
output: torch.Tensor | None,
b_qweight: torch.Tensor,
b_bias: torch.Tensor | None,
b_scales: torch.Tensor,
a_scales: torch.Tensor | None,
global_scale: torch.Tensor | None,
b_qzeros: torch.Tensor | None,
workspace: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_past_padded: torch.Tensor,
topk_weights: torch.Tensor,
moe_block_size: int,
top_k: int,
mul_topk_weights: bool,
b_q_type: ScalarType,
size_m: int,
size_n: int,
size_k: int,
use_atomic_add: bool,
use_fp32_reduce: bool,
is_zp_float: bool,
):
return torch.empty(
(size_m * top_k, size_n), dtype=input.dtype, device=input.device
)
def reshape_and_cache(
key: torch.Tensor,
value: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
slot_mapping: torch.Tensor,
kv_cache_dtype: str,
k_scale: torch.Tensor,
v_scale: torch.Tensor,
) -> None:
torch.ops._C_cache_ops.reshape_and_cache(
key,
value,
key_cache,
value_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
v_scale,
)
def reshape_and_cache_flash(
key: torch.Tensor,
value: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
slot_mapping: torch.Tensor,
kv_cache_dtype: str,
k_scale: torch.Tensor,
v_scale: torch.Tensor,
) -> None:
torch.ops._C_cache_ops.reshape_and_cache_flash(
key,
value,
key_cache,
value_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
v_scale,
)
def fused_minimax_m3_qknorm_rope_kv_insert(
qkv: torch.Tensor,
q_norm_weight: torch.Tensor,
k_norm_weight: torch.Tensor,
cos_sin_cache: torch.Tensor,
positions: torch.Tensor,
num_heads: int,
num_kv_heads: int,
rotary_dim: int,
eps: float,
index_q_norm_weight: torch.Tensor | None = None,
index_k_norm_weight: torch.Tensor | None = None,
num_index_heads: int = 0,
slot_mapping: torch.Tensor | None = None,
index_slot_mapping: torch.Tensor | None = None,
kv_cache: torch.Tensor | None = None,
index_cache: torch.Tensor | None = None,
block_size: int = 0,
q_out: torch.Tensor | None = None,
index_q_out: torch.Tensor | None = None,
kv_cache_dtype: str = "auto",
skip_index_branch: bool = False,
q_fp8_out: torch.Tensor | None = None,
q_fp8_scale: float = 1.0,
) -> None:
"""Fused MiniMax-M3 attention pre-processing (in-place).
Applies Gemma RMSNorm + partial NeoX RoPE to ``qkv`` in place. ``qkv`` is a
single fused tensor:
- dense layer (``num_index_heads == 0``): ``[q | k | v]``;
- sparse layer (``num_index_heads > 0``): ``[q | k | v | index_q |
index_k]`` — the index branch is read straight out of ``qkv``.
When ``kv_cache`` is given (sparse serving), also scatter-inserts the
normed/roped k & v into the paged KV cache by ``slot_mapping`` and the
index key into ``index_cache`` by ``index_slot_mapping``. ``kv_cache_dtype``
selects the cache storage/conversion path. If
``index_slot_mapping`` is omitted, ``slot_mapping`` is used for both caches.
If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N,
niq*128]``) are given, the normed/roped q / index_q are written there
instead of in place — folding the de-interleave into this kernel's store so
callers skip a separate ``.contiguous()`` copy before the SM100 sparse
attention's flat TMA descriptor.
If ``q_fp8_out`` is given, the same normalized q is also written in FP8
E4M3 using ``q_fp8_scale`` as its dequantization scale.
When ``skip_index_branch`` is true, sparse rows still keep their packed
``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v
branches and main KV cache. This is used by MiniMax-M3 index-topk reuse
layers that consume top-k block ids selected by an earlier sparse layer.
"""
torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert(
qkv,
q_norm_weight,
k_norm_weight,
cos_sin_cache,
positions,
num_heads,
num_kv_heads,
rotary_dim,
eps,
index_q_norm_weight,
index_k_norm_weight,
num_index_heads,
slot_mapping,
index_slot_mapping,
kv_cache,
index_cache,
block_size,
q_out,
index_q_out,
kv_cache_dtype,
skip_index_branch,
q_fp8_out,
q_fp8_scale,
)
def fused_kda_decode(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None,
conv_state: torch.Tensor,
raw_g: torch.Tensor,
raw_beta: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
state_indices: torch.Tensor,
state: torch.Tensor,
out: torch.Tensor | None = None,
lower_bound: float | None = None,
output_gate: torch.Tensor | None = None,
norm_weight: torch.Tensor | None = None,
norm_eps: float = 1e-5,
) -> torch.Tensor:
if out is None:
out = torch.empty(
1,
x.shape[0],
raw_g.shape[2],
raw_g.shape[3],
dtype=x.dtype,
device=x.device,
)
torch.ops._C.fused_kda_decode(
x,
weight,
bias,
conv_state,
raw_g,
raw_beta,
A_log,
dt_bias,
state_indices,
state,
out,
lower_bound,
output_gate,
norm_weight,
norm_eps,
)
return out
def fused_gdn_decode_post_conv_mtp(
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
state_indices: torch.Tensor,
cu_seqlens: torch.Tensor,
num_accepted_tokens: torch.Tensor,
state: torch.Tensor,
output_gate: torch.Tensor,
norm_weight: torch.Tensor,
out: torch.Tensor | None = None,
scale: float = 128**-0.5,
norm_eps: float = 1e-5,
output_gate_activation: str = "silu",
) -> torch.Tensor:
if out is None:
out = torch.empty_like(output_gate)
torch.ops._C.fused_gdn_decode_post_conv_mtp(
mixed_qkv,
a,
b,
A_log,
dt_bias,
state_indices,
cu_seqlens,
num_accepted_tokens,
state,
output_gate,
norm_weight,
out,
scale,
norm_eps,
output_gate_activation,
)
return out
def concat_and_cache_mla(
kv_c: torch.Tensor,
k_pe: torch.Tensor,
kv_cache: torch.Tensor,
slot_mapping: torch.Tensor,
kv_cache_dtype: str,
scale: torch.Tensor,
) -> None:
torch.ops._C_cache_ops.concat_and_cache_mla(
kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale
)
def concat_and_cache_mla_grouped(
kv_c: torch.Tensor,
k_pe: torch.Tensor,
kv_cache_ptrs: torch.Tensor,
slot_mapping: torch.Tensor,
block_size: int,
block_stride: int,
entry_stride: int,
kv_scales: torch.Tensor | None = None,
kv_cache_dtype: str = "auto",
) -> None:
torch.ops._C_cache_ops.concat_and_cache_mla_grouped(
kv_c,
k_pe,
kv_cache_ptrs,
slot_mapping,
block_size,
block_stride,
entry_stride,
kv_scales,
kv_cache_dtype,
)
def kimi_k3_attn_res(
prefix: torch.Tensor,
delta: torch.Tensor | None,
blocks: torch.Tensor,
norm_weight: torch.Tensor,
qk_weight: torch.Tensor,
output_norm_weight: torch.Tensor | None,
num_blocks: int,
block_write_idx: int,
eps: float,
output_norm_eps: float,
) -> torch.Tensor:
output = torch.empty_like(prefix)
torch.ops._C.kimi_k3_attn_res(
prefix,
delta,
blocks,
norm_weight,
qk_weight,
output_norm_weight,
output,
num_blocks,
block_write_idx,
eps,
output_norm_eps,
)
return output
def concat_and_cache_mla_rope_fused(
positions: torch.Tensor,
q_pe: torch.Tensor,
k_pe: torch.Tensor,
kv_c: torch.Tensor,
cos_sin_cache: torch.Tensor,
is_neox: bool,
slot_mapping: torch.Tensor,
kv_cache: torch.Tensor,
kv_cache_dtype: str,
kv_cache_scale: torch.Tensor,
) -> None:
torch.ops._C_cache_ops.concat_and_cache_mla_rope_fused(
positions,
q_pe,
k_pe,
kv_c,
cos_sin_cache,
is_neox,
slot_mapping,
kv_cache,
kv_cache_dtype,
kv_cache_scale,
)
def swap_blocks(
src: torch.Tensor,
dst: torch.Tensor,
block_size_in_bytes: int,
block_mapping: torch.Tensor,
) -> None:
"""Copy specific blocks from one tensor to another.
This method assumes each of the two input tensors is composed of
consecutive contiguous blocks, of size block_size_in_bytes.
i.e. the memory layout for each tensor is:
[block0] [block1] ... [block N]
block_mapping determines the subset of blocks to copy of the source tensor,
and their matching destination block number on the destination tensor.
block_mapping is expected to be a tensor of shape (num_blocks_to_copy, 2)
where each block_mapping[i] represents a single copy operation, copying
block #block_mapping[i][0] from the source tensor
to block #block_mapping[i][1] on the destination tensor.
block_mapping should have dtype int64.
The source and the destination tensors can be either on cpu or gpu,
but not both on cpu.
the block mapping tensor must on cpu.
"""
torch.ops._C_cache_ops.swap_blocks(src, dst, block_size_in_bytes, block_mapping)
def swap_blocks_batch(
src_ptrs: torch.Tensor,
dst_ptrs: torch.Tensor,
sizes: torch.Tensor,
is_src_access_order_any: bool = False,
) -> None:
"""Batch version of swap_blocks: submit all copies in a single driver call.
Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i]
of sizes[i] bytes. All three tensors must be CPU tensors with the
platform-appropriate pointer dtype: int64 on CUDA/ROCm (required by
cache_kernels.cu) and uint64 on XPU (required by the XPU DMA engine).
On CUDA 12.8+ this uses cuMemcpyBatchAsync for minimal submission
overhead; on older CUDA it falls back to a loop of cudaMemcpyAsync.
is_src_access_order_any: if True, pass CU_MEMCPY_SRC_ACCESS_ORDER_ANY to
cuMemcpyBatchAsync, letting the DMA engine prefetch source bytes
out of stream order. Only safe when no GPU stream is concurrently
writing to the source. Defaults to False (STREAM ordering), which
is always safe.
"""
if current_platform.is_xpu():
torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes)
else:
torch.ops._C_cache_ops.swap_blocks_batch(
src_ptrs, dst_ptrs, sizes, is_src_access_order_any
)
def convert_fp8(
output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8"
) -> None:
torch.ops._C_cache_ops.convert_fp8(output, input, scale, kv_dtype)
def gather_and_maybe_dequant_cache(
src_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
cu_seq_lens: torch.Tensor,
token_to_seq: torch.Tensor,
num_tokens: int,
kv_cache_dtype: str,
scale: torch.Tensor,
seq_starts: torch.Tensor | None = None,
) -> None:
torch.ops._C_cache_ops.gather_and_maybe_dequant_cache(
src_cache,
dst,
block_table,
cu_seq_lens,
token_to_seq,
num_tokens,
kv_cache_dtype,
scale,
seq_starts,
)
def cp_gather_cache(
src_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
cu_seq_lens: torch.Tensor,
batch_size: int,
seq_starts: torch.Tensor | None = None,
) -> None:
torch.ops._C_cache_ops.cp_gather_cache(
src_cache, dst, block_table, cu_seq_lens, batch_size, seq_starts
)
def cp_gather_and_upconvert_fp8_kv_cache(
src_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
workspace_starts: torch.Tensor,
batch_size: int,
seq_starts: torch.Tensor | None = None,
host_cache: torch.Tensor | None = None,
host_row_ids: torch.Tensor | None = None,
device_row_ids: torch.Tensor | None = None,
) -> None:
"""Gather and upconvert FP8 KV cache to BF16 workspace.
Args:
src_cache: FP8 KV cache [num_blocks, block_size, 656]
dst: BF16 output workspace [total_tokens, 576]
block_table: Block indices [num_reqs, max_blocks]
workspace_starts: Workspace start offsets [num_reqs]
batch_size: Number of requests
seq_starts: Optional source sequence offsets [num_reqs]
host_cache: Optional pinned host rows used for non-resident entries
host_row_ids: Host source row for each remapped cache row
device_row_ids: Device source row, or -1, for each remapped cache row
"""
torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache(
src_cache,
dst,
block_table,
workspace_starts,
batch_size,
seq_starts,
host_cache,
host_row_ids,
device_row_ids,
)
def cp_gather_and_upconvert_nvfp4_kv_cache(
src_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
workspace_starts: torch.Tensor,
batch_size: int,
) -> None:
"""Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace.
Args:
src_cache: NVFP4 KV cache [num_blocks, block_size, 352] (uint8)
dst: BF16 output workspace [total_tokens, 576]
block_table: Block indices [num_reqs, max_blocks]
workspace_starts: Workspace start offsets [num_reqs]
batch_size: Number of requests
"""
torch.ops._C_cache_ops.cp_gather_and_upconvert_nvfp4_kv_cache(
src_cache,
dst,
block_table,
workspace_starts,
batch_size,
)
def concat_mla_q(
ql_nope: torch.Tensor,
q_pe: torch.Tensor,
q_out: torch.Tensor,
) -> None:
"""Concatenate query nope and rope for MLA/DSA attention.
Args:
ql_nope: Query nope component [num_tokens, num_heads, nope_dim]
q_pe: Query rope component [num_tokens, num_heads, rope_dim]
q_out: Output tensor [num_tokens, num_heads, nope_dim + rope_dim]
"""
torch.ops._C_cache_ops.concat_mla_q(ql_nope, q_pe, q_out)
def indexer_k_quant_and_cache(
k: torch.Tensor,
kv_cache: torch.Tensor,
slot_mapping: torch.Tensor,
quant_block_size: int,
kv_cache_dtype: str,
) -> None:
torch.ops._C_cache_ops.indexer_k_quant_and_cache(
k, kv_cache, slot_mapping, quant_block_size, kv_cache_dtype
)
def top_k_per_row_prefill(
logits: torch.Tensor,
cu_seqlen_ks: torch.Tensor,
cu_seqlen_ke: torch.Tensor,
raw_topk_indices: torch.Tensor,
num_rows: int,
stride0: int,
stride1: int,
topk_tokens: int,
) -> None:
torch.ops._C.top_k_per_row_prefill(
logits,
cu_seqlen_ks,
cu_seqlen_ke,
raw_topk_indices,
num_rows,
stride0,
stride1,
topk_tokens,
)
def top_k_per_row_decode(
logits: torch.Tensor,
next_n: int,
seq_lens: torch.Tensor,
raw_topk_indices: torch.Tensor,
num_rows: int,
stride0: int,
stride1: int,
topk_tokens: int,
) -> None:
torch.ops._C.top_k_per_row_decode(
logits,
next_n,
seq_lens,
raw_topk_indices,
num_rows,
stride0,
stride1,
topk_tokens,
)
def cp_gather_indexer_k_quant_cache(
kv_cache: torch.Tensor,
dst_k: torch.Tensor,
dst_scale: torch.Tensor,
block_table: torch.Tensor,
cu_seq_lens: torch.Tensor,
) -> None:
torch.ops._C_cache_ops.cp_gather_indexer_k_quant_cache(
kv_cache, dst_k, dst_scale, block_table, cu_seq_lens
)
def get_device_attribute(attribute: int, device: int) -> int:
return torch.ops._C_cuda_utils.get_device_attribute(attribute, device)
def get_max_shared_memory_per_block_device_attribute(device: int) -> int:
# ruff: noqa: E501
return torch.ops._C_cuda_utils.get_max_shared_memory_per_block_device_attribute(
device
)
# custom ar
def init_custom_ar(
ipc_tensors: list[torch.Tensor],
rank_data: torch.Tensor,
rank: int,
fully_connected: bool,
) -> int:
return torch.ops._C_custom_ar.init_custom_ar(
ipc_tensors, rank_data, rank, fully_connected
)
def all_reduce(
fa: int,
inp: torch.Tensor,
out: torch.Tensor,
reg_buffer: int,
reg_buffer_sz_bytes: int,
) -> None:
torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
def custom_all_gather(
fa: int,
inp: torch.Tensor,
out: torch.Tensor,
reg_buffer: int,
reg_buffer_sz_bytes: int,
) -> None:
torch.ops._C_custom_ar.custom_all_gather(
fa, inp, out, reg_buffer, reg_buffer_sz_bytes
)
def mnnvl_lamport_all_gather(
fa: int,
inp: torch.Tensor,
out: torch.Tensor,
local_buffer: int,
multicast_buffer: int,
epoch_buffer: int,
stage_sz_bytes: int,
) -> None:
torch.ops._C_custom_ar.mnnvl_lamport_all_gather(
fa,
inp,
out,
local_buffer,
multicast_buffer,
epoch_buffer,
stage_sz_bytes,
)
def custom_reduce_scatter(
fa: int,
inp: torch.Tensor,
out: torch.Tensor,
reg_buffer: int,
reg_buffer_sz_bytes: int,
) -> None:
torch.ops._C_custom_ar.custom_reduce_scatter(
fa, inp, out, reg_buffer, reg_buffer_sz_bytes
)
def mnnvl_lamport_reduce_scatter(
fa: int,
inp: torch.Tensor,
out: torch.Tensor,
local_buffer: int,
epoch_buffer: int,
stage_sz_bytes: int,
) -> None:
torch.ops._C_custom_ar.mnnvl_lamport_reduce_scatter(
fa, inp, out, local_buffer, epoch_buffer, stage_sz_bytes
)
def dispose(fa: int) -> None:
torch.ops._C_custom_ar.dispose(fa)
def meta_size() -> int:
return torch.ops._C_custom_ar.meta_size()
def register_buffer(fa: int, ipc_tensors: list[int]) -> None:
return torch.ops._C_custom_ar.register_buffer(fa, ipc_tensors)
def get_graph_buffer_ipc_meta(fa: int) -> tuple[list[int], list[int]]:
return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa)
def register_graph_buffers(
fa: int, handles: list[list[int]], offsets: list[list[int]]
) -> None:
torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets)
def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]:
return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size)
def open_mem_handle(mem_handle: torch.Tensor):
return torch.ops._C_custom_ar.open_mem_handle(mem_handle)
def free_shared_buffer(ptr: int) -> None:
torch.ops._C_custom_ar.free_shared_buffer(ptr)
# quick all reduce
def init_custom_qr(rank: int, world_size: int, qr_max_size: int | None = None) -> int:
return torch.ops._C_custom_ar.init_custom_qr(rank, world_size, qr_max_size)
def qr_destroy(fa: int) -> None:
torch.ops._C_custom_ar.qr_destroy(fa)
def qr_all_reduce(
fa: int,
inp: torch.Tensor,
out: torch.Tensor,
quant_level: int,
cast_bf2half: bool = False,
) -> None:
torch.ops._C_custom_ar.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half)
def qr_get_handle(fa: int) -> torch.Tensor:
return torch.ops._C_custom_ar.qr_get_handle(fa)
def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None:
return torch.ops._C_custom_ar.qr_open_handles(fa, handles)
def qr_max_size() -> int:
return torch.ops._C_custom_ar.qr_max_size()
def sm100_cutlass_mla_decode(
out: torch.Tensor,
lse: torch.Tensor,
q_nope: torch.Tensor,
q_pe: torch.Tensor,
kv_c_and_k_pe_cache: torch.Tensor,
seq_lens: torch.Tensor,
page_table: torch.Tensor,
workspace: torch.Tensor,
scale: float,
num_kv_splits: int,
) -> torch.Tensor:
torch.ops._C.sm100_cutlass_mla_decode(
out,
lse,
q_nope,
q_pe,
kv_c_and_k_pe_cache,
seq_lens,
page_table,
workspace,
scale,
num_kv_splits,
)
return out
def sm100_cutlass_mla_get_workspace_size(
max_seq_len: int, num_batches: int, sm_count: int, num_kv_splits: int
) -> int:
return torch.ops._C.sm100_cutlass_mla_get_workspace_size(
max_seq_len, num_batches, sm_count, num_kv_splits
)
def dsv3_fused_a_gemm(
output: torch.Tensor,
mat_a: torch.Tensor,
mat_b: torch.Tensor,
enable_pdl: bool = False,
) -> None:
"""Low-latency fused-A-style GEMM (SM 9.0+, BF16, 1-16 tokens).
Computes ``output = mat_a @ mat_b`` for the compiled Kimi K3 and
DeepSeek V3 projection shapes. ``mat_a`` and ``output`` are row-major;
``mat_b`` is the column-major transposed weight. ``enable_pdl`` permits
programmatic dependent launch for callers that have validated it.
Requires SM 9.0+.
"""
torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl)
def weight_packed_linear_cpu(
x: torch.Tensor,
packed_weight: torch.Tensor,
n: int,
bias: torch.Tensor | None,
) -> torch.Tensor:
output = torch.empty((*x.shape[0:-1], n), dtype=x.dtype)
torch.ops._C.weight_packed_linear(
output.reshape(-1, n), x.reshape(-1, x.size(-1)), packed_weight, bias, True
)
return output
class CPUQuantMethod(IntEnum):
UNQUANT = 0
INT8_W8A8 = 1
FP8_W8A16 = 2
INT4_W4A8 = 3
MXFP4 = 4
FP8_W8A8 = 5
if hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
@register_fake("_C::dynamic_4bit_int_moe")
def dynamic_4bit_int_moe_fake(
x: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
w13_packed: torch.Tensor,
w2_packed: torch.Tensor,
hidden_size: int,
intermediate_size: int,
group_size: int,
apply_router_weight_on_input: bool,
activation_kind: int,
) -> torch.Tensor:
return x.new_empty((x.size(0), hidden_size))
def fused_experts_cpu(
out: torch.Tensor,
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
moe_comp_method: CPUQuantMethod,
w1_scale: torch.Tensor | None,
w2_scale: torch.Tensor | None,
w1_zero: torch.Tensor | None,
w2_zero: torch.Tensor | None,
block_size: list[int] | None,
w1_bias: torch.Tensor | None = None,
w2_bias: torch.Tensor | None = None,
alpha: float | None = None,
limit: float | None = None,
is_vnni: bool = True,
a1_scale: torch.Tensor | None = None,
) -> None:
torch.ops._C.fused_experts_cpu(
out,
hidden_states,
w1,
w2,
topk_weights,
topk_ids,
moe_comp_method,
w1_scale,
w2_scale,
w1_zero,
w2_zero,
a1_scale,
block_size,
w1_bias,
w2_bias,
alpha,
limit,
is_vnni,
)
if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
@register_fake("_C::int8_scaled_mm_with_quant")
def int8_scaled_mm_with_quant_fake(
mat1: torch.Tensor,
mat2: torch.Tensor,
scales2: torch.Tensor,
bias: torch.Tensor | None,
out_dtype: torch.dtype,
is_vnni: bool,
) -> torch.Tensor:
M = mat1.size(0)
N = mat2.size(0)
return torch.empty((M, N), dtype=out_dtype)
class CPUQuantAlgo(IntEnum):
AWQ = 0
GPTQ = 1
if hasattr(torch.ops._C, "convert_weight_packed_scale_zp"):
@register_fake("_C::convert_weight_packed_scale_zp")
def convert_weight_packed_scale_zp_fake(
qweight: torch.Tensor,
qzeros: torch.Tensor,
scales: torch.Tensor,
quant_method_4bit: CPUQuantAlgo,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return (
torch.empty_like(qweight),
torch.empty_like(qzeros),
torch.empty_like(scales),
)
def convert_weight_packed_scale_zp(
qweight: torch.Tensor,
qzeros: torch.Tensor,
scales: torch.Tensor,
quant_method_4bit: CPUQuantAlgo,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return torch.ops._C.convert_weight_packed_scale_zp(
qweight,
qzeros,
scales,
quant_method_4bit,
)
if hasattr(torch.ops._C, "int4_scaled_mm_cpu"):
@register_fake("_C::int4_scaled_mm_cpu")
def int4_scaled_mm_cpu_fake(
x: torch.Tensor,
w: torch.Tensor,
w_zeros: torch.Tensor,
w_scales: torch.Tensor,
bias: torch.Tensor | None,
) -> torch.Tensor:
N = w_scales.size(0) * w_scales.size(-1)
return torch.empty((x.size(0), N), dtype=x.dtype, device=x.device)
def int4_scaled_mm_cpu(
x: torch.Tensor,
w: torch.Tensor,
w_zeros: torch.Tensor,
w_scales: torch.Tensor,
bias: torch.Tensor | None,
) -> torch.Tensor:
x_shape = x.shape
x_2d = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x
out = torch.ops._C.int4_scaled_mm_cpu(
x_2d,
w,
w_zeros,
w_scales,
bias,
)
out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out
return out
if hasattr(torch.ops._C, "fp8_scaled_mm_cpu"):
@register_fake("_C::fp8_scaled_mm_cpu")
def fp8_scaled_mm_cpu_fake(
mat1: torch.Tensor,
mat2: torch.Tensor,
scales2: torch.Tensor,
block_size: list[int],
bias: torch.Tensor | None,
out_dtype: torch.dtype,
is_vnni: bool,
) -> torch.Tensor:
M = mat1.size(0)
N = mat2.size(0)
return torch.empty((M, N), dtype=out_dtype, device=mat1.device)
_supports_cpu_fp8_w8a16 = bool(hasattr(torch.ops._C, "fp8_scaled_mm_cpu"))
def fp8_scaled_mm_cpu(
mat1: torch.Tensor,
mat2: torch.Tensor,
scales2: torch.Tensor,
block_size: list[int],
bias: torch.Tensor | None,
out_dtype: torch.dtype,
is_vnni: bool,
) -> torch.Tensor:
return torch.ops._C.fp8_scaled_mm_cpu(
mat1, mat2, scales2, block_size, bias, out_dtype, is_vnni
)
# FP8 W8A8 CPU kernels
@cache
def cpu_has_amx_fp8() -> bool:
"""Whether this CPU has native AMX-FP8 MMA support."""
if not hasattr(torch.ops._C, "cpu_has_amx_fp8"):
return False
return bool(torch.ops._C.cpu_has_amx_fp8())
if hasattr(torch.ops._C, "fp8_scaled_mm_with_quant"):
@register_fake("_C::fp8_scaled_mm_with_quant")
def fp8_scaled_mm_with_quant_fake(
act: torch.Tensor,
act_scales: torch.Tensor | None,
channelwise: bool,
weight: torch.Tensor,
weight_scales: torch.Tensor,
bias: torch.Tensor | None,
output_dtype: torch.dtype,
) -> torch.Tensor:
M = act.reshape(-1, act.size(-1)).size(0)
N = weight.size(0) * weight.size(-1)
return torch.empty((M, N), dtype=output_dtype, device=act.device)
def fp8_scaled_mm_with_quant(
act: torch.Tensor,
act_scales: torch.Tensor | None,
channelwise: bool,
weight: torch.Tensor,
weight_scales: torch.Tensor,
bias: torch.Tensor | None,
output_dtype: torch.dtype,
) -> torch.Tensor:
return torch.ops._C.fp8_scaled_mm_with_quant(
act, act_scales, channelwise, weight, weight_scales, bias, output_dtype
)
def chunk_gated_delta_rule_cpu(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
initial_state: torch.Tensor,
output_final_state: bool,
cu_seqlens: torch.Tensor,
head_first: bool,
use_qk_l2norm_in_kernel: bool,
initial_state_indices: torch.Tensor,
eps: float = 1e-5,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.ops._C.chunk_gated_delta_rule_cpu(
query,
key,
value,
g,
beta,
initial_state,
output_final_state,
cu_seqlens,
head_first,
use_qk_l2norm_in_kernel,
initial_state_indices,
eps,
)
def fused_sigmoid_gating_delta_rule_update_cpu(
A_log: torch.Tensor,
dt_bias: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
initial_state_source: torch.Tensor,
initial_state_indices: torch.Tensor,
cu_seqlens: torch.Tensor,
use_qk_l2norm_in_kernel: bool,
softplus_beta: float = 1.0,
softplus_threshold: float = 20.0,
) -> torch.Tensor:
return torch.ops._C.fused_sigmoid_gating_delta_rule_update_cpu(
A_log,
dt_bias,
q,
k,
v,
a,
b,
initial_state_source,
initial_state_indices,
cu_seqlens,
use_qk_l2norm_in_kernel,
softplus_beta,
softplus_threshold,
)
def fused_sigmoid_gating_delta_rule_update_spec_cpu(
A_log: torch.Tensor,
dt_bias: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
initial_state_source: torch.Tensor,
spec_state_indices: torch.Tensor,
num_accepted_tokens: torch.Tensor,
cu_seqlens: torch.Tensor,
use_qk_l2norm_in_kernel: bool,
softplus_beta: float = 1.0,
softplus_threshold: float = 20.0,
) -> torch.Tensor:
return torch.ops._C.fused_sigmoid_gating_delta_rule_update_spec_cpu(
A_log,
dt_bias,
q,
k,
v,
a,
b,
initial_state_source,
spec_state_indices,
num_accepted_tokens,
cu_seqlens,
use_qk_l2norm_in_kernel,
softplus_beta,
softplus_threshold,
)
def fused_gdn_gating_cpu(
A_log: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
dt_bias: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.ops._C.fused_gdn_gating_cpu(
A_log,
a,
b,
dt_bias,
)
def causal_conv1d_weight_pack(
weight: torch.Tensor,
) -> torch.Tensor:
return torch.ops._C.causal_conv1d_weight_pack(
weight,
)
def causal_conv1d_fwd_cpu(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None,
conv_states: torch.Tensor | None,
query_start_loc: torch.Tensor | None,
cache_indices: torch.Tensor | None,
has_initial_state: torch.Tensor | None,
silu_activation: bool,
is_vnni: bool,
) -> torch.Tensor:
return torch.ops._C.causal_conv1d_fwd_cpu(
x,
weight,
bias,
conv_states,
query_start_loc,
cache_indices,
has_initial_state,
silu_activation,
-1,
is_vnni,
)
def causal_conv1d_update_cpu(
x: torch.Tensor,
conv_states: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None,
silu_activation: bool,
conv_state_indices: torch.Tensor | None,
is_vnni: bool,
num_accepted_tokens: torch.Tensor | None = None,
) -> torch.Tensor:
return torch.ops._C.causal_conv1d_update_cpu(
x,
conv_states,
weight,
bias,
silu_activation,
num_accepted_tokens,
conv_state_indices,
-1,
is_vnni,
)
class CPUDNNLGEMMHandler:
def __init__(self) -> None:
self.handler_tensor: torch.Tensor | None = None
self.n = -1
self.k = -1
self.dtor = torch.ops._C.release_dnnl_matmul_handler
def __del__(self):
if self.handler_tensor is not None:
self.dtor(self.handler_tensor.item())
_supports_onednn = bool(hasattr(torch.ops._C, "create_onednn_mm_handler"))
def is_onednn_acl_supported():
return torch.ops._C.is_onednn_acl_supported()
def create_onednn_mm(
weight: torch.Tensor, # [K, N]
primitive_cache_size: int = 128,
) -> CPUDNNLGEMMHandler:
handler = CPUDNNLGEMMHandler()
handler.k, handler.n = weight.size()
# store the handler pointer in a tensor it doesn't get inlined
handler.handler_tensor = torch.tensor(
torch.ops._C.create_onednn_mm_handler(weight, primitive_cache_size),
dtype=torch.int64,
)
return handler
def onednn_mm(
dnnl_handler: CPUDNNLGEMMHandler,
x: torch.Tensor,
bias: torch.Tensor | None,
) -> torch.Tensor:
output = torch.empty((*x.shape[0:-1], dnnl_handler.n), dtype=x.dtype)
torch.ops._C.onednn_mm(
output, x.reshape(-1, dnnl_handler.k), bias, dnnl_handler.handler_tensor
)
return output
def create_onednn_scaled_mm(
weight: torch.Tensor, # [K, N]
weight_scales: torch.Tensor,
output_type: torch.dtype,
dynamic_quant: bool,
use_azp: bool,
primitive_cache_size: int = 128,
) -> CPUDNNLGEMMHandler:
handler = CPUDNNLGEMMHandler()
handler.k, handler.n = weight.size()
# store the handler pointer in a tensor so it doesn't get inlined
handler.handler_tensor = torch.tensor(
torch.ops._C.create_onednn_scaled_mm_handler(
weight,
weight_scales,
output_type,
dynamic_quant,
use_azp,
primitive_cache_size,
),
dtype=torch.int64,
)
return handler
def onednn_scaled_int8_quant(
input: torch.Tensor,
scale: torch.Tensor | None = None,
azp: torch.Tensor | None = None,
symmetric: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
"""Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
Args:
input: The input tensor to be quantized to int8.
scale: Optional scaling factor for the int8 quantization.
When not provided, we invoke dynamic-per-token quantization.
azp: Optional zero-point for the int8 quantization.
Must be provided for asymmetric quantization if `scale` is provided.
symmetric: Whether to use symmetric quantization (scale only, azp ignored).
Returns:
Output int8 tensor, scales, and optionally azp.
"""
output = torch.empty_like(input, dtype=torch.int8)
token_num = input.numel() // input.shape[-1]
input = input.view((token_num, input.shape[-1]))
if scale is not None:
# static-per-tensor quantization.
assert symmetric == (azp is None), (
"azp must only be provided for asymmetric quantization."
)
torch.ops._C.static_scaled_int8_quant(output, input, scale, azp)
return output, scale, azp
# dynamic-per-token quantization.
input_scales = torch.empty((token_num, 1), device=input.device, dtype=torch.float32)
input_azp = None if symmetric else torch.empty_like(input_scales, dtype=torch.int32)
torch.ops._C.dynamic_scaled_int8_quant(output, input, input_scales, input_azp)
return output, input_scales, input_azp
def onednn_scaled_mm(
dnnl_handler: CPUDNNLGEMMHandler,
x: torch.Tensor,
output: torch.Tensor,
input_scale: torch.Tensor | None,
input_zp: torch.Tensor | None,
input_zp_adj: torch.Tensor | None,
bias: torch.Tensor | None,
) -> torch.Tensor:
torch.ops._C.onednn_scaled_mm(
output,
x,
input_scale,
input_zp,
input_zp_adj,
bias,
dnnl_handler.handler_tensor,
)
return output
def cpu_attn_get_scheduler_metadata(
num_reqs: int,
num_heads: int,
num_kv_heads: int,
head_dim: int,
seq_lens: torch.Tensor,
dtype: torch.dtype,
query_start_loc: torch.Tensor,
causal: bool,
sliding_window_size: int,
isa: str,
enable_kv_split: bool,
dynamic_causal: torch.Tensor | None = None,
kv_cache_dtype: str = "auto",
) -> torch.Tensor:
scheduler_metadata = torch.ops._C.get_scheduler_metadata(
num_reqs,
num_heads,
num_kv_heads,
head_dim,
seq_lens,
dtype,
query_start_loc,
causal,
sliding_window_size,
isa,
enable_kv_split,
dynamic_causal,
kv_cache_dtype,
)
return scheduler_metadata
def cpu_attn_reshape_and_cache(
key: torch.Tensor,
value: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
slot_mapping: torch.Tensor,
isa: str,
k_scale: float = 1.0,
v_scale: float = 1.0,
kv_cache_dtype: str = "auto",
) -> None:
torch.ops._C.cpu_attn_reshape_and_cache(
key,
value,
key_cache,
value_cache,
slot_mapping,
isa,
k_scale,
v_scale,
kv_cache_dtype,
)
def cpu_attention_with_kv_cache(
query: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
output: torch.Tensor,
query_start_loc: torch.Tensor,
seq_lens: torch.Tensor,
scale: float,
causal: bool,
alibi_slopes: torch.Tensor | None,
sliding_window: int,
block_table: torch.Tensor,
softcap: float,
scheduler_metadata: torch.Tensor,
s_aux: torch.Tensor | None,
dynamic_causal: torch.Tensor | None = None,
k_scale: float = 1.0,
v_scale: float = 1.0,
kv_cache_dtype: str = "auto",
) -> None:
torch.ops._C.cpu_attention_with_kv_cache(
query,
key_cache,
value_cache,
output,
query_start_loc,
seq_lens,
scale,
causal,
alibi_slopes,
sliding_window,
block_table,
softcap,
scheduler_metadata,
s_aux,
dynamic_causal,
k_scale,
v_scale,
kv_cache_dtype,
)
def cpu_mla_decode(
query: torch.Tensor,
k_buffer: torch.Tensor,
v_buffer: torch.Tensor,
output: torch.Tensor,
key: torch.Tensor | None,
value: torch.Tensor | None,
loc: torch.Tensor | None,
attn_logits: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
sm_scale: float,
logit_cap: float,
is_cross_attn: bool,
sliding_window_size: int,
encoder_lens: torch.Tensor | None,
sinks: torch.Tensor | None,
) -> None:
torch.ops._C.decode_attention_cpu(
query,
k_buffer,
v_buffer,
output,
key,
value,
loc,
attn_logits,
req_to_token,
req_pool_indices,
seq_lens,
sm_scale,
logit_cap,
is_cross_attn,
sliding_window_size,
encoder_lens,
sinks,
)
def cpu_mla_extend(
q_extend: torch.Tensor,
k_extend: torch.Tensor | None,
v_extend: torch.Tensor | None,
o_extend: torch.Tensor,
k_buffer: torch.Tensor,
v_buffer: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
extend_seq_lens: torch.Tensor,
extend_start_loc: torch.Tensor,
max_len_extend: int,
sm_scale: float,
logit_cap: float,
is_cross_attn: bool,
sliding_window_size: int,
encoder_lens: torch.Tensor | None,
sinks: torch.Tensor | None,
tree_mask: torch.Tensor | None = None,
) -> None:
torch.ops._C.extend_attention_cpu(
q_extend,
k_extend,
v_extend,
o_extend,
k_buffer,
v_buffer,
req_to_token,
req_pool_indices,
seq_lens,
extend_seq_lens,
extend_start_loc,
max_len_extend,
sm_scale,
logit_cap,
is_cross_attn,
sliding_window_size,
encoder_lens,
sinks,
tree_mask,
)
def bmm_cpu(
out: torch.Tensor,
mat1: torch.Tensor,
mat2: torch.Tensor,
is_vnni: bool,
scale: torch.Tensor | None = None,
) -> None:
torch.ops._C.bmm_cpu(out, mat1, mat2, is_vnni, scale)
def amx_mla_concat_and_cache(
kv_c_normed: torch.Tensor,
k_pe: torch.Tensor,
kv_cache: torch.Tensor,
slot_mapping: torch.Tensor,
) -> None:
torch.ops._C.concat_and_cache_mla_cpu(kv_c_normed, k_pe, kv_cache, slot_mapping)
if hasattr(torch.ops._C, "hc_pre_fused_cpu"):
@register_fake("_C::hc_pre_fused_cpu")
def hc_pre_fused_cpu_fake(
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
hc_mult: int,
sinkhorn_iters: int,
rms_eps: float,
hc_eps: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
t, _, d = x.shape
y = x.new_empty((t, d))
post = x.new_empty((t, hc_mult), dtype=torch.float32)
comb = x.new_empty((t, hc_mult, hc_mult), dtype=torch.float32)
return y, post, comb
def hc_pre_fused_cpu(
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
hc_mult: int,
sinkhorn_iters: int,
rms_eps: float,
hc_eps: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return torch.ops._C.hc_pre_fused_cpu(
x, hc_fn, hc_scale, hc_base, hc_mult, sinkhorn_iters, rms_eps, hc_eps
)
if hasattr(torch.ops._C, "hc_post_fused_cpu"):
@register_fake("_C::hc_post_fused_cpu")
def hc_post_fused_cpu_fake(
x: torch.Tensor,
residual: torch.Tensor,
post: torch.Tensor,
comb: torch.Tensor,
) -> torch.Tensor:
t, d = x.shape
hc = residual.size(1)
return x.new_empty((t, hc, d))
def hc_post_fused_cpu(
x: torch.Tensor,
residual: torch.Tensor,
post: torch.Tensor,
comb: torch.Tensor,
) -> torch.Tensor:
return torch.ops._C.hc_post_fused_cpu(x, residual, post, comb)
if hasattr(torch.ops._C, "hc_head_fused_cpu"):
@register_fake("_C::hc_head_fused_cpu")
def hc_head_fused_cpu_fake(
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
hc_eps: float,
norm_eps: float,
) -> torch.Tensor:
t, _, d = x.shape
return x.new_empty((t, d))
def hc_head_fused_cpu(
x: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
hc_eps: float,
norm_eps: float,
) -> torch.Tensor:
return torch.ops._C.hc_head_fused_cpu(x, hc_fn, hc_scale, hc_base, hc_eps, norm_eps)
if hasattr(torch.ops._C, "fused_qnorm_rope_kv_insert_cpu"):
@register_fake("_C::fused_qnorm_rope_kv_insert_cpu")
def fused_qnorm_rope_kv_insert_cpu_fake(
q: torch.Tensor,
kv: torch.Tensor,
positions: torch.Tensor,
swa_kv_cache_2d: torch.Tensor,
slot_mapping: torch.Tensor,
cos_sin_cache: torch.Tensor,
q_head_padded: int,
eps: float,
cache_block_size: int,
) -> torch.Tensor:
num_tokens_full, _, head_dim = q.shape
return q.new_empty((num_tokens_full, q_head_padded, head_dim))
def fused_qnorm_rope_kv_insert_cpu(
q: torch.Tensor,
kv: torch.Tensor,
positions: torch.Tensor,
swa_kv_cache_2d: torch.Tensor,
slot_mapping: torch.Tensor,
cos_sin_cache: torch.Tensor,
q_head_padded: int,
eps: float,
cache_block_size: int,
) -> torch.Tensor:
return torch.ops._C.fused_qnorm_rope_kv_insert_cpu(
q,
kv,
positions,
swa_kv_cache_2d,
slot_mapping,
cos_sin_cache,
q_head_padded,
eps,
cache_block_size,
)
def flash_mla_with_kvcache_cpu(
out: torch.Tensor,
q: torch.Tensor,
window_cache_2d: torch.Tensor,
window_slots: torch.Tensor,
window_block_size: int,
compressed_cache_2d: torch.Tensor,
compressed_slots: torch.Tensor,
compressed_block_size: int,
attn_sink: torch.Tensor,
scale: float,
) -> None:
torch.ops._C.flash_mla_with_kvcache_cpu(
out,
q,
window_cache_2d,
window_slots,
window_block_size,
compressed_cache_2d,
compressed_slots,
compressed_block_size,
attn_sink,
scale,
)
def save_partial_states_cpu(
kv: torch.Tensor,
score: torch.Tensor,
ape: torch.Tensor,
positions: torch.Tensor,
state_cache: torch.Tensor,
slot_mapping: torch.Tensor,
) -> None:
torch.ops._C.save_partial_states_cpu(
kv, score, ape, positions, state_cache, slot_mapping
)
def compress_norm_rope_store_cpu(
state_cache: torch.Tensor,
gather_slots: torch.Tensor,
positions: torch.Tensor,
kv_slot_mapping: torch.Tensor,
rms_norm_weight: torch.Tensor,
rms_norm_eps: float,
cos_sin_cache: torch.Tensor,
kv_cache_2d: torch.Tensor,
kv_cache_block_size: int,
compress_ratio: int,
) -> None:
torch.ops._C.compress_norm_rope_store_cpu(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
rms_norm_eps,
cos_sin_cache,
kv_cache_2d,
kv_cache_block_size,
compress_ratio,
)
def compress_norm_rope_store_indexer_cpu(
state_cache: torch.Tensor,
gather_slots: torch.Tensor,
positions: torch.Tensor,
kv_slot_mapping: torch.Tensor,
rms_norm_weight: torch.Tensor,
rms_norm_eps: float,
cos_sin_cache: torch.Tensor,
kv_cache_2d: torch.Tensor,
kv_cache_block_size: int,
compress_ratio: int,
) -> None:
torch.ops._C.compress_norm_rope_store_indexer_cpu(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
rms_norm_eps,
cos_sin_cache,
kv_cache_2d,
kv_cache_block_size,
compress_ratio,
)
if hasattr(torch.ops._C, "fp8_paged_mqa_logits_cpu"):
@register_fake("_C::fp8_paged_mqa_logits_cpu")
def fp8_paged_mqa_logits_cpu_fake(
q_fp8: torch.Tensor,
kvcache_fp8: torch.Tensor,
weight: torch.Tensor,
seq_lens: torch.Tensor,
page_table: torch.Tensor,
block_size: int,
max_seq_len: int,
) -> torch.Tensor:
batch_size = q_fp8.size(0)
return q_fp8.new_empty((batch_size, max_seq_len), dtype=torch.float32)
def fp8_paged_mqa_logits_cpu(
q_fp8: torch.Tensor,
kvcache_fp8: torch.Tensor,
weight: torch.Tensor,
seq_lens: torch.Tensor,
page_table: torch.Tensor,
block_size: int,
max_seq_len: int,
) -> torch.Tensor:
return torch.ops._C.fp8_paged_mqa_logits_cpu(
q_fp8, kvcache_fp8, weight, seq_lens, page_table, block_size, max_seq_len
)
def topk_transform_512_cpu(
scores: torch.Tensor,
seq_lens: torch.Tensor,
page_tables: torch.Tensor,
out_page_indices: torch.Tensor,
page_size: int,
out_raw_indices: torch.Tensor | None = None,
) -> None:
torch.ops._C.topk_transform_512_cpu(
scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices
)
def fused_indexer_q_rope_quant_cpu(
positions: torch.Tensor,
index_q: torch.Tensor,
index_q_cos_sin_cache: torch.Tensor,
index_q_fp8: torch.Tensor,
index_weights: torch.Tensor,
index_weights_softmax_scale: float,
index_weights_head_scale: float,
index_weights_out: torch.Tensor,
) -> None:
torch.ops._C.fused_indexer_q_rope_quant_cpu(
positions,
index_q,
index_q_cos_sin_cache,
index_q_fp8,
index_weights,
index_weights_softmax_scale,
index_weights_head_scale,
index_weights_out,
)
if hasattr(torch.ops._C, "biased_topk_cpu"):
@register_fake("_C::biased_topk_cpu")
def biased_topk_cpu_fake(
hidden_states: torch.Tensor,
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
topk: int,
renormalize: bool,
scoring_func: str,
num_fused_shared_experts: int,
routed_scaling_factor: float | None,
apply_routed_scaling_factor_on_output: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
num_tokens = hidden_states.size(0)
topk_weights = hidden_states.new_empty((num_tokens, topk), dtype=torch.float32)
topk_ids = hidden_states.new_empty((num_tokens, topk), dtype=torch.int32)
return topk_weights, topk_ids
def biased_topk_cpu(
hidden_states: torch.Tensor,
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
topk: int,
renormalize: bool,
scoring_func: str,
num_fused_shared_experts: int,
routed_scaling_factor: float | None,
apply_routed_scaling_factor_on_output: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.ops._C.biased_topk_cpu(
hidden_states,
gating_output,
correction_bias,
topk,
renormalize,
scoring_func,
num_fused_shared_experts,
routed_scaling_factor,
apply_routed_scaling_factor_on_output,
)
if hasattr(torch.ops._C, "hash_topk_cpu"):
@register_fake("_C::hash_topk_cpu")
def hash_topk_cpu_fake(
gating_output: torch.Tensor,
tid2eid: torch.Tensor,
topk: int,
scoring_func: str,
num_fused_shared_experts: int,
num_experts: int,
routed_scaling_factor: float,
) -> tuple[torch.Tensor, torch.Tensor]:
num_tokens = gating_output.size(0)
topk_weights = gating_output.new_empty((num_tokens, topk), dtype=torch.float32)
topk_ids = gating_output.new_empty((num_tokens, topk), dtype=torch.int32)
return topk_weights, topk_ids
def hash_topk_cpu(
gating_output: torch.Tensor,
tid2eid: torch.Tensor,
topk: int,
scoring_func: str,
num_fused_shared_experts: int,
num_experts: int,
routed_scaling_factor: float,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.ops._C.hash_topk_cpu(
gating_output,
tid2eid,
topk,
scoring_func,
num_fused_shared_experts,
num_experts,
routed_scaling_factor,
)
if hasattr(torch.ops._C, "inverse_gptj_rope_o_proj_cpu"):
@register_fake("_C::inverse_gptj_rope_o_proj_cpu")
def inverse_gptj_rope_o_proj_cpu_fake(
o: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
rope_dim: int,
) -> torch.Tensor:
return o.new_empty(o.shape, dtype=torch.float32)
def inverse_gptj_rope_o_proj_cpu(
o: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
rope_dim: int,
) -> torch.Tensor:
return torch.ops._C.inverse_gptj_rope_o_proj_cpu(
o, positions, cos_sin_cache, rope_dim
)
def cpu_gemm_wna16(
input: torch.Tensor,
q_weight: torch.Tensor,
scales: torch.Tensor,
zeros: torch.Tensor | None,
bias: torch.Tensor | None,
pack_factor: int,
isa_hint: str,
) -> torch.Tensor:
output = torch.empty((input.size(0), scales.size(1)), dtype=input.dtype)
torch.ops._C.cpu_gemm_wna16(
input,
q_weight,
output,
scales,
zeros,
bias,
pack_factor,
isa_hint,
)
return output
def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor:
out = torch.empty_like(input)
torch.ops._C.activation_lut_bf16(out, input, activation)
return out
def cpu_prepack_moe_weight(
weight: torch.Tensor,
isa: str,
) -> torch.Tensor:
output = torch.empty_like(weight)
torch.ops._C.prepack_moe_weight(weight, output, isa)
return output
def cpu_prepack_moe_weight_int8(
weight: torch.Tensor,
isa: str,
) -> torch.Tensor:
output = torch.empty_like(weight)
torch.ops._C.prepack_moe_weight_int8(weight, output, isa)
return output
def cpu_fused_moe(
out: torch.Tensor,
input: torch.Tensor,
w13: torch.Tensor,
w2: torch.Tensor,
w13_bias: torch.Tensor | None,
w2_bias: torch.Tensor | None,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
act: str,
isa: str,
skip_weighted: bool = False,
) -> None:
torch.ops._C.cpu_fused_moe(
out,
input,
w13,
w2,
w13_bias,
w2_bias,
topk_weights,
topk_ids,
skip_weighted,
act,
isa,
)
def cpu_fused_moe_int8(
out: torch.Tensor,
input: torch.Tensor,
w13: torch.Tensor,
w2: torch.Tensor,
w13_scale: torch.Tensor,
w2_scale: torch.Tensor,
w13_bias: torch.Tensor | None,
w2_bias: torch.Tensor | None,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
act: str,
isa: str,
skip_weighted: bool = False,
) -> None:
torch.ops._C.cpu_fused_moe_int8(
out,
input,
w13,
w2,
w13_scale,
w2_scale,
w13_bias,
w2_bias,
topk_weights,
topk_ids,
skip_weighted,
act,
isa,
)
if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
@register_fake("_qutlass_C::matmul_mxf4_bf16_tn")
def _fake_matmul_mxf4_bf16_tn(
a: torch.Tensor,
b: torch.Tensor,
a_sf: torch.Tensor,
b_sf: torch.Tensor,
alpha: torch.Tensor,
):
return a.new_empty(*a.shape[:-1], b.shape[0], dtype=torch.bfloat16)
def matmul_mxf4_bf16_tn(
a: torch.Tensor,
b: torch.Tensor,
a_sf: torch.Tensor,
b_sf: torch.Tensor,
alpha: torch.Tensor,
) -> torch.Tensor:
return torch.ops._qutlass_C.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha)
if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxQuest"):
@register_fake("_qutlass_C::fusedQuantizeMxQuest")
def _fake_fused_quantize_mx_quest(
a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor
):
return xh_e2m1, xh_e8m0
if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxAbsMax"):
@register_fake("_qutlass_C::fusedQuantizeMxAbsMax")
def _fake_fused_quantize_mx_absmax(
a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor
):
return xh_e2m1, xh_e8m0
def fusedQuantizeMx(
a: torch.Tensor, b: torch.Tensor, *, method: Literal["quest", "abs_max"] = "quest"
) -> tuple[torch.Tensor, torch.Tensor]:
if a.dim() == 0:
raise ValueError("`a` must have at least 1 dimension.")
if a.size(-1) % 32 != 0:
raise ValueError(f"last dim of `a` must be divisible by 32, got {a.size(-1)}.")
if b.device != a.device:
raise ValueError("`a` and `b` must be on the same device.")
xh_e2m1 = torch.empty(
*a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device
)
rows, cols = a.numel() // a.size(-1), a.size(-1) // 32
n_row_blocks = cdiv(rows, 128)
n_col_blocks = cdiv(cols, 4)
padded_rows = n_row_blocks * 128
padded_cols = n_col_blocks * 4
xh_e8m0 = torch.empty(
padded_rows, padded_cols, dtype=torch.float8_e8m0fnu, device=a.device
)
if not hasattr(torch.ops, "_qutlass_C"):
raise RuntimeError(
"The `_qutlass_C` extension is not loaded. "
"Make sure your custom op library is imported before calling fusedQuantizeMx."
)
if method == "quest":
return torch.ops._qutlass_C.fusedQuantizeMxQuest(a, b, xh_e2m1, xh_e8m0)
elif method == "abs_max":
return torch.ops._qutlass_C.fusedQuantizeMxAbsMax(a, b, xh_e2m1, xh_e8m0)
else:
raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'")
if hasattr(torch.ops._qutlass_C, "fusedQuantizeNvAbsMax"):
@register_fake("_qutlass_C::fusedQuantizeNvAbsMax")
def _fake_fused_quantize_nv_absmax(
a: torch.Tensor,
b: torch.Tensor,
xh_e2m1: torch.Tensor,
xh_e4m3: torch.Tensor,
global_scale: torch.Tensor,
):
return xh_e2m1, xh_e4m3
def fusedQuantizeNv(
a: torch.Tensor, b: torch.Tensor, global_scale: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
xh_e2m1 = torch.empty(
*a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device
)
rows, cols = a.numel() // a.size(-1), a.size(-1) // 16
n_row_blocks = cdiv(rows, 128)
n_col_blocks = cdiv(cols, 4)
padded_rows = n_row_blocks * 128
padded_cols = n_col_blocks * 4
xh_e4m3 = torch.empty(
padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device
)
safeFusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale)
return xh_e2m1, xh_e4m3
@torch.library.custom_op(
"vllm::safeFusedQuantizeNv", mutates_args=("xh_e2m1", "xh_e4m3")
)
def safeFusedQuantizeNv(
a: torch.Tensor,
b: torch.Tensor,
xh_e2m1: torch.Tensor,
xh_e4m3: torch.Tensor,
global_scale: torch.Tensor,
) -> None:
"""Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place
rather than returning them, to prevent torch 2.12+ errors that outputs of custom
operators may not alias any inputs to the custom operator.
"""
torch.ops._qutlass_C.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale)
return
def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
"""Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832)
kernels. Note that these kernels exploit the recursive properties of
Sylvester Hadamards, and therefore do not require transform weight data
Note that sylvester hadamard transforms are also symmetric, which means that
this function is also applies the (transpose <=> inverse) transform.
Args:
x: value to be transformed inplace
inplace: modify value in place
Returns:
value after transformation
"""
return torch.ops._C.hadacore_transform(x, inplace)
if hasattr(torch.ops._C, "hadacore_transform"):
@register_fake("_C::hadacore_transform")
def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor:
return torch.empty_like(x) if not inplace else x
if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
@register_fake("_C::minimax_allreduce_rms_qk")
def _minimax_allreduce_rms_qk_fake(
qkv: torch.Tensor,
norm_weight_q: torch.Tensor,
norm_weight_k: torch.Tensor,
workspace: torch.Tensor,
q_size: int,
kv_size: int,
rank: int,
nranks: int,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
token_num = qkv.shape[0]
return (
torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device),
torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device),
)
Reference in New Issue View Git Blame Copy Permalink