a55caf6ae9 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -0,0 +1,1032 @@

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

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

3 +

4 +

5 +from typing import List # noqa: UP035

6 +

7 +import torch

8 +

9 +from vllm.distributed import (

10 + tensor_model_parallel_all_gather,

11 + tensor_model_parallel_all_reduce,

12 +)

13 +from vllm.triton_utils import tl, triton

14 +from vllm.utils.torch_utils import direct_register_custom_op

15 +

16 +from .utils import supports_pdl

17 +

18 +

19 +@triton.jit

20 +def _get_lora_id(

21 + lora_ids,

22 + token_lora_mapping_ptr,

23 + lora_idx,

24 + pid_m,

25 + top_k_num,

26 + naive_block_assignment: tl.constexpr,

27 +):

28 + """Returns lora_id"""

29 + if naive_block_assignment:

30 + token_idx = pid_m // top_k_num

31 + return tl.load(token_lora_mapping_ptr + token_idx)

32 + else:

33 + return tl.load(lora_ids + lora_idx)

34 +

35 +

36 +@triton.jit

37 +def _get_expert_id(

38 + expert_ids_ptr,

39 + lora_id,

40 + pid_m,

41 + stride_el,

42 + max_loras,

43 + naive_block_assignment: tl.constexpr,

44 +):

45 + """Returns expert_id"""

46 + if naive_block_assignment:

47 + return tl.load(expert_ids_ptr + pid_m)

48 + else:

49 + ind = lora_id * stride_el + pid_m

50 + return tl.load(expert_ids_ptr + ind, ind < max_loras * stride_el, -1)

51 +

52 +

53 +@triton.jit

54 +def _get_token_offs(

55 + sorted_token_ids_ptr,

56 + lora_id,

57 + pid_m,

58 + offs,

59 + stride_tl,

60 + max_loras,

61 + num_valid_tokens,

62 + naive_block_assignment: tl.constexpr,

63 + BLOCK_SIZE_M: tl.constexpr,

64 +):

65 + """Returns token offsets"""

66 + if naive_block_assignment:

67 + return tl.where(offs == 0, pid_m, num_valid_tokens)

68 + else:

69 + offs_token_id = pid_m * BLOCK_SIZE_M + offs

70 + token_ind = stride_tl * lora_id + offs_token_id

71 + return tl.load(

72 + sorted_token_ids_ptr + token_ind, token_ind < max_loras * stride_tl, 0

73 + )

74 +

75 +

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

77 +

78 +

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

80 + """

81 + `_LORA_PTR_DICT` collects the required information during `profile_run`,

82 + After this, it remains constant and subsequent usage is through LUT.

83 + Refer to:

84 + https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py

85 + """

86 + key = tuple(lora_weight.data_ptr() for lora_weight in lora_weights)

87 +

88 + if (ptr_tensor := _LORA_PTR_DICT.get(key)) is not None:

89 + return ptr_tensor

90 +

91 + tensor_ptrs = []

92 + for lora_weight in lora_weights:

93 + tensor_ptrs.append(lora_weight.data_ptr())

94 + ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)

95 +

96 + _LORA_PTR_DICT[key] = ptr_tensor

97 + return _LORA_PTR_DICT.get(key)

98 +

99 +

100 +def _adjust_kernel_inputs(

101 + num_active_loras: int,

102 + sorted_token_ids: torch.Tensor | None,

103 + expert_ids: torch.Tensor,

104 +):

105 + """

106 + helper function to adjust kernel inputs when sorted_token_ids is None

107 + """

108 + if sorted_token_ids is None:

109 + stride_tl = 0

110 + stride_el = 0

111 + grid_lora_dim = 1

112 + else:

113 + stride_tl = sorted_token_ids.stride(0)

114 + stride_el = expert_ids.stride(0)

115 + grid_lora_dim = num_active_loras

116 + return grid_lora_dim, stride_tl, stride_el

117 +

118 +

119 +@triton.jit(

120 + do_not_specialize=[

121 + "num_valid_tokens",

122 + "EM",

123 + "stride_tl",

124 + "stride_el",

125 + "slice_a_size",

126 + "slice_c_size",

127 + ]

128 +)

