Apply best-effort no-cache scoring knobs without importing optional runtimes. Hugging Face-style models decide attention kernels from their config rather than forward kwargs. Setting these attributes keeps the wrapper explicit while still allowing plain PyTorch modules to run uncha
(
model: torch.nn.Module,
config: StatelessForwardConfig,
)
| 333 | "active_completion_tokens": active_tokens, |
| 334 | "device": str(input_ids.device), |
| 335 | "dtype": str(input_ids.dtype).replace("torch.", ""), |
| 336 | "elapsed_ms": elapsed_seconds * 1000.0, |
| 337 | "use_cache": False, |
| 338 | "use_cache_passed": bool(use_cache_passed), |
| 339 | "detached_outputs": bool(detached_outputs), |
| 340 | "zero_kv_cache": kv_cache_summary.tensor_count == 0, |
| 341 | "kv_cache_output_present": kv_cache_summary.tensor_count > 0, |
| 342 | "kv_cache_output_tensors": kv_cache_summary.tensor_count, |
| 343 | "kv_cache_output_bytes": kv_cache_summary.total_bytes, |
| 344 | "kv_cache_output_mb": kv_cache_summary.total_mb, |
| 345 | **dict(no_cache_policy), |
| 346 | } |
| 347 | if cuda_tracking: |
| 348 | device = input_ids.device |
| 349 | metrics["peak_allocated_mb"] = torch.cuda.max_memory_allocated(device) / 1_048_576.0 |
| 350 | metrics["peak_reserved_mb"] = torch.cuda.max_memory_reserved(device) / 1_048_576.0 |
| 351 | return metrics |
| 352 | |
| 353 | |
| 354 | def configure_stateless_model( |
| 355 | model: torch.nn.Module, |
| 356 | config: StatelessForwardConfig, |
| 357 | ) -> dict[str, float | int | str | bool]: |
| 358 | """ |
| 359 | Apply best-effort no-cache scoring knobs without importing optional runtimes. |
| 360 | |
| 361 | Hugging Face-style models decide attention kernels from their config rather |
| 362 | than forward kwargs. Setting these attributes keeps the wrapper explicit |
| 363 | while still allowing plain PyTorch modules to run unchanged. |
| 364 | """ |
| 365 | |
| 366 | use_cache_targets = 0 |
| 367 | attention_targets = 0 |
| 368 | |
| 369 | for target in _model_config_targets(model): |
no test coverage detected