optimum/neuron/models/inference/backend/modules/decoder/modeling_decoder.py

# coding=utf-8
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import logging
import os
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import List, Optional, Tuple, Union

import neuronx_distributed as nxd
import torch
import torch_xla.core.xla_model as xm
from huggingface_hub import HfApi, snapshot_download
from neuronx_distributed.operators.argmax import argmax as nxd_argmax
from neuronx_distributed.parallel_layers.layers import SPMDRank
from neuronx_distributed.parallel_layers.mappings import (
    _gather_along_dim,
    _reduce_scatter_along_dim,
    gather_from_sequence_parallel_region,
)
from torch import nn
from transformers import AutoConfig, PretrainedConfig
from transformers.modeling_outputs import CausalLMOutputWithPast

from ......cache.entries.single_model import SingleModelCacheEntry
from ......cache.hub_cache import hub_neuronx_cache
from ......modeling_decoder import NeuronModelForCausalLM
from ...config import NxDNeuronConfig
from ...pretrained_model import NxDPreTrainedModel
from ...utils.random import set_random_seed
from ..attention import utils as attn_utils
from ..autobucketing import generate_buckets
from ..flashdecode.utils import (
    get_cache_size,
    mask_util,
    turn_2d_mask_to_4d,
)
from ..generation.generation_utils import NxDGenerationMixin
from ..generation.sampling import (
    Sampler,
    mask_padded_logits,
    prepare_sampling_params,
    validate_sampling_params,
)
from ..kvcache.kv_cache_manager import (
    KVCacheManager,
    _slice_kv_cacheline,
)
from .decoder_wrapper import (
    CONTEXT_ENCODING_MODEL_TAG,
    SPECULATION_MODEL_TAG,
    TOKEN_GENERATION_MODEL_TAG,
    NxDDecoderWrapper,
)


logger = logging.getLogger("Neuron")