129 +def _fused_moe_lora_kernel_fp8(

130 + a_ptr,

131 + b_ptr,

132 + c_ptr,

133 + a_scale_ptr,

134 + b_scale_ptr,

135 + topk_weights_ptr,

136 + sorted_token_ids_ptr,

137 + expert_ids_ptr,

138 + num_tokens_post_padded_ptr,

139 + token_lora_mapping_ptr,

140 + # Matrix dimensions

141 + N,

142 + K,

143 + EM,

144 + num_valid_tokens,

145 + num_experts,

146 + top_k_num,

147 + lora_ids,

148 + adapter_enabled,

149 + max_loras, # <<< PR2: rename, used for masks when grid axis-2 != max_loras

150 + # The stride variables represent how much to increase the ptr by when

151 + # moving by 1 element in a particular dimension. E.g. `stride_am` is

152 + # how much to increase `a_ptr` by to get the element one row down

153 + # (A has M rows).

154 + stride_am,

155 + stride_ak,

156 + stride_bl,

157 + stride_be,

158 + stride_bk,

159 + stride_bn,

160 + stride_cm,

161 + stride_cn,

162 + stride_tl,

163 + stride_el,

164 + stride_asm,

165 + stride_ask,

166 + stride_bsl,

167 + stride_bse,

168 + stride_bsk,

169 + stride_bsn,

170 + # block size for block-wise quantization

171 + group_n: tl.constexpr,

172 + group_k: tl.constexpr,

173 + slice_a_size,

174 + slice_c_size,

175 + # Meta-parameters

176 + num_slice_a: tl.constexpr,

177 + num_slice_c: tl.constexpr,

178 + # top_k_num or 1 depending on input token

179 + # is expanded by top_k or not

180 + token_mapping_factor: tl.constexpr,

181 + # whether use naive block assignment

182 + naive_block_assignment: tl.constexpr,

183 + MUL_ROUTED_WEIGHT: tl.constexpr,

184 + ADD_INPUTS: tl.constexpr,

185 + USE_B_L2_CACHE: tl.constexpr, # new, enable .ca load for B

186 + BLOCK_SIZE_M: tl.constexpr,

187 + BLOCK_SIZE_N: tl.constexpr,

188 + BLOCK_SIZE_K: tl.constexpr,

189 + GROUP_SIZE_M: tl.constexpr,

190 + SPLIT_K: tl.constexpr,

191 + USE_GDC: tl.constexpr,

192 + launch_pdl: tl.constexpr,

193 + IS_PRIMARY: tl.constexpr,

194 + use_fp8_w8a8: tl.constexpr,

195 + use_int8_w8a8: tl.constexpr,

196 + use_int8_w8a16: tl.constexpr,

197 + per_channel_quant: tl.constexpr,

198 +):

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

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

201 + grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)

202 +

203 + # calculate pid_m,pid_n

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

205 + pid_sk = pid % SPLIT_K

206 + pid_m_n = pid // SPLIT_K

207 + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)

208 + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

209 +

210 + num_pid_in_group = GROUP_SIZE_M * num_pid_n

211 + group_id = pid_m_n // num_pid_in_group

212 + first_pid_m = group_id * GROUP_SIZE_M

213 + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)

214 + pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)

215 + pid_n = (pid_m_n % num_pid_in_group) // group_size_m

216 +

217 + offs = tl.arange(0, BLOCK_SIZE_M).to(tl.int64)

218 +

219 + # Get lora_id

220 + lora_id = _get_lora_id(

221 + lora_ids,

222 + token_lora_mapping_ptr,

223 + lora_idx,

224 + pid_m,

225 + top_k_num,

226 + naive_block_assignment,

227 + )

228 + if lora_id == -1:

229 + return

230 + moe_enabled = tl.load(adapter_enabled + lora_id)

231 + if moe_enabled == 0:

232 + return

233 + if lora_id >= max_loras:

234 + return

235 +

236 + # Non-naive only: check num_tokens_post_padded

237 + if not naive_block_assignment:

238 + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_id)

239 + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:

240 + return

241 +

242 + # Get expert_id

243 + expert_id = _get_expert_id(

244 + expert_ids_ptr,

245 + lora_id,

246 + pid_m,

247 + stride_el,

248 + max_loras,

249 + naive_block_assignment,

250 + )

251 + if expert_id == -1:

252 + return

253 +

254 + # Get token offsets

255 + offs_token = _get_token_offs(

256 + sorted_token_ids_ptr,

257 + lora_id,

258 + pid_m,

259 + offs,

260 + stride_tl,

261 + max_loras,

262 + num_valid_tokens,

263 + naive_block_assignment,

264 + BLOCK_SIZE_M,

265 + )

266 + # get a_ptr,b_ptr,c_ptr

267 + cur_a_ptr = a_ptr + (slice_id % num_slice_a) * slice_a_size

268 + cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(c_ptr.dtype.element_ty))

269 + cur_c_ptr = c_ptr + (slice_id % num_slice_c) * slice_c_size

270 +

271 + # remove modulo wrap-around

272 + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int32)

273 + offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)

274 + token_mask = offs_token < num_valid_tokens

275 +

276 + # get a_ptrs,b_ptrs

277 + a_ptrs = cur_a_ptr + (

278 + offs_token[:, None] // token_mapping_factor * stride_am

279 + + offs_k[None, :] * stride_ak

280 + )

281 +

282 + b_ptrs = (

283 + cur_b_ptr

284 + + lora_id * stride_bl

285 + + expert_id * stride_be

286 + + offs_k[:, None] * stride_bk

287 + + offs_bn[None, :] * stride_bn

288 + )

289 +

290 + if USE_GDC and IS_PRIMARY:

291 + # GDC launch dependents hints the runtime system to launch dependent kernels.

292 + tl.extra.cuda.gdc_launch_dependents()

293 +

294 + # accumulator

295 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

296 +

297 + if USE_GDC and not IS_PRIMARY:

298 + tl.extra.cuda.gdc_wait()

299 +

300 + for k in range(0, grid_k):

301 + k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)

302 + # GDC wait waits for ALL programs in the prior kernel to complete

303 + # before continuing.

304 + # pre-fetch lora weight

305 + # add (offs_bn < N) mask; optional .ca for B

306 + b_mask = (offs_k[:, None] < k_remaining) & (offs_bn[None, :] < N)

307 + if USE_B_L2_CACHE:

308 + b = tl.load(b_ptrs, mask=b_mask, other=0.0, cache_modifier=".ca")

309 + else:

310 + b = tl.load(b_ptrs, mask=b_mask, other=0.0)

311 +

312 + if USE_GDC and not IS_PRIMARY:

313 + tl.extra.cuda.gdc_wait()

314 + a = tl.load(

315 + a_ptrs,

316 + mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),

317 + other=0.0,

318 + )

319 + accumulator += tl.dot(a, b)

320 + # Advance the ptrs to the next K block.

321 + a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak

322 + b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk

323 +

324 + if MUL_ROUTED_WEIGHT:

325 + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)

326 + accumulator = accumulator * moe_weight[:, None]

327 + accumulator = accumulator.to(c_ptr.dtype.element_ty)

328 + # Write back the block of the output

329 + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

330 + c_ptrs = cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]

331 + c_mask = token_mask[:, None] & (offs_cn[None, :] < N)

332 +

333 + if SPLIT_K == 1:

334 + if ADD_INPUTS:

335 + prev = tl.load(c_ptrs, mask=c_mask, other=0.0)

336 + tl.store(c_ptrs, prev + accumulator, mask=c_mask)

337 + else:

338 + tl.store(c_ptrs, accumulator, mask=c_mask)

339 + else:

340 + tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")

341 +

342 +

343 +@torch.inference_mode()

344 +def _fused_moe_lora_shrink_fp8(

345 + a_intermediate_cache1: torch.Tensor,

346 + # (num_slices, num_tokens, top_k_num, max_lora_rank)

347 + qcurr_hidden_states: torch.Tensor, # (num_tokens, K,)

348 + lora_a_stacked: list[

349 + torch.Tensor

350 + ], # [(max_loras, num_experts, max_lora_rank, K,),...]

351 + topk_weights: torch.Tensor, # (num_tokens, top_k_num)

352 + sorted_token_ids: torch.Tensor | None, # (max_loras, _)

353 + expert_ids: torch.Tensor, # (max_loras, _ ,) or (num_tokens * top_k,)

354 + num_tokens_post_padded: torch.Tensor | None, # (max_loras, )

355 + token_lora_mapping: torch.Tensor,

356 + top_k_num: int,

357 + lora_ids: torch.Tensor,

358 + adapter_enabled: torch.Tensor,

359 + ## adding for kernel

360 + device: torch.device,

361 + N: int,

362 + M: int,

363 + EM: int,

364 + K: int,

365 + num_tokens: int,

366 + num_experts: int,

367 + num_slices: int,

368 + block_size_m: int,

369 + block_size_n: int,

370 + block_size_k: int,

371 + group_size_m: int,

372 + num_warps: int,

373 + num_stages: int,

374 + split_k: int,

375 + num_active_loras: int,

376 + lora_a_scale_stacked: list[torch.Tensor],

377 + mul_routed_weight: bool = False,

378 + use_gdc: bool = False,

379 + act_scale: torch.Tensor | None = None,

380 + use_fp8_w8a8: bool = False,

381 + use_int8_w8a8: bool = False,

382 + use_int8_w8a16: bool = False,

383 + per_channel_quant: bool = False,

384 + block_shape: List[int] | None = None, # noqa: UP006, UP007

385 +) -> None:

386 + if use_fp8_w8a8 or use_int8_w8a8:

387 + assert lora_a_scale_stacked is not None, (

388 + "lora_a_scale_stacked must be provided for w8a8 quantization"

389 + )

390 + assert block_shape is None or triton.cdiv(

391 + lora_a_stacked[0].size(-2), block_shape[0]

392 + ) == lora_a_scale_stacked[0].size(-2), (

393 + "Incompatible block shape for lora_a_scale_stacked.size(-2) "

394 + )

395 + assert block_shape is None or triton.cdiv(

396 + lora_a_stacked[0].size(-1), block_shape[1]

397 + ) == lora_a_scale_stacked[0].size(-1), (

398 + "Incompatible block shape for lora_a_scale_stacked.size(-1) "

399 + )

400 + elif use_int8_w8a16:

401 + assert lora_a_scale_stacked is not None, (

402 + "lora_a_scale_stacked must be provided for w8a16 quantization"

403 + )

404 + assert block_shape is None or block_shape[0] == 0, (

405 + "Block shape for activation must be 0 for w8a16"

406 + )

407 + else:

408 + assert act_scale is None

409 + assert lora_a_scale_stacked is None

410 +

411 + if block_shape is not None:

412 + block_size_k = min(block_size_k, min(block_shape[0], block_shape[1]))

413 +

414 + if lora_a_scale_stacked is not None:

415 + b_scale_ptr = _get_ptr(lora_a_scale_stacked, device)

416 + w1_lora_a_scale_stacked = lora_a_scale_stacked[0]

417 +

418 + w1_lora_a_stacked = lora_a_stacked[0]

419 + shrink_config = {

420 + "BLOCK_SIZE_M": block_size_m,

421 + "BLOCK_SIZE_N": block_size_n,

422 + "BLOCK_SIZE_K": block_size_k,

423 + "GROUP_SIZE_M": group_size_m,

424 + "num_warps": num_warps,

425 + "num_stages": num_stages,

426 + "SPLIT_K": split_k,

427 + "USE_GDC": use_gdc,

428 + "launch_pdl": use_gdc, # triton kernel metadata

429 + }

430 +

431 + b_ptr = _get_ptr(lora_a_stacked, device)

432 +

433 + grid_lora_dim, stride_tl, stride_el = _adjust_kernel_inputs(

434 + num_active_loras, sorted_token_ids, expert_ids

435 + )

436 +

437 + grid = lambda META: (

438 + split_k

439 + * triton.cdiv(EM, META["BLOCK_SIZE_M"])

440 + * triton.cdiv(N, META["BLOCK_SIZE_N"]),

441 + len(lora_a_stacked),

442 + grid_lora_dim,

443 + )

444 + _fused_moe_lora_kernel_fp8[grid](

445 + qcurr_hidden_states,

446 + b_ptr,

447 + a_intermediate_cache1,

448 + act_scale,

449 + b_scale_ptr if lora_a_scale_stacked is not None else None,

450 + topk_weights,

451 + sorted_token_ids,

452 + expert_ids,

453 + num_tokens_post_padded,

454 + token_lora_mapping,

455 + N,

456 + K,

457 + EM,

458 + num_tokens,

459 + num_experts,

460 + top_k_num,

461 + lora_ids,

462 + adapter_enabled,

463 + lora_a_stacked[0].shape[0],

464 + qcurr_hidden_states.stride(0),

465 + qcurr_hidden_states.stride(1),

466 + w1_lora_a_stacked.stride(0),

467 + w1_lora_a_stacked.stride(1),

468 + w1_lora_a_stacked.stride(3),

469 + w1_lora_a_stacked.stride(2),

470 + a_intermediate_cache1.stride(2),

471 + a_intermediate_cache1.stride(3),

472 + stride_tl,

473 + stride_el,

474 + act_scale.stride(0) if act_scale is not None and act_scale.ndim == 2 else 0,

475 + act_scale.stride(1) if act_scale is not None and act_scale.ndim == 2 else 0,

476 + w1_lora_a_scale_stacked.stride(0)

477 + if lora_a_scale_stacked is not None and w1_lora_a_scale_stacked.ndim >= 2

478 + else 0,

479 + w1_lora_a_scale_stacked.stride(1)

480 + if lora_a_scale_stacked is not None and w1_lora_a_scale_stacked.ndim >= 2

481 + else 0,

482 + w1_lora_a_scale_stacked.stride(3)

483 + if lora_a_scale_stacked is not None and w1_lora_a_scale_stacked.ndim == 4

484 + else 0,

485 + w1_lora_a_scale_stacked.stride(2)

486 + if lora_a_scale_stacked is not None and w1_lora_a_scale_stacked.ndim == 4

487 + else 0,

488 + 0 if block_shape is None else block_shape[0],

489 + 0 if block_shape is None else block_shape[1],

490 + slice_a_size=qcurr_hidden_states.numel(),

491 + slice_c_size=a_intermediate_cache1.numel() // num_slices,

492 + num_slice_a=1,

493 + num_slice_c=num_slices,

494 + token_mapping_factor=1 if mul_routed_weight else top_k_num,

495 + naive_block_assignment=sorted_token_ids is None,

496 + MUL_ROUTED_WEIGHT=False,

497 + ADD_INPUTS=False,

498 + USE_B_L2_CACHE=True, # new

499 + IS_PRIMARY=True,

500 + use_fp8_w8a8=use_fp8_w8a8,

501 + use_int8_w8a8=use_int8_w8a8,

502 + use_int8_w8a16=use_int8_w8a16,

503 + per_channel_quant=per_channel_quant,

504 + **shrink_config,

505 + )

506 +

507 +

508 +@torch.inference_mode()

509 +def _fused_moe_lora_expand_fp8(

510 + output: torch.Tensor, # (num_tokens, top_k_num, N*len(lora_a_stacked),)

511 + a_intermediate_cache1: torch.Tensor, # (num_slices, M, top_k_num, max_lora_rank)

512 + lora_b_stacked: list[

513 + torch.Tensor

514 + ], # [(max_loras, num_experts, max_lora_rank, K,),...]

515 + topk_weights: torch.Tensor, # (num_tokens, top_k_num)

516 + sorted_token_ids: torch.Tensor | None, # (max_loras, _)

517 + expert_ids: torch.Tensor, # (max_loras, _ ,) or (num_tokens * top_k,)

518 + num_tokens_post_padded: torch.Tensor | None, # (max_loras, )

519 + token_lora_mapping: torch.Tensor,

520 + top_k_num: int,

521 + lora_ids: torch.Tensor,

522 + adapter_enabled: torch.Tensor,

523 + ## adding for kernel

524 + device: torch.device,

525 + N: int,

526 + M: int,

527 + EM: int,

528 + K: int,

529 + num_tokens: int,

530 + num_experts: int,

531 + num_slices: int,

532 + max_lora_rank: int,

533 + w1_output_dim_size: int,

534 + block_size_m: int,

535 + block_size_n: int,

536 + block_size_k: int,

537 + group_size_m: int,

538 + num_warps: int,

539 + num_stages: int,

540 + split_k: int,

541 + num_active_loras: int,

542 + lora_b_scale_stacked: list[torch.Tensor],

543 + mul_routed_weight: bool = False,

544 + offset: int = 0,

545 + use_gdc: bool = False,

546 + act_scale: torch.Tensor | None = None,

547 + use_fp8_w8a8: bool = False,

548 + use_int8_w8a8: bool = False,

549 + use_int8_w8a16: bool = False,

550 + per_channel_quant: bool = False,

551 + block_shape: List[int] | None = None, # noqa: UP006, UP007

552 +) -> None:

