@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2785c72a14

[Kimi-K3] Extend GEMM-RS to GEMM-AR (#53053)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: wzhao18 <wzhao18.sz@gmail.com>

Thien Tran committed Aug 21, 2026 at 11:44 UTC 2785c72a1497dcd4db53d7c3bc557e379a518e0c
9 files changed +481 -232
.buildkite/test_areas/distributed.yaml
+11 -2
@@ -244,11 +244,20 @@ steps:
244 commands:
245 - pytest -v -s tests/distributed/test_context_parallel.py
246 - pytest -v -s tests/distributed/test_nccl_symm_mem.py
247 - - pytest -v -s tests/kernels/test_kimi_k3_gemm_rs.py
247 - pytest -v -s tests/v1/distributed/test_dbo.py
248 - pytest -v -s tests/distributed/test_mnnvl_alltoall.py
249
251 -
250 +- label: ":nvidia: (B200) Kimi-K3 GEMM-RS/AR"
251 + key: kimi-k3-gemm-rs-ar-2xb200
252 + timeout_in_minutes: 20
253 + device: b200-k8s
254 + working_dir: "/vllm-workspace/"
255 + num_devices: 2
256 + source_file_dependencies:
257 + - vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py
258 + - tests/kernels/test_kimi_k3_gemm_rs_ar.py
259 + commands:
260 + - pytest -v -s tests/kernels/test_kimi_k3_gemm_rs_ar.py
261
262 - label: ":nvidia: (L4) Distributed 2-Node"
263 key: 2-node-test-4-gpus
benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py renamed
+114 -78
@@ -1,11 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Benchmark the SM100 Kimi-K3 GEMM-RS kernel.
3 +"""Benchmark the SM100 Kimi-K3 GEMM-RS/AR kernel.
4
5 All ranks must belong to one NVLink domain. For example, run a TP8 sweep with:
6
7 torchrun --nproc-per-node=8 \
8 - benchmarks/kernels/benchmark_kimi_k3_gemm_rs.py
8 + benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py
9 """
10
11 import argparse
@@ -26,7 +26,7 @@ from vllm.distributed.parallel_state import (
26 init_distributed_environment,
27 initialize_model_parallel,
28 )
29 -from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS
29 +from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr
30
31 # Shared-expert down-proj and attention O-proj.
32 _KIMI_K3_PROJECTION_K = (6144, 12288)
@@ -41,6 +41,12 @@ class Candidate:
41
42 def parse_args() -> argparse.Namespace:
43 parser = argparse.ArgumentParser()
44 + parser.add_argument(
45 + "--mode",
46 + choices=("rs", "ar"),
47 + default="rs",
48 + help="Collective mode to benchmark.",
49 + )
50 parser.add_argument(
51 "--m",
52 type=int,
@@ -138,7 +144,8 @@ def valid_rows(M: int, local_M: int, rank: int) -> int:
144
145
146 def benchmark_shape(
141 - gemm_rs: GemmRS,
147 + gemm_rs_ar: GemmRsAr,
148 + mode: str,
149 M: int,
150 N: int,
151 K: int,
@@ -148,7 +155,8 @@ def benchmark_shape(
155 device_group: dist.ProcessGroup,
156 cpu_group: dist.ProcessGroup,
157 device_barrier: Callable[[], None],
151 -) -> dict[str, float | int]:
158 +) -> dict[str, float | int | str]:
159 + all_reduce = mode == "ar"
160 world_size = dist.get_world_size(device_group)
161 rank = dist.get_rank(device_group)
162 device = torch.device("cuda", torch.accelerator.current_device_index())
@@ -169,17 +177,19 @@ def benchmark_shape(
177 partial = torch.empty((padded_M, N), dtype=torch.bfloat16, device=device)
178 symm_partial = symm_mem.empty((padded_M, N), dtype=torch.bfloat16, device=device)
179 symm_partial_handle = symm_mem.rendezvous(symm_partial, device_group)
172 - rs_inputs = [torch.empty_like(partial) for _ in range(num_workspaces)]
173 - symm_rs_inputs = []
174 - symm_rs_handles = []
180 + collective_inputs = [torch.empty_like(partial) for _ in range(num_workspaces)]
181 + symm_collective_inputs = []
182 + symm_collective_handles = []
183 for _ in range(num_workspaces):
176 - rs_input = symm_mem.empty(
184 + collective_input = symm_mem.empty(
185 (padded_M, N),
186 dtype=torch.bfloat16,
187 device=device,
188 )
181 - symm_rs_inputs.append(rs_input)
182 - symm_rs_handles.append(symm_mem.rendezvous(rs_input, device_group))
189 + symm_collective_inputs.append(collective_input)
190 + symm_collective_handles.append(
191 + symm_mem.rendezvous(collective_input, device_group)
192 + )
193 torch_output = torch.empty((local_M, N), dtype=torch.bfloat16, device=device)
194 symm_output = torch.empty_like(torch_output)
195 gemm_output = torch.empty((M, N), dtype=torch.bfloat16, device=device)
@@ -188,21 +198,27 @@ def benchmark_shape(
198 partial[M:].zero_()
199 symm_partial[M:].zero_()
200
191 - def make_torch_ring_ll_gemm_rs(
201 + def make_torch_ring_ll_gemm_collective(
202 x: torch.Tensor, weight: torch.Tensor
203 ) -> Callable[[], torch.Tensor]:
204 def run() -> torch.Tensor:
205 torch.mm(x, weight.T, out=partial[:M])
206 + if all_reduce:
207 + dist.all_reduce(partial, group=device_group)
208 + return partial[:M]
209 dist.reduce_scatter_single(torch_output, partial, group=device_group)
210 return torch_output
211
212 return run
213
201 - def make_torch_ldmc_gemm_rs(
214 + def make_torch_ldmc_gemm_collective(
215 x: torch.Tensor, weight: torch.Tensor
216 ) -> Callable[[], torch.Tensor]:
217 def run() -> torch.Tensor:
218 torch.mm(x, weight.T, out=symm_partial[:M])
219 + if all_reduce:
220 + dist.all_reduce(symm_partial, group=device_group)
221 + return symm_partial[:M]
222 dist.reduce_scatter_single(
223 symm_output,
224 symm_partial,
@@ -212,11 +228,11 @@ def benchmark_shape(
228
229 return run
230
215 - def make_fused_gemm_rs(
231 + def make_fused_gemm_collective(
232 x: torch.Tensor, weight: torch.Tensor
233 ) -> Callable[[], torch.Tensor]:
234 def run() -> torch.Tensor:
219 - return gemm_rs(x, weight)
235 + return gemm_rs_ar(x, weight)
236
237 return run
238
@@ -229,18 +245,30 @@ def benchmark_shape(
245
246 return run
247
232 - def make_ring_ll_rs(rs_input: torch.Tensor) -> Callable[[], torch.Tensor]:
248 + def make_ring_ll_collective(
249 + collective_input: torch.Tensor,
250 + ) -> Callable[[], torch.Tensor]:
251 def run() -> torch.Tensor:
234 - dist.reduce_scatter_single(torch_output, rs_input, group=device_group)
252 + if all_reduce:
253 + dist.all_reduce(collective_input, group=device_group)
254 + return collective_input
255 + dist.reduce_scatter_single(
256 + torch_output, collective_input, group=device_group
257 + )
258 return torch_output
259
260 return run
261
239 - def make_ldmc_rs(
240 - rs_input: torch.Tensor,
262 + def make_ldmc_collective(
263 + collective_input: torch.Tensor,
264 ) -> Callable[[], torch.Tensor]:
265 def run() -> torch.Tensor:
243 - dist.reduce_scatter_single(symm_output, rs_input, group=device_group)
266 + if all_reduce:
267 + dist.all_reduce(collective_input, group=device_group)
268 + return collective_input
269 + dist.reduce_scatter_single(
270 + symm_output, collective_input, group=device_group
271 + )
272 return symm_output
273
274 return run
@@ -248,15 +276,15 @@ def benchmark_shape(
276 candidates = (
277 Candidate(
278 "ring_ll_us",
251 - [make_torch_ring_ll_gemm_rs(x, w) for x, w in zip(inputs, weights)],
279 + [make_torch_ring_ll_gemm_collective(x, w) for x, w in zip(inputs, weights)],
280 ),
281 Candidate(
282 "ldmc_us",
255 - [make_torch_ldmc_gemm_rs(x, w) for x, w in zip(inputs, weights)],
283 + [make_torch_ldmc_gemm_collective(x, w) for x, w in zip(inputs, weights)],
284 ),
285 Candidate(
258 - "gemm_rs_us",
259 - [make_fused_gemm_rs(x, w) for x, w in zip(inputs, weights)],
286 + "gemm_rs_ar_us",
287 + [make_fused_gemm_collective(x, w) for x, w in zip(inputs, weights)],
288 ),
289 Candidate(
290 "torch_gemm_us",
@@ -264,19 +292,19 @@ def benchmark_shape(
292 check_correctness=False,
293 ),
294 Candidate(
267 - "ring_ll_rs_us",
268 - [make_ring_ll_rs(x) for x in rs_inputs],
295 + "ring_ll_collective_us",
296 + [make_ring_ll_collective(x) for x in collective_inputs],
297 check_correctness=False,
298 ),
299 Candidate(
272 - "ldmc_rs_us",
273 - [make_ldmc_rs(x) for x in symm_rs_inputs],
300 + "ldmc_collective_us",
301 + [make_ldmc_collective(x) for x in symm_collective_inputs],
302 check_correctness=False,
303 ),
304 )
305
306 expected = candidates[0].runs[0]()
279 - rows = valid_rows(M, local_M, rank)
307 + rows = M if all_reduce else valid_rows(M, local_M, rank)
308 for candidate in candidates[1:]:
309 if not candidate.check_correctness:
310 continue
@@ -290,7 +318,7 @@ def benchmark_shape(
318 )
319
320 candidate_graphs = {}
293 - graph_keepalive: list[object] = [symm_partial_handle, *symm_rs_handles]
321 + graph_keepalive: list[object] = [symm_partial_handle, *symm_collective_handles]
322 for candidate in candidates:
323 stream = torch.cuda.Stream()
324 bundles = [capture_graph(run, stream, cpu_group) for run in candidate.runs]
@@ -306,70 +334,71 @@ def benchmark_shape(
334 device_barrier,
335 )
336
309 - best_nccl_rs_us = min(times["ring_ll_rs_us"], times["ldmc_rs_us"])
337 + best_nccl_us = min(times["ring_ll_collective_us"], times["ldmc_collective_us"])
338 return {
339 + "mode": mode.upper(),
340 "M": M,
341 "N": N,
342 "K": K,
343 **times,
315 - "best_nccl_rs_us": best_nccl_rs_us,
316 - "speedup_vs_ring_ll": times["ring_ll_us"] / times["gemm_rs_us"],
317 - "speedup_vs_ldmc": times["ldmc_us"] / times["gemm_rs_us"],
344 + "best_nccl_us": best_nccl_us,
345 + "speedup_vs_ring_ll": times["ring_ll_us"] / times["gemm_rs_ar_us"],
346 + "speedup_vs_ldmc": times["ldmc_us"] / times["gemm_rs_ar_us"],
347 }
348
349
321 -def print_results(results: list[dict[str, float | int]]) -> None:
350 +def print_results(results: list[dict[str, float | int | str]]) -> None:
351 results_df = pd.DataFrame(results)
323 - end_to_end = results_df[
324 - [
325 - "M",
326 - "N",
327 - "K",
328 - "ring_ll_us",
329 - "ldmc_us",
330 - "gemm_rs_us",
331 - "speedup_vs_ring_ll",
332 - "speedup_vs_ldmc",
352 + collective = str(results_df["mode"].iloc[0])
353 + end_to_end = (
354 + results_df[
355 + [
356 + "M",
357 + "N",
358 + "K",
359 + "ring_ll_us",
360 + "ldmc_us",
361 + "gemm_rs_ar_us",
362 + "speedup_vs_ring_ll",
363 + "speedup_vs_ldmc",
364 + ]
365 ]
334 - ].rename(
335 - columns={
336 - "ring_ll_us": "Torch GEMM + NCCL RS (RING_LL) (us)",
337 - "ldmc_us": "Torch GEMM + NCCL RS (LDMC) (us)",
338 - "gemm_rs_us": "GEMM-RS (us)",
339 - "speedup_vs_ring_ll": "Speedup vs RING_LL",
340 - "speedup_vs_ldmc": "Speedup vs LDMC",
341 - }
342 - )
343 - end_to_end = end_to_end.round(
344 - {
345 - "Torch GEMM + NCCL RS (RING_LL) (us)": 2,
346 - "Torch GEMM + NCCL RS (LDMC) (us)": 2,
347 - "GEMM-RS (us)": 2,
348 - "Speedup vs RING_LL": 3,
349 - "Speedup vs LDMC": 3,
350 - }
366 + .rename(
367 + columns={
368 + "ring_ll_us": f"Torch GEMM + NCCL {collective} (RING_LL) (us)",
369 + "ldmc_us": f"Torch GEMM + NCCL {collective} (LDMC) (us)",
370 + "gemm_rs_ar_us": f"GEMM-{collective} (us)",
371 + "speedup_vs_ring_ll": "Speedup vs RING_LL",
372 + "speedup_vs_ldmc": "Speedup vs LDMC",
373 + }
374 + )
375 + .round(3)
376 )
352 - components = results_df[
353 - ["M", "N", "K", "torch_gemm_us", "best_nccl_rs_us", "gemm_rs_us"]
354 - ].rename(
355 - columns={
356 - "torch_gemm_us": "Torch GEMM (us)",
357 - "best_nccl_rs_us": "NCCL RS (best) (us)",
358 - "gemm_rs_us": "GEMM-RS (us)",
359 - }
377 + components = (
378 + results_df[["M", "N", "K", "torch_gemm_us", "best_nccl_us", "gemm_rs_ar_us"]]
379 + .rename(
380 + columns={
381 + "torch_gemm_us": "Torch GEMM (us)",
382 + "best_nccl_us": f"NCCL {collective} (best) (us)",
383 + "gemm_rs_ar_us": f"GEMM-{collective} (us)",
384 + }
385 + )
386 + .round(2)
387 )
361 - components = components.round(2)
388
363 - print("### End-to-end latency")
389 + print(f"### {collective} end-to-end latency")
390 print(end_to_end.to_markdown(index=False))
365 - print("\n### Component latency")
391 + print(f"\n### {collective} component latency")
392 print(components.to_markdown(index=False))
367 - print("\nNCCL RS (best) is the faster of RING_LL and LDMC for each shape.")
393 + print(
394 + f"\nNCCL {collective} (best) is the faster of RING_LL and LDMC "
395 + "for each shape.\n"
396 + )
397
398
399 def main() -> None:
400 args = parse_args()
372 - assert args.m and min(args.m) >= 128
401 + assert args.m and min(args.m) > 0
402 assert args.n % 256 == 0
403 assert args.num_workspaces > 0
404 assert args.warmup_replays >= 0
@@ -401,12 +430,19 @@ def main() -> None:
430 sync_output = torch.empty_like(sync_input)
431
432 def device_barrier() -> None:
433 + # Order the timed launch after a device-side rank rendezvous without
434 + # including the rendezvous itself in the measured event interval.
435 pynccl_comm.all_reduce(sync_input, sync_output)
436
406 - gemm_rs = GemmRS(max_M=max(args.m), N=args.n)
437 + gemm_rs_ar = GemmRsAr(
438 + max_M=max(args.m),
439 + N=args.n,
440 + all_reduce=args.mode == "ar",
441 + )
442 results = [
443 benchmark_shape(
409 - gemm_rs,
444 + gemm_rs_ar,
445 + args.mode,
446 M,
447 args.n,
448 K,
@@ -420,12 +456,12 @@ def main() -> None:
456 for K in K_values
457 for M in args.m
458 ]
459 + del gemm_rs_ar
460
461 if tp_group.rank_in_group == 0:
462 print_results(results)
463
464 dist.barrier(group=tp_group.cpu_group)
428 - del gemm_rs
465 cleanup_dist_env_and_memory()
466
467
tests/kernels/test_kimi_k3_gemm_rs_ar.py renamed
+114 -47
@@ -31,6 +31,7 @@ def _reference(
31 weight: torch.Tensor,
32 world_size: int,
33 group: dist.ProcessGroup,
34 + all_reduce: bool,
35 ) -> torch.Tensor:
36 M = x.shape[0]
37 padded_M = (M + world_size - 1) // world_size * world_size
@@ -43,6 +44,10 @@ def _reference(
44 if padded_M > M:
45 partial[M:].zero_()
46
47 + if all_reduce:
48 + dist.all_reduce(partial, group=group)
49 + return partial[:M]
50 +
51 output = torch.empty(
52 (padded_M // world_size, weight.shape[0]),
53 dtype=x.dtype,
@@ -57,7 +62,11 @@ def _assert_valid_rows_close(
62 expected: torch.Tensor,
63 M: int,
64 rank: int,
65 + all_reduce: bool,
66 ) -> None:
67 + if all_reduce:
68 + torch.testing.assert_close(actual, expected, rtol=5e-2, atol=4.0)
69 + return
70 rank_start = rank * actual.shape[0]
71 valid_rows = min(max(M - rank_start, 0), actual.shape[0])
72 torch.testing.assert_close(
@@ -68,44 +77,24 @@ def _assert_valid_rows_close(
77 )
78
79
71 -def _worker(local_rank: int, world_size: int, master_port: int) -> None:
72 - # This module pulls in cute_dsl, which is unavailable off CUDA, so importing
73 - # it at module level would fail collection. Import it where it is used, as
74 - # `kda.py` does.
75 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS
80 +def _run_mode(
81 + *,
82 + all_reduce: bool,
83 + device: torch.device,
84 + group: dist.ProcessGroup,
85 + rank: int,
86 + world_size: int,
87 + weights: dict[int, torch.Tensor],
88 +) -> None:
89 + # cute_dsl is unavailable off CUDA, so import it only inside the GPU worker.
90 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr
91
77 - device = torch.device("cuda", local_rank)
78 - torch.accelerator.set_device_index(device)
79 - update_environment_variables(
80 - {
81 - "RANK": str(local_rank),
82 - "LOCAL_RANK": str(local_rank),
83 - "WORLD_SIZE": str(world_size),
84 - "MASTER_ADDR": "localhost",
85 - "MASTER_PORT": str(master_port),
86 - }
92 + gemm_rs_ar = GemmRsAr(
93 + max_M=max(M for M, _ in _SHAPES),
94 + N=_N,
95 + all_reduce=all_reduce,
96 )
88 -
89 - init_distributed_environment()
90 - with ensure_current_vllm_config():
91 - initialize_model_parallel(tensor_model_parallel_size=world_size)
92 -
93 - group = dist.group.WORLD
94 - rank = dist.get_rank(group)
95 - gemm_rs = GemmRS(max_M=max(M for M, _ in _SHAPES), N=_N)
96 -
97 - weight_generator = torch.Generator(device=device)
97 input_generator = torch.Generator(device=device)
99 - weights = {}
100 - for K in {K for _, K in _SHAPES}:
101 - weight_generator.manual_seed(1000 + rank * 10 + K)
102 - weights[K] = torch.randn(
103 - _N,
104 - K,
105 - dtype=torch.bfloat16,
106 - device=device,
107 - generator=weight_generator,
108 - )
98
99 # Alternating shapes exercise producer-flag reuse across different grids,
100 # CTA-group choices, and both BN=128 and BN=256 dispatches.
@@ -118,12 +107,29 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None:
107 device=device,
108 generator=input_generator,
109 )
121 - expected = _reference(x, weights[K], world_size, group)
122 - actual = gemm_rs(x, weights[K])
110 + expected = _reference(x, weights[K], world_size, group, all_reduce)
111 + actual = gemm_rs_ar(x, weights[K])
112 + torch.accelerator.synchronize(device)
113 + _assert_valid_rows_close(actual, expected, M, rank, all_reduce)
114 +
115 + if all_reduce:
116 + # AR outputs must remain valid after the symmetric workspace is reused.
117 + lifetime_M, lifetime_K = 257, 768
118 + input_generator.manual_seed(2500)
119 + lifetime_x = torch.randn(
120 + lifetime_M,
121 + lifetime_K,
122 + dtype=torch.bfloat16,
123 + device=device,
124 + generator=input_generator,
125 + )
126 + first_output = gemm_rs_ar(lifetime_x, weights[lifetime_K])
127 + first_snapshot = first_output.clone()
128 + second_output = gemm_rs_ar(-lifetime_x, weights[lifetime_K])
129 torch.accelerator.synchronize(device)
124 - _assert_valid_rows_close(actual, expected, M, rank)
130 + assert first_output.data_ptr() != second_output.data_ptr()
131 + torch.testing.assert_close(first_output, first_snapshot, rtol=0, atol=0)
132
126 - # Exercise GEMM-RS under CUDA graph capture and replay.
133 graph_M, graph_K = 1025, 4224
134 input_generator.manual_seed(3000)
135 graph_x = torch.randn(
@@ -133,42 +139,103 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None:
139 device=device,
140 generator=input_generator,
141 )
136 - graph_expected = _reference(graph_x, weights[graph_K], world_size, group)
142 + graph_expected = _reference(
143 + graph_x,
144 + weights[graph_K],
145 + world_size,
146 + group,
147 + all_reduce,
148 + )
149 graph_output = torch.empty_like(graph_expected)
150
151 capture_stream = torch.cuda.Stream()
152 capture_stream.wait_stream(torch.cuda.current_stream())
153 with torch.cuda.stream(capture_stream):
154 for _ in range(3):
143 - graph_output.copy_(gemm_rs(graph_x, weights[graph_K]))
155 + graph_output.copy_(gemm_rs_ar(graph_x, weights[graph_K]))
156 capture_stream.synchronize()
157 dist.barrier(group=group)
158
159 graph = torch.cuda.CUDAGraph()
160 with torch.cuda.graph(graph, stream=capture_stream):
149 - graph_output.copy_(gemm_rs(graph_x, weights[graph_K]))
161 + graph_output.copy_(gemm_rs_ar(graph_x, weights[graph_K]))
162 torch.cuda.current_stream().wait_stream(capture_stream)
163 dist.barrier(group=group)
164
165 for _ in range(3):
166 graph.replay()
167 torch.accelerator.synchronize(device)
156 - _assert_valid_rows_close(graph_output, graph_expected, graph_M, rank)
168 + _assert_valid_rows_close(
169 + graph_output,
170 + graph_expected,
171 + graph_M,
172 + rank,
173 + all_reduce,
174 + )
175
176 dist.barrier(group=group)
159 - del gemm_rs
177 + del gemm_rs_ar, capture_stream, graph, graph_output
178 +
179 +
180 +def _worker(local_rank: int, world_size: int, master_port: int) -> None:
181 + # This module pulls in cute_dsl, which is unavailable off CUDA, so importing
182 + # it at module level would fail collection. Import it where it is used, as
183 + # `kda.py` does.
184 +
185 + device = torch.device("cuda", local_rank)
186 + torch.accelerator.set_device_index(device)
187 + update_environment_variables(
188 + {
189 + "RANK": str(local_rank),
190 + "LOCAL_RANK": str(local_rank),
191 + "WORLD_SIZE": str(world_size),
192 + "MASTER_ADDR": "localhost",
193 + "MASTER_PORT": str(master_port),
194 + }
195 + )
196 +
197 + init_distributed_environment()
198 + with ensure_current_vllm_config():
199 + initialize_model_parallel(tensor_model_parallel_size=world_size)
200 +
201 + group = dist.group.WORLD
202 + rank = dist.get_rank(group)
203 +
204 + weight_generator = torch.Generator(device=device)
205 + weights = {}
206 + for K in {K for _, K in _SHAPES}:
207 + weight_generator.manual_seed(1000 + rank * 10 + K)
208 + weights[K] = torch.randn(
209 + _N,
210 + K,
211 + dtype=torch.bfloat16,
212 + device=device,
213 + generator=weight_generator,
214 + )
215 +
216 + # Production binds one mode per worker; exercise both mode-bound instances
217 + # sequentially in this test without implying that both are initialized.
218 + for all_reduce in (False, True):
219 + _run_mode(
220 + all_reduce=all_reduce,
221 + device=device,
222 + group=group,
223 + rank=rank,
224 + world_size=world_size,
225 + weights=weights,
226 + )
227 cleanup_dist_env_and_memory()
228
229
230 @pytest.mark.distributed(num_gpus=2)
231 @pytest.mark.skipif(
232 not current_platform.is_device_capability_family(100),
166 - reason="Kimi-K3 GEMM-RS requires SM100",
233 + reason="Kimi-K3 GEMM-RS/AR requires SM100",
234 )
168 -def test_kimi_k3_gemm_rs(monkeypatch: pytest.MonkeyPatch) -> None:
235 +def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch) -> None:
236 world_size = 2
237 if torch.accelerator.device_count() < world_size:
171 - pytest.skip("GEMM-RS requires two GPUs")
238 + pytest.skip("GEMM-RS/AR requires two GPUs")
239
240 monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
241 monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
vllm/envs.py
+4
@@ -210,6 +210,7 @@ if TYPE_CHECKING:
210 VLLM_MOE_SKIP_PADDING: bool = True
211 VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False
212 VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False
213 + VLLM_KIMI_K3_GEMM_AR: bool = True
214 VLLM_KIMI_K3_GEMM_RS: bool = False
215 VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
216 VLLM_USE_FLASHINFER_MOE_INT4: bool = False
@@ -1608,6 +1609,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1609 "VLLM_KIMI_K3_AUX_ATTN_RES_STREAM": lambda: bool(
1610 int(os.getenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "0"))
1611 ),
1612 + # Use the SM100 BF16 GEMM-AR kernel for eligible Kimi-K3 row-parallel
1613 + # attention projections. All TP ranks must belong to one NVLink domain.
1614 + "VLLM_KIMI_K3_GEMM_AR": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_AR", "1"))),
1615 # Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel
1616 # row-parallel projections. All TP ranks must belong to one NVLink domain.
1617 "VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))),
vllm/model_executor/warmup/kernel_warmup.py
+15
@@ -6,6 +6,7 @@ This is useful specifically for JIT'ed kernels as we don't want JIT'ing to
6 happen during model execution.
7 """
8
9 +import sys
10 import time
11 from typing import TYPE_CHECKING
12
@@ -95,6 +96,18 @@ def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None:
96 )
97
98
99 +def _warmup_kimi_k3_gemm_rs_ar() -> None:
100 + # Kimi-K3 model construction imports this module only when GEMM-RS/AR is
101 + # enabled and initializes its singleton before kernel_warmup runs. Avoid
102 + # importing it here so other models do not compile the RS/AR variants.
103 + module = sys.modules.get("vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar")
104 + if module is None:
105 + return
106 + compiled = module.warmup_gemm_rs_ar()
107 + if compiled:
108 + logger.info_once("Warmed up %d Kimi-K3 GEMM-RS/AR variants.", compiled)
109 +
110 +
111 def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
112 from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
113 minimax_m3_msa_warmup,
@@ -146,6 +159,8 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
159 if current_platform.has_device_capability(90):
160 _warmup_ll_bf16_router_gemm(worker.get_model())
161
162 + _warmup_kimi_k3_gemm_rs_ar()
163 +
164 if worker.vllm_config.kernel_config.enable_cutedsl_warmup:
165 # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
166 # to the shared JIT warmup infrastructure.
vllm/models/kimi_k3/nvidia/kda.py
+13 -13
@@ -319,7 +319,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
319 config: KimiLinearConfig,
320 vllm_config: VllmConfig,
321 prefix: str = "",
322 - run_gemm_rs: bool = False,
322 + run_gemm_rs_ar: bool = False,
323 ) -> None:
324 super().__init__(config, vllm_config, prefix)
325 self.use_recoverssm = self.cache_config.use_kda_recoverssm
@@ -491,14 +491,18 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
491 quant_config=self.quant_config,
492 prefix=f"{prefix}.o_proj",
493 )
494 - self.run_gemm_rs = run_gemm_rs
495 - if self.run_gemm_rs:
496 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
494 + self.gemm_rs_ar = None
495 + if run_gemm_rs_ar:
496 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
497 + get_gemm_rs_ar,
498 + )
499
498 - self.run_gemm_rs = get_gemm_rs().can_run(self.o_proj)
499 - if not self.run_gemm_rs:
500 + gemm_rs_ar = get_gemm_rs_ar()
501 + if gemm_rs_ar.can_run(self.o_proj):
502 + self.gemm_rs_ar = gemm_rs_ar
503 + else:
504 logger.warning_once(
501 - "GEMM-RS is disabled for %s due to an incompatible projection.",
505 + "GEMM-RS/AR is disabled for %s due to an incompatible projection.",
506 prefix,
507 )
508 compilation_config = vllm_config.compilation_config
@@ -541,12 +545,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
545 core_attn_out=core_attn_out,
546 )
547 core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
544 - if self.run_gemm_rs:
545 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
546 -
547 - gemm_rs = get_gemm_rs()
548 - if gemm_rs.should_run(core_attn_out):
549 - return gemm_rs(core_attn_out, self.o_proj.weight)
548 + if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(core_attn_out):
549 + return self.gemm_rs_ar(core_attn_out, self.o_proj.weight)
550 return self.o_proj(core_attn_out)[0]
551
552 @eager_break_during_capture
vllm/models/kimi_k3/nvidia/mla.py
+13 -13
@@ -134,7 +134,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
134 aux_stream: torch.cuda.Stream | None = None,
135 use_rope: bool = False,
136 non_causal_multi_token_decode: bool = False,
137 - run_gemm_rs: bool = False,
137 + run_gemm_rs_ar: bool = False,
138 ) -> None:
139 super().__init__()
140 self.hidden_size = hidden_size
@@ -271,14 +271,18 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
271 quant_config=quant_config,
272 prefix=f"{prefix}.o_proj",
273 )
274 - self.run_gemm_rs = run_gemm_rs
275 - if self.run_gemm_rs:
276 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
274 + self.gemm_rs_ar = None
275 + if run_gemm_rs_ar:
276 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
277 + get_gemm_rs_ar,
278 + )
279
278 - self.run_gemm_rs = get_gemm_rs().can_run(self.o_proj)
279 - if not self.run_gemm_rs:
280 + gemm_rs_ar = get_gemm_rs_ar()
281 + if gemm_rs_ar.can_run(self.o_proj):
282 + self.gemm_rs_ar = gemm_rs_ar
283 + else:
284 logger.warning_once(
281 - "GEMM-RS is disabled for %s due to an incompatible projection.",
285 + "GEMM-RS/AR is disabled for %s due to an incompatible projection.",
286 prefix,
287 )
288
@@ -549,12 +553,8 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
553 if gate is not None:
554 attn_out = _gate_sigmoid_mul(attn_out, gate)
555
552 - if self.run_gemm_rs:
553 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
554 -
555 - gemm_rs = get_gemm_rs()
556 - if gemm_rs.should_run(attn_out):
557 - return gemm_rs(attn_out, self.o_proj.weight)
556 + if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(attn_out):
557 + return self.gemm_rs_ar(attn_out, self.o_proj.weight)
558
559 return self.o_proj(attn_out)[0]
560
vllm/models/kimi_k3/nvidia/model.py
+60 -36
@@ -154,15 +154,18 @@ def shard_sequence_parallel_mlp(
154 )
155
156
157 -def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool:
158 - if not envs.VLLM_KIMI_K3_GEMM_RS:
157 +def maybe_init_gemm_rs_ar(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool:
158 + # Both feature flags may be enabled; the worker's static SP topology binds
159 + # its singleton to exactly one mode.
160 + all_reduce = not use_sequence_parallel
161 + mode = "GEMM-AR" if all_reduce else "GEMM-RS"
162 + enabled = envs.VLLM_KIMI_K3_GEMM_AR if all_reduce else envs.VLLM_KIMI_K3_GEMM_RS
163 + if not enabled:
164 return False
165
166 parallel_config = vllm_config.parallel_config
167 tp_size = parallel_config.tensor_parallel_size
163 - if not use_sequence_parallel:
164 - reason = "sequence parallelism is disabled"
165 - elif parallel_config.use_ubatching:
168 + if parallel_config.use_ubatching:
169 reason = "ubatching is enabled"
170 elif vllm_config.model_config.dtype != torch.bfloat16:
171 reason = "the model dtype is not BF16"
@@ -178,17 +181,32 @@ def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) ->
181 reason = None
182
183 if reason is not None:
181 - logger.warning_once("GEMM-RS was requested but is disabled because %s.", reason)
184 + logger.warning_once("%s is disabled because %s.", mode, reason)
185 return False
186
184 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import init_gemm_rs
187 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import init_gemm_rs_ar
188
189 config = vllm_config.model_config.hf_text_config
187 - init_gemm_rs(
188 - max_M=vllm_config.scheduler_config.max_num_batched_tokens,
189 - N=config.hidden_size,
190 - )
191 - logger.info_once("GEMM-RS is enabled.")
190 + try:
191 + init_gemm_rs_ar(
192 + max_M=vllm_config.scheduler_config.max_num_batched_tokens,
193 + N=config.hidden_size,
194 + all_reduce=all_reduce,
195 + )
196 + except RuntimeError as e:
197 + logger.warning_once(
198 + "%s is disabled because initialization failed: %s. This may mean "
199 + "the TP ranks do not share one NVLink domain.",
200 + mode,
201 + e,
202 + )
203 + return False
204 + if all_reduce:
205 + logger.info_once(
206 + "GEMM-AR is enabled. To disable it, set VLLM_KIMI_K3_GEMM_AR=0."
207 + )
208 + else:
209 + logger.info_once("GEMM-RS is enabled.")
210 return True
211
212
@@ -217,7 +235,7 @@ class KimiMLP(nn.Module):
235 reduce_results: bool = True,
236 use_sequence_parallel: bool = False,
237 can_shard_sequence_parallel: bool = False,
220 - run_gemm_rs: bool = False,
238 + run_gemm_rs_ar: bool = False,
239 prefix: str = "",
240 activation_situ_beta: float | None = None,
241 activation_situ_linear_beta: float | None = None,
@@ -230,7 +248,6 @@ class KimiMLP(nn.Module):
248 use_sequence_parallel,
249 can_shard_sequence_parallel,
250 )
233 - self.run_gemm_rs = self.shard_sequence_parallel and run_gemm_rs
251 replicate = use_sequence_parallel and not self.shard_sequence_parallel
252
253 self.gate_up_proj = MergedColumnParallelLinear(
@@ -252,13 +269,22 @@ class KimiMLP(nn.Module):
269 disable_tp=replicate,
270 prefix=f"{prefix}.down_proj",
271 )
255 - if self.run_gemm_rs:
256 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
272 + self.gemm_rs_ar = None
273 + # RS requires sequence sharding; AR operates on replicated tokens.
274 + use_gemm_rs_ar = self.shard_sequence_parallel or (
275 + not use_sequence_parallel and reduce_results
276 + )
277 + if use_gemm_rs_ar and run_gemm_rs_ar:
278 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
279 + get_gemm_rs_ar,
280 + )
281
258 - self.run_gemm_rs = get_gemm_rs().can_run(self.down_proj)
259 - if not self.run_gemm_rs:
282 + gemm_rs_ar = get_gemm_rs_ar()
283 + if gemm_rs_ar.can_run(self.down_proj):
284 + self.gemm_rs_ar = gemm_rs_ar
285 + else:
286 logger.warning_once(
261 - "GEMM-RS is disabled for %s due to an incompatible projection.",
287 + "GEMM-RS/AR is disabled for %s due to an incompatible projection.",
288 prefix,
289 )
290 if hidden_act == "silu":
@@ -284,12 +310,8 @@ class KimiMLP(nn.Module):
310 gate_up, _ = self.gate_up_proj(x)
311 x = self.act_fn(gate_up)
312
287 - if self.run_gemm_rs:
288 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
289 -
290 - gemm_rs = get_gemm_rs()
291 - if gemm_rs.should_run(x):
292 - return gemm_rs(x, self.down_proj.weight)
313 + if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(x):
314 + return self.gemm_rs_ar(x, self.down_proj.weight)
315
316 x, _ = self.down_proj(x)
317 if self.shard_sequence_parallel:
@@ -527,7 +549,7 @@ class KimiMoE(nn.Module):
549 prefix: str = "",
550 layer_idx: int = 0,
551 use_sequence_parallel: bool = False,
530 - run_gemm_rs: bool = False,
552 + run_gemm_rs_ar: bool = False,
553 ):
554 super().__init__()
555 hidden_size = config.hidden_size
@@ -616,7 +638,7 @@ class KimiMoE(nn.Module):
638 # FusedMoE path below hands them to the runner, which fuses
639 # their reduction and assumes the replicated layout.
640 can_shard_sequence_parallel=self.use_mega_moe,
619 - run_gemm_rs=run_gemm_rs,
641 + run_gemm_rs_ar=run_gemm_rs_ar,
642 prefix=f"{prefix}.shared_experts",
643 activation_situ_beta=activation_situ_beta,
644 activation_situ_linear_beta=activation_situ_linear_beta,
@@ -838,7 +860,7 @@ class KimiDecoderLayer(nn.Module):
860 vllm_config: VllmConfig,
861 prefix: str = "",
862 aux_stream: torch.cuda.Stream | None = None,
841 - run_gemm_rs: bool = False,
863 + run_gemm_rs_ar: bool = False,
864 ) -> None:
865 super().__init__()
866 self.hidden_size = config.hidden_size
@@ -873,7 +895,7 @@ class KimiDecoderLayer(nn.Module):
895 config,
896 vllm_config,
897 prefix=f"{prefix}.self_attn",
876 - run_gemm_rs=run_gemm_rs,
898 + run_gemm_rs_ar=run_gemm_rs_ar,
899 )
900 self._self_attn_writes_output = False
901 else:
@@ -910,7 +932,7 @@ class KimiDecoderLayer(nn.Module):
932 quant_config=quant_config,
933 prefix=f"{prefix}.self_attn",
934 aux_stream=aux_stream,
913 - run_gemm_rs=run_gemm_rs,
935 + run_gemm_rs_ar=run_gemm_rs_ar,
936 )
937 self._self_attn_writes_output = False
938
@@ -925,7 +947,7 @@ class KimiDecoderLayer(nn.Module):
947 prefix=f"{prefix}.block_sparse_moe",
948 layer_idx=layer_idx,
949 use_sequence_parallel=self.use_sequence_parallel,
928 - run_gemm_rs=run_gemm_rs,
950 + run_gemm_rs_ar=run_gemm_rs_ar,
951 )
952 self.mlp = self.block_sparse_moe
953 else:
@@ -937,7 +959,7 @@ class KimiDecoderLayer(nn.Module):
959 prefix=f"{prefix}.mlp",
960 use_sequence_parallel=self.use_sequence_parallel,
961 can_shard_sequence_parallel=True,
940 - run_gemm_rs=run_gemm_rs,
962 + run_gemm_rs_ar=run_gemm_rs_ar,
963 activation_situ_beta=config.activation_situ_beta,
964 activation_situ_linear_beta=config.activation_situ_linear_beta,
965 )
@@ -1118,9 +1140,11 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1140
1141 self.vocab_size = config.vocab_size
1142
1121 - # GEMM-RS uses NCCL symmetric-memory multicast, which requires all TP
1122 - # ranks to belong to one NVLink domain.
1123 - self.run_gemm_rs = maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel)
1143 + # GEMM-RS/AR uses NCCL symmetric-memory multicast, which requires all
1144 + # TP ranks to belong to one NVLink domain.
1145 + self.run_gemm_rs_ar = maybe_init_gemm_rs_ar(
1146 + vllm_config, self.use_sequence_parallel
1147 + )
1148
1149 if get_pp_group().is_first_rank:
1150 self.embed_tokens = VocabParallelEmbedding(
@@ -1142,7 +1166,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1166 vllm_config,
1167 prefix,
1168 aux_stream=aux_stream,
1145 - run_gemm_rs=self.run_gemm_rs,
1169 + run_gemm_rs_ar=self.run_gemm_rs_ar,
1170 )
1171
1172 self.start_layer, self.end_layer, self.layers = make_layers(
vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py renamed
+137 -43
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -"""SM100 BF16 GEMM with a fused tensor-parallel reduce-scatter."""
4 +"""SM100 BF16 GEMM with fused tensor-parallel reduce-scatter/all-reduce."""
5
6 # Based on CUTLASS's Blackwell distributed GEMM-RS example at dcf215a.
7 # See https://github.com/NVIDIA/cutlass/issues/3117 for memory semantics.
@@ -32,7 +32,7 @@ def nanosleep(ns: int, *, loc=None, ip=None) -> None:
32
33 @dsl_user_op
34 def multimem_ld_reduce_16B(x: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
35 - # NOTE: assume x is contiguous
35 + # The vector instruction assumes x is contiguous and 16-byte aligned.
36 assert x.element_type == BFloat16
37 vec_type = ".v4.bf16x2"
38
@@ -66,15 +66,40 @@ def multimem_ld_reduce_16B(x: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
66 return cute.recast_tensor(y, x.element_type)
67
68
69 -class Sm100GemmRsBF16:
69 +@dsl_user_op
70 +def multimem_st_16B(dst: cute.Tensor, value: cute.Tensor, *, loc=None, ip=None) -> None:
71 + # The vector instruction assumes dst and value are contiguous BF16x8 vectors
72 + # and dst is 16-byte aligned.
73 + assert dst.element_type == BFloat16
74 + assert value.element_type == BFloat16
75 + ptr = dst.iterator.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
76 + regs = cute.recast_tensor(value, Int32, loc=loc, ip=ip)
77 + llvm.inline_asm(
78 + None,
79 + [ptr] + [regs[i].ir_value(loc=loc, ip=ip) for i in range(4)],
80 + "multimem.st.relaxed.sys.global.v4.f32 [$0], {$1, $2, $3, $4};",
81 + "l,r,r,r,r",
82 + has_side_effects=True,
83 + loc=loc,
84 + ip=ip,
85 + )
86 +
87 +
88 +class Sm100GemmRsArBF16:
89 def __init__(
71 - self, rank: int, num_ranks: int, BN: int = 128, cta_group: int = 1
90 + self,
91 + rank: int,
92 + num_ranks: int,
93 + BN: int = 128,
94 + cta_group: int = 1,
95 + all_reduce: bool = False,
96 ) -> None:
97 self.rank = rank
98 self.num_ranks = num_ranks
99 BM, BK = 128, 64
100 self.cta_tile = (BM, BN, BK)
101 self.cta_group = cta_group
102 + self.all_reduce = all_reduce
103
104 smem_bytes = get_smem_capacity_in_bytes()
105 self.stage_size = (BM + (BN // cta_group)) * BK * 2
@@ -190,6 +215,9 @@ class Sm100GemmRsBF16:
215
216 M, K = A_tma.tma_tensor.shape
217 N, _ = B_tma.tma_tensor.shape
218 + local_M = (
219 + partial_uc.shape[0] // num_ranks if self.all_reduce else output.shape[0]
220 + )
221 grid_m = cute.ceil_div(M, BM)
222 # Keep 2-CTA clusters within a single N tile.
223 grid_m = cute.ceil_div(grid_m, cta_group) * cta_group
@@ -322,7 +350,6 @@ class Sm100GemmRsBF16:
350 # Keep epilogue in warps 0-3 and communication in warps 4-7.
351 # Swapping the warpgroups can hang due to warp scheduling.
352 tid_ = tid % 128
325 - local_M = output.shape[0]
353
354 # Offset in the [M, N] GEMM result.
355 rank_start = self.rank * local_M
@@ -338,7 +365,8 @@ class Sm100GemmRsBF16:
365 max_tile_rows = BM // num_ranks
366 vecs_per_tile = max_tile_rows * vec_cols
367 partial_vecs = cute.zipped_divide(partial_mc, (1, vec_width))
341 - output_vecs = cute.zipped_divide(output, (1, vec_width))
368 + if cutlass.const_expr(not self.all_reduce):
369 + output_vecs = cute.zipped_divide(output, (1, vec_width))
370
371 for tile_id in range(raw_bid, total_tiles, num_bids):
372 rs_bid_m = tile_id % grid_m
@@ -400,7 +428,6 @@ class Sm100GemmRsBF16:
428 reduced_vec.store(tmp.load())
429 reduced_vecs.append(reduced_vec)
430
403 - # Store the result to local L2.
431 for vec_iter in cutlass.range_constexpr(vecs_per_tile // 128):
432 vec_idx = tid_ + vec_iter * 128
433
@@ -409,11 +436,19 @@ class Sm100GemmRsBF16:
436 col = bid_n * vec_cols + vec_idx % vec_cols
437
438 if local_row < local_row_end and global_row < M:
412 - cute.copy(
413 - st_atom,
414 - reduced_vecs[vec_iter],
415 - output_vecs[None, (local_row, col)],
416 - )
439 + if cutlass.const_expr(self.all_reduce):
440 + # Broadcast the result to all ranks.
441 + multimem_st_16B(
442 + partial_vecs[None, (global_row, col)],
443 + reduced_vecs[vec_iter],
444 + )
445 + else:
446 + # Store the result to local L2.
447 + cute.copy(
448 + st_atom,
449 + reduced_vecs[vec_iter],
450 + output_vecs[None, (local_row, col)],
451 + )
452 cute.arch.barrier(barrier_id=BAR_COMM, number_of_threads=128)
453
454 # Release each GEMM tile consumed by this RS tile. The last
@@ -481,19 +516,21 @@ class Sm100GemmRsBF16:
516 num_ranks,
517 )
518
484 - # Exit barrier. GPU scope is sufficient because the kernel writes
485 - # local global memory and only needs to flush it to local L2.
519 + # Exit barrier. GPU scope is sufficient for RS because the kernel
520 + # writes local global memory and only needs to flush it to local L2.
521 + # AR uses system scope to flush the multimem stores to remote L2.
522 cute.arch.barrier(barrier_id=BAR_COMM, number_of_threads=128)
523 if tid_ == 0:
524 exit_flag = total_tiles + raw_bid
525 + scope = "sys" if self.all_reduce else "gpu"
526 utils.distributed.multimem_red_add1(
490 - flags_mc_ptr + exit_flag, order="release", scope="gpu"
527 + flags_mc_ptr + exit_flag, order="release", scope=scope
528 )
529 utils.distributed.spin_lock_atom_cas_acquire_wait(
530 flags_uc_ptr + exit_flag,
531 expected_val=num_ranks,
532 reset_val=0,
496 - scope="gpu",
533 + scope=scope,
534 )
535
536 else:
@@ -564,8 +601,8 @@ class Sm100GemmRsBF16:
601 if tid_ == 0 and bid_m * BM < M:
602 gemm_start = bid_m * BM
603 gemm_end = min(gemm_start + BM, M)
567 - first_owner = gemm_start // output.shape[0]
568 - last_owner = (gemm_end - 1) // output.shape[0]
604 + first_owner = gemm_start // local_M
605 + last_owner = (gemm_end - 1) // local_M
606
607 # signal to all consuming ranks (can be more than 1)
608 for rank in cutlass.range_constexpr(num_ranks):
@@ -595,7 +632,13 @@ class Sm100GemmRsBF16:
632
633 @cache
634 @staticmethod
598 - def compile(rank: int, num_ranks: int, BN: int, cta_group: int):
635 + def compile(
636 + rank: int,
637 + num_ranks: int,
638 + BN: int,
639 + cta_group: int,
640 + all_reduce: bool = False,
641 + ):
642 M = cute.sym_int()
643 padded_M = cute.sym_int()
644 N = cute.sym_int()
@@ -616,18 +659,24 @@ class Sm100GemmRsBF16:
659 assumed_align=32,
660 )
661 partial_mc_ptr = nullptr(BFloat16, cute.AddressSpace.gmem, assumed_align=32)
619 - output = make_fake_tensor(
620 - BFloat16,
621 - (local_M, N),
622 - (cute.sym_int64(divisibility=16), 1),
623 - assumed_align=32,
662 + output = (
663 + None
664 + if all_reduce
665 + else make_fake_tensor(
666 + BFloat16,
667 + (local_M, N),
668 + (cute.sym_int64(divisibility=16), 1),
669 + assumed_align=32,
670 + )
671 )
672 flags = make_fake_tensor(Int32, (num_flags,), (1,), assumed_align=16)
673 flags_mc_ptr = nullptr(Int32, cute.AddressSpace.gmem, assumed_align=16)
674 peer_flag_ptr = nullptr(Int64, cute.AddressSpace.gmem, assumed_align=8)
675
676 stream = make_fake_stream(use_tvm_ffi_env_stream=True)
630 - kernel = Sm100GemmRsBF16(rank, num_ranks, BN, cta_group)
677 + kernel = Sm100GemmRsArBF16(
678 + rank, num_ranks, BN, cta_group, all_reduce=all_reduce
679 + )
680 return cute.compile(
681 kernel,
682 A,
@@ -644,13 +693,20 @@ class Sm100GemmRsBF16:
693 )
694
695
647 -class GemmRS:
648 - """Own the max-sized symmetric workspace for Kimi-K3 GEMM-RS launches.
696 +class GemmRsAr:
697 + """Own the symmetric workspace for Kimi-K3 GEMM-RS/AR launches.
698
699 All TP ranks must belong to one NVLink domain for multimem instructions.
700 +
701 + Each instance is bound to either RS or AR. A vLLM worker has one static
702 + sequence-parallel topology, so the process-wide singleton only needs one
703 + mode. Two independent mode-specific singletons would lift that restriction
704 + but duplicate the large symmetric workspace. A future mixed-mode design
705 + should instead use lightweight RS/AR frontends over one shared multicast
706 + workspace; that is outside this integration's current scope.
707 """
708
653 - def __init__(self, *, max_M: int, N: int) -> None:
709 + def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None:
710 tp_group = get_tp_group()
711 group = tp_group.device_group
712 rank = tp_group.rank_in_group
@@ -667,9 +723,12 @@ class GemmRS:
723 self.max_M = max_M
724 self.N = N
725 self.device = device
726 + self.all_reduce = all_reduce
727
728 self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device)
729 self.partial_handle = symm_mem.rendezvous(self.partial, group)
730 + if self.partial_handle.multicast_ptr == 0:
731 + raise RuntimeError("GEMM-RS/AR requires NVLink multicast memory")
732 self.partial_mc_ptr = make_ptr(
733 BFloat16,
734 self.partial_handle.multicast_ptr,
@@ -684,6 +743,8 @@ class GemmRS:
743 max_flags = grid_m * (N // 128) + self.num_sms
744 self.flags = symm_mem.empty(max_flags, dtype=torch.int32, device=device)
745 self.flags_handle = symm_mem.rendezvous(self.flags, group)
746 + if self.flags_handle.multicast_ptr == 0:
747 + raise RuntimeError("GEMM-RS/AR requires NVLink multicast memory")
748 self.flags.zero_()
749 self.flags_mc_ptr = make_ptr(
750 Int32,
@@ -698,8 +759,6 @@ class GemmRS:
759 assumed_align=8,
760 )
761
701 - assert self.partial_handle.multicast_ptr != 0
702 - assert self.flags_handle.multicast_ptr != 0
762 torch.accelerator.synchronize(device)
763 tp_group.barrier()
764
@@ -721,7 +780,8 @@ class GemmRS:
780 )
781
782 def should_run(self, x: torch.Tensor) -> bool:
724 - # Small-M shapes are supported but faster on the existing LL path.
783 + # Use the same threshold for RS and AR for now. Small-M shapes are
784 + # supported but faster on the existing LL path.
785 return x.shape[0] >= 128
786
787 def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
@@ -752,12 +812,15 @@ class GemmRS:
812 num_ctas = num_ctas // cta_group * cta_group
813 assert self.flags.numel() >= num_tiles + num_ctas
814
755 - output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device)
756 - compiled = Sm100GemmRsBF16.compile(
815 + output = None
816 + if not self.all_reduce:
817 + output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device)
818 + compiled = Sm100GemmRsArBF16.compile(
819 self.rank,
820 self.world_size,
821 BN,
822 cta_group,
823 + self.all_reduce,
824 )
825 compiled(
826 x,
@@ -770,21 +833,52 @@ class GemmRS:
833 self.peer_flag_ptr,
834 num_ctas,
835 )
836 + if self.all_reduce:
837 + # AttnRes may retain output past the next workspace reuse.
838 + # A future kernel could overlap this copy using an extra warp or
839 + # the communication warp.
840 + return self.partial[:M].clone()
841 + assert output is not None
842 return output
843
844
776 -_gemm_rs: GemmRS | None = None
845 +_gemm_rs_ar: GemmRsAr | None = None
846
847
779 -def init_gemm_rs(max_M: int, N: int) -> None:
780 - """Collectively initialize the process-wide GEMM-RS state."""
781 - global _gemm_rs
782 - if _gemm_rs is not None:
783 - assert _gemm_rs.max_M >= max_M and _gemm_rs.N == N
848 +def init_gemm_rs_ar(max_M: int, N: int, *, all_reduce: bool = False) -> None:
849 + """Collectively initialize the process-wide, mode-bound GEMM-RS/AR state."""
850 + global _gemm_rs_ar
851 + if _gemm_rs_ar is not None:
852 + if _gemm_rs_ar.all_reduce != all_reduce:
853 + current = "AR" if _gemm_rs_ar.all_reduce else "RS"
854 + requested = "AR" if all_reduce else "RS"
855 + raise RuntimeError(
856 + f"GEMM-RS/AR is already initialized for {current}; "
857 + f"a worker cannot reinitialize it for {requested}"
858 + )
859 + assert _gemm_rs_ar.max_M >= max_M and _gemm_rs_ar.N == N
860 return
785 - _gemm_rs = GemmRS(max_M=max_M, N=N)
861 + _gemm_rs_ar = GemmRsAr(max_M=max_M, N=N, all_reduce=all_reduce)
862 +
863 +
864 +def warmup_gemm_rs_ar() -> int:
865 + """Compile every reachable dispatch for the initialized GEMM-RS/AR mode."""
866 + # Initialization can be disabled or fail when multicast is unavailable.
867 + if _gemm_rs_ar is None:
868 + return 0
869 + # Keep these profiles in sync with the dispatch in GemmRsAr.__call__.
870 + profiles = ((128, 1), (128, 2), (256, 2))
871 + for BN, cta_group in profiles:
872 + Sm100GemmRsArBF16.compile(
873 + _gemm_rs_ar.rank,
874 + _gemm_rs_ar.world_size,
875 + BN,
876 + cta_group,
877 + _gemm_rs_ar.all_reduce,
878 + )
879 + return len(profiles)
880
881
788 -def get_gemm_rs() -> GemmRS:
789 - assert _gemm_rs is not None, "GEMM-RS is not initialized"
790 - return _gemm_rs
882 +def get_gemm_rs_ar() -> GemmRsAr:
883 + assert _gemm_rs_ar is not None, "GEMM-RS/AR is not initialized"
884 + return _gemm_rs_ar