Make single-batch overlap compatible with offloading (#11614)

This commit is contained in:
fzyzcjy
2025-10-18 08:45:54 +08:00
committed by GitHub
parent dcb8f090ad
commit 33e9bbec35
3 changed files with 33 additions and 23 deletions
+3 -5
View File
@@ -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):