chore(sdk): upgrade to Neuron SDK 2.31

.dockerignore

@@ -0,0 +1,14 @@
+.git
+.github
+.venv
+aws_neuron_venv_pytorch
+__pycache__
+*.pyc
+.ruff_cache
+.mypy_cache
+.pytest_cache
+build
+*.egg-info
+docs
+notebooks
+benchmark

.github/actions/install_neuronx_runtime/action.yml

@@ -12,7 +12,7 @@ runs:
           EOF
           wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
           sudo apt-get update -y
-          sudo apt-get install aws-neuronx-tools=2.26.14.0 aws-neuronx-runtime-lib=2.28.23.0-dd5879008 aws-neuronx-collectives=2.28.27.0-bc30ece58 -y
+          sudo apt-get install aws-neuronx-tools=2.31.13.0-a9e473f33 aws-neuronx-runtime-lib=2.33.10.0-3dcef56f0 aws-neuronx-collectives=2.33.10.0-068180c7a libpython3.12 libarchive13 -y
           export PATH=/opt/aws/neuron/bin:$PATH
           dpkg -l | grep neuron
       - name: Display driver version

.github/actions/sanity-check/action.yml

@@ -33,9 +33,9 @@ runs:
         MODEL_ID: llamafactory/tiny-random-qwen3
       run: |
         # Manually install torch to force CPU-only installation and speed up installation
-        uv venv --python 3.11 on-no-neuronx
+        uv venv --python 3.12 on-no-neuronx
         source on-no-neuronx/bin/activate
-        uv pip install torch==2.8.0 torchvision~=0.23 --index-url https://download.pytorch.org/whl/cpu
+        uv pip install torch==2.9.1 torchvision~=0.24 --index-url https://download.pytorch.org/whl/cpu
         uv pip install .
         # Check that the model is cached
         HF_TOKEN=${{ inputs.hf_token }} \

.github/actions/setup_venv/action.yml

@@ -10,7 +10,7 @@ runs:
       - name: Prepare venv and install Optimum Neuron python package
         shell: bash
         run: |
-          uv venv --python 3.11 aws_neuron_venv_pytorch
+          uv venv --python 3.12 aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           uv pip install .[neuronx,tests]
           # Enable high performance with Xet for all workflows that use this action

.github/workflows/cache_diffusion.yml

@@ -25,7 +25,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/cache_llm.yml

@@ -25,7 +25,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/disabled/test_trainium_training.yml

@@ -27,7 +27,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v4

.github/workflows/doc-build.yml

@@ -17,7 +17,7 @@ on:
 
 jobs:
   build_documentation:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     env:
       COMMIT_SHA: ${{ github.event.pull_request.head.sha }}
       PR_NUMBER: ${{ github.event.number }}

.github/workflows/doc-pr-build.yml

@@ -16,7 +16,7 @@ concurrency:
 
 jobs:
   build_documentation:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     env:
       COMMIT_SHA: ${{ github.event.pull_request.head.sha }}
       PR_NUMBER: ${{ github.event.number }}

.github/workflows/test_cpu_only.yml

@@ -55,7 +55,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6
@@ -67,7 +67,7 @@ jobs:
   test-cpu:
     name: Run CPU Only Tests
     needs: sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/test_inf2_diffusers.yml

@@ -39,7 +39,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/test_inf2_export.yml

@@ -47,7 +47,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/test_inf2_llm.yml

@@ -67,7 +67,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/test_inf2_seq2seq.yml

@@ -35,7 +35,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/test_inf2_slow.yml

@@ -47,7 +47,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/test_inf2_transformers.yml

@@ -45,7 +45,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/test_inf2_vllm.yml

@@ -69,7 +69,7 @@ concurrency:
 jobs:
   sanity:
     name: Sanity
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout
         uses: actions/checkout@v6

.github/workflows/test_sagemaker.yml

@@ -18,7 +18,7 @@ on:
 jobs:
   do-the-job:
     name: Run Sagemaker Related Tests
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     permissions:
       id-token: write   # required for OIDC
       contents: read

.gitignore

@@ -137,3 +137,6 @@ neuronxcc*/
 
 # Ignore claude settings
 .claude/
+
+# Neuron compiler artifact generated under SDK 2.31
+tests/PostSPMDPassesExecutionDuration.txt

docker/vllm/Dockerfile

@@ -1,4 +1,4 @@
-FROM ubuntu:22.04 AS base
+FROM ubuntu:24.04 AS base
 
 # Install system prerequisites
 RUN apt-get update -y \
@@ -11,23 +11,36 @@ RUN apt-get update -y \
     wget \
     libexpat1 \
     libpython3-dev \
+    libarchive13 \
     && rm -rf /var/lib/apt/lists/* \
     && apt-get clean
 
 # Install uv at a specific version on a given path
 RUN curl -LsSf https://astral.sh/uv/0.9.27/install.sh | XDG_BIN_HOME=/usr/local/bin sh
 
+# Provision a standalone Python 3.12 interpreter via uv.
+# Use 3.12, not 3.11: neuronx-cc pins numpy<2 for python_full_version < '3.12', which
+# conflicts with vllm's numpy>=2 requirement.
+RUN uv venv --python 3.12 /opt/venv
+ENV VIRTUAL_ENV=/opt/venv
+ENV PATH="/opt/venv/bin:${PATH}"
+# torch_xla's compiled extension dynamically links against libpython, which uv's
+# standalone interpreter ships but doesn't register with the dynamic linker.
+RUN dirname "$(find /root/.local/share/uv/python -name 'libpython3.12.so.1.0')" \
+    > /etc/ld.so.conf.d/uv-python.conf \
+    && ldconfig
+
 # Setup neuronx repository
-RUN echo "deb https://apt.repos.neuron.amazonaws.com jammy main" > /etc/apt/sources.list.d/neuron.list
+RUN echo "deb https://apt.repos.neuron.amazonaws.com noble main" > /etc/apt/sources.list.d/neuron.list
 RUN wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
 
 # Install neuronx packages
 RUN apt-get update -y \
     && apt-get install -y --no-install-recommends \
-    aws-neuronx-dkms=2.24.7.0 \
-    aws-neuronx-collectives=2.28.27.0-bc30ece58 \
-    aws-neuronx-runtime-lib=2.28.23.0-dd5879008 \
-    aws-neuronx-tools=2.26.14.0 \
+    aws-neuronx-dkms=2.29.0.0 \
+    aws-neuronx-collectives=2.33.10.0-068180c7a \
+    aws-neuronx-runtime-lib=2.33.10.0-3dcef56f0 \
+    aws-neuronx-tools=2.31.13.0-a9e473f33 \
     && rm -rf /var/lib/apt/lists/* \
     && apt-get clean
 
@@ -38,7 +51,7 @@ RUN mkdir optimum-neuron
 COPY optimum optimum-neuron/optimum
 COPY pyproject.toml optimum-neuron/pyproject.toml
 RUN ls optimum-neuron
-RUN cd optimum-neuron && uv pip install --system .[neuronx,vllm]
+RUN cd optimum-neuron && uv pip install .[neuronx,vllm]
 
 # HF base env
 ENV HUGGINGFACE_HUB_CACHE=/tmp \

docs/source/contribute/dev_environment.mdx

@@ -27,7 +27,7 @@ $ python3 -m venv .venv
 $ source .venv/bin/activate
 ```
 
-Note: `optimum-neuron` requires at least python 3.10
+Note: `optimum-neuron` requires Python 3.12 (see `pyproject.toml` for the exact range)
 
 ## Install development tools
 

optimum/commands/neuron/serve.py

@@ -38,7 +38,7 @@
 
     from vllm.entrypoints.openai.api_server import run_server
     from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
-    from vllm.utils import FlexibleArgumentParser
+    from vllm.utils.argparse_utils import FlexibleArgumentParser
 
     from ...neuron.vllm.model_loader import VLLM_2_TRANSFORMERS_TASK_MAPPING
     from ...neuron.vllm.reverse_proxy import RoundRobinProxy
@@ -290,13 +290,16 @@ def run(self):
                 )
 
         # Build the vLLM command arguments.
+        # vLLM's --task flag was replaced by --runner (+ --convert, left at its
+        # "auto" default): "generate" maps directly, "embed" is a pooling runner.
+        vllm_runner = "pooling" if self.args.task == "embed" else self.args.task
         vllm_command = [
             "--model",
             self.args.model,
             "--served_model_name",
             model_id,
-            "--task",
-            self.args.task,
+            "--runner",
+            vllm_runner,
             "--tensor-parallel-size",
             str(tensor_parallel_size),
             "--max-num-seqs",

optimum/exporters/neuron/model_wrappers.py

@@ -463,10 +463,15 @@ def update_past(self, past_key_values):
         return new_past_sa, new_past_ca
 
     def reorder_cache(self, past_key_values, beam_idx):
-        for i in range(len(past_key_values)):
-            gather_index = beam_idx.view([beam_idx.shape[0], 1, 1, 1]).expand_as(past_key_values[i])
-            past_key_values[i] = torch.gather(past_key_values[i], dim=0, index=gather_index)
-        return past_key_values
+        # Do not assign into `past_key_values`: it is the module `ParameterList` holding the KV
+        # cache, and overwriting its entries during the trace unregisters the parameters. The
+        # tracer would then be unable to restore them, and would leak XLA placeholder tensors into
+        # the input/output aliases, which cannot be sent back to the parent process.
+        reordered = []
+        for past_key_value in past_key_values:
+            gather_index = beam_idx.view([beam_idx.shape[0], 1, 1, 1]).expand_as(past_key_value)
+            reordered.append(torch.gather(past_key_value, dim=0, index=gather_index))
+        return reordered
 
     def forward(
         self,

optimum/neuron/models/inference/backend/modules/attention/flash_attention_nki.py

@@ -156,8 +156,14 @@ def _flash_attention_core_large_d(
     o_previous_scaled = nl.ndarray((par_dim(B_P_SIZE), d), dtype=o_buffer.dtype)
     o_previous_scaled[...] = nl.multiply(o_buffer[:, :], alpha)
 
-    # Compute exp(QK - max) and partial sums
-    p_local = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ), dtype=kernel_dtype)
+    # Compute exp(QK - max) and partial sums.
+    # The softmax numerator is kept in acc_type (fp32) all the way through the
+    # transpose and into the PV matmul below. Rounding it to kernel_dtype quantizes
+    # every attention weight to 8 mantissa bits, and that error accumulates over the
+    # attended positions: at 5k tokens it is enough to change a sampled token versus
+    # the CPU fp32 reference. Prefill is ~1.6x slower this way, still well ahead of
+    # the compiler-native path.
+    p_local = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ), dtype=acc_type)
     REDUCTION_TILE = min(2048, LARGE_TILE_SZ // 2)
     p_partial_sum = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ // REDUCTION_TILE), dtype=acc_type)
 
@@ -170,13 +176,13 @@ def _flash_attention_core_large_d(
             scale=1.0,
             reduce_op=nl.add,
             reduce_res=p_partial_sum[:, k_r_i],
-            dtype=kernel_dtype,
+            dtype=acc_type,
         )
 
     ps = nl.sum(p_partial_sum, axis=1, dtype=acc_type)
 
     # Transpose p_local for PV matmul
-    p_local_transposed = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ), dtype=kernel_dtype)
+    p_local_transposed = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ), dtype=acc_type)
     _transpose_p_local(
         p_local_transposed=p_local_transposed,
         p_local=p_local,

optimum/neuron/models/inference/backend/modules/generation/generation_utils.py

@@ -168,13 +168,16 @@ def sample_next_tokens(outputs: torch.Tensor, is_ods: bool | None = None) -> tor
             else:
                 next_token_logits = outputs[:, -1, :].clone()
                 next_token_scores = logits_processor(input_ids, next_token_logits)
-                next_token_scores, next_token_indices = fused_logits_warper(next_token_scores)
                 if do_sample:
+                    next_token_scores, next_token_indices = fused_logits_warper(next_token_scores)
                     probs = torch.nn.functional.softmax(next_token_scores, dim=-1)
                     next_tokens = torch.multinomial(probs, num_samples=1)
+                    next_tokens = torch.gather(next_token_indices, 1, next_tokens).squeeze(1)
                 else:
-                    next_tokens = torch.argmax(next_token_scores, dim=-1, keepdim=True)
-                next_tokens = torch.gather(next_token_indices, 1, next_tokens).squeeze(1)
+                    # Greedy: select from the full logits, so that ties are broken towards the
+                    # lowest token id, like transformers does. Going through the fused warper
+                    # would instead break them according to its top-k sort order.
+                    next_tokens = torch.argmax(next_token_scores, dim=-1)
 
             if has_eos_stopping_criteria:
                 next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)

optimum/neuron/models/inference/backend/pretrained_model.py

@@ -16,6 +16,7 @@
 import logging
 import os
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from functools import partial
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -102,6 +103,26 @@ def get_builder(
     return builder
 
 
+@contextmanager
+def _scratch_compile_cwd():
+    """Run neuronx-cc from a throwaway working directory.
+
+    The NxD weight-layout-optimization step compiles NKI kernels via
+    ``torch_neuronx.xla_impl.trace.hlo_compile``, which invokes the compiler
+    with ``subprocess.run(command)`` and no ``cwd=`` (trace.py). The compiler
+    backend (``walrus_driver``) then materializes content-addressed
+    ``neuronxcc.private_nkl.*`` kernel directories in the process CWD. Run the
+    compile from a temp dir so those droppings are discarded.
+    """
+    prev = os.getcwd()
+    with TemporaryDirectory() as tmp:
+        os.chdir(tmp)
+        try:
+            yield
+        finally:
+            os.chdir(prev)
+
+
 class NxDPreTrainedModel(NeuronPreTrainedModel, ABC):
     _STATE_DICT_MODEL_PREFIX = "model."
     _NEW_STATE_DICT_MODEL_PREFIX = ""
@@ -155,7 +176,8 @@ def compile(
         for bundle_name, bundle_builders in graph_builders.items():
             logger.info(f"Compiling bundle '{bundle_name}' with graphs: {list(bundle_builders.keys())}")
             builder = get_builder(neuron_config, bundle_builders, debug=debug, compiler_args=compiler_args)
-            traced_models[bundle_name] = builder.trace(initialize_model_weights=False)
+            with _scratch_compile_cwd():
+                traced_models[bundle_name] = builder.trace(initialize_model_weights=False)
         return traced_models
 
     @staticmethod