Build the KAI-Scheduler and GPU Fractioning Forks

This guide builds two forks from source, pushes the images to a registry you control, and installs them on a Kubernetes cluster. Together they give each pod a hard share of a GPU’s streaming multiprocessors (SMs) and memory, enforced by CUDA Multi-Process Service (MPS) v3 namespaces. You need them for the “KAI + GPU fractions” experiment in Interpreting GPU Sharing Results: 16 Qwen3-4B models on 8 A100 GPUs, each model capped at 50% of the SMs and 50% of the memory of one GPU.

When the install is done, a pod that asks for gpu-fraction: "0.5" lands on a shared GPU in its own MPS namespace. The MPS server enforces the pod’s SM share, and the driver enforces its memory cap, so the pod cannot raise either one. The next step is deploying the models.

What You Are Building

pod (gpu-fraction: "0.5")
  │
  ▼
KAI-Scheduler ── picks a GPU, writes gpu-compute.portion / gpu-memory.request / gpus.devices
  │
  ▼
kai-gpu-fractioning
  ├── fractiond   reads the annotations, injects env + the namespace's MPS pipe directory
  └── mpsd        creates the MPS namespace, sets its active-thread percentage on the server

Pinned Revisions

Both changes are open pull requests against the upstream repositories. The results published for this experiment were produced by exactly these commits, so check them out by SHA rather than following the pull request.

Prerequisites

Build Machine: KAI-Scheduler

Build Machine: kai-gpu-fractioning

The images target linux/amd64 by default (PLATFORM ?= linux/amd64), which matches A100 and H100 nodes. Override PLATFORM if you build on another architecture.

Cluster and Nodes

Check every row on a GPU node before you install anything. The defaults of a stock GPU Operator install are not enough for compute enforcement.

for r in customresourcedefinitions clusterroles clusterrolebindings priorityclasses namespaces; do
  printf '%s: ' "$r"; kubectl auth can-i create "$r" --all-namespaces
done

Installing the r615 Driver

All the packages and images below exist publicly; choose the path that matches how the cluster manages drivers.

• Host-installed driver (Ubuntu 24.04, NVIDIA CUDA apt repository). The nvidia-driver-pinning-615 package exists only in the NVIDIA CUDA repository, so add that repository first. Then install the branch pin, then the driver and, on NVSwitch systems, Fabric Manager:

  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
  sudo dpkg -i cuda-keyring_1.1-1_all.deb
  sudo apt-get update
  sudo apt-get install nvidia-driver-pinning-615
  sudo apt-get install nvidia-driver-open nvidia-fabricmanager
  sudo reboot

  If a different nvidia-driver-pinning-* package is already installed, remove it first, or the 615 pin does not take effect and nvidia-driver-open resolves to the older branch. If the driver install appears to do nothing, check apt-cache policy nvidia-driver-open for an /etc/apt/preferences.d/ file that pins the NVIDIA origin to a negative priority.

• GPU Operator-managed driver. Pass the version when you install or upgrade the operator:

  helm upgrade --install gpu-operator nvidia/gpu-operator \
    -n gpu-operator --create-namespace \
    --version v26.7.1 \
    --set driver.version=615.71.09

  The driver image nvcr.io/nvidia/driver:615.71.09-ubuntu24.04 is published. Add the chart repository first with helm repo add nvidia https://helm.ngc.nvidia.com/nvidia.

Step 1: Clone and Check Out the Pinned Commits

mkdir kai-forks && cd kai-forks
git clone https://github.com/kai-scheduler/KAI-Scheduler.git
git -C KAI-Scheduler fetch origin pull/2368/head
git -C KAI-Scheduler checkout 44d1d0a966465c69b1fd0d24b40b9a1ebb73699e
git clone https://github.com/kai-scheduler/gpu-fractioning.git kai-gpu-fractioning
git -C kai-gpu-fractioning fetch origin pull/147/head
git -C kai-gpu-fractioning checkout af0544e4dbd1d8e03ca14fc1e06d6265fabed99b

Confirm that both trees are clean, so no local edit is baked into an image:

git -C KAI-Scheduler status --porcelain          # prints nothing
git -C kai-gpu-fractioning status --porcelain    # prints nothing

Step 2: Build and Push the Images

Set the registry and one tag for both forks. The tag records the pairing by naming both commits:

export REGISTRY=registry.example.com/team/kai-fork   # replace with a registry you can push to
export TAG=smshare-44d1d0a-af0544e

kai-gpu-fractioning

