Feat/swa suite

.gitignore

@@ -52,6 +52,7 @@ samples.jsonl.ipynb_checkpoints/
 accuracy_outputs.jsonl
 run.log
 *.log
+model.cache
 my_submission/
 mini_result/
 *_backup/
@@ -63,3 +64,4 @@ backup/
 
 # ── Local-only design QA screenshots ────────────────────────────────────────
 .screenshots/
+workspace/nips26_rebuttal/

configs/runner_configs/runner_nvidia_tensorrt_llm_21323970.yaml.example

@@ -0,0 +1,32 @@
+# TensorRT-LLM runner configuration — Example
+# Copy to: configs/runner_configs/runner_nvidia_tensorrt_llm_<hash8>.yaml
+#
+# All fields are optional — defaults are used when fields are absent.
+
+# ── Scheduler ──────────────────────────────────────────────────────────────
+
+# Maximum batch size (concurrent requests). TRT-LLM packs sequences into
+# micro-batches dynamically. For 24 GB GPUs (RTX 3090/4090), reduce to 128.
+# For 80 GB GPUs, 512 is a safe default. Increase to 1024+ for H200.
+max_batch_size: 512
+
+# Maximum tokens in a single forward pass iteration. Higher values improve
+# GPU utilization but increase memory pressure. For long-context workloads
+# (Suite D, ~28K tokens), consider raising to 32768.
+max_num_tokens: 8192
+
+# ── Engine ─────────────────────────────────────────────────────────────────
+
+# Additional keyword arguments passed directly to tensorrt_llm.LLM() and
+# AsyncLLM(). Refer to the TRT-LLM docs for the full list:
+#   https://nvidia.github.io/TensorRT-LLM/
+# Common options:
+#   backend: "pytorch"       — (default) torch.export graph capture
+#   enable_chunked_prefill   — improves long-context prefill throughput
+#   enable_context_fmha_fusion — fuse QKV projection with FlashAttention
+#   free_gpu_memory_fraction  — fraction of free GPU memory to use for KV cache
+engine_kwargs: {}
+  # Example:
+  # backend: "pytorch"
+  # enable_chunked_prefill: true
+  # free_gpu_memory_fraction: 0.90

datasets/sharegpt_swa_v1/README.md

@@ -0,0 +1,17 @@
+# ShareGPT SWA Dataset v1
+
+Long-context benchmark dataset for Suite H (Sliding Window Attention).
+
+## Source
+Prompts from `sharegpt_longctx_v1`, truncated to the last 10,000 tokens of each conversation. This ensures prompts exceed the sliding window (4096 tokens for Mistral-7B) while staying within the model's maximum context (32,768 tokens).
+
+## Purpose
+Suite H uses Mistral-7B-Instruct-v0.1 (`sliding_window=4096`) to test how Sliding Window Attention changes a model's position on the arithmetic intensity spectrum. Prompts must be longer than the sliding window to trigger SWA-specific behavior in the attention computation and KV cache management.
+
+## Statistics
+- 100 prompts
+- Input tokens: 10,000 (all prompts truncated uniformly)
+- Source: sharegpt_longctx_v1 (original prompts 24K-33K tokens)
+
+## Framework Compatibility Note
+Whether the inference framework honors `sliding_window` for KV cache allocation is framework- and version-dependent. This dataset provides the prompts needed to test SWA; the framework's SWA support determines whether the benchmark measures SWA-specific effects or dense-equivalent behavior.

leaderboard/generate.py

@@ -875,6 +875,9 @@ def extract_row(result: dict) -> dict:
         "min_price_usd_per_hr":               min_price,
         "cost_efficiency_toks_per_dollar_hr": cost_efficiency,
         "tokens_per_watt":                    derived.get("tokens_per_sec_per_watt"),
+        "tokens_per_joule":                  derived.get("tokens_per_joule"),
+        "energy_joules":                     derived.get("energy_joules"),
+        # Metadata
         "accuracy_valid":   accuracy.get("valid"),
         "accuracy_score":   accuracy.get("subset_score"),
         "date":             meta.get("date"),

loadgen/loadgen.py

@@ -33,6 +33,7 @@ def tqdm(iterable, **kwargs):
         return iterable
 
 from .types import InferenceResult, SampleRecord
+from .power import PowerSampler
 
 # Import InferenceRequest for type hints — runners pass list[InferenceRequest]
 # to inference_fn_offline and inference_fn_streaming
@@ -319,6 +320,10 @@ async def run_sustained(
         if not requests:
             raise ValueError("run_sustained requires requests to be loaded.")
 
+        # Start power sampling for the full sustained window
+        power_sampler = PowerSampler()
+        power_sampler.start()
+
         samples           = []
         start_time        = _time.perf_counter()
         next_sample_at    = start_time + interval_seconds
@@ -421,6 +426,9 @@ def _fire_request() -> None:
         if pending_tasks:
             await asyncio.wait(pending_tasks, timeout=60)
 
+        # Stop power sampling after the full sustained window
+        power_stats = power_sampler.stop()
+
         # ── Derived metrics ───────────────────────────────────────────────────
         # Exclude warmup samples from scalar metrics
         analysis_samples = [s for s in samples if not s.get("is_warmup")]
@@ -474,6 +482,8 @@ def _fire_request() -> None:
                 "throttle_onset_minute":               throttle_onset_minute,
                 "ttft_p99_drift_ms":                   ttft_p99_drift_ms,
                 "throughput_post_warmup_reliability":  throughput_cv_block,
+                "power_watts_avg":                     power_stats.avg,
+                "power_watts_peak":                    power_stats.peak,
             }
         }
 
@@ -529,6 +539,8 @@ def _run_offline(self, inference_fn: Callable) -> dict:
                 run_total_throughputs = []
                 run_elapsed_times = []
                 run_samples: list[SampleRecord] = []
+                run_power_avgs: list[float] = []
+                run_power_peaks: list[float] = []
 
                 for run_idx in range(total_runs):
                     is_warmup = run_idx < self.warmup_runs
@@ -538,6 +550,10 @@ def _run_offline(self, inference_fn: Callable) -> dict:
                     desc = f"  client_concurrency={client_concurrency} ({cc_idx+1}/{total_concurrency_levels}) {run_label}"
                     tqdm.write(f"{desc} — sending all {len(self.requests)} requests...")
 
+                    # Start power sampling before the timed window
+                    power_sampler = PowerSampler()
+                    power_sampler.start()
+
                     t_start = time.perf_counter()
 
                     # Send ALL requests at once — engine handles internal batching
@@ -558,6 +574,9 @@ def _run_offline(self, inference_fn: Callable) -> dict:
                     t_end = time.perf_counter()
                     elapsed = t_end - t_start
 