553 + if use_fp8_w8a8 or use_int8_w8a8:

554 + assert lora_b_scale_stacked is not None, (

555 + "lora_b_scale_stacked must be provided for w8a8 quantization"

556 + )

557 + assert block_shape is None or triton.cdiv(

558 + lora_b_stacked[0].size(-2), block_shape[0]

559 + ) == lora_b_scale_stacked[0].size(-2), (

560 + "Incompatible block shape for lora_b_scale_stacked.size(-2) "

561 + )

562 + assert block_shape is None or triton.cdiv(

563 + lora_b_stacked[0].size(-1), block_shape[1]

564 + ) == lora_b_scale_stacked[0].size(-1), (

565 + "Incompatible block shape for lora_b_scale_stacked.size(-1) "

566 + )

567 + elif use_int8_w8a16:

568 + assert lora_b_scale_stacked is not None, (

569 + "lora_b_scale_stacked must be provided for w8a16 quantization"

570 + )

571 + assert block_shape is None or block_shape[0] == 0, (

572 + "Block shape for activation must be 0 for w8a16"

573 + )

574 + else:

575 + assert act_scale is None

576 + assert lora_b_scale_stacked is None

577 +

578 + if lora_b_scale_stacked is not None:

579 + b_scale_ptr = _get_ptr(lora_b_scale_stacked, device)

580 + w1_lora_b_scale_stacked = lora_b_scale_stacked[0]

581 +

582 + if block_shape is not None:

583 + block_size_k = min(block_size_k, min(block_shape[0], block_shape[1]))

584 +

585 + b_ptr = _get_ptr(lora_b_stacked, device)

586 + K = max_lora_rank

587 + N = w1_output_dim_size

588 +

589 + w1_lora_b_stacked = lora_b_stacked[0]

590 +

591 + a_intermediate_cache1 = a_intermediate_cache1.view(

592 + -1, a_intermediate_cache1.shape[3]

593 + )

594 +

595 + expand_config = {

596 + "BLOCK_SIZE_M": block_size_m,

597 + "BLOCK_SIZE_N": block_size_n,

598 + "BLOCK_SIZE_K": block_size_k,

599 + "GROUP_SIZE_M": group_size_m,

600 + "num_warps": num_warps,

601 + "num_stages": num_stages,

602 + "SPLIT_K": 1, # Set split_k = 1 for expand calls

603 + "USE_GDC": use_gdc,

604 + "launch_pdl": use_gdc, # triton kernel metadata

605 + }

606 +

607 + grid_lora_dim, stride_tl, stride_el = _adjust_kernel_inputs(

608 + num_active_loras, sorted_token_ids, expert_ids

609 + )

610 +

611 + grid = lambda META: (

612 + triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),

613 + len(lora_b_stacked),

614 + grid_lora_dim,

615 + )

616 +

617 + # Fast path: directly accumulate into the corresponding slice interval of output.

618 + out_view = output[:, :, offset : offset + num_slices * N]

619 + slice_c_size = N * out_view.stride(2)

620 +

621 + _fused_moe_lora_kernel_fp8[grid](

622 + a_intermediate_cache1,

623 + b_ptr,

624 + out_view,

625 + act_scale,

626 + b_scale_ptr if lora_b_scale_stacked is not None else None,

627 + topk_weights,

628 + sorted_token_ids,

629 + expert_ids,

630 + num_tokens_post_padded,

631 + token_lora_mapping,

632 + N,

633 + K,

634 + EM,

635 + num_tokens,

636 + num_experts,

637 + top_k_num,

638 + lora_ids,

639 + adapter_enabled,

640 + lora_b_stacked[0].shape[0],

641 + a_intermediate_cache1.stride(0),

642 + a_intermediate_cache1.stride(1),

643 + w1_lora_b_stacked.stride(0),

644 + w1_lora_b_stacked.stride(1),

645 + w1_lora_b_stacked.stride(3),

646 + w1_lora_b_stacked.stride(2),

647 + out_view.stride(1),

648 + out_view.stride(2),

649 + stride_tl,

650 + stride_el,

651 + act_scale.stride(0) if act_scale is not None and act_scale.ndim == 2 else 0,

652 + act_scale.stride(1) if act_scale is not None and act_scale.ndim == 2 else 0,

653 + w1_lora_b_scale_stacked.stride(0)

654 + if lora_b_scale_stacked is not None and w1_lora_b_scale_stacked.ndim >= 2

655 + else 0,

656 + w1_lora_b_scale_stacked.stride(1)

657 + if lora_b_scale_stacked is not None and w1_lora_b_scale_stacked.ndim >= 2

658 + else 0,

659 + w1_lora_b_scale_stacked.stride(3)

660 + if lora_b_scale_stacked is not None and w1_lora_b_scale_stacked.ndim == 4

661 + else 0,

662 + w1_lora_b_scale_stacked.stride(2)

663 + if lora_b_scale_stacked is not None and w1_lora_b_scale_stacked.ndim == 4

664 + else 0,

665 + 0 if block_shape is None else block_shape[0],

666 + 0 if block_shape is None else block_shape[1],

667 + slice_a_size=a_intermediate_cache1.numel() // num_slices,

668 + slice_c_size=slice_c_size,

669 + num_slice_a=num_slices,

670 + num_slice_c=num_slices,

671 + token_mapping_factor=1,

672 + naive_block_assignment=sorted_token_ids is None,

673 + MUL_ROUTED_WEIGHT=mul_routed_weight,

674 + ADD_INPUTS=True,

675 + USE_B_L2_CACHE=True, # new

676 + IS_PRIMARY=False,

677 + use_fp8_w8a8=use_fp8_w8a8,

678 + use_int8_w8a8=use_int8_w8a8,

679 + use_int8_w8a16=use_int8_w8a16,

680 + per_channel_quant=per_channel_quant,

681 + **expand_config,

682 + )

