FP8 AllGather in FP8 GroupedGEMM + Fix Stream Usage Issue.

qa/L1_jax_distributed_unittest/test.sh

@@ -9,3 +9,4 @@ set -xe
 mkdir -p "$XML_LOG_DIR"
 
 NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+SCRIPT_NAME=test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh

tests/jax/multi_process_launch.sh

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+#!/bin/bash
+
+SCRIPT_NAME="${SCRIPT_NAME:-test.py}"
+
+
+XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
+ --xla_gpu_enable_command_buffer=''"
+
+export XLA_FLAGS="${XLA_BASE_FLAGS}"
+
+NUM_RUNS=$(nvidia-smi --query-gpu=count --format=csv,noheader)
+for ((i=1; i<NUM_RUNS; i++))
+do
+ CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_PROC > /dev/null 2>&1 &
+done
+
+CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_PROC
+
+wait

tests/jax/test_multi_process_distributed_grouped_gemm.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from transformer_engine.jax.dense import grouped_dense as te_grouped_dense
+from transformer_engine.jax.quantize import (
+ QuantizerFactory,
+ ScalingMode,
+)
+
+from utils import assert_allclose
+
+
+N_GROUP = 8
+MESH_AXIS_NAME = "fsdp"
+
+
+def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis):
+ assert kernel_fsdp_axis in [1, 2]
+ x_shape, w_shape = data_shapes
+
+ x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None))
+ w_sharding = (
+ NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME))
+ if kernel_fsdp_axis == 2
+ else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None))
+ )
+ w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None))
+
+ def init_data():
+ x_key = jax.random.PRNGKey(0)
+ w_key = jax.random.PRNGKey(1)
+ x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16)
+ w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16)
+ w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim))
+ return x, w, w, w_amax
+
+ def test_func(outter_x, outter_w, outter_w_amax):
+ in_specs = (x_sharding.spec, w_sharding.spec, None)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w, w_amax):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+
+ output = te_grouped_dense(
+ x_reshaped,
+ w,
+ n_groups,
+ kernel_amax=w_amax,
+ quantizer_set=quantizer_set,
+ kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis),
+ )
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w, w_amax):
+ output = sharded_group_gemm(x, w, w_amax)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax)
+ dx, dw, _ = vjp_fn(output)
+ return output, dx, dw
+
+ def ref_func(outter_x, outter_w):
+
+ in_specs = (x_sharding.spec, w_no_sharding.spec)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+ output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set)
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w):
+ output = sharded_group_gemm(x, w)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w)
+ dx, dw = vjp_fn(output)
+ return output, dx, dw
+
+ init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None))
+ x, w, w_global, w_amax = init_func()
+
+ o_sharding = x_sharding
+ test_func_jitted = jax.jit(
+ test_func,
+ in_shardings=(x_sharding, w_sharding, None),
+ out_shardings=(o_sharding, x_sharding, w_sharding),
+ )
+ ref_func_jitted = jax.jit(
+ ref_func,
+ in_shardings=(x_sharding, w_no_sharding),
+ out_shardings=(o_sharding, x_sharding, w_no_sharding),
+ )
+
+ out, dx, dw = test_func_jitted(x, w, w_amax)
+ ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global)
+
+ assert_allclose(out, ref_out, dtype=jnp.float8_e4m3fn)
+ assert_allclose(dx, ref_dx, dtype=jnp.float8_e5m2)
+ assert_allclose(dw, ref_dw, dtype=jnp.float8_e5m2)
+
+
+if __name__ == "__main__":
+ from jax.sharding import NamedSharding, PartitionSpec
+ from jax.experimental import shard_map
+ import sys
+
+ coord_addr = sys.argv[1]
+ proc_id = int(sys.argv[2])
+ num_procs = int(sys.argv[3])
+
+ jax.distributed.initialize(
+ coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id
+ )
+
+ mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,))
+
+ with mesh:
+ data_shapes = [((4, 16, 128, 7168), (7168, 2048))]
+ for data_shape in data_shapes:
+ for kernel_fsdp_axis in [1, 2]:
+ test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis)

transformer_engine/jax/cpp_extensions/quantization.py

