[anpl1623]

Summary

This PR implements key-value (KV) caching and optimized autoregressive decoding for Kronos inference, significantly speeding up autoregressive multi-step forecasts and multi-path sampling while maintaining exact numerical equivalence with the reference implementation.

Additionally, this PR exposes unaggregated sample paths via predict_paths / return_paths, enabling probabilistic forecasting, dispersion/variance calculations, and quantile-based risk analysis.

---

Key Changes

1. Key-Value Cache & Attention Optimizations (model/module.py)

• LayerKVCache & KVCache: Added per-layer and model-level KV cache containers with support for incremental token append and cache broadcasting via .expand(repeats) (repeat_interleave).
• Rotary Positional Embeddings (RotaryPositionalEmbedding): Added support for positional offset, allowing single query/key tokens at step $t$ to be correctly rotated against cached history.
• MultiHeadAttentionWithRoPE:
  • Integrated layer_cache into self-attention.
  • Adjusted causal masking (is_causal = seq_len == key_len) to safely handle cached single-token queries against accumulated history.

2. Fast Autoregressive Inference (model/kronos.py)

• Shared Prefix Encoding: Encodes the prompt prefix once per series and broadcasts the resulting KV cache across all Monte Carlo sample paths, eliminating redundant prefix passes.
• Incremental Decoding: Subsequent prediction steps feed only the newly sampled tokens into decode_s1 instead of recomputing the full transformer stack over past tokens.
• Logit Projection Optimization (last_only=True): Added last_only flag to decode_s1 and decode_s2 to project only the active step to vocabulary logits, avoiding large vocabulary-sized projections across historical context.
• Safe Fallback: If context_length + pred_len > max_context, auto_regressive_inference warns and falls back to the reference sliding-window implementation rather than returning diverging numerical approximations.

3. Sampling Distribution & API Improvements (model/kronos.py)

• predict_paths(...) / return_paths=True: Added KronosPredictor.predict_paths() and return_paths parameter to generate() and auto_regressive_inference(), returning all sample_count paths [sample_count, pred_len, features] before mean aggregation.
• Predictor Refactoring:
  • Factored series normalization and timestamp parsing into a reusable _prepare_series() helper.
  • Added .feature_cols property on KronosPredictor.
  • Set models to .eval() mode upon initialization in KronosPredictor.

4. Comprehensive Regression Tests (tests/test_kv_cache.py)

• test_cached_matches_reference: Validates exact numerical agreement between cached and uncached paths across different sample_count values.
• test_paths_are_not_pre_averaged: Validates return_paths output shape, distribution dispersion, and consistency with averaged predictions.
• test_cached_path_refuses_eviction: Ensures cached_auto_regressive_inference rejects context lengths requiring eviction.
• test_overflow_falls_back_to_reference: Verifies fallback to the reference path when context overflows max_context.

---

Verification

Ran full test suite using pytest:

PYTHONPATH=. pytest tests/test_kv_cache.py