cd kai-gpu-fractioning
make docker-build docker-push \
  DOCKER_REPO_BASE="$REGISTRY/kai-gpu-fractioning" \
  VERSION="$TAG"
cd ..

This builds and pushes operator, mpsd, fractiond, and metricsd as $REGISTRY/kai-gpu-fractioning/<name>:$TAG. For a multi-architecture build, run make docker-buildx DOCKER_REPO_BASE="$REGISTRY/kai-gpu-fractioning" VERSION="$TAG" DOCKER_BUILD_PLATFORM=linux/amd64,linux/arm64, which pushes directly. Pass DOCKER_REPO_BASE and VERSION again: without them the target tags the images for the upstream ghcr.io/kai-scheduler registry with a git describe version. It also needs a docker-container buildx builder and QEMU emulation for the non-native architecture. Install the emulator with docker run --privileged --rm tonistiigi/binfmt --install arm64 (or amd64 on an Arm host), create a dedicated builder with docker buildx create --name kai-multiarch --driver docker-container, and prefix the make command with BUILDX_BUILDER=kai-multiarch so your default builder is not changed. A docker-container builder runs in its own network namespace, so it cannot push to a registry on localhost; use a registry it can reach over the network.

KAI-Scheduler

make build uses your current buildx builder. Check which one that is before you build:

docker buildx ls    # the current builder is marked with *; note its DRIVER

cd KAI-Scheduler
make build \
  DOCKER_REPO_BASE="$REGISTRY/kai-scheduler" \
  DOCKER_TAG="$TAG" \
  DOCKER_BUILDX_ADDITIONAL_ARGS=--push
cd ..

To build without pushing, drop DOCKER_BUILDX_ADDITIONAL_ARGS=--push and select a docker driver builder so the images land in your local Docker daemon:

cd KAI-Scheduler
BUILDX_BUILDER=default make build \
  DOCKER_REPO_BASE="$REGISTRY/kai-scheduler" \
  DOCKER_TAG="$TAG"
cd ..

The build compiles 15 services in sequence, each for both architectures, so the first run takes a while. Go caches are kept in ~/.cache/go-build-docker-gocache and ~/.cache/go-build-docker-gopath, and the build leaves a builder:1.26.3-bookworm image in your local Docker daemon. Remove them when you no longer need fast rebuilds.

Verify the Images Were Pushed

for svc in operator mpsd fractiond metricsd; do
  docker buildx imagetools inspect "$REGISTRY/kai-gpu-fractioning/$svc:$TAG" >/dev/null \
    && echo "ok   kai-gpu-fractioning/$svc" || echo "MISSING kai-gpu-fractioning/$svc"
done
for svc in podgrouper scheduler binder resourcereservation snapshot-tool scalingpod \
    nodescaleadjuster podgroupcontroller queuecontroller fairshare-simulator admission \
    operator time-based-fairshare-simulator numa-placement-exporter helm-hooks; do
  docker buildx imagetools inspect "$REGISTRY/kai-scheduler/$svc:$TAG" >/dev/null \
    && echo "ok   kai-scheduler/$svc" || echo "MISSING kai-scheduler/$svc"
done

Every line must print ok: four node-agent images and fifteen scheduler images.

Step 3: Install KAI-Scheduler

Fetch the chart dependency, then install with the fork images. The binder.runtimeClassName value gives KAI’s GPU reservation pod the nvidia RuntimeClass at install time. Without it, the reservation pod cannot load NVML and every fractional pod stays Pending for a reason that is invisible in the workload’s own events.

cd KAI-Scheduler
helm dependency build ./deployments/kai-scheduler
helm upgrade --install kai-scheduler ./deployments/kai-scheduler \
  -n kai-scheduler --create-namespace \
  --set global.gpuSharingMode=NvFractions \
  --set global.registry="$REGISTRY/kai-scheduler" \
  --set global.tag="$TAG" \
  --set binder.runtimeClassName=nvidia
cd ..

• global.gpuSharingMode is case-sensitive: NvFractions.
• Leave global.nvFractions.set at its default, false. Setting it to true installs the upstream kai-gpu-fractioning release from ghcr.io as a subchart, which is not your fork. You install the fork as its own release in Step 4. helm dependency build still has to download that subchart, because Helm checks that every declared dependency is present.
• For a private registry, add --set 'global.imagePullSecrets[0].name=<secret>' and create the secret in kai-scheduler first.

Verify before moving on. Wait for the deployments first, because helm upgrade --install returns before the pods are up:

