Instructions to use ml-resources/MobileLLM-Pro with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use ml-resources/MobileLLM-Pro with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="ml-resources/MobileLLM-Pro", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("ml-resources/MobileLLM-Pro", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("ml-resources/MobileLLM-Pro", trust_remote_code=True, device_map="auto")
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=256)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use ml-resources/MobileLLM-Pro with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "ml-resources/MobileLLM-Pro"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "ml-resources/MobileLLM-Pro",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/ml-resources/MobileLLM-Pro

SGLang

How to use ml-resources/MobileLLM-Pro with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "ml-resources/MobileLLM-Pro" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "ml-resources/MobileLLM-Pro",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "ml-resources/MobileLLM-Pro" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "ml-resources/MobileLLM-Pro",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use ml-resources/MobileLLM-Pro with Docker Model Runner:
```
docker model run hf.co/ml-resources/MobileLLM-Pro
```

MobileLLM-Pro / modeling_mobilellm_p1.py

mrfakename

Mirror facebook/MobileLLM-Pro

9d8dbfd verified 18 days ago

Raw History Blame Contribute Delete

17.3 kB

import math
from dataclasses import dataclass
from typing import Callable, Optional, Union
import copy

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.integrations import use_kernel_forward_from_hub
from transformers.masking_utils import (
    create_causal_mask,
    create_sliding_window_causal_mask,
)
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
    ModelOutput,
)
from transformers.modeling_rope_utils import dynamic_rope_update, ROPE_INIT_FUNCTIONS
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel

from transformers.models.llama4.configuration_llama4 import (
    Llama4Config,
    Llama4TextConfig,
    Llama4VisionConfig,
)
from transformers.models.llama4.modeling_llama4 import (
    apply_rotary_emb,
    eager_attention_forward,
    Llama4PreTrainedModel,
    Llama4TextDecoderLayer,
    Llama4TextL2Norm,
    Llama4TextMLP,
    Llama4TextMoe,
    Llama4TextRMSNorm,
    Llama4TextRotaryEmbedding,
    Llama4TextAttention,
    Llama4TextDecoderLayer,
    Llama4ForCausalLM
)
from transformers.processing_utils import Unpack
from transformers.utils import (
    auto_docstring,
    can_return_tuple,
    logging,
    TransformersKwargs,
)
from transformers.utils.deprecation import deprecate_kwarg
from transformers.utils.generic import check_model_inputs

from .configuration_mobilellm_p1 import MobileLLMP1TextConfig

logger = logging.get_logger(__name__)

