mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
9c8c85d83d
- Decode runs eagerly and compiles per-layer regions (attention half, MoE) with stage boundaries (neuron_ops/regions.py, DEEPINFER_NEURON_GRAPH_MODE=regions); layer KV views are sliced from the roots inside the graph (_LazyKVCache). - vllm-neuron#71 NKI kernels (Apache-2.0, vendored under neuron_ops/nki/ with attribution), called via torch_neuronx.nki_hop.wrap_nki: blockquant (fp8/fp4 GEMMs incl. grouped einsum, MoE and shared expert on bf16 activations), sparse_attn / index_score / compress_pool (nki_attention.py), hc_mixes / hc_expand / hc_collapse / hc_attn_out / rmsnorm / swiglu (nki_hc.py, nki_norm.py). - Compile-safe formulations: row-major scales, no dtype bitcasts of graph inputs (uint8 weight tensors), whole-row KV stores, byte assembly in int32 narrowed once at the store, arithmetic int16 scale packing, plain fma/rsqrt and bf16 F.linear under compile; default backend deepinfer_neuron (torch fp8 carried as uint8, avoiding NCC_EOCP001 with the kernels' OCP fp8). - Host: tiny model 9/9 vs golden on all prompts; device: kernel tests pass (13 attention, 88 hc/norm); region decode graph compiles except the attention projection stage (being reduced). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
130 lines
6.2 KiB
Bash
Executable File
130 lines
6.2 KiB
Bash
Executable File
#!/bin/bash
|
|
# Neuron port test suite. Run after tests/neuron/setup_trn2.sh.
|
|
#
|
|
# bash tests/neuron/run_all.sh cpu # host-only: kernel unit tests + end-to-end vs golden (~1 h)
|
|
# CORES=0-7 bash tests/neuron/run_all.sh device # NeuronCore tests on the given logical cores
|
|
# CORES=0-7 bash tests/neuron/run_all.sh all
|
|
# bash tests/neuron/run_all.sh cpu deep_gemm flash_mla # only selected stages
|
|
#
|
|
# Device stages use the first core of CORES for single-core runs and CORES for multi-rank deep_ep
|
|
# (world sizes 2/4/8 limited by the core count). They refuse to start if those cores are in use.
|
|
# Logs: $WORK/test_logs/<timestamp>/<stage>.log; a PASS/FAIL table is printed at the end.
|
|
set -uo pipefail
|
|
|
|
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
|
WORK="${WORK:-$HOME/di}"
|
|
MODE="${1:-cpu}"; shift || true
|
|
ONLY=("$@")
|
|
CORES="${CORES:-0-7}"
|
|
LOGDIR="$WORK/test_logs/$(date +%Y%m%d-%H%M%S)"
|
|
mkdir -p "$LOGDIR"
|
|
# shellcheck disable=SC1091
|
|
source "$WORK/venv/bin/activate"
|
|
cd "$REPO" || exit 1
|
|
VENDOR="$REPO/deepinfer/neuron_ops/vendor"
|
|
T="$REPO/tests/neuron"
|
|
export DEEPINFER_ACCELERATOR=neuron
|
|
# Kernel unit tests compare bit-for-bit against references quantizing to the e4m3fn range (448); the end-to-end
|
|
# stages run the default trn2 e4m3 range (224, deepinfer/neuron_ops/fp8.py) and compare against the golden.
|
|
UNIT_FP8_MAX="${UNIT_FP8_MAX:-448}"
|
|
declare -a SUMMARY=()
|
|
|
|
first_core() { echo "${CORES%%[-,]*}"; }
|
|
core_count() { python3 - "$CORES" <<'EOF'
|
|
import sys
|
|
n = 0
|
|
for part in sys.argv[1].split(","):
|
|
a, _, b = part.partition("-")
|
|
n += (int(b) - int(a) + 1) if b else 1
|
|
print(n)
|
|
EOF
|
|
}
|
|
selected() { [ ${#ONLY[@]} -eq 0 ] && return 0; for s in "${ONLY[@]}"; do [ "$s" = "$1" ] && return 0; done; return 1; }
|
|
|
|
run() { # run <stage> <command...>
|
|
local stage="$1"; shift
|
|
selected "$stage" || return 0
|
|
echo "=== $stage"
|
|
local t0=$SECONDS
|
|
local fp8_max="$UNIT_FP8_MAX"
|
|
case "$stage" in e2e_*|device_diff) fp8_max="${DEEPINFER_NEURON_FP8_MAX:-224}" ;; esac
|
|
if DEEPINFER_NEURON_FP8_MAX="$fp8_max" "$@" > "$LOGDIR/$stage.log" 2>&1; then r=PASS; else r=FAIL; fi
|
|
SUMMARY+=("$(printf '%-6s %-28s %5ss %s' "$r" "$stage" $((SECONDS - t0)) "$LOGDIR/$stage.log")")
|
|
echo " $r ($((SECONDS - t0))s)"
|
|
[ "$r" = FAIL ] && tail -5 "$LOGDIR/$stage.log" | sed 's/^/ | /'
|
|
}
|
|
|
|
cores_free() { # are the NeuronCores in $CORES free? (other cores may be busy)
|
|
python3 - "$CORES" <<'EOF'
|
|
import json, subprocess, sys
|
|
want = set()
|
|
for part in sys.argv[1].split(","):
|
|
a, _, b = part.partition("-")
|
|
want.update(range(int(a), int(b) + 1) if b else [int(a)])
|
|
busy = [d["neuron_device"] for d in json.loads(subprocess.check_output(["neuron-ls", "-j"]))
|
|
if d.get("neuron_processes") and want & set(d["neuroncore_ids"])]
|
|
if busy:
|
|
print(f"NeuronCores in {sys.argv[1]} are in use (devices {busy}); not starting device stages")
|
|
sys.exit(1 if busy else 0)
|
|
EOF
|
|
}
|
|
|
|
cpu_stages() {
|
|
export DEEPINFER_NEURON_DEVICE=cpu
|
|
run deep_gemm env PYTHONPATH="$VENDOR" python "$T/deep_gemm/test_deep_gemm_neuron.py"
|
|
run deep_gemm_tk env PYTHONPATH="$VENDOR" python "$T/deep_gemm/test_interop_tile_kernels.py"
|
|
run flash_mla bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_flash_mla.py"
|
|
run deep_select bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_deep_select.py"
|
|
run tilelang_stub bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_tilelang.py"
|
|
run tile_kernels python -m pytest -q -p no:cacheprovider "$T/tile_kernels/test_tile_kernels.py"
|
|
for ws in 1 2 4; do
|
|
run "deep_ep_ws$ws" python "$T/deep_ep/test_ep.py" --num-processes $ws --port $((29531 + ws))
|
|
run "engram_ws$ws" python "$T/deep_ep/test_engram.py" --num-processes $ws --port $((29541 + ws))
|
|
run "bucket_ws$ws" python "$T/deep_ep/test_bucket.py" --num-processes $ws --port $((29551 + ws))
|
|
done
|
|
run sampling_kat env PYTHONPATH="$REPO" python "$T/sampling/kat.py"
|
|
run sampling_check env PYTHONPATH="$REPO" python "$T/sampling/check2.py"
|
|
# Upstream interface tests of tilelang_ops (compare the Neuron backend against the *_ref implementations).
|
|
for m in norm store_kv_cache load_paged_kv_cache fused_transform compress state_cache get_kv_index kv_index_cat \
|
|
swa_index candidate_blocks flash_mla_utils gather prepare_v4_decode_metadata random_sampling fused_sample \
|
|
sampler sum_smaller_probs structure_output; do
|
|
run "tl_$m" env PYTHONPATH="$T" python -m pytest -q -p no:cacheprovider --noconftest -p nie_skip \
|
|
"deepinfer/model_executor/layers/tilelang_ops/$m.py" -k "not benchmark and not 129280-4096"
|
|
done
|
|
run e2e_cpu python examples/neuron/run_deepinfer_tiny.py --device cpu --threads 32 --num-kv-blocks 32 \
|
|
--out "$WORK/e2e_cpu" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt" --compare
|
|
}
|
|
|
|
device_stages() {
|
|
cores_free || return 0
|
|
local c0 n
|
|
c0="$(first_core)"; n="$(core_count)"
|
|
export DEEPINFER_NEURON_DEVICE=neuron
|
|
run device_primitives env NEURON_RT_VISIBLE_CORES="$c0" python "$T/device/test_device_primitives.py"
|
|
run nki_hc env NEURON_RT_VISIBLE_CORES="$c0" python -m pytest -q -p no:cacheprovider "$T/compile/test_nki_hc.py"
|
|
for ws in 2 4 8 16; do
|
|
[ "$ws" -le "$n" ] || continue
|
|
for t in ep engram bucket; do
|
|
run "${t}_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \
|
|
python "$T/deep_ep/test_$t.py" --num-processes $ws --port $((29600 + 10 * ws + ${#t}))
|
|
done
|
|
done
|
|
run e2e_neuron env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29733 \
|
|
python examples/neuron/run_deepinfer_tiny.py --device neuron --num-kv-blocks 32 \
|
|
--out "$WORK/e2e_neuron" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt" --compare
|
|
run device_diff env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29734 \
|
|
python examples/neuron/device_diff.py --prompts short --num-kv-blocks 32 \
|
|
--out "$WORK/device_diff" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt"
|
|
}
|
|
|
|
case "$MODE" in
|
|
cpu) cpu_stages ;;
|
|
device) device_stages ;;
|
|
all) cpu_stages; device_stages ;;
|
|
*) echo "usage: $0 cpu|device|all [stage...]"; exit 2 ;;
|
|
esac
|
|
|
|
echo; echo "===== summary ($LOGDIR)"
|
|
printf '%s\n' "${SUMMARY[@]}"
|
|
! printf '%s\n' "${SUMMARY[@]}" | grep -q '^FAIL'
|