Module: GPT2KV

Defined in:
lib/toy/llm/engine/gpt2_kv_engine.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.



346
347
348
349
350
351
352
353
354
355
 
# File 'lib/toy/llm/engine/gpt2_kv_engine.rb', line 346

def self.decode_step(kv_cache, token_id, pos)
  TinyNN.tnn_reset_for_rebuild(kv_cache.sess)
  step = kv_cache.build_decode_step(pos)
  TinyNN.tnn_realize(kv_cache.sess, step.kv_step_logits)
  TinyNN.upload_int_array(kv_cache.sess, step.t_token_id, [token_id])
  TinyNN.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.
  TinyNN.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 GPT2KVFFICache. Counterpart to GPT2FFI.upload_from for the KV cache variant. Note: pos_embed is uploaded in FULL (all context_length rows), not sliced.



290
291
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
 
# File 'lib/toy/llm/engine/gpt2_kv_engine.rb', line 290

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

  TinyNN.upload_row_major(sess, kv_cache.t_token_embed, model.token_embed)
  TinyNN.upload_row_major(sess, kv_cache.t_pos_embed,   model.pos_embed)
  TinyNN.tnn_upload_from_float_array(sess, kv_cache.t_ln_f_gamma, model.ln_f_gamma, d_model)
  TinyNN.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]

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

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

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

    li = li + 1
  end
end