calc_runtime_stats

Runtime statistics calculation for NAS subblock benchmarking via vLLM.

Functions

calc_runtime_for_blocks	Benchmark each full decoder block (attention + FFN together) and return block runtimes.
calc_runtime_for_subblocks	Benchmark each unique subblock and return per-subblock runtimes and no-block overhead.
create_benchmark_model	Build a descriptor-specific model from an exact immutable block layout.
enumerate_runtime_block_configs	Return deterministic unique runtime candidates from a converted teacher checkpoint.

calc_runtime_for_blocks(block_config_set, runtime_stats_config, vocab_size, hidden_size, num_attention_heads, num_key_value_heads, descriptor, lm_config, tokenizer_path, prefill_seq_len, generation_seq_len, batch_size, cache_dir=None)

Benchmark each full decoder block (attention + FFN together) and return block runtimes.

Unlike calc_runtime_for_subblocks() which times attention and FFN independently and sums them, this function times the full (attention + FFN) block in a single vLLM call. This is more accurate because it captures kernel-fusion and memory-bandwidth-reuse effects that the sum of independent measurements misses.

Each unique BlockConfig in block_config_set is benchmarked as a repeated-block model (same as the subblock path) and the per-block runtime is derived via the same differencing formula.

Returns (runtime_by_block_dict, no_block_runtime_ms) analogous to calc_runtime_for_subblocks().

Parameters:

• block_config_set (set[BlockConfig])

• runtime_stats_config (DictConfig)

• vocab_size (int)

• hidden_size (int)

• num_attention_heads (int)

• num_key_value_heads (int)

• descriptor (Type[ModelDescriptor])

• lm_config (Any)

• tokenizer_path (str)

• prefill_seq_len (int)

• generation_seq_len (int)

• batch_size (int)

• cache_dir (Path | None)

Return type:

tuple[dict[BlockConfig, float], float]

calc_runtime_for_subblocks(subblock_config_set, runtime_stats_config, vocab_size, hidden_size, num_attention_heads, num_key_value_heads, descriptor, lm_config, tokenizer_path, prefill_seq_len, generation_seq_len, batch_size, cache_dir=None)

Benchmark each unique subblock and return per-subblock runtimes and no-block overhead.

The distinct vLLM benchmarks are enumerated up front and run concurrently across all visible GPUs (with on-disk caching via cache_dir for resume), then the per-subblock runtimes are derived from the cached measurements using the same differencing the sequential version used.

Parameters:

• subblock_config_set (set[SubblockConfig])

• runtime_stats_config (DictConfig)

• vocab_size (int)

• hidden_size (int)

• num_attention_heads (int)

• num_key_value_heads (int)

• descriptor (Type[ModelDescriptor])

• lm_config (Any)

• tokenizer_path (str)

• prefill_seq_len (int)

• generation_seq_len (int)

• batch_size (int)

• cache_dir (Path | None)

Return type:

tuple[dict[SubblockConfig, float], float]

create_benchmark_model(runtime_config, block_layout)

Build a descriptor-specific model from an exact immutable block layout.

Parameters:

• runtime_config (RuntimeConfig)

• block_layout (tuple[BlockConfig, ...])

enumerate_runtime_block_configs(teacher_dir, descriptor, *, search_space=None, include_noops=True)

Return deterministic unique runtime candidates from a converted teacher checkpoint.

Parameters:

• teacher_dir (Path | str)

• descriptor (Type[ModelDescriptor])

• search_space (Mapping[str, Any] | None)

• include_noops (bool)

Return type:

tuple[BlockConfig, …]