diff --git a/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py b/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py index 0ac3c5f62..649cddf03 100644 --- a/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py +++ b/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py @@ -68,11 +68,10 @@ class TestQwen3NextMTPTopk( ] -# TODO(hzh): After merging the PR that fixes specv2 to correctly return log probs, -# add KLDivergenceMixin back. https://github.com/sgl-project/sglang/pull/18645 -class TestQwen3NextMTPV2(GSM8KMixin, DefaultServerBase): +class TestQwen3NextMTPV2(GSM8KMixin, KLDivergenceMixin, DefaultServerBase): model = QWEN3_NEXT_MODEL gsm8k_accuracy_thres = 0.93 + kl_div_thres = 0.0025 other_args = [ "--trust-remote-code", "--speculative-algorithm",