Merge dev: fused MoE kernels, Qwen3 support, model improvements

[drunkcoding]

Summary

• Fused MoE MLP kernel (CUTLASS BF16): 3-GEMM fused path (gate→up w/ SiLU-mul epilogue→down) with small-M and large-K tile dispatch; correctness and performance benchmarks confirm parity with Torch-native reference
• New CUDA benchmark tests: 11 test files in tests/cuda/ covering fused MLP, topk-softmax, masked-select, and activation kernels
• Qwen3 model support and improvements to DeepSeek-V2/V3, Grok, Mixtral, and NLLBMoE
• Router kernel implementation
• C++ core updates: memory allocator, parallel dispatcher, prefetch scheduler, and AIO layer improvements
• CI fixes: CUDA extension guard for CPU-only builds, codespell statics allowlist, Claude Code Review GitHub Actions workflow

Test plan

• CUDA tests pass: cd tests/cuda && cmake -B build && cmake --build build -j$(nproc)
• C++ unit tests pass: cd tests/cpp/unittest/queues && cmake -B build && cmake --build build -j$(nproc) && ctest --test-dir build -V
• Pre-commit hooks green: pre-commit run --all-files
• Inference smoke test: CUDA_VISIBLE_DEVICES=0 python examples/interface_example.py --model_name_or_path deepseek-ai/DeepSeek-V2-Lite-Chat --offload_dir <path>
• GitHub Actions: Formatting + Build Test CI green on this PR

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR merges a large dev branch focused on faster MoE inference (CUTLASS fused kernels + routing), adds Qwen3 MoE support, and modernizes the build/test/CI pipeline around new C++/CUDA core refactors.

Changes:

• Adds CUTLASS BF16 fused MoE MLP path + multiple CUDA benchmark/test binaries under tests/cuda/.
• Introduces Qwen3 MoE integration and refactors expert dispatch to a new dispatch_local(...)/wait_dispatch_local() flow.
• Overhauls build tooling (new top-level CMake, setup.py CUDA-extension gating, Docker verification, CI workflow updates) and adds/extends C++ unit tests.

Reviewed changes

Copilot reviewed 131 out of 161 changed files in this pull request and generated 20 comments.

Show a summary per file