kubectl -n kai-scheduler wait --for=condition=Available deploy --all --timeout=300s
helm status kai-scheduler -n kai-scheduler                   # STATUS: deployed
kubectl -n kai-scheduler get pods                            # all Running or Completed
kubectl -n kai-scheduler get deploy \
  -o jsonpath='{range .items[*]}{.metadata.name}{"\t"}{.spec.template.spec.containers[0].image}{"\n"}{end}'
kubectl get config.kai.scheduler kai-config -o jsonpath='{.spec.global.gpuSharingMode}{"\n"}'   # NvFractions

Until you finish Step 4, kai-config reports DependenciesFulfilled=False with GpuFractioningConfig not found. That is expected at this point.

Every image must show your $REGISTRY and $TAG. Do not continue with a pod that is not running.

Step 4: Install kai-gpu-fractioning

cd kai-gpu-fractioning
helm upgrade --install gpu-fractioning ./operator/charts \
  -n gpu-fractioning --create-namespace \
  --set images.operator.repository="$REGISTRY/kai-gpu-fractioning/operator" \
  --set images.fractiond.repository="$REGISTRY/kai-gpu-fractioning/fractiond" \
  --set images.metricsd.repository="$REGISTRY/kai-gpu-fractioning/metricsd" \
  --set images.mpsd.repository="$REGISTRY/kai-gpu-fractioning/mpsd" \
  --set images.operator.tag="$TAG" \
  --set images.fractiond.tag="$TAG" \
  --set images.metricsd.tag="$TAG" \
  --set images.mpsd.tag="$TAG"
cd ..

Set every image tag. An empty tag falls back to the chart’s appVersion, which is not your build.

Adjust for your environment:

sudo mkdir -p /opt/kai-runtime
sudo ln -sf /var/snap/microk8s/common/run/containerd.sock /opt/kai-runtime/containerd.sock

  --set fractioningAgent.criSocketPath=/opt/kai-runtime/containerd.sock \
  --set fractioningAgent.nriSocketPath=/var/snap/microk8s/common/run/nri.sock

Verify:

Wait for the operator, then for the DaemonSets it creates. The DaemonSets appear a few seconds after the operator starts:

kubectl -n gpu-fractioning wait --for=condition=Available deploy --all --timeout=300s
sleep 15
kubectl -n gpu-fractioning rollout status ds --timeout=300s

helm status gpu-fractioning -n gpu-fractioning               # STATUS: deployed
kubectl -n gpu-fractioning get pods -o wide                  # operator, plus mpsd and fractiond (2/2, with a metricsd sidecar) on each GPU node
kubectl -n gpu-fractioning get ds \
  -o jsonpath='{range .items[*]}{.metadata.name}{"\t"}{.status.desiredNumberScheduled}{"/"}{.status.numberReady}{"\n"}{end}'
kubectl get gpufractioningconfig default \
  -o jsonpath='{range .status.conditions[*]}{.type}={.status} ({.reason}){"\n"}{end}'

For every DaemonSet, desired must equal ready. A mismatch usually means a node failed the NRI or CRI socket check. Then confirm each GPU node reports ready, which is what gates scheduling:

kubectl get nodes -o jsonpath='{range .items[*]}{.metadata.name}{"\t"}{.status.conditions[?(@.type=="gpu-fractioning.nvidia.com/Ready")].status}{"\n"}{end}'

Confirm that MPS started in multi-user mode. Without it, a server accepts only clients with the owner’s uid, which would force every pod to run as that user:

kubectl -n gpu-fractioning logs ds/gpu-fractioning-mpsd | grep -E "multiuser|active thread percentage"

Expect Adding server 'shared' [multiuser enabled] and a line capping the default namespace at 10%. That cap is defense in depth for a pod that receives the server’s pipe directory instead of its own namespace’s.

Step 5: Verify With a Two-Pod Smoke Test

This step proves that the caps are enforced, not just requested. It puts two half-GPU pods on one GPU. The manifests for the full experiment are in examples/gpu-sharing/kai-gpu-fractions, and a ready-made version of this test, with the same pod names, is smoke-test-half-gpu.yaml. The minimal form is below, for readers without a checkout:

for n in a b; do
cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Pod
metadata:
  name: fraction-smoke-$n
  labels:
    kai.scheduler/queue: default-queue
  annotations:
    gpu-fraction: "0.5"
    nvidia.com/container.main.gpu-compute.mode: sm-sharing
