minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU] remove RECORD_FUNCTION (#23528)

Browse Source
This commit is contained in:
Ma Mingfei
2026-04-24 09:18:29 +08:00
committed by GitHub
parent 87e50f20f6
commit 23e4d381f0
21 changed files with 0 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/activation.cpp
-3
View File
@@ -57,7 +57,6 @@ void act_and_mul_kernel_impl(
// input : {num_tokens, 2 * d}
// output : {num_tokens, d}
at::Tensor silu_and_mul_cpu(at::Tensor& input) {
RECORD_FUNCTION("sgl-kernel::silu_and_mul_cpu", std::vector<c10::IValue>({input}));
auto sizes = input.sizes().vec();
int64_t last_dim = input.ndimension() - 1;
int64_t d = sizes[last_dim] / 2;
@@ -79,7 +78,6 @@ at::Tensor silu_and_mul_cpu(at::Tensor& input) {
}
at::Tensor gelu_tanh_and_mul_cpu(const at::Tensor& input) {
RECORD_FUNCTION("sgl-kernel::gelu_tanh_and_mul_cpu", std::vector<c10::IValue>({input}));
auto sizes = input.sizes().vec();
int64_t last_dim = input.ndimension() - 1;
int64_t d = sizes[last_dim] / 2;
@@ -111,7 +109,6 @@ at::Tensor gelu_tanh_and_mul_cpu(const at::Tensor& input) {
}
at::Tensor gelu_and_mul_cpu(const at::Tensor& input) {
RECORD_FUNCTION("sgl-kernel::gelu_and_mul_cpu", std::vector<c10::IValue>({input}));
auto sizes = input.sizes().vec();
int64_t last_dim = input.ndimension() - 1;
int64_t d = sizes[last_dim] / 2;
sgl-kernel/csrc/cpu/bmm.cpp
-2
View File
@@ -137,8 +137,6 @@ void bmm_kernel_impl(
//
void bmm_cpu(
at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale) {
RECORD_FUNCTION("sgl-kernel::bmm_cpu", std::vector<c10::IValue>({out, mat1, mat2}));
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
// input and out could be non-contiguous
sgl-kernel/csrc/cpu/common.h
-1
View File
@@ -2,7 +2,6 @@
#include <ATen/ATen.h>
#include <ATen/Parallel.h>
#include <ATen/record_function.h>
#if defined(_OPENMP)
#include <omp.h>
sgl-kernel/csrc/cpu/conv3d.cpp
-2
View File
@@ -195,8 +195,6 @@ at::Tensor conv3d_embed_weight_pack(const at::Tensor& weight) {
// conv3d mapped to gemm in embedding
at::Tensor conv3d_embed_cpu(const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias, bool is_vnni) {
RECORD_FUNCTION("sgl_kernel::conv3d_embed_cpu", std::vector<c10::IValue>({input, weight, bias}));
auto packed_w = is_vnni ? weight : conv3d_embed_weight_pack(weight);
CHECK_CONTIGUOUS(input);
sgl-kernel/csrc/cpu/decode.cpp
-5
View File
@@ -1555,11 +1555,6 @@ void decode_attention_cpu(
at::Tensor& seq_lens,
double sm_scale,
double logit_cap) {
RECORD_FUNCTION(
"sgl-kernel::decode_attention_cpu",
std::vector<c10::IValue>(
{query, output, k_buffer, v_buffer, attn_logits, req_to_token, req_pool_indices, seq_lens}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer);
sgl-kernel/csrc/cpu/extend.cpp
-16
View File
@@ -356,22 +356,6 @@ void extend_attention_cpu(
int64_t max_len_extend,
double sm_scale,
double logit_cap) {
RECORD_FUNCTION(
"sgl-kernel::extend_attention_cpu",
std::vector<c10::IValue>(
{q_extend,
k_extend,
v_extend,
o_extend,
k_buffer,
v_buffer,
req_to_token,
req_pool_indices,
seq_lens,
extend_seq_lens,
extend_start_loc,
max_len_extend}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(q_extend);
CHECK_INPUT(o_extend);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend);
sgl-kernel/csrc/cpu/flash_attn.cpp
-4
View File
@@ -441,10 +441,6 @@ at::Tensor flash_attn_varlen_func(
int64_t max_seqlen_q,
int64_t max_seqlen_k,
bool causal) {
RECORD_FUNCTION(
"sgl_kernel::flash_attn_varlen_func",
std::vector<c10::IValue>({q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v);
sgl-kernel/csrc/cpu/gemm.cpp
-4
View File
@@ -729,8 +729,6 @@ at::Tensor convert_scale_packed(at::Tensor& scale) {
//
at::Tensor
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
RECORD_FUNCTION("sgl-kernel::weight_packed_linear", std::vector<c10::IValue>({mat1, mat2, bias}));
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
bool use_fma_gemm = false;
if (packed_w.scalar_type() == at::kFloat) {
@@ -804,8 +802,6 @@ at::Tensor fused_linear_sigmoid_mul(
const std::optional<at::Tensor>& bias,
bool is_vnni,
const at::Tensor& post_mul_mat) {
RECORD_FUNCTION("sgl-kernel::fused_linear_sigmoid_mul", std::vector<c10::IValue>({mat1, mat2, bias, post_mul_mat}));
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
TORCH_CHECK(packed_w.scalar_type() == at::kFloat, "fused_linear_sigmoid_mul requires packed float weight")
sgl-kernel/csrc/cpu/gemm_fp8.cpp
-4
View File
@@ -1035,8 +1035,6 @@ at::Tensor fp8_scaled_mm_cpu(
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni) {
RECORD_FUNCTION("sgl-kernel::fp8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, block_size, bias}));
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
@@ -1104,8 +1102,6 @@ at::Tensor fp8_scaled_mm_cpu(
// scales2: [N, K / G], actual layout: [N / BLOCK_N, K / G, BLOCK_N]
at::Tensor mxfp4_scaled_mm_cpu(
at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2, const std::optional<at::Tensor>& bias, bool is_vnni) {
RECORD_FUNCTION("sgl-kernel::mxfp4_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, bias}));
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
CHECK_INPUT(mat1);
sgl-kernel/csrc/cpu/gemm_int4.cpp
-2
View File
@@ -740,8 +740,6 @@ at::Tensor int4_scaled_mm_cpu_with_quant(
const at::Tensor& weight_qzeros,
const std::optional<at::Tensor>& bias,
at::ScalarType output_dtype) {
RECORD_FUNCTION("sgl-kernel::int4_scaled_mm_cpu_with_quant", std::vector<c10::IValue>({input, weight}));
int64_t M_a = input.size(0);
int64_t K_a = input.size(1);
int64_t lda = input.stride(0);
sgl-kernel/csrc/cpu/gemm_int8.cpp
-6
View File
@@ -380,8 +380,6 @@ INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
RECORD_FUNCTION("sgl-kernel::per_token_quant_int8_cpu", std::vector<c10::IValue>({A}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(A);
CHECK_DIM(2, A);
@@ -427,8 +425,6 @@ at::Tensor int8_scaled_mm_cpu(
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni) {
RECORD_FUNCTION("sgl-kernel::int8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales1, scales2, bias}));
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
CHECK_INPUT(mat1);
@@ -485,8 +481,6 @@ at::Tensor int8_scaled_mm_with_quant(
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni) {
RECORD_FUNCTION("sgl-kernel::int8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, bias}));
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
sgl-kernel/csrc/cpu/interface.cpp
-4
View File
@@ -48,8 +48,6 @@ void initialize(int64_t size, int64_t rank) {
}
void shm_allreduce(torch::Tensor& data, int64_t op) {
RECORD_FUNCTION("sgl-kernel::shm_allreduce", std::vector<c10::IValue>({data}));
TORCH_CHECK(op == c10d::ReduceOp::SUM, "Only torch.distributed.ReduceOp.SUM is supported");
auto numel = data.numel();
@@ -60,8 +58,6 @@ void shm_allreduce(torch::Tensor& data, int64_t op) {
}
torch::Tensor shm_allgather(torch::Tensor& data, int64_t dim) {
RECORD_FUNCTION("sgl-kernel::shm_allgather", std::vector<c10::IValue>({data}));
auto numel = data.numel();
int data_size = numel * data.element_size();
if (dim < 0) {
sgl-kernel/csrc/cpu/mamba/conv.cpp
-4
View File
@@ -550,8 +550,6 @@ at::Tensor causal_conv1d_fwd_cpu(
bool silu_activation,
int64_t pad_slot_id,
bool is_vnni) {
RECORD_FUNCTION("sgl-kernel::causal_conv1d_fwd_cpu", std::vector<c10::IValue>({x, weight, bias}));
CHECK_CONTIGUOUS(weight);
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
@@ -657,8 +655,6 @@ at::Tensor causal_conv1d_update_cpu(
const std::optional<at::Tensor>& conv_state_indices,
int64_t pad_slot_id,
bool is_vnni) {
RECORD_FUNCTION("sgl-kernel::causal_conv1d_update_cpu", std::vector<c10::IValue>({x, weight, bias}));
CHECK_CONTIGUOUS(x);
CHECK_CONTIGUOUS(weight);
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
sgl-kernel/csrc/cpu/mamba/fla.cpp
-8
View File
@@ -1058,9 +1058,6 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
bool head_first,
bool use_qk_l2norm_in_kernel,
double eps = 1e-5) {
RECORD_FUNCTION(
"sgl-kernel::chunk_gated_delta_rule_cpu", std::vector<c10::IValue>({query, key, value, g, beta, initial_state}));
TORCH_CHECK(head_first == false, "chunk_gated_delta_rule_cpu does not support head first");
int64_t B = query.size(0);
int64_t global_seq_len = query.size(1);
@@ -1234,10 +1231,6 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
bool use_qk_l2norm_in_kernel,
double softplus_beta = 1.0,
double softplus_threshold = 20.0) {
RECORD_FUNCTION(
"sgl-kernel::fused_sigmoid_gating_delta_rule_update_cpu",
std::vector<c10::IValue>(
{A_log, dt_bias, q, k, v, a, b, initial_state_source, initial_state_indices, cu_seqlens}));
CHECK_DIM(4, q);
CHECK_DIM(4, v);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);
@@ -1312,7 +1305,6 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
// -A_log.float().exp() * F.softplus(a.float() + dt_bias)
std::tuple<at::Tensor, at::Tensor>
fused_gdn_gating_cpu(const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b, const at::Tensor& dt_bias) {
RECORD_FUNCTION("sgl-kernel::fused_gdn_gating_cpu", std::vector<c10::IValue>({A_log, a, b, dt_bias}));
CHECK_DIM(1, A_log);
CHECK_DIM(2, a);
CHECK_DIM(2, b);
sgl-kernel/csrc/cpu/model/qwen3.cpp
-1
View File
@@ -72,7 +72,6 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_re
int64_t num_heads_v,
int64_t head_qk,
int64_t head_v) {
RECORD_FUNCTION("sgl-kernel::fused_qkvzba_split_reshape_cat_cpu", std::vector<c10::IValue>({mixed_qkvz, mixed_ba}));
CHECK_DIM(2, mixed_qkvz);
CHECK_DIM(2, mixed_ba);
CHECK_INPUT(mixed_qkvz);
sgl-kernel/csrc/cpu/moe.cpp
-5
View File
@@ -852,9 +852,6 @@ at::Tensor fused_experts_cpu(
const std::optional<at::Tensor>& w2_zero,
const std::optional<std::vector<int64_t>> block_size,
bool is_vnni) {
RECORD_FUNCTION(
"sgl-kernel::fused_experts_cpu", std::vector<c10::IValue>({hidden_states, w1, w2, topk_weights, topk_ids}));
auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
@@ -1143,8 +1140,6 @@ at::Tensor shared_expert_cpu(
const std::optional<at::Tensor>& w2_scale,
const std::optional<std::vector<int64_t>> block_size,
bool is_vnni) {
RECORD_FUNCTION("sgl-kernel::shared_expert_cpu", std::vector<c10::IValue>({hidden_states, w1, w2}));
auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
sgl-kernel/csrc/cpu/norm.cpp
-17
View File
@@ -538,8 +538,6 @@ void fused_add_layernorm_kernel_impl(
// input : {batch_size, hidden_size}
at::Tensor l2norm_cpu(at::Tensor& input, double eps) {
RECORD_FUNCTION("sgl-kernel::l2norm_cpu", std::vector<c10::IValue>({input}));
CHECK_INPUT(input);
CHECK_DIM(2, input);
int64_t batch_size = input.size(0);
@@ -565,8 +563,6 @@ at::Tensor l2norm_cpu(at::Tensor& input, double eps) {
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight: {hidden_size}
at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
RECORD_FUNCTION("sgl-kernel::rmsnorm_cpu", std::vector<c10::IValue>({input, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
int64_t inp_dim{input.dim()};
@@ -613,8 +609,6 @@ at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
// bias : {hidden_size}
at::Tensor
layernorm_cpu(const at::Tensor& input, const at::Tensor& weight, const std::optional<at::Tensor>& bias, double eps) {
RECORD_FUNCTION("sgl-kernel::layernorm_cpu", std::vector<c10::IValue>({input, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
int64_t inp_dim{input.dim()};
@@ -657,8 +651,6 @@ layernorm_cpu(const at::Tensor& input, const at::Tensor& weight, const std::opti
}
at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
RECORD_FUNCTION("sgl-kernel::gemma_rmsnorm_cpu", std::vector<c10::IValue>({input, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
CHECK_DIM(2, input);
@@ -694,8 +686,6 @@ at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
// input : {batch_size, hidden_size} or {batch_size, num_head, seq_len, head_dim}
// weight: {hidden_size}
at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
RECORD_FUNCTION("sgl-kernel::gemma3_rmsnorm_cpu", std::vector<c10::IValue>({input, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
TORCH_CHECK(
@@ -761,8 +751,6 @@ at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight: {hidden_size}
at::Tensor gemma4_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps, double scale_shift, bool with_scale) {
RECORD_FUNCTION("sgl-kernel::gemma4_rmsnorm_cpu", std::vector<c10::IValue>({input, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
int64_t inp_dim{input.dim()};
@@ -826,8 +814,6 @@ at::Tensor gemma4_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps,
// weight: {hidden_size}
// gate: {batch_size, hidden_size}
at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Tensor& gate, double eps) {
RECORD_FUNCTION("sgl-kernel::fused_rmsnorm_gated_cpu", std::vector<c10::IValue>({input, weight, gate}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
CHECK_INPUT(gate);
@@ -860,7 +846,6 @@ at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Te
// residual: {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight : {hidden_size}
void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) {
RECORD_FUNCTION("sgl-kernel::fused_add_rmsnorm_cpu", std::vector<c10::IValue>({input, residual, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(residual);
CHECK_INPUT(weight);
@@ -909,7 +894,6 @@ void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor&
// residual: {batch_size, hidden_size}
// weight : {hidden_size}
void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) {
RECORD_FUNCTION("sgl-kernel::gemma_fused_add_rmsnorm_cpu", std::vector<c10::IValue>({input, residual, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(residual);
CHECK_INPUT(weight);
@@ -957,7 +941,6 @@ at::Tensor fused_add_layernorm_cpu(
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
double eps) {
RECORD_FUNCTION("sgl-kernel::fused_add_layernorm_cpu", std::vector<c10::IValue>({input, residual, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(residual);
CHECK_INPUT(weight);
sgl-kernel/csrc/cpu/preprocessor.cpp
-2
View File
@@ -282,8 +282,6 @@ std::tuple<at::Tensor, at::Tensor> image_preprocess_cpu(
int64_t merge_size,
bool disable_grouping,
at::ScalarType out_dtype) {
RECORD_FUNCTION("sgl_kernel::image_preprocess_cpu", std::vector<c10::IValue>({}));
// TODO: lift C++ kernel limitations
TORCH_CHECK(interpolation == "bicubic", "image_preprocess_cpu: support only bicubic mode.");
TORCH_CHECK(do_rescale && do_normalize, "image_preprocess_cpu: support only do_rescale and do_normalize.");
sgl-kernel/csrc/cpu/qkv_proj.cpp
-8
View File
@@ -437,10 +437,6 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope(
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size) {
RECORD_FUNCTION(
"sgl-kernel::qkv_proj_with_rope",
std::vector<c10::IValue>({hidden_states, q_a_proj_weight, q_b_proj_weight, kv_a_proj_weight, w_kc}));
const auto st = hidden_states.scalar_type();
CHECK_INPUT(hidden_states);
CHECK_INPUT(positions);
@@ -649,10 +645,6 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope_fused_weight(
int64_t q_lora_rank,
int64_t kv_lora_rank,
int64_t qk_rope_head_dim) {
RECORD_FUNCTION(
"sgl-kernel::qkv_proj_with_rope_fused_weight",
std::vector<c10::IValue>({hidden_states, qkv_a_proj_weight, q_b_proj_weight, w_kc}));
int64_t hidden_size = hidden_states.size(1);
CHECK_EQ(qkv_a_proj_weight.size(0), q_lora_rank + kv_lora_rank + qk_rope_head_dim);
CHECK_EQ(qkv_a_proj_weight.size(1), get_row_size(hidden_size, use_int8_w8a8));
sgl-kernel/csrc/cpu/rope.cpp
-3
View File
@@ -640,7 +640,6 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
int64_t head_size,
at::Tensor& cos_sin_cache,
bool is_neox) {
RECORD_FUNCTION("sgl-kernel::rotary_embedding_cpu", std::vector<c10::IValue>({query, key}));
CHECK_DIM(1, positions);
const auto input_dim = query.dim();
const auto input_dtype = query.scalar_type();
@@ -775,7 +774,6 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
// sin: [num_tokens, head_size]
std::tuple<at::Tensor, at::Tensor>
apply_rotary_pos_emb_cpu(at::Tensor& query, at::Tensor& key, at::Tensor& cos, at::Tensor& sin) {
RECORD_FUNCTION("sgl-kernel::apply_rotary_pos_emb_cpu", std::vector<c10::IValue>({query, key}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(key);
CHECK_INPUT(cos);
@@ -845,7 +843,6 @@ std::tuple<at::Tensor, at::Tensor> multimodal_rotary_embedding_cpu(
const std::optional<std::vector<int64_t>>& mrope_section,
bool mrope_interleaved,
bool is_neox) {
RECORD_FUNCTION("sgl-kernel::multimodal_rotary_embedding_cpu", std::vector<c10::IValue>({query, key}));
TORCH_CHECK(positions.dim() == 1 || positions.dim() == 2, "positions must be a 1D or 2D tensor");
CHECK_DIM(2, query);
CHECK_DIM(2, key);
sgl-kernel/csrc/cpu/topk.cpp
-5
View File
@@ -449,7 +449,6 @@ void biased_grouped_topk_kernel_impl(
std::tuple<at::Tensor, at::Tensor>
topk_sigmoid_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize) {
RECORD_FUNCTION("sgl-kernel::topk_sigmoid_cpu", std::vector<c10::IValue>({hidden_states, gating_output}));
CHECK_INPUT(gating_output);
const auto st = hidden_states.scalar_type();
@@ -503,7 +502,6 @@ topk_sigmoid_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t t
std::tuple<at::Tensor, at::Tensor>
topk_softmax_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize) {
RECORD_FUNCTION("sgl-kernel::topk_softmax_cpu", std::vector<c10::IValue>({hidden_states, gating_output}));
CHECK_INPUT(gating_output);
const auto st = hidden_states.scalar_type();
@@ -587,7 +585,6 @@ std::tuple<at::Tensor, at::Tensor> grouped_topk_cpu(
"num_token_non_padded must be None default value, got: ",
num_token_non_padded.value());
RECORD_FUNCTION("sgl-kernel::grouped_topk_cpu", std::vector<c10::IValue>({hidden_states, gating_output}));
CHECK_INPUT(gating_output);
const auto st = hidden_states.scalar_type();
@@ -660,8 +657,6 @@ std::tuple<at::Tensor, at::Tensor> biased_grouped_topk_cpu(
!num_token_non_padded.has_value(),
"num_token_non_padded must be None default value, got: ",
num_token_non_padded.value());
RECORD_FUNCTION(
"sgl-kernel::biased_grouped_topk_cpu", std::vector<c10::IValue>({hidden_states, gating_output, correction_bias}));
CHECK_INPUT(gating_output);
CHECK_INPUT(correction_bias);