tpu-inference 0.11.1.dev202511180814__py3-none-any.whl → 0.12.0.dev20251213__py3-none-any.whl

tests/kernels/fused_moe_v1_test.py

@@ -1,7 +1,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from absl.testing import absltest
+from absl.testing import absltest, parameterized
 from jax._src import test_util as jtu
 from jax.sharding import Mesh
 
@@ -10,6 +10,15 @@ from tpu_inference.kernels.fused_moe.v1.kernel import fused_ep_moe, ref_moe
 jax.config.parse_flags_with_absl()
 
 
+def cdiv(a, b):
+    assert b != 0
+    return (a + b - 1) // b
+
+
+def align_to(x, a):
+    return cdiv(x, a) * a
+
+
 def gen_moe_inputs(
     dtype,
     top_k,
@@ -19,11 +28,14 @@ def gen_moe_inputs(
     num_tokens,
     *,
     seed=1234,
+    has_bias=False,
 ):
     key = jax.random.key(seed)
-    k0, k1, k2, k4, k5 = jax.random.split(key, 5)
+    k0, k1, k2, k3, k4, k5, k6 = jax.random.split(key, 7)
+
     a = jax.random.normal(k0, (num_tokens, hidden_size),
                           dtype=jnp.float32).astype(dtype) / 10
+
     w1 = (jax.random.normal(
         k1,
         (num_experts, 2, hidden_size, intermediate_size),
@@ -31,21 +43,54 @@ def gen_moe_inputs(
     ) / 10).astype(dtype)
     w2 = (jax.random.normal(k2, (num_experts, intermediate_size, hidden_size),
                             dtype=jnp.float32) / 10).astype(dtype)
+
+    if has_bias:
+        b1 = (jax.random.normal(k3, (num_experts, 2, intermediate_size),
+                                dtype=jnp.float32) / 10).astype(dtype)
+        b2 = (jax.random.normal(k4, (num_experts, hidden_size),
+                                dtype=jnp.float32) / 10).astype(dtype)
+    else:
+        b1 = b2 = None
+
     gating_output = (
-        jax.random.normal(k4, (num_tokens, num_experts), dtype=jnp.float32) +
+        jax.random.normal(k5, (num_tokens, num_experts), dtype=jnp.float32) +
         jnp.arange(num_tokens * num_experts, dtype=jnp.float32).reshape(
             num_tokens, num_experts) / 100)
+
     # To generate unique top-k!
-    top_k_indices = jax.random.randint(k5, (num_tokens, top_k),
+    top_k_indices = jax.random.randint(k6, (num_tokens, top_k),
                                        minval=0,
                                        maxval=num_experts - 1,
                                        dtype=jnp.int32)
+
     one_hot = (jnp.sum(
         jax.nn.one_hot(top_k_indices, num_experts, dtype=jnp.float32),
         axis=1,
-    ) * 10)
+    ) * 30)
+
     gating_output = (gating_output + one_hot).astype(dtype)
-    return a, w1, w2, gating_output
+
+    return a, w1, w2, b1, b2, gating_output
+
+
+def sub_channel_quantize(x, quant_dtype, wsz=256):
+    """Quantizes x with sub-channel quantization on the 2nd minor."""
+    if jnp.issubdtype(quant_dtype, jnp.floating):
+        dtype_info = jnp.finfo(quant_dtype)
+    else:
+        dtype_info = jnp.iinfo(quant_dtype)
+    dtype_max = float(dtype_info.max)
+    w_lst, scale_lst = [], []
+    assert len(x.shape) >= 2
+    assert x.shape[-2] % wsz == 0
+    for i in range(0, x.shape[-2], wsz):
+        y = x[..., i:i + wsz, :]
+        abs_max = jnp.abs(y).max(axis=-2, keepdims=True)
+        scale = (abs_max / dtype_max).astype(jnp.float32)
+        w = (y / scale).astype(quant_dtype)
+        w_lst.append(w)
+        scale_lst.append(scale)
+    return jnp.concat(w_lst, axis=-2), jnp.concat(scale_lst, axis=-2)
 
 
 @jtu.with_config(jax_numpy_dtype_promotion="standard")
@@ -63,42 +108,266 @@ class MoEKernelTest(jtu.JaxTestCase):
         self.mesh = Mesh(np.array(self.mesh_devices).reshape(1, -1),
                          axis_names=("data", "model"))
 
-    def test_basic(self):
-        dtype = jnp.bfloat16
-        top_k = 2
-        num_experts = 16
-        hidden_size = 256
-        intermediate_size = 256
-        num_tokens = 8 * 2
-
-        a, w1, w2, gating_output = gen_moe_inputs(
+    def _test_moe(
+        self,
+        dtype,
+        top_k,
+        num_experts,
+        hidden_size,
+        intermediate_size,
+        num_tokens,
+        seed,
+        renormalize_topk_logits,
+        bt,
+        bf,
+        bd1,
+        bd2,
+        btc,
+        bfc,
+        bd1c,
+        bd2c,
+        act_fn="silu",
+        w_dtype=None,
+        subc_quant_wsz=None,
+        has_bias=False,
+        atol=2e-1,
+        rtol=2e-1,
+    ):
+        a, w1, w2, b1, b2, gating_output = gen_moe_inputs(
             dtype,
             top_k,
             num_experts,
             hidden_size,
             intermediate_size,
             num_tokens,
+            seed=seed,
+            has_bias=has_bias,
         )
+        w1_scale = None
+        w2_scale = None
+        if w_dtype is not None:
+            if subc_quant_wsz is None:
+                subc_quant_wsz = 256
+            w1, w1_scale = sub_channel_quantize(w1, w_dtype, subc_quant_wsz)
+            w2, w2_scale = sub_channel_quantize(w2, w_dtype, subc_quant_wsz)
 
-        actual = jax.block_until_ready(
-            fused_ep_moe(
-                mesh=self.mesh,
-                tokens=a,
-                w1=w1,
-                w2=w2,
-                gating_output=gating_output,
-                top_k=top_k,
-                bt=32,
-                bf=512,
-                bd1=512,
-                bd2=512,
-                btc=32,
-                bfc=256,
-                bd1c=256,
-                bd2c=256,
-            ))
-        expected = ref_moe(a, w1, w2, gating_output, top_k)
-        self.assertAllClose(expected, actual, atol=2e-2, rtol=2e-2)
+        actual = fused_ep_moe(
+            mesh=self.mesh,
+            tokens=a,
+            w1=w1,
+            w2=w2,
+            gating_output=gating_output,
+            top_k=top_k,
+            renormalize_topk_logits=renormalize_topk_logits,
+            act_fn=act_fn,
+            subc_quant_wsz=subc_quant_wsz,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            b1=b1,
+            b2=b2,
+            bt=bt,
+            bf=bf,
+            bd1=bd1,
+            bd2=bd2,
+            btc=btc,
+            bfc=bfc,
+            bd1c=bd1c,
+            bd2c=bd2c,
+        )
+        expected = ref_moe(
+            a,
+            w1,
+            w2,
+            gating_output,
+            top_k,
+            b1=b1,
+            b2=b2,
+            renormalize_topk_logits=renormalize_topk_logits,
+            activation=act_fn,
+            subc_quant_wsz=subc_quant_wsz,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+        )
+        self.assertAllClose(actual, expected, atol=atol, rtol=rtol)
+
+    @parameterized.product(renormalize_topk_logits=[True, False], )
+    def test_basic(self, renormalize_topk_logits):
+        dtype = jnp.bfloat16
+        top_k = 8
+        num_experts = 128
+        hidden_size = 1024
+        intermediate_size = 1024
+        num_tokens = 8 * 32
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=1234,
+            renormalize_topk_logits=renormalize_topk_logits,
+            bt=32,
+            bf=1024,
+            bd1=1024,
+            bd2=1024,
+            btc=32,
+            bfc=256,
+            bd1c=256,
+            bd2c=256,
+        )
+
+    @parameterized.product(act_fn=["silu", "gelu", "swigluoai"], )
+    def test_activation(self, act_fn):
+        dtype = jnp.bfloat16
+        top_k = 8
+        num_experts = 128
+        hidden_size = 1024
+        intermediate_size = 1024
+        num_tokens = 8 * 32
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=1234,
+            renormalize_topk_logits=True,
+            act_fn=act_fn,
+            bt=32,
+            bf=512,
+            bd1=512,
+            bd2=512,
+            btc=32,
+            bfc=256,
+            bd1c=256,
+            bd2c=256,
+        )
+
+    def test_benchmark_qwen_235(self):
+        num_experts = 128
+        top_k = 8
+        hidden_size = 4096
+        intermediate_size = 1536
+        dtype = jnp.bfloat16
+        num_tokens = 8 * 64
+        seed = 54321
+        renormalize_topk_logits = True
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=seed,
+            renormalize_topk_logits=renormalize_topk_logits,
+            bt=64,
+            bf=768,
+            bd1=2048,
+            bd2=2048,
+            btc=64,
+            bfc=768,
+            bd1c=2048,
+            bd2c=2048,
+            act_fn="silu",
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+    def test_benchmark_qwen_30b_a3b(self):
+        num_experts = 128
+        top_k = 8
+        hidden_size = 2048
+        intermediate_size = 768
+        dtype = jnp.bfloat16
+        num_tokens = 512
+        seed = 54321
+        renormalize_topk_logits = True
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=seed,
+            renormalize_topk_logits=renormalize_topk_logits,
+            bt=16,
+            bf=384,
+            bd1=512,
+            bd2=512,
+            btc=16,
+            bfc=384,
+            bd1c=256,
+            bd2c=256,
+            act_fn="silu",
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+    @parameterized.product(
+        w_dtype=[jnp.int8, jnp.float8_e5m2, jnp.float4_e2m1fn], )
+    def test_sub_channel_quantization(self, w_dtype):
+        if w_dtype in (
+                jnp.float8_e5m2,
+                jnp.float4_e2m1fn,
+        ) and not jtu.is_device_tpu_at_least(version=7):
+            self.skipTest("Expect TPUv7+")
+        dtype = jnp.bfloat16
+        top_k = 8
+        num_experts = 128
+        hidden_size = 1024
+        intermediate_size = 1024
+        num_tokens = 8 * 32
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=1234,
+            renormalize_topk_logits=False,
+            w_dtype=w_dtype,
+            subc_quant_wsz=256,
+            bt=32,
+            bf=1024,
+            bd1=1024,
+            bd2=1024,
+            btc=32,
+            bfc=256,
+            bd1c=256,
+            bd2c=256,
+        )
+
+    def test_bias(self):
+        dtype = jnp.bfloat16
+        top_k = 8
+        num_experts = 128
+        hidden_size = 1024
+        intermediate_size = 1024
+        num_tokens = 8 * 32
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=1234,
+            renormalize_topk_logits=False,
+            has_bias=True,
+            bt=32,
+            bf=512,
+            bd1=512,
+            bd2=512,
+            btc=32,
+            bfc=256,
+            bd1c=256,
+            bd2c=256,
+        )
 
 
 if __name__ == "__main__":

