[WIP] Fix stale routed expert accumulation across MoE forwards

[cla7aye15I4nd]

Reusing one MoE instance adds routed expert results from previous forwards because scatter_reduce includes the persistent destination buffer. With valid inputs, the first call matches the PyTorch reference but the second call is wrong.

Clear expert_cache in the routed stream immediately before reduction. Add a regression that reuses the same module for identical and changed inputs with top-k 1 and 2.

Validation:

• Before the fix: both regression cases fail on the second forward (about 98% of elements differ).
• After the fix: both regression cases and the existing fused-MoE example test pass.
• Targeted repository pre-commit hooks pass; independent static review found no actionable issue.

Run with TILELANG_EXECUTION_BACKEND=nvrtc pytest examples/fusedmoe/test_repeated_forward.py examples/fusedmoe/test_example_fusedmoe.py -q.

GPU validation used an RTX 5070, Python 3.12, PyTorch 2.14.1+cu130, and the TileLang 0.1.15 wheel running the changed example files from base 7a8ded9f1d91b0693fa82467656311947cfd8922. This was not a full source build or an all-backend test. NVCC host compilation was blocked by a CUDA/glibc header incompatibility; NVRTC compiled and ran the GPU kernels. Concurrent forwards on one instance and performance were not evaluated.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

---

_{Comment @coderabbitai help to get the list of available commands.}