[Auto Sync] Update test_deterministic.py (20251112) (#13128)

Co-authored-by: Stefan He <hebiaobuaa@gmail.com>

python/sglang/test/test_deterministic.py

@@ -325,32 +325,66 @@ class TokenIdsAndLogprobs:
 
     @classmethod
     def compare(cls, a: "TokenIdsAndLogprobs", b: "TokenIdsAndLogprobs"):
+        import numpy as np
+
         assert len(a.token_ids) == len(b.token_ids)
         token_match = a.token_ids == b.token_ids
         logprobs_match = a.logprobs == b.logprobs
 
         if token_match:
-            print(f"Token match: {a.token_ids}")
+            print(f"✅ Token match")
         else:
-            print(f"❗Token mismatch: {a.token_ids=} {b.token_ids=}")
+            print(f"❌ Token mismatch: {a.token_ids=} {b.token_ids=}")
 
         if logprobs_match:
-            print(f"Logprobs match:", a.logprobs)
+            print(f"✅ Logprobs match:", a.logprobs[:5])
         else:
-            print(f"❗Logprobs mismatch")
+            print(f"❌ Logprobs mismatch")
+            # Only print first 5 elements for readability
+            n_show = 5
+            a_show = a.logprobs[:n_show]
+            b_show = b.logprobs[:n_show]
             print(
                 "    A:   ",
-                [f"{x:.10f}" if x is not None else "None" for x in a.logprobs],
+                [f"{x:.10f}" if x is not None else "None" for x in a_show],
+                f"... ({len(a.logprobs)} total)" if len(a.logprobs) > n_show else "",
             )
             print(
                 "    B:   ",
-                [f"{x:.10f}" if x is not None else "None" for x in b.logprobs],
+                [f"{x:.10f}" if x is not None else "None" for x in b_show],
+                f"... ({len(b.logprobs)} total)" if len(b.logprobs) > n_show else "",
             )
             diff = [
                 abs(x - y) if x is not None else float("nan")
                 for x, y in zip(a.logprobs, b.logprobs)
             ]
-            print("    Diff:", [f"{x:.10e}" for x in diff])
+            print(
+                "    Diff:",
+                [f"{x:.10e}" for x in diff[:n_show]],
+                f"... ({len(diff)} total)" if len(diff) > n_show else "",
+            )
+
+            # Compute KL-divergence using K3 approximation
+            # KL(P||Q) ≈ (exp(log(P) - log(Q)) - 1) - (log(P) - log(Q))
+            # This is based on selected token logprobs only
+            valid_pairs = [
+                (lp_a, lp_b)
+                for lp_a, lp_b in zip(a.logprobs, b.logprobs)
+                if lp_a is not None and lp_b is not None
+            ]
+            if valid_pairs and token_match:
+                logprobs_a = np.array([lp for lp, _ in valid_pairs])
+                logprobs_b = np.array([lp for _, lp in valid_pairs])
+
+                # K3 approximation: KL(A||B) ≈ (exp(logr) - 1) - logr, where logr = log_a - log_b
+                logr = logprobs_a - logprobs_b
+                kl_per_token = (np.exp(logr) - 1) - logr
+                kl_mean = np.mean(kl_per_token)
+                kl_max = np.max(kl_per_token)
+
+                print(f"    KL(A||B) mean: {kl_mean:.10e}")
+                print(f"    KL(A||B) max : {kl_max:.10e}")
+                print(f"    Mean absolute logprob diff: {np.mean(np.abs(logr)):.10e}")
 
         return token_match and logprobs_match