@@ -931,6 +931,7 @@ def grouped_quantize(
  x: jnp.ndarray,
  quantizer: GroupedQuantizer,
  group_sizes: jnp.ndarray = None,
+ amax: jnp.ndarray = None,
  flatten_axis: int = -1,
 ) -> GroupedScaledTensor1x:
  """Quantize a tensor in grouped manner.
@@ -943,6 +944,7 @@ def grouped_quantize(
  x: Input tensor to quantize
  quantizer: The quantizer to use for quantization
  group_sizes: Array of ints containing the size of each group (default: None)
+ amax: The amax of x; if None, it is auto-generated. (default: None)
  flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1)
 
  Returns:
@@ -985,7 +987,10 @@ def grouped_quantize(
  scale = scale.at[i].set(quantizer_i.scale[0])
 
  if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
- row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
+ if amax is not None:
+ row_amax = amax
+ else:
+ row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
  segment_ids = jnp.repeat(
  jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis]
  )

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -285,18 +285,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type
  size_t out_dtype_bytes = te_dtype_bytes(out_dtype);
 
  if (is_tensor_scaling) {
- cudaStream_t stream_0 = nvte_get_compute_stream(0);
  size_t dpitch = tensor_scaling_sinv_aligment;
  size_t spitch = lhs_sinv_dtype_bytes;
  size_t width = lhs_sinv_dtype_bytes;
  size_t height = lhs_sinv_size;
  cudaMemcpy2DAsync(lhs_scatter_aligned_ptr, dpitch, lhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  spitch = rhs_sinv_dtype_bytes;
  width = rhs_sinv_dtype_bytes;
  height = rhs_sinv_size;
  cudaMemcpy2DAsync(rhs_scatter_aligned_ptr, dpitch, rhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  lhs_sinv_ptr = lhs_scatter_aligned_ptr;
  rhs_sinv_ptr = rhs_scatter_aligned_ptr;
  }

qa/L1_jax_distributed_unittest/test.sh

@@ -9,3 +9,4 @@ set -xe
 mkdir -p "$XML_LOG_DIR"
 
 NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+SCRIPT_NAME=test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh

tests/jax/multi_process_launch.sh

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+#!/bin/bash
+
+SCRIPT_NAME="${SCRIPT_NAME:-test.py}"
+
+
+XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
+ --xla_gpu_enable_command_buffer=''"
+
+export XLA_FLAGS="${XLA_BASE_FLAGS}"
+
+NUM_RUNS=$(nvidia-smi --query-gpu=count --format=csv,noheader)
+for ((i=1; i<NUM_RUNS; i++))
+do
+ CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_PROC > /dev/null 2>&1 &
+done
+
+CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_PROC
+
+wait

tests/jax/test_multi_process_distributed_grouped_gemm.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from transformer_engine.jax.dense import grouped_dense as te_grouped_dense
+from transformer_engine.jax.quantize import (
+ QuantizerFactory,
+ ScalingMode,
+)
+
+from utils import assert_allclose
+
+
+N_GROUP = 8
+MESH_AXIS_NAME = "fsdp"
+
+
+def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis):
+ assert kernel_fsdp_axis in [1, 2]
+ x_shape, w_shape = data_shapes
+
+ x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None))
+ w_sharding = (
+ NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME))
+ if kernel_fsdp_axis == 2
+ else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None))
+ )
+ w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None))
+
+ def init_data():
+ x_key = jax.random.PRNGKey(0)
+ w_key = jax.random.PRNGKey(1)
+ x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16)
+ w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16)
+ w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim))
+ return x, w, w, w_amax
+
+ def test_func(outter_x, outter_w, outter_w_amax):
+ in_specs = (x_sharding.spec, w_sharding.spec, None)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w, w_amax):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+
+ output = te_grouped_dense(
+ x_reshaped,
+ w,
+ n_groups,
+ kernel_amax=w_amax,
+ quantizer_set=quantizer_set,
+ kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis),
+ )
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w, w_amax):
+ output = sharded_group_gemm(x, w, w_amax)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax)
+ dx, dw, _ = vjp_fn(output)
+ return output, dx, dw
+
+ def ref_func(outter_x, outter_w):
+
+ in_specs = (x_sharding.spec, w_no_sharding.spec)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+ output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set)
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w):
+ output = sharded_group_gemm(x, w)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w)
+ dx, dw = vjp_fn(output)
+ return output, dx, dw
+
+ init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None))
+ x, w, w_global, w_amax = init_func()
+
+ o_sharding = x_sharding
+ test_func_jitted = jax.jit(
+ test_func,
+ in_shardings=(x_sharding, w_sharding, None),
+ out_shardings=(o_sharding, x_sharding, w_sharding),
+ )
+ ref_func_jitted = jax.jit(
+ ref_func,
+ in_shardings=(x_sharding, w_no_sharding),
+ out_shardings=(o_sharding, x_sharding, w_no_sharding),
+ )
+
+ out, dx, dw = test_func_jitted(x, w, w_amax)
+ ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global)
+
+ assert_allclose(out, ref_out, dtype=jnp.float8_e4m3fn)
+ assert_allclose(dx, ref_dx, dtype=jnp.float8_e5m2)
+ assert_allclose(dw, ref_dw, dtype=jnp.float8_e5m2)
+
+
+if __name__ == "__main__":
+ from jax.sharding import NamedSharding, PartitionSpec
+ from jax.experimental import shard_map
+ import sys
+
+ coord_addr = sys.argv[1]
+ proc_id = int(sys.argv[2])
+ num_procs = int(sys.argv[3])
+
+ jax.distributed.initialize(
+ coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id
+ )
+
+ mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,))
+
+ with mesh:
+ data_shapes = [((4, 16, 128, 7168), (7168, 2048))]
+ for data_shape in data_shapes:
+ for kernel_fsdp_axis in [1, 2]:
+ test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis)

transformer_engine/jax/cpp_extensions/quantization.py

@@ -931,6 +931,7 @@ def grouped_quantize(
  x: jnp.ndarray,
  quantizer: GroupedQuantizer,
  group_sizes: jnp.ndarray = None,
+ amax: jnp.ndarray = None,
  flatten_axis: int = -1,
 ) -> GroupedScaledTensor1x:
  """Quantize a tensor in grouped manner.
@@ -943,6 +944,7 @@ def grouped_quantize(
  x: Input tensor to quantize
  quantizer: The quantizer to use for quantization
  group_sizes: Array of ints containing the size of each group (default: None)
+ amax: The amax of x; if None, it is auto-generated. (default: None)
  flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1)
 
  Returns:
@@ -985,7 +987,10 @@ def grouped_quantize(
  scale = scale.at[i].set(quantizer_i.scale[0])
 
  if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
- row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
+ if amax is not None:
+ row_amax = amax
+ else:
+ row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
  segment_ids = jnp.repeat(
  jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis]
  )

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -285,18 +285,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type
  size_t out_dtype_bytes = te_dtype_bytes(out_dtype);
 
  if (is_tensor_scaling) {
- cudaStream_t stream_0 = nvte_get_compute_stream(0);
  size_t dpitch = tensor_scaling_sinv_aligment;
  size_t spitch = lhs_sinv_dtype_bytes;
  size_t width = lhs_sinv_dtype_bytes;
  size_t height = lhs_sinv_size;
  cudaMemcpy2DAsync(lhs_scatter_aligned_ptr, dpitch, lhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  spitch = rhs_sinv_dtype_bytes;
  width = rhs_sinv_dtype_bytes;
  height = rhs_sinv_size;
  cudaMemcpy2DAsync(rhs_scatter_aligned_ptr, dpitch, rhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  lhs_sinv_ptr = lhs_scatter_aligned_ptr;
  rhs_sinv_ptr = rhs_scatter_aligned_ptr;
  }

qa/L1_jax_distributed_unittest/test.sh

@@ -9,3 +9,4 @@ set -xe
 mkdir -p "$XML_LOG_DIR"
 
 NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+SCRIPT_NAME=test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh

tests/jax/multi_process_launch.sh

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+#!/bin/bash
+
+SCRIPT_NAME="${SCRIPT_NAME:-test.py}"
+
+
+XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
+ --xla_gpu_enable_command_buffer=''"
+
+export XLA_FLAGS="${XLA_BASE_FLAGS}"
+
+NUM_RUNS=$(nvidia-smi --query-gpu=count --format=csv,noheader)
+for ((i=1; i<NUM_RUNS; i++))
+do
+ CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_PROC > /dev/null 2>&1 &
+done
+
+CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_PROC
+
+wait

tests/jax/test_multi_process_distributed_grouped_gemm.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from transformer_engine.jax.dense import grouped_dense as te_grouped_dense
+from transformer_engine.jax.quantize import (
+ QuantizerFactory,
+ ScalingMode,
+)
+
+from utils import assert_allclose
+
+
+N_GROUP = 8
+MESH_AXIS_NAME = "fsdp"
+
+
+def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis):
+ assert kernel_fsdp_axis in [1, 2]
+ x_shape, w_shape = data_shapes
+
+ x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None))
+ w_sharding = (
+ NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME))
+ if kernel_fsdp_axis == 2
+ else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None))
+ )
+ w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None))
+
+ def init_data():
+ x_key = jax.random.PRNGKey(0)
+ w_key = jax.random.PRNGKey(1)
+ x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16)
+ w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16)
+ w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim))
+ return x, w, w, w_amax
+
+ def test_func(outter_x, outter_w, outter_w_amax):
+ in_specs = (x_sharding.spec, w_sharding.spec, None)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w, w_amax):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+
+ output = te_grouped_dense(
+ x_reshaped,
+ w,
+ n_groups,
+ kernel_amax=w_amax,
+ quantizer_set=quantizer_set,
+ kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis),
+ )
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w, w_amax):
+ output = sharded_group_gemm(x, w, w_amax)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax)
+ dx, dw, _ = vjp_fn(output)
+ return output, dx, dw
+
+ def ref_func(outter_x, outter_w):
+
+ in_specs = (x_sharding.spec, w_no_sharding.spec)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+ output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set)
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w):
+ output = sharded_group_gemm(x, w)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w)
+ dx, dw = vjp_fn(output)
+ return output, dx, dw
+
+ init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None))
+ x, w, w_global, w_amax = init_func()
+
+ o_sharding = x_sharding
+ test_func_jitted = jax.jit(
+ test_func,
+ in_shardings=(x_sharding, w_sharding, None),
+ out_shardings=(o_sharding, x_sharding, w_sharding),
+ )
+ ref_func_jitted = jax.jit(
+ ref_func,
+ in_shardings=(x_sharding, w_no_sharding),
+ out_shardings=(o_sharding, x_sharding, w_no_sharding),
+ )
+
+ out, dx, dw = test_func_jitted(x, w, w_amax)
+ ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global)
+
+ assert_allclose(out, ref_out, dtype=jnp.float8_e4m3fn)
+ assert_allclose(dx, ref_dx, dtype=jnp.float8_e5m2)
+ assert_allclose(dw, ref_dw, dtype=jnp.float8_e5m2)
+
+
+if __name__ == "__main__":
+ from jax.sharding import NamedSharding, PartitionSpec
+ from jax.experimental import shard_map
+ import sys
+
+ coord_addr = sys.argv[1]
+ proc_id = int(sys.argv[2])
+ num_procs = int(sys.argv[3])
+
+ jax.distributed.initialize(
+ coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id
+ )
+
+ mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,))
+
+ with mesh:
+ data_shapes = [((4, 16, 128, 7168), (7168, 2048))]
+ for data_shape in data_shapes:
+ for kernel_fsdp_axis in [1, 2]:
+ test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis)

transformer_engine/jax/cpp_extensions/quantization.py

