@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ec5de7fa7d

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 20:45 UTC

[LoRA] Add one shot triton kernel For MoE LoRA (#42290)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed May 26, 2026 at 10:47 UTC ec5de7fa7d1303051ccd6c7316171b6d70bac4b4

14 files changed +1975 -85

benchmarks/kernels/benchmark_fused_moe_lora_one_shot.py new

+465

@@ -0,0 +1,465 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Benchmark the fused MoE-LoRA fast path (one-shot) vs two-kernel baseline.

4 +

5 +The "one_shot" provider goes through `vllm.lora.ops.triton_ops.fused_moe_lora`

6 +which dispatches to the single-kernel one-shot implementation when

7 +fully_sharded=False (the prefill default).

8 +

9 +The "two_kernel" provider drives `fused_moe_lora_shrink` + `fused_moe_lora_expand`

10 +directly, bypassing the dispatch and matching the legacy two-kernel path's

11 +work distribution. This isolates the win from kernel fusion.

12 +

13 +Run:

14 + .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot

15 + .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot \\

16 + --model qwen3moe

17 +"""

18 +

19 +from __future__ import annotations

20 +

21 +import argparse

22 +import os

23 +import random

24 +

25 +import torch

26 +

27 +from vllm import _custom_ops as ops

28 +from vllm.lora.ops.triton_ops import (

29 + fused_moe_lora,

30 + fused_moe_lora_expand,

31 + fused_moe_lora_shrink,

32 +)

33 +from vllm.triton_utils import triton

34 +

35 +DTYPE = torch.bfloat16

36 +DEVICE = "cuda"

37 +

38 +

39 +# ----- input fabrication -----------------------------------------------------

40 +

41 +

42 +def _round_up(x: int, base: int) -> int:

43 + return ((x + base - 1) // base) * base

44 +

45 +

46 +def _ceildiv(x: int, y: int) -> int:

47 + return (x + y - 1) // y

48 +

49 +

50 +def _assign_loras(num_tokens: int, num_sequences: int, max_loras: int) -> torch.Tensor:

51 + tokens_per_seq = num_tokens // num_sequences

52 + rem = num_tokens % num_sequences

53 + out = torch.empty(num_tokens, dtype=torch.int32)

54 + start = 0

55 + for i in range(num_sequences):

56 + end = start + tokens_per_seq + (1 if i < rem else 0)

57 + out[start:end] = random.randint(0, max_loras - 1)

58 + start = end

59 + return out

60 +

61 +

62 +def _assign_experts(num_tokens: int, num_experts: int, top_k: int):

63 + expert_indices = torch.empty((num_tokens, top_k), dtype=torch.int32)

64 + for i in range(num_tokens):

65 + expert_indices[i] = torch.randperm(num_experts)[:top_k]

66 + weights = torch.rand((num_tokens, top_k), dtype=torch.float32)

67 + weights = weights / weights.sum(dim=1, keepdim=True)

68 + return expert_indices, weights

69 +

70 +

71 +def _make_inputs(

72 + M: int,

73 + K: int,

74 + N_per_slice: int,

75 + rank: int,

76 + num_experts: int,

77 + top_k: int,

78 + max_loras: int,

79 + num_slices: int,

80 + block_size_m: int,

81 +):

82 + """Mirrors the production caller's tensor layout."""

83 + torch.manual_seed(0)

84 + random.seed(0)

85 +

86 + num_sequences = max(1, min(M, 8))

87 + topk_ids_cpu, topk_weights_cpu = _assign_experts(M, num_experts, top_k)

88 + token_lora_cpu = _assign_loras(M, num_sequences, max_loras)

89 + lora_ids_cpu = torch.full((max_loras + 1,), -1, dtype=torch.int32)

90 + uniq = torch.unique(token_lora_cpu, sorted=True)

91 + lora_ids_cpu[: uniq.size(0)].copy_(uniq)

92 +

93 + topk_ids = topk_ids_cpu.to(DEVICE)

94 + topk_weights = topk_weights_cpu.to(device=DEVICE, dtype=DTYPE)

95 + token_lora_mapping = token_lora_cpu.to(DEVICE)

96 + lora_ids = lora_ids_cpu.to(DEVICE)

97 + adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32, device=DEVICE)

98 +

99 + lora_a = [

100 + torch.randn((max_loras, num_experts, rank, K), dtype=DTYPE, device=DEVICE)

101 + / max(K, 1) ** 0.5

102 + for _ in range(num_slices)

103 + ]

104 + lora_b = [

105 + torch.randn(

106 + (max_loras, num_experts, N_per_slice, rank),

107 + dtype=DTYPE,

108 + device=DEVICE,

109 + )

110 + / max(rank, 1) ** 0.5

111 + for _ in range(num_slices)

112 + ]

113 + hidden = torch.randn((M, K), dtype=DTYPE, device=DEVICE)

114 + out_template = torch.zeros(

115 + (M, top_k, num_slices * N_per_slice), dtype=DTYPE, device=DEVICE

116 + )

117 +

118 + # Sorted-path metadata (the prefill default).

119 + max_pad = topk_ids.numel() + num_experts * (block_size_m - 1)

120 + max_pad = _round_up(max_pad, block_size_m)

121 + max_blocks = _ceildiv(max_pad, block_size_m)

122 + sorted_token_ids = torch.empty(

123 + (max_loras * max_pad,), dtype=torch.int32, device=DEVICE

124 + )

125 + expert_ids = torch.empty(

126 + (max_loras * max_blocks,), dtype=torch.int32, device=DEVICE

127 + )

128 + num_post = torch.empty((max_loras,), dtype=torch.int32, device=DEVICE)

129 + ops.moe_lora_align_block_size(

130 + topk_ids,

131 + token_lora_mapping,

132 + num_experts,

133 + block_size_m,

134 + max_loras,

135 + max_pad,

136 + max_blocks,

137 + sorted_token_ids,

138 + expert_ids,

139 + num_post,

140 + adapter_enabled,

141 + lora_ids,

142 + )

143 + expert_ids = expert_ids.view(max_loras, -1).contiguous()

144 + sorted_token_ids = sorted_token_ids.view(max_loras, -1).contiguous()

145 + num_active = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")

146 +

147 + return dict(

148 + hidden=hidden,

149 + lora_a=lora_a,

150 + lora_b=lora_b,

151 + topk_weights=topk_weights,

152 + sorted_token_ids=sorted_token_ids,

153 + expert_ids=expert_ids,

154 + num_post=num_post,

155 + token_lora_mapping=token_lora_mapping,

156 + lora_ids=lora_ids,

157 + num_active=num_active,

158 + adapter_enabled=adapter_enabled,

159 + out_template=out_template,

160 + # bookkeeping

161 + M=M,

162 + K=K,

163 + N_per_slice=N_per_slice,

164 + rank=rank,

165 + num_experts=num_experts,

166 + top_k=top_k,

167 + max_loras=max_loras,

168 + num_slices=num_slices,

169 + block_size_m=block_size_m,

170 + )

171 +

172 +

173 +# ----- providers -------------------------------------------------------------

174 +

175 +

176 +def _run_one_shot(inp: dict):

177 + """Drive `fused_moe_lora` with fully_sharded=False -> one-shot fast path."""

178 + out = inp["out_template"].clone()

179 + fused_moe_lora(

180 + out,

181 + inp["hidden"],

182 + inp["lora_a"],

183 + inp["lora_b"],

184 + inp["topk_weights"],

185 + inp["sorted_token_ids"],

186 + inp["expert_ids"],

187 + inp["num_post"],

188 + inp["token_lora_mapping"],

189 + inp["rank"],

190 + inp["top_k"],

191 + inp["lora_ids"],

192 + inp["num_active"],

193 + inp["adapter_enabled"],

194 + inp["block_size_m"],

195 + 64,

196 + 32,

197 + 8,

198 + 4,

199 + 3,

200 + 1,

201 + inp["block_size_m"],

202 + 64,

203 + 32,

204 + 8,

205 + 4,

206 + 3,

207 + 1,

208 + False,

209 + False,

210 + 0,

211 + )

212 + return out

213 +

214 +

215 +def _run_two_kernel(inp: dict):

216 + """Drive `fused_moe_lora_shrink` + `fused_moe_lora_expand` directly,

217 + bypassing the dispatch. Matches the legacy two-kernel work distribution.

218 + """

219 + M = inp["M"]

220 + top_k = inp["top_k"]

221 + rank = inp["rank"]

222 + num_slices = inp["num_slices"]

223 + N_per_slice = inp["N_per_slice"]

224 + K = inp["K"]

225 + num_experts = inp["num_experts"]

226 + block_m = inp["block_size_m"]

227 +

228 + intermediate = torch.zeros((num_slices, M, top_k, rank), dtype=DTYPE, device=DEVICE)

229 + out = inp["out_template"].clone()

230 + EM = inp["sorted_token_ids"].shape[1]

231 + num_tokens = M * top_k

232 +

233 + fused_moe_lora_shrink(

234 + intermediate,

235 + inp["hidden"],

236 + inp["lora_a"],

237 + inp["topk_weights"],

238 + inp["sorted_token_ids"],

239 + inp["expert_ids"],

240 + inp["num_post"],

241 + inp["token_lora_mapping"],

242 + top_k,

243 + inp["lora_ids"],

244 + inp["adapter_enabled"],

245 + torch.device(DEVICE),

246 + rank,

247 + M,

248 + EM,

249 + K,

250 + num_tokens,

251 + num_experts,

252 + num_slices,

253 + block_m,

254 + 64,

255 + 32,

256 + 8,

257 + 4,

258 + 3,

259 + 1,

260 + inp["num_active"],

261 + False,

262 + )

263 + fused_moe_lora_expand(

264 + out,

265 + intermediate,

266 + inp["lora_b"],

267 + inp["topk_weights"],

268 + inp["sorted_token_ids"],

269 + inp["expert_ids"],

270 + inp["num_post"],

271 + inp["token_lora_mapping"],

272 + top_k,

273 + inp["lora_ids"],

274 + inp["adapter_enabled"],

275 + torch.device(DEVICE),

276 + rank,

277 + M,

278 + EM,

279 + K,

280 + num_tokens,

281 + num_experts,

282 + num_slices,

283 + rank,

284 + N_per_slice,

285 + block_m,

286 + 64,

287 + 32,

288 + 8,

289 + 4,

290 + 3,

291 + 1,

292 + inp["num_active"],

293 + False,

294 + 0,

295 + )

296 + return out

297 +

298 +

299 +PROVIDER_FNS = {

300 + "one_shot": _run_one_shot,

301 + "two_kernel": _run_two_kernel,

302 +}

303 +

304 +

305 +# ----- model presets ---------------------------------------------------------

306 +

307 +

308 +MODEL_PRESETS: dict[str, dict] = {

309 + # Mixtral-8x7B style: E=8, top_k=2, hidden=4096, intermediate=14336

310 + "mixtral": dict(

311 + K=4096,

312 + N_per_slice=7168,

313 + num_experts=8,

314 + top_k=2,

315 + max_loras=4,

316 + num_slices=2,

317 + block_size_m=64,

318 + ),

319 + # Qwen3-MoE / DeepSeek-V2 style: E=64, top_k=8, hidden=2048, inter=1408

320 + "qwen3moe": dict(

321 + K=2048,

322 + N_per_slice=1408,

323 + num_experts=64,

324 + top_k=8,

325 + max_loras=4,

326 + num_slices=2,

327 + block_size_m=64,

328 + ),

329 + # GLM-5.1 (zai-org/GLM-5.1-FP8): E=256, top_k=8, hidden=6144,

330 + # moe_intermediate=2048

331 + "glm5_1": dict(

332 + K=6144,

333 + N_per_slice=2048,

334 + num_experts=256,

335 + top_k=8,

336 + max_loras=4,

337 + num_slices=2,

338 + block_size_m=64,

339 + ),

340 +}

341 +

342 +

