refactor(kernels): the Marlin face serves both consumers from one module

pegainfer-kernels/src/ops.rs

@@ -19,6 +19,8 @@ mod k3_tilelang;
 mod kimi_k2;
 mod linear;
 mod lora;
+#[cfg(any(feature = "gemma4", feature = "kimi-k2"))]
+mod marlin_face;
 mod norm;
 mod sampling;
 

pegainfer-kernels/src/ops/gemma4.rs

@@ -6,6 +6,10 @@ use cudarc::driver::DevicePtr;
 use cudarc::driver::DevicePtrMut;
 
 use crate::ffi;
+use crate::ops::marlin_face::MarlinAlignBuffers;
+use crate::ops::marlin_face::MarlinGemmBuffers;
+use crate::ops::marlin_face::launch_marlin_align;
+use crate::ops::marlin_face::launch_marlin_gemm;
 use crate::tensor::DeviceContext;
 use crate::tensor::DeviceVec;
 use crate::tensor::HiddenStates;
@@ -159,44 +163,48 @@ pub fn gemma4_marlin_nvfp4_moe(
         out.hidden_dim,
         rows * dispatch.top_k
     );
-    let (input_ptr, _input_guard) = input.data.device_ptr(&ctx.stream);
-    let (qweight_ptr, _qweight_guard) = qweight.device_ptr(&ctx.stream);
-    let (scales_ptr, _scales_guard) = scales.device_ptr(&ctx.stream);
     let (global_ptr, _global_guard) = global_scale.device_ptr(&ctx.stream);
-    let (sorted_ptr, _sorted_guard) = dispatch.sorted_token_ids.device_ptr(&ctx.stream);
-    let (expert_ptr, _expert_guard) = dispatch.expert_ids.device_ptr(&ctx.stream);
-    let (padded_ptr, _padded_guard) = dispatch.num_tokens_post_padded.device_ptr(&ctx.stream);
-    let (weights_ptr, _weights_guard) = dispatch.topk_weights.device_ptr(&ctx.stream);
     let workspace_len = workspace.len();
     let sorted_len = dispatch.sorted_token_ids.len();
-    let (workspace_ptr, _workspace_guard) = workspace.device_ptr_mut(&ctx.stream);
-    let (c_tmp_ptr, _c_tmp_guard) = c_tmp.device_ptr_mut(&ctx.stream);
-    let (out_ptr, _out_guard) = out.data.device_ptr_mut(&ctx.stream);
-    let result = unsafe {
-        ffi::gemma4_marlin_nvfp4_moe_cuda(
-            input_ptr as *const ffi::Half,
-            out_ptr as *mut ffi::Half,
-            c_tmp_ptr as *mut f32,
-            qweight_ptr as *const u8,
-            scales_ptr as *const u8,
-            global_ptr as *const f32,
-            workspace_ptr as *mut i32,
-            sorted_ptr as *const i32,
-            expert_ptr as *const i32,
-            padded_ptr as *const i32,
-            weights_ptr as *const f32,
-            i32::try_from(workspace_len)?,
-            i32::try_from(sorted_len)?,
-            i32::try_from(dispatch.block_size)?,
-            i32::try_from(dispatch.top_k)?,
-            dispatch.mul_topk_weights,
-            i32::try_from(rows)?,
-            i32::try_from(size_n)?,
-            i32::try_from(size_k)?,
-            0,
-            crate::tensor::active_cu_stream(ctx),
-        )
+    let buffers = MarlinGemmBuffers {
+        input: &input.data,
+        output: &mut out.data,
+        c_tmp,
+        qweight,
+        scales,
+        workspace,
+        sorted_token_ids: dispatch.sorted_token_ids,
+        expert_ids: dispatch.expert_ids,
+        num_tokens_post_padded: dispatch.num_tokens_post_padded,
+        topk_weights: dispatch.topk_weights,
     };
+    let result = launch_marlin_gemm(ctx, buffers, |ptrs| {
+        Ok(unsafe {
+            ffi::gemma4_marlin_nvfp4_moe_cuda(
+                ptrs.input,
+                ptrs.output,
+                ptrs.c_tmp,
+                ptrs.qweight,
+                ptrs.scales,
+                global_ptr as *const f32,
+                ptrs.workspace,
+                ptrs.sorted_token_ids,
+                ptrs.expert_ids,
+                ptrs.num_tokens_post_padded,
+                ptrs.topk_weights,
+                i32::try_from(workspace_len)?,
+                i32::try_from(sorted_len)?,
+                i32::try_from(dispatch.block_size)?,
+                i32::try_from(dispatch.top_k)?,
+                dispatch.mul_topk_weights,
+                i32::try_from(rows)?,
+                i32::try_from(size_n)?,
+                i32::try_from(size_k)?,
+                0,
+                crate::tensor::active_cu_stream(ctx),
+            )
+        })
+    })?;
     result.result()?;
     Ok(())
 }
