@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6014f9e672

[Kimi-K3] Add GEMM-RS for sequence parallelism (#52079)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Aug 14, 2026 at 00:02 UTC 6014f9e67291ea96548e263b8784a33a3f055954
11 files changed +1591 -13
.buildkite/test_areas/distributed.yaml
+1
@@ -242,6 +242,7 @@ steps:
242 commands:
243 - pytest -v -s tests/distributed/test_context_parallel.py
244 - pytest -v -s tests/distributed/test_nccl_symm_mem.py
245 + - pytest -v -s tests/kernels/test_kimi_k3_gemm_rs.py
246 - pytest -v -s tests/v1/distributed/test_dbo.py
247 - pytest -v -s tests/distributed/test_mnnvl_alltoall.py
248
benchmarks/kernels/benchmark_kimi_k3_gemm_rs.py new
+433
@@ -0,0 +1,433 @@
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.
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
9 +"""
10 +
11 +import argparse
12 +import os
13 +import statistics
14 +from collections.abc import Callable
15 +from dataclasses import dataclass
16 +
17 +import pandas as pd
18 +import torch
19 +import torch.distributed as dist
20 +import torch.distributed._symmetric_memory as symm_mem
21 +
22 +from vllm.config import VllmConfig, set_current_vllm_config
23 +from vllm.distributed import cleanup_dist_env_and_memory
24 +from vllm.distributed.parallel_state import (
25 + get_tp_group,
26 + init_distributed_environment,
27 + initialize_model_parallel,
28 +)
29 +from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS
30 +
31 +# Shared-expert down-proj and attention O-proj.
32 +_KIMI_K3_PROJECTION_K = (6144, 12288)
33 +
34 +
35 +@dataclass
36 +class Candidate:
37 + name: str
38 + runs: list[Callable[[], torch.Tensor]]
39 + check_correctness: bool = True
40 +
41 +
42 +def parse_args() -> argparse.Namespace:
43 + parser = argparse.ArgumentParser()
44 + parser.add_argument(
45 + "--m",
46 + type=int,
47 + nargs="+",
48 + default=[128, 512, 2048, 8192, 32768],
49 + help="Global token counts to benchmark.",
50 + )
51 + parser.add_argument(
52 + "--k",
53 + type=int,
54 + nargs="+",
55 + help=(
56 + "Per-rank input dimensions. By default, derive the Kimi-K3 "
57 + "shared-expert down-proj and O-proj dimensions from the TP "
58 + "world size."
59 + ),
60 + )
61 + parser.add_argument("--n", type=int, default=7168)
62 + parser.add_argument(
63 + "--num-workspaces",
64 + type=int,
65 + default=10,
66 + help="Pointer-distinct inputs and CUDA graphs to rotate.",
67 + )
68 + parser.add_argument("--warmup-replays", type=int, default=5)
69 + parser.add_argument("--samples", type=int, default=20)
70 + return parser.parse_args()
71 +
72 +
73 +def capture_graph(
74 + op: Callable[[], torch.Tensor],
75 + stream: torch.cuda.Stream,
76 + cpu_group: dist.ProcessGroup,
77 +) -> tuple[torch.cuda.CUDAGraph, list[torch.Tensor | None]]:
78 + result: list[torch.Tensor | None] = [None]
79 + stream.wait_stream(torch.cuda.current_stream())
80 + dist.barrier(group=cpu_group)
81 + with torch.cuda.stream(stream):
82 + for _ in range(3):
83 + result[0] = op()
84 + stream.synchronize()
85 + dist.barrier(group=cpu_group)
86 +
87 + graph = torch.cuda.CUDAGraph()
88 + with torch.cuda.graph(graph, stream=stream):
89 + result[0] = op()
90 + torch.cuda.current_stream().wait_stream(stream)
91 + dist.barrier(group=cpu_group)
92 + return graph, result
93 +
94 +
95 +def benchmark_graphs(
96 + candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]],
97 + warmup_replays: int,
98 + samples: int,
99 + device_group: dist.ProcessGroup,
100 + device_barrier: Callable[[], None],
101 +) -> dict[str, float]:
102 + candidate_names = list(candidate_graphs)
103 + for round_index in range(warmup_replays):
104 + for candidate_index in range(len(candidate_names)):
105 + candidate_id = (round_index + candidate_index) % len(candidate_names)
106 + name = candidate_names[candidate_id]
107 + graphs = candidate_graphs[name]
108 + device_barrier()
109 + graphs[round_index % len(graphs)].replay()
110 + torch.accelerator.synchronize()
111 +
112 + timings: dict[str, list[float]] = {name: [] for name in candidate_names}
113 + start = torch.cuda.Event(enable_timing=True)
114 + end = torch.cuda.Event(enable_timing=True)
115 + for sample_index in range(samples):
116 + for candidate_index in range(len(candidate_names)):
117 + candidate_id = (sample_index + candidate_index) % len(candidate_names)
118 + name = candidate_names[candidate_id]
119 + graphs = candidate_graphs[name]
120 + device_barrier()
121 + start.record()
122 + graphs[sample_index % len(graphs)].replay()
123 + end.record()
124 + end.synchronize()
125 +
126 + elapsed = torch.tensor(
127 + start.elapsed_time(end) * 1000,
128 + dtype=torch.float64,
129 + device=torch.accelerator.current_device_index(),
130 + )
131 + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
132 + timings[name].append(elapsed.item())
133 + return {name: statistics.median(values) for name, values in timings.items()}
134 +
135 +
136 +def valid_rows(M: int, local_M: int, rank: int) -> int:
137 + return min(max(M - rank * local_M, 0), local_M)
138 +
139 +
140 +def benchmark_shape(
141 + gemm_rs: GemmRS,
142 + M: int,
143 + N: int,
144 + K: int,
145 + num_workspaces: int,
146 + warmup_replays: int,
147 + samples: int,
148 + device_group: dist.ProcessGroup,
149 + cpu_group: dist.ProcessGroup,
150 + device_barrier: Callable[[], None],
151 +) -> dict[str, float | int]:
152 + world_size = dist.get_world_size(device_group)
153 + rank = dist.get_rank(device_group)
154 + device = torch.device("cuda", torch.accelerator.current_device_index())
155 + padded_M = (M + world_size - 1) // world_size * world_size
156 + local_M = padded_M // world_size
157 +
158 + rng = torch.Generator(device=device)
159 + rng.manual_seed(1000 + rank * 10 + M + K)
160 + inputs = [
161 + torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng)
162 + for _ in range(num_workspaces)
163 + ]
164 + weights = [
165 + torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng)
166 + for _ in range(num_workspaces)
167 + ]
168 +
169 + partial = torch.empty((padded_M, N), dtype=torch.bfloat16, device=device)
170 + symm_partial = symm_mem.empty((padded_M, N), dtype=torch.bfloat16, device=device)
171 + 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 = []
175 + for _ in range(num_workspaces):
176 + rs_input = symm_mem.empty(
177 + (padded_M, N),
178 + dtype=torch.bfloat16,
179 + device=device,
180 + )
181 + symm_rs_inputs.append(rs_input)
182 + symm_rs_handles.append(symm_mem.rendezvous(rs_input, device_group))
183 + torch_output = torch.empty((local_M, N), dtype=torch.bfloat16, device=device)
184 + symm_output = torch.empty_like(torch_output)
185 + gemm_output = torch.empty((M, N), dtype=torch.bfloat16, device=device)
186 +
187 + if padded_M > M:
188 + partial[M:].zero_()
189 + symm_partial[M:].zero_()
190 +
191 + def make_torch_ring_ll_gemm_rs(
192 + x: torch.Tensor, weight: torch.Tensor
193 + ) -> Callable[[], torch.Tensor]:
194 + def run() -> torch.Tensor:
195 + torch.mm(x, weight.T, out=partial[:M])
196 + dist.reduce_scatter_single(torch_output, partial, group=device_group)
197 + return torch_output
198 +
199 + return run
200 +
201 + def make_torch_ldmc_gemm_rs(
202 + x: torch.Tensor, weight: torch.Tensor
203 + ) -> Callable[[], torch.Tensor]:
204 + def run() -> torch.Tensor:
205 + torch.mm(x, weight.T, out=symm_partial[:M])
206 + dist.reduce_scatter_single(
207 + symm_output,
208 + symm_partial,
209 + group=device_group,
210 + )
211 + return symm_output
212 +
213 + return run
214 +
215 + def make_fused_gemm_rs(
216 + x: torch.Tensor, weight: torch.Tensor
217 + ) -> Callable[[], torch.Tensor]:
218 + def run() -> torch.Tensor:
219 + return gemm_rs(x, weight)
220 +
221 + return run
222 +
223 + def make_torch_gemm(
224 + x: torch.Tensor,
225 + weight: torch.Tensor,
226 + ) -> Callable[[], torch.Tensor]:
227 + def run() -> torch.Tensor:
228 + return torch.mm(x, weight.T, out=gemm_output)
229 +
230 + return run
231 +
232 + def make_ring_ll_rs(rs_input: torch.Tensor) -> Callable[[], torch.Tensor]:
233 + def run() -> torch.Tensor:
234 + dist.reduce_scatter_single(torch_output, rs_input, group=device_group)
235 + return torch_output
236 +
237 + return run
238 +
239 + def make_ldmc_rs(
240 + rs_input: torch.Tensor,
241 + ) -> Callable[[], torch.Tensor]:
242 + def run() -> torch.Tensor:
243 + dist.reduce_scatter_single(symm_output, rs_input, group=device_group)
244 + return symm_output
245 +
246 + return run
247 +
248 + candidates = (
249 + Candidate(
250 + "ring_ll_us",
251 + [make_torch_ring_ll_gemm_rs(x, w) for x, w in zip(inputs, weights)],
252 + ),
253 + Candidate(
254 + "ldmc_us",
255 + [make_torch_ldmc_gemm_rs(x, w) for x, w in zip(inputs, weights)],
256 + ),
257 + Candidate(
258 + "gemm_rs_us",
259 + [make_fused_gemm_rs(x, w) for x, w in zip(inputs, weights)],
260 + ),
261 + Candidate(
262 + "torch_gemm_us",
263 + [make_torch_gemm(x, w) for x, w in zip(inputs, weights)],
264 + check_correctness=False,
265 + ),
266 + Candidate(
267 + "ring_ll_rs_us",
268 + [make_ring_ll_rs(x) for x in rs_inputs],
269 + check_correctness=False,
270 + ),
271 + Candidate(
272 + "ldmc_rs_us",
273 + [make_ldmc_rs(x) for x in symm_rs_inputs],
274 + check_correctness=False,
275 + ),
276 + )
277 +
278 + expected = candidates[0].runs[0]()
279 + rows = valid_rows(M, local_M, rank)
280 + for candidate in candidates[1:]:
281 + if not candidate.check_correctness:
282 + continue
283 + actual = candidate.runs[0]()
284 + torch.accelerator.synchronize(device)
285 + torch.testing.assert_close(
286 + actual[:rows],
287 + expected[:rows],
288 + rtol=5e-2,
289 + atol=4.0,
290 + )
291 +
292 + candidate_graphs = {}
293 + graph_keepalive: list[object] = [symm_partial_handle, *symm_rs_handles]
294 + for candidate in candidates:
295 + stream = torch.cuda.Stream()
296 + bundles = [capture_graph(run, stream, cpu_group) for run in candidate.runs]
297 + candidate_graphs[candidate.name] = [graph for graph, _ in bundles]
298 + graph_keepalive.extend(bundles)
299 + graph_keepalive.append(stream)
300 +
301 + times = benchmark_graphs(
302 + candidate_graphs,
303 + warmup_replays,
304 + samples,
305 + device_group,
306 + device_barrier,
307 + )
308 +
309 + best_nccl_rs_us = min(times["ring_ll_rs_us"], times["ldmc_rs_us"])
310 + return {
311 + "M": M,
312 + "N": N,
313 + "K": K,
314 + **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"],
318 + }
319 +
320 +
321 +def print_results(results: list[dict[str, float | int]]) -> None:
322 + 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",
333 + ]
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 + }
351 + )
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 + }
360 + )
361 + components = components.round(2)
362 +
363 + print("### End-to-end latency")
364 + print(end_to_end.to_markdown(index=False))
365 + print("\n### Component latency")
366 + print(components.to_markdown(index=False))
367 + print("\nNCCL RS (best) is the faster of RING_LL and LDMC for each shape.")
368 +
369 +
370 +def main() -> None:
371 + args = parse_args()
372 + assert args.m and min(args.m) >= 128
373 + assert args.n % 256 == 0
374 + assert args.num_workspaces > 0
375 + assert args.warmup_replays >= 0
376 + assert args.samples > 0
377 +
378 + local_rank = int(os.environ["LOCAL_RANK"])
379 + torch.accelerator.set_device_index(local_rank)
380 + init_distributed_environment()
381 + world_size = dist.get_world_size()
382 + if args.k is None:
383 + assert all(K % world_size == 0 for K in _KIMI_K3_PROJECTION_K)
384 + K_values = [K // world_size for K in _KIMI_K3_PROJECTION_K]
385 + else:
386 + K_values = args.k
387 + assert all(K % 64 == 0 for K in K_values)
388 + # Reserve symmetric memory for the NCCL-managed benchmark allocations.
389 + os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
390 + # NCCL-managed symmetric allocations select the NVLS/LDMC collective path.
391 + symm_mem.set_backend("NCCL")
392 + with set_current_vllm_config(VllmConfig()):
393 + initialize_model_parallel(tensor_model_parallel_size=world_size)
394 +
395 + tp_group = get_tp_group()
396 + group_warmup = torch.zeros(1, device=torch.accelerator.current_device_index())
397 + dist.all_reduce(group_warmup, group=tp_group.device_group)
398 + pynccl_comm = tp_group.device_communicator.pynccl_comm
399 + assert pynccl_comm is not None
400 + sync_input = torch.zeros(1, device=torch.accelerator.current_device_index())
401 + sync_output = torch.empty_like(sync_input)
402 +
403 + def device_barrier() -> None:
404 + pynccl_comm.all_reduce(sync_input, sync_output)
405 +
406 + gemm_rs = GemmRS(max_M=max(args.m), N=args.n)
407 + results = [
408 + benchmark_shape(
409 + gemm_rs,
410 + M,
411 + args.n,
412 + K,
413 + args.num_workspaces,
414 + args.warmup_replays,
415 + args.samples,
416 + tp_group.device_group,
417 + tp_group.cpu_group,
418 + device_barrier,
419 + )
420 + for K in K_values
421 + for M in args.m
422 + ]
423 +
424 + if tp_group.rank_in_group == 0:
425 + print_results(results)
426 +
427 + dist.barrier(group=tp_group.cpu_group)
428 + del gemm_rs
429 + cleanup_dist_env_and_memory()
430 +
431 +
432 +if __name__ == "__main__":
433 + main()
tests/kernels/test_kimi_k3_gemm_rs.py new
+178
@@ -0,0 +1,178 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +import torch.distributed as dist
7 +import torch.multiprocessing as mp
8 +
9 +from tests.utils import ensure_current_vllm_config
10 +from vllm.distributed import cleanup_dist_env_and_memory
11 +from vllm.distributed.parallel_state import (
12 + init_distributed_environment,
13 + initialize_model_parallel,
14 +)
15 +from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS
16 +from vllm.platforms import current_platform
17 +from vllm.utils.network_utils import get_open_port
18 +from vllm.utils.system_utils import update_environment_variables
19 +
20 +_SHAPES = (
21 + (129, 768),
22 + (257, 768),
23 + (1023, 4224),
24 + (1024, 4224),
25 + (8191, 4224),
26 +)
27 +_N = 512
28 +
29 +
30 +def _reference(
31 + x: torch.Tensor,
32 + weight: torch.Tensor,
33 + world_size: int,
34 + group: dist.ProcessGroup,
35 +) -> torch.Tensor:
36 + M = x.shape[0]
37 + padded_M = (M + world_size - 1) // world_size * world_size
38 + partial = torch.empty(
39 + (padded_M, weight.shape[0]),
40 + dtype=x.dtype,
41 + device=x.device,
42 + )
43 + torch.mm(x, weight.T, out=partial[:M])
44 + if padded_M > M:
45 + partial[M:].zero_()
46 +
47 + output = torch.empty(
48 + (padded_M // world_size, weight.shape[0]),
49 + dtype=x.dtype,
50 + device=x.device,
51 + )
52 + dist.reduce_scatter_single(output, partial, group=group)
53 + return output
54 +
55 +
56 +def _assert_valid_rows_close(
57 + actual: torch.Tensor,
58 + expected: torch.Tensor,
59 + M: int,
60 + rank: int,
61 +) -> None:
62 + rank_start = rank * actual.shape[0]
63 + valid_rows = min(max(M - rank_start, 0), actual.shape[0])
64 + torch.testing.assert_close(
65 + actual[:valid_rows],
66 + expected[:valid_rows],
67 + rtol=5e-2,
68 + atol=4.0,
69 + )
70 +
71 +
72 +def _worker(local_rank: int, world_size: int, master_port: int) -> None:
73 + device = torch.device("cuda", local_rank)
74 + torch.accelerator.set_device_index(device)
75 + update_environment_variables(
76 + {
77 + "RANK": str(local_rank),
78 + "LOCAL_RANK": str(local_rank),
79 + "WORLD_SIZE": str(world_size),
80 + "MASTER_ADDR": "localhost",
81 + "MASTER_PORT": str(master_port),
82 + }
83 + )
84 +
85 + init_distributed_environment()
86 + with ensure_current_vllm_config():
87 + initialize_model_parallel(tensor_model_parallel_size=world_size)
88 +
89 + group = dist.group.WORLD
90 + rank = dist.get_rank(group)
91 + gemm_rs = GemmRS(max_M=max(M for M, _ in _SHAPES), N=_N)
92 +
93 + weight_generator = torch.Generator(device=device)
94 + input_generator = torch.Generator(device=device)
95 + weights = {}
96 + for K in {K for _, K in _SHAPES}:
97 + weight_generator.manual_seed(1000 + rank * 10 + K)
98 + weights[K] = torch.randn(
99 + _N,
100 + K,
101 + dtype=torch.bfloat16,
102 + device=device,
103 + generator=weight_generator,
104 + )
105 +
106 + # Alternating shapes exercise producer-flag reuse across different grids,
107 + # CTA-group choices, and both BN=128 and BN=256 dispatches.
108 + for M, K in (*_SHAPES, *_SHAPES[::-1]):
109 + input_generator.manual_seed(2000 + M + K)
110 + x = torch.randn(
111 + M,
112 + K,
113 + dtype=torch.bfloat16,
114 + device=device,
115 + generator=input_generator,
116 + )
117 + expected = _reference(x, weights[K], world_size, group)
118 + actual = gemm_rs(x, weights[K])
119 + torch.accelerator.synchronize(device)
120 + _assert_valid_rows_close(actual, expected, M, rank)
121 +
122 + # Exercise GEMM-RS under CUDA graph capture and replay.
123 + graph_M, graph_K = 1025, 4224
124 + input_generator.manual_seed(3000)
125 + graph_x = torch.randn(
126 + graph_M,
127 + graph_K,
128 + dtype=torch.bfloat16,
129 + device=device,
130 + generator=input_generator,
131 + )
132 + graph_expected = _reference(graph_x, weights[graph_K], world_size, group)
133 + graph_output = torch.empty_like(graph_expected)
134 +
135 + capture_stream = torch.cuda.Stream()
136 + capture_stream.wait_stream(torch.cuda.current_stream())
137 + with torch.cuda.stream(capture_stream):
138 + for _ in range(3):
139 + graph_output.copy_(gemm_rs(graph_x, weights[graph_K]))
140 + capture_stream.synchronize()
141 + dist.barrier(group=group)
142 +
143 + graph = torch.cuda.CUDAGraph()
144 + with torch.cuda.graph(graph, stream=capture_stream):
145 + graph_output.copy_(gemm_rs(graph_x, weights[graph_K]))
146 + torch.cuda.current_stream().wait_stream(capture_stream)
147 + dist.barrier(group=group)
148 +
149 + for _ in range(3):
150 + graph.replay()
151 + torch.accelerator.synchronize(device)
152 + _assert_valid_rows_close(graph_output, graph_expected, graph_M, rank)
153 +
154 + dist.barrier(group=group)
155 + del gemm_rs
156 + cleanup_dist_env_and_memory()
157 +
158 +
159 +@pytest.mark.distributed(num_gpus=2)
160 +@pytest.mark.skipif(
161 + not current_platform.is_device_capability_family(100),
162 + reason="Kimi-K3 GEMM-RS requires SM100",
163 +)
164 +def test_kimi_k3_gemm_rs(monkeypatch: pytest.MonkeyPatch) -> None:
165 + world_size = 2
166 + if torch.accelerator.device_count() < world_size:
167 + pytest.skip("GEMM-RS requires two GPUs")
168 +
169 + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
170 + monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
171 + try:
172 + mp.spawn(
173 + _worker,
174 + args=(world_size, get_open_port()),
175 + nprocs=world_size,
176 + )
177 + finally:
178 + cleanup_dist_env_and_memory()
tests/models/kimi_k3/test_sequence_parallel.py
+9 -6
@@ -267,19 +267,21 @@ def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch):
267
268
269 @pytest.mark.parametrize(
270 - ("enabled", "use_sequence_parallel", "tp_size", "expected"),
270 + ("enabled", "use_sequence_parallel", "eligible", "tp_size", "expected"),
271 [
272 - (True, True, 8, True),
273 - (False, True, 8, False), # opt-in only
274 - (True, False, 8, False), # replication only exists under SP
275 - (True, True, 1, False), # nothing to shard
276 - (True, True, 5, False), # 6144 % 5 -- would fail divide()
272 + (True, True, True, 8, True),
273 + (False, True, True, 8, False), # opt-in only
274 + (True, False, True, 8, False), # replication only exists under SP
275 + (True, True, False, 8, False), # FusedMoE path owns the reduction
276 + (True, True, True, 1, False), # nothing to shard
277 + (True, True, True, 5, False), # 6144 % 5 -- would fail divide()
278 ],
279 )
280 def test_shard_sequence_parallel_mlp_gating(
281 monkeypatch,
282 enabled: bool,
283 use_sequence_parallel: bool,
284 + eligible: bool,
285 tp_size: int,
286 expected: bool,
287 ):
@@ -293,6 +295,7 @@ def test_shard_sequence_parallel_mlp_gating(
295 hidden_size=7168,
296 intermediate_size=6144,
297 use_sequence_parallel=use_sequence_parallel,
298 + eligible=eligible,
299 )
300 is expected
301 )
vllm/cute_utils/__init__.py
+12
@@ -40,6 +40,18 @@ def recast_val(x, dtype, *, loc=None, ip=None):
40 return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
41
42
43 +@dsl_user_op
44 +def to_cta0_smem(ptr: cute.Pointer, *, loc=None, ip=None):
45 + return cute.make_ptr(
46 + ptr.dtype,
47 + ptr.toint(loc=loc, ip=ip) & 0xFEFF_FFFF,
48 + cute.AddressSpace.smem,
49 + assumed_align=8,
50 + loc=loc,
51 + ip=ip,
52 + )
53 +
54 +
55 def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None):
56 """A simple helper that wraps group_modes() and tma_partition()
57 NOTE: this should be called WITHOUT cute.elect_one()
vllm/cute_utils/mbarrier.py new
+47
@@ -0,0 +1,47 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from cutlass import Int32
5 +from cutlass._mlir.dialects import nvvm
6 +from cutlass.cutlass_dsl import dsl_user_op
7 +
8 +_SPACE_MAP = {
9 + "cta": nvvm.MBarrierSpaceKind.CTA,
10 + "cluster": nvvm.MBarrierSpaceKind.CLUSTER,
11 +}
12 +
13 +_MEMORY_ORDER_MAP = {
14 + "weak": nvvm.MemOrderKind.WEAK,
15 + "relaxed": nvvm.MemOrderKind.RELAXED,
16 + "acquire": nvvm.MemOrderKind.ACQUIRE,
17 + "release": nvvm.MemOrderKind.RELEASE,
18 + "acq_rel": nvvm.MemOrderKind.ACQ_REL,
19 +}
20 +
21 +
22 +@dsl_user_op
23 +def arrive(mbar, space: str = "cta", order: str = "relaxed", *, loc=None, ip=None):
24 + nvvm.mbarrier_txn(
25 + mbar.to_llvm_ptr(loc=loc, ip=ip),
26 + Int32(1).ir_value(loc=loc, ip=ip),
27 + kind=nvvm.MBarrierTxnKind.ARRIVE,
28 + space=_SPACE_MAP[space],
29 + order=_MEMORY_ORDER_MAP[order],
30 + loc=loc,
31 + ip=ip,
32 + )
33 +
34 +
35 +@dsl_user_op
36 +def arrive_expect_tx(
37 + mbar, size, space: str = "cta", order: str = "relaxed", *, loc=None, ip=None
38 +):
39 + nvvm.mbarrier_txn(
40 + mbar.to_llvm_ptr(loc=loc, ip=ip),
41 + Int32(size).ir_value(loc=loc, ip=ip),
42 + kind=nvvm.MBarrierTxnKind.ARRIVE_EXPECT_TX,
43 + space=_SPACE_MAP[space],
44 + order=_MEMORY_ORDER_MAP[order],
45 + loc=loc,
46 + ip=ip,
47 + )
vllm/envs.py
+4
@@ -206,6 +206,7 @@ if TYPE_CHECKING:
206 VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True
207 VLLM_MOE_SKIP_PADDING: bool = True
208 VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False
209 + VLLM_KIMI_K3_GEMM_RS: bool = False
210 VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
211 VLLM_USE_FLASHINFER_MOE_INT4: bool = False
212 VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None
@@ -1594,6 +1595,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1595 "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool(
1596 int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0"))
1597 ),
1598 + # Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel
1599 + # row-parallel projections. All TP ranks must belong to one NVLink domain.
1600 + "VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))),
1601 # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.
1602 # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).
1603 "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(
vllm/models/kimi_k3/nvidia/kda.py
+17 -1
@@ -305,6 +305,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
305 config: KimiLinearConfig,
306 vllm_config: VllmConfig,
307 prefix: str = "",
308 + run_gemm_rs: bool = False,
309 ) -> None:
310 super().__init__(config, vllm_config, prefix)
311
@@ -470,7 +471,16 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
471 quant_config=self.quant_config,
472 prefix=f"{prefix}.o_proj",
473 )
473 -
474 + self.run_gemm_rs = run_gemm_rs
475 + if self.run_gemm_rs:
476 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
477 +
478 + self.run_gemm_rs = get_gemm_rs().can_run(self.o_proj)
479 + if not self.run_gemm_rs:
480 + logger.warning_once(
481 + "GEMM-RS is disabled for %s due to an incompatible projection.",
482 + prefix,
483 + )
484 compilation_config = vllm_config.compilation_config
485 if prefix in compilation_config.static_forward_context:
486 raise ValueError(f"Duplicate layer name: {prefix}")
@@ -511,6 +521,12 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
521 core_attn_out=core_attn_out,
522 )
523 core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
524 + if self.run_gemm_rs:
525 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
526 +
527 + gemm_rs = get_gemm_rs()
528 + if gemm_rs.should_run(core_attn_out):
529 + return gemm_rs(core_attn_out, self.o_proj.weight)
530 return self.o_proj(core_attn_out)[0]
531
532 @eager_break_during_capture
vllm/models/kimi_k3/nvidia/mla.py
+18 -4
@@ -134,6 +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,
138 ) -> None:
139 super().__init__()
140 self.hidden_size = hidden_size
@@ -270,6 +271,16 @@ 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
277 +
278 + self.run_gemm_rs = get_gemm_rs().can_run(self.o_proj)
279 + if not self.run_gemm_rs:
280 + logger.warning_once(
281 + "GEMM-RS is disabled for %s due to an incompatible projection.",
282 + prefix,
283 + )
284
285 # ---- Attention backend / impl / KV cache ----
286 self.quant_config = quant_config
@@ -541,10 +552,13 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
552 if gate is not None:
553 attn_out = _gate_sigmoid_mul(attn_out, gate)
554
544 - # ``o_proj`` (RowParallelLinear + out-of-place all-reduce) returns a
545 - # fresh private tensor, so return it directly rather than copying into a
546 - # caller buffer -- the previous ``output[:] = ...`` convention forced an
547 - # extra [num_tokens, hidden] copy per layer.
555 + if self.run_gemm_rs:
556 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
557 +
558 + gemm_rs = get_gemm_rs()
559 + if gemm_rs.should_run(attn_out):
560 + return gemm_rs(attn_out, self.o_proj.weight)
561 +
562 return self.o_proj(attn_out)[0]
563
564 @eager_break_during_capture
vllm/models/kimi_k3/nvidia/model.py
+82 -2
@@ -133,6 +133,7 @@ def shard_sequence_parallel_mlp(
133 hidden_size: int,
134 intermediate_size: int,
135 use_sequence_parallel: bool,
136 + eligible: bool,
137 ) -> bool:
138 """Whether to TP-shard a sequence-parallel MLP instead of replicating it.
139
@@ -140,7 +141,7 @@ def shard_sequence_parallel_mlp(
141 the trade-off and :mod:`vllm.envs` for when it is worth enabling.
142 """
143 enabled = envs.VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT
143 - if not (use_sequence_parallel and enabled):
144 + if not (use_sequence_parallel and eligible and enabled):
145 return False
146 tp_size = get_tensor_model_parallel_world_size()
147 return (
@@ -148,6 +149,44 @@ def shard_sequence_parallel_mlp(
149 )
150
151
152 +def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool:
153 + if not envs.VLLM_KIMI_K3_GEMM_RS:
154 + return False
155 +
156 + parallel_config = vllm_config.parallel_config
157 + tp_size = parallel_config.tensor_parallel_size
158 + if not use_sequence_parallel:
159 + reason = "sequence parallelism is disabled"
160 + elif parallel_config.use_ubatching:
161 + reason = "ubatching is enabled"
162 + elif vllm_config.model_config.dtype != torch.bfloat16:
163 + reason = "the model dtype is not BF16"
164 + elif not current_platform.is_cuda():
165 + reason = "the device is not CUDA"
166 + elif not current_platform.is_device_capability_family(100):
167 + reason = "the device is not SM100-family"
168 + elif not 1 < tp_size <= 16:
169 + reason = "TP size is not in the supported range 2-16"
170 + elif 128 % tp_size != 0:
171 + reason = "TP size does not divide 128"
172 + else:
173 + reason = None
174 +
175 + if reason is not None:
176 + logger.warning_once("GEMM-RS was requested but is disabled because %s.", reason)
177 + return False
178 +
179 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import init_gemm_rs
180 +
181 + config = vllm_config.model_config.hf_text_config
182 + init_gemm_rs(
183 + max_M=vllm_config.scheduler_config.max_num_batched_tokens,
184 + N=config.hidden_size,
185 + )
186 + logger.info_once("GEMM-RS is enabled.")
187 + return True
188 +
189 +
190 class KimiMLP(nn.Module):
191 """Dense / shared-expert MLP, optionally TP-sharded under sequence parallel.
192
@@ -172,6 +211,8 @@ class KimiMLP(nn.Module):
211 quant_config: QuantizationConfig | None = None,
212 reduce_results: bool = True,
213 use_sequence_parallel: bool = False,
214 + can_shard_sequence_parallel: bool = False,
215 + run_gemm_rs: bool = False,
216 prefix: str = "",
217 activation_situ_beta: float | None = None,
218 activation_situ_linear_beta: float | None = None,
@@ -182,7 +223,9 @@ class KimiMLP(nn.Module):
223 hidden_size,
224 intermediate_size,
225 use_sequence_parallel,
226 + can_shard_sequence_parallel,
227 )
228 + self.run_gemm_rs = self.shard_sequence_parallel and run_gemm_rs
229 replicate = use_sequence_parallel and not self.shard_sequence_parallel
230
231 self.gate_up_proj = MergedColumnParallelLinear(
@@ -204,6 +247,15 @@ class KimiMLP(nn.Module):
247 disable_tp=replicate,
248 prefix=f"{prefix}.down_proj",
249 )
250 + if self.run_gemm_rs:
251 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
252 +
253 + self.run_gemm_rs = get_gemm_rs().can_run(self.down_proj)
254 + if not self.run_gemm_rs:
255 + logger.warning_once(
256 + "GEMM-RS is disabled for %s due to an incompatible projection.",
257 + prefix,
258 + )
259 if hidden_act == "silu":
260 self.act_fn = SiluAndMul()
261 elif hidden_act == "situ":
@@ -226,6 +278,14 @@ class KimiMLP(nn.Module):
278 x = sp_all_gather(x)
279 gate_up, _ = self.gate_up_proj(x)
280 x = self.act_fn(gate_up)
281 +
282 + if self.run_gemm_rs:
283 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs
284 +
285 + gemm_rs = get_gemm_rs()
286 + if gemm_rs.should_run(x):
287 + return gemm_rs(x, self.down_proj.weight)
288 +
289 x, _ = self.down_proj(x)
290 if self.shard_sequence_parallel:
291 x = sp_reduce_scatter(x)
@@ -462,6 +522,7 @@ class KimiMoE(nn.Module):
522 prefix: str = "",
523 layer_idx: int = 0,
524 use_sequence_parallel: bool = False,
525 + run_gemm_rs: bool = False,
526 ):
527 super().__init__()
528 hidden_size = config.hidden_size
@@ -546,6 +607,11 @@ class KimiMoE(nn.Module):
607 quant_config=quant_config,
608 reduce_results=False,
609 use_sequence_parallel=use_sequence_parallel,
610 + # Only the MegaMoE path calls the shared experts directly; the
611 + # FusedMoE path below hands them to the runner, which fuses
612 + # their reduction and assumes the replicated layout.
613 + can_shard_sequence_parallel=self.use_mega_moe,
614 + run_gemm_rs=run_gemm_rs,
615 prefix=f"{prefix}.shared_experts",
616 activation_situ_beta=activation_situ_beta,
617 activation_situ_linear_beta=activation_situ_linear_beta,
@@ -767,6 +833,7 @@ class KimiDecoderLayer(nn.Module):
833 vllm_config: VllmConfig,
834 prefix: str = "",
835 aux_stream: torch.cuda.Stream | None = None,
836 + run_gemm_rs: bool = False,
837 ) -> None:
838 super().__init__()
839 self.hidden_size = config.hidden_size
@@ -801,6 +868,7 @@ class KimiDecoderLayer(nn.Module):
868 config,
869 vllm_config,
870 prefix=f"{prefix}.self_attn",
871 + run_gemm_rs=run_gemm_rs,
872 )
873 self._self_attn_writes_output = False
874 else:
@@ -837,6 +905,7 @@ class KimiDecoderLayer(nn.Module):
905 quant_config=quant_config,
906 prefix=f"{prefix}.self_attn",
907 aux_stream=aux_stream,
908 + run_gemm_rs=run_gemm_rs,
909 )
910 self._self_attn_writes_output = False
911
@@ -851,6 +920,7 @@ class KimiDecoderLayer(nn.Module):
920 prefix=f"{prefix}.block_sparse_moe",
921 layer_idx=layer_idx,
922 use_sequence_parallel=self.use_sequence_parallel,
923 + run_gemm_rs=run_gemm_rs,
924 )
925 self.mlp = self.block_sparse_moe
926 else:
@@ -861,6 +931,8 @@ class KimiDecoderLayer(nn.Module):
931 quant_config=quant_config,
932 prefix=f"{prefix}.mlp",
933 use_sequence_parallel=self.use_sequence_parallel,
934 + can_shard_sequence_parallel=True,
935 + run_gemm_rs=run_gemm_rs,
936 activation_situ_beta=config.activation_situ_beta,
937 activation_situ_linear_beta=config.activation_situ_linear_beta,
938 )
@@ -991,15 +1063,18 @@ class KimiDecoderLayer(nn.Module):
1063 )
1064 assert hidden_states is not None
1065
1066 + M = None
1067 if self.use_sequence_parallel:
1068 hidden_states = sp_all_gather(hidden_states)
1069 # Remove SP padding before attention.
1070 hidden_states = hidden_states[: positions.shape[0]]
1071 + M = hidden_states.shape[0]
1072
1073 # Attention.
1074 hidden_states = self._run_self_attn(positions, hidden_states)
1075
1002 - if self.use_sequence_parallel:
1076 + # GEMM-RS returns the local sequence shard; standard O-proj preserves M.
1077 + if self.use_sequence_parallel and hidden_states.shape[0] == M:
1078 # Add SP padding if needed, and then perform reduce scatter.
1079 hidden_states = sp_reduce_scatter(hidden_states)
1080
@@ -1038,6 +1113,10 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1113
1114 self.vocab_size = config.vocab_size
1115
1116 + # GEMM-RS uses NCCL symmetric-memory multicast, which requires all TP
1117 + # ranks to belong to one NVLink domain.
1118 + self.run_gemm_rs = maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel)
1119 +
1120 if get_pp_group().is_first_rank:
1121 self.embed_tokens = VocabParallelEmbedding(
1122 config.vocab_size,
@@ -1058,6 +1137,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1137 vllm_config,
1138 prefix,
1139 aux_stream=aux_stream,
1140 + run_gemm_rs=self.run_gemm_rs,
1141 )
1142
1143 self.start_layer, self.end_layer, self.layers = make_layers(
vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs.py new
+790
@@ -0,0 +1,790 @@
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."""
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.
8 +
9 +from functools import cache
10 +
11 +import cutlass
12 +import torch
13 +import torch.distributed._symmetric_memory as symm_mem
14 +from cuda.bindings.driver import CUstream
15 +from cutlass import BFloat16, Int32, Int64, Uint16, cute, utils
16 +from cutlass._mlir import ir
17 +from cutlass._mlir.dialects import llvm, nvvm, vector
18 +from cutlass.cute.nvgpu import cpasync, tcgen05
19 +from cutlass.cute.runtime import make_fake_stream, make_fake_tensor, make_ptr, nullptr
20 +from cutlass.cutlass_dsl import dsl_user_op
21 +from cutlass.utils import get_smem_capacity_in_bytes
22 +
23 +from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem
24 +from vllm.distributed import get_tp_group
25 +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
26 +
27 +
28 +@dsl_user_op
29 +def nanosleep(ns: int, *, loc=None, ip=None) -> None:
30 + nvvm.nanosleep(Int32(ns).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
31 +
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
36 + assert x.element_type == BFloat16
37 + vec_type = ".v4.bf16x2"
38 +
39 + ptr = x.iterator.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
40 + asm = (
41 + "multimem.ld_reduce.relaxed.gpu.global.add.acc::f32"
42 + f"{vec_type} {{$0, $1, $2, $3}}, [$4];"
43 + )
44 + struct = llvm.inline_asm(
45 + llvm.StructType.get_literal([Int32.mlir_type] * 4),
46 + [ptr],
47 + asm,
48 + "=r,=r,=r,=r,l",
49 + has_side_effects=True,
50 + loc=loc,
51 + ip=ip,
52 + )
53 + vec = vector.from_elements(
54 + ir.VectorType.get([4], Int32.mlir_type, loc=loc),
55 + [
56 + llvm.extractvalue(Int32.mlir_type, struct, [i], loc=loc, ip=ip)
57 + for i in range(4)
58 + ],
59 + loc=loc,
60 + ip=ip,
61 + )
62 + ssa = cute.TensorSSA(vec, 4, Int32)
63 +
64 + y = cute.make_rmem_tensor(4, Int32)
65 + y.store(ssa)
66 + return cute.recast_tensor(y, x.element_type)
67 +
68 +
69 +class Sm100GemmRsBF16:
70 + def __init__(
71 + self, rank: int, num_ranks: int, BN: int = 128, cta_group: int = 1
72 + ) -> None:
73 + self.rank = rank
74 + self.num_ranks = num_ranks
75 + BM, BK = 128, 64
76 + self.cta_tile = (BM, BN, BK)
77 + self.cta_group = cta_group
78 +
79 + smem_bytes = get_smem_capacity_in_bytes()
80 + self.stage_size = (BM + (BN // cta_group)) * BK * 2
81 + self.num_stages = smem_bytes // self.stage_size
82 +
83 + @cute.jit
84 + def prepare_tma(
85 + self, tensor: cute.Tensor, BM: cutlass.Constexpr, BK: cutlass.Constexpr
86 + ) -> cpasync.TmaInfo:
87 + tma_group = (
88 + tcgen05.CtaGroup.TWO if self.cta_group == 2 else tcgen05.CtaGroup.ONE
89 + )
90 + tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=tma_group)
91 + swizzle_128b = cute.make_swizzle(3, 4, 3)
92 + layout = cute.make_layout(
93 + (BM, BK, self.num_stages),
94 + stride=(BK, 1, BM * BK),
95 + )
96 + layout = cute.make_composed_layout(swizzle_128b, 0, layout)
97 + return cpasync.make_tiled_tma_atom(tma_op, tensor, layout, (BM, BK))
98 +
99 + @cute.jit
100 + def __call__(
101 + self,
102 + A: cute.Tensor,
103 + B: cute.Tensor,
104 + partial_uc: cute.Tensor,
105 + partial_mc_ptr: cute.Pointer,
106 + output: cute.Tensor,
107 + flags_uc: cute.Tensor,
108 + flags_mc_ptr: cute.Pointer,
109 + peer_flag_ptr: cute.Pointer,
110 + grid_size: Int32,
111 + stream: CUstream,
112 + ) -> None:
113 + N = B.shape[0]
114 + BM, BN, BK = self.cta_tile
115 + A_tma = self.prepare_tma(A, BM, BK)
116 + B_tma = self.prepare_tma(B, BN // self.cta_group, BK)
117 + padded_M = partial_uc.shape[0]
118 + partial_mc = cute.make_tensor(
119 + partial_mc_ptr,
120 + cute.make_layout((padded_M, N), stride=(N, 1)),
121 + )
122 + peer_flags = cute.make_tensor(
123 + peer_flag_ptr,
124 + cute.make_layout(self.num_ranks),
125 + )
126 +
127 + grid = (grid_size, 1, 1)
128 + block = (10 * 32, 1, 1)
129 + cluster = (self.cta_group, 1, 1)
130 + self.kernel(
131 + A_tma,
132 + B_tma,
133 + partial_uc,
134 + partial_mc,
135 + output,
136 + flags_uc.iterator,
137 + flags_mc_ptr,
138 + peer_flags,
139 + ).launch(grid=grid, block=block, cluster=cluster, stream=stream)
140 +
141 + @cute.kernel
142 + def kernel(
143 + self,
144 + A_tma: cpasync.TmaInfo,
145 + B_tma: cpasync.TmaInfo,
146 + partial_uc: cute.Tensor,
147 + partial_mc: cute.Tensor,
148 + output: cute.Tensor,
149 + flags_uc_ptr: cute.Pointer,
150 + flags_mc_ptr: cute.Pointer,
151 + peer_flags: cute.Tensor,
152 + ) -> None:
153 + tid, _, _ = cute.arch.thread_idx()
154 + raw_bid, _, _ = cute.arch.block_idx()
155 + num_bids, _, _ = cute.arch.grid_dim()
156 + warp_id = cute.arch.make_warp_uniform(tid // 32)
157 +
158 + BM, BN, BK = self.cta_tile
159 + cta_group = self.cta_group
160 + num_stages = self.num_stages
161 + num_ranks = self.num_ranks
162 +
163 + is_2cta = cta_group == 2
164 + cta_rank = raw_bid % self.cta_group
165 + num_tmem_stages = 512 // BN
166 +
167 + smem = utils.SmemAllocator()
168 + sA = smem.allocate_tensor(
169 + BFloat16,
170 + A_tma.smem_layout.outer,
171 + byte_alignment=128,
172 + swizzle=A_tma.smem_layout.inner,
173 + )
174 + sB = smem.allocate_tensor(
175 + BFloat16,
176 + B_tma.smem_layout.outer,
177 + byte_alignment=128,
178 + swizzle=B_tma.smem_layout.inner,
179 + )
180 + tma_full_mbar = smem.allocate_array(Int64, num_stages)
181 + tma_empty_mbar = smem.allocate_array(Int64, num_stages)
182 + tmem_full_mbar = smem.allocate_array(Int64, num_tmem_stages)
183 + tmem_empty_mbar = smem.allocate_array(Int64, num_tmem_stages)
184 + taddr = smem.allocate(Int32, 4)
185 +
186 + # Named barriers
187 + BAR_TMEM_ALLOC = 1
188 + BAR_EPILOGUE = 2
189 + BAR_COMM = 3
190 +
191 + M, K = A_tma.tma_tensor.shape
192 + N, _ = B_tma.tma_tensor.shape
193 + grid_m = cute.ceil_div(M, BM)
194 + # Keep 2-CTA clusters within a single N tile.
195 + grid_m = cute.ceil_div(grid_m, cta_group) * cta_group
196 + grid_n = cute.ceil_div(N, BN)
197 +
198 + if warp_id == 0:
199 + with cute.arch.elect_one():
200 + for i in cutlass.range_constexpr(num_stages):
201 + cute.arch.mbarrier_init(tma_full_mbar + i, cta_group)
202 + cute.arch.mbarrier_init(tma_empty_mbar + i, 1)
203 + for i in cutlass.range_constexpr(num_tmem_stages):
204 + cute.arch.mbarrier_init(tmem_full_mbar + i, 1)
205 + cute.arch.mbarrier_init(tmem_empty_mbar + i, 128 * cta_group)
206 + cute.arch.mbarrier_init_fence()
207 + elif warp_id == 1:
208 + cpasync.prefetch_descriptor(A_tma.atom)
209 + cpasync.prefetch_descriptor(B_tma.atom)
210 +
211 + if cutlass.const_expr(is_2cta):
212 + cute.arch.cluster_arrive_relaxed()
213 + cute.arch.cluster_wait()
214 + else:
215 + cute.arch.sync_threads()
216 +
217 + total_tiles = grid_m * grid_n
218 +
219 + if warp_id == 9:
220 + # TMA warp
221 + tma_stage = 0
222 + parity = 1
223 +
224 + if cutlass.const_expr(is_2cta):
225 + tma_full_mbar_ = to_cta0_smem(tma_full_mbar)
226 + else:
227 + tma_full_mbar_ = tma_full_mbar
228 +
229 + # Select global-memory tiles.
230 + # [(BM, BK), (M/BM, K/BK)]
231 + gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK))
232 + gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN // cta_group, BK))
233 +
234 + for bid in range(raw_bid, total_tiles, num_bids):
235 + bid_m = bid % grid_m
236 + bid_n = bid // grid_m
237 + if cutlass.const_expr(cta_group == 2):
238 + bid_n = bid_n * cta_group + cta_rank
239 +
240 + for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1):
241 + mbar = tma_full_mbar_ + tma_stage
242 + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, parity)
243 +
244 + with cute.arch.elect_one():
245 + mbarrier.arrive_expect_tx(mbar, self.stage_size, "cluster")
246 + simple_tma_copy(
247 + A_tma.atom,
248 + gA_tiles[None, (bid_m, iter_k)],
249 + sA[None, None, tma_stage],
250 + mbar,
251 + )
252 + simple_tma_copy(
253 + B_tma.atom,
254 + gB_tiles[None, (bid_n, iter_k)],
255 + sB[None, None, tma_stage],
256 + mbar,
257 + )
258 +
259 + tma_stage = (tma_stage + 1) % num_stages
260 + if tma_stage == 0:
261 + parity ^= 1
262 +
263 + elif warp_id == 8:
264 + # MMA warp
265 + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=5 * 32)
266 +
267 + if cta_rank == 0:
268 + tma_stage = 0
269 + tma_full_parity = 0
270 + tmem_stage = 0
271 + tmem_empty_parity = 1
272 +
273 + MMA_M = BM * cta_group
274 + MMA_N = BN
275 + idesc = _tcgen05.make_bf16_idesc(MMA_M, MMA_N)
276 + sdesc = _tcgen05.make_sdesc_128B_swizzle(0)
277 + multicast_mask = Uint16((1 << self.cta_group) - 1)
278 +
279 + for bid in range(raw_bid, total_tiles, num_bids):
280 + cute.arch.mbarrier_wait(
281 + tmem_empty_mbar + tmem_stage, tmem_empty_parity
282 + )
283 + _tcgen05.fence_after_thread_sync()
284 +
285 + for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1):
286 + d_tmem = BN * tmem_stage
287 + a_addr = sA[None, None, tma_stage].iterator.toint()
288 + b_addr = sB[None, None, tma_stage].iterator.toint()
289 + a_desc = sdesc | (a_addr >> 4)
290 + b_desc = sdesc | (b_addr >> 4)
291 +
292 + cute.arch.mbarrier_wait(
293 + tma_full_mbar + tma_stage, tma_full_parity
294 + )
295 + _tcgen05.fence_after_thread_sync()
296 +
297 + for mma_k in cutlass.range_constexpr(BK // 16):
298 + enable_d = iter_k > 0 or mma_k > 0
299 + _tcgen05.mma_f16(
300 + d_tmem, a_desc, b_desc, idesc, enable_d, cta_group
301 + )
302 + a_desc += 32 >> 4
303 + b_desc += 32 >> 4
304 + _tcgen05.commit(
305 + tma_empty_mbar + tma_stage, multicast_mask, cta_group
306 + )
307 +
308 + tma_stage = (tma_stage + 1) % num_stages
309 + if tma_stage == 0:
310 + tma_full_parity ^= 1
311 +
312 + _tcgen05.commit(
313 + tmem_full_mbar + tmem_stage, multicast_mask, cta_group
314 + )
315 +
316 + tmem_stage = (tmem_stage + 1) % num_tmem_stages
317 + if tmem_stage == 0:
318 + tmem_empty_parity ^= 1
319 +
320 + elif warp_id >= 4:
321 + # Communication warps
322 + # Keep epilogue in warps 0-3 and communication in warps 4-7.
323 + # Swapping the warpgroups can hang due to warp scheduling.
324 + tid_ = tid % 128
325 + local_M = output.shape[0]
326 +
327 + # Offset in the [M, N] GEMM result.
328 + rank_start = self.rank * local_M
329 + rank_end = min(rank_start + local_M, M)
330 +
331 + st_atom = cute.make_copy_atom(
332 + cute.nvgpu.CopyUniversalOp(), BFloat16, num_bits_per_copy=128
333 + )
334 +
335 + # Each thread issues one BF16x8 multimem reduction per vector.
336 + vec_width = 8
337 + vec_cols = BN // vec_width
338 + max_tile_rows = BM // num_ranks
339 + vecs_per_tile = max_tile_rows * vec_cols
340 + partial_vecs = cute.zipped_divide(partial_mc, (1, vec_width))
341 + output_vecs = cute.zipped_divide(output, (1, vec_width))
342 +
343 + for tile_id in range(raw_bid, total_tiles, num_bids):
344 + rs_bid_m = tile_id % grid_m
345 + bid_n = tile_id // grid_m
346 +
347 + local_row_start = rs_bid_m * local_M // grid_m
348 + local_row_end = (rs_bid_m + 1) * local_M // grid_m
349 + global_row_start = rank_start + local_row_start
350 + global_row_end = min(rank_start + local_row_end, M)
351 +
352 + if global_row_start < M:
353 + # Map the current RS tile to the required GEMM tiles.
354 + # Since an RS tile is smaller than a GEMM tile, each RS tile
355 + # can overlap at most 2 GEMM tiles.
356 + gemm_bid_m0 = global_row_start // BM
357 + gemm_bid_m1 = max(global_row_end - 1, global_row_start) // BM
358 +
359 + # Poll local L2 with relaxed GPU-scope loads. The following
360 + # multimem reduction reads every rank's L2 directly, so no
361 + # acquire fence or L1 invalidation is needed.
362 + if tid_ == 0:
363 + # Poll the 1st GEMM tile.
364 + flag_ptr = flags_uc_ptr + bid_n * grid_m + gemm_bid_m0
365 + arrivals = cute.arch.load(
366 + flag_ptr, Int32, sem="relaxed", scope="gpu"
367 + )
368 + while arrivals < num_ranks:
369 + nanosleep(64)
370 + arrivals = cute.arch.load(
371 + flag_ptr, Int32, sem="relaxed", scope="gpu"
372 + )
373 + elif tid_ == 32 and gemm_bid_m1 != gemm_bid_m0:
374 + # Poll the 2nd GEMM tile with another warp if needed.
375 + flag_ptr = flags_uc_ptr + bid_n * grid_m + gemm_bid_m1
376 + arrivals = cute.arch.load(
377 + flag_ptr, Int32, sem="relaxed", scope="gpu"
378 + )
379 + while arrivals < num_ranks:
380 + nanosleep(64)
381 + arrivals = cute.arch.load(
382 + flag_ptr, Int32, sem="relaxed", scope="gpu"
383 + )
384 + cute.arch.barrier(barrier_id=BAR_COMM, number_of_threads=128)
385 +
386 + # Issue multimem.ld_reduce before storing any result.
387 + reduced_vecs = []
388 + for vec_iter in cutlass.range_constexpr(vecs_per_tile // 128):
389 + vec_idx = tid_ + vec_iter * 128
390 +
391 + local_row = local_row_start + vec_idx // vec_cols
392 + global_row = rank_start + local_row
393 + col = bid_n * vec_cols + vec_idx % vec_cols
394 +
395 + reduced_vec = cute.make_rmem_tensor(vec_width, BFloat16)
396 + if local_row < local_row_end and global_row < M:
397 + tmp = multimem_ld_reduce_16B(
398 + partial_vecs[None, (global_row, col)]
399 + )
400 + reduced_vec.store(tmp.load())
401 + reduced_vecs.append(reduced_vec)
402 +
403 + # Store the result to local L2.
404 + for vec_iter in cutlass.range_constexpr(vecs_per_tile // 128):
405 + vec_idx = tid_ + vec_iter * 128
406 +
407 + local_row = local_row_start + vec_idx // vec_cols
408 + global_row = rank_start + local_row
409 + col = bid_n * vec_cols + vec_idx % vec_cols
410 +
411 + 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 + )
417 + cute.arch.barrier(barrier_id=BAR_COMM, number_of_threads=128)
418 +
419 + # Release each GEMM tile consumed by this RS tile. The last
420 + # consumer resets its producer flag for the next launch.
421 + def release_gemm_tile(
422 + gemm_bid_m,
423 + tile_M,
424 + logical_M,
425 + rank_start,
426 + rank_end,
427 + local_M,
428 + grid_m,
429 + tile_N,
430 + flags,
431 + num_ranks,
432 + ):
433 + gemm_start = gemm_bid_m * tile_M
434 + gemm_end = min(gemm_start + tile_M, logical_M)
435 + local_start = max(gemm_start, rank_start) - rank_start
436 + local_end = min(gemm_end, rank_end) - rank_start
437 +
438 + # Compute the number of consumers to identify the last
439 + # arrival, which resets the producer flag.
440 + first_rs_tile = (
441 + cute.ceil_div((local_start + 1) * grid_m, local_M) - 1
442 + )
443 + last_rs_tile = cute.ceil_div(local_end * grid_m, local_M) - 1
444 + num_consumers = last_rs_tile - first_rs_tile + 1
445 +
446 + flag_ptr = flags + tile_N * grid_m + gemm_bid_m
447 + old_count = cute.arch.atomic_add(
448 + flag_ptr, Int32(1), sem="relaxed", scope="gpu"
449 + )
450 + if old_count == num_ranks + num_consumers - 1:
451 + cute.arch.store(
452 + flag_ptr, Int32(0), sem="relaxed", scope="gpu"
453 + )
454 +
455 + if tid_ == 0:
456 + # Arrive at the 1st GEMM tile.
457 + release_gemm_tile(
458 + gemm_bid_m0,
459 + BM,
460 + M,
461 + rank_start,
462 + rank_end,
463 + local_M,
464 + grid_m,
465 + bid_n,
466 + flags_uc_ptr,
467 + num_ranks,
468 + )
469 + elif tid_ == 32 and gemm_bid_m1 != gemm_bid_m0:
470 + # Arrive at the 2nd GEMM tile if necessary.
471 + release_gemm_tile(
472 + gemm_bid_m1,
473 + BM,
474 + M,
475 + rank_start,
476 + rank_end,
477 + local_M,
478 + grid_m,
479 + bid_n,
480 + flags_uc_ptr,
481 + num_ranks,
482 + )
483 +
484 + # Exit barrier. GPU scope is sufficient because the kernel writes
485 + # local global memory and only needs to flush it to local L2.
486 + cute.arch.barrier(barrier_id=BAR_COMM, number_of_threads=128)
487 + if tid_ == 0:
488 + exit_flag = total_tiles + raw_bid
489 + utils.distributed.multimem_red_add1(
490 + flags_mc_ptr + exit_flag, order="release", scope="gpu"
491 + )
492 + utils.distributed.spin_lock_atom_cas_acquire_wait(
493 + flags_uc_ptr + exit_flag,
494 + expected_val=num_ranks,
495 + reset_val=0,
496 + scope="gpu",
497 + )
498 +
499 + else:
500 + # Epilogue warps
501 + warp_id_ = warp_id % 4
502 + tid_ = tid % 128
503 +
504 + peer_flag_bases = cute.make_rmem_tensor(num_ranks, Int64)
505 + if tid_ == 0:
506 + for rank in cutlass.range_constexpr(num_ranks):
507 + peer_flag_bases[rank] = cute.arch.load(
508 + (peer_flags.iterator + rank).llvm_ptr,
509 + Int64,
510 + )
511 +
512 + if warp_id_ == 0:
513 + _tcgen05.alloc(taddr, cta_group)
514 + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=5 * 32)
515 +
516 + WIDTH = cutlass.const_expr(16)
517 + partial_vecs = cute.zipped_divide(partial_uc, (1, WIDTH))
518 +
519 + bf16x16_atom = cute.make_copy_atom(
520 + cute.nvgpu.CopyUniversalOp(),
521 + BFloat16,
522 + num_bits_per_copy=256,
523 + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE,
524 + )
525 +
526 + tmem_stage = 0
527 + parity = 0
528 +
529 + if cutlass.const_expr(is_2cta):
530 + tmem_empty_mbar_ = to_cta0_smem(tmem_empty_mbar)
531 + else:
532 + tmem_empty_mbar_ = tmem_empty_mbar
533 +
534 + for bid in range(raw_bid, total_tiles, num_bids):
535 + bid_m = bid % grid_m
536 + bid_n = bid // grid_m
537 +
538 + if warp_id_ == 0:
539 + cute.arch.mbarrier_wait(tmem_full_mbar + tmem_stage, parity)
540 + cute.arch.barrier(barrier_id=BAR_EPILOGUE, number_of_threads=128)
541 + _tcgen05.fence_after_thread_sync()
542 +
543 + for i in cutlass.range_constexpr(BN // WIDTH):
544 + tcol = tmem_stage * BN + i * WIDTH
545 + regs = _tcgen05.ld(warp_id_ * 32, tcol, "32x32b", WIDTH)
546 + _tcgen05.wait_ld()
547 +
548 + if cutlass.const_expr(i == BN // WIDTH - 1):
549 + _tcgen05.fence_before_thread_sync()
550 + mbarrier.arrive(tmem_empty_mbar_ + tmem_stage, "cluster")
551 +
552 + tmp = cute.make_rmem_tensor(WIDTH, BFloat16)
553 + tmp.store(regs.to(BFloat16))
554 +
555 + global_row = bid_m * BM + tid_
556 + if global_row < M:
557 + coord = (global_row, bid_n * (BN // WIDTH) + i)
558 + cute.copy(bf16x16_atom, tmp, partial_vecs[None, coord])
559 +
560 + cute.arch.barrier(barrier_id=BAR_EPILOGUE, number_of_threads=128)
561 +
562 + # Signal GEMM completion. GPU scope is sufficient because the
563 + # partial output only needs to be flushed to local L2.
564 + if tid_ == 0 and bid_m * BM < M:
565 + gemm_start = bid_m * BM
566 + gemm_end = min(gemm_start + BM, M)
567 + first_owner = gemm_start // output.shape[0]
568 + last_owner = (gemm_end - 1) // output.shape[0]
569 +
570 + # signal to all consuming ranks (can be more than 1)
571 + for rank in cutlass.range_constexpr(num_ranks):
572 + if first_owner <= rank and rank <= last_owner:
573 + ptr = cute.make_ptr(
574 + Int32,
575 + peer_flag_bases[rank],
576 + cute.AddressSpace.gmem,
577 + assumed_align=16,
578 + )
579 + ptr += bid_m + bid_n * grid_m
580 + utils.distributed.red_add1(
581 + ptr, order="release", scope="gpu"
582 + )
583 +
584 + tmem_stage = (tmem_stage + 1) % num_tmem_stages
585 + if tmem_stage == 0:
586 + parity ^= 1
587 +
588 + if cutlass.const_expr(is_2cta):
589 + cute.arch.cluster_arrive_relaxed()
590 + cute.arch.cluster_wait()
591 + else:
592 + cute.arch.barrier(barrier_id=BAR_EPILOGUE, number_of_threads=128)
593 + if warp_id_ == 0:
594 + _tcgen05.dealloc(cta_group)
595 +
596 + @cache
597 + @staticmethod
598 + def compile(rank: int, num_ranks: int, BN: int, cta_group: int):
599 + M = cute.sym_int()
600 + padded_M = cute.sym_int()
601 + N = cute.sym_int()
602 + K = cute.sym_int()
603 + local_M = cute.sym_int()
604 + num_flags = cute.sym_int()
605 +
606 + A = make_fake_tensor(
607 + BFloat16, (M, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16
608 + )
609 + B = make_fake_tensor(
610 + BFloat16, (N, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16
611 + )
612 + partial = make_fake_tensor(
613 + BFloat16,
614 + (padded_M, N),
615 + (cute.sym_int64(divisibility=16), 1),
616 + assumed_align=32,
617 + )
618 + 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,
624 + )
625 + flags = make_fake_tensor(Int32, (num_flags,), (1,), assumed_align=16)
626 + flags_mc_ptr = nullptr(Int32, cute.AddressSpace.gmem, assumed_align=16)
627 + peer_flag_ptr = nullptr(Int64, cute.AddressSpace.gmem, assumed_align=8)
628 +
629 + stream = make_fake_stream(use_tvm_ffi_env_stream=True)
630 + kernel = Sm100GemmRsBF16(rank, num_ranks, BN, cta_group)
631 + return cute.compile(
632 + kernel,
633 + A,
634 + B,
635 + partial,
636 + partial_mc_ptr,
637 + output,
638 + flags,
639 + flags_mc_ptr,
640 + peer_flag_ptr,
641 + 128,
642 + stream,
643 + options="--enable-tvm-ffi",
644 + )
645 +
646 +
647 +class GemmRS:
648 + """Own the max-sized symmetric workspace for Kimi-K3 GEMM-RS launches.
649 +
650 + All TP ranks must belong to one NVLink domain for multimem instructions.
651 + """
652 +
653 + def __init__(self, *, max_M: int, N: int) -> None:
654 + tp_group = get_tp_group()
655 + group = tp_group.device_group
656 + rank = tp_group.rank_in_group
657 + world_size = tp_group.world_size
658 + device = torch.device("cuda", torch.accelerator.current_device_index())
659 +
660 + assert 1 < world_size <= 16
661 + assert 128 % world_size == 0
662 + assert max_M >= 128 and N % 128 == 0
663 +
664 + max_M = (max_M + world_size - 1) // world_size * world_size
665 + self.rank = rank
666 + self.world_size = world_size
667 + self.max_M = max_M
668 + self.N = N
669 + self.device = device
670 +
671 + self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device)
672 + self.partial_handle = symm_mem.rendezvous(self.partial, group)
673 + self.partial_mc_ptr = make_ptr(
674 + BFloat16,
675 + self.partial_handle.multicast_ptr,
676 + cute.AddressSpace.gmem,
677 + assumed_align=32,
678 + )
679 +
680 + grid_m = (max_M + 127) // 128
681 + cta_group = 2 if max_M >= 1024 or grid_m % 2 == 0 else 1
682 + grid_m = (grid_m + cta_group - 1) // cta_group * cta_group
683 + self.num_sms = torch.cuda.get_device_properties(device).multi_processor_count
684 + max_flags = grid_m * (N // 128) + self.num_sms
685 + self.flags = symm_mem.empty(max_flags, dtype=torch.int32, device=device)
686 + self.flags_handle = symm_mem.rendezvous(self.flags, group)
687 + self.flags.zero_()
688 + self.flags_mc_ptr = make_ptr(
689 + Int32,
690 + self.flags_handle.multicast_ptr,
691 + cute.AddressSpace.gmem,
692 + assumed_align=16,
693 + )
694 + self.peer_flag_ptr = make_ptr(
695 + Int64,
696 + self.flags_handle.buffer_ptrs_dev,
697 + cute.AddressSpace.gmem,
698 + assumed_align=8,
699 + )
700 +
701 + assert self.partial_handle.multicast_ptr != 0
702 + assert self.flags_handle.multicast_ptr != 0
703 + torch.accelerator.synchronize(device)
704 + tp_group.barrier()
705 +
706 + def can_run(self, linear: LinearBase) -> bool:
707 + # Validate projection-invariant requirements once during model init.
708 + # only supports BF16 for now
709 + if not isinstance(linear.quant_method, UnquantizedLinearMethod):
710 + return False
711 + w = linear.weight
712 + if w.ndim != 2:
713 + return False
714 + K = w.shape[1]
715 + return (
716 + w.shape == (self.N, K)
717 + and K % 64 == 0
718 + and w.dtype == torch.bfloat16
719 + and w.device == self.device
720 + and w.is_contiguous()
721 + )
722 +
723 + def should_run(self, x: torch.Tensor) -> bool:
724 + # Small-M shapes are supported but faster on the existing LL path.
725 + return x.shape[0] >= 128
726 +
727 + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
728 + assert x.ndim == 2
729 + M, K = x.shape
730 + assert 0 < M <= self.max_M
731 + assert w.shape == (self.N, K) and K % 64 == 0
732 + assert w.dtype == torch.bfloat16
733 + assert w.device == self.device
734 + assert w.is_contiguous()
735 + assert x.dtype == torch.bfloat16
736 + assert x.device == self.device
737 + assert x.is_contiguous()
738 + N = w.shape[0]
739 + padded_M = (M + self.world_size - 1) // self.world_size
740 + padded_M *= self.world_size
741 + local_M = padded_M // self.world_size
742 +
743 + grid_m = (M + 127) // 128
744 + # Avoid padding small odd grids; 2-CTA wins consistently for M >= 1024.
745 + cta_group = 2 if M >= 1024 or grid_m % 2 == 0 else 1
746 + grid_m = (grid_m + cta_group - 1) // cta_group * cta_group
747 + BN = 256 if M * K >= 24 * 1024 * 1024 else 128
748 + assert N % BN == 0
749 +
750 + num_tiles = grid_m * (N // BN)
751 + num_ctas = min(num_tiles, self.num_sms)
752 + num_ctas = num_ctas // cta_group * cta_group
753 + assert self.flags.numel() >= num_tiles + num_ctas
754 +
755 + output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device)
756 + compiled = Sm100GemmRsBF16.compile(
757 + self.rank,
758 + self.world_size,
759 + BN,
760 + cta_group,
761 + )
762 + compiled(
763 + x,
764 + w,
765 + self.partial[:padded_M],
766 + self.partial_mc_ptr,
767 + output,
768 + self.flags,
769 + self.flags_mc_ptr,
770 + self.peer_flag_ptr,
771 + num_ctas,
772 + )
773 + return output
774 +
775 +
776 +_gemm_rs: GemmRS | None = None
777 +
778 +
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
784 + return
785 + _gemm_rs = GemmRS(max_M=max_M, N=N)
786 +
787 +
788 +def get_gemm_rs() -> GemmRS:
789 + assert _gemm_rs is not None, "GEMM-RS is not initialized"
790 + return _gemm_rs