From f140825a563ad6ec698c8dbcfbd63058614a3ef6 Mon Sep 17 00:00:00 2001 From: project6 Date: Fri, 7 Aug 2026 09:14:49 +0000 Subject: [PATCH] =?UTF-8?q?arch(moe):=20translate=20CCCL=20sync=5Fhandler.?= =?UTF-8?q?cuh=20=E2=80=94=20register-at-init,=20resolve-on-first-call?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- qwen3_6_scripts/qwen3_5.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 18721d55..630f545f 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -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)