Replacing bit_width() with itemized_bits(). (#1264)

Signed-off-by: Aman Gupta <[email protected]>

Author: aman2930

tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py

@@ -176,7 +176,9 @@ def gen_random(shape, dtype):
         )
         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 @@ def gen_random(shape, dtype):
         )
         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,

tpu_inference/kernels/collectives/all_gather_matmul.py

@@ -540,12 +540,16 @@ def get_vmem_estimate_bytes(
     """Returns the total vmem bytes used by the kernel."""
     m_per_device = m // tp_size
     n_per_device = n // tp_size
-    y_vmem_bytes = n_per_device * k * dtypes.bit_width(y_dtype) // 8
+    y_vmem_bytes = (n_per_device * k * (dtypes.bit_width(y_dtype) if hasattr(
+        dtypes, "bit_width") else dtypes.itemsize_bits(y_dtype)) // 8)
     total_bytes = (
-        2 * m_per_device * k * dtypes.bit_width(x_dtype) //
-        8  # x_vmem_scratch_ref
+        2 * m_per_device * k *
+        (dtypes.bit_width(x_dtype) if hasattr(dtypes, "bit_width") else
+         dtypes.itemsize_bits(x_dtype)) // 8  # x_vmem_scratch_ref
         + y_vmem_bytes  # y_vmem_scratch_ref
-        + 2 * m * bn * dtypes.bit_width(out_dtype) // 8  # o_vmem_scratch_ref
+        + 2 * m * bn *
+        (dtypes.bit_width(out_dtype) if hasattr(dtypes, "bit_width") else
+         dtypes.itemsize_bits(out_dtype)) // 8  # o_vmem_scratch_ref
         + acc_bytes  # acc_vmem_scratch_ref, jnp.float32
     )
     return total_bytes
@@ -639,8 +643,10 @@ def all_gather_matmul(
     # NOTE(chengjiyao): acc buffer is not used in the grid_k == 1 case.
     if grid_k == 1:
         acc_shape = (8, 128)
-    acc_bytes = acc_shape[0] * acc_shape[1] * dtypes.bit_width(
-        jnp.float32) // 8
+    acc_bytes = (
+        acc_shape[0] *
+        acc_shape[1] * (dtypes.bit_width(jnp.float32) if hasattr(
+            dtypes, "bit_width") else dtypes.itemsize_bits(jnp.float32)) // 8)
     y_vmem_shape = (n_per_device, k) if rhs_transpose else (k, n_per_device)
     estimated_vmem_bytes = get_vmem_estimate_bytes(
         m,

tpu_inference/kernels/fused_moe/v1/kernel.py

@@ -19,7 +19,8 @@ def align_to(x, a):
 
 
 def get_dtype_packing(dtype):
-    bits = dtypes.bit_width(dtype)
+    bits = (dtypes.bit_width(dtype)
+            if hasattr(dtypes, "bit_width") else dtypes.itemsize_bits(dtype))
     return 32 // bits
 
 

tpu_inference/kernels/quantized_matmul/kernel.py

@@ -96,11 +96,20 @@ def get_vmem_limit(
     """Calculate VMEM limit for the kernel."""
 
     # Calculate in/out VMEM size.
-    x_size = batch_block_size * in_block_size * dtypes.bit_width(x_dtype)
-    x_abs_max_size = batch_block_size * dtypes.bit_width(scale_dtype)
-    w_q_size = out_block_size * in_block_size * dtypes.bit_width(w_q_dtype)
-    w_scale_size = out_block_size * dtypes.bit_width(scale_dtype)
-    out_size = batch_block_size * out_block_size * dtypes.bit_width(out_dtype)
+    x_size = (batch_block_size *
+              in_block_size * (dtypes.bit_width(x_dtype) if hasattr(
+                  dtypes, "bit_width") else dtypes.itemsize_bits(x_dtype)))
+    x_abs_max_size = (
+        batch_block_size * (dtypes.bit_width(scale_dtype) if hasattr(
+            dtypes, "bit_width") else dtypes.itemsize_bits(scale_dtype)))
+    w_q_size = (out_block_size *
+                in_block_size * (dtypes.bit_width(w_q_dtype) if hasattr(
+                    dtypes, "bit_width") else dtypes.itemsize_bits(w_q_dtype)))
+    w_scale_size = (out_block_size * (dtypes.bit_width(scale_dtype) if hasattr(
+        dtypes, "bit_width") else dtypes.itemsize_bits(scale_dtype)))
+    out_size = (batch_block_size *
+                out_block_size * (dtypes.bit_width(out_dtype) if hasattr(
+                    dtypes, "bit_width") else dtypes.itemsize_bits(out_dtype)))
 
     vmem_in_out = x_size + x_abs_max_size + w_q_size + w_scale_size + out_size
     vmem_in_out *= 2  # Account for compute and vreg spills.
@@ -114,9 +123,15 @@ def get_vmem_limit(
     vmem_in_out += out_size if (n_batch > 1 or n_out > 1) else 0
 
     # Calculate scratch VMEM size.
-    acc_size = batch_block_size * out_block_size * dtypes.bit_width(acc_dtype)
-    x_q_size = batch_block_size * in_block_size * dtypes.bit_width(x_q_dtype)
-    x_scale_size = batch_block_size * dtypes.bit_width(scale_dtype)
+    acc_size = (batch_block_size *
+                out_block_size * (dtypes.bit_width(acc_dtype) if hasattr(
+                    dtypes, "bit_width") else dtypes.itemsize_bits(acc_dtype)))
+    x_q_size = (batch_block_size *
+                in_block_size * (dtypes.bit_width(x_q_dtype) if hasattr(
+                    dtypes, "bit_width") else dtypes.itemsize_bits(x_q_dtype)))
+    x_scale_size = (
+        batch_block_size * (dtypes.bit_width(scale_dtype) if hasattr(
+            dtypes, "bit_width") else dtypes.itemsize_bits(scale_dtype)))
 
     vmem_scratch = acc_size if save_acc else 0
     vmem_scratch += x_q_size + x_scale_size if save_x_q else 0

tpu_inference/kernels/ragged_paged_attention/v2/kernel.py

@@ -655,7 +655,8 @@ def cdiv(a, b):
 
 
 def get_dtype_packing(dtype):
-    bits = dtypes.bit_width(dtype)
+    bits = (dtypes.bit_width(dtype)
+            if hasattr(dtypes, "bit_width") else dtypes.itemsize_bits(dtype))
     return 32 // bits
 
 

tpu_inference/kernels/ragged_paged_attention/v2/ragged_kv_cache_update.py

@@ -200,7 +200,8 @@ def _prev_power_of_2(n: int) -> int:
 def _get_page_size_bytes(block_size: int, num_combined_kv_heads: int,
                          head_size: int, kv_cache_dtype) -> int:
     """Returns the size in bytes of one page of the KV cache."""
-    kv_cache_dtype_bit_size = dtypes.bit_width(kv_cache_dtype)
+    kv_cache_dtype_bit_size = (dtypes.bit_width(kv_cache_dtype) if hasattr(
+        dtypes, "bit_width") else dtypes.itemsize_bits(kv_cache_dtype))
     padded_head_size = _ceil_div(
         head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
 

tpu_inference/kernels/ragged_paged_attention/v3/util.py

@@ -13,7 +13,8 @@ def align_to(x, a):
 
 
 def get_dtype_bitwidth(dtype):
-    return dtypes.bit_width(dtype)
+    return (dtypes.bit_width(dtype)
+            if hasattr(dtypes, "bit_width") else dtypes.itemsize_bits(dtype))
 
 
 def get_dtype_packing(dtype):

tpu_inference/runner/kv_cache.py

@@ -131,8 +131,10 @@ def get_attention_page_size_bytes(mesh: Mesh,
         assert isinstance(kv_cache_spec, AttentionSpec)
 
         dtype = t2j_dtype(kv_cache_spec.dtype)
-        bits = dtypes.bit_width(dtype)
+        bits = (dtypes.bit_width(dtype) if hasattr(dtypes, "bit_width") else
+                dtypes.itemsize_bits(dtype))
         use_mla = isinstance(kv_cache_spec, MLAAttentionSpec)
+
         kv_cache_shape = get_kv_cache_shape_with_mesh(
             mesh=mesh,
             total_num_pages=1,  # Pass 1 to get shape of a single page.

tpu_inference/utils.py

@@ -190,7 +190,8 @@ def get_padded_num_heads(num_heads: int, sharding_size: int) -> int:
 
 
 def get_dtype_packing(dtype):
-    bits = dtypes.bit_width(dtype)
+    bits = (dtypes.bit_width(dtype)
+            if hasattr(dtypes, "bit_width") else dtypes.itemsize_bits(dtype))
     return 32 // bits