File	Description
tests/cuda/test_single_gemm_tiled.cu	New CUTLASS MMA tiled GEMM test binary
tests/cuda/test_fused_mlp_cutlass.cu	Benchmark fused CUTLASS MoE MLP vs Torch reference
tests/cuda/test_fused_mlp.cu	CUTLASS back-to-back GEMM test harness
tests/cuda/test_expert_fusion_v2.cu	Experimental fused dual-GEMM + SiLU*mul kernel
tests/cuda/test_autotune_blocksize.cu	Tile/block-size autotuning demo
tests/cuda/CMakeLists.txt	Adds CUDA test/benchmark build targets
tests/cpp/unittest/utils/test_simple_object_pool.cpp	New unit tests for SimpleObjectPool
tests/cpp/unittest/utils/test_lfu_cache.cpp	New unit tests documenting LFUCache behavior
tests/cpp/unittest/utils/CMakeLists.txt	Adds utils unit test executables
tests/cpp/unittest/queues/test_lockfree_queue.cpp	Adds edge/concurrency tests for queue
tests/cpp/unittest/queues/CMakeLists.txt	Adds queue unit test targets
setup.py	Replaces op_builder flow with CUDAExtensions for _store/_engine
pyproject.toml	Adds ruff-format config and mypy config scoping
op_builder/prefetch.py	Removes old DeepSpeed-derived op builder
op_builder/all_ops.py	Removes op builder registry
op_builder/init.py	Removes op_builder package init logic
moe_infinity/utils/hf_config.py	Adds Qwen3 architecture parsing support
moe_infinity/utils/config.py	Updates ArcherConfig typing + defaults
moe_infinity/runtime/model_offload.py	Switches to _store extension + adds Qwen3 wiring
moe_infinity/runtime/compile.py	Adds TorchScript expert compilation helper
moe_infinity/ops/op_builder	Removes source-tree alias file
moe_infinity/ops/core	Removes source-tree alias file
moe_infinity/models/switch_transformers.py	Switches to async dispatch/wait execution path
moe_infinity/models/qwen.py	Adds Qwen3MoEBlock implementation
moe_infinity/models/nllb_moe.py	Switches to async dispatch/wait execution path
moe_infinity/models/modeling_deepseek_v2/modeling_deepseek.py	Tweaks generation cache handling
moe_infinity/models/mixtral.py	Switches to async dispatch/wait execution path
moe_infinity/models/grok.py	Switches to async dispatch/wait execution path
moe_infinity/models/deepseek.py	Reworks DeepSeek MoE gate routing and dispatch
moe_infinity/models/init.py	Exposes Qwen3MoEBlock
moe_infinity/kernel/router.py	Adds Triton routing kernels (softmax/topk)
moe_infinity/kernel/init.py	Exports router launchers
moe_infinity/entrypoints/openai/protocol.py	Maps OpenAI max_tokens to HF max_new_tokens
moe_infinity/entrypoints/big_modeling.py	Updates DeepSeek rotary hook wiring
moe_infinity/distributed/expert_executor.py	Refactors dispatch API to pass weights + wait method
moe_infinity/common/constants.py	Adds Qwen3 mapping + DeepSeek v2 import path update
extensions/test_io/CMakeLists.txt	Adds test I/O Python module build via CMake
extensions/kernel/utils.h	Adds kernel utilities + arch detection helpers
extensions/kernel/ops.h	Declares activation/topk_softmax kernel APIs
extensions/kernel/grouped_threadblock_swizzle.h	Adds grouped swizzle support for grouped GEMMs
extensions/kernel/fused_moe_mlp.h	Declares CUTLASS fused MoE MLP API
extensions/kernel/epilogue_utils.h	Adds CUTLASS epilogue helpers (SiLU + mul)
extensions/kernel/dispatch_utils.h	Adds dispatch macros for floating types
extensions/kernel/common_device.h	Adds device activation utility implementations
extensions/kernel/b2b_gemm_grouped_problem_visitor.h	Adds grouped B2B GEMM visitor implementation
extensions/kernel/activation_kernels.cu	Adds CUDA activation kernels for gating/activations
extensions/CMakeLists.txt	Adds CMake-based extension build targets
examples/readme_example.py	Makes example configurable and updates default checkpoint
examples/interface_example.py	Updates dataset/prompting and supports Qwen3
docker/verify_build.py	Adds build outputs verification script
docker/Dockerfile	Adds reproducible build/test container with CUTLASS
core/utils/threadsafe_queue.h	Adds NotifyAll to unblock waiters
core/utils/simple_object_pool.h	Adds vector include for getMany support
core/utils/logger.h	Adds vector/enum logging helpers + conditional macros
core/utils/lockfree_queue.h	Replaces lock-free impl with mutex-backed queue
core/utils/cuda_utils.h	Adds CUTLASS error helper + GPU timer
core/utils/cuda_utils.cpp	Handles null pointer in IsDevicePointer
core/utils/cache.h	Fixes LFUCache invalidation/reset + adds touch tracking
core/python/py_archer_prefetch.cpp	Exposes init_moe_layer/topk_softmax + new dispatcher APIs
core/python/expert_gemm.cu	Adds cuBLAS-based expert fused MLP prototype
core/prefetch/archer_prefetch_handle.cpp	Adds MOE_IO_THREADS support to tensor handle
core/parallel/expert_dispatcher.h	Refactors dispatcher inputs/outputs and adds notify/start
core/model/moe.cpp	Adds InitMoELayer/TopKSoftmax bridge
core/model/model_topology.h	Adds SetModuleMemoryFromDisk_Views declaration
core/model/model_topology.cpp	Adds pipelined disk->host->GPU path + view-based setup
core/model/fused_mlp.h	Adds legacy wrapper API header
core/model/fused_mlp.cu	Implements legacy wrapper calling fused CUTLASS path
core/memory/torch_caching_allocator.h	Adds Torch allocator replacement scaffolding
core/memory/torch_caching_allocator.cpp	Instantiates allocator replacer
core/memory/shared_memory.h	Adds shared memory helpers + deleter
core/memory/shared_memory.cpp	Implements shared memory helpers
core/memory/pinned_memory_pool.h	Adds pinned pool abstraction
core/memory/pinned_memory_pool.cpp	Implements pinned pool with register/unregister
core/memory/device_caching_allocator.cpp	Uses CUDA_CHECK for cudaMalloc retry
core/memory/caching_allocator_bk.h	Adds legacy/backup caching allocator header
core/memory/caching_allocator.h	Adds unified caching allocator + Torch hooks
core/core	Removes core path alias file
core/common/types.h	Adds enum/string helpers + utilities
core/common/sync.h	Adds Futex utility class
core/common/pytorch.h	Adds dtype conversion helpers + blob helpers
core/common/generator.h	Adds ID/UUID/time helpers
core/common/context.h	Adds global Context config
core/common/constant.h	Adds size constants
core/aio/archer_tensor_handle.h	Adds partition sizing + IO thread configuration
core/aio/archer_tensor_handle.cpp	Implements partitioned storage rollover
core/aio/archer_prio_aio_handle.h	Adds pinned pool + exit notifications + thread count
core/aio/archer_prio_aio_handle.cpp	Adds pinned-buffer chunked writes + scheduler CV
core/aio/archer_aio_threadpool.h	Adds round-robin counter
core/aio/archer_aio_threadpool.cpp	Replaces rand() scheduling with round-robin
core/aio/archer_aio_thread.h	Adds atomics + CVs for thread coordination
core/aio/archer_aio_thread.cpp	Reworks thread loop to use CVs (no busy-wait)
core/CMakeLists.txt	Adds archer_core static lib build definition
RELEASE.md	Updates build instructions (drops BUILD_OPS)
README.md	Updates benchmark table + citation formatting
MANIFEST.in	Updates packaging includes (kernel headers/cuda)
CMakeLists.txt	Adds top-level CMake build entrypoint
CITATIONS.md	Adds author to citations
.pre-commit-config.yaml	Adds statics allowlist + mypy hook
.github/workflows/publish.yml	Updates runner OS + wheel build command
.github/workflows/publish-test.yml	Updates runner OS + wheel build command
.github/workflows/pre-commit-format.yml	Updates runner OS
.github/workflows/claude.yml	Adds Claude workflow automation
.github/workflows/claude-code-review.yml	Adds Claude Code Review workflow
.github/workflows/build-test.yml	Switches to CPU-only sdist build in CI
.dockerignore	Adds docker ignore patterns

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

