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

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 12, 23:30 UTC

add support for --fully-sharded-loras in fused_moe (#28761)

Signed-off-by: gnovack <gnovack@amazon.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

gnovack committed Nov 19, 2025 at 00:32 UTC d69062c67af46a2e624be92162e9db585eef329b

6 files changed +274 -10

tests/lora/test_fused_moe_lora_kernel.py

+206 -2

@@ -1,13 +1,25 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +import os

4 import random

6 import pytest

7 import torch

9 +from tests.utils import multi_gpu_test

10 from vllm import _custom_ops as ops

11 +from vllm.distributed import (

12 + init_distributed_environment,

13 + initialize_model_parallel,

14 + tensor_model_parallel_all_gather,

15 + tensor_model_parallel_all_reduce,

16 +)

17 +from vllm.distributed.parallel_state import (

18 + get_tensor_model_parallel_world_size,

19 +)

20 from vllm.lora.ops.triton_ops import fused_moe_lora

21 from vllm.platforms import current_platform

22 +from vllm.utils.network_utils import get_open_port

25 @pytest.fixture(autouse=True)

@@ -122,6 +134,8 @@ def use_fused_moe_lora_kernel(

134 max_loras,

135 num_experts,

136 block_size,

137 + fully_sharded=False,

138 + offset=0,

139 ):

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

141 max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)

@@ -195,10 +209,10 @@ def use_fused_moe_lora_kernel(

209 config["NUM_STAGES"],

210 config["SPLIT_K"],

211 mul_routed_weight,

212 + fully_sharded=fully_sharded,

213 + offset=offset,

214 )

215

200 - return output

201 -

216

