Upload 10 files

inference/README.md

@@ -0,0 +1,388 @@
+# Step-5 minimal inference
+
+> 中文版本：[README.zh.md](README.zh.md)
+
+This directory contains the standalone Step-5 inference path. Full-model text
+and image-text generation have been validated with FP8 on **one 8 x H200 node**
+at **TP=EP=8**, and BF16 on **two 8 x H200 nodes** at **TP=EP=16**.
+The vision encoder uses PyTorch and BF16 weights. All commands below run from
+this directory; the inference files can be used independently.
+
+The default sliding-attention path is PyTorch SDPA. QKNorm/RoPE, DSA sparse
+attention, block-scaled FP8 expert GEMM, and the critical MoE primitives use
+Triton/CUDA; model orchestration and NCCL collectives remain PyTorch.
+
+The same ranks provide tensor parallelism for the dense backbone and expert
+parallelism for the 352 routed experts. TP=EP=8 places 44 experts on each rank;
+TP=EP=16 places 22. TP and EP share a process group; there is no independent
+`TP × EP` process grid.
+
+## Contents
+
+| File | What it is |
+|------|------------|
+| `convert.py` | Converts an original Step-5 checkpoint into per-rank TP shards: `tpN/model-r{rank}.safetensors` + matching `config-r{rank}.json`, plus a shared `tokenizer_files/`. |
+| `generate.py` | Greedy text or image-text inference. Loads one shard per `torchrun` rank; accepts a text `--prompt`, or exact prompt IDs with optional prepared `--image-inputs`. |
+| `model.py` | Text/VL configuration, decoder layers, TP/EP execution, and image-feature injection into packed prefill. Imports compute kernels from `kernel.py`. |
+| `vision_encoder.py` | PyTorch perception encoder with 2D RoPE, absolute positions and convolutional downsamplers. |
+| `kernel.py` | The Triton/CUDA compute kernels: the sparse-attention (DSA) indexer + decode/prefill, block-scaled FP8 expert GEMM, the MoE primitives, and fused QKNorm/RoPE. |
+| `test_vision.py` | Configuration, checkpoint-loading and multimodal inference regressions. |
+| `requirements.txt` | Compatible runtime ranges: torch `>=2.10,<2.14`, transformers `>=4.57.6,<6`, safetensors `>=0.7,<1`, and triton `>=3.6,<3.8`. |
+
+Typical flow: `convert.py` (once, to shard a checkpoint) → `generate.py` (to run it).
+`generate.py` loads the model implemented by `model.py`, `kernel.py` and `vision_encoder.py`.
+
+Expert storage follows the source configuration: absent or empty
+`quantization_config` selects BF16; `quant_method="fp8"` selects FP8 with BF16
+exclusions from `modules_to_not_convert`. Conversion preserves weight dtypes.
+
+## Install
+
+Use a CUDA-enabled PyTorch build matching the host driver. FP8 checkpoint inference
+at TP=8 requires eight visible GPUs with FP8 E4M3 support. The validated vision runtime
+uses PyTorch 2.13.0+cu129, Triton 3.7.1, Transformers 5.15.0, and Safetensors 0.8.0.
+The compatible version ranges are listed in `requirements.txt`.
+
+Create a virtual environment and install the dependencies:
+
+```bash
+python3 -m venv --system-site-packages .venv
+source .venv/bin/activate
+python -m pip install -r requirements.txt
+```
+
+If the host does not already provide a compatible CUDA-enabled PyTorch, install
+it from the CUDA wheel index matching the driver before installing the remaining
+requirements:
+
+```bash
+python -m pip install \
+  --index-url "<matching-pytorch-cuda-wheel-index>" \
+  "torch>=2.10,<2.14"
+python -m pip install -r requirements.txt
+```
+
+The release loads the fast tokenizer serialized in local `tokenizer.json` directly.
+Automatic tokenizer class selection in Transformers 5.15 is incompatible with these
+Step-5 assets and can collapse long prompts. Direct loading preserves the
+byte-level BPE input IDs and output decoding validated with Transformers 4.57.6.
+
+## Long-output stability
+
+Generation uses the following long-output execution contract:
+
+- TP partial results are added in rank order with FP32 accumulation.
+- Sliding-window token-by-token decode uses PyTorch's default SDPA backend
+  selection. This keeps the production path on the backend supported by the
+  installed PyTorch/CUDA runtime instead of forcing the slower math backend.
+- A request that emits EOS is removed from later decode steps and its cache row
+  is compacted.
+
+The checks cover row-position stability inside one active batch. They do not
+promise bitwise-identical output across different batch sizes, schedulers, GPU
+models, or software versions.
+
+## Checkpoints and launch modes
+
+Run all commands from this directory. Choose the FP8 or BF16 section for your
+source checkpoint and use a separate output directory for each format.
+Conversion partitions weights and preserves their dtype; it does not quantize
+a BF16 checkpoint into FP8.
+
+The source model directory supplies the authoritative `config.json`.
+`convert.py` writes `tpN/model-r{rank}.safetensors` and matching
+`config-r{rank}.json` files, with `gqa-provider-shared-tp-v2` and
+`expert_parallel_size=N`, plus shared tokenizer assets in `tokenizer_files/`.
+`generate.py` reads these configs from `--tp-dir` and discovers the copied
+tokenizer automatically. The inference directory does not supply a model config.
+Use the weights and configs from the same conversion. If a config omits
+`expert_parallel_size`, the runtime uses `tp_size`.
+
+<a id="fp8"></a>
+
+## FP8: one node with 8 GPUs
+
+The FP8 path has been validated on **8 x H200**, with **TP=EP=8** and 44 routed
+experts per rank. The source config must declare `quant_method="fp8"`;
+excluded expert layers and the vision tower retain their BF16 weights.
+
+Set the source and output directories:
+
+```bash
+export STEP5_FP8_CHECKPOINT="<fp8-checkpoint-directory>"
+export STEP5_FP8_OUT="<fp8-shard-output-directory>"
+```
+
+Check checkpoint keys and shard shapes, then convert once to a fresh output
+directory. `--check-only` reads headers without loading weights or writing shards:
+
+```bash
+python3 convert.py \
+  --checkpoint "$STEP5_FP8_CHECKPOINT" \
+  --tp-size 8 --ep-size 8 --check-only
+
+python3 convert.py \
+  --checkpoint "$STEP5_FP8_CHECKPOINT" \
+  --out-dir "$STEP5_FP8_OUT" \
+  --tp-size 8 --ep-size 8
+```
+
+Launch text generation on the node with eight visible GPUs:
+
+```bash
+torchrun --standalone --nproc-per-node=8 generate.py \
+  --tp-dir "$STEP5_FP8_OUT/tp8" --ep-size 8 \
+  --prompt "请介绍一下张量并行。" \
+  --max-new-tokens 128
+```
+
+For image-text generation, prepare the matching checkpoint's prompt IDs and
+image tensors according to [Vision inputs](#vision-inputs), then run:
+
+```bash
+torchrun --standalone --nproc-per-node=8 generate.py \
+  --tp-dir "$STEP5_FP8_OUT/tp8" --ep-size 8 \
+  --prompt-json prompt-ids.json --image-inputs image-inputs.safetensors \
+  --max-new-tokens 128
+```
+
+<a id="bf16"></a>
+
+## BF16: 16 GPUs across two nodes
+
+The full 92-layer BF16 VL model needs approximately **144.41 GiB per rank** at
+TP/EP=8 for parameters alone, exceeding the validated H200 node's 140.40 GiB per
+GPU. At **TP=EP=16**, loaded parameters occupy **74.93 GiB per rank**, with 22 routed
+experts per rank. These figures include replicated vision weights and a sharded
+output head, exclude MTP, and leave KV caches, RoPE, activations and communication
+as additional memory requirements.
+
+The complete BF16 decoder, vision encoder and projector passed two-node text
+and single-image greedy generation on **16 H200 GPUs**. All ranks produced
+identical tokens and reached EOS, with finite logits throughout. The measured
+peak PyTorch allocation was **76.37 GiB per GPU** for batch size 1 and a
+943-token image prompt, including 898 visual tokens. This excludes allocations
+outside PyTorch, such as NCCL and the CUDA context. Communication used **NCCL TCP
+Socket**; RDMA, long-context workloads and throughput have not been validated.
+
+The full-model run streamed and sliced the original weights using the converter's
+TP rules and executed the released model forward path. The complete BF16
+`convert.py` → `generate.py --tp-dir` file-based workflow below has not been
+separately run end to end; checkpoint key/shape checks and small-model BF16
+conversion/loading/generation tests have passed.
+
+Use two nodes with eight visible H200 GPUs each. Both nodes need the same code
+and compatible runtime, NCCL connectivity, and access to the same converted
+weights and tokenizer files at the configured paths. Image inputs must also be
+available with identical contents on both nodes.
+
+Set these paths on both nodes. The source checkpoint must use BF16 weights and
+have no quantization config, or an empty one:
+
+```bash
+export STEP5_BF16_CHECKPOINT="<bf16-checkpoint-directory>"
+export STEP5_BF16_OUT="<bf16-shard-output-directory>"
+```
+
+On **one node only**, check the source and convert once into 16 complete shards.
+Make the resulting `tp16/` and `tokenizer_files/` directories accessible to both
+nodes before launching:
+
+```bash
+python3 convert.py \
+  --checkpoint "$STEP5_BF16_CHECKPOINT" \
+  --tp-size 16 --ep-size 16 --check-only
+
+python3 convert.py \
+  --checkpoint "$STEP5_BF16_CHECKPOINT" \
+  --out-dir "$STEP5_BF16_OUT" \
+  --tp-size 16 --ep-size 16
+```
+
+On **both nodes**, set the same reachable node-0 address and free rendezvous port:
+
+```bash
+export MASTER_ADDR="<node-0-reachable-IP-or-hostname>"
+export MASTER_PORT=29500
+```
+
+On **node 0**:
+
+```bash
+export NODE_RANK=0
+```
+
+On **node 1**:
+
+```bash
+export NODE_RANK=1
+```
+
+Run the following command on **both nodes**, with their respective `NODE_RANK`,
+for the same launch. It starts eight workers per node and 16 workers in total:
+
+```bash
+torchrun --nnodes=2 --nproc-per-node=8 \
+  --node-rank="$NODE_RANK" \
+  --master-addr="$MASTER_ADDR" --master-port="$MASTER_PORT" \
+  generate.py \
+  --tp-dir "$STEP5_BF16_OUT/tp16" --ep-size 16 \
+  --prompt "请介绍一下张量并行。" \
+  --max-new-tokens 128
+```
+
+Alternatively, for image-text generation, prepare the inputs according to
+[Vision inputs](#vision-inputs) and run this command on **both nodes**:
+
+```bash
+torchrun --nnodes=2 --nproc-per-node=8 \
+  --node-rank="$NODE_RANK" \
+  --master-addr="$MASTER_ADDR" --master-port="$MASTER_PORT" \
+  generate.py \
+  --tp-dir "$STEP5_BF16_OUT/tp16" --ep-size 16 \
+  --prompt-json prompt-ids.json --image-inputs image-inputs.safetensors \
+  --max-new-tokens 128
+```
+
+The launcher sets `WORLD_SIZE=2 x 8=16`; shard configs and `--ep-size` must both
+use 16. `--standalone` is only used by the single-node FP8 example.
+
+For exact text token IDs in either launch mode, replace `--prompt` and its value
+with `--prompt-json <prompt-json-file>`. The file accepts a JSON list or
+`{"prompt_ids": [...]}`. `--prompt-json-batch` accepts a batch of requests.
+
+## Limitations
+
+- Full-model generation has been validated with FP8 on one 8 x H200 node and
+  BF16 on two 8 x H200 nodes using streamed weights. The full BF16 file-based
+  conversion/loading workflow remains unvalidated end to end.
+- `--ep-size` defaults to 8 and currently must equal TP/world size; independent
+  TP and EP process groups such as `TP=4, EP=2` are not implemented.
+  `generate.py` reads TP from the shard config (and checks it against
+  `WORLD_SIZE`); it intentionally has no separate `--tp-size` switch.
+- Sliding attention uses PyTorch SDPA.
+- Generation is greedy. Continuous dynamic batching, paged KV cache, prefix
+  caching and CUDA graphs are not provided. BF16 cross-node validation covers
+  short text and single-image requests over NCCL TCP Socket.
+- MTP layers are not loaded, so speculative decoding is unavailable.
+- The checkpoint's factory FP8/BF16 expert layout is consumed directly; the
+  converter and generated TP shards are checkpoint-layout specific.
+
+## Vision checkpoints (`step5_vl`)
+
+`Step5ForConditionalGeneration` uses `Step5VLConfig` with top-level
+`model_type=step5_vl`, `text_config.model_type=step5`, and
+`vision_config.model_type=perception_encoder`. Text-only checkpoints use
+`Step5Config` and `Step5ForCausalLM`.
+
+The Step-5 perception encoder uses PyTorch operations, non-causal SDPA,
+2D RoPE, absolute position embeddings, LayerScale and two convolutional
+downsamplers. Projected image features enter the packed text decoder during
+prefill; text decoding uses the DSA cache and co-located TP/EP execution.
+
+The validated VL checkpoint has the following geometry:
+
+| Component | Configuration |
+| --- | --- |
+| Text decoder | 92 layers, hidden size 4096; 3 MTP layers excluded from ordinary generation |
+| Vision Transformer | 47 layers, width 1536, 16 attention heads, MLP size 8960 |
+| Patch embedding | RGB, 14×14 kernel and stride 14 |
+| Activation and normalization | QuickGELU, pre-LN, no post-LN, LayerScale 0.1 |
+| Position encoding | Absolute positions and 2D RoPE; no CLS token |
+| Downsamplers | Two 3×3 convolutions, stride 2, padding 1; 1536→3072→6144 |
+| Projector | 6144→4096, no bias |
+| Image placeholder | Token ID 128001 |
+| Vision weights | 667 BF16 tensors; approximately 3.96 GB per TP rank |
+
+Checkpoint keys are `model.*`, `lm_head.*`, `vision_model.*` and
+`vit_large_projector.*`. Vision and projector tensors are replicated on each
+TP rank. The checkpoint index determines the files and dtypes to load.
+The converter excludes MTP tensors at indices `>= text_config.num_hidden_layers`.
+
+Use the [FP8](#fp8) or [BF16](#bf16) checkpoint-check and conversion
+commands above for VL checkpoints as well. Complete shards include the nested
+config, `Step5ForConditionalGeneration` identity, text, vision and projector weights.
+
+Both resumed conversion and generation verify each shard's complete key set,
+shapes and TP/EP/rank metadata before reading weight data. Generation also rejects
+any missing or unexpected keys reported by the Hugging Face loader; incomplete
+vision weights cannot fall back to random initialization.
+
+<a id="vision-inputs"></a>
+
+### Vision inputs
+
+The model API accepts preprocessed tensors:
+
+| Argument | Meaning |
+| --- | --- |
+| `pixel_values` | Overview images, `[images, 3, H, W]`; additional leading batch dimensions are flattened. |
+| `patch_pixel_values` | Optional crops in image order, with the same layout. |
+| `num_patches` | One nonnegative crop count per overview image. Defaults to zero crops. |
+| `image_embeds` | Alternative already projected features, `[..., text_hidden_size]`; excludes all pixel/crop arguments. |
+
+The released geometry produces 169 tokens for each 728×728 overview and 81
+for each 504×504 crop. Each image contributes **its crops first, then its
+overview**. Concatenate images in prompt order across packed requests. The
+number of `image_token_id` placeholders must equal the total feature-token
+count; `<im_start>` / `<im_end>` remain ordinary text embeddings. Image
+encoding runs only during prefill. Decode uses the existing cache.
+
+Image resizing, normalization, crop selection and prompt construction must come
+from the matching checkpoint processor. This standalone path consumes those
+prepared results; it does not define a new PIL/JPEG processing policy. Save a
+prepared batch with `safetensors.torch.save_file`, for example:
+
+```python
+import torch
+from safetensors.torch import save_file
+
+save_file({
+    "pixel_values": pixel_values.contiguous(),
+    "patch_pixel_values": patch_pixel_values.contiguous(),
+    "num_patches": torch.tensor(num_patches, dtype=torch.int64),
+}, "image-inputs.safetensors")
+```
+
+Use the image-text launch command in the [FP8](#fp8) or [BF16](#bf16)
+section with `prompt-ids.json` (including image placeholders) and
+`image-inputs.safetensors`.
+
+`--prompt-json-batch` supports a mixture of text and image requests. Every rank
+reads the same prepared image file. The Python API accepts these arguments in
+`Step5ForConditionalGeneration.forward` and `generate_greedy` as well.
+
+Run the multimodal integration regression with the optional test dependency:
+
+```bash
+python -m pip install pytest
+python -m pytest -q test_vision.py
+```
+
+`vision_select_layer=-1` and language-frequency 2D RoPE are supported.
+Other layer-selection or frequency modes fail explicitly.
+
+### Validation results
+
+The following results were obtained on H200 with PyTorch 2.13.0+cu129,
+Triton 3.7.1 and Transformers 5.15.0:
+
+| Check | Result |
+| --- | --- |
+| Vision integration, checkpoint identity and BF16 loading tests | 20 passed |
+| Configuration, generation, serving, evaluation and cache regressions | 203 passed, 93 subtests passed |
+| GPU operator gates | All 10 passed |
+| Batch and NCCL TP stability | Passed |
+| FP8 checkpoint schema, TP8 | 2653 runtime tensors validated; 51 MTP tensors excluded |
+| BF16 checkpoint schema, TP8 and TP16 | 2398 runtime tensors validated; 51 MTP tensors excluded |
+| BF16 small-model generation | TP1/TP8 pixel and projected-embedding inputs produced identical output token IDs |
+| Full vision numerical validation | Overview and crop outputs matched exactly; maximum absolute error 0 |
+| TP8 mixed-precision fixture | Pixel and projected-embedding inputs produced identical output token IDs |
+| HF file-shard loading, complete real rank 0 | 85.4 GB shard loaded; real-image features `[898,4096]` matched exactly |
+| Full real FP8 checkpoint TP8 image generation, streamed weights | Correctly described a white Samoyed on grass; 27 generated tokens, normal EOS, identical output on all ranks |
+| Full real BF16 checkpoint, two-node TP16/EP16, streamed weights | Text and single-image generation passed; finite logits, normal EOS and identical tokens on all 16 ranks; 74.93 GiB parameters/rank, 76.37 GiB peak PyTorch allocation/GPU |
+
+The image case used one 728×728 overview and nine 504×504 crops, totaling
+898 visual tokens. Image encoding runs in PyTorch with fused GPU attention.
+These results cover functional and numerical correctness for the tested inputs;
+they do not constitute a multimodal benchmark or OCR/video quality evaluation.

inference/README.zh.md

@@ -0,0 +1,371 @@
+# Step-5 最小推理
+
+> English version: [README.md](README.md)
+
+本目录包含独立的 Step-5 推理路径。完整模型的文本和图文生成已分别通过
+FP8 **单节点 8 × H200、TP=EP=8** 与 BF16 **双节点共 16 × H200、TP=EP=16** 验证。
+视觉编码器使用 PyTorch 和 BF16 权重。以下命令均在本目录执行，推理文件可独立使用。
+
+默认的滑窗注意力走 PyTorch SDPA。QKNorm/RoPE、DSA 稀疏注意力、分块缩放 FP8 专家
+矩阵乘法,以及关键的 MoE 基础算子用 Triton/CUDA 实现;模型编排与 NCCL 集合通信仍是
+PyTorch。
+
+同一组进程（rank）同时承担稠密主干的张量并行和 352 个路由专家的专家并行。
+TP=EP=8 时每个 rank 放置 44 个专家，TP=EP=16 时放置 22 个。
+TP 与 EP 共用进程组，没有独立的 `TP × EP` 进程网格。
+
+常用术语:
+
+| 缩写 | 含义 |
+|------|------|
+| TP | Tensor Parallelism，张量并行，把同一层计算拆到多张卡上。 |
+| EP | Expert Parallelism，专家并行，把 MoE 专家分布到多张卡上。 |
+| MoE | Mixture of Experts，混合专家网络，每个 token 只选择部分专家参与计算。 |
+| FP8 | 8 位浮点数格式，用于降低模型权重和计算开销。 |
+| BF16 | Brain Floating Point 16，16 位浮点数格式。 |
+| SDPA | PyTorch 的缩放点积注意力实现。 |
+| DSA | 稀疏注意力实现，先选出最相关的历史区域，再计算注意力。 |
+| QKNorm | 对注意力中的 query 和 key 向量做归一化。 |
+| RoPE | Rotary Position Embedding，旋转位置编码，用于表示 token 的相对位置。 |
+| NCCL | NVIDIA 集合通信库，用于多张 GPU 之间的数据交换。 |
+| KV cache | 保存历史 token 的注意力 key/value，避免解码时重复计算。 |
+| MTP | Multi-Token Prediction，多 token 预测层，可用于辅助投机解码。 |
+| token | 模型处理文本时使用的基本单位。 |
+
+## 目录内容
+
+| 文件 | 是什么 |
+|------|--------|
+| `convert.py` | 把原始 Step-5 模型权重转换成按进程分片的 TP 权重:`tpN/model-r{rank}.safetensors` + 对应的 `config-r{rank}.json`,外加一份共享的 `tokenizer_files/`。 |
+| `generate.py` | 文本或图文贪心推理。每个 `torchrun` 进程加载一个分片；可传文本 `--prompt`，或准确 prompt IDs 与预处理后的 `--image-inputs`。 |
+| `model.py` | 文本/VL 配置、解码层、TP/EP 执行，以及 packed prefill 中的图像特征注入。计算算子从 `kernel.py` 导入。 |
+| `vision_encoder.py` | PyTorch 视觉编码器，包含 2D RoPE、绝对位置编码及卷积下采样。 |
+| `kernel.py` | Triton/CUDA 计算算子:稀疏注意力(DSA)索引器 + 解码/预填充、分块缩放 FP8 专家矩阵乘法、MoE 基础算子,以及融合的 QKNorm/RoPE。 |
+| `test_vision.py` | 配置、权重加载和多模态推理回归。 |
+| `requirements.txt` | 兼容版本范围：torch `>=2.10,<2.14`、transformers `>=4.57.6,<6`、safetensors `>=0.7,<1`、triton `>=3.6,<3.8`。 |
+
+典型流程:`convert.py`(一次性,把模型权重分片)→ `generate.py`(运行)。
+`generate.py` 加载由 `model.py`、`kernel.py` 和 `vision_encoder.py` 实现的模型。
+
+专家权重格式由源配置决定：缺少或为空的 `quantization_config` 使用 BF16；
+`quant_method="fp8"` 使用 FP8，并按 `modules_to_not_convert` 保留指定层的 BF16
+专家。转换过程保留权重 dtype。
+
+## 安装
+
+使用与主机驱动匹配、启用了 CUDA 的 PyTorch。FP8 模型的 TP=8 推理需要 8 张支持
+FP8 E4M3 的可见 GPU。已验证的视觉运行环境为 PyTorch 2.13.0+cu129、Triton 3.7.1、
+Transformers 5.15.0 和 Safetensors 0.8.0；兼容版本范围见 `requirements.txt`。
+
+创建虚拟环境并安装依赖：
+
+```bash
+python3 -m venv --system-site-packages .venv
+source .venv/bin/activate
+python -m pip install -r requirements.txt
+```
+
+如果主机尚未提供兼容的 CUDA 版 PyTorch，先从与驱动匹配的 CUDA 软件包地址安装，
+再安装其余依赖：
+
+```bash
+python -m pip install \
+  --index-url "<匹配的-pytorch-cuda-软件包地址>" \
+  "torch>=2.10,<2.14"
+python -m pip install -r requirements.txt
+```
+
+发布版直接读取本地 `tokenizer.json` 中保存的快速分词规则。Transformers 5.15
+若只按分词器类别信息自动选择实现，会错误地把长文本压缩成很少的 token。
+字节级 BPE 会先把文本表示成字节，再把常见的连续字节合并成文本片段；直接读取
+`tokenizer.json` 可以同时保留 Transformers 4.57.6 下验证过的输入 token 编号和
+正常的输出文本还原。
+
+## 长输出稳定性
+
+生成过程采用以下长输出计算约定：
+
+- 多卡张量并行的部分结果按 rank 顺序逐个相加,相加时使用 FP32。
+- 滑窗逐 token 解码使用 PyTorch 默认选择的 SDPA 注意力后端,也就是由当前
+  PyTorch 和 CUDA 环境选择其支持的实现,不再强制使用较慢的通用数学实现。
+- 请求生成 EOS（���束标记）后,会从后续解码批次中移除,对应缓存行也会同步整理。
+
+这些检查覆盖同一活动批次内不同行位置的稳定性,但不承诺不同批大小、调度方式、
+GPU 型号或软件版本之间逐位完全相同。
+
+## 模型权重与启动方式
+
+以下命令均在本目录执行。按源权重格式选择 FP8 或 BF16 小节，两种格式使用不同的
+输出目录。转换只切分权重并保留 dtype，不会把 BF16 权重量化成 FP8。
+
+以源模型目录中的 `config.json` 为准。`convert.py` 写出
+`tpN/model-r{rank}.safetensors` 及对应的 `config-r{rank}.json`，包含
+`gqa-provider-shared-tp-v2` 布局标记和 `expert_parallel_size=N`；分词器文件保存在
+共享的 `tokenizer_files/` 下。`generate.py` 从 `--tp-dir` 读取配置，并自动发现
+转换时复制的分词器。推理目录不附带模型配置；分片配置必须与同一次转换产生的
+权重配套使用。配置省略 `expert_parallel_size` 时，运行时使用 `tp_size`。
+
+<a id="fp8"></a>
+
+## FP8：单节点 8 卡
+
+FP8 路径已在 **8 × H200** 上验证，使用 **TP=EP=8**，每个 rank 放置 44 个路由
+专家。源配置需声明 `quant_method="fp8"`；排除量化的专家层和视觉编码器保留 BF16
+权重。
+
+设置源权重和输出目录：
+
+```bash
+export STEP5_FP8_CHECKPOINT="<fp8-checkpoint-directory>"
+export STEP5_FP8_OUT="<fp8-shard-output-directory>"
+```
+
+先检查权重名称与分片形状，再转换到新的输出目录，转换只需执行一次。
+`--check-only` 只读取文件头，不加载权重数据、不写分片：
+
+```bash
+python3 convert.py \
+  --checkpoint "$STEP5_FP8_CHECKPOINT" \
+  --tp-size 8 --ep-size 8 --check-only
+
+python3 convert.py \
+  --checkpoint "$STEP5_FP8_CHECKPOINT" \
+  --out-dir "$STEP5_FP8_OUT" \
+  --tp-size 8 --ep-size 8
+```
+
+在可见 8 张 GPU 的节点上启动文本生成：
+
+```bash
+torchrun --standalone --nproc-per-node=8 generate.py \
+  --tp-dir "$STEP5_FP8_OUT/tp8" --ep-size 8 \
+  --prompt "请介绍一下张量并行。" \
+  --max-new-tokens 128
+```
+
+图文生成先按[图像输入约定](#vision-inputs)，使用对应 checkpoint 的 processor
+准备 prompt IDs 和图像张量，再运行：
+
+```bash
+torchrun --standalone --nproc-per-node=8 generate.py \
+  --tp-dir "$STEP5_FP8_OUT/tp8" --ep-size 8 \
+  --prompt-json prompt-ids.json --image-inputs image-inputs.safetensors \
+  --max-new-tokens 128
+```
+
+<a id="bf16"></a>
+
+## BF16：双节点共 16 卡
+
+完整 92 层 BF16 图文模型在 TP/EP=8 下，仅参数就约需 **144.41 GiB/卡**，超过
+已验证 H200 节点的 **140.40 GiB/卡**。采用 **TP=EP=16** 时，实测加载参数约为
+**74.93 GiB/卡**，每个 rank 放置 22 个路由专家。参数占用包含每卡复制的视觉权重和
+已切分的输出层，不包含 MTP；KV 缓存、RoPE、激活和通信还需要额外显存。
+
+完整 BF16 解码器、视觉编码器与投影层已在**双节点 16 张 H200** 上通过文本与单图
+贪心生成验证。所有 rank 输出 token 一致、正常 EOS，全程 logits 均为有限值。
+batch size 为 1、图文 prompt 共 943 个 token（含 898 个视觉 token）时，
+实测 PyTorch 显存分配峰值为 **76.37 GiB/卡**；此数值不包含 NCCL、CUDA context
+等 PyTorch 之外的分配。通信使用 **NCCL TCP Socket**；尚未验证 RDMA、长上下文
+负载或吞吐性能。
+
+完整模型实测使用转换器的 TP 规则流式读取、切分原始权重，并执行发布版模型的
+forward。下方完整 BF16 `convert.py` → `generate.py --tp-dir` 文件分片流程尚未
+单独端到端执行；权重名称与形状检查、BF16 小模型转换/加载/生成测试已通过。
+
+使用两台各有 8 张可见 H200 的节点。两台机器需要相同代码和兼容环境、可用的 NCCL
+网络，以及在配置路径下访问同一套转换后权重和分词器文件的能力。
+图文生成时，两台机器还需能够读取内容相同的图像输入文件。
+
+在两台节点上设置以下路径。源权重应为 BF16，且没有量化配置或量化配置为空：
+
+```bash
+export STEP5_BF16_CHECKPOINT="<bf16-checkpoint-directory>"
+export STEP5_BF16_OUT="<bf16-shard-output-directory>"
+```
+
+**只在一台节点上执行一次**检查与转换，生成完整的 16 份分片。启动前确保两台节点
+都能访问生成的 `tp16/` 和 `tokenizer_files/` 目录：
+
+```bash
+python3 convert.py \
+  --checkpoint "$STEP5_BF16_CHECKPOINT" \
+  --tp-size 16 --ep-size 16 --check-only
+
+python3 convert.py \
+  --checkpoint "$STEP5_BF16_CHECKPOINT" \
+  --out-dir "$STEP5_BF16_OUT" \
+  --tp-size 16 --ep-size 16
+```
+
+在**两台节点上**设置相同、可访问的节点 0 地址，以及空闲的 rendezvous 端口：
+
+```bash
+export MASTER_ADDR="<node-0-reachable-IP-or-hostname>"
+export MASTER_PORT=29500
+```
+
+在**节点 0**上设置：
+
+```bash
+export NODE_RANK=0
+```
+
+在**节点 1**上设置：
+
+```bash
+export NODE_RANK=1
+```
+
+在**两台节点上分别执行**下面的文本生成命令，使用各自的 `NODE_RANK` 加入同一次
+启动。每台运行 8 个进程，总进程数为 16：
+
+```bash
+torchrun --nnodes=2 --nproc-per-node=8 \
+  --node-rank="$NODE_RANK" \
+  --master-addr="$MASTER_ADDR" --master-port="$MASTER_PORT" \
+  generate.py \
+  --tp-dir "$STEP5_BF16_OUT/tp16" --ep-size 16 \
+  --prompt "请介绍一下张量并行。" \
+  --max-new-tokens 128
+```
+
+图文生成则按[图像输入约定](#vision-inputs)准备数据，在**两台节点上分别执行**：
+
+```bash
+torchrun --nnodes=2 --nproc-per-node=8 \
+  --node-rank="$NODE_RANK" \
+  --master-addr="$MASTER_ADDR" --master-port="$MASTER_PORT" \
+  generate.py \
+  --tp-dir "$STEP5_BF16_OUT/tp16" --ep-size 16 \
+  --prompt-json prompt-ids.json --image-inputs image-inputs.safetensors \
+  --max-new-tokens 128
+```
+
+启动器设置 `WORLD_SIZE=2 × 8=16`，分片配置和 `--ep-size` 都必须为 16。
+`--standalone` 仅用于上面的 FP8 单节点示例。
+
+两种启动方式如需传入精确文本 token IDs，可将 `--prompt` 及其值替换为
+`--prompt-json <prompt-json-file>`。文件接受 JSON 列表或
+`{"prompt_ids": [...]}`；批量请求使用 `--prompt-json-batch`。
+
+## 限制
+
+- 完整模型生成已验证 FP8 单节点 8 × H200 和 BF16 双节点共 16 × H200，使用流式
+  权重加载；完整 BF16 文件分片转换/加载流程尚未端到端验证。
+- `--ep-size` 默认 8,且当前必须等于 TP / 进程总数;独立的 TP 与 EP 进程组(例如
+  `TP=4, EP=2`)未实现。`generate.py` 从分片配置读取 TP(并对进程总数 `WORLD_SIZE` 做校验),
+  刻意没有单独的 `--tp-size` 开关。
+- 滑窗注意力使用 PyTorch SDPA。
+- 只支持贪心生成。连续动态批处理、分页 KV 缓存、前缀缓存和 CUDA Graph
+  均未提供。BF16 跨节点验证覆盖 NCCL TCP Socket 下的短文本与单图请求。
+- MTP 层不加载,因此没有投机解码。
+- 直接使用原始模型权重中的 FP8/BF16 专家布局;转换器与生成的 TP 分片和该权重布局绑定。
+
+## 多模态 checkpoint（`step5_vl`）
+
+`Step5ForConditionalGeneration` 使用 `Step5VLConfig`，顶层类型为
+`model_type=step5_vl`，嵌套配置为 `text_config.model_type=step5` 和
+`vision_config.model_type=perception_encoder`。纯文本 checkpoint 使用
+`Step5Config` 和 `Step5ForCausalLM`。
+
+Step-5 视觉编码器使用 PyTorch 算子，包含非因果 SDPA、二维 RoPE、绝对位置编码、
+LayerScale 和两次卷积下采样。投影后的图像特征在 prefill 注入 packed 文本输入，
+文本解码使用 DSA cache 和共置 TP/EP 执行。
+
+已验证的 VL checkpoint 参数如下：
+
+| 组件 | 配置 |
+| --- | --- |
+| 文本解码器 | 92 层、hidden size 4096；普通生成不加载 3 层 MTP |
+| Vision Transformer | 47 层、width 1536、16 个注意力头、MLP size 8960 |
+| Patch embedding | RGB，14×14 卷积核，stride 14 |
+| 激活与归一化 | QuickGELU、pre-LN、无 post-LN、LayerScale 0.1 |
+| 位置编码 | 绝对位置编码和二维 RoPE；无 CLS token |
+| 下采样 | 两个 3×3 卷积，stride 2、padding 1；1536→3072→6144 |
+| Projector | 6144→4096，无 bias |
+| 图像占位符 | Token ID 128001 |
+| 视觉权重 | 667 个 BF16 张量；每个 TP rank 约 3.96 GB |
+
+Checkpoint 权重名为 `model.*`、`lm_head.*`、`vision_model.*` 和
+`vit_large_projector.*`。Vision 与 projector 在每个 TP rank 上完整复制，
+加载文件及 dtype 以 checkpoint 索引为准。转换器排除层编号
+`>= text_config.num_hidden_layers` 的 MTP 权重。
+
+VL checkpoint 同样使用上面的 [FP8](#fp8) 或 [BF16](#bf16) 检查与转换命令。
+完整分片包含嵌套配置、`Step5ForConditionalGeneration` 架构名，以及文本、
+vision 和 projector 权重。
+
+断点恢复和正式加载都会先核对分片的完整 key、shape 以及 TP/EP/rank 元数据。
+正式加载还会拒绝 Hugging Face 报告的缺失或异常 key，避免视觉权重缺失时
+被默认初始化后继续推理。
+
+<a id="vision-inputs"></a>
+
+### 图像输入约定
+
+图像接口如下：
+
+| 参数 | 含义 |
+| --- | --- |
+| `pixel_values` | 预处理后的主图 `[图片数, 3, H, W]`，额外的前导 batch 维会被展平。 |
+| `patch_pixel_values` | 可选裁剪图，按图片顺序排列。 |
+| `num_patches` | 每张主图对应的非负整数裁剪数；省略时全部为 0。 |
+| `image_embeds` | 已完成 projector 的 `[..., text_hidden_size]` 特征，与所有像素/裁剪参数互斥。 |
+
+728×728 主图输出 169 个 token，504×504 裁剪图输出 81 个 token。
+每张图片的顺序是**先全部裁剪图，再主图**；多请求按 packed prompt 顺序拼接。
+`image_token_id` 的占位数量必须与特征 token 总数相同，`<im_start>` / `<im_end>`
+仍使用文本 embedding。只在 prefill 编码图片，后续 decode 复用 cache。
+
+缩放、归一化、裁剪和 prompt 构造使用对应 checkpoint 的 processor；这里接收
+其预处理结果，没有另行定义 JPEG/PIL 预处理规则。可将结果保存为 safetensors：
+
+```python
+import torch
+from safetensors.torch import save_file
+
+save_file({
+    "pixel_values": pixel_values.contiguous(),
+    "patch_pixel_values": patch_pixel_values.contiguous(),
+    "num_patches": torch.tensor(num_patches, dtype=torch.int64),
+}, "image-inputs.safetensors")
+```
+
+准备好包含图像占位 token 的 `prompt-ids.json` 和 `image-inputs.safetensors` 后，
+使用 [FP8](#fp8) 或 [BF16](#bf16) 小节中对应的图文启动命令。
+
+`--prompt-json-batch` 可混合文本和图像请求，各 rank 读取相同的图像输入文件。
+Python 的 `forward` / `generate_greedy` 也接受这些图像参数。
+
+安装可选测试依赖并运行多模态集成回归：
+
+```bash
+python -m pip install pytest
+python -m pytest -q test_vision.py
+```
+
+支持 `vision_select_layer=-1` 和 `lang` 模式的二维 RoPE，其他层选择或频率模式会显式报错。
+
+### 验证结果
+
+以下结果来自 H200，运行环境为 PyTorch 2.13.0+cu129、Triton 3.7.1、Transformers 5.15.0：
+
+| 检查 | 结果 |
+| --- | --- |
+| 视觉集成、配置身份与 BF16 加载测试 | 20 passed |
+| 配置、生成、服务、评测和缓存回归 | 203 passed，93 subtests passed |
+| GPU 算子检查 | 10 项全部通过 |
+| 批次与 NCCL TP 稳定性 | 通过 |
+| FP8 checkpoint 结构，TP8 | 2653 个 runtime 张量通过，排除 51 个 MTP 张量 |
+| BF16 checkpoint 结构，TP8 与 TP16 | 2398 个 runtime 张量通过，排除 51 个 MTP 张量 |
+| BF16 小模型生成 | TP1/TP8 像素与投影后 embedding 输入生成相同 token IDs |
+| 完整视觉数值验证 | 主图和裁剪图输出精确匹配，最大绝对误差 0 |
+| TP8 混合精度测试模型 | 像素与投影后 embedding 输入生成相同 token IDs |
+| HF 文件分片加载，完整真实 rank 0 | 85.4 GB 分片加载通过，真实图像特征 `[898,4096]` 精确匹配 |
+| 完整真实 FP8 checkpoint TP8 图像生成，流式权重加载 | 正确描述草地上的白色萨摩耶犬；生成 27 个 token，正常 EOS，各 rank 输出相同 |
+| 完整真实 BF16 checkpoint，双节点 TP16/EP16，流式权重加载 | 文本与单图生成通过；logits 均为有限值、正常 EOS、16 个 rank 的 token 一致；参数 74.93 GiB/卡，PyTorch 分配峰值 76.37 GiB/卡 |
+
+单图用例包含一张 728×728 主图和九张 504×504 裁剪图，共 898 个视觉 token。
+视觉编码使用 PyTorch 和融合 GPU 注意力。这些结果覆盖所测输入的功能和数值正确性，
+不构成多模态基准或 OCR/视频质量评测。

inference/convert.py

@@ -0,0 +1,599 @@
+"""Convert the original Step-5 checkpoint into validated TP shards.
+
+The converter writes ``OUT/tpN/model-r{rank}.safetensors`` and matching
+``config-r{rank}.json`` files, plus one shared ``OUT/tokenizer_files`` directory.
+The only supported production layout is ``gqa-provider-shared-tp-v2``.
+
+Example::
+
+    python3 convert.py --checkpoint /path/to/Step-5 \
+        --out-dir /path/to/Step-5-TP --tp-size 8 --ep-size 8
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+from typing import Any, Callable, TYPE_CHECKING
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+if TYPE_CHECKING:
+    from model import Step5Config
+
+
+STEP5_TP_LAYOUT_VERSION = "gqa-provider-shared-tp-v2"
+
+
+# Each rule callable maps (full_tensor, layer_idx_or_None, tp_size, rank, config) to the
+# per-rank tensor. Rules are matched by attempting them in order; the first hit wins.
+Rule = Callable[[torch.Tensor, int | None, int, int, Any], torch.Tensor]
+
+
+def _slice_dim(dim: int) -> Rule:
+    """Column or Row: slice tensor along ``dim`` by rank's window, then clone."""
+
+    def apply(t, _layer, tp, rank, _config):
+        size = t.shape[dim]
+        assert size % tp == 0, f"{size} not divisible by tp={tp}"
+        per = size // tp
+        return t.narrow(dim, rank * per, per).clone()
+
+    return apply
+
+
+def _group_window(num_groups: int, tp: int, rank: int) -> tuple[int, int]:
+    """Return ``(first_group, count)`` consumed by rank's contiguous Q-head slice.
+
+    For TP below the group count, groups partition across ranks. For TP above it, one group
+    is duplicated across ``tp / num_groups`` adjacent ranks. The latter is the production
+    TP=8 case: ``_group_window(4, 8, r) == (r // 2, 1)``.
+    """
+    if not 0 <= rank < tp:
+        raise ValueError(f"rank {rank} outside tp_size={tp}")
+    if num_groups % tp == 0:
+        count = num_groups // tp
+        return rank * count, count
+    if tp % num_groups == 0:
+        replicas = tp // num_groups
+        return rank // replicas, 1
+    raise ValueError(
+        f"tp_size={tp} and group count {num_groups} are incompatible: "
+        "one must divide the other"
+    )
+
+
+def _slice_grouped_dim(dim: int, config_attr: str) -> Rule:
+    """Slice whole GQA/provider groups, duplicating a group when TP exceeds group count."""
+
+    def apply(t, _layer, tp, rank, config):
+        owner = config
+        for component in config_attr.split("."):
+            owner = getattr(owner, component)
+        num_groups = int(owner)
+        size = t.shape[dim]
+        if size % num_groups:
+            raise ValueError(
+                f"tensor dim {dim} size {size} is not divisible by {num_groups} groups"
+            )
+        group_width = size // num_groups
+        first, count = _group_window(num_groups, tp, rank)
+        return t.narrow(dim, first * group_width, count * group_width).clone()
+
+    return apply
+
+
+def _replicate(t, _layer, _tp, _rank, _config):
+    return t.clone()
+
+
+# First matching key substring selects the rule. Attention and dense/shared
+# experts use TP dimension slices; routed experts use EP slices of the expert
+# axis. Routers, normalization parameters and vision weights are replicated.
+_SHARD_TABLE: list[tuple[str, Rule]] = [
+    # Attention: q_proj/g_proj ColumnParallel (slice head dim, dim 0); o_proj RowParallel
+    # (slice head dim on dim 1 of [hidden, num_heads*head_dim] = the in dim). K/V follow
+    # the GQA group consumed by the local Q slice. At TP=8 each rank gets one 192-d head,
+    # duplicated across adjacent rank pairs rather than replicating all four heads.
+    (".self_attn.q_proj.weight", _slice_dim(0)),
+    (".self_attn.g_proj.weight", _slice_dim(0)),
+    (".self_attn.o_proj.weight", _slice_dim(1)),
+    (
+        ".self_attn.k_proj.weight",
+        _slice_grouped_dim(0, "num_attention_groups"),
+    ),
+    (
+        ".self_attn.v_proj.weight",
+        _slice_grouped_dim(0, "num_attention_groups"),
+    ),
+    # Dense MLP: gate/up ColumnParallel (slice intermediate, dim 0); down RowParallel
+    # (slice intermediate on dim 1 of [hidden, intermediate] = the in dim).
+    (".mlp.gate_proj.weight", _slice_dim(0)),
+    (".mlp.up_proj.weight", _slice_dim(0)),
+    (".mlp.down_proj.weight", _slice_dim(1)),
+    # MoE routed experts -- EP, not TP-of-inner-dim. Each rank owns n_experts // tp_size
+    # contiguous experts (full inner dim per expert). Slicing the experts dim on dim 0. Block
+    # scales follow their weight onto the experts dim. ``moe.gate_proj.weight`` matches
+    # ``model.layers.N.moe.gate_proj.weight``; the leading ``.`` excludes ``share_expert`` keys.
+    (".moe.gate_proj.weight", _slice_dim(0)),
+    (".moe.up_proj.weight", _slice_dim(0)),
+    (".moe.down_proj.weight", _slice_dim(0)),
+    (".moe.gate_proj.weight_scale_inv", _slice_dim(0)),
+    (".moe.up_proj.weight_scale_inv", _slice_dim(0)),
+    (".moe.down_proj.weight_scale_inv", _slice_dim(0)),
+    # ``moe.gate.weight`` (the router) and ``moe.router_bias`` are replicated so every rank
+    # selects over all experts and sees the same router scores.
+    (".moe.gate.weight", _replicate),
+    (".moe.router_bias", _replicate),
+    # Shared expert is ordinary TP, not a replicate: gate/up slice their output dim and
+    # down slices its input dim. The runtime deliberately leaves the down result local,
+    # combines it with the local EP-routed result in FP32, then all-reduces once.
+    (".share_expert.gate_proj.weight", _slice_dim(0)),
+    (".share_expert.up_proj.weight", _slice_dim(0)),
+    (".share_expert.down_proj.weight", _slice_dim(1)),
+    # The indexer's Q/W heads are grouped by DSA provider group. Provider groups align with
+    # GQA KV groups, so TP=8 gives every rank one group (four Q/W heads), duplicated across
+    # the same adjacent rank pair. K/Z are single-head MQA projections and remain replicated.
+    # These exact matches must precede the generic sparse-indexer replicate rule.
+    (
+        ".sparse_indexer_q.weight",
+        _slice_grouped_dim(0, "sparse_config.num_provider_groups"),
+    ),
+    (
+        ".sparse_indexer_w.weight",
+        _slice_grouped_dim(0, "sparse_config.num_provider_groups"),
+    ),
+    (".sparse_indexer_", _replicate),
+    (".ssmax_s", _replicate),
+    # Attention + indexer RMSNorms: each is a small per-head scale on the routed tensor
+    # (the residual stream is full, so the norm acts on the full vector). Replicated.
+    (".q_norm.weight", _replicate),
+    (".k_norm.weight", _replicate),
+    (".k_norm.bias", _replicate),
+    # Per-layer input/post-attention RMSNorms and the model's final RMSNorm: replicated.
+    (".input_layernorm.weight", _replicate),
+    (".post_attention_layernorm.weight", _replicate),
+    # Embedding, lm_head: replicated (vocab is split at runtime by VocabParallelLinear).
+    ("model.embed_tokens.weight", _replicate),
+    ("model.norm.weight", _replicate),
+    ("lm_head.weight", _replicate),
+]
+
+
+def _shard_for(key: str) -> Rule:
+    if key.startswith(("vision_model.", "vit_large_projector.")):
+        return _replicate
+    for substr, rule in _SHARD_TABLE:
+        if substr in key:
+            return rule
+    raise KeyError(f"no sharding rule for {key}")
+
+
+def _layer_idx_of(key: str) -> int | None:
+    """``model.layers.N.`` -- extract N if present and in-range, else None."""
+    marker = "model.layers."
+    if key.startswith(marker):
+        rest = key[len(marker):]
+        layer_str = rest.split(".", 1)[0]
+        try:
+            return int(layer_str)
+        except ValueError:
+            return None
+    return None
+
+
+def _order_for_inspection(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Put embed / lm_head / norm first, then layers, then any stragglers."""
+    ordered: dict[str, torch.Tensor] = {}
+    for key in ("model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"):
+        if key in state:
+            ordered[key] = state[key]
+    for key in sorted(state.keys()):
+        if key.startswith("model.layers."):
+            ordered[key] = state[key]
+    for key in sorted(state.keys()):
+        if key not in ordered:
+            ordered[key] = state[key]
+    return ordered
+
+
+def _load_config_for(checkpoint: str, tp_size: int) -> Step5Config:
+    from model import config_from_dict
+
+    with open(os.path.join(checkpoint, "config.json")) as handle:
+        raw = json.load(handle)
+    raw["tp_size"] = tp_size
+    raw["expert_parallel_size"] = tp_size
+    raw["tp_layout_version"] = STEP5_TP_LAYOUT_VERSION
+    return config_from_dict(raw)
+
+
+def _write_config(
+    existing_config_path: str,
+    out_path: str,
+    tp_size: int,
+    expert_parallel_size: int,
+) -> None:
+    with open(existing_config_path) as handle:
+        cfg = json.load(handle)
+    cfg["tp_size"] = tp_size
+    cfg["expert_parallel_size"] = expert_parallel_size
+    cfg["tp_layout_version"] = STEP5_TP_LAYOUT_VERSION
+    cfg.pop("use_triton_sliding_attention", None)
+    from model import Step5VLConfig, config_from_dict
+    config = config_from_dict(cfg)
+    cfg = config.to_dict()
+    if isinstance(config, Step5VLConfig):
+        cfg["architectures"] = ["Step5ForConditionalGeneration"]
+        cfg["auto_map"] = {
+            "AutoConfig": "model.Step5VLConfig",
+            "AutoModelForCausalLM": "model.Step5ForConditionalGeneration",
+            "AutoModelForImageTextToText": "model.Step5ForConditionalGeneration",
+        }
+    else:
+        cfg["architectures"] = ["Step5ForCausalLM"]
+        cfg["auto_map"] = {
+            "AutoConfig": "model.Step5Config",
+            "AutoModelForCausalLM": "model.Step5ForCausalLM",
+        }
+    # ``num_attention_heads`` stays at the full value -- the model derives the per-rank view via
+    # ``num_attention_heads // tp_size``. ``tp_layout_version`` deliberately makes shards
+    # produced by either the old replicated-KV/provider rule or the v1 replicated-shared
+    # rule fail closed instead of loading with plausible but wrong shapes/semantics.
+    with open(out_path, "w") as handle:
+        json.dump(cfg, handle, indent=2)
+        handle.write("\n")
+
+
+def _write_readme(
+    out_path: str,
+    tp_size: int,
+    expert_parallel_size: int,
+    checkpoint: str,
+    config: Step5Config,
+) -> None:
+    local_kv = config.num_kv_groups_per_rank
+    local_index_heads = config.sparse_indexer_num_heads_per_rank
+    text = f"""# step5 TP-{tp_size} shards
+
+This directory holds one `model-r{{0..{tp_size-1}}}.safetensors` and one
+`config-r{{0..{tp_size-1}}}.json` file per rank. The sharding matches the tensor-parallel convention in
+`model.py`. Layout marker: `{STEP5_TP_LAYOUT_VERSION}`. Shards without this
+marker use an obsolete replicated-KV/provider or replicated-shared layout and must not be
+resumed or loaded.
+
+The same {tp_size} ranks form the dense tensor-parallel group and the routed-expert
+placement group (`expert_parallel_size={expert_parallel_size}`). Independent TP and EP
+process-group axes are not implemented.
+
+| Tensor | Rule | Per-rank shape (tp={tp_size}) |
+|--------|------|-------------------------------|
+| `self_attn.q_proj.weight` | ColumnParallel, slice dim 0 by `num_heads / tp_size` | `[{64 // tp_size} * 192, 4096]` |
+| `self_attn.g_proj.weight` | ColumnParallel, slice dim 0 by `num_heads / tp_size` | `[{64 // tp_size}, 4096]` |
+| `self_attn.k_proj.weight` | GQA group slice; duplicate only for ranks consuming the same group | `[{local_kv} * 192, 4096]` |
+| `self_attn.v_proj.weight` | GQA group slice; duplicate only for ranks consuming the same group | `[{local_kv} * 192, 4096]` |
+| `self_attn.o_proj.weight` | RowParallel, slice dim 1 by `num_heads * head_dim / tp_size` | `[4096, {64 // tp_size} * 192]` |
+| `mlp.gate_proj.weight` | ColumnParallel, slice dim 0 by `intermediate / tp_size` | `[{13824 // tp_size}, 4096]` |
+| `mlp.up_proj.weight` | ColumnParallel, slice dim 0 by `intermediate / tp_size` | `[{13824 // tp_size}, 4096]` |
+| `mlp.down_proj.weight` | RowParallel, slice dim 1 by `intermediate / tp_size` | `[4096, {13824 // tp_size}]` |
+| `moe.gate.weight` | replicated (router, every rank selects over all experts) | `[352, 4096]` |
+| `moe.gate_proj.weight` | EP, slice dim 0 (experts) by `n_experts / tp_size` | `[{352 // tp_size}, 1536, 4096]` |
+| `moe.up_proj.weight` | EP, slice dim 0 (experts) by `n_experts / tp_size` | `[{352 // tp_size}, 1536, 4096]` |
+| `moe.down_proj.weight` | EP, slice dim 0 (experts) by `n_experts / tp_size` | `[{352 // tp_size}, 4096, 1536]` |
+| `moe.*.weight_scale_inv` | slice the experts dim (dim 0), inner block grid unchanged | `[{352 // tp_size}, ceil(out/128), ceil(in/128)]` |
+| `moe.router_bias` | replicated | `[352]` |
+| `share_expert.gate/up_proj.weight` | ColumnParallel, slice dim 0 by `share_expert_dim / tp_size` | `[{1536 // tp_size}, 4096]` |
+| `share_expert.down_proj.weight` | RowParallel, slice dim 1; reduce deferred until shared+routed FP32 combine | `[4096, {1536 // tp_size}]` |
+| `sparse_indexer_q.weight` | provider-group slice (4 heads/group) | `[{local_index_heads} * 256, 4096]` |
+| `sparse_indexer_w.weight` | provider-group slice (4 heads/group) | `[{local_index_heads}, 4096]` |
+| `sparse_indexer_k/z`, `ssmax_s`, `*_norm.weight`, `k_norm.bias` | replicated | full |
+| `model.embed_tokens.weight` | replicated | `[128896, 4096]` |
+| `lm_head.weight` | slices the vocab dim at runtime (VocabParallelLinear) | `[128896, 4096]` full per rank |
+| `model.norm.weight` | replicated | `[4096]` |
+
+Reference geometry: `num_attention_heads={64}`, `num_attention_groups={4}`, `head_dim={192}`,
+`hidden={4096}`, `intermediate={13824}`, `moe_intermediate={1536}`, `vocab={128896}`,
+`moe_num_experts={352}`, `topk={8}`. Produced from `{checkpoint}` by
+`inference/convert.py`.
+"""
+    with open(out_path, "w") as handle:
+        handle.write(text)
+
+
+_TOKENIZER_FILES = (
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+    "chat_template.jinja",
+    "vocab.json",
+    "merges.txt",
+    "added_tokens.json",
+    "preprocessor_config.json",
+    "processor_config.json",
+)
+
+
+def _build_rank_state(
+    checkpoint: str,
+    keys: list[str],
+    weight_map: dict[str, str],
+    tp_size: int,
+    rank: int,
+    config: Step5Config,
+) -> dict[str, torch.Tensor]:
+    state: dict[str, torch.Tensor] = {}
+    chunks_by_shard: dict[str, list[str]] = {}
+    for key in keys:
+        chunks_by_shard.setdefault(weight_map[key], []).append(key)
+    for shard in sorted(chunks_by_shard):
+        print(f"  rank {rank}: shard {shard}", flush=True)
+        with safe_open(os.path.join(checkpoint, shard), framework="pt", device="cpu") as handle:
+            for key in chunks_by_shard[shard]:
+                full = handle.get_tensor(key)
+                rule = _shard_for(key)
+                state[key] = rule(full, _layer_idx_of(key), tp_size, rank, config)
+    return state
+
+
+def validate_rank_shard(path, config, rank, *, expected=None) -> None:
+    """Reject incompatible or incomplete shards before HF can initialize missing weights.
+
+    Inspect headers only. In particular, text-only v2 shards must not load under
+    a VL config just because their parallel-layout metadata is still current.
+    """
+    from model import model_class_for_config
+
+    if expected is None:
+        with torch.device("meta"):
+            expected = model_class_for_config(config)(config).state_dict()
+    text_config = config.get_text_config()
+    required_metadata = {
+        "tp_size": str(text_config.tp_size),
+        "expert_parallel_size": str(getattr(text_config, "expert_parallel_size", text_config.tp_size)),
+        "rank": str(rank),
+    }
+    with safe_open(str(path), framework="pt", device="cpu") as handle:
+        metadata = handle.metadata() or {}
+        # Tensor layout is unchanged; accept either serialized spelling, but
+        # reject conflicting markers rather than silently choosing one.
+        layout_fields = [name for name in ("step5_tp_layout", "step4_tp_layout") if name in metadata]
+        if not layout_fields:
+            raise ValueError(f"{path}: missing step5_tp_layout metadata")
+        for field in layout_fields:
+            if metadata[field] != STEP5_TP_LAYOUT_VERSION:
+                raise ValueError(
+                    f"{path}: {field} metadata is {metadata[field]!r}, "
+                    f"expected {STEP5_TP_LAYOUT_VERSION!r}"
+                )
+        for field, value in required_metadata.items():
+            if metadata.get(field) != value:
+                raise ValueError(f"{path}: {field} metadata is {metadata.get(field)!r}, expected {value!r}")
+        keys = set(handle.keys())
+        missing, unexpected = set(expected) - keys, keys - set(expected)
+        if missing or unexpected:
+            raise ValueError(
+                f"{path}: shard key mismatch: missing={sorted(missing)[:8]}, "
+                f"unexpected={sorted(unexpected)[:8]}; regenerate complete TP shards"
+            )
+        for key, tensor in expected.items():
+            shape = tuple(handle.get_slice(key).get_shape())
+            if shape != tuple(tensor.shape):
+                raise ValueError(f"{path}: {key} shape {shape} != model {tuple(tensor.shape)}")
+
+
+def _validate_resume_shards(tp_dir, start_rank, *, config, expected) -> None:
+    """Validate retained ranks with the same contract as the runtime loader."""
+    for rank in range(start_rank):
+        path = os.path.join(tp_dir, f"model-r{rank}.safetensors")
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"cannot resume at rank {start_rank}: missing completed shard {path}"
+            )
+        validate_rank_shard(path, config, rank, expected=expected)
+
+
+def pre_shard(
+    checkpoint: str,
+    out_dir: str,
+    tp_size: int,
+    *,
+    expert_parallel_size: int | None = None,
+    dry_run: bool = False,
+    start_rank: int = 0,
+    check_only: bool = False,
+) -> int:
+    """Write ``tp_size`` per-rank safetensors (plus configs and a README) to ``out_dir``.
+
+    ``start_rank`` resumes a partial conversion after checking earlier ranks for
+    matching layout, topology and complete parameter keys/shapes.
+    """
+    expert_parallel_size = (
+        tp_size if expert_parallel_size is None else expert_parallel_size
+    )
+    if tp_size < 1 or expert_parallel_size < 1:
+        raise ValueError(
+            f"parallel sizes must be positive, got TP={tp_size}, "
+            f"EP={expert_parallel_size}"
+        )
+    if expert_parallel_size != tp_size:
+        raise ValueError(
+            "independent TP and EP groups are not implemented: "
+            f"expert_parallel_size={expert_parallel_size} must equal "
+            f"tp_size={tp_size}. The same ranks shard the dense backbone and "
+            "the routed experts."
+        )
+    if not 0 <= start_rank < tp_size:
+        raise ValueError(
+            f"start_rank must be in [0, {tp_size}), got {start_rank}"
+        )
+    index_path = os.path.join(checkpoint, "model.safetensors.index.json")
+    with open(index_path) as handle:
+        weight_map = json.load(handle)["weight_map"]
+
+    config = _load_config_for(checkpoint, tp_size)
+    text_config = config.get_text_config()
+    keys = sorted(k for k in weight_map if (
+        _layer_idx_of(k) is None or _layer_idx_of(k) < text_config.num_hidden_layers
+    ))
+    from model import model_class_for_config
+
+    with torch.device("meta"):
+        expected = model_class_for_config(config)(config).state_dict()
+    missing = set(expected) - set(keys)
+    unexpected = set(keys) - set(expected)
+    if missing or unexpected:
+        raise ValueError(
+            f"checkpoint key mismatch: missing={sorted(missing)[:8]}, "
+            f"unexpected={sorted(unexpected)[:8]}"
+        )
+    # Check every shape before reading weight data or creating outputs. This also
+    # detects vision geometry/configuration drift on the original checkpoint.
+    by_shard = {}
+    for key in keys:
+        by_shard.setdefault(weight_map[key], []).append(key)
+    for shard, shard_keys in sorted(by_shard.items()):
+        with safe_open(os.path.join(checkpoint, shard), framework="pt", device="cpu") as handle:
+            for key in shard_keys:
+                shape = handle.get_slice(key).get_shape()
+                full = torch.empty(shape, device="meta")
+                actual = _shard_for(key)(full, _layer_idx_of(key), tp_size, 0, text_config)
+                if actual.shape != expected[key].shape:
+                    raise ValueError(f"{key}: shard shape {tuple(actual.shape)} != model {tuple(expected[key].shape)}")
+    print(f"validated {len(keys)} tensors; excluded {len(weight_map) - len(keys)} MTP tensors", flush=True)
+    if check_only:
+        print("PASS (checkpoint keys and TP shapes; no weight data read or files written)")
+        return 0
+
+    tp_dir = os.path.join(out_dir, f"tp{tp_size}")
+    os.makedirs(tp_dir, exist_ok=True)
+    if start_rank:
+        _validate_resume_shards(
+            tp_dir,
+            start_rank,
+            config=config,
+            expected=expected,
+        )
+    tokenizer_dst = os.path.join(out_dir, "tokenizer_files")
+    if not dry_run and not os.path.isdir(tokenizer_dst):
+        os.makedirs(tokenizer_dst, exist_ok=True)
+        for name in _TOKENIZER_FILES:
+            src = os.path.join(checkpoint, name)
+            if os.path.exists(src):
+                shutil.copy(src, os.path.join(tokenizer_dst, name))
+
+    sample_keys = [
+        "model.embed_tokens.weight",
+        "model.norm.weight",
+        "lm_head.weight",
+        "model.layers.0.self_attn.q_proj.weight",
+        "model.layers.0.self_attn.k_proj.weight",
+        "model.layers.0.self_attn.o_proj.weight",
+        "model.layers.0.mlp.gate_proj.weight",
+        "model.layers.0.mlp.down_proj.weight",
+        "model.layers.10.moe.gate_proj.weight",
+        "model.layers.10.moe.gate_proj.weight_scale_inv",
+        "model.layers.10.moe.down_proj.weight",
+        "model.layers.10.moe.down_proj.weight_scale_inv",
+        "model.layers.3.self_attn.sparse_indexer_q.weight",
+        "model.layers.3.self_attn.sparse_indexer_w.weight",
+    ]
+    for rank in range(start_rank, tp_size):
+        print(f"=== rank {rank} ===", flush=True)
+        state = _build_rank_state(checkpoint, keys, weight_map, tp_size, rank, text_config)
+        if dry_run:
+            print(f"rank {rank}:")
+            for key in sample_keys:
+                if key in state:
+                    print(f"  {key}: {tuple(state[key].shape)} {state[key].dtype}")
+        else:
+            ordered = _order_for_inspection(state)
+            path = os.path.join(tp_dir, f"model-r{rank}.safetensors")
+            save_file(
+                ordered,
+                path,
+                metadata={
+                    "format": "pt",
+                    "format_version": "1",
+                    "step5_tp_layout": STEP5_TP_LAYOUT_VERSION,
+                    "tp_size": str(tp_size),
+                    "expert_parallel_size": str(expert_parallel_size),
+                    "rank": str(rank),
+                },
+            )
+            print(f"  wrote {path}  ({os.path.getsize(path) / 1e9:.2f} GB)")
+
+    if dry_run:
+        print(f"(dry_run: no files written, would produce tp{tp_size}/ in {out_dir})")
+        return 0
+
+    config_path = os.path.join(tp_dir, "config-r0.json")
+    if start_rank > 0:
+        print(f"resuming from rank {start_rank} (ranks 0..{start_rank - 1} skipped)", flush=True)
+
+    _write_config(
+        os.path.join(checkpoint, "config.json"),
+        config_path,
+        tp_size,
+        expert_parallel_size,
+    )
+    for rank in range(tp_size):
+        dst = os.path.join(tp_dir, f"config-r{rank}.json")
+        if dst != config_path:
+            shutil.copy(config_path, dst)
+    _write_readme(
+        os.path.join(tp_dir, "README.md"),
+        tp_size,
+        expert_parallel_size,
+        checkpoint,
+        text_config,
+    )
+    print(f"  wrote {len(_TOKENIZER_FILES)} tokenizer files to {tokenizer_dst}")
+    print(f"  wrote {tp_size} per-rank configs to {tp_dir}/config-r{{0..{tp_size-1}}}.json")
+    print("PASS")
+    return 0
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--out-dir", "--out_dir", dest="out_dir", default=None)
+    parser.add_argument("--tp-size", "--tp_size", dest="tp_size", type=int, default=8)
+    parser.add_argument(
+        "--ep-size",
+        "--expert-parallel-size",
+        dest="ep_size",
+        type=int,
+        default=8,
+        help="Routed-expert sharding degree (default: 8). The current "
+        "co-located topology requires EP size to equal TP size.",
+    )
+    parser.add_argument("--dry-run", "--dry_run", dest="dry_run", action="store_true")
+    parser.add_argument(
+        "--start-rank",
+        "--start_rank",
+        dest="start_rank",
+        type=int,
+        default=0,
+        help="Resume from this rank; ranks 0..start_rank-1 are assumed already on disk.",
+    )
+    parser.add_argument("--check-only", action="store_true", help="Validate all checkpoint keys and TP shapes without reading tensor data or writing shards")
+    args = parser.parse_args()
+    if not args.check_only and args.out_dir is None:
+        parser.error("--out-dir is required unless --check-only is set")
+    return pre_shard(
+        checkpoint=args.checkpoint,
+        out_dir=args.out_dir,
+        tp_size=args.tp_size,
+        expert_parallel_size=args.ep_size,
+        dry_run=args.dry_run,
+        start_rank=args.start_rank,
+        check_only=args.check_only,
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

inference/generate.py

@@ -0,0 +1,976 @@
+"""Greedy Step5 inference over pre-sharded tensor-parallel weights.
+
+This standalone entry point loads one ``model-r{rank}.safetensors`` file per
+``torchrun`` process and uses :meth:`Step5ForCausalLM.generate_greedy`, whose
+vocabulary-parallel selection all-gathers every rank's contiguous vocabulary
+window before taking argmax.
+
+Example::
+
+    torchrun --standalone --nproc-per-node=8 inference/generate.py \
+        --tp-dir /path/to/step5_tp8/tp8 \
+        --ep-size 8 \
+        --checkpoint /path/to/original_step5_checkpoint \
+        --prompt "请介绍一下张量并行。" \
+        --max-new-tokens 128
+
+``--prompt-json`` accepts either a JSON list of token ids or an object containing
+``prompt_ids``.  This bypasses tokenization entirely and is useful for exact
+parity runs. ``--prompt-json-batch`` accepts either a JSON list of token-id
+lists or ``{"prompts": [{"prompt_ids": [...]}, ...]}`` and sends the packed
+batch through one ``generate_greedy`` call. Multi-request output is
+``{"batch_size": N, "results": [...]}``; a one-request batch retains the
+original single-request JSON shape. The original checkpoint is only a tokenizer
+source here; model weights always come from ``--tp-dir``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+import sys
+from typing import Any, Sequence
+
+import torch
+import torch.distributed as dist
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--tp-dir",
+        required=True,
+        help="Directory containing config-rN.json and model-rN.safetensors.",
+    )
+    prompt_group = parser.add_mutually_exclusive_group(required=True)
+    prompt_group.add_argument(
+        "--prompt",
+        help="User prompt text. A one-turn chat template is applied by default.",
+    )
+    prompt_group.add_argument(
+        "--prompt-json",
+        help="JSON list of exact token ids, or an object containing `prompt_ids`.",
+    )
+    prompt_group.add_argument(
+        "--prompt-json-batch",
+        help="Packed-batch JSON: either a list of token-id lists or an object "
+        "containing `prompts`, whose entries contain `prompt_ids`.",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        "--tokenizer-path",
+        dest="tokenizer",
+        default=None,
+        help="Local tokenizer directory. Overrides --checkpoint.",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        "--ckpt-path",
+        dest="checkpoint",
+        default=None,
+        help="Original local Step5 checkpoint directory used as the tokenizer "
+        "source. TP model weights are still loaded only from --tp-dir.",
+    )
+    parser.add_argument(
+        "--image-inputs",
+        help="Local safetensors containing pixel_values, optional patch_pixel_values/num_patches, "
+             "or already projected image_embeds. Use with exact prompt IDs containing image placeholders.",
+    )
+    parser.add_argument(
+        "--raw-prompt",
+        action="store_true",
+        help="Tokenize --prompt as plain text instead of applying the chat template.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=128,
+        help="Maximum number of greedily decoded tokens (default: 128).",
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda",
+        help="Device type/string; torchrun normally uses cuda (default: cuda).",
+    )
+    parser.add_argument(
+        "--ep-size",
+        "--expert-parallel-size",
+        dest="ep_size",
+        type=int,
+        default=8,
+        help="Routed-expert sharding degree (default: 8). The current "
+        "co-located topology requires EP size to equal TP/world size.",
+    )
+    return parser
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    if args.max_new_tokens < 0:
+        parser.error("--max-new-tokens must be non-negative")
+    if args.ep_size < 1:
+        parser.error("--ep-size must be positive")
+    if args.raw_prompt and args.prompt is None:
+        parser.error("--raw-prompt is only valid with --prompt")
+    if args.image_inputs and args.prompt is not None:
+        parser.error("--image-inputs requires --prompt-json or --prompt-json-batch")
+    return args
+
+
+def validate_parallel_topology(
+    config: Any,
+    *,
+    world: int,
+    requested_ep_size: int,
+    rank: int = 0,
+) -> int:
+    """Validate and return the co-located expert-parallel size.
+
+    Older validated TP shards predate the explicit
+    ``expert_parallel_size`` metadata, so its only supported fallback is the
+    checkpoint TP size. Values read from JSON are checked without coercion:
+    accepting (for example) ``8.9`` as EP=8 would defeat the fail-closed
+    topology contract.
+    """
+    config_tp_size = getattr(config, "tp_size", None)
+    if (
+        isinstance(config_tp_size, bool)
+        or not isinstance(config_tp_size, int)
+        or config_tp_size < 1
+    ):
+        raise RuntimeError(
+            f"rank {rank}: checkpoint tp_size must be a positive integer, "
+            f"got {config_tp_size!r}"
+        )
+    if config_tp_size != world:
+        raise RuntimeError(
+            f"rank {rank}: checkpoint TP={config_tp_size}, "
+            f"torchrun world size={world}"
+        )
+
+    config_ep_size = getattr(config, "expert_parallel_size", config_tp_size)
+    if (
+        isinstance(config_ep_size, bool)
+        or not isinstance(config_ep_size, int)
+        or config_ep_size < 1
+    ):
+        raise RuntimeError(
+            "checkpoint expert_parallel_size must be a positive integer, "
+            f"got {config_ep_size!r}"
+        )
+    if config_ep_size != config_tp_size:
+        raise RuntimeError(
+            "independent TP and EP groups are not implemented: checkpoint "
+            f"EP={config_ep_size} must equal TP={config_tp_size}"
+        )
+    if requested_ep_size != config_ep_size:
+        raise RuntimeError(
+            f"--ep-size={requested_ep_size} does not match checkpoint/runtime "
+            f"EP={config_ep_size}. This release co-locates TP and EP on the "
+            "same ranks."
+        )
+    return config_ep_size
+
+
+def load_prompt_json(path: str | os.PathLike[str]) -> list[int]:
+    """Load and validate the exact token ids in a parity prompt artifact."""
+    payload = json.loads(Path(path).read_text(encoding="utf-8"))
+    values = payload.get("prompt_ids") if isinstance(payload, dict) else payload
+    if not isinstance(values, list):
+        raise ValueError(
+            f"prompt JSON {path} must be a list or an object containing a `prompt_ids` list"
+        )
+    if not values:
+        raise ValueError(f"prompt JSON {path} contains no token ids")
+
+    prompt_ids: list[int] = []
+    for index, value in enumerate(values):
+        # JSON bool is an int subclass in Python, but never a meaningful token id.
+        if isinstance(value, bool) or not isinstance(value, int):
+            raise ValueError(
+                f"prompt JSON {path} token {index} is not an integer: {value!r}"
+            )
+        if value < 0:
+            raise ValueError(
+                f"prompt JSON {path} token {index} is negative: {value}"
+            )
+        prompt_ids.append(value)
+    return prompt_ids
+
+
+def load_prompt_json_batch(path: str | os.PathLike[str]) -> list[list[int]]:
+    """Load exact token ids for a packed batch.
+
+    Accepted schemas are either ``[[1, 2], [3, 4, 5]]`` or
+    ``{"prompts": [{"prompt_ids": [1, 2]}, {"prompt_ids": [3, 4, 5]}]}``.
+    Empty batches and empty requests are rejected because neither can be
+    represented by the packed model/cache contract.
+    """
+    payload = json.loads(Path(path).read_text(encoding="utf-8"))
+    object_schema = isinstance(payload, dict)
+    prompts = payload.get("prompts") if object_schema else payload
+    if not isinstance(prompts, list):
+        raise ValueError(
+            f"prompt batch JSON {path} must be a list of lists or an object "
+            "containing a `prompts` list"
+        )
+    if not prompts:
+        raise ValueError(f"prompt batch JSON {path} contains no prompts")
+
+    batch: list[list[int]] = []
+    for prompt_index, prompt in enumerate(prompts):
+        if object_schema:
+            if not isinstance(prompt, dict):
+                raise ValueError(
+                    f"prompt batch JSON {path} prompt {prompt_index} must be "
+                    "an object containing `prompt_ids`"
+                )
+            values = prompt.get("prompt_ids")
+        else:
+            values = prompt
+        if not isinstance(values, list):
+            raise ValueError(
+                f"prompt batch JSON {path} prompt {prompt_index} must contain "
+                f"a token-id list"
+            )
+        if not values:
+            raise ValueError(
+                f"prompt batch JSON {path} prompt {prompt_index} contains no token ids"
+            )
+
+        prompt_ids: list[int] = []
+        for token_index, value in enumerate(values):
+            if isinstance(value, bool) or not isinstance(value, int):
+                raise ValueError(
+                    f"prompt batch JSON {path} prompt {prompt_index} token "
+                    f"{token_index} is not an integer: {value!r}"
+                )
+            if value < 0:
+                raise ValueError(
+                    f"prompt batch JSON {path} prompt {prompt_index} token "
+                    f"{token_index} is negative: {value}"
+                )
+            prompt_ids.append(value)
+        batch.append(prompt_ids)
+    return batch
+
+
+def resolve_tokenizer_source(
+    *,
+    tp_dir: str | os.PathLike[str],
+    tokenizer: str | None,
+    checkpoint: str | None,
+) -> str | None:
+    """Resolve explicit tokenizer options, then the pre-shard default layout.
+
+    ``convert.py`` writes weights under ``OUT/tp8`` and tokenizer assets under
+    ``OUT/tokenizer_files``. Explicit values are local directories and are
+    returned here before the loader validates their tokenizer files.
+    """
+    if tokenizer:
+        return tokenizer
+    if checkpoint:
+        return checkpoint
+
+    tp_path = Path(tp_dir).expanduser()
+    candidates = (
+        tp_path / "tokenizer_files",
+        tp_path.parent / "tokenizer_files",
+    )
+    for candidate in candidates:
+        if candidate.is_dir():
+            return str(candidate)
+    return None
+
+
+def tokenize_text(tokenizer: Any, text: str, *, raw_prompt: bool) -> list[int]:
+    """Tokenize one prompt using the checkpoint's chat template or raw text."""
+    if raw_prompt:
+        encoded = tokenizer(text, add_special_tokens=True)
+        values = encoded.input_ids
+    else:
+        values = tokenizer.apply_chat_template(
+            [{"role": "user", "content": text}],
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=False,
+        )
+    if hasattr(values, "input_ids"):
+        values = values.input_ids
+    elif isinstance(values, dict) and "input_ids" in values:
+        values = values["input_ids"]
+    if hasattr(values, "ids"):
+        values = values.ids
+    if isinstance(values, list) and len(values) == 1 and isinstance(values[0], list):
+        values = values[0]
+    if isinstance(values, torch.Tensor):
+        values = values.reshape(-1).tolist()
+    if not isinstance(values, list):
+        raise TypeError(
+            "tokenizer returned unsupported token IDs: "
+            f"{type(values).__name__}"
+        )
+    prompt_ids = [int(token) for token in values]
+    if not prompt_ids:
+        raise ValueError("tokenizer produced an empty prompt")
+    if any(token < 0 for token in prompt_ids):
+        raise ValueError("tokenizer produced a negative token id")
+    return prompt_ids
+
+
+def validate_prompt_ids(
+    prompt_ids: Sequence[int],
+    *,
+    vocab_size: int,
+    max_position_embeddings: int,
+    max_new_tokens: int,
+    image_token_id: int | None = None,
+) -> None:
+    if not prompt_ids:
+        raise ValueError("prompt must contain at least one token")
+    for index, token in enumerate(prompt_ids):
+        if token < 0 or (token >= vocab_size and token != image_token_id):
+            raise ValueError(
+                f"prompt token {index}={token} is outside vocabulary [0, {vocab_size})"
+            )
+    requested = len(prompt_ids) + max_new_tokens
+    if requested > max_position_embeddings:
+        raise ValueError(
+            f"prompt ({len(prompt_ids)} tokens) + completion ({max_new_tokens}) "
+            f"exceeds max_position_embeddings={max_position_embeddings}"
+        )
+
+
+def _load_rank_config(tp_dir: str, rank: int):
+    # Lazy imports keep parser/helper tests independent of transformers and Triton.
+    from model import config_from_dict
+
+    path = Path(tp_dir) / f"config-r{rank}.json"
+    if not path.is_file():
+        raise FileNotFoundError(f"rank {rank}: missing TP config {path}")
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    return config_from_dict(payload)
+
+
+def _load_tokenizer(source: str, *, require_chat_template: bool = False):
+    source_path = Path(source).expanduser()
+    if not source_path.is_dir():
+        raise FileNotFoundError(
+            f"tokenizer source must be a local directory: {source_path}"
+        )
+    required_files = (
+        "tokenizer.json",
+        "tokenizer_config.json",
+    )
+    missing = [
+        name for name in required_files if not (source_path / name).is_file()
+    ]
+    if missing:
+        raise FileNotFoundError(
+            f"tokenizer source is missing required file(s): "
+            f"{', '.join(missing)}"
+        )
+
+    from transformers import PreTrainedTokenizerFast
+
+    # Load Step-5's serialized fast tokenizer directly. Automatic class selection
+    # in Transformers 5.15 can collapse long byte-level BPE prompts. The generic
+    # fast loader preserves the tokenizer.json model, pre-tokenizer, and decoder
+    # and is ID-identical to the validated Transformers 4.57.6 loader.
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(
+        str(source_path),
+        trust_remote_code=True,
+        local_files_only=True,
+    )
+    if require_chat_template and not getattr(tokenizer, "chat_template", None):
+        raise ValueError(
+            "tokenizer has no chat template; define `chat_template` in "
+            "tokenizer_config.json or provide chat_template.jinja"
+        )
+    return tokenizer
+
+
+def _load_exact_id_tokenizer(
+    source: str | None,
+    *,
+    explicit_source: bool,
+):
+    """Optionally load the tokenizer used only to decode exact-ID completions."""
+    if source is None:
+        return None
+    try:
+        return _load_tokenizer(source, require_chat_template=False)
+    except FileNotFoundError:
+        if explicit_source:
+            raise
+        return None
+
+
+def resolve_generation_eos_token_ids(
+    config_eos_token_id: Any,
+    tokenizer: Any | None,
+) -> set[int]:
+    """Union checkpoint and tokenizer EOS ids for production stopping.
+
+    Step5's converted TP config can retain base-model EOS ids while the chat
+    tokenizer terminates assistant turns with a different ``<|im_end|>`` id.
+    Passing only the config values to ``generate_greedy`` would continue past
+    the tokenizer's real end-of-turn marker.  Keep both sets rather than
+    replacing either source.
+    """
+
+    result: set[int] = set()
+
+    def add(value: Any, *, label: str) -> None:
+        if value is None:
+            return
+        values = value if isinstance(value, (list, tuple, set)) else [value]
+        for token in values:
+            if isinstance(token, bool) or not isinstance(token, int) or token < 0:
+                raise ValueError(
+                    f"{label} must contain non-negative integer token ids, "
+                    f"got {token!r}"
+                )
+            result.add(int(token))
+
+    add(config_eos_token_id, label="config eos_token_id")
+    if tokenizer is not None:
+        add(
+            getattr(tokenizer, "eos_token_id", None),
+            label="tokenizer eos_token_id",
+        )
+    if not result:
+        raise ValueError("generation has no EOS token ids")
+    return result
+
+
+def _broadcast_error(error: str | None, *, device: torch.device) -> None:
+    """Make a rank-0 prompt preparation error visible before any peer blocks."""
+    payload: list[str | None] = [error]
+    kwargs: dict[str, Any] = {}
+    if device.type == "cuda":
+        kwargs["device"] = device
+    dist.broadcast_object_list(payload, src=0, **kwargs)
+    if payload[0] is not None:
+        raise RuntimeError(f"rank 0 failed to prepare the prompt: {payload[0]}")
+
+
+def broadcast_eos_token_ids(
+    eos_token_ids: set[int] | None,
+    *,
+    rank: int,
+    device: torch.device,
+) -> set[int]:
+    """Broadcast rank 0's effective config+tokenizer EOS set to every TP rank."""
+
+    if rank == 0:
+        if not eos_token_ids:
+            raise ValueError("rank 0 must provide at least one EOS token id")
+        values = sorted(eos_token_ids)
+        length = torch.tensor([len(values)], dtype=torch.int64, device=device)
+    else:
+        values = []
+        length = torch.zeros(1, dtype=torch.int64, device=device)
+    dist.broadcast(length, src=0)
+    count = int(length.item())
+    if count <= 0:
+        raise ValueError(f"broadcast EOS set has invalid length {count}")
+    if rank == 0:
+        tokens = torch.tensor(values, dtype=torch.int64, device=device)
+    else:
+        tokens = torch.empty(count, dtype=torch.int64, device=device)
+    dist.broadcast(tokens, src=0)
+    result = {int(token) for token in tokens.cpu().tolist()}
+    if len(result) != count or any(token < 0 for token in result):
+        raise ValueError(f"broadcast EOS set is invalid: {sorted(result)}")
+    return result
+
+
+def broadcast_prompt_ids(
+    prompt_ids: Sequence[int] | None,
+    *,
+    rank: int,
+    device: torch.device,
+) -> list[int]:
+    """Broadcast a variable-length rank-0 prompt with tensor collectives."""
+    if rank == 0:
+        if prompt_ids is None:
+            raise ValueError("rank 0 must provide prompt ids")
+        length = torch.tensor([len(prompt_ids)], dtype=torch.int64, device=device)
+    else:
+        length = torch.zeros(1, dtype=torch.int64, device=device)
+    dist.broadcast(length, src=0)
+    token_count = int(length.item())
+    if token_count <= 0:
+        raise ValueError(f"broadcast prompt has invalid length {token_count}")
+
+    if rank == 0:
+        tokens = torch.tensor(prompt_ids, dtype=torch.long, device=device)
+    else:
+        tokens = torch.empty(token_count, dtype=torch.long, device=device)
+    dist.broadcast(tokens, src=0)
+    return [int(token) for token in tokens.cpu().tolist()]
+
+
+def broadcast_prompt_batch(
+    prompt_batch: Sequence[Sequence[int]] | None,
+    *,
+    rank: int,
+    device: torch.device,
+) -> list[list[int]]:
+    """Broadcast rank 0's packed batch as lengths plus one flat token tensor."""
+    if rank == 0:
+        if prompt_batch is None:
+            raise ValueError("rank 0 must provide a prompt batch")
+        batch_size = len(prompt_batch)
+        batch_size_tensor = torch.tensor(
+            [batch_size], dtype=torch.int64, device=device
+        )
+    else:
+        batch_size_tensor = torch.zeros(1, dtype=torch.int64, device=device)
+    dist.broadcast(batch_size_tensor, src=0)
+    batch_size = int(batch_size_tensor.item())
+    if batch_size <= 0:
+        raise ValueError(f"broadcast prompt batch has invalid size {batch_size}")
+
+    if rank == 0:
+        assert prompt_batch is not None
+        lengths = torch.tensor(
+            [len(prompt) for prompt in prompt_batch],
+            dtype=torch.int64,
+            device=device,
+        )
+    else:
+        lengths = torch.empty(batch_size, dtype=torch.int64, device=device)
+    dist.broadcast(lengths, src=0)
+    if bool((lengths <= 0).any().item()):
+        raise ValueError(
+            f"broadcast prompt batch has invalid lengths {lengths.cpu().tolist()}"
+        )
+    token_count = int(lengths.sum().item())
+
+    if rank == 0:
+        assert prompt_batch is not None
+        flat_tokens = torch.tensor(
+            [token for prompt in prompt_batch for token in prompt],
+            dtype=torch.long,
+            device=device,
+        )
+        if flat_tokens.numel() != token_count:
+            raise AssertionError(
+                f"rank 0 prompt batch has {flat_tokens.numel()} flat tokens, "
+                f"but lengths sum to {token_count}"
+            )
+    else:
+        flat_tokens = torch.empty(token_count, dtype=torch.long, device=device)
+    dist.broadcast(flat_tokens, src=0)
+
+    host_lengths = [int(length) for length in lengths.cpu().tolist()]
+    host_tokens = [int(token) for token in flat_tokens.cpu().tolist()]
+    result: list[list[int]] = []
+    offset = 0
+    for length in host_lengths:
+        result.append(host_tokens[offset : offset + length])
+        offset += length
+    if offset != len(host_tokens):
+        raise AssertionError(
+            f"broadcast prompt batch consumed {offset} of {len(host_tokens)} tokens"
+        )
+    return result
+
+
+def _resolve_device(device_arg: str, local_rank: int) -> torch.device:
+    requested = torch.device(device_arg)
+    if requested.type == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA was requested but torch.cuda.is_available() is false")
+        # An explicit index remains useful for single-rank debugging.  The normal
+        # torchrun spelling is simply --device cuda, which maps LOCAL_RANK here.
+        device = (
+            torch.device("cuda", local_rank)
+            if requested.index is None
+            else requested
+        )
+        torch.cuda.set_device(device)
+        return device
+    return requested
+
+
+def _load_rank_model(tp_dir: str, rank: int, config: Any, device: torch.device):
+    from convert import validate_rank_shard
+    from model import model_class_for_config
+
+    shard = Path(tp_dir) / f"model-r{rank}.safetensors"
+    if not shard.is_file():
+        raise FileNotFoundError(f"rank {rank}: missing TP weight shard {shard}")
+    validate_rank_shard(shard, config, rank)
+    model, loading_info = model_class_for_config(config).from_pretrained(
+        str(shard), config=config, output_loading_info=True,
+    )
+    failures = {key: loading_info.get(key) for key in (
+        "missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs",
+    ) if loading_info.get(key)}
+    if failures:
+        raise RuntimeError(f"rank {rank}: incomplete checkpoint loading: {failures}")
+    # TP shards intentionally keep a full replicated lm_head on disk.  Transformers 4
+    # sliced it through ``_load_from_state_dict``; Transformers 5 injects tensors directly,
+    # so the module must be narrowed explicitly after strict loading.
+    model.lm_head.shard_(rank)
+    model = model.to(device).eval()
+    model.pack_qkv_()
+    model.model.build_rope(device)
+    return model
+
+
+def generate_completion(
+    model: Any,
+    prompt_ids: list[int],
+    *,
+    max_new_tokens: int,
+    device: torch.device,
+) -> list[int]:
+    """Run the model's TP-safe greedy path for one request.
+
+    ``Step5ForCausalLM.generate_greedy`` calls ``vocab_parallel_argmax`` at
+    every step.  It therefore reconstructs the full vocabulary and returns a
+    global token id, rather than treating a rank-local column as a token id.
+    """
+    completions = generate_completions(
+        model,
+        [prompt_ids],
+        max_new_tokens=max_new_tokens,
+        device=device,
+    )
+    return completions[0]
+
+
+def generate_completions(
+    model: Any,
+    prompt_batch: list[list[int]],
+    *,
+    max_new_tokens: int,
+    device: torch.device,
+    eos_token_ids: set[int] | None = None,
+    image_inputs: dict[str, Any] | None = None,
+) -> list[list[int]]:
+    """Run one packed batch through the model's TP-safe greedy path."""
+    if not prompt_batch:
+        raise ValueError("prompt batch must contain at least one request")
+    generation_error: BaseException | None = None
+    try:
+        generation_kwargs: dict[str, Any] = {
+            "max_new_tokens": max_new_tokens,
+            "device": device,
+        }
+        if eos_token_ids is not None:
+            generation_kwargs["eos_token_ids"] = eos_token_ids
+        if image_inputs:
+            generation_kwargs.update(image_inputs)
+        completions = model.generate_greedy(prompt_batch, **generation_kwargs)
+        if not isinstance(completions, list):
+            raise TypeError(
+                "generate_greedy must return a list of token-ID lists; "
+                f"got {type(completions).__name__}"
+            )
+        if len(completions) != len(prompt_batch):
+            raise RuntimeError(
+                f"expected {len(prompt_batch)} completions, got {len(completions)}"
+            )
+        for completion_index, completion in enumerate(completions):
+            if not isinstance(completion, list):
+                raise TypeError(
+                    f"completion {completion_index} must be a token-ID list; "
+                    f"got {type(completion).__name__}"
+                )
+            for token_index, token in enumerate(completion):
+                if isinstance(token, bool) or not isinstance(token, int):
+                    raise TypeError(
+                        f"completion {completion_index} token {token_index} "
+                        f"must be an integer; got {type(token).__name__}"
+                    )
+                if token < 0:
+                    raise ValueError(
+                        f"completion {completion_index} token {token_index} "
+                        f"must be non-negative; got {token}"
+                    )
+        return completions
+    except BaseException as exc:
+        generation_error = exc
+        raise
+    finally:
+        if getattr(device, "type", None) == "cuda":
+            try:
+                torch.cuda.empty_cache()
+            except BaseException as cleanup_error:
+                if generation_error is None:
+                    raise
+                print(
+                    "[generate] CUDA allocator cleanup failed after generation "
+                    f"error: {type(cleanup_error).__name__}: {cleanup_error}",
+                    file=sys.stderr,
+                    flush=True,
+                )
+
+
+def format_output(
+    *,
+    prompt_token_count: int,
+    completion_ids: Sequence[int],
+    tokenizer: Any | None,
+    eos_token_ids: Sequence[int] | None = None,
+) -> str:
+    completion = (
+        tokenizer.decode(list(completion_ids), skip_special_tokens=False)
+        if tokenizer is not None
+        else None
+    )
+    payload = {
+        "prompt_token_count": prompt_token_count,
+        "completion_token_ids": list(completion_ids),
+        "completion": completion,
+    }
+    if eos_token_ids is not None:
+        payload["eos_token_ids"] = sorted(int(token) for token in eos_token_ids)
+    return json.dumps(payload, ensure_ascii=False)
+
+
+def format_batch_output(
+    *,
+    prompt_batch: Sequence[Sequence[int]],
+    completion_batch: Sequence[Sequence[int]],
+    tokenizer: Any | None,
+    eos_token_ids: Sequence[int] | None = None,
+) -> str:
+    """Format one request exactly as before, or a multi-request result object."""
+    if len(prompt_batch) != len(completion_batch):
+        raise ValueError(
+            f"prompt/completion batch mismatch: {len(prompt_batch)} != "
+            f"{len(completion_batch)}"
+        )
+    if not prompt_batch:
+        raise ValueError("cannot format an empty batch")
+    if len(prompt_batch) == 1:
+        return format_output(
+            prompt_token_count=len(prompt_batch[0]),
+            completion_ids=completion_batch[0],
+            tokenizer=tokenizer,
+            eos_token_ids=eos_token_ids,
+        )
+
+    results = []
+    for prompt_ids, completion_ids in zip(prompt_batch, completion_batch):
+        completion_list = list(completion_ids)
+        completion = (
+            tokenizer.decode(completion_list, skip_special_tokens=False)
+            if tokenizer is not None
+            else None
+        )
+        results.append(
+            {
+                "prompt_token_count": len(prompt_ids),
+                "completion_token_ids": completion_list,
+                "completion": completion,
+            }
+        )
+    return json.dumps(
+        {
+            "batch_size": len(results),
+            **(
+                {"eos_token_ids": sorted(int(token) for token in eos_token_ids)}
+                if eos_token_ids is not None
+                else {}
+            ),
+            "results": results,
+        },
+        ensure_ascii=False,
+    )
+
+
+def load_image_inputs(path, config):
+    """Load the same preprocessed inputs on every rank; propagate failures together."""
+    inputs, error = {}, None
+    if path is not None:
+        try:
+            from model import Step5VLConfig
+            from safetensors.torch import load_file
+            if not isinstance(config, Step5VLConfig):
+                raise ValueError("--image-inputs requires a step5_vl checkpoint")
+            inputs = load_file(path, device="cpu")
+            allowed = {"pixel_values", "patch_pixel_values", "num_patches", "image_embeds"}
+            if not inputs or set(inputs) - allowed:
+                raise ValueError(f"Image input keys must be a nonempty subset of {sorted(allowed)}")
+            if ("pixel_values" in inputs) == ("image_embeds" in inputs):
+                raise ValueError("Provide exactly one of pixel_values or image_embeds")
+            if "image_embeds" in inputs and len(inputs) != 1:
+                raise ValueError("image_embeds cannot be combined with crop inputs")
+        except Exception as exc:
+            error = f"{type(exc).__name__}: {exc}"
+    errors = [error]
+    if dist.is_initialized():
+        errors = [None] * dist.get_world_size()
+        dist.all_gather_object(errors, error)
+    if any(errors):
+        raise ValueError(f"Unable to load image inputs: {next(e for e in errors if e)}")
+    return inputs
+
+
+def _run(args: argparse.Namespace, *, rank: int, world: int, device: torch.device) -> int:
+    model_config = _load_rank_config(args.tp_dir, rank)
+    config = model_config.get_text_config()
+    image_token_id = getattr(model_config, "image_token_id", None) if args.image_inputs else None
+    config_ep_size = validate_parallel_topology(
+        config,
+        world=world,
+        requested_ep_size=args.ep_size,
+        rank=rank,
+    )
+
+    tokenizer_source = resolve_tokenizer_source(
+        tp_dir=args.tp_dir,
+        tokenizer=args.tokenizer,
+        checkpoint=args.checkpoint,
+    )
+    tokenizer = None
+    explicit_tokenizer_source = (
+        args.tokenizer is not None or args.checkpoint is not None
+    )
+    rank0_prompt_batch: list[list[int]] | None = None
+    rank0_eos_token_ids: set[int] | None = None
+    prompt_error: str | None = None
+    if rank == 0:
+        try:
+            if args.prompt_json_batch is not None:
+                rank0_prompt_batch = load_prompt_json_batch(args.prompt_json_batch)
+                # A tokenizer is optional for exact-id input; it is used only to
+                # decode completions when one is available.
+                tokenizer = _load_exact_id_tokenizer(
+                    tokenizer_source,
+                    explicit_source=explicit_tokenizer_source,
+                )
+            elif args.prompt_json is not None:
+                rank0_prompt_batch = [load_prompt_json(args.prompt_json)]
+                # A tokenizer is optional for exact-id input; it is used only to
+                # decode the completion when one was explicitly or automatically found.
+                tokenizer = _load_exact_id_tokenizer(
+                    tokenizer_source,
+                    explicit_source=explicit_tokenizer_source,
+                )
+            else:
+                if tokenizer_source is None:
+                    raise ValueError(
+                        "--prompt requires --tokenizer/--checkpoint, or tokenizer files "
+                        "at <tp-dir>/../tokenizer_files"
+                    )
+                tokenizer = _load_tokenizer(
+                    tokenizer_source,
+                    require_chat_template=not args.raw_prompt,
+                )
+                rank0_prompt_batch = [
+                    tokenize_text(tokenizer, args.prompt, raw_prompt=args.raw_prompt)
+                ]
+            assert rank0_prompt_batch is not None
+            for prompt_ids in rank0_prompt_batch:
+                validate_prompt_ids(
+                    prompt_ids,
+                    vocab_size=config.vocab_size,
+                    max_position_embeddings=config.max_position_embeddings,
+                    max_new_tokens=args.max_new_tokens,
+                    image_token_id=image_token_id,
+                )
+            rank0_eos_token_ids = resolve_generation_eos_token_ids(
+                config.eos_token_id,
+                tokenizer,
+            )
+        except Exception as exc:
+            prompt_error = f"{type(exc).__name__}: {exc}"
+
+    _broadcast_error(prompt_error, device=device)
+    eos_token_ids = broadcast_eos_token_ids(
+        rank0_eos_token_ids,
+        rank=rank,
+        device=device,
+    )
+    prompt_batch = broadcast_prompt_batch(
+        rank0_prompt_batch, rank=rank, device=device
+    )
+    # Check again on all ranks so a corrupted/mismatched rank config fails before
+    # loading hundreds of gigabytes of weights.
+    for prompt_ids in prompt_batch:
+        validate_prompt_ids(
+            prompt_ids,
+            vocab_size=config.vocab_size,
+            max_position_embeddings=config.max_position_embeddings,
+            max_new_tokens=args.max_new_tokens,
+            image_token_id=image_token_id,
+        )
+
+    if rank == 0:
+        prompt_lengths = [len(prompt_ids) for prompt_ids in prompt_batch]
+        prompt_description = (
+            f"{prompt_lengths[0]} prompt tokens"
+            if len(prompt_lengths) == 1
+            else f"batch={len(prompt_lengths)}, prompt tokens={prompt_lengths}"
+        )
+        print(
+            f"loading Step5 TP={world}, co-located EP={config_ep_size} "
+            f"({prompt_description})",
+            file=sys.stderr,
+            flush=True,
+        )
+    image_inputs = load_image_inputs(args.image_inputs, model_config)
+    model = _load_rank_model(args.tp_dir, rank, model_config, device)
+    completion_batch = generate_completions(
+        model,
+        prompt_batch,
+        max_new_tokens=args.max_new_tokens,
+        device=device,
+        eos_token_ids=eos_token_ids,
+        image_inputs=image_inputs,
+    )
+    if rank == 0:
+        print(
+            format_batch_output(
+                prompt_batch=prompt_batch,
+                completion_batch=completion_batch,
+                tokenizer=tokenizer,
+                eos_token_ids=eos_token_ids,
+            ),
+            flush=True,
+        )
+    return 0
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    args = parse_args(argv)
+    if not dist.is_available():
+        raise SystemExit("torch.distributed is unavailable; launch this CLI with torchrun")
+    if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"}.issubset(os.environ):
+        raise SystemExit(
+            "missing torchrun environment (RANK/WORLD_SIZE/LOCAL_RANK); "
+            "launch with torchrun --nproc-per-node=<TP>"
+        )
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    device = _resolve_device(args.device, local_rank)
+    backend = "nccl" if device.type == "cuda" else "gloo"
+    owns_process_group = False
+    try:
+        if not dist.is_initialized():
+            dist.init_process_group(backend=backend)
+            owns_process_group = True
+        return _run(
+            args,
+            rank=dist.get_rank(),
+            world=dist.get_world_size(),
+            device=device,
+        )
+    finally:
+        if owns_process_group and dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

inference/requirements.txt

@@ -0,0 +1,5 @@
+# Install a CUDA-enabled PyTorch build that matches the host driver.
+torch>=2.10,<2.14
+transformers>=4.57.6,<6
+safetensors>=0.7,<1
+triton>=3.6,<3.8

inference/test_vision.py

@@ -0,0 +1,390 @@
+"""Focused regressions for the standalone Step-5 vision integration.
+
+Run from inference/: python -m pytest -q test_vision.py
+"""
+import json
+from unittest.mock import patch
+
+import pytest
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file, save_file
+
+from convert import _shard_for, pre_shard
+from generate import _load_rank_config, _load_rank_model, load_image_inputs, parse_args, validate_prompt_ids
+from model import (
+    STEP5_TP_LAYOUT_VERSION, Step5Config, Step5ForCausalLM,
+    Step5ForConditionalGeneration, Step5VLConfig, config_from_dict,
+)
+from vision_encoder import EncoderRope2D, Step5VisionEncoder
+
+
+def tiny_config(**kwargs):
+    text = dict(
+        hidden_size=32, intermediate_size=64, num_hidden_layers=1,
+        num_nextn_predict_layers=3, vocab_size=64, num_attention_heads=4,
+        num_attention_groups=2, head_dim=8, max_position_embeddings=128,
+        layer_types=["sliding_attention"] * 4, sliding_window=16, yarn_only_types=["full_attention"],
+        use_moe=False, sparse_config={"enabled": False, "num_provider_groups": 2},
+        tp_layout_version=STEP5_TP_LAYOUT_VERSION,
+    )
+    text.update(kwargs.pop("text", {}))
+    vision = dict(width=16, layers=2, heads=2, image_size=32, patch_size=4,
+                  mlp_ratio=2, use_ln_pre=True, use_ln_post=False)
+    vision.update(kwargs.pop("vision", {}))
+    return Step5VLConfig(text_config=text, vision_config=vision,
+                         image_token_id=70, image_token_len=4, patch_token_len=1, **kwargs)
+
+
+def initialized_model(config):
+    model = Step5ForConditionalGeneration(config)
+    # TP linear layers allocate empty checkpoint storage; seed every parameter
+    # explicitly because these tests exercise random fixtures, not trained weights.
+    with torch.no_grad():
+        for name, parameter in model.named_parameters():
+            if "norm" in name and name.endswith("weight"):
+                parameter.zero_()
+            elif name.endswith("weight_scale_inv"):
+                parameter.uniform_(0.5, 1.5)
+            elif parameter.dtype == torch.float8_e4m3fn:
+                parameter.copy_(torch.randn(parameter.shape) * 0.02)
+            else:
+                parameter.normal_(0, 0.02)
+    return model
+
+
+def test_configuration_and_quantization_roundtrip():
+    quant = {"quant_method": "fp8", "modules_to_not_convert": [
+        "vision_model.transformer.resblocks.3.attn.q_proj", "model.layers.4.moe.gate",
+        "model.layers.88.moe.experts.w13_weight", "model.layers.89.moe.up_proj",
+    ]}
+    config = tiny_config(quantization_config=quant, tp_size=2)
+    assert config.text_config.bf16_expert_layers == [88, 89]
+    assert not hasattr(config, "quantization_config")
+    assert not hasattr(config.text_config, "quantization_config")
+    serialized = json.loads(json.dumps(config.to_dict()))
+    restored = config_from_dict(serialized)
+    assert restored.model_type == "step5_vl"
+    assert restored.architectures == ["Step5ForConditionalGeneration"]
+    assert restored.text_config.tp == 2
+    assert restored.text_config.num_nextn_predict_layers == 3
+    assert len(restored.text_config.to_dict()["layer_types"]) == 4
+    assert restored.text_config.bf16_expert_layers == [88, 89]
+    assert restored.to_dict()["quantization_config"] == quant
+    assert isinstance(config_from_dict(restored.text_config.to_dict()), Step5Config)
+
+
+def test_vision_shape_rectangular_cls_and_optional_layer_scale():
+    config = tiny_config(vision={"use_cls_token": True, "ls_init_value": None}).vision_config
+    encoder = Step5VisionEncoder(config).eval()
+    output = encoder(torch.randn(2, 3, 16, 24))
+    assert output.shape == (2, 24, 16)
+    assert torch.isfinite(output).all()
+    with pytest.raises(ValueError, match="positive multiples"):
+        encoder(torch.randn(1, 3, 15, 16))
+    with pytest.raises(ValueError, match="exceeds configured"):
+        encoder(torch.randn(1, 3, 36, 32))
+
+
+def test_rope_coordinates_and_cls_subset():
+    rope = EncoderRope2D(8, 8, 8, use_cls_token=True)
+    q = torch.ones(1, 1, 7, 8)
+    result, _ = rope(q, q, (2, 3))
+    assert torch.equal(result[:, :, 0], q[:, :, 0])
+    inv = 10000 ** (-torch.arange(0, 4, 2).float() / 4)
+    coords = [(0, 0)] + [(x + 1, y + 1) for y in range(2) for x in range(3)]
+    angles = torch.stack([torch.cat([x * inv, y * inv]).repeat_interleave(2) for x, y in coords])
+    sign = torch.tensor([-1., 1.] * 4)
+    expected = angles.cos() + sign * angles.sin()
+    torch.testing.assert_close(result[0, 0], expected)
+
+
+def test_crop_order_and_precomputed_embeddings():
+    model = Step5ForConditionalGeneration(tiny_config()).eval()
+    overview = torch.randn(2, 4, 32)
+    crops = torch.randn(3, 1, 32)
+    with patch.object(model, "_encode_images", side_effect=[overview, crops]):
+        actual = model.get_multimodal_embeddings(
+            pixel_values=torch.zeros(2, 3, 32, 32),
+            patch_pixel_values=torch.zeros(3, 3, 16, 16), num_patches=[2, 1],
+        )
+    expected = torch.cat([crops[0], crops[1], overview[0], crops[2], overview[1]])
+    torch.testing.assert_close(actual, expected)
+    assert model.get_multimodal_embeddings(image_embeds=overview).shape == (8, 32)
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        model.get_multimodal_embeddings(pixel_values=torch.empty(1), image_embeds=overview)
+    with pytest.raises(ValueError, match="sum\\(num_patches\\)"):
+        model.get_multimodal_embeddings(pixel_values=torch.zeros(2, 3, 32, 32), num_patches=[1, 0])
+    with pytest.raises(ValueError, match="nonnegative integer"):
+        model.get_multimodal_embeddings(pixel_values=torch.zeros(2, 3, 32, 32), num_patches=[0.5, 0])
+    with pytest.raises(ValueError, match="overview images"):
+        model.get_multimodal_embeddings(num_patches=[0])
+
+
+def make_checkpoint(path, config):
+    path.mkdir()
+    model = initialized_model(config)
+    state = model.state_dict()
+    for i in range(config.text_config.num_hidden_layers,
+                   config.text_config.num_hidden_layers + config.text_config.num_nextn_predict_layers):
+        state[f"model.layers.{i}.mtp_block.unused.weight"] = torch.ones(1)
+    save_file(state, str(path / "model.safetensors"))
+    (path / "config.json").write_text(json.dumps(config.to_dict()))
+    (path / "model.safetensors.index.json").write_text(json.dumps({
+        "weight_map": {key: "model.safetensors" for key in state},
+    }))
+    return model
+
+
+def test_conversion_and_strict_loader(tmp_path):
+    original = make_checkpoint(tmp_path / "source", tiny_config())
+    assert pre_shard(str(tmp_path / "source"), str(tmp_path / "unwritten"), 1, check_only=True) == 0
+    assert not (tmp_path / "unwritten").exists()
+    assert pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 1) == 0
+    directory = str(tmp_path / "converted" / "tp1")
+    config = _load_rank_config(directory, 0)
+    loaded = _load_rank_model(directory, 0, config, torch.device("cpu"))
+    assert isinstance(loaded, Step5ForConditionalGeneration)
+    for key, tensor in original.state_dict().items():
+        torch.testing.assert_close(loaded.state_dict()[key], tensor)
+    with torch.inference_mode():
+        pixels = torch.randn(1, 3, 32, 32)
+        # Nonpersistent 2D RoPE buffers must be rebuilt after HF meta loading.
+        expected_features = original.get_multimodal_embeddings(pixel_values=pixels)
+        actual_features = loaded.get_multimodal_embeddings(pixel_values=pixels)
+        assert torch.isfinite(actual_features).all()
+        torch.testing.assert_close(actual_features, expected_features, rtol=0, atol=0)
+    weight = original.vision_model.transformer.resblocks[0].attn.out_proj.weight
+    assert _shard_for("vision_model.transformer.resblocks.0.attn.out_proj.weight")(weight, None, 2, 1, config).shape == weight.shape
+
+
+@pytest.mark.parametrize("multimodal", [False, True])
+def test_checkpoint_identity_normalizes_to_step5(tmp_path, multimodal):
+    config = tiny_config()
+    payload = config.to_dict() if multimodal else config.text_config.to_dict()
+    payload["model_type"] = "step4_vl" if multimodal else "step4"
+    payload["architectures"] = ["Step4ForConditionalGeneration" if multimodal else "Step4ForCausalLM"]
+    payload["auto_map"] = {"AutoConfig": "old_config.Step4Config"}
+    if multimodal:
+        payload["text_config"]["model_type"] = "step4"
+        payload["text_config"]["architectures"] = ["Step4ForCausalLM"]
+        payload["text_config"]["auto_map"] = {"AutoConfig": "old_config.Step4Config"}
+    restored = config_from_dict(payload)
+    assert restored.model_type == ("step5_vl" if multimodal else "step5")
+    encoded = json.dumps(restored.to_dict())
+    assert "step4" not in encoded.lower()
+    restored.save_pretrained(tmp_path)
+    reloaded = config_from_dict(json.loads((tmp_path / "config.json").read_text()))
+    assert reloaded.to_dict()["auto_map"] == restored.to_dict()["auto_map"]
+    assert reloaded.to_dict()["architectures"] == restored.to_dict()["architectures"]
+    assert reloaded.get_text_config().num_nextn_predict_layers == 3
+    with pytest.raises(ValueError, match="unsupported model_type"):
+        config_from_dict(dict(payload, model_type="unrelated_model"))
+
+
+def test_legacy_checkpoint_conversion_and_shard_loading(tmp_path):
+    original = make_checkpoint(tmp_path / "source", tiny_config())
+    source_config = tmp_path / "source" / "config.json"
+    source_config.write_text(source_config.read_text().replace("step5", "step4").replace("Step5", "Step4"))
+    pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 1)
+    directory = tmp_path / "converted" / "tp1"
+    emitted = (directory / "config-r0.json").read_text()
+    assert "step4" not in emitted.lower()
+    shard = directory / "model-r0.safetensors"
+    state = load_file(str(shard))
+    with safe_open(str(shard), framework="pt") as handle:
+        metadata = handle.metadata()
+    assert metadata["step5_tp_layout"] == STEP5_TP_LAYOUT_VERSION
+    assert "step4_tp_layout" not in metadata
+    metadata["step4_tp_layout"] = metadata.pop("step5_tp_layout")
+    save_file(state, str(shard), metadata=metadata)
+    # Existing shards and their per-rank configs must remain loadable together.
+    (directory / "config-r0.json").write_text(emitted.replace("step5", "step4").replace("Step5", "Step4"))
+    config = _load_rank_config(str(directory), 0)
+    loaded = _load_rank_model(str(directory), 0, config, torch.device("cpu"))
+    for name, tensor in original.state_dict().items():
+        torch.testing.assert_close(loaded.state_dict()[name], tensor, rtol=0, atol=0)
+    metadata["step5_tp_layout"] = "incompatible-layout"
+    save_file(state, str(shard), metadata=metadata)
+    with pytest.raises(ValueError, match="step5_tp_layout metadata"):
+        _load_rank_model(str(directory), 0, config, torch.device("cpu"))
+    metadata["step5_tp_layout"] = STEP5_TP_LAYOUT_VERSION
+    metadata["step4_tp_layout"] = "incompatible-layout"
+    save_file(state, str(shard), metadata=metadata)
+    with pytest.raises(ValueError, match="step4_tp_layout metadata"):
+        _load_rank_model(str(directory), 0, config, torch.device("cpu"))
+    metadata.pop("step5_tp_layout")
+    metadata.pop("step4_tp_layout")
+    save_file(state, str(shard), metadata=metadata)
+    with pytest.raises(ValueError, match="missing step5_tp_layout"):
+        _load_rank_model(str(directory), 0, config, torch.device("cpu"))
+
+
+def tiny_moe_config(*, quantized=True):
+    return tiny_config(
+        text=dict(hidden_size=256, intermediate_size=256, num_hidden_layers=2,
+                  layer_types=["sliding_attention"] * 5, num_attention_heads=8,
+                  head_dim=32, vocab_size=128, eos_token_id=[127], use_moe=True,
+                  moe_layer_list=[0, 1], moe_num_experts=8, moe_top_k=2,
+                  moe_intermediate_size=256, share_expert_dim=256),
+        quantization_config={"quant_method": "fp8", "modules_to_not_convert": [
+            "model.layers.1.moe.experts.w13_weight",
+            "model.layers.1.moe.experts.w2_weight",
+            "vision_model.transformer.resblocks.0.attn.q_proj",
+        ]} if quantized else None,
+    )
+
+
+@pytest.mark.parametrize("quantized", [False, True], ids=["bf16", "mixed_fp8_bf16"])
+def test_experts_conversion_and_hf_loading(tmp_path, quantized):
+    original = make_checkpoint(tmp_path / "source", tiny_moe_config(quantized=quantized))
+    pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 2)
+    directory = str(tmp_path / "converted" / "tp2")
+    config = _load_rank_config(directory, 0)
+    assert config.text_config.bf16_expert_layers == ([1] if quantized else [0, 1])
+    loaded = _load_rank_model(directory, 0, config, torch.device("cpu"))
+    stored = load_file(str(tmp_path / "converted" / "tp2" / "model-r0.safetensors"))
+    for key, actual in loaded.state_dict().items():
+        expected = stored[key]
+        if key == "lm_head.weight":
+            expected = expected[:64]
+        assert actual.dtype == expected.dtype, key
+        torch.testing.assert_close(actual.float(), expected.float(), rtol=0, atol=0, msg=key)
+    if quantized:
+        assert loaded.model.layers[0].moe.gate_proj.weight.dtype == torch.float8_e4m3fn
+        assert loaded.model.layers[0].moe.gate_proj.weight_scale_inv.dtype == torch.float32
+    else:
+        assert loaded.model.layers[0].moe.gate_proj.weight.dtype == torch.bfloat16
+        assert not any("weight_scale" in key for key in stored)
+        assert loaded.model.layers[0].moe.gate_proj.weight_scale_inv is None
+    assert loaded.model.layers[1].moe.gate_proj.weight.dtype == torch.bfloat16
+    assert loaded.model.layers[1].moe.gate_proj.weight_scale_inv is None
+    # Independently verify that only the expert axis was partitioned.
+    torch.testing.assert_close(loaded.model.layers[0].moe.gate_proj.weight.float(),
+                               original.model.layers[0].moe.gate_proj.weight[:4].float(), rtol=0, atol=0)
+    with torch.inference_mode():
+        pixels = torch.randn(1, 3, 32, 32)
+        torch.testing.assert_close(loaded.get_multimodal_embeddings(pixel_values=pixels),
+                                   original.get_multimodal_embeddings(pixel_values=pixels), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("quantization", [None, {}])
+def test_unquantized_moe_configuration_roundtrip(quantization):
+    raw = tiny_moe_config(quantized=False).to_dict()
+    raw["text_config"]["quantization_config"] = quantization
+    # Derived fields from another checkpoint must not override its quantization metadata.
+    raw["text_config"]["bf16_expert_layers"] = []
+    config = config_from_dict(raw)
+    for restored in (config, config_from_dict(config.to_dict()),
+                     config_from_dict(config.text_config.to_dict())):
+        text = restored.get_text_config()
+        assert text.bf16_expert_layers == [0, 1]
+        assert not any(text.is_fp8_expert_layer(i) for i in range(2))
+        assert not hasattr(text, "quantization_config")
+
+
+def test_unsupported_expert_quantization_is_rejected():
+    with pytest.raises(ValueError, match="quant_method='fp8'"):
+        tiny_config(quantization_config={"quant_method": "int8"})
+
+
+@pytest.mark.parametrize("corruption", ["missing", "unexpected", "shape", "rank"])
+def test_loader_rejects_incomplete_or_wrong_rank_shards(tmp_path, corruption):
+    make_checkpoint(tmp_path / "source", tiny_config())
+    pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 1)
+    directory = tmp_path / "converted" / "tp1"
+    shard = directory / "model-r0.safetensors"
+    with safe_open(str(shard), framework="pt") as handle:
+        metadata = handle.metadata()
+    state = load_file(str(shard))
+    key = "vision_model.conv1.weight"
+    if corruption == "missing":
+        del state[key]
+    elif corruption == "unexpected":
+        state["vision_model.unexpected.weight"] = torch.ones(1)
+    elif corruption == "shape":
+        state[key] = state[key][:1].clone()
+    else:
+        metadata["rank"] = "1"
+    save_file(state, str(shard), metadata=metadata)
+    config = _load_rank_config(str(directory), 0)
+    with patch.object(Step5ForConditionalGeneration, "from_pretrained") as load:
+        with pytest.raises(ValueError, match="key mismatch|shape|rank metadata"):
+            _load_rank_model(str(directory), 0, config, torch.device("cpu"))
+        load.assert_not_called()
+
+
+def test_resume_rejects_old_text_only_shard(tmp_path):
+    make_checkpoint(tmp_path / "source", tiny_config())
+    pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 2)
+    # A matching shard must resume successfully when converting from a TP=1 config.
+    assert pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 2, start_rank=1) == 0
+    shard = tmp_path / "converted" / "tp2" / "model-r0.safetensors"
+    with safe_open(str(shard), framework="pt") as handle:
+        metadata = handle.metadata()
+    state = {key: tensor for key, tensor in load_file(str(shard)).items()
+             if not key.startswith(("vision_model.", "vit_large_projector."))}
+    save_file(state, str(shard), metadata=metadata)
+    with patch("convert._build_rank_state") as build:
+        with pytest.raises(ValueError, match="shard key mismatch"):
+            pre_shard(str(tmp_path / "source"), str(tmp_path / "converted"), 2, start_rank=1)
+        build.assert_not_called()
+
+
+def test_cli_inputs(tmp_path):
+    path = tmp_path / "images.safetensors"
+    save_file({"image_embeds": torch.randn(4, 32)}, str(path))
+    result = load_image_inputs(str(path), tiny_config())
+    assert result["image_embeds"].shape == (4, 32)
+    with pytest.raises(ValueError, match="step5_vl"):
+        load_image_inputs(str(path), tiny_config().text_config)
+    with pytest.raises(SystemExit):
+        parse_args(["--tp-dir", ".", "--prompt", "test", "--image-inputs", str(path)])
+    limits = dict(vocab_size=64, max_position_embeddings=128, max_new_tokens=1)
+    validate_prompt_ids([5, 70, 6], image_token_id=70, **limits)
+    with pytest.raises(ValueError, match="outside vocabulary"):
+        validate_prompt_ids([5, 70, 6], **limits)
+    with pytest.raises(ValueError, match="outside vocabulary"):
+        validate_prompt_ids([5, 71, 6], image_token_id=70, **limits)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for packed text inference")
+def test_cuda_prefill_decode_text_parity_and_single_vision_call():
+    torch.manual_seed(11)
+    device = torch.device("cuda")
+    config = tiny_config()
+    model = initialized_model(config).to(device).eval()
+    text = Step5ForCausalLM(config.text_config).to(device).eval()
+    text.load_state_dict({key: value for key, value in model.state_dict().items()
+                          if key.startswith(("model.", "lm_head."))}, strict=True)
+    pixels = torch.randn(1, 3, 32, 32, device=device)
+    ids = torch.tensor([5, 70, 70, 70, 70, 6, 7], device=device)
+    lengths = torch.tensor([7], dtype=torch.int32, device=device)
+    with torch.inference_mode():
+        cache_a = model.allocate_cache(batch=1, max_tokens=16, device=device)
+        cache_b = text.allocate_cache(batch=1, max_tokens=16, device=device)
+        features = model.get_multimodal_embeddings(pixel_values=pixels)
+        embeds = model.model.embed_tokens(ids.masked_fill(ids == 70, 0))
+        embeds[ids == 70] = features
+        actual = model(ids, cache_a, lengths, pixel_values=pixels).logits
+        assert torch.isfinite(actual).all()
+        expected = text(ids, cache_b, lengths, inputs_embeds=embeds).logits
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        token = actual.argmax(-1)
+        one = torch.ones(1, dtype=torch.int32, device=device)
+        torch.testing.assert_close(model(token, cache_a, one, is_decode=True).logits,
+                                   text(token, cache_b, one, is_decode=True).logits, rtol=0, atol=0)
+        before = cache_a.lengths.clone()
+        with pytest.raises(ValueError, match="placeholders"):
+            model(ids, cache_a, lengths, image_embeds=features[:1])
+        assert torch.equal(cache_a.lengths, before)
+        with pytest.raises(ValueError, match="only accepted during prefill"):
+            model(token, cache_a, one, is_decode=True, pixel_values=pixels)
+        with pytest.raises(ValueError, match="require image inputs"):
+            model(ids, cache_a, lengths)
+        with patch.object(model.vision_model, "forward", wraps=model.vision_model.forward) as encode:
+            output = model.generate_greedy([ids.tolist(), [8, 9]], max_new_tokens=3,
+                                          device=device, eos_token_ids={-1}, pixel_values=pixels)
+            assert [len(row) for row in output] == [3, 3]
+            assert encode.call_count == 1
+        assert text.generate_greedy([[5, 6]], max_new_tokens=2, device=device, eos_token_ids={-1}) == model.generate_greedy([[5, 6]], max_new_tokens=2, device=device, eos_token_ids={-1})

inference/validate_vision.py

@@ -0,0 +1,119 @@
+"""Validate complete Step-5 vision and projector outputs.
+
+Only vision/projector tensors are loaded. Text generation is covered separately
+by test_vision.py and the TP fixture smoke. No serving stack is required.
+"""
+from __future__ import annotations
+
+import argparse
+import copy
+import hashlib
+import importlib.util
+import json
+from pathlib import Path
+import sys
+import time
+import types
+
+import torch
+from safetensors import safe_open
+
+from model import Step5ForConditionalGeneration, config_from_dict
+from vision_encoder import Step5VisionConfig
+
+
+def load_reference(path):
+    # Supply the reference's config import without installing the full HF model.
+    # Its encoder mathematics are imported unchanged from the supplied file.
+    package = types.ModuleType("_step5_vision_reference")
+    package.__path__ = [str(Path(path).parent)]
+    sys.modules[package.__name__] = package
+    config_module = types.ModuleType(package.__name__ + ".configuration_step3p7")
+    config_module.StepRoboticsVisionEncoderConfig = Step5VisionConfig
+    sys.modules[config_module.__name__] = config_module
+    spec = importlib.util.spec_from_file_location(package.__name__ + ".vision_encoder", path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module.StepRoboticsVisionEncoder
+
+
+@torch.inference_mode()
+def validate(checkpoint, reference_encoder, device):
+    started = time.monotonic()
+    root = Path(checkpoint)
+    raw = json.loads((root / "config.json").read_text())
+    config = config_from_dict(raw)
+    # Keep the complete vision config and projection width. Avoid allocating the
+    # 92-layer language model for a vision-only numerical comparison.
+    scaffold = copy.deepcopy(raw)
+    scaffold["text_config"].update(
+        num_hidden_layers=0, num_nextn_predict_layers=0, layer_types=[],
+        use_moe=False, moe_layer_list=[], vocab_size=16,
+    )
+    scaffold_config = config_from_dict(scaffold)
+    reference_cls = load_reference(reference_encoder)
+    with torch.device(device):
+        model = Step5ForConditionalGeneration(scaffold_config).eval()
+        reference = reference_cls(config.vision_config).to(dtype=torch.bfloat16).eval()
+    index = json.loads((root / "model.safetensors.index.json").read_text())["weight_map"]
+    keys = sorted(k for k in index if k.startswith(("vision_model.", "vit_large_projector.")))
+    state = {}
+    for shard in sorted({index[k] for k in keys}):
+        with safe_open(root / shard, framework="pt", device="cpu") as handle:
+            for key in keys:
+                if index[key] == shard:
+                    state[key] = handle.get_tensor(key)
+    vision_state = {k.removeprefix("vision_model."): v for k, v in state.items() if k.startswith("vision_model.")}
+    projector_state = {k.removeprefix("vit_large_projector."): v for k, v in state.items() if k.startswith("vit_large_projector.")}
+    model.vision_model.load_state_dict(vision_state, strict=True)
+    model.vit_large_projector.load_state_dict(projector_state, strict=True)
+    reference.load_state_dict(vision_state, strict=True)
+    print(f"Loaded {len(keys)} real vision/projector tensors", flush=True)
+    del state, vision_state, projector_state
+    torch.manual_seed(1709)
+    results = []
+    for label, side, batch, expected_tokens in (
+        ("overview", config.vision_config.image_size, 1, config.image_token_len),
+        ("crops", 504, 2, config.patch_token_len),
+    ):
+        pixels = torch.randn(batch, 3, side, side, device=device, dtype=torch.bfloat16)
+        expected = reference(pixels)
+        actual = model.vision_model(pixels)
+        encoder_error = float((actual.float() - expected.float()).abs().max())
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        grid = side // config.vision_config.patch_size
+        expected = expected.transpose(1, 2).reshape(batch, -1, grid, grid)
+        expected = reference.vit_downsampler2(reference.vit_downsampler1(expected))
+        expected = model.vit_large_projector(expected.flatten(2).transpose(1, 2))
+        actual = model._encode_images(pixels)
+        assert actual.shape == (batch, expected_tokens, config.text_config.hidden_size)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        assert torch.isfinite(actual).all()
+        row = dict(case=label, input_shape=list(pixels.shape), output_shape=list(actual.shape),
+                   encoder_max_abs_error=encoder_error,
+                   projected_max_abs_error=float((actual.float() - expected.float()).abs().max()))
+        results.append(row)
+        print(json.dumps(row), flush=True)
+    torch.cuda.synchronize(device)
+    return dict(status="PASS", scope="complete real vision tower and projector; no full language-model generation",
+                checkpoint=str(root), config_sha256=hashlib.sha256((root / "config.json").read_bytes()).hexdigest(),
+                reference_sha256=hashlib.sha256(Path(reference_encoder).read_bytes()).hexdigest(),
+                tensor_count=len(keys), dtype="bfloat16", device=torch.cuda.get_device_name(device),
+                cases=results, elapsed_seconds=time.monotonic() - started)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--reference-encoder", required=True)
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--output", required=True)
+    args = parser.parse_args()
+    result = validate(args.checkpoint, args.reference_encoder, torch.device(args.device))
+    Path(args.output).write_text(json.dumps(result, indent=2) + "\n")
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

inference/vision_encoder.py

@@ -0,0 +1,500 @@
+# Copyright 2025 The StepFun and HuggingFace Inc. team. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+"""Step-5 perception encoder; parameter names match the released checkpoint."""
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.activations import ACT2FN
+from transformers.configuration_utils import PretrainedConfig
+
+
+class Step5VisionConfig(PretrainedConfig):
+    model_type = "perception_encoder"
+
+    def __init__(
+        self,
+        width=1536,
+        layers=47,
+        heads=16,
+        num_channels=3,
+        image_size=728,
+        mlp_ratio=8960 / 1536,
+        patch_size=14,
+        hidden_act="quick_gelu",
+        layer_norm_eps=1e-5,
+        ues_cls_token=False,
+        use_cls_token: Optional[bool] = None,
+        use_ln_pre=True,
+        use_ln_post=False,
+        use_abs_posemb=True,
+        use_rope2d=True,
+        ls_init_value=0.1,
+        **kwargs,
+    ):
+        self.width = width
+        self.layers = layers
+        self.heads = heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.mlp_ratio = mlp_ratio
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        if use_cls_token is None:
+            use_cls_token = ues_cls_token
+        self.ues_cls_token = use_cls_token
+        self.use_cls_token = use_cls_token
+        self.use_ln_pre = use_ln_pre
+        self.ls_init_value = ls_init_value
+        self.use_ln_post = use_ln_post
+        self.use_abs_posemb = use_abs_posemb
+        self.use_rope2d = use_rope2d
+        super().__init__(**kwargs)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate adjacent coordinate pairs for the encoder's interleaved RoPE."""
+    x = x.reshape(*x.shape[:-1], -1, 2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return x.reshape(*x.shape[:-2], -1)
+
+
+def apply_rotary_emb(freqs: torch.Tensor,
+                     t: torch.Tensor,
+                     start_index: int = 0,
+                     scale: float = 1.0,
+                     seq_dim: int = -2) -> torch.Tensor:
+    """Apply 2D rotary embeddings to queries / keys."""
+    dtype = t.dtype
+
+    if t.ndim == 3:
+        seq_len = t.shape[seq_dim]
+        freqs = freqs[-seq_len:]
+
+    rot_dim = freqs.shape[-1]
+    end_index = start_index + rot_dim
+    assert rot_dim <= t.shape[-1], (
+        f"feature dimension {t.shape[-1]} is too small for rot_dim {rot_dim}")
+
+    t_left, t, t_right = (
+        t[..., :start_index],
+        t[..., start_index:end_index],
+        t[..., end_index:],
+    )
+    t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
+    out = torch.cat((t_left, t, t_right), dim=-1)
+    return out.type(dtype)
+
+
+class EncoderRope2D(nn.Module):
+    """Cacheable 2D rotary positional embedding."""
+
+    def __init__(
+        self,
+        dim: int,
+        max_grid_height: int,
+        max_grid_width: int,
+        use_cls_token: bool = False,
+        theta: Union[int, float] = 10000,
+        theta_rescale_factor: float = 1.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_grid_height = max_grid_height
+        self.max_grid_width = max_grid_width
+        self.use_cls_token = use_cls_token
+        self.theta = theta * theta_rescale_factor**(dim / (dim - 2))
+        cache = self._compute_2d_freqs()
+        self.register_buffer("freqs_cache", cache, persistent=False)
+        self._cache_signature = None
+
+    def _compute_inv_freq(self, base: Union[int, float],
+                          dim: int, device=None) -> torch.Tensor:
+
+        freqs = 1.0 / (base**(
+            torch.arange(0, dim, 2, device=device)[:(dim // 2)].float() / dim))
+        return freqs
+
+    def _compute_freqs(self, t: torch.Tensor, inv_freq: torch.Tensor):
+        freqs = torch.einsum("..., f -> ... f", t.type(inv_freq.dtype),
+                             inv_freq)
+        freqs = freqs.repeat_interleave(2, dim=-1)
+        return freqs
+
+    def _compute_2d_freqs(self, device=None) -> torch.Tensor:
+        grid_h_range = torch.arange(self.max_grid_height, dtype=torch.float, device=device)
+        grid_w_range = torch.arange(self.max_grid_width, dtype=torch.float, device=device)
+        if self.use_cls_token:
+            grid_h_range += 1
+            grid_w_range += 1
+        inv_freq = self._compute_inv_freq(self.theta, self.dim // 2, device=device)
+        freqs_h = self._compute_freqs(grid_h_range, inv_freq)[:, None].expand(
+            self.max_grid_height, self.max_grid_width, -1)
+        freqs_w = self._compute_freqs(grid_w_range, inv_freq)[None, :].expand(
+            self.max_grid_height, self.max_grid_width, -1)
+        freqs = torch.cat([freqs_w, freqs_h], dim=-1).reshape(
+            self.max_grid_height * self.max_grid_width, -1)
+        if self.use_cls_token:
+            freqs = torch.cat([torch.zeros(1, freqs.shape[-1], device=device), freqs], dim=0)
+        freqs = freqs[None, None, ...]
+        return freqs
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor,
+                grid_hw: tuple[int, int]):
+        # HF meta loading may materialize nonpersistent buffers with empty
+        # storage. Rebuild from coordinates on first use and after device/dtype
+        # changes, independently of how checkpoint parameters were loaded.
+        signature = (q.device, q.dtype)
+        if self._cache_signature != signature:
+            self.freqs_cache = self._compute_2d_freqs(device=q.device).to(q.dtype)
+            self._cache_signature = signature
+        if not (0 < grid_hw[0] <= self.max_grid_height and 0 < grid_hw[1] <= self.max_grid_width):
+            raise ValueError(f"RoPE grid {grid_hw} exceeds configured grid {(self.max_grid_height, self.max_grid_width)}")
+        # If grid matches cached shape we reuse directly to avoid recomputation.
+        if grid_hw[0] != self.max_grid_height or grid_hw[1] != self.max_grid_width:
+            rows = torch.arange(grid_hw[0], device=q.device).view(-1, 1)
+            cols = torch.arange(grid_hw[1], device=q.device).view(1, -1)
+            positions = (rows * self.max_grid_width + cols).reshape(-1).to(
+                torch.long)
+            if self.use_cls_token:
+                positions = torch.cat(
+                    [torch.zeros(1, device=q.device, dtype=torch.long), positions + 1], dim=0)
+            freqs = self.freqs_cache.index_select(2, positions)
+        else:
+            freqs = self.freqs_cache
+        q = apply_rotary_emb(freqs, q)
+        k = apply_rotary_emb(freqs, k)
+        return q, k
+
+
+class EncoderLayerScale(nn.Module):
+    """Per-channel residual scaling used when ls_init_value is set."""
+
+    def __init__(self, dim: int, init_values: float):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.full((dim,), init_values))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:  # (B, L, D)
+        return hidden_states * self.gamma
+
+
+class EncoderMLP(nn.Module):
+    """Feed-forward network used inside each transformer block."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int,
+                 hidden_act: str):
+        super().__init__()
+        self.c_fc = nn.Linear(hidden_size, intermediate_size, bias=True)
+        self.act_fn = ACT2FN[hidden_act]
+        self.c_proj = nn.Linear(intermediate_size, hidden_size, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+
+        hidden_states = self.c_proj(self.act_fn(self.c_fc(hidden_states)))
+        return hidden_states
+
+
+class EncoderVisionAttention(nn.Module):
+    """Multi-head self attention with optional 2D RoPE."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        max_grid_height: int,
+        max_grid_width: int,
+        use_cls_token: bool = False,
+        use_rope2d: bool = True,
+        rope_theta: Union[int, float] = 10000,
+        rope_theta_rescale_factor: float = 1.0,
+        rope_freqs_for: str = "lang",
+    ):
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(
+                f"hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})."
+            )
+        if use_rope2d and (hidden_size // num_heads) % 4:
+            raise ValueError("2D RoPE requires head_dim divisible by four")
+        if use_rope2d and rope_freqs_for != "lang":
+            raise ValueError("Only rope_freqs_for='lang' is implemented")
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.scale = self.head_dim**-0.5
+        self.in_proj_weight = nn.Parameter(torch.zeros(hidden_size * 3, hidden_size))
+        self.in_proj_bias = nn.Parameter(torch.zeros(hidden_size * 3))
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True)
+
+        self.rope = None
+        if use_rope2d:
+            self.rope = EncoderRope2D(
+                dim=self.head_dim,
+                max_grid_height=max_grid_height,
+                max_grid_width=max_grid_width,
+                use_cls_token=use_cls_token,
+                theta=rope_theta,
+                theta_rescale_factor=rope_theta_rescale_factor,
+            )
+
+    def forward(self, hidden_states: torch.Tensor, grid_hw: tuple[int, int]) -> torch.Tensor:
+        bsz, seq_len, _ = hidden_states.shape
+        qkv = F.linear(
+            hidden_states,
+            self.in_proj_weight,
+            self.in_proj_bias,
+        )
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(bsz, seq_len, self.num_heads,
+                   self.head_dim).transpose(1, 2)
+        k = k.view(bsz, seq_len, self.num_heads,
+                   self.head_dim).transpose(1, 2)
+        if self.rope is not None:
+            q, k = self.rope(q, k, grid_hw=grid_hw)
+        v = v.view(bsz, seq_len, self.num_heads,
+                   self.head_dim).transpose(1, 2)
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v, is_causal=False, scale=self.scale)
+        attn_output = attn_output.transpose(1, 2).reshape(
+            bsz, seq_len, self.num_heads * self.head_dim)
+        return self.out_proj(attn_output)
+
+
+class EncoderVisionBlock(nn.Module):
+    """A single Vision Transformer block (self-attention + MLP)."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float,
+        hidden_act: str,
+        layer_norm_eps: float,
+        ls_init_value: Optional[float] = None,
+        max_grid_height: Optional[int] = None,
+        max_grid_width: Optional[int] = None,
+        use_cls_token: bool = False,
+        use_rope2d: bool = True,
+        rope_kwargs: Optional[dict] = None,
+    ):
+        super().__init__()
+        rope_kwargs = rope_kwargs or {}
+        self.attn = EncoderVisionAttention(
+            hidden_size,
+            num_heads,
+            max_grid_height=max_grid_height,
+            max_grid_width=max_grid_width,
+            use_cls_token=use_cls_token,
+            use_rope2d=use_rope2d,
+            **rope_kwargs,
+        )
+        self.ln_1 = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+
+        intermediate = int(hidden_size * mlp_ratio)
+        self.mlp = EncoderMLP(hidden_size, intermediate, hidden_act)
+
+        self.ls_1 = EncoderLayerScale(hidden_size, ls_init_value) if ls_init_value is not None else nn.Identity()
+        self.ls_2 = EncoderLayerScale(hidden_size, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+    def forward(self, hidden_states: torch.Tensor,
+                grid_hw: tuple[int, int]) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.ln_1(hidden_states)
+        hidden_states = self.attn(hidden_states, grid_hw=grid_hw)
+        hidden_states = residual + self.ls_1(hidden_states)
+
+        residual = hidden_states
+        hidden_states = self.ln_2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + self.ls_2(hidden_states)
+        return hidden_states
+
+
+class EncoderVisionTransformer(nn.Module):
+    """Stack of encoder blocks parameterised by Step5VisionConfig."""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        hidden_act: str,
+        layer_norm_eps: float,
+        ls_init_value: Optional[float] = None,
+        max_grid_height: Optional[int] = None,
+        max_grid_width: Optional[int] = None,
+        use_cls_token: bool = False,
+        use_rope2d: bool = True,
+        rope_kwargs: Optional[dict] = None,
+    ):
+        super().__init__()
+        rope_kwargs = rope_kwargs or {}
+        self.resblocks = nn.ModuleList([
+            EncoderVisionBlock(embed_dim, num_heads, mlp_ratio, hidden_act,
+                               layer_norm_eps,
+                               max_grid_height=max_grid_height,
+                               max_grid_width=max_grid_width,
+                               use_cls_token=use_cls_token,
+                               use_rope2d=use_rope2d,
+                               ls_init_value=ls_init_value,
+                               rope_kwargs=rope_kwargs)
+            for _ in range(depth)
+        ])
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                grid_hw: tuple[int, int]) -> torch.Tensor:
+        for block in self.resblocks:
+            hidden_states = block(hidden_states, grid_hw=grid_hw)
+        return hidden_states
+
+
+class Step5VisionEncoder(nn.Module):
+    """
+    Vision encoder built from Step5VisionConfig.
+
+    The encoder performs patch embedding followed by a stack of transformer
+    blocks. Only the config fields defined in Step5VisionConfig (and
+    Step5VLConfig.vision_config) are expected.
+    """
+
+    def __init__(self, config: Step5VisionConfig):
+        super().__init__()
+        self.config = config
+
+        self.hidden_size = config.width
+        self.num_heads = config.heads
+        self.num_hidden_layers = config.layers
+        self.patch_size = config.patch_size
+        self.image_size = config.image_size
+        self.use_cls_token = getattr(config, "use_cls_token", False)
+        self.use_rope2d = getattr(config, "use_rope2d", True)
+        self.use_abs_posemb = getattr(config, "use_abs_posemb", True)
+        self.layer_norm_eps = config.layer_norm_eps
+        self.mlp_ratio = getattr(config, "mlp_ratio", 8960 / 1536)
+        self.ls_init_value = getattr(config, "ls_init_value", None)
+        self.hidden_act = config.hidden_act
+        self.use_ln_pre = getattr(config, "use_ln_pre", False)
+        self.use_ln_post = getattr(config, "use_ln_post", True)
+
+        # Patch embedding.
+        self.conv1 = nn.Conv2d(in_channels=config.num_channels,
+                               out_channels=self.hidden_size,
+                               kernel_size=self.patch_size,
+                               stride=self.patch_size,
+                               bias=False)
+
+        self.ln_pre = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps) if self.use_ln_pre else nn.Identity()
+        self.ln_post = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps) if self.use_ln_post else nn.Identity()
+
+        grid_size = self.image_size // self.patch_size
+        self.base_grid = (grid_size, grid_size)
+
+        if self.use_cls_token:
+            self.class_embedding = nn.Parameter(
+                torch.randn(self.hidden_size) * (self.hidden_size**-0.5))
+        else:
+            self.class_embedding = None
+
+        if self.use_abs_posemb:
+            self.posemb_grid_size = self.image_size // self.patch_size
+            self.positional_embedding = nn.Parameter(
+                (self.hidden_size**-0.5) * torch.randn(
+                    int(self.use_cls_token) + self.posemb_grid_size**2,
+                    self.hidden_size,
+                ))
+
+        self.transformer = EncoderVisionTransformer(
+            embed_dim=self.hidden_size,
+            depth=self.num_hidden_layers,
+            num_heads=self.num_heads,
+            mlp_ratio=self.mlp_ratio,
+            hidden_act=self.hidden_act,
+            layer_norm_eps=self.layer_norm_eps,
+            ls_init_value=self.ls_init_value,
+            max_grid_height=self.base_grid[0],
+            max_grid_width=self.base_grid[1],
+            use_cls_token=self.use_cls_token,
+            use_rope2d=self.use_rope2d,
+            rope_kwargs={
+                "rope_theta": getattr(config, "rope_theta", 10000),
+                "rope_theta_rescale_factor":
+                getattr(config, "rope_theta_rescale_factor", 1.0),
+                "rope_freqs_for": getattr(config, "rope_freqs_for", "lang"),
+            },
+        )
+        self.vit_downsampler1 = nn.Conv2d(self.hidden_size,
+                                          self.hidden_size * 2,
+                                          kernel_size=3,
+                                          stride=2,
+                                          padding=1)
+        self.vit_downsampler2 = nn.Conv2d(self.hidden_size * 2,
+                                          self.hidden_size * 4,
+                                          kernel_size=3,
+                                          stride=2,
+                                          padding=1)
+
+
+    def sample_abs_posemb(self, grid_h: int, grid_w: int):
+        if self.posemb_grid_size == grid_h and self.posemb_grid_size == grid_w:
+            return self.positional_embedding[None, ...]
+
+        pos_embed = self.positional_embedding
+        if self.use_cls_token:
+            cls_token_embed, pos_embed = pos_embed[:1], pos_embed[1:]
+
+        pos_embed = (pos_embed.reshape(1, self.posemb_grid_size,
+                                       self.posemb_grid_size,
+                                       -1).permute(0, 3, 1, 2).contiguous())
+        pos_embed = F.interpolate(pos_embed,
+                                  size=(grid_h, grid_w),
+                                  mode="bilinear",
+                                  align_corners=False)
+        pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
+
+        if self.use_cls_token:
+            pos_embed = torch.cat([cls_token_embed, pos_embed], dim=0)
+
+        return pos_embed[None, ...]
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            pixel_values: Image tensor of shape (B, C, H, W).
+        """
+        if pixel_values.ndim != 4 or pixel_values.shape[1] != self.config.num_channels:
+            raise ValueError("pixel_values must have shape [images, channels, height, width]")
+        bsz, _, height, width = pixel_values.shape
+        if not bsz or min(height, width) < self.patch_size or height % self.patch_size or width % self.patch_size:
+            raise ValueError("Image dimensions must be positive multiples of patch_size, with a nonempty batch")
+        grid_h, grid_w = height // self.patch_size, width // self.patch_size
+
+        hidden_state = self.conv1(pixel_values)  # (B, D, Gh, Gw)
+        hidden_state = hidden_state.flatten(2).transpose(1, 2)  # (B, Gh*Gw, D)
+
+        if self.use_cls_token:
+            cls_token = self.class_embedding.view(1, 1,
+                                                  -1).expand(bsz, -1, -1)
+            hidden_state = torch.cat([cls_token, hidden_state], dim=1)
+
+        if self.use_abs_posemb:
+            pos_emb = self.sample_abs_posemb(grid_h, grid_w)
+            hidden_state = hidden_state + pos_emb
+        hidden_state = self.ln_pre(hidden_state)
+        hidden_state = self.transformer(hidden_state, grid_hw=(grid_h, grid_w))
+
+        if self.use_ln_post:
+            hidden_state = self.ln_post(hidden_state)
+
+        if self.use_cls_token:
+            hidden_state = hidden_state[:, 1:, :]
+
+        return hidden_state