+                    # Stop power sampling after the timed window
+                    power_stats = power_sampler.stop()
+
                     if oom_occurred:
                         results_by_concurrency.append({
                             "client_concurrency": client_concurrency,
@@ -595,6 +614,8 @@ def _run_offline(self, inference_fn: Callable) -> dict:
                     run_throughputs.append(throughput_output_only)
                     run_total_throughputs.append(throughput_total)
                     run_elapsed_times.append(elapsed)
+                    run_power_avgs.append(power_stats.avg)
+                    run_power_peaks.append(power_stats.peak)
 
                     per_chip_str = ""
                     if self.chip_count > 1:
@@ -631,15 +652,22 @@ def _run_offline(self, inference_fn: Callable) -> dict:
                     f"median={median_throughput:.0f} tok/s{per_chip_str}\n"
                 )
 
+                # Aggregate power across non-warmup runs:
+                # avg = mean of per-run avgs, peak = max of per-run peaks
+                _valid_avgs = [v for v in run_power_avgs if v is not None]
+                _valid_peaks = [v for v in run_power_peaks if v is not None]
+                _agg_power_avg = round(sum(_valid_avgs) / len(_valid_avgs), 2) if _valid_avgs else None
+                _agg_power_peak = round(max(_valid_peaks), 2) if _valid_peaks else None
+
                 results_by_concurrency.append({
                     "client_concurrency": client_concurrency,
                     "throughput_tokens_per_sec": round(median_throughput, 2),
                     "throughput_tokens_per_sec_per_chip": round(median_throughput / self.chip_count, 2),
                     "throughput_tokens_per_sec_total": round(median_throughput_total, 2),
                     "elapsed_seconds_median": round(float(np.median(run_elapsed_times)), 1),
                     "peak_memory_gb": None,
-                    "power_watts_avg": None,
-                    "power_watts_peak": None,
+                    "power_watts_avg": _agg_power_avg,
+                    "power_watts_peak": _agg_power_peak,
                     "oom": False,
                     # Per-run throughput reliability: lets the UI show "stable ✓ /
                     # noisy ⚠ / unstable ✗" without forcing the user to download
@@ -731,6 +759,8 @@ async def _run_online_async(self, async_inference_fn) -> dict:
             run_ttfts: list[list[float]] = []
             run_tpots: list[list[float]] = []
             run_elapsed_times: list[float] = []
+            run_power_avgs: list[float] = []
+            run_power_peaks: list[float] = []
 
             for run_idx in range(self.suite["num_runs"]):
                 run_label = f"run {run_idx + 1}/{self.suite['num_runs']}"
@@ -740,6 +770,10 @@ async def _run_online_async(self, async_inference_fn) -> dict:
                 inter_arrivals = [self._rng.expovariate(target_qps) for _ in range(n)]
                 arrival_times = list(itertools.accumulate(inter_arrivals))
 
+                # Start power sampling before the timed window
+                power_sampler = PowerSampler()
+                power_sampler.start()
+
                 t_start = loop.time()
 
                 async def send_request(req: InferenceRequest, t_arrival: float) -> InferenceResult:
@@ -763,6 +797,11 @@ async def send_request(req: InferenceRequest, t_arrival: float) -> InferenceResu
                 t_run_end = loop.time()
                 run_elapsed_times.append(t_run_end - t_start)
 
+                # Stop power sampling after the timed window
+                power_stats = power_sampler.stop()
+                run_power_avgs.append(power_stats.avg)
+                run_power_peaks.append(power_stats.peak)
+
                 ttfts: list[float] = []
                 tpots: list[float] = []
                 for r in all_results:
@@ -799,6 +838,12 @@ async def send_request(req: InferenceRequest, t_arrival: float) -> InferenceResu
             if sla_met:
                 max_valid_qps = target_qps
 
+            # Aggregate power across runs for this QPS level
+            _valid_avgs = [v for v in run_power_avgs if v is not None]
+            _valid_peaks = [v for v in run_power_peaks if v is not None]
+            _agg_power_avg = round(sum(_valid_avgs) / len(_valid_avgs), 2) if _valid_avgs else None
+            _agg_power_peak = round(max(_valid_peaks), 2) if _valid_peaks else None
+
             sla_icon = "✓" if sla_met else "✗"
             chip_str = f"  ({self.chip_count} chips)" if self.chip_count > 1 else ""
             tqdm.write(f"  [online] qps={target_qps} TTFT_p99={ttft_p99:.0f}ms SLA={sla_ms}ms {sla_icon}{chip_str}")
@@ -814,6 +859,8 @@ async def send_request(req: InferenceRequest, t_arrival: float) -> InferenceResu
                 "tpot_ms_p99": round(tpot_p99, 2),
                 "elapsed_seconds_median": round(float(np.median(run_elapsed_times)), 1),
                 "sla_met": sla_met,
+                "power_watts_avg": _agg_power_avg,
+                "power_watts_peak": _agg_power_peak,
                 "ttft_ms_p99_reliability":
                     _reliability_block(ttft_p99_per_run, decimals=2),
             })
@@ -917,8 +964,14 @@ async def send(req, t_arrival):
         # recovery_time_seconds in a single post-processing pass after
         # all cycles complete.
         cycle_data: list[dict] = []
+        cycle_power_avgs: list[float] = []
+        cycle_power_peaks: list[float] = []
 
         for cycle_idx in range(num_runs):
+            # Start power sampling for the full cycle (steady + burst)
+            power_sampler = PowerSampler()
+            power_sampler.start()
+
             tqdm.write(f"[burst] cycle {cycle_idx + 1}/{num_runs} — steady({steady_qps} qps)...")
 
             steady_results, steady_elapsed, steady_arrivals = await fire_window(
@@ -943,6 +996,11 @@ async def send(req, t_arrival):
             ]
             burst_ttfts = [v for _, v in burst_ttfts_pairs]
 
+            # Stop power sampling after the full cycle
+            power_stats = power_sampler.stop()
+            cycle_power_avgs.append(power_stats.avg)
+            cycle_power_peaks.append(power_stats.peak)
+
             all_steady_ttfts.extend(steady_ttfts)
             all_burst_ttfts.extend(burst_ttfts)
 
@@ -1021,6 +1079,12 @@ async def send(req, t_arrival):
             f"  [burst] insufficient data"
         )
 