217 def use_torch(

218 hidden_states,

@@ -317,3 +331,193 @@ def test_fused_moe_lora_kernel(

331 )

332

333 torch.testing.assert_close(output, output2, atol=1e-1, rtol=1e-1)

334 +

335 +

336 +@multi_gpu_test(num_gpus=2)

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

338 +@pytest.mark.parametrize("top_k_num", [6])

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

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

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

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

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

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

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

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

347 +@pytest.mark.parametrize("column_parallel", [True, False])

348 +def test_fused_moe_lora_kernel_fully_sharded(

349 + num_tokens,

350 + top_k_num,

351 + num_experts,

352 + max_loras,

353 + N,

354 + K,

355 + max_lora_rank,

356 + block_size,

357 + dtype,

358 + seed,

359 + column_parallel,

360 +):

361 + current_platform.seed_everything(seed)

362 + # the number of randomly generated sentences.

363 + num_sequences = 10

364 + # generate data

365 + topk_ids, topk_weights, token_lora_mapping = sample_data(

366 + num_tokens, num_sequences, max_loras, num_experts, top_k_num

367 + )

368 +

369 + def run_torch_spawn(fn, nprocs):

370 + torch.multiprocessing.spawn(

371 + fn,

372 + args=(

373 + nprocs,

374 + f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",

375 + dtype,

376 + seed,

377 + N,

378 + K,

379 + num_tokens,

380 + topk_ids,

381 + topk_weights,

382 + token_lora_mapping,

383 + max_lora_rank,

384 + top_k_num,

385 + max_loras,

386 + num_experts,

387 + block_size,

388 + column_parallel,

389 + ),

390 + nprocs=nprocs,

391 + )

392 +

393 + run_torch_spawn(use_fused_moe_lora_kernel_tensor_parallel, nprocs=2)

394 +

395 +

396 +def use_fused_moe_lora_kernel_tensor_parallel(

397 + local_rank,

398 + world_size,

399 + init_method,

400 + dtype,

401 + seed,

402 + N,

403 + K,

404 + num_tokens,

405 + topk_ids,

406 + topk_weights,

407 + token_lora_mapping,

408 + max_lora_rank,

409 + top_k_num,

410 + max_loras,

411 + num_experts,

412 + block_size,

413 + column_parallel,

414 +):

415 + def _get_shard_slice(shard_size):

416 + return slice(local_rank * shard_size, (local_rank + 1) * shard_size)

417 +

418 + current_platform.seed_everything(seed)

419 +

420 + device = torch.device(f"cuda:{local_rank}")

421 + torch.cuda.set_device(device)

422 + torch.set_default_device(device)

423 + torch.set_default_dtype(dtype)

424 +

425 + init_distributed_environment(

426 + world_size=world_size,

427 + rank=local_rank,

428 + local_rank=local_rank,

429 + distributed_init_method=init_method,

430 + )

431 + initialize_model_parallel(world_size, 1)

432 + tp_size = get_tensor_model_parallel_world_size()

433 +

434 + input_dim = K if column_parallel else N

435 + output_dim = N if column_parallel else K

436 +

437 + # init lora weights

438 + lora_a = torch.rand(

439 + (

440 + max_loras,

441 + num_experts,

442 + max_lora_rank,

443 + input_dim,

444 + ),

445 + dtype=dtype,

446 + )

447 + lora_b = torch.rand(

448 + (

449 + max_loras,

450 + num_experts,

451 + output_dim,

452 + max_lora_rank,

453 + ),

454 + dtype=dtype,

455 + )

456 +

457 + hidden_states = torch.rand(

458 + (

459 + num_tokens,

460 + input_dim,

461 + ),

462 + dtype=dtype,

463 + )

464 +

465 + output = torch.zeros((num_tokens, top_k_num, output_dim), dtype=dtype)

466 + topk_ids = topk_ids.to(device)

467 + topk_weights = topk_weights.to(device)

468 + token_lora_mapping = token_lora_mapping.to(device)

469 +

470 + ref_output = use_torch(

471 + hidden_states,

472 + token_lora_mapping,

473 + topk_ids,

474 + [lora_a],

475 + [lora_b],

476 + top_k_num,

477 + )

478 +

479 + if column_parallel:

480 + # Column parallel (e.g. gate_up_proj): LoRA A is sliced along the rank dim,

481 + # and Lora B is sliced along the output dim

482 + lora_a_shard_size = max_lora_rank // tp_size

483 + lora_a = lora_a[:, :, _get_shard_slice(lora_a_shard_size), :]

484 + max_lora_rank = lora_a_shard_size

485 + offset = 0

486 +

487 + lora_b_shard_size = output_dim // tp_size

488 + lora_b = lora_b[:, :, _get_shard_slice(lora_b_shard_size), :]

489 + output = output[:, :, _get_shard_slice(lora_b_shard_size)].contiguous()

490 + else:

491 + # Row parallel (e.g. down proj): LoRA A is sliced along the input dim,

492 + # and LoRA B is sliced along the output dim

493 + lora_a_shard_size = input_dim // tp_size

494 + lora_a = lora_a[:, :, :, _get_shard_slice(lora_a_shard_size)]

495 + hidden_states = hidden_states[:, _get_shard_slice(lora_a_shard_size)]

496 +

497 + lora_b_shard_size = output_dim // tp_size

498 + lora_b = lora_b[:, :, _get_shard_slice(lora_b_shard_size), :]

499 + offset = lora_b_shard_size * local_rank

500 +

501 + use_fused_moe_lora_kernel(

502 + topk_ids,

503 + topk_weights,

504 + token_lora_mapping,

505 + max_lora_rank,

506 + top_k_num,

507 + [lora_a],

508 + [lora_b],

509 + hidden_states,

510 + output,

511 + max_loras,

512 + num_experts,

513 + block_size,

514 + fully_sharded=True,

515 + offset=offset,

516 + )

517 +

518 + if column_parallel:

519 + output = tensor_model_parallel_all_gather(output)

520 + else:

521 + output = tensor_model_parallel_all_reduce(output)

522 +

523 + torch.testing.assert_close(output, ref_output, atol=1e-1, rtol=1e-1)

tests/lora/test_olmoe_tp.py

+8 -2

@@ -2,6 +2,8 @@

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

5 +import pytest

6 +

7 import vllm

8 from vllm.lora.request import LoRARequest

@@ -111,8 +113,9 @@ def test_olmoe_lora_mixed(olmoe_lora_files):

113 generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])

114

115

116 +@pytest.mark.parametrize("fully_sharded_loras", [False, True])

117 @multi_gpu_test(num_gpus=2)

115 -def test_olmoe_lora_tp2(olmoe_lora_files):

118 +def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):

119 llm = vllm.LLM(

120 MODEL_PATH,

121 max_model_len=1024,

@@ -122,14 +125,16 @@ def test_olmoe_lora_tp2(olmoe_lora_files):

125 trust_remote_code=True,

126 enable_chunked_prefill=True,

127 tensor_parallel_size=2,

128 + fully_sharded_loras=fully_sharded_loras,

129 )

130

131 generate_and_test(llm, olmoe_lora_files, lora_id=1)

132 generate_and_test(llm, olmoe_lora_files, lora_id=2)

