@@ -1212,9 +1212,15 @@ def __init__(
12121212 threads : int ,
12131213 device_slot_selection : bool ,
12141214 kimi_topk : bool ,
1215+ row_blocks : int = 1 ,
12151216 ) -> None :
12161217 super ().__init__ (world_size , rank , threads , device_slot_selection )
12171218 self ._kimi_topk = bool (kimi_topk )
1219+ if type (row_blocks ) is not int or not 1 <= row_blocks <= _MAX_BLOCKS :
1220+ raise ValueError ("paired gather row blocks must fit the signal plane" )
1221+ if kimi_topk and row_blocks != 1 :
1222+ raise ValueError ("fused single-row expert selection requires one CTA" )
1223+ self ._row_blocks = row_blocks
12181224
12191225 @cute .jit
12201226 def __call__ (
@@ -1261,6 +1267,7 @@ def __call__(
12611267 first_packs : Int32 ,
12621268 second_packs : Int32 ,
12631269 slot_delta_256b : Int32 ,
1270+ grid_x : Int32 ,
12641271 stream : cuda .CUstream ,
12651272 ) -> None :
12661273 self .kernel (
@@ -1307,7 +1314,7 @@ def __call__(
13071314 second_packs ,
13081315 slot_delta_256b ,
13091316 ).launch (
1310- grid = (1 , 1 , 1 ),
1317+ grid = (grid_x , 1 , 1 ),
13111318 block = (self ._threads , 1 , 1 ),
13121319 max_number_threads = (512 , 1 , 1 ),
13131320 min_blocks_per_mp = 1 ,
@@ -1398,6 +1405,14 @@ def kernel(
13981405 signal15 ,
13991406 )
14001407 tidx , _ , _ = cute .arch .thread_idx ()
1408+ row_group , grid = Int32 (0 ), Int32 (1 )
1409+ if cutlass .const_expr (self ._row_blocks != 1 ):
1410+ row_group , _ , _ = cute .arch .block_idx ()
1411+ grid , _ , _ = cute .arch .grid_dim ()
1412+ # Matching CTAs on every rank publish the same complete rows. A CTA
1413+ # never reads a row staged by a different CTA, so its peer barrier
1414+ # establishes all input dependencies without a local grid barrier.
1415+ owned_rows = (batch - row_group + grid - 1 ) // grid
14011416 slot_offset = Int64 (0 )
14021417 if cutlass .const_expr (self ._device_slot_selection ):
14031418 generation = ld_relaxed_gpu_u32 (
@@ -1409,21 +1424,29 @@ def kernel(
14091424 local_stage = self ._staging_words (staging [self ._rank ] + slot_offset )
14101425
14111426 linear = Int32 (tidx )
1412- while linear < batch * first_packs :
1413- batch_index = linear // first_packs
1414- pack = linear - batch_index * first_packs
1427+ while linear < owned_rows * first_packs :
1428+ local_row = linear // first_packs
1429+ batch_index = local_row * grid + row_group
1430+ pack = linear - local_row * first_packs
1431+ source_pack = Int64 (linear )
1432+ if cutlass .const_expr (self ._row_blocks != 1 ):
1433+ source_pack = Int64 (batch_index ) * Int64 (first_packs ) + Int64 (pack )
14151434 _copy_16b (
1416- local_first + Int64 ( linear ) * Int64 (4 ),
1435+ local_first + source_pack * Int64 (4 ),
14171436 local_stage
14181437 + (Int64 (batch_index ) * Int64 (combined_packs ) + Int64 (pack )) * Int64 (4 ),
14191438 )
14201439 linear += Int32 (self ._threads )
14211440 linear = Int32 (tidx )
1422- while linear < batch * second_packs :
1423- batch_index = linear // second_packs
1424- pack = linear - batch_index * second_packs
1441+ while linear < owned_rows * second_packs :
1442+ local_row = linear // second_packs
1443+ batch_index = local_row * grid + row_group
1444+ pack = linear - local_row * second_packs
1445+ source_pack = Int64 (linear )
1446+ if cutlass .const_expr (self ._row_blocks != 1 ):
1447+ source_pack = Int64 (batch_index ) * Int64 (second_packs ) + Int64 (pack )
14251448 _copy_16b (
1426- local_second + Int64 ( linear ) * Int64 (4 ),
1449+ local_second + source_pack * Int64 (4 ),
14271450 local_stage
14281451 + (
14291452 Int64 (batch_index ) * Int64 (combined_packs )
@@ -1442,11 +1465,12 @@ def kernel(
14421465 max_blocks = _MAX_BLOCKS ,
14431466 )
14441467
1445- first_output_packs = batch * Int32 (self ._world_size ) * first_packs
1468+ first_output_packs = owned_rows * Int32 (self ._world_size ) * first_packs
14461469 linear = Int32 (tidx )
14471470 while linear < first_output_packs :
14481471 batch_index = linear // (Int32 (self ._world_size ) * first_packs )
14491472 row_pack = linear - (batch_index * Int32 (self ._world_size ) * first_packs )
1473+ batch_index = batch_index * grid + row_group
14501474 source_rank = row_pack // first_packs
14511475 pack = row_pack - source_rank * first_packs
14521476 # Default to this rank's own slice so the address is always
@@ -1471,20 +1495,26 @@ def kernel(
14711495 source_address = Int64 (
14721496 (source_words + (source_base + Int64 (pack )) * Int64 (4 )).toint ()
14731497 )
1498+ output_pack = Int64 (linear )
1499+ if cutlass .const_expr (self ._row_blocks != 1 ):
1500+ output_pack = Int64 (batch_index ) * Int64 (self ._world_size ) * Int64 (
1501+ first_packs
1502+ ) + Int64 (row_pack )
14741503 _copy_16b_addr (
14751504 source_address ,
1476- Int64 ((output_first + Int64 ( linear ) * Int64 (4 )).toint ()),
1505+ Int64 ((output_first + output_pack * Int64 (4 )).toint ()),
14771506 )
14781507 linear += Int32 (self ._threads )
14791508
14801509 if cutlass .const_expr (not self ._kimi_topk ):
1481- second_output_packs = batch * Int32 (self ._world_size ) * second_packs
1510+ second_output_packs = owned_rows * Int32 (self ._world_size ) * second_packs
14821511 linear = Int32 (tidx )
14831512 while linear < second_output_packs :
14841513 batch_index = linear // (Int32 (self ._world_size ) * second_packs )
14851514 row_pack = linear - (
14861515 batch_index * Int32 (self ._world_size ) * second_packs
14871516 )
1517+ batch_index = batch_index * grid + row_group
14881518 source_rank = row_pack // second_packs
14891519 pack = row_pack - source_rank * second_packs
14901520 source_address = Int64 (
@@ -1511,9 +1541,14 @@ def kernel(
15111541 source_words + (source_base + Int64 (pack )) * Int64 (4 )
15121542 ).toint ()
15131543 )
1544+ output_pack = Int64 (linear )
1545+ if cutlass .const_expr (self ._row_blocks != 1 ):
1546+ output_pack = Int64 (batch_index ) * Int64 (self ._world_size ) * Int64 (
1547+ second_packs
1548+ ) + Int64 (row_pack )
15141549 _copy_16b_addr (
15151550 source_address ,
1516- Int64 ((output_second + Int64 ( linear ) * Int64 (4 )).toint ()),
1551+ Int64 ((output_second + output_pack * Int64 (4 )).toint ()),
15171552 )
15181553 linear += Int32 (self ._threads )
15191554 else :
@@ -1694,7 +1729,7 @@ class SharedStorage:
16941729 _a2a_graph_epoch_arrive (
16951730 (self_signal + Int64 (_GRAPH_EPOCH_INDEX )).toint (),
16961731 (self_signal + Int64 (_GRAPH_ARRIVED_INDEX )).toint (),
1697- Uint32 (1 ),
1732+ Uint32 (grid ),
16981733 )
16991734
17001735
@@ -2164,6 +2199,7 @@ def _pair_launcher_key(
21642199 threads : int ,
21652200 device_slot_selection : bool ,
21662201 kimi_topk : bool ,
2202+ row_blocks : int = 1 ,
21672203) -> tuple [object , ...]:
21682204 """Return the cache key for a paired all-gather launcher."""
21692205 return (
@@ -2172,6 +2208,7 @@ def _pair_launcher_key(
21722208 int (threads ),
21732209 bool (device_slot_selection ),
21742210 bool (kimi_topk ),
2211+ int (row_blocks ),
21752212 )
21762213
21772214
@@ -2181,6 +2218,7 @@ def is_all_gather_pair_prepared(
21812218 threads : int ,
21822219 device_slot_selection : bool ,
21832220 kimi_topk : bool = False ,
2221+ row_blocks : int = 1 ,
21842222) -> bool :
21852223 return (
21862224 _pair_launcher_key (
@@ -2189,6 +2227,7 @@ def is_all_gather_pair_prepared(
21892227 threads ,
21902228 device_slot_selection ,
21912229 kimi_topk ,
2230+ row_blocks ,
21922231 )
21932232 in _PREPARED_PAIR_LAUNCHERS
21942233 )
@@ -2201,6 +2240,7 @@ def _get_compiled_all_gather_pair(
22012240 threads : int ,
22022241 device_slot_selection : bool ,
22032242 kimi_topk : bool = False ,
2243+ row_blocks : int = 1 ,
22042244) -> Callable :
22052245 """Compile and retain one paired all-gather launcher specialization."""
22062246 launch = _AllGatherPairLaunch (
@@ -2209,13 +2249,15 @@ def _get_compiled_all_gather_pair(
22092249 threads ,
22102250 device_slot_selection ,
22112251 kimi_topk ,
2252+ row_blocks ,
22122253 )
22132254 key = _pair_launcher_key (
22142255 world_size ,
22152256 rank ,
22162257 threads ,
22172258 device_slot_selection ,
22182259 kimi_topk ,
2260+ row_blocks ,
22192261 )
22202262 raise_if_kernel_resolution_frozen ("cute.compile" , target = launch , cache_key = key )
22212263 p_u32 = _u32_ptr (16 , align = 4 )
@@ -2236,6 +2278,7 @@ def _get_compiled_all_gather_pair(
22362278 1 ,
22372279 1 ,
22382280 1 ,
2281+ 1 ,
22392282 current_cuda_stream (),
22402283 compile_spec = KernelCompileSpec .from_key (
22412284 (
@@ -2251,6 +2294,7 @@ def _get_compiled_all_gather_pair(
22512294 "threads" ,
22522295 "device_slot_selection" ,
22532296 "kimi_topk" ,
2297+ "row_blocks" ,
22542298 ),
22552299 ),
22562300 )
@@ -2285,6 +2329,7 @@ def run(
22852329 int (first_packs ),
22862330 int (second_packs ),
22872331 int (slot_delta_256b ),
2332+ min (int (batch ), row_blocks ),
22882333 current_cuda_stream (),
22892334 )
22902335
0 commit comments