@@ -931,6 +931,7 @@ def grouped_quantize(
  x: jnp.ndarray,
  quantizer: GroupedQuantizer,
  group_sizes: jnp.ndarray = None,
+ amax: jnp.ndarray = None,
  flatten_axis: int = -1,
 ) -> GroupedScaledTensor1x:
  """Quantize a tensor in grouped manner.
@@ -943,6 +944,7 @@ def grouped_quantize(
  x: Input tensor to quantize
  quantizer: The quantizer to use for quantization
  group_sizes: Array of ints containing the size of each group (default: None)
+ amax: The amax of x; if None, it is auto-generated. (default: None)
  flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1)
 
  Returns:
@@ -985,7 +987,10 @@ def grouped_quantize(
  scale = scale.at[i].set(quantizer_i.scale[0])
 
  if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
- row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
+ if amax is not None:
+ row_amax = amax
+ else:
+ row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
  segment_ids = jnp.repeat(
  jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis]
  )

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -285,18 +285,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type
  size_t out_dtype_bytes = te_dtype_bytes(out_dtype);
 
  if (is_tensor_scaling) {
- cudaStream_t stream_0 = nvte_get_compute_stream(0);
  size_t dpitch = tensor_scaling_sinv_aligment;
  size_t spitch = lhs_sinv_dtype_bytes;
  size_t width = lhs_sinv_dtype_bytes;
  size_t height = lhs_sinv_size;
  cudaMemcpy2DAsync(lhs_scatter_aligned_ptr, dpitch, lhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  spitch = rhs_sinv_dtype_bytes;
  width = rhs_sinv_dtype_bytes;
  height = rhs_sinv_size;
  cudaMemcpy2DAsync(rhs_scatter_aligned_ptr, dpitch, rhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  lhs_sinv_ptr = lhs_scatter_aligned_ptr;
  rhs_sinv_ptr = rhs_scatter_aligned_ptr;
  }

qa/L1_jax_distributed_unittest/test.sh

@@ -9,3 +9,4 @@ set -xe
 mkdir -p "$XML_LOG_DIR"
 
 NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+SCRIPT_NAME=test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh

tests/jax/multi_process_launch.sh

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+#!/bin/bash
+
+SCRIPT_NAME="${SCRIPT_NAME:-test.py}"
+
+
+XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
+ --xla_gpu_enable_command_buffer=''"
+
+export XLA_FLAGS="${XLA_BASE_FLAGS}"
+
+NUM_RUNS=$(nvidia-smi --query-gpu=count --format=csv,noheader)
+for ((i=1; i<NUM_RUNS; i++))
+do
+ CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_PROC > /dev/null 2>&1 &
+done
+
+CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_PROC
+
+wait

tests/jax/test_multi_process_distributed_grouped_gemm.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from transformer_engine.jax.dense import grouped_dense as te_grouped_dense
+from transformer_engine.jax.quantize import (
+ QuantizerFactory,
+ ScalingMode,
+)
+
+from utils import assert_allclose
+
+
+N_GROUP = 8
+MESH_AXIS_NAME = "fsdp"
+
+
+def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis):
+ assert kernel_fsdp_axis in [1, 2]
+ x_shape, w_shape = data_shapes
+
+ x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None))
+ w_sharding = (
+ NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME))
+ if kernel_fsdp_axis == 2
+ else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None))
+ )
+ w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None))
+
+ def init_data():
+ x_key = jax.random.PRNGKey(0)
+ w_key = jax.random.PRNGKey(1)
+ x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16)
+ w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16)
+ w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim))
+ return x, w, w, w_amax
+
+ def test_func(outter_x, outter_w, outter_w_amax):
+ in_specs = (x_sharding.spec, w_sharding.spec, None)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w, w_amax):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+
+ output = te_grouped_dense(
+ x_reshaped,
+ w,
+ n_groups,
+ kernel_amax=w_amax,
+ quantizer_set=quantizer_set,
+ kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis),
+ )
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w, w_amax):
+ output = sharded_group_gemm(x, w, w_amax)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax)
+ dx, dw, _ = vjp_fn(output)
+ return output, dx, dw
+
+ def ref_func(outter_x, outter_w):
+
+ in_specs = (x_sharding.spec, w_no_sharding.spec)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+ output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set)
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w):
+ output = sharded_group_gemm(x, w)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w)
+ dx, dw = vjp_fn(output)
+ return output, dx, dw
+
+ init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None))
+ x, w, w_global, w_amax = init_func()
+
+ o_sharding = x_sharding
+ test_func_jitted = jax.jit(
+ test_func,
+ in_shardings=(x_sharding, w_sharding, None),
+ out_shardings=(o_sharding, x_sharding, w_sharding),
+ )
+ ref_func_jitted = jax.jit(
+ ref_func,
+ in_shardings=(x_sharding, w_no_sharding),
+ out_shardings=(o_sharding, x_sharding, w_no_sharding),
+ )
+
+ out, dx, dw = test_func_jitted(x, w, w_amax)
+ ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global)
+
+ assert_allclose(out, ref_out, dtype=jnp.float8_e4m3fn)
+ assert_allclose(dx, ref_dx, dtype=jnp.float8_e5m2)
+ assert_allclose(dw, ref_dw, dtype=jnp.float8_e5m2)
+
+
+if __name__ == "__main__":
+ from jax.sharding import NamedSharding, PartitionSpec
+ from jax.experimental import shard_map
+ import sys
+
+ coord_addr = sys.argv[1]
+ proc_id = int(sys.argv[2])
+ num_procs = int(sys.argv[3])
+
+ jax.distributed.initialize(
+ coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id
+ )
+
+ mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,))
+
+ with mesh:
+ data_shapes = [((4, 16, 128, 7168), (7168, 2048))]
+ for data_shape in data_shapes:
+ for kernel_fsdp_axis in [1, 2]:
+ test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis)

transformer_engine/jax/cpp_extensions/quantization.py

@@ -931,6 +931,7 @@ def grouped_quantize(
  x: jnp.ndarray,
  quantizer: GroupedQuantizer,
  group_sizes: jnp.ndarray = None,
+ amax: jnp.ndarray = None,
  flatten_axis: int = -1,
 ) -> GroupedScaledTensor1x:
  """Quantize a tensor in grouped manner.
@@ -943,6 +944,7 @@ def grouped_quantize(
  x: Input tensor to quantize
  quantizer: The quantizer to use for quantization
  group_sizes: Array of ints containing the size of each group (default: None)
+ amax: The amax of x; if None, it is auto-generated. (default: None)
  flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1)
 
  Returns:
