Porting Pallas pl.load and pl.load to modern indexing for Jax 0.8.3.dev.

Author: camiloCienet

axlearn/common/flash_attention/gpu_decoding.py

@@ -95,7 +95,7 @@ def _compute(block_kv_start_idx, block_kv_seqlen, o, m_i, l_i):
         # Load q: it will stay in L1 throughout. Indices form a matrix because we
         # read, compute, and write all in 2d chunks. 1 element ~= 1 CUDA thread index.
         # q tile has shape [block_h, head_dim].
-        q = pl.load(q_ref, (slice(None), slice(None)), mask=q_mask) * softmax_scale
+        q = jnp.where(q_mask, q_ref[...], 0.0) * softmax_scale
 
         mask_indices = jnp.arange(block_k)
 
@@ -110,13 +110,11 @@ def body(start_k, carry):
 
             def compute():
                 curr_k_slice = pl.ds(start_k * block_k, block_k)
-                k = pl.load(k_ref, (curr_k_slice, slice(None)), mask=mask[:, None], other=0.0)
+                k = jnp.where(mask[:, None], k_ref[curr_k_slice, :], 0.0)
                 k = k.astype(q.dtype)
                 qk = pl.dot(q, k.T, precision=precision)  # [block_h, block_k]
                 if bias_ref is not None:
-                    qk += pl.load(
-                        bias_ref, (slice(None), curr_k_slice), mask=mask[None, :], other=0.0
-                    )
+                    qk += jnp.where(mask[None, :], bias_ref[:, curr_k_slice], 0.0)
 
                 qk = jnp.where(logits_mask[None, :], qk, NEG_INF)
 
@@ -128,7 +126,7 @@ def compute():
                 s_curr = jnp.exp(qk - m_next[:, None])
                 l_curr = s_curr.sum(axis=-1)
                 l_next = l_prev_corr + l_curr
-                v = pl.load(v_ref, (curr_k_slice, slice(None)), mask=mask[:, None], other=0.0)
+                v = jnp.where(mask[:, None], v_ref[curr_k_slice, :], 0.0)
                 v = v.astype(q.dtype)
                 o_curr = pl.dot(s_curr.astype(v.dtype), v, precision=precision)
 
@@ -155,7 +153,8 @@ def no_compute():
     o = jnp.zeros((block_h, head_dim), dtype=jnp.float32)
 
     block_kv_start_idx = prog_j * split_k_seq_len
-    kv_seq_len = pl.load(kv_seq_len_ref, ())
+    kv_seq_len = kv_seq_len_ref[...]
+
     block_kv_seqlen = jnp.minimum((prog_j + 1) * split_k_seq_len, kv_seq_len)
 
     # Skip padding in seq dim.
@@ -167,9 +166,9 @@ def no_compute():
 
     # Write output to HBM.
     vec_q_mask = q_mask.reshape(-1)
-    pl.store(l_ref, slice(None), l_i, mask=vec_q_mask)
-    pl.store(m_ref, slice(None), m_i, mask=vec_q_mask)
-    pl.store(o_ref, (slice(None), slice(None)), o, mask=q_mask)
+    l_ref[...] = jnp.where(vec_q_mask, l_i, l_ref[...])
+    m_ref[...] = jnp.where(vec_q_mask, m_i, m_ref[...])
+    o_ref[...] = jnp.where(q_mask, o, o_ref[...])
 
 
 def _get_sm_count() -> int:

axlearn/common/flash_attention/gpu_paged_attention.py

@@ -82,13 +82,13 @@ def _compute(
             Tuple of (output, max_logits, sum_exponentials).
         """
         q_slice = pl.ds(0, block_h)
-        q = pl.load(q_ref, (q_slice, slice(None))) * softmax_scale
+        q = q_ref[q_slice,:] * softmax_scale
 
         def body(start_k, carry):
             o_prev, m_prev, l_prev = carry
 
             page_tables_slice = pl.ds(start_k * pages_per_compute_block, pages_per_compute_block)
-            page_tables = pl.load(page_tables_ref, page_tables_slice)
+            page_tables = page_tables_ref[page_tables_slice]
             k = key_ref[page_tables].reshape(block_k, head_dim)
             v = value_ref[page_tables].reshape(block_k, head_dim)
             logits = pl.dot(q, k.T, precision=precision)  # [block_h, block_k]
@@ -100,7 +100,7 @@ def body(start_k, carry):
 
             if bias_ref is not None:
                 bias_slice = pl.ds(start_k * block_k, block_k)
-                bias = pl.load(bias_ref, (slice(None), bias_slice))
+                bias = bias_ref[:,bias_slice]
                 logits += bias
 
             if mask_fn is not None:
@@ -130,7 +130,7 @@ def body(start_k, carry):
     m_i = jnp.zeros(block_h, dtype=jnp.float32) + jnp.finfo(jnp.float32).min
     l_i = jnp.zeros(block_h, dtype=jnp.float32)
     o = jnp.zeros((block_h, head_dim), dtype=jnp.float32)
-    length = pl.load(lengths_ref, ())
+    length = lengths_ref[...]
 
     start_page_idx = partition_idx * pages_per_partition
     end_page_idx = start_page_idx + pages_per_partition