spec:
  schedulerName: kai-scheduler
  runtimeClassName: nvidia
  restartPolicy: Never
  securityContext:
    runAsNonRoot: true
    runAsUser: 1000
  containers:
  - name: main
    image: nvidia/cuda:13.3.1-base-ubuntu24.04
    command: ["sleep", "3600"]
    securityContext:
      allowPrivilegeEscalation: false
      capabilities:
        drop: ["ALL"]
EOF
done
kubectl wait --for=condition=Ready pod/fraction-smoke-a pod/fraction-smoke-b --timeout=180s

The container must be named main because the annotation names the container it applies to. The chart creates default-queue for you.

Run these four checks. All four must hold:

# 1. What the binder decided: portion, memory, and the physical GPU UUID
for n in a b; do
  kubectl get pod fraction-smoke-$n -o jsonpath='{.metadata.annotations}' \
    | grep -o '"nvidia\.com/[^"]*":"[^"]*"' | sort -u
done
# 2. What fractiond injected into the container
kubectl exec fraction-smoke-a -- env | grep -E 'CUDA_MPS|NVIDIA_GPU_MEMORY|NVIDIA_VISIBLE_DEVICES'
# 3. What the MPS server enforces (ground truth, not the client's environment)
kubectl -n gpu-fractioning exec ds/gpu-fractioning-mpsd -- \
  nvidia-cuda-mps-control -p 3 namespace list --server=shared
# 4. The pod's own view of memory, with a hostile client-side override
kubectl exec fraction-smoke-a -- env CUDA_MPS_ACTIVE_THREAD_PERCENTAGE=100 \
  nvidia-smi --query-gpu=memory.total --format=csv,noheader

Clean up:

kubectl delete pod fraction-smoke-a fraction-smoke-b

Measure the Effective SM Share

The checks above show the caps that were requested and configured. To show that the SM share is enforced, run CUDA’s deviceQuery sample as a half-GPU pod that tries to override its own share to 100%:

cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Pod
metadata:
  name: fraction-devicequery
  labels:
    kai.scheduler/queue: default-queue
  annotations:
    gpu-fraction: "0.5"
    nvidia.com/container.main.gpu-compute.mode: sm-sharing
spec:
  schedulerName: kai-scheduler
  runtimeClassName: nvidia
  restartPolicy: Never
  securityContext:
    runAsNonRoot: true
    runAsUser: 1000
  containers:
  - name: main
    image: nvcr.io/nvidia/k8s/cuda-sample:devicequery
    env:
    - name: CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
      value: "100"
    securityContext:
      allowPrivilegeEscalation: false
      capabilities:
        drop: ["ALL"]
EOF
kubectl wait --for=jsonpath='{.status.phase}'=Succeeded pod/fraction-devicequery --timeout=180s
kubectl logs fraction-devicequery | grep -E 'Multiprocessors|Result'
kubectl delete pod fraction-devicequery

The Multiprocessors count must be half the GPU’s SM count, rounded down, despite the override: for example, (054) Multiprocessors on an A100 with 108 SMs. fractiond replaces the variable when the container is created, and the MPS server caps the namespace whatever the client asks for. deviceQuery still reports the physical memory total, as the note below explains.

Troubleshooting

Uninstall

Delete the default GpuFractioningConfig first, while the operator is still running. The operator then stops the mpsd and fractiond DaemonSets and clears the node conditions. If you uninstall the chart first, the DaemonSets keep running, MPS stays on the GPUs, and the object can no longer be deleted: its gpu-fractioning.nvidia.com/cleanup-node-conditions finalizer is removed only by the operator.

kubectl delete gpufractioningconfig default --wait
kubectl -n gpu-fractioning wait --for=delete pod -l app.kubernetes.io/component=mpsd --timeout=180s
helm uninstall gpu-fractioning -n gpu-fractioning
helm uninstall kai-scheduler -n kai-scheduler

The mpsd pods take about a minute to terminate. Helm keeps the default-queue and default-parent-queue queues, the default SchedulingShard, the CRDs, and the kai-resource-reservation, kai-scheduler, and gpu-fractioning namespaces. For a clean slate, remove them too. Delete the CRDs only if nothing else in the cluster uses them:

kubectl get crd -o name | grep -E 'kai\.scheduler|scheduling\.run\.ai' | xargs kubectl delete
kubectl delete namespace kai-resource-reservation kai-scheduler gpu-fractioning

Next Steps

• Deploy the 16-model experiment on the cluster you just built.
• Per-GPU Compute and Memory Limits with MPS covers the MPS v3 mechanism that these forks drive for you.
• Interpreting GPU Sharing Results explains how to read the benchmark.