Quantize activations before EP all gather

[Shuwen-Fang]

Quantize activations before EP all-gather (ring-of-experts)

This PR quantizes EP all gather collective and activation sorting when qwix quantization is enabled. This is free because gmm is done in fp8.

Result
xprof (tbd)

Testing strategy

• Numerical equivalence (test_quantize_before_ep_all_gather_equivalence): compares loss + full gradient tree with quantize_before_ep_all_gather=True vs False to confirmsmoving quantization earlier doesn't change training math.
• AOT compile checks on real target configs

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces support for Qwix quantization (pre-quantized activations) across GMM and ragged sort/gather kernels, including integration in the MoE layer and a new unit test. While these changes enable quantized execution paths, several critical issues must be addressed: a runtime TypeError in the new unit test due to an unsupported lhs_scale argument in gmm_v2, a potential ZeroDivisionError in ragged_gather_reduce_v2.py when num_cores is 2, and an AttributeError from accessing jnp.bfloat16.dtype. Additionally, shape checks are required in ragged_sort.py before gathering scales to prevent runtime failures on scalar or per-channel scales, private imports in ops.py should be replaced with standard JAX broadcasting, and dead code in ragged_gather_reduce_v2.py should be cleaned up.

[gemini-code-assist]

The gmm_v2 function does not accept lhs_scale as a parameter. Passing lhs_scale=lhs_scale here will raise a TypeError at runtime, causing the unit test to fail. Instead, the output of gmm_v2 should be scaled by lhs_scale manually.

Suggested change

	actual_prequantized = gmm_backend.gmm_v2(
	lhs_q,
	rhs_q,
	group_sizes,
	rhs_scale=rhs_scale,
	lhs_scale=lhs_scale,
	group_offset=group_offset_arr,
	maybe_quantize_lhs=False,
	)
	actual_prequantized = gmm_backend.gmm_v2(
	lhs_q,
	rhs_q,
	group_sizes,
	rhs_scale=rhs_scale,
	group_offset=group_offset_arr,
	maybe_quantize_lhs=False,
	)
	actual_prequantized *= lhs_scale.astype(actual_prequantized.dtype)

[gemini-code-assist]

When num_cores == 2, num_column_partitions can grow larger than num_cores (e.g., to 4 or 8), which causes num_row_partitions = num_cores // num_column_partitions to evaluate to 0. A value of 0 for num_row_partitions leads to a ZeroDivisionError in _preprocess and main_kernel (e.g., valid_rows_mask.shape[0] // num_row_partitions). Please clamp num_row_partitions to at least 1 using max(1, ...).

Suggested change

	num_row_partitions = num_cores // num_column_partitions
	num_row_partitions = max(1, num_cores // num_column_partitions)

[gemini-code-assist]

In standard JAX/NumPy, jnp.bfloat16 is a type object and does not have a .dtype attribute. Accessing jnp.bfloat16.dtype will raise an AttributeError at runtime. Please use jnp.bfloat16 directly.

Suggested change

	if out_dtype is None or jnp.issubdtype(out_dtype, jnp.float8_e4m3fn):
	out_dtype = jnp.bfloat16.dtype
	if out_dtype is None or jnp.issubdtype(out_dtype, jnp.float8_e4m3fn):
	out_dtype = jnp.bfloat16

[gemini-code-assist]

If hidden_states_local is a QArray with a per-tensor (scalar) or per-channel scale, calling ragged_gather on hidden_states_local.scale will fail at runtime because the scale dimensions do not match the token dimension. Please add a check to only gather the scale if its first dimension matches the token dimension, similar to the logic used in moe.py.

      if isinstance(hidden_states_local, qpl.QArray):
        x_qval = ragged_gather(
            hidden_states_local.qvalue,
            token_indices_sorted,
            shard_output_start[None],
            shard_output_end[None],
            enforce_fallback=enforce_gather_fallback,
            flops_override=gather_flops_override,
            bytes_accessed_override=gather_bytes_accessed_override,
        )
        if hidden_states_local.scale.shape[0] == hidden_states_local.qvalue.shape[0]:
          x_scale = ragged_gather(
              hidden_states_local.scale,
              token_indices_sorted,
              shard_output_start[None],
              shard_output_end[None],
              enforce_fallback=enforce_gather_fallback,
              flops_override=gather_flops_override,
              bytes_accessed_override=gather_bytes_accessed_override,
          )
        else:
          x_scale = hidden_states_local.scale
        x = qpl.QArray(qvalue=x_qval, scale=x_scale)

[gemini-code-assist]

If hidden_states_local is a QArray with a per-tensor (scalar) or per-channel scale, calling ragged_gather on hidden_states_local.scale will fail at runtime because the scale dimensions do not match the token dimension. Please add a check to only gather the scale if its first dimension matches the token dimension, similar to the logic used in moe.py.

      if isinstance(hidden_states_local, qpl.QArray):
        x_qval = ragged_gather(
            hidden_states_local.qvalue,
            sliced_indices,
            jnp.int32(0)[None],
            gather_end[None],
            enforce_fallback=enforce_gather_fallback,
            flops_override=gather_flops_override,
            bytes_accessed_override=gather_bytes_accessed_override,
        )
        if hidden_states_local.scale.shape[0] == hidden_states_local.qvalue.shape[0]:
          x_scale = ragged_gather(
              hidden_states_local.scale,
              sliced_indices,
              jnp.int32(0)[None],
              gather_end[None],
              enforce_fallback=enforce_gather_fallback,
              flops_override=gather_flops_override,
              bytes_accessed_override=gather_bytes_accessed_override,
          )
        else:
          x_scale = hidden_states_local.scale
        x = qpl.QArray(qvalue=x_qval, scale=x_scale)

[gemini-code-assist]

Disabling the MAX_ITERATIONS check with if False and ... introduces dead code and is a code smell. If the check is no longer desired, please remove or comment it out cleanly.

Suggested change

	if False and num_iterations > _CostModelConstants.MAX_ITERATIONS:
	break
	# TODO: Determine if MAX_ITERATIONS check is still needed or can be removed.

[codecov]

Codecov Report

❌ Patch coverage is 41.17647% with 40 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	32.14%	13 Missing and 6 partials ⚠️
src/maxtext/kernels/ragged/ragged_sort.py	18.18%	8 Missing and 1 partial ⚠️
src/maxtext/kernels/megablox/ops.py	65.00%	3 Missing and 4 partials ⚠️
...ernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py	44.44%	2 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!