perf(kimi-k2): compact local top-1 readback into one pinned transfer

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<<<blocks, threads, 0, stream>>>(
+      ids, values, reinterpret_cast<Top1Packet*>(packets), rows);
+}
 }

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;
 }

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},

openinfer-kimi-k2/src/runner/worker/runtime.rs

@@ -327,27 +327,73 @@ pub(super) fn read_local_top1_batch_values(
     active_rows: usize,
     top1_values: &mut CudaSlice<half::bf16>,
     out: &mut CudaSlice<i32>,
+    packets: &mut CudaSlice<u8>,
+    packets_host: &mut PinnedHostSlice<u8>,
 ) -> Result<Vec<(u32, f32)>> {
-    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 {}",
             top_id,
             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)
 }

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) {

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<i32>,
     pub(crate) top1_partial_values: CudaSlice<f32>,
     pub(crate) top1_partial_indices: CudaSlice<i32>,
+    /// Compact top-1 readback packets: 8 bytes/row {i32 id, bf16 value} (#716).
+    pub(crate) top1_packets: CudaSlice<u8>,
+    pub(crate) top1_packets_host: PinnedHostSlice<u8>,
     /// 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<Self> {
         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,
         })