fix: preserve agentic delay bounds after context scaling

docs/tutorials/agentic-code-generator.md

@@ -96,6 +96,37 @@ If the sampled session would hit `max_prompt_tokens` before reaching the target:
 - With `allow_truncation: true`, the generator returns the partial session and
   marks it as `forced_retire`.
 
+## Inter-Turn Delays
+
+`inter_turn_delay.context_scaling` defaults to `true` for compatibility with
+existing workloads. After sampling the agentic/human mixture and applying
+`inter_turn_delay.max`, the generator multiplies the delay by:
+
+`max(0.2, 1 - 0.8 * previous_input / max_prompt_tokens)`
+
+`previous_input` is the previous turn's cumulative input length. This is a
+workload heuristic: fuller contexts produce shorter gaps. It is not an
+inference-latency measurement. Set `"context_scaling": false` inside the
+`inter_turn_delay` object to disable it. The setting is saved in the manifest
+and generated config schema, and appears as 0/1 in the quality config summary.
+
+Component `min`/`max` bounds apply during rejection sampling. After context
+scaling, the selected component's `min` is enforced again. The global
+`inter_turn_delay.max` is the final upper cap and takes precedence if it is
+lower than that component's minimum; a global maximum of zero permits zero
+delay. Turn 0 always has zero delay.
+
+The random draws and cap-before-scaling order are unchanged. With scaling
+enabled and no component minima, an existing config and seed produce the same
+trace. Configs with minima may now produce longer gaps where scaling previously
+violated the bound. Disabling scaling intentionally changes delays and timestamps.
+
+Component means/medians describe the underlying lognormal distributions, not
+targets for the final mixture. Rejection sampling, caps, and context scaling
+change the output distribution. `quality.json` reports observed delay statistics
+with null aggregate mean/median targets; the component values shown in reports
+should not be interpreted as those aggregate targets.
+
 ## Generate a Dataset
 
 Create a dataset with the built-in default configuration:

src/aiperf/dataset/agentic_code_gen/configs/spec.json

@@ -147,6 +147,12 @@
           "title": "Agentic Fraction",
           "type": "number"
         },
