@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f92b78f6ef

[Kernel][DSV4.1] Fuse MXFP8 wo_b GEMM with sequence-parallel reduce-scatter (#57428)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

Canlin Guo committed Sep 22, 2026 at 16:07 UTC f92b78f6ef9c9b28f60668da77af5b65645b1a45
18 files changed +1036 -623
.buildkite/test_areas/distributed.yaml
+6 -5
@@ -387,17 +387,18 @@ steps:
387 - pytest -v -s tests/v1/distributed/test_dbo.py
388 - pytest -v -s tests/distributed/test_mnnvl_alltoall.py
389
390 -- label: ":nvidia: (B200) Kimi-K3 GEMM-RS/AR"
391 - key: kimi-k3-gemm-rs-ar-2xb200
390 +- label: ":nvidia: (B200) GEMM-RS/AR"
391 + key: gemm-rs-ar-2xb200
392 timeout_in_minutes: 20
393 device: b200-k8s
394 working_dir: "/vllm-workspace/"
395 num_devices: 2
396 source_file_dependencies:
397 - - vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py
398 - - tests/kernels/test_kimi_k3_gemm_rs_ar.py
397 + - vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
398 + - vllm/cute_utils/
399 + - tests/kernels/test_gemm_rs_ar.py
400 commands:
400 - - pytest -v -s tests/kernels/test_kimi_k3_gemm_rs_ar.py
401 + - pytest -v -s tests/kernels/test_gemm_rs_ar.py
402
403 - label: ":nvidia: (L4) Distributed 2-Node"
404 key: 2-node-test-4-gpus
benchmarks/kernels/benchmark_gemm_rs_ar.py new
+407
@@ -0,0 +1,407 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Benchmark the SM100 GEMM-RS/AR kernel against the model's unfused path.
4 +
5 +The baseline is what a model runs without the fusion: the projection's own
6 +linear kernel followed by vLLM's collective (``sp_reduce_scatter`` for RS,
7 +``tensor_model_parallel_all_reduce`` for AR), with the collective backend
8 +chosen by vLLM's dispatch exactly as in serving. The same two kernels are also
9 +timed bare (the GEMM alone, the collective alone on a precomputed partial);
10 +their max is the floor a perfectly overlapped fused kernel could reach.
11 +
12 +All ranks must belong to one NVLink domain. Run a TP8 sweep over the Kimi-K3
13 +projections with:
14 +
15 + torchrun --nproc-per-node=8 -- benchmarks/kernels/benchmark_gemm_rs_ar.py
16 +
17 +or the DeepSeek-V4.1 MXFP8 ``wo_b`` projection at TP4 with:
18 +
19 + torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \
20 + --model deepseek_v41 --backend flashinfer_cutedsl
21 +"""
22 +
23 +import argparse
24 +import json
25 +import os
26 +import statistics
27 +from collections.abc import Callable
28 +
29 +import torch
30 +import torch.distributed as dist
31 +
32 +from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
33 +from vllm.distributed import (
34 + cleanup_dist_env_and_memory,
35 + tensor_model_parallel_all_reduce,
36 +)
37 +from vllm.distributed.parallel_state import (
38 + get_tp_group,
39 + init_distributed_environment,
40 + initialize_model_parallel,
41 +)
42 +from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
43 +from vllm.model_executor.layers.linear import RowParallelLinear
44 +from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter
45 +
46 +# Per model: (hidden size N, global input widths K). Kimi-K3 lists the
47 +# shared-expert down-proj and attention O-proj; DeepSeek-V4.1 lists the
48 +# ``wo_b`` output projection (o_groups * o_lora_rank).
49 +_MODEL_PROJECTIONS = {
50 + "kimi_k3": (7168, (6144, 12288)),
51 + "deepseek_v41": (5120, (8192,)),
52 +}
53 +
54 +
55 +def parse_args() -> argparse.Namespace:
56 + parser = argparse.ArgumentParser()
57 + parser.add_argument(
58 + "--mode",
59 + choices=("rs", "ar"),
60 + default="rs",
61 + help="Collective mode to benchmark.",
62 + )
63 + parser.add_argument(
64 + "--m",
65 + type=int,
66 + nargs="+",
67 + default=[128, 512, 2048, 8192, 32768],
68 + help="Global token counts to benchmark.",
69 + )
70 + parser.add_argument(
71 + "--model",
72 + choices=tuple(_MODEL_PROJECTIONS),
73 + default="kimi_k3",
74 + help="Model whose projection shapes set the default --n and --k.",
75 + )
76 + parser.add_argument(
77 + "--backend",
78 + choices=("bf16", "flashinfer_cutlass", "flashinfer_cutedsl"),
79 + default="bf16",
80 + help="Weight dtype/layout: BF16 or an online-MXFP8 FlashInfer kernel.",
81 + )
82 + parser.add_argument(
83 + "--k",
84 + type=int,
85 + nargs="+",
86 + help=(
87 + "Per-rank input dimensions. By default, derive the selected "
88 + "model's projection dimensions from the TP world size."
89 + ),
90 + )
91 + parser.add_argument(
92 + "--n",
93 + type=int,
94 + help="Output width. Defaults to the selected model's hidden size.",
95 + )
96 + parser.add_argument(
97 + "--num-workspaces",
98 + type=int,
99 + default=10,
100 + help="Pointer-distinct inputs and CUDA graphs to rotate.",
101 + )
102 + parser.add_argument("--warmup-replays", type=int, default=5)
103 + parser.add_argument("--samples", type=int, default=20)
104 + parser.add_argument("--json", type=str, help="Also write the results here.")
105 + return parser.parse_args()
106 +
107 +
108 +def make_projection(
109 + gemm_rs_ar: GemmRsAr,
110 + weight: torch.Tensor,
111 + world_size: int,
112 + backend: str,
113 +) -> RowParallelLinear:
114 + """Wrap a BF16 [N, K] weight the way the model would hand it to GEMM-RS.
115 +
116 + For the MXFP8 backends this quantizes online through the FlashInfer
117 + kernel's post-load step, so the fused kernel and the unfused baseline see
118 + the production weight and scale layouts.
119 + """
120 + N, K = weight.shape
121 + quant_config = None
122 + if backend != "bf16":
123 + from vllm.config.quantization import QuantizationConfigArgs
124 + from vllm.model_executor.layers.quantization.online.base import (
125 + OnlineQuantizationConfig,
126 + )
127 +
128 + quant_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8"))
129 + with torch.device(weight.device):
130 + linear = RowParallelLinear(
131 + K * world_size,
132 + N,
133 + bias=False,
134 + params_dtype=weight.dtype,
135 + quant_config=quant_config,
136 + # The baseline reduces the local partial itself, like the models.
137 + reduce_results=False,
138 + return_bias=False,
139 + )
140 + # Select the fused path before online quantization replaces the weights.
141 + assert gemm_rs_ar.can_run(linear)
142 + linear.weight = torch.nn.Parameter(weight, requires_grad=False)
143 + linear.quant_method.process_weights_after_loading(linear)
144 + return linear
145 +
146 +
147 +def capture_graph(
148 + op: Callable[[], torch.Tensor],
149 + stream: torch.cuda.Stream,
150 + cpu_group: dist.ProcessGroup,
151 +) -> torch.cuda.CUDAGraph:
152 + stream.wait_stream(torch.cuda.current_stream())
153 + dist.barrier(group=cpu_group)
154 + with torch.cuda.stream(stream):
155 + for _ in range(3):
156 + op()
157 + stream.synchronize()
158 + dist.barrier(group=cpu_group)
159 +
160 + graph = torch.cuda.CUDAGraph()
161 + with torch.cuda.graph(graph, stream=stream):
162 + op()
163 + torch.cuda.current_stream().wait_stream(stream)
164 + dist.barrier(group=cpu_group)
165 + return graph
166 +
167 +
168 +def benchmark_graphs(
169 + candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]],
170 + warmup_replays: int,
171 + samples: int,
172 + device_group: dist.ProcessGroup,
173 + device_barrier: Callable[[], None],
174 +) -> dict[str, float]:
175 + """Median over samples of the slowest rank's replay time, in us."""
176 + names = list(candidate_graphs)
177 + for round_index in range(warmup_replays):
178 + for offset in range(len(names)):
179 + name = names[(round_index + offset) % len(names)]
180 + graphs = candidate_graphs[name]
181 + device_barrier()
182 + graphs[round_index % len(graphs)].replay()
183 + torch.accelerator.synchronize()
184 +
185 + timings: dict[str, list[float]] = {name: [] for name in names}
186 + start = torch.cuda.Event(enable_timing=True)
187 + end = torch.cuda.Event(enable_timing=True)
188 + for sample_index in range(samples):
189 + for offset in range(len(names)):
190 + name = names[(sample_index + offset) % len(names)]
191 + graphs = candidate_graphs[name]
192 + device_barrier()
193 + start.record()
194 + graphs[sample_index % len(graphs)].replay()
195 + end.record()
196 + end.synchronize()
197 +
198 + elapsed = torch.tensor(
199 + start.elapsed_time(end) * 1000,
200 + dtype=torch.float64,
201 + device=torch.accelerator.current_device_index(),
202 + )
203 + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
204 + timings[name].append(elapsed.item())
205 + return {name: statistics.median(values) for name, values in timings.items()}
206 +
207 +
208 +def benchmark_shape(
209 + gemm_rs_ar: GemmRsAr,
210 + mode: str,
211 + backend: str,
212 + M: int,
213 + N: int,
214 + K: int,
215 + num_workspaces: int,
216 + warmup_replays: int,
217 + samples: int,
218 + device_group: dist.ProcessGroup,
219 + cpu_group: dist.ProcessGroup,
220 + device_barrier: Callable[[], None],
221 +) -> dict[str, float | int | str]:
222 + all_reduce = mode == "ar"
223 + world_size = dist.get_world_size(device_group)
224 + rank = dist.get_rank(device_group)
225 + device = torch.device("cuda", torch.accelerator.current_device_index())
226 + local_M = (M + world_size - 1) // world_size
227 +
228 + rng = torch.Generator(device=device)
229 + rng.manual_seed(1000 + rank * 10 + M + K)
230 + inputs = [
231 + torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng)
232 + for _ in range(num_workspaces)
233 + ]
234 + projections = [
235 + make_projection(
236 + gemm_rs_ar,
237 + torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng),
238 + world_size,
239 + backend,
240 + )
241 + for _ in range(num_workspaces)
242 + ]
243 +
244 + def collective(partial: torch.Tensor) -> torch.Tensor:
245 + if all_reduce:
246 + return tensor_model_parallel_all_reduce(partial)
247 + return sp_reduce_scatter(partial)
248 +
249 + def gemm(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
250 + return linear(x)
251 +
252 + def unfused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
253 + return collective(linear(x))
254 +
255 + def fused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
256 + return gemm_rs_ar.apply(x, linear)
257 +
258 + # The bare collective runs on the GEMM's own output so it sees the same
259 + # message size and layout the unfused path hands to the collective.
260 + partials = [linear(x) for x, linear in zip(inputs, projections)]
261 +
262 + expected = unfused(inputs[0], projections[0])
263 + actual = fused(inputs[0], projections[0])
264 + torch.accelerator.synchronize(device)
265 + rows = M if all_reduce else min(max(M - rank * local_M, 0), local_M)
266 + torch.testing.assert_close(
267 + actual[:rows], expected[:rows], rtol=5e-2, atol=4.0, msg=f"M={M}, K={K}"
268 + )
269 +
270 + candidates: dict[str, list[Callable[[], torch.Tensor]]] = {
271 + "gemm": [lambda x=x, w=w: gemm(x, w) for x, w in zip(inputs, projections)],
272 + "collective": [lambda p=p: collective(p) for p in partials],
273 + "vllm": [lambda x=x, w=w: unfused(x, w) for x, w in zip(inputs, projections)],
274 + "gemm_rs_ar": [
275 + lambda x=x, w=w: fused(x, w) for x, w in zip(inputs, projections)
276 + ],
277 + }
278 +
279 + candidate_graphs = {}
280 + keepalive: list[object] = []
281 + for name, ops in candidates.items():
282 + stream = torch.cuda.Stream()
283 + candidate_graphs[name] = [capture_graph(op, stream, cpu_group) for op in ops]
284 + keepalive.append(stream)
285 +
286 + times = benchmark_graphs(
287 + candidate_graphs, warmup_replays, samples, device_group, device_barrier
288 + )
289 + floor = max(times["gemm"], times["collective"])
290 + return {
291 + "mode": mode.upper(),
292 + "backend": backend,
293 + "M": M,
294 + "N": N,
295 + "K": K,
296 + "gemm_us": times["gemm"],
297 + "collective_us": times["collective"],
298 + "floor_us": floor,
299 + "vllm_us": times["vllm"],
300 + "gemm_rs_ar_us": times["gemm_rs_ar"],
301 + "latency_change": (times["gemm_rs_ar"] - times["vllm"]) / times["vllm"],
302 + "floor_gap": (times["gemm_rs_ar"] - floor) / floor,
303 + }
304 +
305 +
306 +def print_results(results: list[dict[str, float | int | str]]) -> None:
307 + collective = results[0]["mode"]
308 + backend = results[0]["backend"]
309 + print(f"### GEMM-{collective} vs vLLM unfused path ({backend} weights)")
310 + print(
311 + f"| M | N | K | GEMM (us) | {collective} (us) | floor (us) | vLLM (us) | "
312 + f"GEMM-{collective} (us) | vs vLLM | vs floor |"
313 + )
314 + print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|")
315 + for r in results:
316 + print(
317 + f"| {r['M']} | {r['N']} | {r['K']} | {r['gemm_us']:.1f} | "
318 + f"{r['collective_us']:.1f} | {r['floor_us']:.1f} | {r['vllm_us']:.1f} | "
319 + f"{r['gemm_rs_ar_us']:.1f} | {r['latency_change']:+.1%} | "
320 + f"{r['floor_gap']:+.1%} |"
321 + )
322 +
323 +
324 +def main() -> None:
325 + args = parse_args()
326 + default_n, default_k = _MODEL_PROJECTIONS[args.model]
327 + if args.n is None:
328 + args.n = default_n
329 + assert args.m and min(args.m) > 0
330 + assert args.n % 128 == 0
331 + assert args.num_workspaces > 0
332 + assert args.warmup_replays >= 0
333 + assert args.samples > 0
334 +
335 + local_rank = int(os.environ["LOCAL_RANK"])
336 + torch.accelerator.set_device_index(local_rank)
337 + init_distributed_environment()
338 + world_size = dist.get_world_size()
339 + if args.k is None:
340 + assert all(K % world_size == 0 for K in default_k)
341 + K_values = [K // world_size for K in default_k]
342 + else:
343 + K_values = args.k
344 + k_alignment = 64 if args.backend == "bf16" else 128
345 + assert all(K % k_alignment == 0 for K in K_values)
346 +
347 + config = VllmConfig()
348 + config.model_config = ModelConfig(dtype="bfloat16")
349 + # Pin the online-MXFP8 kernel the projections quantize through.
350 + config.kernel_config.linear_backend = (
351 + "auto" if args.backend == "bf16" else args.backend
352 + )
353 + with set_current_vllm_config(config):
354 + # Builds the TP communicator as in serving, so the baseline's
355 + # collective is whatever vLLM's dispatch selects (FlashInfer all-reduce
356 + # first, then NCCL symmetric memory, custom all-reduce, pynccl).
357 + initialize_model_parallel(tensor_model_parallel_size=world_size)
358 +
359 + tp_group = get_tp_group()
360 + pynccl_comm = tp_group.device_communicator.pynccl_comm
361 + assert pynccl_comm is not None
362 + sync_input = torch.zeros(1, device=torch.accelerator.current_device_index())
363 + sync_output = torch.empty_like(sync_input)
364 +
365 + def device_barrier() -> None:
366 + # Order the timed launch after a device-side rank rendezvous without
367 + # including the rendezvous itself in the measured event interval.
368 + pynccl_comm.all_reduce(sync_input, sync_output)
369 +
370 + gemm_rs_ar = GemmRsAr(
371 + max_M=max(args.m),
372 + N=args.n,
373 + all_reduce=args.mode == "ar",
374 + )
375 + with set_current_vllm_config(config):
376 + results = [
377 + benchmark_shape(
378 + gemm_rs_ar,
379 + args.mode,
380 + args.backend,
381 + M,
382 + args.n,
383 + K,
384 + args.num_workspaces,
385 + args.warmup_replays,
386 + args.samples,
387 + tp_group.device_group,
388 + tp_group.cpu_group,
389 + device_barrier,
390 + )
391 + for K in K_values
392 + for M in args.m
393 + ]
394 + del gemm_rs_ar
395 +
396 + if tp_group.rank_in_group == 0:
397 + print_results(results)
398 + if args.json:
399 + with open(args.json, "w") as f:
400 + json.dump(results, f, indent=2)
401 +
402 + dist.barrier(group=tp_group.cpu_group)
403 + cleanup_dist_env_and_memory()
404 +
405 +
406 +if __name__ == "__main__":
407 + main()
benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py deleted
-469
@@ -1,469 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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_ar.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_ar import GemmRsAr
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 - "--mode",
46 - choices=("rs", "ar"),
47 - default="rs",
48 - help="Collective mode to benchmark.",
49 - )
50 - parser.add_argument(
51 - "--m",
52 - type=int,
53 - nargs="+",
54 - default=[128, 512, 2048, 8192, 32768],
55 - help="Global token counts to benchmark.",
56 - )
57 - parser.add_argument(
58 - "--k",
59 - type=int,
60 - nargs="+",
61 - help=(
62 - "Per-rank input dimensions. By default, derive the Kimi-K3 "
63 - "shared-expert down-proj and O-proj dimensions from the TP "
64 - "world size."
65 - ),
66 - )
67 - parser.add_argument("--n", type=int, default=7168)
68 - parser.add_argument(
69 - "--num-workspaces",
70 - type=int,
71 - default=10,
72 - help="Pointer-distinct inputs and CUDA graphs to rotate.",
73 - )
74 - parser.add_argument("--warmup-replays", type=int, default=5)
75 - parser.add_argument("--samples", type=int, default=20)
76 - return parser.parse_args()
77 -
78 -
79 -def capture_graph(
80 - op: Callable[[], torch.Tensor],
81 - stream: torch.cuda.Stream,
82 - cpu_group: dist.ProcessGroup,
83 -) -> tuple[torch.cuda.CUDAGraph, list[torch.Tensor | None]]:
84 - result: list[torch.Tensor | None] = [None]
85 - stream.wait_stream(torch.cuda.current_stream())
86 - dist.barrier(group=cpu_group)
87 - with torch.cuda.stream(stream):
88 - for _ in range(3):
89 - result[0] = op()
90 - stream.synchronize()
91 - dist.barrier(group=cpu_group)
92 -
93 - graph = torch.cuda.CUDAGraph()
94 - with torch.cuda.graph(graph, stream=stream):
95 - result[0] = op()
96 - torch.cuda.current_stream().wait_stream(stream)
97 - dist.barrier(group=cpu_group)
98 - return graph, result
99 -
100 -
101 -def benchmark_graphs(
102 - candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]],
103 - warmup_replays: int,
104 - samples: int,
105 - device_group: dist.ProcessGroup,
106 - device_barrier: Callable[[], None],
107 -) -> dict[str, float]:
108 - candidate_names = list(candidate_graphs)
109 - for round_index in range(warmup_replays):
110 - for candidate_index in range(len(candidate_names)):
111 - candidate_id = (round_index + candidate_index) % len(candidate_names)
112 - name = candidate_names[candidate_id]
113 - graphs = candidate_graphs[name]
114 - device_barrier()
115 - graphs[round_index % len(graphs)].replay()
116 - torch.accelerator.synchronize()
117 -
118 - timings: dict[str, list[float]] = {name: [] for name in candidate_names}
119 - start = torch.cuda.Event(enable_timing=True)
120 - end = torch.cuda.Event(enable_timing=True)
121 - for sample_index in range(samples):
122 - for candidate_index in range(len(candidate_names)):
123 - candidate_id = (sample_index + candidate_index) % len(candidate_names)
124 - name = candidate_names[candidate_id]
125 - graphs = candidate_graphs[name]
126 - device_barrier()
127 - start.record()
128 - graphs[sample_index % len(graphs)].replay()
129 - end.record()
130 - end.synchronize()
131 -
132 - elapsed = torch.tensor(
133 - start.elapsed_time(end) * 1000,
134 - dtype=torch.float64,
135 - device=torch.accelerator.current_device_index(),
136 - )
137 - dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
138 - timings[name].append(elapsed.item())
139 - return {name: statistics.median(values) for name, values in timings.items()}
140 -
141 -
142 -def valid_rows(M: int, local_M: int, rank: int) -> int:
143 - return min(max(M - rank * local_M, 0), local_M)
144 -
145 -
146 -def benchmark_shape(
147 - gemm_rs_ar: GemmRsAr,
148 - mode: str,
149 - M: int,
150 - N: int,
151 - K: int,
152 - num_workspaces: int,
153 - warmup_replays: int,
154 - samples: int,
155 - device_group: dist.ProcessGroup,
156 - cpu_group: dist.ProcessGroup,
157 - device_barrier: Callable[[], None],
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())
163 - padded_M = (M + world_size - 1) // world_size * world_size
164 - local_M = padded_M // world_size
165 -
166 - rng = torch.Generator(device=device)
167 - rng.manual_seed(1000 + rank * 10 + M + K)
168 - inputs = [
169 - torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng)
170 - for _ in range(num_workspaces)
171 - ]
172 - weights = [
173 - torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng)
174 - for _ in range(num_workspaces)
175 - ]
176 -
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)
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):
184 - collective_input = symm_mem.empty(
185 - (padded_M, N),
186 - dtype=torch.bfloat16,
187 - device=device,
188 - )
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)
196 -
197 - if padded_M > M:
198 - partial[M:].zero_()
199 - symm_partial[M:].zero_()
200 -
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 -
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,
225 - group=device_group,
226 - )
227 - return symm_output
228 -
229 - return run
230 -
231 - def make_fused_gemm_collective(
232 - x: torch.Tensor, weight: torch.Tensor
233 - ) -> Callable[[], torch.Tensor]:
234 - def run() -> torch.Tensor:
235 - return gemm_rs_ar(x, weight)
236 -
237 - return run
238 -
239 - def make_torch_gemm(
240 - x: torch.Tensor,
241 - weight: torch.Tensor,
242 - ) -> Callable[[], torch.Tensor]:
243 - def run() -> torch.Tensor:
244 - return torch.mm(x, weight.T, out=gemm_output)
245 -
246 - return run
247 -
248 - def make_ring_ll_collective(
249 - collective_input: torch.Tensor,
250 - ) -> Callable[[], torch.Tensor]:
251 - def run() -> torch.Tensor:
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 -
262 - def make_ldmc_collective(
263 - collective_input: torch.Tensor,
264 - ) -> Callable[[], torch.Tensor]:
265 - def run() -> torch.Tensor:
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
275 -
276 - candidates = (
277 - Candidate(
278 - "ring_ll_us",
279 - [make_torch_ring_ll_gemm_collective(x, w) for x, w in zip(inputs, weights)],
280 - ),
281 - Candidate(
282 - "ldmc_us",
283 - [make_torch_ldmc_gemm_collective(x, w) for x, w in zip(inputs, weights)],
284 - ),
285 - Candidate(
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",
291 - [make_torch_gemm(x, w) for x, w in zip(inputs, weights)],
292 - check_correctness=False,
293 - ),
294 - Candidate(
295 - "ring_ll_collective_us",
296 - [make_ring_ll_collective(x) for x in collective_inputs],
297 - check_correctness=False,
298 - ),
299 - Candidate(
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]()
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
311 - actual = candidate.runs[0]()
312 - torch.accelerator.synchronize(device)
313 - torch.testing.assert_close(
314 - actual[:rows],
315 - expected[:rows],
316 - rtol=5e-2,
317 - atol=4.0,
318 - )
319 -
320 - candidate_graphs = {}
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]
325 - candidate_graphs[candidate.name] = [graph for graph, _ in bundles]
326 - graph_keepalive.extend(bundles)
327 - graph_keepalive.append(stream)
328 -
329 - times = benchmark_graphs(
330 - candidate_graphs,
331 - warmup_replays,
332 - samples,
333 - device_group,
334 - device_barrier,
335 - )
336 -
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,
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 -
350 -def print_results(results: list[dict[str, float | int | str]]) -> None:
351 - results_df = pd.DataFrame(results)
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 - ]
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 - )
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 - )
388 -
389 - print(f"### {collective} end-to-end latency")
390 - print(end_to_end.to_markdown(index=False))
391 - print(f"\n### {collective} component latency")
392 - print(components.to_markdown(index=False))
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()
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
405 - assert args.samples > 0
406 -
407 - local_rank = int(os.environ["LOCAL_RANK"])
408 - torch.accelerator.set_device_index(local_rank)
409 - init_distributed_environment()
410 - world_size = dist.get_world_size()
411 - if args.k is None:
412 - assert all(K % world_size == 0 for K in _KIMI_K3_PROJECTION_K)
413 - K_values = [K // world_size for K in _KIMI_K3_PROJECTION_K]
414 - else:
415 - K_values = args.k
416 - assert all(K % 64 == 0 for K in K_values)
417 - # Reserve symmetric memory for the NCCL-managed benchmark allocations.
418 - os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
419 - # NCCL-managed symmetric allocations select the NVLS/LDMC collective path.
420 - symm_mem.set_backend("NCCL")
421 - with set_current_vllm_config(VllmConfig()):
422 - initialize_model_parallel(tensor_model_parallel_size=world_size)
423 -
424 - tp_group = get_tp_group()
425 - group_warmup = torch.zeros(1, device=torch.accelerator.current_device_index())
426 - dist.all_reduce(group_warmup, group=tp_group.device_group)
427 - pynccl_comm = tp_group.device_communicator.pynccl_comm
428 - assert pynccl_comm is not None
429 - sync_input = torch.zeros(1, device=torch.accelerator.current_device_index())
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 -
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(
444 - gemm_rs_ar,
445 - args.mode,
446 - M,
447 - args.n,
448 - K,
449 - args.num_workspaces,
450 - args.warmup_replays,
451 - args.samples,
452 - tp_group.device_group,
453 - tp_group.cpu_group,
454 - device_barrier,
455 - )
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)
465 - cleanup_dist_env_and_memory()
466 -
467 -
468 -if __name__ == "__main__":
469 - main()
tests/kernels/test_gemm_rs_ar.py renamed
+87 -26
@@ -6,7 +6,7 @@ import torch
6 import torch.distributed as dist
7 import torch.multiprocessing as mp
8
9 -from tests.utils import ensure_current_vllm_config
9 +from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
10 from vllm.distributed import cleanup_dist_env_and_memory
11 from vllm.distributed.parallel_state import (
12 init_distributed_environment,
@@ -17,6 +17,7 @@ from vllm.utils.network_utils import get_open_port
17 from vllm.utils.system_utils import update_environment_variables
18
19 _SHAPES = (
20 + (129, 512),
21 (129, 768),
22 (257, 768),
23 (1023, 4224),
@@ -29,6 +30,7 @@ _N = 512
30 def _reference(
31 x: torch.Tensor,
32 weight: torch.Tensor,
33 + weight_scale: torch.Tensor | None,
34 world_size: int,
35 group: dist.ProcessGroup,
36 all_reduce: bool,
@@ -40,7 +42,21 @@ def _reference(
42 dtype=x.dtype,
43 device=x.device,
44 )
43 - torch.mm(x, weight.T, out=partial[:M])
45 + if weight_scale is None:
46 + torch.mm(x, weight.T, out=partial[:M])
47 + else:
48 + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
49 + mxfp8_e4m3_quantize,
50 + )
51 +
52 + x_q, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
53 + partial[:M] = torch._scaled_mm(
54 + x_q,
55 + weight.T,
56 + x_sf.view(torch.float8_e8m0fnu),
57 + weight_scale.view(torch.float8_e8m0fnu),
58 + out_dtype=x.dtype,
59 + )
60 if padded_M > M:
61 partial[M:].zero_()
62
@@ -75,6 +91,10 @@ def _assert_valid_rows_close(
91 rtol=5e-2,
92 atol=4.0,
93 )
94 + # Sequence-parallel callers may pass an unpadded M and read the padding
95 + # rows back (they flow through per-token kernels before being dropped),
96 + # so RS must leave them zero like ``sp_reduce_scatter`` does.
97 + assert not actual[valid_rows:].any()
98
99
100 def _run_mode(
@@ -85,19 +105,52 @@ def _run_mode(
105 rank: int,
106 world_size: int,
107 weights: dict[int, torch.Tensor],
108 + backend: str,
109 ) -> None:
110 # 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
111 + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
112
113 gemm_rs_ar = GemmRsAr(
114 max_M=max(M for M, _ in _SHAPES),
115 N=_N,
116 all_reduce=all_reduce,
117 )
118 + from vllm.config.quantization import QuantizationConfigArgs
119 + from vllm.model_executor.layers.linear import RowParallelLinear
120 + from vllm.model_executor.layers.quantization.online.base import (
121 + OnlineQuantizationConfig,
122 + )
123 +
124 + quant_config = (
125 + None
126 + if backend == "bf16"
127 + else OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8"))
128 + )
129 + projections = {}
130 + operands = {}
131 + for K, w in weights.items():
132 + with torch.device(device):
133 + linear = RowParallelLinear(
134 + K * world_size,
135 + _N,
136 + bias=False,
137 + params_dtype=torch.bfloat16,
138 + quant_config=quant_config,
139 + )
140 + # Select the fused path before online quantization replaces meta weights.
141 + assert gemm_rs_ar.can_run(linear)
142 + linear.weight = torch.nn.Parameter(w.clone(), requires_grad=False)
143 + linear.quant_method.process_weights_after_loading(linear)
144 + projections[K] = linear
145 + weight = linear.weight
146 + if backend == "flashinfer_cutedsl":
147 + weight = weight.t()
148 + operands[K] = (weight, getattr(linear, "weight_scale", None))
149 +
150 input_generator = torch.Generator(device=device)
151
99 - # Alternating shapes exercise producer-flag reuse across different grids,
100 - # CTA-group choices, and both BN=128 and BN=256 dispatches.
152 + # Alternating shapes exercise producer-flag reuse across different grids
153 + # and CTA-group choices.
154 for M, K in (*_SHAPES, *_SHAPES[::-1]):
155 input_generator.manual_seed(2000 + M + K)
156 x = torch.randn(
@@ -107,8 +160,8 @@ def _run_mode(
160 device=device,
161 generator=input_generator,
162 )
110 - expected = _reference(x, weights[K], world_size, group, all_reduce)
111 - actual = gemm_rs_ar(x, weights[K])
163 + expected = _reference(x, *operands[K], world_size, group, all_reduce)
164 + actual = gemm_rs_ar.apply(x, projections[K])
165 torch.accelerator.synchronize(device)
166 _assert_valid_rows_close(actual, expected, M, rank, all_reduce)
167
@@ -123,9 +176,9 @@ def _run_mode(
176 device=device,
177 generator=input_generator,
178 )
126 - first_output = gemm_rs_ar(lifetime_x, weights[lifetime_K])
179 + first_output = gemm_rs_ar.apply(lifetime_x, projections[lifetime_K])
180 first_snapshot = first_output.clone()
128 - second_output = gemm_rs_ar(-lifetime_x, weights[lifetime_K])
181 + second_output = gemm_rs_ar.apply(-lifetime_x, projections[lifetime_K])
182 torch.accelerator.synchronize(device)
183 assert first_output.data_ptr() != second_output.data_ptr()
184 torch.testing.assert_close(first_output, first_snapshot, rtol=0, atol=0)
@@ -141,7 +194,7 @@ def _run_mode(
194 )
195 graph_expected = _reference(
196 graph_x,
144 - weights[graph_K],
197 + *operands[graph_K],
198 world_size,
199 group,
200 all_reduce,
@@ -152,13 +205,13 @@ def _run_mode(
205 capture_stream.wait_stream(torch.cuda.current_stream())
206 with torch.cuda.stream(capture_stream):
207 for _ in range(3):
155 - graph_output.copy_(gemm_rs_ar(graph_x, weights[graph_K]))
208 + graph_output.copy_(gemm_rs_ar.apply(graph_x, projections[graph_K]))
209 capture_stream.synchronize()
210 dist.barrier(group=group)
211
212 graph = torch.cuda.CUDAGraph()
213 with torch.cuda.graph(graph, stream=capture_stream):
161 - graph_output.copy_(gemm_rs_ar(graph_x, weights[graph_K]))
214 + graph_output.copy_(gemm_rs_ar.apply(graph_x, projections[graph_K]))
215 torch.cuda.current_stream().wait_stream(capture_stream)
216 dist.barrier(group=group)
217
@@ -177,7 +230,7 @@ def _run_mode(
230 del gemm_rs_ar, capture_stream, graph, graph_output
231
232
180 -def _worker(local_rank: int, world_size: int, master_port: int) -> None:
233 +def _worker(local_rank: int, world_size: int, master_port: int, backend: str) -> None:
234 # This module pulls in cute_dsl, which is unavailable off CUDA, so importing
235 # it at module level would fail collection. Import it where it is used, as
236 # `kda.py` does.
@@ -195,7 +248,10 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None:
248 )
249
250 init_distributed_environment()
198 - with ensure_current_vllm_config():
251 + config = VllmConfig()
252 + config.model_config = ModelConfig(dtype="bfloat16")
253 + config.kernel_config.linear_backend = "auto" if backend == "bf16" else backend
254 + with set_current_vllm_config(config):
255 initialize_model_parallel(tensor_model_parallel_size=world_size)
256
257 group = dist.group.WORLD
@@ -215,24 +271,29 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None:
271
272 # Production binds one mode per worker; exercise both mode-bound instances
273 # 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 - )
274 + with set_current_vllm_config(config):
275 + for all_reduce in (False, True):
276 + _run_mode(
277 + all_reduce=all_reduce,
278 + device=device,
279 + group=group,
280 + rank=rank,
281 + world_size=world_size,
282 + weights=weights,
283 + backend=backend,
284 + )
285 cleanup_dist_env_and_memory()
286
287
288 +@pytest.mark.parametrize(
289 + "backend", ["bf16", "flashinfer_cutlass", "flashinfer_cutedsl"]
290 +)
291 @pytest.mark.distributed(num_gpus=2)
292 @pytest.mark.skipif(
293 not current_platform.is_device_capability_family(100),
233 - reason="Kimi-K3 GEMM-RS/AR requires SM100",
294 + reason="GEMM-RS/AR requires SM100",
295 )
235 -def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch) -> None:
296 +def test_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, backend: str) -> None:
297 world_size = 2
298 if torch.accelerator.device_count() < world_size:
299 pytest.skip("GEMM-RS/AR requires two GPUs")
@@ -242,7 +303,7 @@ def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch) -> None:
303 try:
304 mp.spawn(
305 _worker,
245 - args=(world_size, get_open_port()),
306 + args=(world_size, get_open_port(), backend),
307 nprocs=world_size,
308 )
309 finally:
vllm/cute_utils/_tcgen05.py
+71
@@ -78,6 +78,10 @@ def make_sdesc_128B_swizzle(LBO: int):
78 return Uint64((LBO >> 4 << 16) | (SBO >> 4 << 32) | (1 << 46) | (2 << 61))
79
80
81 +def make_mxfp8_idesc(MMA_M: int, MMA_N: int):
82 + return Uint32((MMA_N >> 3 << 17) | (1 << 23) | (MMA_M >> 7 << 27))
83 +
84 +
85 @dsl_user_op
86 def mma_f16(
87 d_tmem,
@@ -130,6 +134,73 @@ def mma_ts_f16(
134 )
135
136
137 +@dsl_user_op
138 +def mma_mxfp8(
139 + d_tmem,
140 + a_desc,
141 + b_desc,
142 + idesc,
143 + sfa_tmem,
144 + sfb_tmem,
145 + enable_input_d,
146 + cta_group: int = 1,
147 + *,
148 + loc=None,
149 + ip=None,
150 +) -> None:
151 + with cute.arch.elect_one():
152 + nvvm.tcgen05_mma_block_scale(
153 + nvvm.Tcgen05MMAKind.MXF8F6F4,
154 + NVVM_CTA_GROUP_MAP[cta_group],
155 + _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
156 + Uint64(a_desc).ir_value(loc=loc, ip=ip),
157 + Uint64(b_desc).ir_value(loc=loc, ip=ip),
158 + Int32(idesc).ir_value(loc=loc, ip=ip),
159 + Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
160 + _make_tmem_llvm_ptr(sfa_tmem, loc=loc, ip=ip),
161 + _make_tmem_llvm_ptr(sfb_tmem, loc=loc, ip=ip),
162 + scale_vec_size=nvvm.Tcgen05MMAScaleVecSize.X1,
163 + loc=loc,
164 + ip=ip,
165 + )
166 +
167 +
168 +@dsl_user_op
169 +def cp(
170 + tmem,
171 + sdesc,
172 + shape: str,
173 + mcast: str,
174 + cta_group: int = 1,
175 + *,
176 + loc=None,
177 + ip=None,
178 +) -> None:
179 + SHAPE_MAP = {
180 + "128x256b": nvvm.Tcgen05CpShape.SHAPE_128x256b,
181 + "4x256b": nvvm.Tcgen05CpShape.SHAPE_4x256b,
182 + "128x128b": nvvm.Tcgen05CpShape.SHAPE_128x128b,
183 + "64x128b": nvvm.Tcgen05CpShape.SHAPE_64x128b,
184 + "32x128b": nvvm.Tcgen05CpShape.SHAPE_32x128b,
185 + }
186 + MCAST_MAP = {
187 + "none": nvvm.Tcgen05CpMulticast.NONE,
188 + "warpx2::02_13": nvvm.Tcgen05CpMulticast.WARPX2_02_13,
189 + "warpx2::01_23": nvvm.Tcgen05CpMulticast.WARPX2_01_23,
190 + "warpx4": nvvm.Tcgen05CpMulticast.WARPX4,
191 + }
192 + with cute.arch.elect_one():
193 + nvvm.tcgen05_cp(
194 + SHAPE_MAP[shape],
195 + _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip),
196 + Uint64(sdesc).ir_value(loc=loc, ip=ip),
197 + group=NVVM_CTA_GROUP_MAP[cta_group],
198 + multicast=MCAST_MAP[mcast],
199 + loc=loc,
200 + ip=ip,
201 + )
202 +
203 +
204 @dsl_user_op
205 def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None):
206 mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip)
vllm/envs.py
+6 -4
@@ -213,7 +213,7 @@ if TYPE_CHECKING:
213 VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False
214 VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False
215 VLLM_KIMI_K3_GEMM_AR: bool = True
216 - VLLM_KIMI_K3_GEMM_RS: bool = False
216 + VLLM_ENABLE_GEMM_RS: bool = False
217 VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
218 VLLM_USE_FLASHINFER_MOE_INT4: bool = False
219 VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None
@@ -1629,9 +1629,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1629 # Use the SM100 BF16 GEMM-AR kernel for eligible Kimi-K3 row-parallel
1630 # attention projections. All TP ranks must belong to one NVLink domain.
1631 "VLLM_KIMI_K3_GEMM_AR": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_AR", "1"))),
1632 - # Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel
1633 - # row-parallel projections. All TP ranks must belong to one NVLink domain.
1634 - "VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))),
1632 + # Fuse eligible sequence-parallel row-parallel projections with their TP
1633 + # reduce-scatter using the SM100 BF16/MXFP8 GEMM-RS kernel (Kimi-K3
1634 + # attention/shared-expert projections, DeepSeek-V4.1 ``wo_b``). All TP
1635 + # ranks must belong to one NVLink domain.
1636 + "VLLM_ENABLE_GEMM_RS": lambda: bool(int(os.getenv("VLLM_ENABLE_GEMM_RS", "0"))),
1637 # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.
1638 # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).
1639 "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(
vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py renamed
+340 -57
@@ -1,18 +1,36 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -"""SM100 BF16 GEMM with fused tensor-parallel reduce-scatter/all-reduce."""
4 +"""SM100 BF16/MXFP8 GEMM with fused tensor-parallel reduce-scatter/all-reduce.
5 +
6 +Model-neutral: any row-parallel projection whose output is reduced across TP
7 +(Kimi-K3 ``o_proj``/``down_proj``, DeepSeek-V4.1 ``wo_b``) can bind to the
8 +process-wide workspace through ``maybe_init_gemm_rs_ar`` + ``GemmRsAr.apply``.
9 +"""
10
11 # Based on CUTLASS's Blackwell distributed GEMM-RS example at dcf215a.
12 # See https://github.com/NVIDIA/cutlass/issues/3117 for memory semantics.
13 +# The MXFP8 path follows the block-scaled SM100 GEMM template from
14 +# https://github.com/gau-nernst/gn-kernels (sm100_mm_mxfp8.py).
15
16 from functools import cache
17 +from typing import TYPE_CHECKING
18
19 import cutlass
20 import torch
21 import torch.distributed._symmetric_memory as symm_mem
22 from cuda.bindings.driver import CUstream
15 -from cutlass import BFloat16, Int32, Int64, Uint16, cute, utils
23 +from cutlass import (
24 + BFloat16,
25 + Float8E4M3FN,
26 + Float8E8M0FNU,
27 + Int32,
28 + Int64,
29 + Uint16,
30 + Uint64,
31 + cute,
32 + utils,
33 +)
34 from cutlass._mlir import ir
35 from cutlass._mlir.dialects import llvm, nvvm, vector
36 from cutlass.cute.nvgpu import cpasync, tcgen05
@@ -24,6 +42,10 @@ from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem
42 from vllm.distributed import get_tp_group
43 from vllm.logger import init_logger
44 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
45 +from vllm.platforms import current_platform
46 +
47 +if TYPE_CHECKING:
48 + from vllm.config import VllmConfig
49
50 logger = init_logger(__name__)
51
@@ -88,24 +110,31 @@ def multimem_st_16B(dst: cute.Tensor, value: cute.Tensor, *, loc=None, ip=None)
110 )
111
112
91 -class Sm100GemmRsArBF16:
113 +class Sm100GemmRsAr:
114 def __init__(
115 self,
116 rank: int,
117 num_ranks: int,
96 - BN: int = 128,
118 cta_group: int = 1,
119 all_reduce: bool = False,
120 + dtype: str = "bf16",
121 ) -> None:
122 + assert dtype in ("bf16", "mxfp8")
123 self.rank = rank
124 self.num_ranks = num_ranks
102 - BM, BK = 128, 64
125 + BM = 128
126 + BN = 128
127 + BK = {"bf16": 64, "mxfp8": 128}[dtype] # 128B
128 self.cta_tile = (BM, BN, BK)
129 self.cta_group = cta_group
130 self.all_reduce = all_reduce
131 + self.dtype = dtype
132
133 smem_bytes = get_smem_capacity_in_bytes()
108 - self.stage_size = (BM + (BN // cta_group)) * BK * 2
134 + if dtype == "bf16":
135 + self.stage_size = (BM + (BN // cta_group)) * BK * 2
136 + elif dtype == "mxfp8":
137 + self.stage_size = (BM + (BN // cta_group)) * BK + (BM + BN) * (BK // 32)
138 self.num_stages = smem_bytes // self.stage_size
139
140 @cute.jit
@@ -124,11 +153,36 @@ class Sm100GemmRsArBF16:
153 layout = cute.make_composed_layout(swizzle_128b, 0, layout)
154 return cpasync.make_tiled_tma_atom(tma_op, tensor, layout, (BM, BK))
155
156 + @cute.jit
157 + def prepare_sf_tma(
158 + self,
159 + SF: cute.Tensor,
160 + R: Int32,
161 + Ksf: Int32,
162 + BR: cutlass.Constexpr,
163 + BKsf: cutlass.Constexpr,
164 + ) -> cpasync.TmaInfo:
165 + tma_group = (
166 + tcgen05.CtaGroup.TWO if self.cta_group == 2 else tcgen05.CtaGroup.ONE
167 + )
168 + tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=tma_group)
169 + # NVIDIA SF gmem layout: an [R, Ksf] scale grid padded to R % 128 == 0
170 + # and permuted as [R/128, Ksf/4, 32, 4, 4]. The trailing [32, 4, 4]
171 + # atom is the opaque 512B block tcgen05 expects, so treat it as bytes.
172 + # The Int64 recast keeps the TMA boxDim <= 256.
173 + g_layout = cute.make_layout((Ksf * 128, cute.ceil_div(R, 128)))
174 + SF_i64 = cute.recast_tensor(cute.make_tensor(SF.iterator, g_layout), Int64)
175 + tiler = (128 * BKsf // 8, BR // 128)
176 + layout = cute.make_layout((*tiler, self.num_stages))
177 + return cpasync.make_tiled_tma_atom(tma_op, SF_i64, layout, tiler)
178 +
179 @cute.jit
180 def __call__(
181 self,
182 A: cute.Tensor,
183 B: cute.Tensor,
184 + SFA: cute.Tensor,
185 + SFB: cute.Tensor,
186 partial_uc: cute.Tensor,
187 partial_mc_ptr: cute.Pointer,
188 output: cute.Tensor,
@@ -138,10 +192,17 @@ class Sm100GemmRsArBF16:
192 grid_size: Int32,
193 stream: CUstream,
194 ) -> None:
195 + M, K = A.shape
196 N = B.shape[0]
197 BM, BN, BK = self.cta_tile
198 A_tma = self.prepare_tma(A, BM, BK)
199 B_tma = self.prepare_tma(B, BN // self.cta_group, BK)
200 + if cutlass.const_expr(self.dtype == "mxfp8"):
201 + SFA_tma = self.prepare_sf_tma(SFA, M, K // 32, BM, BK // 32)
202 + SFB_tma = self.prepare_sf_tma(SFB, N, K // 32, BN, BK // 32)
203 + else:
204 + SFA_tma = None
205 + SFB_tma = None
206 padded_M = partial_uc.shape[0]
207 partial_mc = cute.make_tensor(
208 partial_mc_ptr,
@@ -158,6 +219,8 @@ class Sm100GemmRsArBF16:
219 self.kernel(
220 A_tma,
221 B_tma,
222 + SFA_tma,
223 + SFB_tma,
224 partial_uc,
225 partial_mc,
226 output,
@@ -171,6 +234,8 @@ class Sm100GemmRsArBF16:
234 self,
235 A_tma: cpasync.TmaInfo,
236 B_tma: cpasync.TmaInfo,
237 + SFA_tma: cpasync.TmaInfo,
238 + SFB_tma: cpasync.TmaInfo,
239 partial_uc: cute.Tensor,
240 partial_mc: cute.Tensor,
241 output: cute.Tensor,
@@ -190,25 +255,33 @@ class Sm100GemmRsArBF16:
255
256 is_2cta = cta_group == 2
257 cta_rank = raw_bid % self.cta_group
193 - num_tmem_stages = 512 // BN
258 +
259 + # MXFP8 reserves the last 16 tmem columns (496-512) for the A/B scale
260 + # factors.
261 + num_tmem_stages = (496 if self.dtype == "mxfp8" else 512) // BN
262
263 smem = utils.SmemAllocator()
264 sA = smem.allocate_tensor(
197 - BFloat16,
265 + A_tma.atom.value_type,
266 A_tma.smem_layout.outer,
267 byte_alignment=128,
268 swizzle=A_tma.smem_layout.inner,
269 )
270 sB = smem.allocate_tensor(
203 - BFloat16,
271 + B_tma.atom.value_type,
272 B_tma.smem_layout.outer,
273 byte_alignment=128,
274 swizzle=B_tma.smem_layout.inner,
275 )
276 + if cutlass.const_expr(self.dtype == "mxfp8"):
277 + sSFA = smem.allocate_tensor(Int64, SFA_tma.smem_layout, byte_alignment=128)
278 + sSFB = smem.allocate_tensor(Int64, SFB_tma.smem_layout, byte_alignment=128)
279 +
280 tma_full_mbar = smem.allocate_array(Int64, num_stages)
281 tma_empty_mbar = smem.allocate_array(Int64, num_stages)
282 tmem_full_mbar = smem.allocate_array(Int64, num_tmem_stages)
283 tmem_empty_mbar = smem.allocate_array(Int64, num_tmem_stages)
284 +
285 taddr = smem.allocate(Int32, 4)
286
287 # Named barriers
@@ -238,6 +311,9 @@ class Sm100GemmRsArBF16:
311 elif warp_id == 1:
312 cpasync.prefetch_descriptor(A_tma.atom)
313 cpasync.prefetch_descriptor(B_tma.atom)
314 + if cutlass.const_expr(self.dtype == "mxfp8"):
315 + cpasync.prefetch_descriptor(SFA_tma.atom)
316 + cpasync.prefetch_descriptor(SFB_tma.atom)
317
318 if cutlass.const_expr(is_2cta):
319 cute.arch.cluster_arrive_relaxed()
@@ -261,12 +337,22 @@ class Sm100GemmRsArBF16:
337 # [(BM, BK), (M/BM, K/BK)]
338 gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK))
339 gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN // cta_group, BK))
340 + if cutlass.const_expr(self.dtype == "mxfp8"):
341 + BKsf = BK // 32
342 + gSFA_tiles = cute.zipped_divide(
343 + SFA_tma.tma_tensor, (128 * BKsf // 8, BM // 128)
344 + )
345 + gSFB_tiles = cute.zipped_divide(
346 + SFB_tma.tma_tensor, (128 * BKsf // 8, BN // 128)
347 + )
348
349 for bid in range(raw_bid, total_tiles, num_bids):
350 bid_m = bid % grid_m
351 bid_n = bid // grid_m
352 if cutlass.const_expr(cta_group == 2):
269 - bid_n = bid_n * cta_group + cta_rank
353 + bid_n_b = bid_n * cta_group + cta_rank
354 + else:
355 + bid_n_b = bid_n
356
357 for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1):
358 mbar = tma_full_mbar_ + tma_stage
@@ -282,10 +368,23 @@ class Sm100GemmRsArBF16:
368 )
369 simple_tma_copy(
370 B_tma.atom,
285 - gB_tiles[None, (bid_n, iter_k)],
371 + gB_tiles[None, (bid_n_b, iter_k)],
372 sB[None, None, tma_stage],
373 mbar,
374 )
375 + if cutlass.const_expr(self.dtype == "mxfp8"):
376 + simple_tma_copy(
377 + SFA_tma.atom,
378 + gSFA_tiles[None, (iter_k, bid_m)],
379 + sSFA[None, None, tma_stage],
380 + mbar,
381 + )
382 + simple_tma_copy(
383 + SFB_tma.atom,
384 + gSFB_tiles[None, (iter_k, bid_n)],
385 + sSFB[None, None, tma_stage],
386 + mbar,
387 + )
388
389 tma_stage = (tma_stage + 1) % num_stages
390 if tma_stage == 0:
@@ -303,8 +402,15 @@ class Sm100GemmRsArBF16:
402
403 MMA_M = BM * cta_group
404 MMA_N = BN
306 - idesc = _tcgen05.make_bf16_idesc(MMA_M, MMA_N)
405 sdesc = _tcgen05.make_sdesc_128B_swizzle(0)
406 + if cutlass.const_expr(self.dtype == "mxfp8"):
407 + # SF smem descriptor: no swizzle, SBO = 8 * 16 bytes.
408 + sdesc_sf = Uint64((8 << 32) | (1 << 46))
409 + sfa_tmem = 496
410 + sfb_tmem = sfa_tmem + 4
411 + idesc = _tcgen05.make_mxfp8_idesc(MMA_M, MMA_N)
412 + else:
413 + idesc = _tcgen05.make_bf16_idesc(MMA_M, MMA_N)
414 multicast_mask = Uint16((1 << self.cta_group) - 1)
415
416 for bid in range(raw_bid, total_tiles, num_bids):
@@ -319,17 +425,54 @@ class Sm100GemmRsArBF16:
425 b_addr = sB[None, None, tma_stage].iterator.toint()
426 a_desc = sdesc | (a_addr >> 4)
427 b_desc = sdesc | (b_addr >> 4)
428 + if cutlass.const_expr(self.dtype == "mxfp8"):
429 + sfa_desc = sdesc_sf | (
430 + sSFA[None, None, tma_stage].iterator.toint() >> 4
431 + )
432 + sfb_desc = sdesc_sf | (
433 + sSFB[None, None, tma_stage].iterator.toint() >> 4
434 + )
435
436 cute.arch.mbarrier_wait(
437 tma_full_mbar + tma_stage, tma_full_parity
438 )
439 _tcgen05.fence_after_thread_sync()
440
328 - for mma_k in cutlass.range_constexpr(BK // 16):
329 - enable_d = iter_k > 0 or mma_k > 0
330 - _tcgen05.mma_f16(
331 - d_tmem, a_desc, b_desc, idesc, enable_d, cta_group
441 + # load SF from smem->tmem
442 + if cutlass.const_expr(self.dtype == "mxfp8"):
443 + _tcgen05.cp(
444 + sfa_tmem, sfa_desc, "32x128b", "warpx4", cta_group
445 )
446 + for j in cutlass.range_constexpr(BN // 128):
447 + _tcgen05.cp(
448 + sfb_tmem + j * 4,
449 + sfb_desc,
450 + "32x128b",
451 + "warpx4",
452 + cta_group,
453 + )
454 + sfb_desc += (128 * 4) >> 4
455 +
456 + # 4 = 128B / 32B
457 + for mma_k in cutlass.range_constexpr(4):
458 + enable_d = iter_k > 0 or mma_k > 0
459 + if cutlass.const_expr(self.dtype == "mxfp8"):
460 + # select SF IDs
461 + idesc_k = idesc + ((mma_k << 4) | (mma_k << 29))
462 + _tcgen05.mma_mxfp8(
463 + d_tmem,
464 + a_desc,
465 + b_desc,
466 + idesc_k,
467 + sfa_tmem,
468 + sfb_tmem,
469 + enable_d,
470 + cta_group,
471 + )
472 + else:
473 + _tcgen05.mma_f16(
474 + d_tmem, a_desc, b_desc, idesc, enable_d, cta_group
475 + )
476 a_desc += 32 >> 4
477 b_desc += 32 >> 4
478 _tcgen05.commit(
@@ -638,9 +781,9 @@ class Sm100GemmRsArBF16:
781 def compile(
782 rank: int,
783 num_ranks: int,
641 - BN: int,
784 cta_group: int,
785 all_reduce: bool = False,
786 + dtype: str = "bf16",
787 ):
788 M = cute.sym_int()
789 padded_M = cute.sym_int()
@@ -649,12 +792,37 @@ class Sm100GemmRsArBF16:
792 local_M = cute.sym_int()
793 num_flags = cute.sym_int()
794
795 + ab_dtype = Float8E4M3FN if dtype == "mxfp8" else BFloat16
796 + # Row stride must keep rows 16B-aligned: 8 BF16 or 16 FP8 elements.
797 + row_div = 16 if dtype == "mxfp8" else 8
798 A = make_fake_tensor(
653 - BFloat16, (M, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16
799 + ab_dtype,
800 + (M, K),
801 + (cute.sym_int64(divisibility=row_div), 1),
802 + assumed_align=16,
803 )
804 B = make_fake_tensor(
656 - BFloat16, (N, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16
805 + ab_dtype,
806 + (N, K),
807 + (cute.sym_int64(divisibility=row_div), 1),
808 + assumed_align=16,
809 )
810 + if dtype == "mxfp8":
811 + SFA = make_fake_tensor(
812 + Float8E8M0FNU,
813 + (cute.sym_int(divisibility=512),),
814 + (1,),
815 + assumed_align=16,
816 + )
817 + SFB = make_fake_tensor(
818 + Float8E8M0FNU,
819 + (cute.sym_int(divisibility=512),),
820 + (1,),
821 + assumed_align=16,
822 + )
823 + else:
824 + SFA = None
825 + SFB = None
826 partial = make_fake_tensor(
827 BFloat16,
828 (padded_M, N),
@@ -677,13 +845,15 @@ class Sm100GemmRsArBF16:
845 peer_flag_ptr = nullptr(Int64, cute.AddressSpace.gmem, assumed_align=8)
846
847 stream = make_fake_stream(use_tvm_ffi_env_stream=True)
680 - kernel = Sm100GemmRsArBF16(
681 - rank, num_ranks, BN, cta_group, all_reduce=all_reduce
848 + kernel = Sm100GemmRsAr(
849 + rank, num_ranks, cta_group, all_reduce=all_reduce, dtype=dtype
850 )
851 return cute.compile(
852 kernel,
853 A,
854 B,
855 + SFA,
856 + SFB,
857 partial,
858 partial_mc_ptr,
859 output,
@@ -697,7 +867,7 @@ class Sm100GemmRsArBF16:
867
868
869 class GemmRsAr:
700 - """Own the symmetric workspace for Kimi-K3 GEMM-RS/AR launches.
870 + """Own the symmetric workspace for GEMM-RS/AR launches.
871
872 All TP ranks must belong to one NVLink domain for multimem instructions.
873
@@ -727,6 +897,7 @@ class GemmRsAr:
897 self.N = N
898 self.device = device
899 self.all_reduce = all_reduce
900 + self.dtypes: set[str] = set()
901
902 self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device)
903 self.partial_handle = symm_mem.rendezvous(self.partial, group)
@@ -766,21 +937,40 @@ class GemmRsAr:
937 tp_group.barrier()
938
939 def can_run(self, linear: LinearBase) -> bool:
769 - # Validate projection-invariant requirements once during model init.
770 - # only supports BF16 for now
771 - if not isinstance(linear.quant_method, UnquantizedLinearMethod):
772 - return False
940 + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
941 + FlashInferCutedslMxfp8LinearKernel,
942 + FlashInferCutlassMxfp8LinearKernel,
943 + )
944 +
945 + # Called before weight loading; online quantization starts on meta.
946 + if isinstance(linear.quant_method, UnquantizedLinearMethod):
947 + dtype, k_alignment = "bf16", 64
948 + if linear.weight.dtype != torch.bfloat16:
949 + return False
950 + else:
951 + method = getattr(linear, "scheme", linear.quant_method)
952 + kernel = getattr(method, "kernel", None)
953 + if not isinstance(
954 + kernel,
955 + (
956 + FlashInferCutedslMxfp8LinearKernel,
957 + FlashInferCutlassMxfp8LinearKernel,
958 + ),
959 + ):
960 + return False
961 + dtype, k_alignment = "mxfp8", 128
962 +
963 w = linear.weight
774 - if w.ndim != 2:
964 + if (
965 + w.ndim != 2
966 + or w.shape[0] != self.N
967 + or w.shape[1] % k_alignment != 0
968 + or (w.device != self.device and not w.is_meta)
969 + or not w.is_contiguous()
970 + ):
971 return False
776 - K = w.shape[1]
777 - return (
778 - w.shape == (self.N, K)
779 - and K % 64 == 0
780 - and w.dtype == torch.bfloat16
781 - and w.device == self.device
782 - and w.is_contiguous()
783 - )
972 + self.dtypes.add(dtype)
973 + return True
974
975 def warn_incompatible_projection(self) -> None:
976 logger.warning_once(
@@ -794,18 +984,51 @@ class GemmRsAr:
984 # supported but faster on the existing LL path.
985 return x.shape[0] >= 128
986
797 - def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
987 + def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor:
988 + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
989 + FlashInferCutedslMxfp8LinearKernel,
990 + )
991 +
992 + method = getattr(linear, "scheme", linear.quant_method)
993 + w = linear.weight
994 + if isinstance(
995 + getattr(method, "kernel", None), FlashInferCutedslMxfp8LinearKernel
996 + ):
997 + # This backend stores a column-major [K, N] view after loading.
998 + w = w.t()
999 + w_sf = getattr(linear, "weight_scale", None)
1000 +
1001 assert x.ndim == 2
1002 M, K = x.shape
1003 assert 0 < M <= self.max_M
801 - assert w.shape == (self.N, K) and K % 64 == 0
802 - assert w.dtype == torch.bfloat16
803 - assert w.device == self.device
804 - assert w.is_contiguous()
1004 assert x.dtype == torch.bfloat16
1005 assert x.device == self.device
1006 assert x.is_contiguous()
808 - N = w.shape[0]
1007 + assert w.device == self.device
1008 + N = self.N
1009 + dtype = "mxfp8" if w.dtype == torch.float8_e4m3fn else "bf16"
1010 + x_sf = None
1011 + if dtype == "mxfp8":
1012 + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1013 + mxfp8_e4m3_quantize,
1014 + )
1015 +
1016 + assert w_sf is not None
1017 + assert w_sf.dtype == torch.uint8 and w_sf.device == self.device
1018 + assert w_sf.ndim == 1 and w_sf.is_contiguous()
1019 + assert w_sf.numel() == N * (K // 32)
1020 + assert w.shape == (N, K) and K % 128 == 0
1021 + assert w.is_contiguous()
1022 + x, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
1023 + # The kernel declares the scales as E8M0; the stored/quantized
1024 + # tensors carry the same bits as uint8.
1025 + x_sf = x_sf.view(torch.float8_e8m0fnu)
1026 + w_sf = w_sf.view(torch.float8_e8m0fnu)
1027 + else:
1028 + assert w_sf is None
1029 + assert w.shape == (N, K) and K % 64 == 0
1030 + assert w.dtype == torch.bfloat16
1031 + assert w.is_contiguous()
1032 padded_M = (M + self.world_size - 1) // self.world_size
1033 padded_M *= self.world_size
1034 local_M = padded_M // self.world_size
@@ -814,10 +1037,8 @@ class GemmRsAr:
1037 # Avoid padding small odd grids; 2-CTA wins consistently for M >= 1024.
1038 cta_group = 2 if M >= 1024 or grid_m % 2 == 0 else 1
1039 grid_m = (grid_m + cta_group - 1) // cta_group * cta_group
817 - BN = 256 if M * K >= 24 * 1024 * 1024 else 128
818 - assert N % BN == 0
1040
820 - num_tiles = grid_m * (N // BN)
1041 + num_tiles = grid_m * (N // 128)
1042 num_ctas = min(num_tiles, self.num_sms)
1043 num_ctas = num_ctas // cta_group * cta_group
1044 assert self.flags.numel() >= num_tiles + num_ctas
@@ -825,16 +1046,26 @@ class GemmRsAr:
1046 output = None
1047 if not self.all_reduce:
1048 output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device)
828 - compiled = Sm100GemmRsArBF16.compile(
1049 + # The kernel never touches rows at or beyond M. Callers that hand
1050 + # in an unpadded M (DeepSeek-V4.1 slices the SP all-gather back to
1051 + # the real token count) expect the same zero padding rows that
1052 + # ``sp_reduce_scatter`` produces, so clear them here. This is a
1053 + # no-op whenever M is already a multiple of the TP size.
1054 + valid_rows = min(max(M - self.rank * local_M, 0), local_M)
1055 + if valid_rows < local_M:
1056 + output[valid_rows:].zero_()
1057 + compiled = Sm100GemmRsAr.compile(
1058 self.rank,
1059 self.world_size,
831 - BN,
1060 cta_group,
1061 self.all_reduce,
1062 + dtype=dtype,
1063 )
1064 compiled(
1065 x,
1066 w,
1067 + x_sf,
1068 + w_sf,
1069 self.partial[:padded_M],
1070 self.partial_mc_ptr,
1071 output,
@@ -871,22 +1102,74 @@ def init_gemm_rs_ar(max_M: int, N: int, *, all_reduce: bool = False) -> None:
1102 _gemm_rs_ar = GemmRsAr(max_M=max_M, N=N, all_reduce=all_reduce)
1103
1104
1105 +def maybe_init_gemm_rs_ar(
1106 + vllm_config: "VllmConfig", *, N: int, all_reduce: bool
1107 +) -> bool:
1108 + """Collectively initialize the mode-bound GEMM-RS/AR state if supported.
1109 +
1110 + Callers gate on their own feature flag first; this checks the worker
1111 + topology (SM100-family CUDA, BF16 model dtype, no ubatching, a TP size in
1112 + 2-16 that divides 128) and the NVLink multicast rendezvous. Returns whether
1113 + projections may bind to the kernel through ``get_gemm_rs_ar().can_run``.
1114 + """
1115 + mode = "GEMM-AR" if all_reduce else "GEMM-RS"
1116 + parallel_config = vllm_config.parallel_config
1117 + tp_size = parallel_config.tensor_parallel_size
1118 + if parallel_config.use_ubatching:
1119 + reason = "ubatching is enabled"
1120 + elif vllm_config.model_config.dtype != torch.bfloat16:
1121 + reason = "the model dtype is not BF16"
1122 + elif not current_platform.is_cuda():
1123 + reason = "the device is not CUDA"
1124 + elif not current_platform.is_device_capability_family(100):
1125 + reason = "the device is not SM100-family"
1126 + elif not 1 < tp_size <= 16:
1127 + reason = "TP size is not in the supported range 2-16"
1128 + elif 128 % tp_size != 0:
1129 + reason = "TP size does not divide 128"
1130 + elif N % 128 != 0:
1131 + reason = f"the output width {N} is not a multiple of 128"
1132 + else:
1133 + reason = None
1134 +
1135 + if reason is not None:
1136 + logger.warning_once("%s is disabled because %s.", mode, reason)
1137 + return False
1138 +
1139 + try:
1140 + init_gemm_rs_ar(
1141 + max_M=vllm_config.scheduler_config.max_num_batched_tokens,
1142 + N=N,
1143 + all_reduce=all_reduce,
1144 + )
1145 + except RuntimeError as e:
1146 + logger.warning_once(
1147 + "%s is disabled because initialization failed: %s. This may mean "
1148 + "the TP ranks do not share one NVLink domain.",
1149 + mode,
1150 + e,
1151 + )
1152 + return False
1153 + logger.info_once("%s is enabled.", mode)
1154 + return True
1155 +
1156 +
1157 def warmup_gemm_rs_ar() -> int:
1158 """Compile every reachable dispatch for the initialized GEMM-RS/AR mode."""
1159 # Initialization can be disabled or fail when multicast is unavailable.
1160 if _gemm_rs_ar is None:
1161 return 0
879 - # Keep these profiles in sync with the dispatch in GemmRsAr.__call__.
880 - profiles = ((128, 1), (128, 2), (256, 2))
881 - for BN, cta_group in profiles:
882 - Sm100GemmRsArBF16.compile(
883 - _gemm_rs_ar.rank,
884 - _gemm_rs_ar.world_size,
885 - BN,
886 - cta_group,
887 - _gemm_rs_ar.all_reduce,
888 - )
889 - return len(profiles)
1162 +
1163 + for dtype in sorted(_gemm_rs_ar.dtypes):
1164 + for cta_group in (1, 2):
1165 + Sm100GemmRsAr.compile(
1166 + _gemm_rs_ar.rank,
1167 + _gemm_rs_ar.world_size,
1168 + cta_group,
1169 + _gemm_rs_ar.all_reduce,
1170 + dtype=dtype,
1171 + )
1172 + return 2 * len(_gemm_rs_ar.dtypes)
1173
1174
1175 def get_gemm_rs_ar() -> GemmRsAr:
vllm/model_executor/warmup/kernel_warmup.py
+8 -7
@@ -136,16 +136,17 @@ def _warmup_bf16x3_router_gemm(
136 logger.info_once("Warmed up BF16x3 router GEMM configs: %s.", configs)
137
138
139 -def _warmup_kimi_k3_gemm_rs_ar() -> None:
140 - # Kimi-K3 model construction imports this module only when GEMM-RS/AR is
141 - # enabled and initializes its singleton before kernel_warmup runs. Avoid
142 - # importing it here so other models do not compile the RS/AR variants.
143 - module = sys.modules.get("vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar")
139 +def _warmup_gemm_rs_ar() -> None:
140 + # Model construction (Kimi-K3, DeepSeek-V4.1) imports this module only
141 + # when GEMM-RS/AR is enabled and initializes its singleton before
142 + # kernel_warmup runs. Avoid importing it here so other models do not
143 + # compile the RS/AR variants.
144 + module = sys.modules.get("vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar")
145 if module is None:
146 return
147 compiled = module.warmup_gemm_rs_ar()
148 if compiled:
148 - logger.info_once("Warmed up %d Kimi-K3 GEMM-RS/AR variants.", compiled)
149 + logger.info_once("Warmed up %d GEMM-RS/AR variants.", compiled)
150
151
152 def _autotune_kimi_k3_kda_qkvg(model: torch.nn.Module) -> None:
@@ -208,7 +209,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
209 if current_platform.has_device_capability(90):
210 _warmup_ll_bf16_router_gemm(worker.get_model())
211
211 - _warmup_kimi_k3_gemm_rs_ar()
212 + _warmup_gemm_rs_ar()
213
214 if worker.vllm_config.kernel_config.enable_cutedsl_warmup:
215 # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
vllm/models/deepseek_v4/nvidia/ops/o_proj.py
+6 -2
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 +
5 import torch
6 import torch.nn as nn
7
@@ -31,7 +33,7 @@ def deep_gemm_fp8_o_proj(
33 positions: torch.Tensor,
34 cos_sin_cache: torch.Tensor,
35 wo_a: nn.Module,
34 - wo_b: nn.Module,
36 + wo_b: Callable[[torch.Tensor], torch.Tensor],
37 *,
38 n_groups: int,
39 heads_per_group: int,
@@ -44,7 +46,9 @@ def deep_gemm_fp8_o_proj(
46 """O projection: inverse RoPE + grouped wo_a + wo_b.
47
48 Shared by the FlashMLA and FlashInfer CUDA backends. The attention
47 - layer selects the recipe at initialization.
49 + layer selects the recipe at initialization. ``wo_b`` is any callable over
50 + the flattened ``z``: the projection module itself, or a wrapper that also
51 + reduce-scatters its output (DeepSeek-V4.1 GEMM-RS).
52 """
53 use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn
54 o_proj_input, o_scale = fused_inv_rope_fp8_quant(
vllm/models/deepseek_v41/attention.py
+40
@@ -26,6 +26,7 @@ from vllm.model_executor.layers.linear import (
26 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
27 from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer
28 from vllm.models.common.ops import fused_q_kv_rmsnorm
29 +from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter
30 from vllm.models.deepseek_v41.common.ops import (
31 MXFP4_BLOCK_SIZE,
32 fused_indexer_q_rope_quant,
@@ -33,6 +34,7 @@ from vllm.models.deepseek_v41.common.ops import (
34 )
35
36 if TYPE_CHECKING:
37 + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
38 from vllm.v1.attention.backends.mla.sparse_swa import (
39 DeepseekSparseSWAMetadata,
40 )
@@ -389,6 +391,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
391 return_bias=False,
392 prefix=f"{prefix}.wo_b",
393 )
394 + # Set by ``bind_gemm_rs`` when the decoder layer runs sequence
395 + # parallel and the fused GEMM + reduce-scatter kernel accepts wo_b.
396 + self.gemm_rs: GemmRsAr | None = None
397
398 # Initialize rotary embedding before the indexer/compressor consume it.
399 self.rotary_emb = build_deepseek_v4_rope(
@@ -672,6 +677,41 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
677 )
678 return self._o_proj(attn_out, positions)
679
680 + def bind_gemm_rs(self) -> None:
681 + """Fuse ``wo_b`` with the sequence-parallel TP reduce-scatter.
682 +
683 + The decoder layer calls this after the model initialized the
684 + process-wide GEMM-RS workspace and before weights load: the
685 + eligibility check inspects the projection's linear kernel and weight
686 + shape, which online quantization may later re-layout.
687 + """
688 + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
689 + get_gemm_rs_ar,
690 + )
691 +
692 + # The layer owns the reduction under sequence parallel.
693 + assert not self.wo_b.reduce_results
694 + gemm_rs = get_gemm_rs_ar()
695 + if gemm_rs.can_run(self.wo_b):
696 + self.gemm_rs = gemm_rs
697 + else:
698 + gemm_rs.warn_incompatible_projection()
699 +
700 + def _wo_b_proj(self, z: torch.Tensor) -> torch.Tensor:
701 + """Apply ``wo_b``; with GEMM-RS bound, also reduce-scatter the result.
702 +
703 + Every ``_o_proj`` implementation projects through this so the decoder
704 + layer sees one contract: when ``gemm_rs`` is bound the output is
705 + already the local sequence-parallel shard, otherwise it is the
706 + unreduced TP partial.
707 + """
708 + if self.gemm_rs is None:
709 + return self.wo_b(z)
710 + if self.gemm_rs.should_run(z):
711 + return self.gemm_rs.apply(z, self.wo_b)
712 + # Small batches stay on the unfused path, which is faster there.
713 + return sp_reduce_scatter(self.wo_b(z))
714 +
715 def _alloc_attn_out(
716 self, num_tokens: int, hidden_states: torch.Tensor
717 ) -> "torch.Tensor | QuantizedActivation":
vllm/models/deepseek_v41/nvidia/dspark.py
+7
@@ -58,6 +58,7 @@ from .model import (
58 _linear_scale_param_name,
59 _use_sequence_parallel,
60 make_deepseek_v4_expert_params_mapping,
61 + maybe_init_gemm_rs,
62 prepare_mega_gate_routing_metadata,
63 )
64
@@ -113,12 +114,18 @@ class DSparkDeepseekV4Model(nn.Module):
114 )
115
116 current_vllm_config = get_current_vllm_config()
117 + # The target model already holds the GEMM-RS workspace (same hidden
118 + # size, same TP group); this only re-checks and binds the draft layers.
119 + run_gemm_rs = maybe_init_gemm_rs(
120 + current_vllm_config, self.use_sequence_parallel
121 + )
122 self.layers = nn.ModuleList(
123 [
124 DeepseekV4DecoderLayer(
125 current_vllm_config,
126 prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
127 topk_indices_buffer=self.topk_indices_buffer,
128 + run_gemm_rs=run_gemm_rs,
129 )
130 for i in range(self.num_dspark_layers)
131 ]
vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
+1 -1
@@ -225,7 +225,7 @@ class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention):
225 z,
226 recipe=self._einsum_recipe,
227 )
228 - return self.wo_b(z.flatten(1))
228 + return self._wo_b_proj(z.flatten(1))
229
230 # ---- weights -----------------------------------------------------------
231
vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
+2 -2
@@ -248,7 +248,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
248 positions,
249 self.rotary_emb.cos_sin_cache,
250 self.wo_a,
251 - self.wo_b,
251 + self._wo_b_proj,
252 n_groups=self.n_local_groups,
253 heads_per_group=self.n_local_heads // self.n_local_groups,
254 nope_dim=self.nope_head_dim,
@@ -627,7 +627,7 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
627 positions,
628 self.rotary_emb.cos_sin_cache,
629 self.wo_a,
630 - self.wo_b,
630 + self._wo_b_proj,
631 n_groups=self.n_local_groups,
632 heads_per_group=self.n_local_heads // self.n_local_groups,
633 nope_dim=self.nope_head_dim,
vllm/models/deepseek_v41/nvidia/flashmla.py
+1 -1
@@ -65,7 +65,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
65 positions,
66 self.rotary_emb.cos_sin_cache,
67 self.wo_a,
68 - self.wo_b,
68 + self._wo_b_proj,
69 n_groups=self.n_local_groups,
70 heads_per_group=self.n_local_heads // self.n_local_groups,
71 nope_dim=self.nope_head_dim,
vllm/models/deepseek_v41/nvidia/model.py
+37 -1
@@ -193,6 +193,30 @@ def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
193 )
194
195
196 +def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool:
197 + """Set up the fused ``wo_b`` GEMM + reduce-scatter when opted in.
198 +
199 + Sequence parallel is the only topology where ``wo_b`` ends in a
200 + reduce-scatter, so the kernel is bound to RS mode. The workspace is a
201 + process-wide singleton shared with any other model code in this worker
202 + (the DSpark drafter reuses it), and its NVLink multicast rendezvous is
203 + collective, so every TP rank must take the same decision here.
204 + """
205 + if not (use_sequence_parallel and envs.VLLM_ENABLE_GEMM_RS):
206 + return False
207 +
208 + # The kernel module pulls in cute_dsl, so import it only once opted in.
209 + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
210 + maybe_init_gemm_rs_ar,
211 + )
212 +
213 + hidden_size = vllm_config.model_config.hf_config.hidden_size
214 + if not maybe_init_gemm_rs_ar(vllm_config, N=hidden_size, all_reduce=False):
215 + return False
216 + logger.info_once("To disable DeepSeek-V4.1 GEMM-RS, set VLLM_ENABLE_GEMM_RS=0.")
217 + return True
218 +
219 +
220 class DeepseekV4DecoderLayer(nn.Module):
221 def __init__(
222 self,
@@ -202,6 +226,7 @@ class DeepseekV4DecoderLayer(nn.Module):
226 aux_stream_list: list[torch.cuda.Stream] | None = None,
227 candidate_block_buffer: torch.Tensor | None = None,
228 engram_layout: EngramLayout | None = None,
229 + run_gemm_rs: bool = False,
230 mhc_stream: torch.cuda.Stream | None = None,
231 fuse_mhc_all_reduce: bool = False,
232 ):
@@ -236,6 +261,10 @@ class DeepseekV4DecoderLayer(nn.Module):
261 )
262 if self.use_sequence_parallel or fuse_mhc_all_reduce:
263 self.attn.wo_b.reduce_results = False
264 + if run_gemm_rs:
265 + # Binds only when wo_b's kernel and shape qualify; otherwise
266 + # forward keeps the separate reduce-scatter below.
267 + self.attn.bind_gemm_rs()
268 self.ffn = DeepseekV4MoE(
269 vllm_config,
270 prefix=f"{prefix}.ffn",
@@ -492,7 +521,9 @@ class DeepseekV4DecoderLayer(nn.Module):
521 x = sp_all_gather(x)[: positions.shape[0]]
522
523 x = self.attn(positions, x, None)
495 - if self.use_sequence_parallel:
524 + # With GEMM-RS bound, the attention output is already this rank's
525 + # sequence-parallel shard (see DeepseekV4Attention._wo_b_proj).
526 + if self.use_sequence_parallel and self.attn.gemm_rs is None:
527 x = sp_reduce_scatter(x)
528
529 if mhc_stream is not None:
@@ -597,6 +628,10 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
628 ):
629 engram_config.dp_shared_memory = can_share_engram_tables(self.engram_layout)
630
631 + # GEMM-RS uses NCCL symmetric-memory multicast, which requires all TP
632 + # ranks to belong to one NVLink domain. Collective: run before layers.
633 + self.run_gemm_rs = maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel)
634 +
635 self.start_layer, self.end_layer, self.layers = make_layers(
636 config.num_hidden_layers,
637 lambda prefix: DeepseekV4DecoderLayer(
@@ -606,6 +641,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
641 aux_stream_list=aux_stream_list,
642 candidate_block_buffer=self.candidate_block_buffer,
643 engram_layout=self.engram_layout,
644 + run_gemm_rs=self.run_gemm_rs,
645 mhc_stream=mhc_stream,
646 fuse_mhc_all_reduce=self.fuse_mhc_all_reduce,
647 ),
vllm/models/kimi_k3/nvidia/kda.py
+2 -2
@@ -724,7 +724,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
724 )
725 self.gemm_rs_ar = None
726 if run_gemm_rs_ar:
727 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
727 + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
728 get_gemm_rs_ar,
729 )
730
@@ -809,7 +809,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
809 )
810 core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
811 if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(core_attn_out):
812 - return self.gemm_rs_ar(core_attn_out, self.o_proj.weight)
812 + return self.gemm_rs_ar.apply(core_attn_out, self.o_proj)
813 return self.o_proj(core_attn_out)[0]
814
815 @eager_break_during_capture
vllm/models/kimi_k3/nvidia/mla.py
+2 -2
@@ -301,7 +301,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
301 )
302 self.gemm_rs_ar = None
303 if run_gemm_rs_ar:
304 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
304 + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
305 get_gemm_rs_ar,
306 )
307
@@ -631,7 +631,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
631 attn_out = _gate_sigmoid_mul(attn_out, gate)
632
633 if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(attn_out):
634 - return self.gemm_rs_ar(attn_out, self.o_proj.weight)
634 + return self.gemm_rs_ar.apply(attn_out, self.o_proj)
635
636 return self.o_proj(attn_out)[0]
637
vllm/models/kimi_k3/nvidia/model.py
+13 -44
@@ -160,54 +160,23 @@ def maybe_init_gemm_rs_ar(vllm_config: VllmConfig, use_sequence_parallel: bool)
160 # its singleton to exactly one mode.
161 all_reduce = not use_sequence_parallel
162 mode = "GEMM-AR" if all_reduce else "GEMM-RS"
163 - enabled = envs.VLLM_KIMI_K3_GEMM_AR if all_reduce else envs.VLLM_KIMI_K3_GEMM_RS
163 + enabled = envs.VLLM_KIMI_K3_GEMM_AR if all_reduce else envs.VLLM_ENABLE_GEMM_RS
164 if not enabled:
165 return False
166
167 - parallel_config = vllm_config.parallel_config
168 - tp_size = parallel_config.tensor_parallel_size
169 - if parallel_config.use_ubatching:
170 - reason = "ubatching is enabled"
171 - elif vllm_config.model_config.dtype != torch.bfloat16:
172 - reason = "the model dtype is not BF16"
173 - elif not current_platform.is_cuda():
174 - reason = "the device is not CUDA"
175 - elif not current_platform.is_device_capability_family(100):
176 - reason = "the device is not SM100-family"
177 - elif not 1 < tp_size <= 16:
178 - reason = "TP size is not in the supported range 2-16"
179 - elif 128 % tp_size != 0:
180 - reason = "TP size does not divide 128"
181 - else:
182 - reason = None
183 -
184 - if reason is not None:
185 - logger.warning_once("%s is disabled because %s.", mode, reason)
186 - return False
187 -
188 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import init_gemm_rs_ar
167 + # The kernel module pulls in cute_dsl, so import it only once the flag
168 + # asks for it.
169 + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
170 + maybe_init_gemm_rs_ar as maybe_init_shared_gemm_rs_ar,
171 + )
172
173 config = vllm_config.model_config.hf_text_config
191 - try:
192 - init_gemm_rs_ar(
193 - max_M=vllm_config.scheduler_config.max_num_batched_tokens,
194 - N=config.hidden_size,
195 - all_reduce=all_reduce,
196 - )
197 - except RuntimeError as e:
198 - logger.warning_once(
199 - "%s is disabled because initialization failed: %s. This may mean "
200 - "the TP ranks do not share one NVLink domain.",
201 - mode,
202 - e,
203 - )
174 + if not maybe_init_shared_gemm_rs_ar(
175 + vllm_config, N=config.hidden_size, all_reduce=all_reduce
176 + ):
177 return False
205 - if all_reduce:
206 - logger.info_once(
207 - "GEMM-AR is enabled. To disable it, set VLLM_KIMI_K3_GEMM_AR=0."
208 - )
209 - else:
210 - logger.info_once("GEMM-RS is enabled.")
178 + flag = "VLLM_KIMI_K3_GEMM_AR" if all_reduce else "VLLM_ENABLE_GEMM_RS"
179 + logger.info_once("To disable %s, set %s=0.", mode, flag)
180 return True
181
182
@@ -276,7 +245,7 @@ class KimiMLP(nn.Module):
245 not use_sequence_parallel and reduce_results
246 )
247 if use_gemm_rs_ar and run_gemm_rs_ar:
279 - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
248 + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
249 get_gemm_rs_ar,
250 )
251
@@ -309,7 +278,7 @@ class KimiMLP(nn.Module):
278 x = self.act_fn(gate_up)
279
280 if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(x):
312 - return self.gemm_rs_ar(x, self.down_proj.weight)
281 + return self.gemm_rs_ar.apply(x, self.down_proj)
282
283 x, _ = self.down_proj(x)
284 if self.shard_sequence_parallel: