[RFC]: Supporting multi-modality models in AIBrix

[happyandslow]

Summary

AIBrix has been successfully adopted as a cloud-native orchestration framework for LLM inference engines (e.g., vLLM, SGLang). To expand its scope into multimodal AI inference, new capabilities are required across the gateway, engine integration, data handling, caching, metrics, and pre/post-processing layers. This RFC motivates and details the prioritized tasks needed to extend AIBrix to fully support multimodal scenarios.

Motivation

By extending AIBrix’s architecture in these ways, we could enable AIBrix to be a unified inference orchestration platform—handling not just text LLMs, but richer, multimodal workloads—while retaining its strengths in routing, cost-efficiency, and scalability.

Proposed Change

1. Multimodal Engine Integration

• Many multimodal models (e.g., Qwen-VL, LLaVA, GPT-4o) accept text + image input and return text output.
• OpenAI’s /chat/completions and /completions APIs already support mixed input payloads ({"type": "text"} + {"type": "image_url"}), which are widely adopted.

TODO

• Verify AIBrix API Gateway to accept different forms of modality input. Fully support Open AI API through vLLM. Refer to examples in https://github.com/vllm-project/vllm/blob/main/examples/online_serving/openai_chat_completion_client_for_multimodal.py for further details
  • image/video url
  • image/video/audio encoding
  • raw data
• Deliver a demo deployment with a text+image model (e.g., Qwen-VL, ) running through AIBrix.
  • Single image + text
  • Multi images + text
  • Video
  • Audio
• Verify AIBrix's support for embedding input. Refer to examples in https://github.com/vllm-project/vllm/blob/main/examples/online_serving/openai_chat_embedding_client_for_multimodal.py for further details.
  • https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#embeddings-api_1
  • https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#transcriptions-api_1
• Support embedding endpoint
  • [API] add support for embeddings api #1208 @varungup90

2. Engine Adaptation Beyond LLMs

• Current AIBrix engine integrations (vLLM, SGLang) assume an OpenAI-compatible LLM interface.
• Diffusion models (e.g., xDiT, https://github.com/xdit-project/xDiT) expose non-OpenAI APIs with different input/output formats.
• Supporting such engines would expand AIBrix beyond text and vision-language models into generative image/video/audio domains.

TODO

• Evaluate candidate API standards: [RFC]: xDiT Video Generation API #1595
  • Provide a pluggable API gateway layer to support alternative APIs when required.
• Define a capability registry per engine (supports: [text, image, audio], API type). [RFC]: xDiT Video Generation API #1595
• Prototype integration with multi-modality engine (e.g., xDiT) to test the gateway + orchestration path.
  • Deliver a demo deployment with video generation model: https://huggingface.co/Wan-AI

3. Large File Download / Access Support

• Multimodal models often need to handle large input files (video/audio) or large outputs (image generations, video clips).
• Current API payloads are insufficient for multi-100MB content.
• Industry practice (e.g., OpenAI file uploads, AWS Bedrock) relies on remote storage integration (S3/GCS/Azure Blob).

TODO

• Extend AIBrix Gateway to:
  • Support file references (URIs) instead of inline payloads.
  • Integrate with cloud storage (e.g., S3) -compatible object storage for upload/download.
• Define policies for temporary storage, caching, and cleanup of large multimodal assets.
• Add data movement monitoring to avoid bottlenecks.

4. KV Cache and Beyond

• AIBrix currently implements a distributed KV cache to accelerate LLM inference by storing attention states.
• LMCache added support for identical images so that image reuse could be identified (and prevent different images to reused the same KV due to the use of placeholder tokens)
  • https://blog.lmcache.ai/2025-07-03-multimodal-models/
• (Beyond KV cache) In multimodal models, cacheable elements may include:
  • Visual embeddings (CLIP features).
  • Latent representations for image/video/audio.
• Extending caching semantics will improve throughput and latency for repeated multimodal contexts.

TODO

• Validate and adding support for AIBrix to identify image/video/audio reuse.
  • Benchmark benefits for common workloads (e.g., repeated image captioning over the same video frames).
• Extend cache API to allow multimodal engines to store/retrieve intermediate representations.
  • See some techniques mentioned here: https://github.com/xdit-project/xDiT?tab=readme-ov-file#cache-acceleration
    • https://liewfeng.github.io/TeaCache/
  • Generalize KV cache service to support arbitrary embedding/latent types.

5. Modality-Specific Metrics

• Current AIBrix metrics (e.g., tokens/sec, request latency) are LLM-centric.
• Multimodal workloads require new observability dimensions:
  • Image/video size processed (pixels/frame count).
  • Audio duration.
  • Preprocessing latency (I/O, decoding, embedding).
• These metrics are crucial for autoscaling and SLA enforcement.

TODO

• Extend telemetry pipeline with modality-aware metrics exporters.
• Define new autoscaling triggers, examples:
  • Scale on average pixels/sec.
  • Scale on average audio seconds/sec.
• Update dashboards to visualize cross-modality utilization trends.

6. Pre/Post Processing Pipelines

• Multimodal inference typically requires non-trivial preprocessing and postprocessing:
  • Image decoding/normalization, resizing.
  • Audio resampling or spectrogram generation.
  • Video frame extraction.
  • Postprocessing for outputs (e.g., base64→PNG, WAV encoding).

vLLM's separation of preprocessing (e.g., encoding) pipeline separation is under development

• [RFC]: Prototype Separating Vision Encoder to Its Own Worker vllm#20799

• [Core] Encoder separation for Encode-Prefill-Decode Disaggregation vllm#21740

• Three ways of servings:

  1. Frozen LLM + External Encoder (Encoder → projector → embeddings as prefix tokens to frozen LLM)
  
  2. Unified Discrete Tokenization (Non-text → discrete tokens in shared vocab; train LLM jointly)
  
  3. Orchestration / Tool Use (LLM calls external encoder services → text results fed back in)- Centralizing these within AIBrix prevents duplication across engines and enables standardized observability and scaling.
  

TODO

• Introduce pre/post-processing sidecar components or pipeline stages.
• Allow pluggable user-defined preprocessing modules for custom use cases.

Alternatives Considered

None

[happyandslow]

Tracking API compatibility with vision/video/audio language model in #1512

[happyandslow]

Demoing image/video/audio language models deployment files using vLLM as example engines: https://github.com/happyandslow/aibrix/tree/lexu/test-multimodality-scenarios/test/regression/v0.5.0/vllm

[Jeffwan]

@happyandslow overall roadmap looks good to me. this is a very solid plan, let's keep moving!

[happyandslow]

2. Engine Adaptation Beyond LLMs

• Current AIBrix engine integrations (vLLM, SGLang) assume an OpenAI-compatible LLM interface.
• Diffusion models (e.g., xDiT, https://github.com/xdit-project/xDiT) expose non-OpenAI APIs with different input/output formats.
• Supporting such engines would expand AIBrix beyond text and vision-language models into generative image/video/audio domains.

TODO

• Evaluate candidate API standards:
  • Provide a pluggable API gateway layer to support alternative APIs when required.
• Define a capability registry per engine (supports: [text, image, audio], API type).
• Prototype integration with multi-modality engine (e.g., xDiT) to test the gateway + orchestration path.
  • Deliver a demo deployment with video generation model: https://huggingface.co/Wan-AI

Option 1: Supporting ComfyUI

• ComfyUI API: https://github.com/comfyanonymous/ComfyUI/blob/master/server.py
  /prompt, /history, /history/{prompt_id}, etc.
• Running customized models on top of ComfyUI: http://docs.comfy.org/development/core-concepts/models#portable%2Fmanual-install
• POST and GET (applies to both image and video)

  • POST to /prompt endpoint expects a json format input (workflow). See example: https://raw.githubusercontent.com/Comfy-Org/workflow_templates/refs/heads/main/templates/video_wan2_2_5B_ti2v.json

  • Example POST

    curl -X POST http://127.0.0.1:8188/prompt \
      -H "Content-Type: application/json" \
      -d '{
        "prompt": {
          "1": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": { "ckpt_name": "v1-5-pruned.ckpt" }
          },
          "2": {
            "class_type": "CLIPTextEncode",
            "inputs": {
              "text": "A futuristic city skyline at sunset",
              "clip": ["1", 1]
            }
          },
          "3": {
            "class_type": "KSampler",
            "inputs": {
              "seed": 42,
              "steps": 20,
              "cfg": 7,
              "sampler_name": "euler",
              "scheduler": "normal",
              "denoise": 1.0,
              "model": ["1", 0],
              "positive": ["2", 0],
              "latent_image": ["1", 2]
            }
          },
          "4": {
            "class_type": "VAEDecode",
            "inputs": { "samples": ["3", 0], "vae": ["1", 3] }
          },
          "5": {
            "class_type": "SaveImage",
            "inputs": { "filename_prefix": "city_skyline" }
          }
        }
      }'
    
    

    Response

    {
      "prompt_id": "2e1d5e91-bc34-49e3-9ac3-7f51c8f0e1c2",
      "number": 0,
      "node_errors": {}
    }
    

  • Expected output { "prompt_id": "b8e9f8d2-..." } and result is retrieve by GET GET /view/{uuid}/5/0

  curl http://127.0.0.1:8188/view/ea2b8c58-f31d-42de-9a3c-6e04a6b8a4b7/5/0 --output result.png
  

Option 2: Supporting OpenAI's multi-modality API:

• Image generation: https://platform.openai.com/docs/guides/image-generation?image-generation-model=gpt-image-1&api=image
• POST and GET
  • Example POST

  curl https://api.openai.com/v1/images/generations \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer $OPENAI_API_KEY" \
    -d '{
      "model": "gpt-image-1",
      "prompt": "A cute baby sea otter",
      "n": 1,
      "size": "1024x1024"
    }'
  

  Response

  {
    "created": 1713833628,
    "data": [
      {
        "b64_json": "..."
      }
    ],
    "usage": {
      "total_tokens": 100,
      "input_tokens": 50,
      "output_tokens": 50,
      "input_tokens_details": {
        "text_tokens": 10,
        "image_tokens": 40
      }
    }
  }
  
  

• OpenAI API does not support text-to-video yet.

Option 3: Creare custom endpoint based on targeted engine (e.g., xDiT)

• xDit supports a simple /generate API that generate image and video. See http server here: https://github.com/xdit-project/xDiT/tree/main/entrypoints
• The output is stored in a node local path
• Example request

  curl -X POST "http://localhost:6000/generate" \
       -H "Content-Type: application/json" \
       -d '{
             "prompt": "a cute rabbit",
             "num_inference_steps": 50,
             "seed": 42,
             "cfg": 7.5, 
             "save_disk_path": "/tmp"
           }'
  

@Jeffwan @varungup90 I'd like some input on which directions we should move to in order to better aligned with the development goal.

[Jeffwan]

@happyandslow option 1's priority is low. comfyui should be considered after aibrix supports enough multi-modality models
I personally think we should provide an api adapter or aggregation layer (xDit, Batch API ) etc on our side. underneath, it could adapt xDit based on engine type and forward the request

[happyandslow]

3. Large File Download / Access Support

• Multimodal models often need to handle large input files (video/audio) or large outputs (image generations, video clips).
• Current API payloads are insufficient for multi-100MB content.
• Industry practice (e.g., OpenAI file uploads, AWS Bedrock) relies on remote storage integration (S3/GCS/Azure Blob).

TODO

• Extend AIBrix Gateway to:
  • Support file references (URIs) instead of inline payloads.
  • Integrate with cloud storage (e.g., S3) -compatible object storage for upload/download.
• Define policies for temporary storage, caching, and cleanup of large multimodal assets.
• Add data movement monitoring to avoid bottlenecks.

Currently vLLM supports reading files stored from a local filesystem. This could be useful if file size is large and we can probably enable many source of multi-modality data. This probably requires AIBrix runtime to support download files from cloud storage like S3/TOS etc, similar to downloading models from cloud storage. However, the difference here is that the multi-modality data (image/video/audio etc.) needs to be downloaded upon request arrival instead of during the container initialization stage.

There might be three ways of achieving this.

Option 1: File downloaded by runtime, request forwarded by runtime to engine.

• This requires runtime and engine shares the same file system
• The benefits are:
  • The gateway is only forward the request once
  • The gateway can be stateless (without storing the request)
• The downside are:
  • Gateway should be aware of whether to dispatch request to runtime or to the engine (and that the request path would be different based on the request content), which I doubt would be a good design

Option 2: File downloaded by runtime, request dispatched by gateway to engine.

• This requires runtime and engine shares the same file system
• The benefits:
  • The gateway could separate handler for runtime and engine message cleanly
• Downside:
  • The gateway needs to be stateful by storing the request
  • Longer latency by one more trip

Option 3: File downloaded by gateway, request dispatched by gateway to engine.

• The benefit
  • No need to involve runtime and file system sharing
  • Shorter latency as it does not involve request forwarding
• The downside:
  • Probably would have scalability issue or gateway becomes a bottleneck since all downloads happen at gateway
  • Complex gateway design

@Jeffwan @varungup90 I'd like some input on which directions we should move to in order to better aligned with the development goal.

[Jeffwan]

@happyandslow it depends on how to implement the apiserver. did you talk with @varungup90 on the aggregated server idea?

[happyandslow]

@happyandslow it depends on how to implement the apiserver. did you talk with @varungup90 on the aggregated server idea?

Yeah he suggested me to bring this issue up for discussion cc @varungup90