This PR integrates fused CUTLASS MoE MLP kernels and router support, adds Qwen3 MoE model support, and refactors the build/test tooling (CMake, setup.py, CI, Docker) to support the updated C++/CUDA core and new benchmarks/tests.

Changes:

• Adds/updates CUDA kernels and benchmarks (fused MoE MLP + router/topk-softmax/activations) and related build plumbing.
• Adds Qwen3 MoE integration and updates several model wrappers to use the new expert dispatch interface.
• Refactors core I/O/AIO, memory pools, dispatch infrastructure, and introduces new C++ unit tests + CI/Docker verification.

Reviewed changes

Copilot reviewed 131 out of 161 changed files in this pull request and generated 14 comments.

Show a summary per file

File	Description
tests/cuda/test_single_gemm_tiled.cu	Adds a CUTLASS tiled GEMM test program.
tests/cuda/test_fused_mlp_cutlass.cu	Adds Torch-vs-CUTLASS fused MLP benchmark and correctness check.
tests/cuda/test_fused_mlp.cu	Adds CUTLASS back-to-back GEMM test driver.
tests/cuda/test_expert_fusion_v2.cu	Adds an experimental fully fused expert GEMM kernel test.
tests/cuda/test_autotune_blocksize.cu	Adds a tile/blocksize autotuning example binary.
tests/cuda/CMakeLists.txt	Adds CUDA tests build configuration and Torch/CUTLASS wiring.
tests/cpp/unittest/utils/test_simple_object_pool.cpp	Adds gtest coverage for SimpleObjectPool.
tests/cpp/unittest/utils/test_lfu_cache.cpp	Adds gtest coverage for LFUCache behavior and reset/eviction.
tests/cpp/unittest/utils/CMakeLists.txt	Adds CMake targets for utils unit tests.
tests/cpp/unittest/queues/test_lockfree_queue.cpp	Adds edge-case and concurrency tests for queue implementation.
tests/cpp/unittest/queues/CMakeLists.txt	Adds CMake targets for queue unit tests.
setup.py	Refactors extension build to new _store/_engine CUDAExtensions + CUTLASS includes.
pyproject.toml	Adds ruff format settings and mypy config.
op_builder/prefetch.py	Removes legacy DeepSpeed-derived op builder file.
op_builder/all_ops.py	Removes legacy op discovery/build system.
op_builder/init.py	Removes legacy op-builder package init/closures.
moe_infinity/utils/hf_config.py	Adds Qwen3 parsing support and typing improvements.
moe_infinity/utils/config.py	Adjusts ArcherConfig types/defaults (trace_path optional, threads default).
moe_infinity/runtime/model_offload.py	Switches to loading new _store extension; adds Qwen3 patching and new dispatch path.
moe_infinity/runtime/compile.py	Adds TorchScript compilation helper for expert modules.
moe_infinity/ops/op_builder	Removes symlink-like path entry.
moe_infinity/ops/core	Removes symlink-like path entry.
moe_infinity/models/switch_transformers.py	Migrates dispatch to new executor API (dispatch_local + wait).
moe_infinity/models/qwen.py	Adds Qwen3MoEBlock implementation using executor + fused routing.
moe_infinity/models/nllb_moe.py	Migrates dispatch to new executor API and output handling.
moe_infinity/models/modeling_deepseek_v2/modeling_deepseek.py	Alters generation cache max length handling.
moe_infinity/models/mixtral.py	Migrates dispatch to new executor API and reshaping/casting.
moe_infinity/models/grok.py	Migrates dispatch to new executor API and casting.
moe_infinity/models/deepseek.py	Reworks DeepSeek MoE gating/routing and executor dispatch integration.
moe_infinity/models/init.py	Exports Qwen3MoEBlock.
moe_infinity/kernel/router.py	Adds Triton router kernels (softmax + fused softmax/topk).
moe_infinity/kernel/init.py	Exposes router kernel launch helpers.
moe_infinity/entrypoints/openai/protocol.py	Maps OpenAI max_tokens to HF max_new_tokens; raises default max_tokens.
moe_infinity/entrypoints/big_modeling.py	Updates deepseek module patching to deepseek_v2/deepseek_v3 paths.
moe_infinity/distributed/expert_executor.py	Changes executor API to pass weights; adds wait_dispatch_local.
moe_infinity/common/constants.py	Adds Qwen3 model mapping and updates deepseek import path.
extensions/test_io/CMakeLists.txt	Adds CMake to build/install a test IO Python module.
extensions/kernel/utils.h	Adds CUDA/CUTLASS arch detection helpers and tile calculator.
extensions/kernel/ops.h	Adds kernel function declarations (activations, topk_softmax).
extensions/kernel/grouped_threadblock_swizzle.h	Adds CUTLASS grouped swizzle/visitor for grouped GEMMs.
extensions/kernel/fused_moe_mlp.h	Adds fused CUTLASS MoE MLP interface header.
extensions/kernel/epilogue_utils.h	Adds custom CUTLASS epilogue utilities for fused SiLU+mul.
extensions/kernel/dispatch_utils.h	Adds ATen dispatch helper macros.
extensions/kernel/common_device.h	Adds device activation helpers (relu/silu/gelu) + warp activation.
extensions/kernel/b2b_gemm_grouped_problem_visitor.h	Adds grouped problem visitor for B2B GEMMs.
extensions/kernel/activation_kernels.cu	Adds CUDA activation/gating kernels and PyTorch bindings.
extensions/CMakeLists.txt	Adds CMake build/install for prefetch_op and test_io modules.
examples/readme_example.py	Makes README example configurable via CLI and updates default model.
examples/interface_example.py	Reworks dataset/prompting and moves MoE init earlier; adds Qwen3 handling.
docker/verify_build.py	Adds build verification script for produced shared libs and refactor artifacts.
docker/Dockerfile	Adds full build/test Dockerfile including CUTLASS build and extension verification.
core/utils/threadsafe_queue.h	Adds NotifyAll to wake all waiters.
core/utils/simple_object_pool.h	Adds include for new APIs/tests.
core/utils/logger.h	Adds LogStream streaming for vector and enum classes + conditional warn/fatal macros.
core/utils/lockfree_queue.h	Replaces lock-free queue with mutex-backed queue implementation (API rename).
core/utils/cuda_utils.h	Adds CUTLASS error checking and a GPU timer utility.
core/utils/cuda_utils.cpp	Adds nullptr handling to IsDevicePointer.
core/utils/cache.h	Fixes LFU cache iterator invalidation and reset behavior; adds touch counter.
core/python/py_archer_prefetch.cpp	Extends Python bindings with init_moe_layer/topk_softmax and wait_hidden_states.
core/python/expert_gemm.cu	Adds cuBLAS-based expert fused MLP kernels and PyBind module.
core/prefetch/archer_prefetch_handle.cpp	Adds MOE_IO_THREADS support and passes IO thread count to tensor handle.
core/parallel/expert_dispatcher.h	Refactors dispatcher inputs/outputs, threading/queues, adds notify hooks and JIT path.
core/model/moe.cpp	Adds MoELayer init/topk-softmax entrypoints for Python.
core/model/model_topology.h	Adds SetModuleMemoryFromDisk_Views declaration.
core/model/model_topology.cpp	Adds pipelined disk→host→GPU path and view-only host tensor setup.
core/model/fused_mlp.h	Adds legacy wrapper declaration pointing at new fused_moe_ffn_into path.
core/model/fused_mlp.cu	Implements legacy wrapper delegating to fused_moe_ffn_into.
core/memory/torch_caching_allocator.h	Adds Torch allocator integration to use caching allocator on load.
core/memory/torch_caching_allocator.cpp	Instantiates allocator replacement singleton.
core/memory/shared_memory.h	Adds shared memory attach/detach helpers and RAII deleter.
core/memory/shared_memory.cpp	Implements shared memory mmap attach/detach functions.
core/memory/pinned_memory_pool.h	Adds pinned host memory pool interface.
core/memory/pinned_memory_pool.cpp	Implements pinned pool allocation and acquire/release.
core/memory/device_caching_allocator.cpp	Improves cudaMalloc error handling via CUDA_CHECK.
core/memory/caching_allocator_bk.h	Adds backup templated caching allocator (CUDA/UM/host variants).
core/memory/caching_allocator.h	Adds unified caching allocator API incl. SHM/PIN/CUDA and Torch hooks.
core/core	Removes path entry file.
core/common/types.h	Adds enum/string helpers, constexpr utilities, and pointer helpers.
core/common/sync.h	Adds a Futex wrapper utility.
core/common/pytorch.h	Adds blob→tensor helpers and dtype conversion utilities.
core/common/generator.h	Adds UUID/time/id generators.
core/common/context.h	Adds global Context container for config values.
core/common/constant.h	Adds size constants and warp size constant.
core/aio/archer_tensor_handle.h	Adds partition size constant and IO-thread configurable ctor.
core/aio/archer_tensor_handle.cpp	Adds storage partition rollover and passes IO thread config to prio handle.
core/aio/archer_prio_aio_handle.h	Refactors prio AIO to use atomic exit + thread pool + pinned pool.
core/aio/archer_prio_aio_handle.cpp	Implements multi-thread IO scheduling, pinned pool writes, clean shutdown, notifications.
core/aio/archer_aio_threadpool.h	Adds round-robin enqueue counter.
core/aio/archer_aio_threadpool.cpp	Uses round-robin instead of rand for thread selection.
core/aio/archer_aio_thread.h	Adds atomics + condition variables for proper waiting.
core/aio/archer_aio_thread.cpp	Implements cv-based worker loop and done notification (removes busy-wait).
core/CMakeLists.txt	Adds archer_core static library build with CUDA sources and linking.
RELEASE.md	Updates release steps to remove BUILD_OPS usage.
README.md	Updates benchmark table and citation formatting/authors.
MANIFEST.in	Changes packaging includes (drops op_builder, adds extensions/kernel sources).
CMakeLists.txt	Adds top-level CMake build for core + extensions, Torch/CUTLASS discovery.
CITATIONS.md	Adds an author to citations.
.pre-commit-config.yaml	Adds mypy hook and updates codespell allowlist.
.github/workflows/publish.yml	Updates runner OS and removes BUILD_OPS from build step.
.github/workflows/publish-test.yml	Updates runner OS and removes BUILD_OPS from build step.
.github/workflows/pre-commit-format.yml	Updates runner OS.
.github/workflows/claude.yml	Adds Claude Code workflow.
.github/workflows/claude-code-review.yml	Adds Claude Code Review workflow.
.github/workflows/build-test.yml	Updates CI to build sdist on ubuntu-22.04 with CPU torch.
.dockerignore	Adds docker ignore rules for build artifacts and large files.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[drunkcoding]

