merge v0.4.0

docs/source/dev/kernel/paged_attention.rst

+vLLM Paged Attention
+====================
+
+-  Currently, vLLM utilizes its own implementation of a multi-head query
+   attention kernel (``csrc/attention/attention_kernels.cu``). 
+   This kernel is designed to be compatible with
+   vLLM's paged KV caches, where the key and value cache are stored in
+   separate blocks (note that this block concept differs from the GPU
+   thread block. So in a later document, I will refer to vLLM paged
+   attention block as "block", while refer to GPU thread block as
+   "thread block").
+-  To achieve high performance, this kernel relies on a specially
+   designed memory layout and access method, specifically when threads
+   read data from global memory to shared memory. The purpose of this
+   document is to provide a high-level explanation of the kernel
+   implementation step by step, aiding those who wish to learn about the
+   vLLM multi-head query attention kernel. After going through this 
+   document, users will likely have a better understanding and feel easier
+   to follow the actual implementation.
+-  Please note that this document may not cover all details, such as how
+   to calculate the correct index for the corresponding data or the dot
+   multiplication implementation. However, after reading this document
+   and becoming familiar with the high-level logic flow, it should be
+   easier for you to read the actual code and understand the details.
+
+Inputs
+------
+
+-  The kernel function takes a list of arguments for the current thread
+   to perform its assigned work. The three most important arguments are
+   the input pointers ``q``, ``k_cache``, and ``v_cache``, which point
+   to query, key, and value data on global memory that need to be read
+   and processed. The output pointer ``out`` points to global memory
+   where the result should be written. These four pointers actually
+   refer to multi-dimensional arrays, but each thread only accesses the
+   portion of data assigned to it. I have omitted all other runtime
+   parameters here for simplicity.
+
+   .. code:: cpp
+
+      template<
+      typename scalar_t,
+      int HEAD_SIZE,
+      int BLOCK_SIZE,
+      int NUM_THREADS,
+      int PARTITION_SIZE = 0>
+      __device__ void paged_attention_kernel(
+      ... // Other side args.
+      const scalar_t* __restrict__ out,       // [num_seqs, num_heads, max_num_partitions, head_size]
+      const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
+      const scalar_t* __restrict__ k_cache,   // [num_blocks, num_kv_heads, head_size/x, block_size, x]
+      const scalar_t* __restrict__ v_cache,   // [num_blocks, num_kv_heads, head_size, block_size]
+      ... // Other side args.
+      )
+
+-  There are also a list of template arguments above the function
+   signature that are determined during compilation time. ``scalar_t``
+   represents the data type of the query, key, and value data elements,
+   such as FP16. ``HEAD_SIZE`` indicates the number of elements in each
+   head. ``BLOCK_SIZE`` refers to the number of tokens in each block.
+   ``NUM_THREADS`` denotes the number of threads in each thread block.
+   ``PARTITION_SIZE`` represents the number of tensor parallel GPUs (For
+   simplicity, we assume this is 0 and tensor parallel is disabled).
+-  With these arguments, we need to perform a sequence of preparations.
+   This includes calculating the current head index, block index, and
+   other necessary variables. However, for now, we can ignore these
+   preparations and proceed directly to the actual calculations. It will
+   be easier to understand them once we grasp the entire flow.
+
+Concepts
+--------
+
+-  Just before we dive into the calculation flow, I want to describe a
+   few concepts that are needed for later sections. However, you may
+   skip this section and return later if you encounter any confusing
+   terminologies.
+-  **Sequence**: A sequence represents a client request. For example,
+   the data pointed to by ``q`` has a shape of
+   ``[num_seqs, num_heads, head_size]``. That represents there are total
+   ``num_seqs`` of query sequence data are pointed by ``q``. Since this 
+   kernel is a single query attention kernel, each sequence only has one
+   query token. Hence, the ``num_seqs`` equals the total number of tokens 
+   that are processed in the batch.
+-  **Context**: The context consists of the generated tokens from the
+   sequence. For instance, ``["What", "is", "your"]`` are the context
+   tokens, and the input query token is ``"name"``. The model might
+   generate the token ``"?"``.
+-  **Vec**: The vec is a list of elements that are fetched and
+   calculated together. For query and key data, the vec size
+   (``VEC_SIZE``) is determined so that each thread group can fetch and
+   calculate 16 bytes of data at a time. For value data, the vec size
+   (``V_VEC_SIZE``) is determined so that each thread can fetch and
+   calculate 16 bytes of data at a time. For example, if the
+   ``scalar_t`` is FP16 (2 bytes) and ``THREAD_GROUP_SIZE`` is 2, the 
+   ``VEC_SIZE`` will be 4, while the ``V_VEC_SIZE`` will be 8.
+-  **Thread group**: The thread group is a small group of
+   threads(\ ``THREAD_GROUP_SIZE``) that fetches and calculates one
+   query token and one key token at a time. Each thread handles only a
+   portion of the token data. The total number of elements processed by
+   one thread group is referred as ``x``. For example, if the thread
+   group contains 2 threads and the head size is 8, then thread 0
+   handles the query and key elements at index 0, 2, 4, 6, while thread
+   1 handles the elements at index 1, 3, 5, 7.
+-  **Block**: The key and value cache data in vLLM are split into
+   blocks. Each block stores data for a fixed number(\ ``BLOCK_SIZE``)
+   of tokens at one head. Each block may contain only a portion of the
+   whole context tokens. For example, if the block size is 16 and the
+   head size is 128, then for one head, one block can store 16 \* 128 =
+   2048 elements.
+-  **Warp**: A warp is a group of 32 threads(\ ``WARP_SIZE``) that
+   execute simultaneously on a stream multiprocessor (SM). In this
+   kernel, each warp processes the calculation between one query token
+   and key tokens of one entire block at a time (it may process multiple
+   blocks in multiple iterations). For example, if there are 4 warps and
+   6 blocks for one context, the assignment would be like warp 0 handles
+   the 0th, 4th blocks, warp 1 handles the 1st, 5th blocks, warp 2
+   handles the 2nd block and warp 3 handles the 3rd block.
+-  **Thread block**: A thread block is a group of
+   threads(\ ``NUM_THREADS``) that can access the same shared memory.
+   Each thread block contains multiple warps(\ ``NUM_WARPS``), and in
+   this kernel, each thread block processes the calculation between one
+   query token and key tokens of a whole context.
+-  **Grid**: A grid is a collection of thread blocks and defines the
+   shape of the collection. In this kernel, the shape is
+   ``(num_heads, num_seqs, max_num_partitions)``. Therefore, each thread
+   block only handles the calculation for one head, one sequence, and
+   one partition.
+
+Query
+-----
+
+-  This section will introduce how query data is stored in memory and
+   fetched by each thread. As mentioned above, each thread group fetches
+   one query token data, while each thread itself only handles a part of
+   one query token data. Within each warp, every thread group will fetch
+   the same query token data, but will multiply it with different key
+   token data.
+
+   .. code:: cpp
+
+      const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
+
+   .. figure:: ../../assets/kernel/query.png
+      :alt: query
+      :width: 70%
+      :align: center
+
+      Query data of one token at one head
+
+-  Each thread defines its own ``q_ptr`` which points to the assigned
+   query token data on global memory. For example, if ``VEC_SIZE`` is 4
+   and ``HEAD_SIZE`` is 128, the ``q_ptr`` points to data that contains
+   total of 128 elements divided into 128 / 4 = 32 vecs.
+
+   .. figure:: ../../assets/kernel/q_vecs.png
+      :alt: q_vecs
+      :width: 70%
+      :align: center
+
+      ``q_vecs`` for one thread group
+
+   .. code:: cpp
+
+      __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
+
+-  Next, we need to read the global memory data pointed to by ``q_ptr``
+   into shared memory as ``q_vecs``. It is important to note that each
+   vecs is assigned to a different row. For example, if the
+   ``THREAD_GROUP_SIZE`` is 2, thread 0 will handle the 0th row vecs,
+   while thread 1 handles the 1st row vecs. By reading the query data in
+   this way, neighboring threads like thread 0 and thread 1 can read
+   neighbor memory, achieving the memory coalescing to improve
+   performance.
+
+Key
+---
+
+-  Similar to the "Query" section, this section introduces memory layout
+   and assignment for keys. While each thread group only handle one
+   query token one kernel run, it may handle multiple key tokens across
+   multiple iterations. Meanwhile, each warp will process multiple blocks
+   of key tokens in multiple iterations, ensuring that all context
+   tokens are processed by the entire thread group after the kernel run.
+   In this context, "handle" refers to performing the dot multiplication
+   between query data and key data.
+
+   .. code:: cpp
+
+      const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride
+                          + kv_head_idx * kv_head_stride
+                          + physical_block_offset * x;
+
+-  Unlike to ``q_ptr``, ``k_ptr`` in each thread will point to different
+   key token at different iterations. As shown above, that ``k_ptr``
+   points to key token data based on ``k_cache`` at assigned block,
+   assigned head and assigned token.
+
+   .. figure:: ../../assets/kernel/key.png
+      :alt: key
+      :width: 70%
+      :align: center
+
+      Key data of all context tokens at one head
+
+-  The diagram above illustrates the memory layout for key data. It
+   assumes that the ``BLOCK_SIZE`` is 16, ``HEAD_SIZE`` is 128, ``x`` is
+   8, ``THREAD_GROUP_SIZE`` is 2, and there are a total of 4 warps. Each
+   rectangle represents all the elements for one key token at one head,
+   which will be processed by one thread group. The left half shows the
+   total 16 blocks of key token data for warp 0, while the right half
+   represents the remaining key token data for other warps or
+   iterations. Inside each rectangle, there are a total 32 vecs (128
+   elements for one token) that will be processed by 2 threads (one
+   thread group) separately.
+
+   .. figure:: ../../assets/kernel/k_vecs.png
+      :alt: k_vecs
+      :width: 70%
+      :align: center
+
+      ``k_vecs`` for one thread
+
+   .. code:: cpp
+
+      K_vec k_vecs[NUM_VECS_PER_THREAD]
+
+-  Next, we need to read the key token data from ``k_ptr`` and store
+   them on register memory as ``k_vecs``. We use register memory for
+   ``k_vecs`` because it will only be accessed by one thread once,
+   whereas ``q_vecs`` will be accessed by multiple threads multiple
+   times. Each ``k_vecs`` will contain multiple vectors for later
+   calculation. Each vec will be set at each inner iteration. The
+   assignment of vecs allows neighboring threads in a warp to read
+   neighboring memory together, which again promotes the memory
+   coalescing. For instance, thread 0 will read vec 0, while thread 1
+   will read vec 1. In the next inner loop, thread 0 will read vec 2,
+   while thread 1 will read vec 3, and so on.
+-  You may still be a little confused about the overall flow. Don't
+   worry, please keep reading the next "QK" section. It will illustrate
+   the query and key calculation flow in a clearer and higher-level
+   manner.
+
+QK
+---
+
+-  As shown the pseudo code below, before the entire for loop block, we
+   fetch the query data for one token and store it in ``q_vecs``. Then,
+   in the outer for loop, we iterate through different ``k_ptrs`` that
+   point to different tokens and prepare the ``k_vecs`` in the inner for
+   loop. Finally, we perform the dot multiplication between the
+   ``q_vecs`` and each ``k_vecs``.
+
+   .. code:: cpp
+
+      q_vecs = ...
+      for ... {
+         k_ptr = ...
+         for ... {
+            k_vecs[i] = ...
+         }
+         ...
+         float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(q_vecs[thread_group_offset], k_vecs);
+      }
+
+-  As mentioned before, for each thread, it only fetches part of the
+   query and key token data at a time. However, there will be a cross
+   thread group reduction happen in the ``Qk_dot<>::dot`` . So ``qk``
+   returned here is not just between part of the query and key token dot
+   multiplication, but actually a full result between entire query and
+   key token data.
+-  For example, if the value of ``HEAD_SIZE`` is 128 and
+   ``THREAD_GROUP_SIZE`` is 2, each thread's ``k_vecs`` will contain
+   total 64 elements. However, the returned ``qk`` is actually the
+   result of dot multiplication between 128 query elements and 128 key
+   elements. If you want to learn more about the details of the dot
+   multiplication and reduction, you may refer to the implementation of
+   ``Qk_dot<>::dot``. However, for the sake of simplicity, I will not
+   cover it in this document.
+
+Softmax
+-------
+
+-  Next, we need to calculate the normalized softmax for all ``qk``\ s,
+   as shown above, where each :math:`x` represents a ``qk``. To do this,
+   we must obtain the reduced value of ``qk_max``\ (:math:`m(x)`) and
+   the ``exp_sum``\ (:math:`\ell(x)`) of all ``qk``\ s. The reduction
+   should be performed across the entire thread block, encompassing
+   results between the query token and all context key tokens.
+
+   .. math::
+      :nowrap:
+
+      \begin{gather*}
+      m(x):=\max _i \quad x_i \\ \quad f(x):=\left[\begin{array}{lll}e^{x_1-m(x)} & \ldots & e^{x_B-m(x)}\end{array}\right]\\ \quad \ell(x):=\sum_i f(x)_i \\
+      \quad \operatorname{softmax}(x):=\frac{f(x)}{\ell(x)}
+      \end{gather*}
+
+``qk_max`` and ``logits``
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+-  Just right after we get the ``qk`` result, we can set the temporary
+   ``logits`` result with ``qk`` (In the end, the ``logits`` should
+   store the normalized softmax result). Also we can compare and collect
+   the ``qk_max`` for all ``qk``\ s that are calculated by current
+   thread group.
+
+   .. code:: cpp
+
+      if (thread_group_offset == 0) {
+         const bool mask = token_idx >= context_len;
+         logits[token_idx - start_token_idx] = mask ? 0.f : qk;
+         qk_max = mask ? qk_max : fmaxf(qk_max, qk);
+      }
+
+-  Please note that the ``logits`` here is on shared memory, so each
+   thread group will set the fields for its own assigned context tokens.
+   Overall, the size of logits should be number of context tokens.
+
+   .. code:: cpp
+
+      for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
+          qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
+      }
+
+      if (lane == 0) {
+         red_smem[warp_idx] = qk_max;
+      }
+
+-  Then we need to get the reduced ``qk_max`` across each warp. The main
+   idea is to make threads in warp to communicate with each other and
+   get the final max ``qk`` .
+
+   .. code:: cpp
+
+      for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
+          qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
+      }
+      qk_max = VLLM_SHFL_SYNC(qk_max, 0);
+
+-  Finally, we can get the reduced ``qk_max`` from whole thread block by
+   compare the ``qk_max`` from all warps in this thread block. Then we
+   need to broadcast the final result to each thread.
+
+``exp_sum``
+~~~~~~~~~~~
+
+-  Similar to ``qk_max``, we need to get the reduced sum value from the
+   entire thread block too.
+
+   .. code:: cpp
+
+      for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
+          float val = __expf(logits[i] - qk_max);
+          logits[i] = val;
+          exp_sum += val;
+      }
+      ...
+      exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);
+
+-  Firstly, sum all exp values from each thread group, and meanwhile,
+   convert each entry of ``logits`` from ``qk`` to ``exp(qk - qk_max)``.
+   Please note, the ``qk_max`` here is already the max ``qk`` across the
+   whole thread block. And then we can do reduction for ``exp_sum``
+   across whole thread block just like the ``qk_max``.
+
+   .. code:: cpp
+
+      const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
+      for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
+         logits[i] *= inv_sum;
+      }
+
+-  Finally, with the reduced ``qk_max`` and ``exp_sum``, we can obtain
+   the final normalized softmax result as ``logits``. This ``logits``
+   variable will be used for dot multiplication with the value data in
+   later steps. Now, it should store the normalized softmax result of
+   ``qk`` for all assigned context tokens.
+
+Value
+-----
+
+.. figure:: ../../assets/kernel/value.png
+   :alt: value
+   :width: 70%
+   :align: center
+
+   Value data of all context tokens at one head
+
+.. figure:: ../../assets/kernel/logits_vec.png
+   :alt: logits_vec
+   :width: 50%
+   :align: center
+
+   ``logits_vec`` for one thread
+
+.. figure:: ../../assets/kernel/v_vec.png
+   :alt: v_vec
+   :width: 70%
+   :align: center
+
+   List of ``v_vec`` for one thread
+
+-  Now we need to retrieve the value data and perform dot multiplication
+   with ``logits``. Unlike query and key, there is no thread group
+   concept for value data. As shown in diagram, different from key token
+   memory layout, elements from the same column correspond to the same
+   value token. For one block of value data, there are ``HEAD_SIZE`` of
+   rows and ``BLOCK_SIZE`` of columns that are split into multiple
+   ``v_vecs``.
+-  Each thread always fetches ``V_VEC_SIZE`` elements from the same
+   ``V_VEC_SIZE`` of tokens at a time. As a result, a single thread
+   retrieves multiple ``v_vec``\ s from different rows and the same
+   columns through multiple inner iterations. For each ``v_vec``, it
+   needs to be dot multiplied with the corresponding ``logits_vec``,
+   which is also ``V_VEC_SIZE`` elements from ``logits``. Overall, with
+   multiple inner iterations, each warp will process one block of value
+   tokens. And with multiple outer iterations, the whole context value
+   tokens are processd
+
+   .. code:: cpp
+
+      float accs[NUM_ROWS_PER_THREAD];
+      for ... { // Iteration over different blocks.
+          logits_vec = ...
+          for ... { // Iteration over different rows.
+              v_vec = ...
+              ...
+              accs[i] += dot(logits_vec, v_vec);
+          }
+      }
+
+-  As shown in the above pseudo code, in the outer loop, similar to
+   ``k_ptr``, ``logits_vec`` iterates over different blocks and reads
+   ``V_VEC_SIZE`` elements from ``logits``. In the inner loop, each
+   thread reads ``V_VEC_SIZE`` elements from the same tokens as a
+   ``v_vec`` and performs dot multiplication. It is important to note
+   that in each inner iteration, the thread fetches different head
+   position elements for the same tokens. The dot result is then
+   accumulated in ``accs``. Therefore, each entry of ``accs`` is mapped
+   to a head position assigned to the current thread.
+-  For example, if ``BLOCK_SIZE`` is 16 and ``V_VEC_SIZE`` is 8, each
+   thread fetches 8 value elements for 8 tokens at a time. Each element
+   is from different tokens at the same head position. If ``HEAD_SIZE``
+   is 128 and ``WARP_SIZE`` is 32, for each inner loop, a warp needs to
+   fetch ``WARP_SIZE * V_VEC_SIZE = 256`` elements. This means there are
+   a total of 128 \* 16 / 256 = 8 inner iterations for a warp to handle
+   a whole block of value tokens. And each ``accs`` in each thread
+   contains 8 elements that accumulated at 8 different head positions.
+   For the thread 0, the ``accs`` variable will have 8 elements, which
+   are 0th, 32th … 224th elements of a value head that are accumulated
+   from all assigned 8 tokens.
+
+LV
+---
+-  Now, we need to perform reduction for ``accs`` within each warp. This
+   process allows each thread to accumulate the ``accs`` for the
+   assigned head positions of all tokens in one block.
+
+   .. code:: cpp
+
+      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+         float acc = accs[i];
+         for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
+            acc += VLLM_SHFL_XOR_SYNC(acc, mask);
+         }
+         accs[i] = acc;
+      }
+
+-  Next, we perform reduction for ``accs`` across all warps, allowing
+   each thread to have the accumulation of ``accs`` for the assigned
+   head positions of all context tokens. Please note that each ``accs``
+   in every thread only stores the accumulation for a portion of
+   elements of the entire head for all context tokens. However, overall,
+   all results for output have been calculated but are just stored in
+   different thread register memory.
+
+   .. code:: cpp
+
+      float* out_smem = reinterpret_cast<float*>(shared_mem);
+      for (int i = NUM_WARPS; i > 1; i /= 2) {
+          // Upper warps write to shared memory.
+          ...
+              float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
+              for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+                      ...
+              dst[row_idx] = accs[i];
+          }
+
+          // Lower warps update the output.
+              const float* src = &out_smem[warp_idx * HEAD_SIZE];
+          for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+                      ...
+              accs[i] += src[row_idx];
+          }
+
+              // Write out the accs.
+      }
+
+Output
+------
+
+-  Now we can write all of calculated result from local register memory
+   to final output global memory.
+
+   .. code:: cpp
+
+      scalar_t* out_ptr = out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE
+                      + head_idx * max_num_partitions * HEAD_SIZE
+                      + partition_idx * HEAD_SIZE;
+
+-  First, we need to define the ``out_ptr`` variable, which points to
+   the start address of the assigned sequence and assigned head.
+
+   .. code:: cpp
+
+      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
+      if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
+          from_float(*(out_ptr + row_idx), accs[i]);
+      }
+      }
+
+-  Finally, we need to iterate over different assigned head positions
+   and write out the corresponding accumulated result based on the
+   ``out_ptr``.

docs/source/dev/sampling_params.rst

+Sampling Params
+===============
+
+.. automodule:: vllm.sampling_params.SamplingParams
\ No newline at end of file

docs/source/getting_started/amd-installation.rst

@@ -100,7 +100,7 @@ You can build and install vLLM from source:
 
 Build a docker image from `Dockerfile.rocm`, and launch a docker container.
 
-The `Dokerfile.rocm` is designed to support both ROCm 5.7 and ROCm 6.0 and later versions. It provides flexibility to customize the build of docker image using the following arguments:
+The `Dockerfile.rocm` is designed to support both ROCm 5.7 and ROCm 6.0 and later versions. It provides flexibility to customize the build of docker image using the following arguments:
 
 * `BASE_IMAGE`: specifies the base image used when running ``docker build``, specifically the PyTorch on ROCm base image. We have tested ROCm 5.7 and ROCm 6.0. The default is `rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1`
 * `FX_GFX_ARCHS`: specifies the GFX architecture that is used to build flash-attention, for example, `gfx90a;gfx942` for MI200 and MI300. The default is `gfx90a;gfx942`

docs/source/getting_started/installation.rst

@@ -60,6 +60,15 @@ You can also build and install vLLM from source:
     $ cd vllm
     $ pip install -e .  # This may take 5-10 minutes.
 
