from quark.torch import LLMTemplate, ModelQuantizer

qwen3_5_moe_text_template = LLMTemplate(
    model_type="qwen3_5_moe_text",
    kv_layers_name=["*k_proj", "*v_proj"],
    q_layer_name="*q_proj"
)
LLMTemplate.register_template(qwen3_5_moe_text_template)

# === NOTE: Change the following source and output dir as needed ===
src_ckpt_path = "Qwen/Qwen3.8-2.4T-A95B-FP8"
output_dir = "amd/Qwen3.8-2.4T-A95B-Quark-MXFP4"

quant_scheme = "mxfp4"
exclude_layers = [
    "lm_head",
    "*self_attn*",
    "*linear_attn*",
    "*mlp.gate",
    "*shared_expert*",
    "mtp.fc",
]
template = LLMTemplate.get("qwen3_5_moe_text")

quant_config = template.get_config(scheme=quant_scheme, exclude_layers=exclude_layers)

# Quantize with File2File mode to avoid loading full model
quantizer = ModelQuantizer(quant_config)
quantizer.direct_quantize_checkpoint(
    pretrained_model_path=src_ckpt_path,
    save_path=output_dir,
)

 python3 -m sglang.launch_server \
    --model-path amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 \
    --served-model-name amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 \
    --tp-size 8 \
    --attention-backend aiter \
    --page-size 1 \
    --kv-cache-dtype auto \
    --chunked-prefill-size 16384 \
    --watchdog-timeout 1200 \
    --mem-fraction-static 0.9 \
    --model-loader-extra-config '{"enable_multithread_load": true}' \
    --trust-remote-code \
    --host 0.0.0.0 \
    --port 9001

python3 -m sglang.test.run_eval \
    --port 9001 \
    --model amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 \
    --eval-name gsm8k \
    --num-examples 1319 \
    --num-threads 512 \
    --max-tokens 2048 \
    --chat-template-kwargs '{"enable_thinking": false}'