343 +M_RANGE = [16, 64, 256, 1024, 4096, 16384]

344 +RANK_RANGE = [8, 16, 32, 64]

345 +

346 +

347 +def get_benchmark(model: str, max_loras: int | None = None):

348 + preset = dict(MODEL_PRESETS[model])

349 + if max_loras is not None:

350 + preset["max_loras"] = max_loras

351 +

352 + @triton.testing.perf_report(

353 + triton.testing.Benchmark(

354 + x_names=["M", "rank"],

355 + x_vals=[(M, R) for M in M_RANGE for R in RANK_RANGE],

356 + line_arg="provider",

357 + line_vals=list(PROVIDER_FNS.keys()),

358 + line_names=["one_shot (fused)", "two_kernel (legacy)"],

359 + styles=[("red", "-"), ("blue", "-")],

360 + ylabel="ms",

361 + plot_name=f"fused_moe_lora-{model}-loras{preset['max_loras']}",

362 + args={"preset": preset},

363 + )

364 + )

365 + def benchmark(M, rank, provider, preset):

366 + inp = _make_inputs(

367 + M=M,

368 + K=preset["K"],

369 + N_per_slice=preset["N_per_slice"],

370 + rank=rank,

371 + num_experts=preset["num_experts"],

372 + top_k=preset["top_k"],

373 + max_loras=preset["max_loras"],

374 + num_slices=preset["num_slices"],

375 + block_size_m=preset["block_size_m"],

376 + )

377 + fn = PROVIDER_FNS[provider]

378 + quantiles = [0.5, 0.2, 0.8]

379 + ms, min_ms, max_ms = triton.testing.do_bench(

380 + lambda: fn(inp), quantiles=quantiles

381 + )

382 + return ms, max_ms, min_ms

383 +

384 + return benchmark

385 +

386 +

387 +# ----- correctness sanity ---------------------------------------------------

388 +

389 +

390 +def calculate_diff(model: str, M: int, rank: int, max_loras: int | None = None):

391 + preset = dict(MODEL_PRESETS[model])

392 + if max_loras is not None:

393 + preset["max_loras"] = max_loras

394 + inp = _make_inputs(

395 + M=M,

396 + K=preset["K"],

397 + N_per_slice=preset["N_per_slice"],

398 + rank=rank,

399 + num_experts=preset["num_experts"],

400 + top_k=preset["top_k"],

401 + max_loras=preset["max_loras"],

402 + num_slices=preset["num_slices"],

403 + block_size_m=preset["block_size_m"],

404 + )

405 + out_one = _run_one_shot(inp)

406 + out_two = _run_two_kernel(inp)

407 + max_abs = (out_one.float() - out_two.float()).abs().max().item()

408 + print(

409 + f" model={model:<9} M={M:<6} rank={rank:<3} "

410 + f"max|one_shot - two_kernel|={max_abs:.4g} "

411 + f"ref|max|={out_two.float().abs().max().item():.3g}"

412 + )

413 + if max_abs <= 5e-2:

414 + print(" ✅ outputs match within bf16 tolerance")

415 + else:

416 + print(" ❌ outputs differ beyond expected bf16 noise")

417 +

418 +

419 +# ----- main ------------------------------------------------------------------

420 +

421 +

422 +if __name__ == "__main__":

423 + parser = argparse.ArgumentParser()

424 + parser.add_argument(

425 + "--model",

426 + type=str,

427 + default="mixtral",

428 + choices=list(MODEL_PRESETS.keys()),

429 + help="Model preset to sweep",

430 + )

431 + parser.add_argument(

432 + "--save-path",

433 + type=str,

434 + default="./configs/fused_moe_lora_one_shot/",

435 + help="Directory to save benchmark results",

436 + )

437 + parser.add_argument(

438 + "--check-only",

439 + action="store_true",

440 + help="Run correctness sanity check only, no perf sweep",

441 + )

442 + parser.add_argument(

443 + "--max-loras",

444 + type=int,

445 + default=None,

446 + help="Override max_loras in the model preset (number of LoRA adapters "

447 + "active in the batch). Defaults to the preset's value.",

448 + )

449 + args = parser.parse_args()

450 +

451 + print(f"Correctness check ({args.model}):")

452 + calculate_diff(args.model, M=256, rank=32, max_loras=args.max_loras)

453 + if args.check_only:

454 + raise SystemExit(0)

455 +

456 + effective_max_loras = (

457 + args.max_loras

458 + if args.max_loras is not None

459 + else MODEL_PRESETS[args.model]["max_loras"]

460 + )

461 + print(f"\nGPU: {torch.cuda.get_device_name()}")

462 + print(f"Model preset: {args.model} max_loras={effective_max_loras}\n")

463 + benchmark = get_benchmark(args.model, max_loras=args.max_loras)

464 + os.makedirs(args.save_path, exist_ok=True)

465 + benchmark.run(print_data=True, save_path=args.save_path)

tests/lora/test_fused_moe_lora_kernel.py

+495

@@ -141,6 +141,7 @@ def use_fused_moe_lora_kernel(

141 block_size,

142 fully_sharded=False,

143 offset=0,

144 + add_inputs=True,

145 ):

146 max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)

147 max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)

@@ -222,6 +223,7 @@ def use_fused_moe_lora_kernel(

223 mul_routed_weight,

224 fully_sharded=fully_sharded,

225 offset=offset,

226 + add_inputs=add_inputs,

227 )

228

229

@@ -371,6 +373,7 @@ def use_fused_moe_lora_kernel_naive(

373 block_size,

374 fully_sharded=False,

375 offset=0,

376 + add_inputs=True,

377 ):

378 """

379 Test helper for naive_block_assignment path.

@@ -435,6 +438,7 @@ def use_fused_moe_lora_kernel_naive(

438 mul_routed_weight=mul_routed_weight,

439 fully_sharded=fully_sharded,

440 offset=offset,

441 + add_inputs=add_inputs,

442 )

443

444

@@ -745,3 +749,494 @@ def use_fused_moe_lora_kernel_tensor_parallel(

749 output = tensor_model_parallel_all_reduce(output)

750

751 torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)

752 +

753 +

754 +# -- one-shot fast-path coverage --------------------------------------------

755 +# The fused shrink+expand one-shot kernel pads `BLOCK_R` to next_pow2(rank),

756 +# with a floor of 16 (tensor-core minimum). Small ranks (4, 8) exercise the

757 +# rank-dim masking and are not covered by the original tests, which start at

758 +# rank=16. The legacy two-kernel path additionally fails on rank=4 in TMA

759 +# mode because the rank-dim stride (rank * elem_size) is not 16-byte

760 +# aligned; the one-shot fast path takes precedence whenever fully_sharded

761 +# is False so this regression is hidden in normal use, but the test still

762 +# ensures the one-shot logic is correct against the pytorch reference.

763 +

764 +

765 +@pytest.mark.parametrize("num_tokens", [16, 100])

766 +@pytest.mark.parametrize("top_k_num", [2])

767 +@pytest.mark.parametrize("num_experts", [8, 64])

768 +@pytest.mark.parametrize("max_loras", [4])

769 +@pytest.mark.parametrize("N", [1408])

770 +@pytest.mark.parametrize("K", [2048])

771 +@pytest.mark.parametrize("max_lora_rank", [4, 8])

772 +@pytest.mark.parametrize("block_size", [16, 64])

773 +@pytest.mark.parametrize("num_slices", [1, 2])

774 +@pytest.mark.parametrize("dtype", DTYPES)

775 +@pytest.mark.parametrize("device", DEVICES)

776 +@pytest.mark.parametrize("seed", SEED)

777 +def test_fused_moe_lora_kernel_small_rank(

778 + num_tokens,

779 + top_k_num,

780 + num_experts,

781 + max_loras,

782 + N,

783 + K,

784 + max_lora_rank,

785 + block_size,

786 + num_slices,

787 + dtype,

788 + device,

789 + seed,

790 +):

791 + """One-shot fast path covering rank<16 (padded to BLOCK_R=16 inside kernel)."""

792 + torch.set_default_device(device)

793 + set_random_seed(seed)

794 + num_sequences = max(1, min(num_tokens, 8))

795 + topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(

796 + num_tokens, num_sequences, max_loras, num_experts, top_k_num

797 + )

798 +

799 + lora_a_stacked = [

800 + torch.rand(

801 + (max_loras, num_experts, max_lora_rank, K),

802 + dtype=dtype,

803 + )

804 + for _ in range(num_slices)

805 + ]

806 + lora_b_stacked = [

807 + torch.rand(

808 + (max_loras, num_experts, N // num_slices, max_lora_rank),

809 + dtype=dtype,

810 + )

811 + for _ in range(num_slices)

812 + ]

813 + hidden_states = torch.rand((num_tokens, K), dtype=dtype)

814 +

815 + output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)

816 + use_fused_moe_lora_kernel(

817 + topk_ids,

818 + topk_weights,

819 + token_lora_mapping,

820 + max_lora_rank,

821 + top_k_num,

822 + lora_ids,

823 + lora_a_stacked,

824 + lora_b_stacked,

825 + hidden_states,

826 + output,

827 + max_loras,

828 + num_experts,

829 + block_size,

830 + )

831 + output_ref = use_torch(

832 + hidden_states,

833 + token_lora_mapping,

834 + topk_ids,

835 + lora_a_stacked,

836 + lora_b_stacked,

837 + top_k_num,

838 + num_slices,

839 + )

840 + torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)

841 +

842 +

843 +@pytest.mark.parametrize("num_tokens", [16, 64])

844 +@pytest.mark.parametrize("top_k_num", [2])

845 +@pytest.mark.parametrize("num_experts", [8])

846 +@pytest.mark.parametrize("max_loras", [4])

847 +@pytest.mark.parametrize("N", [2048])

848 +@pytest.mark.parametrize("K", [4096])

849 +@pytest.mark.parametrize("max_lora_rank", [8, 16, 32, 64])

850 +@pytest.mark.parametrize("block_size", [64])

851 +@pytest.mark.parametrize("num_slices", [2])

852 +@pytest.mark.parametrize("dtype", [torch.bfloat16])

853 +@pytest.mark.parametrize("device", DEVICES)

854 +@pytest.mark.parametrize("seed", SEED)

855 +def test_fused_moe_lora_kernel_npid_path(

856 + num_tokens,

857 + top_k_num,

858 + num_experts,

859 + max_loras,

860 + N,

861 + K,

862 + max_lora_rank,

863 + block_size,

864 + num_slices,

865 + dtype,

866 + device,

867 + seed,

868 +):

869 + """Exercise the small-batch / NPID > 1 branch of the one-shot fast path.

870 +

871 + With these sizes the one-shot wrapper computes NPID_FACTOR > 1 (base CTA count

872 + < SM count), so each program covers only an outer chunk of N. The

873 + cross-outer-block write mask is the correctness-critical bit.

