fix: handle 3D KV tensors in prefix cache for Qwen3.5 (PR waybarrios#144)

Qwen3.5 produces 3D tensors (n_kv_heads, seq_len, head_dim) instead of 4D.

Author: seanpianka

vllm_mlx/prefix_cache.py

@@ -652,9 +652,13 @@ def _extract_block_tensor_slice(
 
                 keys, values = layer_state["state"]
 
-                # KV cache shape: (batch, n_kv_heads, seq_len, head_dim)
-                # Slice along seq_len dimension (axis 2)
-                seq_len = keys.shape[2] if hasattr(keys, "shape") else 0
+                # KV cache shape varies by model architecture:
+                #   4D: (batch, n_kv_heads, seq_len, head_dim) — most models
+                #   3D: (n_kv_heads, seq_len, head_dim)        — e.g. Qwen3.5
+                # Determine the sequence dimension dynamically
+                ndim = keys.ndim if hasattr(keys, "ndim") else len(keys.shape)
+                seq_axis = ndim - 2  # seq_len is always second-to-last
+                seq_len = keys.shape[seq_axis] if hasattr(keys, "shape") else 0
 
                 if end_idx > seq_len:
                     # Requested range extends beyond available data
@@ -665,11 +669,20 @@ def _extract_block_tensor_slice(
                     actual_end = min(end_idx, seq_len)
                     if start_idx >= actual_end:
                         continue
-                    keys_slice = keys[:, :, start_idx:actual_end, :]
-                    values_slice = values[:, :, start_idx:actual_end, :]
+                    # Build a dynamic slice that works for both 3D and 4D tensors
+                    slices = tuple(
+                        slice(start_idx, actual_end) if i == seq_axis else slice(None)
+                        for i in range(ndim)
+                    )
+                    keys_slice = keys[slices]
+                    values_slice = values[slices]
                 else:
-                    keys_slice = keys[:, :, start_idx:end_idx, :]
-                    values_slice = values[:, :, start_idx:end_idx, :]
+                    slices = tuple(
+                        slice(start_idx, end_idx) if i == seq_axis else slice(None)
+                        for i in range(ndim)
+                    )
+                    keys_slice = keys[slices]
+                    values_slice = values[slices]
 
                 block_slices.append((keys_slice, values_slice))
 
@@ -821,10 +834,12 @@ def reconstruct_cache(
                 if not layer_keys:
                     continue
 
-                # Concatenate along sequence dimension (axis 2)
-                # Shape: (batch, n_kv_heads, seq_len, head_dim)
-                concat_keys = mx.concatenate(layer_keys, axis=2)
-                concat_values = mx.concatenate(layer_values, axis=2)
+                # Concatenate along sequence dimension
+                # Shape varies: 4D (batch, heads, seq, dim) or 3D (heads, seq, dim)
+                ndim = layer_keys[0].ndim if hasattr(layer_keys[0], "ndim") else len(layer_keys[0].shape)
+                seq_axis = ndim - 2  # seq_len is always second-to-last
+                concat_keys = mx.concatenate(layer_keys, axis=seq_axis)
+                concat_values = mx.concatenate(layer_values, axis=seq_axis)
 
                 # Create KVCache object
                 # Try to use mlx_lm's KVCache.from_state if available
@@ -833,7 +848,7 @@ def reconstruct_cache(
 
                     # Create new cache and set its state
                     cache = KVCache()
-                    seq_len = concat_keys.shape[2]
+                    seq_len = concat_keys.shape[seq_axis]
 
                     # Set internal state directly
                     # KVCache stores keys/values and offset
@@ -849,7 +864,8 @@ class SimpleKVCache:
                         def __init__(self, keys, values):
                             self.keys = keys
                             self.values = values
-                            self.offset = keys.shape[2]
+                            ndim = keys.ndim if hasattr(keys, "ndim") else len(keys.shape)
+                            self.offset = keys.shape[ndim - 2]
 
                         @property
                         def state(self):