Module: GPT2KVCuda

Defined in:
lib/toy/llm/engine/gpt2_kv_engine_cuda.rb

Class Method Summary collapse

Class Method Details

.decode_step(kv_cache, token_id, pos) ⇒ Object

Decode one new token at position ‘pos`. Writes K, V[:, pos] as a side effect, returns the (vocab,) logits Mat for the new position. The caller can argmax (greedy) or sample.



348
349
350
351
352
353
354
355
356
357
 
# File 'lib/toy/llm/engine/gpt2_kv_engine_cuda.rb', line 348

def self.decode_step(kv_cache, token_id, pos)
  TinyNNCuda.tnn_reset_for_rebuild(kv_cache.sess)
  step = kv_cache.build_decode_step(pos)
  TinyNNCuda.tnn_realize(kv_cache.sess, step.kv_step_logits)
  TinyNNCuda.upload_int_array(kv_cache.sess, step.t_token_id, [token_id])
  TinyNNCuda.tnn_compute(kv_cache.sess)
  # Logits ne=[vocab, 1]. Download as (1, vocab) row-major — same
  # layout as a single-row Mat with vocab columns.
  TinyNNCuda.download_row_major(kv_cache.sess, step.kv_step_logits, 1, kv_cache.vocab_size)
end

.upload_from(kv_cache, model) ⇒ Object

Upload all GPT-2 weights (+ zero-init the K/V buffers) into a realized GPT2KVFFICacheCuda. Counterpart to GPT2FFI.upload_from for the KV cache variant. Note: pos_embed is uploaded in FULL (all context_length rows), not sliced.



292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
 
# File 'lib/toy/llm/engine/gpt2_kv_engine_cuda.rb', line 292

def self.upload_from(kv_cache, model)
  sess    = kv_cache.sess
  n       = kv_cache.n_layers
  n_heads = kv_cache.n_heads
  d_model = kv_cache.d_model
  d_head  = kv_cache.d_head
  max_T   = kv_cache.max_T

  TinyNNCuda.upload_row_major(sess, kv_cache.t_token_embed, model.token_embed)
  TinyNNCuda.upload_row_major(sess, kv_cache.t_pos_embed,   model.pos_embed)
  TinyNNCuda.tnn_upload_from_float_array(sess, kv_cache.t_ln_f_gamma, model.ln_f_gamma, d_model)
  TinyNNCuda.tnn_upload_from_float_array(sess, kv_cache.t_ln_f_beta,  model.ln_f_beta,  d_model)

  # Zero buffers for K/V (ggml_backend_alloc_ctx_tensors typically
  # zeros, but be explicit so reuse across multiple decode runs has
  # a clean starting state).
  kv_zero_k = Mat.new(max_T,  d_head)
  kv_zero_v = Mat.new(d_head, max_T)

  li = 0
  while li < n
    blk_n = model.gpt2_blocks[li]
    blk_f = kv_cache.kv_blocks_ffi[li]

    TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_ln1_gamma, blk_n.ln1_gamma, d_model)
    TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_ln1_beta,  blk_n.ln1_beta,  d_model)
    TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_ln2_gamma, blk_n.ln2_gamma, d_model)
    TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_ln2_beta,  blk_n.ln2_beta,  d_model)

    h = 0
    while h < n_heads
      TinyNNCuda.stage_transposed_and_upload(sess, blk_f.t_w_q[h], blk_n.w_q[h])
      TinyNNCuda.stage_transposed_and_upload(sess, blk_f.t_w_k[h], blk_n.w_k[h])
      TinyNNCuda.stage_transposed_and_upload(sess, blk_f.t_w_v[h], blk_n.w_v[h])
      TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_b_q[h], blk_n.b_q[h], d_head)
      TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_b_k[h], blk_n.b_k[h], d_head)
      TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_b_v[h], blk_n.b_v[h], d_head)
      TinyNNCuda.upload_row_major(sess, blk_f.t_K[h], kv_zero_k)
      TinyNNCuda.upload_row_major(sess, blk_f.t_V[h], kv_zero_v)
      h = h + 1
    end

    TinyNNCuda.stage_transposed_and_upload(sess, blk_f.t_w_o,   blk_n.w_o)
    TinyNNCuda.stage_transposed_and_upload(sess, blk_f.t_w_ff1, blk_n.w_ff1)
    TinyNNCuda.stage_transposed_and_upload(sess, blk_f.t_w_ff2, blk_n.w_ff2)
    TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_b_o,   blk_n.b_o,   d_model)
    TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_b_ff1, blk_n.b_ff1, kv_cache.d_ff)
    TinyNNCuda.tnn_upload_from_float_array(sess, blk_f.t_b_ff2, blk_n.b_ff2, d_model)

    li = li + 1
  end
end