874 + """

875 + torch.set_default_device(device)

876 + set_random_seed(seed)

877 + num_sequences = max(1, min(num_tokens, 4))

878 + topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(

879 + num_tokens, num_sequences, max_loras, num_experts, top_k_num

880 + )

881 +

882 + lora_a_stacked = [

883 + torch.rand(

884 + (max_loras, num_experts, max_lora_rank, K),

885 + dtype=dtype,

886 + )

887 + for _ in range(num_slices)

888 + ]

889 + lora_b_stacked = [

890 + torch.rand(

891 + (max_loras, num_experts, N // num_slices, max_lora_rank),

892 + dtype=dtype,

893 + )

894 + for _ in range(num_slices)

895 + ]

896 + hidden_states = torch.rand((num_tokens, K), dtype=dtype)

897 +

898 + output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)

899 + use_fused_moe_lora_kernel(

900 + topk_ids,

901 + topk_weights,

902 + token_lora_mapping,

903 + max_lora_rank,

904 + top_k_num,

905 + lora_ids,

906 + lora_a_stacked,

907 + lora_b_stacked,

908 + hidden_states,

909 + output,

910 + max_loras,

911 + num_experts,

912 + block_size,

913 + )

914 + output_ref = use_torch(

915 + hidden_states,

916 + token_lora_mapping,

917 + topk_ids,

918 + lora_a_stacked,

919 + lora_b_stacked,

920 + top_k_num,

921 + num_slices,

922 + )

923 + torch.testing.assert_close(output, output_ref, atol=2e-2, rtol=2e-2)

924 +

925 +

926 +# -- one-shot corner-case coverage ------------------------------------------

927 +# Each of the following exercises a path where the kernel is launched but

928 +# every program early-exits, leaving the output unchanged. The contract is

929 +# additive (`output += contribution`), so an empty contribution must leave

930 +# the input residual untouched.

931 +

932 +

933 +def _build_one_shot_inputs(

934 + num_tokens,

935 + top_k_num,

936 + num_experts,

937 + max_loras,

938 + max_lora_rank,

939 + K,

940 + N,

941 + num_slices,

942 + block_size,

943 + dtype,

944 +):

945 + """Common scaffolding for the corner-case tests below."""

946 + num_sequences = max(1, min(num_tokens, 4)) if num_tokens > 0 else 1

947 + if num_tokens > 0:

948 + topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(

949 + num_tokens, num_sequences, max_loras, num_experts, top_k_num

950 + )

951 + else:

952 + # M=0 path: caller may still hand us empty tensors with the right shape.

953 + topk_ids = torch.empty((0, top_k_num), dtype=torch.int32)

954 + topk_weights = torch.empty((0, top_k_num), dtype=torch.float32)

955 + token_lora_mapping = torch.empty((0,), dtype=torch.int32)

956 + lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)

957 +

958 + lora_a_stacked = [

959 + torch.rand((max_loras, num_experts, max_lora_rank, K), dtype=dtype)

960 + for _ in range(num_slices)

961 + ]

962 + lora_b_stacked = [

963 + torch.rand(

964 + (max_loras, num_experts, N // num_slices, max_lora_rank), dtype=dtype

965 + )

966 + for _ in range(num_slices)

967 + ]

968 + hidden_states = torch.rand((max(num_tokens, 0), K), dtype=dtype)

969 + return (

970 + topk_ids,

971 + topk_weights.to(dtype),

972 + token_lora_mapping,

973 + lora_ids,

974 + lora_a_stacked,

975 + lora_b_stacked,

976 + hidden_states,

977 + )

978 +

979 +

980 +def _call_one_shot(

981 + output,

982 + hidden_states,

983 + lora_a_stacked,

984 + lora_b_stacked,

985 + topk_weights,

986 + sorted_token_ids,

987 + expert_ids,

988 + num_tokens_post_padded,

989 + token_lora_mapping,

990 + max_lora_rank,

991 + top_k_num,

992 + lora_ids,

993 + num_active_loras,

994 + adapter_enabled,

995 + block_size,

996 + add_inputs=True,

997 +):

998 + """Direct call into fused_moe_lora with one-shot-routed defaults."""

999 + from vllm.lora.ops.triton_ops import fused_moe_lora as _op

1000 +

1001 + _op(

1002 + output,

1003 + hidden_states,

1004 + lora_a_stacked,

1005 + lora_b_stacked,

1006 + topk_weights,

1007 + sorted_token_ids,

1008 + expert_ids,

1009 + num_tokens_post_padded,

1010 + token_lora_mapping,

1011 + max_lora_rank,

1012 + top_k_num,

1013 + lora_ids,

1014 + num_active_loras,

1015 + adapter_enabled,

1016 + block_size,

1017 + 32,

1018 + 64,

1019 + 1,

1020 + 4,

1021 + 3,

1022 + 1,

1023 + block_size,

1024 + 32,

1025 + 64,

1026 + 1,

1027 + 4,

1028 + 3,

1029 + 1,

1030 + False,

1031 + False,

1032 + 0,

1033 + add_inputs,

1034 + )

1035 +

1036 +

1037 +@pytest.mark.parametrize(

1038 + "trigger",

1039 + ["sorted_lora_ids_neg", "naive_mapping_neg", "naive_all_disabled"],

1040 +)

1041 +@pytest.mark.parametrize("device", DEVICES)

1042 +def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device):

1043 + """one-shot must leave the residual byte-identical when every program

1044 + must early-exit. Three trigger conditions are covered:

1045 +

1046 + - "sorted_lora_ids_neg": sorted path, lora_ids all -1 (lora_id<0 check)

1047 + - "naive_mapping_neg": naive path, token_lora_mapping all -1

1048 + - "naive_all_disabled": naive path, adapter_enabled all 0

1049 + """

1050 + torch.set_default_device(device)

1051 + set_random_seed(0)

1052 +

1053 + # Per-trigger shapes: naive_mapping_neg needs the naive dispatch gate

1054 + # `num_tokens*top_k*8 <= num_experts*max_loras` to hold, hence the

1055 + # larger E/max_loras and smaller num_tokens.

1056 + if trigger == "naive_mapping_neg":

1057 + num_tokens, top_k, E, max_loras, R = 8, 2, 64, 8, 16

1058 + elif trigger == "naive_all_disabled":

1059 + num_tokens, top_k, E, max_loras, R = 32, 2, 8, 4, 32

1060 + else: # sorted_lora_ids_neg

1061 + num_tokens, top_k, E, max_loras, R = 32, 2, 8, 4, 16

1062 + K, N = 1024, 1024

1063 + block_size, num_slices, dtype = 16, 2, torch.bfloat16

1064 +

1065 + (

1066 + topk_ids,

1067 + topk_weights,

1068 + token_lora_mapping,

1069 + lora_ids,

1070 + lora_a_stacked,

1071 + lora_b_stacked,

1072 + hidden_states,

1073 + ) = _build_one_shot_inputs(

1074 + num_tokens, top_k, E, max_loras, R, K, N, num_slices, block_size, dtype

1075 + )

1076 +

1077 + adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)

1078 + num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")

1079 +

1080 + if trigger == "sorted_lora_ids_neg":

1081 + lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)

1082 + max_pad = topk_ids.numel() + E * (block_size - 1)

1083 + max_pad = round_up(max_pad, block_size)

1084 + max_blocks = CEILDIV(max_pad, block_size)

1085 + sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)

1086 + expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)

1087 + num_post = torch.zeros((max_loras,), dtype=torch.int32)

1088 + else:

1089 + sorted_token_ids = None

1090 + expert_ids = topk_ids.reshape(-1).contiguous()

1091 + num_post = None

1092 + if trigger == "naive_mapping_neg":

1093 + token_lora_mapping = torch.full((num_tokens,), -1, dtype=torch.int32)

1094 + lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)

1095 + else: # naive_all_disabled

1096 + adapter_enabled = torch.zeros(max_loras + 1, dtype=torch.int32)

1097 +

1098 + residual = torch.randn((num_tokens, top_k, N), dtype=dtype) * 0.1

1099 + output = residual.clone()

1100 +

1101 + _call_one_shot(

1102 + output,

1103 + hidden_states,

1104 + lora_a_stacked,

1105 + lora_b_stacked,

1106 + topk_weights,

1107 + sorted_token_ids,

1108 + expert_ids,

1109 + num_post,

1110 + token_lora_mapping,

1111 + R,

1112 + top_k,

1113 + lora_ids,

1114 + num_active_loras,

1115 + adapter_enabled,

1116 + block_size,

1117 + )

1118 + torch.testing.assert_close(output, residual, atol=0, rtol=0)

1119 +

1120 +

1121 +@pytest.mark.parametrize("device", DEVICES)

1122 +def test_fused_moe_lora_kernel_zero_grid_no_crash(device):

1123 + """num_active_loras=0 (or num_slices=0) would otherwise launch a grid

1124 + with a zero dimension. one-shot wrapper must short-circuit before launch."""

1125 + torch.set_default_device(device)

1126 + set_random_seed(0)

1127 + num_tokens, top_k, E, max_loras, R, K, N = 8, 2, 8, 4, 16, 1024, 1024

1128 + block_size, num_slices, dtype = 16, 2, torch.bfloat16

1129 +

1130 + (

1131 + topk_ids,

1132 + topk_weights,

1133 + token_lora_mapping,

1134 + lora_ids,

1135 + lora_a_stacked,

1136 + lora_b_stacked,

1137 + hidden_states,

1138 + ) = _build_one_shot_inputs(

1139 + num_tokens,

1140 + top_k,

1141 + E,

1142 + max_loras,

1143 + R,

1144 + K,

1145 + N,

1146 + num_slices,

1147 + block_size,

1148 + dtype,

1149 + )

1150 + adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)

1151 + num_active_loras = torch.tensor([0], dtype=torch.int32, device="cpu")

1152 + residual = torch.randn((num_tokens, top_k, N), dtype=dtype) * 0.1

1153 + output = residual.clone()

1154 +

1155 + # sorted path is the one that uses num_active_loras for grid axis 2

1156 + max_pad = topk_ids.numel() + E * (block_size - 1)

1157 + max_pad = round_up(max_pad, block_size)

1158 + max_blocks = CEILDIV(max_pad, block_size)

1159 + sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)

1160 + expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)

1161 + num_post = torch.zeros((max_loras,), dtype=torch.int32)

1162 + _call_one_shot(

1163 + output,

1164 + hidden_states,

1165 + lora_a_stacked,

1166 + lora_b_stacked,

1167 + topk_weights,

1168 + sorted_token_ids,

1169 + expert_ids,

1170 + num_post,

1171 + token_lora_mapping,

1172 + R,

1173 + top_k,

1174 + lora_ids,

1175 + num_active_loras,

1176 + adapter_enabled,

1177 + block_size,

1178 + )

1179 + torch.testing.assert_close(output, residual, atol=0, rtol=0)

1180 +

1181 +

1182 +@pytest.mark.parametrize("device", DEVICES)

1183 +def test_fused_moe_lora_kernel_rejects_bad_block_size_m(device):

1184 + """one-shot must surface a clear assertion when shrink_block_size_m is not

1185 + a power of 2 / less than 16, instead of the cryptic Triton compile

