[Kernel] Remove KV masking by performing full bkv fetches in the first 2 steps (#1240)

Signed-off-by: Chengji Yao <chengjiyao@google.com>
Co-authored-by: Kyuyeun Kim <kyuyeunk@google.com>

Author: yaochengji

tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py

@@ -99,7 +99,7 @@ def gen_random(shape, dtype):
                     (0, 0),
                     (0, 0),
                 ),
-                constant_values=jnp.nan,
+                constant_values=0,
             ).reshape(
                 -1,
                 page_size,
@@ -122,7 +122,7 @@ def gen_random(shape, dtype):
             kv_cache,
             ((0, num_pages - kv_cache.shape[0]), (0, 0), (0, 0), (0, 0),
              (0, 0)),
-            constant_values=jnp.nan,
+            constant_values=0,
         )
         page_indices = jnp.stack(page_indices_list, axis=0)
         page_indices = jnp.pad(

tpu_inference/kernels/ragged_paged_attention/v3/kernel_hd64.py

@@ -456,7 +456,12 @@ def _async_copy(src, dst, sem, wait):
         else:
             cp.start()
 
-    def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, wait=False):
+    def _fetch_bkv(seq_idx,
+                   bkv_idx,
+                   bkv_sem_idx,
+                   *,
+                   is_full_fetch=False,
+                   wait=False):
         sem = sems.at[0, bkv_sem_idx]
         vmem_ref = bkv_x2_ref.at[bkv_sem_idx]
 
@@ -534,10 +539,29 @@ def _fetch_bkv_from_new_kv():
                     wait,
                 )
 
+            # NOTE(chengjiyao): This condition is true for the first two bkv fetches.
+            # We need to ensure the bkv_x2_ref VMEM buffer is fully initialized to
+            # avoid potential NaN values in regions not overwritten by actual data.
+            # This is done by padding the remaining parts of the buffer with data
+            # from the KV cache. This special handling is only strictly necessary
+            # until both buffers in the double buffer (bkv_x2_ref) have been written
+            # to at least once.
+            @pl.when(is_full_fetch)
+            def _make_sure_bkv_vmem_is_not_nan():
+                effective_sz = offset + bkv_sz_frm_new
+                remaining_sz = bkv_sz - effective_sz
+                _async_copy(
+                    cache_hbm_ref.at[pl.ds(0, remaining_sz)],
+                    vmem_ref.at[pl.ds(effective_sz, remaining_sz)],
+                    sem,
+                    wait,
+                )
+
             return kv_len_start + offset, bkv_sz_frm_new
         else:
             offset = jnp.minimum(kv_left_frm_cache, page_size * bkv_p)
-            dst = vmem_ref.at[pl.ds(0, offset + bkv_sz_frm_new)]
+            sz = lax.select(is_full_fetch, bkv_sz, offset + bkv_sz_frm_new)
+            dst = vmem_ref.at[pl.ds(0, sz)]
             _async_copy(
                 src=dst,
                 dst=dst,
@@ -664,11 +688,18 @@ def _send_bo(seq_idx, bo_idx, bo_sem_idx, *, wait=False):
             wait,
         )
 
-    def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx):
-        return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx)
+    def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, is_full_fetch=False):
+        return _fetch_bkv(seq_idx,
+                          bkv_idx,
+                          bkv_sem_idx,
+                          is_full_fetch=is_full_fetch)
 
-    def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx):
-        return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, wait=True)
+    def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, is_full_fetch=False):
+        return _fetch_bkv(seq_idx,
+                          bkv_idx,
+                          bkv_sem_idx,
+                          is_full_fetch=is_full_fetch,
+                          wait=True)
 
     def start_fetch_bq(seq_idx, bq_idx, bq_sem_idx):
         return _fetch_bq(seq_idx, bq_idx, bq_sem_idx)
@@ -726,7 +757,7 @@ def strided_load(ref, start, step):
         vec = ref[start::step]
         return vec
 
