Language Models and Search Space Reduction

Introduction

Since publishing our paper and discussing the algorithm with other engineers and researchers, one recurring question has been why we chose to use a language model at the core of the proposal engine. Given the nature of the optimization problem, why not formulate it as an integer linear program (ILP)?

In this post, I want to examine the problem more formally, explain the limitations of a purely ILP-based approach, and describe why we chose a language model for configuration search.

The distinction comes down to two separate problems: identifying promising deployment configurations and selecting an allocation from a known set of candidates. An ILP is well suited to the latter. The former is considerably more difficult: the configuration space is combinatorial, serving performance depends on nonlinear interactions between parameters, and the relationship between deployment decisions and performance outcomes is not fully known a priori.

One of our primary challenges was therefore that of search-space reduction: how do we efficiently identify a small set of promising configurations from an enormous space of possibilities, within the time constraints of a continuously operating cluster planner?

Koi addresses this by combining a language model for candidate generation, a causal structure to guide its search, and deterministic components to predict, score, and select cluster-wide deployment plans.

1. The parameter space

Deploying an LLM for inference involves substantially more than selecting a GPU and loading model weights. Across the different layers of the inference stack, a deployment configuration spans 140+ parameters, including:

Each parameter introduces additional possible configurations. If parameter i has kᵢ possible values, the size of the unconstrained configuration space is:

$$N = \prod_i k_i$$

The search space grows multiplicatively, not additively. Even under the simplified assumption that each of the 140 parameters has only two possible values, the Cartesian product contains 2¹⁴⁰ possible combinations.

This is an illustrative counting argument. Not every combination is physically valid, and some parameters only apply when particular features are enabled. Nevertheless, the feasible configuration space is far too large to enumerate exhaustively.

At cluster scale, the problem grows further. For a cluster running J jobs, each with a candidate configuration set Cⱼ, the unconstrained number of possible deployment plans is:

$$N_{\text{cluster}} = \prod_{j=1}^{J} |C_j|$$

These decisions are also coupled through shared GPU capacity, quotas, placement restrictions, and competing workload objectives. A configuration that is optimal for one model in isolation may be suboptimal for the cluster if it consumes hardware required by another model, which is usually the case!

2. Parameters are coupled

The difficulty is not simply the number of combinations. Configuration parameters interact, and their effects on serving performance are often nonlinear.

Consider several common parameters that an operator might adjust when deploying a model.

Increasing tensor parallelism (TP) distributes model weights and computation across more GPUs, reducing per-device memory requirements but introducing drawbacks in the form of communication overhead. Whether increasing TP improves throughput depends on the interconnect, model architecture, GPU memory bandwidth, and batch size, amongst other things.

Pipeline parallelism (PP) similarly introduces different tradeoffs. Increasing PP reduces the number of layers assigned to each stage but can introduce pipeline bubbles and stage imbalance. The resulting efficiency depends on microbatching, communication overhead, and workload characteristics.

Batching interacts with both. Larger batches can improve accelerator utilization but increase memory pressure and queueing latency. The resulting KV-cache pressure depends on request concurrency, context lengths, cache precision, and the memory which remains after allocation to the model weights.

Quantization further changes the feasible configuration space. Reducing weight precision may allow a model to fit on fewer GPUs, enabling parallelism configurations that were previously infeasible. However, its performance depends on kernel support, quantization method, hardware, and workload shape.

These interactions make independent parameter optimization insufficient. The best TP configuration on an NVLink-connected node may not be the best configuration across nodes connected by a slower fabric. Similarly, a configuration optimized for compute-heavy prefill may perform poorly under memory-bandwidth-intensive decoding.

We observed these differences in our experiments with Llama 3 70B on AWS A10G GPUs. Across five feasible parallelism configurations, throughput varied by 4.4× and cost per token by 7.5× for the same model and workload. The search space is therefore not only combinatorial. Its performance landscape is defined by interacting parameters whose effects vary across hardware, software, and workload conditions.

3. Why not an ILP?

Integer linear programming (ILP) is a mathematical optimization technique for finding the best combination of discrete decisions subject to linear constraints. It is widely used in scheduling, resource allocation, and placement problems, making it a natural candidate for inference orchestration.

An ILP requires three components: decision variables, an objective function to maximize or minimize, and constraints defining which combinations are feasible. Given these inputs, a solver searches for the best feasible assignment. There is already research demonstrating the effectiveness of this approach for inference orchestration. One example is SageServe: Serving Models, Fast and Slow out of Microsoft, which uses an ILP to optimize GPU instance allocation across models and cloud regions.

How ILP works for inference orchestration

SageServe addresses a resource allocation problem: given a fleet of GPUs across multiple regions and changing interactive and batch workloads, how many instances of each model should run on each GPU type, and where?

The system uses an AutoRegressive Integrated Moving Average (ARIMA) model to forecast inference demand in tokens per second. Its ILP then determines how many model instances to add or remove in each region, minimizing provisioning overhead while maintaining sufficient capacity for forecasted interactive demand.

Consider a simplified example with illustrative throughput and cost values.

