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:
project6
2026-08-07 09:14:49 +00:00
parent 43ede018a1
commit f140825a56

View File

@@ -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)