Spaces:
Running
Running
| from __future__ import annotations | |
| import math | |
| from .models import AcceleratorProfile, ModelProfile | |
| from .profiles import QUANTIZATION_BYTES, QUANTIZATION_COMPUTE_MULTIPLIER | |
| class AnalyticalLatencyModel: | |
| """Hardware-aware analytical proxy for prefill/decode latency. | |
| This model intentionally does *not* claim benchmark-grade accuracy. It uses | |
| model architecture, accelerator peak specs, conservative efficiency factors, | |
| and a roofline-style max(compute_time, memory_time) approximation. Public | |
| profiles are tagged `analytical-reference` throughout the app. | |
| """ | |
| def __init__( | |
| self, | |
| model: ModelProfile, | |
| accelerator: AcceleratorProfile, | |
| quantization: str = "fp16", | |
| prefill_scale: float = 1.0, | |
| decode_scale: float = 1.0, | |
| ): | |
| if quantization not in QUANTIZATION_BYTES: | |
| raise ValueError(f"Unsupported quantization: {quantization}") | |
| self.model = model | |
| self.accelerator = accelerator | |
| self.quantization = quantization | |
| self.weight_bytes_per_param = QUANTIZATION_BYTES[quantization] | |
| self.compute_overhead = QUANTIZATION_COMPUTE_MULTIPLIER[quantization] | |
| self.prefill_scale = max(float(prefill_scale), 1e-6) | |
| self.decode_scale = max(float(decode_scale), 1e-6) | |
| def model_weight_gb(self) -> float: | |
| return self.model.params_b * self.weight_bytes_per_param | |
| def kv_bytes_per_token(self) -> float: | |
| # K + V, all layers, KV heads only. KV state is assumed fp16 in the current model. | |
| return ( | |
| 2 | |
| * self.model.layers | |
| * self.model.kv_heads | |
| * self.model.head_dim | |
| * 2.0 | |
| ) | |
| def _compute_efficiency(self, batch_size: int, tokens: int) -> float: | |
| scale = 1.0 + 0.12 * math.log2(max(batch_size, 1)) + 0.035 * math.log2(max(tokens, 1)) | |
| return min(0.88, self.accelerator.compute_efficiency * scale) | |
| def _bandwidth_efficiency(self, batch_size: int) -> float: | |
| scale = 1.0 + 0.06 * math.log2(max(batch_size, 1)) | |
| return min(0.91, self.accelerator.bandwidth_efficiency * scale) | |
| def prefill_seconds(self, token_counts: list[int]) -> float: | |
| if not token_counts: | |
| return 0.0 | |
| batch = len(token_counts) | |
| total_tokens = sum(token_counts) | |
| max_seq = max(token_counts) | |
| dense_flops = 2.0 * self.model.params_b * 1e9 * total_tokens | |
| # Approximate quadratic attention component. It is small for short | |
| # contexts but becomes visible at long prompts. | |
| attention_flops = ( | |
| 4.0 | |
| * self.model.layers | |
| * self.model.hidden_size | |
| * sum(t * t for t in token_counts) | |
| ) | |
| flops = (dense_flops + attention_flops) * self.compute_overhead | |
| compute = flops / ( | |
| self.accelerator.peak_tflops_fp16 * 1e12 * self._compute_efficiency(batch, max_seq) | |
| ) | |
| weight_bytes = self.model.params_b * 1e9 * self.weight_bytes_per_param | |
| activation_bytes = total_tokens * self.model.hidden_size * self.model.layers * 2.0 * 0.18 | |
| memory = (weight_bytes + activation_bytes) / ( | |
| self.accelerator.bandwidth_gbps * 1e9 * self._bandwidth_efficiency(batch) | |
| ) | |
| # Kernel launch / scheduling proxy prevents implausibly tiny times. | |
| launch = 0.0018 + 0.00008 * batch | |
| return (max(compute, memory * 0.28) + launch) * self.prefill_scale | |
| def decode_step_seconds(self, context_lengths: list[int]) -> float: | |
| if not context_lengths: | |
| return 0.0 | |
| batch = len(context_lengths) | |
| avg_context = sum(context_lengths) / batch | |
| dense_flops = 2.0 * self.model.params_b * 1e9 * batch | |
| attention_flops = ( | |
| 4.0 | |
| * self.model.layers | |
| * self.model.hidden_size | |
| * sum(context_lengths) | |
| ) | |
| flops = (dense_flops + attention_flops) * self.compute_overhead | |
| compute = flops / ( | |
| self.accelerator.peak_tflops_fp16 * 1e12 * self._compute_efficiency(batch, 1) | |
| ) | |
| # Decode is commonly memory-bound. We model one shared weight stream plus | |
| # KV reads that scale with batch and context length. | |
| weight_bytes = self.model.params_b * 1e9 * self.weight_bytes_per_param | |
| kv_read_bytes = self.kv_bytes_per_token() * sum(context_lengths) | |
| memory = (weight_bytes + kv_read_bytes) / ( | |
| self.accelerator.bandwidth_gbps * 1e9 * self._bandwidth_efficiency(batch) | |
| ) | |
| launch = 0.0012 + 0.000035 * batch + 0.00000003 * avg_context | |
| return (max(compute, memory) + launch) * self.decode_scale | |