43ede018a1ba10d6f7366814151defb11cda421c
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)
project_6
Description
Languages
C++
41.8%
Cuda
31.6%
Python
22.2%
C
2.1%
CMake
1.1%
Other
1.1%