dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

arch(moe): translate CCCL sync_handler.cuh — register-at-init, resolve-on-first-call

Browse Source 
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
1 changed files with 11 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
qwen3_6_scripts/qwen3_5.py
View File

@@ -921,6 +921,14 @@ class Qwen3_5MoeSparseBlock(nn.Module):
self.shared_expert_gate = ReplicatedLinear( self.shared_expert_gate = ReplicatedLinear(
hidden_size, 1, bias=False, quant_config=quant_config) 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( def _pure_pytorch_experts(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
@@ -1076,9 +1084,9 @@ class Qwen3_5MoeSparseBlock(nn.Module):
# ixf_F.vllm_invoke_fused_moe_kernel # ixf_F.vllm_invoke_fused_moe_kernel
# The original comment "ixformer lacks MoE kernels" may have been # The original comment "ixformer lacks MoE kernels" may have been
# wrong or outdated. Try native first, catch and fallback if it fails. # wrong or outdated. Try native first, catch and fallback if it fails.
# cc_dispatch pattern: _hw_policy detected MoE kernel availability at # cc_dispatch + sync_handler: strategy resolved on first call,
# module load. Skip native attempt entirely if we know it will fail. # pre-registered field checked as None (no hasattr overhead).
if not hasattr(self, '_use_native_moe'): if self._use_native_moe is None:
_hw_policy.detect(hidden_states.device) _hw_policy.detect(hidden_states.device)
# Only try native if at least align+invoke are available # Only try native if at least align+invoke are available
# (topk_softmax has PyTorch fallback in _custom_ops.py) # (topk_softmax has PyTorch fallback in _custom_ops.py)