Quantize activations before EP all gather

[Shuwen-Fang]

PiperOrigin-RevId: 959862142

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

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]

Importing from a private module qwix._src.core.qarray is a maintainability risk as these internal APIs are unstable and can change without notice. Since both out and lhs_scale are standard JAX arrays, standard JAX/NumPy broadcasting (e.g., out * lhs_scale) is fully supported and should be preferred.

Suggested change

	if lhs_scale is not None:
	out = call_with_generic_broadcast(jnp.multiply, out, lhs_scale.astype(out.dtype))
	if lhs_scale is not None:
	out = out * lhs_scale.astype(out.dtype)

[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 40.84507% with 42 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	29.03%	15 Missing and 7 partials ⚠️
src/maxtext/kernels/ragged/ragged_sort.py	27.27%	6 Missing and 2 partials ⚠️
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!