@@ -985,7 +987,10 @@ def grouped_quantize(
  scale = scale.at[i].set(quantizer_i.scale[0])
 
  if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
- row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
+ if amax is not None:
+ row_amax = amax
+ else:
+ row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
  segment_ids = jnp.repeat(
  jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis]
  )

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -285,18 +285,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type
  size_t out_dtype_bytes = te_dtype_bytes(out_dtype);
 
  if (is_tensor_scaling) {
- cudaStream_t stream_0 = nvte_get_compute_stream(0);
  size_t dpitch = tensor_scaling_sinv_aligment;
  size_t spitch = lhs_sinv_dtype_bytes;
  size_t width = lhs_sinv_dtype_bytes;
  size_t height = lhs_sinv_size;
  cudaMemcpy2DAsync(lhs_scatter_aligned_ptr, dpitch, lhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  spitch = rhs_sinv_dtype_bytes;
  width = rhs_sinv_dtype_bytes;
  height = rhs_sinv_size;
  cudaMemcpy2DAsync(rhs_scatter_aligned_ptr, dpitch, rhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  lhs_sinv_ptr = lhs_scatter_aligned_ptr;
  rhs_sinv_ptr = rhs_scatter_aligned_ptr;
  }

qa/L1_jax_distributed_unittest/test.sh

@@ -9,3 +9,4 @@ set -xe
 mkdir -p "$XML_LOG_DIR"
 
 NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+SCRIPT_NAME=test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh

tests/jax/multi_process_launch.sh

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+#!/bin/bash
+
+SCRIPT_NAME="${SCRIPT_NAME:-test.py}"
+
+
+XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
+ --xla_gpu_enable_command_buffer=''"
+
+export XLA_FLAGS="${XLA_BASE_FLAGS}"
+
+NUM_RUNS=$(nvidia-smi --query-gpu=count --format=csv,noheader)
+for ((i=1; i<NUM_RUNS; i++))
+do
+ CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_PROC > /dev/null 2>&1 &
+done
+
+CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_PROC
+
+wait

tests/jax/test_multi_process_distributed_grouped_gemm.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from transformer_engine.jax.dense import grouped_dense as te_grouped_dense
+from transformer_engine.jax.quantize import (
+ QuantizerFactory,
+ ScalingMode,
+)
+
+from utils import assert_allclose
+
+
+N_GROUP = 8
+MESH_AXIS_NAME = "fsdp"
+
+
+def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis):
+ assert kernel_fsdp_axis in [1, 2]
+ x_shape, w_shape = data_shapes
+
+ x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None))
+ w_sharding = (
+ NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME))
+ if kernel_fsdp_axis == 2
+ else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None))
+ )
+ w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None))
+
+ def init_data():
+ x_key = jax.random.PRNGKey(0)
+ w_key = jax.random.PRNGKey(1)
+ x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16)
+ w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16)
+ w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim))
+ return x, w, w, w_amax
+
+ def test_func(outter_x, outter_w, outter_w_amax):
+ in_specs = (x_sharding.spec, w_sharding.spec, None)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w, w_amax):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+
+ output = te_grouped_dense(
+ x_reshaped,
+ w,
+ n_groups,
+ kernel_amax=w_amax,
+ quantizer_set=quantizer_set,
+ kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis),
+ )
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w, w_amax):
+ output = sharded_group_gemm(x, w, w_amax)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax)
+ dx, dw, _ = vjp_fn(output)
+ return output, dx, dw
+
+ def ref_func(outter_x, outter_w):
+
+ in_specs = (x_sharding.spec, w_no_sharding.spec)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+ output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set)
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w):
+ output = sharded_group_gemm(x, w)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w)
+ dx, dw = vjp_fn(output)
+ return output, dx, dw
+
+ init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None))
+ x, w, w_global, w_amax = init_func()
+
+ o_sharding = x_sharding
+ test_func_jitted = jax.jit(
+ test_func,
+ in_shardings=(x_sharding, w_sharding, None),
+ out_shardings=(o_sharding, x_sharding, w_sharding),
+ )
+ ref_func_jitted = jax.jit(
+ ref_func,
+ in_shardings=(x_sharding, w_no_sharding),
+ out_shardings=(o_sharding, x_sharding, w_no_sharding),
+ )
+
+ out, dx, dw = test_func_jitted(x, w, w_amax)
+ ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global)
+
+ assert_allclose(out, ref_out, dtype=jnp.float8_e4m3fn)
+ assert_allclose(dx, ref_dx, dtype=jnp.float8_e5m2)
+ assert_allclose(dw, ref_dw, dtype=jnp.float8_e5m2)
+
+
+if __name__ == "__main__":
+ from jax.sharding import NamedSharding, PartitionSpec
+ from jax.experimental import shard_map
+ import sys
+
+ coord_addr = sys.argv[1]
+ proc_id = int(sys.argv[2])
+ num_procs = int(sys.argv[3])
+
+ jax.distributed.initialize(
+ coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id
+ )
+
+ mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,))
+
+ with mesh:
+ data_shapes = [((4, 16, 128, 7168), (7168, 2048))]
+ for data_shape in data_shapes:
+ for kernel_fsdp_axis in [1, 2]:
+ test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis)

transformer_engine/jax/cpp_extensions/quantization.py

@@ -931,6 +931,7 @@ def grouped_quantize(
  x: jnp.ndarray,
  quantizer: GroupedQuantizer,
  group_sizes: jnp.ndarray = None,
+ amax: jnp.ndarray = None,
  flatten_axis: int = -1,
 ) -> GroupedScaledTensor1x:
  """Quantize a tensor in grouped manner.
@@ -943,6 +944,7 @@ def grouped_quantize(
  x: Input tensor to quantize
  quantizer: The quantizer to use for quantization
  group_sizes: Array of ints containing the size of each group (default: None)
+ amax: The amax of x; if None, it is auto-generated. (default: None)
  flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1)
 
  Returns:
@@ -985,7 +987,10 @@ def grouped_quantize(
  scale = scale.at[i].set(quantizer_i.scale[0])
 
  if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
- row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
+ if amax is not None:
+ row_amax = amax
+ else:
+ row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
  segment_ids = jnp.repeat(
  jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis]
  )

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -285,18 +285,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type
  size_t out_dtype_bytes = te_dtype_bytes(out_dtype);
 
  if (is_tensor_scaling) {
- cudaStream_t stream_0 = nvte_get_compute_stream(0);
  size_t dpitch = tensor_scaling_sinv_aligment;
  size_t spitch = lhs_sinv_dtype_bytes;
  size_t width = lhs_sinv_dtype_bytes;
  size_t height = lhs_sinv_size;
  cudaMemcpy2DAsync(lhs_scatter_aligned_ptr, dpitch, lhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  spitch = rhs_sinv_dtype_bytes;
  width = rhs_sinv_dtype_bytes;
  height = rhs_sinv_size;
  cudaMemcpy2DAsync(rhs_scatter_aligned_ptr, dpitch, rhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  lhs_sinv_ptr = lhs_scatter_aligned_ptr;
  rhs_sinv_ptr = rhs_scatter_aligned_ptr;
  }

qa/L1_jax_distributed_unittest/test.sh

@@ -9,3 +9,4 @@ set -xe
 mkdir -p "$XML_LOG_DIR"
 
 NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+SCRIPT_NAME=test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh

tests/jax/multi_process_launch.sh

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+#!/bin/bash
+
+SCRIPT_NAME="${SCRIPT_NAME:-test.py}"
+
+
+XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
+ --xla_gpu_enable_command_buffer=''"
+
+export XLA_FLAGS="${XLA_BASE_FLAGS}"
+
+NUM_RUNS=$(nvidia-smi --query-gpu=count --format=csv,noheader)
+for ((i=1; i<NUM_RUNS; i++))
+do
+ CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_PROC > /dev/null 2>&1 &
+done
+
+CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_PROC
+
+wait

tests/jax/test_multi_process_distributed_grouped_gemm.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from transformer_engine.jax.dense import grouped_dense as te_grouped_dense
+from transformer_engine.jax.quantize import (
+ QuantizerFactory,
+ ScalingMode,
+)
+
+from utils import assert_allclose
+
+
+N_GROUP = 8
+MESH_AXIS_NAME = "fsdp"
+
+
+def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis):
+ assert kernel_fsdp_axis in [1, 2]
+ x_shape, w_shape = data_shapes
+
+ x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None))
+ w_sharding = (
+ NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME))
+ if kernel_fsdp_axis == 2
+ else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None))
+ )
+ w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None))
+
+ def init_data():
+ x_key = jax.random.PRNGKey(0)
+ w_key = jax.random.PRNGKey(1)
+ x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16)
+ w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16)
+ w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim))
+ return x, w, w, w_amax
+
+ def test_func(outter_x, outter_w, outter_w_amax):
+ in_specs = (x_sharding.spec, w_sharding.spec, None)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w, w_amax):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+
+ output = te_grouped_dense(
+ x_reshaped,
+ w,
+ n_groups,
+ kernel_amax=w_amax,
+ quantizer_set=quantizer_set,
+ kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis),
+ )
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w, w_amax):
+ output = sharded_group_gemm(x, w, w_amax)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax)
+ dx, dw, _ = vjp_fn(output)
+ return output, dx, dw
+
+ def ref_func(outter_x, outter_w):
+
+ in_specs = (x_sharding.spec, w_no_sharding.spec)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+ output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set)
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w):
+ output = sharded_group_gemm(x, w)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w)
+ dx, dw = vjp_fn(output)
+ return output, dx, dw
+
+ init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None))
+ x, w, w_global, w_amax = init_func()
+
+ o_sharding = x_sharding
+ test_func_jitted = jax.jit(
+ test_func,
+ in_shardings=(x_sharding, w_sharding, None),
+ out_shardings=(o_sharding, x_sharding, w_sharding),
+ )
+ ref_func_jitted = jax.jit(
+ ref_func,
+ in_shardings=(x_sharding, w_no_sharding),
+ out_shardings=(o_sharding, x_sharding, w_no_sharding),
+ )
+
+ out, dx, dw = test_func_jitted(x, w, w_amax)
+ ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global)
+
+ assert_allclose(out, ref_out, dtype=jnp.float8_e4m3fn)
+ assert_allclose(dx, ref_dx, dtype=jnp.float8_e5m2)
+ assert_allclose(dw, ref_dw, dtype=jnp.float8_e5m2)
+
+
+if __name__ == "__main__":
+ from jax.sharding import NamedSharding, PartitionSpec
+ from jax.experimental import shard_map
+ import sys
+
+ coord_addr = sys.argv[1]
+ proc_id = int(sys.argv[2])
+ num_procs = int(sys.argv[3])
+
+ jax.distributed.initialize(
+ coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id
+ )
+
+ mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,))
+
+ with mesh:
+ data_shapes = [((4, 16, 128, 7168), (7168, 2048))]
+ for data_shape in data_shapes:
+ for kernel_fsdp_axis in [1, 2]:
+ test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis)

transformer_engine/jax/cpp_extensions/quantization.py

@@ -931,6 +931,7 @@ def grouped_quantize(
  x: jnp.ndarray,
  quantizer: GroupedQuantizer,
  group_sizes: jnp.ndarray = None,
+ amax: jnp.ndarray = None,
  flatten_axis: int = -1,
 ) -> GroupedScaledTensor1x:
  """Quantize a tensor in grouped manner.