class NxDDecoderModel(nn.Module):
    """
    Base model that NeuronXXXModel classes inherit from.

    The forward() function will be traced and compiled by NxD.
    """

    def __init__(self, config: PretrainedConfig, neuron_config: NxDNeuronConfig):
        super().__init__()

        self.config = config
        self.sampler = None
        self.kv_mgr = None
        self.neuron_config = neuron_config
        self.batch_size = neuron_config.batch_size
        self.n_positions = neuron_config.sequence_length
        self.vocab_size = config.vocab_size
        self.speculation_length = neuron_config.speculation_length
        self.padding_side = neuron_config.padding_side
        self.max_length = neuron_config.sequence_length
        self.sequence_parallel_enabled = neuron_config.sequence_parallel_enabled
        self.sequence_dimension = 1 if self.sequence_parallel_enabled else None
        self.rank_util = SPMDRank(world_size=neuron_config.tp_degree)
        self.num_cores_per_group = neuron_config.num_cores_per_group
        if neuron_config.on_device_sampling:
            # Instantiate a multinomial Sampler (it can still be used for greedy by passing topk=1)
            self.sampler = Sampler(neuron_config, do_sample=True)
        self.kv_mgr = KVCacheManager(config, neuron_config, num_kv_head=config.num_key_value_heads)

    def initialize_process_group(self, seed: int = 0):
        if not torch.dist.is_initialized():
            torch.dist.init_process_group(backend="xla")
        else:
            logging.warning("torch.distributed was already initialized, skipping...")

        if not nxd.parallel_layers.parallel_state.model_parallel_is_initialized():
            nxd.parallel_layers.initialize_model_parallel(
                tensor_model_parallel_size=self.neuron_config.tp_degree,
                pipeline_model_parallel_size=self.neuron_config.pp_degree,
                expert_model_parallel_size=self.neuron_config.ep_degree,
            )
        else:
            logging.warning("NxD was already initialized, skipping...")

        # set seed
        set_random_seed(seed)

    def _is_context_encoding(self, input_ids: torch.Tensor):
        return input_ids.shape[-1] > 1 and input_ids.shape[-1] != self.speculation_length

    def _is_for_speculation(self, input_ids: torch.Tensor):
        return input_ids.shape[-1] == self.speculation_length

    def _create_context_attn_mask(self, attention_mask, **kwargs):
        # Block diagonal causal mask for chunked prefill
        if self.neuron_config.is_chunked_prefill:
            return self._create_chunked_prefill_attn_mask(**kwargs)

        # Lower triangle causal mask for classic attention
        mask = torch.full((self.n_positions, self.n_positions), True, device=attention_mask.device).tril(diagonal=0)
        mask = mask[None, None, :, :].expand(self.batch_size, 1, self.n_positions, self.n_positions)

        if self.padding_side == "right":
            return mask
        else:
            expanded_mask = (
                attention_mask[:, None, None, :]
                .expand(self.batch_size, 1, self.n_positions, self.n_positions)
                .to(torch.bool)
            )
            return torch.logical_and(mask, expanded_mask)

    def _create_chunked_prefill_attn_mask(
        self,
        query_lens: torch.Tensor,
        key_lens: torch.Tensor,
        max_query_len: int,
        max_key_len: int,
        **kwargs,
    ) -> torch.Tensor:
        return attn_utils.create_block_diagonal_attn_mask(query_lens, key_lens, max_query_len, max_key_len, **kwargs)

    def _create_spec_attn_mask(self, attention_mask):
        return (
            attention_mask[:, None, None, :]
            .expand(self.batch_size, 1, self.speculation_length, self.n_positions)
            .to(torch.bool)
        )

    def _create_simple_attn_mask(self, attention_mask):
        return attention_mask[:, None, None, :].expand(self.batch_size, 1, 1, self.n_positions).to(torch.bool)

    def create_attn_mask(self, attention_mask, is_for_context_encoding, is_for_speculation, **kwargs):
        if is_for_context_encoding:
            return self._create_context_attn_mask(attention_mask, **kwargs)
        elif is_for_speculation:
            return self._create_spec_attn_mask(attention_mask)
        else:
            return self._create_simple_attn_mask(attention_mask)

    def _slice_kv_cache(self, kv_cache, n_positions):
        past_key_values = []
        for idx in range(len(kv_cache)):
            k_cache = _slice_kv_cacheline(self.neuron_config.padding_side, n_positions, kv_cache[idx][0])
            v_cache = _slice_kv_cacheline(self.neuron_config.padding_side, n_positions, kv_cache[idx][1])
            past_key_values.append([k_cache, v_cache])
        return past_key_values

    def _is_reorder_needed(self, is_for_context_encoding, is_for_speculation):
        return not is_for_context_encoding and not is_for_speculation and self.neuron_config.continuous_batching

    def forward(
        self,
        input_ids,
        attention_mask,
        position_ids,
        seq_ids,
        sampling_params,
        scatter_index=None,
        # In llava context encoding model, input_embeds is precomputed
        inputs_embeds: Optional[torch.FloatTensor] = None,
        kv_cache: Optional[torch.Tensor] = None,
    ):
        is_for_context_encoding = self._is_context_encoding(input_ids)
        is_for_speculation = self._is_for_speculation(input_ids)

        cache_size = (
            get_cache_size(self.n_positions, self.num_cores_per_group, is_for_context_encoding)
            if self.neuron_config.flash_decoding_enabled
            else self.n_positions
        )

        # It is either for context encoding or for token generation
        if is_for_context_encoding:
            past_key_values = None
        else:
            if kv_cache is None:
                past_key_values = self.kv_mgr.get_cache(cache_size)
            else:
                past_key_values = self._slice_kv_cache(kv_cache, cache_size)

        # Prepare attention mask(s)
        attention_mask = self.create_attn_mask(
            attention_mask,
            is_for_context_encoding,
            is_for_speculation,
        )
        active_mask = None
        if is_for_speculation:
            active_mask = torch.full(
                (self.speculation_length, self.speculation_length),
                True,
                device=attention_mask.device,
            ).tril(diagonal=0)
            active_mask = active_mask[None, None, :, :].expand(
                self.batch_size, 1, self.speculation_length, self.speculation_length
            )

        # FD masks
        active_mask_2d = None
        if self.neuron_config.flash_decoding_enabled and not is_for_context_encoding:
            rank_id = self.rank_util.get_rank()
            active_mask_2d, attention_mask_2d = mask_util(
                pos_ids=position_ids,
                rank_id=rank_id,
                num_cores_per_group=self.num_cores_per_group,
                cache_size=cache_size,
            )
            active_mask = turn_2d_mask_to_4d(active_mask_2d, n_positions=1, batch_size=self.batch_size)
            attention_mask = turn_2d_mask_to_4d(attention_mask_2d, n_positions=cache_size, batch_size=self.batch_size)

        hidden_states, past_key_values = self.get_model_output(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            active_mask=active_mask,
            inputs_embeds=inputs_embeds,
        )

        updated_kv_cache = self.kv_mgr.update_cache(
            is_for_context_encoding=is_for_context_encoding,
            seq_ids=seq_ids,
            position_ids=position_ids,
            new_key_values=past_key_values,
            seq_len=cache_size,
            scatter_index=scatter_index,
            active_mask=active_mask_2d,
            kvcache_buffer=kv_cache,
        )

        batch_size, num_tokens, hidden_size = hidden_states.shape
        if self.padding_side == "left":
            index = torch.tensor([num_tokens - 1], device=hidden_states.device)
            index = index.unsqueeze(1).expand(batch_size, 1, hidden_size)
            hidden_states = torch.gather(hidden_states, dim=1, index=index)
        else:
            if not (position_ids.shape[-1] == self.speculation_length or position_ids.shape[-1] == 1):
                # context encoding
                index = torch.max(position_ids, dim=1, keepdim=True).indices
                index = index.unsqueeze(1).expand(batch_size, 1, hidden_size)
                hidden_states = torch.gather(hidden_states, dim=1, index=index)

        logits = self.lm_head(hidden_states)
        logits = logits.float()

        if hasattr(self.lm_head, "pad_size"):
            if self.lm_head.gather_output:
                rank_id = torch.tensor(0, device=logits.device, dtype=torch.int32)
                world_size = 1
            else:
                rank_id = self.rank_util.get_rank()
                world_size = torch.distributed.get_world_size(group=self.lm_head.tensor_parallel_group)
            logits = mask_padded_logits(logits, rank_id, world_size, pad_size=self.lm_head.pad_size)

        res = logits
        if self.neuron_config.on_device_sampling:
            # perform sampling on Neuron to get tokens
            # FIXME, logits[:, -1, :] is not correct for speculation model, this is a tempory fix.
            if is_for_speculation:
                res = nxd_argmax(tensor=logits, dim=2, gather_dim=2, keepdim=False)
            else:
                res = self.sampler(logits[:, -1, :], sampling_params)

        outputs = [res]
        if self.neuron_config.output_logits:
            logits = _gather_along_dim(
                logits,
                partition_dim=2,
            )
            outputs += [logits]
        outputs += updated_kv_cache

        return outputs

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def get_model_output(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        active_mask: Optional[List[torch.FloatTensor]] = None,
        # In llava context encoding model, input_embeds is precomputed
        inputs_embeds: Optional[torch.FloatTensor] = None,
    ):
        batch_size, seq_length = input_ids.shape[:2]

        past_key_values_length = 0
        if past_key_values is not None:
            past_key_values_length = past_key_values[0][0].shape[2]

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if position_ids is None:
            device = input_ids.device if input_ids is not None else inputs_embeds.device  # noqa
            position_ids = torch.arange(
                past_key_values_length,
                seq_length + past_key_values_length,
                dtype=torch.long,
                device=device,
            )
            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
        else:
            position_ids = position_ids.view(-1, seq_length).long()

        # NeuronLlamaModel class manages the KV cache. So the attention_mask will be generated and passed
        # through to LlamaModel. We override the HF's code that generates attention mask because HF does
        # not support left aligned RHS padding. This enables Neuron to achieve higher performance and
        # extensibility.
        #
        # 4d mask is passed through the layers
        # attention_mask = _prepare_4d_causal_attention_mask(
        #     attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
        # )

        # embed positions
        if self.sequence_parallel_enabled:
            # TODO: Replace this with rankid + scatter call once supported
            hidden_states = _reduce_scatter_along_dim(
                inputs_embeds,
                self.sequence_dimension,
                xm.REDUCE_MAX,
            )
        else:
            hidden_states = inputs_embeds

        # decoder layers
        next_decoder_cache = ()
        cos_cache = None
        sin_cache = None
        for idx, decoder_layer in enumerate(self.layers):
            past_key_value = past_key_values[idx] if past_key_values is not None else None

            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_value=past_key_value,
                active_mask=active_mask,
                cos_cache=cos_cache,
                sin_cache=sin_cache,
            )

            hidden_states = layer_outputs[0]
            next_decoder_cache += (layer_outputs[1],)
            cos_cache, sin_cache = layer_outputs[2:]

        hidden_states = self.norm(hidden_states)

        if self.sequence_parallel_enabled:
            hidden_states = gather_from_sequence_parallel_region(
                hidden_states,
                self.sequence_dimension,
            )

        return (hidden_states, next_decoder_cache)