+.. tip::
+    To avoid your system being overloaded, you can limit the number of compilation jobs
+    to be run simultaneously, via the environment variable `MAX_JOBS`. For example:
+
+    .. code-block:: console
+
+        $ export MAX_JOBS=6
+        $ pip install -e .
+
 .. tip::
     If you have trouble building vLLM, we recommend using the NVIDIA PyTorch Docker image.
 

docs/source/getting_started/neuron-installation.rst

+.. _installation_neuron:
+
+Installation with Neuron
+========================
+
+vLLM 0.3.3 onwards supports model inferencing and serving on AWS Trainium/Inferentia with Neuron SDK.
+At the moment Paged Attention is not supported in Neuron SDK, but naive continuous batching is supported in transformers-neuronx.
+Data types currently supported in Neuron SDK are FP16 and BF16.
+
+Requirements
+------------
+
+* OS: Linux
+* Python: 3.8 -- 3.11
+* Accelerator: NeuronCore_v2 (in trn1/inf2 instances)
+* Pytorch 2.0.1/2.1.1
+* AWS Neuron SDK 2.16/2.17 (Verified on python 3.8)
+
+Installation steps:
+
+- :ref:`Build from source <build_from_source_neuron>`
+
+  - :ref:`Step 0. Launch Trn1/Inf2 instances <launch_instances>`
+  - :ref:`Step 1. Install drivers and tools <install_drivers>`
+  - :ref:`Step 2. Install transformers-neuronx and its dependencies <install_tnx>`
+  - :ref:`Step 3. Install vLLM from source <install_vllm>`
+
+.. _build_from_source_neuron:
+
+Build from source
+-----------------
+
+Following instructions are applicable to Neuron SDK 2.16 and beyond.
+
+.. _launch_instances:
+
+Step 0. Launch Trn1/Inf2 instances
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Here are the steps to launch trn1/inf2 instances, in order to install `PyTorch Neuron ("torch-neuronx") Setup on Ubuntu 22.04 LTS <https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/neuron-setup/pytorch/neuronx/ubuntu/torch-neuronx-ubuntu22.html>`_.
+
+- Please follow the instructions at `launch an Amazon EC2 Instance <https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/EC2_GetStarted.html#ec2-launch-instance>`_ to launch an instance. When choosing the instance type at the EC2 console, please make sure to select the correct instance type.
+- To get more information about instances sizes and pricing see: `Trn1 web page <https://aws.amazon.com/ec2/instance-types/trn1/>`_, `Inf2 web page <https://aws.amazon.com/ec2/instance-types/inf2/>`_
+- Select Ubuntu Server 22.04 TLS AMI
+- When launching a Trn1/Inf2, please adjust your primary EBS volume size to a minimum of 512GB.
+- After launching the instance, follow the instructions in `Connect to your instance <https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AccessingInstancesLinux.html>`_ to connect to the instance
+
+.. _install_drivers:
+
+Step 1. Install drivers and tools
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The installation of drivers and tools wouldn't be necessary, if `Deep Learning AMI Neuron <https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html>`_ is installed. In case the drivers and tools are not installed on the operating system, follow the steps below:
+
+.. code-block:: console
+
+    # Configure Linux for Neuron repository updates
+    . /etc/os-release
+    sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
+    deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main
+    EOF
+    wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
+
+    # Update OS packages
+    sudo apt-get update -y
+
+    # Install OS headers
+    sudo apt-get install linux-headers-$(uname -r) -y
+
+    # Install git
+    sudo apt-get install git -y
+
+    # install Neuron Driver
+    sudo apt-get install aws-neuronx-dkms=2.* -y
+
+    # Install Neuron Runtime
+    sudo apt-get install aws-neuronx-collectives=2.* -y
+    sudo apt-get install aws-neuronx-runtime-lib=2.* -y
+
+    # Install Neuron Tools
+    sudo apt-get install aws-neuronx-tools=2.* -y
+
+    # Add PATH
+    export PATH=/opt/aws/neuron/bin:$PATH
+
+
+.. _install_tnx:
+
+Step 2. Install transformers-neuronx and its dependencies
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+`transformers-neuronx <https://github.com/aws-neuron/transformers-neuronx>`_ will be the backend to support inference on trn1/inf2 instances.
+Follow the steps below to install transformer-neuronx package and its dependencies.
+
+.. code-block:: console
+
+    # Install Python venv
+    sudo apt-get install -y python3.10-venv g++
+
+    # Create Python venv
+    python3.10 -m venv aws_neuron_venv_pytorch
+
+    # Activate Python venv
+    source aws_neuron_venv_pytorch/bin/activate
+
+    # Install Jupyter notebook kernel
+    pip install ipykernel
+    python3.10 -m ipykernel install --user --name aws_neuron_venv_pytorch --display-name "Python (torch-neuronx)"
+    pip install jupyter notebook
+    pip install environment_kernels
+
+    # Set pip repository pointing to the Neuron repository
+    python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
+
+    # Install wget, awscli
+    python -m pip install wget
+    python -m pip install awscli
+
+    # Update Neuron Compiler and Framework
+    python -m pip install --upgrade neuronx-cc==2.* --pre torch-neuronx==2.1.* torchvision transformers-neuronx
+
+.. _install_vllm:
+
+Step 3. Install vLLM from source
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Once neuronx-cc and transformers-neuronx packages are installed, we will be able to install vllm as follows:
+
+.. code-block:: console
+
+    $ git clone https://github.com/vllm-project/vllm.git
+    $ cd vllm
+    $ pip install -U -r requirements-neuron.txt
+    $ pip install .
+
+If neuron packages are detected correctly in the installation process, ``vllm-0.3.0+neuron212`` will be installed.

