Skip to main content

Module kvgpu

Module kvgpu 

Source
Expand description

Host side of the GPU KV decoder: packs GPU-layout TKV streams (crate::kvcodec, gpu = true) into the flat tables the device kernel reads (tang-compute tkv_decode, MSL source in tang-compute/src/kernels/kvdec_msl.rs). No GPU dependency here.

One kernel thread per lane = one (stream, block, head) arithmetic-coded substream. Tables (all u32):

  • stream: every stream’s bytes, concatenated, little-endian words;
  • lanes: LANE_WORDS words per lane: byte offset, byte length, first token, tokens, channel-table base, output base (elements), output token stride, scratch base (words), unit (f32 bits), sink, sink divisor, lane width (channels), RoPE flag, RoPE theta (f32 bits), first RoPE pair of the lane, head dim;
  • chans: 4 words per channel, in lane-local order per (stream, head, channel group): mean (i32), rho, class, dense class index within the group;
  • prior: crate::kvcodec::gpu_prior_table;
  • scratch: scratch_words words of adaptive-model state.

Structs§

Batch
Flat decoder tables for a set of streams.
Placement
Where a stream’s output goes: element index of (head, token) row start and the token stride, plus optional RoPE (K given pre-RoPE).

Constants§

LANE_WORDS
Words per lane descriptor.