wehub-resource-sync
3811 lines
110 KiB
Python
3811 lines
110 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from enum import IntEnum
|
|
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)
|
|
# Must be zero-initialized: the swizzled scale buffer is padded to
|
|
# (round_up(m, 128), round_up(scale_n, 4) // 4) but the NVFP4 quant
|
|
# kernel does not write every padded element that the downstream
|
|
# NVFP4 GEMM reads. torch.empty leaves those padded scale factors
|
|
# uninitialized, which corrupts dequantization and causes a severe
|
|
# Blackwell NVFP4 decode throughput/output-length regression.
|
|
return torch.zeros(
|
|
(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)
|
|
|
|
@register_fake("_C::scaled_fp4_quant.out")
|
|
def _scaled_fp4_quant_out_fake(
|
|
input: torch.Tensor,
|
|
input_scale: torch.Tensor,
|
|
is_sf_swizzled_layout: bool,
|
|
*,
|
|
output: torch.Tensor,
|
|
output_scale: torch.Tensor,
|
|
) -> None:
|
|
return None
|
|
|
|
|
|
# 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)
|
|
|
|
|
|
# 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)
|
|
|
|
|
|
# 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,
|
|
b_g_idx: 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,
|
|
b_g_idx,
|
|
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,
|
|
b_g_idx: 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, q_perm: torch.Tensor, bit: int) -> None:
|
|
torch.ops._C.gptq_shuffle(q_weight, q_perm, bit)
|
|
|
|
|
|
def gptq_gemm_rdna3(
|
|
a: torch.Tensor,
|
|
b_q_weight: torch.Tensor,
|
|
b_qzeros: torch.Tensor,
|
|
b_scales: torch.Tensor,
|
|
b_g_idx: torch.Tensor,
|
|
use_v2_format: bool,
|
|
) -> torch.Tensor:
|
|
return torch.ops._rocm_C.gptq_gemm_rdna3(
|
|
a, b_q_weight, b_qzeros, b_scales, b_g_idx, 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,
|
|
b_g_idx: 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,
|
|
b_g_idx: 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, "_rocm_C") and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3"):
|
|
|
|
@register_fake("_rocm_C::moe_gptq_gemm_rdna3")
|
|
def _moe_gptq_gemm_rdna3_fake(
|
|
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:
|
|
return
|
|
|
|
|
|
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 FusedMoE 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,
|
|
perm: 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, perm, 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,
|
|
perm: 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,
|
|
perm: 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], perm[e], size_k, size_n, num_bits, is_a_8bit
|
|
)
|
|
return output
|
|
|
|
|
|
def awq_marlin_moe_repack(
|
|
b_q_weight: torch.Tensor,
|
|
perm: 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,
|
|
g_idx: torch.Tensor | None,
|
|
perm: torch.Tensor | None,
|
|
workspace: torch.Tensor,
|
|
b_q_type: ScalarType,
|
|
size_m: int,
|
|
size_n: int,
|
|
size_k: int,
|
|
is_k_full: bool = True,
|
|
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,
|
|
g_idx,
|
|
perm,
|
|
workspace,
|
|
b_q_type.id,
|
|
size_m,
|
|
size_n,
|
|
size_k,
|
|
is_k_full,
|
|
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,
|
|
g_idx: torch.Tensor | None,
|
|
perm: torch.Tensor | None,
|
|
workspace: torch.Tensor,
|
|
b_q_type_id: int,
|
|
size_m: torch.SymInt,
|
|
size_n: torch.SymInt,
|
|
size_k: torch.SymInt,
|
|
is_k_full: bool = True,
|
|
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,
|
|
):
|
|
"""
|
|
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/b/c/group_scale_strides:
|
|
Strides describing the memory layout of the input tensors.
|
|
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)
|
|
|
|
|
|
def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
|
|
return torch.ops._C.permute_cols(a, perm)
|
|
|
|
|
|
if hasattr(torch.ops._C, "permute_cols"):
|
|
|
|
@register_fake("_C::permute_cols")
|
|
def _permute_cols_fake(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
|
|
return torch.empty_like(a)
|
|
|
|
|
|
# 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.
|
|
use_8x4_sf_layout: Whether to use the 8x4 or 128x4 layout for the scaling
|
|
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
|
|
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.
|
|
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:
|
|
tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp.
|
|
"""
|
|
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,
|
|
)
|
|
|
|
|
|
# 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 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):
|
|
torch.ops._moe_C.moe_sum(input, output)
|
|
|
|
|
|
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,
|
|
) -> None:
|
|
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,
|
|
)
|
|
|
|
|
|
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 dsv3_router_gemm(
|
|
hidden_states: torch.Tensor,
|
|
router_weight: torch.Tensor,
|
|
output_dtype: torch.dtype,
|
|
) -> torch.Tensor:
|
|
output = torch.empty(
|
|
hidden_states.shape[0],
|
|
router_weight.shape[0],
|
|
device=hidden_states.device,
|
|
dtype=output_dtype,
|
|
)
|
|
torch.ops._moe_C.dsv3_router_gemm(output, hidden_states, router_weight)
|
|
return output
|
|
|
|
|
|
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
|
|
|
|
|
|
if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):
|
|
|
|
@register_fake("_C::fp32_router_gemm")
|
|
def fp32_router_gemm_fake(
|
|
output: torch.Tensor,
|
|
mat_a: torch.Tensor,
|
|
mat_b: torch.Tensor,
|
|
) -> None:
|
|
return
|
|
|
|
|
|
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,
|
|
) -> None:
|
|
torch.ops._moe_C.topk_softmax(
|
|
topk_weights,
|
|
topk_ids,
|
|
token_expert_indices,
|
|
gating_output,
|
|
renormalize,
|
|
e_score_correction_bias,
|
|
)
|
|
|
|
|
|
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,
|
|
) -> None:
|
|
torch.ops._moe_C.topk_sigmoid(
|
|
topk_weights,
|
|
topk_ids,
|
|
token_expert_indices,
|
|
gating_output,
|
|
renormalize,
|
|
e_score_correction_bias,
|
|
routed_scaling_factor,
|
|
)
|
|
|
|
|
|
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,
|
|
) -> 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,
|
|
)
|
|
|
|
|
|
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():
|
|
raise NotImplementedError(
|
|
"The fused grouped_topk kernel is only available on CUDA platforms"
|
|
)
|
|
return torch.ops._moe_C.grouped_topk(
|
|
scores,
|
|
num_expert_group,
|
|
topk_group,
|
|
topk,
|
|
renormalize,
|
|
routed_scaling_factor,
|
|
bias,
|
|
scoring_func,
|
|
)
|
|
|
|
|
|
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,
|
|
g_idx: torch.Tensor | None,
|
|
perm: 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,
|
|
is_k_full: bool,
|
|
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,
|
|
g_idx,
|
|
perm,
|
|
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,
|
|
is_k_full,
|
|
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,
|
|
g_idx: torch.Tensor | None,
|
|
perm: 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,
|
|
is_k_full: bool,
|
|
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,
|
|
) -> 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.
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
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_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,
|
|
seq_lens: torch.Tensor,
|
|
workspace_starts: torch.Tensor,
|
|
batch_size: int,
|
|
) -> 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]
|
|
seq_lens: Sequence lengths [num_reqs]
|
|
workspace_starts: Workspace start offsets [num_reqs]
|
|
batch_size: Number of requests
|
|
"""
|
|
torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache(
|
|
src_cache, dst, block_table, seq_lens, 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 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,
|
|
) -> None:
|
|
"""DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
|
|
|
|
Computes output = mat_a @ mat_b.T where:
|
|
mat_a: [num_tokens, 7168] row-major bf16 (hidden states)
|
|
mat_b: [7168, 2112] column-major bf16 (weight transposed)
|
|
output: [num_tokens, 2112] row-major bf16
|
|
|
|
Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes.
|
|
Requires SM 9.0+ (Hopper).
|
|
"""
|
|
torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b)
|
|
|
|
|
|
if hasattr(torch.ops._C, "weight_packed_linear"):
|
|
|
|
@register_fake("_C::weight_packed_linear")
|
|
def weight_packed_linear_fake(
|
|
mat1: torch.Tensor,
|
|
mat2: torch.Tensor,
|
|
bias: torch.Tensor | None,
|
|
is_vnni: bool,
|
|
) -> torch.Tensor:
|
|
return torch.empty(
|
|
(mat1.size(0), mat2.size(0)), dtype=mat1.dtype, device=mat2.device
|
|
)
|
|
|
|
|
|
class CPUQuantMethod(IntEnum):
|
|
UNQUANT = 0
|
|
INT8_W8A8 = 1
|
|
FP8_W8A16 = 2
|
|
INT4_W4A8 = 3
|
|
MXFP4 = 4
|
|
|
|
|
|
if hasattr(torch.ops._C, "fused_experts_cpu"):
|
|
|
|
@register_fake("_C::fused_experts_cpu")
|
|
def fused_experts_cpu_fake(
|
|
hidden_states: torch.Tensor,
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
inplace: bool,
|
|
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,
|
|
w2_bias: torch.Tensor | None,
|
|
alpha: float | None,
|
|
limit: float | None,
|
|
is_vnni: bool,
|
|
) -> torch.Tensor:
|
|
return torch.empty_like(hidden_states)
|
|
|
|
|
|
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(
|
|
hidden_states: torch.Tensor,
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
inplace: bool,
|
|
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,
|
|
) -> torch.Tensor:
|
|
return torch.ops._C.fused_experts_cpu(
|
|
hidden_states,
|
|
w1,
|
|
w2,
|
|
topk_weights,
|
|
topk_ids,
|
|
inplace,
|
|
moe_comp_method,
|
|
w1_scale,
|
|
w2_scale,
|
|
w1_zero,
|
|
w2_zero,
|
|
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
|
|
)
|
|
|
|
|
|
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,
|
|
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,
|
|
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,
|
|
) -> torch.Tensor:
|
|
return torch.ops._C.causal_conv1d_update_cpu(
|
|
x,
|
|
conv_states,
|
|
weight,
|
|
bias,
|
|
silu_activation,
|
|
None,
|
|
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,
|
|
):
|
|
"""
|
|
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:
|
|
tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : 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,
|
|
) -> 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,
|
|
)
|
|
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_gemm_wna16(
|
|
input: torch.Tensor,
|
|
q_weight: torch.Tensor,
|
|
scales: torch.Tensor,
|
|
zeros: torch.Tensor | None,
|
|
g_idx: 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,
|
|
g_idx,
|
|
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_fused_moe(
|
|
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,
|
|
) -> torch.Tensor:
|
|
output = torch.empty_like(input)
|
|
torch.ops._C.cpu_fused_moe(
|
|
output,
|
|
input,
|
|
w13,
|
|
w2,
|
|
w13_bias,
|
|
w2_bias,
|
|
topk_weights,
|
|
topk_ids,
|
|
skip_weighted,
|
|
act,
|
|
isa,
|
|
)
|
|
return output
|
|
|
|
|
|
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, "fusedQuantizeNv"):
|
|
|
|
@register_fake("_qutlass_C::fusedQuantizeNv")
|
|
def _fake_fused_quantize_nv(
|
|
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
|
|
)
|
|
|
|
return torch.ops._qutlass_C.fusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale)
|
|
|
|
|
|
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),
|
|
)
|