@@ -943,6 +944,7 @@ def grouped_quantize(
  x: Input tensor to quantize
  quantizer: The quantizer to use for quantization
  group_sizes: Array of ints containing the size of each group (default: None)
+ amax: The amax of x; if None, it is auto-generated. (default: None)
  flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1)
 
  Returns:
@@ -985,7 +987,10 @@ def grouped_quantize(
  scale = scale.at[i].set(quantizer_i.scale[0])
 
  if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
- row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
+ if amax is not None:
+ row_amax = amax
+ else:
+ row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
  segment_ids = jnp.repeat(
  jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis]
  )

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -285,18 +285,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type
  size_t out_dtype_bytes = te_dtype_bytes(out_dtype);
 
  if (is_tensor_scaling) {
- cudaStream_t stream_0 = nvte_get_compute_stream(0);
  size_t dpitch = tensor_scaling_sinv_aligment;
  size_t spitch = lhs_sinv_dtype_bytes;
  size_t width = lhs_sinv_dtype_bytes;
  size_t height = lhs_sinv_size;
  cudaMemcpy2DAsync(lhs_scatter_aligned_ptr, dpitch, lhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  spitch = rhs_sinv_dtype_bytes;
  width = rhs_sinv_dtype_bytes;
  height = rhs_sinv_size;
  cudaMemcpy2DAsync(rhs_scatter_aligned_ptr, dpitch, rhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  lhs_sinv_ptr = lhs_scatter_aligned_ptr;
  rhs_sinv_ptr = rhs_scatter_aligned_ptr;
  }

qa/L1_jax_distributed_unittest/test.sh

@@ -9,3 +9,4 @@ set -xe
 mkdir -p "$XML_LOG_DIR"
 
 NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+SCRIPT_NAME=test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh

tests/jax/multi_process_launch.sh

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+#!/bin/bash
+
+SCRIPT_NAME="${SCRIPT_NAME:-test.py}"
+
+
+XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
+ --xla_gpu_enable_command_buffer=''"
+
+export XLA_FLAGS="${XLA_BASE_FLAGS}"
+
+NUM_RUNS=$(nvidia-smi --query-gpu=count --format=csv,noheader)
+for ((i=1; i<NUM_RUNS; i++))
+do
+ CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_PROC > /dev/null 2>&1 &
+done
+
+CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_PROC
+
+wait

tests/jax/test_multi_process_distributed_grouped_gemm.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from transformer_engine.jax.dense import grouped_dense as te_grouped_dense
+from transformer_engine.jax.quantize import (
+ QuantizerFactory,
+ ScalingMode,
+)
+
+from utils import assert_allclose
+
+
+N_GROUP = 8
+MESH_AXIS_NAME = "fsdp"
+
+
+def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis):
+ assert kernel_fsdp_axis in [1, 2]
+ x_shape, w_shape = data_shapes
+
+ x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None))
+ w_sharding = (
+ NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME))
+ if kernel_fsdp_axis == 2
+ else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None))
+ )
+ w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None))
+
+ def init_data():
+ x_key = jax.random.PRNGKey(0)
+ w_key = jax.random.PRNGKey(1)
+ x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16)
+ w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16)
+ w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim))
+ return x, w, w, w_amax
+
+ def test_func(outter_x, outter_w, outter_w_amax):
+ in_specs = (x_sharding.spec, w_sharding.spec, None)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w, w_amax):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+
+ output = te_grouped_dense(
+ x_reshaped,
+ w,
+ n_groups,
+ kernel_amax=w_amax,
+ quantizer_set=quantizer_set,
+ kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis),
+ )
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w, w_amax):
+ output = sharded_group_gemm(x, w, w_amax)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax)
+ dx, dw, _ = vjp_fn(output)
+ return output, dx, dw
+
+ def ref_func(outter_x, outter_w):
+
+ in_specs = (x_sharding.spec, w_no_sharding.spec)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+ output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set)
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w):
+ output = sharded_group_gemm(x, w)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w)
+ dx, dw = vjp_fn(output)
+ return output, dx, dw
+
+ init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None))
+ x, w, w_global, w_amax = init_func()
+
+ o_sharding = x_sharding
+ test_func_jitted = jax.jit(
+ test_func,
+ in_shardings=(x_sharding, w_sharding, None),
+ out_shardings=(o_sharding, x_sharding, w_sharding),
+ )
+ ref_func_jitted = jax.jit(
+ ref_func,
+ in_shardings=(x_sharding, w_no_sharding),
+ out_shardings=(o_sharding, x_sharding, w_no_sharding),
+ )
+
+ out, dx, dw = test_func_jitted(x, w, w_amax)
+ ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global)
+
+ assert_allclose(out, ref_out, dtype=jnp.float8_e4m3fn)
+ assert_allclose(dx, ref_dx, dtype=jnp.float8_e5m2)
+ assert_allclose(dw, ref_dw, dtype=jnp.float8_e5m2)
+
+
+if __name__ == "__main__":
+ from jax.sharding import NamedSharding, PartitionSpec
+ from jax.experimental import shard_map
+ import sys
+
+ coord_addr = sys.argv[1]
+ proc_id = int(sys.argv[2])
+ num_procs = int(sys.argv[3])
+
+ jax.distributed.initialize(
+ coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id
+ )
+
+ mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,))
+
+ with mesh:
+ data_shapes = [((4, 16, 128, 7168), (7168, 2048))]
+ for data_shape in data_shapes:
+ for kernel_fsdp_axis in [1, 2]:
+ test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis)