+        "context_scaling": {
+          "default": true,
+          "description": "Shorten delays as context fills using max(0.2, 1 - 0.8 * previous_input / max_prompt_tokens); disable to keep sampled delays. Component minima are enforced after scaling, with the global max taking precedence. Enabled by default for workload compatibility.",
+          "title": "Context Scaling",
+          "type": "boolean"
+        },
         "human_delay": {
           "$ref": "#/$defs/LognormalParams",
           "description": "Slow delay distribution (human think time)"

src/aiperf/dataset/agentic_code_gen/distributions.py

@@ -88,17 +88,33 @@ def sample_lognormal(
 
 
 def sample_mixture_delay(
-    config: MixtureDelayConfig, rng: Generator, size: int = 1
+    config: MixtureDelayConfig,
+    rng: Generator,
+    size: int = 1,
+    *,
+    delay_scale: float = 1.0,
 ) -> np.ndarray:
     """Sample from the two-component mixture delay model.
 
     For each sample, a Bernoulli draw selects agentic (fast) vs human (slow),
-    then the corresponding lognormal is sampled.
+    then the corresponding lognormal is sampled. The global cap is applied
+    before scaling to preserve historical workloads. Component minima are
+    restored after scaling; the global cap overrides a conflicting minimum.
     """
+    if not 0 < delay_scale <= 1:
+        raise ValueError("delay_scale must be in (0, 1]")
     is_agentic = rng.random(size=size) < config.agentic_fraction
     agentic_samples = sample_lognormal(config.agentic_delay, rng, size=size)
     human_samples = sample_lognormal(config.human_delay, rng, size=size)
     samples = np.where(is_agentic, agentic_samples, human_samples)
+    if config.max is not None:
+        samples = np.minimum(samples, config.max)
+    samples *= delay_scale
+    if config.agentic_delay.min is not None or config.human_delay.min is not None:
+        minimums = np.where(
+            is_agentic, config.agentic_delay.min or 0.0, config.human_delay.min or 0.0
+        )
+        samples = np.maximum(samples, minimums)
     if config.max is not None:
         samples = np.minimum(samples, config.max)
     return samples

src/aiperf/dataset/agentic_code_gen/models.py

@@ -144,6 +144,15 @@ class MixtureDelayConfig(AIPerfBaseModel):
     A Bernoulli draw selects which component to sample from.
     """
 
+    context_scaling: bool = Field(
+        default=True,
+        description=(
+            "Shorten delays as context fills using max(0.2, 1 - 0.8 * "
+            "previous_input / max_prompt_tokens); disable to keep sampled delays. "
+            "Component minima are enforced after scaling, with the global max "
+            "taking precedence. Enabled by default for workload compatibility."
+        ),
+    )
     agentic_fraction: float = Field(
         default=0.7,
         ge=0.0,

src/aiperf/dataset/agentic_code_gen/reporting/metrics.py

@@ -330,6 +330,9 @@ def config_summary(config: SessionDistributionConfig) -> dict[str, float | int]:
         "inter_turn_delay_agentic_fraction": config.inter_turn_delay.agentic_fraction,
         "inter_turn_delay_agentic_mean_ms": config.inter_turn_delay.agentic_delay.mean,
         "inter_turn_delay_human_mean_ms": config.inter_turn_delay.human_delay.mean,
+        "inter_turn_delay_context_scaling": int(
+            config.inter_turn_delay.context_scaling
+        ),
         "turn_mode_enabled": int(config.turns is not None),
     }
     if config.reset is not None:

src/aiperf/dataset/agentic_code_gen/session_synthesizer.py

@@ -140,14 +140,14 @@ def _sample_output_length(self) -> int:
         )
 
     def _sample_delay_ms(self, prev_input: int) -> float:
-        delay_ms = float(
-            sample_mixture_delay(self._config.inter_turn_delay, self._rng, size=1)[0]
-        )
+        config = self._config.inter_turn_delay
         context_ratio = prev_input / self._config.max_prompt_tokens
-        delay_ms *= max(0.2, 1.0 - 0.8 * context_ratio)
-        if self._config.inter_turn_delay.max is not None:
-            delay_ms = min(delay_ms, self._config.inter_turn_delay.max)
-        return delay_ms
+        delay_scale = (
+            max(0.2, 1.0 - 0.8 * context_ratio) if config.context_scaling else 1.0
+        )
+        return float(
+            sample_mixture_delay(config, self._rng, size=1, delay_scale=delay_scale)[0]
+        )
 
     def _sample_new_tokens(self) -> int:
         new_tokens = int(

tests/unit/dataset/agentic_code_gen/test_distributions.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import pytest
+from pytest import param
 
 from aiperf.dataset.agentic_code_gen.distributions import (
     fit_from_samples,
@@ -150,6 +151,48 @@ def test_clipping_works(self) -> None:
 
 
 class TestSampleMixtureDelay:
+    def test_scaled_mixture_applies_each_selected_component_minimum(self) -> None:
+        config = MixtureDelayConfig(
+            agentic_fraction=0.5,
+            agentic_delay=LognormalParams(mean=1000, median=1000, min=900),
+            human_delay=LognormalParams(mean=5000, median=5000, min=4500),
+        )
+
+        samples = sample_mixture_delay(
+            config, np.random.default_rng(42), size=100, delay_scale=0.2
+        )
+
+        assert set(samples) == {900.0, 4500.0}
+
+    def test_zero_global_cap_overrides_component_minima(self) -> None:
+        config = MixtureDelayConfig(
+            agentic_delay=LognormalParams(mean=1000, median=1000, min=900),
+            human_delay=LognormalParams(mean=5000, median=5000, min=4500),
+            max=0,
+        )
+
+        samples = sample_mixture_delay(
+            config, np.random.default_rng(42), size=100, delay_scale=0.2
+        )
+
+        assert np.array_equal(samples, np.zeros(100))
+
+    @pytest.mark.parametrize(
+        "delay_scale",
+        [
+            param(0, id="zero"),
+            param(-1, id="negative"),
+            param(1.1, id="above-one"),
+            param(float("nan"), id="nan"),
+            param(float("inf"), id="inf"),
+        ],
+    )  # fmt: skip
+    def test_invalid_delay_scale_is_rejected(self, delay_scale: float) -> None:
+        with pytest.raises(ValueError, match="delay_scale must be in"):
+            sample_mixture_delay(
+                MixtureDelayConfig(), np.random.default_rng(42), delay_scale=delay_scale
+            )
+
     def test_returns_correct_shape(self) -> None:
         config = MixtureDelayConfig(
             agentic_fraction=0.7,

tests/unit/dataset/agentic_code_gen/test_roundtrip.py

@@ -10,12 +10,15 @@
 
 import orjson
 import pytest
+from pytest import param
 
 from aiperf.dataset.agentic_code_gen.config import load_config
 from aiperf.dataset.agentic_code_gen.distributions import lognormal_from_mean_median
 from aiperf.dataset.agentic_code_gen.models import (
     CacheLayerConfig,
     Layer15GroupConfig,
+    LognormalParams,
+    MixtureDelayConfig,
     ResetConfig,
     SessionDistributionConfig,
 )
@@ -25,6 +28,43 @@
 
 
 class TestRoundtrip:
+    @pytest.mark.parametrize(
+        "context_scaling, expected_delay",
+        [param(True, 900, id="scaled"), param(False, 1000, id="unscaled")],
+    )  # fmt: skip
+    def test_delay_config_survives_manifest_and_trace_roundtrip(
+        self,
+        tmp_path: Path,
+        turns_config: SessionDistributionConfig,
+        context_scaling: bool,
+        expected_delay: float,
+    ) -> None:
+        turns_config.cache.layer1_tokens = 2000
+        turns_config.inter_turn_delay = MixtureDelayConfig(
+            context_scaling=context_scaling,
+            agentic_fraction=1.0,
+            agentic_delay=LognormalParams(mean=1000, median=1000, min=900),
+        )
+        sessions = SessionSynthesizer(turns_config, seed=42).synthesize_sessions(3)
+
+        jsonl, manifest, quality = write_dataset(
+            sessions, tmp_path / "run", turns_config, seed=42
+        )
+
+        restored_config = load_config(str(manifest))
+        assert restored_config.inter_turn_delay.context_scaling is context_scaling
+        replay = SessionSynthesizer(restored_config, seed=42).synthesize_sessions(3)
+        assert replay == sessions
+        turns = [
+            MooncakeTrace(**orjson.loads(line))
+            for line in jsonl.read_bytes().splitlines()
+        ]
+        delays = [turn.delay for turn in turns if turn.delay is not None]
+        assert len(delays) == 9
+        assert delays == pytest.approx([expected_delay] * 9)
+        summary = orjson.loads(quality.read_bytes())["config_summary"]
+        assert summary["inter_turn_delay_context_scaling"] == int(context_scaling)
+
     def test_config_to_mooncake_roundtrip(self, tmp_path: Path) -> None:
         config = SessionDistributionConfig()
         synth = SessionSynthesizer(config, seed=42)

tests/unit/dataset/agentic_code_gen/test_session_synthesizer.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import pytest
+from pytest import param
 
 from aiperf.dataset.agentic_code_gen.distributions import lognormal_from_mean_median
 from aiperf.dataset.agentic_code_gen.models import (
@@ -160,6 +161,68 @@ def test_sessions_have_at_least_one_turn(
             assert len(session.turns) >= 1
 
 
+class TestInterTurnDelay:
+    @pytest.mark.parametrize(
+        "agentic_fraction, expected",
+        [param(1.0, 900, id="agentic"), param(0.0, 4500, id="human")],
+    )  # fmt: skip
+    @pytest.mark.parametrize(
+        "context_ratio",
+        [param(0.75, id="partly-full"), param(1.0, id="full")],
+    )  # fmt: skip
+    def test_context_scaling_preserves_selected_component_minimum(
+        self, agentic_fraction: float, expected: float, context_ratio: float
+    ) -> None:
+        config = SessionDistributionConfig(
+            inter_turn_delay=MixtureDelayConfig(
+                agentic_fraction=agentic_fraction,
+                agentic_delay=LognormalParams(mean=1000, median=1000, min=900),
+                human_delay=LognormalParams(mean=5000, median=5000, min=4500),
+            )
+        )
+        synth = SessionSynthesizer(config, seed=42)
+
+        delay = synth._sample_delay_ms(int(config.max_prompt_tokens * context_ratio))
+
+        assert delay == pytest.approx(expected)
+
+    @pytest.mark.parametrize(
+        "context_ratio",
+        [
+            param(0.0, id="empty"),
+            param(0.75, id="partly-full"),
+            param(1.0, id="full"),
+        ],
+    )  # fmt: skip
+    def test_disabled_context_scaling_preserves_sampled_delay(
+        self, context_ratio: float
+    ) -> None:
+        config = SessionDistributionConfig(
+            inter_turn_delay=MixtureDelayConfig(
+                context_scaling=False,
+                agentic_fraction=1.0,
+                agentic_delay=LognormalParams(mean=1000, median=1000),
+            )
+        )
+        synth = SessionSynthesizer(config, seed=42)
+
+        delay = synth._sample_delay_ms(int(config.max_prompt_tokens * context_ratio))
+
+        assert delay == pytest.approx(1000)
+
+    def test_global_cap_takes_precedence_over_component_minimum(self) -> None:
+        config = SessionDistributionConfig(
+            inter_turn_delay=MixtureDelayConfig(
+                agentic_fraction=0.0,
+                human_delay=LognormalParams(mean=5000, median=5000, min=4500),
+                max=600,
+            )
+        )
+        synth = SessionSynthesizer(config, seed=42)
+
+        assert synth._sample_delay_ms(config.max_prompt_tokens) == 600
+
+
 class TestExplicitTurnMode:
     def _overflowing_turn_mode_config(
         self, allow_truncation: bool