@@ -309,29 +317,37 @@ pub fn marlin_moe_align_block_size(
         "marlin_moe_align_block_size: {max_padded} padded slots cannot hold {routes} routes with \
          one part filled block per expert"
     );
-    let (idx_ptr, _idx_guard) = topk_idx.device_ptr(&ctx.stream);
-    let (sorted_ptr, _sorted_guard) = out.sorted_token_ids.device_ptr_mut(&ctx.stream);
-    let (expert_ptr, _expert_guard) = out.expert_ids.device_ptr_mut(&ctx.stream);
-    let (padded_ptr, _padded_guard) = out.num_tokens_post_padded.device_ptr_mut(&ctx.stream);
-    let (offsets_ptr, _offsets_guard) = out.expert_offsets.device_ptr_mut(&ctx.stream);
-    let result = unsafe {
-        ffi::marlin_moe_align_block_size_cuda(
-            idx_ptr as *const i32,
-            sorted_ptr as *mut i32,
-            expert_ptr as *mut i32,
-            padded_ptr as *mut i32,
-            offsets_ptr as *mut u32,
-            std::ptr::null_mut(),
-            i32::try_from(rows)?,
-            i32::try_from(top_k)?,
-            0,
-            i32::try_from(experts)?,
-            i32::try_from(block_size)?,
-            i32::try_from(max_padded)?,
-            i32::try_from(max_blocks)?,
-            crate::tensor::active_cu_stream(ctx),
-        )
+    let buffers = MarlinAlignBuffers {
+        topk_idx,
+        sorted_token_ids: out.sorted_token_ids,
+        expert_ids: out.expert_ids,
+        num_tokens_post_padded: out.num_tokens_post_padded,
+        expert_offsets: out.expert_offsets,
     };
+    let result = launch_marlin_align(
+        ctx,
+        buffers,
+        |idx_ptr, sorted_ptr, expert_ptr, padded_ptr, offsets_ptr| {
+            Ok(unsafe {
+                ffi::marlin_moe_align_block_size_cuda(
+                    idx_ptr,
+                    sorted_ptr,
+                    expert_ptr,
+                    padded_ptr,
+                    offsets_ptr,
+                    std::ptr::null_mut(),
+                    i32::try_from(rows)?,
+                    i32::try_from(top_k)?,
+                    0,
+                    i32::try_from(experts)?,
+                    i32::try_from(block_size)?,
+                    i32::try_from(max_padded)?,
+                    i32::try_from(max_blocks)?,
+                    crate::tensor::active_cu_stream(ctx),
+                )
+            })
+        },
+    )?;
     result.result()?;
     Ok(())
 }

pegainfer-kernels/src/ops/kimi_k2/experts.rs

@@ -6,6 +6,10 @@ use cudarc::driver::DevicePtrMut;
 use half::bf16;
 
 use crate::ffi;
+use crate::ops::marlin_face::MarlinAlignBuffers;
+use crate::ops::marlin_face::MarlinGemmBuffers;
+use crate::ops::marlin_face::launch_marlin_align;
+use crate::ops::marlin_face::launch_marlin_gemm;
 #[cfg(test)]
 use crate::tensor::AxisSpec;
 use crate::tensor::DeviceContext;
