062cb160cf
- 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
53 lines
1.2 KiB
Bash
Executable File
53 lines
1.2 KiB
Bash
Executable File
#!/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 ==="
|