RFC: Better API and control over deciding which worker is selected.

[michaelfeil]

We are looking to improve the routing logic.

Specifically the lines:

dynamo/lib/llm/src/kv_router/scheduler.rs

Lines 470 to 521 in 0b7e127

	for (worker_id, config) in workers.iter() {
	// Get data_parallel_size from runtime config
	// data_parallel_size defaults to 1 in ModelRuntimeConfig
	let data_parallel_size = config.as_ref().map(|c| c.data_parallel_size).unwrap_or(1); // Fallback if config is None
	// Iterate over all dp_ranks for this worker
	for dp_rank in 0..data_parallel_size {
	let worker = WorkerWithDpRank::new(*worker_id, dp_rank);
	// Get overlap for this worker (defaults to 0 if not in overlaps)
	let overlap = *overlaps.get(&worker).unwrap_or(&0);
	// this is the number of prefill tokens the worker would have if the request were scheduled there
	let prefill_token = *prefill_tokens.get(&worker).unwrap_or(&isl);
	let potential_prefill_block = (prefill_token as f64) / (block_size as f64);
	// this is the number of decode blocks the worker would have if the request were scheduled there
	let decode_block = *decode_blocks
	.get(&worker)
	.unwrap_or(&(potential_prefill_block.floor() as usize))
	as f64;
	// Use override if provided, otherwise use default config
	let overlap_weight = request
	.router_config_override
	.as_ref()
	.and_then(|cfg| cfg.overlap_score_weight)
	.unwrap_or(self.kv_router_config.overlap_score_weight);
	// Calculate logit (lower is better)
	let logit = overlap_weight * potential_prefill_block + decode_block;
	worker_logits.insert(worker, logit);
	tracing::info!(
	"Formula for worker_id={} dp_rank={:?} with {overlap} cached blocks: {logit:.3} \
	= {overlap_weight:.1} * prefill_blocks + decode_blocks \
	= {overlap_weight:.1} * {potential_prefill_block:.3} + {decode_block:.3}",
	worker.worker_id,
	worker.dp_rank
	);
	}
	}
	// Use softmax sampling to select worker(s)
	// Use override if provided, otherwise use default config
	let temperature = request
	.router_config_override
	.as_ref()
	.and_then(|cfg| cfg.router_temperature)
	.unwrap_or(self.kv_router_config.router_temperature);
	let candidates = softmax_sample(&worker_logits, temperature);

• Under e.g. high-kv-cache hit rates, the router is super flickery and will shift traffic around.
• Will over-subscribe a Worker, because since dynamo 0.6, we added DP aware routing. The result is that e.g. all dp_ranks are seen individual and are not getting traffic.
  Understood, there is a full API for the worker_selector, but the e.g. softmax sample and the following lines are helpful.
  If I want to implement various versions of it, it would be helpful to

[rmccorm4]

CC @PeaBrane @nnshah1