Suppose we have one LLM that can run on two GPU types:

• An A100 instance provides 400 tokens/sec at a provisioning cost of $2.
• An H100 instance provides 700 tokens/sec at a provisioning cost of $5.

Assume we already have one A100 instance running and forecast demand increasing to 1,000 tokens/sec.

We could provision two additional A100 instances, bringing total capacity to 1,200 tokens/sec at an additional cost of $4. Alternatively, we could provision one H100 alongside the existing A100, providing 1,100 tokens/sec at a cost of $5.

An ILP can determine that the first option minimizes provisioning cost while satisfying the capacity requirement.

SageServe generalizes this across multiple models, GPU types, and regions. Its decision variables represent changes in instance counts, its objective minimizes VM acquisition and model-loading overhead, and its constraints ensure sufficient capacity for forecasted traffic. The ILP runs hourly, while separate routing and scaling heuristics respond to shorter-term changes.

This formulation works because SageServe assumes that the throughput of a particular model on a particular GPU type is already known. It represents this capacity as a fixed coefficient, allowing the solver to calculate aggregate serving capacity by summing the contributions of individual instances.

The important distinction is that SageServe's ILP optimizes how many instances to allocate, not how those instances should be configured internally. Koi operates over a substantially larger decision space. We are not only determining how many instances to allocate, but also their hardware, parallelism, serving engine, batching behavior, memory configuration, and other deployment parameters outlined earlier.

There are two principal ways we could attempt to formulate this problem as an ILP.

The problem of candidate generation

The first approach would be to represent individual configuration parameters directly as decision variables: GPU type, TP, PP, batch size, KV-cache settings, quantization, and so on.

However, deployment performance is not a linear combination of these variables. Increasing TP may improve throughput on one interconnect and degrade it on another. Quantization changes memory requirements, potentially enabling different parallelism configurations, which in turn alter communication overhead and batching capacity.

To capture these dependencies in an ILP, we would need additional variables and constraints to represent parameter interactions, or construct piecewise-linear approximations of the underlying performance relationships. This is possible in principle, but requires a sufficiently accurate model of how these parameters interact and affect serving performance.

The second approach would be to enumerate complete deployment configurations beforehand and treat each as a discrete candidate. For every candidate, we would estimate its GPU requirements, throughput, latency, and cost. The ILP would then select the best combination across jobs.

This avoids explicitly modeling parameter interactions within the ILP, but shifts the difficulty to generating and evaluating the candidate catalog itself (which brings us full circle to the problem we would employ ILP to solve).

As discussed earlier, the configuration space grows multiplicatively across approximately 140 parameters. Exhaustive enumeration is impractical, even after excluding physically invalid combinations. Hardware heterogeneity, new accelerators, and serving-stack updates continually introduce additional combinations.

An ILP solver does not necessarily enumerate every possible cluster-wide allocation. Modern solvers use techniques such as branch-and-bound to reduce the search. However, they still require a useful representation of the available choices and estimates of their performance.

The problem of performance estimation

Even with a manageable candidate catalog, we still need reliable performance estimates for each configuration.

Consider a 70B model deployed with TP4/PP2 versus TP8/PP1. An ILP can calculate which deployment consumes fewer GPUs or incurs a lower reservation cost. It cannot independently determine which will achieve better throughput, TTFT, or TPOT.

These outcomes depend on collective communication, pipeline efficiency, KV-cache pressure, batching dynamics, and queueing. Their relationships vary across model architectures, accelerators, serving engines, and workload distributions. Furthermore, these estimates are not stationary. Software updates, changes in request distributions, and shifts in cluster conditions can invalidate previously measured performance characteristics.

This is a separate modeling problem. In Koi, we address it using a deterministic performance surrogate calibrated against observed deployment behavior, a topic that warrants its own technical discussion and something we are actively working on (open positions for this on our careers page if interested).

Where ILP remains useful

None of this makes ILP unsuitable for inference orchestration. SageServe demonstrates its usefulness when the primary decisions are instance counts and placement, and the relevant performance coefficients are available.

Similarly, once we generate a small set of promising deployment configurations, estimate their performance, and represent their resource requirements, integer optimization becomes a reasonable approach to cluster-wide allocation.

We evaluated this approach in the paper through a rolling CP-SAT baseline. At each planning interval, it enumerates policy-valid placement candidates and optimizes priority-weighted admission, reservation cost, and GPU footprint under capacity constraints.

Unlike Koi, this baseline does not learn deployment performance from runtime evidence or optimize the same full set of serving-quality objectives. Its performance therefore cannot be attributed solely to the choice of integer optimization.

The limitation we are addressing is more specific: an ILP can optimize allocation among known candidates, but it does not independently solve the problem of discovering promising configurations or predicting how they will perform.

This brings us back to the initial search problem. Given an enormous space of interacting deployment parameters, how do we identify promising configurations within a practical planning budget?

Our architectural choice was to use a language model for this initial search reduction. Rather than exhaustively enumerating configurations, Koi uses a language model to reason over workload characteristics and a causal graph connecting deployment decisions to runtime bottlenecks and serving outcomes.