class MobileLLMP1TextAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: MobileLLMP1TextConfig, layer_idx):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.is_sliding = config.layer_types[layer_idx] == "sliding_attention"
        self.head_dim = getattr(
            config, "head_dim", config.hidden_size // config.num_attention_heads
        )
        self.num_attention_heads = config.num_attention_heads
        self.num_key_value_groups = (
            config.num_attention_heads // config.num_key_value_heads
        )
        self.num_key_value_heads = config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attn_scale = config.attn_scale
        self.floor_scale = config.floor_scale
        self.attn_temperature_tuning = config.attn_temperature_tuning
        self.attention_dropout = config.attention_dropout
        self.is_causal = True
        self.use_rope = config.no_rope_layers[layer_idx]
        self.sliding_window = config.sliding_window if self.is_sliding else None
        self.q_proj = nn.Linear(
            config.hidden_size,
            config.num_attention_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.k_proj = nn.Linear(
            config.hidden_size,
            config.num_key_value_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.v_proj = nn.Linear(
            config.hidden_size,
            config.num_key_value_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.o_proj = nn.Linear(
            config.num_attention_heads * self.head_dim,
            config.hidden_size,
            bias=config.attention_bias,
        )
        if self.config.use_qk_norm and self.use_rope:
            self.qk_norm = Llama4TextL2Norm(config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_proj(hidden_states).view(hidden_shape)
        key_states = self.k_proj(hidden_states).view(*input_shape, -1, self.head_dim)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        if self.use_rope:  # the 16E model skips rope for long context on certain layers
            query_states, key_states = apply_rotary_emb(
                query_states, key_states, position_embeddings.to(query_states.device)
            )

        if hasattr(self, "qk_norm"):  # the 128E model does not use qk_norm
            query_states = self.qk_norm(query_states)
            key_states = self.qk_norm(key_states)

        # Use temperature tuning from https://huggingface.co/papers/2501.19399) to NoROPE layers
        if self.attn_temperature_tuning and not self.use_rope:
            attn_scales = (
                torch.log(
                    torch.floor((cache_position.float() + 1.0) / self.floor_scale) + 1.0
                )
                * self.attn_scale
                + 1.0
            )
            attn_scales = attn_scales.view((1, input_shape[-1], 1, 1)).expand(
                (*input_shape, 1, 1)
            )  # batch size > 1
            query_states = (query_states * attn_scales).to(query_states.dtype)

        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)

        if past_key_values is not None:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {"cache_position": cache_position}
            key_states, value_states = past_key_values.update(
                key_states, value_states, self.layer_idx, cache_kwargs
            )

        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[
                self.config._attn_implementation
            ]

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            sliding_window=self.sliding_window,
            **kwargs,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class MobileLLMP1TextDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.layer_idx = layer_idx
        self.attention_type = config.layer_types[layer_idx]
        self.self_attn = MobileLLMP1TextAttention(config, layer_idx)
        self.is_moe_layer = layer_idx in config.moe_layers
        if self.is_moe_layer:  # the 128E model interleaves dense / sparse
            self.feed_forward = Llama4TextMoe(config)
        else:
            self.feed_forward = Llama4TextMLP(
                config, intermediate_size=config.intermediate_size_mlp
            )

        self.input_layernorm = Llama4TextRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_attention_layernorm = Llama4TextRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[
        torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]
    ]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        attention_states, _ = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            **kwargs,
        )
        hidden_states = residual + attention_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.feed_forward(hidden_states)
        if self.is_moe_layer:
            hidden_states, _ = hidden_states
        hidden_states = residual + hidden_states.view(residual.shape)
        return hidden_states

class MobileLLMP1TextModel(Llama4PreTrainedModel):
    _no_split_modules = ["MobileLLMP1TextDecoderLayer"]
    base_model_prefix = "model"
    config: MobileLLMP1TextConfig
    _can_record_outputs = {
        "attentions": MobileLLMP1TextAttention,
        "hidden_states": MobileLLMP1TextDecoderLayer,
        "router_logits": Llama4TextMoe,
    }

    def __init__(self, config: MobileLLMP1TextConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [
                MobileLLMP1TextDecoderLayer(config, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
        self.norm = Llama4TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = Llama4TextRotaryEmbedding(config=config)
        self.gradient_checkpointing = False

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Union[tuple, BaseModelOutputWithPast]:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
                "You must specify exactly one of input_ids or inputs_embeds"
            )

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

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        if cache_position is None:
            past_seen_tokens = (
                past_key_values.get_seq_length() if past_key_values is not None else 0
            )
            cache_position = torch.arange(
                past_seen_tokens,
                past_seen_tokens + inputs_embeds.shape[1],
                device=inputs_embeds.device,
            )

        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        # It may already have been prepared by e.g. `generate`
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            # Prepare mask arguments
            mask_kwargs = {
                "config": self.config,
                "input_embeds": inputs_embeds,
                "attention_mask": attention_mask,
                "cache_position": cache_position,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }
            sliding_mask_kwargs = mask_kwargs.copy()
            del sliding_mask_kwargs['position_ids']

            # Create the masks
            causal_mask_mapping = {
                "full_attention": create_causal_mask(**mask_kwargs),
                "sliding_attention": create_sliding_window_causal_mask(
                    **sliding_mask_kwargs
                ),
            }

        hidden_states = inputs_embeds

        # create position embeddings to be shared across the decoder layers
        freq_cis = self.rotary_emb(hidden_states, position_ids)

        # found = False
        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                cache_position=cache_position,
                position_embeddings=freq_cis,
                **kwargs,
            )
        hidden_states = self.norm(hidden_states)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
        )


class MobileLLMP1ForCausalLM(Llama4PreTrainedModel, GenerationMixin):
    _no_split_modules = ["MobileLLMP1TextDecoderLayer"]
    base_model_prefix = "language_model"
    _tied_weights_keys = ["lm_head.weight"]
    _tp_plan = {"lm_head": "colwise_rep"}
    config: MobileLLMP1TextConfig

    def __init__(self, config: MobileLLMP1TextConfig):
        super().__init__(config)
        self.model = MobileLLMP1TextModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Union[tuple, CausalLMOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
        Example:
        ```python
        >>> from transformers import AutoTokenizer, Llama4ForCausalLM
        >>> model = Llama4ForCausalLM.from_pretrained("meta-llama4/Llama4-2-7b-hf")
        >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama4/Llama4-2-7b-hf")
        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")
        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            cache_position=cache_position,
            **kwargs,
        )

        hidden_states = outputs[0]
        # Only compute necessary logixts, and do not upcast them to float if we are not computing the loss
        slice_indices = (
            slice(-logits_to_keep, None)
            if isinstance(logits_to_keep, int)
            else logits_to_keep
        )
        logits = self.lm_head(hidden_states[:, slice_indices, :])
        loss = None
        if labels is not None:
            loss = self.loss_function(
                logits=logits,
                labels=labels,
                vocab_size=self.config.vocab_size,
                **kwargs,
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

__all__ = [
    "MobileLLMP1ForCausalLM",
    "MobileLLMP1TextModel",
    "MobileLLMP1TextDecoderLayer",
    "MobileLLMP1TextAttention",
]

	import math
	from dataclasses import dataclass
	from typing import Callable, Optional, Union
	import copy

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	from transformers.activations import ACT2FN
	from transformers.cache_utils import Cache, DynamicCache
	from transformers.generation import GenerationMixin
	from transformers.integrations import use_kernel_forward_from_hub
	from transformers.masking_utils import (
	create_causal_mask,
	create_sliding_window_causal_mask,
	)
	from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
	from transformers.modeling_layers import GradientCheckpointingLayer
	from transformers.modeling_outputs import (
	BaseModelOutput,
	BaseModelOutputWithPast,
	CausalLMOutputWithPast,
	ModelOutput,
	)
	from transformers.modeling_rope_utils import dynamic_rope_update, ROPE_INIT_FUNCTIONS
	from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel

	from transformers.models.llama4.configuration_llama4 import (
	Llama4Config,
	Llama4TextConfig,
	Llama4VisionConfig,
	)
	from transformers.models.llama4.modeling_llama4 import (
	apply_rotary_emb,
	eager_attention_forward,
	Llama4PreTrainedModel,
	Llama4TextDecoderLayer,
	Llama4TextL2Norm,
	Llama4TextMLP,
	Llama4TextMoe,
	Llama4TextRMSNorm,
	Llama4TextRotaryEmbedding,
	Llama4TextAttention,
	Llama4TextDecoderLayer,
	Llama4ForCausalLM
	)
	from transformers.processing_utils import Unpack
	from transformers.utils import (
	auto_docstring,
	can_return_tuple,
	logging,
	TransformersKwargs,
	)
	from transformers.utils.deprecation import deprecate_kwarg
	from transformers.utils.generic import check_model_inputs

	from .configuration_mobilellm_p1 import MobileLLMP1TextConfig

	logger = logging.get_logger(__name__)

	class MobileLLMP1TextAttention(nn.Module):
	"""Multi-headed attention from 'Attention Is All You Need' paper"""

	def __init__(self, config: MobileLLMP1TextConfig, layer_idx):
	super().__init__()
	self.config = config
	self.layer_idx = layer_idx
	self.is_sliding = config.layer_types[layer_idx] == "sliding_attention"
	self.head_dim = getattr(
	config, "head_dim", config.hidden_size // config.num_attention_heads
	)
	self.num_attention_heads = config.num_attention_heads
	self.num_key_value_groups = (
	config.num_attention_heads // config.num_key_value_heads
	)
	self.num_key_value_heads = config.num_key_value_heads
	self.scaling = self.head_dim**-0.5
	self.attn_scale = config.attn_scale
	self.floor_scale = config.floor_scale
	self.attn_temperature_tuning = config.attn_temperature_tuning
	self.attention_dropout = config.attention_dropout
	self.is_causal = True
	self.use_rope = config.no_rope_layers[layer_idx]
	self.sliding_window = config.sliding_window if self.is_sliding else None
	self.q_proj = nn.Linear(
	config.hidden_size,
	config.num_attention_heads * self.head_dim,
	bias=config.attention_bias,
	)
	self.k_proj = nn.Linear(
	config.hidden_size,
	config.num_key_value_heads * self.head_dim,
	bias=config.attention_bias,
	)
	self.v_proj = nn.Linear(
	config.hidden_size,
	config.num_key_value_heads * self.head_dim,
	bias=config.attention_bias,
	)
	self.o_proj = nn.Linear(
	config.num_attention_heads * self.head_dim,
	config.hidden_size,
	bias=config.attention_bias,
	)
	if self.config.use_qk_norm and self.use_rope:
	self.qk_norm = Llama4TextL2Norm(config.rms_norm_eps)

	def forward(
	self,
	hidden_states: torch.Tensor,
	position_embeddings: tuple[torch.Tensor, torch.Tensor],
	attention_mask: Optional[torch.Tensor],
	past_key_values: Optional[Cache] = None,
	cache_position: Optional[torch.LongTensor] = None,
	**kwargs: Unpack[FlashAttentionKwargs],
	) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
	input_shape = hidden_states.shape[:-1]
	hidden_shape = (*input_shape, -1, self.head_dim)

	query_states = self.q_proj(hidden_states).view(hidden_shape)
	key_states = self.k_proj(hidden_states).view(*input_shape, -1, self.head_dim)
	value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

	if self.use_rope: # the 16E model skips rope for long context on certain layers
	query_states, key_states = apply_rotary_emb(
	query_states, key_states, position_embeddings.to(query_states.device)
	)

	if hasattr(self, "qk_norm"): # the 128E model does not use qk_norm
	query_states = self.qk_norm(query_states)
	key_states = self.qk_norm(key_states)

	# Use temperature tuning from https://huggingface.co/papers/2501.19399) to NoROPE layers
	if self.attn_temperature_tuning and not self.use_rope:
	attn_scales = (
	torch.log(
	torch.floor((cache_position.float() + 1.0) / self.floor_scale) + 1.0
	)
	* self.attn_scale
	+ 1.0
	)
	attn_scales = attn_scales.view((1, input_shape[-1], 1, 1)).expand(
	(*input_shape, 1, 1)
	) # batch size > 1
	query_states = (query_states * attn_scales).to(query_states.dtype)

	query_states = query_states.transpose(1, 2)
	key_states = key_states.transpose(1, 2)

	if past_key_values is not None:
	# sin and cos are specific to RoPE models; cache_position needed for the static cache
	cache_kwargs = {"cache_position": cache_position}
	key_states, value_states = past_key_values.update(
	key_states, value_states, self.layer_idx, cache_kwargs
	)

	attention_interface: Callable = eager_attention_forward
	if self.config._attn_implementation != "eager":
	attention_interface = ALL_ATTENTION_FUNCTIONS[
	self.config._attn_implementation
	]

	attn_output, attn_weights = attention_interface(
	self,
	query_states,
	key_states,
	value_states,
	attention_mask,
	dropout=0.0 if not self.training else self.attention_dropout,
	scaling=self.scaling,
	sliding_window=self.sliding_window,
	**kwargs,
	)

	attn_output = attn_output.reshape(*input_shape, -1).contiguous()
	attn_output = self.o_proj(attn_output)
	return attn_output, attn_weights


	class MobileLLMP1TextDecoderLayer(GradientCheckpointingLayer):
	def __init__(self, config, layer_idx):
	super().__init__()
	self.hidden_size = config.hidden_size
	self.layer_idx = layer_idx
	self.attention_type = config.layer_types[layer_idx]
	self.self_attn = MobileLLMP1TextAttention(config, layer_idx)
	self.is_moe_layer = layer_idx in config.moe_layers
	if self.is_moe_layer: # the 128E model interleaves dense / sparse
	self.feed_forward = Llama4TextMoe(config)
	else:
	self.feed_forward = Llama4TextMLP(
	config, intermediate_size=config.intermediate_size_mlp
	)

	self.input_layernorm = Llama4TextRMSNorm(
	config.hidden_size, eps=config.rms_norm_eps
	)
	self.post_attention_layernorm = Llama4TextRMSNorm(
	config.hidden_size, eps=config.rms_norm_eps
	)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[Cache] = None,
	use_cache: Optional[bool] = False,
	cache_position: Optional[torch.LongTensor] = None,
	position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
	**kwargs: Unpack[FlashAttentionKwargs],
	) -> tuple[
	torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]
	]:
	residual = hidden_states
	hidden_states = self.input_layernorm(hidden_states)

	# Self Attention
	attention_states, _ = self.self_attn(
	hidden_states=hidden_states,
	position_embeddings=position_embeddings,
	attention_mask=attention_mask,
	past_key_values=past_key_values,
	use_cache=use_cache,
	cache_position=cache_position,
	**kwargs,
	)
	hidden_states = residual + attention_states

	# Fully Connected
	residual = hidden_states
	hidden_states = self.post_attention_layernorm(hidden_states)
	hidden_states = self.feed_forward(hidden_states)
	if self.is_moe_layer:
	hidden_states, _ = hidden_states
	hidden_states = residual + hidden_states.view(residual.shape)
	return hidden_states

	class MobileLLMP1TextModel(Llama4PreTrainedModel):
	_no_split_modules = ["MobileLLMP1TextDecoderLayer"]
	base_model_prefix = "model"
	config: MobileLLMP1TextConfig
	_can_record_outputs = {
	"attentions": MobileLLMP1TextAttention,
	"hidden_states": MobileLLMP1TextDecoderLayer,
	"router_logits": Llama4TextMoe,
	}

	def __init__(self, config: MobileLLMP1TextConfig):
	super().__init__(config)
	self.padding_idx = config.pad_token_id
	self.vocab_size = config.vocab_size

	self.embed_tokens = nn.Embedding(
	config.vocab_size, config.hidden_size, self.padding_idx
	)
	self.layers = nn.ModuleList(
	[
	MobileLLMP1TextDecoderLayer(config, layer_idx)
	for layer_idx in range(config.num_hidden_layers)
	]
	)
	self.norm = Llama4TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.rotary_emb = Llama4TextRotaryEmbedding(config=config)
	self.gradient_checkpointing = False

	# Initialize weights and apply final processing
	self.post_init()

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[Cache] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	use_cache: Optional[bool] = None,
	cache_position: Optional[torch.LongTensor] = None,
	**kwargs: Unpack[TransformersKwargs],
	) -> Union[tuple, BaseModelOutputWithPast]:
	if (input_ids is None) ^ (inputs_embeds is not None):
	raise ValueError(
	"You must specify exactly one of input_ids or inputs_embeds"
	)

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

	if use_cache and past_key_values is None:
	past_key_values = DynamicCache(config=self.config)

	if cache_position is None:
	past_seen_tokens = (
	past_key_values.get_seq_length() if past_key_values is not None else 0
	)
	cache_position = torch.arange(
	past_seen_tokens,
	past_seen_tokens + inputs_embeds.shape[1],
	device=inputs_embeds.device,
	)

	if position_ids is None:
	position_ids = cache_position.unsqueeze(0)

	# It may already have been prepared by e.g. `generate`
	if not isinstance(causal_mask_mapping := attention_mask, dict):
	# Prepare mask arguments
	mask_kwargs = {
	"config": self.config,
	"input_embeds": inputs_embeds,
	"attention_mask": attention_mask,
	"cache_position": cache_position,
	"past_key_values": past_key_values,
	"position_ids": position_ids,
	}
	sliding_mask_kwargs = mask_kwargs.copy()
	del sliding_mask_kwargs['position_ids']

	# Create the masks
	causal_mask_mapping = {
	"full_attention": create_causal_mask(**mask_kwargs),
	"sliding_attention": create_sliding_window_causal_mask(
	**sliding_mask_kwargs
	),
	}

	hidden_states = inputs_embeds

	# create position embeddings to be shared across the decoder layers
	freq_cis = self.rotary_emb(hidden_states, position_ids)

	# found = False
	for decoder_layer in self.layers[: self.config.num_hidden_layers]:
	hidden_states = decoder_layer(
	hidden_states,
	attention_mask=causal_mask_mapping[decoder_layer.attention_type],
	position_ids=position_ids,
	past_key_values=past_key_values,
	use_cache=use_cache,
	cache_position=cache_position,
	position_embeddings=freq_cis,
	**kwargs,
	)
	hidden_states = self.norm(hidden_states)

	return BaseModelOutputWithPast(
	last_hidden_state=hidden_states,
	past_key_values=past_key_values if use_cache else None,
	)


	class MobileLLMP1ForCausalLM(Llama4PreTrainedModel, GenerationMixin):
	_no_split_modules = ["MobileLLMP1TextDecoderLayer"]
	base_model_prefix = "language_model"
	_tied_weights_keys = ["lm_head.weight"]
	_tp_plan = {"lm_head": "colwise_rep"}
	config: MobileLLMP1TextConfig

	def __init__(self, config: MobileLLMP1TextConfig):
	super().__init__(config)
	self.model = MobileLLMP1TextModel(config)
	self.vocab_size = config.vocab_size
	self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

	# Initialize weights and apply final processing
	self.post_init()

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	labels: Optional[torch.LongTensor] = None,
	use_cache: Optional[bool] = None,
	cache_position: Optional[torch.LongTensor] = None,
	logits_to_keep: Union[int, torch.Tensor] = 0,
	**kwargs: Unpack[TransformersKwargs],
	) -> Union[tuple, CausalLMOutputWithPast]:
	r"""
	labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, optional):
	Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
	config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
	(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

	Example:

	```python
	>>> from transformers import AutoTokenizer, Llama4ForCausalLM

	>>> model = Llama4ForCausalLM.from_pretrained("meta-llama4/Llama4-2-7b-hf")
	>>> tokenizer = AutoTokenizer.from_pretrained("meta-llama4/Llama4-2-7b-hf")

	>>> prompt = "Hey, are you conscious? Can you talk to me?"
	>>> inputs = tokenizer(prompt, return_tensors="pt")

	>>> # Generate
	>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
	>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
	"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
	```"""
	outputs = self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	use_cache=use_cache,
	cache_position=cache_position,
	**kwargs,
	)

	hidden_states = outputs[0]
	# Only compute necessary logixts, and do not upcast them to float if we are not computing the loss
	slice_indices = (
	slice(-logits_to_keep, None)
	if isinstance(logits_to_keep, int)
	else logits_to_keep
	)
	logits = self.lm_head(hidden_states[:, slice_indices, :])
	loss = None
	if labels is not None:
	loss = self.loss_function(
	logits=logits,
	labels=labels,
	vocab_size=self.config.vocab_size,
	**kwargs,
	)

	return CausalLMOutputWithPast(
	loss=loss,
	logits=logits,
	past_key_values=outputs.past_key_values,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	)

	__all__ = [
	"MobileLLMP1ForCausalLM",
	"MobileLLMP1TextModel",
	"MobileLLMP1TextDecoderLayer",
	"MobileLLMP1TextAttention",
	]