[Perf] Size D2D producer grids proportionally to transfer bytes

native/csrc/ring/producer.cu

@@ -656,21 +656,14 @@ __global__ void record_producer_segmented_pack_kernel(
 // Host launchers
 // ---------------------------------------------------------------------------
 
-static inline uint32_t pick_tier_blocks(uint64_t nbytes) {
-    if      (nbytes <= TIER1_THRESHOLD) return TIER0_BLOCKS;
-    else if (nbytes <= TIER2_THRESHOLD) return TIER1_BLOCKS;
-    else if (nbytes <= TIER3_THRESHOLD) return TIER2_BLOCKS;
-    else                                return TIER3_BLOCKS;
-}
-
 void launch_producer_static(
     const RingState& ring,
     const uint8_t*   d_src,
     uint64_t         nbytes,
     uint32_t         hook_type,
     cudaStream_t     stream)
 {
-    const uint32_t n_blocks = pick_tier_blocks(nbytes);
+    const uint32_t n_blocks = producer_grid_blocks(nbytes);
     producer_static_kernel<<<n_blocks, PRODUCER_BLOCK_DIM, 0, stream>>>(
         ring, d_src, nbytes, hook_type);
 }
@@ -684,7 +677,7 @@ void launch_producer_prefix(
     uint32_t         hook_type,
     cudaStream_t     stream)
 {
-    const uint32_t n_blocks = pick_tier_blocks(nbytes_upper);
+    const uint32_t n_blocks = producer_grid_blocks(nbytes_upper);
     producer_prefix_kernel<<<n_blocks, PRODUCER_BLOCK_DIM, 0, stream>>>(
         ring, d_src, nbytes_upper, row_count_dev_ptr, row_bytes, hook_type);
 }
@@ -698,11 +691,8 @@ void launch_producer_chunked(
     uint32_t         hook_type,
     cudaStream_t     stream)
 {
-    uint32_t n_blocks = pick_tier_blocks(nbytes_upper);
-    // Round up so n_blocks % K == 0 (clean block-to-chunk partition).
-    if (n_blocks < K) n_blocks = K;
-    const uint32_t r = n_blocks % K;
-    if (r != 0) n_blocks += (K - r);
+    // A multiple of K (clean block-to-chunk partition) within the cap.
+    const uint32_t n_blocks = producer_chunked_grid_blocks(nbytes_upper, K);
     producer_chunked_kernel<<<n_blocks, PRODUCER_BLOCK_DIM, 0, stream>>>(
         ring, d_src, nbytes_upper, chunk_bytes_dev_ptr, K, hook_type);
 }
@@ -714,7 +704,7 @@ void launch_record_producer_static(
     const int32_t*   emit_gate,
     int32_t          emit_value,
     cudaStream_t     stream) {
-    const uint32_t blocks = pick_tier_blocks(nbytes);
+    const uint32_t blocks = producer_grid_blocks(nbytes);
     record_producer_static_kernel<<<blocks, PRODUCER_BLOCK_DIM, 0, stream>>>(
         ring, d_src, nbytes, emit_gate, emit_value);
 }
@@ -732,7 +722,7 @@ void launch_record_producer_prefix(
         throw std::invalid_argument(
             "record prefix producer requires positive row bytes");
     }
-    const uint32_t blocks = pick_tier_blocks(nbytes_upper);
+    const uint32_t blocks = producer_grid_blocks(nbytes_upper);
     record_producer_prefix_kernel<<<blocks, PRODUCER_BLOCK_DIM, 0, stream>>>(
         ring, d_src, nbytes_upper, row_count_dev_ptr, row_bytes,
         emit_gate, emit_value);
@@ -752,10 +742,8 @@ void launch_record_producer_chunked(
         throw std::invalid_argument(
             "record chunked producer requires a valid evenly divided chunk count");
     }
-    uint32_t blocks = pick_tier_blocks(nbytes_upper);
-    if (blocks < chunk_count) blocks = chunk_count;
-    const uint32_t remainder = blocks % chunk_count;
-    if (remainder != 0) blocks += chunk_count - remainder;
+    const uint32_t blocks =
+        producer_chunked_grid_blocks(nbytes_upper, chunk_count);
     record_producer_chunked_kernel<<<blocks, PRODUCER_BLOCK_DIM, 0, stream>>>(
         ring, d_src, nbytes_upper, chunk_bytes_dev_ptr, chunk_count,
         emit_gate, emit_value);

native/csrc/ring/producer.cuh

@@ -28,6 +28,62 @@
 
 namespace ring {
 
+// Grid geometry for the D2D producers: one block per
+// PRODUCER_BYTES_PER_BLOCK of the bytes the launch may copy (the
+// capture-time upper bound for the prefix and chunked producers), capped at
+// PRODUCER_MAX_BLOCKS.  Host-only code (and CPU unit tests) can use these
+// helpers without a CUDA compiler.
+//
+// 16 KiB per block is four uint4 stores per thread at 256 threads; 8, 16 and
+// 32 KiB per block measured within 0.5 us of each other, and 16 KiB keeps a
+// graph captured for many rows cheap when a replay copies only a few.
+//
+// PRODUCER_MAX_BLOCKS is a bounded heuristic cap, not a residency
+// guarantee: 264 is sized for about two resident 256-thread blocks per SM on
+// 128/132-SM parts.  On smaller GPUs, or when other kernels occupy SMs, a
+// capped grid can still take more than one wave.  Correctness never depends
+// on the grid fitting in one wave: the kernels stride over the bytes they
+// actually copy, so a grid sized by the upper bound stays correct when a
+// replay copies fewer bytes, and surplus blocks only join the arrival count.
+static constexpr uint64_t PRODUCER_BYTES_PER_BLOCK = 16 * 1024;
+static constexpr uint32_t PRODUCER_MAX_BLOCKS      = 264;
+
+#if defined(__CUDACC__)
+#define RING_PRODUCER_HD __host__ __device__
+#else
+#define RING_PRODUCER_HD
+#endif
+
+// Overflow-free ceiling division: nbytes + PRODUCER_BYTES_PER_BLOCK - 1
+// would wrap for byte counts near UINT64_MAX.
+RING_PRODUCER_HD inline uint32_t producer_grid_blocks(uint64_t nbytes) {
+    if (nbytes <= PRODUCER_BYTES_PER_BLOCK) return 1;
+    const uint64_t blocks = (nbytes - 1) / PRODUCER_BYTES_PER_BLOCK + 1;
+    return blocks > PRODUCER_MAX_BLOCKS
+        ? PRODUCER_MAX_BLOCKS
+        : static_cast<uint32_t>(blocks);
+}
+
+// Grid for the chunked producers, which partition blocks evenly across
+// `nchunks` chunks and therefore need a multiple of `nchunks` blocks.  The
+// proportional grid is rounded up to that multiple while it stays within
+// PRODUCER_MAX_BLOCKS; once rounding up would pass the cap, the largest
+// multiple of `nchunks` not above the cap is used instead.  Only when
+// `nchunks` itself exceeds the cap does the grid exceed it (one block per
+// chunk is the minimum the partition needs).  `nchunks == 0` is rejected by
+// the launchers; it maps to the unchunked grid here.
+RING_PRODUCER_HD inline uint32_t producer_chunked_grid_blocks(
+    uint64_t nbytes_upper, uint32_t nchunks) {
+    const uint32_t raw = producer_grid_blocks(nbytes_upper);
+    if (nchunks <= 1) return raw;
+    if (nchunks >= PRODUCER_MAX_BLOCKS || raw <= nchunks) return nchunks;
+    const uint32_t rounded_up = ((raw - 1) / nchunks + 1) * nchunks;
+    if (rounded_up <= PRODUCER_MAX_BLOCKS) return rounded_up;
+    return (PRODUCER_MAX_BLOCKS / nchunks) * nchunks;
+}
+
+#undef RING_PRODUCER_HD
+
 // Upper bound on K (chunk count) for shared-memory sizing in the
 // chunked kernel and host-side argument validation.
 static constexpr uint32_t PRODUCER_MAX_K = 64;
@@ -43,17 +99,6 @@ namespace ring {
 
 static constexpr uint32_t PRODUCER_BLOCK_DIM = 256;
 
-// Size-tier thresholds for grid selection (bytes).
-static constexpr uint64_t TIER1_THRESHOLD =   64 * 1024;       //  64 KB
-static constexpr uint64_t TIER2_THRESHOLD =  4 * 1024 * 1024;  //   4 MB
-static constexpr uint64_t TIER3_THRESHOLD = 32 * 1024 * 1024;  //  32 MB
-
-// Block counts per tier.
-static constexpr uint32_t TIER0_BLOCKS =  1;  // <= 64 KB
-static constexpr uint32_t TIER1_BLOCKS =  4;  // <= 4 MB
-static constexpr uint32_t TIER2_BLOCKS = 16;  // <= 32 MB
-static constexpr uint32_t TIER3_BLOCKS = 64;  // > 32 MB
-
 // Record-only row packers use one block per logical output row and cap the
 // grid so a conservative upper bound cannot create an unbounded launch.
 static constexpr uint32_t RECORD_PACK_MAX_BLOCKS = 1024;

tests/native/ring/Makefile

@@ -69,7 +69,8 @@ TARGETS := $(BUILD)/test_rings $(BUILD)/test_producer \
            $(BUILD)/test_clickhouse_record_sink
 
 .PHONY: all test_rings test_producer test_ring_engine test_null_mode \
-        test_record_consumer test_clickhouse_record_sink run clean FORCE
+        test_record_consumer test_clickhouse_record_sink run clean FORCE \
+        bench_producer_grid
 
 all: $(TARGETS)
 
@@ -125,6 +126,13 @@ test_record_consumer: $(BUILD)/test_record_consumer
 test_clickhouse_record_sink: $(BUILD)/test_clickhouse_record_sink
 	$(BUILD)/test_clickhouse_record_sink
 
+# Manual microbenchmark; not part of `all` or `run`.
+$(BUILD)/bench_producer_grid: bench_producer_grid.cu $(PRODUCER_CU) $(CUDA_CONFIG_STAMP) | $(BUILD)
+	$(CUDA_NVCC) $(NVCC_FLAGS) bench_producer_grid.cu -o $@
+
+bench_producer_grid: $(BUILD)/bench_producer_grid
+	$(BUILD)/bench_producer_grid
+
 run: all
 	$(BUILD)/test_rings
 	$(BUILD)/test_producer

tests/native/ring/bench_producer_grid.cu

@@ -0,0 +1,128 @@
+// Microbenchmark: D2D producer copy time under the previous size-tier grid
+// (1/4/16/64 blocks at 64 KB / 4 MB / 32 MB) and the proportional grid
+// (ring::producer_grid_blocks).  Not part of `make run`; build and run with
+//   make -C tests/native/ring bench_producer_grid
+//
+// Each configuration replays a CUDA graph of 64 back-to-back producer
+// launches so launch overhead is amortized the same way it is in a captured
+// decode step; the reported time is per launch.
+
+#include "../../../native/csrc/ring/producer.cu"
+#include "ring/ring_alloc.h"
+
+#include <cuda_runtime.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+#define CUDA_CHECK(expr)                                                    \
+    do {                                                                    \
+        cudaError_t error_ = (expr);                                        \
+        if (error_ != cudaSuccess) {                                        \
+            std::fprintf(stderr, "CUDA error at %s:%d: %s\n",             \
+                         __FILE__, __LINE__, cudaGetErrorString(error_));    \
+            std::exit(1);                                                   \
+        }                                                                   \
+    } while (0)
+
+static uint32_t previous_tier_blocks(uint64_t nbytes) {
+    if (nbytes <= 64ull * 1024) return 1;
+    if (nbytes <= 4ull * 1024 * 1024) return 4;
+    if (nbytes <= 32ull * 1024 * 1024) return 16;
+    return 64;
+}
+
+static constexpr int kLaunchesPerGraph = 64;
+static constexpr int kReplays = 50;
+
+// Returns microseconds per producer launch.
+static double time_static(ring::RingState& state, const uint8_t* src,
+                          uint64_t nbytes, uint32_t blocks,
+                          cudaStream_t stream) {
+    cudaGraph_t graph;
+    cudaGraphExec_t exec;
+    CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
+    for (int i = 0; i < kLaunchesPerGraph; ++i) {
+        ring::producer_static_kernel<<<blocks, ring::PRODUCER_BLOCK_DIM, 0,
+                                       stream>>>(state, src, nbytes, 0);
+    }
+    CUDA_CHECK(cudaStreamEndCapture(stream, &graph));
+    CUDA_CHECK(cudaGraphInstantiate(&exec, graph, 0));
+
+    cudaEvent_t start, stop;
+    CUDA_CHECK(cudaEventCreate(&start));
+    CUDA_CHECK(cudaEventCreate(&stop));
+    // Warm up, then reset the ring heads between replays from the host
+    // (outside the timed region) so every replay writes the same slots.
+    for (int warm = 0; warm < 3; ++warm) {
+        *state.task_head = 0;
+        *state.payload_head = 0;
+        CUDA_CHECK(cudaGraphLaunch(exec, stream));
+        CUDA_CHECK(cudaStreamSynchronize(stream));
+    }
+    float total_ms = 0.0f;
+    for (int replay = 0; replay < kReplays; ++replay) {
+        *state.task_head = 0;
+        *state.payload_head = 0;
+        CUDA_CHECK(cudaEventRecord(start, stream));
+        CUDA_CHECK(cudaGraphLaunch(exec, stream));
+        CUDA_CHECK(cudaEventRecord(stop, stream));
+        CUDA_CHECK(cudaEventSynchronize(stop));
+        float ms = 0.0f;
+        CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
+        total_ms += ms;
+    }
+    CUDA_CHECK(cudaEventDestroy(start));
+    CUDA_CHECK(cudaEventDestroy(stop));
+    CUDA_CHECK(cudaGraphExecDestroy(exec));
+    CUDA_CHECK(cudaGraphDestroy(graph));
+    return 1000.0 * total_ms / (kReplays * kLaunchesPerGraph);
+}
+
+int main() {
+    setbuf(stdout, nullptr);
+    ring::set_ring_null_mode(false);
+    cudaDeviceProp prop;
+    CUDA_CHECK(cudaGetDeviceProperties(&prop, 0));
+    std::printf("device: %s (%d SMs)\n", prop.name, prop.multiProcessorCount);
+
+    const std::vector<uint64_t> sizes = {
+        8ull * 1024,           // one decode row of a 4096-wide bf16 hidden state
+        64ull * 1024,          // 8 rows
+        256ull * 1024,         // 32 rows
+        1200ull * 1024,        // ~150 rows
+        4ull * 1024 * 1024,    // 512 rows
+        16ull * 1024 * 1024,   // a 2048-token prefill chunk
+    };
+    uint64_t largest = sizes.back();
+    ring::RingConfig cfg;
+    cfg.task_ring_entries = 128;
+    cfg.payload_ring_bytes = largest * (kLaunchesPerGraph + 1);
+    cfg.pinned_staging_bytes = 4096;
+    ring::AllocatedRing allocated(cfg);
+    allocated.init();
+    ring::RingState& state = allocated.state();
+
+    uint8_t* src = nullptr;
+    CUDA_CHECK(cudaMalloc(&src, largest));
+    CUDA_CHECK(cudaMemset(src, 0x5A, largest));
+    cudaStream_t stream;
+    CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+
+    std::printf("%12s %10s %12s %10s %12s %8s\n", "bytes", "old_blocks",
+                "old_us", "new_blocks", "new_us", "speedup");
+    for (uint64_t nbytes : sizes) {
+        const uint32_t old_blocks = previous_tier_blocks(nbytes);
+        const uint32_t new_blocks = ring::producer_grid_blocks(nbytes);
+        const double old_us = time_static(state, src, nbytes, old_blocks, stream);
+        const double new_us = time_static(state, src, nbytes, new_blocks, stream);
+        std::printf("%12llu %10u %12.2f %10u %12.2f %7.2fx  (%.0f -> %.0f GB/s)\n",
+                    static_cast<unsigned long long>(nbytes), old_blocks, old_us,
+                    new_blocks, new_us, old_us / new_us,
+                    nbytes / (old_us * 1e3), nbytes / (new_us * 1e3));
+    }
+    CUDA_CHECK(cudaStreamDestroy(stream));
+    CUDA_CHECK(cudaFree(src));
+    return 0;
+}