[pallas:triton] Migrated example kernels to plgpu.{load,store}

PiperOrigin-RevId: 786856202

Author: superbobry

jax/experimental/pallas/ops/gpu/attention.py

@@ -108,11 +108,9 @@ def mha_forward_kernel(
   # q tile has shape [block_q, head_dim_padded], head_dim_padded >= head_dim.
   curr_q_slice = pl.dslice(start_q * block_q, block_q)
   head_mask = (jnp.arange(head_dim_padded) < head_dim)[None, :]
-  q = pl.load(q_ref, (slice(None), slice(None)), mask=head_mask, other=0.0)
+  q = plgpu.load(q_ref, mask=head_mask, other=0.0)
   q_segment_ids = (
-      None
-      if segment_ids_ref is None
-      else pl.load(segment_ids_ref, (curr_q_slice,))
+      None if segment_ids_ref is None else segment_ids_ref[curr_q_slice]
   )
   # In FlashAttention algorithm 1 there are 2 loops: slow over tiles of kv (size
   # (Bc == block_k here), and fast over blocks of q (size Br == block_q here).
@@ -122,7 +120,7 @@ def body(start_k, carry):
     o_prev, m_prev, l_prev = carry
     curr_k_slice = pl.dslice(start_k * block_k, block_k)
 
-    k = pl.load(k_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
+    k = plgpu.load(k_ref.at[curr_k_slice, :], mask=head_mask, other=0.0)
     qk = pl.dot(q, k.T)   # [block_q, block_k]
 
     # Scale logits to convert from base-2 to the natural log domain.
@@ -140,7 +138,7 @@ def body(start_k, carry):
     if causal or segment_ids_ref is not None:
       mask = None
       if segment_ids_ref is not None:
-        kv_segment_ids = pl.load(segment_ids_ref, (curr_k_slice,))
+        kv_segment_ids = segment_ids_ref[curr_k_slice]
         mask = segment_mask(q_segment_ids, kv_segment_ids)
       if causal:
         span_q = start_q * block_q + jnp.arange(block_q)
@@ -162,7 +160,7 @@ def body(start_k, carry):
     l_curr = s_curr.sum(axis=-1)
     l_next = l_prev_corr + l_curr
     o_prev_corr = correction[:, None] * o_prev
-    v = pl.load(v_ref, (curr_k_slice, slice(None)), mask=head_mask)
+    v = plgpu.load(v_ref.at[curr_k_slice, :], mask=head_mask)
     o_curr = pl.dot(s_curr.astype(v.dtype), v)
 
     o_next = o_prev_corr + o_curr
@@ -183,8 +181,7 @@ def body(start_k, carry):
     lse_ref = residual_refs[0]
     lse_ref[...] = m_i + jnp.log2(l_i)
   # Write output to dram.
-  pl.store(o_ref, (slice(None), slice(o.shape[-1])), o.astype(o_ref.dtype),
-           mask=head_mask)
+  plgpu.store(o_ref.at[:, : o.shape[-1]], o.astype(o_ref.dtype), mask=head_mask)
 
 def segment_mask(
     q_segment_ids: jax.Array,
@@ -328,8 +325,8 @@ def _mha_forward(
 def _preprocess_backward_kernel(out_ref, dout_ref, delta_ref, head_dim: int):
   # load
   head_mask = (jnp.arange(out_ref.shape[-1]) < head_dim)[None, :]
-  o = pl.load(out_ref, (slice(None), slice(None)), mask=head_mask, other=0.0)
-  do = pl.load(dout_ref, (slice(None), slice(None)), mask=head_mask, other=0.0)
+  o = plgpu.load(out_ref, mask=head_mask, other=0.0)
+  do = plgpu.load(dout_ref, mask=head_mask, other=0.0)
   # compute
   delta = jnp.sum(o * do, axis=1)
   # write-back
@@ -402,20 +399,18 @@ def mha_backward_kernel(
   dk = jnp.zeros([block_kv_dkv, head_dim_padded], dtype=jnp.float32)
 
   head_mask = (jnp.arange(head_dim_padded) < head_dim)[None, :]
-  v = pl.load(v_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
-  k = pl.load(k_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
+  v = plgpu.load(v_ref.at[curr_k_slice, :], mask=head_mask, other=0.0)
+  k = plgpu.load(k_ref.at[curr_k_slice, :], mask=head_mask, other=0.0)
   span_k = start_k * block_kv_dkv + jnp.arange(block_kv_dkv)
   kv_segment_ids = (
-      None
-      if segment_ids_ref is None
-      else pl.load(segment_ids_ref, (curr_k_slice,))
+      None if segment_ids_ref is None else segment_ids_ref[curr_k_slice]
   )
 
   def inner_loop_dkdv(start_q, carry):
     dv, dk = carry
     curr_q_slice = pl.dslice(start_q * block_q_dkv, block_q_dkv)
 
-    q = pl.load(q_ref, (curr_q_slice, slice(None)), mask=head_mask, other=0.0)
+    q = plgpu.load(q_ref.at[curr_q_slice, :], mask=head_mask, other=0.0)
     qk = pl.dot(q, k.T)
     qk_scale = math.log2(math.e)
     if sm_scale != 1.:
@@ -425,7 +420,7 @@ def inner_loop_dkdv(start_q, carry):
     if causal or segment_ids_ref is not None:
       mask = None
       if segment_ids_ref is not None:
-        q_segment_ids = pl.load(segment_ids_ref, (curr_q_slice,))
+        q_segment_ids = segment_ids_ref[curr_q_slice]
         mask = segment_mask(q_segment_ids, kv_segment_ids)
 
       if causal:
@@ -436,10 +431,11 @@ def inner_loop_dkdv(start_q, carry):
         )
       qk = jnp.where(mask, qk, DEFAULT_MASK_VALUE)
 
-    lse = pl.load(lse_ref, (curr_q_slice,))
-    di = pl.load(delta_ref, (curr_q_slice,))
-    do = pl.load(do_scaled_ref, (curr_q_slice, slice(None)), mask=head_mask,
-                 other=0.0)
+    lse = lse_ref[curr_q_slice]
+    di = delta_ref[curr_q_slice]
+    do = plgpu.load(
+        do_scaled_ref.at[curr_q_slice, :], mask=head_mask, other=0.0
+    )
 
     p = jnp.exp2(qk - lse[:, None])
     dv = dv + pl.dot(p.astype(do.dtype).T, do)
@@ -456,10 +452,12 @@ def inner_loop_dkdv(start_q, carry):
   dv, dk = lax.fori_loop(
       lower_bound, pl.cdiv(q_seq_len, block_q_dkv), inner_loop_dkdv, (dv, dk)
   )
-  pl.store(dv_ref, (slice(None), slice(dv.shape[-1])), dv.astype(dv_ref.dtype),
-           mask=head_mask)
-  pl.store(dk_ref, (slice(None), slice(dk.shape[-1])), dk.astype(dk_ref.dtype),
-           mask=head_mask)
+  plgpu.store(
+      dv_ref.at[:, : dv.shape[-1]], dv.astype(dv_ref.dtype), mask=head_mask
+  )
+  plgpu.store(
+      dk_ref.at[:, : dk.shape[-1]], dk.astype(dk_ref.dtype), mask=head_mask
+  )
 
   # Scan #2: dQ
   #   1. Load a block of Q of size (block_q_dq, head_dim) in SMEM.
@@ -470,21 +468,18 @@ def inner_loop_dkdv(start_q, carry):
   span_q = start_q * block_q_dq + jnp.arange(block_q_dq)
   dq = jnp.zeros([block_q_dq, head_dim_padded], dtype=jnp.float32)
 
-  q = pl.load(q_ref, (curr_q_slice, slice(None)), mask=head_mask, other=0.0)
+  q = plgpu.load(q_ref.at[curr_q_slice, :], mask=head_mask, other=0.0)
   q_segment_ids = (
-      None
-      if segment_ids_ref is None
-      else pl.load(segment_ids_ref, (curr_q_slice,))
+      None if segment_ids_ref is None else segment_ids_ref[curr_q_slice]
   )
-  lse = pl.load(lse_ref, (curr_q_slice,))
-  do = pl.load(do_scaled_ref, (curr_q_slice, slice(None)), mask=head_mask,
-               other=0.0)
-  di = pl.load(delta_ref, (curr_q_slice,))
+  lse = lse_ref[curr_q_slice]
+  do = plgpu.load(do_scaled_ref.at[curr_q_slice, :], mask=head_mask, other=0.0)
+  di = delta_ref[curr_q_slice]
 
   def inner_loop_dq(start_k, dq):
     curr_k_slice = pl.dslice(start_k * block_kv_dq, block_kv_dq)
-    k = pl.load(k_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
-    v = pl.load(v_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
+    k = plgpu.load(k_ref.at[curr_k_slice, :], mask=head_mask, other=0.0)
+    v = plgpu.load(v_ref.at[curr_k_slice, :], mask=head_mask, other=0.0)
 
     qk = pl.dot(q, k.T)
     qk_scale = math.log2(math.e)
@@ -495,7 +490,7 @@ def inner_loop_dq(start_k, dq):
     if causal or segment_ids_ref is not None:
       mask = None
       if segment_ids_ref is not None:
-        kv_segment_ids = pl.load(segment_ids_ref, (curr_k_slice,))
+        kv_segment_ids = segment_ids_ref[curr_k_slice]
         mask = segment_mask(q_segment_ids, kv_segment_ids)
 
       if causal:
@@ -523,8 +518,9 @@ def inner_loop_dq(start_k, dq):
     upper_bound = pl.cdiv(kv_seq_len, block_kv_dq)
 
   dq = lax.fori_loop(0, upper_bound, inner_loop_dq, (dq))
-  pl.store(dq_ref, (slice(None), slice(dq.shape[-1])), dq.astype(dq_ref.dtype),
-           mask=head_mask)
+  plgpu.store(
+      dq_ref.at[:, : dq.shape[-1]], dq.astype(dq_ref.dtype), mask=head_mask
+  )
 
 
 def _mha_backward(sm_scale: float, causal: bool, block_sizes: BlockSizes,

jax/experimental/pallas/ops/gpu/decode_attention.py

@@ -50,7 +50,7 @@ def _compute(start_idx, kv_seq_len, 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, (q_slice, pl.ds(None)), mask=q_mask)
+    q = plgpu.load(q_ref.at[q_slice, :], mask=q_mask)
 
     def _dot(a, b):
       # if a.shape[0] == 1:
@@ -66,7 +66,7 @@ def body(start_k, carry):
       o_prev, m_prev, l_prev = carry
       curr_k_slice = pl.ds(start_k * block_k, block_k)
 
-      k = pl.load(k_ref, (curr_k_slice, slice(None)))
+      k = k_ref[curr_k_slice, :]
       qk = _dot(q, k.T)  # [block_h, block_k]
       if sm_scale != 1.0:
         qk *= sm_scale  # [block_h, block_k]
@@ -86,7 +86,7 @@ def body(start_k, carry):
       )  # Use m_next instead of m_curr to avoid a correction on l_curr
       l_curr = s_curr.sum(axis=-1)
       l_next = l_prev_corr + l_curr
-      v = pl.load(v_ref, (curr_k_slice, slice(None)))
+      v = v_ref[curr_k_slice, :]
       o_curr = _dot(s_curr.astype(v.dtype), v)
 
       # flash2 unscaled_o
@@ -106,10 +106,10 @@ def body(start_k, carry):
 
   start_idx = split_k_seq_len * prog_j
   if start_idx_ref is not None:
-    start_idx = jnp.maximum(start_idx, pl.load(start_idx_ref, ()))
+    start_idx = jnp.maximum(start_idx, start_idx_ref[()])
   kv_seq_len = (prog_j + 1) * split_k_seq_len  # lower bound on actual k_seq_len
   if kv_seq_len_ref is not None:
-    kv_seq_len = jnp.minimum(kv_seq_len, pl.load(kv_seq_len_ref, ()))
+    kv_seq_len = jnp.minimum(kv_seq_len, kv_seq_len_ref[()])
 
   if start_idx_ref is None and kv_seq_len is None:
     o, m_i, l_i = _compute(start_idx, kv_seq_len, o, m_i, l_i)
@@ -122,10 +122,10 @@ def body(start_k, carry):
   if residual_refs:
     l_ref, m_ref = residual_refs
     vec_q_mask = q_mask.reshape(-1) if q_mask is not None else None
-    pl.store(l_ref, q_slice, l_i, mask=vec_q_mask)
-    pl.store(m_ref, q_slice, m_i, mask=vec_q_mask)
+    plgpu.store(l_ref.at[q_slice], l_i, mask=vec_q_mask)
+    plgpu.store(m_ref.at[q_slice], m_i, mask=vec_q_mask)
   o = o.astype(o_ref.dtype)
-  pl.store(o_ref, (q_slice, pl.ds(None)), o, mask=q_mask)
+  plgpu.store(o_ref.at[q_slice, :], o, mask=q_mask)
 
 
 def decode_attn_unbatched(

jax/experimental/pallas/ops/gpu/layer_norm.py

@@ -35,17 +35,19 @@ def layer_norm_forward_kernel(
   def mean_body(i, acc_ref):
     col_idx = i * block_size + jnp.arange(block_size)
     mask = col_idx < n_col
-    a = pl.load(x_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_last").astype(jnp.float32)
+    a = plgpu.load(
+        x_ref.at[col_idx], mask=mask, other=0.0, eviction_policy="evict_last"
+    ).astype(jnp.float32)
     acc_ref[:] += a
   mean = for_loop(pl.cdiv(n_col, block_size), mean_body,
                   jnp.zeros(block_size)).sum() / n_col
 
   def var_body(i, acc_ref):
     col_idx = i * block_size + jnp.arange(block_size)
     mask = col_idx < n_col
-    a = pl.load(x_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_last").astype(jnp.float32)
+    a = plgpu.load(
+        x_ref.at[col_idx], mask=mask, other=0.0, eviction_policy="evict_last"
+    ).astype(jnp.float32)
     a = jnp.where(mask, a - mean, 0.)
     acc_ref[:] += a * a
   var = for_loop(pl.cdiv(n_col, block_size), var_body,
@@ -59,12 +61,13 @@ def var_body(i, acc_ref):
   def body(i, _):
     col_idx = i * block_size + jnp.arange(block_size)
     mask = col_idx < n_col
-    weight = pl.load(weight_ref, (col_idx,), mask=mask)
-    bias = pl.load(bias_ref, (col_idx,), mask=mask)
-    x = pl.load(x_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_first").astype(jnp.float32)
+    weight = plgpu.load(weight_ref.at[col_idx], mask=mask)
+    bias = plgpu.load(bias_ref.at[col_idx], mask=mask)
+    x = plgpu.load(
+        x_ref.at[col_idx], mask=mask, other=0.0, eviction_policy="evict_first"
+    ).astype(jnp.float32)
     out = (x - mean) * rstd * weight + bias
-    pl.store(o_ref, (col_idx,), out.astype(o_ref.dtype), mask=mask)
+    plgpu.store(o_ref.at[col_idx], out.astype(o_ref.dtype), mask=mask)
   for_loop(pl.cdiv(n_col, block_size), body, ())
 
 
@@ -119,12 +122,18 @@ def layer_norm_backward_kernel_dx(
   def mean_body(i, acc_ref):
     col_idx = i * block_size + jnp.arange(block_size)
     mask = col_idx < n_col
-    a = pl.load(x_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_last").astype(jnp.float32)
-    dout = pl.load(do_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_last").astype(jnp.float32)
-    weight = pl.load(weight_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_last").astype(jnp.float32)
+    a = plgpu.load(
+        x_ref.at[col_idx], mask=mask, other=0.0, eviction_policy="evict_last"
+    ).astype(jnp.float32)
+    dout = plgpu.load(
+        do_ref.at[col_idx], mask=mask, other=0.0, eviction_policy="evict_last"
+    ).astype(jnp.float32)
+    weight = plgpu.load(
+        weight_ref.at[col_idx],
+        mask=mask,
+        other=0.0,
+        eviction_policy="evict_last",
+    ).astype(jnp.float32)
     a_hat = (a - mean_ref[...]) * rstd_ref[...]
     wdout = weight * dout
     mean1_acc_ref, mean2_acc_ref = acc_ref
@@ -139,12 +148,18 @@ def mean_body(i, acc_ref):
   def dx_body(i, acc_ref):
     col_idx = i * block_size + jnp.arange(block_size)
     mask = col_idx < n_col
-    a = pl.load(x_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_last").astype(jnp.float32)
-    dout = pl.load(do_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_last").astype(jnp.float32)
-    weight = pl.load(weight_ref, (col_idx,), mask=mask, other=0.,
-                eviction_policy="evict_last").astype(jnp.float32)
+    a = plgpu.load(
+        x_ref.at[col_idx], mask=mask, other=0.0, eviction_policy="evict_last"
+    ).astype(jnp.float32)
+    dout = plgpu.load(
+        do_ref.at[col_idx], mask=mask, other=0.0, eviction_policy="evict_last"
+    ).astype(jnp.float32)
+    weight = plgpu.load(
+        weight_ref.at[col_idx],
+        mask=mask,
+        other=0.0,
+        eviction_policy="evict_last",
+    ).astype(jnp.float32)
     a_hat = (a - mean_ref[...]) * rstd_ref[...]
     wdout = weight * dout
     da = (wdout - (a_hat * mean1 + mean2)) * rstd_ref[...]
@@ -168,21 +183,25 @@ def body(i, acc_ref):
     row_idx = i * block_m + jnp.arange(block_m)
     row_mask = row_idx < m
     mask = row_mask[:, None] & col_mask[None, :]
-    a = pl.load(
-        x_ref, (row_idx[:, None], col_idx[None]), mask=mask, other=0.0
+    a = plgpu.load(
+        x_ref.at[row_idx[:, None], col_idx[None]], mask=mask, other=0.0
     ).astype(jnp.float32)
-    dout = pl.load(
-        do_ref, (row_idx[:, None], col_idx[None]), mask=mask, other=0.0
+    dout = plgpu.load(
+        do_ref.at[row_idx[:, None], col_idx[None]], mask=mask, other=0.0
     ).astype(jnp.float32)
-    mean = pl.load(mean_ref, (row_idx,), mask=row_mask, other=0.).astype(jnp.float32)
-    rstd = pl.load(rstd_ref, (row_idx,), mask=row_mask, other=0.).astype(jnp.float32)
+    mean = plgpu.load(mean_ref.at[row_idx], mask=row_mask, other=0.0).astype(
+        jnp.float32
+    )
+    rstd = plgpu.load(rstd_ref.at[row_idx], mask=row_mask, other=0.0).astype(
+        jnp.float32
+    )
     a_hat = (a - mean[:, None]) * rstd[:, None]
     dw_acc_ref, db_acc_ref = acc_ref
     dw_acc_ref[:] += (dout * a_hat).sum(axis=0)
     db_acc_ref[:] += dout.sum(axis=0)
   dw_acc, db_acc = for_loop(pl.cdiv(m, block_m), body, (jnp.zeros(block_n), jnp.zeros(block_n)))
-  pl.store(dw_ref, (col_idx,), dw_acc.astype(dw_ref.dtype), mask=col_mask)
-  pl.store(db_ref, (col_idx,), db_acc.astype(db_ref.dtype), mask=col_mask)
+  plgpu.store(dw_ref.at[col_idx], dw_acc.astype(dw_ref.dtype), mask=col_mask)
+  plgpu.store(db_ref.at[col_idx], db_acc.astype(db_ref.dtype), mask=col_mask)
 
 
 def layer_norm_backward(

jax/experimental/pallas/ops/gpu/paged_attention.py

@@ -54,7 +54,7 @@ def paged_attention_kernel(
 
   def _compute(start_page_idx, end_page_idx, o, m_i, l_i):
     q_slice = pl.ds(0, block_h)
-    q = pl.load(q_ref, (q_slice, slice(None)))
+    q = q_ref[q_slice, :]
 
     # Loop over blocks of pages to process a entire page sequence partition.
     # Grid loops over q blocks over num_heads.
@@ -64,7 +64,7 @@ def body(start_k, carry):
       block_tables_slice = pl.ds(
           start_k * pages_per_compute_block, pages_per_compute_block
       )
-      block_tables = pl.load(block_tables_ref, block_tables_slice)
+      block_tables = block_tables_ref[block_tables_slice]
       k = k_pages_ref[block_tables].reshape(block_k, head_dim)
       v = v_pages_ref[block_tables].reshape(block_k, head_dim)
       if k_scales_pages_ref is not None: