[RFC]: Neuron Support for V1 Engine

[aarondou]

Motivation.

This RFC discusses AWS Neuron’s plan to integrate NeuronX Distributed (NxD) Inference with V1 architecture on vLLM.

with: @elaineyz @mrinalks @rgrandhiamzn

Proposed Changes.

Change 1: Chunked prefill support — the default mechanism in V1.

We have upgraded core Neuron interfaces to comply with V1 Engine and enable Chunked Prefill. No modifications to vLLM core interfaces were needed. This change offers a native implementation of Chunked Prefill — a required feature on V1.

• Neuron Worker compliant with V1 Model Loader, and Model Runner.
• Chunked Prefill Worker that enables chunked prefill on NxD Inference.
• No modifications to vLLM Core Interfaces.

Change 2: Optionally disable Chunked Prefill — for higher performance on Neuron.

For customers who see a higher performance on Neuron without Chunked Prefill, we offer an option to disable Chunked Prefill.

Option 1: New scheduler mode in V1 that disables chunked prefill (RECOMMENDED)

• While we understand that chunked prefill is a defining characteristic of the V1 architecture, we believe that allowing the flexibility to optionally disable it can make vLLM more easily adoptable by more hardware providers and users.
• A simple, non-intrusive way to turn off chunked prefill on top of the V1 scheduler has already been developed by IBM Spyre’s implementation. Concretely, we can override the schedule() function to intercept running requests (i.e. in practice this would be decode phase requests), check whether there are requests in the waiting queue (i.e. in practice this would be new requests ready to be prefilled), and temporarily purge the running queue and have the base scheduler perform scheduling as usually (on the waiting requests first). After the base scheduler is finished, the override will then restore the running queue.

Option 2: Override V1 scheduler via a scheduler plugin repository

• We can maintain a lightweight plugin repository that exposes the scheduler extension logic described in Option 1.
  Note: There is no additional code in-tree to achieve this behavior.
• To opt-in to this implementation, users only need to pip install this repository in addition to vllm, and may continue to invoke vLLM exactly the same as they normally would (e.g. python -m vllm.entrypoints.openai.api_server)
• Users can still disable the scheduler override after installing the plugin via an environment variable.
• We will plan to house the plugin repository on the Aws Neuron GitHub.

Feedback Period.

1 week

CC List.

@aarondou
@aws-navyadhara
@aws-satyajith
@aws-yishanm
@elaineyz
@liangfu
@mrinalks
@rgrandhiamzn
@robertgshaw2-redhat
@simon-mo
@WoosukKwon
@DarkLight1337

Any Other Things.

• The Transformers NeuronX (TNx) framework will be deprecated upon Neuron’s migration to V1.
• Follow-up RFCs:
  • Proof-of-concept native integration on vLLM modeling code.
  • Performant Chunked Prefill on Neuron

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[simon-mo]

What happens to the persistent batch when chunked prefill is off?

[aarondou]

Could you please clarify a little on what do you mean by persistent batch here?

[mrinalks]

@aarondou nit/ replace "Proposed Change" with "Proposed Changes"; replace "Change # 1 and # 2" with bullets 1. and 2. to avoid linking to PR # 1 and PR # 2.

[aarondou]

nit/ replace "Proposed Change" with "Proposed Changes"; replace "Change # 1 and # 2" with bullets 1. and 2. to avoid linking to PR # 1 and PR # 2.

Done

[aarondou]

What happens to the persistent batch when chunked prefill is off?

Hi @simon-mo , did you mean to ask about Change 2 Option 1 for New scheduler mode in V1 that disables chunked prefill (RECOMMENDED)?

The core objective is to support chunked-prefill-disabled mode directly in-tree. We are considering two possible approaches:

• In-tree scheduling override: Implement a scheduling override mechanism similar to IBM Spyre’s implementation, which introduces a holdback queue to manage requests between the waiting, running, and holdback states on top of the V1 scheduler. The key difference from Spyre’s approach is that this implementation would reside in-tree rather than as an out-of-tree plugin.

• Reimplementing scheduling logic: Alternatively, instead of building on top of the V1 scheduler, we could reimplement parts of the scheduling logic to support continuous batching natively for chunked-prefill-disabled mode. This direction is less defined, and we would like to discuss with the maintainers and community to explore it further.

[elaineyz]

What happens to the persistent batch when chunked prefill is off?

Hi @simon-mo, if by "persistent batch" you mean the InputBatch object, then the answer is that we continue to use the same interface and respect all contracts even when chunked prefill is off. Our goal is to achieve disabling chunked prefill in a nonintrusive way to the rest of vllm implementations.

Refer to this section in the RFC for details on our thinking. All changes would be confined within the scheduler:

Concretely, we can override the schedule() function to intercept running requests (i.e. in practice this would be decode phase requests), check whether there are requests in the waiting queue (i.e. in practice this would be new requests ready to be prefilled), and temporarily purge the running queue and have the base scheduler perform scheduling as usually (on the waiting requests first). After the base scheduler is finished, the override will then restore the running queue.

[aarondou]

We will not implement the Neuron vLLM V1 integration directly in-tree. Instead, we will follow the vLLM community’s recommendation by building the integration as a Neuron plugin under the vllm-project.