Last updated: 2026-09-07

Sampling

Compute Argmax

All API Sections

Wrap the GPU argmax reduction used for greedy token sampling.

This helper owns the compute pipeline for the two-phase argmax shader and records the reduction dispatches that pick the next token entirely on GPU.

1 exports 5 methods src/compute/argmax.zig

1 exports shown

struct

ArgmaxDispatch

# 
pub const ArgmaxDispatch = struct

GPU-accelerated two-phase argmax reduction for greedy token sampling.

src/compute/argmax.zig:20

Methods

5

method

ArgmaxDispatch.init

# 
pub fn init( instance: *const Instance, shader_dir: []const u8, allocator: std.mem.Allocator, ) !ArgmaxDispatch

Create the argmax compute pipeline and descriptor pool on the given Vulkan instance.

Parameters
instance
Vulkan instance that owns the device used for all Vulkan calls.
shader_dir
Directory path searched for `argmax.spv`; if the shader is missing the pipeline is set to null and a warning is logged.
allocator
Allocator used internally by the pipeline creation helper.
Returns

An initialised `ArgmaxDispatch`; the caller must call `deinit` to release GPU resources.

src/compute/argmax.zig:30

method

ArgmaxDispatch.record

# 
pub fn record( self: *const ArgmaxDispatch, cmd: *CommandBuffer, descriptor_set: vk.c.VkDescriptorSet, n_logits: u32, phase0_workgroups: u32, ) !void

Record the two-phase argmax reduction into a command buffer.

Phase 0 dispatches `phase0_workgroups` workgroups that each reduce a slice of the logit vector and write partial (value, index) results; phase 1 dispatches a single workgroup that reduces those partials to the final winner. A compute barrier is inserted between the two phases.

Parameters
cmd
Command buffer to record dispatches into.
descriptor_set
Descriptor set with logits, partials, and result buffers already bound.
n_logits
Total number of logits in the input buffer (vocabulary size).
phase0_workgroups
Number of workgroups launched in phase 0; also the number of partial results consumed by phase 1.

src/compute/argmax.zig:78

method

ArgmaxDispatch.allocDescriptorSet

# 
pub fn allocDescriptorSet(self: *const ArgmaxDispatch) !vk.c.VkDescriptorSet

Allocate a descriptor set from the argmax descriptor pool.

Returns

A freshly allocated `VkDescriptorSet` using the pipeline's layout, or an error if allocation fails or the shader was not loaded.

Notes

The set must be freed back to the pool before calling `deinit`.

src/compute/argmax.zig:104

method

ArgmaxDispatch.writeDescriptorSet

# 
pub fn writeDescriptorSet( self: *const ArgmaxDispatch, descriptor_set: vk.c.VkDescriptorSet, logits_buf: vk.c.VkBuffer, logits_size: vk.c.VkDeviceSize, partials_buf: vk.c.VkBuffer, partials_size: vk.c.VkDeviceSize, result_buf: vk.c.VkBuffer, result_size: vk.c.VkDeviceSize, ) void

Bind the logits, partials, and result buffers to a descriptor set via `vkUpdateDescriptorSets`.

The three buffers map to shader bindings 0, 1, and 2 respectively.

Parameters
descriptor_set
Target descriptor set to update (must have been allocated via `allocDescriptorSet`).
logits_buf
Storage buffer containing the raw logit values (shader binding 0).
logits_size
Byte range of `logits_buf` to expose to the shader.
partials_buf
Intermediate storage buffer for phase-0 partial results (shader binding 1).
partials_size
Byte range of `partials_buf` to expose to the shader.
result_buf
Output storage buffer that receives the winning token index after phase 1 (shader binding 2).
result_size
Byte range of `result_buf` to expose to the shader.

src/compute/argmax.zig:128