docs/source/index.rst

@@ -62,19 +62,19 @@ Documentation
 
    getting_started/installation
    getting_started/amd-installation
+   getting_started/neuron-installation
    getting_started/quickstart
 
 .. toctree::
    :maxdepth: 1
    :caption: Serving
 
-   serving/distributed_serving
-   serving/run_on_sky
-   serving/deploying_with_kserve
-   serving/deploying_with_triton
+   serving/openai_compatible_server
    serving/deploying_with_docker
-   serving/serving_with_langchain
+   serving/distributed_serving
    serving/metrics
+   serving/usage_stats
+   serving/integrations
 
 .. toctree::
    :maxdepth: 1
@@ -96,7 +96,9 @@ Documentation
    :maxdepth: 2
    :caption: Developer Documentation
 
+   dev/sampling_params
    dev/engine/engine_index
+   dev/kernel/paged_attention
 
 Indices and tables
 ==================

docs/source/models/adding_model.rst

@@ -56,8 +56,8 @@ Next, you need to rewrite the :code:`forward` methods of your model by following
     -    return_dict: Optional[bool] = None,
     -) -> Union[Tuple, CausalLMOutputWithPast]:
     +    positions: torch.Tensor,
-    +    kv_caches: List[KVCache],
-    +    input_metadata: InputMetadata,
+    +    kv_caches: List[torch.Tensor],
+    +    attn_metadata: AttentionMetadata,
     +) -> Optional[SamplerOutput]:
 
 1. Update the code by considering that :code:`input_ids` and :code:`positions` are now flattened tensors.

