From fb72a37fdefa80ac00b23096130d5f815bbe2136 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sun, 9 Aug 2026 22:22:27 +0800 Subject: [PATCH] Add deterministic logprob-consistency test for inkling-small nvfp4 (#34168) --- .../models_e2e/test_inkling_small_nvfp4.py | 102 +++++++++++++++++- 1 file changed, 98 insertions(+), 4 deletions(-) diff --git a/test/registered/models_e2e/test_inkling_small_nvfp4.py b/test/registered/models_e2e/test_inkling_small_nvfp4.py index 1d00d5cb7..e11f0f550 100644 --- a/test/registered/models_e2e/test_inkling_small_nvfp4.py +++ b/test/registered/models_e2e/test_inkling_small_nvfp4.py @@ -18,6 +18,17 @@ from urllib.parse import urlparse from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci + +# Aliased so pytest does not collect the imported `test_`-prefixed helpers as tests. +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_decode_cache_hit_helper as assert_logprobs_match_decode_cache_hit, +) +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_helper as assert_logprobs_match, +) +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_prefill_cache_hit_helper as assert_logprobs_match_prefill_cache_hit, +) from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, @@ -25,16 +36,27 @@ from sglang.test.test_utils import ( popen_launch_server, ) -register_cuda_ci(est_time=600, stage="extra-b", runner_config="4-gpu-b200") +register_cuda_ci(est_time=1200, stage="extra-b", runner_config="4-gpu-b200") _MODEL_PATH = os.environ.get( "INKLING_SMALL_TEST_MODEL_PATH", "thinkingmachines/Inkling-Small-NVFP4" ) # Measured 0.900 (10-shot, 200 questions, tp=4, invalid=0.000) -- completion, -# so no thinking. The floor sits ~2.5 sigma of the 200-question sampling noise +# so no thinking. The floor sits ~4.5 sigma of the 200-question sampling noise # below that: a real accuracy collapse trips it, the sample spread does not. -GSM8K_THRESHOLD = 0.85 +GSM8K_THRESHOLD = 0.80 + +# All three helpers measure exactly 0 on this config -- every logprob matches +# bit for bit. The floor is only there to keep a stray ulp from failing the +# run; the divergence a state-reuse bug produces lands orders of magnitude +# above it. +KL_DIV_THRESHOLD = 1e-9 + +# Past the 512-token sliding window, so decode carries the window through the +# handover from prompt tokens to generated ones -- where a stale conv/mamba +# checkpoint or a mis-restored prefix would surface. +KL_MAX_NEW_TOKENS = 1024 class TestInklingSmallNvfp4(CustomTestCase): @@ -89,7 +111,7 @@ class TestInklingSmallNvfp4(CustomTestCase): num_shots=10, data_path=None, num_questions=200, - max_new_tokens=16000, + max_new_tokens=512, parallel=128, host=f"http://{url.hostname}", port=int(url.port), @@ -99,5 +121,77 @@ class TestInklingSmallNvfp4(CustomTestCase): self.assertGreaterEqual(metrics["accuracy"], GSM8K_THRESHOLD) +class TestInklingSmallNvfp4Deterministic(CustomTestCase): + """Prefill and decode must score a token identically once every kernel on + the path is batch-invariant, which is what deterministic inference buys. + Drift here is then a state-reuse bug -- a stale conv/mamba checkpoint, or a + prefix restored from the radix cache that does not reproduce a fresh + prefill -- rather than the float noise a loose threshold would hide. + + Runs its own server: the accuracy case above has to stay on the production + numerics, so it cannot share this one. + """ + + @classmethod + def setUpClass(cls): + cls.model = _MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp", + "4", + "--trust-remote-code", + "--quantization", + "modelopt_fp4", + "--attention-backend", + "fa4", + "--page-size", + "128", + "--fp4-gemm-backend", + "flashinfer_trtllm", + "--moe-runner-backend", + "flashinfer_trtllm_routed", + "--mamba-radix-cache-strategy", + "extra_buffer", + "--swa-full-tokens-ratio", + "0.1", + "--mamba-full-memory-ratio", + "0.1", + "--mem-fraction-static", + "0.85", + "--enable-deterministic-inference", + "--disable-prefill-cuda-graph", + ], + env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + + @classmethod + def tearDownClass(cls): + if getattr(cls, "process", None) is not None: + kill_process_tree(cls.process.pid) + + def _run(self, helper): + helper( + self.base_url, + {self.model: {"kl_div": KL_DIV_THRESHOLD}}, + self.model, + max_samples=32, + max_new_tokens=KL_MAX_NEW_TOKENS, + trust_remote_code=True, + ) + + def test_input_output_logprobs_match(self): + self._run(assert_logprobs_match) + + def test_input_output_logprobs_match_prefill_cache_hit(self): + self._run(assert_logprobs_match_prefill_cache_hit) + + def test_input_output_logprobs_match_decode_cache_hit(self): + self._run(assert_logprobs_match_decode_cache_hit) + + if __name__ == "__main__": unittest.main()