@@ -749,30 +753,37 @@ pub fn kimi_moe_marlin_align_block_size<'a>(
     );
 
     {
-        let (topk_ptr, _topk_guard) = topk_idx.device_ptr(&ctx.stream);
-        let (sorted_ptr, _sorted_guard) = workspace.sorted_token_ids.device_ptr_mut(&ctx.stream);
-        let (expert_ids_ptr, _expert_ids_guard) = workspace.expert_ids.device_ptr_mut(&ctx.stream);
-        let (num_tokens_ptr, _num_tokens_guard) =
-            workspace.num_tokens_post_padded.device_ptr_mut(&ctx.stream);
-        let (offsets_ptr, _offsets_guard) = workspace.expert_offsets.device_ptr_mut(&ctx.stream);
-        let result = unsafe {
-            ffi::kimi_moe_marlin_align_block_size_cuda(
-                topk_ptr as *const i32,
-                sorted_ptr as *mut i32,
-                expert_ids_ptr as *mut i32,
-                num_tokens_ptr as *mut i32,
-                offsets_ptr as *mut u32,
-                std::ptr::null_mut(),
-                active_tokens as i32,
-                KIMI_K2_TOPK as i32,
-                global_expert_start as i32,
-                KIMI_K2_LOCAL_EXPERTS as i32,
-                workspace.block_size as i32,
-                workspace.max_padded_tokens as i32,
-                workspace.max_m_blocks as i32,
-                ctx.stream.cu_stream(),
-            )
+        let buffers = MarlinAlignBuffers {
+            topk_idx,
+            sorted_token_ids: &mut workspace.sorted_token_ids,
+            expert_ids: &mut workspace.expert_ids,
+            num_tokens_post_padded: &mut workspace.num_tokens_post_padded,
+            expert_offsets: &mut workspace.expert_offsets,
         };
+        let result = launch_marlin_align(
+            ctx,
+            buffers,
+            |topk_ptr, sorted_ptr, expert_ids_ptr, num_tokens_ptr, offsets_ptr| {
+                Ok(unsafe {
+                    ffi::kimi_moe_marlin_align_block_size_cuda(
+                        topk_ptr,
+                        sorted_ptr,
+                        expert_ids_ptr,
+                        num_tokens_ptr,
+                        offsets_ptr,
+                        std::ptr::null_mut(),
+                        active_tokens as i32,
+                        KIMI_K2_TOPK as i32,
+                        global_expert_start as i32,
+                        KIMI_K2_LOCAL_EXPERTS as i32,
+                        workspace.block_size as i32,
+                        workspace.max_padded_tokens as i32,
+                        workspace.max_m_blocks as i32,
+                        ctx.stream.cu_stream(),
+                    )
+                })
+            },
+        )?;
         result.result()?;
     }
     Ok(KimiMarlinRouting {
@@ -1369,43 +1380,46 @@ fn launch_marlin_wna16_gemm(
         "Kimi Marlin WNA16 weight package must be non-empty"
     );
     let lock_len = workspace.locks.len();
-    let (input_ptr, _input_guard) = input.device_ptr(&ctx.stream);
-    let (output_ptr, _output_guard) = output.device_ptr_mut(&ctx.stream);
-    let (c_tmp_ptr, _c_tmp_guard) = workspace.c_tmp.device_ptr_mut(&ctx.stream);
-    let (weight_ptr, _weight_guard) = weight_packed_uint4b8.device_ptr(&ctx.stream);
-    let (scale_ptr, _scale_guard) = weight_scale_permuted.device_ptr(&ctx.stream);
-    let (locks_ptr, _locks_guard) = workspace.locks.device_ptr_mut(&ctx.stream);
-    let (sorted_ptr, _sorted_guard) = routing.sorted_token_ids.device_ptr(&ctx.stream);
-    let (expert_ids_ptr, _expert_ids_guard) = routing.expert_ids.device_ptr(&ctx.stream);
-    let (num_tokens_ptr, _num_tokens_guard) =
-        routing.num_tokens_post_padded.device_ptr(&ctx.stream);
-    let (topk_ptr, _topk_guard) = topk_weight.device_ptr(&ctx.stream);
-    let result = unsafe {
-        ffi::kimi_marlin_wna16_gemm_cuda(
-            input_ptr as *const ffi::Half,
-            output_ptr as *mut ffi::Half,
-            c_tmp_ptr as *mut f32,
-            weight_ptr as *const u8,
-            scale_ptr as *const ffi::Half,
-            locks_ptr as *mut i32,
-            sorted_ptr as *const i32,
-            expert_ids_ptr as *const i32,
-            num_tokens_ptr as *const i32,
-            topk_ptr as *const f32,
-            lock_len as i32,
-            routing.max_padded_tokens as i32,
-            routing.block_size as i32,
-            top_k as i32,
-            mul_topk_weights,
-            size_m as i32,
-            size_n as i32,
-            size_k as i32,
-            KIMI_K2_LOCAL_EXPERTS as i32,
-            KIMI_K2_INT4_GROUP_SIZE as i32,
-            0,
-            ctx.stream.cu_stream(),
-        )
+    let buffers = MarlinGemmBuffers {
+        input,
+        output,
+        c_tmp: &mut workspace.c_tmp,
+        qweight: weight_packed_uint4b8,
+        scales: weight_scale_permuted,
+        workspace: &mut workspace.locks,
+        sorted_token_ids: routing.sorted_token_ids,
+        expert_ids: routing.expert_ids,
+        num_tokens_post_padded: routing.num_tokens_post_padded,
+        topk_weights: topk_weight,
     };
+    let result = launch_marlin_gemm(ctx, buffers, |ptrs| {
+        Ok(unsafe {
+            ffi::kimi_marlin_wna16_gemm_cuda(
+                ptrs.input,
+                ptrs.output,
+                ptrs.c_tmp,
+                ptrs.qweight,
+                ptrs.scales as *const ffi::Half,
+                ptrs.workspace,
+                ptrs.sorted_token_ids,
+                ptrs.expert_ids,
+                ptrs.num_tokens_post_padded,
+                ptrs.topk_weights,
+                lock_len as i32,
+                routing.max_padded_tokens as i32,
+                routing.block_size as i32,
+                top_k as i32,
+                mul_topk_weights,
+                size_m as i32,
+                size_n as i32,
+                size_k as i32,
+                KIMI_K2_LOCAL_EXPERTS as i32,
+                KIMI_K2_INT4_GROUP_SIZE as i32,
+                0,
+                ctx.stream.cu_stream(),
+            )
+        })
+    })?;
     result.result()?;
     Ok(())
 }