style: format Rust workspace

crates/xserv-kernels/src/moe.rs

@@ -1,66 +1,113 @@
 use std::ffi::c_void;
 use xserv_tensor::{DType, Tensor};
 
-use crate::gemm::{cublas_handle, CublasHandle};
+use crate::gemm::{CublasHandle, cublas_handle};
 
 unsafe extern "C" {
     fn launch_moe_topk_softmax_bf16(
         router_logits: *const c_void,
-        topk_ids: *mut c_void, topk_weights: *mut c_void,
-        num_tokens: i32, num_experts: i32, top_k: i32,
+        topk_ids: *mut c_void,
+        topk_weights: *mut c_void,
+        num_tokens: i32,
+        num_experts: i32,
+        top_k: i32,
         stream: *mut c_void,
     );
     fn launch_moe_replicate_bf16(
-        x: *const c_void, x_rep: *mut c_void,
-        num_tokens: i32, hidden: i32, local_experts: i32,
+        x: *const c_void,
+        x_rep: *mut c_void,
+        num_tokens: i32,
+        hidden: i32,
+        local_experts: i32,
         stream: *mut c_void,
     );
     fn launch_moe_bias_add_3d_bf16(
-        x: *mut c_void, bias: *const c_void,
-        batch: i32, num_tokens: i32, dim: i32,
+        x: *mut c_void,
+        bias: *const c_void,
+        batch: i32,
+        num_tokens: i32,
+        dim: i32,
         stream: *mut c_void,
     );
     fn launch_moe_weighted_sum_bf16(
         expert_out: *const c_void,
-        topk_ids: *const c_void, topk_weights: *const c_void,
+        topk_ids: *const c_void,
+        topk_weights: *const c_void,
         out: *mut c_void,
-        num_tokens: i32, hidden: i32, top_k: i32,
-        expert_start: i32, local_experts: i32,
+        num_tokens: i32,
+        hidden: i32,
+        top_k: i32,
+        expert_start: i32,
+        local_experts: i32,
         stream: *mut c_void,
     );
 
     fn launch_moe_sparse_gemv_fp8_bf16(
-        x: *const c_void, w: *const c_void, w_scales: *const c_void,
-        bias: *const c_void, topk_ids: *const c_void, y: *mut c_void,
-        num_tokens: i32, n: i32, k: i32, top_k: i32,
-        expert_start: i32, local_experts: i32, x_per_slot: i32,
+        x: *const c_void,
+        w: *const c_void,
+        w_scales: *const c_void,
+        bias: *const c_void,
+        topk_ids: *const c_void,
+        y: *mut c_void,
+        num_tokens: i32,
+        n: i32,
+        k: i32,
+        top_k: i32,
+        expert_start: i32,
+        local_experts: i32,
+        x_per_slot: i32,
         stream: *mut c_void,
     );
     fn launch_moe_sparse_gemv_mxfp4_bf16(
-        x: *const c_void, w_packed: *const c_void, w_scales: *const c_void,
-        bias: *const c_void, topk_ids: *const c_void, y: *mut c_void,
-        num_tokens: i32, n: i32, k: i32, top_k: i32,
-        expert_start: i32, local_experts: i32, x_per_slot: i32,
+        x: *const c_void,
+        w_packed: *const c_void,
+        w_scales: *const c_void,
+        bias: *const c_void,
+        topk_ids: *const c_void,
+        y: *mut c_void,
+        num_tokens: i32,
+        n: i32,
+        k: i32,
+        top_k: i32,
+        expert_start: i32,
+        local_experts: i32,
+        x_per_slot: i32,
         stream: *mut c_void,
     );
     fn launch_moe_weighted_sum_sparse_bf16(
         down: *const c_void,
-        topk_ids: *const c_void, topk_weights: *const c_void,
+        topk_ids: *const c_void,
+        topk_weights: *const c_void,
         out: *mut c_void,
-        num_tokens: i32, hidden: i32, top_k: i32,
-        expert_start: i32, local_experts: i32,
+        num_tokens: i32,
+        hidden: i32,
+        top_k: i32,
+        expert_start: i32,
+        local_experts: i32,
         stream: *mut c_void,
     );
 
     fn cublasGemmStridedBatchedEx(
         handle: CublasHandle,
-        transa: i32, transb: i32,
-        m: i32, n: i32, k: i32,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
         alpha: *const c_void,
-        a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
-        b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
+        a: *const c_void,
+        a_type: i32,
+        lda: i32,
+        stride_a: i64,
+        b: *const c_void,
+        b_type: i32,
+        ldb: i32,
+        stride_b: i64,
         beta: *const c_void,
-        c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
+        c: *mut c_void,
+        c_type: i32,
+        ldc: i32,
+        stride_c: i64,
         batch_count: i32,
         compute_type: i32,
         algo: i32,
@@ -99,7 +146,9 @@ pub fn moe_topk_softmax(
             router_logits.data_ptr() as *const c_void,
             topk_ids.data_ptr() as *mut c_void,
             topk_weights.data_ptr() as *mut c_void,
-            num_tokens as i32, num_experts as i32, top_k as i32,
+            num_tokens as i32,
+            num_experts as i32,
+            top_k as i32,
             xserv_cuda::current_stream_raw(),
         );
     }
@@ -114,13 +163,19 @@ pub fn moe_replicate(x: &Tensor, local_experts: usize) -> Tensor {
     assert!(x.is_contiguous());
     let num_tokens = x.shape()[0];
     let hidden = x.shape()[1];
-    let out = Tensor::empty(&[local_experts, num_tokens, hidden], DType::BF16, x.device());
+    let out = Tensor::empty(
+        &[local_experts, num_tokens, hidden],
+        DType::BF16,
+        x.device(),
+    );
 
     unsafe {
         launch_moe_replicate_bf16(
             x.data_ptr() as *const c_void,
             out.data_ptr() as *mut c_void,
-            num_tokens as i32, hidden as i32, local_experts as i32,
+            num_tokens as i32,
+            hidden as i32,
+            local_experts as i32,
             xserv_cuda::current_stream_raw(),
         );
     }
@@ -143,7 +198,9 @@ pub fn moe_bias_add_3d(x: &Tensor, bias: &Tensor) {
         launch_moe_bias_add_3d_bf16(
             x.data_ptr() as *mut c_void,
             bias.data_ptr() as *const c_void,
-            batch as i32, num_tokens as i32, dim as i32,
+            batch as i32,
+            num_tokens as i32,
+            dim as i32,
             xserv_cuda::current_stream_raw(),
         );
     }
@@ -175,8 +232,11 @@ pub fn moe_weighted_sum(
             topk_ids.data_ptr() as *const c_void,
             topk_weights.data_ptr() as *const c_void,
             out.data_ptr() as *mut c_void,
-            num_tokens as i32, hidden as i32, top_k as i32,
-            expert_start as i32, local_experts as i32,
+            num_tokens as i32,
+            hidden as i32,
+            top_k as i32,
+            expert_start as i32,
+            local_experts as i32,
             xserv_cuda::current_stream_raw(),
         );
     }
@@ -198,9 +258,16 @@ pub fn moe_weighted_sum(
 /// consumer must skip them (see moe_weighted_sum_sparse).
 #[allow(clippy::too_many_arguments)]
 pub fn moe_sparse_gemv_fp8(
-    x: &Tensor, w_fp8_t: &Tensor, w_scales: &Tensor, bias: &Tensor,
-    topk_ids: &Tensor, num_tokens: usize, top_k: usize,
-    expert_start: usize, local_experts: usize, x_per_slot: bool,
+    x: &Tensor,
+    w_fp8_t: &Tensor,
+    w_scales: &Tensor,
+    bias: &Tensor,
+    topk_ids: &Tensor,
+    num_tokens: usize,
+    top_k: usize,
+    expert_start: usize,
+    local_experts: usize,
+    x_per_slot: bool,
 ) -> Tensor {
     assert_eq!(x.dtype(), DType::BF16);
     assert!(x.is_contiguous());
@@ -211,7 +278,14 @@ pub fn moe_sparse_gemv_fp8(
     // silently skip a K%16 tail.
     assert_eq!(k % 16, 0, "sparse FP8 GEMV requires K % 16 == 0, got {k}");
     assert_eq!(x.shape()[x.ndim() - 1], k);
-    assert_eq!(x.shape()[0], if x_per_slot { num_tokens * top_k } else { num_tokens });
+    assert_eq!(
+        x.shape()[0],
+        if x_per_slot {
+            num_tokens * top_k
+        } else {
+            num_tokens
+        }
+    );
 
     let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
     unsafe {
@@ -222,8 +296,13 @@ pub fn moe_sparse_gemv_fp8(
             bias.data_ptr() as *const c_void,
             topk_ids.data_ptr() as *const c_void,
             y.data_ptr() as *mut c_void,
-            num_tokens as i32, n as i32, k as i32, top_k as i32,
-            expert_start as i32, local_experts as i32, x_per_slot as i32,
+            num_tokens as i32,
+            n as i32,
+            k as i32,
+            top_k as i32,
+            expert_start as i32,
+            local_experts as i32,
+            x_per_slot as i32,
             xserv_cuda::current_stream_raw(),
         );
     }
@@ -234,16 +313,32 @@ pub fn moe_sparse_gemv_fp8(
 /// with packed 4-bit weights [E, N, K/2] + UE8M0 block scales [E, N, K/32].
 #[allow(clippy::too_many_arguments)]
 pub fn moe_sparse_gemv_mxfp4(
-    x: &Tensor, w_packed: &Tensor, w_scales: &Tensor, bias: &Tensor,
-    topk_ids: &Tensor, num_tokens: usize, top_k: usize, n: usize, k: usize,
-    expert_start: usize, local_experts: usize, x_per_slot: bool,
+    x: &Tensor,
+    w_packed: &Tensor,
+    w_scales: &Tensor,
+    bias: &Tensor,
+    topk_ids: &Tensor,
+    num_tokens: usize,
+    top_k: usize,
+    n: usize,
+    k: usize,
+    expert_start: usize,
+    local_experts: usize,
+    x_per_slot: bool,
 ) -> Tensor {
     assert_eq!(x.dtype(), DType::BF16);
     assert!(x.is_contiguous());
     // 32-element MXFP4 blocks, read as uint4 (32 nibbles) per lane.
     assert_eq!(k % 32, 0, "sparse MXFP4 GEMV requires K % 32 == 0, got {k}");
     assert_eq!(x.shape()[x.ndim() - 1], k);
-    assert_eq!(x.shape()[0], if x_per_slot { num_tokens * top_k } else { num_tokens });
+    assert_eq!(
+        x.shape()[0],
+        if x_per_slot {
+            num_tokens * top_k
+        } else {
+            num_tokens
+        }
+    );
 
     let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
     unsafe {
@@ -254,8 +349,13 @@ pub fn moe_sparse_gemv_mxfp4(
             bias.data_ptr() as *const c_void,
             topk_ids.data_ptr() as *const c_void,
             y.data_ptr() as *mut c_void,
-            num_tokens as i32, n as i32, k as i32, top_k as i32,
-            expert_start as i32, local_experts as i32, x_per_slot as i32,
+            num_tokens as i32,
+            n as i32,
+            k as i32,
+            top_k as i32,
+            expert_start as i32,
+            local_experts as i32,
+            x_per_slot as i32,
             xserv_cuda::current_stream_raw(),
         );
     }
@@ -286,8 +386,11 @@ pub fn moe_weighted_sum_sparse(
             topk_ids.data_ptr() as *const c_void,
             topk_weights.data_ptr() as *const c_void,
             out.data_ptr() as *mut c_void,
-            num_tokens as i32, hidden as i32, top_k as i32,
-            expert_start as i32, local_experts as i32,
+            num_tokens as i32,
+            hidden as i32,
+            top_k as i32,
+            expert_start as i32,
+            local_experts as i32,
             xserv_cuda::current_stream_raw(),
         );
     }
@@ -341,13 +444,25 @@ pub fn batched_gemm_strided(a: &Tensor, b: &Tensor) -> Tensor {
         cublasSetStream_v2(handle, xserv_cuda::current_stream_raw());
         let status = cublasGemmStridedBatchedEx(
             handle,
-            0, 0, // CUBLAS_OP_N, CUBLAS_OP_N
-            n as i32, m as i32, k as i32,
+            0,
+            0, // CUBLAS_OP_N, CUBLAS_OP_N
+            n as i32,
+            m as i32,
+            k as i32,
             &alpha as *const f32 as *const c_void,
-            b.data_ptr() as *const c_void, CUDA_R_16BF, n as i32, stride_b,
-            a.data_ptr() as *const c_void, CUDA_R_16BF, k as i32, stride_a,
+            b.data_ptr() as *const c_void,
+            CUDA_R_16BF,
+            n as i32,
+            stride_b,
+            a.data_ptr() as *const c_void,
+            CUDA_R_16BF,
+            k as i32,
+            stride_a,
             &beta as *const f32 as *const c_void,
-            c.data_ptr() as *mut c_void, CUDA_R_16BF, n as i32, stride_c,
+            c.data_ptr() as *mut c_void,
+            CUDA_R_16BF,
+            n as i32,
+            stride_c,
             batch as i32,
             CUBLAS_COMPUTE_32F,
             CUBLAS_GEMM_DEFAULT,