683 +

684 +

685 +@torch.inference_mode()

686 +def _fused_moe_lora_fp8(

687 + output: torch.Tensor, # (num_tokens, top_k_num, N*len(lora_a_stacked),)

688 + qcurr_hidden_states: torch.Tensor, # (num_tokens, K,)

689 + lora_a_stacked: list[

690 + torch.Tensor

691 + ], # [(max_loras, num_experts, max_lora_rank, K,),...]

692 + lora_b_stacked: list[

693 + torch.Tensor

694 + ], # [(max_loras, num_experts, N, max_lora_rank,),...]

695 + topk_weights: torch.Tensor, # (num_tokens, top_k_num)

696 + sorted_token_ids: torch.Tensor | None, # (max_loras, _)

697 + expert_ids: torch.Tensor, # (max_loras, _ ,) or (num_tokens * top_k,)

698 + num_tokens_post_padded: torch.Tensor | None, # (max_loras, )

699 + token_lora_mapping: torch.Tensor,

700 + max_lora_rank: int,

701 + top_k_num: int,

702 + lora_ids: torch.Tensor,

703 + num_active_loras: int,

704 + adapter_enabled: torch.Tensor,

705 + shrink_block_size_m: int,

706 + shrink_block_size_n: int,

707 + shrink_block_size_k: int,

708 + shrink_group_size_m: int,

709 + shrink_num_warps: int,

710 + shrink_num_stages: int,

711 + shrink_split_k: int,

712 + expand_block_size_m: int,

713 + expand_block_size_n: int,

714 + expand_block_size_k: int,

715 + expand_group_size_m: int,

716 + expand_num_warps: int,

717 + expand_num_stages: int,

718 + expand_split_k: int,

719 + lora_a_scale_stacked: list[torch.Tensor],

720 + lora_b_scale_stacked: list[torch.Tensor],

721 + shrink_act_scale: torch.Tensor | None = None,

722 + expand_act_scale: torch.Tensor | None = None,

723 + mul_routed_weight: bool = False,

724 + fully_sharded: bool = False,

725 + offset: int = 0,

726 + use_fp8_w8a8: bool = False,

727 + use_int8_w8a8: bool = False,

728 + use_int8_w8a16: bool = False,

729 + per_channel_quant: bool = False,

730 + block_shape: List[int] | None = None, # noqa: UP006, UP007

731 +) -> None:

732 + assert len(lora_a_stacked) == len(lora_b_stacked) > 0

733 + assert topk_weights.dim() == qcurr_hidden_states.dim() == 2

734 + if sorted_token_ids is None:

735 + assert expert_ids.dim() == 1

736 + else:

737 + assert sorted_token_ids is not None

738 + assert num_tokens_post_padded is not None

739 + assert (

740 + sorted_token_ids.dim()

741 + == expert_ids.dim()

742 + == topk_weights.dim()

743 + == qcurr_hidden_states.dim()

744 + == 2

745 + )

746 + assert (

747 + sorted_token_ids.shape[0]

748 + == expert_ids.shape[0]

749 + == num_tokens_post_padded.shape[0]

750 + )

751 + assert output.shape[0] == topk_weights.shape[0]

752 + assert top_k_num == topk_weights.shape[1]

753 + device = qcurr_hidden_states.device

754 + num_slices = len(lora_a_stacked)

755 + w1_lora_b_stacked = lora_b_stacked[0]

756 + num_experts = lora_a_stacked[0].shape[1]

757 + N = max_lora_rank

758 + M = topk_weights.shape[0]

759 + K = qcurr_hidden_states.shape[1]

760 + num_tokens = M * top_k_num

761 + w1_output_dim_size = w1_lora_b_stacked.shape[2]

762 + assert shrink_block_size_m == expand_block_size_m

