diff --git a/frontier/cc_backend/backends/vidur_cc_backend.py b/frontier/cc_backend/backends/vidur_cc_backend.py index ca1983a..0c57f05 100644 --- a/frontier/cc_backend/backends/vidur_cc_backend.py +++ b/frontier/cc_backend/backends/vidur_cc_backend.py @@ -882,2 +882,21 @@ class VidurCCBackend(BaseCCBackend): - # Fallback to analytical if not in cache - logger.debug(f"num_tokens={num_tokens} not in cache, using analytical fallback") + # The precomputed lookup is capped at 100k elements, while realistic + # TP payloads are commonly much larger. A cache miss does not mean the + # measured-data model is unavailable: predict on demand and memoize the + # exact payload instead of silently switching model families. + with self._cache_lock: + model = self._models.get("all_reduce") + if model is not None: + features = pd.DataFrame({"num_tokens": [num_tokens]}) + result = float(model.predict(features)[0]) + with self._cache_lock: + self._predictions["all_reduce"][(num_tokens,)] = result + logger.debug( + f"predict_allreduce: data_size={data_size_bytes}, num_tokens={num_tokens}, " + f"result={result:.6f} ms (ML model, on-demand cache miss)" + ) + return max(0.0, result) + + logger.debug( + f"num_tokens={num_tokens} not in cache and model unavailable, " + "using analytical fallback" + ) diff --git a/tests/unit/test_vidur_cc_large_payload.py b/tests/unit/test_vidur_cc_large_payload.py new file mode 100644 index 0000000..7e87aa7 --- /dev/null +++ b/tests/unit/test_vidur_cc_large_payload.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import threading +import unittest + +import numpy as np + +from frontier.cc_backend.backends.vidur_cc_backend import VidurCCBackend + + +class RecordingModel: + def __init__(self, value: float) -> None: + self.value = value + self.features = [] + + def predict(self, features): + self.features.append(features.copy()) + return np.array([self.value]) + + +class VidurCCLargePayloadTest(unittest.TestCase): + def test_cache_miss_uses_measured_model_and_memoizes(self) -> None: + backend = object.__new__(VidurCCBackend) + backend._cache_lock = threading.RLock() + backend._num_devices = 2 + backend._predictions = {"all_reduce": {(100000,): 0.1}} + model = RecordingModel(0.321) + backend._models = {"all_reduce": model} + backend._analytical_fallback_allreduce = lambda *_: self.fail( + "analytical fallback must not run when the measured model exists" + ) + + data_size_bytes = 2048 * 2048 * 2 + expected_elements = data_size_bytes // 2 + first = backend.predict_allreduce(data_size_bytes, num_devices=2) + second = backend.predict_allreduce(data_size_bytes, num_devices=2) + + self.assertEqual(first, 0.321) + self.assertEqual(second, 0.321) + self.assertEqual(len(model.features), 1) + self.assertEqual( + int(model.features[0].iloc[0]["num_tokens"]), expected_elements + ) + self.assertEqual( + backend._predictions["all_reduce"][(expected_elements,)], 0.321 + ) + + +if __name__ == "__main__": + unittest.main()