feat(cute_dsl/moe): add `moe_output_memset_inplace` dense memset wrapper (#3328)
## ๐ Description Adds `moe_output_memset_inplace`, a thin C++ binding around `cudaMemsetAsync` for zeroing the active MoE output slice before GEMM2 finalize, and uses it to replace the two `moe_output.zero_()` call sites in `_moe_core_impl`. Mirrors TRT-LLM's `moe_output_memset_inplace` Path A (`cuteDslMoeUtilsOp.cpp` at the `!enable_alltoall || ep_size <= top_k` branch). Functionally equivalent to `tensor.zero_()` but with lower per-call launch overhead (one `cudaMemsetAsync` vs PyTorch's `FillFunctor<c10::BFloat16>` kernel) โ saves ~2-3 ยตs per call at the cells where memset cost is visible. **Stream selection**: the C++ binding takes an explicit `cuda_stream_ptr` parameter (PyTorch's current stream); the Python wrapper passes `_get_cuda_stream_ptr()`. This is required because the underlying `get_current_stream()` C++ helper resolves through `TVMFFIEnvGetStream`, not `at::cuda::getCurrentCUDAStream()` โ so the Python `torch.cuda.stream(...)` context manager would otherwise not propagate to the `cudaMemsetAsync` call and aux-stream memset overlap with surrounding GEMM work would be silently nullified. Same pattern as `moe_sort` / `flashinfer_moe_sort` in this same file. **Scope**: this entry point exposes only Path A. Current callers of the monolithic CuteDSL MoE API handle all-to-all outside this function, so TRT-LLM's internal-alltoall Path B (the sparse `moeOutputMemset` kernel) is not part of this API. The existing sparse `moe_output_memset` bindings remain available if a future internal-alltoall integration needs them. ## ๐ Related Issues <!-- Link any related issues here --> ## ๐ 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 - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## ๐งช Tests * `TestMoeOutputMemsetInplace` (GPU): parametrized correctness across bf16/fp16 ร 2 shapes, unsupported-dtype guard, and CUDA-graph capture/replay verification (which would fail if the wrapper routed to any stream other than the capture stream). * `TestMoeOutputMemsetInplaceContract` (CPU/mock-only): deterministic contract test that monkeypatches the FFI dispatch + stream-ptr getter and asserts the wrapper passes `_get_cuda_stream_ptr()` as the 4th FFI argument. - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added an optimized, dense-only asynchronous buffer-clear operation for mixture-of-experts outputs with CUDA-graph-compatible stream handling. * **Tests** * Added tests verifying zeroing behavior across supported dtypes/shapes, erroring for unsupported dtypes, and correct behavior under CUDA graph capture. * **Documentation** * Clarified generated module docstring to better describe existing sparse zeroing behavior and retention for future integrations. <!-- review_stack_entry_start --> [](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3328) <!-- review_stack_entry_end --> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
L
Lee Nau committed
34fe1ff09b3e157c63df7476cbed344045f94b7c
Parent: ce43023
Committed by GitHub <noreply@github.com>
on 5/18/2026, 9:17:38 PM