#!/usr/bin/env bash
# frank-medium container entrypoint: vllm-radiance MXFP4 on 2x Radeon AI PRO
# R9700 (gfx1201), serving just1moremodel/Qwen3.8-27B-Uncensored-MXFP4-awq.
#
# This is a fork of the patch-and-serve block inside radiance-vllm-mxfp4's
# serve-mxfp4.sh (clone at /home/dank/GitPackages/radiance-vllm-mxfp4, pinned
# commit 9735329348ae1c9319f01cbe5f7004ec5f4dba63). The published image alone
# is NOT sufficient for this model: the shipped libr4d predates the
# gated-delta-net overflow fix and NaNs the output, and the MXFP4 gates need
# relaxing, so the repo's patches are applied at every container start from
# the /patches mount. It runs INSIDE the container (executed as the image
# command by the frank-medium quadlet); the launcher's host-side GPU/TP/KV
# detection is not needed here because this fleet is fixed 2x gfx1201 and the
# KV pin is pinned explicitly in the quadlet.
#
# KEEP IN SYNC: if the radiance clone is updated, re-diff this file against
# the new serve-mxfp4.sh container command and against the DRY_RUN output.
#
# Env contract: every knob below is ${VAR:-default}; the quadlet sets the
# fleet-specific ones (SNAP, PORT, SERVED_NAMES, TP, KV_MEM, ...).
set -euo pipefail

# ---------------------------------------------------------------- serving shape
: "${SNAP:=/models/Qwen3.8-27B-Uncensored-MXFP4-awq}"
: "${DRAFTER:=/models/Qwen3.8-27B-DFlash2-FP8}"
: "${CHAT_TEMPLATE:=/patches/qwen-fixed-v22.3.jinja}"
: "${PORT:=7770}"
: "${SERVED_NAMES:=frank-medium}"
: "${MAXLEN:=262144}"
: "${MAXSEQS:=8}"
# KV_MEM=0 is passed THROUGH as --kv-cache-memory 0, which the image's
# patched KV sizing reads as "profile the shape" (the launcher's rule).
# Blanking the variable instead omits the flag and makes vLLM size the cache
# from --gpu-memory-utilization alone: measured on the single-GPU replicas
# that overshoots and the allocation dies with torch.OutOfMemoryError right
# after the weights load (2.16 GiB requested, 116 MiB free). Do not "fix"
# this into an empty expansion.
: "${KV_MEM:=18563072000}"   # shipped kv-profiles.tsv row 2x7551-32624 @ 8/8192/262144/dflash
: "${SPEC_METHOD:=dflash}"
if [ "$SPEC_METHOD" = dflash ]; then : "${SPEC:=7}"; else : "${SPEC:=4}"; fi
: "${MIN_M:=0}"              # 0 = the W4A8 kernel serves every M (MIN_M=1 still routes M=1 to aiter)
: "${R4D_ATTN:=1}"
: "${ASYNC:=0}"

if [ "$R4D_ATTN" = "1" ]; then ATTN=R4D; else ATTN=ROCM_AITER_UNIFIED_ATTN; fi
if [ "$ASYNC" = 1 ]; then ASYNC_FLAG="--async-scheduling"; UNPAD=false; else
  ASYNC_FLAG="--no-async-scheduling"; UNPAD=true; fi

# Drafting config. dflash: the DFlash2-FP8 block-diffusion drafter (its
# attention backend must support full cuda graphs -> TRITON_ATTN). The
# disable_padded_drafter_batch unpad is the single-stream lever (~+50%).
if [ "$SPEC_METHOD" = dflash ]; then
  if [ ! -f "$DRAFTER/config.json" ]; then
    echo "frank-medium: no dflash drafter at $DRAFTER" >&2
    exit 1
  fi
  DRAFT_SAMPLE=${DRAFT_SAMPLE:-greedy}
  SPEC_CFG="{\"method\":\"dflash\",\"model\":\"$DRAFTER\",\"num_speculative_tokens\":$SPEC,\"attention_backend\":\"TRITON_ATTN\",\"disable_padded_drafter_batch\":$UNPAD,\"draft_sample_method\":\"$DRAFT_SAMPLE\"}"
  RADIANCE_DRAFT_RERANK=${RADIANCE_DRAFT_RERANK:-80}
  RADIANCE_VERIFY_HEAD=${RADIANCE_VERIFY_HEAD:-1}
else
  SPEC_CFG="{\"method\":\"mtp\",\"num_speculative_tokens\":$SPEC,\"attention_backend\":\"$ATTN\",\"disable_padded_drafter_batch\":$UNPAD}"
  RADIANCE_DRAFT_RERANK=${RADIANCE_DRAFT_RERANK:-32}
  RADIANCE_VERIFY_HEAD=${RADIANCE_VERIFY_HEAD:-0}
fi

# The AR size gate compares raw bf16 bytes: CHUNK x hidden(5120) x 2. Derived
# so changing CHUNK cannot silently drop prefill back onto RCCL.
export RADIANCE_AR_MAX_KB=$(( (CHUNK * 5120 * 2) / 1024 + 4096 ))

# ---------------------------------------------------------------- runtime env
export HF_HUB_OFFLINE=1
export VLLM_LOGGING_LEVEL=${VLLM_LOGGING_LEVEL:-INFO}
export VLLM_NO_USAGE_STATS=${VLLM_NO_USAGE_STATS:-1}
export VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1
export VLLM_ROCM_USE_AITER_MHA=0 VLLM_ROCM_USE_AITER_MLA=0 VLLM_ROCM_USE_AITER_MOE=0
export VLLM_ROCM_USE_AITER_LINEAR=0 VLLM_ROCM_USE_AITER_FP8BMM=0
export VLLM_ROCM_USE_AITER_FP4BMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0
export NCCL_PROTO=Simple
export RADIANCE_USE_R4D=${RADIANCE_USE_R4D:-1} RADIANCE_USE_R4D_AR=${RADIANCE_USE_R4D_AR:-1}
export RADIANCE_DFLASH_SELECTOR_TOPK="${RADIANCE_DFLASH_SELECTOR_TOPK:-}"
export RADIANCE_R4D_REPORT=1
export RADIANCE_PRESHUFFLE=${RADIANCE_PRESHUFFLE:-1}
export RADIANCE_MXFP4_SANITIZE="${RADIANCE_MXFP4_SANITIZE:-0}"
export RADIANCE_MXFP4=1 RADIANCE_MXFP4_W4A8=1 RADIANCE_MXFP4_W4A8_MIN_M="$MIN_M"
export RADIANCE_FAST_DRAFT=${RADIANCE_FAST_DRAFT:-1}
export RADIANCE_DRAFT_TAU=${RADIANCE_DRAFT_TAU:-0.20}
export RADIANCE_DRAFT_RERANK
export RADIANCE_VERIFY_HEAD
export RADIANCE_VERIFY_HEAD_MAX_M=${RADIANCE_VERIFY_HEAD_MAX_M:-32}
export RADIANCE_MXFP4_DEBUG=${RADIANCE_MXFP4_DEBUG:-0}
export RADIANCE_MXFP4_PUREQUANT=${RADIANCE_MXFP4_PUREQUANT:-0}
export RADIANCE_MXFP4_SYNC=${RADIANCE_MXFP4_SYNC:-0}
export RADIANCE_MXFP4_CLONE=${RADIANCE_MXFP4_CLONE:-0}
export RADIANCE_MXFP4_CHECKX=${RADIANCE_MXFP4_CHECKX:-0}
export RADIANCE_MXFP4_PADOUT=${RADIANCE_MXFP4_PADOUT:-0}
export RADIANCE_MXFP4_TN4_MIN_M=${RADIANCE_MXFP4_TN4_MIN_M:-2048}
# The decode-kernel band must cover MAXSEQS x (SPEC+1) rows: 64 covers 8x8.
if [ "${MAXSEQS}" -gt 8 ]; then
  export RADIANCE_MXFP4_DECODE_MAX_M=${RADIANCE_MXFP4_DECODE_MAX_M:-128}