1186 + failure (`arange's range must be a power of 2`)."""

1187 + torch.set_default_device(device)

1188 + set_random_seed(0)

1189 + num_tokens, top_k, E, max_loras, R, K, N = 32, 2, 8, 4, 16, 1024, 1024

1190 + num_slices, dtype = 2, torch.bfloat16

1191 + block_size = 24 # NOT a power of 2

1192 +

1193 + (

1194 + topk_ids,

1195 + topk_weights,

1196 + token_lora_mapping,

1197 + lora_ids,

1198 + lora_a_stacked,

1199 + lora_b_stacked,

1200 + hidden_states,

1201 + ) = _build_one_shot_inputs(

1202 + num_tokens,

1203 + top_k,

1204 + E,

1205 + max_loras,

1206 + R,

1207 + K,

1208 + N,

1209 + num_slices,

1210 + 16,

1211 + dtype,

1212 + )

1213 + # Build sorted-mode metadata at block_size=16 so shapes are sane,

1214 + # but pass block_size=24 to the op (the buggy combination).

1215 + max_pad = topk_ids.numel() + E * (16 - 1)

1216 + max_pad = round_up(max_pad, 16)

1217 + max_blocks = CEILDIV(max_pad, 16)

1218 + sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)

1219 + expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)

1220 + num_post = torch.zeros((max_loras,), dtype=torch.int32)

1221 + adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)

1222 + num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")

1223 + output = torch.zeros((num_tokens, top_k, N), dtype=dtype)

1224 +

1225 + with pytest.raises(AssertionError, match="shrink_block_size_m"):

1226 + _call_one_shot(

1227 + output,

1228 + hidden_states,

1229 + lora_a_stacked,

1230 + lora_b_stacked,

1231 + topk_weights,

1232 + sorted_token_ids,

1233 + expert_ids,

1234 + num_post,

1235 + token_lora_mapping,

1236 + R,

1237 + top_k,

1238 + lora_ids,

1239 + num_active_loras,

1240 + adapter_enabled,

1241 + block_size,

1242 + )

vllm/envs.py

@@ -1971,6 +1971,8 @@ environment_variables: dict[str, Callable[[], Any]] = {

1971 int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0"))

1972 ),

1973 # Whether to enable dual cuda streams for LoRA computation

1974 + # (used by both BaseLinearLayerWithLoRA and FusedMoEWithLoRA to

1975 + # overlap the base layer compute with the LoRA fast path).

1976 "VLLM_LORA_ENABLE_DUAL_STREAM": lambda: bool(

1977 int(os.getenv("VLLM_LORA_ENABLE_DUAL_STREAM", "0"))

1978 ),

vllm/lora/layers/base_linear.py

+1 -8

@@ -25,16 +25,9 @@ from vllm.utils.multi_stream_utils import maybe_execute_in_parallel

25 from vllm.utils.torch_utils import direct_register_custom_op

27 from .base import BaseLayerWithLoRA

28 -from .utils import _get_lora_device

28 +from .utils import _get_lora_aux_cuda_stream, _get_lora_device

30 if envs.VLLM_LORA_ENABLE_DUAL_STREAM:

31 - _lora_aux_cuda_stream: torch.cuda.Stream | None = None

32 -

33 - def _get_lora_aux_cuda_stream() -> torch.cuda.Stream | None:

34 - global _lora_aux_cuda_stream

35 - if _lora_aux_cuda_stream is None and current_platform.is_cuda_alike():

36 - _lora_aux_cuda_stream = torch.cuda.Stream()

37 - return _lora_aux_cuda_stream

32 def lora_linear_async(

33 layer_name: str,

vllm/lora/layers/fused_moe.py

+25 -1

@@ -19,8 +19,9 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

19 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

20 MoEPrepareAndFinalizeNoDPEPModular,

21 )

22 +from vllm.platforms import current_platform

23 -from .utils import _get_lora_device

24 +from .utils import _get_lora_aux_cuda_stream, _get_lora_device

27 class FusedMoEWithLoRA(BaseLayerWithLoRA):

@@ -34,6 +35,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

35 self.tp_size = self.base_layer.tp_size

36 self.tp_rank = self.base_layer.tp_rank

37 self.device = _get_lora_device(base_layer)

38 +

39 + self._enable_aux_cuda_stream = envs.VLLM_LORA_ENABLE_DUAL_STREAM

40 + self._init_lora_stream_context()

41 # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed

42 # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3)

43 self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1

@@ -65,7 +69,25 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

69 FusedMoEModularMethod(self.base_layer.quant_method, moe_kernel)

70 )

72 + def _init_lora_stream_context(self) -> None:

73 + self._lora_stream: torch.cuda.Stream | None = None

74 + self._events: tuple[torch.cuda.Event, ...] | None = None

75 + if not self._enable_aux_cuda_stream:

76 + return

77 + if not current_platform.is_cuda_alike():

78 + return

79 + self._lora_stream = _get_lora_aux_cuda_stream()

80 + # 4 events: 2 per (base GEMM, LoRA) pair so w13 and w2 don't reuse

81 + # the same event objects; reuse-within-a-pair is fine because the

82 + # second pair starts only after intermediate_cache1.add_() has joined.

83 + self._events = tuple(torch.cuda.Event() for _ in range(4))

84 +

85 def _build_lora_context(self):

86 + use_dual_stream = (

87 + self._enable_aux_cuda_stream

88 + and not self.fully_sharded

89 + and self._lora_stream is not None

90 + )

91 return MoELoRAContext(

92 w13_lora_a_stacked=self.w13_lora_a_stacked,

93 w13_lora_b_stacked=self.w13_lora_b_stacked,

@@ -81,6 +103,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

103 local_num_experts=self.base_layer.local_num_experts,

104 punica_wrapper=self.punica_wrapper,

105 use_tuned_config=bool(envs.VLLM_TUNED_CONFIG_FOLDER),

106 + aux_stream=self._lora_stream if use_dual_stream else None,

107 + events=self._events if use_dual_stream else None,

108 )

109

110 def _create_lora_a_weights(

vllm/lora/layers/utils.py

+13

@@ -7,9 +7,22 @@ from enum import Enum

7 import torch

8 import torch.nn as nn

10 +from vllm import envs

11 from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config

12 +from vllm.platforms import current_platform

13 from vllm.utils.math_utils import next_power_of_2

15 +_lora_aux_cuda_stream: torch.cuda.Stream | None = None

16 +

17 +

18 +def _get_lora_aux_cuda_stream() -> torch.cuda.Stream | None:

19 + if not envs.VLLM_LORA_ENABLE_DUAL_STREAM:

20 + return None

21 + global _lora_aux_cuda_stream

22 + if _lora_aux_cuda_stream is None and current_platform.is_cuda_alike():

23 + _lora_aux_cuda_stream = torch.cuda.Stream()

24 + return _lora_aux_cuda_stream

25 +

27 class LoRAMappingType(Enum):

28 LANGUAGE = 1

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

+791

@@ -105,6 +105,742 @@ def _get_c_ptrs(

105 _LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}

106

107

108 +# ---------------------------------------------------------------------------

109 +# Fully-fused MoE-LoRA kernel (one-shot): shrink + expand combined into a single

110 +# launch with the rank-dim intermediate kept in registers. Used by the fast

111 +# path of `_fused_moe_lora` for `fully_sharded=False`. The legacy two-kernel

112 +# path (`_fused_moe_lora_kernel` above) is retained for `fully_sharded=True`

113 +# because that path needs to materialise the intermediate cache for an

114 +# all_reduce / all_gather between shrink and expand.

115 +# ---------------------------------------------------------------------------

116 +

117 +

118 +@triton.heuristics({"EVEN_K": lambda args: args["K"] % args["BLOCK_K"] == 0})

119 +@triton.jit

120 +def _fused_moe_lora_one_shot_kernel(

121 + # ---- pointers ----

122 + x_ptr,

123 + A_ptrs,

124 + B_ptrs,

125 + out_ptr,

126 + topk_weights_ptr,

127 + sorted_token_ids_ptr,

128 + expert_ids_ptr,

129 + num_tokens_post_padded_ptr,

130 + token_lora_mapping_ptr,

131 + lora_ids_ptr,

132 + adapter_enabled_ptr,

133 + # ---- dims ----

134 + N,

135 + K,

136 + num_valid_tokens,

137 + top_k_num,

138 + max_loras,

139 + # ---- strides ----

140 + stride_xm,

141 + stride_xk,

142 + stride_A_lora,

143 + stride_A_expert,

144 + stride_A_r,

145 + stride_A_k,

146 + stride_B_lora,

147 + stride_B_expert,

148 + stride_B_n,

149 + stride_B_r,

150 + stride_om,

151 + stride_on,

152 + stride_tl_,

153 + stride_el,

154 + # ---- scalar ----

155 + slice_n_offset,

156 + # ---- constexpr (set per call) ----

157 + token_mapping_factor: tl.constexpr,

158 + naive_block_assignment: tl.constexpr,

159 + MUL_ROUTED_WEIGHT: tl.constexpr,

160 + BLOCK_M: tl.constexpr,

161 + BLOCK_R: tl.constexpr,

162 + actual_rank: tl.constexpr,

163 + NPID_FACTOR: tl.constexpr,

164 + BLOCK_N: tl.constexpr,

165 + BLOCK_K: tl.constexpr,

166 + EVEN_K: tl.constexpr,

167 + ADD_INPUTS: tl.constexpr,

168 +):

169 + pid_full = tl.program_id(axis=0)

170 + pid_m = pid_full // NPID_FACTOR

171 + pid_n_outer = pid_full % NPID_FACTOR

172 + slice_id = tl.program_id(axis=1)

173 + lora_idx = tl.program_id(axis=2)

174 +

175 + # Resolve lora_id.

176 + if naive_block_assignment:

177 + token_idx_for_lora = pid_m // top_k_num

178 + lora_id = tl.load(token_lora_mapping_ptr + token_idx_for_lora)

179 + else:

180 + lora_id = tl.load(lora_ids_ptr + lora_idx)

181 + if lora_id < 0:

182 + return

183 + if lora_id >= max_loras:

184 + return

185 + enabled = tl.load(adapter_enabled_ptr + lora_id)

186 + if enabled == 0:

187 + return

188 +

189 + if not naive_block_assignment:

190 + ntpp = tl.load(num_tokens_post_padded_ptr + lora_id)

191 + if pid_m * BLOCK_M >= ntpp:

192 + return

193 +

194 + # Resolve expert_id.

195 + if naive_block_assignment:

196 + expert_id = tl.load(expert_ids_ptr + pid_m)

197 + else:

198 + ind = lora_id * stride_el + pid_m

199 + expert_id = tl.load(

200 + expert_ids_ptr + ind, mask=ind < max_loras * stride_el, other=-1

201 + )

202 + if expert_id < 0:

203 + return

204 +

205 + # Compute offs_token (flat token ids).

206 + offs = tl.arange(0, BLOCK_M).to(tl.int64)

207 + if naive_block_assignment:

208 + offs_token = tl.where(offs == 0, pid_m, num_valid_tokens)

209 + else:

210 + offs_token_id = pid_m * BLOCK_M + offs

211 + token_ind = stride_tl_ * lora_id + offs_token_id

212 + offs_token = tl.load(

213 + sorted_token_ids_ptr + token_ind,

214 + mask=token_ind < max_loras * stride_tl_,

215 + other=num_valid_tokens,

216 + )

217 + token_mask = offs_token < num_valid_tokens

218 +

219 + # N range owned by this program. Splitting [0, N) into NPID_FACTOR

220 + # contiguous outer blocks lets us scale parallelism for small batches.

221 + n_per_outer = tl.cdiv(N, NPID_FACTOR)

222 + n_lo = pid_n_outer * n_per_outer

223 + n_hi = tl.minimum((pid_n_outer + 1) * n_per_outer, N)

224 + if n_lo >= N:

225 + return

226 +

227 + # Slice pointers.

228 + cur_A_ptr = tl.load(A_ptrs + slice_id).to(tl.pointer_type(out_ptr.dtype.element_ty))

229 + cur_B_ptr = tl.load(B_ptrs + slice_id).to(tl.pointer_type(out_ptr.dtype.element_ty))

230 +

231 + A_base = cur_A_ptr + lora_id * stride_A_lora + expert_id * stride_A_expert

232 + B_base = cur_B_ptr + lora_id * stride_B_lora + expert_id * stride_B_expert

233 +

234 + # SHRINK: tmp[BLOCK_M, BLOCK_R] = x @ A^T, accumulated in fp32 registers.

235 + offs_r = tl.arange(0, BLOCK_R)

236 + rank_mask = offs_r < actual_rank

237 + # Clamp rank offsets so OOB rows of A / B map to address 0; the mask

238 + # zeros the loaded values. Required when BLOCK_R > actual_rank

239 + # (e.g. rank=4 padded to 16) -- without clamping, tl.load would address

240 + # the next expert's memory.

241 + safe_offs_r = tl.where(rank_mask, offs_r, 0)

242 + offs_k = tl.arange(0, BLOCK_K)

243 +

244 + offs_x_row = offs_token // token_mapping_factor

245 + x_ptrs = x_ptr + offs_x_row[:, None] * stride_xm + offs_k[None, :] * stride_xk

246 + a_ptrs = A_base + offs_k[:, None] * stride_A_k + safe_offs_r[None, :] * stride_A_r

247 +

248 + tmp = tl.zeros((BLOCK_M, BLOCK_R), dtype=tl.float32)

249 + if EVEN_K:

250 + for _ in range(0, K, BLOCK_K):

251 + x = tl.load(x_ptrs, mask=token_mask[:, None], other=0.0)

252 + a = tl.load(a_ptrs, mask=rank_mask[None, :], other=0.0)

253 + tmp += tl.dot(x, a)

254 + x_ptrs += BLOCK_K * stride_xk

255 + a_ptrs += BLOCK_K * stride_A_k

256 + else:

257 + for kb in range(0, K, BLOCK_K):

258 + k_remain = K - kb

259 + k_mask = offs_k < k_remain

260 + x = tl.load(x_ptrs, mask=token_mask[:, None] & k_mask[None, :], other=0.0)

261 + a = tl.load(a_ptrs, mask=k_mask[:, None] & rank_mask[None, :], other=0.0)

262 + tmp += tl.dot(x, a)

263 + x_ptrs += BLOCK_K * stride_xk

264 + a_ptrs += BLOCK_K * stride_A_k

265 +

266 + tmp_typed = tmp.to(out_ptr.dtype.element_ty)

267 +

268 + # EXPAND: out[tokens, n] += tmp @ B^T, looped over BLOCK_N tiles within

269 + # this program's [n_lo, n_hi). The (offs_n < n_hi) mask is required

270 + # whenever BLOCK_N > n_per_outer to keep adjacent outer blocks from

271 + # writing into each other's columns.

272 + if MUL_ROUTED_WEIGHT:

273 + moe_w = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0).to(

274 + tl.float32

275 + )

276 +

277 + out_slice_base = out_ptr + slice_id * slice_n_offset

278 +

279 + for n_start in range(n_lo, n_hi, BLOCK_N):

280 + offs_n = n_start + tl.arange(0, BLOCK_N)

281 + n_mask = (offs_n < N) & (offs_n < n_hi)

282 +

283 + b_ptrs = (

284 + B_base + safe_offs_r[:, None] * stride_B_r + offs_n[None, :] * stride_B_n

285 + )

286 + b = tl.load(b_ptrs, mask=rank_mask[:, None] & n_mask[None, :], other=0.0)

287 +

288 + acc = tl.dot(tmp_typed, b) # (BLOCK_M, BLOCK_N) fp32

289 + if MUL_ROUTED_WEIGHT:

290 + acc = acc * moe_w[:, None]

291 +

292 + out_ptrs = (

293 + out_slice_base

294 + + offs_token[:, None] * stride_om

295 + + offs_n[None, :] * stride_on

296 + )

297 + out_mask = token_mask[:, None] & n_mask[None, :]

298 + if ADD_INPUTS:

299 + prev = tl.load(out_ptrs, mask=out_mask, other=0.0)

300 + tl.store(out_ptrs, prev + acc.to(out_ptr.dtype.element_ty), mask=out_mask)

301 + else:

302 + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=out_mask)

303 +

304 +

305 +def _run_fused_moe_lora_one_shot(

306 + output: torch.Tensor,

307 + qcurr_hidden_states: torch.Tensor,

308 + lora_a_stacked: list[torch.Tensor],

309 + lora_b_stacked: list[torch.Tensor],

310 + topk_weights: torch.Tensor,

311 + sorted_token_ids: torch.Tensor | None,

312 + expert_ids: torch.Tensor,

313 + num_tokens_post_padded: torch.Tensor | None,

314 + token_lora_mapping: torch.Tensor,

315 + max_lora_rank: int,

316 + top_k_num: int,

317 + lora_ids: torch.Tensor,

318 + num_active_loras: torch.Tensor,

319 + adapter_enabled: torch.Tensor,

320 + mul_routed_weight: bool,

321 + block_size_m: int,

322 + add_inputs: bool = True,

323 +) -> None:

324 + """Fast-path wrapper: launches one fused shrink+expand kernel.

