Remove excessive CuteDslMoEWrapper memory allocation - #3404
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between 3a40dce24a505fe1a002782386b663faa1605abd and cf4ac21. 📒 Files selected for processing (3)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThis PR replaces per-layer preallocated MOE/GEMM buffers with persistent CUDA Stream/Event resources for CUDA-graph compatibility, deprecates ChangesCUDA-graph buffer preallocation removal
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the CuteDslMoEWrapper class to remove pre-allocated buffers, shifting instead to holding persistent CUDA stream and event resources for async-memset overlap and CUDA graph compatibility. The max_num_tokens parameter is deprecated and ignored. The associated tests are updated to reflect these changes. The review feedback highlights a high-severity concurrency risk where multi-threaded environments could suffer from race conditions on shared global CUDA resources when use_cuda_graph=False. Additionally, a high-severity concern is raised regarding CUDA graph memory safety due to dynamic allocation of moe_output inside run(), suggesting allowing user-managed output buffers. Finally, a deprecation warning is recommended for the deprecated max_num_tokens parameter.
| if use_cuda_graph: | ||
| self._allocate_buffers() | ||
|
|
||
| def _allocate_buffers(self) -> None: | ||
| """Pre-allocate all buffers for CUDA graph compatibility. | ||
|
|
||
| Buffers are sized to fit *any* ``tile_size in VALID_TILE_SIZES``, | ||
| not just ``self.tile_size``. Two distinct buffer-shape concerns: | ||
|
|
||
| - ``max_num_permuted_tokens`` is monotonically *increasing* in | ||
| ``tile_size`` (the ``(tile - 1) * num_local_experts`` padding | ||
| term grows faster than ``max_num_tiles`` shrinks), so | ||
| permuted-token-indexed buffers (``_gemm1_output``, | ||
| ``_gemm1_output_scale``, | ||
| ``out_permuted_idx_to_expanded_idx``) must be sized using | ||
| ``max(VALID_TILE_SIZES)``. | ||
| - ``max_num_tiles`` is monotonically *decreasing* in | ||
| ``tile_size``, so the tile-count-indexed moe_sort buffers | ||
| (``out_tile_idx_to_expert_idx``, | ||
| ``out_tile_idx_to_mn_limit``) must be sized using | ||
| ``min(VALID_TILE_SIZES)``. | ||
|
|
||
| Sizing this way means the ``use_prealloc`` gate at | ||
| ``_forward_with_tactic`` doesn't need a ``tile_size == | ||
| self.tile_size`` check at runtime: whichever tactic the | ||
| autotuner picked, the prealloc fits. This preserves the | ||
| wrapper's CUDA-graph contract (``run()`` is graph-safe with | ||
| ``use_cuda_graph=True``) regardless of which ``tile_size`` the | ||
| autotuner ends up choosing. | ||
| """ | ||
| smallest_tile = min(VALID_TILE_SIZES) | ||
| largest_tile = max(VALID_TILE_SIZES) | ||
| max_permuted_across_tiles = get_max_num_permuted_tokens( | ||
| self.max_num_tokens, self.top_k, self.num_local_experts, largest_tile | ||
| ) | ||
|
|
||
| # moe_sort buffers — allocate using smallest_tile so the | ||
| # tile-count-indexed buffers (out_tile_idx_to_expert_idx, | ||
| # out_tile_idx_to_mn_limit) are large enough for any tile_size, | ||
| # then override out_permuted_idx_to_expanded_idx (which scales | ||
| # with tile_size in the opposite direction) to fit the largest | ||
| # tile_size's max_num_permuted_tokens. | ||
| self._moe_sort_buffers = allocate_moe_sort_buffers( | ||
| num_tokens=self.max_num_tokens, | ||
| num_experts=self.num_experts, | ||
| top_k=self.top_k, | ||
| num_local_experts=self.num_local_experts, | ||
| tile_tokens_dim=smallest_tile, | ||
| device=self.device, | ||
| ) | ||
| self._moe_sort_buffers["out_permuted_idx_to_expanded_idx"] = torch.empty( | ||
| (max_permuted_across_tiles,), dtype=torch.int32, device=self.device | ||
| ) | ||
|
|
||
| # GEMM1 output (FP4 quantized) | ||
| self._gemm1_output = torch.empty( | ||
| (max_permuted_across_tiles, self.intermediate_size // 2), | ||
| dtype=torch.uint8, | ||
| device=self.device, | ||
| ) | ||
|
|
||
| # GEMM1 output scale | ||
| scale_size = max_permuted_across_tiles * ( | ||
| self.intermediate_size // self.sf_vec_size | ||
| ) | ||
| self._gemm1_output_scale = torch.empty( | ||
| (scale_size,), dtype=torch.uint8, device=self.device | ||
| ) | ||
|
|
||
| # Final output | ||
| self._moe_output = torch.empty( | ||
| (self.max_num_tokens, self.hidden_size), | ||
| dtype=self.output_dtype, | ||
| device=self.device, | ||
| ) | ||
|
|
||
| # CUDA resources | ||
| self._aux_stream = torch.cuda.Stream(device=self.device) | ||
| self._main_event = torch.cuda.Event() | ||
| self._memset_event = torch.cuda.Event() | ||
| self._aux_stream = torch.cuda.Stream(device=self.device) | ||
| self._main_event = torch.cuda.Event() | ||
| self._memset_event = torch.cuda.Event() |
There was a problem hiding this comment.
Concurrency Risk / Race Condition in Multi-threaded Environments
When use_cuda_graph=False, self._aux_stream, self._main_event, and self._memset_event are left as None. During execution, _moe_core_impl falls back to the module-level _get_cuda_graph_resources(), which returns a single globally shared set of stream and event resources.
If multiple threads concurrently execute CuteDslMoEWrapper instances with use_cuda_graph=False, they will share the same CUDA stream and events. This leads to severe race conditions (e.g., concurrent record() and wait() calls on the same events), resulting in premature synchronization, undefined behavior, or silent data corruption.
To ensure thread safety, we should always initialize instance-specific stream and event resources in __init__ regardless of use_cuda_graph.
self._aux_stream = torch.cuda.Stream(device=self.device)
self._main_event = torch.cuda.Event()
self._memset_event = torch.cuda.Event()| moe_output = torch.empty( | ||
| (num_tokens, self.hidden_size), | ||
| dtype=self.output_dtype, | ||
| device=x.device, | ||
| ) |
There was a problem hiding this comment.
CUDA Graph Memory Safety & API Inconsistency
With the removal of pre-allocated buffers, moe_output is now always allocated dynamically inside run() via torch.empty(). When capturing a CUDA graph, this dynamic allocation is captured. If the user does not manually keep the returned tensor alive across graph replays, PyTorch's caching allocator may reuse its memory for other eager-mode tensors, leading to silent memory corruption during replay.
To allow safe user-managed output buffers (and to align with the functional API cute_dsl_fused_moe_nvfp4 which already accepts moe_output), we should add moe_output: Optional[torch.Tensor] = None to the signature of CuteDslMoEWrapper.run() and only allocate it if it is not provided.
| moe_output = torch.empty( | |
| (num_tokens, self.hidden_size), | |
| dtype=self.output_dtype, | |
| device=x.device, | |
| ) | |
| if moe_output is None: | |
| moe_output = torch.empty( | |
| (num_tokens, self.hidden_size), | |
| dtype=self.output_dtype, | |
| device=x.device, | |
| ) |
| @@ -392,7 +396,6 @@ def __init__( | |||
| self.hidden_size = hidden_size | |||
| self.intermediate_size = intermediate_size | |||
| self.use_cuda_graph = use_cuda_graph | |||
There was a problem hiding this comment.
Missing Deprecation Warning for max_num_tokens
Since max_num_tokens is now deprecated and ignored, we should raise a DeprecationWarning if the user passes a non-None value. This ensures users are aware of the deprecation and can clean up their codebases accordingly.
| self.use_cuda_graph = use_cuda_graph | |
| self.use_cuda_graph = use_cuda_graph | |
| if max_num_tokens is not None: | |
| import warnings | |
| warnings.warn( | |
| "max_num_tokens is deprecated and ignored.", | |
| DeprecationWarning, | |
| stacklevel=2, | |
| ) |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/fused_moe/cute_dsl/fused_moe.py (1)
89-98:⚠️ Potential issue | 🟠 Major | ⚡ Quick winKey
_cuda_graph_resourcesby CUDA device (avoid global stream/event reuse across GPUs).
_cuda_graph_resourcesmemoizes one stream/event trio process-wide, and_moe_core_impl()falls back to it wheneveraux_stream,main_event, ormemset_eventis missing. The functional API never suppliesmain_event/memset_event, and the wrapper only supplies these whenuse_cuda_graph=True; otherwise it also hits the fallback. Becausetorch.cuda.Event()/torch.cuda.Stream()are created without an explicit device, they bind to the current CUDA device at first call and can be reused whenxis on a different GPU. Scope the cache by device (e.g.,x.device) and create events/streams on that device explicitly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/fused_moe.py` around lines 89 - 98, The _cuda_graph_resources cache is process-wide and can return events/streams created on a different GPU; change _get_cuda_graph_resources to key the cache by CUDA device and create events/streams explicitly on that device (use x.device or torch.cuda.current_device() when called), e.g., store per-device entries like _cuda_graph_resources[device] = {"main_event": torch.cuda.Event(device=device), "memset_event": torch.cuda.Event(device=device), "aux_stream": torch.cuda.Stream(device=device)}; update callers (e.g., _moe_core_impl and any functional wrappers that rely on aux_stream, main_event, memset_event) to pass the tensor/device to _get_cuda_graph_resources so they always receive device-local events/streams and avoid cross-GPU reuse.
🧹 Nitpick comments (1)
flashinfer/fused_moe/cute_dsl/fused_moe.py (1)
365-385: ⚡ Quick winWarn when
max_num_tokensis still provided.The docs now say this parameter is deprecated and ignored, but callers get no signal that the old sizing contract disappeared. Emitting a
DeprecationWarninghere makes the API change discoverable before downstream code keeps passing a dead config knob.♻️ Suggested change
+import warnings + from typing import Any, Dict, Optional, Tuple @@ """Initialize the MoE wrapper. @@ self.output_dtype = output_dtype self.device = device self.enable_pdl = enable_pdl + if max_num_tokens is not None: + warnings.warn( + "`max_num_tokens` is deprecated and ignored by " + "`CuteDslMoEWrapper`.", + DeprecationWarning, + stacklevel=2, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/fused_moe.py` around lines 365 - 385, In the MoE wrapper initializer where the parameter max_num_tokens is accepted (the constructor in fused_moe.py that currently documents "max_num_tokens" as deprecated), add a runtime deprecation warning: when max_num_tokens is not None, call warnings.warn with a clear message that max_num_tokens is deprecated and ignored, use category=DeprecationWarning and set an appropriate stacklevel (e.g., stacklevel=2) so callers see the source; import the warnings module if not already present.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@flashinfer/fused_moe/cute_dsl/fused_moe.py`:
- Around line 89-98: The _cuda_graph_resources cache is process-wide and can
return events/streams created on a different GPU; change
_get_cuda_graph_resources to key the cache by CUDA device and create
events/streams explicitly on that device (use x.device or
torch.cuda.current_device() when called), e.g., store per-device entries like
_cuda_graph_resources[device] = {"main_event": torch.cuda.Event(device=device),
"memset_event": torch.cuda.Event(device=device), "aux_stream":
torch.cuda.Stream(device=device)}; update callers (e.g., _moe_core_impl and any
functional wrappers that rely on aux_stream, main_event, memset_event) to pass
the tensor/device to _get_cuda_graph_resources so they always receive
device-local events/streams and avoid cross-GPU reuse.
---
Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/fused_moe.py`:
- Around line 365-385: In the MoE wrapper initializer where the parameter
max_num_tokens is accepted (the constructor in fused_moe.py that currently
documents "max_num_tokens" as deprecated), add a runtime deprecation warning:
when max_num_tokens is not None, call warnings.warn with a clear message that
max_num_tokens is deprecated and ignored, use category=DeprecationWarning and
set an appropriate stacklevel (e.g., stacklevel=2) so callers see the source;
import the warnings module if not already present.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4cad15ad-f5a8-447c-b1fd-596b20c33147
📥 Commits
Reviewing files that changed from the base of the PR and between 56d537a and b3f413dc51ead4d22de3a08a6e57fc36017deb0a.
📒 Files selected for processing (3)
flashinfer/fused_moe/cute_dsl/fused_moe.pyflashinfer/fused_moe/cute_dsl/tuner.pytests/moe/test_cute_dsl_fused_moe.py
|
/bot run |
|
/bot run |
|
@nv-yunzheq could you also review this? and cc @leejnau for vis |
|
@nvjullin could you rebase so that I can trigger another pipeline? thanks! |
b3f413d to
3a40dce
Compare
|
Rebased onto main |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/fused_moe/cute_dsl/fused_moe.py (1)
365-365: 💤 Low valueConsider emitting a deprecation warning when
max_num_tokensis provided.Users passing a non-None
max_num_tokensvalue may not notice it's being ignored unless they read the updated docstring. Awarnings.warn(..., DeprecationWarning)would surface this more explicitly.This is optional since the docstring already documents the deprecation clearly.
Also applies to: 381-385
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/fused_moe.py` at line 365, If callers pass a non-None max_num_tokens (which is deprecated) add an explicit warnings.warn call to surface this; in the functions/methods where max_num_tokens appears (e.g., the signature containing max_num_tokens: Optional[int] = None around the fused_moe.py region and the other occurrences at the 381-385 block), import the warnings module if needed and at the start of each function check "if max_num_tokens is not None:" then call warnings.warn("max_num_tokens is deprecated and will be ignored; please remove it.", DeprecationWarning, stacklevel=2) so users see the deprecation when they pass a value.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/fused_moe.py`:
- Line 365: If callers pass a non-None max_num_tokens (which is deprecated) add
an explicit warnings.warn call to surface this; in the functions/methods where
max_num_tokens appears (e.g., the signature containing max_num_tokens:
Optional[int] = None around the fused_moe.py region and the other occurrences at
the 381-385 block), import the warnings module if needed and at the start of
each function check "if max_num_tokens is not None:" then call
warnings.warn("max_num_tokens is deprecated and will be ignored; please remove
it.", DeprecationWarning, stacklevel=2) so users see the deprecation when they
pass a value.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: eb774c81-b335-496e-b565-e2a9a432e573
📥 Commits
Reviewing files that changed from the base of the PR and between b3f413dc51ead4d22de3a08a6e57fc36017deb0a and 3a40dce24a505fe1a002782386b663faa1605abd.
📒 Files selected for processing (3)
flashinfer/fused_moe/cute_dsl/fused_moe.pyflashinfer/fused_moe/cute_dsl/tuner.pytests/moe/test_cute_dsl_fused_moe.py
💤 Files with no reviewable changes (1)
- tests/moe/test_cute_dsl_fused_moe.py
✅ Files skipped from review due to trivial changes (1)
- flashinfer/fused_moe/cute_dsl/tuner.py
|
/bot run |
|
[FAILED] Pipeline #53509427: 12/20 passed |
First failure is known issue #3470. |
|
@nvjullin could you fix the conflicts? |
3a40dce to
cf4ac21
Compare
|
Done, it was a docstring style rewrite causing conflicts. |
|
/bot run tests/moe |
|
[FAILED] Pipeline #54060962: 11/20 passed |
samuellees
left a comment
samuellees
left a comment
There was a problem hiding this comment.
LGTM, thanks for the contribution!
Follows the NVFP4 wrapper (flashinfer-ai#3404): CUDA graph capture records allocations from its private pool, so pre-sizing workspace for a maximum batch never helped capture and only cost memory. Drop the output, routing, intermediate and scale buffers and keep just the persistent stream and events, which do have to exist before capture. max_num_tokens is now accepted but ignored, matching CuteDslMoEWrapper. This also removes the per-batch-size buffer caches, which grew without bound when max_num_tokens was unset, and makes run() return a caller-owned tensor instead of a view of wrapper storage that the next call overwrites. AI-assisted with Claude Code. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
5 participants
📌 Description
Resolves #3308.
Previously,
max_num_tokenswere only checked in cuda graph path, but now deprecated and unused.Dropped tests that test for pre-allocated buffers.
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
API Changes
max_num_tokensis now optional/deprecated (ignored at runtime); existing calls remain compatible.Performance Improvements
Tests
Documentation
max_num_tokens.