[mxfp8 moe training] increase num_warps in mxfp8 a2a comms kernel (pytorch#3087)

danielvegamyhre · web-flow · commit 0d3217d07ec1 · 2025-09-27T08:29:08.000-07:00
[mxfp8 moe training] increase num_warps in mxfp8 a2a kernel
diff --git a/torchao/prototype/moe_training/kernels/mxfp8/comms.py b/torchao/prototype/moe_training/kernels/mxfp8/comms.py
@@ -275,7 +275,7 @@ def _mxfp8_on_device_all_to_all_v(
         world_size=input_hdl.world_size,
         BLOCKS_PER_REMOTE_RANK=BLOCKS_PER_REMOTE_RANK,
         BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=1,
+        num_warps=16,
     )
 
     return output

Original file line number	Diff line number	Diff line change
`@@ -275,7 +275,7 @@ def _mxfp8_on_device_all_to_all_v(`
`275`	`275`	`world_size=input_hdl.world_size,`
`276`	`276`	`BLOCKS_PER_REMOTE_RANK=BLOCKS_PER_REMOTE_RANK,`
`277`	`277`	`BLOCK_SIZE=BLOCK_SIZE,`
`278`		`- num_warps=1,`
	`278`	`+ num_warps=16,`
`279`	`279`	`)`
`280`	`280`
`281`	`281`	`return output`