diff --git a/src/parallax/metal/paged_attention/kernel.py b/src/parallax/metal/paged_attention/kernel.py
deleted file mode 100644
index 9ee315bb..00000000
--- a/src/parallax/metal/paged_attention/kernel.py
+++ /dev/null
@@ -1,249 +0,0 @@
-import os
-from typing import Dict, List, Optional
-
-import mlx.core as mx
-
-# Cache for compiled kernels
-_KERNELS: Dict[str, object] = {}
-
-
-def _get_metal_source(filename):
-    path = os.path.join(os.path.dirname(__file__), filename)
-    with open(path, "r") as f:
-        return f.read()
-
-
-def _type_to_string(dtype: mx.Dtype) -> str:
-    if dtype == mx.float32:
-        return "float"
-    elif dtype == mx.float16:
-        return "half"
-    elif dtype == mx.bfloat16:
-        # Metal 3.1+ supports bfloat, typically via bfloat16_t or using half
-        # For now we map to bfloat16_t assuming compiler support
-        return "bfloat16_t"
-    else:
-        raise ValueError(f"Unsupported dtype for paged attention: {dtype}")
-
-
-def _get_kernel(
-    name: str,
-    filename: str,
-    input_names: List[str],
-    output_names: List[str],
-    dtype: mx.Dtype = mx.float32,
-):
-    type_str = _type_to_string(dtype)
-    kernel_key = f"{name}_{type_str}"
-
-    if kernel_key not in _KERNELS:
-        source = _get_metal_source(filename)
-        # Simple template substitution
-        source = source.replace("{{T}}", type_str)
-
-        header = """
-#include <metal_stdlib>
-using namespace metal;
-"""
-        _KERNELS[kernel_key] = mx.fast.metal_kernel(
-            name=name,  # Internal name for MLX JIT cache (not used for dispatch if we hold the object)
-            input_names=input_names,
-            output_names=output_names,
-            source=source,
-            header=header,
-        )
-    return _KERNELS[kernel_key]
-
-
-def paged_attention(
-    queries: mx.array,
-    key_cache: mx.array,
-    value_cache: mx.array,
-    block_tables: mx.array,
-    context_lengths: mx.array,
-    block_size: int,
-    scale: float,
-    num_kv_heads: int,
-    v_head_dim: Optional[int] = None,
-    top_k_indices: Optional[mx.array] = None,
-    window_size: Optional[int] = None,
-    sinks: Optional[mx.array] = None,
-) -> mx.array:
-    """
-    Paged Attention using Metal Kernel.
-    """
-    batch_size = queries.shape[0]
-    num_heads = queries.shape[1]
-    dtype = queries.dtype
-
-    if queries.ndim == 4:
-        if queries.shape[2] != 1:
-            pass
-        queries = queries.squeeze(2)
-
-    k_head_dim = queries.shape[2]
-    if v_head_dim is None:
-        v_head_dim = k_head_dim
-
-    num_layers = key_cache.shape[0]
-    num_total_blocks = key_cache.shape[1]
-    max_blocks = block_tables.shape[1]
-
-    # Prepare Constants
-    def mk_int(val):
-        return mx.array(val, dtype=mx.int32)
-
-    c_num_heads = mk_int(num_heads)
-    c_num_kv_heads = mk_int(num_kv_heads)
-    c_k_head_dim = mk_int(k_head_dim)
-    c_v_head_dim = mk_int(v_head_dim)
-    c_block_size = mk_int(block_size)
-    c_max_blocks = mk_int(max_blocks)
-    c_num_layers = mk_int(num_layers)
-    c_num_total_blocks = mk_int(num_total_blocks)
-    c_scale = mx.array(scale, dtype=queries.dtype)
-
-    if top_k_indices is not None:
-        index_topk = top_k_indices.shape[1]
-        c_index_topk = mk_int(index_topk)
-
-        inputs = [
-            queries,
-            key_cache,
-            value_cache,
-            block_tables,
-            context_lengths,
-            top_k_indices,
-            c_num_heads,
-            c_num_kv_heads,
-            c_k_head_dim,
-            c_v_head_dim,
-            c_block_size,
-            c_max_blocks,
-            c_num_layers,
-            c_num_total_blocks,
-            c_scale,
-            c_index_topk,
-        ]
-
-        input_names = [
-            "queries",
-            "key_cache",
-            "value_cache",
-            "block_tables",
-            "context_lengths",
-            "top_k_indices",
-            "num_heads",
-            "num_kv_heads",
-            "k_head_dim",
-            "v_head_dim",
-            "block_size",
-            "max_blocks",
-            "num_layers",
-            "num_total_blocks",
-            "scale",
-            "index_topk",
-        ]
-        kernel_name = "paged_attention_deepseek_v32_kernel"
-        filename = "paged_attention_deepseek_v32.metal"
-    elif window_size is not None and sinks is not None:
-        c_window_size = mk_int(window_size)
-        c_sinks = sinks
-        inputs = [
-            queries,
-            key_cache,
-            value_cache,
-            block_tables,
-            context_lengths,
-            c_sinks,
-            c_num_heads,
-            c_num_kv_heads,
-            c_k_head_dim,
-            c_v_head_dim,
-            c_block_size,
-            c_max_blocks,
-            c_num_layers,
-            c_num_total_blocks,
-            c_scale,
-            c_window_size,
-        ]
-
-        input_names = [
-            "queries",
-            "key_cache",
-            "value_cache",
-            "block_tables",
-            "context_lengths",
-            "sinks",
-            "num_heads",
-            "num_kv_heads",
-            "k_head_dim",
-            "v_head_dim",
-            "block_size",
-            "max_blocks",
-            "num_layers",
-            "num_total_blocks",
-            "scale",
-            "window_size",
-        ]
-        kernel_name = "paged_attention_gpt_oss_kernel"
-        filename = "paged_attention_gpt_oss.metal"
-    else:
-        inputs = [
-            queries,
-            key_cache,
-            value_cache,
-            block_tables,
-            context_lengths,
-            c_num_heads,
-            c_num_kv_heads,
-            c_k_head_dim,
-            c_v_head_dim,
-            c_block_size,
-            c_max_blocks,
-            c_num_layers,
-            c_num_total_blocks,
-            c_scale,
-        ]
-
-        input_names = [
-            "queries",
-            "key_cache",
-            "value_cache",
-            "block_tables",
-            "context_lengths",
-            "num_heads",
-            "num_kv_heads",
-            "k_head_dim",
-            "v_head_dim",
-            "block_size",
-            "max_blocks",
-            "num_layers",
-            "num_total_blocks",
-            "scale",
-        ]
-        kernel_name = "paged_attention_kernel"
-        filename = "paged_attention.metal"
-
-    kernel = _get_kernel(
-        name=kernel_name,
-        filename=filename,
-        input_names=input_names,
-        output_names=["output"],
-        dtype=dtype,  # This will generate paged_attention_kernel_half etc.
-    )
-
-    grid = (num_heads * 32, batch_size, 1)
-    thread_group = (32, 1, 1)
-
-    outputs = kernel(
-        inputs=inputs,
-        grid=grid,
-        threadgroup=thread_group,
-        output_shapes=[(batch_size, num_heads, v_head_dim)],
-        output_dtypes=[dtype],  # Output matches input dtype
-        verbose=False,
-    )
-
-    out = outputs[0]
-    return out[:, :, None, :]
diff --git a/src/parallax/metal/paged_attention/paged_attention.metal b/src/parallax/metal/paged_attention/paged_attention.metal
deleted file mode 100644
index 6780a411..00000000
--- a/src/parallax/metal/paged_attention/paged_attention.metal
+++ /dev/null
@@ -1,127 +0,0 @@
-
-// Inputs:
-// queries, key_cache, value_cache, block_tables, context_lengths
-// output (output array)
-// num_heads, num_kv_heads, k_head_dim, v_head_dim, block_size, max_blocks,
-// num_layers, num_total_blocks, scale (All pointers)
-
-uint3 gid = thread_position_in_grid;
-uint3 tid = thread_position_in_threadgroup;
-
-// Each threadgroup handles one head.
-// Assuming threadgroup size is 32 (SIMD width).
-// head_idx comes from the group index.
-// Or we can calculate it from gid if grid is linear.
-// grid.x = num_heads * 32.
-
-int head_idx = gid.x / 32;
-int batch_idx = gid.y;
-
-// Dereference constants
-int _num_heads = num_heads;
-int _num_kv_heads = num_kv_heads;
-int _k_head_dim = k_head_dim;
-int _v_head_dim = v_head_dim;
-int _block_size = block_size;
-int _max_blocks = max_blocks;
-int _num_total_blocks = num_total_blocks;
-float _scale = scale;
-
-if (head_idx >= _num_heads)
-  return;
-
-int kv_head_idx = head_idx / (_num_heads / _num_kv_heads);
-
-// Load Query
-// Q: [batch, num_heads, k_head_dim]
-// Thread i loads elements i, i+32, ...
-
-// Support up to 256 head dim (8 * 32)
-float q_vec[8] = {0.0f};
-
-int q_offset = batch_idx * _num_heads * _k_head_dim + head_idx * _k_head_dim;
-
-for (int i = tid.x; i < _k_head_dim; i += 32) {
-  if (i < 256) {
-    q_vec[i / 32] = queries[q_offset + i];
-  }
-}
-
-// Running statistics for Softmax
-float m_i = -INFINITY;
-float l_i = 0.0f;
-float acc_vec[8] = {0.0f};
-
-int context_len = context_lengths[batch_idx];
-int num_context_blocks = (context_len + _block_size - 1) / _block_size;
-
-// Strides for Key
-long k_block_stride = _num_kv_heads * _block_size * _k_head_dim;
-long k_head_stride = _block_size * _k_head_dim;
-
-// Strides for Value
-long v_block_stride = _num_kv_heads * _block_size * _v_head_dim;
-long v_head_stride = _block_size * _v_head_dim;
-
-// Iterate over blocks
-for (int b = 0; b < num_context_blocks; b++) {
-  int block_num = block_tables[batch_idx * _max_blocks + b];
-
-  long k_block_base =
-      block_num * k_block_stride + kv_head_idx * k_head_stride;
-  long v_block_base =
-      block_num * v_block_stride + kv_head_idx * v_head_stride;
-
-  int tokens_in_block = _block_size;
-  if (b == num_context_blocks - 1) {
-    tokens_in_block = context_len % _block_size;
-    if (tokens_in_block == 0)
-      tokens_in_block = _block_size;
-  }
-
-  for (int t = 0; t < tokens_in_block; t++) {
-    // Compute Dot Product Q * K[t]
-    float score = 0.0f;
-    for (int i = tid.x; i < _k_head_dim; i += 32) {
-      // offset inside block: t * k_head_dim + i
-      float k_val = key_cache[k_block_base + t * _k_head_dim + i];
-
-      if (i < 256) {
-        score += q_vec[i / 32] * k_val;
-      }
-    }
-
-    // SIMD Reduction for score
-    score = simd_sum(score);
-    score *= _scale;
-
-    // Softmax update
-    float m_prev = m_i;
-    m_i = max(m_prev, score);
-    float alpha = exp(m_prev - m_i);
-    float beta = exp(score - m_i);
-
-    l_i = l_i * alpha + beta;
-
-    // Accumulate V
-    for (int i = tid.x; i < _v_head_dim; i += 32) {
-      float v_val = value_cache[v_block_base + t * _v_head_dim + i];
-      if (i < 256) {
-        acc_vec[i / 32] = acc_vec[i / 32] * alpha + v_val * beta;
-      }
-    }
-  }
-}
-
-// Finalize Output
-for (int i = 0; i < 8; i++) {
-  acc_vec[i] /= l_i;
-}
-
-int out_offset = batch_idx * _num_heads * _v_head_dim + head_idx * _v_head_dim;
-
-for (int i = tid.x; i < _v_head_dim; i += 32) {
-  if (i < 256) {
-    output[out_offset + i] = ({{T}})acc_vec[i / 32];
-  }
-}
diff --git a/src/parallax/metal/paged_attention/paged_attention_deepseek_v32.metal b/src/parallax/metal/paged_attention/paged_attention_deepseek_v32.metal
deleted file mode 100644
index 900fa104..00000000
--- a/src/parallax/metal/paged_attention/paged_attention_deepseek_v32.metal
+++ /dev/null
@@ -1,180 +0,0 @@
-// Inputs:
-// queries, key_cache, value_cache, block_tables, context_lengths, top_k_indices
-// output (output array)
-// num_heads, num_kv_heads, k_head_dim, v_head_dim, block_size, max_blocks,
-// num_layers, num_total_blocks, scale, index_topk (All pointers)
-
-uint3 gid = thread_position_in_grid;
-uint3 tid = thread_position_in_threadgroup;
-
-// Each threadgroup handles one head.
-// Assuming threadgroup size is 32 (SIMD width).
-// head_idx comes from the group index.
-// Or we can calculate it from gid if grid is linear.
-// grid.x = num_heads * 32.
-
-int head_idx = gid.x / 32;
-int batch_idx = gid.y;
-
-// Dereference constants
-int _num_heads = num_heads;
-int _num_kv_heads = num_kv_heads;
-int _k_head_dim = k_head_dim;
-int _v_head_dim = v_head_dim;
-int _block_size = block_size;
-int _max_blocks = max_blocks;
-int _num_total_blocks = num_total_blocks;
-float _scale = scale;
-int _index_topk = index_topk;
-
-if (head_idx >= _num_heads)
-  return;
-
-int kv_head_idx = head_idx / (_num_heads / _num_kv_heads);
-
-// Load Query
-// Q: [batch, num_heads, k_head_dim]
-// Thread i loads elements i, i+32, ...
-
-// Support up to 256 head dim (8 * 32)
-float q_vec[8] = {0.0f};
-
-int q_offset = batch_idx * _num_heads * _k_head_dim + head_idx * _k_head_dim;
-
-for (int i = tid.x; i < _k_head_dim; i += 32) {
-  if (i < 256) {
-    q_vec[i / 32] = queries[q_offset + i];
-  }
-}
-
-// Running statistics for Softmax
-float m_i = -INFINITY;
-float l_i = 0.0f;
-float acc_vec[8] = {0.0f};
-
-int context_len = context_lengths[batch_idx];
-
-// Strides for Key
-long k_block_stride = _num_kv_heads * _block_size * _k_head_dim;
-long k_head_stride = _block_size * _k_head_dim;
-
-// Strides for Value
-long v_block_stride = _num_kv_heads * _block_size * _v_head_dim;
-long v_head_stride = _block_size * _v_head_dim;
-
-// Check if we use Full Attention or Sparse Attention
-// We check the first element of top_k_indices for this batch
-int first_topk_idx = top_k_indices[batch_idx * _index_topk];
-
-if (first_topk_idx == -1) {
-    // -------------------------------------------------------------------------
-    // Full Attention (Standard Paged Attention)
-    // -------------------------------------------------------------------------
-    int num_context_blocks = (context_len + _block_size - 1) / _block_size;
-
-    for (int b = 0; b < num_context_blocks; b++) {
-      int block_num = block_tables[batch_idx * _max_blocks + b];
-
-      long k_block_base =
-          block_num * k_block_stride + kv_head_idx * k_head_stride;
-      long v_block_base =
-          block_num * v_block_stride + kv_head_idx * v_head_stride;
-
-      int tokens_in_block = _block_size;
-      if (b == num_context_blocks - 1) {
-        tokens_in_block = context_len % _block_size;
-        if (tokens_in_block == 0)
-          tokens_in_block = _block_size;
-      }
-
-      for (int t = 0; t < tokens_in_block; t++) {
-        // Compute Dot Product Q * K[t]
-        float score = 0.0f;
-        for (int i = tid.x; i < _k_head_dim; i += 32) {
-          float k_val = key_cache[k_block_base + t * _k_head_dim + i];
-          if (i < 256) {
-            score += q_vec[i / 32] * k_val;
-          }
-        }
-
-        // SIMD Reduction for score
-        score = simd_sum(score);
-        score *= _scale;
-
-        // Softmax update
-        float m_prev = m_i;
-        m_i = max(m_prev, score);
-        float alpha = exp(m_prev - m_i);
-        float beta = exp(score - m_i);
-
-        l_i = l_i * alpha + beta;
-
-        // Accumulate V
-        for (int i = tid.x; i < _v_head_dim; i += 32) {
-          float v_val = value_cache[v_block_base + t * _v_head_dim + i];
-          if (i < 256) {
-            acc_vec[i / 32] = acc_vec[i / 32] * alpha + v_val * beta;
-          }
-        }
-      }
-    }
-} else {
-    // -------------------------------------------------------------------------
-    // Sparse Attention (Top-K Indices)
-    // -------------------------------------------------------------------------
-    for (int k = 0; k < _index_topk; k++) {
-        int token_idx = top_k_indices[batch_idx * _index_topk + k];
-
-        // Map logical token index to physical location
-        int block_idx_in_table = token_idx / _block_size;
-        int block_offset = token_idx % _block_size;
-        int block_num = block_tables[batch_idx * _max_blocks + block_idx_in_table];
-
-        long k_block_base =
-            block_num * k_block_stride + kv_head_idx * k_head_stride;
-        long v_block_base =
-            block_num * v_block_stride + kv_head_idx * v_head_stride;
-
-        // Compute Dot Product Q * K[token_idx]
-        float score = 0.0f;
-        for (int i = tid.x; i < _k_head_dim; i += 32) {
-            float k_val = key_cache[k_block_base + block_offset * _k_head_dim + i];
-            if (i < 256) {
-                score += q_vec[i / 32] * k_val;
-            }
-        }
-
-        // SIMD Reduction for score
-        score = simd_sum(score);
-        score *= _scale;
-
-        // Softmax update
-        float m_prev = m_i;
-        m_i = max(m_prev, score);
-        float alpha = exp(m_prev - m_i);
-        float beta = exp(score - m_i);
-
-        l_i = l_i * alpha + beta;
-
-        // Accumulate V
-        for (int i = tid.x; i < _v_head_dim; i += 32) {
-            float v_val = value_cache[v_block_base + block_offset * _v_head_dim + i];
-            if (i < 256) {
-                acc_vec[i / 32] = acc_vec[i / 32] * alpha + v_val * beta;
-            }
-        }
-    }
-}
-
-// Finalize Output
-for (int i = 0; i < 8; i++) {
-  acc_vec[i] /= l_i;
-}
-
-int out_offset = batch_idx * _num_heads * _v_head_dim + head_idx * _v_head_dim;
-
-for (int i = tid.x; i < _v_head_dim; i += 32) {
-  if (i < 256) {
-    output[out_offset + i] = ({{T}})acc_vec[i / 32];
-  }
-}
diff --git a/src/parallax/metal/paged_attention/paged_attention_gpt_oss.metal b/src/parallax/metal/paged_attention/paged_attention_gpt_oss.metal
deleted file mode 100644
index 4a63d62a..00000000
--- a/src/parallax/metal/paged_attention/paged_attention_gpt_oss.metal
+++ /dev/null
@@ -1,177 +0,0 @@
-
-// Inputs:
-// queries, key_cache, value_cache, block_tables, context_lengths, sinks
-// output (output array)
-// num_heads, num_kv_heads, k_head_dim, v_head_dim, block_size, max_blocks,
-// num_layers, num_total_blocks, scale, window_size (All pointers)
-
-uint3 gid = thread_position_in_grid;
-uint3 tid = thread_position_in_threadgroup;
-
-// Each threadgroup handles one head.
-// Assuming threadgroup size is 32 (SIMD width).
-// head_idx comes from the group index.
-// Or we can calculate it from gid if grid is linear.
-// grid.x = num_heads * 32.
-
-int head_idx = gid.x / 32;
-int batch_idx = gid.y;
-
-// Dereference constants
-int _num_heads = num_heads;
-int _num_kv_heads = num_kv_heads;
-int _k_head_dim = k_head_dim;
-int _v_head_dim = v_head_dim;
-int _block_size = block_size;
-int _max_blocks = max_blocks;
-int _num_total_blocks = num_total_blocks;
-float _scale = scale;
-int _window_size = window_size;
-
-if (head_idx >= _num_heads)
-  return;
-
-int kv_head_idx = head_idx / (_num_heads / _num_kv_heads);
-
-// Load Query
-// Q: [batch, num_heads, k_head_dim]
-// Thread i loads elements i, i+32, ...
-
-// Support up to 256 head dim (8 * 32)
-float q_vec[8] = {0.0f};
-
-int q_offset = batch_idx * _num_heads * _k_head_dim + head_idx * _k_head_dim;
-
-for (int i = tid.x; i < _k_head_dim; i += 32) {
-  if (i < 256) {
-    q_vec[i / 32] = queries[q_offset + i];
-  }
-}
-
-// Running statistics for Softmax
-float m_i = -INFINITY;
-float l_i = 0.0f;
-float acc_vec[8] = {0.0f};
-
-// -------------------------------------------------------------------------
-// Handle Implicit Sink Token (KV = 0, Not in Block Table)
-// -------------------------------------------------------------------------
-// The sink token is assumed to be an extra token (logical pos 0? or just extra)
-// that is NOT stored in the paged KV blocks.
-// Its Key is 0, Value is 0.
-// Its Attention Score = (Q * K_sink) * scale + sinks_bias
-//                     = (Q * 0) * scale + sinks[head_idx]
-//                     = sinks[head_idx]
-
-float sink_score = sinks[head_idx];
-
-// Initialize Softmax stats with Sink Token if it's not masked out (-inf)
-// We use a threshold slightly larger than -INFINITY to avoid precision issues
-if (sink_score > -3.40282e+38f) { // roughly > -FLT_MAX
-    m_i = sink_score;
-    l_i = 1.0f;
-    // acc_vec remains 0.0f because V_sink is 0.
-}
-
-int context_len = context_lengths[batch_idx];
-int num_context_blocks = (context_len + _block_size - 1) / _block_size;
-
-// Strides for Key
-long k_block_stride = _num_kv_heads * _block_size * _k_head_dim;
-long k_head_stride = _block_size * _k_head_dim;
-
-// Strides for Value
-long v_block_stride = _num_kv_heads * _block_size * _v_head_dim;
-long v_head_stride = _block_size * _v_head_dim;
-
-// Iterate over blocks
-for (int b = 0; b < num_context_blocks; b++) {
-  int block_num = block_tables[batch_idx * _max_blocks + b];
-
-  // Calculate logic position range for this block
-  int block_start_pos = b * _block_size;
-  int block_end_pos = block_start_pos + _block_size; // exclusive
-
-  // Handle Window Attention Skip
-  // Current query position is at (context_len - 1)
-  // Window range is [context_len - 1 - window_size, context_len - 1]
-
-  // NOTE: Sink token is handled separately above, so we don't need to protect block 0
-  // unless block 0 also contains other vital tokens (it usually does).
-
-  if (_window_size > 0) {
-      int window_start = context_len - 1 - _window_size;
-      if (block_end_pos <= window_start) {
-          // Entire block is out of window
-          continue;
-      }
-  }
-
-  long k_block_base =
-      block_num * k_block_stride + kv_head_idx * k_head_stride;
-  long v_block_base =
-      block_num * v_block_stride + kv_head_idx * v_head_stride;
-
-  int tokens_in_block = _block_size;
-  if (b == num_context_blocks - 1) {
-    tokens_in_block = context_len % _block_size;
-    if (tokens_in_block == 0)
-      tokens_in_block = _block_size;
-  }
-
-  for (int t = 0; t < tokens_in_block; t++) {
-    int token_pos = block_start_pos + t;
-
-    // Check if token is within window
-    if (_window_size > 0) {
-        int window_start = context_len - 1 - _window_size;
-        if (token_pos < window_start) {
-            continue;
-        }
-    }
-
-    // Compute Dot Product Q * K[t]
-    float score = 0.0f;
-    for (int i = tid.x; i < _k_head_dim; i += 32) {
-      // offset inside block: t * k_head_dim + i
-      float k_val = key_cache[k_block_base + t * _k_head_dim + i];
-
-      if (i < 256) {
-        score += q_vec[i / 32] * k_val;
-      }
-    }
-
-    // SIMD Reduction for score
-    score = simd_sum(score);
-    score *= _scale;
-
-    // Softmax update
-    float m_prev = m_i;
-    m_i = max(m_prev, score);
-    float alpha = exp(m_prev - m_i);
-    float beta = exp(score - m_i);
-
-    l_i = l_i * alpha + beta;
-
-    // Accumulate V
-    for (int i = tid.x; i < _v_head_dim; i += 32) {
-      float v_val = value_cache[v_block_base + t * _v_head_dim + i];
-      if (i < 256) {
-        acc_vec[i / 32] = acc_vec[i / 32] * alpha + v_val * beta;
-      }
-    }
-  }
-}
-
-// Finalize Output
-for (int i = 0; i < 8; i++) {
-  acc_vec[i] /= l_i;
-}
-
-int out_offset = batch_idx * _num_heads * _v_head_dim + head_idx * _v_head_dim;
-
-for (int i = tid.x; i < _v_head_dim; i += 32) {
-  if (i < 256) {
-    output[out_offset + i] = ({{T}})acc_vec[i / 32];
-  }
-}
diff --git a/src/parallax/models/deepseek_v2.py b/src/parallax/models/deepseek_v2.py
deleted file mode 100644
index b8b4774b..00000000
--- a/src/parallax/models/deepseek_v2.py
+++ /dev/null
@@ -1,200 +0,0 @@
-"""
-hidden_dimefines the Qwen3 model.
-"""
-
-from typing import Any, List, Optional
-
-from parallax_utils.logging_config import get_logger
-
-logger = get_logger(__name__)
-
-import mlx.core as mx
-from mlx_lm.models.base import scaled_dot_product_attention
-from mlx_lm.models.deepseek_v2 import DeepseekV2Attention as MLXDeepseekV2Attention
-from mlx_lm.models.deepseek_v2 import DeepseekV2DecoderLayer as MLXDeepseekV2Block
-from mlx_lm.models.deepseek_v2 import ModelArgs
-
-from parallax.metal.paged_attention.kernel import paged_attention
-from parallax.server.cache.base import BaseCache
-from parallax.utils.prefix_cache_utils import compute_attention_with_prefix_cache
-from parallax_extensions.ops import reshape_and_cache
-
-
-class ParallaxDeepSeekV2Attention(MLXDeepseekV2Attention):
-    """A custom attention module for Parallax, extending the DeepseekV2 Attention class.
-
-    We apply explicit KV cache handling and passing in `offset` directly from Request.
-    This version returns the new K and V states for external caching.
-    """
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[BaseCache] = None,
-        block_tables: Optional[mx.array] = None,
-        context_lengths: Optional[mx.array] = None,
-        slot_mapping: Optional[mx.array] = None,
-        prefix_lens: Optional[mx.array] = None,
-        **kwargs,
-    ) -> mx.array:
-        """
-        Attention forward pass with explicit KV cache handling.
-
-        Args:
-            x: (batch, target_len, hidden_dim) - Input hidden states for the current query segment.
-            mask: (batch, n_q_heads, target_len, source_len)
-            cache: BaseCache object containing the layer cache.
-            block_tables: (batch, max_blocks) - PagedKV block tables.
-            context_lengths: (batch,) - PagedKV sequence lengths.
-            slot_mapping: (batch * target_len,) - Flattened slot mapping.
-            prefix_lens: (batch,) - Number of prefix tokens already cached (for RoPE offset).
-
-        Returns:
-            output_h: (batch, target_len, hidden_dim) - Output hidden states.
-        """
-        batch, target_len, _ = x.shape
-        if self.q_lora_rank is None:
-            q = self.q_proj(x)
-        else:
-            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
-
-        q = q.reshape(batch, target_len, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
-        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
-        compressed_kv = self.kv_a_proj_with_mqa(x)
-        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
-        k_pe = k_pe.reshape(batch, target_len, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-        kv = kv.reshape(batch, target_len, self.num_heads, -1)
-
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
-        k_nope = k_nope.transpose(0, 2, 1, 3)
-
-        key_cache_global, value_cache_global = cache.get_cache()
-
-        # Compute RoPE offsets for all batch items at once
-        if target_len == 1:
-            # Decode phase: position is context_length - 1
-            current_pos = context_lengths - 1
-        elif prefix_lens is not None:
-            # Prefill phase - start from prefix_len if using prefix cache
-            current_pos = prefix_lens
-        else:
-            # Prefill phase - no prefix cache
-            current_pos = 0
-
-        # Apply RoPE to all batch items at once
-        q_pe = self.rope(q_pe, offset=current_pos)
-        k_pe = self.rope(k_pe, offset=current_pos)
-
-        k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
-
-        # Construct full keys
-        keys = mx.concatenate([k_nope, k_pe], axis=-1)
-
-        block_size = key_cache_global.shape[3]
-
-        reshape_and_cache(
-            keys.transpose(0, 2, 1, 3),
-            values,
-            key_cache_global,
-            value_cache_global,
-            block_tables,
-            context_lengths,
-            block_size,
-            slot_mapping=slot_mapping,
-        )
-
-        if target_len == 1:
-            # Decode Phase: Use Paged Attention Kernel
-            output = paged_attention(
-                queries,
-                key_cache_global,
-                value_cache_global,
-                block_tables,
-                context_lengths,
-                block_size,
-                self.scale,
-                self.num_heads,  # num_kv_heads (MQA/MLA, here num_heads == num_kv_heads effectively after repeat?)
-                v_head_dim=values.shape[-1],
-            )
-            output = output.transpose(0, 2, 1, 3).reshape(batch, target_len, -1)
-        else:
-            # Prefill Phase: Need to attend to both cached prefix and new tokens
-            # Check if any request has prefix cache
-            has_prefix_cache = prefix_lens is not None and bool(mx.any(prefix_lens > 0))
-
-            if has_prefix_cache:
-                # Use shared prefix cache handling with batch processing
-                k_new = keys  # (batch, num_heads, target_len, head_dim)
-                v_new = values.transpose(0, 2, 1, 3)  # (batch, num_heads, target_len, head_dim)
-                output = compute_attention_with_prefix_cache(
-                    queries,  # (batch, num_heads, target_len, head_dim)
-                    k_new,
-                    v_new,
-                    cache,
-                    block_tables,
-                    prefix_lens,
-                    target_len,
-                    self.scale,
-                    self.num_heads,  # num_kv_heads (in DeepSeekV2, num_heads == num_kv_heads after repeat)
-                    mask=mask,
-                )
-            else:
-                # No prefix cache, use standard self-attention on local data only
-                output = scaled_dot_product_attention(
-                    queries,
-                    keys,
-                    values.transpose(0, 2, 1, 3),
-                    scale=self.scale,
-                    mask=mask,
-                    cache=None,
-                )
-                output = output.transpose(0, 2, 1, 3).reshape(batch, target_len, -1)
-
-        return self.o_proj(output)
-
-
-class ParallaxDeepSeekV2Block(MLXDeepseekV2Block):
-    """A custom transformer block for Parallax, extending the Qwen3 Block class.
-    This version handles the KV cache explicitly and returns new K and V states.
-    """
-
-    def __init__(self, args: ModelArgs, layer_idx: int, local_layer_idx: int):
-        super().__init__(args, layer_idx=layer_idx)
-        self.self_attn = ParallaxDeepSeekV2Attention(args)
-        self.layer_idx = layer_idx
-        self.local_layer_idx = local_layer_idx
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[List[Any]] = None,
-        block_tables: Optional[mx.array] = None,
-        context_lengths: Optional[mx.array] = None,
-        slot_mapping: Optional[mx.array] = None,
-        **kwargs,
-    ):
-        r = self.self_attn(
-            self.input_layernorm(x),
-            mask,
-            cache[self.local_layer_idx],
-            block_tables=block_tables,
-            context_lengths=context_lengths,
-            slot_mapping=slot_mapping,
-            **kwargs,
-        )
-        h = x + r
-        r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r
-        return out
-
-    @classmethod
-    def get_architecture(cls):
-        """Get the architecture name for the block."""
-        return "DeepseekV2ForCausalLM"
-
-
-EntryClass = ParallaxDeepSeekV2Block
diff --git a/src/parallax/models/deepseek_v3.py b/src/parallax/models/deepseek_v3.py
index 2d34b2d1..2a4ccfbb 100644
--- a/src/parallax/models/deepseek_v3.py
+++ b/src/parallax/models/deepseek_v3.py
@@ -14,10 +14,9 @@
 from mlx_lm.models.deepseek_v3 import DeepseekV3DecoderLayer as MLXDeepseekV3Block
 from mlx_lm.models.deepseek_v3 import ModelArgs
 
-from parallax.metal.paged_attention.kernel import paged_attention
 from parallax.server.cache.base import BaseCache
-from parallax.utils.prefix_cache_utils import compute_attention_with_prefix_cache
-from parallax_extensions.ops import reshape_and_cache
+from parallax.server.cache.dsa_cache import DeepSeekSparseCache
+from parallax_extensions.ops import mla_paged_attention, reshape_and_cache
 
 
 class ParallaxDeepSeekV3Attention(MLXDeepseekV3Attention):
@@ -27,6 +26,62 @@ class ParallaxDeepSeekV3Attention(MLXDeepseekV3Attention):
     This version returns the new K and V states for external caching.
     """
 
+    def _read_mla_batch(
+        self,
+        cache: DeepSeekSparseCache,
+        block_tables: mx.array,
+        lengths: mx.array,
+        max_len: int,
+        dtype: mx.Dtype,
+    ) -> tuple[mx.array, mx.array]:
+        batch = lengths.shape[0]
+        latent = mx.zeros((batch, 1, max_len, self.kv_lora_rank), dtype=dtype)
+        rope = mx.zeros((batch, 1, max_len, self.qk_rope_head_dim), dtype=dtype)
+
+        for i in range(batch):
+            length = int(lengths[i])
+            if length <= 0:
+                continue
+            latent_i, rope_i = cache.read_prefix_mla(block_tables[i], length)
+            latent[i, :, :length, :] = latent_i
+            rope[i, :, :length, :] = rope_i
+        return latent, rope
+
+    def _mla_attention(
+        self,
+        q_nope: mx.array,
+        q_pe: mx.array,
+        kv_latent: mx.array,
+        k_pe: mx.array,
+        mask: Optional[mx.array],
+    ) -> mx.array:
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = pe_scores + mask.astype(pe_scores.dtype)
+
+        if q_nope.shape[2] == 1:
+            q_latent = self.embed_q(q_nope)
+            output = scaled_dot_product_attention(
+                q_latent,
+                kv_latent,
+                kv_latent,
+                scale=self.scale,
+                mask=pe_scores,
+                cache=None,
+            )
+            return self.unembed_out(output)
+
+        k_nope = self.embed_q(kv_latent, transpose=False)
+        values = self.unembed_out(kv_latent)
+        return scaled_dot_product_attention(
+            q_nope,
+            k_nope,
+            values,
+            scale=self.scale,
+            mask=pe_scores,
+            cache=None,
+        )
+
     def __call__(
         self,
         x: mx.array,
@@ -56,6 +111,9 @@ def __call__(
             output_h: (batch, target_len, hidden_dim) - Output hidden states.
         """
         batch, target_len, _ = x.shape
+        if not isinstance(cache, DeepSeekSparseCache):
+            raise TypeError("ParallaxDeepSeekV3Attention requires DeepSeekSparseCache.")
+
         if self.q_lora_rank is None:
             q = self.q_proj(x)
         else:
@@ -66,13 +124,8 @@ def __call__(
         compressed_kv = self.kv_a_proj_with_mqa(x)
         compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
         k_pe = k_pe.reshape(batch, target_len, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-
-        kv = kv.reshape(batch, target_len, self.num_heads, -1)
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
-        k_nope = k_nope.transpose(0, 2, 1, 3)
-
-        key_cache_global, value_cache_global = cache.get_cache()
+        kv_latent = self.kv_a_layernorm(compressed_kv)
+        kv_latent = kv_latent[:, None, :, :]
 
         # Compute RoPE offsets using array operations instead of loops
         if target_len == 1:
@@ -89,73 +142,87 @@ def __call__(
         q_pe = self.rope(q_pe, offset=current_pos)
         k_pe = self.rope(k_pe, offset=current_pos)
 
-        k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
-        keys = mx.concatenate([k_nope, k_pe], axis=-1)
-
-        block_size = key_cache_global.shape[3]
+        latent_cache, rope_cache = cache.get_cache()
 
         reshape_and_cache(
-            keys.transpose(0, 2, 1, 3),
-            values,
-            key_cache_global,
-            value_cache_global,
+            kv_latent.transpose(0, 2, 1, 3),
+            k_pe.transpose(0, 2, 1, 3),
+            latent_cache,
+            rope_cache,
             block_tables,
             context_lengths,
-            block_size,
+            cache.block_size,
             slot_mapping=slot_mapping,
         )
 
         if target_len == 1:
-            # Decode phase: Use Paged Attention
-            output = paged_attention(
-                queries,
-                key_cache_global,
-                value_cache_global,
+            q_latent = self.embed_q(q_nope)
+            output = mla_paged_attention(
+                q_latent,
+                q_pe,
+                latent_cache,
+                rope_cache,
                 block_tables,
                 context_lengths,
-                block_size,
+                cache.block_size,
                 self.scale,
-                self.num_heads,
-                v_head_dim=values.shape[-1],
             )
+            output = self.unembed_out(output)
             output = output.transpose(0, 2, 1, 3).reshape(batch, target_len, -1)
         else:
-            # Prefill Phase: Need to attend to both cached prefix and new tokens
-            # Check if any request has prefix cache
             has_prefix_cache = prefix_lens is not None and bool(mx.any(prefix_lens > 0))
 
             if has_prefix_cache:
-                # Use shared prefix cache handling with batch processing
-                # keys: (batch, num_heads, target_len, head_dim)
-                # values: (batch, target_len, num_heads, head_dim) -> transpose to (batch, num_heads, target_len, head_dim)
-                k_new = keys  # (batch, num_heads, target_len, head_dim)
-                v_new = values.transpose(0, 2, 1, 3)  # (batch, num_heads, target_len, head_dim)
-                output = compute_attention_with_prefix_cache(
-                    queries,  # (batch, num_heads, target_len, head_dim)
-                    k_new,
-                    v_new,
-                    cache,
-                    block_tables,
-                    prefix_lens,
-                    target_len,
-                    self.scale,
-                    self.num_heads,  # In deepseek_v3, num_heads equals n_kv_heads after MQA processing
-                    mask=mask,
-                )
+                max_prefix_len = int(mx.max(prefix_lens))
+                if max_prefix_len > 0:
+                    prefix_latent, prefix_k_pe = self._read_mla_batch(
+                        cache,
+                        block_tables,
+                        prefix_lens,
+                        max_prefix_len,
+                        kv_latent.dtype,
+                    )
+                    kv_full = mx.concatenate([prefix_latent, kv_latent], axis=2)
+                    k_pe_full = mx.concatenate([prefix_k_pe, k_pe], axis=2)
+                else:
+                    kv_full = kv_latent
+                    k_pe_full = k_pe
+
+                row_indices = mx.arange(target_len, dtype=mx.int32)
+                new_lens = context_lengths - prefix_lens
+                q_positions = prefix_lens[:, None] + row_indices[None, :]
+
+                if max_prefix_len > 0:
+                    prefix_positions = mx.arange(max_prefix_len, dtype=mx.int32)
+                    prefix_positions = mx.broadcast_to(
+                        prefix_positions[None, :], (batch, max_prefix_len)
+                    )
+                    prefix_valid = prefix_positions < prefix_lens[:, None]
+                else:
+                    prefix_positions = mx.zeros((batch, 0), dtype=mx.int32)
+                    prefix_valid = mx.zeros((batch, 0), dtype=mx.bool_)
+
+                new_positions = prefix_lens[:, None] + row_indices[None, :]
+                new_valid = row_indices[None, :] < new_lens[:, None]
+                key_positions = mx.concatenate([prefix_positions, new_positions], axis=1)
+                key_valid = mx.concatenate([prefix_valid, new_valid], axis=1)
+                valid = key_valid[:, None, :] & new_valid[:, :, None]
+                valid = valid & (key_positions[:, None, :] <= q_positions[:, :, None])
+                mask = mx.where(valid[:, None, :, :], 0.0, -float("inf")).astype(q_nope.dtype)
+
+                output = self._mla_attention(q_nope, q_pe, kv_full, k_pe_full, mask)
+                output = output.transpose(0, 2, 1, 3).reshape(batch, target_len, -1)
             else:
-                # No prefix cache, use standard self-attention on local data only
                 if mask is not None:
-                    mask = mx.array(mask, dtype=queries.dtype)
-
-                output = scaled_dot_product_attention(
-                    queries,
-                    keys,
-                    values.transpose(0, 2, 1, 3),
-                    scale=self.scale,
-                    mask=mask,
-                    cache=None,
-                )
+                    if mask.ndim == 2:
+                        mask = mask[None, None, :, :]
+                    elif mask.ndim == 3:
+                        mask = mask[:, None, :, :]
+                    if mask.dtype == mx.bool_:
+                        mask = mx.where(mask, 0.0, -float("inf"))
+                    mask = mask.astype(q_nope.dtype)
+
+                output = self._mla_attention(q_nope, q_pe, kv_latent, k_pe, mask)
                 output = output.transpose(0, 2, 1, 3).reshape(batch, target_len, -1)
 
         return self.o_proj(output)
diff --git a/src/parallax/server/cache/dsa_cache.py b/src/parallax/server/cache/dsa_cache.py
index 1f61f627..c903576c 100644
--- a/src/parallax/server/cache/dsa_cache.py
+++ b/src/parallax/server/cache/dsa_cache.py
@@ -25,10 +25,10 @@ def __init__(
         head_dim: int,
         head_dim_v: int,
         dtype: mx.Dtype,
-        index_head_dim: int,
-        index_n_heads: int,
         kv_lora_rank: int,
         qk_rope_head_dim: int,
+        index_head_dim: Optional[int] = None,
+        index_n_heads: Optional[int] = None,
         index_key_heads: int = 1,
     ):
         self.num_blocks = num_blocks
@@ -39,6 +39,9 @@ def __init__(
         self.dtype = dtype
         self.kv_lora_rank = kv_lora_rank
         self.qk_rope_head_dim = qk_rope_head_dim
+        if (index_head_dim is None) != (index_n_heads is None):
+            raise ValueError("index_head_dim and index_n_heads must be provided together.")
+
         self.index_head_dim = index_head_dim
         self.index_n_heads = index_n_heads
         self.index_key_heads = index_key_heads
@@ -51,17 +54,21 @@ def __init__(
             (1, num_blocks, 1, block_size, qk_rope_head_dim),
             dtype=dtype,
         )
-        self.indexer_key_cache = mx.zeros(
-            (
-                1,
-                num_blocks,
-                index_key_heads,
-                block_size,
-                index_head_dim,
-            ),
-            dtype=dtype,
-        )
-        mx.eval(self.latent_cache, self.rope_cache, self.indexer_key_cache)
+        self.indexer_key_cache = None
+        if index_head_dim is not None:
+            self.indexer_key_cache = mx.zeros(
+                (
+                    1,
+                    num_blocks,
+                    index_key_heads,
+                    block_size,
+                    index_head_dim,
+                ),
+                dtype=dtype,
+            )
+            mx.eval(self.latent_cache, self.rope_cache, self.indexer_key_cache)
+        else:
+            mx.eval(self.latent_cache, self.rope_cache)
 
     def get_cache(self) -> Tuple[mx.array, mx.array]:
         return self.latent_cache, self.rope_cache
@@ -133,6 +140,9 @@ def read_index_k(
         Returns:
             index_k: (index_heads, context_len, index_head_dim)
         """
+        if self.indexer_key_cache is None:
+            raise ValueError("DeepSeekSparseCache was created without an indexer cache.")
+
         positions = mx.arange(context_len)
         block_indices = positions // self.block_size
         offsets = positions % self.block_size
diff --git a/src/parallax/server/cache_manager.py b/src/parallax/server/cache_manager.py
index 6f2df219..30ea93aa 100644
--- a/src/parallax/server/cache_manager.py
+++ b/src/parallax/server/cache_manager.py
@@ -9,6 +9,14 @@
 from parallax.server.cache.kv_cache import KVCachePacked
 from parallax.server.cache.linear_cache import LinearCache
 from parallax.server.cache.minimax_m3_cache import MiniMaxM3SparseCache
+from parallax.utils.layer_types import (
+    ATTENTION,
+    ATTENTION_LAYER_TYPES,
+    DSA_ATTENTION,
+    LINEAR,
+    MLA_ATTENTION,
+    MSA_ATTENTION,
+)
 from parallax_utils.logging_config import get_logger
 
 logger = get_logger(__name__)
@@ -35,8 +43,7 @@ def __init__(
         index_key_heads: Optional[int] = None,
         kv_lora_rank: Optional[int] = None,
         qk_rope_head_dim: Optional[int] = None,
-        sparse_cache_type: Optional[str] = None,
-        # Hybrid Config: List of 'attention' or 'linear' or None (default 'attention')
+        # Hybrid Config: List of cache layer types or None (default 'attention')
         layer_types: Optional[List[str]] = None,
         # Linear Model / State Cache Params
         conv_dim: Optional[int] = None,
@@ -59,7 +66,6 @@ def __init__(
         self.index_key_heads = index_key_heads
         self.kv_lora_rank = kv_lora_rank
         self.qk_rope_head_dim = qk_rope_head_dim
-        self.sparse_cache_type = sparse_cache_type
         self.dtype = dtype
         self.block_size = block_size
         self.max_num_seqs = max_num_seqs
@@ -82,14 +88,14 @@ def __init__(
 
         # Determine layer types
         if layer_types is None:
-            self.layer_types = ["attention"] * num_layers
+            self.layer_types = [ATTENTION] * num_layers
         else:
             assert len(layer_types) == num_layers, "layer_types length must match num_layers"
             self.layer_types = layer_types
 
         # Check if we need blocks (any attention layer) and slots (any linear layer)
-        self.needs_blocks = any(t == "attention" for t in self.layer_types)
-        self.needs_slots = any(t == "linear" for t in self.layer_types)
+        self.needs_blocks = any(t in ATTENTION_LAYER_TYPES for t in self.layer_types)
+        self.needs_slots = any(t == LINEAR for t in self.layer_types)
 
         self.num_gpu_blocks, self.num_linear_prefix_slots = self._calculate_cache_allocation(
             self.cache_memory_fraction, self.dtype
@@ -115,12 +121,12 @@ def __init__(
 
         if self.needs_blocks:
             logger.info(
-                f"Allocated Paged KV Cache for {self.layer_types.count('attention')} layers: "
+                f"Allocated Paged KV Cache for {self._num_attention_layers()} layers: "
                 f"{self.num_gpu_blocks} blocks, {self.block_size} block_size, max_tokens: {self.num_gpu_blocks * self.block_size}"
             )
         if self.needs_slots:
             logger.info(
-                f"Allocated Linear State Cache for {self.layer_types.count('linear')} layers: "
+                f"Allocated Linear State Cache for {self._num_linear_layers()} layers: "
                 f"{self.max_num_seqs} active slots, "
                 f"{self.num_linear_prefix_slots} prefix slots"
             )
@@ -162,50 +168,82 @@ def _on_prefix_linear_slot_evict(self, slot: int):
             self.prefix_slot_allocator.free(slot)
             logger.debug(f"Freed evicted prefix linear slot: {slot}")
 
+    def _num_attention_layers(self) -> int:
+        return sum(1 for t in self.layer_types if t in ATTENTION_LAYER_TYPES)
+
+    def _num_linear_layers(self) -> int:
+        return sum(1 for t in self.layer_types if t == LINEAR)
+
+    def _validate_mla_cache_params(self, layer_type: str):
+        if self.kv_lora_rank is None or self.qk_rope_head_dim is None:
+            raise ValueError(
+                f"{layer_type} requires kv_lora_rank and qk_rope_head_dim "
+                "for compressed MLA storage."
+            )
+
+    def _validate_index_cache_params(self, layer_type: str):
+        if self.index_head_dim is None or self.index_n_heads is None:
+            raise ValueError(f"{layer_type} requires index_head_dim and index_n_heads.")
+
     def _create_cache(self, layer_type: str) -> BaseCache:
-        if layer_type == "attention":
-            if self.index_head_dim is not None and self.index_n_heads is not None:
-                if self.sparse_cache_type == "minimax_m3":
-                    return MiniMaxM3SparseCache(
-                        num_blocks=self.num_gpu_blocks,
-                        block_size=self.block_size,
-                        num_kv_heads=self.num_kv_heads,
-                        head_dim=self.head_dim,
-                        head_dim_v=self.head_dim_v,
-                        dtype=self.dtype,
-                        index_head_dim=self.index_head_dim,
-                        index_n_heads=self.index_n_heads,
-                        index_key_heads=self.index_key_heads or 1,
-                    )
-                if self.kv_lora_rank is None or self.qk_rope_head_dim is None:
-                    raise ValueError(
-                        "DeepSeek/GLM DSA cache requires kv_lora_rank and qk_rope_head_dim "
-                        "for compressed MLA storage."
-                    )
-                return DeepSeekSparseCache(
-                    num_blocks=self.num_gpu_blocks,
-                    block_size=self.block_size,
-                    num_kv_heads=self.num_kv_heads,
-                    head_dim=self.head_dim,
-                    head_dim_v=self.head_dim_v,
-                    dtype=self.dtype,
-                    index_head_dim=self.index_head_dim,
-                    index_n_heads=self.index_n_heads,
-                    kv_lora_rank=self.kv_lora_rank,
-                    qk_rope_head_dim=self.qk_rope_head_dim,
-                    index_key_heads=self.index_key_heads or 1,
-                )
-            else:
-                return KVCachePacked(
-                    num_blocks=self.num_gpu_blocks,
-                    block_size=self.block_size,
-                    num_kv_heads=self.num_kv_heads,
-                    head_dim=self.head_dim,
-                    head_dim_v=self.head_dim_v,
-                    dtype=self.dtype,
-                )
+        if layer_type == ATTENTION:
+            return KVCachePacked(
+                num_blocks=self.num_gpu_blocks,
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_dim=self.head_dim,
+                head_dim_v=self.head_dim_v,
+                dtype=self.dtype,
+            )
+
+        if layer_type == MLA_ATTENTION:
+            self._validate_mla_cache_params(layer_type)
+            return DeepSeekSparseCache(
+                num_blocks=self.num_gpu_blocks,
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_dim=self.head_dim,
+                head_dim_v=self.head_dim_v,
+                dtype=self.dtype,
+                index_head_dim=None,
+                index_n_heads=None,
+                kv_lora_rank=self.kv_lora_rank,
+                qk_rope_head_dim=self.qk_rope_head_dim,
+                index_key_heads=None,
+            )
+
+        if layer_type == DSA_ATTENTION:
+            self._validate_mla_cache_params(layer_type)
+            self._validate_index_cache_params(layer_type)
+            return DeepSeekSparseCache(
+                num_blocks=self.num_gpu_blocks,
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_dim=self.head_dim,
+                head_dim_v=self.head_dim_v,
+                dtype=self.dtype,
+                index_head_dim=self.index_head_dim,
+                index_n_heads=self.index_n_heads,
+                kv_lora_rank=self.kv_lora_rank,
+                qk_rope_head_dim=self.qk_rope_head_dim,
+                index_key_heads=self.index_key_heads or 1,
+            )
 
-        elif layer_type == "linear":
+        if layer_type == MSA_ATTENTION:
+            self._validate_index_cache_params(layer_type)
+            return MiniMaxM3SparseCache(
+                num_blocks=self.num_gpu_blocks,
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_dim=self.head_dim,
+                head_dim_v=self.head_dim_v,
+                dtype=self.dtype,
+                index_head_dim=self.index_head_dim,
+                index_n_heads=self.index_n_heads,
+                index_key_heads=self.index_key_heads or 1,
+            )
+
+        if layer_type == LINEAR:
             # We assume uniform linear config for all linear layers for now
             return LinearCache(
                 max_num_seqs=self.max_linear_slots,
@@ -217,8 +255,8 @@ def _create_cache(self, layer_type: str) -> BaseCache:
                 linear_num_v_heads=self.linear_num_v_heads,
                 dtype=self.dtype,
             )
-        else:
-            raise ValueError(f"Unknown layer type: {layer_type}")
+
+        raise ValueError(f"Unknown layer type: {layer_type}")
 
     def _needs_prefix_linear_slots(self) -> bool:
         return self.enable_prefix_cache and self.needs_blocks and self.needs_slots
@@ -228,7 +266,7 @@ def _dtype_size(self, dtype: mx.Dtype) -> int:
 
     def _calculate_linear_slot_bytes(self, dtype_size: int) -> int:
         """Calculate memory needed for one linear slot across all linear layers."""
-        num_linear_layers = self.layer_types.count("linear")
+        num_linear_layers = self._num_linear_layers()
         if num_linear_layers == 0:
             return 0
 
@@ -257,32 +295,39 @@ def _calculate_linear_cache_bytes(self, dtype_size: int, num_slots: int) -> int:
 
     def _calculate_kv_block_bytes(self, dtype_size: int) -> int:
         """Calculate memory needed for one KV block across all attention layers."""
-        if (
-            self.index_head_dim is not None
-            and self.index_n_heads is not None
-            and self.sparse_cache_type != "minimax_m3"
-            and self.kv_lora_rank is not None
-            and self.qk_rope_head_dim is not None
-        ):
-            one_layer_block_bytes = (
-                self.block_size * (self.kv_lora_rank + self.qk_rope_head_dim) * dtype_size
-            )
-        else:
-            one_layer_block_bytes = (
-                self.num_kv_heads * self.block_size * (self.head_dim + self.head_dim_v) * dtype_size
-            )
-        if self.index_head_dim is not None and self.index_n_heads is not None:
-            if self.sparse_cache_type == "minimax_m3":
-                index_heads = self.index_key_heads or 1
-            elif self.kv_lora_rank is not None and self.qk_rope_head_dim is not None:
-                index_heads = self.index_key_heads or 1
+        total_bytes = 0
+        standard_block_bytes = (
+            self.num_kv_heads * self.block_size * (self.head_dim + self.head_dim_v) * dtype_size
+        )
+        for layer_type in self.layer_types:
+            if layer_type == ATTENTION:
+                total_bytes += standard_block_bytes
+            elif layer_type == MLA_ATTENTION:
+                self._validate_mla_cache_params(layer_type)
+                total_bytes += (
+                    self.block_size * (self.kv_lora_rank + self.qk_rope_head_dim) * dtype_size
+                )
+            elif layer_type == DSA_ATTENTION:
+                self._validate_mla_cache_params(layer_type)
+                self._validate_index_cache_params(layer_type)
+                total_bytes += (
+                    self.block_size * (self.kv_lora_rank + self.qk_rope_head_dim) * dtype_size
+                )
+                total_bytes += (
+                    (self.index_key_heads or 1) * self.block_size * self.index_head_dim * dtype_size
+                )
+            elif layer_type == MSA_ATTENTION:
+                self._validate_index_cache_params(layer_type)
+                total_bytes += standard_block_bytes
+                total_bytes += (
+                    (self.index_key_heads or 1) * self.block_size * self.index_head_dim * dtype_size
+                )
+            elif layer_type == LINEAR:
+                continue
             else:
-                index_heads = self.index_n_heads
-            one_layer_block_bytes += (
-                index_heads * self.block_size * self.index_head_dim * dtype_size
-            )
+                raise ValueError(f"Unknown layer type: {layer_type}")
 
-        return one_layer_block_bytes * self.layer_types.count("attention")
+        return total_bytes
 
     def _calculate_prefix_linear_bytes(
         self,
@@ -325,14 +370,14 @@ def _calculate_cache_allocation(
         )
 
         # Total bytes per block = Sum over all attention layers
-        num_attention_layers = self.layer_types.count("attention")
+        num_attention_layers = self._num_attention_layers()
         total_block_bytes = self._calculate_kv_block_bytes(dtype_size)
 
         if total_block_bytes == 0:
             if active_linear_cache_bytes > 0:
                 logger.info(
                     f"Linear cache will use {active_linear_cache_bytes / 1024**3:.2f} GB "
-                    f"for {self.layer_types.count('linear')} layers "
+                    f"for {self._num_linear_layers()} layers "
                     f"({self.max_num_seqs} active slots, 0 prefix slots)"
                 )
             return 0, 0
@@ -368,7 +413,7 @@ def _calculate_cache_allocation(
             )
             logger.info(
                 f"Linear cache will use {total_linear_cache_bytes / 1024**3:.2f} GB "
-                f"for {self.layer_types.count('linear')} layers "
+                f"for {self._num_linear_layers()} layers "
                 f"({self.max_num_seqs} active slots, "
                 f"{num_linear_prefix_slots} prefix slots)"
             )
diff --git a/src/parallax/server/executor/base_executor.py b/src/parallax/server/executor/base_executor.py
index 456099b5..cc56488c 100755
--- a/src/parallax/server/executor/base_executor.py
+++ b/src/parallax/server/executor/base_executor.py
@@ -516,13 +516,6 @@ def recv_requests_from_peer(self) -> Tuple[List[Request], str]:
                         self.weight_version = int(recv_req[2].decode("ascii"))
                     else:
                         raise ValueError(f"Unknown request type: {recv_req[0]}")
-                    # First peer is responsible for tokenization
-                    # if self.is_first_peer and isinstance(recv_req, InitialRequest):
-                    #     recv_req.input_ids = self.tokenizer.encode(recv_req.prompt)
-                    #     recv_req.prompt_len = len(recv_req.input_ids)
-                    #     recv_req.max_total_length = min(
-                    #         recv_req.max_total_length, recv_req.prompt_len + recv_req.max_new_tokens
-                    #     )
 
                 except zmq.ZMQError:
                     break
diff --git a/src/parallax/server/executor/mlx_executor.py b/src/parallax/server/executor/mlx_executor.py
index 33c1646b..6fd20831 100755
--- a/src/parallax/server/executor/mlx_executor.py
+++ b/src/parallax/server/executor/mlx_executor.py
@@ -24,6 +24,7 @@
     complete_local_middle_chunk,
     filter_middle_chunk_next_batch,
 )
+from parallax.utils.layer_types import INDEX_CACHE_LAYER_TYPES, MLA_CACHE_LAYER_TYPES
 from parallax.utils.mac_prefill_adder import AddReqResult, MACPrefillAdder
 from parallax.utils.utils import (
     combine_padding_and_causal_masks,
@@ -131,8 +132,6 @@ def __init__(
             f"MLX sharded model loaded in {(time.time() - t0) * 1000:.1f} ms; num_layers={self.config.get('num_hidden_layers')}"
         )
 
-        is_minimax_m3 = self.config.get("model_type") == "minimax_m3"
-
         # TODO: Duplicate code to BaseExecutor since num_shard_layers and dtype are needed for initializing kv cache
         self.num_shard_layers = end_layer - start_layer
         self.dtype = get_device_dtype(dtype, device)
@@ -173,16 +172,11 @@ def __init__(
         index_head_dim = self.config.get("index_head_dim", None)
         index_n_heads = self.config.get("index_n_heads", None)
         kv_lora_rank = self.config.get("kv_lora_rank", None)
-        is_mla_dsa = (
-            index_head_dim is not None
-            and index_n_heads is not None
-            and not is_minimax_m3
-            and kv_lora_rank is not None
-            and qk_rope_head_dim is not None
-        )
-        index_key_heads = 1 if (is_mla_dsa or is_minimax_m3) else None
 
         layer_types = get_layer_types(self.config, start_layer, end_layer)
+        uses_mla_cache = any(t in MLA_CACHE_LAYER_TYPES for t in layer_types)
+        uses_index_cache = any(t in INDEX_CACHE_LAYER_TYPES for t in layer_types)
+        index_key_heads = 1 if uses_index_cache else None
         sliding_window = self.config.get("sliding_window", None)
         use_sliding_window = self.config.get("use_sliding_window", None)
         if use_sliding_window is False:
@@ -227,12 +221,11 @@ def __init__(
             block_size=kv_block_size,
             cache_memory_fraction=kv_cache_memory_fraction,
             head_dim_v=v_head_dim,
-            index_head_dim=index_head_dim,
-            index_n_heads=index_n_heads,
+            index_head_dim=index_head_dim if uses_index_cache else None,
+            index_n_heads=index_n_heads if uses_index_cache else None,
             index_key_heads=index_key_heads,
-            kv_lora_rank=kv_lora_rank if is_mla_dsa else None,
-            qk_rope_head_dim=qk_rope_head_dim if is_mla_dsa else None,
-            sparse_cache_type="minimax_m3" if is_minimax_m3 else None,
+            kv_lora_rank=kv_lora_rank if uses_mla_cache else None,
+            qk_rope_head_dim=qk_rope_head_dim if uses_mla_cache else None,
             layer_types=layer_types,
             max_num_seqs=max_batch_size // micro_batch_ratio,
             conv_dim=conv_dim,
diff --git a/src/parallax/server/scheduler.py b/src/parallax/server/scheduler.py
index df7fb7a6..ff163b71 100644
--- a/src/parallax/server/scheduler.py
+++ b/src/parallax/server/scheduler.py
@@ -124,16 +124,8 @@ def get_running_request(self, request_id: str) -> Optional[Request]:
         """Gets a request that is currently in the running state."""
         return self._running_requests.get(request_id)
 
-    def _prompt_string_to_request(self, request_str: str) -> InitialRequest:
-        """Convert the prompt string to InitialRequest."""
-        assert self.is_first_peer, "Only first peer can enqueue InitialRequest."
-        input_ids = self.tokenizer.encode(request_str)
-        return InitialRequest.from_prompt_ids(input_ids, self.max_new_tokens, self.max_total_length)
-
-    def enque_request(self, request: Request | str):
+    def enque_request(self, request: Request):
         """Enque a request to the scheduler's wait queue."""
-        if isinstance(request, str):
-            request = self._prompt_string_to_request(request)
 
         if request.is_finished:
             logger.warning(
diff --git a/src/parallax/utils/layer_types.py b/src/parallax/utils/layer_types.py
new file mode 100644
index 00000000..38adccbe
--- /dev/null
+++ b/src/parallax/utils/layer_types.py
@@ -0,0 +1,20 @@
+"""Cache layer type labels shared by executor and cache manager."""
+
+from typing import Final
+
+ATTENTION: Final = "attention"
+MLA_ATTENTION: Final = "mla_attention"
+DSA_ATTENTION: Final = "dsa_attention"
+MSA_ATTENTION: Final = "msa_attention"
+LINEAR: Final = "linear"
+
+ATTENTION_LAYER_TYPES: Final = frozenset(
+    {
+        ATTENTION,
+        MLA_ATTENTION,
+        DSA_ATTENTION,
+        MSA_ATTENTION,
+    }
+)
+MLA_CACHE_LAYER_TYPES: Final = frozenset({MLA_ATTENTION, DSA_ATTENTION})
+INDEX_CACHE_LAYER_TYPES: Final = frozenset({DSA_ATTENTION, MSA_ATTENTION})
diff --git a/src/parallax/utils/tokenizer_utils.py b/src/parallax/utils/tokenizer_utils.py
index 1d5f9d6c..2fb1e082 100755
--- a/src/parallax/utils/tokenizer_utils.py
+++ b/src/parallax/utils/tokenizer_utils.py
@@ -1,223 +1,18 @@
-"""
-Implements parallax detokenizers for performance.
-"""
+"""Tokenizer loading helpers."""
 
-import json
-import uuid
-from dataclasses import dataclass
-from functools import partial
-from json import JSONDecodeError
-from typing import Any, Callable, Dict, List, Optional, Tuple
-
-from mlx_lm.tokenizer_utils import (
-    BPEStreamingDetokenizer,
-    NaiveStreamingDetokenizer,
-    SPMStreamingDetokenizer,
-    _is_bpe_decoder,
-    _is_spm_decoder,
-    _is_spm_decoder_no_space,
-)
 from mlx_lm.tokenizer_utils import load as _mlx_load_tokenizer
 
 
-class ParallaxNaiveStreamingDetokenizer(NaiveStreamingDetokenizer):
-    """A custom BPE streaming detokenizer that add an argument 'tokenizer'"""
-
-    def __init__(self, tokenizer, tokenmap):
-        self._tokenizer = tokenizer
-        self._tokenizer.decode([0])
-        self.reset()
-
-
-class ParallaxBPEStreamingDetokenizer(BPEStreamingDetokenizer):
-    """A custom BPE streaming detokenizer that skips initializing tokenmap"""
-
-    def __init__(self, tokenizer, tokenmap):
-        self.clean_spaces = tokenizer.clean_up_tokenization_spaces
-        self.tokenmap = tokenmap
-        self.reset()
-        self.make_byte_decoder()
-
-
-class ParallaxSPMStreamingDetokenizer(SPMStreamingDetokenizer):
-    """A custom SPM streaming detokenizer that skips initializing tokenmap"""
-
-    def __init__(self, tokenizer, tokenmap, trim_space=True):
-        self.trim_space = trim_space
-        self._sep = "\u2581".encode()
-        self.tokenmap = tokenmap
-        self.reset()
-
-
-def _get_spm_tokenmap(tokenizer):
-    """Initialize spm tokenmap for reuse"""
-    # Extract the tokens in a list from id to text
-    tokenmap = [""] * (max(tokenizer.vocab.values()) + 1)
-    for value, tokenid in tokenizer.vocab.items():
-        if value.startswith("<0x"):
-            # Replace bytes with their value
-            tokenmap[tokenid] = bytes([int(value[3:5], 16)])
-        else:
-            tokenmap[tokenid] = value.encode()
-    return tokenmap
-
-
-def _get_bpe_tokenmap(tokenizer):
-    """Initialize bpe tokenmap for reuse"""
-    # Extract the tokens in a list from id to text
-    tokenmap = [None] * len(tokenizer.vocab)
-    for value, tokenid in tokenizer.vocab.items():
-        tokenmap[tokenid] = value
-    return tokenmap
-
-
-def load_detokenizer(model_path, tokenizer):
-    """Load a huggingface tokenizer and try to infer the type of streaming
-    detokenizer to use.
-
-    Note, to use a fast streaming tokenizer, pass a local file path rather than
-    a Hugging Face repo ID.
-    """
-    detokenizer_class = ParallaxNaiveStreamingDetokenizer
-    tokenmap = None
-
-    tokenizer_file = model_path / "tokenizer.json"
-    if tokenizer_file.exists():
-        with open(tokenizer_file, "r", encoding="utf-8") as fid:
-            try:
-                tokenizer_content = json.load(fid)
-            except JSONDecodeError as e:
-                raise JSONDecodeError("Failed to parse tokenizer.json", e.doc, e.pos)
-
-        if "decoder" in tokenizer_content:
-            if _is_spm_decoder(tokenizer_content["decoder"]):
-                detokenizer_class = ParallaxSPMStreamingDetokenizer
-                tokenmap = _get_spm_tokenmap(tokenizer)
-            elif _is_spm_decoder_no_space(tokenizer_content["decoder"]):
-                detokenizer_class = partial(ParallaxSPMStreamingDetokenizer, trim_space=False)
-                tokenmap = _get_spm_tokenmap(tokenizer)
-            elif _is_bpe_decoder(tokenizer_content["decoder"]):
-                detokenizer_class = ParallaxBPEStreamingDetokenizer
-                tokenmap = _get_bpe_tokenmap(tokenizer)
-
-    return detokenizer_class, tokenmap
-
-
 def load_tokenizer(model_path, trust_remote_code=True, tokenizer_config_extra=None, **kwargs):
     """
     Wrapper function for MLX load_tokenizer that defaults trust_remote_code to True.
     This is needed for models like Kimi-K2 that contain custom code.
-
-    Args:
-        model_path: Path to the model
-        trust_remote_code: Whether to trust remote code (defaults to True)
-        tokenizer_config_extra: Extra config to pass to AutoTokenizer.from_pretrained
-        **kwargs: Additional arguments to pass to the original load_tokenizer
-
-    Returns:
-        The loaded tokenizer
     """
     if tokenizer_config_extra is None:
         tokenizer_config_extra = {}
 
-    # Add trust_remote_code to the tokenizer config
     if trust_remote_code:
         tokenizer_config_extra = tokenizer_config_extra.copy()
         tokenizer_config_extra["trust_remote_code"] = True
 
     return _mlx_load_tokenizer(model_path, tokenizer_config_extra=tokenizer_config_extra, **kwargs)
-
-
-@dataclass
-class ToolCallState:
-    has_tool_calling: bool
-    tool_call_start: Optional[str]
-    tool_call_end: Optional[str]
-    tool_parser: Optional[Callable[[str, Any], Dict[str, Any]]]
-    tools: Optional[List[Any]]
-    stream: bool
-    in_tool_call: bool = False
-    tool_text: str = ""
-    tool_call_idx: int = 0
-    made_tool_call: bool = False
-
-    @classmethod
-    def from_tokenizer(cls, tokenizer, tools: Optional[List[Any]], stream: bool):
-        has_tool_calling = bool(getattr(tokenizer, "has_tool_calling", False))
-        tool_parser = getattr(tokenizer, "tool_parser", None)
-        tool_call_start = getattr(tokenizer, "tool_call_start", None)
-        tool_call_end = getattr(tokenizer, "tool_call_end", None)
-        if not (has_tool_calling and tool_parser and tool_call_start and tool_call_end):
-            return cls(
-                has_tool_calling=False,
-                tool_call_start=None,
-                tool_call_end=None,
-                tool_parser=None,
-                tools=tools,
-                stream=stream,
-            )
-        return cls(
-            has_tool_calling=True,
-            tool_call_start=tool_call_start,
-            tool_call_end=tool_call_end,
-            tool_parser=tool_parser,
-            tools=tools,
-            stream=stream,
-        )
-
-    def _format_tool_call(self, tool_call: Dict[str, Any]):
-        tool_call_id = tool_call.pop("id", None) or str(uuid.uuid4())
-        tool_call["arguments"] = json.dumps(tool_call["arguments"], ensure_ascii=False)
-        out = {
-            "function": tool_call,
-            "type": "function",
-            "id": tool_call_id,
-        }
-        if self.stream:
-            out["index"] = self.tool_call_idx
-            self.tool_call_idx += 1
-        return out
-
-    def _parse_tool_text(self, tool_text: str) -> Tuple[List[Dict[str, Any]], Optional[str]]:
-        try:
-            parsed = self.tool_parser(tool_text, self.tools)
-        except Exception:
-            fallback_text = f"{self.tool_call_start}{tool_text}{self.tool_call_end}"
-            return [], fallback_text
-        if isinstance(parsed, list):
-            return [self._format_tool_call(tc) for tc in parsed], None
-        return [self._format_tool_call(parsed)], None
-
-    def extract_from_segment(self, segment: str) -> Tuple[str, List[Dict[str, Any]]]:
-        if not self.has_tool_calling or not segment:
-            return segment, []
-        output_chunks: List[str] = []
-        new_tool_calls: List[Dict[str, Any]] = []
-        idx = 0
-        while idx < len(segment):
-            if not self.in_tool_call:
-                start_pos = segment.find(self.tool_call_start, idx)
-                if start_pos == -1:
-                    output_chunks.append(segment[idx:])
-                    break
-                if start_pos > idx:
-                    output_chunks.append(segment[idx:start_pos])
-                self.in_tool_call = True
-                self.made_tool_call = True
-                self.tool_text = ""
-                idx = start_pos + len(self.tool_call_start)
-            else:
-                end_pos = segment.find(self.tool_call_end, idx)
-                if end_pos == -1:
-                    self.tool_text += segment[idx:]
-                    break
-                self.tool_text += segment[idx:end_pos]
-                parsed_calls, fallback_text = self._parse_tool_text(self.tool_text)
-                if parsed_calls:
-                    new_tool_calls.extend(parsed_calls)
-                if fallback_text:
-                    output_chunks.append(fallback_text)
-                self.tool_text = ""
-                self.in_tool_call = False
-                idx = end_pos + len(self.tool_call_end)
-        return "".join(output_chunks), new_tool_calls
diff --git a/src/parallax/utils/utils.py b/src/parallax/utils/utils.py
index 33d10d8e..9addb90e 100644
--- a/src/parallax/utils/utils.py
+++ b/src/parallax/utils/utils.py
@@ -12,6 +12,13 @@
 import torch
 import zmq
 
+from parallax.utils.layer_types import (
+    ATTENTION,
+    DSA_ATTENTION,
+    LINEAR,
+    MLA_ATTENTION,
+    MSA_ATTENTION,
+)
 from parallax.utils.model_download import download_model_file
 
 
@@ -420,8 +427,27 @@ def initialize_nccl_port():
     return nccl_port
 
 
+def _attention_cache_layer_type(config: dict) -> str:
+    model_type = config.get("model_type")
+    if model_type == "minimax_m3":
+        return MSA_ATTENTION
+
+    has_mla_cache = (
+        config.get("kv_lora_rank") is not None and config.get("qk_rope_head_dim") is not None
+    )
+    has_dsa_index = (
+        config.get("index_head_dim") is not None and config.get("index_n_heads") is not None
+    )
+    if has_mla_cache and has_dsa_index:
+        return DSA_ATTENTION
+    if has_mla_cache and model_type in {"deepseek_v3", "kimi_k2"}:
+        return MLA_ATTENTION
+    return ATTENTION
+
+
 def get_layer_types(config: dict, start_layer: int, end_layer: int) -> List[str]:
     num_shard_layers = end_layer - start_layer
+    attention_type = _attention_cache_layer_type(config)
 
     # Case 1: Explicit layer types (e.g., DeepSeek with layers_block_type)
     layer_types = config.get("layers_block_type", None)
@@ -429,7 +455,7 @@ def get_layer_types(config: dict, start_layer: int, end_layer: int) -> List[str]
         if len(layer_types) >= end_layer:
             layer_types = layer_types[start_layer:end_layer]
         return [
-            "linear" if t in ["mamba", "linear_attention"] else "attention" for t in layer_types
+            LINEAR if t in ["mamba", "linear_attention"] else attention_type for t in layer_types
         ]
 
     # Case 2: linear_attn_config with full_attn_layers (e.g., Kimi)
@@ -439,9 +465,9 @@ def get_layer_types(config: dict, start_layer: int, end_layer: int) -> List[str]
         layer_types = []
         for i in range(start_layer, end_layer):
             if i in full_attn_layers:
-                layer_types.append("attention")
+                layer_types.append(attention_type)
             else:
-                layer_types.append("linear")
+                layer_types.append(LINEAR)
         return layer_types
 
     # Case 3: full_attention_interval (e.g., Qwen3Next)
@@ -450,8 +476,8 @@ def get_layer_types(config: dict, start_layer: int, end_layer: int) -> List[str]
         layer_types = []
         for i in range(start_layer, end_layer):
             is_linear = (i + 1) % full_attention_interval != 0
-            layer_types.append("linear" if is_linear else "attention")
+            layer_types.append(LINEAR if is_linear else attention_type)
         return layer_types
 
     # Default: all attention layers
-    return ["attention"] * num_shard_layers
+    return [attention_type] * num_shard_layers
diff --git a/src/parallax_extensions/CMakelists.txt b/src/parallax_extensions/CMakelists.txt
index 41630ff0..7bf25e4b 100755
--- a/src/parallax_extensions/CMakelists.txt
+++ b/src/parallax_extensions/CMakelists.txt
@@ -36,6 +36,7 @@ target_sources(
   parallax_ext
   PUBLIC ${CMAKE_CURRENT_LIST_DIR}/kernels/dsa_paged_attention.cpp
          ${CMAKE_CURRENT_LIST_DIR}/kernels/dsa_indexer.cpp
+         ${CMAKE_CURRENT_LIST_DIR}/kernels/mla_paged_attention.cpp
          ${CMAKE_CURRENT_LIST_DIR}/kernels/paged_attention.cpp
          ${CMAKE_CURRENT_LIST_DIR}/kernels/sparse_indexer.cpp
          ${CMAKE_CURRENT_LIST_DIR}/kernels/sparse_paged_attention.cpp
@@ -62,6 +63,7 @@ if(MLX_BUILD_METAL)
     ${CMAKE_CURRENT_LIST_DIR}/kernels/float8.metal
     ${CMAKE_CURRENT_LIST_DIR}/kernels/dsa_paged_attention.metal
     ${CMAKE_CURRENT_LIST_DIR}/kernels/dsa_indexer.metal
+    ${CMAKE_CURRENT_LIST_DIR}/kernels/mla_paged_attention.metal
     ${CMAKE_CURRENT_LIST_DIR}/kernels/paged_attention.metal
     ${CMAKE_CURRENT_LIST_DIR}/kernels/sparse_indexer.metal
     ${CMAKE_CURRENT_LIST_DIR}/kernels/sparse_paged_attention.metal
diff --git a/src/parallax_extensions/bindings.cpp b/src/parallax_extensions/bindings.cpp
index 4ecd37b3..edfc3deb 100755
--- a/src/parallax_extensions/bindings.cpp
+++ b/src/parallax_extensions/bindings.cpp
@@ -3,6 +3,7 @@
 
 #include "kernels/dsa_indexer.h"
 #include "kernels/dsa_paged_attention.h"
+#include "kernels/mla_paged_attention.h"
 #include "kernels/paged_attention.h"
 #include "kernels/reshape_and_cache.h"
 #include "kernels/sparse_indexer.h"
@@ -34,6 +35,37 @@ NB_MODULE(_ext, m) {
             array: ``Dummy output``
       )");
 
+  m.def(
+      "mla_paged_attention",
+      &parallax_ext::mla_paged_attention,
+      "q_latent"_a,
+      "q_pe"_a,
+      "latent_cache"_a,
+      "rope_cache"_a,
+      "block_tables"_a,
+      "seq_lens"_a,
+      "block_size"_a,
+      "scale"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"(
+        Decode dense MLA paged attention over compressed MLA cache.
+
+        Args:
+            q_latent (array): MLA latent query [num_seqs, num_heads, latent_dim].
+            q_pe (array): RoPE query [num_seqs, num_heads, rope_dim].
+            latent_cache (array): Paged MLA latent cache [1, num_blocks, 1, block_size, latent_dim].
+            rope_cache (array): Paged RoPE key cache [1, num_blocks, 1, block_size, rope_dim].
+            block_tables (array): Input array [num_seqs, max_num_blocks_per_seq].
+            seq_lens (array): Input array [num_seqs].
+            block_size (int): Cache block size.
+            scale (float): Attention scale.
+            stream (Stream or Device): Stream on which to schedule the operation.
+
+        Returns:
+            array: latent-space attention result [num_seqs, num_heads, latent_dim]
+      )");
+
   m.def(
       "dsa_paged_attention",
       &parallax_ext::dsa_paged_attention,
diff --git a/src/parallax_extensions/kernels/mla_paged_attention.cpp b/src/parallax_extensions/kernels/mla_paged_attention.cpp
new file mode 100644
index 00000000..73815a35
--- /dev/null
+++ b/src/parallax_extensions/kernels/mla_paged_attention.cpp
@@ -0,0 +1,138 @@
+#include <sstream>
+#include <string>
+
+#include "mla_paged_attention.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/utils.h"
+#include "utils.h"
+
+namespace parallax_ext {
+
+mx::array mla_paged_attention(
+    const mx::array& q_latent,
+    const mx::array& q_pe,
+    const mx::array& latent_cache,
+    const mx::array& rope_cache,
+    const mx::array& block_tables,
+    const mx::array& seq_lens,
+    const int64_t block_size,
+    const float scale,
+    mx::StreamOrDevice s /* = {} */
+) {
+    auto out_dtype = q_latent.dtype();
+    auto out_shape = q_latent.shape();
+    const std::vector<mx::array> inputs = {
+        q_latent, q_pe, latent_cache, rope_cache, block_tables, seq_lens};
+    return mx::array(
+        out_shape,
+        out_dtype,
+        std::make_shared<MLAPagedAttention>(to_stream(s), block_size, scale),
+        inputs);
+}
+
+void MLAPagedAttention::eval_cpu(
+  const std::vector<mx::array>& inputs,
+  std::vector<mx::array>& outputs) {
+    return;
+}
+
+void MLAPagedAttention::eval_gpu(
+  const std::vector<mx::array>& inputs,
+  std::vector<mx::array>& outputs) {
+    assert(inputs.size() == 6);
+    auto& q_latent = inputs[0];
+    auto& q_pe = inputs[1];
+    auto& latent_cache = inputs[2];
+    auto& rope_cache = inputs[3];
+    auto& block_tables = inputs[4];
+    auto& seq_lens = inputs[5];
+    auto& out = outputs[0];
+
+    const int latent_dim = q_latent.shape(2);
+    const int rope_dim = q_pe.shape(2);
+    if (latent_dim > 512) {
+      std::ostringstream msg;
+      msg << "MLAPagedAttention supports latent_dim <= 512, got " << latent_dim;
+      throw std::runtime_error(msg.str());
+    }
+    if (rope_dim > 128) {
+      std::ostringstream msg;
+      msg << "MLAPagedAttention supports rope_dim <= 128, got " << rope_dim;
+      throw std::runtime_error(msg.str());
+    }
+    if (q_pe.shape(0) != q_latent.shape(0) || q_pe.shape(1) != q_latent.shape(1)) {
+      throw std::runtime_error("MLAPagedAttention requires q_latent and q_pe batch/heads to match");
+    }
+    if (latent_cache.dtype() != q_latent.dtype() || rope_cache.dtype() != q_latent.dtype()) {
+      throw std::runtime_error(
+          "MLAPagedAttention requires q_latent, latent_cache, and rope_cache to share dtype");
+    }
+    if (latent_cache.shape(4) != latent_dim || rope_cache.shape(4) != rope_dim) {
+      throw std::runtime_error("MLAPagedAttention cache dims must match query dims");
+    }
+
+    auto& s = stream();
+    auto& d = mx::metal::device(s.device);
+    auto lib = d.get_library("parallax_ext", current_binary_dir());
+    auto& compute_encoder = mx::metal::get_command_encoder(s);
+
+    out.set_data(mlx::core::allocator::malloc(out.nbytes()));
+
+    constexpr int num_threads = 32;
+    constexpr int num_simd_lanes = 32;
+    const int64_t num_seqs = q_latent.shape(0);
+    const int64_t num_heads = q_latent.shape(1);
+    const int64_t max_num_blocks_per_seq = block_tables.shape(1);
+
+    std::string kname = "mla_paged_attention_" + get_type_string(out.dtype());
+    kname += "_cache_" + get_type_string(latent_cache.dtype());
+    kname += "_ld" + std::to_string(latent_dim);
+    kname += "_rd" + std::to_string(rope_dim);
+    kname += "_nt" + std::to_string(num_threads);
+    kname += "_nsl" + std::to_string(num_simd_lanes);
+
+    auto kernel = d.get_kernel(kname, lib);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    float scale_32 = static_cast<float>(scale_);
+    int32_t block_size_32 = static_cast<int32_t>(block_size_);
+    int32_t max_num_blocks_per_seq_32 = static_cast<int32_t>(max_num_blocks_per_seq);
+    int32_t q_latent_seq_stride_32 = static_cast<int32_t>(q_latent.strides(0));
+    int32_t q_latent_head_stride_32 = static_cast<int32_t>(q_latent.strides(1));
+    int32_t q_pe_seq_stride_32 = static_cast<int32_t>(q_pe.strides(0));
+    int32_t q_pe_head_stride_32 = static_cast<int32_t>(q_pe.strides(1));
+    int32_t latent_block_stride_32 = static_cast<int32_t>(latent_cache.strides(1));
+    int32_t latent_token_stride_32 = static_cast<int32_t>(latent_cache.strides(3));
+    int32_t rope_block_stride_32 = static_cast<int32_t>(rope_cache.strides(1));
+    int32_t rope_token_stride_32 = static_cast<int32_t>(rope_cache.strides(3));
+
+    compute_encoder.set_output_array(out, 0);
+    compute_encoder.set_input_array(q_latent, 1);
+    compute_encoder.set_input_array(q_pe, 2);
+    compute_encoder.set_input_array(latent_cache, 3);
+    compute_encoder.set_input_array(rope_cache, 4);
+    compute_encoder.set_input_array(block_tables, 5);
+    compute_encoder.set_input_array(seq_lens, 6);
+    compute_encoder.set_bytes(scale_32, 7);
+    compute_encoder.set_bytes(block_size_32, 8);
+    compute_encoder.set_bytes(max_num_blocks_per_seq_32, 9);
+    compute_encoder.set_bytes(q_latent_seq_stride_32, 10);
+    compute_encoder.set_bytes(q_latent_head_stride_32, 11);
+    compute_encoder.set_bytes(q_pe_seq_stride_32, 12);
+    compute_encoder.set_bytes(q_pe_head_stride_32, 13);
+    compute_encoder.set_bytes(latent_block_stride_32, 14);
+    compute_encoder.set_bytes(latent_token_stride_32, 15);
+    compute_encoder.set_bytes(rope_block_stride_32, 16);
+    compute_encoder.set_bytes(rope_token_stride_32, 17);
+
+    MTL::Size grid = MTL::Size(num_heads, num_seqs, 1);
+    MTL::Size threadgroup = MTL::Size(num_threads, 1, 1);
+    compute_encoder.dispatch_threadgroups(grid, threadgroup);
+}
+
+bool MLAPagedAttention::is_equivalent(const mx::Primitive& other) const {
+  const MLAPagedAttention& r_other = static_cast<const MLAPagedAttention&>(other);
+  return block_size_ == r_other.block_size_ && scale_ == r_other.scale_;
+}
+
+} // namespace parallax_ext
diff --git a/src/parallax_extensions/kernels/mla_paged_attention.h b/src/parallax_extensions/kernels/mla_paged_attention.h
new file mode 100644
index 00000000..20416a00
--- /dev/null
+++ b/src/parallax_extensions/kernels/mla_paged_attention.h
@@ -0,0 +1,42 @@
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mx = mlx::core;
+
+namespace parallax_ext {
+
+mx::array mla_paged_attention(
+    const mx::array& q_latent,
+    const mx::array& q_pe,
+    const mx::array& latent_cache,
+    const mx::array& rope_cache,
+    const mx::array& block_tables,
+    const mx::array& seq_lens,
+    const int64_t block_size,
+    const float scale,
+    mx::StreamOrDevice s = {});
+
+class MLAPagedAttention : public mx::Primitive {
+  public:
+    explicit MLAPagedAttention(mx::Stream stream, int64_t block_size, float scale)
+        : mx::Primitive(stream), block_size_(block_size), scale_(scale) {};
+
+    void eval_cpu(
+        const std::vector<mx::array>& inputs,
+        std::vector<mx::array>& outputs) override;
+    void eval_gpu(
+        const std::vector<mx::array>& inputs,
+        std::vector<mx::array>& outputs) override;
+
+    const char* name() const override {
+      return "MLAPagedAttention";
+    }
+
+    bool is_equivalent(const mx::Primitive& other) const override;
+
+  private:
+    int64_t block_size_;
+    float scale_;
+};
+
+} // namespace parallax_ext
diff --git a/src/parallax_extensions/kernels/mla_paged_attention.metal b/src/parallax_extensions/kernels/mla_paged_attention.metal
new file mode 100644
index 00000000..2437d3fa
--- /dev/null
+++ b/src/parallax_extensions/kernels/mla_paged_attention.metal
@@ -0,0 +1,163 @@
+#include "./utils.metal"
+#include <metal_stdlib>
+
+using namespace metal;
+
+#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
+
+template <typename T, typename CACHE_T, int LATENT_DIM, int ROPE_DIM,
+          int NUM_THREADS, int NUM_SIMD_LANES>
+[[kernel]] void mla_paged_attention(
+    device T *out [[buffer(0)]],
+    device const T *q_latent [[buffer(1)]],
+    device const T *q_pe [[buffer(2)]],
+    device const CACHE_T *latent_cache [[buffer(3)]],
+    device const CACHE_T *rope_cache [[buffer(4)]],
+    device const int32_t *block_tables [[buffer(5)]],
+    device const int32_t *context_lens [[buffer(6)]],
+    const constant float &scale [[buffer(7)]],
+    const constant int &block_size [[buffer(8)]],
+    const constant int &max_num_blocks_per_seq [[buffer(9)]],
+    const constant int &q_latent_seq_stride [[buffer(10)]],
+    const constant int &q_latent_head_stride [[buffer(11)]],
+    const constant int &q_pe_seq_stride [[buffer(12)]],
+    const constant int &q_pe_head_stride [[buffer(13)]],
+    const constant int &latent_block_stride [[buffer(14)]],
+    const constant int &latent_token_stride [[buffer(15)]],
+    const constant int &rope_block_stride [[buffer(16)]],
+    const constant int &rope_token_stride [[buffer(17)]],
+    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]],
+    uint3 threadgroups_per_grid [[threadgroups_per_grid]],
+    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]]) {
+  const int head_idx = threadgroup_position_in_grid.x;
+  const int seq_idx = threadgroup_position_in_grid.y;
+  const int tid = thread_position_in_threadgroup.x;
+  const int context_len = context_lens[seq_idx];
+  const int num_context_blocks = DIVIDE_ROUND_UP(context_len, block_size);
+
+  const device T *q_latent_ptr =
+      q_latent + seq_idx * q_latent_seq_stride + head_idx * q_latent_head_stride;
+  const device T *q_pe_ptr =
+      q_pe + seq_idx * q_pe_seq_stride + head_idx * q_pe_head_stride;
+
+  float q_latent_vec[16] = {0.0f};
+  float q_pe_vec[4] = {0.0f};
+  float acc_vec[16] = {0.0f};
+
+  for (int i = tid; i < LATENT_DIM; i += NUM_SIMD_LANES) {
+    q_latent_vec[i / NUM_SIMD_LANES] = (float)q_latent_ptr[i];
+  }
+  for (int i = tid; i < ROPE_DIM; i += NUM_SIMD_LANES) {
+    q_pe_vec[i / NUM_SIMD_LANES] = (float)q_pe_ptr[i];
+  }
+
+  float m_i = -FLT_MAX;
+  float l_i = 0.0f;
+  const device int32_t *block_table =
+      block_tables + seq_idx * max_num_blocks_per_seq;
+
+  for (int block_idx = 0; block_idx < num_context_blocks; ++block_idx) {
+    const int physical_block_number = block_table[block_idx];
+    const int block_start = block_idx * block_size;
+    const int tokens_in_block = min(block_size, context_len - block_start);
+
+    const device CACHE_T *latent_block =
+        latent_cache + (int64_t)physical_block_number * latent_block_stride;
+    const device CACHE_T *rope_block =
+        rope_cache + (int64_t)physical_block_number * rope_block_stride;
+
+    for (int token_offset = 0; token_offset < tokens_in_block; ++token_offset) {
+      const device CACHE_T *latent_ptr =
+          latent_block + token_offset * latent_token_stride;
+      const device CACHE_T *rope_ptr =
+          rope_block + token_offset * rope_token_stride;
+
+      float score = 0.0f;
+      for (int i = tid; i < LATENT_DIM; i += NUM_SIMD_LANES) {
+        score += q_latent_vec[i / NUM_SIMD_LANES] * (float)latent_ptr[i];
+      }
+      for (int i = tid; i < ROPE_DIM; i += NUM_SIMD_LANES) {
+        score += q_pe_vec[i / NUM_SIMD_LANES] * (float)rope_ptr[i];
+      }
+      score = simd_sum(score) * scale;
+
+      const float m_prev = m_i;
+      m_i = max(m_prev, score);
+      const float alpha = exp(m_prev - m_i);
+      const float beta = exp(score - m_i);
+      l_i = l_i * alpha + beta;
+
+      for (int i = tid; i < LATENT_DIM; i += NUM_SIMD_LANES) {
+        const int idx = i / NUM_SIMD_LANES;
+        acc_vec[idx] = acc_vec[idx] * alpha + (float)latent_ptr[i] * beta;
+      }
+    }
+  }
+
+  device T *out_ptr =
+      out + (seq_idx * threadgroups_per_grid.x + head_idx) * LATENT_DIM;
+
+  const float inv_l = l_i > 0.0f ? 1.0f / l_i : 0.0f;
+  for (int i = tid; i < LATENT_DIM; i += NUM_SIMD_LANES) {
+    out_ptr[i] = T(acc_vec[i / NUM_SIMD_LANES] * inv_l);
+  }
+}
+
+#define instantiate_mla_paged_attention_inner(                                \
+    type, cache_type, latent_dim, rope_dim, num_threads, num_simd_lanes)       \
+  template [[host_name("mla_paged_attention_" #type "_cache_" #cache_type     \
+                       "_ld" #latent_dim "_rd" #rope_dim "_nt" #num_threads  \
+                       "_nsl" #num_simd_lanes)]] [[kernel]] void              \
+  mla_paged_attention<type, cache_type, latent_dim, rope_dim, num_threads,    \
+                      num_simd_lanes>(                                        \
+      device type *out [[buffer(0)]], device const type *q_latent             \
+      [[buffer(1)]], device const type *q_pe [[buffer(2)]],                   \
+      device const cache_type *latent_cache [[buffer(3)]],                    \
+      device const cache_type *rope_cache [[buffer(4)]],                      \
+      device const int32_t *block_tables [[buffer(5)]],                       \
+      device const int32_t *context_lens [[buffer(6)]],                       \
+      const constant float &scale [[buffer(7)]],                              \
+      const constant int &block_size [[buffer(8)]],                           \
+      const constant int &max_num_blocks_per_seq [[buffer(9)]],               \
+      const constant int &q_latent_seq_stride [[buffer(10)]],                 \
+      const constant int &q_latent_head_stride [[buffer(11)]],                \
+      const constant int &q_pe_seq_stride [[buffer(12)]],                     \
+      const constant int &q_pe_head_stride [[buffer(13)]],                    \
+      const constant int &latent_block_stride [[buffer(14)]],                 \
+      const constant int &latent_token_stride [[buffer(15)]],                 \
+      const constant int &rope_block_stride [[buffer(16)]],                   \
+      const constant int &rope_token_stride [[buffer(17)]],                   \
+      uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]],     \
+      uint3 threadgroups_per_grid [[threadgroups_per_grid]],                  \
+      uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]]);
+
+#define instantiate_mla_rope_dims(type, cache_type, latent_dim, num_threads,  \
+                                  num_simd_lanes)                             \
+  instantiate_mla_paged_attention_inner(type, cache_type, latent_dim, 2,      \
+                                        num_threads, num_simd_lanes);         \
+  instantiate_mla_paged_attention_inner(type, cache_type, latent_dim, 32,     \
+                                        num_threads, num_simd_lanes);         \
+  instantiate_mla_paged_attention_inner(type, cache_type, latent_dim, 64,     \
+                                        num_threads, num_simd_lanes);         \
+  instantiate_mla_paged_attention_inner(type, cache_type, latent_dim, 128,    \
+                                        num_threads, num_simd_lanes);
+
+#define instantiate_mla_latent_dims(type, cache_type, num_threads,            \
+                                    num_simd_lanes)                           \
+  instantiate_mla_rope_dims(type, cache_type, 4, num_threads,                 \
+                            num_simd_lanes);                                  \
+  instantiate_mla_rope_dims(type, cache_type, 64, num_threads,                \
+                            num_simd_lanes);                                  \
+  instantiate_mla_rope_dims(type, cache_type, 128, num_threads,               \
+                            num_simd_lanes);                                  \
+  instantiate_mla_rope_dims(type, cache_type, 256, num_threads,               \
+                            num_simd_lanes);                                  \
+  instantiate_mla_rope_dims(type, cache_type, 512, num_threads,               \
+                            num_simd_lanes);
+
+#define instantiate_mla_paged_attention(type, cache_type, num_simd_lanes)     \
+  instantiate_mla_latent_dims(type, cache_type, 32, num_simd_lanes);
+
+instantiate_mla_paged_attention(float, float, 32);
+instantiate_mla_paged_attention(bfloat16_t, bfloat16_t, 32);
+instantiate_mla_paged_attention(half, half, 32);
diff --git a/src/parallax_extensions/lib/_ext.cpython-312-darwin.so b/src/parallax_extensions/lib/_ext.cpython-312-darwin.so
index a47bf235..62fb6eac 100755
Binary files a/src/parallax_extensions/lib/_ext.cpython-312-darwin.so and b/src/parallax_extensions/lib/_ext.cpython-312-darwin.so differ
diff --git a/src/parallax_extensions/lib/libparallax_ext.dylib b/src/parallax_extensions/lib/libparallax_ext.dylib
index ff3c7463..0bac2a74 100755
Binary files a/src/parallax_extensions/lib/libparallax_ext.dylib and b/src/parallax_extensions/lib/libparallax_ext.dylib differ
diff --git a/src/parallax_extensions/lib/parallax_ext.metallib b/src/parallax_extensions/lib/parallax_ext.metallib
index e12aa3ca..6f19dc7a 100644
Binary files a/src/parallax_extensions/lib/parallax_ext.metallib and b/src/parallax_extensions/lib/parallax_ext.metallib differ
diff --git a/src/parallax_extensions/ops.py b/src/parallax_extensions/ops.py
index 5cd435d8..856177e3 100644
--- a/src/parallax_extensions/ops.py
+++ b/src/parallax_extensions/ops.py
@@ -46,6 +46,7 @@ def load_extension_module() -> ModuleType:
 
 _ext = load_extension_module()
 _ext_dsa_indexer_scores_with_update = _ext.dsa_indexer_scores_with_update
+_ext_mla_paged_attention = _ext.mla_paged_attention
 _ext_dsa_paged_attention = _ext.dsa_paged_attention
 _ext_paged_attention_v1 = _ext.paged_attention_v1
 _ext_paged_attention_v2 = _ext.paged_attention_v2
@@ -59,6 +60,57 @@ def load_extension_module() -> ModuleType:
 _PAGED_ATTENTION_V1_MAX_LENGTH = 8192
 
 
+def mla_paged_attention(
+    q_latent: mx.array,
+    q_pe: mx.array,
+    latent_cache: mx.array,
+    rope_cache: mx.array,
+    block_tables: mx.array,
+    context_lengths: mx.array,
+    block_size: int,
+    scale: float,
+) -> mx.array:
+    """
+    Dense MLA decode attention over compressed paged cache.
+
+    Computes:
+      softmax(scale * (q_latent @ latent_cache.T + q_pe @ rope_cache.T)) @ latent_cache
+    over the full logical context for each sequence.
+    """
+    if q_latent.ndim == 4:
+        if q_latent.shape[2] != 1:
+            raise ValueError("mla_paged_attention only supports one query token.")
+        q_latent = q_latent.squeeze(2)
+    if q_pe.ndim == 4:
+        if q_pe.shape[2] != 1:
+            raise ValueError("mla_paged_attention only supports one query token.")
+        q_pe = q_pe.squeeze(2)
+    if q_latent.ndim != 3 or q_pe.ndim != 3:
+        raise ValueError("q_latent and q_pe must be shaped (batch, heads, dim).")
+    if q_latent.shape[:2] != q_pe.shape[:2]:
+        raise ValueError("q_latent and q_pe batch/head dimensions must match.")
+    if latent_cache.ndim != 5 or rope_cache.ndim != 5:
+        raise ValueError("latent_cache and rope_cache must be paged cache tensors.")
+    if latent_cache.shape[2] != 1 or rope_cache.shape[2] != 1:
+        raise ValueError("mla_paged_attention expects one MLA cache head.")
+    if latent_cache.shape[-1] != q_latent.shape[-1]:
+        raise ValueError("latent cache dim must match q_latent dim.")
+    if rope_cache.shape[-1] != q_pe.shape[-1]:
+        raise ValueError("rope cache dim must match q_pe dim.")
+
+    output = _ext_mla_paged_attention(
+        mx.contiguous(q_latent),
+        mx.contiguous(q_pe),
+        latent_cache,
+        rope_cache,
+        mx.contiguous(block_tables.astype(mx.int32)),
+        mx.contiguous(context_lengths.astype(mx.int32)),
+        block_size,
+        scale,
+    )
+    return output[:, :, None, :]
+
+
 def store_indexer_cache(
     key: mx.array,
     key_cache: mx.array,
diff --git a/tests/parallax_extensions_tests/test_dsa_paged_attention.py b/tests/parallax_extensions_tests/test_dsa_paged_attention.py
index f4af54d5..d7056bab 100644
--- a/tests/parallax_extensions_tests/test_dsa_paged_attention.py
+++ b/tests/parallax_extensions_tests/test_dsa_paged_attention.py
@@ -4,7 +4,11 @@
 import pytest
 
 from parallax.server.cache.dsa_cache import DeepSeekSparseCache
-from parallax_extensions.ops import dsa_paged_attention, reshape_and_cache
+from parallax_extensions.ops import (
+    dsa_paged_attention,
+    mla_paged_attention,
+    reshape_and_cache,
+)
 
 pytestmark = pytest.mark.skipif(sys.platform != "darwin", reason="MLX tests require macOS")
 
@@ -146,6 +150,85 @@ def test_dsa_paged_attention_dense_minus_one_row_match_reference():
     assert mx.allclose(out[0], ref, rtol=1e-4, atol=1e-4)
 
 
+def test_mla_paged_attention_dense_matches_reference():
+    latent_seq = mx.array(
+        [
+            [0.1, -0.2, 0.3, -0.4],
+            [0.5, -0.6, 0.7, -0.8],
+            [0.9, -1.0, 1.1, -1.2],
+            [1.3, -1.4, 1.5, -1.6],
+            [1.7, -1.8, 1.9, -2.0],
+            [2.1, -2.2, 2.3, -2.4],
+        ],
+        dtype=mx.float32,
+    )
+    rope_seq = mx.array(
+        [
+            [0.2, 0.1],
+            [-0.3, 0.4],
+            [0.5, -0.6],
+            [0.7, 0.8],
+            [-0.9, 1.0],
+            [1.1, -1.2],
+        ],
+        dtype=mx.float32,
+    )
+    block_size = 4
+    cache = DeepSeekSparseCache(
+        num_blocks=2,
+        block_size=block_size,
+        num_kv_heads=1,
+        head_dim=latent_seq.shape[-1],
+        head_dim_v=latent_seq.shape[-1],
+        dtype=mx.float32,
+        kv_lora_rank=latent_seq.shape[-1],
+        qk_rope_head_dim=rope_seq.shape[-1],
+    )
+    block_tables = mx.array([[0, 1]], dtype=mx.int32)
+    context_lengths = mx.array([5], dtype=mx.int32)
+    reshape_and_cache(
+        latent_seq[None, :, None, :],
+        rope_seq[None, :, None, :],
+        cache.latent_cache,
+        cache.rope_cache,
+        block_tables,
+        mx.array([latent_seq.shape[0]], dtype=mx.int32),
+        block_size,
+        slot_mapping=mx.arange(latent_seq.shape[0], dtype=mx.int64),
+    )
+    mx.eval(cache.latent_cache, cache.rope_cache)
+
+    q_latent = mx.array(
+        [[[0.2, 0.1, -0.5, 0.4], [-0.3, 0.5, 0.2, -0.1]]],
+        dtype=mx.float32,
+    )
+    q_pe = mx.array([[[0.3, -0.2], [0.6, 0.1]]], dtype=mx.float32)
+    scale = 0.5
+
+    out = mla_paged_attention(
+        q_latent,
+        q_pe,
+        cache.latent_cache,
+        cache.rope_cache,
+        block_tables,
+        context_lengths,
+        block_size,
+        scale,
+    ).squeeze(2)
+    ref = _reference(
+        q_latent,
+        q_pe,
+        latent_seq,
+        rope_seq,
+        mx.arange(int(context_lengths[0]), dtype=mx.int32),
+        scale,
+    )
+    mx.eval(out, ref)
+
+    assert cache.get_indexer_cache() is None
+    assert mx.allclose(out[0], ref, rtol=1e-4, atol=1e-4)
+
+
 def test_dsa_paged_attention_glm_dimensions_smoke():
     latent_dim = 512
     rope_dim = 64
diff --git a/tests/parallax_extensions_tests/test_paged_attention_v1.py b/tests/parallax_extensions_tests/test_paged_attention_v1.py
index 31468166..95911a6f 100644
--- a/tests/parallax_extensions_tests/test_paged_attention_v1.py
+++ b/tests/parallax_extensions_tests/test_paged_attention_v1.py
@@ -5,15 +5,12 @@
 import numpy as np
 import pytest
 
-from parallax.metal.paged_attention.kernel import paged_attention as old_paged_attention
 from parallax_extensions.ops import (
     paged_attention_v1,
     reshape_and_cache,
     sparse_paged_attention,
 )
 
-old_reshape_and_cache = reshape_and_cache
-
 
 def get_packing_factor(dtype):
     if dtype == mx.float32:
@@ -129,57 +126,6 @@ def reference_paged_attention(
     return output
 
 
-def _bench_old(
-    q_old,
-    k_cache,
-    v_cache,
-    block_tables,
-    context_lengths,
-    block_size,
-    scale,
-    num_kv_heads,
-    iters,
-    warmup,
-    window_size=None,
-    sinks=None,
-):
-    kwargs = {}
-    if window_size is not None:
-        kwargs["window_size"] = window_size
-    if sinks is not None:
-        kwargs["sinks"] = sinks
-    for _ in range(warmup):
-        out = old_paged_attention(
-            q_old,
-            k_cache,
-            v_cache,
-            block_tables,
-            context_lengths,
-            block_size,
-            scale,
-            num_kv_heads,
-            **kwargs,
-        )
-        mx.eval(out)
-    mx.synchronize()
-    start = time.perf_counter()
-    for _ in range(iters):
-        out = old_paged_attention(
-            q_old,
-            k_cache,
-            v_cache,
-            block_tables,
-            context_lengths,
-            block_size,
-            scale,
-            num_kv_heads,
-            **kwargs,
-        )
-        mx.eval(out)
-    mx.synchronize()
-    return (time.perf_counter() - start) / iters * 1000.0
-
-
 def _bench_new(
     q,
     k_cache,
@@ -558,20 +504,6 @@ def test_reference_cases(self):
         k_prefill = k_seq.transpose(0, 2, 1, 3)
         v_prefill = v_seq.transpose(0, 2, 1, 3)
 
-        old_key_cache = mx.zeros((1, num_blocks, num_kv_heads, block_size, head_dim), dtype=dtype)
-        old_value_cache = mx.zeros((1, num_blocks, num_kv_heads, block_size, head_dim), dtype=dtype)
-
-        old_reshape_and_cache(
-            k_prefill,
-            v_prefill,
-            old_key_cache,
-            old_value_cache,
-            block_tables,
-            context_lengths,
-            block_size,
-            slot_mapping=slot_mapping,
-        )
-
         x = get_packing_factor(dtype)
         new_key_cache = mx.zeros(
             (num_blocks, num_kv_heads, head_dim // x, block_size, x), dtype=dtype
@@ -589,12 +521,12 @@ def test_reference_cases(self):
             slot_mapping=slot_mapping,
         )
 
-        mx.eval(old_key_cache, old_value_cache, new_key_cache, new_value_cache)
+        mx.eval(new_key_cache, new_value_cache)
         mx.synchronize()
 
         sinks = mx.array(np.random.uniform(-0.5, 0.5, size=(num_heads,)), dtype=mx.float32)
         window_sizes = [block_size // 2, block_size]
-        cases = [{"name": "baseline", "window_size": 0, "sinks": None, "use_old": True}]
+        cases = [{"name": "baseline", "window_size": 0, "sinks": None}]
         for ws in window_sizes:
             cases.append({"name": f"window-{ws}", "window_size": ws, "sinks": None})
         cases.append({"name": "sink-only", "window_size": 0, "sinks": sinks})
@@ -604,7 +536,6 @@ def test_reference_cases(self):
         for case in cases:
             window_size = case["window_size"]
             sinks_case = case["sinks"]
-            use_old = case.get("use_old", False)
 
             ref_output = reference_paged_attention(
                 q,
@@ -619,24 +550,6 @@ def test_reference_cases(self):
             mx.eval(ref_output)
 
             atol = 1e-2 if sinks_case is not None else 1e-3
-            if use_old:
-                old_output = old_paged_attention(
-                    q,
-                    old_key_cache,
-                    old_value_cache,
-                    block_tables,
-                    context_lengths,
-                    block_size,
-                    scale,
-                    num_kv_heads,
-                )
-                mx.eval(old_output)
-                diff_old = mx.abs(ref_output.astype(mx.float32) - old_output.astype(mx.float32))
-                max_diff_old = mx.max(diff_old).item()
-                assert mx.allclose(
-                    ref_output.astype(mx.float32), old_output.astype(mx.float32), atol=1e-3
-                ).item(), f"Old kernel mismatch (max diff {max_diff_old:.6f})"
-
             new_output = paged_attention_v1(
                 q,
                 new_key_cache,
@@ -681,11 +594,6 @@ def test_benchmark_paged_vs_native(self):
         q = mx.random.normal(shape=(bs, n_heads, dim)).astype(dtype)
         k_cont = mx.random.normal(shape=(bs, n_kv_heads, seq_len, dim)).astype(dtype)
         v_cont = mx.random.normal(shape=(bs, n_kv_heads, seq_len, dim)).astype(dtype)
-        q_old = q[:, :, None, :]
-
-        old_k_cache = mx.random.normal((1, total_blocks, n_kv_heads, block_size, dim)).astype(dtype)
-        old_v_cache = mx.random.normal((1, total_blocks, n_kv_heads, block_size, dim)).astype(dtype)
-        mx.eval(old_k_cache, old_v_cache)
 
         print(f"\n[Benchmark BS={bs}, Len={seq_len}, Float16]")
 
@@ -714,21 +622,6 @@ def test_benchmark_paged_vs_native(self):
             sinks_case = case["sinks"]
             label = case["name"]
 
-            old_ws = None if ws == 0 else ws
-            old_ms = _bench_old(
-                q_old,
-                old_k_cache,
-                old_v_cache,
-                block_tables,
-                seq_lens,
-                block_size,
-                scale,
-                n_kv_heads,
-                iters=100,
-                warmup=10,
-                window_size=old_ws,
-                sinks=sinks_case,
-            )
             new_ms = _bench_new(
                 q,
                 key_cache,
@@ -743,7 +636,4 @@ def test_benchmark_paged_vs_native(self):
                 window_size=ws,
                 sinks=sinks_case,
             )
-            speedup = old_ms / new_ms if new_ms > 0 else float("inf")
-            print(f"Old kernel ({label}): {old_ms:.3f} ms")
             print(f"Paged Attn V1 ({label}): {new_ms:.3f} ms")
-            print(f"Speedup ({label}): {speedup:.2f}x")
diff --git a/tests/test_deepseek_v32.py b/tests/test_deepseek_v32.py
index 8ed9ebc3..b55d4bb8 100644
--- a/tests/test_deepseek_v32.py
+++ b/tests/test_deepseek_v32.py
@@ -14,7 +14,9 @@
     derive_indexer_types,
 )
 from parallax.server.cache.dsa_cache import DeepSeekSparseCache
+from parallax.server.cache.kv_cache import KVCachePacked
 from parallax.server.cache_manager import CacheManager
+from parallax.utils.layer_types import DSA_ATTENTION, MLA_ATTENTION
 from parallax.utils.utils import create_causal_mask
 
 
@@ -167,6 +169,7 @@ def test_cache_manager_counts_compressed_mla_and_unexpanded_indexer_bytes(monkey
         index_key_heads=1,
         kv_lora_rank=4,
         qk_rope_head_dim=2,
+        layer_types=[DSA_ATTENTION, DSA_ATTENTION],
     )
 
     dtype_size = 4
@@ -174,6 +177,62 @@ def test_cache_manager_counts_compressed_mla_and_unexpanded_indexer_bytes(monkey
     assert manager._calculate_kv_block_bytes(dtype_size) == expected
 
 
+def test_cache_manager_counts_dense_mla_without_indexer(monkeypatch):
+    monkeypatch.setattr(
+        mx.metal,
+        "device_info",
+        lambda: {"max_recommended_working_set_size": 1024 * 1024},
+    )
+    monkeypatch.setattr(mx, "get_active_memory", lambda: 0)
+
+    manager = CacheManager(
+        num_layers=2,
+        num_kv_heads=2,
+        head_dim=4,
+        head_dim_v=4,
+        dtype=mx.float32,
+        block_size=4,
+        cache_memory_fraction=1.0,
+        kv_lora_rank=4,
+        qk_rope_head_dim=2,
+        layer_types=[MLA_ATTENTION, MLA_ATTENTION],
+    )
+
+    dtype_size = 4
+    expected = 2 * 4 * (4 + 2) * dtype_size
+    assert manager._calculate_kv_block_bytes(dtype_size) == expected
+    assert isinstance(manager.caches[0], DeepSeekSparseCache)
+    assert manager.caches[0].get_indexer_cache() is None
+
+
+def test_cache_manager_does_not_infer_cache_type_from_sparse_params(monkeypatch):
+    monkeypatch.setattr(
+        mx.metal,
+        "device_info",
+        lambda: {"max_recommended_working_set_size": 1024 * 1024},
+    )
+    monkeypatch.setattr(mx, "get_active_memory", lambda: 0)
+
+    manager = CacheManager(
+        num_layers=1,
+        num_kv_heads=2,
+        head_dim=4,
+        head_dim_v=4,
+        dtype=mx.float32,
+        block_size=4,
+        cache_memory_fraction=1.0,
+        index_head_dim=4,
+        index_n_heads=2,
+        kv_lora_rank=4,
+        qk_rope_head_dim=2,
+    )
+
+    dtype_size = 4
+    expected = 2 * 4 * (4 + 4) * dtype_size
+    assert manager._calculate_kv_block_bytes(dtype_size) == expected
+    assert isinstance(manager.caches[0], KVCachePacked)
+
+
 def test_indexer_types_default_to_all_full():
     assert derive_indexer_types(4) == ["full", "full", "full", "full"]
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 5c35988b..79ddca8f 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -2,6 +2,12 @@
 from unittest.mock import patch
 
 from parallax.utils import utils
+from parallax.utils.layer_types import (
+    DSA_ATTENTION,
+    LINEAR,
+    MLA_ATTENTION,
+    MSA_ATTENTION,
+)
 
 
 def test_is_metal_available_uses_mlx_metal_is_available(monkeypatch):
@@ -56,3 +62,50 @@ def test_load_config_only_downloads_config_json(tmp_path):
         filename="config.json",
         local_files_only=True,
     )
+
+
+def test_get_layer_types_marks_dense_mla_attention():
+    config = {
+        "model_type": "deepseek_v3",
+        "kv_lora_rank": 512,
+        "qk_rope_head_dim": 64,
+    }
+
+    assert utils.get_layer_types(config, 0, 2) == [MLA_ATTENTION, MLA_ATTENTION]
+
+
+def test_get_layer_types_marks_dsa_attention():
+    config = {
+        "model_type": "glm_moe_dsa",
+        "kv_lora_rank": 512,
+        "qk_rope_head_dim": 64,
+        "index_head_dim": 128,
+        "index_n_heads": 4,
+    }
+
+    assert utils.get_layer_types(config, 0, 2) == [DSA_ATTENTION, DSA_ATTENTION]
+
+
+def test_get_layer_types_preserves_hybrid_linear_layers_with_mla_cache():
+    config = {
+        "model_type": "kimi_k2",
+        "kv_lora_rank": 512,
+        "qk_rope_head_dim": 64,
+        "linear_attn_config": {"full_attn_layers": [1, 3]},
+    }
+
+    assert utils.get_layer_types(config, 0, 4) == [
+        LINEAR,
+        MLA_ATTENTION,
+        LINEAR,
+        MLA_ATTENTION,
+    ]
+
+
+def test_get_layer_types_marks_msa_attention_cache():
+    config = {"model_type": "minimax_m3"}
+
+    assert utils.get_layer_types(config, 0, 2) == [
+        MSA_ATTENTION,
+        MSA_ATTENTION,
+    ]