[mxfp8 moe training] integrate triton quant/dequant kernels into mxfp8 all to all

Author: danielvegamyhre

benchmarks/prototype/moe_training/mxfp8/bench_all_to_all_v.py

@@ -42,8 +42,10 @@ class ExperimentConfig:
 
 @dataclass(frozen=True)
 class ExperimentResult:
-    bf16_ms: float
-    mxfp8_ms: float
+    bf16_fwd_ms: float
+    mxfp8_fwd_ms: float
+    bf16_bwd_ms: float
+    mxfp8_bwd_ms: float
 
 
 @dataclass(frozen=True)
@@ -67,7 +69,7 @@ def get_configs() -> List[ExperimentConfig]:
     return configs
 
 
-def default_a2a_fwd_bwd(
+def default_a2a_fwd(
     routed_input: torch.Tensor,
     labels: torch.Tensor,
     output_splits_list: list[int],
@@ -81,15 +83,11 @@ def default_a2a_fwd_bwd(
         device_mesh.get_group(),
     )
     routed_input = torch.ops._c10d_functional.wait_tensor(routed_input)
-
-    loss = F.mse_loss(routed_input, labels)
-    loss.backward()
-
     torch.cuda.synchronize()
     return routed_input
 
 
-def mxfp8_a2a_fwd_bwd(
+def mxfp8_a2a_fwd(
     routed_input: torch.Tensor,
     labels: torch.Tensor,
     output_splits_list: list[int],
@@ -102,16 +100,17 @@ def mxfp8_a2a_fwd_bwd(
         input_splits_list,
         device_mesh.get_group(),
     )
-
-    loss = F.mse_loss(routed_input, labels)
-    loss.backward()
     torch.cuda.synchronize()
     return routed_input
 
 
-# Compile target funcs
-default_a2a_sync_compiled = torch.compile(default_a2a_fwd_bwd)
-mxfp8_a2a_sync_compiled = torch.compile(mxfp8_a2a_fwd_bwd)
+def mse_loss_and_bwd(
+    routed_input: torch.Tensor,
+    labels: torch.Tensor,
+):
+    loss = F.mse_loss(routed_input, labels)
+    loss.backward()
+    torch.cuda.synchronize()
 
 
 def run_experiment(
@@ -149,62 +148,78 @@ def warmup(func_no_args):
 
     # Bench default a2a (exclude d2h sync from preparing input splits_list and output_splits_list)
     warmup(
-        lambda: default_a2a_sync_compiled(
+        lambda: default_a2a_fwd(
             ref_x, labels, output_splits_list, input_splits_list, mesh
         )
     )
     start_sec = time.perf_counter()
-    default_a2a_sync_compiled(
+    bf16_output = default_a2a_fwd(
         ref_x, labels, output_splits_list, input_splits_list, mesh
     )
     end_sec = time.perf_counter()
-    bf16_ms = (end_sec - start_sec) * 1e3
+    bf16_fwd_ms = (end_sec - start_sec) * 1e3
     if args.profile:
         profile_fn(
-            default_a2a_sync_compiled,
+            default_a2a_fwd,
             ref_x,
             labels,
             output_splits_list,
             input_splits_list,
             mesh,
             distributed=True,
-            profile_name="all_to_all_single_autograd",
+            profile_name="default_a2a_fwd",
         )
 
+    # Bench default a2a bwd
+    warmup(lambda: mse_loss_and_bwd(bf16_output, labels))
+    start_sec = time.perf_counter()
+    mse_loss_and_bwd(bf16_output, labels)
+    end_sec = time.perf_counter()
+    bf16_bwd_ms = (end_sec - start_sec) * 1e3
+
     # Bench mxfp8 sync a2a (exclude d2h sync from preparing input splits_list and output_splits_list)
     warmup(
-        lambda: mxfp8_a2a_sync_compiled(
-            x, labels, output_splits_list, input_splits_list, mesh
-        )
+        lambda: mxfp8_a2a_fwd(x, labels, output_splits_list, input_splits_list, mesh)
     )
     start_sec = time.perf_counter()
-    mxfp8_a2a_sync_compiled(x, labels, output_splits_list, input_splits_list, mesh)
+    mxfp8_output = mxfp8_a2a_fwd(x, labels, output_splits_list, input_splits_list, mesh)
     end_sec = time.perf_counter()
     mxfp8_ms = (end_sec - start_sec) * 1e3
     if args.profile:
         profile_fn(
-            mxfp8_a2a_sync_compiled,
+            mxfp8_a2a_fwd,
             x,
             labels,
             output_splits_list,
             input_splits_list,
             mesh,
             distributed=True,
-            profile_name="to_mxfp8_a2a_dequant",
+            profile_name="mxfp8_a2a_fwd",
         )
 
+    # Bench mxfp8 a2a bwd
+    warmup(lambda: mse_loss_and_bwd(mxfp8_output, labels))
+    start_sec = time.perf_counter()
+    mse_loss_and_bwd(mxfp8_output, labels)
+    end_sec = time.perf_counter()
+    mxfp8_bwd_ms = (end_sec - start_sec) * 1e3
+
     return ExperimentResult(
-        bf16_ms=bf16_ms,
-        mxfp8_ms=mxfp8_ms,
+        bf16_fwd_ms=bf16_fwd_ms,
+        mxfp8_fwd_ms=mxfp8_ms,
+        bf16_bwd_ms=bf16_bwd_ms,
+        mxfp8_bwd_ms=mxfp8_bwd_ms,
     )
 
 
 def print_results(experiments: List[Experiment]):
     headers = [
         "input_shape",
         "num_splits",
-        "bf16_ms",
-        "mxfp8_ms",
+        "fwd_bf16_ms",
+        "fwd_mxfp8_ms",
+        "bwd_bf16_ms",
+        "bwd_mxfp8_ms",
     ]
     rows = []
     num_splits = dist.get_world_size()
@@ -213,8 +228,10 @@ def print_results(experiments: List[Experiment]):
             [
                 str(experiment.config.input_shape),
                 num_splits,
-                experiment.result.bf16_ms,
-                experiment.result.mxfp8_ms,
+                experiment.result.bf16_fwd_ms,
+                experiment.result.mxfp8_fwd_ms,
+                experiment.result.bf16_bwd_ms,
+                experiment.result.mxfp8_bwd_ms,
             ]
         )
     print(tabulate(rows, headers=headers))

torchao/prototype/moe_training/kernels/mxfp8/comms.py

@@ -11,7 +11,10 @@
     blockwise_barrier,
     sync_threads,
 )
-from torchao.prototype.mx_formats.config import ScaleCalculationMode
+from torchao.prototype.mx_formats.kernels import (
+    triton_mxfp8_dequant_dim0,
+    triton_to_mxfp8_dim0,
+)
 from torchao.prototype.mx_formats.mx_tensor import to_dtype, to_mx
 
 
@@ -473,11 +476,9 @@ def forward(
         """
         # Quantize input
         block_size = 32
-        input_scales, input_data = to_mx(
+        input_data, input_scales = triton_to_mxfp8_dim0(
             input,
-            elem_dtype=torch.float8_e4m3fn,
-            block_size=block_size,
-            scaling_mode=ScaleCalculationMode.RCEIL,
+            inner_block_size=block_size,
         )
 
         # Dispatch data (async)
@@ -501,20 +502,17 @@ def forward(
         output_data = torch.ops._c10d_functional.wait_tensor(output_data)
 
         # Dequantize output
-        lowp_dtype = output_data.dtype
         hp_dtype = input.dtype
-        hp_output = to_dtype(
+        triton_hp_output = triton_mxfp8_dequant_dim0(
             output_data,
-            output_scales.view(torch.float8_e8m0fnu),
-            lowp_dtype,
-            block_size,
+            output_scales,
             hp_dtype,
+            block_size,
         )
-
         ctx.input_splits = input_splits
         ctx.output_splits = output_splits
         ctx.group = group
-        return hp_output
+        return triton_hp_output
 
     @staticmethod
     def backward(ctx, grad_output_hp):
@@ -529,11 +527,9 @@ def backward(ctx, grad_output_hp):
 
         # Quantize grad_output
         block_size = 32
-        grad_out_scales, grad_out_data = to_mx(
+        grad_out_data, grad_out_scales = triton_to_mxfp8_dim0(
             grad_output_hp,
-            elem_dtype=torch.float8_e4m3fn,
-            block_size=block_size,
-            scaling_mode=ScaleCalculationMode.RCEIL,
+            inner_block_size=block_size,
         )
 
         # Dispatch data (async)
@@ -557,13 +553,11 @@ def backward(ctx, grad_output_hp):
         grad_input_scales = torch.ops._c10d_functional.wait_tensor(grad_input_scales)
 
         hp_dtype = grad_output_hp.dtype
-        lowp_dtype = grad_input_data.dtype
-        grad_input_hp = to_dtype(
+        grad_input_hp = triton_mxfp8_dequant_dim0(
             grad_input_data,
-            grad_input_scales.view(torch.float8_e8m0fnu),
-            lowp_dtype,
-            block_size,
+            grad_input_scales,
             hp_dtype,
+            block_size,
         )
         return grad_input_hp, None, None, None