Katkat
Back

Schedule Large and Small Model Inference on an Eight-GPU Pool

AnthropicMachine Learning Engineer · Technical Screen
HardMachine Learning EngineerTechnical ScreenML System Design

Last updated

Design an inference scheduler for a shared 8-GPU node that co-locates single-GPU small models and all-reduce 8-GPU large models. The design must address gang scheduling, prefix cache reuse, deadline-aware priorities, and overload mitigation.

Scenario

You are tasked with designing the node-level execution and scheduling layer for an inference worker equipped with 8 interconnected GPUs (e.g., 8x H100s connected via NVLink/NVSwitch).

The node serves two types of models:

  1. Small Models (SM): Optimized for 1 GPU each (TP=1). Up to 8 SM instances can run concurrently, or multiple SM requests can be packed onto available individual GPUs.
  2. Large Models (LM): Require all 8 GPUs running in lockstep using Tensor Parallelism (TP=8).

Because LM requires all 8 GPUs simultaneously, it requires a gang scheduling reservation. However, constantly reserving all 8 GPUs starves SM workloads, while naive greedy scheduling of SM workloads starves LM requests indefinitely.

Requirements

Functional Requirements

  • Gang Reservations & Co-scheduling: Coordinate execution of TP=8 LM tasks alongside independent TP=1 SM tasks without deadlock or perpetual starvation.
  • Cache-Safe Scheduling: Support prefix-cache locality for both model tiers to maximize prompt KV-cache reuse across batched requests.
  • Deadline & SLA Awareness: Support latency-sensitive requests with strict deadlines (e.g., time-to-first-token [TTFT] <= 200ms) and background/throughput queries.
  • Overload Handling & Admission Control: Apply backpressure, request shed, or priority preemption when the system saturates.

Non-Functional Requirements & Scale

  • Node Capacity: 8 GPUs, 80 GB VRAM each.
  • Throughput: Support peak load of ~1,500 SM queries/sec or ~100 LM queries/sec.
  • LM Iteration Time: 20-50ms per decoding step.
  • SM Iteration Time: 5-15ms per decoding step.
  • Context Lengths: Up to 32k tokens per sequence. Efficient PagedAttention / chunked prefill integration.

What Your Design Should Cover

  1. Queue Architecture & Scheduling Algorithm: How jobs are enqueued, prioritized, and dispatched. Explain how you prevent head-of-line blocking and model starvation (e.g., aging, reservation windows, or virtual-time round-robin).
  2. State & Preemption Management: How the scheduler handles running SM jobs when an LM gang reservation is triggered (e.g., drain vs. pause/swap vs. recompute).
  3. Cache & Memory Routing: How KV-cache affinity interacts with the scheduling policy (e.g., dispatching SM requests to specific GPUs holding matching prefix trees).
  4. Failure Modes & Overload Policy: Mechanics of rate-limiting, degradation modes, and backpressure propagation back to the cluster routing layer.

Hints

Hint 1

Consider modeling LM execution as periodic fixed-horizon gang leases and using chunked prefill so running SM batches can cleanly yield or drain quickly without aborting.

Hint 2

Prefix caching requires sticky routing at the node dispatcher; balance the tradeoff between routing to a warm-cache GPU versus routing to an idle GPU to minimize total TTFT.

Hint 3

To prevent starvation without completely pausing SM traffic, explore token-budgeting or virtual finish-time scheduling between the LM queue and the aggregate SM pool.

Follow-up questions

  • How does chunked prefill change the preemption trade-off between draining an SM job versus checkpointing its KV-cache to host RAM?
  • If prefix cache hits reduce prefill time by 80%, how do you formulate the threshold at which waiting for a specific busy GPU is faster than scheduling immediately on a cold GPU?
  • How would your scheduler adapt if network routing across nodes enabled swapping SM workloads to other hosts in the cluster during an LM reservation burst?