perf(moe): translate CCCL smem_resource_raw.cuh — buffer reuse for MoE output
smem_resource_raw.cuh entire design (180 lines): Manages shared memory as multi-stage pipeline resources. Core idea: one memory region, multiple stages, barrier-synchronized. - mStageCount stages share the same SMEM base pointer - data() returns mPtrBase + mStageCurrent * mStride (stage rotation) - incrementStage() rotates, parity flips on wraparound - release/acquire protocol for producer-consumer sync Key insight: allocate once, reuse forever via stage rotation + zeroing. Translation to MoE _pure_pytorch_experts: Previous: torch.zeros_like(hidden_states) every call — GPU malloc + memset. Now: class-level _moe_out_buf, resized only when shape changes, .zero_() in-place (memset only, no malloc). On BI-V100 without async allocator, this eliminates a synchronous cudaMalloc per MoE layer per forward pass. With 28 MoE layers × 2 calls/step (prefill+decode), that is 56 fewer allocations per step. CCCL source: cub/cub/detail/warpspeed/resource/smem_resource_raw.cuh Maps to: qwen3_6_scripts/qwen3_5.py (Qwen3_5MoeSparseBlock)
This commit is contained in:
@@ -988,7 +988,17 @@ class Qwen3_5MoeSparseBlock(nn.Module):
|
||||
# to enable batched GEMM across expert groups (CCCL segmented_reduce pattern).
|
||||
# TODO: implement when we have benchmark data showing this path is hot.
|
||||
|
||||
out = torch.zeros_like(hidden_states)
|
||||
# smem_resource_raw.cuh: reuse buffer across calls.
|
||||
# CCCL manages SMEM as multi-stage ping-pong: same memory, different
|
||||
# stages. We do the same: keep a class-level buffer, resize only if
|
||||
# shape changes, zero in-place instead of allocating.
|
||||
_buf_key = (T, hidden_states.shape[-1])
|
||||
if not hasattr(self, '_moe_out_buf') or self._moe_out_buf_key != _buf_key:
|
||||
self._moe_out_buf = torch.zeros_like(hidden_states)
|
||||
self._moe_out_buf_key = _buf_key
|
||||
else:
|
||||
self._moe_out_buf.zero_()
|
||||
out = self._moe_out_buf
|
||||
|
||||
# Flatten all (token, expert) assignments: (T*top_k,) pairs
|
||||
flat_eids = topk_ids.view(-1) # (T*K,)
|
||||
|
||||
Reference in New Issue
Block a user