Make single-batch overlap compatible with offloading (#11614)
This commit is contained in:
@@ -42,7 +42,7 @@ class CombineOverlapArgs:
|
||||
wait_event: torch.cuda.Event
|
||||
num_sms: int
|
||||
signal: Optional[torch.Tensor] = None
|
||||
threshold: int = -1
|
||||
threshold: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -61,8 +61,6 @@ def execute_sbo(
|
||||
forward_batch: ForwardBatch,
|
||||
alt_stream: Optional = None,
|
||||
):
|
||||
shared_output = None
|
||||
|
||||
dispatch_output = experts.dispatch(
|
||||
hidden_states, topk_idx, topk_weights, forward_batch
|
||||
)
|
||||
@@ -82,7 +80,7 @@ def execute_sbo(
|
||||
with deep_gemm_wrapper.configure_deep_gemm_num_sms(
|
||||
meta_overlap_args["compute_num_sms"]
|
||||
):
|
||||
shared_output = forward_shared_experts()
|
||||
forward_shared_experts()
|
||||
|
||||
hidden_states = experts.combine(
|
||||
hidden_states,
|
||||
@@ -92,7 +90,7 @@ def execute_sbo(
|
||||
overlap_args=combine_overlap_args,
|
||||
)
|
||||
|
||||
return hidden_states, shared_output
|
||||
return hidden_states
|
||||
|
||||
|
||||
def _compute_overlap_args(dispatch_output, alt_stream):
|
||||
|
||||
Reference in New Issue
Block a user