Fix: Dynamic RoPE Cache Expansion to Prevent Position-ID Out-of-Bounds in EAGLE + Long-Sequence Workloads (#10788)

This commit is contained in:
YAMY
2025-10-19 11:37:43 +08:00
committed by GitHub
parent b288f4f440
commit 80407b0493
4 changed files with 103 additions and 0 deletions
@@ -147,6 +147,36 @@ class RotaryEmbedding(CustomOp):
cache = torch.cat((cos, sin), dim=-1)
return cache
def _ensure_cos_sin_cache_length(self, needed_max_pos: int):
"""Ensure cos_sin_cache length > needed_max_pos."""
cur_len = int(self.cos_sin_cache.shape[0])
if needed_max_pos < cur_len:
return
# Align to 128 to reduce realloc frequency
new_len = ((needed_max_pos + 128) // 128) * 128
device = self.cos_sin_cache.device
dtype = self.cos_sin_cache.dtype
# Compute inv_freq on same device
inv_freq = self._compute_inv_freq(self.base).to(device=device)
# Incremental computation for new positions only
start = cur_len
t_new = torch.arange(start, new_len, dtype=inv_freq.dtype, device=device)
if t_new.numel() == 0:
return
freqs_new = torch.einsum("i,j->ij", t_new, inv_freq)
cos_new = freqs_new.cos()
sin_new = freqs_new.sin()
new_rows = torch.cat((cos_new, sin_new), dim=-1).to(dtype=dtype)
# Update cache with new rows
self.cos_sin_cache = torch.cat((self.cos_sin_cache, new_rows), dim=0).to(
device=device, dtype=dtype
)
def forward_native(
self,
positions: torch.Tensor,