Clean up sgl kernel (#12413)

Co-authored-by: Byron Hsu <byronhsu1230@gmail.com>

sgl-kernel/include/sgl_kernel_ops.h

@@ -191,32 +191,6 @@ void fast_topk_transform_ragged_interface(
 void gelu_quick(at::Tensor& out, const at::Tensor& input);
 #endif
 
-/*
- * From gguf quantization
- */
-torch::Tensor
-ggml_dequantize(torch::Tensor W, int64_t type, int64_t m, int64_t n, std::optional<at::ScalarType> const& dtype);
-
-torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W, torch::Tensor X, int64_t type, int64_t row);
-
-torch::Tensor ggml_mul_mat_a8(torch::Tensor W, torch::Tensor X, int64_t type, int64_t row);
-
-torch::Tensor ggml_moe_a8(
-    torch::Tensor X,
-    torch::Tensor W,
-    torch::Tensor sorted_token_ids,
-    torch::Tensor expert_ids,
-    torch::Tensor num_tokens_post_padded,
-    int64_t type,
-    int64_t row,
-    int64_t top_k,
-    int64_t tokens);
-
-torch::Tensor ggml_moe_a8_vec(
-    torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, int64_t top_k, int64_t type, int64_t row, int64_t tokens);
-
-int64_t ggml_moe_get_block_size(int64_t type);
-
 /*
  * From csrc/gemm
  */
@@ -333,7 +307,12 @@ void moe_align_block_size(
     bool pad_sorted_token_ids);
 
 void topk_softmax(
-    torch::Tensor& topk_weights, torch::Tensor& topk_indices, torch::Tensor& gating_output, bool renormalize);
+    torch::Tensor& topk_weights,
+    torch::Tensor& topk_indices,
+    torch::Tensor& gating_output,
+    bool renormalize,
+    double moe_softcapping,
+    const c10::optional<torch::Tensor>& correction_bias);
 
 void moe_sum_reduce(at::Tensor& input, at::Tensor& output, double routed_scaling_factor);
 
@@ -417,6 +396,7 @@ void silu_and_mul_scaled_fp4_experts_quant(
     torch::Tensor const& input_global_scale,
     torch::Tensor const& mask,
     bool use_silu_and_mul);
+
 /*
  * From csrc/moe/cutlass_moe/w4a8
  */
@@ -445,7 +425,9 @@ void cutlass_w4a8_moe_mm(
     torch::Tensor const& s_strides,
     int64_t chunk_size,
     int64_t topk);
-
+/*
+ * From csrc/moe/marlin_moe_wna16
+ */
 torch::Tensor moe_wna16_marlin_gemm(
     torch::Tensor& a,
     std::optional<torch::Tensor> const& c_or_none,
@@ -680,6 +662,11 @@ void transfer_kv_all_layer_direct_lf_pf(
     const at::Tensor& dst_indices,
     int64_t page_size);
 
+/*
+ * From csrc/memory
+ */
+void store_kv_cache(at::Tensor k_cache, at::Tensor v_cache, at::Tensor out_loc, at::Tensor k, at::Tensor v);
+
 /*
  * From FlashInfer
  */
@@ -798,12 +785,12 @@ void convert_vertical_slash_indexes_mergehead(
     bool causal);
 
 /*
- * From XGrammar
+ * From csrc/grammar
  */
 void ApplyTokenBitmaskInplace(at::Tensor logits, at::Tensor bitmask, at::optional<at::Tensor> indices = at::nullopt);
 
 /*
- * From QServe
+ * From csrc/gemm (QServe)
  */
 void qserve_w4a8_per_chn_gemm(
     const torch::Tensor& _in_feats,
@@ -823,16 +810,37 @@ void qserve_w4a8_per_group_gemm(
     const torch::Tensor& _ascales,
     torch::Tensor& _out_feats);
 
+/*
+ * From csrc/quantization/gguf
+ */
+torch::Tensor
+ggml_dequantize(torch::Tensor W, int64_t type, int64_t m, int64_t n, std::optional<at::ScalarType> const& dtype);
+
+torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W, torch::Tensor X, int64_t type, int64_t row);
+
+torch::Tensor ggml_mul_mat_a8(torch::Tensor W, torch::Tensor X, int64_t type, int64_t row);
+
+torch::Tensor ggml_moe_a8(
+    torch::Tensor X,
+    torch::Tensor W,
+    torch::Tensor sorted_token_ids,
+    torch::Tensor expert_ids,
+    torch::Tensor num_tokens_post_padded,
+    int64_t type,
+    int64_t row,
+    int64_t top_k,
+    int64_t tokens);
+
+torch::Tensor ggml_moe_a8_vec(
+    torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, int64_t top_k, int64_t type, int64_t row, int64_t tokens);
+
+int64_t ggml_moe_get_block_size(int64_t type);
+
 /*
  * From csrc/spatial
  */
 std::vector<int64_t> create_greenctx_stream_by_value(int64_t smA, int64_t smB, int64_t device);
 
-/*
- * From csrc/memory
- */
-void store_kv_cache(at::Tensor k_cache, at::Tensor v_cache, at::Tensor out_loc, at::Tensor k, at::Tensor v);
-
 /*
  * From csrc/mamba
  */
@@ -883,7 +891,7 @@ torch::Tensor fast_hadamard_transform_28N(torch::Tensor& x, double scale);
 torch::Tensor fast_hadamard_transform_40N(torch::Tensor& x, double scale);
 
 /*
- * From csrc/fastertransformer
+ * From flashmla
  */
 std::vector<at::Tensor> get_mla_decoding_metadata(
     at::Tensor& seqlens_k,