Ragged gather reduce output in bf16 (previously f32)

src/maxtext/kernels/ragged/ragged_gather_reduce.py

@@ -85,6 +85,7 @@ def main_kernel(
     subcore_axis_name: str,
     num_row_partitions: int,
     num_column_partitions: int,
+    is_bf16: bool,
 ):
   """Main Pallas kernel for ragged gather and reduction on SparseCore."""
   tpu_info = pltpu.get_tpu_info()
@@ -168,12 +169,16 @@ def row_loop(row_block_id):
       dst_indices = dst_indices_vmem_ref[...]
       topk_weights = topk_weights_vmem_ref[...]
 
-      in_dtype = in_hbm_ref.dtype
+      in_dtype = jnp.bfloat16 if is_bf16 else jnp.float32
       input_dtype_bits = jax.dtypes.itemsize_bits(in_dtype)
       input_packing = 32 // input_dtype_bits
 
-      in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32)
-      out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32)
+      if is_bf16:
+        in_32b_hbm_ref = in_hbm_ref
+        out_32b_hbm_ref = out_hbm_ref
+      else:
+        in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32)
+        out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32)
 
       for col_vmem_start in range(0, col_size, num_lanes):
         col_hbm_start = col_start + col_vmem_start
@@ -249,7 +254,11 @@ def dma_write_loop(col_vmem_start, carry):
                 data,
             )
             previous_accumulated_data = accumulated_data
-            data_to_write = jax.lax.bitcast_convert_type(accumulated_data, jnp.uint32)
+            data_u32 = jax.lax.bitcast_convert_type(accumulated_data, jnp.uint32)
+            if in_dtype == jnp.bfloat16:
+              data_to_write = jnp.bitwise_right_shift(data_u32, 16)
+            else:
+              data_to_write = data_u32
             out_vmem_ref[row_src, col_slice] = data_to_write
 
             # We write the last row (within a row tile)'s accumulated data to
@@ -260,15 +269,7 @@ def dma_write_loop(col_vmem_start, carry):
             if row_src == num_simd_lanes - 1:
               prev_iter_last_row_vmem_ref[0, col_slice] = data_to_write
 
-        # Start dma write.
-        # When there are multiple sources rows in the current row_tile that
-        # contribute to the same destination row, the accumulated data is
-        # stored in the last row's idx in `out_vmem_ref`.
-        # Logically, we could skip all the source rows that are not the last
-        # for each destination row, but we want to avoid using `pl.when` for
-        # efficiency. We just repeat the write of latest accumulated data
-        # multiple times.
-        # `src_row_idx_in_vmem` tracks the right idx in vmem for each hbm write.
+        # Calculate src_row_idx_in_vmem and row_valid_vec
         src_row_idx_in_vmem = []
         row_valid_vec = []
         for row_vmem_idx in reversed(range(num_simd_lanes)):
@@ -291,14 +292,11 @@ def dma_write_loop(col_vmem_start, carry):
         src_row_idx_in_vmem.reverse()
         row_valid_vec.reverse()
 
-        # There must be at least one valid row to write in the current row_tile.
-        # When num valid writes is not a multiple of row_tile_size, we repeat
-        # the last valid write to avoid valid data in hbm being overwritten
-        # (and to avoid using `pl.when`).
+        # Trigger DMA writes (original unpacked DMA loop)
         last_valid_src_row_vmem = -1
         last_valid_dst_row_hbm = -1
-        for i, (src_row_idx_in_vmem, row_valid) in enumerate(zip(src_row_idx_in_vmem, row_valid_vec, strict=True)):
-          src_row_vmem = jnp.where(row_valid, src_row_idx_in_vmem, last_valid_src_row_vmem)
+        for i, (src_row_idx, row_valid) in enumerate(zip(src_row_idx_in_vmem, row_valid_vec, strict=True)):
+          src_row_vmem = jnp.where(row_valid, src_row_idx, last_valid_src_row_vmem)
           dst_row_hbm = jnp.where(row_valid, dst_indices[i], last_valid_dst_row_hbm)
           pltpu.make_async_copy(
               out_vmem_ref.at[src_row_vmem, pl.ds(col_vmem_start, num_lanes)],
@@ -451,6 +449,8 @@ def ragged_gather_reduce(
   aligned_hidden_size = _align_to(hidden_size, 128 * num_column_partitions)
   col_size = aligned_hidden_size // num_column_partitions
   row_tile_size = num_simd_lanes
+
+  is_bf16 = x.dtype == jnp.bfloat16
   padded_input_size = _align_to(
       input_size,
       math.lcm(num_rows_partitions * row_tile_size, reduce_group_size),
@@ -487,6 +487,16 @@ def ragged_gather_reduce(
       core_axis_name="core",
       subcore_axis_name="subcore",
   )
+
+  if is_bf16:
+    x_input = jax.lax.bitcast_convert_type(x.reshape(-1, 2), jnp.uint32).reshape(padded_input_size // 2, aligned_hidden_size)
+    out_shape = (padded_input_size // reduce_group_size, aligned_hidden_size)
+    out_dtype = jnp.uint32
+  else:
+    x_input = x
+    out_shape = (padded_input_size // reduce_group_size, aligned_hidden_size)
+    out_dtype = x.dtype
+
   # Each output row from `main_kernel` will be of type float32, and then casted
   # to the input dtype when doing the filter operation.
   out = pl.kernel(  # pytype: disable=wrong-keyword-args
@@ -496,17 +506,15 @@ def ragged_gather_reduce(
           subcore_axis_name=vector_mesh.subcore_axis_name,
           num_row_partitions=num_rows_partitions,
           num_column_partitions=num_column_partitions,
+          is_bf16=is_bf16,
       ),
       compiler_params=pltpu.CompilerParams(  # pytype: disable=wrong-keyword-args
           **_COMPILER_PARAMS,
       ),
       mesh=vector_mesh,
       name="sc_ragged_gather_reduce",
       **{
-          _OUT_KW: jax.ShapeDtypeStruct(
-              (padded_input_size // reduce_group_size, aligned_hidden_size),
-              jnp.float32,
-          ),
+          _OUT_KW: jax.ShapeDtypeStruct(out_shape, out_dtype),
           _SCRATCH_KW: dict(  # pylint: disable=use-dict-literal
               num_rows_per_row_partition_vmem_ref=pltpu.VMEM((num_simd_lanes,), jnp.int32),
               out_vmem_ref=pltpu.VMEM((num_simd_lanes, col_size), jnp.uint32),
@@ -517,7 +525,10 @@ def ragged_gather_reduce(
               sem_ref=pltpu.SemaphoreType.DMA((2,)),
           ),
       },
-  )(num_src_rows_per_row_partition, x, src_indices, dst_indices, topk_weights)
+  )(num_src_rows_per_row_partition, x_input, src_indices, dst_indices, topk_weights)
+
+  if is_bf16:
+    out = jax.lax.bitcast_convert_type(out.reshape(-1), jnp.bfloat16)[..., 0].reshape(padded_input_size // reduce_group_size, aligned_hidden_size)
 
   # If there is no valid source row in a reduce group, set that group's output
   # to zero.