class NxDModelForCausalLM(NxDGenerationMixin, NxDPreTrainedModel, NeuronModelForCausalLM):
    _model_cls = None

    def __init__(
        self,
        config: PretrainedConfig,
        neuron_config: NxDNeuronConfig,
        traced_model: torch.jit.ScriptModule,
        context_encoding_model: NxDDecoderWrapper,
        token_generation_model: NxDDecoderWrapper = None,
        speculation_model: NxDDecoderWrapper = None,
    ):
        self.context_encoding_model = context_encoding_model
        self.token_generation_model = token_generation_model
        self.speculation_model = speculation_model
        # Model wrappers are used by the parent class to assign weights to the model.
        model_wrappers = [self.context_encoding_model]
        if self.token_generation_model is not None:
            model_wrappers.append(self.token_generation_model)
        if self.speculation_model is not None:
            model_wrappers.append(self.speculation_model)
        super().__init__(
            config=config, neuron_config=neuron_config, traced_model=traced_model, model_wrappers=model_wrappers
        )

        self.text_config = self.config.get_text_config()
        self.vocab_size = self.text_config.vocab_size
        self.padding_side = self.neuron_config.padding_side
        self.kv_cache_populated = False

        # async related
        self.async_mode = self.neuron_config.async_mode
        self.next_cpu_inputs = None
        self.prior_outputs = None
        self.unequal_batching = self.neuron_config.ctx_batch_size != self.neuron_config.tkg_batch_size
        if self.async_mode:
            os.environ["NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS"] = "2"

        self.sampler = None

    @staticmethod
    def create_context_encoding_wrapper(model_cls, config, neuron_config, **model_init_kwargs):
        new_neuron_config = copy.deepcopy(neuron_config)
        new_neuron_config.batch_size = neuron_config.ctx_batch_size
        new_neuron_config.n_active_tokens = neuron_config.max_context_length

        if new_neuron_config.enable_bucketing:
            buckets = generate_buckets(128, new_neuron_config.max_context_length)
        else:
            buckets = generate_buckets(
                new_neuron_config.max_context_length,
                new_neuron_config.max_context_length,
            )

        return NxDDecoderWrapper(
            config=config,
            neuron_config=new_neuron_config,
            buckets=buckets,
            bucket_n_active_tokens=True,
            model_cls=model_cls,
            tag=CONTEXT_ENCODING_MODEL_TAG,
            model_init_kwargs=model_init_kwargs,
        )

    @staticmethod
    def create_token_generation_wrapper(
        model_cls, config, neuron_config, enable_wlt_optimization: bool = True, **model_init_kwargs
    ):
        new_neuron_config = copy.deepcopy(neuron_config)
        new_neuron_config.batch_size = neuron_config.tkg_batch_size
        new_neuron_config.n_active_tokens = 1
        new_neuron_config.sequence_parallel_enabled = False

        if new_neuron_config.enable_bucketing:
            buckets = generate_buckets(128, neuron_config.sequence_length)
        else:
            buckets = generate_buckets(neuron_config.sequence_length, neuron_config.sequence_length)

        # shouldn't be used in token gen models
        new_neuron_config.sequence_parallel_enabled = False

        return NxDDecoderWrapper(
            config=config,
            neuron_config=new_neuron_config,
            buckets=buckets,
            bucket_n_active_tokens=False,
            model_cls=model_cls,
            tag=TOKEN_GENERATION_MODEL_TAG,
            priority_model_idx=0 if enable_wlt_optimization else None,  # to turn on weight layout optimization
            model_init_kwargs=model_init_kwargs,
        )

    @staticmethod
    def create_speculation_wrapper(model_cls, config, neuron_config, **model_init_kwargs):
        new_neuron_config = copy.deepcopy(neuron_config)
        new_neuron_config.batch_size = neuron_config.tkg_batch_size
        new_neuron_config.n_active_tokens = neuron_config.speculation_length

        new_neuron_config.sequence_parallel_enabled = False

        if new_neuron_config.enable_bucketing:
            buckets = generate_buckets(128, neuron_config.sequence_length)
        else:
            buckets = generate_buckets(neuron_config.sequence_length, neuron_config.sequence_length)

        return NxDDecoderWrapper(
            config=config,
            neuron_config=new_neuron_config,
            buckets=buckets,
            bucket_n_active_tokens=False,
            model_cls=model_cls,
            tag=SPECULATION_MODEL_TAG,
            priority_model_idx=0,  # to turn on weight layout optimization
            model_init_kwargs=model_init_kwargs,
        )

    @staticmethod
    def create_model_wrappers(model_cls, config, neuron_config, **model_init_kwargs):
        context_encoding_model = NxDModelForCausalLM.create_context_encoding_wrapper(
            model_cls,
            config,
            neuron_config,
            **model_init_kwargs,
        )
        token_generation_model = NxDModelForCausalLM.create_token_generation_wrapper(
            model_cls,
            config,
            neuron_config,
            **model_init_kwargs,
        )
        speculation_model = (
            NxDModelForCausalLM.create_speculation_wrapper(
                model_cls,
                config,
                neuron_config,
                **model_init_kwargs,
            )
            if neuron_config.speculation_length > 0
            else None
        )
        return context_encoding_model, token_generation_model, speculation_model

    def forward(
        self,
        input_ids: torch.LongTensor,
        position_ids: Optional[torch.LongTensor],
        seq_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        sampling_params: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        if self.async_mode:
            # derive future cpu inputs from current cpu inputs
            if position_ids.shape[1] == input_ids.shape[1]:
                next_position_ids = torch.amax(position_ids, 1, keepdim=True)
            else:
                next_position_ids = position_ids

            next_position_ids = next_position_ids + 1
            next_attention_mask = self._infer_attention_mask(next_position_ids)
            self.next_cpu_inputs = {
                "attention_mask": next_attention_mask,
                "position_ids": next_position_ids,
            }

        # infer attention_mask from position_ids if not provided
        if attention_mask is None:
            attention_mask = self._infer_attention_mask(position_ids)

        if seq_ids is None:
            seq_ids = torch.arange(input_ids.shape[0])

        if sampling_params is None:
            if self.neuron_config.on_device_sampling:
                raise ValueError("The sampling params tensor is required for on-device sampling.")
            # Just pass a dummy tensor to the model, it will be ignored
            sampling_params = prepare_sampling_params(seq_ids.shape[0])
        elif self.neuron_config.on_device_sampling:
            validate_sampling_params(sampling_params, self.neuron_config.max_topk)

        output_attentions, output_hidden_states, return_dict = self._setup_func_config(
            output_attentions, output_hidden_states, return_dict
        )

        logits_or_next_tokens = self._get_model_outputs(
            input_ids,
            attention_mask,
            position_ids,
            seq_ids,
            sampling_params,
        )

        logging.debug("---output---")
        logging.debug(
            f"{'tokens' if self.neuron_config.on_device_sampling else 'logits'} = %s, ",
            logits_or_next_tokens,
        )

        return self._construct_output(logits_or_next_tokens)

    def _setup_func_config(self, output_attentions, output_hidden_states, return_dict):
        output_attentions = output_attentions if output_attentions is not None else self.text_config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.text_config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else getattr(self.config, "use_return_dict", None)
        return output_attentions, output_hidden_states, return_dict

    def _infer_attention_mask(self, position_ids):
        assert position_ids is not None, "need to call forward with position_ids if attention_mask is not provided"
        batch_size, seq_len = position_ids.shape
        if position_ids.shape[-1] == 1:
            seq_len = self.neuron_config.sequence_length
            position_ids_to_compare = position_ids.expand(batch_size, seq_len) - 1
        else:
            seq_len = position_ids.shape[-1]
            position_ids_to_compare = position_ids
        mask = torch.arange(seq_len).view(1, -1).expand(batch_size, seq_len)
        attention_mask = (position_ids_to_compare >= mask).to(dtype=position_ids.dtype)
        return attention_mask

    def _get_async_output(
        self,
        ranked_async_tensor,
    ):
        outputs = [[async_tensor[0].cpu()] for async_tensor in ranked_async_tensor]
        return outputs[0][0]

    def _get_model_outputs(
        self,
        input_ids,
        attention_mask,
        position_ids,
        seq_ids,
        sampling_params,
    ):
        # casting inputs to int32
        input_ids = input_ids.to(torch.int32)
        attention_mask = attention_mask.to(torch.int32)
        position_ids = position_ids.to(torch.int32)
        seq_ids = seq_ids.to(torch.int32)

        if input_ids.shape[-1] > 1 and not position_ids.min().item():
            outputs = self.context_encoding_model(
                input_ids,
                attention_mask,
                position_ids,
                seq_ids,
                sampling_params,
            )

            self.kv_cache_populated = True
            if self.async_mode:
                if not self.unequal_batching:
                    # for now only cte + tkg flow is supported with async (this will be enforced at config level)
                    next_outputs = self.token_generation_model(
                        outputs,
                        self.next_cpu_inputs["attention_mask"],
                        self.next_cpu_inputs["position_ids"],
                        seq_ids,
                        sampling_params,
                    )
                    outputs = self._get_async_output(outputs)  # block on cte call
                    self.prior_outputs = next_outputs
                else:
                    if isinstance(
                        outputs, list
                    ):  # in case the outputs weren't passed through `torch.cat` in model_wrapper.py
                        outputs = self._get_async_output(outputs)  # block on cte call

                    self.prior_outputs = None

        elif input_ids.shape[-1] == self.neuron_config.speculation_length:
            outputs = self.speculation_model(
                input_ids,
                attention_mask,
                position_ids,
                seq_ids,
                sampling_params,
            )
        else:
            if (
                self.next_cpu_inputs is not None and self.prior_outputs is not None
            ):  # this is never not None and not in async mode
                _input_ids = self.prior_outputs
                _attention_mask = self.next_cpu_inputs["attention_mask"]
                _position_ids = self.next_cpu_inputs["position_ids"]
            else:
                _input_ids = input_ids
                _attention_mask = attention_mask
                _position_ids = position_ids

            next_outputs = self.token_generation_model(
                _input_ids,
                _attention_mask,
                _position_ids,
                seq_ids,
                sampling_params,
            )
            if self.async_mode:
                if self.prior_outputs is None:  # this means that next_outputs is processing token to be returned
                    self.prior_outputs = next_outputs
                    next_outputs = self.token_generation_model(  # submit future token request
                        next_outputs,
                        self.next_cpu_inputs["attention_mask"],
                        self.next_cpu_inputs["position_ids"],
                        seq_ids,
                        sampling_params,
                    )
                outputs = self.prior_outputs
                if isinstance(outputs, list):
                    outputs = self._get_async_output(
                        self.prior_outputs
                    )  # block on prior (sometimes current) token gen request

                self.prior_outputs = next_outputs
            else:
                outputs = next_outputs

        return outputs

    def _construct_output(self, logits_or_next_tokens):
        next_tokens = logits_or_next_tokens

        OutputParams = CausalLMOutputWithPast(
            logits=None if self.neuron_config.on_device_sampling else logits_or_next_tokens,
            hidden_states=logits_or_next_tokens,
            attentions=None,
        )

        OutputParams.tokens = next_tokens

        return OutputParams

    def reset(self):
        # We need to reset the KV cache flag for a new batch of inference.
        # When the flag is reset, the subsequent run will invoke the
        # context encoding model.
        self.kv_cache_populated = False

    def get_required_kwargs(self) -> List[str]:
        """The list of required kwargs to the model's forward"""
        return []

    @classmethod
    def get_compiler_args(cls, neuron_config: NxDNeuronConfig) -> str:
        tensorizer_options = (
            "--enable-ccop-compute-overlap "
            f"--cc-pipeline-tiling-factor={neuron_config.cc_pipeline_tiling_factor} "
            "--vectorize-dge-dma "
            "--vectorize-strided-dma "
        )

        compiler_args = (
            "--auto-cast=none --model-type=transformer "
            f"--tensorizer-options='{tensorizer_options}'"
            " -O2 "
            f" --internal-num-neuroncores-per-sengine={neuron_config.logical_nc_config}"
        )

        if neuron_config.target:
            compiler_args += f" --target {neuron_config.target}"

        logging.info(f"neuronx-cc compiler_args are: {compiler_args}")
        return compiler_args

    # NeuronModelForCausalLM methods
    @classmethod
    def _from_pretrained(
        cls,
        model_id: Union[str, "Path"],
        config: "PretrainedConfig",
        revision: Optional[str] = None,
        token: Optional[Union[bool, str]] = None,
        cache_dir: Optional[str] = None,
        force_download: Optional[bool] = False,
        subfolder: Optional[str] = "",
        local_files_only: Optional[bool] = False,
        trust_remote_code: Optional[bool] = False,
        **kwargs,
    ) -> "NeuronModelForCausalLM":
        if len(kwargs) > 0:
            logger.warning("Ignoring the following kwargs as they are not supported by neuron: %s", kwargs.keys())
        neuron_config = cls.get_neuron_config_cls().from_pretrained(model_id)
        context_encoding_model, token_generation_model, speculation_model = cls.create_model_wrappers(
            model_cls=cls._model_cls,
            config=config,
            neuron_config=neuron_config,
        )
        if not os.path.exists(model_id):
            # The model_id is a model hub id: download the model from the hub.
            with TemporaryDirectory() as tmpdir:
                snapshot_download(
                    repo_id=model_id,
                    revision=revision,
                    cache_dir=cache_dir,
                    local_dir=tmpdir,
                    force_download=force_download,
                    local_files_only=local_files_only,
                    token=token,
                    allow_patterns=[cls.COMPILED_MODEL_FILE_NAME],
                )
                traced_model = torch.jit.load(os.path.join(tmpdir, cls.COMPILED_MODEL_FILE_NAME))
        else:
            traced_model = torch.jit.load(os.path.join(model_id, cls.COMPILED_MODEL_FILE_NAME))
        model = cls(
            config=config,
            neuron_config=neuron_config,
            traced_model=traced_model,
            context_encoding_model=context_encoding_model,
            token_generation_model=token_generation_model,
            speculation_model=speculation_model,
        )
        model.load_weights(
            model_id,
            cache_dir=cache_dir,
            force_download=force_download,
            local_files_only=local_files_only,
            token=token,
        )
        return model

    @classmethod
    def export(
        cls,
        model_id: str,
        config: "PretrainedConfig",
        neuron_config: "NxDNeuronConfig",
        token: Optional[Union[bool, str]] = None,
        revision: Optional[str] = None,
        cache_dir: Optional[str] = None,
        force_download: Optional[bool] = False,
        subfolder: Optional[str] = "",
        local_files_only: Optional[bool] = False,
        trust_remote_code: Optional[bool] = False,
        load_weights: bool = True,
        **kwargs,
    ) -> "NeuronModelForCausalLM":
        if len(kwargs) > 0:
            logger.warning("Ignoring the following kwargs as they are not supported by neuron: %s", kwargs.keys())
        config = AutoConfig.from_pretrained(
            model_id,
            token=token,
            revision=revision,
            cache_dir=cache_dir,
            force_download=force_download,
            trust_remote_code=trust_remote_code,
        )
        # Override torch_dtype in config as it is used by the neuronx_distributed code to cast weights to the correct type
        config.torch_dtype = neuron_config.torch_dtype
        context_encoding_model, token_generation_model, speculation_model = cls.create_model_wrappers(
            model_cls=cls._model_cls,
            config=config,
            neuron_config=neuron_config,
        )
        model_wrappers = []
        for wrapper in context_encoding_model, token_generation_model, speculation_model:
            if wrapper is not None:
                model_wrappers.append(wrapper)

        # The model NEFF files will be cached locally, but if the model_id corresponds
        # to a hub model, we also create a cache entry for it.
        cache_entry = (
            None
            if os.path.exists(model_id)
            else SingleModelCacheEntry(model_id, task="text-generation", config=config, neuron_config=neuron_config)
        )
        with hub_neuronx_cache(entry=cache_entry):
            traced_model = NxDPreTrainedModel.compile(
                neuron_config=neuron_config,
                model_wrappers=model_wrappers,
                compiler_args=cls.get_compiler_args(neuron_config),
            )
        model = cls(
            config=config,
            neuron_config=neuron_config,
            traced_model=traced_model,
            context_encoding_model=context_encoding_model,
            token_generation_model=token_generation_model,
            speculation_model=speculation_model,
        )
        if load_weights:
            model.load_weights(
                model_id,
                cache_dir=cache_dir,
                force_download=force_download,
                local_files_only=local_files_only,
                token=token,
            )
        return model

    def _save_pretrained(self, save_directory: Union[str, Path]):
        model_name_or_path = getattr(self.config, "_name_or_path")
        # If the model was exported from a local path, we need to save the checkpoint (not that we also shard it)
        weight_path = model_name_or_path if os.path.isdir(model_name_or_path) else None
        self.save(save_directory, weight_path=weight_path)

    def push_to_hub(
        self,
        save_directory: str,
        repository_id: str,
        private: Optional[bool] = None,
        revision: Optional[str] = None,
        token: Union[bool, str] = True,
        endpoint: Optional[str] = None,
    ) -> str:
        api = HfApi(endpoint=endpoint)

        api.create_repo(
            token=token,
            repo_id=repository_id,
            exist_ok=True,
            private=private,
        )
        ignore_patterns = []
        checkpoint_id = self.neuron_config.checkpoint_id
        if checkpoint_id is not None:
            # Avoid uploading checkpoints when the original model is available on the hub
            ignore_patterns = [self.CHECKPOINT_DIR + "/*"]
        api.upload_folder(
            repo_id=repository_id,
            folder_path=save_directory,
            token=token,
            revision=revision,
            ignore_patterns=ignore_patterns,
        )