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

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 15, 21:00 UTC

[V1] EP/TP MoE + DP Attention (#13931)

Tyler Michael Smith committed Mar 5, 2025 at 00:27 UTC 72c62eae5f010e45e027f51593d68798f84ee789

17 files changed +250 -75

examples/offline_inference/data_parallel.py

+10 -7

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -# usage: VLLM_USE_V1=1 python examples/offline_inference/data_parallel.py

2 +# usage:

3 +# VLLM_TEST_ENABLE_EP=1 VLLM_USE_V1=1 \

4 +# python examples/offline_inference/data_parallel.py

5 # we need to have a launcher to create multiple data parallel

6 # ranks. And each rank will create a vLLM instance to process its own prompts.

7 import os

@@ -7,6 +9,9 @@ import os

9 from vllm import LLM, SamplingParams

10 from vllm.utils import get_open_port

12 +GPUs_per_dp_rank = 2

13 +DP_size = 2

14 +

16 def main(dp_size, dp_rank, dp_master_ip, dp_master_port, GPUs_per_dp_rank):

17 os.environ["VLLM_DP_RANK"] = str(dp_rank)

@@ -48,8 +53,8 @@ def main(dp_size, dp_rank, dp_master_ip, dp_master_port, GPUs_per_dp_rank):

53 max_tokens=16 * (dp_rank + 1))

55 # Create an LLM.

51 - llm = LLM(model="facebook/opt-125m",

52 - tensor_parallel_size=2,

56 + llm = LLM(model="ibm-research/PowerMoE-3b",

57 + tensor_parallel_size=GPUs_per_dp_rank,

58 enforce_eager=True)

59 outputs = llm.generate(prompts, sampling_params)

60 # Print the outputs.

@@ -62,14 +67,12 @@ def main(dp_size, dp_rank, dp_master_ip, dp_master_port, GPUs_per_dp_rank):

68 if __name__ == "__main__":

69 from multiprocessing import Process

65 - dp_size = 2

66 - GPUs_per_dp_rank = 2

70 dp_master_ip = "127.0.0.1"

71 dp_master_port = get_open_port()

72 procs = []

70 - for i in range(dp_size):

73 + for i in range(DP_size):

74 proc = Process(target=main,

72 - args=(dp_size, i, dp_master_ip, dp_master_port,

75 + args=(DP_size, i, dp_master_ip, dp_master_port,

76 GPUs_per_dp_rank))

77 proc.start()

78 procs.append(proc)

tests/kernels/test_moe.py

@@ -217,6 +217,7 @@ def test_mixtral_moe(dtype: torch.dtype):

217 intermediate_size=config.intermediate_size,

218 params_dtype=dtype,

219 tp_size=1,

220 + dp_size=1,

221 ).cuda()

222

223 # Load the weights

vllm/attention/layer.py

+2 -2

@@ -324,7 +324,7 @@ def unified_attention(

324 ) -> torch.Tensor:

325 forward_context: ForwardContext = get_forward_context()

326 attn_metadata = forward_context.attn_metadata

327 - self = forward_context.attn_layers[layer_name]

327 + self = forward_context.no_compile_layers[layer_name]

328 kv_cache = self.kv_cache[forward_context.virtual_engine]

329 return self.impl.forward(self, query, key, value, kv_cache, attn_metadata)

330

@@ -356,7 +356,7 @@ def unified_attention_with_output(

356 ) -> None:

357 forward_context: ForwardContext = get_forward_context()

358 attn_metadata = forward_context.attn_metadata

359 - self = forward_context.attn_layers[layer_name]

359 + self = forward_context.no_compile_layers[layer_name]

360 kv_cache = self.kv_cache[forward_context.virtual_engine]

361 self.impl.forward(self,

362 query,

vllm/compilation/backends.py

+3 -2

@@ -396,8 +396,9 @@ class VllmBackend:

396

397 cache_dir = self.compilation_config.cache_dir

398 os.makedirs(cache_dir, exist_ok=True)

399 - local_cache_dir = os.path.join(

400 - cache_dir, f"rank_{vllm_config.parallel_config.rank}")

399 + rank = vllm_config.parallel_config.rank

400 + dp_rank = vllm_config.parallel_config.data_parallel_rank

401 + local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")

402 self.compilation_config.local_cache_dir = local_cache_dir

403

404 disable_cache = envs.VLLM_DISABLE_COMPILE_CACHE

vllm/forward_context.py

+15 -7

@@ -25,16 +25,22 @@ batchsize_logging_interval: float = envs.VLLM_LOG_BATCHSIZE_INTERVAL

25 batchsize_forward_time: defaultdict = defaultdict(list)

28 +@dataclass

29 +class DPMetadata:

30 + num_tokens_across_dp: list[int]

31 + cu_tokens_across_dp_cpu: torch.Tensor

32 +

33 +

34 @dataclass

35 class ForwardContext:

36 # copy from vllm_config.compilation_config.static_forward_context

31 - attn_layers: dict[str, Any]

37 + no_compile_layers: dict[str, Any]

38 # TODO: extend to support per-layer dynamic forward context

39 attn_metadata: "AttentionMetadata" # set dynamically for each forward pass

40 # TODO: remove after making all virtual_engines share the same kv cache

41 virtual_engine: int # set dynamically for each forward pass

36 - num_tokens_across_dp: Optional[

37 - list[int]] = None # set dynamically for each forward pass

42 + # set dynamically for each forward pass

43 + dp_metadata: Optional[DPMetadata] = None

46 _forward_context: Optional[ForwardContext] = None

@@ -61,7 +67,7 @@ def set_forward_context(attn_metadata: Any,

67 need_to_track_batchsize = track_batchsize and attn_metadata is not None

68 if need_to_track_batchsize:

69 forward_start_time = time.perf_counter()

64 - num_tokens_across_dp = None

70 + dp_metadata: Optional[DPMetadata] = None

71 if vllm_config.parallel_config.data_parallel_size > 1:

72 dp_size = vllm_config.parallel_config.data_parallel_size

73 dp_rank = vllm_config.parallel_config.data_parallel_rank

@@ -82,15 +88,17 @@ def set_forward_context(attn_metadata: Any,

88 dtype=torch.int32)

89 from vllm.distributed.parallel_state import get_dp_group

90 dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group)

85 - num_tokens_across_dp = num_tokens_tensor.tolist()

91 + cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_tensor, dim=0)

92 + dp_metadata = DPMetadata(num_tokens_across_dp, cu_tokens_across_dp_cpu)

94 global _forward_context

95 prev_context = _forward_context

96 _forward_context = ForwardContext(

90 - attn_layers=vllm_config.compilation_config.static_forward_context,

97 + no_compile_layers=vllm_config.compilation_config.

98 + static_forward_context,

99 virtual_engine=virtual_engine,

100 attn_metadata=attn_metadata,

93 - num_tokens_across_dp=num_tokens_across_dp)

101 + dp_metadata=dp_metadata)

