awslabs / awslabs/keys_values

Support dynamic RoPE scaling during inference call

Open
#12 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
18
Forks
8
PR merge metrics
No merged PRs in 30d

Description

**Is your feature request related to a problem? Please describe.**

RoPE scaling by YaRN is supported. But the sequence length needs to be fixed up front. For inference with several turns (tool calling), we may not know the sequence length up front.

**Describe the solution you'd like**

Allow to call YaRN scaling several times during an inference run. This is for inference runs with are stopped (say, for a tool call) and continued. A main difficulty is that KV caches then have to store non-encoded keys, and we need to support encoding of keys on the fly.

**Describe alternatives you've considered**

This is maybe not needed right now, but only when tool-calling inference is used.

Contributor guide

Open the contributing guide

Research direction

Start by examining the existing YaRN RoPE scaling and KV-cache handling in the repository; the issue does not name specific files or tests. Done would mean supporting repeated scaling during a stopped and resumed inference run, including storing non-encoded keys and encoding them on the fly.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.