325 +

326 + The shape contract matches `_fused_moe_lora`. `output` has shape

327 + `(num_tokens, top_k_num, num_slices * N_per_slice)`. When

328 + `add_inputs=True` (default) the kernel reads-modifies-writes `output`

329 + in place; when `add_inputs=False` the kernel overwrites `output` with

330 + the LoRA delta only. The latter is used by the dual-stream path that

331 + sums LoRA into the base output on a separate stream.

332 + """

333 + num_slices = len(lora_a_stacked)

334 + device = qcurr_hidden_states.device

335 +

336 + A0 = lora_a_stacked[0]

337 + B0 = lora_b_stacked[0]

338 + max_loras_w = A0.shape[0]

339 + rank = A0.shape[2]

340 + K = A0.shape[3]

341 + N_per_slice = B0.shape[2]

342 +

343 + # rank padding is to next pow2 with a floor of 16 (tensor-core minimum

344 + # K-dim). Beyond 128 the (BLOCK_M, BLOCK_R) accumulator outgrows the

345 + # register file; rank tiling would be needed but is out of scope for

346 + # this kernel. Tried floor=32 to double MMA density per K-step but it

347 + # regressed across all M (+8 to +40%): the (64,32) fp32 accumulator +

348 + # widened B tile pushed register count past spill threshold, lowering

349 + # occupancy by more than the MMA gain saved.

350 + assert rank <= 128, (

351 + f"fused_moe_lora_one_shot supports max_lora_rank<=128; got rank={rank}"

352 + )

353 + BLOCK_R = max(triton.next_power_of_2(rank), 16)

354 +

355 + num_experts = A0.shape[1]

356 + naive = sorted_token_ids is None

357 + if sorted_token_ids is None:

358 + EM_grid = topk_weights.numel()

359 + BLOCK_M = 16

360 + stride_tl_ = 0

361 + stride_el = 0

362 + grid_lora_dim = 1

363 + else:

364 + EM_grid = sorted_token_ids.shape[1]

365 + # BLOCK_M must equal moe_lora_align_block_size's block_size. The

366 + # caller passes that explicitly; deriving it from tensor shapes is

367 + # unsafe because sorted_token_ids.shape[1] is the raw padded length

368 + # (not necessarily a multiple of block_size — e.g. OLMoE prefill

369 + # produces sorted=139200 with expert_ids=1088 and block_size=128).

370 + # tl.arange and tl.dot need block_size_m to be a power of 2 and at

371 + # least 16. The Python-side assertion gives a clearer error than

372 + # the cryptic Triton compile failure.

373 + assert block_size_m >= 16 and (block_size_m & (block_size_m - 1)) == 0, (

374 + f"shrink_block_size_m must be a power of 2 and >=16; got {block_size_m}"

375 + )

376 + BLOCK_M = block_size_m

377 + stride_tl_ = sorted_token_ids.stride(0)

378 + stride_el = expert_ids.stride(0)

379 + grid_lora_dim = int(num_active_loras.item())

380 +

381 + # Empty-work guards: the grid would otherwise have a zero dimension,

382 + # which Triton rejects. None of these is a hot path in production -- a

383 + # batch with zero tokens, an EM_grid of zero, or zero active LoRAs all

384 + # mean there's nothing to add to `output`.

385 + if EM_grid == 0 or grid_lora_dim == 0 or num_slices == 0:

386 + return

387 +

388 + token_mapping_factor = 1 if mul_routed_weight else top_k_num

389 +

390 + A_ptrs = _get_ptr(lora_a_stacked, device)

391 + B_ptrs = _get_ptr(lora_b_stacked, device)

392 +

393 + # Flatten (num_tokens, top_k) → flat_token axis. The kernel addresses

394 + # output via offs_token * stride_om, which is correct iff the dim-0 /

395 + # dim-1 strides collapse cleanly: stride(0) == top_k * stride(1). All

396 + # production callers pass contiguous output, so this always holds; the

397 + # explicit check guards against future regressions where a non-trivial

398 + # view (e.g. permute) would silently break in-place accumulation.

399 + assert output.dim() == 3, f"output must be 3-D, got {output.shape}"

400 + assert output.stride(0) == output.shape[1] * output.stride(1), (

401 + "fused_moe_lora_one_shot requires output.stride(0) == top_k*stride(1); "

402 + f"got shape={output.shape} strides={output.stride()}"

403 + )

404 + out_view = output.view(-1, output.shape[-1])

405 + M_blocks = triton.cdiv(EM_grid, BLOCK_M) if not naive else EM_grid

406 +

407 + # NPID_FACTOR heuristic: scale N-axis parallelism when base CTA count is

408 + # short of saturating the SM array. Cap by the cost of redundant shrink.

409 + sm_count = torch.cuda.get_device_properties(device).multi_processor_count

410 + base_programs = max(M_blocks * num_slices * grid_lora_dim, 1)

411 + shrink_ratio = K / max(K + N_per_slice, 1)

412 + max_npid_by_budget = max(1, int(1.5 / max(shrink_ratio, 1e-3)) + 1)

413 + target = 2 * sm_count

414 + if base_programs >= int(1.5 * sm_count):

415 + npid = 1

416 + else:

417 + npid_occ = max(1, min(16, (target + base_programs - 1) // base_programs))

418 + npid = min(npid_occ, max_npid_by_budget)

419 + npid = max(1, min(npid, max(1, N_per_slice // 128)))

420 +

421 + # Robust defaults across the prefill regime (H100/H200/B200, bf16/fp16).

422 + # NPID > 1 is the small-M / under-saturated path -- more warps help

423 + # amortise the inner-N expand loop. ns=3 instead of 4: GB200 ncu showed

424 + # the 4-stage pipeline pushed register count to 168/thread and capped

425 + # achieved occupancy at ~17% (3 blocks/SM, register-bound); ns=3 frees

426 + # ~30 regs/thread which keeps a 4th block resident on small grids.

427 + # Tried BLOCK_N=64 for w13 (N=192) to avoid the half-wasted second

428 + # tile: regressed 11-29% because the "waste" was just masked stores

429 + # (cheap) and the extra iteration added load + index overhead.

430 + if npid > 1:

431 + block_n, nw, ns = 128, 8, 3

432 + else:

433 + block_n, nw, ns = 128, 4, 3

434 + # BLOCK_K choice: for hidden-sized K (≥256, i.e. the K=hidden_size

435 + # shrink input on w13) force BLOCK_K=128 -- the wider tile halves the

436 + # K-loop trip count and removes the scoreboard stalls that dominated

437 + # M=16-64 on GB200 (kernel time -13% to -37% vs the work_per_expert

438 + # heuristic which picked 64 for low-tokens-per-expert ratios). For

439 + # small-K shapes (e.g. w2 with K=192 where the down-proj reads the

440 + # MoE intermediate) keep the work_per_expert heuristic: BLOCK_K=128

441 + # would force the EVEN_K=False masked path and add no K-loop savings

442 + # (K/64=3 vs K/128=2 masked) while inflating per-program startup.

443 + if K >= 256:

444 + block_k = 128

445 + else:

446 + work_per_expert = topk_weights.numel() / max(num_experts, 1)

447 + block_k = 128 if work_per_expert >= 16 else 64

448 +

449 + grid = (M_blocks * npid, num_slices, grid_lora_dim)

450 +

451 + _fused_moe_lora_one_shot_kernel[grid](

452 + qcurr_hidden_states,

453 + A_ptrs,

454 + B_ptrs,

455 + out_view,

456 + topk_weights,

457 + sorted_token_ids,

458 + expert_ids,

459 + num_tokens_post_padded,

460 + token_lora_mapping,

461 + lora_ids,

462 + adapter_enabled,

463 + N_per_slice,

464 + K,

465 + topk_weights.numel(),

466 + top_k_num,

467 + max_loras_w,

468 + qcurr_hidden_states.stride(0),

469 + qcurr_hidden_states.stride(1),

470 + A0.stride(0),

471 + A0.stride(1),

472 + A0.stride(2),

473 + A0.stride(3),

474 + B0.stride(0),

475 + B0.stride(1),

476 + B0.stride(2),

477 + B0.stride(3),

478 + out_view.stride(0),

479 + out_view.stride(1),

480 + stride_tl_,

481 + stride_el,

482 + N_per_slice,

483 + token_mapping_factor=token_mapping_factor,

484 + naive_block_assignment=naive,

485 + MUL_ROUTED_WEIGHT=mul_routed_weight,

486 + BLOCK_M=BLOCK_M,

487 + BLOCK_R=BLOCK_R,

488 + actual_rank=rank,

489 + NPID_FACTOR=npid,

490 + BLOCK_N=block_n,

491 + BLOCK_K=block_k,

492 + ADD_INPUTS=add_inputs,

493 + num_warps=nw,

494 + num_stages=ns,

495 + )

496 +

497 +

498 +# ---------------------------------------------------------------------------

499 +# Small-batch (decode-style) fused MoE-LoRA kernel — sub-path of the

500 +# one_shot fast path.

501 +# ---------------------------------------------------------------------------

502 +

503 +

504 +@triton.heuristics({"EVEN_K": lambda args: args["K"] % args["BLOCK_K"] == 0})

505 +@triton.jit

506 +def _fused_moe_lora_small_batch_kernel(

507 + # ---- pointers ----

508 + x_ptr,

509 + A_ptrs,

510 + B_ptrs,

511 + out_ptr,

512 + topk_weights_ptr,

513 + expert_ids_ptr, # (num_tokens * top_k_num,)

514 + token_lora_mapping_ptr, # (num_tokens,)

515 + adapter_enabled_ptr,

516 + # ---- dims ----

517 + N,

518 + K,

519 + top_k_num,

520 + max_loras,

521 + work_total, # = pair_slices * n_chunks_per_pair_slice

522 + pair_slices, # = num_tokens * top_k_num * NUM_SLICES

523 + # ---- strides ----

524 + stride_xm,

525 + stride_xk,

526 + stride_A_lora,

527 + stride_A_expert,

528 + stride_A_r,

529 + stride_A_k,

530 + stride_B_lora,

531 + stride_B_expert,

532 + stride_B_n,

533 + stride_B_r,

534 + stride_om,

535 + stride_on,

536 + # ---- scalar (runtime ints, NOT constexpr) ----

537 + # n_tiles_per_program / n_chunks_per_pair_slice are deliberately

538 + # runtime: each distinct value would otherwise trigger a fresh Triton

539 + # compile -> fresh kernel binary -> fresh CUDA graph instance per

540 + # batch size. Production traces showed that variant explosion adding

541 + # ~5.9k graph instantiations on top of legacy. Runtime args mean one

542 + # shared binary across all chunk sizes.

543 + slice_n_offset,

544 + n_tiles_per_program,

545 + n_chunks_per_pair_slice,

546 + # ---- constexpr ----

547 + token_mapping_factor: tl.constexpr,

548 + MUL_ROUTED_WEIGHT: tl.constexpr,

549 + ADD_INPUTS: tl.constexpr,

550 + BLOCK_R: tl.constexpr,

551 + actual_rank: tl.constexpr,

552 + BLOCK_N: tl.constexpr,

553 + BLOCK_K: tl.constexpr,

554 + NUM_SLICES: tl.constexpr,

555 + EVEN_K: tl.constexpr,

556 +):

557 + """Persistent fused MoE-LoRA kernel for naive_block_assignment inputs.