133

134

135 +@pytest.mark.parametrize("fully_sharded_loras", [False, True])

136 @multi_gpu_test(num_gpus=4)

132 -def test_olmoe_lora_tp4(olmoe_lora_files):

137 +def test_olmoe_lora_tp4(olmoe_lora_files, fully_sharded_loras):

138 llm = vllm.LLM(

139 MODEL_PATH,

140 max_model_len=1024,

@@ -139,6 +144,7 @@ def test_olmoe_lora_tp4(olmoe_lora_files):

144 trust_remote_code=True,

145 enable_chunked_prefill=True,

146 tensor_parallel_size=4,

147 + fully_sharded_loras=fully_sharded_loras,

148 )

149

150 generate_and_test(llm, olmoe_lora_files, lora_id=1)

vllm/lora/layers/fused_moe.py

+32 -4

@@ -12,6 +12,7 @@ from vllm.distributed.parallel_state import (

12 get_tensor_model_parallel_rank,

13 get_tensor_model_parallel_world_size,

14 )

15 +from vllm.distributed.utils import divide

16 from vllm.lora.layers.base import BaseLayerWithLoRA

17 from vllm.lora.ops.triton_ops.utils import get_lora_op_configs

18 from vllm.model_executor.layers.fused_moe import FusedMoE

@@ -205,6 +206,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

206 shrink_config, ## pass the shrink config

207 expand_config, ## pass the expand config

208 self.adapter_enabled,

209 + fully_sharded=self.fully_sharded,

210 )

211

212 result = func(*args, **kwargs)

@@ -250,7 +252,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

252 sorted_token_ids_lora = sorted_token_ids_lora.view(max_loras, -1)

253 intermediate_cache2 = moe_state_dict["intermediate_cache2"]

254 intermediate_cache3 = args[0]

253 - max_lora_rank = self.w1_lora_a_stacked.shape[-2]

255 + max_lora_rank = self.w2_lora_a_stacked.shape[-2]

256 +

257 + shard_size_w2 = divide(self.base_layer.hidden_size, self.tp_size)

258 +

259 self.punica_wrapper.add_lora_fused_moe(

260 intermediate_cache3,

261 intermediate_cache2,

@@ -266,6 +271,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

271 expand_config, ## pass the expand config

272 self.adapter_enabled,

273 True,

274 + fully_sharded=self.fully_sharded,

275 + offset=shard_size_w2 * self.tp_rank if self.fully_sharded else 0,

276 )

277

278 result = func(*args, **kwargs)

@@ -294,6 +301,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

301 model_config: PretrainedConfig | None = None,

302 ) -> None:

303 """Initializes lora matrices."""

304 + self.fully_sharded = lora_config.fully_sharded_loras

305

