diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index 28f59996a8..97c01de692 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/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,6 +506,7 @@ 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, @@ -503,10 +514,7 @@ def ragged_gather_reduce( 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.