Updated in local branch: wrapped the device-only specialization with #ifdef __CUDA_ARCH__ so host compilation no longer evaluates __CUDA_ARCH__ expressions.

[drunkcoding]

Fixed in local branch by adding missing <string> and <stdexcept> includes directly in core/common/context.h.

[drunkcoding]

Fixed in local branch: switched time_to_exit_ checks to time_to_exit_.load() in both the wait predicate and immediate exit branch.

[drunkcoding]

Adjusted in local branch: BLOCK_E is now triton.next_power_of_2(E) and invalid expert lanes are masked to -inf, so routing is no longer capped at a fixed 32-expert block.

[drunkcoding]

Addressed in the same local router update: removed the hard-coded BLOCK_E=128 path and switched to triton.next_power_of_2(E) to avoid the constant cap.

[drunkcoding]

Updated in local branch: removed hard-coded _GLIBCXX_USE_CXX11_ABI=0 and switched to Torch-provided compile flags via add_compile_options(${TORCH_CXX_FLAGS}).

[drunkcoding]

Not addressed in the current local diff for this round. I will keep this thread open for a follow-up cleanup of stale test commentary.

[drunkcoding]

Addressed locally: the stale "known bugs" wording in tests/cpp/unittest/utils/test_lfu_cache.cpp has been updated to regression-coverage framing, so the comment no longer implies the bug is still expected.

