47 lines
1.1 KiB
Bash
Executable File
47 lines
1.1 KiB
Bash
Executable File
#!/bin/bash
|
|
# Setup and run the layer test on B200 — no Docker, no vLLM
|
|
# Compiles the kernel raw and runs the comparison test
|
|
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
REPO_DIR="$(dirname "$SCRIPT_DIR")"
|
|
VENV_DIR="$REPO_DIR/tests/.venv"
|
|
|
|
echo "=== NVFP4 Layer Test Setup ==="
|
|
echo "Repo: $REPO_DIR"
|
|
echo ""
|
|
|
|
# 1. Create venv
|
|
if [ ! -d "$VENV_DIR" ]; then
|
|
echo "[1/4] Creating venv..."
|
|
python3 -m venv "$VENV_DIR"
|
|
else
|
|
echo "[1/4] Venv already exists, skipping"
|
|
fi
|
|
|
|
source "$VENV_DIR/bin/activate"
|
|
|
|
# 2. Install dependencies
|
|
echo "[2/4] Installing Python deps..."
|
|
pip install --upgrade pip -q
|
|
pip install -r "$SCRIPT_DIR/requirements.txt" -q
|
|
|
|
# 3. Build and install the kernel from source
|
|
echo "[3/4] Building kernel from source (this takes a few minutes)..."
|
|
cd "$REPO_DIR"
|
|
pip install -e . --no-build-isolation -q
|
|
|
|
# Build the CUTLASS NVFP4 GEMM C++ extension
|
|
echo " Building CUTLASS C++ extension..."
|
|
cd "$REPO_DIR/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm"
|
|
python setup.py install
|
|
|
|
# 4. Run the test
|
|
echo "[4/4] Running layer test..."
|
|
cd "$SCRIPT_DIR"
|
|
python3 layertest.py
|
|
|
|
echo ""
|
|
echo "=== Done ==="
|