tests/kernels/mla_v1_test.py

@@ -42,6 +42,7 @@ class MlaRaggedPagedAttentionKernelTest(jtu.JaxTestCase):
 
         padded_r_dim = align_to(r_dim, 128)
         padded_lkv_dim = align_to(lkv_dim, 128)
+        padded_kv_dim = padded_lkv_dim + padded_r_dim
         packing = get_dtype_packing(kv_dtype)
         q_lens = [s[0] for s in seq_lens]
         kv_lens_list = [s[1] for s in seq_lens]
@@ -69,13 +70,10 @@ class MlaRaggedPagedAttentionKernelTest(jtu.JaxTestCase):
         new_kv_c = gen_random((total_q_len, lkv_dim), kv_dtype)
         new_k_pe = gen_random((total_q_len, r_dim), kv_dtype)
 
-        cache_kv_c = gen_random(
-            (total_num_pages, page_size // packing, packing, padded_lkv_dim),
+        cache_kv = gen_random(
+            (total_num_pages, page_size // packing, packing, padded_kv_dim),
             kv_dtype,
         )
-        cache_k_pe = gen_random(
-            (total_num_pages, page_size // packing, packing, padded_r_dim),
-            kv_dtype)
         kv_lens = jnp.array(kv_lens_list, dtype=jnp.int32)
         page_indices = jnp.array(page_indices_list, dtype=jnp.int32)
         cu_q_lens = jnp.array(cu_q_lens_list, dtype=jnp.int32)
@@ -84,14 +82,13 @@ class MlaRaggedPagedAttentionKernelTest(jtu.JaxTestCase):
         ql_nope_for_kernel = ql_nope.copy()
         q_pe_for_kernel = q_pe.copy()
 
-        expected_out, expected_updated_kv_c, expeceted_updated_k_pe = (
+        expected_out, expected_updated_kv = (
             mla.ref_mla_ragged_paged_attention(
                 ql_nope,
                 q_pe,
                 new_kv_c,
                 new_k_pe,
-                cache_kv_c.copy(),
-                cache_k_pe.copy(),
+                cache_kv.copy(),
                 kv_lens,
                 page_indices,
                 cu_q_lens,
@@ -101,50 +98,141 @@ class MlaRaggedPagedAttentionKernelTest(jtu.JaxTestCase):
                 soft_cap=soft_cap,
             ))
 
-        kernel_out, kernel_updated_kv_c, kernel_updated_k_pe = (
-            mla.mla_ragged_paged_attention(
-                ql_nope_for_kernel,
-                q_pe_for_kernel,
-                new_kv_c,
-                new_k_pe,
-                cache_kv_c.copy(),
-                cache_k_pe.copy(),
-                kv_lens,
-                page_indices,
-                cu_q_lens,
-                distribution,
-                sm_scale=sm_scale,
-                sliding_window=sliding_window,
-                soft_cap=soft_cap,
-                num_kv_pages_per_block=num_kv_pages_per_block,
-                num_queries_per_block=num_queries_per_block,
-                vmem_limit_bytes=vmem_limit_bytes,
-            ))
+        kernel_out, kernel_updated_kv = (mla.mla_ragged_paged_attention(
+            ql_nope_for_kernel,
+            q_pe_for_kernel,
+            new_kv_c,
+            new_k_pe,
+            cache_kv.copy(),
+            kv_lens,
+            page_indices,
+            cu_q_lens,
+            distribution,
+            sm_scale=sm_scale,
+            sliding_window=sliding_window,
+            soft_cap=soft_cap,
+            num_kv_pages_per_block=num_kv_pages_per_block,
+            num_queries_per_block=num_queries_per_block,
+            vmem_limit_bytes=vmem_limit_bytes,
+        ))
 
         self.assertEqual(expected_out.shape,
                          (total_q_len, num_heads, padded_lkv_dim))
         self.assertEqual(
-            expected_updated_kv_c.shape,
-            (total_num_pages, page_size // packing, packing, padded_lkv_dim),
-        )
-        self.assertEqual(
-            expeceted_updated_k_pe.shape,
-            (total_num_pages, page_size // packing, packing, padded_r_dim),
+            expected_updated_kv.shape,
+            (total_num_pages, page_size // packing, packing, padded_kv_dim),
         )
         self.assertEqual(expected_out.dtype, kv_dtype)
-        self.assertEqual(expected_updated_kv_c.dtype, kv_dtype)
-        self.assertEqual(expeceted_updated_k_pe.dtype, kv_dtype)
+        self.assertEqual(expected_updated_kv.dtype, kv_dtype)
 
         self.assertAllClose(expected_out, kernel_out, atol=0.2, rtol=0.2)
-        self.assertAllClose(expected_updated_kv_c,
-                            kernel_updated_kv_c,
-                            atol=0.2,
-                            rtol=0.2)
-        self.assertAllClose(expeceted_updated_k_pe,
-                            kernel_updated_k_pe,
+        self.assertAllClose(expected_updated_kv,
+                            kernel_updated_kv,
                             atol=0.2,
                             rtol=0.2)
 
+    def test_update_kv_cache(self):
+        lkv_dim = 4
+        r_dim = 4
+        padded_lkv_dim = align_to(lkv_dim, 128)
+        padded_r_dim = align_to(r_dim, 128)
+        kv_dtype = jnp.bfloat16
+        new_kv_c = jnp.arange(16, dtype=kv_dtype).reshape((4, lkv_dim))
+        new_k_pe = (jnp.arange(16, dtype=kv_dtype).reshape((4, r_dim)) + 100)
+        total_num_pages = 2
+        page_size = 4
+        cache_kv_shape = mla.get_kv_cache_shape(
+            total_num_pages,
+            page_size,
+            padded_lkv_dim + padded_r_dim,
+            kv_dtype,
+        )
+        cache_kv = jnp.zeros(cache_kv_shape, dtype=kv_dtype)
+
+        # two sequences, first with 3 tokens, second with 1 token
+        kv_lens = jnp.array([3, 1], dtype=jnp.int32)
+        # first seq uses page 0, second uses page 1
+        page_indices = jnp.array([0, -1, 1, -1], dtype=jnp.int32)
+        # three tokens for first seq, one for second
+        cu_q_lens = jnp.array([0, 3, 4], dtype=jnp.int32)
+        distribution = jnp.array([0, 0, 2], dtype=jnp.int32)
+
+        # manually compute the expected cache
+        padded_new_kv_c = jnp.pad(new_kv_c,
+                                  ((0, 0), (0, padded_lkv_dim - lkv_dim)),
+                                  constant_values=0)
+        padded_new_k_pe = jnp.pad(new_k_pe,
+                                  ((0, 0), (0, padded_r_dim - r_dim)),
+                                  constant_values=0)
+
+        expected_cache = cache_kv
+        # First sequence
+        # token 0
+        page_idx, row, col = 0, 0, 0
+        expected_cache = expected_cache.at[page_idx, row,
+                                           col, :padded_lkv_dim].set(
+                                               padded_new_kv_c[0])
+        expected_cache = expected_cache.at[page_idx, row, col,
+                                           padded_lkv_dim:padded_lkv_dim +
+                                           padded_r_dim].set(
+                                               padded_new_k_pe[0])
+        # token 1
+        page_idx, row, col = 0, 0, 1
+        expected_cache = expected_cache.at[page_idx, row,
+                                           col, :padded_lkv_dim].set(
+                                               padded_new_kv_c[1])
+        expected_cache = expected_cache.at[page_idx, row, col,
+                                           padded_lkv_dim:padded_lkv_dim +
+                                           padded_r_dim].set(
+                                               padded_new_k_pe[1])
+        # token 2
+        page_idx, row, col = 0, 1, 0
+        expected_cache = expected_cache.at[page_idx, row,
+                                           col, :padded_lkv_dim].set(
+                                               padded_new_kv_c[2])
+        expected_cache = expected_cache.at[page_idx, row, col,
+                                           padded_lkv_dim:padded_lkv_dim +
+                                           padded_r_dim].set(
+                                               padded_new_k_pe[2])
+
+        # Second sequence
+        # token 0
+        page_idx, row, col = 1, 0, 0
+        expected_cache = expected_cache.at[page_idx, row,
+                                           col, :padded_lkv_dim].set(
+                                               padded_new_kv_c[3])
+        expected_cache = expected_cache.at[page_idx, row, col,
+                                           padded_lkv_dim:padded_lkv_dim +
+                                           padded_r_dim].set(
+                                               padded_new_k_pe[3])
+
+        updated_cache = mla.update_kv_cache(
+            new_kv_c,
+            new_k_pe,
+            cache_kv,
+            kv_lens,
+            page_indices,
+            cu_q_lens,
+            distribution,
+        )
+
+        self.assertAllClose(updated_cache, expected_cache)
+
+    def test_get_kv_cache_shape(self):
+        total_num_pages = 10
+        page_size = 16
+        lkv_dim = 128
+        kv_dtype = jnp.bfloat16
+        # The calculation for the expected shape is as follows:
+        # kv_packing is determined by the dtype, which is 2 for bfloat16.
+        # The second dimension is page_size / kv_packing = 16 / 2 = 8
+        # The third dimension is kv_packing = 2
+        # The fourth dimension is lkv_dim aligned to 128, which is 128
+        expected_shape = (10, 8, 2, 128)
+        self.assertEqual(
+            mla.get_kv_cache_shape(total_num_pages, page_size, lkv_dim,
+                                   kv_dtype), expected_shape)
+
     def test_ragged_paged_attention_basic(self):
         dtype = jnp.bfloat16
         seq_lens = [(192, 328), (128, 180), (64, 255)]

tests/kernels/quantized_matmul_kernel_test.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import functools
-
 import jax
 import jax.numpy as jnp
 from absl.testing import absltest, parameterized
@@ -10,6 +8,7 @@ from jax._src import test_util as jtu
 from tpu_inference.kernels.quantized_matmul import (kernel, tuned_block_sizes,
                                                     util)
 
+xla_quantized_matmul = kernel.xla_quantized_matmul
 quantized_matmul_kernel = kernel.quantized_matmul_kernel
 quantize_tensor = util.quantize_tensor
 get_tuned_block_sizes = tuned_block_sizes.get_tuned_block_sizes
@@ -17,37 +16,6 @@ get_tuned_block_sizes = tuned_block_sizes.get_tuned_block_sizes
 jax.config.parse_flags_with_absl()
 
 
-@functools.partial(jax.jit, static_argnames=["quantize_activation"])
-def reference_quantized_matmul(
-    x: jax.Array,
-    w_q: jax.Array,
-    w_scale: jax.Array,
-    quantize_activation=True,
-):
-    if quantize_activation:
-        acc_dtype = jnp.float32
-        if quantize_activation and jnp.issubdtype(w_q.dtype, jnp.integer):
-            acc_dtype = jnp.int32
-
-        x_q, x_scale = quantize_tensor(x, w_q.dtype)
-        out = jax.lax.dot_general(
-            x_q,
-            w_q,
-            dimension_numbers=(((1, ), (1, )), ((), ())),
-            preferred_element_type=acc_dtype,
-        ).astype(jnp.float32)
-        out *= x_scale
-    else:
-        out = jax.lax.dot_general(
-            x,
-            w_q,
-            dimension_numbers=(((1, ), (1, )), ((), ())),
-            preferred_element_type=jnp.float32,
-        )
-    out *= jnp.expand_dims(w_scale, 0)
-    return out.astype(x.dtype)
-
-
 @jtu.with_config(jax_numpy_dtype_promotion="standard")
 class QuantizedMatmulKernelTest(jtu.JaxTestCase):
 
@@ -94,7 +62,7 @@ class QuantizedMatmulKernelTest(jtu.JaxTestCase):
             x_q_dtype=x_q_dtype,
             tuned_value=tuned_value,
         )
-        expected = reference_quantized_matmul(
+        expected = xla_quantized_matmul(
             x, w_q, w_scale, quantize_activation=quantize_activation)
 
         self.assertAllClose(output,

tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py

@@ -176,7 +176,9 @@ class RaggedPagedAttentionHeadDim64KernelTest(jtu.JaxTestCase):
         )
         output = output[:cu_q_lens[distribution[-1]]]
 
-        dtype_bits = dtypes.bit_width(jnp.dtype(kv_dtype))
+        dtype_bits = (dtypes.bit_width(jnp.dtype(kv_dtype)) if hasattr(
+            dtypes, "bit_width") else dtypes.itemsize_bits(
+                jnp.dtype(kv_dtype)))
         tols = {
             32: 0.15,
             16: 0.2,

tests/kernels/ragged_paged_attention_kernel_v3_test.py

@@ -162,7 +162,9 @@ class RaggedPagedAttentionKernelTest(jtu.JaxTestCase):
         )
         output = output[:cu_q_lens[distribution[-1]]]
 
-        dtype_bits = dtypes.bit_width(jnp.dtype(kv_dtype))
+        dtype_bits = (dtypes.bit_width(jnp.dtype(kv_dtype)) if hasattr(
+            dtypes, "bit_width") else dtypes.itemsize_bits(
+                jnp.dtype(kv_dtype)))
         tols = {
             32: 0.15,
             16: 0.2,