[CRITICAL] batch_quantized_cache extend() silently truncates KV data #67

New issue

Closed

opened 2026-05-20 01:45:28 +02:00 by sleepy · 0 comments

sleepy commented

2026-05-20 01:45:28 +02:00

Owner

Severity: CRITICAL - Data loss during continuous batching.

Location: omlx/cache/batch_quantized_cache.py, extend() method, lines ~180-215

Problem: The pad helper computes right = max_size - c.keys[0].shape[2] - left. When left + c.keys[0].shape[2] exceeds max_size, right becomes negative. The code then truncates actual KV tokens from the end.

Example: Cache A has _idx=100 but buffer shape[2]=200. Cache B has _idx=150. For A: left=50, right=200-200-50=-50. A is trimmed by 50 tokens.

Fix: max_size should account for (left + buffer_size) across both caches.

**Severity:** CRITICAL - Data loss during continuous batching. **Location:** `omlx/cache/batch_quantized_cache.py`, `extend()` method, lines ~180-215 **Problem:** The `pad` helper computes `right = max_size - c.keys[0].shape[2] - left`. When `left + c.keys[0].shape[2]` exceeds `max_size`, `right` becomes negative. The code then truncates actual KV tokens from the end. **Example:** Cache A has `_idx=100` but buffer `shape[2]=200`. Cache B has `_idx=150`. For A: `left=50`, `right=200-200-50=-50`. A is trimmed by 50 tokens. **Fix:** `max_size` should account for `(left + buffer_size)` across both caches.

sleepy referenced this issue from a commit

2026-05-20 01:57:22 +02:00

fix(batch_quantized_cache): prevent extend() from truncating KV data

sleepy referenced this issue

2026-05-20 01:58:11 +02:00

[cache] Fix batch_quantized_cache extend() truncation bug (#67) #78

sleepy referenced this issue from a commit

2026-05-20 02:23:16 +02:00

[cache] Fix batch_quantized_cache extend() truncation bug (#67) (#78)

sleepy closed this issue

2026-05-20 02:28:03 +02:00