transformer_engine/jax/cpp_extensions/quantization.py

@@ -931,6 +931,7 @@ def grouped_quantize(
  x: jnp.ndarray,
  quantizer: GroupedQuantizer,
  group_sizes: jnp.ndarray = None,
+ amax: jnp.ndarray = None,
  flatten_axis: int = -1,
 ) -> GroupedScaledTensor1x:
  """Quantize a tensor in grouped manner.
@@ -943,6 +944,7 @@ def grouped_quantize(
  x: Input tensor to quantize
  quantizer: The quantizer to use for quantization
  group_sizes: Array of ints containing the size of each group (default: None)
+ amax: The amax of x; if None, it is auto-generated. (default: None)
  flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1)
 
  Returns:
@@ -985,7 +987,10 @@ def grouped_quantize(
  scale = scale.at[i].set(quantizer_i.scale[0])
 
  if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
- row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
+ if amax is not None:
+ row_amax = amax
+ else:
+ row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
  segment_ids = jnp.repeat(
  jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis]
  )

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -285,18 +285,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type
  size_t out_dtype_bytes = te_dtype_bytes(out_dtype);
 
  if (is_tensor_scaling) {
- cudaStream_t stream_0 = nvte_get_compute_stream(0);
  size_t dpitch = tensor_scaling_sinv_aligment;
  size_t spitch = lhs_sinv_dtype_bytes;
  size_t width = lhs_sinv_dtype_bytes;
  size_t height = lhs_sinv_size;
  cudaMemcpy2DAsync(lhs_scatter_aligned_ptr, dpitch, lhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  spitch = rhs_sinv_dtype_bytes;
  width = rhs_sinv_dtype_bytes;
  height = rhs_sinv_size;
  cudaMemcpy2DAsync(rhs_scatter_aligned_ptr, dpitch, rhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  lhs_sinv_ptr = lhs_scatter_aligned_ptr;
  rhs_sinv_ptr = rhs_scatter_aligned_ptr;
  }

qa/L1_jax_distributed_unittest/test.sh

@@ -9,3 +9,4 @@ set -xe
 mkdir -p "$XML_LOG_DIR"
 
 NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+SCRIPT_NAME=test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh

tests/jax/multi_process_launch.sh

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+#!/bin/bash
+
+SCRIPT_NAME="${SCRIPT_NAME:-test.py}"
+
+
+XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
+ --xla_gpu_enable_command_buffer=''"
+
+export XLA_FLAGS="${XLA_BASE_FLAGS}"
+
+NUM_RUNS=$(nvidia-smi --query-gpu=count --format=csv,noheader)
+for ((i=1; i<NUM_RUNS; i++))
+do
+ CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_PROC > /dev/null 2>&1 &
+done
+
+CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_PROC
+
+wait

tests/jax/test_multi_process_distributed_grouped_gemm.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from transformer_engine.jax.dense import grouped_dense as te_grouped_dense
+from transformer_engine.jax.quantize import (
+ QuantizerFactory,
+ ScalingMode,
+)
+
+from utils import assert_allclose
+
+
+N_GROUP = 8
+MESH_AXIS_NAME = "fsdp"
+
+
+def test_grouped_gemm_fp8_allgather(data_shapes, kernel_fsdp_axis):
+ assert kernel_fsdp_axis in [1, 2]
+ x_shape, w_shape = data_shapes
+
+ x_sharding = NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None, None, None))
+ w_sharding = (
+ NamedSharding(mesh, PartitionSpec(None, None, MESH_AXIS_NAME))
+ if kernel_fsdp_axis == 2
+ else NamedSharding(mesh, PartitionSpec(None, MESH_AXIS_NAME, None))
+ )
+ w_no_sharding = NamedSharding(mesh, PartitionSpec(None, None, None))
+
+ def init_data():
+ x_key = jax.random.PRNGKey(0)
+ w_key = jax.random.PRNGKey(1)
+ x = jax.random.normal(x_key, shape=(N_GROUP, *x_shape), dtype=jnp.bfloat16)
+ w = jax.random.normal(w_key, shape=(N_GROUP, *w_shape), dtype=jnp.bfloat16)
+ w_amax = jnp.max(jnp.abs(w), axis=range(1, w.ndim))
+ return x, w, w, w_amax
+
+ def test_func(outter_x, outter_w, outter_w_amax):
+ in_specs = (x_sharding.spec, w_sharding.spec, None)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w, w_amax):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+
+ output = te_grouped_dense(
+ x_reshaped,
+ w,
+ n_groups,
+ kernel_amax=w_amax,
+ quantizer_set=quantizer_set,
+ kernel_fsdp_info=(MESH_AXIS_NAME, kernel_fsdp_axis),
+ )
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w, w_amax):
+ output = sharded_group_gemm(x, w, w_amax)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w, outter_w_amax)
+ dx, dw, _ = vjp_fn(output)
+ return output, dx, dw
+
+ def ref_func(outter_x, outter_w):
+
+ in_specs = (x_sharding.spec, w_no_sharding.spec)
+ out_specs = x_sharding.spec
+
+ @partial(
+ shard_map.shard_map,
+ mesh=mesh,
+ in_specs=in_specs,
+ out_specs=out_specs,
+ check_rep=False,
+ )
+ def sharded_group_gemm(x, w):
+ group_size = x.shape[0]
+ x_reshaped = x.reshape(-1, x.shape[-1])
+ n_groups = jnp.full(group_size, x_reshaped.shape[0] // group_size)
+
+ quantizer_set = QuantizerFactory.create_set(
+ scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+ fwd_dtype=jnp.float8_e4m3fn,
+ bwd_dtype=jnp.float8_e5m2,
+ is_2x2x=True,
+ n_groups=group_size,
+ )
+ output = te_grouped_dense(x_reshaped, w, n_groups, quantizer_set=quantizer_set)
+ output = output.reshape(*x.shape[:-1], -1)
+ return output
+
+ def run(x, w):
+ output = sharded_group_gemm(x, w)
+ return output
+
+ output, vjp_fn = jax.vjp(run, outter_x, outter_w)
+ dx, dw = vjp_fn(output)
+ return output, dx, dw
+
+ init_func = jax.jit(init_data, out_shardings=(x_sharding, w_sharding, w_no_sharding, None))
+ x, w, w_global, w_amax = init_func()
+
+ o_sharding = x_sharding
+ test_func_jitted = jax.jit(
+ test_func,
+ in_shardings=(x_sharding, w_sharding, None),
+ out_shardings=(o_sharding, x_sharding, w_sharding),
+ )
+ ref_func_jitted = jax.jit(
+ ref_func,
+ in_shardings=(x_sharding, w_no_sharding),
+ out_shardings=(o_sharding, x_sharding, w_no_sharding),
+ )
+
+ out, dx, dw = test_func_jitted(x, w, w_amax)
+ ref_out, ref_dx, ref_dw = ref_func_jitted(x, w_global)
+
+ assert_allclose(out, ref_out, dtype=jnp.float8_e4m3fn)
+ assert_allclose(dx, ref_dx, dtype=jnp.float8_e5m2)
+ assert_allclose(dw, ref_dw, dtype=jnp.float8_e5m2)
+
+
+if __name__ == "__main__":
+ from jax.sharding import NamedSharding, PartitionSpec
+ from jax.experimental import shard_map
+ import sys
+
+ coord_addr = sys.argv[1]
+ proc_id = int(sys.argv[2])
+ num_procs = int(sys.argv[3])
+
+ jax.distributed.initialize(
+ coordinator_address=coord_addr, num_processes=num_procs, process_id=proc_id
+ )
+
+ mesh = jax.make_mesh((num_procs,), (MESH_AXIS_NAME,))
+
+ with mesh:
+ data_shapes = [((4, 16, 128, 7168), (7168, 2048))]
+ for data_shape in data_shapes:
+ for kernel_fsdp_axis in [1, 2]:
+ test_grouped_gemm_fp8_allgather(data_shape, kernel_fsdp_axis)

transformer_engine/jax/cpp_extensions/quantization.py

@@ -931,6 +931,7 @@ def grouped_quantize(
  x: jnp.ndarray,
  quantizer: GroupedQuantizer,
  group_sizes: jnp.ndarray = None,
+ amax: jnp.ndarray = None,
  flatten_axis: int = -1,
 ) -> GroupedScaledTensor1x:
  """Quantize a tensor in grouped manner.