558 +

559 + Each program owns one (pair × slice × n_chunk) work item. A "chunk"

560 + covers `n_tiles_per_program` consecutive output-N tiles, all of which

561 + share a single shrink — so the rank-vector is computed once per

562 + program and the A weights for that (lora, expert, slice) are loaded

563 + once instead of n_tiles_per_program times.

564 +

565 + The wrapper picks `n_tiles_per_program` to keep the grid close to

566 + 2*SM_count: at very small batch (work_total ≤ SM_count) the chunk

567 + size collapses to 1 and behaviour matches a per-tile GEMV; as batch

568 + grows the chunk grows so we trade some N-axis parallelism for shrink

569 + reuse. When `work_total` exceeds the launched grid, the outer stride

570 + loop drains the leftover work units serially.

571 + """

572 + pid = tl.program_id(axis=0)

573 + num_programs = tl.num_programs(axis=0)

574 +

575 + offs_r = tl.arange(0, BLOCK_R)

576 + rank_mask = offs_r < actual_rank

577 + # Clamp OOB rank lanes so they address row 0 of A/B; the mask zeros

578 + # the loaded values. Required when BLOCK_R > actual_rank (e.g. rank=4

579 + # padded to 16) -- without clamping, tl.load would address the next

580 + # expert's memory.

581 + safe_offs_r = tl.where(rank_mask, offs_r, 0)

582 + offs_k = tl.arange(0, BLOCK_K)

583 +

584 + # Persistent stride loop: when grid < work_total each program walks

585 + # multiple work items. When grid == work_total the loop runs exactly

586 + # once and the kernel degenerates to the per-tile GEMV.

587 + for work_id in range(pid, work_total, num_programs):

588 + n_chunk_idx = work_id % n_chunks_per_pair_slice

589 + pair_slice_idx = work_id // n_chunks_per_pair_slice

590 + # NUM_SLICES is constexpr (typ. 1 or 2) so divmod folds.

591 + pair_idx = pair_slice_idx // NUM_SLICES

592 + slice_id = pair_slice_idx % NUM_SLICES

593 +

594 + # Resolve lora_id / expert_id; skip the body for inactive lanes.

595 + # Using a single `valid` flag instead of early `return` keeps the

596 + # outer stride loop alive — `return` would exit the whole program

597 + # and skip later work items assigned to this SM.

598 + token_idx = pair_idx // top_k_num

599 + lora_id = tl.load(token_lora_mapping_ptr + token_idx)

600 + valid = (lora_id >= 0) & (lora_id < max_loras)

601 + enabled = tl.load(adapter_enabled_ptr + tl.where(valid, lora_id, 0))

602 + valid = valid & (enabled != 0)

603 + expert_id = tl.load(expert_ids_ptr + pair_idx)

604 + valid = valid & (expert_id >= 0)

605 +

606 + if valid:

607 + cur_A_ptr = tl.load(A_ptrs + slice_id).to(

608 + tl.pointer_type(out_ptr.dtype.element_ty)

609 + )

610 + cur_B_ptr = tl.load(B_ptrs + slice_id).to(

611 + tl.pointer_type(out_ptr.dtype.element_ty)

612 + )

613 + A_base = cur_A_ptr + lora_id * stride_A_lora + expert_id * stride_A_expert

614 + B_base = cur_B_ptr + lora_id * stride_B_lora + expert_id * stride_B_expert

615 +

616 + x_row = pair_idx // token_mapping_factor

617 + x_row_ptr = x_ptr + x_row * stride_xm

618 +

619 + # SHRINK GEMV (once per program; reused across n_tiles_per_program

620 + # expand tiles below). Sum-reduction over BLOCK_K with fp32

621 + # accumulator — same precision path as the one_shot kernel.

622 + rank_vec = tl.zeros((BLOCK_R,), dtype=tl.float32)

623 + if EVEN_K:

624 + for kb in range(0, K, BLOCK_K):

625 + cur_k = kb + offs_k

626 + x_tile = tl.load(x_row_ptr + cur_k * stride_xk).to(tl.float32)

627 + a_tile = tl.load(

628 + A_base

629 + + safe_offs_r[:, None] * stride_A_r

630 + + cur_k[None, :] * stride_A_k,

631 + mask=rank_mask[:, None],

632 + other=0.0,

633 + ).to(tl.float32)

634 + rank_vec += tl.sum(a_tile * x_tile[None, :], axis=1)

635 + else:

636 + for kb in range(0, K, BLOCK_K):

637 + cur_k = kb + offs_k

638 + k_mask = cur_k < K

639 + x_tile = tl.load(

640 + x_row_ptr + cur_k * stride_xk, mask=k_mask, other=0.0

641 + ).to(tl.float32)

642 + a_tile = tl.load(

643 + A_base

644 + + safe_offs_r[:, None] * stride_A_r

645 + + cur_k[None, :] * stride_A_k,

646 + mask=rank_mask[:, None] & k_mask[None, :],

647 + other=0.0,

648 + ).to(tl.float32)

649 + rank_vec += tl.sum(a_tile * x_tile[None, :], axis=1)

650 +

651 + # EXPAND: walk n_tiles_per_program consecutive output-N tiles

652 + # using the same rank_vec. The loop is a runtime range (not

653 + # tl.static_range) so a single compiled kernel handles every

654 + # chunk size — see the note on the kernel signature.

655 + n_tile_start = n_chunk_idx * n_tiles_per_program

656 + out_row_ptr = out_ptr + slice_id * slice_n_offset + pair_idx * stride_om

657 +

658 + if MUL_ROUTED_WEIGHT:

659 + moe_w = tl.load(topk_weights_ptr + pair_idx).to(tl.float32)

660 +

661 + for nt in range(n_tiles_per_program):

662 + n_lo = (n_tile_start + nt) * BLOCK_N

663 + if n_lo < N:

664 + offs_n = n_lo + tl.arange(0, BLOCK_N)

665 + n_mask = offs_n < N

666 + b_tile = tl.load(

667 + B_base

668 + + offs_n[:, None] * stride_B_n

669 + + safe_offs_r[None, :] * stride_B_r,

670 + mask=n_mask[:, None] & rank_mask[None, :],

671 + other=0.0,

672 + ).to(tl.float32)

673 + out_tile = tl.sum(b_tile * rank_vec[None, :], axis=1)

674 +

675 + if MUL_ROUTED_WEIGHT:

676 + out_tile = out_tile * moe_w

677 +

678 + out_ptrs = out_row_ptr + offs_n * stride_on

679 + if ADD_INPUTS:

680 + prev = tl.load(out_ptrs, mask=n_mask, other=0.0).to(tl.float32)

681 + tl.store(

682 + out_ptrs,

683 + (prev + out_tile).to(out_ptr.dtype.element_ty),

684 + mask=n_mask,

685 + )

686 + else:

687 + tl.store(

688 + out_ptrs,

689 + out_tile.to(out_ptr.dtype.element_ty),

690 + mask=n_mask,

691 + )

692 +

693 +

694 +def _pick_small_batch_chunk(pair_slices: int, N_tiles: int, sm_count: int) -> int:

695 + """Pick `n_tiles_per_program` so the launched grid stays near

696 + 2*SM_count.

697 +

698 + Sizes for occupancy first (more programs in flight → better latency

699 + hiding for the K-loop A/x loads). Once the per-tile grid already

700 + exceeds 2*SM_count we increase the chunk size to amortise the shrink

701 + cost — at that point the GPU is saturated by per-program work and

702 + packing more tiles per program lets the rank_vec be reused.

703 + """

704 + target_grid = max(1, 2 * sm_count)

705 + total_work = pair_slices * N_tiles

706 + if total_work <= target_grid:

707 + return 1

708 + ntpp = (total_work + target_grid - 1) // target_grid

709 + return min(ntpp, N_tiles)

710 +

711 +

712 +def _run_fused_moe_lora_small_batch(

713 + output: torch.Tensor,

714 + qcurr_hidden_states: torch.Tensor,

715 + lora_a_stacked: list[torch.Tensor],

716 + lora_b_stacked: list[torch.Tensor],

717 + topk_weights: torch.Tensor,

718 + expert_ids_flat: torch.Tensor, # (num_tokens * top_k_num,)

719 + token_lora_mapping: torch.Tensor,

720 + top_k_num: int,

721 + adapter_enabled: torch.Tensor,

722 + mul_routed_weight: bool,

723 + add_inputs: bool = True,

724 +) -> None:

725 + """Small-batch GEMV-style wrapper. Naive-block-assignment inputs only.

726 +

727 + Shape contract matches `_run_fused_moe_lora_one_shot`: `output` is

728 + `(num_tokens, top_k_num, num_slices * N_per_slice)` with

729 + contiguous-style strides, `expert_ids_flat` is the flattened

730 + `topk_ids` of shape `(num_tokens * top_k_num,)`, and the

731 + rank-padded LoRA weights live in `lora_a_stacked` /

