43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
from sglang.test.kl_test_utils import (
|
|
test_input_output_logprobs_match_decode_cache_hit_helper,
|
|
test_input_output_logprobs_match_prefill_cache_hit_helper,
|
|
)
|
|
|
|
|
|
class KLDivergenceMixin:
|
|
kl_div_thres: float
|
|
kl_div_thres_decode: float | None = None
|
|
kl_div_thres_prefill: float | None = None
|
|
kl_div_max_samples: int = 32
|
|
kl_div_prefill_max_new_tokens: int = 512
|
|
kl_div_decode_max_new_tokens: int = 512
|
|
|
|
@classmethod
|
|
def _build_acc_thresholds(cls, threshold):
|
|
"""Build an ACC_THRESHOLDS dict compatible with kl_test_utils."""
|
|
return {cls.model: {"kl_div": threshold}}
|
|
|
|
@classmethod
|
|
def test_input_output_logprobs_match_prefill_cache_hit(cls):
|
|
test_input_output_logprobs_match_prefill_cache_hit_helper(
|
|
base_url=cls.base_url,
|
|
ACC_THRESHOLDS=cls._build_acc_thresholds(
|
|
cls.kl_div_thres_prefill or cls.kl_div_thres
|
|
),
|
|
model_name=cls.model,
|
|
max_samples=cls.kl_div_max_samples,
|
|
max_new_tokens=cls.kl_div_prefill_max_new_tokens,
|
|
)
|
|
|
|
@classmethod
|
|
def test_input_output_logprobs_match_decode_cache_hit(cls):
|
|
test_input_output_logprobs_match_decode_cache_hit_helper(
|
|
base_url=cls.base_url,
|
|
ACC_THRESHOLDS=cls._build_acc_thresholds(
|
|
cls.kl_div_thres_decode or cls.kl_div_thres
|
|
),
|
|
model_name=cls.model,
|
|
max_samples=cls.kl_div_max_samples,
|
|
max_new_tokens=cls.kl_div_decode_max_new_tokens,
|
|
)
|