From 6a3b469fb4512aca71318634945d05893bc6bd89 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:22:18 +0000 Subject: [PATCH 01/11] Fix b/527066611: Implement bf16 output packing for ragged_gather_reduce Pallas SC kernel --- .../kernels/ragged/ragged_gather_reduce.py | 96 ++++++++++++++----- 1 file changed, 70 insertions(+), 26 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index 28f59996a8..cc5c8d95cb 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -260,15 +260,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,22 +283,74 @@ 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`). - 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) - 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)], - out_32b_hbm_ref.at[dst_row_hbm, pl.ds(col_hbm_start, num_lanes)], - send_sem, - ).start() - last_valid_src_row_vmem = src_row_vmem - last_valid_dst_row_hbm = dst_row_hbm + # If output is bf16, pack the accumulated float32 results into bf16 (packed as uint32) + if in_dtype == jnp.bfloat16: + base_packed_row = dst_indices[0] // 2 + for col_compute_offset in range(0, num_lanes, num_simd_lanes): + col_slice = pl.ds(col_vmem_start + col_compute_offset, num_simd_lanes) + + packed_out = jnp.zeros((9, 16), dtype=jnp.uint32) + for i in range(num_simd_lanes): + vmem_lane = src_row_idx_in_vmem[i] + dst_row = dst_indices[i] + + # Load accumulated float32 data + data_u32 = out_vmem_ref[vmem_lane, col_slice] + data_f32 = jax.lax.bitcast_convert_type(data_u32, jnp.float32) + + # Convert to bf16 and cast to uint32 (lower 16 bits) + data_bf16_u32 = jax.lax.bitcast_convert_type(data_f32.astype(jnp.bfloat16), jnp.uint32) + data_bf16_u32 = jnp.bitwise_and(data_bf16_u32, 0xffff) + + # Shift: even -> lower 16, odd -> upper 16 + shift = jnp.where(dst_row % 2 == 0, 0, 16) + shifted_data = jnp.bitwise_left_shift(data_bf16_u32, shift) + + # Local packed row index + local_packed_row = (dst_row // 2) - base_packed_row + + packed_out = packed_out.at[local_packed_row].or_(shifted_data) + + # Write packed rows back to the first 9 rows of out_vmem_ref + for p in range(9): + out_vmem_ref[p, col_slice] = packed_out[p] + + # Trigger DMA writes + if in_dtype == jnp.bfloat16: + last_valid_local_row = -1 + last_valid_dst_row_packed = -1 + base_packed_row = dst_indices[0] // 2 + + for i, row_valid in enumerate(row_valid_vec): + dst_row = dst_indices[i] + local_packed_row = (dst_row // 2) - base_packed_row + dst_row_packed = dst_row // 2 + + src_row_vmem = jnp.where(row_valid, local_packed_row, last_valid_local_row) + dst_row_hbm = jnp.where(row_valid, dst_row_packed, last_valid_dst_row_packed) + + pltpu.make_async_copy( + out_vmem_ref.at[src_row_vmem, pl.ds(col_vmem_start, num_lanes)], + out_32b_hbm_ref.at[dst_row_hbm, pl.ds(col_hbm_start, num_lanes)], + send_sem, + ).start() + + last_valid_local_row = jnp.where(row_valid, local_packed_row, last_valid_local_row) + last_valid_dst_row_packed = jnp.where(row_valid, dst_row_packed, last_valid_dst_row_packed) + else: + # Original unpacked DMA write loop (same as clean code) + last_valid_src_row_vmem = -1 + last_valid_dst_row_hbm = -1 + 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)], + out_32b_hbm_ref.at[dst_row_hbm, pl.ds(col_hbm_start, num_lanes)], + send_sem, + ).start() + last_valid_src_row_vmem = src_row_vmem + last_valid_dst_row_hbm = dst_row_hbm return carry @@ -505,7 +549,7 @@ def ragged_gather_reduce( **{ _OUT_KW: jax.ShapeDtypeStruct( (padded_input_size // reduce_group_size, aligned_hidden_size), - jnp.float32, + x.dtype, ), _SCRATCH_KW: dict( # pylint: disable=use-dict-literal num_rows_per_row_partition_vmem_ref=pltpu.VMEM((num_simd_lanes,), jnp.int32), From 4ed85b4d51705acae4782d43f59983b64309f309 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:30:44 +0000 Subject: [PATCH 02/11] Fix bitcast shape mismatch by doing two-stage cast (bf16 -> uint16 -> uint32) --- src/maxtext/kernels/ragged/ragged_gather_reduce.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index cc5c8d95cb..e134d122a2 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -299,7 +299,9 @@ def dma_write_loop(col_vmem_start, carry): data_f32 = jax.lax.bitcast_convert_type(data_u32, jnp.float32) # Convert to bf16 and cast to uint32 (lower 16 bits) - data_bf16_u32 = jax.lax.bitcast_convert_type(data_f32.astype(jnp.bfloat16), jnp.uint32) + data_bf16 = data_f32.astype(jnp.bfloat16) + data_u16 = jax.lax.bitcast_convert_type(data_bf16, jnp.uint16) + data_bf16_u32 = data_u16.astype(jnp.uint32) data_bf16_u32 = jnp.bitwise_and(data_bf16_u32, 0xffff) # Shift: even -> lower 16, odd -> upper 16 From c3d1d5b831497f74b619846b38a80e58183cfd44 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:32:38 +0000 Subject: [PATCH 03/11] Fix bitwise OR syntax by using .apply(lambda val: jnp.bitwise_or(val, shifted_data)) --- src/maxtext/kernels/ragged/ragged_gather_reduce.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index e134d122a2..f12e11624e 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -311,7 +311,7 @@ def dma_write_loop(col_vmem_start, carry): # Local packed row index local_packed_row = (dst_row // 2) - base_packed_row - packed_out = packed_out.at[local_packed_row].or_(shifted_data) + packed_out = packed_out.at[local_packed_row].apply(lambda val: jnp.bitwise_or(val, shifted_data)) # Write packed rows back to the first 9 rows of out_vmem_ref for p in range(9): From 10cb7a347e7ca4a5c4a27abfbb3498907bbd255f Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:34:18 +0000 Subject: [PATCH 04/11] Fix bitwise OR syntax by using .set(jnp.bitwise_or(packed_out[local_packed_row], shifted_data)) --- src/maxtext/kernels/ragged/ragged_gather_reduce.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index f12e11624e..f904526eba 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -311,7 +311,7 @@ def dma_write_loop(col_vmem_start, carry): # Local packed row index local_packed_row = (dst_row // 2) - base_packed_row - packed_out = packed_out.at[local_packed_row].apply(lambda val: jnp.bitwise_or(val, shifted_data)) + packed_out = packed_out.at[local_packed_row].set(jnp.bitwise_or(packed_out[local_packed_row], shifted_data)) # Write packed rows back to the first 9 rows of out_vmem_ref for p in range(9): From 10e2d141e946a0a1612bf49d237bc803544a5380 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:35:32 +0000 Subject: [PATCH 05/11] Fix packed_out shape by using num_simd_lanes instead of hardcoded 16 --- src/maxtext/kernels/ragged/ragged_gather_reduce.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index f904526eba..e4fd0ea046 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -289,7 +289,7 @@ def dma_write_loop(col_vmem_start, carry): for col_compute_offset in range(0, num_lanes, num_simd_lanes): col_slice = pl.ds(col_vmem_start + col_compute_offset, num_simd_lanes) - packed_out = jnp.zeros((9, 16), dtype=jnp.uint32) + packed_out = jnp.zeros((9, num_simd_lanes), dtype=jnp.uint32) for i in range(num_simd_lanes): vmem_lane = src_row_idx_in_vmem[i] dst_row = dst_indices[i] From 6412f35b7b5918c1259a936d9df34d243fd1fdce Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:37:06 +0000 Subject: [PATCH 06/11] Fix dynamic_slice Pallas SC limitation by packing directly into VMEM Ref out_vmem_ref --- .../kernels/ragged/ragged_gather_reduce.py | 43 ++++++++----------- 1 file changed, 18 insertions(+), 25 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index e4fd0ea046..a96467deff 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -286,59 +286,52 @@ def dma_write_loop(col_vmem_start, carry): # If output is bf16, pack the accumulated float32 results into bf16 (packed as uint32) if in_dtype == jnp.bfloat16: base_packed_row = dst_indices[0] // 2 + num_packed_rows = num_simd_lanes // 2 + 1 for col_compute_offset in range(0, num_lanes, num_simd_lanes): col_slice = pl.ds(col_vmem_start + col_compute_offset, num_simd_lanes) - packed_out = jnp.zeros((9, num_simd_lanes), dtype=jnp.uint32) + # 1. Read all accumulated data into a local static buffer to avoid hazards + accum_data = jnp.zeros((num_simd_lanes, num_simd_lanes), dtype=jnp.uint32) for i in range(num_simd_lanes): vmem_lane = src_row_idx_in_vmem[i] + accum_data = accum_data.at[i].set(out_vmem_ref[vmem_lane, col_slice]) + + # 2. Initialize the packed rows in out_vmem_ref to 0 (since we will OR into them) + for p in range(num_packed_rows): + out_vmem_ref[p, col_slice] = jnp.zeros((num_simd_lanes,), dtype=jnp.uint32) + + # 3. Pack and accumulate into out_vmem_ref (dynamic indexing on VMEM Ref is allowed) + for i in range(num_simd_lanes): dst_row = dst_indices[i] - # Load accumulated float32 data - data_u32 = out_vmem_ref[vmem_lane, col_slice] + data_u32 = accum_data[i] data_f32 = jax.lax.bitcast_convert_type(data_u32, jnp.float32) - # Convert to bf16 and cast to uint32 (lower 16 bits) data_bf16 = data_f32.astype(jnp.bfloat16) data_u16 = jax.lax.bitcast_convert_type(data_bf16, jnp.uint16) data_bf16_u32 = data_u16.astype(jnp.uint32) data_bf16_u32 = jnp.bitwise_and(data_bf16_u32, 0xffff) - # Shift: even -> lower 16, odd -> upper 16 shift = jnp.where(dst_row % 2 == 0, 0, 16) shifted_data = jnp.bitwise_left_shift(data_bf16_u32, shift) - # Local packed row index local_packed_row = (dst_row // 2) - base_packed_row - packed_out = packed_out.at[local_packed_row].set(jnp.bitwise_or(packed_out[local_packed_row], shifted_data)) - - # Write packed rows back to the first 9 rows of out_vmem_ref - for p in range(9): - out_vmem_ref[p, col_slice] = packed_out[p] + current_packed = out_vmem_ref[local_packed_row, col_slice] + out_vmem_ref[local_packed_row, col_slice] = jnp.bitwise_or(current_packed, shifted_data) # Trigger DMA writes if in_dtype == jnp.bfloat16: - last_valid_local_row = -1 - last_valid_dst_row_packed = -1 base_packed_row = dst_indices[0] // 2 + num_packed_rows = num_simd_lanes // 2 + 1 - for i, row_valid in enumerate(row_valid_vec): - dst_row = dst_indices[i] - local_packed_row = (dst_row // 2) - base_packed_row - dst_row_packed = dst_row // 2 - - src_row_vmem = jnp.where(row_valid, local_packed_row, last_valid_local_row) - dst_row_hbm = jnp.where(row_valid, dst_row_packed, last_valid_dst_row_packed) - + for p in range(num_packed_rows): + dst_row_hbm = base_packed_row + p pltpu.make_async_copy( - out_vmem_ref.at[src_row_vmem, pl.ds(col_vmem_start, num_lanes)], + out_vmem_ref.at[p, pl.ds(col_vmem_start, num_lanes)], out_32b_hbm_ref.at[dst_row_hbm, pl.ds(col_hbm_start, num_lanes)], send_sem, ).start() - - last_valid_local_row = jnp.where(row_valid, local_packed_row, last_valid_local_row) - last_valid_dst_row_packed = jnp.where(row_valid, dst_row_packed, last_valid_dst_row_packed) else: # Original unpacked DMA write loop (same as clean code) last_valid_src_row_vmem = -1 From 689b384f72fb15dae484185ef290629f9ad09660 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:38:17 +0000 Subject: [PATCH 07/11] Fix scatter Pallas SC limitation by using Python list of Tracers for accum_data --- src/maxtext/kernels/ragged/ragged_gather_reduce.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index a96467deff..01c5d62163 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -290,11 +290,11 @@ def dma_write_loop(col_vmem_start, carry): for col_compute_offset in range(0, num_lanes, num_simd_lanes): col_slice = pl.ds(col_vmem_start + col_compute_offset, num_simd_lanes) - # 1. Read all accumulated data into a local static buffer to avoid hazards - accum_data = jnp.zeros((num_simd_lanes, num_simd_lanes), dtype=jnp.uint32) + # 1. Read all accumulated data into a Python list of Tracers to avoid hazards and JAX scatter/dynamic_slice + accum_data = [None] * num_simd_lanes for i in range(num_simd_lanes): vmem_lane = src_row_idx_in_vmem[i] - accum_data = accum_data.at[i].set(out_vmem_ref[vmem_lane, col_slice]) + accum_data[i] = out_vmem_ref[vmem_lane, col_slice] # 2. Initialize the packed rows in out_vmem_ref to 0 (since we will OR into them) for p in range(num_packed_rows): From fa003d394e68ff040acf2d4393f8deccafdac5a9 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:40:46 +0000 Subject: [PATCH 08/11] Move bf16 bitcasting outside of Pallas kernel to bypass TPU compiler reinterpret_cast limitation --- .../kernels/ragged/ragged_gather_reduce.py | 42 +++++++++++++------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index 01c5d62163..7ce5bd3668 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 @@ -490,10 +495,12 @@ 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 - padded_input_size = _align_to( - input_size, - math.lcm(num_rows_partitions * row_tile_size, reduce_group_size), - ) + + is_bf16 = x.dtype == jnp.bfloat16 + alignment_lcm = math.lcm(num_rows_partitions * row_tile_size, reduce_group_size) + if is_bf16: + alignment_lcm = math.lcm(num_rows_partitions * row_tile_size, reduce_group_size * 2) + padded_input_size = _align_to(input_size, alignment_lcm) pad_input_size = padded_input_size - input_size x = jnp.pad( @@ -526,6 +533,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), jnp.uint32).reshape(padded_input_size // 2, aligned_hidden_size) + out_shape = (padded_input_size // reduce_group_size // 2, 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 @@ -535,6 +552,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, @@ -542,10 +560,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), - x.dtype, - ), + _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), @@ -556,7 +571,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).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. From fcf45ef12f27650ac58199e740fbb427f6200bf2 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:42:28 +0000 Subject: [PATCH 09/11] Fix JAX bitcast shape collapse rule by reshaping x to (-1, 2) before bitcasting --- src/maxtext/kernels/ragged/ragged_gather_reduce.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index 7ce5bd3668..f23e368d10 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -535,7 +535,7 @@ def ragged_gather_reduce( ) if is_bf16: - x_input = jax.lax.bitcast_convert_type(x.reshape(-1), jnp.uint32).reshape(padded_input_size // 2, aligned_hidden_size) + 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 // 2, aligned_hidden_size) out_dtype = jnp.uint32 else: From 045399824596bc07b60dd5764268aaa726749826 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:45:08 +0000 Subject: [PATCH 10/11] Implement elegant unpacked HBM writes for bf16, completely eliminating dynamic VMEM indexing and write collisions --- .../kernels/ragged/ragged_gather_reduce.py | 95 +++++-------------- 1 file changed, 25 insertions(+), 70 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index f23e368d10..b03c962c41 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -254,7 +254,12 @@ 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) + if in_dtype == jnp.bfloat16: + data_bf16 = accumulated_data.astype(jnp.bfloat16) + data_u16 = jax.lax.bitcast_convert_type(data_bf16, jnp.uint16) + data_to_write = data_u16.astype(jnp.uint32) + else: + data_to_write = jax.lax.bitcast_convert_type(accumulated_data, jnp.uint32) out_vmem_ref[row_src, col_slice] = data_to_write # We write the last row (within a row tile)'s accumulated data to @@ -288,69 +293,19 @@ def dma_write_loop(col_vmem_start, carry): src_row_idx_in_vmem.reverse() row_valid_vec.reverse() - # If output is bf16, pack the accumulated float32 results into bf16 (packed as uint32) - if in_dtype == jnp.bfloat16: - base_packed_row = dst_indices[0] // 2 - num_packed_rows = num_simd_lanes // 2 + 1 - for col_compute_offset in range(0, num_lanes, num_simd_lanes): - col_slice = pl.ds(col_vmem_start + col_compute_offset, num_simd_lanes) - - # 1. Read all accumulated data into a Python list of Tracers to avoid hazards and JAX scatter/dynamic_slice - accum_data = [None] * num_simd_lanes - for i in range(num_simd_lanes): - vmem_lane = src_row_idx_in_vmem[i] - accum_data[i] = out_vmem_ref[vmem_lane, col_slice] - - # 2. Initialize the packed rows in out_vmem_ref to 0 (since we will OR into them) - for p in range(num_packed_rows): - out_vmem_ref[p, col_slice] = jnp.zeros((num_simd_lanes,), dtype=jnp.uint32) - - # 3. Pack and accumulate into out_vmem_ref (dynamic indexing on VMEM Ref is allowed) - for i in range(num_simd_lanes): - dst_row = dst_indices[i] - - data_u32 = accum_data[i] - data_f32 = jax.lax.bitcast_convert_type(data_u32, jnp.float32) - - data_bf16 = data_f32.astype(jnp.bfloat16) - data_u16 = jax.lax.bitcast_convert_type(data_bf16, jnp.uint16) - data_bf16_u32 = data_u16.astype(jnp.uint32) - data_bf16_u32 = jnp.bitwise_and(data_bf16_u32, 0xffff) - - shift = jnp.where(dst_row % 2 == 0, 0, 16) - shifted_data = jnp.bitwise_left_shift(data_bf16_u32, shift) - - local_packed_row = (dst_row // 2) - base_packed_row - - current_packed = out_vmem_ref[local_packed_row, col_slice] - out_vmem_ref[local_packed_row, col_slice] = jnp.bitwise_or(current_packed, shifted_data) - - # Trigger DMA writes - if in_dtype == jnp.bfloat16: - base_packed_row = dst_indices[0] // 2 - num_packed_rows = num_simd_lanes // 2 + 1 - - for p in range(num_packed_rows): - dst_row_hbm = base_packed_row + p - pltpu.make_async_copy( - out_vmem_ref.at[p, pl.ds(col_vmem_start, num_lanes)], - out_32b_hbm_ref.at[dst_row_hbm, pl.ds(col_hbm_start, num_lanes)], - send_sem, - ).start() - else: - # Original unpacked DMA write loop (same as clean code) - last_valid_src_row_vmem = -1 - last_valid_dst_row_hbm = -1 - 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)], - out_32b_hbm_ref.at[dst_row_hbm, pl.ds(col_hbm_start, num_lanes)], - send_sem, - ).start() - last_valid_src_row_vmem = src_row_vmem - last_valid_dst_row_hbm = dst_row_hbm + # Trigger DMA writes (original unpacked DMA loop) + last_valid_src_row_vmem = -1 + last_valid_dst_row_hbm = -1 + 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)], + out_32b_hbm_ref.at[dst_row_hbm, pl.ds(col_hbm_start, num_lanes)], + send_sem, + ).start() + last_valid_src_row_vmem = src_row_vmem + last_valid_dst_row_hbm = dst_row_hbm return carry @@ -497,10 +452,10 @@ def ragged_gather_reduce( row_tile_size = num_simd_lanes is_bf16 = x.dtype == jnp.bfloat16 - alignment_lcm = math.lcm(num_rows_partitions * row_tile_size, reduce_group_size) - if is_bf16: - alignment_lcm = math.lcm(num_rows_partitions * row_tile_size, reduce_group_size * 2) - padded_input_size = _align_to(input_size, alignment_lcm) + padded_input_size = _align_to( + input_size, + math.lcm(num_rows_partitions * row_tile_size, reduce_group_size), + ) pad_input_size = padded_input_size - input_size x = jnp.pad( @@ -536,7 +491,7 @@ def ragged_gather_reduce( 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 // 2, aligned_hidden_size) + out_shape = (padded_input_size // reduce_group_size, aligned_hidden_size) out_dtype = jnp.uint32 else: x_input = x @@ -574,7 +529,7 @@ def ragged_gather_reduce( )(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).reshape(padded_input_size // reduce_group_size, aligned_hidden_size) + 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. From 70a333687533dd55fbafdb10c5787d926dde9029 Mon Sep 17 00:00:00 2001 From: gobbleturk Date: Tue, 23 Jun 2026 18:46:43 +0000 Subject: [PATCH 11/11] Optimize bf16 register conversion using 32-bit bitcast and right-shift, avoiding unsupported 16-bit register bitcasts --- src/maxtext/kernels/ragged/ragged_gather_reduce.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index b03c962c41..97c01de692 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -254,12 +254,11 @@ def dma_write_loop(col_vmem_start, carry): data, ) previous_accumulated_data = accumulated_data + data_u32 = jax.lax.bitcast_convert_type(accumulated_data, jnp.uint32) if in_dtype == jnp.bfloat16: - data_bf16 = accumulated_data.astype(jnp.bfloat16) - data_u16 = jax.lax.bitcast_convert_type(data_bf16, jnp.uint16) - data_to_write = data_u16.astype(jnp.uint32) + data_to_write = jnp.bitwise_right_shift(data_u32, 16) else: - data_to_write = jax.lax.bitcast_convert_type(accumulated_data, jnp.uint32) + 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