From 4a70c8790c59837069cc965c005e55f124a0ff37 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 31 Aug 2026 11:49:47 +0100 Subject: [PATCH] refactor(kernels): the Marlin face serves both consumers from one module Signed-off-by: Feathbow --- pegainfer-kernels/src/ops.rs | 2 + pegainfer-kernels/src/ops/gemma4.rs | 128 ++++++++++-------- pegainfer-kernels/src/ops/kimi_k2/experts.rs | 132 ++++++++++--------- pegainfer-kernels/src/ops/marlin_face.rs | 108 +++++++++++++++ 4 files changed, 255 insertions(+), 115 deletions(-) create mode 100644 pegainfer-kernels/src/ops/marlin_face.rs diff --git a/pegainfer-kernels/src/ops.rs b/pegainfer-kernels/src/ops.rs index 085e9e53f..88ca47469 100644 --- a/pegainfer-kernels/src/ops.rs +++ b/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; diff --git a/pegainfer-kernels/src/ops/gemma4.rs b/pegainfer-kernels/src/ops/gemma4.rs index 027f86fa0..c3865a818 100644 --- a/pegainfer-kernels/src/ops/gemma4.rs +++ b/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(()) } diff --git a/pegainfer-kernels/src/ops/kimi_k2/experts.rs b/pegainfer-kernels/src/ops/kimi_k2/experts.rs index a02c5e24d..ae4746ceb 100644 --- a/pegainfer-kernels/src/ops/kimi_k2/experts.rs +++ b/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(()) } diff --git a/pegainfer-kernels/src/ops/marlin_face.rs b/pegainfer-kernels/src/ops/marlin_face.rs new file mode 100644 index 000000000..095adc726 --- /dev/null +++ b/pegainfer-kernels/src/ops/marlin_face.rs @@ -0,0 +1,108 @@ +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use cudarc::driver::sys::CUresult; + +use crate::tensor::DeviceContext; + +pub(super) struct MarlinGemmBuffers<'a, Scale> { + pub input: &'a CudaSlice, + pub output: &'a mut CudaSlice, + pub c_tmp: &'a mut CudaSlice, + pub qweight: &'a CudaSlice, + pub scales: &'a CudaSlice, + pub workspace: &'a mut CudaSlice, + pub sorted_token_ids: &'a CudaSlice, + pub expert_ids: &'a CudaSlice, + pub num_tokens_post_padded: &'a CudaSlice, + pub topk_weights: &'a CudaSlice, +} + +pub(super) struct MarlinGemmPointers { + pub input: *const u16, + pub output: *mut u16, + pub c_tmp: *mut f32, + pub qweight: *const u8, + pub scales: *const Scale, + pub workspace: *mut i32, + pub sorted_token_ids: *const i32, + pub expert_ids: *const i32, + pub num_tokens_post_padded: *const i32, + pub topk_weights: *const f32, +} + +pub(super) fn launch_marlin_gemm( + ctx: &DeviceContext, + buffers: MarlinGemmBuffers<'_, Scale>, + launch: impl FnOnce(MarlinGemmPointers) -> anyhow::Result, +) -> anyhow::Result { + let MarlinGemmBuffers { + input, + output, + c_tmp, + qweight, + scales, + workspace, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_weights, + } = buffers; + let (input, _input_guard) = input.device_ptr(&ctx.stream); + let (output, _output_guard) = output.device_ptr_mut(&ctx.stream); + let (c_tmp, _c_tmp_guard) = c_tmp.device_ptr_mut(&ctx.stream); + let (qweight, _qweight_guard) = qweight.device_ptr(&ctx.stream); + let (scales, _scales_guard) = scales.device_ptr(&ctx.stream); + let (workspace, _workspace_guard) = workspace.device_ptr_mut(&ctx.stream); + let (sorted_token_ids, _sorted_guard) = sorted_token_ids.device_ptr(&ctx.stream); + let (expert_ids, _expert_guard) = expert_ids.device_ptr(&ctx.stream); + let (num_tokens_post_padded, _padded_guard) = num_tokens_post_padded.device_ptr(&ctx.stream); + let (topk_weights, _weights_guard) = topk_weights.device_ptr(&ctx.stream); + launch(MarlinGemmPointers { + input: input as *const u16, + output: output as *mut u16, + c_tmp: c_tmp as *mut f32, + qweight: qweight as *const u8, + scales: scales as *const Scale, + workspace: workspace as *mut i32, + sorted_token_ids: sorted_token_ids as *const i32, + expert_ids: expert_ids as *const i32, + num_tokens_post_padded: num_tokens_post_padded as *const i32, + topk_weights: topk_weights as *const f32, + }) +} + +pub(super) struct MarlinAlignBuffers<'a> { + pub topk_idx: &'a CudaSlice, + pub sorted_token_ids: &'a mut CudaSlice, + pub expert_ids: &'a mut CudaSlice, + pub num_tokens_post_padded: &'a mut CudaSlice, + pub expert_offsets: &'a mut CudaSlice, +} + +pub(super) fn launch_marlin_align( + ctx: &DeviceContext, + buffers: MarlinAlignBuffers<'_>, + launch: impl FnOnce(*const i32, *mut i32, *mut i32, *mut i32, *mut u32) -> anyhow::Result, +) -> anyhow::Result { + let MarlinAlignBuffers { + topk_idx, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + expert_offsets, + } = buffers; + let (topk_idx, _topk_guard) = topk_idx.device_ptr(&ctx.stream); + let (sorted_token_ids, _sorted_guard) = sorted_token_ids.device_ptr_mut(&ctx.stream); + let (expert_ids, _expert_guard) = expert_ids.device_ptr_mut(&ctx.stream); + let (num_tokens_post_padded, _padded_guard) = + num_tokens_post_padded.device_ptr_mut(&ctx.stream); + let (expert_offsets, _offsets_guard) = expert_offsets.device_ptr_mut(&ctx.stream); + launch( + topk_idx as *const i32, + sorted_token_ids as *mut i32, + expert_ids as *mut i32, + num_tokens_post_padded as *mut i32, + expert_offsets as *mut u32, + ) +}