[drunkcoding]

Not addressed in the current local diff for this round. I will track this futex typing/portability concern in a separate follow-up patch.

[drunkcoding]

Addressed locally in core/common/sync.h: Futex is now constrained to futex-compatible 32-bit integral storage via static_asserts, and futex syscalls use explicit futex-address/expected-value casts (reinterpret_cast<int*>(&value_), static_cast(...)).

[drunkcoding]

Not addressed in the current local diff for this round. I will include the explicit <tuple> include in a follow-up header hygiene pass.

[drunkcoding]

Addressed locally: added explicit #include in core/memory/shared_memory.h for the std::tuple<void*, int> AttachSharedMemory declaration.

[drunkcoding]

Not addressed in the current local diff for this round. I will handle these missing standard includes in a follow-up source hygiene patch.

[drunkcoding]

Addressed locally: added explicit #include <errno.h>, #include <fcntl.h>, and #include <string.h> in core/memory/shared_memory.cpp to remove transitive-include reliance for errno/O_RDWR/strerror usage.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Pull request overview

Copilot reviewed 131 out of 161 changed files in this pull request and generated 13 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

This uses Python if / break control flow inside a @triton.jit kernel based on runtime values (p > top_vals[j]), which Triton does not support. Rewrite top-k selection using vectorized tl.where-based logic (e.g., repeated tl.max with masking for TOPK steps, or sorting via supported primitives), so the kernel can compile and execute correctly.