fi
export RADIANCE_MXFP4_DECODE_NT=${RADIANCE_MXFP4_DECODE_NT:-1}
export RADIANCE_MXFP4_A_TILED_MIN_M=${RADIANCE_MXFP4_A_TILED_MIN_M:-513}
export RADIANCE_MXFP4_WPERM=${RADIANCE_MXFP4_WPERM:-1}
export RADIANCE_TP_PAD=${RADIANCE_TP_PAD:-0}   # TP=2 here; 3 is the only non-zero value upstream supports
export RADIANCE_TP_PAD_DRAFTER=${RADIANCE_TP_PAD_DRAFTER:-1}
export RADIANCE_TP_PAD_STRICT=${RADIANCE_TP_PAD_STRICT:-1}
export RADIANCE_GDN_MERGE_INPROJ=${RADIANCE_GDN_MERGE_INPROJ:-1}
export RADIANCE_GDN_NORM_QUANT=${RADIANCE_GDN_NORM_QUANT:-1}
export RADIANCE_GDN_STRIDED_GATES=${RADIANCE_GDN_STRIDED_GATES:-0}
export RADIANCE_GDN_EMPTY_OUT=${RADIANCE_GDN_EMPTY_OUT:-0}
export R4D_ATTN_FP8=${R4D_ATTN_FP8:-3}
export RADIANCE_DFLASH_CALIB="${RADIANCE_DFLASH_CALIB:-}"
export RADIANCE_DFLASH_CALIB_TOKENS="${RADIANCE_DFLASH_CALIB_TOKENS:-200000}"
export RADIANCE_MXFP4_SHADOW="${RADIANCE_MXFP4_SHADOW:-}"
export RADIANCE_MXFP4_KERNEL_N="${RADIANCE_MXFP4_KERNEL_N:-}"
export RADIANCE_MXFP4_KERNEL_NK="${RADIANCE_MXFP4_KERNEL_NK:-}"
export RADIANCE_MXFP4_CHECKALL="${RADIANCE_MXFP4_CHECKALL:-}"
export RADIANCE_MXFP4_DECODE_KS="${RADIANCE_MXFP4_DECODE_KS:-}"
export RADIANCE_MXFP4_DECODE_BK="${RADIANCE_MXFP4_DECODE_BK:-}"
export RADIANCE_MXFP4_PERBLOCK_NK="${RADIANCE_MXFP4_PERBLOCK_NK:-}"
export RADIANCE_MXFP4_REFLINEAR="${RADIANCE_MXFP4_REFLINEAR:-0}"
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-}"
export RADIANCE_MXFP4_EPIFAST="${RADIANCE_MXFP4_EPIFAST:-1}"
export RADIANCE_AR_OVERLAP=${RADIANCE_AR_OVERLAP:-0}
export RADIANCE_GDN_FUSED_UPDATE=${RADIANCE_GDN_FUSED_UPDATE:-1}
export RADIANCE_DYNAMIC_WIDTH=${RADIANCE_DYNAMIC_WIDTH:-1}
export RADIANCE_DYNW_ALPHA=${RADIANCE_DYNW_ALPHA:-0.35}
export RADIANCE_DYNW_MARGIN=${RADIANCE_DYNW_MARGIN:-2}
export RADIANCE_DYNW_MIN=${RADIANCE_DYNW_MIN:-2}
export RADIANCE_DYNW_MIN_BATCH=${RADIANCE_DYNW_MIN_BATCH:-3}
export RADIANCE_AR_QNB=${RADIANCE_AR_QNB:-96}
export RADIANCE_AR_QNT=${RADIANCE_AR_QNT:-1024}
export MXFP4_CUMODE=${MXFP4_CUMODE:-0}
export RADIANCE_STEP_TRACE=${RADIANCE_STEP_TRACE:-0}
export RADIANCE_AR_OVERLAP_MIN_M=${RADIANCE_AR_OVERLAP_MIN_M:-2048}
export RADIANCE_AR_OVERLAP_SLICES=${RADIANCE_AR_OVERLAP_SLICES:-4}
export RADIANCE_MXFP4_EPIFAST=${RADIANCE_MXFP4_EPIFAST:-1}
export RADIANCE_MXFP4_R4D_DECODE_MAX_M=${RADIANCE_MXFP4_R4D_DECODE_MAX_M:-0}
export RADIANCE_TOPK_TRITON_MIN_ROWS=${RADIANCE_TOPK_TRITON_MIN_ROWS:-1}
export RADIANCE_SKINNY_GEMM=${RADIANCE_SKINNY_GEMM:-1}
export RADIANCE_MXFP4_HOIST_QUANT=${RADIANCE_MXFP4_HOIST_QUANT:-1}
export RADIANCE_MXFP4_TRACED_QUANT=${RADIANCE_MXFP4_TRACED_QUANT:-1}
export RADIANCE_FP8_STREAM=${RADIANCE_FP8_STREAM:-1}
export RADIANCE_RMS_QUANT_FUSION=${RADIANCE_RMS_QUANT_FUSION:-1}
export RADIANCE_NVFP4_MXFP4=${RADIANCE_NVFP4_MXFP4:-0}
export RADIANCE_NVFP4_EXP=${RADIANCE_NVFP4_EXP:-mse}
export RADIANCE_NVFP4_FP8_LAYERS=${RADIANCE_NVFP4_FP8_LAYERS:-mxfp4}
export RADIANCE_NVFP4_BF16_LAYERS=${RADIANCE_NVFP4_BF16_LAYERS:-in_proj_ba}
export RADIANCE_NVFP4_LMHEAD=${RADIANCE_NVFP4_LMHEAD:-bf16}
export RADIANCE_GDN_PATHS=${RADIANCE_GDN_PATHS:-both}
export RADIANCE_GDN_NANTRACE=${RADIANCE_GDN_NANTRACE:-0}
export RADIANCE_MXFP4_CHECK_MAX_M=${RADIANCE_MXFP4_CHECK_MAX_M:-128}
export RADIANCE_MXFP4_MHIST=${RADIANCE_MXFP4_MHIST:-0}
export VLLM_CACHE_ROOT=/cache/vllm
export TORCHINDUCTOR_CACHE_DIR=/cache/inductor
export TRITON_CACHE_DIR=/cache/triton
export AITER_ROOT_DIR=/cache/aiter
export TRITON_CACHE_AUTOTUNING=1
# The patched libr4d built by setup-mxfp4.sh, bind-mounted at /r4d.
R4D_SO=${R4D_SO:-/r4d}