763 + EM = (

764 + sorted_token_ids.shape[1]

765 + if sorted_token_ids is not None

766 + else num_tokens * shrink_block_size_m

767 + )

768 +

769 + a_intermediate_cache1 = torch.zeros(

770 + (num_slices, M, top_k_num, max_lora_rank),

771 + dtype=output.dtype,

772 + device=device,

773 + )

774 +

775 + use_gdc = supports_pdl(device) and not fully_sharded

776 + _fused_moe_lora_shrink_fp8(

777 + a_intermediate_cache1,

778 + qcurr_hidden_states,

779 + lora_a_stacked,

780 + topk_weights,

781 + sorted_token_ids,

782 + expert_ids,

783 + num_tokens_post_padded,

784 + token_lora_mapping,

785 + top_k_num,

786 + lora_ids,

787 + adapter_enabled,

788 + ## adding for kernel

789 + device,

790 + N,

791 + M,

792 + EM,

793 + K,

794 + num_tokens,

795 + num_experts,

796 + num_slices,

797 + shrink_block_size_m,

798 + shrink_block_size_n,

799 + shrink_block_size_k,

800 + shrink_group_size_m,

801 + shrink_num_warps,

802 + shrink_num_stages,

803 + shrink_split_k,

804 + num_active_loras,

805 + lora_a_scale_stacked,

806 + mul_routed_weight=mul_routed_weight,

807 + use_gdc=use_gdc,

808 + act_scale=shrink_act_scale,

809 + use_fp8_w8a8=use_fp8_w8a8,

810 + use_int8_w8a8=use_int8_w8a8,

811 + use_int8_w8a16=use_int8_w8a16,

812 + per_channel_quant=per_channel_quant,

813 + block_shape=block_shape,

814 + )

815 +

816 + if fully_sharded:

817 + if max_lora_rank == w1_lora_b_stacked.shape[-1]:

818 + a_intermediate_cache1 = tensor_model_parallel_all_reduce(

819 + a_intermediate_cache1

820 + )

821 + else:

822 + a_intermediate_cache1 = tensor_model_parallel_all_gather(

823 + a_intermediate_cache1

824 + )

825 +

826 + # reset max_lora_rank to the full rank after allgather

827 + max_lora_rank = a_intermediate_cache1.shape[-1]

828 +

829 + _fused_moe_lora_expand_fp8(

830 + output,

831 + a_intermediate_cache1,

832 + lora_b_stacked,

833 + topk_weights,

834 + sorted_token_ids,

835 + expert_ids,

836 + num_tokens_post_padded,

837 + token_lora_mapping,

838 + top_k_num,

839 + lora_ids,

840 + adapter_enabled,

841 + ## adding for kernel

842 + device,

843 + N,

844 + M,

845 + EM,

846 + K,

847 + num_tokens,

848 + num_experts,

849 + num_slices,

850 + max_lora_rank,

851 + w1_output_dim_size,

852 + expand_block_size_m,

853 + expand_block_size_n,

854 + expand_block_size_k,

855 + expand_group_size_m,

856 + expand_num_warps,

857 + expand_num_stages,

858 + expand_split_k,

859 + num_active_loras,

860 + lora_b_scale_stacked,

861 + mul_routed_weight=mul_routed_weight,

862 + offset=offset,

863 + use_gdc=use_gdc,

864 + act_scale=expand_act_scale,

865 + use_fp8_w8a8=use_fp8_w8a8,

866 + use_int8_w8a8=use_int8_w8a8,

867 + use_int8_w8a16=use_int8_w8a16,

868 + per_channel_quant=per_channel_quant,

869 + block_shape=block_shape,

870 + )

871 +

872 +

873 +def _fused_moe_lora_fp8_fake(

874 + output: torch.Tensor,

875 + qcurr_hidden_states: torch.Tensor,

876 + lora_a_stacked: list[torch.Tensor],

877 + lora_b_stacked: list[torch.Tensor],

878 + topk_weights: torch.Tensor,

879 + sorted_token_ids: torch.Tensor | None,

880 + expert_ids: torch.Tensor,

881 + num_tokens_post_padded: torch.Tensor | None,

882 + token_lora_mapping: torch.Tensor,

883 + max_lora_rank: int,

884 + top_k_num: int,

885 + lora_ids: torch.Tensor,

886 + num_active_loras: int,

887 + adapter_enabled: torch.Tensor,

888 + shrink_block_size_m: int,

889 + shrink_block_size_n: int,

890 + shrink_block_size_k: int,

891 + shrink_group_size_m: int,

892 + shrink_num_warps: int,

893 + shrink_num_stages: int,

894 + shrink_split_k: int,

895 + expand_block_size_m: int,

896 + expand_block_size_n: int,

897 + expand_block_size_k: int,

898 + expand_group_size_m: int,

899 + expand_num_warps: int,

900 + expand_num_stages: int,

901 + expand_split_k: int,

902 + lora_a_scale_stacked: list[torch.Tensor],

903 + lora_b_scale_stacked: list[torch.Tensor],

904 + mul_routed_weight: bool = False,

905 + fully_sharded: bool = False,

906 + offset: int = 0,

907 + shrink_act_scale: torch.Tensor | None = None,

908 + expand_act_scale: torch.Tensor | None = None,

909 + use_fp8_w8a8: bool = False,

910 + use_int8_w8a8: bool = False,

911 + use_int8_w8a16: bool = False,

912 + per_channel_quant: bool = False,

913 + block_shape: List[int] | None = None, # noqa: UP006, UP007

914 +) -> None:

