#!/bin/sh
# Aiden production-1.6 launch wrapper for sm121 (per-rank Engram row ranges).
#
# The model's two Engram tables (layers 1 and 14) hold a different number of
# rows each (384,006,168 / 384,016,682 for this snapshot), and the overlay
# refuses to start unless the ranks' ranges tile [0, N) exactly per layer.
# A full snapshot has no manifest, so the ranges come from SPARK_ENGRAM_ROWS;
# a full range on every rank would double-count row contributions in the
# reduce, so each rank gets its own contiguous slice. The per-layer row count
# is read from the checkpoint's safetensors headers at launch (the table is
# 384M rows of fp8 e4m3 + scale, so the per-layer N is not in config.json).
set -e
DIR="$(dirname "$(readlink -f "$0" 2>/dev/null || echo /)")/models--deepseek-ai--DeepSeek-V4.1-Flash/snapshots"
SNAP="$(ls -d /root/.cache/huggingface/hub/models--deepseek-ai--DeepSeek-V4.1-Flash/snapshots/* | head -1)"
R="${NODE_RANK:-0}"
export SPARK_ENGRAM_ROWS="$(python3 - "$SNAP" "$R" <<'PY'
import json, struct, sys
snap, rank = sys.argv[1], int(sys.argv[2])
def rows(path, key):
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        h = json.loads(f.read(n))
    return h[key]["shape"][0]
n1 = rows(snap + "/model-00047-of-00048.safetensors", "layers.1.engram.embed.weight")
n14 = rows(snap + "/model-00048-of-00048.safetensors", "layers.14.engram.embed.weight")
def tile(n, r, parts=4):
    base, rem = divmod(n, parts)
    sizes = [base + (1 if i < rem else 0) for i in range(parts)]
    lo = sum(sizes[:r])
    return lo, lo + sizes[r]
lo1, hi1 = tile(n1, rank)
lo14, hi14 = tile(n14, rank)
print(f"1:{lo1}:{hi1},14:{lo14}:{hi14}")
PY
)"
# the b12x RoCE runtime takes the same GID index the launcher resolved for NCCL
[ -n "${NCCL_IB_GID_INDEX:-}" ] && export B12X_ROCE_GID_INDEX="$NCCL_IB_GID_INDEX"
exec python3 -m sglang.launch_server "$@"