Fix no-padding CUDA graph admission (#31273)

This commit is contained in:
Mohammad Miadh Angkad
2026-07-19 19:28:35 -07:00
committed by GitHub
parent 8bf2ab9be9
commit 35f2d4f761
4 changed files with 8 additions and 6 deletions
@@ -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
)