docs/source/models/engine_args.rst

@@ -81,6 +81,10 @@ Below, you can find an explanation of every engine argument for vLLM:
 
     Token block size for contiguous chunks of tokens.
 
+.. option:: --enable-prefix-caching
+
+    Enables automatic prefix caching
+
 .. option:: --seed <seed>
 
     Random seed for operations.

docs/source/models/lora.rst

@@ -92,7 +92,8 @@ LoRA adapter requests if they were provided and ``max_loras`` is set high enough
 
 The following is an example request
 
-.. code-block::bash 
+.. code-block:: bash
+
     curl http://localhost:8000/v1/completions \
         -H "Content-Type: application/json" \
         -d '{

docs/source/models/supported_models.rst

@@ -8,84 +8,125 @@ The following is the list of model architectures that are currently supported by
 Alongside each architecture, we include some popular models that use it.
 
 .. list-table::
-  :widths: 25 25 50
+  :widths: 25 25 50 5
   :header-rows: 1
 
   * - Architecture
     - Models
     - Example HuggingFace Models
+    - :ref:`LoRA <lora>`
   * - :code:`AquilaForCausalLM`
     - Aquila
     - :code:`BAAI/Aquila-7B`, :code:`BAAI/AquilaChat-7B`, etc.
+    - ✅︎
   * - :code:`BaiChuanForCausalLM`
     - Baichuan
     - :code:`baichuan-inc/Baichuan2-13B-Chat`, :code:`baichuan-inc/Baichuan-7B`, etc.
+    - ✅︎
   * - :code:`ChatGLMModel`
     - ChatGLM
     - :code:`THUDM/chatglm2-6b`, :code:`THUDM/chatglm3-6b`, etc.
+    - ✅︎
+  * - :code:`CohereForCausalLM`
+    - Command-R
+    - :code:`CohereForAI/c4ai-command-r-v01`, etc.
+    - 
+  * - :code:`DbrxForCausalLM`
+    - DBRX
+    - :code:`databricks/dbrx-base`, :code:`databricks/dbrx-instruct`, etc.
+    - 
   * - :code:`DeciLMForCausalLM`
     - DeciLM
     - :code:`Deci/DeciLM-7B`, :code:`Deci/DeciLM-7B-instruct`, etc.
