Fix returning reference to temporary in moe gemm (#3332)
<!-- .github/pull_request_template.md -->
## 📌 Description
<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->
This issue was observed in compiler warnings and can cause runtime
crashes due to UB.
Example compiler warning:
```
[38/560] /usr/local/cuda/bin/nvcc --generate-dependencies-with-compile --dependency-output /workspace/flashinfer/build/aot/cached_ops/fused_moe_100/100_cutlass_kernel_file_gemm_grouped_sm100_M128_group13.generated.cuda.o.d -DPy_LIMITED_API=0x03090000 -D_GLIBCXX_USE_CXX11_ABI=1 -I/workspace/flashinfer/csrc/nv_internal -I/workspace/flashinfer/csrc/nv_internal/include -I/workspace/flashinfer/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include -I/workspace/flashinfer/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include -I/workspace/flashinfer/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels -I/workspace/flashinfer/build/aot/generated/cutlass_instantiations/100 -I/workspace/flashinfer/3rdparty/cccl/cub -I/workspace/flashinfer/3rdparty/cccl/libcudacxx/include -I/workspace/flashinfer/3rdparty/cccl/thrust -isystem /usr/include/python3.12 -isystem /usr/local/cuda/include -isystem /usr/local/lib/python3.12/dist-packages/tvm_ffi/include -isystem /usr/local/lib/python3.12/dist-packages/tvm_ffi/include -isystem /workspace/flashinfer/include -isystem /workspace/flashinfer/csrc -isystem /workspace/flashinfer/3rdparty/cutlass/include -isystem /workspace/flashinfer/3rdparty/cutlass/tools/util/include -isystem /workspace/flashinfer/3rdparty/spdlog/include --compiler-options=-fPIC --expt-relaxed-constexpr -static-global-template-stub=false -DFLASHINFER_ENABLE_FP8_E8M0 -DFLASHINFER_ENABLE_FP4_E2M1 -std=c++17 --threads=2 -use_fast_math -Xfatbin=-compress-all -DFLASHINFER_ENABLE_F16 -DFLASHINFER_ENABLE_BF16 -DFLASHINFER_ENABLE_FP8_E4M3 -DFLASHINFER_ENABLE_FP8_E5M2 -g -O0 -DCUTLASS_DEBUG_TRACE_LEVEL=2 -DCOMPILE_BLACKWELL_TMA_GEMMS -DCOMPILE_BLACKWELL_TMA_GROUPED_GEMMS -DENABLE_BF16 -DENABLE_FP8 -DENABLE_FP4 -DUSING_OSS_CUTLASS_MOE_GEMM -DCUTLASS_ENABLE_GDC_FOR_SM100=1 -gencode=arch=compute_100a,code=sm_100a -DFLASHINFER_ENABLE_FP8_E8M0 -DFLASHINFER_ENABLE_FP4_E2M1 -c /workspace/flashinfer/build/aot/generated/cutlass_instantiations/100/gemm_grouped/100/cutlass_kernel_file_gemm_grouped_sm100_M128_group13.generated.cu -o /workspace/flashinfer/build/aot/cached_ops/fused_moe_100/100_cutlass_kernel_file_gemm_grouped_sm100_M128_group13.generated.cuda.o
/workspace/flashinfer/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl: In instantiation of ‘ReturnType tensorrt_llm::kernels::cutlass_kernels_oss::construct_if_true(Args&& ...) [with bool FLAG = true; ReturnType = cutlass::epilogue::fusion::FusionCallbacks<cutlass::epilogue::Sm90PtrArrayTmaWarpSpecialized<4, 2, 32, false, false, 1>, cutlass::epilogue::fusion::ScaledAccPerRowBiasPerColScaleScatter<cutlass::layout::ColumnMajor, cutlass::half_t, float, cutlass::half_t, float>, cute::tuple<cute::C<128>, cute::C<256>, cute::C<64> >, cute::tuple<cute::Layout<cute::C<128>, cute::C<1> >, cute::Layout<cute::C<32>, cute::C<1> > >, cute::ComposedLayout<cute::Swizzle<3, 4, 3>, cute::smem_ptr_flag_bits<16>, cute::Layout<cute::tuple<cute::C<64>, cute::C<8> >, cute::tuple<cute::C<1>, cute::C<64> > > >, cute::SM90_U16x8_STSM_T>::Arguments&&; Args = {float, std::nullptr_t, const float**&, cute::tuple<cute::C<0>, cute::C<0>, long int>, const cutlass::half_t* const*, cute::tuple<cute::C<1>, cute::C<0>, long int>, const float**&, cute::tuple<cute::C<0>, cute::C<1>, long int>, cutlass::half_t*, cute::tuple<cute::C<1>, long int, long int>&, const int**&, int&, int&, bool&}]’:
/workspace/flashinfer/build/aot/generated/cutlass_instantiations/100/gemm_grouped/100/cutlass_kernel_file_gemm_grouped_sm100_M128_group13.generated.cu:11:0: required from here
/workspace/flashinfer/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl:64:48: warning: returning reference to temporary [-Wreturn-local-addr]
64 | return ReturnType{std::forward<Args>(args)...};
| ^
/workspace/flashinfer/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl: In instantiation of ‘ReturnType tensorrt_llm::kernels::cutlass_kernels_oss::construct_if_true(Args&& ...) [with bool FLAG = true; ReturnType = cutlass::epilogue::fusion::FusionCallbacks<cutlass::epilogue::Sm90PtrArrayTmaWarpSpecialized<4, 2, 32, false, false, 1>, cutlass::epilogue::fusion::ScaledAccPerColBiasPerRowScaleScatter<cutlass::layout::RowMajor, cutlass::half_t, float, cutlass::half_t, float>, cute::tuple<cute::C<128>, cute::C<256>, cute::C<64> >, cute::tuple<cute::Layout<cute::C<128>, cute::C<1> >, cute::Layout<cute::C<32>, cute::C<1> > >, cute::ComposedLayout<cute::Swizzle<2, 4, 3>, cute::smem_ptr_flag_bits<16>, cute::Layout<cute::tuple<cute::C<8>, cute::C<32> >, cute::tuple<cute::C<32>, cute::C<1> > > >, cute::AutoVectorizingCopyWithAssumedAlignment<128> >::Arguments&&; Args = {float, std::nullptr_t, const float**&, cute::tuple<cute::C<0>, cute::C<0>, long int>, const cutlass::half_t* const*, cute::tuple<cute::C<0>, cute::C<1>, long int>, const float**&, cute::tuple<cute::C<1>, cute::C<0>, long int>, cutlass::half_t*, cute::tuple<long int, cute::C<1>, long int>&, const int**&, int&, int&, bool&}]’:
/workspace/flashinfer/build/aot/generated/cutlass_instantiations/100/gemm_grouped/100/cutlass_kernel_file_gemm_grouped_sm100_M128_group13.generated.cu:18:0: required from here
/workspace/flashinfer/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl:64:48: warning: returning reference to temporary [-Wreturn-local-addr]
```
## 🔍 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
- [ ] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [ ] I have installed the hooks with `pre-commit install`.
- [ ] 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
- [ ] Tests have been added or updated as needed.
- [ ] 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
* **Refactor**
* Improved type normalization and return handling in GPU kernel
utilities to ensure value-category decay and safer return types,
enhancing type safety and compatibility within the tensor optimization
layer.
<!-- review_stack_entry_start -->
[](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3332)
<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai --> B
Ben Barsdell committed
150f103a261984d88943d1a7962ac01454053638
Parent: bff85f3
Committed by GitHub <noreply@github.com>
on 5/25/2026, 10:55:38 PM