[BugFix] AssertionError: Do not capture num_reqs > max_num_reqs for uniform batch

[LucasWilkinson]

Fix for #25494 , optionally let the dispatcher decide the cudagraph_mode in the dummy run and remove assert

[ywang96]

Is this a duplicate of #25498?

[LucasWilkinson]

Is this a duplicate of #25498?

No not quite; #25498 will regress the fix #25407

[gemini-code-assist]

Code Review

This pull request addresses an AssertionError that occurs during dummy runs for CUDA graph capturing, specifically when the number of requests (num_reqs) exceeds the configured maximum (max_num_reqs). The fix involves removing the assertion and allowing the cudagraph_runtime_mode to be determined dynamically by the dispatcher. This is a sensible change, as dummy runs for profiling or graph capturing may need to simulate scenarios that don't adhere to normal request limits. The code modifications are clean and directly address the issue. I have one suggestion to improve the logic for determining the cudagraph_runtime_mode.

[gemini-code-assist]

The logic for determining cudagraph_runtime_mode has been refactored. While the new implementation is cleaner, it introduces a redundant call to self.cudagraph_dispatcher.dispatch in cases where cudagraph_runtime_mode is already CUDAGraphMode.NONE. In the previous version, this call was correctly skipped. To optimize this, we can check for CUDAGraphMode.NONE upfront and only call the dispatcher when necessary.

            if cudagraph_runtime_mode == CUDAGraphMode.NONE:
                _cg_mode = CUDAGraphMode.NONE
                batch_descriptor = None
            else:
                # filter out the valid batch descriptor
                _cg_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
                    BatchDescriptor(num_tokens=num_tokens,
                                    uniform_decode=uniform_decode))
                if cudagraph_runtime_mode is not None:
                    # sanity check
                    assert cudagraph_runtime_mode == _cg_mode, (
                        f"Cudagraph runtime mode mismatch at dummy_run. "
                        f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
                    )
                else:
                    cudagraph_runtime_mode = _cg_mode

[yewentao256]

Thanks for the work! Just question

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[yewentao256]

Verified that this also could solve #25494

[yewentao256]

Suggested change

	f"Cudagraph runtime mode mismatch at dummy_run. "
	assert cudagraph_runtime_mode in [CUDAGraphMode.NONE, _cg_mode], (