DotCache-Arena / dotcache /__init__.py
DeanoCalver's picture
Initial DotCache Arena Space upload
751ad26 verified
Raw
History Blame Contribute Delete
11.8 kB
from .attention_reference import (
explicit_dequantized_attention,
mix_page_ref,
run_attention_reference,
score_page_ref,
)
from .attention_runtime import decode_step, mix_page, prepare_page, prepare_pages, score_page
from .config import DotCacheConfig
from .evaluation_protocol import EvaluationMetadata, EvaluationRecord, build_evaluation_record, derive_standard_metrics
from .page_cache import PreparedPageCache
from .page_oracle import (
OracleDatasetSplitSummary,
OracleDatasetSplitManifestEntry,
OracleDatasetSplitSuiteResult,
OracleDatasetSplitSuiteSpec,
OracleLabelRecord,
OracleLabelingResult,
OracleSelectorCandidateTrainingRow,
OracleSelectorTrainingRow,
OracleBatchReplayResult,
OracleBatchTraceResult,
OracleCandidateResult,
OracleReplayResult,
OracleThresholds,
build_oracle_label_records,
build_selector_candidate_training_rows,
build_selector_training_rows,
PageTraceRecord,
default_candidate_specs,
load_oracle_dataset_split_manifest,
load_oracle_label_records,
load_page_trace,
load_page_trace_manifest,
materialize_oracle_dataset_split,
materialize_oracle_dataset_split_suite,
merge_page_trace_manifests,
run_oracle_labeling,
run_oracle_batch_replay,
run_oracle_replay,
save_oracle_labels,
save_oracle_dataset_split_manifest,
save_page_trace_manifest,
save_selector_candidate_training_rows,
save_selector_training_rows,
save_page_trace,
select_page_trace_paths,
upsert_oracle_dataset_split_manifest_entry,
)
from .encode import encode_page
from .planner import choose_mode
from .session_runtime import PagedDecodeSession
from .tracing import ExecutionTrace
from .types import EncodedPage, PageHeader
from .model_kv_cache import ModelPagedKVCache
from .kv_quant_registry import KvQuantBaselineSpec, get_kv_quant_baseline, list_kv_quant_baselines
from .model_registry import ModelSpec, get_model_spec, list_model_specs
from .planner import LayerPolicy, PageModeSpec, PageStats, choose_page_mode, observe_page
from .state_cache_sim import (
CapturedStateSample,
StateAblationResult,
StateLayerRecord,
StateSimResult,
StateTileSpec,
load_captured_state_sample,
simulate_state_codec,
simulate_state_sequence,
)
from .selector_baselines import (
CandidateSafeLinearSelectorModel,
LinearSelectorModel,
RUNTIME_SELECTOR_FEATURE_NAMES,
SelectorCandidateExample,
SelectorEvaluationSummary,
SelectorExample,
SelectorPrediction,
SelectorSplit,
StaticRuleSelectorModel,
evaluate_candidate_selector_model,
evaluate_selector_model,
adjust_linear_selector_model_logits,
discover_selector_split_dirs,
load_selector_candidate_examples,
load_selector_examples,
load_selector_split_examples,
render_selector_fixed_split_batch_markdown,
render_selector_bakeoff_markdown,
render_selector_aggregate_markdown,
run_selector_baseline_bakeoff,
run_selector_fixed_split_batch_bakeoff,
run_selector_fixed_split_bakeoff,
run_selector_leave_prompt_family_layer_out_bakeoff,
run_selector_leave_prompt_family_out_bakeoff,
run_selector_leave_prompt_variant_out_bakeoff,
run_selector_leave_layer_out_bakeoff,
run_selector_multiseed_bakeoff,
save_linear_selector_model,
split_selector_examples,
train_candidate_safe_linear_selector,
train_linear_selector,
train_runtime_linear_selector,
train_static_rule_selector,
)
from .selector_profiles import SelectorProfileResolution, resolve_learned_page_selector_profile
__all__ = [
"DotCacheConfig",
"EncodedPage",
"EvaluationMetadata",
"EvaluationRecord",
"ExecutionTrace",
"CapturedStateSample",
"CandidateSafeLinearSelectorModel",
"KvQuantBaselineSpec",
"LayerPolicy",
"LinearSelectorModel",
"ModelSpec",
"ModelPagedKVCache",
"OracleDatasetSplitManifestEntry",
"OracleDatasetSplitSummary",
"OracleDatasetSplitSuiteResult",
"OracleDatasetSplitSuiteSpec",
"OracleLabelRecord",
"OracleLabelingResult",
"OracleSelectorCandidateTrainingRow",
"OracleSelectorTrainingRow",
"OracleBatchReplayResult",
"OracleBatchTraceResult",
"OracleCandidateResult",
"OracleReplayResult",
"OracleThresholds",
"PageTraceRecord",
"PageModeSpec",
"PageStats",
"PageHeader",
"PagedDecodeSession",
"PreparedPageCache",
"RUNTIME_SELECTOR_FEATURE_NAMES",
"SelectorCandidateExample",
"SelectorEvaluationSummary",
"SelectorExample",
"SelectorPrediction",
"SelectorProfileResolution",
"SelectorSplit",
"StaticRuleSelectorModel",
"StateAblationResult",
"StateLayerRecord",
"StateSimResult",
"StateTileSpec",
"choose_mode",
"choose_page_mode",
"build_evaluation_record",
"build_oracle_label_records",
"build_selector_candidate_training_rows",
"build_selector_training_rows",
"default_candidate_specs",
"decode_step",
"derive_standard_metrics",
"adjust_linear_selector_model_logits",
"discover_selector_split_dirs",
"encode_page",
"evaluate_candidate_selector_model",
"evaluate_selector_model",
"explicit_dequantized_attention",
"get_kv_quant_baseline",
"get_model_spec",
"list_model_specs",
"list_kv_quant_baselines",
"load_oracle_dataset_split_manifest",
"load_oracle_label_records",
"load_selector_candidate_examples",
"load_selector_examples",
"load_selector_split_examples",
"load_page_trace",
"materialize_oracle_dataset_split",
"materialize_oracle_dataset_split_suite",
"merge_page_trace_manifests",
"mix_page",
"mix_page_ref",
"prepare_page",
"prepare_pages",
"run_attention_reference",
"score_page",
"score_page_ref",
"load_captured_state_sample",
"simulate_state_codec",
"simulate_state_sequence",
"observe_page",
"render_selector_aggregate_markdown",
"render_selector_bakeoff_markdown",
"render_selector_fixed_split_batch_markdown",
"load_page_trace_manifest",
"run_selector_baseline_bakeoff",
"run_selector_fixed_split_batch_bakeoff",
"run_selector_fixed_split_bakeoff",
"run_selector_leave_prompt_family_layer_out_bakeoff",
"run_selector_leave_prompt_family_out_bakeoff",
"run_selector_leave_prompt_variant_out_bakeoff",
"run_selector_leave_layer_out_bakeoff",
"run_selector_multiseed_bakeoff",
"run_oracle_labeling",
"run_oracle_batch_replay",
"run_oracle_replay",
"resolve_learned_page_selector_profile",
"save_oracle_labels",
"save_oracle_dataset_split_manifest",
"save_page_trace_manifest",
"save_selector_candidate_training_rows",
"save_selector_training_rows",
"save_page_trace",
"save_linear_selector_model",
"select_page_trace_paths",
"upsert_oracle_dataset_split_manifest_entry",
"split_selector_examples",
"train_candidate_safe_linear_selector",
"train_linear_selector",
"train_runtime_linear_selector",
"train_static_rule_selector",
]
try: # pragma: no cover - optional HF path
from .integrations import (
LlamaDotCacheHarness,
LlamaDotCacheModelAdapter,
LlamaReplayRecord,
build_llama_page_trace_records,
build_llama_prefill_page_trace_records,
export_llama_page_traces,
Qwen2DotCacheHarness,
Qwen2DotCacheModelAdapter,
Qwen35AttentionSubsetDotCacheHarness,
Qwen35AttentionSubsetDotCacheModelAdapter,
Qwen35AttentionSubsetHarness,
Qwen35AttentionSubsetModelAdapter,
Qwen35DeltaNetStateRecord,
Qwen35DeltaNetStateHarness,
Qwen35DeltaNetStateModelAdapter,
build_qwen35_deltanet_state_sample,
build_attention_subset_page_trace_records,
build_attention_subset_prefill_page_trace_records,
capture_qwen35_deltanet_state_sample,
Qwen35TextHarness,
Qwen35TextModelAdapter,
export_attention_subset_page_traces,
VllmAdapterConfig,
VllmDotCacheModelAdapter,
VllmPagedKVCache,
VLLM_V1_MULTIPROCESSING_ENV,
configure_vllm_inprocess_runtime,
get_vllm_version,
install_dotcache_on_vllm_model,
install_dotcache_on_vllm_runtime,
require_supported_vllm_version,
inspect_qwen35_deltanet_state,
inspect_qwen35_hybrid_state,
load_qwen35_text_only_from_pretrained,
run_qwen35_attention_subset_prefill_ablation_harness,
run_qwen35_attention_subset_page_trace_capture_harness,
run_qwen35_attention_subset_dotcache_harness,
run_qwen35_attention_subset_replay_harness,
run_qwen35_deltanet_state_ablation_harness,
run_llama_generation_harness,
run_llama_page_trace_capture_harness,
run_llama_replay_harness,
run_qwen2_generation_harness,
run_qwen2_loss_harness,
run_qwen2_replay_harness,
run_qwen35_text_generation_harness,
run_qwen35_text_loss_harness,
save_qwen35_deltanet_state_sample,
summarize_qwen35_dotcache_fit,
summarize_qwen35_hybrid_state,
transformers_available,
vllm_available,
)
except ImportError: # pragma: no cover - exercised when optional deps are absent
pass
else:
__all__.extend(
[
"LlamaDotCacheHarness",
"LlamaDotCacheModelAdapter",
"LlamaReplayRecord",
"build_llama_page_trace_records",
"build_llama_prefill_page_trace_records",
"export_llama_page_traces",
"Qwen2DotCacheHarness",
"Qwen2DotCacheModelAdapter",
"Qwen35AttentionSubsetDotCacheHarness",
"Qwen35AttentionSubsetDotCacheModelAdapter",
"Qwen35AttentionSubsetHarness",
"Qwen35AttentionSubsetModelAdapter",
"Qwen35DeltaNetStateRecord",
"Qwen35DeltaNetStateHarness",
"Qwen35DeltaNetStateModelAdapter",
"build_attention_subset_page_trace_records",
"build_attention_subset_prefill_page_trace_records",
"build_qwen35_deltanet_state_sample",
"capture_qwen35_deltanet_state_sample",
"Qwen35TextHarness",
"Qwen35TextModelAdapter",
"export_attention_subset_page_traces",
"VllmAdapterConfig",
"VllmDotCacheModelAdapter",
"VllmPagedKVCache",
"VLLM_V1_MULTIPROCESSING_ENV",
"configure_vllm_inprocess_runtime",
"get_vllm_version",
"install_dotcache_on_vllm_model",
"install_dotcache_on_vllm_runtime",
"require_supported_vllm_version",
"inspect_qwen35_deltanet_state",
"inspect_qwen35_hybrid_state",
"load_qwen35_text_only_from_pretrained",
"run_qwen35_attention_subset_prefill_ablation_harness",
"run_qwen35_attention_subset_page_trace_capture_harness",
"run_qwen35_attention_subset_dotcache_harness",
"run_qwen35_attention_subset_replay_harness",
"run_qwen35_deltanet_state_ablation_harness",
"run_llama_generation_harness",
"run_llama_page_trace_capture_harness",
"run_llama_replay_harness",
"run_qwen2_generation_harness",
"run_qwen2_loss_harness",
"run_qwen2_replay_harness",
"run_qwen35_text_generation_harness",
"run_qwen35_text_loss_harness",
"save_qwen35_deltanet_state_sample",
"summarize_qwen35_dotcache_fit",
"summarize_qwen35_hybrid_state",
"transformers_available",
"vllm_available",
]
)