[Copilot]

The kernel computes logits via a fully unrolled scalar loop over H (and later over BLOCK_E), which will explode compile time and generate very large kernels for typical MoE hidden sizes (e.g., 2k–8k). Consider blocking over H (e.g., tl.arange tiles + partial reductions) or using a matmul-style approach (or call into an existing optimized kernel) to keep compile time and runtime reasonable.

[Copilot]

Importing moe_infinity._store at module import time will raise ImportError (and break importing moe_infinity.models.qwen) in CPU-only installs or builds where the CUDA extension is not present. Since the block already uses a runtime-injected self.lib, defer importing _store (or remove this import entirely) and load it lazily through the same _load_prefetch_lib() mechanism used elsewhere.

[Copilot]

The previous logic initialized next_states = torch.zeros_like(hidden_states) and then selectively added expert outputs, so treating 0 as a sentinel for 'unrouted' tokens was valid. Now next_states is produced by the executor and may legitimately contain zeros for routed tokens, which would be incorrectly overwritten by the source hidden_states. Use an explicit token participation mask (e.g., derived from router_mask) to merge routed vs unrouted tokens, or ensure the executor returns hidden_states for unrouted tokens internally.

[Copilot]

The previous logic initialized next_states = torch.zeros_like(hidden_states) and then selectively added expert outputs, so treating 0 as a sentinel for 'unrouted' tokens was valid. Now next_states is produced by the executor and may legitimately contain zeros for routed tokens, which would be incorrectly overwritten by the source hidden_states. Use an explicit token participation mask (e.g., derived from router_mask) to merge routed vs unrouted tokens, or ensure the executor returns hidden_states for unrouted tokens internally.

[Copilot]

cutlass::HostTensor typically requires explicit sync_device() before using device_data() to ensure device allocations are populated. As written, X/Wg/Wu are filled on host but never synced to device before launching the kernel, so the kernel can read uninitialized device buffers. Call X.sync_device(); Wg.sync_device(); Wu.sync_device(); (and ensure C is allocated on device as needed) before the launch.

[Copilot]

This header uses memcpy but does not include <cstring> (or another header guaranteeing the declaration). Add #include <cstring> to avoid compilation failures on stricter toolchains.

[Copilot]

This header uses memcpy but does not include <cstring> (or another header guaranteeing the declaration). Add #include <cstring> to avoid compilation failures on stricter toolchains.

[Copilot]

This header unconditionally includes Linux-only futex headers, which will fail to build on non-Linux platforms. If cross-platform builds are expected, guard this implementation with #if defined(__linux__) and provide a fallback (e.g., condition_variable-based) or move this file behind a Linux-only build option.

[Copilot]

The constant name appears to be a typo: kWrapSize should likely be kWarpSize (CUDA warp size). Renaming avoids confusion and improves searchability.