xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-12 12:39:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

wip

Browse Source 
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
This commit is contained in:
Lucas Wilkinson 2025-05-23 03:31:49 +00:00
parent 00f526f55b
commit 18bf91e6a8
4 changed files with 17 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
examples/basic-ub.py
View File

@ -40,11 +40,11 @@ def main():
max_model_len=1024, max_model_len=1024,
#load_format="dummy", #load_format="dummy",
############### ###############
tensor_parallel_size=1, #tensor_parallel_size=1,
#data_parallel_size=2, data_parallel_size=2,
enable_expert_parallel=False, enable_expert_parallel=True,
############### ###############
enable_microbatching=True, #enable_microbatching=True,
) )
# Generate texts from the prompts. # Generate texts from the prompts.
# The output is a list of RequestOutput objects # The output is a list of RequestOutput objects

2
vllm/config.py
View File

@ -4332,7 +4332,7 @@ class VllmConfig:
logger.warning_once( logger.warning_once(
"Piecewise compilation is not supported with " "Piecewise compilation is not supported with "
"microbatching. Disabling piecewiseching compilation.") "microbatching. Disabling piecewiseching compilation.")
self.compilation_config.level = CompilationLevel.DYNAMO_ONCE self.compilation_config.level = CompilationLevel.NO_COMPILATION
if self.model_config and self.model_config.use_mla and \ if self.model_config and self.model_config.use_mla and \

8
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
View File

@ -7,6 +7,7 @@ import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.model_executor.layers.fused_moe.utils import ( from vllm.model_executor.layers.fused_moe.utils import (
moe_kernel_quantize_input) moe_kernel_quantize_input)
from vllm.v1.worker.ubatching import get_current_ubatch_context, yield_impl
# Note use: layer.get_all_to_all() to get an AllToAll instance # Note use: layer.get_all_to_all() to get an AllToAll instance
@ -117,7 +118,11 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
do_send=send, do_send=send,
do_recv=not send, do_recv=not send,
) )
# if ubatch_ctx is not None:
# ubatch_ctx.gpu_stream_wait()
dispatch(True) # Send dispatch(True) # Send
yield_impl(gpu_wait=False)
dispatch(False) # Recv dispatch(False) # Recv
return expert_x, expert_x_scale, expert_num_tokens return expert_x, expert_x_scale, expert_num_tokens
@ -155,5 +160,8 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
do_send=send, do_send=send,
do_recv=not send, do_recv=not send,
) )
# if ubatch_ctx is not None:
# ubatch_ctx.gpu_stream_wait()
combine(True) combine(True)
yield_impl(gpu_wait=False)
combine(False) combine(False)

4
vllm/model_executor/models/deepseek_v2.py
View File

@ -49,6 +49,7 @@ from vllm.model_executor.model_loader.weight_utils import (
default_weight_loader, maybe_remap_kv_scale_name) default_weight_loader, maybe_remap_kv_scale_name)
from vllm.model_executor.sampling_metadata import SamplingMetadata from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import IntermediateTensors from vllm.sequence import IntermediateTensors
from vllm.v1.worker.ubatching import get_current_ubatch_context
from .interfaces import SupportsPP from .interfaces import SupportsPP
from .utils import (PPMissingLayer, is_pp_missing_parameter, from .utils import (PPMissingLayer, is_pp_missing_parameter,
@ -656,6 +657,9 @@ class DeepseekV2Model(nn.Module):
intermediate_tensors: Optional[IntermediateTensors], intermediate_tensors: Optional[IntermediateTensors],
inputs_embeds: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, IntermediateTensors]: ) -> Union[torch.Tensor, IntermediateTensors]:
if ubatch_ctx := get_current_ubatch_context() is not None:
print("in forward, ubatch:", ubatch_ctx.id)
if get_pp_group().is_first_rank: if get_pp_group().is_first_rank:
if inputs_embeds is not None: if inputs_embeds is not None:
hidden_states = inputs_embeds hidden_states = inputs_embeds