732 + `lora_b_stacked`.

733 +

734 + The kernel is persistent over (pair × slice × n_chunk) work items —

735 + each program does one shrink and reuses the rank vector across

736 + `n_tiles_per_program` expand tiles. The chunk size scales with the

737 + pair-slice count so very small batches keep per-tile parallelism

738 + while medium batches cut redundant shrinks.

739 + """

740 + num_slices = len(lora_a_stacked)

741 + device = qcurr_hidden_states.device

742 +

743 + A0 = lora_a_stacked[0]

744 + B0 = lora_b_stacked[0]

745 + max_loras_w = A0.shape[0]

746 + rank = A0.shape[2]

747 + K = A0.shape[3]

748 + N_per_slice = B0.shape[2]

749 +

750 + # Rank padding: floor 16 (tensor-core min K), ceil to next pow2. The

751 + # ≤64 cap is set conservatively for the prototype: at rank 64 the

752 + # per-program register footprint is rank_vec(64 fp32) + b_tile(BLOCK_N

753 + # × 64 fp32) = e.g. 128*64*4 = 32 KiB, comfortably within the 64 KiB

754 + # register file even with num_warps=8. Doubling to 128 would push us

755 + # against the limit and require shared-memory staging.

756 + assert rank <= 64, f"fused_moe_lora_small_batch supports rank<=64; got rank={rank}"

757 + BLOCK_R = max(triton.next_power_of_2(rank), 16)

758 +

759 + num_tokens = topk_weights.shape[0]

760 + M_grid = num_tokens * top_k_num

761 + if M_grid == 0 or num_slices == 0:

762 + return

763 +

764 + token_mapping_factor = 1 if mul_routed_weight else top_k_num

765 +

766 + A_ptrs = _get_ptr(lora_a_stacked, device)

767 + B_ptrs = _get_ptr(lora_b_stacked, device)

768 +

769 + assert output.dim() == 3, f"output must be 3-D, got {output.shape}"

770 + assert output.stride(0) == output.shape[1] * output.stride(1), (

771 + "fused_moe_lora_small_batch requires output.stride(0) == "

772 + f"top_k*stride(1); got shape={output.shape} strides={output.stride()}"

773 + )

774 + out_view = output.view(-1, output.shape[-1])

775 +

776 + # Block sizes. BLOCK_N=128 matches the one_shot's expand tile and gives

777 + # 6-24 N tiles for typical N ∈ [768, 3072], enough to saturate the SM

778 + # array once M_grid * num_slices reaches ~SM_count. BLOCK_K=128 halves

779 + # the K-loop trip count vs 64 and pays for itself once K ≥ 1024 (the

780 + # only regime we care about — hidden sizes are always large here).

781 + BLOCK_N = 128

782 + BLOCK_K = 128

783 + nw = 4

784 + ns = 3

785 +

786 + N_tiles = triton.cdiv(N_per_slice, BLOCK_N)

787 + pair_slices = M_grid * num_slices

788 +

789 + sm_count = torch.cuda.get_device_properties(device).multi_processor_count

790 + n_tiles_per_program = _pick_small_batch_chunk(pair_slices, N_tiles, sm_count)

791 + n_chunks = triton.cdiv(N_tiles, n_tiles_per_program)

792 + work_total = pair_slices * n_chunks

793 +

794 + # Grid sizing: keep parallelism uncapped when work_total is small (so

795 + # very small batches still spread across SMs); cap at 2*SM_count once

796 + # we have plenty of work, letting the in-kernel stride loop drain the

797 + # remainder.

798 + grid_size = min(work_total, max(1, 2 * sm_count))

799 + grid = (grid_size,)

800 +

801 + _fused_moe_lora_small_batch_kernel[grid](

802 + qcurr_hidden_states,

803 + A_ptrs,

804 + B_ptrs,

805 + out_view,

806 + topk_weights,

807 + expert_ids_flat,

808 + token_lora_mapping,

809 + adapter_enabled,

810 + N_per_slice,

811 + K,

812 + top_k_num,

813 + max_loras_w,

814 + work_total,

815 + pair_slices,

816 + qcurr_hidden_states.stride(0),

817 + qcurr_hidden_states.stride(1),

818 + A0.stride(0),

819 + A0.stride(1),

820 + A0.stride(2),

821 + A0.stride(3),

822 + B0.stride(0),

823 + B0.stride(1),

824 + B0.stride(2),

825 + B0.stride(3),

826 + out_view.stride(0),

827 + out_view.stride(1),

828 + N_per_slice,

829 + n_tiles_per_program,

830 + n_chunks,

831 + token_mapping_factor=token_mapping_factor,

832 + MUL_ROUTED_WEIGHT=mul_routed_weight,

833 + ADD_INPUTS=add_inputs,

834 + BLOCK_R=BLOCK_R,

835 + actual_rank=rank,

836 + BLOCK_N=BLOCK_N,

837 + BLOCK_K=BLOCK_K,

838 + NUM_SLICES=num_slices,

839 + num_warps=nw,

840 + num_stages=ns,

841 + )

842 +

843 +

844 def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):

845 """

846 `_LORA_PTR_DICT` collects the required information during `profile_run`,

@@ -706,6 +1442,7 @@ def _fused_moe_lora(

1442 mul_routed_weight: bool = False,

1443 fully_sharded: bool = False,

1444 offset: int = 0,

1445 + add_inputs: bool = True,

1446 ) -> None:

1447 assert len(lora_a_stacked) == len(lora_b_stacked) > 0

1448 assert topk_weights.dim() == qcurr_hidden_states.dim() == 2

@@ -728,6 +1465,59 @@ def _fused_moe_lora(

1465 )

1466 assert output.shape[0] == topk_weights.shape[0]

1467 assert top_k_num == topk_weights.shape[1]

1468 +

1469 + # Fast path: single fused kernel

1470 + if not fully_sharded:

1471 + M_pairs = topk_weights.numel()

1472 + if (

1473 + sorted_token_ids is None

1474 + and max_lora_rank <= 64

1475 + and M_pairs * max_lora_rank <= 1024

1476 + ):

1477 + _run_fused_moe_lora_small_batch(

1478 + output,

1479 + qcurr_hidden_states,

1480 + lora_a_stacked,

1481 + lora_b_stacked,

1482 + topk_weights,

1483 + expert_ids,

1484 + token_lora_mapping,

1485 + top_k_num,

1486 + adapter_enabled,

1487 + mul_routed_weight,

1488 + add_inputs=add_inputs,

1489 + )

1490 + return

1491 + # shrink/expand BLOCK_SIZE_M must match the block_size that

1492 + # moe_lora_align_block_size used; both shrink and expand pass the

1493 + # same value (asserted by `shrink_block_size_m == expand_block_size_m`

1494 + # below).

1495 + _run_fused_moe_lora_one_shot(

1496 + output,

1497 + qcurr_hidden_states,

1498 + lora_a_stacked,

1499 + lora_b_stacked,

1500 + topk_weights,

1501 + sorted_token_ids,

1502 + expert_ids,

1503 + num_tokens_post_padded,

1504 + token_lora_mapping,

1505 + max_lora_rank,

1506 + top_k_num,

1507 + lora_ids,

1508 + num_active_loras,

1509 + adapter_enabled,

1510 + mul_routed_weight,

1511 + shrink_block_size_m,

1512 + add_inputs=add_inputs,

1513 + )

1514 + return

1515 +

1516 + assert add_inputs, (

1517 + "fused_moe_lora(add_inputs=False) is only supported on the "

1518 + "fully_sharded=False fast path"

1519 + )

1520 +

1521 device = qcurr_hidden_states.device

1522 num_slices = len(lora_a_stacked)

1523 w1_lora_b_stacked = lora_b_stacked[0]

@@ -894,6 +1684,7 @@ def _fused_moe_lora_fake(

1684 mul_routed_weight: bool = False,

1685 fully_sharded: bool = False,

1686 offset: int = 0,

1687 + add_inputs: bool = True,

1688 ) -> None:

1689 return

1690

vllm/lora/punica_wrapper/punica_base.py

@@ -514,6 +514,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

514 num_slices: int,

515 fully_sharded: bool,

516 use_tuned_config: bool,

517 + add_inputs: bool = True,

518 token_lora_mapping: torch.Tensor | None = None,

519 ) -> tuple[

520 torch.Tensor | None,

@@ -554,6 +555,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

555 fully_sharded: bool,

556 tp_rank: int,

557 use_tuned_config: bool,

558 + add_inputs: bool = True,

559 ) -> None:

560 """Apply w2 LoRA to y (intermediate_cache3) in-place before moe_sum.

561

vllm/lora/punica_wrapper/punica_gpu.py

@@ -435,6 +435,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

435 fully_sharded: bool = False,

436 offset: int = 0,

437 token_lora_mapping: torch.Tensor | None = None,

438 + add_inputs: bool = True,

439 ):

440 """

441 Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.

@@ -484,6 +485,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

485 mul_routed_weight,

486 fully_sharded,

487 offset,

488 + add_inputs,

489 )

490

491 def add_lora_w13(

@@ -506,6 +508,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

508 num_slices: int,

509 fully_sharded: bool,

510 use_tuned_config: bool,

511 + add_inputs: bool = True,

512 token_lora_mapping: torch.Tensor | None = None,

513 ) -> tuple[

514 torch.Tensor | None,

@@ -610,6 +613,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

613 adapter_enabled,

614 fully_sharded=fully_sharded,

615 token_lora_mapping=token_lora_mapping,

616 + add_inputs=add_inputs,

617 )

618

619 return (

@@ -640,6 +644,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

644 fully_sharded: bool,

645 tp_rank: int,

646 use_tuned_config: bool,

647 + add_inputs: bool = True,

648 ) -> None:

649 import functools

650

@@ -722,4 +727,5 @@ class PunicaWrapperGPU(PunicaWrapperBase):

727 fully_sharded=fully_sharded,

728 offset=offset,

729 token_lora_mapping=token_lora_mapping,

730 + add_inputs=add_inputs,

731 )

vllm/lora/punica_wrapper/punica_xpu.py

@@ -390,6 +390,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

390 fully_sharded: bool = False,

391 offset: int = 0,

392 token_lora_mapping: torch.Tensor | None = None,

393 + add_inputs: bool = True,

394 ):

