arch(moe): translate CCCL sync_handler.cuh — register-at-init, resolve-on-first-call
sync_handler.cuh entire design (140 lines):
Centralized synchronization resource manager for GPU kernels.
Two-phase lifecycle:
Phase 1 (host, constexpr): registerResource(numStages) + registerPhase()
Declares what resources are needed. No allocation yet.
Phase 2 (device, once): clusterInitSync()
Initializes all mbarriers in one pass. After this, no more registration.
Key properties:
- Non-copyable, non-movable (single source of truth)
- Fixed-size arrays (mMaxNumResources=10) — no dynamic allocation
- Destructor asserts mHasInitialized (catch forgotten init)
- Block-strided barrier init (all warps participate)
Translation to MoeSparseBlock:
Previous: hasattr() checks in forward hot path to lazy-init _use_native_moe
Now: Pre-declare _use_native_moe=None in __init__ (Phase 1: registration)
First forward resolves it via _hw_policy (Phase 2: initialization)
Subsequent forwards: None-check is faster than hasattr()
Also pre-declare _moe_out_buf fields to avoid attribute creation in forward.
CCCL source: cub/cub/detail/warpspeed/sync_handler.cuh
Maps to: qwen3_6_scripts/qwen3_5.py (Qwen3_5MoeSparseBlock)
This commit is contained in:
@@ -921,6 +921,14 @@ class Qwen3_5MoeSparseBlock(nn.Module):
|
||||
self.shared_expert_gate = ReplicatedLinear(
|
||||
hidden_size, 1, bias=False, quant_config=quant_config)
|
||||
|
||||
# sync_handler.cuh: register resources at init, initialize once.
|
||||
# Pre-declare MoE strategy here (resolved on first forward when device
|
||||
# is known). _use_native_moe is set to None = "not yet decided".
|
||||
# This avoids hasattr() checks in the forward hot path.
|
||||
self._use_native_moe: Optional[bool] = None
|
||||
self._moe_out_buf: Optional[torch.Tensor] = None
|
||||
self._moe_out_buf_key: Optional[tuple] = None
|
||||
|
||||
def _pure_pytorch_experts(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -1076,9 +1084,9 @@ class Qwen3_5MoeSparseBlock(nn.Module):
|
||||
# ixf_F.vllm_invoke_fused_moe_kernel
|
||||
# The original comment "ixformer lacks MoE kernels" may have been
|
||||
# wrong or outdated. Try native first, catch and fallback if it fails.
|
||||
# cc_dispatch pattern: _hw_policy detected MoE kernel availability at
|
||||
# module load. Skip native attempt entirely if we know it will fail.
|
||||
if not hasattr(self, '_use_native_moe'):
|
||||
# cc_dispatch + sync_handler: strategy resolved on first call,
|
||||
# pre-registered field checked as None (no hasattr overhead).
|
||||
if self._use_native_moe is None:
|
||||
_hw_policy.detect(hidden_states.device)
|
||||
# Only try native if at least align+invoke are available
|
||||
# (topk_softmax has PyTorch fallback in _custom_ops.py)
|
||||
|
||||
Reference in New Issue
Block a user