Phase 0: SM90 MegaMoE design doc, reference baseline, nsys script
- MEGAMOE_SM90_DESIGN.md: complete design document with finalized decisions (fused single kernel, cooperative + single-WG, dynamic BLOCK_M, etc.) - tests/test_mega_moe_sm90.py: PyTorch FP32/BF16 reference implementation for dispatch → L1 GEMM → SwiGLU → L2 GEMM → combine pipeline - scripts/run_nsys_mega_moe_sm90.sh: nsys profiling wrapper script - megamoe-research-reports/: research analysis of PR304/323/347/352/357/360
This commit is contained in:
52
scripts/run_nsys_mega_moe_sm90.sh
Executable file
52
scripts/run_nsys_mega_moe_sm90.sh
Executable file
@@ -0,0 +1,52 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
NUM_RANKS=${NUM_RANKS:-2}
|
||||
NUM_TOKENS=${NUM_TOKENS:-128}
|
||||
HIDDEN=${HIDDEN:-4096}
|
||||
INTERMEDIATE=${INTERMEDIATE:-2048}
|
||||
NUM_EXPERTS=${NUM_EXPERTS:-16}
|
||||
NUM_TOPK=${NUM_TOPK:-6}
|
||||
OUTPUT_DIR=${OUTPUT_DIR:-./nsys-results}
|
||||
|
||||
mkdir -p "$OUTPUT_DIR"
|
||||
|
||||
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
|
||||
REPORT_NAME="nsys_sm90_megamoe_t${NUM_TOKENS}_r${NUM_RANKS}_${TIMESTAMP}"
|
||||
|
||||
echo "=== SM90 MegaMoE nsys profiling ==="
|
||||
echo " Ranks: $NUM_RANKS"
|
||||
echo " Tokens: $NUM_TOKENS"
|
||||
echo " Hidden: $HIDDEN"
|
||||
echo " Intermediate: $INTERMEDIATE"
|
||||
echo " Experts: $NUM_EXPERTS"
|
||||
echo " TopK: $NUM_TOPK"
|
||||
echo " Output: $OUTPUT_DIR/$REPORT_NAME"
|
||||
echo ""
|
||||
|
||||
NSYS_FLAGS=(
|
||||
--trace=cuda,nvtx
|
||||
--cuda-memory-usage=true
|
||||
--force-overwrite=true
|
||||
--output="$OUTPUT_DIR/$REPORT_NAME"
|
||||
--stats=true
|
||||
)
|
||||
|
||||
TORCH_FLAGS=(
|
||||
--nproc_per_node="$NUM_RANKS"
|
||||
)
|
||||
|
||||
TEST_FLAGS=(
|
||||
--num-tokens="$NUM_TOKENS"
|
||||
--hidden="$HIDDEN"
|
||||
--intermediate-hidden="$INTERMEDIATE"
|
||||
--num-experts="$NUM_EXPERTS"
|
||||
--num-topk="$NUM_TOPK"
|
||||
)
|
||||
|
||||
nsys profile "${NSYS_FLAGS[@]}" \
|
||||
torchrun "${TORCH_FLAGS[@]}" \
|
||||
tests/test_mega_moe_sm90.py "${TEST_FLAGS[@]}"
|
||||
|
||||
echo ""
|
||||
echo "=== Done. Report: $OUTPUT_DIR/$REPORT_NAME.nsys-rep ==="
|
||||
Reference in New Issue
Block a user