# The norm+quant fusion turns on vLLM's passes (one --compilation-config
# instance; argparse keeps the last, so entries accumulate into one flag).
CC_ITEMS=""
if [ "${RADIANCE_NORMQUANT_FUSION:-1}" != "0" ]; then
  CC_ITEMS="\"pass_config\":{\"fuse_norm_quant\":true,\"fuse_act_quant\":true}"
fi
EXTRA=${EXTRA:-}
if [ -n "$CC_ITEMS" ]; then
  EXTRA="$EXTRA --compilation-config {$CC_ITEMS}"
fi

# ---------------------------------------------------------------- patch block
# Byte-for-byte the launcher's container-start block. The patches come from
# the /patches mount (the pinned radiance clone). SP must match the image's
# venv layout; the image tag is pinned in the quadlet.
SP=/opt/vllm/lib/python3.12/site-packages
cd /patches
python3 patch_quark_mxfp4.py
python3 patch_nvfp4_mxfp4.py
python3 patch_tp3_pad.py
python3 patch_ar_maxbytes.py
python3 patch_topk_triton_rows.py
python3 patch_dflash_calib.py
python3 patch_dflash_mxfp4_kv.py
python3 patch_rmsquant_fusion.py
python3 patch_verify_head.py
python3 patch_kv_group_size.py
python3 patch_topk_composite.py
python3 patch_gdn_shared_build.py
python3 patch_dflash_selector_topk.py
python3 patch_gdn_merge_inproj.py
python3 patch_dynwidth.py
python3 patch_async_dynwidth.py
python3 patch_step_trace.py
python3 patch_ar_geometry.py
python3 patch_ar_3rank.py
python3 patch_gdn_glue.py
# Non-fatal: fixes content=null on thinking-off requests; not required to serve.
python3 patch_qwen3_thinkoff.py \
  || echo "[radiance] WARNING: thinkoff patch did not apply; thinking-off requests will return empty content"
cp mxfp4-configs/*.json "$SP"/aiter/ops/triton/configs/gemm/
cp radiance_preamble.py /opt/radiance_preamble.py
cp radiance_nvfp4.py radiance_mxfp4.py radiance_gdn.py radiance_rmsquant.py radiance_drafthead.py \
   radiance_verifyhead.py radiance_gdnmerge.py radiance_aroverlap.py radiance_topk.py \
   radiance_arnq.py radiance_tp3pad.py "$SP"/
hipcc -O3 -w -std=c++17 -fPIC -shared --offload-arch=gfx1201 $([ "${MXFP4_CUMODE:-0}" = 1 ] && echo -mcumode) \
  $(python3 -m pybind11 --includes) radiance_mxfp4_fp8.hip -o "$SP"/radiance_mxfp4_fp8.so
if [ -n "${R4D_SO:-}" ] && [ -f /r4d/r4d.so ]; then
  cp /r4d/r4d.so "$SP"/r4d.so
  echo "[radiance] using patched r4d.so from $R4D_SO"
fi
# Leave /patches before exec: a stale radiance_mxfp4_fp8.so there would shadow
# the one just compiled into site-packages (cwd precedes it on sys.path).
cd /
exec /opt/radiance_entrypoint.sh "$SNAP" --served-model-name $SERVED_NAMES \
  --host 0.0.0.0 --port "$PORT" \
  --kv-cache-dtype fp8 --tensor-parallel-size "$TP" \
  --gpu-memory-utilization "$GPU_UTIL" \
  ${KV_MEM:+--kv-cache-memory "$KV_MEM"} \
  --max-model-len "$MAXLEN" --max-num-seqs "$MAXSEQS" --max-num-batched-tokens "$CHUNK" \
  --attention-backend "$ATTN" \
  --speculative-config "$SPEC_CFG" \
  $ASYNC_FLAG $EXTRA \
  --enable-prefix-caching --mamba-cache-mode align --enable-auto-tool-choice \
  --tool-call-parser qwen3_coder --reasoning-parser qwen3 \
  --override-generation-config '{"temperature":0.7,"top_p":0.95,"top_k":20}' \
  --chat-template "$CHAT_TEMPLATE"