306 self.adapter_enabled = torch.tensor(

307 [0] * (max_loras + 1), dtype=torch.int, device=self.device

@@ -303,7 +311,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

311 (

312 max_loras,

313 self.base_layer.local_num_experts,

306 - lora_config.max_lora_rank,

314 + lora_config.max_lora_rank

315 + if not self.fully_sharded

316 + else divide(lora_config.max_lora_rank, self.tp_size),

317 self.base_layer.hidden_size,

318 ),

319 dtype=lora_config.lora_dtype,

@@ -334,7 +344,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

344 (

345 max_loras,

346 self.base_layer.local_num_experts,

337 - self.base_layer.hidden_size,

347 + self.base_layer.hidden_size

348 + if not self.fully_sharded

349 + else divide(self.base_layer.hidden_size, self.tp_size),

350 lora_config.max_lora_rank,

351 ),

352 dtype=lora_config.lora_dtype,

@@ -345,7 +357,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

357 (

358 max_loras,

359 self.base_layer.local_num_experts,

348 - lora_config.max_lora_rank,

360 + lora_config.max_lora_rank

361 + if not self.fully_sharded

362 + else divide(lora_config.max_lora_rank, self.tp_size),

363 self.base_layer.hidden_size,

364 ),

365 dtype=lora_config.lora_dtype,

@@ -419,6 +433,20 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

433 w3_lora_b = w3_lora_b[start_idx:end_idx, :]

434 w2_lora_a = w2_lora_a[:, start_idx:end_idx]

435

436 + if self.fully_sharded:

437 + # Based on S-LoRA, we slice W1 and W3 A along the rank dim,

438 + # and W2 B along the hidden_size dim.

439 + w13_shard_size = self.w1_lora_a_stacked[index, eid].shape[0]

440 + w13_start_idx = self.tp_rank * w13_shard_size

441 + w13_end_idx = (self.tp_rank + 1) * w13_shard_size

442 + w1_lora_a = w1_lora_a[w13_start_idx:w13_end_idx, :]

443 + w3_lora_a = w3_lora_a[w13_start_idx:w13_end_idx, :]

444 +

445 + w2_shard_size = self.w2_lora_b_stacked[index, eid].shape[0]

446 + w2_start_idx = self.tp_rank * w2_shard_size

447 + w2_end_idx = (self.tp_rank + 1) * w2_shard_size

448 + w2_lora_b = w2_lora_b[w2_start_idx:w2_end_idx, :]

449 +

450 self.w1_lora_a_stacked[

451 index, eid, : w1_lora_a.shape[0], : w1_lora_a.shape[1]

452 ].copy_(w1_lora_a, non_blocking=True)

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

+22 -2

@@ -3,6 +3,10 @@

4 import torch

6 +from vllm.distributed import (

7 + tensor_model_parallel_all_gather,

8 + tensor_model_parallel_all_reduce,

9 +)

10 from vllm.triton_utils import tl, triton

11 from vllm.utils.torch_utils import direct_register_custom_op

@@ -311,6 +315,7 @@ def _fused_moe_lora_expand(

315 num_stages: int,

316 split_k: int,

317 mul_routed_weight: bool = False,

318 + offset: int = 0,

319 ) -> None:

320 b_ptr = _get_ptr(lora_b_stacked, device)

321 K = max_lora_rank

@@ -380,7 +385,7 @@ def _fused_moe_lora_expand(

385 **expand_config,

386 )

387 for i in range(num_slices):

383 - output[:, :, i * N : (i + 1) * N] += b_intermediate_cache1[i]

388 + output[:, :, i * N + offset : (i + 1) * N + offset] += b_intermediate_cache1[i]

389

390

391 @torch.inference_mode()

@@ -416,6 +421,8 @@ def _fused_moe_lora(

421 expand_num_stages: int,

422 expand_split_k: int,

423 mul_routed_weight: bool = False,

424 + fully_sharded: bool = False,

425 + offset: int = 0,

426 ) -> None:

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

428 assert (

@@ -430,7 +437,6 @@ def _fused_moe_lora(

437 == expert_ids.shape[0]

438 == num_tokens_post_padded.shape[0]

439 )

433 - assert len(lora_b_stacked) * lora_b_stacked[0].shape[-2] == output.shape[-1]

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

441 assert top_k_num == topk_weights.shape[1]

442 device = qcurr_hidden_states.device

@@ -480,6 +486,19 @@ def _fused_moe_lora(

486 mul_routed_weight,

487 )

488

489 + if fully_sharded:

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

491 + a_intermediate_cache1 = tensor_model_parallel_all_reduce(

492 + a_intermediate_cache1

493 + )

494 + else:

495 + a_intermediate_cache1 = tensor_model_parallel_all_gather(

496 + a_intermediate_cache1

497 + )

498 +

499 + # reset max_lora_rank to the full rank after allgather

500 + max_lora_rank = a_intermediate_cache1.shape[-1]

501 +

502 _fused_moe_lora_expand(

503 output,

504 a_intermediate_cache1,

@@ -510,6 +529,7 @@ def _fused_moe_lora(

529 expand_num_stages,

530 expand_split_k,

531 mul_routed_weight,

532 + offset,

533 )

534

535

vllm/lora/punica_wrapper/punica_base.py

@@ -483,6 +483,8 @@ class PunicaWrapperBase(PunicaWrapperABC):

483 expand_config,

484 adapter_enabled: torch.Tensor,

485 mul_routed_weight=False,

486 + fully_sharded: bool = False,

487 + offset: int = 0,

488 ):

489 """

490 Performs a fused forward computation for LoRA of

vllm/lora/punica_wrapper/punica_gpu.py

@@ -375,6 +375,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):

375 expand_config,

376 adapter_enabled: torch.Tensor,

377 mul_routed_weight=False,

378 + fully_sharded: bool = False,

379 + offset: int = 0,

380 ):

381 """

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

@@ -408,4 +410,6 @@ class PunicaWrapperGPU(PunicaWrapperBase):

410 expand_config.get("NUM_STAGES", 3),

411 expand_config.get("SPLIT_K", 1),

412 mul_routed_weight,

413 + fully_sharded,

414 + offset,

415 )