diff --git a/python/sglang/test/kits/fwd_occupancy_kit.py b/python/sglang/test/kits/fwd_occupancy_kit.py index 6385a61e3..1743c001e 100644 --- a/python/sglang/test/kits/fwd_occupancy_kit.py +++ b/python/sglang/test/kits/fwd_occupancy_kit.py @@ -95,7 +95,7 @@ class FwdOccupancyMixin: ) def _fwd_occupancy_fire(self, prompt: str, max_new_tokens: int): - """Fire one /generate, return (completion_tokens, wall_time). + """Fire one /generate, return (meta_info, wall_time). Must not be called concurrently -- that would break the single-batch invariant.""" t0 = time.perf_counter() @@ -115,13 +115,12 @@ class FwdOccupancyMixin: except requests.RequestException: # Final stats-vs-threshold is the signal; individual fire # failure isn't. - return 0, 0.0 + return {}, 0.0 elapsed = time.perf_counter() - t0 try: - tokens = resp.json().get("meta_info", {}).get("completion_tokens", 0) + return resp.json().get("meta_info", {}), elapsed except ValueError: # non-JSON body - tokens = 0 - return tokens, elapsed + return {}, elapsed def _fwd_occupancy_warmup(self): """Fill cuda graphs + step the device-timer past its first NaN @@ -134,19 +133,16 @@ class FwdOccupancyMixin: def _fwd_occupancy_measure(self): """Background-fire one long single-batch request, scrape - /metrics on the foreground; return (non-NaN samples, - token_tps).""" + /metrics on the foreground; return (non-NaN samples, perf).""" samples = [] request_done = threading.Event() - result = {"completion_tokens": 0, "elapsed": 0.0} + result = {"meta_info": {}, "elapsed": 0.0} def fire_one(): try: - result["completion_tokens"], result["elapsed"] = ( - self._fwd_occupancy_fire( - self.fwd_occupancy_prompt, - self.fwd_occupancy_max_new_tokens, - ) + result["meta_info"], result["elapsed"] = self._fwd_occupancy_fire( + self.fwd_occupancy_prompt, + self.fwd_occupancy_max_new_tokens, ) finally: request_done.set() @@ -161,56 +157,28 @@ class FwdOccupancyMixin: time.sleep(self.fwd_occupancy_scrape_interval) firer.join(timeout=_GENERATE_REQUEST_TIMEOUT) - token_tps = ( - result["completion_tokens"] / result["elapsed"] - if result["elapsed"] > 0 - else 0.0 - ) - return samples, token_tps + return samples, self._fwd_occupancy_perf(result) + + def _fwd_occupancy_perf(self, result): + """Aggregate per-request perf metrics from the fire result + (input/output tokens, decode tps, mean inter-token latency, + wall-clock tps).""" + meta = result["meta_info"] or {} + elapsed = result["elapsed"] + out = meta.get("completion_tokens", 0) or 0 + decode_tps = meta.get("decode_throughput", 0.0) or 0.0 + return { + "input_tokens": meta.get("prompt_tokens", 0) or 0, + "output_tokens": out, + "decode_tps": decode_tps, + "mean_itl_ms": (1000.0 / decode_tps) if decode_tps > 0 else 0.0, + "wall_tps": (out / elapsed) if elapsed > 0 else 0.0, + } def test_fwd_occupancy(self): self._assert_metrics_device_timer_enabled() self._fwd_occupancy_warmup() - samples, token_tps = self._fwd_occupancy_measure() - - self.assertGreaterEqual( - len(samples), - self.fwd_occupancy_min_samples, - f"only {len(samples)} non-NaN occupancy samples collected " - f"(need >= {self.fwd_occupancy_min_samples}); the measurement " - "window may be too short or the gauge stuck at NaN", - ) - - # Median is the steady-state signal; peak / p10 included in the - # assertion message for triage. - samples_sorted = sorted(samples) - median = statistics.median(samples_sorted) - peak = samples_sorted[-1] - p10_idx = min(len(samples_sorted) - 1, max(0, len(samples_sorted) // 10)) - p10 = samples_sorted[p10_idx] - print( - "\n" - + tabulate.tabulate( - [ - ["samples (n)", len(samples)], - ["median", f"{median:.2f}"], - ["peak", f"{peak:.2f}"], - ["p10", f"{p10:.2f}"], - ["threshold", f"{self.fwd_occupancy_threshold:.2f}"], - ["token tps", f"{token_tps:.2f}"], - ], - headers=["fwd_occupancy", "value"], - tablefmt="github", - ) - ) - - self.assertGreater( - median, - self.fwd_occupancy_threshold, - f"sglang:fwd_occupancy median={median:.2f} did not exceed " - f"threshold {self.fwd_occupancy_threshold} " - f"(peak={peak:.2f}, p10={p10:.2f}, n={len(samples)})", - ) + samples, perf = self._fwd_occupancy_measure() # The 2048-token decode above populates the spec running average # if a spec algorithm is enabled; absent otherwise (vanilla @@ -222,8 +190,65 @@ class FwdOccupancyMixin: avg_accept = info["internal_states"][0].get("avg_spec_accept_length") except (requests.RequestException, KeyError, IndexError): avg_accept = None + + # Median is the steady-state signal; peak / p10 included in the + # assertion message for triage. Both tables print before any + # assertion so the numbers surface even on assertion failure. + samples_sorted = sorted(samples) + if samples_sorted: + median = statistics.median(samples_sorted) + peak = samples_sorted[-1] + p10_idx = min(len(samples_sorted) - 1, max(0, len(samples_sorted) // 10)) + p10 = samples_sorted[p10_idx] + else: + median = peak = p10 = float("nan") + + perf_rows = [ + ["input tokens", perf["input_tokens"]], + ["output tokens", perf["output_tokens"]], + ["decode tps", f"{perf['decode_tps']:.2f}"], + ["mean itl (ms)", f"{perf['mean_itl_ms']:.2f}"], + ["wall tps", f"{perf['wall_tps']:.2f}"], + ] + if avg_accept is not None: + perf_rows.append(["avg spec accept", f"{avg_accept:.3f}"]) + print( + "\n" + + tabulate.tabulate( + perf_rows, headers=["perf metric", "value"], tablefmt="github" + ) + ) + + print( + "\n\n" + + tabulate.tabulate( + [ + ["samples (n)", len(samples)], + ["median", f"{median:.2f}"], + ["peak", f"{peak:.2f}"], + ["p10", f"{p10:.2f}"], + ["threshold", f"{self.fwd_occupancy_threshold:.2f}"], + ], + headers=["fwd_occupancy", "value"], + tablefmt="github", + ) + ) + + self.assertGreaterEqual( + len(samples), + self.fwd_occupancy_min_samples, + f"only {len(samples)} non-NaN occupancy samples collected " + f"(need >= {self.fwd_occupancy_min_samples}); the measurement " + "window may be too short or the gauge stuck at NaN", + ) + self.assertGreater( + median, + self.fwd_occupancy_threshold, + f"sglang:fwd_occupancy median={median:.2f} did not exceed " + f"threshold {self.fwd_occupancy_threshold} " + f"(peak={peak:.2f}, p10={p10:.2f}, n={len(samples)})", + ) if avg_accept is not None: - print(f"avg_spec_accept_length = {avg_accept:.3f}") self.assertGreater( avg_accept, self.fwd_occupancy_acc_length_threshold,