llmgw/internal/router/scheduler/profiles.go

package scheduler

// ProfileConfig holds thresholds shared by the built-in profiles.
type ProfileConfig struct {
	// TargetWaiting is the queue scorer's "healthy buffer" target.
	TargetWaiting float64
	// TargetInflight normalizes the inflight scorer.
	TargetInflight float64
	// KVCeiling is the KV occupancy ceiling.
	KVCeiling float64
	// SessionLoadFactor is the load ceiling multiplier for session affinity.
	SessionLoadFactor float64

	// Prefix, when set, adds the approximate prefix tree to the profiles that
	// benefit from it. Nil leaves them exactly as they were — session affinity
	// alone, which is tier 0 and already most of the win.
	Prefix *PrefixScorer
}

// DefaultProfileConfig returns sensible starting values.
// These numbers should be measured with ctl bench (§10); these are only seeds.
func DefaultProfileConfig() ProfileConfig {
	return ProfileConfig{
		TargetWaiting:     4,
		TargetInflight:    8,
		KVCeiling:         0.9,
		SessionLoadFactor: 1.25,
	}
}

func (c *ProfileConfig) normalize() {
	d := DefaultProfileConfig()
	if c.TargetWaiting <= 0 {
		c.TargetWaiting = d.TargetWaiting
	}
	if c.TargetInflight <= 0 {
		c.TargetInflight = d.TargetInflight
	}
	if c.KVCeiling <= 0 || c.KVCeiling > 1 {
		c.KVCeiling = d.KVCeiling
	}
	if c.SessionLoadFactor < 1 {
		c.SessionLoadFactor = d.SessionLoadFactor
	}
}

// Chat is the profile for multi-turn conversations and agent workloads.
// Session affinity dominates, to preserve prefix cache reuse.
func Chat(c ProfileConfig) *Profile {
	c.normalize()
	return &Profile{
		Name:    "chat",
		Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling), AdapterAffinity(c.TargetWaiting, c.KVCeiling)},
		Scorers: c.withPrefix([]WeightedScorer{
			{Session(c.SessionLoadFactor), 3.0},
			{Queue(c.TargetWaiting), 1.0},
			{KV(c.KVCeiling), 1.0},
			{Inflight(c.TargetInflight), 1.0},
			{Weight(), 0.5},
		}, 2.0),
		Picker: P2C(),
	}
}

// withPrefix appends the prefix scorer when one is configured.
//
// Weighted below session affinity on purpose. A session id is a statement from
// the client that these requests belong together; a prefix match is the
// router's guess that they might. When the two disagree the client is more
// likely to be right.
func (c ProfileConfig) withPrefix(scorers []WeightedScorer, weight float64) []WeightedScorer {
	if c.Prefix == nil {
		return scorers
	}
	return append(scorers, WeightedScorer{c.Prefix, weight})
}

// Batch is the profile for bulk work: throughput matters, TTFT does not, so
// finding the shortest queue outranks stickiness.
func Batch(c ProfileConfig) *Profile {
	c.normalize()
	return &Profile{
		Name:    "batch",
		Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling), AdapterAffinity(c.TargetWaiting, c.KVCeiling)},
		Scorers: []WeightedScorer{
			{Queue(c.TargetWaiting), 3.0},
			{KV(c.KVCeiling), 1.0},
			{Load(30), 1.0},
			{Inflight(c.TargetInflight), 1.0},
			{Weight(), 1.0},
		},
		Picker: MaxScore(),
	}
}

// LongContext is the profile for long-prompt workloads. KV pressure is the
// bottleneck there, so the kv scorer dominates.
func LongContext(c ProfileConfig) *Profile {
	c.normalize()
	return &Profile{
		Name:    "long-context",
		Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling), AdapterAffinity(c.TargetWaiting, c.KVCeiling)},
		// Prefix matching earns more here than anywhere else: a long prompt is
		// mostly prefill, and prefill is what a cache hit skips.
		Scorers: c.withPrefix([]WeightedScorer{
			{KV(c.KVCeiling), 3.0},
			{Session(c.SessionLoadFactor), 2.0},
			{Queue(c.TargetWaiting), 1.0},
			{Inflight(c.TargetInflight), 1.0},
		}, 3.0),
		Picker: Softmax(0.2),
	}
}

// Debug is a round-robin profile for troubleshooting.
func Debug() *Profile {
	return &Profile{
		Name:    "debug",
		Filters: []Filter{BreakerOpen()},
		Picker:  RoundRobin(),
	}
}

// Builtin returns a built-in profile by name, or nil if the name is unknown.
func Builtin(name string, c ProfileConfig) *Profile {
	switch name {
	case "chat":
		return Chat(c)
	case "batch":
		return Batch(c)
	case "long-context":
		return LongContext(c)
	case "debug":
		return Debug()
	default:
		return nil
	}
}