Skip to content

Commit d2c7873

Browse files
committed
feat: expose sin_cache / cos_cache accessors on nn::RoPE
Add public `sin_cache()` and `cos_cache()` accessors that return the pre-computed rotary embedding tables. Downstream code that wants to apply RoPE without going through `RoPE::forward` (for example, a dispatch that routes the computation to a different operator library) needs these to reproduce the cached scaling — particularly for LongRoPE, whose per-dimension short/long factors are non-trivial to recompute outside this class. No change to behavior; the new accessors expose the existing buffers read-only.
1 parent 3f0a98c commit d2c7873

1 file changed

Lines changed: 5 additions & 0 deletions

File tree

include/infinicore/nn/rope.hpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,11 @@ class RoPE : public Module {
120120
Algo algo() const { return algo_; }
121121
DataType dtype() const { return dtype_; }
122122

123+
// Accessors for the pre-computed sin/cos cache tables.
124+
// Shape: `[max_seq_len, head_dim / 2]`, dtype: `dtype()`.
125+
Tensor sin_cache() const { return sin_cache_; }
126+
Tensor cos_cache() const { return cos_cache_; }
127+
123128
// String representation
124129
std::string extra_repr() const;
125130

0 commit comments

Comments
 (0)