[ATOM SGL] GLM 5.2 MTP

[ZhiweiYan-96]

Motivation

Enable glm5.2 MTP support.

Test Plan

export AITER_QUICK_REDUCE_QUANTIZATION=INT4
export AITER_USE_FLYDSL_MOE_SORTING=1
export SGLANG_USE_AITER=1
export SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
export PYTHONPATH=/home/qichu_qle/zhiwei/dsv4/atom:/app/sglang

MODEL_PATH=/workspace/shared/data/amd_int/models/GLM-5.2-MXFP4
TP=4
PORT=8015
MODEL_LOADER_EXTRA_CONFIG='{"online_quant_config":{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","*.mlp.gate","*expert*"]}}'

export CUDA_VISIBLE_DEVICES=4,5,6,7
TORCHINDUCTOR_COMPILE_THREADS=128 \
python3 -m sglang.launch_server \
    --model-path "${MODEL_PATH}" \
    --host localhost \
    --port "${PORT}" \
    --trust-remote-code \
    --tp-size "${TP}" \
    --mem-fraction-static 0.8 \
    --disable-radix-cache \
    --kv-cache-dtype fp8_e4m3 \
    --model-loader-extra-config "${MODEL_LOADER_EXTRA_CONFIG}" \
    --speculative-algorithm EAGLE \
    --speculative-num-steps 3 \
    --speculative-eagle-topk 1 \
    --speculative-num-draft-tokens 4 \
    2>&1 | tee glm-server-mxfp4-tp4-mtp.log

gsm8k script, --gen_kwargs '{"chat_template_kwargs":{"enable_thinking":false}}' is added, avoiding too long thinking.

#!/usr/bin/env bash
set -x

addr=localhost
port=8015
url=http://${addr}:${port}/v1/chat/completions
# model=/workspace/shared/data/amd_int/models/MiniMax-M3-MXFP8
model=/workspace/shared/data/amd_int/models/GLM-5.2-MXFP4
task=gsm8k
limit=${1:-}
limit_args=()
if [[ -n "${limit}" ]]; then
        limit_args+=(--limit "${limit}")
fi

lm_eval --model local-chat-completions \
        --model_args model=${model},base_url=${url},num_concurrent=32,max_retries=1,max_gen_toks=16384 \
        --tasks ${task} \
        --num_fewshot 5 \
        --batch_size 16 \
        --apply_chat_template \
        --fewshot_as_multiturn \
        --gen_kwargs '{"chat_template_kwargs":{"enable_thinking":false}}' \
        "${limit_args[@]}" \
        2>&1 | tee log.m3_mxfp8_sglang.gsm8k.log

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.