diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index d8efa17359..ade638473e 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -457,6 +457,7 @@ o_lora_rank: 0 # Output LoRA rank for Compressed Attention.
 o_groups: 0 # Output groups for Compressed Attention.
 compress_ratios: [] # Per-layer compression ratios (0, 4, 128, etc).
 compressed_rope_max_timescale: 160_000 # If positive, used for Compressed Sparse/Heavy Attention.
+use_csa_streamindex_kernel: false # Whether to use Pallas TPU kernel for CSA StreamIndex score computation.
 
 # QK-Clip (Muon Clip) Configuration
 use_qk_clip: false  # Enable QK-Clip (supported in MLA with DotProduct or Tokamax Splash)
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index db9a68d1cd..c97cc808fa 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -712,6 +712,10 @@ class CompressedAttention(BaseModel):
   compressed_rope_max_timescale: int = Field(
       160000, description="If positive, used for Compressed Sparse/Heavy Attention."
   )
+  use_csa_streamindex_kernel: bool = Field(
+      False,
+      description="Whether to use Pallas TPU kernel for CSA StreamIndex score computation.",
+  )
 
 
 class AttentionIndexer(BaseModel):
diff --git a/src/maxtext/kernels/attention/__init__.py b/src/maxtext/kernels/attention/__init__.py
index 3bf51ab4d6..bb28108248 100644
--- a/src/maxtext/kernels/attention/__init__.py
+++ b/src/maxtext/kernels/attention/__init__.py
@@ -14,4 +14,5 @@
 
 """Attention kernels."""
 
+from maxtext.kernels.attention import csa_streamindex
 from maxtext.kernels.attention import splash_attention_kernel
