How vllm-project/vllm-neuron works
vLLM Neuron is a beta vLLM platform plugin that runs supported language, embedding, and multimodal models on AWS Trainium. A request enters through vLLM’s compatible serving APIs, is scheduled and executed by Neuron workers, and produces generated tokens or embeddings. The repository also contains distributed execution, disaggregated encoder routing, speculative decoding, accuracy-debugging and snapshot capture, and observability components. The README documents these capabilities; sampled source excerpts are partial, so uncertain internal wiring is omitted.
Main entry points: neuron_model_runner.py Model runner · neuron_worker.py Neuron worker · router.py Encoder router · platform.py Neuron platform
Components 23
Serving and execution
- Neuron platform
vllm_neuron/vllm/platform.py - Request scheduler
vllm_neuron/vllm/core/scheduler.py - Neuron worker
vllm_neuron/vllm/worker/neuron_worker.py - Model runner
vllm_neuron/vllm/worker/neuron_model_runner.py - EAGLE decoding
vllm_neuron/vllm/spec_decode/eagle.py
Models and compute
- Model registry
vllm_neuron/model/registry.py - Model implementations
vllm_neuron/model - Neuron operations
vllm_neuron/functional - Sampling
vllm_neuron/nn/sampler.py - KV cache
vllm_neuron/model/kv_cache.py
Distributed inference
- Parallel state
vllm_neuron/parallel/neuron_parallel_state.py - Collectives
vllm_neuron/parallel/all2all.py - Encoder router
vllm_neuron/vllm/disaggregated_encoder/router.py
Debug and observability
- Accuracy debugger
vllm_neuron/accuracy/accuracy_debugger/api.py - Logit validation
vllm_neuron/accuracy/logit_validation.py - Tensor snapshots
vllm_neuron/snapshot/capture.py - Serving metrics
vllm_neuron/metrics.py - Neuron profiler
vllm_neuron/vllm/worker/neuron_profiler.py
Other components
- API client
- vLLM serving APIs
- Vision encoders
- Prefill/decode servers
- AWS Trainium
Connections 20
- API client → vLLM serving APIs: sends requests
- vLLM serving APIs → Neuron platform: selects backend
- Neuron platform → Request scheduler: integrates scheduling
- Request scheduler → Neuron worker: dispatches work
- Neuron worker → Model runner: loads and runs
- Model runner → Model registry: resolves model
- Model registry → Model implementations: selects implementation
- Model implementations → Neuron operations: uses operations
- Model runner → Sampling: samples outputs
- Model runner → KV cache: binds cache
- Model runner → Parallel state: uses groups
- Parallel state → Collectives: coordinates ranks
- Model runner → AWS Trainium: executes on
- Encoder router → Vision encoders: routes encoding
- Encoder router → Prefill/decode servers: forwards requests
- Model runner → EAGLE decoding: supports decoding
- Accuracy debugger → Logit validation: runs validation
- Model runner → Tensor snapshots: publishes context
- Neuron worker → Neuron profiler: creates profiler
- Model runner → Serving metrics: reports metrics