102 try:

103 yield

104 finally:

vllm/model_executor/layers/fused_moe/layer.py

+155 -32

@@ -8,9 +8,11 @@ import torch

8 from torch.nn.parameter import UninitializedParameter

10 import vllm.envs as envs

11 -from vllm.distributed import (get_tensor_model_parallel_rank,

11 +from vllm.config import get_current_vllm_config

12 +from vllm.distributed import (get_dp_group, get_tensor_model_parallel_rank,

13 get_tensor_model_parallel_world_size,

14 tensor_model_parallel_all_reduce)

15 +from vllm.forward_context import ForwardContext, get_forward_context

16 from vllm.logger import init_logger

17 from vllm.model_executor.custom_op import CustomOp

18 from vllm.model_executor.layers.quantization.base_config import (

@@ -18,6 +20,7 @@ from vllm.model_executor.layers.quantization.base_config import (

20 from vllm.model_executor.utils import set_weight_attrs

21 from vllm.platforms import current_platform

22 from vllm.platforms.interface import CpuArchEnum

23 +from vllm.utils import direct_register_custom_op

25 if current_platform.is_cuda_alike():

26 from .fused_moe import fused_experts

@@ -246,6 +249,51 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

249 forward_native = forward_cuda

250

251

252 +def determine_expert_map(

253 + ep_size: int, ep_rank: int,

254 + global_num_experts: int) -> Tuple[int, Optional[torch.Tensor]]:

255 + """

256 + Calculates how many experts should be assigned to each rank for EP and

257 + creates a mapping from global to local expert index. Experts are

258 + distributed evenly across ranks. Any remaining are assigned to the

259 + last rank.

260 +

261 + Args:

262 + ep_size (int): The size of the expert parallel group

263 + global_num_experts (int): The total number of experts in the model.

264 +

265 + Returns:

266 + Tuple[int, Optional[torch.Tensor]]: A tuple containing:

267 + - local_num_experts (int): The number of experts assigned

268 + to the current rank.

269 + - expert_map (Optional[torch.Tensor]): A tensor of shape

270 + (global_num_experts,) mapping from global to local index.

271 + Contains -1 for experts not assigned to the current rank.

272 + Returns None if ep_size is 1.

273 + """

274 + assert ep_size > 0

275 + if ep_size == 1:

276 + return (global_num_experts, None)

277 +

278 + local_num_experts = global_num_experts // ep_size

279 +

280 + # Create a tensor of size num_experts filled with -1

281 + expert_map = torch.full((global_num_experts, ), -1, dtype=torch.int32)

282 + # Create a expert map for the local experts

283 + if ep_rank < (ep_size - 1):

284 + # Each non-last rank gets local_num_experts experts.

285 + expert_map[ep_rank * local_num_experts:

286 + (ep_rank + 1) * local_num_experts] = \

287 + torch.arange(0, local_num_experts, dtype=torch.int32)

288 + else:

289 + # All remaining experts are assigned to the last rank.

290 + local_num_experts = (global_num_experts - ep_rank * local_num_experts)

291 +

292 + expert_map[-local_num_experts:] = \

293 + torch.arange(0, local_num_experts, dtype=torch.int32)

294 + return (local_num_experts, expert_map)

295 +

296 +

297 class FusedMoE(torch.nn.Module):

298 """FusedMoE layer for MoE models.

299

@@ -282,6 +330,7 @@ class FusedMoE(torch.nn.Module):

330 quant_config: Optional[QuantizationConfig] = None,

331 tp_size: Optional[int] = None,

332 ep_size: Optional[int] = None,

333 + dp_size: Optional[int] = None,

334 prefix: str = "",

335 custom_routing_function: Optional[Callable] = None,

336 scoring_func: str = "softmax",

@@ -293,16 +342,48 @@ class FusedMoE(torch.nn.Module):

342 if params_dtype is None:

343 params_dtype = torch.get_default_dtype()

344

345 + # For smuggling this layer into the fused moe custom op

346 + compilation_config = get_current_vllm_config().compilation_config

347 + if prefix in compilation_config.static_forward_context:

348 + raise ValueError("Duplicate layer name: {}".format(prefix))

349 + compilation_config.static_forward_context[prefix] = self

350 + self.layer_name = prefix

351 + self.use_direct_call = not envs.VLLM_TEST_ENABLE_EP

352 +

353 + # Note: here we guard against accessing the TP and DP groups when

354 + # uninitialized (this happens when testing)

355 self.tp_size = (tp_size if tp_size is not None else

356 get_tensor_model_parallel_world_size())

357 + tp_rank = 0 if self.tp_size == 1 else get_tensor_model_parallel_rank()

358 + self.dp_size = (dp_size

359 + if dp_size is not None else get_dp_group().world_size)

360 + self.dp_rank = (0

361 + if self.dp_size == 1 else get_dp_group().rank_in_group)

362 + self.global_num_experts = num_experts

363 +

364 if envs.VLLM_TEST_ENABLE_EP:

299 - self.ep_size = self.tp_size

365 + # Set TP size to 1 to adjust for EP and adjust EP size and rank

366 + # for DP attention.

367 + self.ep_rank = tp_rank + self.tp_size * self.dp_rank

368 + self.tp_rank = 0

369 + self.ep_size = self.tp_size * self.dp_size

370 self.tp_size = 1

371 +

372 + self.local_num_experts, self.expert_map = determine_expert_map(

373 + ep_size=self.ep_size,

374 + ep_rank=self.ep_rank,

375 + global_num_experts=self.global_num_experts)

376 else:

377 + # Adjust TP size for DP attention

378 + self.tp_rank = tp_rank + self.tp_size * self.dp_rank

379 + self.ep_rank = 0

380 + self.tp_size = self.tp_size * self.dp_size

381 self.ep_size = 1

382 + self.local_num_experts = self.global_num_experts

383 + self.expert_map = None

384 self.top_k = top_k

385 self.global_num_experts = num_experts

305 - self.local_num_experts = self.global_num_experts // self.ep_size

386 +

387 assert intermediate_size % self.tp_size == 0

388 self.intermediate_size_per_partition = intermediate_size // self.tp_size

389 self.reduce_results = reduce_results

@@ -316,26 +397,6 @@ class FusedMoE(torch.nn.Module):

397 self.scoring_func = scoring_func

398 self.e_score_correction_bias = e_score_correction_bias

399 self.activation = activation

319 - self.expert_map = None

320 -

321 - if self.ep_size > 1:

322 - # Create a tensor of size num_experts filled with -1

323 - self.expert_map = torch.full((self.global_num_experts, ),

324 - -1,

325 - dtype=torch.int32)

326 - # Create a expert map for the local experts

327 - ep_rank = get_tensor_model_parallel_rank()

328 - if ep_rank < (self.ep_size - 1):

329 - # Each non-last rank gets local_num_experts experts.

330 - self.expert_map[ep_rank * self.local_num_experts:

331 - (ep_rank + 1) * self.local_num_experts] = \

332 - torch.arange(0, self.local_num_experts, dtype=torch.int32)

333 - else:

334 - # All remaining experts are assigned to the last rank.

335 - self.local_num_experts = (self.global_num_experts -

336 - ep_rank * self.local_num_experts)

337 - self.expert_map[-self.local_num_experts:] = \

338 - torch.arange(0, self.local_num_experts, dtype=torch.int32)

400

401 if self.scoring_func != "softmax" and not self.use_grouped_topk:

402 raise ValueError("Only softmax scoring function is supported for "

@@ -493,9 +554,6 @@ class FusedMoE(torch.nn.Module):

554 if expert_id == -1:

555 return

556

496 - # TP rank is set to 0 if EP is enabled

497 - tp_rank = 0 if self.ep_size > 1 else get_tensor_model_parallel_rank()

498 -

557 # compressed-tensors checkpoints with packed weights are stored flipped

558 # TODO (mgoin): check self.quant_method.quant_config.quant_format

559 # against known CompressionFormat enum values that have this quality

@@ -539,8 +597,7 @@ class FusedMoE(torch.nn.Module):

597 final_shape = list(loaded_weight.shape)

598 if shard_id in ["w1", "w3"]:

599 final_shape[1] *= 2

542 - final_shape[shard_dim] = final_shape[

543 - shard_dim] // get_tensor_model_parallel_world_size()

600 + final_shape[shard_dim] = final_shape[shard_dim] // self.tp_size

601 param.materialize(final_shape, dtype=loaded_weight.dtype)

602

603 expert_data = param.data if full_load else param.data[expert_id]

@@ -567,7 +624,7 @@ class FusedMoE(torch.nn.Module):

624 shard_id=shard_id,

625 loaded_weight=loaded_weight,

626 expert_data=expert_data,

570 - tp_rank=tp_rank)

627 + tp_rank=self.tp_rank)

628 return

629

630 # Case weight scales and zero_points

@@ -584,7 +641,7 @@ class FusedMoE(torch.nn.Module):

641 shard_dim=shard_dim,

642 loaded_weight=loaded_weight,

643 expert_data=expert_data,

587 - tp_rank=tp_rank)

644 + tp_rank=self.tp_rank)

645 elif quant_method in [

646 FusedMoeWeightScaleSupported.GROUP.value,

647 FusedMoeWeightScaleSupported.BLOCK.value,

@@ -594,7 +651,7 @@ class FusedMoE(torch.nn.Module):

651 shard_dim=shard_dim,

652 loaded_weight=loaded_weight,

653 expert_data=expert_data,

597 - tp_rank=tp_rank,

654 + tp_rank=self.tp_rank,

655 load_full_w2=getattr(param, "load_full_w2", False))

656 elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value:

657 self._load_per_tensor_weight_scale(shard_id=shard_id,

@@ -621,7 +678,7 @@ class FusedMoE(torch.nn.Module):

678 shard_dim=shard_dim,

679 loaded_weight=loaded_weight,

680 expert_data=expert_data,

624 - tp_rank=tp_rank)

681 + tp_rank=self.tp_rank)

682 return

683

684 @staticmethod

@@ -665,10 +722,45 @@ class FusedMoE(torch.nn.Module):

722

723 return topk_weights, topk_ids

724

725 + def naive_multicast(self, x: torch.Tensor,

726 + cu_tokens_across_dp_cpu: torch.Tensor):

727 + assert (len(x.shape) == 2)

728 + buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),

729 + device=x.device,

730 + dtype=x.dtype)

731 +

732 + start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[

733 + self.dp_rank - 1]

734 + end = cu_tokens_across_dp_cpu[self.dp_rank]

735 + buffer[start:end, :].copy_(x)

736 + for idx in range(get_dp_group().world_size):

737 + start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]

738 + end = cu_tokens_across_dp_cpu[idx]

739 + get_dp_group().broadcast(buffer[start:end, :], idx)

740 +

741 + return buffer

742 +

743 def forward(self, hidden_states: torch.Tensor,

744 router_logits: torch.Tensor):

745 + if self.use_direct_call:

746 + return self.forward_impl(hidden_states, router_logits)

747 + else:

748 + return torch.ops.vllm.moe_forward(hidden_states, router_logits,

749 + self.layer_name)

750 +

751 + def forward_impl(self, hidden_states: torch.Tensor,

752 + router_logits: torch.Tensor):

753 assert self.quant_method is not None

754

755 + if self.dp_size > 1:

756 + cu_tokens_across_dp_cpu = get_forward_context(

757 + ).dp_metadata.cu_tokens_across_dp_cpu

758 +

759 + hidden_states = self.naive_multicast(hidden_states,

760 + cu_tokens_across_dp_cpu)

761 + router_logits = self.naive_multicast(router_logits,

762 + cu_tokens_across_dp_cpu)

763 +

764 # Matrix multiply.

765 final_hidden_states = self.quant_method.apply(

766 layer=self,

@@ -687,6 +779,14 @@ class FusedMoE(torch.nn.Module):

779 activation=self.activation,

780 )

781

782 + if self.dp_size > 1:

783 + start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[

784 + self.dp_rank - 1]

785 + end = cu_tokens_across_dp_cpu[self.dp_rank]

786 +

787 + all_hidden_states = get_dp_group().all_reduce(final_hidden_states)

788 + final_hidden_states = all_hidden_states[start:end, :]

789 +

790 if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):

791 # Default set to False. (May have to add shared expert outputs.)

792 final_hidden_states = tensor_model_parallel_all_reduce(

@@ -757,3 +857,26 @@ class FusedMoE(torch.nn.Module):

857 s += f", scoring_func='{self.scoring_func}', activation='{self.activation}'" # noqa: E501

858

859 return s

860 +

861 +

862 +def moe_forward(hidden_states: torch.Tensor, router_logits: torch.Tensor,

863 + layer_name: str) -> torch.Tensor:

864 + forward_context: ForwardContext = get_forward_context()

865 + self = forward_context.no_compile_layers[layer_name]

866 + assert self.quant_method is not None

867 +

868 + return self.forward_impl(hidden_states, router_logits)

869 +

870 +

871 +def moe_forward_fake(hidden_states: torch.Tensor, router_logits: torch.Tensor,

872 + layer_name: str) -> torch.Tensor:

873 + return torch.empty_like(hidden_states)

874 +

875 +

876 +direct_register_custom_op(

877 + op_name="moe_forward",

878 + op_func=moe_forward,

879 + mutates_args=[],

880 + fake_impl=moe_forward_fake,

881 + dispatch_key=current_platform.dispatch_key,

882 +)

vllm/model_executor/models/aria.py

+11 -7

@@ -46,7 +46,7 @@ class AriaImagePixelInputs(TypedDict):

46 pixel_values: torch.Tensor

47 pixel_mask: Optional[torch.Tensor]

48 """

49 - Shape:

49 + Shape:

50 pixel_values: `(batch_size * num_images, num_channels, height, width)`

51 pixel_mask: `(batch_size * num_images, height, width)`

52 """

@@ -135,11 +135,11 @@ class AriaProjector(nn.Module):

135 query numbers,

136 e.g., {1225: 128, 4900: 256}. This allows for different query sizes

137 based on image resolution.

138 - embed_dim (int): Embedding dimension.

139 - num_heads (int): Number of attention heads.

140 - kv_dim (int): Dimension of key and value.

141 - ff_dim (int): Hidden dimension of the feed-forward network.

142 - output_dim (int): Output dimension.

138 + embed_dim (int): Embedding dimension.

139 + num_heads (int): Number of attention heads.

140 + kv_dim (int): Dimension of key and value.

141 + ff_dim (int): Hidden dimension of the feed-forward network.

142 + output_dim (int): Output dimension.

143 norm_layer (nn.Module): Normalization layer. Default is nn.LayerNorm.

144

145 Outputs:

@@ -239,6 +239,7 @@ class AriaTextMoELayer(nn.Module):

239 self,

240 config: AriaTextConfig,

241 quant_config: Optional[QuantizationConfig],

242 + prefix: str = "",

243 ) -> None:

244 super().__init__()

245 self.config = config

@@ -254,6 +255,7 @@ class AriaTextMoELayer(nn.Module):

255 intermediate_size=config.intermediate_size,

256 quant_config=quant_config,

257 reduce_results=True,

258 + prefix=f"{prefix}.experts",

259 )

260 self.shared_experts = LlamaMLP(

261 config.hidden_size,

@@ -301,7 +303,9 @@ class AriaTextDecoderLayer(LlamaDecoderLayer):

303 prefix: str = "",

304 ) -> None:

305 super().__init__(config, cache_config, quant_config, prefix)

304 - self.mlp = AriaTextMoELayer(config, quant_config=quant_config)

306 + self.mlp = AriaTextMoELayer(config,

307 + quant_config=quant_config,

308 + prefix=f"{prefix}.mlp")

309

310

311 class AriaTextModel(LlamaModel, SupportsQuant):

vllm/model_executor/models/dbrx.py

+6 -2

@@ -65,6 +65,7 @@ class DbrxExperts(FusedMoE):

65 config: DbrxConfig,

66 quant_config: Optional[QuantizationConfig] = None,

67 params_dtype: Optional[torch.dtype] = None,

68 + prefix: str = "",

69 ):

70 super().__init__(

71 num_experts=config.ffn_config.moe_num_experts,

@@ -76,6 +77,7 @@ class DbrxExperts(FusedMoE):

77 renormalize=True,

78 quant_config=quant_config,

79 tp_size=get_tensor_model_parallel_world_size(),

80 + prefix=prefix,

81 )

82 self.config = config

83 self.tp_size = get_tensor_model_parallel_world_size()

@@ -139,6 +141,7 @@ class DbrxMoE(nn.Module):

141 config: DbrxConfig,

142 quant_config: Optional[QuantizationConfig] = None,

143 params_dtype: Optional[torch.dtype] = None,

144 + prefix: str = "",

145 ):

146 super().__init__()

147 self.d_model = config.d_model

@@ -150,7 +153,8 @@ class DbrxMoE(nn.Module):

153

154 self.experts = DbrxExperts(config=config,

155 quant_config=quant_config,

153 - params_dtype=self.params_dtype)

156 + params_dtype=self.params_dtype,

157 + prefix=f"{prefix}.experts")

158

159 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

160 orig_shape = hidden_states.shape

@@ -291,7 +295,7 @@ class DbrxBlock(nn.Module):

295 cache_config,

296 quant_config,

297 prefix=f"{prefix}.norm_attn_norm")

294 - self.ffn = DbrxMoE(config, quant_config)

298 + self.ffn = DbrxMoE(config, quant_config, prefix=f"{prefix}.ffn")

299

300 def forward(

301 self,

vllm/model_executor/models/jamba.py

+15 -6

@@ -47,7 +47,8 @@ class JambaMoE(nn.Module):

47 top_k: Optional[int] = None,

48 params_dtype: Optional[torch.dtype] = None,

49 tp_size: Optional[int] = None,

50 - quant_config: Optional[QuantizationConfig] = None):

50 + quant_config: Optional[QuantizationConfig] = None,

51 + prefix: str = ""):

52 super().__init__()

53 self.num_total_experts = num_experts or config.num_experts

54 self.top_k = top_k or config.num_experts_per_tok

@@ -70,7 +71,8 @@ class JambaMoE(nn.Module):

71 reduce_results=True,

72 renormalize=False,

73 use_grouped_topk=False,

73 - quant_config=quant_config)

74 + quant_config=quant_config,

75 + prefix=f"{prefix}.experts")

77 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

78 orig_shape = hidden_states.shape

@@ -92,13 +94,15 @@ class JambaMLP(JambaMoE):

94 config: JambaConfig,

95 params_dtype: Optional[torch.dtype] = None,

96 tp_size: Optional[int] = None,

95 - quant_config: Optional[QuantizationConfig] = None):

97 + quant_config: Optional[QuantizationConfig] = None,

98 + prefix: str = ""):

99 super().__init__(config,

100 num_experts=1,

101 top_k=1,

102 params_dtype=params_dtype,

103 tp_size=tp_size,

101 - quant_config=quant_config)

104 + quant_config=quant_config,

105 + prefix=prefix)

106

107

108 class JambaMambaDecoderLayer(nn.Module):

@@ -109,6 +113,7 @@ class JambaMambaDecoderLayer(nn.Module):

113 cache_config: Optional[CacheConfig] = None,

114 quant_config: Optional[QuantizationConfig] = None,

115 is_lora_enabled: Optional[bool] = False,

116 + prefix: str = "",

117 **kwargs) -> None:

118 super().__init__()

119 self.config = config

@@ -129,7 +134,9 @@ class JambaMambaDecoderLayer(nn.Module):

134

135 num_experts = config.layers_num_experts[layer_idx]

136 ffn_layer_class = JambaMoE if num_experts > 1 else JambaMLP

132 - self.feed_forward = ffn_layer_class(config, quant_config=quant_config)

137 + self.feed_forward = ffn_layer_class(config,

138 + quant_config=quant_config,

139 + prefix=f"{prefix}.feed_forward")

140 self.input_layernorm = RMSNorm(config.hidden_size,

141 eps=config.rms_norm_eps)

142 self.pre_ff_layernorm = RMSNorm(config.hidden_size,

@@ -211,7 +218,9 @@ class JambaAttentionDecoderLayer(nn.Module):

218

219 num_experts = config.layers_num_experts[layer_idx]

220 ffn_layer_class = JambaMoE if num_experts > 1 else JambaMLP

214 - self.feed_forward = ffn_layer_class(config, quant_config=quant_config)

221 + self.feed_forward = ffn_layer_class(config,

222 + quant_config=quant_config,

223 + prefix=f"{prefix}.feed_forward")

224 self.input_layernorm = RMSNorm(config.hidden_size,

225 eps=config.rms_norm_eps)

226 self.pre_ff_layernorm = RMSNorm(config.hidden_size,

vllm/model_executor/models/mixtral.py

@@ -71,6 +71,7 @@ class MixtralMoE(nn.Module):

71 params_dtype: Optional[torch.dtype] = None,

72 quant_config: Optional[QuantizationConfig] = None,

73 tp_size: Optional[int] = None,

74 + dp_size: Optional[int] = None,

75 prefix: str = ""):

76 super().__init__()

77 self.hidden_size = hidden_size

@@ -93,6 +94,7 @@ class MixtralMoE(nn.Module):

94 renormalize=True,

95 quant_config=quant_config,

96 tp_size=tp_size,

97 + dp_size=dp_size,

98 prefix=f"{prefix}.experts")

100 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

vllm/model_executor/models/olmoe.py

+3 -1

@@ -80,7 +80,8 @@ class OlmoeMoE(nn.Module):

80 reduce_results=True,

81 renormalize=False,

82 quant_config=quant_config,

83 - tp_size=tp_size)

83 + tp_size=tp_size,

84 + prefix=f"{prefix}.experts")

86 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

87 # NOTE: hidden_states can have either 1D or 2D shape.

@@ -212,6 +213,7 @@ class OlmoeDecoderLayer(nn.Module):

213 hidden_size=config.hidden_size,

214 intermediate_size=config.intermediate_size,

215 quant_config=quant_config,

216 + prefix=f"{prefix}.mlp",

217 )

218 self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5)

219 self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5)

vllm/model_executor/models/phimoe.py

+4 -1

@@ -249,6 +249,7 @@ class PhiMoE(nn.Module):

249 params_dtype: Optional[torch.dtype] = None,

250 quant_config: Optional[QuantizationConfig] = None,

251 tp_size: Optional[int] = None,

252 + prefix: str = "",

253 ):

254 super().__init__()

255 self.hidden_size = hidden_size

@@ -272,7 +273,8 @@ class PhiMoE(nn.Module):

273 renormalize=False,

274 quant_config=quant_config,

275 tp_size=tp_size,

275 - custom_routing_function=phimoe_routing_function)

276 + custom_routing_function=phimoe_routing_function,

277 + prefix=f"{prefix}.experts")

278

279 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

280 # NOTE: hidden_states can have either 1D or 2D shape.

@@ -396,6 +398,7 @@ class PhiMoEDecoderLayer(nn.Module):

398 hidden_size=config.hidden_size,

399 intermediate_size=config.intermediate_size,

400 quant_config=quant_config,

401 + prefix=f"{prefix}.block_sparse_moe",

402 )

403 self.input_layernorm = nn.LayerNorm(config.hidden_size,

404 eps=config.rms_norm_eps,

vllm/model_executor/models/qwen2_moe.py

+5 -2

@@ -100,6 +100,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):

100 self,

101 config: PretrainedConfig,

102 quant_config: Optional[QuantizationConfig] = None,

103 + prefix: str = "",

104 ):

105 super().__init__()

106 self.tp_size = get_tensor_model_parallel_world_size()

@@ -115,7 +116,8 @@ class Qwen2MoeSparseMoeBlock(nn.Module):

116 intermediate_size=config.moe_intermediate_size,

117 reduce_results=False,

118 renormalize=config.norm_topk_prob,

118 - quant_config=quant_config)

119 + quant_config=quant_config,

120 + prefix=f"{prefix}.experts")

121

122 self.gate = ReplicatedLinear(config.hidden_size,

123 config.num_experts,

@@ -277,7 +279,8 @@ class Qwen2MoeDecoderLayer(nn.Module):

279 config.num_experts > 0 and

280 (layer_idx + 1) % config.decoder_sparse_step == 0):

281 self.mlp = Qwen2MoeSparseMoeBlock(config=config,

280 - quant_config=quant_config)

282 + quant_config=quant_config,

283 + prefix=f"{prefix}.mlp")

284 else:

285 self.mlp = Qwen2MoeMLP(

286 hidden_size=config.hidden_size,

vllm/platforms/cuda.py

@@ -111,6 +111,7 @@ class CudaPlatformBase(Platform):

111 def check_and_update_config(cls, vllm_config: VllmConfig) -> None:

112 parallel_config = vllm_config.parallel_config

113 scheduler_config = vllm_config.scheduler_config

114 + compilation_config = vllm_config.compilation_config

115

116 if parallel_config.worker_cls == "auto":

117 if scheduler_config.is_multi_step:

@@ -150,6 +151,14 @@ class CudaPlatformBase(Platform):

151 "FlashMLA: Forcing kv cache block size to 64 since this"

152 " is currently the only block size supported by the kernel.")

153

154 + if (parallel_config.data_parallel_size > 1

155 + and compilation_config.use_cudagraph):

156 + logger.info(

157 + "Data Parallel: Forcing enforce eager to be True since DP is "

158 + "currently not supported with CUDA Graphs.")

159 + vllm_config.model_config.enforce_eager = True

160 + compilation_config.use_cudagraph = False

161 +

162 @classmethod

163 def get_current_memory_usage(cls,

164 device: Optional[torch.types.Device] = None

vllm/utils.py

+2 -2

@@ -2196,8 +2196,8 @@ def bind_kv_cache(

2196 from vllm.model_executor.models.utils import extract_layer_index

2197 layer_need_kv_cache = [

2198 layer_name for layer_name in ctx

2199 - if ctx[layer_name].attn_type in (AttentionType.DECODER,

2200 - AttentionType.ENCODER_DECODER)

2199 + if (hasattr(ctx[layer_name], 'attn_type') and ctx[layer_name].attn_type

2200 + in (AttentionType.DECODER, AttentionType.ENCODER_DECODER))

2201 ]

2202 layer_index_sorted = sorted(

2203 set(

vllm/v1/engine/core.py

-1

@@ -149,7 +149,6 @@ class EngineCore:

149 if not self.scheduler.has_unfinished_requests():

150 return EngineCoreOutputs(

151 outputs=[], scheduler_stats=self.scheduler.make_stats())

152 -

152 scheduler_output = self.scheduler.schedule()

153 output = self.model_executor.execute_model(scheduler_output)

154 engine_core_outputs = self.scheduler.update_from_output(

vllm/v1/worker/gpu_model_runner.py

+7 -3

@@ -17,6 +17,7 @@ from vllm.distributed.parallel_state import get_pp_group, graph_capture

17 from vllm.forward_context import set_forward_context

18 from vllm.inputs import INPUT_REGISTRY

19 from vllm.logger import init_logger

20 +from vllm.model_executor.layers.fused_moe import FusedMoE

21 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

22 from vllm.model_executor.model_loader import get_model

23 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs

@@ -1357,7 +1358,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1358 """

1359 Initialize KV cache based on `kv_cache_config`.

1360 Args:

1360 - kv_cache_config: Configuration for the KV cache, including the KV

1361 + kv_cache_config: Configuration for the KV cache, including the KV

1362 cache size of each layer

1363 """

1364 if len(kv_cache_config.groups) > 1:

@@ -1389,10 +1390,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1390

1391 def get_kv_cache_spec(self) -> KVCacheSpec:

1392 """

1392 - Generates the KVCacheSpec by parsing the kv cache format from each

1393 + Generates the KVCacheSpec by parsing the kv cache format from each

1394 Attention module in the static forward context.

1395 Returns:

1395 - KVCacheSpec: A dictionary mapping layer names to their KV cache

1396 + KVCacheSpec: A dictionary mapping layer names to their KV cache

1397 format. Layers that do not need KV cache are not included.

1398 """

1399

@@ -1400,6 +1401,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1401 block_size = self.vllm_config.cache_config.block_size

1402 kv_cache_spec: KVCacheSpec = {}

1403 for layer_name, attn_module in forward_ctx.items():

1404 + if isinstance(attn_module, FusedMoE):

1405 + continue

1406 +

1407 # TODO: Support other attention modules, e.g., sliding window,

1408 # cross-attention, MLA.

1409 assert isinstance(attn_module, Attention)