diff --git a/src/maxtext/kernels/attention/csa_streamindex.py b/src/maxtext/kernels/attention/csa_streamindex.py
new file mode 100644
index 0000000000..5696f1042c
--- /dev/null
+++ b/src/maxtext/kernels/attention/csa_streamindex.py
@@ -0,0 +1,233 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Fused Pallas TPU kernel for DeepSeek-V4 CSA StreamIndex Score Computation.
+
+Computes:
+  index_scores = sum_h(ReLU(q_h @ comp^T) * softmax_scale * w_h)
+with optional in-VMEM causal future masking.
+
+The kernel fuses dot-product, ReLU activation, softmax scaling, head-weight
+contraction, and causal future masking into on-chip TPU VMEM registers, avoiding
+the materialization of the intermediate [B, H, S, W] 4D tensor in HBM.
+"""
+
+import functools
+import jax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+
+
+def csa_streamindex_score_kernel(
+    q_ref,        # [num_heads, block_q, head_dim]
+    k_ref,        # [block_w, head_dim]
+    w_ref,        # [block_q, num_heads]
+    out_ref,      # [block_q, block_w]
+    *,
+    softmax_scale: float,
+    compress_rate: int = 0,
+):
+  """Fused Pallas TPU kernel with 2D MXU matmul."""
+  num_heads, block_q, head_dim = q_ref.shape
+  block_w, _ = k_ref.shape
+
+  # Reshape Q to 2D: (num_heads * block_q, head_dim) for native 2D systolic array MXU contraction
+  q_2d = q_ref[...].reshape(num_heads * block_q, head_dim)
+  k_2d = k_ref[...]
+
+  # 2D MXU matmul: (num_heads * block_q, head_dim) @ (block_w, head_dim)^T -> (num_heads * block_q, block_w)
+  s_2d = jnp.einsum("nd,md->nm", q_2d, k_2d, preferred_element_type=jnp.float32)
+
+  # Reshape to (num_heads, block_q, block_w) and apply ReLU
+  s = s_2d.reshape(num_heads, block_q, block_w)
+  s = jnp.maximum(s, 0.0)
+
+  # Multiply by weights and sum across heads in VMEM
+  w = w_ref[...].astype(jnp.float32).transpose(1, 0)[:, :, None]
+  s_weighted = jnp.sum(s * w, axis=0) * softmax_scale
+
+  # In-VMEM causal future masking
+  if compress_rate > 0:
+    i = pl.program_id(1)
+    j = pl.program_id(2)
+    q_indices = i * block_q + jnp.arange(block_q, dtype=jnp.int32)[:, None]
+    k_indices = (j * block_w + jnp.arange(block_w, dtype=jnp.int32)[None, :]) * compress_rate
+    future_mask = (k_indices + compress_rate) > (q_indices + 1)
+    s_weighted = jnp.where(future_mask, -1e9, s_weighted)
+
+  out_ref[...] = s_weighted.astype(out_ref.dtype)
+
+
+def _csa_streamindex_score_pallas_fwd(
+    q: jax.Array,
+    compressed: jax.Array,
+    weights: jax.Array,
+    *,
+    softmax_scale: float,
+    compress_rate: int = 0,
+    block_q: int | None = None,
+    block_w: int | None = None,
+    interpret: bool = False,
+) -> jax.Array:
+  """Forward implementation using fused Pallas TPU kernel for head-major [B, H, S, D] q."""
+  batch_size, num_heads, seq_len, head_dim = q.shape
+  _, compressed_len, comp_head_dim = compressed.shape
+  assert comp_head_dim == head_dim, f"{comp_head_dim=} != {head_dim=}"
+  assert weights.shape == (batch_size, seq_len, num_heads), f"{weights.shape=} != {(batch_size, seq_len, num_heads)=}"
+
+  if block_q is None:
+    block_q = 128 if num_heads >= 32 else 256
+  if block_w is None:
+    block_w = 1024 if num_heads >= 32 else 2048
+
+  padded_s = ((seq_len + block_q - 1) // block_q) * block_q
+  padded_w = ((compressed_len + block_w - 1) // block_w) * block_w
+
+  if padded_s > seq_len:
+    pad_s = padded_s - seq_len
+    q = jnp.pad(q, ((0, 0), (0, 0), (0, pad_s), (0, 0)))
+    weights = jnp.pad(weights, ((0, 0), (0, pad_s), (0, 0)))
+  if padded_w > compressed_len:
+    pad_w = padded_w - compressed_len
+    compressed = jnp.pad(compressed, ((0, 0), (0, pad_w), (0, 0)))
+
+  grid = (batch_size, padded_s // block_q, padded_w // block_w)
+
+  in_specs = [
+      pl.BlockSpec((None, num_heads, block_q, head_dim), lambda b, i, j: (b, 0, i, 0)),
+      pl.BlockSpec((None, block_w, head_dim), lambda b, i, j: (b, j, 0)),
+      pl.BlockSpec((None, block_q, num_heads), lambda b, i, j: (b, i, 0)),
+  ]
+  out_specs = pl.BlockSpec((None, block_q, block_w), lambda b, i, j: (b, i, j))
+
+  out = pl.pallas_call(
+      functools.partial(
+          csa_streamindex_score_kernel,
+          softmax_scale=softmax_scale,
+          compress_rate=compress_rate,
+      ),
+      in_specs=in_specs,
+      out_specs=out_specs,
+      grid=grid,
+      compiler_params=pltpu.CompilerParams(
+          dimension_semantics=("parallel", "parallel", "arbitrary"),
+      ),
+      out_shape=jax.ShapeDtypeStruct((batch_size, padded_s, padded_w), jnp.float32),
+      interpret=interpret,
+  )(q, compressed, weights)
+
+  return out[:, :seq_len, :compressed_len]
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(3, 4, 5, 6, 7))
+def csa_streamindex_score(
+    q: jax.Array,
+    compressed: jax.Array,
+    weights: jax.Array,
+    softmax_scale: float,
+    compress_rate: int = 0,
+    block_q: int | None = None,
+    block_w: int | None = None,
+    interpret: bool = False,
+) -> jax.Array:
+  """Computes CSA StreamIndex scores using head-major [B, H, S, D] q layout with 2D MXU matmul."""
+  return _csa_streamindex_score_pallas_fwd(
+      q,
+      compressed,
+      weights,
+      softmax_scale=softmax_scale,
+      compress_rate=compress_rate,
+      block_q=block_q,
+      block_w=block_w,
+      interpret=interpret,
+  )
+
+
+def _csa_streamindex_score_fwd(
+    q: jax.Array,
+    compressed: jax.Array,
+    weights: jax.Array,
+    softmax_scale: float,
+    compress_rate: int = 0,
+    block_q: int | None = None,
+    block_w: int | None = None,
+    interpret: bool = False,
+) -> tuple[jax.Array, tuple[jax.Array, jax.Array, jax.Array]]:
+  out = _csa_streamindex_score_pallas_fwd(
+      q,
+      compressed,
+      weights,
+      softmax_scale=softmax_scale,
+      compress_rate=compress_rate,
+      block_q=block_q,
+      block_w=block_w,
+      interpret=interpret,
+  )
+  return out, (q, compressed, weights)
+
+
+def _csa_streamindex_score_bwd(
+    softmax_scale: float,
+    compress_rate: int,
+    block_q: int | None,
+    block_w: int | None,
+    interpret: bool,
+    res: tuple[jax.Array, jax.Array, jax.Array],
+    g: jax.Array,
+) -> tuple[jax.Array, jax.Array, jax.Array]:
+  del block_q, block_w, interpret
+  q, compressed, weights = res
+  _, vjp_fn = jax.vjp(
+      functools.partial(
+          reference_csa_streamindex_score,
+          softmax_scale=softmax_scale,
+          compress_rate=compress_rate,
+      ),
+      q,
+      compressed,
+      weights,
+  )
+  dq, dk, dw = vjp_fn(g)
+  return dq, dk, dw
+
+
+csa_streamindex_score.defvjp(
+    _csa_streamindex_score_fwd, _csa_streamindex_score_bwd
+)
+
+
+def reference_csa_streamindex_score(
+    q: jax.Array,
+    compressed: jax.Array,
+    weights: jax.Array,
+    *,
+    softmax_scale: float,
+    compress_rate: int = 0,
+) -> jax.Array:
+  """Reference score computation matching the pure JAX einsum path."""
+  scores = jnp.einsum("bhsd,bwd->bhsw", q.astype(jnp.float32), compressed.astype(jnp.float32))
+  scores = jax.nn.relu(scores) * softmax_scale
+  index_scores = jnp.einsum("bhsw,bsh->bsw", scores, weights.astype(jnp.float32))
+  if compress_rate > 0:
+    seq_len = q.shape[2]
+    compressed_len = compressed.shape[1]
+    position_ids = jnp.arange(seq_len, dtype=jnp.int32)[None, :]
+    usable_len = compressed_len * compress_rate
+    block_positions = position_ids[:, :usable_len:compress_rate]
+    future_mask = (block_positions[:, None, :] + compress_rate) > (position_ids[:, :, None] + 1)
+    index_scores = jnp.where(future_mask, -1e9, index_scores)
+  return index_scores
+
+
diff --git a/src/maxtext/layers/attention_compressed.py b/src/maxtext/layers/attention_compressed.py
index 6530a8b0fb..1b93e9f808 100644
--- a/src/maxtext/layers/attention_compressed.py
+++ b/src/maxtext/layers/attention_compressed.py
@@ -45,6 +45,7 @@
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.inference.kvcache import KVQuant
 from maxtext.inference import kvcache
+from maxtext.kernels.attention import csa_streamindex
 
 
 class CSAPoolingConfig(enum.IntEnum):
@@ -874,20 +875,28 @@ def indexer_compressor_fn(buf_kv, buf_gate):
       return jnp.zeros((batch_size, seq_len, min(self.index_topk, compressed_len)), dtype=jnp.int32)
 
     # --- TOP-K ROUTING MATH (Executes in both Prefill and AR) ---
-    compressed_kv = jnp.expand_dims(compressed, axis=1)
-    compressed_kv = jnp.broadcast_to(compressed_kv, (batch_size, self.index_n_heads, compressed_len, self.index_head_dim))
-
     q = self.q_proj(q_latent).reshape((batch_size, seq_len, self.index_n_heads, self.index_head_dim))
     q = jnp.transpose(q, (0, 2, 1, 3))
     q = self.rotary_emb(q, position_ids, unsqueeze_dim=1)
-
-    q = q.astype(jnp.float32)
-    compressed_kv = compressed_kv.astype(jnp.float32)
-
-    scores = jnp.einsum("bhsd,bhwd->bhsw", q, compressed_kv)
-    scores = jax.nn.relu(scores) * self.softmax_scale
     weights = self.weights_proj(hidden_states).astype(jnp.float32) * self.weights_scaling
-    index_scores = jnp.einsum("bhsw,bsh->bsw", scores, weights)
+    if self.config.use_csa_streamindex_kernel:
+      index_scores = csa_streamindex.csa_streamindex_score(
+          q=q,
+          compressed=compressed,
+          weights=weights,
+          softmax_scale=self.softmax_scale,
+          compress_rate=self.compress_rate,
+      )
+    else:
+      compressed_kv = jnp.expand_dims(compressed, axis=1)
+      compressed_kv = jnp.broadcast_to(
+          compressed_kv, (batch_size, self.index_n_heads, compressed_len, self.index_head_dim)
+      )
+      q = q.astype(jnp.float32)
+      compressed_kv = compressed_kv.astype(jnp.float32)
+      scores = jnp.einsum("bhsd,bhwd->bhsw", q, compressed_kv)
+      scores = jax.nn.relu(scores) * self.softmax_scale
+      index_scores = jnp.einsum("bhsw,bsh->bsw", scores, weights)
 
     k = min(self.index_topk, compressed_len)
 
@@ -897,8 +906,9 @@ def indexer_compressor_fn(buf_kv, buf_gate):
       block_positions = position_ids[:, : usable_len : self.compress_rate]
       future_mask = (block_positions[:, None, :] + self.compress_rate) > (position_ids[:, :, None] + 1)
 
-    # Apply the mask to the scores
-    index_scores = jnp.where(future_mask, jnp.full_like(index_scores, -jnp.inf), index_scores)
+    # Apply the mask to the scores if not already applied by the kernel
+    if not self.config.use_csa_streamindex_kernel:
+      index_scores = jnp.where(future_mask, jnp.full_like(index_scores, -jnp.inf), index_scores)
 
     combined_invalid = future_mask
     if attention_mask is not None:
diff --git a/tests/unit/csa_streamindex_test.py b/tests/unit/csa_streamindex_test.py
new file mode 100644
index 0000000000..1516554760
--- /dev/null
+++ b/tests/unit/csa_streamindex_test.py
@@ -0,0 +1,275 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for DeepSeek-V4 CSA StreamIndex Pallas TPU score kernel."""
+
+import unittest
+from unittest import mock
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+import numpy as np
+
+from maxtext.configs.pyconfig import initialize
+from maxtext.kernels.attention import csa_streamindex
+from maxtext.layers.attention_compressed import DeepseekV4Indexer
+from maxtext.layers.embeddings import DeepSeekV4RotaryEmbedding
+from tests.utils.test_helpers import get_test_config_path
+
+
+class TestCsaStreamIndexScoreKernel(unittest.TestCase):
+  """Unit tests for Pallas CSA StreamIndex score kernel."""
+
+  def setUp(self):
+    self.key = jax.random.PRNGKey(42)
+
+  def test_parity_exact_multiple(self):
+    """Verifies numerical parity matching reference einsum on exact block multiples."""
+    key1, key2, key3 = jax.random.split(self.key, 3)
+    b, h, s, w, d = 2, 64, 256, 128, 128
+    q = jax.random.normal(key1, (b, h, s, d), dtype=jnp.bfloat16)
+    compressed = jax.random.normal(key2, (b, w, d), dtype=jnp.bfloat16)
+    weights = jax.random.normal(key3, (b, s, h), dtype=jnp.float32)
+    softmax_scale = d**-0.5
+
+    expected = csa_streamindex.reference_csa_streamindex_score(
+        q, compressed, weights, softmax_scale=softmax_scale
+    )
+    actual = csa_streamindex.csa_streamindex_score(
+        q,
+        compressed,
+        weights,
+        softmax_scale=softmax_scale,
+        block_q=128,
+        block_w=128,
+        interpret=True,
+    )
+    np.testing.assert_allclose(actual, expected, rtol=1e-1, atol=1e-1)
+
+  def test_parity_non_multiples(self):
+    """Verifies padding handling when seq_len and compressed_len are not multiples of block sizes."""
+    key1, key2, key3 = jax.random.split(self.key, 3)
+    b, h, s, w, d = 2, 32, 150, 70, 64
+    q = jax.random.normal(key1, (b, h, s, d), dtype=jnp.bfloat16)
+    compressed = jax.random.normal(key2, (b, w, d), dtype=jnp.bfloat16)
+    weights = jax.random.normal(key3, (b, s, h), dtype=jnp.float32)
+    softmax_scale = d**-0.5
+
+    expected = csa_streamindex.reference_csa_streamindex_score(
+        q, compressed, weights, softmax_scale=softmax_scale
+    )
+    actual = csa_streamindex.csa_streamindex_score(
+        q,
+        compressed,
+        weights,
+        softmax_scale=softmax_scale,
+        block_q=128,
+        block_w=128,
+        interpret=True,
+    )
+    np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-2)
+
+  def test_causal_mask_parity(self):
+    """Verifies numerical parity with in-VMEM causal masking."""
+    key1, key2, key3 = jax.random.split(self.key, 3)
+    b, h, s, w, d = 1, 4, 256, 64, 64
+    q = jax.random.normal(key1, (b, h, s, d), dtype=jnp.bfloat16)
+    compressed = jax.random.normal(key2, (b, w, d), dtype=jnp.bfloat16)
+    weights = jax.random.normal(key3, (b, s, h), dtype=jnp.float32)
+    softmax_scale = d**-0.5
+    compress_rate = 4
+
+    expected = csa_streamindex.reference_csa_streamindex_score(
+        q, compressed, weights, softmax_scale=softmax_scale, compress_rate=compress_rate
+    )
+    actual = csa_streamindex.csa_streamindex_score(
+        q,
+        compressed,
+        weights,
+        softmax_scale=softmax_scale,
+        compress_rate=compress_rate,
+        block_q=128,
+        block_w=128,
+        interpret=True,
+    )
+    np.testing.assert_allclose(actual, expected, rtol=1e-1, atol=1e-1)
+
+  def test_gradient_parity(self):
+    """Verifies that custom_vjp gradients match reference autograd."""
+    b, h, s, w, d = 1, 4, 128, 128, 32
+    key1, key2, key3 = jax.random.split(self.key, 3)
+    q = jax.random.normal(key1, (b, h, s, d), dtype=jnp.bfloat16)
+    comp = jax.random.normal(key2, (b, w, d), dtype=jnp.bfloat16)
+    weights = jax.random.normal(key3, (b, s, h), dtype=jnp.float32)
+    scale = 32.0**-0.5
+
+    def loss_kernel(q, comp, weights):
+      return jnp.sum(
+          csa_streamindex.csa_streamindex_score(
+              q, comp, weights, softmax_scale=scale, block_q=128, block_w=128, interpret=True
+          )
+      )
+
+    def loss_ref(q, comp, weights):
+      return jnp.sum(
+          csa_streamindex.reference_csa_streamindex_score(
+              q, comp, weights, softmax_scale=scale
+          )
+      )
+
+    g_q_k, g_c_k, g_w_k = jax.grad(loss_kernel, argnums=(0, 1, 2))(q, comp, weights)
+    g_q_r, g_c_r, g_w_r = jax.grad(loss_ref, argnums=(0, 1, 2))(q, comp, weights)
+
+    np.testing.assert_allclose(g_q_k, g_q_r, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(g_c_k, g_c_r, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(g_w_k, g_w_r, rtol=1e-3, atol=1e-3)
+
+  def test_tpu_compile_smoke_production_tiles(self):
+    """Compiles and executes with interpret=False on TPU hardware (DeepSeek-V4 production shapes)."""
+    if jax.default_backend() != "tpu":
+      self.skipTest("TPU hardware required for Mosaic compilation smoke test.")
+    b, h, s, d = 1, 64, 4096, 128
+    w = s // 4
+    scale = d**-0.5
+    key1, key2, key3 = jax.random.split(self.key, 3)
+    q = jax.random.normal(key1, (b, h, s, d), dtype=jnp.bfloat16)
+    compressed = jax.random.normal(key2, (b, w, d), dtype=jnp.bfloat16)
+    weights = jax.random.normal(key3, (b, s, h), dtype=jnp.float32)
+
+    fn = jax.jit(
+        lambda q, k, w: csa_streamindex.csa_streamindex_score(
+            q, k, w, softmax_scale=scale, block_q=128, block_w=512, interpret=False
+        )
+    )
+    out = fn(q, compressed, weights).block_until_ready()
+    self.assertEqual(out.shape, (b, s, w))
+
+
+class TestDeepseekv4IndexerIntegration(unittest.TestCase):
+  """Integration tests for Deepseekv4Indexer with CSA StreamIndex kernel dispatch."""
+
+  def setUp(self):
+    self.mesh = Mesh(jax.devices(), ("data",))
+    self.rotary = DeepSeekV4RotaryEmbedding(
+        head_dim=64,
+        partial_rotary_factor=16.0 / 64.0,
+        mesh=self.mesh,
+    )
+
+  def _get_config(self, use_csa_streamindex_kernel: bool = False):
+    with mock.patch("maxtext.utils.max_utils.maybe_initialize_jax_distributed_system"):
+      return initialize(
+          [
+              None,
+              get_test_config_path(),
+              "model_name=deepseek4-284b",
+              "attention=dot_product",
+              "qk_rope_head_dim=16",
+              "v_head_dim=16",
+              "qk_nope_head_dim=16",
+              "indexer_n_heads=16",
+              "indexer_head_dim=64",
+              "indexer_topk=32",
+              "override_model_config=True",
+              f"use_csa_streamindex_kernel={use_csa_streamindex_kernel}",
+          ]
+      )
+
+  def test_indexer_kernel_vs_einsum_output_parity(self):
+    """Verifies that Deepseekv4Indexer outputs match whether kernel or einsum is used."""
+    config_einsum = self._get_config(use_csa_streamindex_kernel=False)
+    config_kernel = self._get_config(use_csa_streamindex_kernel=True)
+    b, s, emb_dim, q_lora = 1, 128, config_einsum.emb_dim, config_einsum.q_lora_rank
+
+    indexer_einsum = DeepseekV4Indexer(
+        config=config_einsum,
+        compress_ratio=4,
+        rotary_embedding=self.rotary,
+        rngs=nnx.Rngs(0),
+    )
+    indexer_kernel = DeepseekV4Indexer(
+        config=config_kernel,
+        compress_ratio=4,
+        rotary_embedding=self.rotary,
+        rngs=nnx.Rngs(0),
+    )
+
+    key1, key2 = jax.random.split(jax.random.PRNGKey(0))
+    hidden = jax.random.normal(key1, (b, s, emb_dim), dtype=jnp.bfloat16)
+    q_latent = jax.random.normal(key2, (b, s, q_lora), dtype=jnp.bfloat16)
+    pos = jnp.arange(s, dtype=jnp.int32)[None, :]
+
+    out_einsum = indexer_einsum(hidden, q_latent, pos)
+
+    real_kernel_fn = csa_streamindex.csa_streamindex_score
+
+    def interpret_kernel_fn(*args, **kwargs):
+      kwargs["interpret"] = True
+      return real_kernel_fn(*args, **kwargs)
+
+    with mock.patch.object(csa_streamindex, "csa_streamindex_score", side_effect=interpret_kernel_fn):
+      out_kernel = indexer_kernel(hidden, q_latent, pos)
+
+    np.testing.assert_array_equal(out_kernel, out_einsum)
+
+  def test_indexer_einsum_when_disabled(self):
+    """Verifies that when use_csa_streamindex_kernel=False, einsum path is used."""
+    config_einsum = self._get_config(use_csa_streamindex_kernel=False)
+    b, s, emb_dim, q_lora = 1, 128, config_einsum.emb_dim, config_einsum.q_lora_rank
+
+    indexer = DeepseekV4Indexer(
+        config=config_einsum,
+        compress_ratio=4,
+        rotary_embedding=self.rotary,
+        rngs=nnx.Rngs(0),
+    )
+
+    hidden = jnp.ones((b, s, emb_dim), dtype=jnp.bfloat16)
+    q_latent = jnp.ones((b, s, q_lora), dtype=jnp.bfloat16)
+    pos = jnp.arange(s, dtype=jnp.int32)[None, :]
+
+    with mock.patch.object(csa_streamindex, "csa_streamindex_score") as mock_kernel:
+      out = indexer(hidden, q_latent, pos)
+      mock_kernel.assert_not_called()
+      self.assertEqual(out.shape, (b, s, min(32, s // 4)))
+
+  def test_jaxpr_verification(self):
+    """Verifies that jaxpr contains pallas_call when enabled and dot_general when disabled."""
+    q = jnp.zeros((1, 4, 256, 64), dtype=jnp.bfloat16)
+    compressed = jnp.zeros((1, 128, 64), dtype=jnp.bfloat16)
+    weights = jnp.zeros((1, 256, 4), dtype=jnp.float32)
+    scale = 64.0**-0.5
+
+    def compute_scores(q, compressed, weights, use_kernel):
+      if use_kernel:
+        return csa_streamindex.csa_streamindex_score(
+            q, compressed, weights, softmax_scale=scale, block_q=128, block_w=128
+        )
+      else:
+        return csa_streamindex.reference_csa_streamindex_score(
+            q, compressed, weights, softmax_scale=scale
+        )
+
+    jaxpr_kernel = jax.make_jaxpr(compute_scores, static_argnums=(3,))(q, compressed, weights, True)
+    self.assertIn("pallas_call", str(jaxpr_kernel))
+
+    jaxpr_einsum = jax.make_jaxpr(compute_scores, static_argnums=(3,))(q, compressed, weights, False)
+    self.assertNotIn("pallas_call", str(jaxpr_einsum))
+    self.assertIn("dot_general", str(jaxpr_einsum))
+
+
+if __name__ == "__main__":
+  unittest.main()