@@ -943,6 +944,7 @@ def grouped_quantize(
  x: Input tensor to quantize
  quantizer: The quantizer to use for quantization
  group_sizes: Array of ints containing the size of each group (default: None)
+ amax: The amax of x; if None, it is auto-generated. (default: None)
  flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1)
 
  Returns:
@@ -985,7 +987,10 @@ def grouped_quantize(
  scale = scale.at[i].set(quantizer_i.scale[0])
 
  if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
- row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
+ if amax is not None:
+ row_amax = amax
+ else:
+ row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim))
  segment_ids = jnp.repeat(
  jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis]
  )

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -285,18 +285,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type
  size_t out_dtype_bytes = te_dtype_bytes(out_dtype);
 
  if (is_tensor_scaling) {
- cudaStream_t stream_0 = nvte_get_compute_stream(0);
  size_t dpitch = tensor_scaling_sinv_aligment;
  size_t spitch = lhs_sinv_dtype_bytes;
  size_t width = lhs_sinv_dtype_bytes;
  size_t height = lhs_sinv_size;
  cudaMemcpy2DAsync(lhs_scatter_aligned_ptr, dpitch, lhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  spitch = rhs_sinv_dtype_bytes;
  width = rhs_sinv_dtype_bytes;
  height = rhs_sinv_size;
  cudaMemcpy2DAsync(rhs_scatter_aligned_ptr, dpitch, rhs_sinv_ptr, spitch, width, height,
- cudaMemcpyDeviceToDevice, stream_0);
+ cudaMemcpyDeviceToDevice, stream);
  lhs_sinv_ptr = lhs_scatter_aligned_ptr;
  rhs_sinv_ptr = rhs_scatter_aligned_ptr;
  }