diff --git a/.gitignore b/.gitignore index 0381dbf2..3382b367 100644 --- a/.gitignore +++ b/.gitignore @@ -168,7 +168,10 @@ cache settings.json # Task memory and repair receipts are local-only; do not vendor them. -task_memory/ +# Narrow exceptions publish the review records of the two stacked PRs below. +task_memory/* +!task_memory/task_2026-09-21_oversized_module_split/ +!task_memory/task_2026-09-21_issue26_correctness_pr/ repairs/ # Local linked worktrees used for isolated feature implementation. diff --git a/frontier/config/cluster_config.py b/frontier/config/cluster_config.py new file mode 100644 index 00000000..52bc2537 --- /dev/null +++ b/frontier/config/cluster_config.py @@ -0,0 +1,1888 @@ +"""The cluster topology configuration and its flat per-role field surface. + +`ClusterConfig` declares one flat field for every role-specific override the +CLI exposes, validates their combinations, and sets up either the single +monolithic cluster or the disaggregated role clusters. Construction of the +per-role configurations lives in `cluster_role_config`, topology reporting in +`cluster_topology_summary`. +""" + +from __future__ import annotations + +from dataclasses import MISSING, dataclass, field +from typing import List, Optional + +from frontier.config.cluster_role_config import ( + ClusterRoleConfigBuilder, + _get_cc_backend_configs, +) +from frontier.config.cluster_scheduler_config import ( + BaseClusterSchedulerConfig, + RoundRobinClusterSchedulerConfig, +) +from frontier.config.cluster_topology_summary import ClusterTopologySummary +from frontier.config.execution_time_predictor_config import ( + BaseExecutionTimePredictorConfig, + RandomForrestExecutionTimePredictorConfig, +) +from frontier.config.parallel_semantics import ( + FrontierParallelismMapping, + validate_frontier_shared_parallel_domains, +) +from frontier.config.release_guards import ( + AICONFIGURATOR_BACKEND_RELEASE_ERROR, + DISAGGREGATED_CLUSTER_FIELD_NAMES, + DISAGGREGATED_CLUSTER_FIELD_PREFIXES, +) +from frontier.config.replica_config import ReplicaConfig +from frontier.config.replica_scheduler_config import ( + BaseReplicaSchedulerConfig, + SarathiSchedulerConfig, +) +from frontier.logger import init_logger +from frontier.types import ClusterSchedulerType, ClusterType + +logger = init_logger(__name__) + + +@dataclass +class ClusterConfig(ClusterRoleConfigBuilder, ClusterTopologySummary): + # === Common fields for all modes === + cluster_scheduler_config: BaseClusterSchedulerConfig = field( + default_factory=RoundRobinClusterSchedulerConfig, + metadata={ + "help": "Cluster scheduler config.", + }, + ) + replica_scheduler_config: BaseReplicaSchedulerConfig = field( + default_factory=SarathiSchedulerConfig, + metadata={"help": "Replica scheduler config."}, + ) + cluster_type: ClusterType = field( + default=None, + metadata={ + "help": "Type of the cluster: monolithic, prefill, decode-attn, or decode-ffn." + }, + ) + execution_time_predictor_config: BaseExecutionTimePredictorConfig = field( + default_factory=RandomForrestExecutionTimePredictorConfig, + metadata={"help": "Execution time predictor config."}, + ) + cc_backend_config: BaseCCBackendConfig = field( + default_factory=lambda: _get_cc_backend_configs()[5](), # AstraSimAnalyticalCCBackendConfig + metadata={ + "help": "CC (Collective Communication) backend config for communication latency prediction." + }, + ) + + # === co-location/Monolithic mode fields === + num_replicas: Optional[int] = field( + default=1, + metadata={ + "help": "Number of replicas", + }, + ) + replica_config: Optional[ReplicaConfig] = field( + default_factory=lambda: ReplicaConfig(model_name="meta-llama/Llama-2-7b-hf"), + metadata={ + "help": "Replica configuration", + }, + ) + + # === Disaggregated mode fields === + prefill_cluster_num_replicas: Optional[int] = field( + default=None, + metadata={ + "help": "Number of replicas for prefill cluster. Used only in pd-af-disaggregation mode.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cluster_num_replicas: Optional[int] = field( + default=None, + metadata={ + "help": "Number of replicas for decode attention cluster. Used only in pd-af-disaggregation mode.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cluster_num_replicas: Optional[int] = field( + default=None, + metadata={ + "help": "Number of replicas for decode FFN cluster. " + "Each replica is an independent FFN serving copy; MoE EP lanes are " + "scoped inside the selected replica. Used only in pd-af-disaggregation mode.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_cluster_num_replicas: Optional[int] = field( + default=None, + metadata={ + "help": "Number of replicas for unified decode cluster. Used only in pd-disaggregation mode.", + "mode_dependency": "pd-disaggregation", + }, + ) + prefill_replica_config_memory_margin_fraction: Optional[float] = field( + default=None, + metadata={ + "help": "Memory margin fraction for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_num_pipeline_stages: Optional[int] = field( + default=None, + metadata={ + "help": "Number of pipeline stages for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_attn_tensor_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "Attention tensor parallel size for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_moe_tensor_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "MoE tensor parallel size for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_moe_expert_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "MoE expert parallel size for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_total_expert_num: Optional[int] = field( + default=None, + metadata={ + "help": "Total expert number for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_local_expert_num: Optional[int] = field( + default=None, + metadata={ + "help": "Local expert number for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_router_load_balancing_type: Optional[str] = field( + default=None, + metadata={ + "help": "MOE router load balancing type for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_router_topk: Optional[int] = field( + default=None, + metadata={ + "help": "Router topk for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_device: Optional[str] = field( + default=None, + metadata={ + "help": "Device for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + prefill_replica_config_network_device: Optional[str] = field( + default=None, + metadata={ + "help": "Network device for prefill cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_config_memory_margin_fraction: Optional[float] = field( + default=None, + metadata={ + "help": "Memory margin fraction for decode attention cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_config_num_pipeline_stages: Optional[int] = field( + default=None, + metadata={ + "help": "Number of pipeline stages for decode attention cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_config_attn_tensor_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "Attention tensor parallel size for decode attention cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_config_device: Optional[str] = field( + default=None, + metadata={ + "help": "Device for decode attention cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_config_network_device: Optional[str] = field( + default=None, + metadata={ + "help": "Network device for decode attention cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_memory_margin_fraction: Optional[float] = field( + default=None, + metadata={ + "help": "Memory margin fraction for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_num_pipeline_stages: Optional[int] = field( + default=None, + metadata={ + "help": "Number of pipeline stages for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_moe_tensor_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "MoE tensor parallel size for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_moe_expert_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "MoE expert parallel size for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_total_expert_num: Optional[int] = field( + default=None, + metadata={ + "help": "Total expert number for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_local_expert_num: Optional[int] = field( + default=None, + metadata={ + "help": "Local expert number for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_router_load_balancing_type: Optional[str] = field( + default=None, + metadata={ + "help": "MOE router load balancing type for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_router_topk: Optional[int] = field( + default=None, + metadata={ + "help": "Router topk for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_device: Optional[str] = field( + default=None, + metadata={ + "help": "Device for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_config_network_device: Optional[str] = field( + default=None, + metadata={ + "help": "Network device for decode FFN cluster.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + + # === PD-Disaggregation Mode: Unified DECODE Cluster Configuration === + decode_replica_config_memory_margin_fraction: Optional[float] = field( + default=None, + metadata={ + "help": "Memory margin fraction for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_num_pipeline_stages: Optional[int] = field( + default=None, + metadata={ + "help": "Number of pipeline stages for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_attn_tensor_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "Attention tensor parallel size for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_moe_tensor_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "MoE tensor parallel size for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_moe_expert_parallel_size: Optional[int] = field( + default=None, + metadata={ + "help": "MoE expert parallel size for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_total_expert_num: Optional[int] = field( + default=None, + metadata={ + "help": "Total expert number for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_local_expert_num: Optional[int] = field( + default=None, + metadata={ + "help": "Local expert number for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_router_load_balancing_type: Optional[str] = field( + default=None, + metadata={ + "help": "MOE router load balancing type for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_router_topk: Optional[int] = field( + default=None, + metadata={ + "help": "Router topk for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_device: Optional[str] = field( + default=None, + metadata={ + "help": "Device for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_config_network_device: Optional[str] = field( + default=None, + metadata={ + "help": "Network device for unified decode cluster.", + "mode_dependency": "pd-disaggregation", + }, + ) + + # === AF Pipeline Configuration === + # This field is for internal use by the created cluster-specific configs + af_pipeline_num_micro_batch: int = field( + default=-1, + metadata={ + "help": "Internal field for the number of micro-batches. Should be set via cluster-specific parameters below.", + }, + ) + + # User-facing parameters for setting micro-batch number in decode clusters + decode_attn_af_pipeline_num_micro_batch: Optional[int] = field( + default=None, + metadata={"help": "Number of micro-batches for the decode_attn cluster."}, + ) + decode_ffn_af_pipeline_num_micro_batch: Optional[int] = field( + default=None, + metadata={"help": "Number of micro-batches for the decode_ffn cluster."}, + ) + + # User-facing parameter for setting micro-batch SIZE specifically for decode-attn + decode_attn_micro_batch_size: Optional[int] = field( + default=None, + metadata={ + "help": "Target micro-batch SIZE for decode-attn cluster (per (replica, dp)).", + }, + ) + + # User-facing parameter for setting request allocation threshold for decode-attn + decode_attn_request_allocation_threshold: Optional[int] = field( + default=None, + metadata={ + "help": "Request accumulation threshold for decode-attn cluster. " + "Only trigger allocation when accumulated requests reach this threshold. " + "Default: None (equals total number of requests in offline mode).", + }, + ) + + # === AFD CUDA Graph Configuration === + # Aligned with StepFun-vLLM's cudagraph_batch_sizes for AFD attention server + decode_attn_use_cuda_graph: bool = field( + default=False, + metadata={ + "help": "Deprecated. Use SimulationConfig.use_cuda_graph instead. " + "CUDA Graph is now a global setting for pd-af-disaggregation.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cudagraph_capture_sizes: Optional[List[int]] = field( + default=None, + metadata={ + "help": "Deprecated. Use SimulationConfig.cudagraph_capture_sizes instead. " + "CUDA Graph capture sizes are now shared across decode-attn and decode-ffn.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + + decode_attn_replica_id_start_for_ffn: Optional[int] = field( + default=None, + metadata={ + "help": "Derived first global replica id for DECODE_ATTN lanes; used by DECODE_FFN grouping.", + }, + ) + + # === Per-Cluster Replica Scheduler Configuration === + # These fields allow per-cluster-type customization of replica scheduler parameters + # If not set, they fall back to the base replica_scheduler_config values + + # PREFILL cluster scheduler configuration + prefill_replica_scheduler_config_type: Optional[str] = field( + default=None, + metadata={ + "help": "Replica scheduler type for prefill cluster. Overrides base replica_scheduler_config_type.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_replica_scheduler_config_batch_size_cap: Optional[int] = field( + default=None, + metadata={ + "help": "Batch size cap (max_num_seqs) for prefill cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_replica_scheduler_config_max_tokens_in_batch: Optional[int] = field( + default=None, + metadata={ + "help": "Max tokens in batch (max_num_batched_tokens) for prefill cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_replica_scheduler_config_enable_chunked_prefill: Optional[bool] = field( + default=None, + metadata={ + "help": "Enable Chunked Prefill for the prefill cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_replica_scheduler_config_long_prefill_token_threshold: Optional[int] = ( + field( + default=None, + metadata={ + "help": "Long-prefill token threshold for the prefill cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + ) + prefill_replica_scheduler_config_num_blocks: Optional[int] = field( + default=None, + metadata={ + "help": "Number of blocks for prefill cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_replica_scheduler_config_block_size: Optional[int] = field( + default=None, + metadata={ + "help": "Block size for prefill cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_replica_scheduler_config_watermark_blocks_fraction: Optional[float] = field( + default=None, + metadata={ + "help": "Watermark blocks fraction for prefill cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + + # DECODE cluster scheduler configuration (for unified decode in pd-disaggregation mode) + decode_replica_scheduler_config_type: Optional[str] = field( + default=None, + metadata={ + "help": "Replica scheduler type for decode cluster. Overrides base replica_scheduler_config_type.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_scheduler_config_batch_size_cap: Optional[int] = field( + default=None, + metadata={ + "help": "Batch size cap (max_num_seqs) for decode cluster replica scheduler.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_scheduler_config_max_tokens_in_batch: Optional[int] = field( + default=None, + metadata={ + "help": "Max tokens in batch (max_num_batched_tokens) for decode cluster replica scheduler.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_scheduler_config_num_blocks: Optional[int] = field( + default=None, + metadata={ + "help": "Number of blocks for decode cluster replica scheduler.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_scheduler_config_block_size: Optional[int] = field( + default=None, + metadata={ + "help": "Block size for decode cluster replica scheduler.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_replica_scheduler_config_watermark_blocks_fraction: Optional[float] = field( + default=None, + metadata={ + "help": "Watermark blocks fraction for decode cluster replica scheduler.", + "mode_dependency": "pd-disaggregation", + }, + ) + + # DECODE_ATTN cluster scheduler configuration (for pd-af-disaggregation mode) + decode_attn_replica_scheduler_config_type: Optional[str] = field( + default=None, + metadata={ + "help": "Replica scheduler type for decode attention cluster. Overrides base replica_scheduler_config_type.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_scheduler_config_batch_size_cap: Optional[int] = field( + default=None, + metadata={ + "help": "Batch size cap (max_num_seqs) for decode attention cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_scheduler_config_max_tokens_in_batch: Optional[int] = field( + default=None, + metadata={ + "help": "Max tokens in batch (max_num_batched_tokens) for decode attention cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_scheduler_config_num_blocks: Optional[int] = field( + default=None, + metadata={ + "help": "Number of blocks for decode attention cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_scheduler_config_block_size: Optional[int] = field( + default=None, + metadata={ + "help": "Block size for decode attention cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_replica_scheduler_config_watermark_blocks_fraction: Optional[float] = ( + field( + default=None, + metadata={ + "help": "Watermark blocks fraction for decode attention cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + ) + + # DECODE_FFN cluster scheduler configuration (for pd-af-disaggregation mode) + decode_ffn_replica_scheduler_config_type: Optional[str] = field( + default=None, + metadata={ + "help": "Replica scheduler type for decode FFN cluster. Overrides base replica_scheduler_config_type.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_scheduler_config_batch_size_cap: Optional[int] = field( + default=None, + metadata={ + "help": "Batch size cap (max_num_seqs) for decode FFN cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_scheduler_config_max_tokens_in_batch: Optional[int] = field( + default=None, + metadata={ + "help": "Max tokens in batch (max_num_batched_tokens) for decode FFN cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_scheduler_config_num_blocks: Optional[int] = field( + default=None, + metadata={ + "help": "Number of blocks for decode FFN cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_scheduler_config_block_size: Optional[int] = field( + default=None, + metadata={ + "help": "Block size for decode FFN cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_replica_scheduler_config_watermark_blocks_fraction: Optional[float] = ( + field( + default=None, + metadata={ + "help": "Watermark blocks fraction for decode FFN cluster replica scheduler.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + ) + + # === Per-Cluster CC Backend Configuration === + # These fields allow per-cluster-type customization of CC backend parameters + # If not set, they fall back to the base cc_backend_config values + + # PREFILL cluster CC backend configuration + prefill_cc_backend_config_type: Optional[str] = field( + default=None, + metadata={ + "help": "CC backend type for prefill cluster. Options: 'vidur', 'analytical', 'collective_sim', 'astra_sim_analytical'. Overrides base cc_backend_config type.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_network_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Network bandwidth in Gbps for prefill cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_network_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Network latency in microseconds for prefill cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_intra_node_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-node bandwidth in Gbps for prefill cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_repo_root: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication backend repo root for prefill cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_system: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication backend system for prefill cluster CC backend (internal-only mode). Empty means infer from device.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_source_backend: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication source backend for prefill cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_source_version: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication source version for prefill cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_database_mode: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication database mode for prefill cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_tp_allreduce_impl: Optional[str] = field( + default=None, + metadata={ + "help": "TP allreduce implementation for prefill cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_custom_allreduce_variant: Optional[str] = field( + default=None, + metadata={ + "help": "Custom allreduce runtime label for prefill cluster CC backend when internal communication backend raw data has multiple variants.", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_prediction_cache_size: Optional[int] = field( + default=None, + metadata={ + "help": "Prediction cache size for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_placement_order: Optional[str] = field( + default=None, + metadata={ + "help": "Rank placement order for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_intra_server_topology: Optional[str] = field( + default=None, + metadata={ + "help": "Intra-server topology for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_inter_server_topology: Optional[str] = field( + default=None, + metadata={ + "help": "Inter-server topology for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_intra_server_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-server bandwidth in Gbps for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_intra_server_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-server latency in microseconds for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_inter_server_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Inter-server bandwidth in Gbps for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_inter_server_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Inter-server latency in microseconds for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_p2p_src_index: Optional[int] = field( + default=None, + metadata={ + "help": "P2P source participant index for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_p2p_dst_index: Optional[int] = field( + default=None, + metadata={ + "help": "P2P destination participant index for prefill cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_cc_backend_config_nvlink_allreduce_launch_overhead_us: Optional[float] = ( + field( + default=None, + metadata={ + "help": ( + "Per-step intra-server allreduce launch overhead in microseconds " + "for prefill cluster collective-sim backend." + ), + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + ) + prefill_execution_time_predictor_config_mlp_up_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override mlp_up_proj calibration scale for the prefill cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_execution_time_predictor_config_attn_pre_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override attn_pre_proj calibration scale for the prefill cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_execution_time_predictor_config_attn_post_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override attn_post_proj calibration scale for the prefill cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_execution_time_predictor_config_attn_decode_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override attn_decode calibration scale for the prefill cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_execution_time_predictor_config_attn_kv_cache_save_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override attn_kv_cache_save calibration scale for the prefill cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + prefill_execution_time_predictor_config_mlp_down_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override mlp_down_proj calibration scale for the prefill cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-af-disaggregation,pd-disaggregation", + }, + ) + + # DECODE cluster CC backend configuration (for unified decode in pd-disaggregation mode) + decode_cc_backend_config_type: Optional[str] = field( + default=None, + metadata={ + "help": "CC backend type for decode cluster. Options: 'vidur', 'analytical', 'collective_sim', 'astra_sim_analytical'. Overrides base cc_backend_config type.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_network_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Network bandwidth in Gbps for decode cluster CC backend (analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_network_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Network latency in microseconds for decode cluster CC backend (analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_intra_node_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-node bandwidth in Gbps for decode cluster CC backend (analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_repo_root: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication backend repo root for decode cluster CC backend (internal-only mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_system: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication backend system for decode cluster CC backend (internal-only mode). Empty means infer from device.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_source_backend: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication source backend for decode cluster CC backend (internal-only mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_source_version: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication source version for decode cluster CC backend (internal-only mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_database_mode: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication database mode for decode cluster CC backend (internal-only mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_tp_allreduce_impl: Optional[str] = field( + default=None, + metadata={ + "help": "TP allreduce implementation for decode cluster CC backend (internal-only mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_custom_allreduce_variant: Optional[str] = field( + default=None, + metadata={ + "help": "Custom allreduce runtime label for decode cluster CC backend when internal communication backend raw data has multiple variants.", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_prediction_cache_size: Optional[int] = field( + default=None, + metadata={ + "help": "Prediction cache size for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_placement_order: Optional[str] = field( + default=None, + metadata={ + "help": "Rank placement order for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_intra_server_topology: Optional[str] = field( + default=None, + metadata={ + "help": "Intra-server topology for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_inter_server_topology: Optional[str] = field( + default=None, + metadata={ + "help": "Inter-server topology for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_intra_server_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-server bandwidth in Gbps for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_intra_server_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-server latency in microseconds for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_inter_server_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Inter-server bandwidth in Gbps for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_inter_server_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Inter-server latency in microseconds for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_p2p_src_index: Optional[int] = field( + default=None, + metadata={ + "help": "P2P source participant index for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_p2p_dst_index: Optional[int] = field( + default=None, + metadata={ + "help": "P2P destination participant index for decode cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-disaggregation", + }, + ) + decode_cc_backend_config_nvlink_allreduce_launch_overhead_us: Optional[float] = ( + field( + default=None, + metadata={ + "help": ( + "Per-step intra-server allreduce launch overhead in microseconds " + "for decode cluster collective-sim backend." + ), + "mode_dependency": "pd-disaggregation", + }, + ) + ) + decode_execution_time_predictor_config_mlp_up_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override mlp_up_proj calibration scale for the decode cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-disaggregation", + }, + ) + decode_execution_time_predictor_config_attn_pre_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override attn_pre_proj calibration scale for the decode cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-disaggregation", + }, + ) + decode_execution_time_predictor_config_attn_post_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override attn_post_proj calibration scale for the decode cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-disaggregation", + }, + ) + decode_execution_time_predictor_config_attn_decode_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override attn_decode calibration scale for the decode cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-disaggregation", + }, + ) + decode_execution_time_predictor_config_attn_kv_cache_save_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override attn_kv_cache_save calibration scale for the decode cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-disaggregation", + }, + ) + decode_execution_time_predictor_config_mlp_down_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override mlp_down_proj calibration scale for the decode cluster " + "execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-disaggregation", + }, + ) + decode_execution_time_predictor_config_decode_phase_mlp_down_proj_calibration_scale: Optional[ + float + ] = field( + default=None, + metadata={ + "help": ( + "Override decode-phase-only mlp_down_proj calibration scale for the " + "decode cluster execution-time predictor. Must be > 0." + ), + "mode_dependency": "pd-disaggregation", + }, + ) + + # DECODE_ATTN cluster CC backend configuration (for pd-af-disaggregation mode) + decode_attn_cc_backend_config_type: Optional[str] = field( + default=None, + metadata={ + "help": "CC backend type for decode attention cluster. Options: 'vidur', 'analytical', 'collective_sim', 'astra_sim_analytical'. Overrides base cc_backend_config type.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_network_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Network bandwidth in Gbps for decode attention cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_network_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Network latency in microseconds for decode attention cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_intra_node_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-node bandwidth in Gbps for decode attention cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_repo_root: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication backend repo root for decode attention cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_system: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication backend system for decode attention cluster CC backend (internal-only mode). Empty means infer from device.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_source_backend: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication source backend for decode attention cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_source_version: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication source version for decode attention cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_database_mode: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication database mode for decode attention cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_tp_allreduce_impl: Optional[str] = field( + default=None, + metadata={ + "help": "TP allreduce implementation for decode attention cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_custom_allreduce_variant: Optional[str] = field( + default=None, + metadata={ + "help": "Custom allreduce runtime label for decode attention cluster CC backend when internal communication backend raw data has multiple variants.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_prediction_cache_size: Optional[int] = field( + default=None, + metadata={ + "help": "Prediction cache size for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_placement_order: Optional[str] = field( + default=None, + metadata={ + "help": "Rank placement order for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_intra_server_topology: Optional[str] = field( + default=None, + metadata={ + "help": "Intra-server topology for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_inter_server_topology: Optional[str] = field( + default=None, + metadata={ + "help": "Inter-server topology for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_intra_server_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-server bandwidth in Gbps for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_intra_server_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-server latency in microseconds for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_inter_server_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Inter-server bandwidth in Gbps for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_inter_server_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Inter-server latency in microseconds for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_p2p_src_index: Optional[int] = field( + default=None, + metadata={ + "help": "P2P source participant index for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_attn_cc_backend_config_p2p_dst_index: Optional[int] = field( + default=None, + metadata={ + "help": "P2P destination participant index for decode attention cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + + # DECODE_FFN cluster CC backend configuration (for pd-af-disaggregation mode) + decode_ffn_cc_backend_config_type: Optional[str] = field( + default=None, + metadata={ + "help": "CC backend type for decode FFN cluster. Options: 'vidur', 'analytical', 'collective_sim', 'astra_sim_analytical'. Overrides base cc_backend_config type.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_network_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Network bandwidth in Gbps for decode FFN cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_network_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Network latency in microseconds for decode FFN cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_intra_node_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-node bandwidth in Gbps for decode FFN cluster CC backend (analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_repo_root: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication backend repo root for decode FFN cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_system: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication backend system for decode FFN cluster CC backend (internal-only mode). Empty means infer from device.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_source_backend: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication source backend for decode FFN cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_source_version: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication source version for decode FFN cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_database_mode: Optional[str] = field( + default=None, + metadata={ + "help": "Internal-only communication database mode for decode FFN cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_tp_allreduce_impl: Optional[str] = field( + default=None, + metadata={ + "help": "TP allreduce implementation for decode FFN cluster CC backend (internal-only mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_custom_allreduce_variant: Optional[str] = field( + default=None, + metadata={ + "help": "Custom allreduce runtime label for decode FFN cluster CC backend when internal communication backend raw data has multiple variants.", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_prediction_cache_size: Optional[int] = field( + default=None, + metadata={ + "help": "Prediction cache size for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_placement_order: Optional[str] = field( + default=None, + metadata={ + "help": "Rank placement order for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_intra_server_topology: Optional[str] = field( + default=None, + metadata={ + "help": "Intra-server topology for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_inter_server_topology: Optional[str] = field( + default=None, + metadata={ + "help": "Inter-server topology for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_intra_server_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-server bandwidth in Gbps for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_intra_server_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Intra-server latency in microseconds for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_inter_server_bandwidth_gbps: Optional[float] = field( + default=None, + metadata={ + "help": "Inter-server bandwidth in Gbps for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_inter_server_latency_us: Optional[float] = field( + default=None, + metadata={ + "help": "Inter-server latency in microseconds for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_p2p_src_index: Optional[int] = field( + default=None, + metadata={ + "help": "P2P source participant index for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + decode_ffn_cc_backend_config_p2p_dst_index: Optional[int] = field( + default=None, + metadata={ + "help": "P2P destination participant index for decode FFN cluster CC backend (astra_sim_analytical mode).", + "mode_dependency": "pd-af-disaggregation", + }, + ) + + def __post_init__(self): + self._validate_open_source_release_cc_backend_guard() + + # check and set args only in first init (not in Cluster()) + if self.cluster_type is None: + # Early validation based on mode + self._validate_mode_consistency() + + # Basic validation for micro-batch size if provided + if self.decode_attn_micro_batch_size is not None: + assert ( + self.decode_attn_micro_batch_size >= 1 + ), f"decode_attn_micro_batch_size must be >=1, got {self.decode_attn_micro_batch_size}" + + # Ensure micro_batch_size equals batch_size_cap for DECODE_ATTN + # In DECODE_ATTN, micro-batch and batch are semantically equivalent + if ( + self.decode_attn_micro_batch_size is not None + and self.decode_attn_replica_scheduler_config_batch_size_cap is not None + ): + # Both specified: enforce equality + if ( + self.decode_attn_micro_batch_size + != self.decode_attn_replica_scheduler_config_batch_size_cap + ): + raise ValueError( + f"DECODE_ATTN micro_batch_size ({self.decode_attn_micro_batch_size}) " + f"must equal batch_size_cap ({self.decode_attn_replica_scheduler_config_batch_size_cap}). " + f"Reason: In DECODE_ATTN, micro-batch and batch are semantically equivalent." + ) + elif self.decode_attn_micro_batch_size is not None: + # Only micro_batch_size specified: propagate to batch_size_cap + self.decode_attn_replica_scheduler_config_batch_size_cap = ( + self.decode_attn_micro_batch_size + ) + elif self.decode_attn_replica_scheduler_config_batch_size_cap is not None: + # Only batch_size_cap specified: propagate to micro_batch_size + self.decode_attn_micro_batch_size = ( + self.decode_attn_replica_scheduler_config_batch_size_cap + ) + # If neither is specified, keep both as None (use defaults later) + + if self._has_disaggregation_params_set(): + self._setup_disaggregated_configs() + + # Add a check to ensure af_pipeline_num_micro_batch is consistent (only for PD+AF mode) + is_pd_af_mode = ( + self.decode_attn_cluster_num_replicas is not None + and self.decode_ffn_cluster_num_replicas is not None + ) + if is_pd_af_mode: + attn_mb = self.decode_attn_af_pipeline_num_micro_batch + ffn_mb = self.decode_ffn_af_pipeline_num_micro_batch + + assert ( + attn_mb is not None and ffn_mb is not None + ), "In PD+AF disaggregated mode, both decode_attn_af_pipeline_num_micro_batch and decode_ffn_af_pipeline_num_micro_batch must be set." + + assert ( + attn_mb == ffn_mb + ), "The af_pipeline_num_micro_batch must be the same for both decode_attn and decode_ffn clusters." + + # AFD Divisibility Validation (Fail Fast Strategy) + # Aligned with StepFun-vLLM's requirement that batch sizes be divisible by num_stages + # Unlike StepFun which silently rounds down, we fail fast to help users identify + # configuration issues early. + num_stages = attn_mb + if num_stages > 1: + self._validate_afd_divisibility(num_stages) + else: + self._setup_monolithic_config() + + self._validate_prefix_cache_spec_decode_compatibility() + + def _validate_prefix_cache_spec_decode_compatibility(self) -> None: + from frontier.spec_decode.runtime import ( + method_requires_prefix_matching_disabled, + ) + + prefix_enabled = bool( + getattr(self.replica_scheduler_config, "enable_prefix_caching", False) + ) + if not prefix_enabled: + return + + replica_configs = [ + ("replica_config", getattr(self, "replica_config", None)), + ("prefill_replica_config", getattr(self, "prefill_replica_config", None)), + ("decode_replica_config", getattr(self, "decode_replica_config", None)), + ( + "decode_attn_replica_config", + getattr(self, "decode_attn_replica_config", None), + ), + ( + "decode_ffn_replica_config", + getattr(self, "decode_ffn_replica_config", None), + ), + ] + for replica_config_name, replica_config in replica_configs: + if replica_config is None: + continue + spec_decode_config = getattr( + replica_config, "speculative_decoding_config", None + ) + if spec_decode_config is None or not spec_decode_config.enabled: + continue + method = str(getattr(spec_decode_config, "method", "")).strip() + if method and method_requires_prefix_matching_disabled(method): + raise ValueError( + "Speculative decoding method " + f"{method!r} requires prefix caching to be disabled, " + f"but replica_scheduler_config.enable_prefix_caching=True " + f"for {replica_config_name}." + ) + + def _validate_afd_divisibility(self, num_stages: int): + """Validate that key batch size parameters are divisible by num_stages. + + Currently relaxed — no divisibility enforcement. + """ + return None + + def _validate_mode_consistency(self): + """Validate that configuration is consistent with the intended mode.""" + + has_disaggregated_fields = ( + self.prefill_cluster_num_replicas is not None + or self.decode_attn_cluster_num_replicas is not None + or self.decode_ffn_cluster_num_replicas is not None + or self.decode_cluster_num_replicas is not None + ) + + # The `num_replicas` field is exclusively for monolithic mode. + # `replica_config` can be used as a template in disaggregated mode, so its presence is not a conflict. + has_monolithic_exclusive_field = self.num_replicas is not None + + if has_disaggregated_fields and has_monolithic_exclusive_field: + logger.warning( + "Both disaggregated and monolithic configuration fields are set. " + "The 'num_replicas' field (for monolithic mode) was provided but will be ignored in disaggregated mode. " + "Please use cluster-specific replica counts like 'prefill_cluster_num_replicas'." + ) + + def _validate_open_source_release_cc_backend_guard(self) -> None: + from frontier.cc_backend.cc_backend_config import AiconfiguratorCCBackendConfig + + if isinstance(self.cc_backend_config, AiconfiguratorCCBackendConfig): + raise ValueError(AICONFIGURATOR_BACKEND_RELEASE_ERROR) + + def _setup_monolithic_config(self): + """Setup configuration for monolithic (co-location) mode.""" + # Ensure required fields are set for monolithic mode + assert self.num_replicas != None, "Num replicas must be set" + assert self.replica_config != None, "Replica config must be set" + + # Set cluster type for monolithic mode + self.cluster_type = ClusterType.MONOLITHIC + + # Predictor routing details are indexed by serving Replica identity. + # Keep that capacity dimension explicit instead of deriving it from + # attention-DP lanes. + self.replica_config.cluster_num_replicas = int(self.num_replicas) + + # Reuse the same parallel-domain validation used by disaggregated clusters so + # monolithic MoE layouts fail fast when attention and MoE domains disagree. + self._validate_replica_config(self.replica_config, "monolithic") + self.world_size = self.replica_config.world_size * self.num_replicas + + # Clear disaggregated fields to avoid confusion + self.prefill_replica_config = None + self.decode_attn_replica_config = None + self.decode_ffn_replica_config = None + self.prefill_cluster_num_replicas = None + self.decode_attn_cluster_num_replicas = None + self.decode_ffn_cluster_num_replicas = None + + # else: + # assert self.replica_config.expert_parallel_size == self.replica_config.tensor_parallel_size, "For local MoE, expert_parallel_size must be equal to tensor_parallel_size" + + def _setup_disaggregated_configs(self): + """Setup configuration for disaggregated mode (PD or PD+AF).""" + # Clear monolithic fields since they're not used + # self.replica_config = None + self.num_replicas = None + + # Determine disaggregation mode + is_pd_af_mode = ( + self.decode_attn_cluster_num_replicas is not None + and self.decode_ffn_cluster_num_replicas is not None + ) + is_pd_mode = self.decode_cluster_num_replicas is not None + + # Ensure required disaggregated fields are set + assert ( + self.prefill_cluster_num_replicas != None + ), "Prefill cluster num replicas must be set" + + # Requirement 10.4: Validate that replica counts are positive + if ( + self.prefill_cluster_num_replicas is not None + and self.prefill_cluster_num_replicas <= 0 + ): + raise ValueError( + f"prefill_cluster_num_replicas must be positive, got {self.prefill_cluster_num_replicas}" + ) + + if is_pd_af_mode: + # PD+AF disaggregation mode + assert ( + self.decode_attn_cluster_num_replicas != None + ), "Decode attention cluster num replicas must be set" + assert ( + self.decode_ffn_cluster_num_replicas != None + ), "Decode FFN cluster num replicas must be set" + assert ( + not is_pd_mode + ), "Cannot set both PD and PD+AF disaggregation parameters" + + # Requirement 10.4: Validate that replica counts are positive (PD+AF mode) + if self.decode_attn_cluster_num_replicas <= 0: + raise ValueError( + f"decode_attn_cluster_num_replicas must be positive, got {self.decode_attn_cluster_num_replicas}" + ) + if self.decode_ffn_cluster_num_replicas <= 0: + raise ValueError( + f"decode_ffn_cluster_num_replicas must be positive, got {self.decode_ffn_cluster_num_replicas}" + ) + + # DECODE_FFN grouping semantics are implemented only in the + # RoundRobinClusterScheduler path. + cluster_scheduler_type = self.cluster_scheduler_config.get_type() + if cluster_scheduler_type != ClusterSchedulerType.ROUND_ROBIN: + raise ValueError( + "PD+AF mode requires RoundRobin cluster scheduler when DECODE_FFN is enabled. " + f"Got cluster_scheduler_config_type={cluster_scheduler_type}." + ) + + for field_name, decode_role in ( + ( + "decode_attn_replica_config_num_pipeline_stages", + "decode_attn", + ), + ( + "decode_ffn_replica_config_num_pipeline_stages", + "decode_ffn", + ), + ): + pipeline_stages = getattr(self, field_name) + if pipeline_stages not in (None, 1): + raise ValueError( + f"{field_name} must be 1 for {decode_role}, " + f"got {pipeline_stages}." + ) + + # Create ReplicaConfig objects from flattened fields + self.prefill_replica_config = self._create_replica_config_from_fields( + "prefill" + ) + self.decode_attn_replica_config = self._create_replica_config_from_fields( + "decode_attn" + ) + self.decode_ffn_replica_config = self._create_replica_config_from_fields( + "decode_ffn" + ) + + # Enforce disaggregation constraints + assert ( + self.decode_attn_replica_config.num_pipeline_stages == 1 + ), "Decode attention cluster must have 1 pipeline stage" + assert ( + self.decode_ffn_replica_config.num_pipeline_stages == 1 + ), "Decode FFN cluster must have 1 pipeline stage" + + # Validate each cluster + self._validate_replica_config(self.prefill_replica_config, "prefill") + self._validate_replica_config( + self.decode_attn_replica_config, "decode_attn" + ) + self._validate_replica_config(self.decode_ffn_replica_config, "decode_ffn") + + # Calculate world sizes + self.prefill_world_size = ( + self.prefill_cluster_num_replicas + * self.prefill_replica_config.world_size + ) + self.decode_attn_world_size = ( + self.decode_attn_cluster_num_replicas + * self.decode_attn_replica_config.world_size + ) + self.decode_ffn_world_size = ( + self.decode_ffn_cluster_num_replicas + * self.decode_ffn_replica_config.world_size + ) + self.world_size = ( + self.prefill_world_size + + self.decode_attn_world_size + + self.decode_ffn_world_size + ) + + elif is_pd_mode: + # PD disaggregation mode + assert ( + self.decode_cluster_num_replicas != None + ), "Decode cluster num replicas must be set" + assert ( + not is_pd_af_mode + ), "Cannot set both PD and PD+AF disaggregation parameters" + + # Requirement 10.4: Validate that replica counts are positive (PD mode) + if self.decode_cluster_num_replicas <= 0: + raise ValueError( + f"decode_cluster_num_replicas must be positive, got {self.decode_cluster_num_replicas}" + ) + + # Create ReplicaConfig objects from flattened fields + self.prefill_replica_config = self._create_replica_config_from_fields( + "prefill" + ) + self.decode_replica_config = self._create_replica_config_from_fields( + "decode" + ) + + # Validate each cluster + self._validate_replica_config(self.prefill_replica_config, "prefill") + self._validate_replica_config(self.decode_replica_config, "decode") + + # Calculate world sizes + self.prefill_world_size = ( + self.prefill_cluster_num_replicas + * self.prefill_replica_config.world_size + ) + self.decode_world_size = ( + self.decode_cluster_num_replicas * self.decode_replica_config.world_size + ) + self.world_size = self.prefill_world_size + self.decode_world_size + + else: + raise ValueError( + "Invalid disaggregation configuration: must set either PD or PD+AF parameters" + ) + + # Ensure consistent dummy mode configuration across all clusters + self._ensure_consistent_dummy_mode() + + print(f"Total world size: {self.world_size}") + + def _ensure_consistent_dummy_mode(self): + """Ensure all clusters use the same dummy mode configuration.""" + # Get the main execution_time_predictor_config dummy mode settings + main_config = self.execution_time_predictor_config + main_dummy_mode = main_config.enable_dummy_mode + main_dummy_time = main_config.dummy_execution_time_ms + + if main_dummy_mode: + print(f"Applying dummy mode (time={main_dummy_time}ms) to all clusters") + + # Apply dummy mode settings to all cluster configs + self.execution_time_predictor_config.enable_dummy_mode = True + self.execution_time_predictor_config.dummy_execution_time_ms = ( + main_dummy_time + ) + + # Note: In the current architecture, all clusters share the same execution_time_predictor_config + # This ensures consistency across all clusters + + def _field_is_set_to_non_default(self, field_def) -> bool: + value = getattr(self, field_def.name) + if field_def.default is not MISSING: + return value != field_def.default + if field_def.default_factory is not MISSING: + return value != field_def.default_factory() + return value is not None + + def _has_disaggregation_params_set(self) -> bool: + """Check if any disaggregation-specific cluster fields have been set.""" + for field_def in self.__dataclass_fields__.values(): + if ( + not field_def.name.startswith(DISAGGREGATED_CLUSTER_FIELD_PREFIXES) + and field_def.name not in DISAGGREGATED_CLUSTER_FIELD_NAMES + ): + continue + if self._field_is_set_to_non_default(field_def): + return True + return False + + def _validate_replica_config( + self, replica_config: ReplicaConfig, cluster_name: str + ): + """Validate replica configuration for specific cluster.""" + # Validate pipeline parallelism configuration (double-check, should already be validated in __post_init__) + if ( + replica_config.model_config.num_layers % replica_config.num_pipeline_stages + != 0 + ): + raise ValueError( + f"Pipeline parallelism configuration error in {cluster_name} cluster: " + f"num_layers ({replica_config.model_config.num_layers}) must be evenly divisible by " + f"num_pipeline_stages ({replica_config.num_pipeline_stages}). " + f"Current configuration would result in uneven layer distribution across pipeline stages." + ) + + # Validate dense model configuration in disaggregated modes. + is_dense_model = not replica_config.model_config.is_moe + if is_dense_model and cluster_name in ["prefill", "decode"]: + # For dense models in PD-disaggregation mode, enforce attn_dp = 1 + if replica_config.attn_dp != 1: + raise ValueError( + f"Dense models in PD-disaggregation mode require attn_dp=1 " + f"in {cluster_name} cluster, got {replica_config.attn_dp}. " + f"Dense models do not support attn data parallelism in disaggregated mode." + ) + # Ensure MoE parallelism is disabled for dense models + if replica_config.moe_expert_parallel_size != 1: + raise ValueError( + f"Dense models require moe_expert_parallel_size=1 in {cluster_name} cluster, " + f"got {replica_config.moe_expert_parallel_size}. " + f"Dense models do not have expert parallelism." + ) + + if is_dense_model and cluster_name == "decode_attn": + if replica_config.attn_dp != 1: + raise ValueError( + "Dense models require attn_dp=1 in " + f"{cluster_name} cluster, got " + f"{replica_config.attn_dp}." + ) + + if is_dense_model and cluster_name == "decode_ffn": + if replica_config.moe_expert_parallel_size != 1: + raise ValueError( + "Dense models require moe_expert_parallel_size=1 in " + f"{cluster_name} cluster, got " + f"{replica_config.moe_expert_parallel_size}." + ) + if replica_config.router_topk != 1: + raise ValueError( + f"Dense models require router_topk=1 in {cluster_name} " + f"cluster, got {replica_config.router_topk}." + ) + + normalized_cluster_name = str(cluster_name).strip().lower() + if normalized_cluster_name in {"prefill", "decode", "monolithic"} and replica_config.model_config.is_moe: + validate_frontier_shared_parallel_domains( + FrontierParallelismMapping( + cluster_num_replicas=1, + attn_tensor_parallel_size=replica_config.attn_tensor_parallel_size, + attn_dp=replica_config.attn_dp, + moe_tensor_parallel_size=replica_config.moe_tensor_parallel_size, + moe_expert_parallel_size=replica_config.moe_expert_parallel_size, + ) + ) + + if cluster_name != "decode_attn": + pass + # else: + # assert replica_config.moe_expert_parallel_size == replica_config.moe_tensor_parallel_size, f"For local MoE in {cluster_name} cluster, moe_expert_parallel_size must be equal to moe_tensor_parallel_size" + else: + assert ( + replica_config.moe_expert_parallel_size == 0 + and replica_config.local_expert_num == 0 + ), "For decode attention cluster, moe_expert_parallel_size and local_expert_num must be 0" diff --git a/frontier/config/cluster_role_config.py b/frontier/config/cluster_role_config.py new file mode 100644 index 00000000..2393949c --- /dev/null +++ b/frontier/config/cluster_role_config.py @@ -0,0 +1,603 @@ +"""Per-role configuration construction for a disaggregated cluster. + +`ClusterConfig` owns a single flat field surface covering every role the +release supports. The methods here turn that flat surface into the concrete +`ReplicaConfig`, execution-time predictor configuration and communication-cost +backend configuration each role needs. + +They are a mixin rather than free functions so that the split is exactly a +move: the bodies, the method names and every call site are unchanged. +""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple + +from frontier.config.execution_time_predictor_config import ( + BaseExecutionTimePredictorConfig, +) +from frontier.config.replica_config import ReplicaConfig +from frontier.types import ClusterType + +if TYPE_CHECKING: + from frontier.cc_backend.cc_backend_config import BaseCCBackendConfig + from frontier.config.cluster_config import ClusterConfig + + +def _get_cc_backend_configs(): + """Lazily import CC backend config classes to avoid circular imports.""" + from frontier.cc_backend.cc_backend_config import ( + BaseCCBackendConfig, + VidurCCBackendConfig, + AnalyticalCCBackendConfig, + CollectiveSimCCBackendConfig, + AiconfiguratorCCBackendConfig, + AstraSimAnalyticalCCBackendConfig, + ) + + return ( + BaseCCBackendConfig, + VidurCCBackendConfig, + AnalyticalCCBackendConfig, + CollectiveSimCCBackendConfig, + AiconfiguratorCCBackendConfig, + AstraSimAnalyticalCCBackendConfig, + ) + + +class ClusterRoleConfigBuilder: + """Builds the per-role configurations a disaggregated cluster needs.""" + + def _create_replica_config_from_fields(self, cluster_prefix: str) -> ReplicaConfig: + """Create ReplicaConfig object from flattened fields.""" + # Get default values from main replica_config or use ReplicaConfig defaults + main_config = self.replica_config if self.replica_config else ReplicaConfig() + + # Extract cluster-specific fields using getattr with fallback to main config + def get_field_value(field_name: str): + cluster_field_name = f"{cluster_prefix}_replica_config_{field_name}" + cluster_value = getattr(self, cluster_field_name, None) + if cluster_value is not None: + return cluster_value + return getattr(main_config, field_name) + + if cluster_prefix == "decode_attn": + moe_expert_parallel_size = 0 + moe_tensor_parallel_size = 0 + total_expert_num = 0 + local_expert_num = 0 + num_pipeline_stages = 1 + router_load_balancing_type = None + router_topk = None + moe_routing_distribution_type = get_field_value( + "moe_routing_distribution_type" + ) + attn_tensor_parallel_size = get_field_value("attn_tensor_parallel_size") + attn_dp = get_field_value("attn_dp") + else: + attn_tensor_parallel_size = get_field_value("attn_tensor_parallel_size") + attn_dp = get_field_value("attn_dp") + moe_tensor_parallel_size = get_field_value("moe_tensor_parallel_size") + moe_expert_parallel_size = get_field_value("moe_expert_parallel_size") + total_expert_num = get_field_value("total_expert_num") + local_expert_num = get_field_value("local_expert_num") + router_load_balancing_type = get_field_value("router_load_balancing_type") + router_topk = get_field_value("router_topk") + moe_routing_distribution_type = get_field_value( + "moe_routing_distribution_type" + ) + if cluster_prefix == "decode_ffn": + num_pipeline_stages = 1 + else: + num_pipeline_stages = get_field_value("num_pipeline_stages") + + return ReplicaConfig( + model_name=get_field_value("model_name"), + memory_margin_fraction=get_field_value("memory_margin_fraction"), + num_pipeline_stages=num_pipeline_stages, + attn_tensor_parallel_size=attn_tensor_parallel_size, + attn_dp=attn_dp, + moe_tensor_parallel_size=moe_tensor_parallel_size, + moe_expert_parallel_size=moe_expert_parallel_size, + total_expert_num=total_expert_num, + local_expert_num=local_expert_num, + router_load_balancing_type=router_load_balancing_type, + router_topk=router_topk, + moe_routing_seed=get_field_value("moe_routing_seed"), + moe_routing_trace_path=get_field_value("moe_routing_trace_path"), + decode_attn_initial_lane_trace_path=get_field_value( + "decode_attn_initial_lane_trace_path" + ), + decode_attn_steady_state_snapshot_path=get_field_value( + "decode_attn_steady_state_snapshot_path" + ), + decode_attn_steady_state_measurement_report_path=get_field_value( + "decode_attn_steady_state_measurement_report_path" + ), + moe_routing_distribution_type=moe_routing_distribution_type, + device=get_field_value("device"), + network_device=get_field_value("network_device"), + cluster_prefix=cluster_prefix, + speculative_decoding_config=main_config.speculative_decoding_config, + ) + + def get_cluster_configs_for_disaggregation( + self, + ) -> Dict[ClusterType, "ClusterConfig"]: + """Generate cluster configurations for disaggregated mode.""" + # Lazy import: ClusterConfig inherits this mixin, so importing it at + # module level would close a cycle. + from frontier.config.cluster_config import ClusterConfig + + if not self._has_disaggregation_params_set(): + return {ClusterType.MONOLITHIC: self} + + cluster_configs = {} + + # Prefill cluster + if self.prefill_cluster_num_replicas: + prefill_config = ClusterConfig( + cluster_type=ClusterType.PREFILL, + num_replicas=self.prefill_cluster_num_replicas, + replica_config=self.prefill_replica_config + or self._create_replica_config_copy(), + cluster_scheduler_config=self.cluster_scheduler_config, + replica_scheduler_config=self.replica_scheduler_config, + execution_time_predictor_config=( + self._create_execution_time_predictor_config_for_cluster( + "prefill" + ) + ), + cc_backend_config=self._create_cc_backend_config_for_cluster("prefill"), + # Propagate cluster-specific replica scheduler config parameters + prefill_replica_scheduler_config_type=self.prefill_replica_scheduler_config_type, + prefill_replica_scheduler_config_batch_size_cap=self.prefill_replica_scheduler_config_batch_size_cap, + prefill_replica_scheduler_config_max_tokens_in_batch=self.prefill_replica_scheduler_config_max_tokens_in_batch, + prefill_replica_scheduler_config_enable_chunked_prefill=self.prefill_replica_scheduler_config_enable_chunked_prefill, + prefill_replica_scheduler_config_long_prefill_token_threshold=self.prefill_replica_scheduler_config_long_prefill_token_threshold, + prefill_replica_scheduler_config_num_blocks=self.prefill_replica_scheduler_config_num_blocks, + prefill_replica_scheduler_config_block_size=self.prefill_replica_scheduler_config_block_size, + prefill_replica_scheduler_config_watermark_blocks_fraction=self.prefill_replica_scheduler_config_watermark_blocks_fraction, + ) + cluster_configs[ClusterType.PREFILL] = prefill_config + + # Decode Attention cluster + if self.decode_attn_cluster_num_replicas: + # Determine micro-batch SIZE for decode-attn + # NOTE: Only decode_attn_micro_batch_size exists (no generic fallback) + _da_mbs = self.decode_attn_micro_batch_size + decode_attn_config = ClusterConfig( + cluster_type=ClusterType.DECODE_ATTN, + num_replicas=self.decode_attn_cluster_num_replicas, + replica_config=self.decode_attn_replica_config + or self._create_replica_config_copy(), + cluster_scheduler_config=self.cluster_scheduler_config, + replica_scheduler_config=self.replica_scheduler_config, + execution_time_predictor_config=( + self._create_execution_time_predictor_config_for_cluster( + "decode_attn" + ) + ), + cc_backend_config=self._create_cc_backend_config_for_cluster( + "decode_attn" + ), + af_pipeline_num_micro_batch=self.decode_attn_af_pipeline_num_micro_batch, + decode_attn_micro_batch_size=_da_mbs, + decode_attn_request_allocation_threshold=self.decode_attn_request_allocation_threshold, + # Propagate cluster-specific replica scheduler config parameters + decode_attn_replica_scheduler_config_type=self.decode_attn_replica_scheduler_config_type, + decode_attn_replica_scheduler_config_batch_size_cap=self.decode_attn_replica_scheduler_config_batch_size_cap, + decode_attn_replica_scheduler_config_max_tokens_in_batch=self.decode_attn_replica_scheduler_config_max_tokens_in_batch, + decode_attn_replica_scheduler_config_num_blocks=self.decode_attn_replica_scheduler_config_num_blocks, + decode_attn_replica_scheduler_config_block_size=self.decode_attn_replica_scheduler_config_block_size, + decode_attn_replica_scheduler_config_watermark_blocks_fraction=self.decode_attn_replica_scheduler_config_watermark_blocks_fraction, + ) + cluster_configs[ClusterType.DECODE_ATTN] = decode_attn_config + + # Decode FFN cluster + if self.decode_ffn_cluster_num_replicas: + decode_ffn_config = ClusterConfig( + cluster_type=ClusterType.DECODE_FFN, + num_replicas=self.decode_ffn_cluster_num_replicas, + replica_config=self.decode_ffn_replica_config + or self._create_replica_config_copy(), + cluster_scheduler_config=self.cluster_scheduler_config, + replica_scheduler_config=self.replica_scheduler_config, + execution_time_predictor_config=( + self._create_execution_time_predictor_config_for_cluster( + "decode_ffn" + ) + ), + cc_backend_config=self._create_cc_backend_config_for_cluster( + "decode_ffn" + ), + af_pipeline_num_micro_batch=self.decode_ffn_af_pipeline_num_micro_batch, + # Propagate cluster-specific replica scheduler config parameters + decode_ffn_replica_scheduler_config_type=self.decode_ffn_replica_scheduler_config_type, + decode_ffn_replica_scheduler_config_batch_size_cap=self.decode_ffn_replica_scheduler_config_batch_size_cap, + decode_ffn_replica_scheduler_config_max_tokens_in_batch=self.decode_ffn_replica_scheduler_config_max_tokens_in_batch, + decode_ffn_replica_scheduler_config_num_blocks=self.decode_ffn_replica_scheduler_config_num_blocks, + decode_ffn_replica_scheduler_config_block_size=self.decode_ffn_replica_scheduler_config_block_size, + decode_ffn_replica_scheduler_config_watermark_blocks_fraction=self.decode_ffn_replica_scheduler_config_watermark_blocks_fraction, + decode_attn_cluster_num_replicas=self.decode_attn_cluster_num_replicas, + ) + # Propagate only the source Attention-Replica capacity. AFD + # grouping is Replica-level; it must not manufacture DP lanes. + decode_ffn_config.decode_attn_replica_id_start_for_ffn = int( + self.prefill_cluster_num_replicas + ) + + cluster_configs[ClusterType.DECODE_FFN] = decode_ffn_config + + # Unified Decode cluster (PD-disaggregation mode) + if self.decode_cluster_num_replicas: + decode_config = ClusterConfig( + cluster_type=ClusterType.DECODE, + num_replicas=self.decode_cluster_num_replicas, + replica_config=self.decode_replica_config + or self._create_replica_config_copy(), + cluster_scheduler_config=self.cluster_scheduler_config, + replica_scheduler_config=self.replica_scheduler_config, + execution_time_predictor_config=( + self._create_execution_time_predictor_config_for_cluster("decode") + ), + cc_backend_config=self._create_cc_backend_config_for_cluster("decode"), + # Propagate cluster-specific replica scheduler config parameters + decode_replica_scheduler_config_type=self.decode_replica_scheduler_config_type, + decode_replica_scheduler_config_batch_size_cap=self.decode_replica_scheduler_config_batch_size_cap, + decode_replica_scheduler_config_max_tokens_in_batch=self.decode_replica_scheduler_config_max_tokens_in_batch, + decode_replica_scheduler_config_num_blocks=self.decode_replica_scheduler_config_num_blocks, + decode_replica_scheduler_config_block_size=self.decode_replica_scheduler_config_block_size, + decode_replica_scheduler_config_watermark_blocks_fraction=self.decode_replica_scheduler_config_watermark_blocks_fraction, + ) + cluster_configs[ClusterType.DECODE] = decode_config + + return cluster_configs + + def _create_execution_time_predictor_config_for_cluster( + self, cluster_prefix: str + ) -> BaseExecutionTimePredictorConfig: + """Create cluster-specific execution-time predictor config overrides.""" + base_config = self.execution_time_predictor_config + override_values = {} + for calibration_field in ( + "attn_pre_proj_calibration_scale", + "attn_post_proj_calibration_scale", + "attn_decode_calibration_scale", + "attn_kv_cache_save_calibration_scale", + "mlp_up_proj_calibration_scale", + "mlp_down_proj_calibration_scale", + "decode_phase_mlp_down_proj_calibration_scale", + ): + override_field = ( + f"{cluster_prefix}_execution_time_predictor_config_" + f"{calibration_field}" + ) + override_value = getattr(self, override_field, None) + if override_value is None: + continue + override_value = float(override_value) + if override_value <= 0.0: + raise ValueError( + f"ClusterConfig.{override_field} must be > 0, got={override_value!r}" + ) + override_values[calibration_field] = override_value + + if not override_values: + return base_config + + return replace(base_config, **override_values) + + def _create_cc_backend_config_for_cluster( + self, cluster_prefix: str + ) -> BaseCCBackendConfig: + """ + Create CC backend configuration for a specific cluster. + + This method creates a cluster-specific CC backend configuration by: + 1. Checking for cluster-specific override values + 2. Falling back to the base cc_backend_config values if not overridden + + Args: + cluster_prefix: Cluster prefix (e.g., "prefill", "decode", "decode_attn", "decode_ffn") + + Returns: + CC backend configuration for the specified cluster + """ + creators = self._cc_backend_creators() + + cluster_type_str = getattr( + self, f"{cluster_prefix}_cc_backend_config_type", None + ) + if cluster_type_str is not None: + cluster_type_key = cluster_type_str.lower() + for type_key, _, create in creators: + if type_key == cluster_type_key: + return create(cluster_prefix) + raise ValueError(f"Unknown CC backend type: {cluster_type_str}") + + base_config = self.cc_backend_config + for _, config_class, create in creators: + if isinstance(base_config, config_class): + return create(cluster_prefix) + raise ValueError( + "Unsupported base CC backend config type for cluster-specific " + f"construction: {type(base_config).__name__}" + ) + + def _cc_backend_creators( + self, + ) -> Tuple[Tuple[str, type, Callable[[str], "BaseCCBackendConfig"]], ...]: + """Return the supported CC backends as (type key, config class, creator). + + One ordered table serves both selection paths: an explicit + ``_cc_backend_config_type`` string, and, when that is absent, + the concrete class of the base ``cc_backend_config``. + """ + # Lazy import to avoid circular imports + ( + _, + VidurCCBackendConfig, + AnalyticalCCBackendConfig, + CollectiveSimCCBackendConfig, + AiconfiguratorCCBackendConfig, + AstraSimAnalyticalCCBackendConfig, + ) = _get_cc_backend_configs() + + return ( + ( + "analytical", + AnalyticalCCBackendConfig, + self._create_analytical_cc_backend_config, + ), + ("vidur", VidurCCBackendConfig, self._create_vidur_cc_backend_config), + ( + "collective_sim", + CollectiveSimCCBackendConfig, + self._create_collective_sim_cc_backend_config, + ), + ( + "aiconfigurator", + AiconfiguratorCCBackendConfig, + self._create_aiconfigurator_cc_backend_config, + ), + ( + "astra_sim_analytical", + AstraSimAnalyticalCCBackendConfig, + self._create_astra_sim_analytical_cc_backend_config, + ), + ) + + def _cc_backend_value_reader( + self, cluster_prefix: str, base_config: "BaseCCBackendConfig", config_class: type + ) -> Callable[[str, Any], Any]: + """Return a reader resolving one backend field for a cluster. + + Resolution order is the cluster-specific flat field, then the base + configuration when it already is of this backend type, then the default + the caller supplies. + """ + + def get_value(field_name: str, default_value: Any) -> Any: + cluster_value = getattr( + self, f"{cluster_prefix}_cc_backend_config_{field_name}", None + ) + if cluster_value is not None: + return cluster_value + if isinstance(base_config, config_class): + return getattr(base_config, field_name, default_value) + return default_value + + return get_value + + @staticmethod + + def _shared_cc_backend_fields(base_config: "BaseCCBackendConfig") -> Dict[str, Any]: + """Return the fields every CC backend inherits from the base config.""" + return { + "profiling_data_dir": getattr( + base_config, "profiling_data_dir", "data/profiling/network" + ), + "cache_dir": getattr(base_config, "cache_dir", "cache"), + "no_cache": getattr(base_config, "no_cache", False), + } + + def _create_analytical_cc_backend_config( + self, cluster_prefix: str + ) -> AnalyticalCCBackendConfig: + """Create analytical CC backend config with cluster-specific overrides.""" + # Lazy import to avoid circular imports + _, _, AnalyticalCCBackendConfig, _, _, _ = _get_cc_backend_configs() + + base_config = self.cc_backend_config + + get_value = self._cc_backend_value_reader( + cluster_prefix, base_config, AnalyticalCCBackendConfig + ) + + return AnalyticalCCBackendConfig( + **self._shared_cc_backend_fields(base_config), + network_bandwidth_gbps=get_value("network_bandwidth_gbps", 100.0), + network_latency_us=get_value("network_latency_us", 1.0), + intra_node_bandwidth_gbps=get_value("intra_node_bandwidth_gbps", 600.0), + ) + + def _create_vidur_cc_backend_config( + self, cluster_prefix: str + ) -> VidurCCBackendConfig: + """Create Vidur CC backend config with cluster-specific overrides.""" + # Lazy import to avoid circular imports + _, VidurCCBackendConfig, _, _, _, _ = _get_cc_backend_configs() + + base_config = self.cc_backend_config + + # For Vidur config, we mainly use the base config values + # as Vidur-specific parameters are typically shared across clusters + if isinstance(base_config, VidurCCBackendConfig): + return VidurCCBackendConfig( + profiling_data_dir=base_config.profiling_data_dir, + cache_dir=base_config.cache_dir, + no_cache=base_config.no_cache, + all_reduce_input_file=base_config.all_reduce_input_file, + send_recv_input_file=base_config.send_recv_input_file, + k_fold_cv_splits=base_config.k_fold_cv_splits, + num_training_job_threads=base_config.num_training_job_threads, + ) + else: + # Create default Vidur config + return VidurCCBackendConfig() + + def _create_collective_sim_cc_backend_config( + self, cluster_prefix: str + ) -> "CollectiveSimCCBackendConfig": + """Create collective-sim CC backend config with cluster-specific overrides.""" + ( + _, + _, + _, + CollectiveSimCCBackendConfig, + _, + _, + ) = _get_cc_backend_configs() + + from pathlib import Path + + base_config = self.cc_backend_config + if not isinstance(base_config, CollectiveSimCCBackendConfig): + base_config = CollectiveSimCCBackendConfig() + + get_value = self._cc_backend_value_reader( + cluster_prefix, base_config, CollectiveSimCCBackendConfig + ) + + if base_config.runner_out_dir: + cluster_out_dir = str(Path(base_config.runner_out_dir) / cluster_prefix) + return replace( + base_config, + runner_out_dir=cluster_out_dir, + nvlink_allreduce_launch_overhead_us=get_value( + "nvlink_allreduce_launch_overhead_us", + 50.0, + ), + ) + + return replace( + base_config, + nvlink_allreduce_launch_overhead_us=get_value( + "nvlink_allreduce_launch_overhead_us", + 50.0, + ), + ) + + def _create_aiconfigurator_cc_backend_config( + self, cluster_prefix: str + ) -> "AiconfiguratorCCBackendConfig": + """Create aiconfigurator CC backend config with cluster-specific overrides.""" + ( + _, + _, + _, + _, + AiconfiguratorCCBackendConfig, + _, + ) = _get_cc_backend_configs() + + base_config = self.cc_backend_config + + get_value = self._cc_backend_value_reader( + cluster_prefix, base_config, AiconfiguratorCCBackendConfig + ) + + return AiconfiguratorCCBackendConfig( + **self._shared_cc_backend_fields(base_config), + repo_root=get_value("repo_root", "sota-infer-engine/aiconfigurator"), + system=get_value("system", ""), + source_backend=get_value("source_backend", "vllm"), + source_version=get_value("source_version", ""), + database_mode=get_value("database_mode", "silicon"), + tp_allreduce_impl=get_value("tp_allreduce_impl", "custom_allreduce"), + custom_allreduce_variant=get_value("custom_allreduce_variant", None), + ) + + def _create_astra_sim_analytical_cc_backend_config( + self, cluster_prefix: str + ) -> "AstraSimAnalyticalCCBackendConfig": + """Create astra-sim analytical CC backend config with cluster-specific overrides.""" + ( + _, + _, + _, + _, + _, + AstraSimAnalyticalCCBackendConfig, + ) = _get_cc_backend_configs() + + base_config = self.cc_backend_config + + get_value = self._cc_backend_value_reader( + cluster_prefix, base_config, AstraSimAnalyticalCCBackendConfig + ) + + return AstraSimAnalyticalCCBackendConfig( + **self._shared_cc_backend_fields(base_config), + prediction_cache_size=get_value("prediction_cache_size", 4096), + placement_order=get_value("placement_order", "TP,CP,DP,EP"), + intra_server_topology=get_value( + "intra_server_topology", "FullyConnected" + ), + inter_server_topology=get_value( + "inter_server_topology", "FullyConnected" + ), + intra_server_bandwidth_gbps=get_value( + "intra_server_bandwidth_gbps", 600.0 + ), + intra_server_latency_us=get_value("intra_server_latency_us", 1.0), + inter_server_bandwidth_gbps=get_value( + "inter_server_bandwidth_gbps", 100.0 + ), + inter_server_latency_us=get_value("inter_server_latency_us", 1.0), + ring_bidirectional=( + base_config.ring_bidirectional + if isinstance(base_config, AstraSimAnalyticalCCBackendConfig) + else True + ), + p2p_src_index=get_value("p2p_src_index", 0), + p2p_dst_index=get_value("p2p_dst_index", 1), + ) + + def _create_replica_config_copy(self) -> ReplicaConfig: + """Create a copy of the main replica config for disaggregated clusters.""" + # Note: This method now needs to be called before replica_config is cleared + # We need to preserve the original config temporarily + original_config = ( + self.replica_config if self.replica_config else ReplicaConfig() + ) + + return ReplicaConfig( + model_name=original_config.model_name, + memory_margin_fraction=original_config.memory_margin_fraction, + num_pipeline_stages=original_config.num_pipeline_stages, + attn_tensor_parallel_size=original_config.attn_tensor_parallel_size, + attn_dp=original_config.attn_dp, + moe_tensor_parallel_size=original_config.moe_tensor_parallel_size, + moe_expert_parallel_size=original_config.moe_expert_parallel_size, + total_expert_num=original_config.total_expert_num, + router_load_balancing_type=original_config.router_load_balancing_type, + router_topk=original_config.router_topk, + moe_routing_seed=original_config.moe_routing_seed, + moe_routing_distribution_type=original_config.moe_routing_distribution_type, + moe_routing_trace_path=original_config.moe_routing_trace_path, + decode_attn_initial_lane_trace_path=( + original_config.decode_attn_initial_lane_trace_path + ), + decode_attn_steady_state_snapshot_path=( + original_config.decode_attn_steady_state_snapshot_path + ), + decode_attn_steady_state_measurement_report_path=( + original_config.decode_attn_steady_state_measurement_report_path + ), + device=original_config.device, + network_device=original_config.network_device, + speculative_decoding_config=original_config.speculative_decoding_config, + ) diff --git a/frontier/config/cluster_scheduler_config.py b/frontier/config/cluster_scheduler_config.py new file mode 100644 index 00000000..d757a1c4 --- /dev/null +++ b/frontier/config/cluster_scheduler_config.py @@ -0,0 +1,48 @@ +"""Cluster scheduler configuration: how a cluster picks a replica.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from frontier.config.base_poly_config import BasePolyConfig +from frontier.types import ClusterSchedulerType + + +@dataclass +class BaseClusterSchedulerConfig(BasePolyConfig): + pass + + +@dataclass +class RandomClusterSchedulerConfig(BaseClusterSchedulerConfig): + @staticmethod + def get_type(): + return ClusterSchedulerType.RANDOM + + +@dataclass +class RoundRobinClusterSchedulerConfig(BaseClusterSchedulerConfig): + @staticmethod + def get_type(): + return ClusterSchedulerType.ROUND_ROBIN + + +@dataclass +class LORClusterSchedulerConfig(BaseClusterSchedulerConfig): + @staticmethod + def get_type(): + return ClusterSchedulerType.LOR + + +@dataclass +class StickyRoundRobinClusterSchedulerConfig(BaseClusterSchedulerConfig): + @staticmethod + def get_type(): + return ClusterSchedulerType.STICKY_ROUND_ROBIN + + +@dataclass +class StickyLORClusterSchedulerConfig(BaseClusterSchedulerConfig): + @staticmethod + def get_type(): + return ClusterSchedulerType.STICKY_LOR diff --git a/frontier/config/cluster_topology_summary.py b/frontier/config/cluster_topology_summary.py new file mode 100644 index 00000000..24ab6e78 --- /dev/null +++ b/frontier/config/cluster_topology_summary.py @@ -0,0 +1,219 @@ +"""Cluster topology accounting: replica inventory, server counts, statistics. + +These methods read already-resolved cluster configurations and report on them. +They neither construct nor validate configuration, which is why they sit apart +from both the field surface and the per-role builders. +""" + +from __future__ import annotations + +from typing import Dict, List, Tuple + +from frontier.config.cluster_role_config import _get_cc_backend_configs +from frontier.config.parallel_semantics import ( + resolve_collective_sim_physical_topology, +) +from frontier.config.replica_config import ReplicaConfig + + +class ClusterTopologySummary: + """Reports the replica inventory and server counts of a cluster.""" + + def _collect_cluster_info(self) -> List[Tuple[str, int, ReplicaConfig]]: + clusters_info = [] + + if self._has_disaggregation_params_set(): + if self.prefill_cluster_num_replicas and self.prefill_replica_config: + clusters_info.append( + ( + "PREFILL", + self.prefill_cluster_num_replicas, + self.prefill_replica_config, + ) + ) + + if ( + self.decode_attn_cluster_num_replicas + and self.decode_attn_replica_config + ): + clusters_info.append( + ( + "DECODE_ATTN", + self.decode_attn_cluster_num_replicas, + self.decode_attn_replica_config, + ) + ) + + if self.decode_ffn_cluster_num_replicas and self.decode_ffn_replica_config: + clusters_info.append( + ( + "DECODE_FFN", + self.decode_ffn_cluster_num_replicas, + self.decode_ffn_replica_config, + ) + ) + + if self.decode_cluster_num_replicas and self.decode_replica_config: + clusters_info.append( + ( + "DECODE", + self.decode_cluster_num_replicas, + self.decode_replica_config, + ) + ) + else: + clusters_info.append(("MONOLITHIC", self.num_replicas, self.replica_config)) + + return clusters_info + + def get_server_count_metadata(self, sys_arch: str) -> Dict[str, int]: + clusters_info = self._collect_cluster_info() + server_counts_by_cluster = {} + _, _, _, CollectiveSimCCBackendConfig, _, _ = _get_cc_backend_configs() + cluster_prefix_by_name = { + "PREFILL": "prefill", + "DECODE_ATTN": "decode_attn", + "DECODE_FFN": "decode_ffn", + "DECODE": "decode", + } + + for cluster_name, num_replicas, replica_config in clusters_info: + cluster_total_devices = int(num_replicas) * int(replica_config.world_size) + num_devices_per_node = int(replica_config.node_config.num_devices_per_node) + if cluster_total_devices <= 0: + raise ValueError( + "cluster_total_devices must be positive when computing " + f"server-count metadata, got {cluster_total_devices}" + ) + if num_devices_per_node <= 0: + raise ValueError( + "num_devices_per_node must be positive when computing " + f"server-count metadata, got {num_devices_per_node}" + ) + if cluster_name == "MONOLITHIC": + cc_backend_config = self.cc_backend_config + else: + cc_backend_config = self._create_cc_backend_config_for_cluster( + cluster_prefix_by_name[cluster_name] + ) + if isinstance(cc_backend_config, CollectiveSimCCBackendConfig): + physical_topology = resolve_collective_sim_physical_topology( + cluster_total_devices=cluster_total_devices, + num_devices_per_node=num_devices_per_node, + scenario_profile=getattr( + cc_backend_config, + "scenario_profile", + None, + ), + ) + server_counts_by_cluster[cluster_name] = int(physical_topology.servers) + continue + server_counts_by_cluster[cluster_name] = ( + cluster_total_devices + num_devices_per_node - 1 + ) // num_devices_per_node + + if sys_arch == "co-location": + if "MONOLITHIC" not in server_counts_by_cluster: + raise ValueError("Missing MONOLITHIC cluster for co-location mode.") + return {"server_count": server_counts_by_cluster["MONOLITHIC"]} + if sys_arch == "pd-disaggregation": + if "PREFILL" not in server_counts_by_cluster or "DECODE" not in server_counts_by_cluster: + raise ValueError( + "Missing PREFILL or DECODE cluster for pd-disaggregation mode." + ) + return { + "prefill_server_count": server_counts_by_cluster["PREFILL"], + "decode_server_count": server_counts_by_cluster["DECODE"], + } + if sys_arch == "pd-af-disaggregation": + required = ["PREFILL", "DECODE_ATTN", "DECODE_FFN"] + if any(name not in server_counts_by_cluster for name in required): + raise ValueError( + "Missing PREFILL, DECODE_ATTN, or DECODE_FFN cluster for pd-af-disaggregation mode." + ) + return { + "prefill_server_count": server_counts_by_cluster["PREFILL"], + "decode_attn_server_count": server_counts_by_cluster["DECODE_ATTN"], + "decode_ffn_server_count": server_counts_by_cluster["DECODE_FFN"], + } + + raise ValueError(f"Unknown system architecture: {sys_arch}") + + def print_cluster_statistics(self, simulation_mode: str, sys_arch: str): + """Calculate and print statistics for all clusters (called from SimulationConfig).""" + clusters_info = self._collect_cluster_info() + + # Calculate total statistics + self.total_clusters = len(clusters_info) + self.cluster_world_sizes = {} + + # Calculate world_size if not already set + if not hasattr(self, "world_size") or self.world_size is None: + self.world_size = sum( + num_replicas * replica_config.world_size + for _, num_replicas, replica_config in clusters_info + ) + + # Print cluster configuration summary + print("\n" + "=" * 70) + print("CLUSTER CONFIGURATION SUMMARY") + print("=" * 70) + print("Simulation mode: ", simulation_mode) + print("System architecture: ", sys_arch) + print("=" * 70) + print(f"Total Clusters: {self.total_clusters}") + print(f"Total World Size: {self.world_size}") + server_count_metadata = self.get_server_count_metadata(sys_arch) + if sys_arch == "co-location": + print(f"Server count: {server_count_metadata['server_count']}") + elif sys_arch == "pd-disaggregation": + print( + f"Prefill server count: {server_count_metadata['prefill_server_count']}" + ) + print( + f"Decode server count: {server_count_metadata['decode_server_count']}" + ) + elif sys_arch == "pd-af-disaggregation": + print( + f"Prefill server count: {server_count_metadata['prefill_server_count']}" + ) + print( + f"Decode-Attn server count: {server_count_metadata['decode_attn_server_count']}" + ) + print( + f"Decode-FFN server count: {server_count_metadata['decode_ffn_server_count']}" + ) + print() + + for cluster_name, num_replicas, replica_config in clusters_info: + cluster_world_size = num_replicas * replica_config.world_size + self.cluster_world_sizes[cluster_name] = cluster_world_size + + print(f"Cluster Type: {cluster_name}") + print(f" Cluster World Size: {cluster_world_size}") + print(f" Num Replicas (Instances): {num_replicas}") + print(f" Replica World Size: {replica_config.world_size}") + if ( + cluster_name == "PREFILL" + or cluster_name == "MONOLITHIC" + or cluster_name == "DECODE" + ): + print( + f" Configuration: PP{replica_config.num_pipeline_stages} × (Attn_TP{replica_config.attn_tensor_parallel_size} x Attn_DP{replica_config.attn_dp}) | (MoE_TP{replica_config.moe_tensor_parallel_size} x MoE_EP{replica_config.moe_expert_parallel_size})" + ) + print(f" Total Expert Num: {replica_config.total_expert_num}") + print(f" Local Expert Num: {replica_config.local_expert_num}") + elif cluster_name == "DECODE_ATTN": + print( + f" Configuration: PP{replica_config.num_pipeline_stages} × Attn_TP{replica_config.attn_tensor_parallel_size} x Attn_DP{replica_config.attn_dp}" + ) + elif cluster_name == "DECODE_FFN": + print( + f" Configuration: PP{replica_config.num_pipeline_stages} × MoE_TP{replica_config.moe_tensor_parallel_size} x MoE_EP{replica_config.moe_expert_parallel_size}" + ) + print(f" Total Expert Num: {replica_config.total_expert_num}") + print(f" Local Expert Num: {replica_config.local_expert_num}") + + print("-" * 50) + + print("=" * 70 + "\n") diff --git a/frontier/config/config.py b/frontier/config/config.py index 4d2cc147..cb2b3986 100644 --- a/frontier/config/config.py +++ b/frontier/config/config.py @@ -1,3 +1,10 @@ +"""Top-level simulation configuration and the cluster topology it owns. + +The workload, scheduler, metrics, speculative-decoding, replica and predictor +configuration families live in sibling modules and are re-exported here so that +the historical import surface is unchanged. +""" + from __future__ import annotations from abc import ABC @@ -5,7 +12,7 @@ from datetime import datetime import json import os -from typing import List, Optional, Dict, Tuple, TYPE_CHECKING +from typing import Any, Callable, List, Optional, Dict, Tuple, TYPE_CHECKING from frontier.config.base_poly_config import BasePolyConfig from frontier.config.device_sku_config import BaseDeviceSKUConfig @@ -24,16 +31,10 @@ if TYPE_CHECKING: pass -# NOTE: CC backend configs are imported lazily via _get_cc_backend_configs() function -# to avoid circular imports. Do NOT add direct imports here. -# The lazy import is used in: -# - cc_backend_config field default_factory -# - _create_cc_backend_config_for_cluster() -# - _create_analytical_cc_backend_config() -# - _create_vidur_cc_backend_config() -# - _create_collective_sim_cc_backend_config() -# - _create_aiconfigurator_cc_backend_config() -# - _create_astra_sim_analytical_cc_backend_config() +# NOTE: CC backend configs are imported lazily via _get_cc_backend_configs() to +# avoid a circular import. Do NOT add direct imports here. The lazy import is +# used by the ClusterConfig cc_backend_config default factory and by the +# per-backend creators below. from frontier.config.model_config import BaseModelConfig from frontier.config.node_sku_config import BaseNodeSKUConfig @@ -63,5005 +64,72 @@ validate_run_id, ) -logger = init_logger(__name__) - - -DISAGGREGATED_ARCHITECTURE_RELEASE_ERROR = ( - "Error: Disaggregated architecture support is currently being optimized and is not included in this release. " - "It will be available in an upcoming version. Please use the co-located architecture for current usage and testing." -) - -PD_DISAGGREGATION_PARALLEL_CLUSTER_RELEASE_ERROR = ( - "Error: pd-disaggregation public release support requires " - "--no-enable_parallel_clusters. Parallel PDD is excluded from " - "pre-release-v0.3 because post-ISSUE-022 five-pair MoE-64 measurements " - "were slower than sequential: Simulator.run() by 35.29% and shell E2E " - "by 24.81% (paired medians). The implementation remains available only " - "to internal correctness tests." -) - -PD_AF_DISAGGREGATION_PARALLEL_CLUSTER_RELEASE_ERROR = ( - "Error: pd-af-disaggregation v0.3 requires " - "--no-enable_parallel_clusters. Parallel cluster processing for " - "pd-af-disaggregation is deferred in this release." -) - -PD_AF_PREFIX_CACHING_RELEASE_ERROR = ( - "Prefix caching is excluded for pd-af-disaggregation in v0.3. " - "Disable replica_scheduler_config.enable_prefix_caching." +# The configuration families below were split out of this module. They are +# imported here both because ClusterConfig and SimulationConfig reference them +# and because `frontier/config/__init__.py` re-exports this module with a star +# import, so every historical `from frontier.config import X` and +# `from frontier.config.config import X` keeps resolving. +from frontier.config.release_guards import ( + AICONFIGURATOR_BACKEND_RELEASE_ERROR, + DISAGGREGATED_ARCHITECTURE_RELEASE_ERROR, + DISAGGREGATED_CLUSTER_FIELD_NAMES, + DISAGGREGATED_CLUSTER_FIELD_PREFIXES, + PD_AF_DISAGGREGATION_PARALLEL_CLUSTER_RELEASE_ERROR, + PD_AF_PREFIX_CACHING_RELEASE_ERROR, + PD_AF_TRACE_REPLAY_DEFERRED_ERROR, + PD_DISAGGREGATION_PARALLEL_CLUSTER_RELEASE_ERROR, ) - -AICONFIGURATOR_BACKEND_RELEASE_ERROR = ( - "Error: The aiconfigurator communication backend is not included in this release. " - "Please use collective_sim, astra_sim_analytical, analytical, or vidur for current usage and testing." +from frontier.config.request_generator_config import ( + BaseRequestGeneratorConfig, + BaseRequestIntervalGeneratorConfig, + BaseRequestLengthGeneratorConfig, + FixedRequestLengthGeneratorConfig, + GammaRequestIntervalGeneratorConfig, + PoissonRequestIntervalGeneratorConfig, + StaticRequestIntervalGeneratorConfig, + SyntheticRequestGeneratorConfig, + TraceRequestGeneratorConfig, + TraceRequestIntervalGeneratorConfig, + TraceRequestLengthGeneratorConfig, + UniformRequestLengthGeneratorConfig, + ZipfRequestLengthGeneratorConfig, ) - -PD_AF_TRACE_REPLAY_DEFERRED_ERROR = ( - "Error: pd-af-disaggregation v0.3 trace-replay is deferred. " - "The configured trace-driven fields are public stubs only and are not " - "implemented in this release." +from frontier.config.replica_scheduler_config import ( + BaseReplicaSchedulerConfig, + FasterTransformerSchedulerConfig, + LightllmSchedulerConfig, + OrcaSchedulerConfig, + SarathiSchedulerConfig, + SglangSchedulerConfig, + Sj2QBoundedCarryoverSchedulerConfig, + Sj2QFastServeLiteSchedulerConfig, + Sj2QPenaltyOnlySchedulerConfig, + Sj2qBoundedCarryoverSchedulerConfig, + Sj2qFastserveLiteSchedulerConfig, + Sj2qPenaltyOnlySchedulerConfig, + VllmSchedulerConfig, + VllmV1SchedulerConfig, ) - -DISAGGREGATED_CLUSTER_FIELD_PREFIXES = ( - "prefill_", - "decode_", - "decode_attn_", - "decode_ffn_", +from frontier.config.metrics_config import MetricsConfig +from frontier.config.speculative_decoding_config import SpeculativeDecodingConfig +from frontier.config.replica_config import ReplicaConfig +from frontier.config.cluster_scheduler_config import ( + BaseClusterSchedulerConfig, + LORClusterSchedulerConfig, + RandomClusterSchedulerConfig, + RoundRobinClusterSchedulerConfig, + StickyLORClusterSchedulerConfig, + StickyRoundRobinClusterSchedulerConfig, ) - -DISAGGREGATED_CLUSTER_FIELD_NAMES = frozenset( - { - "af_pipeline_num_micro_batch", - } +from frontier.config.execution_time_predictor_config import ( + BaseExecutionTimePredictorConfig, + LinearRegressionExecutionTimePredictorConfig, + RandomForrestExecutionTimePredictorConfig, ) +from frontier.config.cluster_config import ClusterConfig -# Lazy import helper for cc_backend_config to avoid circular imports -def _get_cc_backend_configs(): - """Lazily import CC backend config classes to avoid circular imports.""" - from frontier.cc_backend.cc_backend_config import ( - BaseCCBackendConfig, - VidurCCBackendConfig, - AnalyticalCCBackendConfig, - CollectiveSimCCBackendConfig, - AiconfiguratorCCBackendConfig, - AstraSimAnalyticalCCBackendConfig, - ) - - return ( - BaseCCBackendConfig, - VidurCCBackendConfig, - AnalyticalCCBackendConfig, - CollectiveSimCCBackendConfig, - AiconfiguratorCCBackendConfig, - AstraSimAnalyticalCCBackendConfig, - ) - - -@dataclass -class BaseRequestIntervalGeneratorConfig(BasePolyConfig): - seed: int = field( - default=42, - metadata={"help": "Seed for the random number generator."}, - ) - - -@dataclass -class BaseRequestLengthGeneratorConfig(BasePolyConfig): - seed: int = field( - default=42, - metadata={"help": "Seed for the random number generator."}, - ) - max_tokens: int = field( - default=4096, - metadata={"help": "Maximum tokens."}, - ) - - -@dataclass -class TraceRequestIntervalGeneratorConfig(BaseRequestIntervalGeneratorConfig): - trace_file: str = field( - default="data/processed_traces/AzureFunctionsInvocationTraceForTwoWeeksJan2021Processed.csv", - metadata={"help": "Path to the trace request interval generator file."}, - ) - start_time: str = field( - default="1970-01-04 12:00:00", - metadata={"help": "Start time of the trace request interval generator."}, - ) - end_time: str = field( - default="1970-01-04 15:00:00", - metadata={"help": "End time of the trace request interval generator."}, - ) - time_scale_factor: float = field( - default=1.0, - metadata={ - "help": "Time scale factor for the trace request interval generator." - }, - ) - - @staticmethod - def get_type(): - return RequestIntervalGeneratorType.TRACE - - -@dataclass -class PoissonRequestIntervalGeneratorConfig(BaseRequestIntervalGeneratorConfig): - qps: float = field( - default=0.5, - metadata={"help": "Queries per second for Poisson Request Interval Generator."}, - ) - - @staticmethod - def get_type(): - return RequestIntervalGeneratorType.POISSON - - -@dataclass -class GammaRequestIntervalGeneratorConfig(BaseRequestIntervalGeneratorConfig): - qps: float = field( - default=0.2, - metadata={"help": "Queries per second for Gamma Request Interval Generator."}, - ) - cv: float = field( - default=0.5, - metadata={ - "help": "Coefficient of variation for Gamma Request Interval Generator." - }, - ) - - @staticmethod - def get_type(): - return RequestIntervalGeneratorType.GAMMA - - -@dataclass -class StaticRequestIntervalGeneratorConfig(BaseRequestIntervalGeneratorConfig): - @staticmethod - def get_type(): - return RequestIntervalGeneratorType.STATIC - - -@dataclass -class TraceRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): - trace_file: str = field( - default="data/processed_traces/sharegpt_8k_filtered_stats_llama2_tokenizer.csv", - metadata={"help": "Path to the trace request length generator file."}, - ) - prefill_scale_factor: float = field( - default=1, - metadata={ - "help": "Prefill scale factor for the trace request length generator." - }, - ) - decode_scale_factor: float = field( - default=1, - metadata={ - "help": "Decode scale factor for the trace request length generator." - }, - ) - - @staticmethod - def get_type(): - return RequestLengthGeneratorType.TRACE - - -@dataclass -class ZipfRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): - theta: float = field( - default=0.6, - metadata={"help": "Theta for Zipf Request Length Generator."}, - ) - scramble: bool = field( - default=False, - metadata={"help": "Scramble for Zipf Request Length Generator."}, - ) - min_tokens: int = field( - default=1024, - metadata={"help": "Minimum tokens for Zipf Request Length Generator."}, - ) - prefill_to_decode_ratio: float = field( - default=20.0, - metadata={"help": "Prefill to decode ratio for Zipf Request Length Generator."}, - ) - - @staticmethod - def get_type(): - return RequestLengthGeneratorType.ZIPF - - -@dataclass -class UniformRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): - min_tokens: int = field( - default=1024, - metadata={"help": "Minimum tokens for Uniform Request Length Generator."}, - ) - prefill_to_decode_ratio: float = field( - default=20.0, - metadata={ - "help": "Prefill to decode ratio for Uniform Request Length Generator." - }, - ) - - @staticmethod - def get_type(): - return RequestLengthGeneratorType.UNIFORM - - -@dataclass -class FixedRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): - prefill_tokens: int = field( - default=2048, - metadata={"help": "Prefill tokens for Fixed Request Length Generator."}, - ) - decode_tokens: int = field( - default=512, - metadata={"help": "Decode tokens for Fixed Request Length Generator."}, - ) - - @staticmethod - def get_type(): - return RequestLengthGeneratorType.FIXED - - def __post_init__(self): - if self.decode_tokens < 1: - raise ValueError(f"decode_tokens must be >= 1, got {self.decode_tokens}") - if self.prefill_tokens < 2: - raise ValueError(f"prefill_tokens must be >1, got {self.prefill_tokens}") - - -@dataclass -class BaseRequestGeneratorConfig(BasePolyConfig): - seed: int = field( - default=42, - metadata={"help": "Seed for the random number generator."}, - ) - num_decode_bound_requests: Optional[int] = field( - default=None, - metadata={ - "help": "Number of generated requests that require decode-cluster work. " - "Derived by request generation and used by offline pd-disaggregation scheduling." - }, - ) - - -@dataclass -class SyntheticRequestGeneratorConfig(BaseRequestGeneratorConfig): - length_generator_config: BaseRequestLengthGeneratorConfig = field( - default_factory=FixedRequestLengthGeneratorConfig, - metadata={"help": "Length generator config for Synthetic Request Generator."}, - ) - interval_generator_config: BaseRequestIntervalGeneratorConfig = field( - default_factory=PoissonRequestIntervalGeneratorConfig, - metadata={"help": "Interval generator config for Synthetic Request Generator."}, - ) - num_requests: Optional[int] = field( - default=128, - metadata={"help": "Number of requests for Synthetic Request Generator."}, - ) - duration: Optional[float] = field( - default=None, - metadata={"help": "Duration of the synthetic request generator."}, - ) - default_priority: int = field( - default=0, - metadata={ - "help": "Default priority for all generated requests. " - "Lower value = higher priority (0 = highest). " - "Matches vLLM v1 semantics." - }, - ) - - def __post_init__(self): - self.max_tokens = self.length_generator_config.max_tokens - - @staticmethod - def get_type(): - return RequestGeneratorType.SYNTHETIC - - -@dataclass -class TraceRequestGeneratorConfig(BaseRequestGeneratorConfig): - trace_file: str = field( - default="data/processed_traces/splitwise_conv.csv", - metadata={"help": "Path to the trace request generator file."}, - ) - prefill_scale_factor: float = field( - default=1.0, - metadata={"help": "Prefill scale factor for the trace request generator."}, - ) - decode_scale_factor: float = field( - default=1.0, - metadata={"help": "Decode scale factor for the trace request generator."}, - ) - time_scale_factor: float = field( - default=1.0, - metadata={"help": "Time scale factor for the trace request generator."}, - ) - max_tokens: int = field( - default=4096, - metadata={"help": "Maximum tokens for the trace request generator."}, - ) - - @staticmethod - def get_type(): - return RequestGeneratorType.TRACE_REPLAY - - -@dataclass -class BaseReplicaSchedulerConfig(BasePolyConfig): - batch_size_cap: int = field( - default=128, - metadata={"help": "Maximum batch size cap (max_num_seqs in vLLM)"}, - ) - block_size: int = field( - default=16, - metadata={"help": "Block size."}, - ) - watermark_blocks_fraction: float = field( - default=0.01, - metadata={"help": "Watermark blocks fraction."}, - ) - num_blocks: Optional[int] = field( - default=106596, - metadata={"help": "Number of blocks."}, - ) - - -@dataclass -class VllmSchedulerConfig(BaseReplicaSchedulerConfig): - max_tokens_in_batch: int = field( - default=4096, - metadata={"help": "Maximum tokens (max_num_batched_tokens) in batch for vLLM."}, - ) - - @staticmethod - def get_type(): - return ReplicaSchedulerType.VLLM - - -@dataclass -class LightllmSchedulerConfig(BaseReplicaSchedulerConfig): - max_tokens_in_batch: int = field( - default=4096, - metadata={"help": "Maximum tokens in batch for LightLLM."}, - ) - max_waiting_iters: int = field( - default=10, - metadata={"help": "Maximum waiting iterations for LightLLM."}, - ) - - @staticmethod - def get_type(): - return ReplicaSchedulerType.LIGHTLLM - - -@dataclass -class OrcaSchedulerConfig(BaseReplicaSchedulerConfig): - @staticmethod - def get_type(): - return ReplicaSchedulerType.ORCA - - -@dataclass -class FasterTransformerSchedulerConfig(BaseReplicaSchedulerConfig): - @staticmethod - def get_type(): - return ReplicaSchedulerType.FASTER_TRANSFORMER - - -@dataclass -class SarathiSchedulerConfig(BaseReplicaSchedulerConfig): - chunk_size: int = field( - default=512, - metadata={"help": "Chunk size for Sarathi."}, - ) - - @staticmethod - def get_type(): - return ReplicaSchedulerType.SARATHI - - -@dataclass -class VllmV1SchedulerConfig(BaseReplicaSchedulerConfig): - """ - Configuration for the vLLM v1 engine replica scheduler. - - This scheduler simulates the admission control behavior of vLLM v1 engine, - including two-phase scheduling, token budget management, and preemption. - - Note: Class name uses 'VllmV1' (not 'VLLMv1') to generate clean CLI parameter - names: --vllm_v1_scheduler_config_* instead of --v_l_l_mv1_scheduler_config_* - """ - - max_tokens_in_batch: int = field( - default=16384, - metadata={ - "help": "Maximum tokens per scheduling iteration (max_num_batched_tokens in vLLM v1)." - }, - ) - scheduling_policy: str = field( - default="fcfs", - metadata={ - "help": "Scheduling policy: 'fcfs' (First-Come-First-Served) or 'priority'." - }, - ) - enable_preemption: bool = field( - default=True, - metadata={ - "help": "Enable preemption when memory is insufficient for running requests." - }, - ) - enable_chunked_prefill: bool = field( - default=False, - metadata={ - "help": "Enable chunked prefill admission when waiting prefill requests exceed current token budget." - }, - ) - enable_phase_aware_thinking_profile: bool = field( - default=False, - metadata={ - "help": "Enable an iteration-scoped hidden-round/final-round scheduler profile override for Thinking Mode home queues." - }, - ) - hidden_phase_max_tokens_in_batch: Optional[int] = field( - default=None, - metadata={ - "help": "Optional hidden-round override for max_tokens_in_batch when enable_phase_aware_thinking_profile=True." - }, - ) - hidden_phase_enable_chunked_prefill: Optional[bool] = field( - default=None, - metadata={ - "help": "Optional hidden-round override for enable_chunked_prefill when enable_phase_aware_thinking_profile=True." - }, - ) - hidden_phase_batch_size_cap: Optional[int] = field( - default=None, - metadata={ - "help": "Optional hidden-round override for batch_size_cap when enable_phase_aware_thinking_profile=True." - }, - ) - final_phase_max_tokens_in_batch: Optional[int] = field( - default=None, - metadata={ - "help": "Optional final-round override for max_tokens_in_batch when enable_phase_aware_thinking_profile=True." - }, - ) - final_phase_enable_chunked_prefill: Optional[bool] = field( - default=None, - metadata={ - "help": "Optional final-round override for enable_chunked_prefill when enable_phase_aware_thinking_profile=True." - }, - ) - final_phase_batch_size_cap: Optional[int] = field( - default=None, - metadata={ - "help": "Optional final-round override for batch_size_cap when enable_phase_aware_thinking_profile=True." - }, - ) - final_prefill_reserved_slots: int = field( - default=0, - metadata={ - "help": "Per-iteration PREFILL admission slots reserved for final-round prefill requests. Hidden requests may borrow idle reserved slots." - }, - ) - final_prefill_reserved_tokens: int = field( - default=0, - metadata={ - "help": "Per-iteration PREFILL token budget reserved for final-round prefill requests. Hidden requests may borrow idle reserved tokens." - }, - ) - final_decode_reserved_slots: int = field( - default=0, - metadata={ - "help": "Per-iteration DECODE running/admission slots reserved for final-round decode requests. Hidden requests may borrow idle reserved slots." - }, - ) - enable_final_running_request_reclaim: bool = field( - default=False, - metadata={ - "help": "When final backlog appears, reclaim hidden requests that have borrowed final reserved running slots so the final slice becomes active running capacity." - }, - ) - enable_final_round_priority_boost: bool = field( - default=False, - metadata={ - "help": "Promote re-entered final-round Thinking Mode requests into a higher-priority band under priority scheduling." - }, - ) - final_round_priority_value: int = field( - default=-1, - metadata={ - "help": "Priority value assigned to promoted final-round requests. Lower values mean higher priority." - }, - ) - enable_prefix_caching: bool = field( - default=False, - metadata={ - "help": "Enable block-hash-based prefix matching and KV cache reuse." - }, - ) - prefix_caching_hash_algo: str = field( - default="builtin", - metadata={ - "help": "Hash algorithm label for explicit prefix block hashes. Supported: 'builtin', 'sha256'." - }, - ) - num_preallocate_tokens: int = field( - default=0, - metadata={ - "help": "Number of tokens worth of KV cache blocks to preallocate for each request." - }, - ) - long_prefill_token_threshold: int = field( - default=0, - metadata={ - "help": "Optional upper bound on per-iteration prefill tokens for each request. 0 disables threshold." - }, - ) - num_blocks: Optional[int] = field( - default=0, - metadata={ - "help": "Number of KV cache blocks. Use 0 to auto-derive from the memory planner in planner modes." - }, - ) - num_blocks_mode: str = field( - default="memory_planner_profiled", - metadata={ - "help": "How to initialize num_blocks: 'memory_planner' (auto-derive with parameter-only estimate), 'memory_planner_profiled' (auto-derive with calibrated non-KV overhead), or 'explicit' (require num_blocks>0)." - }, - ) - gpu_memory_utilization: Optional[float] = field( - default=None, - metadata={ - "help": "vLLM-style GPU memory utilization ratio used by memory_planner mode. If unset, fallback to 1 - replica memory_margin_fraction." - }, - ) - non_kv_cache_overhead_bytes: int = field( - default=0, - metadata={ - "help": "Calibrated non-KV memory overhead in bytes for memory_planner_profiled mode." - }, - ) - runtime_weights_memory_source: str = field( - default="param_counter", - metadata={ - "help": "Weights memory source for runtime non-KV profiling: 'param_counter' (estimated bytes) or 'runtime_model_load' (measure loaded model parameter bytes)." - }, - ) - enable_runtime_non_kv_cache_overhead_profiling: bool = field( - default=False, - metadata={ - "help": "Enable runtime single-rank profiling to auto-estimate non_kv_cache_overhead_bytes during scheduler initialization. Requires num_blocks_mode=memory_planner_profiled." - }, - ) - nccl_buffer_comm_base_overhead_bytes: int = field( - default=100 * 1024 * 1024, - metadata={ - "help": "Per-communicator fixed NCCL overhead in bytes (proxy, queues). " - "Default 100 MiB, calibrated for A800." - }, - ) - nccl_buffer_per_peer_overhead_bytes: int = field( - default=15 * 1024 * 1024, - metadata={ - "help": "Per-peer NCCL transport buffer overhead in bytes. " - "Default 15 MiB, calibrated for A800 intra-node." - }, - ) - nccl_buffer_custom_ar_enabled: bool = field( - default=False, - metadata={ - "help": "Enable CustomAllreduce buffer estimation. " - "False for A800 (compute 8.0), True for H100 (9.0+)." - }, - ) - nccl_buffer_vllm_worker_base_extra_bytes: int = field( - default=0, - metadata={ - "help": "Domain-aware vLLM worker-process non-torch addend in bytes " - "for runtime non-KV profiling. Default 0; pass validated " - "case-local values explicitly." - }, - ) - nccl_buffer_pp_final_stage_extra_bytes: int = field( - default=0, - metadata={ - "help": "Additional final pipeline-stage vLLM worker non-torch addend " - "in bytes for runtime non-KV profiling. Default 0." - }, - ) - nccl_buffer_dp_communicator_extra_bytes: int = field( - default=0, - metadata={ - "help": "Additional data-parallel communicator non-torch addend in " - "bytes for runtime non-KV profiling. Default 0." - }, - ) - nccl_buffer_ep_all2all_extra_bytes: int = field( - default=0, - metadata={ - "help": "Additional MoE expert-parallel all-to-all non-torch addend " - "in bytes for runtime non-KV profiling. Default 0." - }, - ) - use_analytical_param_memory: bool = field( - default=False, - metadata={ - "help": "When runtime non-KV profiling is enabled in memory_planner_profiled mode, keep analytical ParamCounter param memory for planner calculation. Default False uses runtime-profiled param memory." - }, - ) - - def __post_init__(self) -> None: - allowed_modes = {"memory_planner", "memory_planner_profiled", "explicit"} - if self.num_blocks_mode not in allowed_modes: - raise ValueError( - "VllmV1SchedulerConfig.num_blocks_mode must be one of " - f"{sorted(allowed_modes)}, got={self.num_blocks_mode!r}" - ) - - if self.gpu_memory_utilization is not None: - if self.gpu_memory_utilization <= 0 or self.gpu_memory_utilization > 1.0: - raise ValueError( - "VllmV1SchedulerConfig.gpu_memory_utilization must be in (0, 1], got=" - f"{self.gpu_memory_utilization!r}" - ) - - if self.non_kv_cache_overhead_bytes < 0: - raise ValueError( - "VllmV1SchedulerConfig.non_kv_cache_overhead_bytes must be >= 0, got=" - f"{self.non_kv_cache_overhead_bytes!r}" - ) - - allowed_hash_algorithms = {"builtin", "sha256"} - if self.prefix_caching_hash_algo not in allowed_hash_algorithms: - raise ValueError( - "VllmV1SchedulerConfig.prefix_caching_hash_algo must be one of " - f"{sorted(allowed_hash_algorithms)}, got={self.prefix_caching_hash_algo!r}" - ) - - if self.num_preallocate_tokens < 0: - raise ValueError( - "VllmV1SchedulerConfig.num_preallocate_tokens must be >= 0, got=" - f"{self.num_preallocate_tokens!r}" - ) - - if self.long_prefill_token_threshold < 0: - raise ValueError( - "VllmV1SchedulerConfig.long_prefill_token_threshold must be >= 0, got=" - f"{self.long_prefill_token_threshold!r}" - ) - if ( - self.long_prefill_token_threshold > 0 - and not self.enable_chunked_prefill - ): - raise ValueError( - "VllmV1SchedulerConfig.long_prefill_token_threshold > 0 " - "requires enable_chunked_prefill=True" - ) - - phase_override_values = ( - self.hidden_phase_max_tokens_in_batch, - self.hidden_phase_enable_chunked_prefill, - self.hidden_phase_batch_size_cap, - self.final_phase_max_tokens_in_batch, - self.final_phase_enable_chunked_prefill, - self.final_phase_batch_size_cap, - ) - if not self.enable_phase_aware_thinking_profile and any( - value is not None for value in phase_override_values - ): - raise ValueError( - "VllmV1SchedulerConfig phase-aware override fields require " - "enable_phase_aware_thinking_profile=True" - ) - if self.enable_phase_aware_thinking_profile and all( - value is None for value in phase_override_values - ): - raise ValueError( - "VllmV1SchedulerConfig.enable_phase_aware_thinking_profile=True " - "requires at least one hidden/final override field" - ) - - for field_name in ( - "hidden_phase_max_tokens_in_batch", - "hidden_phase_batch_size_cap", - "final_phase_max_tokens_in_batch", - "final_phase_batch_size_cap", - ): - field_value = getattr(self, field_name) - if field_value is not None and field_value <= 0: - raise ValueError( - f"VllmV1SchedulerConfig.{field_name} must be > 0 when set, " - f"got={field_value!r}" - ) - - for field_name in ( - "final_prefill_reserved_slots", - "final_prefill_reserved_tokens", - "final_decode_reserved_slots", - ): - field_value = getattr(self, field_name) - if field_value < 0: - raise ValueError( - f"VllmV1SchedulerConfig.{field_name} must be >= 0, " - f"got={field_value!r}" - ) - - if ( - self.long_prefill_token_threshold > 0 - and self.enable_phase_aware_thinking_profile - ): - if self.hidden_phase_enable_chunked_prefill is False: - raise ValueError( - "VllmV1SchedulerConfig.hidden_phase_enable_chunked_prefill=False " - "is incompatible with long_prefill_token_threshold > 0" - ) - if self.final_phase_enable_chunked_prefill is False: - raise ValueError( - "VllmV1SchedulerConfig.final_phase_enable_chunked_prefill=False " - "is incompatible with long_prefill_token_threshold > 0" - ) - - if self.nccl_buffer_comm_base_overhead_bytes < 0: - raise ValueError( - "VllmV1SchedulerConfig.nccl_buffer_comm_base_overhead_bytes must be >= 0, got=" - f"{self.nccl_buffer_comm_base_overhead_bytes!r}" - ) - - if self.nccl_buffer_per_peer_overhead_bytes < 0: - raise ValueError( - "VllmV1SchedulerConfig.nccl_buffer_per_peer_overhead_bytes must be >= 0, got=" - f"{self.nccl_buffer_per_peer_overhead_bytes!r}" - ) - - for field_name in ( - "nccl_buffer_vllm_worker_base_extra_bytes", - "nccl_buffer_pp_final_stage_extra_bytes", - "nccl_buffer_dp_communicator_extra_bytes", - "nccl_buffer_ep_all2all_extra_bytes", - ): - field_value = getattr(self, field_name) - if field_value < 0: - raise ValueError( - f"VllmV1SchedulerConfig.{field_name} must be >= 0, " - f"got={field_value!r}" - ) - - allowed_weights_sources = {"param_counter", "runtime_model_load"} - if self.runtime_weights_memory_source not in allowed_weights_sources: - raise ValueError( - "VllmV1SchedulerConfig.runtime_weights_memory_source must be one of " - f"{sorted(allowed_weights_sources)}, got={self.runtime_weights_memory_source!r}" - ) - - if ( - self.enable_runtime_non_kv_cache_overhead_profiling - and self.num_blocks_mode != "memory_planner_profiled" - ): - raise ValueError( - "VllmV1SchedulerConfig.enable_runtime_non_kv_cache_overhead_profiling " - "requires num_blocks_mode=memory_planner_profiled, got=" - f"{self.num_blocks_mode!r}" - ) - - if ( - self.use_analytical_param_memory - and not self.enable_runtime_non_kv_cache_overhead_profiling - ): - raise ValueError( - "VllmV1SchedulerConfig.use_analytical_param_memory " - "requires enable_runtime_non_kv_cache_overhead_profiling=True" - ) - - enable_thinking_round_priority: bool = field( - default=False, - metadata={ - "help": "When enabled, final-round thinking requests are prioritized " - "over non-final-round requests in the waiting queue.", - }, - ) - - @staticmethod - def get_type(): - return ReplicaSchedulerType.VLLM_V1 - - -@dataclass -class Sj2qFastserveLiteSchedulerConfig(VllmV1SchedulerConfig): - """ - Configuration for the SJ-2Q / FastServe-lite scheduler. - - Note: Class name uses 'Sj2qFastserve' (not 'Sj2QFastServe') to generate clean - CLI parameter names: --sj2q_fastserve_lite_scheduler_config_* instead of - --sj2_q_fast_serve_lite_scheduler_config_*. - """ - - long_round_new_prompt_threshold: int = field( - default=2048, - metadata={ - "help": "Rounds whose new prompt tokens exceed this threshold enter QL and mark long_history." - }, - ) - short_round_boost_threshold: int = field( - default=512, - metadata={ - "help": "Tiny-prefill threshold used for QH prioritization and the prefill-release-only boost when long_history is already true." - }, - ) - boost_credit_token_budget: int = field( - default=2048, - metadata={ - "help": "Deprecated compatibility field retained for CLI stability; current prefill-release-only boost demotes on prefill completion instead of token-budget exhaustion." - }, - ) - enable_aging: bool = field( - default=False, - metadata={ - "help": "Enable optional aging-based QL promotion back into QH. The UC3 v2 enhancement lane keeps this disabled." - }, - ) - aging_wait_threshold_ms: float = field( - default=7.5, - metadata={ - "help": "QL waiting-time threshold in milliseconds for a temporary aging-based QH boost." - }, - ) - aging_boost_token_budget: int = field( - default=512, - metadata={ - "help": "Token budget granted when an aged QL session is temporarily promoted into QH." - }, - ) - - def __post_init__(self) -> None: - super().__post_init__() - - if self.enable_phase_aware_thinking_profile: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig does not allow phase-aware oracle scheduling." - ) - if self.enable_thinking_round_priority: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig does not allow final-round priority override." - ) - if ( - self.final_prefill_reserved_slots != 0 - or self.final_prefill_reserved_tokens != 0 - or self.final_decode_reserved_slots != 0 - ): - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig requires all final reserved slot/token settings to remain 0." - ) - if self.enable_final_running_request_reclaim: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig does not allow final running-request reclaim." - ) - if self.enable_final_round_priority_boost: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig does not allow final-round priority boost." - ) - - if self.long_round_new_prompt_threshold <= 0: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig.long_round_new_prompt_threshold must be > 0." - ) - if self.short_round_boost_threshold <= 0: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig.short_round_boost_threshold must be > 0." - ) - if ( - self.short_round_boost_threshold - > self.long_round_new_prompt_threshold - ): - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig.short_round_boost_threshold must be <= long_round_new_prompt_threshold." - ) - if self.boost_credit_token_budget <= 0: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig.boost_credit_token_budget must be > 0." - ) - if self.enable_aging and self.aging_wait_threshold_ms <= 0: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig.aging_wait_threshold_ms must be > 0 when aging is enabled." - ) - if self.aging_boost_token_budget <= 0: - raise ValueError( - "Sj2QFastserveLiteSchedulerConfig.aging_boost_token_budget must be > 0." - ) - - @staticmethod - def get_type(): - return ReplicaSchedulerType.SJ2Q_FASTSERVE_LITE - - -Sj2QFastServeLiteSchedulerConfig = Sj2qFastserveLiteSchedulerConfig - - -@dataclass -class Sj2qPenaltyOnlySchedulerConfig(VllmV1SchedulerConfig): - """ - Configuration for the penalty-only SJ-2Q scheduler. - - Note: Class name uses 'Sj2q' to generate clean CLI parameter names like - --sj2q_penalty_only_scheduler_config_*. - """ - - long_round_new_prompt_threshold: int = field( - default=4096, - metadata={ - "help": "Rounds whose new prompt tokens exceed this threshold immediately enter Qlong and mark long_history." - }, - ) - service_cap_tokens: int = field( - default=8192, - metadata={ - "help": "Session-level cumulative new-token service cap after which the session stays in Qlong." - }, - ) - long_liveness_quota: int = field( - default=32, - metadata={ - "help": "Maximum consecutive Qshort slices allowed before forcing one Qlong slice when Qlong is non-empty." - }, - ) - - def __post_init__(self) -> None: - super().__post_init__() - - if self.enable_phase_aware_thinking_profile: - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig does not allow phase-aware oracle scheduling." - ) - if self.enable_thinking_round_priority: - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig does not allow final-round priority override." - ) - if ( - self.final_prefill_reserved_slots != 0 - or self.final_prefill_reserved_tokens != 0 - or self.final_decode_reserved_slots != 0 - ): - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig requires all final reserved slot/token settings to remain 0." - ) - if self.enable_final_running_request_reclaim: - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig does not allow final running-request reclaim." - ) - if self.enable_final_round_priority_boost: - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig does not allow final-round priority boost." - ) - if self.long_round_new_prompt_threshold <= 0: - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig.long_round_new_prompt_threshold must be > 0." - ) - if self.service_cap_tokens <= 0: - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig.service_cap_tokens must be > 0." - ) - if self.service_cap_tokens < self.long_round_new_prompt_threshold: - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig.service_cap_tokens must be >= long_round_new_prompt_threshold." - ) - if self.long_liveness_quota <= 0: - raise ValueError( - "Sj2qPenaltyOnlySchedulerConfig.long_liveness_quota must be > 0." - ) - - @staticmethod - def get_type(): - return ReplicaSchedulerType.SJ2Q_PENALTY_ONLY - - -Sj2QPenaltyOnlySchedulerConfig = Sj2qPenaltyOnlySchedulerConfig - - -@dataclass -class Sj2qBoundedCarryoverSchedulerConfig(Sj2qPenaltyOnlySchedulerConfig): - """ - Configuration for the bounded-carryover SJ-2Q scheduler. - - Note: Class name uses 'Sj2q' to generate clean CLI parameter names like - --sj2q_bounded_carryover_scheduler_config_*. - """ - - @staticmethod - def get_type(): - return ReplicaSchedulerType.SJ2Q_BOUNDED_CARRYOVER - - -Sj2QBoundedCarryoverSchedulerConfig = Sj2qBoundedCarryoverSchedulerConfig - - -@dataclass -class SglangSchedulerConfig(VllmV1SchedulerConfig): - """ - Thin config wrapper for the Frontier SGLang-style replica scheduler. - - This intentionally reuses the vLLM v1 scheduler fields and only changes - the scheduler type to keep the integration surface minimal. - """ - - @staticmethod - def get_type(): - return ReplicaSchedulerType.SGLANG - - -@dataclass -class MetricsConfig: - """Metric configuration.""" - - write_metrics: bool = field( - default=True, - metadata={"help": "Whether to write metrics."}, - ) - write_json_trace: bool = field( - default=False, - metadata={"help": "Whether to write json trace."}, - ) - wandb_project: Optional[str] = field( - default=None, - metadata={"help": "Weights & Biases project name."}, - ) - wandb_group: Optional[str] = field( - default=None, - metadata={"help": "Weights & Biases group name."}, - ) - wandb_run_name: Optional[str] = field( - default=None, - metadata={"help": "Weights & Biases run name."}, - ) - wandb_sweep_id: Optional[str] = field( - default=None, - metadata={"help": "Weights & Biases sweep id."}, - ) - wandb_run_id: Optional[str] = field( - default=None, - metadata={"help": "Weights & Biases run id."}, - ) - enable_chrome_trace: bool = field( - default=True, - metadata={"help": "Enable Chrome tracing."}, - ) - - # Op-Level Tracing - enable_op_level_tracing: bool = field( - default=False, - metadata={"help": "Enable detailed op-level tracing (output to JSONL)."}, - ) - trace_output_file: str = field( - default="op_traces.jsonl", - metadata={"help": "Output filename for op-level traces."}, - ) - enable_metrics_ground_truth_trace: bool = field( - default=False, - metadata={ - "help": "Enable explicit request-level metrics ground-truth JSONL output." - }, - ) - metrics_ground_truth_trace_file: str = field( - default="metrics_ground_truth.jsonl", - metadata={"help": "Output filename for metrics ground-truth request traces."}, - ) - enable_per_layer_expansion: bool = field( - default=False, - metadata={ - "help": "Enable per-layer trace expansion. When enabled, traces show " - "individual layer operations instead of aggregated spans." - }, - ) - num_requests_to_trace_per_layer: int = field( - default=5, - metadata={ - "help": "Number of requests to capture with per-layer expansion. " - "Only applies when enable_per_layer_expansion is True." - }, - ) - - save_table_to_wandb: bool = field( - default=False, - metadata={"help": "Whether to save table to wandb."}, - ) - store_plots: bool = field( - default=True, - metadata={"help": "Whether to store plots."}, - ) - enable_memory_time_series: bool = field( - default=False, - metadata={ - "help": "Enable memory usage time series output. " - "Only valid when log_level is 'debug'." - }, - ) - store_operation_metrics: bool = field( - default=False, - metadata={"help": "Whether to store operation metrics."}, - ) - store_token_completion_metrics: bool = field( - default=False, - metadata={"help": "Whether to store token completion metrics."}, - ) - store_request_metrics: bool = field( - default=True, - metadata={"help": "Whether to store request metrics."}, - ) - store_batch_metrics: bool = field( - default=True, - metadata={"help": "Whether to store batch metrics."}, - ) - store_utilization_metrics: bool = field( - default=True, - metadata={"help": "Whether to store utilization metrics."}, - ) - keep_individual_batch_metrics: bool = field( - default=False, - metadata={"help": "Whether to keep individual batch metrics."}, - ) - store_frontier_stage_batch_ledger: bool = field( - default=True, - metadata={"help": "Whether to write the full Frontier stage-batch ledger."}, - ) - store_frontier_stage_batch_ledger_summary: bool = field( - default=False, - metadata={ - "help": "Whether to write a bounded Frontier stage-batch ledger summary." - }, - ) - subsamples: Optional[int] = field( - default=None, - metadata={"help": "Subsamples."}, - ) - min_batch_index: Optional[int] = field( - default=None, - metadata={"help": "Minimum batch index."}, - ) - max_batch_index: Optional[int] = field( - default=None, - metadata={"help": "Maximum batch index."}, - ) - output_dir: str = field( - default="outputs/metrics", - metadata={"help": "Metrics output root directory."}, - ) - cache_dir: str = field( - default="cache", - metadata={"help": "Cache directory."}, - ) - run_id: Optional[str] = field( - default=None, - metadata={ - "help": "Metrics run id used under outputs/metrics///." - }, - ) - - def __post_init__(self): - if self.run_id is None: - self.run_id = f"run_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S-%f')}" - self.run_id = validate_run_id(self.run_id) - self.trace_output_file = validate_output_filename( - self.trace_output_file, "trace_output_file" - ) - self.metrics_ground_truth_trace_file = validate_output_filename( - self.metrics_ground_truth_trace_file, "metrics_ground_truth_trace_file" - ) - os.makedirs(self.output_dir, exist_ok=True) - - -@dataclass -class SpeculativeDecodingConfig: - enabled: bool = field( - default=False, - metadata={"help": "Enable speculative decoding simulation."}, - ) - method: str = field( - default="eagle", - metadata={ - "help": "Speculative decoding method. Must match vLLM method names." - }, - ) - spec_model_name: str = field( - default="", - metadata={ - "help": "Optional draft/spec model name for methods whose proposer " - "decoder comes from a separate draft model (for example draft-model MTP)." - }, - ) - num_speculative_tokens: int = field( - default=4, - metadata={"help": "Number of draft tokens planned per speculative iteration."}, - ) - committed_tokens_per_iteration: int = field( - default=2, - metadata={ - "help": "Deterministic committed token count per speculative iteration " - "(includes 1 target token + accepted drafts)." - }, - ) - acceptance_trace_file: str = field( - default="", - metadata={ - "help": "Optional deterministic acceptance trace JSON file. Supported " - "formats: list[int] or {'committed_tokens_per_iteration': list[int], " - "'scheduled_draft_tokens_per_iteration': optional list[int], " - "'per_request_committed_tokens_per_iteration': optional dict[str, list[int]], " - "'per_request_scheduled_draft_tokens_per_iteration': optional dict[str, list[int]]}. " - "When set, trace overrides committed_tokens_per_iteration and can " - "optionally override planned draft widths per iteration." - }, - ) - proposer_overhead_ms_by_method: Dict[str, float] = field( - default_factory=dict, - metadata={ - "help": "Method-aware proposer overhead in milliseconds per speculative " - "verify request (method -> overhead_ms >= 0)." - }, - ) - decode_draft_proposer_latency_profile_file: str = field( - default="", - metadata={ - "help": "Optional structured latency profile JSON for decode draft " - "proposer overhead. Expected workload key: " - "(method, model_name, attn_tp_size, num_speculative_tokens, " - "spec_verify_request_count)." - }, - ) - mtp_n_predict: int = field( - default=0, - metadata={ - "help": "Optional MTP capability metadata. Number of tokens predicted " - "per MTP block. Only valid for MTP methods." - }, - ) - mtp_num_layers: int = field( - default=0, - metadata={ - "help": "Optional MTP capability metadata. Number of MTP layers. " - "Only valid for MTP methods." - }, - ) - trace_calibration_file: str = field( - default="", - metadata={ - "help": "Optional calibration JSON file. Supported keys: " - "proposer_overhead_ms_by_method and metadata." - }, - ) - - @staticmethod - def _validate_method_float_map( - *, - map_name: str, - raw_map: Optional[Dict[str, float]], - supported_methods: set[str], - min_value: float, - inclusive_min: bool, - ) -> Dict[str, float]: - if raw_map is None: - return {} - if not isinstance(raw_map, dict): - raise ValueError( - f"SpeculativeDecodingConfig.{map_name} must be a dict, " - f"got={type(raw_map).__name__}" - ) - - validated: Dict[str, float] = {} - for method_name, value in raw_map.items(): - if method_name not in supported_methods: - raise ValueError( - f"SpeculativeDecodingConfig.{map_name} contains unsupported method " - f"{method_name!r}; supported={sorted(supported_methods)}" - ) - numeric_value = float(value) - if inclusive_min: - if numeric_value < min_value: - raise ValueError( - f"SpeculativeDecodingConfig.{map_name}[{method_name!r}] " - f"must be >= {min_value}, got={numeric_value!r}" - ) - elif numeric_value <= min_value: - raise ValueError( - f"SpeculativeDecodingConfig.{map_name}[{method_name!r}] " - f"must be > {min_value}, got={numeric_value!r}" - ) - validated[method_name] = numeric_value - return validated - - @staticmethod - def _load_trace_calibration_payload( - trace_calibration_file: str, - ) -> Dict[str, Dict[str, float]]: - if not trace_calibration_file: - return {} - if not os.path.isfile(trace_calibration_file): - raise ValueError( - "SpeculativeDecodingConfig.trace_calibration_file does not exist: " - f"{trace_calibration_file!r}" - ) - try: - with open(trace_calibration_file, "r", encoding="utf-8") as f: - payload = json.load(f) - except json.JSONDecodeError as exc: - raise ValueError( - "SpeculativeDecodingConfig.trace_calibration_file must be valid JSON: " - f"{trace_calibration_file!r}" - ) from exc - - if not isinstance(payload, dict): - raise ValueError( - "SpeculativeDecodingConfig.trace_calibration_file must contain a JSON " - f"object, got={type(payload).__name__}" - ) - return payload - - @staticmethod - def _load_acceptance_trace_payload( - *, - acceptance_trace_file: str, - ): - if not acceptance_trace_file: - return None - if not os.path.isfile(acceptance_trace_file): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file does not exist: " - f"{acceptance_trace_file!r}" - ) - try: - with open(acceptance_trace_file, "r", encoding="utf-8") as f: - payload = json.load(f) - except json.JSONDecodeError as exc: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file must be valid JSON: " - f"{acceptance_trace_file!r}" - ) from exc - - if not isinstance(payload, (list, dict)): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file must be list or dict, " - f"got={type(payload).__name__}" - ) - return payload - - @staticmethod - def _load_committed_tokens_trace( - *, - acceptance_trace_payload, - max_committed_tokens: int, - ) -> Optional[List[int]]: - if acceptance_trace_payload is None: - return None - - if isinstance(acceptance_trace_payload, list): - committed_tokens_trace_raw = acceptance_trace_payload - else: - if "committed_tokens_per_iteration" not in acceptance_trace_payload: - return None - committed_tokens_trace_raw = acceptance_trace_payload[ - "committed_tokens_per_iteration" - ] - - if not isinstance(committed_tokens_trace_raw, list): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file committed token trace " - f"must be a list, got={type(committed_tokens_trace_raw).__name__}" - ) - if len(committed_tokens_trace_raw) == 0: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file committed token trace " - "must be non-empty." - ) - - committed_tokens_trace: List[int] = [] - for idx, value in enumerate(committed_tokens_trace_raw): - committed = int(value) - if committed < 0: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file values must be >= 0, " - f"got index={idx}, value={value!r}" - ) - if committed > max_committed_tokens: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file values must be <= " - f"1 + num_speculative_tokens ({max_committed_tokens}), " - f"got index={idx}, value={value!r}" - ) - committed_tokens_trace.append(committed) - return committed_tokens_trace - - @staticmethod - def _load_per_request_committed_tokens_trace( - *, - acceptance_trace_payload, - max_committed_tokens: int, - ) -> Optional[Dict[str, List[int]]]: - if acceptance_trace_payload is None or not isinstance( - acceptance_trace_payload, dict - ): - return None - if "per_request_committed_tokens_per_iteration" not in acceptance_trace_payload: - return None - - raw_trace_map = acceptance_trace_payload[ - "per_request_committed_tokens_per_iteration" - ] - if not isinstance(raw_trace_map, dict): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file " - "per_request_committed_tokens_per_iteration must be a dict, " - f"got={type(raw_trace_map).__name__}" - ) - if len(raw_trace_map) == 0: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file " - "per_request_committed_tokens_per_iteration must be non-empty." - ) - - per_request_trace: Dict[str, List[int]] = {} - for raw_request_id, raw_trace in raw_trace_map.items(): - request_id = str(raw_request_id) - if not request_id: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file per-request " - "trace keys must be non-empty strings." - ) - if request_id in per_request_trace: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file contains " - f"duplicate request_id={request_id!r} after normalization." - ) - if not isinstance(raw_trace, list): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file " - "per-request committed token trace must be a list, " - f"got request_id={request_id!r}, type={type(raw_trace).__name__}" - ) - if len(raw_trace) == 0: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file per-request " - f"committed token trace must be non-empty, request_id={request_id!r}" - ) - - validated_trace: List[int] = [] - for idx, value in enumerate(raw_trace): - committed = int(value) - if committed < 0: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file per-request " - "committed token values must be >= 0, " - f"got request_id={request_id!r}, index={idx}, value={value!r}" - ) - if committed > max_committed_tokens: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file per-request " - "committed token values must be <= 1 + num_speculative_tokens " - f"({max_committed_tokens}), got request_id={request_id!r}, " - f"index={idx}, value={value!r}" - ) - validated_trace.append(committed) - per_request_trace[request_id] = validated_trace - return per_request_trace - - @staticmethod - def _load_scheduled_draft_tokens_trace( - *, - acceptance_trace_payload, - max_scheduled_draft_tokens: int, - committed_trace_length: int, - ) -> Optional[List[int]]: - if acceptance_trace_payload is None or not isinstance(acceptance_trace_payload, dict): - return None - if "scheduled_draft_tokens_per_iteration" not in acceptance_trace_payload: - return None - - scheduled_draft_tokens_trace_raw = acceptance_trace_payload[ - "scheduled_draft_tokens_per_iteration" - ] - if not isinstance(scheduled_draft_tokens_trace_raw, list): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file scheduled draft token " - f"trace must be a list, got={type(scheduled_draft_tokens_trace_raw).__name__}" - ) - if len(scheduled_draft_tokens_trace_raw) != committed_trace_length: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file " - "scheduled_draft_tokens_per_iteration length must match " - "committed_tokens_per_iteration length." - ) - - scheduled_draft_tokens_trace: List[int] = [] - for idx, value in enumerate(scheduled_draft_tokens_trace_raw): - scheduled_drafts = int(value) - if scheduled_drafts < 0: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file scheduled draft " - "trace values must be >= 0, " - f"got index={idx}, value={value!r}" - ) - if scheduled_drafts > max_scheduled_draft_tokens: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file scheduled draft " - "trace values must be <= num_speculative_tokens " - f"({max_scheduled_draft_tokens}), got index={idx}, value={value!r}" - ) - scheduled_draft_tokens_trace.append(scheduled_drafts) - return scheduled_draft_tokens_trace - - @staticmethod - def _load_per_request_scheduled_draft_tokens_trace( - *, - acceptance_trace_payload, - max_scheduled_draft_tokens: int, - per_request_committed_trace: Optional[Dict[str, List[int]]], - ) -> Optional[Dict[str, List[int]]]: - if acceptance_trace_payload is None or not isinstance( - acceptance_trace_payload, dict - ): - return None - if ( - "per_request_scheduled_draft_tokens_per_iteration" - not in acceptance_trace_payload - ): - return None - if per_request_committed_trace is None: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file " - "per_request_scheduled_draft_tokens_per_iteration requires " - "per_request_committed_tokens_per_iteration." - ) - - raw_trace_map = acceptance_trace_payload[ - "per_request_scheduled_draft_tokens_per_iteration" - ] - if not isinstance(raw_trace_map, dict): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file " - "per_request_scheduled_draft_tokens_per_iteration must be a dict, " - f"got={type(raw_trace_map).__name__}" - ) - - normalized_keys = {str(request_id) for request_id in raw_trace_map.keys()} - committed_keys = set(per_request_committed_trace.keys()) - if normalized_keys != committed_keys: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file " - "per_request_scheduled_draft_tokens_per_iteration keys must match " - "per_request_committed_tokens_per_iteration keys." - ) - - per_request_trace: Dict[str, List[int]] = {} - for request_id, committed_trace in per_request_committed_trace.items(): - raw_trace = raw_trace_map[request_id] - if not isinstance(raw_trace, list): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file per-request " - "scheduled draft token trace must be a list, " - f"got request_id={request_id!r}, type={type(raw_trace).__name__}" - ) - if len(raw_trace) != len(committed_trace): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file " - "per_request_scheduled_draft_tokens_per_iteration length must " - "match per_request_committed_tokens_per_iteration length, " - f"request_id={request_id!r}" - ) - - validated_trace: List[int] = [] - for idx, value in enumerate(raw_trace): - scheduled_drafts = int(value) - if scheduled_drafts < 0: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file per-request " - "scheduled draft token values must be >= 0, " - f"got request_id={request_id!r}, index={idx}, value={value!r}" - ) - if scheduled_drafts > max_scheduled_draft_tokens: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file per-request " - "scheduled draft token values must be <= num_speculative_tokens " - f"({max_scheduled_draft_tokens}), got request_id={request_id!r}, " - f"index={idx}, value={value!r}" - ) - validated_trace.append(scheduled_drafts) - per_request_trace[request_id] = validated_trace - return per_request_trace - - def __post_init__(self) -> None: - supported_methods = { - "ngram", - "medusa", - "eagle", - "eagle3", - "deepseek_mtp", - "ernie_mtp", - "qwen3_moe_mtp", - "qwen3_next_mtp", - } - mtp_methods = { - "deepseek_mtp", - "ernie_mtp", - "qwen3_moe_mtp", - "qwen3_next_mtp", - } - if self.enabled and self.method not in supported_methods: - raise ValueError( - "SpeculativeDecodingConfig.method must match vLLM method names, " - f"got={self.method!r}, supported={sorted(supported_methods)}" - ) - if self.enabled and self.method in mtp_methods and self.mtp_n_predict <= 0: - raise ValueError( - "MTP methods require mtp_n_predict > 0 when enabled=True, " - f"got method={self.method!r}, mtp_n_predict={self.mtp_n_predict!r}" - ) - if self.enabled and self.method in mtp_methods and self.mtp_num_layers <= 0: - raise ValueError( - "MTP methods require mtp_num_layers > 0 when enabled=True, " - f"got method={self.method!r}, mtp_num_layers={self.mtp_num_layers!r}" - ) - if self.mtp_n_predict < 0: - raise ValueError( - "SpeculativeDecodingConfig.mtp_n_predict must be >= 0, " - f"got={self.mtp_n_predict!r}" - ) - if self.mtp_num_layers < 0: - raise ValueError( - "SpeculativeDecodingConfig.mtp_num_layers must be >= 0, " - f"got={self.mtp_num_layers!r}" - ) - if self.mtp_n_predict > 0 and self.method not in mtp_methods: - raise ValueError( - "SpeculativeDecodingConfig.mtp_n_predict is only valid for MTP " - f"methods, got method={self.method!r}" - ) - if self.mtp_num_layers > 0 and self.method not in mtp_methods: - raise ValueError( - "SpeculativeDecodingConfig.mtp_num_layers is only valid for MTP " - f"methods, got method={self.method!r}" - ) - if self.enabled and self.num_speculative_tokens <= 0: - raise ValueError( - "SpeculativeDecodingConfig.num_speculative_tokens must be > 0 when " - f"enabled=True, got={self.num_speculative_tokens}" - ) - if ( - self.method in mtp_methods - and self.mtp_n_predict > 0 - and self.num_speculative_tokens % self.mtp_n_predict != 0 - ): - raise ValueError( - "SpeculativeDecodingConfig.num_speculative_tokens must be divisible " - "by mtp_n_predict when mtp_n_predict > 0 for MTP methods, " - f"got num_speculative_tokens={self.num_speculative_tokens}, " - f"mtp_n_predict={self.mtp_n_predict}" - ) - max_committed_tokens = int(self.num_speculative_tokens) + 1 - if self.committed_tokens_per_iteration < 1: - raise ValueError( - "SpeculativeDecodingConfig.committed_tokens_per_iteration must be >= 1, " - f"got={self.committed_tokens_per_iteration!r}" - ) - if self.committed_tokens_per_iteration > max_committed_tokens: - raise ValueError( - "SpeculativeDecodingConfig.committed_tokens_per_iteration must be <= " - f"1 + num_speculative_tokens ({max_committed_tokens}), " - f"got={self.committed_tokens_per_iteration!r}" - ) - acceptance_trace_payload = self._load_acceptance_trace_payload( - acceptance_trace_file=self.acceptance_trace_file, - ) - self._committed_tokens_trace = self._load_committed_tokens_trace( - acceptance_trace_payload=acceptance_trace_payload, - max_committed_tokens=max_committed_tokens, - ) - self._per_request_committed_tokens_trace = ( - self._load_per_request_committed_tokens_trace( - acceptance_trace_payload=acceptance_trace_payload, - max_committed_tokens=max_committed_tokens, - ) - ) - if ( - acceptance_trace_payload is not None - and self._committed_tokens_trace is None - and self._per_request_committed_tokens_trace is None - ): - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file JSON object must " - "contain key 'committed_tokens_per_iteration' or " - "'per_request_committed_tokens_per_iteration'." - ) - self._scheduled_draft_tokens_trace = self._load_scheduled_draft_tokens_trace( - acceptance_trace_payload=acceptance_trace_payload, - max_scheduled_draft_tokens=int(self.num_speculative_tokens), - committed_trace_length=( - len(self._committed_tokens_trace) - if self._committed_tokens_trace is not None - else 0 - ), - ) - self._per_request_scheduled_draft_tokens_trace = ( - self._load_per_request_scheduled_draft_tokens_trace( - acceptance_trace_payload=acceptance_trace_payload, - max_scheduled_draft_tokens=int(self.num_speculative_tokens), - per_request_committed_trace=self._per_request_committed_tokens_trace, - ) - ) - if self._scheduled_draft_tokens_trace is not None: - for idx, (committed_tokens, scheduled_draft_tokens) in enumerate( - zip( - self._committed_tokens_trace, - self._scheduled_draft_tokens_trace, - ) - ): - if committed_tokens > 1 + scheduled_draft_tokens: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file committed " - "tokens must be <= 1 + scheduled_draft_tokens_per_iteration, " - f"got index={idx}, committed={committed_tokens}, " - f"scheduled_draft_tokens={scheduled_draft_tokens}" - ) - if self._per_request_scheduled_draft_tokens_trace is not None: - for request_id, committed_trace in ( - self._per_request_committed_tokens_trace.items() - ): - scheduled_trace = self._per_request_scheduled_draft_tokens_trace[ - request_id - ] - for idx, (committed_tokens, scheduled_draft_tokens) in enumerate( - zip(committed_trace, scheduled_trace) - ): - if committed_tokens > 1 + scheduled_draft_tokens: - raise ValueError( - "SpeculativeDecodingConfig.acceptance_trace_file per-request " - "committed tokens must be <= 1 + " - "per_request_scheduled_draft_tokens_per_iteration, " - f"got request_id={request_id!r}, index={idx}, " - f"committed={committed_tokens}, " - f"scheduled_draft_tokens={scheduled_draft_tokens}" - ) - - trace_payload = self._load_trace_calibration_payload(self.trace_calibration_file) - supported_trace_keys = { - "proposer_overhead_ms_by_method", - "metadata", - } - unexpected_keys = sorted(set(trace_payload.keys()) - supported_trace_keys) - if unexpected_keys: - raise ValueError( - "Unsupported keys in trace calibration file: " - f"{unexpected_keys}, supported={sorted(supported_trace_keys)}" - ) - - trace_proposer_overheads = self._validate_method_float_map( - map_name="proposer_overhead_ms_by_method", - raw_map=trace_payload.get("proposer_overhead_ms_by_method", {}), - supported_methods=supported_methods, - min_value=0.0, - inclusive_min=True, - ) - config_proposer_overheads = self._validate_method_float_map( - map_name="proposer_overhead_ms_by_method", - raw_map=self.proposer_overhead_ms_by_method, - supported_methods=supported_methods, - min_value=0.0, - inclusive_min=True, - ) - - # Config-driven values override trace-derived values for deterministic control. - self.proposer_overhead_ms_by_method = { - **trace_proposer_overheads, - **config_proposer_overheads, - } - self._decode_draft_proposer_latency_profile = ( - load_decode_draft_proposer_latency_profile( - profile_file=self.decode_draft_proposer_latency_profile_file, - supported_methods=supported_methods, - ) - ) - - -@dataclass -class ReplicaConfig: - memory_margin_fraction: float = field( - default=0.1, - metadata={"help": "Memory margin fraction."}, - ) - num_pipeline_stages: int = field( - default=1, - metadata={"help": "Number of pipeline stages (pp size)."}, - ) - attn_tensor_parallel_size: int = field( - default=1, - metadata={"help": "Attention tensor parallel size (attn_tp size)."}, - ) - attn_dp: int = field( - default=1, - metadata={ - "help": "Attention data-parallel lanes owned by one Replica.", - }, - ) - moe_tensor_parallel_size: int = field( - default=1, - metadata={"help": "MoE tensor parallel size (moe_tp size)."}, - ) - moe_expert_parallel_size: int = field( - default=1, - metadata={"help": "MoE expert parallel size (moe_ep size)."}, - ) - total_expert_num: int = field( - default=1, - metadata={"help": "Total expert number."}, - ) - router_load_balancing_type: str = field( - default="None", - metadata={"help": "MOE router load balancing type."}, - ) - router_topk: int = field( - default=0, - metadata={"help": "Router topk. Set to 0 to inherit from model config."}, - ) - moe_routing_seed: int = field( - default=42, - metadata={ - "help": "Random seed for deterministic MoE routing distribution generation. " - "Must be a non-negative integer." - }, - ) - moe_routing_trace_path: str = field( - default="", - metadata={ - "help": "Deferred StepFun merged trace JSONL for unsupported trace replay. " - "A non-empty path fails fast at the architecture boundary." - }, - ) - decode_attn_initial_lane_trace_path: str = field( - default="", - metadata={ - "help": "Optional StepFun attention trace JSONL for trace-driven " - "decode-attn initial lane occupancy and warmup replay." - }, - ) - decode_attn_steady_state_snapshot_path: str = field( - default="", - metadata={ - "help": "Optional StepFun attention trace JSONL for explicit " - "decode-attn steady-state snapshot hydration." - }, - ) - decode_attn_steady_state_measurement_report_path: str = field( - default="", - metadata={ - "help": "Optional StepFun measurement JSON for post-boundary " - "decode-attn request arrival replay." - }, - ) - moe_routing_distribution_type: str = field( - default="balanced", - metadata={ - "help": "MoE expert-load distribution for disaggregated routing simulation. " - "Valid values: 'balanced', 'random', 'skewed', or 'zipf'. This controls " - "token-to-expert load skew without changing router_topk/model semantics." - }, - ) - device: str = field( - default="a100", - metadata={"help": "Device."}, - ) - network_device: str = field( - default="a100_pairwise_nvlink", - metadata={"help": "Network device."}, - ) - speculative_decoding_config: SpeculativeDecodingConfig = field( - default_factory=SpeculativeDecodingConfig, - metadata={"help": "Speculative decoding simulation configuration."}, - ) - - # configs should be set by the user - cluster_prefix: str = None - local_expert_num: int = None - model_name: str = "meta-llama/Llama-2-7b-hf" - - def __post_init__(self): - if type(self.attn_dp) is not int or self.attn_dp <= 0: - raise ValueError( - "attn_dp must be a positive integer, " - f"got {self.attn_dp!r}" - ) - if self.cluster_prefix == "decode_attn" and self.attn_dp != 1: - raise ValueError( - "DECODE_ATTN requires attn_dp=1 because it is the PD-AF attention role" - ) - # Load model and device configs first (needed for validation) - self.model_config: BaseModelConfig = BaseModelConfig.create_from_name( - self.model_name - ) - self.device_config: BaseDeviceSKUConfig = ( - BaseDeviceSKUConfig.create_from_type_string(self.device) - ) - self.node_config: BaseNodeSKUConfig = BaseNodeSKUConfig.create_from_type_string( - self.network_device - ) - - # Auto-set total_expert_num from model config if not explicitly set and model is MoE - if ( - self.total_expert_num == 1 - and self.model_config.is_moe - and self.model_config.num_experts > 0 - ): - self.total_expert_num = self.model_config.num_experts - - # Align router_topk with model config when not explicitly set. - if self.model_config.is_moe: - if self.router_topk is None or int(self.router_topk) <= 0: - if self.model_config.num_experts_per_tok > 0: - self.router_topk = int(self.model_config.num_experts_per_tok) - else: - raise ValueError( - "router_topk is not set and model_config.num_experts_per_tok is missing" - ) - else: - if self.router_topk is None or int(self.router_topk) <= 0: - self.router_topk = 1 - - valid_moe_routing_distribution_types = { - "balanced", - "random", - "skewed", - "zipf", - } - self.moe_routing_distribution_type = str( - self.moe_routing_distribution_type - ).strip().lower() - if self.moe_routing_distribution_type not in valid_moe_routing_distribution_types: - raise ValueError( - "moe_routing_distribution_type must be one of " - f"{sorted(valid_moe_routing_distribution_types)}, " - f"got {self.moe_routing_distribution_type!r}" - ) - - # Validate pipeline parallelism configuration early - if self.model_config.num_layers % self.num_pipeline_stages != 0: - raise ValueError( - f"Pipeline parallelism configuration error: " - f"num_layers ({self.model_config.num_layers}) must be evenly divisible by " - f"num_pipeline_stages ({self.num_pipeline_stages}). " - f"Current configuration would result in uneven layer distribution across pipeline stages. " - f"Please adjust num_pipeline_stages to be a divisor of {self.model_config.num_layers}." - ) - - # Note: this world_size only limits in replica dimension. - if self.cluster_prefix == "prefill": - self.world_size = ( - self.num_pipeline_stages - * self.attn_tensor_parallel_size - * self.attn_dp - ) - elif self.cluster_prefix == "decode_attn": - self.world_size = ( - self.num_pipeline_stages - * self.attn_tensor_parallel_size - * self.attn_dp - ) - elif self.cluster_prefix == "decode_ffn": - self.world_size = ( - self.num_pipeline_stages - * self.moe_tensor_parallel_size - * self.moe_expert_parallel_size - ) - elif self.cluster_prefix == "decode": - # Unified decode cluster (PD-disaggregation): similar to prefill, includes both Attention and FFN - self.world_size = ( - self.num_pipeline_stages - * self.attn_tensor_parallel_size - * self.attn_dp - ) - else: # Monolithic - self.world_size = ( - self.num_pipeline_stages - * self.attn_tensor_parallel_size - * self.attn_dp - ) - - # Validate expert parallelism configuration for MoE models - # Use model_config.is_moe for MoE detection - NOT total_expert_num - if self.cluster_prefix != "decode_attn" and self.model_config.is_moe: - if self.total_expert_num > 1: - assert ( - self.total_expert_num % self.moe_expert_parallel_size == 0 - ), "total_expert_num must be divisible by moe_expert_parallel_size" - self.local_expert_num = ( - self.total_expert_num // self.moe_expert_parallel_size - ) - - if ( - self.speculative_decoding_config.enabled - and self.cluster_prefix in {"decode_attn", "decode_ffn"} - ): - raise ValueError( - "Speculative decoding Phase 1 supports only co-location and " - "pd-disaggregation decode path. decode_attn/decode_ffn are not " - f"supported, cluster_prefix={self.cluster_prefix!r}." - ) - - from frontier.attention.gdn.guards import validate_gdn_runtime_support - - validate_gdn_runtime_support( - self.model_config, - speculative_enabled=bool(self.speculative_decoding_config.enabled), - num_pipeline_stages=self.num_pipeline_stages, - moe_expert_parallel_size=self.moe_expert_parallel_size, - attn_dp=self.attn_dp, - cross_node=( - int(self.world_size) > int(self.node_config.num_devices_per_node) - ), - ) - - -@dataclass -class BaseClusterSchedulerConfig(BasePolyConfig): - pass - - -@dataclass -class RandomClusterSchedulerConfig(BaseClusterSchedulerConfig): - @staticmethod - def get_type(): - return ClusterSchedulerType.RANDOM - - -@dataclass -class RoundRobinClusterSchedulerConfig(BaseClusterSchedulerConfig): - @staticmethod - def get_type(): - return ClusterSchedulerType.ROUND_ROBIN - - -@dataclass -class LORClusterSchedulerConfig(BaseClusterSchedulerConfig): - @staticmethod - def get_type(): - return ClusterSchedulerType.LOR - - -@dataclass -class StickyRoundRobinClusterSchedulerConfig(BaseClusterSchedulerConfig): - @staticmethod - def get_type(): - return ClusterSchedulerType.STICKY_ROUND_ROBIN - - -@dataclass -class StickyLORClusterSchedulerConfig(BaseClusterSchedulerConfig): - @staticmethod - def get_type(): - return ClusterSchedulerType.STICKY_LOR - - -@dataclass -class BaseExecutionTimePredictorConfig(BasePolyConfig): - linear_op_input_file: str = field( - default="./data/profiling/compute/{DEVICE}/{MODEL}/linear_op.csv", - metadata={"help": "Path to the linear operation profiling input file."}, - ) - # Backward compatibility alias - mlp_input_file: str = field( - default="", - metadata={"help": "[DEPRECATED] Use linear_op_input_file instead."}, - ) - atten_input_file: str = field( - default="./data/profiling/compute/{DEVICE}/{MODEL}/attention.csv", - metadata={"help": "Path to the attention input file."}, - ) - gdn_input_file: str = field( - default="./data/profiling/compute/{DEVICE}/{MODEL}/gdn.csv", - metadata={"help": "Path to the standard GDN profiling input file."}, - ) - all_reduce_input_file: str = field( - default="./data/profiling/network/{NETWORK_DEVICE}/all_reduce.csv", - metadata={"help": "Path to the all reduce input file."}, - ) - send_recv_input_file: str = field( - default="./data/profiling/network/{NETWORK_DEVICE}/send_recv.csv", - metadata={"help": "Path to the send recv input file."}, - ) - cpu_overhead_input_file: str = field( - default="./data/profiling/cpu_overhead/{NETWORK_DEVICE}/{MODEL}/cpu_overheads.csv", - metadata={"help": "Path to the cpu overhead input file."}, - ) - cpu_overhead_kernel_only_input_file: str = field( - default="./data/profiling/cpu_overhead/{NETWORK_DEVICE}/{MODEL}/cpu_overheads_kernel_only.csv", - metadata={"help": "Path to the kernel-only cpu overhead input file."}, - ) - pp_stage_boundary_input_file: str = field( - default="./data/profiling/other_overhead/{DEVICE}/{MODEL}/pp_stage_boundary.csv", - metadata={"help": "Path to the pipeline stage-boundary overhead input file."}, - ) - pp_receiver_head_input_file: str = field( - default="./data/profiling/other_overhead/{DEVICE}/{MODEL}/pp_receiver_head.csv", - metadata={"help": "Path to the PP receiver-head overhead input file."}, - ) - pp_producer_send_path_input_file: str = field( - default="./data/profiling/other_overhead/{DEVICE}/{MODEL}/pp_producer_send_path.csv", - metadata={"help": "Path to the PP producer send-path overhead input file."}, - ) - pp_prefill_consumer_active_input_file: str = field( - default="./data/profiling/other_overhead/{DEVICE}/{MODEL}/pp_prefill_consumer_active.csv", - metadata={ - "help": "Path to the PP prefill consumer-active overhead input file." - }, - ) - moe_input_file: str = field( - default="./data/profiling/compute/{DEVICE}/{MODEL}/moe.csv", - metadata={"help": "Path to the MoE profiling input file."}, - ) - linear_op_kernel_only_input_file: str = field( - default="./data/profiling/compute/{DEVICE}/{MODEL}/linear_op_kernel_only.csv", - metadata={"help": "Path to the kernel-only linear operation profiling input file."}, - ) - atten_kernel_only_input_file: str = field( - default="./data/profiling/compute/{DEVICE}/{MODEL}/attention_kernel_only.csv", - metadata={"help": "Path to the kernel-only attention input file."}, - ) - moe_kernel_only_input_file: str = field( - default="./data/profiling/compute/{DEVICE}/{MODEL}/moe_kernel_only.csv", - metadata={"help": "Path to the kernel-only MoE profiling input file."}, - ) - k_fold_cv_splits: int = field( - default=10, - metadata={"help": "Number of k fold cross validation splits."}, - ) - no_cache: bool = field( - default=False, - metadata={"help": "Whether to cache prediction models."}, - ) - kv_cache_prediction_granularity: int = field( - default=64, - metadata={"help": "KV cache prediction granularity."}, - ) - prediction_max_prefill_chunk_size: int = field( - default=4096, - metadata={"help": "Max prefill chunk size for prediction."}, - ) - prediction_max_batch_size: int = field( - default=128, - metadata={"help": "Max batch size for prediction."}, - ) - prediction_max_tokens_per_request: int = field( - default=4096, - metadata={"help": "Max tokens per request for prediction."}, - ) - attention_decode_batching_overhead_fraction: float = field( - default=0.1, - metadata={"help": "Attention decode batching overhead fraction."}, - ) - attention_prefill_batching_overhead_fraction: float = field( - default=0.1, - metadata={"help": "Attention prefill batching overhead fraction."}, - ) - attn_pre_proj_calibration_scale: float = field( - default=1.0, - metadata={ - "help": "Multiplicative calibration scale for attn_pre_proj prediction. Must be > 0." - }, - ) - prefill_phase_attn_pre_proj_calibration_scale: Optional[float] = field( - default=None, - metadata={ - "help": ( - "Optional multiplicative calibration scale for attn_pre_proj " - "prediction when the batch includes prefill tokens. Must be > 0." - ) - }, - ) - attn_post_proj_calibration_scale: float = field( - default=1.0, - metadata={ - "help": "Multiplicative calibration scale for attn_post_proj prediction. Must be > 0." - }, - ) - prefill_phase_attn_post_proj_calibration_scale: Optional[float] = field( - default=None, - metadata={ - "help": ( - "Optional multiplicative calibration scale for attn_post_proj " - "prediction when the batch includes prefill tokens. Must be > 0." - ) - }, - ) - attn_decode_calibration_scale: float = field( - default=1.0, - metadata={ - "help": "Multiplicative calibration scale for attn_decode prediction. Must be > 0." - }, - ) - attn_decode_in_mixed_calibration_scale: Optional[float] = field( - default=None, - metadata={ - "help": ( - "Optional multiplicative calibration scale for attn_decode_in_mixed " - "prediction when a co-location batch contains both prefill and decode " - "tokens. Must be > 0." - ) - }, - ) - late_decode_attn_decode_calibration_scale: Optional[float] = field( - default=None, - metadata={ - "help": ( - "Optional multiplicative calibration scale for attn_decode " - "prediction when every decode request in the batch has already " - "completed the first pure decode token. Must be > 0." - ) - }, - ) - attn_kv_cache_save_calibration_scale: float = field( - default=1.0, - metadata={ - "help": "Multiplicative calibration scale for attn_kv_cache_save prediction. Must be > 0." - }, - ) - prefill_phase_attn_kv_cache_save_calibration_scale: Optional[float] = field( - default=None, - metadata={ - "help": ( - "Optional multiplicative calibration scale for attn_kv_cache_save " - "prediction when the batch includes prefill tokens. Must be > 0." - ) - }, - ) - mlp_up_proj_calibration_scale: float = field( - default=1.0, - metadata={ - "help": "Multiplicative calibration scale for mlp_up_proj prediction. Must be > 0." - }, - ) - prefill_phase_mlp_up_proj_calibration_scale: Optional[float] = field( - default=None, - metadata={ - "help": ( - "Optional multiplicative calibration scale for mlp_up_proj " - "prediction when the batch includes prefill tokens. Must be > 0." - ) - }, - ) - mlp_down_proj_calibration_scale: float = field( - default=1.0, - metadata={ - "help": "Multiplicative calibration scale for mlp_down_proj prediction. Must be > 0." - }, - ) - decode_phase_mlp_down_proj_calibration_scale: Optional[float] = field( - default=None, - metadata={ - "help": ( - "Optional multiplicative calibration scale for mlp_down_proj " - "prediction when the batch contains decode tokens but no " - "prefill tokens. Must be > 0." - ) - }, - ) - nccl_cpu_launch_overhead_ms: float = field( - default=0.02, - metadata={"help": "NCCL CPU launch overhead in ms."}, - ) - nccl_cpu_skew_overhead_per_device_ms: float = field( - default=0.0, - metadata={"help": "NCCL CPU skew overhead per device in ms."}, - ) - num_training_job_threads: int = field( - default=-1, - metadata={"help": "Number of training job threads."}, - ) - skip_cpu_overhead_modeling: bool = field( - default=True, - metadata={"help": "Whether to skip CPU overhead modeling."}, - ) - - # Dummy mode configuration for fast testing and development - enable_dummy_mode: bool = field( - default=False, - metadata={ - "help": "Enable dummy mode to skip ML model training and return fixed execution times." - }, - ) - dummy_execution_time_ms: float = field( - default=1.0, - metadata={ - "help": "Fixed execution time in milliseconds to return in dummy mode." - }, - ) - - def __post_init__(self) -> None: - for field_name in ( - "attn_pre_proj_calibration_scale", - "prefill_phase_attn_pre_proj_calibration_scale", - "attn_post_proj_calibration_scale", - "prefill_phase_attn_post_proj_calibration_scale", - "attn_decode_calibration_scale", - "attn_decode_in_mixed_calibration_scale", - "late_decode_attn_decode_calibration_scale", - "attn_kv_cache_save_calibration_scale", - "prefill_phase_attn_kv_cache_save_calibration_scale", - "mlp_up_proj_calibration_scale", - "prefill_phase_mlp_up_proj_calibration_scale", - "mlp_down_proj_calibration_scale", - "decode_phase_mlp_down_proj_calibration_scale", - ): - raw_value = getattr(self, field_name) - if raw_value is None: - continue - value = float(raw_value) - if value <= 0.0: - raise ValueError( - f"{self.__class__.__name__}.{field_name} must be > 0, got={value!r}" - ) - def validate_linear_op_input(self) -> None: - """Validate linear_op_input_file configuration. - - Raises: - ValueError: If mlp.csv path is used or linear_op_input_file is empty. - """ - # Reject mlp.csv paths - if self.linear_op_input_file and "mlp.csv" in self.linear_op_input_file: - raise ValueError( - f"mlp.csv is forbidden in linear_op_input_file. " - f"Use linear_op.csv instead. Got: {self.linear_op_input_file}" - ) - - # Reject empty path when validation is explicitly called - if not self.linear_op_input_file: - raise ValueError( - "linear_op_input_file must be set to a valid linear_op.csv path." - ) - - # Warn if deprecated mlp_input_file is used - if self.mlp_input_file: - import logging - logger = logging.getLogger(__name__) - logger.warning( - "mlp_input_file is deprecated and will be ignored. " - "Use linear_op_input_file with linear_op.csv instead." - ) - - - -@dataclass -class LinearRegressionExecutionTimePredictorConfig(BaseExecutionTimePredictorConfig): - polynomial_degree: List[int] = field( - default_factory=lambda: list(range(1, 6)), - metadata={"help": "Polynomial degree for linear regression."}, - ) - polynomial_include_bias: List[bool] = field( - default_factory=lambda: [True, False], - metadata={"help": "Polynomial include bias for linear regression."}, - ) - polynomial_interaction_only: List[bool] = field( - default_factory=lambda: [True, False], - metadata={"help": "Polynomial interaction only for linear regression."}, - ) - fit_intercept: List[bool] = field( - default_factory=lambda: [True, False], - metadata={"help": "Fit intercept for linear regression."}, - ) - - @staticmethod - def get_type(): - return ExecutionTimePredictorType.LINEAR_REGRESSION - - -@dataclass -class RandomForrestExecutionTimePredictorConfig(BaseExecutionTimePredictorConfig): - num_estimators: List[int] = field( - default_factory=lambda: [250, 500, 750], - metadata={"help": "Number of estimators for random forest."}, - ) - max_depth: List[int] = field( - default_factory=lambda: [8, 16, 32], - metadata={"help": "Maximum depth for random forest."}, - ) - min_samples_split: List[int] = field( - default_factory=lambda: [2, 5, 10], - metadata={"help": "Minimum samples split for random forest."}, - ) - - @staticmethod - def get_type(): - return ExecutionTimePredictorType.RANDOM_FORREST - - -@dataclass -class ClusterConfig: - # === Common fields for all modes === - cluster_scheduler_config: BaseClusterSchedulerConfig = field( - default_factory=RoundRobinClusterSchedulerConfig, - metadata={ - "help": "Cluster scheduler config.", - }, - ) - replica_scheduler_config: BaseReplicaSchedulerConfig = field( - default_factory=SarathiSchedulerConfig, - metadata={"help": "Replica scheduler config."}, - ) - cluster_type: ClusterType = field( - default=None, - metadata={ - "help": "Type of the cluster: monolithic, prefill, decode-attn, or decode-ffn." - }, - ) - execution_time_predictor_config: BaseExecutionTimePredictorConfig = field( - default_factory=RandomForrestExecutionTimePredictorConfig, - metadata={"help": "Execution time predictor config."}, - ) - cc_backend_config: BaseCCBackendConfig = field( - default_factory=lambda: _get_cc_backend_configs()[5](), # AstraSimAnalyticalCCBackendConfig - metadata={ - "help": "CC (Collective Communication) backend config for communication latency prediction." - }, - ) - - # === co-location/Monolithic mode fields === - num_replicas: Optional[int] = field( - default=1, - metadata={ - "help": "Number of replicas", - }, - ) - replica_config: Optional[ReplicaConfig] = field( - default_factory=lambda: ReplicaConfig(model_name="meta-llama/Llama-2-7b-hf"), - metadata={ - "help": "Replica configuration", - }, - ) - - # === Disaggregated mode fields === - prefill_cluster_num_replicas: Optional[int] = field( - default=None, - metadata={ - "help": "Number of replicas for prefill cluster. Used only in pd-af-disaggregation mode.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cluster_num_replicas: Optional[int] = field( - default=None, - metadata={ - "help": "Number of replicas for decode attention cluster. Used only in pd-af-disaggregation mode.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cluster_num_replicas: Optional[int] = field( - default=None, - metadata={ - "help": "Number of replicas for decode FFN cluster. " - "Each replica is an independent FFN serving copy; MoE EP lanes are " - "scoped inside the selected replica. Used only in pd-af-disaggregation mode.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_cluster_num_replicas: Optional[int] = field( - default=None, - metadata={ - "help": "Number of replicas for unified decode cluster. Used only in pd-disaggregation mode.", - "mode_dependency": "pd-disaggregation", - }, - ) - prefill_replica_config_memory_margin_fraction: Optional[float] = field( - default=None, - metadata={ - "help": "Memory margin fraction for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_num_pipeline_stages: Optional[int] = field( - default=None, - metadata={ - "help": "Number of pipeline stages for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_attn_tensor_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "Attention tensor parallel size for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_moe_tensor_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "MoE tensor parallel size for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_moe_expert_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "MoE expert parallel size for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_total_expert_num: Optional[int] = field( - default=None, - metadata={ - "help": "Total expert number for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_local_expert_num: Optional[int] = field( - default=None, - metadata={ - "help": "Local expert number for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_router_load_balancing_type: Optional[str] = field( - default=None, - metadata={ - "help": "MOE router load balancing type for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_router_topk: Optional[int] = field( - default=None, - metadata={ - "help": "Router topk for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_device: Optional[str] = field( - default=None, - metadata={ - "help": "Device for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - prefill_replica_config_network_device: Optional[str] = field( - default=None, - metadata={ - "help": "Network device for prefill cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_config_memory_margin_fraction: Optional[float] = field( - default=None, - metadata={ - "help": "Memory margin fraction for decode attention cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_config_num_pipeline_stages: Optional[int] = field( - default=None, - metadata={ - "help": "Number of pipeline stages for decode attention cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_config_attn_tensor_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "Attention tensor parallel size for decode attention cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_config_device: Optional[str] = field( - default=None, - metadata={ - "help": "Device for decode attention cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_config_network_device: Optional[str] = field( - default=None, - metadata={ - "help": "Network device for decode attention cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_memory_margin_fraction: Optional[float] = field( - default=None, - metadata={ - "help": "Memory margin fraction for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_num_pipeline_stages: Optional[int] = field( - default=None, - metadata={ - "help": "Number of pipeline stages for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_moe_tensor_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "MoE tensor parallel size for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_moe_expert_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "MoE expert parallel size for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_total_expert_num: Optional[int] = field( - default=None, - metadata={ - "help": "Total expert number for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_local_expert_num: Optional[int] = field( - default=None, - metadata={ - "help": "Local expert number for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_router_load_balancing_type: Optional[str] = field( - default=None, - metadata={ - "help": "MOE router load balancing type for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_router_topk: Optional[int] = field( - default=None, - metadata={ - "help": "Router topk for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_device: Optional[str] = field( - default=None, - metadata={ - "help": "Device for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_config_network_device: Optional[str] = field( - default=None, - metadata={ - "help": "Network device for decode FFN cluster.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - - # === PD-Disaggregation Mode: Unified DECODE Cluster Configuration === - decode_replica_config_memory_margin_fraction: Optional[float] = field( - default=None, - metadata={ - "help": "Memory margin fraction for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_num_pipeline_stages: Optional[int] = field( - default=None, - metadata={ - "help": "Number of pipeline stages for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_attn_tensor_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "Attention tensor parallel size for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_moe_tensor_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "MoE tensor parallel size for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_moe_expert_parallel_size: Optional[int] = field( - default=None, - metadata={ - "help": "MoE expert parallel size for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_total_expert_num: Optional[int] = field( - default=None, - metadata={ - "help": "Total expert number for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_local_expert_num: Optional[int] = field( - default=None, - metadata={ - "help": "Local expert number for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_router_load_balancing_type: Optional[str] = field( - default=None, - metadata={ - "help": "MOE router load balancing type for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_router_topk: Optional[int] = field( - default=None, - metadata={ - "help": "Router topk for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_device: Optional[str] = field( - default=None, - metadata={ - "help": "Device for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_config_network_device: Optional[str] = field( - default=None, - metadata={ - "help": "Network device for unified decode cluster.", - "mode_dependency": "pd-disaggregation", - }, - ) - - # === AF Pipeline Configuration === - # This field is for internal use by the created cluster-specific configs - af_pipeline_num_micro_batch: int = field( - default=-1, - metadata={ - "help": "Internal field for the number of micro-batches. Should be set via cluster-specific parameters below.", - }, - ) - - # User-facing parameters for setting micro-batch number in decode clusters - decode_attn_af_pipeline_num_micro_batch: Optional[int] = field( - default=None, - metadata={"help": "Number of micro-batches for the decode_attn cluster."}, - ) - decode_ffn_af_pipeline_num_micro_batch: Optional[int] = field( - default=None, - metadata={"help": "Number of micro-batches for the decode_ffn cluster."}, - ) - - # User-facing parameter for setting micro-batch SIZE specifically for decode-attn - decode_attn_micro_batch_size: Optional[int] = field( - default=None, - metadata={ - "help": "Target micro-batch SIZE for decode-attn cluster (per (replica, dp)).", - }, - ) - - # User-facing parameter for setting request allocation threshold for decode-attn - decode_attn_request_allocation_threshold: Optional[int] = field( - default=None, - metadata={ - "help": "Request accumulation threshold for decode-attn cluster. " - "Only trigger allocation when accumulated requests reach this threshold. " - "Default: None (equals total number of requests in offline mode).", - }, - ) - - # === AFD CUDA Graph Configuration === - # Aligned with StepFun-vLLM's cudagraph_batch_sizes for AFD attention server - decode_attn_use_cuda_graph: bool = field( - default=False, - metadata={ - "help": "Deprecated. Use SimulationConfig.use_cuda_graph instead. " - "CUDA Graph is now a global setting for pd-af-disaggregation.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cudagraph_capture_sizes: Optional[List[int]] = field( - default=None, - metadata={ - "help": "Deprecated. Use SimulationConfig.cudagraph_capture_sizes instead. " - "CUDA Graph capture sizes are now shared across decode-attn and decode-ffn.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - - decode_attn_replica_id_start_for_ffn: Optional[int] = field( - default=None, - metadata={ - "help": "Derived first global replica id for DECODE_ATTN lanes; used by DECODE_FFN grouping.", - }, - ) - - # === Per-Cluster Replica Scheduler Configuration === - # These fields allow per-cluster-type customization of replica scheduler parameters - # If not set, they fall back to the base replica_scheduler_config values - - # PREFILL cluster scheduler configuration - prefill_replica_scheduler_config_type: Optional[str] = field( - default=None, - metadata={ - "help": "Replica scheduler type for prefill cluster. Overrides base replica_scheduler_config_type.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_replica_scheduler_config_batch_size_cap: Optional[int] = field( - default=None, - metadata={ - "help": "Batch size cap (max_num_seqs) for prefill cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_replica_scheduler_config_max_tokens_in_batch: Optional[int] = field( - default=None, - metadata={ - "help": "Max tokens in batch (max_num_batched_tokens) for prefill cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_replica_scheduler_config_enable_chunked_prefill: Optional[bool] = field( - default=None, - metadata={ - "help": "Enable Chunked Prefill for the prefill cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_replica_scheduler_config_long_prefill_token_threshold: Optional[int] = ( - field( - default=None, - metadata={ - "help": "Long-prefill token threshold for the prefill cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - ) - prefill_replica_scheduler_config_num_blocks: Optional[int] = field( - default=None, - metadata={ - "help": "Number of blocks for prefill cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_replica_scheduler_config_block_size: Optional[int] = field( - default=None, - metadata={ - "help": "Block size for prefill cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_replica_scheduler_config_watermark_blocks_fraction: Optional[float] = field( - default=None, - metadata={ - "help": "Watermark blocks fraction for prefill cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - - # DECODE cluster scheduler configuration (for unified decode in pd-disaggregation mode) - decode_replica_scheduler_config_type: Optional[str] = field( - default=None, - metadata={ - "help": "Replica scheduler type for decode cluster. Overrides base replica_scheduler_config_type.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_scheduler_config_batch_size_cap: Optional[int] = field( - default=None, - metadata={ - "help": "Batch size cap (max_num_seqs) for decode cluster replica scheduler.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_scheduler_config_max_tokens_in_batch: Optional[int] = field( - default=None, - metadata={ - "help": "Max tokens in batch (max_num_batched_tokens) for decode cluster replica scheduler.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_scheduler_config_num_blocks: Optional[int] = field( - default=None, - metadata={ - "help": "Number of blocks for decode cluster replica scheduler.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_scheduler_config_block_size: Optional[int] = field( - default=None, - metadata={ - "help": "Block size for decode cluster replica scheduler.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_replica_scheduler_config_watermark_blocks_fraction: Optional[float] = field( - default=None, - metadata={ - "help": "Watermark blocks fraction for decode cluster replica scheduler.", - "mode_dependency": "pd-disaggregation", - }, - ) - - # DECODE_ATTN cluster scheduler configuration (for pd-af-disaggregation mode) - decode_attn_replica_scheduler_config_type: Optional[str] = field( - default=None, - metadata={ - "help": "Replica scheduler type for decode attention cluster. Overrides base replica_scheduler_config_type.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_scheduler_config_batch_size_cap: Optional[int] = field( - default=None, - metadata={ - "help": "Batch size cap (max_num_seqs) for decode attention cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_scheduler_config_max_tokens_in_batch: Optional[int] = field( - default=None, - metadata={ - "help": "Max tokens in batch (max_num_batched_tokens) for decode attention cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_scheduler_config_num_blocks: Optional[int] = field( - default=None, - metadata={ - "help": "Number of blocks for decode attention cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_scheduler_config_block_size: Optional[int] = field( - default=None, - metadata={ - "help": "Block size for decode attention cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_replica_scheduler_config_watermark_blocks_fraction: Optional[float] = ( - field( - default=None, - metadata={ - "help": "Watermark blocks fraction for decode attention cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - ) - - # DECODE_FFN cluster scheduler configuration (for pd-af-disaggregation mode) - decode_ffn_replica_scheduler_config_type: Optional[str] = field( - default=None, - metadata={ - "help": "Replica scheduler type for decode FFN cluster. Overrides base replica_scheduler_config_type.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_scheduler_config_batch_size_cap: Optional[int] = field( - default=None, - metadata={ - "help": "Batch size cap (max_num_seqs) for decode FFN cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_scheduler_config_max_tokens_in_batch: Optional[int] = field( - default=None, - metadata={ - "help": "Max tokens in batch (max_num_batched_tokens) for decode FFN cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_scheduler_config_num_blocks: Optional[int] = field( - default=None, - metadata={ - "help": "Number of blocks for decode FFN cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_scheduler_config_block_size: Optional[int] = field( - default=None, - metadata={ - "help": "Block size for decode FFN cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_replica_scheduler_config_watermark_blocks_fraction: Optional[float] = ( - field( - default=None, - metadata={ - "help": "Watermark blocks fraction for decode FFN cluster replica scheduler.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - ) - - # === Per-Cluster CC Backend Configuration === - # These fields allow per-cluster-type customization of CC backend parameters - # If not set, they fall back to the base cc_backend_config values - - # PREFILL cluster CC backend configuration - prefill_cc_backend_config_type: Optional[str] = field( - default=None, - metadata={ - "help": "CC backend type for prefill cluster. Options: 'vidur', 'analytical', 'collective_sim', 'astra_sim_analytical'. Overrides base cc_backend_config type.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_network_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Network bandwidth in Gbps for prefill cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_network_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Network latency in microseconds for prefill cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_intra_node_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-node bandwidth in Gbps for prefill cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_repo_root: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication backend repo root for prefill cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_system: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication backend system for prefill cluster CC backend (internal-only mode). Empty means infer from device.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_source_backend: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication source backend for prefill cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_source_version: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication source version for prefill cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_database_mode: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication database mode for prefill cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_tp_allreduce_impl: Optional[str] = field( - default=None, - metadata={ - "help": "TP allreduce implementation for prefill cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_custom_allreduce_variant: Optional[str] = field( - default=None, - metadata={ - "help": "Custom allreduce runtime label for prefill cluster CC backend when internal communication backend raw data has multiple variants.", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_prediction_cache_size: Optional[int] = field( - default=None, - metadata={ - "help": "Prediction cache size for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_placement_order: Optional[str] = field( - default=None, - metadata={ - "help": "Rank placement order for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_intra_server_topology: Optional[str] = field( - default=None, - metadata={ - "help": "Intra-server topology for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_inter_server_topology: Optional[str] = field( - default=None, - metadata={ - "help": "Inter-server topology for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_intra_server_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-server bandwidth in Gbps for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_intra_server_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-server latency in microseconds for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_inter_server_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Inter-server bandwidth in Gbps for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_inter_server_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Inter-server latency in microseconds for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_p2p_src_index: Optional[int] = field( - default=None, - metadata={ - "help": "P2P source participant index for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_p2p_dst_index: Optional[int] = field( - default=None, - metadata={ - "help": "P2P destination participant index for prefill cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_cc_backend_config_nvlink_allreduce_launch_overhead_us: Optional[float] = ( - field( - default=None, - metadata={ - "help": ( - "Per-step intra-server allreduce launch overhead in microseconds " - "for prefill cluster collective-sim backend." - ), - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - ) - prefill_execution_time_predictor_config_mlp_up_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override mlp_up_proj calibration scale for the prefill cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_execution_time_predictor_config_attn_pre_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override attn_pre_proj calibration scale for the prefill cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_execution_time_predictor_config_attn_post_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override attn_post_proj calibration scale for the prefill cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_execution_time_predictor_config_attn_decode_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override attn_decode calibration scale for the prefill cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_execution_time_predictor_config_attn_kv_cache_save_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override attn_kv_cache_save calibration scale for the prefill cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - prefill_execution_time_predictor_config_mlp_down_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override mlp_down_proj calibration scale for the prefill cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-af-disaggregation,pd-disaggregation", - }, - ) - - # DECODE cluster CC backend configuration (for unified decode in pd-disaggregation mode) - decode_cc_backend_config_type: Optional[str] = field( - default=None, - metadata={ - "help": "CC backend type for decode cluster. Options: 'vidur', 'analytical', 'collective_sim', 'astra_sim_analytical'. Overrides base cc_backend_config type.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_network_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Network bandwidth in Gbps for decode cluster CC backend (analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_network_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Network latency in microseconds for decode cluster CC backend (analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_intra_node_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-node bandwidth in Gbps for decode cluster CC backend (analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_repo_root: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication backend repo root for decode cluster CC backend (internal-only mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_system: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication backend system for decode cluster CC backend (internal-only mode). Empty means infer from device.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_source_backend: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication source backend for decode cluster CC backend (internal-only mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_source_version: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication source version for decode cluster CC backend (internal-only mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_database_mode: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication database mode for decode cluster CC backend (internal-only mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_tp_allreduce_impl: Optional[str] = field( - default=None, - metadata={ - "help": "TP allreduce implementation for decode cluster CC backend (internal-only mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_custom_allreduce_variant: Optional[str] = field( - default=None, - metadata={ - "help": "Custom allreduce runtime label for decode cluster CC backend when internal communication backend raw data has multiple variants.", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_prediction_cache_size: Optional[int] = field( - default=None, - metadata={ - "help": "Prediction cache size for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_placement_order: Optional[str] = field( - default=None, - metadata={ - "help": "Rank placement order for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_intra_server_topology: Optional[str] = field( - default=None, - metadata={ - "help": "Intra-server topology for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_inter_server_topology: Optional[str] = field( - default=None, - metadata={ - "help": "Inter-server topology for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_intra_server_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-server bandwidth in Gbps for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_intra_server_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-server latency in microseconds for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_inter_server_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Inter-server bandwidth in Gbps for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_inter_server_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Inter-server latency in microseconds for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_p2p_src_index: Optional[int] = field( - default=None, - metadata={ - "help": "P2P source participant index for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_p2p_dst_index: Optional[int] = field( - default=None, - metadata={ - "help": "P2P destination participant index for decode cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-disaggregation", - }, - ) - decode_cc_backend_config_nvlink_allreduce_launch_overhead_us: Optional[float] = ( - field( - default=None, - metadata={ - "help": ( - "Per-step intra-server allreduce launch overhead in microseconds " - "for decode cluster collective-sim backend." - ), - "mode_dependency": "pd-disaggregation", - }, - ) - ) - decode_execution_time_predictor_config_mlp_up_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override mlp_up_proj calibration scale for the decode cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-disaggregation", - }, - ) - decode_execution_time_predictor_config_attn_pre_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override attn_pre_proj calibration scale for the decode cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-disaggregation", - }, - ) - decode_execution_time_predictor_config_attn_post_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override attn_post_proj calibration scale for the decode cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-disaggregation", - }, - ) - decode_execution_time_predictor_config_attn_decode_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override attn_decode calibration scale for the decode cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-disaggregation", - }, - ) - decode_execution_time_predictor_config_attn_kv_cache_save_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override attn_kv_cache_save calibration scale for the decode cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-disaggregation", - }, - ) - decode_execution_time_predictor_config_mlp_down_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override mlp_down_proj calibration scale for the decode cluster " - "execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-disaggregation", - }, - ) - decode_execution_time_predictor_config_decode_phase_mlp_down_proj_calibration_scale: Optional[ - float - ] = field( - default=None, - metadata={ - "help": ( - "Override decode-phase-only mlp_down_proj calibration scale for the " - "decode cluster execution-time predictor. Must be > 0." - ), - "mode_dependency": "pd-disaggregation", - }, - ) - - # DECODE_ATTN cluster CC backend configuration (for pd-af-disaggregation mode) - decode_attn_cc_backend_config_type: Optional[str] = field( - default=None, - metadata={ - "help": "CC backend type for decode attention cluster. Options: 'vidur', 'analytical', 'collective_sim', 'astra_sim_analytical'. Overrides base cc_backend_config type.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_network_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Network bandwidth in Gbps for decode attention cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_network_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Network latency in microseconds for decode attention cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_intra_node_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-node bandwidth in Gbps for decode attention cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_repo_root: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication backend repo root for decode attention cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_system: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication backend system for decode attention cluster CC backend (internal-only mode). Empty means infer from device.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_source_backend: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication source backend for decode attention cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_source_version: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication source version for decode attention cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_database_mode: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication database mode for decode attention cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_tp_allreduce_impl: Optional[str] = field( - default=None, - metadata={ - "help": "TP allreduce implementation for decode attention cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_custom_allreduce_variant: Optional[str] = field( - default=None, - metadata={ - "help": "Custom allreduce runtime label for decode attention cluster CC backend when internal communication backend raw data has multiple variants.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_prediction_cache_size: Optional[int] = field( - default=None, - metadata={ - "help": "Prediction cache size for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_placement_order: Optional[str] = field( - default=None, - metadata={ - "help": "Rank placement order for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_intra_server_topology: Optional[str] = field( - default=None, - metadata={ - "help": "Intra-server topology for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_inter_server_topology: Optional[str] = field( - default=None, - metadata={ - "help": "Inter-server topology for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_intra_server_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-server bandwidth in Gbps for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_intra_server_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-server latency in microseconds for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_inter_server_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Inter-server bandwidth in Gbps for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_inter_server_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Inter-server latency in microseconds for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_p2p_src_index: Optional[int] = field( - default=None, - metadata={ - "help": "P2P source participant index for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_attn_cc_backend_config_p2p_dst_index: Optional[int] = field( - default=None, - metadata={ - "help": "P2P destination participant index for decode attention cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - - # DECODE_FFN cluster CC backend configuration (for pd-af-disaggregation mode) - decode_ffn_cc_backend_config_type: Optional[str] = field( - default=None, - metadata={ - "help": "CC backend type for decode FFN cluster. Options: 'vidur', 'analytical', 'collective_sim', 'astra_sim_analytical'. Overrides base cc_backend_config type.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_network_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Network bandwidth in Gbps for decode FFN cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_network_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Network latency in microseconds for decode FFN cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_intra_node_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-node bandwidth in Gbps for decode FFN cluster CC backend (analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_repo_root: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication backend repo root for decode FFN cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_system: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication backend system for decode FFN cluster CC backend (internal-only mode). Empty means infer from device.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_source_backend: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication source backend for decode FFN cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_source_version: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication source version for decode FFN cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_database_mode: Optional[str] = field( - default=None, - metadata={ - "help": "Internal-only communication database mode for decode FFN cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_tp_allreduce_impl: Optional[str] = field( - default=None, - metadata={ - "help": "TP allreduce implementation for decode FFN cluster CC backend (internal-only mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_custom_allreduce_variant: Optional[str] = field( - default=None, - metadata={ - "help": "Custom allreduce runtime label for decode FFN cluster CC backend when internal communication backend raw data has multiple variants.", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_prediction_cache_size: Optional[int] = field( - default=None, - metadata={ - "help": "Prediction cache size for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_placement_order: Optional[str] = field( - default=None, - metadata={ - "help": "Rank placement order for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_intra_server_topology: Optional[str] = field( - default=None, - metadata={ - "help": "Intra-server topology for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_inter_server_topology: Optional[str] = field( - default=None, - metadata={ - "help": "Inter-server topology for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_intra_server_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-server bandwidth in Gbps for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_intra_server_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Intra-server latency in microseconds for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_inter_server_bandwidth_gbps: Optional[float] = field( - default=None, - metadata={ - "help": "Inter-server bandwidth in Gbps for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_inter_server_latency_us: Optional[float] = field( - default=None, - metadata={ - "help": "Inter-server latency in microseconds for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_p2p_src_index: Optional[int] = field( - default=None, - metadata={ - "help": "P2P source participant index for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - decode_ffn_cc_backend_config_p2p_dst_index: Optional[int] = field( - default=None, - metadata={ - "help": "P2P destination participant index for decode FFN cluster CC backend (astra_sim_analytical mode).", - "mode_dependency": "pd-af-disaggregation", - }, - ) - - def __post_init__(self): - self._validate_open_source_release_cc_backend_guard() - - # check and set args only in first init (not in Cluster()) - if self.cluster_type is None: - # Early validation based on mode - self._validate_mode_consistency() - - # Basic validation for micro-batch size if provided - if self.decode_attn_micro_batch_size is not None: - assert ( - self.decode_attn_micro_batch_size >= 1 - ), f"decode_attn_micro_batch_size must be >=1, got {self.decode_attn_micro_batch_size}" - - # Ensure micro_batch_size equals batch_size_cap for DECODE_ATTN - # In DECODE_ATTN, micro-batch and batch are semantically equivalent - if ( - self.decode_attn_micro_batch_size is not None - and self.decode_attn_replica_scheduler_config_batch_size_cap is not None - ): - # Both specified: enforce equality - if ( - self.decode_attn_micro_batch_size - != self.decode_attn_replica_scheduler_config_batch_size_cap - ): - raise ValueError( - f"DECODE_ATTN micro_batch_size ({self.decode_attn_micro_batch_size}) " - f"must equal batch_size_cap ({self.decode_attn_replica_scheduler_config_batch_size_cap}). " - f"Reason: In DECODE_ATTN, micro-batch and batch are semantically equivalent." - ) - elif self.decode_attn_micro_batch_size is not None: - # Only micro_batch_size specified: propagate to batch_size_cap - self.decode_attn_replica_scheduler_config_batch_size_cap = ( - self.decode_attn_micro_batch_size - ) - elif self.decode_attn_replica_scheduler_config_batch_size_cap is not None: - # Only batch_size_cap specified: propagate to micro_batch_size - self.decode_attn_micro_batch_size = ( - self.decode_attn_replica_scheduler_config_batch_size_cap - ) - # If neither is specified, keep both as None (use defaults later) - - if self._has_disaggregation_params_set(): - self._setup_disaggregated_configs() - - # Add a check to ensure af_pipeline_num_micro_batch is consistent (only for PD+AF mode) - is_pd_af_mode = ( - self.decode_attn_cluster_num_replicas is not None - and self.decode_ffn_cluster_num_replicas is not None - ) - if is_pd_af_mode: - attn_mb = self.decode_attn_af_pipeline_num_micro_batch - ffn_mb = self.decode_ffn_af_pipeline_num_micro_batch - - assert ( - attn_mb is not None and ffn_mb is not None - ), "In PD+AF disaggregated mode, both decode_attn_af_pipeline_num_micro_batch and decode_ffn_af_pipeline_num_micro_batch must be set." - - assert ( - attn_mb == ffn_mb - ), "The af_pipeline_num_micro_batch must be the same for both decode_attn and decode_ffn clusters." - - # AFD Divisibility Validation (Fail Fast Strategy) - # Aligned with StepFun-vLLM's requirement that batch sizes be divisible by num_stages - # Unlike StepFun which silently rounds down, we fail fast to help users identify - # configuration issues early. - num_stages = attn_mb - if num_stages > 1: - self._validate_afd_divisibility(num_stages) - else: - self._setup_monolithic_config() - - self._validate_prefix_cache_spec_decode_compatibility() - - def _validate_prefix_cache_spec_decode_compatibility(self) -> None: - from frontier.spec_decode.runtime import ( - method_requires_prefix_matching_disabled, - ) - - prefix_enabled = bool( - getattr(self.replica_scheduler_config, "enable_prefix_caching", False) - ) - if not prefix_enabled: - return - - replica_configs = [ - ("replica_config", getattr(self, "replica_config", None)), - ("prefill_replica_config", getattr(self, "prefill_replica_config", None)), - ("decode_replica_config", getattr(self, "decode_replica_config", None)), - ( - "decode_attn_replica_config", - getattr(self, "decode_attn_replica_config", None), - ), - ( - "decode_ffn_replica_config", - getattr(self, "decode_ffn_replica_config", None), - ), - ] - for replica_config_name, replica_config in replica_configs: - if replica_config is None: - continue - spec_decode_config = getattr( - replica_config, "speculative_decoding_config", None - ) - if spec_decode_config is None or not spec_decode_config.enabled: - continue - method = str(getattr(spec_decode_config, "method", "")).strip() - if method and method_requires_prefix_matching_disabled(method): - raise ValueError( - "Speculative decoding method " - f"{method!r} requires prefix caching to be disabled, " - f"but replica_scheduler_config.enable_prefix_caching=True " - f"for {replica_config_name}." - ) - - def _validate_afd_divisibility(self, num_stages: int): - """Validate that key batch size parameters are divisible by num_stages. - - Currently relaxed — no divisibility enforcement. - """ - return None - - def _validate_mode_consistency(self): - """Validate that configuration is consistent with the intended mode.""" - - has_disaggregated_fields = ( - self.prefill_cluster_num_replicas is not None - or self.decode_attn_cluster_num_replicas is not None - or self.decode_ffn_cluster_num_replicas is not None - or self.decode_cluster_num_replicas is not None - ) - - # The `num_replicas` field is exclusively for monolithic mode. - # `replica_config` can be used as a template in disaggregated mode, so its presence is not a conflict. - has_monolithic_exclusive_field = self.num_replicas is not None - - if has_disaggregated_fields and has_monolithic_exclusive_field: - logger.warning( - "Both disaggregated and monolithic configuration fields are set. " - "The 'num_replicas' field (for monolithic mode) was provided but will be ignored in disaggregated mode. " - "Please use cluster-specific replica counts like 'prefill_cluster_num_replicas'." - ) - - def _validate_open_source_release_cc_backend_guard(self) -> None: - from frontier.cc_backend.cc_backend_config import AiconfiguratorCCBackendConfig - - if isinstance(self.cc_backend_config, AiconfiguratorCCBackendConfig): - raise ValueError(AICONFIGURATOR_BACKEND_RELEASE_ERROR) - - def _setup_monolithic_config(self): - """Setup configuration for monolithic (co-location) mode.""" - # Ensure required fields are set for monolithic mode - assert self.num_replicas != None, "Num replicas must be set" - assert self.replica_config != None, "Replica config must be set" - - # Set cluster type for monolithic mode - self.cluster_type = ClusterType.MONOLITHIC - - # Predictor routing details are indexed by serving Replica identity. - # Keep that capacity dimension explicit instead of deriving it from - # attention-DP lanes. - self.replica_config.cluster_num_replicas = int(self.num_replicas) - - # Reuse the same parallel-domain validation used by disaggregated clusters so - # monolithic MoE layouts fail fast when attention and MoE domains disagree. - self._validate_replica_config(self.replica_config, "monolithic") - self.world_size = self.replica_config.world_size * self.num_replicas - - # Clear disaggregated fields to avoid confusion - self.prefill_replica_config = None - self.decode_attn_replica_config = None - self.decode_ffn_replica_config = None - self.prefill_cluster_num_replicas = None - self.decode_attn_cluster_num_replicas = None - self.decode_ffn_cluster_num_replicas = None - - # else: - # assert self.replica_config.expert_parallel_size == self.replica_config.tensor_parallel_size, "For local MoE, expert_parallel_size must be equal to tensor_parallel_size" - - def _setup_disaggregated_configs(self): - """Setup configuration for disaggregated mode (PD or PD+AF).""" - # Clear monolithic fields since they're not used - # self.replica_config = None - self.num_replicas = None - - # Determine disaggregation mode - is_pd_af_mode = ( - self.decode_attn_cluster_num_replicas is not None - and self.decode_ffn_cluster_num_replicas is not None - ) - is_pd_mode = self.decode_cluster_num_replicas is not None - - # Ensure required disaggregated fields are set - assert ( - self.prefill_cluster_num_replicas != None - ), "Prefill cluster num replicas must be set" - - # Requirement 10.4: Validate that replica counts are positive - if ( - self.prefill_cluster_num_replicas is not None - and self.prefill_cluster_num_replicas <= 0 - ): - raise ValueError( - f"prefill_cluster_num_replicas must be positive, got {self.prefill_cluster_num_replicas}" - ) - - if is_pd_af_mode: - # PD+AF disaggregation mode - assert ( - self.decode_attn_cluster_num_replicas != None - ), "Decode attention cluster num replicas must be set" - assert ( - self.decode_ffn_cluster_num_replicas != None - ), "Decode FFN cluster num replicas must be set" - assert ( - not is_pd_mode - ), "Cannot set both PD and PD+AF disaggregation parameters" - - # Requirement 10.4: Validate that replica counts are positive (PD+AF mode) - if self.decode_attn_cluster_num_replicas <= 0: - raise ValueError( - f"decode_attn_cluster_num_replicas must be positive, got {self.decode_attn_cluster_num_replicas}" - ) - if self.decode_ffn_cluster_num_replicas <= 0: - raise ValueError( - f"decode_ffn_cluster_num_replicas must be positive, got {self.decode_ffn_cluster_num_replicas}" - ) - - # DECODE_FFN grouping semantics are implemented only in the - # RoundRobinClusterScheduler path. - cluster_scheduler_type = self.cluster_scheduler_config.get_type() - if cluster_scheduler_type != ClusterSchedulerType.ROUND_ROBIN: - raise ValueError( - "PD+AF mode requires RoundRobin cluster scheduler when DECODE_FFN is enabled. " - f"Got cluster_scheduler_config_type={cluster_scheduler_type}." - ) - - for field_name, decode_role in ( - ( - "decode_attn_replica_config_num_pipeline_stages", - "decode_attn", - ), - ( - "decode_ffn_replica_config_num_pipeline_stages", - "decode_ffn", - ), - ): - pipeline_stages = getattr(self, field_name) - if pipeline_stages not in (None, 1): - raise ValueError( - f"{field_name} must be 1 for {decode_role}, " - f"got {pipeline_stages}." - ) - - # Create ReplicaConfig objects from flattened fields - self.prefill_replica_config = self._create_replica_config_from_fields( - "prefill" - ) - self.decode_attn_replica_config = self._create_replica_config_from_fields( - "decode_attn" - ) - self.decode_ffn_replica_config = self._create_replica_config_from_fields( - "decode_ffn" - ) - - # Enforce disaggregation constraints - assert ( - self.decode_attn_replica_config.num_pipeline_stages == 1 - ), "Decode attention cluster must have 1 pipeline stage" - assert ( - self.decode_ffn_replica_config.num_pipeline_stages == 1 - ), "Decode FFN cluster must have 1 pipeline stage" - - # Validate each cluster - self._validate_replica_config(self.prefill_replica_config, "prefill") - self._validate_replica_config( - self.decode_attn_replica_config, "decode_attn" - ) - self._validate_replica_config(self.decode_ffn_replica_config, "decode_ffn") - - # Calculate world sizes - self.prefill_world_size = ( - self.prefill_cluster_num_replicas - * self.prefill_replica_config.world_size - ) - self.decode_attn_world_size = ( - self.decode_attn_cluster_num_replicas - * self.decode_attn_replica_config.world_size - ) - self.decode_ffn_world_size = ( - self.decode_ffn_cluster_num_replicas - * self.decode_ffn_replica_config.world_size - ) - self.world_size = ( - self.prefill_world_size - + self.decode_attn_world_size - + self.decode_ffn_world_size - ) - - elif is_pd_mode: - # PD disaggregation mode - assert ( - self.decode_cluster_num_replicas != None - ), "Decode cluster num replicas must be set" - assert ( - not is_pd_af_mode - ), "Cannot set both PD and PD+AF disaggregation parameters" - - # Requirement 10.4: Validate that replica counts are positive (PD mode) - if self.decode_cluster_num_replicas <= 0: - raise ValueError( - f"decode_cluster_num_replicas must be positive, got {self.decode_cluster_num_replicas}" - ) - - # Create ReplicaConfig objects from flattened fields - self.prefill_replica_config = self._create_replica_config_from_fields( - "prefill" - ) - self.decode_replica_config = self._create_replica_config_from_fields( - "decode" - ) - - # Validate each cluster - self._validate_replica_config(self.prefill_replica_config, "prefill") - self._validate_replica_config(self.decode_replica_config, "decode") - - # Calculate world sizes - self.prefill_world_size = ( - self.prefill_cluster_num_replicas - * self.prefill_replica_config.world_size - ) - self.decode_world_size = ( - self.decode_cluster_num_replicas * self.decode_replica_config.world_size - ) - self.world_size = self.prefill_world_size + self.decode_world_size - - else: - raise ValueError( - "Invalid disaggregation configuration: must set either PD or PD+AF parameters" - ) - - # Ensure consistent dummy mode configuration across all clusters - self._ensure_consistent_dummy_mode() - - print(f"Total world size: {self.world_size}") - - def _ensure_consistent_dummy_mode(self): - """Ensure all clusters use the same dummy mode configuration.""" - # Get the main execution_time_predictor_config dummy mode settings - main_config = self.execution_time_predictor_config - main_dummy_mode = main_config.enable_dummy_mode - main_dummy_time = main_config.dummy_execution_time_ms - - if main_dummy_mode: - print(f"Applying dummy mode (time={main_dummy_time}ms) to all clusters") - - # Apply dummy mode settings to all cluster configs - self.execution_time_predictor_config.enable_dummy_mode = True - self.execution_time_predictor_config.dummy_execution_time_ms = ( - main_dummy_time - ) - - # Note: In the current architecture, all clusters share the same execution_time_predictor_config - # This ensures consistency across all clusters - - def _field_is_set_to_non_default(self, field_def) -> bool: - value = getattr(self, field_def.name) - if field_def.default is not MISSING: - return value != field_def.default - if field_def.default_factory is not MISSING: - return value != field_def.default_factory() - return value is not None - - def _has_disaggregation_params_set(self) -> bool: - """Check if any disaggregation-specific cluster fields have been set.""" - for field_def in self.__dataclass_fields__.values(): - if ( - not field_def.name.startswith(DISAGGREGATED_CLUSTER_FIELD_PREFIXES) - and field_def.name not in DISAGGREGATED_CLUSTER_FIELD_NAMES - ): - continue - if self._field_is_set_to_non_default(field_def): - return True - return False - - def _create_replica_config_from_fields(self, cluster_prefix: str) -> ReplicaConfig: - """Create ReplicaConfig object from flattened fields.""" - # Get default values from main replica_config or use ReplicaConfig defaults - main_config = self.replica_config if self.replica_config else ReplicaConfig() - - # Extract cluster-specific fields using getattr with fallback to main config - def get_field_value(field_name: str): - cluster_field_name = f"{cluster_prefix}_replica_config_{field_name}" - cluster_value = getattr(self, cluster_field_name, None) - if cluster_value is not None: - return cluster_value - return getattr(main_config, field_name) - - if cluster_prefix == "decode_attn": - moe_expert_parallel_size = 0 - moe_tensor_parallel_size = 0 - total_expert_num = 0 - local_expert_num = 0 - num_pipeline_stages = 1 - router_load_balancing_type = None - router_topk = None - moe_routing_distribution_type = get_field_value( - "moe_routing_distribution_type" - ) - attn_tensor_parallel_size = get_field_value("attn_tensor_parallel_size") - attn_dp = get_field_value("attn_dp") - else: - attn_tensor_parallel_size = get_field_value("attn_tensor_parallel_size") - attn_dp = get_field_value("attn_dp") - moe_tensor_parallel_size = get_field_value("moe_tensor_parallel_size") - moe_expert_parallel_size = get_field_value("moe_expert_parallel_size") - total_expert_num = get_field_value("total_expert_num") - local_expert_num = get_field_value("local_expert_num") - router_load_balancing_type = get_field_value("router_load_balancing_type") - router_topk = get_field_value("router_topk") - moe_routing_distribution_type = get_field_value( - "moe_routing_distribution_type" - ) - if cluster_prefix == "decode_ffn": - num_pipeline_stages = 1 - else: - num_pipeline_stages = get_field_value("num_pipeline_stages") - - return ReplicaConfig( - model_name=get_field_value("model_name"), - memory_margin_fraction=get_field_value("memory_margin_fraction"), - num_pipeline_stages=num_pipeline_stages, - attn_tensor_parallel_size=attn_tensor_parallel_size, - attn_dp=attn_dp, - moe_tensor_parallel_size=moe_tensor_parallel_size, - moe_expert_parallel_size=moe_expert_parallel_size, - total_expert_num=total_expert_num, - local_expert_num=local_expert_num, - router_load_balancing_type=router_load_balancing_type, - router_topk=router_topk, - moe_routing_seed=get_field_value("moe_routing_seed"), - moe_routing_trace_path=get_field_value("moe_routing_trace_path"), - decode_attn_initial_lane_trace_path=get_field_value( - "decode_attn_initial_lane_trace_path" - ), - decode_attn_steady_state_snapshot_path=get_field_value( - "decode_attn_steady_state_snapshot_path" - ), - decode_attn_steady_state_measurement_report_path=get_field_value( - "decode_attn_steady_state_measurement_report_path" - ), - moe_routing_distribution_type=moe_routing_distribution_type, - device=get_field_value("device"), - network_device=get_field_value("network_device"), - cluster_prefix=cluster_prefix, - speculative_decoding_config=main_config.speculative_decoding_config, - ) - - def _validate_replica_config( - self, replica_config: ReplicaConfig, cluster_name: str - ): - """Validate replica configuration for specific cluster.""" - # Validate pipeline parallelism configuration (double-check, should already be validated in __post_init__) - if ( - replica_config.model_config.num_layers % replica_config.num_pipeline_stages - != 0 - ): - raise ValueError( - f"Pipeline parallelism configuration error in {cluster_name} cluster: " - f"num_layers ({replica_config.model_config.num_layers}) must be evenly divisible by " - f"num_pipeline_stages ({replica_config.num_pipeline_stages}). " - f"Current configuration would result in uneven layer distribution across pipeline stages." - ) - - # Validate dense model configuration in disaggregated modes. - is_dense_model = not replica_config.model_config.is_moe - if is_dense_model and cluster_name in ["prefill", "decode"]: - # For dense models in PD-disaggregation mode, enforce attn_dp = 1 - if replica_config.attn_dp != 1: - raise ValueError( - f"Dense models in PD-disaggregation mode require attn_dp=1 " - f"in {cluster_name} cluster, got {replica_config.attn_dp}. " - f"Dense models do not support attn data parallelism in disaggregated mode." - ) - # Ensure MoE parallelism is disabled for dense models - if replica_config.moe_expert_parallel_size != 1: - raise ValueError( - f"Dense models require moe_expert_parallel_size=1 in {cluster_name} cluster, " - f"got {replica_config.moe_expert_parallel_size}. " - f"Dense models do not have expert parallelism." - ) - - if is_dense_model and cluster_name == "decode_attn": - if replica_config.attn_dp != 1: - raise ValueError( - "Dense models require attn_dp=1 in " - f"{cluster_name} cluster, got " - f"{replica_config.attn_dp}." - ) - - if is_dense_model and cluster_name == "decode_ffn": - if replica_config.moe_expert_parallel_size != 1: - raise ValueError( - "Dense models require moe_expert_parallel_size=1 in " - f"{cluster_name} cluster, got " - f"{replica_config.moe_expert_parallel_size}." - ) - if replica_config.router_topk != 1: - raise ValueError( - f"Dense models require router_topk=1 in {cluster_name} " - f"cluster, got {replica_config.router_topk}." - ) - - normalized_cluster_name = str(cluster_name).strip().lower() - if normalized_cluster_name in {"prefill", "decode", "monolithic"} and replica_config.model_config.is_moe: - validate_frontier_shared_parallel_domains( - FrontierParallelismMapping( - cluster_num_replicas=1, - attn_tensor_parallel_size=replica_config.attn_tensor_parallel_size, - attn_dp=replica_config.attn_dp, - moe_tensor_parallel_size=replica_config.moe_tensor_parallel_size, - moe_expert_parallel_size=replica_config.moe_expert_parallel_size, - ) - ) - - if cluster_name != "decode_attn": - pass - # else: - # assert replica_config.moe_expert_parallel_size == replica_config.moe_tensor_parallel_size, f"For local MoE in {cluster_name} cluster, moe_expert_parallel_size must be equal to moe_tensor_parallel_size" - else: - assert ( - replica_config.moe_expert_parallel_size == 0 - and replica_config.local_expert_num == 0 - ), "For decode attention cluster, moe_expert_parallel_size and local_expert_num must be 0" - - def _collect_cluster_info(self) -> List[Tuple[str, int, ReplicaConfig]]: - clusters_info = [] - - if self._has_disaggregation_params_set(): - if self.prefill_cluster_num_replicas and self.prefill_replica_config: - clusters_info.append( - ( - "PREFILL", - self.prefill_cluster_num_replicas, - self.prefill_replica_config, - ) - ) - - if ( - self.decode_attn_cluster_num_replicas - and self.decode_attn_replica_config - ): - clusters_info.append( - ( - "DECODE_ATTN", - self.decode_attn_cluster_num_replicas, - self.decode_attn_replica_config, - ) - ) - - if self.decode_ffn_cluster_num_replicas and self.decode_ffn_replica_config: - clusters_info.append( - ( - "DECODE_FFN", - self.decode_ffn_cluster_num_replicas, - self.decode_ffn_replica_config, - ) - ) - - if self.decode_cluster_num_replicas and self.decode_replica_config: - clusters_info.append( - ( - "DECODE", - self.decode_cluster_num_replicas, - self.decode_replica_config, - ) - ) - else: - clusters_info.append(("MONOLITHIC", self.num_replicas, self.replica_config)) - - return clusters_info - - def get_server_count_metadata(self, sys_arch: str) -> Dict[str, int]: - clusters_info = self._collect_cluster_info() - server_counts_by_cluster = {} - _, _, _, CollectiveSimCCBackendConfig, _, _ = _get_cc_backend_configs() - cluster_prefix_by_name = { - "PREFILL": "prefill", - "DECODE_ATTN": "decode_attn", - "DECODE_FFN": "decode_ffn", - "DECODE": "decode", - } - - for cluster_name, num_replicas, replica_config in clusters_info: - cluster_total_devices = int(num_replicas) * int(replica_config.world_size) - num_devices_per_node = int(replica_config.node_config.num_devices_per_node) - if cluster_total_devices <= 0: - raise ValueError( - "cluster_total_devices must be positive when computing " - f"server-count metadata, got {cluster_total_devices}" - ) - if num_devices_per_node <= 0: - raise ValueError( - "num_devices_per_node must be positive when computing " - f"server-count metadata, got {num_devices_per_node}" - ) - if cluster_name == "MONOLITHIC": - cc_backend_config = self.cc_backend_config - else: - cc_backend_config = self._create_cc_backend_config_for_cluster( - cluster_prefix_by_name[cluster_name] - ) - if isinstance(cc_backend_config, CollectiveSimCCBackendConfig): - physical_topology = resolve_collective_sim_physical_topology( - cluster_total_devices=cluster_total_devices, - num_devices_per_node=num_devices_per_node, - scenario_profile=getattr( - cc_backend_config, - "scenario_profile", - None, - ), - ) - server_counts_by_cluster[cluster_name] = int(physical_topology.servers) - continue - server_counts_by_cluster[cluster_name] = ( - cluster_total_devices + num_devices_per_node - 1 - ) // num_devices_per_node - - if sys_arch == "co-location": - if "MONOLITHIC" not in server_counts_by_cluster: - raise ValueError("Missing MONOLITHIC cluster for co-location mode.") - return {"server_count": server_counts_by_cluster["MONOLITHIC"]} - if sys_arch == "pd-disaggregation": - if "PREFILL" not in server_counts_by_cluster or "DECODE" not in server_counts_by_cluster: - raise ValueError( - "Missing PREFILL or DECODE cluster for pd-disaggregation mode." - ) - return { - "prefill_server_count": server_counts_by_cluster["PREFILL"], - "decode_server_count": server_counts_by_cluster["DECODE"], - } - if sys_arch == "pd-af-disaggregation": - required = ["PREFILL", "DECODE_ATTN", "DECODE_FFN"] - if any(name not in server_counts_by_cluster for name in required): - raise ValueError( - "Missing PREFILL, DECODE_ATTN, or DECODE_FFN cluster for pd-af-disaggregation mode." - ) - return { - "prefill_server_count": server_counts_by_cluster["PREFILL"], - "decode_attn_server_count": server_counts_by_cluster["DECODE_ATTN"], - "decode_ffn_server_count": server_counts_by_cluster["DECODE_FFN"], - } - - raise ValueError(f"Unknown system architecture: {sys_arch}") - - def print_cluster_statistics(self, simulation_mode: str, sys_arch: str): - """Calculate and print statistics for all clusters (called from SimulationConfig).""" - clusters_info = self._collect_cluster_info() - - # Calculate total statistics - self.total_clusters = len(clusters_info) - self.cluster_world_sizes = {} - - # Calculate world_size if not already set - if not hasattr(self, "world_size") or self.world_size is None: - self.world_size = sum( - num_replicas * replica_config.world_size - for _, num_replicas, replica_config in clusters_info - ) - - # Print cluster configuration summary - print("\n" + "=" * 70) - print("CLUSTER CONFIGURATION SUMMARY") - print("=" * 70) - print("Simulation mode: ", simulation_mode) - print("System architecture: ", sys_arch) - print("=" * 70) - print(f"Total Clusters: {self.total_clusters}") - print(f"Total World Size: {self.world_size}") - server_count_metadata = self.get_server_count_metadata(sys_arch) - if sys_arch == "co-location": - print(f"Server count: {server_count_metadata['server_count']}") - elif sys_arch == "pd-disaggregation": - print( - f"Prefill server count: {server_count_metadata['prefill_server_count']}" - ) - print( - f"Decode server count: {server_count_metadata['decode_server_count']}" - ) - elif sys_arch == "pd-af-disaggregation": - print( - f"Prefill server count: {server_count_metadata['prefill_server_count']}" - ) - print( - f"Decode-Attn server count: {server_count_metadata['decode_attn_server_count']}" - ) - print( - f"Decode-FFN server count: {server_count_metadata['decode_ffn_server_count']}" - ) - print() - - for cluster_name, num_replicas, replica_config in clusters_info: - cluster_world_size = num_replicas * replica_config.world_size - self.cluster_world_sizes[cluster_name] = cluster_world_size - - print(f"Cluster Type: {cluster_name}") - print(f" Cluster World Size: {cluster_world_size}") - print(f" Num Replicas (Instances): {num_replicas}") - print(f" Replica World Size: {replica_config.world_size}") - if ( - cluster_name == "PREFILL" - or cluster_name == "MONOLITHIC" - or cluster_name == "DECODE" - ): - print( - f" Configuration: PP{replica_config.num_pipeline_stages} × (Attn_TP{replica_config.attn_tensor_parallel_size} x Attn_DP{replica_config.attn_dp}) | (MoE_TP{replica_config.moe_tensor_parallel_size} x MoE_EP{replica_config.moe_expert_parallel_size})" - ) - print(f" Total Expert Num: {replica_config.total_expert_num}") - print(f" Local Expert Num: {replica_config.local_expert_num}") - elif cluster_name == "DECODE_ATTN": - print( - f" Configuration: PP{replica_config.num_pipeline_stages} × Attn_TP{replica_config.attn_tensor_parallel_size} x Attn_DP{replica_config.attn_dp}" - ) - elif cluster_name == "DECODE_FFN": - print( - f" Configuration: PP{replica_config.num_pipeline_stages} × MoE_TP{replica_config.moe_tensor_parallel_size} x MoE_EP{replica_config.moe_expert_parallel_size}" - ) - print(f" Total Expert Num: {replica_config.total_expert_num}") - print(f" Local Expert Num: {replica_config.local_expert_num}") - - print("-" * 50) - - print("=" * 70 + "\n") - - def get_cluster_configs_for_disaggregation( - self, - ) -> Dict[ClusterType, "ClusterConfig"]: - """Generate cluster configurations for disaggregated mode.""" - if not self._has_disaggregation_params_set(): - return {ClusterType.MONOLITHIC: self} - - cluster_configs = {} - - # Prefill cluster - if self.prefill_cluster_num_replicas: - prefill_config = ClusterConfig( - cluster_type=ClusterType.PREFILL, - num_replicas=self.prefill_cluster_num_replicas, - replica_config=self.prefill_replica_config - or self._create_replica_config_copy(), - cluster_scheduler_config=self.cluster_scheduler_config, - replica_scheduler_config=self.replica_scheduler_config, - execution_time_predictor_config=( - self._create_execution_time_predictor_config_for_cluster( - "prefill" - ) - ), - cc_backend_config=self._create_cc_backend_config_for_cluster("prefill"), - # Propagate cluster-specific replica scheduler config parameters - prefill_replica_scheduler_config_type=self.prefill_replica_scheduler_config_type, - prefill_replica_scheduler_config_batch_size_cap=self.prefill_replica_scheduler_config_batch_size_cap, - prefill_replica_scheduler_config_max_tokens_in_batch=self.prefill_replica_scheduler_config_max_tokens_in_batch, - prefill_replica_scheduler_config_enable_chunked_prefill=self.prefill_replica_scheduler_config_enable_chunked_prefill, - prefill_replica_scheduler_config_long_prefill_token_threshold=self.prefill_replica_scheduler_config_long_prefill_token_threshold, - prefill_replica_scheduler_config_num_blocks=self.prefill_replica_scheduler_config_num_blocks, - prefill_replica_scheduler_config_block_size=self.prefill_replica_scheduler_config_block_size, - prefill_replica_scheduler_config_watermark_blocks_fraction=self.prefill_replica_scheduler_config_watermark_blocks_fraction, - ) - cluster_configs[ClusterType.PREFILL] = prefill_config - - # Decode Attention cluster - if self.decode_attn_cluster_num_replicas: - # Determine micro-batch SIZE for decode-attn - # NOTE: Only decode_attn_micro_batch_size exists (no generic fallback) - _da_mbs = self.decode_attn_micro_batch_size - decode_attn_config = ClusterConfig( - cluster_type=ClusterType.DECODE_ATTN, - num_replicas=self.decode_attn_cluster_num_replicas, - replica_config=self.decode_attn_replica_config - or self._create_replica_config_copy(), - cluster_scheduler_config=self.cluster_scheduler_config, - replica_scheduler_config=self.replica_scheduler_config, - execution_time_predictor_config=( - self._create_execution_time_predictor_config_for_cluster( - "decode_attn" - ) - ), - cc_backend_config=self._create_cc_backend_config_for_cluster( - "decode_attn" - ), - af_pipeline_num_micro_batch=self.decode_attn_af_pipeline_num_micro_batch, - decode_attn_micro_batch_size=_da_mbs, - decode_attn_request_allocation_threshold=self.decode_attn_request_allocation_threshold, - # Propagate cluster-specific replica scheduler config parameters - decode_attn_replica_scheduler_config_type=self.decode_attn_replica_scheduler_config_type, - decode_attn_replica_scheduler_config_batch_size_cap=self.decode_attn_replica_scheduler_config_batch_size_cap, - decode_attn_replica_scheduler_config_max_tokens_in_batch=self.decode_attn_replica_scheduler_config_max_tokens_in_batch, - decode_attn_replica_scheduler_config_num_blocks=self.decode_attn_replica_scheduler_config_num_blocks, - decode_attn_replica_scheduler_config_block_size=self.decode_attn_replica_scheduler_config_block_size, - decode_attn_replica_scheduler_config_watermark_blocks_fraction=self.decode_attn_replica_scheduler_config_watermark_blocks_fraction, - ) - cluster_configs[ClusterType.DECODE_ATTN] = decode_attn_config - - # Decode FFN cluster - if self.decode_ffn_cluster_num_replicas: - decode_ffn_config = ClusterConfig( - cluster_type=ClusterType.DECODE_FFN, - num_replicas=self.decode_ffn_cluster_num_replicas, - replica_config=self.decode_ffn_replica_config - or self._create_replica_config_copy(), - cluster_scheduler_config=self.cluster_scheduler_config, - replica_scheduler_config=self.replica_scheduler_config, - execution_time_predictor_config=( - self._create_execution_time_predictor_config_for_cluster( - "decode_ffn" - ) - ), - cc_backend_config=self._create_cc_backend_config_for_cluster( - "decode_ffn" - ), - af_pipeline_num_micro_batch=self.decode_ffn_af_pipeline_num_micro_batch, - # Propagate cluster-specific replica scheduler config parameters - decode_ffn_replica_scheduler_config_type=self.decode_ffn_replica_scheduler_config_type, - decode_ffn_replica_scheduler_config_batch_size_cap=self.decode_ffn_replica_scheduler_config_batch_size_cap, - decode_ffn_replica_scheduler_config_max_tokens_in_batch=self.decode_ffn_replica_scheduler_config_max_tokens_in_batch, - decode_ffn_replica_scheduler_config_num_blocks=self.decode_ffn_replica_scheduler_config_num_blocks, - decode_ffn_replica_scheduler_config_block_size=self.decode_ffn_replica_scheduler_config_block_size, - decode_ffn_replica_scheduler_config_watermark_blocks_fraction=self.decode_ffn_replica_scheduler_config_watermark_blocks_fraction, - decode_attn_cluster_num_replicas=self.decode_attn_cluster_num_replicas, - ) - # Propagate only the source Attention-Replica capacity. AFD - # grouping is Replica-level; it must not manufacture DP lanes. - decode_ffn_config.decode_attn_replica_id_start_for_ffn = int( - self.prefill_cluster_num_replicas - ) - - cluster_configs[ClusterType.DECODE_FFN] = decode_ffn_config - - # Unified Decode cluster (PD-disaggregation mode) - if self.decode_cluster_num_replicas: - decode_config = ClusterConfig( - cluster_type=ClusterType.DECODE, - num_replicas=self.decode_cluster_num_replicas, - replica_config=self.decode_replica_config - or self._create_replica_config_copy(), - cluster_scheduler_config=self.cluster_scheduler_config, - replica_scheduler_config=self.replica_scheduler_config, - execution_time_predictor_config=( - self._create_execution_time_predictor_config_for_cluster("decode") - ), - cc_backend_config=self._create_cc_backend_config_for_cluster("decode"), - # Propagate cluster-specific replica scheduler config parameters - decode_replica_scheduler_config_type=self.decode_replica_scheduler_config_type, - decode_replica_scheduler_config_batch_size_cap=self.decode_replica_scheduler_config_batch_size_cap, - decode_replica_scheduler_config_max_tokens_in_batch=self.decode_replica_scheduler_config_max_tokens_in_batch, - decode_replica_scheduler_config_num_blocks=self.decode_replica_scheduler_config_num_blocks, - decode_replica_scheduler_config_block_size=self.decode_replica_scheduler_config_block_size, - decode_replica_scheduler_config_watermark_blocks_fraction=self.decode_replica_scheduler_config_watermark_blocks_fraction, - ) - cluster_configs[ClusterType.DECODE] = decode_config - - return cluster_configs - - def _create_execution_time_predictor_config_for_cluster( - self, cluster_prefix: str - ) -> BaseExecutionTimePredictorConfig: - """Create cluster-specific execution-time predictor config overrides.""" - base_config = self.execution_time_predictor_config - override_values = {} - for calibration_field in ( - "attn_pre_proj_calibration_scale", - "attn_post_proj_calibration_scale", - "attn_decode_calibration_scale", - "attn_kv_cache_save_calibration_scale", - "mlp_up_proj_calibration_scale", - "mlp_down_proj_calibration_scale", - "decode_phase_mlp_down_proj_calibration_scale", - ): - override_field = ( - f"{cluster_prefix}_execution_time_predictor_config_" - f"{calibration_field}" - ) - override_value = getattr(self, override_field, None) - if override_value is None: - continue - override_value = float(override_value) - if override_value <= 0.0: - raise ValueError( - f"ClusterConfig.{override_field} must be > 0, got={override_value!r}" - ) - override_values[calibration_field] = override_value - - if not override_values: - return base_config - - return replace(base_config, **override_values) - - def _create_cc_backend_config_for_cluster( - self, cluster_prefix: str - ) -> BaseCCBackendConfig: - """ - Create CC backend configuration for a specific cluster. - - This method creates a cluster-specific CC backend configuration by: - 1. Checking for cluster-specific override values - 2. Falling back to the base cc_backend_config values if not overridden - - Args: - cluster_prefix: Cluster prefix (e.g., "prefill", "decode", "decode_attn", "decode_ffn") - - Returns: - CC backend configuration for the specified cluster - """ - # Lazy import to avoid circular imports - ( - _, - VidurCCBackendConfig, - AnalyticalCCBackendConfig, - CollectiveSimCCBackendConfig, - AiconfiguratorCCBackendConfig, - AstraSimAnalyticalCCBackendConfig, - ) = _get_cc_backend_configs() - - # Get cluster-specific type override - type_field = f"{cluster_prefix}_cc_backend_config_type" - cluster_type_str = getattr(self, type_field, None) - - # Determine which config type to use - if cluster_type_str is not None: - # Use cluster-specific type - cluster_type_key = cluster_type_str.lower() - if cluster_type_key == "analytical": - return self._create_analytical_cc_backend_config(cluster_prefix) - elif cluster_type_key == "vidur": - return self._create_vidur_cc_backend_config(cluster_prefix) - elif cluster_type_key == "collective_sim": - return self._create_collective_sim_cc_backend_config(cluster_prefix) - elif cluster_type_key == "aiconfigurator": - return self._create_aiconfigurator_cc_backend_config(cluster_prefix) - elif cluster_type_key == "astra_sim_analytical": - return self._create_astra_sim_analytical_cc_backend_config( - cluster_prefix - ) - else: - raise ValueError(f"Unknown CC backend type: {cluster_type_str}") - else: - # Use base config type - base_config = self.cc_backend_config - if isinstance(base_config, AnalyticalCCBackendConfig): - return self._create_analytical_cc_backend_config(cluster_prefix) - elif isinstance(base_config, VidurCCBackendConfig): - return self._create_vidur_cc_backend_config(cluster_prefix) - elif isinstance(base_config, CollectiveSimCCBackendConfig): - return self._create_collective_sim_cc_backend_config(cluster_prefix) - elif isinstance(base_config, AiconfiguratorCCBackendConfig): - return self._create_aiconfigurator_cc_backend_config(cluster_prefix) - elif isinstance(base_config, AstraSimAnalyticalCCBackendConfig): - return self._create_astra_sim_analytical_cc_backend_config( - cluster_prefix - ) - else: - raise ValueError( - "Unsupported base CC backend config type for cluster-specific " - f"construction: {type(base_config).__name__}" - ) - - def _create_analytical_cc_backend_config( - self, cluster_prefix: str - ) -> AnalyticalCCBackendConfig: - """Create analytical CC backend config with cluster-specific overrides.""" - # Lazy import to avoid circular imports - _, _, AnalyticalCCBackendConfig, _, _, _ = _get_cc_backend_configs() - - base_config = self.cc_backend_config - - # Get cluster-specific values with fallback to base config - def get_value(field_name: str, default_value): - cluster_field = f"{cluster_prefix}_cc_backend_config_{field_name}" - cluster_value = getattr(self, cluster_field, None) - if cluster_value is not None: - return cluster_value - if isinstance(base_config, AnalyticalCCBackendConfig): - return getattr(base_config, field_name, default_value) - return default_value - - return AnalyticalCCBackendConfig( - profiling_data_dir=( - base_config.profiling_data_dir - if hasattr(base_config, "profiling_data_dir") - else "data/profiling/network" - ), - cache_dir=( - base_config.cache_dir if hasattr(base_config, "cache_dir") else "cache" - ), - no_cache=( - base_config.no_cache if hasattr(base_config, "no_cache") else False - ), - network_bandwidth_gbps=get_value("network_bandwidth_gbps", 100.0), - network_latency_us=get_value("network_latency_us", 1.0), - intra_node_bandwidth_gbps=get_value("intra_node_bandwidth_gbps", 600.0), - ) - - def _create_vidur_cc_backend_config( - self, cluster_prefix: str - ) -> VidurCCBackendConfig: - """Create Vidur CC backend config with cluster-specific overrides.""" - # Lazy import to avoid circular imports - _, VidurCCBackendConfig, _, _, _, _ = _get_cc_backend_configs() - - base_config = self.cc_backend_config - - # For Vidur config, we mainly use the base config values - # as Vidur-specific parameters are typically shared across clusters - if isinstance(base_config, VidurCCBackendConfig): - return VidurCCBackendConfig( - profiling_data_dir=base_config.profiling_data_dir, - cache_dir=base_config.cache_dir, - no_cache=base_config.no_cache, - all_reduce_input_file=base_config.all_reduce_input_file, - send_recv_input_file=base_config.send_recv_input_file, - k_fold_cv_splits=base_config.k_fold_cv_splits, - num_training_job_threads=base_config.num_training_job_threads, - ) - else: - # Create default Vidur config - return VidurCCBackendConfig() - - def _create_collective_sim_cc_backend_config( - self, cluster_prefix: str - ) -> "CollectiveSimCCBackendConfig": - """Create collective-sim CC backend config with cluster-specific overrides.""" - ( - _, - _, - _, - CollectiveSimCCBackendConfig, - _, - _, - ) = _get_cc_backend_configs() - - from dataclasses import replace - from pathlib import Path - - base_config = self.cc_backend_config - if not isinstance(base_config, CollectiveSimCCBackendConfig): - base_config = CollectiveSimCCBackendConfig() - - def get_value(field_name: str, default_value): - cluster_field = f"{cluster_prefix}_cc_backend_config_{field_name}" - cluster_value = getattr(self, cluster_field, None) - if cluster_value is not None: - return cluster_value - return getattr(base_config, field_name, default_value) - - if base_config.runner_out_dir: - cluster_out_dir = str(Path(base_config.runner_out_dir) / cluster_prefix) - return replace( - base_config, - runner_out_dir=cluster_out_dir, - nvlink_allreduce_launch_overhead_us=get_value( - "nvlink_allreduce_launch_overhead_us", - 50.0, - ), - ) - - return replace( - base_config, - nvlink_allreduce_launch_overhead_us=get_value( - "nvlink_allreduce_launch_overhead_us", - 50.0, - ), - ) - - def _create_aiconfigurator_cc_backend_config( - self, cluster_prefix: str - ) -> "AiconfiguratorCCBackendConfig": - """Create aiconfigurator CC backend config with cluster-specific overrides.""" - ( - _, - _, - _, - _, - AiconfiguratorCCBackendConfig, - _, - ) = _get_cc_backend_configs() - - base_config = self.cc_backend_config - - def get_value(field_name: str, default_value): - cluster_field = f"{cluster_prefix}_cc_backend_config_{field_name}" - cluster_value = getattr(self, cluster_field, None) - if cluster_value is not None: - return cluster_value - if isinstance(base_config, AiconfiguratorCCBackendConfig): - return getattr(base_config, field_name, default_value) - return default_value - - return AiconfiguratorCCBackendConfig( - profiling_data_dir=( - base_config.profiling_data_dir - if hasattr(base_config, "profiling_data_dir") - else "data/profiling/network" - ), - cache_dir=( - base_config.cache_dir if hasattr(base_config, "cache_dir") else "cache" - ), - no_cache=( - base_config.no_cache if hasattr(base_config, "no_cache") else False - ), - repo_root=get_value("repo_root", "sota-infer-engine/aiconfigurator"), - system=get_value("system", ""), - source_backend=get_value("source_backend", "vllm"), - source_version=get_value("source_version", ""), - database_mode=get_value("database_mode", "silicon"), - tp_allreduce_impl=get_value("tp_allreduce_impl", "custom_allreduce"), - custom_allreduce_variant=get_value("custom_allreduce_variant", None), - ) - - def _create_astra_sim_analytical_cc_backend_config( - self, cluster_prefix: str - ) -> "AstraSimAnalyticalCCBackendConfig": - """Create astra-sim analytical CC backend config with cluster-specific overrides.""" - ( - _, - _, - _, - _, - _, - AstraSimAnalyticalCCBackendConfig, - ) = _get_cc_backend_configs() - - base_config = self.cc_backend_config - - def get_value(field_name: str, default_value): - cluster_field = f"{cluster_prefix}_cc_backend_config_{field_name}" - cluster_value = getattr(self, cluster_field, None) - if cluster_value is not None: - return cluster_value - if isinstance(base_config, AstraSimAnalyticalCCBackendConfig): - return getattr(base_config, field_name, default_value) - return default_value - - return AstraSimAnalyticalCCBackendConfig( - profiling_data_dir=( - base_config.profiling_data_dir - if hasattr(base_config, "profiling_data_dir") - else "data/profiling/network" - ), - cache_dir=( - base_config.cache_dir if hasattr(base_config, "cache_dir") else "cache" - ), - no_cache=( - base_config.no_cache if hasattr(base_config, "no_cache") else False - ), - prediction_cache_size=get_value("prediction_cache_size", 4096), - placement_order=get_value("placement_order", "TP,CP,DP,EP"), - intra_server_topology=get_value( - "intra_server_topology", "FullyConnected" - ), - inter_server_topology=get_value( - "inter_server_topology", "FullyConnected" - ), - intra_server_bandwidth_gbps=get_value( - "intra_server_bandwidth_gbps", 600.0 - ), - intra_server_latency_us=get_value("intra_server_latency_us", 1.0), - inter_server_bandwidth_gbps=get_value( - "inter_server_bandwidth_gbps", 100.0 - ), - inter_server_latency_us=get_value("inter_server_latency_us", 1.0), - ring_bidirectional=( - base_config.ring_bidirectional - if isinstance(base_config, AstraSimAnalyticalCCBackendConfig) - else True - ), - p2p_src_index=get_value("p2p_src_index", 0), - p2p_dst_index=get_value("p2p_dst_index", 1), - ) - - def _create_replica_config_copy(self) -> ReplicaConfig: - """Create a copy of the main replica config for disaggregated clusters.""" - # Note: This method now needs to be called before replica_config is cleared - # We need to preserve the original config temporarily - original_config = ( - self.replica_config if self.replica_config else ReplicaConfig() - ) - - return ReplicaConfig( - model_name=original_config.model_name, - memory_margin_fraction=original_config.memory_margin_fraction, - num_pipeline_stages=original_config.num_pipeline_stages, - attn_tensor_parallel_size=original_config.attn_tensor_parallel_size, - attn_dp=original_config.attn_dp, - moe_tensor_parallel_size=original_config.moe_tensor_parallel_size, - moe_expert_parallel_size=original_config.moe_expert_parallel_size, - total_expert_num=original_config.total_expert_num, - router_load_balancing_type=original_config.router_load_balancing_type, - router_topk=original_config.router_topk, - moe_routing_seed=original_config.moe_routing_seed, - moe_routing_distribution_type=original_config.moe_routing_distribution_type, - moe_routing_trace_path=original_config.moe_routing_trace_path, - decode_attn_initial_lane_trace_path=( - original_config.decode_attn_initial_lane_trace_path - ), - decode_attn_steady_state_snapshot_path=( - original_config.decode_attn_steady_state_snapshot_path - ), - decode_attn_steady_state_measurement_report_path=( - original_config.decode_attn_steady_state_measurement_report_path - ), - device=original_config.device, - network_device=original_config.network_device, - speculative_decoding_config=original_config.speculative_decoding_config, - ) +logger = init_logger(__name__) @dataclass diff --git a/frontier/config/execution_time_predictor_config.py b/frontier/config/execution_time_predictor_config.py new file mode 100644 index 00000000..7ca20cf6 --- /dev/null +++ b/frontier/config/execution_time_predictor_config.py @@ -0,0 +1,314 @@ +"""Execution-time predictor configuration and its calibration scales.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List, Optional + +from frontier.config.base_poly_config import BasePolyConfig +from frontier.types import ExecutionTimePredictorType + + +@dataclass +class BaseExecutionTimePredictorConfig(BasePolyConfig): + linear_op_input_file: str = field( + default="./data/profiling/compute/{DEVICE}/{MODEL}/linear_op.csv", + metadata={"help": "Path to the linear operation profiling input file."}, + ) + # Backward compatibility alias + mlp_input_file: str = field( + default="", + metadata={"help": "[DEPRECATED] Use linear_op_input_file instead."}, + ) + atten_input_file: str = field( + default="./data/profiling/compute/{DEVICE}/{MODEL}/attention.csv", + metadata={"help": "Path to the attention input file."}, + ) + gdn_input_file: str = field( + default="./data/profiling/compute/{DEVICE}/{MODEL}/gdn.csv", + metadata={"help": "Path to the standard GDN profiling input file."}, + ) + all_reduce_input_file: str = field( + default="./data/profiling/network/{NETWORK_DEVICE}/all_reduce.csv", + metadata={"help": "Path to the all reduce input file."}, + ) + send_recv_input_file: str = field( + default="./data/profiling/network/{NETWORK_DEVICE}/send_recv.csv", + metadata={"help": "Path to the send recv input file."}, + ) + cpu_overhead_input_file: str = field( + default="./data/profiling/cpu_overhead/{NETWORK_DEVICE}/{MODEL}/cpu_overheads.csv", + metadata={"help": "Path to the cpu overhead input file."}, + ) + cpu_overhead_kernel_only_input_file: str = field( + default="./data/profiling/cpu_overhead/{NETWORK_DEVICE}/{MODEL}/cpu_overheads_kernel_only.csv", + metadata={"help": "Path to the kernel-only cpu overhead input file."}, + ) + pp_stage_boundary_input_file: str = field( + default="./data/profiling/other_overhead/{DEVICE}/{MODEL}/pp_stage_boundary.csv", + metadata={"help": "Path to the pipeline stage-boundary overhead input file."}, + ) + pp_receiver_head_input_file: str = field( + default="./data/profiling/other_overhead/{DEVICE}/{MODEL}/pp_receiver_head.csv", + metadata={"help": "Path to the PP receiver-head overhead input file."}, + ) + pp_producer_send_path_input_file: str = field( + default="./data/profiling/other_overhead/{DEVICE}/{MODEL}/pp_producer_send_path.csv", + metadata={"help": "Path to the PP producer send-path overhead input file."}, + ) + pp_prefill_consumer_active_input_file: str = field( + default="./data/profiling/other_overhead/{DEVICE}/{MODEL}/pp_prefill_consumer_active.csv", + metadata={ + "help": "Path to the PP prefill consumer-active overhead input file." + }, + ) + moe_input_file: str = field( + default="./data/profiling/compute/{DEVICE}/{MODEL}/moe.csv", + metadata={"help": "Path to the MoE profiling input file."}, + ) + linear_op_kernel_only_input_file: str = field( + default="./data/profiling/compute/{DEVICE}/{MODEL}/linear_op_kernel_only.csv", + metadata={"help": "Path to the kernel-only linear operation profiling input file."}, + ) + atten_kernel_only_input_file: str = field( + default="./data/profiling/compute/{DEVICE}/{MODEL}/attention_kernel_only.csv", + metadata={"help": "Path to the kernel-only attention input file."}, + ) + moe_kernel_only_input_file: str = field( + default="./data/profiling/compute/{DEVICE}/{MODEL}/moe_kernel_only.csv", + metadata={"help": "Path to the kernel-only MoE profiling input file."}, + ) + k_fold_cv_splits: int = field( + default=10, + metadata={"help": "Number of k fold cross validation splits."}, + ) + no_cache: bool = field( + default=False, + metadata={"help": "Whether to cache prediction models."}, + ) + kv_cache_prediction_granularity: int = field( + default=64, + metadata={"help": "KV cache prediction granularity."}, + ) + prediction_max_prefill_chunk_size: int = field( + default=4096, + metadata={"help": "Max prefill chunk size for prediction."}, + ) + prediction_max_batch_size: int = field( + default=128, + metadata={"help": "Max batch size for prediction."}, + ) + prediction_max_tokens_per_request: int = field( + default=4096, + metadata={"help": "Max tokens per request for prediction."}, + ) + attention_decode_batching_overhead_fraction: float = field( + default=0.1, + metadata={"help": "Attention decode batching overhead fraction."}, + ) + attention_prefill_batching_overhead_fraction: float = field( + default=0.1, + metadata={"help": "Attention prefill batching overhead fraction."}, + ) + attn_pre_proj_calibration_scale: float = field( + default=1.0, + metadata={ + "help": "Multiplicative calibration scale for attn_pre_proj prediction. Must be > 0." + }, + ) + prefill_phase_attn_pre_proj_calibration_scale: Optional[float] = field( + default=None, + metadata={ + "help": ( + "Optional multiplicative calibration scale for attn_pre_proj " + "prediction when the batch includes prefill tokens. Must be > 0." + ) + }, + ) + attn_post_proj_calibration_scale: float = field( + default=1.0, + metadata={ + "help": "Multiplicative calibration scale for attn_post_proj prediction. Must be > 0." + }, + ) + prefill_phase_attn_post_proj_calibration_scale: Optional[float] = field( + default=None, + metadata={ + "help": ( + "Optional multiplicative calibration scale for attn_post_proj " + "prediction when the batch includes prefill tokens. Must be > 0." + ) + }, + ) + attn_decode_calibration_scale: float = field( + default=1.0, + metadata={ + "help": "Multiplicative calibration scale for attn_decode prediction. Must be > 0." + }, + ) + attn_decode_in_mixed_calibration_scale: Optional[float] = field( + default=None, + metadata={ + "help": ( + "Optional multiplicative calibration scale for attn_decode_in_mixed " + "prediction when a co-location batch contains both prefill and decode " + "tokens. Must be > 0." + ) + }, + ) + late_decode_attn_decode_calibration_scale: Optional[float] = field( + default=None, + metadata={ + "help": ( + "Optional multiplicative calibration scale for attn_decode " + "prediction when every decode request in the batch has already " + "completed the first pure decode token. Must be > 0." + ) + }, + ) + attn_kv_cache_save_calibration_scale: float = field( + default=1.0, + metadata={ + "help": "Multiplicative calibration scale for attn_kv_cache_save prediction. Must be > 0." + }, + ) + prefill_phase_attn_kv_cache_save_calibration_scale: Optional[float] = field( + default=None, + metadata={ + "help": ( + "Optional multiplicative calibration scale for attn_kv_cache_save " + "prediction when the batch includes prefill tokens. Must be > 0." + ) + }, + ) + mlp_up_proj_calibration_scale: float = field( + default=1.0, + metadata={ + "help": "Multiplicative calibration scale for mlp_up_proj prediction. Must be > 0." + }, + ) + prefill_phase_mlp_up_proj_calibration_scale: Optional[float] = field( + default=None, + metadata={ + "help": ( + "Optional multiplicative calibration scale for mlp_up_proj " + "prediction when the batch includes prefill tokens. Must be > 0." + ) + }, + ) + mlp_down_proj_calibration_scale: float = field( + default=1.0, + metadata={ + "help": "Multiplicative calibration scale for mlp_down_proj prediction. Must be > 0." + }, + ) + decode_phase_mlp_down_proj_calibration_scale: Optional[float] = field( + default=None, + metadata={ + "help": ( + "Optional multiplicative calibration scale for mlp_down_proj " + "prediction when the batch contains decode tokens but no " + "prefill tokens. Must be > 0." + ) + }, + ) + nccl_cpu_launch_overhead_ms: float = field( + default=0.02, + metadata={"help": "NCCL CPU launch overhead in ms."}, + ) + nccl_cpu_skew_overhead_per_device_ms: float = field( + default=0.0, + metadata={"help": "NCCL CPU skew overhead per device in ms."}, + ) + num_training_job_threads: int = field( + default=-1, + metadata={"help": "Number of training job threads."}, + ) + skip_cpu_overhead_modeling: bool = field( + default=True, + metadata={"help": "Whether to skip CPU overhead modeling."}, + ) + + # Dummy mode configuration for fast testing and development + enable_dummy_mode: bool = field( + default=False, + metadata={ + "help": "Enable dummy mode to skip ML model training and return fixed execution times." + }, + ) + dummy_execution_time_ms: float = field( + default=1.0, + metadata={ + "help": "Fixed execution time in milliseconds to return in dummy mode." + }, + ) + + def __post_init__(self) -> None: + for field_name in ( + "attn_pre_proj_calibration_scale", + "prefill_phase_attn_pre_proj_calibration_scale", + "attn_post_proj_calibration_scale", + "prefill_phase_attn_post_proj_calibration_scale", + "attn_decode_calibration_scale", + "attn_decode_in_mixed_calibration_scale", + "late_decode_attn_decode_calibration_scale", + "attn_kv_cache_save_calibration_scale", + "prefill_phase_attn_kv_cache_save_calibration_scale", + "mlp_up_proj_calibration_scale", + "prefill_phase_mlp_up_proj_calibration_scale", + "mlp_down_proj_calibration_scale", + "decode_phase_mlp_down_proj_calibration_scale", + ): + raw_value = getattr(self, field_name) + if raw_value is None: + continue + value = float(raw_value) + if value <= 0.0: + raise ValueError( + f"{self.__class__.__name__}.{field_name} must be > 0, got={value!r}" + ) + + + +@dataclass +class LinearRegressionExecutionTimePredictorConfig(BaseExecutionTimePredictorConfig): + polynomial_degree: List[int] = field( + default_factory=lambda: list(range(1, 6)), + metadata={"help": "Polynomial degree for linear regression."}, + ) + polynomial_include_bias: List[bool] = field( + default_factory=lambda: [True, False], + metadata={"help": "Polynomial include bias for linear regression."}, + ) + polynomial_interaction_only: List[bool] = field( + default_factory=lambda: [True, False], + metadata={"help": "Polynomial interaction only for linear regression."}, + ) + fit_intercept: List[bool] = field( + default_factory=lambda: [True, False], + metadata={"help": "Fit intercept for linear regression."}, + ) + + @staticmethod + def get_type(): + return ExecutionTimePredictorType.LINEAR_REGRESSION + + +@dataclass +class RandomForrestExecutionTimePredictorConfig(BaseExecutionTimePredictorConfig): + num_estimators: List[int] = field( + default_factory=lambda: [250, 500, 750], + metadata={"help": "Number of estimators for random forest."}, + ) + max_depth: List[int] = field( + default_factory=lambda: [8, 16, 32], + metadata={"help": "Maximum depth for random forest."}, + ) + min_samples_split: List[int] = field( + default_factory=lambda: [2, 5, 10], + metadata={"help": "Minimum samples split for random forest."}, + ) + + @staticmethod + def get_type(): + return ExecutionTimePredictorType.RANDOM_FORREST diff --git a/frontier/config/metrics_config.py b/frontier/config/metrics_config.py new file mode 100644 index 00000000..613f9575 --- /dev/null +++ b/frontier/config/metrics_config.py @@ -0,0 +1,173 @@ +"""Metrics collection, output taxonomy and cache locations.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +import os +from typing import Optional + +from frontier.utils.output_paths import ( + validate_output_filename, + validate_run_id, +) + + +@dataclass +class MetricsConfig: + """Metric configuration.""" + + write_metrics: bool = field( + default=True, + metadata={"help": "Whether to write metrics."}, + ) + write_json_trace: bool = field( + default=False, + metadata={"help": "Whether to write json trace."}, + ) + wandb_project: Optional[str] = field( + default=None, + metadata={"help": "Weights & Biases project name."}, + ) + wandb_group: Optional[str] = field( + default=None, + metadata={"help": "Weights & Biases group name."}, + ) + wandb_run_name: Optional[str] = field( + default=None, + metadata={"help": "Weights & Biases run name."}, + ) + wandb_sweep_id: Optional[str] = field( + default=None, + metadata={"help": "Weights & Biases sweep id."}, + ) + wandb_run_id: Optional[str] = field( + default=None, + metadata={"help": "Weights & Biases run id."}, + ) + enable_chrome_trace: bool = field( + default=True, + metadata={"help": "Enable Chrome tracing."}, + ) + + # Op-Level Tracing + enable_op_level_tracing: bool = field( + default=False, + metadata={"help": "Enable detailed op-level tracing (output to JSONL)."}, + ) + trace_output_file: str = field( + default="op_traces.jsonl", + metadata={"help": "Output filename for op-level traces."}, + ) + enable_metrics_ground_truth_trace: bool = field( + default=False, + metadata={ + "help": "Enable explicit request-level metrics ground-truth JSONL output." + }, + ) + metrics_ground_truth_trace_file: str = field( + default="metrics_ground_truth.jsonl", + metadata={"help": "Output filename for metrics ground-truth request traces."}, + ) + enable_per_layer_expansion: bool = field( + default=False, + metadata={ + "help": "Enable per-layer trace expansion. When enabled, traces show " + "individual layer operations instead of aggregated spans." + }, + ) + num_requests_to_trace_per_layer: int = field( + default=5, + metadata={ + "help": "Number of requests to capture with per-layer expansion. " + "Only applies when enable_per_layer_expansion is True." + }, + ) + + save_table_to_wandb: bool = field( + default=False, + metadata={"help": "Whether to save table to wandb."}, + ) + store_plots: bool = field( + default=True, + metadata={"help": "Whether to store plots."}, + ) + enable_memory_time_series: bool = field( + default=False, + metadata={ + "help": "Enable memory usage time series output. " + "Only valid when log_level is 'debug'." + }, + ) + store_operation_metrics: bool = field( + default=False, + metadata={"help": "Whether to store operation metrics."}, + ) + store_token_completion_metrics: bool = field( + default=False, + metadata={"help": "Whether to store token completion metrics."}, + ) + store_request_metrics: bool = field( + default=True, + metadata={"help": "Whether to store request metrics."}, + ) + store_batch_metrics: bool = field( + default=True, + metadata={"help": "Whether to store batch metrics."}, + ) + store_utilization_metrics: bool = field( + default=True, + metadata={"help": "Whether to store utilization metrics."}, + ) + keep_individual_batch_metrics: bool = field( + default=False, + metadata={"help": "Whether to keep individual batch metrics."}, + ) + store_frontier_stage_batch_ledger: bool = field( + default=True, + metadata={"help": "Whether to write the full Frontier stage-batch ledger."}, + ) + store_frontier_stage_batch_ledger_summary: bool = field( + default=False, + metadata={ + "help": "Whether to write a bounded Frontier stage-batch ledger summary." + }, + ) + subsamples: Optional[int] = field( + default=None, + metadata={"help": "Subsamples."}, + ) + min_batch_index: Optional[int] = field( + default=None, + metadata={"help": "Minimum batch index."}, + ) + max_batch_index: Optional[int] = field( + default=None, + metadata={"help": "Maximum batch index."}, + ) + output_dir: str = field( + default="outputs/metrics", + metadata={"help": "Metrics output root directory."}, + ) + cache_dir: str = field( + default="cache", + metadata={"help": "Cache directory."}, + ) + run_id: Optional[str] = field( + default=None, + metadata={ + "help": "Metrics run id used under outputs/metrics///." + }, + ) + + def __post_init__(self): + if self.run_id is None: + self.run_id = f"run_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S-%f')}" + self.run_id = validate_run_id(self.run_id) + self.trace_output_file = validate_output_filename( + self.trace_output_file, "trace_output_file" + ) + self.metrics_ground_truth_trace_file = validate_output_filename( + self.metrics_ground_truth_trace_file, "metrics_ground_truth_trace_file" + ) + os.makedirs(self.output_dir, exist_ok=True) diff --git a/frontier/config/release_guards.py b/frontier/config/release_guards.py new file mode 100644 index 00000000..1e19de5c --- /dev/null +++ b/frontier/config/release_guards.py @@ -0,0 +1,56 @@ +"""Release-guard messages and the disaggregated cluster field tables. + +These are plain data shared by the configuration families, the schedulers, +the events and the metrics store. Keeping them in a leaf module lets every +configuration module import them without importing one another. +""" + + +DISAGGREGATED_ARCHITECTURE_RELEASE_ERROR = ( + "Error: Disaggregated architecture support is currently being optimized and is not included in this release. " + "It will be available in an upcoming version. Please use the co-located architecture for current usage and testing." +) + +PD_DISAGGREGATION_PARALLEL_CLUSTER_RELEASE_ERROR = ( + "Error: pd-disaggregation public release support requires " + "--no-enable_parallel_clusters. Parallel PDD is excluded from " + "pre-release-v0.3 because post-ISSUE-022 five-pair MoE-64 measurements " + "were slower than sequential: Simulator.run() by 35.29% and shell E2E " + "by 24.81% (paired medians). The implementation remains available only " + "to internal correctness tests." +) + +PD_AF_DISAGGREGATION_PARALLEL_CLUSTER_RELEASE_ERROR = ( + "Error: pd-af-disaggregation v0.3 requires " + "--no-enable_parallel_clusters. Parallel cluster processing for " + "pd-af-disaggregation is deferred in this release." +) + +PD_AF_PREFIX_CACHING_RELEASE_ERROR = ( + "Prefix caching is excluded for pd-af-disaggregation in v0.3. " + "Disable replica_scheduler_config.enable_prefix_caching." +) + +AICONFIGURATOR_BACKEND_RELEASE_ERROR = ( + "Error: The aiconfigurator communication backend is not included in this release. " + "Please use collective_sim, astra_sim_analytical, analytical, or vidur for current usage and testing." +) + +PD_AF_TRACE_REPLAY_DEFERRED_ERROR = ( + "Error: pd-af-disaggregation v0.3 trace-replay is deferred. " + "The configured trace-driven fields are public stubs only and are not " + "implemented in this release." +) + +DISAGGREGATED_CLUSTER_FIELD_PREFIXES = ( + "prefill_", + "decode_", + "decode_attn_", + "decode_ffn_", +) + +DISAGGREGATED_CLUSTER_FIELD_NAMES = frozenset( + { + "af_pipeline_num_micro_batch", + } +) diff --git a/frontier/config/replica_config.py b/frontier/config/replica_config.py new file mode 100644 index 00000000..6514f9d9 --- /dev/null +++ b/frontier/config/replica_config.py @@ -0,0 +1,249 @@ +"""Per-replica hardware, parallelism and model configuration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from frontier.config.device_sku_config import BaseDeviceSKUConfig +from frontier.config.model_config import BaseModelConfig +from frontier.config.node_sku_config import BaseNodeSKUConfig +from frontier.config.speculative_decoding_config import ( + SpeculativeDecodingConfig, +) + + +@dataclass +class ReplicaConfig: + memory_margin_fraction: float = field( + default=0.1, + metadata={"help": "Memory margin fraction."}, + ) + num_pipeline_stages: int = field( + default=1, + metadata={"help": "Number of pipeline stages (pp size)."}, + ) + attn_tensor_parallel_size: int = field( + default=1, + metadata={"help": "Attention tensor parallel size (attn_tp size)."}, + ) + attn_dp: int = field( + default=1, + metadata={ + "help": "Attention data-parallel lanes owned by one Replica.", + }, + ) + moe_tensor_parallel_size: int = field( + default=1, + metadata={"help": "MoE tensor parallel size (moe_tp size)."}, + ) + moe_expert_parallel_size: int = field( + default=1, + metadata={"help": "MoE expert parallel size (moe_ep size)."}, + ) + total_expert_num: int = field( + default=1, + metadata={"help": "Total expert number."}, + ) + router_load_balancing_type: str = field( + default="None", + metadata={"help": "MOE router load balancing type."}, + ) + router_topk: int = field( + default=0, + metadata={"help": "Router topk. Set to 0 to inherit from model config."}, + ) + moe_routing_seed: int = field( + default=42, + metadata={ + "help": "Random seed for deterministic MoE routing distribution generation. " + "Must be a non-negative integer." + }, + ) + moe_routing_trace_path: str = field( + default="", + metadata={ + "help": "Deferred StepFun merged trace JSONL for unsupported trace replay. " + "A non-empty path fails fast at the architecture boundary." + }, + ) + decode_attn_initial_lane_trace_path: str = field( + default="", + metadata={ + "help": "Optional StepFun attention trace JSONL for trace-driven " + "decode-attn initial lane occupancy and warmup replay." + }, + ) + decode_attn_steady_state_snapshot_path: str = field( + default="", + metadata={ + "help": "Optional StepFun attention trace JSONL for explicit " + "decode-attn steady-state snapshot hydration." + }, + ) + decode_attn_steady_state_measurement_report_path: str = field( + default="", + metadata={ + "help": "Optional StepFun measurement JSON for post-boundary " + "decode-attn request arrival replay." + }, + ) + moe_routing_distribution_type: str = field( + default="balanced", + metadata={ + "help": "MoE expert-load distribution for disaggregated routing simulation. " + "Valid values: 'balanced', 'random', 'skewed', or 'zipf'. This controls " + "token-to-expert load skew without changing router_topk/model semantics." + }, + ) + device: str = field( + default="a100", + metadata={"help": "Device."}, + ) + network_device: str = field( + default="a100_pairwise_nvlink", + metadata={"help": "Network device."}, + ) + speculative_decoding_config: SpeculativeDecodingConfig = field( + default_factory=SpeculativeDecodingConfig, + metadata={"help": "Speculative decoding simulation configuration."}, + ) + + # configs should be set by the user + cluster_prefix: str = None + local_expert_num: int = None + model_name: str = "meta-llama/Llama-2-7b-hf" + + def __post_init__(self): + if type(self.attn_dp) is not int or self.attn_dp <= 0: + raise ValueError( + "attn_dp must be a positive integer, " + f"got {self.attn_dp!r}" + ) + if self.cluster_prefix == "decode_attn" and self.attn_dp != 1: + raise ValueError( + "DECODE_ATTN requires attn_dp=1 because it is the PD-AF attention role" + ) + # Load model and device configs first (needed for validation) + self.model_config: BaseModelConfig = BaseModelConfig.create_from_name( + self.model_name + ) + self.device_config: BaseDeviceSKUConfig = ( + BaseDeviceSKUConfig.create_from_type_string(self.device) + ) + self.node_config: BaseNodeSKUConfig = BaseNodeSKUConfig.create_from_type_string( + self.network_device + ) + + # Auto-set total_expert_num from model config if not explicitly set and model is MoE + if ( + self.total_expert_num == 1 + and self.model_config.is_moe + and self.model_config.num_experts > 0 + ): + self.total_expert_num = self.model_config.num_experts + + # Align router_topk with model config when not explicitly set. + if self.model_config.is_moe: + if self.router_topk is None or int(self.router_topk) <= 0: + if self.model_config.num_experts_per_tok > 0: + self.router_topk = int(self.model_config.num_experts_per_tok) + else: + raise ValueError( + "router_topk is not set and model_config.num_experts_per_tok is missing" + ) + else: + if self.router_topk is None or int(self.router_topk) <= 0: + self.router_topk = 1 + + valid_moe_routing_distribution_types = { + "balanced", + "random", + "skewed", + "zipf", + } + self.moe_routing_distribution_type = str( + self.moe_routing_distribution_type + ).strip().lower() + if self.moe_routing_distribution_type not in valid_moe_routing_distribution_types: + raise ValueError( + "moe_routing_distribution_type must be one of " + f"{sorted(valid_moe_routing_distribution_types)}, " + f"got {self.moe_routing_distribution_type!r}" + ) + + # Validate pipeline parallelism configuration early + if self.model_config.num_layers % self.num_pipeline_stages != 0: + raise ValueError( + f"Pipeline parallelism configuration error: " + f"num_layers ({self.model_config.num_layers}) must be evenly divisible by " + f"num_pipeline_stages ({self.num_pipeline_stages}). " + f"Current configuration would result in uneven layer distribution across pipeline stages. " + f"Please adjust num_pipeline_stages to be a divisor of {self.model_config.num_layers}." + ) + + # Note: this world_size only limits in replica dimension. + if self.cluster_prefix == "prefill": + self.world_size = ( + self.num_pipeline_stages + * self.attn_tensor_parallel_size + * self.attn_dp + ) + elif self.cluster_prefix == "decode_attn": + self.world_size = ( + self.num_pipeline_stages + * self.attn_tensor_parallel_size + * self.attn_dp + ) + elif self.cluster_prefix == "decode_ffn": + self.world_size = ( + self.num_pipeline_stages + * self.moe_tensor_parallel_size + * self.moe_expert_parallel_size + ) + elif self.cluster_prefix == "decode": + # Unified decode cluster (PD-disaggregation): similar to prefill, includes both Attention and FFN + self.world_size = ( + self.num_pipeline_stages + * self.attn_tensor_parallel_size + * self.attn_dp + ) + else: # Monolithic + self.world_size = ( + self.num_pipeline_stages + * self.attn_tensor_parallel_size + * self.attn_dp + ) + + # Validate expert parallelism configuration for MoE models + # Use model_config.is_moe for MoE detection - NOT total_expert_num + if self.cluster_prefix != "decode_attn" and self.model_config.is_moe: + if self.total_expert_num > 1: + assert ( + self.total_expert_num % self.moe_expert_parallel_size == 0 + ), "total_expert_num must be divisible by moe_expert_parallel_size" + self.local_expert_num = ( + self.total_expert_num // self.moe_expert_parallel_size + ) + + if ( + self.speculative_decoding_config.enabled + and self.cluster_prefix in {"decode_attn", "decode_ffn"} + ): + raise ValueError( + "Speculative decoding Phase 1 supports only co-location and " + "pd-disaggregation decode path. decode_attn/decode_ffn are not " + f"supported, cluster_prefix={self.cluster_prefix!r}." + ) + + from frontier.attention.gdn.guards import validate_gdn_runtime_support + + validate_gdn_runtime_support( + self.model_config, + speculative_enabled=bool(self.speculative_decoding_config.enabled), + num_pipeline_stages=self.num_pipeline_stages, + moe_expert_parallel_size=self.moe_expert_parallel_size, + attn_dp=self.attn_dp, + cross_node=( + int(self.world_size) > int(self.node_config.num_devices_per_node) + ), + ) diff --git a/frontier/config/replica_scheduler_config.py b/frontier/config/replica_scheduler_config.py new file mode 100644 index 00000000..3839c59d --- /dev/null +++ b/frontier/config/replica_scheduler_config.py @@ -0,0 +1,711 @@ +"""Replica scheduler configuration for every supported batching policy.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Optional + +from frontier.config.base_poly_config import BasePolyConfig +from frontier.types import ReplicaSchedulerType + + +@dataclass +class BaseReplicaSchedulerConfig(BasePolyConfig): + batch_size_cap: int = field( + default=128, + metadata={"help": "Maximum batch size cap (max_num_seqs in vLLM)"}, + ) + block_size: int = field( + default=16, + metadata={"help": "Block size."}, + ) + watermark_blocks_fraction: float = field( + default=0.01, + metadata={"help": "Watermark blocks fraction."}, + ) + num_blocks: Optional[int] = field( + default=106596, + metadata={"help": "Number of blocks."}, + ) + + +@dataclass +class VllmSchedulerConfig(BaseReplicaSchedulerConfig): + max_tokens_in_batch: int = field( + default=4096, + metadata={"help": "Maximum tokens (max_num_batched_tokens) in batch for vLLM."}, + ) + + @staticmethod + def get_type(): + return ReplicaSchedulerType.VLLM + + +@dataclass +class LightllmSchedulerConfig(BaseReplicaSchedulerConfig): + max_tokens_in_batch: int = field( + default=4096, + metadata={"help": "Maximum tokens in batch for LightLLM."}, + ) + max_waiting_iters: int = field( + default=10, + metadata={"help": "Maximum waiting iterations for LightLLM."}, + ) + + @staticmethod + def get_type(): + return ReplicaSchedulerType.LIGHTLLM + + +@dataclass +class OrcaSchedulerConfig(BaseReplicaSchedulerConfig): + @staticmethod + def get_type(): + return ReplicaSchedulerType.ORCA + + +@dataclass +class FasterTransformerSchedulerConfig(BaseReplicaSchedulerConfig): + @staticmethod + def get_type(): + return ReplicaSchedulerType.FASTER_TRANSFORMER + + +@dataclass +class SarathiSchedulerConfig(BaseReplicaSchedulerConfig): + chunk_size: int = field( + default=512, + metadata={"help": "Chunk size for Sarathi."}, + ) + + @staticmethod + def get_type(): + return ReplicaSchedulerType.SARATHI + + +@dataclass +class VllmV1SchedulerConfig(BaseReplicaSchedulerConfig): + """ + Configuration for the vLLM v1 engine replica scheduler. + + This scheduler simulates the admission control behavior of vLLM v1 engine, + including two-phase scheduling, token budget management, and preemption. + + Note: Class name uses 'VllmV1' (not 'VLLMv1') to generate clean CLI parameter + names: --vllm_v1_scheduler_config_* instead of --v_l_l_mv1_scheduler_config_* + """ + + max_tokens_in_batch: int = field( + default=16384, + metadata={ + "help": "Maximum tokens per scheduling iteration (max_num_batched_tokens in vLLM v1)." + }, + ) + scheduling_policy: str = field( + default="fcfs", + metadata={ + "help": "Scheduling policy: 'fcfs' (First-Come-First-Served) or 'priority'." + }, + ) + enable_preemption: bool = field( + default=True, + metadata={ + "help": "Enable preemption when memory is insufficient for running requests." + }, + ) + enable_chunked_prefill: bool = field( + default=False, + metadata={ + "help": "Enable chunked prefill admission when waiting prefill requests exceed current token budget." + }, + ) + enable_phase_aware_thinking_profile: bool = field( + default=False, + metadata={ + "help": "Enable an iteration-scoped hidden-round/final-round scheduler profile override for Thinking Mode home queues." + }, + ) + hidden_phase_max_tokens_in_batch: Optional[int] = field( + default=None, + metadata={ + "help": "Optional hidden-round override for max_tokens_in_batch when enable_phase_aware_thinking_profile=True." + }, + ) + hidden_phase_enable_chunked_prefill: Optional[bool] = field( + default=None, + metadata={ + "help": "Optional hidden-round override for enable_chunked_prefill when enable_phase_aware_thinking_profile=True." + }, + ) + hidden_phase_batch_size_cap: Optional[int] = field( + default=None, + metadata={ + "help": "Optional hidden-round override for batch_size_cap when enable_phase_aware_thinking_profile=True." + }, + ) + final_phase_max_tokens_in_batch: Optional[int] = field( + default=None, + metadata={ + "help": "Optional final-round override for max_tokens_in_batch when enable_phase_aware_thinking_profile=True." + }, + ) + final_phase_enable_chunked_prefill: Optional[bool] = field( + default=None, + metadata={ + "help": "Optional final-round override for enable_chunked_prefill when enable_phase_aware_thinking_profile=True." + }, + ) + final_phase_batch_size_cap: Optional[int] = field( + default=None, + metadata={ + "help": "Optional final-round override for batch_size_cap when enable_phase_aware_thinking_profile=True." + }, + ) + final_prefill_reserved_slots: int = field( + default=0, + metadata={ + "help": "Per-iteration PREFILL admission slots reserved for final-round prefill requests. Hidden requests may borrow idle reserved slots." + }, + ) + final_prefill_reserved_tokens: int = field( + default=0, + metadata={ + "help": "Per-iteration PREFILL token budget reserved for final-round prefill requests. Hidden requests may borrow idle reserved tokens." + }, + ) + final_decode_reserved_slots: int = field( + default=0, + metadata={ + "help": "Per-iteration DECODE running/admission slots reserved for final-round decode requests. Hidden requests may borrow idle reserved slots." + }, + ) + enable_final_running_request_reclaim: bool = field( + default=False, + metadata={ + "help": "When final backlog appears, reclaim hidden requests that have borrowed final reserved running slots so the final slice becomes active running capacity." + }, + ) + enable_final_round_priority_boost: bool = field( + default=False, + metadata={ + "help": "Promote re-entered final-round Thinking Mode requests into a higher-priority band under priority scheduling." + }, + ) + final_round_priority_value: int = field( + default=-1, + metadata={ + "help": "Priority value assigned to promoted final-round requests. Lower values mean higher priority." + }, + ) + enable_prefix_caching: bool = field( + default=False, + metadata={ + "help": "Enable block-hash-based prefix matching and KV cache reuse." + }, + ) + prefix_caching_hash_algo: str = field( + default="builtin", + metadata={ + "help": "Hash algorithm label for explicit prefix block hashes. Supported: 'builtin', 'sha256'." + }, + ) + num_preallocate_tokens: int = field( + default=0, + metadata={ + "help": "Number of tokens worth of KV cache blocks to preallocate for each request." + }, + ) + long_prefill_token_threshold: int = field( + default=0, + metadata={ + "help": "Optional upper bound on per-iteration prefill tokens for each request. 0 disables threshold." + }, + ) + num_blocks: Optional[int] = field( + default=0, + metadata={ + "help": "Number of KV cache blocks. Use 0 to auto-derive from the memory planner in planner modes." + }, + ) + num_blocks_mode: str = field( + default="memory_planner_profiled", + metadata={ + "help": "How to initialize num_blocks: 'memory_planner' (auto-derive with parameter-only estimate), 'memory_planner_profiled' (auto-derive with calibrated non-KV overhead), or 'explicit' (require num_blocks>0)." + }, + ) + gpu_memory_utilization: Optional[float] = field( + default=None, + metadata={ + "help": "vLLM-style GPU memory utilization ratio used by memory_planner mode. If unset, fallback to 1 - replica memory_margin_fraction." + }, + ) + non_kv_cache_overhead_bytes: int = field( + default=0, + metadata={ + "help": "Calibrated non-KV memory overhead in bytes for memory_planner_profiled mode." + }, + ) + runtime_weights_memory_source: str = field( + default="param_counter", + metadata={ + "help": "Weights memory source for runtime non-KV profiling: 'param_counter' (estimated bytes) or 'runtime_model_load' (measure loaded model parameter bytes)." + }, + ) + enable_runtime_non_kv_cache_overhead_profiling: bool = field( + default=False, + metadata={ + "help": "Enable runtime single-rank profiling to auto-estimate non_kv_cache_overhead_bytes during scheduler initialization. Requires num_blocks_mode=memory_planner_profiled." + }, + ) + nccl_buffer_comm_base_overhead_bytes: int = field( + default=100 * 1024 * 1024, + metadata={ + "help": "Per-communicator fixed NCCL overhead in bytes (proxy, queues). " + "Default 100 MiB, calibrated for A800." + }, + ) + nccl_buffer_per_peer_overhead_bytes: int = field( + default=15 * 1024 * 1024, + metadata={ + "help": "Per-peer NCCL transport buffer overhead in bytes. " + "Default 15 MiB, calibrated for A800 intra-node." + }, + ) + nccl_buffer_custom_ar_enabled: bool = field( + default=False, + metadata={ + "help": "Enable CustomAllreduce buffer estimation. " + "False for A800 (compute 8.0), True for H100 (9.0+)." + }, + ) + nccl_buffer_vllm_worker_base_extra_bytes: int = field( + default=0, + metadata={ + "help": "Domain-aware vLLM worker-process non-torch addend in bytes " + "for runtime non-KV profiling. Default 0; pass validated " + "case-local values explicitly." + }, + ) + nccl_buffer_pp_final_stage_extra_bytes: int = field( + default=0, + metadata={ + "help": "Additional final pipeline-stage vLLM worker non-torch addend " + "in bytes for runtime non-KV profiling. Default 0." + }, + ) + nccl_buffer_dp_communicator_extra_bytes: int = field( + default=0, + metadata={ + "help": "Additional data-parallel communicator non-torch addend in " + "bytes for runtime non-KV profiling. Default 0." + }, + ) + nccl_buffer_ep_all2all_extra_bytes: int = field( + default=0, + metadata={ + "help": "Additional MoE expert-parallel all-to-all non-torch addend " + "in bytes for runtime non-KV profiling. Default 0." + }, + ) + use_analytical_param_memory: bool = field( + default=False, + metadata={ + "help": "When runtime non-KV profiling is enabled in memory_planner_profiled mode, keep analytical ParamCounter param memory for planner calculation. Default False uses runtime-profiled param memory." + }, + ) + + def __post_init__(self) -> None: + allowed_modes = {"memory_planner", "memory_planner_profiled", "explicit"} + if self.num_blocks_mode not in allowed_modes: + raise ValueError( + "VllmV1SchedulerConfig.num_blocks_mode must be one of " + f"{sorted(allowed_modes)}, got={self.num_blocks_mode!r}" + ) + + if self.gpu_memory_utilization is not None: + if self.gpu_memory_utilization <= 0 or self.gpu_memory_utilization > 1.0: + raise ValueError( + "VllmV1SchedulerConfig.gpu_memory_utilization must be in (0, 1], got=" + f"{self.gpu_memory_utilization!r}" + ) + + if self.non_kv_cache_overhead_bytes < 0: + raise ValueError( + "VllmV1SchedulerConfig.non_kv_cache_overhead_bytes must be >= 0, got=" + f"{self.non_kv_cache_overhead_bytes!r}" + ) + + allowed_hash_algorithms = {"builtin", "sha256"} + if self.prefix_caching_hash_algo not in allowed_hash_algorithms: + raise ValueError( + "VllmV1SchedulerConfig.prefix_caching_hash_algo must be one of " + f"{sorted(allowed_hash_algorithms)}, got={self.prefix_caching_hash_algo!r}" + ) + + if self.num_preallocate_tokens < 0: + raise ValueError( + "VllmV1SchedulerConfig.num_preallocate_tokens must be >= 0, got=" + f"{self.num_preallocate_tokens!r}" + ) + + if self.long_prefill_token_threshold < 0: + raise ValueError( + "VllmV1SchedulerConfig.long_prefill_token_threshold must be >= 0, got=" + f"{self.long_prefill_token_threshold!r}" + ) + if ( + self.long_prefill_token_threshold > 0 + and not self.enable_chunked_prefill + ): + raise ValueError( + "VllmV1SchedulerConfig.long_prefill_token_threshold > 0 " + "requires enable_chunked_prefill=True" + ) + + phase_override_values = ( + self.hidden_phase_max_tokens_in_batch, + self.hidden_phase_enable_chunked_prefill, + self.hidden_phase_batch_size_cap, + self.final_phase_max_tokens_in_batch, + self.final_phase_enable_chunked_prefill, + self.final_phase_batch_size_cap, + ) + if not self.enable_phase_aware_thinking_profile and any( + value is not None for value in phase_override_values + ): + raise ValueError( + "VllmV1SchedulerConfig phase-aware override fields require " + "enable_phase_aware_thinking_profile=True" + ) + if self.enable_phase_aware_thinking_profile and all( + value is None for value in phase_override_values + ): + raise ValueError( + "VllmV1SchedulerConfig.enable_phase_aware_thinking_profile=True " + "requires at least one hidden/final override field" + ) + + for field_name in ( + "hidden_phase_max_tokens_in_batch", + "hidden_phase_batch_size_cap", + "final_phase_max_tokens_in_batch", + "final_phase_batch_size_cap", + ): + field_value = getattr(self, field_name) + if field_value is not None and field_value <= 0: + raise ValueError( + f"VllmV1SchedulerConfig.{field_name} must be > 0 when set, " + f"got={field_value!r}" + ) + + for field_name in ( + "final_prefill_reserved_slots", + "final_prefill_reserved_tokens", + "final_decode_reserved_slots", + ): + field_value = getattr(self, field_name) + if field_value < 0: + raise ValueError( + f"VllmV1SchedulerConfig.{field_name} must be >= 0, " + f"got={field_value!r}" + ) + + if ( + self.long_prefill_token_threshold > 0 + and self.enable_phase_aware_thinking_profile + ): + if self.hidden_phase_enable_chunked_prefill is False: + raise ValueError( + "VllmV1SchedulerConfig.hidden_phase_enable_chunked_prefill=False " + "is incompatible with long_prefill_token_threshold > 0" + ) + if self.final_phase_enable_chunked_prefill is False: + raise ValueError( + "VllmV1SchedulerConfig.final_phase_enable_chunked_prefill=False " + "is incompatible with long_prefill_token_threshold > 0" + ) + + if self.nccl_buffer_comm_base_overhead_bytes < 0: + raise ValueError( + "VllmV1SchedulerConfig.nccl_buffer_comm_base_overhead_bytes must be >= 0, got=" + f"{self.nccl_buffer_comm_base_overhead_bytes!r}" + ) + + if self.nccl_buffer_per_peer_overhead_bytes < 0: + raise ValueError( + "VllmV1SchedulerConfig.nccl_buffer_per_peer_overhead_bytes must be >= 0, got=" + f"{self.nccl_buffer_per_peer_overhead_bytes!r}" + ) + + for field_name in ( + "nccl_buffer_vllm_worker_base_extra_bytes", + "nccl_buffer_pp_final_stage_extra_bytes", + "nccl_buffer_dp_communicator_extra_bytes", + "nccl_buffer_ep_all2all_extra_bytes", + ): + field_value = getattr(self, field_name) + if field_value < 0: + raise ValueError( + f"VllmV1SchedulerConfig.{field_name} must be >= 0, " + f"got={field_value!r}" + ) + + allowed_weights_sources = {"param_counter", "runtime_model_load"} + if self.runtime_weights_memory_source not in allowed_weights_sources: + raise ValueError( + "VllmV1SchedulerConfig.runtime_weights_memory_source must be one of " + f"{sorted(allowed_weights_sources)}, got={self.runtime_weights_memory_source!r}" + ) + + if ( + self.enable_runtime_non_kv_cache_overhead_profiling + and self.num_blocks_mode != "memory_planner_profiled" + ): + raise ValueError( + "VllmV1SchedulerConfig.enable_runtime_non_kv_cache_overhead_profiling " + "requires num_blocks_mode=memory_planner_profiled, got=" + f"{self.num_blocks_mode!r}" + ) + + if ( + self.use_analytical_param_memory + and not self.enable_runtime_non_kv_cache_overhead_profiling + ): + raise ValueError( + "VllmV1SchedulerConfig.use_analytical_param_memory " + "requires enable_runtime_non_kv_cache_overhead_profiling=True" + ) + + enable_thinking_round_priority: bool = field( + default=False, + metadata={ + "help": "When enabled, final-round thinking requests are prioritized " + "over non-final-round requests in the waiting queue.", + }, + ) + + @staticmethod + def get_type(): + return ReplicaSchedulerType.VLLM_V1 + + +@dataclass +class Sj2qFastserveLiteSchedulerConfig(VllmV1SchedulerConfig): + """ + Configuration for the SJ-2Q / FastServe-lite scheduler. + + Note: Class name uses 'Sj2qFastserve' (not 'Sj2QFastServe') to generate clean + CLI parameter names: --sj2q_fastserve_lite_scheduler_config_* instead of + --sj2_q_fast_serve_lite_scheduler_config_*. + """ + + long_round_new_prompt_threshold: int = field( + default=2048, + metadata={ + "help": "Rounds whose new prompt tokens exceed this threshold enter QL and mark long_history." + }, + ) + short_round_boost_threshold: int = field( + default=512, + metadata={ + "help": "Tiny-prefill threshold used for QH prioritization and the prefill-release-only boost when long_history is already true." + }, + ) + boost_credit_token_budget: int = field( + default=2048, + metadata={ + "help": "Deprecated compatibility field retained for CLI stability; current prefill-release-only boost demotes on prefill completion instead of token-budget exhaustion." + }, + ) + enable_aging: bool = field( + default=False, + metadata={ + "help": "Enable optional aging-based QL promotion back into QH. The UC3 v2 enhancement lane keeps this disabled." + }, + ) + aging_wait_threshold_ms: float = field( + default=7.5, + metadata={ + "help": "QL waiting-time threshold in milliseconds for a temporary aging-based QH boost." + }, + ) + aging_boost_token_budget: int = field( + default=512, + metadata={ + "help": "Token budget granted when an aged QL session is temporarily promoted into QH." + }, + ) + + def __post_init__(self) -> None: + super().__post_init__() + + if self.enable_phase_aware_thinking_profile: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig does not allow phase-aware oracle scheduling." + ) + if self.enable_thinking_round_priority: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig does not allow final-round priority override." + ) + if ( + self.final_prefill_reserved_slots != 0 + or self.final_prefill_reserved_tokens != 0 + or self.final_decode_reserved_slots != 0 + ): + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig requires all final reserved slot/token settings to remain 0." + ) + if self.enable_final_running_request_reclaim: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig does not allow final running-request reclaim." + ) + if self.enable_final_round_priority_boost: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig does not allow final-round priority boost." + ) + + if self.long_round_new_prompt_threshold <= 0: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig.long_round_new_prompt_threshold must be > 0." + ) + if self.short_round_boost_threshold <= 0: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig.short_round_boost_threshold must be > 0." + ) + if ( + self.short_round_boost_threshold + > self.long_round_new_prompt_threshold + ): + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig.short_round_boost_threshold must be <= long_round_new_prompt_threshold." + ) + if self.boost_credit_token_budget <= 0: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig.boost_credit_token_budget must be > 0." + ) + if self.enable_aging and self.aging_wait_threshold_ms <= 0: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig.aging_wait_threshold_ms must be > 0 when aging is enabled." + ) + if self.aging_boost_token_budget <= 0: + raise ValueError( + "Sj2QFastserveLiteSchedulerConfig.aging_boost_token_budget must be > 0." + ) + + @staticmethod + def get_type(): + return ReplicaSchedulerType.SJ2Q_FASTSERVE_LITE + + +Sj2QFastServeLiteSchedulerConfig = Sj2qFastserveLiteSchedulerConfig + + +@dataclass +class Sj2qPenaltyOnlySchedulerConfig(VllmV1SchedulerConfig): + """ + Configuration for the penalty-only SJ-2Q scheduler. + + Note: Class name uses 'Sj2q' to generate clean CLI parameter names like + --sj2q_penalty_only_scheduler_config_*. + """ + + long_round_new_prompt_threshold: int = field( + default=4096, + metadata={ + "help": "Rounds whose new prompt tokens exceed this threshold immediately enter Qlong and mark long_history." + }, + ) + service_cap_tokens: int = field( + default=8192, + metadata={ + "help": "Session-level cumulative new-token service cap after which the session stays in Qlong." + }, + ) + long_liveness_quota: int = field( + default=32, + metadata={ + "help": "Maximum consecutive Qshort slices allowed before forcing one Qlong slice when Qlong is non-empty." + }, + ) + + def __post_init__(self) -> None: + super().__post_init__() + + if self.enable_phase_aware_thinking_profile: + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig does not allow phase-aware oracle scheduling." + ) + if self.enable_thinking_round_priority: + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig does not allow final-round priority override." + ) + if ( + self.final_prefill_reserved_slots != 0 + or self.final_prefill_reserved_tokens != 0 + or self.final_decode_reserved_slots != 0 + ): + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig requires all final reserved slot/token settings to remain 0." + ) + if self.enable_final_running_request_reclaim: + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig does not allow final running-request reclaim." + ) + if self.enable_final_round_priority_boost: + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig does not allow final-round priority boost." + ) + if self.long_round_new_prompt_threshold <= 0: + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig.long_round_new_prompt_threshold must be > 0." + ) + if self.service_cap_tokens <= 0: + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig.service_cap_tokens must be > 0." + ) + if self.service_cap_tokens < self.long_round_new_prompt_threshold: + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig.service_cap_tokens must be >= long_round_new_prompt_threshold." + ) + if self.long_liveness_quota <= 0: + raise ValueError( + "Sj2qPenaltyOnlySchedulerConfig.long_liveness_quota must be > 0." + ) + + @staticmethod + def get_type(): + return ReplicaSchedulerType.SJ2Q_PENALTY_ONLY + + +Sj2QPenaltyOnlySchedulerConfig = Sj2qPenaltyOnlySchedulerConfig + + +@dataclass +class Sj2qBoundedCarryoverSchedulerConfig(Sj2qPenaltyOnlySchedulerConfig): + """ + Configuration for the bounded-carryover SJ-2Q scheduler. + + Note: Class name uses 'Sj2q' to generate clean CLI parameter names like + --sj2q_bounded_carryover_scheduler_config_*. + """ + + @staticmethod + def get_type(): + return ReplicaSchedulerType.SJ2Q_BOUNDED_CARRYOVER + + +Sj2QBoundedCarryoverSchedulerConfig = Sj2qBoundedCarryoverSchedulerConfig + + +@dataclass +class SglangSchedulerConfig(VllmV1SchedulerConfig): + """ + Thin config wrapper for the Frontier SGLang-style replica scheduler. + + This intentionally reuses the vLLM v1 scheduler fields and only changes + the scheduler type to keep the integration surface minimal. + """ + + @staticmethod + def get_type(): + return ReplicaSchedulerType.SGLANG diff --git a/frontier/config/request_generator_config.py b/frontier/config/request_generator_config.py new file mode 100644 index 00000000..3fe2b5ac --- /dev/null +++ b/frontier/config/request_generator_config.py @@ -0,0 +1,262 @@ +"""Workload generation configuration: arrival intervals, lengths, generators.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Optional + +from frontier.config.base_poly_config import BasePolyConfig +from frontier.types import ( + RequestGeneratorType, + RequestIntervalGeneratorType, + RequestLengthGeneratorType, +) + + +@dataclass +class BaseRequestIntervalGeneratorConfig(BasePolyConfig): + seed: int = field( + default=42, + metadata={"help": "Seed for the random number generator."}, + ) + + +@dataclass +class BaseRequestLengthGeneratorConfig(BasePolyConfig): + seed: int = field( + default=42, + metadata={"help": "Seed for the random number generator."}, + ) + max_tokens: int = field( + default=4096, + metadata={"help": "Maximum tokens."}, + ) + + +@dataclass +class TraceRequestIntervalGeneratorConfig(BaseRequestIntervalGeneratorConfig): + trace_file: str = field( + default="data/processed_traces/AzureFunctionsInvocationTraceForTwoWeeksJan2021Processed.csv", + metadata={"help": "Path to the trace request interval generator file."}, + ) + start_time: str = field( + default="1970-01-04 12:00:00", + metadata={"help": "Start time of the trace request interval generator."}, + ) + end_time: str = field( + default="1970-01-04 15:00:00", + metadata={"help": "End time of the trace request interval generator."}, + ) + time_scale_factor: float = field( + default=1.0, + metadata={ + "help": "Time scale factor for the trace request interval generator." + }, + ) + + @staticmethod + def get_type(): + return RequestIntervalGeneratorType.TRACE + + +@dataclass +class PoissonRequestIntervalGeneratorConfig(BaseRequestIntervalGeneratorConfig): + qps: float = field( + default=0.5, + metadata={"help": "Queries per second for Poisson Request Interval Generator."}, + ) + + @staticmethod + def get_type(): + return RequestIntervalGeneratorType.POISSON + + +@dataclass +class GammaRequestIntervalGeneratorConfig(BaseRequestIntervalGeneratorConfig): + qps: float = field( + default=0.2, + metadata={"help": "Queries per second for Gamma Request Interval Generator."}, + ) + cv: float = field( + default=0.5, + metadata={ + "help": "Coefficient of variation for Gamma Request Interval Generator." + }, + ) + + @staticmethod + def get_type(): + return RequestIntervalGeneratorType.GAMMA + + +@dataclass +class StaticRequestIntervalGeneratorConfig(BaseRequestIntervalGeneratorConfig): + @staticmethod + def get_type(): + return RequestIntervalGeneratorType.STATIC + + +@dataclass +class TraceRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): + trace_file: str = field( + default="data/processed_traces/sharegpt_8k_filtered_stats_llama2_tokenizer.csv", + metadata={"help": "Path to the trace request length generator file."}, + ) + prefill_scale_factor: float = field( + default=1, + metadata={ + "help": "Prefill scale factor for the trace request length generator." + }, + ) + decode_scale_factor: float = field( + default=1, + metadata={ + "help": "Decode scale factor for the trace request length generator." + }, + ) + + @staticmethod + def get_type(): + return RequestLengthGeneratorType.TRACE + + +@dataclass +class ZipfRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): + theta: float = field( + default=0.6, + metadata={"help": "Theta for Zipf Request Length Generator."}, + ) + scramble: bool = field( + default=False, + metadata={"help": "Scramble for Zipf Request Length Generator."}, + ) + min_tokens: int = field( + default=1024, + metadata={"help": "Minimum tokens for Zipf Request Length Generator."}, + ) + prefill_to_decode_ratio: float = field( + default=20.0, + metadata={"help": "Prefill to decode ratio for Zipf Request Length Generator."}, + ) + + @staticmethod + def get_type(): + return RequestLengthGeneratorType.ZIPF + + +@dataclass +class UniformRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): + min_tokens: int = field( + default=1024, + metadata={"help": "Minimum tokens for Uniform Request Length Generator."}, + ) + prefill_to_decode_ratio: float = field( + default=20.0, + metadata={ + "help": "Prefill to decode ratio for Uniform Request Length Generator." + }, + ) + + @staticmethod + def get_type(): + return RequestLengthGeneratorType.UNIFORM + + +@dataclass +class FixedRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): + prefill_tokens: int = field( + default=2048, + metadata={"help": "Prefill tokens for Fixed Request Length Generator."}, + ) + decode_tokens: int = field( + default=512, + metadata={"help": "Decode tokens for Fixed Request Length Generator."}, + ) + + @staticmethod + def get_type(): + return RequestLengthGeneratorType.FIXED + + def __post_init__(self): + if self.decode_tokens < 1: + raise ValueError(f"decode_tokens must be >= 1, got {self.decode_tokens}") + if self.prefill_tokens < 2: + raise ValueError(f"prefill_tokens must be >1, got {self.prefill_tokens}") + + +@dataclass +class BaseRequestGeneratorConfig(BasePolyConfig): + seed: int = field( + default=42, + metadata={"help": "Seed for the random number generator."}, + ) + num_decode_bound_requests: Optional[int] = field( + default=None, + metadata={ + "help": "Number of generated requests that require decode-cluster work. " + "Derived by request generation and used by offline pd-disaggregation scheduling." + }, + ) + + +@dataclass +class SyntheticRequestGeneratorConfig(BaseRequestGeneratorConfig): + length_generator_config: BaseRequestLengthGeneratorConfig = field( + default_factory=FixedRequestLengthGeneratorConfig, + metadata={"help": "Length generator config for Synthetic Request Generator."}, + ) + interval_generator_config: BaseRequestIntervalGeneratorConfig = field( + default_factory=PoissonRequestIntervalGeneratorConfig, + metadata={"help": "Interval generator config for Synthetic Request Generator."}, + ) + num_requests: Optional[int] = field( + default=128, + metadata={"help": "Number of requests for Synthetic Request Generator."}, + ) + duration: Optional[float] = field( + default=None, + metadata={"help": "Duration of the synthetic request generator."}, + ) + default_priority: int = field( + default=0, + metadata={ + "help": "Default priority for all generated requests. " + "Lower value = higher priority (0 = highest). " + "Matches vLLM v1 semantics." + }, + ) + + def __post_init__(self): + self.max_tokens = self.length_generator_config.max_tokens + + @staticmethod + def get_type(): + return RequestGeneratorType.SYNTHETIC + + +@dataclass +class TraceRequestGeneratorConfig(BaseRequestGeneratorConfig): + trace_file: str = field( + default="data/processed_traces/splitwise_conv.csv", + metadata={"help": "Path to the trace request generator file."}, + ) + prefill_scale_factor: float = field( + default=1.0, + metadata={"help": "Prefill scale factor for the trace request generator."}, + ) + decode_scale_factor: float = field( + default=1.0, + metadata={"help": "Decode scale factor for the trace request generator."}, + ) + time_scale_factor: float = field( + default=1.0, + metadata={"help": "Time scale factor for the trace request generator."}, + ) + max_tokens: int = field( + default=4096, + metadata={"help": "Maximum tokens for the trace request generator."}, + ) + + @staticmethod + def get_type(): + return RequestGeneratorType.TRACE_REPLAY diff --git a/frontier/config/speculative_decoding_config.py b/frontier/config/speculative_decoding_config.py new file mode 100644 index 00000000..16c7be9f --- /dev/null +++ b/frontier/config/speculative_decoding_config.py @@ -0,0 +1,622 @@ +"""Speculative decoding and MTP configuration, including trace admission.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +import os +from typing import Dict, List, Optional + +from frontier.spec_decode.proposer_profile import ( + load_decode_draft_proposer_latency_profile, +) + + +@dataclass +class SpeculativeDecodingConfig: + enabled: bool = field( + default=False, + metadata={"help": "Enable speculative decoding simulation."}, + ) + method: str = field( + default="eagle", + metadata={ + "help": "Speculative decoding method. Must match vLLM method names." + }, + ) + spec_model_name: str = field( + default="", + metadata={ + "help": "Optional draft/spec model name for methods whose proposer " + "decoder comes from a separate draft model (for example draft-model MTP)." + }, + ) + num_speculative_tokens: int = field( + default=4, + metadata={"help": "Number of draft tokens planned per speculative iteration."}, + ) + committed_tokens_per_iteration: int = field( + default=2, + metadata={ + "help": "Deterministic committed token count per speculative iteration " + "(includes 1 target token + accepted drafts)." + }, + ) + acceptance_trace_file: str = field( + default="", + metadata={ + "help": "Optional deterministic acceptance trace JSON file. Supported " + "formats: list[int] or {'committed_tokens_per_iteration': list[int], " + "'scheduled_draft_tokens_per_iteration': optional list[int], " + "'per_request_committed_tokens_per_iteration': optional dict[str, list[int]], " + "'per_request_scheduled_draft_tokens_per_iteration': optional dict[str, list[int]]}. " + "When set, trace overrides committed_tokens_per_iteration and can " + "optionally override planned draft widths per iteration." + }, + ) + proposer_overhead_ms_by_method: Dict[str, float] = field( + default_factory=dict, + metadata={ + "help": "Method-aware proposer overhead in milliseconds per speculative " + "verify request (method -> overhead_ms >= 0)." + }, + ) + decode_draft_proposer_latency_profile_file: str = field( + default="", + metadata={ + "help": "Optional structured latency profile JSON for decode draft " + "proposer overhead. Expected workload key: " + "(method, model_name, attn_tp_size, num_speculative_tokens, " + "spec_verify_request_count)." + }, + ) + mtp_n_predict: int = field( + default=0, + metadata={ + "help": "Optional MTP capability metadata. Number of tokens predicted " + "per MTP block. Only valid for MTP methods." + }, + ) + mtp_num_layers: int = field( + default=0, + metadata={ + "help": "Optional MTP capability metadata. Number of MTP layers. " + "Only valid for MTP methods." + }, + ) + trace_calibration_file: str = field( + default="", + metadata={ + "help": "Optional calibration JSON file. Supported keys: " + "proposer_overhead_ms_by_method and metadata." + }, + ) + + @staticmethod + def _validate_method_float_map( + *, + map_name: str, + raw_map: Optional[Dict[str, float]], + supported_methods: set[str], + min_value: float, + inclusive_min: bool, + ) -> Dict[str, float]: + if raw_map is None: + return {} + if not isinstance(raw_map, dict): + raise ValueError( + f"SpeculativeDecodingConfig.{map_name} must be a dict, " + f"got={type(raw_map).__name__}" + ) + + validated: Dict[str, float] = {} + for method_name, value in raw_map.items(): + if method_name not in supported_methods: + raise ValueError( + f"SpeculativeDecodingConfig.{map_name} contains unsupported method " + f"{method_name!r}; supported={sorted(supported_methods)}" + ) + numeric_value = float(value) + if inclusive_min: + if numeric_value < min_value: + raise ValueError( + f"SpeculativeDecodingConfig.{map_name}[{method_name!r}] " + f"must be >= {min_value}, got={numeric_value!r}" + ) + elif numeric_value <= min_value: + raise ValueError( + f"SpeculativeDecodingConfig.{map_name}[{method_name!r}] " + f"must be > {min_value}, got={numeric_value!r}" + ) + validated[method_name] = numeric_value + return validated + + @staticmethod + def _load_trace_calibration_payload( + trace_calibration_file: str, + ) -> Dict[str, Dict[str, float]]: + if not trace_calibration_file: + return {} + if not os.path.isfile(trace_calibration_file): + raise ValueError( + "SpeculativeDecodingConfig.trace_calibration_file does not exist: " + f"{trace_calibration_file!r}" + ) + try: + with open(trace_calibration_file, "r", encoding="utf-8") as f: + payload = json.load(f) + except json.JSONDecodeError as exc: + raise ValueError( + "SpeculativeDecodingConfig.trace_calibration_file must be valid JSON: " + f"{trace_calibration_file!r}" + ) from exc + + if not isinstance(payload, dict): + raise ValueError( + "SpeculativeDecodingConfig.trace_calibration_file must contain a JSON " + f"object, got={type(payload).__name__}" + ) + return payload + + @staticmethod + def _load_acceptance_trace_payload( + *, + acceptance_trace_file: str, + ): + if not acceptance_trace_file: + return None + if not os.path.isfile(acceptance_trace_file): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file does not exist: " + f"{acceptance_trace_file!r}" + ) + try: + with open(acceptance_trace_file, "r", encoding="utf-8") as f: + payload = json.load(f) + except json.JSONDecodeError as exc: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file must be valid JSON: " + f"{acceptance_trace_file!r}" + ) from exc + + if not isinstance(payload, (list, dict)): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file must be list or dict, " + f"got={type(payload).__name__}" + ) + return payload + + @staticmethod + def _load_committed_tokens_trace( + *, + acceptance_trace_payload, + max_committed_tokens: int, + ) -> Optional[List[int]]: + if acceptance_trace_payload is None: + return None + + if isinstance(acceptance_trace_payload, list): + committed_tokens_trace_raw = acceptance_trace_payload + else: + if "committed_tokens_per_iteration" not in acceptance_trace_payload: + return None + committed_tokens_trace_raw = acceptance_trace_payload[ + "committed_tokens_per_iteration" + ] + + if not isinstance(committed_tokens_trace_raw, list): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file committed token trace " + f"must be a list, got={type(committed_tokens_trace_raw).__name__}" + ) + if len(committed_tokens_trace_raw) == 0: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file committed token trace " + "must be non-empty." + ) + + committed_tokens_trace: List[int] = [] + for idx, value in enumerate(committed_tokens_trace_raw): + committed = int(value) + if committed < 0: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file values must be >= 0, " + f"got index={idx}, value={value!r}" + ) + if committed > max_committed_tokens: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file values must be <= " + f"1 + num_speculative_tokens ({max_committed_tokens}), " + f"got index={idx}, value={value!r}" + ) + committed_tokens_trace.append(committed) + return committed_tokens_trace + + @staticmethod + def _load_per_request_committed_tokens_trace( + *, + acceptance_trace_payload, + max_committed_tokens: int, + ) -> Optional[Dict[str, List[int]]]: + if acceptance_trace_payload is None or not isinstance( + acceptance_trace_payload, dict + ): + return None + if "per_request_committed_tokens_per_iteration" not in acceptance_trace_payload: + return None + + raw_trace_map = acceptance_trace_payload[ + "per_request_committed_tokens_per_iteration" + ] + if not isinstance(raw_trace_map, dict): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file " + "per_request_committed_tokens_per_iteration must be a dict, " + f"got={type(raw_trace_map).__name__}" + ) + if len(raw_trace_map) == 0: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file " + "per_request_committed_tokens_per_iteration must be non-empty." + ) + + per_request_trace: Dict[str, List[int]] = {} + for raw_request_id, raw_trace in raw_trace_map.items(): + request_id = str(raw_request_id) + if not request_id: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file per-request " + "trace keys must be non-empty strings." + ) + if request_id in per_request_trace: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file contains " + f"duplicate request_id={request_id!r} after normalization." + ) + if not isinstance(raw_trace, list): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file " + "per-request committed token trace must be a list, " + f"got request_id={request_id!r}, type={type(raw_trace).__name__}" + ) + if len(raw_trace) == 0: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file per-request " + f"committed token trace must be non-empty, request_id={request_id!r}" + ) + + validated_trace: List[int] = [] + for idx, value in enumerate(raw_trace): + committed = int(value) + if committed < 0: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file per-request " + "committed token values must be >= 0, " + f"got request_id={request_id!r}, index={idx}, value={value!r}" + ) + if committed > max_committed_tokens: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file per-request " + "committed token values must be <= 1 + num_speculative_tokens " + f"({max_committed_tokens}), got request_id={request_id!r}, " + f"index={idx}, value={value!r}" + ) + validated_trace.append(committed) + per_request_trace[request_id] = validated_trace + return per_request_trace + + @staticmethod + def _load_scheduled_draft_tokens_trace( + *, + acceptance_trace_payload, + max_scheduled_draft_tokens: int, + committed_trace_length: int, + ) -> Optional[List[int]]: + if acceptance_trace_payload is None or not isinstance(acceptance_trace_payload, dict): + return None + if "scheduled_draft_tokens_per_iteration" not in acceptance_trace_payload: + return None + + scheduled_draft_tokens_trace_raw = acceptance_trace_payload[ + "scheduled_draft_tokens_per_iteration" + ] + if not isinstance(scheduled_draft_tokens_trace_raw, list): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file scheduled draft token " + f"trace must be a list, got={type(scheduled_draft_tokens_trace_raw).__name__}" + ) + if len(scheduled_draft_tokens_trace_raw) != committed_trace_length: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file " + "scheduled_draft_tokens_per_iteration length must match " + "committed_tokens_per_iteration length." + ) + + scheduled_draft_tokens_trace: List[int] = [] + for idx, value in enumerate(scheduled_draft_tokens_trace_raw): + scheduled_drafts = int(value) + if scheduled_drafts < 0: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file scheduled draft " + "trace values must be >= 0, " + f"got index={idx}, value={value!r}" + ) + if scheduled_drafts > max_scheduled_draft_tokens: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file scheduled draft " + "trace values must be <= num_speculative_tokens " + f"({max_scheduled_draft_tokens}), got index={idx}, value={value!r}" + ) + scheduled_draft_tokens_trace.append(scheduled_drafts) + return scheduled_draft_tokens_trace + + @staticmethod + def _load_per_request_scheduled_draft_tokens_trace( + *, + acceptance_trace_payload, + max_scheduled_draft_tokens: int, + per_request_committed_trace: Optional[Dict[str, List[int]]], + ) -> Optional[Dict[str, List[int]]]: + if acceptance_trace_payload is None or not isinstance( + acceptance_trace_payload, dict + ): + return None + if ( + "per_request_scheduled_draft_tokens_per_iteration" + not in acceptance_trace_payload + ): + return None + if per_request_committed_trace is None: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file " + "per_request_scheduled_draft_tokens_per_iteration requires " + "per_request_committed_tokens_per_iteration." + ) + + raw_trace_map = acceptance_trace_payload[ + "per_request_scheduled_draft_tokens_per_iteration" + ] + if not isinstance(raw_trace_map, dict): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file " + "per_request_scheduled_draft_tokens_per_iteration must be a dict, " + f"got={type(raw_trace_map).__name__}" + ) + + normalized_keys = {str(request_id) for request_id in raw_trace_map.keys()} + committed_keys = set(per_request_committed_trace.keys()) + if normalized_keys != committed_keys: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file " + "per_request_scheduled_draft_tokens_per_iteration keys must match " + "per_request_committed_tokens_per_iteration keys." + ) + + per_request_trace: Dict[str, List[int]] = {} + for request_id, committed_trace in per_request_committed_trace.items(): + raw_trace = raw_trace_map[request_id] + if not isinstance(raw_trace, list): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file per-request " + "scheduled draft token trace must be a list, " + f"got request_id={request_id!r}, type={type(raw_trace).__name__}" + ) + if len(raw_trace) != len(committed_trace): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file " + "per_request_scheduled_draft_tokens_per_iteration length must " + "match per_request_committed_tokens_per_iteration length, " + f"request_id={request_id!r}" + ) + + validated_trace: List[int] = [] + for idx, value in enumerate(raw_trace): + scheduled_drafts = int(value) + if scheduled_drafts < 0: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file per-request " + "scheduled draft token values must be >= 0, " + f"got request_id={request_id!r}, index={idx}, value={value!r}" + ) + if scheduled_drafts > max_scheduled_draft_tokens: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file per-request " + "scheduled draft token values must be <= num_speculative_tokens " + f"({max_scheduled_draft_tokens}), got request_id={request_id!r}, " + f"index={idx}, value={value!r}" + ) + validated_trace.append(scheduled_drafts) + per_request_trace[request_id] = validated_trace + return per_request_trace + + def __post_init__(self) -> None: + supported_methods = { + "ngram", + "medusa", + "eagle", + "eagle3", + "deepseek_mtp", + "ernie_mtp", + "qwen3_moe_mtp", + "qwen3_next_mtp", + } + mtp_methods = { + "deepseek_mtp", + "ernie_mtp", + "qwen3_moe_mtp", + "qwen3_next_mtp", + } + if self.enabled and self.method not in supported_methods: + raise ValueError( + "SpeculativeDecodingConfig.method must match vLLM method names, " + f"got={self.method!r}, supported={sorted(supported_methods)}" + ) + if self.enabled and self.method in mtp_methods and self.mtp_n_predict <= 0: + raise ValueError( + "MTP methods require mtp_n_predict > 0 when enabled=True, " + f"got method={self.method!r}, mtp_n_predict={self.mtp_n_predict!r}" + ) + if self.enabled and self.method in mtp_methods and self.mtp_num_layers <= 0: + raise ValueError( + "MTP methods require mtp_num_layers > 0 when enabled=True, " + f"got method={self.method!r}, mtp_num_layers={self.mtp_num_layers!r}" + ) + if self.mtp_n_predict < 0: + raise ValueError( + "SpeculativeDecodingConfig.mtp_n_predict must be >= 0, " + f"got={self.mtp_n_predict!r}" + ) + if self.mtp_num_layers < 0: + raise ValueError( + "SpeculativeDecodingConfig.mtp_num_layers must be >= 0, " + f"got={self.mtp_num_layers!r}" + ) + if self.mtp_n_predict > 0 and self.method not in mtp_methods: + raise ValueError( + "SpeculativeDecodingConfig.mtp_n_predict is only valid for MTP " + f"methods, got method={self.method!r}" + ) + if self.mtp_num_layers > 0 and self.method not in mtp_methods: + raise ValueError( + "SpeculativeDecodingConfig.mtp_num_layers is only valid for MTP " + f"methods, got method={self.method!r}" + ) + if self.enabled and self.num_speculative_tokens <= 0: + raise ValueError( + "SpeculativeDecodingConfig.num_speculative_tokens must be > 0 when " + f"enabled=True, got={self.num_speculative_tokens}" + ) + if ( + self.method in mtp_methods + and self.mtp_n_predict > 0 + and self.num_speculative_tokens % self.mtp_n_predict != 0 + ): + raise ValueError( + "SpeculativeDecodingConfig.num_speculative_tokens must be divisible " + "by mtp_n_predict when mtp_n_predict > 0 for MTP methods, " + f"got num_speculative_tokens={self.num_speculative_tokens}, " + f"mtp_n_predict={self.mtp_n_predict}" + ) + max_committed_tokens = int(self.num_speculative_tokens) + 1 + if self.committed_tokens_per_iteration < 1: + raise ValueError( + "SpeculativeDecodingConfig.committed_tokens_per_iteration must be >= 1, " + f"got={self.committed_tokens_per_iteration!r}" + ) + if self.committed_tokens_per_iteration > max_committed_tokens: + raise ValueError( + "SpeculativeDecodingConfig.committed_tokens_per_iteration must be <= " + f"1 + num_speculative_tokens ({max_committed_tokens}), " + f"got={self.committed_tokens_per_iteration!r}" + ) + acceptance_trace_payload = self._load_acceptance_trace_payload( + acceptance_trace_file=self.acceptance_trace_file, + ) + self._committed_tokens_trace = self._load_committed_tokens_trace( + acceptance_trace_payload=acceptance_trace_payload, + max_committed_tokens=max_committed_tokens, + ) + self._per_request_committed_tokens_trace = ( + self._load_per_request_committed_tokens_trace( + acceptance_trace_payload=acceptance_trace_payload, + max_committed_tokens=max_committed_tokens, + ) + ) + if ( + acceptance_trace_payload is not None + and self._committed_tokens_trace is None + and self._per_request_committed_tokens_trace is None + ): + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file JSON object must " + "contain key 'committed_tokens_per_iteration' or " + "'per_request_committed_tokens_per_iteration'." + ) + self._scheduled_draft_tokens_trace = self._load_scheduled_draft_tokens_trace( + acceptance_trace_payload=acceptance_trace_payload, + max_scheduled_draft_tokens=int(self.num_speculative_tokens), + committed_trace_length=( + len(self._committed_tokens_trace) + if self._committed_tokens_trace is not None + else 0 + ), + ) + self._per_request_scheduled_draft_tokens_trace = ( + self._load_per_request_scheduled_draft_tokens_trace( + acceptance_trace_payload=acceptance_trace_payload, + max_scheduled_draft_tokens=int(self.num_speculative_tokens), + per_request_committed_trace=self._per_request_committed_tokens_trace, + ) + ) + if self._scheduled_draft_tokens_trace is not None: + for idx, (committed_tokens, scheduled_draft_tokens) in enumerate( + zip( + self._committed_tokens_trace, + self._scheduled_draft_tokens_trace, + ) + ): + if committed_tokens > 1 + scheduled_draft_tokens: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file committed " + "tokens must be <= 1 + scheduled_draft_tokens_per_iteration, " + f"got index={idx}, committed={committed_tokens}, " + f"scheduled_draft_tokens={scheduled_draft_tokens}" + ) + if self._per_request_scheduled_draft_tokens_trace is not None: + for request_id, committed_trace in ( + self._per_request_committed_tokens_trace.items() + ): + scheduled_trace = self._per_request_scheduled_draft_tokens_trace[ + request_id + ] + for idx, (committed_tokens, scheduled_draft_tokens) in enumerate( + zip(committed_trace, scheduled_trace) + ): + if committed_tokens > 1 + scheduled_draft_tokens: + raise ValueError( + "SpeculativeDecodingConfig.acceptance_trace_file per-request " + "committed tokens must be <= 1 + " + "per_request_scheduled_draft_tokens_per_iteration, " + f"got request_id={request_id!r}, index={idx}, " + f"committed={committed_tokens}, " + f"scheduled_draft_tokens={scheduled_draft_tokens}" + ) + + trace_payload = self._load_trace_calibration_payload(self.trace_calibration_file) + supported_trace_keys = { + "proposer_overhead_ms_by_method", + "metadata", + } + unexpected_keys = sorted(set(trace_payload.keys()) - supported_trace_keys) + if unexpected_keys: + raise ValueError( + "Unsupported keys in trace calibration file: " + f"{unexpected_keys}, supported={sorted(supported_trace_keys)}" + ) + + trace_proposer_overheads = self._validate_method_float_map( + map_name="proposer_overhead_ms_by_method", + raw_map=trace_payload.get("proposer_overhead_ms_by_method", {}), + supported_methods=supported_methods, + min_value=0.0, + inclusive_min=True, + ) + config_proposer_overheads = self._validate_method_float_map( + map_name="proposer_overhead_ms_by_method", + raw_map=self.proposer_overhead_ms_by_method, + supported_methods=supported_methods, + min_value=0.0, + inclusive_min=True, + ) + + # Config-driven values override trace-derived values for deterministic control. + self.proposer_overhead_ms_by_method = { + **trace_proposer_overheads, + **config_proposer_overheads, + } + self._decode_draft_proposer_latency_profile = ( + load_decode_draft_proposer_latency_profile( + profile_file=self.decode_draft_proposer_latency_profile_file, + supported_methods=supported_methods, + ) + ) diff --git a/frontier/execution_time_predictor/layer_contract_resolution.py b/frontier/execution_time_predictor/layer_contract_resolution.py new file mode 100644 index 00000000..8306f092 --- /dev/null +++ b/frontier/execution_time_predictor/layer_contract_resolution.py @@ -0,0 +1,495 @@ +"""Resolution of the typed layer contract a training or lookup call applies. + +A profiling row is admitted for a model only when its typed operator +contract matches the layer the caller asked about, at the tensor- and +expert-parallel sizes that layer actually uses. These methods derive those +keys and the contract signature that identifies a trained model. +""" + +import hashlib +import json +import pandas as pd + +from frontier.execution_time_predictor.attention_tp_policy import ( + resolve_effective_attention_tp_size, +) +from frontier.execution_time_predictor.prediction_model_identity import ( + _resolve_model_architecture_profile, + _resolve_profile_typed_family_for_query, + _serialize_selected_layer_cache_identity, +) +from frontier.model_architectures import LayerKind, ResolvedLayerContract +from frontier.moe_routing_runtime import ( + filter_moe_gating_routing_topk_rows, + resolve_moe_gating_routing_runtime_path, +) +from frontier.operators.binding import resolve_operator_query_tp_mode +from frontier.operators.families import ( + MOE_FAMILY, + get_operator_family, + is_moe_operator_ep_agnostic, + resolve_moe_operator_tp_key, +) +from frontier.operators.spec import TensorParallelMode +from frontier.operators.typed_contracts import TYPED_OPERATOR_CONTRACTS_COLUMN +from frontier.spec_decode.mtp_registry import ( + get_target_embedded_mtp_linear_ops, + is_target_embedded_mtp_same_tp_linear_op, +) +from frontier.spec_decode.runtime import is_target_embedded_mtp_enabled +from frontier.types import ClusterType +from typing import List, Optional, Tuple + + +class LayerContractResolution: + """Typed layer contract and parallel-size resolution.""" + + def _get_ffn_tp_key(self, cluster_type: ClusterType, replica_config, is_moe_model: bool) -> int: + if cluster_type == ClusterType.DECODE_FFN: + # In the FFN-only PD-AF cluster, dense FFN tensor parallelism is + # carried by moe_tensor_parallel_size as the FFN-domain TP field. + # attn_tensor_parallel_size can remain at its default because this + # cluster owns no attention weights. Use the FFN-domain TP for + # both dense and MoE DECODE_FFN profiling selection. + return replica_config.moe_tensor_parallel_size + if ( + is_moe_model + and cluster_type in { + ClusterType.PREFILL, + ClusterType.DECODE, + ClusterType.MONOLITHIC, + } + ): + return replica_config.moe_tensor_parallel_size + return replica_config.attn_tensor_parallel_size + + def _resolve_typed_layer_contract( + self, + op_name: str, + cluster_type: ClusterType, + replica_config, + *, + is_moe_model: bool, + layer_id: Optional[int] = None, + ) -> Optional[ResolvedLayerContract]: + """Resolve a typed FFN contract through the architecture profile.""" + + model_config = getattr(replica_config, "model_config", None) + architecture_profile = _resolve_model_architecture_profile(model_config) + if architecture_profile is None: + return None + + typed_family = _resolve_profile_typed_family_for_query( + architecture_profile, op_name + ) + if typed_family is None: + return None + typed_family_id, _ = typed_family + + # DECODE_ATTN is attention-only. Its exact zero domain is a deliberate + # sentinel; any non-zero value indicates a malformed configuration. + if cluster_type == ClusterType.DECODE_ATTN: + zero_fields = ( + "attn_tensor_parallel_size", + "moe_tensor_parallel_size", + "moe_expert_parallel_size", + ) + invalid = { + field_name: getattr(replica_config, field_name, None) + for field_name in zero_fields + if getattr(replica_config, field_name, None) != 0 + } + if invalid: + raise ValueError( + "DECODE_ATTN typed FFN resolution requires exact zero " + f"parallel sizes, got {invalid!r}" + ) + return None + + from frontier.operators.binding import bind_operator_query + + binding = bind_operator_query(op_name, family_id=typed_family_id) + if binding.family_id != typed_family_id: + raise ValueError( + f"Operator query {op_name!r} resolved to family " + f"{binding.family_id!r}, expected {typed_family_id!r}" + ) + + moe_tp_size = getattr(replica_config, "moe_tensor_parallel_size", None) + attention_tp_size = getattr( + replica_config, "attn_tensor_parallel_size", None + ) + if cluster_type == ClusterType.DECODE_FFN: + # The FFN-only role stores its domain size in the existing MoE TP + # field, while the profile still owns the semantic TP mode. + attention_tp_size = moe_tp_size + ffn_tp_size = self._get_ffn_tp_key( + cluster_type, replica_config, is_moe_model + ) + return architecture_profile.resolve_layer_contract( + model_config, + layer_id=layer_id, + operator_name=op_name, + attention_tp_size=attention_tp_size, + moe_tp_size=moe_tp_size, + ffn_tp_size=ffn_tp_size, + expert_parallel_size=getattr( + replica_config, "moe_expert_parallel_size", None + ), + ) + + def _resolve_ffn_layer_contracts( + self, + cluster_type: ClusterType, + replica_config, + is_moe_model: bool, + ) -> Tuple[Tuple[str, ResolvedLayerContract], ...]: + """Resolve each profile-owned FFN domain used by one training pass.""" + + if cluster_type == ClusterType.DECODE_ATTN: + zero_fields = ( + "attn_tensor_parallel_size", + "moe_tensor_parallel_size", + "moe_expert_parallel_size", + ) + invalid = { + field_name: getattr(replica_config, field_name, None) + for field_name in zero_fields + if getattr(replica_config, field_name, None) != 0 + } + if invalid: + raise ValueError( + "DECODE_ATTN FFN contract resolution requires exact zero " + f"parallel sizes, got {invalid!r}" + ) + return () + + model_config = getattr(replica_config, "model_config", None) + if model_config is None: + return () + architecture_profile = _resolve_model_architecture_profile(model_config) + if architecture_profile is None: + return () + if bool(is_moe_model) != bool(getattr(model_config, "is_moe", False)): + raise ValueError( + "is_moe_model does not match model configuration while resolving " + "typed FFN contracts" + ) + + contracts: list[Tuple[str, ResolvedLayerContract]] = [] + for spec in architecture_profile.iter_active_layer_contracts(model_config): + family_is_moe = spec.layer_kind is not LayerKind.DENSE + for family_id in spec.operator_family_ids: + family = get_operator_family(family_id) + profiling_ops = tuple(family.profiling_ops()) + if not profiling_ops: + raise ValueError( + f"Typed operator family {family_id!r} has no profiling operators" + ) + contract = self._resolve_typed_layer_contract( + profiling_ops[0].name, + cluster_type, + replica_config, + is_moe_model=family_is_moe, + ) + if contract is None: + raise ValueError( + f"Missing typed layer contract for operator family {family_id!r}" + ) + if contract.operator_family_id != family_id: + raise ValueError( + f"Operator family {family_id!r} resolved to " + f"{contract.operator_family_id!r}" + ) + contracts.append((family_id, contract)) + return tuple(contracts) + + def _get_ffn_contract_signature( + self, + cluster_type: ClusterType, + replica_config, + is_moe_model: bool, + ) -> str: + """Return a deterministic signature for the active FFN domains.""" + + entries = self._resolve_ffn_layer_contracts( + cluster_type, replica_config, is_moe_model + ) + if not entries: + return "none" + payload = [] + for family_id, contract in entries: + family = get_operator_family(family_id) + profiling_ops = tuple(family.profiling_ops()) + if not profiling_ops: + raise ValueError( + f"Typed operator family {family_id!r} has no profiling operators" + ) + + # The first operator is the compatibility representative returned + # by _resolve_ffn_layer_contracts(). Include every sibling as + # well: a family may mix EP-agnostic routing operators with an + # EP-sensitive grouped GEMM, and the cache signature must retain + # both semantics. + for operator in profiling_ops: + operator_contract = contract + if operator is not profiling_ops[0]: + operator_contract = self._resolve_typed_layer_contract( + operator.profiling_name(), + cluster_type, + replica_config, + is_moe_model=contract.layer_kind is not LayerKind.DENSE, + ) + if operator_contract is None: + raise ValueError( + "Missing typed layer contract for operator " + f"{operator.profiling_name()!r} in family {family_id!r}" + ) + payload.append( + { + "family_id": family_id, + "operator_name": operator.profiling_name(), + "identity": _serialize_selected_layer_cache_identity( + operator_contract + ), + } + ) + serialized = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest()[:16] + + @staticmethod + + def _is_mixed_layer_moe_model(model_config, is_moe_model: bool) -> bool: + """Return whether a model needs both MoE and dense FFN predictors. + + Some MoE architectures keep dense FFN layers at the model boundaries. + Their runtime dispatch is layer-specific, so model-level ``is_moe`` is + insufficient to decide which predictor families must be materialized. + Keep the legacy pure-MoE path unchanged when the layer-count contract + is unavailable. + """ + if not is_moe_model or model_config is None: + return False + get_num_moe_layers = getattr(model_config, "get_num_moe_layers", None) + num_layers = getattr(model_config, "num_layers", None) + if callable(get_num_moe_layers) and isinstance(num_layers, int): + return int(get_num_moe_layers()) < int(num_layers) + return False + + def _get_linear_op_tp_key(self, op_name: str, cluster_type: ClusterType, replica_config, is_moe_model: bool) -> int: + model_config = getattr(replica_config, "model_config", None) + # Lightweight configs retain the scalar FFN compatibility path, but + # generic linear attention names still require a profile declaration + # for TP-mode lookup. + architecture_profile = _resolve_model_architecture_profile( + model_config, + allow_generic=True, + ) + if op_name in get_target_embedded_mtp_linear_ops(): + return resolve_effective_attention_tp_size( + op_name="attn_pre_proj", + requested_tp_size=replica_config.attn_tensor_parallel_size, + num_kv_heads=replica_config.model_config.num_kv_heads, + cluster_type=cluster_type, + warning_cache=getattr(self, "_attention_tp_warning_cache", None), + include_linear_ops=True, + ) + + try: + tp_mode = resolve_operator_query_tp_mode( + op_name, + architecture_profile=architecture_profile, + ) + except (TypeError, ValueError) as exc: + raise ValueError(f"Unsupported linear op for TP mapping: {op_name}") from exc + + typed_contract = self._resolve_typed_layer_contract( + op_name, + cluster_type, + replica_config, + is_moe_model=is_moe_model, + ) + if ( + typed_contract is not None + and typed_contract.tensor_parallel_size is not None + ): + return typed_contract.tensor_parallel_size + + if tp_mode is TensorParallelMode.REPLICATED: + if ( + is_target_embedded_mtp_enabled( + getattr(replica_config, "speculative_decoding_config", None) + ) + and is_target_embedded_mtp_same_tp_linear_op(op_name) + ): + return resolve_effective_attention_tp_size( + op_name="attn_pre_proj", + requested_tp_size=replica_config.attn_tensor_parallel_size, + num_kv_heads=replica_config.model_config.num_kv_heads, + cluster_type=cluster_type, + warning_cache=getattr(self, "_attention_tp_warning_cache", None), + include_linear_ops=True, + ) + return 1 + + if tp_mode is TensorParallelMode.FFN_TP: + return self._get_ffn_tp_key(cluster_type, replica_config, is_moe_model) + + if tp_mode is TensorParallelMode.ATTENTION_TP: + return resolve_effective_attention_tp_size( + op_name=op_name, + requested_tp_size=replica_config.attn_tensor_parallel_size, + num_kv_heads=replica_config.model_config.num_kv_heads, + cluster_type=cluster_type, + warning_cache=getattr(self, "_attention_tp_warning_cache", None), + include_linear_ops=True, + ) + + raise ValueError(f"Unsupported linear op for TP mapping: {op_name}") + + @staticmethod + + def _get_moe_op_tp_key( + op_name: str, + replica_config, + cluster_type: ClusterType | None = None, + ) -> int: + try: + return resolve_moe_operator_tp_key( + op_name, + moe_tp_size=replica_config.moe_tensor_parallel_size, + cluster_type=cluster_type, + family=MOE_FAMILY, + ) + except ValueError as exc: + if str(exc).startswith("Unsupported MoE op:"): + raise ValueError( + f"Unsupported MoE op for TP mapping: {op_name}" + ) from exc + raise + + @staticmethod + + def _is_moe_op_ep_agnostic(op_name: str) -> bool: + try: + return is_moe_operator_ep_agnostic(op_name, family=MOE_FAMILY) + except ValueError as exc: + if str(exc).startswith("Unsupported MoE op:"): + raise ValueError( + f"Unsupported MoE op for EP mapping: {op_name}" + ) from exc + raise + + def _validate_moe_dataset_contract( + self, + file_path: str, + replica_config, + model_names: List[str], + cluster_type: ClusterType, + layer_contract: Optional[ResolvedLayerContract] = None, + ) -> None: + """Validate op-level MoE profiling key coverage before model training.""" + df = pd.read_csv(file_path) + required_columns = [ + "num_experts", + "router_topk", + "hidden_dim", + "expert_hidden_dim", + "num_tensor_parallel_workers", + "expert_parallel_size", + ] + missing_columns = [col for col in required_columns if col not in df.columns] + if missing_columns: + raise ValueError( + f"MoE dataset contract validation failed for {file_path}: " + f"missing required columns {missing_columns}." + ) + + model_config = replica_config.model_config + if layer_contract is None: + # A legacy caller has no profile-owned contract to validate. Keep + # the historical scalar admission rule, while refusing to guess + # when the dataset advertises typed metadata. + if TYPED_OPERATOR_CONTRACTS_COLUMN in df.columns: + raise ValueError( + "typed MoE profiling data requires an explicit routed layer contract" + ) + expected_expert_width = getattr(model_config, "mlp_hidden_dim", None) + if type(expected_expert_width) is not int or expected_expert_width <= 0: + raise ValueError( + "legacy MoE dataset validation requires a positive model_config.mlp_hidden_dim" + ) + else: + if layer_contract.layer_kind is not LayerKind.ROUTED: + raise ValueError( + "MoE dataset contract validation requires a routed layer contract" + ) + expected_expert_width = layer_contract.effective_ffn_width + base_df = df[ + (df["num_experts"] == model_config.num_experts) + & (df["router_topk"] == model_config.num_experts_per_tok) + & (df["hidden_dim"] == model_config.embedding_dim) + & (df["expert_hidden_dim"] == expected_expert_width) + ] + + if len(base_df) == 0: + raise ValueError( + "MoE dataset contract validation failed: no rows match model configuration in " + f"{file_path}. Required: num_experts={model_config.num_experts}, " + f"router_topk={model_config.num_experts_per_tok}, hidden_dim={model_config.embedding_dim}, " + f"expert_hidden_dim={expected_expert_width}." + ) + + available_pairs = sorted( + { + (int(tp), int(ep)) + for tp, ep in base_df[ + ["num_tensor_parallel_workers", "expert_parallel_size"] + ].drop_duplicates().itertuples(index=False, name=None) + } + ) + requested_routing_runtime_path = resolve_moe_gating_routing_runtime_path( + getattr(replica_config, "moe_routing_distribution_type", "balanced") + ) + + missing_requirements: List[str] = [] + for model_name in model_names: + tp_key = self._get_moe_op_tp_key( + model_name, + replica_config, + cluster_type, + ) + if self._is_moe_op_ep_agnostic(model_name): + op_df = base_df[ + base_df["num_tensor_parallel_workers"] == tp_key + ] + requirement = f"TP={tp_key}, EP=ANY" + else: + ep_key = replica_config.moe_expert_parallel_size + op_df = base_df[ + (base_df["num_tensor_parallel_workers"] == tp_key) + & (base_df["expert_parallel_size"] == ep_key) + ] + requirement = f"TP={tp_key}, EP={ep_key}" + if model_name == "moe_gating_routing_topk": + op_df = filter_moe_gating_routing_topk_rows( + op_df, + requested_runtime_path=requested_routing_runtime_path, + source_name=file_path, + ) + requirement = ( + f"{requirement}, routing_runtime_path=" + f"{requested_routing_runtime_path}" + ) + if len(op_df) == 0: + missing_requirements.append(f"{model_name} requires {requirement}") + + if missing_requirements: + requirement_text = "\n - ".join(missing_requirements) + raise ValueError( + "MoE dataset contract validation failed before training.\n" + f"File: {file_path}\n" + "Missing op-level key coverage:\n" + f" - {requirement_text}\n" + f"Available (TP, EP) pairs for matched model rows: {available_pairs}" + ) diff --git a/frontier/execution_time_predictor/moe_dataset_training.py b/frontier/execution_time_predictor/moe_dataset_training.py new file mode 100644 index 00000000..72f2f68c --- /dev/null +++ b/frontier/execution_time_predictor/moe_dataset_training.py @@ -0,0 +1,438 @@ +"""Admission and training of the MoE profiling dataset. + +The dataset contract decides which profiling rows a given replica may train +on, at its tensor- and expert-parallel sizes and its routing runtime. The +trainer then fits one model per MoE operator from the admitted rows. +""" + +import numpy as np +import os +import pandas as pd + +from frontier.execution_time_predictor.moe_predictor_helpers import ( + _get_moe_family_model_names, + _get_prefill_hot_moe_gating_model_names, + _is_moe_gating_family_model_name, + _validate_moe_columns, +) +from frontier.logger import init_logger +from frontier.moe_gating_runtime import ( + DEFAULT_MOE_GATING_RUNTIME_CONTEXT, + PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT, + PrefillHotRowsUnavailableError, + filter_moe_gating_rows_by_runtime_context, + get_moe_gating_base_model_name, + has_prefill_hot_moe_gating_rows, + should_enable_prefill_hot_moe_gating_contract, +) +from frontier.moe_routing_runtime import filter_moe_gating_routing_topk_rows +from frontier.operators.typed_contracts import ( + TYPED_OPERATOR_CONTRACTS_COLUMN, + validate_typed_operator_contracts, +) +from sklearn.base import BaseEstimator +from typing import Any, Dict, List, Optional + + +logger = init_logger(__name__) + + +class MoeDatasetTraining: + """MoE profiling dataset admission and per-operator training.""" + + + # Load imbalance feature columns used for MoE training (aligned with SharedPredictionModelManager) + # Reference: frontier/training/moe_trainer.py lines 224-239 (authoritative source) + MOE_LOAD_IMBALANCE_FEATURES = [ + # Config features (6) - describe model configuration + "total_routed_tokens", # Total tokens after routing (num_tokens * router_topk) + "num_experts_per_device", # Number of experts per device after EP sharding + "hidden_dim", # Model hidden dimension + "expert_hidden_dim", # Expert FFN hidden dimension + "router_topk", # Number of experts each token is routed to + "model_expansion_ratio", # expert_hidden_dim / hidden_dim + # Derived features (2) - derived from config and routing + "tokens_per_expert_avg", # Average tokens per expert + "tokens_to_experts_ratio", # tokens / num_experts ratio + # Load features (6) - describe load distribution characteristics + "expert_utilization", # Proportion of experts with non-zero load + "min_load_ratio", # Min load / average load + "load_imbalance_cv", # Coefficient of Variation: std/mean, key imbalance metric + "max_load_ratio", # Max load / average load + "load_entropy", # Entropy of load distribution (higher = more uniform) + "load_gini_coefficient", # Gini coefficient: 0=equality, 1=inequality + ] + + def _validate_moe_dataset_contract( + self, + moe_df: pd.DataFrame, + moe_input_file: str, + model_names: List[str], + moe_tp_size: int, + moe_ep_size: int, + ) -> pd.DataFrame: + """Validate op-level MoE key coverage and return model-filtered dataframe.""" + _validate_moe_columns(moe_df) + required_columns = [ + "num_experts", + "router_topk", + "hidden_dim", + "expert_hidden_dim", + "num_tensor_parallel_workers", + "expert_parallel_size", + ] + missing_columns = [col for col in required_columns if col not in moe_df.columns] + if missing_columns: + raise ValueError( + f"MoE dataset contract validation failed for {moe_input_file}: " + f"missing required columns {missing_columns}." + ) + + model_config = self._model_config + base_df = moe_df[ + (moe_df["num_experts"] == model_config.num_experts) + & (moe_df["router_topk"] == model_config.num_experts_per_tok) + & (moe_df["hidden_dim"] == model_config.embedding_dim) + & (moe_df["expert_hidden_dim"] == model_config.mlp_hidden_dim) + ].copy() + + if len(base_df) == 0: + raise ValueError( + "MoE dataset contract validation failed: no rows match model configuration in " + f"{moe_input_file}. Required: num_experts={model_config.num_experts}, " + f"router_topk={model_config.num_experts_per_tok}, hidden_dim={model_config.embedding_dim}, " + f"expert_hidden_dim={model_config.mlp_hidden_dim}." + ) + + available_pairs = sorted( + { + (int(tp), int(ep)) + for tp, ep in base_df[ + ["num_tensor_parallel_workers", "expert_parallel_size"] + ].drop_duplicates().itertuples(index=False, name=None) + } + ) + requested_routing_runtime_path = ( + self._get_requested_moe_gating_routing_runtime_path() + ) + + missing_requirements: List[str] = [] + for model_name in model_names: + base_model_name = get_moe_gating_base_model_name(model_name) + tp_key = self._get_moe_op_tp_key( + base_model_name, + moe_tp_size, + cluster_type=getattr(self, "_cluster_type", None), + ) + requirement_parts = [f"TP={tp_key}"] + if self._is_moe_op_ep_agnostic(base_model_name): + op_df = base_df[base_df["num_tensor_parallel_workers"] == tp_key] + requirement_parts.append("EP=ANY") + else: + op_df = base_df[ + (base_df["num_tensor_parallel_workers"] == tp_key) + & (base_df["expert_parallel_size"] == moe_ep_size) + ] + requirement_parts.append(f"EP={moe_ep_size}") + if base_model_name == "moe_gating_routing_topk": + op_df = filter_moe_gating_routing_topk_rows( + op_df, + requested_runtime_path=requested_routing_runtime_path, + source_name=moe_input_file, + ) + requirement_parts.append( + f"routing_runtime_path={requested_routing_runtime_path}" + ) + if _is_moe_gating_family_model_name(base_model_name): + op_df = filter_moe_gating_rows_by_runtime_context( + op_df, + requested_context=DEFAULT_MOE_GATING_RUNTIME_CONTEXT, + source_name=moe_input_file, + ) + requirement_parts.append( + "gating_runtime_context=" + f"{DEFAULT_MOE_GATING_RUNTIME_CONTEXT}" + ) + requirement = ", ".join(requirement_parts) + if len(op_df) == 0: + missing_requirements.append(f"{model_name} requires {requirement}") + continue + target_col = f"time_stats.{base_model_name}.median" + if op_df[target_col].dropna().empty: + missing_requirements.append( + f"{model_name} requires {requirement}, target={target_col} " + "to contain at least one non-NaN row" + ) + + if missing_requirements: + requirement_text = "\n - ".join(missing_requirements) + raise ValueError( + "MoE dataset contract validation failed before training.\n" + f"File: {moe_input_file}\n" + "Missing op-level key coverage:\n" + f" - {requirement_text}\n" + f"Available (TP, EP) pairs for matched model rows: {available_pairs}" + ) + + return base_df + + def _train_moe_models(self) -> Dict[str, BaseEstimator]: + """Train MoE-specific models (gating, shuffling, grouped_gemm) for independent training mode. + + For moe_grouped_gemm, uses 14 load-imbalance features if available in the profiling data. + This enables simulation mode with per-expert token allocation. + Other MoE models (gating_linear, gating_routing_topk, shuffling) use only num_tokens. + """ + models = {} + moe_input_file = getattr(self, "_moe_input_file", "/synthetic/moe.csv") + + if not os.path.exists(moe_input_file): + logger.warning(f"MoE input file does not exist: {moe_input_file}") + return models + + moe_df = pd.read_csv(moe_input_file) + if TYPED_OPERATOR_CONTRACTS_COLUMN in moe_df.columns: + # Validate every row before scalar or TP/EP filtering can hide a + # malformed typed contract. + moe_df[TYPED_OPERATOR_CONTRACTS_COLUMN].map( + lambda raw_contracts: validate_typed_operator_contracts( + raw_contracts, + model_config=self._model_config, + ) + ) + + metadata = self._get_profiling_metadata(moe_df, moe_input_file) + self._validate_active_measurement_type(metadata, moe_input_file) + + tp_col = "num_tensor_parallel_workers" + ep_col = "expert_parallel_size" + moe_tp_size = self._replica_config.moe_tensor_parallel_size + moe_ep_size = self._replica_config.moe_expert_parallel_size + + if tp_col not in moe_df.columns: + raise ValueError( + f"Required column '{tp_col}' is missing in {moe_input_file}. " + "Re-run MoE profiling with TP metadata enabled." + ) + if ep_col not in moe_df.columns: + raise ValueError( + f"Required column '{ep_col}' is missing in {moe_input_file}. " + "Re-run MoE profiling with EP metadata enabled." + ) + + base_model_names = _get_moe_family_model_names() + model_names = list(base_model_names) + model_filtered_df = self._validate_moe_dataset_contract( + moe_df, + moe_input_file, + base_model_names, + moe_tp_size, + moe_ep_size, + ) + if should_enable_prefill_hot_moe_gating_contract( + model_config=self._model_config, + ): + if has_prefill_hot_moe_gating_rows(model_filtered_df): + model_names.extend(_get_prefill_hot_moe_gating_model_names()) + else: + logger.warning( + "Prefill-hot gating contract is enabled for model=%s, but " + "dataset %s has no usable prefill_hot rows; skipping " + "__prefill_hot pseudo-model training.", + self._replica_config.model_name, + moe_input_file, + ) + + self._register_profiling_metadata_for_ops( + model_names, metadata, moe_input_file + ) + + requested_routing_runtime_path = ( + self._get_requested_moe_gating_routing_runtime_path() + ) + moe_df_cache: Dict[ + tuple[int, Optional[int], Optional[str], Optional[str]], pd.DataFrame + ] = {} + + def _get_moe_df_for_op( + model_name: str, + ) -> tuple[pd.DataFrame, int, Optional[int]]: + base_model_name = get_moe_gating_base_model_name(model_name) + tp_key = self._get_moe_op_tp_key( + base_model_name, + moe_tp_size, + cluster_type=getattr(self, "_cluster_type", None), + ) + ep_key: Optional[int] + if self._is_moe_op_ep_agnostic(base_model_name): + ep_key = None + else: + ep_key = moe_ep_size + runtime_path_key: Optional[str] = None + if base_model_name == "moe_gating_routing_topk": + runtime_path_key = requested_routing_runtime_path + gating_context_key: Optional[str] = None + if _is_moe_gating_family_model_name(base_model_name): + gating_context_key = DEFAULT_MOE_GATING_RUNTIME_CONTEXT + if model_name.endswith("__prefill_hot"): + gating_context_key = PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT + cache_key = (tp_key, ep_key, runtime_path_key, gating_context_key) + if cache_key not in moe_df_cache: + filtered_df = model_filtered_df[ + model_filtered_df[tp_col] == tp_key + ].copy() + if ep_key is not None: + filtered_df = filtered_df[ + filtered_df[ep_col] == ep_key + ].copy() + if runtime_path_key is not None: + filtered_df = filter_moe_gating_routing_topk_rows( + filtered_df, + requested_runtime_path=runtime_path_key, + source_name=moe_input_file, + ) + if gating_context_key is not None: + filtered_df = filter_moe_gating_rows_by_runtime_context( + filtered_df, + requested_context=gating_context_key, + source_name=moe_input_file, + ) + if len(filtered_df) == 0: + ep_desc = "ANY" if ep_key is None else str(ep_key) + raise ValueError( + f"No MoE data after filtering for TP={tp_key}, EP={ep_desc}. " + f"Requested by op-level TP mapping in {moe_input_file}." + ) + filtered_df["num_tokens_rounded"] = filtered_df["num_tokens"].apply( + lambda x: max(1, round(x / 8) * 8) + ) + moe_df_cache[cache_key] = filtered_df + return moe_df_cache[cache_key], tp_key, ep_key + + for model_name in model_names: + try: + op_df, moe_tp_key, moe_ep_key = _get_moe_df_for_op(model_name) + except PrefillHotRowsUnavailableError as exc: + logger.warning( + "Skipping %s because prefill-hot gating rows are unavailable " + "for the requested TP/EP slice (%s).", + model_name, + exc, + ) + continue + target_op_name = get_moe_gating_base_model_name(model_name) + target_col = f"time_stats.{target_op_name}.median" + if target_col not in op_df.columns: + ep_desc = "ANY" if moe_ep_key is None else str(moe_ep_key) + raise ValueError( + f"Column '{target_col}' not found in MoE dataframe for TP={moe_tp_key}, EP={ep_desc}. " + "Re-run MoE profiling with split gating columns." + ) + + # Per-operation feature selection (aligned with SharedPredictionModelManager). + if model_name == "moe_grouped_gemm": + available_load_features = [ + f for f in self.MOE_LOAD_IMBALANCE_FEATURES if f in op_df.columns + ] + has_load_imbalance_features = len(available_load_features) == len( + self.MOE_LOAD_IMBALANCE_FEATURES + ) + if 0 < len(available_load_features) < len(self.MOE_LOAD_IMBALANCE_FEATURES): + missing_features = [ + f + for f in self.MOE_LOAD_IMBALANCE_FEATURES + if f not in op_df.columns + ] + raise ValueError( + f"Partial load imbalance features found ({len(available_load_features)}/" + f"{len(self.MOE_LOAD_IMBALANCE_FEATURES)}) for TP={moe_tp_key}. " + f"Missing: {missing_features}." + ) + if has_load_imbalance_features: + feature_cols = available_load_features + logger.info( + f" {model_name}: Using load imbalance features ({len(feature_cols)} features, TP={moe_tp_key})" + ) + else: + feature_cols = ["num_tokens"] + logger.info( + f" {model_name}: Load imbalance features not found; using num_tokens only (TP={moe_tp_key})" + ) + elif model_name == "moe_shuffling": + available_load_features = [ + f for f in self.MOE_LOAD_IMBALANCE_FEATURES if f in op_df.columns + ] + if len(available_load_features) == len(self.MOE_LOAD_IMBALANCE_FEATURES): + feature_cols = available_load_features + logger.info( + f" {model_name}: Using load imbalance features ({len(feature_cols)} features, TP={moe_tp_key})" + ) + else: + feature_cols = ["num_tokens"] + logger.info( + f" {model_name}: Full load imbalance features unavailable; using num_tokens only (TP={moe_tp_key})" + ) + else: + feature_cols = ["num_tokens"] + logger.info(f" {model_name}: Using num_tokens only (1 feature, TP={moe_tp_key})") + + models[model_name] = self._train_model( + model_name=model_name, + df=op_df, + feature_cols=feature_cols, + target_col=target_col, + ) + logger.info(f"Trained MoE model: {model_name}") + + return models + + def _register_additional_profiling_metadata_from_files(self) -> None: + moe_input_file = self._moe_input_file + model_names = _get_moe_family_model_names() + if should_enable_prefill_hot_moe_gating_contract( + model_config=self._model_config, + ): + include_prefill_hot_models = False + try: + moe_df = pd.read_csv(moe_input_file) + except FileNotFoundError: + moe_df = None + if moe_df is not None: + if TYPED_OPERATOR_CONTRACTS_COLUMN in moe_df.columns: + moe_df[TYPED_OPERATOR_CONTRACTS_COLUMN].map( + lambda raw_contracts: validate_typed_operator_contracts( + raw_contracts, + model_config=self._model_config, + ) + ) + include_prefill_hot_models = has_prefill_hot_moe_gating_rows(moe_df) + if include_prefill_hot_models: + model_names.extend(_get_prefill_hot_moe_gating_model_names()) + self._register_profiling_metadata_from_file(moe_input_file, model_names) + + def _train_models(self) -> Dict[str, BaseEstimator]: + """Override to include MoE model training for independent training mode.""" + models = super()._train_models() + + if self._model_manager is None: + moe_models = self._train_moe_models() + models.update(moe_models) + logger.info(f"Trained MoE models independently: {list(moe_models.keys())}") + else: + logger.info("MoE models loaded from ExecutionTimePredictionModelManager.") + + return models + + def _predict_for_compute_models(self) -> Dict[str, Any]: + predictions = super()._predict_for_compute_models() + extra_model_names = _get_prefill_hot_moe_gating_model_names() + num_token_range = np.arange(1, self._max_tokens + 1) + X = pd.DataFrame({"num_tokens": num_token_range}) + for model_name in extra_model_names: + if model_name not in self._models: + continue + model = self._models[model_name] + predictions[model_name] = self._get_model_prediction( + model_name, model, X + ) + return predictions diff --git a/frontier/execution_time_predictor/moe_mtp_replay.py b/frontier/execution_time_predictor/moe_mtp_replay.py new file mode 100644 index 00000000..02c944cb --- /dev/null +++ b/frontier/execution_time_predictor/moe_mtp_replay.py @@ -0,0 +1,216 @@ +"""MoE timing for speculative-decoding MTP replay rows. + +An MTP iteration replays decoder layers for the draft tokens it proposed, so +its MoE time is an aggregate over lanes rather than a single layer lookup. +""" + +import math + +from frontier.entities import Batch, ExecutionTime +from frontier.types import ClusterType + + +class MoeMtpReplay: + """MoE time for MTP replay rows and terminal overshoot.""" + + def _predict_mtp_moe_lane_phase_aggregate( + self, + *, + predictor, + batch: Batch, + pipeline_stage: int, + cluster_type: ClusterType, + layer_id: int, + num_layers: int, + ) -> tuple[ExecutionTime, tuple[float, float, float, float, float]]: + """Return one shared attention result and the five lane barriers. + + ``predictor`` is explicit because structural MTP may run against a + secondary predictor owned by this parent. The attention probe is kept + at one layer: pipeline and CPU overhead are batch-level terms, while + the returned physical phase barriers are the only values scaled by + ``num_layers`` at the caller. + """ + + if type(num_layers) is not int or num_layers < 1: + raise ValueError(f"num_layers must be a positive integer, got {num_layers!r}") + + attention_execution_time = predictor.predict_stage_execution_time( + batch=batch, + stage_id=pipeline_stage, + cluster_type=cluster_type, + num_layers=1, + layer_id=layer_id, + include_ffn=False, + ) + attention_time_ms = float(attention_execution_time.model_time_ms) + if not math.isfinite(attention_time_ms) or attention_time_ms < 0: + raise ValueError( + "MTP structural attention time must be finite and non-negative, " + f"got {attention_time_ms}" + ) + + workload = predictor._materialize_layer_ep_workload( + batch=batch, + cluster_type=cluster_type, + layer_id=layer_id, + ) + participant_ep_ids = tuple(workload.participant_ep_ids) + if not participant_ep_ids: + raise ValueError("MTP MoE replay produced no EP participants") + + effective_tokens = int( + batch.get_effective_total_tokens_for_compute(cluster_type) + ) + if effective_tokens <= 0: + raise ValueError( + "MTP MoE replay requires positive pre-routing effective tokens, " + f"got {effective_tokens}" + ) + + phase_values: list[list[float]] = [] + for ep_id in participant_ep_ids: + lane_workload = workload.lane(int(ep_id)) + lane_phases = predictor.predict_moe_lane_phase_times( + batch=batch, + lane_workload=lane_workload, + pipeline_stage=pipeline_stage, + cluster_type=cluster_type, + ) + if len(lane_phases) != 5: + raise ValueError( + "MTP MoE lane phase API must return five values, " + f"got ep_id={ep_id}, values={lane_phases!r}" + ) + normalized_phases = [float(value) for value in lane_phases] + if any( + not math.isfinite(value) or value < 0 + for value in normalized_phases + ): + raise ValueError( + "MTP MoE lane phase times must be finite and non-negative, " + f"got ep_id={ep_id}, values={normalized_phases}" + ) + phase_values.append(normalized_phases) + + phase_maxima = tuple( + max(values[index] for values in phase_values) for index in range(5) + ) + return attention_execution_time, phase_maxima + + def _predict_mtp_terminal_row_time_ms( + self, + *, + batch: Batch, + stage_id: int, + cluster_type: ClusterType, + num_layers: int, + layer_id: int, + ) -> float: + """Predict a terminal MTP row with physical EP barriers when required.""" + + model_config = getattr(self, "_model_config", None) + if model_config is None or not bool(getattr(model_config, "is_moe", False)): + return super()._predict_mtp_terminal_row_time_ms( + batch=batch, + stage_id=stage_id, + cluster_type=cluster_type, + num_layers=num_layers, + layer_id=layer_id, + ) + is_moe_layer = getattr(model_config, "is_moe_layer", None) + if not callable(is_moe_layer): + raise ValueError( + "MTP terminal MoE prediction requires model_config.is_moe_layer" + ) + if not bool(is_moe_layer(layer_id)): + return super()._predict_mtp_terminal_row_time_ms( + batch=batch, + stage_id=stage_id, + cluster_type=cluster_type, + num_layers=num_layers, + layer_id=layer_id, + ) + if cluster_type not in (ClusterType.MONOLITHIC, ClusterType.DECODE): + return super()._predict_mtp_terminal_row_time_ms( + batch=batch, + stage_id=stage_id, + cluster_type=cluster_type, + num_layers=num_layers, + layer_id=layer_id, + ) + if int(getattr(self, "_moe_ep_size", 1)) <= 1: + return super()._predict_mtp_terminal_row_time_ms( + batch=batch, + stage_id=stage_id, + cluster_type=cluster_type, + num_layers=num_layers, + layer_id=layer_id, + ) + + attention_execution_time, phase_maxima = ( + self._predict_mtp_moe_lane_phase_aggregate( + predictor=self, + batch=batch, + pipeline_stage=stage_id, + cluster_type=cluster_type, + layer_id=layer_id, + num_layers=num_layers, + ) + ) + attention_time_ms = float(attention_execution_time.total_time * 1e3) + if not math.isfinite(attention_time_ms) or attention_time_ms < 0: + raise ValueError( + "MTP terminal attention time must be finite and non-negative, " + f"got {attention_time_ms}" + ) + return attention_time_ms + sum(phase_maxima) * int(num_layers) + + def _predict_mtp_decoder_layer_time_ms( + self, + *, + predictor, + batch: Batch, + ) -> float: + layer_id = 0 + model_config = getattr(predictor, "_model_config", None) + if model_config is None: + raise ValueError( + "MTP structural decoder prediction requires model_config" + ) + if not bool(getattr(model_config, "is_moe", False)): + return super()._predict_mtp_decoder_layer_time_ms( + predictor=predictor, + batch=batch, + ) + + is_moe_layer = getattr(model_config, "is_moe_layer", None) + if not callable(is_moe_layer): + raise ValueError( + "MTP structural MoE decoder prediction requires " + "model_config.is_moe_layer" + ) + if not bool(is_moe_layer(layer_id)): + return super()._predict_mtp_decoder_layer_time_ms( + predictor=predictor, + batch=batch, + ) + + cluster_type = getattr(predictor, "_cluster_type", None) + if not isinstance(cluster_type, ClusterType): + raise ValueError( + "MTP structural MoE decoder prediction requires a valid cluster_type" + ) + + attention_execution_time, phase_maxima = ( + self._predict_mtp_moe_lane_phase_aggregate( + predictor=predictor, + batch=batch, + pipeline_stage=0, + cluster_type=cluster_type, + layer_id=layer_id, + num_layers=1, + ) + ) + attention_time_ms = float(attention_execution_time.model_time_ms) + return attention_time_ms + sum(phase_maxima) diff --git a/frontier/execution_time_predictor/moe_operator_times.py b/frontier/execution_time_predictor/moe_operator_times.py new file mode 100644 index 00000000..d13a2914 --- /dev/null +++ b/frontier/execution_time_predictor/moe_operator_times.py @@ -0,0 +1,713 @@ +"""Predicted time of each MoE operator in a layer. + +Gating, routing top-k, shuffling, the grouped expert GEMM and the +expert-parallel collective, plus the token-count resolution each of them +needs and the model selection that decides which trained estimator answers. +""" + +from frontier.config import get_quantization_manager +from frontier.entities import Batch, ExecutionTime +from frontier.entities.time_components import MoETime +from frontier.execution_time_predictor.moe_predictor_helpers import ( + _MOE_GATING_OPERATOR_NAMES, +) +from frontier.logger import init_logger +from frontier.moe_ep_workload import EPLaneWorkload, resolve_ep_lane_workload +from frontier.moe_gating_runtime import ( + DEFAULT_MOE_GATING_RUNTIME_CONTEXT, + PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT, + get_moe_gating_prediction_model_name, + should_use_prefill_hot_moe_gating_context, +) +from frontier.operators.families import ( + MOE_FAMILY, + get_comm_operator, + is_moe_operator_ep_agnostic, + resolve_moe_operator_tp_key, +) +from frontier.types import ClusterType +from typing import Dict, Mapping, Optional + + +logger = init_logger(__name__) + + +class MoeOperatorTimes: + """Per-operator MoE time prediction and its inputs.""" + + @staticmethod + + def _get_dummy_shared_domain_moe_scope_time( + execution_time: ExecutionTime, + ) -> float: + """Return the fixed per-operator MoE scope used by shared-domain decode. + + The generic dummy ``ExecutionTime`` keeps the deprecated aggregate + ``moe_gating_time`` contract by splitting that baseline across the two + structured gating fields. The shared-domain decode contract models + each named gating operator as one fixed structural slot, matching its + historical five-operator scope. Resolve that compatibility at this + boundary from the MoE family registry; all other operators retain the + descriptor-aware structured timing, including zero-lane routed work. + """ + + moe_time = execution_time.moe_or_mlp_time_component + if not isinstance(moe_time, MoETime): + raise ValueError( + "shared-domain dummy timing requires a MoE execution component" + ) + operator_times = moe_time.operator_times + if operator_times is None: + raise ValueError( + "shared-domain dummy timing requires structured MoE operator times" + ) + + scope_time = 0.0 + for operator_name, operator_time in operator_times.op_times.items(): + if operator_name in _MOE_GATING_OPERATOR_NAMES: + # ``moe_gating_time`` is the one fixed dummy baseline for each + # named gating operator; the structured fields store its + # compatibility split as 0.5 * baseline each. + scope_time += float(moe_time.moe_gating_time) + else: + scope_time += float(operator_time) + return scope_time + + def predict_monolithic_decode_shared_domain_lane_moe_times_ms( + self, + batch: Batch, + layer_id: int, + ) -> Dict[int, float]: + """Estimate per-EP-lane pre-collective MoE time for monolithic pure decode. + + Returns per-lane post-attention MoE compute in milliseconds. The result is + used by the MONOLITHIC decode sync path to model shared-domain readiness skew + before `expert_parallel_allreduce`. + """ + if self._enable_dummy_mode: + lane_workloads = self._resolve_shared_domain_lane_workloads( + batch, + cluster_type=ClusterType.MONOLITHIC, + layer_id=layer_id, + ) + lane_times_ms: Dict[int, float] = {} + for lane_workload in lane_workloads: + # This helper returns only the per-layer MoE scope. The dummy + # execution seam needs a stage value for its complete object, + # but no stage-boundary term is included in the component total. + execution_time = self._get_dummy_execution_time( + batch, + pipeline_stage=0, + include_attention=False, + lane_workload=lane_workload, + ) + lane_times_ms[lane_workload.ep_id] = ( + self._get_dummy_shared_domain_moe_scope_time(execution_time) + ) + return lane_times_ms + + lane_workloads = self._resolve_shared_domain_lane_workloads( + batch, + cluster_type=ClusterType.MONOLITHIC, + layer_id=layer_id, + ) + + post_attention_layernorm_time = self._get_mlp_norm_layer_act_execution_time(batch) + gating_linear_time = self._get_gating_linear_time(batch) + gating_routing_topk_time = self._get_gating_routing_topk_time(batch) + share_expert_total_time = 0.0 + if self._model_config.supports_share_expert(): + share_expert_total_time = ( + self._get_share_expert_up_proj_execution_time(batch) + + self._get_share_expert_down_proj_execution_time(batch) + + self._get_share_expert_act_execution_time(batch) + ) + + lane_times_ms: Dict[int, float] = {} + for lane_workload in lane_workloads: + lane_id = lane_workload.ep_id + shuffling_time = self._get_moe_shuffling_time( + batch, + moe_tokens_input=lane_workload, + ) + grouped_gemm_time = self._get_grouped_gemm_time( + lane_workload, + batch=batch, + ) + + lane_times_ms[lane_id] = ( + post_attention_layernorm_time + + gating_linear_time + + gating_routing_topk_time + + shuffling_time + + grouped_gemm_time + + share_expert_total_time + ) + + return lane_times_ms + + @staticmethod + + def _get_moe_op_tp_key( + op_name: str, + moe_tp_size: int, + cluster_type: ClusterType | None = None, + ) -> int: + try: + return resolve_moe_operator_tp_key( + op_name, + moe_tp_size=moe_tp_size, + cluster_type=cluster_type, + family=MOE_FAMILY, + ) + except ValueError as exc: + if str(exc).startswith("Unsupported MoE op:"): + raise ValueError( + f"Unsupported MoE op for TP mapping: {op_name}" + ) from exc + raise + + @staticmethod + + def _is_moe_op_ep_agnostic(op_name: str) -> bool: + try: + return is_moe_operator_ep_agnostic(op_name, family=MOE_FAMILY) + except ValueError as exc: + if str(exc).startswith("Unsupported MoE op:"): + raise ValueError( + f"Unsupported MoE op for EP mapping: {op_name}" + ) from exc + raise + + def _select_moe_gating_prediction_model_name( + self, + base_model_name: str, + batch: Batch, + ) -> str: + requested_context = DEFAULT_MOE_GATING_RUNTIME_CONTEXT + if should_use_prefill_hot_moe_gating_context( + model_config=self._model_config, + batch=batch, + ): + requested_context = PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT + candidate_model_name = get_moe_gating_prediction_model_name( + base_model_name, + requested_context=requested_context, + ) + if candidate_model_name in self._predictions: + return candidate_model_name + return base_model_name + + def _use_expert_parallel_alltoall_path(self, batch: Batch) -> bool: + moe_ep_size = int(getattr(self, "_moe_ep_size", 1)) + if moe_ep_size <= 1: + return False + # EP is replica-local and is independent of the retired attention-DP + # lane concept. A full batch on any MoE serving role therefore uses + # the EP communication/accounting path whenever EP>1. + return True + + def _predict_expert_parallel_phase_operator_times( + self, + batch: Batch, + *, + lane_workload: Optional[EPLaneWorkload] = None, + ) -> dict[str, float]: + """Predict exact dispatch and combine collectives for one MoE layer.""" + + if self._moe_ep_size <= 1: + return { + "expert_parallel_alltoall_dispatch": 0.0, + "expert_parallel_alltoall_combine": 0.0, + } + if not self._use_expert_parallel_alltoall_path(batch): + raise ValueError( + "Canonical MoE EP execution requires named all-to-all dispatch " + "and combine phases" + ) + return { + op_name: self._predict_comm_operator( + get_comm_operator(op_name), + batch, + lane_workload=lane_workload, + ) + for op_name in ( + "expert_parallel_alltoall_dispatch", + "expert_parallel_alltoall_combine", + ) + } + + def _get_effective_moe_total_tokens(self, batch: Batch) -> int: + effective_tokens = int( + batch.get_effective_total_tokens_rounded(self._cluster_type) + ) + if effective_tokens < 0: + raise ValueError( + f"effective MoE tokens must be non-negative, got {effective_tokens}" + ) + return effective_tokens + + def _get_moe_pre_routing_token_count(self, batch: Optional[Batch]) -> int: + """Return the source-batch width used by pre-routing MoE models. + + A physical EP lane carries only an assignment subset, so its routed + count cannot identify the source width. Callers that need the + one-feature profiling domain must provide the source batch explicitly. + """ + + if batch is None: + raise ValueError( + "MoE pre-routing token lookup requires the source batch; " + "an EPLaneWorkload cannot supply that width" + ) + return self._get_effective_moe_total_tokens(batch) + + def _get_local_ep_routed_tokens( + self, + batch: Batch, + *, + lane_workload: Optional[EPLaneWorkload] = None, + ) -> int: + source = batch if lane_workload is None else lane_workload + resolved_lane_workload = resolve_ep_lane_workload(source, required=True) + assert resolved_lane_workload is not None + return resolved_lane_workload.routed_token_count + + def _get_moe_tokens_input( + self, batch: Batch, layer_id: int = 0 + ) -> EPLaneWorkload | int: + """ + Unified entry point to get MoE tokens input for grouped GEMM prediction. + + EP lane batches carry the canonical physical descriptor. A regular + non-lane batch may use the scalar pre-routing token path for legacy + one-feature models; load-aware models require an explicit descriptor. + + Args: + batch: The batch being processed + layer_id: The layer ID for which to get token allocation (default 0) + + Returns: + - In load-imbalance mode: ``EPLaneWorkload`` + - In single-token-count profiling mode: pre-routing token count + + Raises: + ValueError: If the selected routing mode is not supported by the active predictor + """ + lane_workload = resolve_ep_lane_workload(batch, required=False) + if lane_workload is not None: + if lane_workload.router_topk != int(self._router_topk): + raise ValueError( + "EPLaneWorkload router_topk does not match predictor topology: " + f"descriptor={lane_workload.router_topk}, predictor={self._router_topk}" + ) + return lane_workload + + load_aware = any( + isinstance(prediction, dict) + and prediction.get("_on_demand_prediction", False) + for prediction in ( + getattr(self, "_predictions", {}).get("moe_shuffling"), + getattr(self, "_predictions", {}).get("moe_grouped_gemm"), + ) + ) + if load_aware: + cluster_type = getattr(self, "_cluster_type", None) + if not isinstance(cluster_type, ClusterType): + raise ValueError( + "load-aware MoE prediction requires an initialized cluster_type" + ) + workload = self._materialize_layer_ep_workload( + batch=batch, + cluster_type=cluster_type, + layer_id=layer_id, + ) + if len(workload.participant_ep_ids) != int(self._moe_ep_size): + raise ValueError( + "materialized EP lane count does not match predictor topology: " + f"descriptors={len(workload.participant_ep_ids)}, " + f"predictor={self._moe_ep_size}" + ) + if int(self._moe_ep_size) != 1: + raise ValueError( + "load-aware regular-batch prediction requires an explicit " + "physical EP lane for EP>1" + ) + return workload.lane(0) + return self._get_effective_moe_total_tokens(batch) + + def _get_gating_time(self, batch: Batch) -> float: + """ + Get total MoE gating network execution time (linear + routing_topk). + + The gating network determines which experts each token should be routed to. + Prediction is based on num_tokens feature from profiling data. + + Returns: + Total gating time (sum of linear and routing_topk times) + """ + return self._get_gating_linear_time(batch) + self._get_gating_routing_topk_time( + batch + ) + + def _get_gating_linear_time(self, batch: Batch) -> float: + """ + Get MoE gating linear layer execution time. + + The gating linear layer computes logits from hidden states (hidden_dim -> num_experts). + """ + if not self._supports_operation("moe_gating_linear"): + raise NotImplementedError( + "MoE gating linear is not supported for cluster type" + ) + model_name = self._select_moe_gating_prediction_model_name( + "moe_gating_linear", + batch, + ) + if model_name not in self._predictions: + raise NotImplementedError( + "MoE gating linear is not supported for cluster type" + ) + effective_tokens = batch.get_effective_total_tokens_rounded(self._cluster_type) + return self._get_prediction_for_features( + model_name, + {"num_tokens": effective_tokens}, + feature_names=("num_tokens",), + ) + + def _get_gating_routing_topk_time(self, batch: Batch) -> float: + """ + Get MoE gating routing topk execution time. + + The routing topk operation selects top-K experts and applies softmax normalization. + """ + if not self._supports_operation("moe_gating_routing_topk"): + raise NotImplementedError( + "MoE gating routing topk is not supported for cluster type" + ) + model_name = self._select_moe_gating_prediction_model_name( + "moe_gating_routing_topk", + batch, + ) + if model_name not in self._predictions: + raise NotImplementedError( + "MoE gating routing topk is not supported for cluster type" + ) + effective_tokens = batch.get_effective_total_tokens_rounded(self._cluster_type) + return self._get_prediction_for_features( + model_name, + {"num_tokens": effective_tokens}, + feature_names=("num_tokens",), + ) + + def _resolve_shuffling_per_expert_tokens( + self, + batch: Batch, + moe_tokens_input: Optional[EPLaneWorkload] = None, + ) -> EPLaneWorkload: + source = batch if moe_tokens_input is None else moe_tokens_input + lane_workload = resolve_ep_lane_workload(source, required=True) + assert lane_workload is not None + return lane_workload + + def _get_moe_shuffling_time( + self, + batch: Batch, + moe_tokens_input: Optional[EPLaneWorkload] = None, + ) -> float: + """ + Get MoE token shuffling execution time using trained prediction model. + + Shuffling involves dispatching tokens to assigned experts. When the model is + trained with load-imbalance features, use on-demand prediction driven by + per-expert allocation; otherwise use the legacy num_tokens lookup table. + """ + if not self._supports_operation("moe_shuffling"): + raise NotImplementedError("MoE shuffling is not supported for cluster type") + if "moe_shuffling" not in self._predictions: + raise NotImplementedError("MoE shuffling is not supported for cluster type") + if moe_tokens_input is not None and not isinstance( + moe_tokens_input, EPLaneWorkload + ): + raise TypeError( + "MoE shuffling requires an EPLaneWorkload descriptor when an " + "explicit workload is supplied" + ) + + prediction_cache = self._predictions["moe_shuffling"] + if isinstance(prediction_cache, dict) and prediction_cache.get( + "_on_demand_prediction", False + ): + lane_workload = self._resolve_shuffling_per_expert_tokens( + batch, + moe_tokens_input=moe_tokens_input, + ) + if lane_workload.routed_token_count == 0: + raw_time = 0.0 + else: + features = self._build_moe_load_imbalance_features( + lane_workload, + batch=batch, + ) + raw_time = self._get_on_demand_prediction( + "moe_shuffling", features + ) + else: + lane_workload = resolve_ep_lane_workload(batch, required=False) + if moe_tokens_input is not None: + lane_workload = resolve_ep_lane_workload( + moe_tokens_input, + required=True, + ) + if lane_workload is not None: + if lane_workload.routed_token_count == 0: + return 0.0 + effective_tokens = self._get_moe_pre_routing_token_count(batch) + else: + effective_tokens = batch.get_effective_total_tokens_rounded( + self._cluster_type + ) + raw_time = self._get_prediction_for_features( + "moe_shuffling", + {"num_tokens": effective_tokens}, + feature_names=("num_tokens",), + ) + + return raw_time + + def _get_expert_parallel_communication_time( + self, + batch: Batch, + *, + lane_workload: Optional[EPLaneWorkload] = None, + ) -> float: + """ + Get expert parallel communication time. + + Shared-domain MoE execution (monolithic / prefill / decode) uses + expert-parallel all-reduce when EP is enabled without all-to-all routing. + Post-routing EP batches (e.g. DECODE_FFN) and flattened multi-DP MoE + paths keep the all-to-all communication model. + """ + if self._moe_ep_size <= 1: + return 0.0 + + uses_alltoall = self._use_expert_parallel_alltoall_path(batch) + resolved_lane_workload = None + if uses_alltoall: + resolved_lane_workload = resolve_ep_lane_workload( + batch if lane_workload is None else lane_workload, + required=True, + ) + assert resolved_lane_workload is not None + + if self._cc_backend is not None: + quant_manager = get_quantization_manager() + + if uses_alltoall: + routed_tokens = self._get_local_ep_routed_tokens( + batch, + lane_workload=resolved_lane_workload, + ) + data_size_bytes = self._model_config.embedding_dim * 2 * routed_tokens + data_size_bytes = quant_manager.adjust_tensor_size( + "expert_parallel_communication", data_size_bytes, self._cluster_type + ) + result = self._cc_backend.predict_all_to_all( + data_size_bytes=data_size_bytes, + num_devices=self._moe_ep_size, + cluster_type=self._cluster_type, + comm_domain="EP", + ) + logger.debug( + f"_get_expert_parallel_communication_time: using EP all-to-all, " + f"data_size={data_size_bytes}, num_devices={self._moe_ep_size}, " + f"result={result:.6f} ms" + ) + return result + + effective_tokens = batch.get_effective_total_tokens_rounded(self._cluster_type) + data_size_bytes = self._model_config.embedding_dim * 2 * effective_tokens + data_size_bytes = quant_manager.adjust_tensor_size( + "allreduce", data_size_bytes, self._cluster_type + ) + result = self._cc_backend.predict_allreduce( + data_size_bytes=data_size_bytes, + num_devices=self._moe_ep_size, + cluster_type=self._cluster_type, + comm_domain="EP", + ) + result = self._strip_collective_sim_allreduce_launch_overhead_if_needed( + batch=batch, + predicted_ms=result, + num_devices=self._moe_ep_size, + comm_domain="EP", + ) + logger.debug( + f"_get_expert_parallel_communication_time: using EP all-reduce, " + f"data_size={data_size_bytes}, num_devices={self._moe_ep_size}, " + f"result={result:.6f} ms" + ) + return result + + if self._enable_dummy_mode: + logger.debug( + f"_get_expert_parallel_communication_time: CC Backend not available, " + f"using dummy mode value={self._dummy_execution_time} ms" + ) + return self._dummy_execution_time + + raise RuntimeError( + f"CC Backend is required for expert parallel communication prediction " + f"but was not provided. Either:\n" + f" 1. Configure a CC Backend (e.g., --cc_backend vidur or --cc_backend analytical)\n" + f" 2. Enable dummy mode explicitly (--enable_dummy_mode)\n" + f"Current state: cc_backend=None, enable_dummy_mode={self._enable_dummy_mode}" + ) + + def _get_grouped_gemm_time( + self, + num_tokens_or_allocation, + batch: Optional[Batch] = None, + ) -> float: + """ + Calculate grouped GEMM time using trained prediction model. + + Args: + num_tokens_or_allocation: An ``EPLaneWorkload`` for EP-aware + prediction, or an integer for the legacy + one-feature non-lane path. + + Returns: + Total grouped GEMM execution time + """ + if not self._supports_operation("moe_grouped_gemm"): + raise NotImplementedError( + "MoE grouped_gemm is not supported for cluster type" + ) + + if "moe_grouped_gemm" not in self._predictions: + raise NotImplementedError( + "MoE grouped_gemm is not supported for cluster type" + ) + + prediction_cache = self._predictions["moe_grouped_gemm"] + + if isinstance(num_tokens_or_allocation, Mapping): + raise TypeError( + "MoE grouped_gemm requires an EPLaneWorkload descriptor; raw " + "expert-token maps are not a predictor workload contract" + ) + lane_workload = ( + resolve_ep_lane_workload(num_tokens_or_allocation, required=True) + if isinstance(num_tokens_or_allocation, EPLaneWorkload) + else None + ) + + # Check if this model uses on-demand prediction (trained with load imbalance features) + if isinstance(prediction_cache, dict) and prediction_cache.get( + "_on_demand_prediction" + ): + # On-demand prediction mode: model was trained with load imbalance features. + # We must provide the full feature set computed from per-expert token distribution. + if lane_workload is None: + raise ValueError( + "moe_grouped_gemm is in load-imbalance (on-demand) mode and " + "requires an EPLaneWorkload descriptor" + ) + + if lane_workload.routed_token_count == 0: + return 0.0 + + features = self._build_moe_load_imbalance_features( + lane_workload, + batch=batch, + ) + return self._get_on_demand_prediction("moe_grouped_gemm", features) + + # Standard cache lookup mode (trained with num_tokens only) + if lane_workload is not None: + if lane_workload.routed_token_count == 0: + return 0.0 + source_num_tokens = self._get_moe_pre_routing_token_count(batch) + raw_time = self._get_prediction_for_features( + "moe_grouped_gemm", + {"num_tokens": source_num_tokens}, + feature_names=("num_tokens",), + ) + return raw_time + + # Backward compatibility: single number of tokens + num_tokens = num_tokens_or_allocation + if isinstance(num_tokens, bool) or not isinstance(num_tokens, (int, float)): + raise TypeError( + "MoE grouped_gemm requires an EPLaneWorkload descriptor or a " + "numeric token count" + ) + if num_tokens <= 0: + return 0.0 + raw_time = self._get_prediction_for_features( + "moe_grouped_gemm", + {"num_tokens": num_tokens}, + feature_names=("num_tokens",), + ) + return raw_time + + @staticmethod + + def _resolve_moe_execution_inputs( + *, + moe_tokens_input: object, + lane_workload: Optional[EPLaneWorkload], + include_moe: bool, + ) -> tuple[object, Optional[EPLaneWorkload]]: + """Resolve one canonical MoE input and its optional physical lane. + + ``moe_tokens_input`` is retained for the legacy scalar one-feature + lookup, while ``lane_workload`` carries the physical routed domain. + A physical call must use one descriptor for both roles; allowing a + scalar or a second descriptor alongside it would let communication and + routed compute describe different workloads. + """ + + if isinstance(moe_tokens_input, Mapping): + raise TypeError( + "moe_tokens_input cannot be a raw expert-token map; provide an " + "EPLaneWorkload descriptor" + ) + + explicit_lane = ( + resolve_ep_lane_workload(lane_workload, required=True) + if lane_workload is not None + else None + ) + input_lane = ( + resolve_ep_lane_workload(moe_tokens_input, required=True) + if isinstance(moe_tokens_input, EPLaneWorkload) + else None + ) + + if explicit_lane is not None: + if input_lane is not None: + if input_lane != explicit_lane: + raise ValueError( + "moe_tokens_input and lane_workload must refer to the " + "same EPLaneWorkload descriptor" + ) + return explicit_lane, explicit_lane + if moe_tokens_input is not None: + raise TypeError( + "cannot combine a scalar moe_tokens_input with an " + "explicit lane_workload" + ) + return explicit_lane, explicit_lane + + if input_lane is not None: + return input_lane, input_lane + + if include_moe and moe_tokens_input is None: + raise ValueError( + "moe_tokens_input is required when include_moe=True. " + "Provide a scalar token count or an EPLaneWorkload descriptor." + ) + return moe_tokens_input, None diff --git a/frontier/execution_time_predictor/moe_predictor_helpers.py b/frontier/execution_time_predictor/moe_predictor_helpers.py new file mode 100644 index 00000000..b9310285 --- /dev/null +++ b/frontier/execution_time_predictor/moe_predictor_helpers.py @@ -0,0 +1,176 @@ +"""Module-level helpers shared by the MoE execution-time predictor parts.""" + +import math +import pandas as pd + +from frontier.entities.time_components import MoEOperatorTimes +from frontier.moe_gating_runtime import get_moe_gating_base_model_name +from frontier.operators.families import MOE_FAMILY, get_family_profiling_names +from typing import Mapping + + +def _normalize_routing_details_for_trace( + routing_details: Mapping[int, Mapping[int, Mapping[int, float]]], +) -> dict[str, dict[str, dict[str, float]]]: + """Return a strict JSON-safe copy of runtime routing details. + + The matrix checker compares this emitted object with an independently + materialized sidecar. The trace must therefore contain the actual + predictor-owned map, not a derived token allocation or a digest. + """ + + if not isinstance(routing_details, Mapping) or not routing_details: + raise ValueError("routing_details trace payload must be a non-empty mapping") + normalized: dict[str, dict[str, dict[str, float]]] = {} + for replica_id, per_layer in routing_details.items(): + if type(replica_id) is not int or replica_id < 0: + raise ValueError( + "routing_details trace replica IDs must be non-negative integers" + ) + if not isinstance(per_layer, Mapping) or not per_layer: + raise ValueError( + f"routing_details trace replica {replica_id} has no layer map" + ) + normalized_layers: dict[str, dict[str, float]] = {} + for layer_id, per_expert in per_layer.items(): + if type(layer_id) is not int or layer_id < 0: + raise ValueError( + "routing_details trace layer IDs must be non-negative integers" + ) + if not isinstance(per_expert, Mapping) or not per_expert: + raise ValueError( + f"routing_details trace layer {layer_id} has no expert map" + ) + normalized_experts: dict[str, float] = {} + for expert_id, ratio in per_expert.items(): + if type(expert_id) is not int or expert_id < 0: + raise ValueError( + "routing_details trace expert IDs must be non-negative integers" + ) + value = float(ratio) + if not math.isfinite(value) or value < 0.0: + raise ValueError( + "routing_details trace ratios must be finite and non-negative" + ) + normalized_experts[str(expert_id)] = value + ratio_sum = sum(normalized_experts.values()) + if not math.isclose(ratio_sum, 1.0, rel_tol=0.0, abs_tol=1e-12): + raise ValueError( + "routing_details trace ratios must sum to one " + f"for replica={replica_id} layer={layer_id}, got {ratio_sum}" + ) + normalized_layers[str(layer_id)] = normalized_experts + normalized[str(replica_id)] = normalized_layers + return normalized + + +def _get_moe_family_model_names() -> list[str]: + return list(get_family_profiling_names(MOE_FAMILY)) + + +def _get_moe_family_operator_by_model_name(model_name: str): + moe_ops = { + operator.profiling_name(): operator + for operator in MOE_FAMILY.profiling_ops() + } + if model_name not in moe_ops: + raise ValueError(f"Unsupported MoE op: {model_name}") + return moe_ops[model_name] + + +def _get_moe_gating_family_model_names() -> list[str]: + return [ + operator.profiling_name() + for operator in MOE_FAMILY.profiling_ops() + if operator.precision_name() == "moe_gating" + ] + + +_MOE_GATING_OPERATOR_NAMES = frozenset( + operator.name + for operator in MOE_FAMILY.profiling_ops() + if operator.precision_name() == "moe_gating" +) + + +def _get_prefill_hot_moe_gating_model_names() -> list[str]: + return [ + f"{model_name}__prefill_hot" + for model_name in _get_moe_gating_family_model_names() + ] + + +def _is_moe_gating_family_model_name(model_name: str) -> bool: + base_model_name = get_moe_gating_base_model_name(model_name) + return _get_moe_family_operator_by_model_name( + base_model_name + ).precision_name() == "moe_gating" + + +def _build_moe_operator_times( + *, + mlp_norm_time: float, + moe_gating_linear_time: float, + moe_gating_routing_topk_time: float, + moe_shuffling_time: float, + moe_grouped_gemm_time: float, + share_expert_up_proj_time: float = 0.0, + share_expert_act_time: float = 0.0, + share_expert_down_proj_time: float = 0.0, + include_share_expert: bool = False, +) -> MoEOperatorTimes: + op_times = { + "post_attention_layernorm": mlp_norm_time, + "moe_gating_linear": moe_gating_linear_time, + "moe_gating_routing_topk": moe_gating_routing_topk_time, + "moe_shuffling": moe_shuffling_time, + "moe_grouped_gemm": moe_grouped_gemm_time, + } + if include_share_expert: + op_times.update( + { + "share_expert_up_proj": share_expert_up_proj_time, + "share_expert_act": share_expert_act_time, + "share_expert_down_proj": share_expert_down_proj_time, + } + ) + return MoEOperatorTimes(op_times=op_times) + + +def _validate_moe_columns(moe_df: pd.DataFrame) -> None: + """ + Validate that MoE DataFrame contains required split gating columns. + + This function enforces fail-fast behavior by rejecting legacy moe_gating + column format and requiring the split columns (moe_gating_linear and + moe_gating_routing_topk). + + Args: + moe_df: DataFrame containing MoE profiling data + + Raises: + ValueError: If required split columns are missing or if legacy + moe_gating column is present without split columns + """ + required_columns = [ + f"time_stats.{operator_name}.median" + for operator_name in get_family_profiling_names(MOE_FAMILY) + ] + + missing_columns = [col for col in required_columns if col not in moe_df.columns] + + if missing_columns: + # Check if legacy moe_gating column exists (for better error message) + legacy_col = "time_stats.moe_gating.median" + if legacy_col in moe_df.columns: + raise ValueError( + f"Missing required MoE columns: {missing_columns}. " + f"Found legacy '{legacy_col}' column which is no longer supported. " + f"Re-run MoE profiling with split gating scopes enabled to generate " + f"'moe_gating_linear' and 'moe_gating_routing_topk' columns." + ) + else: + raise ValueError( + f"Missing required MoE columns: {missing_columns}. " + f"Re-run MoE profiling with split gating scopes enabled." + ) diff --git a/frontier/execution_time_predictor/moe_routing_workload.py b/frontier/execution_time_predictor/moe_routing_workload.py new file mode 100644 index 00000000..0f6bc138 --- /dev/null +++ b/frontier/execution_time_predictor/moe_routing_workload.py @@ -0,0 +1,583 @@ +"""Expert-load distribution and the per-lane workload it produces. + +Routing decides how many of a batch's tokens each expert receives. These +methods draw that distribution, turn it into the per-lane token counts an +expert-parallel domain sees, and admit the aggregate a lane may process. +""" + +import json + +from frontier.config import ReplicaConfig +from frontier.entities import Batch +from frontier.execution_time_predictor.moe_predictor_helpers import ( + _normalize_routing_details_for_trace, +) +from frontier.logger import init_logger +from frontier.moe_ep_workload import ( + EPLaneWorkload, + LayerEPWorkload, + build_contiguous_expert_ownership, + generate_moe_routing_ratios, + materialize_layer_ep_workload, + resolve_ep_lane_workload, + resolve_routing_details, +) +from frontier.moe_routing_runtime import resolve_moe_gating_routing_runtime_path +from frontier.types import ClusterType +from typing import Dict, List, Mapping, Optional + + +logger = init_logger(__name__) + + +class MoeRoutingWorkload: + """Expert-load distribution and per-lane routed workload.""" + + @staticmethod + + def _emit_routing_details_snapshot( + cluster_type: ClusterType, + routing_details: Mapping[int, Mapping[int, Mapping[int, float]]], + ) -> None: + """Emit the exact predictor-owned routing map for external validation.""" + + normalized = _normalize_routing_details_for_trace(routing_details) + payload = { + "schema_version": 1, + "cluster": cluster_type.name, + "routing_details": normalized, + } + logger.info( + "[ROUTING-SNAPSHOT] %s", + json.dumps(payload, sort_keys=True, separators=(",", ":")), + ) + + def _get_requested_moe_gating_routing_runtime_path(self) -> str: + return resolve_moe_gating_routing_runtime_path( + getattr(self, "_moe_routing_distribution_type", "balanced") + ) + + @staticmethod + + def _get_ep_lane_routed_token_count( + batch: Batch, + lane_workload: Optional[EPLaneWorkload] = None, + ) -> Optional[int]: + """Return an EP lane's routed-token count, or ``None`` for full batches. + + Dummy mode still models the same five-phase EP contract as the + profiling-backed path. The lane-local routed compute therefore has + to depend on the materialized expert map even when the other dummy + components use fixed structural timings. + """ + + if lane_workload is None: + lane_workload = resolve_ep_lane_workload(batch, required=False) + if lane_workload is None: + return None + return lane_workload.routed_token_count + + def _admit_routed_ep_aggregate( + self, + batch: Batch, + *, + routed_moe: bool, + ep_size: Optional[int] = None, + router_topk: Optional[int] = None, + lane_workload: Optional[EPLaneWorkload] = None, + conservation_context: str = "routed MoE admission", + ) -> Optional[EPLaneWorkload]: + """Admit a concrete routed MoE call at the public predictor boundary. + + Concrete predictors own the semantic classification of a call. Once + that classification is routed MoE, an EP>1 call must identify one + physical lane before any mode-specific timing or lookup work begins. + The helper also validates the descriptor against the active predictor + topology and the source/lane token ledger. Workload construction + remains owned by the scheduler/materializer path. + """ + if type(routed_moe) is not bool: + raise ValueError("routed_moe must be a bool") + if not routed_moe: + return None + + if ep_size is None: + configured_ep_size = getattr(self, "_moe_ep_size", None) + if configured_ep_size is None: + replica_config = getattr(self, "_replica_config", None) + configured_ep_size = getattr( + replica_config, + "moe_expert_parallel_size", + None, + ) + else: + configured_ep_size = ep_size + if type(configured_ep_size) is not int or configured_ep_size < 1: + raise ValueError( + "routed MoE admission requires a positive integer EP size, got " + f"{configured_ep_size!r}" + ) + + batch_lane_workload = resolve_ep_lane_workload(batch, required=False) + explicit_lane_workload = ( + resolve_ep_lane_workload(lane_workload, required=True) + if lane_workload is not None + else None + ) + if ( + batch_lane_workload is not None + and explicit_lane_workload is not None + and batch_lane_workload != explicit_lane_workload + ): + raise ValueError( + "batch and lane_workload must refer to the same " + "EPLaneWorkload descriptor" + ) + resolved_lane_workload = explicit_lane_workload or batch_lane_workload + if configured_ep_size > 1 and resolved_lane_workload is None: + raise ValueError( + "Routed MoE prediction with EP>1 requires an EPLaneWorkload " + "descriptor" + ) + if resolved_lane_workload is None: + return None + + configured_router_topk = ( + getattr(self, "_router_topk", None) + if router_topk is None + else router_topk + ) + if configured_router_topk is not None: + if ( + type(configured_router_topk) is not int + or configured_router_topk < 1 + ): + raise ValueError( + "routed MoE admission requires a positive integer router top-k, " + f"got {configured_router_topk!r}" + ) + if resolved_lane_workload.router_topk != configured_router_topk: + raise ValueError( + "lane_workload router_topk does not match predictor topology: " + f"descriptor={resolved_lane_workload.router_topk}, " + f"predictor={configured_router_topk}" + ) + else: + configured_router_topk = resolved_lane_workload.router_topk + + if ( + resolved_lane_workload.moe_expert_parallel_size + != configured_ep_size + ): + raise ValueError( + "lane_workload EP size does not match predictor topology: " + f"descriptor={resolved_lane_workload.moe_expert_parallel_size}, " + f"predictor={configured_ep_size}" + ) + + # A descriptor attached to an EP lane entity already contains routed + # assignments, so its count must match that entity's physical width. + # An explicit descriptor paired with an ordinary source batch represents + # one assignment subset of the aggregate. The aggregate materializer, + # rather than this predictor boundary, owns its conservation ledger. + source_total_num_tokens = getattr(batch, "total_num_tokens", None) + if source_total_num_tokens is not None: + if ( + type(source_total_num_tokens) is not int + or source_total_num_tokens < 0 + ): + raise ValueError( + "routed MoE admission requires batch.total_num_tokens to be " + "a non-negative integer, got " + f"{source_total_num_tokens!r}" + ) + if batch_lane_workload is not None: + expected_routed_token_count = source_total_num_tokens + if ( + resolved_lane_workload.routed_token_count + != expected_routed_token_count + ): + raise ValueError( + f"Token conservation violated in {conservation_context}: " + f"allocated {resolved_lane_workload.routed_token_count}, " + f"expected {expected_routed_token_count}" + ) + + return resolved_lane_workload + + def _init_global_routing_allocations(self) -> Dict[int, Dict[int, float]]: + """Pre-compute global expert allocation ratios for shared-domain EP sync. + + Monolithic decode with EP enabled needs a global view across all experts to + derive per-lane post-MoE arrival skew before the shared-domain all-reduce. + """ + total_experts = self._replica_config.total_expert_num + if self._cluster_type == ClusterType.DECODE_ATTN or not self._model_config.is_moe: + return {} + num_layers = self._model_config.num_layers + + if type(total_experts) is not int or total_experts <= 0: + raise ValueError( + "total_expert_num must be an exact positive int for routing; " + f"got {total_experts!r}" + ) + if type(self._moe_ep_size) is not int or self._moe_ep_size <= 0: + raise ValueError( + "moe_expert_parallel_size must be an exact positive int for routing; " + f"got {self._moe_ep_size!r}" + ) + if total_experts % self._moe_ep_size != 0: + raise ValueError( + "total_expert_num must be divisible by moe_expert_parallel_size; " + f"got total_expert_num={total_experts}, " + f"moe_expert_parallel_size={self._moe_ep_size}" + ) + + distribution_type = self._moe_routing_distribution_type + allocations: Dict[int, Dict[int, float]] = {} + for layer_id in range(num_layers): + allocations[layer_id] = generate_moe_routing_ratios( + total_expert_num=total_experts, + distribution_type=distribution_type, + seed=self._moe_routing_seed, + layer_id=layer_id, + ) + + return allocations + + def _build_shared_routing_details( + self, + ) -> Dict[int, Dict[int, Dict[int, float]]]: + """Expose one immutable-shape routing source for monolithic schedulers. + + ``_global_routing_allocations`` is generated once per model layer and is + intentionally replica-independent. The scheduler, however, performs an + exact ``(replica_id, global_layer_id)`` lookup. Materialize that lookup + shape here without generating a second distribution or assigning tokens + to requests. Integer token accounting and EP ownership splitting remain + the responsibility of the shared per-layer materializer. + + The current monolithic predictor is constructed from ``ReplicaConfig`` + rather than ``ClusterConfig``. The canonical cluster capacity is + bound to ``ReplicaConfig.cluster_num_replicas`` before this method is + called. When the simulator supplies ``_actual_replica_ids``, those + process-global IDs are used as the outer map keys; otherwise local + ``range(replica_count)`` keys support standalone predictor construction. + A missing capacity is an invalid topology, not a condition to infer + from an attention-DP field. + """ + replica_count = self._replica_config.cluster_num_replicas + if type(replica_count) is not int or replica_count <= 0: + raise ValueError( + "A positive cluster replica count is required to build shared " + f"routing details; got {replica_count!r}" + ) + + actual_replica_ids = self._actual_replica_ids + if actual_replica_ids is None: + replica_ids = list(range(replica_count)) + else: + if not isinstance(actual_replica_ids, (list, tuple)): + raise ValueError( + "actual_replica_ids must be a list or tuple when provided" + ) + replica_ids = list(actual_replica_ids) + if len(replica_ids) != replica_count: + raise ValueError( + "actual_replica_ids length must match cluster replica count; " + f"got {len(replica_ids)} for {replica_count} replicas" + ) + if any( + type(replica_id) is not int or replica_id < 0 + for replica_id in replica_ids + ): + raise ValueError( + "actual_replica_ids must contain exact non-negative integers" + ) + if len(set(replica_ids)) != len(replica_ids): + raise ValueError("actual_replica_ids must be unique") + + return { + replica_id: { + layer_id: dict(expert_ratios) + for layer_id, expert_ratios in self._global_routing_allocations.items() + } + for replica_id in replica_ids + } + + def _get_routing_details_for_cluster(self, cluster_type: ClusterType): + """Return the exact pre-generated routing map for one serving role.""" + attribute_by_cluster = { + ClusterType.MONOLITHIC: "_monolithic_routing_details", + ClusterType.PREFILL: "_prefill_routing_details", + ClusterType.DECODE: "_decode_routing_details", + ClusterType.DECODE_FFN: "_decode_ffn_routing_details", + } + attribute_name = attribute_by_cluster.get(cluster_type) + if attribute_name is None: + raise ValueError( + f"MoE routing materialization does not support cluster_type={cluster_type}" + ) + routing_details = getattr(self, attribute_name, None) + if routing_details is None: + raise ValueError( + f"Missing pre-generated routing details for cluster_type={cluster_type}" + ) + return routing_details + + def _get_cluster_replica_config(self, cluster_type: ClusterType) -> ReplicaConfig: + """Return the serving replica; disaggregated predictors override by role.""" + return self._replica_config + + def _materialize_layer_ep_workload( + self, batch: Batch, cluster_type: ClusterType, layer_id: int + ) -> LayerEPWorkload: + """Materialize one exact Replica-local EP workload for a MoE layer.""" + # Routing tables are built once by the constructor and have no runtime + # mutation API. Topology and exact replica/layer/token identity are in + # the key; the frozen workload can be shared across repeated EP waves. + workload_cache = self._layer_workload_cache + cache_capacity = self._layer_workload_cache_capacity + cluster_replica_config = self._get_cluster_replica_config(cluster_type) + routing_details = self._get_routing_details_for_cluster(cluster_type) + target_replica_id = int(batch.replica_id) + global_layer_id = int(layer_id) + routing_token_count = int(batch.total_num_tokens) + router_topk = int(cluster_replica_config.router_topk) + total_expert_num = int(cluster_replica_config.total_expert_num) + moe_ep_size = int(cluster_replica_config.moe_expert_parallel_size) + cache_key = ( + cluster_type, + target_replica_id, + global_layer_id, + routing_token_count, + router_topk, + total_expert_num, + moe_ep_size, + ) + cached_workload = workload_cache.get(cache_key) + if cached_workload is not None: + workload_cache.move_to_end(cache_key) + return cached_workload + workload = materialize_layer_ep_workload( + routing_ratios=resolve_routing_details( + routing_details, + target_replica_id=target_replica_id, + global_layer_id=global_layer_id, + ), + target_replica_id=target_replica_id, + global_layer_id=global_layer_id, + routing_token_count=routing_token_count, + router_topk=router_topk, + total_expert_num=total_expert_num, + moe_expert_parallel_size=moe_ep_size, + expert_to_ep=build_contiguous_expert_ownership( + total_expert_num, + moe_ep_size, + ), + ) + workload_cache[cache_key] = workload + workload_cache.move_to_end(cache_key) + while len(workload_cache) > cache_capacity: + workload_cache.popitem(last=False) + return workload + + def _resolve_layer_lane_workload( + self, + batch: Batch, + *, + cluster_type: ClusterType, + layer_id: int, + ) -> EPLaneWorkload: + """Resolve one physical lane descriptor at a predictor boundary. + + Scheduler-created lane entities already carry the descriptor. A + regular batch may be materialized into one lane only for EP=1; an EP>1 + aggregate must be expanded by the scheduler's lane wave first so the + predictor never guesses which local expert domain a global map denotes. + """ + + lane_workload = resolve_ep_lane_workload(batch, required=False) + if lane_workload is not None: + return lane_workload + + layer_workload = self._materialize_layer_ep_workload( + batch=batch, + cluster_type=cluster_type, + layer_id=layer_id, + ) + participant_ep_ids = tuple(layer_workload.participant_ep_ids) + if len(participant_ep_ids) != int(self._moe_ep_size): + raise ValueError( + "materialized EP lane count does not match predictor topology: " + f"descriptors={len(participant_ep_ids)}, predictor={self._moe_ep_size}" + ) + if len(participant_ep_ids) != 1: + raise ValueError( + "regular aggregate MoE prediction requires a physical EP lane " + "for EP>1; scheduler lane materialization is required" + ) + return layer_workload.lane(participant_ep_ids[0]) + + def _resolve_shared_domain_lane_workloads( + self, + batch: Batch, + *, + cluster_type: ClusterType, + layer_id: int, + ) -> tuple[EPLaneWorkload, ...]: + """Resolve every physical lane for a shared-domain MoE timing probe.""" + + lane_count = int(self._moe_ep_size) + if lane_count <= 0: + raise ValueError( + "shared-domain MoE lane resolution requires a positive EP size, " + f"got {lane_count}" + ) + + lane_workload = resolve_ep_lane_workload(batch, required=False) + if lane_workload is not None: + if lane_workload.moe_expert_parallel_size != lane_count: + raise ValueError( + "batch lane workload EP size does not match predictor: " + f"descriptor={lane_workload.moe_expert_parallel_size}, " + f"predictor={lane_count}" + ) + lane_workloads = (lane_workload,) + else: + workload = self._materialize_layer_ep_workload( + batch=batch, + cluster_type=cluster_type, + layer_id=layer_id, + ) + lane_workloads = tuple( + workload.lane(ep_id) for ep_id in workload.participant_ep_ids + ) + + if len(lane_workloads) != lane_count: + raise ValueError( + "materialized EP lane count does not match predictor topology: " + f"descriptors={len(lane_workloads)}, predictor={lane_count}" + ) + return lane_workloads + + def _build_moe_load_imbalance_features( + self, + lane_workload: EPLaneWorkload, + *, + batch: Optional[Batch] = None, + ) -> Dict[str, float]: + lane_workload = resolve_ep_lane_workload(lane_workload, required=True) + assert lane_workload is not None + + from frontier.moe_load_imbalance import MoELoadImbalanceInput + + expert_token_counts = [int(v) for v in lane_workload.local_token_counts] + + total_routed_tokens = int(sum(expert_token_counts)) + if lane_workload.router_topk <= 0: + raise ValueError(f"Invalid router_topk={lane_workload.router_topk}") + + source_num_tokens = self._get_moe_pre_routing_token_count(batch) + + load_input = MoELoadImbalanceInput( + num_tokens=source_num_tokens, + num_experts_per_device=lane_workload.local_expert_width, + hidden_dim=int(self._model_config.embedding_dim), + expert_hidden_dim=int(self._model_config.mlp_hidden_dim), + router_topk=int(lane_workload.router_topk), + expert_token_counts=expert_token_counts, + load_distribution="runtime", + ) + features = load_input.to_features_dict() + features.pop("load_distribution", None) + features.pop("seed", None) + missing_features = [ + name + for name in self.MOE_LOAD_IMBALANCE_FEATURES + if name not in features + ] + if missing_features: + raise ValueError( + "MoE load-imbalance feature construction is missing canonical " + f"features: {missing_features}" + ) + unexpected_features = sorted( + set(features) - set(self.MOE_LOAD_IMBALANCE_FEATURES) + ) + if unexpected_features: + raise ValueError( + "MoE load-imbalance feature construction produced unexpected " + f"features: {unexpected_features}" + ) + return { + name: features[name] + for name in self.MOE_LOAD_IMBALANCE_FEATURES + } + + def _simulate_routing_per_layer( + self, batches: List[Batch], stage_id: int + ) -> Dict[int, Dict[str, Dict[int, float]]]: + """ + Simulate routing for each layer in the stage. + Returns: {layer_id: {replica_id: {moe_component: time_value}}} + """ + del stage_id + cluster_type = getattr(self, "_cluster_type", None) + if not isinstance(cluster_type, ClusterType): + raise ValueError( + "layer routing prediction requires an initialized cluster_type" + ) + + # Routing materialization is stage-local and follows the canonical + # aggregate-to-lane seam. Predictor consumers receive only physical + # lane descriptors, even when this legacy helper returns one result per + # source replica. + num_layers = self._num_layers_per_pipeline_stage + layer_routing_results = {} + + for layer_id in range(num_layers): + layer_routing_results[layer_id] = {} + + for batch in batches: + replica_id = int(batch.replica_id) + layer_workload = self._materialize_layer_ep_workload( + batch=batch, + cluster_type=cluster_type, + layer_id=layer_id, + ) + lane_workloads = tuple( + layer_workload.lane(ep_id) + for ep_id in layer_workload.participant_ep_ids + ) + if not lane_workloads: + raise ValueError( + "layer routing materialization produced no EP lanes: " + f"replica_id={replica_id}, layer_id={layer_id}" + ) + grouped_gemm_time = max( + self._get_grouped_gemm_time(lane_workload, batch=batch) + for lane_workload in lane_workloads + ) + shuffling_time = max( + self._get_moe_shuffling_time( + batch, + moe_tokens_input=lane_workload, + ) + for lane_workload in lane_workloads + ) + communication_time = max( + self._get_expert_parallel_communication_time( + batch, + lane_workload=lane_workload, + ) + for lane_workload in lane_workloads + ) + layer_routing_results[layer_id][replica_id] = { + "moe_grouped_gemm_time": grouped_gemm_time, + "expert_parallel_communication_time": communication_time, + "moe_gating_time": self._get_gating_time(batch), + "moe_shuffling_time": shuffling_time, + } + + return layer_routing_results diff --git a/frontier/execution_time_predictor/prediction_family_trainers.py b/frontier/execution_time_predictor/prediction_family_trainers.py new file mode 100644 index 00000000..da19f7d1 --- /dev/null +++ b/frontier/execution_time_predictor/prediction_family_trainers.py @@ -0,0 +1,1700 @@ +"""Per-family training of the execution-time prediction models. + +One method per operator family: dense FFN and MoE, dense MLP, attention, +latent MLA attention, residual, pipeline- and tensor-parallel communication, +and CPU overhead. Each selects the rows its family owns, builds the feature +frame, and hands one model at a time to the shared fitting routine. +""" + +import os +import pandas as pd + +from frontier.attention.families import ( + DENSE_ATTENTION_FAMILY, + LATENT_MLA_ATTENTION_FAMILY, +) +from frontier.attention.model_binding import resolve_runtime_attention_family +from frontier.attention.ops import AttentionOperatorRole +from frontier.attention.profiling_mapping import ( + get_enabled_predictor_median_columns, + get_enabled_predictor_metric_name_by_role, + get_enabled_predictor_metric_names, + get_enabled_shared_predictor_feature_columns, + validate_attention_profiling_dataframe, +) +from frontier.attention.string_coercion import coerce_truthy_int +from frontier.execution_time_predictor.prediction_model_identity import ( + _add_layer_contract_to_training_context, + _build_exact_feature_lookup, + _get_contract_hash, + _get_moe_family_model_names, + _get_prefill_hot_moe_gating_model_names, + _is_moe_gating_family_model_name, + _layer_contract_kwargs, + _normalize_layer_contract_context, + _resolve_model_architecture_profile, + _resolve_model_architecture_profile_id, + _serialize_selected_layer_cache_identity, +) +from frontier.logger import init_logger +from frontier.model_architectures import ResolvedLayerContract +from frontier.moe_gating_runtime import ( + DEFAULT_MOE_GATING_RUNTIME_CONTEXT, + PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT, + PrefillHotRowsUnavailableError, + filter_moe_gating_rows_by_runtime_context, + get_moe_gating_base_model_name, + has_prefill_hot_moe_gating_rows, + should_enable_prefill_hot_moe_gating_contract, +) +from frontier.moe_routing_runtime import ( + filter_moe_gating_routing_topk_rows, + resolve_moe_gating_routing_runtime_path, +) +from frontier.operators.families import ( + FFN_FAMILY, + SHARE_EXPERT_FAMILY, + get_family_profiling_names, +) +from frontier.spec_decode.runtime import is_target_embedded_mtp_enabled +from frontier.types import ClusterType, MeasurementType +from sklearn.base import BaseEstimator +from sklearn.model_selection import GridSearchCV +from typing import Any, Dict, List, Optional, Tuple + + +logger = init_logger(__name__) + + +class PredictionFamilyTrainers: + """Per-family model training for the execution-time predictor.""" + + def _train_ffn_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, + linear_ops_file: str, moe_file: str, + is_moe_model: bool, trained_model_signatures: set) -> Dict[str, BaseEstimator]: + """ + Train FFN/MoE models for a specific cluster. + + This function handles FFN-related operations in the Transformer layer: + - FFN core operations (from linear_op.csv): mlp_up_proj, mlp_down_proj, mlp_act + - MoE core operations (from moe.csv): moe_gating_linear, moe_gating_routing_topk, moe_shuffling, moe_grouped_gemm + - Pre-FFN normalization (from linear_op.csv): post_attention_layernorm + + Transformer layer context: + ... → Attention → add → [post_attention_layernorm] → [FFN/MoE] → add → ... + """ + models = {} + + ffn_tp_key = self._get_ffn_tp_key(cluster_type, replica_config, is_moe_model) + tp_size = ffn_tp_key + + # Create a signature for this FFN model configuration. + model_config = replica_config.model_config + model_arch = model_config.get_model_arch() if model_config is not None else "generic" + architecture_profile_id = _resolve_model_architecture_profile_id(model_config) + primary_contract = self._resolve_typed_layer_contract( + "moe_grouped_gemm" if is_moe_model else "mlp_up_proj", + cluster_type, + replica_config, + is_moe_model=is_moe_model, + ) + typed_contract_hash = self._get_ffn_contract_signature( + cluster_type, + replica_config, + is_moe_model, + ) + active_measurement_type = getattr( + self, "_active_measurement_type", MeasurementType.CUDA_EVENT + ) + ffn_signature = ( + f"ffn_{replica_config.device}_{replica_config.model_name}_{tp_size}" + f"_moe{is_moe_model}_arch_profile{architecture_profile_id}" + f"_layer_contracts{typed_contract_hash}" + f"_family{self._measurement_family_name(active_measurement_type)}" + ) + + if ffn_signature in trained_model_signatures: + logger.info(f"Skipping FFN models training for {cluster_type} - already trained with signature {ffn_signature}") + return models + + # Build training context for error messages + training_context = { + 'cluster_type': str(cluster_type), + 'device': replica_config.device, + 'model_name': replica_config.model_name, + 'tensor_parallel_size': tp_size, + 'is_moe_model': is_moe_model, + 'model_arch': model_arch, + 'model_architecture_profile': architecture_profile_id, + 'use_qk_norm': bool(getattr(model_config, 'use_qk_norm', False)), + } + + # Choose input file based on model type + if is_moe_model: + moe_input_file = moe_file + if not os.path.exists(moe_input_file): + raise FileNotFoundError(f"MoE input file {moe_input_file} not found") + logger.info(f"Loading MoE data for {cluster_type} from: {moe_input_file}") + training_context['input_file'] = moe_input_file + + # MoE core operations with per-operation feature selection + # Split gating into moe_gating_linear and moe_gating_routing_topk (Step 1.6) + # Aligned with frontier/training/moe_trainer.py _get_feature_cols() method + base_moe_model_names = _get_moe_family_model_names() + moe_model_names = list(base_moe_model_names) + if should_enable_prefill_hot_moe_gating_contract( + model_config=model_config, + model_arch=model_arch, + model_name=replica_config.model_name, + ): + prefill_hot_probe_df = pd.read_csv(moe_input_file) + include_prefill_hot_models = has_prefill_hot_moe_gating_rows( + prefill_hot_probe_df + ) + + if include_prefill_hot_models: + moe_model_names.extend(_get_prefill_hot_moe_gating_model_names()) + else: + logger.warning( + "Prefill-hot gating contract enabled for model=%s, but " + "dataset %s has no usable prefill_hot rows; skipping " + "__prefill_hot pseudo-models in shared-manager training.", + replica_config.model_name, + moe_input_file, + ) + self._validate_moe_dataset_contract( + moe_input_file, + replica_config, + base_moe_model_names, + cluster_type, + **_layer_contract_kwargs(primary_contract), + ) + requested_routing_runtime_path = resolve_moe_gating_routing_runtime_path( + getattr(replica_config, "moe_routing_distribution_type", "balanced") + ) + + moe_df_cache: Dict[ + Tuple[ + int, + Optional[int], + Optional[str], + Optional[str], + Optional[str], + ], + pd.DataFrame, + ] = {} + + def _get_moe_df_for_op( + model_name: str, + ) -> Tuple[ + pd.DataFrame, + int, + Optional[int], + Optional[ResolvedLayerContract], + ]: + base_model_name = get_moe_gating_base_model_name(model_name) + op_layer_contract = self._resolve_typed_layer_contract( + base_model_name, + cluster_type, + replica_config, + is_moe_model=True, + ) + tp_key = self._get_moe_op_tp_key( + base_model_name, + replica_config, + cluster_type, + ) + if tp_key <= 0: + raise ValueError( + f"Invalid TP key for MoE training: {tp_key} (op={model_name})" + ) + + ep_key: Optional[int] + if self._is_moe_op_ep_agnostic(base_model_name): + ep_key = None + else: + ep_key = replica_config.moe_expert_parallel_size + + runtime_path_key: Optional[str] = None + if base_model_name == "moe_gating_routing_topk": + runtime_path_key = requested_routing_runtime_path + + gating_context_key: Optional[str] = None + if _is_moe_gating_family_model_name(base_model_name): + gating_context_key = DEFAULT_MOE_GATING_RUNTIME_CONTEXT + if model_name.endswith("__prefill_hot"): + gating_context_key = PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT + + contract_identity = _serialize_selected_layer_cache_identity( + op_layer_contract + ) + cache_key = ( + tp_key, + ep_key, + runtime_path_key, + gating_context_key, + contract_identity, + ) + if cache_key not in moe_df_cache: + op_df = self._load_moe_df( + moe_input_file, + replica_config, + load_imbalance=False, + tensor_parallel_size=tp_key, + expert_parallel_size=ep_key, + **_layer_contract_kwargs( + op_layer_contract, + operator_name=base_model_name, + ), + ) + if runtime_path_key is not None: + op_df = filter_moe_gating_routing_topk_rows( + op_df, + requested_runtime_path=runtime_path_key, + source_name=moe_input_file, + ) + if gating_context_key is not None: + op_df = filter_moe_gating_rows_by_runtime_context( + op_df, + requested_context=gating_context_key, + source_name=moe_input_file, + ) + moe_df_cache[cache_key] = op_df + ep_desc = "ANY" if ep_key is None else str(ep_key) + logger.info( + f"Loaded {len(moe_df_cache[cache_key])} rows for MoE training " + f"(op={model_name}, tp_key={tp_key}, ep_key={ep_desc}, " + f"routing_runtime_path={runtime_path_key or 'ANY'}, " + f"gating_runtime_context={gating_context_key or 'ANY'}, " + "auto feature mode)" + ) + return moe_df_cache[cache_key], tp_key, ep_key, op_layer_contract + + for model_name in moe_model_names: + model_signature = f"{model_name}_{ffn_signature}" + if model_signature not in trained_model_signatures: + try: + ( + op_moe_df, + moe_tp_key, + moe_ep_key, + op_layer_contract, + ) = _get_moe_df_for_op(model_name) + except PrefillHotRowsUnavailableError as exc: + logger.warning( + "Skipping %s because prefill-hot gating rows are unavailable " + "for the requested TP/EP slice (%s).", + model_name, + exc, + ) + continue + op_training_context = _add_layer_contract_to_training_context( + training_context, + op_layer_contract, + ) + op_training_context['tensor_parallel_size'] = moe_tp_key + op_training_context['expert_parallel_size'] = ( + "ANY" if moe_ep_key is None else moe_ep_key + ) + + # Per-operation feature selection. + if model_name == "moe_grouped_gemm": + available_load_features = [ + f for f in self.MOE_LOAD_IMBALANCE_FEATURES + if f in op_moe_df.columns + ] + has_load_imbalance_features = ( + len(available_load_features) + == len(self.MOE_LOAD_IMBALANCE_FEATURES) + ) + if 0 < len(available_load_features) < len(self.MOE_LOAD_IMBALANCE_FEATURES): + missing_features = [ + f for f in self.MOE_LOAD_IMBALANCE_FEATURES + if f not in op_moe_df.columns + ] + raise ValueError( + f"Partial load imbalance features found ({len(available_load_features)}/" + f"{len(self.MOE_LOAD_IMBALANCE_FEATURES)}) for {model_name} at TP={moe_tp_key}. " + f"Missing: {missing_features}." + ) + + if has_load_imbalance_features: + op_feature_cols = available_load_features + logger.info( + f" {model_name}: Using load imbalance features " + f"({len(op_feature_cols)} features, TP={moe_tp_key})" + ) + else: + op_feature_cols = ["num_tokens"] + logger.info( + f" {model_name}: Load imbalance features not found; " + f"using num_tokens only (TP={moe_tp_key})." + ) + elif model_name == "moe_shuffling": + available_load_features = [ + f for f in self.MOE_LOAD_IMBALANCE_FEATURES + if f in op_moe_df.columns + ] + if len(available_load_features) == len(self.MOE_LOAD_IMBALANCE_FEATURES): + op_feature_cols = available_load_features + logger.info( + f" {model_name}: Using load imbalance features " + f"({len(op_feature_cols)} features, TP={moe_tp_key})" + ) + else: + # For shuffling we allow partial/legacy datasets and fall back to + # num_tokens-only training when the full load feature set is absent. + op_feature_cols = ["num_tokens"] + logger.info( + f" {model_name}: Full load imbalance features unavailable; " + f"using num_tokens only (TP={moe_tp_key})." + ) + else: + op_feature_cols = ["num_tokens"] + logger.info( + f" {model_name}: Using num_tokens only (1 feature, TP={moe_tp_key})" + ) + + # Store feature_cols in training_context for this specific operation + op_training_context['feature_cols'] = op_feature_cols + + target_op_name = get_moe_gating_base_model_name(model_name) + train_kwargs: Dict[str, Any] = dict( + model_name=model_name, + df=op_moe_df, + feature_cols=op_feature_cols, + target_col=f"time_stats.{target_op_name}.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=op_training_context, + ) + train_kwargs.update(_layer_contract_kwargs(op_layer_contract)) + models[model_name] = self._train_single_model( + **train_kwargs, + ) + trained_model_signatures.add(model_signature) + logger.info(f"Trained {model_name} for {cluster_type} with features: {op_feature_cols}") + + # Step2Mini/Step3 share_expert operations (forward_3: shared expert alongside routed experts) + model_config = replica_config.model_config + if model_config is not None and model_config.supports_share_expert(): + # share_expert operations are trained from linear_op.csv (not moe.csv) + if not os.path.exists(linear_ops_file): + raise FileNotFoundError( + f"Linear ops input file {linear_ops_file} not found for share_expert" + ) + + step2mini_share_expert_model_names = list( + get_family_profiling_names(SHARE_EXPERT_FAMILY) + ) + if not step2mini_share_expert_model_names: + raise ValueError("Shared-expert operator family has no profiling names") + share_expert_tp_key = self._get_linear_op_tp_key( + step2mini_share_expert_model_names[0], + cluster_type, + replica_config, + is_moe_model, + ) + shared_layer_contract = self._resolve_typed_layer_contract( + step2mini_share_expert_model_names[0], + cluster_type, + replica_config, + is_moe_model=True, + ) + if ( + shared_layer_contract is None + and _resolve_model_architecture_profile(model_config) is not None + ): + raise ValueError( + "Missing shared layer contract for share-expert training" + ) + share_expert_linear_ops_df = self._load_linear_op_df( + linear_ops_file, + share_expert_tp_key, + **_layer_contract_kwargs( + shared_layer_contract, + operator_name=step2mini_share_expert_model_names[0], + ), + ) + logger.info(f"Loaded {len(share_expert_linear_ops_df)} rows for share_expert training") + + for model_name in step2mini_share_expert_model_names: + model_signature = f"{model_name}_{ffn_signature}" + if model_signature not in trained_model_signatures: + # Update training context to reflect linear_op.csv source. + shared_training_context = _add_layer_contract_to_training_context( + training_context, + shared_layer_contract, + ) + shared_training_context['input_file'] = linear_ops_file + shared_training_context['tensor_parallel_size'] = share_expert_tp_key + target_col = f"time_stats.{model_name}.median" + if target_col not in share_expert_linear_ops_df.columns: + raise ValueError( + f"share_expert operation '{model_name}' column '{target_col}' not found in profiling data. " + f"Ensure profiling was run with a model architecture that includes share_expert. " + f"Available columns: {list(share_expert_linear_ops_df.columns)}" + ) + train_kwargs: Dict[str, Any] = dict( + model_name=model_name, + df=share_expert_linear_ops_df, + feature_cols=["num_tokens"], + target_col=target_col, + execution_time_predictor_config=execution_time_predictor_config, + training_context=shared_training_context, + ) + train_kwargs.update( + _layer_contract_kwargs(shared_layer_contract) + ) + models[model_name] = self._train_single_model( + **train_kwargs, + ) + trained_model_signatures.add(model_signature) + logger.info(f"Trained {model_name} for {cluster_type}") + + # Mixed-layer MoE models (for example step-moe-noquant) also have + # dense boundary layers. Train these additions after the legacy + # MoE/share-expert families so RandomForest training order remains + # compatible with the historical predictor artifact contract. + if self._is_mixed_layer_moe_model(model_config, is_moe_model): + dense_ffn_tp_key = self._get_ffn_tp_key( + cluster_type, replica_config, is_moe_model=False + ) + dense_layer_contract = self._resolve_typed_layer_contract( + "mlp_up_proj", + cluster_type, + replica_config, + is_moe_model=False, + ) + dense_ffn_signature = ( + f"ffn_{replica_config.device}_{replica_config.model_name}_{dense_ffn_tp_key}" + f"_moeFalse_arch_profile{architecture_profile_id}" + f"_layer_contracts{_get_contract_hash(dense_layer_contract)}" + f"_family{self._measurement_family_name(active_measurement_type)}" + ) + dense_training_context = dict(training_context) + dense_training_context["is_moe_model"] = False + dense_training_context["tensor_parallel_size"] = dense_ffn_tp_key + dense_training_context = _add_layer_contract_to_training_context( + dense_training_context, + dense_layer_contract, + ) + self._train_dense_mlp_models_for_cluster( + cluster_type=cluster_type, + replica_config=replica_config, + execution_time_predictor_config=execution_time_predictor_config, + linear_ops_file=linear_ops_file, + ffn_signature=dense_ffn_signature, + ffn_tp_key=dense_ffn_tp_key, + training_context=dense_training_context, + trained_model_signatures=trained_model_signatures, + models=models, + layer_contract=dense_layer_contract, + ) + else: + self._train_dense_mlp_models_for_cluster( + cluster_type=cluster_type, + replica_config=replica_config, + execution_time_predictor_config=execution_time_predictor_config, + linear_ops_file=linear_ops_file, + ffn_signature=ffn_signature, + ffn_tp_key=ffn_tp_key, + training_context=training_context, + trained_model_signatures=trained_model_signatures, + models=models, + layer_contract=primary_contract, + ) + + # Pre-FFN normalization (post_attention_layernorm) - always from linear_op.csv + if not os.path.exists(linear_ops_file): + raise FileNotFoundError(f"Linear ops input file {linear_ops_file} not found for post_attention_layernorm") + layernorm_tp_key = self._get_linear_op_tp_key( + "post_attention_layernorm", + cluster_type, + replica_config, + is_moe_model, + ) + linear_ops_df = self._load_linear_op_df(linear_ops_file, layernorm_tp_key) + layernorm_context = dict(training_context) + layernorm_context["input_file"] = linear_ops_file + layernorm_context["tensor_parallel_size"] = layernorm_tp_key + + layernorm_model_name = "post_attention_layernorm" + layernorm_signature = f"{layernorm_model_name}_{ffn_signature}" + if layernorm_signature not in trained_model_signatures: + models[layernorm_model_name] = self._train_single_model( + model_name=layernorm_model_name, + df=linear_ops_df, + feature_cols=["num_tokens"], + target_col=f"time_stats.{layernorm_model_name}.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=layernorm_context, + ) + trained_model_signatures.add(layernorm_signature) + logger.info(f"Trained {layernorm_model_name} for {cluster_type}") + + # Mark this FFN configuration as trained + trained_model_signatures.add(ffn_signature) + return models + + def _train_dense_mlp_models_for_cluster( + self, + *, + cluster_type: ClusterType, + replica_config, + execution_time_predictor_config, + linear_ops_file: str, + ffn_signature: str, + ffn_tp_key: int, + training_context: Dict[str, Any], + trained_model_signatures: set, + models: Dict[str, BaseEstimator], + layer_contract: Optional[ResolvedLayerContract] = None, + ) -> None: + """Materialize dense MLP predictors from the linear-op profile.""" + if not os.path.exists(linear_ops_file): + raise FileNotFoundError(f"Linear ops input file {linear_ops_file} not found") + if ( + layer_contract is None + and _resolve_model_architecture_profile( + getattr(replica_config, "model_config", None) + ) + is not None + ): + layer_contract = self._resolve_typed_layer_contract( + "mlp_up_proj", + cluster_type, + replica_config, + is_moe_model=False, + ) + if ( + layer_contract is not None + and layer_contract.tensor_parallel_size is not None + and layer_contract.tensor_parallel_size != ffn_tp_key + ): + raise ValueError( + "Dense FFN training TP conflicts with its typed layer contract: " + f"ffn_tp_key={ffn_tp_key}, " + f"contract_tp={layer_contract.tensor_parallel_size}" + ) + logger.info(f"Loading MLP data for {cluster_type} from: {linear_ops_file}") + dense_model_names = tuple(get_family_profiling_names(FFN_FAMILY)) + if not dense_model_names: + raise ValueError("FFN operator family has no profiling names") + linear_ops_df = self._load_linear_op_df( + linear_ops_file, + ffn_tp_key, + **_layer_contract_kwargs( + layer_contract, + operator_name=dense_model_names[0], + ), + ) + logger.info(f"Loaded {len(linear_ops_df)} rows for MLP training") + dense_training_context = _add_layer_contract_to_training_context( + training_context, + layer_contract, + ) + dense_training_context["input_file"] = linear_ops_file + dense_training_context["tensor_parallel_size"] = ffn_tp_key + + missing_standard_columns = [ + f"time_stats.{model_name}.median" + for model_name in dense_model_names + if f"time_stats.{model_name}.median" not in linear_ops_df.columns + ] + if missing_standard_columns: + model_config = getattr(replica_config, "model_config", None) + supports_share_expert = bool( + model_config is not None + and model_config.supports_share_expert() + ) + if supports_share_expert: + logger.info( + "Skipping standard dense MLP training for %s: profile provides " + "shared-expert operations instead; missing columns=%s", + cluster_type, + missing_standard_columns, + ) + return + raise ValueError( + "Dense MLP profiling data is incomplete; missing columns: " + + ", ".join(missing_standard_columns) + ) + + for model_name in dense_model_names: + model_signature = f"{model_name}_{ffn_signature}" + if model_signature in trained_model_signatures: + continue + train_kwargs: Dict[str, Any] = dict( + model_name=model_name, + df=linear_ops_df, + feature_cols=["num_tokens"], + target_col=f"time_stats.{model_name}.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=dense_training_context, + ) + train_kwargs.update(_layer_contract_kwargs(layer_contract)) + models[model_name] = self._train_single_model(**train_kwargs) + trained_model_signatures.add(model_signature) + logger.info(f"Trained {model_name} for {cluster_type}") + + def _train_attn_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, replica_scheduler_config, linear_ops_file: str, attn_file: str, trained_model_signatures: set) -> Dict[str, BaseEstimator]: + """ + Train attention-related models for a cluster. + + This function handles Attention-related operations in the Transformer layer: + - Pre-attention normalization (from linear_op.csv): input_layernorm + - Attention projections (from linear_op.csv): attn_pre_proj, attn_post_proj, attn_rope + - Attention core operations (from attention.csv): attn_kv_cache_save, attn_prefill, attn_decode + + Transformer layer context: + Input → [input_layernorm] → [attn_pre_proj → attn_rope → attn_prefill/decode → attn_kv_cache_save → attn_post_proj] → add → ... + """ + models = {} + tp_size = replica_config.attn_tensor_parallel_size + + model_config = replica_config.model_config + model_arch = model_config.get_model_arch() if model_config is not None else "generic" + architecture_profile_id = _resolve_model_architecture_profile_id(model_config) + attention_signature = ( + f"attention_{replica_config.device}_{replica_config.model_name}_{tp_size}" + f"_{replica_scheduler_config.block_size}_arch_profile{architecture_profile_id}" + f"_family{self._measurement_family_name(self._active_measurement_type)}" + ) + + if attention_signature in trained_model_signatures: + logger.info(f"Skipping attention models training for {cluster_type} - already trained") + return models + + # Build training context for error messages + training_context = { + 'cluster_type': str(cluster_type), + 'device': replica_config.device, + 'model_name': replica_config.model_name, + 'tensor_parallel_size': tp_size, + 'block_size': replica_scheduler_config.block_size, + 'model_arch': model_arch, + 'model_architecture_profile': architecture_profile_id, + 'use_qk_norm': bool(getattr(model_config, 'use_qk_norm', False)), + } + + # ========== Part 1: Linear operations from linear_op.csv ========== + # These include: input_layernorm, attn_pre_proj, attn_post_proj, attn_rope + if not os.path.exists(linear_ops_file): + raise FileNotFoundError(f"Linear ops input file {linear_ops_file} not found") + + logger.info(f"Loading sharded attention linear-op data from: {linear_ops_file}") + attn_tp_key = self._get_linear_op_tp_key( + "attn_pre_proj", + cluster_type, + replica_config, + is_moe_model=False, + ) + required_columns = self._get_required_attn_linear_op_columns(model_config) + attn_linear_ops_df = self._load_linear_op_df( + linear_ops_file, + attn_tp_key, + required_columns=required_columns, + training_context=training_context, + ) + logger.info( + f"Loaded {len(attn_linear_ops_df)} rows for sharded attention ops training" + ) + + # Pre-attention normalization: input_layernorm + input_layernorm_tp_key = self._get_linear_op_tp_key( + "input_layernorm", + cluster_type, + replica_config, + is_moe_model=False, + ) + input_layernorm_df = self._load_linear_op_df( + linear_ops_file, + input_layernorm_tp_key, + required_columns=["time_stats.input_layernorm.median"], + training_context=training_context, + ) + input_layernorm_context = dict(training_context) + input_layernorm_context["input_file"] = linear_ops_file + input_layernorm_context["tensor_parallel_size"] = input_layernorm_tp_key + + layernorm_model_name = "input_layernorm" + layernorm_signature = f"{layernorm_model_name}_{attention_signature}" + if layernorm_signature not in trained_model_signatures: + models[layernorm_model_name] = self._train_single_model( + model_name=layernorm_model_name, + df=input_layernorm_df, + feature_cols=["num_tokens"], + target_col=f"time_stats.{layernorm_model_name}.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=input_layernorm_context, + ) + trained_model_signatures.add(layernorm_signature) + logger.info(f"Trained {layernorm_model_name} for {cluster_type}") + + # Attention projections: attn_pre_proj, attn_post_proj, attn_rope + attn_proj_context = dict(training_context) + attn_proj_context["input_file"] = linear_ops_file + attn_proj_context["tensor_parallel_size"] = attn_tp_key + attn_proj_model_names = ["attn_pre_proj", "attn_post_proj", "attn_rope"] + for model_name in attn_proj_model_names: + model_signature = f"{model_name}_{attention_signature}" + if model_signature not in trained_model_signatures: + models[model_name] = self._train_single_model( + model_name=model_name, + df=attn_linear_ops_df, + feature_cols=["num_tokens"], + target_col=f"time_stats.{model_name}.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=attn_proj_context, + ) + trained_model_signatures.add(model_signature) + logger.info(f"Trained {model_name} for {cluster_type}") + + if is_target_embedded_mtp_enabled( + getattr(replica_config, "speculative_decoding_config", None) + ): + required_mtp_columns = ( + self._get_required_target_embedded_mtp_linear_op_columns() + ) + missing_mtp_columns = [ + col for col in required_mtp_columns if col not in attn_linear_ops_df.columns + ] + all_nan_mtp_columns = [ + col + for col in required_mtp_columns + if col in attn_linear_ops_df.columns + and attn_linear_ops_df[col].isna().all() + ] + if missing_mtp_columns or all_nan_mtp_columns: + raise ValueError( + "target-embedded MTP compute profiling columns are missing or all-NaN in " + f"{linear_ops_file}. " + f"Missing columns: {missing_mtp_columns}. " + f"All-NaN columns: {all_nan_mtp_columns}. " + "Re-run linear-op profiling with --include_target_embedded_mtp." + ) + for model_name in ["mtp_fusion_proj", "lm_head_linear"]: + model_signature = f"{model_name}_{attention_signature}" + if model_signature not in trained_model_signatures: + models[model_name] = self._train_single_model( + model_name=model_name, + df=attn_linear_ops_df, + feature_cols=["num_tokens"], + target_col=f"time_stats.{model_name}.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=attn_proj_context, + ) + trained_model_signatures.add(model_signature) + logger.info( + "Trained %s for %s (target-embedded MTP)", + model_name, + cluster_type, + ) + + model_config = replica_config.model_config + architecture_profile = _resolve_model_architecture_profile(model_config) + predictor_attention_extra_ops = ( + architecture_profile.predictor_attention_extra_ops + if architecture_profile is not None + else () + ) + for model_name in predictor_attention_extra_ops: + model_signature = f"{model_name}_{attention_signature}" + if model_signature not in trained_model_signatures: + target_col = f"time_stats.{model_name}.median" + if target_col not in attn_linear_ops_df.columns: + raise ValueError( + f"Architecture-profile operation '{model_name}' column '{target_col}' not found in profiling data. " + f"Ensure profiling was run with the selected model architecture profile. " + f"Available columns: {list(attn_linear_ops_df.columns)}" + ) + models[model_name] = self._train_single_model( + model_name=model_name, + df=attn_linear_ops_df, + feature_cols=["num_tokens"], + target_col=target_col, + execution_time_predictor_config=execution_time_predictor_config, + training_context=attn_proj_context, + ) + trained_model_signatures.add(model_signature) + logger.info("Trained architecture-profile %s for %s", model_name, cluster_type) + + # ========== Part 2: Attention core operations from attention.csv ========== + if not os.path.exists(attn_file): + raise FileNotFoundError(f"Attention input file {attn_file} not found") + + logger.info(f"Loading attention data from: {attn_file}") + attention_df = self._load_attention_df( + attn_file, + replica_config, + replica_scheduler_config, + cluster_type=cluster_type, + ) + training_context['input_file'] = attn_file + + # Family-aware attention-core training. Latent-MLA profiles carry six + # ``attn_mla_*`` operators with a structural layout the dense block cannot + # consume; route them through the MLA branch before the dense derive (which + # assumes dense feature columns such as ``prefill_chunk_size``). + if self._is_mla_family(replica_config.model_config): + attention_df = self._get_mla_attention_df_with_derived_features( + attention_df + ) + logger.info( + f"Loaded {len(attention_df)} rows for latent-MLA attention core training" + ) + models.update( + self._train_mla_attention_core_models( + attention_df=attention_df, + attention_signature=attention_signature, + cluster_type=cluster_type, + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + trained_model_signatures=trained_model_signatures, + ) + ) + trained_model_signatures.add(attention_signature) + return models + + attention_df = self._get_attention_df_with_derived_features(attention_df) + logger.info(f"Loaded {len(attention_df)} rows for attention core training") + measurement_type = self._active_measurement_type + dense_attention_model_names = get_enabled_predictor_metric_names( + DENSE_ATTENTION_FAMILY + ) + dense_attention_target_columns = dict( + zip( + dense_attention_model_names, + get_enabled_predictor_median_columns(DENSE_ATTENTION_FAMILY), + ) + ) + dense_attention_feature_columns = get_enabled_shared_predictor_feature_columns( + DENSE_ATTENTION_FAMILY + ) + + # Train kv_cache_save model + kv_cache_model_name = get_enabled_predictor_metric_name_by_role( + DENSE_ATTENTION_FAMILY, + AttentionOperatorRole.CACHE_WRITE, + ) + kv_cache_model_signature = f"{kv_cache_model_name}_{attention_signature}" + if kv_cache_model_signature not in trained_model_signatures: + kv_cache_feature_cols = list( + dense_attention_feature_columns[kv_cache_model_name] + ) + missing_cols = [ + col for col in kv_cache_feature_cols if col not in attention_df.columns + ] + if missing_cols: + raise ValueError( + f"Missing columns for {kv_cache_model_name} training: {missing_cols}. " + "Re-run attention profiling with mixed-batch metadata." + ) + models[kv_cache_model_name] = self._train_single_model( + model_name=kv_cache_model_name, + df=attention_df, + feature_cols=kv_cache_feature_cols, + target_col=dense_attention_target_columns[kv_cache_model_name], + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + persist_exact_lookup=True, + ) + trained_model_signatures.add(kv_cache_model_signature) + logger.info(f"Trained {kv_cache_model_name} for {cluster_type}") + + # Split data for prefill and decode. + # Mixed-batch prefill rows in attention_combined.csv use prefill_chunk_size=0, + # so standard prefill training must keep only rows with positive chunk size. + true_mixed_df = attention_df[attention_df["is_true_mixed_batch"]].copy() + standard_df = attention_df[~attention_df["is_true_mixed_batch"]].copy() + prefill_df = standard_df[~standard_df["is_decode"]].copy() + decode_df = standard_df[standard_df["is_decode"]].copy() + standard_prefill_df = pd.DataFrame() + if measurement_type in (MeasurementType.CUDA_EVENT, MeasurementType.DEVICE_EVENT): + if "prefill_chunk_size" not in prefill_df.columns: + raise ValueError( + "Missing required column 'prefill_chunk_size' in attention profiling data." + ) + standard_prefill_df = prefill_df[prefill_df["prefill_chunk_size"] > 0].copy() + + prefill_model_name = get_enabled_predictor_metric_name_by_role( + DENSE_ATTENTION_FAMILY, + AttentionOperatorRole.PREFILL_KERNEL, + ) + prefill_model_signature = f"{prefill_model_name}_{attention_signature}" + if prefill_model_signature not in trained_model_signatures: + if len(standard_prefill_df) == 0: + raise ValueError( + "No standard prefill rows (prefill_chunk_size > 0) found in eager attention profiling data." + ) + models[prefill_model_name] = self._train_single_model( + model_name=prefill_model_name, + df=standard_prefill_df, + feature_cols=list(dense_attention_feature_columns[prefill_model_name]), + target_col=dense_attention_target_columns[prefill_model_name], + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + ) + trained_model_signatures.add(prefill_model_signature) + logger.info(f"Trained {prefill_model_name} for {cluster_type}") + + decode_model_name = get_enabled_predictor_metric_name_by_role( + DENSE_ATTENTION_FAMILY, + AttentionOperatorRole.DECODE_KERNEL, + ) + decode_model_signature = f"{decode_model_name}_{attention_signature}" + if decode_model_signature not in trained_model_signatures: + if len(decode_df) == 0: + logger.info( + "Skipping eager %s training for %s - no standard decode rows", + decode_model_name, + cluster_type, + ) + else: + decode_feature_cols = list( + dense_attention_feature_columns[decode_model_name] + ) + missing_decode_cols = [ + col + for col in [ + *decode_feature_cols, + dense_attention_target_columns[decode_model_name], + ] + if col not in decode_df.columns + ] + if missing_decode_cols: + logger.info( + "Skipping eager %s training for %s - missing decode feature columns %s", + decode_model_name, + cluster_type, + missing_decode_cols, + ) + else: + models[decode_model_name] = self._train_single_model( + model_name=decode_model_name, + df=decode_df, + feature_cols=decode_feature_cols, + target_col=dense_attention_target_columns[decode_model_name], + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + ) + trained_model_signatures.add(decode_model_signature) + logger.info(f"Trained eager {decode_model_name} for {cluster_type}") + elif measurement_type == MeasurementType.KERNEL_ONLY: + decode_model_name = get_enabled_predictor_metric_name_by_role( + DENSE_ATTENTION_FAMILY, + AttentionOperatorRole.DECODE_KERNEL, + ) + decode_model_signature = f"{decode_model_name}_{attention_signature}" + if decode_model_signature not in trained_model_signatures: + if len(decode_df) == 0: + raise ValueError( + "No standard decode rows found in kernel-only attention profiling data." + ) + models[decode_model_name] = self._train_single_model( + model_name=decode_model_name, + df=decode_df, + feature_cols=list(dense_attention_feature_columns[decode_model_name]), + target_col=dense_attention_target_columns[decode_model_name], + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + ) + trained_model_signatures.add(decode_model_signature) + logger.info(f"Trained {decode_model_name} for {cluster_type}") + else: + raise ValueError(f"Unsupported measurement_type={measurement_type!r}") + + # ========== Part 3: Mixed-batch prefill model (optional, high-dimensional) ========== + # attn_prefill_mixed uses 12 features and requires on-demand prediction at runtime + # Check if profiling data contains mixed-batch features + mixed_batch_model_signature = f"attn_prefill_mixed_{attention_signature}" + if measurement_type in (MeasurementType.CUDA_EVENT, MeasurementType.DEVICE_EVENT) and mixed_batch_model_signature not in trained_model_signatures: + # Check for mixed-batch specific columns in the dataframe + required_mixed_features = self.ATTN_PREFILL_MIXED_FEATURES + has_mixed_batch_data = all(feat in prefill_df.columns for feat in required_mixed_features) + + if has_mixed_batch_data: + logger.info(f"Training attn_prefill_mixed with {len(required_mixed_features)} features for {cluster_type}") + + # Filter for mixed-prefill rows (exclude true mixed prefill+decode rows) + mixed_batch_df = prefill_df[ + prefill_df["is_mixed_batch"] | (prefill_df["batch_size"] > 1) + ].copy() + + if len(mixed_batch_df) > 0: + models["attn_prefill_mixed"] = self._train_single_model( + model_name="attn_prefill_mixed", + df=mixed_batch_df, + feature_cols=required_mixed_features, + target_col="time_stats.attn_prefill.median", # Same target column as attn_prefill + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + persist_exact_lookup=True, + ) + trained_model_signatures.add(mixed_batch_model_signature) + logger.info(f"Trained attn_prefill_mixed with {len(mixed_batch_df)} samples for {cluster_type}") + else: + logger.warning(f"No mixed-batch data (batch_size > 1) available for attn_prefill_mixed in {cluster_type}") + else: + missing_features = [f for f in required_mixed_features if f not in prefill_df.columns] + logger.info(f"Skipping attn_prefill_mixed for {cluster_type} - missing features: {missing_features}") + + decode_in_mixed_signature = f"attn_decode_in_mixed_{attention_signature}" + if measurement_type in (MeasurementType.CUDA_EVENT, MeasurementType.DEVICE_EVENT) and decode_in_mixed_signature not in trained_model_signatures: + required_decode_mixed_features = self.ATTN_DECODE_IN_MIXED_FEATURES + has_decode_mixed_data = all( + feat in true_mixed_df.columns for feat in required_decode_mixed_features + ) + if has_decode_mixed_data: + if len(true_mixed_df) > 0: + models["attn_decode_in_mixed"] = self._train_single_model( + model_name="attn_decode_in_mixed", + df=true_mixed_df, + feature_cols=required_decode_mixed_features, + target_col="time_stats.attn_decode.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + persist_exact_lookup=True, + ) + trained_model_signatures.add(decode_in_mixed_signature) + logger.info( + f"Trained attn_decode_in_mixed with {len(true_mixed_df)} samples for {cluster_type}" + ) + else: + logger.info( + f"Skipping attn_decode_in_mixed for {cluster_type} - no true mixed rows" + ) + else: + missing_features = [ + f for f in required_decode_mixed_features if f not in true_mixed_df.columns + ] + logger.info( + f"Skipping attn_decode_in_mixed for {cluster_type} - missing features: {missing_features}" + ) + + trained_model_signatures.add(attention_signature) + return models + + @staticmethod + + def _is_mla_family(model_config) -> bool: + """Return True when the model binds to the latent-MLA attention family.""" + if model_config is None: + return False + return ( + resolve_runtime_attention_family(model_config).family_id + == LATENT_MLA_ATTENTION_FAMILY.family_id + ) + + def _get_mla_attention_df_with_derived_features( + self, df: pd.DataFrame + ) -> pd.DataFrame: + """Derive latent-MLA attention features (normalize ``is_prefill`` to int). + + Mirrors the monolithic ``SklearnExecutionTimePredictor`` MLA early-return: + latent-MLA training keys on the imported structural columns directly and must + NOT add the dense ``num_tokens`` / ``prefill_chunk_size`` derived features. + """ + df_with_derived_features = df.copy() + if "is_prefill" in df_with_derived_features.columns: + df_with_derived_features["is_prefill"] = coerce_truthy_int( + df_with_derived_features["is_prefill"] + ) + return df_with_derived_features + + def _filter_mla_attention_df( + self, + df: pd.DataFrame, + file_path: str, + replica_config, + replica_scheduler_config, + ) -> pd.DataFrame: + """Filter an imported latent-MLA profile to the requested structural layout. + + Verbatim port of the monolithic ``_filter_mla_attention_df`` (adapted to the + shared-manager's per-cluster ``replica_config`` / ``replica_scheduler_config`` + instead of instance state). Fail-fast on missing structural columns or an empty + post-filter frame per §7 (no silent fallback). + """ + validate_attention_profiling_dataframe( + df, + LATENT_MLA_ATTENTION_FAMILY, + measurement_type=self._active_measurement_type, + ) + + model_config = replica_config.model_config + expected_values = { + "n_q_head": int(getattr(model_config, "num_q_heads")), + "n_kv_head": int( + model_config.get_runtime_num_kv_heads() + if hasattr(model_config, "get_runtime_num_kv_heads") + else 1 + ), + "head_size": int( + model_config.get_runtime_head_size() + if hasattr(model_config, "get_runtime_head_size") + else int(getattr(model_config, "kv_lora_rank")) + + int(getattr(model_config, "qk_rope_head_dim")) + ), + "qk_nope_head_dim": int(getattr(model_config, "qk_nope_head_dim")), + "qk_rope_head_dim": int(getattr(model_config, "qk_rope_head_dim")), + "qk_head_dim": int(model_config.get_qk_head_dim()), + "kv_lora_rank": int(getattr(model_config, "kv_lora_rank")), + "v_head_dim": int(getattr(model_config, "v_head_dim")), + "block_size": int(replica_scheduler_config.block_size), + "num_tensor_parallel_workers": int( + replica_config.attn_tensor_parallel_size + ), + } + missing_columns = [ + column for column in expected_values if column not in df.columns + ] + if missing_columns: + raise ValueError( + "MLA attention profiling data is missing structural columns: " + f"{missing_columns}. file={file_path}" + ) + + filtered = df.copy() + for column, expected_value in expected_values.items(): + filtered = filtered[filtered[column].astype(int) == expected_value] + + if filtered.empty: + raise ValueError( + "No MLA attention profiling rows remain after structural filtering. " + f"file={file_path}, expected={expected_values}" + ) + return filtered + + def _train_mla_attention_core_models( + self, + attention_df: pd.DataFrame, + attention_signature: str, + cluster_type: ClusterType, + execution_time_predictor_config, + training_context: Dict[str, Any], + trained_model_signatures: set, + ) -> Dict[str, BaseEstimator]: + """Train the six latent-MLA attention-core operators (training-only A2 fix). + + Mirrors the monolithic ``_train_mla_attention_layer_models`` (sparse-by-target + row filtering + exact-row memoization). The on-demand consumer pairs each + estimator's ``_frontier_exact_lookup`` with the same module-level builder, so + the disaggregation prediction path works unchanged once these models exist. + """ + model_names = list( + get_enabled_predictor_metric_names(LATENT_MLA_ATTENTION_FAMILY) + ) + target_columns = dict( + zip( + model_names, + get_enabled_predictor_median_columns(LATENT_MLA_ATTENTION_FAMILY), + ) + ) + feature_columns = get_enabled_shared_predictor_feature_columns( + LATENT_MLA_ATTENTION_FAMILY + ) + + models: Dict[str, BaseEstimator] = {} + for model_name in model_names: + model_signature = f"{model_name}_{attention_signature}" + if model_signature in trained_model_signatures: + continue + + feature_cols = list(feature_columns[model_name]) + target_col = target_columns[model_name] + required_columns = [*feature_cols, target_col] + missing_columns = [ + column + for column in required_columns + if column not in attention_df.columns + ] + all_nan_columns = [ + column + for column in required_columns + if column in attention_df.columns + and attention_df[column].isna().all() + ] + if missing_columns or all_nan_columns: + raise ValueError( + "MLA attention profiling data cannot train " + f"{model_name}." + f"\nMissing columns: {missing_columns}" + f"\nAll-NaN columns: {all_nan_columns}" + ) + + op_attention_df = attention_df.dropna(subset=[target_col]).copy() + if op_attention_df.empty: + raise ValueError( + "MLA attention profiling data cannot train " + f"{model_name}: target column {target_col!r} has no " + "observed timing rows." + ) + nan_feature_columns = [ + column + for column in feature_cols + if op_attention_df[column].isna().any() + ] + if nan_feature_columns: + raise ValueError( + "MLA attention profiling data cannot train " + f"{model_name}: feature columns contain NaN after " + f"target filtering: {nan_feature_columns}" + ) + + model = self._train_single_model( + model_name=model_name, + df=op_attention_df, + feature_cols=feature_cols, + target_col=target_col, + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + persist_exact_lookup=True, + ) + if not hasattr(model, "_frontier_exact_lookup"): + model._frontier_exact_lookup = _build_exact_feature_lookup( + op_attention_df, + feature_cols, + target_col, + ) + models[model_name] = model + trained_model_signatures.add(model_signature) + logger.info(f"Trained {model_name} for {cluster_type}") + + return models + + def _train_residual_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, + linear_ops_file: str, trained_model_signatures: set) -> Dict[str, BaseEstimator]: + """ + Train residual connection models for a cluster. + + This function handles residual connection operations in the Transformer layer: + - Residual add operation (from linear_op.csv): add + + Transformer layer context: + ... → Attention → [add] → LayerNorm → FFN/MoE → [add] → ... + + The residual add operation is used after both Attention and FFN blocks, + making it a common operation that serves both sub-layers. + """ + models = {} + + model_config = replica_config.model_config + + # RMSNorm: add is fused into layernorm, no separate add model needed + if model_config is not None and model_config.uses_fused_add_norm: + logger.info(f"Skipping residual add model training for {cluster_type} " + f"— model uses fused add+norm (RMSNorm)") + return models + + is_moe_model = model_config is not None and model_config.is_moe + tp_size = self._get_linear_op_tp_key( + "add", + cluster_type, + replica_config, + is_moe_model, + ) + + # Create a signature for this residual model configuration + residual_signature = f"residual_{replica_config.device}_{replica_config.model_name}_{tp_size}_family{self._measurement_family_name(self._active_measurement_type)}" + + if residual_signature in trained_model_signatures: + logger.info(f"Skipping residual models training for {cluster_type} - already trained with signature {residual_signature}") + return models + + if not os.path.exists(linear_ops_file): + raise FileNotFoundError(f"Linear ops input file {linear_ops_file} not found for residual models") + + logger.info(f"Loading linear ops data for residual models from: {linear_ops_file}") + linear_ops_df = self._load_linear_op_df(linear_ops_file, tp_size) + logger.info(f"Loaded {len(linear_ops_df)} rows for residual training") + + # Build training context for error messages + training_context = { + 'cluster_type': str(cluster_type), + 'device': replica_config.device, + 'model_name': replica_config.model_name, + 'tensor_parallel_size': tp_size, + 'input_file': linear_ops_file, + } + + # Train the residual add model + add_model_name = "add" + add_signature = f"{add_model_name}_{residual_signature}" + if add_signature not in trained_model_signatures: + models[add_model_name] = self._train_single_model( + model_name=add_model_name, + df=linear_ops_df, + feature_cols=["num_tokens"], + target_col=f"time_stats.{add_model_name}.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + ) + trained_model_signatures.add(add_signature) + logger.info(f"Trained {add_model_name} for {cluster_type}") + + # Mark this residual configuration as trained + trained_model_signatures.add(residual_signature) + return models + + def _train_pipeline_parallel_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, trained_model_signatures: set) -> Dict[str, BaseEstimator]: + """Train pipeline parallel communication models for a cluster.""" + models = {} + + _, _, _, send_recv_input_file, _, _ = self._get_input_files_for_config(replica_config, execution_time_predictor_config) + + pp_signature = f"send_recv_{replica_config.network_device}_{replica_config.num_pipeline_stages}_{replica_config.attn_tensor_parallel_size}_family{self._measurement_family_name(self._active_measurement_type)}" + + if pp_signature in trained_model_signatures: + logger.info(f"Skipping send_recv model training for {cluster_type} - already trained") + return models + + send_recv_df = self._load_send_recv_df(send_recv_input_file, replica_config) + send_recv_df = self._get_send_recv_df_with_derived_features(send_recv_df, replica_config) + + # Build training context for error messages + training_context = { + 'cluster_type': str(cluster_type), + 'device': replica_config.device, + 'model_name': replica_config.model_name, + 'pipeline_stages': replica_config.num_pipeline_stages, + 'tensor_parallel_size': replica_config.attn_tensor_parallel_size, + 'network_device': replica_config.network_device, + 'input_file': send_recv_input_file, + } + + models["send_recv"] = self._train_single_model( + model_name="send_recv", + df=send_recv_df, + feature_cols=["num_tokens"], + target_col="time_stats.send_recv.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + ) + + trained_model_signatures.add(pp_signature) + return models + + def _train_tensor_parallel_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, use_attn_tp: bool, trained_model_signatures: set) -> Dict[str, BaseEstimator]: + """Train tensor parallel communication models for a cluster.""" + models = {} + + _, _, all_reduce_input_file, _, _, _ = self._get_input_files_for_config(replica_config, execution_time_predictor_config) + + # Use different tensor parallel size based on cluster type + tp_size = replica_config.attn_tensor_parallel_size if use_attn_tp else replica_config.moe_tensor_parallel_size + + tp_signature = f"all_reduce_{replica_config.network_device}_{tp_size}_family{self._measurement_family_name(self._active_measurement_type)}" + + if tp_signature in trained_model_signatures: + logger.info(f"Skipping all_reduce model training for {cluster_type} - already trained") + return models + + # 添加详细的上下文信息 + training_context = { + 'cluster_type': cluster_type, + 'device': replica_config.device, + 'model_name': replica_config.model_name, + 'tensor_parallel_size': tp_size, + 'network_device': replica_config.network_device, + 'input_file': all_reduce_input_file, + 'use_attn_tp': use_attn_tp + } + + logger.info(f"Loading all_reduce data for {cluster_type}: file={all_reduce_input_file}, tp_size={tp_size}") + + all_reduce_df = self._load_all_reduce_df(all_reduce_input_file, replica_config, tp_size) + logger.info(f"Loaded {len(all_reduce_df)} rows for all_reduce training") + + all_reduce_df = self._get_all_reduce_df_with_derived_features(all_reduce_df, replica_config) + logger.info(f"After feature engineering: {len(all_reduce_df)} rows") + + models["all_reduce"] = self._train_single_model( + model_name="all_reduce", + df=all_reduce_df, + feature_cols=["num_tokens"], + target_col="time_stats.all_reduce.median", + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context + ) + + trained_model_signatures.add(tp_signature) + return models + + def _train_cpu_overhead_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, trained_model_signatures: set) -> Dict[str, BaseEstimator]: + """Train CPU overhead models for a cluster.""" + models = {} + + if execution_time_predictor_config.skip_cpu_overhead_modeling: + return models + + _, _, _, _, cpu_overhead_input_file, _ = self._get_input_files_for_config(replica_config, execution_time_predictor_config) + + cpu_signature = f"cpu_overhead_{replica_config.network_device}_{replica_config.model_name}_{replica_config.attn_tensor_parallel_size}_family{self._measurement_family_name(self._active_measurement_type)}" + + if cpu_signature in trained_model_signatures: + logger.info(f"Skipping CPU overhead models training for {cluster_type} - already trained") + return models + + cpu_overhead_df = self._load_cpu_overhead_df(cpu_overhead_input_file, replica_config) + if cpu_overhead_df.empty: + logger.warning( + "Skipping CPU overhead model training for cluster %s due to missing/empty CPU overhead profiling data. file=%s", + cluster_type, + cpu_overhead_input_file, + ) + trained_model_signatures.add(cpu_signature) + return models + + # Build training context for error messages + training_context = { + 'cluster_type': str(cluster_type), + 'device': replica_config.device, + 'model_name': replica_config.model_name, + 'tensor_parallel_size': replica_config.attn_tensor_parallel_size, + 'network_device': replica_config.network_device, + 'input_file': cpu_overhead_input_file, + } + + model_names = [ + "schedule", + "sampler_e2e", + "prepare_inputs_e2e", + "process_model_outputs", + "ray_comm_time", + ] + + for model_name in model_names: + target_col = "ray_comm_time_mean" if model_name == "ray_comm_time" else f"{model_name}_median" + + model_signature = f"{model_name}_{cpu_signature}" + if model_signature not in trained_model_signatures: + feature_cols = [ + "batch_size", + "num_prefill_tokens", + "num_decode_tokens", + ] + model = self._train_single_model( + model_name=model_name, + df=cpu_overhead_df, + feature_cols=feature_cols, + target_col=target_col, + execution_time_predictor_config=execution_time_predictor_config, + training_context=training_context, + persist_exact_lookup=True, + ) + if not hasattr(model, "_frontier_exact_lookup"): + model._frontier_exact_lookup = _build_exact_feature_lookup( + cpu_overhead_df, + feature_cols, + target_col, + ) + models[model_name] = model + trained_model_signatures.add(model_signature) + + trained_model_signatures.add(cpu_signature) + return models + + def _train_single_model( + self, + model_name: str, + df: pd.DataFrame, + feature_cols: List[str], + target_col: str, + execution_time_predictor_config, + training_context: Optional[Dict[str, Any]] = None, + persist_exact_lookup: bool = True, + layer_contract: Optional[ResolvedLayerContract] = None, + ) -> BaseEstimator: + """Train a single model with given data and configuration.""" + layer_contract, training_context = _normalize_layer_contract_context( + training_context, + explicit_layer_contract=layer_contract, + ) + if len(df) == 0: + # 提供详细的错误信息，以便调试 + context_info = "" + if training_context: + context_info = f""" + Training Context: + - Cluster Type: {training_context.get('cluster_type', 'Unknown')} + - Device: {training_context.get('device', 'Unknown')} + - Model Name: {training_context.get('model_name', 'Unknown')} + - Pipeline Stages: {training_context.get('pipeline_stages', 'Unknown')} + - Network Device: {training_context.get('network_device', 'Unknown')} + - Tensor Parallel Size: {training_context.get('tensor_parallel_size', 'Unknown')} + - Input File: {training_context.get('input_file', 'Unknown')} + - Block Size: {training_context.get('block_size', 'Unknown')} + - Feature Columns: {feature_cols} + - Target Column: {target_col} + """ + + raise Exception(f"Training data for model {model_name} is empty.{context_info}") + + required_cols = feature_cols + [target_col] + nan_row_mask = df[required_cols].isna().any(axis=1) + nan_row_count = int(nan_row_mask.sum()) + if nan_row_count > 0: + logger.warning( + "Dropping %d/%d rows with NaN feature/target values before training %s " + "(target=%s).", + nan_row_count, + len(df), + model_name, + target_col, + ) + df = df.loc[~nan_row_mask].copy() + if len(df) == 0: + raise ValueError( + f"Training data for model {model_name} is empty after dropping NaN rows " + f"(target={target_col})." + ) + + profiling_precision = self._get_profiling_precision_from_df(df) + measurement_type = self._validate_active_measurement_type(df) + hash_args = ( + model_name, + df, + execution_time_predictor_config, + profiling_precision, + measurement_type, + ) + model_hash = self._get_model_hash( + *hash_args, + **_layer_contract_kwargs(layer_contract), + ) + cached_model = self._load_model_from_cache(model_name, model_hash) + if cached_model is not None: + if layer_contract is not None: + requested_identity = _serialize_selected_layer_cache_identity( + layer_contract + ) + if requested_identity is None: + raise ValueError( + "resolved layer contract did not produce a cache identity" + ) + self._validate_cached_layer_cache_identity( + model_name=model_name, + model=cached_model, + requested_identity=requested_identity, + ) + if persist_exact_lookup: + self._ensure_exact_lookup_metadata( + model_name=model_name, + model_hash=model_hash, + model=cached_model, + df=df, + feature_cols=feature_cols, + target_col=target_col, + ) + self._store_model_precision( + model_name, + profiling_precision, + cached_model, + **_layer_contract_kwargs(layer_contract), + ) + return cached_model + + # ============================================================ + # CACHE MISS: Model not found in cache + # ============================================================ + # When running in production mode (non-dummy mode), we expect all models + # to be pre-trained using the standalone training module and cached. + # If a model is not found in cache, it indicates a configuration mismatch + # or missing profiling/training step. + # + # To train models, use the standalone training workflow: + # 1. Run profiling: tests/test_pd_af_profiling.sh + # 2. Run training: tests/test_pd_af_training.sh + # 3. Run simulation: tests/test_small_scale_pd_af_disaggregation_cluster_parallel.sh + # ============================================================ + + error_msg = f""" + ❌ MODEL CACHE MISS ERROR ❌ + + Model '{model_name}' with hash '{model_hash}' not found in cache directory: {self._cache_dir} + + Configuration Details: + - Model Name: {model_name} + - Cache Hash: {model_hash} + - Cache Directory: {self._cache_dir} + - Expected Cache File: {self._cache_dir}/{model_name}_{model_hash}.pkl + """ + + if training_context: + error_msg += f""" + Training Context: + - Cluster Type: {training_context.get('cluster_type', 'Unknown')} + - Device: {training_context.get('device', 'Unknown')} + - Model Name: {training_context.get('model_name', 'Unknown')} + - Tensor Parallel Size: {training_context.get('tensor_parallel_size', 'Unknown')} + - Expert Parallel Size: {training_context.get('moe_expert_parallel_size', 'N/A')} + - Input File: {training_context.get('input_file', 'Unknown')} + - Feature Columns: {feature_cols} + - Target Column: {target_col} + """ + + error_msg += f""" + + ⚠️ REQUIRED ACTION ⚠️ + + This error indicates that the required model has not been pre-trained. + Please follow the complete workflow: + + ============================================================ + + NOTE: Real-time training is TEMPORARILY ENABLED for cache generation. + """ + + logger.warning(error_msg) + logger.info(f"CACHE MISS: Training model '{model_name}' with hash '{model_hash}' in real-time...") + + # ============================================================ + # TEMPORARILY ENABLED: Real-time training code + # ============================================================ + # This code performs real-time model training during simulation + # initialization to generate missing cache files. + # ============================================================ + + estimator, grid_search_params = self._create_estimator_and_params(execution_time_predictor_config) + + cv = min(execution_time_predictor_config.k_fold_cv_splits, len(df)) if len(df) >= 2 else 2 + + grid_search = GridSearchCV( + estimator=estimator, + param_grid=grid_search_params, + scoring=self._get_scorer(), + cv=cv, + n_jobs=execution_time_predictor_config.num_training_job_threads, + ) + + X, y = df[feature_cols], df[target_col] + grid_search.fit(X, y) + score = grid_search.score(X, y) + + logger.info(f"✓ Trained model {model_name} with MAPE {-score}%") + + best_estimator = grid_search.best_estimator_ + # Persist feature metadata for runtime on-demand prediction (e.g., moe_grouped_gemm load imbalance mode). + setattr(best_estimator, "_frontier_feature_names", list(feature_cols)) + setattr(best_estimator, "_frontier_target_col", target_col) + # Tie the trained estimator to its cache hash so prediction caches can include model identity. + setattr(best_estimator, "_frontier_model_hash", model_hash) + if layer_contract is not None: + self._model_contract_identity(best_estimator, layer_contract) + + if persist_exact_lookup: + setattr( + best_estimator, + "_frontier_exact_lookup", + _build_exact_feature_lookup(df, feature_cols, target_col), + ) + + self._store_model_in_cache(model_name, model_hash, best_estimator) + self._store_model_precision( + model_name, + profiling_precision, + best_estimator, + **_layer_contract_kwargs(layer_contract), + ) + return best_estimator diff --git a/frontier/execution_time_predictor/prediction_model_identity.py b/frontier/execution_time_predictor/prediction_model_identity.py new file mode 100644 index 00000000..7153053a --- /dev/null +++ b/frontier/execution_time_predictor/prediction_model_identity.py @@ -0,0 +1,324 @@ +"""Model, operator-family and layer-contract identity helpers. + +These pure functions answer identity questions the predictor stack asks in +several places: which operator family a profiling model name belongs to, +which architecture profile a replica resolves to, and what the cache +identity of a resolved layer contract is. They hold no state, so they sit +in a leaf module every predictor component can import. +""" + +import hashlib +import json +import pandas as pd + +from frontier.model_architectures import ( + LayerKind, + ModelArchitectureProfile, + ResolvedLayerContract, + get_model_architecture_profile, +) +from frontier.moe_gating_runtime import get_moe_gating_base_model_name +from frontier.operators.families import ( + MOE_FAMILY, + get_family_profiling_names, + get_operator_family, +) +from frontier.operators.typed_contracts import ( + TYPED_OPERATOR_CONTRACTS_COLUMN, + matches_resolved_layer_contract, +) +from typing import Any, Dict, List, Mapping, Optional, Tuple, cast + +MIGRATION_HELP_COMMAND = ( + "python -m frontier.profiling.migrate_csv_metadata --help" +) +def _get_moe_family_model_names() -> List[str]: + return list(get_family_profiling_names(MOE_FAMILY)) + + +def _get_moe_family_operator_by_model_name(model_name: str): + moe_ops = { + operator.profiling_name(): operator + for operator in MOE_FAMILY.profiling_ops() + } + if model_name not in moe_ops: + raise ValueError(f"Unsupported MoE op: {model_name}") + return moe_ops[model_name] + + +def _get_moe_gating_family_model_names() -> List[str]: + return [ + operator.profiling_name() + for operator in MOE_FAMILY.profiling_ops() + if operator.precision_name() == "moe_gating" + ] + + +def _get_prefill_hot_moe_gating_model_names() -> List[str]: + return [ + f"{model_name}__prefill_hot" + for model_name in _get_moe_gating_family_model_names() + ] + + +def _resolve_model_architecture_profile( + model_config: Any, + *, + allow_generic: bool = False, +) -> Optional[ModelArchitectureProfile]: + if model_config is None: + return None + getter = getattr(model_config, "get_model_architecture_profile", None) + if callable(getter): + return cast(Optional[ModelArchitectureProfile], getter()) + + # Lightweight test/config adapters that predate the typed contract do not + # expose a profile accessor or typed widths. Keep those callers on the + # scalar compatibility path instead of treating the generic fallback as a + # complete typed declaration. An explicit profile or typed width opts the + # adapter into strict profile-owned resolution. + typed_fields = ( + "model_architecture_profile", + "dense_mlp_hidden_dim", + "routed_mlp_hidden_dim", + "share_expert_dim", + ) + if not allow_generic and not any( + getattr(model_config, field_name, None) is not None + for field_name in typed_fields + ): + return None + return get_model_architecture_profile(model_config) + + +def _resolve_model_architecture_profile_id(model_config) -> str: + architecture_profile = _resolve_model_architecture_profile(model_config) + if architecture_profile is None: + return "generic" + return architecture_profile.profile_id + + +def _resolve_profile_typed_family_for_query( + architecture_profile: ModelArchitectureProfile, + op_name: str, +) -> Optional[Tuple[str, LayerKind]]: + """Resolve a query to the profile-owned typed operator family.""" + + if not isinstance(op_name, str) or not op_name: + raise ValueError("typed operator query name must be a non-empty string") + matches: list[Tuple[str, LayerKind]] = [] + for layer_contract in architecture_profile.layer_contracts: + for family_id in layer_contract.operator_family_ids: + family = get_operator_family(family_id) + if any( + op_name == operator.name + or op_name == operator.profiling_name() + for operator in family.operators + ): + matches.append((family_id, layer_contract.layer_kind)) + if len(matches) > 1: + raise ValueError( + f"Operator query {op_name!r} belongs to multiple typed layer families: " + f"{sorted(family_id for family_id, _ in matches)}" + ) + return matches[0] if matches else None +def _serialize_selected_layer_cache_identity( + layer_contract: Optional[ResolvedLayerContract], +) -> Optional[str]: + """Serialize the selected semantic domain used by a model cache. + + Physical ``layer_id`` and producer-side domain envelopes do not identify a + trained estimator. Keep only the selected fields that affect estimator + admission and reuse, in deterministic JSON form. + """ + + if layer_contract is None: + return None + if not isinstance(layer_contract, ResolvedLayerContract): + raise TypeError( + "layer_contract must be a ResolvedLayerContract when provided" + ) + metadata = layer_contract.typed_metadata_identity() + selected_fields = ( + "profile_id", + "operator_family_id", + "layer_kind", + "dimension_source", + "effective_ffn_width", + "tensor_parallel_mode", + "expert_parallel_mode", + "selected_expert_parallel_size", + "selected_tensor_parallel_size", + "selected_padded_ffn_width", + ) + return json.dumps( + {field_name: metadata[field_name] for field_name in selected_fields}, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ) + + +def _get_contract_hash( + layer_contract: Optional[ResolvedLayerContract], +) -> str: + """Return a short deterministic hash for an optional selected contract.""" + + identity = _serialize_selected_layer_cache_identity(layer_contract) + if identity is None: + return "none" + return hashlib.sha256(identity.encode("utf-8")).hexdigest()[:16] + + +def _validate_typed_parallel_selection( + layer_contract: ResolvedLayerContract, + *, + tensor_parallel_size: Optional[int] = None, + expert_parallel_size: Optional[int] = None, +) -> None: + """Validate explicit loader selectors against a resolved contract.""" + + if not isinstance(layer_contract, ResolvedLayerContract): + raise TypeError("layer_contract must be a ResolvedLayerContract") + contract_tp = layer_contract.tensor_parallel_size + if ( + contract_tp is not None + and tensor_parallel_size is not None + and contract_tp != tensor_parallel_size + ): + raise ValueError( + f"typed layer contract TP {contract_tp} conflicts with " + f"tensor_parallel_size {tensor_parallel_size}" + ) + contract_ep = layer_contract.expert_parallel_size + if ( + contract_ep is not None + and expert_parallel_size is not None + and contract_ep != expert_parallel_size + ): + raise ValueError( + f"typed layer contract EP {contract_ep} conflicts with " + f"expert_parallel_size {expert_parallel_size}" + ) + + +def _typed_row_matches_contract( + raw_contracts: Any, + layer_contract: ResolvedLayerContract, + *, + operator_name: Optional[str], +) -> bool: + """Match one parsed or serialized row to its exact typed operator contract.""" + + if not isinstance(operator_name, str) or not operator_name: + raise ValueError( + "typed profiling loading requires a non-empty operator_name when " + f"the canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" + ) + if not isinstance(layer_contract.operator_family_id, str) or not layer_contract.operator_family_id: + raise ValueError( + "typed profiling loading requires a layer contract with an operator family id" + ) + return matches_resolved_layer_contract( + raw_contracts, + layer_contract, + operator_name=operator_name, + ) + + +def _normalize_layer_contract_context( + training_context: Optional[Mapping[str, Any]], + explicit_layer_contract: Optional[ResolvedLayerContract] = None, +) -> Tuple[Optional[ResolvedLayerContract], Dict[str, Any]]: + """Resolve one contract and keep every context representation consistent.""" + + context = dict(training_context or {}) + context_contract = context.get("layer_contract") + if context_contract is not None and not isinstance( + context_contract, ResolvedLayerContract + ): + raise TypeError( + "training_context['layer_contract'] must be a ResolvedLayerContract" + ) + if explicit_layer_contract is not None and not isinstance( + explicit_layer_contract, ResolvedLayerContract + ): + raise TypeError("layer_contract must be a ResolvedLayerContract") + + if context_contract is not None and explicit_layer_contract is not None: + if not context_contract.is_semantically_equivalent(explicit_layer_contract): + raise ValueError( + "conflicting layer_contract values were provided through the " + "explicit argument and training_context" + ) + + resolved_contract = explicit_layer_contract or context_contract + context_identity = context.get("layer_contract_identity") + if context_identity is not None and not isinstance(context_identity, str): + raise TypeError( + "training_context['layer_contract_identity'] must be a string" + ) + selected_identity = _serialize_selected_layer_cache_identity(resolved_contract) + if context_identity is not None and context_identity != selected_identity: + raise ValueError( + "training_context['layer_contract_identity'] does not match the " + "supplied layer_contract" + ) + if resolved_contract is None: + return None, context + + context["layer_contract"] = resolved_contract + context["layer_contract_identity"] = selected_identity + context["layer_kind"] = resolved_contract.layer_kind.value + context["effective_ffn_width"] = resolved_contract.effective_ffn_width + context["tensor_parallel_mode"] = resolved_contract.tensor_parallel_mode.value + context["expert_parallel_mode"] = resolved_contract.expert_parallel_mode.value + return resolved_contract, context + + +def _add_layer_contract_to_training_context( + training_context: Mapping[str, Any], + layer_contract: Optional[ResolvedLayerContract], +) -> Dict[str, Any]: + """Copy a training context and attach a resolved typed contract.""" + + _, context = _normalize_layer_contract_context( + training_context, + explicit_layer_contract=layer_contract, + ) + return context + + +def _layer_contract_kwargs( + layer_contract: Optional[ResolvedLayerContract], + *, + operator_name: Optional[str] = None, +) -> Dict[str, Any]: + """Return typed keyword arguments only for an opted-in contract path.""" + + if layer_contract is None: + return {} + kwargs: Dict[str, Any] = {"layer_contract": layer_contract} + if operator_name is not None: + kwargs["operator_name"] = operator_name + return kwargs +def _is_moe_gating_family_model_name(model_name: str) -> bool: + base_model_name = get_moe_gating_base_model_name(model_name) + return _get_moe_family_operator_by_model_name( + base_model_name + ).precision_name() == "moe_gating" +def _build_exact_feature_lookup( + df: pd.DataFrame, + feature_cols: List[str], + target_col: str, +) -> Dict[Tuple[float, ...], float]: + """Build exact profiling-row lookups before falling back to regression.""" + if df.empty: + return {} + grouped = df.groupby(feature_cols, dropna=False)[target_col].mean() + lookup: Dict[Tuple[float, ...], float] = {} + for key, value in grouped.items(): + key_tuple = key if isinstance(key, tuple) else (key,) + lookup[tuple(float(item) for item in key_tuple)] = float(value) + return lookup diff --git a/frontier/execution_time_predictor/prediction_model_registry.py b/frontier/execution_time_predictor/prediction_model_registry.py new file mode 100644 index 00000000..bc250b85 --- /dev/null +++ b/frontier/execution_time_predictor/prediction_model_registry.py @@ -0,0 +1,587 @@ +"""The trained-model registry, its cache keys and the persistent cache. + +A trained estimator is identified by the configuration that produced it, the +profiling rows it saw, its precision and measurement family, and the layer +contract it covers. These methods compute that identity, store and look up +estimators under it, and read and write the on-disk cache. +""" + +import hashlib +import os +import pandas as pd +import pickle + +from fasteners import InterProcessReaderWriterLock +from frontier.execution_time_predictor.cache_io import atomic_pickle_dump +from frontier.execution_time_predictor.prediction_model_identity import ( + MIGRATION_HELP_COMMAND, + _build_exact_feature_lookup, + _resolve_model_architecture_profile, + _resolve_profile_typed_family_for_query, + _serialize_selected_layer_cache_identity, +) +from frontier.logger import init_logger +from frontier.model_architectures import LayerKind, ResolvedLayerContract +from frontier.moe_gating_runtime import get_moe_gating_base_model_name +from frontier.types import ClusterType, MeasurementType +from sklearn.base import BaseEstimator +from typing import Any, Dict, List, Mapping, Optional, Tuple + + +logger = init_logger(__name__) + + +class PredictionModelRegistry: + """Trained-model identity, in-memory registry and persistent cache.""" + + def _get_hash_relevant_config(self, config) -> Dict[str, Any]: + """ + Extract only the configuration parameters that affect model performance. + + Parameters that should be included: + - Profiling data paths (determine input data source) + - Prediction range parameters (determine prediction cache scope) + - Performance adjustment parameters (affect predicted values) + - ML hyperparameters (affect model structure) + + Parameters that should be excluded: + - Training process parameters (k_fold_cv_splits, num_training_job_threads) + - Runtime configuration (no_cache, skip_cpu_overhead_modeling, enable_dummy_mode, dummy_execution_time_ms) + """ + hash_relevant_params = { + # Category 1: Profiling data paths + 'linear_op_input_file': config.linear_op_input_file, + 'atten_input_file': config.atten_input_file, + 'all_reduce_input_file': config.all_reduce_input_file, + 'send_recv_input_file': config.send_recv_input_file, + 'moe_input_file': config.moe_input_file, + 'linear_op_kernel_only_input_file': config.linear_op_kernel_only_input_file, + 'atten_kernel_only_input_file': config.atten_kernel_only_input_file, + 'moe_kernel_only_input_file': config.moe_kernel_only_input_file, + 'cpu_overhead_input_file': config.cpu_overhead_input_file, + 'cpu_overhead_kernel_only_input_file': getattr( + config, + 'cpu_overhead_kernel_only_input_file', + config.cpu_overhead_input_file, + ), + + # Category 2: Prediction range parameters + 'kv_cache_prediction_granularity': config.kv_cache_prediction_granularity, + 'prediction_max_prefill_chunk_size': config.prediction_max_prefill_chunk_size, + 'prediction_max_batch_size': config.prediction_max_batch_size, + 'prediction_max_tokens_per_request': config.prediction_max_tokens_per_request, + + # Category 3: Performance adjustment parameters + 'attention_decode_batching_overhead_fraction': config.attention_decode_batching_overhead_fraction, + 'attention_prefill_batching_overhead_fraction': config.attention_prefill_batching_overhead_fraction, + 'attn_pre_proj_calibration_scale': config.attn_pre_proj_calibration_scale, + 'prefill_phase_attn_pre_proj_calibration_scale': config.prefill_phase_attn_pre_proj_calibration_scale, + 'attn_post_proj_calibration_scale': config.attn_post_proj_calibration_scale, + 'prefill_phase_attn_post_proj_calibration_scale': config.prefill_phase_attn_post_proj_calibration_scale, + 'attn_decode_calibration_scale': config.attn_decode_calibration_scale, + 'attn_decode_in_mixed_calibration_scale': config.attn_decode_in_mixed_calibration_scale, + 'late_decode_attn_decode_calibration_scale': config.late_decode_attn_decode_calibration_scale, + 'attn_kv_cache_save_calibration_scale': config.attn_kv_cache_save_calibration_scale, + 'prefill_phase_attn_kv_cache_save_calibration_scale': config.prefill_phase_attn_kv_cache_save_calibration_scale, + 'mlp_up_proj_calibration_scale': config.mlp_up_proj_calibration_scale, + 'prefill_phase_mlp_up_proj_calibration_scale': config.prefill_phase_mlp_up_proj_calibration_scale, + 'mlp_down_proj_calibration_scale': config.mlp_down_proj_calibration_scale, + 'decode_phase_mlp_down_proj_calibration_scale': config.decode_phase_mlp_down_proj_calibration_scale, + 'nccl_cpu_launch_overhead_ms': config.nccl_cpu_launch_overhead_ms, + 'nccl_cpu_skew_overhead_per_device_ms': config.nccl_cpu_skew_overhead_per_device_ms, + } + + # Category 4: ML Hyperparameters (type-specific) + if hasattr(config, 'num_estimators'): # Random Forest + hash_relevant_params['num_estimators'] = config.num_estimators + hash_relevant_params['max_depth'] = config.max_depth + hash_relevant_params['min_samples_split'] = config.min_samples_split + elif hasattr(config, 'polynomial_degree'): # Linear Regression + hash_relevant_params['polynomial_degree'] = config.polynomial_degree + hash_relevant_params['polynomial_include_bias'] = config.polynomial_include_bias + hash_relevant_params['polynomial_interaction_only'] = config.polynomial_interaction_only + hash_relevant_params['fit_intercept'] = config.fit_intercept + + return hash_relevant_params + + def _get_model_hash( + self, + model_name: str, + df: pd.DataFrame, + execution_time_predictor_config, + profiling_precision: str, + measurement_type: MeasurementType, + layer_contract: Optional[ResolvedLayerContract] = None, + ) -> str: + """ + Calculate hash for model caching based on configuration and data. + + Hash is calculated from: + 1. Hash-relevant configuration parameters (excluding runtime/training process params) + 2. Model name + 3. DataFrame content hash + + This ensures that only changes to parameters that affect model performance + will invalidate the cache. + """ + # Extract only hash-relevant parameters + hash_relevant_config = self._get_hash_relevant_config(execution_time_predictor_config) + config_str = str(sorted(hash_relevant_config.items())) # Sort for deterministic ordering + + # Calculate DataFrame hash + df_hash_str = hashlib.md5(df.to_json().encode("utf-8")).hexdigest() + + selected_identity = _serialize_selected_layer_cache_identity(layer_contract) + contract_component = ( + f"_{selected_identity}" if selected_identity is not None else "" + ) + + # Combine all components. The selected semantic domain is part of a + # typed key; physical layer occurrence is intentionally absent. + combined_str = ( + f"{config_str}_{model_name}_{df_hash_str}_{profiling_precision}_" + f"{measurement_type.value}{contract_component}" + ) + hash_value = hashlib.md5(combined_str.encode("utf-8")).hexdigest()[0:8] + + # Debug output for hash calculation + if model_name == "attn_pre_proj": + logger.info(f"[DEBUG] Hash calculation for {model_name}:") + logger.info(f" - DataFrame shape: {df.shape}") + logger.info(f" - DataFrame hash: {df_hash_str[:16]}...") + logger.info(f" - Hash-relevant config keys: {sorted(hash_relevant_config.keys())}") + logger.info(f" - Final hash: {hash_value}") + + return hash_value + + def _get_profiling_precision_from_df(self, df: pd.DataFrame) -> str: + """Extract profiling precision from DataFrame. + + FAIL-FAST: Raises ValueError if profiling_precision column is missing or invalid. + This enforces strict metadata requirements and prevents silent fallbacks. + """ + if "profiling_precision" not in df.columns: + raise ValueError( + "profiling_precision column is missing from profiling data. " + f"Run '{MIGRATION_HELP_COMMAND}' to add required metadata columns to legacy CSV files." + ) + + precision_values = df["profiling_precision"].dropna().unique().tolist() + if not precision_values: + raise ValueError("profiling_precision column is empty") + if len(precision_values) > 1: + raise ValueError( + f"Multiple profiling_precision values found: {precision_values}" + ) + return str(precision_values[0]).upper() + + def _get_measurement_type_from_df(self, df: pd.DataFrame) -> MeasurementType: + if "measurement_type" not in df.columns: + raise ValueError( + "measurement_type column is missing from profiling data. " + f"Run '{MIGRATION_HELP_COMMAND}' to add required metadata columns to legacy CSV files." + ) + + measurement_values = df["measurement_type"].dropna().unique().tolist() + if not measurement_values: + raise ValueError("measurement_type column is empty") + if len(measurement_values) > 1: + raise ValueError( + f"Multiple measurement_type values found: {measurement_values}" + ) + return MeasurementType.from_string(str(measurement_values[0])) + + def _validate_active_measurement_type(self, df: pd.DataFrame) -> MeasurementType: + measurement_type = self._get_measurement_type_from_df(df) + if measurement_type != self._active_measurement_type: + raise ValueError( + f"measurement_type mismatch: expected {self._active_measurement_type.value} " + f"but found {measurement_type.value}." + ) + return measurement_type + + @staticmethod + + def _validate_cached_layer_cache_identity( + *, + model_name: str, + model: BaseEstimator, + requested_identity: str, + ) -> None: + """Reject a typed cache entry whose selected domain does not match.""" + + cached_identity = getattr(model, "_frontier_layer_cache_identity", None) + if cached_identity is None: + raise ValueError( + f"Cached model {model_name!r} is missing selected layer cache identity" + ) + if not isinstance(cached_identity, str): + raise ValueError( + f"Cached model {model_name!r} has an invalid selected layer cache " + f"identity of type {type(cached_identity).__name__}" + ) + if cached_identity != requested_identity: + raise ValueError( + f"Cached model {model_name!r} selected layer cache identity mismatch: " + f"cached={cached_identity!r}, requested={requested_identity!r}" + ) + + @staticmethod + + def _model_contract_identity( + model: BaseEstimator, + layer_contract: Optional[ResolvedLayerContract], + ) -> Optional[str]: + """Attach and return the selected semantic identity for a model.""" + + requested_identity = _serialize_selected_layer_cache_identity(layer_contract) + attached_identity = getattr(model, "_frontier_layer_cache_identity", None) + if attached_identity is not None and not isinstance(attached_identity, str): + raise TypeError( + "_frontier_layer_cache_identity must be a string when present" + ) + if ( + requested_identity is not None + and attached_identity is not None + and requested_identity != attached_identity + ): + raise ValueError( + "model selected layer cache identity conflicts with the requested " + "contract" + ) + identity = requested_identity or attached_identity + if identity is not None: + setattr(model, "_frontier_layer_cache_identity", identity) + return identity + + def _contract_model_registry( + self, family_name: str + ) -> Dict[Tuple[str, Optional[str]], BaseEstimator]: + registry_attr = { + "eager": "_trained_models_eager_by_contract", + "device_event": "_trained_models_device_event_by_contract", + "kernel_only": "_trained_models_kernel_only_by_contract", + }.get(family_name) + if registry_attr is None: + raise ValueError(f"Unsupported family_name={family_name!r}") + return getattr(self, registry_attr) + + def _contract_precision_registry( + self, family_name: str + ) -> Dict[str, Dict[Tuple[str, Optional[str]], BaseEstimator]]: + registry_attr = { + "eager": "_models_by_precision_eager_by_contract", + "device_event": "_models_by_precision_device_event_by_contract", + "kernel_only": "_models_by_precision_kernel_only_by_contract", + }.get(family_name) + if registry_attr is None: + raise ValueError(f"Unsupported family_name={family_name!r}") + return getattr(self, registry_attr) + + def _legacy_model_registry(self, family_name: str) -> Dict[str, BaseEstimator]: + registry_attr = { + "eager": "_trained_models_eager", + "device_event": "_trained_models_device_event", + "kernel_only": "_trained_models_kernel_only", + }.get(family_name) + if registry_attr is None: + raise ValueError(f"Unsupported family_name={family_name!r}") + return getattr(self, registry_attr) + + def _legacy_precision_registry( + self, family_name: str + ) -> Dict[str, Dict[str, BaseEstimator]]: + registry_attr = { + "eager": "_models_by_precision_eager", + "device_event": "_models_by_precision_device_event", + "kernel_only": "_models_by_precision_kernel_only", + }.get(family_name) + if registry_attr is None: + raise ValueError(f"Unsupported family_name={family_name!r}") + return getattr(self, registry_attr) + + def _legacy_precision_bucket( + self, family_name: str, precision_key: str + ) -> Dict[str, BaseEstimator]: + if not isinstance(precision_key, str) or not precision_key: + raise ValueError( + f"precision_key must be a non-empty string, got {precision_key!r}" + ) + registry = self._legacy_precision_registry(family_name) + canonical_key = precision_key.upper() + bucket = registry.get(canonical_key) + if bucket is not None: + return bucket + for stored_key, stored_bucket in registry.items(): + if str(stored_key).upper() == canonical_key: + return stored_bucket + return {} + + def _store_model_precision( + self, + model_name: str, + precision: str, + model: BaseEstimator, + layer_contract: Optional[ResolvedLayerContract] = None, + ) -> None: + if not isinstance(precision, str) or not precision.strip(): + raise ValueError(f"precision must be a non-empty string, got {precision!r}") + precision_key = precision.upper() + family_name = self._measurement_family_name(self._active_measurement_type) + identity = self._model_contract_identity(model, layer_contract) + if identity is None: + self._legacy_model_registry(family_name)[model_name] = model + self._legacy_precision_registry(family_name).setdefault( + precision_key, {} + )[model_name] = model + return + + model_key = (model_name, identity) + self._contract_model_registry(family_name)[model_key] = model + self._contract_precision_registry(family_name).setdefault( + precision_key, {} + )[model_key] = model + + def _get_family_model( + self, + family_name: str, + model_name: str, + *, + precision_key: Optional[str] = None, + requested_identity: Optional[str] = None, + ) -> Optional[BaseEstimator]: + if precision_key is not None: + precision_key = precision_key.upper() + source = self._contract_precision_registry(family_name).get( + precision_key, {} + ) + else: + source = self._contract_model_registry(family_name) + + typed_candidates = { + identity: candidate + for (candidate_name, identity), candidate in source.items() + if candidate_name == model_name + } + legacy = ( + self._legacy_precision_bucket(family_name, precision_key).get(model_name) + if precision_key is not None + else self._legacy_model_registry(family_name).get(model_name) + ) + legacy_identity = ( + getattr(legacy, "_frontier_layer_cache_identity", None) + if legacy is not None + else None + ) + + if requested_identity is not None: + model = typed_candidates.get(requested_identity) + if model is not None: + return model + if legacy is not None and legacy_identity == requested_identity: + return legacy + return None + + if len(typed_candidates) > 1: + identities = sorted( + "" if identity is None else identity + for identity in typed_candidates + ) + raise ValueError( + f"Model '{model_name}' has multiple layer contracts; provide " + f"layer_contract. Available identities: {identities}" + ) + if len(typed_candidates) == 1: + typed_identity = next(iter(typed_candidates)) + if legacy is not None and legacy_identity != typed_identity: + raise ValueError( + f"Model '{model_name}' has multiple layer contracts; provide " + f"layer_contract. Available identities: " + f"[{typed_identity!r}, {legacy_identity or ''!r}]" + ) + return next(iter(typed_candidates.values())) + return legacy + + def _resolve_cluster_model_contract( + self, + cluster_type: Optional[ClusterType], + model_name: str, + ) -> Optional[ResolvedLayerContract]: + """Resolve the typed domain requested by one cluster view.""" + + if cluster_type is None: + return None + cluster_configs = getattr(self, "_cluster_configs", None) or {} + cluster_config = cluster_configs.get(cluster_type) + if cluster_config is None: + return None + replica_config = getattr(cluster_config, "replica_config", None) + model_config = getattr(replica_config, "model_config", None) + if replica_config is None or model_config is None: + return None + architecture_profile = _resolve_model_architecture_profile(model_config) + if architecture_profile is None: + return None + base_name = get_moe_gating_base_model_name(model_name) + typed_family = _resolve_profile_typed_family_for_query( + architecture_profile, base_name + ) + if typed_family is None: + return None + _, layer_kind = typed_family + return self._resolve_typed_layer_contract( + base_name, + cluster_type, + replica_config, + is_moe_model=layer_kind is not LayerKind.DENSE, + ) + + def _is_ffn_typed_model_for_cluster( + self, + cluster_type: Optional[ClusterType], + model_name: str, + ) -> bool: + """Return whether a model belongs to an FFN domain excluded by a view.""" + + if cluster_type != ClusterType.DECODE_ATTN: + return False + cluster_config = self._cluster_configs.get(cluster_type) + replica_config = getattr(cluster_config, "replica_config", None) + model_config = getattr(replica_config, "model_config", None) + architecture_profile = _resolve_model_architecture_profile(model_config) + if architecture_profile is None: + return False + base_name = get_moe_gating_base_model_name(model_name) + return _resolve_profile_typed_family_for_query( + architecture_profile, base_name + ) is not None + + def _models_view_for_family( + self, + family_name: str, + cluster_type: Optional[ClusterType] = None, + ) -> Dict[str, BaseEstimator]: + """Project one measurement family's canonical registries to model names.""" + + names = set(self._legacy_model_registry(family_name)) + names.update( + model_name + for model_name, _identity in self._contract_model_registry(family_name) + ) + models: Dict[str, BaseEstimator] = {} + for model_name in sorted(names): + if self._is_ffn_typed_model_for_cluster(cluster_type, model_name): + continue + contract = self._resolve_cluster_model_contract(cluster_type, model_name) + identity = _serialize_selected_layer_cache_identity(contract) + model = self._get_family_model( + family_name, + model_name, + requested_identity=identity, + ) + if identity is not None and model is None: + raise ValueError( + f"No trained model for {model_name!r} matches the typed contract " + f"requested by cluster {cluster_type!r}: {identity}" + ) + if model is not None: + models[model_name] = model + return models + + def get_model( + self, + model_name: str, + precision: Optional[str] = None, + layer_contract: Optional[ResolvedLayerContract] = None, + ) -> Optional[BaseEstimator]: + """Get a model by name, precision, and optional typed contract.""" + + if self._all_dummy_mode: + return None + requested_identity = _serialize_selected_layer_cache_identity(layer_contract) + precision_key = precision.upper() if precision else None + for family_name in ("eager", "device_event", "kernel_only"): + model = self._get_family_model( + family_name, + model_name, + precision_key=precision_key, + requested_identity=requested_identity, + ) + if model is not None: + return model + + if precision_key is not None: + available_precisions = sorted( + { + str(value).upper() + for value in self._contract_precision_registry("eager") + } + | { + str(value).upper() + for value in self._contract_precision_registry("device_event") + } + | { + str(value).upper() + for value in self._contract_precision_registry("kernel_only") + } + | { + str(value).upper() + for value in self._legacy_precision_registry("eager") + } + | { + str(value).upper() + for value in self._legacy_precision_registry("device_event") + } + | { + str(value).upper() + for value in self._legacy_precision_registry("kernel_only") + } + ) + raise ValueError( + f"Model '{model_name}' not available for precision '{precision_key}'. " + f"Available precisions: {available_precisions}. " + "Ensure profiling data matches the requested precision." + ) + return None + + def _load_model_from_cache(self, model_name: str, model_hash: str) -> BaseEstimator: + with InterProcessReaderWriterLock(f"{self._cache_dir}/{model_hash}_model_lock.file").read_lock(): + cache_file = f"{self._cache_dir}/{model_name}_{model_hash}.pkl" + if not os.path.exists(cache_file): + return None + logger.info(f"✓ Loaded pre-trained model '{model_name}' from cache (hash: {model_hash})") + logger.info(f" Cache file: {cache_file}") + return pickle.load(open(cache_file, "rb")) + + def _store_model_in_cache(self, model_name: str, model_hash: str, model: BaseEstimator) -> None: + with InterProcessReaderWriterLock(f"{self._cache_dir}/{model_hash}_model_lock.file").write_lock(): + cache_file = f"{self._cache_dir}/{model_name}_{model_hash}.pkl" + atomic_pickle_dump(model, cache_file) + logger.info(f"✓ Saved trained model '{model_name}' to cache (hash: {model_hash})") + logger.info(f" Cache file: {cache_file}") + + def _ensure_exact_lookup_metadata( + self, + *, + model_name: str, + model_hash: str, + model: BaseEstimator, + df: pd.DataFrame, + feature_cols: List[str], + target_col: str, + ) -> None: + """Persist exact measured rows before an on-demand model cache is stored.""" + if hasattr(model, "_frontier_exact_lookup"): + exact_lookup = getattr(model, "_frontier_exact_lookup") + if not isinstance(exact_lookup, Mapping): + raise ValueError( + f"Exact lookup metadata for {model_name} must be a mapping" + ) + return + + setattr( + model, + "_frontier_exact_lookup", + _build_exact_feature_lookup(df, feature_cols, target_col), + ) + self._store_model_in_cache(model_name, model_hash, model) diff --git a/frontier/execution_time_predictor/profiling_dataframe_loaders.py b/frontier/execution_time_predictor/profiling_dataframe_loaders.py new file mode 100644 index 00000000..3a628e27 --- /dev/null +++ b/frontier/execution_time_predictor/profiling_dataframe_loaders.py @@ -0,0 +1,985 @@ +"""Loading and validation of the profiling CSVs the predictors train on. + +Each loader reads one file of the canonical profiling taxonomy, validates the +columns its consumers require, applies the typed operator contract filter +when the file carries one, and derives the feature columns training needs. +""" + +import os +import pandas as pd + +from frontier.attention.families import DENSE_ATTENTION_FAMILY +from frontier.attention.ops import AttentionOperatorRole +from frontier.attention.profiling_mapping import ( + get_enabled_predictor_median_column_by_role, + get_enabled_predictor_metric_name_by_role, +) +from frontier.attention.string_coercion import coerce_truthy_bool +from frontier.execution_time_predictor.attention_dataset_contract import ( + enforce_mixed_attention_input_contract, +) +from frontier.execution_time_predictor.attention_tp_policy import ( + resolve_effective_attention_tp_size, +) +from frontier.execution_time_predictor.prediction_model_identity import ( + _resolve_model_architecture_profile, + _typed_row_matches_contract, + _validate_typed_parallel_selection, +) +from frontier.execution_time_predictor.profiling_metadata import ( + infer_single_runtime_model_config, + infer_single_runtime_profile, + validate_model_architecture_profile, +) +from frontier.logger import init_logger +from frontier.model_architectures import ResolvedLayerContract +from frontier.operators.typed_contracts import ( + TYPED_OPERATOR_CONTRACTS_COLUMN, + validate_typed_operator_contracts, +) +from frontier.profiling.cpu_overhead.validation import ( + apply_cpu_overhead_schema_v2_defaults, + validate_cpu_overhead_dataframe, +) +from frontier.types import ClusterType +from typing import Any, Dict, List, Optional, cast + + +logger = init_logger(__name__) + + +class ProfilingDataFrameLoaders: + """Profiling CSV loading, validation and derived features.""" + + def _load_linear_op_df( + self, + file_path: str, + tensor_parallel_size: int, + required_columns: Optional[List[str]] = None, + training_context: Optional[Dict[str, Any]] = None, + layer_contract: Optional[ResolvedLayerContract] = None, + operator_name: Optional[str] = None, + ) -> pd.DataFrame: + """ + Load linear operation dataframe (linear_op.csv or mlp.csv) with tensor parallel filtering. + + This function loads profiling data for linear operations including: + - Attention projections: attn_pre_proj, attn_post_proj, attn_rope + - MLP operations: mlp_up_proj, mlp_down_proj, mlp_act + - LayerNorm operations: input_layernorm, post_attention_layernorm + - Residual operations: add + + Note: This function is for linear_op.csv data only. For MoE data, use _load_moe_df(). + + Args: + file_path: Path to the profiling CSV file (linear_op.csv or mlp.csv) + tensor_parallel_size: Required tensor parallel size for filtering + + Returns: + Filtered DataFrame + + Raises: + FileNotFoundError: If the input file does not exist + ValueError: If required columns are missing or no data matches filtering criteria + """ + if layer_contract is not None: + _validate_typed_parallel_selection( + layer_contract, + tensor_parallel_size=tensor_parallel_size, + ) + + # Check file existence + if not os.path.exists(file_path): + raise FileNotFoundError( + f"Linear ops input file does not exist: {file_path}\n" + f"Please run profiling first to generate this file.\n" + f"Suggested command: bash frontier/profiling/example/test_profiling_linear_op.sh" + ) + + df = pd.read_csv(file_path) + logger.info(f"Original linear ops data: {len(df)} rows, {len(df.columns)} columns") + expected_profile = ( + layer_contract.profile_id + if layer_contract is not None + else infer_single_runtime_profile(self) + ) + if expected_profile is not None and "model_architecture_profile" in df.columns: + validate_model_architecture_profile( + df, + file_path=file_path, + expected_profile=expected_profile, + ) + + # Check required column + if 'num_tensor_parallel_workers' not in df.columns: + raise ValueError( + f"Column 'num_tensor_parallel_workers' not found in {file_path}\n" + f"Available columns: {list(df.columns)}\n" + f"This may indicate a corrupted or incompatible profiling file." + ) + + has_typed_contracts = TYPED_OPERATOR_CONTRACTS_COLUMN in df.columns + parsed_typed_contracts: Optional[pd.Series] = None + if has_typed_contracts: + if not operator_name and layer_contract is not None: + raise ValueError( + "typed profiling loading requires operator_name when the " + f"canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" + ) + if operator_name is not None and layer_contract is None: + raise ValueError( + "typed profiling loading requires layer_contract when the " + f"canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" + ) + # Parse every row before applying scalar filters so malformed metadata + # cannot be hidden by an unrelated TP or width selector. + parsed_typed_contracts = cast( + pd.Series, + df[TYPED_OPERATOR_CONTRACTS_COLUMN].map( + lambda raw_contracts: validate_typed_operator_contracts( + raw_contracts, + model_config=infer_single_runtime_model_config(self), + ) + ), + ) + + # Show filtering conditions + available_tp = sorted(df['num_tensor_parallel_workers'].unique()) + logger.info(f"Filtering conditions:") + logger.info(f" - num_tensor_parallel_workers == {tensor_parallel_size}") + logger.info(f" - Available num_tensor_parallel_workers: {available_tp}") + + # Apply filtering + filtered_df: pd.DataFrame = cast( + pd.DataFrame, + df[df["num_tensor_parallel_workers"] == tensor_parallel_size], + ) + if parsed_typed_contracts is not None and layer_contract is not None: + selected_layer_contract = layer_contract + if not isinstance(operator_name, str) or not operator_name: + raise ValueError( + "typed profiling loading requires a non-empty operator_name " + "for contract matching" + ) + typed_mask = parsed_typed_contracts.loc[filtered_df.index].map( + lambda raw_contracts: _typed_row_matches_contract( + raw_contracts, + selected_layer_contract, + operator_name=operator_name, + ) + ) + filtered_df = cast(pd.DataFrame, filtered_df[typed_mask]) + if filtered_df.empty: + raise ValueError( + "No linear-op rows match the selected typed layer contract " + f"for operator={operator_name!r}, TP={tensor_parallel_size} " + f"in {file_path}" + ) + elif layer_contract is not None: + if "n_expanded_embd" not in filtered_df.columns: + raise ValueError( + "Legacy linear-op profiling data is missing 'n_expanded_embd' " + f"for typed contract loading in {file_path}" + ) + filtered_df = cast( + pd.DataFrame, + filtered_df[ + filtered_df["n_expanded_embd"] + == layer_contract.effective_ffn_width + ], + ) + logger.info(f"After filtering: {len(filtered_df)} rows") + + expected_use_qk_norm = None + if training_context is not None and "use_qk_norm" in training_context: + expected_use_qk_norm = bool(training_context["use_qk_norm"]) + + if expected_use_qk_norm is True and "use_qk_norm" not in filtered_df.columns: + raise ValueError( + "linear_op profiling data is missing 'use_qk_norm' column for a model " + "that requires QK-norm-aware filtering. " + f"file={file_path}, model={training_context.get('model_name') if training_context else 'unknown'}" + ) + + if expected_use_qk_norm is not None and "use_qk_norm" in filtered_df.columns: + filtered_df = filtered_df[ + filtered_df["use_qk_norm"].astype(bool) == expected_use_qk_norm + ] + logger.info( + "After use_qk_norm filtering: %s rows (expected_use_qk_norm=%s)", + len(filtered_df), + expected_use_qk_norm, + ) + + if len(filtered_df) == 0: + width_requirement = ( + layer_contract.effective_ffn_width + if layer_contract is not None + else "legacy model width" + ) + raise ValueError( + f"No data matches the filtering criteria in {file_path}\n" + f"Required tensor_parallel_size: {tensor_parallel_size}\n" + f"Available tensor_parallel_sizes: {available_tp}\n" + f"Required effective_ffn_width: {width_requirement}\n" + f"Please run profiling with the correct configuration." + ) + + if required_columns: + self._validate_required_linear_op_columns( + filtered_df, + required_columns, + file_path, + training_context=training_context, + ) + + return filtered_df + + def _get_required_attn_linear_op_columns(self, model_config) -> List[str]: + required_columns = [ + "time_stats.attn_pre_proj.median", + "time_stats.attn_post_proj.median", + "time_stats.attn_rope.median", + ] + if model_config is not None and bool(getattr(model_config, "use_qk_norm", False)): + required_columns.append("use_qk_norm") + architecture_profile = _resolve_model_architecture_profile(model_config) + if architecture_profile is not None: + required_columns.extend( + f"time_stats.{op_name}.median" + for op_name in architecture_profile.predictor_attention_extra_ops + ) + return required_columns + + @staticmethod + + def _get_required_target_embedded_mtp_linear_op_columns() -> List[str]: + return [ + "time_stats.mtp_fusion_proj.median", + "time_stats.lm_head_linear.median", + ] + + @staticmethod + + def _validate_required_linear_op_columns( + df: pd.DataFrame, + required_columns: List[str], + file_path: str, + training_context: Optional[Dict[str, Any]] = None, + ) -> None: + missing_columns = [col for col in required_columns if col not in df.columns] + all_nan_columns = [ + col + for col in required_columns + if col in df.columns and df[col].isna().all() + ] + + if missing_columns or all_nan_columns: + context_text = "" + if training_context: + context_text = f"\nTraining context: {training_context}" + + raise ValueError( + "Required attention linear op columns are missing or all-NaN in " + f"{file_path}." + f"\nMissing columns: {missing_columns}" + f"\nAll-NaN columns: {all_nan_columns}" + f"{context_text}" + ) + + def _load_attention_df( + self, + file_path: str, + replica_config, + replica_scheduler_config, + cluster_type: Optional[ClusterType] = None, + ) -> pd.DataFrame: + """ + Load attention dataframe (attention.csv) with model configuration filtering. + + Args: + file_path: Path to the attention profiling CSV file + replica_config: Replica configuration for filtering + replica_scheduler_config: Replica scheduler configuration for block size + cluster_type: Cluster type for policy warning context + + Returns: + Filtered DataFrame + + Raises: + FileNotFoundError: If the input file does not exist + ValueError: If no data matches filtering criteria + """ + # Check file existence + if not os.path.exists(file_path): + raise FileNotFoundError( + f"Attention input file does not exist: {file_path}\n" + f"Please run attention profiling first to generate this file.\n" + f"Suggested command: bash frontier/profiling/example/test_profiling_attention.sh" + ) + + df = pd.read_csv(file_path) + df = df.drop_duplicates() + logger.info(f"Original attention data: {len(df)} rows, {len(df.columns)} columns") + + enforce_mixed_attention_input_contract( + attention_file_path=file_path, + available_columns=df.columns, + ) + + # Latent-MLA profiles use a distinct structural schema (runtime kv heads = 1, + # head size = kv_lora_rank + qk_rope_head_dim); route them to the MLA + # structural filter before the dense cache-write fill / dense filter. + model_config = replica_config.model_config + if self._is_mla_family(model_config): + return self._filter_mla_attention_df( + df, file_path, replica_config, replica_scheduler_config + ) + + # Fill missing cache-write column for older attention profiling CSVs. + cache_write_median_column = get_enabled_predictor_median_column_by_role( + DENSE_ATTENTION_FAMILY, + AttentionOperatorRole.CACHE_WRITE, + ) + for column in [cache_write_median_column]: + if column not in df.columns: + df[column] = 0 + else: + df.fillna({column: 0}, inplace=True) + + model_config = replica_config.model_config + requested_tp = replica_config.attn_tensor_parallel_size + prefill_op_name = get_enabled_predictor_metric_name_by_role( + DENSE_ATTENTION_FAMILY, + AttentionOperatorRole.PREFILL_KERNEL, + ) + effective_tp = resolve_effective_attention_tp_size( + op_name=prefill_op_name, + requested_tp_size=requested_tp, + num_kv_heads=model_config.num_kv_heads, + cluster_type=cluster_type, + warning_cache=getattr(self, "_attention_tp_warning_cache", None), + include_linear_ops=False, + ) + + # Show filtering conditions + logger.info(f"Filtering conditions:") + logger.info(f" - n_embd == {model_config.embedding_dim}") + logger.info(f" - n_q_head == {model_config.num_q_heads}") + logger.info(f" - n_kv_head == {model_config.num_kv_heads}") + logger.info(f" - block_size == {replica_scheduler_config.block_size}") + logger.info( + " - num_tensor_parallel_workers == %s (requested_tp=%s)", + effective_tp, + requested_tp, + ) + + filtered_df = df[ + (df["n_embd"] == model_config.embedding_dim) + & (df["n_q_head"] == model_config.num_q_heads) + & (df["n_kv_head"] == model_config.num_kv_heads) + & (df["block_size"] == replica_scheduler_config.block_size) + & (df["num_tensor_parallel_workers"] == effective_tp) + ] + + logger.info(f"After filtering: {len(filtered_df)} rows") + + if len(filtered_df) == 0: + # Surface what is available to make debugging explicit. + available = { + "n_embd": sorted(df["n_embd"].unique().tolist()) if "n_embd" in df else [], + "n_q_head": sorted(df["n_q_head"].unique().tolist()) if "n_q_head" in df else [], + "n_kv_head": sorted(df["n_kv_head"].unique().tolist()) if "n_kv_head" in df else [], + "block_size": sorted(df["block_size"].unique().tolist()) if "block_size" in df else [], + "num_tensor_parallel_workers": sorted(df["num_tensor_parallel_workers"].unique().tolist()) if "num_tensor_parallel_workers" in df else [], + } + + logger.error( + "Attention profiling rows are missing for the requested configuration. " + "Available values: %s", available + ) + + raise ValueError( + f"No data matches the filtering criteria in {file_path}\n" + f"Required configuration:\n" + f" - n_embd: {model_config.embedding_dim}\n" + f" - n_q_head: {model_config.num_q_heads}\n" + f" - n_kv_head: {model_config.num_kv_heads}\n" + f" - block_size: {replica_scheduler_config.block_size}\n" + f" - tensor_parallel_size(requested): {requested_tp}\n" + f" - tensor_parallel_size(effective): {effective_tp}\n" + f"Available values: {available}\n" + f"Please run attention profiling with the correct configuration." + ) + + return filtered_df + + def _load_all_reduce_df(self, file_path: str, replica_config, tensor_parallel_size: int) -> pd.DataFrame: + """ + Load all_reduce dataframe with cluster-specific tensor parallel size. + + Args: + file_path: Path to the communication profiling CSV file + replica_config: Replica configuration + tensor_parallel_size: Required tensor parallel size for filtering + + Returns: + Filtered DataFrame + + Raises: + FileNotFoundError: If the input file does not exist + ValueError: If no data matches filtering criteria + """ + if not os.path.exists(file_path): + raise FileNotFoundError( + f"All-reduce input file does not exist: {file_path}\n" + f"Please run communication profiling first.\n" + f"Suggested command: bash frontier/profiling/example/test_profiling_communication.sh" + ) + + df = pd.read_csv(file_path) + logger.info(f"Original all_reduce data: {len(df)} rows") + + # Show filtering conditions + logger.info(f"Filtering conditions:") + logger.info(f" - num_workers == {tensor_parallel_size}") + logger.info(f" - devices_per_node == {tensor_parallel_size}") + logger.info(f" - collective == 'all_reduce'") + + filtered_df = df[ + (df["num_workers"] == tensor_parallel_size) + & (df["devices_per_node"] == tensor_parallel_size) + & (df["collective"] == "all_reduce") + ] + + logger.info(f"After filtering: {len(filtered_df)} rows") + + if len(filtered_df) == 0: + available_info = "" + if len(df) > 0: + available_info = ( + f"Available values in file:\n" + f" - num_workers: {sorted(df['num_workers'].unique())}\n" + f" - devices_per_node: {sorted(df['devices_per_node'].unique())}\n" + f" - collective: {sorted(df['collective'].unique())}" + ) + raise ValueError( + f"No data matches the filtering criteria in {file_path}\n" + f"Required: num_workers={tensor_parallel_size}, devices_per_node={tensor_parallel_size}, collective='all_reduce'\n" + f"{available_info}" + ) + + return filtered_df + + def _load_send_recv_df(self, file_path: str, replica_config) -> pd.DataFrame: + """ + Load send_recv dataframe for pipeline parallel communication. + + Args: + file_path: Path to the communication profiling CSV file + replica_config: Replica configuration + + Returns: + Filtered DataFrame + + Raises: + FileNotFoundError: If the input file does not exist + """ + if not os.path.exists(file_path): + raise FileNotFoundError( + f"Send/recv input file does not exist: {file_path}\n" + f"Please run communication profiling first.\n" + f"Suggested command: bash frontier/profiling/example/test_profiling_communication.sh" + ) + + num_workers = replica_config.num_pipeline_stages * replica_config.attn_tensor_parallel_size + devices_per_node = replica_config.node_config.num_devices_per_node + is_multi_node = num_workers > devices_per_node + + if is_multi_node: + devices_per_node = 1 + else: + devices_per_node = 2 + + df = pd.read_csv(file_path) + logger.info(f"Original send_recv data: {len(df)} rows") + logger.info(f"Filtering conditions: collective='send_recv', devices_per_node={devices_per_node}") + + filtered_df = df[ + (df["collective"] == "send_recv") + & (df["devices_per_node"] == devices_per_node) + ] + + logger.info(f"After filtering: {len(filtered_df)} rows") + return filtered_df + + def _load_cpu_overhead_df(self, file_path: str, replica_config) -> pd.DataFrame: + """ + Load CPU overhead dataframe with model configuration filtering. + + Args: + file_path: Path to the CPU overhead profiling CSV file + replica_config: Replica configuration + + Returns: + Filtered DataFrame + + Raises: + FileNotFoundError: If the input file does not exist + """ + if not os.path.exists(file_path): + logger.warning( + "CPU overhead input file does not exist: %s. " + "Skipping CPU overhead model training for this cluster.", + file_path, + ) + return pd.DataFrame() + + df = pd.read_csv(file_path) + if df.empty: + logger.warning( + "CPU overhead input file is empty: %s. " + "Skipping CPU overhead model training for this cluster.", + file_path, + ) + return pd.DataFrame() + + df = apply_cpu_overhead_schema_v2_defaults( + df, + warn_fn=logger.warning, + context=file_path, + ) + df = validate_cpu_overhead_dataframe(df) + + model_config = replica_config.model_config + + logger.info(f"Original CPU overhead data: {len(df)} rows") + logger.info(f"Filtering conditions: model_name='{model_config.get_name()}', tensor_parallel_degree={replica_config.attn_tensor_parallel_size}") + + filtered_df = df[ + (df["model_name"] == model_config.get_name()) + & (df["tensor_parallel_degree"] == replica_config.attn_tensor_parallel_size) + ] + + logger.info(f"After filtering: {len(filtered_df)} rows") + if filtered_df.empty: + logger.warning( + "No CPU overhead profiling rows found for model_name='%s', " + "tensor_parallel_degree=%s in file '%s'.", + model_config.get_name(), + replica_config.attn_tensor_parallel_size, + file_path, + ) + return filtered_df + + # Load imbalance feature columns used for MoE training + # These features describe the load distribution across experts + # Reference: frontier/training/moe_trainer.py lines 224-239 (authoritative source) + # Reference: frontier/profiling/moe/LOAD_IMBALANCE_GUIDE.md + MOE_LOAD_IMBALANCE_FEATURES = [ + # Config features (6) - describe model configuration + "total_routed_tokens", # Total tokens after routing (num_tokens * router_topk) + "num_experts_per_device", # Number of experts per device after EP sharding + "hidden_dim", # Model hidden dimension + "expert_hidden_dim", # Expert FFN hidden dimension + "router_topk", # Number of experts each token is routed to + "model_expansion_ratio", # expert_hidden_dim / hidden_dim + # Derived features (2) - derived from config and routing + "tokens_per_expert_avg", # Average tokens per expert + "tokens_to_experts_ratio", # tokens / num_experts ratio + # Load features (6) - describe load distribution characteristics + "expert_utilization", # Proportion of experts with non-zero load + "min_load_ratio", # Min load / average load + "load_imbalance_cv", # Coefficient of Variation: std/mean, key imbalance metric + "max_load_ratio", # Max load / average load + "load_entropy", # Entropy of load distribution (higher = more uniform) + "load_gini_coefficient", # Gini coefficient: 0=equality, 1=inequality + ] + + # Feature columns for mixed-batch attention prefill model + # These features capture batch heterogeneity characteristics together with + # the uniform KV-cache context used by MixedAttentionInput profiling. + # Reference: frontier/training/attention_trainer.py lines 362-375 (authoritative source) + ATTN_PREFILL_MIXED_FEATURES = [ + # Core features (7) + "batch_size", # Number of sequences in batch + "kv_cache_size", # Uniform KV cache context for the mixed batch + "total_tokens", # Total tokens across all sequences + "avg_seq_len", # Average sequence length + "min_seq_len", # Minimum sequence length + "max_seq_len", # Maximum sequence length + "total_tokens_squared", # Computational complexity proxy + # Heterogeneity features (3) + "seq_len_variance", # Variance of sequence lengths + "seq_len_cv", # Coefficient of variation (std/mean) + "seq_len_range", # max_seq_len - min_seq_len + # Interaction features (2) + "batch_variance_interaction", # batch_size * seq_len_variance + "batch_cv_interaction", # batch_size * seq_len_cv + ] + + ATTN_DECODE_IN_MIXED_FEATURES = [ + "decode_batch_size", + "decode_avg_kv_cache_size", + "num_prefill_seqs", + "total_prefill_tokens", + "total_batch_size", + "batch_composition_ratio", + "total_tokens", + ] + + def _load_moe_df( + self, + file_path: str, + replica_config, + load_imbalance: bool = True, + tensor_parallel_size: Optional[int] = None, + expert_parallel_size: Optional[int] = None, + layer_contract: Optional[ResolvedLayerContract] = None, + operator_name: Optional[str] = None, + ) -> pd.DataFrame: + """ + Load MoE dataframe with cluster-specific configuration filtering. + + This function loads and filters MoE profiling data based on the model configuration + and parallelism settings. It supports two training modes controlled by `load_imbalance`: + + 1. **Load Imbalance Mode (default, load_imbalance=True)**: + - Uses profiling data that includes load imbalance features + - Training will use features like `load_imbalance_cv`, `load_gini_coefficient`, etc. + - Recommended for accurate MoE execution time prediction under real-world scenarios + - Requires profiling with `--enable_load_imbalance` flag + + 2. **Standard Mode (load_imbalance=False)**: + - Uses basic profiling data without load imbalance features + - Training only uses `num_tokens` as feature + - Simpler but less accurate for imbalanced workloads + - Compatible with legacy profiling data + + The difference is in the **training features used**, not data row filtering. + Load imbalance mode uses additional features to capture expert load distribution. + + Reference: frontier/profiling/moe/LOAD_IMBALANCE_GUIDE.md + + Args: + file_path: Path to the MoE profiling CSV file + replica_config: Replica configuration containing model and parallelism settings + load_imbalance: Training mode flag: + - True (default): Load imbalance mode - use load imbalance features + - False: Standard mode - only use basic num_tokens feature + tensor_parallel_size: Optional TP override for op-specific MoE training. + If None, uses replica_config.moe_tensor_parallel_size. + expert_parallel_size: Optional EP filter for op-specific MoE training. + If None, EP filtering is skipped (used for EP-agnostic replicated ops). + + Returns: + Filtered DataFrame ready for MoE model training + + Raises: + FileNotFoundError: If the input file does not exist + ValueError: If no data matches filtering criteria or required features are missing + """ + if layer_contract is not None: + _validate_typed_parallel_selection( + layer_contract, + tensor_parallel_size=tensor_parallel_size, + expert_parallel_size=expert_parallel_size, + ) + + if not os.path.exists(file_path): + raise FileNotFoundError( + f"MoE input file does not exist: {file_path}\n" + f"Please run MoE profiling first.\n" + f"Suggested command: bash frontier/profiling/example/test_profiling_moe.sh" + ) + + df = pd.read_csv(file_path) + logger.info(f"Original MoE data: {len(df)} rows, {len(df.columns)} columns") + expected_profile = ( + layer_contract.profile_id + if layer_contract is not None + else infer_single_runtime_profile(self) + ) + if expected_profile is not None and "model_architecture_profile" in df.columns: + validate_model_architecture_profile( + df, + file_path=file_path, + expected_profile=expected_profile, + ) + + has_typed_contracts = TYPED_OPERATOR_CONTRACTS_COLUMN in df.columns + parsed_typed_contracts: Optional[pd.Series] = None + if has_typed_contracts: + if not operator_name: + raise ValueError( + "typed profiling loading requires operator_name when the " + f"canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" + ) + if layer_contract is None: + raise ValueError( + "typed profiling loading requires layer_contract when the " + f"canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" + ) + # Parse every row before applying scalar filters so malformed metadata + # cannot be hidden by an unrelated TP, EP, or width selector. + parsed_typed_contracts = cast( + pd.Series, + df[TYPED_OPERATOR_CONTRACTS_COLUMN].map( + lambda raw_contracts: validate_typed_operator_contracts( + raw_contracts, + model_config=replica_config.model_config, + ) + ), + ) + + model_config = replica_config.model_config + training_mode = "load_imbalance (load_imbalance=True)" if load_imbalance else "standard (load_imbalance=False)" + if tensor_parallel_size is None: + tensor_parallel_size = replica_config.moe_tensor_parallel_size + if tensor_parallel_size <= 0: + raise ValueError( + f"Invalid tensor_parallel_size for MoE data loading: {tensor_parallel_size}" + ) + + # Display filtering conditions + logger.info(f"Filtering conditions:") + logger.info(f" - num_experts == {model_config.num_experts}") + logger.info(f" - router_topk == {model_config.num_experts_per_tok}") + logger.info(f" - hidden_dim == {model_config.embedding_dim}") + expected_expert_width = ( + layer_contract.effective_ffn_width + if layer_contract is not None + else model_config.mlp_hidden_dim + ) + logger.info(f" - expert_hidden_dim == {expected_expert_width}") + logger.info(f" - num_tensor_parallel_workers == {tensor_parallel_size}") + if expert_parallel_size is None: + logger.info(" - expert_parallel_size == ANY (EP-agnostic op)") + else: + logger.info(f" - expert_parallel_size == {expert_parallel_size}") + logger.info(f" - training_mode: {training_mode}") + + # Display available values in the dataset + available_info = [] + if len(df) > 0: + if 'num_experts' in df.columns: + available_info.append(f" - Available num_experts: {sorted(df['num_experts'].unique())}") + if 'router_topk' in df.columns: + available_info.append(f" - Available router_topk: {sorted(df['router_topk'].unique())}") + if 'num_tensor_parallel_workers' in df.columns: + available_info.append(f" - Available num_tensor_parallel_workers: {sorted(df['num_tensor_parallel_workers'].unique())}") + if 'expert_parallel_size' in df.columns: + available_info.append(f" - Available expert_parallel_size: {sorted(df['expert_parallel_size'].unique())}") + if 'load_distribution' in df.columns: + available_info.append(f" - Available load_distribution: {sorted(df['load_distribution'].unique())}") + + for info in available_info: + logger.info(info) + + # Apply filtering based on MoE configuration + filtered_df = cast(pd.DataFrame, df[ + (df["num_experts"] == model_config.num_experts) + & (df["router_topk"] == model_config.num_experts_per_tok) + & (df["hidden_dim"] == model_config.embedding_dim) + & (df["num_tensor_parallel_workers"] == tensor_parallel_size) + ]) + if not has_typed_contracts: + filtered_df = filtered_df[ + filtered_df["expert_hidden_dim"] == expected_expert_width + ] + else: + if parsed_typed_contracts is None: + raise RuntimeError( + "typed MoE metadata column was detected but could not be parsed" + ) + if layer_contract is None: + raise ValueError( + "typed MoE filtering requires a resolved layer contract" + ) + selected_layer_contract = layer_contract + typed_mask = parsed_typed_contracts.loc[filtered_df.index].map( + lambda raw_contracts: _typed_row_matches_contract( + raw_contracts, + selected_layer_contract, + operator_name=operator_name, + ) + ) + filtered_df = cast(pd.DataFrame, filtered_df[typed_mask]) + filtered_df = cast(pd.DataFrame, filtered_df) + if expert_parallel_size is not None: + if "expert_parallel_size" not in filtered_df.columns: + raise ValueError( + "MoE profiling data is missing 'expert_parallel_size' while " + f"EP={expert_parallel_size} is required in {file_path}" + ) + filtered_df = filtered_df[ + filtered_df["expert_parallel_size"] == expert_parallel_size + ] + + logger.info(f"After config filtering: {len(filtered_df)} rows") + + # Check for load imbalance features if load_imbalance mode is enabled + if load_imbalance: + missing_features = [ + f for f in self.MOE_LOAD_IMBALANCE_FEATURES + if f not in filtered_df.columns + ] + if missing_features: + logger.warning( + f"Load imbalance mode requested but missing features: {missing_features}\n" + f"Available columns: {list(filtered_df.columns)}\n" + f"Please run MoE profiling with --enable_load_imbalance flag.\n" + f"Use load_imbalance=False (standard mode) explicitly if you want to train without load imbalance features." + ) + raise ValueError("Missing load imbalance features") + # Note: We don't change load_imbalance here, caller should handle feature selection + else: + logger.info(f"Load imbalance features available: {self.MOE_LOAD_IMBALANCE_FEATURES}") + + if len(filtered_df) == 0: + ep_requirement = "ANY" if expert_parallel_size is None else expert_parallel_size + available_info_text = "\n".join(available_info) + message = ( + f"No data matches the filtering criteria in {file_path}\n" + f"Required MoE configuration:\n" + f" - num_experts: {model_config.num_experts}\n" + f" - router_topk: {model_config.num_experts_per_tok}\n" + f" - hidden_dim: {model_config.embedding_dim}\n" + f" - expert_hidden_dim: {expected_expert_width}\n" + f" - tensor_parallel_size: {tensor_parallel_size}\n" + f" - expert_parallel_size: {ep_requirement}\n" + f" - training_mode: {training_mode}\n" + ) + if has_typed_contracts: + message += ( + f" - typed operator: {operator_name!r}\n" + " - typed layer contract admission: required\n" + ) + if available_info_text: + message += available_info_text + raise ValueError( + message + ) + + return filtered_df + + def _get_attention_df_with_derived_features(self, df: pd.DataFrame) -> pd.DataFrame: + """Add derived features to attention dataframe. + + Standard features for attn_prefill and attn_decode: + - num_tokens: max(prefill_chunk_size, batch_size) + - is_decode: derived from is_prefill when available, else prefill_chunk_size == 0 + - prefill_chunk_size_squared: prefill_chunk_size ** 2 + + Mixed-batch features for attn_prefill_mixed (12 features): + Reference: frontier/training/attention_trainer.py lines 362-375 + These features capture batch heterogeneity for accurate prefill time prediction. + """ + df_with_derived_features = df.copy() + + # Standard attention features + df_with_derived_features["num_tokens"] = df_with_derived_features[["prefill_chunk_size", "batch_size"]].max(axis=1) + if "is_prefill" in df_with_derived_features.columns: + normalized_prefill_values = coerce_truthy_bool( + df_with_derived_features["is_prefill"] + ) + df_with_derived_features["is_decode"] = ~normalized_prefill_values + else: + df_with_derived_features["is_decode"] = (df_with_derived_features["prefill_chunk_size"] == 0) + df_with_derived_features["prefill_chunk_size_squared"] = (df_with_derived_features["prefill_chunk_size"] ** 2) + + def _normalize_bool_series(series: pd.Series) -> pd.Series: + return coerce_truthy_bool(series) + + if "is_mixed_batch" in df_with_derived_features.columns: + df_with_derived_features["is_mixed_batch"] = _normalize_bool_series( + df_with_derived_features["is_mixed_batch"] + ) + else: + df_with_derived_features["is_mixed_batch"] = False + + if "is_true_mixed_batch" in df_with_derived_features.columns: + df_with_derived_features["is_true_mixed_batch"] = _normalize_bool_series( + df_with_derived_features["is_true_mixed_batch"] + ) + else: + df_with_derived_features["is_true_mixed_batch"] = False + + # Mixed-batch features for attn_prefill_mixed (if applicable) + # Check if the profiling data contains mixed-batch specific columns + has_mixed_batch_data = "total_tokens" in df_with_derived_features.columns + + if has_mixed_batch_data: + logger.info("Adding mixed-batch derived features for attn_prefill_mixed") + + # total_tokens_squared for computational complexity + if "total_tokens" in df_with_derived_features.columns: + df_with_derived_features["total_tokens_squared"] = ( + df_with_derived_features["total_tokens"] ** 2 + ) + + # seq_len_range = max_seq_len - min_seq_len + if "max_seq_len" in df_with_derived_features.columns and "min_seq_len" in df_with_derived_features.columns: + df_with_derived_features["seq_len_range"] = ( + df_with_derived_features["max_seq_len"] - + df_with_derived_features["min_seq_len"] + ) + + # Interaction features: batch_size * heterogeneity metrics + if "seq_len_variance" in df_with_derived_features.columns: + df_with_derived_features["batch_variance_interaction"] = ( + df_with_derived_features["batch_size"] * + df_with_derived_features["seq_len_variance"] + ) + + if "seq_len_cv" in df_with_derived_features.columns: + df_with_derived_features["batch_cv_interaction"] = ( + df_with_derived_features["batch_size"] * + df_with_derived_features["seq_len_cv"] + ) + + if { + "num_prefill_seqs", + "num_decode_seqs", + }.issubset(df_with_derived_features.columns) and ( + "total_batch_size" not in df_with_derived_features.columns + ): + df_with_derived_features["total_batch_size"] = ( + df_with_derived_features["num_prefill_seqs"] + + df_with_derived_features["num_decode_seqs"] + ) + + if { + "num_prefill_seqs", + "total_batch_size", + }.issubset(df_with_derived_features.columns) and ( + "batch_composition_ratio" not in df_with_derived_features.columns + ): + total_batch_size = df_with_derived_features["total_batch_size"].replace(0, pd.NA) + df_with_derived_features["batch_composition_ratio"] = ( + df_with_derived_features["num_prefill_seqs"] / total_batch_size + ).fillna(0.0) + + if ( + "num_decode_seqs" in df_with_derived_features.columns + and "decode_batch_size" not in df_with_derived_features.columns + ): + df_with_derived_features["decode_batch_size"] = df_with_derived_features[ + "num_decode_seqs" + ] + + return df_with_derived_features + + def _get_all_reduce_df_with_derived_features(self, df: pd.DataFrame, replica_config) -> pd.DataFrame: + df_with_derived_features = df.copy() + df_with_derived_features["num_tokens"] = ( + df_with_derived_features["size"] / replica_config.model_config.embedding_dim / 2 + ) + return df_with_derived_features + + def _get_send_recv_df_with_derived_features(self, df: pd.DataFrame, replica_config) -> pd.DataFrame: + df_with_derived_features = df.copy() + df_with_derived_features["num_tokens"] = ( + df_with_derived_features["size"] / replica_config.model_config.embedding_dim / 2 + ) + return df_with_derived_features diff --git a/frontier/execution_time_predictor/shared_prediction_model_manager.py b/frontier/execution_time_predictor/shared_prediction_model_manager.py index 9256ffca..c4327219 100644 --- a/frontier/execution_time_predictor/shared_prediction_model_manager.py +++ b/frontier/execution_time_predictor/shared_prediction_model_manager.py @@ -98,311 +98,47 @@ is_target_embedded_mtp_same_tp_linear_op, ) -logger = init_logger(__name__) -MIGRATION_HELP_COMMAND = ( - "python -m frontier.profiling.migrate_csv_metadata --help" -) - - -def _get_moe_family_model_names() -> List[str]: - return list(get_family_profiling_names(MOE_FAMILY)) - - -def _get_moe_family_operator_by_model_name(model_name: str): - moe_ops = { - operator.profiling_name(): operator - for operator in MOE_FAMILY.profiling_ops() - } - if model_name not in moe_ops: - raise ValueError(f"Unsupported MoE op: {model_name}") - return moe_ops[model_name] - - -def _get_moe_gating_family_model_names() -> List[str]: - return [ - operator.profiling_name() - for operator in MOE_FAMILY.profiling_ops() - if operator.precision_name() == "moe_gating" - ] - - -def _get_prefill_hot_moe_gating_model_names() -> List[str]: - return [ - f"{model_name}__prefill_hot" - for model_name in _get_moe_gating_family_model_names() - ] - - -def _resolve_model_architecture_profile( - model_config: Any, - *, - allow_generic: bool = False, -) -> Optional[ModelArchitectureProfile]: - if model_config is None: - return None - getter = getattr(model_config, "get_model_architecture_profile", None) - if callable(getter): - return cast(Optional[ModelArchitectureProfile], getter()) - - # Lightweight test/config adapters that predate the typed contract do not - # expose a profile accessor or typed widths. Keep those callers on the - # scalar compatibility path instead of treating the generic fallback as a - # complete typed declaration. An explicit profile or typed width opts the - # adapter into strict profile-owned resolution. - typed_fields = ( - "model_architecture_profile", - "dense_mlp_hidden_dim", - "routed_mlp_hidden_dim", - "share_expert_dim", - ) - if not allow_generic and not any( - getattr(model_config, field_name, None) is not None - for field_name in typed_fields - ): - return None - return get_model_architecture_profile(model_config) - - -def _resolve_model_architecture_profile_id(model_config) -> str: - architecture_profile = _resolve_model_architecture_profile(model_config) - if architecture_profile is None: - return "generic" - return architecture_profile.profile_id - - -def _resolve_profile_typed_family_for_query( - architecture_profile: ModelArchitectureProfile, - op_name: str, -) -> Optional[Tuple[str, LayerKind]]: - """Resolve a query to the profile-owned typed operator family.""" - - if not isinstance(op_name, str) or not op_name: - raise ValueError("typed operator query name must be a non-empty string") - matches: list[Tuple[str, LayerKind]] = [] - for layer_contract in architecture_profile.layer_contracts: - for family_id in layer_contract.operator_family_ids: - family = get_operator_family(family_id) - if any( - op_name == operator.name - or op_name == operator.profiling_name() - for operator in family.operators - ): - matches.append((family_id, layer_contract.layer_kind)) - if len(matches) > 1: - raise ValueError( - f"Operator query {op_name!r} belongs to multiple typed layer families: " - f"{sorted(family_id for family_id, _ in matches)}" - ) - return matches[0] if matches else None - - -def _serialize_selected_layer_cache_identity( - layer_contract: Optional[ResolvedLayerContract], -) -> Optional[str]: - """Serialize the selected semantic domain used by a model cache. - - Physical ``layer_id`` and producer-side domain envelopes do not identify a - trained estimator. Keep only the selected fields that affect estimator - admission and reuse, in deterministic JSON form. - """ - - if layer_contract is None: - return None - if not isinstance(layer_contract, ResolvedLayerContract): - raise TypeError( - "layer_contract must be a ResolvedLayerContract when provided" - ) - metadata = layer_contract.typed_metadata_identity() - selected_fields = ( - "profile_id", - "operator_family_id", - "layer_kind", - "dimension_source", - "effective_ffn_width", - "tensor_parallel_mode", - "expert_parallel_mode", - "selected_expert_parallel_size", - "selected_tensor_parallel_size", - "selected_padded_ffn_width", - ) - return json.dumps( - {field_name: metadata[field_name] for field_name in selected_fields}, - sort_keys=True, - separators=(",", ":"), - ensure_ascii=True, - allow_nan=False, - ) - - -def _get_contract_hash( - layer_contract: Optional[ResolvedLayerContract], -) -> str: - """Return a short deterministic hash for an optional selected contract.""" - - identity = _serialize_selected_layer_cache_identity(layer_contract) - if identity is None: - return "none" - return hashlib.sha256(identity.encode("utf-8")).hexdigest()[:16] - - -def _validate_typed_parallel_selection( - layer_contract: ResolvedLayerContract, - *, - tensor_parallel_size: Optional[int] = None, - expert_parallel_size: Optional[int] = None, -) -> None: - """Validate explicit loader selectors against a resolved contract.""" - - if not isinstance(layer_contract, ResolvedLayerContract): - raise TypeError("layer_contract must be a ResolvedLayerContract") - contract_tp = layer_contract.tensor_parallel_size - if ( - contract_tp is not None - and tensor_parallel_size is not None - and contract_tp != tensor_parallel_size - ): - raise ValueError( - f"typed layer contract TP {contract_tp} conflicts with " - f"tensor_parallel_size {tensor_parallel_size}" - ) - contract_ep = layer_contract.expert_parallel_size - if ( - contract_ep is not None - and expert_parallel_size is not None - and contract_ep != expert_parallel_size - ): - raise ValueError( - f"typed layer contract EP {contract_ep} conflicts with " - f"expert_parallel_size {expert_parallel_size}" - ) - - -def _typed_row_matches_contract( - raw_contracts: Any, - layer_contract: ResolvedLayerContract, - *, - operator_name: Optional[str], -) -> bool: - """Match one parsed or serialized row to its exact typed operator contract.""" - - if not isinstance(operator_name, str) or not operator_name: - raise ValueError( - "typed profiling loading requires a non-empty operator_name when " - f"the canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" - ) - if not isinstance(layer_contract.operator_family_id, str) or not layer_contract.operator_family_id: - raise ValueError( - "typed profiling loading requires a layer contract with an operator family id" - ) - return matches_resolved_layer_contract( - raw_contracts, - layer_contract, - operator_name=operator_name, - ) - - -def _normalize_layer_contract_context( - training_context: Optional[Mapping[str, Any]], - explicit_layer_contract: Optional[ResolvedLayerContract] = None, -) -> Tuple[Optional[ResolvedLayerContract], Dict[str, Any]]: - """Resolve one contract and keep every context representation consistent.""" - - context = dict(training_context or {}) - context_contract = context.get("layer_contract") - if context_contract is not None and not isinstance( - context_contract, ResolvedLayerContract - ): - raise TypeError( - "training_context['layer_contract'] must be a ResolvedLayerContract" - ) - if explicit_layer_contract is not None and not isinstance( - explicit_layer_contract, ResolvedLayerContract - ): - raise TypeError("layer_contract must be a ResolvedLayerContract") - - if context_contract is not None and explicit_layer_contract is not None: - if not context_contract.is_semantically_equivalent(explicit_layer_contract): - raise ValueError( - "conflicting layer_contract values were provided through the " - "explicit argument and training_context" - ) - - resolved_contract = explicit_layer_contract or context_contract - context_identity = context.get("layer_contract_identity") - if context_identity is not None and not isinstance(context_identity, str): - raise TypeError( - "training_context['layer_contract_identity'] must be a string" - ) - selected_identity = _serialize_selected_layer_cache_identity(resolved_contract) - if context_identity is not None and context_identity != selected_identity: - raise ValueError( - "training_context['layer_contract_identity'] does not match the " - "supplied layer_contract" - ) - if resolved_contract is None: - return None, context - - context["layer_contract"] = resolved_contract - context["layer_contract_identity"] = selected_identity - context["layer_kind"] = resolved_contract.layer_kind.value - context["effective_ffn_width"] = resolved_contract.effective_ffn_width - context["tensor_parallel_mode"] = resolved_contract.tensor_parallel_mode.value - context["expert_parallel_mode"] = resolved_contract.expert_parallel_mode.value - return resolved_contract, context - - -def _add_layer_contract_to_training_context( - training_context: Mapping[str, Any], - layer_contract: Optional[ResolvedLayerContract], -) -> Dict[str, Any]: - """Copy a training context and attach a resolved typed contract.""" - - _, context = _normalize_layer_contract_context( - training_context, - explicit_layer_contract=layer_contract, - ) - return context - - -def _layer_contract_kwargs( - layer_contract: Optional[ResolvedLayerContract], - *, - operator_name: Optional[str] = None, -) -> Dict[str, Any]: - """Return typed keyword arguments only for an opted-in contract path.""" - - if layer_contract is None: - return {} - kwargs: Dict[str, Any] = {"layer_contract": layer_contract} - if operator_name is not None: - kwargs["operator_name"] = operator_name - return kwargs - -def _is_moe_gating_family_model_name(model_name: str) -> bool: - base_model_name = get_moe_gating_base_model_name(model_name) - return _get_moe_family_operator_by_model_name( - base_model_name - ).precision_name() == "moe_gating" - - -def _build_exact_feature_lookup( - df: pd.DataFrame, - feature_cols: List[str], - target_col: str, -) -> Dict[Tuple[float, ...], float]: - """Build exact profiling-row lookups before falling back to regression.""" - if df.empty: - return {} - grouped = df.groupby(feature_cols, dropna=False)[target_col].mean() - lookup: Dict[Tuple[float, ...], float] = {} - for key, value in grouped.items(): - key_tuple = key if isinstance(key, tuple) else (key,) - lookup[tuple(float(item) for item in key_tuple)] = float(value) - return lookup +from frontier.execution_time_predictor.layer_contract_resolution import ( + LayerContractResolution, +) +from frontier.execution_time_predictor.prediction_family_trainers import ( + PredictionFamilyTrainers, +) +from frontier.execution_time_predictor.prediction_model_identity import ( + MIGRATION_HELP_COMMAND, + _add_layer_contract_to_training_context, + _build_exact_feature_lookup, + _get_contract_hash, + _get_moe_family_model_names, + _get_moe_family_operator_by_model_name, + _get_moe_gating_family_model_names, + _get_prefill_hot_moe_gating_model_names, + _is_moe_gating_family_model_name, + _layer_contract_kwargs, + _normalize_layer_contract_context, + _resolve_model_architecture_profile, + _resolve_model_architecture_profile_id, + _resolve_profile_typed_family_for_query, + _serialize_selected_layer_cache_identity, + _typed_row_matches_contract, + _validate_typed_parallel_selection, +) +from frontier.execution_time_predictor.prediction_model_registry import ( + PredictionModelRegistry, +) +from frontier.execution_time_predictor.profiling_dataframe_loaders import ( + ProfilingDataFrameLoaders, +) +logger = init_logger(__name__) -class ExecutionTimePredictionModelManager: +class ExecutionTimePredictionModelManager( + PredictionFamilyTrainers, + ProfilingDataFrameLoaders, + PredictionModelRegistry, + LayerContractResolution, +): """ Centralized manager for training and caching ML models used for execution time prediction. Analyzes all cluster configurations to determine the union of required prediction models. @@ -541,6 +277,7 @@ def _analyze_cluster_requirements(self) -> Dict[str, Any]: return capabilities @staticmethod + def _measurement_family_name(measurement_type: MeasurementType) -> str: if measurement_type == MeasurementType.CUDA_EVENT: return "eager" @@ -551,6 +288,7 @@ def _measurement_family_name(measurement_type: MeasurementType) -> str: raise ValueError(f"Unsupported measurement_type={measurement_type!r}") @staticmethod + def _event_measurement_type_for_replica(replica_config) -> MeasurementType: """Select an event family from the configured device metadata only.""" @@ -560,6 +298,7 @@ def _set_active_measurement_type(self, measurement_type: MeasurementType) -> Non self._active_measurement_type = measurement_type @staticmethod + def _is_kernel_only_measurement_enabled_for_cluster( cluster_type: ClusterType, ) -> bool: @@ -680,6 +419,7 @@ def _create_estimator_and_params(self, execution_time_predictor_config): return estimator, grid_search_params @staticmethod + def mean_absolute_percentage_error(y_true: np.array, y_pred: np.array) -> float: y_true, y_pred = np.array(y_true), np.array(y_pred) zero_true_mask = y_true == 0 @@ -890,3725 +630,93 @@ def get_gdn_predictor(self, cluster_type: ClusterType) -> Any | None: return self._gdn_predictors.get(cluster_type) - def _get_ffn_tp_key(self, cluster_type: ClusterType, replica_config, is_moe_model: bool) -> int: - if cluster_type == ClusterType.DECODE_FFN: - # In the FFN-only PD-AF cluster, dense FFN tensor parallelism is - # carried by moe_tensor_parallel_size as the FFN-domain TP field. - # attn_tensor_parallel_size can remain at its default because this - # cluster owns no attention weights. Use the FFN-domain TP for - # both dense and MoE DECODE_FFN profiling selection. - return replica_config.moe_tensor_parallel_size - if ( - is_moe_model - and cluster_type in { - ClusterType.PREFILL, - ClusterType.DECODE, - ClusterType.MONOLITHIC, - } - ): - return replica_config.moe_tensor_parallel_size - return replica_config.attn_tensor_parallel_size - - def _resolve_typed_layer_contract( - self, - op_name: str, - cluster_type: ClusterType, - replica_config, - *, - is_moe_model: bool, - layer_id: Optional[int] = None, - ) -> Optional[ResolvedLayerContract]: - """Resolve a typed FFN contract through the architecture profile.""" - - model_config = getattr(replica_config, "model_config", None) - architecture_profile = _resolve_model_architecture_profile(model_config) - if architecture_profile is None: - return None + def get_models(self) -> Dict[str, Dict[str, BaseEstimator]]: + """Return the trained models grouped by measurement family.""" + if self._all_dummy_mode: + logger.debug("Returning empty models dict for dummy mode") + return {"eager": {}, "kernel_only": {}} + models = { + "eager": self._models_view_for_family("eager"), + "kernel_only": self._models_view_for_family("kernel_only"), + } + device_event_models = self._models_view_for_family("device_event") + if device_event_models: + models["device_event"] = device_event_models + return models - typed_family = _resolve_profile_typed_family_for_query( - architecture_profile, op_name + def _event_family_for_cluster(self, cluster_type: ClusterType) -> str: + cluster_config = (getattr(self, "_cluster_configs", None) or {}).get( + cluster_type ) - if typed_family is None: - return None - typed_family_id, _ = typed_family - - # DECODE_ATTN is attention-only. Its exact zero domain is a deliberate - # sentinel; any non-zero value indicates a malformed configuration. - if cluster_type == ClusterType.DECODE_ATTN: - zero_fields = ( - "attn_tensor_parallel_size", - "moe_tensor_parallel_size", - "moe_expert_parallel_size", - ) - invalid = { - field_name: getattr(replica_config, field_name, None) - for field_name in zero_fields - if getattr(replica_config, field_name, None) != 0 - } - if invalid: - raise ValueError( - "DECODE_ATTN typed FFN resolution requires exact zero " - f"parallel sizes, got {invalid!r}" - ) - return None + replica_config = getattr(cluster_config, "replica_config", None) + measurement_type = self._event_measurement_type_for_replica(replica_config) + return self._measurement_family_name(measurement_type) - from frontier.operators.binding import bind_operator_query + def get_models_for_cluster(self, cluster_type: ClusterType) -> Dict[str, Dict[str, BaseEstimator]]: + """Return a cluster-specific view of trained models grouped by measurement family.""" + if self._all_dummy_mode: + return {"eager": {}, "kernel_only": {}} - binding = bind_operator_query(op_name, family_id=typed_family_id) - if binding.family_id != typed_family_id: - raise ValueError( - f"Operator query {op_name!r} resolved to family " - f"{binding.family_id!r}, expected {typed_family_id!r}" - ) + event_family = self._event_family_for_cluster(cluster_type) - moe_tp_size = getattr(replica_config, "moe_tensor_parallel_size", None) - attention_tp_size = getattr( - replica_config, "attn_tensor_parallel_size", None - ) - if cluster_type == ClusterType.DECODE_FFN: - # The FFN-only role stores its domain size in the existing MoE TP - # field, while the profile still owns the semantic TP mode. - attention_tp_size = moe_tp_size - ffn_tp_size = self._get_ffn_tp_key( - cluster_type, replica_config, is_moe_model - ) - return architecture_profile.resolve_layer_contract( - model_config, - layer_id=layer_id, - operator_name=op_name, - attention_tp_size=attention_tp_size, - moe_tp_size=moe_tp_size, - ffn_tp_size=ffn_tp_size, - expert_parallel_size=getattr( - replica_config, "moe_expert_parallel_size", None - ), - ) + def _event_models() -> Dict[str, BaseEstimator]: + return self._models_view_for_family(event_family, cluster_type) - def _resolve_ffn_layer_contracts( - self, - cluster_type: ClusterType, - replica_config, - is_moe_model: bool, - ) -> Tuple[Tuple[str, ResolvedLayerContract], ...]: - """Resolve each profile-owned FFN domain used by one training pass.""" + event_key = "eager" if event_family == "eager" else event_family - if cluster_type == ClusterType.DECODE_ATTN: - zero_fields = ( - "attn_tensor_parallel_size", - "moe_tensor_parallel_size", - "moe_expert_parallel_size", - ) - invalid = { - field_name: getattr(replica_config, field_name, None) - for field_name in zero_fields - if getattr(replica_config, field_name, None) != 0 + if cluster_type == ClusterType.PREFILL: + models = { + event_key: _event_models(), + "kernel_only": {}, } - if invalid: - raise ValueError( - "DECODE_ATTN FFN contract resolution requires exact zero " - f"parallel sizes, got {invalid!r}" - ) - return () - - model_config = getattr(replica_config, "model_config", None) - if model_config is None: - return () - architecture_profile = _resolve_model_architecture_profile(model_config) - if architecture_profile is None: - return () - if bool(is_moe_model) != bool(getattr(model_config, "is_moe", False)): - raise ValueError( - "is_moe_model does not match model configuration while resolving " - "typed FFN contracts" - ) - - contracts: list[Tuple[str, ResolvedLayerContract]] = [] - for spec in architecture_profile.iter_active_layer_contracts(model_config): - family_is_moe = spec.layer_kind is not LayerKind.DENSE - for family_id in spec.operator_family_ids: - family = get_operator_family(family_id) - profiling_ops = tuple(family.profiling_ops()) - if not profiling_ops: - raise ValueError( - f"Typed operator family {family_id!r} has no profiling operators" - ) - contract = self._resolve_typed_layer_contract( - profiling_ops[0].name, - cluster_type, - replica_config, - is_moe_model=family_is_moe, - ) - if contract is None: - raise ValueError( - f"Missing typed layer contract for operator family {family_id!r}" - ) - if contract.operator_family_id != family_id: - raise ValueError( - f"Operator family {family_id!r} resolved to " - f"{contract.operator_family_id!r}" - ) - contracts.append((family_id, contract)) - return tuple(contracts) - - def _get_ffn_contract_signature( - self, - cluster_type: ClusterType, - replica_config, - is_moe_model: bool, - ) -> str: - """Return a deterministic signature for the active FFN domains.""" - - entries = self._resolve_ffn_layer_contracts( - cluster_type, replica_config, is_moe_model - ) - if not entries: - return "none" - payload = [] - for family_id, contract in entries: - family = get_operator_family(family_id) - profiling_ops = tuple(family.profiling_ops()) - if not profiling_ops: - raise ValueError( - f"Typed operator family {family_id!r} has no profiling operators" - ) - - # The first operator is the compatibility representative returned - # by _resolve_ffn_layer_contracts(). Include every sibling as - # well: a family may mix EP-agnostic routing operators with an - # EP-sensitive grouped GEMM, and the cache signature must retain - # both semantics. - for operator in profiling_ops: - operator_contract = contract - if operator is not profiling_ops[0]: - operator_contract = self._resolve_typed_layer_contract( - operator.profiling_name(), - cluster_type, - replica_config, - is_moe_model=contract.layer_kind is not LayerKind.DENSE, - ) - if operator_contract is None: - raise ValueError( - "Missing typed layer contract for operator " - f"{operator.profiling_name()!r} in family {family_id!r}" - ) - payload.append( - { - "family_id": family_id, - "operator_name": operator.profiling_name(), - "identity": _serialize_selected_layer_cache_identity( - operator_contract - ), - } - ) - serialized = json.dumps(payload, sort_keys=True, separators=(",", ":")) - return hashlib.sha256(serialized.encode("utf-8")).hexdigest()[:16] - - @staticmethod - def _is_mixed_layer_moe_model(model_config, is_moe_model: bool) -> bool: - """Return whether a model needs both MoE and dense FFN predictors. - - Some MoE architectures keep dense FFN layers at the model boundaries. - Their runtime dispatch is layer-specific, so model-level ``is_moe`` is - insufficient to decide which predictor families must be materialized. - Keep the legacy pure-MoE path unchanged when the layer-count contract - is unavailable. - """ - if not is_moe_model or model_config is None: - return False - get_num_moe_layers = getattr(model_config, "get_num_moe_layers", None) - num_layers = getattr(model_config, "num_layers", None) - if callable(get_num_moe_layers) and isinstance(num_layers, int): - return int(get_num_moe_layers()) < int(num_layers) - return False - - def _get_linear_op_tp_key(self, op_name: str, cluster_type: ClusterType, replica_config, is_moe_model: bool) -> int: - model_config = getattr(replica_config, "model_config", None) - # Lightweight configs retain the scalar FFN compatibility path, but - # generic linear attention names still require a profile declaration - # for TP-mode lookup. - architecture_profile = _resolve_model_architecture_profile( - model_config, - allow_generic=True, - ) - if op_name in get_target_embedded_mtp_linear_ops(): - return resolve_effective_attention_tp_size( - op_name="attn_pre_proj", - requested_tp_size=replica_config.attn_tensor_parallel_size, - num_kv_heads=replica_config.model_config.num_kv_heads, - cluster_type=cluster_type, - warning_cache=getattr(self, "_attention_tp_warning_cache", None), - include_linear_ops=True, - ) - - try: - tp_mode = resolve_operator_query_tp_mode( - op_name, - architecture_profile=architecture_profile, - ) - except (TypeError, ValueError) as exc: - raise ValueError(f"Unsupported linear op for TP mapping: {op_name}") from exc - - typed_contract = self._resolve_typed_layer_contract( - op_name, - cluster_type, - replica_config, - is_moe_model=is_moe_model, - ) - if ( - typed_contract is not None - and typed_contract.tensor_parallel_size is not None - ): - return typed_contract.tensor_parallel_size - - if tp_mode is TensorParallelMode.REPLICATED: + return models + if cluster_type in [ClusterType.DECODE, ClusterType.DECODE_ATTN, ClusterType.DECODE_FFN]: if ( - is_target_embedded_mtp_enabled( - getattr(replica_config, "speculative_decoding_config", None) - ) - and is_target_embedded_mtp_same_tp_linear_op(op_name) + global_vars.get_sys_arch() == "pd-af-disaggregation" + and cluster_type == ClusterType.DECODE_ATTN ): - return resolve_effective_attention_tp_size( - op_name="attn_pre_proj", - requested_tp_size=replica_config.attn_tensor_parallel_size, - num_kv_heads=replica_config.model_config.num_kv_heads, - cluster_type=cluster_type, - warning_cache=getattr(self, "_attention_tp_warning_cache", None), - include_linear_ops=True, - ) - return 1 - - if tp_mode is TensorParallelMode.FFN_TP: - return self._get_ffn_tp_key(cluster_type, replica_config, is_moe_model) - - if tp_mode is TensorParallelMode.ATTENTION_TP: - return resolve_effective_attention_tp_size( - op_name=op_name, - requested_tp_size=replica_config.attn_tensor_parallel_size, - num_kv_heads=replica_config.model_config.num_kv_heads, - cluster_type=cluster_type, - warning_cache=getattr(self, "_attention_tp_warning_cache", None), - include_linear_ops=True, - ) - - raise ValueError(f"Unsupported linear op for TP mapping: {op_name}") - - @staticmethod - def _get_moe_op_tp_key( - op_name: str, - replica_config, - cluster_type: ClusterType | None = None, - ) -> int: - try: - return resolve_moe_operator_tp_key( - op_name, - moe_tp_size=replica_config.moe_tensor_parallel_size, - cluster_type=cluster_type, - family=MOE_FAMILY, - ) - except ValueError as exc: - if str(exc).startswith("Unsupported MoE op:"): - raise ValueError( - f"Unsupported MoE op for TP mapping: {op_name}" - ) from exc - raise - - @staticmethod - def _is_moe_op_ep_agnostic(op_name: str) -> bool: - try: - return is_moe_operator_ep_agnostic(op_name, family=MOE_FAMILY) - except ValueError as exc: - if str(exc).startswith("Unsupported MoE op:"): - raise ValueError( - f"Unsupported MoE op for EP mapping: {op_name}" - ) from exc - raise - - def _validate_moe_dataset_contract( - self, - file_path: str, - replica_config, - model_names: List[str], - cluster_type: ClusterType, - layer_contract: Optional[ResolvedLayerContract] = None, - ) -> None: - """Validate op-level MoE profiling key coverage before model training.""" - df = pd.read_csv(file_path) - required_columns = [ - "num_experts", - "router_topk", - "hidden_dim", - "expert_hidden_dim", - "num_tensor_parallel_workers", - "expert_parallel_size", - ] - missing_columns = [col for col in required_columns if col not in df.columns] - if missing_columns: - raise ValueError( - f"MoE dataset contract validation failed for {file_path}: " - f"missing required columns {missing_columns}." - ) - - model_config = replica_config.model_config - if layer_contract is None: - # A legacy caller has no profile-owned contract to validate. Keep - # the historical scalar admission rule, while refusing to guess - # when the dataset advertises typed metadata. - if TYPED_OPERATOR_CONTRACTS_COLUMN in df.columns: - raise ValueError( - "typed MoE profiling data requires an explicit routed layer contract" - ) - expected_expert_width = getattr(model_config, "mlp_hidden_dim", None) - if type(expected_expert_width) is not int or expected_expert_width <= 0: - raise ValueError( - "legacy MoE dataset validation requires a positive model_config.mlp_hidden_dim" - ) - else: - if layer_contract.layer_kind is not LayerKind.ROUTED: - raise ValueError( - "MoE dataset contract validation requires a routed layer contract" - ) - expected_expert_width = layer_contract.effective_ffn_width - base_df = df[ - (df["num_experts"] == model_config.num_experts) - & (df["router_topk"] == model_config.num_experts_per_tok) - & (df["hidden_dim"] == model_config.embedding_dim) - & (df["expert_hidden_dim"] == expected_expert_width) - ] - - if len(base_df) == 0: - raise ValueError( - "MoE dataset contract validation failed: no rows match model configuration in " - f"{file_path}. Required: num_experts={model_config.num_experts}, " - f"router_topk={model_config.num_experts_per_tok}, hidden_dim={model_config.embedding_dim}, " - f"expert_hidden_dim={expected_expert_width}." - ) - - available_pairs = sorted( - { - (int(tp), int(ep)) - for tp, ep in base_df[ - ["num_tensor_parallel_workers", "expert_parallel_size"] - ].drop_duplicates().itertuples(index=False, name=None) + models = { + event_key: _event_models(), + "kernel_only": self._models_view_for_family( + "kernel_only", cluster_type + ), + } + return models + if not self._is_kernel_only_measurement_enabled_for_cluster(cluster_type): + return { + event_key: _event_models(), + "kernel_only": {}, + } + return { + "eager": {}, + "kernel_only": self._models_view_for_family( + "kernel_only", cluster_type + ), } - ) - requested_routing_runtime_path = resolve_moe_gating_routing_runtime_path( - getattr(replica_config, "moe_routing_distribution_type", "balanced") - ) - - missing_requirements: List[str] = [] - for model_name in model_names: - tp_key = self._get_moe_op_tp_key( - model_name, - replica_config, - cluster_type, - ) - if self._is_moe_op_ep_agnostic(model_name): - op_df = base_df[ - base_df["num_tensor_parallel_workers"] == tp_key - ] - requirement = f"TP={tp_key}, EP=ANY" - else: - ep_key = replica_config.moe_expert_parallel_size - op_df = base_df[ - (base_df["num_tensor_parallel_workers"] == tp_key) - & (base_df["expert_parallel_size"] == ep_key) - ] - requirement = f"TP={tp_key}, EP={ep_key}" - if model_name == "moe_gating_routing_topk": - op_df = filter_moe_gating_routing_topk_rows( - op_df, - requested_runtime_path=requested_routing_runtime_path, - source_name=file_path, - ) - requirement = ( - f"{requirement}, routing_runtime_path=" - f"{requested_routing_runtime_path}" + if cluster_type == ClusterType.MONOLITHIC: + kernel_only_models = {} + if self._is_kernel_only_measurement_enabled_for_cluster(cluster_type): + kernel_only_models = self._models_view_for_family( + "kernel_only", cluster_type ) - if len(op_df) == 0: - missing_requirements.append(f"{model_name} requires {requirement}") - - if missing_requirements: - requirement_text = "\n - ".join(missing_requirements) - raise ValueError( - "MoE dataset contract validation failed before training.\n" - f"File: {file_path}\n" - "Missing op-level key coverage:\n" - f" - {requirement_text}\n" - f"Available (TP, EP) pairs for matched model rows: {available_pairs}" - ) - - def _train_ffn_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, - linear_ops_file: str, moe_file: str, - is_moe_model: bool, trained_model_signatures: set) -> Dict[str, BaseEstimator]: - """ - Train FFN/MoE models for a specific cluster. - - This function handles FFN-related operations in the Transformer layer: - - FFN core operations (from linear_op.csv): mlp_up_proj, mlp_down_proj, mlp_act - - MoE core operations (from moe.csv): moe_gating_linear, moe_gating_routing_topk, moe_shuffling, moe_grouped_gemm - - Pre-FFN normalization (from linear_op.csv): post_attention_layernorm + return { + event_key: _event_models(), + "kernel_only": kernel_only_models, + } + raise ValueError(f"Unsupported cluster_type={cluster_type!r}") - Transformer layer context: - ... → Attention → add → [post_attention_layernorm] → [FFN/MoE] → add → ... - """ - models = {} + def get_training_file_paths(self, cluster_type: ClusterType) -> Dict[str, str]: + """Get the resolved profiling file paths for a specific cluster type.""" + if cluster_type not in self._cluster_configs: + return {} - ffn_tp_key = self._get_ffn_tp_key(cluster_type, replica_config, is_moe_model) - tp_size = ffn_tp_key + cluster_config = self._cluster_configs[cluster_type] + replica_config = cluster_config.replica_config + execution_time_predictor_config = cluster_config.execution_time_predictor_config - # Create a signature for this FFN model configuration. - model_config = replica_config.model_config - model_arch = model_config.get_model_arch() if model_config is not None else "generic" - architecture_profile_id = _resolve_model_architecture_profile_id(model_config) - primary_contract = self._resolve_typed_layer_contract( - "moe_grouped_gemm" if is_moe_model else "mlp_up_proj", - cluster_type, - replica_config, - is_moe_model=is_moe_model, - ) - typed_contract_hash = self._get_ffn_contract_signature( - cluster_type, - replica_config, - is_moe_model, - ) - active_measurement_type = getattr( - self, "_active_measurement_type", MeasurementType.CUDA_EVENT - ) - ffn_signature = ( - f"ffn_{replica_config.device}_{replica_config.model_name}_{tp_size}" - f"_moe{is_moe_model}_arch_profile{architecture_profile_id}" - f"_layer_contracts{typed_contract_hash}" - f"_family{self._measurement_family_name(active_measurement_type)}" + return resolve_training_file_paths( + execution_time_predictor_config, + device=replica_config.device, + model=replica_config.model_config.get_name(), + network_device=replica_config.network_device, ) - - if ffn_signature in trained_model_signatures: - logger.info(f"Skipping FFN models training for {cluster_type} - already trained with signature {ffn_signature}") - return models - - # Build training context for error messages - training_context = { - 'cluster_type': str(cluster_type), - 'device': replica_config.device, - 'model_name': replica_config.model_name, - 'tensor_parallel_size': tp_size, - 'is_moe_model': is_moe_model, - 'model_arch': model_arch, - 'model_architecture_profile': architecture_profile_id, - 'use_qk_norm': bool(getattr(model_config, 'use_qk_norm', False)), - } - - # Choose input file based on model type - if is_moe_model: - moe_input_file = moe_file - if not os.path.exists(moe_input_file): - raise FileNotFoundError(f"MoE input file {moe_input_file} not found") - logger.info(f"Loading MoE data for {cluster_type} from: {moe_input_file}") - training_context['input_file'] = moe_input_file - - # MoE core operations with per-operation feature selection - # Split gating into moe_gating_linear and moe_gating_routing_topk (Step 1.6) - # Aligned with frontier/training/moe_trainer.py _get_feature_cols() method - base_moe_model_names = _get_moe_family_model_names() - moe_model_names = list(base_moe_model_names) - if should_enable_prefill_hot_moe_gating_contract( - model_config=model_config, - model_arch=model_arch, - model_name=replica_config.model_name, - ): - prefill_hot_probe_df = pd.read_csv(moe_input_file) - include_prefill_hot_models = has_prefill_hot_moe_gating_rows( - prefill_hot_probe_df - ) - - if include_prefill_hot_models: - moe_model_names.extend(_get_prefill_hot_moe_gating_model_names()) - else: - logger.warning( - "Prefill-hot gating contract enabled for model=%s, but " - "dataset %s has no usable prefill_hot rows; skipping " - "__prefill_hot pseudo-models in shared-manager training.", - replica_config.model_name, - moe_input_file, - ) - self._validate_moe_dataset_contract( - moe_input_file, - replica_config, - base_moe_model_names, - cluster_type, - **_layer_contract_kwargs(primary_contract), - ) - requested_routing_runtime_path = resolve_moe_gating_routing_runtime_path( - getattr(replica_config, "moe_routing_distribution_type", "balanced") - ) - - moe_df_cache: Dict[ - Tuple[ - int, - Optional[int], - Optional[str], - Optional[str], - Optional[str], - ], - pd.DataFrame, - ] = {} - - def _get_moe_df_for_op( - model_name: str, - ) -> Tuple[ - pd.DataFrame, - int, - Optional[int], - Optional[ResolvedLayerContract], - ]: - base_model_name = get_moe_gating_base_model_name(model_name) - op_layer_contract = self._resolve_typed_layer_contract( - base_model_name, - cluster_type, - replica_config, - is_moe_model=True, - ) - tp_key = self._get_moe_op_tp_key( - base_model_name, - replica_config, - cluster_type, - ) - if tp_key <= 0: - raise ValueError( - f"Invalid TP key for MoE training: {tp_key} (op={model_name})" - ) - - ep_key: Optional[int] - if self._is_moe_op_ep_agnostic(base_model_name): - ep_key = None - else: - ep_key = replica_config.moe_expert_parallel_size - - runtime_path_key: Optional[str] = None - if base_model_name == "moe_gating_routing_topk": - runtime_path_key = requested_routing_runtime_path - - gating_context_key: Optional[str] = None - if _is_moe_gating_family_model_name(base_model_name): - gating_context_key = DEFAULT_MOE_GATING_RUNTIME_CONTEXT - if model_name.endswith("__prefill_hot"): - gating_context_key = PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT - - contract_identity = _serialize_selected_layer_cache_identity( - op_layer_contract - ) - cache_key = ( - tp_key, - ep_key, - runtime_path_key, - gating_context_key, - contract_identity, - ) - if cache_key not in moe_df_cache: - op_df = self._load_moe_df( - moe_input_file, - replica_config, - load_imbalance=False, - tensor_parallel_size=tp_key, - expert_parallel_size=ep_key, - **_layer_contract_kwargs( - op_layer_contract, - operator_name=base_model_name, - ), - ) - if runtime_path_key is not None: - op_df = filter_moe_gating_routing_topk_rows( - op_df, - requested_runtime_path=runtime_path_key, - source_name=moe_input_file, - ) - if gating_context_key is not None: - op_df = filter_moe_gating_rows_by_runtime_context( - op_df, - requested_context=gating_context_key, - source_name=moe_input_file, - ) - moe_df_cache[cache_key] = op_df - ep_desc = "ANY" if ep_key is None else str(ep_key) - logger.info( - f"Loaded {len(moe_df_cache[cache_key])} rows for MoE training " - f"(op={model_name}, tp_key={tp_key}, ep_key={ep_desc}, " - f"routing_runtime_path={runtime_path_key or 'ANY'}, " - f"gating_runtime_context={gating_context_key or 'ANY'}, " - "auto feature mode)" - ) - return moe_df_cache[cache_key], tp_key, ep_key, op_layer_contract - - for model_name in moe_model_names: - model_signature = f"{model_name}_{ffn_signature}" - if model_signature not in trained_model_signatures: - try: - ( - op_moe_df, - moe_tp_key, - moe_ep_key, - op_layer_contract, - ) = _get_moe_df_for_op(model_name) - except PrefillHotRowsUnavailableError as exc: - logger.warning( - "Skipping %s because prefill-hot gating rows are unavailable " - "for the requested TP/EP slice (%s).", - model_name, - exc, - ) - continue - op_training_context = _add_layer_contract_to_training_context( - training_context, - op_layer_contract, - ) - op_training_context['tensor_parallel_size'] = moe_tp_key - op_training_context['expert_parallel_size'] = ( - "ANY" if moe_ep_key is None else moe_ep_key - ) - - # Per-operation feature selection. - if model_name == "moe_grouped_gemm": - available_load_features = [ - f for f in self.MOE_LOAD_IMBALANCE_FEATURES - if f in op_moe_df.columns - ] - has_load_imbalance_features = ( - len(available_load_features) - == len(self.MOE_LOAD_IMBALANCE_FEATURES) - ) - if 0 < len(available_load_features) < len(self.MOE_LOAD_IMBALANCE_FEATURES): - missing_features = [ - f for f in self.MOE_LOAD_IMBALANCE_FEATURES - if f not in op_moe_df.columns - ] - raise ValueError( - f"Partial load imbalance features found ({len(available_load_features)}/" - f"{len(self.MOE_LOAD_IMBALANCE_FEATURES)}) for {model_name} at TP={moe_tp_key}. " - f"Missing: {missing_features}." - ) - - if has_load_imbalance_features: - op_feature_cols = available_load_features - logger.info( - f" {model_name}: Using load imbalance features " - f"({len(op_feature_cols)} features, TP={moe_tp_key})" - ) - else: - op_feature_cols = ["num_tokens"] - logger.info( - f" {model_name}: Load imbalance features not found; " - f"using num_tokens only (TP={moe_tp_key})." - ) - elif model_name == "moe_shuffling": - available_load_features = [ - f for f in self.MOE_LOAD_IMBALANCE_FEATURES - if f in op_moe_df.columns - ] - if len(available_load_features) == len(self.MOE_LOAD_IMBALANCE_FEATURES): - op_feature_cols = available_load_features - logger.info( - f" {model_name}: Using load imbalance features " - f"({len(op_feature_cols)} features, TP={moe_tp_key})" - ) - else: - # For shuffling we allow partial/legacy datasets and fall back to - # num_tokens-only training when the full load feature set is absent. - op_feature_cols = ["num_tokens"] - logger.info( - f" {model_name}: Full load imbalance features unavailable; " - f"using num_tokens only (TP={moe_tp_key})." - ) - else: - op_feature_cols = ["num_tokens"] - logger.info( - f" {model_name}: Using num_tokens only (1 feature, TP={moe_tp_key})" - ) - - # Store feature_cols in training_context for this specific operation - op_training_context['feature_cols'] = op_feature_cols - - target_op_name = get_moe_gating_base_model_name(model_name) - train_kwargs: Dict[str, Any] = dict( - model_name=model_name, - df=op_moe_df, - feature_cols=op_feature_cols, - target_col=f"time_stats.{target_op_name}.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=op_training_context, - ) - train_kwargs.update(_layer_contract_kwargs(op_layer_contract)) - models[model_name] = self._train_single_model( - **train_kwargs, - ) - trained_model_signatures.add(model_signature) - logger.info(f"Trained {model_name} for {cluster_type} with features: {op_feature_cols}") - - # Step2Mini/Step3 share_expert operations (forward_3: shared expert alongside routed experts) - model_config = replica_config.model_config - if model_config is not None and model_config.supports_share_expert(): - # share_expert operations are trained from linear_op.csv (not moe.csv) - if not os.path.exists(linear_ops_file): - raise FileNotFoundError( - f"Linear ops input file {linear_ops_file} not found for share_expert" - ) - - step2mini_share_expert_model_names = list( - get_family_profiling_names(SHARE_EXPERT_FAMILY) - ) - if not step2mini_share_expert_model_names: - raise ValueError("Shared-expert operator family has no profiling names") - share_expert_tp_key = self._get_linear_op_tp_key( - step2mini_share_expert_model_names[0], - cluster_type, - replica_config, - is_moe_model, - ) - shared_layer_contract = self._resolve_typed_layer_contract( - step2mini_share_expert_model_names[0], - cluster_type, - replica_config, - is_moe_model=True, - ) - if ( - shared_layer_contract is None - and _resolve_model_architecture_profile(model_config) is not None - ): - raise ValueError( - "Missing shared layer contract for share-expert training" - ) - share_expert_linear_ops_df = self._load_linear_op_df( - linear_ops_file, - share_expert_tp_key, - **_layer_contract_kwargs( - shared_layer_contract, - operator_name=step2mini_share_expert_model_names[0], - ), - ) - logger.info(f"Loaded {len(share_expert_linear_ops_df)} rows for share_expert training") - - for model_name in step2mini_share_expert_model_names: - model_signature = f"{model_name}_{ffn_signature}" - if model_signature not in trained_model_signatures: - # Update training context to reflect linear_op.csv source. - shared_training_context = _add_layer_contract_to_training_context( - training_context, - shared_layer_contract, - ) - shared_training_context['input_file'] = linear_ops_file - shared_training_context['tensor_parallel_size'] = share_expert_tp_key - target_col = f"time_stats.{model_name}.median" - if target_col not in share_expert_linear_ops_df.columns: - raise ValueError( - f"share_expert operation '{model_name}' column '{target_col}' not found in profiling data. " - f"Ensure profiling was run with a model architecture that includes share_expert. " - f"Available columns: {list(share_expert_linear_ops_df.columns)}" - ) - train_kwargs: Dict[str, Any] = dict( - model_name=model_name, - df=share_expert_linear_ops_df, - feature_cols=["num_tokens"], - target_col=target_col, - execution_time_predictor_config=execution_time_predictor_config, - training_context=shared_training_context, - ) - train_kwargs.update( - _layer_contract_kwargs(shared_layer_contract) - ) - models[model_name] = self._train_single_model( - **train_kwargs, - ) - trained_model_signatures.add(model_signature) - logger.info(f"Trained {model_name} for {cluster_type}") - - # Mixed-layer MoE models (for example step-moe-noquant) also have - # dense boundary layers. Train these additions after the legacy - # MoE/share-expert families so RandomForest training order remains - # compatible with the historical predictor artifact contract. - if self._is_mixed_layer_moe_model(model_config, is_moe_model): - dense_ffn_tp_key = self._get_ffn_tp_key( - cluster_type, replica_config, is_moe_model=False - ) - dense_layer_contract = self._resolve_typed_layer_contract( - "mlp_up_proj", - cluster_type, - replica_config, - is_moe_model=False, - ) - dense_ffn_signature = ( - f"ffn_{replica_config.device}_{replica_config.model_name}_{dense_ffn_tp_key}" - f"_moeFalse_arch_profile{architecture_profile_id}" - f"_layer_contracts{_get_contract_hash(dense_layer_contract)}" - f"_family{self._measurement_family_name(active_measurement_type)}" - ) - dense_training_context = dict(training_context) - dense_training_context["is_moe_model"] = False - dense_training_context["tensor_parallel_size"] = dense_ffn_tp_key - dense_training_context = _add_layer_contract_to_training_context( - dense_training_context, - dense_layer_contract, - ) - self._train_dense_mlp_models_for_cluster( - cluster_type=cluster_type, - replica_config=replica_config, - execution_time_predictor_config=execution_time_predictor_config, - linear_ops_file=linear_ops_file, - ffn_signature=dense_ffn_signature, - ffn_tp_key=dense_ffn_tp_key, - training_context=dense_training_context, - trained_model_signatures=trained_model_signatures, - models=models, - layer_contract=dense_layer_contract, - ) - else: - self._train_dense_mlp_models_for_cluster( - cluster_type=cluster_type, - replica_config=replica_config, - execution_time_predictor_config=execution_time_predictor_config, - linear_ops_file=linear_ops_file, - ffn_signature=ffn_signature, - ffn_tp_key=ffn_tp_key, - training_context=training_context, - trained_model_signatures=trained_model_signatures, - models=models, - layer_contract=primary_contract, - ) - - # Pre-FFN normalization (post_attention_layernorm) - always from linear_op.csv - if not os.path.exists(linear_ops_file): - raise FileNotFoundError(f"Linear ops input file {linear_ops_file} not found for post_attention_layernorm") - layernorm_tp_key = self._get_linear_op_tp_key( - "post_attention_layernorm", - cluster_type, - replica_config, - is_moe_model, - ) - linear_ops_df = self._load_linear_op_df(linear_ops_file, layernorm_tp_key) - layernorm_context = dict(training_context) - layernorm_context["input_file"] = linear_ops_file - layernorm_context["tensor_parallel_size"] = layernorm_tp_key - - layernorm_model_name = "post_attention_layernorm" - layernorm_signature = f"{layernorm_model_name}_{ffn_signature}" - if layernorm_signature not in trained_model_signatures: - models[layernorm_model_name] = self._train_single_model( - model_name=layernorm_model_name, - df=linear_ops_df, - feature_cols=["num_tokens"], - target_col=f"time_stats.{layernorm_model_name}.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=layernorm_context, - ) - trained_model_signatures.add(layernorm_signature) - logger.info(f"Trained {layernorm_model_name} for {cluster_type}") - - # Mark this FFN configuration as trained - trained_model_signatures.add(ffn_signature) - return models - - def _train_dense_mlp_models_for_cluster( - self, - *, - cluster_type: ClusterType, - replica_config, - execution_time_predictor_config, - linear_ops_file: str, - ffn_signature: str, - ffn_tp_key: int, - training_context: Dict[str, Any], - trained_model_signatures: set, - models: Dict[str, BaseEstimator], - layer_contract: Optional[ResolvedLayerContract] = None, - ) -> None: - """Materialize dense MLP predictors from the linear-op profile.""" - if not os.path.exists(linear_ops_file): - raise FileNotFoundError(f"Linear ops input file {linear_ops_file} not found") - if ( - layer_contract is None - and _resolve_model_architecture_profile( - getattr(replica_config, "model_config", None) - ) - is not None - ): - layer_contract = self._resolve_typed_layer_contract( - "mlp_up_proj", - cluster_type, - replica_config, - is_moe_model=False, - ) - if ( - layer_contract is not None - and layer_contract.tensor_parallel_size is not None - and layer_contract.tensor_parallel_size != ffn_tp_key - ): - raise ValueError( - "Dense FFN training TP conflicts with its typed layer contract: " - f"ffn_tp_key={ffn_tp_key}, " - f"contract_tp={layer_contract.tensor_parallel_size}" - ) - logger.info(f"Loading MLP data for {cluster_type} from: {linear_ops_file}") - dense_model_names = tuple(get_family_profiling_names(FFN_FAMILY)) - if not dense_model_names: - raise ValueError("FFN operator family has no profiling names") - linear_ops_df = self._load_linear_op_df( - linear_ops_file, - ffn_tp_key, - **_layer_contract_kwargs( - layer_contract, - operator_name=dense_model_names[0], - ), - ) - logger.info(f"Loaded {len(linear_ops_df)} rows for MLP training") - dense_training_context = _add_layer_contract_to_training_context( - training_context, - layer_contract, - ) - dense_training_context["input_file"] = linear_ops_file - dense_training_context["tensor_parallel_size"] = ffn_tp_key - - missing_standard_columns = [ - f"time_stats.{model_name}.median" - for model_name in dense_model_names - if f"time_stats.{model_name}.median" not in linear_ops_df.columns - ] - if missing_standard_columns: - model_config = getattr(replica_config, "model_config", None) - supports_share_expert = bool( - model_config is not None - and model_config.supports_share_expert() - ) - if supports_share_expert: - logger.info( - "Skipping standard dense MLP training for %s: profile provides " - "shared-expert operations instead; missing columns=%s", - cluster_type, - missing_standard_columns, - ) - return - raise ValueError( - "Dense MLP profiling data is incomplete; missing columns: " - + ", ".join(missing_standard_columns) - ) - - for model_name in dense_model_names: - model_signature = f"{model_name}_{ffn_signature}" - if model_signature in trained_model_signatures: - continue - train_kwargs: Dict[str, Any] = dict( - model_name=model_name, - df=linear_ops_df, - feature_cols=["num_tokens"], - target_col=f"time_stats.{model_name}.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=dense_training_context, - ) - train_kwargs.update(_layer_contract_kwargs(layer_contract)) - models[model_name] = self._train_single_model(**train_kwargs) - trained_model_signatures.add(model_signature) - logger.info(f"Trained {model_name} for {cluster_type}") - - def _train_attn_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, replica_scheduler_config, linear_ops_file: str, attn_file: str, trained_model_signatures: set) -> Dict[str, BaseEstimator]: - """ - Train attention-related models for a cluster. - - This function handles Attention-related operations in the Transformer layer: - - Pre-attention normalization (from linear_op.csv): input_layernorm - - Attention projections (from linear_op.csv): attn_pre_proj, attn_post_proj, attn_rope - - Attention core operations (from attention.csv): attn_kv_cache_save, attn_prefill, attn_decode - - Transformer layer context: - Input → [input_layernorm] → [attn_pre_proj → attn_rope → attn_prefill/decode → attn_kv_cache_save → attn_post_proj] → add → ... - """ - models = {} - tp_size = replica_config.attn_tensor_parallel_size - - model_config = replica_config.model_config - model_arch = model_config.get_model_arch() if model_config is not None else "generic" - architecture_profile_id = _resolve_model_architecture_profile_id(model_config) - attention_signature = ( - f"attention_{replica_config.device}_{replica_config.model_name}_{tp_size}" - f"_{replica_scheduler_config.block_size}_arch_profile{architecture_profile_id}" - f"_family{self._measurement_family_name(self._active_measurement_type)}" - ) - - if attention_signature in trained_model_signatures: - logger.info(f"Skipping attention models training for {cluster_type} - already trained") - return models - - # Build training context for error messages - training_context = { - 'cluster_type': str(cluster_type), - 'device': replica_config.device, - 'model_name': replica_config.model_name, - 'tensor_parallel_size': tp_size, - 'block_size': replica_scheduler_config.block_size, - 'model_arch': model_arch, - 'model_architecture_profile': architecture_profile_id, - 'use_qk_norm': bool(getattr(model_config, 'use_qk_norm', False)), - } - - # ========== Part 1: Linear operations from linear_op.csv ========== - # These include: input_layernorm, attn_pre_proj, attn_post_proj, attn_rope - if not os.path.exists(linear_ops_file): - raise FileNotFoundError(f"Linear ops input file {linear_ops_file} not found") - - logger.info(f"Loading sharded attention linear-op data from: {linear_ops_file}") - attn_tp_key = self._get_linear_op_tp_key( - "attn_pre_proj", - cluster_type, - replica_config, - is_moe_model=False, - ) - required_columns = self._get_required_attn_linear_op_columns(model_config) - attn_linear_ops_df = self._load_linear_op_df( - linear_ops_file, - attn_tp_key, - required_columns=required_columns, - training_context=training_context, - ) - logger.info( - f"Loaded {len(attn_linear_ops_df)} rows for sharded attention ops training" - ) - - # Pre-attention normalization: input_layernorm - input_layernorm_tp_key = self._get_linear_op_tp_key( - "input_layernorm", - cluster_type, - replica_config, - is_moe_model=False, - ) - input_layernorm_df = self._load_linear_op_df( - linear_ops_file, - input_layernorm_tp_key, - required_columns=["time_stats.input_layernorm.median"], - training_context=training_context, - ) - input_layernorm_context = dict(training_context) - input_layernorm_context["input_file"] = linear_ops_file - input_layernorm_context["tensor_parallel_size"] = input_layernorm_tp_key - - layernorm_model_name = "input_layernorm" - layernorm_signature = f"{layernorm_model_name}_{attention_signature}" - if layernorm_signature not in trained_model_signatures: - models[layernorm_model_name] = self._train_single_model( - model_name=layernorm_model_name, - df=input_layernorm_df, - feature_cols=["num_tokens"], - target_col=f"time_stats.{layernorm_model_name}.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=input_layernorm_context, - ) - trained_model_signatures.add(layernorm_signature) - logger.info(f"Trained {layernorm_model_name} for {cluster_type}") - - # Attention projections: attn_pre_proj, attn_post_proj, attn_rope - attn_proj_context = dict(training_context) - attn_proj_context["input_file"] = linear_ops_file - attn_proj_context["tensor_parallel_size"] = attn_tp_key - attn_proj_model_names = ["attn_pre_proj", "attn_post_proj", "attn_rope"] - for model_name in attn_proj_model_names: - model_signature = f"{model_name}_{attention_signature}" - if model_signature not in trained_model_signatures: - models[model_name] = self._train_single_model( - model_name=model_name, - df=attn_linear_ops_df, - feature_cols=["num_tokens"], - target_col=f"time_stats.{model_name}.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=attn_proj_context, - ) - trained_model_signatures.add(model_signature) - logger.info(f"Trained {model_name} for {cluster_type}") - - if is_target_embedded_mtp_enabled( - getattr(replica_config, "speculative_decoding_config", None) - ): - required_mtp_columns = ( - self._get_required_target_embedded_mtp_linear_op_columns() - ) - missing_mtp_columns = [ - col for col in required_mtp_columns if col not in attn_linear_ops_df.columns - ] - all_nan_mtp_columns = [ - col - for col in required_mtp_columns - if col in attn_linear_ops_df.columns - and attn_linear_ops_df[col].isna().all() - ] - if missing_mtp_columns or all_nan_mtp_columns: - raise ValueError( - "target-embedded MTP compute profiling columns are missing or all-NaN in " - f"{linear_ops_file}. " - f"Missing columns: {missing_mtp_columns}. " - f"All-NaN columns: {all_nan_mtp_columns}. " - "Re-run linear-op profiling with --include_target_embedded_mtp." - ) - for model_name in ["mtp_fusion_proj", "lm_head_linear"]: - model_signature = f"{model_name}_{attention_signature}" - if model_signature not in trained_model_signatures: - models[model_name] = self._train_single_model( - model_name=model_name, - df=attn_linear_ops_df, - feature_cols=["num_tokens"], - target_col=f"time_stats.{model_name}.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=attn_proj_context, - ) - trained_model_signatures.add(model_signature) - logger.info( - "Trained %s for %s (target-embedded MTP)", - model_name, - cluster_type, - ) - - model_config = replica_config.model_config - architecture_profile = _resolve_model_architecture_profile(model_config) - predictor_attention_extra_ops = ( - architecture_profile.predictor_attention_extra_ops - if architecture_profile is not None - else () - ) - for model_name in predictor_attention_extra_ops: - model_signature = f"{model_name}_{attention_signature}" - if model_signature not in trained_model_signatures: - target_col = f"time_stats.{model_name}.median" - if target_col not in attn_linear_ops_df.columns: - raise ValueError( - f"Architecture-profile operation '{model_name}' column '{target_col}' not found in profiling data. " - f"Ensure profiling was run with the selected model architecture profile. " - f"Available columns: {list(attn_linear_ops_df.columns)}" - ) - models[model_name] = self._train_single_model( - model_name=model_name, - df=attn_linear_ops_df, - feature_cols=["num_tokens"], - target_col=target_col, - execution_time_predictor_config=execution_time_predictor_config, - training_context=attn_proj_context, - ) - trained_model_signatures.add(model_signature) - logger.info("Trained architecture-profile %s for %s", model_name, cluster_type) - - # ========== Part 2: Attention core operations from attention.csv ========== - if not os.path.exists(attn_file): - raise FileNotFoundError(f"Attention input file {attn_file} not found") - - logger.info(f"Loading attention data from: {attn_file}") - attention_df = self._load_attention_df( - attn_file, - replica_config, - replica_scheduler_config, - cluster_type=cluster_type, - ) - training_context['input_file'] = attn_file - - # Family-aware attention-core training. Latent-MLA profiles carry six - # ``attn_mla_*`` operators with a structural layout the dense block cannot - # consume; route them through the MLA branch before the dense derive (which - # assumes dense feature columns such as ``prefill_chunk_size``). - if self._is_mla_family(replica_config.model_config): - attention_df = self._get_mla_attention_df_with_derived_features( - attention_df - ) - logger.info( - f"Loaded {len(attention_df)} rows for latent-MLA attention core training" - ) - models.update( - self._train_mla_attention_core_models( - attention_df=attention_df, - attention_signature=attention_signature, - cluster_type=cluster_type, - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - trained_model_signatures=trained_model_signatures, - ) - ) - trained_model_signatures.add(attention_signature) - return models - - attention_df = self._get_attention_df_with_derived_features(attention_df) - logger.info(f"Loaded {len(attention_df)} rows for attention core training") - measurement_type = self._active_measurement_type - dense_attention_model_names = get_enabled_predictor_metric_names( - DENSE_ATTENTION_FAMILY - ) - dense_attention_target_columns = dict( - zip( - dense_attention_model_names, - get_enabled_predictor_median_columns(DENSE_ATTENTION_FAMILY), - ) - ) - dense_attention_feature_columns = get_enabled_shared_predictor_feature_columns( - DENSE_ATTENTION_FAMILY - ) - - # Train kv_cache_save model - kv_cache_model_name = get_enabled_predictor_metric_name_by_role( - DENSE_ATTENTION_FAMILY, - AttentionOperatorRole.CACHE_WRITE, - ) - kv_cache_model_signature = f"{kv_cache_model_name}_{attention_signature}" - if kv_cache_model_signature not in trained_model_signatures: - kv_cache_feature_cols = list( - dense_attention_feature_columns[kv_cache_model_name] - ) - missing_cols = [ - col for col in kv_cache_feature_cols if col not in attention_df.columns - ] - if missing_cols: - raise ValueError( - f"Missing columns for {kv_cache_model_name} training: {missing_cols}. " - "Re-run attention profiling with mixed-batch metadata." - ) - models[kv_cache_model_name] = self._train_single_model( - model_name=kv_cache_model_name, - df=attention_df, - feature_cols=kv_cache_feature_cols, - target_col=dense_attention_target_columns[kv_cache_model_name], - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - persist_exact_lookup=True, - ) - trained_model_signatures.add(kv_cache_model_signature) - logger.info(f"Trained {kv_cache_model_name} for {cluster_type}") - - # Split data for prefill and decode. - # Mixed-batch prefill rows in attention_combined.csv use prefill_chunk_size=0, - # so standard prefill training must keep only rows with positive chunk size. - true_mixed_df = attention_df[attention_df["is_true_mixed_batch"]].copy() - standard_df = attention_df[~attention_df["is_true_mixed_batch"]].copy() - prefill_df = standard_df[~standard_df["is_decode"]].copy() - decode_df = standard_df[standard_df["is_decode"]].copy() - standard_prefill_df = pd.DataFrame() - if measurement_type in (MeasurementType.CUDA_EVENT, MeasurementType.DEVICE_EVENT): - if "prefill_chunk_size" not in prefill_df.columns: - raise ValueError( - "Missing required column 'prefill_chunk_size' in attention profiling data." - ) - standard_prefill_df = prefill_df[prefill_df["prefill_chunk_size"] > 0].copy() - - prefill_model_name = get_enabled_predictor_metric_name_by_role( - DENSE_ATTENTION_FAMILY, - AttentionOperatorRole.PREFILL_KERNEL, - ) - prefill_model_signature = f"{prefill_model_name}_{attention_signature}" - if prefill_model_signature not in trained_model_signatures: - if len(standard_prefill_df) == 0: - raise ValueError( - "No standard prefill rows (prefill_chunk_size > 0) found in eager attention profiling data." - ) - models[prefill_model_name] = self._train_single_model( - model_name=prefill_model_name, - df=standard_prefill_df, - feature_cols=list(dense_attention_feature_columns[prefill_model_name]), - target_col=dense_attention_target_columns[prefill_model_name], - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - ) - trained_model_signatures.add(prefill_model_signature) - logger.info(f"Trained {prefill_model_name} for {cluster_type}") - - decode_model_name = get_enabled_predictor_metric_name_by_role( - DENSE_ATTENTION_FAMILY, - AttentionOperatorRole.DECODE_KERNEL, - ) - decode_model_signature = f"{decode_model_name}_{attention_signature}" - if decode_model_signature not in trained_model_signatures: - if len(decode_df) == 0: - logger.info( - "Skipping eager %s training for %s - no standard decode rows", - decode_model_name, - cluster_type, - ) - else: - decode_feature_cols = list( - dense_attention_feature_columns[decode_model_name] - ) - missing_decode_cols = [ - col - for col in [ - *decode_feature_cols, - dense_attention_target_columns[decode_model_name], - ] - if col not in decode_df.columns - ] - if missing_decode_cols: - logger.info( - "Skipping eager %s training for %s - missing decode feature columns %s", - decode_model_name, - cluster_type, - missing_decode_cols, - ) - else: - models[decode_model_name] = self._train_single_model( - model_name=decode_model_name, - df=decode_df, - feature_cols=decode_feature_cols, - target_col=dense_attention_target_columns[decode_model_name], - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - ) - trained_model_signatures.add(decode_model_signature) - logger.info(f"Trained eager {decode_model_name} for {cluster_type}") - elif measurement_type == MeasurementType.KERNEL_ONLY: - decode_model_name = get_enabled_predictor_metric_name_by_role( - DENSE_ATTENTION_FAMILY, - AttentionOperatorRole.DECODE_KERNEL, - ) - decode_model_signature = f"{decode_model_name}_{attention_signature}" - if decode_model_signature not in trained_model_signatures: - if len(decode_df) == 0: - raise ValueError( - "No standard decode rows found in kernel-only attention profiling data." - ) - models[decode_model_name] = self._train_single_model( - model_name=decode_model_name, - df=decode_df, - feature_cols=list(dense_attention_feature_columns[decode_model_name]), - target_col=dense_attention_target_columns[decode_model_name], - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - ) - trained_model_signatures.add(decode_model_signature) - logger.info(f"Trained {decode_model_name} for {cluster_type}") - else: - raise ValueError(f"Unsupported measurement_type={measurement_type!r}") - - # ========== Part 3: Mixed-batch prefill model (optional, high-dimensional) ========== - # attn_prefill_mixed uses 12 features and requires on-demand prediction at runtime - # Check if profiling data contains mixed-batch features - mixed_batch_model_signature = f"attn_prefill_mixed_{attention_signature}" - if measurement_type in (MeasurementType.CUDA_EVENT, MeasurementType.DEVICE_EVENT) and mixed_batch_model_signature not in trained_model_signatures: - # Check for mixed-batch specific columns in the dataframe - required_mixed_features = self.ATTN_PREFILL_MIXED_FEATURES - has_mixed_batch_data = all(feat in prefill_df.columns for feat in required_mixed_features) - - if has_mixed_batch_data: - logger.info(f"Training attn_prefill_mixed with {len(required_mixed_features)} features for {cluster_type}") - - # Filter for mixed-prefill rows (exclude true mixed prefill+decode rows) - mixed_batch_df = prefill_df[ - prefill_df["is_mixed_batch"] | (prefill_df["batch_size"] > 1) - ].copy() - - if len(mixed_batch_df) > 0: - models["attn_prefill_mixed"] = self._train_single_model( - model_name="attn_prefill_mixed", - df=mixed_batch_df, - feature_cols=required_mixed_features, - target_col="time_stats.attn_prefill.median", # Same target column as attn_prefill - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - persist_exact_lookup=True, - ) - trained_model_signatures.add(mixed_batch_model_signature) - logger.info(f"Trained attn_prefill_mixed with {len(mixed_batch_df)} samples for {cluster_type}") - else: - logger.warning(f"No mixed-batch data (batch_size > 1) available for attn_prefill_mixed in {cluster_type}") - else: - missing_features = [f for f in required_mixed_features if f not in prefill_df.columns] - logger.info(f"Skipping attn_prefill_mixed for {cluster_type} - missing features: {missing_features}") - - decode_in_mixed_signature = f"attn_decode_in_mixed_{attention_signature}" - if measurement_type in (MeasurementType.CUDA_EVENT, MeasurementType.DEVICE_EVENT) and decode_in_mixed_signature not in trained_model_signatures: - required_decode_mixed_features = self.ATTN_DECODE_IN_MIXED_FEATURES - has_decode_mixed_data = all( - feat in true_mixed_df.columns for feat in required_decode_mixed_features - ) - if has_decode_mixed_data: - if len(true_mixed_df) > 0: - models["attn_decode_in_mixed"] = self._train_single_model( - model_name="attn_decode_in_mixed", - df=true_mixed_df, - feature_cols=required_decode_mixed_features, - target_col="time_stats.attn_decode.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - persist_exact_lookup=True, - ) - trained_model_signatures.add(decode_in_mixed_signature) - logger.info( - f"Trained attn_decode_in_mixed with {len(true_mixed_df)} samples for {cluster_type}" - ) - else: - logger.info( - f"Skipping attn_decode_in_mixed for {cluster_type} - no true mixed rows" - ) - else: - missing_features = [ - f for f in required_decode_mixed_features if f not in true_mixed_df.columns - ] - logger.info( - f"Skipping attn_decode_in_mixed for {cluster_type} - missing features: {missing_features}" - ) - - trained_model_signatures.add(attention_signature) - return models - - @staticmethod - def _is_mla_family(model_config) -> bool: - """Return True when the model binds to the latent-MLA attention family.""" - if model_config is None: - return False - return ( - resolve_runtime_attention_family(model_config).family_id - == LATENT_MLA_ATTENTION_FAMILY.family_id - ) - - def _get_mla_attention_df_with_derived_features( - self, df: pd.DataFrame - ) -> pd.DataFrame: - """Derive latent-MLA attention features (normalize ``is_prefill`` to int). - - Mirrors the monolithic ``SklearnExecutionTimePredictor`` MLA early-return: - latent-MLA training keys on the imported structural columns directly and must - NOT add the dense ``num_tokens`` / ``prefill_chunk_size`` derived features. - """ - df_with_derived_features = df.copy() - if "is_prefill" in df_with_derived_features.columns: - df_with_derived_features["is_prefill"] = coerce_truthy_int( - df_with_derived_features["is_prefill"] - ) - return df_with_derived_features - - def _filter_mla_attention_df( - self, - df: pd.DataFrame, - file_path: str, - replica_config, - replica_scheduler_config, - ) -> pd.DataFrame: - """Filter an imported latent-MLA profile to the requested structural layout. - - Verbatim port of the monolithic ``_filter_mla_attention_df`` (adapted to the - shared-manager's per-cluster ``replica_config`` / ``replica_scheduler_config`` - instead of instance state). Fail-fast on missing structural columns or an empty - post-filter frame per §7 (no silent fallback). - """ - validate_attention_profiling_dataframe( - df, - LATENT_MLA_ATTENTION_FAMILY, - measurement_type=self._active_measurement_type, - ) - - model_config = replica_config.model_config - expected_values = { - "n_q_head": int(getattr(model_config, "num_q_heads")), - "n_kv_head": int( - model_config.get_runtime_num_kv_heads() - if hasattr(model_config, "get_runtime_num_kv_heads") - else 1 - ), - "head_size": int( - model_config.get_runtime_head_size() - if hasattr(model_config, "get_runtime_head_size") - else int(getattr(model_config, "kv_lora_rank")) - + int(getattr(model_config, "qk_rope_head_dim")) - ), - "qk_nope_head_dim": int(getattr(model_config, "qk_nope_head_dim")), - "qk_rope_head_dim": int(getattr(model_config, "qk_rope_head_dim")), - "qk_head_dim": int(model_config.get_qk_head_dim()), - "kv_lora_rank": int(getattr(model_config, "kv_lora_rank")), - "v_head_dim": int(getattr(model_config, "v_head_dim")), - "block_size": int(replica_scheduler_config.block_size), - "num_tensor_parallel_workers": int( - replica_config.attn_tensor_parallel_size - ), - } - missing_columns = [ - column for column in expected_values if column not in df.columns - ] - if missing_columns: - raise ValueError( - "MLA attention profiling data is missing structural columns: " - f"{missing_columns}. file={file_path}" - ) - - filtered = df.copy() - for column, expected_value in expected_values.items(): - filtered = filtered[filtered[column].astype(int) == expected_value] - - if filtered.empty: - raise ValueError( - "No MLA attention profiling rows remain after structural filtering. " - f"file={file_path}, expected={expected_values}" - ) - return filtered - - def _train_mla_attention_core_models( - self, - attention_df: pd.DataFrame, - attention_signature: str, - cluster_type: ClusterType, - execution_time_predictor_config, - training_context: Dict[str, Any], - trained_model_signatures: set, - ) -> Dict[str, BaseEstimator]: - """Train the six latent-MLA attention-core operators (training-only A2 fix). - - Mirrors the monolithic ``_train_mla_attention_layer_models`` (sparse-by-target - row filtering + exact-row memoization). The on-demand consumer pairs each - estimator's ``_frontier_exact_lookup`` with the same module-level builder, so - the disaggregation prediction path works unchanged once these models exist. - """ - model_names = list( - get_enabled_predictor_metric_names(LATENT_MLA_ATTENTION_FAMILY) - ) - target_columns = dict( - zip( - model_names, - get_enabled_predictor_median_columns(LATENT_MLA_ATTENTION_FAMILY), - ) - ) - feature_columns = get_enabled_shared_predictor_feature_columns( - LATENT_MLA_ATTENTION_FAMILY - ) - - models: Dict[str, BaseEstimator] = {} - for model_name in model_names: - model_signature = f"{model_name}_{attention_signature}" - if model_signature in trained_model_signatures: - continue - - feature_cols = list(feature_columns[model_name]) - target_col = target_columns[model_name] - required_columns = [*feature_cols, target_col] - missing_columns = [ - column - for column in required_columns - if column not in attention_df.columns - ] - all_nan_columns = [ - column - for column in required_columns - if column in attention_df.columns - and attention_df[column].isna().all() - ] - if missing_columns or all_nan_columns: - raise ValueError( - "MLA attention profiling data cannot train " - f"{model_name}." - f"\nMissing columns: {missing_columns}" - f"\nAll-NaN columns: {all_nan_columns}" - ) - - op_attention_df = attention_df.dropna(subset=[target_col]).copy() - if op_attention_df.empty: - raise ValueError( - "MLA attention profiling data cannot train " - f"{model_name}: target column {target_col!r} has no " - "observed timing rows." - ) - nan_feature_columns = [ - column - for column in feature_cols - if op_attention_df[column].isna().any() - ] - if nan_feature_columns: - raise ValueError( - "MLA attention profiling data cannot train " - f"{model_name}: feature columns contain NaN after " - f"target filtering: {nan_feature_columns}" - ) - - model = self._train_single_model( - model_name=model_name, - df=op_attention_df, - feature_cols=feature_cols, - target_col=target_col, - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - persist_exact_lookup=True, - ) - if not hasattr(model, "_frontier_exact_lookup"): - model._frontier_exact_lookup = _build_exact_feature_lookup( - op_attention_df, - feature_cols, - target_col, - ) - models[model_name] = model - trained_model_signatures.add(model_signature) - logger.info(f"Trained {model_name} for {cluster_type}") - - return models - - def _train_residual_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, - linear_ops_file: str, trained_model_signatures: set) -> Dict[str, BaseEstimator]: - """ - Train residual connection models for a cluster. - - This function handles residual connection operations in the Transformer layer: - - Residual add operation (from linear_op.csv): add - - Transformer layer context: - ... → Attention → [add] → LayerNorm → FFN/MoE → [add] → ... - - The residual add operation is used after both Attention and FFN blocks, - making it a common operation that serves both sub-layers. - """ - models = {} - - model_config = replica_config.model_config - - # RMSNorm: add is fused into layernorm, no separate add model needed - if model_config is not None and model_config.uses_fused_add_norm: - logger.info(f"Skipping residual add model training for {cluster_type} " - f"— model uses fused add+norm (RMSNorm)") - return models - - is_moe_model = model_config is not None and model_config.is_moe - tp_size = self._get_linear_op_tp_key( - "add", - cluster_type, - replica_config, - is_moe_model, - ) - - # Create a signature for this residual model configuration - residual_signature = f"residual_{replica_config.device}_{replica_config.model_name}_{tp_size}_family{self._measurement_family_name(self._active_measurement_type)}" - - if residual_signature in trained_model_signatures: - logger.info(f"Skipping residual models training for {cluster_type} - already trained with signature {residual_signature}") - return models - - if not os.path.exists(linear_ops_file): - raise FileNotFoundError(f"Linear ops input file {linear_ops_file} not found for residual models") - - logger.info(f"Loading linear ops data for residual models from: {linear_ops_file}") - linear_ops_df = self._load_linear_op_df(linear_ops_file, tp_size) - logger.info(f"Loaded {len(linear_ops_df)} rows for residual training") - - # Build training context for error messages - training_context = { - 'cluster_type': str(cluster_type), - 'device': replica_config.device, - 'model_name': replica_config.model_name, - 'tensor_parallel_size': tp_size, - 'input_file': linear_ops_file, - } - - # Train the residual add model - add_model_name = "add" - add_signature = f"{add_model_name}_{residual_signature}" - if add_signature not in trained_model_signatures: - models[add_model_name] = self._train_single_model( - model_name=add_model_name, - df=linear_ops_df, - feature_cols=["num_tokens"], - target_col=f"time_stats.{add_model_name}.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - ) - trained_model_signatures.add(add_signature) - logger.info(f"Trained {add_model_name} for {cluster_type}") - - # Mark this residual configuration as trained - trained_model_signatures.add(residual_signature) - return models - - def _train_pipeline_parallel_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, trained_model_signatures: set) -> Dict[str, BaseEstimator]: - """Train pipeline parallel communication models for a cluster.""" - models = {} - - _, _, _, send_recv_input_file, _, _ = self._get_input_files_for_config(replica_config, execution_time_predictor_config) - - pp_signature = f"send_recv_{replica_config.network_device}_{replica_config.num_pipeline_stages}_{replica_config.attn_tensor_parallel_size}_family{self._measurement_family_name(self._active_measurement_type)}" - - if pp_signature in trained_model_signatures: - logger.info(f"Skipping send_recv model training for {cluster_type} - already trained") - return models - - send_recv_df = self._load_send_recv_df(send_recv_input_file, replica_config) - send_recv_df = self._get_send_recv_df_with_derived_features(send_recv_df, replica_config) - - # Build training context for error messages - training_context = { - 'cluster_type': str(cluster_type), - 'device': replica_config.device, - 'model_name': replica_config.model_name, - 'pipeline_stages': replica_config.num_pipeline_stages, - 'tensor_parallel_size': replica_config.attn_tensor_parallel_size, - 'network_device': replica_config.network_device, - 'input_file': send_recv_input_file, - } - - models["send_recv"] = self._train_single_model( - model_name="send_recv", - df=send_recv_df, - feature_cols=["num_tokens"], - target_col="time_stats.send_recv.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - ) - - trained_model_signatures.add(pp_signature) - return models - - def _train_tensor_parallel_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, use_attn_tp: bool, trained_model_signatures: set) -> Dict[str, BaseEstimator]: - """Train tensor parallel communication models for a cluster.""" - models = {} - - _, _, all_reduce_input_file, _, _, _ = self._get_input_files_for_config(replica_config, execution_time_predictor_config) - - # Use different tensor parallel size based on cluster type - tp_size = replica_config.attn_tensor_parallel_size if use_attn_tp else replica_config.moe_tensor_parallel_size - - tp_signature = f"all_reduce_{replica_config.network_device}_{tp_size}_family{self._measurement_family_name(self._active_measurement_type)}" - - if tp_signature in trained_model_signatures: - logger.info(f"Skipping all_reduce model training for {cluster_type} - already trained") - return models - - # 添加详细的上下文信息 - training_context = { - 'cluster_type': cluster_type, - 'device': replica_config.device, - 'model_name': replica_config.model_name, - 'tensor_parallel_size': tp_size, - 'network_device': replica_config.network_device, - 'input_file': all_reduce_input_file, - 'use_attn_tp': use_attn_tp - } - - logger.info(f"Loading all_reduce data for {cluster_type}: file={all_reduce_input_file}, tp_size={tp_size}") - - all_reduce_df = self._load_all_reduce_df(all_reduce_input_file, replica_config, tp_size) - logger.info(f"Loaded {len(all_reduce_df)} rows for all_reduce training") - - all_reduce_df = self._get_all_reduce_df_with_derived_features(all_reduce_df, replica_config) - logger.info(f"After feature engineering: {len(all_reduce_df)} rows") - - models["all_reduce"] = self._train_single_model( - model_name="all_reduce", - df=all_reduce_df, - feature_cols=["num_tokens"], - target_col="time_stats.all_reduce.median", - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context - ) - - trained_model_signatures.add(tp_signature) - return models - - def _train_cpu_overhead_models_for_cluster(self, cluster_type: ClusterType, replica_config, execution_time_predictor_config, trained_model_signatures: set) -> Dict[str, BaseEstimator]: - """Train CPU overhead models for a cluster.""" - models = {} - - if execution_time_predictor_config.skip_cpu_overhead_modeling: - return models - - _, _, _, _, cpu_overhead_input_file, _ = self._get_input_files_for_config(replica_config, execution_time_predictor_config) - - cpu_signature = f"cpu_overhead_{replica_config.network_device}_{replica_config.model_name}_{replica_config.attn_tensor_parallel_size}_family{self._measurement_family_name(self._active_measurement_type)}" - - if cpu_signature in trained_model_signatures: - logger.info(f"Skipping CPU overhead models training for {cluster_type} - already trained") - return models - - cpu_overhead_df = self._load_cpu_overhead_df(cpu_overhead_input_file, replica_config) - if cpu_overhead_df.empty: - logger.warning( - "Skipping CPU overhead model training for cluster %s due to missing/empty CPU overhead profiling data. file=%s", - cluster_type, - cpu_overhead_input_file, - ) - trained_model_signatures.add(cpu_signature) - return models - - # Build training context for error messages - training_context = { - 'cluster_type': str(cluster_type), - 'device': replica_config.device, - 'model_name': replica_config.model_name, - 'tensor_parallel_size': replica_config.attn_tensor_parallel_size, - 'network_device': replica_config.network_device, - 'input_file': cpu_overhead_input_file, - } - - model_names = [ - "schedule", - "sampler_e2e", - "prepare_inputs_e2e", - "process_model_outputs", - "ray_comm_time", - ] - - for model_name in model_names: - target_col = "ray_comm_time_mean" if model_name == "ray_comm_time" else f"{model_name}_median" - - model_signature = f"{model_name}_{cpu_signature}" - if model_signature not in trained_model_signatures: - feature_cols = [ - "batch_size", - "num_prefill_tokens", - "num_decode_tokens", - ] - model = self._train_single_model( - model_name=model_name, - df=cpu_overhead_df, - feature_cols=feature_cols, - target_col=target_col, - execution_time_predictor_config=execution_time_predictor_config, - training_context=training_context, - persist_exact_lookup=True, - ) - if not hasattr(model, "_frontier_exact_lookup"): - model._frontier_exact_lookup = _build_exact_feature_lookup( - cpu_overhead_df, - feature_cols, - target_col, - ) - models[model_name] = model - trained_model_signatures.add(model_signature) - - trained_model_signatures.add(cpu_signature) - return models - - def _train_single_model( - self, - model_name: str, - df: pd.DataFrame, - feature_cols: List[str], - target_col: str, - execution_time_predictor_config, - training_context: Optional[Dict[str, Any]] = None, - persist_exact_lookup: bool = True, - layer_contract: Optional[ResolvedLayerContract] = None, - ) -> BaseEstimator: - """Train a single model with given data and configuration.""" - layer_contract, training_context = _normalize_layer_contract_context( - training_context, - explicit_layer_contract=layer_contract, - ) - if len(df) == 0: - # 提供详细的错误信息，以便调试 - context_info = "" - if training_context: - context_info = f""" - Training Context: - - Cluster Type: {training_context.get('cluster_type', 'Unknown')} - - Device: {training_context.get('device', 'Unknown')} - - Model Name: {training_context.get('model_name', 'Unknown')} - - Pipeline Stages: {training_context.get('pipeline_stages', 'Unknown')} - - Network Device: {training_context.get('network_device', 'Unknown')} - - Tensor Parallel Size: {training_context.get('tensor_parallel_size', 'Unknown')} - - Input File: {training_context.get('input_file', 'Unknown')} - - Block Size: {training_context.get('block_size', 'Unknown')} - - Feature Columns: {feature_cols} - - Target Column: {target_col} - """ - - raise Exception(f"Training data for model {model_name} is empty.{context_info}") - - required_cols = feature_cols + [target_col] - nan_row_mask = df[required_cols].isna().any(axis=1) - nan_row_count = int(nan_row_mask.sum()) - if nan_row_count > 0: - logger.warning( - "Dropping %d/%d rows with NaN feature/target values before training %s " - "(target=%s).", - nan_row_count, - len(df), - model_name, - target_col, - ) - df = df.loc[~nan_row_mask].copy() - if len(df) == 0: - raise ValueError( - f"Training data for model {model_name} is empty after dropping NaN rows " - f"(target={target_col})." - ) - - profiling_precision = self._get_profiling_precision_from_df(df) - measurement_type = self._validate_active_measurement_type(df) - hash_args = ( - model_name, - df, - execution_time_predictor_config, - profiling_precision, - measurement_type, - ) - model_hash = self._get_model_hash( - *hash_args, - **_layer_contract_kwargs(layer_contract), - ) - cached_model = self._load_model_from_cache(model_name, model_hash) - if cached_model is not None: - if layer_contract is not None: - requested_identity = _serialize_selected_layer_cache_identity( - layer_contract - ) - if requested_identity is None: - raise ValueError( - "resolved layer contract did not produce a cache identity" - ) - self._validate_cached_layer_cache_identity( - model_name=model_name, - model=cached_model, - requested_identity=requested_identity, - ) - if persist_exact_lookup: - self._ensure_exact_lookup_metadata( - model_name=model_name, - model_hash=model_hash, - model=cached_model, - df=df, - feature_cols=feature_cols, - target_col=target_col, - ) - self._store_model_precision( - model_name, - profiling_precision, - cached_model, - **_layer_contract_kwargs(layer_contract), - ) - return cached_model - - # ============================================================ - # CACHE MISS: Model not found in cache - # ============================================================ - # When running in production mode (non-dummy mode), we expect all models - # to be pre-trained using the standalone training module and cached. - # If a model is not found in cache, it indicates a configuration mismatch - # or missing profiling/training step. - # - # To train models, use the standalone training workflow: - # 1. Run profiling: tests/test_pd_af_profiling.sh - # 2. Run training: tests/test_pd_af_training.sh - # 3. Run simulation: tests/test_small_scale_pd_af_disaggregation_cluster_parallel.sh - # ============================================================ - - error_msg = f""" - ❌ MODEL CACHE MISS ERROR ❌ - - Model '{model_name}' with hash '{model_hash}' not found in cache directory: {self._cache_dir} - - Configuration Details: - - Model Name: {model_name} - - Cache Hash: {model_hash} - - Cache Directory: {self._cache_dir} - - Expected Cache File: {self._cache_dir}/{model_name}_{model_hash}.pkl - """ - - if training_context: - error_msg += f""" - Training Context: - - Cluster Type: {training_context.get('cluster_type', 'Unknown')} - - Device: {training_context.get('device', 'Unknown')} - - Model Name: {training_context.get('model_name', 'Unknown')} - - Tensor Parallel Size: {training_context.get('tensor_parallel_size', 'Unknown')} - - Expert Parallel Size: {training_context.get('moe_expert_parallel_size', 'N/A')} - - Input File: {training_context.get('input_file', 'Unknown')} - - Feature Columns: {feature_cols} - - Target Column: {target_col} - """ - - error_msg += f""" - - ⚠️ REQUIRED ACTION ⚠️ - - This error indicates that the required model has not been pre-trained. - Please follow the complete workflow: - - ============================================================ - - NOTE: Real-time training is TEMPORARILY ENABLED for cache generation. - """ - - logger.warning(error_msg) - logger.info(f"CACHE MISS: Training model '{model_name}' with hash '{model_hash}' in real-time...") - - # ============================================================ - # TEMPORARILY ENABLED: Real-time training code - # ============================================================ - # This code performs real-time model training during simulation - # initialization to generate missing cache files. - # ============================================================ - - estimator, grid_search_params = self._create_estimator_and_params(execution_time_predictor_config) - - cv = min(execution_time_predictor_config.k_fold_cv_splits, len(df)) if len(df) >= 2 else 2 - - grid_search = GridSearchCV( - estimator=estimator, - param_grid=grid_search_params, - scoring=self._get_scorer(), - cv=cv, - n_jobs=execution_time_predictor_config.num_training_job_threads, - ) - - X, y = df[feature_cols], df[target_col] - grid_search.fit(X, y) - score = grid_search.score(X, y) - - logger.info(f"✓ Trained model {model_name} with MAPE {-score}%") - - best_estimator = grid_search.best_estimator_ - # Persist feature metadata for runtime on-demand prediction (e.g., moe_grouped_gemm load imbalance mode). - setattr(best_estimator, "_frontier_feature_names", list(feature_cols)) - setattr(best_estimator, "_frontier_target_col", target_col) - # Tie the trained estimator to its cache hash so prediction caches can include model identity. - setattr(best_estimator, "_frontier_model_hash", model_hash) - if layer_contract is not None: - self._model_contract_identity(best_estimator, layer_contract) - - if persist_exact_lookup: - setattr( - best_estimator, - "_frontier_exact_lookup", - _build_exact_feature_lookup(df, feature_cols, target_col), - ) - - self._store_model_in_cache(model_name, model_hash, best_estimator) - self._store_model_precision( - model_name, - profiling_precision, - best_estimator, - **_layer_contract_kwargs(layer_contract), - ) - return best_estimator - - # ======================================================================== - # Data Loading Methods - # ======================================================================== - # These methods load profiling data from CSV files and apply filtering. - # - # Data source mapping: - # - linear_op.csv (or mlp.csv for backward compatibility): - # - Attention projections: attn_pre_proj, attn_post_proj, attn_rope - # - MLP operations: mlp_up_proj, mlp_down_proj, mlp_act - # - LayerNorm operations: input_layernorm, post_attention_layernorm - # - Residual operations: add - # - # - attention.csv: - # - Attention core: attn_kv_cache_save, attn_prefill, attn_decode - # - # - moe.csv: - # - MoE operations: moe_gating_linear, moe_gating_routing_topk, moe_shuffling, moe_grouped_gemm - # ======================================================================== - - def _load_linear_op_df( - self, - file_path: str, - tensor_parallel_size: int, - required_columns: Optional[List[str]] = None, - training_context: Optional[Dict[str, Any]] = None, - layer_contract: Optional[ResolvedLayerContract] = None, - operator_name: Optional[str] = None, - ) -> pd.DataFrame: - """ - Load linear operation dataframe (linear_op.csv or mlp.csv) with tensor parallel filtering. - - This function loads profiling data for linear operations including: - - Attention projections: attn_pre_proj, attn_post_proj, attn_rope - - MLP operations: mlp_up_proj, mlp_down_proj, mlp_act - - LayerNorm operations: input_layernorm, post_attention_layernorm - - Residual operations: add - - Note: This function is for linear_op.csv data only. For MoE data, use _load_moe_df(). - - Args: - file_path: Path to the profiling CSV file (linear_op.csv or mlp.csv) - tensor_parallel_size: Required tensor parallel size for filtering - - Returns: - Filtered DataFrame - - Raises: - FileNotFoundError: If the input file does not exist - ValueError: If required columns are missing or no data matches filtering criteria - """ - if layer_contract is not None: - _validate_typed_parallel_selection( - layer_contract, - tensor_parallel_size=tensor_parallel_size, - ) - - # Check file existence - if not os.path.exists(file_path): - raise FileNotFoundError( - f"Linear ops input file does not exist: {file_path}\n" - f"Please run profiling first to generate this file.\n" - f"Suggested command: bash frontier/profiling/example/test_profiling_linear_op.sh" - ) - - df = pd.read_csv(file_path) - logger.info(f"Original linear ops data: {len(df)} rows, {len(df.columns)} columns") - expected_profile = ( - layer_contract.profile_id - if layer_contract is not None - else infer_single_runtime_profile(self) - ) - if expected_profile is not None and "model_architecture_profile" in df.columns: - validate_model_architecture_profile( - df, - file_path=file_path, - expected_profile=expected_profile, - ) - - # Check required column - if 'num_tensor_parallel_workers' not in df.columns: - raise ValueError( - f"Column 'num_tensor_parallel_workers' not found in {file_path}\n" - f"Available columns: {list(df.columns)}\n" - f"This may indicate a corrupted or incompatible profiling file." - ) - - has_typed_contracts = TYPED_OPERATOR_CONTRACTS_COLUMN in df.columns - parsed_typed_contracts: Optional[pd.Series] = None - if has_typed_contracts: - if not operator_name and layer_contract is not None: - raise ValueError( - "typed profiling loading requires operator_name when the " - f"canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" - ) - if operator_name is not None and layer_contract is None: - raise ValueError( - "typed profiling loading requires layer_contract when the " - f"canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" - ) - # Parse every row before applying scalar filters so malformed metadata - # cannot be hidden by an unrelated TP or width selector. - parsed_typed_contracts = cast( - pd.Series, - df[TYPED_OPERATOR_CONTRACTS_COLUMN].map( - lambda raw_contracts: validate_typed_operator_contracts( - raw_contracts, - model_config=infer_single_runtime_model_config(self), - ) - ), - ) - - # Show filtering conditions - available_tp = sorted(df['num_tensor_parallel_workers'].unique()) - logger.info(f"Filtering conditions:") - logger.info(f" - num_tensor_parallel_workers == {tensor_parallel_size}") - logger.info(f" - Available num_tensor_parallel_workers: {available_tp}") - - # Apply filtering - filtered_df: pd.DataFrame = cast( - pd.DataFrame, - df[df["num_tensor_parallel_workers"] == tensor_parallel_size], - ) - if parsed_typed_contracts is not None and layer_contract is not None: - selected_layer_contract = layer_contract - if not isinstance(operator_name, str) or not operator_name: - raise ValueError( - "typed profiling loading requires a non-empty operator_name " - "for contract matching" - ) - typed_mask = parsed_typed_contracts.loc[filtered_df.index].map( - lambda raw_contracts: _typed_row_matches_contract( - raw_contracts, - selected_layer_contract, - operator_name=operator_name, - ) - ) - filtered_df = cast(pd.DataFrame, filtered_df[typed_mask]) - if filtered_df.empty: - raise ValueError( - "No linear-op rows match the selected typed layer contract " - f"for operator={operator_name!r}, TP={tensor_parallel_size} " - f"in {file_path}" - ) - elif layer_contract is not None: - if "n_expanded_embd" not in filtered_df.columns: - raise ValueError( - "Legacy linear-op profiling data is missing 'n_expanded_embd' " - f"for typed contract loading in {file_path}" - ) - filtered_df = cast( - pd.DataFrame, - filtered_df[ - filtered_df["n_expanded_embd"] - == layer_contract.effective_ffn_width - ], - ) - logger.info(f"After filtering: {len(filtered_df)} rows") - - expected_use_qk_norm = None - if training_context is not None and "use_qk_norm" in training_context: - expected_use_qk_norm = bool(training_context["use_qk_norm"]) - - if expected_use_qk_norm is True and "use_qk_norm" not in filtered_df.columns: - raise ValueError( - "linear_op profiling data is missing 'use_qk_norm' column for a model " - "that requires QK-norm-aware filtering. " - f"file={file_path}, model={training_context.get('model_name') if training_context else 'unknown'}" - ) - - if expected_use_qk_norm is not None and "use_qk_norm" in filtered_df.columns: - filtered_df = filtered_df[ - filtered_df["use_qk_norm"].astype(bool) == expected_use_qk_norm - ] - logger.info( - "After use_qk_norm filtering: %s rows (expected_use_qk_norm=%s)", - len(filtered_df), - expected_use_qk_norm, - ) - - if len(filtered_df) == 0: - width_requirement = ( - layer_contract.effective_ffn_width - if layer_contract is not None - else "legacy model width" - ) - raise ValueError( - f"No data matches the filtering criteria in {file_path}\n" - f"Required tensor_parallel_size: {tensor_parallel_size}\n" - f"Available tensor_parallel_sizes: {available_tp}\n" - f"Required effective_ffn_width: {width_requirement}\n" - f"Please run profiling with the correct configuration." - ) - - if required_columns: - self._validate_required_linear_op_columns( - filtered_df, - required_columns, - file_path, - training_context=training_context, - ) - - return filtered_df - - def _get_required_attn_linear_op_columns(self, model_config) -> List[str]: - required_columns = [ - "time_stats.attn_pre_proj.median", - "time_stats.attn_post_proj.median", - "time_stats.attn_rope.median", - ] - if model_config is not None and bool(getattr(model_config, "use_qk_norm", False)): - required_columns.append("use_qk_norm") - architecture_profile = _resolve_model_architecture_profile(model_config) - if architecture_profile is not None: - required_columns.extend( - f"time_stats.{op_name}.median" - for op_name in architecture_profile.predictor_attention_extra_ops - ) - return required_columns - - @staticmethod - def _get_required_target_embedded_mtp_linear_op_columns() -> List[str]: - return [ - "time_stats.mtp_fusion_proj.median", - "time_stats.lm_head_linear.median", - ] - - @staticmethod - def _validate_required_linear_op_columns( - df: pd.DataFrame, - required_columns: List[str], - file_path: str, - training_context: Optional[Dict[str, Any]] = None, - ) -> None: - missing_columns = [col for col in required_columns if col not in df.columns] - all_nan_columns = [ - col - for col in required_columns - if col in df.columns and df[col].isna().all() - ] - - if missing_columns or all_nan_columns: - context_text = "" - if training_context: - context_text = f"\nTraining context: {training_context}" - - raise ValueError( - "Required attention linear op columns are missing or all-NaN in " - f"{file_path}." - f"\nMissing columns: {missing_columns}" - f"\nAll-NaN columns: {all_nan_columns}" - f"{context_text}" - ) - - def _load_attention_df( - self, - file_path: str, - replica_config, - replica_scheduler_config, - cluster_type: Optional[ClusterType] = None, - ) -> pd.DataFrame: - """ - Load attention dataframe (attention.csv) with model configuration filtering. - - Args: - file_path: Path to the attention profiling CSV file - replica_config: Replica configuration for filtering - replica_scheduler_config: Replica scheduler configuration for block size - cluster_type: Cluster type for policy warning context - - Returns: - Filtered DataFrame - - Raises: - FileNotFoundError: If the input file does not exist - ValueError: If no data matches filtering criteria - """ - # Check file existence - if not os.path.exists(file_path): - raise FileNotFoundError( - f"Attention input file does not exist: {file_path}\n" - f"Please run attention profiling first to generate this file.\n" - f"Suggested command: bash frontier/profiling/example/test_profiling_attention.sh" - ) - - df = pd.read_csv(file_path) - df = df.drop_duplicates() - logger.info(f"Original attention data: {len(df)} rows, {len(df.columns)} columns") - - enforce_mixed_attention_input_contract( - attention_file_path=file_path, - available_columns=df.columns, - ) - - # Latent-MLA profiles use a distinct structural schema (runtime kv heads = 1, - # head size = kv_lora_rank + qk_rope_head_dim); route them to the MLA - # structural filter before the dense cache-write fill / dense filter. - model_config = replica_config.model_config - if self._is_mla_family(model_config): - return self._filter_mla_attention_df( - df, file_path, replica_config, replica_scheduler_config - ) - - # Fill missing cache-write column for older attention profiling CSVs. - cache_write_median_column = get_enabled_predictor_median_column_by_role( - DENSE_ATTENTION_FAMILY, - AttentionOperatorRole.CACHE_WRITE, - ) - for column in [cache_write_median_column]: - if column not in df.columns: - df[column] = 0 - else: - df.fillna({column: 0}, inplace=True) - - model_config = replica_config.model_config - requested_tp = replica_config.attn_tensor_parallel_size - prefill_op_name = get_enabled_predictor_metric_name_by_role( - DENSE_ATTENTION_FAMILY, - AttentionOperatorRole.PREFILL_KERNEL, - ) - effective_tp = resolve_effective_attention_tp_size( - op_name=prefill_op_name, - requested_tp_size=requested_tp, - num_kv_heads=model_config.num_kv_heads, - cluster_type=cluster_type, - warning_cache=getattr(self, "_attention_tp_warning_cache", None), - include_linear_ops=False, - ) - - # Show filtering conditions - logger.info(f"Filtering conditions:") - logger.info(f" - n_embd == {model_config.embedding_dim}") - logger.info(f" - n_q_head == {model_config.num_q_heads}") - logger.info(f" - n_kv_head == {model_config.num_kv_heads}") - logger.info(f" - block_size == {replica_scheduler_config.block_size}") - logger.info( - " - num_tensor_parallel_workers == %s (requested_tp=%s)", - effective_tp, - requested_tp, - ) - - filtered_df = df[ - (df["n_embd"] == model_config.embedding_dim) - & (df["n_q_head"] == model_config.num_q_heads) - & (df["n_kv_head"] == model_config.num_kv_heads) - & (df["block_size"] == replica_scheduler_config.block_size) - & (df["num_tensor_parallel_workers"] == effective_tp) - ] - - logger.info(f"After filtering: {len(filtered_df)} rows") - - if len(filtered_df) == 0: - # Surface what is available to make debugging explicit. - available = { - "n_embd": sorted(df["n_embd"].unique().tolist()) if "n_embd" in df else [], - "n_q_head": sorted(df["n_q_head"].unique().tolist()) if "n_q_head" in df else [], - "n_kv_head": sorted(df["n_kv_head"].unique().tolist()) if "n_kv_head" in df else [], - "block_size": sorted(df["block_size"].unique().tolist()) if "block_size" in df else [], - "num_tensor_parallel_workers": sorted(df["num_tensor_parallel_workers"].unique().tolist()) if "num_tensor_parallel_workers" in df else [], - } - - logger.error( - "Attention profiling rows are missing for the requested configuration. " - "Available values: %s", available - ) - - raise ValueError( - f"No data matches the filtering criteria in {file_path}\n" - f"Required configuration:\n" - f" - n_embd: {model_config.embedding_dim}\n" - f" - n_q_head: {model_config.num_q_heads}\n" - f" - n_kv_head: {model_config.num_kv_heads}\n" - f" - block_size: {replica_scheduler_config.block_size}\n" - f" - tensor_parallel_size(requested): {requested_tp}\n" - f" - tensor_parallel_size(effective): {effective_tp}\n" - f"Available values: {available}\n" - f"Please run attention profiling with the correct configuration." - ) - - return filtered_df - - def _load_all_reduce_df(self, file_path: str, replica_config, tensor_parallel_size: int) -> pd.DataFrame: - """ - Load all_reduce dataframe with cluster-specific tensor parallel size. - - Args: - file_path: Path to the communication profiling CSV file - replica_config: Replica configuration - tensor_parallel_size: Required tensor parallel size for filtering - - Returns: - Filtered DataFrame - - Raises: - FileNotFoundError: If the input file does not exist - ValueError: If no data matches filtering criteria - """ - if not os.path.exists(file_path): - raise FileNotFoundError( - f"All-reduce input file does not exist: {file_path}\n" - f"Please run communication profiling first.\n" - f"Suggested command: bash frontier/profiling/example/test_profiling_communication.sh" - ) - - df = pd.read_csv(file_path) - logger.info(f"Original all_reduce data: {len(df)} rows") - - # Show filtering conditions - logger.info(f"Filtering conditions:") - logger.info(f" - num_workers == {tensor_parallel_size}") - logger.info(f" - devices_per_node == {tensor_parallel_size}") - logger.info(f" - collective == 'all_reduce'") - - filtered_df = df[ - (df["num_workers"] == tensor_parallel_size) - & (df["devices_per_node"] == tensor_parallel_size) - & (df["collective"] == "all_reduce") - ] - - logger.info(f"After filtering: {len(filtered_df)} rows") - - if len(filtered_df) == 0: - available_info = "" - if len(df) > 0: - available_info = ( - f"Available values in file:\n" - f" - num_workers: {sorted(df['num_workers'].unique())}\n" - f" - devices_per_node: {sorted(df['devices_per_node'].unique())}\n" - f" - collective: {sorted(df['collective'].unique())}" - ) - raise ValueError( - f"No data matches the filtering criteria in {file_path}\n" - f"Required: num_workers={tensor_parallel_size}, devices_per_node={tensor_parallel_size}, collective='all_reduce'\n" - f"{available_info}" - ) - - return filtered_df - - def _load_send_recv_df(self, file_path: str, replica_config) -> pd.DataFrame: - """ - Load send_recv dataframe for pipeline parallel communication. - - Args: - file_path: Path to the communication profiling CSV file - replica_config: Replica configuration - - Returns: - Filtered DataFrame - - Raises: - FileNotFoundError: If the input file does not exist - """ - if not os.path.exists(file_path): - raise FileNotFoundError( - f"Send/recv input file does not exist: {file_path}\n" - f"Please run communication profiling first.\n" - f"Suggested command: bash frontier/profiling/example/test_profiling_communication.sh" - ) - - num_workers = replica_config.num_pipeline_stages * replica_config.attn_tensor_parallel_size - devices_per_node = replica_config.node_config.num_devices_per_node - is_multi_node = num_workers > devices_per_node - - if is_multi_node: - devices_per_node = 1 - else: - devices_per_node = 2 - - df = pd.read_csv(file_path) - logger.info(f"Original send_recv data: {len(df)} rows") - logger.info(f"Filtering conditions: collective='send_recv', devices_per_node={devices_per_node}") - - filtered_df = df[ - (df["collective"] == "send_recv") - & (df["devices_per_node"] == devices_per_node) - ] - - logger.info(f"After filtering: {len(filtered_df)} rows") - return filtered_df - - def _load_cpu_overhead_df(self, file_path: str, replica_config) -> pd.DataFrame: - """ - Load CPU overhead dataframe with model configuration filtering. - - Args: - file_path: Path to the CPU overhead profiling CSV file - replica_config: Replica configuration - - Returns: - Filtered DataFrame - - Raises: - FileNotFoundError: If the input file does not exist - """ - if not os.path.exists(file_path): - logger.warning( - "CPU overhead input file does not exist: %s. " - "Skipping CPU overhead model training for this cluster.", - file_path, - ) - return pd.DataFrame() - - df = pd.read_csv(file_path) - if df.empty: - logger.warning( - "CPU overhead input file is empty: %s. " - "Skipping CPU overhead model training for this cluster.", - file_path, - ) - return pd.DataFrame() - - df = apply_cpu_overhead_schema_v2_defaults( - df, - warn_fn=logger.warning, - context=file_path, - ) - df = validate_cpu_overhead_dataframe(df) - - model_config = replica_config.model_config - - logger.info(f"Original CPU overhead data: {len(df)} rows") - logger.info(f"Filtering conditions: model_name='{model_config.get_name()}', tensor_parallel_degree={replica_config.attn_tensor_parallel_size}") - - filtered_df = df[ - (df["model_name"] == model_config.get_name()) - & (df["tensor_parallel_degree"] == replica_config.attn_tensor_parallel_size) - ] - - logger.info(f"After filtering: {len(filtered_df)} rows") - if filtered_df.empty: - logger.warning( - "No CPU overhead profiling rows found for model_name='%s', " - "tensor_parallel_degree=%s in file '%s'.", - model_config.get_name(), - replica_config.attn_tensor_parallel_size, - file_path, - ) - return filtered_df - - # Load imbalance feature columns used for MoE training - # These features describe the load distribution across experts - # Reference: frontier/training/moe_trainer.py lines 224-239 (authoritative source) - # Reference: frontier/profiling/moe/LOAD_IMBALANCE_GUIDE.md - MOE_LOAD_IMBALANCE_FEATURES = [ - # Config features (6) - describe model configuration - "total_routed_tokens", # Total tokens after routing (num_tokens * router_topk) - "num_experts_per_device", # Number of experts per device after EP sharding - "hidden_dim", # Model hidden dimension - "expert_hidden_dim", # Expert FFN hidden dimension - "router_topk", # Number of experts each token is routed to - "model_expansion_ratio", # expert_hidden_dim / hidden_dim - # Derived features (2) - derived from config and routing - "tokens_per_expert_avg", # Average tokens per expert - "tokens_to_experts_ratio", # tokens / num_experts ratio - # Load features (6) - describe load distribution characteristics - "expert_utilization", # Proportion of experts with non-zero load - "min_load_ratio", # Min load / average load - "load_imbalance_cv", # Coefficient of Variation: std/mean, key imbalance metric - "max_load_ratio", # Max load / average load - "load_entropy", # Entropy of load distribution (higher = more uniform) - "load_gini_coefficient", # Gini coefficient: 0=equality, 1=inequality - ] - - # Feature columns for mixed-batch attention prefill model - # These features capture batch heterogeneity characteristics together with - # the uniform KV-cache context used by MixedAttentionInput profiling. - # Reference: frontier/training/attention_trainer.py lines 362-375 (authoritative source) - ATTN_PREFILL_MIXED_FEATURES = [ - # Core features (7) - "batch_size", # Number of sequences in batch - "kv_cache_size", # Uniform KV cache context for the mixed batch - "total_tokens", # Total tokens across all sequences - "avg_seq_len", # Average sequence length - "min_seq_len", # Minimum sequence length - "max_seq_len", # Maximum sequence length - "total_tokens_squared", # Computational complexity proxy - # Heterogeneity features (3) - "seq_len_variance", # Variance of sequence lengths - "seq_len_cv", # Coefficient of variation (std/mean) - "seq_len_range", # max_seq_len - min_seq_len - # Interaction features (2) - "batch_variance_interaction", # batch_size * seq_len_variance - "batch_cv_interaction", # batch_size * seq_len_cv - ] - - ATTN_DECODE_IN_MIXED_FEATURES = [ - "decode_batch_size", - "decode_avg_kv_cache_size", - "num_prefill_seqs", - "total_prefill_tokens", - "total_batch_size", - "batch_composition_ratio", - "total_tokens", - ] - - def _load_moe_df( - self, - file_path: str, - replica_config, - load_imbalance: bool = True, - tensor_parallel_size: Optional[int] = None, - expert_parallel_size: Optional[int] = None, - layer_contract: Optional[ResolvedLayerContract] = None, - operator_name: Optional[str] = None, - ) -> pd.DataFrame: - """ - Load MoE dataframe with cluster-specific configuration filtering. - - This function loads and filters MoE profiling data based on the model configuration - and parallelism settings. It supports two training modes controlled by `load_imbalance`: - - 1. **Load Imbalance Mode (default, load_imbalance=True)**: - - Uses profiling data that includes load imbalance features - - Training will use features like `load_imbalance_cv`, `load_gini_coefficient`, etc. - - Recommended for accurate MoE execution time prediction under real-world scenarios - - Requires profiling with `--enable_load_imbalance` flag - - 2. **Standard Mode (load_imbalance=False)**: - - Uses basic profiling data without load imbalance features - - Training only uses `num_tokens` as feature - - Simpler but less accurate for imbalanced workloads - - Compatible with legacy profiling data - - The difference is in the **training features used**, not data row filtering. - Load imbalance mode uses additional features to capture expert load distribution. - - Reference: frontier/profiling/moe/LOAD_IMBALANCE_GUIDE.md - - Args: - file_path: Path to the MoE profiling CSV file - replica_config: Replica configuration containing model and parallelism settings - load_imbalance: Training mode flag: - - True (default): Load imbalance mode - use load imbalance features - - False: Standard mode - only use basic num_tokens feature - tensor_parallel_size: Optional TP override for op-specific MoE training. - If None, uses replica_config.moe_tensor_parallel_size. - expert_parallel_size: Optional EP filter for op-specific MoE training. - If None, EP filtering is skipped (used for EP-agnostic replicated ops). - - Returns: - Filtered DataFrame ready for MoE model training - - Raises: - FileNotFoundError: If the input file does not exist - ValueError: If no data matches filtering criteria or required features are missing - """ - if layer_contract is not None: - _validate_typed_parallel_selection( - layer_contract, - tensor_parallel_size=tensor_parallel_size, - expert_parallel_size=expert_parallel_size, - ) - - if not os.path.exists(file_path): - raise FileNotFoundError( - f"MoE input file does not exist: {file_path}\n" - f"Please run MoE profiling first.\n" - f"Suggested command: bash frontier/profiling/example/test_profiling_moe.sh" - ) - - df = pd.read_csv(file_path) - logger.info(f"Original MoE data: {len(df)} rows, {len(df.columns)} columns") - expected_profile = ( - layer_contract.profile_id - if layer_contract is not None - else infer_single_runtime_profile(self) - ) - if expected_profile is not None and "model_architecture_profile" in df.columns: - validate_model_architecture_profile( - df, - file_path=file_path, - expected_profile=expected_profile, - ) - - has_typed_contracts = TYPED_OPERATOR_CONTRACTS_COLUMN in df.columns - parsed_typed_contracts: Optional[pd.Series] = None - if has_typed_contracts: - if not operator_name: - raise ValueError( - "typed profiling loading requires operator_name when the " - f"canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" - ) - if layer_contract is None: - raise ValueError( - "typed profiling loading requires layer_contract when the " - f"canonical {TYPED_OPERATOR_CONTRACTS_COLUMN!r} column is present" - ) - # Parse every row before applying scalar filters so malformed metadata - # cannot be hidden by an unrelated TP, EP, or width selector. - parsed_typed_contracts = cast( - pd.Series, - df[TYPED_OPERATOR_CONTRACTS_COLUMN].map( - lambda raw_contracts: validate_typed_operator_contracts( - raw_contracts, - model_config=replica_config.model_config, - ) - ), - ) - - model_config = replica_config.model_config - training_mode = "load_imbalance (load_imbalance=True)" if load_imbalance else "standard (load_imbalance=False)" - if tensor_parallel_size is None: - tensor_parallel_size = replica_config.moe_tensor_parallel_size - if tensor_parallel_size <= 0: - raise ValueError( - f"Invalid tensor_parallel_size for MoE data loading: {tensor_parallel_size}" - ) - - # Display filtering conditions - logger.info(f"Filtering conditions:") - logger.info(f" - num_experts == {model_config.num_experts}") - logger.info(f" - router_topk == {model_config.num_experts_per_tok}") - logger.info(f" - hidden_dim == {model_config.embedding_dim}") - expected_expert_width = ( - layer_contract.effective_ffn_width - if layer_contract is not None - else model_config.mlp_hidden_dim - ) - logger.info(f" - expert_hidden_dim == {expected_expert_width}") - logger.info(f" - num_tensor_parallel_workers == {tensor_parallel_size}") - if expert_parallel_size is None: - logger.info(" - expert_parallel_size == ANY (EP-agnostic op)") - else: - logger.info(f" - expert_parallel_size == {expert_parallel_size}") - logger.info(f" - training_mode: {training_mode}") - - # Display available values in the dataset - available_info = [] - if len(df) > 0: - if 'num_experts' in df.columns: - available_info.append(f" - Available num_experts: {sorted(df['num_experts'].unique())}") - if 'router_topk' in df.columns: - available_info.append(f" - Available router_topk: {sorted(df['router_topk'].unique())}") - if 'num_tensor_parallel_workers' in df.columns: - available_info.append(f" - Available num_tensor_parallel_workers: {sorted(df['num_tensor_parallel_workers'].unique())}") - if 'expert_parallel_size' in df.columns: - available_info.append(f" - Available expert_parallel_size: {sorted(df['expert_parallel_size'].unique())}") - if 'load_distribution' in df.columns: - available_info.append(f" - Available load_distribution: {sorted(df['load_distribution'].unique())}") - - for info in available_info: - logger.info(info) - - # Apply filtering based on MoE configuration - filtered_df = cast(pd.DataFrame, df[ - (df["num_experts"] == model_config.num_experts) - & (df["router_topk"] == model_config.num_experts_per_tok) - & (df["hidden_dim"] == model_config.embedding_dim) - & (df["num_tensor_parallel_workers"] == tensor_parallel_size) - ]) - if not has_typed_contracts: - filtered_df = filtered_df[ - filtered_df["expert_hidden_dim"] == expected_expert_width - ] - else: - if parsed_typed_contracts is None: - raise RuntimeError( - "typed MoE metadata column was detected but could not be parsed" - ) - if layer_contract is None: - raise ValueError( - "typed MoE filtering requires a resolved layer contract" - ) - selected_layer_contract = layer_contract - typed_mask = parsed_typed_contracts.loc[filtered_df.index].map( - lambda raw_contracts: _typed_row_matches_contract( - raw_contracts, - selected_layer_contract, - operator_name=operator_name, - ) - ) - filtered_df = cast(pd.DataFrame, filtered_df[typed_mask]) - filtered_df = cast(pd.DataFrame, filtered_df) - if expert_parallel_size is not None: - if "expert_parallel_size" not in filtered_df.columns: - raise ValueError( - "MoE profiling data is missing 'expert_parallel_size' while " - f"EP={expert_parallel_size} is required in {file_path}" - ) - filtered_df = filtered_df[ - filtered_df["expert_parallel_size"] == expert_parallel_size - ] - - logger.info(f"After config filtering: {len(filtered_df)} rows") - - # Check for load imbalance features if load_imbalance mode is enabled - if load_imbalance: - missing_features = [ - f for f in self.MOE_LOAD_IMBALANCE_FEATURES - if f not in filtered_df.columns - ] - if missing_features: - logger.warning( - f"Load imbalance mode requested but missing features: {missing_features}\n" - f"Available columns: {list(filtered_df.columns)}\n" - f"Please run MoE profiling with --enable_load_imbalance flag.\n" - f"Use load_imbalance=False (standard mode) explicitly if you want to train without load imbalance features." - ) - raise ValueError("Missing load imbalance features") - # Note: We don't change load_imbalance here, caller should handle feature selection - else: - logger.info(f"Load imbalance features available: {self.MOE_LOAD_IMBALANCE_FEATURES}") - - if len(filtered_df) == 0: - ep_requirement = "ANY" if expert_parallel_size is None else expert_parallel_size - available_info_text = "\n".join(available_info) - message = ( - f"No data matches the filtering criteria in {file_path}\n" - f"Required MoE configuration:\n" - f" - num_experts: {model_config.num_experts}\n" - f" - router_topk: {model_config.num_experts_per_tok}\n" - f" - hidden_dim: {model_config.embedding_dim}\n" - f" - expert_hidden_dim: {expected_expert_width}\n" - f" - tensor_parallel_size: {tensor_parallel_size}\n" - f" - expert_parallel_size: {ep_requirement}\n" - f" - training_mode: {training_mode}\n" - ) - if has_typed_contracts: - message += ( - f" - typed operator: {operator_name!r}\n" - " - typed layer contract admission: required\n" - ) - if available_info_text: - message += available_info_text - raise ValueError( - message - ) - - return filtered_df - - def _get_attention_df_with_derived_features(self, df: pd.DataFrame) -> pd.DataFrame: - """Add derived features to attention dataframe. - - Standard features for attn_prefill and attn_decode: - - num_tokens: max(prefill_chunk_size, batch_size) - - is_decode: derived from is_prefill when available, else prefill_chunk_size == 0 - - prefill_chunk_size_squared: prefill_chunk_size ** 2 - - Mixed-batch features for attn_prefill_mixed (12 features): - Reference: frontier/training/attention_trainer.py lines 362-375 - These features capture batch heterogeneity for accurate prefill time prediction. - """ - df_with_derived_features = df.copy() - - # Standard attention features - df_with_derived_features["num_tokens"] = df_with_derived_features[["prefill_chunk_size", "batch_size"]].max(axis=1) - if "is_prefill" in df_with_derived_features.columns: - normalized_prefill_values = coerce_truthy_bool( - df_with_derived_features["is_prefill"] - ) - df_with_derived_features["is_decode"] = ~normalized_prefill_values - else: - df_with_derived_features["is_decode"] = (df_with_derived_features["prefill_chunk_size"] == 0) - df_with_derived_features["prefill_chunk_size_squared"] = (df_with_derived_features["prefill_chunk_size"] ** 2) - - def _normalize_bool_series(series: pd.Series) -> pd.Series: - return coerce_truthy_bool(series) - - if "is_mixed_batch" in df_with_derived_features.columns: - df_with_derived_features["is_mixed_batch"] = _normalize_bool_series( - df_with_derived_features["is_mixed_batch"] - ) - else: - df_with_derived_features["is_mixed_batch"] = False - - if "is_true_mixed_batch" in df_with_derived_features.columns: - df_with_derived_features["is_true_mixed_batch"] = _normalize_bool_series( - df_with_derived_features["is_true_mixed_batch"] - ) - else: - df_with_derived_features["is_true_mixed_batch"] = False - - # Mixed-batch features for attn_prefill_mixed (if applicable) - # Check if the profiling data contains mixed-batch specific columns - has_mixed_batch_data = "total_tokens" in df_with_derived_features.columns - - if has_mixed_batch_data: - logger.info("Adding mixed-batch derived features for attn_prefill_mixed") - - # total_tokens_squared for computational complexity - if "total_tokens" in df_with_derived_features.columns: - df_with_derived_features["total_tokens_squared"] = ( - df_with_derived_features["total_tokens"] ** 2 - ) - - # seq_len_range = max_seq_len - min_seq_len - if "max_seq_len" in df_with_derived_features.columns and "min_seq_len" in df_with_derived_features.columns: - df_with_derived_features["seq_len_range"] = ( - df_with_derived_features["max_seq_len"] - - df_with_derived_features["min_seq_len"] - ) - - # Interaction features: batch_size * heterogeneity metrics - if "seq_len_variance" in df_with_derived_features.columns: - df_with_derived_features["batch_variance_interaction"] = ( - df_with_derived_features["batch_size"] * - df_with_derived_features["seq_len_variance"] - ) - - if "seq_len_cv" in df_with_derived_features.columns: - df_with_derived_features["batch_cv_interaction"] = ( - df_with_derived_features["batch_size"] * - df_with_derived_features["seq_len_cv"] - ) - - if { - "num_prefill_seqs", - "num_decode_seqs", - }.issubset(df_with_derived_features.columns) and ( - "total_batch_size" not in df_with_derived_features.columns - ): - df_with_derived_features["total_batch_size"] = ( - df_with_derived_features["num_prefill_seqs"] - + df_with_derived_features["num_decode_seqs"] - ) - - if { - "num_prefill_seqs", - "total_batch_size", - }.issubset(df_with_derived_features.columns) and ( - "batch_composition_ratio" not in df_with_derived_features.columns - ): - total_batch_size = df_with_derived_features["total_batch_size"].replace(0, pd.NA) - df_with_derived_features["batch_composition_ratio"] = ( - df_with_derived_features["num_prefill_seqs"] / total_batch_size - ).fillna(0.0) - - if ( - "num_decode_seqs" in df_with_derived_features.columns - and "decode_batch_size" not in df_with_derived_features.columns - ): - df_with_derived_features["decode_batch_size"] = df_with_derived_features[ - "num_decode_seqs" - ] - - return df_with_derived_features - - def _get_all_reduce_df_with_derived_features(self, df: pd.DataFrame, replica_config) -> pd.DataFrame: - df_with_derived_features = df.copy() - df_with_derived_features["num_tokens"] = ( - df_with_derived_features["size"] / replica_config.model_config.embedding_dim / 2 - ) - return df_with_derived_features - - def _get_send_recv_df_with_derived_features(self, df: pd.DataFrame, replica_config) -> pd.DataFrame: - df_with_derived_features = df.copy() - df_with_derived_features["num_tokens"] = ( - df_with_derived_features["size"] / replica_config.model_config.embedding_dim / 2 - ) - return df_with_derived_features - - def _get_moe_df_with_derived_features(self, df: pd.DataFrame) -> pd.DataFrame: - """ - Add derived features to MoE dataframe. - - The MoE profiling data already contains num_tokens as a direct column, - so we just ensure it exists and return the dataframe. - Additional derived features can be added here if needed in the future. - """ - df_with_derived_features = df.copy() - - # Verify that num_tokens column exists (it should be in the profiling output) - if "num_tokens" not in df_with_derived_features.columns: - logger.warning("num_tokens column not found in MoE dataframe") - logger.warning(f"Available columns: {list(df_with_derived_features.columns)}") - # If num_tokens is missing, we cannot proceed with training - raise ValueError("MoE profiling data must contain 'num_tokens' column") - - return df_with_derived_features - - def _get_hash_relevant_config(self, config) -> Dict[str, Any]: - """ - Extract only the configuration parameters that affect model performance. - - Parameters that should be included: - - Profiling data paths (determine input data source) - - Prediction range parameters (determine prediction cache scope) - - Performance adjustment parameters (affect predicted values) - - ML hyperparameters (affect model structure) - - Parameters that should be excluded: - - Training process parameters (k_fold_cv_splits, num_training_job_threads) - - Runtime configuration (no_cache, skip_cpu_overhead_modeling, enable_dummy_mode, dummy_execution_time_ms) - """ - hash_relevant_params = { - # Category 1: Profiling data paths - 'linear_op_input_file': config.linear_op_input_file, - 'atten_input_file': config.atten_input_file, - 'all_reduce_input_file': config.all_reduce_input_file, - 'send_recv_input_file': config.send_recv_input_file, - 'moe_input_file': config.moe_input_file, - 'linear_op_kernel_only_input_file': config.linear_op_kernel_only_input_file, - 'atten_kernel_only_input_file': config.atten_kernel_only_input_file, - 'moe_kernel_only_input_file': config.moe_kernel_only_input_file, - 'cpu_overhead_input_file': config.cpu_overhead_input_file, - 'cpu_overhead_kernel_only_input_file': getattr( - config, - 'cpu_overhead_kernel_only_input_file', - config.cpu_overhead_input_file, - ), - - # Category 2: Prediction range parameters - 'kv_cache_prediction_granularity': config.kv_cache_prediction_granularity, - 'prediction_max_prefill_chunk_size': config.prediction_max_prefill_chunk_size, - 'prediction_max_batch_size': config.prediction_max_batch_size, - 'prediction_max_tokens_per_request': config.prediction_max_tokens_per_request, - - # Category 3: Performance adjustment parameters - 'attention_decode_batching_overhead_fraction': config.attention_decode_batching_overhead_fraction, - 'attention_prefill_batching_overhead_fraction': config.attention_prefill_batching_overhead_fraction, - 'attn_pre_proj_calibration_scale': config.attn_pre_proj_calibration_scale, - 'prefill_phase_attn_pre_proj_calibration_scale': config.prefill_phase_attn_pre_proj_calibration_scale, - 'attn_post_proj_calibration_scale': config.attn_post_proj_calibration_scale, - 'prefill_phase_attn_post_proj_calibration_scale': config.prefill_phase_attn_post_proj_calibration_scale, - 'attn_decode_calibration_scale': config.attn_decode_calibration_scale, - 'attn_decode_in_mixed_calibration_scale': config.attn_decode_in_mixed_calibration_scale, - 'late_decode_attn_decode_calibration_scale': config.late_decode_attn_decode_calibration_scale, - 'attn_kv_cache_save_calibration_scale': config.attn_kv_cache_save_calibration_scale, - 'prefill_phase_attn_kv_cache_save_calibration_scale': config.prefill_phase_attn_kv_cache_save_calibration_scale, - 'mlp_up_proj_calibration_scale': config.mlp_up_proj_calibration_scale, - 'prefill_phase_mlp_up_proj_calibration_scale': config.prefill_phase_mlp_up_proj_calibration_scale, - 'mlp_down_proj_calibration_scale': config.mlp_down_proj_calibration_scale, - 'decode_phase_mlp_down_proj_calibration_scale': config.decode_phase_mlp_down_proj_calibration_scale, - 'nccl_cpu_launch_overhead_ms': config.nccl_cpu_launch_overhead_ms, - 'nccl_cpu_skew_overhead_per_device_ms': config.nccl_cpu_skew_overhead_per_device_ms, - } - - # Category 4: ML Hyperparameters (type-specific) - if hasattr(config, 'num_estimators'): # Random Forest - hash_relevant_params['num_estimators'] = config.num_estimators - hash_relevant_params['max_depth'] = config.max_depth - hash_relevant_params['min_samples_split'] = config.min_samples_split - elif hasattr(config, 'polynomial_degree'): # Linear Regression - hash_relevant_params['polynomial_degree'] = config.polynomial_degree - hash_relevant_params['polynomial_include_bias'] = config.polynomial_include_bias - hash_relevant_params['polynomial_interaction_only'] = config.polynomial_interaction_only - hash_relevant_params['fit_intercept'] = config.fit_intercept - - return hash_relevant_params - - def _get_model_hash( - self, - model_name: str, - df: pd.DataFrame, - execution_time_predictor_config, - profiling_precision: str, - measurement_type: MeasurementType, - layer_contract: Optional[ResolvedLayerContract] = None, - ) -> str: - """ - Calculate hash for model caching based on configuration and data. - - Hash is calculated from: - 1. Hash-relevant configuration parameters (excluding runtime/training process params) - 2. Model name - 3. DataFrame content hash - - This ensures that only changes to parameters that affect model performance - will invalidate the cache. - """ - # Extract only hash-relevant parameters - hash_relevant_config = self._get_hash_relevant_config(execution_time_predictor_config) - config_str = str(sorted(hash_relevant_config.items())) # Sort for deterministic ordering - - # Calculate DataFrame hash - df_hash_str = hashlib.md5(df.to_json().encode("utf-8")).hexdigest() - - selected_identity = _serialize_selected_layer_cache_identity(layer_contract) - contract_component = ( - f"_{selected_identity}" if selected_identity is not None else "" - ) - - # Combine all components. The selected semantic domain is part of a - # typed key; physical layer occurrence is intentionally absent. - combined_str = ( - f"{config_str}_{model_name}_{df_hash_str}_{profiling_precision}_" - f"{measurement_type.value}{contract_component}" - ) - hash_value = hashlib.md5(combined_str.encode("utf-8")).hexdigest()[0:8] - - # Debug output for hash calculation - if model_name == "attn_pre_proj": - logger.info(f"[DEBUG] Hash calculation for {model_name}:") - logger.info(f" - DataFrame shape: {df.shape}") - logger.info(f" - DataFrame hash: {df_hash_str[:16]}...") - logger.info(f" - Hash-relevant config keys: {sorted(hash_relevant_config.keys())}") - logger.info(f" - Final hash: {hash_value}") - - return hash_value - - def _get_profiling_precision_from_df(self, df: pd.DataFrame) -> str: - """Extract profiling precision from DataFrame. - - FAIL-FAST: Raises ValueError if profiling_precision column is missing or invalid. - This enforces strict metadata requirements and prevents silent fallbacks. - """ - if "profiling_precision" not in df.columns: - raise ValueError( - "profiling_precision column is missing from profiling data. " - f"Run '{MIGRATION_HELP_COMMAND}' to add required metadata columns to legacy CSV files." - ) - - precision_values = df["profiling_precision"].dropna().unique().tolist() - if not precision_values: - raise ValueError("profiling_precision column is empty") - if len(precision_values) > 1: - raise ValueError( - f"Multiple profiling_precision values found: {precision_values}" - ) - return str(precision_values[0]).upper() - - def _get_measurement_type_from_df(self, df: pd.DataFrame) -> MeasurementType: - if "measurement_type" not in df.columns: - raise ValueError( - "measurement_type column is missing from profiling data. " - f"Run '{MIGRATION_HELP_COMMAND}' to add required metadata columns to legacy CSV files." - ) - - measurement_values = df["measurement_type"].dropna().unique().tolist() - if not measurement_values: - raise ValueError("measurement_type column is empty") - if len(measurement_values) > 1: - raise ValueError( - f"Multiple measurement_type values found: {measurement_values}" - ) - return MeasurementType.from_string(str(measurement_values[0])) - - def _validate_active_measurement_type(self, df: pd.DataFrame) -> MeasurementType: - measurement_type = self._get_measurement_type_from_df(df) - if measurement_type != self._active_measurement_type: - raise ValueError( - f"measurement_type mismatch: expected {self._active_measurement_type.value} " - f"but found {measurement_type.value}." - ) - return measurement_type - - @staticmethod - def _validate_cached_layer_cache_identity( - *, - model_name: str, - model: BaseEstimator, - requested_identity: str, - ) -> None: - """Reject a typed cache entry whose selected domain does not match.""" - - cached_identity = getattr(model, "_frontier_layer_cache_identity", None) - if cached_identity is None: - raise ValueError( - f"Cached model {model_name!r} is missing selected layer cache identity" - ) - if not isinstance(cached_identity, str): - raise ValueError( - f"Cached model {model_name!r} has an invalid selected layer cache " - f"identity of type {type(cached_identity).__name__}" - ) - if cached_identity != requested_identity: - raise ValueError( - f"Cached model {model_name!r} selected layer cache identity mismatch: " - f"cached={cached_identity!r}, requested={requested_identity!r}" - ) - - @staticmethod - def _model_contract_identity( - model: BaseEstimator, - layer_contract: Optional[ResolvedLayerContract], - ) -> Optional[str]: - """Attach and return the selected semantic identity for a model.""" - - requested_identity = _serialize_selected_layer_cache_identity(layer_contract) - attached_identity = getattr(model, "_frontier_layer_cache_identity", None) - if attached_identity is not None and not isinstance(attached_identity, str): - raise TypeError( - "_frontier_layer_cache_identity must be a string when present" - ) - if ( - requested_identity is not None - and attached_identity is not None - and requested_identity != attached_identity - ): - raise ValueError( - "model selected layer cache identity conflicts with the requested " - "contract" - ) - identity = requested_identity or attached_identity - if identity is not None: - setattr(model, "_frontier_layer_cache_identity", identity) - return identity - - def _contract_model_registry( - self, family_name: str - ) -> Dict[Tuple[str, Optional[str]], BaseEstimator]: - registry_attr = { - "eager": "_trained_models_eager_by_contract", - "device_event": "_trained_models_device_event_by_contract", - "kernel_only": "_trained_models_kernel_only_by_contract", - }.get(family_name) - if registry_attr is None: - raise ValueError(f"Unsupported family_name={family_name!r}") - return getattr(self, registry_attr) - - def _contract_precision_registry( - self, family_name: str - ) -> Dict[str, Dict[Tuple[str, Optional[str]], BaseEstimator]]: - registry_attr = { - "eager": "_models_by_precision_eager_by_contract", - "device_event": "_models_by_precision_device_event_by_contract", - "kernel_only": "_models_by_precision_kernel_only_by_contract", - }.get(family_name) - if registry_attr is None: - raise ValueError(f"Unsupported family_name={family_name!r}") - return getattr(self, registry_attr) - - def _legacy_model_registry(self, family_name: str) -> Dict[str, BaseEstimator]: - registry_attr = { - "eager": "_trained_models_eager", - "device_event": "_trained_models_device_event", - "kernel_only": "_trained_models_kernel_only", - }.get(family_name) - if registry_attr is None: - raise ValueError(f"Unsupported family_name={family_name!r}") - return getattr(self, registry_attr) - - def _legacy_precision_registry( - self, family_name: str - ) -> Dict[str, Dict[str, BaseEstimator]]: - registry_attr = { - "eager": "_models_by_precision_eager", - "device_event": "_models_by_precision_device_event", - "kernel_only": "_models_by_precision_kernel_only", - }.get(family_name) - if registry_attr is None: - raise ValueError(f"Unsupported family_name={family_name!r}") - return getattr(self, registry_attr) - - def _legacy_precision_bucket( - self, family_name: str, precision_key: str - ) -> Dict[str, BaseEstimator]: - if not isinstance(precision_key, str) or not precision_key: - raise ValueError( - f"precision_key must be a non-empty string, got {precision_key!r}" - ) - registry = self._legacy_precision_registry(family_name) - canonical_key = precision_key.upper() - bucket = registry.get(canonical_key) - if bucket is not None: - return bucket - for stored_key, stored_bucket in registry.items(): - if str(stored_key).upper() == canonical_key: - return stored_bucket - return {} - - def _store_model_precision( - self, - model_name: str, - precision: str, - model: BaseEstimator, - layer_contract: Optional[ResolvedLayerContract] = None, - ) -> None: - if not isinstance(precision, str) or not precision.strip(): - raise ValueError(f"precision must be a non-empty string, got {precision!r}") - precision_key = precision.upper() - family_name = self._measurement_family_name(self._active_measurement_type) - identity = self._model_contract_identity(model, layer_contract) - if identity is None: - self._legacy_model_registry(family_name)[model_name] = model - self._legacy_precision_registry(family_name).setdefault( - precision_key, {} - )[model_name] = model - return - - model_key = (model_name, identity) - self._contract_model_registry(family_name)[model_key] = model - self._contract_precision_registry(family_name).setdefault( - precision_key, {} - )[model_key] = model - - def _get_family_model( - self, - family_name: str, - model_name: str, - *, - precision_key: Optional[str] = None, - requested_identity: Optional[str] = None, - ) -> Optional[BaseEstimator]: - if precision_key is not None: - precision_key = precision_key.upper() - source = self._contract_precision_registry(family_name).get( - precision_key, {} - ) - else: - source = self._contract_model_registry(family_name) - - typed_candidates = { - identity: candidate - for (candidate_name, identity), candidate in source.items() - if candidate_name == model_name - } - legacy = ( - self._legacy_precision_bucket(family_name, precision_key).get(model_name) - if precision_key is not None - else self._legacy_model_registry(family_name).get(model_name) - ) - legacy_identity = ( - getattr(legacy, "_frontier_layer_cache_identity", None) - if legacy is not None - else None - ) - - if requested_identity is not None: - model = typed_candidates.get(requested_identity) - if model is not None: - return model - if legacy is not None and legacy_identity == requested_identity: - return legacy - return None - - if len(typed_candidates) > 1: - identities = sorted( - "" if identity is None else identity - for identity in typed_candidates - ) - raise ValueError( - f"Model '{model_name}' has multiple layer contracts; provide " - f"layer_contract. Available identities: {identities}" - ) - if len(typed_candidates) == 1: - typed_identity = next(iter(typed_candidates)) - if legacy is not None and legacy_identity != typed_identity: - raise ValueError( - f"Model '{model_name}' has multiple layer contracts; provide " - f"layer_contract. Available identities: " - f"[{typed_identity!r}, {legacy_identity or ''!r}]" - ) - return next(iter(typed_candidates.values())) - return legacy - - def _resolve_cluster_model_contract( - self, - cluster_type: Optional[ClusterType], - model_name: str, - ) -> Optional[ResolvedLayerContract]: - """Resolve the typed domain requested by one cluster view.""" - - if cluster_type is None: - return None - cluster_configs = getattr(self, "_cluster_configs", None) or {} - cluster_config = cluster_configs.get(cluster_type) - if cluster_config is None: - return None - replica_config = getattr(cluster_config, "replica_config", None) - model_config = getattr(replica_config, "model_config", None) - if replica_config is None or model_config is None: - return None - architecture_profile = _resolve_model_architecture_profile(model_config) - if architecture_profile is None: - return None - base_name = get_moe_gating_base_model_name(model_name) - typed_family = _resolve_profile_typed_family_for_query( - architecture_profile, base_name - ) - if typed_family is None: - return None - _, layer_kind = typed_family - return self._resolve_typed_layer_contract( - base_name, - cluster_type, - replica_config, - is_moe_model=layer_kind is not LayerKind.DENSE, - ) - - def _is_ffn_typed_model_for_cluster( - self, - cluster_type: Optional[ClusterType], - model_name: str, - ) -> bool: - """Return whether a model belongs to an FFN domain excluded by a view.""" - - if cluster_type != ClusterType.DECODE_ATTN: - return False - cluster_config = self._cluster_configs.get(cluster_type) - replica_config = getattr(cluster_config, "replica_config", None) - model_config = getattr(replica_config, "model_config", None) - architecture_profile = _resolve_model_architecture_profile(model_config) - if architecture_profile is None: - return False - base_name = get_moe_gating_base_model_name(model_name) - return _resolve_profile_typed_family_for_query( - architecture_profile, base_name - ) is not None - - def _models_view_for_family( - self, - family_name: str, - cluster_type: Optional[ClusterType] = None, - ) -> Dict[str, BaseEstimator]: - """Project one measurement family's canonical registries to model names.""" - - names = set(self._legacy_model_registry(family_name)) - names.update( - model_name - for model_name, _identity in self._contract_model_registry(family_name) - ) - models: Dict[str, BaseEstimator] = {} - for model_name in sorted(names): - if self._is_ffn_typed_model_for_cluster(cluster_type, model_name): - continue - contract = self._resolve_cluster_model_contract(cluster_type, model_name) - identity = _serialize_selected_layer_cache_identity(contract) - model = self._get_family_model( - family_name, - model_name, - requested_identity=identity, - ) - if identity is not None and model is None: - raise ValueError( - f"No trained model for {model_name!r} matches the typed contract " - f"requested by cluster {cluster_type!r}: {identity}" - ) - if model is not None: - models[model_name] = model - return models - - def get_model( - self, - model_name: str, - precision: Optional[str] = None, - layer_contract: Optional[ResolvedLayerContract] = None, - ) -> Optional[BaseEstimator]: - """Get a model by name, precision, and optional typed contract.""" - - if self._all_dummy_mode: - return None - requested_identity = _serialize_selected_layer_cache_identity(layer_contract) - precision_key = precision.upper() if precision else None - for family_name in ("eager", "device_event", "kernel_only"): - model = self._get_family_model( - family_name, - model_name, - precision_key=precision_key, - requested_identity=requested_identity, - ) - if model is not None: - return model - - if precision_key is not None: - available_precisions = sorted( - { - str(value).upper() - for value in self._contract_precision_registry("eager") - } - | { - str(value).upper() - for value in self._contract_precision_registry("device_event") - } - | { - str(value).upper() - for value in self._contract_precision_registry("kernel_only") - } - | { - str(value).upper() - for value in self._legacy_precision_registry("eager") - } - | { - str(value).upper() - for value in self._legacy_precision_registry("device_event") - } - | { - str(value).upper() - for value in self._legacy_precision_registry("kernel_only") - } - ) - raise ValueError( - f"Model '{model_name}' not available for precision '{precision_key}'. " - f"Available precisions: {available_precisions}. " - "Ensure profiling data matches the requested precision." - ) - return None - - def _load_model_from_cache(self, model_name: str, model_hash: str) -> BaseEstimator: - with InterProcessReaderWriterLock(f"{self._cache_dir}/{model_hash}_model_lock.file").read_lock(): - cache_file = f"{self._cache_dir}/{model_name}_{model_hash}.pkl" - if not os.path.exists(cache_file): - return None - logger.info(f"✓ Loaded pre-trained model '{model_name}' from cache (hash: {model_hash})") - logger.info(f" Cache file: {cache_file}") - return pickle.load(open(cache_file, "rb")) - - def _store_model_in_cache(self, model_name: str, model_hash: str, model: BaseEstimator) -> None: - with InterProcessReaderWriterLock(f"{self._cache_dir}/{model_hash}_model_lock.file").write_lock(): - cache_file = f"{self._cache_dir}/{model_name}_{model_hash}.pkl" - atomic_pickle_dump(model, cache_file) - logger.info(f"✓ Saved trained model '{model_name}' to cache (hash: {model_hash})") - logger.info(f" Cache file: {cache_file}") - - def _ensure_exact_lookup_metadata( - self, - *, - model_name: str, - model_hash: str, - model: BaseEstimator, - df: pd.DataFrame, - feature_cols: List[str], - target_col: str, - ) -> None: - """Persist exact measured rows before an on-demand model cache is stored.""" - if hasattr(model, "_frontier_exact_lookup"): - exact_lookup = getattr(model, "_frontier_exact_lookup") - if not isinstance(exact_lookup, Mapping): - raise ValueError( - f"Exact lookup metadata for {model_name} must be a mapping" - ) - return - - setattr( - model, - "_frontier_exact_lookup", - _build_exact_feature_lookup(df, feature_cols, target_col), - ) - self._store_model_in_cache(model_name, model_hash, model) - - def get_models(self) -> Dict[str, Dict[str, BaseEstimator]]: - """Return the trained models grouped by measurement family.""" - if self._all_dummy_mode: - logger.debug("Returning empty models dict for dummy mode") - return {"eager": {}, "kernel_only": {}} - models = { - "eager": self._models_view_for_family("eager"), - "kernel_only": self._models_view_for_family("kernel_only"), - } - device_event_models = self._models_view_for_family("device_event") - if device_event_models: - models["device_event"] = device_event_models - return models - - def _event_family_for_cluster(self, cluster_type: ClusterType) -> str: - cluster_config = (getattr(self, "_cluster_configs", None) or {}).get( - cluster_type - ) - replica_config = getattr(cluster_config, "replica_config", None) - measurement_type = self._event_measurement_type_for_replica(replica_config) - return self._measurement_family_name(measurement_type) - - def get_models_for_cluster(self, cluster_type: ClusterType) -> Dict[str, Dict[str, BaseEstimator]]: - """Return a cluster-specific view of trained models grouped by measurement family.""" - if self._all_dummy_mode: - return {"eager": {}, "kernel_only": {}} - - event_family = self._event_family_for_cluster(cluster_type) - - def _event_models() -> Dict[str, BaseEstimator]: - return self._models_view_for_family(event_family, cluster_type) - - event_key = "eager" if event_family == "eager" else event_family - - if cluster_type == ClusterType.PREFILL: - models = { - event_key: _event_models(), - "kernel_only": {}, - } - return models - if cluster_type in [ClusterType.DECODE, ClusterType.DECODE_ATTN, ClusterType.DECODE_FFN]: - if ( - global_vars.get_sys_arch() == "pd-af-disaggregation" - and cluster_type == ClusterType.DECODE_ATTN - ): - models = { - event_key: _event_models(), - "kernel_only": self._models_view_for_family( - "kernel_only", cluster_type - ), - } - return models - if not self._is_kernel_only_measurement_enabled_for_cluster(cluster_type): - return { - event_key: _event_models(), - "kernel_only": {}, - } - return { - "eager": {}, - "kernel_only": self._models_view_for_family( - "kernel_only", cluster_type - ), - } - if cluster_type == ClusterType.MONOLITHIC: - kernel_only_models = {} - if self._is_kernel_only_measurement_enabled_for_cluster(cluster_type): - kernel_only_models = self._models_view_for_family( - "kernel_only", cluster_type - ) - return { - event_key: _event_models(), - "kernel_only": kernel_only_models, - } - raise ValueError(f"Unsupported cluster_type={cluster_type!r}") - - def get_required_capabilities(self) -> Dict[str, Any]: - """Return the analyzed requirements.""" - return self._required_capabilities - - def get_training_file_paths(self, cluster_type: ClusterType) -> Dict[str, str]: - """Get the resolved profiling file paths for a specific cluster type.""" - if cluster_type not in self._cluster_configs: - return {} - - cluster_config = self._cluster_configs[cluster_type] - replica_config = cluster_config.replica_config - execution_time_predictor_config = cluster_config.execution_time_predictor_config - - return resolve_training_file_paths( - execution_time_predictor_config, - device=replica_config.device, - model=replica_config.model_config.get_name(), - network_device=replica_config.network_device, - ) - - def get_training_context(self, cluster_type: ClusterType) -> Dict[str, Any]: - """ - Get comprehensive training context for a specific cluster type. - - Args: - cluster_type: The cluster type to get context for - - Returns: - Dictionary containing training context information - """ - if cluster_type not in self._cluster_configs: - return {} - - cluster_config = self._cluster_configs[cluster_type] - replica_config = cluster_config.replica_config - replica_scheduler_config = cluster_config.replica_scheduler_config - - file_paths = self.get_training_file_paths(cluster_type) - - return { - 'cluster_type': cluster_type, - 'device': replica_config.device, - 'model_name': replica_config.model_name, - 'attn_tensor_parallel_size': replica_config.attn_tensor_parallel_size, - 'moe_tensor_parallel_size': replica_config.moe_tensor_parallel_size, - 'num_pipeline_stages': replica_config.num_pipeline_stages, - 'network_device': replica_config.network_device, - 'block_size': replica_scheduler_config.block_size, - 'file_paths': file_paths, - 'max_tokens': getattr(replica_config, 'max_tokens', None), - 'max_batch_size': getattr(replica_config, 'max_batch_size', None) - } diff --git a/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py b/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py index ab46612c..92c8d21c 100644 --- a/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py +++ b/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py @@ -84,351 +84,42 @@ logger = init_logger(__name__) -def _normalize_routing_details_for_trace( - routing_details: Mapping[int, Mapping[int, Mapping[int, float]]], -) -> dict[str, dict[str, dict[str, float]]]: - """Return a strict JSON-safe copy of runtime routing details. - - The matrix checker compares this emitted object with an independently - materialized sidecar. The trace must therefore contain the actual - predictor-owned map, not a derived token allocation or a digest. - """ - - if not isinstance(routing_details, Mapping) or not routing_details: - raise ValueError("routing_details trace payload must be a non-empty mapping") - normalized: dict[str, dict[str, dict[str, float]]] = {} - for replica_id, per_layer in routing_details.items(): - if type(replica_id) is not int or replica_id < 0: - raise ValueError( - "routing_details trace replica IDs must be non-negative integers" - ) - if not isinstance(per_layer, Mapping) or not per_layer: - raise ValueError( - f"routing_details trace replica {replica_id} has no layer map" - ) - normalized_layers: dict[str, dict[str, float]] = {} - for layer_id, per_expert in per_layer.items(): - if type(layer_id) is not int or layer_id < 0: - raise ValueError( - "routing_details trace layer IDs must be non-negative integers" - ) - if not isinstance(per_expert, Mapping) or not per_expert: - raise ValueError( - f"routing_details trace layer {layer_id} has no expert map" - ) - normalized_experts: dict[str, float] = {} - for expert_id, ratio in per_expert.items(): - if type(expert_id) is not int or expert_id < 0: - raise ValueError( - "routing_details trace expert IDs must be non-negative integers" - ) - value = float(ratio) - if not math.isfinite(value) or value < 0.0: - raise ValueError( - "routing_details trace ratios must be finite and non-negative" - ) - normalized_experts[str(expert_id)] = value - ratio_sum = sum(normalized_experts.values()) - if not math.isclose(ratio_sum, 1.0, rel_tol=0.0, abs_tol=1e-12): - raise ValueError( - "routing_details trace ratios must sum to one " - f"for replica={replica_id} layer={layer_id}, got {ratio_sum}" - ) - normalized_layers[str(layer_id)] = normalized_experts - normalized[str(replica_id)] = normalized_layers - return normalized - - -def _get_moe_family_model_names() -> list[str]: - return list(get_family_profiling_names(MOE_FAMILY)) - - -def _get_moe_family_operator_by_model_name(model_name: str): - moe_ops = { - operator.profiling_name(): operator - for operator in MOE_FAMILY.profiling_ops() - } - if model_name not in moe_ops: - raise ValueError(f"Unsupported MoE op: {model_name}") - return moe_ops[model_name] - - -def _get_moe_gating_family_model_names() -> list[str]: - return [ - operator.profiling_name() - for operator in MOE_FAMILY.profiling_ops() - if operator.precision_name() == "moe_gating" - ] - - -_MOE_GATING_OPERATOR_NAMES = frozenset( - operator.name - for operator in MOE_FAMILY.profiling_ops() - if operator.precision_name() == "moe_gating" +from frontier.execution_time_predictor.moe_dataset_training import ( + MoeDatasetTraining, +) +from frontier.execution_time_predictor.moe_mtp_replay import MoeMtpReplay +from frontier.execution_time_predictor.moe_operator_times import MoeOperatorTimes +from frontier.execution_time_predictor.moe_predictor_helpers import ( + _build_moe_operator_times, + _get_moe_family_model_names, + _get_moe_family_operator_by_model_name, + _get_moe_gating_family_model_names, + _get_prefill_hot_moe_gating_model_names, + _is_moe_gating_family_model_name, + _normalize_routing_details_for_trace, + _validate_moe_columns, +) +from frontier.execution_time_predictor.moe_routing_workload import ( + MoeRoutingWorkload, ) -def _get_prefill_hot_moe_gating_model_names() -> list[str]: - return [ - f"{model_name}__prefill_hot" - for model_name in _get_moe_gating_family_model_names() - ] - - -def _is_moe_gating_family_model_name(model_name: str) -> bool: - base_model_name = get_moe_gating_base_model_name(model_name) - return _get_moe_family_operator_by_model_name( - base_model_name - ).precision_name() == "moe_gating" - - -def _build_moe_operator_times( - *, - mlp_norm_time: float, - moe_gating_linear_time: float, - moe_gating_routing_topk_time: float, - moe_shuffling_time: float, - moe_grouped_gemm_time: float, - share_expert_up_proj_time: float = 0.0, - share_expert_act_time: float = 0.0, - share_expert_down_proj_time: float = 0.0, - include_share_expert: bool = False, -) -> MoEOperatorTimes: - op_times = { - "post_attention_layernorm": mlp_norm_time, - "moe_gating_linear": moe_gating_linear_time, - "moe_gating_routing_topk": moe_gating_routing_topk_time, - "moe_shuffling": moe_shuffling_time, - "moe_grouped_gemm": moe_grouped_gemm_time, - } - if include_share_expert: - op_times.update( - { - "share_expert_up_proj": share_expert_up_proj_time, - "share_expert_act": share_expert_act_time, - "share_expert_down_proj": share_expert_down_proj_time, - } - ) - return MoEOperatorTimes(op_times=op_times) - - -def _validate_moe_columns(moe_df: pd.DataFrame) -> None: - """ - Validate that MoE DataFrame contains required split gating columns. - - This function enforces fail-fast behavior by rejecting legacy moe_gating - column format and requiring the split columns (moe_gating_linear and - moe_gating_routing_topk). - - Args: - moe_df: DataFrame containing MoE profiling data - - Raises: - ValueError: If required split columns are missing or if legacy - moe_gating column is present without split columns - """ - required_columns = [ - f"time_stats.{operator_name}.median" - for operator_name in get_family_profiling_names(MOE_FAMILY) - ] - - missing_columns = [col for col in required_columns if col not in moe_df.columns] - - if missing_columns: - # Check if legacy moe_gating column exists (for better error message) - legacy_col = "time_stats.moe_gating.median" - if legacy_col in moe_df.columns: - raise ValueError( - f"Missing required MoE columns: {missing_columns}. " - f"Found legacy '{legacy_col}' column which is no longer supported. " - f"Re-run MoE profiling with split gating scopes enabled to generate " - f"'moe_gating_linear' and 'moe_gating_routing_topk' columns." - ) - else: - raise ValueError( - f"Missing required MoE columns: {missing_columns}. " - f"Re-run MoE profiling with split gating scopes enabled." - ) - - -class SklearnMoEExecutionTimePredictor(SklearnExecutionTimePredictor): +class SklearnMoEExecutionTimePredictor( + MoeOperatorTimes, + MoeRoutingWorkload, + MoeDatasetTraining, + MoeMtpReplay, + SklearnExecutionTimePredictor, +): # Layer routing is deterministic for a predictor and the resulting # workload is immutable. Keep a bounded cache because the same # replica/layer/token shape is revisited across EP waves, while a # long-lived predictor must not retain an unbounded request history. _LAYER_WORKLOAD_CACHE_CAPACITY = 256 - @staticmethod - def _emit_routing_details_snapshot( - cluster_type: ClusterType, - routing_details: Mapping[int, Mapping[int, Mapping[int, float]]], - ) -> None: - """Emit the exact predictor-owned routing map for external validation.""" - - normalized = _normalize_routing_details_for_trace(routing_details) - payload = { - "schema_version": 1, - "cluster": cluster_type.name, - "routing_details": normalized, - } - logger.info( - "[ROUTING-SNAPSHOT] %s", - json.dumps(payload, sort_keys=True, separators=(",", ":")), - ) - - def _get_requested_moe_gating_routing_runtime_path(self) -> str: - return resolve_moe_gating_routing_runtime_path( - getattr(self, "_moe_routing_distribution_type", "balanced") - ) @staticmethod - def _get_ep_lane_routed_token_count( - batch: Batch, - lane_workload: Optional[EPLaneWorkload] = None, - ) -> Optional[int]: - """Return an EP lane's routed-token count, or ``None`` for full batches. - - Dummy mode still models the same five-phase EP contract as the - profiling-backed path. The lane-local routed compute therefore has - to depend on the materialized expert map even when the other dummy - components use fixed structural timings. - """ - - if lane_workload is None: - lane_workload = resolve_ep_lane_workload(batch, required=False) - if lane_workload is None: - return None - return lane_workload.routed_token_count - - def _admit_routed_ep_aggregate( - self, - batch: Batch, - *, - routed_moe: bool, - ep_size: Optional[int] = None, - router_topk: Optional[int] = None, - lane_workload: Optional[EPLaneWorkload] = None, - conservation_context: str = "routed MoE admission", - ) -> Optional[EPLaneWorkload]: - """Admit a concrete routed MoE call at the public predictor boundary. - - Concrete predictors own the semantic classification of a call. Once - that classification is routed MoE, an EP>1 call must identify one - physical lane before any mode-specific timing or lookup work begins. - The helper also validates the descriptor against the active predictor - topology and the source/lane token ledger. Workload construction - remains owned by the scheduler/materializer path. - """ - if type(routed_moe) is not bool: - raise ValueError("routed_moe must be a bool") - if not routed_moe: - return None - - if ep_size is None: - configured_ep_size = getattr(self, "_moe_ep_size", None) - if configured_ep_size is None: - replica_config = getattr(self, "_replica_config", None) - configured_ep_size = getattr( - replica_config, - "moe_expert_parallel_size", - None, - ) - else: - configured_ep_size = ep_size - if type(configured_ep_size) is not int or configured_ep_size < 1: - raise ValueError( - "routed MoE admission requires a positive integer EP size, got " - f"{configured_ep_size!r}" - ) - - batch_lane_workload = resolve_ep_lane_workload(batch, required=False) - explicit_lane_workload = ( - resolve_ep_lane_workload(lane_workload, required=True) - if lane_workload is not None - else None - ) - if ( - batch_lane_workload is not None - and explicit_lane_workload is not None - and batch_lane_workload != explicit_lane_workload - ): - raise ValueError( - "batch and lane_workload must refer to the same " - "EPLaneWorkload descriptor" - ) - resolved_lane_workload = explicit_lane_workload or batch_lane_workload - if configured_ep_size > 1 and resolved_lane_workload is None: - raise ValueError( - "Routed MoE prediction with EP>1 requires an EPLaneWorkload " - "descriptor" - ) - if resolved_lane_workload is None: - return None - - configured_router_topk = ( - getattr(self, "_router_topk", None) - if router_topk is None - else router_topk - ) - if configured_router_topk is not None: - if ( - type(configured_router_topk) is not int - or configured_router_topk < 1 - ): - raise ValueError( - "routed MoE admission requires a positive integer router top-k, " - f"got {configured_router_topk!r}" - ) - if resolved_lane_workload.router_topk != configured_router_topk: - raise ValueError( - "lane_workload router_topk does not match predictor topology: " - f"descriptor={resolved_lane_workload.router_topk}, " - f"predictor={configured_router_topk}" - ) - else: - configured_router_topk = resolved_lane_workload.router_topk - - if ( - resolved_lane_workload.moe_expert_parallel_size - != configured_ep_size - ): - raise ValueError( - "lane_workload EP size does not match predictor topology: " - f"descriptor={resolved_lane_workload.moe_expert_parallel_size}, " - f"predictor={configured_ep_size}" - ) - - # A descriptor attached to an EP lane entity already contains routed - # assignments, so its count must match that entity's physical width. - # An explicit descriptor paired with an ordinary source batch represents - # one assignment subset of the aggregate. The aggregate materializer, - # rather than this predictor boundary, owns its conservation ledger. - source_total_num_tokens = getattr(batch, "total_num_tokens", None) - if source_total_num_tokens is not None: - if ( - type(source_total_num_tokens) is not int - or source_total_num_tokens < 0 - ): - raise ValueError( - "routed MoE admission requires batch.total_num_tokens to be " - "a non-negative integer, got " - f"{source_total_num_tokens!r}" - ) - if batch_lane_workload is not None: - expected_routed_token_count = source_total_num_tokens - if ( - resolved_lane_workload.routed_token_count - != expected_routed_token_count - ): - raise ValueError( - f"Token conservation violated in {conservation_context}: " - f"allocated {resolved_lane_workload.routed_token_count}, " - f"expected {expected_routed_token_count}" - ) - - return resolved_lane_workload - @staticmethod def _resolve_moe_layer_classification( model_config: Any, *, @@ -645,1519 +336,124 @@ def _get_dummy_execution_time( attention_layer_post_proj_execution_time=( base_time if include_attention else 0.0 ), - attn_norm_time=base_time if include_attention else 0.0, - mlp_norm_time=base_time if include_ffn else 0.0, - add_time=add_time, - add_attn_residual_time=add_attn_residual_time, - add_ffn_residual_time=add_ffn_residual_time, - tensor_parallel_communication_time=attn_tp_allreduce_time, - attn_tensor_parallel_allreduce_time=attn_tp_allreduce_time, - moe_tensor_parallel_allreduce_time=ffn_tp_allreduce_time, - pipeline_parallel_communication_time=base_time if include_stage_owned else 0.0, - expert_parallel_communication_time=expert_parallel_comm_time, - moe_gating_time=base_time if is_moe else 0.0, - moe_shuffling_time=( - 0.0 if zero_routed_ep_lane else base_time - ) if is_moe else 0.0, - schedule_time=base_time if include_stage_owned else 0.0, - sampler_e2e_time=base_time if include_stage_owned else 0.0, - prepare_inputs_e2e_time=base_time if include_stage_owned else 0.0, - process_model_outputs_time=base_time if include_stage_owned else 0.0, - ray_comm_time=base_time if include_stage_owned else 0.0, - pp_stage_boundary_handoff_time=pp_stage_boundary_handoff_time, - is_moe=is_moe, - mlp_layer_up_proj_execution_time=( - base_time if include_ffn and not is_moe else 0.0 - ), - mlp_layer_down_proj_execution_time=( - base_time if include_ffn and not is_moe else 0.0 - ), - mlp_layer_act_execution_time=( - base_time if include_ffn and not is_moe else 0.0 - ), - moe_grouped_gemm_time=moe_grouped_gemm_time, - share_expert_up_proj_time=share_expert_time, - share_expert_down_proj_time=share_expert_time, - share_expert_act_time=share_expert_time, - tensor_parallel_allgather_time=ffn_tp_allgather_time, - share_expert_tensor_parallel_allreduce_time=share_expert_tp_allreduce_time, - dp_input_allreduce_time=dp_input_allreduce_time, - dp_output_allreduce_time=dp_output_allreduce_time, - communication_operator_times=communication_operator_times, - mlp_operator_times=mlp_operator_times, - moe_operator_times=moe_operator_times, - ) - - def __init__( - self, - predictor_config: BaseExecutionTimePredictorConfig, - replica_config: ReplicaConfig, - replica_scheduler_config: BaseReplicaSchedulerConfig, - metrics_config: MetricsConfig, - model_manager: ExecutionTimePredictionModelManager = None, - cluster_type: ClusterType = None, - training_file_paths: Dict[str, str] = None, - cc_backend: Optional["BaseCCBackend"] = None, - actual_replica_ids: Optional[list] = None, - ) -> None: - self._is_moe = True - self._router_topk = replica_config.router_topk - self._moe_tp_size = replica_config.moe_tensor_parallel_size - self._moe_ep_size = replica_config.moe_expert_parallel_size - self._actual_replica_ids = actual_replica_ids - self._attention_query_cache_hits = 0 - self._attention_query_cache_misses = 0 - self._layer_workload_cache_capacity = self._LAYER_WORKLOAD_CACHE_CAPACITY - self._layer_workload_cache = OrderedDict() - - # Initialize the canonical distribution selector before parent init so - # profiling paths choose matching gating-runtime metadata. - self._moe_routing_distribution_type = str( - getattr(replica_config, "moe_routing_distribution_type", "balanced") - ).strip().lower() - valid_distribution_types = {"balanced", "random", "skewed", "zipf"} - if self._moe_routing_distribution_type not in valid_distribution_types: - raise ValueError( - "moe_routing_distribution_type must be one of " - f"{sorted(valid_distribution_types)}, got " - f"{self._moe_routing_distribution_type!r}" - ) - self._moe_routing_seed = getattr(replica_config, "moe_routing_seed", 42) - if type(self._moe_routing_seed) is not int or self._moe_routing_seed < 0: - raise ValueError( - "moe_routing_seed must be an exact non-negative int, " - f"got {self._moe_routing_seed}." - ) - self._moe_gating_routing_runtime_path = ( - resolve_moe_gating_routing_runtime_path( - self._moe_routing_distribution_type - ) - ) - - super().__init__( - predictor_config, - replica_config, - replica_scheduler_config, - metrics_config, - model_manager, - cluster_type, - training_file_paths, - cc_backend, - ) - - # Pre-compute one global routing source. EP ownership is applied later - # by the shared per-layer materializer. - self._monolithic_routing_details = None - self._global_routing_allocations = self._init_global_routing_allocations() - if self._cluster_type == ClusterType.MONOLITHIC and self._model_config.is_moe: - self._monolithic_routing_details = self._build_shared_routing_details() - self._emit_routing_details_snapshot( - ClusterType.MONOLITHIC, - self._monolithic_routing_details, - ) - logger.info( - "[MoE Routing] Initialized global routing allocations: " - "distribution=%s, seed=%s, num_layers=%s", - self._moe_routing_distribution_type, - self._moe_routing_seed, - len(self._global_routing_allocations), - ) - - def _init_global_routing_allocations(self) -> Dict[int, Dict[int, float]]: - """Pre-compute global expert allocation ratios for shared-domain EP sync. - - Monolithic decode with EP enabled needs a global view across all experts to - derive per-lane post-MoE arrival skew before the shared-domain all-reduce. - """ - total_experts = self._replica_config.total_expert_num - if self._cluster_type == ClusterType.DECODE_ATTN or not self._model_config.is_moe: - return {} - num_layers = self._model_config.num_layers - - if type(total_experts) is not int or total_experts <= 0: - raise ValueError( - "total_expert_num must be an exact positive int for routing; " - f"got {total_experts!r}" - ) - if type(self._moe_ep_size) is not int or self._moe_ep_size <= 0: - raise ValueError( - "moe_expert_parallel_size must be an exact positive int for routing; " - f"got {self._moe_ep_size!r}" - ) - if total_experts % self._moe_ep_size != 0: - raise ValueError( - "total_expert_num must be divisible by moe_expert_parallel_size; " - f"got total_expert_num={total_experts}, " - f"moe_expert_parallel_size={self._moe_ep_size}" - ) - - distribution_type = self._moe_routing_distribution_type - allocations: Dict[int, Dict[int, float]] = {} - for layer_id in range(num_layers): - allocations[layer_id] = generate_moe_routing_ratios( - total_expert_num=total_experts, - distribution_type=distribution_type, - seed=self._moe_routing_seed, - layer_id=layer_id, - ) - - return allocations - - def _build_shared_routing_details( - self, - ) -> Dict[int, Dict[int, Dict[int, float]]]: - """Expose one immutable-shape routing source for monolithic schedulers. - - ``_global_routing_allocations`` is generated once per model layer and is - intentionally replica-independent. The scheduler, however, performs an - exact ``(replica_id, global_layer_id)`` lookup. Materialize that lookup - shape here without generating a second distribution or assigning tokens - to requests. Integer token accounting and EP ownership splitting remain - the responsibility of the shared per-layer materializer. - - The current monolithic predictor is constructed from ``ReplicaConfig`` - rather than ``ClusterConfig``. The canonical cluster capacity is - bound to ``ReplicaConfig.cluster_num_replicas`` before this method is - called. When the simulator supplies ``_actual_replica_ids``, those - process-global IDs are used as the outer map keys; otherwise local - ``range(replica_count)`` keys support standalone predictor construction. - A missing capacity is an invalid topology, not a condition to infer - from an attention-DP field. - """ - replica_count = self._replica_config.cluster_num_replicas - if type(replica_count) is not int or replica_count <= 0: - raise ValueError( - "A positive cluster replica count is required to build shared " - f"routing details; got {replica_count!r}" - ) - - actual_replica_ids = self._actual_replica_ids - if actual_replica_ids is None: - replica_ids = list(range(replica_count)) - else: - if not isinstance(actual_replica_ids, (list, tuple)): - raise ValueError( - "actual_replica_ids must be a list or tuple when provided" - ) - replica_ids = list(actual_replica_ids) - if len(replica_ids) != replica_count: - raise ValueError( - "actual_replica_ids length must match cluster replica count; " - f"got {len(replica_ids)} for {replica_count} replicas" - ) - if any( - type(replica_id) is not int or replica_id < 0 - for replica_id in replica_ids - ): - raise ValueError( - "actual_replica_ids must contain exact non-negative integers" - ) - if len(set(replica_ids)) != len(replica_ids): - raise ValueError("actual_replica_ids must be unique") - - return { - replica_id: { - layer_id: dict(expert_ratios) - for layer_id, expert_ratios in self._global_routing_allocations.items() - } - for replica_id in replica_ids - } - - - def _get_routing_details_for_cluster(self, cluster_type: ClusterType): - """Return the exact pre-generated routing map for one serving role.""" - attribute_by_cluster = { - ClusterType.MONOLITHIC: "_monolithic_routing_details", - ClusterType.PREFILL: "_prefill_routing_details", - ClusterType.DECODE: "_decode_routing_details", - ClusterType.DECODE_FFN: "_decode_ffn_routing_details", - } - attribute_name = attribute_by_cluster.get(cluster_type) - if attribute_name is None: - raise ValueError( - f"MoE routing materialization does not support cluster_type={cluster_type}" - ) - routing_details = getattr(self, attribute_name, None) - if routing_details is None: - raise ValueError( - f"Missing pre-generated routing details for cluster_type={cluster_type}" - ) - return routing_details - - def _get_cluster_replica_config(self, cluster_type: ClusterType) -> ReplicaConfig: - """Return the serving replica; disaggregated predictors override by role.""" - return self._replica_config - - def _materialize_layer_ep_workload( - self, batch: Batch, cluster_type: ClusterType, layer_id: int - ) -> LayerEPWorkload: - """Materialize one exact Replica-local EP workload for a MoE layer.""" - # Routing tables are built once by the constructor and have no runtime - # mutation API. Topology and exact replica/layer/token identity are in - # the key; the frozen workload can be shared across repeated EP waves. - workload_cache = self._layer_workload_cache - cache_capacity = self._layer_workload_cache_capacity - cluster_replica_config = self._get_cluster_replica_config(cluster_type) - routing_details = self._get_routing_details_for_cluster(cluster_type) - target_replica_id = int(batch.replica_id) - global_layer_id = int(layer_id) - routing_token_count = int(batch.total_num_tokens) - router_topk = int(cluster_replica_config.router_topk) - total_expert_num = int(cluster_replica_config.total_expert_num) - moe_ep_size = int(cluster_replica_config.moe_expert_parallel_size) - cache_key = ( - cluster_type, - target_replica_id, - global_layer_id, - routing_token_count, - router_topk, - total_expert_num, - moe_ep_size, - ) - cached_workload = workload_cache.get(cache_key) - if cached_workload is not None: - workload_cache.move_to_end(cache_key) - return cached_workload - workload = materialize_layer_ep_workload( - routing_ratios=resolve_routing_details( - routing_details, - target_replica_id=target_replica_id, - global_layer_id=global_layer_id, - ), - target_replica_id=target_replica_id, - global_layer_id=global_layer_id, - routing_token_count=routing_token_count, - router_topk=router_topk, - total_expert_num=total_expert_num, - moe_expert_parallel_size=moe_ep_size, - expert_to_ep=build_contiguous_expert_ownership( - total_expert_num, - moe_ep_size, - ), - ) - workload_cache[cache_key] = workload - workload_cache.move_to_end(cache_key) - while len(workload_cache) > cache_capacity: - workload_cache.popitem(last=False) - return workload - - def _resolve_layer_lane_workload( - self, - batch: Batch, - *, - cluster_type: ClusterType, - layer_id: int, - ) -> EPLaneWorkload: - """Resolve one physical lane descriptor at a predictor boundary. - - Scheduler-created lane entities already carry the descriptor. A - regular batch may be materialized into one lane only for EP=1; an EP>1 - aggregate must be expanded by the scheduler's lane wave first so the - predictor never guesses which local expert domain a global map denotes. - """ - - lane_workload = resolve_ep_lane_workload(batch, required=False) - if lane_workload is not None: - return lane_workload - - layer_workload = self._materialize_layer_ep_workload( - batch=batch, - cluster_type=cluster_type, - layer_id=layer_id, - ) - participant_ep_ids = tuple(layer_workload.participant_ep_ids) - if len(participant_ep_ids) != int(self._moe_ep_size): - raise ValueError( - "materialized EP lane count does not match predictor topology: " - f"descriptors={len(participant_ep_ids)}, predictor={self._moe_ep_size}" - ) - if len(participant_ep_ids) != 1: - raise ValueError( - "regular aggregate MoE prediction requires a physical EP lane " - "for EP>1; scheduler lane materialization is required" - ) - return layer_workload.lane(participant_ep_ids[0]) - - def _resolve_shared_domain_lane_workloads( - self, - batch: Batch, - *, - cluster_type: ClusterType, - layer_id: int, - ) -> tuple[EPLaneWorkload, ...]: - """Resolve every physical lane for a shared-domain MoE timing probe.""" - - lane_count = int(self._moe_ep_size) - if lane_count <= 0: - raise ValueError( - "shared-domain MoE lane resolution requires a positive EP size, " - f"got {lane_count}" - ) - - lane_workload = resolve_ep_lane_workload(batch, required=False) - if lane_workload is not None: - if lane_workload.moe_expert_parallel_size != lane_count: - raise ValueError( - "batch lane workload EP size does not match predictor: " - f"descriptor={lane_workload.moe_expert_parallel_size}, " - f"predictor={lane_count}" - ) - lane_workloads = (lane_workload,) - else: - workload = self._materialize_layer_ep_workload( - batch=batch, - cluster_type=cluster_type, - layer_id=layer_id, - ) - lane_workloads = tuple( - workload.lane(ep_id) for ep_id in workload.participant_ep_ids - ) - - if len(lane_workloads) != lane_count: - raise ValueError( - "materialized EP lane count does not match predictor topology: " - f"descriptors={len(lane_workloads)}, predictor={lane_count}" - ) - return lane_workloads - - @staticmethod - def _get_dummy_shared_domain_moe_scope_time( - execution_time: ExecutionTime, - ) -> float: - """Return the fixed per-operator MoE scope used by shared-domain decode. - - The generic dummy ``ExecutionTime`` keeps the deprecated aggregate - ``moe_gating_time`` contract by splitting that baseline across the two - structured gating fields. The shared-domain decode contract models - each named gating operator as one fixed structural slot, matching its - historical five-operator scope. Resolve that compatibility at this - boundary from the MoE family registry; all other operators retain the - descriptor-aware structured timing, including zero-lane routed work. - """ - - moe_time = execution_time.moe_or_mlp_time_component - if not isinstance(moe_time, MoETime): - raise ValueError( - "shared-domain dummy timing requires a MoE execution component" - ) - operator_times = moe_time.operator_times - if operator_times is None: - raise ValueError( - "shared-domain dummy timing requires structured MoE operator times" - ) - - scope_time = 0.0 - for operator_name, operator_time in operator_times.op_times.items(): - if operator_name in _MOE_GATING_OPERATOR_NAMES: - # ``moe_gating_time`` is the one fixed dummy baseline for each - # named gating operator; the structured fields store its - # compatibility split as 0.5 * baseline each. - scope_time += float(moe_time.moe_gating_time) - else: - scope_time += float(operator_time) - return scope_time - - def predict_monolithic_decode_shared_domain_lane_moe_times_ms( - self, - batch: Batch, - layer_id: int, - ) -> Dict[int, float]: - """Estimate per-EP-lane pre-collective MoE time for monolithic pure decode. - - Returns per-lane post-attention MoE compute in milliseconds. The result is - used by the MONOLITHIC decode sync path to model shared-domain readiness skew - before `expert_parallel_allreduce`. - """ - if self._enable_dummy_mode: - lane_workloads = self._resolve_shared_domain_lane_workloads( - batch, - cluster_type=ClusterType.MONOLITHIC, - layer_id=layer_id, - ) - lane_times_ms: Dict[int, float] = {} - for lane_workload in lane_workloads: - # This helper returns only the per-layer MoE scope. The dummy - # execution seam needs a stage value for its complete object, - # but no stage-boundary term is included in the component total. - execution_time = self._get_dummy_execution_time( - batch, - pipeline_stage=0, - include_attention=False, - lane_workload=lane_workload, - ) - lane_times_ms[lane_workload.ep_id] = ( - self._get_dummy_shared_domain_moe_scope_time(execution_time) - ) - return lane_times_ms - - lane_workloads = self._resolve_shared_domain_lane_workloads( - batch, - cluster_type=ClusterType.MONOLITHIC, - layer_id=layer_id, - ) - - post_attention_layernorm_time = self._get_mlp_norm_layer_act_execution_time(batch) - gating_linear_time = self._get_gating_linear_time(batch) - gating_routing_topk_time = self._get_gating_routing_topk_time(batch) - share_expert_total_time = 0.0 - if self._model_config.supports_share_expert(): - share_expert_total_time = ( - self._get_share_expert_up_proj_execution_time(batch) - + self._get_share_expert_down_proj_execution_time(batch) - + self._get_share_expert_act_execution_time(batch) - ) - - lane_times_ms: Dict[int, float] = {} - for lane_workload in lane_workloads: - lane_id = lane_workload.ep_id - shuffling_time = self._get_moe_shuffling_time( - batch, - moe_tokens_input=lane_workload, - ) - grouped_gemm_time = self._get_grouped_gemm_time( - lane_workload, - batch=batch, - ) - - lane_times_ms[lane_id] = ( - post_attention_layernorm_time - + gating_linear_time - + gating_routing_topk_time - + shuffling_time - + grouped_gemm_time - + share_expert_total_time - ) - - return lane_times_ms - - # Load imbalance feature columns used for MoE training (aligned with SharedPredictionModelManager) - # Reference: frontier/training/moe_trainer.py lines 224-239 (authoritative source) - MOE_LOAD_IMBALANCE_FEATURES = [ - # Config features (6) - describe model configuration - "total_routed_tokens", # Total tokens after routing (num_tokens * router_topk) - "num_experts_per_device", # Number of experts per device after EP sharding - "hidden_dim", # Model hidden dimension - "expert_hidden_dim", # Expert FFN hidden dimension - "router_topk", # Number of experts each token is routed to - "model_expansion_ratio", # expert_hidden_dim / hidden_dim - # Derived features (2) - derived from config and routing - "tokens_per_expert_avg", # Average tokens per expert - "tokens_to_experts_ratio", # tokens / num_experts ratio - # Load features (6) - describe load distribution characteristics - "expert_utilization", # Proportion of experts with non-zero load - "min_load_ratio", # Min load / average load - "load_imbalance_cv", # Coefficient of Variation: std/mean, key imbalance metric - "max_load_ratio", # Max load / average load - "load_entropy", # Entropy of load distribution (higher = more uniform) - "load_gini_coefficient", # Gini coefficient: 0=equality, 1=inequality - ] - - @staticmethod - def _get_moe_op_tp_key( - op_name: str, - moe_tp_size: int, - cluster_type: ClusterType | None = None, - ) -> int: - try: - return resolve_moe_operator_tp_key( - op_name, - moe_tp_size=moe_tp_size, - cluster_type=cluster_type, - family=MOE_FAMILY, - ) - except ValueError as exc: - if str(exc).startswith("Unsupported MoE op:"): - raise ValueError( - f"Unsupported MoE op for TP mapping: {op_name}" - ) from exc - raise - - @staticmethod - def _is_moe_op_ep_agnostic(op_name: str) -> bool: - try: - return is_moe_operator_ep_agnostic(op_name, family=MOE_FAMILY) - except ValueError as exc: - if str(exc).startswith("Unsupported MoE op:"): - raise ValueError( - f"Unsupported MoE op for EP mapping: {op_name}" - ) from exc - raise - - def _validate_moe_dataset_contract( - self, - moe_df: pd.DataFrame, - moe_input_file: str, - model_names: List[str], - moe_tp_size: int, - moe_ep_size: int, - ) -> pd.DataFrame: - """Validate op-level MoE key coverage and return model-filtered dataframe.""" - _validate_moe_columns(moe_df) - required_columns = [ - "num_experts", - "router_topk", - "hidden_dim", - "expert_hidden_dim", - "num_tensor_parallel_workers", - "expert_parallel_size", - ] - missing_columns = [col for col in required_columns if col not in moe_df.columns] - if missing_columns: - raise ValueError( - f"MoE dataset contract validation failed for {moe_input_file}: " - f"missing required columns {missing_columns}." - ) - - model_config = self._model_config - base_df = moe_df[ - (moe_df["num_experts"] == model_config.num_experts) - & (moe_df["router_topk"] == model_config.num_experts_per_tok) - & (moe_df["hidden_dim"] == model_config.embedding_dim) - & (moe_df["expert_hidden_dim"] == model_config.mlp_hidden_dim) - ].copy() - - if len(base_df) == 0: - raise ValueError( - "MoE dataset contract validation failed: no rows match model configuration in " - f"{moe_input_file}. Required: num_experts={model_config.num_experts}, " - f"router_topk={model_config.num_experts_per_tok}, hidden_dim={model_config.embedding_dim}, " - f"expert_hidden_dim={model_config.mlp_hidden_dim}." - ) - - available_pairs = sorted( - { - (int(tp), int(ep)) - for tp, ep in base_df[ - ["num_tensor_parallel_workers", "expert_parallel_size"] - ].drop_duplicates().itertuples(index=False, name=None) - } - ) - requested_routing_runtime_path = ( - self._get_requested_moe_gating_routing_runtime_path() - ) - - missing_requirements: List[str] = [] - for model_name in model_names: - base_model_name = get_moe_gating_base_model_name(model_name) - tp_key = self._get_moe_op_tp_key( - base_model_name, - moe_tp_size, - cluster_type=getattr(self, "_cluster_type", None), - ) - requirement_parts = [f"TP={tp_key}"] - if self._is_moe_op_ep_agnostic(base_model_name): - op_df = base_df[base_df["num_tensor_parallel_workers"] == tp_key] - requirement_parts.append("EP=ANY") - else: - op_df = base_df[ - (base_df["num_tensor_parallel_workers"] == tp_key) - & (base_df["expert_parallel_size"] == moe_ep_size) - ] - requirement_parts.append(f"EP={moe_ep_size}") - if base_model_name == "moe_gating_routing_topk": - op_df = filter_moe_gating_routing_topk_rows( - op_df, - requested_runtime_path=requested_routing_runtime_path, - source_name=moe_input_file, - ) - requirement_parts.append( - f"routing_runtime_path={requested_routing_runtime_path}" - ) - if _is_moe_gating_family_model_name(base_model_name): - op_df = filter_moe_gating_rows_by_runtime_context( - op_df, - requested_context=DEFAULT_MOE_GATING_RUNTIME_CONTEXT, - source_name=moe_input_file, - ) - requirement_parts.append( - "gating_runtime_context=" - f"{DEFAULT_MOE_GATING_RUNTIME_CONTEXT}" - ) - requirement = ", ".join(requirement_parts) - if len(op_df) == 0: - missing_requirements.append(f"{model_name} requires {requirement}") - continue - target_col = f"time_stats.{base_model_name}.median" - if op_df[target_col].dropna().empty: - missing_requirements.append( - f"{model_name} requires {requirement}, target={target_col} " - "to contain at least one non-NaN row" - ) - - if missing_requirements: - requirement_text = "\n - ".join(missing_requirements) - raise ValueError( - "MoE dataset contract validation failed before training.\n" - f"File: {moe_input_file}\n" - "Missing op-level key coverage:\n" - f" - {requirement_text}\n" - f"Available (TP, EP) pairs for matched model rows: {available_pairs}" - ) - - return base_df - - def _train_moe_models(self) -> Dict[str, BaseEstimator]: - """Train MoE-specific models (gating, shuffling, grouped_gemm) for independent training mode. - - For moe_grouped_gemm, uses 14 load-imbalance features if available in the profiling data. - This enables simulation mode with per-expert token allocation. - Other MoE models (gating_linear, gating_routing_topk, shuffling) use only num_tokens. - """ - models = {} - moe_input_file = getattr(self, "_moe_input_file", "/synthetic/moe.csv") - - if not os.path.exists(moe_input_file): - logger.warning(f"MoE input file does not exist: {moe_input_file}") - return models - - moe_df = pd.read_csv(moe_input_file) - if TYPED_OPERATOR_CONTRACTS_COLUMN in moe_df.columns: - # Validate every row before scalar or TP/EP filtering can hide a - # malformed typed contract. - moe_df[TYPED_OPERATOR_CONTRACTS_COLUMN].map( - lambda raw_contracts: validate_typed_operator_contracts( - raw_contracts, - model_config=self._model_config, - ) - ) - - metadata = self._get_profiling_metadata(moe_df, moe_input_file) - self._validate_active_measurement_type(metadata, moe_input_file) - - tp_col = "num_tensor_parallel_workers" - ep_col = "expert_parallel_size" - moe_tp_size = self._replica_config.moe_tensor_parallel_size - moe_ep_size = self._replica_config.moe_expert_parallel_size - - if tp_col not in moe_df.columns: - raise ValueError( - f"Required column '{tp_col}' is missing in {moe_input_file}. " - "Re-run MoE profiling with TP metadata enabled." - ) - if ep_col not in moe_df.columns: - raise ValueError( - f"Required column '{ep_col}' is missing in {moe_input_file}. " - "Re-run MoE profiling with EP metadata enabled." - ) - - base_model_names = _get_moe_family_model_names() - model_names = list(base_model_names) - model_filtered_df = self._validate_moe_dataset_contract( - moe_df, - moe_input_file, - base_model_names, - moe_tp_size, - moe_ep_size, - ) - if should_enable_prefill_hot_moe_gating_contract( - model_config=self._model_config, - ): - if has_prefill_hot_moe_gating_rows(model_filtered_df): - model_names.extend(_get_prefill_hot_moe_gating_model_names()) - else: - logger.warning( - "Prefill-hot gating contract is enabled for model=%s, but " - "dataset %s has no usable prefill_hot rows; skipping " - "__prefill_hot pseudo-model training.", - self._replica_config.model_name, - moe_input_file, - ) - - self._register_profiling_metadata_for_ops( - model_names, metadata, moe_input_file - ) - - requested_routing_runtime_path = ( - self._get_requested_moe_gating_routing_runtime_path() - ) - moe_df_cache: Dict[ - tuple[int, Optional[int], Optional[str], Optional[str]], pd.DataFrame - ] = {} - - def _get_moe_df_for_op( - model_name: str, - ) -> tuple[pd.DataFrame, int, Optional[int]]: - base_model_name = get_moe_gating_base_model_name(model_name) - tp_key = self._get_moe_op_tp_key( - base_model_name, - moe_tp_size, - cluster_type=getattr(self, "_cluster_type", None), - ) - ep_key: Optional[int] - if self._is_moe_op_ep_agnostic(base_model_name): - ep_key = None - else: - ep_key = moe_ep_size - runtime_path_key: Optional[str] = None - if base_model_name == "moe_gating_routing_topk": - runtime_path_key = requested_routing_runtime_path - gating_context_key: Optional[str] = None - if _is_moe_gating_family_model_name(base_model_name): - gating_context_key = DEFAULT_MOE_GATING_RUNTIME_CONTEXT - if model_name.endswith("__prefill_hot"): - gating_context_key = PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT - cache_key = (tp_key, ep_key, runtime_path_key, gating_context_key) - if cache_key not in moe_df_cache: - filtered_df = model_filtered_df[ - model_filtered_df[tp_col] == tp_key - ].copy() - if ep_key is not None: - filtered_df = filtered_df[ - filtered_df[ep_col] == ep_key - ].copy() - if runtime_path_key is not None: - filtered_df = filter_moe_gating_routing_topk_rows( - filtered_df, - requested_runtime_path=runtime_path_key, - source_name=moe_input_file, - ) - if gating_context_key is not None: - filtered_df = filter_moe_gating_rows_by_runtime_context( - filtered_df, - requested_context=gating_context_key, - source_name=moe_input_file, - ) - if len(filtered_df) == 0: - ep_desc = "ANY" if ep_key is None else str(ep_key) - raise ValueError( - f"No MoE data after filtering for TP={tp_key}, EP={ep_desc}. " - f"Requested by op-level TP mapping in {moe_input_file}." - ) - filtered_df["num_tokens_rounded"] = filtered_df["num_tokens"].apply( - lambda x: max(1, round(x / 8) * 8) - ) - moe_df_cache[cache_key] = filtered_df - return moe_df_cache[cache_key], tp_key, ep_key - - for model_name in model_names: - try: - op_df, moe_tp_key, moe_ep_key = _get_moe_df_for_op(model_name) - except PrefillHotRowsUnavailableError as exc: - logger.warning( - "Skipping %s because prefill-hot gating rows are unavailable " - "for the requested TP/EP slice (%s).", - model_name, - exc, - ) - continue - target_op_name = get_moe_gating_base_model_name(model_name) - target_col = f"time_stats.{target_op_name}.median" - if target_col not in op_df.columns: - ep_desc = "ANY" if moe_ep_key is None else str(moe_ep_key) - raise ValueError( - f"Column '{target_col}' not found in MoE dataframe for TP={moe_tp_key}, EP={ep_desc}. " - "Re-run MoE profiling with split gating columns." - ) - - # Per-operation feature selection (aligned with SharedPredictionModelManager). - if model_name == "moe_grouped_gemm": - available_load_features = [ - f for f in self.MOE_LOAD_IMBALANCE_FEATURES if f in op_df.columns - ] - has_load_imbalance_features = len(available_load_features) == len( - self.MOE_LOAD_IMBALANCE_FEATURES - ) - if 0 < len(available_load_features) < len(self.MOE_LOAD_IMBALANCE_FEATURES): - missing_features = [ - f - for f in self.MOE_LOAD_IMBALANCE_FEATURES - if f not in op_df.columns - ] - raise ValueError( - f"Partial load imbalance features found ({len(available_load_features)}/" - f"{len(self.MOE_LOAD_IMBALANCE_FEATURES)}) for TP={moe_tp_key}. " - f"Missing: {missing_features}." - ) - if has_load_imbalance_features: - feature_cols = available_load_features - logger.info( - f" {model_name}: Using load imbalance features ({len(feature_cols)} features, TP={moe_tp_key})" - ) - else: - feature_cols = ["num_tokens"] - logger.info( - f" {model_name}: Load imbalance features not found; using num_tokens only (TP={moe_tp_key})" - ) - elif model_name == "moe_shuffling": - available_load_features = [ - f for f in self.MOE_LOAD_IMBALANCE_FEATURES if f in op_df.columns - ] - if len(available_load_features) == len(self.MOE_LOAD_IMBALANCE_FEATURES): - feature_cols = available_load_features - logger.info( - f" {model_name}: Using load imbalance features ({len(feature_cols)} features, TP={moe_tp_key})" - ) - else: - feature_cols = ["num_tokens"] - logger.info( - f" {model_name}: Full load imbalance features unavailable; using num_tokens only (TP={moe_tp_key})" - ) - else: - feature_cols = ["num_tokens"] - logger.info(f" {model_name}: Using num_tokens only (1 feature, TP={moe_tp_key})") - - models[model_name] = self._train_model( - model_name=model_name, - df=op_df, - feature_cols=feature_cols, - target_col=target_col, - ) - logger.info(f"Trained MoE model: {model_name}") - - return models - - def _register_additional_profiling_metadata_from_files(self) -> None: - moe_input_file = self._moe_input_file - model_names = _get_moe_family_model_names() - if should_enable_prefill_hot_moe_gating_contract( - model_config=self._model_config, - ): - include_prefill_hot_models = False - try: - moe_df = pd.read_csv(moe_input_file) - except FileNotFoundError: - moe_df = None - if moe_df is not None: - if TYPED_OPERATOR_CONTRACTS_COLUMN in moe_df.columns: - moe_df[TYPED_OPERATOR_CONTRACTS_COLUMN].map( - lambda raw_contracts: validate_typed_operator_contracts( - raw_contracts, - model_config=self._model_config, - ) - ) - include_prefill_hot_models = has_prefill_hot_moe_gating_rows(moe_df) - if include_prefill_hot_models: - model_names.extend(_get_prefill_hot_moe_gating_model_names()) - self._register_profiling_metadata_from_file(moe_input_file, model_names) - - def _train_models(self) -> Dict[str, BaseEstimator]: - """Override to include MoE model training for independent training mode.""" - models = super()._train_models() - - if self._model_manager is None: - moe_models = self._train_moe_models() - models.update(moe_models) - logger.info(f"Trained MoE models independently: {list(moe_models.keys())}") - else: - logger.info("MoE models loaded from ExecutionTimePredictionModelManager.") - - return models - - def _predict_for_compute_models(self) -> Dict[str, Any]: - predictions = super()._predict_for_compute_models() - extra_model_names = _get_prefill_hot_moe_gating_model_names() - num_token_range = np.arange(1, self._max_tokens + 1) - X = pd.DataFrame({"num_tokens": num_token_range}) - for model_name in extra_model_names: - if model_name not in self._models: - continue - model = self._models[model_name] - predictions[model_name] = self._get_model_prediction( - model_name, model, X - ) - return predictions - - def _select_moe_gating_prediction_model_name( - self, - base_model_name: str, - batch: Batch, - ) -> str: - requested_context = DEFAULT_MOE_GATING_RUNTIME_CONTEXT - if should_use_prefill_hot_moe_gating_context( - model_config=self._model_config, - batch=batch, - ): - requested_context = PREFILL_HOT_MOE_GATING_RUNTIME_CONTEXT - candidate_model_name = get_moe_gating_prediction_model_name( - base_model_name, - requested_context=requested_context, - ) - if candidate_model_name in self._predictions: - return candidate_model_name - return base_model_name - - def _use_expert_parallel_alltoall_path(self, batch: Batch) -> bool: - moe_ep_size = int(getattr(self, "_moe_ep_size", 1)) - if moe_ep_size <= 1: - return False - # EP is replica-local and is independent of the retired attention-DP - # lane concept. A full batch on any MoE serving role therefore uses - # the EP communication/accounting path whenever EP>1. - return True - - def _predict_expert_parallel_phase_operator_times( - self, - batch: Batch, - *, - lane_workload: Optional[EPLaneWorkload] = None, - ) -> dict[str, float]: - """Predict exact dispatch and combine collectives for one MoE layer.""" - - if self._moe_ep_size <= 1: - return { - "expert_parallel_alltoall_dispatch": 0.0, - "expert_parallel_alltoall_combine": 0.0, - } - if not self._use_expert_parallel_alltoall_path(batch): - raise ValueError( - "Canonical MoE EP execution requires named all-to-all dispatch " - "and combine phases" - ) - return { - op_name: self._predict_comm_operator( - get_comm_operator(op_name), - batch, - lane_workload=lane_workload, - ) - for op_name in ( - "expert_parallel_alltoall_dispatch", - "expert_parallel_alltoall_combine", - ) - } - - def _get_effective_moe_total_tokens(self, batch: Batch) -> int: - effective_tokens = int( - batch.get_effective_total_tokens_rounded(self._cluster_type) - ) - if effective_tokens < 0: - raise ValueError( - f"effective MoE tokens must be non-negative, got {effective_tokens}" - ) - return effective_tokens - - def _get_moe_pre_routing_token_count(self, batch: Optional[Batch]) -> int: - """Return the source-batch width used by pre-routing MoE models. - - A physical EP lane carries only an assignment subset, so its routed - count cannot identify the source width. Callers that need the - one-feature profiling domain must provide the source batch explicitly. - """ - - if batch is None: - raise ValueError( - "MoE pre-routing token lookup requires the source batch; " - "an EPLaneWorkload cannot supply that width" - ) - return self._get_effective_moe_total_tokens(batch) - - def _get_local_ep_routed_tokens( - self, - batch: Batch, - *, - lane_workload: Optional[EPLaneWorkload] = None, - ) -> int: - source = batch if lane_workload is None else lane_workload - resolved_lane_workload = resolve_ep_lane_workload(source, required=True) - assert resolved_lane_workload is not None - return resolved_lane_workload.routed_token_count - - def _get_moe_tokens_input( - self, batch: Batch, layer_id: int = 0 - ) -> EPLaneWorkload | int: - """ - Unified entry point to get MoE tokens input for grouped GEMM prediction. - - EP lane batches carry the canonical physical descriptor. A regular - non-lane batch may use the scalar pre-routing token path for legacy - one-feature models; load-aware models require an explicit descriptor. - - Args: - batch: The batch being processed - layer_id: The layer ID for which to get token allocation (default 0) - - Returns: - - In load-imbalance mode: ``EPLaneWorkload`` - - In single-token-count profiling mode: pre-routing token count - - Raises: - ValueError: If the selected routing mode is not supported by the active predictor - """ - lane_workload = resolve_ep_lane_workload(batch, required=False) - if lane_workload is not None: - if lane_workload.router_topk != int(self._router_topk): - raise ValueError( - "EPLaneWorkload router_topk does not match predictor topology: " - f"descriptor={lane_workload.router_topk}, predictor={self._router_topk}" - ) - return lane_workload - - load_aware = any( - isinstance(prediction, dict) - and prediction.get("_on_demand_prediction", False) - for prediction in ( - getattr(self, "_predictions", {}).get("moe_shuffling"), - getattr(self, "_predictions", {}).get("moe_grouped_gemm"), - ) - ) - if load_aware: - cluster_type = getattr(self, "_cluster_type", None) - if not isinstance(cluster_type, ClusterType): - raise ValueError( - "load-aware MoE prediction requires an initialized cluster_type" - ) - workload = self._materialize_layer_ep_workload( - batch=batch, - cluster_type=cluster_type, - layer_id=layer_id, - ) - if len(workload.participant_ep_ids) != int(self._moe_ep_size): - raise ValueError( - "materialized EP lane count does not match predictor topology: " - f"descriptors={len(workload.participant_ep_ids)}, " - f"predictor={self._moe_ep_size}" - ) - if int(self._moe_ep_size) != 1: - raise ValueError( - "load-aware regular-batch prediction requires an explicit " - "physical EP lane for EP>1" - ) - return workload.lane(0) - return self._get_effective_moe_total_tokens(batch) - - def _get_gating_time(self, batch: Batch) -> float: - """ - Get total MoE gating network execution time (linear + routing_topk). - - The gating network determines which experts each token should be routed to. - Prediction is based on num_tokens feature from profiling data. - - Returns: - Total gating time (sum of linear and routing_topk times) - """ - return self._get_gating_linear_time(batch) + self._get_gating_routing_topk_time( - batch - ) - - def _get_gating_linear_time(self, batch: Batch) -> float: - """ - Get MoE gating linear layer execution time. - - The gating linear layer computes logits from hidden states (hidden_dim -> num_experts). - """ - if not self._supports_operation("moe_gating_linear"): - raise NotImplementedError( - "MoE gating linear is not supported for cluster type" - ) - model_name = self._select_moe_gating_prediction_model_name( - "moe_gating_linear", - batch, - ) - if model_name not in self._predictions: - raise NotImplementedError( - "MoE gating linear is not supported for cluster type" - ) - effective_tokens = batch.get_effective_total_tokens_rounded(self._cluster_type) - return self._get_prediction_for_features( - model_name, - {"num_tokens": effective_tokens}, - feature_names=("num_tokens",), - ) - - def _get_gating_routing_topk_time(self, batch: Batch) -> float: - """ - Get MoE gating routing topk execution time. - - The routing topk operation selects top-K experts and applies softmax normalization. - """ - if not self._supports_operation("moe_gating_routing_topk"): - raise NotImplementedError( - "MoE gating routing topk is not supported for cluster type" - ) - model_name = self._select_moe_gating_prediction_model_name( - "moe_gating_routing_topk", - batch, - ) - if model_name not in self._predictions: - raise NotImplementedError( - "MoE gating routing topk is not supported for cluster type" - ) - effective_tokens = batch.get_effective_total_tokens_rounded(self._cluster_type) - return self._get_prediction_for_features( - model_name, - {"num_tokens": effective_tokens}, - feature_names=("num_tokens",), - ) - - def _resolve_shuffling_per_expert_tokens( - self, - batch: Batch, - moe_tokens_input: Optional[EPLaneWorkload] = None, - ) -> EPLaneWorkload: - source = batch if moe_tokens_input is None else moe_tokens_input - lane_workload = resolve_ep_lane_workload(source, required=True) - assert lane_workload is not None - return lane_workload - - def _build_moe_load_imbalance_features( - self, - lane_workload: EPLaneWorkload, - *, - batch: Optional[Batch] = None, - ) -> Dict[str, float]: - lane_workload = resolve_ep_lane_workload(lane_workload, required=True) - assert lane_workload is not None - - from frontier.moe_load_imbalance import MoELoadImbalanceInput - - expert_token_counts = [int(v) for v in lane_workload.local_token_counts] - - total_routed_tokens = int(sum(expert_token_counts)) - if lane_workload.router_topk <= 0: - raise ValueError(f"Invalid router_topk={lane_workload.router_topk}") - - source_num_tokens = self._get_moe_pre_routing_token_count(batch) - - load_input = MoELoadImbalanceInput( - num_tokens=source_num_tokens, - num_experts_per_device=lane_workload.local_expert_width, - hidden_dim=int(self._model_config.embedding_dim), - expert_hidden_dim=int(self._model_config.mlp_hidden_dim), - router_topk=int(lane_workload.router_topk), - expert_token_counts=expert_token_counts, - load_distribution="runtime", - ) - features = load_input.to_features_dict() - features.pop("load_distribution", None) - features.pop("seed", None) - missing_features = [ - name - for name in self.MOE_LOAD_IMBALANCE_FEATURES - if name not in features - ] - if missing_features: - raise ValueError( - "MoE load-imbalance feature construction is missing canonical " - f"features: {missing_features}" - ) - unexpected_features = sorted( - set(features) - set(self.MOE_LOAD_IMBALANCE_FEATURES) - ) - if unexpected_features: - raise ValueError( - "MoE load-imbalance feature construction produced unexpected " - f"features: {unexpected_features}" - ) - return { - name: features[name] - for name in self.MOE_LOAD_IMBALANCE_FEATURES - } - - def _get_moe_shuffling_time( - self, - batch: Batch, - moe_tokens_input: Optional[EPLaneWorkload] = None, - ) -> float: - """ - Get MoE token shuffling execution time using trained prediction model. - - Shuffling involves dispatching tokens to assigned experts. When the model is - trained with load-imbalance features, use on-demand prediction driven by - per-expert allocation; otherwise use the legacy num_tokens lookup table. - """ - if not self._supports_operation("moe_shuffling"): - raise NotImplementedError("MoE shuffling is not supported for cluster type") - if "moe_shuffling" not in self._predictions: - raise NotImplementedError("MoE shuffling is not supported for cluster type") - if moe_tokens_input is not None and not isinstance( - moe_tokens_input, EPLaneWorkload - ): - raise TypeError( - "MoE shuffling requires an EPLaneWorkload descriptor when an " - "explicit workload is supplied" - ) - - prediction_cache = self._predictions["moe_shuffling"] - if isinstance(prediction_cache, dict) and prediction_cache.get( - "_on_demand_prediction", False - ): - lane_workload = self._resolve_shuffling_per_expert_tokens( - batch, - moe_tokens_input=moe_tokens_input, - ) - if lane_workload.routed_token_count == 0: - raw_time = 0.0 - else: - features = self._build_moe_load_imbalance_features( - lane_workload, - batch=batch, - ) - raw_time = self._get_on_demand_prediction( - "moe_shuffling", features - ) - else: - lane_workload = resolve_ep_lane_workload(batch, required=False) - if moe_tokens_input is not None: - lane_workload = resolve_ep_lane_workload( - moe_tokens_input, - required=True, - ) - if lane_workload is not None: - if lane_workload.routed_token_count == 0: - return 0.0 - effective_tokens = self._get_moe_pre_routing_token_count(batch) - else: - effective_tokens = batch.get_effective_total_tokens_rounded( - self._cluster_type - ) - raw_time = self._get_prediction_for_features( - "moe_shuffling", - {"num_tokens": effective_tokens}, - feature_names=("num_tokens",), - ) - - return raw_time - - def _get_expert_parallel_communication_time( - self, - batch: Batch, - *, - lane_workload: Optional[EPLaneWorkload] = None, - ) -> float: - """ - Get expert parallel communication time. - - Shared-domain MoE execution (monolithic / prefill / decode) uses - expert-parallel all-reduce when EP is enabled without all-to-all routing. - Post-routing EP batches (e.g. DECODE_FFN) and flattened multi-DP MoE - paths keep the all-to-all communication model. - """ - if self._moe_ep_size <= 1: - return 0.0 - - uses_alltoall = self._use_expert_parallel_alltoall_path(batch) - resolved_lane_workload = None - if uses_alltoall: - resolved_lane_workload = resolve_ep_lane_workload( - batch if lane_workload is None else lane_workload, - required=True, - ) - assert resolved_lane_workload is not None - - if self._cc_backend is not None: - quant_manager = get_quantization_manager() - - if uses_alltoall: - routed_tokens = self._get_local_ep_routed_tokens( - batch, - lane_workload=resolved_lane_workload, - ) - data_size_bytes = self._model_config.embedding_dim * 2 * routed_tokens - data_size_bytes = quant_manager.adjust_tensor_size( - "expert_parallel_communication", data_size_bytes, self._cluster_type - ) - result = self._cc_backend.predict_all_to_all( - data_size_bytes=data_size_bytes, - num_devices=self._moe_ep_size, - cluster_type=self._cluster_type, - comm_domain="EP", - ) - logger.debug( - f"_get_expert_parallel_communication_time: using EP all-to-all, " - f"data_size={data_size_bytes}, num_devices={self._moe_ep_size}, " - f"result={result:.6f} ms" - ) - return result - - effective_tokens = batch.get_effective_total_tokens_rounded(self._cluster_type) - data_size_bytes = self._model_config.embedding_dim * 2 * effective_tokens - data_size_bytes = quant_manager.adjust_tensor_size( - "allreduce", data_size_bytes, self._cluster_type - ) - result = self._cc_backend.predict_allreduce( - data_size_bytes=data_size_bytes, - num_devices=self._moe_ep_size, - cluster_type=self._cluster_type, - comm_domain="EP", - ) - result = self._strip_collective_sim_allreduce_launch_overhead_if_needed( - batch=batch, - predicted_ms=result, - num_devices=self._moe_ep_size, - comm_domain="EP", - ) - logger.debug( - f"_get_expert_parallel_communication_time: using EP all-reduce, " - f"data_size={data_size_bytes}, num_devices={self._moe_ep_size}, " - f"result={result:.6f} ms" - ) - return result - - if self._enable_dummy_mode: - logger.debug( - f"_get_expert_parallel_communication_time: CC Backend not available, " - f"using dummy mode value={self._dummy_execution_time} ms" - ) - return self._dummy_execution_time - - raise RuntimeError( - f"CC Backend is required for expert parallel communication prediction " - f"but was not provided. Either:\n" - f" 1. Configure a CC Backend (e.g., --cc_backend vidur or --cc_backend analytical)\n" - f" 2. Enable dummy mode explicitly (--enable_dummy_mode)\n" - f"Current state: cc_backend=None, enable_dummy_mode={self._enable_dummy_mode}" - ) - - def _is_grouped_gemm_on_demand_mode(self) -> bool: - """ - Check if moe_grouped_gemm predictor is in on-demand (load-imbalance) mode. - - Returns: - True if the model was trained with 14 load-imbalance features (requires Dict input), - False if trained with 1 feature (num_tokens only, accepts int input). - """ - if "moe_grouped_gemm" not in self._predictions: - return False - prediction_cache = self._predictions["moe_grouped_gemm"] - return isinstance(prediction_cache, dict) and prediction_cache.get( - "_on_demand_prediction", False + attn_norm_time=base_time if include_attention else 0.0, + mlp_norm_time=base_time if include_ffn else 0.0, + add_time=add_time, + add_attn_residual_time=add_attn_residual_time, + add_ffn_residual_time=add_ffn_residual_time, + tensor_parallel_communication_time=attn_tp_allreduce_time, + attn_tensor_parallel_allreduce_time=attn_tp_allreduce_time, + moe_tensor_parallel_allreduce_time=ffn_tp_allreduce_time, + pipeline_parallel_communication_time=base_time if include_stage_owned else 0.0, + expert_parallel_communication_time=expert_parallel_comm_time, + moe_gating_time=base_time if is_moe else 0.0, + moe_shuffling_time=( + 0.0 if zero_routed_ep_lane else base_time + ) if is_moe else 0.0, + schedule_time=base_time if include_stage_owned else 0.0, + sampler_e2e_time=base_time if include_stage_owned else 0.0, + prepare_inputs_e2e_time=base_time if include_stage_owned else 0.0, + process_model_outputs_time=base_time if include_stage_owned else 0.0, + ray_comm_time=base_time if include_stage_owned else 0.0, + pp_stage_boundary_handoff_time=pp_stage_boundary_handoff_time, + is_moe=is_moe, + mlp_layer_up_proj_execution_time=( + base_time if include_ffn and not is_moe else 0.0 + ), + mlp_layer_down_proj_execution_time=( + base_time if include_ffn and not is_moe else 0.0 + ), + mlp_layer_act_execution_time=( + base_time if include_ffn and not is_moe else 0.0 + ), + moe_grouped_gemm_time=moe_grouped_gemm_time, + share_expert_up_proj_time=share_expert_time, + share_expert_down_proj_time=share_expert_time, + share_expert_act_time=share_expert_time, + tensor_parallel_allgather_time=ffn_tp_allgather_time, + share_expert_tensor_parallel_allreduce_time=share_expert_tp_allreduce_time, + dp_input_allreduce_time=dp_input_allreduce_time, + dp_output_allreduce_time=dp_output_allreduce_time, + communication_operator_times=communication_operator_times, + mlp_operator_times=mlp_operator_times, + moe_operator_times=moe_operator_times, ) - def _get_grouped_gemm_time( + def __init__( self, - num_tokens_or_allocation, - batch: Optional[Batch] = None, - ) -> float: - """ - Calculate grouped GEMM time using trained prediction model. - - Args: - num_tokens_or_allocation: An ``EPLaneWorkload`` for EP-aware - prediction, or an integer for the legacy - one-feature non-lane path. + predictor_config: BaseExecutionTimePredictorConfig, + replica_config: ReplicaConfig, + replica_scheduler_config: BaseReplicaSchedulerConfig, + metrics_config: MetricsConfig, + model_manager: ExecutionTimePredictionModelManager = None, + cluster_type: ClusterType = None, + training_file_paths: Dict[str, str] = None, + cc_backend: Optional["BaseCCBackend"] = None, + actual_replica_ids: Optional[list] = None, + ) -> None: + self._is_moe = True + self._router_topk = replica_config.router_topk + self._moe_tp_size = replica_config.moe_tensor_parallel_size + self._moe_ep_size = replica_config.moe_expert_parallel_size + self._actual_replica_ids = actual_replica_ids + self._attention_query_cache_hits = 0 + self._attention_query_cache_misses = 0 + self._layer_workload_cache_capacity = self._LAYER_WORKLOAD_CACHE_CAPACITY + self._layer_workload_cache = OrderedDict() - Returns: - Total grouped GEMM execution time - """ - if not self._supports_operation("moe_grouped_gemm"): - raise NotImplementedError( - "MoE grouped_gemm is not supported for cluster type" + # Initialize the canonical distribution selector before parent init so + # profiling paths choose matching gating-runtime metadata. + self._moe_routing_distribution_type = str( + getattr(replica_config, "moe_routing_distribution_type", "balanced") + ).strip().lower() + valid_distribution_types = {"balanced", "random", "skewed", "zipf"} + if self._moe_routing_distribution_type not in valid_distribution_types: + raise ValueError( + "moe_routing_distribution_type must be one of " + f"{sorted(valid_distribution_types)}, got " + f"{self._moe_routing_distribution_type!r}" ) - - if "moe_grouped_gemm" not in self._predictions: - raise NotImplementedError( - "MoE grouped_gemm is not supported for cluster type" + self._moe_routing_seed = getattr(replica_config, "moe_routing_seed", 42) + if type(self._moe_routing_seed) is not int or self._moe_routing_seed < 0: + raise ValueError( + "moe_routing_seed must be an exact non-negative int, " + f"got {self._moe_routing_seed}." ) - - prediction_cache = self._predictions["moe_grouped_gemm"] - - if isinstance(num_tokens_or_allocation, Mapping): - raise TypeError( - "MoE grouped_gemm requires an EPLaneWorkload descriptor; raw " - "expert-token maps are not a predictor workload contract" + self._moe_gating_routing_runtime_path = ( + resolve_moe_gating_routing_runtime_path( + self._moe_routing_distribution_type ) - lane_workload = ( - resolve_ep_lane_workload(num_tokens_or_allocation, required=True) - if isinstance(num_tokens_or_allocation, EPLaneWorkload) - else None ) - # Check if this model uses on-demand prediction (trained with load imbalance features) - if isinstance(prediction_cache, dict) and prediction_cache.get( - "_on_demand_prediction" - ): - # On-demand prediction mode: model was trained with load imbalance features. - # We must provide the full feature set computed from per-expert token distribution. - if lane_workload is None: - raise ValueError( - "moe_grouped_gemm is in load-imbalance (on-demand) mode and " - "requires an EPLaneWorkload descriptor" - ) - - if lane_workload.routed_token_count == 0: - return 0.0 - - features = self._build_moe_load_imbalance_features( - lane_workload, - batch=batch, - ) - return self._get_on_demand_prediction("moe_grouped_gemm", features) - - # Standard cache lookup mode (trained with num_tokens only) - if lane_workload is not None: - if lane_workload.routed_token_count == 0: - return 0.0 - source_num_tokens = self._get_moe_pre_routing_token_count(batch) - raw_time = self._get_prediction_for_features( - "moe_grouped_gemm", - {"num_tokens": source_num_tokens}, - feature_names=("num_tokens",), - ) - return raw_time - - # Backward compatibility: single number of tokens - num_tokens = num_tokens_or_allocation - if isinstance(num_tokens, bool) or not isinstance(num_tokens, (int, float)): - raise TypeError( - "MoE grouped_gemm requires an EPLaneWorkload descriptor or a " - "numeric token count" - ) - if num_tokens <= 0: - return 0.0 - raw_time = self._get_prediction_for_features( - "moe_grouped_gemm", - {"num_tokens": num_tokens}, - feature_names=("num_tokens",), + super().__init__( + predictor_config, + replica_config, + replica_scheduler_config, + metrics_config, + model_manager, + cluster_type, + training_file_paths, + cc_backend, ) - return raw_time - @staticmethod - def _resolve_moe_execution_inputs( - *, - moe_tokens_input: object, - lane_workload: Optional[EPLaneWorkload], - include_moe: bool, - ) -> tuple[object, Optional[EPLaneWorkload]]: - """Resolve one canonical MoE input and its optional physical lane. - - ``moe_tokens_input`` is retained for the legacy scalar one-feature - lookup, while ``lane_workload`` carries the physical routed domain. - A physical call must use one descriptor for both roles; allowing a - scalar or a second descriptor alongside it would let communication and - routed compute describe different workloads. - """ - - if isinstance(moe_tokens_input, Mapping): - raise TypeError( - "moe_tokens_input cannot be a raw expert-token map; provide an " - "EPLaneWorkload descriptor" + # Pre-compute one global routing source. EP ownership is applied later + # by the shared per-layer materializer. + self._monolithic_routing_details = None + self._global_routing_allocations = self._init_global_routing_allocations() + if self._cluster_type == ClusterType.MONOLITHIC and self._model_config.is_moe: + self._monolithic_routing_details = self._build_shared_routing_details() + self._emit_routing_details_snapshot( + ClusterType.MONOLITHIC, + self._monolithic_routing_details, ) - - explicit_lane = ( - resolve_ep_lane_workload(lane_workload, required=True) - if lane_workload is not None - else None - ) - input_lane = ( - resolve_ep_lane_workload(moe_tokens_input, required=True) - if isinstance(moe_tokens_input, EPLaneWorkload) - else None + logger.info( + "[MoE Routing] Initialized global routing allocations: " + "distribution=%s, seed=%s, num_layers=%s", + self._moe_routing_distribution_type, + self._moe_routing_seed, + len(self._global_routing_allocations), ) - if explicit_lane is not None: - if input_lane is not None: - if input_lane != explicit_lane: - raise ValueError( - "moe_tokens_input and lane_workload must refer to the " - "same EPLaneWorkload descriptor" - ) - return explicit_lane, explicit_lane - if moe_tokens_input is not None: - raise TypeError( - "cannot combine a scalar moe_tokens_input with an " - "explicit lane_workload" - ) - return explicit_lane, explicit_lane - - if input_lane is not None: - return input_lane, input_lane - - if include_moe and moe_tokens_input is None: - raise ValueError( - "moe_tokens_input is required when include_moe=True. " - "Provide a scalar token count or an EPLaneWorkload descriptor." - ) - return moe_tokens_input, None - def _predict_attention_layer_time_with_query_cache( self, *, @@ -2193,6 +489,7 @@ def _predict_attention_layer_time_with_query_cache( return result @staticmethod + def _clone_attention_time(value: AttentionTime) -> AttentionTime: """Clone scalar attention timings without recursive object copying.""" @@ -2206,7 +503,6 @@ def _clone_attention_time(value: AttentionTime) -> AttentionTime: operator_times = type(operator_times)(dict(operator_times.op_times)) return replace(value, operator_times=operator_times) - # This is now a private method used internally for MoE-specific logic def _get_execution_time_internal( self, batch: Batch, @@ -2554,80 +850,6 @@ def _get_execution_time_internal( ), ) - def _simulate_routing_per_layer( - self, batches: List[Batch], stage_id: int - ) -> Dict[int, Dict[str, Dict[int, float]]]: - """ - Simulate routing for each layer in the stage. - Returns: {layer_id: {replica_id: {moe_component: time_value}}} - """ - del stage_id - cluster_type = getattr(self, "_cluster_type", None) - if not isinstance(cluster_type, ClusterType): - raise ValueError( - "layer routing prediction requires an initialized cluster_type" - ) - - # Routing materialization is stage-local and follows the canonical - # aggregate-to-lane seam. Predictor consumers receive only physical - # lane descriptors, even when this legacy helper returns one result per - # source replica. - num_layers = self._num_layers_per_pipeline_stage - layer_routing_results = {} - - for layer_id in range(num_layers): - layer_routing_results[layer_id] = {} - - for batch in batches: - replica_id = int(batch.replica_id) - layer_workload = self._materialize_layer_ep_workload( - batch=batch, - cluster_type=cluster_type, - layer_id=layer_id, - ) - lane_workloads = tuple( - layer_workload.lane(ep_id) - for ep_id in layer_workload.participant_ep_ids - ) - if not lane_workloads: - raise ValueError( - "layer routing materialization produced no EP lanes: " - f"replica_id={replica_id}, layer_id={layer_id}" - ) - grouped_gemm_time = max( - self._get_grouped_gemm_time(lane_workload, batch=batch) - for lane_workload in lane_workloads - ) - shuffling_time = max( - self._get_moe_shuffling_time( - batch, - moe_tokens_input=lane_workload, - ) - for lane_workload in lane_workloads - ) - communication_time = max( - self._get_expert_parallel_communication_time( - batch, - lane_workload=lane_workload, - ) - for lane_workload in lane_workloads - ) - layer_routing_results[layer_id][replica_id] = { - "moe_grouped_gemm_time": grouped_gemm_time, - "expert_parallel_communication_time": communication_time, - "moe_gating_time": self._get_gating_time(batch), - "moe_shuffling_time": shuffling_time, - } - - return layer_routing_results - - # Phase 2.5: Removed deprecated get_moe_stage_execution_details() method - # MoE models now use predict_moe_layer_time() and other fine-grained APIs - - # ======================================================================== - # New unified API implementation (Phase 0) - MoE extensions - # ======================================================================== - def predict_moe_lane_phase_times( self, *, @@ -2987,7 +1209,6 @@ def predict_allgather_time( return result raise NotImplementedError("MoE all-gather prediction not implemented") - # return self._dummy_execution_time def predict_alltoall_time( self, @@ -3026,209 +1247,6 @@ def predict_alltoall_time( return result raise NotImplementedError("MoE all-to-all prediction not implemented") - # return self._dummy_execution_time - - def _predict_mtp_moe_lane_phase_aggregate( - self, - *, - predictor, - batch: Batch, - pipeline_stage: int, - cluster_type: ClusterType, - layer_id: int, - num_layers: int, - ) -> tuple[ExecutionTime, tuple[float, float, float, float, float]]: - """Return one shared attention result and the five lane barriers. - - ``predictor`` is explicit because structural MTP may run against a - secondary predictor owned by this parent. The attention probe is kept - at one layer: pipeline and CPU overhead are batch-level terms, while - the returned physical phase barriers are the only values scaled by - ``num_layers`` at the caller. - """ - - if type(num_layers) is not int or num_layers < 1: - raise ValueError(f"num_layers must be a positive integer, got {num_layers!r}") - - attention_execution_time = predictor.predict_stage_execution_time( - batch=batch, - stage_id=pipeline_stage, - cluster_type=cluster_type, - num_layers=1, - layer_id=layer_id, - include_ffn=False, - ) - attention_time_ms = float(attention_execution_time.model_time_ms) - if not math.isfinite(attention_time_ms) or attention_time_ms < 0: - raise ValueError( - "MTP structural attention time must be finite and non-negative, " - f"got {attention_time_ms}" - ) - - workload = predictor._materialize_layer_ep_workload( - batch=batch, - cluster_type=cluster_type, - layer_id=layer_id, - ) - participant_ep_ids = tuple(workload.participant_ep_ids) - if not participant_ep_ids: - raise ValueError("MTP MoE replay produced no EP participants") - - effective_tokens = int( - batch.get_effective_total_tokens_for_compute(cluster_type) - ) - if effective_tokens <= 0: - raise ValueError( - "MTP MoE replay requires positive pre-routing effective tokens, " - f"got {effective_tokens}" - ) - - phase_values: list[list[float]] = [] - for ep_id in participant_ep_ids: - lane_workload = workload.lane(int(ep_id)) - lane_phases = predictor.predict_moe_lane_phase_times( - batch=batch, - lane_workload=lane_workload, - pipeline_stage=pipeline_stage, - cluster_type=cluster_type, - ) - if len(lane_phases) != 5: - raise ValueError( - "MTP MoE lane phase API must return five values, " - f"got ep_id={ep_id}, values={lane_phases!r}" - ) - normalized_phases = [float(value) for value in lane_phases] - if any( - not math.isfinite(value) or value < 0 - for value in normalized_phases - ): - raise ValueError( - "MTP MoE lane phase times must be finite and non-negative, " - f"got ep_id={ep_id}, values={normalized_phases}" - ) - phase_values.append(normalized_phases) - - phase_maxima = tuple( - max(values[index] for values in phase_values) for index in range(5) - ) - return attention_execution_time, phase_maxima - - def _predict_mtp_terminal_row_time_ms( - self, - *, - batch: Batch, - stage_id: int, - cluster_type: ClusterType, - num_layers: int, - layer_id: int, - ) -> float: - """Predict a terminal MTP row with physical EP barriers when required.""" - - model_config = getattr(self, "_model_config", None) - if model_config is None or not bool(getattr(model_config, "is_moe", False)): - return super()._predict_mtp_terminal_row_time_ms( - batch=batch, - stage_id=stage_id, - cluster_type=cluster_type, - num_layers=num_layers, - layer_id=layer_id, - ) - is_moe_layer = getattr(model_config, "is_moe_layer", None) - if not callable(is_moe_layer): - raise ValueError( - "MTP terminal MoE prediction requires model_config.is_moe_layer" - ) - if not bool(is_moe_layer(layer_id)): - return super()._predict_mtp_terminal_row_time_ms( - batch=batch, - stage_id=stage_id, - cluster_type=cluster_type, - num_layers=num_layers, - layer_id=layer_id, - ) - if cluster_type not in (ClusterType.MONOLITHIC, ClusterType.DECODE): - return super()._predict_mtp_terminal_row_time_ms( - batch=batch, - stage_id=stage_id, - cluster_type=cluster_type, - num_layers=num_layers, - layer_id=layer_id, - ) - if int(getattr(self, "_moe_ep_size", 1)) <= 1: - return super()._predict_mtp_terminal_row_time_ms( - batch=batch, - stage_id=stage_id, - cluster_type=cluster_type, - num_layers=num_layers, - layer_id=layer_id, - ) - - attention_execution_time, phase_maxima = ( - self._predict_mtp_moe_lane_phase_aggregate( - predictor=self, - batch=batch, - pipeline_stage=stage_id, - cluster_type=cluster_type, - layer_id=layer_id, - num_layers=num_layers, - ) - ) - attention_time_ms = float(attention_execution_time.total_time * 1e3) - if not math.isfinite(attention_time_ms) or attention_time_ms < 0: - raise ValueError( - "MTP terminal attention time must be finite and non-negative, " - f"got {attention_time_ms}" - ) - return attention_time_ms + sum(phase_maxima) * int(num_layers) - - def _predict_mtp_decoder_layer_time_ms( - self, - *, - predictor, - batch: Batch, - ) -> float: - layer_id = 0 - model_config = getattr(predictor, "_model_config", None) - if model_config is None: - raise ValueError( - "MTP structural decoder prediction requires model_config" - ) - if not bool(getattr(model_config, "is_moe", False)): - return super()._predict_mtp_decoder_layer_time_ms( - predictor=predictor, - batch=batch, - ) - - is_moe_layer = getattr(model_config, "is_moe_layer", None) - if not callable(is_moe_layer): - raise ValueError( - "MTP structural MoE decoder prediction requires " - "model_config.is_moe_layer" - ) - if not bool(is_moe_layer(layer_id)): - return super()._predict_mtp_decoder_layer_time_ms( - predictor=predictor, - batch=batch, - ) - - cluster_type = getattr(predictor, "_cluster_type", None) - if not isinstance(cluster_type, ClusterType): - raise ValueError( - "MTP structural MoE decoder prediction requires a valid cluster_type" - ) - - attention_execution_time, phase_maxima = ( - self._predict_mtp_moe_lane_phase_aggregate( - predictor=predictor, - batch=batch, - pipeline_stage=0, - cluster_type=cluster_type, - layer_id=layer_id, - num_layers=1, - ) - ) - attention_time_ms = float(attention_execution_time.model_time_ms) - return attention_time_ms + sum(phase_maxima) def predict_stage_execution_time( self, diff --git a/frontier/scheduler/replica_scheduler/vllm_v1_decision_log.py b/frontier/scheduler/replica_scheduler/vllm_v1_decision_log.py new file mode 100644 index 00000000..38f0f425 --- /dev/null +++ b/frontier/scheduler/replica_scheduler/vllm_v1_decision_log.py @@ -0,0 +1,44 @@ +"""Optional JSONL log of vLLM V1 scheduling decisions. + +The log is enabled by ``FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH`` and is a +debugging aid only; when the variable is unset every call is a no-op. It lives +in its own module so that the handler is configured exactly once, whichever +scheduler component logs first. +""" + +import json +import logging +import os +from typing import Any, Dict, Optional + + +_FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH = os.environ.get( + "FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH", "" +) +_frontier_vllm_v1_sched_decision_logger: Optional[logging.Logger] = None + +if _FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH: + _frontier_vllm_v1_sched_decision_logger = logging.getLogger( + "frontier.vllm_v1_sched_decision" + ) + _frontier_vllm_v1_sched_decision_logger.setLevel(logging.INFO) + _frontier_vllm_v1_sched_decision_logger.propagate = False + + _decision_log_dir = os.path.dirname(_FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH) + if _decision_log_dir: + os.makedirs(_decision_log_dir, exist_ok=True) + + _decision_handler = logging.FileHandler(_FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH) + _decision_handler.setFormatter(logging.Formatter("%(message)s")) + _frontier_vllm_v1_sched_decision_logger.addHandler(_decision_handler) + + +def _log_frontier_vllm_v1_schedule_decision(event: Dict[str, Any]) -> None: + if _frontier_vllm_v1_sched_decision_logger is None: + return + _frontier_vllm_v1_sched_decision_logger.info(json.dumps(event)) + + +def schedule_decision_logging_enabled() -> bool: + """Return whether the decision log is configured for this process.""" + return _frontier_vllm_v1_sched_decision_logger is not None diff --git a/frontier/scheduler/replica_scheduler/vllm_v1_decode_attn_cohort.py b/frontier/scheduler/replica_scheduler/vllm_v1_decode_attn_cohort.py new file mode 100644 index 00000000..2fa32a7f --- /dev/null +++ b/frontier/scheduler/replica_scheduler/vllm_v1_decode_attn_cohort.py @@ -0,0 +1,219 @@ +"""Decode-attention cohort state for the PD-AF decode-attention role. + +A PD-AF ``DECODE_ATTN`` replica processes its requests in cohorts that move +through the attention and FFN stages together. These methods own the cohort +identity, its stage slots and its phase. +""" + +from typing import Dict, List, Optional + +from frontier.config import global_vars +from frontier.entities.batch import Batch, Request +from frontier.logger import get_cluster_logger +from frontier.types import ClusterType + + +class DecodeAttentionCohort: + """Cohort identity, stage slots and phase for the decode-attention role.""" + + def _get_decode_attn_active_cohort_states(self) -> Dict[int, Dict[str, object]]: + cohort_states = getattr(self, "_decode_attn_active_cohort_states", None) + if cohort_states is None: + cohort_states = {} + self._decode_attn_active_cohort_states = cohort_states + return cohort_states + + def _allocate_decode_attn_cohort_id(self) -> int: + if self._cluster_type != ClusterType.DECODE_ATTN: + raise ValueError( + "DECODE_ATTN cohort IDs can only be allocated by a DECODE_ATTN scheduler" + ) + cohort_id = int(getattr(self, "_decode_attn_next_cohort_id", 0)) + self._decode_attn_next_cohort_id = cohort_id + 1 + return cohort_id + + @staticmethod + + def _validate_decode_attn_wave_stages( + cohort_state: Dict[str, object], + ) -> tuple[set[int], Dict[int, str], Dict[int, int]]: + if type(cohort_state) is not dict: + raise RuntimeError("DECODE_ATTN active cohort state must be an exact dict") + + active_stage_indices = cohort_state.get("active_stage_indices") + if type(active_stage_indices) is not set or not active_stage_indices: + raise RuntimeError( + "DECODE_ATTN active_stage_indices must be a non-empty exact set" + ) + for stage_idx in active_stage_indices: + if type(stage_idx) is not int or stage_idx < 0: + raise RuntimeError( + "DECODE_ATTN active stage index must be an exact non-negative " + f"int, got {stage_idx!r}" + ) + + stage_phases = cohort_state.get("stage_phases") + if type(stage_phases) is not dict: + raise RuntimeError("DECODE_ATTN stage phases must be an exact dict") + for stage_idx, stage_phase in stage_phases.items(): + if type(stage_idx) is not int or stage_idx < 0: + raise RuntimeError( + "DECODE_ATTN stage phase index must be an exact non-negative " + f"int, got {stage_idx!r}" + ) + if type(stage_phase) is not str or stage_phase not in { + "local_attn", + "ffn_inflight", + }: + raise RuntimeError( + "DECODE_ATTN stage phase must be local_attn or ffn_inflight, " + f"got {stage_phase!r}" + ) + if set(stage_phases) != active_stage_indices: + raise RuntimeError( + "DECODE_ATTN stage phase key set must exactly match active stages: " + f"phase_keys={sorted(stage_phases)}, " + f"active={sorted(active_stage_indices)}" + ) + + stage_layers = cohort_state.get("stage_current_layer_ids") + if type(stage_layers) is not dict: + raise RuntimeError("DECODE_ATTN stage layers must be an exact dict") + for stage_idx, stage_layer in stage_layers.items(): + if type(stage_idx) is not int or stage_idx < 0: + raise RuntimeError( + "DECODE_ATTN stage layer index must be an exact non-negative " + f"int, got {stage_idx!r}" + ) + if type(stage_layer) is not int or stage_layer < 0: + raise RuntimeError( + "DECODE_ATTN stage layer must be an exact non-negative int, " + f"got {stage_layer!r}" + ) + if set(stage_layers) != active_stage_indices: + raise RuntimeError( + "DECODE_ATTN stage layer key set must exactly match active stages: " + f"layer_keys={sorted(stage_layers)}, " + f"active={sorted(active_stage_indices)}" + ) + + return active_stage_indices, stage_phases, stage_layers + + def get_decode_attn_active_stage_slots( + self, + *, + phase: str | None = None, + layer_id: int | None = None, + ) -> tuple[int, ...]: + if self._cluster_type != ClusterType.DECODE_ATTN: + return tuple() + if phase is not None and type(phase) is not str: + raise ValueError( + "DECODE_ATTN active-stage phase filter must be an exact str, " + f"got {phase!r}" + ) + if layer_id is not None and ( + type(layer_id) is not int or layer_id < 0 + ): + raise ValueError( + "DECODE_ATTN active-stage layer filter must be an exact " + f"non-negative int, got {layer_id!r}" + ) + + active_stage_slots: set[int] = set() + cohort_states = getattr(self, "_decode_attn_active_cohort_states", {}) + if type(cohort_states) is not dict: + raise RuntimeError( + "DECODE_ATTN active cohort states must be an exact dict" + ) + for cohort_id, state in cohort_states.items(): + if type(cohort_id) is not int or cohort_id < 0: + raise RuntimeError( + "DECODE_ATTN active cohort ID must be an exact non-negative " + f"int, got {cohort_id!r}" + ) + if type(state) is not dict: + raise RuntimeError( + "DECODE_ATTN active cohort state must be an exact dict" + ) + pending_request_ids = state.get("pending_request_ids") + if type(pending_request_ids) is not set: + raise RuntimeError( + "DECODE_ATTN pending request IDs must be an exact set, " + f"got {pending_request_ids!r}" + ) + for request_id in pending_request_ids: + if type(request_id) is not int or request_id < 0: + raise RuntimeError( + "DECODE_ATTN pending request ID must be an exact " + f"non-negative int, got {request_id!r}" + ) + if not pending_request_ids: + continue + + active_indices, stage_phases, stage_layers = ( + self._validate_decode_attn_wave_stages(state) + ) + for stage_idx in active_indices: + stage_layer = stage_layers[stage_idx] + if ( + layer_id is not None + and stage_layer != layer_id + ): + continue + stage_phase = stage_phases[stage_idx] + if phase is not None and stage_phase != phase: + continue + active_stage_slots.add(stage_idx) + + return tuple(sorted(active_stage_slots)) + + def set_decode_attn_cohort_phase_for_batch( + self, + batch: Batch, + *, + phase: str, + layer_id: int | None = None, + ) -> None: + if self._cluster_type != ClusterType.DECODE_ATTN: + return + + cohort_id = batch.decode_attn_cohort_id + if cohort_id is None: + return + + normalized_phase = str(phase) + if normalized_phase not in {"local_attn", "ffn_inflight"}: + raise ValueError( + f"Unsupported DECODE_ATTN cohort phase: {normalized_phase}" + ) + + cohort_state = self._get_decode_attn_active_cohort_states().get( + int(cohort_id) + ) + if cohort_state is None: + return + + stage_idx = batch.afd_stage_idx + if stage_idx is None: + cohort_state["af_phase"] = normalized_phase + if layer_id is not None: + cohort_state["current_layer_id"] = int(layer_id) + return + + active_stage_indices, stage_phases, stage_layers = ( + self._validate_decode_attn_wave_stages(cohort_state) + ) + normalized_stage_idx = int(stage_idx) + if normalized_stage_idx not in active_stage_indices: + raise ValueError( + "DECODE_ATTN cohort stage is not active: " + f"stage={normalized_stage_idx}, active={sorted(active_stage_indices)}" + ) + stage_phases[normalized_stage_idx] = normalized_phase + if layer_id is not None: + stage_layers[normalized_stage_idx] = int(layer_id) + cohort_state["current_layer_id"] = int(layer_id) + + phases = set(stage_phases.values()) + cohort_state["af_phase"] = phases.pop() if len(phases) == 1 else "mixed" diff --git a/frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py b/frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py index dab40b77..a9baefda 100644 --- a/frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py +++ b/frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py @@ -17,97 +17,51 @@ """ from collections import deque -from dataclasses import dataclass, replace -import json -import logging -from math import ceil -import os -import time +from dataclasses import replace from typing import Any, Dict, List, Optional, Sequence, Tuple from frontier.config import global_vars from frontier.attention.gdn.guards import model_has_gdn, validate_gdn_runtime_support from frontier.attention.gdn.state import GatedDeltaNetStateSlotManager -from frontier.entities.batch import ( - Batch, - DecodeCudaGraphMetadata, - Request, - SpecDecodeBatchMetadata, -) -from frontier.kv_cache.base_kv_cache_manager import KVCacheAllocationResult -from frontier.kv_cache.kv_cache_block import KVCacheBlock, KVCacheBlockBinding +from frontier.entities.batch import Batch, Request from frontier.kv_cache.replica_kv_cache_manager import ReplicaKVCacheManager from frontier.logger import get_cluster_logger from frontier.scheduler.replica_scheduler.base_replica_scheduler import ( BaseReplicaScheduler, ) -from frontier.spec_decode import ( - compute_iteration_outcome, - get_planned_draft_tokens, - is_spec_decode_enabled, - method_uses_lookahead_slots, +from frontier.scheduler.replica_scheduler.vllm_v1_decision_log import ( + _log_frontier_vllm_v1_schedule_decision, +) +from frontier.scheduler.replica_scheduler.vllm_v1_decode_attn_cohort import ( + DecodeAttentionCohort, +) +from frontier.scheduler.replica_scheduler.vllm_v1_iteration_policy import ( + IterationSchedulingPolicy, +) +from frontier.scheduler.replica_scheduler.vllm_v1_kv_allocation import KvBlockAllocation +from frontier.scheduler.replica_scheduler.vllm_v1_mtp_wait import ( + TargetEmbeddedMtpWaitPolicy, ) +from frontier.scheduler.replica_scheduler.vllm_v1_prefix_cache import ( + PrefixCacheAdmission, + PrefixCacheLedger, +) +from frontier.scheduler.replica_scheduler.vllm_v1_role_schedules import ( + DisaggregatedRoleScheduling, +) +from frontier.spec_decode import is_spec_decode_enabled, method_uses_lookahead_slots from frontier.types import ClusterType -_REQUEST_PROGRESS_PRESERVING_PREEMPTION_CLUSTER_TYPES = frozenset( - { - ClusterType.DECODE, - ClusterType.DECODE_ATTN, - } -) - -_FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH = os.environ.get( - "FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH", "" -) -_frontier_vllm_v1_sched_decision_logger: Optional[logging.Logger] = None - -if _FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH: - _frontier_vllm_v1_sched_decision_logger = logging.getLogger( - "frontier.vllm_v1_sched_decision" - ) - _frontier_vllm_v1_sched_decision_logger.setLevel(logging.INFO) - _frontier_vllm_v1_sched_decision_logger.propagate = False - - _decision_log_dir = os.path.dirname(_FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH) - if _decision_log_dir: - os.makedirs(_decision_log_dir, exist_ok=True) - - _decision_handler = logging.FileHandler(_FRONTIER_VLLM_V1_SCHED_DECISION_LOG_PATH) - _decision_handler.setFormatter(logging.Formatter("%(message)s")) - _frontier_vllm_v1_sched_decision_logger.addHandler(_decision_handler) - - -def _log_frontier_vllm_v1_schedule_decision(event: Dict[str, Any]) -> None: - if _frontier_vllm_v1_sched_decision_logger is None: - return - _frontier_vllm_v1_sched_decision_logger.info(json.dumps(event)) - - -@dataclass(frozen=True) -class PrefixCacheAdmission: - raw_hit_blocks: tuple[KVCacheBlock, ...] - effective_hit_blocks: tuple[KVCacheBlock, ...] - raw_hit_bindings: tuple[KVCacheBlockBinding, ...] - effective_hit_bindings: tuple[KVCacheBlockBinding, ...] - raw_cached_tokens: int - effective_cached_tokens: int - num_new_tokens: int - full_hit_backoff_applied: bool - - -def _serialize_prefix_cache_binding( - binding: KVCacheBlockBinding, -) -> Dict[str, Any]: - return { - "block_hash": binding.block_hash, - "block_id": int(binding.block_id), - "creator_request_id": str(binding.creator_request_id), - "binding_epoch": int(binding.binding_epoch), - } - - -class VLLMv1EngineReplicaScheduler(BaseReplicaScheduler): +class VLLMv1EngineReplicaScheduler( + IterationSchedulingPolicy, + KvBlockAllocation, + PrefixCacheLedger, + TargetEmbeddedMtpWaitPolicy, + DecodeAttentionCohort, + DisaggregatedRoleScheduling, + BaseReplicaScheduler, +): """ Replica scheduler that simulates vLLM v1 engine admission control. @@ -330,68 +284,6 @@ def _create_batch(self, requests: List[Request], num_tokens: List[int]) -> Batch self._mark_batch_requests_active(batch) return batch - def _refresh_target_embedded_mtp_prefill_boundary_state( - self, batch: Batch, request: Request - ) -> None: - metadata = batch.spec_decode_metadata - if metadata is not None: - for _, batch_request in enumerate(batch.requests): - if batch_request is request: - # The metadata row is authoritative for this batch. A - # positive verify width was already recorded when metadata - # was built; a zero verify width means the request did not - # participate in this spec-decode iteration. - return - if self._cluster_type != ClusterType.MONOLITHIC: - return - if not getattr(request, "spec_decode_enabled", False): - return - if not getattr(request, "spec_method_is_target_embedded_mtp", False): - return - if not getattr(request, "is_prefill_complete", False): - return - spec_decode_config = getattr(self, "_spec_decode_config", None) - if spec_decode_config is None: - raise ValueError("Speculative decoding config is not initialized") - if int(getattr(request, "spec_total_iterations", 0)) == 0: - planned_drafts = get_planned_draft_tokens( - spec_decode_config, - request.remaining_decode_tokens, - iteration_index=0, - request_id=str(request.id), - ) - outcome = compute_iteration_outcome( - spec_decode_config, - request.remaining_decode_tokens, - planned_draft_tokens=planned_drafts, - iteration_index=0, - request_id=str(request.id), - ) - if planned_drafts == 0 and outcome.committed_tokens == 1: - # Some vLLM target-embedded MTP requests emit a real prefill - # commit row: one sampled token, no scheduled drafts. Frontier - # already advanced that token at the prefill boundary, so only - # the trace cursor and spec stats need to catch up here. - request.record_spec_decode_iteration( - verify_tokens=outcome.verify_tokens, - accepted_drafts=outcome.accepted_draft_tokens, - rejected_drafts=outcome.rejected_draft_tokens, - committed_tokens=outcome.committed_tokens, - ) - else: - request.set_spec_next_planned_draft_tokens(planned_drafts) - return - if int(getattr(request, "spec_total_iterations", 0)) != 1: - return - request.set_spec_next_planned_draft_tokens( - get_planned_draft_tokens( - spec_decode_config, - request.remaining_decode_tokens, - iteration_index=request.spec_total_iterations, - request_id=str(request.id), - ) - ) - def _get_active_batch_request_counts(self) -> Dict[int, int]: active_counts = getattr(self, "_active_batch_request_counts", None) if active_counts is None: @@ -416,3859 +308,752 @@ def _release_batch_requests_active(self, batch: Batch) -> None: def _is_request_active_in_batch(self, request: Request) -> bool: return self._get_active_batch_request_counts().get(request.id, 0) > 0 - def _get_decode_attn_active_cohort_states(self) -> Dict[int, Dict[str, object]]: - cohort_states = getattr(self, "_decode_attn_active_cohort_states", None) - if cohort_states is None: - cohort_states = {} - self._decode_attn_active_cohort_states = cohort_states - return cohort_states - - def _allocate_decode_attn_cohort_id(self) -> int: - if self._cluster_type != ClusterType.DECODE_ATTN: - raise ValueError( - "DECODE_ATTN cohort IDs can only be allocated by a DECODE_ATTN scheduler" - ) - cohort_id = int(getattr(self, "_decode_attn_next_cohort_id", 0)) - self._decode_attn_next_cohort_id = cohort_id + 1 - return cohort_id - - @staticmethod - def _validate_decode_attn_wave_stages( - cohort_state: Dict[str, object], - ) -> tuple[set[int], Dict[int, str], Dict[int, int]]: - if type(cohort_state) is not dict: - raise RuntimeError("DECODE_ATTN active cohort state must be an exact dict") - - active_stage_indices = cohort_state.get("active_stage_indices") - if type(active_stage_indices) is not set or not active_stage_indices: - raise RuntimeError( - "DECODE_ATTN active_stage_indices must be a non-empty exact set" - ) - for stage_idx in active_stage_indices: - if type(stage_idx) is not int or stage_idx < 0: - raise RuntimeError( - "DECODE_ATTN active stage index must be an exact non-negative " - f"int, got {stage_idx!r}" + def complete_kv_transfer_for_requests( + self, requests: Sequence[Request] + ) -> None: + for request in requests: + if request.id not in self._pending_kv_transfer_requests: + raise ValueError( + "KV transfer completion for request without pending transfer state: " + f"request_id={request.id}, " + f"source_cluster={self._cluster_type.name}, " + f"source_replica={self._replica_id}, " + f"source_dp={self._replica_local_id}" ) - stage_phases = cohort_state.get("stage_phases") - if type(stage_phases) is not dict: - raise RuntimeError("DECODE_ATTN stage phases must be an exact dict") - for stage_idx, stage_phase in stage_phases.items(): - if type(stage_idx) is not int or stage_idx < 0: - raise RuntimeError( - "DECODE_ATTN stage phase index must be an exact non-negative " - f"int, got {stage_idx!r}" - ) - if type(stage_phase) is not str or stage_phase not in { - "local_attn", - "ffn_inflight", - }: - raise RuntimeError( - "DECODE_ATTN stage phase must be local_attn or ffn_inflight, " - f"got {stage_phase!r}" - ) - if set(stage_phases) != active_stage_indices: - raise RuntimeError( - "DECODE_ATTN stage phase key set must exactly match active stages: " - f"phase_keys={sorted(stage_phases)}, " - f"active={sorted(active_stage_indices)}" - ) + if request.id in self._allocation_map: + self._free_request_resources(request) + self._pending_kv_transfer_requests.discard(request.id) - stage_layers = cohort_state.get("stage_current_layer_ids") - if type(stage_layers) is not dict: - raise RuntimeError("DECODE_ATTN stage layers must be an exact dict") - for stage_idx, stage_layer in stage_layers.items(): - if type(stage_idx) is not int or stage_idx < 0: - raise RuntimeError( - "DECODE_ATTN stage layer index must be an exact non-negative " - f"int, got {stage_idx!r}" - ) - if type(stage_layer) is not int or stage_layer < 0: - raise RuntimeError( - "DECODE_ATTN stage layer must be an exact non-negative int, " - f"got {stage_layer!r}" - ) - if set(stage_layers) != active_stage_indices: - raise RuntimeError( - "DECODE_ATTN stage layer key set must exactly match active stages: " - f"layer_keys={sorted(stage_layers)}, " - f"active={sorted(active_stage_indices)}" - ) + def on_batch_end(self, batch: Batch) -> None: + """ + Handle batch completion - update running requests state. - return active_stage_indices, stage_phases, stage_layers - - def get_decode_attn_active_stage_slots( - self, - *, - phase: str | None = None, - layer_id: int | None = None, - ) -> tuple[int, ...]: - if self._cluster_type != ClusterType.DECODE_ATTN: - return tuple() - if phase is not None and type(phase) is not str: - raise ValueError( - "DECODE_ATTN active-stage phase filter must be an exact str, " - f"got {phase!r}" - ) - if layer_id is not None and ( - type(layer_id) is not int or layer_id < 0 - ): - raise ValueError( - "DECODE_ATTN active-stage layer filter must be an exact " - f"non-negative int, got {layer_id!r}" - ) + For completed requests: free resources and remove from running list. + For ongoing requests: keep in running list for next iteration. - active_stage_slots: set[int] = set() - cohort_states = getattr(self, "_decode_attn_active_cohort_states", {}) - if type(cohort_states) is not dict: - raise RuntimeError( - "DECODE_ATTN active cohort states must be an exact dict" - ) - for cohort_id, state in cohort_states.items(): - if type(cohort_id) is not int or cohort_id < 0: - raise RuntimeError( - "DECODE_ATTN active cohort ID must be an exact non-negative " - f"int, got {cohort_id!r}" - ) - if type(state) is not dict: - raise RuntimeError( - "DECODE_ATTN active cohort state must be an exact dict" + Special handling for PREFILL cluster in disaggregated mode: + - Requests are transferred to DECODE cluster after prefill completion + - Partially-prefilled requests stay in PREFILL running list for next chunk + + Args: + batch: The batch that has completed execution + """ + self._num_running_batches -= 1 + + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + self._release_batch_requests_active(batch) + + for request in batch.requests: + self._refresh_target_embedded_mtp_prefill_boundary_state(batch, request) + if self._cluster_type == ClusterType.DECODE_ATTN: + decode_attn_cohort_id = getattr( + batch, + "decode_attn_cohort_id", + None, ) - pending_request_ids = state.get("pending_request_ids") - if type(pending_request_ids) is not set: - raise RuntimeError( - "DECODE_ATTN pending request IDs must be an exact set, " - f"got {pending_request_ids!r}" + if decode_attn_cohort_id is None: + self._decode_attn_active_request_ids.discard(request.id) + else: + cohort_states = self._get_decode_attn_active_cohort_states() + cohort_state = cohort_states.get(int(decode_attn_cohort_id)) + if cohort_state is None: + self._decode_attn_active_request_ids.discard(request.id) + else: + cohort_state["pending_request_ids"].discard(request.id) + if not cohort_state["pending_request_ids"]: + for cohort_request_id in cohort_state["all_request_ids"]: + self._decode_attn_active_request_ids.discard( + cohort_request_id + ) + cohort_states.pop(int(decode_attn_cohort_id), None) + if ( + getattr(self, "_decode_attn_open_cohort_id", None) + == decode_attn_cohort_id + ): + self._decode_attn_open_cohort_id = None + + if request.completed: + extra_release_iters = ( + self._get_monolithic_pp_extra_terminal_release_iters() ) - for request_id in pending_request_ids: - if type(request_id) is not int or request_id < 0: - raise RuntimeError( - "DECODE_ATTN pending request ID must be an exact " - f"non-negative int, got {request_id!r}" + if extra_release_iters > 0: + pending_release_iters = ( + self._get_monolithic_pp_pending_terminal_release_iters() + ) + pending_release_iters[request.id] = max( + pending_release_iters.get(request.id, 0), + extra_release_iters, + ) + logger.debug( + "[VLLMv1Engine] Request %s completed, deferring free for " + "%s extra MONOLITHIC+PP terminal iteration(s)", + request.id, + extra_release_iters, ) - if not pending_request_ids: - continue - - active_indices, stage_phases, stage_layers = ( - self._validate_decode_attn_wave_stages(state) - ) - for stage_idx in active_indices: - stage_layer = stage_layers[stage_idx] - if ( - layer_id is not None - and stage_layer != layer_id - ): - continue - stage_phase = stage_phases[stage_idx] - if phase is not None and stage_phase != phase: continue - active_stage_slots.add(stage_idx) + # Request finished - free resources and remove from running + self._free_request_resources(request) + self._scheduled_num_computed_tokens_by_request.pop(request.id, None) + if request in self._running_requests: + self._running_requests.remove(request) + logger.debug( + f"[VLLMv1Engine] Request {request.id} completed, " + f"freed resources, running_reqs={len(self._running_requests)}" + ) + elif self._cluster_type == ClusterType.PREFILL: + # PREFILL cluster in disaggregated mode: + # Requests are transferred to DECODE cluster after prefill completion + # MODIFIED: Do NOT free KV cache here - it will be freed when transfer completes + # This matches vLLM v1 behavior (scheduler.py:1480-1501) + # where blocks are freed on finished_sending event - return tuple(sorted(active_stage_slots)) + if request.is_prefill_complete: + # Remove from running list only after prefill is fully complete. + self._scheduled_num_computed_tokens_by_request.pop(request.id, None) + if request in self._running_requests: + self._running_requests.remove(request) - def set_decode_attn_cohort_phase_for_batch( - self, - batch: Batch, - *, - phase: str, - layer_id: int | None = None, - ) -> None: - if self._cluster_type != ClusterType.DECODE_ATTN: - return + # Track that this request's KV cache is pending transfer. + self._pending_kv_transfer_requests.add(request.id) - cohort_id = batch.decode_attn_cohort_id - if cohort_id is None: - return + logger.info( + f"[VLLMv1Engine] Request {request.id} prefill complete, " + f"KV cache retained for transfer (blocks={self._allocation_map.get(request.id, 0)}), " + f"running_reqs={len(self._running_requests)}" + ) + else: + # Partial prefill: keep request in running queue for the next chunk. + logger.debug( + f"[VLLMv1Engine] Request {request.id} partial prefill complete, " + f"processed_tokens={request.num_processed_tokens}, " + f"running_reqs={len(self._running_requests)}" + ) + elif self._cluster_type == ClusterType.DECODE_ATTN: + # DECODE_ATTN in PD-AF mode: + # This method is called ONLY by GlobalBatchEndEvent (decode step complete) + # NOT called for intermediate layers (those go through _af_immediate_batch_queue) - normalized_phase = str(phase) - if normalized_phase not in {"local_attn", "ffn_inflight"}: - raise ValueError( - f"Unsupported DECODE_ATTN cohort phase: {normalized_phase}" - ) + # Note: _num_running_batches already decremented at method start (line 151) + # This is correct - decode step completed, release pipeline slot - cohort_state = self._get_decode_attn_active_cohort_states().get( - int(cohort_id) - ) - if cohort_state is None: - return - - stage_idx = batch.afd_stage_idx - if stage_idx is None: - cohort_state["af_phase"] = normalized_phase - if layer_id is not None: - cohort_state["current_layer_id"] = int(layer_id) - return - - active_stage_indices, stage_phases, stage_layers = ( - self._validate_decode_attn_wave_stages(cohort_state) - ) - normalized_stage_idx = int(stage_idx) - if normalized_stage_idx not in active_stage_indices: - raise ValueError( - "DECODE_ATTN cohort stage is not active: " - f"stage={normalized_stage_idx}, active={sorted(active_stage_indices)}" - ) - stage_phases[normalized_stage_idx] = normalized_phase - if layer_id is not None: - stage_layers[normalized_stage_idx] = int(layer_id) - cohort_state["current_layer_id"] = int(layer_id) - - phases = set(stage_phases.values()) - cohort_state["af_phase"] = phases.pop() if len(phases) == 1 else "mixed" - - def _get_monolithic_pp_waiting_admission_delay_iters(self) -> Dict[int, int]: - delay_iters = getattr( - self, - "_monolithic_pp_waiting_admission_delay_iters", - None, - ) - if delay_iters is None: - delay_iters = {} - self._monolithic_pp_waiting_admission_delay_iters = delay_iters - return delay_iters + # For completed requests: free resources and remove from running list + # For ongoing requests: keep in _running_requests for next decode step + # + # Note: request._completed_layer_count is already reset to 0 by request.on_batch_end() + # This ensures layer-consistent grouping in next _schedule_decode_attn_only() call - def _is_target_embedded_mtp_request(self, request: Request) -> bool: - return ( - bool(getattr(request, "spec_decode_enabled", False)) - and bool(getattr(request, "spec_method_is_target_embedded_mtp", False)) - ) + if request.completed: + # Request finished all decode tokens - free resources and remove + self._free_request_resources(request) + if request in self._running_requests: + self._running_requests.remove(request) + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Request {request.id} completed, " + f"freed resources, running_reqs={len(self._running_requests)}" + ) + else: + # Request continues with next decode token - keep in _running_requests + # Phase 1 of _schedule_decode_attn_only() will pick this up + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Request {request.id} continues to next decode step, " + f"processed_tokens={request.num_processed_tokens}" + ) + else: + # Request continues - keep in running list + # (will be scheduled again in next iteration) + if self._should_apply_monolithic_pp_mtp_output_wait(request): + output_wait_iters = ( + self._get_monolithic_pp_mtp_output_wait_iters_for_request( + request + ) + ) + if output_wait_iters > 0: + self._add_monolithic_pp_mtp_output_wait( + request.id, + wait_iters=output_wait_iters, + ) + logger.debug( + f"[VLLMv1Engine] Request {request.id} continues, " + f"processed_tokens={request.num_processed_tokens}" + ) - def _has_future_planned_draft_tokens(self, request: Request) -> bool: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, - ) - per_request_trace = getattr( - spec_config, - "_per_request_scheduled_draft_tokens_trace", - None, - ) - if per_request_trace is None: - return False - request_id = str(request.id) - if request_id not in per_request_trace: - raise ValueError( - "per-request scheduled draft trace missing request_id=" - f"{request_id!r}" + def _schedule_running_requests( + self, token_budget: int, preempted_requests: List[Request] + ) -> Tuple[int, List[Request], List[int]]: + """ + Phase 1: Schedule requests currently in RUNNING state. + + Iterate through running requests and try to allocate memory for + their next tokens. May trigger preemption if memory is insufficient. + + Args: + token_budget: Remaining token budget for this iteration + preempted_requests: List to track preempted requests + + Returns: + Tuple of (remaining_budget, scheduled_requests, num_tokens_list) + """ + logger = get_cluster_logger(__name__, self._cluster_type.name) + scheduled = [] + num_tokens_list = [] + waiting_final_prefill_count = ( + self._count_final_fast_lane_requests( + self._preempted_requests + self._request_queue, + final_predicate=self._is_final_prefill_fast_lane_request, ) - next_iteration = int(getattr(request, "spec_total_iterations", 0)) + 1 - return any(int(tokens) > 0 for tokens in per_request_trace[request_id][next_iteration:]) - - def _get_target_embedded_mtp_terminal_overshoot_rows( - self, - request: Request, - *, - start_iteration_index: int, - ) -> List[Tuple[int, int, int, int, int]]: - if self._cluster_type != ClusterType.MONOLITHIC: - return [] - if self._num_stages <= 1: - return [] - if not getattr(request, "spec_decode_enabled", False): - return [] - if not getattr(request, "spec_method_is_target_embedded_mtp", False): - return [] - - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - raise ValueError("Speculative decoding config is not initialized") - per_request_planned_trace = getattr( - spec_config, - "_per_request_scheduled_draft_tokens_trace", - None, - ) - per_request_committed_trace = getattr( - spec_config, - "_per_request_committed_tokens_trace", - None, + if self._cluster_type == ClusterType.PREFILL + else 0 ) - if per_request_planned_trace is None and per_request_committed_trace is None: - return [] - if per_request_planned_trace is None or per_request_committed_trace is None: - raise ValueError( - "terminal target-embedded MTP overshoot modeling requires both " - "per-request scheduled draft and committed token traces" - ) - request_id = str(request.id) - if request_id not in per_request_planned_trace: - raise ValueError( - "per-request scheduled draft trace missing request_id=" - f"{request_id!r}" - ) - if request_id not in per_request_committed_trace: - raise ValueError( - "per-request acceptance trace missing request_id=" - f"{request_id!r}" + self._current_iteration_token_budget = token_budget + req_index = 0 + while req_index < len(self._running_requests) and token_budget > 0: + self._current_iteration_token_budget = token_budget + request = self._running_requests[req_index] + is_final_prefill_running_request = ( + self._cluster_type == ClusterType.PREFILL + and self._is_final_prefill_fast_lane_request(request) ) - - planned_trace = per_request_planned_trace[request_id] - committed_trace = per_request_committed_trace[request_id] - if len(planned_trace) != len(committed_trace): - raise ValueError( - "per-request MTP trace length mismatch: " - f"request_id={request_id!r}, planned_len={len(planned_trace)}, " - f"committed_len={len(committed_trace)}" + is_hidden_prefill_running_request = ( + self._cluster_type == ClusterType.PREFILL + and not request.is_prefill_complete + and not is_final_prefill_running_request ) - start_idx = int(start_iteration_index) - if start_idx < 0: - raise ValueError( - f"start_iteration_index must be >= 0, got={start_idx}" + if ( + request.id + in self._get_monolithic_pp_pending_terminal_release_iters() + ): + req_index += 1 + continue + + continuation_request_ids = getattr( + self, "_continuation_request_ids", set() ) - if start_idx >= len(planned_trace): - return [] - - terminal_rows: List[Tuple[int, int, int, int, int]] = [] - for idx in range(start_idx, len(planned_trace)): - planned_drafts = int(planned_trace[idx]) - raw_committed = int(committed_trace[idx]) - if planned_drafts < 0: - raise ValueError( - "terminal scheduled draft tokens must be >= 0, " - f"request_id={request_id!r}, iteration_index={idx}, " - f"got={planned_drafts}" - ) - if raw_committed < 0: - raise ValueError( - "terminal committed tokens must be >= 0, " - f"request_id={request_id!r}, iteration_index={idx}, " - f"got={raw_committed}" - ) - trace_verify_tokens = 1 + planned_drafts - if raw_committed > trace_verify_tokens: - raise ValueError( - "terminal committed tokens cannot exceed verify window: " - f"request_id={request_id!r}, iteration_index={idx}, " - f"committed={raw_committed}, " - f"verify_tokens={trace_verify_tokens}" + if request.id in continuation_request_ids: + logger.debug( + "[VLLMv1Engine] Phase 1: skipping req=%s " + "(already scheduled in current cycle)", + request.id, ) - if planned_drafts == 0 and raw_committed == 0: + req_index += 1 continue - # Once the logical response is complete, vLLM's online scheduler no - # longer replays the full forced-acceptance scheduled-draft window - # for request latency. The diagnostic acceptance trace can still - # contain scheduled_draft_tokens=32 for the next trace row, while - # the clean scheduler batch log exposes only a one-token cleanup row - # for that completed request. Model that terminal cleanup as one - # target token and keep the raw committed count only for audit - # provenance. Replaying the full trace window here over-extends - # short decode-tail request latency and violates the clean online - # metric scope. - terminal_cleanup_verify_tokens = 1 - terminal_rows.append( - ( - 0, - terminal_cleanup_verify_tokens, - 0, - 0, - raw_committed, + if ( + self._cluster_type == ClusterType.MONOLITHIC + and self._num_stages > 1 + and request.id + in self._get_monolithic_pp_mtp_output_wait_request_ids() + ): + logger.debug( + "[VLLMv1Engine][MONOLITHIC] Phase 1: delaying req=%s " + "for one PP output-visible MTP scheduler step", + request.id, ) - ) - return terminal_rows + req_index += 1 + continue - def _should_delay_monolithic_pp_waiting_admission_on_add( - self, request: Request - ) -> bool: - if self._cluster_type != ClusterType.MONOLITHIC: - return False - if self._num_stages <= 1: - return False - if not self._is_target_embedded_mtp_request(request): - return False - planned_drafts = int(getattr(request, "spec_next_planned_draft_tokens", 0)) - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, + active_in_pp_batch = ( + self._cluster_type in {ClusterType.MONOLITHIC, ClusterType.DECODE} + and self._num_stages > 1 + and self._is_request_active_in_batch(request) ) - num_speculative_tokens = int( - getattr(spec_config, "num_speculative_tokens", 0) - ) - if planned_drafts <= 0: - # A zero first scheduled-draft step still has PP lookahead admission - # visibility cost when a long decode will enter later MTP draft steps. - if num_speculative_tokens > 2: - return False - if int(getattr(request, "num_decode_tokens", 0)) < 128: - return False - if not self._has_future_planned_draft_tokens(request): - return False - else: - block_size = int(getattr(self._config, "block_size", 16)) - if num_speculative_tokens >= block_size: - # Full-block-or-wider target-embedded MTP request admission is - # already protected by output-visible guards after it starts - # running. Adding a pre-admission PP boundary here makes late - # online arrivals miss the vLLM-visible scheduler slot and - # under-batches wide verify traces relative to vLLM. - return False + if active_in_pp_batch: + if self._cluster_type == ClusterType.MONOLITHIC: + reserved_tokens = ( + self._get_monolithic_pp_mtp_visible_budget_reservation_tokens( + request, + token_budget, + ) + ) + if reserved_tokens > 0: + token_budget -= reserved_tokens + self._current_iteration_token_budget = token_budget + logger.debug( + "[VLLMv1Engine][MONOLITHIC] Phase 1: reserving " + "%s token(s) for active output-visible MTP req=%s", + reserved_tokens, + request.id, + ) + logger.debug( + "[VLLMv1Engine][%s] Phase 1: skipping req=%s " + "(already active in a PP batch)", + self._cluster_type.name, + request.id, + ) + req_index += 1 + continue + if ( - num_speculative_tokens > 2 - and int(getattr(request, "num_prefill_tokens", 0)) - >= int(self._max_num_scheduled_tokens) - + max(1, int(self._max_num_scheduled_tokens) // 2) + self._cluster_type in {ClusterType.MONOLITHIC, ClusterType.DECODE} + and self._num_stages > 1 + and getattr(request, "completed_layer_count", 0) != 0 ): - # Long multi-chunk prefills have enough remaining prefill work - # to expose the next PP boundary through the prefill chunks - # themselves. Adding a separate half-block MTP admission delay - # over-queues these arrivals and inflates p90 TTFT; keep the - # delay for shorter prefills where r106 showed global - # half-block skipping over-corrects TPOT tails. - return False - return ( - self._num_running_batches > 0 - or bool(self._running_requests) - or bool(self._get_active_batch_request_counts()) - ) + logger.debug( + "[VLLMv1Engine][%s] Phase 1: skipping in-flight req=%s " + "with layer_count=%s (PP continuation still active)", + self._cluster_type.name, + request.id, + getattr(request, "completed_layer_count", None), + ) + req_index += 1 + continue - def _add_monolithic_pp_waiting_admission_delay( - self, request_id: int, *, wait_iters: Optional[int] = None - ) -> None: - resolved_wait_iters = int( - wait_iters if wait_iters is not None else max(1, self._num_stages - 1) - ) - if resolved_wait_iters <= 0: - raise ValueError( - f"wait_iters must be positive, got={resolved_wait_iters}" - ) - delay_iters = self._get_monolithic_pp_waiting_admission_delay_iters() - delay_iters[request_id] = max( - int(delay_iters.get(request_id, 0)), - resolved_wait_iters, - ) + # Calculate number of new tokens to process + num_new_tokens = self._get_request_next_num_tokens(request) - def _should_defer_monolithic_pp_waiting_admission( - self, request: Request - ) -> bool: - delay_iters = self._get_monolithic_pp_waiting_admission_delay_iters() - remaining = int(delay_iters.get(request.id, 0)) - if remaining <= 0: - delay_iters.pop(request.id, None) - return False - if self._cluster_type != ClusterType.MONOLITHIC or self._num_stages <= 1: - delay_iters.pop(request.id, None) - return False - if not self._is_target_embedded_mtp_request(request): - delay_iters.pop(request.id, None) - return False - if ( - not self._running_requests - and self._num_running_batches <= 0 - and not self._get_active_batch_request_counts() - ): - # No active PP work remains to provide a future output-visible - # scheduler boundary; fail open to avoid deadlocking the queue. - delay_iters.pop(request.id, None) - return False - remaining -= 1 - if remaining > 0: - delay_iters[request.id] = remaining - else: - delay_iters.pop(request.id, None) - return True - - def _get_monolithic_pp_mtp_near_full_prefill_request_ids(self) -> set[int]: - request_ids = getattr( - self, - "_monolithic_pp_mtp_near_full_prefill_request_ids", - None, - ) - if request_ids is None: - request_ids = set() - self._monolithic_pp_mtp_near_full_prefill_request_ids = request_ids - return request_ids - - def _get_monolithic_pp_mtp_single_output_wait_request_ids(self) -> set[int]: - request_ids = getattr( - self, - "_monolithic_pp_mtp_single_output_wait_request_ids", - None, - ) - if request_ids is None: - request_ids = set() - self._monolithic_pp_mtp_single_output_wait_request_ids = request_ids - return request_ids - - def _get_target_embedded_mtp_request_acceptance_ratio( - self, request: Request - ) -> Optional[float]: - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, - ) - if spec_config is None: - return None - committed_trace_map = getattr( - spec_config, - "_per_request_committed_tokens_trace", - None, - ) - scheduled_trace_map = getattr( - spec_config, - "_per_request_scheduled_draft_tokens_trace", - None, - ) - if committed_trace_map is None and scheduled_trace_map is None: - return None - if committed_trace_map is None or scheduled_trace_map is None: - raise ValueError( - "MTP request acceptance audit requires both per-request " - "committed and scheduled-draft traces" - ) - request_id = str(request.id) - if request_id not in committed_trace_map: - raise ValueError( - "per-request acceptance trace missing request_id=" - f"{request_id!r}" - ) - if request_id not in scheduled_trace_map: - raise ValueError( - "per-request scheduled draft trace missing request_id=" - f"{request_id!r}" + # Apply max_model_len limit + scheduler_num_computed_tokens = self._get_scheduler_num_computed_tokens( + request ) - committed_trace = committed_trace_map[request_id] - scheduled_trace = scheduled_trace_map[request_id] - if len(committed_trace) != len(scheduled_trace): - raise ValueError( - "MTP request acceptance audit trace length mismatch: " - f"request_id={request_id!r}, " - f"committed_len={len(committed_trace)}, " - f"scheduled_len={len(scheduled_trace)}" + max_allowed = self._max_model_len - scheduler_num_computed_tokens + num_new_tokens = min(num_new_tokens, max_allowed) + num_new_tokens = self._apply_long_prefill_token_threshold( + request, num_new_tokens ) - accepted_drafts = 0 - scheduled_drafts = 0 - for committed_tokens, planned_drafts in zip( - committed_trace, - scheduled_trace, - ): - planned = int(planned_drafts) - if planned <= 0: + + # Apply token budget limit + effective_token_budget = token_budget + if ( + is_hidden_prefill_running_request + and waiting_final_prefill_count > 0 + and self._prefill_iteration_reserved_tokens_remaining > 0 + ): + effective_token_budget = max( + token_budget + - min( + self._prefill_iteration_reserved_tokens_remaining, + token_budget, + ), + 0, + ) + if effective_token_budget <= 0: + req_index += 1 + continue + num_new_tokens = min(num_new_tokens, effective_token_budget) + + if num_new_tokens <= 0: + req_index += 1 continue - committed = int(committed_tokens) - accepted_drafts += min(max(committed - 1, 0), planned) - scheduled_drafts += planned - if scheduled_drafts <= 0: - return None - return accepted_drafts / scheduled_drafts - - def _get_monolithic_pp_mtp_output_wait_prefill_threshold(self) -> int: - max_scheduled_tokens = int(self._max_num_scheduled_tokens) - block_size = int(getattr(self._config, "block_size", 16)) - headroom_tokens = 4 * block_size - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, + + # Try to allocate with preemption + preempted_count_before = len(preempted_requests) + can_schedule = self._try_allocate_with_preemption( + request, + num_new_tokens, + preempted_requests, + scheduler_num_computed_tokens=scheduler_num_computed_tokens, ) - num_speculative_tokens = int( - getattr(spec_config, "num_speculative_tokens", 0) - ) - if num_speculative_tokens >= block_size: - # Wide target-embedded MTP carries a larger PP lookahead payload than - # the narrow-window cases that established the original four-block - # threshold. Reserve window-proportional headroom so long chunked - # prefill slices enter the same output-visible continuation lane - # instead of being treated as ordinary prefill chunks. - headroom_tokens = max(headroom_tokens, 8 * num_speculative_tokens) - return max(1, max_scheduled_tokens - headroom_tokens) - - def _get_monolithic_pp_mtp_output_wait_iters(self) -> int: - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, + token_budget = self._rollback_current_iteration_preempted_requests( + scheduled_requests=scheduled, + scheduled_num_tokens=num_tokens_list, + newly_preempted_requests=preempted_requests[ + preempted_count_before: + ], + token_budget=token_budget, ) - if spec_config is None: - return 2 - block_size = int(getattr(self._config, "block_size", 16)) - num_speculative_tokens = int( - getattr(spec_config, "num_speculative_tokens", 0) - ) - if num_speculative_tokens < block_size: - return 2 + if can_schedule: + self._advance_scheduler_num_computed_tokens(request, num_new_tokens) + scheduled.append(request) + num_tokens_list.append(num_new_tokens) + token_budget -= num_new_tokens + self._current_iteration_token_budget = token_budget + if is_final_prefill_running_request: + self._prefill_iteration_reserved_tokens_remaining = max( + self._prefill_iteration_reserved_tokens_remaining + - num_new_tokens, + 0, + ) + req_index += 1 - committed_trace_map = getattr( - spec_config, - "_per_request_committed_tokens_trace", - None, - ) - scheduled_trace_map = getattr( - spec_config, - "_per_request_scheduled_draft_tokens_trace", - None, - ) - accepted_drafts = 0 - scheduled_drafts = 0 - if committed_trace_map is not None or scheduled_trace_map is not None: - if committed_trace_map is None or scheduled_trace_map is None: - raise ValueError( - "MTP output-wait acceptance audit requires both " - "per-request committed and scheduled-draft traces" + # Flow validation: log RUNNING request scheduled + logger.info( + f"[RUNNING_SCHEDULED] req={request.id}, " + f"num_new_tokens={num_new_tokens}, " + f"blocks_allocated={self._allocation_map.get(request.id, 0)}" ) - if set(committed_trace_map.keys()) != set(scheduled_trace_map.keys()): - raise ValueError( - "MTP output-wait acceptance audit trace keys mismatch" + available_blocks_running = int( + self._config.num_blocks - self._num_allocated_blocks ) - for request_id, committed_trace in committed_trace_map.items(): - scheduled_trace = scheduled_trace_map[request_id] - if len(committed_trace) != len(scheduled_trace): - raise ValueError( - "MTP output-wait acceptance audit trace length mismatch: " - f"request_id={request_id!r}, " - f"committed_len={len(committed_trace)}, " - f"scheduled_len={len(scheduled_trace)}" - ) - for committed_tokens, planned_drafts in zip( - committed_trace, - scheduled_trace, - ): - planned = int(planned_drafts) - if planned <= 0: - continue - committed = int(committed_tokens) - accepted_drafts += min(max(committed - 1, 0), planned) - scheduled_drafts += planned - else: - committed_trace = getattr( - spec_config, - "_committed_tokens_trace", - None, - ) - scheduled_trace = getattr( - spec_config, - "_scheduled_draft_tokens_trace", - None, - ) - if committed_trace is None or scheduled_trace is None: - return 2 - if len(committed_trace) != len(scheduled_trace): - raise ValueError( - "MTP output-wait acceptance audit global trace length mismatch: " - f"committed_len={len(committed_trace)}, " - f"scheduled_len={len(scheduled_trace)}" + self._emit_schedule_decision_event( + event="decision", + decision_result="RUNNING_SCHEDULED", + request_id=request.id, + token_budget=token_budget, + available_blocks=available_blocks_running, + num_tokens=num_new_tokens, ) - for committed_tokens, planned_drafts in zip( - committed_trace, - scheduled_trace, - ): - planned = int(planned_drafts) - if planned <= 0: - continue - committed = int(committed_tokens) - accepted_drafts += min(max(committed - 1, 0), planned) - scheduled_drafts += planned - - if scheduled_drafts <= 0: - return 2 - acceptance_ratio = accepted_drafts / scheduled_drafts - if acceptance_ratio >= 0.5: - if self._has_monolithic_pp_visible_waiting_requests(): - # High-acceptance wide MTP should still preserve the extra - # output-visible boundary while fresh prefill admissions are - # visible. Otherwise decode continuation can consume the - # online token budget before late-arriving prefills enter, - # inflating TTFT tails. Once no waiting prefill/resume request - # is visible, shorten the wait to protect short decode tails. - return 2 - # High-acceptance wide target-embedded MTP has fewer decode - # scheduler turns per request. A single output-visible turn is - # enough to expose PP continuation without repeatedly holding - # short tail requests behind terminal trace rows. - return 1 - return 2 - - def _get_monolithic_pp_mtp_output_wait_iters_for_request( - self, request: Request - ) -> int: - wait_iters = self._get_monolithic_pp_mtp_output_wait_iters() - if request.id in self._get_monolithic_pp_mtp_single_output_wait_request_ids(): - return min(wait_iters, 1) - if self._should_apply_monolithic_pp_mtp_fractional_extra_output_wait( - request - ): - counts = self._get_monolithic_pp_mtp_fractional_output_wait_counts() - previous_count = int(counts.get(request.id, 0)) - counts[request.id] = previous_count + 1 - if previous_count == 0: - return 0 - if previous_count % 2 == 1: - return wait_iters - return min(wait_iters, 1) - if self._should_extend_monolithic_pp_mtp_long_decode_output_wait(request): - return wait_iters + 1 - return wait_iters - - def _is_monolithic_pp_mtp_half_block_low_acceptance_request( - self, - request: Request, - *, - acceptance_ratio_limit: float, - ) -> bool: - block_size = int(getattr(self._config, "block_size", 16)) - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, - ) - num_speculative_tokens = int( - getattr(spec_config, "num_speculative_tokens", 0) - ) - if num_speculative_tokens < max(1, block_size // 2): - return False - if num_speculative_tokens >= block_size: - return False + else: + # Request was preempted, stop processing running requests + break - acceptance_ratio = self._get_target_embedded_mtp_request_acceptance_ratio( - request - ) - if acceptance_ratio is None: - return False - return acceptance_ratio < acceptance_ratio_limit - - def _is_monolithic_pp_mtp_mid_prefill_request( - self, - request: Request, - ) -> bool: - block_size = int(getattr(self._config, "block_size", 16)) - num_prefill_tokens = int(getattr(request, "num_prefill_tokens", 0)) - max_scheduled_tokens = int(self._max_num_scheduled_tokens) - if num_prefill_tokens < max(1, max_scheduled_tokens - 4 * block_size): - return False - if num_prefill_tokens > ( - max_scheduled_tokens + max(1, max_scheduled_tokens // 2) - ): - return False - return True + return token_budget, scheduled, num_tokens_list - def _should_apply_monolithic_pp_mtp_fractional_extra_output_wait( - self, - request: Request, - ) -> bool: - if not self._is_target_embedded_mtp_request(request): - return False - if not self._is_monolithic_pp_mtp_mid_prefill_request(request): - return False + def _get_sorted_waiting_queue(self) -> List[Request]: + """ + Get waiting requests sorted by scheduling policy. - block_size = int(getattr(self._config, "block_size", 16)) - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, - ) - num_speculative_tokens = int( - getattr(spec_config, "num_speculative_tokens", 0) - ) - if num_speculative_tokens < max(1, block_size // 2): - return False - if num_speculative_tokens >= block_size: - return False + FCFS: Original queue order (first arrived first). + Priority: Sorted by (priority, arrival_time) ascending. + Thinking-round priority: Final-round requests first, then by + existing policy within each tier. - acceptance_ratio = self._get_target_embedded_mtp_request_acceptance_ratio( - request - ) - if acceptance_ratio is None: - return False + Returns: + List of requests in scheduling order + """ + # Combine main queue and preempted requests + # Preempted requests should be prioritized (at front of queue) + combined = self._preempted_requests + self._request_queue - block_size = int(getattr(self._config, "block_size", 16)) - num_decode_tokens = int(getattr(request, "num_decode_tokens", 0)) - if num_decode_tokens < 16 * block_size: - return False - if num_decode_tokens < 24 * block_size: - if acceptance_ratio >= 0.5: - return False - else: - if acceptance_ratio >= 0.55: - return False - - # The v8/a0.3 request-level RCA shows that medium-length decode tails - # and long decode tails need about one and a half PP output-visible - # wait turns after the first visible decode result. Skip the first - # wait so TTFT remains a prefill/first-token metric, then alternate - # the regular low-acceptance wait with a single-turn wait. This keeps - # the correction as a scheduler visibility family rather than an - # op-runtime calibration scale. - return True - - def _should_extend_monolithic_pp_mtp_long_decode_output_wait( - self, request: Request - ) -> bool: - if not self._is_monolithic_pp_mtp_half_block_low_acceptance_request( - request, - acceptance_ratio_limit=0.25, + if getattr( + getattr(self, "_config", None), "enable_thinking_round_priority", False ): - return False - if not self._is_monolithic_pp_mtp_mid_prefill_request(request): - return False + # Final-round requests first, then by priority, then FIFO + return sorted( + combined, + key=lambda r: ( + 0 if r.is_final_thinking_round else 1, + r.priority, + r.arrived_at, + ), + ) + elif self._scheduling_policy == "priority": + # Sort by priority (ascending) then arrival time (ascending) + return sorted(combined, key=lambda r: (r.priority, r.arrived_at)) + else: + # FCFS: maintain insertion order (preempted first) + return combined - block_size = int(getattr(self._config, "block_size", 16)) - num_decode_tokens = int(getattr(request, "num_decode_tokens", 0)) - if num_decode_tokens < 24 * block_size: - return False - - # Low-acceptance half-block MTP keeps many long decode continuations - # alive after a near-full prefill admission. vLLM exposes an additional - # PP-visible output boundary for the very long decode tail; model that - # boundary as one extra scheduler wait turn instead of - # hiding the residual in compute calibration scale. - return True - - def _record_monolithic_pp_mtp_near_full_prefill_slices(self, batch: Batch) -> None: - if self._cluster_type != ClusterType.MONOLITHIC: - return - if self._num_stages <= 1: - return - near_full_prefill_threshold = ( - self._get_monolithic_pp_mtp_output_wait_prefill_threshold() - ) - for request, num_tokens in zip(batch.requests, batch.num_tokens): - if getattr(request, "is_prefill_complete", False): - continue - if not getattr(request, "spec_decode_enabled", False): - continue - if not getattr(request, "spec_method_is_target_embedded_mtp", False): - continue - if int(num_tokens) < near_full_prefill_threshold: - if not self._should_record_monolithic_pp_mtp_subthreshold_single_wait_prefill( - request, - num_tokens=int(num_tokens), - near_full_prefill_threshold=near_full_prefill_threshold, - ): - if not self._should_record_monolithic_pp_mtp_subthreshold_long_decode_prefill( - request, - num_tokens=int(num_tokens), - near_full_prefill_threshold=near_full_prefill_threshold, - ): - continue - else: - self._get_monolithic_pp_mtp_single_output_wait_request_ids().add( - request.id - ) - self._get_monolithic_pp_mtp_near_full_prefill_request_ids().add( - request.id - ) + def _set_waiting_queues_from_ordered_requests( + self, ordered_requests: List[Request] + ) -> None: + """Rebuild waiting queues from ordered requests. - def _should_record_monolithic_pp_mtp_subthreshold_single_wait_prefill( - self, - request: Request, - *, - num_tokens: int, - near_full_prefill_threshold: int, - ) -> bool: - block_size = int(getattr(self._config, "block_size", 16)) - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, - ) - num_speculative_tokens = int( - getattr(spec_config, "num_speculative_tokens", 0) - ) - if num_speculative_tokens >= block_size: - return False + Requests with `_preempted=True` stay in `_preempted_requests` to keep + preemption recovery semantics and queue priority. + """ + self._preempted_requests = [] + self._request_queue = [] + for request in ordered_requests: + if getattr(request, "_preempted", False): + self._preempted_requests.append(request) + else: + self._request_queue.append(request) - max_scheduled_tokens = int(self._max_num_scheduled_tokens) - min_subthreshold_tokens = max( - 1, - max_scheduled_tokens - 8 * block_size, - ) - if int(num_tokens) < min_subthreshold_tokens: - return False - if int(num_tokens) >= int(near_full_prefill_threshold): - return False - if int(getattr(request, "num_prefill_tokens", 0)) < ( - max_scheduled_tokens + max(1, max_scheduled_tokens // 2) - ): - return False - if int(getattr(request, "num_decode_tokens", 0)) > 4 * block_size: - return False + def _schedule_waiting_requests( + self, token_budget: int + ) -> Tuple[int, List[Request], List[int]]: + """ + Phase 2: Schedule requests in WAITING state. - acceptance_ratio = self._get_target_embedded_mtp_request_acceptance_ratio( - request - ) - if acceptance_ratio is None: - return False - if acceptance_ratio >= 0.5: - return False + Only called when no preemption occurred in Phase 1. + Attempts to admit new requests from the waiting queue. - # Low-acceptance narrow MTP preserves more decode turns than the - # high-acceptance case, so a two-turn output wait over-delays the - # terminal request tail. A single PP-visible wait models the missing - # output boundary without absorbing the residual into compute scale. - return True - - def _should_record_monolithic_pp_mtp_subthreshold_long_decode_prefill( - self, - request: Request, - *, - num_tokens: int, - near_full_prefill_threshold: int, - ) -> bool: - block_size = int(getattr(self._config, "block_size", 16)) - max_scheduled_tokens = int(self._max_num_scheduled_tokens) - min_subthreshold_tokens = max( - 1, - max_scheduled_tokens - 8 * block_size, - ) - if int(num_tokens) < min_subthreshold_tokens: - return False - if int(num_tokens) >= int(near_full_prefill_threshold): - return False - if not ( - self._should_apply_monolithic_pp_mtp_fractional_extra_output_wait( - request - ) - or self._should_extend_monolithic_pp_mtp_long_decode_output_wait( - request - ) - ): - return False + Args: + token_budget: Remaining token budget for this iteration - # These subthreshold prefill slices are close enough to the online - # max-token boundary to expose the same PP output-visible behavior as - # near-full chunks, but only for the low-acceptance half-block - # medium/long decode slices identified by the request-level TPOT RCA. - return True + Returns: + Tuple of (remaining_budget, scheduled_requests, num_tokens_list) + """ + logger = get_cluster_logger(__name__, self._cluster_type.name) + scheduled = [] + num_tokens_list = [] - def _is_prefix_caching_enabled(self) -> bool: - return getattr(self, "_kv_cache_manager", None) is not None + fast_lane_prefill_enabled = self._cluster_type == ClusterType.PREFILL and ( + self._final_prefill_reserved_slots > 0 + or self._final_prefill_reserved_tokens > 0 + ) - def _sync_prefix_cache_allocation_state( - self, request: Optional[Request] = None - ) -> None: - if not self._is_prefix_caching_enabled(): - return - assert self._kv_cache_manager is not None - self._num_allocated_blocks = int(self._kv_cache_manager.num_used_blocks) - if request is not None: - num_blocks = int(self._kv_cache_manager.get_num_blocks_for_request(request)) - if num_blocks > 0: - self._allocation_map[request.id] = num_blocks - else: - self._allocation_map.pop(request.id, None) - - def _find_request_by_id(self, request_id: int) -> Optional[Request]: - request_groups = [ - getattr(self, "_running_requests", []), - getattr(self, "_request_queue", []), - getattr(self, "_preempted_requests", []), - getattr(self, "_waiting_requests", []), - ] - for requests in request_groups: - for request in requests: - if request.id == request_id: - return request - return None + # Get sorted waiting queue based on policy + waiting_queue = ( + self._build_prefill_waiting_queue() + if fast_lane_prefill_enabled + else deque(self._get_sorted_waiting_queue()) + ) + skipped_waiting_requests: deque[Request] = deque() - def complete_kv_transfer_for_requests( - self, requests: Sequence[Request] - ) -> None: - for request in requests: - if request.id not in self._pending_kv_transfer_requests: - raise ValueError( - "KV transfer completion for request without pending transfer state: " - f"request_id={request.id}, " - f"source_cluster={self._cluster_type.name}, " - f"source_replica={self._replica_id}, " - f"source_dp={self._replica_local_id}" + self._current_iteration_token_budget = token_budget + while waiting_queue and token_budget > 0: + self._current_iteration_token_budget = token_budget + final_waiting_count = ( + self._count_final_fast_lane_requests( + waiting_queue, + final_predicate=self._is_final_prefill_fast_lane_request, ) - - if request.id in self._allocation_map: - self._free_request_resources(request) - self._pending_kv_transfer_requests.discard(request.id) - - def _free_request_resources(self, request: Request) -> None: - self._get_monolithic_pp_mtp_near_full_prefill_request_ids().discard( - request.id - ) - self._get_monolithic_pp_mtp_single_output_wait_request_ids().discard( - request.id - ) - self._get_monolithic_pp_mtp_fractional_output_wait_counts().pop( - request.id, - None, - ) - self._get_monolithic_pp_mtp_output_wait_remaining_iters().pop( - request.id, - None, - ) - self._get_monolithic_pp_mtp_output_wait_request_ids().discard(request.id) - self._get_monolithic_pp_waiting_admission_delay_iters().pop( - request.id, - None, - ) - if self._is_prefix_caching_enabled(): - assert self._kv_cache_manager is not None - self._kv_cache_manager.free(request) - self._allocation_map.pop(request.id, None) - self._sync_prefix_cache_allocation_state() - gdn_slot_manager = self._gdn_state_slot_manager - if gdn_slot_manager is not None: - gdn_slot_manager.release(request.id) - return - self.free(request.id) - gdn_slot_manager = self._gdn_state_slot_manager - if gdn_slot_manager is not None: - gdn_slot_manager.release(request.id) - - def _free_request_resources_by_id(self, request_id: int) -> None: - request = self._find_request_by_id(request_id) - if request is not None: - self._free_request_resources(request) - return - self.free(request_id) - # Completion/cancellation callbacks may arrive after the request has - # left every scheduler queue. Release an orphaned ownership token as - # part of the same idempotent cleanup boundary so a slot cannot leak. - gdn_slot_manager = self._gdn_state_slot_manager - if gdn_slot_manager is not None: - gdn_slot_manager.release(request_id) - - def _prepare_prefix_cache_admission( - self, request: Request - ) -> PrefixCacheAdmission: - if not self._is_prefix_caching_enabled(): - return PrefixCacheAdmission( - raw_hit_blocks=(), - effective_hit_blocks=(), - raw_hit_bindings=(), - effective_hit_bindings=(), - raw_cached_tokens=0, - effective_cached_tokens=0, - num_new_tokens=self._get_request_next_num_tokens(request), - full_hit_backoff_applied=False, - ) - if request.block_hash_ids is None: - raise ValueError( - "block_hash_ids are required when enable_prefix_caching=True" + if fast_lane_prefill_enabled + else 0 ) - assert self._kv_cache_manager is not None - computed_blocks, num_computed_tokens = self._kv_cache_manager.get_computed_blocks( - request - ) - raw_hit_blocks = tuple(computed_blocks) - raw_hit_bindings: list[KVCacheBlockBinding] = [] - query_hashes = list(request.block_hash_ids) - for query_index, block in enumerate(raw_hit_blocks): - binding = block.binding - if binding is None: - raise ValueError( - f"Prefix cache hit block {block.block_id} has no binding identity." - ) - if binding.block_hash != query_hashes[query_index]: - raise ValueError( - "Prefix cache hit binding disagrees with ordered query hash: " - f"query_index={query_index}, " - f"query_hash={query_hashes[query_index]!r}, " - f"binding_hash={binding.block_hash!r}" - ) - raw_hit_bindings.append(binding) - raw_cached_tokens = int(num_computed_tokens) - num_new_tokens = int(request.num_prefill_tokens) - int(num_computed_tokens) - full_hit_backoff_applied = False - if num_new_tokens == 0 and computed_blocks: - num_computed_tokens -= int(self._config.block_size) - num_new_tokens = int(self._config.block_size) - computed_blocks = list(computed_blocks[:-1]) - self._kv_cache_manager.prefix_cache_stats.hits -= 1 - full_hit_backoff_applied = True - return PrefixCacheAdmission( - raw_hit_blocks=raw_hit_blocks, - effective_hit_blocks=tuple(computed_blocks), - raw_hit_bindings=tuple(raw_hit_bindings), - effective_hit_bindings=tuple( - raw_hit_bindings[: len(computed_blocks)] - ), - raw_cached_tokens=raw_cached_tokens, - effective_cached_tokens=int(num_computed_tokens), - num_new_tokens=int(num_new_tokens), - full_hit_backoff_applied=full_hit_backoff_applied, - ) + has_final_waiting = final_waiting_count > 0 + # Check max concurrent requests limit + if len(self._running_requests) >= self._max_num_running_reqs: + break - def _prefix_cache_identity_event_base( - self, - *, - event: str, - request: Request, - ) -> Dict[str, Any]: - event_seq = int(self._prefix_cache_identity_event_seq) - self._prefix_cache_identity_event_seq = event_seq + 1 - cluster_name = ( - self._cluster_type.name - if self._cluster_type is not None - else ClusterType.MONOLITHIC.name - ) - replica_local_id = self._replica_local_id - if replica_local_id is not None and ( - type(replica_local_id) is not int or replica_local_id < 0 - ): - raise ValueError( - "Prefix cache identity replica_local_id must be None or an " - f"exact non-negative int, got {replica_local_id!r}" - ) - return { - "event": event, - "prefix_identity_schema_version": 1, - "source": "frontier", - "scheduler": "vllm_v1", - "cluster_type": cluster_name, - "replica_id": int(self._replica_id), - "replica_local_id": replica_local_id, - "iteration_id": int(self._active_schedule_iteration_id), - "identity_event_seq": event_seq, - "request_id": str(request.id), - "prefix_cache_block_size": int(self._config.block_size), - "simulation_time": float(self._current_schedule_time), - "simulation_time_semantics": "frontier_event_time_seconds", - } - - def _serialize_prefix_cache_hit_bindings( - self, - *, - request: Request, - bindings: Sequence[KVCacheBlockBinding], - ) -> List[Dict[str, Any]]: - query_hashes = list(request.block_hash_ids or []) - if len(bindings) > len(query_hashes): - raise ValueError( - "Prefix cache hit count exceeds the ordered query hash count." - ) - rows: List[Dict[str, Any]] = [] - for query_index, binding in enumerate(bindings): - if binding.block_hash != query_hashes[query_index]: - raise ValueError( - "Prefix cache hit binding disagrees with ordered query hash: " - f"query_index={query_index}, " - f"query_hash={query_hashes[query_index]!r}, " - f"binding_hash={binding.block_hash!r}" + request = waiting_queue[0] + if self._should_defer_monolithic_pp_waiting_admission(request): + logger.debug( + "[VLLMv1Engine][MONOLITHIC] Phase 2: delaying req=%s " + "until a PP output-visible scheduler boundary", + request.id, ) - rows.append( - { - "query_index": query_index, - **_serialize_prefix_cache_binding(binding), - } + break + + is_final_prefill_request = fast_lane_prefill_enabled and ( + self._is_final_prefill_fast_lane_request(request) ) - return rows - - def _emit_prefix_cache_identity_events( - self, - *, - request: Request, - num_new_tokens: int, - allocation: KVCacheAllocationResult, - admission: Optional[PrefixCacheAdmission], - ) -> None: - if admission is not None: - admission_payload = self._prefix_cache_identity_event_base( - event="prefix_cache_admission", - request=request, + is_hidden_prefill_request = ( + fast_lane_prefill_enabled + and not request.is_prefill_complete + and not is_final_prefill_request ) - admission_payload.update( - { - "query_hashes": list(request.block_hash_ids or []), - "raw_hit_blocks": self._serialize_prefix_cache_hit_bindings( - request=request, - bindings=admission.raw_hit_bindings, - ), - "admitted_hit_blocks": self._serialize_prefix_cache_hit_bindings( - request=request, - bindings=admission.effective_hit_bindings, - ), - "raw_cached_tokens": int(admission.raw_cached_tokens), - "admitted_cached_tokens": int( - admission.effective_cached_tokens - ), - "num_new_tokens": int(num_new_tokens), - "full_hit_backoff_applied": bool( - admission.full_hit_backoff_applied - ), - } + computed_blocks = None + prefix_cached_tokens = 0 + prefix_cache_admission: Optional[PrefixCacheAdmission] = None + scheduler_num_computed_tokens = self._get_scheduler_num_computed_tokens( + request ) - _log_frontier_vllm_v1_schedule_decision(admission_payload) - allocation_payload = self._prefix_cache_identity_event_base( - event="prefix_cache_allocation", - request=request, - ) - reused_blocks: List[Dict[str, Any]] = [] - for block in allocation.reused_blocks: - binding = block.binding - if binding is None: - raise ValueError( - f"Reused Prefix cache block {block.block_id} has no binding identity." + # Calculate number of new tokens to process + if self._is_prefix_caching_enabled() and not request.is_prefill_complete: + prefix_cache_admission = self._prepare_prefix_cache_admission( + request ) - reused_blocks.append(_serialize_prefix_cache_binding(binding)) - allocation_payload.update( - { - "num_new_tokens": int(num_new_tokens), - "reused_blocks": reused_blocks, - "new_block_ids": [ - int(block.block_id) for block in allocation.new_blocks - ], - "evicted_bindings": [ - _serialize_prefix_cache_binding(binding) - for binding in allocation.evicted_bindings - ], - "new_bindings": [ - _serialize_prefix_cache_binding(binding) - for binding in allocation.new_bindings - ], - } - ) - _log_frontier_vllm_v1_schedule_decision(allocation_payload) - - def _build_decode_cuda_graph_metadata( - self, batch: Batch - ) -> Optional[DecodeCudaGraphMetadata]: - if self._cluster_type not in (ClusterType.MONOLITHIC, ClusterType.DECODE): - return None - if ( - getattr(self, "_spec_decode_enabled", False) - and not global_vars.get_allow_spec_decode_cuda_graph_diagnostic() - ): - # Phase 2+ baseline: speculative decoding always runs in eager mode. - # We intentionally disable decode CUDA graph modeling for all - # speculative batches to reduce alignment complexity; future work - # can reintroduce method-specific CUDA graph semantics. - return None - - config_mode = global_vars.get_decode_cuda_graph_mode() - if config_mode == "none": - return None + computed_blocks = list( + prefix_cache_admission.effective_hit_blocks + ) + prefix_cached_tokens = int( + prefix_cache_admission.effective_cached_tokens + ) + num_new_tokens = int(prefix_cache_admission.num_new_tokens) + max_allowed = self._max_model_len - prefix_cached_tokens + else: + num_new_tokens = self._get_request_next_num_tokens(request) + max_allowed = self._max_model_len - scheduler_num_computed_tokens - capture_hit, capture_size = self._resolve_decode_cuda_graph_capture_size( - batch.total_num_tokens - ) - decode_query_lens = [ - int(num_tokens) - for request, num_tokens in zip(batch.requests, batch.num_tokens) - if request.is_prefill_complete - ] - original_decode_batch_size = len(decode_query_lens) - - # Align with vLLM's uniform_decode_query_len semantics. - # When speculative decoding is enabled, FULL decode cudagraphs are only - # valid for uniform batches whose query_len matches - # 1 + num_speculative_tokens. Non-uniform speculative verify batches - # must dispatch to mixed/piecewise graphs or fall back to eager. - uniform_decode_query_len = 1 - if getattr(self, "_spec_decode_enabled", False): - spec_decode_config = getattr(self, "_spec_decode_config", None) - if spec_decode_config is None: - raise ValueError("Speculative decoding config is not initialized") - uniform_decode_query_len += int(spec_decode_config.num_speculative_tokens) - - is_uniform_decode_batch = ( - bool(decode_query_lens) - and len(decode_query_lens) == len(batch.requests) - and all( - query_len == uniform_decode_query_len - for query_len in decode_query_lens + # Apply max_model_len limit + num_new_tokens = min(num_new_tokens, max_allowed) + num_new_tokens = self._apply_long_prefill_token_threshold( + request, num_new_tokens ) - ) - is_mixed_batch = not is_uniform_decode_batch - - runtime_mode = "NONE" - if config_mode == "full_decode_only": - if is_uniform_decode_batch and capture_hit: - runtime_mode = "FULL" - elif config_mode == "piecewise" and capture_hit: - runtime_mode = "PIECEWISE" - - if runtime_mode == "NONE": - capture_hit = False - capture_size = batch.total_num_tokens - - padded_decode_batch_size = ( - capture_size if capture_hit else original_decode_batch_size - ) - padded_total_tokens = capture_size if capture_hit else batch.total_num_tokens - - return DecodeCudaGraphMetadata( - config_mode=config_mode, - runtime_mode=runtime_mode, - capture_hit=capture_hit, - is_mixed_batch=is_mixed_batch, - original_total_tokens=batch.total_num_tokens, - padded_total_tokens=padded_total_tokens, - original_decode_batch_size=original_decode_batch_size, - padded_decode_batch_size=padded_decode_batch_size, - ) - def _resolve_decode_cuda_graph_capture_size(self, total_tokens: int) -> Tuple[bool, int]: - cudagraph_capture_sizes = global_vars.get_cudagraph_capture_sizes() - if cudagraph_capture_sizes is None: - max_num_seqs = getattr( - self, - "_max_num_running_reqs", - getattr(self, "_max_batch_size", total_tokens), - ) - max_num_seqs = max(int(max_num_seqs), total_tokens) - cudagraph_capture_sizes = [1, 2, 4] + [ - 8 * i for i in range(1, max_num_seqs // 8 + 1) - ] - - for capture_size in sorted(cudagraph_capture_sizes): - if total_tokens <= capture_size: - return True, int(capture_size) - return False, int(total_tokens) - - def _build_spec_decode_batch_metadata( - self, batch: Batch - ) -> Optional[SpecDecodeBatchMetadata]: - if not getattr(self, "_spec_decode_enabled", False): - return None - if self._cluster_type not in (ClusterType.MONOLITHIC, ClusterType.DECODE): - return None - if batch.num_decode_tokens <= 0: - return None - spec_decode_config = getattr(self, "_spec_decode_config", None) - if spec_decode_config is None: - raise ValueError("Speculative decoding config is not initialized") - - planned_drafts_list: List[int] = [] - verify_tokens_list: List[int] = [] - accepted_drafts_list: List[int] = [] - rejected_drafts_list: List[int] = [] - committed_tokens_list: List[int] = [] - terminal_planned_drafts_list: List[List[int]] = [] - terminal_verify_tokens_list: List[List[int]] = [] - terminal_accepted_drafts_list: List[List[int]] = [] - terminal_rejected_drafts_list: List[List[int]] = [] - terminal_raw_committed_tokens_list: List[List[int]] = [] - per_request_outcomes: Dict[int, Tuple[int, Any, List[Tuple[int, int, int, int, int]]]] = {} - - for request, scheduled_tokens in zip(batch.requests, batch.num_tokens): - if not getattr(request, "is_prefill_complete", False) or not getattr( - request, "spec_decode_enabled", False + effective_token_budget = token_budget + if ( + is_hidden_prefill_request + and has_final_waiting + and self._prefill_iteration_reserved_slots_remaining > 0 + and len(self._running_requests) + >= ( + self._max_num_running_reqs + - self._prefill_iteration_reserved_slots_remaining + ) ): - planned_drafts_list.append(0) - verify_tokens_list.append(0) - accepted_drafts_list.append(0) - rejected_drafts_list.append(0) - committed_tokens_list.append(int(scheduled_tokens)) - terminal_planned_drafts_list.append([]) - terminal_verify_tokens_list.append([]) - terminal_accepted_drafts_list.append([]) - terminal_rejected_drafts_list.append([]) - terminal_raw_committed_tokens_list.append([]) + waiting_queue.popleft() + skipped_waiting_requests.append(request) continue - - request_id = int(request.id) - scheduled_tokens_int = int(scheduled_tokens) - if request_id in per_request_outcomes: - ( - recorded_scheduled_tokens, - recorded_outcome, - recorded_terminal_rows, - ) = per_request_outcomes[request_id] - if recorded_scheduled_tokens != scheduled_tokens_int: - raise ValueError( - "Inconsistent scheduled_tokens for duplicated request in the " - "same batch: " - f"request_id={request_id}, " - f"first={recorded_scheduled_tokens}, " - f"current={scheduled_tokens_int}" - ) - outcome = recorded_outcome - terminal_rows = recorded_terminal_rows - else: - if getattr(request, "spec_method_is_target_embedded_mtp", False): - planned_drafts = int(request.spec_next_planned_draft_tokens) - else: - planned_drafts = max(scheduled_tokens_int - 1, 0) - remaining_decode = request.remaining_decode_tokens - outcome = compute_iteration_outcome( - spec_decode_config, - remaining_decode, - planned_draft_tokens=planned_drafts, - iteration_index=request.spec_total_iterations, - request_id=str(request.id), - ) - request.record_spec_decode_iteration( - verify_tokens=outcome.verify_tokens, - accepted_drafts=outcome.accepted_draft_tokens, - rejected_drafts=outcome.rejected_draft_tokens, - committed_tokens=outcome.committed_tokens, - ) - - next_remaining_decode = max( - remaining_decode - outcome.committed_tokens, + if ( + is_hidden_prefill_request + and has_final_waiting + and self._prefill_iteration_reserved_tokens_remaining > 0 + ): + effective_token_budget = max( + token_budget + - min( + self._prefill_iteration_reserved_tokens_remaining, + token_budget, + ), 0, ) - terminal_rows: List[Tuple[int, int, int, int, int]] = [] - if next_remaining_decode == 0: - terminal_rows = ( - self._get_target_embedded_mtp_terminal_overshoot_rows( - request, - start_iteration_index=request.spec_total_iterations, - ) - ) - request.set_spec_next_planned_draft_tokens( - get_planned_draft_tokens( - spec_decode_config, - next_remaining_decode, - iteration_index=request.spec_total_iterations, - request_id=str(request.id), - ) - ) - per_request_outcomes[request_id] = ( - scheduled_tokens_int, - outcome, - terminal_rows, - ) - - planned_drafts_list.append(outcome.planned_draft_tokens) - verify_tokens_list.append(outcome.verify_tokens) - accepted_drafts_list.append(outcome.accepted_draft_tokens) - rejected_drafts_list.append(outcome.rejected_draft_tokens) - committed_tokens_list.append(outcome.committed_tokens) - terminal_planned_drafts_list.append( - [int(row[0]) for row in terminal_rows] - ) - terminal_verify_tokens_list.append([int(row[1]) for row in terminal_rows]) - terminal_accepted_drafts_list.append( - [int(row[2]) for row in terminal_rows] - ) - terminal_rejected_drafts_list.append( - [int(row[3]) for row in terminal_rows] - ) - terminal_raw_committed_tokens_list.append( - [int(row[4]) for row in terminal_rows] - ) - - metadata = SpecDecodeBatchMetadata( - method=spec_decode_config.method, - planned_draft_tokens_per_request=planned_drafts_list, - verify_tokens_per_request=verify_tokens_list, - accepted_draft_tokens_per_request=accepted_drafts_list, - rejected_draft_tokens_per_request=rejected_drafts_list, - committed_tokens_per_request=committed_tokens_list, - uses_lookahead_slots=getattr( - self, "_spec_method_uses_lookahead_slots", False - ), - terminal_overshoot_planned_draft_tokens_per_request=( - terminal_planned_drafts_list - ), - terminal_overshoot_verify_tokens_per_request=( - terminal_verify_tokens_list - ), - terminal_overshoot_accepted_draft_tokens_per_request=( - terminal_accepted_drafts_list - ), - terminal_overshoot_rejected_draft_tokens_per_request=( - terminal_rejected_drafts_list - ), - terminal_overshoot_raw_committed_tokens_per_request=( - terminal_raw_committed_tokens_list - ), - ) - metadata.validate(len(batch.requests)) - return metadata - - # ========== Scheduling Policy Selection ========== - - def _get_scheduling_policy(self) -> str: - """ - Get the scheduling policy to use. + if effective_token_budget <= 0: + waiting_queue.popleft() + skipped_waiting_requests.append(request) + continue - This method provides a clean interface for policy selection that can be - easily extended in future work to support command-line parameter control. + # When chunked prefill is disabled, waiting prefills that exceed token + # budget are skipped for this iteration. + if ( + not self._enable_chunked_prefill + and not request.is_prefill_complete + and num_new_tokens > effective_token_budget + ): + waiting_queue.popleft() + skipped_waiting_requests.append(request) + continue - Returns: - str: The scheduling policy ('fcfs' or 'priority') - """ - # Use the scheduling policy from configuration - return self._config.scheduling_policy + # Apply token budget limit after chunked-prefill guard + num_new_tokens = min(num_new_tokens, effective_token_budget) - def _get_iteration_phase_aware_waiting_requests(self) -> List[Request]: - if self._cluster_type not in (ClusterType.MONOLITHIC, ClusterType.PREFILL): - return [] - return list(self._preempted_requests) + list(self._request_queue) + if num_new_tokens <= 0: + waiting_queue.popleft() + continue - def _resolve_iteration_round_class(self) -> Optional[str]: - if not self._enable_phase_aware_thinking_profile: - return None - - waiting_requests = self._get_iteration_phase_aware_waiting_requests() - thinking_requests = [ - request - for request in waiting_requests - if getattr(request, "is_thinking_mode_enabled", False) - ] - if not thinking_requests: - return None - if any(request.is_final_thinking_round for request in thinking_requests): - return "final" - return "hidden" - - def _get_iteration_scheduler_profile(self) -> Dict[str, Any]: - round_class = self._resolve_iteration_round_class() - profile = { - "round_class": round_class, - "max_num_running_reqs": int(self._config.batch_size_cap), - "max_num_scheduled_tokens": int(self._config.max_tokens_in_batch), - "enable_chunked_prefill": bool( - getattr(self._config, "enable_chunked_prefill", False) - ), - } - if round_class is None: - return profile - - prefix = f"{round_class}_phase_" - max_tokens_override = getattr(self._config, f"{prefix}max_tokens_in_batch") - chunked_override = getattr( - self._config, f"{prefix}enable_chunked_prefill" - ) - batch_size_override = getattr(self._config, f"{prefix}batch_size_cap") - if max_tokens_override is not None: - profile["max_num_scheduled_tokens"] = int(max_tokens_override) - if chunked_override is not None: - profile["enable_chunked_prefill"] = bool(chunked_override) - if batch_size_override is not None: - profile["max_num_running_reqs"] = int(batch_size_override) - return profile - - def _refresh_iteration_scheduler_profile(self) -> None: - profile = self._get_iteration_scheduler_profile() - self._active_iteration_round_class = profile["round_class"] - self._max_num_running_reqs = int(profile["max_num_running_reqs"]) - self._max_num_scheduled_tokens = int(profile["max_num_scheduled_tokens"]) - self._enable_chunked_prefill = bool(profile["enable_chunked_prefill"]) - - def _maybe_promote_final_round_priority(self, request: Request) -> None: - if not self._enable_final_round_priority_boost: - return - if not getattr(request, "is_thinking_mode_enabled", False): - return - if not request.is_final_thinking_round or request.completed_thinking_rounds <= 0: - return - request.set_priority(min(request.priority, self._final_round_priority_value)) - - def _is_final_prefill_fast_lane_request(self, request: Request) -> bool: - return bool( - getattr(request, "is_thinking_mode_enabled", False) - and request.is_final_thinking_round - and not request.is_prefill_complete - ) - - def _is_final_decode_fast_lane_request(self, request: Request) -> bool: - return bool( - getattr(request, "is_thinking_mode_enabled", False) - and request.is_final_thinking_round - and request.is_prefill_complete - and not request.completed - ) - - def _ordered_requests_with_final_lane( - self, - requests: List[Request], - *, - final_predicate, - ) -> List[Request]: - final_requests: List[Request] = [] - non_final_requests: List[Request] = [] - for request in requests: - if final_predicate(request): - final_requests.append(request) - else: - non_final_requests.append(request) - if not final_requests: - return requests - return final_requests + non_final_requests - - def _build_prefill_waiting_queue(self) -> deque[Request]: - ordered_requests = self._get_sorted_waiting_queue() - if ( - self._cluster_type == ClusterType.PREFILL - and ( - self._final_prefill_reserved_slots > 0 - or self._final_prefill_reserved_tokens > 0 - ) - ): - ordered_requests = self._ordered_requests_with_final_lane( - ordered_requests, - final_predicate=self._is_final_prefill_fast_lane_request, - ) - return deque(ordered_requests) - - def _build_decode_waiting_queue(self) -> deque[Request]: - ordered_requests = list(self._waiting_requests) - if getattr( - getattr(self, "_config", None), "enable_thinking_round_priority", False - ): - ordered_requests.sort( - key=lambda r: ( - 0 if r.is_final_thinking_round else 1, - r.priority, - r.arrived_at, - ) - ) - elif self._scheduling_policy == "priority": - ordered_requests.sort(key=lambda r: (r.priority, r.arrived_at)) - - if ( - self._cluster_type == ClusterType.DECODE - and self._final_decode_reserved_slots > 0 - ): - ordered_requests = self._ordered_requests_with_final_lane( - ordered_requests, - final_predicate=self._is_final_decode_fast_lane_request, - ) - return deque(ordered_requests) - - def _count_final_fast_lane_requests( - self, - requests: List[Request] | deque[Request], - *, - final_predicate, - ) -> int: - return sum(1 for request in requests if final_predicate(request)) - - def _select_final_running_reclaim_victim( - self, - *, - final_predicate, - ) -> Optional[Request]: - candidates = [ - request - for request in self._running_requests - if not final_predicate(request) - ] - if not candidates: - return None - if self._scheduling_policy == "priority": - return max(candidates, key=lambda r: (r.priority, r.arrived_at)) - return candidates[-1] - - def _reclaim_borrowed_final_running_slots( - self, - *, - waiting_requests: List[Request] | deque[Request], - final_predicate, - reserved_slots: int, - lane_name: str, - ) -> List[Request]: - if not self._enable_final_running_request_reclaim or reserved_slots <= 0: - return [] - - final_waiting_count = self._count_final_fast_lane_requests( - waiting_requests, - final_predicate=final_predicate, - ) - if final_waiting_count <= 0: - return [] - - final_running_count = self._count_final_fast_lane_requests( - self._running_requests, - final_predicate=final_predicate, - ) - remaining_reserved_slots = max( - reserved_slots - min(final_running_count, reserved_slots), - 0, - ) - target_new_final_admissions = min( - final_waiting_count, - remaining_reserved_slots, - ) - if target_new_final_admissions <= 0: - return [] - - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - reclaimed_requests: List[Request] = [] - while ( - max(self._max_num_running_reqs - len(self._running_requests), 0) - < target_new_final_admissions - ): - victim = self._select_final_running_reclaim_victim( - final_predicate=final_predicate, - ) - if victim is None: + # Try to allocate (no preemption for waiting requests in Phase 2) + if not self._can_allocate_request( + request, + num_new_tokens, + new_computed_blocks=computed_blocks, + scheduler_num_computed_tokens=scheduler_num_computed_tokens, + ): + # Flow validation: log memory pressure for waiting queue admission + available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) logger.info( - "[FINAL-SLICE-RECLAIM] lane=%s stopped_without_victim " - "target_new_final_admissions=%s running_count=%s", - lane_name, - target_new_final_admissions, - len(self._running_requests), + f"[MEMORY_PRESSURE] trigger=waiting_allocation_failed, " + f"requesting_req={request.id}, " + f"requested_tokens={num_new_tokens}, " + f"available_blocks={available_blocks}, " + f"running_queue_size={len(self._running_requests)}, " + f"waiting_queue_size={len(waiting_queue)}" ) + # Cannot allocate - stop scheduling new requests break - logger.info( - "[FINAL-SLICE-RECLAIM] lane=%s reclaiming_hidden_req=%s " - "target_new_final_admissions=%s running_count_before=%s", - lane_name, - victim.id, - target_new_final_admissions, - len(self._running_requests), - ) - self._preempt_request(victim, reclaimed_requests) - - if reclaimed_requests: - logger.info( - "[FINAL-SLICE-RECLAIM] lane=%s reclaimed_count=%s " - "running_count_after=%s", - lane_name, - len(reclaimed_requests), - len(self._running_requests), - ) - return reclaimed_requests - - def _get_num_waiting_reqs_for_decision_log(self) -> int: - if self._cluster_type in (ClusterType.DECODE, ClusterType.DECODE_ATTN): - return len(self._waiting_requests) - return len(self._request_queue) + len(self._preempted_requests) - - def _apply_long_prefill_token_threshold( - self, request: Request, num_new_tokens: int - ) -> int: - """Apply long prefill threshold only for prefill-phase requests.""" - if request.is_prefill_complete or self._long_prefill_token_threshold <= 0: - return num_new_tokens - return min(num_new_tokens, self._long_prefill_token_threshold) - - def _emit_schedule_decision_event( - self, - *, - event: str, - decision_result: Optional[str], - request_id: Optional[int], - token_budget: int, - num_tokens: int, - available_blocks: Optional[int] = None, - batch_request_ids: Optional[List[int]] = None, - request_num_tokens: Optional[List[int]] = None, - batch_size: int = 0, - batch_num_tokens: int = 0, - ) -> None: - if _frontier_vllm_v1_sched_decision_logger is None: - return - - if available_blocks is None: - available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) - - cluster_name = self._cluster_type.name if self._cluster_type else "MONOLITHIC" - payload: Dict[str, Any] = { - "event": event, - "source": "frontier", - "scheduler": "vllm_v1", - "cluster_type": cluster_name, - "iteration_id": int(self._active_schedule_iteration_id), - "decision_result": decision_result, - "request_id": None if request_id is None else str(request_id), - "token_budget": int(token_budget), - "available_blocks": int(available_blocks), - "num_tokens": int(num_tokens), - "num_running_reqs": len(self._running_requests), - "num_waiting_reqs": self._get_num_waiting_reqs_for_decision_log(), - "max_num_running_reqs": int(self._max_num_running_reqs), - "max_num_scheduled_tokens": int(self._max_num_scheduled_tokens), - "batch_request_ids": [str(req_id) for req_id in (batch_request_ids or [])], - "request_num_tokens": [int(v) for v in (request_num_tokens or [])], - "batch_size": int(batch_size), - "batch_num_tokens": int(batch_num_tokens), - "timestamp": time.time(), - "timestamp_semantics": "wall_clock_epoch_seconds", - "simulation_time": float(self._current_schedule_time), - "simulation_time_semantics": "frontier_event_time_seconds", - } - if self._kv_cache_manager is not None: - prefix_cache_stats = self._kv_cache_manager.prefix_cache_stats - payload.update( - { - "prefix_cache_metric_semantics": "block_level", - "prefix_cache_unit": "blocks", - "prefix_cache_block_size": int(self._config.block_size), - "prefix_cache_requests": int(prefix_cache_stats.requests), - "prefix_cache_queries": int(prefix_cache_stats.queries), - "prefix_cache_hits": int(prefix_cache_stats.hits), - } - ) - _log_frontier_vllm_v1_schedule_decision(payload) - - # ========== Batch Completion Handling ========== - - def on_batch_end(self, batch: Batch) -> None: - """ - Handle batch completion - update running requests state. - - For completed requests: free resources and remove from running list. - For ongoing requests: keep in running list for next iteration. - - Special handling for PREFILL cluster in disaggregated mode: - - Requests are transferred to DECODE cluster after prefill completion - - Partially-prefilled requests stay in PREFILL running list for next chunk - - Args: - batch: The batch that has completed execution - """ - self._num_running_batches -= 1 - - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - self._release_batch_requests_active(batch) - - for request in batch.requests: - self._refresh_target_embedded_mtp_prefill_boundary_state(batch, request) - if self._cluster_type == ClusterType.DECODE_ATTN: - decode_attn_cohort_id = getattr( - batch, - "decode_attn_cohort_id", - None, - ) - if decode_attn_cohort_id is None: - self._decode_attn_active_request_ids.discard(request.id) - else: - cohort_states = self._get_decode_attn_active_cohort_states() - cohort_state = cohort_states.get(int(decode_attn_cohort_id)) - if cohort_state is None: - self._decode_attn_active_request_ids.discard(request.id) - else: - cohort_state["pending_request_ids"].discard(request.id) - if not cohort_state["pending_request_ids"]: - for cohort_request_id in cohort_state["all_request_ids"]: - self._decode_attn_active_request_ids.discard( - cohort_request_id - ) - cohort_states.pop(int(decode_attn_cohort_id), None) - if ( - getattr(self, "_decode_attn_open_cohort_id", None) - == decode_attn_cohort_id - ): - self._decode_attn_open_cohort_id = None - - if request.completed: - extra_release_iters = ( - self._get_monolithic_pp_extra_terminal_release_iters() - ) - if extra_release_iters > 0: - pending_release_iters = ( - self._get_monolithic_pp_pending_terminal_release_iters() - ) - pending_release_iters[request.id] = max( - pending_release_iters.get(request.id, 0), - extra_release_iters, - ) - logger.debug( - "[VLLMv1Engine] Request %s completed, deferring free for " - "%s extra MONOLITHIC+PP terminal iteration(s)", - request.id, - extra_release_iters, - ) - continue - # Request finished - free resources and remove from running - self._free_request_resources(request) - self._scheduled_num_computed_tokens_by_request.pop(request.id, None) - if request in self._running_requests: - self._running_requests.remove(request) - logger.debug( - f"[VLLMv1Engine] Request {request.id} completed, " - f"freed resources, running_reqs={len(self._running_requests)}" - ) - elif self._cluster_type == ClusterType.PREFILL: - # PREFILL cluster in disaggregated mode: - # Requests are transferred to DECODE cluster after prefill completion - # MODIFIED: Do NOT free KV cache here - it will be freed when transfer completes - # This matches vLLM v1 behavior (scheduler.py:1480-1501) - # where blocks are freed on finished_sending event - - if request.is_prefill_complete: - # Remove from running list only after prefill is fully complete. - self._scheduled_num_computed_tokens_by_request.pop(request.id, None) - if request in self._running_requests: - self._running_requests.remove(request) - - # Track that this request's KV cache is pending transfer. - self._pending_kv_transfer_requests.add(request.id) - - logger.info( - f"[VLLMv1Engine] Request {request.id} prefill complete, " - f"KV cache retained for transfer (blocks={self._allocation_map.get(request.id, 0)}), " - f"running_reqs={len(self._running_requests)}" - ) - else: - # Partial prefill: keep request in running queue for the next chunk. - logger.debug( - f"[VLLMv1Engine] Request {request.id} partial prefill complete, " - f"processed_tokens={request.num_processed_tokens}, " - f"running_reqs={len(self._running_requests)}" - ) - elif self._cluster_type == ClusterType.DECODE_ATTN: - # DECODE_ATTN in PD-AF mode: - # This method is called ONLY by GlobalBatchEndEvent (decode step complete) - # NOT called for intermediate layers (those go through _af_immediate_batch_queue) - - # Note: _num_running_batches already decremented at method start (line 151) - # This is correct - decode step completed, release pipeline slot - # For completed requests: free resources and remove from running list - # For ongoing requests: keep in _running_requests for next decode step - # - # Note: request._completed_layer_count is already reset to 0 by request.on_batch_end() - # This ensures layer-consistent grouping in next _schedule_decode_attn_only() call - - if request.completed: - # Request finished all decode tokens - free resources and remove - self._free_request_resources(request) - if request in self._running_requests: - self._running_requests.remove(request) - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Request {request.id} completed, " - f"freed resources, running_reqs={len(self._running_requests)}" - ) - else: - # Request continues with next decode token - keep in _running_requests - # Phase 1 of _schedule_decode_attn_only() will pick this up - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Request {request.id} continues to next decode step, " - f"processed_tokens={request.num_processed_tokens}" - ) - else: - # Request continues - keep in running list - # (will be scheduled again in next iteration) - if self._should_apply_monolithic_pp_mtp_output_wait(request): - output_wait_iters = ( - self._get_monolithic_pp_mtp_output_wait_iters_for_request( - request - ) + self._allocate_request( + request, + num_new_tokens, + new_computed_blocks=computed_blocks, + prefix_cache_admission=( + replace( + prefix_cache_admission, + num_new_tokens=int(num_new_tokens), ) - if output_wait_iters > 0: - self._add_monolithic_pp_mtp_output_wait( - request.id, - wait_iters=output_wait_iters, - ) - logger.debug( - f"[VLLMv1Engine] Request {request.id} continues, " - f"processed_tokens={request.num_processed_tokens}" - ) - - def _get_monolithic_pp_pending_terminal_release_iters(self) -> Dict[int, int]: - pending = getattr( - self, - "_monolithic_pp_pending_terminal_release_iters", - None, - ) - if pending is None: - pending = {} - self._monolithic_pp_pending_terminal_release_iters = pending - return pending - - def _get_monolithic_pp_extra_terminal_release_iters(self) -> int: - if self._cluster_type != ClusterType.MONOLITHIC: - return 0 - - pp = int(getattr(self._replica_config, "num_pipeline_stages", 1)) - if pp <= 1: - return 0 - - # Frontier's last-stage batch-end already accounts for one terminal - # drain iteration. Deeper PP still needs the sampled-token-return - # boundary to reach the scheduler before blocks can be released. - return max(pp // 2 - 1, 0) - - def _has_monolithic_pp_pending_terminal_release(self) -> bool: - return bool(self._get_monolithic_pp_pending_terminal_release_iters()) - - def _has_monolithic_pp_visible_waiting_requests(self) -> bool: - return bool(self._request_queue or self._preempted_requests) - - def _get_monolithic_pp_iteration_start_release_threshold(self) -> int: - if self._cluster_type != ClusterType.MONOLITHIC: - return 1 - - pp = int(getattr(self._replica_config, "num_pipeline_stages", 1)) - if pp <= 4: - return 1 - - # Once terminal release is materialized at iteration_start, deeper - # MONOLITHIC+PP pipelines expose the release boundary earlier than the - # old end-of-iteration bookkeeping. The validated scheduler-visible - # contracts are pp4->1 and pp8->2, so keep the threshold PP-depth - # aware instead of assuming a single remaining hop for every PP size. - return max(pp // 4, 1) - - def _advance_monolithic_pp_terminal_release_boundary(self) -> None: - pending_release_iters = ( - self._get_monolithic_pp_pending_terminal_release_iters() - ) - if not pending_release_iters: - return - - release_visible_threshold = ( - self._get_monolithic_pp_iteration_start_release_threshold() - ) - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - - ready_request_ids: List[int] = [] - for request_id, remaining_iters in list(pending_release_iters.items()): - if ( - remaining_iters <= release_visible_threshold - and self._has_monolithic_pp_visible_waiting_requests() - and request_id - not in self._monolithic_pp_waiting_sensitive_release_extensions - ): - self._monolithic_pp_waiting_sensitive_release_extensions.add(request_id) - pending_release_iters[request_id] = 1 - logger.debug( - "[VLLMv1Engine] Delaying MONOLITHIC+PP terminal release for " - "request %s by one extra empty iteration because waiting " - "requests are already visible", - request_id, - ) - continue - if remaining_iters <= 1: - ready_request_ids.append(request_id) - pending_release_iters.pop(request_id, None) - else: - pending_release_iters[request_id] = remaining_iters - 1 - - if not ready_request_ids: - if pending_release_iters: - self._monolithic_pp_terminal_release_followup_poll_pending = True - logger.debug( - "[VLLMv1Engine] Keeping MONOLITHIC+PP terminal release self-driven " - "with one follow-up schedule poll while pending state remains: %s", - dict(pending_release_iters), - ) - return - - ready_request_id_set = set(ready_request_ids) - for request_id in ready_request_ids: - self._free_request_resources_by_id(request_id) - self._scheduled_num_computed_tokens_by_request.pop(request_id, None) - self._monolithic_pp_waiting_sensitive_release_extensions.discard(request_id) - - self._running_requests = [ - request - for request in self._running_requests - if request.id not in ready_request_id_set - ] - self._monolithic_pp_terminal_release_followup_poll_pending = bool( - pending_release_iters - ) or self._has_monolithic_pp_visible_waiting_requests() - - logger.debug( - "[VLLMv1Engine] Released %s MONOLITHIC+PP terminal request(s) " - "after sampled-token-return-equivalent boundary: %s", - len(ready_request_ids), - ready_request_ids, - ) + if prefix_cache_admission is not None + else None + ), + scheduler_num_computed_tokens=scheduler_num_computed_tokens, + ) - def _materialize_monolithic_pp_terminal_release_before_iteration_start( - self, - ) -> None: - pending_release_iters = ( - self._get_monolithic_pp_pending_terminal_release_iters() - ) - if not pending_release_iters: - return - if self._has_monolithic_pp_visible_waiting_requests(): - return - - release_visible_threshold = ( - self._get_monolithic_pp_iteration_start_release_threshold() - ) - ready_request_ids = [ - request_id - for request_id, remaining_iters in list(pending_release_iters.items()) - if remaining_iters <= release_visible_threshold - ] - if not ready_request_ids: - return - - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - ready_request_id_set = set(ready_request_ids) - for request_id in ready_request_ids: - pending_release_iters.pop(request_id, None) - self._free_request_resources_by_id(request_id) - self._scheduled_num_computed_tokens_by_request.pop(request_id, None) - self._monolithic_pp_waiting_sensitive_release_extensions.discard( - request_id - ) - - self._running_requests = [ - request - for request in self._running_requests - if request.id not in ready_request_id_set - ] - logger.debug( - "[VLLMv1Engine] Materialized %s MONOLITHIC+PP terminal release(s) " - "before iteration_start because no waiting request is visible: %s", - len(ready_request_ids), - ready_request_ids, - ) - - def consume_monolithic_pp_terminal_release_followup_poll(self) -> bool: - pending = bool( - getattr( - self, - "_monolithic_pp_terminal_release_followup_poll_pending", - False, - ) - ) - self._monolithic_pp_terminal_release_followup_poll_pending = False - return pending - - def _get_monolithic_pp_mtp_output_wait_request_ids(self) -> set[int]: - request_ids = getattr( - self, - "_monolithic_pp_mtp_output_wait_request_ids", - None, - ) - if request_ids is None: - request_ids = set() - self._monolithic_pp_mtp_output_wait_request_ids = request_ids - return request_ids - - def _get_monolithic_pp_mtp_output_wait_remaining_iters(self) -> Dict[int, int]: - remaining_iters = getattr( - self, - "_monolithic_pp_mtp_output_wait_remaining_iters", - None, - ) - if remaining_iters is None: - remaining_iters = {} - self._monolithic_pp_mtp_output_wait_remaining_iters = remaining_iters - return remaining_iters - - def _get_monolithic_pp_mtp_fractional_output_wait_counts(self) -> Dict[int, int]: - counts = getattr( - self, - "_monolithic_pp_mtp_fractional_output_wait_counts", - None, - ) - if counts is None: - counts = {} - self._monolithic_pp_mtp_fractional_output_wait_counts = counts - return counts - - def _add_monolithic_pp_mtp_output_wait( - self, request_id: int, *, wait_iters: int = 2 - ) -> None: - if wait_iters <= 0: - raise ValueError(f"wait_iters must be positive, got={wait_iters}") - self._get_monolithic_pp_mtp_output_wait_request_ids().add(request_id) - remaining_iters = self._get_monolithic_pp_mtp_output_wait_remaining_iters() - remaining_iters[request_id] = max( - int(remaining_iters.get(request_id, 0)), - int(wait_iters), - ) - - def _should_apply_monolithic_pp_mtp_output_wait( - self, request: Request - ) -> bool: - if self._cluster_type != ClusterType.MONOLITHIC: - return False - if self._num_stages <= 1: - return False - if not getattr(request, "spec_decode_enabled", False): - return False - if not getattr(request, "spec_method_is_target_embedded_mtp", False): - return False - if not getattr(request, "is_prefill_complete", False): - return False - if ( - request.id - not in self._get_monolithic_pp_mtp_near_full_prefill_request_ids() - ): - return False - processed_decode_tokens = int( - getattr(request, "num_processed_decode_tokens", 0) - ) - if processed_decode_tokens <= 0: - return False - if self._get_monolithic_pp_mtp_output_wait_iters() == 1: - block_size = int(getattr(self._config, "block_size", 16)) - remaining_decode_tokens = ( - int(getattr(request, "num_decode_tokens", 0)) - - processed_decode_tokens - ) - if remaining_decode_tokens <= block_size: - # High-acceptance wide-MTP short tails have no useful future - # prefill visibility to protect once the request is within the - # final block-sized decode window. Skipping this idle turn - # avoids a terminal PP output-wait residual without changing - # CUDA op calibration. - return False - return True - - def _has_monolithic_pp_mtp_output_wait(self) -> bool: - return bool(self._get_monolithic_pp_mtp_output_wait_request_ids()) - - def _should_reserve_monolithic_pp_mtp_visible_budget( - self, request: Request - ) -> bool: - if self._cluster_type != ClusterType.MONOLITHIC: - return False - if self._num_stages <= 1: - return False - if not self._is_target_embedded_mtp_request(request): - return False - if not getattr(request, "is_prefill_complete", False): - return False - block_size = int(getattr(self._config, "block_size", 16)) - active_verify_window_tokens = int( - getattr(request, "spec_current_verify_tokens", 0) - ) - if active_verify_window_tokens <= 0: - spec_config = getattr(self, "_spec_decode_config", None) - if spec_config is None: - spec_config = getattr( - getattr(self, "_replica_config", None), - "speculative_decoding_config", - None, - ) - active_verify_window_tokens = int( - getattr(spec_config, "num_speculative_tokens", 0) - ) - if active_verify_window_tokens < block_size: - # Narrow target-embedded MTP verify windows do not consume a full - # cache block of output-visible scheduler budget. Reserving a - # synthetic block-sized budget for them under MONOLITHIC+PP - # under-batches waiting prefill relative to vLLM clean traces. - return False - return self._has_monolithic_pp_visible_waiting_requests() - - def _get_monolithic_pp_mtp_visible_budget_reservation_tokens( - self, request: Request, token_budget: int - ) -> int: - if not self._should_reserve_monolithic_pp_mtp_visible_budget(request): - return 0 - reserved_tokens = max( - int(getattr(request, "spec_current_verify_tokens", 1)), - self._get_request_next_num_tokens(request), - ) - return min(max(reserved_tokens, 0), token_budget) - - def _clear_monolithic_pp_mtp_output_wait(self) -> None: - request_ids = self._get_monolithic_pp_mtp_output_wait_request_ids() - remaining_iters = self._get_monolithic_pp_mtp_output_wait_remaining_iters() - if not remaining_iters: - request_ids.clear() - return - next_waiting_request_ids: set[int] = set() - for request_id in list(request_ids): - remaining = int(remaining_iters.get(request_id, 1)) - 1 - if remaining > 0: - remaining_iters[request_id] = remaining - next_waiting_request_ids.add(request_id) - else: - remaining_iters.pop(request_id, None) - request_ids.clear() - request_ids.update(next_waiting_request_ids) - - def consume_monolithic_pp_mtp_output_wait_followup_poll(self) -> bool: - pending = bool( - getattr( - self, - "_monolithic_pp_mtp_output_wait_followup_poll_pending", - False, - ) - ) - self._monolithic_pp_mtp_output_wait_followup_poll_pending = False - return pending - - # ========== Memory Allocation Helpers ========== - - def _get_explicit_scheduler_num_computed_tokens( - self, request: Request - ) -> Optional[int]: - scheduled_frontier = getattr( - self, "_scheduled_num_computed_tokens_by_request", {} - ).get(request.id) - if scheduled_frontier is None: - return None - return int(scheduled_frontier) - - def _get_scheduler_num_computed_tokens(self, request: Request) -> int: - """Return the scheduler-visible computed frontier for a request.""" - scheduled_frontier = self._get_explicit_scheduler_num_computed_tokens(request) - if scheduled_frontier is not None: - return scheduled_frontier - - processed_tokens = int(request.num_processed_tokens) - if ( - getattr(self, "_cluster_type", None) == ClusterType.MONOLITHIC - and request.is_prefill_complete - and processed_tokens > int(request.num_prefill_tokens) - ): - # MONOLITHIC request metrics grant the first decode token at the - # prefill-complete boundary, but vLLM's scheduler frontier does not - # advance to that token until the first decode scheduling step. - return max(int(request.num_prefill_tokens), processed_tokens - 1) - return processed_tokens - - def _advance_scheduler_num_computed_tokens( - self, request: Request, num_scheduled_tokens: int - ) -> None: - if num_scheduled_tokens < 0: - raise ValueError( - f"num_scheduled_tokens must be >= 0, got {num_scheduled_tokens}" - ) - self._scheduled_num_computed_tokens_by_request[request.id] = ( - self._get_scheduler_num_computed_tokens(request) + int(num_scheduled_tokens) - ) - - def _get_kv_accounted_processed_tokens(self, request: Request) -> int: - """Return processed tokens used for KV block accounting. - - In MONOLITHIC mode we intentionally count the first generated token at - prefill boundary for request-level progression parity. However, vLLM's - KV block growth does not advance at that boundary; it advances when the - first decode scheduling step is executed. To align block semantics, KV - accounting excludes that boundary token. - """ - explicit_scheduler_frontier = self._get_explicit_scheduler_num_computed_tokens( - request - ) - if getattr(self, "_cluster_type", None) != ClusterType.MONOLITHIC: - if explicit_scheduler_frontier is not None: - return explicit_scheduler_frontier - return int(request.num_processed_tokens) - if not getattr(request, "is_prefill_complete", False): - if explicit_scheduler_frontier is not None: - return explicit_scheduler_frontier - return int(request.num_processed_tokens) - processed_tokens = int(request.num_processed_tokens) - inflight_verify_tokens = 1 - if ( - getattr(request, "spec_decode_enabled", False) - and getattr(request, "spec_method_uses_lookahead_slots", False) - ): - inflight_verify_tokens = max( - 1, int(getattr(request, "spec_current_verify_tokens", 1)) - ) - decode_boundary_adjusted_tokens = max( - int(request.num_prefill_tokens), processed_tokens - inflight_verify_tokens - ) - if explicit_scheduler_frontier is None: - return decode_boundary_adjusted_tokens - return max(explicit_scheduler_frontier, decode_boundary_adjusted_tokens) - - def _get_request_next_num_tokens(self, request: Request) -> int: - assert not request.completed - - computed_tokens = self._get_scheduler_num_computed_tokens(request) - cluster_type = getattr(self, "_cluster_type", None) - - if request.is_prefill_complete: - if getattr(request, "spec_decode_enabled", False): - if getattr(request, "spec_method_is_target_embedded_mtp", False): - planned_drafts = int( - getattr(request, "spec_next_planned_draft_tokens", 0) - ) - if ( - cluster_type == ClusterType.MONOLITHIC - and int(getattr(request, "num_processed_decode_tokens", 0)) - == 1 - and computed_tokens <= int(request.num_prefill_tokens) - ): - return max(planned_drafts, 1) - return 1 + int(getattr(request, "spec_next_planned_draft_tokens", 0)) - if cluster_type == ClusterType.MONOLITHIC: - # In MONOLITHIC mode, request.num_processed_tokens includes the - # post-prefill decode bonus. A new decode step is schedulable - # only when request-side progress has advanced beyond the - # scheduler-visible frontier, mirroring vLLM's - # num_tokens_with_spec/num_computed_tokens gating under PP. - return max(int(request.num_processed_tokens) - computed_tokens, 0) - return 1 - - remaining_prefill_tokens = int(request.num_prefill_tokens) - computed_tokens - return max(remaining_prefill_tokens, 0) - - def _get_num_tokens_for_kv_reservation( - self, request: Request, scheduled_tokens: int - ) -> int: - reserved_tokens = int(scheduled_tokens) - if reserved_tokens <= 0: - raise ValueError( - f"scheduled_tokens must be > 0, got={scheduled_tokens}" - ) - if not getattr(request, "is_prefill_complete", False): - return reserved_tokens - if not getattr(request, "spec_decode_enabled", False): - return reserved_tokens - if getattr(request, "spec_method_uses_lookahead_slots", False): - return reserved_tokens - # ngram/medusa path: no lookahead slot reservation in Phase 1. - # Strategy A keeps allocation simple and lets future decode iterations - # amortize accepted-token KV growth without immediate draft-slot reserves. - return 1 - - def _get_initial_allocation_num_blocks( - self, request: Request, reserved_tokens: int - ) -> int: - """Return the block count used to check and commit a first allocation.""" - kv_accounted_tokens = self._get_kv_accounted_processed_tokens(request) - total_tokens = min( - kv_accounted_tokens + reserved_tokens, self._max_model_len - ) - return ceil(total_tokens / self._config.block_size) - - def _can_allocate_request( - self, - request: Request, - num_new_tokens: int = 1, - new_computed_blocks=None, - *, - scheduler_num_computed_tokens: Optional[int] = None, - ) -> bool: - """ - Check if memory can be allocated for a request. - - For new requests: check if prefill blocks can be allocated. - For running requests: check if at least one block is available. - - Args: - request: The request to check allocation for - num_new_tokens: Number of new tokens to allocate (used for decode) - - Returns: - bool: True if allocation is possible - """ - gdn_slot_manager = self._gdn_state_slot_manager - if ( - gdn_slot_manager is not None - and request.id not in self._allocation_map - and not gdn_slot_manager.has_slot(request.id) - and not gdn_slot_manager.has_available_slot - ): - return False - if self._is_prefix_caching_enabled(): - assert self._kv_cache_manager is not None - return self._kv_cache_manager.can_allocate_slots( - request, - num_new_tokens, - new_computed_blocks=new_computed_blocks, - scheduler_num_computed_tokens=( - self._get_scheduler_num_computed_tokens(request) - if scheduler_num_computed_tokens is None - else scheduler_num_computed_tokens - ), - ) - - reserved_tokens = self._get_num_tokens_for_kv_reservation( - request, num_new_tokens - ) - if request.id not in self._allocation_map: - # New request - estimate blocks from current token frontier - # (already processed + newly scheduled in this iteration), then - # clamp by max_model_len to keep allocation semantics consistent - # with chunked prefill scheduling. - num_required_blocks = self._get_initial_allocation_num_blocks( - request, reserved_tokens - ) - available_blocks = ( - self._config.num_blocks - - self._num_allocated_blocks - - num_required_blocks - ) - return available_blocks >= self._watermark_blocks - - # Running request - check if we need additional blocks for decode - num_tokens_reserved = self._allocation_map[request.id] * self._config.block_size - kv_accounted_tokens = self._get_kv_accounted_processed_tokens(request) - num_tokens_required = max( - 0, kv_accounted_tokens + reserved_tokens - num_tokens_reserved - ) - - if num_tokens_required <= 0: - return True - - # Need additional blocks - num_additional_blocks = ceil(num_tokens_required / self._config.block_size) - return self.can_allocate(num_additional_blocks) - - def _allocate_request( - self, - request: Request, - num_new_tokens: int = 1, - new_computed_blocks=None, - prefix_cache_admission: Optional[PrefixCacheAdmission] = None, - *, - scheduler_num_computed_tokens: Optional[int] = None, - ) -> Optional[KVCacheAllocationResult]: - """ - Allocate memory blocks for a request. - - For new requests: allocate blocks for prefill tokens + first decode token. - For running requests: allocate additional blocks if needed. - - Args: - request: The request to allocate for - num_new_tokens: Number of new tokens being processed - """ - if type(num_new_tokens) is not int or num_new_tokens <= 0: - raise ValueError( - "num_new_tokens must be a positive integer, " - f"got {num_new_tokens!r}" - ) - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - if self._is_prefix_caching_enabled(): - assert self._kv_cache_manager is not None - computed_blocks = list(new_computed_blocks or []) - if prefix_cache_admission is not None: - expected_blocks = prefix_cache_admission.effective_hit_blocks - if len(computed_blocks) != len(expected_blocks) or any( - actual is not expected - for actual, expected in zip(computed_blocks, expected_blocks) - ): - raise ValueError( - "Committed Prefix cache admission blocks differ from " - "the allocation input." - ) - if int(num_new_tokens) != int(prefix_cache_admission.num_new_tokens): - raise ValueError( - "Committed Prefix cache admission token count differs " - "from the allocation input." - ) - allocation = self._kv_cache_manager.allocate_slots( - request, - num_new_tokens, - new_computed_blocks=computed_blocks, - scheduler_num_computed_tokens=( - self._get_scheduler_num_computed_tokens(request) - if scheduler_num_computed_tokens is None - else scheduler_num_computed_tokens - ), - ) - if allocation is None: - raise ValueError( - f"Failed to allocate prefix-cache-managed KV blocks for request {request.id}" - ) - self._sync_prefix_cache_allocation_state(request) - self._emit_prefix_cache_identity_events( - request=request, - num_new_tokens=num_new_tokens, - allocation=allocation, - admission=prefix_cache_admission, - ) - return allocation - - reserved_tokens = self._get_num_tokens_for_kv_reservation( - request, num_new_tokens - ) - - if request.id not in self._allocation_map: - # Commit the exact block count used by the allocation preflight. - # This includes a transferred token frontier for a first - # DECODE_ATTN allocation after the PREFILL handoff. - num_required_blocks = self._get_initial_allocation_num_blocks( - request, reserved_tokens - ) - self.allocate(request.id, num_required_blocks) - gdn_slot_manager = self._gdn_state_slot_manager - if gdn_slot_manager is not None: - try: - gdn_slot_manager.allocate(request.id) - except Exception: - # KV and state ownership must commit atomically from the - # scheduler's perspective. Roll back the KV allocation - # before exposing the slot failure to admission. - self.free(request.id) - raise - logger.debug( - f"[VLLMv1Engine] Allocated {num_required_blocks} blocks for request {request.id} " - f"(scheduled_tokens={num_new_tokens}, reserved_tokens={reserved_tokens})" - ) - return None - - # Running request - check if additional blocks needed - gdn_slot_manager = self._gdn_state_slot_manager - if gdn_slot_manager is not None: - if not gdn_slot_manager.has_slot(request.id): - raise RuntimeError( - f"GDN state slot missing for admitted request {request.id}" - ) - gdn_slot_manager.resume(request.id) - num_tokens_reserved = self._allocation_map[request.id] * self._config.block_size - kv_accounted_tokens = self._get_kv_accounted_processed_tokens(request) - num_tokens_required = max( - 0, kv_accounted_tokens + reserved_tokens - num_tokens_reserved - ) - - if num_tokens_required <= 0: - return None - - # Allocate additional blocks - num_additional_blocks = ceil(num_tokens_required / self._config.block_size) - self.allocate(request.id, num_additional_blocks) - return None - - # ========== Preemption Logic ========== - - def _select_preemption_victim( - self, exclude: Optional[Request] = None - ) -> Optional[Request]: - """ - Select a victim request for preemption based on scheduling policy. - - FCFS policy: Preempt the most recently added request (queue tail). - Priority policy: Preempt the request with lowest priority - (highest priority value, then latest arrival). - - Args: - exclude: Optional request to exclude from victim selection (typically the requesting request) - - Returns: - Optional[Request]: The victim request, or None if no victims available - """ - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - - if not self._running_requests: - return None - - # Filter out excluded request - candidates = ( - [r for r in self._running_requests if r != exclude] - if exclude - else self._running_requests - ) - - if not candidates: - return None - - if self._scheduling_policy == "priority": - # Priority policy: preempt request with highest priority value (lowest priority) - # Tie-breaker: latest arrival time - victim = max(candidates, key=lambda r: (r.priority, r.arrived_at)) - - # Flow validation: log victim selection - logger.info( - f"[VICTIM_SELECTION] policy=PRIORITY, " - f"victim={victim.id}, " - f"priority={victim.priority}, " - f"reason=highest_priority_value" - ) - return victim - else: - # FCFS policy: preempt most recently added (queue tail) - # If exclude is specified and is the tail, select the second-to-last - if exclude and candidates and candidates[-1] != self._running_requests[-1]: - # exclude was the tail, use candidates[-1] which is second-to-last - victim = candidates[-1] - else: - victim = candidates[-1] if candidates else None - - if victim is None: - return None - - # Flow validation: log victim selection - logger.info( - f"[VICTIM_SELECTION] policy=FCFS, " - f"victim={victim.id}, " - f"position=tail, " - f"reason=last_in_running_queue" - ) - return victim - - def _preempt_request( - self, victim: Request, preempted_requests: List[Request] - ) -> None: - """ - Preempt a request - free its resources and move to waiting queue. - - Args: - victim: The request to preempt - preempted_requests: List to track preempted requests for this iteration - """ - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - - # GDN state cannot be dropped and restored by the simulator. Reject - # before touching request counters, allocations, or queue membership. - validate_gdn_runtime_support( - self._replica_config.model_config, - preemption_requires_state_drop=True, - ) - - # Capture state before modification - num_computed_tokens_before = self._get_scheduler_num_computed_tokens(victim) - freed_blocks = self._allocation_map.get(victim.id, 0) - running_count_before = len(self._running_requests) - queue_position_before = ( - self._running_requests.index(victim) - if victim in self._running_requests - else -1 - ) - - # Record preemption statistics in the request entity - # This must be done BEFORE resetting num_processed_tokens - victim.record_preemption(self._cluster_type, num_computed_tokens_before) - victim.advance_runtime_epoch() - - # Remove from running requests - if victim in self._running_requests: - self._running_requests.remove(victim) - - # Free allocated blocks - if victim.id in self._allocation_map: - self._free_request_resources(victim) - - # Mark as preempted and reset the scheduler-visible computed frontier. - # Disaggregated decode requests arrive after PREFILL has completed and - # the prompt KV frontier has transferred to the decode-side cluster. - # Their Request-level token lifecycle must survive memory preemption; - # only scheduler-local computed state and KV allocation are restarted. - victim._preempted = True - if ( - self._cluster_type - not in _REQUEST_PROGRESS_PRESERVING_PREEMPTION_CLUSTER_TYPES - ): - victim._num_processed_tokens = 0 # Reset computed tokens as in vLLM v1 - self._scheduled_num_computed_tokens_by_request.pop(victim.id, None) - - # Record re-entry to waiting queue for waiting time tracking after the - # lifecycle decision above and before adding the request to the queue. - victim.on_enter_waiting_queue(self._current_schedule_time, self._cluster_type) - - # Add to front of appropriate waiting queue (prepend) - # DECODE and DECODE_ATTN clusters use _waiting_requests, others use _request_queue - if self._cluster_type in [ClusterType.DECODE, ClusterType.DECODE_ATTN]: - self._waiting_requests.insert(0, victim) - else: - self._request_queue.insert(0, victim) - - # Track for this iteration - preempted_requests.append(victim) - - logger.info( - f"[VLLMv1Engine] Preempted request {victim.id} " - f"(policy={self._scheduling_policy}), " - f"running_reqs={len(self._running_requests)}" - ) - - # Flow validation: log preemption event - logger.info( - f"[PREEMPTION] req={victim.id} preempted, " - f"policy={self._scheduling_policy}, " - f"freed_blocks={freed_blocks}" - ) - available_blocks_preempt = int(self._config.num_blocks - self._num_allocated_blocks) - self._emit_schedule_decision_event( - event="decision", - decision_result="PREEMPTED", - request_id=victim.id, - token_budget=self._current_iteration_token_budget, - available_blocks=available_blocks_preempt, - num_tokens=0, - ) - - # Flow validation: log detailed preemption info - victim_selection_reason = ( - "lowest_priority" - if self._scheduling_policy == "priority" - else "tail_of_running_queue" - ) - logger.info( - f"[PREEMPTION_DETAIL] req={victim.id}, " - f"num_computed_tokens_before={num_computed_tokens_before}, " - f"freed_blocks={freed_blocks}, " - f"policy={self._scheduling_policy}, " - f"victim_selection_reason={victim_selection_reason}, " - f"queue_position_before={queue_position_before}, " - f"running_count_before={running_count_before}, " - f"running_count_after={len(self._running_requests)}" - ) - - def _try_allocate_with_preemption( - self, - request: Request, - num_new_tokens: int, - preempted_requests: List[Request], - *, - scheduler_num_computed_tokens: Optional[int] = None, - ) -> bool: - """ - Try to allocate memory for a request, preempting other requests if necessary. - - This implements the core preemption loop from vLLM v1 scheduler. - - Args: - request: The request to allocate for - num_new_tokens: Number of new tokens to process - preempted_requests: List to track preempted requests - - Returns: - bool: True if allocation succeeded (possibly after preemption) - """ - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - - while True: - if scheduler_num_computed_tokens is None: - can_allocate = self._can_allocate_request(request, num_new_tokens) - else: - can_allocate = self._can_allocate_request( - request, - num_new_tokens, - scheduler_num_computed_tokens=scheduler_num_computed_tokens, - ) - if can_allocate: - if scheduler_num_computed_tokens is None: - self._allocate_request(request, num_new_tokens) - else: - self._allocate_request( - request, - num_new_tokens, - scheduler_num_computed_tokens=scheduler_num_computed_tokens, - ) - return True - - if not self._enable_preemption: - return False - - # Flow validation: log memory pressure - available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) - logger.info( - f"[MEMORY_PRESSURE] trigger=allocation_failed, " - f"requesting_req={request.id}, " - f"requested_tokens={num_new_tokens}, " - f"available_blocks={available_blocks}, " - f"running_queue_size={len(self._running_requests)}" - ) - - # Select victim for preemption (exclude current request) - victim = self._select_preemption_victim(exclude=request) - - if victim is None: - # No victims available (all other requests have higher priority or no other requests) - # Preempt self and move to waiting queue - self._preempt_request(request, preempted_requests) - return False - - # Preempt victim and try again - self._preempt_request(victim, preempted_requests) - - # ========== Phase 1: RUNNING Requests Scheduling ========== - - def _schedule_running_requests( - self, token_budget: int, preempted_requests: List[Request] - ) -> Tuple[int, List[Request], List[int]]: - """ - Phase 1: Schedule requests currently in RUNNING state. - - Iterate through running requests and try to allocate memory for - their next tokens. May trigger preemption if memory is insufficient. - - Args: - token_budget: Remaining token budget for this iteration - preempted_requests: List to track preempted requests - - Returns: - Tuple of (remaining_budget, scheduled_requests, num_tokens_list) - """ - logger = get_cluster_logger(__name__, self._cluster_type.name) - scheduled = [] - num_tokens_list = [] - waiting_final_prefill_count = ( - self._count_final_fast_lane_requests( - self._preempted_requests + self._request_queue, - final_predicate=self._is_final_prefill_fast_lane_request, - ) - if self._cluster_type == ClusterType.PREFILL - else 0 - ) - - self._current_iteration_token_budget = token_budget - req_index = 0 - while req_index < len(self._running_requests) and token_budget > 0: - self._current_iteration_token_budget = token_budget - request = self._running_requests[req_index] - is_final_prefill_running_request = ( - self._cluster_type == ClusterType.PREFILL - and self._is_final_prefill_fast_lane_request(request) - ) - is_hidden_prefill_running_request = ( - self._cluster_type == ClusterType.PREFILL - and not request.is_prefill_complete - and not is_final_prefill_running_request - ) - - if ( - request.id - in self._get_monolithic_pp_pending_terminal_release_iters() - ): - req_index += 1 - continue - - continuation_request_ids = getattr( - self, "_continuation_request_ids", set() - ) - if request.id in continuation_request_ids: - logger.debug( - "[VLLMv1Engine] Phase 1: skipping req=%s " - "(already scheduled in current cycle)", - request.id, - ) - req_index += 1 - continue - - if ( - self._cluster_type == ClusterType.MONOLITHIC - and self._num_stages > 1 - and request.id - in self._get_monolithic_pp_mtp_output_wait_request_ids() - ): - logger.debug( - "[VLLMv1Engine][MONOLITHIC] Phase 1: delaying req=%s " - "for one PP output-visible MTP scheduler step", - request.id, - ) - req_index += 1 - continue - - active_in_pp_batch = ( - self._cluster_type in {ClusterType.MONOLITHIC, ClusterType.DECODE} - and self._num_stages > 1 - and self._is_request_active_in_batch(request) - ) - if active_in_pp_batch: - if self._cluster_type == ClusterType.MONOLITHIC: - reserved_tokens = ( - self._get_monolithic_pp_mtp_visible_budget_reservation_tokens( - request, - token_budget, - ) - ) - if reserved_tokens > 0: - token_budget -= reserved_tokens - self._current_iteration_token_budget = token_budget - logger.debug( - "[VLLMv1Engine][MONOLITHIC] Phase 1: reserving " - "%s token(s) for active output-visible MTP req=%s", - reserved_tokens, - request.id, - ) - logger.debug( - "[VLLMv1Engine][%s] Phase 1: skipping req=%s " - "(already active in a PP batch)", - self._cluster_type.name, - request.id, - ) - req_index += 1 - continue - - if ( - self._cluster_type in {ClusterType.MONOLITHIC, ClusterType.DECODE} - and self._num_stages > 1 - and getattr(request, "completed_layer_count", 0) != 0 - ): - logger.debug( - "[VLLMv1Engine][%s] Phase 1: skipping in-flight req=%s " - "with layer_count=%s (PP continuation still active)", - self._cluster_type.name, - request.id, - getattr(request, "completed_layer_count", None), - ) - req_index += 1 - continue - - # Calculate number of new tokens to process - num_new_tokens = self._get_request_next_num_tokens(request) - - # Apply max_model_len limit - scheduler_num_computed_tokens = self._get_scheduler_num_computed_tokens( - request - ) - max_allowed = self._max_model_len - scheduler_num_computed_tokens - num_new_tokens = min(num_new_tokens, max_allowed) - num_new_tokens = self._apply_long_prefill_token_threshold( - request, num_new_tokens - ) - - # Apply token budget limit - effective_token_budget = token_budget - if ( - is_hidden_prefill_running_request - and waiting_final_prefill_count > 0 - and self._prefill_iteration_reserved_tokens_remaining > 0 - ): - effective_token_budget = max( - token_budget - - min( - self._prefill_iteration_reserved_tokens_remaining, - token_budget, - ), - 0, - ) - if effective_token_budget <= 0: - req_index += 1 - continue - num_new_tokens = min(num_new_tokens, effective_token_budget) - - if num_new_tokens <= 0: - req_index += 1 - continue - - # Try to allocate with preemption - preempted_count_before = len(preempted_requests) - can_schedule = self._try_allocate_with_preemption( - request, - num_new_tokens, - preempted_requests, - scheduler_num_computed_tokens=scheduler_num_computed_tokens, - ) - token_budget = self._rollback_current_iteration_preempted_requests( - scheduled_requests=scheduled, - scheduled_num_tokens=num_tokens_list, - newly_preempted_requests=preempted_requests[ - preempted_count_before: - ], - token_budget=token_budget, - ) - - if can_schedule: - self._advance_scheduler_num_computed_tokens(request, num_new_tokens) - scheduled.append(request) - num_tokens_list.append(num_new_tokens) - token_budget -= num_new_tokens - self._current_iteration_token_budget = token_budget - if is_final_prefill_running_request: - self._prefill_iteration_reserved_tokens_remaining = max( - self._prefill_iteration_reserved_tokens_remaining - - num_new_tokens, - 0, - ) - req_index += 1 - - # Flow validation: log RUNNING request scheduled - logger.info( - f"[RUNNING_SCHEDULED] req={request.id}, " - f"num_new_tokens={num_new_tokens}, " - f"blocks_allocated={self._allocation_map.get(request.id, 0)}" - ) - available_blocks_running = int( - self._config.num_blocks - self._num_allocated_blocks - ) - self._emit_schedule_decision_event( - event="decision", - decision_result="RUNNING_SCHEDULED", - request_id=request.id, - token_budget=token_budget, - available_blocks=available_blocks_running, - num_tokens=num_new_tokens, - ) - else: - # Request was preempted, stop processing running requests - break - - return token_budget, scheduled, num_tokens_list - - def _rollback_current_iteration_preempted_requests( - self, - *, - scheduled_requests: List[Request], - scheduled_num_tokens: List[int], - newly_preempted_requests: List[Request], - token_budget: int, - ) -> int: - if not newly_preempted_requests: - return token_budget - - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - preempted_request_ids = { - int(request.id) for request in newly_preempted_requests - } - if not preempted_request_ids: - return token_budget - - kept_requests: List[Request] = [] - kept_num_tokens: List[int] = [] - refunded_tokens = 0 - - for scheduled_request, scheduled_tokens in zip( - scheduled_requests, scheduled_num_tokens - ): - if int(scheduled_request.id) in preempted_request_ids: - refunded_tokens += int(scheduled_tokens) - logger.info( - "[RUNNING-SCHEDULE-ROLLBACK] req=%s removed from current iteration " - "after same-iteration preemption, refunded_tokens=%s", - scheduled_request.id, - scheduled_tokens, - ) - continue - kept_requests.append(scheduled_request) - kept_num_tokens.append(int(scheduled_tokens)) - - if refunded_tokens == 0: - return token_budget - - scheduled_requests[:] = kept_requests - scheduled_num_tokens[:] = kept_num_tokens - token_budget += refunded_tokens - self._current_iteration_token_budget = token_budget - return token_budget - - # ========== Phase 2: WAITING Requests Scheduling ========== - - def _get_sorted_waiting_queue(self) -> List[Request]: - """ - Get waiting requests sorted by scheduling policy. - - FCFS: Original queue order (first arrived first). - Priority: Sorted by (priority, arrival_time) ascending. - Thinking-round priority: Final-round requests first, then by - existing policy within each tier. - - Returns: - List of requests in scheduling order - """ - # Combine main queue and preempted requests - # Preempted requests should be prioritized (at front of queue) - combined = self._preempted_requests + self._request_queue - - if getattr( - getattr(self, "_config", None), "enable_thinking_round_priority", False - ): - # Final-round requests first, then by priority, then FIFO - return sorted( - combined, - key=lambda r: ( - 0 if r.is_final_thinking_round else 1, - r.priority, - r.arrived_at, - ), - ) - elif self._scheduling_policy == "priority": - # Sort by priority (ascending) then arrival time (ascending) - return sorted(combined, key=lambda r: (r.priority, r.arrived_at)) - else: - # FCFS: maintain insertion order (preempted first) - return combined - - def _set_waiting_queues_from_ordered_requests( - self, ordered_requests: List[Request] - ) -> None: - """Rebuild waiting queues from ordered requests. - - Requests with `_preempted=True` stay in `_preempted_requests` to keep - preemption recovery semantics and queue priority. - """ - self._preempted_requests = [] - self._request_queue = [] - for request in ordered_requests: - if getattr(request, "_preempted", False): - self._preempted_requests.append(request) - else: - self._request_queue.append(request) - - def _schedule_waiting_requests( - self, token_budget: int - ) -> Tuple[int, List[Request], List[int]]: - """ - Phase 2: Schedule requests in WAITING state. - - Only called when no preemption occurred in Phase 1. - Attempts to admit new requests from the waiting queue. - - Args: - token_budget: Remaining token budget for this iteration - - Returns: - Tuple of (remaining_budget, scheduled_requests, num_tokens_list) - """ - logger = get_cluster_logger(__name__, self._cluster_type.name) - scheduled = [] - num_tokens_list = [] - - fast_lane_prefill_enabled = self._cluster_type == ClusterType.PREFILL and ( - self._final_prefill_reserved_slots > 0 - or self._final_prefill_reserved_tokens > 0 - ) - - # Get sorted waiting queue based on policy - waiting_queue = ( - self._build_prefill_waiting_queue() - if fast_lane_prefill_enabled - else deque(self._get_sorted_waiting_queue()) - ) - skipped_waiting_requests: deque[Request] = deque() - - self._current_iteration_token_budget = token_budget - while waiting_queue and token_budget > 0: - self._current_iteration_token_budget = token_budget - final_waiting_count = ( - self._count_final_fast_lane_requests( - waiting_queue, - final_predicate=self._is_final_prefill_fast_lane_request, - ) - if fast_lane_prefill_enabled - else 0 - ) - has_final_waiting = final_waiting_count > 0 - # Check max concurrent requests limit - if len(self._running_requests) >= self._max_num_running_reqs: - break - - request = waiting_queue[0] - if self._should_defer_monolithic_pp_waiting_admission(request): - logger.debug( - "[VLLMv1Engine][MONOLITHIC] Phase 2: delaying req=%s " - "until a PP output-visible scheduler boundary", - request.id, - ) - break - - is_final_prefill_request = fast_lane_prefill_enabled and ( - self._is_final_prefill_fast_lane_request(request) - ) - is_hidden_prefill_request = ( - fast_lane_prefill_enabled - and not request.is_prefill_complete - and not is_final_prefill_request - ) - computed_blocks = None - prefix_cached_tokens = 0 - prefix_cache_admission: Optional[PrefixCacheAdmission] = None - scheduler_num_computed_tokens = self._get_scheduler_num_computed_tokens( - request - ) - - # Calculate number of new tokens to process - if self._is_prefix_caching_enabled() and not request.is_prefill_complete: - prefix_cache_admission = self._prepare_prefix_cache_admission( - request - ) - computed_blocks = list( - prefix_cache_admission.effective_hit_blocks - ) - prefix_cached_tokens = int( - prefix_cache_admission.effective_cached_tokens - ) - num_new_tokens = int(prefix_cache_admission.num_new_tokens) - max_allowed = self._max_model_len - prefix_cached_tokens - else: - num_new_tokens = self._get_request_next_num_tokens(request) - max_allowed = self._max_model_len - scheduler_num_computed_tokens - - # Apply max_model_len limit - num_new_tokens = min(num_new_tokens, max_allowed) - num_new_tokens = self._apply_long_prefill_token_threshold( - request, num_new_tokens - ) - - effective_token_budget = token_budget - if ( - is_hidden_prefill_request - and has_final_waiting - and self._prefill_iteration_reserved_slots_remaining > 0 - and len(self._running_requests) - >= ( - self._max_num_running_reqs - - self._prefill_iteration_reserved_slots_remaining - ) - ): - waiting_queue.popleft() - skipped_waiting_requests.append(request) - continue - if ( - is_hidden_prefill_request - and has_final_waiting - and self._prefill_iteration_reserved_tokens_remaining > 0 - ): - effective_token_budget = max( - token_budget - - min( - self._prefill_iteration_reserved_tokens_remaining, - token_budget, - ), - 0, - ) - if effective_token_budget <= 0: - waiting_queue.popleft() - skipped_waiting_requests.append(request) - continue - - # When chunked prefill is disabled, waiting prefills that exceed token - # budget are skipped for this iteration. - if ( - not self._enable_chunked_prefill - and not request.is_prefill_complete - and num_new_tokens > effective_token_budget - ): - waiting_queue.popleft() - skipped_waiting_requests.append(request) - continue - - # Apply token budget limit after chunked-prefill guard - num_new_tokens = min(num_new_tokens, effective_token_budget) - - if num_new_tokens <= 0: - waiting_queue.popleft() - continue - - # Try to allocate (no preemption for waiting requests in Phase 2) - if not self._can_allocate_request( - request, - num_new_tokens, - new_computed_blocks=computed_blocks, - scheduler_num_computed_tokens=scheduler_num_computed_tokens, - ): - # Flow validation: log memory pressure for waiting queue admission - available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) - logger.info( - f"[MEMORY_PRESSURE] trigger=waiting_allocation_failed, " - f"requesting_req={request.id}, " - f"requested_tokens={num_new_tokens}, " - f"available_blocks={available_blocks}, " - f"running_queue_size={len(self._running_requests)}, " - f"waiting_queue_size={len(waiting_queue)}" - ) - # Cannot allocate - stop scheduling new requests - break - - self._allocate_request( - request, - num_new_tokens, - new_computed_blocks=computed_blocks, - prefix_cache_admission=( - replace( - prefix_cache_admission, - num_new_tokens=int(num_new_tokens), - ) - if prefix_cache_admission is not None - else None - ), - scheduler_num_computed_tokens=scheduler_num_computed_tokens, - ) - - # Commit queue ownership only after KV and state admission succeeds - waiting_queue.popleft() - was_preempted = request in self._preempted_requests - if request in self._preempted_requests: - self._preempted_requests.remove(request) - if request in self._request_queue: - self._request_queue.remove(request) - self._get_monolithic_pp_waiting_admission_delay_iters().pop( - request.id, None - ) - - # Record leaving waiting queue for waiting time tracking - request.on_leave_waiting_queue( - self._current_schedule_time, self._cluster_type - ) - - if prefix_cached_tokens > 0: - request.on_cache_hit(prefix_cached_tokens) - self._advance_scheduler_num_computed_tokens(request, num_new_tokens) - - # Add to running requests - self._running_requests.append(request) - - # Clear preempted flag if set - request._preempted = False - - scheduled.append(request) - num_tokens_list.append(num_new_tokens) - token_budget -= num_new_tokens - self._current_iteration_token_budget = token_budget - if is_final_prefill_request: - self._prefill_iteration_reserved_slots_remaining = max( - self._prefill_iteration_reserved_slots_remaining - 1, - 0, - ) - self._prefill_iteration_reserved_tokens_remaining = max( - self._prefill_iteration_reserved_tokens_remaining - num_new_tokens, - 0, - ) - - # Flow validation: log WAITING request admission - logger.info( - f"[ADMISSION] req={request.id} admitted, " - f"num_tokens={num_new_tokens}, " - f"running_count={len(self._running_requests)}, " - f"token_budget_remaining={token_budget}" - ) - available_blocks_admission = int( - self._config.num_blocks - self._num_allocated_blocks - ) - self._emit_schedule_decision_event( - event="decision", - decision_result="ADMISSION", - request_id=request.id, - token_budget=token_budget, - available_blocks=available_blocks_admission, - num_tokens=num_new_tokens, - ) - - # Flow validation: log preemption recovery if applicable - if was_preempted: - # Preempted requests need full recomputation from prefill tokens - recompute_tokens = request.num_prefill_tokens - logger.info( - f"[PREEMPTION_RECOVERY] req={request.id}, " - f"was_preempted=True, " - f"recompute_tokens={recompute_tokens}" - ) - - # vLLM parity for skipped waiting requests: - # prepend skipped queue back to waiting queue. - if skipped_waiting_requests: - if self._scheduling_policy == "priority": - merged_requests = list(waiting_queue) + list(skipped_waiting_requests) - waiting_queue = deque( - sorted(merged_requests, key=lambda r: (r.priority, r.arrived_at)) - ) - else: - waiting_queue.extend(skipped_waiting_requests) - - self._set_waiting_queues_from_ordered_requests(list(waiting_queue)) - - return token_budget, scheduled, num_tokens_list - - # ========== Main Scheduling Entry Point ========== - - def _get_next_batch(self, is_micro_batch: bool = False) -> Optional[Batch]: - """ - Build the next batch using vLLM v1 two-phase scheduling algorithm. - - Phase 1: Schedule RUNNING requests (decode phase) - Phase 2: Schedule WAITING requests (prefill phase) - only if no preemption - - This method handles cluster-type-specific behavior: - - MONOLITHIC: Full two-phase scheduling - - PREFILL: Two-phase scheduling for running partial-prefill + waiting admission - - DECODE: Only Phase 1 (scheduling running requests) - - Args: - is_micro_batch: Whether this is for micro-batch (ignored in vLLM v1) - - Returns: - Optional[Batch]: The next batch to execute, or None if no work - """ - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - self._active_schedule_iteration_id = self._schedule_iteration_id - self._schedule_iteration_id += 1 - self._refresh_iteration_scheduler_profile() - logger.info( - "[ITERATION_PROFILE] round_class=%s max_tokens=%s batch_size_cap=%s chunked_prefill=%s", - self._active_iteration_round_class, - self._max_num_scheduled_tokens, - self._max_num_running_reqs, - self._enable_chunked_prefill, - ) - - # Route to cluster-specific scheduling - if self._cluster_type == ClusterType.PREFILL: - return self._schedule_prefill_only() - elif self._cluster_type == ClusterType.DECODE: - return self._schedule_decode_only() - elif self._cluster_type == ClusterType.DECODE_ATTN: - return self._schedule_decode_attn_only(is_micro_batch) - else: - # MONOLITHIC or other: full two-phase scheduling - return self._schedule_two_phase() - - def _schedule_two_phase(self) -> Optional[Batch]: - """ - Full two-phase scheduling for MONOLITHIC cluster. - - Returns: - Optional[Batch]: The scheduled batch - """ - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - - all_scheduled_requests: List[Request] = [] - all_num_tokens: List[int] = [] - preempted_requests: List[Request] = [] - waiting_scheduled: List[Request] = [] - waiting_tokens: List[int] = [] - self._materialize_monolithic_pp_terminal_release_before_iteration_start() - token_budget = self._max_num_scheduled_tokens - available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) - waiting_count = len(self._request_queue) + len(self._preempted_requests) - waiting_final_prefill_count = self._count_final_fast_lane_requests( - self._preempted_requests + self._request_queue, - final_predicate=self._is_final_prefill_fast_lane_request, - ) - self._prefill_iteration_reserved_slots_remaining = ( - self._final_prefill_reserved_slots - if ( - self._enable_final_running_request_reclaim - and waiting_final_prefill_count > 0 - ) - else 0 - ) - self._prefill_iteration_reserved_tokens_remaining = ( - self._final_prefill_reserved_tokens - if ( - self._enable_final_running_request_reclaim - and waiting_final_prefill_count > 0 - ) - else 0 - ) - waiting_final_prefill_count = self._count_final_fast_lane_requests( - self._preempted_requests + self._request_queue, - final_predicate=self._is_final_prefill_fast_lane_request, - ) - self._prefill_iteration_reserved_slots_remaining = ( - self._final_prefill_reserved_slots if waiting_final_prefill_count > 0 else 0 - ) - self._prefill_iteration_reserved_tokens_remaining = ( - self._final_prefill_reserved_tokens if waiting_final_prefill_count > 0 else 0 - ) - - # Flow validation: log iteration start - logger.info( - f"[ITERATION_START] token_budget={token_budget}, " - f"running_count={len(self._running_requests)}, " - f"waiting_count={waiting_count}, " - f"available_blocks={available_blocks}, " - f"max_running_reqs={self._max_num_running_reqs}" - ) - self._emit_schedule_decision_event( - event="iteration_start", - decision_result=None, - request_id=None, - token_budget=token_budget, - available_blocks=available_blocks, - num_tokens=0, - ) - - # Flow validation: log memory state - total_blocks = int(self._config.num_blocks) - allocated_blocks = int(self._num_allocated_blocks) - usage_ratio = allocated_blocks / total_blocks if total_blocks > 0 else 0.0 - watermark = self._watermark_blocks - logger.info( - f"[MEMORY_STATE] total_blocks={total_blocks}, " - f"allocated_blocks={allocated_blocks}, " - f"free_blocks={available_blocks}, " - f"usage_ratio={usage_ratio:.4f}, " - f"watermark_blocks={watermark}" - ) - - if self._monolithic_pp_terminal_release_followup_poll_pending: - self._emit_schedule_decision_event( - event="iteration_end", - decision_result=None, - request_id=None, - token_budget=token_budget, - num_tokens=0, - available_blocks=int( - self._config.num_blocks - self._num_allocated_blocks - ), - batch_request_ids=[], - request_num_tokens=[], - batch_size=0, - batch_num_tokens=0, - ) - return None - - # Flow validation: log Phase 1 start - logger.info( - f"[PHASE1_START] running_count={len(self._running_requests)}, " - f"token_budget={token_budget}" - ) - - # === Phase 1: Schedule RUNNING requests === - token_budget, running_scheduled, running_tokens = ( - self._schedule_running_requests(token_budget, preempted_requests) - ) - all_scheduled_requests.extend(running_scheduled) - all_num_tokens.extend(running_tokens) - - # Flow validation: log Phase 1 end - available_blocks_p1 = int(self._config.num_blocks - self._num_allocated_blocks) - logger.info( - f"[PHASE1_END] scheduled_count={len(running_scheduled)}, " - f"preempted_count={len(preempted_requests)}, " - f"token_budget_remaining={token_budget}, " - f"available_blocks={available_blocks_p1}" - ) - - # === Phase 2: Schedule WAITING requests (only if no preemption) === - if not preempted_requests and not self._has_monolithic_pp_pending_terminal_release(): - # Flow validation: log Phase 2 start - waiting_count_p2 = len(self._request_queue) + len(self._preempted_requests) - logger.info( - f"[PHASE2_START] waiting_count={waiting_count_p2}, " - f"token_budget={token_budget}, " - f"running_count={len(self._running_requests)}" - ) - - token_budget, waiting_scheduled, waiting_tokens = ( - self._schedule_waiting_requests(token_budget) - ) - all_scheduled_requests.extend(waiting_scheduled) - all_num_tokens.extend(waiting_tokens) - - # Flow validation: log Phase 2 end - available_blocks_p2 = int( - self._config.num_blocks - self._num_allocated_blocks - ) - logger.info( - f"[PHASE2_END] admitted_count={len(waiting_scheduled)}, " - f"token_budget_remaining={token_budget}, " - f"available_blocks={available_blocks_p2}, " - f"running_count={len(self._running_requests)}" - ) - elif not preempted_requests and self._has_monolithic_pp_pending_terminal_release(): - logger.info( - "[PHASE2_SKIPPED] waiting admission blocked by pending " - "MONOLITHIC+PP terminal release boundary" - ) - - if not all_scheduled_requests: - if self._has_monolithic_pp_mtp_output_wait(): - self._clear_monolithic_pp_mtp_output_wait() - self._monolithic_pp_mtp_output_wait_followup_poll_pending = True - self._advance_monolithic_pp_terminal_release_boundary() - self._emit_schedule_decision_event( - event="iteration_end", - decision_result=None, - request_id=None, - token_budget=token_budget, - num_tokens=0, - available_blocks=int(self._config.num_blocks - self._num_allocated_blocks), - batch_request_ids=[], - request_num_tokens=[], - batch_size=0, - batch_num_tokens=0, - ) - return None - - # Match vLLM v1 output order: new/resumed admissions first, then running. - ordered_scheduled_requests = waiting_scheduled + running_scheduled - ordered_num_tokens = waiting_tokens + running_tokens - - # Flow validation: log batch formation - total_tokens = sum(all_num_tokens) - new_admitted = len( - [r for r in all_scheduled_requests if r not in running_scheduled] - ) - resumed = len( - [r for r in all_scheduled_requests if getattr(r, "_preempted", False)] - ) - running_continued = len(running_scheduled) - batch_size = len(all_scheduled_requests) - - logger.info( - f"[BATCH_FORMATION] total_tokens={total_tokens}, " - f"new_admitted={new_admitted}, " - f"resumed={resumed}, " - f"running_continued={running_continued}, " - f"batch_size={batch_size}" - ) - self._emit_schedule_decision_event( - event="iteration_end", - decision_result=None, - request_id=None, - token_budget=token_budget, - num_tokens=total_tokens, - available_blocks=int(self._config.num_blocks - self._num_allocated_blocks), - batch_request_ids=[request.id for request in ordered_scheduled_requests], - request_num_tokens=ordered_num_tokens, - batch_size=batch_size, - batch_num_tokens=total_tokens, - ) - self._advance_monolithic_pp_terminal_release_boundary() - - return self._create_batch(ordered_scheduled_requests, ordered_num_tokens) - - def _schedule_prefill_only(self) -> Optional[Batch]: - """ - Scheduling for PREFILL cluster. - - In PD-disaggregation, the prefill cluster only handles new requests - that need prefill computation. With chunked prefill enabled, running - partial-prefill requests are also scheduled in Phase 1. - - Returns: - Optional[Batch]: The scheduled batch - """ - logger = get_cluster_logger( - __name__, self._cluster_type.name if self._cluster_type else None - ) - - token_budget = self._max_num_scheduled_tokens - available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) - waiting_count = len(self._request_queue) + len(self._preempted_requests) - - # Flow validation: log iteration start - logger.info( - f"[ITERATION_START] token_budget={token_budget}, " - f"running_count={len(self._running_requests)}, " - f"waiting_count={waiting_count}, " - f"available_blocks={available_blocks}, " - f"max_running_reqs={self._max_num_running_reqs}" - ) - self._emit_schedule_decision_event( - event="iteration_start", - decision_result=None, - request_id=None, - token_budget=token_budget, - available_blocks=available_blocks, - num_tokens=0, - ) + # Commit queue ownership only after KV and state admission succeeds + waiting_queue.popleft() + was_preempted = request in self._preempted_requests + if request in self._preempted_requests: + self._preempted_requests.remove(request) + if request in self._request_queue: + self._request_queue.remove(request) + self._get_monolithic_pp_waiting_admission_delay_iters().pop( + request.id, None + ) - # Flow validation: log memory state - total_blocks = int(self._config.num_blocks) - allocated_blocks = int(self._num_allocated_blocks) - usage_ratio = allocated_blocks / total_blocks if total_blocks > 0 else 0.0 - watermark = self._watermark_blocks - logger.info( - f"[MEMORY_STATE] total_blocks={total_blocks}, " - f"allocated_blocks={allocated_blocks}, " - f"free_blocks={available_blocks}, " - f"usage_ratio={usage_ratio:.4f}, " - f"watermark_blocks={watermark}" - ) - reclaimed_requests = self._reclaim_borrowed_final_running_slots( - waiting_requests=self._preempted_requests + self._request_queue, - final_predicate=self._is_final_prefill_fast_lane_request, - reserved_slots=self._final_prefill_reserved_slots, - lane_name="prefill", - ) - if reclaimed_requests: - waiting_count = len(self._request_queue) + len(self._preempted_requests) + # Record leaving waiting queue for waiting time tracking + request.on_leave_waiting_queue( + self._current_schedule_time, self._cluster_type + ) - all_scheduled_requests: List[Request] = [] - all_num_tokens: List[int] = [] - preempted_requests: List[Request] = [] - waiting_scheduled: List[Request] = [] - waiting_tokens: List[int] = [] + if prefix_cached_tokens > 0: + request.on_cache_hit(prefix_cached_tokens) + self._advance_scheduler_num_computed_tokens(request, num_new_tokens) - # Phase 1: schedule running requests (partial prefill continuation) - logger.info( - f"[PHASE1_START] running_count={len(self._running_requests)}, " - f"token_budget={token_budget}, " - f"waiting_count={waiting_count}" - ) - token_budget, running_scheduled, running_tokens = self._schedule_running_requests( - token_budget, preempted_requests - ) - all_scheduled_requests.extend(running_scheduled) - all_num_tokens.extend(running_tokens) + # Add to running requests + self._running_requests.append(request) - available_blocks_p1 = int(self._config.num_blocks - self._num_allocated_blocks) - logger.info( - f"[PHASE1_END] scheduled_count={len(running_scheduled)}, " - f"preempted_count={len(preempted_requests)}, " - f"token_budget_remaining={token_budget}, " - f"available_blocks={available_blocks_p1}" - ) + # Clear preempted flag if set + request._preempted = False + + scheduled.append(request) + num_tokens_list.append(num_new_tokens) + token_budget -= num_new_tokens + self._current_iteration_token_budget = token_budget + if is_final_prefill_request: + self._prefill_iteration_reserved_slots_remaining = max( + self._prefill_iteration_reserved_slots_remaining - 1, + 0, + ) + self._prefill_iteration_reserved_tokens_remaining = max( + self._prefill_iteration_reserved_tokens_remaining - num_new_tokens, + 0, + ) - # Phase 2: schedule waiting requests only when Phase 1 has no preemption - if not preempted_requests: + # Flow validation: log WAITING request admission logger.info( - f"[PHASE2_START] waiting_count={waiting_count}, " - f"token_budget={token_budget}, " - f"running_count={len(self._running_requests)}" - ) - token_budget, waiting_scheduled, waiting_tokens = ( - self._schedule_waiting_requests(token_budget) + f"[ADMISSION] req={request.id} admitted, " + f"num_tokens={num_new_tokens}, " + f"running_count={len(self._running_requests)}, " + f"token_budget_remaining={token_budget}" ) - all_scheduled_requests.extend(waiting_scheduled) - all_num_tokens.extend(waiting_tokens) - - available_blocks_p2 = int( + available_blocks_admission = int( self._config.num_blocks - self._num_allocated_blocks ) - logger.info( - f"[PHASE2_END] admitted_count={len(waiting_scheduled)}, " - f"token_budget_remaining={token_budget}, " - f"available_blocks={available_blocks_p2}, " - f"running_count={len(self._running_requests)}" - ) - - if not all_scheduled_requests: self._emit_schedule_decision_event( - event="iteration_end", - decision_result=None, - request_id=None, + event="decision", + decision_result="ADMISSION", + request_id=request.id, token_budget=token_budget, - num_tokens=0, - available_blocks=int(self._config.num_blocks - self._num_allocated_blocks), - batch_request_ids=[], - request_num_tokens=[], - batch_size=0, - batch_num_tokens=0, + available_blocks=available_blocks_admission, + num_tokens=num_new_tokens, ) - return None - ordered_scheduled_requests = waiting_scheduled + running_scheduled - ordered_num_tokens = waiting_tokens + running_tokens + # Flow validation: log preemption recovery if applicable + if was_preempted: + # Preempted requests need full recomputation from prefill tokens + recompute_tokens = request.num_prefill_tokens + logger.info( + f"[PREEMPTION_RECOVERY] req={request.id}, " + f"was_preempted=True, " + f"recompute_tokens={recompute_tokens}" + ) - # Flow validation: log batch formation - total_tokens = sum(all_num_tokens) - new_admitted = len(waiting_scheduled) - resumed = len( - [r for r in all_scheduled_requests if getattr(r, "_preempted", False)] - ) - running_continued = len(running_scheduled) - batch_size = len(all_scheduled_requests) + # vLLM parity for skipped waiting requests: + # prepend skipped queue back to waiting queue. + if skipped_waiting_requests: + if self._scheduling_policy == "priority": + merged_requests = list(waiting_queue) + list(skipped_waiting_requests) + waiting_queue = deque( + sorted(merged_requests, key=lambda r: (r.priority, r.arrived_at)) + ) + else: + waiting_queue.extend(skipped_waiting_requests) - logger.info( - f"[BATCH_FORMATION] total_tokens={total_tokens}, " - f"new_admitted={new_admitted}, " - f"resumed={resumed}, " - f"running_continued={running_continued}, " - f"batch_size={batch_size}" - ) - self._emit_schedule_decision_event( - event="iteration_end", - decision_result=None, - request_id=None, - token_budget=token_budget, - num_tokens=total_tokens, - available_blocks=int(self._config.num_blocks - self._num_allocated_blocks), - batch_request_ids=[request.id for request in ordered_scheduled_requests], - request_num_tokens=ordered_num_tokens, - batch_size=batch_size, - batch_num_tokens=total_tokens, - ) + self._set_waiting_queues_from_ordered_requests(list(waiting_queue)) - return self._create_batch(ordered_scheduled_requests, ordered_num_tokens) + return token_budget, scheduled, num_tokens_list + + def _get_next_batch(self, is_micro_batch: bool = False) -> Optional[Batch]: + """ + Build the next batch using vLLM v1 two-phase scheduling algorithm. + + Phase 1: Schedule RUNNING requests (decode phase) + Phase 2: Schedule WAITING requests (prefill phase) - only if no preemption + + This method handles cluster-type-specific behavior: + - MONOLITHIC: Full two-phase scheduling + - PREFILL: Two-phase scheduling for running partial-prefill + waiting admission + - DECODE: Only Phase 1 (scheduling running requests) + + Args: + is_micro_batch: Whether this is for micro-batch (ignored in vLLM v1) - def _schedule_decode_only(self) -> Optional[Batch]: + Returns: + Optional[Batch]: The next batch to execute, or None if no work """ - Scheduling for DECODE cluster - two-phase scheduling matching vLLM v1. + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + self._active_schedule_iteration_id = self._schedule_iteration_id + self._schedule_iteration_id += 1 + self._refresh_iteration_scheduler_profile() + logger.info( + "[ITERATION_PROFILE] round_class=%s max_tokens=%s batch_size_cap=%s chunked_prefill=%s", + self._active_iteration_round_class, + self._max_num_scheduled_tokens, + self._max_num_running_reqs, + self._enable_chunked_prefill, + ) - Phase 1: Schedule RUNNING requests (ongoing decode iterations) - Phase 2: Admit WAITING requests (new arrivals from prefill cluster) + # Route to cluster-specific scheduling + if self._cluster_type == ClusterType.PREFILL: + return self._schedule_prefill_only() + elif self._cluster_type == ClusterType.DECODE: + return self._schedule_decode_only() + elif self._cluster_type == ClusterType.DECODE_ATTN: + return self._schedule_decode_attn_only(is_micro_batch) + else: + # MONOLITHIC or other: full two-phase scheduling + return self._schedule_two_phase() - This matches vLLM v1's scheduling algorithm where requests must be - admitted from waiting queue to running queue before generating tokens. + def _schedule_two_phase(self) -> Optional[Batch]: + """ + Full two-phase scheduling for MONOLITHIC cluster. Returns: Optional[Batch]: The scheduled batch @@ -4282,21 +1067,46 @@ def _schedule_decode_only(self) -> Optional[Batch]: preempted_requests: List[Request] = [] waiting_scheduled: List[Request] = [] waiting_tokens: List[int] = [] + self._materialize_monolithic_pp_terminal_release_before_iteration_start() token_budget = self._max_num_scheduled_tokens available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) - waiting_final_decode_count = self._count_final_fast_lane_requests( - self._waiting_requests, - final_predicate=self._is_final_decode_fast_lane_request, + waiting_count = len(self._request_queue) + len(self._preempted_requests) + waiting_final_prefill_count = self._count_final_fast_lane_requests( + self._preempted_requests + self._request_queue, + final_predicate=self._is_final_prefill_fast_lane_request, + ) + self._prefill_iteration_reserved_slots_remaining = ( + self._final_prefill_reserved_slots + if ( + self._enable_final_running_request_reclaim + and waiting_final_prefill_count > 0 + ) + else 0 + ) + self._prefill_iteration_reserved_tokens_remaining = ( + self._final_prefill_reserved_tokens + if ( + self._enable_final_running_request_reclaim + and waiting_final_prefill_count > 0 + ) + else 0 + ) + waiting_final_prefill_count = self._count_final_fast_lane_requests( + self._preempted_requests + self._request_queue, + final_predicate=self._is_final_prefill_fast_lane_request, + ) + self._prefill_iteration_reserved_slots_remaining = ( + self._final_prefill_reserved_slots if waiting_final_prefill_count > 0 else 0 ) - self._decode_iteration_reserved_slots_remaining = ( - self._final_decode_reserved_slots if waiting_final_decode_count > 0 else 0 + self._prefill_iteration_reserved_tokens_remaining = ( + self._final_prefill_reserved_tokens if waiting_final_prefill_count > 0 else 0 ) # Flow validation: log iteration start logger.info( f"[ITERATION_START] token_budget={token_budget}, " f"running_count={len(self._running_requests)}, " - f"waiting_count={len(self._waiting_requests)}, " + f"waiting_count={waiting_count}, " f"available_blocks={available_blocks}, " f"max_running_reqs={self._max_num_running_reqs}" ) @@ -4321,12 +1131,23 @@ def _schedule_decode_only(self) -> Optional[Batch]: f"usage_ratio={usage_ratio:.4f}, " f"watermark_blocks={watermark}" ) - self._reclaim_borrowed_final_running_slots( - waiting_requests=self._waiting_requests, - final_predicate=self._is_final_decode_fast_lane_request, - reserved_slots=self._final_decode_reserved_slots, - lane_name="decode", - ) + + if self._monolithic_pp_terminal_release_followup_poll_pending: + self._emit_schedule_decision_event( + event="iteration_end", + decision_result=None, + request_id=None, + token_budget=token_budget, + num_tokens=0, + available_blocks=int( + self._config.num_blocks - self._num_allocated_blocks + ), + batch_request_ids=[], + request_num_tokens=[], + batch_size=0, + batch_num_tokens=0, + ) + return None # Flow validation: log Phase 1 start logger.info( @@ -4350,17 +1171,18 @@ def _schedule_decode_only(self) -> Optional[Batch]: f"available_blocks={available_blocks_p1}" ) - # === Phase 2: Admit WAITING requests (only if no preemption) === - if not preempted_requests: + # === Phase 2: Schedule WAITING requests (only if no preemption) === + if not preempted_requests and not self._has_monolithic_pp_pending_terminal_release(): # Flow validation: log Phase 2 start + waiting_count_p2 = len(self._request_queue) + len(self._preempted_requests) logger.info( - f"[PHASE2_START] waiting_count={len(self._waiting_requests)}, " + f"[PHASE2_START] waiting_count={waiting_count_p2}, " f"token_budget={token_budget}, " f"running_count={len(self._running_requests)}" ) token_budget, waiting_scheduled, waiting_tokens = ( - self._schedule_decode_waiting_requests(token_budget) + self._schedule_waiting_requests(token_budget) ) all_scheduled_requests.extend(waiting_scheduled) all_num_tokens.extend(waiting_tokens) @@ -4375,8 +1197,17 @@ def _schedule_decode_only(self) -> Optional[Batch]: f"available_blocks={available_blocks_p2}, " f"running_count={len(self._running_requests)}" ) + elif not preempted_requests and self._has_monolithic_pp_pending_terminal_release(): + logger.info( + "[PHASE2_SKIPPED] waiting admission blocked by pending " + "MONOLITHIC+PP terminal release boundary" + ) if not all_scheduled_requests: + if self._has_monolithic_pp_mtp_output_wait(): + self._clear_monolithic_pp_mtp_output_wait() + self._monolithic_pp_mtp_output_wait_followup_poll_pending = True + self._advance_monolithic_pp_terminal_release_boundary() self._emit_schedule_decision_event( event="iteration_end", decision_result=None, @@ -4391,7 +1222,7 @@ def _schedule_decode_only(self) -> Optional[Batch]: ) return None - # Match vLLM v1 output order: new admissions first, then running. + # Match vLLM v1 output order: new/resumed admissions first, then running. ordered_scheduled_requests = waiting_scheduled + running_scheduled ordered_num_tokens = waiting_tokens + running_tokens @@ -4400,7 +1231,9 @@ def _schedule_decode_only(self) -> Optional[Batch]: new_admitted = len( [r for r in all_scheduled_requests if r not in running_scheduled] ) - resumed = 0 # DECODE doesn't handle preempted requests (they come from prefill) + resumed = len( + [r for r in all_scheduled_requests if getattr(r, "_preempted", False)] + ) running_continued = len(running_scheduled) batch_size = len(all_scheduled_requests) @@ -4423,597 +1256,12 @@ def _schedule_decode_only(self) -> Optional[Batch]: batch_size=batch_size, batch_num_tokens=total_tokens, ) + self._advance_monolithic_pp_terminal_release_boundary() return self._create_batch(ordered_scheduled_requests, ordered_num_tokens) - def _schedule_decode_waiting_requests( - self, token_budget: int - ) -> Tuple[int, List[Request], List[int]]: - """ - Phase 2 for DECODE cluster: Admit requests from waiting queue. - - This method handles requests that have arrived from the prefill cluster - and are waiting to be admitted to the running queue for decode iterations. - Matches vLLM v1's Phase 2 scheduling behavior. - - Args: - token_budget: Remaining token budget for this iteration - - Returns: - Tuple of (remaining_budget, scheduled_requests, num_tokens_list) - """ - logger = get_cluster_logger(__name__, self._cluster_type.name) - scheduled: List[Request] = [] - num_tokens_list: List[int] = [] - - fast_lane_decode_enabled = self._cluster_type == ClusterType.DECODE and ( - self._final_decode_reserved_slots > 0 - ) - waiting_queue = self._build_decode_waiting_queue() - skipped_waiting_requests: deque[Request] = deque() - - self._current_iteration_token_budget = token_budget - while waiting_queue and token_budget > 0: - self._current_iteration_token_budget = token_budget - final_waiting_count = ( - self._count_final_fast_lane_requests( - waiting_queue, - final_predicate=self._is_final_decode_fast_lane_request, - ) - if fast_lane_decode_enabled - else 0 - ) - has_final_waiting = final_waiting_count > 0 - # Check max concurrent requests limit - if len(self._running_requests) >= self._max_num_running_reqs: - logger.debug( - f"[VLLMv1Engine][DECODE] Phase 2: max running requests " - f"reached ({self._max_num_running_reqs}), stopping admission" - ) - break - - request = waiting_queue[0] - is_final_decode_request = fast_lane_decode_enabled and ( - self._is_final_decode_fast_lane_request(request) - ) - is_hidden_decode_request = ( - fast_lane_decode_enabled and not is_final_decode_request - ) - - if ( - is_hidden_decode_request - and has_final_waiting - and self._decode_iteration_reserved_slots_remaining > 0 - and len(self._running_requests) - >= ( - self._max_num_running_reqs - - self._decode_iteration_reserved_slots_remaining - ) - ): - waiting_queue.popleft() - skipped_waiting_requests.append(request) - continue - - num_new_tokens = self._get_request_next_num_tokens(request) - - # Apply max_model_len limit - scheduler_num_computed_tokens = self._get_scheduler_num_computed_tokens( - request - ) - max_allowed = self._max_model_len - scheduler_num_computed_tokens - num_new_tokens = min(num_new_tokens, max_allowed) - - # Apply token budget limit - num_new_tokens = min(num_new_tokens, token_budget) - - if num_new_tokens <= 0: - # Request has reached max length, remove from queue - waiting_queue.popleft() - logger.debug( - f"[VLLMv1Engine][DECODE] Phase 2: req={request.id} " - f"reached max length, removing from waiting queue" - ) - continue - - # Try to allocate (no preemption for waiting requests in Phase 2) - if not self._can_allocate_request( - request, - num_new_tokens, - scheduler_num_computed_tokens=scheduler_num_computed_tokens, - ): - # Cannot allocate - stop admitting new requests - logger.debug( - f"[VLLMv1Engine][DECODE] Phase 2: cannot allocate " - f"req={request.id}, stopping admission" - ) - break - - # Check if this request was previously preempted - was_preempted = getattr(request, "_preempted", False) - - # Remove from waiting queue and allocate - waiting_queue.popleft() - - # Record leaving waiting queue for waiting time tracking - request.on_leave_waiting_queue( - self._current_schedule_time, self._cluster_type - ) - - self._allocate_request( - request, - num_new_tokens, - scheduler_num_computed_tokens=scheduler_num_computed_tokens, - ) - self._advance_scheduler_num_computed_tokens(request, num_new_tokens) - - # Add to running requests - self._running_requests.append(request) - - # Clear preempted flag if set - if was_preempted: - request._preempted = False - - scheduled.append(request) - num_tokens_list.append(num_new_tokens) - token_budget -= num_new_tokens - self._current_iteration_token_budget = token_budget - if is_final_decode_request: - self._decode_iteration_reserved_slots_remaining = max( - self._decode_iteration_reserved_slots_remaining - 1, - 0, - ) - - # Flow validation: log ADMISSION event (matching vLLM v1) - logger.info( - f"[ADMISSION] req={request.id} admitted, " - f"num_tokens={num_new_tokens}, " - f"running_count={len(self._running_requests)}, " - f"token_budget_remaining={token_budget}" - ) - available_blocks_admission = int( - self._config.num_blocks - self._num_allocated_blocks - ) - self._emit_schedule_decision_event( - event="decision", - decision_result="ADMISSION", - request_id=request.id, - token_budget=token_budget, - available_blocks=available_blocks_admission, - num_tokens=num_new_tokens, - ) - - # Flow validation: log preemption recovery if applicable - if was_preempted: - # For DECODE cluster, preempted requests need full recomputation - # from their original prefill tokens - recompute_tokens = request.num_prefill_tokens - logger.info( - f"[PREEMPTION_RECOVERY] req={request.id}, " - f"was_preempted=True, " - f"recompute_tokens={recompute_tokens}" - ) - - if skipped_waiting_requests: - waiting_queue.extend(skipped_waiting_requests) - self._waiting_requests = list(waiting_queue) - - return token_budget, scheduled, num_tokens_list - - def _should_use_dense_decode_attn_metadata_wave(self) -> bool: - """Return whether dense PP=1 PDAF needs one DES macro-wave batch.""" - if self._cluster_type != ClusterType.DECODE_ATTN: - return False - if self._replica_is_moe: - return False - return ( - self._num_stages == 1 - and self._af_pipeline_num_micro_batch > 1 - ) - - def _schedule_decode_attn_only( - self, is_micro_batch: bool = True - ) -> Optional[Batch]: - """ - Scheduling for DECODE_ATTN cluster in PD-AF disaggregation mode. - - This method is called ONLY for Priority 2 scheduling (new micro-batch formation). - Priority 1 (AF immediate inflight batches) is handled by on_schedule() directly. - - Two-level scheduling strategy based on decode step: - - Incomplete decode step (is_mb_last_layer=False): batch-level, via _af_immediate_batch_queue - - Complete decode step (is_mb_last_layer=True): request-level, via this method - - Phase 1: Schedule running requests (ongoing decode from _running_requests) - - For each request in _running_requests: - - Calculate new tokens to process (usually 1 for decode) - - Allocate memory for new tokens - - If allocation fails: trigger preemption following vLLM v1 behavior - - Add to scheduled batch - - Phase 2: Admit new requests from _waiting_requests (if Phase 1 had no preemption) - - Check memory budget and token budget - - Form micro-batch with layer-consistent grouping (fix: do we need it? all requests are layer-0) - - All new requests start at layer 0, so naturally layer-consistent - - Layer-consistent grouping is implicitly guaranteed: - - Running requests have _completed_layer_count = 0 (reset after decode step completion) - - New requests also start at layer 0 - - Therefore, all requests in a micro-batch are layer-consistent - - Note on initial state: - - On first scheduling, _running_requests is empty, so Phase 1 produces no output - - Phase 2 will admit new requests from _waiting_requests to _running_requests - - Subsequent decode steps will have Phase 1 populated from previous on_batch_end() - - Args: - is_micro_batch: Should always be True for DECODE_ATTN - - Returns: - Optional[Batch]: The scheduled micro-batch, or None if no requests available - """ - logger = get_cluster_logger(__name__, self._cluster_type.name) - - if is_micro_batch and self._af_pending_micro_batches: - return self._af_pending_micro_batches.popleft() - - # Enable preemption for DECODE_ATTN to handle memory pressure - # Preemption logic follows vLLM v1 behavior for running requests - preemption_enabled = True - preempted_requests: List[Request] = [] - - # Phase 1: Schedule running requests - scheduled_requests = [] - scheduled_tokens = [] - - # Get request IDs to exclude (already scheduled in inflight batches) - continuation_request_ids = getattr(self, "_continuation_request_ids", set()) - # _running_requests in inclued reqs: inflight(layer!=0) req, completed req (really?) - - for request in self._running_requests: - # ISSUE-008 FIX: Check batch size limit at start of Phase 1 loop. - # This prevents scheduling more requests than _micro_batch_size allows, - # ensuring proper batch size enforcement in DECODE_ATTN cluster. - if len(scheduled_requests) >= self._micro_batch_size: - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 1: reached micro_batch_size limit " - f"({self._micro_batch_size}), stopping" - ) - break - - if request.completed: - # why would a running request be completed but still in _running_requests? - raise ValueError(f"Request {request.id} is already completed") - continue - - # CRITICAL FIX: Only schedule requests ready for new decode step (layer_count = 0) - # Requests with layer_count > 0 are still in-flight (mid-layer processing) - # and should NOT be re-scheduled until their current decode step completes. - # This ensures layer-consistent grouping in micro-batches. - if request.completed_layer_count != 0: - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 1: skipping in-flight req={request.id} " - f"with layer_count={request.completed_layer_count} (not ready for new decode step)" - ) - continue - - # Requests in active A->F->A roundtrip must not be re-scheduled until - # F->A transfer end clears the in-flight marker. - if request.af_roundtrip_inflight: - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 1: skipping req={request.id} " - f"(AF roundtrip still in-flight)" - ) - continue - - # CRITICAL FIX: Skip requests already scheduled in continuation batches (Priority 1) - # This prevents the same request from being scheduled into multiple batches - if request.id in continuation_request_ids: - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 1: skipping req={request.id} " - f"(already in continuation batch from Priority 1)" - ) - continue - - # Calculate tokens for decode: usually 1 - num_new_tokens = 1 - - # Try to allocate memory - if self._can_allocate_request(request, num_new_tokens): - self._allocate_request(request, num_new_tokens) - scheduled_requests.append(request) - scheduled_tokens.append(num_new_tokens) - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 1: scheduled running req={request.id}, " - f"num_tokens={num_new_tokens}" - ) - else: - # Memory pressure - try allocation with preemption - if not preemption_enabled: - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 1: cannot allocate req={request.id}, " - f"preemption disabled, skipping" - ) - continue - - # Try to allocate with preemption (follows vLLM v1 behavior) - preempted_count_before = len(preempted_requests) - success = self._try_allocate_with_preemption( - request, num_new_tokens, preempted_requests - ) - self._current_iteration_token_budget = ( - self._rollback_current_iteration_preempted_requests( - scheduled_requests=scheduled_requests, - scheduled_num_tokens=scheduled_tokens, - newly_preempted_requests=preempted_requests[ - preempted_count_before: - ], - token_budget=self._current_iteration_token_budget, - ) - ) - if success: - scheduled_requests.append(request) - scheduled_tokens.append(num_new_tokens) - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 1: scheduled req={request.id} " - f"after preemption, num_tokens={num_new_tokens}" - ) - else: - # Request itself was preempted or no victim available - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 1: req={request.id} " - f"preempted or allocation failed" - ) - - # Check micro-batch size limit - remaining_slots = self._micro_batch_size - len(scheduled_requests) - - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] After Phase 1: scheduled={len(scheduled_requests)}, " - f"remaining_slots={remaining_slots}, micro_batch_size={self._micro_batch_size}" - ) - - # Phase 2: Admit new requests (only if no preemption occurred) - if len(preempted_requests) == 0 and remaining_slots > 0: - for request in list(self._waiting_requests): - if remaining_slots <= 0: - break - - # New requests start at layer 0 - naturally layer-consistent - assert request.completed_layer_count == 0, ( - f"New request {request.id} should have completed_layer_count=0, got {request.completed_layer_count}" - ) - - # Allocate decode token - num_tokens = 1 - if self._can_allocate_request(request, num_tokens): - self._waiting_requests.remove(request) - request.on_leave_waiting_queue( - self._current_schedule_time, self._cluster_type - ) - self._allocate_request(request, num_tokens) - self._running_requests.append(request) - scheduled_requests.append(request) - scheduled_tokens.append(num_tokens) - remaining_slots -= 1 - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 2: admitted new req={request.id}, " - f"num_tokens={num_tokens}, running_count={len(self._running_requests)}" - ) - else: - logger.debug( - f"[VLLMv1Engine][DECODE_ATTN] Phase 2: cannot allocate req={request.id}, " - f"stopping admission" - ) - break - - # (scheduled_requests, scheduled_tokens) is the scheduler's output - # we should use scheduler_output to creat microbatch for pd-af - - # Create batch if we have scheduled requests - if scheduled_requests: - logger.info( - f"[VLLMv1Engine][DECODE_ATTN] Created micro-batch with {len(scheduled_requests)} requests" - ) - - num_reqs = len(scheduled_requests) - num_stages = self._af_pipeline_num_micro_batch - if num_stages is None or num_stages <= 0: - raise ValueError( - "af_pipeline_num_micro_batch must be positive for DECODE_ATTN" - ) - - replay_decode_token_index = int( - scheduled_requests[0].current_decode_token_index - ) - decode_attn_cohort_id = self._allocate_decode_attn_cohort_id() - decode_attn_cohort_request_ids = tuple( - request.id for request in scheduled_requests - ) - cohort_state = self._get_decode_attn_active_cohort_states().setdefault( - decode_attn_cohort_id, - { - "all_request_ids": set(), - "pending_request_ids": set(), - "af_phase": "local_attn", - "active_stage_indices": set(), - "stage_phases": {}, - "stage_current_layer_ids": {}, - }, - ) - cohort_state["all_request_ids"].update( - decode_attn_cohort_request_ids - ) - cohort_state["pending_request_ids"].update( - decode_attn_cohort_request_ids - ) - cohort_state["current_layer_id"] = int( - scheduled_requests[0].completed_layer_count - ) - - # StepFun-vLLM partitioning: split requests by stage - if num_reqs >= num_stages: - num_reqs_per_stage = num_reqs // num_stages - stage_reqs_start_loc = [ - num_reqs_per_stage * i for i in range(num_stages + 1) - ] - stage_reqs_start_loc[-1] = num_reqs - else: - stage_reqs_start_loc = list(range(num_reqs + 1)) - - afd_stage_metadata = None - if self._cluster_type == ClusterType.DECODE_ATTN and num_stages > 0: - from frontier.config import global_vars - from frontier.entities.batch import AFDStageMetadata - - use_cuda_graph = global_vars.get_use_cuda_graph() - cudagraph_capture_sizes = global_vars.get_cudagraph_capture_sizes() - if use_cuda_graph and cudagraph_capture_sizes is None: - max_num_seqs = ( - self._micro_batch_size - if hasattr(self, "_micro_batch_size") - else 64 - ) - cudagraph_capture_sizes = [1, 2, 4] + [ - 8 * i for i in range(1, max_num_seqs // 8 + 1) - ] - - afd_stage_metadata = AFDStageMetadata.from_batch_params( - num_reqs=num_reqs, - num_tokens_per_req=scheduled_tokens, - num_stages=num_stages, - dp_stage_max_tokens=None, - use_cuda_graph=use_cuda_graph, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ffn_use_cuda_graph=use_cuda_graph, - ffn_cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - - first_micro_batch = None - if self._should_use_dense_decode_attn_metadata_wave(): - macro_batch = self._create_batch( - scheduled_requests, - scheduled_tokens, - ) - macro_batch.afd_stage_idx = 0 - macro_batch.afd_stage_represents_all_stages = True - macro_batch.replay_decode_token_index = replay_decode_token_index - macro_batch.decode_attn_cohort_id = decode_attn_cohort_id - macro_batch.decode_attn_cohort_request_ids = ( - decode_attn_cohort_request_ids - ) - if afd_stage_metadata is not None: - macro_batch.afd_stage_metadata = afd_stage_metadata - cohort_state["active_stage_indices"].add(0) - cohort_state["stage_phases"][0] = "local_attn" - cohort_state["stage_current_layer_ids"][0] = int( - scheduled_requests[0].completed_layer_count - ) - return macro_batch - - shared_decode_attn_global_id = int(self._batch_creation_counter) - for stage_idx in range(len(stage_reqs_start_loc) - 1): - start_idx = stage_reqs_start_loc[stage_idx] - end_idx = stage_reqs_start_loc[stage_idx + 1] - stage_requests = scheduled_requests[start_idx:end_idx] - stage_tokens = scheduled_tokens[start_idx:end_idx] - micro_batch = self._create_batch(stage_requests, stage_tokens) - micro_batch.set_global_id(shared_decode_attn_global_id) - micro_batch.afd_stage_idx = stage_idx - micro_batch.replay_decode_token_index = int( - stage_requests[0].current_decode_token_index - ) - micro_batch.decode_attn_cohort_id = decode_attn_cohort_id - micro_batch.decode_attn_cohort_request_ids = ( - decode_attn_cohort_request_ids - ) - if afd_stage_metadata is not None: - micro_batch.afd_stage_metadata = afd_stage_metadata - normalized_stage_idx = int(stage_idx) - cohort_state["active_stage_indices"].add(normalized_stage_idx) - cohort_state["stage_phases"][normalized_stage_idx] = "local_attn" - cohort_state["stage_current_layer_ids"][normalized_stage_idx] = int( - scheduled_requests[0].completed_layer_count - ) - if first_micro_batch is None: - first_micro_batch = micro_batch - else: - self._af_pending_micro_batches.append(micro_batch) - - return first_micro_batch - - - logger.debug("[VLLMv1Engine][DECODE_ATTN] No requests to schedule") - return None - - def _attach_afd_metadata_if_needed(self, batch: Batch) -> Batch: - """Attach AFD stage metadata to batch if num_stages > 1. - - This method generates AFDStageMetadata following StepFun-vLLM's three-layer - padding strategy. The metadata is used for compute time and communication - volume prediction. - - Note: DP padding (Layer 2) requires per-stage max tokens across DP ranks. - In simulator, we can compute this at cluster scheduler level where all DP - lanes are visible. For now, we skip DP padding here and let cluster scheduler - handle it when aggregating batches. - - Args: - batch: The batch to attach metadata to - - Returns: - The batch with afd_stage_metadata attached (if applicable) - """ - from frontier.entities.batch import AFDStageMetadata - - # Get num_stages from cluster config (af_pipeline_num_micro_batch) - num_stages = self._af_pipeline_num_micro_batch - if num_stages is None or num_stages <= 0: - return batch - - if self._cluster_type != ClusterType.DECODE_ATTN: - return batch - - # Get CUDA Graph configuration from global simulation config - from frontier.config import global_vars - - use_cuda_graph = global_vars.get_use_cuda_graph() - cudagraph_capture_sizes = global_vars.get_cudagraph_capture_sizes() - - # Generate default capture sizes if not specified - # Aligned with StepFun-vLLM's default: [1, 2, 4] + [8 * i for i in range(1, max_num_seqs // 8 + 1)] - if use_cuda_graph and cudagraph_capture_sizes is None: - max_num_seqs = self._micro_batch_size if hasattr(self, "_micro_batch_size") else 64 - cudagraph_capture_sizes = [1, 2, 4] + [ - 8 * i for i in range(1, max_num_seqs // 8 + 1) - ] - - # Compute AFD metadata - # Note: dp_stage_max_tokens is None here - cluster scheduler will handle DP padding - batch.afd_stage_metadata = AFDStageMetadata.from_batch_params( - num_reqs=len(batch.requests), - num_tokens_per_req=batch.num_tokens, - num_stages=num_stages, - dp_stage_max_tokens=None, # Cluster scheduler handles DP padding - use_cuda_graph=use_cuda_graph, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ffn_use_cuda_graph=use_cuda_graph, - ffn_cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - - logger = get_cluster_logger(__name__, self._cluster_type.name) - logger.debug( - f"[AFD-METADATA] batch={batch.id} " - f"num_stages={num_stages} " - f"original_tokens={batch.afd_stage_metadata.original_total_tokens} " - f"padded_tokens={batch.afd_stage_metadata.padded_total_tokens} " - f"padding_overhead={batch.afd_stage_metadata.num_pad_tokens}" - ) - - return batch - - # ========== Property Overrides ========== - @property + def num_pending_requests(self) -> int: """ Return total schedulable pending requests for this cluster. diff --git a/frontier/scheduler/replica_scheduler/vllm_v1_iteration_policy.py b/frontier/scheduler/replica_scheduler/vllm_v1_iteration_policy.py new file mode 100644 index 00000000..731251fa --- /dev/null +++ b/frontier/scheduler/replica_scheduler/vllm_v1_iteration_policy.py @@ -0,0 +1,597 @@ +"""Per-iteration scheduling policy, fast lanes, CUDA graph and spec-decode metadata. + +These methods decide the shape of one scheduling iteration: which policy orders +the waiting queue, whether a request takes a final-round fast lane, which CUDA +graph capture size the decode batch maps onto, and what speculative-decoding +metadata the batch carries. +""" + +from collections import deque +import time +from typing import Any, Dict, List, Optional, Tuple + +from frontier.config import global_vars +from frontier.entities.batch import ( + Batch, + DecodeCudaGraphMetadata, + Request, + SpecDecodeBatchMetadata, +) +from frontier.logger import get_cluster_logger +from frontier.scheduler.replica_scheduler.vllm_v1_decision_log import ( + _log_frontier_vllm_v1_schedule_decision, + schedule_decision_logging_enabled, +) +from frontier.spec_decode import compute_iteration_outcome, get_planned_draft_tokens +from frontier.types import ClusterType + + +class IterationSchedulingPolicy: + """Iteration ordering, fast lanes, CUDA graph sizing and batch metadata.""" + + def _build_decode_cuda_graph_metadata( + self, batch: Batch + ) -> Optional[DecodeCudaGraphMetadata]: + if self._cluster_type not in (ClusterType.MONOLITHIC, ClusterType.DECODE): + return None + if ( + getattr(self, "_spec_decode_enabled", False) + and not global_vars.get_allow_spec_decode_cuda_graph_diagnostic() + ): + # Phase 2+ baseline: speculative decoding always runs in eager mode. + # We intentionally disable decode CUDA graph modeling for all + # speculative batches to reduce alignment complexity; future work + # can reintroduce method-specific CUDA graph semantics. + return None + + config_mode = global_vars.get_decode_cuda_graph_mode() + if config_mode == "none": + return None + + capture_hit, capture_size = self._resolve_decode_cuda_graph_capture_size( + batch.total_num_tokens + ) + decode_query_lens = [ + int(num_tokens) + for request, num_tokens in zip(batch.requests, batch.num_tokens) + if request.is_prefill_complete + ] + original_decode_batch_size = len(decode_query_lens) + + # Align with vLLM's uniform_decode_query_len semantics. + # When speculative decoding is enabled, FULL decode cudagraphs are only + # valid for uniform batches whose query_len matches + # 1 + num_speculative_tokens. Non-uniform speculative verify batches + # must dispatch to mixed/piecewise graphs or fall back to eager. + uniform_decode_query_len = 1 + if getattr(self, "_spec_decode_enabled", False): + spec_decode_config = getattr(self, "_spec_decode_config", None) + if spec_decode_config is None: + raise ValueError("Speculative decoding config is not initialized") + uniform_decode_query_len += int(spec_decode_config.num_speculative_tokens) + + is_uniform_decode_batch = ( + bool(decode_query_lens) + and len(decode_query_lens) == len(batch.requests) + and all( + query_len == uniform_decode_query_len + for query_len in decode_query_lens + ) + ) + is_mixed_batch = not is_uniform_decode_batch + + runtime_mode = "NONE" + if config_mode == "full_decode_only": + if is_uniform_decode_batch and capture_hit: + runtime_mode = "FULL" + elif config_mode == "piecewise" and capture_hit: + runtime_mode = "PIECEWISE" + + if runtime_mode == "NONE": + capture_hit = False + capture_size = batch.total_num_tokens + + padded_decode_batch_size = ( + capture_size if capture_hit else original_decode_batch_size + ) + padded_total_tokens = capture_size if capture_hit else batch.total_num_tokens + + return DecodeCudaGraphMetadata( + config_mode=config_mode, + runtime_mode=runtime_mode, + capture_hit=capture_hit, + is_mixed_batch=is_mixed_batch, + original_total_tokens=batch.total_num_tokens, + padded_total_tokens=padded_total_tokens, + original_decode_batch_size=original_decode_batch_size, + padded_decode_batch_size=padded_decode_batch_size, + ) + + def _resolve_decode_cuda_graph_capture_size(self, total_tokens: int) -> Tuple[bool, int]: + cudagraph_capture_sizes = global_vars.get_cudagraph_capture_sizes() + if cudagraph_capture_sizes is None: + max_num_seqs = getattr( + self, + "_max_num_running_reqs", + getattr(self, "_max_batch_size", total_tokens), + ) + max_num_seqs = max(int(max_num_seqs), total_tokens) + cudagraph_capture_sizes = [1, 2, 4] + [ + 8 * i for i in range(1, max_num_seqs // 8 + 1) + ] + + for capture_size in sorted(cudagraph_capture_sizes): + if total_tokens <= capture_size: + return True, int(capture_size) + return False, int(total_tokens) + + def _build_spec_decode_batch_metadata( + self, batch: Batch + ) -> Optional[SpecDecodeBatchMetadata]: + if not getattr(self, "_spec_decode_enabled", False): + return None + if self._cluster_type not in (ClusterType.MONOLITHIC, ClusterType.DECODE): + return None + if batch.num_decode_tokens <= 0: + return None + spec_decode_config = getattr(self, "_spec_decode_config", None) + if spec_decode_config is None: + raise ValueError("Speculative decoding config is not initialized") + + planned_drafts_list: List[int] = [] + verify_tokens_list: List[int] = [] + accepted_drafts_list: List[int] = [] + rejected_drafts_list: List[int] = [] + committed_tokens_list: List[int] = [] + terminal_planned_drafts_list: List[List[int]] = [] + terminal_verify_tokens_list: List[List[int]] = [] + terminal_accepted_drafts_list: List[List[int]] = [] + terminal_rejected_drafts_list: List[List[int]] = [] + terminal_raw_committed_tokens_list: List[List[int]] = [] + per_request_outcomes: Dict[int, Tuple[int, Any, List[Tuple[int, int, int, int, int]]]] = {} + + for request, scheduled_tokens in zip(batch.requests, batch.num_tokens): + if not getattr(request, "is_prefill_complete", False) or not getattr( + request, "spec_decode_enabled", False + ): + planned_drafts_list.append(0) + verify_tokens_list.append(0) + accepted_drafts_list.append(0) + rejected_drafts_list.append(0) + committed_tokens_list.append(int(scheduled_tokens)) + terminal_planned_drafts_list.append([]) + terminal_verify_tokens_list.append([]) + terminal_accepted_drafts_list.append([]) + terminal_rejected_drafts_list.append([]) + terminal_raw_committed_tokens_list.append([]) + continue + + request_id = int(request.id) + scheduled_tokens_int = int(scheduled_tokens) + if request_id in per_request_outcomes: + ( + recorded_scheduled_tokens, + recorded_outcome, + recorded_terminal_rows, + ) = per_request_outcomes[request_id] + if recorded_scheduled_tokens != scheduled_tokens_int: + raise ValueError( + "Inconsistent scheduled_tokens for duplicated request in the " + "same batch: " + f"request_id={request_id}, " + f"first={recorded_scheduled_tokens}, " + f"current={scheduled_tokens_int}" + ) + outcome = recorded_outcome + terminal_rows = recorded_terminal_rows + else: + if getattr(request, "spec_method_is_target_embedded_mtp", False): + planned_drafts = int(request.spec_next_planned_draft_tokens) + else: + planned_drafts = max(scheduled_tokens_int - 1, 0) + remaining_decode = request.remaining_decode_tokens + outcome = compute_iteration_outcome( + spec_decode_config, + remaining_decode, + planned_draft_tokens=planned_drafts, + iteration_index=request.spec_total_iterations, + request_id=str(request.id), + ) + request.record_spec_decode_iteration( + verify_tokens=outcome.verify_tokens, + accepted_drafts=outcome.accepted_draft_tokens, + rejected_drafts=outcome.rejected_draft_tokens, + committed_tokens=outcome.committed_tokens, + ) + + next_remaining_decode = max( + remaining_decode - outcome.committed_tokens, + 0, + ) + terminal_rows: List[Tuple[int, int, int, int, int]] = [] + if next_remaining_decode == 0: + terminal_rows = ( + self._get_target_embedded_mtp_terminal_overshoot_rows( + request, + start_iteration_index=request.spec_total_iterations, + ) + ) + request.set_spec_next_planned_draft_tokens( + get_planned_draft_tokens( + spec_decode_config, + next_remaining_decode, + iteration_index=request.spec_total_iterations, + request_id=str(request.id), + ) + ) + per_request_outcomes[request_id] = ( + scheduled_tokens_int, + outcome, + terminal_rows, + ) + + planned_drafts_list.append(outcome.planned_draft_tokens) + verify_tokens_list.append(outcome.verify_tokens) + accepted_drafts_list.append(outcome.accepted_draft_tokens) + rejected_drafts_list.append(outcome.rejected_draft_tokens) + committed_tokens_list.append(outcome.committed_tokens) + terminal_planned_drafts_list.append( + [int(row[0]) for row in terminal_rows] + ) + terminal_verify_tokens_list.append([int(row[1]) for row in terminal_rows]) + terminal_accepted_drafts_list.append( + [int(row[2]) for row in terminal_rows] + ) + terminal_rejected_drafts_list.append( + [int(row[3]) for row in terminal_rows] + ) + terminal_raw_committed_tokens_list.append( + [int(row[4]) for row in terminal_rows] + ) + + metadata = SpecDecodeBatchMetadata( + method=spec_decode_config.method, + planned_draft_tokens_per_request=planned_drafts_list, + verify_tokens_per_request=verify_tokens_list, + accepted_draft_tokens_per_request=accepted_drafts_list, + rejected_draft_tokens_per_request=rejected_drafts_list, + committed_tokens_per_request=committed_tokens_list, + uses_lookahead_slots=getattr( + self, "_spec_method_uses_lookahead_slots", False + ), + terminal_overshoot_planned_draft_tokens_per_request=( + terminal_planned_drafts_list + ), + terminal_overshoot_verify_tokens_per_request=( + terminal_verify_tokens_list + ), + terminal_overshoot_accepted_draft_tokens_per_request=( + terminal_accepted_drafts_list + ), + terminal_overshoot_rejected_draft_tokens_per_request=( + terminal_rejected_drafts_list + ), + terminal_overshoot_raw_committed_tokens_per_request=( + terminal_raw_committed_tokens_list + ), + ) + metadata.validate(len(batch.requests)) + return metadata + + def _get_scheduling_policy(self) -> str: + """ + Get the scheduling policy to use. + + This method provides a clean interface for policy selection that can be + easily extended in future work to support command-line parameter control. + + Returns: + str: The scheduling policy ('fcfs' or 'priority') + """ + # Use the scheduling policy from configuration + return self._config.scheduling_policy + + def _get_iteration_phase_aware_waiting_requests(self) -> List[Request]: + if self._cluster_type not in (ClusterType.MONOLITHIC, ClusterType.PREFILL): + return [] + return list(self._preempted_requests) + list(self._request_queue) + + def _resolve_iteration_round_class(self) -> Optional[str]: + if not self._enable_phase_aware_thinking_profile: + return None + + waiting_requests = self._get_iteration_phase_aware_waiting_requests() + thinking_requests = [ + request + for request in waiting_requests + if getattr(request, "is_thinking_mode_enabled", False) + ] + if not thinking_requests: + return None + if any(request.is_final_thinking_round for request in thinking_requests): + return "final" + return "hidden" + + def _get_iteration_scheduler_profile(self) -> Dict[str, Any]: + round_class = self._resolve_iteration_round_class() + profile = { + "round_class": round_class, + "max_num_running_reqs": int(self._config.batch_size_cap), + "max_num_scheduled_tokens": int(self._config.max_tokens_in_batch), + "enable_chunked_prefill": bool( + getattr(self._config, "enable_chunked_prefill", False) + ), + } + if round_class is None: + return profile + + prefix = f"{round_class}_phase_" + max_tokens_override = getattr(self._config, f"{prefix}max_tokens_in_batch") + chunked_override = getattr( + self._config, f"{prefix}enable_chunked_prefill" + ) + batch_size_override = getattr(self._config, f"{prefix}batch_size_cap") + if max_tokens_override is not None: + profile["max_num_scheduled_tokens"] = int(max_tokens_override) + if chunked_override is not None: + profile["enable_chunked_prefill"] = bool(chunked_override) + if batch_size_override is not None: + profile["max_num_running_reqs"] = int(batch_size_override) + return profile + + def _refresh_iteration_scheduler_profile(self) -> None: + profile = self._get_iteration_scheduler_profile() + self._active_iteration_round_class = profile["round_class"] + self._max_num_running_reqs = int(profile["max_num_running_reqs"]) + self._max_num_scheduled_tokens = int(profile["max_num_scheduled_tokens"]) + self._enable_chunked_prefill = bool(profile["enable_chunked_prefill"]) + + def _maybe_promote_final_round_priority(self, request: Request) -> None: + if not self._enable_final_round_priority_boost: + return + if not getattr(request, "is_thinking_mode_enabled", False): + return + if not request.is_final_thinking_round or request.completed_thinking_rounds <= 0: + return + request.set_priority(min(request.priority, self._final_round_priority_value)) + + def _is_final_prefill_fast_lane_request(self, request: Request) -> bool: + return bool( + getattr(request, "is_thinking_mode_enabled", False) + and request.is_final_thinking_round + and not request.is_prefill_complete + ) + + def _is_final_decode_fast_lane_request(self, request: Request) -> bool: + return bool( + getattr(request, "is_thinking_mode_enabled", False) + and request.is_final_thinking_round + and request.is_prefill_complete + and not request.completed + ) + + def _ordered_requests_with_final_lane( + self, + requests: List[Request], + *, + final_predicate, + ) -> List[Request]: + final_requests: List[Request] = [] + non_final_requests: List[Request] = [] + for request in requests: + if final_predicate(request): + final_requests.append(request) + else: + non_final_requests.append(request) + if not final_requests: + return requests + return final_requests + non_final_requests + + def _build_prefill_waiting_queue(self) -> deque[Request]: + ordered_requests = self._get_sorted_waiting_queue() + if ( + self._cluster_type == ClusterType.PREFILL + and ( + self._final_prefill_reserved_slots > 0 + or self._final_prefill_reserved_tokens > 0 + ) + ): + ordered_requests = self._ordered_requests_with_final_lane( + ordered_requests, + final_predicate=self._is_final_prefill_fast_lane_request, + ) + return deque(ordered_requests) + + def _build_decode_waiting_queue(self) -> deque[Request]: + ordered_requests = list(self._waiting_requests) + if getattr( + getattr(self, "_config", None), "enable_thinking_round_priority", False + ): + ordered_requests.sort( + key=lambda r: ( + 0 if r.is_final_thinking_round else 1, + r.priority, + r.arrived_at, + ) + ) + elif self._scheduling_policy == "priority": + ordered_requests.sort(key=lambda r: (r.priority, r.arrived_at)) + + if ( + self._cluster_type == ClusterType.DECODE + and self._final_decode_reserved_slots > 0 + ): + ordered_requests = self._ordered_requests_with_final_lane( + ordered_requests, + final_predicate=self._is_final_decode_fast_lane_request, + ) + return deque(ordered_requests) + + def _count_final_fast_lane_requests( + self, + requests: List[Request] | deque[Request], + *, + final_predicate, + ) -> int: + return sum(1 for request in requests if final_predicate(request)) + + def _select_final_running_reclaim_victim( + self, + *, + final_predicate, + ) -> Optional[Request]: + candidates = [ + request + for request in self._running_requests + if not final_predicate(request) + ] + if not candidates: + return None + if self._scheduling_policy == "priority": + return max(candidates, key=lambda r: (r.priority, r.arrived_at)) + return candidates[-1] + + def _reclaim_borrowed_final_running_slots( + self, + *, + waiting_requests: List[Request] | deque[Request], + final_predicate, + reserved_slots: int, + lane_name: str, + ) -> List[Request]: + if not self._enable_final_running_request_reclaim or reserved_slots <= 0: + return [] + + final_waiting_count = self._count_final_fast_lane_requests( + waiting_requests, + final_predicate=final_predicate, + ) + if final_waiting_count <= 0: + return [] + + final_running_count = self._count_final_fast_lane_requests( + self._running_requests, + final_predicate=final_predicate, + ) + remaining_reserved_slots = max( + reserved_slots - min(final_running_count, reserved_slots), + 0, + ) + target_new_final_admissions = min( + final_waiting_count, + remaining_reserved_slots, + ) + if target_new_final_admissions <= 0: + return [] + + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + reclaimed_requests: List[Request] = [] + while ( + max(self._max_num_running_reqs - len(self._running_requests), 0) + < target_new_final_admissions + ): + victim = self._select_final_running_reclaim_victim( + final_predicate=final_predicate, + ) + if victim is None: + logger.info( + "[FINAL-SLICE-RECLAIM] lane=%s stopped_without_victim " + "target_new_final_admissions=%s running_count=%s", + lane_name, + target_new_final_admissions, + len(self._running_requests), + ) + break + logger.info( + "[FINAL-SLICE-RECLAIM] lane=%s reclaiming_hidden_req=%s " + "target_new_final_admissions=%s running_count_before=%s", + lane_name, + victim.id, + target_new_final_admissions, + len(self._running_requests), + ) + self._preempt_request(victim, reclaimed_requests) + + if reclaimed_requests: + logger.info( + "[FINAL-SLICE-RECLAIM] lane=%s reclaimed_count=%s " + "running_count_after=%s", + lane_name, + len(reclaimed_requests), + len(self._running_requests), + ) + return reclaimed_requests + + def _get_num_waiting_reqs_for_decision_log(self) -> int: + if self._cluster_type in (ClusterType.DECODE, ClusterType.DECODE_ATTN): + return len(self._waiting_requests) + return len(self._request_queue) + len(self._preempted_requests) + + def _apply_long_prefill_token_threshold( + self, request: Request, num_new_tokens: int + ) -> int: + """Apply long prefill threshold only for prefill-phase requests.""" + if request.is_prefill_complete or self._long_prefill_token_threshold <= 0: + return num_new_tokens + return min(num_new_tokens, self._long_prefill_token_threshold) + + def _emit_schedule_decision_event( + self, + *, + event: str, + decision_result: Optional[str], + request_id: Optional[int], + token_budget: int, + num_tokens: int, + available_blocks: Optional[int] = None, + batch_request_ids: Optional[List[int]] = None, + request_num_tokens: Optional[List[int]] = None, + batch_size: int = 0, + batch_num_tokens: int = 0, + ) -> None: + if not schedule_decision_logging_enabled(): + return + + if available_blocks is None: + available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) + + cluster_name = self._cluster_type.name if self._cluster_type else "MONOLITHIC" + payload: Dict[str, Any] = { + "event": event, + "source": "frontier", + "scheduler": "vllm_v1", + "cluster_type": cluster_name, + "iteration_id": int(self._active_schedule_iteration_id), + "decision_result": decision_result, + "request_id": None if request_id is None else str(request_id), + "token_budget": int(token_budget), + "available_blocks": int(available_blocks), + "num_tokens": int(num_tokens), + "num_running_reqs": len(self._running_requests), + "num_waiting_reqs": self._get_num_waiting_reqs_for_decision_log(), + "max_num_running_reqs": int(self._max_num_running_reqs), + "max_num_scheduled_tokens": int(self._max_num_scheduled_tokens), + "batch_request_ids": [str(req_id) for req_id in (batch_request_ids or [])], + "request_num_tokens": [int(v) for v in (request_num_tokens or [])], + "batch_size": int(batch_size), + "batch_num_tokens": int(batch_num_tokens), + "timestamp": time.time(), + "timestamp_semantics": "wall_clock_epoch_seconds", + "simulation_time": float(self._current_schedule_time), + "simulation_time_semantics": "frontier_event_time_seconds", + } + if self._kv_cache_manager is not None: + prefix_cache_stats = self._kv_cache_manager.prefix_cache_stats + payload.update( + { + "prefix_cache_metric_semantics": "block_level", + "prefix_cache_unit": "blocks", + "prefix_cache_block_size": int(self._config.block_size), + "prefix_cache_requests": int(prefix_cache_stats.requests), + "prefix_cache_queries": int(prefix_cache_stats.queries), + "prefix_cache_hits": int(prefix_cache_stats.hits), + } + ) + _log_frontier_vllm_v1_schedule_decision(payload) diff --git a/frontier/scheduler/replica_scheduler/vllm_v1_kv_allocation.py b/frontier/scheduler/replica_scheduler/vllm_v1_kv_allocation.py new file mode 100644 index 00000000..bfc2aeba --- /dev/null +++ b/frontier/scheduler/replica_scheduler/vllm_v1_kv_allocation.py @@ -0,0 +1,713 @@ +"""KV block accounting, allocation and preemption for the vLLM V1 scheduler. + +This is the memory side of admission: how many tokens a request has already had +accounted, how many blocks its next step needs, whether those blocks are +available, and which request to preempt when they are not. +""" + +from math import ceil +from typing import List, Optional + +from frontier.attention.gdn.guards import validate_gdn_runtime_support +from frontier.entities.batch import Request +from frontier.kv_cache.base_kv_cache_manager import KVCacheAllocationResult +from frontier.logger import get_cluster_logger +from frontier.scheduler.replica_scheduler.vllm_v1_prefix_cache import ( + PrefixCacheAdmission, +) +from frontier.types import ClusterType + + +_REQUEST_PROGRESS_PRESERVING_PREEMPTION_CLUSTER_TYPES = frozenset( + { + ClusterType.DECODE, + ClusterType.DECODE_ATTN, + } +) + + +class KvBlockAllocation: + """Token accounting, KV block allocation and preemption.""" + + def _find_request_by_id(self, request_id: int) -> Optional[Request]: + request_groups = [ + getattr(self, "_running_requests", []), + getattr(self, "_request_queue", []), + getattr(self, "_preempted_requests", []), + getattr(self, "_waiting_requests", []), + ] + for requests in request_groups: + for request in requests: + if request.id == request_id: + return request + return None + + def _free_request_resources(self, request: Request) -> None: + self._get_monolithic_pp_mtp_near_full_prefill_request_ids().discard( + request.id + ) + self._get_monolithic_pp_mtp_single_output_wait_request_ids().discard( + request.id + ) + self._get_monolithic_pp_mtp_fractional_output_wait_counts().pop( + request.id, + None, + ) + self._get_monolithic_pp_mtp_output_wait_remaining_iters().pop( + request.id, + None, + ) + self._get_monolithic_pp_mtp_output_wait_request_ids().discard(request.id) + self._get_monolithic_pp_waiting_admission_delay_iters().pop( + request.id, + None, + ) + if self._is_prefix_caching_enabled(): + assert self._kv_cache_manager is not None + self._kv_cache_manager.free(request) + self._allocation_map.pop(request.id, None) + self._sync_prefix_cache_allocation_state() + gdn_slot_manager = self._gdn_state_slot_manager + if gdn_slot_manager is not None: + gdn_slot_manager.release(request.id) + return + self.free(request.id) + gdn_slot_manager = self._gdn_state_slot_manager + if gdn_slot_manager is not None: + gdn_slot_manager.release(request.id) + + def _free_request_resources_by_id(self, request_id: int) -> None: + request = self._find_request_by_id(request_id) + if request is not None: + self._free_request_resources(request) + return + self.free(request_id) + # Completion/cancellation callbacks may arrive after the request has + # left every scheduler queue. Release an orphaned ownership token as + # part of the same idempotent cleanup boundary so a slot cannot leak. + gdn_slot_manager = self._gdn_state_slot_manager + if gdn_slot_manager is not None: + gdn_slot_manager.release(request_id) + + def _get_explicit_scheduler_num_computed_tokens( + self, request: Request + ) -> Optional[int]: + scheduled_frontier = getattr( + self, "_scheduled_num_computed_tokens_by_request", {} + ).get(request.id) + if scheduled_frontier is None: + return None + return int(scheduled_frontier) + + def _get_scheduler_num_computed_tokens(self, request: Request) -> int: + """Return the scheduler-visible computed frontier for a request.""" + scheduled_frontier = self._get_explicit_scheduler_num_computed_tokens(request) + if scheduled_frontier is not None: + return scheduled_frontier + + processed_tokens = int(request.num_processed_tokens) + if ( + getattr(self, "_cluster_type", None) == ClusterType.MONOLITHIC + and request.is_prefill_complete + and processed_tokens > int(request.num_prefill_tokens) + ): + # MONOLITHIC request metrics grant the first decode token at the + # prefill-complete boundary, but vLLM's scheduler frontier does not + # advance to that token until the first decode scheduling step. + return max(int(request.num_prefill_tokens), processed_tokens - 1) + return processed_tokens + + def _advance_scheduler_num_computed_tokens( + self, request: Request, num_scheduled_tokens: int + ) -> None: + if num_scheduled_tokens < 0: + raise ValueError( + f"num_scheduled_tokens must be >= 0, got {num_scheduled_tokens}" + ) + self._scheduled_num_computed_tokens_by_request[request.id] = ( + self._get_scheduler_num_computed_tokens(request) + int(num_scheduled_tokens) + ) + + def _get_kv_accounted_processed_tokens(self, request: Request) -> int: + """Return processed tokens used for KV block accounting. + + In MONOLITHIC mode we intentionally count the first generated token at + prefill boundary for request-level progression parity. However, vLLM's + KV block growth does not advance at that boundary; it advances when the + first decode scheduling step is executed. To align block semantics, KV + accounting excludes that boundary token. + """ + explicit_scheduler_frontier = self._get_explicit_scheduler_num_computed_tokens( + request + ) + if getattr(self, "_cluster_type", None) != ClusterType.MONOLITHIC: + if explicit_scheduler_frontier is not None: + return explicit_scheduler_frontier + return int(request.num_processed_tokens) + if not getattr(request, "is_prefill_complete", False): + if explicit_scheduler_frontier is not None: + return explicit_scheduler_frontier + return int(request.num_processed_tokens) + processed_tokens = int(request.num_processed_tokens) + inflight_verify_tokens = 1 + if ( + getattr(request, "spec_decode_enabled", False) + and getattr(request, "spec_method_uses_lookahead_slots", False) + ): + inflight_verify_tokens = max( + 1, int(getattr(request, "spec_current_verify_tokens", 1)) + ) + decode_boundary_adjusted_tokens = max( + int(request.num_prefill_tokens), processed_tokens - inflight_verify_tokens + ) + if explicit_scheduler_frontier is None: + return decode_boundary_adjusted_tokens + return max(explicit_scheduler_frontier, decode_boundary_adjusted_tokens) + + def _get_request_next_num_tokens(self, request: Request) -> int: + assert not request.completed + + computed_tokens = self._get_scheduler_num_computed_tokens(request) + cluster_type = getattr(self, "_cluster_type", None) + + if request.is_prefill_complete: + if getattr(request, "spec_decode_enabled", False): + if getattr(request, "spec_method_is_target_embedded_mtp", False): + planned_drafts = int( + getattr(request, "spec_next_planned_draft_tokens", 0) + ) + if ( + cluster_type == ClusterType.MONOLITHIC + and int(getattr(request, "num_processed_decode_tokens", 0)) + == 1 + and computed_tokens <= int(request.num_prefill_tokens) + ): + return max(planned_drafts, 1) + return 1 + int(getattr(request, "spec_next_planned_draft_tokens", 0)) + if cluster_type == ClusterType.MONOLITHIC: + # In MONOLITHIC mode, request.num_processed_tokens includes the + # post-prefill decode bonus. A new decode step is schedulable + # only when request-side progress has advanced beyond the + # scheduler-visible frontier, mirroring vLLM's + # num_tokens_with_spec/num_computed_tokens gating under PP. + return max(int(request.num_processed_tokens) - computed_tokens, 0) + return 1 + + remaining_prefill_tokens = int(request.num_prefill_tokens) - computed_tokens + return max(remaining_prefill_tokens, 0) + + def _get_num_tokens_for_kv_reservation( + self, request: Request, scheduled_tokens: int + ) -> int: + reserved_tokens = int(scheduled_tokens) + if reserved_tokens <= 0: + raise ValueError( + f"scheduled_tokens must be > 0, got={scheduled_tokens}" + ) + if not getattr(request, "is_prefill_complete", False): + return reserved_tokens + if not getattr(request, "spec_decode_enabled", False): + return reserved_tokens + if getattr(request, "spec_method_uses_lookahead_slots", False): + return reserved_tokens + # ngram/medusa path: no lookahead slot reservation in Phase 1. + # Strategy A keeps allocation simple and lets future decode iterations + # amortize accepted-token KV growth without immediate draft-slot reserves. + return 1 + + def _get_initial_allocation_num_blocks( + self, request: Request, reserved_tokens: int + ) -> int: + """Return the block count used to check and commit a first allocation.""" + kv_accounted_tokens = self._get_kv_accounted_processed_tokens(request) + total_tokens = min( + kv_accounted_tokens + reserved_tokens, self._max_model_len + ) + return ceil(total_tokens / self._config.block_size) + + def _can_allocate_request( + self, + request: Request, + num_new_tokens: int = 1, + new_computed_blocks=None, + *, + scheduler_num_computed_tokens: Optional[int] = None, + ) -> bool: + """ + Check if memory can be allocated for a request. + + For new requests: check if prefill blocks can be allocated. + For running requests: check if at least one block is available. + + Args: + request: The request to check allocation for + num_new_tokens: Number of new tokens to allocate (used for decode) + + Returns: + bool: True if allocation is possible + """ + gdn_slot_manager = self._gdn_state_slot_manager + if ( + gdn_slot_manager is not None + and request.id not in self._allocation_map + and not gdn_slot_manager.has_slot(request.id) + and not gdn_slot_manager.has_available_slot + ): + return False + if self._is_prefix_caching_enabled(): + assert self._kv_cache_manager is not None + return self._kv_cache_manager.can_allocate_slots( + request, + num_new_tokens, + new_computed_blocks=new_computed_blocks, + scheduler_num_computed_tokens=( + self._get_scheduler_num_computed_tokens(request) + if scheduler_num_computed_tokens is None + else scheduler_num_computed_tokens + ), + ) + + reserved_tokens = self._get_num_tokens_for_kv_reservation( + request, num_new_tokens + ) + if request.id not in self._allocation_map: + # New request - estimate blocks from current token frontier + # (already processed + newly scheduled in this iteration), then + # clamp by max_model_len to keep allocation semantics consistent + # with chunked prefill scheduling. + num_required_blocks = self._get_initial_allocation_num_blocks( + request, reserved_tokens + ) + available_blocks = ( + self._config.num_blocks + - self._num_allocated_blocks + - num_required_blocks + ) + return available_blocks >= self._watermark_blocks + + # Running request - check if we need additional blocks for decode + num_tokens_reserved = self._allocation_map[request.id] * self._config.block_size + kv_accounted_tokens = self._get_kv_accounted_processed_tokens(request) + num_tokens_required = max( + 0, kv_accounted_tokens + reserved_tokens - num_tokens_reserved + ) + + if num_tokens_required <= 0: + return True + + # Need additional blocks + num_additional_blocks = ceil(num_tokens_required / self._config.block_size) + return self.can_allocate(num_additional_blocks) + + def _allocate_request( + self, + request: Request, + num_new_tokens: int = 1, + new_computed_blocks=None, + prefix_cache_admission: Optional[PrefixCacheAdmission] = None, + *, + scheduler_num_computed_tokens: Optional[int] = None, + ) -> Optional[KVCacheAllocationResult]: + """ + Allocate memory blocks for a request. + + For new requests: allocate blocks for prefill tokens + first decode token. + For running requests: allocate additional blocks if needed. + + Args: + request: The request to allocate for + num_new_tokens: Number of new tokens being processed + """ + if type(num_new_tokens) is not int or num_new_tokens <= 0: + raise ValueError( + "num_new_tokens must be a positive integer, " + f"got {num_new_tokens!r}" + ) + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + if self._is_prefix_caching_enabled(): + assert self._kv_cache_manager is not None + computed_blocks = list(new_computed_blocks or []) + if prefix_cache_admission is not None: + expected_blocks = prefix_cache_admission.effective_hit_blocks + if len(computed_blocks) != len(expected_blocks) or any( + actual is not expected + for actual, expected in zip(computed_blocks, expected_blocks) + ): + raise ValueError( + "Committed Prefix cache admission blocks differ from " + "the allocation input." + ) + if int(num_new_tokens) != int(prefix_cache_admission.num_new_tokens): + raise ValueError( + "Committed Prefix cache admission token count differs " + "from the allocation input." + ) + allocation = self._kv_cache_manager.allocate_slots( + request, + num_new_tokens, + new_computed_blocks=computed_blocks, + scheduler_num_computed_tokens=( + self._get_scheduler_num_computed_tokens(request) + if scheduler_num_computed_tokens is None + else scheduler_num_computed_tokens + ), + ) + if allocation is None: + raise ValueError( + f"Failed to allocate prefix-cache-managed KV blocks for request {request.id}" + ) + self._sync_prefix_cache_allocation_state(request) + self._emit_prefix_cache_identity_events( + request=request, + num_new_tokens=num_new_tokens, + allocation=allocation, + admission=prefix_cache_admission, + ) + return allocation + + reserved_tokens = self._get_num_tokens_for_kv_reservation( + request, num_new_tokens + ) + + if request.id not in self._allocation_map: + # Commit the exact block count used by the allocation preflight. + # This includes a transferred token frontier for a first + # DECODE_ATTN allocation after the PREFILL handoff. + num_required_blocks = self._get_initial_allocation_num_blocks( + request, reserved_tokens + ) + self.allocate(request.id, num_required_blocks) + gdn_slot_manager = self._gdn_state_slot_manager + if gdn_slot_manager is not None: + try: + gdn_slot_manager.allocate(request.id) + except Exception: + # KV and state ownership must commit atomically from the + # scheduler's perspective. Roll back the KV allocation + # before exposing the slot failure to admission. + self.free(request.id) + raise + logger.debug( + f"[VLLMv1Engine] Allocated {num_required_blocks} blocks for request {request.id} " + f"(scheduled_tokens={num_new_tokens}, reserved_tokens={reserved_tokens})" + ) + return None + + # Running request - check if additional blocks needed + gdn_slot_manager = self._gdn_state_slot_manager + if gdn_slot_manager is not None: + if not gdn_slot_manager.has_slot(request.id): + raise RuntimeError( + f"GDN state slot missing for admitted request {request.id}" + ) + gdn_slot_manager.resume(request.id) + num_tokens_reserved = self._allocation_map[request.id] * self._config.block_size + kv_accounted_tokens = self._get_kv_accounted_processed_tokens(request) + num_tokens_required = max( + 0, kv_accounted_tokens + reserved_tokens - num_tokens_reserved + ) + + if num_tokens_required <= 0: + return None + + # Allocate additional blocks + num_additional_blocks = ceil(num_tokens_required / self._config.block_size) + self.allocate(request.id, num_additional_blocks) + return None + + def _select_preemption_victim( + self, exclude: Optional[Request] = None + ) -> Optional[Request]: + """ + Select a victim request for preemption based on scheduling policy. + + FCFS policy: Preempt the most recently added request (queue tail). + Priority policy: Preempt the request with lowest priority + (highest priority value, then latest arrival). + + Args: + exclude: Optional request to exclude from victim selection (typically the requesting request) + + Returns: + Optional[Request]: The victim request, or None if no victims available + """ + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + + if not self._running_requests: + return None + + # Filter out excluded request + candidates = ( + [r for r in self._running_requests if r != exclude] + if exclude + else self._running_requests + ) + + if not candidates: + return None + + if self._scheduling_policy == "priority": + # Priority policy: preempt request with highest priority value (lowest priority) + # Tie-breaker: latest arrival time + victim = max(candidates, key=lambda r: (r.priority, r.arrived_at)) + + # Flow validation: log victim selection + logger.info( + f"[VICTIM_SELECTION] policy=PRIORITY, " + f"victim={victim.id}, " + f"priority={victim.priority}, " + f"reason=highest_priority_value" + ) + return victim + else: + # FCFS policy: preempt most recently added (queue tail) + # If exclude is specified and is the tail, select the second-to-last + if exclude and candidates and candidates[-1] != self._running_requests[-1]: + # exclude was the tail, use candidates[-1] which is second-to-last + victim = candidates[-1] + else: + victim = candidates[-1] if candidates else None + + if victim is None: + return None + + # Flow validation: log victim selection + logger.info( + f"[VICTIM_SELECTION] policy=FCFS, " + f"victim={victim.id}, " + f"position=tail, " + f"reason=last_in_running_queue" + ) + return victim + + def _preempt_request( + self, victim: Request, preempted_requests: List[Request] + ) -> None: + """ + Preempt a request - free its resources and move to waiting queue. + + Args: + victim: The request to preempt + preempted_requests: List to track preempted requests for this iteration + """ + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + + # GDN state cannot be dropped and restored by the simulator. Reject + # before touching request counters, allocations, or queue membership. + validate_gdn_runtime_support( + self._replica_config.model_config, + preemption_requires_state_drop=True, + ) + + # Capture state before modification + num_computed_tokens_before = self._get_scheduler_num_computed_tokens(victim) + freed_blocks = self._allocation_map.get(victim.id, 0) + running_count_before = len(self._running_requests) + queue_position_before = ( + self._running_requests.index(victim) + if victim in self._running_requests + else -1 + ) + + # Record preemption statistics in the request entity + # This must be done BEFORE resetting num_processed_tokens + victim.record_preemption(self._cluster_type, num_computed_tokens_before) + victim.advance_runtime_epoch() + + # Remove from running requests + if victim in self._running_requests: + self._running_requests.remove(victim) + + # Free allocated blocks + if victim.id in self._allocation_map: + self._free_request_resources(victim) + + # Mark as preempted and reset the scheduler-visible computed frontier. + # Disaggregated decode requests arrive after PREFILL has completed and + # the prompt KV frontier has transferred to the decode-side cluster. + # Their Request-level token lifecycle must survive memory preemption; + # only scheduler-local computed state and KV allocation are restarted. + victim._preempted = True + if ( + self._cluster_type + not in _REQUEST_PROGRESS_PRESERVING_PREEMPTION_CLUSTER_TYPES + ): + victim._num_processed_tokens = 0 # Reset computed tokens as in vLLM v1 + self._scheduled_num_computed_tokens_by_request.pop(victim.id, None) + + # Record re-entry to waiting queue for waiting time tracking after the + # lifecycle decision above and before adding the request to the queue. + victim.on_enter_waiting_queue(self._current_schedule_time, self._cluster_type) + + # Add to front of appropriate waiting queue (prepend) + # DECODE and DECODE_ATTN clusters use _waiting_requests, others use _request_queue + if self._cluster_type in [ClusterType.DECODE, ClusterType.DECODE_ATTN]: + self._waiting_requests.insert(0, victim) + else: + self._request_queue.insert(0, victim) + + # Track for this iteration + preempted_requests.append(victim) + + logger.info( + f"[VLLMv1Engine] Preempted request {victim.id} " + f"(policy={self._scheduling_policy}), " + f"running_reqs={len(self._running_requests)}" + ) + + # Flow validation: log preemption event + logger.info( + f"[PREEMPTION] req={victim.id} preempted, " + f"policy={self._scheduling_policy}, " + f"freed_blocks={freed_blocks}" + ) + available_blocks_preempt = int(self._config.num_blocks - self._num_allocated_blocks) + self._emit_schedule_decision_event( + event="decision", + decision_result="PREEMPTED", + request_id=victim.id, + token_budget=self._current_iteration_token_budget, + available_blocks=available_blocks_preempt, + num_tokens=0, + ) + + # Flow validation: log detailed preemption info + victim_selection_reason = ( + "lowest_priority" + if self._scheduling_policy == "priority" + else "tail_of_running_queue" + ) + logger.info( + f"[PREEMPTION_DETAIL] req={victim.id}, " + f"num_computed_tokens_before={num_computed_tokens_before}, " + f"freed_blocks={freed_blocks}, " + f"policy={self._scheduling_policy}, " + f"victim_selection_reason={victim_selection_reason}, " + f"queue_position_before={queue_position_before}, " + f"running_count_before={running_count_before}, " + f"running_count_after={len(self._running_requests)}" + ) + + def _try_allocate_with_preemption( + self, + request: Request, + num_new_tokens: int, + preempted_requests: List[Request], + *, + scheduler_num_computed_tokens: Optional[int] = None, + ) -> bool: + """ + Try to allocate memory for a request, preempting other requests if necessary. + + This implements the core preemption loop from vLLM v1 scheduler. + + Args: + request: The request to allocate for + num_new_tokens: Number of new tokens to process + preempted_requests: List to track preempted requests + + Returns: + bool: True if allocation succeeded (possibly after preemption) + """ + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + + while True: + if scheduler_num_computed_tokens is None: + can_allocate = self._can_allocate_request(request, num_new_tokens) + else: + can_allocate = self._can_allocate_request( + request, + num_new_tokens, + scheduler_num_computed_tokens=scheduler_num_computed_tokens, + ) + if can_allocate: + if scheduler_num_computed_tokens is None: + self._allocate_request(request, num_new_tokens) + else: + self._allocate_request( + request, + num_new_tokens, + scheduler_num_computed_tokens=scheduler_num_computed_tokens, + ) + return True + + if not self._enable_preemption: + return False + + # Flow validation: log memory pressure + available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) + logger.info( + f"[MEMORY_PRESSURE] trigger=allocation_failed, " + f"requesting_req={request.id}, " + f"requested_tokens={num_new_tokens}, " + f"available_blocks={available_blocks}, " + f"running_queue_size={len(self._running_requests)}" + ) + + # Select victim for preemption (exclude current request) + victim = self._select_preemption_victim(exclude=request) + + if victim is None: + # No victims available (all other requests have higher priority or no other requests) + # Preempt self and move to waiting queue + self._preempt_request(request, preempted_requests) + return False + + # Preempt victim and try again + self._preempt_request(victim, preempted_requests) + + def _rollback_current_iteration_preempted_requests( + self, + *, + scheduled_requests: List[Request], + scheduled_num_tokens: List[int], + newly_preempted_requests: List[Request], + token_budget: int, + ) -> int: + if not newly_preempted_requests: + return token_budget + + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + preempted_request_ids = { + int(request.id) for request in newly_preempted_requests + } + if not preempted_request_ids: + return token_budget + + kept_requests: List[Request] = [] + kept_num_tokens: List[int] = [] + refunded_tokens = 0 + + for scheduled_request, scheduled_tokens in zip( + scheduled_requests, scheduled_num_tokens + ): + if int(scheduled_request.id) in preempted_request_ids: + refunded_tokens += int(scheduled_tokens) + logger.info( + "[RUNNING-SCHEDULE-ROLLBACK] req=%s removed from current iteration " + "after same-iteration preemption, refunded_tokens=%s", + scheduled_request.id, + scheduled_tokens, + ) + continue + kept_requests.append(scheduled_request) + kept_num_tokens.append(int(scheduled_tokens)) + + if refunded_tokens == 0: + return token_budget + + scheduled_requests[:] = kept_requests + scheduled_num_tokens[:] = kept_num_tokens + token_budget += refunded_tokens + self._current_iteration_token_budget = token_budget + return token_budget diff --git a/frontier/scheduler/replica_scheduler/vllm_v1_mtp_wait.py b/frontier/scheduler/replica_scheduler/vllm_v1_mtp_wait.py new file mode 100644 index 00000000..4e38876e --- /dev/null +++ b/frontier/scheduler/replica_scheduler/vllm_v1_mtp_wait.py @@ -0,0 +1,1142 @@ +"""Output-wait policy for target-embedded MTP under monolithic pipeline parallelism. + +A target-embedded MTP request can finish its draft tokens before the pipeline +has drained, so admitting the next request immediately would let the simulated +engine run ahead of what the real one can do. These methods decide how long to +hold admission and terminal release, and how much visible budget to reserve. +""" + +from typing import Dict, List, Optional, Tuple + +from frontier.entities.batch import Batch, Request +from frontier.logger import get_cluster_logger +from frontier.spec_decode import compute_iteration_outcome, get_planned_draft_tokens +from frontier.types import ClusterType + + +class TargetEmbeddedMtpWaitPolicy: + """Admission, output-wait and terminal-release timing for embedded MTP.""" + + def _refresh_target_embedded_mtp_prefill_boundary_state( + self, batch: Batch, request: Request + ) -> None: + metadata = batch.spec_decode_metadata + if metadata is not None: + for _, batch_request in enumerate(batch.requests): + if batch_request is request: + # The metadata row is authoritative for this batch. A + # positive verify width was already recorded when metadata + # was built; a zero verify width means the request did not + # participate in this spec-decode iteration. + return + if self._cluster_type != ClusterType.MONOLITHIC: + return + if not getattr(request, "spec_decode_enabled", False): + return + if not getattr(request, "spec_method_is_target_embedded_mtp", False): + return + if not getattr(request, "is_prefill_complete", False): + return + spec_decode_config = getattr(self, "_spec_decode_config", None) + if spec_decode_config is None: + raise ValueError("Speculative decoding config is not initialized") + if int(getattr(request, "spec_total_iterations", 0)) == 0: + planned_drafts = get_planned_draft_tokens( + spec_decode_config, + request.remaining_decode_tokens, + iteration_index=0, + request_id=str(request.id), + ) + outcome = compute_iteration_outcome( + spec_decode_config, + request.remaining_decode_tokens, + planned_draft_tokens=planned_drafts, + iteration_index=0, + request_id=str(request.id), + ) + if planned_drafts == 0 and outcome.committed_tokens == 1: + # Some vLLM target-embedded MTP requests emit a real prefill + # commit row: one sampled token, no scheduled drafts. Frontier + # already advanced that token at the prefill boundary, so only + # the trace cursor and spec stats need to catch up here. + request.record_spec_decode_iteration( + verify_tokens=outcome.verify_tokens, + accepted_drafts=outcome.accepted_draft_tokens, + rejected_drafts=outcome.rejected_draft_tokens, + committed_tokens=outcome.committed_tokens, + ) + else: + request.set_spec_next_planned_draft_tokens(planned_drafts) + return + if int(getattr(request, "spec_total_iterations", 0)) != 1: + return + request.set_spec_next_planned_draft_tokens( + get_planned_draft_tokens( + spec_decode_config, + request.remaining_decode_tokens, + iteration_index=request.spec_total_iterations, + request_id=str(request.id), + ) + ) + + def _get_monolithic_pp_waiting_admission_delay_iters(self) -> Dict[int, int]: + delay_iters = getattr( + self, + "_monolithic_pp_waiting_admission_delay_iters", + None, + ) + if delay_iters is None: + delay_iters = {} + self._monolithic_pp_waiting_admission_delay_iters = delay_iters + return delay_iters + + def _is_target_embedded_mtp_request(self, request: Request) -> bool: + return ( + bool(getattr(request, "spec_decode_enabled", False)) + and bool(getattr(request, "spec_method_is_target_embedded_mtp", False)) + ) + + def _has_future_planned_draft_tokens(self, request: Request) -> bool: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + per_request_trace = getattr( + spec_config, + "_per_request_scheduled_draft_tokens_trace", + None, + ) + if per_request_trace is None: + return False + request_id = str(request.id) + if request_id not in per_request_trace: + raise ValueError( + "per-request scheduled draft trace missing request_id=" + f"{request_id!r}" + ) + next_iteration = int(getattr(request, "spec_total_iterations", 0)) + 1 + return any(int(tokens) > 0 for tokens in per_request_trace[request_id][next_iteration:]) + + def _get_target_embedded_mtp_terminal_overshoot_rows( + self, + request: Request, + *, + start_iteration_index: int, + ) -> List[Tuple[int, int, int, int, int]]: + if self._cluster_type != ClusterType.MONOLITHIC: + return [] + if self._num_stages <= 1: + return [] + if not getattr(request, "spec_decode_enabled", False): + return [] + if not getattr(request, "spec_method_is_target_embedded_mtp", False): + return [] + + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + raise ValueError("Speculative decoding config is not initialized") + per_request_planned_trace = getattr( + spec_config, + "_per_request_scheduled_draft_tokens_trace", + None, + ) + per_request_committed_trace = getattr( + spec_config, + "_per_request_committed_tokens_trace", + None, + ) + if per_request_planned_trace is None and per_request_committed_trace is None: + return [] + if per_request_planned_trace is None or per_request_committed_trace is None: + raise ValueError( + "terminal target-embedded MTP overshoot modeling requires both " + "per-request scheduled draft and committed token traces" + ) + + request_id = str(request.id) + if request_id not in per_request_planned_trace: + raise ValueError( + "per-request scheduled draft trace missing request_id=" + f"{request_id!r}" + ) + if request_id not in per_request_committed_trace: + raise ValueError( + "per-request acceptance trace missing request_id=" + f"{request_id!r}" + ) + + planned_trace = per_request_planned_trace[request_id] + committed_trace = per_request_committed_trace[request_id] + if len(planned_trace) != len(committed_trace): + raise ValueError( + "per-request MTP trace length mismatch: " + f"request_id={request_id!r}, planned_len={len(planned_trace)}, " + f"committed_len={len(committed_trace)}" + ) + + start_idx = int(start_iteration_index) + if start_idx < 0: + raise ValueError( + f"start_iteration_index must be >= 0, got={start_idx}" + ) + if start_idx >= len(planned_trace): + return [] + + terminal_rows: List[Tuple[int, int, int, int, int]] = [] + for idx in range(start_idx, len(planned_trace)): + planned_drafts = int(planned_trace[idx]) + raw_committed = int(committed_trace[idx]) + if planned_drafts < 0: + raise ValueError( + "terminal scheduled draft tokens must be >= 0, " + f"request_id={request_id!r}, iteration_index={idx}, " + f"got={planned_drafts}" + ) + if raw_committed < 0: + raise ValueError( + "terminal committed tokens must be >= 0, " + f"request_id={request_id!r}, iteration_index={idx}, " + f"got={raw_committed}" + ) + trace_verify_tokens = 1 + planned_drafts + if raw_committed > trace_verify_tokens: + raise ValueError( + "terminal committed tokens cannot exceed verify window: " + f"request_id={request_id!r}, iteration_index={idx}, " + f"committed={raw_committed}, " + f"verify_tokens={trace_verify_tokens}" + ) + if planned_drafts == 0 and raw_committed == 0: + continue + + # Once the logical response is complete, vLLM's online scheduler no + # longer replays the full forced-acceptance scheduled-draft window + # for request latency. The diagnostic acceptance trace can still + # contain scheduled_draft_tokens=32 for the next trace row, while + # the clean scheduler batch log exposes only a one-token cleanup row + # for that completed request. Model that terminal cleanup as one + # target token and keep the raw committed count only for audit + # provenance. Replaying the full trace window here over-extends + # short decode-tail request latency and violates the clean online + # metric scope. + terminal_cleanup_verify_tokens = 1 + terminal_rows.append( + ( + 0, + terminal_cleanup_verify_tokens, + 0, + 0, + raw_committed, + ) + ) + return terminal_rows + + def _should_delay_monolithic_pp_waiting_admission_on_add( + self, request: Request + ) -> bool: + if self._cluster_type != ClusterType.MONOLITHIC: + return False + if self._num_stages <= 1: + return False + if not self._is_target_embedded_mtp_request(request): + return False + planned_drafts = int(getattr(request, "spec_next_planned_draft_tokens", 0)) + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + num_speculative_tokens = int( + getattr(spec_config, "num_speculative_tokens", 0) + ) + if planned_drafts <= 0: + # A zero first scheduled-draft step still has PP lookahead admission + # visibility cost when a long decode will enter later MTP draft steps. + if num_speculative_tokens > 2: + return False + if int(getattr(request, "num_decode_tokens", 0)) < 128: + return False + if not self._has_future_planned_draft_tokens(request): + return False + else: + block_size = int(getattr(self._config, "block_size", 16)) + if num_speculative_tokens >= block_size: + # Full-block-or-wider target-embedded MTP request admission is + # already protected by output-visible guards after it starts + # running. Adding a pre-admission PP boundary here makes late + # online arrivals miss the vLLM-visible scheduler slot and + # under-batches wide verify traces relative to vLLM. + return False + if ( + num_speculative_tokens > 2 + and int(getattr(request, "num_prefill_tokens", 0)) + >= int(self._max_num_scheduled_tokens) + + max(1, int(self._max_num_scheduled_tokens) // 2) + ): + # Long multi-chunk prefills have enough remaining prefill work + # to expose the next PP boundary through the prefill chunks + # themselves. Adding a separate half-block MTP admission delay + # over-queues these arrivals and inflates p90 TTFT; keep the + # delay for shorter prefills where r106 showed global + # half-block skipping over-corrects TPOT tails. + return False + return ( + self._num_running_batches > 0 + or bool(self._running_requests) + or bool(self._get_active_batch_request_counts()) + ) + + def _add_monolithic_pp_waiting_admission_delay( + self, request_id: int, *, wait_iters: Optional[int] = None + ) -> None: + resolved_wait_iters = int( + wait_iters if wait_iters is not None else max(1, self._num_stages - 1) + ) + if resolved_wait_iters <= 0: + raise ValueError( + f"wait_iters must be positive, got={resolved_wait_iters}" + ) + delay_iters = self._get_monolithic_pp_waiting_admission_delay_iters() + delay_iters[request_id] = max( + int(delay_iters.get(request_id, 0)), + resolved_wait_iters, + ) + + def _should_defer_monolithic_pp_waiting_admission( + self, request: Request + ) -> bool: + delay_iters = self._get_monolithic_pp_waiting_admission_delay_iters() + remaining = int(delay_iters.get(request.id, 0)) + if remaining <= 0: + delay_iters.pop(request.id, None) + return False + if self._cluster_type != ClusterType.MONOLITHIC or self._num_stages <= 1: + delay_iters.pop(request.id, None) + return False + if not self._is_target_embedded_mtp_request(request): + delay_iters.pop(request.id, None) + return False + if ( + not self._running_requests + and self._num_running_batches <= 0 + and not self._get_active_batch_request_counts() + ): + # No active PP work remains to provide a future output-visible + # scheduler boundary; fail open to avoid deadlocking the queue. + delay_iters.pop(request.id, None) + return False + remaining -= 1 + if remaining > 0: + delay_iters[request.id] = remaining + else: + delay_iters.pop(request.id, None) + return True + + def _get_monolithic_pp_mtp_near_full_prefill_request_ids(self) -> set[int]: + request_ids = getattr( + self, + "_monolithic_pp_mtp_near_full_prefill_request_ids", + None, + ) + if request_ids is None: + request_ids = set() + self._monolithic_pp_mtp_near_full_prefill_request_ids = request_ids + return request_ids + + def _get_monolithic_pp_mtp_single_output_wait_request_ids(self) -> set[int]: + request_ids = getattr( + self, + "_monolithic_pp_mtp_single_output_wait_request_ids", + None, + ) + if request_ids is None: + request_ids = set() + self._monolithic_pp_mtp_single_output_wait_request_ids = request_ids + return request_ids + + def _get_target_embedded_mtp_request_acceptance_ratio( + self, request: Request + ) -> Optional[float]: + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + if spec_config is None: + return None + committed_trace_map = getattr( + spec_config, + "_per_request_committed_tokens_trace", + None, + ) + scheduled_trace_map = getattr( + spec_config, + "_per_request_scheduled_draft_tokens_trace", + None, + ) + if committed_trace_map is None and scheduled_trace_map is None: + return None + if committed_trace_map is None or scheduled_trace_map is None: + raise ValueError( + "MTP request acceptance audit requires both per-request " + "committed and scheduled-draft traces" + ) + request_id = str(request.id) + if request_id not in committed_trace_map: + raise ValueError( + "per-request acceptance trace missing request_id=" + f"{request_id!r}" + ) + if request_id not in scheduled_trace_map: + raise ValueError( + "per-request scheduled draft trace missing request_id=" + f"{request_id!r}" + ) + committed_trace = committed_trace_map[request_id] + scheduled_trace = scheduled_trace_map[request_id] + if len(committed_trace) != len(scheduled_trace): + raise ValueError( + "MTP request acceptance audit trace length mismatch: " + f"request_id={request_id!r}, " + f"committed_len={len(committed_trace)}, " + f"scheduled_len={len(scheduled_trace)}" + ) + accepted_drafts = 0 + scheduled_drafts = 0 + for committed_tokens, planned_drafts in zip( + committed_trace, + scheduled_trace, + ): + planned = int(planned_drafts) + if planned <= 0: + continue + committed = int(committed_tokens) + accepted_drafts += min(max(committed - 1, 0), planned) + scheduled_drafts += planned + if scheduled_drafts <= 0: + return None + return accepted_drafts / scheduled_drafts + + def _get_monolithic_pp_mtp_output_wait_prefill_threshold(self) -> int: + max_scheduled_tokens = int(self._max_num_scheduled_tokens) + block_size = int(getattr(self._config, "block_size", 16)) + headroom_tokens = 4 * block_size + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + num_speculative_tokens = int( + getattr(spec_config, "num_speculative_tokens", 0) + ) + if num_speculative_tokens >= block_size: + # Wide target-embedded MTP carries a larger PP lookahead payload than + # the narrow-window cases that established the original four-block + # threshold. Reserve window-proportional headroom so long chunked + # prefill slices enter the same output-visible continuation lane + # instead of being treated as ordinary prefill chunks. + headroom_tokens = max(headroom_tokens, 8 * num_speculative_tokens) + return max(1, max_scheduled_tokens - headroom_tokens) + + def _get_monolithic_pp_mtp_output_wait_iters(self) -> int: + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + if spec_config is None: + return 2 + + block_size = int(getattr(self._config, "block_size", 16)) + num_speculative_tokens = int( + getattr(spec_config, "num_speculative_tokens", 0) + ) + if num_speculative_tokens < block_size: + return 2 + + committed_trace_map = getattr( + spec_config, + "_per_request_committed_tokens_trace", + None, + ) + scheduled_trace_map = getattr( + spec_config, + "_per_request_scheduled_draft_tokens_trace", + None, + ) + accepted_drafts = 0 + scheduled_drafts = 0 + if committed_trace_map is not None or scheduled_trace_map is not None: + if committed_trace_map is None or scheduled_trace_map is None: + raise ValueError( + "MTP output-wait acceptance audit requires both " + "per-request committed and scheduled-draft traces" + ) + if set(committed_trace_map.keys()) != set(scheduled_trace_map.keys()): + raise ValueError( + "MTP output-wait acceptance audit trace keys mismatch" + ) + for request_id, committed_trace in committed_trace_map.items(): + scheduled_trace = scheduled_trace_map[request_id] + if len(committed_trace) != len(scheduled_trace): + raise ValueError( + "MTP output-wait acceptance audit trace length mismatch: " + f"request_id={request_id!r}, " + f"committed_len={len(committed_trace)}, " + f"scheduled_len={len(scheduled_trace)}" + ) + for committed_tokens, planned_drafts in zip( + committed_trace, + scheduled_trace, + ): + planned = int(planned_drafts) + if planned <= 0: + continue + committed = int(committed_tokens) + accepted_drafts += min(max(committed - 1, 0), planned) + scheduled_drafts += planned + else: + committed_trace = getattr( + spec_config, + "_committed_tokens_trace", + None, + ) + scheduled_trace = getattr( + spec_config, + "_scheduled_draft_tokens_trace", + None, + ) + if committed_trace is None or scheduled_trace is None: + return 2 + if len(committed_trace) != len(scheduled_trace): + raise ValueError( + "MTP output-wait acceptance audit global trace length mismatch: " + f"committed_len={len(committed_trace)}, " + f"scheduled_len={len(scheduled_trace)}" + ) + for committed_tokens, planned_drafts in zip( + committed_trace, + scheduled_trace, + ): + planned = int(planned_drafts) + if planned <= 0: + continue + committed = int(committed_tokens) + accepted_drafts += min(max(committed - 1, 0), planned) + scheduled_drafts += planned + + if scheduled_drafts <= 0: + return 2 + acceptance_ratio = accepted_drafts / scheduled_drafts + if acceptance_ratio >= 0.5: + if self._has_monolithic_pp_visible_waiting_requests(): + # High-acceptance wide MTP should still preserve the extra + # output-visible boundary while fresh prefill admissions are + # visible. Otherwise decode continuation can consume the + # online token budget before late-arriving prefills enter, + # inflating TTFT tails. Once no waiting prefill/resume request + # is visible, shorten the wait to protect short decode tails. + return 2 + # High-acceptance wide target-embedded MTP has fewer decode + # scheduler turns per request. A single output-visible turn is + # enough to expose PP continuation without repeatedly holding + # short tail requests behind terminal trace rows. + return 1 + return 2 + + def _get_monolithic_pp_mtp_output_wait_iters_for_request( + self, request: Request + ) -> int: + wait_iters = self._get_monolithic_pp_mtp_output_wait_iters() + if request.id in self._get_monolithic_pp_mtp_single_output_wait_request_ids(): + return min(wait_iters, 1) + if self._should_apply_monolithic_pp_mtp_fractional_extra_output_wait( + request + ): + counts = self._get_monolithic_pp_mtp_fractional_output_wait_counts() + previous_count = int(counts.get(request.id, 0)) + counts[request.id] = previous_count + 1 + if previous_count == 0: + return 0 + if previous_count % 2 == 1: + return wait_iters + return min(wait_iters, 1) + if self._should_extend_monolithic_pp_mtp_long_decode_output_wait(request): + return wait_iters + 1 + return wait_iters + + def _is_monolithic_pp_mtp_half_block_low_acceptance_request( + self, + request: Request, + *, + acceptance_ratio_limit: float, + ) -> bool: + block_size = int(getattr(self._config, "block_size", 16)) + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + num_speculative_tokens = int( + getattr(spec_config, "num_speculative_tokens", 0) + ) + if num_speculative_tokens < max(1, block_size // 2): + return False + if num_speculative_tokens >= block_size: + return False + + acceptance_ratio = self._get_target_embedded_mtp_request_acceptance_ratio( + request + ) + if acceptance_ratio is None: + return False + return acceptance_ratio < acceptance_ratio_limit + + def _is_monolithic_pp_mtp_mid_prefill_request( + self, + request: Request, + ) -> bool: + block_size = int(getattr(self._config, "block_size", 16)) + num_prefill_tokens = int(getattr(request, "num_prefill_tokens", 0)) + max_scheduled_tokens = int(self._max_num_scheduled_tokens) + if num_prefill_tokens < max(1, max_scheduled_tokens - 4 * block_size): + return False + if num_prefill_tokens > ( + max_scheduled_tokens + max(1, max_scheduled_tokens // 2) + ): + return False + return True + + def _should_apply_monolithic_pp_mtp_fractional_extra_output_wait( + self, + request: Request, + ) -> bool: + if not self._is_target_embedded_mtp_request(request): + return False + if not self._is_monolithic_pp_mtp_mid_prefill_request(request): + return False + + block_size = int(getattr(self._config, "block_size", 16)) + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + num_speculative_tokens = int( + getattr(spec_config, "num_speculative_tokens", 0) + ) + if num_speculative_tokens < max(1, block_size // 2): + return False + if num_speculative_tokens >= block_size: + return False + + acceptance_ratio = self._get_target_embedded_mtp_request_acceptance_ratio( + request + ) + if acceptance_ratio is None: + return False + + block_size = int(getattr(self._config, "block_size", 16)) + num_decode_tokens = int(getattr(request, "num_decode_tokens", 0)) + if num_decode_tokens < 16 * block_size: + return False + if num_decode_tokens < 24 * block_size: + if acceptance_ratio >= 0.5: + return False + else: + if acceptance_ratio >= 0.55: + return False + + # The v8/a0.3 request-level RCA shows that medium-length decode tails + # and long decode tails need about one and a half PP output-visible + # wait turns after the first visible decode result. Skip the first + # wait so TTFT remains a prefill/first-token metric, then alternate + # the regular low-acceptance wait with a single-turn wait. This keeps + # the correction as a scheduler visibility family rather than an + # op-runtime calibration scale. + return True + + def _should_extend_monolithic_pp_mtp_long_decode_output_wait( + self, request: Request + ) -> bool: + if not self._is_monolithic_pp_mtp_half_block_low_acceptance_request( + request, + acceptance_ratio_limit=0.25, + ): + return False + if not self._is_monolithic_pp_mtp_mid_prefill_request(request): + return False + + block_size = int(getattr(self._config, "block_size", 16)) + num_decode_tokens = int(getattr(request, "num_decode_tokens", 0)) + if num_decode_tokens < 24 * block_size: + return False + + # Low-acceptance half-block MTP keeps many long decode continuations + # alive after a near-full prefill admission. vLLM exposes an additional + # PP-visible output boundary for the very long decode tail; model that + # boundary as one extra scheduler wait turn instead of + # hiding the residual in compute calibration scale. + return True + + def _record_monolithic_pp_mtp_near_full_prefill_slices(self, batch: Batch) -> None: + if self._cluster_type != ClusterType.MONOLITHIC: + return + if self._num_stages <= 1: + return + near_full_prefill_threshold = ( + self._get_monolithic_pp_mtp_output_wait_prefill_threshold() + ) + for request, num_tokens in zip(batch.requests, batch.num_tokens): + if getattr(request, "is_prefill_complete", False): + continue + if not getattr(request, "spec_decode_enabled", False): + continue + if not getattr(request, "spec_method_is_target_embedded_mtp", False): + continue + if int(num_tokens) < near_full_prefill_threshold: + if not self._should_record_monolithic_pp_mtp_subthreshold_single_wait_prefill( + request, + num_tokens=int(num_tokens), + near_full_prefill_threshold=near_full_prefill_threshold, + ): + if not self._should_record_monolithic_pp_mtp_subthreshold_long_decode_prefill( + request, + num_tokens=int(num_tokens), + near_full_prefill_threshold=near_full_prefill_threshold, + ): + continue + else: + self._get_monolithic_pp_mtp_single_output_wait_request_ids().add( + request.id + ) + self._get_monolithic_pp_mtp_near_full_prefill_request_ids().add( + request.id + ) + + def _should_record_monolithic_pp_mtp_subthreshold_single_wait_prefill( + self, + request: Request, + *, + num_tokens: int, + near_full_prefill_threshold: int, + ) -> bool: + block_size = int(getattr(self._config, "block_size", 16)) + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + num_speculative_tokens = int( + getattr(spec_config, "num_speculative_tokens", 0) + ) + if num_speculative_tokens >= block_size: + return False + + max_scheduled_tokens = int(self._max_num_scheduled_tokens) + min_subthreshold_tokens = max( + 1, + max_scheduled_tokens - 8 * block_size, + ) + if int(num_tokens) < min_subthreshold_tokens: + return False + if int(num_tokens) >= int(near_full_prefill_threshold): + return False + if int(getattr(request, "num_prefill_tokens", 0)) < ( + max_scheduled_tokens + max(1, max_scheduled_tokens // 2) + ): + return False + if int(getattr(request, "num_decode_tokens", 0)) > 4 * block_size: + return False + + acceptance_ratio = self._get_target_embedded_mtp_request_acceptance_ratio( + request + ) + if acceptance_ratio is None: + return False + if acceptance_ratio >= 0.5: + return False + + # Low-acceptance narrow MTP preserves more decode turns than the + # high-acceptance case, so a two-turn output wait over-delays the + # terminal request tail. A single PP-visible wait models the missing + # output boundary without absorbing the residual into compute scale. + return True + + def _should_record_monolithic_pp_mtp_subthreshold_long_decode_prefill( + self, + request: Request, + *, + num_tokens: int, + near_full_prefill_threshold: int, + ) -> bool: + block_size = int(getattr(self._config, "block_size", 16)) + max_scheduled_tokens = int(self._max_num_scheduled_tokens) + min_subthreshold_tokens = max( + 1, + max_scheduled_tokens - 8 * block_size, + ) + if int(num_tokens) < min_subthreshold_tokens: + return False + if int(num_tokens) >= int(near_full_prefill_threshold): + return False + if not ( + self._should_apply_monolithic_pp_mtp_fractional_extra_output_wait( + request + ) + or self._should_extend_monolithic_pp_mtp_long_decode_output_wait( + request + ) + ): + return False + + # These subthreshold prefill slices are close enough to the online + # max-token boundary to expose the same PP output-visible behavior as + # near-full chunks, but only for the low-acceptance half-block + # medium/long decode slices identified by the request-level TPOT RCA. + return True + + def _get_monolithic_pp_pending_terminal_release_iters(self) -> Dict[int, int]: + pending = getattr( + self, + "_monolithic_pp_pending_terminal_release_iters", + None, + ) + if pending is None: + pending = {} + self._monolithic_pp_pending_terminal_release_iters = pending + return pending + + def _get_monolithic_pp_extra_terminal_release_iters(self) -> int: + if self._cluster_type != ClusterType.MONOLITHIC: + return 0 + + pp = int(getattr(self._replica_config, "num_pipeline_stages", 1)) + if pp <= 1: + return 0 + + # Frontier's last-stage batch-end already accounts for one terminal + # drain iteration. Deeper PP still needs the sampled-token-return + # boundary to reach the scheduler before blocks can be released. + return max(pp // 2 - 1, 0) + + def _has_monolithic_pp_pending_terminal_release(self) -> bool: + return bool(self._get_monolithic_pp_pending_terminal_release_iters()) + + def _has_monolithic_pp_visible_waiting_requests(self) -> bool: + return bool(self._request_queue or self._preempted_requests) + + def _get_monolithic_pp_iteration_start_release_threshold(self) -> int: + if self._cluster_type != ClusterType.MONOLITHIC: + return 1 + + pp = int(getattr(self._replica_config, "num_pipeline_stages", 1)) + if pp <= 4: + return 1 + + # Once terminal release is materialized at iteration_start, deeper + # MONOLITHIC+PP pipelines expose the release boundary earlier than the + # old end-of-iteration bookkeeping. The validated scheduler-visible + # contracts are pp4->1 and pp8->2, so keep the threshold PP-depth + # aware instead of assuming a single remaining hop for every PP size. + return max(pp // 4, 1) + + def _advance_monolithic_pp_terminal_release_boundary(self) -> None: + pending_release_iters = ( + self._get_monolithic_pp_pending_terminal_release_iters() + ) + if not pending_release_iters: + return + + release_visible_threshold = ( + self._get_monolithic_pp_iteration_start_release_threshold() + ) + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + + ready_request_ids: List[int] = [] + for request_id, remaining_iters in list(pending_release_iters.items()): + if ( + remaining_iters <= release_visible_threshold + and self._has_monolithic_pp_visible_waiting_requests() + and request_id + not in self._monolithic_pp_waiting_sensitive_release_extensions + ): + self._monolithic_pp_waiting_sensitive_release_extensions.add(request_id) + pending_release_iters[request_id] = 1 + logger.debug( + "[VLLMv1Engine] Delaying MONOLITHIC+PP terminal release for " + "request %s by one extra empty iteration because waiting " + "requests are already visible", + request_id, + ) + continue + if remaining_iters <= 1: + ready_request_ids.append(request_id) + pending_release_iters.pop(request_id, None) + else: + pending_release_iters[request_id] = remaining_iters - 1 + + if not ready_request_ids: + if pending_release_iters: + self._monolithic_pp_terminal_release_followup_poll_pending = True + logger.debug( + "[VLLMv1Engine] Keeping MONOLITHIC+PP terminal release self-driven " + "with one follow-up schedule poll while pending state remains: %s", + dict(pending_release_iters), + ) + return + + ready_request_id_set = set(ready_request_ids) + for request_id in ready_request_ids: + self._free_request_resources_by_id(request_id) + self._scheduled_num_computed_tokens_by_request.pop(request_id, None) + self._monolithic_pp_waiting_sensitive_release_extensions.discard(request_id) + + self._running_requests = [ + request + for request in self._running_requests + if request.id not in ready_request_id_set + ] + self._monolithic_pp_terminal_release_followup_poll_pending = bool( + pending_release_iters + ) or self._has_monolithic_pp_visible_waiting_requests() + + logger.debug( + "[VLLMv1Engine] Released %s MONOLITHIC+PP terminal request(s) " + "after sampled-token-return-equivalent boundary: %s", + len(ready_request_ids), + ready_request_ids, + ) + + def _materialize_monolithic_pp_terminal_release_before_iteration_start( + self, + ) -> None: + pending_release_iters = ( + self._get_monolithic_pp_pending_terminal_release_iters() + ) + if not pending_release_iters: + return + if self._has_monolithic_pp_visible_waiting_requests(): + return + + release_visible_threshold = ( + self._get_monolithic_pp_iteration_start_release_threshold() + ) + ready_request_ids = [ + request_id + for request_id, remaining_iters in list(pending_release_iters.items()) + if remaining_iters <= release_visible_threshold + ] + if not ready_request_ids: + return + + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + ready_request_id_set = set(ready_request_ids) + for request_id in ready_request_ids: + pending_release_iters.pop(request_id, None) + self._free_request_resources_by_id(request_id) + self._scheduled_num_computed_tokens_by_request.pop(request_id, None) + self._monolithic_pp_waiting_sensitive_release_extensions.discard( + request_id + ) + + self._running_requests = [ + request + for request in self._running_requests + if request.id not in ready_request_id_set + ] + logger.debug( + "[VLLMv1Engine] Materialized %s MONOLITHIC+PP terminal release(s) " + "before iteration_start because no waiting request is visible: %s", + len(ready_request_ids), + ready_request_ids, + ) + + def consume_monolithic_pp_terminal_release_followup_poll(self) -> bool: + pending = bool( + getattr( + self, + "_monolithic_pp_terminal_release_followup_poll_pending", + False, + ) + ) + self._monolithic_pp_terminal_release_followup_poll_pending = False + return pending + + def _get_monolithic_pp_mtp_output_wait_request_ids(self) -> set[int]: + request_ids = getattr( + self, + "_monolithic_pp_mtp_output_wait_request_ids", + None, + ) + if request_ids is None: + request_ids = set() + self._monolithic_pp_mtp_output_wait_request_ids = request_ids + return request_ids + + def _get_monolithic_pp_mtp_output_wait_remaining_iters(self) -> Dict[int, int]: + remaining_iters = getattr( + self, + "_monolithic_pp_mtp_output_wait_remaining_iters", + None, + ) + if remaining_iters is None: + remaining_iters = {} + self._monolithic_pp_mtp_output_wait_remaining_iters = remaining_iters + return remaining_iters + + def _get_monolithic_pp_mtp_fractional_output_wait_counts(self) -> Dict[int, int]: + counts = getattr( + self, + "_monolithic_pp_mtp_fractional_output_wait_counts", + None, + ) + if counts is None: + counts = {} + self._monolithic_pp_mtp_fractional_output_wait_counts = counts + return counts + + def _add_monolithic_pp_mtp_output_wait( + self, request_id: int, *, wait_iters: int = 2 + ) -> None: + if wait_iters <= 0: + raise ValueError(f"wait_iters must be positive, got={wait_iters}") + self._get_monolithic_pp_mtp_output_wait_request_ids().add(request_id) + remaining_iters = self._get_monolithic_pp_mtp_output_wait_remaining_iters() + remaining_iters[request_id] = max( + int(remaining_iters.get(request_id, 0)), + int(wait_iters), + ) + + def _should_apply_monolithic_pp_mtp_output_wait( + self, request: Request + ) -> bool: + if self._cluster_type != ClusterType.MONOLITHIC: + return False + if self._num_stages <= 1: + return False + if not getattr(request, "spec_decode_enabled", False): + return False + if not getattr(request, "spec_method_is_target_embedded_mtp", False): + return False + if not getattr(request, "is_prefill_complete", False): + return False + if ( + request.id + not in self._get_monolithic_pp_mtp_near_full_prefill_request_ids() + ): + return False + processed_decode_tokens = int( + getattr(request, "num_processed_decode_tokens", 0) + ) + if processed_decode_tokens <= 0: + return False + if self._get_monolithic_pp_mtp_output_wait_iters() == 1: + block_size = int(getattr(self._config, "block_size", 16)) + remaining_decode_tokens = ( + int(getattr(request, "num_decode_tokens", 0)) + - processed_decode_tokens + ) + if remaining_decode_tokens <= block_size: + # High-acceptance wide-MTP short tails have no useful future + # prefill visibility to protect once the request is within the + # final block-sized decode window. Skipping this idle turn + # avoids a terminal PP output-wait residual without changing + # CUDA op calibration. + return False + return True + + def _has_monolithic_pp_mtp_output_wait(self) -> bool: + return bool(self._get_monolithic_pp_mtp_output_wait_request_ids()) + + def _should_reserve_monolithic_pp_mtp_visible_budget( + self, request: Request + ) -> bool: + if self._cluster_type != ClusterType.MONOLITHIC: + return False + if self._num_stages <= 1: + return False + if not self._is_target_embedded_mtp_request(request): + return False + if not getattr(request, "is_prefill_complete", False): + return False + block_size = int(getattr(self._config, "block_size", 16)) + active_verify_window_tokens = int( + getattr(request, "spec_current_verify_tokens", 0) + ) + if active_verify_window_tokens <= 0: + spec_config = getattr(self, "_spec_decode_config", None) + if spec_config is None: + spec_config = getattr( + getattr(self, "_replica_config", None), + "speculative_decoding_config", + None, + ) + active_verify_window_tokens = int( + getattr(spec_config, "num_speculative_tokens", 0) + ) + if active_verify_window_tokens < block_size: + # Narrow target-embedded MTP verify windows do not consume a full + # cache block of output-visible scheduler budget. Reserving a + # synthetic block-sized budget for them under MONOLITHIC+PP + # under-batches waiting prefill relative to vLLM clean traces. + return False + return self._has_monolithic_pp_visible_waiting_requests() + + def _get_monolithic_pp_mtp_visible_budget_reservation_tokens( + self, request: Request, token_budget: int + ) -> int: + if not self._should_reserve_monolithic_pp_mtp_visible_budget(request): + return 0 + reserved_tokens = max( + int(getattr(request, "spec_current_verify_tokens", 1)), + self._get_request_next_num_tokens(request), + ) + return min(max(reserved_tokens, 0), token_budget) + + def _clear_monolithic_pp_mtp_output_wait(self) -> None: + request_ids = self._get_monolithic_pp_mtp_output_wait_request_ids() + remaining_iters = self._get_monolithic_pp_mtp_output_wait_remaining_iters() + if not remaining_iters: + request_ids.clear() + return + next_waiting_request_ids: set[int] = set() + for request_id in list(request_ids): + remaining = int(remaining_iters.get(request_id, 1)) - 1 + if remaining > 0: + remaining_iters[request_id] = remaining + next_waiting_request_ids.add(request_id) + else: + remaining_iters.pop(request_id, None) + request_ids.clear() + request_ids.update(next_waiting_request_ids) + + def consume_monolithic_pp_mtp_output_wait_followup_poll(self) -> bool: + pending = bool( + getattr( + self, + "_monolithic_pp_mtp_output_wait_followup_poll_pending", + False, + ) + ) + self._monolithic_pp_mtp_output_wait_followup_poll_pending = False + return pending diff --git a/frontier/scheduler/replica_scheduler/vllm_v1_prefix_cache.py b/frontier/scheduler/replica_scheduler/vllm_v1_prefix_cache.py new file mode 100644 index 00000000..b2b5066a --- /dev/null +++ b/frontier/scheduler/replica_scheduler/vllm_v1_prefix_cache.py @@ -0,0 +1,253 @@ +"""Prefix-cache admission and its identity ledger for the vLLM V1 scheduler. + +Admission decides how many cached blocks a request may reuse; the ledger emits +the identity events that let a run be checked for block-reuse correctness. +""" + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence + +from frontier.entities.batch import Request +from frontier.kv_cache.base_kv_cache_manager import KVCacheAllocationResult +from frontier.kv_cache.kv_cache_block import KVCacheBlock, KVCacheBlockBinding +from frontier.scheduler.replica_scheduler.vllm_v1_decision_log import ( + _log_frontier_vllm_v1_schedule_decision, +) +from frontier.types import ClusterType + + +@dataclass(frozen=True) +class PrefixCacheAdmission: + raw_hit_blocks: tuple[KVCacheBlock, ...] + effective_hit_blocks: tuple[KVCacheBlock, ...] + raw_hit_bindings: tuple[KVCacheBlockBinding, ...] + effective_hit_bindings: tuple[KVCacheBlockBinding, ...] + raw_cached_tokens: int + effective_cached_tokens: int + num_new_tokens: int + full_hit_backoff_applied: bool + + +def _serialize_prefix_cache_binding( + binding: KVCacheBlockBinding, +) -> Dict[str, Any]: + return { + "block_hash": binding.block_hash, + "block_id": int(binding.block_id), + "creator_request_id": str(binding.creator_request_id), + "binding_epoch": int(binding.binding_epoch), + } + + +class PrefixCacheLedger: + """Prefix-cache admission and identity-event emission.""" + + def _is_prefix_caching_enabled(self) -> bool: + return getattr(self, "_kv_cache_manager", None) is not None + + def _sync_prefix_cache_allocation_state( + self, request: Optional[Request] = None + ) -> None: + if not self._is_prefix_caching_enabled(): + return + assert self._kv_cache_manager is not None + self._num_allocated_blocks = int(self._kv_cache_manager.num_used_blocks) + if request is not None: + num_blocks = int(self._kv_cache_manager.get_num_blocks_for_request(request)) + if num_blocks > 0: + self._allocation_map[request.id] = num_blocks + else: + self._allocation_map.pop(request.id, None) + + def _prepare_prefix_cache_admission( + self, request: Request + ) -> PrefixCacheAdmission: + if not self._is_prefix_caching_enabled(): + return PrefixCacheAdmission( + raw_hit_blocks=(), + effective_hit_blocks=(), + raw_hit_bindings=(), + effective_hit_bindings=(), + raw_cached_tokens=0, + effective_cached_tokens=0, + num_new_tokens=self._get_request_next_num_tokens(request), + full_hit_backoff_applied=False, + ) + if request.block_hash_ids is None: + raise ValueError( + "block_hash_ids are required when enable_prefix_caching=True" + ) + assert self._kv_cache_manager is not None + computed_blocks, num_computed_tokens = self._kv_cache_manager.get_computed_blocks( + request + ) + raw_hit_blocks = tuple(computed_blocks) + raw_hit_bindings: list[KVCacheBlockBinding] = [] + query_hashes = list(request.block_hash_ids) + for query_index, block in enumerate(raw_hit_blocks): + binding = block.binding + if binding is None: + raise ValueError( + f"Prefix cache hit block {block.block_id} has no binding identity." + ) + if binding.block_hash != query_hashes[query_index]: + raise ValueError( + "Prefix cache hit binding disagrees with ordered query hash: " + f"query_index={query_index}, " + f"query_hash={query_hashes[query_index]!r}, " + f"binding_hash={binding.block_hash!r}" + ) + raw_hit_bindings.append(binding) + raw_cached_tokens = int(num_computed_tokens) + num_new_tokens = int(request.num_prefill_tokens) - int(num_computed_tokens) + full_hit_backoff_applied = False + if num_new_tokens == 0 and computed_blocks: + num_computed_tokens -= int(self._config.block_size) + num_new_tokens = int(self._config.block_size) + computed_blocks = list(computed_blocks[:-1]) + self._kv_cache_manager.prefix_cache_stats.hits -= 1 + full_hit_backoff_applied = True + return PrefixCacheAdmission( + raw_hit_blocks=raw_hit_blocks, + effective_hit_blocks=tuple(computed_blocks), + raw_hit_bindings=tuple(raw_hit_bindings), + effective_hit_bindings=tuple( + raw_hit_bindings[: len(computed_blocks)] + ), + raw_cached_tokens=raw_cached_tokens, + effective_cached_tokens=int(num_computed_tokens), + num_new_tokens=int(num_new_tokens), + full_hit_backoff_applied=full_hit_backoff_applied, + ) + + def _prefix_cache_identity_event_base( + self, + *, + event: str, + request: Request, + ) -> Dict[str, Any]: + event_seq = int(self._prefix_cache_identity_event_seq) + self._prefix_cache_identity_event_seq = event_seq + 1 + cluster_name = ( + self._cluster_type.name + if self._cluster_type is not None + else ClusterType.MONOLITHIC.name + ) + replica_local_id = self._replica_local_id + if replica_local_id is not None and ( + type(replica_local_id) is not int or replica_local_id < 0 + ): + raise ValueError( + "Prefix cache identity replica_local_id must be None or an " + f"exact non-negative int, got {replica_local_id!r}" + ) + return { + "event": event, + "prefix_identity_schema_version": 1, + "source": "frontier", + "scheduler": "vllm_v1", + "cluster_type": cluster_name, + "replica_id": int(self._replica_id), + "replica_local_id": replica_local_id, + "iteration_id": int(self._active_schedule_iteration_id), + "identity_event_seq": event_seq, + "request_id": str(request.id), + "prefix_cache_block_size": int(self._config.block_size), + "simulation_time": float(self._current_schedule_time), + "simulation_time_semantics": "frontier_event_time_seconds", + } + + def _serialize_prefix_cache_hit_bindings( + self, + *, + request: Request, + bindings: Sequence[KVCacheBlockBinding], + ) -> List[Dict[str, Any]]: + query_hashes = list(request.block_hash_ids or []) + if len(bindings) > len(query_hashes): + raise ValueError( + "Prefix cache hit count exceeds the ordered query hash count." + ) + rows: List[Dict[str, Any]] = [] + for query_index, binding in enumerate(bindings): + if binding.block_hash != query_hashes[query_index]: + raise ValueError( + "Prefix cache hit binding disagrees with ordered query hash: " + f"query_index={query_index}, " + f"query_hash={query_hashes[query_index]!r}, " + f"binding_hash={binding.block_hash!r}" + ) + rows.append( + { + "query_index": query_index, + **_serialize_prefix_cache_binding(binding), + } + ) + return rows + + def _emit_prefix_cache_identity_events( + self, + *, + request: Request, + num_new_tokens: int, + allocation: KVCacheAllocationResult, + admission: Optional[PrefixCacheAdmission], + ) -> None: + if admission is not None: + admission_payload = self._prefix_cache_identity_event_base( + event="prefix_cache_admission", + request=request, + ) + admission_payload.update( + { + "query_hashes": list(request.block_hash_ids or []), + "raw_hit_blocks": self._serialize_prefix_cache_hit_bindings( + request=request, + bindings=admission.raw_hit_bindings, + ), + "admitted_hit_blocks": self._serialize_prefix_cache_hit_bindings( + request=request, + bindings=admission.effective_hit_bindings, + ), + "raw_cached_tokens": int(admission.raw_cached_tokens), + "admitted_cached_tokens": int( + admission.effective_cached_tokens + ), + "num_new_tokens": int(num_new_tokens), + "full_hit_backoff_applied": bool( + admission.full_hit_backoff_applied + ), + } + ) + _log_frontier_vllm_v1_schedule_decision(admission_payload) + + allocation_payload = self._prefix_cache_identity_event_base( + event="prefix_cache_allocation", + request=request, + ) + reused_blocks: List[Dict[str, Any]] = [] + for block in allocation.reused_blocks: + binding = block.binding + if binding is None: + raise ValueError( + f"Reused Prefix cache block {block.block_id} has no binding identity." + ) + reused_blocks.append(_serialize_prefix_cache_binding(binding)) + allocation_payload.update( + { + "num_new_tokens": int(num_new_tokens), + "reused_blocks": reused_blocks, + "new_block_ids": [ + int(block.block_id) for block in allocation.new_blocks + ], + "evicted_bindings": [ + _serialize_prefix_cache_binding(binding) + for binding in allocation.evicted_bindings + ], + "new_bindings": [ + _serialize_prefix_cache_binding(binding) + for binding in allocation.new_bindings + ], + } + ) + _log_frontier_vllm_v1_schedule_decision(allocation_payload) diff --git a/frontier/scheduler/replica_scheduler/vllm_v1_role_schedules.py b/frontier/scheduler/replica_scheduler/vllm_v1_role_schedules.py new file mode 100644 index 00000000..f53ec307 --- /dev/null +++ b/frontier/scheduler/replica_scheduler/vllm_v1_role_schedules.py @@ -0,0 +1,858 @@ +"""Single-role scheduling entry points for disaggregated clusters. + +Co-location schedules prefill and decode together through the two-phase path +that stays on the scheduler itself. A disaggregated cluster instead drives one +role per replica, and each of those roles has its own entry point here. +""" + +from collections import deque +from typing import List, Optional, Tuple + +from frontier.config import global_vars +from frontier.entities.batch import Batch, Request +from frontier.logger import get_cluster_logger +from frontier.types import ClusterType + + +class DisaggregatedRoleScheduling: + """Prefill-only, decode-only and decode-attention scheduling entry points.""" + + def _schedule_prefill_only(self) -> Optional[Batch]: + """ + Scheduling for PREFILL cluster. + + In PD-disaggregation, the prefill cluster only handles new requests + that need prefill computation. With chunked prefill enabled, running + partial-prefill requests are also scheduled in Phase 1. + + Returns: + Optional[Batch]: The scheduled batch + """ + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + + token_budget = self._max_num_scheduled_tokens + available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) + waiting_count = len(self._request_queue) + len(self._preempted_requests) + + # Flow validation: log iteration start + logger.info( + f"[ITERATION_START] token_budget={token_budget}, " + f"running_count={len(self._running_requests)}, " + f"waiting_count={waiting_count}, " + f"available_blocks={available_blocks}, " + f"max_running_reqs={self._max_num_running_reqs}" + ) + self._emit_schedule_decision_event( + event="iteration_start", + decision_result=None, + request_id=None, + token_budget=token_budget, + available_blocks=available_blocks, + num_tokens=0, + ) + + # Flow validation: log memory state + total_blocks = int(self._config.num_blocks) + allocated_blocks = int(self._num_allocated_blocks) + usage_ratio = allocated_blocks / total_blocks if total_blocks > 0 else 0.0 + watermark = self._watermark_blocks + logger.info( + f"[MEMORY_STATE] total_blocks={total_blocks}, " + f"allocated_blocks={allocated_blocks}, " + f"free_blocks={available_blocks}, " + f"usage_ratio={usage_ratio:.4f}, " + f"watermark_blocks={watermark}" + ) + reclaimed_requests = self._reclaim_borrowed_final_running_slots( + waiting_requests=self._preempted_requests + self._request_queue, + final_predicate=self._is_final_prefill_fast_lane_request, + reserved_slots=self._final_prefill_reserved_slots, + lane_name="prefill", + ) + if reclaimed_requests: + waiting_count = len(self._request_queue) + len(self._preempted_requests) + + all_scheduled_requests: List[Request] = [] + all_num_tokens: List[int] = [] + preempted_requests: List[Request] = [] + waiting_scheduled: List[Request] = [] + waiting_tokens: List[int] = [] + + # Phase 1: schedule running requests (partial prefill continuation) + logger.info( + f"[PHASE1_START] running_count={len(self._running_requests)}, " + f"token_budget={token_budget}, " + f"waiting_count={waiting_count}" + ) + token_budget, running_scheduled, running_tokens = self._schedule_running_requests( + token_budget, preempted_requests + ) + all_scheduled_requests.extend(running_scheduled) + all_num_tokens.extend(running_tokens) + + available_blocks_p1 = int(self._config.num_blocks - self._num_allocated_blocks) + logger.info( + f"[PHASE1_END] scheduled_count={len(running_scheduled)}, " + f"preempted_count={len(preempted_requests)}, " + f"token_budget_remaining={token_budget}, " + f"available_blocks={available_blocks_p1}" + ) + + # Phase 2: schedule waiting requests only when Phase 1 has no preemption + if not preempted_requests: + logger.info( + f"[PHASE2_START] waiting_count={waiting_count}, " + f"token_budget={token_budget}, " + f"running_count={len(self._running_requests)}" + ) + token_budget, waiting_scheduled, waiting_tokens = ( + self._schedule_waiting_requests(token_budget) + ) + all_scheduled_requests.extend(waiting_scheduled) + all_num_tokens.extend(waiting_tokens) + + available_blocks_p2 = int( + self._config.num_blocks - self._num_allocated_blocks + ) + logger.info( + f"[PHASE2_END] admitted_count={len(waiting_scheduled)}, " + f"token_budget_remaining={token_budget}, " + f"available_blocks={available_blocks_p2}, " + f"running_count={len(self._running_requests)}" + ) + + if not all_scheduled_requests: + self._emit_schedule_decision_event( + event="iteration_end", + decision_result=None, + request_id=None, + token_budget=token_budget, + num_tokens=0, + available_blocks=int(self._config.num_blocks - self._num_allocated_blocks), + batch_request_ids=[], + request_num_tokens=[], + batch_size=0, + batch_num_tokens=0, + ) + return None + + ordered_scheduled_requests = waiting_scheduled + running_scheduled + ordered_num_tokens = waiting_tokens + running_tokens + + # Flow validation: log batch formation + total_tokens = sum(all_num_tokens) + new_admitted = len(waiting_scheduled) + resumed = len( + [r for r in all_scheduled_requests if getattr(r, "_preempted", False)] + ) + running_continued = len(running_scheduled) + batch_size = len(all_scheduled_requests) + + logger.info( + f"[BATCH_FORMATION] total_tokens={total_tokens}, " + f"new_admitted={new_admitted}, " + f"resumed={resumed}, " + f"running_continued={running_continued}, " + f"batch_size={batch_size}" + ) + self._emit_schedule_decision_event( + event="iteration_end", + decision_result=None, + request_id=None, + token_budget=token_budget, + num_tokens=total_tokens, + available_blocks=int(self._config.num_blocks - self._num_allocated_blocks), + batch_request_ids=[request.id for request in ordered_scheduled_requests], + request_num_tokens=ordered_num_tokens, + batch_size=batch_size, + batch_num_tokens=total_tokens, + ) + + return self._create_batch(ordered_scheduled_requests, ordered_num_tokens) + + def _schedule_decode_only(self) -> Optional[Batch]: + """ + Scheduling for DECODE cluster - two-phase scheduling matching vLLM v1. + + Phase 1: Schedule RUNNING requests (ongoing decode iterations) + Phase 2: Admit WAITING requests (new arrivals from prefill cluster) + + This matches vLLM v1's scheduling algorithm where requests must be + admitted from waiting queue to running queue before generating tokens. + + Returns: + Optional[Batch]: The scheduled batch + """ + logger = get_cluster_logger( + __name__, self._cluster_type.name if self._cluster_type else None + ) + + all_scheduled_requests: List[Request] = [] + all_num_tokens: List[int] = [] + preempted_requests: List[Request] = [] + waiting_scheduled: List[Request] = [] + waiting_tokens: List[int] = [] + token_budget = self._max_num_scheduled_tokens + available_blocks = int(self._config.num_blocks - self._num_allocated_blocks) + waiting_final_decode_count = self._count_final_fast_lane_requests( + self._waiting_requests, + final_predicate=self._is_final_decode_fast_lane_request, + ) + self._decode_iteration_reserved_slots_remaining = ( + self._final_decode_reserved_slots if waiting_final_decode_count > 0 else 0 + ) + + # Flow validation: log iteration start + logger.info( + f"[ITERATION_START] token_budget={token_budget}, " + f"running_count={len(self._running_requests)}, " + f"waiting_count={len(self._waiting_requests)}, " + f"available_blocks={available_blocks}, " + f"max_running_reqs={self._max_num_running_reqs}" + ) + self._emit_schedule_decision_event( + event="iteration_start", + decision_result=None, + request_id=None, + token_budget=token_budget, + available_blocks=available_blocks, + num_tokens=0, + ) + + # Flow validation: log memory state + total_blocks = int(self._config.num_blocks) + allocated_blocks = int(self._num_allocated_blocks) + usage_ratio = allocated_blocks / total_blocks if total_blocks > 0 else 0.0 + watermark = self._watermark_blocks + logger.info( + f"[MEMORY_STATE] total_blocks={total_blocks}, " + f"allocated_blocks={allocated_blocks}, " + f"free_blocks={available_blocks}, " + f"usage_ratio={usage_ratio:.4f}, " + f"watermark_blocks={watermark}" + ) + self._reclaim_borrowed_final_running_slots( + waiting_requests=self._waiting_requests, + final_predicate=self._is_final_decode_fast_lane_request, + reserved_slots=self._final_decode_reserved_slots, + lane_name="decode", + ) + + # Flow validation: log Phase 1 start + logger.info( + f"[PHASE1_START] running_count={len(self._running_requests)}, " + f"token_budget={token_budget}" + ) + + # === Phase 1: Schedule RUNNING requests === + token_budget, running_scheduled, running_tokens = ( + self._schedule_running_requests(token_budget, preempted_requests) + ) + all_scheduled_requests.extend(running_scheduled) + all_num_tokens.extend(running_tokens) + + # Flow validation: log Phase 1 end + available_blocks_p1 = int(self._config.num_blocks - self._num_allocated_blocks) + logger.info( + f"[PHASE1_END] scheduled_count={len(running_scheduled)}, " + f"preempted_count={len(preempted_requests)}, " + f"token_budget_remaining={token_budget}, " + f"available_blocks={available_blocks_p1}" + ) + + # === Phase 2: Admit WAITING requests (only if no preemption) === + if not preempted_requests: + # Flow validation: log Phase 2 start + logger.info( + f"[PHASE2_START] waiting_count={len(self._waiting_requests)}, " + f"token_budget={token_budget}, " + f"running_count={len(self._running_requests)}" + ) + + token_budget, waiting_scheduled, waiting_tokens = ( + self._schedule_decode_waiting_requests(token_budget) + ) + all_scheduled_requests.extend(waiting_scheduled) + all_num_tokens.extend(waiting_tokens) + + # Flow validation: log Phase 2 end + available_blocks_p2 = int( + self._config.num_blocks - self._num_allocated_blocks + ) + logger.info( + f"[PHASE2_END] admitted_count={len(waiting_scheduled)}, " + f"token_budget_remaining={token_budget}, " + f"available_blocks={available_blocks_p2}, " + f"running_count={len(self._running_requests)}" + ) + + if not all_scheduled_requests: + self._emit_schedule_decision_event( + event="iteration_end", + decision_result=None, + request_id=None, + token_budget=token_budget, + num_tokens=0, + available_blocks=int(self._config.num_blocks - self._num_allocated_blocks), + batch_request_ids=[], + request_num_tokens=[], + batch_size=0, + batch_num_tokens=0, + ) + return None + + # Match vLLM v1 output order: new admissions first, then running. + ordered_scheduled_requests = waiting_scheduled + running_scheduled + ordered_num_tokens = waiting_tokens + running_tokens + + # Flow validation: log batch formation + total_tokens = sum(all_num_tokens) + new_admitted = len( + [r for r in all_scheduled_requests if r not in running_scheduled] + ) + resumed = 0 # DECODE doesn't handle preempted requests (they come from prefill) + running_continued = len(running_scheduled) + batch_size = len(all_scheduled_requests) + + logger.info( + f"[BATCH_FORMATION] total_tokens={total_tokens}, " + f"new_admitted={new_admitted}, " + f"resumed={resumed}, " + f"running_continued={running_continued}, " + f"batch_size={batch_size}" + ) + self._emit_schedule_decision_event( + event="iteration_end", + decision_result=None, + request_id=None, + token_budget=token_budget, + num_tokens=total_tokens, + available_blocks=int(self._config.num_blocks - self._num_allocated_blocks), + batch_request_ids=[request.id for request in ordered_scheduled_requests], + request_num_tokens=ordered_num_tokens, + batch_size=batch_size, + batch_num_tokens=total_tokens, + ) + + return self._create_batch(ordered_scheduled_requests, ordered_num_tokens) + + def _schedule_decode_waiting_requests( + self, token_budget: int + ) -> Tuple[int, List[Request], List[int]]: + """ + Phase 2 for DECODE cluster: Admit requests from waiting queue. + + This method handles requests that have arrived from the prefill cluster + and are waiting to be admitted to the running queue for decode iterations. + Matches vLLM v1's Phase 2 scheduling behavior. + + Args: + token_budget: Remaining token budget for this iteration + + Returns: + Tuple of (remaining_budget, scheduled_requests, num_tokens_list) + """ + logger = get_cluster_logger(__name__, self._cluster_type.name) + scheduled: List[Request] = [] + num_tokens_list: List[int] = [] + + fast_lane_decode_enabled = self._cluster_type == ClusterType.DECODE and ( + self._final_decode_reserved_slots > 0 + ) + waiting_queue = self._build_decode_waiting_queue() + skipped_waiting_requests: deque[Request] = deque() + + self._current_iteration_token_budget = token_budget + while waiting_queue and token_budget > 0: + self._current_iteration_token_budget = token_budget + final_waiting_count = ( + self._count_final_fast_lane_requests( + waiting_queue, + final_predicate=self._is_final_decode_fast_lane_request, + ) + if fast_lane_decode_enabled + else 0 + ) + has_final_waiting = final_waiting_count > 0 + # Check max concurrent requests limit + if len(self._running_requests) >= self._max_num_running_reqs: + logger.debug( + f"[VLLMv1Engine][DECODE] Phase 2: max running requests " + f"reached ({self._max_num_running_reqs}), stopping admission" + ) + break + + request = waiting_queue[0] + is_final_decode_request = fast_lane_decode_enabled and ( + self._is_final_decode_fast_lane_request(request) + ) + is_hidden_decode_request = ( + fast_lane_decode_enabled and not is_final_decode_request + ) + + if ( + is_hidden_decode_request + and has_final_waiting + and self._decode_iteration_reserved_slots_remaining > 0 + and len(self._running_requests) + >= ( + self._max_num_running_reqs + - self._decode_iteration_reserved_slots_remaining + ) + ): + waiting_queue.popleft() + skipped_waiting_requests.append(request) + continue + + num_new_tokens = self._get_request_next_num_tokens(request) + + # Apply max_model_len limit + scheduler_num_computed_tokens = self._get_scheduler_num_computed_tokens( + request + ) + max_allowed = self._max_model_len - scheduler_num_computed_tokens + num_new_tokens = min(num_new_tokens, max_allowed) + + # Apply token budget limit + num_new_tokens = min(num_new_tokens, token_budget) + + if num_new_tokens <= 0: + # Request has reached max length, remove from queue + waiting_queue.popleft() + logger.debug( + f"[VLLMv1Engine][DECODE] Phase 2: req={request.id} " + f"reached max length, removing from waiting queue" + ) + continue + + # Try to allocate (no preemption for waiting requests in Phase 2) + if not self._can_allocate_request( + request, + num_new_tokens, + scheduler_num_computed_tokens=scheduler_num_computed_tokens, + ): + # Cannot allocate - stop admitting new requests + logger.debug( + f"[VLLMv1Engine][DECODE] Phase 2: cannot allocate " + f"req={request.id}, stopping admission" + ) + break + + # Check if this request was previously preempted + was_preempted = getattr(request, "_preempted", False) + + # Remove from waiting queue and allocate + waiting_queue.popleft() + + # Record leaving waiting queue for waiting time tracking + request.on_leave_waiting_queue( + self._current_schedule_time, self._cluster_type + ) + + self._allocate_request( + request, + num_new_tokens, + scheduler_num_computed_tokens=scheduler_num_computed_tokens, + ) + self._advance_scheduler_num_computed_tokens(request, num_new_tokens) + + # Add to running requests + self._running_requests.append(request) + + # Clear preempted flag if set + if was_preempted: + request._preempted = False + + scheduled.append(request) + num_tokens_list.append(num_new_tokens) + token_budget -= num_new_tokens + self._current_iteration_token_budget = token_budget + if is_final_decode_request: + self._decode_iteration_reserved_slots_remaining = max( + self._decode_iteration_reserved_slots_remaining - 1, + 0, + ) + + # Flow validation: log ADMISSION event (matching vLLM v1) + logger.info( + f"[ADMISSION] req={request.id} admitted, " + f"num_tokens={num_new_tokens}, " + f"running_count={len(self._running_requests)}, " + f"token_budget_remaining={token_budget}" + ) + available_blocks_admission = int( + self._config.num_blocks - self._num_allocated_blocks + ) + self._emit_schedule_decision_event( + event="decision", + decision_result="ADMISSION", + request_id=request.id, + token_budget=token_budget, + available_blocks=available_blocks_admission, + num_tokens=num_new_tokens, + ) + + # Flow validation: log preemption recovery if applicable + if was_preempted: + # For DECODE cluster, preempted requests need full recomputation + # from their original prefill tokens + recompute_tokens = request.num_prefill_tokens + logger.info( + f"[PREEMPTION_RECOVERY] req={request.id}, " + f"was_preempted=True, " + f"recompute_tokens={recompute_tokens}" + ) + + if skipped_waiting_requests: + waiting_queue.extend(skipped_waiting_requests) + self._waiting_requests = list(waiting_queue) + + return token_budget, scheduled, num_tokens_list + + def _should_use_dense_decode_attn_metadata_wave(self) -> bool: + """Return whether dense PP=1 PDAF needs one DES macro-wave batch.""" + if self._cluster_type != ClusterType.DECODE_ATTN: + return False + if self._replica_is_moe: + return False + return ( + self._num_stages == 1 + and self._af_pipeline_num_micro_batch > 1 + ) + + def _schedule_decode_attn_only( + self, is_micro_batch: bool = True + ) -> Optional[Batch]: + """ + Scheduling for DECODE_ATTN cluster in PD-AF disaggregation mode. + + This method is called ONLY for Priority 2 scheduling (new micro-batch formation). + Priority 1 (AF immediate inflight batches) is handled by on_schedule() directly. + + Two-level scheduling strategy based on decode step: + - Incomplete decode step (is_mb_last_layer=False): batch-level, via _af_immediate_batch_queue + - Complete decode step (is_mb_last_layer=True): request-level, via this method + + Phase 1: Schedule running requests (ongoing decode from _running_requests) + - For each request in _running_requests: + - Calculate new tokens to process (usually 1 for decode) + - Allocate memory for new tokens + - If allocation fails: trigger preemption following vLLM v1 behavior + - Add to scheduled batch + + Phase 2: Admit new requests from _waiting_requests (if Phase 1 had no preemption) + - Check memory budget and token budget + - Form micro-batch with layer-consistent grouping (fix: do we need it? all requests are layer-0) + - All new requests start at layer 0, so naturally layer-consistent + + Layer-consistent grouping is implicitly guaranteed: + - Running requests have _completed_layer_count = 0 (reset after decode step completion) + - New requests also start at layer 0 + - Therefore, all requests in a micro-batch are layer-consistent + + Note on initial state: + - On first scheduling, _running_requests is empty, so Phase 1 produces no output + - Phase 2 will admit new requests from _waiting_requests to _running_requests + - Subsequent decode steps will have Phase 1 populated from previous on_batch_end() + + Args: + is_micro_batch: Should always be True for DECODE_ATTN + + Returns: + Optional[Batch]: The scheduled micro-batch, or None if no requests available + """ + logger = get_cluster_logger(__name__, self._cluster_type.name) + + if is_micro_batch and self._af_pending_micro_batches: + return self._af_pending_micro_batches.popleft() + + # Enable preemption for DECODE_ATTN to handle memory pressure + # Preemption logic follows vLLM v1 behavior for running requests + preemption_enabled = True + preempted_requests: List[Request] = [] + + # Phase 1: Schedule running requests + scheduled_requests = [] + scheduled_tokens = [] + + # Get request IDs to exclude (already scheduled in inflight batches) + continuation_request_ids = getattr(self, "_continuation_request_ids", set()) + # _running_requests in inclued reqs: inflight(layer!=0) req, completed req (really?) + + for request in self._running_requests: + # ISSUE-008 FIX: Check batch size limit at start of Phase 1 loop. + # This prevents scheduling more requests than _micro_batch_size allows, + # ensuring proper batch size enforcement in DECODE_ATTN cluster. + if len(scheduled_requests) >= self._micro_batch_size: + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 1: reached micro_batch_size limit " + f"({self._micro_batch_size}), stopping" + ) + break + + if request.completed: + # why would a running request be completed but still in _running_requests? + raise ValueError(f"Request {request.id} is already completed") + continue + + # CRITICAL FIX: Only schedule requests ready for new decode step (layer_count = 0) + # Requests with layer_count > 0 are still in-flight (mid-layer processing) + # and should NOT be re-scheduled until their current decode step completes. + # This ensures layer-consistent grouping in micro-batches. + if request.completed_layer_count != 0: + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 1: skipping in-flight req={request.id} " + f"with layer_count={request.completed_layer_count} (not ready for new decode step)" + ) + continue + + # Requests in active A->F->A roundtrip must not be re-scheduled until + # F->A transfer end clears the in-flight marker. + if request.af_roundtrip_inflight: + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 1: skipping req={request.id} " + f"(AF roundtrip still in-flight)" + ) + continue + + # CRITICAL FIX: Skip requests already scheduled in continuation batches (Priority 1) + # This prevents the same request from being scheduled into multiple batches + if request.id in continuation_request_ids: + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 1: skipping req={request.id} " + f"(already in continuation batch from Priority 1)" + ) + continue + + # Calculate tokens for decode: usually 1 + num_new_tokens = 1 + + # Try to allocate memory + if self._can_allocate_request(request, num_new_tokens): + self._allocate_request(request, num_new_tokens) + scheduled_requests.append(request) + scheduled_tokens.append(num_new_tokens) + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 1: scheduled running req={request.id}, " + f"num_tokens={num_new_tokens}" + ) + else: + # Memory pressure - try allocation with preemption + if not preemption_enabled: + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 1: cannot allocate req={request.id}, " + f"preemption disabled, skipping" + ) + continue + + # Try to allocate with preemption (follows vLLM v1 behavior) + preempted_count_before = len(preempted_requests) + success = self._try_allocate_with_preemption( + request, num_new_tokens, preempted_requests + ) + self._current_iteration_token_budget = ( + self._rollback_current_iteration_preempted_requests( + scheduled_requests=scheduled_requests, + scheduled_num_tokens=scheduled_tokens, + newly_preempted_requests=preempted_requests[ + preempted_count_before: + ], + token_budget=self._current_iteration_token_budget, + ) + ) + if success: + scheduled_requests.append(request) + scheduled_tokens.append(num_new_tokens) + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 1: scheduled req={request.id} " + f"after preemption, num_tokens={num_new_tokens}" + ) + else: + # Request itself was preempted or no victim available + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 1: req={request.id} " + f"preempted or allocation failed" + ) + + # Check micro-batch size limit + remaining_slots = self._micro_batch_size - len(scheduled_requests) + + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] After Phase 1: scheduled={len(scheduled_requests)}, " + f"remaining_slots={remaining_slots}, micro_batch_size={self._micro_batch_size}" + ) + + # Phase 2: Admit new requests (only if no preemption occurred) + if len(preempted_requests) == 0 and remaining_slots > 0: + for request in list(self._waiting_requests): + if remaining_slots <= 0: + break + + # New requests start at layer 0 - naturally layer-consistent + assert request.completed_layer_count == 0, ( + f"New request {request.id} should have completed_layer_count=0, got {request.completed_layer_count}" + ) + + # Allocate decode token + num_tokens = 1 + if self._can_allocate_request(request, num_tokens): + self._waiting_requests.remove(request) + request.on_leave_waiting_queue( + self._current_schedule_time, self._cluster_type + ) + self._allocate_request(request, num_tokens) + self._running_requests.append(request) + scheduled_requests.append(request) + scheduled_tokens.append(num_tokens) + remaining_slots -= 1 + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 2: admitted new req={request.id}, " + f"num_tokens={num_tokens}, running_count={len(self._running_requests)}" + ) + else: + logger.debug( + f"[VLLMv1Engine][DECODE_ATTN] Phase 2: cannot allocate req={request.id}, " + f"stopping admission" + ) + break + + # (scheduled_requests, scheduled_tokens) is the scheduler's output + # we should use scheduler_output to creat microbatch for pd-af + + # Create batch if we have scheduled requests + if scheduled_requests: + logger.info( + f"[VLLMv1Engine][DECODE_ATTN] Created micro-batch with {len(scheduled_requests)} requests" + ) + + num_reqs = len(scheduled_requests) + num_stages = self._af_pipeline_num_micro_batch + if num_stages is None or num_stages <= 0: + raise ValueError( + "af_pipeline_num_micro_batch must be positive for DECODE_ATTN" + ) + + replay_decode_token_index = int( + scheduled_requests[0].current_decode_token_index + ) + decode_attn_cohort_id = self._allocate_decode_attn_cohort_id() + decode_attn_cohort_request_ids = tuple( + request.id for request in scheduled_requests + ) + cohort_state = self._get_decode_attn_active_cohort_states().setdefault( + decode_attn_cohort_id, + { + "all_request_ids": set(), + "pending_request_ids": set(), + "af_phase": "local_attn", + "active_stage_indices": set(), + "stage_phases": {}, + "stage_current_layer_ids": {}, + }, + ) + cohort_state["all_request_ids"].update( + decode_attn_cohort_request_ids + ) + cohort_state["pending_request_ids"].update( + decode_attn_cohort_request_ids + ) + cohort_state["current_layer_id"] = int( + scheduled_requests[0].completed_layer_count + ) + + # StepFun-vLLM partitioning: split requests by stage + if num_reqs >= num_stages: + num_reqs_per_stage = num_reqs // num_stages + stage_reqs_start_loc = [ + num_reqs_per_stage * i for i in range(num_stages + 1) + ] + stage_reqs_start_loc[-1] = num_reqs + else: + stage_reqs_start_loc = list(range(num_reqs + 1)) + + afd_stage_metadata = None + if self._cluster_type == ClusterType.DECODE_ATTN and num_stages > 0: + from frontier.config import global_vars + from frontier.entities.batch import AFDStageMetadata + + use_cuda_graph = global_vars.get_use_cuda_graph() + cudagraph_capture_sizes = global_vars.get_cudagraph_capture_sizes() + if use_cuda_graph and cudagraph_capture_sizes is None: + max_num_seqs = ( + self._micro_batch_size + if hasattr(self, "_micro_batch_size") + else 64 + ) + cudagraph_capture_sizes = [1, 2, 4] + [ + 8 * i for i in range(1, max_num_seqs // 8 + 1) + ] + + afd_stage_metadata = AFDStageMetadata.from_batch_params( + num_reqs=num_reqs, + num_tokens_per_req=scheduled_tokens, + num_stages=num_stages, + dp_stage_max_tokens=None, + use_cuda_graph=use_cuda_graph, + cudagraph_capture_sizes=cudagraph_capture_sizes, + ffn_use_cuda_graph=use_cuda_graph, + ffn_cudagraph_capture_sizes=cudagraph_capture_sizes, + ) + + first_micro_batch = None + if self._should_use_dense_decode_attn_metadata_wave(): + macro_batch = self._create_batch( + scheduled_requests, + scheduled_tokens, + ) + macro_batch.afd_stage_idx = 0 + macro_batch.afd_stage_represents_all_stages = True + macro_batch.replay_decode_token_index = replay_decode_token_index + macro_batch.decode_attn_cohort_id = decode_attn_cohort_id + macro_batch.decode_attn_cohort_request_ids = ( + decode_attn_cohort_request_ids + ) + if afd_stage_metadata is not None: + macro_batch.afd_stage_metadata = afd_stage_metadata + cohort_state["active_stage_indices"].add(0) + cohort_state["stage_phases"][0] = "local_attn" + cohort_state["stage_current_layer_ids"][0] = int( + scheduled_requests[0].completed_layer_count + ) + return macro_batch + + shared_decode_attn_global_id = int(self._batch_creation_counter) + for stage_idx in range(len(stage_reqs_start_loc) - 1): + start_idx = stage_reqs_start_loc[stage_idx] + end_idx = stage_reqs_start_loc[stage_idx + 1] + stage_requests = scheduled_requests[start_idx:end_idx] + stage_tokens = scheduled_tokens[start_idx:end_idx] + micro_batch = self._create_batch(stage_requests, stage_tokens) + micro_batch.set_global_id(shared_decode_attn_global_id) + micro_batch.afd_stage_idx = stage_idx + micro_batch.replay_decode_token_index = int( + stage_requests[0].current_decode_token_index + ) + micro_batch.decode_attn_cohort_id = decode_attn_cohort_id + micro_batch.decode_attn_cohort_request_ids = ( + decode_attn_cohort_request_ids + ) + if afd_stage_metadata is not None: + micro_batch.afd_stage_metadata = afd_stage_metadata + normalized_stage_idx = int(stage_idx) + cohort_state["active_stage_indices"].add(normalized_stage_idx) + cohort_state["stage_phases"][normalized_stage_idx] = "local_attn" + cohort_state["stage_current_layer_ids"][normalized_stage_idx] = int( + scheduled_requests[0].completed_layer_count + ) + if first_micro_batch is None: + first_micro_batch = micro_batch + else: + self._af_pending_micro_batches.append(micro_batch) + + return first_micro_batch + + + logger.debug("[VLLMv1Engine][DECODE_ATTN] No requests to schedule") + return None diff --git a/task_memory/task_2026-09-21_issue26_correctness_pr/plan.md b/task_memory/task_2026-09-21_issue26_correctness_pr/plan.md new file mode 100644 index 00000000..614ab7ba --- /dev/null +++ b/task_memory/task_2026-09-21_issue26_correctness_pr/plan.md @@ -0,0 +1,841 @@ +# Issue 26 Correctness PR — Plan + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Landed the execution specification verbatim (Section "Execution Specification" below) and recorded the amendments agreed with the user before Step 0. | + +## Amendments (authoritative where they differ from the specification below) + +Facts were verified against `origin/main` at `1f694f7c549aa3aeeb7c5bbae04e119c09167a77`, `origin/bug/ttft-check` at `a7b3320fe9b8b083ee86b91dae3d6838f4443d91`, and the host `kun-workspace-vgen2` on 2026-09-21. Each row records the specification clause, the verified fact or user decision, and the resulting rule. + +| # | Spec clause | Verified fact / decision | Amended rule | +| --- | --- | --- | --- | +| A1 | §3.1, §3.2 `docs/development/issue26-correctness-pr/` | User decision: task records live in `task_memory/` of the PR worktree; no new `docs/` tree. `task_memory/` is ignored by `.gitignore`, and Git cannot re-include a child of an excluded directory. | Records live in `task_memory/task_2026-09-21_issue26_correctness_pr/`. `.gitignore` uses `task_memory/*` plus a narrow `!task_memory//` exception for this directory and the refactor directory in A3. Only Markdown and small text evidence are placed there; generated JSON/CSV stay in the scratch root. The CLAUDE.md files `requirements.md`, `plan.md`, `progress.md`, `summary.md` are kept; the specification's `review.md` and `validation.md` are supporting files in the same directory. | +| A2 | §5 execution order, §4.1 oversized modules | User decision (Q4=c, Q8=a, Q9=b, Q10=b): the four touched modules above 2,000 lines (`frontier/config/config.py` 5720, `frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py` 5138, `frontier/execution_time_predictor/shared_prediction_model_manager.py` 4614, `frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py` 3539) receive a full cleanup-first and functional split. | Stacked PRs. `refactor/oversized-module-split` branches from main and carries only behavior-preserving cleanup and splits. `fix/issue26-correctness-pr` branches from the refactor branch. Both draft PRs are opened; the correctness PR's base is the refactor branch and is retargeted to main after the refactor PR merges. Steps 2–7 of this specification start only after the refactor branch's fidelity matrix passes. The other five modules above 2,000 lines are out of scope. | +| A3 | §3.2 four documents | Decision Q12=b. | The refactor PR has its own records in `task_memory/task_2026-09-21_oversized_module_split/`. This plan references that directory instead of duplicating its content. | +| A4 | §6 baseline file `tests/unit/test_open_source_release_arch_guard.py` | The file does not exist on main (`git ls-files` returns nothing). `AGENTS.md` still references it; recorded as a deferred documentation drift, not fixed here. | Baseline selection: `tests/unit/test_cluster_scheduler_dp_lanes.py`, `tests/unit/test_colocation_release_review_contracts.py`, `tests/unit/test_config_owned_contracts.py`, `tests/unit/test_stage_execution_time.py`, `tests/unit/test_stage_finalized_contract.py`, `tests/unit/test_moe_routing_runtime.py`, plus the co-location and PDD dense offline example wrappers. | +| A5 | §6 "use the existing CPU environment" | No conda and no Frontier-capable Python environment existed on this host. `uv` is available at `/data/ycfeng/toolchain/bin/uv`. | Task environment: `/data/ycfeng/envs/frontier-py310` created with `uv venv --python 3.10` and `uv pip install -e ".[test]"` from the refactor worktree. `PYTHONPATH` points at the active worktree. Scratch root: `FRONTIER_TMP_ROOT=/data/ycfeng/tmp/issue26-correctness-pr`. | +| A6 | §7.2 separate read-only vLLM checkout | No local vLLM-BS checkout existed. User decision: clone under `.real-engine/`. | Reference checkout: `/.real-engine/vLLM-BS` detached at `ea95f571e20937c7c908c6d59ddd1cd6bf9268f1`; `.real-engine/` is listed in `.git/info/exclude` (not in the tracked `.gitignore`). Read-only. | +| A7 | §13 Step 7 | `fwyc0573/frontier-htsim` exposes only `main` at `b8518afcc310f0fe0e3ce52ba6b4f0bf57a3be04`; commit `e564935d…` returns HTTP 422 (not found). No local clone containing it exists on this host (the historical calibration worktree is absent). User decision Q2=c: decide at Step 7. | Step 7 starts from these facts. Reconstruction in a companion branch requires the user's authorization at that point; otherwise the package is `EXCLUDED`. | +| A8 | §12 Step 6 native validation, §1 execution environment | `/data/ycfeng/stepfun-env-handbook/stepmind-python-rjob.md` exists and StepMind `RJobBackend` GPU workers are available. User rule: simulator runs stay on the local CPU; GPU workers are used only for the Step 6 native numerical check and for any additional profiling CSVs. | Unchanged Step 6 requirements; the GPU path is StepMind per the runbook. | +| A9 | §3.3 publication, §3.5 approvals | User authorized (Q5) pushing both branches to `origin` (`NetX-lab/Frontier`) and creating/updating the two draft PRs for the whole task. Not authorized: merge, force-push, history rewrite, closing Issue 26. | Checkpoint pushes need no further approval. Q6=a: this session stops after the Step 0 push for user review. | +| A10 | §2.1 candidate snapshot | Verified: merge-base is `d71ad80b…`; the final candidate commit `a7b3320` touches only `task_memory/` (1006 files, no source); the candidate gitlink is `e564935d…`; donor design documents exist at the archive top level; `tests/unit/test_moe_routing_runtime.py` already exists on main and is modified by the candidate. | No change to the specification; facts recorded for Step 1. | +| A11 | §7.1 candidate scope | The candidate adds ~80 `tests/e2e|integration|performance/issue26_*` experiment scripts. | Default disposition `DROP` unless Step 1 finds a specific reusable helper (the specification already names `tests/integration/issue26_dp_coordinator_reference.py` for review). | + +## Execution Specification (verbatim copy of `.local-draft/Frontier_Issue26_Correctness_PR_Execution_Spec_2026-09-21.md`) + +# Frontier Issue 26: Correctness Fixes for Main + +## Execution specification for Claude Code + +**Prepared:** 2026-09-21 +**Task type:** A new, independently reviewable pull request against `NetX-lab/Frontier:main` +**Execution environment:** CPU master by default; narrowly scoped GPU-worker tests only when native kernel correctness requires them +**Status:** Ready to begin implementation. This document records a source-reviewed plan, not completed implementation or newly executed tests. + +> Start a new worktree and branch from freshly fetched `origin/main`. Extract and improve the justified fixes from `bug/ttft-check`; do not merge that branch wholesale. Read the relevant Frontier and vLLM 0.10.2 code before editing. Complete the steps below in order, publish code, tests, and task documentation after each meaningful checkpoint, and provide a clear user-facing report. Do not resume TTFT calibration or run a vLLM serving benchmark. + +## 1. Objective and boundaries + +Deliver a small, readable set of correctness fixes that stands on its own without a Frontier-versus-vLLM latency comparison. The source branch is evidence and a source of candidate changes, not an implementation to copy without review. + +The intended fixes are: + +| Work package | Intended outcome | Inclusion rule | +| --- | --- | --- | +| Round-robin DP placement | Preserve DP rotation across separate scheduling calls. | Core scope. | +| Shared monolithic forward execution | Allow prefill, decode, and mixed source batches to complete one shared EP forward without splitting synchronization identity or borrowing another source's timing. | Core scope; port as one coherent behavior change. | +| vLLM-style DP request placement | Offer an opt-in, explicitly bounded implementation of vLLM 0.10.2 request-count selection and delayed count publication. | Include after source and event-integration checks. Keep existing defaults. | +| Routing implementation identity | Keep expert-load distribution separate from the implementation used to predict routing cost; prevent model/cache sharing across incompatible implementations. | Core scope; carry the identity through the complete selection chain. | +| Legacy fused-MoE profiling | Execute the missing gated activation and local output reduction, with accurate measurement ownership. | Include after CPU checks and required native numerical checks; do not overwrite main's newer backend support. | +| Optional collective backend input handling | Distinguish zero payload from missing payload; reject negative input; preserve explicit CLI precedence. | Conditional on a reviewable, remotely fetchable submodule fix and real CPU runner tests. | + +**Success is semantic correctness, regression safety, maintainability, and reviewability. A TTFT error threshold is not an acceptance criterion.** + +### Explicitly excluded + +Do not launch vLLM servers, replay the historical 100-request calibration workload, compare serving TTFT/TPOT/E2E results, or conduct Nsight/Kineto investigations. Do not retune collective latency, bandwidth, CPU overhead, or compute scaling. Do not add an empirical residual to make simulated latency match a measurement. + +Do not implement a new physical communication execution model, change the default ideal EP accounting, adopt task-local profiling CSVs as official datasets, or force-add the old task archive. Do not rewrite unrelated scheduler, profiling, or training infrastructure. Do not change the vLLM reference to make a test pass. Do not merge the final PR or close Issue 26: this PR addresses selected correctness defects, not the entire calibration issue. + +A complete **Frontier-only CPU simulation** is allowed and required for relevant lifecycle checks. That is different from a vLLM serving comparison. + +## 2. Source snapshot and recovered context + +### 2.1 Revisions reviewed when this specification was prepared + +| Repository / reference | Reviewed revision | Role | +| --- | --- | --- | +| `NetX-lab/Frontier`, `main` | `1f694f7c549aa3aeeb7c5bbae04e119c09167a77` | Integration baseline at specification time. | +| `NetX-lab/Frontier`, `bug/ttft-check` | `a7b3320fe9b8b083ee86b91dae3d6838f4443d91` | Candidate fixes and historical evidence. | +| Candidate's parent before the large evidence import | `b7f8d055461d9208b8ae57eceeb6c0246cbc8d3c` | Convenient source-code comparison point; verify rather than assume that the final commit is documentation-only. | +| Reviewed common ancestor | `d71ad80b0800880808a0857fd30477e6d96592c6` | Separates candidate changes from subsequent main development. | +| `fwyc0573/vLLM-BS`, `feature/frontier-comparison-instrumentation` | `ea95f571e20937c7c908c6d59ddd1cd6bf9268f1` | User-designated vLLM reference. | +| Frontier optional backend in reviewed main | `b8518afcc310f0fe0e3ce52ba6b4f0bf57a3be04` | Existing `collective-sim` gitlink. | +| Optional backend in candidate | `e564935d3874d8c71b52a554ab7c9a72e5e19f68` | Candidate gitlink; a fresh remote API lookup could not resolve this commit. Recheck in Step 7. | + +Fetch again at execution start. Use the latest fetched main as the new branch base and record its actual SHA. Pin the candidate and reference revisions used in the audit. If either has advanced, inspect the changes; do not silently move a reference during an implementation step. + +Main and the candidate have diverged. Main contains newer execution-time reporting, model support, profiling API selection, and accelerator handling. A change appearing in the candidate does not justify reverting those main changes. Use both a common-ancestor diff and a direct main-to-candidate comparison. [R1–R3] + +### 2.2 What is already known, and what is not + +The old task found useful independent defects: DP placement could restart at lane zero; shared monolithic work could split by local request phase; a legacy MoE profiling path omitted real arithmetic; routing runtime identity was insufficiently separated from load distribution. The task also accumulated timing experiments whose measurement conditions differ. Numerical calibration remains unresolved. [R2, R4] + +Do not copy historical `PASS` counts into this task as fresh validation. Old tests are design references. Re-run or improve them on the new branch. In particular, some shared-forward tests replace ownership or wave behavior with stubs; those are not sufficient evidence for the real event lifecycle. The old MoE native test covers a limited BF16 case, not every dtype or backend. [R6, R7] + +The old experiment checkout also used local vLLM overlays beyond the remote reference. They are not prerequisites for this PR. Only recover an overlay when a narrowly scoped native test genuinely needs a compatibility fix, and record exactly what it changes. Never use an undocumented overlay as the reference implementation. + +## 3. Working agreement + +### 3.1 New worktree, unchanged existing work + +Create a **new named branch in a new worktree**. Do not switch, reset, clean, rebase, or reuse the calibration worktree. Do not delete another agent's files or modify its environment. A dirty existing checkout is a reason to preserve its state, not permission to clean it. + +Suggested names: + +```text +Branch: fix/issue26-correctness-pr +Worktree: /.worktrees/issue26-correctness-pr +Docs: docs/development/issue26-correctness-pr/ +``` + +These are task setup names, not paths to embed in production code. If a name already exists, inspect it. Resume it only when it is demonstrably this task's already-created worktree; otherwise choose a new descriptive suffix. Do not overwrite it. + +### 3.2 Four tracked task documents + +Keep the task's authoritative, GitHub-reviewable state in four files: + +| File | Contents | +| --- | --- | +| `plan.md` | This specification, copied in full, with a short amendment history when decisions change. | +| `progress.md` | Current branch/base, step status, most recent checkpoint, blockers, exact next action, and concise chronological updates. | +| `review.md` | Source references, candidate-change dispositions, design decisions, and the final code-review findings. | +| `validation.md` | Commands, environments, test selections, results, baseline failures, native-test evidence, and validation limits. | + +Main ignores `task_memory/` and `repairs/`. Do not rely on those paths for reviewable progress or alter their global ignore policy. Put this task's records under `docs/` and verify that they are tracked. Do not create an additional reporting framework or a separate document for every tiny action. [R5] + +Store reproducible test inputs in the existing test layout. The repository also broadly ignores generated JSON/CSV files under `tests/`; verify that any intentionally committed small fixture is actually tracked. Prefer in-test data or a narrowly scoped ignore exception over force-adding generated output directories. + +### 3.3 Checkpoint = implementation + validation + documentation + publication + +At every completed key step: + +1. Review the actual diff, run the required checks, and distinguish new failures from baseline failures. +2. Update the four documents as applicable. Include what changed, why, what was observed, and remaining limits. +3. Commit the implementation, tests, and documentation together, or as a small adjacent set of coherent commits. +4. Push **all those commits to the code repository branch**, not merely documentation or a separate notes repository. +5. Verify the remote branch points to the pushed local HEAD and provide GitHub commit/branch links to the user. + +The initial setup checkpoint and final documentation-only corrections may naturally have no production changes. An implementation checkpoint must not be described as published when its code is still uncommitted or local-only. + +Do not repeatedly amend or force-push commits the user may already be reviewing. Keep published history additive. A final history rewrite or PR merge requires explicit user approval. + +A commit cannot contain its own final SHA in a tracked receipt. Report the new SHA in the user-facing update, and reference it in subsequent records when useful. Do not generate recursive receipt-only commits merely to embed a document's own commit hash. + +### 3.4 Required user-facing update + +After each key step, after a blocker, and before ending a work session, report: + +```text +Completed: +Reason: +Evidence: +Conclusion: +Published: +Next: +``` + +Use complete sentences and actual findings, not a command transcript or a generic statement that progress was made. During longer work, give short updates at meaningful boundaries. A failed or skipped test is not a pass; a successful local commit is not a successful push. + +If push fails, retain the local changes, record the error without exposing credentials, and report `LOCAL_ONLY` publication state. Resolve repository access before treating the checkpoint as delivered. Do not accumulate several completed but unpublished work packages without informing the user. + +### 3.5 High-value decisions: question the user when evidence cannot decide + +Do not ask the user to resolve facts available in source, tests, or existing records. Investigate those first. Do not ask routine permission to perform the work already authorized here. + +When a consequential choice remains unresolved, ask **one focused question at a time**. State the relevant evidence, the alternatives, the effect on correctness/scope/compatibility, your recommended choice, and what remains blocked. Continue independent work without silently implementing the disputed choice. + +Examples that justify a decision checkpoint: + +| Trigger | Decision to put to the user | +| --- | --- | +| Correcting the MoE measurement changes accepted dataset semantics, and existing metadata cannot distinguish obsolete rows. | Approve a narrowly scoped data-contract change and migration, or defer that work package? Do not silently reinterpret old data. | +| The new DP strategy advertises configurations for which the available step identity is not valid. | Support those configurations with a small source-backed change, or explicitly narrow the opt-in capability? Do not quietly make it Qwen/DP2-specific. | +| Required native validation is unavailable or exposes a precision/backend incompatibility. | Keep the affected package blocked in a draft PR, or split it into a later PR? Do not waive numerical validation implicitly. | +| The optional submodule fix cannot be published to its configured remote. | Omit that conditional package, or authorize the necessary companion-repository work? Do not publish an unreachable gitlink. | +| A safe fix requires substantially widening the agreed public API or changing defaults. | Present the smallest viable alternatives before proceeding. | + +Do not stop the whole task for a local naming or formatting choice. Do not turn this mechanism into approval requests after every step. + +## 4. Implementation standards + +### 4.1 Read the owning code before extracting a helper + +For every changed runtime path, inspect its producer, consumer, lifecycle owner, existing helper, and focused tests. In `review.md`, briefly explain why the selected ownership is correct and which candidate mechanisms were removed or simplified. + +Use existing configuration types, event handling, request transitions, predictor contracts, metrics ownership, cache utilities, and temporary-output helpers. Do not introduce another scheduler framework, duplicate model registry, generic compatibility layer, or task-specific runtime switch. + +Respect repository `AGENTS.md` and applicable nested guidance. For a touched oversized module, identify the relevant responsibility and an existing or sensible extraction boundary. Prefer a small cohesive extraction when it simplifies the fix. Do not split unrelated code solely to hit a line-count target. [R1] + +### 4.2 Validate at boundaries; trust established internal state + +Validate external configuration, loaded artifacts, native backend capabilities, and untrusted event inputs at their owning boundary. Once required state is established, access it directly. + +Avoid repeated `getattr(..., None)`, `hasattr`, fallback assignments, and `if x is None` ladders for fields that must exist. Do not turn a missing required execution record into zero work, an idle batch, an empty list, or a default routing implementation. + +Retain legitimate optional states: an absent previous count snapshot, an optional expert map, or a genuinely disabled reporting path can be meaningful. Do not mechanically delete `None` handling or assert optional state always exists. Document the invariant rather than adding defensive checks at every call site. + +Do not modify production behavior to accommodate incomplete test stubs. Build an appropriate test fixture or use the real object. + +### 4.3 Names and structure + +Use names that describe requests, DP lanes, shared forward steps, load reports, expert outputs, and measurement identity. Reuse existing terminology where it is part of Frontier's API. Prefer short functions with one state transition or computation. + +Examples of suitable names are `RequestLoad`, `select_dp_lane`, `publish_request_counts`, `complete_shared_forward`, and `routing_runtime_path`; these are examples, not instructions to rename established APIs unnecessarily. + +Do not introduce opaque labels such as `arm` or `protocol` for new functions, variables, flags, or work packages. Do not use task numbers, dates, model names, or calibration experiment names in production identifiers. Avoid vague suffixes such as `manager`, `adapter`, or `context` unless the object genuinely owns that responsibility. + +### 4.4 Constants and supported scope + +Derive DP/TP/EP sizes, replica IDs, layer bounds, tensor shapes, and dtype from existing contracts. Do not embed H200, Qwen, DP2, TP4, EP8, 4096 tokens, or historical request IDs in production logic. + +A version-defined behavior is different from a calibration constant. The reference's waiting-count weight and publication intervals may appear as clearly named source-backed constants. Do not expose tuning flags just to eliminate a literal. Explicit restrictions such as one frontend and PP1 must be honest capability boundaries, not hidden assumptions. + +### 4.5 Preserve current-main behavior outside each fix + +Retain main's supported co-location, sequential PDD, and sequential PD-AF behavior; supported model families; existing accelerator selection; functional and legacy fused-MoE entry points; and reporting/measurement families. + +In particular, do not restore an obsolete execution-time builder from the candidate over main's current `StageExecutionTime` and demand-driven reporting logic. Keep reporting disabled paths cheap. Do not add per-layer rescans or redundant predictor queries when existing owned state is sufficient. + +Use a small paired CPU-runtime smoke as a regression signal after scheduler changes. Record event counts and reproducible conditions; investigate a clear regression. This is not a simulator-throughput benchmark campaign or a fixed percentage performance gate. + +## 5. Execution order + +| Step | Work | Main execution surface | Publication milestone | +| --- | --- | --- | --- | +| 0 | Create worktree, establish references and baseline | CPU / Git | New branch and tracked plan pushed. | +| 1 | Audit candidate changes and vLLM semantics | CPU / source | Reviewed scope and draft PR published. | +| 2 | Fix RR DP rotation | CPU | Source, tests, and results pushed. | +| 3 | Fix shared monolithic forward lifecycle | CPU | Full coherent runtime fix pushed. | +| 4 | Add bounded opt-in vLLM DP placement | CPU | Selection and event-integration evidence pushed. | +| 5 | Separate routing implementation identity | CPU | Config/data/model/cache chain pushed. | +| 6 | Repair legacy fused-MoE profiling | CPU + focused GPU worker | Code and CPU results pushed; native result published separately when complete. | +| 7 | Resolve optional zero-payload backend fix | CPU / companion repository | Included and remotely fetchable, or explicitly excluded. | +| 8 | Run combined regression, simplify the diff, hand off PR | CPU; native results reused only for unchanged code | Complete review-ready draft and final status pushed. | + +Steps 2–4 come first because they affect runtime execution. Do not block them on optional backend availability or GPU allocation. Do not call Step 6 complete until its native acceptance requirements are satisfied, or the user explicitly removes it from this PR. + +## 6. Step 0 — Create the worktree and establish a reproducible baseline + +### Actions + +Inspect the existing repository, remotes, worktrees, and local instructions. Confirm `origin` is the intended Frontier repository. Do not initialize candidate submodules merely to read source: the candidate contains an unresolved gitlink. + +The following is a starting sequence from an existing Frontier checkout. Resolve name collisions before running `git worktree add`. + +```bash +ROOT="$(git rev-parse --show-toplevel)" +git -C "$ROOT" remote -v +git -C "$ROOT" worktree list --porcelain +git -C "$ROOT" status --short + +git -C "$ROOT" fetch --no-recurse-submodules origin +BASE="$(git -C "$ROOT" rev-parse 'origin/main^{commit}')" +CANDIDATE="$(git -C "$ROOT" rev-parse 'origin/bug/ttft-check^{commit}')" +MERGE_BASE="$(git -C "$ROOT" merge-base "$BASE" "$CANDIDATE")" +BRANCH="fix/issue26-correctness-pr" +WORKTREE="$ROOT/.worktrees/issue26-correctness-pr" + +# Inspect any pre-existing branch or directory; never overwrite either. +git -C "$ROOT" worktree add -b "$BRANCH" "$WORKTREE" "$BASE" +cd "$WORKTREE" +mkdir -p docs/development/issue26-correctness-pr +``` + +Copy this document into `plan.md`; create the other three task records. Record full SHAs, repository URLs, branch/worktree paths, the chosen Python executable, and the first planned validation commands. + +Use the existing CPU environment when suitable. Keep GPU dependencies out of the minimal simulator environment. Avoid changing a shared editable installation to point at this worktree: use a task-specific environment or an explicit worktree `PYTHONPATH`. Record and verify the actual imported `frontier` path. + +Use Frontier's existing scratch-root support for test output. If `FRONTIER_TMP_ROOT` is used, point it to a new task-owned directory and record it. Do not reuse, delete, or overwrite old calibration caches. + +### Baseline validation + +Read `AGENTS.md`, applicable nested instructions, `pyproject.toml`, existing test configuration, and the relevant test helpers before selecting commands. Run a small existing CPU suite for the areas this PR will touch. Candidate baseline files include: + +```text +tests/unit/test_open_source_release_arch_guard.py +tests/unit/test_cluster_scheduler_dp_lanes.py +tests/unit/test_colocation_release_review_contracts.py +tests/unit/test_config_owned_contracts.py +tests/unit/test_stage_execution_time.py +tests/unit/test_stage_finalized_contract.py +``` + +Check their existence and dependency requirements at the fetched revision. Use an explicit file selection, not an unreviewed repository-wide command that launches native dependencies. Record any import, collection, or baseline failure with its actual exception. Do not modify production code to conceal a pre-existing environment problem. + +Run an existing small Frontier-only CPU smoke with synthetic/dummy execution times and the built-in analytical communication backend. Record the command and expected completion criteria. No trained historical profiles or GPU service are needed for this baseline. + +### Exit and publication + +The new worktree is based on the recorded main SHA; existing worktrees are unchanged; task docs are tracked; the baseline result and limitations are recorded. Commit and push the setup checkpoint. Verify that GitHub shows the branch and the full plan. A baseline failure must be explicit; it does not automatically prevent a source audit. + +## 7. Step 1 — Audit candidate changes and the reference behavior + +### 7.1 Three-way source review + +Use all three views below, then read complete functions and their call sites rather than relying on patch context alone: + +```bash +git diff --name-status "$MERGE_BASE" "$CANDIDATE" -- frontier tests docs +git diff "$MERGE_BASE" "$CANDIDATE" -- +git diff "$MERGE_BASE" "$BASE" -- +git diff "$BASE" "$CANDIDATE" -- +``` + +Read the donor task's `design_dp_load_balancing.md`, `design_shared_forward_sync.md`, `review_shared_forward_sync.md`, and relevant focused test reports. Consult `summary.md` or `handoff.md` only to resolve provenance or an implementation decision; do not restart the calibration investigation. The complete historical archive is not required reading for this PR. + +For each meaningful candidate hunk, record one disposition in `review.md`: + +| Disposition | Meaning | +| --- | --- | +| `PORT` | The behavior is justified and still missing from main. | +| `ADAPT` | Preserve the intent but implement it using main's current ownership or API. | +| `ALREADY_PRESENT` | Main already has an equivalent correction; add a missing regression only when useful. | +| `DROP` | Redundant, experimental, overcomplicated, unrelated, or incorrect. | +| `BLOCKED` | An evidence or user decision prevents inclusion. | + +Record the old defect, main behavior, reference behavior when applicable, chosen owner, planned test, and donor commit/path. Do not cherry-pick a large commit solely because its message says it is a fix. + +### 7.2 Required vLLM 0.10.2 source reading + +Create or use a separate read-only reference checkout pinned to the designated vLLM revision. Establish its relationship to vLLM 0.10.2 from version/release history and relevant source differences. The fork branch name or installed package version alone is not proof. When comparison to the upstream tag is necessary, pin that tag's resolved commit as well and record meaningful fork differences; do not substitute a newer vLLM release. + +| Reference source | Questions the audit must answer | +| --- | --- | +| `vllm/v1/engine/core_client.py` | How are DP engines ordered and scored? What does one frontend reserve locally? How are count updates applied? What changes with multiple frontends or explicit DP rank selection? | +| `vllm/v1/engine/coordinator.py` | When are changed counts published? What do the previous-step snapshot, minimum collection wait, unchanged heartbeat, and wave/step ordering mean? | +| `vllm/v1/engine/core.py` and scheduler stats producers | Which events cause a report, at what point relative to request state updates, and which reports are suppressed? | +| `vllm/v1/core/sched/scheduler.py` | What belongs to waiting versus running, including admitted-but-unscheduled work, preemption, and completion? | +| `vllm/v1/worker/gpu_model_runner.py`, `vllm/forward_context.py` | Why can DP source batches have different local phases while taking part in shared expert work? How are token populations and dummy participants represented? | +| `vllm/model_executor/models/qwen3_moe.py` | What is the model's gated expert computation and reduction path? Use this as a concrete case, not a production model-name condition. | +| `vllm/model_executor/layers/fused_moe/fused_moe.py`, `layer.py`, related activation/output helpers | What is the exact low-level arithmetic, routing-weight placement, workspace use, expert-map behavior, and local output reduction? | +| `vllm/distributed/device_communicators/all2all.py` and group helpers | Which reductions are local expert aggregation versus distributed communication? Avoid accounting for either twice. | + +Produce a concise reference-behavior table with pinned file/symbol references. Explain why Frontier represents the relevant behavior the way it does. No live vLLM server is needed. + +### 7.3 Main integration points to inspect + +Inspect request admission, `GlobalBatchEndEvent`, shared forward identity, waiting rooms, stage ownership, layer advancement, `Batch`/`Request` completion, `StageExecutionTime`, and metrics reporting before the scheduler fixes. + +Inspect `ReplicaConfig` creation/copying, routing-runtime resolution, model training signatures, model registries, family/precision selection, persistent cache loading, and current profiling backend selection before predictor/profiler edits. + +Prefer existing test utilities such as the repository's scratch-root and predictor-cache fixtures. Inspect any donor reference-comparison helper before adopting it. Do not create a second general testing harness. + +### Validation and publication + +Finalize the work-package dispositions and exact CPU test selections. Identify any decision checkpoint rather than inventing an answer. Create a **draft PR** after the branch and scope audit have been pushed, using the existing repository template when present. + +Suggested title: `Fix DP request placement, shared forwards, and MoE profiling contracts`. + +Use `Related to #26`, not `Fixes #26`. State that latency calibration and vLLM E2E comparison are out of scope. Keep the draft open and update it at subsequent checkpoints. Do not request merging before the user has reviewed it. + +## 8. Step 2 — Preserve RR DP rotation across scheduling calls + +### Source scope + +Primary implementation: `frontier/scheduler/cluster_scheduler/round_robin_cluster_scheduler.py::_schedule_batch_mode`. Start from main's implementation and the donor's small correction. Existing regression starting point: `tests/unit/test_cluster_scheduler_dp_lanes.py`. + +### Required behavior + +Request assignment must not depend on how an identical ordered request stream is divided across calls to `schedule()`. + +The donor's intended order is: + +```python +ordinal = completed_assignments + request_index +replica_index = ordinal % num_replicas +dp_lane = (ordinal // num_replicas) % dp_size +``` + +Reuse the existing persistent counter if its meaning matches this calculation. Do not create a second counter or cache without a demonstrated need. Preserve the existing order across replicas and the return structure. Use actual replica IDs, not an assumption that IDs equal list indices. + +Keep dedicated decode/PD-AF allocation logic unchanged unless the same defect is independently demonstrated there. This is an RR correction, not vLLM load-balancer emulation. + +### Required CPU tests + +Test the production scheduling path with one request at a time, a single burst, and different batch boundaries for the same sequence. Cover DP1, more than two DP lanes, multiple replicas, non-contiguous replica IDs, an empty scheduling call, and continuation after that empty call. A small parameterized table is sufficient; no arbitrary scenario count is required. + +A minimal regression must fail on the old implementation for the intended reason: repeatedly starting lane assignment from zero. Compare placement by request identity and preserve order within each destination lane, with every request assigned exactly once. Do not require the flattened return-list order of one burst to equal that of many calls: the existing implementation may group each call's results by replica. + +Run the existing DP-lane and RR tests plus the small CPU smoke. Compare unaffected DP1 behavior against the recorded baseline. + +### Exit and publication + +The same request sequence has the same assignment regardless of call partitioning; the new test is demonstrably sensitive to the old defect; unrelated allocation behavior is unchanged. Update docs, commit source and tests, push, and report the result and next step. + +## 9. Step 3 — Repair shared monolithic forward completion + +### Source scope + +Review these donor paths together, then adapt them to main: + +```text +frontier/scheduler/cluster_scheduler/base_cluster_scheduler.py +frontier/scheduler/utils/forward_sync_state.py +frontier/scheduler/utils/sync_state.py +frontier/scheduler/utils/forward_collective.py # donor addition +frontier/scheduler/utils/ep_wave_inputs.py +frontier/scheduler/utils/ep_wave_schedule.py +frontier/scheduler/utils/prefill_collective.py +frontier/scheduler/utils/decode_collective.py +frontier/events/replica_stage_schedule_event.py +``` + +Also read the current sync-entry and stage-ownership helpers, `PrefillSyncEvent`, `DecodeSyncEvent`, `BatchStageEndEvent`, and request completion consumers. The listed paths are an audit map, not permission to overwrite each file. + +### Invariants to implement + +**One shared forward identity.** In a monolithic model requiring shared expert execution, participating source lanes join the same forward step regardless of local prefill/decode classification. Do not add a parallel hierarchy of event classes. Preserve existing separation for disaggregated execution where it has real meaning. + +**One source owner per request.** Preserve original source batches and their request/token vectors. A real request cannot appear in two source lanes of the same shared step. Validate ownership at the appropriate group-formation boundary rather than repeatedly scanning unchanged membership at every operator. + +**One completion of the shared expert work.** The completed shared work must be consumed once, and full-stage ownership must be restored once for the participating group. Understand what the restore helper's return value means. Do not conflate “not attempted” with “not applicable” through fallback calls that restore owners twice. + +**Source-local continuation.** Each source uses its own next-attention inputs, context lengths, stage tail, and request phase. Do not choose a single sample batch and apply its predicted duration to all source lanes. An idle participant may be required for synchronization but must not become a completed user request. + +**Correct layer advancement.** In a mixed source, advance the decode subset exactly once per completed layer. Pure-decode sources follow their existing completion path; prefill requests must not gain a duplicate decode-layer increment. Respect pipeline stage bounds and supported speculative-decoding state. + +**Single accounting owner.** Reuse main's current execution-time and reporting ownership. Shared waiting contributes to elapsed stage time; it must not also be invented as CPU work or charged twice as an operator. Include shared expert work in the appropriate source model-time accounting without fabricating a second metrics system. Preserve demand-driven reporting when it is disabled. + +Do not change same-layer admission timing for dense layers within a MoE model merely because it can be refactored at the same time. Fix source-local continuation and bookkeeping without expanding the scheduling policy. + +### Required CPU tests + +Use the donor's `test_monolithic_mixed_forward_sync.py` as a starting point, not the full acceptance suite. Add behavior tests covering: + +| Case | Required assertion | +| --- | --- | +| Prefill/prefill, decode/decode, prefill/decode, and mixed-source combinations | Each admitted group reaches one shared completion without deadlock. | +| Reversed source arrival order | The same shared identity and correct outcome; no dependence on which phase arrives first. | +| Unequal source tokens/context lengths | Each lane queries/uses its own continuation; a deliberate unequal-duration fixture detects borrowed timing. | +| A real source plus idle participants | Idle participation does not create fake requests, progress, or terminal callbacks. | +| Duplicate request ownership | A clear failure at the owning boundary, before producing inconsistent state. | +| Multiple layers and successive forward steps | Owners, open-step bindings, waiting rooms, and pending events are consumed or released correctly. | +| Mixed-batch decode requests | One layer increment and one final token credit where appropriate; prior TTFT remains unchanged. | +| Supported dense-layer transitions within a MoE model | No invented EP collective for dense work; local continuation remains correct. | +| Reporting on versus off | Same simulated execution outcome; reporting does not introduce extra mandatory runtime work. | + +At least one integration test must run the **real Frontier event loop and real admission/ownership/completion code** through a small multi-request case with prefill overlapping decode. Inject deterministic execution times only at the predictor/backend boundary. Do not replace the synchronization, ownership restoration, or terminal callbacks being tested. + +Assert exact request/token conservation, unique completion, no stranded source batch, no live ownership record after completion, and no unfinished request when the event queue drains. A nonempty metrics file is not sufficient. + +Run relevant existing tests for forward identity, EP-wave materialization, stage ownership, stage execution reporting, sequential PDD/PD-AF, and supported prefix/speculative paths affected by the call chain. Resolve concrete test file names in Step 1 and record the executed list. No GPU is needed. + +### Exit and publication + +Publish the shared lifecycle fix as a coherent unit, not an intermediate state where waiting rooms are shared but completion ownership is still phase-specific. Tests must distinguish the old mixed-phase defect and detect a deliberately wrong source-timing implementation. Record any unsupported path explicitly. Commit, push, and provide the runtime-invariant evidence. + +## 10. Step 4 — Add bounded, opt-in vLLM-style DP placement + +### Source scope + +Candidate components to review: + +```text +frontier/config/cluster_scheduler_config.py # donor addition +frontier/config/config.py +frontier/types/cluster_scheduler_type.py +frontier/scheduler/request_load.py # donor addition +frontier/scheduler/utils/vllm_dp_load_balancer.py # donor addition +frontier/scheduler/cluster_scheduler/vllm_load_balancing_cluster_scheduler.py +frontier/scheduler/cluster_scheduler/cluster_scheduler_registry.py +frontier/scheduler/cluster_scheduler/base_cluster_scheduler.py +frontier/scheduler/replica_scheduler/base_replica_scheduler.py +frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py +frontier/events/cluster_schedule_event.py +frontier/events/global_batch_end_event.py +``` + +The config-family extraction is optional unless needed to keep the change cohesive. Preserve existing imports, config discovery, CLI flattening, and serialization. Do not add a broad config migration merely to register a new strategy. + +### Required behavior and boundaries + +Preserve the reference's separation between engine counts, the last publishable snapshot, and frontend estimates. For the intended one-frontend mode, select the first engine with minimum `4 * waiting + running`, and reserve one local waiting request. A fresh published snapshot replaces the frontend estimates; it is not an incremental adjustment to their locally reserved counts. [R3] + +Model the changed-count publication, previous-step snapshot, minimum collection wait, and unchanged heartbeat according to the pinned reference. The expected source constants are 4, 50 ms, 100 ms, and 5000 ms; confirm them in the actual source before using them. Explain the initialization epoch and equal-time ordering in code comments or the reference table, not as hidden timing adjustments. + +Use one small state owner for the count-publication/selection behavior and the existing cluster scheduler for DES integration. Pass simulation time explicitly through a clean interface; avoid a mutable “last routing time” bridge if the existing scheduling API can carry time directly without a broad change. If an adapter is necessary for older policies, keep it thin and document its contract. + +Report load at the correct post-step point, after the lane's real request-state transition. A load accessor must count admitted-but-unscheduled running requests and preempted waiting work correctly. Reuse that accessor in existing diagnostics rather than maintain two definitions. + +Lazy timer advancement is acceptable when it preserves all observable publication/selection behavior. Do not create perpetual heartbeat events that keep a drained simulation alive. Verify suppressed reports and same-time routing/report events, not just standalone timeout arithmetic. + +Initial support remains one serving Replica, one modeled frontend, co-location/MONOLITHIC, V1, and PP1. Keep the strategy opt-in and retain all existing defaults. Do not claim support for multi-frontend delivery, IPC timing, elastic scaling, or exact warm-start publication phase. + +The donor uses `ForwardSyncState.get_step_id(batch)` as a report-order key. **Do not assume this is valid for every configuration accepted by the new strategy.** Check that its order/equality properties match the reference's relevant wave/step semantics for the admitted paths. Check dense models, DP1, idle participants, successive request waves, and real batch-end hooks. Provide a correct existing identity, or propose a clearly bounded capability restriction to the user; do not manufacture a guessed step number or silently use a request ID. + +### Required CPU tests + +Start from `test_vllm_dp_load_balancer.py`, then add the missing integration coverage: + +| Layer | Cases and assertions | +| --- | --- | +| Pure count selection | Weighted counts, deterministic ties, local reservations, replacement by a new snapshot, empty initial counts. | +| Publication state | Same-step reports, a newer step while changes are pending, unchanged/suppressed reports, heartbeat replacement, delayed frontend observation. | +| Ordering | Report at a deadline, select at a deadline, both at the same timestamp, multiple events within the same millisecond, monotonic time validation. | +| Request populations | Waiting admission, running but not scheduled, preemption, completion, and the resulting reports from actual lane state. | +| Real DES integration | `ClusterScheduleEvent` supplies time; `GlobalBatchEndEvent` observes post-step state; the strategy does not keep an otherwise finished run alive. | +| Capability and defaults | Allowed topology works; unsupported topology fails before simulation; RR remains the default where it was previously selected. | +| Configuration | Enum/registry discovery, old imports, generated CLI selection, config copying, and round-trip representation. | + +Use a compact reference-driven trace test with explicit times and expected publications/selections. Prefer exercising relevant pinned reference methods through a minimal test-only fake clock/poller when practical. Otherwise derive the expected trace independently from the source and document the mapping. Inspect the existing donor `tests/integration/issue26_dp_coordinator_reference.py` before deciding whether it is reusable. + +Do not simply duplicate the new balancer into the expected-result generator. Source-string assertions and tests that assign final state directly cannot replace behavioral tests. + +### Exit and publication + +The selection and publication rules are source-backed; actual Frontier events exercise them; supported scope is truthful; current defaults and unrelated strategies are unchanged. Publish the implementation and tests with an explicit statement that source-level behavior is validated but no vLLM serving-placement or timing equivalence is claimed. + +## 11. Step 5 — Separate routing load distribution from implementation identity + +### Source scope + +```text +frontier/config/config.py +frontier/moe_routing_runtime.py +frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py +frontier/execution_time_predictor/shared_prediction_model_manager.py +``` + +Inspect current model-cache helpers and any query caches reached by these paths. Existing donor tests include `test_moe_routing_runtime.py` and `test_moe_routing_runtime_model_sharing.py`; reuse main's predictor-cache fixtures where applicable. + +### Required behavior + +The configured expert-load distribution and the routing implementation used for timing prediction are separate dimensions. Retain existing default resolution, but support the explicit validated runtime override using Frontier's normal config mechanisms. + +Follow the value through: + +```text +Configuration creation / copying + -> resolved routing runtime + -> dataset row selection and validation + -> training identity + -> trained-model registry and precision/measurement-family selection + -> persistent cache load + -> runtime model lookup / query cache, where relevant +``` + +Do not implement only the CLI flag. Two routing implementations with the same layer shape must not select each other's routing-cost model. Conversely, equivalent resolved configurations should retain sharing; do not separate models merely because one configuration used an explicit default and another left it implicit. + +Only include runtime identity where it changes the measured semantics. Avoid unnecessarily retraining or duplicating unrelated expert-compute models when their inputs and contracts are identical. Keep current layer, precision, and measurement-family identity intact, including main's device-event support. + +Resolve required configuration once at its owner. Avoid nested `getattr` fallbacks on every query or reattaching mutable identity to a cached estimator without provenance. An artifact with unknown or conflicting routing identity must not be silently relabeled as the caller's requested runtime. Any recovery of missing metadata needs proof from the existing cache key and source dataset; otherwise reject it clearly or use an approved migration. + +### Required CPU tests + +Use tiny synthetic rows and existing training/cache paths, not historical profiling data or a large fitting run. + +Verify defaults and explicit overrides; invalid runtime rejection; config reconstruction/copying; dataset filtering; rejection of incompatible or ambiguous rows; two runtimes with otherwise identical features; equivalent explicit/implicit configurations; precision/family separation; cache miss and cache hit; serialized reload in a fresh manager instance; and rejection of conflicting cached metadata. + +Include at least one test that goes through actual minimal training or the existing trainer's real storage path and then a cache reload. Registry-only fake-estimator tests do not demonstrate persistent-cache isolation. Test unchanged ordinary model sharing as well as the newly separated routing models. + +### Exit and publication + +The override reaches the actual model selection; mismatched runtimes cannot collide; existing compatible sharing remains intact; fresh and cached paths agree. Publish source, tests, and the identity rules. Do not claim that timing prediction has become more accurate without a separate calibration task. + +## 12. Step 6 — Repair legacy fused-MoE profiling without regressing main + +### Source scope and implementation + +Primary source: `frontier/profiling/moe/moe_vllm_kernel.py`. Read its callers, measurement export, training targets, and current backend selection before editing. + +The legacy path must perform the actual supported gated expert computation: + +```text +First expert matrix multiplication + -> gated SiLU activation + -> activation quantization, when the selected supported path needs it + -> second expert matrix multiplication with the correct routing weights + -> local reduction of the top-k expert outputs +``` + +The donor adds `silu_and_mul`, local `moe_sum`, an activation buffer, and shared output workspace. Port the justified arithmetic and layout into main's current legacy entry point. Do not replace the file or revert functional fused-experts, ROCm/MXFP4 handling, profile-method validation, or shared timer statistics. Main already tests population statistics and finite single-sample behavior; preserve them. [R7, R8] + +Local top-k output reduction is not a distributed collective. Do not add or remove DP/TP/EP communication costs as part of this arithmetic repair. + +Use explicit existing backend selection. Avoid broad import/exception fallback paths that silently switch the implementation. Keep optional native dependencies out of CPU simulator imports. Do not change a public profiling return contract just to expose an intermediate tensor to a test when the existing low-level helper can be tested directly. + +### Measurement ownership is part of the fix + +Specify exactly which operations belong to the corrected measured target. In particular, identify whether local output reduction is contained in the grouped-expert measurement or separately represented. The full local computation must be counted exactly once in training and runtime queries. + +Do not silently keep using old rows or cached models as if their measured scope had changed. Prefer the repository's existing artifact/measurement identity mechanism for a narrowly scoped distinction. Preserve unaffected measurement families and data. Do not delete shared caches or require regeneration of all profiling data. + +If existing metadata cannot safely distinguish the old incomplete target from the corrected target, stop that part for a user decision. Present the smallest metadata/migration change and the alternative of deferring the profiling package. “The old cache probably will not be used” is not an acceptable compatibility rule. + +Do not generalize from the Qwen BF16 case to arbitrary activation or quantization. Follow the actual model contract. If an advertised path is unsupported, expose that fact and discuss the compatibility impact before narrowing it; do not silently compute gated SiLU for a non-gated model. + +### CPU validation + +Extend the existing test layout rather than duplicating the native test framework. Starting points include: + +```text +tests/unit/test_moe_fused_event_contract.py +tests/unit/test_moe_native_admission.py +tests/unit/test_moe_mxfp4_increment10.py +tests/unit/test_device_timer_contract.py +tests/unit/test_timer_owner_lifecycle.py +``` + +Some profiling-boundary tests require CPU-importable Torch/native-related Python packages even though no GPU kernel runs. Inspect their import requirements and use the existing suitable test environment. Do not make those optional packages new mandatory dependencies of the simulator. + +Required checks include the supported call order, activation/output dimensions, routing-weight placement, expert-map propagation, workspace alias lifetime, measurement ownership, old-target/cache admission, and unchanged functional-backend dispatch. Add a small explicit CPU arithmetic example that distinguishes a gated activation from merely slicing the first matrix output. Mark mock-based dispatch tests as boundary validation, not native numerical parity. + +Maintain existing event-method/platform checks and statistics tests. Do not restore the donor's older timing/statistics implementation while moving its arithmetic fix. + +### Native numerical validation: narrow but required + +Run a **local-expert numerical test on a GPU worker**, not a serving or timing benchmark. The CPU master must not execute CUDA work. Reuse the approved job mechanism and existing suitable image/environment; record actual source/import paths, binary versions, device, command, and results. Do not invent a new image, download model weights, allocate an eight-GPU server, or build a full serving environment unless the minimal test genuinely requires it and the user authorizes that expansion. + +Use identical generated inputs, weights, routing weights/IDs, expert map, dtype, and kernel configuration for the repaired Frontier helper and the pinned vLLM reference. Compare real outputs, including local output reduction. A single compatible GPU can ordinarily test different local EP partitions sequentially; the logical EP size is not a requirement for a distributed serving run. + +Minimum BF16 acceptance: + +| Case | Purpose | +| --- | --- | +| Existing Qwen-shaped 4096 and 4097 token cases, using the checked-in model config and local expert-map layout | Preserve the donor test's important production-shaped coverage. | +| More than one local expert partition | Detect incorrect global/local expert handling without launching an EP cluster. | +| A smaller boundary case with a different valid top-k and uneven expert occupancy | Detect shape-specific correctness and missing-local-expert output errors. | + +The donor test is `tests/unit/test_moe_fused_expert_numerical_parity.py`. Adapt its useful checks into the repository's appropriate native-test location and capability gating. Do not import the entire GPU stack into the ordinary CPU unit suite merely to preserve its old filename. + +For identical BF16 kernels and operation order, retain the donor's zero-tolerance comparison unless an identified supported implementation difference justifies a documented tolerance. Do not loosen tolerance merely to make a failure disappear. Check finite outputs and repeated invocation with different inputs so workspace reuse cannot leak stale results. + +Also validate any precision/backend path whose advertised execution is changed by this patch. A CPU shape test does not prove FP8 arithmetic. Keep the GPU matrix limited to affected paths, but do not claim an untested changed native path is proven correct. Resolve an unavailable required path with the user rather than launching an open-ended GPU campaign. + +### Exit and publication + +Push the code, CPU tests, and documentation before waiting on GPU availability; mark native validation `NOT_RUN` or `BLOCKED` explicitly. Once native checks finish, publish their concise results and any required fixes. Step 6 is `PASS` only when arithmetic, measurement ownership, artifact compatibility, and the required native checks are all satisfied. + +The unrelated runtime fixes may continue while this step is blocked. The draft PR must not be presented as ready to merge with a required numerical check silently skipped. + +## 13. Step 7 — Resolve the conditional zero-payload backend change + +### Decision boundary + +Inspect `.gitmodules` and the actual gitlink in the new branch. The configured optional backend repository is `fwyc0573/frontier-htsim`. At specification time, the candidate's `e564935d...` commit could not be resolved through the configured remote API. Recheck repository access and commit reachability; an access problem and an unpublished commit are different findings. Do not repeatedly try guessed repositories. [R9] + +If the fix is available, inspect its diff against main's pinned backend and ensure the parent gitlink does not pull unrelated changes. If it exists only locally, verify that exact local source and tests before publishing it. If necessary, reconstruct the small input-handling fix from source and failing tests; do not treat a report as a substitute for code review. + +Keep the backend package conditional. Its absence must not block RR, shared-forward, routing, or other analytical-backend CPU tests. + +### Required behavior + +Explicit zero is valid input where the existing backend accepts an empty transfer. Missing required payload is still an error. Negative payload is rejected. An explicit CLI zero overrides a positive JSON value. Positive-payload behavior remains unchanged. + +Do not equate zero payload with zero total collective time. Preserve the backend's existing synchronization/latency semantics for the selected collective. Do not fit or change its timing parameters. + +### Required CPU validation and publication + +Start from donor `tests/unit/test_collective_sim_zero_payload.py`. Test the real scenario serialization, CLI runner, and built CPU simulation executable with zero, missing, negative, and positive input, plus CLI-versus-JSON precedence. A mocked subprocess result is insufficient. + +If a new companion-repository change must be created or published, obtain the user's decision on that additional repository scope first; continue independent Frontier work while it is pending. Once authorized, create a dedicated backend branch, commit tests and source, push to the correct permitted remote, and open a companion draft PR when appropriate. Do not push to its main branch. Publish the backend commit **before** updating Frontier's gitlink. + +Validate the resulting Frontier checkout with a clean temporary checkout/clone and `git submodule update --init` for this backend. This must fetch the exact intended SHA without relying on another worktree's local Git object store. Build and execute the CPU runner tests in the recorded environment. + +Then commit and push Frontier's gitlink, parent-side tests, and docs. Report both repositories and their commits/dependency status. + +If the package cannot safely be included, leave Frontier on its current-main backend revision and record `EXCLUDED` with the reason. Do not vendor the backend, alter `.gitmodules` to an unapproved destination, or publish an unreachable gitlink. Ask the user when companion-repository scope or access needs a decision. + +## 14. Step 8 — Combined regression and final PR review + +### 14.1 Validate the integrated branch + +Run the selected focused suites together, not only each package in isolation. Use fresh task-owned outputs and caches where the test concerns first-load behavior. Also exercise cache-hit paths intentionally rather than accidentally inheriting a donor cache. + +The combined CPU selection must cover: + +| Area | Required result | +| --- | --- | +| RR and opt-in DP placement | Persistent assignment, reference count semantics, real event hooks, and unchanged defaults. | +| Shared forwards | Correct mixed-source progress, source-local timing, ownership release, no dangling waiting rooms, and request/token conservation. | +| Existing architectures | Supported co-location, sequential PDD, and sequential PD-AF CPU regressions remain valid. Preserve their existing unsupported-mode guards. | +| Stage reporting | Current `StageExecutionTime`/metrics behavior survives; enabling reporting does not change simulated completion. | +| Routing models | End-to-end config-to-query identity, compatible sharing, and persistent-cache isolation. | +| Profiling boundaries | Correct legacy arithmetic structure and measurement ownership; main's functional/accelerator/timer dispatch tests remain intact. | +| Optional backend, when included | A fresh remotely fetchable submodule checkout passes the actual CPU runner tests. | +| Public interfaces | CLI/config registration, imports, supported backend selection, and relevant docs are consistent. | + +Run small deterministic Frontier-only simulations with real event processing: an RR arrival stream, an opt-in DP-placement stream, and a stream that creates overlapping prefill/decode source batches. Reuse existing fixtures. Their purpose is lifecycle validation, not matching historical H200 latency. Include heterogeneous/dense-layer or pipeline cases where the changed shared code applies, even though the new DP strategy itself remains PP1-only. + +Before handoff, fetch main again. If it advanced, inspect overlap and assess whether integration validation is needed. Prefer additive integration into the published branch when necessary; do not rebase or force-push away the user's review anchors. Rerun the affected tests after any integration change. + +Do not run native profiling suites on the CPU master. Do not launch vLLM E2E comparisons as a final confidence check; they are not part of this task. + +### 14.2 Review the complete diff against the actual PR base + +Perform a final line-by-line review of changed code and meaningful surrounding context. Record the reviewed revision and actual review method; do not describe a self-review as independent review. Use these questions: + +- Is every change linked to a demonstrated defect, an essential regression test, or a directly related simplification? +- Are state ownership and initialization explicit? Are repeated fallback checks or parallel state representations still present without a reason? +- Do source batches keep their own shape and progress? Are terminal events and ownership transitions unique? +- Are config, layer, routing-runtime, precision, and measurement identities preserved through both fresh and cached paths? +- Is any operation omitted, counted twice, or relabeled without compatible metadata? +- Did the patch preserve current-main model/backend support and demand-driven reporting? +- Are tests checking production behavior rather than copying the implementation or replacing the behavior under test with a stub? +- Are any workstation paths, credentials, datasets, weights, generated traces, caches, or unreachable submodule references staged? +- Can a reader understand the names and functions without the historical calibration conversation? + +Remove candidate scaffolding that no longer has a purpose. Do not add speculative abstractions in the cleanup pass. Re-run the affected tests after cleanup. + +A broader CPU test selection may expose pre-existing failures. Reproduce them on the recorded base with the same environment where practical; report them as baseline failures with evidence. Do not relabel an unexplained failure as pre-existing or expand into unrelated repairs without a decision. + +### 14.3 Draft PR contents + +Update the existing draft PR rather than creating duplicates. Its description should contain the problem statement, included fixes and deliberate exclusions, important design choices, test commands/results, native-test status, compatibility/data-scope implications, source reference revision, and any companion PR dependency. + +Link the tracked `plan.md`, `progress.md`, `review.md`, and `validation.md`, plus the relevant implementation commits. Explain that no vLLM serving/TTFT comparison was performed or required. Do not paste large logs into the PR body. + +Keep the PR draft until the user reviews it and decides the next GitHub action. Report whether technical acceptance is complete independently of GitHub's draft status. Do not merge, squash, delete branches, or close Issue 26. + +## 15. Validation records and completion rules + +### 15.1 Keep four independent statuses + +Do not collapse code status, test status, publication status, and user approval into one `PASS` label. + +```text +Work package: NOT_STARTED | IN_PROGRESS | PASS | FAIL | BLOCKED | EXCLUDED +Test result: PASS | FAIL | SKIPPED | NOT_RUN +Publication: LOCAL_ONLY | PUSHED_VERIFIED +User review: NOT_REVIEWED | CHANGES_REQUESTED | APPROVED +``` + +`EXCLUDED` must state whether exclusion was permitted by this specification or explicitly approved by the user. `SKIPPED` is never evidence that a native check passed. A job's platform status is not a substitute for test output and exit status. + +### 15.2 Minimal per-checkpoint evidence + +Use concise tables in `validation.md`: + +| Field | Record | +| --- | --- | +| Step and purpose | Which correctness claim the test supports. | +| Source under test | Base/code revision and the exact candidate diff tested, or the committed implementation revision. | +| Environment | Python executable/version; relevant package versions; working directory and import path; native details only when relevant. | +| Command | Executable command with the selected test files and options. | +| Outcome | Pass/fail/skip counts, process exit code, and important assertions or numerical result. | +| Baseline comparison | Whether the defect test fails before the fix and how unrelated baseline failures were handled. | +| Limits | What was mocked, not run, unsupported, or not established by the result. | + +Record tested code accurately. If testing precedes the checkpoint commit, keep the production/test diff unchanged between that run and commit; describe it as the tested checkpoint diff rather than claim the unmodified parent SHA was tested. Rerun affected checks after additional code edits. The user-facing update supplies the resulting commit SHA. + +Keep enough small evidence in the tracked documents to review results in GitHub without access to local logs. Large raw output remains in the task-owned scratch location and is referenced as supplementary material, not the sole proof of completion. Redact secrets and avoid committing private environment dumps. + +### 15.3 Final acceptance checklist + +- [ ] A new worktree and branch were created from a freshly fetched and recorded main revision. +- [ ] The candidate/main/source audit records which changes were ported, adapted, already present, dropped, or blocked. +- [ ] Relevant vLLM 0.10.2 behavior is explained from pinned source, not assumed from branch names or old reports. +- [ ] RR DP assignment is independent of scheduling-call partitioning. +- [ ] Real shared-forward event processing preserves ownership, source-local timing, progress, and unique completion. +- [ ] The opt-in DP strategy is correctly integrated, explicitly bounded, and does not change defaults. +- [ ] Routing-runtime identity is correct through config, data, training, registries, persistent caches, and lookup. +- [ ] Any included legacy MoE change has correct arithmetic, single measurement ownership, compatible artifact handling, and required native numerical evidence. +- [ ] Any included backend gitlink is reachable from the configured remote and tested from a fresh checkout; otherwise the conditional package is explicitly excluded. +- [ ] Relevant current-main architecture, reporting, model, and accelerator regressions are preserved. +- [ ] Required tests are actually executed, with skipped/unrun work visible rather than converted to a pass. +- [ ] The final diff contains no calibration constants, unsupported accuracy claim, task archive, generated cache, or unrelated rewrite. +- [ ] Source, tests, and task docs are committed and pushed, and remote HEAD verification succeeds. +- [ ] The draft PR links all review material and distinguishes technical completion from user approval. + +No hidden GPU or vLLM E2E requirement may appear at the end. The only native work required here is the explicitly described numerical validation for included changed native execution paths. + +## 16. Publication commands and interruption recovery + +### 16.1 Normal checkpoint publication + +Use explicit paths when staging. Do not use an unreviewed `git add -A` in a worktree containing experiment outputs. + +```bash +git diff --check +git status --short +# Replace the placeholders with the reviewed source, tests, and docs paths. +git add -- docs/development/issue26-correctness-pr/ +git diff --cached --stat +git diff --cached --check +git commit -m "" + +BRANCH="$(git branch --show-current)" +git push --set-upstream origin "$BRANCH" +LOCAL_HEAD="$(git rev-parse HEAD)" +REMOTE_HEAD="$(git ls-remote origin "refs/heads/$BRANCH" | cut -f1)" +test "$LOCAL_HEAD" = "$REMOTE_HEAD" +git status --short +``` + +The sequence is a template, not a substitute for reviewing the staged diff. Inspect the code repository's GitHub commit/files view after pushing. For companion-repository changes, repeat the corresponding checks in that repository and identify both SHAs in the report. + +### 16.2 Resume after a terminal failure or a new session + +First read `progress.md`, then the relevant part of `plan.md`, `review.md`, and `validation.md`. Confirm the current worktree, branch, recorded base, local status, recent commits, upstream branch, and remote HEAD. Inspect pending changes before deciding whether a step completed. + +Treat code edits, tests, commits, and pushes as distinct stages. A completed test with no commit requires reviewing and committing the tested diff; a local commit with a failed push requires publication, not reimplementation. A test result from different code or a changed native environment cannot be reused without checking the difference. + +Do not delete partial work, submit duplicate GPU jobs, recreate the branch, or rerun a large task solely because the terminal disconnected. Inspect the recorded native job and existing artifacts first. Keep the next action specific, for example: “Run the real mixed-source lifecycle test after the ownership change, then commit and push Step 3.” + +Before stopping, leave the exact current state and next command/action in `progress.md`, publish the checkpoint when possible, and send the required user-facing update. Do not promise unattended work after the session ends. + +## 17. Source index + +These are reading anchors for the execution agent. They establish the reviewed snapshot and candidate intent; they are not test results for the future PR branch. Use pinned file/symbol links in the new task's review records after rechecking references. + +**[R1] Current-main repository and contributor guidance** + +- `https://github.com/NetX-lab/Frontier/tree/1f694f7c549aa3aeeb7c5bbae04e119c09167a77` +- `https://github.com/NetX-lab/Frontier/blob/1f694f7c549aa3aeeb7c5bbae04e119c09167a77/AGENTS.md` +- `https://github.com/NetX-lab/Frontier/blob/1f694f7c549aa3aeeb7c5bbae04e119c09167a77/pyproject.toml` + +**[R2] Candidate branch and task archive** + +- `https://github.com/NetX-lab/Frontier/tree/a7b3320fe9b8b083ee86b91dae3d6838f4443d91` +- `https://github.com/NetX-lab/Frontier/tree/a7b3320fe9b8b083ee86b91dae3d6838f4443d91/task_memory/task_2026-09-07_issue26_ttft_h200` + +**[R3] User-designated vLLM reference** + +- Branch: `https://github.com/fwyc0573/vLLM-BS/tree/feature/frontier-comparison-instrumentation` +- Reviewed source: `https://github.com/fwyc0573/vLLM-BS/tree/ea95f571e20937c7c908c6d59ddd1cd6bf9268f1` +- Read the files listed in Section 7.2 at this pinned revision. Record the verified 0.10.2 relationship and any material fork changes. + +**[R4] Candidate design and recovery documents** + +Under the pinned task archive in [R2]: `design_dp_load_balancing.md`, `design_shared_forward_sync.md`, `review_shared_forward_sync.md`, `summary.md`, and `handoff.md`. Prefer the focused design documents over the chronological experiment history. + +**[R5] Ignore policy affecting reviewable records** + +- `https://github.com/NetX-lab/Frontier/blob/1f694f7c549aa3aeeb7c5bbae04e119c09167a77/.gitignore` + +**[R6] Candidate shared-forward and DP tests** + +Under `tests/unit/` at [R2]: `test_cluster_scheduler_dp_lanes.py`, `test_monolithic_mixed_forward_sync.py`, and `test_vllm_dp_load_balancer.py`. The donor's `tests/integration/issue26_dp_coordinator_reference.py` is a possible reference-test helper, subject to review. + +**[R7] Candidate MoE and routing tests** + +Under `tests/unit/` at [R2]: `test_moe_fused_expert_numerical_parity.py`, `test_moe_routing_runtime.py`, and `test_moe_routing_runtime_model_sharing.py`. + +**[R8] Main profiling and reporting integration anchors** + +Under [R1]: `frontier/profiling/moe/moe_vllm_kernel.py`, `frontier/entities/stage_execution_time.py`, `frontier/scheduler/utils/prefill_collective.py`, `frontier/scheduler/utils/decode_collective.py`, and `tests/unit/test_moe_fused_event_contract.py`. + +**[R9] Optional backend and historical zero-payload evidence** + +- Parent configuration: `https://github.com/NetX-lab/Frontier/blob/1f694f7c549aa3aeeb7c5bbae04e119c09167a77/.gitmodules` +- Configured repository: `https://github.com/fwyc0573/frontier-htsim` +- Candidate report: `task_memory/task_2026-09-07_issue26_ttft_h200/test_report_2026-09-08_collective_zero_payload.md` at [R2]. +- Candidate parent regression: `tests/unit/test_collective_sim_zero_payload.py` at [R2]. + +--- + +**First action for a fresh Claude Code session:** perform Step 0, publish the new worktree branch and tracked plan, then complete the source audit in Step 1. Do not begin by cherry-picking the calibration branch or launching a GPU job. diff --git a/task_memory/task_2026-09-21_issue26_correctness_pr/progress.md b/task_memory/task_2026-09-21_issue26_correctness_pr/progress.md new file mode 100644 index 00000000..d28f42fd --- /dev/null +++ b/task_memory/task_2026-09-21_issue26_correctness_pr/progress.md @@ -0,0 +1,37 @@ +# Issue 26 Correctness PR — Progress + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Step 0 started: records landed, environment created, baseline pending. | + +## Status + +| Field | Value | +| --- | --- | +| Correctness branch | `fix/issue26-correctness-pr` (to be created from `refactor/oversized-module-split` after the refactor branch's Step 0 push) | +| Base at creation | pending | +| Prerequisite | Refactor PR fidelity matrix PASS (see `task_memory/task_2026-09-21_oversized_module_split/progress.md`) | +| Current step | Step 0 (records landed; branch creation pending) | +| Publication | LOCAL_ONLY | +| Next action | Create the correctness worktree from the refactor branch, push, and stop for user review (Q6=a). | + +## Step status + +| Step | Work package | Status | Test | Publication | User review | +| --- | --- | --- | --- | --- | --- | +| 0 | Worktree, references, baseline | IN_PROGRESS | NOT_RUN | LOCAL_ONLY | NOT_REVIEWED | +| 1 | Candidate/vLLM audit | NOT_STARTED | — | — | — | +| 2 | RR DP rotation | NOT_STARTED | — | — | — | +| 3 | Shared monolithic forward | NOT_STARTED | — | — | — | +| 4 | Opt-in vLLM DP placement | NOT_STARTED | — | — | — | +| 5 | Routing implementation identity | NOT_STARTED | — | — | — | +| 6 | Legacy fused-MoE profiling | NOT_STARTED | — | — | — | +| 7 | Optional zero-payload backend | NOT_STARTED (facts in `plan.md` A7) | — | — | — | +| 8 | Combined regression, PR hand-off | NOT_STARTED | — | — | — | + +## Chronological updates + +- 2026-09-21: Baseline on the shared base recorded in the refactor task's Step 0 report; vLLM reference cloned (no tags in the fork; upstream `v0.10.2` comparison deferred to Step 1). +- 2026-09-21: Draft specification analyzed; eleven facts verified against main, the candidate, the submodule remote, and the host; planning interview settled twelve decisions (see `requirements.md`). Records landed under `task_memory/`, `.gitignore` narrowed, `plan.md` carries the Amendments table. diff --git a/task_memory/task_2026-09-21_issue26_correctness_pr/requirements.md b/task_memory/task_2026-09-21_issue26_correctness_pr/requirements.md new file mode 100644 index 00000000..f1b63387 --- /dev/null +++ b/task_memory/task_2026-09-21_issue26_correctness_pr/requirements.md @@ -0,0 +1,28 @@ +# Issue 26 Correctness PR — Requirements + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Recorded the original request, the specification hand-off, and the decisions from the planning interview. | + +## [Original Request] 2026-09-21 + +"this is a new task and current is plan mode, we need to discuss and land key docs first. plz read and analysis draft `.local-draft/Frontier_Issue26_Correctness_PR_Execution_Spec_2026-09-21.md` and land doc. if there is something need to consult me, use grill-me skill." + +The draft specification is landed verbatim in `plan.md` together with an Amendments table. + +## Decisions from the planning interview (2026-09-21) + +| Question | Decision | +| --- | --- | +| Q1/Q7 Task document location | `task_memory/` inside the PR worktree; narrow `.gitignore` exception so the records are published with the PR. No `docs/development/` tree. | +| Q2 Step 7 collective-sim zero-payload | Decide at Step 7 (facts: candidate gitlink `e564935d…` unreachable on the configured remote; no local clone). | +| Q3 GPU path | StepMind runbook exists; GPU workers available. Simulator runs on local CPU; GPU only for Step 6 native check and additional profiling CSVs. | +| Q4/Q8 2,000-line gate | Full cleanup + split for the four touched oversized modules only. | +| Q9/Q10 PR organization | Stacked: `refactor/oversized-module-split` first, `fix/issue26-correctness-pr` based on it; both draft PRs open, correctness base = refactor branch, retarget after merge. | +| Q11 Refactor acceptance | ≥50 scenarios; per scenario `request_metrics.csv` value-identical and `system_metrics.json` identical after removing timestamps/run ids; any difference is FAIL unless an explicitly approved fidelity fix. | +| Q12 Task directories | Two: `task_2026-09-21_oversized_module_split` and `task_2026-09-21_issue26_correctness_pr`. | +| Q5 Publication | Authorized: push both branches to `origin`, create/update draft PRs. Not authorized: merge, force-push, history rewrite, closing Issue 26. | +| Q6 Stop point | Stop after the Step 0 push for user review. | +| vLLM reference | Clone into `.real-engine/` (local exclude), pinned to `ea95f57`. | diff --git a/task_memory/task_2026-09-21_issue26_correctness_pr/review.md b/task_memory/task_2026-09-21_issue26_correctness_pr/review.md new file mode 100644 index 00000000..8d74f19a --- /dev/null +++ b/task_memory/task_2026-09-21_issue26_correctness_pr/review.md @@ -0,0 +1,35 @@ +# Issue 26 Correctness PR — Review Records + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Created with the pinned source snapshot. Dispositions are filled in Step 1. | + +## Pinned source snapshot + +| Repository / reference | Revision | Role | +| --- | --- | --- | +| `NetX-lab/Frontier` `main` | `1f694f7c549aa3aeeb7c5bbae04e119c09167a77` | Integration baseline (verified 2026-09-21 after `git fetch`). | +| `NetX-lab/Frontier` `bug/ttft-check` | `a7b3320fe9b8b083ee86b91dae3d6838f4443d91` | Candidate. Final commit touches only `task_memory/`. | +| Candidate parent before evidence import | `b7f8d055461d9208b8ae57eceeb6c0246cbc8d3c` | Source comparison point (verified equal source tree to `a7b3320`). | +| Merge base | `d71ad80b0800880808a0857fd30477e6d96592c6` | Verified with `git merge-base`. | +| `fwyc0573/vLLM-BS` | `ea95f571e20937c7c908c6d59ddd1cd6bf9268f1` | vLLM reference, `.real-engine/vLLM-BS`. | +| `fwyc0573/frontier-htsim` main gitlink | `b8518afcc310f0fe0e3ce52ba6b4f0bf57a3be04` | Current optional backend. | +| Candidate gitlink | `e564935d3874d8c71b52a554ab7c9a72e5e19f68` | Not reachable on the configured remote (HTTP 422, 2026-09-21). | + +## Candidate change dispositions + +To be completed in Step 1. Format: path, donor hunk summary, old defect, main behavior, reference behavior, disposition (`PORT`/`ADAPT`/`ALREADY_PRESENT`/`DROP`/`BLOCKED`), chosen owner, planned test. + +| Path | Disposition | Notes | +| --- | --- | --- | +| `tests/e2e/issue26_*`, `tests/integration/issue26_dp_*rca*`, `tests/performance/issue26_*` | DROP (default, per `plan.md` A11) | Experiment scripts; re-evaluate only a named helper. | + +## Design decisions + +Pending Step 1. + +## Final code-review findings + +Pending Step 8. diff --git a/task_memory/task_2026-09-21_issue26_correctness_pr/summary.md b/task_memory/task_2026-09-21_issue26_correctness_pr/summary.md new file mode 100644 index 00000000..e85c75c8 --- /dev/null +++ b/task_memory/task_2026-09-21_issue26_correctness_pr/summary.md @@ -0,0 +1,9 @@ +# Issue 26 Correctness PR — Summary + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Placeholder created at Step 0. | + +Completion archive; to be written at Step 8. Until then `progress.md` is authoritative. diff --git a/task_memory/task_2026-09-21_issue26_correctness_pr/validation.md b/task_memory/task_2026-09-21_issue26_correctness_pr/validation.md new file mode 100644 index 00000000..6743f2bd --- /dev/null +++ b/task_memory/task_2026-09-21_issue26_correctness_pr/validation.md @@ -0,0 +1,23 @@ +# Issue 26 Correctness PR — Validation Records + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Created. Environment recorded; baseline results are recorded in the refactor task's Step 0 report because both branches share the same base commit. | + +## Environment + +| Field | Value | +| --- | --- | +| Host | `kun-workspace-vgen2` (CPU master) | +| Python | `/data/ycfeng/envs/frontier-py310/bin/python` (uv-managed CPython 3.10.6) | +| Install | `uv pip install -e ".[test]"` from the active worktree | +| Scratch root | `FRONTIER_TMP_ROOT=/data/ycfeng/tmp/issue26-correctness-pr` | +| vLLM reference | `.real-engine/vLLM-BS` at `ea95f57` | + +## Per-checkpoint evidence + +| Step | Purpose | Source under test | Command | Outcome | Baseline comparison | Limits | +| --- | --- | --- | --- | --- | --- | --- | +| 0 | Baseline for touched areas | `1f694f7` | See `task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_step0_baseline.md` | 84 passed / 10 failed (pre-existing or environmental, listed in the report); two dummy smokes PASS | Base itself | Shared with the refactor branch (same base). | diff --git a/task_memory/task_2026-09-21_oversized_module_split/issues.md b/task_memory/task_2026-09-21_oversized_module_split/issues.md new file mode 100644 index 00000000..4398039d --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/issues.md @@ -0,0 +1,116 @@ +# Oversized Module Split — Issues and Resolutions + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Created with the issues found during the `config.py` split. | +| 2026-09-21 | Added I5 to I8, the four regressions the scheduler split introduced. | +| 2026-09-21 | Added I9 and I10, the test-binding changes the predictor splits required. | + +## I1 — `ClusterConfig` is constructed at runtime by a method that moved out + +| Field | Record | +| --- | --- | +| Found by | `tests/unit/test_pdaf_config_contract.py` and `tests/unit/test_simulator_transfer_predictor_lifecycle.py`, 6 failures | +| Symptom | `NameError: name 'ClusterConfig' is not defined` at `frontier/config/cluster_role_config.py:135` | +| Cause | `get_cluster_configs_for_disaggregation` builds one `ClusterConfig` per role. When that method moved into the `ClusterRoleConfigBuilder` mixin, the name was imported only under `TYPE_CHECKING`, because the survey had recorded it as an annotation. It is a real constructor call, and a module-level import would close a cycle: `cluster_config` imports the mixin. | +| Fix | A lazy import of `ClusterConfig` inside the method, matching the existing `_get_cc_backend_configs` pattern in the same package. | +| Prevention | Any name a moved method *constructs* or *calls* must be imported at runtime, not under `TYPE_CHECKING`. The static free-name analysis cannot distinguish the two; only executing the code does. | + +This is the reason the split is gated by the unit suites in addition to the fidelity matrix: no fidelity case had reached this path before the unit run, because the disaggregated example wrappers exercise it but the failure surfaced first in the faster unit selection. + +## I2 — Long-prefill co-location case exceeded the model context + +| Field | Record | +| --- | --- | +| Found by | The first fidelity baseline capture | +| Symptom | `Sequential simulation ended with non-empty scheduler state`, exit code 1, at simulated time 7.296 ms | +| Cause | The case asked for 4096 prefill tokens plus 16 decode tokens on Llama-2-7b, whose context is 4096. No request can ever be admitted, so the queue never drains. | +| Resolution | Case parameters corrected to 3584 prefill tokens. This is a property of the configuration, not a simulator defect, so nothing in `frontier/` changed. | + +## I3 — MoE fidelity cases violated the shared-domain invariant + +| Field | Record | +| --- | --- | +| Found by | The first fidelity baseline capture | +| Symptom | `ValueError: Frontier shared attention/MoE parallel domain requires attn_tp*attn_dp == moe_tp*moe_ep`, and the wrapper's own stricter guard `ATTN_TP == MOE_TP * MOE_EP` | +| Cause | Two invalid case topologies. | +| Resolution | The EP1 case now uses `ATTN_TP=1`. The MoE attention-DP case was replaced by an EP4 topology: the MoE wrapper's guard ignores `attn_dp`, so it cannot express `attn_dp > 1`, and the dense matrix already covers DP lanes. | + +## I4 — Predictor cache names differed for reasons unrelated to the refactor + +| Field | Record | +| --- | --- | +| Found by | The cache-file-name check in the fidelity comparator | +| Symptom | 138 of 426 cache entries had different hashes between the two checkouts while every simulated output was identical | +| Cause | `ExecutionTimePredictionModelManager._get_hash_relevant_config` includes the profiling input file paths in the model hash, and the two `examples/profiling/smoke_simulator_*_csv.sh` wrappers default their data base to an absolute path under their own repository root. | +| Resolution | Both cases now pass `DATA_DIR_BASE=data/profiling`, which resolves identically from either checkout. After recapture the difference is zero, so the check can now detect a real training-identity change. | + +## I5 — A moved method referenced a module-level logger object + +| Field | Record | +| --- | --- | +| Found by | The fidelity matrix, 66 of 67 cases | +| Symptom | `NameError: name '_frontier_vllm_v1_sched_decision_logger' is not defined`, raised from `vllm_v1_iteration_policy.py` inside `_emit_schedule_decision_event` | +| Cause | The decision-log guard tests the module-level logger object directly, not only the logging function. The extraction imported the function into the new module but not the object, which stayed behind in `vllm_v1_decision_log.py`. | +| Fix | `schedule_decision_logging_enabled()` added to `vllm_v1_decision_log.py`; the guard calls it. No module outside `vllm_v1_decision_log.py` now names the logger object. | + +## I6 — `deque` missing from the main scheduler module + +| Field | Record | +| --- | --- | +| Found by | `tests/unit/test_simulator_transfer_predictor_lifecycle.py` | +| Symptom | `NameError: name 'deque' is not defined` in `__init__` | +| Cause | The rebuilt import header dropped a name the retained code still uses. | +| Fix | Import restored. | + +## I7 — `validate_gdn_runtime_support` missing from the main scheduler module + +| Field | Record | +| --- | --- | +| Found by | `tests/unit/test_gdn_scheduler_slots.py` and five other GDN suites, 16 failures | +| Symptom | `NameError` in `__init__` | +| Cause | The guard is used both by the retained constructor and by the extracted allocation methods. The extraction moved the import instead of duplicating it. | +| Fix | Import restored in both modules. | +| Prevention | A static check now parses each split module and reports every name that is loaded but neither imported, defined locally, nor a builtin. It found this one. Its only remaining hit is `BaseCCBackendConfig` in `frontier/config/cluster_config.py`, which is a string annotation the flat CLI generator resolves through its own lazy-import special case, exactly as the pre-split `config.py` did. | + +## I8 — A test patched the decision logger on the module that no longer calls it + +| Field | Record | +| --- | --- | +| Found by | `tests/unit/test_prefix_cache_identity_ledger.py` | +| Symptom | The test captured zero events where it expected two | +| Cause | The test monkeypatches `_log_frontier_vllm_v1_schedule_decision` on the main scheduler module. The prefix-cache methods now live in `vllm_v1_prefix_cache.py` and resolve the name in that module's namespace, so the patch no longer intercepts. | +| Resolution | The patch target moved with the methods. This is the one test change in the scheduler step; the test's subject, the ledger's emitted events, is unchanged. | + +## I9 — Tests that patch a module-level name must follow the code that reads it + +| Field | Record | +| --- | --- | +| Found by | The unit selection, 8 failures across 6 files | +| Symptom | A patched helper had no effect, so a test observed an empty list or the unpatched value | +| Cause | These tests monkeypatch a module-level name on `shared_prediction_model_manager` and then drive a method that reads it. After the split the method resolves that name in the module it moved to, so the patch no longer intercepts. This is the same class as I8. | +| Resolution | Each patch target moved to the module that now owns the method under test: five files to `prediction_family_trainers`, one to `profiling_dataframe_loaders`, and one test to `prediction_model_identity` because `MOE_FAMILY` is read by `_get_moe_family_model_names`, which resolves it in its own module. | + +One of the eight is different in kind and deserves separate attention in review. `test_raw_model_profile_resolution_callsites_are_allowlisted` is a governance gate: it pins the exact set of functions permitted to resolve a raw model architecture profile, and how many times each may do so. Exactly one line of that allowlist changed. The entry keeps the same function name, the same kind, and the same expected call count of 1; only the owning file goes from `shared_prediction_model_manager.py` to `prediction_model_identity.py`, and the allowlist still holds 11 entries. No entry was added, removed, or given a higher count, so the property the gate exists to protect is unchanged. + +## I10 — A patched name that is now read in two modules + +| Field | Record | +| --- | --- | +| Found by | `tests/unit/test_moe_share_expert_operator_families.py`, 3 failures | +| Symptom | After repointing the patch to the module that owns the function being called, training still raised `Unsupported MoE op for TP mapping`, because part of the path still saw the real operator family | +| Cause | `MOE_FAMILY` was one binding in one module before the split. It is now imported by both `moe_predictor_helpers` and `moe_operator_times`, and the code path under test reads it in both. | +| Resolution | The three tests now install the fake family in both modules, one added line each. This is the only place in this branch where a test gained a line instead of having one changed, and it is a direct consequence of one name becoming two bindings. | +| Follow-up | The first attempt at this edit inserted the added line twice in two of the three tests. The script applied a literal replacement and then a regular expression that inspected only the `setattr` block, not the line following it, so the two blocks the literal pass had already handled were extended again. The duplicates were inert, since the second call set the same attribute to the same value, but they are a copy-paste artifact in a diff whose claim is a reviewed pure move. Removed in a follow-up commit rather than by amending the pushed one. | +| Note for review | A reviewer checking that the split preserved behavior should read this as evidence that it did: the test still asserts the same thing, and it needed the second patch precisely because the production code reads the name in both places. | + +## I11 — A pytest selection that aborts at collection proves nothing + +| Field | Record | +| --- | --- | +| Found by | The first unit parity run for the MoE predictor split | +| Symptom | Both sides reported the same 7 collection errors and no test results, and the comparison said IDENTICAL | +| Cause | Widening the selection pulled in seven modules that import `torch`, which the minimal CPU environment deliberately excludes. Pytest stops at collection errors by default, so nothing ran, and comparing two empty result sets trivially agreed. | +| Resolution | `--continue-on-collection-errors`, after which both sides run 3016 tests. The lesson is that a parity comparison has to assert that tests actually ran, not only that the two sides agree. | diff --git a/task_memory/task_2026-09-21_oversized_module_split/module_survey.md b/task_memory/task_2026-09-21_oversized_module_split/module_survey.md new file mode 100644 index 00000000..f46e5d28 --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/module_survey.md @@ -0,0 +1,229 @@ +# Oversized Module Split — Structural Surveys + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Collected four read-only structural surveys at `1f694f7`. | + +These surveys were produced by read-only, grep-backed inspection of `origin/main` at `1f694f7c549aa3aeeb7c5bbae04e119c09167a77`. Line numbers refer to that revision. Every "dead" or "unreferenced" claim must be re-verified with a fresh grep immediately before the corresponding cleanup edit in Step 2; the surveys are planning input, not deletion authority. + +--- + +## S1. `frontier/config/config.py` (5720 lines) + +### Inventory + +| Lines | Contents | +| --- | --- | +| 1–67 | Imports | +| 69–117 | Five `*_RELEASE_ERROR` strings, `DISAGGREGATED_CLUSTER_FIELD_PREFIXES`, `DISAGGREGATED_CLUSTER_FIELD_NAMES` | +| 120–138 | `_get_cc_backend_configs()` lazy-import shim (6-tuple) | +| 141–310 | Request interval/length generator configs (10 dataclasses) | +| 312–389 | Request generator configs (Base, Synthetic, Trace) | +| 390–464 | Small replica scheduler configs (Base, Vllm, Lightllm, Orca, FasterTransformer, Sarathi) | +| 465–869 | `VllmV1SchedulerConfig` (~400 lines; `__post_init__` 695) | +| 870–1091 | Sj2q family + Sglang (4 subclasses of VllmV1) | +| 1092–1251 | `MetricsConfig` | +| 1252–1861 | `SpeculativeDecodingConfig`, incl. six trace/JSON loader staticmethods 1371–1668 | +| 1862–2098 | `ReplicaConfig` (`__post_init__` 1963) | +| 2099–2138 | Cluster scheduler configs (6 small dataclasses) | +| 2139–2470 | Execution-time predictor configs (Base 2140–2426, LinearRegression 2428, RandomForrest 2452) | +| 2471–5066 | `ClusterConfig`: fields 2501–3850 (co-location 2501, disaggregated 2515, PD unified decode 2729, AF pipeline 2808, AFD CUDA graph 2845, per-cluster replica scheduler 2871, per-cluster CC backend 3071); methods 3851–5065 | +| 5067–5720 | `SimulationConfig(ABC)`: fields 5068–5299, `__post_init__` 5300, validators 5361–5645, accessors 5646–5702, `create_from_cli_args` 5703, `to_dict` 5710, `write_config_to_file` 5717 | + +`ClusterConfig` methods: `__post_init__` 3851; validators 3925/3966/3973/3994; monolithic setup 4000, disaggregated setup 4030; 4197–4234 dummy-mode and field-set introspection; `_create_replica_config_from_fields` 4236; `_validate_replica_config` 4309; cluster info/stats/printing 4386/4433/4506; `get_cluster_configs_for_disaggregation` 4585; predictor-per-cluster 4714; six `_create_*_cc_backend_config` 4748–5029; `_create_replica_config_copy` 5030. + +### Coupling + +- `ClusterConfig` holds defaults from nearly every other family (`RoundRobinClusterSchedulerConfig`, `SarathiSchedulerConfig`, `BaseExecutionTimePredictorConfig`, `ReplicaConfig`, lazy CC backend). +- `get_cluster_configs_for_disaggregation` (4585–4713) builds per-`ClusterType` copies and calls predictor/CC-backend factories. +- `SimulationConfig.__post_init__` (5300–5360) sets `global_vars` (5309–5336), prints cluster statistics (5355), normalizes the metrics dir (5358), writes the config file (5359): side effects inside dataclass init. +- Polymorphism: `BasePolyConfig` + `get_type()` on 24 classes keyed by `frontier.types` enums. Field-prefix dispatch through `DISAGGREGATED_CLUSTER_FIELD_PREFIXES`. +- CLI coupling: `create_flat_dataclass(cls)` 5705; `to_dict` depends on `__flat_config__` 5711. + +### Importers + +35 `from frontier.config.config import ...` sites in 29 files. `frontier/config/__init__.py` ends with `from .config import *`, so `from frontier.config import X` is also a public path. Most-imported names: `ReplicaConfig` (~44), `MetricsConfig` (~26), `ClusterConfig` (~20), `SimulationConfig` (~19), `RandomForrestExecutionTimePredictorConfig` (~15), `VllmV1SchedulerConfig` (~14), `BaseReplicaSchedulerConfig` (~9), `DISAGGREGATED_ARCHITECTURE_RELEASE_ERROR` (~7). Any split must preserve both import paths. + +### Cleanup candidates + +| Item | Lines | Evidence | +| --- | --- | --- | +| `BaseExecutionTimePredictorConfig.validate_linear_op_input` | 2397 | def only, zero call sites | +| `print_cluster_statistics` | 4506 | only caller 5355 | +| `write_config_to_file` | 5717 | only caller 5359 | +| Identical `hasattr(base_config, ...)` triplets | 4839–4846, 4951–4958, 4996–5003 | repeated 3x across CC-backend creators | +| Six near-parallel CC-backend creators | 4748–5029 | same shape, table-driven candidate | +| `hasattr` on declared dataclass fields | 4515, 5679, 5711 | defensive on guaranteed attributes | +| Two release-guard ladders from the same constants | 3994–3999, 5465–5481 | `DISAGGREGATED_ARCHITECTURE_RELEASE_ERROR` never raised in-file | +| Spec-decode trace loaders | 1371–1668 | file IO inside a config dataclass; each used once | +| Method-local re-imports | 4892–4893, 5309 | `replace` already imported at line 4 | +| 55 `getattr`/`hasattr` uses | — | concentrated in the disaggregation/CC-backend region | + +--- + +## S2. `frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py` (5138 lines) + +Module level: env-driven JSONL decision logger set up at import (58–79), `_log_frontier_vllm_v1_schedule_decision` 81–86, `PrefixCacheAdmission` 87–98, `_serialize_prefix_cache_binding` 99–109. `VLLMv1EngineReplicaScheduler(BaseReplicaScheduler)` 110–5138. + +### Inventory + +| Group | Lines | +| --- | --- | +| Construction / flag+state init (`__init__`, ~28 `getattr(self._config, ...)`) | 129–320 | +| Batch formation / active-set bookkeeping | 321–332, 395–418, 4948–5015 | +| Decode-attn cohort / wave (PD-AF) | 419–619, 4602–4612, 4613–4947 | +| Spec-decode / MTP monolithic-PP wait heuristics (largest theme, ~1300 lines) | 333–394, 620–1351, 2365–2696 | +| Prefix caching | 1352–1368, 1446–1638 | +| Request lookup / resource release | 1369–1445 | +| CUDA-graph capture sizing | 1639–1734 | +| Spec-decode batch metadata | 1735–1889 | +| Policy / iteration profile / fast lanes | 1890–2135 | +| Chunked-prefill budget / decision-log hooks | 2136–2209 | +| `on_batch_end` | 2210–2364 | +| Token ledger / KV accounting | 2697–2832 | +| KV block allocation | 2833–3024 | +| Preemption | 3025–3273, 3489–3538 | +| Phase 1 running | 3276–3488 | +| Phase 2 waiting / admission | 3539–3855 | +| Scheduling entry points (`_get_next_batch`, `_schedule_two_phase`, prefill/decode-only) | 3856–4601 | +| Public overrides (`num_pending_requests`, `peek_waiting_requests`, `is_empty`, `add_request`) | 5017–5138 | + +No DP-lane code exists in this file. GDN state-slot handling is threaded through init/free/allocate (30 hits). + +### Existing helpers not reused + +| Responsibility | Existing helper | Relationship | +| --- | --- | --- | +| Prefix caching | `scheduler/utils/prefix_cache.py` | not imported; local admission/ledger 1446–1638 | +| Idle diagnostics | `scheduler/utils/scheduler_diagnostics.py` | not imported; `is_empty` 5043–5079 near-copies base 689–710 | +| AFD metadata | `scheduler/utils/afd_metadata.py::aggregate_afd_metadata` | not imported; local `_attach_afd_metadata_if_needed` 4948–5015 | +| Batch building | `scheduler/utils/batch_builders.py` | not imported | +| MTP metrics | `scheduler/utils/mtp_metrics.py` | not imported; inline in `on_batch_end` | +| Spec decode | `frontier/spec_decode` | genuinely delegated | + +### Contract + +Overrides of `BaseReplicaScheduler` (9): `_create_batch`, `_get_request_next_num_tokens`, `complete_kv_transfer_for_requests`, `on_batch_end`, `_get_next_batch`, `num_pending_requests`, `peek_waiting_requests`, `is_empty`, `add_request`. + +Subclasses (`SGLangStyleReplicaScheduler`, three `SJ2Q*` schedulers) override private members: `_schedule_two_phase`, `_schedule_running_requests`, `_get_sorted_waiting_queue`, `_build_decode_waiting_queue`, `_get_request_next_num_tokens`, `_emit_schedule_decision_event`, `_resolve_iteration_round_class`, `_get_iteration_scheduler_profile`, `_maybe_promote_final_round_priority`, `_is_final_{prefill,decode}_fast_lane_request`, `on_batch_end`, `add_request`. These names must stay on the class (or be re-bound) after any split. + +External callers: `complete_kv_transfer_for_requests` (`events/kv_cache_transfer_end_event.py:62`), `consume_monolithic_pp_*_followup_poll` (`events/replica_schedule_event.py:101–122`, behind `hasattr`), `get_decode_attn_active_stage_slots` (`scheduler/utils/pdaf_attention.py:107`, via `getattr`). Seven unit-test files bind private methods via `object.__new__`. + +### Cleanup candidates + +| Finding | Lines | Evidence | +| --- | --- | --- | +| Dead `_attach_afd_metadata_if_needed` (68 lines) | 4948–5015 | zero references besides def; duplicates `afd_metadata.aggregate_afd_metadata` | +| Duplicated `is_empty` body vs base | 5043–5079 | copy of base 689–710 plus two terms | +| `hasattr` on state set in `__init__` | 5056–5060 | 3 `hasattr` total | +| 173 `getattr(` calls (42 `self`, 28 `self._config`, 51 `request`) | `__init__` 129–320, MTP block 620–1351 | defaults unreachable for own attributes | +| Import-time side effects (log dir + handler) | 58–79 | | +| Method-local re-imports | 4865–4866, 4966, 4977 | `global_vars` already imported at 28 | + +--- + +## S3. `frontier/execution_time_predictor/shared_prediction_model_manager.py` (4614 lines) + +### Inventory + +| Responsibility | Symbols (lines) | +| --- | --- | +| MoE family name helpers | 107–135, 382 | +| Model-architecture profile checks | 136–198 | +| Typed operator contract validation / identity | 199–381; methods 912–1052, 1107–1123 | +| Exact-lookup (query time) | 389–404, 4462–4487 | +| Class ctor / cluster requirement analysis | `ExecutionTimePredictionModelManager` 405, `__init__` 421–469 (18 registry attrs 424–449), 470–543 | +| Measurement family / device-event timer | 544–649, 4502 | +| sklearn estimator / scoring | 650–698 | +| Training orchestration | `_train_all_required_models` 699–842, GDN 843–892 | +| TP/EP key resolution | 893, 1124–1225 | +| Training signature | `_get_ffn_contract_signature` 1053–1106 | +| Per-family trainers | FFN 1340–1806, dense MLP 1807–1907, attention 1908–2346, MLA 2347–2532, residual 2533–2606, PP 2607–2644, TP 2645–2690, CPU overhead 2691–2764, `_train_single_model` 2765–2987 | +| MoE dataset contract validation | 1226–1339 | +| CSV loading + column validation | 2988–3797 | +| Derived features | 3798–3937 | +| Cache-key / identity | `_get_hash_relevant_config` 3938–4007, `_get_model_hash` 4008–4057 | +| Precision / measurement from df | 4058–4104 | +| Estimator registry & sharing | 4105–4389 | +| Query-time lookup | `get_model` 4390–4445 | +| Persistent cache + locking | 4446–4461 | +| Public API | `get_models` 4488, `get_models_for_cluster` 4510–4562, `get_required_capabilities` 4563, `get_training_file_paths` 4567, `get_training_context` 4583–4614 | + +### Consumers + +`sklearn_execution_time_predictor.py`: `get_gdn_predictor` (437), `get_models` (572), `get_models_for_cluster` (574). `simulator.py`: constructor (157), `get_training_file_paths` (177). No other production callers; the MoE predictor holds only a type annotation. Zero external references anywhere: `get_required_capabilities`, `get_training_context`, `get_model`. About 15 unit tests reach private methods via `object.__new__`. + +### Overlapping modules + +`cache_io.py` (atomic dumps, dataset fingerprint), `measurement_input_paths.py` (manager 616–649 and 4567 are thin wrappers), `profiling_metadata.py`, `attention_dataset_contract.py`, `attention_tp_policy.py`, `execution_time_predictor_registry.py`, `frontier/operators/typed_contracts.py`, `frontier/model_architectures.py`, `frontier/moe_routing_runtime.py`, `frontier/moe_gating_runtime.py`, `frontier/training/*_trainer.py` (`base_trainer.py:285` notes its MAPE must match this file: duplicated logic). + +### Routing-runtime identity today (input to correctness Step 5) + +| Site | Fact | +| --- | --- | +| 1294–1296, 1440–1442 | `resolve_moe_gating_routing_runtime_path(getattr(replica_config, "moe_routing_distribution_type", "balanced"))` | +| 1488–1502 | `runtime_path_key` is part of the per-call `moe_df_cache` tuple only when `base_model_name == "moe_gating_routing_topk"`; not persisted | +| 1377–1382 | `ffn_signature` has no routing-runtime component | +| 1053–1104 | `_get_ffn_contract_signature` has no routing-runtime component | +| 4048–4051 | `_get_model_hash` has no routing-runtime term; it enters only indirectly through `df_hash_str` of the filtered frame | +| 4105–4155 | cached identity validation compares layer cache identity only | + +### Cleanup candidates + +| Finding | Evidence | +| --- | --- | +| Dead `_get_moe_df_with_derived_features` (3919–3937) | def only | +| Unreferenced public `get_required_capabilities` (4563), `get_training_context` (4583) | def only | +| Test-only `_validate_cached_layer_cache_identity` (4105) | 2 in-file, 2 in one test | +| Duplicated MAPE (683) vs `sklearn_execution_time_predictor.py`, `training/base_trainer.py`, `vidur_cc_backend.py` | 4 definitions | +| Duplicated derived-feature helpers (3905/3912) | 3 files each | +| Four near-identical registry accessors 4156/4168/4180/4190 | `return getattr(self, attr)` bodies | +| `_serialize_selected_layer_cache_identity` recomputed at 242, 336, 1098, 1496, 2836, 4035, 4136, 4375, 4400 | same contract, four phases | +| 57 `getattr` + 7 `hasattr`; repeated `getattr(replica_config, "model_config", None)` at 923, 1013, 1125, 1827, 1873, 4318, 4349; default-`"balanced"` duplicated at 1295 and 1441 | defensive ladders | + +--- + +## S4. `frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py` (3539 lines) + +Single class `SklearnMoEExecutionTimePredictor(SklearnExecutionTimePredictor)` 254–3539 plus 8 module helpers 87–251. + +### Inventory + +| Responsibility | Members (lines) | +| --- | --- | +| Module helpers / operator-family naming | 87–251 | +| Dataset loading / filtering / contract | `_validate_moe_dataset_contract` 1186–1298, `_train_moe_models` 1299–1508 (`moe_df_cache` 1373–1430), 1509–1559 | +| Routing distribution / details generation | 261–278, 766–886, `_simulate_routing_per_layer` 2557–2630 | +| Expert-load / lane workload | 284–430, 887–1022, `_build_moe_load_imbalance_features` 1791–1844 | +| EP sync / shared-domain cost | 1023–1154, 1910–1999, 2953–3030 | +| Routing-cost model selection | 279–283, 1155–1185, 1560–1587 | +| Token-count resolution | 1618–1716, 2103–2160 | +| Gating / shuffling / grouped-GEMM compute | 1588–1617, 1717–1790, 1845–1909, 2000–2102 | +| Attention query caching | 2161–2209 | +| Layer/stage orchestration | 431–765, 2210–2556, 2631–2952, 3233–3539 | +| MTP replay | 3031–3232 | +| Diagnostics | 274, 758–763, 2809–2838, 2882–2925, 3332–3439 | + +### Inheritance and delegation + +Overrides (10): `__init__`, `_train_models`, `_predict_for_compute_models`, `_register_additional_profiling_metadata_from_files`, `predict_moe_layer_time`, `predict_allgather_time`, `predict_alltoall_time`, `predict_stage_execution_time`, `_predict_mtp_terminal_row_time_ms`, `_predict_mtp_decoder_layer_time_ms`. 43 new methods. Zero calls into `shared_prediction_model_manager` (annotation only). Subclassed by `sklearn_disaggregation_execution_time_predictor.py` (overrides `predict_stage_execution_time`; calls `_resolve_layer_lane_workload`, `_get_cluster_replica_config`). + +### Routing identity today (input to correctness Step 5) + +`_moe_routing_distribution_type` (715–724) is both the expert-load shape fed to `generate_moe_routing_ratios` (794–801) and the sole input selecting the profiling runtime path (731–734, 279–283). `_moe_gating_routing_runtime_path` (731–734) is assigned and never read. `_build_moe_load_imbalance_features` hardcodes `load_distribution="runtime"` (1817) and pops it (1820). Load distribution and routing implementation identity are conflated in one scalar. + +### Cleanup candidates + +| Finding | Evidence | +| --- | --- | +| `_is_grouped_gemm_on_demand_mode` (2000–2014) | zero references | +| `_simulate_routing_per_layer` (2557–2630) | test-only; self-described legacy | +| `predict_monolithic_decode_shared_domain_lane_moe_times_ms` (1060–1154) | test-only public API | +| `_moe_gating_routing_runtime_path` field | written, never read | +| 29 `getattr` + 5 `hasattr` on own attributes set in `__init__` (e.g. 1245, 1384, 1693, 2565; 328, 1580, 3157; 1688–1689) | unreachable defaults | +| Duplicated MoE-layer classification | 431–475 vs 3127–3136 and 3191–3202 | +| Duplicated share-expert triples | 1103–1107 and 2330–2333 | +| Per-layer re-queries of gating/shuffling/grouped-GEMM/EP-comm without memoization | 3243–3255; only attention is cached 2161–2194 | +| INFO-level diagnostic blocks on every prediction | 2882–2925, 3405–3439 | diff --git a/task_memory/task_2026-09-21_oversized_module_split/plan.md b/task_memory/task_2026-09-21_oversized_module_split/plan.md new file mode 100644 index 00000000..d4798ef4 --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/plan.md @@ -0,0 +1,221 @@ +# Oversized Module Split — Plan + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Initial plan: scope, proposed split boundaries, sequencing, fidelity matrix design, acceptance criteria. Boundaries are proposals derived from `module_survey.md`; each is confirmed against the code before its step starts. | +| 2026-09-21 | Recorded the boundaries actually implemented for `config.py` in section 3.1, which differ from the proposal. | +| 2026-09-21 | Recorded the boundaries actually implemented for the vLLM V1 replica scheduler in section 3.2. | +| 2026-09-21 | Recorded the boundaries actually implemented for the shared prediction model manager in section 3.3 and the MoE predictor in 3.4. | + +## 1. Scope and result + +Branch `refactor/oversized-module-split` (base `1f694f7c549aa3aeeb7c5bbae04e119c09167a77`) brings the four modules below under the `AGENTS.md` 2,000-line gate through a cleanup-first pass and a functional split, with **no behavior or numeric change**. The Issue 26 correctness branch is based on this branch, so every boundary chosen here must leave the correctness fixes with a clear owner (see §3). + +| Module | Lines | Target after split | +| --- | --- | --- | +| `frontier/config/config.py` | 5720 | ≤ 2,000 per child module | +| `frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py` | 5138 | ≤ 2,000 per child module | +| `frontier/execution_time_predictor/shared_prediction_model_manager.py` | 4614 | ≤ 2,000 per child module | +| `frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py` | 3539 | ≤ 2,000 per child module | + +Out of scope: the other five modules above 2,000 lines (`sklearn_execution_time_predictor.py` 8263, `metrics_store.py` 5586, `sklearn_disaggregation_execution_time_predictor.py` 2985, `profiling/attention/main.py` 2157, `entities/request.py` 2125), any behavior fix, any change to public CLI flags, config field names, metrics schema, or profiling CSV contracts. + +## 2. Rules for every edit + +1. **Cleanup before split.** Remove only code whose lack of references is re-verified by grep at edit time. Replace `getattr(self, "_x", default)` on attributes assigned in `__init__` with direct access only when the assignment is unconditional; otherwise leave it and record the reason. +2. **Move, do not rewrite.** A split moves functions and classes verbatim (imports adjusted). Renames are limited to module paths; class and public method names stay. Private methods overridden by subclasses or bound by tests (listed in `module_survey.md`) stay on the class or are re-exposed under the same name. +3. **Preserve import surfaces.** `frontier/config/__init__.py` star re-export and `from frontier.config.config import X` keep working through re-exports in the original module. Same for the three other modules. +4. **Fidelity gate after each step.** The matrix in §5 must be value-identical to the main baseline after every step, not only at the end. A step that changes any number is reverted or explained as an approved fidelity fix. +5. **Public draft PR after Step 1.** Each later step is committed and pushed as a coherent unit with its matrix result recorded in `progress.md` and `validation.md`. + +## 3. Proposed boundaries (to confirm at each step) + +### 3.1 `frontier/config/` (implemented) + +The proposal below was revised during implementation. Two things forced the change. First, `ClusterConfig` alone is 2,588 lines, so leaving it in `config.py` would have kept that file above the gate no matter which leaf families moved out. Second, `flat_dataclass` resolves string annotations in the *defining module's* namespace, so every module must carry the imports its own annotations need; that is a correctness constraint on the split, not a style choice. + +What was implemented: + +| Module | Lines | Content | +| --- | --- | --- | +| `config.py` | 788 | `SimulationConfig`, the lazy CC-backend import note, and the re-export block that keeps `from frontier.config[.config] import X` working for all 36 names other modules import | +| `cluster_config.py` | 1888 | `ClusterConfig`: the flat per-role field surface, `__post_init__`, the validators, monolithic and disaggregated setup, `_validate_replica_config` | +| `cluster_role_config.py` | 586 | `_get_cc_backend_configs` and `ClusterRoleConfigBuilder`: the 13 methods that build the per-role replica, predictor and CC-backend configurations | +| `cluster_topology_summary.py` | 217 | `ClusterTopologySummary`: `_collect_cluster_info`, `get_server_count_metadata`, `print_cluster_statistics` | +| `release_guards.py` | 56 | The six release-guard messages and the two disaggregated field-name tables | +| `request_generator_config.py` | 262 | Arrival interval, request length and request generator families | +| `replica_scheduler_config.py` | 711 | Every replica scheduler configuration, including the vLLM V1 family and its Sj2q and SGLang subclasses | +| `metrics_config.py` | 173 | `MetricsConfig` | +| `speculative_decoding_config.py` | 622 | `SpeculativeDecodingConfig` and its trace loaders | +| `replica_config.py` | 250 | `ReplicaConfig` | +| `cluster_scheduler_config.py` | 48 | The cluster scheduler configuration family | +| `execution_time_predictor_config.py` | 314 | The predictor configuration family and its calibration scales | + +The two groups extracted from `ClusterConfig` are **mixins that `ClusterConfig` inherits**, not free functions taking the config. That keeps the split a pure move: the method bodies, the method names and every call site are unchanged, including the four methods that external code and tests call directly (`get_cluster_configs_for_disaggregation`, `get_server_count_metadata`, `_validate_replica_config`, `_create_replica_config_copy`). Rewriting them as free functions would have touched every line of 731 moved lines and made the diff unreviewable, with the fidelity matrix as the only remaining check. + +Naming note for review: `ClusterRoleConfigBuilder` and `ClusterTopologySummary` describe what each group produces. If a reviewer prefers different names, renaming them is mechanical and affects only three files. + +### 3.1a Original proposal (superseded) + +| Child module | Content (survey lines) | Approx. lines | +| --- | --- | --- | +| `config.py` (kept) | `SimulationConfig`, `ClusterConfig` core fields and validators, re-exports of everything below | ~1,900 after moves | +| `request_generator_config.py` | interval/length/request generator dataclasses (141–389) | ~250 | +| `replica_scheduler_config.py` | small scheduler configs + `VllmV1SchedulerConfig` + Sj2q/Sglang (390–1091) | ~700 | +| `speculative_decoding_config.py` | `SpeculativeDecodingConfig` (1252–1861) | ~610 | +| `replica_config.py` | `ReplicaConfig` (1862–2098) + cluster scheduler configs (2099–2138) | ~280 | +| `execution_time_predictor_config.py` | predictor configs (2139–2470) | ~330 | +| `metrics_config.py` | `MetricsConfig` (1092–1251) | ~160 | +| `cluster_config_factories.py` | `get_cluster_configs_for_disaggregation`, `_create_replica_config_from_fields`, `_create_replica_config_copy`, the six `_create_*_cc_backend_config` (4236–4308, 4585–4713, 4748–5066) as functions taking the `ClusterConfig` | ~900 | + +Cleanup first: the three duplicated `hasattr(base_config, ...)` triplets and the six parallel CC-backend creators become one table-driven helper; `validate_linear_op_input` (dead), method-local re-imports, and `hasattr` on declared dataclass fields are removed. Existing `frontier/config/cluster_scheduler_config.py` in the donor branch is **not** copied; the boundary above is chosen from main's structure. + +Correctness-branch owner after split: the opt-in DP placement config (Step 4) lands in `replica_config.py` next to the cluster scheduler configs; routing runtime override (Step 5) lands in `replica_config.py`. + +### 3.2 `frontier/scheduler/replica_scheduler/` (implemented) + +Same mixin mechanism as the configuration split, for the same reason: the move stays a move, and the private methods four subclasses override keep resolving correctly because the mixins sit before `BaseReplicaScheduler` in the base list. That ordering was verified for every overridden name. + +| Module | Lines | Content | +| --- | --- | --- | +| `vllm_v1_engine_replica_scheduler.py` | 1386 | The class shell, `__init__`, batch creation and active-set bookkeeping, `complete_kv_transfer_for_requests`, `on_batch_end`, phase 1 and phase 2 scheduling, the two-phase entry point, and the public overrides | +| `vllm_v1_mtp_wait.py` | 1142 | `TargetEmbeddedMtpWaitPolicy`: admission delay, output wait and terminal release timing for target-embedded MTP under monolithic pipeline parallelism | +| `vllm_v1_role_schedules.py` | 858 | `DisaggregatedRoleScheduling`: the prefill-only, decode-only, decode-waiting and decode-attention entry points a disaggregated cluster drives | +| `vllm_v1_kv_allocation.py` | 713 | `KvBlockAllocation`: token accounting, block allocation, preemption and resource release | +| `vllm_v1_iteration_policy.py` | 597 | `IterationSchedulingPolicy`: scheduling policy, fast lanes, CUDA graph capture sizing, speculative-decoding batch metadata, decision-log emission | +| `vllm_v1_prefix_cache.py` | 253 | `PrefixCacheAdmission`, `PrefixCacheLedger`: admission and identity events | +| `vllm_v1_decode_attn_cohort.py` | 219 | `DecodeAttentionCohort`: cohort identity, stage slots and phase for the PD-AF decode-attention role | +| `vllm_v1_decision_log.py` | 44 | The optional JSONL decision log and its enabled predicate | + +Cleanup first: removed `_attach_afd_metadata_if_needed`, 65 lines with no caller anywhere, which duplicated `frontier/scheduler/utils/afd_metadata.py`. + +Four regressions were introduced and fixed; `issues.md` I5 to I8 record each one, what found it, and what prevents the class of defect from recurring. + +### 3.2a Original proposal (superseded) + +| Child module | Content (survey lines) | Approx. lines | +| --- | --- | --- | +| `vllm_v1_engine_replica_scheduler.py` (kept) | class shell, `__init__`, scheduling entry points, phase 1/2, admission, `on_batch_end`, public overrides; delegates to the helpers below | ~1,900 | +| `vllm_v1_kv_allocation.py` | token ledger / KV accounting, block allocation, preemption (2697–3273, 3489–3538) | ~700 | +| `vllm_v1_mtp_wait.py` | target-embedded MTP monolithic-PP wait heuristics and terminal release (333–394, 620–1351, 2365–2696) | ~1,300 | +| `vllm_v1_prefix_cache.py` | prefix-cache admission/ledger and identity events (87–109, 1352–1368, 1446–1638) | ~250 | +| `vllm_v1_decode_attn_cohort.py` | PD-AF decode-attn cohort/wave logic (419–619, 4602–4947) | ~550 | +| `vllm_v1_iteration_policy.py` | policy / iteration profile / fast lanes / CUDA-graph capture sizing / spec-decode batch metadata (1639–2135) | ~500 | + +Mechanism: helpers become plain functions or small stateless classes that take the scheduler instance; the class keeps thin methods with the original names so subclass overrides and `object.__new__` tests keep binding. Cleanup first: dead `_attach_afd_metadata_if_needed`, duplicated `is_empty` body, method-local re-imports, `getattr(self, ...)` on attributes assigned unconditionally in `__init__`. The import-time log handler (58–79) is kept but moved behind a function called from `__init__` only if the current behavior (env-var gated) is preserved exactly; otherwise it stays. + +Correctness-branch owner after split: request-load accounting for Step 4 (`waiting` / `running` counts) is added to the kept class as one accessor; no change to the helper modules. + +### 3.3 `frontier/execution_time_predictor/` (manager implemented) + +| Module | Lines | Content | +| --- | --- | --- | +| `shared_prediction_model_manager.py` | 722 | Construction, cluster requirement analysis, measurement-family and input-file selection, the estimator and scorer factory, training orchestration, the GDN predictor, and the public API | +| `prediction_family_trainers.py` | 1700 | `PredictionFamilyTrainers`: one method per operator family, plus the shared fitting routine | +| `profiling_dataframe_loaders.py` | 985 | `ProfilingDataFrameLoaders`: one loader per profiling CSV, its column validation, the feature-column constants and the derived features | +| `prediction_model_registry.py` | 587 | `PredictionModelRegistry`: cache keys, trained-model identity, the in-memory registries and the persistent cache | +| `layer_contract_resolution.py` | 495 | `LayerContractResolution`: typed layer contract and TP/EP key resolution, the FFN contract signature, the MoE dataset contract | +| `prediction_model_identity.py` | 324 | Module-level identity helpers with no state: operator-family names, architecture-profile resolution, layer cache identity, typed contract matching | + +Cleanup first, all three verified definition-only repo-wide: `get_required_capabilities`, `get_training_context` and `_get_moe_df_with_derived_features`. + +The seam this split has to respect is named in the stacked PR's `review.md` under W5: the routing-implementation identity is lost because `ffn_signature` carries no routing term and `trained_model_signatures` is shared across clusters. The signature therefore lands in `layer_contract_resolution.py` and the cache keys and registry in `prediction_model_registry.py`, so the later fix edits those two modules and not the loaders. + +Eight tests needed their patch target moved; `issues.md` I9 records why and what the one governance-allowlist change does and does not alter. + +### 3.3a Original proposal (superseded) + +| Child module | Content (survey lines) | Approx. lines | +| --- | --- | --- | +| `shared_prediction_model_manager.py` (kept) | `ExecutionTimePredictionModelManager` ctor, cluster requirement analysis, training orchestration, public API, query-time `get_models*` | ~900 | +| `prediction_model_registry.py` | estimator registry & sharing, precision buckets, cache identity validation, `get_model` (4105–4445) | ~350 | +| `prediction_model_cache.py` | `_get_hash_relevant_config`, `_get_model_hash`, persistent load/store with locking (3938–4057, 4446–4487), reusing `cache_io.py` | ~200 | +| `profiling_dataframe_loaders.py` | CSV loaders, column validation, derived features (2988–3937) | ~950 | +| `family_trainers.py` | per-family trainers and `_train_single_model` (1340–2987) | ~1,650 | +| `layer_contract_resolution.py` | typed operator contract helpers, TP/EP key resolution, FFN signature, MoE dataset contract (199–381, 893–1339) | ~600 | + +Cleanup first: dead `_get_moe_df_with_derived_features`, unreferenced `get_required_capabilities` / `get_training_context`, the four registry accessor bodies, duplicated MAPE (keep one definition and import it where `base_trainer.py` says it must match), duplicated derived-feature helpers, repeated `getattr(replica_config, "model_config", None)` ladders. + +Correctness-branch owner after split: routing-runtime identity (Step 5) enters `layer_contract_resolution.py` (signature) and `prediction_model_cache.py` (hash) and `prediction_model_registry.py` (lookup) at one clearly named point each. + +### 3.4 `sklearn_moe_execution_time_predictor.py` (implemented) + +| Module | Lines | Content | +| --- | --- | --- | +| `sklearn_moe_execution_time_predictor.py` | 1557 | The class shell, dummy-mode timing, layer classification, the attention query cache, the layer and stage orchestration and the public prediction entry points | +| `moe_operator_times.py` | 710 | `MoeOperatorTimes`: gating, routing top-k, shuffling, grouped expert GEMM, the expert-parallel collective, and the token-count resolution each needs | +| `moe_routing_workload.py` | 583 | `MoeRoutingWorkload`: the expert-load distribution and the per-lane routed workload it produces | +| `moe_dataset_training.py` | 438 | `MoeDatasetTraining`: dataset admission, per-operator training, and the load-imbalance feature list | +| `moe_mtp_replay.py` | 216 | `MoeMtpReplay`: MoE time for speculative-decoding replay rows | +| `moe_predictor_helpers.py` | 176 | The module-level helpers, in a leaf module so the mixins can use them without importing the predictor | + +Cleanup removed `_is_grouped_gemm_on_demand_mode`, 14 lines with no reference anywhere. + +Three tests needed a second patch target rather than a moved one: `MOE_FAMILY` is now read in two modules, so a fake family has to be installed in both for it to be seen end to end. That is recorded in `issues.md` I10, and it is the one case in this branch where a test gained a line rather than changing one. + +### 3.4a Original proposal (superseded) + +| Child module | Content (survey lines) | Approx. lines | +| --- | --- | --- | +| `sklearn_moe_execution_time_predictor.py` (kept) | class shell, `__init__`, layer/stage orchestration, MTP replay, overrides (431–765, 2210–2556, 2631–2952, 3031–3539) | ~1,700 | +| `moe_routing_workload.py` | routing distribution/details generation, expert-load and lane workload (261–430, 766–1022, 1791–1844, 2557–2630) | ~600 | +| `moe_operator_times.py` | gating / shuffling / grouped-GEMM / EP-communication / token-count helpers and module-level operator-family helpers (87–251, 1155–1185, 1560–1790, 1845–2160) | ~800 | +| `moe_dataset_contract.py` | `_validate_moe_dataset_contract`, `_train_moe_models` dataset filtering (1186–1559) | ~380 | + +Cleanup first: `_is_grouped_gemm_on_demand_mode` (dead), the never-read `_moe_gating_routing_runtime_path` field (**kept until Step 5 of the correctness PR decides its owner**; recorded, not removed here), `getattr(self, ...)` on own attributes, duplicated MoE-layer classification and share-expert triples. The INFO-level diagnostic blocks are left untouched in this PR because log output is not part of the fidelity gate and changing them is not a size problem. + +Correctness-branch owner after split: separating load distribution from routing implementation identity (Step 5) is a change inside `moe_routing_workload.py` (load shape) and `moe_dataset_contract.py` (runtime row selection). + +## 4. Sequencing and dependencies + +```text +Step 0 (worktree, records, env, baseline) + -> Step 1 (fidelity harness + main baseline capture) + -> Step 2 (cleanup pass, all four modules; matrix) + -> {Step 3 config split, Step 4 scheduler split} # independent, may run in parallel worktrees + -> {Step 5 manager split, Step 6 MoE predictor split} # 6 depends on 5 only for shared helper placement + -> Step 7 (full matrix + unit suites + PR hand-off) +``` + +Steps 3 and 4 touch disjoint packages and can be parallelized; Steps 5 and 6 share `frontier/execution_time_predictor/` and are sequenced. Each step ends with matrix PASS, commit, push. + +## 5. Fidelity matrix design (Step 1 deliverable) + +- **Harness location:** `tests/e2e/refactor_fidelity/` with a small driver that (a) generates a deterministic case manifest, (b) runs each case through the checked-in example wrappers or `python -m frontier.main` in two worktrees (baseline main at `1f694f7` and the refactor branch) with the same environment, (c) compares outputs. It reuses `tests/scratch_root.py`, the example wrappers, and case-generation ideas from `tests/e2e/moe_ep_non_dummy_matrix.py`, but does not depend on that 7,900-line harness or its pinned baseline commit. +- **Comparison rule:** `request_metrics.csv` compared value-by-value (exact equality; both runs are deterministic CPU simulations); `system_metrics.json` compared after removing timestamps, run ids, wall-clock durations, host/Python metadata, and output paths. The removal list is explicit in the harness and recorded in `validation.md`. +- **Coverage (≥ 50 cases):** + +| Dimension | Values | +| --- | --- | +| Architecture | co-location, sequential PDD, sequential PD-AF | +| Model | dense (`llama2_7b_dense_example`), MoE (`Qwen3-30B-A3B-tiny`, `Phi-tiny-MoE-instruct`), PD-AF EP=2 topology | +| Predictor | dummy mode; checked-in `data/profiling/compute/h800/*` CSVs for dense and MoE | +| Mode | offline, online (Poisson arrivals) | +| Requests | 4 / 16 / 64 requests; prompt lengths 128 / 1024 / 4096; decode 16 / 128 | +| QPS (online) | 0.5, 2, 8 | +| Features | Chunked Prefill on/off, `decode_cuda_graph_mode` none/full_decode_only, prefix caching (fixture trace), speculative decoding (MoE recipe), Thinking Mode | +| Parallelism | TP1/TP2, PP1/PP2, `attn_dp` 1/2, EP 1/2 | + +The manifest fixes the exact combinations (not a full cross product) so the total is 50–70 cases with predictable wall time. Cases that main itself fails are recorded as baseline failures and excluded from the pass count with their exception. + +- **Baseline capture:** the main-side outputs are produced once from a read-only checkout of `1f694f7` (a separate worktree `.worktrees/fidelity-baseline-main`, never edited) and stored under `FRONTIER_TMP_ROOT/refactor-fidelity/baseline/`; their manifest hash is recorded in `validation.md`. + +## 6. Acceptance criteria + +- Each of the four modules and every new child module is ≤ 2,000 lines; the count is recorded in `validation.md` per step. +- Fidelity matrix: all non-baseline-failing cases identical after every step and at the end. +- Existing unit suites that touch the four modules pass or fail exactly as on main (baseline-failure list recorded at Step 0). +- No public import path breaks: a test imports every name listed in `module_survey.md` §S1 "Importers" from both `frontier.config` and `frontier.config.config`, and the corresponding names for the other three modules. +- Draft PR opened against `main` with the matrix summary, module line counts, and links to these records; no merge without user approval. + +## 7. Validation commands (recorded per step in `validation.md`) + +```bash +export PYTHONPATH=/data/ycfeng/Frontier/.worktrees/oversized-module-split +export FRONTIER_TMP_ROOT=/data/ycfeng/tmp/issue26-correctness-pr +PY=/data/ycfeng/envs/frontier-py310/bin/python +$PY -m pytest -q -p no:cacheprovider +$PY tests/e2e/refactor_fidelity/run_matrix.py --baseline-worktree .worktrees/fidelity-baseline-main --candidate-worktree . --manifest tests/e2e/refactor_fidelity/manifest.json +``` diff --git a/task_memory/task_2026-09-21_oversized_module_split/progress.md b/task_memory/task_2026-09-21_oversized_module_split/progress.md new file mode 100644 index 00000000..95f05bcb --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/progress.md @@ -0,0 +1,60 @@ +# Oversized Module Split — Progress + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Step 0: worktree and branch created, records landed, module surveys collected, baseline pending. | +| 2026-09-22 | Status header corrected (it still described Step 0 after Steps 1-6 had landed). Checkpoints A and B recorded. | +| 2026-09-22 | External review finding C34-01 applied: the predictor-cache name comparison now requires each side's manifest `cases_executed_in_last_run` to equal the full case table, so `--start`/`--limit` continuations no longer count as clean full runs; four gate tests added; the Checkpoint B verdict re-derived with the corrected rule and unchanged. | + +## Status + +| Field | Value | +| --- | --- | +| Branch | `refactor/oversized-module-split` | +| Base | `1f694f7c549aa3aeeb7c5bbae04e119c09167a77` (`origin/main`, fetched 2026-09-21) | +| Worktree | `/data/ycfeng/Frontier/.worktrees/oversized-module-split` | +| Python | `/data/ycfeng/envs/frontier-py310/bin/python` | +| Current step | Checkpoints A and B complete; C34-01 correction landed 2026-09-22. PR #34 is open for review (GitHub API state: not draft). | +| Publication | PUSHED_VERIFIED through the C34-01 correction (SHA in the commit log); PR #34 description synchronized | +| Next action | Checkpoint C, on the correctness branch: rebase PR #35 onto the corrected harness, re-measure W2 with source and harness at one revision, and strengthen the W2 placement tests (R35-01). Coordinate first: that worktree is shared with the W3 owner. | + +## Steps + +| Step | Work | Status | +| --- | --- | --- | +| 0 | Worktree, records, environment, baseline | PASS | +| 1 | Fidelity matrix harness and main baseline capture | PASS (67 cases, baseline 67/67) | +| 2 | Cleanup pass per module, matrix re-run | PASS (all four modules) | +| 3 | Split `config.py` | PASS (12 modules, largest 1888 lines) | +| 4 | Split `vllm_v1_engine_replica_scheduler.py` | PASS (8 modules, largest 1386 lines) | +| 5 | Split `shared_prediction_model_manager.py` | PASS (6 modules, largest 1700 lines) | +| 6 | Split `sklearn_moe_execution_time_predictor.py` | PASS (6 modules, largest 1557 lines) | +| 7 | Full matrix, unit suites, draft PR hand-off | PASS (final record in the Checkpoint B report) | +| A | Review comments R34-01 / R34-02: fidelity gate correctness and provenance | PASS | +| A+ | External review C34-01: cache-name comparison eligibility decided by the executed case list, not the filter field alone | PASS (`test_report_2026-09-22_cache_eligibility_correction.md`) | +| B | Review comment R34-03 / R34-04: final evidence record and retained checks | PASS | + +## Chronological updates + +- 2026-09-22: External review C34-01. `compare_labels` treated `cache_clean_before_run` plus an empty `case_filter` as proof of one clean full run, but `--start` and `--limit` narrow the executed selection without setting a filter, so a clean partial run merged onto full retained results passed the predicate. The predicate now also requires `set(cases_executed_in_last_run) == full case table` on both manifests, and a manifest without that field is ineligible. Four gate tests cover `--limit`, `--start`, symmetric partial caches and the missing field; the healthy case now asserts `predictor_cache_compared: true`. The Checkpoint B labels (`baseline_v2`, `candidate_bb582a4`) both record 71 executed cases, and rerunning `compare` on the retained outputs with the corrected rule gives the same verdict: 71 of 71 identical, `predictor_cache_populated_cleanly: true`, `predictor_cache_compared: true`, 0 cache differences. Report: `test_report_2026-09-22_cache_eligibility_correction.md`. +- 2026-09-22: Checkpoint B. The review's cheaper remedy for R34-03 was checked and is not available: `db15e64..5ef96b5` touches only `cases.py`, so the production tree was unchanged at the tip, but `candidate_db15e64` held 67 records with no DP case, and the old `baseline` label was an assembled partial run. Both sides were therefore recaptured as single clean full 71-case runs from detached checkouts, driven by the harness committed at `bb582a4`: baseline_v2 at `1f694f7` and candidate_bb582a4 at `bb582a4`, both clean, unfiltered, cache-cleaned, 71 executed, 426 cache files each. Result: 71 of 71 compared, 71 identical, 0 mismatched, 0 failures of any kind, 0 missing evidence, 0 definition differences, no provenance findings, and the predictor cache names compared and matched. Whole `tests/unit`: 84 failed / 3679 passed / 49 skipped here against 84 / 3644 / 49 on `1f694f7`, with the 84 failure identities byte-identical and the 35 extra passes being exactly the tests added in Checkpoints A and B; ten modules are excluded on both sides for missing `torch` or `matplotlib`. R34-04 landed as 13 tests in `tests/unit/test_module_split_boundaries.py`. Separately, all 142 estimators in the baseline-produced cache unpickle and run under the split code. See `test_report_2026-09-22_checkpoint_b_final_evidence.md`. +- 2026-09-22: One finding pinned rather than fixed. `typing.get_type_hints(ClusterConfig)` raises `NameError: BaseCCBackendConfig`, because the annotation names a class the module does not import at runtime. Verified to fail identically on `1f694f7`, so it predates the split; the generated CLI is unaffected. Recorded in `KNOWN_UNRESOLVED_CONFIG_ANNOTATIONS` so a new occurrence fails the test. + +- 2026-09-22: Checkpoint A. The fidelity gate could return 0 without comparing anything: `baseline_failures` was absent from the failure predicate and completeness tested case-id presence rather than successful comparison, and `list_artifacts` made two deleted artifact directories compare equal. Completeness now means compared; every reason a case was not compared is a failure; an absent directory is reported; each case record carries `source_revision`, `source_dirty`, `harness_revision` and a case-definition digest; a continuation whose retained records disagree is refused before any case runs; `measure_commit` refuses a dirty reused checkout without cleaning it. Four false successes reproduced against the pre-fix harness and none against the fixed one, with both controls unchanged. 22 new tests in `tests/unit/test_refactor_fidelity_gate.py`; 30 passed across every test that mentions the harness. No file under `frontier/` was touched. See `test_report_2026-09-22_checkpoint_a_fidelity_gate.md`. +- 2026-09-22: Consequence to carry into Checkpoint B: the provenance stamp is new, so every label captured before today fails the gate for want of provenance, and the `baseline` label was in any case an assembled partial run (last execution a filtered `dp_` run of four cases merged onto 67 without a cache clean, `case_count` 72 over 71 result lines). Both sides must be recaptured as single clean full runs. + +- 2026-09-21: Branch created from `1f694f7`. `.gitignore` narrowed to `task_memory/*` with exceptions for the two task directories. Structural surveys of the four modules collected into `module_survey.md` (read-only inspection, grep-backed). +- 2026-09-21: Environment `/data/ycfeng/envs/frontier-py310` created (CPython 3.10.6). Baseline: 84 passed / 10 failed (all in `test_colocation_release_review_contracts.py`, caused by debug scripts absent from main and a bare `python` executable missing on PATH); co-location and PDD dense dummy smokes PASS. See `test_report_2026-09-21_step0_baseline.md`. +- 2026-09-21: Step 1 complete. Harness `tests/e2e/refactor_fidelity/` built: 67 cases across co-location, sequential PDD and sequential PD-AF, dense and MoE, offline and online, dummy and checked-in-CSV predictors. Self-check (same code twice) 13/13 identical. Baseline capture 67/67 after correcting two invalid case topologies. See `test_report_2026-09-21_step1_fidelity_matrix.md`. +- 2026-09-21: Step 2 started with `config.py`. Removed the unreferenced `validate_linear_op_input`; replaced the two CC-backend dispatch ladders with one ordered table; deduplicated four field-resolution closures and three identical field triplets; dropped a redundant method-local import. 5720 to 5687 lines. Generated CLI flag set identical (753 flags). Matrix 67/67 identical. A 138-name predictor cache difference was traced to the two profiling smoke wrappers defaulting their data base to an absolute path under the repository root, which enters the model hash; both cases now use a repository-relative base and both sides were recaptured. +- 2026-09-21: Step 3 started. Extracted eight leaf modules from `config.py`: `release_guards.py`, `request_generator_config.py`, `replica_scheduler_config.py`, `metrics_config.py`, `speculative_decoding_config.py`, `replica_config.py`, `cluster_scheduler_config.py`, `execution_time_predictor_config.py`. `flat_dataclass` resolves string annotations in the defining module's namespace, so each module carries the imports its own annotations need; the compatibility gate is the unchanged CLI flag set plus the 36 names other modules import from `frontier.config`. +- 2026-09-21: `config.py` split complete and gated. Twelve modules, largest 1888 lines. Gates: generated CLI flag set identical (753 flags); all 36 externally imported names resolve from both `frontier.config` and `frontier.config.config`; `ClusterConfig` keeps 181 dataclass fields and gains the two mixins in its MRO; the 62-file configuration unit selection gives 10 failed / 671 passed on both sides with byte-identical failure identities, all pre-existing drift; fidelity matrix 67 identical, 0 mismatched, 0 predictor cache name differences. +- 2026-09-21: One real defect was introduced and caught by the unit selection before the matrix reached it: `get_cluster_configs_for_disaggregation` constructs `ClusterConfig` at runtime and the extraction had imported the name only under `TYPE_CHECKING`. Recorded as I1 in `issues.md`, fixed with a lazy import. +- 2026-09-21: vLLM V1 replica scheduler split complete and gated. `vllm_v1_engine_replica_scheduler.py` went from 5138 lines to eight modules, largest 1386. Cleanup removed `_attach_afd_metadata_if_needed`, 65 lines with no caller that duplicated `scheduler/utils/afd_metadata.py`. Gates: CLI flag set identical; method resolution verified for every private method the four subclasses override, and for the two the base class also defines; unit selection of 73 files across `tests/unit` and `tests/integration` gives 13 failed / 1570 passed / 19 skipped / 5 errors on both sides with identical failure identities; fidelity matrix 67 identical, 0 mismatched, 0 predictor cache name differences. +- 2026-09-21: Four regressions were introduced by that split and fixed; see `issues.md` I5 to I8. Three were names the moved or retained code loads at runtime without a runtime binding, which is now checked mechanically: a script parses each split module and reports names loaded but neither imported, defined locally, nor builtin, counting `TYPE_CHECKING` imports as unbound. Its only remaining hit is a string annotation the flat CLI generator resolves through its own lazy-import special case. +- 2026-09-21: Independent confirmation of commit `99922d2` from a second session, measured in an isolated detached worktree at that commit rather than through the shared tree: 67/67 identical, 0 mismatched, 0 cache name differences. +- 2026-09-21: Prediction model manager split complete. 4614 lines to six modules, largest 1700. Cleanup removed three definition-only members. Unit selection widened to the predictor tests, 2300-odd tests over both sides: identical results, including identical failure identities. Eight tests needed their patch target moved with the code they exercise; see `issues.md` I9. +- 2026-09-21: MoE predictor split complete. 3539 lines to six modules, largest 1557. Cleanup removed `_is_grouped_gemm_on_demand_mode`, unreferenced. Unit selection widened again to 183 files: identical on both sides at 28 failed / 3016 passed / 55 skipped / 12 errors. Three tests needed the fake operator family installed in a second module; see `issues.md` I10. The first attempt at this selection aborted at collection and was discarded; see I11. +- 2026-09-21: Independent measurement of `3b00f16`, the manager split, from a detached checkout at that commit: 67/67 identical, 0 mismatched, and a clean predictor cache report with 0 rekeyed models. That is the direct evidence that moving the model hash, the cache keys, the registry and the FFN contract signature into different modules changed no training identity. diff --git a/task_memory/task_2026-09-21_oversized_module_split/requirements.md b/task_memory/task_2026-09-21_oversized_module_split/requirements.md new file mode 100644 index 00000000..04c862aa --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/requirements.md @@ -0,0 +1,24 @@ +# Oversized Module Split — Requirements + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Created from the Issue 26 planning interview. | + +## [Original Request] 2026-09-21 + +This task was created as the prerequisite of the Issue 26 correctness PR (`task_memory/task_2026-09-21_issue26_correctness_pr/`). During the planning interview the user decided: + +- Q4=c: apply the `AGENTS.md` 2,000-line gate as a full cleanup-first pass and functional split, not as a documentation-only analysis. +- Q8=a: scope is limited to the four modules the correctness PR must touch: `frontier/config/config.py` (5720 lines), `frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py` (5138), `frontier/execution_time_predictor/shared_prediction_model_manager.py` (4614), `frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py` (3539). The other five modules above 2,000 lines are out of scope. +- Q9=b, Q10=b: the refactor is a separate PR from branch `refactor/oversized-module-split` (base `origin/main` `1f694f7c549aa3aeeb7c5bbae04e119c09167a77`). The correctness branch is based on it. Both draft PRs are opened; the correctness PR's base is this branch until this PR merges. +- Q11: acceptance is a fidelity matrix of at least 50 scenarios. For each scenario the refactor branch and main must produce value-identical `request_metrics.csv` and identical `system_metrics.json` after removing timestamps and run ids. Any difference is FAIL unless it is an explicitly approved fidelity fix. Coverage: dense/MoE, offline/online, co-location / sequential PDD / sequential PD-AF, dummy predictor and checked-in profiling CSVs, varied request lengths, counts, and QPS. +- Q5: pushing this branch to `origin` and creating/updating its draft PR are authorized. Merge, force-push, history rewrite are not. +- Q7=b: task records are published with the PR through a narrow `.gitignore` exception for this directory. + +## Constraints carried from `AGENTS.md` + +- Cleanup first: remove redundant, dead, or overly defensive code before splitting; document the reason for anything that remains above 2,000 lines. +- Record proposed boundaries and sequencing before implementing a split. +- Plain ML-system names for new modules; no behavior change, no numeric change, no public import surface break (`frontier/config/__init__.py` star re-export must keep working). diff --git a/task_memory/task_2026-09-21_oversized_module_split/summary.md b/task_memory/task_2026-09-21_oversized_module_split/summary.md new file mode 100644 index 00000000..9baadb47 --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/summary.md @@ -0,0 +1,84 @@ +# Oversized Module Split — Summary + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Placeholder created at Step 0. | +| 2026-09-21 | Completion archive written after the fourth and last module split. | +| 2026-09-22 | Synchronized with the maintainer review: case table is 71, the final acceptance record is the clean recapture in `test_report_2026-09-22_checkpoint_b_final_evidence.md`, and the fidelity gate itself was corrected first. | +| 2026-09-22 | External review C34-01: cache-name comparison eligibility now requires the executed case list to equal the full table; Checkpoint B verdict re-derived and unchanged. | + +## Overview + +Four Python modules that the Issue 26 correctness fixes must edit were far above the 2,000-line maintainability gate in `AGENTS.md`. Editing them in the same pull request as the behavior fixes would have made that diff unreadable, because a reviewer could not tell a moved line from a changed one. This branch brings all four under the gate through a cleanup-first pass and a functional split, with no behavior change and no numeric change, and the correctness branch is stacked on top of it. + +| Module | Before | After (largest child) | Modules | +| --- | --- | --- | --- | +| `frontier/config/config.py` | 5720 | 1888 | 12 | +| `frontier/scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py` | 5138 | 1386 | 8 | +| `frontier/execution_time_predictor/shared_prediction_model_manager.py` | 4614 | 1700 | 6 | +| `frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py` | 3539 | 1557 | 6 | + +Five modules remain above 2,000 lines. All five were out of scope from the start and are named in `plan.md` section 1: `sklearn_execution_time_predictor.py` (8263), `metrics_store.py` (5586), `sklearn_disaggregation_execution_time_predictor.py` (2985), `profiling/attention/main.py` (2157), `entities/request.py` (2125). + +## Deliverables + +| Path | Contents | +| --- | --- | +| `frontier/config/` | 12 modules: the simulation config and its cluster topology, the cluster role builders and topology summary, and one module per configuration family | +| `frontier/scheduler/replica_scheduler/vllm_v1_*.py` | 8 modules: the scheduler plus MTP wait policy, role schedules, KV allocation, iteration policy, prefix cache, decode-attention cohort and the decision log | +| `frontier/execution_time_predictor/` | 12 modules across the two predictors: family trainers, dataframe loaders, model registry, layer contract resolution, identity helpers, and the MoE operator times, routing workload, dataset training, MTP replay and helpers | +| `tests/e2e/refactor_fidelity/` | The 71-case fidelity matrix that gates the branch, and the gate corrections from review comments R34-01 and R34-02 | +| `tests/unit/test_refactor_fidelity_gate.py` | 22 tests pinning that the gate cannot pass without successful comparisons, and that one label describes one measurement | +| `tests/unit/test_module_split_boundaries.py` | 13 tests for the annotation, re-export, MRO and cache-loading boundaries the split could break | +| `task_memory/task_2026-09-21_oversized_module_split/` | `plan.md`, `module_survey.md`, `issues.md`, `progress.md`, the Step 0 baseline report, the Step 1 matrix report, and the two Checkpoint reports of 2026-09-22 | + +## How the no-change claim is supported + +**The fidelity matrix.** 71 cases, each run through a checked-in example wrapper on both the branch and a read-only checkout of the base commit. Every artifact a run writes is compared and the file sets must match: `request_metrics.csv`, `system_metrics.json`, `frontier_stage_batch_ledger.jsonl`, `op_precision_metadata.csv` and `config.json`. The gate is exact equality with no tolerance, because a behavior-preserving refactor has no reason to change a simulated number. Only three run-specific absolute paths are normalized, and that list was checked against the real artifacts rather than assumed. + +Coverage: co-location, sequential PDD and sequential PD-AF; dense and MoE; offline and online; the dummy predictor and the checked-in profiling CSVs; request counts from 4 to 64; prompts from 128 to 3584 tokens; three arrival rates; chunked prefill, all three decode CUDA graph modes, prefix caching, speculative decoding and thinking mode; TP, PP, attention DP and EP variations. + +**The predictor cache names.** Retraining from the same CSV reproduces the same numbers, so a changed training identity or cache key would leave no trace in the outputs. The comparison therefore also checks the names of the predictor cache files each side produces, and classifies any difference as rekeyed, baseline-only or candidate-only. This is the specific evidence that moving the model hash, the cache keys, the registry and the FFN contract signature into different modules changed no training identity. A side's cache listing enters that comparison only when its manifest shows one clean, unfiltered run whose `cases_executed_in_last_run` equals the full case table; `--start`/`--limit` continuations are excluded (corrected 2026-09-22 after external review finding C34-01, verdict unchanged). + +**The unit suites.** A selection that grows with each step, ending with the whole `tests/unit` directory, run on both sides with failure identities compared rather than counts. Identities matter: a count comparison would have hidden several of the defects below, because the same files also carry pre-existing failures. + +**The gate itself.** The three claims above are only worth as much as the comparator that produces them. The maintainer review found that the comparator could return success having compared nothing, so it was corrected and given its own regression tests before the acceptance run was taken. The four false successes are reproduced against the pre-fix harness in `test_report_2026-09-22_checkpoint_a_fidelity_gate.md`. + +**The CLI surface.** The generated flag set, 753 flags, compared after every step, plus a check that all 36 names other modules import from `frontier.config` still resolve from both import paths. + +**A static check.** Every split module is parsed and any name loaded at runtime without a runtime binding is reported, counting `TYPE_CHECKING` imports as unbound. It found one defect the matrix had not yet reached. + +## Observed validation results + +| Step | Module | Unit selection | Fidelity matrix | +| --- | --- | --- | --- | +| 3 | `config.py` | 62 files, 10 failed / 671 passed, identical identities | 67 identical, 0 mismatched, 0 cache differences | +| 4 | vLLM V1 replica scheduler | 73 files, 13 failed / 1570 passed, identical identities | 67 identical, 0 mismatched, 0 cache differences | +| 5 | prediction model manager | 2300-odd tests, identical | 67 identical, 0 mismatched, 0 rekeyed models | +| 6 | MoE execution-time predictor | 183 files, 28 failed / 3016 passed, identical identities | 67 identical, 0 mismatched | +| **Final** | whole branch at `bb582a4` | whole `tests/unit`: 84 failed / 3679 passed / 49 skipped, failure identities byte-identical to `1f694f7` (84 / 3644 / 49) | **71 compared, 71 identical, 0 mismatched, 0 failures, 0 missing evidence, 0 cache differences** | + +The final row is the acceptance record, taken after the gate corrections landed, with both sides run as single clean full captures from detached checkouts. Its provenance, the excluded modules and its limits are in `test_report_2026-09-22_checkpoint_b_final_evidence.md`. The per-step rows above it were taken with the pre-correction gate and are retained as history, not as the acceptance evidence. + +The config and manager splits were additionally measured from detached checkouts at their own commits by a second session, so the tree measured was the commit and nothing else. + +Pre-existing failures fall into two buckets and are not caused by this branch: ten in `test_colocation_release_review_contracts.py`, which open `tests/debug/` scripts that are not tracked on `main` at all, and several that raise `ModuleNotFoundError` on `torch`, which the minimal CPU environment deliberately excludes. + +## What a reviewer should look at + +The eleven entries in `issues.md` are the substance of the review. Eight are defects this branch introduced and fixed, and each records what found it. Three of those were names the moved or retained code loads at runtime without a runtime binding, which is now checked mechanically rather than by inspection. + +The test changes divide into three kinds, and only the second is a judgment call: + +1. A patch target moved with the code it patches. A test that patches a module-level name stops intercepting once the method reading it resolves that name in another module. Seven test files, no assertion changed. +2. One governance allowlist entry. `test_raw_model_profile_resolution_callsites_are_allowlisted` pins which functions may resolve a raw model architecture profile and how many times each may do so. Exactly one line changed, a file path; the function name, the kind, the expected count of 1 and the total of 11 entries are unchanged, so the property the gate protects is intact. +3. Three tests gained a line. `MOE_FAMILY` was one binding and is now imported by two modules that the path under test both read, so a fake operator family has to be installed in both. The assertions are unchanged, and the second patch is needed precisely because the production code genuinely reads the name in both places. + +## Open and deferred work + +- None outstanding for the matrix: the final measurement is the clean 71-case recapture of `bb582a4` recorded in the Checkpoint B report. +- The naming of the extracted mixin classes is a review point, not a settled decision. Renaming any of them is mechanical and affects three files at most. +- Two modules keep a name that is reported by the static check and is correct: `BaseCCBackendConfig` in `frontier/config/cluster_config.py` is a string annotation the flat CLI generator resolves through its own lazy-import special case, exactly as the pre-split `config.py` did. `typing.get_type_hints(ClusterConfig)` therefore raises on this branch — and equally on `1f694f7`, which was verified. It is pinned in `KNOWN_UNRESOLVED_CONFIG_ANNOTATIONS` so a new occurrence fails. +- `AGENTS.md` still points readers at `tests/unit/test_open_source_release_arch_guard.py` and two `tests/debug/` scripts that do not exist on `main`. Recorded as deferred; fixing it is not in this branch's scope. diff --git a/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_config_split_isolated_validation.md b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_config_split_isolated_validation.md new file mode 100644 index 00000000..3e6a409f --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_config_split_isolated_validation.md @@ -0,0 +1,53 @@ +# Test Report 2026-09-21 — Isolated Validation of the Config Split (`99922d2`) + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Initial report. | + +## Why this run exists + +The branch worktree was being edited while the config split's own matrix run was in flight, so that run measured a mixture of the committed config split and an in-progress scheduler split and had to be discarded. This report records a run that cannot have that problem: the simulator under test is a **read-only detached checkout of the commit itself**, not the shared working tree. + +It is an independent confirmation of `99922d2`, not a replacement for the gate that commit already passed. + +## Method + +| Field | Value | +| --- | --- | +| Commit under test | `99922d26365e5639bfdbcc2fe59e94d0750a6d27` (`refactor(config): split config.py into one module per configuration family`) | +| Checkout under test | `/data/ycfeng/Frontier/.worktrees/fidelity-candidate-99922d2`, detached at that commit, no source edits | +| Only local change | `tests/e2e/refactor_fidelity/run_matrix.py`, the harness hardening described below. No file under `frontier/` differs from the commit. | +| Baseline | `/data/ycfeng/tmp/issue26-correctness-pr/refactor-fidelity/baseline`, captured earlier from `.worktrees/fidelity-baseline-main` at `1f694f7c549aa3aeeb7c5bbae04e119c09167a77` | +| Python | `/data/ycfeng/envs/frontier-py310/bin/python`, CPython 3.10.6 | +| Command | `run_matrix.py run --repo-root --label candidate_99922d2 --output-root --jobs 4 --clean-cache --continue-on-failure`, then `compare --baseline-label baseline --candidate-label candidate_99922d2` | + +## Result + +| Measure | Expected | Actual | +| --- | --- | --- | +| Cases producing artifacts | 67 of 67 | 67 of 67 | +| Identical cases | 67 | **67** | +| Mismatched cases | 0 | **0** | +| Baseline failures excluded | 0 | 0 | +| Candidate-only failures | 0 | **0** | +| Predictor cache file names differing | 0 | **0** | + +**Result: PASS.** Every compared artifact matched exactly: `request_metrics.csv`, `system_metrics.json`, `frontier_stage_batch_ledger.jsonl`, `op_precision_metadata.csv` and `config.json`, across all 67 cases, with only run-specific absolute paths normalized. + +The predictor cache result is worth stating separately. All 426 cache file names matched, which means the split changed no training identity and no model cache key. Output equality alone would not have shown that, because retraining from the same CSV reproduces the same numbers. + +## Harness hardening applied for this run + +Two changes to `tests/e2e/refactor_fidelity/run_matrix.py`, both about diagnosability rather than about what counts as a pass: + +1. **Bounded retry on artifact discovery.** A case that exits zero but appears to have written nothing is retried for up to ten seconds before being recorded as a failure. A directory listing can lag the child process, and a spurious failure is worse than waiting. A genuinely empty run still fails, only later. +2. **Failure log capture.** The last twenty lines of a failing case's log are stored in `results.jsonl` and printed by `compare`. Without this the evidence is lost as soon as the case is re-run, because a re-run overwrites `run.log`. + +Both were motivated by a concrete incident: two cases were reported as failures in an earlier run, passed immediately when re-run individually, and by then their logs had been overwritten, so the cause could not be established from the record. + +## Limits + +- This confirms that 67 supported configurations produce identical output at `99922d2`. It says nothing about configurations outside the matrix, and it makes no accuracy claim. +- The isolated checkout remains on disk at `.worktrees/fidelity-candidate-99922d2` as a pre-scheduler-split reference. It is detached and ignored by Git; delete it once the branch's later splits are validated. diff --git a/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_step0_baseline.md b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_step0_baseline.md new file mode 100644 index 00000000..120e0ec5 --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_step0_baseline.md @@ -0,0 +1,71 @@ +# Test Report 2026-09-21 — Step 0 Baseline (`1f694f7`) + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Initial report. | + +## Environment + +| Field | Value | +| --- | --- | +| Host | `kun-workspace-vgen2` (CPU only) | +| Worktree | `/data/ycfeng/Frontier/.worktrees/oversized-module-split` at `1f694f7c549aa3aeeb7c5bbae04e119c09167a77` (before any source change) | +| Python | `/data/ycfeng/envs/frontier-py310/bin/python`, CPython 3.10.6 (uv venv; `uv pip install -e ".[test]"`) | +| Packages | numpy 2.2.6, pandas 2.3.3, scikit-learn 1.7.2, scipy 1.15.3, plotly 7.1.0, pytest 9.1.1 | +| Import path | `frontier` resolved from the worktree (`frontier/main.py` path verified); `frontier/` is a namespace package (no `__init__.py`) | +| Env vars | `PYTHONPATH=`, `FRONTIER_TMP_ROOT=/data/ycfeng/tmp/issue26-correctness-pr`, `WANDB_DISABLED=true`, `VIDUR_DISABLE_WANDB=1` | +| Raw logs | `/data/ycfeng/tmp/issue26-correctness-pr/step0/{pytest_baseline.log,coloc_dense.log,pdd_dense.log}` (scratch, not committed) | + +## 1. Unit selection + +Command: + +```bash +python -m pytest tests/unit/test_cluster_scheduler_dp_lanes.py \ + tests/unit/test_colocation_release_review_contracts.py \ + tests/unit/test_config_owned_contracts.py \ + tests/unit/test_stage_execution_time.py \ + tests/unit/test_stage_finalized_contract.py \ + tests/unit/test_moe_routing_runtime.py -q -p no:cacheprovider +``` + +Expected: collection succeeds; failures, if any, are attributable to the unmodified base or the environment. + +Actual: `10 failed, 84 passed in 3.42s`, exit code 1. **Result: PASS as a baseline record** (all failures are pre-existing on main or environmental; none involve the modules in scope). + +| Failing test (`tests/unit/test_colocation_release_review_contracts.py`) | Cause (observed) | Class | +| --- | --- | --- | +| `test_config_optimizer_help_prints_without_argparse_percent_crash` | `FileNotFoundError: 'python'`: the test spawns the bare `python` executable, which is not on this host's PATH | environment | +| `test_debug_e2e_base_does_not_default_to_private_conda_path` | `tests/debug/e2e-level/monolith_mode/scripts/test_base.sh` does not exist on main (`git ls-files tests/debug` is empty) | pre-existing on main | +| `test_debug_e2e_conda_activation_temporarily_disables_nounset` | same missing `test_base.sh` | pre-existing on main | +| `test_debug_e2e_base_uses_safe_pythonpath_expansion` | same missing `test_base.sh` | pre-existing on main | +| `test_readme_debug_scripts_do_not_use_post_increment_under_set_e` | missing `test_dense_tp2_pp2_dummy.sh` | pre-existing on main | +| `test_debug_e2e_base_resolves_latest_canonical_metrics_run_dir` | missing `test_base.sh` (bash exit 127) | pre-existing on main | +| `test_release_debug_scripts_use_canonical_metrics_resolver` | missing `test_dense_tp2_pp2_dummy.sh` | pre-existing on main | +| `test_readme_moe_debug_script_uses_valid_shared_parallel_domain` | missing `test_moe_tp2_ep2_pp2_dummy.sh` | pre-existing on main | +| `test_readme_moe_debug_script_header_uses_current_tp_terms` | missing `test_moe_tp2_ep2_pp2_dummy.sh` | pre-existing on main | +| `test_readme_moe_debug_script_satisfies_shared_parallel_domain` | missing `test_moe_tp2_ep2_pp2_dummy.sh` | pre-existing on main | + +Observation: `AGENTS.md` (§Tests) still tells readers to run the two missing debug scripts and `tests/unit/test_open_source_release_arch_guard.py`, which is also absent. Recorded as documentation drift; not fixed by either PR. + +Per-file result: `test_cluster_scheduler_dp_lanes.py`, `test_config_owned_contracts.py`, `test_stage_execution_time.py`, `test_stage_finalized_contract.py`, `test_moe_routing_runtime.py` all passed; the 10 failures are confined to `test_colocation_release_review_contracts.py`. + +## 2. Frontier-only CPU smokes (dummy predictor, analytical backend) + +| Script | Command | Expected | Actual | Result | +| --- | --- | --- | --- | --- | +| Co-location dense offline | `PYTHON_BIN=$PY METRICS_OUTPUT_DIR=$FRONTIER_TMP_ROOT/step0/coloc_dense bash examples/architecture/co-location/offline/dense_model_basic.sh` | exit 0, `request_metrics.csv` and `system_metrics.json` written | exit 0; `request_metrics.csv` 16 rows; `system_metrics.json` sections `simulation_metadata, quantization_config, ttft_statistics, tpot_statistics, request_e2e_time_statistics, throughput_metrics, spec_decode_statistics, preemption_statistics, ...` | PASS | +| PDD dense offline (sequential) | `PYTHON_BIN=$PY METRICS_OUTPUT_DIR=$FRONTIER_TMP_ROOT/step0/pdd_dense bash examples/architecture/pdd/offline/dense_model_basic.sh` | exit 0, same artifacts | exit 0; `request_metrics.csv` 8 rows; sections include `kv_cache_transfer_statistics` | PASS | + +Outputs were redirected to the scratch root; the worktree shows no generated files (`git status` lists only `.gitignore` and `task_memory/`). + +## 3. Reference checkout + +`fwyc0573/vLLM-BS` cloned to `/data/ycfeng/Frontier/.real-engine/vLLM-BS` (`.git/info/exclude`), detached at `ea95f571e20937c7c908c6d59ddd1cd6bf9268f1` ("Instrument vLLM attention ops for Frontier calibration"), contained in `origin/feature/frontier-comparison-instrumentation`. The fork exposes **no tags**, so the relationship to upstream `v0.10.2` could not be established from tags; Step 1 of the correctness task must add the upstream remote read-only and compare against the resolved `v0.10.2` commit, as the specification requires. + +## Limits + +- Dummy-mode smokes verify lifecycle and artifact production only, not numerical parity. +- The unit selection is a baseline record; the 10 failures were not reproduced with a different Python to separate PATH effects further, because their assertion messages already name the missing files or executable. diff --git a/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_step1_fidelity_matrix.md b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_step1_fidelity_matrix.md new file mode 100644 index 00000000..6ca08f66 --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-21_step1_fidelity_matrix.md @@ -0,0 +1,160 @@ +# Test Report 2026-09-21 — Step 1 Fidelity Matrix and Step 2 Cleanup + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-21 | Initial report: harness self-check, baseline capture, and the first cleanup comparison. | +| 2026-09-21 | Added section 6, the `config.py` split. | + +## Environment + +| Field | Value | +| --- | --- | +| Host | `kun-workspace-vgen2`, CPU only | +| Python | `/data/ycfeng/envs/frontier-py310/bin/python`, CPython 3.10.6 | +| Packages | numpy 2.2.6, pandas 2.3.3, scikit-learn 1.7.2, scipy 1.15.3, plotly 7.1.0, pytest 9.1.1 | +| Baseline checkout | `/data/ycfeng/Frontier/.worktrees/fidelity-baseline-main`, detached at `1f694f7c549aa3aeeb7c5bbae04e119c09167a77`, never edited | +| Candidate checkout | `/data/ycfeng/Frontier/.worktrees/oversized-module-split` | +| Output root | `/data/ycfeng/tmp/issue26-correctness-pr/refactor-fidelity` (scratch, not committed) | +| Harness | `tests/e2e/refactor_fidelity/` (`cases.py`, `compare.py`, `run_matrix.py`) | + +## 1. What the matrix compares + +Each case runs a checked-in example wrapper with the working directory set to the checkout under test, which is what makes the shipped configuration's relative profiling-data and `cache` paths resolve. The driver always comes from the candidate checkout, so one case table and one comparator measure both sides. + +Every file the run writes into the normalized metrics directory is compared, and the file sets must match: + +| Artifact | Comparison | +| --- | --- | +| `request_metrics.csv` | Every cell, exactly. First differing cell is reported with its column name. | +| `system_metrics.json` | Every leaf, exactly. | +| `frontier_stage_batch_ledger.jsonl` | Record count, then every leaf of every record. | +| `op_precision_metadata.csv` | Every cell, exactly. | +| `config.json` | Every leaf, exactly, after path substitution. | + +The only normalization is the substitution of three run-specific absolute paths (the checkout root, the output root, the label root). That list is short because it was checked against the real artifacts: on `1f694f7` the metrics files carry no timestamps, wall-clock durations, hostnames or paths, and `config.json` carries exactly one absolute path, the metrics output directory. Anything the substitution list fails to cover shows up as a difference rather than being silently accepted. + +There is no tolerance. A behavior-preserving refactor has no reason to change a simulated number. + +In addition, the names of the predictor cache files each side produces are compared. Retraining from the same CSV reproduces the same numbers, so a changed training identity or cache key would otherwise leave no trace in the outputs. + +## 2. Case coverage + +67 cases, all driven through checked-in wrappers: + +| Group | Cases | What it covers | +| --- | --- | --- | +| `trained_predictor` | 6 | Dummy mode disabled; checked-in CSVs for dense and MoE on two devices; the kernel-only CUDA graph path | +| `colocation_offline_dense` | 13 | Request counts 4/8/16/64, prompts 128/512/1024/3584, TP1/TP2, PP1/PP2, one and two replicas, attention DP2, all three decode CUDA graph modes, chunked prefill on and off, the Sarathi and SGLang schedulers, two dummy latencies | +| `colocation_offline_moe` | 11 | EP1/EP2/EP4, MoE TP1/TP2, all four routing distributions, top-k 2 and 4, long prefill, 32 requests, CUDA graph off | +| `colocation_online` | 6 | Dense and MoE at 0.5, 2 and 8 requests per second | +| `colocation_features` | 7 | Thinking mode, speculative decoding (two token counts), prefix caching, each offline and online | +| `pd_disaggregation` | 14 | Dense and MoE, replica counts per role, TP2, chunked prefill off, EP1, skewed routing, thinking mode, speculative decoding, prefix caching, two online arrival rates | +| `pd_af_disaggregation` | 10 | Dense and MoE, EP1 and EP2, the global CUDA graph contract, offline and online | + +## 3. Harness self-check + +Running the same 13 co-location dense cases twice from the **same** checkout, in two separate processes, must produce identical artifacts; otherwise the comparator cannot distinguish a refactor regression from ordinary run-to-run noise. + +| Field | Value | +| --- | --- | +| Command | `run_matrix.py run --repo-root --label selfcheck_a|selfcheck_b --case-filter coloc_dense_offline` then `compare` | +| Expected | 13 identical, 0 mismatched | +| Actual | identical 13, mismatched 0, cache file names identical | +| Result | **PASS** | + +This covers the 343 KB stage ledger and `config.json` as well as the metrics files. + +## 4. Baseline capture + +| Field | Value | +| --- | --- | +| Command | `run_matrix.py run --repo-root --label baseline --output-root --jobs 6 --clean-cache` | +| Expected | Every case completes and writes a metrics directory | +| Actual (first attempt) | 65/67. Two cases failed. | +| Actual (after case corrections) | **67/67**, zero failures | +| Result | **PASS** | + +Both first-attempt failures were invalid case definitions, not simulator defects, and were corrected in the case table before any source change: + +| Case | Diagnosis | Correction | +| --- | --- | --- | +| `coloc_dense_offline_long_prefill` | 4096 prefill tokens plus 16 decode tokens exceeds the 4096-token context of Llama-2-7b, so no request is ever admitted and the run ends with a non-empty scheduler state | prefill reduced to 3584 tokens | +| `coloc_moe_offline_ep1`, `coloc_moe_offline_attn_dp2` | Violated the shared-domain invariant `attn_tp * attn_dp == moe_tp * moe_ep`; the MoE wrapper additionally enforces the stricter `ATTN_TP == MOE_TP * MOE_EP`, which cannot express `attn_dp > 1` | EP1 case set to `ATTN_TP=1`; the MoE DP case was replaced by an EP4 topology, since the dense matrix already covers DP lanes | + +## 5. Cleanup comparison (Step 2, first part) + +Source changes under test, in `frontier/config/config.py` only: + +1. Removed `BaseExecutionTimePredictorConfig.validate_linear_op_input`, which had no caller anywhere in `frontier/`, `tests/`, `examples/` or `docs/` (re-verified by grep immediately before the edit). +2. Replaced the two parallel CC-backend dispatch ladders with one ordered `(type key, config class, creator)` table, preserving the original evaluation order. The five backend config classes are siblings of `BaseCCBackendConfig` with no cross-inheritance, so the order is not load-bearing either way. +3. Replaced four copies of a field-resolution closure with one `_cc_backend_value_reader`, and three identical `hasattr`-guarded field triplets with one `_shared_cc_backend_fields`. `hasattr(x, "a") and x.a or default` and `getattr(x, "a", default)` are equivalent, so this is a textual deduplication. +4. Removed a method-local re-import of `dataclasses.replace`, which the module already imports. + +Net effect: 5720 to 5687 lines. + +| Check | Expected | Actual | Result | +| --- | --- | --- | --- | +| Generated CLI flag set | Identical to the baseline | 753 flags, `diff` empty | PASS | +| Fidelity matrix | 67 identical, 0 mismatched | 67 identical, 0 mismatched, 0 candidate-only failures | PASS | +| Predictor cache file names | Identical | 138 of 426 names differed; see below | Investigated, harness cause, corrected | + +### The cache-name difference was a harness artifact + +The differing names were the `.pkl`, `_predictions.csv` and lock files of the two `examples/profiling/smoke_simulator_*_csv.sh` cases. Cause: `ExecutionTimePredictionModelManager._get_hash_relevant_config` includes the profiling input file paths in the model hash, and those wrappers default `DATA_DIR_BASE` to an absolute path under their own repository root, which differs between the two checkouts by construction. The architecture wrappers were unaffected because they leave the relative path templates in place. + +This is not a training-identity change: the simulated outputs of both smoke cases were already identical, and the CSV contents are the same file in both checkouts. + +Correction: both smoke cases now pass `DATA_DIR_BASE=data/profiling`, a repository-relative base that resolves identically from either checkout, so the cache-name check becomes meaningful again. Both sides were recaptured with a cleaned cache after the correction. + +## Limits + +- The matrix establishes that the outputs of 67 supported configurations do not change. It does not prove that no unsupported or unexercised configuration changes, and it makes no accuracy claim. +- Dummy-mode cases exercise structure and lifecycle, not realistic latency. Six cases use the checked-in profiling CSVs, which is what covers dataset loading, training identity and the persistent cache. +- The predictor cache comparison compares file names, not pickle contents, because pickle bytes are not required to be reproducible. + +## 6. `config.py` split (Step 3) + +`config.py` went from one 5,687-line module to twelve modules, the largest 1,888 lines. The full boundary table is in `plan.md` section 3.1; `issues.md` I1 records the one real defect this step introduced and how it was found. + +### 6.1 Compatibility gates + +| Check | Expected | Actual | Result | +| --- | --- | --- | --- | +| Generated CLI flag set | Identical to the base commit | 753 flags, `diff` empty | PASS | +| Public import surface | All 36 names other modules import resolve from both `frontier.config` and `frontier.config.config` | none missing from either | PASS | +| `ClusterConfig` dataclass fields | Unchanged | 181 fields; the two mixins contribute none, as intended | PASS | +| `ClusterConfig` MRO | `ClusterConfig`, `ClusterRoleConfigBuilder`, `ClusterTopologySummary`, `object` | as expected | PASS | + +### 6.2 Unit suites + +Selection: every file under `tests/unit/` that mentions `ClusterConfig`, `frontier.config.config` or `from frontier.config import`, 62 files. + +```bash +python -m pytest $(grep -rln --include="*.py" \ + "ClusterConfig\|frontier.config.config\|from frontier.config import" tests/unit | sort) \ + -q -p no:cacheprovider --tb=no +``` + +| Side | Result | +| --- | --- | +| Base commit `1f694f7` | 10 failed, 671 passed | +| This branch | 10 failed, 671 passed | + +The failing test identities are **byte-identical** between the two sides: all ten are in `tests/unit/test_colocation_release_review_contracts.py` and are the pre-existing drift recorded in the Step 0 baseline report (the `tests/debug/` scripts they open are not tracked on `main`, and one test spawns a bare `python` executable absent from this host's PATH). + +The first run of this selection found **7 additional failures** that the fidelity matrix would not have caught at that point. They are recorded as I1 in `issues.md`: `get_cluster_configs_for_disaggregation` constructs `ClusterConfig` objects at runtime, and the extraction had imported the name only under `TYPE_CHECKING`. Fixed with a lazy import inside the method, matching the package's existing `_get_cc_backend_configs` pattern. + +### 6.3 Fidelity matrix + +Both sides recaptured with a cleaned predictor cache. + +| Check | Expected | Actual | Result | +| --- | --- | --- | --- | +| Cases producing artifacts | 67/67 on both sides | 67/67 | PASS | +| Cases identical | 67 identical, 0 mismatched | identical 67, mismatched 0 | PASS | +| Candidate-only failures | none | 0 | PASS | +| Predictor cache file names | Identical | 0 baseline-only, 0 candidate-only | PASS | + +Compared revisions: baseline `1f694f7c549a`, candidate `56efdf765ee0` plus the working tree of this step. diff --git a/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_cache_eligibility_correction.md b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_cache_eligibility_correction.md new file mode 100644 index 00000000..44c68875 --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_cache_eligibility_correction.md @@ -0,0 +1,54 @@ +# Test Report 2026-09-22 — C34-01: predictor-cache comparison eligibility + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-22 | Created: defect, correction, gate tests, re-derivation of the Checkpoint B verdict. | + +## Finding (external review of PR #34 / PR #35, 2026-09-22, C34-01, P1) + +`compare_labels` in `tests/e2e/refactor_fidelity/run_matrix.py` compared the two +sides' predictor-cache file names when both result sets covered the full case +table and both manifests had `cache_clean_before_run: true` with an empty +`case_filter`. `_select_cases` also honors `--start` and `--limit`, which narrow +the executed selection without writing a filter. A run that cleaned the cache, +executed a slice, and merged its records onto full retained results therefore +carried a partially populated cache while satisfying the predicate, and a +matching partial cache on the other side would have compared equal for the +wrong reason. The manifest already records `cases_executed_in_last_run`; the +predicate did not read it. + +## Correction + +| File | Change | +| --- | --- | +| `tests/e2e/refactor_fidelity/run_matrix.py` | `populated_by_one_clean_full_run(manifest)` requires `cache_clean_before_run`, no `case_filter`, and `set(cases_executed_in_last_run) == full case table`; a manifest without the field is ineligible. The console line names the reason. Report keys unchanged (`predictor_cache_populated_cleanly`, `predictor_cache_compared`). | +| `tests/unit/test_refactor_fidelity_gate.py` | `_write_side` gains `cases_executed` and `record_cases_executed`; the healthy case asserts both cache keys true; four new tests: `--limit`-narrowed, `--start`-narrowed, symmetric partial caches, manifest without the executed list. | + +No production module changed. The manifest schema is unchanged. + +## Verification + +Environment: `/data/ycfeng/envs/frontier-py310/bin/python` (Python 3.10), +`PYTHONPATH` at the worktree root, run from +`/data/ycfeng/Frontier/.worktrees/oversized-module-split`. + +| # | Check | Command | Expected | Actual | Result | +| --- | --- | --- | --- | --- | --- | +| 1 | Gate unit tests | `python -m pytest tests/unit/test_refactor_fidelity_gate.py -q -p no:cacheprovider` | all pass, including the four new cases | 26 passed in 1.05 s | PASS | +| 2 | New cases fail on the old predicate | same file imported against the unmodified runner (observed when the tests were first run with the PR35 copy of `run_matrix.py` on `sys.path`) | the four new cases fail, the 22 existing pass | 4 failed, 22 passed | PASS (negative control) | +| 3 | Checkpoint B verdict re-derived | `python tests/e2e/refactor_fidelity/run_matrix.py compare --output-root /data/ycfeng/tmp/issue26-correctness-pr/refactor-fidelity --baseline-label baseline_v2 --candidate-label candidate_bb582a4` | same verdict as the Checkpoint B report | `cases compared: 71 of 71`, `identical: 71`, `mismatched: 0`, `predictor cache file names differ: 0 baseline-only, 0 candidate-only`, exit 0; `comparison.json`: `predictor_cache_populated_cleanly: true`, `predictor_cache_compared: true`, `complete_comparison: true` | PASS | + +Both Checkpoint B manifests record `cache_clean_before_run: true`, `case_filter: +null`, 71 entries in `cases_executed_in_last_run` and 426 cache files, so the +corrected rule admits them; the previous `comparison.json` was saved before the +rerun and the rerun rewrote it with the same verdict. + +## Limits + +- Observation, not inference: the retained output root is scratch and not + committed; the rerun's `comparison.json` lives there. +- The rule reads the manifest the runner writes. A manifest edited by hand can + still claim a full run; the gate protects against the harness's own `--start` + and `--limit` paths, which is the reviewed gap. diff --git a/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_checkpoint_a_fidelity_gate.md b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_checkpoint_a_fidelity_gate.md new file mode 100644 index 00000000..8658423f --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_checkpoint_a_fidelity_gate.md @@ -0,0 +1,138 @@ +# Test Report 2026-09-22 — Checkpoint A: the fidelity gate's false successes + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-22 | Initial report: R34-01 and R34-02 corrections, with the before/after reproduction. | + +## What this closes + +Maintainer review comments **R34-01** (the gate can report success without +successful comparisons) and **R34-02** (partial reruns and reused worktrees can +mix source identities), from +`.local-draft/Frontier_PR34_PR35_Review_and_D1_D2_Decisions_2026-09-22.md`. + +No file under `frontier/` was touched. The change is confined to +`tests/e2e/refactor_fidelity/` plus one new test module. + +## Environment + +| Field | Value | +| --- | --- | +| Host | `kun-workspace-vgen2`, CPU only | +| Python | `/data/ycfeng/envs/frontier-py310/bin/python`, CPython 3.10.6 | +| Worktree | `/data/ycfeng/Frontier/.worktrees/oversized-module-split` | +| Pre-fix harness | extracted from `5ef96b5` with `git show` into a scratch directory, so no worktree was created or deleted for the comparison | + +## 1. The defects, as found in source + +### R34-01, first path: a full result table of failed runs + +`compare_labels` skipped a case whose baseline execution failed, appending it to +`baseline_failures` — and `baseline_failures` was absent from the `failed` +predicate (`run_matrix.py:535` before the change). `complete` (`:437`) tested +only whether both result tables contained every case id, which a table of +failure records satisfies. A matrix in which every case failed on both sides +therefore compared nothing and returned 0, which `measure_commit.py:139` prints +as `VERDICT: IDENTICAL`. + +The same branch hid a *new* candidate failure whenever the baseline had failed +on the same case. + +### R34-01, second path: deleted artifact directories + +`list_artifacts` returns `[]` for a path that does not exist +(`compare.py:83-84`), so `compare_artifact_directories` found no differences +between two absent directories and recorded the case as identical. + +### R34-02: one label could describe several measurements + +`run_label` merged previously recorded cases into `results.jsonl` and then +rewrote the label-wide manifest with the current run's `git_head`, environment +and cache listing. Nothing checked that the retained records came from the same +source. `measure_commit.py` reused an existing checkout after comparing `HEAD` +only; a detached worktree is not read-only. + +## 2. Corrections + +| Area | Change | +| --- | --- | +| Completeness | Completeness now means *compared*, not *present in the table*: `compared_ids` is built from the cases that reached a content comparison, and `not_compared = full_case_set - compared_ids`. | +| Failure predicate | `baseline_failures`, `missing_evidence`, `definition_mismatches`, `provenance_findings` and `unexplained` were added. `--allow-partial` now waives exactly one thing: cases nobody attempted on either side. | +| Structural guard | `unexplained` lists any case that was neither compared nor reported under a specific finding, so a future code path cannot drop a case silently and still pass. | +| Missing evidence | `missing_evidence_for` checks that a successful record still names an artifact directory, that the directory exists, and that the files on disk match the recorded inventory. `compare_artifact_directories` reports an absent directory instead of treating it as empty. | +| Case identity | `FidelityCase.definition_digest()` digests script, env, extra args and the serial/parallel flag. A case id whose digest differs between the two sides is not compared. | +| Source provenance | Every case record carries `source_revision`, `source_dirty` and `harness_revision`. `check_retained_records` refuses a continuation whose retained records disagree, **before** running anything, and reports retained ids that have left the case table. | +| Manifest | `case_count` counts the lines actually written. An earlier label reported 72 over 71 result lines because a stale retained id was counted but not written. | +| Cache comparison | Cache file names are compared only when both sides ran the whole table **and** each side's manifest records a clean cache with no case filter. | +| Reused checkouts | `measure_commit.reuse_blocked_reason` refuses a reused checkout with tracked modifications or untracked files. It reports and stops; it does not clean the checkout. | + +## 3. Before/after reproduction + +Driver: `scratchpad/driver/negative_control.py`, run once with the pre-fix +harness on `PYTHONPATH` and once with the corrected harness. Identical synthetic +inputs in both runs; the driver lives outside both trees so the script directory +cannot shadow the import. + +| Scenario | Expected | Before (`5ef96b5`) | After | +| --- | --- | --- | --- | +| All 71 cases failed on both sides | fail | **exit 0**, compared 0 | exit 1, compared 0 | +| Baseline failure hiding a candidate failure | fail | **exit 0**, compared 70 | exit 1, compared 70 | +| Baseline failure, candidate succeeded | fail | **exit 0**, compared 70 | exit 1, compared 70 | +| Artifact directories deleted on both sides | fail | **exit 0**, compared 71 | exit 1, compared 70 | +| Healthy identical sides (control) | pass | exit 0, compared 71 | exit 0, compared 71 | +| Healthy content mismatch (control) | fail | exit 1, compared 71 | exit 1, compared 71 | + +Four false successes before; none after; both controls unchanged. The fourth row +is the clearest: before the fix the deleted case was *counted as compared and +identical*; after it is excluded as missing evidence, which is why `compared` +drops to 70. + +## 4. Committed regression tests + +`tests/unit/test_refactor_fidelity_gate.py`, 22 tests, all passing: + +```bash +PYTHONPATH=$PWD python -m pytest tests/unit/test_refactor_fidelity_gate.py -q -p no:cacheprovider +# 22 passed in 0.85s +``` + +They cover the six R34-01 scenarios the review listed, the five R34-02 +scenarios, and the three reused-checkout states. One asserts that the dirty +check leaves the modified file untouched, so the check cannot start "fixing" +the checkout to pass itself. + +Every test that mentions the harness: + +```bash +PYTHONPATH=$PWD python -m pytest $(grep -rln --include="*.py" refactor_fidelity tests/ | sort) \ + -q -p no:cacheprovider +# 30 passed in 25.77s +``` + +## 5. Live verification of the run side + +One real case, then a continuation, then a tampered continuation, against this +worktree (which was deliberately dirty at the time, and was recorded as such): + +| Step | Command | Result | +| --- | --- | --- | +| Run one case | `run --label smoke --case-filter coloc_dense_offline_small` | `ok`, 5 artifacts; record stamped `source_revision=5ef96b5…`, `source_dirty=True`, `harness_revision=5ef96b5…`, `case_digest=a1e285a457d62d46` | +| Continue, same source | `run --label smoke --case-filter coloc_dense_offline_default` | allowed; 2 records | +| Continue after rewriting the first record's `source_revision` to `000…` | same command | **refused, exit 2**, naming the conflicting case and both provenance tuples | + +The smoke output directory was removed afterwards. + +## Limits + +- This report is about the gate, not about the refactor. It does not re-establish + that the production tree is unchanged; that is Checkpoint B. +- The provenance stamp is new, so **every label captured before this change now + fails the gate for want of provenance**. That is intended: those labels cannot + be shown to describe one revision, and the baseline label in particular was an + assembled partial run. Both sides are recaptured in Checkpoint B. +- `missing_evidence_for` compares the recorded file *names* against disk, not + their digests. The content comparison that follows reads the same files, so a + changed file is caught there; a file replaced between recording and comparison + with one of the same name is not separately flagged. diff --git a/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_checkpoint_b_final_evidence.md b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_checkpoint_b_final_evidence.md new file mode 100644 index 00000000..e97d4cdb --- /dev/null +++ b/task_memory/task_2026-09-21_oversized_module_split/test_report_2026-09-22_checkpoint_b_final_evidence.md @@ -0,0 +1,210 @@ +# Test Report 2026-09-22 — Checkpoint B: the final acceptance record for PR #34 + +## Modification History + +| Date | Change | +| --- | --- | +| 2026-09-22 | Initial report: clean recapture of both matrix sides, unit-suite regression comparison, retained split checks. | +| 2026-09-22 | C34-01: the cache-comparison row now names the two report keys it rests on, and section 2 records the re-derivation of the verdict with the corrected eligibility rule (executed case list, not filter alone). | + +## What this closes + +Maintainer review comments **R34-03** (publish an evidence record for the final +source and current matrix) and **R34-04** (retain the checks that specifically +protect the module split). + +This report supersedes the earlier chronological runs as the acceptance record. +Those runs are not withdrawn — their per-case verdicts were genuine — but they +cannot serve as the final record, for reasons stated in section 1. + +## Why the earlier evidence could not simply be cited + +The review suggested proving that existing receipts already covered the final +tree rather than rerunning. That path was checked and is not available: + +| Fact | Consequence | +| --- | --- | +| `db15e64..5ef96b5` is one commit touching only `tests/e2e/refactor_fidelity/cases.py` | The production tree was indeed unchanged at the tip. | +| `candidate_db15e64` held 67 records and **no DP-placement case** | The four cases that took the table from 67 to 71 had never run against the refactor tip. | +| The old `baseline` label's last execution was a filtered `dp_` run of four cases merged onto 67, with `cache_clean_before_run: false` | It was an assembled label: its cache listing did not come from one clean full run, and its `case_count` said 72 over 71 result lines because a retained record named a case id that had left the table. | +| Checkpoint A added per-case provenance | Every label captured before it now fails the gate for want of provenance, by design. | + +Both sides were therefore recaptured as single clean full runs. + +## Environment + +| Field | Value | +| --- | --- | +| Host | `kun-workspace-vgen2`, CPU only | +| Python | `/data/ycfeng/envs/frontier-py310/bin/python`, CPython 3.10.6 | +| Packages | numpy 2.2.6, pandas 2.3.3, scikit-learn 1.7.2, scipy 1.15.3, plotly 7.1.0, pytest 9.1.1 | +| Output root | `/data/ycfeng/tmp/issue26-correctness-pr/refactor-fidelity` (scratch, not committed) | + +## 1. Fidelity matrix — the final record + +Both sides were run by the harness committed at `bb582a4`, from a clean +detached checkout of that commit, so the measurement and the code being +measured agree and neither side was measured from a working tree. + +```bash +# from .worktrees/fidelity-candidate-bb582a4, PYTHONPATH=$PWD +python tests/e2e/refactor_fidelity/run_matrix.py run \ + --repo-root .worktrees/fidelity-baseline-main --label baseline_v2 \ + --output-root "$OUT" --jobs 6 --clean-cache --continue-on-failure +python tests/e2e/refactor_fidelity/run_matrix.py run \ + --repo-root . --label candidate_bb582a4 \ + --output-root "$OUT" --jobs 6 --clean-cache --continue-on-failure +python tests/e2e/refactor_fidelity/run_matrix.py compare \ + --output-root "$OUT" --baseline-label baseline_v2 --candidate-label candidate_bb582a4 +``` + +### Provenance of the two sides + +| Field | `baseline_v2` | `candidate_bb582a4` | +| --- | --- | --- | +| Source revision | `1f694f7c549aa3aeeb7c5bbae04e119c09167a77` | `bb582a41702eb31edf4a6477fcf76e4f32d0f3ef` | +| Source working tree | clean | clean | +| Harness revision | `bb582a4` | `bb582a4` | +| Case filter | none | none | +| Cases executed in this run | 71 | 71 | +| `case_count` in the manifest | 71 | 71 | +| Cache cleaned before the run | yes | yes | +| Cache files produced | 426 | 426 | + +### Result + +| Measure | Expected | Actual | Result | +| --- | --- | --- | --- | +| Case table size | 71 | 71 | — | +| Cases compared | 71 | **71** | PASS | +| Identical | 71 | **71** | PASS | +| Mismatched | 0 | **0** | PASS | +| Baseline failures | 0 | **0** | PASS | +| Candidate-only failures | 0 | **0** | PASS | +| Cases missing from one side | 0 | **0** | PASS | +| Cases with missing evidence | 0 | **0** | PASS | +| Cases with differing definitions | 0 | **0** | PASS | +| Cases not compared | 0 | **0** | PASS | +| Cases not compared without explanation | 0 | **0** | PASS | +| Provenance findings | none | **none** | PASS | +| Predictor cache names compared | yes | **yes** (`predictor_cache_populated_cleanly: true`, `predictor_cache_compared: true`; both manifests list all 71 cases in `cases_executed_in_last_run`) | PASS | +| Predictor cache differences | 0 | **0 baseline-only, 0 candidate-only** | PASS | + +**Comparison exit code 0.** This is the first run of this matrix in which +`cases_compared` equals the case-table size *and* the gate that checks it is +the corrected one, so the number means what it says. + +The cache result is worth stating separately: all 426 cache file names matched, +across a table that includes the six trained-predictor cases, which means the +split changed no training identity and no model cache key. Output equality alone +could not show that, because retraining from the same CSV reproduces the same +numbers. + +**Re-derived 2026-09-22 (C34-01).** The eligibility rule behind the cache row +was found too weak: it accepted `cache_clean_before_run` plus an empty +`case_filter`, which a `--start`/`--limit` continuation also satisfies. The rule +now requires each manifest's `cases_executed_in_last_run` to equal the full +case table. Both labels here record 71 executed cases, and rerunning +`run_matrix.py compare` on the retained output root with the corrected harness +reproduces this table exactly: 71 of 71 identical, `predictor_cache_compared: +true`, 0 baseline-only and 0 candidate-only cache files, exit code 0. The +acceptance verdict therefore stands under the corrected rule; see +`test_report_2026-09-22_cache_eligibility_correction.md`. + +## 2. Unit suite — regression comparison + +Whole `tests/unit` directory, both sides, same command, same exclusions. + +**Excluded, not run:** ten modules fail at collection on both sides for missing +optional dependencies (`torch`, `matplotlib`), which the minimal release +environment deliberately omits: + +```text +test_mla_native_profiling_wrapper test_moe_routing_input_contract +test_moe_fused_event_contract test_native_profiling_model_type_policy +test_moe_gating_constructor_boundary test_profiling_timing_stats_contract +test_moe_load_distribution_contract test_sim_walltime_scaling_plot +test_moe_native_admission test_vllm_rocm_attention_wrapper_increment9 +``` + +| Side | Failed | Passed | Skipped | Collection errors | +| --- | --- | --- | --- | --- | +| Baseline `1f694f7` | 84 | 3644 | 49 | 10 (excluded above) | +| This branch | 84 | **3679** | 49 | 10 (excluded above) | + +**The 84 failing test identities are byte-identical between the two sides** +(`diff` of the sorted `FAILED` lists is empty). The 35 additional passing tests +are exactly the 35 added in Checkpoints A and B. + +Inherited failures by module, all pre-existing on `main`: + +| Module | Count | Cause | +| --- | --- | --- | +| `test_pdaf_parity_reference_observer_bootstrap.py` | 51 | Requires the pinned PD-AF reference checkout, absent on this host | +| `test_colocation_release_review_contracts.py` | 10 | Opens `tests/debug/` scripts not tracked on `main`; one spawns a bare `python` absent from PATH | +| `test_profiling_governance_minimal_red.py` | 5 | Profiling dependencies | +| `test_moe_mxfp4_increment10.py` | 5 | Profiling dependencies | +| `test_pdd_public_surface_docs.py` | 3 | Documentation contract drift on `main` | +| `test_pdaf_examples.py` | 3 | PD-AF reference checkout | +| `test_examples_documentation_contracts.py`, `test_collectives_increment11.py` | 2 each | Documentation / optional backend | +| three attention-modeling modules | 1 each | Profiling dependencies | + +This is a regression comparison. It is **not** a claim that these 84 tests pass. + +## 3. Retained checks for the split (R34-04) + +`tests/unit/test_module_split_boundaries.py`, 13 tests, all passing. These ran +as task-local commands during the split; committing them is the difference +between a check that happened once and one that keeps happening. + +| Review priority | Test | +| --- | --- | +| Generated CLI/config behavior beyond flag names | `test_no_new_config_dataclass_loses_its_annotations`, `test_the_flat_cli_can_still_be_generated` | +| Public re-exports and annotation resolution | `test_public_config_names_resolve_from_the_entry_point_that_is_imported`, `test_the_split_modules_stay_reachable_through_the_package` | +| Subclass override and `super()` resolution | `test_vllm_v1_scheduler_reaches_every_extracted_mixin`, `test_split_classes_keep_their_mixin_order`, `test_cluster_config_fields_come_only_from_the_owning_class`, `test_sglang_still_reaches_the_extracted_decision_log_helper` | +| Persistent estimator loading in a fresh manager | `test_a_cached_estimator_loads_into_a_fresh_registry` | +| Importability and the names that previously failed | `test_every_split_module_imports_in_a_fresh_interpreter_order`, `test_runtime_only_names_are_not_hidden_behind_type_checking` | + +### One finding, pinned rather than fixed + +`typing.get_type_hints(ClusterConfig)` raises `NameError: BaseCCBackendConfig`. +The same lookup fails on `1f694f7`, where the class still lived in the single +`config.py`, so this **predates the split**: `cc_backend_config` is annotated +with a name the module does not import at runtime, presumably to avoid a +circular import. The generated CLI is unaffected, which +`test_the_flat_cli_can_still_be_generated` asserts directly. It is recorded in +`KNOWN_UNRESOLVED_CONFIG_ANNOTATIONS` so that a *new* unresolvable annotation +fails the test, and so that fixing this one also fails it until the pin is +dropped. + +## 4. Cross-revision cache loading + +The committed unit test writes and reads a cache within one process. The review +asked for more: a **baseline-produced** cache loaded by the split code. The +recapture produced exactly that artifact, so it was used. + +| Check | Result | +| --- | --- | +| Pickled estimators in `.worktrees/fidelity-baseline-main/cache` (written by `1f694f7`) | 142 | +| Unpickled by the split code | **142 / 142** | +| Also usable — `predict()` ran on those exposing `n_features_in_` | 86 | +| Loaded through `PredictionModelRegistry._load_model_from_cache` | ok | +| Failures | **0** | + +A pickle records the module path of the class it holds, and moving code is +exactly what invalidates that. Matching cache *names* would not have caught it. + +## Limits + +- The matrix establishes that 71 supported configurations produce identical + output. It says nothing about configurations outside the table and makes no + accuracy claim. +- Dummy-mode cases exercise structure and lifecycle, not realistic latency. Six + cases use the checked-in profiling CSVs, which is what covers dataset loading, + training identity and the persistent cache. +- The predictor cache comparison compares file names, not pickle bytes, which + are not required to be reproducible. Section 4 covers loadability separately. +- Section 4 loads the estimators the 71-case table happens to train. It is not + an exhaustive inventory of every artifact Frontier can cache. +- Ten unit modules were not executed at all. Their coverage is unknown on both + sides, not passing. diff --git a/tests/e2e/refactor_fidelity/cases.py b/tests/e2e/refactor_fidelity/cases.py new file mode 100644 index 00000000..c765ce1a --- /dev/null +++ b/tests/e2e/refactor_fidelity/cases.py @@ -0,0 +1,598 @@ +"""Deterministic case table for the refactor fidelity matrix. + +The matrix answers one question: does a behavior-preserving refactor change any +simulator output? Every case therefore runs a checked-in example wrapper, so +the matrix exercises the same configuration surface the release documents, and +a case cannot silently drift away from a supported recipe. + +Cases are written out explicitly instead of being generated from a full cross +product. A reviewer can read the tables below and see exactly which +configuration each case pins. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from typing import Mapping, Sequence + + +COLOCATION_OFFLINE_DENSE = "examples/architecture/co-location/offline/dense_model_basic.sh" +COLOCATION_OFFLINE_MOE = "examples/architecture/co-location/offline/moe_model_basic.sh" +COLOCATION_OFFLINE_THINKING = "examples/architecture/co-location/offline/thinking_mode_basic.sh" +COLOCATION_OFFLINE_SPEC_DEC = "examples/architecture/co-location/offline/moe_spec_dec.sh" +COLOCATION_OFFLINE_PREFIX_CACHE = "examples/architecture/co-location/offline/moe_prefix_caching.sh" +COLOCATION_ONLINE_DENSE = "examples/architecture/co-location/online/dense_model_basic_online.sh" +COLOCATION_ONLINE_MOE = "examples/architecture/co-location/online/moe_model_basic_online.sh" +COLOCATION_ONLINE_THINKING = "examples/architecture/co-location/online/thinking_mode_basic_online.sh" +COLOCATION_ONLINE_SPEC_DEC = "examples/architecture/co-location/online/moe_spec_dec_online.sh" +COLOCATION_ONLINE_PREFIX_CACHE = "examples/architecture/co-location/online/moe_prefix_caching_online.sh" + +PDD_OFFLINE_DENSE = "examples/architecture/pdd/offline/dense_model_basic.sh" +PDD_OFFLINE_MOE = "examples/architecture/pdd/offline/moe_model_basic.sh" +PDD_OFFLINE_THINKING = "examples/architecture/pdd/offline/thinking_mode_basic.sh" +PDD_OFFLINE_SPEC_DEC = "examples/architecture/pdd/offline/moe_spec_dec.sh" +PDD_OFFLINE_PREFIX_CACHE = "examples/architecture/pdd/offline/moe_prefix_caching.sh" +PDD_ONLINE_DENSE = "examples/architecture/pdd/online/dense_model_basic_online.sh" +PDD_ONLINE_MOE = "examples/architecture/pdd/online/moe_model_basic_online.sh" + +PDAF_OFFLINE_DENSE = "examples/architecture/pd-af-disagg/offline/dense_model_basic.sh" +PDAF_OFFLINE_MOE = "examples/architecture/pd-af-disagg/offline/moe_model_basic.sh" +PDAF_OFFLINE_MOE_EP = "examples/architecture/pd-af-disagg/offline/moe_model_ep.sh" +PDAF_OFFLINE_DENSE_CUDA_GRAPH = "examples/architecture/pd-af-disagg/offline/dense_cuda_graph.sh" +PDAF_OFFLINE_MOE_CUDA_GRAPH = "examples/architecture/pd-af-disagg/offline/moe_cuda_graph.sh" +PDAF_ONLINE_DENSE = "examples/architecture/pd-af-disagg/online/dense_model_basic_online.sh" +PDAF_ONLINE_MOE = "examples/architecture/pd-af-disagg/online/moe_model_basic_online.sh" +PDAF_ONLINE_MOE_EP = "examples/architecture/pd-af-disagg/online/moe_model_ep_online.sh" +PDAF_ONLINE_MOE_CUDA_GRAPH = "examples/architecture/pd-af-disagg/online/moe_cuda_graph_online.sh" + +PROFILING_SMOKE_DENSE_CSV = "examples/profiling/smoke_simulator_dense_csv.sh" +PROFILING_SMOKE_MOE_CSV = "examples/profiling/smoke_simulator_moe_csv.sh" + + +@dataclass(frozen=True) +class FidelityCase: + """One reproducible simulator run. + + ``env`` overrides the wrapper's uppercase variables. ``extra_args`` is + appended after the wrapper's ``--`` separator and therefore reaches + ``frontier.main`` directly, which is how the matrix covers flags a wrapper + does not expose. + """ + + case_id: str + group: str + script: str + purpose: str + env: Mapping[str, str] = field(default_factory=dict) + extra_args: Sequence[str] = field(default_factory=tuple) + #: Non-dummy cases train predictors into the label-wide cache directory. + #: They are run serially and before the parallel cases so that the cache is + #: populated once per label rather than raced by several writers. + uses_trained_predictor: bool = False + + def as_record(self) -> dict: + return { + "case_id": self.case_id, + "group": self.group, + "script": self.script, + "purpose": self.purpose, + "env": dict(self.env), + "extra_args": list(self.extra_args), + "uses_trained_predictor": self.uses_trained_predictor, + "case_digest": self.definition_digest(), + } + + def definition_digest(self) -> str: + """Digest the fields that decide what this case actually runs. + + Two sides are only comparable when the same ``case_id`` meant the same + run on both of them. The digest covers the wrapper, the environment + overrides, the extra arguments and the serial/parallel classification, + and deliberately omits ``group`` and ``purpose``, which are + documentation and change nothing about the subprocess. + """ + + payload = json.dumps( + { + "script": self.script, + "env": dict(self.env), + "extra_args": list(self.extra_args), + "uses_trained_predictor": self.uses_trained_predictor, + }, + sort_keys=True, + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:16] + + +def _colocation_dense_offline() -> list[FidelityCase]: + """Co-location offline dense: request population, parallelism, runtime toggles.""" + + rows: list[tuple[str, str, dict[str, str], tuple[str, ...]]] = [ + ( + "coloc_dense_offline_small", + "smallest deterministic dense run", + {"NUM_REQUESTS": "4", "PREFILL_TOKENS": "128", "DECODE_TOKENS": "16", + "ATTN_TP": "1", "DECODE_CUDA_GRAPH_MODE": "none"}, + (), + ), + ( + "coloc_dense_offline_default", + "the published example defaults", + {}, + (), + ), + ( + "coloc_dense_offline_long_prefill", + "3584-token prefill through chunked prefill, just inside the 4096 context", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "3584", "DECODE_TOKENS": "16", + "MAX_TOKENS_IN_BATCH": "2048", "LONG_PREFILL_TOKEN_THRESHOLD": "512"}, + (), + ), + ( + "coloc_dense_offline_no_chunked_prefill", + "chunked prefill disabled", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "1024", "DECODE_TOKENS": "32", + "ENABLE_CHUNKED_PREFILL": "false", "LONG_PREFILL_TOKEN_THRESHOLD": "0", + "DECODE_CUDA_GRAPH_MODE": "none"}, + (), + ), + ( + "coloc_dense_offline_pp2", + "two pipeline stages", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "512", "DECODE_TOKENS": "32", + "PP": "2", "DECODE_CUDA_GRAPH_MODE": "none"}, + (), + ), + ( + "coloc_dense_offline_tp1", + "no attention tensor parallelism", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "512", "DECODE_TOKENS": "32", + "ATTN_TP": "1", "DECODE_CUDA_GRAPH_MODE": "none"}, + (), + ), + ( + "coloc_dense_offline_two_replicas", + "cluster scheduler distributes over two replicas", + {"NUM_REQUESTS": "16", "PREFILL_TOKENS": "256", "DECODE_TOKENS": "32", + "NUM_REPLICAS": "2"}, + (), + ), + ( + "coloc_dense_offline_many_requests", + "64 short requests", + {"NUM_REQUESTS": "64", "PREFILL_TOKENS": "128", "DECODE_TOKENS": "16"}, + (), + ), + ( + "coloc_dense_offline_attn_dp2", + "two attention DP lanes", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "512", "DECODE_TOKENS": "32", + "ATTN_TP": "2", "DECODE_CUDA_GRAPH_MODE": "none"}, + ("--replica_config_attn_dp", "2"), + ), + ( + "coloc_dense_offline_cuda_graph_piecewise", + "piecewise decode CUDA graph mode", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "512", "DECODE_TOKENS": "32", + "DECODE_CUDA_GRAPH_MODE": "piecewise"}, + (), + ), + ( + "coloc_dense_offline_sarathi", + "Sarathi replica scheduler", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "512", "DECODE_TOKENS": "32", + "REPLICA_SCHEDULER": "sarathi", "DECODE_CUDA_GRAPH_MODE": "none"}, + (), + ), + ( + "coloc_dense_offline_sglang", + "SGLang-style prefill-first scheduler", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "512", "DECODE_TOKENS": "32", + "REPLICA_SCHEDULER": "sglang", "DECODE_CUDA_GRAPH_MODE": "none"}, + (), + ), + ( + "coloc_dense_offline_dummy_time_quarter", + "different dummy operator latency", + {"NUM_REQUESTS": "8", "PREFILL_TOKENS": "512", "DECODE_TOKENS": "32", + "DUMMY_EXEC_TIME_MS": "0.25"}, + (), + ), + ] + return [ + FidelityCase(case_id, "colocation_offline_dense", COLOCATION_OFFLINE_DENSE, + purpose, env, extra) + for case_id, purpose, env, extra in rows + ] + + +def _colocation_moe_offline() -> list[FidelityCase]: + """Co-location offline MoE: EP/TP domains, routing distributions, topk.""" + + rows: list[tuple[str, str, dict[str, str], tuple[str, ...]]] = [ + ("coloc_moe_offline_default", "the published MoE example defaults", {}, ()), + ( + "coloc_moe_offline_ep1", + "single expert-parallel domain (attn_tp*attn_dp == moe_tp*moe_ep)", + {"MOE_EP": "1", "MOE_TP": "1", "ATTN_TP": "1"}, + (), + ), + ( + "coloc_moe_offline_moe_tp1", + "expert parallelism without intra-expert TP", + {"MOE_TP": "1", "MOE_EP": "2", "ATTN_TP": "2"}, + (), + ), + ( + "coloc_moe_offline_routing_random", + "random expert-load distribution", + {"MOE_ROUTING_DISTRIBUTION_TYPE": "random"}, + (), + ), + ( + "coloc_moe_offline_routing_skewed", + "skewed expert-load distribution", + {"MOE_ROUTING_DISTRIBUTION_TYPE": "skewed"}, + (), + ), + ( + "coloc_moe_offline_routing_zipf", + "zipf expert-load distribution", + {"MOE_ROUTING_DISTRIBUTION_TYPE": "zipf"}, + (), + ), + ( + "coloc_moe_offline_topk4", + "router top-k of four", + {"ROUTER_TOPK": "4"}, + (), + ), + ( + "coloc_moe_offline_long_prefill", + "2048-token prefill with MoE layers", + {"NUM_REQUESTS": "4", "PREFILL_TOKENS": "2048", "DECODE_TOKENS": "16", + "MAX_TOKENS_IN_BATCH": "2048", "LONG_PREFILL_TOKEN_THRESHOLD": "256"}, + (), + ), + ( + "coloc_moe_offline_many_requests", + "32 short MoE requests", + {"NUM_REQUESTS": "32", "PREFILL_TOKENS": "128", "DECODE_TOKENS": "16"}, + (), + ), + ( + "coloc_moe_offline_cuda_graph_none", + "MoE without decode CUDA graph modeling", + {"DECODE_CUDA_GRAPH_MODE": "none"}, + (), + ), + ( + # The MoE wrapper enforces ATTN_TP == MOE_TP * MOE_EP, so it cannot + # express an attn_dp > 1 shared domain; the dense matrix covers DP + # lanes and this case widens expert parallelism instead. + "coloc_moe_offline_ep4", + "four expert-parallel domains over eight experts", + {"ATTN_TP": "4", "MOE_TP": "1", "MOE_EP": "4", + "DECODE_CUDA_GRAPH_MODE": "none"}, + (), + ), + ] + return [ + FidelityCase(case_id, "colocation_offline_moe", COLOCATION_OFFLINE_MOE, + purpose, env, extra) + for case_id, purpose, env, extra in rows + ] + + +def _colocation_online() -> list[FidelityCase]: + """Co-location online arrivals across three arrival rates.""" + + cases: list[FidelityCase] = [] + for qps in ("0.5", "2.0", "8.0"): + tag = qps.replace(".", "p") + cases.append(FidelityCase( + case_id=f"coloc_dense_online_qps{tag}", + group="colocation_online", + script=COLOCATION_ONLINE_DENSE, + purpose=f"dense online arrivals at {qps} qps", + env={"NUM_REQUESTS": "16", "PREFILL_TOKENS": "512", + "DECODE_TOKENS": "64", "QPS": qps}, + )) + cases.append(FidelityCase( + case_id=f"coloc_moe_online_qps{tag}", + group="colocation_online", + script=COLOCATION_ONLINE_MOE, + purpose=f"MoE online arrivals at {qps} qps", + env={"NUM_REQUESTS": "16", "PREFILL_TOKENS": "256", + "DECODE_TOKENS": "32", "QPS": qps}, + )) + return cases + + +def _colocation_features() -> list[FidelityCase]: + """Thinking mode, speculative decoding, and prefix caching, offline and online.""" + + rows = [ + ("coloc_thinking_offline", COLOCATION_OFFLINE_THINKING, + "thinking mode with tool-call rounds", {}), + ("coloc_thinking_online", COLOCATION_ONLINE_THINKING, + "thinking mode with online arrivals", {}), + ("coloc_spec_dec_offline", COLOCATION_OFFLINE_SPEC_DEC, + "speculative decoding on a MoE model", {}), + ("coloc_spec_dec_online", COLOCATION_ONLINE_SPEC_DEC, + "speculative decoding with online arrivals", {}), + ("coloc_prefix_cache_offline", COLOCATION_OFFLINE_PREFIX_CACHE, + "prefix caching over the shared-session fixture", {}), + ("coloc_prefix_cache_online", COLOCATION_ONLINE_PREFIX_CACHE, + "prefix caching with online arrivals", {}), + ("coloc_spec_dec_offline_ntokens4", COLOCATION_OFFLINE_SPEC_DEC, + "four speculative tokens per iteration", + {"NUM_SPECULATIVE_TOKENS": "4", "COMMITTED_TOKENS_PER_ITERATION": "4"}), + ] + return [ + FidelityCase(case_id, "colocation_features", script, purpose, env) + for case_id, script, purpose, env in rows + ] + + +def _pdd() -> list[FidelityCase]: + """Sequential prefill/decode disaggregation.""" + + rows = [ + ("pdd_dense_offline_default", PDD_OFFLINE_DENSE, + "the published sequential PDD dense defaults", {}), + ("pdd_dense_offline_two_prefill_replicas", PDD_OFFLINE_DENSE, + "two prefill replicas feeding one decode replica", + {"PREFILL_REPLICAS": "2", "NUM_REQUESTS": "16", "PREFILL_TOKENS": "256"}), + ("pdd_dense_offline_two_decode_replicas", PDD_OFFLINE_DENSE, + "one prefill replica feeding two decode replicas", + {"DECODE_REPLICAS": "2", "NUM_REQUESTS": "16", "DECODE_TOKENS": "32"}), + ("pdd_dense_offline_tp2", PDD_OFFLINE_DENSE, + "tensor parallelism on both roles", + {"PREFILL_ATTN_TP": "2", "DECODE_ATTN_TP": "2", "NUM_REQUESTS": "8"}), + ("pdd_dense_offline_no_chunked_prefill", PDD_OFFLINE_DENSE, + "chunked prefill disabled", + {"ENABLE_CHUNKED_PREFILL": "false", "LONG_PREFILL_TOKEN_THRESHOLD": "0"}), + ("pdd_moe_offline_default", PDD_OFFLINE_MOE, + "the published sequential PDD MoE defaults", {}), + ("pdd_moe_offline_ep1", PDD_OFFLINE_MOE, + "PDD MoE with a single expert-parallel domain", + {"PREFILL_MOE_EP": "1", "DECODE_MOE_EP": "1", + "PREFILL_ATTN_TP": "1", "DECODE_ATTN_TP": "1"}), + ("pdd_moe_offline_routing_skewed", PDD_OFFLINE_MOE, + "PDD MoE with a skewed expert-load distribution", + {"MOE_ROUTING_DISTRIBUTION_TYPE": "skewed"}), + ("pdd_thinking_offline", PDD_OFFLINE_THINKING, + "PDD thinking mode with KV handoffs", {}), + ("pdd_spec_dec_offline", PDD_OFFLINE_SPEC_DEC, + "PDD speculative decoding", {}), + ("pdd_prefix_cache_offline", PDD_OFFLINE_PREFIX_CACHE, + "PDD prefix caching over the shared-session fixture", {}), + ("pdd_dense_online_qps2", PDD_ONLINE_DENSE, + "PDD dense online arrivals", {"QPS": "2.0", "NUM_REQUESTS": "16"}), + ("pdd_moe_online_qps2", PDD_ONLINE_MOE, + "PDD MoE online arrivals", {"QPS": "2.0", "NUM_REQUESTS": "16"}), + ("pdd_dense_online_qps8", PDD_ONLINE_DENSE, + "PDD dense at a higher arrival rate", + {"QPS": "8.0", "NUM_REQUESTS": "16", "PREFILL_TOKENS": "256"}), + ] + return [ + FidelityCase(case_id, "pd_disaggregation", script, purpose, env) + for case_id, script, purpose, env in rows + ] + + +def _pdaf() -> list[FidelityCase]: + """Sequential attention/FFN disaggregation with KV and M2N transfers.""" + + rows = [ + ("pdaf_dense_offline_default", PDAF_OFFLINE_DENSE, + "the published sequential PD-AF dense defaults", {}), + ("pdaf_moe_offline_default", PDAF_OFFLINE_MOE, + "PD-AF MoE with one expert-parallel domain", {}), + ("pdaf_moe_offline_ep2", PDAF_OFFLINE_MOE_EP, + "PD-AF MoE with two expert-parallel domains", {}), + ("pdaf_moe_offline_ep2_many", PDAF_OFFLINE_MOE_EP, + "PD-AF MoE EP with more requests", + {"NUM_REQUESTS": "16", "PREFILL_TOKENS": "128", "DECODE_TOKENS": "16"}), + ("pdaf_dense_offline_cuda_graph", PDAF_OFFLINE_DENSE_CUDA_GRAPH, + "PD-AF dense with the global CUDA graph contract", {}), + ("pdaf_moe_offline_cuda_graph", PDAF_OFFLINE_MOE_CUDA_GRAPH, + "PD-AF MoE with the global CUDA graph contract", {}), + ("pdaf_dense_online_default", PDAF_ONLINE_DENSE, + "PD-AF dense online arrivals", {}), + ("pdaf_moe_online_default", PDAF_ONLINE_MOE, + "PD-AF MoE online arrivals", {}), + ("pdaf_moe_online_ep2", PDAF_ONLINE_MOE_EP, + "PD-AF MoE EP online arrivals", {}), + ("pdaf_moe_online_cuda_graph", PDAF_ONLINE_MOE_CUDA_GRAPH, + "PD-AF MoE online with the global CUDA graph contract", {}), + ] + return [ + FidelityCase(case_id, "pd_af_disaggregation", script, purpose, env) + for case_id, script, purpose, env in rows + ] + + +def _trained_predictor() -> list[FidelityCase]: + """Runs with the dummy predictor disabled, fed by checked-in profiling CSVs. + + These exercise dataset loading, training-identity computation, the estimator + registry, and the persistent predictor cache, which the dummy-mode cases + never reach. + """ + + rows: list[tuple[str, str, str, dict[str, str], tuple[str, ...]]] = [ + # These two wrappers default DATA_DIR_BASE to an absolute path under + # the repository root, and the predictor cache key includes the + # profiling input paths. A repository-relative base keeps the key + # identical across two checkouts so that a real change of training + # identity is visible in the cache file names. + ( + "trained_dense_csv_smoke", + PROFILING_SMOKE_DENSE_CSV, + "checked-in dense CSVs fed straight into the simulator", + {"DATA_DIR_BASE": "data/profiling"}, + (), + ), + ( + "trained_moe_csv_smoke", + PROFILING_SMOKE_MOE_CSV, + "checked-in MoE CSVs fed straight into the simulator", + {"DATA_DIR_BASE": "data/profiling"}, + (), + ), + ( + "trained_coloc_dense_h800", + COLOCATION_OFFLINE_DENSE, + "co-location dense on h800 profiles without dummy mode", + {"ENABLE_DUMMY_MODE": "false", "DEVICE": "h800", + "MODEL_NAME": "llama2_7b_dense_example", "ATTN_TP": "1", + "NUM_REQUESTS": "4", "PREFILL_TOKENS": "128", "DECODE_TOKENS": "8", + "DECODE_CUDA_GRAPH_MODE": "none", "MAX_TOKENS_IN_BATCH": "256", + "LONG_PREFILL_TOKEN_THRESHOLD": "64"}, + ("--random_forrest_execution_time_predictor_config_skip_cpu_overhead_modeling",), + ), + ( + "trained_coloc_dense_h800_cuda_graph", + COLOCATION_OFFLINE_DENSE, + "same dense profiles through the kernel-only CUDA graph path", + {"ENABLE_DUMMY_MODE": "false", "DEVICE": "h800", + "MODEL_NAME": "llama2_7b_dense_example", "ATTN_TP": "1", + "NUM_REQUESTS": "4", "PREFILL_TOKENS": "128", "DECODE_TOKENS": "8", + "DECODE_CUDA_GRAPH_MODE": "full_decode_only", + "MAX_TOKENS_IN_BATCH": "256", "LONG_PREFILL_TOKEN_THRESHOLD": "64"}, + ("--random_forrest_execution_time_predictor_config_skip_cpu_overhead_modeling",), + ), + ( + "trained_coloc_moe_phi_h800", + COLOCATION_OFFLINE_MOE, + "co-location MoE on the tiny Phi profiles without dummy mode", + {"ENABLE_DUMMY_MODE": "false", "MODEL_NAME": "Phi-tiny-MoE-instruct", + "ATTN_TP": "1", "MOE_TP": "1", "MOE_EP": "1", + "MOE_ROUTING_DISTRIBUTION_TYPE": "random", + "NUM_REQUESTS": "4", "PREFILL_TOKENS": "128", "DECODE_TOKENS": "8", + "DECODE_CUDA_GRAPH_MODE": "none", "MAX_TOKENS_IN_BATCH": "256", + "LONG_PREFILL_TOKEN_THRESHOLD": "64"}, + ("--replica_config_device", "h800", + "--random_forrest_execution_time_predictor_config_skip_cpu_overhead_modeling"), + ), + ( + "trained_coloc_moe_qwen3_h800", + COLOCATION_OFFLINE_MOE, + "co-location MoE on the tiny Qwen3 profiles without dummy mode", + {"ENABLE_DUMMY_MODE": "false", "MODEL_NAME": "Qwen3-30B-A3B-tiny", + "ATTN_TP": "1", "MOE_TP": "1", "MOE_EP": "1", + "MOE_ROUTING_DISTRIBUTION_TYPE": "random", + "NUM_REQUESTS": "4", "PREFILL_TOKENS": "128", "DECODE_TOKENS": "8", + "DECODE_CUDA_GRAPH_MODE": "none", "MAX_TOKENS_IN_BATCH": "256", + "LONG_PREFILL_TOKEN_THRESHOLD": "64"}, + ("--replica_config_device", "h800", + "--random_forrest_execution_time_predictor_config_skip_cpu_overhead_modeling"), + ), + ] + return [ + FidelityCase(case_id, "trained_predictor", script, purpose, env, extra, + uses_trained_predictor=True) + for case_id, script, purpose, env, extra in rows + ] + + +def _dp_placement() -> list[FidelityCase]: + """Cases whose outcome depends on how requests are placed on DP lanes. + + The rest of the matrix runs one attention DP lane, where lane placement + cannot vary, so it gives a change to DP placement almost no coverage. These + cases exist to give the next such change a real blast radius. + + Two properties matter and are covered separately. Placement must not depend + on how an identical request stream is divided across scheduling calls, + which needs a run that enters the scheduler many times, so most of these + are online. And placement across replicas and across lanes interact, which + needs more than one replica as well as more than one lane. + + Only dense co-location cases appear here, and that is a coverage limit + worth stating rather than a choice. The prefill role reaches the same + placement path as the monolithic role, but no shipped recipe can give it + more than one lane: a dense model in a disaggregated architecture is + rejected with "Dense models do not support attn data parallelism in + disaggregated mode", and the MoE wrappers require + ``ATTN_TP == MOE_TP * MOE_EP`` while the runtime requires + ``attn_tp * attn_dp == moe_tp * moe_ep``, which have no common solution + above one lane. Placement changes affecting the prefill role therefore + have to be validated by unit tests, not by this matrix. + """ + + rows: list[tuple[str, str, str, dict[str, str], tuple[str, ...]]] = [ + ( + "dp_dense_online_lanes2", + COLOCATION_ONLINE_DENSE, + "two lanes with arrivals spread over many scheduling calls", + {"NUM_REQUESTS": "16", "PREFILL_TOKENS": "256", "DECODE_TOKENS": "32", + "QPS": "2.0", "ATTN_TP": "2", "DECODE_CUDA_GRAPH_MODE": "none"}, + ("--replica_config_attn_dp", "2"), + ), + ( + "dp_dense_online_lanes4", + COLOCATION_ONLINE_DENSE, + "four lanes, so the lane index wraps more than once", + {"NUM_REQUESTS": "16", "PREFILL_TOKENS": "256", "DECODE_TOKENS": "32", + "QPS": "2.0", "ATTN_TP": "4", "DECODE_CUDA_GRAPH_MODE": "none"}, + ("--replica_config_attn_dp", "4"), + ), + ( + "dp_dense_offline_lanes2_replicas2", + COLOCATION_OFFLINE_DENSE, + "two lanes over two replicas, where the replica and lane terms interact", + {"NUM_REQUESTS": "16", "PREFILL_TOKENS": "256", "DECODE_TOKENS": "32", + "NUM_REPLICAS": "2", "ATTN_TP": "2", "DECODE_CUDA_GRAPH_MODE": "none"}, + ("--replica_config_attn_dp", "2"), + ), + ( + "dp_dense_online_lanes2_replicas2", + COLOCATION_ONLINE_DENSE, + "two lanes over two replicas with arrivals spread across calls", + {"NUM_REQUESTS": "24", "PREFILL_TOKENS": "256", "DECODE_TOKENS": "16", + "QPS": "4.0", "NUM_REPLICAS": "2", "ATTN_TP": "2", + "DECODE_CUDA_GRAPH_MODE": "none"}, + ("--replica_config_attn_dp", "2"), + ), + ] + return [ + FidelityCase(case_id, "dp_placement", script, purpose, env, extra) + for case_id, script, purpose, env, extra in rows + ] + + +def build_cases() -> list[FidelityCase]: + """Return the full ordered matrix. + + Order is stable so that a partial run selected with ``--start``/``--limit`` + covers the same cases on both sides of a comparison. + """ + + cases: list[FidelityCase] = [] + cases.extend(_trained_predictor()) + cases.extend(_colocation_dense_offline()) + cases.extend(_colocation_moe_offline()) + cases.extend(_colocation_online()) + cases.extend(_colocation_features()) + cases.extend(_pdd()) + cases.extend(_pdaf()) + cases.extend(_dp_placement()) + + seen: set[str] = set() + for case in cases: + if case.case_id in seen: + raise ValueError(f"duplicate fidelity case id: {case.case_id}") + seen.add(case.case_id) + return cases + + +def group_counts(cases: Sequence[FidelityCase]) -> dict[str, int]: + counts: dict[str, int] = {} + for case in cases: + counts[case.group] = counts.get(case.group, 0) + 1 + return counts diff --git a/tests/e2e/refactor_fidelity/compare.py b/tests/e2e/refactor_fidelity/compare.py new file mode 100644 index 00000000..3489e297 --- /dev/null +++ b/tests/e2e/refactor_fidelity/compare.py @@ -0,0 +1,356 @@ +"""Artifact comparison for the refactor fidelity matrix. + +The gate is exact equality. A behavior-preserving refactor has no reason to +change a simulated number, so this module does not implement tolerances. The +only normalization is the substitution of run-specific absolute paths, which +differ between the two checkouts by construction and are not simulator output. + +Observed on ``1f694f7``: ``system_metrics.json``, ``request_metrics.csv``, +``op_precision_metadata.csv`` and ``frontier_stage_batch_ledger.jsonl`` carry no +timestamps, wall-clock durations, hostnames or paths, and ``config.json`` +carries exactly one absolute path (the metrics output directory). The +substitution list below is therefore small on purpose; anything it fails to +cover shows up as a difference rather than being silently accepted. +""" + +from __future__ import annotations + +import csv +import hashlib +import io +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Sequence + + +MAX_REPORTED_DIFFERENCES_PER_ARTIFACT = 5 +MAX_REPORTED_CACHE_FINDINGS_PER_KIND = 20 + + +@dataclass(frozen=True) +class PathSubstitution: + """One literal replaced by a stable token before comparison.""" + + literal: str + token: str + + +@dataclass(frozen=True) +class ArtifactDifference: + artifact: str + kind: str + detail: str + + def as_record(self) -> dict: + return {"artifact": self.artifact, "kind": self.kind, "detail": self.detail} + + +def substitutions_for(repo_root: Path, output_root: Path, label: str) -> list[PathSubstitution]: + """Build the substitution list for one side of the comparison. + + Longer literals come first so that a nested path is replaced before its + parent directory. + """ + + pairs = [ + (str(Path(output_root).resolve() / label), ""), + (str(Path(output_root).resolve()), ""), + (str(Path(repo_root).resolve()), ""), + ] + pairs.sort(key=lambda pair: len(pair[0]), reverse=True) + return [PathSubstitution(literal, token) for literal, token in pairs] + + +def apply_substitutions(text: str, substitutions: Sequence[PathSubstitution]) -> str: + for substitution in substitutions: + text = text.replace(substitution.literal, substitution.token) + return text + + +def file_digest(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def list_artifacts(directory: Path) -> list[str]: + """Return the sorted relative paths of every file under ``directory``.""" + + if not directory.is_dir(): + return [] + return sorted( + str(path.relative_to(directory)) + for path in directory.rglob("*") + if path.is_file() + ) + + +def _read_text(path: Path, substitutions: Sequence[PathSubstitution]) -> str: + return apply_substitutions(path.read_text(encoding="utf-8"), substitutions) + + +def _json_leaf_differences( + baseline: Any, candidate: Any, path: str = "" +) -> list[str]: + """Return human-readable descriptions of the first differing leaves.""" + + differences: list[str] = [] + if type(baseline) is not type(candidate) and not ( + isinstance(baseline, (int, float)) and isinstance(candidate, (int, float)) + ): + return [f"{path or '/'}: type {type(baseline).__name__} != {type(candidate).__name__}"] + + if isinstance(baseline, dict): + baseline_keys = set(baseline) + candidate_keys = set(candidate) + for key in sorted(baseline_keys - candidate_keys): + differences.append(f"{path}/{key}: missing in candidate") + for key in sorted(candidate_keys - baseline_keys): + differences.append(f"{path}/{key}: missing in baseline") + for key in sorted(baseline_keys & candidate_keys): + if len(differences) >= MAX_REPORTED_DIFFERENCES_PER_ARTIFACT: + break + differences.extend( + _json_leaf_differences(baseline[key], candidate[key], f"{path}/{key}") + ) + return differences[:MAX_REPORTED_DIFFERENCES_PER_ARTIFACT] + + if isinstance(baseline, list): + if len(baseline) != len(candidate): + return [f"{path or '/'}: list length {len(baseline)} != {len(candidate)}"] + for index, (left, right) in enumerate(zip(baseline, candidate)): + if len(differences) >= MAX_REPORTED_DIFFERENCES_PER_ARTIFACT: + break + differences.extend(_json_leaf_differences(left, right, f"{path}[{index}]")) + return differences[:MAX_REPORTED_DIFFERENCES_PER_ARTIFACT] + + if baseline != candidate: + return [f"{path or '/'}: {baseline!r} != {candidate!r}"] + return [] + + +def _compare_json( + name: str, baseline: Path, candidate: Path, substitutions_pair: tuple[Sequence[PathSubstitution], Sequence[PathSubstitution]] +) -> list[ArtifactDifference]: + baseline_subs, candidate_subs = substitutions_pair + baseline_value = json.loads(_read_text(baseline, baseline_subs)) + candidate_value = json.loads(_read_text(candidate, candidate_subs)) + differences = _json_leaf_differences(baseline_value, candidate_value) + return [ArtifactDifference(name, "content", detail) for detail in differences] + + +def _compare_jsonl( + name: str, baseline: Path, candidate: Path, substitutions_pair: tuple[Sequence[PathSubstitution], Sequence[PathSubstitution]] +) -> list[ArtifactDifference]: + baseline_subs, candidate_subs = substitutions_pair + baseline_lines = _read_text(baseline, baseline_subs).splitlines() + candidate_lines = _read_text(candidate, candidate_subs).splitlines() + if len(baseline_lines) != len(candidate_lines): + return [ArtifactDifference( + name, "content", + f"record count {len(baseline_lines)} != {len(candidate_lines)}", + )] + + differences: list[ArtifactDifference] = [] + for index, (left, right) in enumerate(zip(baseline_lines, candidate_lines)): + if left == right: + continue + leaf_differences = _json_leaf_differences(json.loads(left), json.loads(right)) + for detail in leaf_differences: + differences.append(ArtifactDifference(name, "content", f"record {index}: {detail}")) + if len(differences) >= MAX_REPORTED_DIFFERENCES_PER_ARTIFACT: + break + return differences[:MAX_REPORTED_DIFFERENCES_PER_ARTIFACT] + + +def _compare_csv( + name: str, baseline: Path, candidate: Path, substitutions_pair: tuple[Sequence[PathSubstitution], Sequence[PathSubstitution]] +) -> list[ArtifactDifference]: + baseline_subs, candidate_subs = substitutions_pair + baseline_rows = list(csv.reader(io.StringIO(_read_text(baseline, baseline_subs)))) + candidate_rows = list(csv.reader(io.StringIO(_read_text(candidate, candidate_subs)))) + if len(baseline_rows) != len(candidate_rows): + return [ArtifactDifference( + name, "content", + f"row count {len(baseline_rows)} != {len(candidate_rows)}", + )] + + differences: list[ArtifactDifference] = [] + header = baseline_rows[0] if baseline_rows else [] + for row_index, (left_row, right_row) in enumerate(zip(baseline_rows, candidate_rows)): + if left_row == right_row: + continue + if len(left_row) != len(right_row): + differences.append(ArtifactDifference( + name, "content", + f"row {row_index}: column count {len(left_row)} != {len(right_row)}", + )) + else: + for column_index, (left, right) in enumerate(zip(left_row, right_row)): + if left == right: + continue + column = ( + header[column_index] + if row_index > 0 and column_index < len(header) + else f"column {column_index}" + ) + differences.append(ArtifactDifference( + name, "content", + f"row {row_index} [{column}]: {left!r} != {right!r}", + )) + if len(differences) >= MAX_REPORTED_DIFFERENCES_PER_ARTIFACT: + break + if len(differences) >= MAX_REPORTED_DIFFERENCES_PER_ARTIFACT: + break + return differences[:MAX_REPORTED_DIFFERENCES_PER_ARTIFACT] + + +def _compare_bytes(name: str, baseline: Path, candidate: Path) -> list[ArtifactDifference]: + if baseline.read_bytes() == candidate.read_bytes(): + return [] + return [ArtifactDifference( + name, "content", + f"binary contents differ (sha256 {file_digest(baseline)[:16]} != {file_digest(candidate)[:16]})", + )] + + +def compare_artifact_directories( + baseline_dir: Path, + candidate_dir: Path, + baseline_substitutions: Sequence[PathSubstitution], + candidate_substitutions: Sequence[PathSubstitution], + ignored_artifacts: Iterable[str] = (), +) -> list[ArtifactDifference]: + """Compare every artifact under the two directories. + + A file present on one side only is a difference; so is any content + difference that survives path substitution. + + An absent directory is reported rather than treated as an empty one. + ``list_artifacts`` returns an empty list for a path that does not exist, so + without this check two deleted artifact directories would produce two equal + empty inventories and the case would be recorded as identical. A + comparison with nothing to compare is missing evidence, not agreement. + """ + + differences: list[ArtifactDifference] = [] + for side, directory in (("baseline", baseline_dir), ("candidate", candidate_dir)): + if not directory.is_dir(): + differences.append(ArtifactDifference( + directory.name, + "missing_directory", + f"the {side} artifact directory does not exist: {directory}", + )) + if differences: + return differences + + ignored = set(ignored_artifacts) + baseline_names = [name for name in list_artifacts(baseline_dir) if name not in ignored] + candidate_names = [name for name in list_artifacts(candidate_dir) if name not in ignored] + + for name in sorted(set(baseline_names) - set(candidate_names)): + differences.append(ArtifactDifference(name, "missing_in_candidate", "produced only by the baseline")) + for name in sorted(set(candidate_names) - set(baseline_names)): + differences.append(ArtifactDifference(name, "missing_in_baseline", "produced only by the candidate")) + + substitutions_pair = (baseline_substitutions, candidate_substitutions) + for name in sorted(set(baseline_names) & set(candidate_names)): + baseline_path = baseline_dir / name + candidate_path = candidate_dir / name + suffix = Path(name).suffix.lower() + try: + if suffix == ".json": + differences.extend(_compare_json(name, baseline_path, candidate_path, substitutions_pair)) + elif suffix == ".jsonl": + differences.extend(_compare_jsonl(name, baseline_path, candidate_path, substitutions_pair)) + elif suffix == ".csv": + differences.extend(_compare_csv(name, baseline_path, candidate_path, substitutions_pair)) + else: + differences.extend(_compare_bytes(name, baseline_path, candidate_path)) + except (UnicodeDecodeError, json.JSONDecodeError, csv.Error) as error: + differences.append(ArtifactDifference( + name, "unreadable", f"{type(error).__name__}: {error}", + )) + return differences + + +CACHE_NAME_PATTERN = re.compile( + r"^(?P.*?)_?(?P[0-9a-f]{8,})(?P[^/]*)$" +) + + +@dataclass(frozen=True) +class CacheNameFinding: + """One classified difference between two sides' predictor cache files.""" + + kind: str + stem: str + detail: str + + def as_record(self) -> dict: + return {"kind": self.kind, "stem": self.stem, "detail": self.detail} + + +def _split_cache_name(name: str) -> tuple[str, str, str]: + """Split a cache file name into (stem, digest, suffix). + + Model artifacts are named ``{model_name}_{hash}{suffix}`` and lock files + ``{hash}_{kind}_lock.file``, so a name without a stem is a lock file. + """ + + match = CACHE_NAME_PATTERN.match(name) + if not match: + return name, "", "" + return match.group("stem"), match.group("digest"), match.group("suffix") + + +def classify_cache_differences( + baseline_names: Sequence[str], candidate_names: Sequence[str] +) -> list[CacheNameFinding]: + """Explain why two sides' predictor cache file names differ. + + The distinction that matters is between a *re-keyed* artifact, where the + same model name appears on both sides under a different hash, and an + artifact that exists on one side only. The first means a training + identity or cache key changed, which a refactor must not do; the second + means a model started or stopped being trained. + """ + + only_baseline = sorted(set(baseline_names) - set(candidate_names)) + only_candidate = sorted(set(candidate_names) - set(baseline_names)) + if not only_baseline and not only_candidate: + return [] + + def by_stem(names: Sequence[str]) -> dict[str, list[str]]: + grouped: dict[str, list[str]] = {} + for name in names: + stem, digest, suffix = _split_cache_name(name) + key = f"{stem}{suffix}" if stem else f"{suffix}" + grouped.setdefault(key, []).append(digest) + return grouped + + baseline_by_stem = by_stem(only_baseline) + candidate_by_stem = by_stem(only_candidate) + + findings: list[CacheNameFinding] = [] + for stem in sorted(set(baseline_by_stem) & set(candidate_by_stem)): + findings.append(CacheNameFinding( + "rekeyed", stem, + f"baseline {', '.join(sorted(baseline_by_stem[stem]))} -> " + f"candidate {', '.join(sorted(candidate_by_stem[stem]))}", + )) + for stem in sorted(set(baseline_by_stem) - set(candidate_by_stem)): + findings.append(CacheNameFinding( + "only_in_baseline", stem, + f"hashes {', '.join(sorted(baseline_by_stem[stem]))}", + )) + for stem in sorted(set(candidate_by_stem) - set(baseline_by_stem)): + findings.append(CacheNameFinding( + "only_in_candidate", stem, + f"hashes {', '.join(sorted(candidate_by_stem[stem]))}", + )) + return findings diff --git a/tests/e2e/refactor_fidelity/measure_commit.py b/tests/e2e/refactor_fidelity/measure_commit.py new file mode 100644 index 00000000..91670f25 --- /dev/null +++ b/tests/e2e/refactor_fidelity/measure_commit.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +"""Measure one commit against the captured baseline, from a detached checkout. + + PYTHONPATH=$PWD python tests/e2e/refactor_fidelity/measure_commit.py \\ + --sha 99922d2 --output-root "$FRONTIER_TMP_ROOT/refactor-fidelity" + +Running the matrix against a shared working tree measures whatever happens to +be on disk at that moment. When two people work on one branch, that is a +mixture, and a mixture is not a verdict on anything: one such run reported 66 +failures that belonged to somebody else's half-finished edit. + +This driver removes that failure mode. It checks the commit out into its own +detached worktree, runs the matrix there under a label named after the commit, +and compares against the baseline captured earlier. The harness it runs is the +one committed at that revision, so the measurement and the code being measured +always agree. + +The worktree is kept by default: a validated checkout of an earlier split is +the reference you want when a later split reports a difference and you need to +attribute it. Pass --remove-worktree when you no longer need it. +""" + +from __future__ import annotations + +import argparse +import subprocess +import sys +from pathlib import Path +from typing import Sequence + + +def _git(*args: str, cwd: Path | None = None) -> str: + result = subprocess.run( + ["git", *args], cwd=str(cwd) if cwd else None, + capture_output=True, text=True, check=True, + ) + return result.stdout.strip() + + +def _main_worktree(start: Path) -> Path: + """Return the repository's main worktree, given any worktree inside it.""" + + common_dir = Path(_git("rev-parse", "--path-format=absolute", "--git-common-dir", cwd=start)) + return common_dir.parent + + +def reuse_blocked_reason(checkout: Path, full_sha: str) -> str | None: + """Return why an existing checkout may not be reused, or None if it may. + + A detached worktree is not read-only. Matching HEAD says which commit was + checked out, not what is on disk now, and this driver's whole claim is that + the measurement belongs to one commit. + """ + + existing = _git("rev-parse", "HEAD", cwd=checkout) + if existing != full_sha: + return f"it is at {existing}, not {full_sha}" + dirty = _git("status", "--porcelain", cwd=checkout) + if dirty: + return "its working tree has been modified:\n" + "\n".join( + f" {line}" for line in dirty.splitlines() + ) + return None + + +def _run(command: Sequence[str], cwd: Path) -> int: + print(f"\n$ {' '.join(command)}", flush=True) + return subprocess.run(command, cwd=str(cwd), check=False).returncode + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--sha", required=True, help="commit to measure") + parser.add_argument("--output-root", required=True) + parser.add_argument("--baseline-label", default="baseline") + parser.add_argument("--label", default=None, + help="defaults to candidate_") + parser.add_argument("--python-bin", default=sys.executable) + parser.add_argument("--jobs", type=int, default=4) + parser.add_argument("--limit", type=int, default=None, + help="measure only the first N cases, for a plumbing check") + parser.add_argument("--case-filter", default=None) + parser.add_argument("--remove-worktree", action="store_true", + help="delete the detached checkout after comparing") + args = parser.parse_args(argv) + + here = Path.cwd() + main_worktree = _main_worktree(here) + full_sha = _git("rev-parse", f"{args.sha}^{{commit}}", cwd=here) + short_sha = full_sha[:7] + label = args.label or f"candidate_{short_sha}" + checkout = main_worktree / ".worktrees" / f"fidelity-candidate-{short_sha}" + + print(f"commit {full_sha}") + print(f"subject {_git('log', '-1', '--format=%s', full_sha, cwd=here)}") + print(f"label {label}") + print(f"checkout {checkout}") + + if checkout.exists(): + # Report and stop; do not clean the modifications away to make the + # check pass, because whoever made them has not been asked. + blocked = reuse_blocked_reason(checkout, full_sha) + if blocked is not None: + print(f"refusing to reuse {checkout}: {blocked}", file=sys.stderr) + print( + " restore the checkout yourself, or remove it and let this " + "driver create a fresh one.", + file=sys.stderr, + ) + return 2 + print("reusing the existing detached checkout at the same commit") + else: + _git("worktree", "add", "--detach", str(checkout), full_sha, cwd=main_worktree) + + driver = checkout / "tests" / "e2e" / "refactor_fidelity" / "run_matrix.py" + if not driver.is_file(): + print(f"the harness is not present at {full_sha}: {driver}", file=sys.stderr) + return 2 + + run_command = [ + args.python_bin, str(driver), "run", + "--repo-root", str(checkout), + "--label", label, + "--output-root", args.output_root, + "--python-bin", args.python_bin, + "--jobs", str(args.jobs), + "--clean-cache", + "--continue-on-failure", + ] + if args.limit is not None: + run_command += ["--limit", str(args.limit)] + if args.case_filter: + run_command += ["--case-filter", args.case_filter] + + # The driver imports tests.e2e.refactor_fidelity, so it needs its own + # checkout on the path, not the one this script was invoked from. + import os + env_note = f"PYTHONPATH={checkout}" + print(f"\n({env_note})") + os.environ["PYTHONPATH"] = str(checkout) + + if _run(run_command, cwd=checkout) != 0: + print("\nthe run reported failing cases; comparing anyway", file=sys.stderr) + + compare_command = [ + args.python_bin, str(driver), "compare", + "--output-root", args.output_root, + "--baseline-label", args.baseline_label, + "--candidate-label", label, + ] + verdict = _run(compare_command, cwd=checkout) + + if args.remove_worktree: + _git("worktree", "remove", "--force", str(checkout), cwd=main_worktree) + print(f"removed {checkout}") + else: + print(f"\nkept {checkout} as a reference; " + f"remove it with: git -C {main_worktree} worktree remove {checkout}") + + print("\nVERDICT: " + ("IDENTICAL" if verdict == 0 else "DIFFERENCES FOUND")) + return verdict + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/e2e/refactor_fidelity/run_matrix.py b/tests/e2e/refactor_fidelity/run_matrix.py new file mode 100644 index 00000000..a8190e8d --- /dev/null +++ b/tests/e2e/refactor_fidelity/run_matrix.py @@ -0,0 +1,846 @@ +#!/usr/bin/env python3 +"""Driver for the refactor fidelity matrix. + +Capture one side, then the other, then compare: + + PYTHONPATH=$PWD python tests/e2e/refactor_fidelity/run_matrix.py run \\ + --repo-root /path/to/baseline/worktree --label baseline \\ + --output-root "$FRONTIER_TMP_ROOT/refactor-fidelity" --clean-cache + PYTHONPATH=$PWD python tests/e2e/refactor_fidelity/run_matrix.py run \\ + --repo-root "$PWD" --label candidate \\ + --output-root "$FRONTIER_TMP_ROOT/refactor-fidelity" --clean-cache + PYTHONPATH=$PWD python tests/e2e/refactor_fidelity/run_matrix.py compare \\ + --output-root "$FRONTIER_TMP_ROOT/refactor-fidelity" + +The driver always comes from the checkout it is invoked in, while ``--repo-root`` +selects the simulator under test, so both sides are measured by one case table +and one comparator. + +Each case runs a checked-in example wrapper with the working directory set to +the repository under test, which is what makes the relative profiling-data and +``cache`` paths in the shipped configuration resolve. Each side therefore keeps +its own predictor cache inside its own checkout. The names of the cache files a +side produces are recorded and compared as well, because a changed training +identity or cache key would otherwise be invisible: retraining from the same CSV +yields the same numbers. +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +from tests.e2e.refactor_fidelity.cases import FidelityCase, build_cases, group_counts +from tests.e2e.refactor_fidelity.compare import ( + MAX_REPORTED_CACHE_FINDINGS_PER_KIND, + classify_cache_differences, + compare_artifact_directories, + file_digest, + list_artifacts, + substitutions_for, +) + + +DEFAULT_TIMEOUT_SECONDS = 1800 + + +@dataclass +class CaseOutcome: + case: FidelityCase + returncode: int + duration_seconds: float + artifact_dir: str | None + artifacts: list[dict] + error: str | None + log_tail: str | None = None + + def as_record(self) -> dict: + record = self.case.as_record() + record.update({ + "returncode": self.returncode, + "duration_seconds": round(self.duration_seconds, 3), + "artifact_dir": self.artifact_dir, + "artifacts": self.artifacts, + "error": self.error, + "log_tail": self.log_tail, + }) + return record + + +def _git(repo_root: Path, *args: str) -> str: + result = subprocess.run( + ["git", "-C", str(repo_root), *args], + capture_output=True, text=True, check=False, + ) + return result.stdout.strip() + + +def _package_versions(python_bin: str) -> dict[str, str]: + script = ( + "import json,platform;" + "mods=['numpy','pandas','sklearn','scipy','plotly','fasteners','ddsketch'];" + "out={'python':platform.python_version()};" + "\nfor m in mods:\n" + " try:\n" + " out[m]=__import__(m).__version__\n" + " except Exception as e:\n" + " out[m]='unavailable: %s' % type(e).__name__\n" + "print(json.dumps(out))" + ) + result = subprocess.run([python_bin, "-c", script], capture_output=True, text=True, check=False) + try: + return json.loads(result.stdout.strip().splitlines()[-1]) + except (ValueError, IndexError): + return {"error": result.stderr.strip()[:400]} + + +ARTIFACT_DISCOVERY_RETRY_SECONDS = 10.0 + + +def _find_artifact_dir(metrics_root: Path, wait_seconds: float = 0.0) -> Path | None: + """Locate the single normalized metrics directory a run produced. + + A successful run that appears to have written nothing is retried for a + bounded window: the directory listing can lag the child process on a + networked filesystem, and reporting a spurious failure would be worse than + waiting. A genuinely empty run still fails, only later. + """ + + deadline = time.monotonic() + wait_seconds + while True: + candidates = sorted( + path.parent for path in metrics_root.rglob("system_metrics.json") + ) + if len(candidates) == 1: + return candidates[0] + if len(candidates) > 1: + raise RuntimeError( + f"expected one metrics directory under {metrics_root}, " + f"found {len(candidates)}: " + + ", ".join(str(path) for path in candidates) + ) + if time.monotonic() >= deadline: + return None + time.sleep(0.5) + + +def _log_tail(log_path: Path, max_lines: int = 20) -> str: + """Return the last lines of a run log, for a failure record.""" + + try: + lines = log_path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError as error: + return f"" + return "\n".join(lines[-max_lines:]) + + +def _run_case( + case: FidelityCase, + repo_root: Path, + label_root: Path, + python_bin: str, + timeout_seconds: int, +) -> CaseOutcome: + case_root = label_root / "cases" / case.case_id + metrics_root = case_root / "metrics" + if case_root.exists(): + shutil.rmtree(case_root) + metrics_root.mkdir(parents=True) + + script_path = repo_root / case.script + if not script_path.is_file(): + return CaseOutcome(case, 127, 0.0, None, [], f"missing script: {script_path}") + + env = os.environ.copy() + env.pop("PYTHONSTARTUP", None) + env.update({ + "PYTHONPATH": str(repo_root), + "PYTHONHASHSEED": "0", + "PYTHONDONTWRITEBYTECODE": "1", + "PYTHON_BIN": python_bin, + "WANDB_DISABLED": "true", + "VIDUR_DISABLE_WANDB": "1", + "METRICS_OUTPUT_DIR": str(metrics_root), + "RUN_ID": case.case_id, + }) + env.update(case.env) + + command = ["bash", str(script_path)] + if case.extra_args: + command.append("--") + command.extend(case.extra_args) + + log_path = case_root / "run.log" + started = time.monotonic() + with log_path.open("w", encoding="utf-8") as log: + log.write(f"# command: {' '.join(command)}\n") + log.write(f"# cwd: {repo_root}\n") + log.write("# env overrides: " + json.dumps(dict(case.env), sort_keys=True) + "\n\n") + log.flush() + try: + completed = subprocess.run( + command, cwd=str(repo_root), env=env, + stdout=log, stderr=subprocess.STDOUT, + timeout=timeout_seconds, check=False, + ) + returncode = completed.returncode + error = None + except subprocess.TimeoutExpired: + returncode = 124 + error = f"timed out after {timeout_seconds}s" + duration = time.monotonic() - started + + artifact_dir: Path | None = None + artifacts: list[dict] = [] + try: + artifact_dir = _find_artifact_dir( + metrics_root, + wait_seconds=ARTIFACT_DISCOVERY_RETRY_SECONDS if returncode == 0 else 0.0, + ) + except RuntimeError as failure: + error = str(failure) + if artifact_dir is not None: + for name in list_artifacts(artifact_dir): + path = artifact_dir / name + artifacts.append({ + "name": name, + "sha256": file_digest(path), + "bytes": path.stat().st_size, + }) + elif returncode == 0 and error is None: + error = "run reported success but produced no system_metrics.json" + + return CaseOutcome( + case=case, + returncode=returncode, + duration_seconds=duration, + artifact_dir=( + str(artifact_dir.relative_to(label_root)) if artifact_dir is not None else None + ), + artifacts=artifacts, + error=error, + # Keep the evidence with the record: a later re-run of the same case + # overwrites run.log, which would otherwise erase why it failed. + log_tail=_log_tail(log_path) if (returncode != 0 or error) else None, + ) + + +def _select_cases( + cases: Sequence[FidelityCase], case_filter: str | None, start: int, limit: int | None +) -> list[FidelityCase]: + selected = [case for case in cases if not case_filter or case_filter in case.case_id] + selected = selected[start:] + if limit is not None: + selected = selected[:limit] + return selected + + +def source_provenance(repo_root: Path) -> dict: + """Identify the source a case was measured against. + + Every case record carries this, not just the label-wide manifest, because + a filtered run rewrites the manifest with its own revision while keeping + the records of the cases it did not execute. Without a per-case stamp + there is nothing left to show that those retained cases ran on the same + source. + """ + + return { + "source_revision": _git(repo_root, "rev-parse", "HEAD"), + "source_dirty": bool(_git(repo_root, "status", "--porcelain")), + "harness_revision": _git(Path(__file__).resolve().parents[3], "rev-parse", "HEAD"), + } + + +def check_retained_records( + results_path: Path, provenance: dict, executed_ids: set[str], known_ids: set[str] +) -> tuple[dict[str, dict], list[str], list[str]]: + """Decide which previously recorded cases this run may keep. + + Returns the retained records, the case ids whose provenance conflicts with + this run, and the case ids that are no longer in the case table. A + conflicting record is never silently dropped or silently kept: the caller + refuses the run so that one label always describes one source. + """ + + retained: dict[str, dict] = {} + if results_path.is_file(): + for line in results_path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + record = json.loads(line) + retained[record["case_id"]] = record + + conflicts = [ + case_id + for case_id, record in sorted(retained.items()) + if case_id not in executed_ids + and any(record.get(key) != value for key, value in provenance.items()) + ] + stale = sorted(case_id for case_id in retained if case_id not in known_ids) + return retained, conflicts, stale + + +def run_label(args: argparse.Namespace) -> int: + repo_root = Path(args.repo_root).resolve() + output_root = Path(args.output_root).resolve() + label_root = output_root / args.label + label_root.mkdir(parents=True, exist_ok=True) + + cases = _select_cases(build_cases(), args.case_filter, args.start, args.limit) + if not cases: + print("no cases selected", file=sys.stderr) + return 2 + + provenance = source_provenance(repo_root) + results_path = label_root / "results.jsonl" + all_cases = build_cases() + retained, conflicts, stale = check_retained_records( + results_path, + provenance, + {case.case_id for case in cases}, + {case.case_id for case in all_cases}, + ) + # Refuse before running anything rather than after, so a rejected label + # costs a second instead of a full matrix. + if conflicts: + print( + f"refusing to add to label {args.label!r}: {len(conflicts)} retained " + "case records describe a different source, a different harness, or " + "carry no provenance at all, and this run would leave the label " + "describing a mixture.", + file=sys.stderr, + ) + for case_id in conflicts[:10]: + recorded = {key: retained[case_id].get(key) for key in provenance} + print(f" {case_id}: recorded {recorded}", file=sys.stderr) + if len(conflicts) > 10: + print(f" ... and {len(conflicts) - 10} more", file=sys.stderr) + print(f" this run: {provenance}", file=sys.stderr) + print(" write to a new label, or re-run the whole matrix.", file=sys.stderr) + return 2 + if stale: + print( + f"dropping {len(stale)} retained record(s) for case ids that are no " + f"longer in the case table: {', '.join(stale)}" + ) + + cache_dir = repo_root / "cache" + if args.clean_cache and cache_dir.exists(): + shutil.rmtree(cache_dir) + + print(f"label={args.label} repo_root={repo_root} cases={len(cases)}") + print(f"source {provenance['source_revision'][:12]} " + f"dirty={provenance['source_dirty']} " + f"harness {provenance['harness_revision'][:12]}") + print(f"groups: {json.dumps(group_counts(cases), sort_keys=True)}") + + serial_cases = [case for case in cases if case.uses_trained_predictor] + parallel_cases = [case for case in cases if not case.uses_trained_predictor] + + outcomes: dict[str, CaseOutcome] = {} + + def record(outcome: CaseOutcome) -> None: + outcomes[outcome.case.case_id] = outcome + status = "ok" if outcome.returncode == 0 and outcome.error is None else "FAIL" + detail = f" ({outcome.error})" if outcome.error else "" + print( + f" [{status:>4}] {outcome.case.case_id} rc={outcome.returncode} " + f"{outcome.duration_seconds:.1f}s artifacts={len(outcome.artifacts)}{detail}", + flush=True, + ) + + for case in serial_cases: + record(_run_case(case, repo_root, label_root, args.python_bin, args.timeout_seconds)) + + if parallel_cases: + with ThreadPoolExecutor(max_workers=max(1, args.jobs)) as pool: + futures = [ + pool.submit(_run_case, case, repo_root, label_root, args.python_bin, args.timeout_seconds) + for case in parallel_cases + ] + for future in futures: + record(future.result()) + + ordered = [outcomes[case.case_id] for case in cases] + # A filtered or partial run refreshes only the cases it executed and keeps + # the records of the cases it skipped, so the two sides stay comparable. + # The guard above has already established that the retained records + # describe this same source and harness. + merged = dict(retained) + for outcome in ordered: + merged[outcome.case.case_id] = {**outcome.as_record(), **provenance} + case_order = [case.case_id for case in all_cases] + written = [case_id for case_id in case_order if case_id in merged] + with results_path.open("w", encoding="utf-8") as handle: + for case_id in written: + handle.write(json.dumps(merged[case_id], sort_keys=True) + "\n") + + cache_files = sorted( + str(path.relative_to(cache_dir)) for path in cache_dir.rglob("*") if path.is_file() + ) if cache_dir.is_dir() else [] + + manifest = { + "label": args.label, + "repo_root": str(repo_root), + "git_head": _git(repo_root, "rev-parse", "HEAD"), + "git_describe_branch": _git(repo_root, "rev-parse", "--abbrev-ref", "HEAD"), + "git_dirty_paths": _git(repo_root, "status", "--short"), + "python_bin": args.python_bin, + "package_versions": _package_versions(args.python_bin), + # What the results file holds, not what the merge dictionary held: a + # retained record for a case id that has since left the table is + # dropped, and counting it here is what made an earlier label report + # 72 cases over 71 result lines. + "case_count": len(written), + "cases_executed_in_last_run": [case.case_id for case in cases], + "cases_dropped_as_stale": stale, + **provenance, + "case_filter": args.case_filter, + "cache_dir": str(cache_dir), + "cache_clean_before_run": bool(args.clean_cache), + "cache_files": cache_files, + "driver_checkout": str(Path(__file__).resolve().parents[3]), + } + (label_root / "manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + failures = [o for o in ordered if o.returncode != 0 or o.error is not None] + print(f"\n{len(ordered) - len(failures)}/{len(ordered)} cases produced artifacts") + if failures: + print("cases that did not produce artifacts:") + for outcome in failures: + print(f" {outcome.case.case_id}: rc={outcome.returncode} {outcome.error or ''}") + print(f"results: {results_path}") + return 0 if not failures or args.continue_on_failure else 1 + + +def _load_side(output_root: Path, label: str) -> tuple[dict, dict[str, dict]]: + label_root = output_root / label + manifest = json.loads((label_root / "manifest.json").read_text(encoding="utf-8")) + results: dict[str, dict] = {} + for line in (label_root / "results.jsonl").read_text(encoding="utf-8").splitlines(): + if line.strip(): + record = json.loads(line) + results[record["case_id"]] = record + return manifest, results + + +PROVENANCE_KEYS = ("source_revision", "source_dirty", "harness_revision") + + +def side_provenance_findings(label: str, results: dict[str, dict]) -> list[str]: + """Report anything that stops this label from describing one measurement.""" + + findings: list[str] = [] + without = sorted( + case_id for case_id, record in results.items() + if any(key not in record for key in PROVENANCE_KEYS) + ) + if without: + findings.append( + f"{label}: {len(without)} case record(s) carry no source provenance, " + f"so the label cannot be shown to describe one revision " + f"(first: {', '.join(without[:3])}). Recapture this side." + ) + stamped = { + tuple(record[key] for key in PROVENANCE_KEYS) + for record in results.values() + if all(key in record for key in PROVENANCE_KEYS) + } + if len(stamped) > 1: + findings.append( + f"{label}: case records describe {len(stamped)} different " + f"source/harness combinations: {sorted(stamped)}" + ) + for revision, dirty, _harness in sorted(stamped): + if dirty: + findings.append( + f"{label}: measured against a modified working tree at " + f"{revision[:12]}, which is not a commit-specific result" + ) + return findings + + +def missing_evidence_for(record: dict, label_root: Path) -> str | None: + """Return why a successful case's recorded artifacts cannot be compared. + + A record that says a case succeeded is not evidence on its own. The files + it named have to still be on disk, because the comparison reads the + directory and an absent directory otherwise reads as an empty one. + """ + + if not record.get("artifact_dir"): + return "the record names no artifact directory" + recorded = {entry["name"] for entry in record.get("artifacts", [])} + if not recorded: + return "the record lists no artifacts" + directory = label_root / record["artifact_dir"] + if not directory.is_dir(): + return f"the recorded artifact directory is gone: {directory}" + on_disk = set(list_artifacts(directory)) + if on_disk != recorded: + lost = sorted(recorded - on_disk) + extra = sorted(on_disk - recorded) + return ( + f"the directory no longer matches the record " + f"({len(lost)} missing, {len(extra)} unexpected)" + + (f"; missing {lost[:3]}" if lost else "") + ) + return None + + +def compare_labels(args: argparse.Namespace) -> int: + output_root = Path(args.output_root).resolve() + baseline_manifest, baseline_results = _load_side(output_root, args.baseline_label) + candidate_manifest, candidate_results = _load_side(output_root, args.candidate_label) + + baseline_root = output_root / args.baseline_label + candidate_root = output_root / args.candidate_label + baseline_subs = substitutions_for( + Path(baseline_manifest["repo_root"]), output_root, args.baseline_label + ) + candidate_subs = substitutions_for( + Path(candidate_manifest["repo_root"]), output_root, args.candidate_label + ) + + identical: list[str] = [] + mismatched: list[dict] = [] + baseline_failures: list[dict] = [] + candidate_only_failures: list[dict] = [] + missing: list[str] = [] + missing_evidence: list[dict] = [] + definition_mismatches: list[dict] = [] + + for case_id in sorted(set(baseline_results) | set(candidate_results)): + baseline_record = baseline_results.get(case_id) + candidate_record = candidate_results.get(case_id) + if baseline_record is None or candidate_record is None: + missing.append(case_id) + continue + + # The two sides are joined by case id, so the id has to have meant the + # same run on both of them. Without this an edited case definition + # compares two different experiments under one name. + baseline_digest = baseline_record.get("case_digest") + candidate_digest = candidate_record.get("case_digest") + if baseline_digest != candidate_digest: + definition_mismatches.append({ + "case_id": case_id, + "baseline_case_digest": baseline_digest, + "candidate_case_digest": candidate_digest, + }) + continue + + baseline_ok = baseline_record["returncode"] == 0 and not baseline_record["error"] + candidate_ok = candidate_record["returncode"] == 0 and not candidate_record["error"] + if not baseline_ok: + baseline_failures.append({ + "case_id": case_id, + "baseline_returncode": baseline_record["returncode"], + "baseline_error": baseline_record["error"], + "candidate_returncode": candidate_record["returncode"], + "candidate_error": candidate_record["error"], + "candidate_also_failed": not candidate_ok, + }) + continue + if not candidate_ok: + candidate_only_failures.append({ + "case_id": case_id, + "candidate_returncode": candidate_record["returncode"], + "candidate_error": candidate_record["error"], + "candidate_log_tail": candidate_record.get("log_tail"), + }) + continue + + evidence_problems = { + side: reason + for side, reason in ( + ("baseline", missing_evidence_for(baseline_record, baseline_root)), + ("candidate", missing_evidence_for(candidate_record, candidate_root)), + ) + if reason is not None + } + if evidence_problems: + missing_evidence.append({"case_id": case_id, **evidence_problems}) + continue + + differences = compare_artifact_directories( + baseline_root / baseline_record["artifact_dir"], + candidate_root / candidate_record["artifact_dir"], + baseline_subs, + candidate_subs, + ) + if differences: + mismatched.append({ + "case_id": case_id, + "differences": [difference.as_record() for difference in differences], + }) + else: + identical.append(case_id) + + # The predictor cache is populated by the cases that actually ran, so the + # cache comparison only means something when both sides ran the same full + # matrix. A filtered or partial run would otherwise report every model the + # other side trained as a difference. + full_case_set = {case.case_id for case in build_cases()} + # A cache listing is only a fair comparison when each side's cache was + # populated by one clean run of the whole table. A label assembled from a + # filtered continuation without --clean-cache carries models trained by an + # earlier case selection, which is not what the other side has. The + # filter alone does not settle it: --start and --limit also narrow the + # executed selection and leave no filter in the manifest, so the executed + # case ids are checked by name. A manifest without that list cannot show + # a full run and is not compared. + def populated_by_one_clean_full_run(manifest: dict) -> bool: + executed = manifest.get("cases_executed_in_last_run") + return bool( + manifest.get("cache_clean_before_run") + and not manifest.get("case_filter") + and executed is not None + and set(executed) == full_case_set + ) + + cache_populated_cleanly = all( + populated_by_one_clean_full_run(manifest) + for manifest in (baseline_manifest, candidate_manifest) + ) + cache_comparable = ( + set(baseline_results) == full_case_set + and set(candidate_results) == full_case_set + and cache_populated_cleanly + ) + baseline_cache = baseline_manifest.get("cache_files", []) + candidate_cache = candidate_manifest.get("cache_files", []) + if cache_comparable: + cache_only_in_baseline = sorted(set(baseline_cache) - set(candidate_cache)) + cache_only_in_candidate = sorted(set(candidate_cache) - set(baseline_cache)) + cache_findings = classify_cache_differences(baseline_cache, candidate_cache) + else: + cache_only_in_baseline = [] + cache_only_in_candidate = [] + cache_findings = [] + + # Completeness is about what was compared, not about which case ids appear + # in the two result tables. A table can be full of records that all + # describe failed runs, and comparing none of them is not agreement. + compared_ids = set(identical) | {entry["case_id"] for entry in mismatched} + compared = len(compared_ids) + not_compared = sorted(full_case_set - compared_ids) + complete = not not_compared + provenance_findings = ( + side_provenance_findings(args.baseline_label, baseline_results) + + side_provenance_findings(args.candidate_label, candidate_results) + ) + # Every case that was not compared must be accounted for by one of the + # specific findings above. Anything left over means a path through this + # function dropped a case silently, and the gate must not pass on it. + explained = ( + set(missing) + | {entry["case_id"] for entry in baseline_failures} + | {entry["case_id"] for entry in candidate_only_failures} + | {entry["case_id"] for entry in missing_evidence} + | {entry["case_id"] for entry in definition_mismatches} + ) + absent_from_both = full_case_set - set(baseline_results) - set(candidate_results) + unexplained = sorted(set(not_compared) - explained - absent_from_both) + + report = { + "baseline": { + "label": args.baseline_label, + "repo_root": baseline_manifest["repo_root"], + "git_head": baseline_manifest["git_head"], + }, + "candidate": { + "label": args.candidate_label, + "repo_root": candidate_manifest["repo_root"], + "git_head": candidate_manifest["git_head"], + }, + "identical_cases": identical, + "mismatched_cases": mismatched, + "baseline_failures": baseline_failures, + "candidate_only_failures": candidate_only_failures, + "cases_missing_from_one_side": missing, + "cases_with_missing_evidence": missing_evidence, + "cases_with_differing_definitions": definition_mismatches, + "cases_not_compared": not_compared, + "cases_not_compared_without_explanation": unexplained, + "cases_absent_from_both_sides": sorted(absent_from_both), + "provenance_findings": provenance_findings, + "predictor_cache_populated_cleanly": cache_populated_cleanly, + "predictor_cache_files_only_in_baseline": cache_only_in_baseline, + "predictor_cache_files_only_in_candidate": cache_only_in_candidate, + "predictor_cache_findings": [f.as_record() for f in cache_findings], + "predictor_cache_compared": cache_comparable, + "cases_compared": compared, + "case_table_size": len(full_case_set), + "complete_comparison": complete, + } + report_path = output_root / "comparison.json" + report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + print(f"baseline {args.baseline_label} @ {baseline_manifest['git_head'][:12]}") + print(f"candidate {args.candidate_label} @ {candidate_manifest['git_head'][:12]}") + print(f"cases compared: {compared} of {len(full_case_set)} in the case table") + print(f"identical: {len(identical)}") + print(f"mismatched: {len(mismatched)}") + print(f"baseline failures: {len(baseline_failures)}") + print(f"candidate-only failures: {len(candidate_only_failures)}") + print(f"cases missing from one side: {len(missing)}") + print(f"cases with missing evidence: {len(missing_evidence)}") + print(f"cases with differing definitions: {len(definition_mismatches)}") + print(f"cases not compared: {len(not_compared)}") + if not cache_comparable: + print( + "predictor cache file names: not compared, because at least one " + "side's cache was not populated by one clean run of the complete " + "case set" + ) + else: + print( + "predictor cache file names differ: " + f"{len(cache_only_in_baseline)} baseline-only, " + f"{len(cache_only_in_candidate)} candidate-only" + ) + if cache_findings: + rekeyed = [f for f in cache_findings if f.kind == "rekeyed"] + if rekeyed: + print( + f" {len(rekeyed)} artifacts kept their model name but changed hash, " + "which means a training identity or cache key changed:" + ) + for finding in rekeyed[:MAX_REPORTED_CACHE_FINDINGS_PER_KIND]: + print(f" {finding.stem}: {finding.detail}") + if len(rekeyed) > MAX_REPORTED_CACHE_FINDINGS_PER_KIND: + print(f" ... and {len(rekeyed) - MAX_REPORTED_CACHE_FINDINGS_PER_KIND} more") + for kind, label in (("only_in_baseline", "baseline"), + ("only_in_candidate", "candidate")): + entries = [f for f in cache_findings if f.kind == kind] + if entries: + print(f" {len(entries)} artifacts exist only on the {label} side:") + for finding in entries[:MAX_REPORTED_CACHE_FINDINGS_PER_KIND]: + print(f" {finding.stem}: {finding.detail}") + if len(entries) > MAX_REPORTED_CACHE_FINDINGS_PER_KIND: + print( + f" ... and {len(entries) - MAX_REPORTED_CACHE_FINDINGS_PER_KIND} more" + ) + for entry in mismatched: + print(f"\nMISMATCH {entry['case_id']}") + for difference in entry["differences"]: + print(f" {difference['artifact']} [{difference['kind']}] {difference['detail']}") + for entry in candidate_only_failures: + print(f"\nCANDIDATE-ONLY FAILURE {entry['case_id']}: rc={entry['candidate_returncode']} " + f"{entry['candidate_error'] or ''}") + if entry.get("candidate_log_tail"): + print(" last log lines:") + for line in entry["candidate_log_tail"].splitlines(): + print(f" {line}") + for entry in baseline_failures: + print(f"\nBASELINE FAILURE {entry['case_id']}: rc={entry['baseline_returncode']} " + f"{entry['baseline_error'] or ''}" + + (" (the candidate failed too)" if entry["candidate_also_failed"] else "")) + for entry in missing_evidence: + sides = ", ".join(f"{side}: {reason}" for side, reason in entry.items() + if side != "case_id") + print(f"\nMISSING EVIDENCE {entry['case_id']}: {sides}") + for entry in definition_mismatches: + print(f"\nDIFFERENT DEFINITION {entry['case_id']}: " + f"baseline {entry['baseline_case_digest']} != " + f"candidate {entry['candidate_case_digest']}") + for finding in provenance_findings: + print(f"\nPROVENANCE {finding}") + if unexplained: + print(f"\nUNEXPLAINED: {len(unexplained)} case(s) were neither compared " + f"nor reported as a specific finding: {', '.join(unexplained[:10])}") + print(f"\nreport: {report_path}") + + # A comparison passes only when it actually compared the whole case table. + # Two sides that both ran nothing agree trivially, and two sides whose runs + # all failed have no artifacts to disagree about; neither is evidence about + # the branch. Every reason a case was not compared is therefore a failure + # in its own right, and --allow-partial waives exactly one of them: cases + # nobody attempted on either side. + incomplete = bool(absent_from_both) and not args.allow_partial + if incomplete: + print( + "\nINCOMPLETE: this comparison did not cover the whole case table, " + "so it is not evidence that the branch is unchanged. " + "Re-run both sides without a filter, or pass --allow-partial to " + "accept a deliberate subset." + ) + failed = bool( + mismatched + or baseline_failures + or candidate_only_failures + or missing + or missing_evidence + or definition_mismatches + or unexplained + or provenance_findings + or cache_only_in_baseline + or cache_only_in_candidate + or incomplete + ) + return 1 if failed else 0 + + +def list_cases(_: argparse.Namespace) -> int: + cases = build_cases() + print(f"{len(cases)} cases") + for group, count in sorted(group_counts(cases).items()): + print(f" {group}: {count}") + print() + for case in cases: + marker = "T" if case.uses_trained_predictor else " " + print(f"{marker} {case.case_id:<44} {case.script}") + return 0 + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + subparsers = parser.add_subparsers(dest="command", required=True) + + run_parser = subparsers.add_parser("run", help="capture one side of the matrix") + run_parser.add_argument("--repo-root", required=True, help="checkout of the simulator under test") + run_parser.add_argument("--label", required=True, help="name of this side, e.g. baseline or candidate") + run_parser.add_argument("--output-root", required=True) + run_parser.add_argument("--python-bin", default=sys.executable) + run_parser.add_argument("--jobs", type=int, default=4) + run_parser.add_argument("--case-filter", default=None, help="substring of case_id") + run_parser.add_argument("--start", type=int, default=0) + run_parser.add_argument("--limit", type=int, default=None) + run_parser.add_argument("--timeout-seconds", type=int, default=DEFAULT_TIMEOUT_SECONDS) + run_parser.add_argument("--clean-cache", action="store_true", + help="remove /cache before running") + run_parser.add_argument("--continue-on-failure", action="store_true", + help="exit 0 even when some cases produced no artifacts") + run_parser.set_defaults(func=run_label) + + compare_parser = subparsers.add_parser("compare", help="compare two captured sides") + compare_parser.add_argument("--output-root", required=True) + compare_parser.add_argument("--baseline-label", default="baseline") + compare_parser.add_argument("--candidate-label", default="candidate") + compare_parser.add_argument( + "--allow-partial", action="store_true", + help="accept a comparison that does not cover the whole case table", + ) + compare_parser.set_defaults(func=compare_labels) + + list_parser = subparsers.add_parser("list", help="print the case table") + list_parser.set_defaults(func=list_cases) + + args = parser.parse_args(argv) + return args.func(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_attention_tp_effective_mapping.py b/tests/unit/test_attention_tp_effective_mapping.py index 559db65f..0aaacabb 100644 --- a/tests/unit/test_attention_tp_effective_mapping.py +++ b/tests/unit/test_attention_tp_effective_mapping.py @@ -7,7 +7,7 @@ import pytest from frontier.attention.ops import AttentionOperatorRole -from frontier.execution_time_predictor import shared_prediction_model_manager +from frontier.execution_time_predictor import profiling_dataframe_loaders from frontier.execution_time_predictor import sklearn_execution_time_predictor from frontier.execution_time_predictor.attention_tp_policy import ( get_attention_non_linear_tp_policy_ops, @@ -207,7 +207,7 @@ def test_shared_model_manager_load_attention_df_uses_role_derived_dense_loader_k manager._attention_tp_warning_cache = set() monkeypatch.setattr( - shared_prediction_model_manager, + profiling_dataframe_loaders, "get_enabled_predictor_median_column_by_role", lambda _family, role: { AttentionOperatorRole.CACHE_WRITE: "time_stats.role_cache.median", @@ -215,7 +215,7 @@ def test_shared_model_manager_load_attention_df_uses_role_derived_dense_loader_k raising=False, ) monkeypatch.setattr( - shared_prediction_model_manager, + profiling_dataframe_loaders, "get_enabled_predictor_metric_name_by_role", lambda _family, role: { AttentionOperatorRole.PREFILL_KERNEL: "role_prefill", @@ -228,7 +228,7 @@ def _resolve_effective_tp(**kwargs): return 8 if kwargs["op_name"] == "role_prefill" else 1 monkeypatch.setattr( - shared_prediction_model_manager, + profiling_dataframe_loaders, "resolve_effective_attention_tp_size", _resolve_effective_tp, ) diff --git a/tests/unit/test_model_architecture_registry.py b/tests/unit/test_model_architecture_registry.py index b9fff1ee..d018e86e 100644 --- a/tests/unit/test_model_architecture_registry.py +++ b/tests/unit/test_model_architecture_registry.py @@ -515,7 +515,7 @@ def test_raw_model_profile_resolution_callsites_are_allowlisted() -> None: ("frontier/config/model_config.py", "BaseModelConfig.__post_init__", "helper"): 1, # Config-like prediction adapters may not own a BaseModelConfig snapshot. ( - "frontier/execution_time_predictor/shared_prediction_model_manager.py", + "frontier/execution_time_predictor/prediction_model_identity.py", "_resolve_model_architecture_profile", "helper", ): 1, diff --git a/tests/unit/test_module_split_boundaries.py b/tests/unit/test_module_split_boundaries.py new file mode 100644 index 00000000..f548d3a6 --- /dev/null +++ b/tests/unit/test_module_split_boundaries.py @@ -0,0 +1,323 @@ +"""Checks that specifically protect the four module splits. + +The fidelity matrix proves that supported configurations produce the same +numbers. It cannot see the failure modes a move-heavy refactor actually has, +because those break at import or construction time, before any simulation runs: +an annotation that no longer resolves in its defining module, a name that used +to be re-exported, a mixin whose method the owning class no longer reaches, a +pickled estimator whose module path moved. + +Each test here was run as a task-local command during the split. Committing them +is the difference between a check that happened once and a check that keeps +happening. +""" + +from __future__ import annotations + +import dataclasses +import importlib +import pkgutil +import sys +import typing +from pathlib import Path + +import pytest + +import frontier.config +from frontier.config.config import SimulationConfig +from frontier.config.flat_dataclass import create_flat_dataclass + + +# --- annotation resolution in each defining module -------------------------- + + +def _config_modules() -> list[str]: + package = frontier.config + return sorted( + f"frontier.config.{info.name}" + for info in pkgutil.iter_modules(package.__path__) + if not info.ispkg + ) + + +#: ``ClusterConfig.cc_backend_config`` is annotated ``BaseCCBackendConfig``, +#: which the module does not import at runtime. This predates the split: the +#: same lookup fails on ``1f694f7``, where the class still lived in the single +#: ``config.py``. It is pinned rather than fixed so that a *new* unresolvable +#: annotation, which is what the split could plausibly introduce, fails here. +KNOWN_UNRESOLVED_CONFIG_ANNOTATIONS = {"frontier.config.cluster_config.ClusterConfig"} + + +def test_no_new_config_dataclass_loses_its_annotations() -> None: + """``flat_dataclass`` resolves string annotations in the defining module. + + ``from __future__ import annotations`` makes every annotation a string, and + the CLI generator resolves each one against the namespace of the module + that defines the dataclass. Splitting one module into twelve therefore + means each new module must import the names its own annotations mention, + not merely the names its code calls. A missing import is invisible until + something asks for the type. + """ + + unresolved: dict[str, str] = {} + for module_name in _config_modules(): + module = importlib.import_module(module_name) + for name, value in vars(module).items(): + if not dataclasses.is_dataclass(value) or value.__module__ != module_name: + continue + try: + typing.get_type_hints(value) + except Exception as error: # noqa: BLE001 - the message is the finding + unresolved[f"{module_name}.{name}"] = f"{type(error).__name__}: {error}" + + new = {key: reason for key, reason in unresolved.items() + if key not in KNOWN_UNRESOLVED_CONFIG_ANNOTATIONS} + assert not new, "annotations that no longer resolve:\n" + "\n".join( + f" {key}: {reason}" for key, reason in sorted(new.items()) + ) + fixed = KNOWN_UNRESOLVED_CONFIG_ANNOTATIONS - set(unresolved) + assert not fixed, ( + f"these now resolve, so drop them from the pinned set: {sorted(fixed)}" + ) + + +def test_the_flat_cli_can_still_be_generated() -> None: + """The end-to-end consequence of the test above, through the real generator.""" + + flat = create_flat_dataclass(SimulationConfig) + assert len(dataclasses.fields(flat)) > 700 + + +# --- public re-exports ------------------------------------------------------ + + +def _names_imported_from_frontier_config(repo_root: Path) -> dict[str, set[str]]: + """Collect what other modules import, keyed by the module they import from. + + ``frontier.config`` and ``frontier.config.config`` are different entry + points with different contents: ``PrecisionType`` and the quantization + helpers live on the package, never on ``config.py``. Each import site is + therefore checked against the module it actually names. + """ + + import ast + + wanted: dict[str, set[str]] = {"frontier.config": set(), "frontier.config.config": set()} + for path in sorted((repo_root / "frontier").rglob("*.py")): + try: + tree = ast.parse(path.read_text(encoding="utf-8")) + except SyntaxError: + continue + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module in wanted: + wanted[node.module].update( + alias.name for alias in node.names if alias.name != "*" + ) + return wanted + + +def test_public_config_names_resolve_from_the_entry_point_that_is_imported() -> None: + repo_root = Path(__file__).resolve().parents[2] + holders = { + "frontier.config": frontier.config, + "frontier.config.config": importlib.import_module("frontier.config.config"), + } + + missing: list[str] = [] + for module_name, names in _names_imported_from_frontier_config(repo_root).items(): + for name in sorted(names): + if not hasattr(holders[module_name], name): + missing.append(f"{module_name}.{name}") + assert not missing, "names other modules import but cannot reach: " + ", ".join(missing) + + +def test_the_split_modules_stay_reachable_through_the_package() -> None: + """Everything ``config.py`` re-exports must also come off the package. + + Callers use both spellings, and the split moved the definitions out from + under both of them at once. + """ + + config_module = importlib.import_module("frontier.config.config") + exported = [ + name for name, value in vars(config_module).items() + if not name.startswith("_") and dataclasses.is_dataclass(value) + ] + assert len(exported) > 20, "the re-export surface collapsed" + unreachable = [name for name in exported if not hasattr(frontier.config, name)] + assert not unreachable, ( + "config.py exports these but the package does not: " + ", ".join(sorted(unreachable)) + ) + + +# --- mixin method resolution ------------------------------------------------ + + +def test_vllm_v1_scheduler_reaches_every_extracted_mixin() -> None: + from frontier.scheduler.replica_scheduler.base_replica_scheduler import ( + BaseReplicaScheduler, + ) + from frontier.scheduler.replica_scheduler.vllm_v1_engine_replica_scheduler import ( + VLLMv1EngineReplicaScheduler, + ) + + mro = VLLMv1EngineReplicaScheduler.__mro__ + names = [cls.__name__ for cls in mro] + expected_order = [ + "VLLMv1EngineReplicaScheduler", + "IterationSchedulingPolicy", + "KvBlockAllocation", + "PrefixCacheLedger", + "TargetEmbeddedMtpWaitPolicy", + "DecodeAttentionCohort", + "DisaggregatedRoleScheduling", + ] + assert names[: len(expected_order)] == expected_order + # The base must stay behind every mixin, or a mixin override would lose to + # the base implementation it was extracted from. + assert names.index("BaseReplicaScheduler") > names.index("DisaggregatedRoleScheduling") + assert issubclass(VLLMv1EngineReplicaScheduler, BaseReplicaScheduler) + + +def test_sglang_still_reaches_the_extracted_decision_log_helper() -> None: + """The donor branch deleted this method while this caller kept calling it.""" + + from frontier.scheduler.replica_scheduler.sglang_style_replica_scheduler import ( + SGLangStyleReplicaScheduler, + ) + + method = getattr(SGLangStyleReplicaScheduler, "_get_num_waiting_reqs_for_decision_log") + assert method.__module__ == ( + "frontier.scheduler.replica_scheduler.vllm_v1_iteration_policy" + ), "the helper moved; check that both decision-log callers still resolve it" + + +@pytest.mark.parametrize( + ("class_path", "expected_head"), + ( + ( + "frontier.execution_time_predictor.shared_prediction_model_manager." + "ExecutionTimePredictionModelManager", + ["ExecutionTimePredictionModelManager", "PredictionFamilyTrainers", + "ProfilingDataFrameLoaders", "PredictionModelRegistry", + "LayerContractResolution"], + ), + ( + "frontier.execution_time_predictor.sklearn_moe_execution_time_predictor." + "SklearnMoEExecutionTimePredictor", + ["SklearnMoEExecutionTimePredictor", "MoeOperatorTimes", "MoeRoutingWorkload", + "MoeDatasetTraining", "MoeMtpReplay", "SklearnExecutionTimePredictor"], + ), + ( + "frontier.config.cluster_config.ClusterConfig", + ["ClusterConfig", "ClusterRoleConfigBuilder", "ClusterTopologySummary"], + ), + ), +) +def test_split_classes_keep_their_mixin_order(class_path: str, expected_head: list[str]) -> None: + module_name, _, class_name = class_path.rpartition(".") + owner = getattr(importlib.import_module(module_name), class_name) + assert [cls.__name__ for cls in owner.__mro__][: len(expected_head)] == expected_head + + +def test_cluster_config_fields_come_only_from_the_owning_class() -> None: + """The two extracted mixins hold behavior, not state.""" + + from frontier.config.cluster_config import ClusterConfig + from frontier.config.cluster_role_config import ClusterRoleConfigBuilder + from frontier.config.cluster_topology_summary import ClusterTopologySummary + + for mixin in (ClusterRoleConfigBuilder, ClusterTopologySummary): + assert not dataclasses.is_dataclass(mixin), ( + f"{mixin.__name__} became a dataclass, so it now contributes fields " + "to ClusterConfig and changes the generated CLI" + ) + assert len(dataclasses.fields(ClusterConfig)) == 181 + + +# --- estimator cache loading ------------------------------------------------ + + +def test_a_cached_estimator_loads_into_a_fresh_registry(tmp_path: Path) -> None: + """Cache *names* matching does not prove a cached object still loads. + + A pickle records the module path of the class it holds. Moving code is + exactly what invalidates that, and a cold run that retrains from scratch + would never notice, because it writes the file it then reads. + """ + + from sklearn.tree import DecisionTreeRegressor + + from frontier.execution_time_predictor.prediction_model_registry import ( + PredictionModelRegistry, + ) + + class _Holder(PredictionModelRegistry): + def __init__(self, cache_dir: str) -> None: # noqa: D107 - test scaffold + self._cache_dir = cache_dir + + estimator = DecisionTreeRegressor(random_state=0).fit([[0.0], [1.0]], [0.0, 2.0]) + + writer = _Holder(str(tmp_path)) + writer._store_model_in_cache("attn_prefill", "abc123", estimator) + + reader = _Holder(str(tmp_path)) + loaded = reader._load_model_from_cache("attn_prefill", "abc123") + + assert loaded is not None, "a fresh registry could not read what another wrote" + assert loaded.predict([[1.0]]).tolist() == estimator.predict([[1.0]]).tolist() + assert reader._load_model_from_cache("attn_prefill", "not_this_hash") is None + + +def test_every_split_module_imports_in_a_fresh_interpreter_order() -> None: + """Import each new module first, with nothing else loaded from its package. + + The config split's one real defect was a name available only under + ``TYPE_CHECKING`` while a method constructed it at runtime. Importing a + module in isolation is what surfaces that class of mistake. + """ + + split_modules = [ + *_config_modules(), + "frontier.scheduler.replica_scheduler.vllm_v1_decision_log", + "frontier.scheduler.replica_scheduler.vllm_v1_decode_attn_cohort", + "frontier.scheduler.replica_scheduler.vllm_v1_iteration_policy", + "frontier.scheduler.replica_scheduler.vllm_v1_kv_allocation", + "frontier.scheduler.replica_scheduler.vllm_v1_mtp_wait", + "frontier.scheduler.replica_scheduler.vllm_v1_prefix_cache", + "frontier.scheduler.replica_scheduler.vllm_v1_role_schedules", + "frontier.execution_time_predictor.layer_contract_resolution", + "frontier.execution_time_predictor.moe_dataset_training", + "frontier.execution_time_predictor.moe_mtp_replay", + "frontier.execution_time_predictor.moe_operator_times", + "frontier.execution_time_predictor.moe_predictor_helpers", + "frontier.execution_time_predictor.moe_routing_workload", + "frontier.execution_time_predictor.prediction_family_trainers", + "frontier.execution_time_predictor.prediction_model_identity", + "frontier.execution_time_predictor.prediction_model_registry", + "frontier.execution_time_predictor.profiling_dataframe_loaders", + ] + failures: list[str] = [] + for module_name in split_modules: + try: + importlib.import_module(module_name) + except Exception as error: # noqa: BLE001 - the message is the finding + failures.append(f"{module_name}: {type(error).__name__}: {error}") + assert not failures, "modules that do not import:\n" + "\n".join(failures) + + +def test_runtime_only_names_are_not_hidden_behind_type_checking() -> None: + """The I1 defect, pinned: a runtime constructor needs a runtime import.""" + + from frontier.config.cluster_role_config import ClusterRoleConfigBuilder + + source = Path( + sys.modules[ClusterRoleConfigBuilder.__module__].__file__ + ).read_text(encoding="utf-8") + builder = "get_cluster_configs_for_disaggregation" + assert builder in source + body = source[source.index(f"def {builder}"):] + assert "from frontier.config.cluster_config import ClusterConfig" in body, ( + "this method constructs ClusterConfig at runtime, so the name has to be " + "imported outside TYPE_CHECKING" + ) diff --git a/tests/unit/test_moe_share_expert_operator_families.py b/tests/unit/test_moe_share_expert_operator_families.py index 4279d1f3..e44d0b2e 100644 --- a/tests/unit/test_moe_share_expert_operator_families.py +++ b/tests/unit/test_moe_share_expert_operator_families.py @@ -202,7 +202,10 @@ def test_moe_auxiliary_tp_key_preserves_deferred_pdd_legacy_scope() -> None: def test_moe_column_validation_uses_moe_family_profiling_names(monkeypatch) -> None: import pandas as pd - import frontier.execution_time_predictor.sklearn_moe_execution_time_predictor as moe_module + import frontier.execution_time_predictor.moe_predictor_helpers as moe_module + # MOE_FAMILY is read in two modules after the split, so both bindings + # have to be patched for the fake family to be seen end to end. + import frontier.execution_time_predictor.moe_operator_times as moe_operator_module def _operator(name: str): return SimpleNamespace(name=name, profiling_name=lambda: name) @@ -217,6 +220,7 @@ def _operator(name: str): ) ), ) + monkeypatch.setattr(moe_operator_module, "MOE_FAMILY", moe_module.MOE_FAMILY) with pytest.raises(ValueError, match="time_stats.moe_family_second.median"): moe_module._validate_moe_columns( @@ -229,7 +233,10 @@ def test_sklearn_moe_training_uses_moe_family_profiling_names( tmp_path, ) -> None: import pandas as pd - import frontier.execution_time_predictor.sklearn_moe_execution_time_predictor as moe_module + import frontier.execution_time_predictor.moe_predictor_helpers as moe_module + # MOE_FAMILY is read in two modules after the split, so both bindings + # have to be patched for the fake family to be seen end to end. + import frontier.execution_time_predictor.moe_operator_times as moe_operator_module def _operator(name: str): return SimpleNamespace( @@ -250,6 +257,7 @@ def _operator(name: str): ) ), ) + monkeypatch.setattr(moe_operator_module, "MOE_FAMILY", moe_module.MOE_FAMILY) csv_path = tmp_path / "moe.csv" pd.DataFrame( @@ -314,7 +322,10 @@ def test_sklearn_moe_dataset_contract_filters_gating_context_from_operator_preci monkeypatch, ) -> None: import pandas as pd - import frontier.execution_time_predictor.sklearn_moe_execution_time_predictor as moe_module + import frontier.execution_time_predictor.moe_predictor_helpers as moe_module + # MOE_FAMILY is read in two modules after the split, so both bindings + # have to be patched for the fake family to be seen end to end. + import frontier.execution_time_predictor.moe_operator_times as moe_operator_module def _operator(name: str): return SimpleNamespace( @@ -330,6 +341,7 @@ def _operator(name: str): "MOE_FAMILY", SimpleNamespace(profiling_ops=lambda: (_operator("moe_family_gate"),)), ) + monkeypatch.setattr(moe_operator_module, "MOE_FAMILY", moe_module.MOE_FAMILY) predictor = object.__new__(_ConcreteSklearnMoEExecutionTimePredictor) predictor._model_config = SimpleNamespace( @@ -367,7 +379,9 @@ def test_shared_manager_validates_moe_training_names_from_moe_family( monkeypatch, tmp_path, ) -> None: - import frontier.execution_time_predictor.shared_prediction_model_manager as manager_module + # MOE_FAMILY is read by _get_moe_family_model_names, which resolves it in + # its own module, so the patch has to target that module. + import frontier.execution_time_predictor.prediction_model_identity as manager_module class _StopAfterValidation(Exception): pass diff --git a/tests/unit/test_prefix_cache_identity_ledger.py b/tests/unit/test_prefix_cache_identity_ledger.py index c88a6b8c..c9b1a750 100644 --- a/tests/unit/test_prefix_cache_identity_ledger.py +++ b/tests/unit/test_prefix_cache_identity_ledger.py @@ -10,7 +10,7 @@ from frontier.kv_cache.base_kv_cache_manager import KVCacheManager from frontier.kv_cache.kv_cache_block_pool import BlockPool from frontier.scheduler.replica_scheduler import ( - vllm_v1_engine_replica_scheduler as scheduler_module, + vllm_v1_prefix_cache as prefix_cache_module, ) from frontier.scheduler.replica_scheduler.vllm_v1_engine_replica_scheduler import ( VLLMv1EngineReplicaScheduler, @@ -104,7 +104,7 @@ def test_committed_full_hit_admission_records_reuse_eviction_and_rebinding( events: list[dict[str, object]] = [] monkeypatch.setattr( - scheduler_module, + prefix_cache_module, "_log_frontier_vllm_v1_schedule_decision", lambda event: events.append(dict(event)), ) diff --git a/tests/unit/test_profiling_governance_minimal_red.py b/tests/unit/test_profiling_governance_minimal_red.py index 260a3217..a6ed0ffb 100644 --- a/tests/unit/test_profiling_governance_minimal_red.py +++ b/tests/unit/test_profiling_governance_minimal_red.py @@ -1378,7 +1378,7 @@ def _train_single(**kwargs): def test_routed_and_shared_training_carry_distinct_contracts(tmp_path) -> None: """Routed and shared-expert training use their own typed domains.""" - import frontier.execution_time_predictor.shared_prediction_model_manager as module + import frontier.execution_time_predictor.prediction_family_trainers as module linear_file = tmp_path / "linear.csv" moe_file = tmp_path / "moe.csv" diff --git a/tests/unit/test_refactor_fidelity_gate.py b/tests/unit/test_refactor_fidelity_gate.py new file mode 100644 index 00000000..e5b9b379 --- /dev/null +++ b/tests/unit/test_refactor_fidelity_gate.py @@ -0,0 +1,496 @@ +"""The fidelity gate must not report success without successful comparisons. + +These tests exist because it did. `compare_labels` counted a case as covered +when its id appeared in both result tables, which is not the same as the case +having run, produced artifacts, and been compared. A matrix whose executions +all failed therefore satisfied every condition the gate checked and exited +zero, which `measure_commit.py` prints as ``VERDICT: IDENTICAL``. + +The second group covers provenance: one label has to describe one measurement, +because a filtered run keeps the records of the cases it did not execute while +rewriting the label-wide manifest with its own revision. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +import json +import subprocess +from pathlib import Path + +import pytest + +from tests.e2e.refactor_fidelity.cases import build_cases +from tests.e2e.refactor_fidelity.measure_commit import reuse_blocked_reason +from tests.e2e.refactor_fidelity.run_matrix import check_retained_records, compare_labels + + +ALL_CASES = build_cases() +BASELINE_REVISION = "a" * 40 +CANDIDATE_REVISION = "c" * 40 +HARNESS_REVISION = "h" * 40 + + +def _write_side( + output_root: Path, + label: str, + *, + revision: str, + failed: frozenset[str] = frozenset(), + content: str = '{"value": 1}', + content_overrides: dict[str, str] | None = None, + delete_artifacts_for: frozenset[str] = frozenset(), + digest_overrides: dict[str, str] | None = None, + dirty: bool = False, + stamped: bool = True, + clean_cache: bool = True, + case_filter: str | None = None, + cases_executed: Sequence[str] | None = None, + record_cases_executed: bool = True, +) -> Path: + """Write one synthetic label directory shaped like a real matrix run. + + `cases_executed` narrows what the manifest says the last run executed, the + way `--start`/`--limit` do without leaving a filter behind; results still + hold every case, as a merged continuation's would. `record_cases_executed` + False drops the field, as a manifest from before it existed. + """ + + label_root = output_root / label + label_root.mkdir(parents=True, exist_ok=True) + records = [] + for case in ALL_CASES: + artifact_dir = f"cases/{case.case_id}/metrics/run" + record = { + **case.as_record(), + "returncode": 1 if case.case_id in failed else 0, + "duration_seconds": 0.1, + "artifact_dir": artifact_dir, + "error": "boom" if case.case_id in failed else None, + "log_tail": "traceback" if case.case_id in failed else None, + } + if stamped: + record.update({ + "source_revision": revision, + "source_dirty": dirty, + "harness_revision": HARNESS_REVISION, + }) + if digest_overrides and case.case_id in digest_overrides: + record["case_digest"] = digest_overrides[case.case_id] + + if case.case_id in failed: + record["artifact_dir"] = None + record["artifacts"] = [] + else: + body = (content_overrides or {}).get(case.case_id, content) + record["artifacts"] = [{"name": "system_metrics.json", + "sha256": "0" * 64, "bytes": len(body)}] + if case.case_id not in delete_artifacts_for: + target = label_root / artifact_dir + target.mkdir(parents=True, exist_ok=True) + (target / "system_metrics.json").write_text(body, encoding="utf-8") + records.append(record) + + (label_root / "results.jsonl").write_text( + "".join(json.dumps(record, sort_keys=True) + "\n" for record in records), + encoding="utf-8", + ) + manifest = { + "label": label, + "repo_root": str(output_root / f"{label}-checkout"), + "git_head": revision, + "git_dirty_paths": "M frontier/config/config.py" if dirty else "", + "case_count": len(records), + "case_filter": case_filter, + "cases_executed_in_last_run": ( + list(cases_executed) if cases_executed is not None + else [case.case_id for case in ALL_CASES] + ), + "cache_clean_before_run": clean_cache, + "cache_files": ["model_abc.pkl"], + "source_revision": revision, + "source_dirty": dirty, + "harness_revision": HARNESS_REVISION, + } + if not record_cases_executed: + del manifest["cases_executed_in_last_run"] + (label_root / "manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return label_root + + +def _compare(output_root: Path, allow_partial: bool = False) -> tuple[int, dict]: + exit_code = compare_labels(argparse.Namespace( + output_root=str(output_root), + baseline_label="baseline", + candidate_label="candidate", + allow_partial=allow_partial, + )) + report = json.loads((output_root / "comparison.json").read_text(encoding="utf-8")) + return exit_code, report + + +# --- R34-01: the gate cannot pass without successful comparisons ------------- + + +def test_healthy_identical_sides_pass(tmp_path: Path) -> None: + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 0 + assert report["cases_compared"] == len(ALL_CASES) + assert report["mismatched_cases"] == [] + assert report["complete_comparison"] is True + # An unfiltered full run on both sides is the one shape whose cache + # listings are compared by name. + assert report["predictor_cache_populated_cleanly"] is True + assert report["predictor_cache_compared"] is True + + +def test_healthy_content_mismatch_fails(tmp_path: Path) -> None: + changed = ALL_CASES[0].case_id + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION, + content_overrides={changed: '{"value": 2}'}) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 1 + assert [entry["case_id"] for entry in report["mismatched_cases"]] == [changed] + assert report["cases_compared"] == len(ALL_CASES) + + +def test_every_case_failing_on_both_sides_is_not_a_pass(tmp_path: Path) -> None: + """The exact false success this gate was missing.""" + + everything = frozenset(case.case_id for case in ALL_CASES) + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION, failed=everything) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION, failed=everything) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 1 + assert report["cases_compared"] == 0 + assert report["complete_comparison"] is False + assert len(report["baseline_failures"]) == len(ALL_CASES) + + +def test_baseline_failure_hiding_a_different_candidate_failure_fails(tmp_path: Path) -> None: + broken = frozenset({ALL_CASES[0].case_id}) + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION, failed=broken) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION, failed=broken) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 1 + assert report["cases_compared"] == len(ALL_CASES) - 1 + assert report["baseline_failures"][0]["candidate_also_failed"] is True + + +def test_baseline_failure_with_candidate_success_is_not_full_equality(tmp_path: Path) -> None: + broken = frozenset({ALL_CASES[0].case_id}) + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION, failed=broken) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 1 + assert report["cases_not_compared"] == [ALL_CASES[0].case_id] + assert report["baseline_failures"][0]["candidate_also_failed"] is False + + +@pytest.mark.parametrize("sides", [("baseline",), ("candidate",), ("baseline", "candidate")]) +def test_missing_artifact_directory_is_missing_evidence(tmp_path: Path, sides) -> None: + gone = frozenset({ALL_CASES[0].case_id}) + for label, revision in (("baseline", BASELINE_REVISION), ("candidate", CANDIDATE_REVISION)): + _write_side(tmp_path, label, revision=revision, + delete_artifacts_for=gone if label in sides else frozenset()) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 1 + assert [entry["case_id"] for entry in report["cases_with_missing_evidence"]] == list(gone) + assert report["cases_compared"] == len(ALL_CASES) - 1 + assert report["cases_not_compared_without_explanation"] == [] + + +def test_a_case_absent_from_both_sides_needs_allow_partial(tmp_path: Path) -> None: + for label, revision in (("baseline", BASELINE_REVISION), ("candidate", CANDIDATE_REVISION)): + label_root = _write_side(tmp_path, label, revision=revision) + results = label_root / "results.jsonl" + kept = [line for line in results.read_text(encoding="utf-8").splitlines() + if json.loads(line)["case_id"] != ALL_CASES[0].case_id] + results.write_text("\n".join(kept) + "\n", encoding="utf-8") + + assert _compare(tmp_path)[0] == 1 + exit_code, report = _compare(tmp_path, allow_partial=True) + assert exit_code == 0 + assert report["cases_absent_from_both_sides"] == [ALL_CASES[0].case_id] + + +# --- R34-02: one label describes one measurement ----------------------------- + + +def test_same_case_id_with_a_different_definition_is_not_compared(tmp_path: Path) -> None: + renamed = ALL_CASES[0].case_id + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION, + digest_overrides={renamed: "deadbeefdeadbeef"}) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 1 + assert [entry["case_id"] for entry in report["cases_with_differing_definitions"]] == [renamed] + assert renamed not in report["identical_cases"] + + +def test_a_side_measured_dirty_is_reported(tmp_path: Path) -> None: + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION, dirty=True) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 1 + assert any("modified working tree" in finding for finding in report["provenance_findings"]) + + +def test_a_side_without_provenance_is_reported(tmp_path: Path) -> None: + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION, stamped=False) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 1 + assert any("no source provenance" in finding for finding in report["provenance_findings"]) + + +def test_an_assembled_cache_is_not_compared(tmp_path: Path) -> None: + """A filtered continuation without a cache clean cannot be compared by name.""" + + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION, + clean_cache=False, case_filter="dp_") + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 0, "an assembled cache is not by itself a fidelity failure" + assert report["predictor_cache_compared"] is False + assert report["predictor_cache_populated_cleanly"] is False + + +def _retained(tmp_path: Path, records: list[dict]) -> Path: + path = tmp_path / "results.jsonl" + path.write_text("".join(json.dumps(r, sort_keys=True) + "\n" for r in records), + encoding="utf-8") + return path + + +def _record(case_id: str, **overrides) -> dict: + record = { + "case_id": case_id, + "source_revision": BASELINE_REVISION, + "source_dirty": False, + "harness_revision": HARNESS_REVISION, + } + record.update(overrides) + return record + + +FIRST_THREE = [case.case_id for case in ALL_CASES[:3]] +ALL_BUT_FIRST_FIVE = [case.case_id for case in ALL_CASES[5:]] + + +@pytest.mark.parametrize( + "executed", + [FIRST_THREE, ALL_BUT_FIRST_FIVE], + ids=["limit_narrowed", "start_narrowed"], +) +def test_a_continuation_narrowed_without_a_filter_is_not_compared( + tmp_path: Path, executed: list[str] +) -> None: + """`--limit` and `--start` leave no filter behind, so the executed list decides. + + The label is clean and unfiltered by every other field, and its results + hold the whole table from the merge; only the executed ids say the cache + was populated by part of it. + """ + + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION, + cases_executed=executed) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 0 + assert report["cases_compared"] == len(ALL_CASES) + assert report["predictor_cache_populated_cleanly"] is False + assert report["predictor_cache_compared"] is False + + +def test_two_sides_narrowed_the_same_way_are_still_not_compared(tmp_path: Path) -> None: + """Symmetric partial caches would compare equal for the wrong reason.""" + + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION, + cases_executed=FIRST_THREE) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION, + cases_executed=FIRST_THREE) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 0 + assert report["predictor_cache_compared"] is False + assert report["predictor_cache_files_only_in_baseline"] == [] + assert report["predictor_cache_files_only_in_candidate"] == [] + + +def test_a_manifest_without_the_executed_list_is_not_compared(tmp_path: Path) -> None: + """An older manifest cannot show a full run, so it is not given the benefit.""" + + _write_side(tmp_path, "baseline", revision=BASELINE_REVISION, + record_cases_executed=False) + _write_side(tmp_path, "candidate", revision=CANDIDATE_REVISION) + + exit_code, report = _compare(tmp_path) + + assert exit_code == 0 + assert report["predictor_cache_populated_cleanly"] is False + assert report["predictor_cache_compared"] is False + + +def test_filtered_continuation_on_the_same_source_is_allowed(tmp_path: Path) -> None: + known = {case.case_id for case in ALL_CASES} + first, second = ALL_CASES[0].case_id, ALL_CASES[1].case_id + path = _retained(tmp_path, [_record(first), _record(second)]) + + retained, conflicts, stale = check_retained_records( + path, + {"source_revision": BASELINE_REVISION, "source_dirty": False, + "harness_revision": HARNESS_REVISION}, + {second}, + known, + ) + + assert conflicts == [] + assert stale == [] + assert set(retained) == {first, second} + + +def test_continuation_on_a_different_revision_conflicts(tmp_path: Path) -> None: + known = {case.case_id for case in ALL_CASES} + first, second = ALL_CASES[0].case_id, ALL_CASES[1].case_id + path = _retained(tmp_path, [_record(first), _record(second)]) + + _, conflicts, _ = check_retained_records( + path, + {"source_revision": CANDIDATE_REVISION, "source_dirty": False, + "harness_revision": HARNESS_REVISION}, + {second}, + known, + ) + + assert conflicts == [first], "the case this run did not re-execute is the one at risk" + + +def test_continuation_after_the_source_became_dirty_conflicts(tmp_path: Path) -> None: + known = {case.case_id for case in ALL_CASES} + first, second = ALL_CASES[0].case_id, ALL_CASES[1].case_id + path = _retained(tmp_path, [_record(first), _record(second)]) + + _, conflicts, _ = check_retained_records( + path, + {"source_revision": BASELINE_REVISION, "source_dirty": True, + "harness_revision": HARNESS_REVISION}, + {second}, + known, + ) + + assert conflicts == [first] + + +def test_unstamped_retained_records_conflict(tmp_path: Path) -> None: + known = {case.case_id for case in ALL_CASES} + first = ALL_CASES[0].case_id + path = _retained(tmp_path, [{"case_id": first}]) + + _, conflicts, _ = check_retained_records( + path, + {"source_revision": BASELINE_REVISION, "source_dirty": False, + "harness_revision": HARNESS_REVISION}, + set(), + known, + ) + + assert conflicts == [first] + + +def test_a_retained_record_outside_the_case_table_is_stale(tmp_path: Path) -> None: + known = {case.case_id for case in ALL_CASES} + path = _retained(tmp_path, [_record("coloc_dense_offline_renamed_away")]) + + _, conflicts, stale = check_retained_records( + path, + {"source_revision": BASELINE_REVISION, "source_dirty": False, + "harness_revision": HARNESS_REVISION}, + set(), + known, + ) + + assert stale == ["coloc_dense_offline_renamed_away"] + assert conflicts == [] + + +# --- R34-02: a reused detached checkout must be clean ------------------------ + + +@pytest.fixture() +def tiny_repo(tmp_path: Path) -> tuple[Path, str]: + repo = tmp_path / "repo" + repo.mkdir() + def git(*args: str) -> str: + return subprocess.run(["git", "-C", str(repo), *args], + capture_output=True, text=True, check=True).stdout.strip() + git("init", "-q") + git("config", "user.email", "t@example.com") + git("config", "user.name", "t") + (repo / "source.py").write_text("value = 1\n", encoding="utf-8") + git("add", "source.py") + git("commit", "-q", "-m", "initial") + return repo, git("rev-parse", "HEAD") + + +def test_a_clean_checkout_at_the_right_commit_may_be_reused(tiny_repo) -> None: + repo, head = tiny_repo + assert reuse_blocked_reason(repo, head) is None + + +def test_a_checkout_at_another_commit_is_refused(tiny_repo) -> None: + repo, _ = tiny_repo + reason = reuse_blocked_reason(repo, "f" * 40) + assert reason is not None and "not " + "f" * 40 in reason + + +def test_a_modified_checkout_is_refused_even_at_the_right_commit(tiny_repo) -> None: + repo, head = tiny_repo + (repo / "source.py").write_text("value = 2\n", encoding="utf-8") + + reason = reuse_blocked_reason(repo, head) + + assert reason is not None + assert "working tree has been modified" in reason + assert "source.py" in reason + assert (repo / "source.py").read_text(encoding="utf-8") == "value = 2\n", ( + "the check must not clean the checkout to make itself pass" + ) + + +def test_an_untracked_source_file_also_refuses_reuse(tiny_repo) -> None: + repo, head = tiny_repo + (repo / "extra_module.py").write_text("value = 3\n", encoding="utf-8") + + reason = reuse_blocked_reason(repo, head) + + assert reason is not None and "extra_module.py" in reason diff --git a/tests/unit/test_shared_prediction_model_manager_eager_attention_decode.py b/tests/unit/test_shared_prediction_model_manager_eager_attention_decode.py index f8d295ed..d535946e 100644 --- a/tests/unit/test_shared_prediction_model_manager_eager_attention_decode.py +++ b/tests/unit/test_shared_prediction_model_manager_eager_attention_decode.py @@ -8,7 +8,7 @@ from frontier.attention.families import DENSE_ATTENTION_FAMILY from frontier.attention.ops import AttentionOperatorRole from frontier.config import ReplicaConfig -from frontier.execution_time_predictor import shared_prediction_model_manager +from frontier.execution_time_predictor import prediction_family_trainers from frontier.execution_time_predictor import sklearn_execution_time_predictor from frontier.execution_time_predictor.shared_prediction_model_manager import ( ExecutionTimePredictionModelManager, @@ -158,7 +158,7 @@ def test_shared_manager_dense_physical_attention_models_follow_family_mapping( manager._get_attention_df_with_derived_features = lambda df: df # type: ignore[attr-defined] monkeypatch.setattr( - shared_prediction_model_manager, + prediction_family_trainers, "get_enabled_predictor_metric_names", lambda family: ( ("catalog_kv", "catalog_prefill", "catalog_decode") @@ -167,7 +167,7 @@ def test_shared_manager_dense_physical_attention_models_follow_family_mapping( ), ) monkeypatch.setattr( - shared_prediction_model_manager, + prediction_family_trainers, "get_enabled_predictor_metric_name_by_role", lambda family, role: ( { @@ -180,7 +180,7 @@ def test_shared_manager_dense_physical_attention_models_follow_family_mapping( ), ) monkeypatch.setattr( - shared_prediction_model_manager, + prediction_family_trainers, "get_enabled_predictor_median_columns", lambda family: ( ( @@ -193,7 +193,7 @@ def test_shared_manager_dense_physical_attention_models_follow_family_mapping( ), ) monkeypatch.setattr( - shared_prediction_model_manager, + prediction_family_trainers, "get_enabled_shared_predictor_feature_columns", lambda family: ( { @@ -293,7 +293,7 @@ def test_shared_manager_dense_physical_attention_roles_do_not_depend_on_family_o manager._get_attention_df_with_derived_features = lambda df: df # type: ignore[attr-defined] monkeypatch.setattr( - shared_prediction_model_manager, + prediction_family_trainers, "get_enabled_predictor_metric_names", lambda family: ( ("role_prefill", "role_decode", "role_cache") @@ -302,7 +302,7 @@ def test_shared_manager_dense_physical_attention_roles_do_not_depend_on_family_o ), ) monkeypatch.setattr( - shared_prediction_model_manager, + prediction_family_trainers, "get_enabled_predictor_metric_name_by_role", lambda family, role: ( { @@ -315,7 +315,7 @@ def test_shared_manager_dense_physical_attention_roles_do_not_depend_on_family_o ), ) monkeypatch.setattr( - shared_prediction_model_manager, + prediction_family_trainers, "get_enabled_predictor_median_columns", lambda family: ( ( @@ -328,7 +328,7 @@ def test_shared_manager_dense_physical_attention_roles_do_not_depend_on_family_o ), ) monkeypatch.setattr( - shared_prediction_model_manager, + prediction_family_trainers, "get_enabled_shared_predictor_feature_columns", lambda family: ( { diff --git a/tests/unit/test_shared_prediction_model_manager_mixed_layer_moe.py b/tests/unit/test_shared_prediction_model_manager_mixed_layer_moe.py index 8ace3473..147fffe4 100644 --- a/tests/unit/test_shared_prediction_model_manager_mixed_layer_moe.py +++ b/tests/unit/test_shared_prediction_model_manager_mixed_layer_moe.py @@ -15,7 +15,7 @@ def test_mixed_layer_moe_materializes_moe_and_dense_mlp_predictors( monkeypatch, tmp_path ) -> None: """Mixed-layer MoE training must expose both runtime FFN branches.""" - import frontier.execution_time_predictor.shared_prediction_model_manager as manager_module + import frontier.execution_time_predictor.prediction_family_trainers as manager_module linear_file = tmp_path / "linear_op.csv" moe_file = tmp_path / "moe.csv" @@ -154,7 +154,7 @@ def test_mixed_layer_dense_training_does_not_reorder_legacy_share_expert_models( monkeypatch, tmp_path ) -> None: """Mixed-layer additions must not perturb legacy RF training order.""" - import frontier.execution_time_predictor.shared_prediction_model_manager as manager_module + import frontier.execution_time_predictor.prediction_family_trainers as manager_module linear_file = tmp_path / "linear_op.csv" moe_file = tmp_path / "moe.csv"