RL Inference
at GB300 Scale

Descriptor Ring — ring.h

The central dispatch structure. The CPU (producer) enqueues work descriptors; GPU SMs (consumers) dequeue them with load-acquire semantics. The ring lives in NVLink-C2C coherent memory and requires no cache flushes.

/* Producer-consumer index pair, cache-line padded */
typedef struct __attribute__((packed)) {
  volatile uint32_t head;
  uint32_t           tail;
  uint8_t            pad[56];
} RingHeadTail;

typedef struct {
  RingHeadTail  hta __attribute__((aligned(64)));
  RingHeadTail  htb __attribute__((aligned(64)));
  Descriptor    slots[RING_SIZE] __attribute__((aligned(128)));
} CommandRing;

/* Acquire n contiguous slots; returns UINT32_MAX if full */
static inline uint32_t
ring_acquire(CommandRing *ring, uint32_t n) {
  uint32_t h = __atomic_load_n(&ring->hta.head, __ATOMIC_ACQUIRE);
  uint32_t t = __atomic_load_n(&ring->hta.tail, __ATOMIC_RELAXED);
  if (RING_SIZE - (h - t) < n) return UINT32_MAX;
  uint32_t pos = h & (RING_SIZE - 1);
  __atomic_store_n(&ring->hta.head, h + n, __ATOMIC_RELEASE);
  return pos;
}

Work Descriptor — descriptor.h

Each descriptor encodes a single decode step: which sequence, which KV block range, which attention parameters, and where to write the output token.

typedef struct __attribute__((packed)) {
  uint64_t  seq_id;
  uint32_t  kv_block_offset;
  uint16_t  num_kv_blocks;
  uint8_t   attention_flags;
  uint8_t   pad;
  uint32_t  output_token_offset;
  uint64_t  reward_cookie;
} Descriptor;

static_assert(sizeof(Descriptor) == 24, "Descriptor must be 24 bytes");

Hugepage KV Arena — arena.h

The KV cache is pre-allocated as a contiguous hugepage-backed slab at process startup. Each trajectory receives a fixed-size block. The arena uses a bitmap allocator with O(1) acquire/release.

#define GB(n) ((size_t)(n) << 30)

typedef struct {
  uint8_t  *base;
  size_t    capacity;
  size_t    block_size;
  uint64_t  bitmap[BITMAP_WORDS];
} KVArena;

/* Pre-fault all pages at init; guarantees no page faults at runtime */
void
arena_init(KVArena *a, size_t capacity, size_t block_size) {
  a->capacity   = capacity;
  a->block_size = block_size;
  a->base = mmap(NULL, capacity, PROT_READ | PROT_WRITE,
                 MAP_ANONYMOUS | MAP_PRIVATE | MAP_HUGETLB, -1, 0);
  if (a->base == MAP_FAILED) abort();
  memset(a->base, 0, capacity);
  memset(a->bitmap, 0, sizeof(a->bitmap));
}

/* Acquire one block; returns index or -1 */
int64_t
arena_acquire(KVArena *a) {
  for (int w = 0; w < BITMAP_WORDS; w++) {
    uint64_t bits = ~a->bitmap[w];
    if (bits) {
      int b = __builtin_ctzll(bits);
      a->bitmap[w] |= (1ULL << b);
      return w * 64 + b;
    }
  }
  return -1;
}

Persistent GPU Worker — worker.cu

The persistent decode loop runs on every participating SM. It polls the command ring, processes descriptors, and writes completion tokens. No host interaction occurs between start-of-trajectory and end-of-trajectory.

// Persistent decode worker kernel (one warp per SM)
__global__ void
decode_worker(CommandRing *ring, KVArena *arena,
              OutputRing *out, uint64_t *step_count) {
  uint32_t lane = threadIdx.x & 31;
  while (true) {
    if (lane == 0) {
      uint32_t t = __atomic_load_n(&ring->htb.tail, __ATOMIC_ACQUIRE);
      uint32_t h = __atomic_load_n(&ring->hta.head, __ATOMIC_RELAXED);
      if (t < h) {
        uint32_t pos = t & (RING_SIZE - 1);
        Descriptor d = ring->slots[pos];
        __atomic_store_n(&ring->htb.tail, t + 1, __ATOMIC_RELEASE);
        process_descriptor(d, arena, out);
        if (step_count) atomicAdd(step_count, 1ULL);
      }
    }
  }
}

NUMA-Pinned Allocation — numa.c

All runtime memory is bound to the nearest NUMA node. The helper below allocates and faults hugepages on a specific node.

#include <numa.h>

void *
numa_alloc_hugepages(int node, size_t size) {
  void *p = mmap(NULL, size, PROT_READ | PROT_WRITE,
                  MAP_ANONYMOUS | MAP_PRIVATE | MAP_HUGETLB,
                  -1, 0);
  if (p == MAP_FAILED) return NULL;
  struct bitmask *mask = numa_allocate_cpumask();
  numa_bitmask_setbit(mask, node);
  mbind(p, size, MPOL_BIND, mask->maskp, mask->size, MPOL_MF_MOVE | MPOL_MF_STRICT);
  memset(p, 0, size);
  numa_free_cpumask(mask);
  return p;
}

GPU-NIC Command Ring — atomic.h

When RL trajectories require network I/O (e.g., fetching checkpoints or sending rewards), the GPU communicates with the NIC via a dedicated atomic doorbell ring. The NIC polls a producer index; the GPU atomically advances it.

/* NIC-visible doorbell: GPU increments to notify NIC of new work */
typedef struct {
  volatile uint64_t doorbell __attribute__((aligned(64)));
  uint8_t            pad[56];
} NicDoorbell;

__device__ void
nic_ring_push(NicDoorbell *db, uint64_t cmd_desc_paddr) {
  *(uint64_t *)nic_cmd_slot = cmd_desc_paddr;
  __threadfence();
  atomicAdd((unsigned long long *)&db->doorbell, 1);
}