+    - 
   * - :code:`BloomForCausalLM`
     - BLOOM, BLOOMZ, BLOOMChat
     - :code:`bigscience/bloom`, :code:`bigscience/bloomz`, etc.
+    - 
   * - :code:`FalconForCausalLM`
     - Falcon
     - :code:`tiiuae/falcon-7b`, :code:`tiiuae/falcon-40b`, :code:`tiiuae/falcon-rw-7b`, etc.
+    - 
   * - :code:`GemmaForCausalLM`
     - Gemma
     - :code:`google/gemma-2b`, :code:`google/gemma-7b`, etc.
+    - ✅︎
   * - :code:`GPT2LMHeadModel`
     - GPT-2
     - :code:`gpt2`, :code:`gpt2-xl`, etc.
+    - 
   * - :code:`GPTBigCodeForCausalLM`
     - StarCoder, SantaCoder, WizardCoder
     - :code:`bigcode/starcoder`, :code:`bigcode/gpt_bigcode-santacoder`, :code:`WizardLM/WizardCoder-15B-V1.0`, etc.
+    - 
   * - :code:`GPTJForCausalLM`
     - GPT-J
     - :code:`EleutherAI/gpt-j-6b`, :code:`nomic-ai/gpt4all-j`, etc.
+    - 
   * - :code:`GPTNeoXForCausalLM`
     - GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM
     - :code:`EleutherAI/gpt-neox-20b`, :code:`EleutherAI/pythia-12b`, :code:`OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, :code:`databricks/dolly-v2-12b`, :code:`stabilityai/stablelm-tuned-alpha-7b`, etc.
+    - 
   * - :code:`InternLMForCausalLM`
     - InternLM
     - :code:`internlm/internlm-7b`, :code:`internlm/internlm-chat-7b`, etc.
+    - ✅︎
   * - :code:`InternLM2ForCausalLM`
     - InternLM2
     - :code:`internlm/internlm2-7b`, :code:`internlm/internlm2-chat-7b`, etc.
+    -
+  * - :code:`JAISLMHeadModel`
+    - Jais
+    - :code:`core42/jais-13b`, :code:`core42/jais-13b-chat`, :code:`core42/jais-30b-v3`, :code:`core42/jais-30b-chat-v3`, etc.
+    -
   * - :code:`LlamaForCausalLM`
     - LLaMA, LLaMA-2, Vicuna, Alpaca, Yi
     - :code:`meta-llama/Llama-2-13b-hf`, :code:`meta-llama/Llama-2-70b-hf`, :code:`openlm-research/open_llama_13b`, :code:`lmsys/vicuna-13b-v1.3`, :code:`01-ai/Yi-6B`, :code:`01-ai/Yi-34B`, etc.
+    - ✅︎
   * - :code:`MistralForCausalLM`
     - Mistral, Mistral-Instruct
     - :code:`mistralai/Mistral-7B-v0.1`, :code:`mistralai/Mistral-7B-Instruct-v0.1`, etc.
+    - ✅︎
   * - :code:`MixtralForCausalLM`
     - Mixtral-8x7B, Mixtral-8x7B-Instruct
     - :code:`mistralai/Mixtral-8x7B-v0.1`, :code:`mistralai/Mixtral-8x7B-Instruct-v0.1`, etc.
+    - ✅︎
   * - :code:`MPTForCausalLM`
     - MPT, MPT-Instruct, MPT-Chat, MPT-StoryWriter
     - :code:`mosaicml/mpt-7b`, :code:`mosaicml/mpt-7b-storywriter`, :code:`mosaicml/mpt-30b`, etc.
+    - 
   * - :code:`OLMoForCausalLM`
     - OLMo
     - :code:`allenai/OLMo-1B`, :code:`allenai/OLMo-7B`, etc.
+    - 
   * - :code:`OPTForCausalLM`
     - OPT, OPT-IML
     - :code:`facebook/opt-66b`, :code:`facebook/opt-iml-max-30b`, etc.
+    - 
   * - :code:`OrionForCausalLM`
     - Orion
     - :code:`OrionStarAI/Orion-14B-Base`, :code:`OrionStarAI/Orion-14B-Chat`, etc.
+    - 
   * - :code:`PhiForCausalLM`
     - Phi
     - :code:`microsoft/phi-1_5`, :code:`microsoft/phi-2`, etc.
+    - 
   * - :code:`QWenLMHeadModel`
     - Qwen
     - :code:`Qwen/Qwen-7B`, :code:`Qwen/Qwen-7B-Chat`, etc.
+    - 
   * - :code:`Qwen2ForCausalLM`
     - Qwen2
     - :code:`Qwen/Qwen2-beta-7B`, :code:`Qwen/Qwen2-beta-7B-Chat`, etc.
+    - ✅︎
+  * - :code:`Qwen2MoeForCausalLM`
+    - Qwen2MoE
+    - :code:`Qwen/Qwen1.5-MoE-A2.7B`, :code:`Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc.
+    - 
   * - :code:`StableLmForCausalLM`
     - StableLM
     - :code:`stabilityai/stablelm-3b-4e1t/` , :code:`stabilityai/stablelm-base-alpha-7b-v2`, etc.
+    - 
 
 If your model uses one of the above model architectures, you can seamlessly run your model with vLLM.
 Otherwise, please refer to :ref:`Adding a New Model <adding_a_new_model>` for instructions on how to implement support for your model.

docs/source/serving/deploying_with_bentoml.rst

+.. _deploying_with_bentoml:
+
+Deploying with BentoML
+======================
+
+`BentoML <https://github.com/bentoml/BentoML>`_ allows you to deploy a large language model (LLM) server with vLLM as the backend, which exposes OpenAI-compatible endpoints. You can serve the model locally or containerize it as an OCI-complicant image and deploy it on Kubernetes.
+
+For details, see the tutorial `vLLM inference in the BentoML documentation <https://docs.bentoml.com/en/latest/use-cases/large-language-models/vllm.html>`_.
\ No newline at end of file

docs/source/serving/integrations.rst

+Integrations
+------------
+
+.. toctree::
+   :maxdepth: 1
+
+   run_on_sky
+   deploying_with_kserve
+   deploying_with_triton
+   deploying_with_bentoml
+   serving_with_langchain

docs/source/serving/openai_compatible_server.md

+# OpenAI Compatible Server
+
+vLLM provides an HTTP server that implements OpenAI's [Completions](https://platform.openai.com/docs/api-reference/completions) and [Chat](https://platform.openai.com/docs/api-reference/chat) API.
+
+You can start the server using Python, or using [Docker](deploying_with_docker.rst):
+```bash
+python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-2-7b-hf --dtype float32 --api-key token-abc123
+```
+
+To call the server, you can use the official OpenAI Python client library, or any other HTTP client.
+```python
+from openai import OpenAI
+client = OpenAI(
+    base_url="http://localhost:8000/v1",
+    api_key="token-abc123",
+)
+
+completion = client.chat.completions.create(
+  model="meta-llama/Llama-2-7b-hf",
+  messages=[
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Hello!"}
+  ]
+)
+
+print(completion.choices[0].message)
+```
+
+## API Reference
+Please see the [OpenAI API Reference](https://platform.openai.com/docs/api-reference) for more information on the API. We support all parameters except:
+- Chat: `tools`, and `tool_choice`.
+- Completions: `suffix`.
+
+## Extra Parameters
+vLLM supports a set of parameters that are not part of the OpenAI API.
+In order to use them, you can pass them as extra parameters in the OpenAI client.
+Or directly merge them into the JSON payload if you are using HTTP call directly.
+
+```python
+completion = client.chat.completions.create(
+  model="meta-llama/Llama-2-7b-hf",
+  messages=[
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Classify this sentiment: vLLM is wonderful!"}
+  ],
+  extra_body={
+    "guided_choice": ["positive", "negative"]
+  }
+)
+```
+
+### Extra Parameters for Chat API
+The following [sampling parameters (click through to see documentation)](../dev/sampling_params.rst) are supported.
+
+```{literalinclude} ../../../vllm/entrypoints/openai/protocol.py
+:language: python
+:start-after: begin-chat-completion-sampling-params
+:end-before: end-chat-completion-sampling-params
+```
+
+The following extra parameters are supported:
+
+```{literalinclude} ../../../vllm/entrypoints/openai/protocol.py
+:language: python
+:start-after: begin-chat-completion-extra-params
+:end-before: end-chat-completion-extra-params
+```
+
+### Extra Parameters for Completions API
+The following [sampling parameters (click through to see documentation)](../dev/sampling_params.rst) are supported.
+
+```{literalinclude} ../../../vllm/entrypoints/openai/protocol.py
+:language: python
+:start-after: begin-completion-sampling-params
+:end-before: end-completion-sampling-params
+```
+
+The following extra parameters are supported:
+
+```{literalinclude} ../../../vllm/entrypoints/openai/protocol.py
+:language: python
+:start-after: begin-completion-extra-params
+:end-before: end-completion-extra-params
+```
+
+## Chat Template
+
+In order for the language model to support chat protocol, vLLM requires the model to include
+a chat template in its tokenizer configuration. The chat template is a Jinja2 template that
+specifies how are roles, messages, and other chat-specific tokens are encoded in the input.
+
+An example chat template for `meta-llama/Llama-2-7b-chat-hf` can be found [here](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf/blob/09bd0f49e16738cdfaa6e615203e126038736eb0/tokenizer_config.json#L12)
+
+Some models do not provide a chat template even though they are instruction/chat fine-tuned. For those model,
+you can manually specify their chat template in the `--chat-template` parameter with the file path to the chat
+template, or the template in string form. Without a chat template, the server will not be able to process chat
+and all chat requests will error.
+
+```bash
+python -m vllm.entrypoints.openai.api_server \
+  --model ... \
+  --chat-template ./path-to-chat-template.jinja
+```
+
+vLLM community provides a set of chat templates for popular models. You can find them in the examples
+directory [here](https://github.com/vllm-project/vllm/tree/main/examples/)
+
+## Command line arguments for the server
+
+```{argparse}
+:module: vllm.entrypoints.openai.cli_args
+:func: make_arg_parser
+:prog: vllm-openai-server
+```
\ No newline at end of file

docs/source/serving/usage_stats.md

+# Usage Stats Collection
+
+vLLM collects anonymous usage data by default to help the engineering team better understand which hardware and model configurations are widely used. This data allows them to prioritize their efforts on the most common workloads. The collected data is transparent, does not contain any sensitive information, and will be publicly released for the community's benefit.
+
+## What data is collected?
+
+You can see the up to date list of data collected by vLLM in the [usage_lib.py](https://github.com/vllm-project/vllm/blob/main/vllm/usage/usage_lib.py).
+
+Here is an example as of v0.4.0:
+
+```json
+{
+  "uuid": "fbe880e9-084d-4cab-a395-8984c50f1109",
+  "provider": "GCP",
+  "num_cpu": 24,
+  "cpu_type": "Intel(R) Xeon(R) CPU @ 2.20GHz",
+  "cpu_family_model_stepping": "6,85,7",
+  "total_memory": 101261135872,
+  "architecture": "x86_64",
+  "platform": "Linux-5.10.0-28-cloud-amd64-x86_64-with-glibc2.31",
+  "gpu_count": 2,
+  "gpu_type": "NVIDIA L4",
+  "gpu_memory_per_device": 23580639232,
+  "model_architecture": "OPTForCausalLM",
+  "vllm_version": "0.3.2+cu123",
+  "context": "LLM_CLASS",
+  "log_time": 1711663373492490000,
+  "source": "production",
+  "dtype": "torch.float16",
+  "tensor_parallel_size": 1,
+  "block_size": 16,
+  "gpu_memory_utilization": 0.9,
+  "quantization": null,
+  "kv_cache_dtype": "auto",
+  "enable_lora": false,
+  "enable_prefix_caching": false,
+  "enforce_eager": false,
+  "disable_custom_all_reduce": true
+}
+```
+
+You can preview the collected data by running the following command:
+
+```bash
+tail ~/.config/vllm/usage_stats.json
+```
+
+## Opt-out of Usage Stats Collection
+
+You can opt-out of usage stats collection by setting the VLLM_NO_USAGE_STATS or DO_NOT_TRACK environment variable, or by creating a ~/.config/vllm/do_not_track file:
+
+```bash
+# Any of the following methods can disable usage stats collection
+export VLLM_NO_USAGE_STATS=1
+export DO_NOT_TRACK=1
+mkdir -p ~/.config/vllm && touch ~/.config/vllm/do_not_track
+```

examples/gradio_openai_chatbot_webserver.py

 import argparse
-from openai import OpenAI
+
 import gradio as gr
+from openai import OpenAI
 
 # Argument parser setup
 parser = argparse.ArgumentParser(

examples/llava_example.py

+import argparse
+import os
+import subprocess
+
+import torch
+
+from vllm import LLM
+from vllm.sequence import MultiModalData
+
+# The assets are located at `s3://air-example-data-2/vllm_opensource_llava/`.
+
+
+def run_llava_pixel_values():
+    llm = LLM(
+        model="llava-hf/llava-1.5-7b-hf",
+        image_input_type="pixel_values",
+        image_token_id=32000,
+        image_input_shape="1,3,336,336",
+        image_feature_size=576,
+    )
+
+    prompt = "<image>" * 576 + (
+        "\nUSER: What is the content of this image?\nASSISTANT:")
+
+    # This should be provided by another online or offline component.
+    images = torch.load("images/stop_sign_pixel_values.pt")
+
+    outputs = llm.generate(prompt,
+                           multi_modal_data=MultiModalData(
+                               type=MultiModalData.Type.IMAGE, data=images))
+    for o in outputs:
+        generated_text = o.outputs[0].text
+        print(generated_text)
+
+
+def run_llava_image_features():
+    llm = LLM(
+        model="llava-hf/llava-1.5-7b-hf",
+        image_input_type="image_features",
+        image_token_id=32000,
+        image_input_shape="1,576,1024",
+        image_feature_size=576,
+    )
+
+    prompt = "<image>" * 576 + (
+        "\nUSER: What is the content of this image?\nASSISTANT:")
+
+    # This should be provided by another online or offline component.
+    images = torch.load("images/stop_sign_image_features.pt")
+
+    outputs = llm.generate(prompt,
+                           multi_modal_data=MultiModalData(
+                               type=MultiModalData.Type.IMAGE, data=images))
+    for o in outputs:
+        generated_text = o.outputs[0].text
+        print(generated_text)
+
+
+def main(args):
+    if args.type == "pixel_values":
+        run_llava_pixel_values()
+    else:
+        run_llava_image_features()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Demo on Llava")
+    parser.add_argument("--type",
+                        type=str,
+                        choices=["pixel_values", "image_features"],
+                        default="pixel_values",
+                        help="image input type")
+    args = parser.parse_args()
+    # Download from s3
+    s3_bucket_path = "s3://air-example-data-2/vllm_opensource_llava/"
+    local_directory = "images"
+
+    # Make sure the local directory exists or create it
+    os.makedirs(local_directory, exist_ok=True)
+
+    # Use AWS CLI to sync the directory, assume anonymous access
+    subprocess.check_call([
+        "aws",
+        "s3",
+        "sync",
+        s3_bucket_path,
+        local_directory,
+        "--no-sign-request",
+    ])
+    main(args)

examples/llm_engine_example.py

 import argparse
 from typing import List, Tuple
 
-from vllm import EngineArgs, LLMEngine, SamplingParams, RequestOutput
+from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
 
 
 def create_test_prompts() -> List[Tuple[str, SamplingParams]]:

examples/multilora_inference.py

 """
-This example shows how to use the multi-LoRA functionality for offline inference.
+This example shows how to use the multi-LoRA functionality
+for offline inference.
 
 Requires HuggingFace credentials for access to Llama2.
 """
 
-from typing import Optional, List, Tuple
+from typing import List, Optional, Tuple
 
 from huggingface_hub import snapshot_download
 
-from vllm import EngineArgs, LLMEngine, SamplingParams, RequestOutput
+from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
 from vllm.lora.request import LoRARequest
 
 
@@ -34,14 +35,16 @@ def create_test_prompts(
                         top_k=5,
                         presence_penalty=0.2,
                         max_tokens=128), None),
-        ("[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]",
+        (
+            "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]",  # noqa: E501
             SamplingParams(temperature=0.0,
                            logprobs=1,
                            prompt_logprobs=1,
                            max_tokens=128,
                            stop_token_ids=[32003]),
             LoRARequest("sql-lora", 1, lora_path)),
-        ("[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]",
+        (
+            "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]",  # noqa: E501
             SamplingParams(n=3,
                            best_of=3,
                            use_beam_search=True,
@@ -49,14 +52,16 @@ def create_test_prompts(
                            max_tokens=128,
                            stop_token_ids=[32003]),
             LoRARequest("sql-lora", 1, lora_path)),
-        ("[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]",
+        (
+            "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]",  # noqa: E501
             SamplingParams(temperature=0.0,
                            logprobs=1,
                            prompt_logprobs=1,
                            max_tokens=128,
                            stop_token_ids=[32003]),
             LoRARequest("sql-lora2", 2, lora_path)),
-        ("[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]",
+        (
+            "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]",  # noqa: E501
             SamplingParams(n=3,
                            best_of=3,
                            use_beam_search=True,

examples/offline_inference_distributed.py

@@ -5,11 +5,13 @@ distributively on a multi-nodes cluster.
 Learn more about Ray Data in https://docs.ray.io/en/latest/data/data.html
 """
 
-from vllm import LLM, SamplingParams
 from typing import Dict
+
 import numpy as np
 import ray
 
+from vllm import LLM, SamplingParams
+
 # Create a sampling params object.
 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 

examples/offline_inference_neuron.py

@@ -12,17 +12,20 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
 # Create an LLM.
 llm = LLM(
-    model="openlm-research/open_llama_3b",
+    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
     max_num_seqs=8,
-    # The max_model_len and block_size arguments are required to be same as max sequence length,
-    # when targeting neuron device. Currently, this is a known limitation in continuous batching
-    # support in transformers-neuronx.
+    # The max_model_len and block_size arguments are required to be same as
+    # max sequence length when targeting neuron device.
+    # Currently, this is a known limitation in continuous batching support
+    # in transformers-neuronx.
     # TODO(liangfu): Support paged-attention in transformers-neuronx.
     max_model_len=128,
     block_size=128,
     # The device can be automatically detected when AWS Neuron SDK is installed.
-    # The device argument can be either unspecified for automated detection, or explicitly assigned.
-    device="neuron")
+    # The device argument can be either unspecified for automated detection,
+    # or explicitly assigned.
+    device="neuron",
+    tensor_parallel_size=2)
 # Generate texts from the prompts. The output is a list of RequestOutput objects
 # that contain the prompt, generated text, and other information.
 outputs = llm.generate(prompts, sampling_params)