Skip to content

Fix DeepEP TMA constraint violation for MoE CUDA graph batch sizes #1267

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kevalmorabia97 wants to merge 1 commit into from
Closed

Fix DeepEP TMA constraint violation for MoE CUDA graph batch sizes #1267

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions modelopt/deploy/llm/generate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,9 +126,16 @@ def _find_max_position_embeddings(cfg: dict) -> int | None:

cuda_graph_config = None
if max_batch_size > 0:
batch_sizes = [2**i for i in range(int((max_batch_size - 1).bit_length()))] + [
max_batch_size
]
if ep > 1:
# DeepEP low-latency dispatch requires (num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0.
# num_max_dispatch_tokens_per_rank equals the CUDA-graph batch size, so filter out
# any batch size that would violate the constraint.
batch_sizes = [b for b in batch_sizes if (ep * b) % 4 == 0]
cuda_graph_config = CudaGraphConfig(
batch_sizes=[2**i for i in range(int((max_batch_size - 1).bit_length()))]
+ [max_batch_size],
batch_sizes=batch_sizes,
max_batch_size=max_batch_size,
enable_padding=True,
)
Comment on lines +136 to 141

@coderabbitai coderabbitai Bot Apr 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🌐 Web query:

In TensorRT-LLM (llmapi.CudaGraphConfig), is batch_sizes=[] valid input? If valid, what exact runtime behavior is expected? If invalid, what exception/error is raised and in which versions?

💡 Result:

No, batch_sizes=[] is not a valid input for llmapi.CudaGraphConfig in TensorRT-LLM. It raises a ValueError during validation in the @model_validator because max(self.batch_sizes) on an empty list fails. The validator sorts the list and computes derived_max = max(self.batch_sizes), which errors on empty input. To disable CUDA graphs or use defaults, pass batch_sizes=None (default) or max_batch_size=0. With batch_sizes=None and max_batch_size=0 (also default), it falls back to generating batch sizes up to 128. Setting max_batch_size=0 with batch_sizes=None likely results in no graphs captured, using eager execution (runtime behavior seen in issues like #4811, #5634 where max_batch_size=0 or cuda_graph_config: null disables graphs). This behavior is from recent versions (post-2025 commits at hash 656091bf), consistent across source code and issues. No version-specific differences found; validation prevents empty lists uniformly.

Citations:

🏁 Script executed:

# First, let's find and examine the file
fd -t f "generate.py" | grep -E "deploy/llm"

Repository: NVIDIA/Model-Optimizer

Length of output: 97

🏁 Script executed:

# Then examine the specific lines mentioned in the review
head -150 modelopt/deploy/llm/generate.py | tail -30

Repository: NVIDIA/Model-Optimizer

Length of output: 1401

🏁 Script executed:

# Get more context around the issue to understand the full function
sed -n '100,160p' modelopt/deploy/llm/generate.py

Repository: NVIDIA/Model-Optimizer

Length of output: 2378

Add a guard to check that batch_sizes is not empty before creating CudaGraphConfig.

The filter at line 136 can produce an empty list when ep > 1 and no batch sizes satisfy the (ep * b) % 4 == 0 constraint. Passing batch_sizes=[] to CudaGraphConfig raises a ValueError because its validator calls max(self.batch_sizes) on the empty list.

Suggested fix 
             batch_sizes = [b for b in batch_sizes if (ep * b) % 4 == 0]
-            cuda_graph_config = CudaGraphConfig(
-                batch_sizes=batch_sizes,
-                max_batch_size=max_batch_size,
-                enable_padding=True,
-            )
+            if batch_sizes:
+                cuda_graph_config = CudaGraphConfig(
+                    batch_sizes=batch_sizes,
+                    max_batch_size=max_batch_size,
+                    enable_padding=True,
+                )
+            else:
+                warnings.warn(
+                    "No CUDA graph batch sizes satisfy DeepEP TMA constraint "
+                    f"(ep={ep}, max_batch_size={max_batch_size}); disabling cuda graphs."
+                )
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@modelopt/deploy/llm/generate.py` around lines 136 - 141, The filtered
batch_sizes list can become empty (when ep > 1) causing CudaGraphConfig(...) to
raise when it calls max(self.batch_sizes); capture the original batch_sizes
before filtering, then after applying batch_sizes = [b for b in batch_sizes if
(ep * b) % 4 == 0] add a guard: if not batch_sizes, raise a clear ValueError (or
log + bail) that includes ep and the original batch_sizes so the caller knows
why no valid batch sizes exist; ensure this guard sits immediately before
constructing CudaGraphConfig so CudaGraphConfig(...) never receives an empty
list.

Expand Down
Loading