SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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 -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](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