-    def strided_load_bkv(bkv_sem_idx, start, step, *, bkv_mask):
+    def strided_load_bkv(bkv_sem_idx, start, step):
         assert start % kv_packing == 0
         assert step % kv_packing == 0
         start //= kv_packing
@@ -735,7 +766,6 @@ def strided_load_bkv(bkv_sem_idx, start, step, *, bkv_mask):
             bkv_sz * step, actual_head_dim_x2))
 
         kv = strided_load(kv_ref, start, step)
-        kv = lax.select(bkv_mask, kv, jnp.zeros_like(kv))
         bitwidth = 32 // kv_packing
         repack_ty = jnp.dtype(f"uint{bitwidth}")
         lst = []
@@ -809,31 +839,36 @@ def prefetch_next_bq():
             def compute_with_bkv(bkv_idx, _):
                 # Create bitmask for KV.
                 assert bkv_sz % kv_packing == 0
-                actual_bkv_sz = jnp.minimum(bkv_sz, kv_len - bkv_idx * bkv_sz)
-                bkv_shape = (bkv_sz, actual_head_dim_x2)
-                bkv_mask = lax.broadcasted_iota(jnp.int32, bkv_shape,
-                                                0) < actual_bkv_sz
 
                 # Get next bkv ids.
                 bkv_sem_idx = sem_ids_ref[1]
-                next_seq_idx, _, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids(
-                    seq_idx, bq_idx, bkv_idx, bkv_sem_idx)
+                next_seq_idx, next_bq_idx_for_kv, next_bkv_idx, next_bkv_sem_idx = (
+                    get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx))
 
                 # Prefetch next bkv
                 @pl.when(next_seq_idx < num_seqs)
                 def prefetch_next_bkv():
                     sem_ids_ref[1] = next_bkv_sem_idx
-                    start_fetch_bkv(next_seq_idx, next_bkv_idx,
-                                    next_bkv_sem_idx)
+                    start_fetch_bkv(
+                        next_seq_idx,
+                        next_bkv_idx,
+                        next_bkv_sem_idx,
+                        is_full_fetch=next_seq_idx + next_bq_idx_for_kv +
+                        next_bkv_idx < 2,
+                    )
 
                 # Wait for cur bq if not ready yet
                 @pl.when(bkv_idx == bkv_idx_start)
                 def wait_cur_bq():
                     wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx)
 
                 # Wait for cur bkv
-                offset, update_sz = wait_fetch_bkv(seq_idx, bkv_idx,
-                                                   bkv_sem_idx)
+                offset, update_sz = wait_fetch_bkv(
+                    seq_idx,
+                    bkv_idx,
+                    bkv_sem_idx,
+                    is_full_fetch=seq_idx + bq_idx + bkv_idx < 2,
+                )
 
                 # Start updating bkv to kv cache if applicable.
                 # Only needed in first bq loop.
@@ -862,7 +897,6 @@ def update_cur_bkv_to_cache():
                         bkv_sem_idx,
                         kv_head_start,
                         num_kv_heads,
-                        bkv_mask=bkv_mask,
                     )
                     assert len(bkv_lst) == kv_packing
                     for i in range(kv_packing):
@@ -946,7 +980,7 @@ def update_cur_bkv_to_cache():
     @pl.when(seq_idx == 0)
     def prologue():
         start_fetch_bq(0, 0, 0)
-        start_fetch_bkv(0, bkv_idx_start, 0)
+        start_fetch_bkv(0, bkv_idx_start, 0, is_full_fetch=True)
 
     @pl.when(seq_idx < decode_end)
     def process_decode():

tpu_inference/runner/kv_cache.py

@@ -82,7 +82,7 @@ def create_kv_caches(
                       ShardingAxisName.ATTN_HEAD))
 
     def _allocate() -> jax.Array:
-        return jnp.empty(
+        return jnp.zeros(
             shape=cache_shape,
             dtype=cache_dtype,
         )