+        # Aggregate power across all cycles
+        _valid_avgs = [v for v in cycle_power_avgs if v is not None]
+        _valid_peaks = [v for v in cycle_power_peaks if v is not None]
+        _agg_power_avg = round(sum(_valid_avgs) / len(_valid_avgs), 2) if _valid_avgs else None
+        _agg_power_peak = round(max(_valid_peaks), 2) if _valid_peaks else None
+
         self._write_samples(all_samples)
         return {"burst": {
             "sla_ttft_ms": sla_ms,
@@ -1040,6 +1104,8 @@ async def send(req, t_arrival):
             "recovery_time_seconds_per_cycle": [
                 round(v, 2) for v in cycle_recovery_times
             ] if cycle_recovery_times else [],
+            "power_watts_avg": _agg_power_avg,
+            "power_watts_peak": _agg_power_peak,
             "_recovery_definition": (
                 "Median seconds within the post-burst steady window before "
                 "rolling TTFT p99 drops below 1.5x the long-term steady baseline. "
@@ -1075,6 +1141,10 @@ async def _run_interactive_async(self, async_inference_fn) -> dict:
         run_elapsed_times: list[float] = []
         ttft_p99_per_run: list[float] = []
 
+        # Start power sampling across the full interactive window
+        power_sampler = PowerSampler()
+        power_sampler.start()
+
         total_runs = self.warmup_runs + self.suite["num_runs"]
 
         for run_idx in range(total_runs):
@@ -1135,6 +1205,9 @@ async def _run_interactive_async(self, async_inference_fn) -> dict:
                     f"({run_elapsed:.0f}s)"
                 )
 
+        # Stop power sampling after all interactive runs complete
+        power_stats = power_sampler.stop()
+
         sampled = self._rng.sample(all_samples, min(MAX_SAMPLES_PER_CONFIG, len(all_samples)))
         self._write_samples(sampled)
 
@@ -1149,6 +1222,8 @@ async def _run_interactive_async(self, async_inference_fn) -> dict:
             "elapsed_seconds_median": round(float(np.median(run_elapsed_times)), 1) if run_elapsed_times else None,
             "ttft_ms_p99_reliability":
                 _reliability_block(ttft_p99_per_run, decimals=2),
+            "power_watts_avg": power_stats.avg,
+            "power_watts_peak": power_stats.peak,
         }}
 
     # ------------------------------------------------------------------