The resulting candidates are passed to deterministic components for feasibility checking, performance prediction, scoring, and cluster-wide selection. The language model does not replace mathematical optimization. It provides a search heuristic for reducing a poorly characterized configuration space into a tractable set of candidates.

4. Causal-guided search

This raises an important question which I’m sure you’re asking: is Koi simply asking a language model to propose configurations based on its existing systems and domain knowledge? Not quite. Koi separates candidate generation from deterministic evaluation and uses an explicit causal structure to guide the language model's search.

Rather than enumerating configurations or asking a language model to predict their performance directly, we represent how deployment decisions influence serving outcomes through measurable runtime behavior.

We model the inference system as a three-layer directed acyclic graph (DAG) containing:

• Decisions (X): GPU selection, parallelism, batching, quantization, engine configuration.
• Runtime mediators (V): KV-cache pressure, pipeline bubbles, communication latency, memory pressure, SM utilization.
• Outcomes (Y): Throughput, TTFT, TPOT, cost per token, SLO margin.

The graph is bipartite and follows the structure: X → V → Y

This representation connects configuration decisions to performance outcomes through runtime mechanisms. For example, increasing pipeline parallelism may introduce pipeline bubbles that reduce throughput. Increasing batch concurrency may raise KV-cache utilization, potentially degrading TPOT. Koi maintains confidence and uncertainty for relationships in this graph. Related edges form scoped mechanisms: testable hypotheses about serving behavior that apply to specific models, hardware, workloads, and operating conditions.

The language model can query this graph backward from a serving objective to identify relevant configuration decisions.

For example, if a workload is violating its TPOT target, the planner can trace the runtime mediators associated with TPOT and identify which deployment parameters influence them. Rather than searching across every configuration parameter, it can focus on decisions relevant to the suspected bottleneck.

This is where the language model becomes useful. It combines qualitative systems knowledge, structured causal relationships, workload characteristics, and observed behavior to generate a small set of plausible deployment configurations within a limited search budget.

Koi organizes this process through a Recursive Language Model (RLM) hierarchy. A root planner considers cluster-wide resource availability and assigns soft hardware budgets to individual jobs. Per-job specialist agents then use workload characteristics and applicable causal mechanisms to propose configurations within those budgets.

The language model is not guaranteed to find the globally optimal configuration. Its purpose is to identify promising candidates efficiently in a large, changing search space, leaving their evaluation to deterministic components.

5. Proposal vs decision

A critical distinction in Koi's architecture is that the language model does not determine which configurations are deployed. In fact, we argue that a language model has no ability to predict serving outcomes, and reliance on domain knowledge alone is insufficient for this task. Thus, its proposals are advisory and can be rejected by deterministic evaluation.

Each candidate passes through components responsible for validating feasibility, predicting performance, and selecting a cluster-wide allocation.

First, candidates are checked for physical feasibility, including VRAM requirements, parallelism compatibility, available capacity, and deployment policies.

Next, a performance surrogate built on NVIDIA AIConfigurator and execution replay predicts serving behavior, including TTFT, TPOT, and throughput. Accurate performance modeling remains an open challenge, and improving these surrogates is an active area of our work.

These predictions are combined with allocation costs and job-specific objectives. Because surrogate estimates are imperfect, Koi tracks historical prediction errors and accounts for uncertainty when estimating SLO violation risk.

A cluster-wide selector evaluates proposed configurations using four components:

• Exploitation: Predicted performance against each job's weighted serving objectives.
• Exploration: The value of obtaining evidence about uncertain causal relationships.
• Risk: Estimated probability of violating serving constraints.
• Reconfiguration cost: Model loading, parallel execution during transitions, draining, and transient performance degradation.

Selection is subject to hard resource, physical feasibility, SLO-risk, swap-budget, and job-policy constraints. The language model cannot override these constraints or the deterministic scores.

At each planning interval, Koi evaluates proposed cluster plans and returns diagnostics for invalid or non-improving candidates. The planner can use this feedback to generate additional proposals within its search budget. Koi then commits the highest-scoring feasible plan found, rather than claiming a globally optimal solution.

Once deployments are running, Koi collects telemetry and compares observed behavior against surrogate predictions. Persistent deviations are detected using CUSUM tests, providing evidence that supports or contradicts applicable causal mechanisms.

These observations update the causal evidence state and inform subsequent planning intervals. The result is a closed-loop planner in which every deployment contributes evidence that can improve future configuration search.

Conclusion

The choice between a language model and an ILP is not a choice between probabilistic reasoning and mathematical optimization. They address different parts of the problem.

Integer optimization is effective when selecting among candidates with known constraints and reliable objective estimates. Koi's more difficult problem is identifying which configurations are worth considering when the search space is combinatorial and the performance relationships are nonlinear, uncertain, and changing.

The additional constraint is computational: the planner must generate and evaluate candidates within a bounded time while continuously adapting to new deployment evidence.

We use a language model to reduce the search space through structured causal reasoning, deterministic components to evaluate and select deployment plans, and runtime evidence to continuously improve both the predictions and the search.