diff --git a/openinfer-kernels/csrc/shared/argmax.cu b/openinfer-kernels/csrc/shared/argmax.cu index 851f643fb..b26981876 100644 --- a/openinfer-kernels/csrc/shared/argmax.cu +++ b/openinfer-kernels/csrc/shared/argmax.cu @@ -429,4 +429,27 @@ void markov_step_argmax_cuda(const __nv_bfloat16* base, partial_indices, out_tokens, sampled_tokens, block_size, step, rows, tiles_per_row); } + +struct Top1Packet { + int id; + __nv_bfloat16 value; +}; + +__global__ void pack_top1_packets_kernel(const int* ids, + const __nv_bfloat16* values, + Top1Packet* packets, int rows) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row < rows) { + packets[row].id = ids[row]; + packets[row].value = values[row]; + } +} + +void pack_top1_packets_cuda(const int* ids, const __nv_bfloat16* values, + void* packets, int rows, cudaStream_t stream) { + int threads = 128; + int blocks = (rows + threads - 1) / threads; + pack_top1_packets_kernel<<>>( + ids, values, reinterpret_cast(packets), rows); +} } diff --git a/openinfer-kernels/src/ffi/shared.rs b/openinfer-kernels/src/ffi/shared.rs index eb1caafe5..8017e7945 100644 --- a/openinfer-kernels/src/ffi/shared.rs +++ b/openinfer-kernels/src/ffi/shared.rs @@ -740,6 +740,16 @@ unsafe extern "C" { } +unsafe extern "C" { + pub fn pack_top1_packets_cuda( + ids: *const i32, + values: *const Half, + packets: *mut core::ffi::c_void, + rows: i32, + stream: CUstream, + ); +} + unsafe extern "C" { pub fn openinfer_kernels_last_error() -> *const std::os::raw::c_char; } diff --git a/openinfer-kimi-k2/src/runner/worker.rs b/openinfer-kimi-k2/src/runner/worker.rs index 51fc2352e..50898a499 100644 --- a/openinfer-kimi-k2/src/runner/worker.rs +++ b/openinfer-kimi-k2/src/runner/worker.rs @@ -8,7 +8,7 @@ use std::{ use anyhow::{Context, Result, ensure}; use bytesize::ByteSize; use crossbeam_channel::{Receiver, Sender, bounded, unbounded}; -use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut}; +use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut, HostSlice, PinnedHostSlice}; use cudarc::nccl::{ ReduceOp, safe::{Comm, Id}, diff --git a/openinfer-kimi-k2/src/runner/worker/runtime.rs b/openinfer-kimi-k2/src/runner/worker/runtime.rs index ddc33c637..8663db7d4 100644 --- a/openinfer-kimi-k2/src/runner/worker/runtime.rs +++ b/openinfer-kimi-k2/src/runner/worker/runtime.rs @@ -327,19 +327,65 @@ pub(super) fn read_local_top1_batch_values( active_rows: usize, top1_values: &mut CudaSlice, out: &mut CudaSlice, + packets: &mut CudaSlice, + packets_host: &mut PinnedHostSlice, ) -> Result> { - ctx.sync()?; - let top_ids = ctx - .stream - .clone_dtoh(&*out) - .map_err(|err| anyhow::anyhow!("D2H Kimi batched top1 ids read failed: {err}"))?; - let top_values = ctx - .stream - .clone_dtoh(&*top1_values) - .map_err(|err| anyhow::anyhow!("D2H Kimi batched top1 values read failed: {err}"))?; + ensure!( + active_rows <= logits.seq_len, + "read_local_top1_batch_values: active_rows {} exceeds logits seq_len {}", + active_rows, + logits.seq_len + ); + ensure!( + out.len() >= active_rows && top1_values.len() >= active_rows, + "read_local_top1_batch_values: scratch len {} / {} < active_rows {}", + out.len(), + top1_values.len(), + active_rows + ); + ensure!( + packets.len() >= active_rows * 8, + "read_local_top1_batch_values: packets len {} < active_rows * 8 {}", + packets.len(), + active_rows * 8 + ); + + { + let (ids_ptr, _ids_guard) = out.device_ptr_mut(&ctx.stream); + let (values_ptr, _values_guard) = top1_values.device_ptr_mut(&ctx.stream); + let (packets_ptr, _packets_guard) = packets.device_ptr_mut(&ctx.stream); + unsafe { + ffi::pack_top1_packets_cuda( + ids_ptr as *const i32, + values_ptr as *const ffi::Half, + packets_ptr as *mut core::ffi::c_void, + active_rows as i32, + ctx.stream.cu_stream(), + ); + } + } + + let active_bytes = active_rows * 8; + { + let (host_slice, guard) = unsafe { packets_host.stream_synced_mut_slice(&ctx.stream) }; + ctx.stream + .memcpy_dtoh( + &packets.slice(0..active_bytes), + &mut host_slice[..active_bytes], + ) + .map_err(|err| anyhow::anyhow!("D2H Kimi batched top1 packet read failed: {err}"))?; + drop(guard); // records the pinned buffer's event on the stream + } + let host = packets_host + .as_slice() + .map_err(|err| anyhow::anyhow!("Kimi batched top1 packet wait failed: {err}"))?; + let mut rows = Vec::with_capacity(active_rows); for row in 0..active_rows { - let top_id = top_ids[row]; + let base = row * 8; + let top_id = + i32::from_le_bytes([host[base], host[base + 1], host[base + 2], host[base + 3]]); + let value_bits = u16::from_le_bytes([host[base + 4], host[base + 5]]); ensure!( top_id >= 0 && (top_id as usize) < logits.hidden_dim, "Kimi batched local top1 id {} at row {} out of logits range {}", @@ -347,7 +393,7 @@ pub(super) fn read_local_top1_batch_values( row, logits.hidden_dim ); - rows.push((top_id as u32, top_values[row].to_f32())); + rows.push((top_id as u32, half::bf16::from_bits(value_bits).to_f32())); } Ok(rows) } diff --git a/openinfer-kimi-k2/src/runner/worker/state.rs b/openinfer-kimi-k2/src/runner/worker/state.rs index e832392e3..c4716652e 100644 --- a/openinfer-kimi-k2/src/runner/worker/state.rs +++ b/openinfer-kimi-k2/src/runner/worker/state.rs @@ -386,6 +386,8 @@ impl KimiRankThreadState { active_len, &mut decode_arena.scratch.sampling.top1_value_scratch, &mut decode_arena.scratch.sampling.top1_out, + &mut decode_arena.scratch.sampling.top1_packets, + &mut decode_arena.scratch.sampling.top1_packets_host, )?; let mut picks: Vec<(u32, f32)> = local_top1; for (sampling_row, token) in sampling_rows.iter().zip(&sampled) { diff --git a/openinfer-kimi-k2/src/typed_scratch.rs b/openinfer-kimi-k2/src/typed_scratch.rs index cd69921e6..e3d56215b 100644 --- a/openinfer-kimi-k2/src/typed_scratch.rs +++ b/openinfer-kimi-k2/src/typed_scratch.rs @@ -1,7 +1,7 @@ //! Typed Kimi decode scratch buffers. use anyhow::{Result, ensure}; -use cudarc::driver::CudaSlice; +use cudarc::driver::{CudaSlice, PinnedHostSlice}; use openinfer_kernels::gpu_buffers; use openinfer_kernels::tensor::{DeviceContext, GpuTensor, HiddenStates}; @@ -146,6 +146,9 @@ pub(crate) struct SamplingScratch { pub(crate) top1_out: CudaSlice, pub(crate) top1_partial_values: CudaSlice, pub(crate) top1_partial_indices: CudaSlice, + /// Compact top-1 readback packets: 8 bytes/row {i32 id, bf16 value} (#716). + pub(crate) top1_packets: CudaSlice, + pub(crate) top1_packets_host: PinnedHostSlice, /// Buffers for non-greedy rows (f32 probs are batch x vocab, ~42 MB at /// batch 64) — allocated on the first sampling request so greedy-only /// serving pays nothing. @@ -156,11 +159,14 @@ pub(crate) struct SamplingScratch { impl SamplingScratch { pub(crate) fn new(ctx: &DeviceContext, batch_size: usize) -> Result { let partials = argmax_batch_bf16_split_partials_len(batch_size, KIMI_K2_VOCAB); + let packet_bytes = batch_size * 8; // sizeof {i32 id, bf16 value} per row Ok(Self { top1_value_scratch: ctx.stream.alloc_zeros(batch_size)?, top1_out: ctx.stream.alloc_zeros(batch_size)?, top1_partial_values: ctx.stream.alloc_zeros(partials)?, top1_partial_indices: ctx.stream.alloc_zeros(partials)?, + top1_packets: ctx.stream.alloc_zeros(packet_bytes)?, + top1_packets_host: unsafe { ctx.ctx.alloc_pinned(packet_bytes)? }, batch_sampling: None, batch_size, })