Fix no-padding CUDA graph admission (#31273)
This commit is contained in:
@@ -532,9 +532,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
|
||||
else:
|
||||
cuda_graph_bs = forward_batch.batch_size
|
||||
|
||||
graph_key = cuda_graph_bs
|
||||
if self.enable_pdmux:
|
||||
graph_key = f"{get_current_stream_idx()}_{cuda_graph_bs}"
|
||||
graph_key = self._make_graph_key(
|
||||
cuda_graph_bs,
|
||||
stream_idx=get_current_stream_idx() if self.enable_pdmux else None,
|
||||
variant_label=self._resolve_lora_variant(forward_batch),
|
||||
)
|
||||
|
||||
is_bs_supported = (
|
||||
self.backend.can_run(forward_batch, graph_key)
|
||||
|
||||
@@ -310,7 +310,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
|
||||
cuda_graph_bs = forward_batch.batch_size
|
||||
|
||||
is_bs_supported = (
|
||||
self.backend.can_run(forward_batch, cuda_graph_bs)
|
||||
self.backend.can_run(forward_batch, self._make_graph_key(cuda_graph_bs))
|
||||
if self.disable_padding
|
||||
else cuda_graph_bs <= self.max_bs
|
||||
)
|
||||
|
||||
@@ -307,7 +307,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
|
||||
cuda_graph_bs = forward_batch.seq_lens.numel()
|
||||
|
||||
is_bs_supported = (
|
||||
self.backend.can_run(forward_batch, cuda_graph_bs)
|
||||
self.backend.can_run(forward_batch, self._make_graph_key(cuda_graph_bs))
|
||||
if self.disable_padding
|
||||
else cuda_graph_bs <= self.max_bs
|
||||
)
|
||||
|
||||
@@ -218,7 +218,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
|
||||
cuda_graph_bs = forward_batch.seq_lens.numel()
|
||||
|
||||
is_bs_supported = (
|
||||
self.backend.can_run(forward_batch, cuda_graph_bs)
|
||||
self.backend.can_run(forward_batch, self._make_graph_key(cuda_graph_bs))
|
||||
if self.disable_padding
|
||||
else cuda_graph_bs <= self.max_bs
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user