Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+step-bar-chart.png filter=lfs diff=lfs merge=lfs -text
+model.safetensors.index.json filter=lfs diff=lfs merge=lfs -text
+quantization_config.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,581 @@
+---
+license: apache-2.0
+base_model:
+- stepfun-ai/step-3.5-flash
+library_name: transformers
+---
+
+# Step 3.5 Flash
+
+<div align="center">
+  
+<div align="center" style="display: flex; justify-content: center; align-items: center;">
+  <img src="stepfun.svg" width="25" style="margin-right: 10px;"/>
+  <h1 style="margin: 0; border-bottom: none;">Step 3.5 Flash</h1>
+</div>
+
+[![GitHub](https://img.shields.io/badge/GitHub-181717?style=flat&logo=github&logoColor=white)](https://github.com/stepfun-ai/Step-3.5-Flash)
+[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20HF-StepFun/STEP3p5-preview)](https://huggingface.co/stepfun-ai/Step-3.5-Flash)
+[![ModelScope](https://img.shields.io/badge/ModelScope-StepFun/STEP3p5-preview)](https://modelscope.cn/models/stepfun-ai/Step-3.5-Flash)
+[![Discord](https://img.shields.io/badge/Discord-Join-5865F2?logo=discord&logoColor=white)](https://discord.gg/RcMJhNVAQc)
+[![Webpage](https://img.shields.io/badge/Webpage-Blog-blue)](https://static.stepfun.com/blog/step-3.5-flash/)
+[![Paper](https://img.shields.io/badge/Arxiv-TechReport-red)](https://arxiv.org/abs/2602.10604)
+[![License](https://img.shields.io/badge/License-Apache%202.0-green)]()
+[![Chat with the model on OpenRouter](https://img.shields.io/badge/Chat%20with%20the%20model-OpenRouter-5B3DF5?logo=chatbot&logoColor=white)](https://openrouter.ai/chat?models=stepfun/step-3.5-flash:free)
+[![Chat with the model on HuggingfaceSpace](https://img.shields.io/badge/Chat%20with%20the%20model-HuggingfaceSpace-5B3DF5?logo=chatbot&logoColor=white)](https://huggingface.co/spaces/stepfun-ai/Step-3.5-Flash)
+</div>
+
+
+
+## 1. Introduction 
+
+**Step 3.5 Flash** ([visit website](https://static.stepfun.com/blog/step-3.5-flash/)) is our most capable open-source foundation model, engineered to deliver frontier reasoning and agentic capabilities with exceptional efficiency. Built on a sparse Mixture of Experts (MoE) architecture, it selectively activates only 11B of its 196B parameters per token. This "intelligence density" allows it to rival the reasoning depth of top-tier proprietary models, while maintaining the agility required for real-time interaction.
+
+## 2. Key Capabilities
+
+- **Deep Reasoning at Speed**: While chatbots are built for reading, agents must reason fast. Powered by 3-way Multi-Token Prediction (MTP-3), Step 3.5 Flash achieves a generation throughput of **100–300 tok/s** in typical usage (peaking at **350 tok/s** for single-stream coding tasks). This allows for complex, multi-step reasoning chains with immediate responsiveness.
+
+- **A Robust Engine for Coding & Agents**: Step 3.5 Flash is purpose-built for agentic tasks, integrating a scalable RL framework that drives consistent self-improvement. It achieves **74.4% on SWE-bench Verified** and **51.0% on Terminal-Bench 2.0**, proving its ability to handle sophisticated, long-horizon tasks with unwavering stability.
+
+- **Efficient Long Context**: The model supports a cost-efficient **256K context window** by employing a 3:1 Sliding Window Attention (SWA) ratio—integrating three SWA layers for every full-attention layer. This hybrid approach ensures consistent performance across massive datasets or long codebases while significantly reducing the computational overhead typical of standard long-context models.
+
+- **Accessible Local Deployment**: Optimized for accessibility, Step 3.5 Flash brings elite-level intelligence to local environments. It runs securely on high-end consumer hardware (e.g., Mac Studio M4 Max, NVIDIA DGX Spark), ensuring data privacy without sacrificing performance.
+
+## 3. Performance
+
+Step 3.5 Flash delivers performance parity with leading closed-source systems while remaining open and efficient.
+
+![](step-bar-chart.png)
+
+Performance of Step 3.5 Flash measured across **Reasoning**, **Coding**, and **Agentic Abilities**. Open-source models (left) are sorted by their total parameter count, while top-tier proprietary models are shown on the right. xbench-DeepSearch scores are sourced from [official publications](https://xbench.org/agi/aisearch) for consistency. The shadowed bars represent the enhanced performance of Step 3.5 Flash using [Parallel Thinking](https://arxiv.org/pdf/2601.05593).
+
+### Detailed Benchmarks
+
+| Benchmark | Step 3.5 Flash | DeepSeek V3.2 | Kimi K2 Thinking / K2.5 | GLM-4.7 | MiniMax M2.1 | MiMo-V2 Flash |
+| --- | --- | --- | --- | --- | --- | --- |
+| # Activated Params | 11B | 37B | 32B | 32B | 10B | 15B |
+| # Total Params (MoE) | 196B | 671B | 1T | 355B | 230B | 309B |
+| Est. decoding cost @ 128K context, Hopper GPU** | **1.0x**<br>100 tok/s, MTP-3, EP8 | **6.0x**<br>33 tok/s, MTP-1, EP32 | **18.9x**<br>33 tok/s, no MTP, EP32 | **18.9x**<br>100 tok/s, MTP-3, EP8 | **3.9x**<br>100 tok/s, MTP-3, EP8 | **1.2x**<br>100 tok/s, MTP-3, EP8 |
+| | | | **Agent** | | | |
+| τ²-Bench | 88.2 | 80.3 (85.2*) | 74.3*/85.4* | 87.4 | 86.6* | 80.3 (84.1*) |
+| BrowseComp | 51.6 | 51.4 | 41.5* / 60.6 | 52.0 | 47.4 | 45.4 |
+| BrowseComp (w/ Context Manager) | 69.0 | 67.6 | 60.2/74.9 | 67.5 | 62.0 | 58.3 |
+| BrowseComp-ZH | 66.9 | 65.0 | 62.3 / 62.3* | 66.6 | 47.8* | 51.2* |
+| BrowseComp-ZH (w/ Context Manager) | 73.7 | — | —/— | — | — | — |
+| GAIA (no file) | 84.5 | 75.1* | 75.6*/75.9* | 61.9* | 64.3* | 78.2* |
+| xbench-DeepSearch (2025.05) | 83.7 | 78.0* | 76.0*/76.7* | 72.0* | 68.7* | 69.3* |
+| xbench-DeepSearch (2025.10) | 56.3 | 55.7* | —/40+ | 52.3* | 43.0* | 44.0* |
+| ResearchRubrics | 65.3 | 55.8* | 56.2*/59.5* | 62.0* | 60.2* | 54.3* |
+| | | | **Reasoning** | | | |
+| AIME 2025 | 97.3 | 93.1 | 94.5/96.1 | 95.7 | 83.0 | 94.1 (95.1*) |
+| HMMT 2025 (Feb.) | 98.4 | 92.5 | 89.4/95.4 | 97.1 | 71.0* | 84.4 (95.4*) |
+| HMMT 2025 (Nov.) | 94.0 | 90.2 | 89.2*/— | 93.5 | 74.3* | 91.0* |
+| IMOAnswerBench | 85.4 | 78.3 | 78.6/81.8 | 82.0 | 60.4* | 80.9* |
+| | | | **Coding** | | | |
+| LiveCodeBench-V6 | 86.4 | 83.3 | 83.1/85.0 | 84.9 | — | 80.6 (81.6*) |
+| SWE-bench Verified | 74.4 | 73.1 | 71.3/76.8 | 73.8 | 74.0 | 73.4 |
+| Terminal-Bench 2.0 | 51.0 | 46.4 | 35.7*/50.8 | 41.0 | 47.9 | 38.5 |
+
+**Notes**:
+1. "—" indicates the score is not publicly available or not tested.
+2. "*" indicates the original score was inaccessible or lower than our reproduced, so we report the evaluation under the same test conditions as Step 3.5 Flash to ensure fair comparability.
+3. **BrowseComp (with Context Manager)**: When the effective context length exceeds a predefined threshold, the agent resets the context and restarts the agent loop. By contrast, Kimi K2.5 and DeepSeek-V3.2 used a "discard-all" strategy.
+4. **Decoding Cost**: Estimates are based on a methodology similar to, but more accurate than, the approach described arxiv.org/abs/2507.19427
+
+### Recommended Inference Parameters
+1. For general chat domain, we suggest: `temperature=0.6, top_p=0.95`
+2. For reasoning / agent scenario, we recommend: `temperature=1.0, top_p=0.95`.
+
+## 4. Architecture Details
+
+Step 3.5 Flash is built on a **Sparse Mixture-of-Experts (MoE)** transformer architecture, optimized for high throughput and low VRAM usage during inference.
+
+### 4.1 Technical Specifications
+
+| Component | Specification |
+| :--- | :--- |
+| **Backbone** | 45-layer Transformer (4,096 hidden dim) |
+| **Context Window** | 256K |
+| **Vocabulary** | 128,896 tokens |
+| **Total Parameters** | **196.81B** (196B Backbone + 0.81B Head) |
+| **Active Parameters** | **~11B** (per token generation) |
+
+### 4.2 Mixture of Experts (MoE) Routing
+
+Unlike traditional dense models, Step 3.5 Flash uses a fine-grained routing strategy to maximize efficiency:
+- **Fine-Grained Experts**: 288 routed experts per layer + 1 shared expert (always active).
+- **Sparse Activation**: Only the Top-8 experts are selected per token.
+- **Result**: The model retains the "memory" of a 196B parameter model but executes with the speed of an 11B model. 
+
+### 4.3 Multi-Token Prediction (MTP)
+
+To improve inference speed, we utilize a specialized MTP Head consisting of a sliding-window attention mechanism and a dense Feed-Forward Network (FFN). This module predicts 4 tokens simultaneously in a single forward pass, significantly accelerating inference without degrading quality.
+
+## 5. Quick Start
+
+You can get started with Step 3.5 Flash in minutes using Cloud API via our supported providers.
+
+### 5.1 Get Your API Key.
+
+Sign up at [OpenRouter](https://openrouter.ai) or [platform.stepfun.ai](https://platform.stepfun.ai), and grab your API key. 
+
+> OpenRouter now offers free trial for Step 3.5 Flash.
+
+| Provider | Website | Base URL |
+| :--- | :--- | :--- |
+| OpenRouter | https://openrouter.ai | https://openrouter.ai/api/v1 |
+| StepFun | https://platform.stepfun.ai | https://api.stepfun.ai/v1 |
+
+### 5.2 Setup
+
+Install the standard OpenAI SDK (compatible with both platforms).
+
+```bash
+pip install --upgrade "openai>=1.0"
+```
+
+Note: OpenRouter supports multiple SDKs. Learn more [here](https://openrouter.ai/docs/quickstart).
+
+### 5.3 Implementation Example
+
+This example shows starting a chat with Step 3.5 Flash.
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+    api_key="YOUR_API_KEY",
+    base_url="https://api.stepfun.ai/v1", # or "https://openrouter.ai/api/v1"
+    # Optional: OpenRouter headers for app rankings
+    default_headers={
+        "HTTP-Referer": "<YOUR_SITE_URL>", 
+        "X-Title": "<YOUR_SITE_NAME>",
+    }
+)
+
+completion = client.chat.completions.create(
+    model="step-3.5-flash", # Use "stepfun/step-3.5-flash" for OpenRouter
+    messages=[
+        {
+            "role": "system",
+            "content": "You are an AI chat assistant provided by StepFun. You are good at Chinese, English, and many other languages.",
+        },
+        {
+            "role": "user",
+            "content": "Introduce StepFun's artificial intelligence capabilities."
+        },
+    ],
+)
+
+print(completion.choices[0].message.content)
+```
+
+## 6. Local Deployment
+
+Step 3.5 Flash is optimized for local inference and supports industry-standard backends including vLLM, SGLang, Hugging Face Transformers and llama.cpp.
+
+### 6.1 vLLM 
+We recommend using the latest nightly build of vLLM.
+1. Install vLLM.
+
+```bash
+# via Docker
+docker pull vllm/vllm-openai:nightly
+
+# or via pip (nightly wheels)
+pip install -U vllm --pre \
+  --index-url https://pypi.org/simple \
+  --extra-index-url https://wheels.vllm.ai/nightly
+```
+2. Launch the server.
+
+**Note**: Full MTP3 support is not yet available in vLLM. We are actively working on a Pull Request to integrate this feature, which is expected to significantly enhance decoding performance.
+
+  - For fp8 model
+```bash  
+vllm serve <MODEL_PATH_OR_HF_ID> \
+  --served-model-name step3p5-flash \
+  --tensor-parallel-size 8 \
+  --enable-expert-parallel \
+  --disable-cascade-attn \
+  --reasoning-parser step3p5 \
+  --enable-auto-tool-choice \
+  --tool-call-parser step3p5 \
+  --hf-overrides '{"num_nextn_predict_layers": 1}' \
+  --speculative_config '{"method": "step3p5_mtp", "num_speculative_tokens": 1}' \
+  --trust-remote-code \
+  --quantization fp8
+```
+
+  - For bf16 model
+```bash
+vllm serve <MODEL_PATH_OR_HF_ID> \
+  --served-model-name step3p5-flash \
+  --tensor-parallel-size 8 \
+  --enable-expert-parallel \
+  --disable-cascade-attn \
+  --reasoning-parser step3p5 \
+  --enable-auto-tool-choice \
+  --tool-call-parser step3p5 \
+  --hf-overrides '{"num_nextn_predict_layers": 1}' \
+  --speculative_config '{"method": "step3p5_mtp", "num_speculative_tokens": 1}' \
+  --trust-remote-code    
+```
+You can also refer to the [Step-3.5-Flash](https://github.com/vllm-project/recipes/blob/main/StepFun/Step-3.5-Flash.md) recipe.
+
+### 6.2 SGLang
+
+1. Install SGLang.
+```bash
+# via Docker
+docker pull lmsysorg/sglang:dev-pr-18084
+# or from source (pip)
+pip install "sglang[all] @ git+https://github.com/sgl-project/sglang.git"
+```
+
+2. Launch the server.
+  - For bf16 model
+
+```bash
+sglang serve --model-path <MODEL_PATH_OR_HF_ID> \
+  --served-model-name step3p5-flash \
+  --tp-size 8 \
+  --tool-call-parser step3p5 \
+  --reasoning-parser step3p5 \
+  --speculative-algorithm EAGLE \
+  --speculative-num-steps 3 \
+  --speculative-eagle-topk 1 \
+  --speculative-num-draft-tokens 4 \
+  --enable-multi-layer-eagle \
+  --host 0.0.0.0 \
+  --port 8000
+```
+  - For fp8 model
+```bash
+sglang serve --model-path <MODEL_PATH_OR_HF_ID> \
+  --served-model-name step3p5-flash \
+  --tp-size 8 \
+  --ep-size 8 \
+  --tool-call-parser step3p5 \
+  --reasoning-parser step3p5 \
+  --speculative-algorithm EAGLE \
+  --speculative-num-steps 3 \
+  --speculative-eagle-topk 1 \
+  --speculative-num-draft-tokens 4 \
+  --enable-multi-layer-eagle \
+  --host 0.0.0.0 \
+  --port 8000
+```
+
+### 6.3 Transformers (Debug / Verification)
+
+Use this snippet for quick functional verification. For high-throughput serving, use vLLM or SGLang.
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL_PATH = "<MODEL_PATH_OR_HF_ID>"
+
+# 1. Setup
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    torch_dtype="auto",
+    device_map="auto",
+)
+
+# 2. Prepare Input
+messages = [{"role": "user", "content": "Explain the significance of the number 42."}]
+inputs = tokenizer.apply_chat_template(
+    messages,
+    tokenize=True,
+    add_generation_prompt=True,
+    return_dict=True,
+    return_tensors="pt",
+).to(model.device)
+
+# 3. Generate
+generated_ids = model.generate(**inputs, max_new_tokens=128, do_sample=False)
+output_text = tokenizer.decode(generated_ids[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+
+print(output_text)
+```
+
+### 6.4 llama.cpp
+
+#### System Requirements
+- GGUF Model Weights(int4): 111.5 GB
+- Runtime Overhead: ~7 GB
+- Minimum VRAM: 120 GB (e.g., Mac studio, DGX-Spark, AMD Ryzen AI Max+ 395)
+- Recommended: 128GB unified memory
+#### Steps
+1. Use official llama.cpp:
+> the folder `Step-3.5-Flash/tree/main/llama.cpp` is **obsolete**
+```bash
+git clone https://github.com/ggml-org/llama.cpp
+cd llama.cpp
+```
+2. Build llama.cpp on Mac:
+```bash
+cmake -S . -B build-macos \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DGGML_METAL=ON \
+  -DGGML_ACCELERATE=ON \
+  -DLLAMA_BUILD_EXAMPLES=ON \
+  -DLLAMA_BUILD_COMMON=ON \
+  -DGGML_LTO=ON
+cmake --build build-macos -j8
+```
+3. Build llama.cpp on DGX-Spark:
+```bash
+cmake -S . -B build-cuda \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DGGML_CUDA=ON \
+  -DGGML_CUDA_GRAPHS=ON \
+  -DLLAMA_CURL=OFF \
+  -DLLAMA_BUILD_EXAMPLES=ON \
+  -DLLAMA_BUILD_COMMON=ON
+cmake --build build-cuda -j8
+```
+4. Build llama.cpp on AMD Windows:
+```bash
+cmake -S . -B build-vulkan \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DLLAMA_CURL=OFF \
+  -DGGML_OPENMP=ON \
+  -DGGML_VULKAN=ON
+cmake --build build-vulkan -j8
+```
+5. Run with llama-cli
+```bash
+./llama-cli -m step3.5_flash_Q4_K_S.gguf -c 16384 -b 2048 -ub 2048 -fa on --temp 1.0 -p "What's your name?"
+```
+6. Test performance with llama-batched-bench:
+```bash
+./llama-batched-bench -m step3.5_flash_Q4_K_S.gguf -c 32768 -b 2048 -ub 2048 -npp 0,2048,8192,16384,32768 -ntg 128 -npl 1
+```
+
+## 7. Using Step 3.5 Flash on Agent Platforms
+
+### 7.1 Claude Code & Codex
+It's straightforward to add Step 3.5 Flash to the list of models in most coding environments. See below for the instructions for configuring Claude Code and Codex to use Step 3.5 Flash.
+
+#### 7.1.1 Prerequisites
+Sign up at StepFun.ai or OpenRouter and grab an API key, as mentioned in the Quick Start.
+
+#### 7.1.2 Environment setup
+Claude Code and Codex rely on Node.js. We recommend installing Node.js version > v20. You can install Node via nvm.
+
+**Mac/Linux**:
+```bash
+# Install nvm on Mac/Linux via curl：
+# Step 1
+curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.0/install.sh | bash
+
+# Copy the full command
+export NVM_DIR="$HOME/.nvm"
+[ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"  # This loads nvm
+[ -s "$NVM_DIR/bash_completion" ] && \. "$NVM_DIR/bash_completion"
+
+# Users in China can set up npm mirror
+config set registry https://registry.npmmirror.com
+
+# Step 2
+nvm install v22
+
+# Make sure Node.js is installed
+node --version
+
+npm --version
+```
+
+**Windows**:
+You can download the installation file (`nvm-setup.exe`) from [https://github.com/coreybutler/nvm-windows/releases](https://github.com/coreybutler/nvm-windows/releases). Follow the instructions to install nvm. Run nvm commands to make sure it is installed.
+
+#### 7.1.3 Use Step 3.5 Flash on Claude Code
+
+1. Install Claude Code.
+```bash
+# install claude code via npm
+npm install -g @anthropic-ai/claude-code
+
+# test if the installation is successful
+claude --version 
+```
+
+2. Configure Claude Code.
+
+  To accommodate diverse workflows in Claude Code, we support both **Anthropic-style** and **OpenAI-style** APIs. 
+
+  **Option A: Anthropic API style**:
+
+  > If you intend to use the **OpenRouter** API, refer to the OpenRouter integration guide.  
+
+  Step 1: Edit Claude Settings. Update `~/.claude/settings.json`.
+  > You only need to modify the fields shown below. Leave the rest of the file unchanged.
+
+  ```json
+  {
+  "env": {
+    "ANTHROPIC_API_KEY": "API_KEY_from_StepFun",
+    "ANTHROPIC_BASE_URL": "https://api.stepfun.ai/"
+  },
+  "model": "step-3.5-flash"
+  }
+  ```
+  Step 2: Start Claude Code.
+  
+  Save the file, and then start Claude Code. Run `/status` to confirm the model and base URL.
+
+  ```txt
+  ❯ /status
+  ─────────────────────────────────────────────────────────────────────────────────
+  Settings:  Status   Config   Usage  (←/→ or tab to cycle)
+
+  Version: 2.1.1
+  Session name: /rename to add a name
+  Session ID: 676dae61-259d-4eef-8c2f-0f1641600553
+  cwd: /Users/step-test/
+  Auth token: none
+  API key: ANTHROPIC_API_KEY
+  Anthropic base URL: https://api.stepfun.ai/
+
+  Model: step-3.5-flash
+  Setting sources: User settings
+  ```
+
+  **Option B: OpenAI API style**
+
+  > Note: OpenAI API style here refers to the `chat/completions/` format.
+
+  > We recommend using `claude-code-router`. For details, see [https://github.com/musistudio/claude-code-router](https://github.com/musistudio/claude-code-router).
+
+  After Claude Code is installed, install `claude-code-router` :
+
+  ```bash
+  # install ccr via npm
+  npm install -g @musistudio/claude-code-router
+
+  # validate it is installed
+  ccr -v
+  ```
+
+  Add the following configurations to `~/.claude-code-router/config.json`.
+
+  ```json
+  {
+  "PORT": 3456,
+  "Providers": [
+    {
+      "name": "stepfun-api",
+      "api_base_url": "https://api.stepfun.com/v1/chat/completions",
+      "api_key": "StepFun_API_KEY",
+      "models": ["step-3.5-flash"],
+      "transformer":{
+           "step-3.5-flash": { "use": ["OpenAI"]}
+      }
+    }
+  ],
+  "Router": {
+    "default": "stepfun-api,step-3.5-flash",
+    "background": "stepfun-api,step-3.5-flash",
+    "think": "stepfun-api,step-3.5-flash",
+    "longContext": "stepfun-api,step-3.5-flash",
+    "webSearch": "stepfun-api,step-3.5-flash"
+  }
+  }
+  ```
+  You can now start Claude Code:
+
+  ```bash
+  # Start Claude
+  ccr code 
+
+  # restart ccr if configs are changed
+  ccr restart 
+  ```
+
+#### 7.1.4 Use Step 3.5 Flash on Codex
+1. Install Codex
+```bash
+# Install codex via npm
+npm install -g @openai/codex
+
+# Test if it is installed
+codex --version
+```
+
+2. Configure Codex
+Add the following settings to `~/.codex/config.toml`, keeping the rest of the settings as they are.
+
+```json
+model="step-3.5-flash"
+model_provider = "stepfun-chat"
+preferred_auth_method = "apikey"
+
+# configure the provider
+[model_providers.stepfun-chat]
+name = "OpenAI using response"
+base_url = "https://api.stepfun.com/v1"
+env_key = "OPENAI_API_KEY"
+wire_api = "chat"
+query_params = {}
+```
+
+For Codex, `wire_api` only supports `chat` . If you use the `responses` mode, you'll need to change to `chat`. Please also switch `model_provider` to the newly configured `stepfun-chat`.
+
+When finishing the configuration, run codex in a new Terminal window to start Codex. Run `/status` to check the configuration.
+
+```bash
+/status
+📂 Workspace
+  • Path: /Users/step-test/
+  • Approval Mode: on-request
+  • Sandbox: workspace-write
+  • AGENTS files: (none)
+
+🧠 Model
+  • Name: step-3.5-flash
+  • Provider: Stepfun-chat
+
+💻 Client
+  • CLI Version: 0.40.0
+```
+
+#### 7.1.5 Use Step 3.5 Flash on Step-DeepResearch (DeepResearch)
+1. Use the reference environment setup below and configure `MODEL_NAME` to `Step-3.5-Flash`. [https://github.com/stepfun-ai/StepDeepResearch?tab=readme-ov-file#1-environment-setup](https://github.com/stepfun-ai/StepDeepResearch?tab=readme-ov-file#1-environment-setup)
+
+
+## 8. Known Issues and Future Directions
+
+1. **Token Efficiency**. Step 3.5 Flash achieves frontier-level agentic intelligence but currently relies on longer generation trajectories than Gemini 3.0 Pro to reach comparable quality.
+2. **Efficient Universal Mastery**. We aim to unify generalist versatility with deep domain expertise. To achieve this efficiently, we are advancing variants of on-policy distillation, allowing the model to internalize expert behaviors with higher sample efficiency.
+3. **RL for More Agentic Tasks**. While Step 3.5 Flash demonstrates competitive performance on academic agentic benchmarks, the next frontier of agentic AI necessitates the application of RL to intricate, expert-level tasks found in professional work, engineering, and research.
+4. **Operational Scope and Constraints**. Step 3.5 Flash is tailored for coding and work-centric tasks, but may experience reduced stability during distribution shifts. This typically occurs in highly specialized domains or long-horizon, multi-turn dialogues, where the model may exhibit repetitive reasoning, mixed-language outputs, or inconsistencies in time and identity awareness.
+
+## 9. Co-Developing the Future
+
+We view our roadmap as a living document, evolving continuously based on real-world usage and developer feedback.
+As we work to shape the future of AGI by expanding broad model capabilities, we want to ensure we are solving the right problems. We invite you to be part of this continuous feedback loop—your insights directly influence our priorities.
+
+- **Join the Conversation**: Our Discord community is the primary hub for brainstorming future architectures, proposing capabilities, and getting early access updates 🚀
+- **Report Friction**: Encountering limitations? You can open an issue on GitHub or flag it directly in our Discord support channels.
+
+## 📜 Citation
+
+If you find this project useful in your research, please cite our technical report:
+
+```tex
+@misc{huang2026step35flashopen,
+      title={Step 3.5 Flash: Open Frontier-Level Intelligence with 11B Active Parameters}, 
+      author={Ailin Huang and Ang Li and Aobo Kong and Bin Wang and Binxing Jiao and Bo Dong and Bojun Wang and Boyu Chen and Brian Li and Buyun Ma and Chang Su and Changxin Miao and Changyi Wan and Chao Lou and Chen Hu and Chen Xu and Chenfeng Yu and Chengting Feng and Chengyuan Yao and Chunrui Han and Dan Ma and Dapeng Shi and Daxin Jiang and Dehua Ma and Deshan Sun and Di Qi and Enle Liu and Fajie Zhang and Fanqi Wan and Guanzhe Huang and Gulin Yan and Guoliang Cao and Guopeng Li and Han Cheng and Hangyu Guo and Hanshan Zhang and Hao Nie and Haonan Jia and Haoran Lv and Hebin Zhou and Hekun Lv and Heng Wang and Heung-Yeung Shum and Hongbo Huang and Hongbo Peng and Hongyu Zhou and Hongyuan Wang and Houyong Chen and Huangxi Zhu and Huimin Wu and Huiyong Guo and Jia Wang and Jian Zhou and Jianjian Sun and Jiaoren Wu and Jiaran Zhang and Jiashu Lv and Jiashuo Liu and Jiayi Fu and Jiayu Liu and Jie Cheng and Jie Luo and Jie Yang and Jie Zhou and Jieyi Hou and Jing Bai and Jingcheng Hu and Jingjing Xie and Jingwei Wu and Jingyang Zhang and Jishi Zhou and Junfeng Liu and Junzhe Lin and Ka Man Lo and Kai Liang and Kaibo Liu and Kaijun Tan and Kaiwen Yan and Kaixiang Li and Kang An and Kangheng Lin and Lei Yang and Liang Lv and Liang Zhao and Liangyu Chen and Lieyu Shi and Liguo Tan and Lin Lin and Lina Chen and Luck Ma and Mengqiang Ren and Michael Li and Ming Li and Mingliang Li and Mingming Zhang and Mingrui Chen and Mitt Huang and Na Wang and Peng Liu and Qi Han and Qian Zhao and Qinglin He and Qinxin Du and Qiuping Wu and Quan Sun and Rongqiu Yang and Ruihang Miao and Ruixin Han and Ruosi Wan and Ruyan Guo and Shan Wang and Shaoliang Pang and Shaowen Yang and Shengjie Fan and Shijie Shang and Shiliang Yang and Shiwei Li and Shuangshuang Tian and Siqi Liu and Siye Wu and Siyu Chen and Song Yuan and Tiancheng Cao and Tianchi Yue and Tianhao Cheng and Tianning Li and Tingdan Luo and Wang You and Wei Ji and Wei Yuan and Wei Zhang and Weibo Wu and Weihao Xie and Wen Sun and Wenjin Deng and Wenzhen Zheng and Wuxun Xie and Xiangfeng Wang and Xiangwen Kong and Xiangyu Liu and Xiangyu Zhang and Xiaobo Yang and Xiaojia Liu and Xiaolan Yuan and Xiaoran Jiao and Xiaoxiao Ren and Xiaoyun Zhang and Xin Li and Xin Liu and Xin Wu and Xing Chen and Xingping Yang and Xinran Wang and Xu Zhao and Xuan He and Xuanti Feng and Xuedan Cai and Xuqiang Zhou and Yanbo Yu and Yang Li and Yang Xu and Yanlin Lai and Yanming Xu and Yaoyu Wang and Yeqing Shen and Yibo Zhu and Yichen Lv and Yicheng Cao and Yifeng Gong and Yijing Yang and Yikun Yang and Yin Zhao and Yingxiu Zhao and Yinmin Zhang and Yitong Zhang and Yixuan Zhang and Yiyang Chen and Yongchi Zhao and Yongshen Long and Yongyao Wang and Yousong Guan and Yu Zhou and Yuang Peng and Yuanhao Ding and Yuantao Fan and Yuanzhen Yang and Yuchu Luo and Yudi Zhao and Yue Peng and Yueqiang Lin and Yufan Lu and Yuling Zhao and Yunzhou Ju and Yurong Zhang and Yusheng Li and Yuxiang Yang and Yuyang Chen and Yuzhu Cai and Zejia Weng and Zetao Hong and Zexi Li and Zhe Xie and Zheng Ge and Zheng Gong and Zheng Zeng and Zhenyi Lu and Zhewei Huang and Zhichao Chang and Zhiguo Huang and Zhiheng Hu and Zidong Yang and Zili Wang and Ziqi Ren and Zixin Zhang and Zixuan Wang},
+      year={2026},
+      eprint={2602.10604},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2602.10604}, 
+}
+```
+
+## License
+This project is open-sourced under the [Apache 2.0 License](https://www.apache.org/licenses/LICENSE-2.0).

chat_template.jinja

@@ -0,0 +1,80 @@
+{% macro render_content(content) %}{% if content is none %}{{- '' }}{% elif content is string %}{{- content }}{% elif content is mapping %}{{- content['value'] if 'value' in content else content['text'] }}{% elif content is iterable %}{% for item in content %}{% if item.type == 'text' %}{{- item['value'] if 'value' in item else item['text'] }}{% elif item.type == 'image' %}<im_patch>{% endif %}{% endfor %}{% endif %}{% endmacro %}
+{{bos_token}}{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- render_content(messages[0].content) + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou have access to the following functions in JSONSchema format:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson(ensure_ascii=False) }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...>\n...\n</function> block must be nested within <tool_call>\n...\n</tool_call> XML tags\n- Required parameters MUST be specified\n</IMPORTANT><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + render_content(messages[0].content) + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and render_content(message.content) is string and not(render_content(message.content).startswith('<tool_response>') and render_content(message.content).endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content) %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {%- set role_name = 'observation' if (message.role == "system" and not loop.first and message.name == 'observation') else message.role %}
+        {{- '<|im_start|>' + role_name + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = render_content(message.reasoning_content) %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- else %}
+                {%- set reasoning_content = '' %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- if tool_call.arguments is defined %}
+                    {%- set arguments = tool_call.arguments %}
+                    {%- for args_name, args_value in arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson(ensure_ascii=False) | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>tool_response\n' }}
+        {%- endif %}
+        {{- '<tool_response>' }}
+        {{- content }}
+        {{- '</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n<think>\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,327 @@
+{
+    "architectures": [
+        "Step3p5ForCausalLM"
+    ],
+    "model_type": "step3p5",
+    "auto_map": {
+        "AutoConfig": "configuration_step3p5.Step3p5Config",
+        "AutoModelForCausalLM": "modeling_step3p5.Step3p5ForCausalLM"
+    },
+    "rope_scaling": {
+        "rope_type": "llama3",
+        "factor": 2.0,
+        "original_max_position_embeddings": 131072,
+        "low_freq_factor": 1.0,
+        "high_freq_factor": 32.0
+    },
+    "yarn_only_types": [
+        "full_attention"
+    ],
+    "hidden_size": 4096,
+    "intermediate_size": 11264,
+    "num_hidden_layers": 45,
+    "max_seq_len": 262144,
+    "vocab_size": 128896,
+    "torch_dtype": "bfloat16",
+    "use_qk_norm": true,
+    "moe_layers_enum": "3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44",
+    "num_attention_heads": 64,
+    "num_attention_groups": 8,
+    "head_dim": 128,
+    "use_moe": true,
+    "moe_num_experts": 288,
+    "moe_top_k": 8,
+    "moe_intermediate_size": 1280,
+    "share_expert_dim": 1280,
+    "moe_layer_offset": 0,
+    "moe_every_n_layer": 1,
+    "norm_expert_weight": true,
+    "moe_router_activation": "sigmoid",
+    "moe_router_scaling_factor": 3.0,
+    "att_impl_type": "GQA",
+    "rope_theta": [
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0,
+        5000000.0,
+        10000.0,
+        10000.0,
+        10000.0
+    ],
+    "use_head_wise_attn_gate": true,
+    "sliding_window": 512,
+    "use_moe_router_bias": true,
+    "need_fp32_gate": true,
+    "sink": false,
+    "layer_types": [
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention"
+    ],
+    "use_rope_layers": [],
+    "num_nextn_predict_layers": 3,
+    "partial_rotary_factors": [
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0,
+        0.5,
+        1.0,
+        1.0,
+        1.0
+    ],
+    "eos_token_id": [
+        1,
+        2,
+        128007
+    ],
+    "bos_token_id": 0,
+    "pad_token_id": 0,
+    "attention_other_setting": {
+        "attention_type": "sliding_attention",
+        "num_attention_heads": 96,
+        "num_attention_groups": 8,
+        "head_dim": 128,
+        "true_head_dim": 128
+    },
+    "swiglu_limits": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        7,
+        7,
+        0.0,
+        0.0,
+        0.0
+    ],
+    "swiglu_limits_shared": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        16,
+        0.0,
+        0.0,
+        0.0
+    ],
+    "zero_centered": true,
+    "max_position_embeddings": 262144,
+    "quantization_config": {
+        "quant_method": "exl3",
+        "version": "0.0.22",
+        "bits": 3.0,
+        "head_bits": 6,
+        "calibration": {
+            "rows": 250,
+            "cols": 2048
+        },
+        "out_scales": "always",
+        "codebook": "mcg"
+    }
+}

configuration_step3p5.py

@@ -0,0 +1,60 @@
+from typing import Any, Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+
+class Step3p5Config(PretrainedConfig):
+    model_type = "step3p5"
+    architectures = ["Step3p5ForCausalLM"]
+
+    def __init__(
+        self,
+        hidden_size: int = 4096,
+        intermediate_size: int = 11264,
+        num_attention_heads: int = 64,
+        num_attention_groups: int = 8,
+        num_hidden_layers: int = 45,
+        max_seq_len: int = 128000,
+        vocab_size: int = 128815,
+        rms_norm_eps: float = 1e-5,
+        moe_intermediate_size: int = 1280,
+        moe_num_experts: int = 288,
+        moe_top_k: int = 8,
+        rope_theta: float = 10000,
+        rope_scaling: Optional[dict[str, Any]] = None,
+        max_position_embeddings: int = 128000,
+        share_expert_dims: int = 1280,
+        head_dim: int = 128,
+        norm_expert_weight: bool = True,
+        layer_types: list[str] = None,
+        sliding_window: Optional[int] = None,
+        moe_layers_enum: tuple[int] = (3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
+                                       15, 16, 17, 18, 19, 20, 21, 22, 23, 24,
+                                       25, 26, 27, 28, 29, 30, 31, 32, 33, 34,
+                                       35, 36, 37, 38, 39, 40, 41, 42, 43, 44),
+        **kwargs,
+    ) -> None:
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.num_attention_groups = num_attention_groups
+        self.num_hidden_layers = num_hidden_layers
+        self.max_seq_len = max_seq_len
+        self.vocab_size = vocab_size
+        self.rms_norm_eps = rms_norm_eps
+        self.moe_intermediate_size = moe_intermediate_size
+        self.moe_num_experts = moe_num_experts
+        self.moe_top_k = moe_top_k
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.max_position_embeddings = max_position_embeddings
+        self.share_expert_dim = share_expert_dims
+        self.head_dim = head_dim
+        self.norm_expert_weight = norm_expert_weight
+        self.moe_layers_enum = moe_layers_enum
+        self.layer_types = layer_types
+        self.sliding_window = sliding_window
+        self.use_cache = True
+        super().__init__(**kwargs)
+

model-00001-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df8e12f0932a196c7c1c3a000c1cf61dc30b31b48cba12340b2e4ba8e7bdf5cd
+size 8345423429

model-00002-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eacc0fa73340dbcaf5d44bc2d68a5faf7264a9d628c6f54ecb51757abf75ff3c
+size 7021015425

model-00003-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c92ee7816e7ffafc3f8181d3946d206c673406cf8e0cbe1c8dc01bbbe3ab491e
+size 7021025898

model-00004-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fb089a6aea4c503587606216ebc55096ff9fddd5bb5cc95b7a1a69040855423
+size 7021025898

model-00005-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d459a78da9947ac7137cb03560ea96f980f23f2e531b79369b84ba26441a19a
+size 7021025898

model-00006-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d86dac4c2eedbd6b880d398b836db5b74c68d5d92af0f6740d45ad2f1f5d8eb6
+size 7021025898

model-00007-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7f477e02ac7976ac922158c149e5f54e4a6ff2530bf79334857fc8d78f61d4b
+size 7021025898

model-00008-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59e4e130cd44396626eba4efec086790241b15502b7ddf9833617ddbbce720a6
+size 7021025898

model-00009-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b6031033e0f7384307c67a2b56a0b9a8cb0c4278f0b9c01c702fa36530d7769
+size 7021025898

model-00010-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:888980ee35bfc7b35aae5b5030269697bedbf7b661f44d503827a3d929fde4a0
+size 7021025898

model-00011-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c94a7ddc675f9c4c2ef2124de16dcb7c8ae7f993523e03f4a76c49f2a2d5d5bd
+size 3900323099

model.safetensors.index.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3d066152700a34baa6c114d2aa7c6420a5dcf08a3cde4f2fcae3f8529855105
+size 13604905

modeling_step3p5.py

@@ -0,0 +1,899 @@
+# # Copyright 2025 The LLAMA4 and 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.
+from dataclasses import dataclass
+from typing import Callable, Optional, Tuple, Union
+
+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.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 BaseModelOutputWithPast, ModelOutput
+from transformers.modeling_rope_utils import (ROPE_INIT_FUNCTIONS,
+                                              dynamic_rope_update)
+from transformers.modeling_utils import (ALL_ATTENTION_FUNCTIONS,
+                                         PreTrainedModel)
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, can_return_tuple, logging
+
+from .configuration_step3p5 import Step3p5Config
+
+logger = logging.get_logger(__name__)
+
+__all__ = ["Step3p5Model", "Step3p5ForCausalLM"]
+
+class Step3p5RotaryEmbedding(nn.Module):
+
+    def __init__(self, config: Step3p5Config, device=None, layer_idx=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        self.layer_idx = layer_idx
+        if config.rope_parameters is not None:
+            self.rope_type = config.rope_parameters.get(
+                "rope_type", config.rope_parameters.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        partial_rotary_factors = getattr(config, "partial_rotary_factors",
+                                         None)
+        if partial_rotary_factors is not None:
+            config.partial_rotary_factor = partial_rotary_factors[
+                self.layer_idx]
+        else:
+            config.partial_rotary_factor = 1.0
+
+        self.rope_theta = config.rope_theta
+        if isinstance(config.rope_theta, list):
+            self.rope_theta = config.rope_theta.copy()
+            config.rope_theta = self.rope_theta[self.layer_idx]
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = self.rope_init_fn(
+            self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+        config.rope_theta = self.rope_theta
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
+            position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float().to(x.device)
+
+        device_type = x.device.type if isinstance(
+            x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type,
+                            enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float()
+                     @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+    # Apply rotary embeddings on the first half or full tensor
+    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+    # Concatenate back to full shape
+    q_embed = torch.cat([q_embed, q_pass], dim=-1)
+    k_embed = torch.cat([k_embed, k_pass], dim=-1)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :,
+                                  None, :, :].expand(batch,
+                                                     num_key_value_heads,
+                                                     n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen,
+                                 head_dim)
+
+
+# Adapted from transformers.models.llama.modeling_llama.eager_attention_forward -> llama4 doesn't cast attn weights to fp32
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+    # breakpoint()
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, :key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+    attn_weights = nn.functional.dropout(attn_weights,
+                                         p=dropout,
+                                         training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+@dataclass
+class Step3p5CausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[list[torch.FloatTensor]] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+
+
+class Step3p5MLP(nn.Module):
+
+    def __init__(self, config, intermediate_size=None, swiglu_limit=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size,
+                                   self.intermediate_size,
+                                   bias=False)
+        self.up_proj = nn.Linear(self.hidden_size,
+                                 self.intermediate_size,
+                                 bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size,
+                                   self.hidden_size,
+                                   bias=False)
+        self.act_fn = ACT2FN["silu"]
+        self.limit = swiglu_limit
+
+    def forward(self, x):
+        up = self.up_proj(x)
+        gate = self.act_fn(self.gate_proj(x))
+        if self.limit is not None:
+            gate = gate.clamp(min=None, max=self.limit)
+            up = up.clamp(min=-self.limit, max=self.limit)
+
+        return self.down_proj(gate * up)
+
+
+def sigmoid_routing_function(gating_output: torch.Tensor, topk: int,
+                             renormalize: bool):
+    gating_output = gating_output.float()
+    gate_prob = torch.sigmoid(gating_output)
+    gate_prob = gate_prob / gate_prob.sum(dim=-1, keepdim=True)
+    topk_prob, indices = torch.topk(gate_prob, k=topk, dim=1)
+    expert_topk_weight = topk_prob
+    if renormalize:
+        expert_topk_weight = expert_topk_weight / torch.sum(
+            expert_topk_weight, dim=-1, keepdim=True)
+    return expert_topk_weight, indices
+
+
+def softmax_routing_function(gating_output: torch.Tensor, top_k: int,
+                             renormalize: bool):
+    gating_output = gating_output.float()
+    gate_prob = torch.softmax(gating_output, dim=-1)
+    gate_prob = gate_prob / gate_prob.sum(dim=-1, keepdim=True)
+    topk_prob, indices = torch.topk(gate_prob, k=top_k, dim=1)
+    expert_topk_weight = topk_prob
+    if renormalize:
+        expert_topk_weight = expert_topk_weight / torch.sum(
+            expert_topk_weight, dim=-1, keepdim=True)
+    return expert_topk_weight, indices.to(torch.int32)
+
+
+class MoELinear(nn.Module):
+
+    def __init__(self, num_experts, in_features, out_features):
+        super().__init__()
+        self.num_experts = num_experts
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = nn.Parameter(
+            torch.empty(num_experts, out_features, in_features))
+
+    def forward(self, x, expert_id):
+        x = F.linear(x.float(), self.weight[expert_id].float())
+        return x
+
+
+class Step3p5MoEMLP(nn.Module):
+
+    def __init__(self, config, swiglu_limit=None):
+        super().__init__()
+        self.num_experts = config.moe_num_experts
+        self.top_k = config.moe_top_k
+        self.hidden_size = config.hidden_size
+        self.moe_intermediate_size = config.moe_intermediate_size
+
+        self.use_moe_router_bias = config.use_moe_router_bias
+        if self.use_moe_router_bias:
+            self.router_bias = nn.Parameter(torch.zeros(config.moe_num_experts,
+                                                        dtype=torch.float32),
+                                            requires_grad=False)
+            self.custom_routing_function = self.router_bias_func
+        elif config.moe_router_activation == "sigmoid":
+            self.custom_routing_function = sigmoid_routing_function
+        else:
+            self.custom_routing_function = None
+        self.need_fp32_gate = config.need_fp32_gate
+        self.routed_scaling_factor = getattr(config,
+                                             "moe_router_scaling_factor", 1.0)
+        
+        # gating
+        self.gate = nn.Linear(self.hidden_size, self.num_experts, bias=False)
+            
+        self.act_fn = ACT2FN["silu"]
+        self.limit = swiglu_limit
+
+        self.up_proj = MoELinear(self.num_experts, self.hidden_size,
+                                 self.moe_intermediate_size)
+        self.gate_proj = MoELinear(self.num_experts, self.hidden_size,
+                                   self.moe_intermediate_size)
+        self.down_proj = MoELinear(self.num_experts,
+                                   self.moe_intermediate_size,
+                                   self.hidden_size)
+
+    def router_bias_func(self, gating_output: torch.Tensor, topk: int,
+                         renormalize: bool):
+        gate_prob = torch.sigmoid(gating_output.float())
+        gate_prob_with_bias = gate_prob + self.router_bias.unsqueeze(0)
+        _, indices = torch.topk(gate_prob_with_bias, k=topk, dim=1)
+        topk_prob = torch.gather(gate_prob, 1, indices)
+        expert_topk_weight = topk_prob
+        if renormalize:
+            expert_topk_weight = expert_topk_weight / (
+                torch.sum(expert_topk_weight, dim=-1, keepdim=True) + 1e-20)
+        return expert_topk_weight, indices
+
+    def get_expert_output(self, inputs: torch.Tensor, expert_id):
+        #if self.limit is None:
+        up = self.up_proj(inputs, expert_id)
+        gate = self.act_fn(self.gate_proj(inputs, expert_id))
+        if self.limit is not None:
+            gate = gate.clamp(min=None, max=self.limit)
+            up = up.clamp(min=-self.limit, max=self.limit)
+
+        return self.down_proj(gate * up, expert_id)
+
+    def forward(self, hidden_states):
+        """ """
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        if self.need_fp32_gate:
+            router_logits = torch.matmul(hidden_states.to(torch.float32), self.gate.weight.t().to(torch.float32))
+        else:
+            # router_logits: (batch * sequence_length, n_experts)
+            router_logits = self.gate(hidden_states)
+        
+        if self.custom_routing_function:
+            routing_weights, selected_experts = self.custom_routing_function(
+                router_logits, self.top_k, renormalize=True)
+        else:
+            routing_weights = F.softmax(router_logits,
+                                        dim=1,
+                                        dtype=torch.float)
+            routing_weights, selected_experts = torch.topk(routing_weights,
+                                                           self.top_k,
+                                                           dim=-1)
+
+        routing_weights = routing_weights * self.routed_scaling_factor
+
+        final_hidden_states = torch.zeros(
+            (batch_size * sequence_length, hidden_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device)
+
+        # One hot encode the selected experts to create an expert mask
+        # this will be used to easily index which expert is going to be sollicitated
+        expert_mask = torch.nn.functional.one_hot(
+            selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+        # Loop over all available experts in the model and perform the computation on each expert
+        for expert_idx in range(self.num_experts):
+            idx, top_x = torch.where(expert_mask[expert_idx])
+
+            # Index the correct hidden states and compute the expert hidden state for
+            # the current expert. We need to make sure to multiply the output hidden
+            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
+            current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+            current_hidden_states = (
+                self.get_expert_output(current_state, expert_idx) *
+                routing_weights[top_x, idx, None])
+
+            # However `index_add_` only support torch tensors for indexing so we'll use
+            # the `top_x` tensor here.
+            final_hidden_states.index_add_(
+                0, top_x, current_hidden_states.to(hidden_states.dtype))
+        final_hidden_states = final_hidden_states.reshape(
+            batch_size, sequence_length, hidden_dim)
+        return final_hidden_states
+
+
+class Step3p5RMSNorm(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        normed = x * torch.rsqrt(variance + self.variance_epsilon)
+        normed = normed * (self.weight.float() + 1)
+        return normed.to(dtype)
+class Step3p5Attention(nn.Module):
+
+    def __init__(self, config: Step3p5Config, layer_idx):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.num_attention_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_attention_groups
+
+        layer_types = getattr(config, "layer_types", [])
+        if layer_types:
+            enable_sliding_window = layer_types[
+                self.layer_idx] == "sliding_attention"
+        else:
+            enable_sliding_window = self.layer_idx % 2 == 0
+        
+        if hasattr(config, "yarn_only_types") and layer_types[
+                self.layer_idx] not in config.yarn_only_types:
+            config.rope_parameters = None
+        else:
+            config.rope_parameters = getattr(config, "rope_scaling", None)
+
+        self.sliding_window = config.sliding_window
+        if enable_sliding_window:
+            self.num_attention_heads = config.attention_other_setting[
+                "num_attention_heads"]
+            self.num_key_value_heads = config.attention_other_setting[
+                "num_attention_groups"]
+
+        if self.sliding_window is not None and enable_sliding_window:
+            self.sliding_window = (self.sliding_window)
+        else:
+            self.sliding_window = None
+        self.head_dim = getattr(config, "head_dim",
+                        config.hidden_size // self.num_attention_heads)
+        self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
+
+        self.rotary_emb = Step3p5RotaryEmbedding(config, layer_idx=layer_idx)
+
+        self.q_size = self.num_attention_heads * self.head_dim
+        self.kv_size = self.num_key_value_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(config.hidden_size, self.q_size, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
+        self.o_proj = nn.Linear(self.q_size, config.hidden_size, bias=False)
+        self.q_norm = Step3p5RMSNorm(self.head_dim,
+                                    eps=config.rms_norm_eps)
+        self.k_norm = Step3p5RMSNorm(self.head_dim,
+                                    eps=config.rms_norm_eps)
+
+        self.use_head_wise_attn_gate = config.use_head_wise_attn_gate
+        if self.use_head_wise_attn_gate:
+            self.g_proj = nn.Linear(config.hidden_size,
+                                    self.num_attention_heads,
+                                    bias=False)
+
+        self.use_rope = True
+        use_rope_layers = getattr(config, "use_rope_layers", None)
+        if use_rope_layers:
+            self.use_rope = use_rope_layers[self.layer_idx]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: 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_norm(
+            self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(
+            self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(
+            1, 2)
+        if self.use_head_wise_attn_gate:
+            gate_states = self.g_proj(hidden_states)
+        cos, sin = self.rotary_emb(hidden_states, position_ids)
+
+        # cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin)
+
+        # query_states, key_states = apply_rotary_pos_emb(query_norm_states, key_norm_states, cos, sin)
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; position_ids needed for the static cache
+            cache_kwargs = {
+                "sin": sin,
+                "cos": cos,
+                "cache_position": cache_position
+            }
+            key_states, value_states = past_key_value.update(
+                key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        # TODO: considering FP8；
+        # RuntimeError: Expected attn_mask dtype to be bool or float or to match query dtype,
+        # but got attn_mask.dtype: long int and  query.dtype: c10::BFloat16 instead.
+        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,  # main diff with Llama
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(*input_shape, -1)
+        if self.use_head_wise_attn_gate:
+            output = attn_output.view(
+                *attn_output.shape[:-1], self.num_attention_heads,
+                self.head_dim) * gate_states.unsqueeze(-1).sigmoid()
+            attn_output = output.view(*attn_output.shape)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class Step3p5DecoderLayer(GradientCheckpointingLayer):
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.layer_idx = layer_idx
+        self.self_attn = Step3p5Attention(config, layer_idx)
+        self.attention_type = config.layer_types[layer_idx]
+
+        moe_layers_enum = getattr(config, "moe_layers_enum", None)
+        if moe_layers_enum is not None:
+            moe_layers_idx = [
+                int(i) for i in moe_layers_enum.strip().split(',')
+            ]
+        else:
+            moe_layers_idx = [i for i in range(1, config.num_hidden_layers)]
+        self.is_moe_layer = layer_idx in moe_layers_idx
+        self.use_moe = False
+
+        if config.swiglu_limits_shared and config.swiglu_limits_shared[
+                layer_idx] is not None and config.swiglu_limits_shared[
+                    layer_idx] != 0:
+            swiglu_limit_shared = config.swiglu_limits_shared[layer_idx]
+        else:
+            swiglu_limit_shared = None
+        if config.swiglu_limits and config.swiglu_limits[
+                layer_idx] is not None and config.swiglu_limits[layer_idx] != 0:
+            swiglu_limit = config.swiglu_limits[layer_idx]
+        else:
+            swiglu_limit = None
+        if self.is_moe_layer:
+            self.moe = Step3p5MoEMLP(config, swiglu_limit=swiglu_limit)  #
+            self.share_expert = Step3p5MLP(
+                config,
+                intermediate_size=config.share_expert_dim,
+                swiglu_limit=swiglu_limit_shared)
+            self.use_moe = True
+        else:
+            self.mlp = Step3p5MLP(config,
+                                 intermediate_size=config.intermediate_size,
+                                 swiglu_limit=swiglu_limit_shared)
+
+        self.input_layernorm = Step3p5RMSNorm(
+            config.hidden_size,
+            eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Step3p5RMSNorm(
+            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_value: Optional[tuple[torch.Tensor]] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+        
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        if self.use_moe:
+            share_output = self.share_expert(hidden_states)
+            moe_output = self.moe(hidden_states)
+            ffn_output = moe_output + share_output
+        else:
+            ffn_output = self.mlp(hidden_states)
+        if isinstance(ffn_output, tuple):
+            hidden_states, _ = ffn_output
+        else:
+            hidden_states = ffn_output
+
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class Step3p5PreTrainedModel(PreTrainedModel):
+    # Link this model family to its configuration class so PreTrainedModel.from_pretrained
+    # can load the config instead of failing with a NoneType error.
+    config_class = Step3p5Config
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = ["past_key_values"]
+    _keys_to_ignore_on_load_unexpected = [
+        r"model\.layers\.45\.*",
+        r"model\.layers\.46\.*",
+        r"model\.layers\.47\.*"
+    ]
+    _supports_flash_attn = False
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+
+class Step3p5Model(Step3p5PreTrainedModel, GenerationMixin):
+    _no_split_modules = ["Step3p5DecoderLayer"]
+    base_model_prefix = "model"
+    _tied_weights_keys = ["lm_head.weight"]
+    config: Step3p5Config
+    def __init__(self, config: Step3p5Config):
+        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([
+            Step3p5DecoderLayer(config, layer_idx)
+            for layer_idx in range(config.num_hidden_layers)
+        ])
+        self.norm = Step3p5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self, input_ids):
+        return self.embed_tokens(input_ids)
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: 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,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        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()
+
+        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)
+
+        hidden_states = inputs_embeds
+
+        # 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,
+            }
+            # Create the masks
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+            }
+
+            # The sliding window alternating layers are not always activated depending on the config
+            if self.has_sliding_layers:
+                causal_mask_mapping[
+                    "sliding_attention"] = create_sliding_window_causal_mask(
+                        **mask_kwargs)
+
+        # # create position embeddings to be shared across the decoder layers
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        for decoder_layer in self.layers[:self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states, )
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[
+                    decoder_layer.attention_type],
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+            hidden_states = layer_outputs
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class Step3p5ForCausalLM(Step3p5PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    config: Step3p5Config
+
+    def __init__(self, config: Step3p5Config):
+        super().__init__(config)
+        self.model = Step3p5Model(config)
+        self.lm_head = nn.Linear(config.hidden_size,
+                                 config.vocab_size,
+                                 bias=False)
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.model.get_decoder()
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        num_patches=None,
+        patch_pixel_values=None,
+        patch_newline_mask=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,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, Step3p5CausalLMOutputWithPast]:
+        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."
+        ```"""
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        # breakpoint()
+        outputs = self.model(
+            input_ids=input_ids,
+            num_patches=num_patches,
+            patch_pixel_values=patch_pixel_values,
+            patch_newline_mask=patch_newline_mask,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+
+        return Step3p5CausalLMOutputWithPast(logits=logits, )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        inputs_embeds=None,
+        pixel_values=None,
+        attention_mask=None,
+        cache_position=None,
+        logits_to_keep=None,
+        **kwargs,
+    ):
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            **kwargs,
+        )
+
+        if cache_position[0] == 0:
+            # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+            # Otherwise we need pixel values to be passed to model
+            model_inputs["pixel_values"] = pixel_values
+
+        return model_inputs
+
+    def _fix_state_dict_key_on_load(self, key: str) -> tuple[str, bool]:
+        if key.startswith("language_model."):
+            return key[len("language_model."):], True
+
+        return key, False

quantization_config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:257f6665a7d41fcbb7e6cad3c563ee3579e97bc145ba05ffae71bc59eeb51a7a
+size 44970995

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<｜begin▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}