Skip to content

Fix predictor lookup cache miss fallback#4

Open
tiger3rose wants to merge 1 commit into
NetX-lab:mainfrom
tiger3rose:fix/sklearn-lookup-cache-miss
Open

Fix predictor lookup cache miss fallback#4
tiger3rose wants to merge 1 commit into
NetX-lab:mainfrom
tiger3rose:fix/sklearn-lookup-cache-miss

Conversation

@tiger3rose

Copy link
Copy Markdown

Summary

Fix sklearn execution-time predictor lookup-cache misses by falling back to on-demand sklearn prediction and caching the result at runtime.

Runtime batch aggregation can produce shapes outside the precomputed prediction grid, for example larger aggregated KV cache sizes in chunked prefill. Previously, lookup-backed predictor paths directly indexed the precomputed cache and raised KeyError even though a trained sklearn model was available.

Changes

  • Add _get_lookup_or_predict() for lookup-backed predictor models.
  • Preserve existing cache-hit behavior.
  • On cache miss, run the trained sklearn model on demand.
  • Clamp negative predictions to zero and cache the runtime result.
  • Replace direct lookup-cache indexing across dense linear/norm/add/rope, KV cache save, decode, prefill, and speculative decode paths.

Validation

  • python3 -m py_compile frontier/execution_time_predictor/sklearn_execution_time_predictor.py
  • git diff --check -- frontier/execution_time_predictor/sklearn_execution_time_predictor.py
  • Verified the original attn_prefill lookup miss path progresses with on-demand fallback.
  • Verified full online simulation after collecting true mixed attention profiling.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant