Skip to content

Add GPU / TPU / accelerator targeting and resource injection #18

Description

@tartkri-bot

Problem

The predictor topology reads comp.Resources but does not provide any GPU (NVIDIA GPU / TPU / AWS Inferentia / Trainium) resource helpers, and does not auto-inject or validate GPU-specific requirements. KServe users frequently need:

  1. GPU resource requests: Explicitly requesting nvidia.com/gpu: "1" or amd.com/gpu in the container resources.
  2. Auto-GPU detection: Based on model format or selected ServingRuntime, optionally suggest/add GPU resources automatically.
  3. TPU support (Google): Selector for cloud.google.com/gke-accelerator and toleration for cloud.google.com/gke-tpu.
  4. AWS Inferentia / Trainium: Resources like aws.amazon.com/neuron and node selectors for Neuron nodes.
  5. GPU-specific environment variables: Auto-inject CUDA_VISIBLE_DEVICES, NVIDIA_VISIBLE_DEVICES, or runtime-specific vars when GPU is requested.

Currently, users must manually set comp.Resources with the exact GPU resource name and quantity, which is Kubernetes-specific and error-prone.

Suggested approach

Add a gpu or accelerator block to ModelServerCustomSpec (and LLM parameters):

  • type"nvidia", "tpu", "inferentia", "trainium", "amd", etc.
  • count — number of accelerators.
  • autoInjectEnv — boolean, defaults to true. When enabled, inject common env vars for the selected accelerator type.

Alternatively, keep it simpler: add nvidiaGPU, tpu, neuron as optional typed fields with count + autoEnv.

Acceptance criteria

  • GPU resources are correctly injected into the serving container's resource requests/limits.
  • Node selectors / tolerations for the selected accelerator type are auto-applied where applicable.
  • Accelerator-specific env vars are injected when requested.
  • Multi-GPU scenarios (tensor parallelism across GPUs on the same node) are validated.
  • Chainsaw tests verify GPU scheduling on clusters with GPU nodes (or mock with node labels for CI).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions