Module: Toy::LLM::Primitives::RoPE

Defined in:
lib/toy/llm/primitives/rope.rb,
lib/toy/llm/primitives/rope_cuda.rb,
lib/toy/llm/primitives/rope_metal.rb

Defined Under Namespace

Classes: Cfg

Constant Summary collapse

NAME = 
:rope

Class Method Summary collapse

Class Method Details

.apply_2d(sess, t_pre, positions, freq_factors, cfg, t_seq, t_batch) ⇒ Object

Shape-lift + rope + un-lift, identical to both seq-forward call sites (K path + Q path). t_pre is ne=[d_head, T*B] (ne==1); ggml_rope_ext requires a->ne == positions-> ne, so lift to ne=[d_head, 1, T*B], rope, then reshape back to 2D. Reshape is metadata-only on contiguous tensors; at T=1,B=1 it is a no-op (1 == 1). Returns the rotated 2D tensor handle.



55
56
57
58
59
60
61
62
63
64
 
# File 'lib/toy/llm/primitives/rope.rb', line 55

def self.apply_2d(sess, t_pre, positions, freq_factors, cfg, t_seq, t_batch)
  tb = t_seq * t_batch
  t_pre3 = TinyNN.tnn_reshape_3d(sess, t_pre, cfg.d_head, 1, tb)
  t3     = TinyNN.tnn_rope_ext(sess, t_pre3, positions,
                               cfg.d_head, cfg.base,
                               cfg.freq_scale, cfg.ext_factor,
                               cfg.attn_factor, cfg.beta_fast,
                               cfg.beta_slow, freq_factors)
  TinyNN.tnn_reshape_2d(sess, t3, cfg.d_head, tb)
end