915 + return

916 +

917 +

918 +def _fused_moe_lora_shrink_fp8_fake(

919 + a_intermediate_cache1: torch.Tensor,

920 + qcurr_hidden_states: torch.Tensor,

921 + lora_a_stacked: list[torch.Tensor],

922 + topk_weights: torch.Tensor,

923 + sorted_token_ids: torch.Tensor | None,

924 + expert_ids: torch.Tensor,

925 + num_tokens_post_padded: torch.Tensor | None,

926 + token_lora_mapping: torch.Tensor,

927 + top_k_num: int,

928 + lora_ids: torch.Tensor,

929 + adapter_enabled: torch.Tensor,

930 + device: torch.device,

931 + N: int,

932 + M: int,

933 + EM: int,

934 + K: int,

935 + num_tokens: int,

936 + num_experts: int,

937 + num_slices: int,

938 + block_size_m: int,

939 + block_size_n: int,

940 + block_size_k: int,

941 + group_size_m: int,

942 + num_warps: int,

943 + num_stages: int,

944 + split_k: int,

945 + num_active_loras: int,

946 + lora_a_scale_stacked: list[torch.Tensor],

947 + mul_routed_weight: bool = False,

948 + use_gdc: bool = False,

949 + act_scale: torch.Tensor | None = None,

950 + use_fp8_w8a8: bool = False,

951 + use_int8_w8a8: bool = False,

952 + use_int8_w8a16: bool = False,

953 + per_channel_quant: bool = False,

954 + block_shape: List[int] | None = None, # noqa: UP006, UP007

955 +) -> None:

956 + return

957 +

958 +

959 +def _fused_moe_lora_expand_fp8_fake(

960 + output: torch.Tensor,

961 + a_intermediate_cache1: torch.Tensor,

962 + lora_b_stacked: list[torch.Tensor],

963 + topk_weights: torch.Tensor,

964 + sorted_token_ids: torch.Tensor | None,

965 + expert_ids: torch.Tensor,

966 + num_tokens_post_padded: torch.Tensor | None,

967 + token_lora_mapping: torch.Tensor,

968 + top_k_num: int,

969 + lora_ids: torch.Tensor,

970 + adapter_enabled: torch.Tensor,

971 + device: torch.device,

972 + N: int,

973 + M: int,

974 + EM: int,

975 + K: int,

976 + num_tokens: int,

977 + num_experts: int,

978 + num_slices: int,

979 + max_lora_rank: int,

980 + w1_output_dim_size: int,

981 + block_size_m: int,

982 + block_size_n: int,

983 + block_size_k: int,

984 + group_size_m: int,

985 + num_warps: int,

986 + num_stages: int,

987 + split_k: int,

988 + num_active_loras: int,

989 + act_scale: torch.Tensor,

990 + lora_b_scale_stacked: list[torch.Tensor],

991 + mul_routed_weight: bool = False,

992 + offset: int = 0,

993 + use_fp8_w8a8: bool = False,

994 + use_int8_w8a8: bool = False,

995 + use_int8_w8a16: bool = False,

996 + per_channel_quant: bool = False,

997 + block_shape: List[int] | None = None, # noqa: UP006, UP007

998 + use_gdc: bool = False,

999 +) -> None:

1000 + return

1001 +

1002 +

1003 +try:

1004 + direct_register_custom_op(

1005 + op_name="fused_moe_lora_fp8",

1006 + op_func=_fused_moe_lora_fp8,

1007 + mutates_args=["output"],

1008 + fake_impl=_fused_moe_lora_fp8_fake,

1009 + )

1010 +

1011 + direct_register_custom_op(

1012 + op_name="fused_moe_lora_shrink_fp8",

1013 + op_func=_fused_moe_lora_shrink_fp8,

1014 + mutates_args=["a_intermediate_cache1"],

1015 + fake_impl=_fused_moe_lora_shrink_fp8_fake,

1016 + )

1017 +

1018 + direct_register_custom_op(

1019 + op_name="fused_moe_lora_expand_fp8",

1020 + op_func=_fused_moe_lora_expand_fp8,

1021 + mutates_args=["output"],

1022 + fake_impl=_fused_moe_lora_expand_fp8_fake,

1023 + )

1024 +

1025 + fused_moe_lora_fp8 = torch.ops.vllm.fused_moe_lora_fp8

1026 + fused_moe_lora_shrink_fp8 = torch.ops.vllm.fused_moe_lora_shrink_fp8

1027 + fused_moe_lora_expand_fp8 = torch.ops.vllm.fused_moe_lora_expand_fp8

1028 +

1029 +except AttributeError:

1030 + fused_moe_lora_fp8 = _fused_moe_lora_fp8

1031 + fused_moe_lora_shrink_fp8 = _fused_moe_lora_shrink_fp8

1032 + fused_moe_lora_expand_fp8 = _fused_moe_lora_expand_fp8

[LoRA] Support Quantized Adapters (#30286)