[Apple Silicon] Add Metal kernel support in sgl-kernel (#23449)
Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
This commit is contained in:
R0CKSTAR
@@ -1,212 +1,224 @@
|
||||
import torch
|
||||
from sgl_kernel.debug_utils import maybe_wrap_debug_kernel
|
||||
from sgl_kernel.load_utils import _load_architecture_specific_ops, _preload_cuda_library
|
||||
import platform
|
||||
import sys
|
||||
|
||||
# Initialize the ops library based on current GPU
|
||||
common_ops = _load_architecture_specific_ops()
|
||||
from sgl_kernel.version import __version__ # noqa: F401
|
||||
|
||||
# Preload the CUDA library to avoid the issue of libcudart.so.12 not found
|
||||
if torch.version.cuda is not None:
|
||||
_preload_cuda_library()
|
||||
|
||||
|
||||
from sgl_kernel.allreduce import *
|
||||
from sgl_kernel.attention import (
|
||||
cutlass_mla_decode,
|
||||
cutlass_mla_get_workspace_size,
|
||||
merge_state_v2,
|
||||
)
|
||||
from sgl_kernel.cutlass_moe import cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_data
|
||||
from sgl_kernel.elementwise import (
|
||||
concat_mla_absorb_q,
|
||||
concat_mla_k,
|
||||
copy_to_gpu_no_ce,
|
||||
fused_add_rmsnorm,
|
||||
gelu_and_mul,
|
||||
gelu_tanh_and_mul,
|
||||
gemma_fused_add_rmsnorm,
|
||||
gemma_rmsnorm,
|
||||
rmsnorm,
|
||||
rotary_embedding,
|
||||
silu_and_mul,
|
||||
)
|
||||
from sgl_kernel.expert_specialization import (
|
||||
es_fp8_blockwise_scaled_grouped_mm,
|
||||
es_sm100_mxfp8_blockscaled_grouped_mm,
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant,
|
||||
)
|
||||
from sgl_kernel.gemm import (
|
||||
awq_dequantize,
|
||||
bmm_fp8,
|
||||
dsv3_fused_a_gemm,
|
||||
dsv3_router_gemm,
|
||||
fp8_blockwise_scaled_mm,
|
||||
fp8_scaled_mm,
|
||||
gptq_gemm,
|
||||
gptq_shuffle,
|
||||
int8_scaled_mm,
|
||||
qserve_w4a8_per_chn_gemm,
|
||||
qserve_w4a8_per_group_gemm,
|
||||
sgl_per_token_group_quant_8bit,
|
||||
sgl_per_token_group_quant_fp8,
|
||||
sgl_per_token_group_quant_int8,
|
||||
sgl_per_token_quant_fp8,
|
||||
shuffle_rows,
|
||||
)
|
||||
from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda
|
||||
from sgl_kernel.kvcacheio import (
|
||||
transfer_kv_all_layer,
|
||||
transfer_kv_all_layer_mla,
|
||||
transfer_kv_per_layer,
|
||||
transfer_kv_per_layer_mla,
|
||||
)
|
||||
from sgl_kernel.mamba import (
|
||||
causal_conv1d_fn_cpu,
|
||||
causal_conv1d_fwd,
|
||||
causal_conv1d_update,
|
||||
causal_conv1d_update_cpu,
|
||||
chunk_gated_delta_rule_cpu,
|
||||
)
|
||||
from sgl_kernel.memory import weak_ref_tensor
|
||||
from sgl_kernel.moe import (
|
||||
apply_shuffle_mul_sum,
|
||||
fp8_blockwise_scaled_grouped_mm,
|
||||
fused_qk_norm_rope,
|
||||
kimi_k2_moe_fused_gate,
|
||||
moe_align_block_size,
|
||||
moe_fused_gate,
|
||||
moe_sum,
|
||||
moe_sum_reduce,
|
||||
prepare_moe_input,
|
||||
topk_sigmoid,
|
||||
topk_softmax,
|
||||
)
|
||||
from sgl_kernel.quantization import (
|
||||
ggml_dequantize,
|
||||
ggml_moe_a8,
|
||||
ggml_moe_a8_vec,
|
||||
ggml_moe_get_block_size,
|
||||
ggml_mul_mat_a8,
|
||||
ggml_mul_mat_vec_a8,
|
||||
)
|
||||
from sgl_kernel.sampling import (
|
||||
top_k_renorm_prob,
|
||||
top_p_renorm_prob,
|
||||
)
|
||||
from sgl_kernel.speculative import (
|
||||
build_tree_kernel_efficient,
|
||||
reconstruct_indices_from_tree_mask,
|
||||
segment_packbits,
|
||||
tree_speculative_sampling_target_only,
|
||||
verify_tree_greedy,
|
||||
)
|
||||
from sgl_kernel.top_k import (
|
||||
fast_topk,
|
||||
fast_topk_transform_fused,
|
||||
fast_topk_transform_ragged_fused,
|
||||
fast_topk_v2,
|
||||
)
|
||||
from sgl_kernel.version import __version__
|
||||
|
||||
if torch.version.hip is not None:
|
||||
from sgl_kernel.elementwise import gelu_quick
|
||||
|
||||
if hasattr(torch.version, "musa") and torch.version.musa is not None:
|
||||
from sgl_kernel.musa import (
|
||||
musa_batched_rotary_embedding_contiguous,
|
||||
musa_fused_gemv,
|
||||
musa_fused_moe_gemv,
|
||||
musa_fused_mul_add,
|
||||
musa_rotary_embedding_contiguous,
|
||||
# On macOS only the Metal extension is shipped; skip CUDA op loading and
|
||||
# re-exports so those symbols are not exposed on Apple Silicon.
|
||||
if sys.platform == "darwin" and platform.machine() == "arm64":
|
||||
from sgl_kernel.metal import *
|
||||
else:
|
||||
import torch
|
||||
from sgl_kernel.debug_utils import maybe_wrap_debug_kernel
|
||||
from sgl_kernel.load_utils import (
|
||||
_load_architecture_specific_ops,
|
||||
_preload_cuda_library,
|
||||
)
|
||||
|
||||
# Initialize the ops library based on current GPU
|
||||
common_ops = _load_architecture_specific_ops()
|
||||
|
||||
_DEBUG_EXPORT_NAMES = [
|
||||
"apply_shuffle_mul_sum",
|
||||
"apply_token_bitmask_inplace_cuda",
|
||||
"awq_dequantize",
|
||||
"bmm_fp8",
|
||||
"build_tree_kernel_efficient",
|
||||
"causal_conv1d_fwd",
|
||||
"causal_conv1d_update",
|
||||
"concat_mla_absorb_q",
|
||||
"concat_mla_k",
|
||||
"copy_to_gpu_no_ce",
|
||||
"cutlass_mla_decode",
|
||||
"cutlass_mla_get_workspace_size",
|
||||
"dsv3_fused_a_gemm",
|
||||
"dsv3_router_gemm",
|
||||
"es_fp8_blockwise_scaled_grouped_mm",
|
||||
"es_sm100_mxfp8_blockscaled_grouped_mm",
|
||||
"es_sm100_mxfp8_blockscaled_grouped_quant",
|
||||
"fast_topk",
|
||||
"fast_topk_transform_fused",
|
||||
"fast_topk_transform_ragged_fused",
|
||||
"fast_topk_v2",
|
||||
"fp8_blockwise_scaled_grouped_mm",
|
||||
"fp8_blockwise_scaled_mm",
|
||||
"fp8_scaled_mm",
|
||||
"fused_add_rmsnorm",
|
||||
"fused_qk_norm_rope",
|
||||
"gelu_and_mul",
|
||||
"gelu_tanh_and_mul",
|
||||
"gemma_fused_add_rmsnorm",
|
||||
"gemma_rmsnorm",
|
||||
"gptq_gemm",
|
||||
"gptq_shuffle",
|
||||
"int8_scaled_mm",
|
||||
"kimi_k2_moe_fused_gate",
|
||||
"merge_state_v2",
|
||||
"moe_align_block_size",
|
||||
"moe_fused_gate",
|
||||
"moe_sum",
|
||||
"moe_sum_reduce",
|
||||
"prepare_moe_input",
|
||||
"qserve_w4a8_per_chn_gemm",
|
||||
"qserve_w4a8_per_group_gemm",
|
||||
"reconstruct_indices_from_tree_mask",
|
||||
"rmsnorm",
|
||||
"rotary_embedding",
|
||||
"segment_packbits",
|
||||
"sgl_per_token_group_quant_8bit",
|
||||
"sgl_per_token_group_quant_fp8",
|
||||
"sgl_per_token_group_quant_int8",
|
||||
"sgl_per_token_quant_fp8",
|
||||
"shuffle_rows",
|
||||
"silu_and_mul",
|
||||
"top_k_renorm_prob",
|
||||
"top_p_renorm_prob",
|
||||
"topk_sigmoid",
|
||||
"topk_softmax",
|
||||
"transfer_kv_all_layer",
|
||||
"transfer_kv_all_layer_mla",
|
||||
"transfer_kv_per_layer",
|
||||
"transfer_kv_per_layer_mla",
|
||||
"tree_speculative_sampling_target_only",
|
||||
"verify_tree_greedy",
|
||||
"weak_ref_tensor",
|
||||
]
|
||||
# Preload the CUDA library to avoid the issue of libcudart.so.12 not found
|
||||
if torch.version.cuda is not None:
|
||||
_preload_cuda_library()
|
||||
|
||||
if torch.version.hip is not None:
|
||||
_DEBUG_EXPORT_NAMES.append("gelu_quick")
|
||||
from sgl_kernel.allreduce import *
|
||||
from sgl_kernel.attention import (
|
||||
cutlass_mla_decode,
|
||||
cutlass_mla_get_workspace_size,
|
||||
merge_state_v2,
|
||||
)
|
||||
from sgl_kernel.cutlass_moe import (
|
||||
cutlass_w4a8_moe_mm,
|
||||
get_cutlass_w4a8_moe_mm_data,
|
||||
)
|
||||
from sgl_kernel.elementwise import (
|
||||
concat_mla_absorb_q,
|
||||
concat_mla_k,
|
||||
copy_to_gpu_no_ce,
|
||||
fused_add_rmsnorm,
|
||||
gelu_and_mul,
|
||||
gelu_tanh_and_mul,
|
||||
gemma_fused_add_rmsnorm,
|
||||
gemma_rmsnorm,
|
||||
rmsnorm,
|
||||
rotary_embedding,
|
||||
silu_and_mul,
|
||||
)
|
||||
from sgl_kernel.expert_specialization import (
|
||||
es_fp8_blockwise_scaled_grouped_mm,
|
||||
es_sm100_mxfp8_blockscaled_grouped_mm,
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant,
|
||||
)
|
||||
from sgl_kernel.gemm import (
|
||||
awq_dequantize,
|
||||
bmm_fp8,
|
||||
dsv3_fused_a_gemm,
|
||||
dsv3_router_gemm,
|
||||
fp8_blockwise_scaled_mm,
|
||||
fp8_scaled_mm,
|
||||
gptq_gemm,
|
||||
gptq_shuffle,
|
||||
int8_scaled_mm,
|
||||
qserve_w4a8_per_chn_gemm,
|
||||
qserve_w4a8_per_group_gemm,
|
||||
sgl_per_token_group_quant_8bit,
|
||||
sgl_per_token_group_quant_fp8,
|
||||
sgl_per_token_group_quant_int8,
|
||||
sgl_per_token_quant_fp8,
|
||||
shuffle_rows,
|
||||
)
|
||||
from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda
|
||||
from sgl_kernel.kvcacheio import (
|
||||
transfer_kv_all_layer,
|
||||
transfer_kv_all_layer_mla,
|
||||
transfer_kv_per_layer,
|
||||
transfer_kv_per_layer_mla,
|
||||
)
|
||||
from sgl_kernel.mamba import (
|
||||
causal_conv1d_fn_cpu,
|
||||
causal_conv1d_fwd,
|
||||
causal_conv1d_update,
|
||||
causal_conv1d_update_cpu,
|
||||
chunk_gated_delta_rule_cpu,
|
||||
)
|
||||
from sgl_kernel.memory import weak_ref_tensor
|
||||
from sgl_kernel.moe import (
|
||||
apply_shuffle_mul_sum,
|
||||
fp8_blockwise_scaled_grouped_mm,
|
||||
fused_qk_norm_rope,
|
||||
kimi_k2_moe_fused_gate,
|
||||
moe_align_block_size,
|
||||
moe_fused_gate,
|
||||
moe_sum,
|
||||
moe_sum_reduce,
|
||||
prepare_moe_input,
|
||||
topk_sigmoid,
|
||||
topk_softmax,
|
||||
)
|
||||
from sgl_kernel.quantization import (
|
||||
ggml_dequantize,
|
||||
ggml_moe_a8,
|
||||
ggml_moe_a8_vec,
|
||||
ggml_moe_get_block_size,
|
||||
ggml_mul_mat_a8,
|
||||
ggml_mul_mat_vec_a8,
|
||||
)
|
||||
from sgl_kernel.sampling import (
|
||||
top_k_renorm_prob,
|
||||
top_p_renorm_prob,
|
||||
)
|
||||
from sgl_kernel.speculative import (
|
||||
build_tree_kernel_efficient,
|
||||
reconstruct_indices_from_tree_mask,
|
||||
segment_packbits,
|
||||
tree_speculative_sampling_target_only,
|
||||
verify_tree_greedy,
|
||||
)
|
||||
from sgl_kernel.top_k import (
|
||||
fast_topk,
|
||||
fast_topk_transform_fused,
|
||||
fast_topk_transform_ragged_fused,
|
||||
fast_topk_v2,
|
||||
)
|
||||
from sgl_kernel.version import __version__
|
||||
|
||||
for _name in _DEBUG_EXPORT_NAMES:
|
||||
if _name in globals():
|
||||
globals()[_name] = maybe_wrap_debug_kernel(
|
||||
globals()[_name], f"sgl_kernel.{_name}"
|
||||
if torch.version.hip is not None:
|
||||
from sgl_kernel.elementwise import gelu_quick
|
||||
|
||||
if hasattr(torch.version, "musa") and torch.version.musa is not None:
|
||||
from sgl_kernel.musa import (
|
||||
musa_batched_rotary_embedding_contiguous,
|
||||
musa_fused_gemv,
|
||||
musa_fused_moe_gemv,
|
||||
musa_fused_mul_add,
|
||||
musa_rotary_embedding_contiguous,
|
||||
)
|
||||
|
||||
del _name
|
||||
del _DEBUG_EXPORT_NAMES
|
||||
_DEBUG_EXPORT_NAMES = [
|
||||
"apply_shuffle_mul_sum",
|
||||
"apply_token_bitmask_inplace_cuda",
|
||||
"awq_dequantize",
|
||||
"bmm_fp8",
|
||||
"build_tree_kernel_efficient",
|
||||
"causal_conv1d_fwd",
|
||||
"causal_conv1d_update",
|
||||
"concat_mla_absorb_q",
|
||||
"concat_mla_k",
|
||||
"copy_to_gpu_no_ce",
|
||||
"cutlass_mla_decode",
|
||||
"cutlass_mla_get_workspace_size",
|
||||
"dsv3_fused_a_gemm",
|
||||
"dsv3_router_gemm",
|
||||
"es_fp8_blockwise_scaled_grouped_mm",
|
||||
"es_sm100_mxfp8_blockscaled_grouped_mm",
|
||||
"es_sm100_mxfp8_blockscaled_grouped_quant",
|
||||
"fast_topk",
|
||||
"fast_topk_transform_fused",
|
||||
"fast_topk_transform_ragged_fused",
|
||||
"fast_topk_v2",
|
||||
"fp8_blockwise_scaled_grouped_mm",
|
||||
"fp8_blockwise_scaled_mm",
|
||||
"fp8_scaled_mm",
|
||||
"fused_add_rmsnorm",
|
||||
"fused_qk_norm_rope",
|
||||
"gelu_and_mul",
|
||||
"gelu_tanh_and_mul",
|
||||
"gemma_fused_add_rmsnorm",
|
||||
"gemma_rmsnorm",
|
||||
"gptq_gemm",
|
||||
"gptq_shuffle",
|
||||
"int8_scaled_mm",
|
||||
"kimi_k2_moe_fused_gate",
|
||||
"merge_state_v2",
|
||||
"moe_align_block_size",
|
||||
"moe_fused_gate",
|
||||
"moe_sum",
|
||||
"moe_sum_reduce",
|
||||
"prepare_moe_input",
|
||||
"qserve_w4a8_per_chn_gemm",
|
||||
"qserve_w4a8_per_group_gemm",
|
||||
"reconstruct_indices_from_tree_mask",
|
||||
"rmsnorm",
|
||||
"rotary_embedding",
|
||||
"segment_packbits",
|
||||
"sgl_per_token_group_quant_8bit",
|
||||
"sgl_per_token_group_quant_fp8",
|
||||
"sgl_per_token_group_quant_int8",
|
||||
"sgl_per_token_quant_fp8",
|
||||
"shuffle_rows",
|
||||
"silu_and_mul",
|
||||
"top_k_renorm_prob",
|
||||
"top_p_renorm_prob",
|
||||
"topk_sigmoid",
|
||||
"topk_softmax",
|
||||
"transfer_kv_all_layer",
|
||||
"transfer_kv_all_layer_mla",
|
||||
"transfer_kv_per_layer",
|
||||
"transfer_kv_per_layer_mla",
|
||||
"tree_speculative_sampling_target_only",
|
||||
"verify_tree_greedy",
|
||||
"weak_ref_tensor",
|
||||
]
|
||||
|
||||
if torch.version.hip is not None:
|
||||
_DEBUG_EXPORT_NAMES.append("gelu_quick")
|
||||
|
||||
def create_greenctx_stream_by_value(*args, **kwargs):
|
||||
from sgl_kernel.spatial import create_greenctx_stream_by_value as _impl
|
||||
for _name in _DEBUG_EXPORT_NAMES:
|
||||
if _name in globals():
|
||||
globals()[_name] = maybe_wrap_debug_kernel(
|
||||
globals()[_name], f"sgl_kernel.{_name}"
|
||||
)
|
||||
|
||||
return _impl(*args, **kwargs)
|
||||
del _name
|
||||
del _DEBUG_EXPORT_NAMES
|
||||
|
||||
def create_greenctx_stream_by_value(*args, **kwargs):
|
||||
from sgl_kernel.spatial import create_greenctx_stream_by_value as _impl
|
||||
|
||||
def get_sm_available(*args, **kwargs):
|
||||
from sgl_kernel.spatial import get_sm_available as _impl
|
||||
return _impl(*args, **kwargs)
|
||||
|
||||
return _impl(*args, **kwargs)
|
||||
def get_sm_available(*args, **kwargs):
|
||||
from sgl_kernel.spatial import get_sm_available as _impl
|
||||
|
||||
return _impl(*args, **kwargs)
|
||||
|
||||
Reference in New Issue
Block a user