Skip to content

Commit 7bb37d6

Browse files
voipmonitorclaude
andauthored
Partition paired DCP gathers across row blocks (#99)
* Partition paired DCP gathers across row blocks Co-Authored-By: Kimi <noreply@anthropic.com> * Validate paired gather row blocks without coercion The public query factory converted row_blocks to int before validation, allowing booleans, fractional values and strings to select a different grid. Pass the original value to the existing strict integer/range validator. Also correct the byte-fixture dimensions in the CUDA replay test so its exact comparison reaches the transport implementation. Invalid grid metadata now fails before compilation. Valid integer plans retain the same kernels and launch geometry; the byte-level reference and guard assertions are unchanged. Focused CPU coverage passes 100 tests with 13 CUDA compile skips. Paired gather CUDA tests pass on SM120 at TP4 and TP16, including queued changing rows, input/guard preservation and epoch rollover. For TP16 P8 with local BF16 [8,224] and FP32 [8,56], alternating graph microbenchmarks measured 33.3148 us with one row block and 31.0025 us with two. This measures the collective only and does not establish whole-model speedup. Signed-off-by: Martin Vit <martin@voipmonitor.org> --------- Signed-off-by: Martin Vit <martin@voipmonitor.org> Co-authored-by: Kimi <noreply@anthropic.com>
1 parent de9d64a commit 7bb37d6

7 files changed

Lines changed: 481 additions & 15 deletions

File tree

‎flashinfer/experimental/b12x/comm/pcie/_dcp_a2a_cute.py‎

Lines changed: 59 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -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

‎flashinfer/experimental/b12x/comm/pcie/_dcp_preparation.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,8 @@ def query_from_runtime(runtime, *, surface, call) -> PcieQuery:
135135
and (out is None or out.data_ptr() % 16 == 0),
136136
)
137137
elif surface.endswith("all_gather_pair"):
138+
from .pcie_dcp_a2a import _pair_outputs_disjoint
139+
138140
first, second = (
139141
_call_value(call, "local_first"),
140142
_call_value(call, "local_second"),
@@ -144,6 +146,10 @@ def query_from_runtime(runtime, *, surface, call) -> PcieQuery:
144146
second_dtype=str(second.dtype),
145147
first_shape=tuple(first.shape),
146148
second_shape=tuple(second.shape),
149+
row_blocks=call.get("row_blocks", 1),
150+
disjoint_output=_pair_outputs_disjoint(
151+
first, second, call.get("out_first"), call.get("out_second")
152+
),
147153
)
148154
elif surface.endswith("all_gather_pair_kimi_topk"):
149155
metadata["threads"] = 512
@@ -227,7 +233,12 @@ def compile_dcp_surface(query_payload, ordinal):
227233
if surface.endswith("all_gather_pair"):
228234
return {
229235
slot: cute._get_compiled_all_gather_pair(
230-
query.world_size, query.rank, int(call["threads"]), slot, False
236+
query.world_size,
237+
query.rank,
238+
int(call["threads"]),
239+
slot,
240+
False,
241+
int(call.get("row_blocks", 1)),
231242
)
232243
for slot in (False, True)
233244
}

‎flashinfer/experimental/b12x/comm/pcie/_tuning.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,12 @@ def _validate_query(query: PcieQuery, device) -> None:
140140
raise ValueError(
141141
"word exchange requires disjoint aligned data and owned TP2/TP4 state"
142142
)
143+
if query.surface.endswith(".all_gather_pair"):
144+
blocks = query.call.get("row_blocks", 1)
145+
if type(blocks) is not int or not 1 <= blocks <= 64:
146+
raise ValueError("paired gather row blocks must fit the signal plane")
147+
if blocks > 1 and not query.call.get("disjoint_output", False):
148+
raise ValueError("parallel paired gather requires disjoint outputs")
143149
if query.surface == "kimi_topk16":
144150
if query.topology != "local":
145151
raise ValueError("stateless kimi_topk16 uses one local device")

‎flashinfer/experimental/b12x/comm/pcie/pcie_dcp_a2a.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,28 @@ def lse_reduce_scatter_reference(
282282
return (outputs * weights.unsqueeze(-1)).sum(dim=0).to(partial_outputs.dtype)
283283

284284

285+
def _pair_outputs_disjoint(first, second, out_first, out_second):
286+
"""Check raw contiguous spans; omitted outputs are freshly allocated."""
287+
288+
def overlap(a, b):
289+
return (
290+
a is not None
291+
and b is not None
292+
and a.device == b.device
293+
and a.data_ptr() < b.data_ptr() + b.nbytes
294+
and b.data_ptr() < a.data_ptr() + a.nbytes
295+
)
296+
297+
return not (
298+
overlap(out_first, out_second)
299+
or any(
300+
overlap(source, out)
301+
for source in (first, second)
302+
for out in (out_first, out_second)
303+
)
304+
)
305+
306+
285307
class PCIeDCPA2A:
286308
"""One ordered IPC channel for DCP attention collectives."""
287309

@@ -1251,6 +1273,13 @@ def _all_gather_pair_on_device(
12511273
if not output.is_contiguous():
12521274
raise ValueError(f"{name} output must be contiguous")
12531275
threads, _ = self._resolve_launch_config(threads=threads, block_limit=1)
1276+
row_blocks = state.query.call.get("row_blocks", 1)
1277+
if row_blocks > 1 and not _pair_outputs_disjoint(
1278+
local_first, local_second, out_first, out_second
1279+
):
1280+
raise ValueError(
1281+
"parallel paired gather outputs must not overlap its inputs or each other"
1282+
)
12541283
capturing = _is_current_stream_capturing(self.device)
12551284
if capturing:
12561285
from ._dcp_a2a_cute import is_all_gather_pair_prepared
@@ -1261,6 +1290,7 @@ def _all_gather_pair_on_device(
12611290
threads,
12621291
True,
12631292
False,
1293+
row_blocks,
12641294
):
12651295
raise RuntimeError(
12661296
"cold PCIe DCP paired gather CUDA graph capture is not "

0 commit comments

Comments
 (0)