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.



335
336
337
338
339
340
341
342
343
344
 
# File 'lib/toy/llm/engine/gpt2_kv_engine.rb', line 335

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.



279
280
281
282
283
284
285
286
287
288
289
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
 
# File 'lib/toy/llm/engine/gpt2_kv_engine.rb', line 279

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