395 """

396 Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.

@@ -439,6 +440,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

440 mul_routed_weight,

441 fully_sharded,

442 offset,

443 + add_inputs,

444 )

445

446 def add_lora_w13(

@@ -461,6 +463,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

463 num_slices: int,

464 fully_sharded: bool,

465 use_tuned_config: bool,

466 + add_inputs: bool = True,

467 token_lora_mapping: torch.Tensor | None = None,

468 ) -> tuple[

469 torch.Tensor | None,

@@ -594,6 +597,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

597 fully_sharded: bool,

598 tp_rank: int,

599 use_tuned_config: bool,

600 + add_inputs: bool = True,

601 ) -> None:

602 import functools

603

vllm/model_executor/layers/fused_moe/experts/lora_context.py

+10

@@ -43,6 +43,16 @@ class MoELoRAContext:

43 # try_get_optimal_moe_lora_config for Triton kernel tile configs.

44 use_tuned_config: bool

46 + # Optional dual-stream support for overlapping each (base GEMM, LoRA)

47 + # pair. When aux_stream is None, the experts.apply() path runs the

48 + # original sequential schedule. When set, base GEMM runs on the default

49 + # stream and the LoRA fast-path writes the delta into a fresh buffer on

50 + # aux_stream, which the default stream sums in afterwards.

51 + # Events are paired one-per-overlap-pair: events[0,1] for w13,

52 + # events[2,3] for w2, so the two pairs do not race on the same event.

53 + aux_stream: torch.cuda.Stream | None = None

54 + events: tuple[torch.cuda.Event, ...] | None = None

55 +

56 # Per-rank token→LoRA mapping after EP dispatch. Set by

57 # FusedMoEPrepareAndFinalizeModular.prepare() when EP+LoRA is active, read

58 # by LoRAExpertsMixin helpers in place of punica_wrapper's global mapping.

vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py

@@ -45,6 +45,7 @@ class LoRAExpertsMixin:

45 w2: torch.Tensor,

46 num_tokens: int,

47 top_k_num: int,

48 + add_inputs: bool = True,

49 ) -> tuple[

50 torch.Tensor | None,

51 torch.Tensor | None,

@@ -70,6 +71,7 @@ class LoRAExpertsMixin:

71 lora_context.w13_num_slices,

72 lora_context.fully_sharded,

73 lora_context.use_tuned_config,

74 + add_inputs=add_inputs,

75 token_lora_mapping=lora_context.local_token_lora_mapping,

76 )

@@ -88,6 +90,7 @@ class LoRAExpertsMixin:

90 w1: torch.Tensor,

91 w2: torch.Tensor,

92 top_k_num: int,

93 + add_inputs: bool = True,

94 ) -> None:

95 lora_context.punica_wrapper.add_lora_w2(

96 y,

@@ -109,4 +112,5 @@ class LoRAExpertsMixin:

112 lora_context.fully_sharded,

113 lora_context.tp_rank,

114 lora_context.use_tuned_config,

115 + add_inputs=add_inputs,

116 )

vllm/model_executor/layers/fused_moe/experts/triton_moe.py

+155 -76

@@ -48,6 +48,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

48 )

49 from vllm.platforms import current_platform

50 from vllm.triton_utils import tl

51 +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel

54 class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):

@@ -247,55 +248,99 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):

248 )

249 )

250

250 - invoke_fused_moe_triton_kernel(

251 - hidden_states,

252 - w1,

253 - intermediate_cache1,

254 - a1q_scale if a1q_scale is not None else self.a1_scale,

255 - self.w1_scale,

256 - None, # topk_weights

257 - sorted_token_ids,

258 - expert_ids,

259 - num_tokens_post_padded,

260 - False, # mul_routed_weights

261 - top_k_num,

262 - config,

263 - compute_type=compute_type,

264 - use_fp8_w8a8=self.quant_config.use_fp8_w8a8,

265 - use_int8_w8a8=self.quant_config.use_int8_w8a8,

266 - use_int8_w8a16=self.quant_config.use_int8_w8a16,

267 - use_int4_w4a16=self.quant_config.use_int4_w4a16,

268 - per_channel_quant=self.per_act_token_quant,

269 - block_shape=self.block_shape,

270 - B_bias=self.w1_bias,

271 - )

251 + # LoRA w13: applied to intermediate_cache1 before activation. When

252 + # the LoRA layer requested a dual-stream schedule, we run base w13

253 + # GEMM on the default stream and the LoRA fast-path on aux_stream;

254 + # the LoRA writes its delta into a fresh zero buffer (add_inputs=

255 + # False) and we sum it into intermediate_cache1 after both finish.

256

273 - # LoRA w13: applied to intermediate_cache1 before activation, using

274 - # hidden_states as the lora_a input. moe_lora_align_block_size is

275 - # called once here and results reused for the w2 LoRA below.

257 sorted_token_ids_lora = None

258 expert_ids_lora = None

259 num_tokens_post_padded_lora = None

260 token_lora_mapping = None

261 lora_context = self._lora_context

281 - if lora_context is not None:

262 +

263 + def _base_w13_fn():

264 + invoke_fused_moe_triton_kernel(

265 + hidden_states,

266 + w1,

267 + intermediate_cache1,

268 + a1q_scale if a1q_scale is not None else self.a1_scale,

269 + self.w1_scale,

270 + None, # topk_weights

271 + sorted_token_ids,

272 + expert_ids,

273 + num_tokens_post_padded,

274 + False, # mul_routed_weights

275 + top_k_num,

276 + config,

277 + compute_type=compute_type,

278 + use_fp8_w8a8=self.quant_config.use_fp8_w8a8,

279 + use_int8_w8a8=self.quant_config.use_int8_w8a8,

280 + use_int8_w8a16=self.quant_config.use_int8_w8a16,

281 + use_int4_w4a16=self.quant_config.use_int4_w4a16,

282 + per_channel_quant=self.per_act_token_quant,

283 + block_shape=self.block_shape,

284 + B_bias=self.w1_bias,

285 + )

286 +

287 + if lora_context is not None and lora_context.aux_stream is not None:

288 + # add_inputs=False: kernel overwrites lora_delta_w13. zeros (not

289 + # empty) so untouched rows -- e.g. blocks where every program

290 + # early-exits because lora_id<0 -- stay at zero and the trailing

291 + # add_() is a no-op there.

292 + lora_delta_w13 = torch.zeros_like(intermediate_cache1)

293 +

294 + def _lora_w13_fn():

295 + return self.apply_w13_lora(

296 + lora_context,

297 + y=lora_delta_w13,

298 + x=hidden_states,

299 + topk_ids=topk_ids,

300 + topk_weights=topk_weights,

301 + expert_map=expert_map,

302 + w1=w1,

303 + w2=w2,

304 + num_tokens=num_tokens,

305 + top_k_num=top_k_num,

306 + add_inputs=False,

307 + )

308 +

309 + assert lora_context.events is not None

310 + _, lora_meta = maybe_execute_in_parallel(

311 + _base_w13_fn,

312 + _lora_w13_fn,

313 + lora_context.events[0],

314 + lora_context.events[1],

315 + lora_context.aux_stream,

316 + )

317 (

318 sorted_token_ids_lora,

319 expert_ids_lora,

320 num_tokens_post_padded_lora,

321 token_lora_mapping,

287 - ) = self.apply_w13_lora(

288 - lora_context,

289 - y=intermediate_cache1,

290 - x=hidden_states,

291 - topk_ids=topk_ids,

292 - topk_weights=topk_weights,

293 - expert_map=expert_map,

294 - w1=w1,

295 - w2=w2,

296 - num_tokens=num_tokens,

297 - top_k_num=top_k_num,

298 - )

322 + ) = lora_meta

323 + intermediate_cache1.add_(lora_delta_w13)

324 + else:

325 + _base_w13_fn()

326 + if lora_context is not None:

327 + (

328 + sorted_token_ids_lora,

329 + expert_ids_lora,

330 + num_tokens_post_padded_lora,

331 + token_lora_mapping,

332 + ) = self.apply_w13_lora(

333 + lora_context,

334 + y=intermediate_cache1,

335 + x=hidden_states,

336 + topk_ids=topk_ids,

337 + topk_weights=topk_weights,

338 + expert_map=expert_map,

339 + w1=w1,

340 + w2=w2,

341 + num_tokens=num_tokens,

342 + top_k_num=top_k_num,

343 + )

344

345 a2q_scale: torch.Tensor | None = None

346

@@ -328,48 +373,82 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):

373 quantization_emulation=self.quantization_emulation,

374 )

375

331 - invoke_fused_moe_triton_kernel(

332 - qintermediate_cache2,

333 - w2,

334 - intermediate_cache3,

335 - a2q_scale,

336 - self.w2_scale,

337 - topk_weights,

338 - sorted_token_ids,

339 - expert_ids,

340 - num_tokens_post_padded,

341 - not apply_router_weight_on_input,

342 - 1,

343 - config,

344 - compute_type=compute_type,

345 - use_fp8_w8a8=self.quant_config.use_fp8_w8a8,

346 - use_int8_w8a8=self.quant_config.use_int8_w8a8,

347 - use_int8_w8a16=self.quant_config.use_int8_w8a16,

348 - use_int4_w4a16=self.quant_config.use_int4_w4a16,

349 - per_channel_quant=self.per_act_token_quant,

350 - block_shape=self.block_shape,

351 - B_bias=self.w2_bias,

352 - )

353 -

376 # LoRA w2: applied to intermediate_cache3 before moe_sum, using the

377 # unquantized intermediate_cache2 as the lora_a input. Reuses the

356 - # sorted_token_ids_lora computed above.

357 - if lora_context is not None:

358 - self.apply_w2_lora(

359 - lora_context,

360 - y=intermediate_cache3,

361 - x=intermediate_cache2,

362 - topk_weights=topk_weights,

363 - sorted_token_ids_lora=sorted_token_ids_lora,

364 - expert_ids_lora=expert_ids_lora,

365 - num_tokens_post_padded_lora=num_tokens_post_padded_lora,

366 - token_lora_mapping=token_lora_mapping,

367 - num_tokens=num_tokens,

368 - w1=w1,

369 - w2=w2,

370 - top_k_num=top_k_num,

378 + # sorted_token_ids_lora computed above. Same dual-stream pattern as

379 + # the w13 pair: base GEMM on default stream, LoRA delta on aux,

380 + # join via .add_() into intermediate_cache3.

381 + def _base_w2_fn():

382 + invoke_fused_moe_triton_kernel(

383 + qintermediate_cache2,

384 + w2,

385 + intermediate_cache3,

386 + a2q_scale,

387 + self.w2_scale,

388 + topk_weights,

389 + sorted_token_ids,

390 + expert_ids,

391 + num_tokens_post_padded,

392 + not apply_router_weight_on_input,

393 + 1,

394 + config,

395 + compute_type=compute_type,

396 + use_fp8_w8a8=self.quant_config.use_fp8_w8a8,

397 + use_int8_w8a8=self.quant_config.use_int8_w8a8,

398 + use_int8_w8a16=self.quant_config.use_int8_w8a16,

399 + use_int4_w4a16=self.quant_config.use_int4_w4a16,

400 + per_channel_quant=self.per_act_token_quant,

401 + block_shape=self.block_shape,

402 + B_bias=self.w2_bias,

403 )

404

405 + if lora_context is not None and lora_context.aux_stream is not None:

406 + lora_delta_w2 = torch.zeros_like(intermediate_cache3)

407 +

408 + def _lora_w2_fn():

409 + self.apply_w2_lora(

410 + lora_context,

411 + y=lora_delta_w2,

412 + x=intermediate_cache2,

413 + topk_weights=topk_weights,

414 + sorted_token_ids_lora=sorted_token_ids_lora,

415 + expert_ids_lora=expert_ids_lora,

416 + num_tokens_post_padded_lora=num_tokens_post_padded_lora,

417 + token_lora_mapping=token_lora_mapping,

418 + num_tokens=num_tokens,

419 + w1=w1,

420 + w2=w2,

421 + top_k_num=top_k_num,

422 + add_inputs=False,

423 + )

424 +

425 + assert lora_context.events is not None

426 + maybe_execute_in_parallel(

427 + _base_w2_fn,

428 + _lora_w2_fn,

429 + lora_context.events[2],

430 + lora_context.events[3],

431 + lora_context.aux_stream,

432 + )

433 + intermediate_cache3.add_(lora_delta_w2)

434 + else:

435 + _base_w2_fn()

436 + if lora_context is not None:

437 + self.apply_w2_lora(

438 + lora_context,

439 + y=intermediate_cache3,

440 + x=intermediate_cache2,

441 + topk_weights=topk_weights,

442 + sorted_token_ids_lora=sorted_token_ids_lora,

443 + expert_ids_lora=expert_ids_lora,

444 + num_tokens_post_padded_lora=num_tokens_post_padded_lora,

445 + token_lora_mapping=token_lora_mapping,

446 + num_tokens=num_tokens,

447 + w1=w1,

448 + w2=w2,

449 + top_k_num=top_k_num,

450 + )

451 +

452 # separate function is required for MoE + LoRA

453 self.moe_sum(intermediate_cache3, output)

454

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -164,6 +164,8 @@ def _moe_forward_shared_fake(

164 return shared_out, fused_out

165

166

167 +# NOTE: `moe_forward` and `moe_forward_shared` being opaque custom ops is a

168 +# load-bearing assumption for the MoE-LoRA dual-stream path.

169 direct_register_custom_op(

170 op_name="moe_forward",

171 op_func=_moe_forward,